diff --git a/.githooks/prepare-commit-msg b/.githooks/prepare-commit-msg new file mode 100755 index 0000000000000..a091376b8c0d8 --- /dev/null +++ b/.githooks/prepare-commit-msg @@ -0,0 +1,147 @@ +#!/usr/bin/env python + + +# 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. + +# A hook script to prepare the commit log message. +# Automatically Called by "git commit" (including git cherry-pick). +# The hook's purpose is to edit the commit +# message. If the hook fails with a non-zero status, +# the commit is aborted. Users are allowed to further edit commit message +# generated by this hook. + +# git examples to invoke the hook: +# +# Cherry-pick (please use -x): +# git cherry-pick -e -x +# +# Hotfix: +# git commit ... + +import sys +import re + +# 0: Hotfix 1: cherry-pick +commit_type = 0 +keywords = ["[LI-HOTFIX]", "[LI-CHERRY-PICK]"] + +commit_msg_file = sys.argv[1] +commit_source = None +if (len(sys.argv) >= 3): + commit_source = sys.argv[2] + +SHORT_HASH_LEN = 8 +Original_Description_Key = "ORIGINAL_DESCRIPTION" +LI_Description_Key = "LI_DESCRIPTION" +LI_Description_Value = [""] +Ticket_Key = "TICKET" +Exit_Criteria_Key = "EXIT_CRITERIA" +Exit_Criteria_Value = ["MANUAL [\"\"]"] +Exit_Criteria_Help = ["# EXIT_CRITERIA = \n", "# e.g., \n", \ + "# when the specified hash(s) is presented in the history, this commit is no longer needed:\n# EXIT_CRITERIA = HASH [, ...] \n", \ + "# When the specified tickets are closed and there are patches with these tickets in the title in the commit history, this commit is no longer needed:\n# EXIT_CRITERIA = TICKET [, ...]\n", \ + "# The exit criteria for this commit requires manual investigation:\n# EXIT_CRITERIA = MANUAL []\n"] + +def writeKeyValue(fh, key, value, multiline): + fh.write(key + " = ") + if (multiline): + fh.write("\n") + for line in value: + fh.write(line) + fh.write("\n") + return + +with open (commit_msg_file, "r") as fh: + oldMsg = fh.readlines() + +firstline = "" + +if (len(oldMsg) > 0): + firstline = oldMsg[0].lstrip() + +# skip already formatted commits +for word in keywords: + if (firstline.startswith(word)): + sys.exit(0) + +actual_msg = [] +comment_msg = [] +for line in oldMsg[1:]: + if (line.startswith("#")): + comment_msg.append(line) + else: + actual_msg.append(line) + +for line in oldMsg: + cherry_pick_hash = re.search(r"\(cherry picked from commit ([^)]+)\)", line) + if (cherry_pick_hash is not None): + break +ticket = "" +if (cherry_pick_hash is not None): + hash = cherry_pick_hash.group(1) + print ("hash:" + hash) + + commit_type = 1 + + match_ticket = re.search(r"(KAFKA-\d+)", firstline, re.I) + + if (match_ticket is not None): + ticket = match_ticket.group(1) + print(ticket) + + if (len(hash) > SHORT_HASH_LEN): + short_hash = hash[:SHORT_HASH_LEN] + else: + short_hash = hash + title = keywords[commit_type] + " [" + short_hash + "] " + firstline + Exit_Criteria_Value = "HASH [" + hash + "]" + +else: + if not firstline: + firstline = "\n" + title = keywords[commit_type] + " " + firstline + +# populate formatted commit message +with open(commit_msg_file, 'w') as f: + f.write(title) + writeKeyValue(f, Ticket_Key, [ticket], False) + + if (commit_type == 0 and len(actual_msg) > 0): + LI_Description_Value = actual_msg + writeKeyValue(f, LI_Description_Key, LI_Description_Value, True) + + writeKeyValue(f, Exit_Criteria_Key, Exit_Criteria_Value, False) + + if commit_source is None: + for line in Exit_Criteria_Help: + f.write(line) + + # for cherry-pick, preserve the original commit message + if (commit_type == 1): + writeKeyValue(f, Original_Description_Key, actual_msg, True) + + for line in comment_msg: + f.write(line) + +sys.exit(0) + + + + + + + + diff --git a/.githooks/setup.sh b/.githooks/setup.sh new file mode 100755 index 0000000000000..57cb8432b1805 --- /dev/null +++ b/.githooks/setup.sh @@ -0,0 +1,17 @@ +#!/bin/bash +# 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. + +cp $(dirname $0)/prepare-commit-msg $(dirname $0)/../.git/hooks/ diff --git a/.githooks/uninstall.sh b/.githooks/uninstall.sh new file mode 100755 index 0000000000000..d99c67bb674b0 --- /dev/null +++ b/.githooks/uninstall.sh @@ -0,0 +1,17 @@ +#!/bin/bash +# 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. + +rm $(dirname $0)/../.git/hooks/prepare-commit-msg diff --git a/.gitignore b/.gitignore index a31643f00af6e..ff7f9f07b5208 100644 --- a/.gitignore +++ b/.gitignore @@ -34,8 +34,7 @@ Vagrantfile.local config/server-* config/zookeeper-* core/data/* -gradle/wrapper/* -gradlew +gradle/wrapper/*.jar gradlew.bat results @@ -53,3 +52,7 @@ kafkatest.egg-info/ systest/ *.swp clients/src/generated +clients/src/generated-test +streams/src/generated/ +jmh-benchmarks/generated +jmh-benchmarks/src/main/generated diff --git a/.travis.yml b/.travis.yml index 8a22c9bc4cd4a..4e9f19ce4f14c 100644 --- a/.travis.yml +++ b/.travis.yml @@ -14,41 +14,33 @@ sudo: required dist: trusty language: java -env: - - _DUCKTAPE_OPTIONS="--subset 0 --subsets 15" - - _DUCKTAPE_OPTIONS="--subset 1 --subsets 15" - - _DUCKTAPE_OPTIONS="--subset 2 --subsets 15" - - _DUCKTAPE_OPTIONS="--subset 3 --subsets 15" - - _DUCKTAPE_OPTIONS="--subset 4 --subsets 15" - - _DUCKTAPE_OPTIONS="--subset 5 --subsets 15" - - _DUCKTAPE_OPTIONS="--subset 6 --subsets 15" - - _DUCKTAPE_OPTIONS="--subset 7 --subsets 15" - - _DUCKTAPE_OPTIONS="--subset 8 --subsets 15" - - _DUCKTAPE_OPTIONS="--subset 9 --subsets 15" - - _DUCKTAPE_OPTIONS="--subset 10 --subsets 15" - - _DUCKTAPE_OPTIONS="--subset 11 --subsets 15" - - _DUCKTAPE_OPTIONS="--subset 12 --subsets 15" - - _DUCKTAPE_OPTIONS="--subset 13 --subsets 15" - - _DUCKTAPE_OPTIONS="--subset 14 --subsets 15" - jdk: - oraclejdk8 + +# By default, the install step will run ./gradlew assemble, which actually builds. Skip this step and do the actual +# build during the build step. +install: true -before_install: - - gradle wrapper - +# exclude the streams test and connect test script: - - ./gradlew rat - - ./gradlew systemTestLibs && /bin/bash ./tests/docker/run_tests.sh - -services: - - docker + - ./gradlew cleanTest :clients:test :core:test --no-daemon -PxmlFindBugsReport=true -PtestLoggingEvents=started,passed,skipped,failed before_cache: - rm -f $HOME/.gradle/caches/modules-2/modules-2.lock - rm -fr $HOME/.gradle/caches/*/plugin-resolution/ + cache: directories: - "$HOME/.m2/repository" - "$HOME/.gradle/caches/" - "$HOME/.gradle/wrapper/" + +# This will triger publishing artifacts to bintray upon the creation of a new tag in the "x.y.z.w" format. +deploy: + provider: script + script: ./gradlew -Pversion=$TRAVIS_TAG :clients:uploadArchives :core:uploadArchives --no-daemon + skip_cleanup: true + on: + tags: true + all_branches: true + condition: $TRAVIS_TAG =~ ^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+$ diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index e9a4d25e87ed5..07fd85711e005 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -1,6 +1,6 @@ ## Contributing to Kafka -*Before opening a pull request*, review the [Contributing](http://kafka.apache.org/contributing.html) and [Contributing Code Changes](https://cwiki.apache.org/confluence/display/KAFKA/Contributing+Code+Changes) pages. +*Before opening a pull request*, review the [Contributing](https://kafka.apache.org/contributing.html) and [Contributing Code Changes](https://cwiki.apache.org/confluence/display/KAFKA/Contributing+Code+Changes) pages. It lists steps that are required before creating a PR. diff --git a/NOTICE b/NOTICE index 8a9b248cfa7ec..f5505f39ad215 100644 --- a/NOTICE +++ b/NOTICE @@ -1,8 +1,8 @@ Apache Kafka -Copyright 2019 The Apache Software Foundation. +Copyright 2020 The Apache Software Foundation. This product includes software developed at -The Apache Software Foundation (http://www.apache.org/). +The Apache Software Foundation (https://www.apache.org/). This distribution has a binary dependency on jersey, which is available under the CDDL License. The source code of jersey can be found at https://github.com/jersey/jersey/. diff --git a/README.md b/README.md index a1e742dc2128d..55c3212a03ec7 100644 --- a/README.md +++ b/README.md @@ -1,25 +1,60 @@ -Apache Kafka +LinkedIn Branch of Apache Kafka ================= -See our [web site](http://kafka.apache.org) for details on the project. -You need to have [Gradle](http://www.gradle.org/installation) and [Java](http://www.oracle.com/technetwork/java/javase/downloads/index.html) installed. +This is the version of Kafka running at LinkedIn. -Kafka requires Gradle 5.0 or higher. +Kafka was born at LinkedIn. We run thousands of brokers to deliver trillions of +messages per day. We run a slightly modified version of Apache Kafka trunk. +This branch contains the LinkedIn Kafka release. -Java 8 should be used for building in order to support both Java 8 and Java 11 at runtime. +This branch is made up of: -Scala 2.12 is used by default, see below for how to use a different Scala version or all of the supported Scala versions. +* Apache Kafka trunk (upstream) up to some branch point, see *-li* branch name for base version, you'll be able to get the exact commit from git +* Cherry-picked commits from upstream after branch point +* Patches that are on their way upstream but we have deployed internally in the meantime +* Patches that are of no interest to upstream + +We are making this branch available for people interested. We will be +documenting the changes in the near future with some more detailed explanations +in the [LinkedIn Engineering Blog](https://engineering.linkedin.com/blog). + +If you are interested in learning more, we invite you to our [Streaming +Meetup](https://www.meetup.com/Stream-Processing-Meetup-LinkedIn/) where we +discuss streaming technologies like [Kafka](http://kafka.apache.org) and +[Samza](http://samza.apache.org). + +You are encouraged to check out other Kafka projects from LinkedIn: + +* [Cruise Control](https://github.com/linkedin/cruise-control) +* [Li-Apache-Kafka-Clients](https://github.com/linkedin/li-apache-kafka-clients) +* [Burrow](https://github.com/linkedin/Burrow) +* [Kafka Monitor](https://github.com/linkedin/kafka-monitor) + +### Contributing ### + +At this moment we are not accepting external contributions directly. Please +contribute to [Apache Kafka](http://kafka.apache.org). + +For security issues with this branch please review +[LinkedIn Security +Guidelines](https://www.linkedin.com/help/linkedin/answer/62924/security-vulnerabilities?lang=en). +General Kafka issues should be communicated via the Kafka community. + + +Apache Kafka +================= +See our [web site](https://kafka.apache.org) for details on the project. -### First bootstrap and download the wrapper ### - cd kafka_source_dir - gradle +You need to have [Java](http://www.oracle.com/technetwork/java/javase/downloads/index.html) installed. -Now everything else will work. +Java 8 should be used for building in order to support both Java 8 and Java 11 at runtime. + +Scala 2.12 is used by default, see below for how to use a different Scala version or all of the supported Scala versions. ### Build a jar and run it ### ./gradlew jar -Follow instructions in http://kafka.apache.org/documentation.html#quickstart +Follow instructions in https://kafka.apache.org/documentation.html#quickstart ### Build source jar ### ./gradlew srcJar @@ -77,7 +112,7 @@ The release file can be found inside `./core/build/distributions/`. ### Cleaning the build ### ./gradlew clean -### Running a task with a particular version of Scala (either 2.11.x or 2.12.x) ### +### Running a task with one of the Scala versions available (2.11.x, 2.12.x or 2.13.x) ### *Note that if building the jars with a version other than 2.12.x, you need to set the `SCALA_VERSION` variable or change it in `bin/kafka-run-class.sh` to run the quick start.* You can pass either the major version (eg 2.12) or the full version (eg 2.12.7): @@ -86,7 +121,7 @@ You can pass either the major version (eg 2.12) or the full version (eg 2.12.7): ./gradlew -PscalaVersion=2.12 test ./gradlew -PscalaVersion=2.12 releaseTarGz -### Running a task with all scala versions ### +### Running a task with all the scala versions enabled by default ### Append `All` to the task name: @@ -114,13 +149,20 @@ build directory (`${project_dir}/bin`) clashes with Kafka's scripts directory an to avoid known issues with this configuration. ### Publishing the jar for all version of Scala and for all projects to maven ### - ./gradlew uploadArchivesAll + ./gradlew -Pversion= uploadArchivesAll + +By default, this command will publish artifacts to a Bintray repository named "kafka" under an account specified by the BINTRAY_USER environment variable. The BINTRAY_KEY +environment variable is used for the password for that account. -Please note for this to work you should create/update `${GRADLE_USER_HOME}/gradle.properties` (typically, `~/.gradle/gradle.properties`) and assign the following variables +If you want to override this to use a different maven repository, you should create/update `${GRADLE_USER_HOME}/gradle.properties` (typically, `~/.gradle/gradle.properties`) +and assign the following variables mavenUrl= mavenUsername= mavenPassword= + +Signing is disabled by default. If you need signing, please set the following variables in `gradle.properties` as well: + signing.keyId= signing.password= signing.secretKeyRingFile= @@ -200,6 +242,22 @@ The following options should be set with a `-P` switch, for example `./gradlew - * `testLoggingEvents`: unit test events to be logged, separated by comma. For example `./gradlew -PtestLoggingEvents=started,passed,skipped,failed test`. * `xmlSpotBugsReport`: enable XML reports for spotBugs. This also disables HTML reports as only one can be enabled at a time. +### Dependency Analysis ### + +The gradle [dependency debugging documentation](https://docs.gradle.org/current/userguide/viewing_debugging_dependencies.html) mentions using the `dependencies` or `dependencyInsight` tasks to debug dependencies for the root project or individual subprojects. + +Alternatively, use the `allDeps` or `allDepInsight` tasks for recursively iterating through all subprojects: + + ./gradlew allDeps + + ./gradlew allDepInsight --configuration runtime --dependency com.fasterxml.jackson.core:jackson-databind + +These take the same arguments as the builtin variants. + +### Running system tests ### + +See [tests/README.md](tests/README.md). + ### Running in Vagrant ### See [vagrant/README.md](vagrant/README.md). @@ -209,4 +267,4 @@ See [vagrant/README.md](vagrant/README.md). Apache Kafka is interested in building the community; we would welcome any thoughts or [patches](https://issues.apache.org/jira/browse/KAFKA). You can reach us [on the Apache mailing lists](http://kafka.apache.org/contact.html). To contribute follow the instructions here: - * http://kafka.apache.org/contributing.html + * https://kafka.apache.org/contributing.html diff --git a/TROGDOR.md b/TROGDOR.md index b551773ab98ab..3891857562f93 100644 --- a/TROGDOR.md +++ b/TROGDOR.md @@ -16,7 +16,7 @@ Running Kafka: > ./bin/kafka-server-start.sh ./config/server.properties &> /tmp/kafka.log & -Then, we want to run a Trogdor Agent, plus a Trogdor broker. +Then, we want to run a Trogdor Agent, plus a Trogdor Coordinator. To run the Trogdor Agent: @@ -46,6 +46,7 @@ We can run showTask to see what the task's status is: Task bar of type org.apache.kafka.trogdor.workload.ProduceBenchSpec is DONE. FINISHED at 2019-01-09T20:38:22.039-08:00 after 6s To see the results, we use showTask with --show-status: + > ./bin/trogdor.sh client showTask -t localhost:8889 -i produce0 --show-status Task bar of type org.apache.kafka.trogdor.workload.ProduceBenchSpec is DONE. FINISHED at 2019-01-09T20:38:22.039-08:00 after 6s Status: { @@ -86,6 +87,36 @@ The task specification is usually written as JSON. For example, this task speci "durationMs": 30000, "partitions": [["node1", "node2"], ["node3"]] } + +This task runs a simple ProduceBench test on a cluster with one producer node, 5 topics, and 10,000 messages per second. +The keys are generated sequentially and the configured partitioner (DefaultPartitioner) is used. + + { + "class": "org.apache.kafka.trogdor.workload.ProduceBenchSpec", + "durationMs": 10000000, + "producerNode": "node0", + "bootstrapServers": "localhost:9092", + "targetMessagesPerSec": 10000, + "maxMessages": 50000, + "activeTopics": { + "foo[1-3]": { + "numPartitions": 10, + "replicationFactor": 1 + } + }, + "inactiveTopics": { + "foo[4-5]": { + "numPartitions": 10, + "replicationFactor": 1 + } + }, + "keyGenerator": { + "type": "sequential", + "size": 8, + "offset": 1 + }, + "useConfiguredPartitioner": true + } Tasks are submitted to the coordinator. Once the coordinator determines that it is time for the task to start, it creates workers on agent processes. The workers run until the task is done. @@ -105,6 +136,7 @@ Trogdor can run several workloads. Workloads perform operations on the cluster ### ProduceBench ProduceBench starts a Kafka producer on a single agent node, producing to several partitions. The workload measures the average produce latency, as well as the median, 95th percentile, and 99th percentile latency. +It can be configured to use a transactional producer which can commit transactions based on a set time interval or number of messages. ### RoundTripWorkload RoundTripWorkload tests both production and consumption. The workload starts a Kafka producer and consumer on a single node. The consumer will read back the messages that were produced by the producer. @@ -123,6 +155,27 @@ ProcessStopFault stops a process by sending it a SIGSTOP signal. When the fault ### NetworkPartitionFault NetworkPartitionFault sets up an artificial network partition between one or more sets of nodes. Currently, this is implemented using iptables. The iptables rules are set up on the outbound traffic from the affected nodes. Therefore, the affected nodes should still be reachable from outside the cluster. +External Processes +======================================== +Trogdor supports running arbitrary commands in external processes. This is a generic way to run any configurable command in the Trogdor framework - be it a Python program, bash script, docker image, etc. + +### ExternalCommandWorker +ExternalCommandWorker starts an external command defined by the ExternalCommandSpec. It essentially allows you to run any command on any Trogdor agent node. +The worker communicates with the external process via its stdin, stdout and stderr in a JSON protocol. It uses stdout for any actionable communication and only logs what it sees in stderr. +On startup the worker will first send a message describing the workload to the external process in this format: +``` +{"id":, "workload":} +``` +and will then listen for messages from the external process, again in a JSON format. +Said JSON can contain the following fields: +- status: If the object contains this field, the status of the worker will be set to the given value. +- error: If the object contains this field, the error of the worker will be set to the given value. Once an error occurs, the external process will be terminated. +- log: If the object contains this field, a log message will be issued with this text. +An example: +```json +{"log": "Finished successfully.", "status": {"p99ProduceLatency": "100ms", "messagesSent": 10000}} +``` + Exec Mode ======================================== Sometimes, you just want to run a test quickly on a single node. In this case, you can use "exec mode." This mode allows you to run a single Trogdor Agent without a Coordinator. diff --git a/bin/connect-mirror-maker.sh b/bin/connect-mirror-maker.sh new file mode 100755 index 0000000000000..a2c040dad297e --- /dev/null +++ b/bin/connect-mirror-maker.sh @@ -0,0 +1,45 @@ +#!/bin/sh +# 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. + +if [ $# -lt 1 ]; +then + echo "USAGE: $0 [-daemon] mm2.properties" + exit 1 +fi + +base_dir=$(dirname $0) + +if [ "x$KAFKA_LOG4J_OPTS" = "x" ]; then + export KAFKA_LOG4J_OPTS="-Dlog4j.configuration=file:$base_dir/../config/connect-log4j.properties" +fi + +if [ "x$KAFKA_HEAP_OPTS" = "x" ]; then + export KAFKA_HEAP_OPTS="-Xms256M -Xmx2G" +fi + +EXTRA_ARGS=${EXTRA_ARGS-'-name mirrorMaker'} + +COMMAND=$1 +case $COMMAND in + -daemon) + EXTRA_ARGS="-daemon "$EXTRA_ARGS + shift + ;; + *) + ;; +esac + +exec $(dirname $0)/kafka-run-class.sh $EXTRA_ARGS org.apache.kafka.connect.mirror.MirrorMaker "$@" diff --git a/bin/kafka-leader-election.sh b/bin/kafka-leader-election.sh new file mode 100755 index 0000000000000..88baef398de95 --- /dev/null +++ b/bin/kafka-leader-election.sh @@ -0,0 +1,17 @@ +#!/bin/bash +# 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. + +exec $(dirname $0)/kafka-run-class.sh kafka.admin.LeaderElectionCommand "$@" diff --git a/bin/kafka-run-class.sh b/bin/kafka-run-class.sh index c60419bb4733e..018e52f8388fa 100755 --- a/bin/kafka-run-class.sh +++ b/bin/kafka-run-class.sh @@ -20,7 +20,7 @@ then exit 1 fi -# CYGINW == 1 if Cygwin is detected, else 0. +# CYGWIN == 1 if Cygwin is detected, else 0. if [[ $(uname -a) =~ "CYGWIN" ]]; then CYGWIN=1 else @@ -48,7 +48,7 @@ should_include_file() { base_dir=$(dirname $0)/.. if [ -z "$SCALA_VERSION" ]; then - SCALA_VERSION=2.12.8 + SCALA_VERSION=2.12.10 fi if [ -z "$SCALA_BINARY_VERSION" ]; then @@ -57,10 +57,12 @@ fi # run ./gradlew copyDependantLibs to get all dependant jars in a local dir shopt -s nullglob -for dir in "$base_dir"/core/build/dependant-libs-${SCALA_VERSION}*; -do - CLASSPATH="$CLASSPATH:$dir/*" -done +if [ -z "$UPGRADE_KAFKA_STREAMS_TEST_VERSION" ]; then + for dir in "$base_dir"/core/build/dependant-libs-${SCALA_VERSION}*; + do + CLASSPATH="$CLASSPATH:$dir/*" + done +fi for file in "$base_dir"/examples/build/libs/kafka-examples*.jar; do @@ -110,6 +112,14 @@ else CLASSPATH="$file":"$CLASSPATH" fi done + if [ "$SHORT_VERSION_NO_DOTS" = "0100" ]; then + CLASSPATH="/opt/kafka-$UPGRADE_KAFKA_STREAMS_TEST_VERSION/libs/zkclient-0.8.jar":"$CLASSPATH" + CLASSPATH="/opt/kafka-$UPGRADE_KAFKA_STREAMS_TEST_VERSION/libs/zookeeper-3.4.6.jar":"$CLASSPATH" + fi + if [ "$SHORT_VERSION_NO_DOTS" = "0101" ]; then + CLASSPATH="/opt/kafka-$UPGRADE_KAFKA_STREAMS_TEST_VERSION/libs/zkclient-0.9.jar":"$CLASSPATH" + CLASSPATH="/opt/kafka-$UPGRADE_KAFKA_STREAMS_TEST_VERSION/libs/zookeeper-3.4.8.jar":"$CLASSPATH" + fi fi for file in "$rocksdb_lib_dir"/rocksdb*.jar; @@ -129,7 +139,7 @@ do CLASSPATH="$CLASSPATH:$dir/*" done -for cc_pkg in "api" "transforms" "runtime" "file" "json" "tools" "basic-auth-extension" +for cc_pkg in "api" "transforms" "runtime" "file" "mirror" "mirror-client" "json" "tools" "basic-auth-extension" do for file in "$base_dir"/connect/${cc_pkg}/build/libs/connect-${cc_pkg}*.jar; do @@ -238,13 +248,6 @@ if [ -z "$KAFKA_JVM_PERFORMANCE_OPTS" ]; then KAFKA_JVM_PERFORMANCE_OPTS="-server -XX:+UseG1GC -XX:MaxGCPauseMillis=20 -XX:InitiatingHeapOccupancyPercent=35 -XX:+ExplicitGCInvokesConcurrent -Djava.awt.headless=true" fi -# version option -for args in "$@" ; do - if [ "$args" = "--version" ]; then - exec $JAVA $KAFKA_HEAP_OPTS $KAFKA_JVM_PERFORMANCE_OPTS $KAFKA_GC_LOG_OPTS $KAFKA_JMX_OPTS $KAFKA_LOG4J_OPTS -cp $CLASSPATH $KAFKA_OPTS "kafka.utils.VersionInfo" - fi -done - while [ $# -gt 0 ]; do COMMAND=$1 case $COMMAND in diff --git a/bin/windows/kafka-delete-records.bat b/bin/windows/kafka-delete-records.bat new file mode 100644 index 0000000000000..d07e05f88a22b --- /dev/null +++ b/bin/windows/kafka-delete-records.bat @@ -0,0 +1,17 @@ +@echo off +rem Licensed to the Apache Software Foundation (ASF) under one or more +rem contributor license agreements. See the NOTICE file distributed with +rem this work for additional information regarding copyright ownership. +rem The ASF licenses this file to You under the Apache License, Version 2.0 +rem (the "License"); you may not use this file except in compliance with +rem the License. You may obtain a copy of the License at +rem +rem http://www.apache.org/licenses/LICENSE-2.0 +rem +rem Unless required by applicable law or agreed to in writing, software +rem distributed under the License is distributed on an "AS IS" BASIS, +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.DeleteRecordsCommand %* diff --git a/bin/windows/kafka-leader-election.bat b/bin/windows/kafka-leader-election.bat new file mode 100644 index 0000000000000..0432a99b6e413 --- /dev/null +++ b/bin/windows/kafka-leader-election.bat @@ -0,0 +1,17 @@ +@echo off +rem Licensed to the Apache Software Foundation (ASF) under one or more +rem contributor license agreements. See the NOTICE file distributed with +rem this work for additional information regarding copyright ownership. +rem The ASF licenses this file to You under the Apache License, Version 2.0 +rem (the "License"); you may not use this file except in compliance with +rem the License. You may obtain a copy of the License at +rem +rem http://www.apache.org/licenses/LICENSE-2.0 +rem +rem Unless required by applicable law or agreed to in writing, software +rem distributed under the License is distributed on an "AS IS" BASIS, +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.LeaderElectionCommand %* diff --git a/bin/windows/kafka-log-dirs.bat b/bin/windows/kafka-log-dirs.bat new file mode 100644 index 0000000000000..b490d47feaed6 --- /dev/null +++ b/bin/windows/kafka-log-dirs.bat @@ -0,0 +1,17 @@ +@echo off +rem Licensed to the Apache Software Foundation (ASF) under one or more +rem contributor license agreements. See the NOTICE file distributed with +rem this work for additional information regarding copyright ownership. +rem The ASF licenses this file to You under the Apache License, Version 2.0 +rem (the "License"); you may not use this file except in compliance with +rem the License. You may obtain a copy of the License at +rem +rem http://www.apache.org/licenses/LICENSE-2.0 +rem +rem Unless required by applicable law or agreed to in writing, software +rem distributed under the License is distributed on an "AS IS" BASIS, +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.LogDirsCommand %* diff --git a/bin/windows/kafka-run-class.bat b/bin/windows/kafka-run-class.bat index 3bec536e67a73..ad7cbb882ce08 100755 --- a/bin/windows/kafka-run-class.bat +++ b/bin/windows/kafka-run-class.bat @@ -27,7 +27,7 @@ set BASE_DIR=%CD% popd IF ["%SCALA_VERSION%"] EQU [""] ( - set SCALA_VERSION=2.12.8 + set SCALA_VERSION=2.12.10 ) IF ["%SCALA_BINARY_VERSION%"] EQU [""] ( diff --git a/bin/windows/kafka-streams-application-reset.bat b/bin/windows/kafka-streams-application-reset.bat new file mode 100644 index 0000000000000..1cfb6f518c824 --- /dev/null +++ b/bin/windows/kafka-streams-application-reset.bat @@ -0,0 +1,23 @@ +@echo off +rem Licensed to the Apache Software Foundation (ASF) under one or more +rem contributor license agreements. See the NOTICE file distributed with +rem this work for additional information regarding copyright ownership. +rem The ASF licenses this file to You under the Apache License, Version 2.0 +rem (the "License"); you may not use this file except in compliance with +rem the License. You may obtain a copy of the License at +rem +rem http://www.apache.org/licenses/LICENSE-2.0 +rem +rem Unless required by applicable law or agreed to in writing, software +rem distributed under the License is distributed on an "AS IS" BASIS, +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. + +SetLocal +IF ["%KAFKA_HEAP_OPTS%"] EQU [""] ( + set KAFKA_HEAP_OPTS=-Xmx512M +) + +"%~dp0kafka-run-class.bat" kafka.tools.StreamsResetter %* +EndLocal diff --git a/build.gradle b/build.gradle index ff316c829c32b..60840bf20b9f7 100644 --- a/build.gradle +++ b/build.gradle @@ -16,6 +16,8 @@ import org.ajoberstar.grgit.Grgit import java.nio.charset.StandardCharsets +import javax.xml.bind.DatatypeConverter +import java.security.MessageDigest buildscript { repositories { @@ -26,16 +28,17 @@ buildscript { } } apply from: file('gradle/buildscript.gradle'), to: buildscript + apply from: "$rootDir/gradle/dependencies.gradle" dependencies { // For Apache Rat plugin to ignore non-Git files - classpath "org.ajoberstar:grgit:1.9.3" - classpath 'com.github.ben-manes:gradle-versions-plugin:0.20.0' - classpath 'org.scoverage:gradle-scoverage:2.5.0' - classpath 'com.github.jengelman.gradle.plugins:shadow:4.0.3' - classpath 'org.owasp:dependency-check-gradle:4.0.2' - classpath "com.diffplug.spotless:spotless-plugin-gradle:3.17.0" - classpath "gradle.plugin.com.github.spotbugs:spotbugs-gradle-plugin:1.6.9" + classpath "org.ajoberstar.grgit:grgit-core:$versions.grgit" + classpath "com.github.ben-manes:gradle-versions-plugin:$versions.gradleVersionsPlugin" + classpath "org.scoverage:gradle-scoverage:$versions.scoveragePlugin" + classpath "com.github.jengelman.gradle.plugins:shadow:$versions.shadowPlugin" + classpath "org.owasp:dependency-check-gradle:$versions.owaspDepCheckPlugin" + classpath "com.diffplug.spotless:spotless-plugin-gradle:$versions.spotlessPlugin" + classpath "gradle.plugin.com.github.spotbugs:spotbugs-gradle-plugin:$versions.spotbugsPlugin" } } @@ -43,18 +46,17 @@ apply plugin: "com.diffplug.gradle.spotless" spotless { scala { target 'streams/**/*.scala' - scalafmt('1.5.1').configFile('checkstyle/.scalafmt.conf') + scalafmt("$versions.scalafmt").configFile('checkstyle/.scalafmt.conf') } } -apply from: "$rootDir/gradle/dependencies.gradle" allprojects { repositories { mavenCentral() } - + apply plugin: 'idea' apply plugin: 'org.owasp.dependencycheck' apply plugin: 'com.github.ben-manes.versions' @@ -91,7 +93,7 @@ allprojects { } ext { - gradleVersion = "5.1.1" + gradleVersion = "$versions.gradle" minJavaVersion = "8" buildVersionFileName = "kafka-version.properties" @@ -100,9 +102,14 @@ ext { skipSigning = project.hasProperty('skipSigning') && skipSigning.toBoolean() shouldSign = !skipSigning && !version.endsWith("SNAPSHOT") && project.gradle.startParameter.taskNames.any { it.contains("upload") } - mavenUrl = project.hasProperty('mavenUrl') ? project.mavenUrl : '' - mavenUsername = project.hasProperty('mavenUsername') ? project.mavenUsername : '' - mavenPassword = project.hasProperty('mavenPassword') ? project.mavenPassword : '' + bintrayUsername = System.getenv('BINTRAY_USER') + bintrayKey=System.getenv('BINTRAY_KEY') + bintrayUrl = bintrayApiBaseUrl + '/maven/linkedin/' + bintrayRepo + '/' + bintrayPackage + '/;publish=1' + + // By default, publish to Bintray. + mavenUrl = project.hasProperty('mavenUrl') ? project.mavenUrl : bintrayUrl + mavenUsername = project.hasProperty('mavenUsername') ? project.mavenUsername : bintrayUsername + mavenPassword = project.hasProperty('mavenPassword') ? project.mavenPassword : bintrayKey userShowStandardStreams = project.hasProperty("showStandardStreams") ? showStandardStreams : null @@ -120,7 +127,7 @@ if (file('.git').exists()) { rat { // Exclude everything under the directory that git should be ignoring via .gitignore or that isn't checked in. These // restrict us only to files that are checked in or are staged. - def repo = Grgit.open(project.getRootDir()) + def repo = Grgit.open(currentDir: project.getRootDir()) excludes = new ArrayList(repo.clean(ignore: false, directories: true, dryRun: true)) // And some of the files that we have checked in should also be excluded from this check excludes.addAll([ @@ -136,13 +143,22 @@ if (file('.git').exists()) { '**/id_rsa.pub', 'checkstyle/suppressions.xml', 'streams/quickstart/java/src/test/resources/projects/basic/goal.txt', - 'streams/streams-scala/logs/*' + 'streams/streams-scala/logs/*', + '**/generated/**' ]) } } subprojects { + + // enable running :dependencies task recursively on all subprojects + // eg: ./gradlew allDeps + task allDeps(type: DependencyReportTask) {} + // enable running :dependencyInsight task recursively on all subprojects + // eg: ./gradlew allDepInsight --configuration runtime --dependency com.fasterxml.jackson.core:jackson-databind + task allDepInsight(type: DependencyInsightReportTask) doLast {} + apply plugin: 'java' // apply the eclipse plugin only to subprojects that hold code. 'connect' is just a folder. if (!project.name.equals('connect')) { @@ -187,13 +203,13 @@ subprojects { afterEvaluate { pom.artifactId = "${archivesBaseName}" pom.project { - name 'Apache Kafka' + name 'LinkedIn fork of Apache Kafka' packaging 'jar' - url 'http://kafka.apache.org' + url 'http://github.com/linkedin/kafka' licenses { license { name 'The Apache Software License, Version 2.0' - url 'http://www.apache.org/licenses/LICENSE-2.0.txt' + url 'https://www.apache.org/licenses/LICENSE-2.0.txt' distribution 'repo' } } @@ -237,7 +253,13 @@ subprojects { // code in the test (best guess is that it is tail output from last test). We won't have // an output file for these, so simply ignore them. If they become critical for debugging, // they can be seen with showStandardStreams. - if (td.name == td.className) { + if (td.name == td.className || td.className == null) { + // silently ignore output unrelated to specific test methods + return + } else if (logStreams.get(tid) == null) { + println "WARNING: unexpectedly got output for a test [${tid}]" + + " that we didn't previously see in the beforeTest hook." + + " Message for debugging: [" + toe.message + "]." return } try { @@ -321,6 +343,50 @@ subprojects { } jar { + + doFirst { + //version embedding at build time + //we create a file called __Versioning__[md5] and set its payload to contain versioning info + // + //the file name is not fixed, but contains an md5 hash of the project group name and version. + //its not fixed so that if multiple libraries are repackaged (think fat/shaded jar) the versioning files + //are unlikely to conflict or get overwritten. the md5 hash part is chosen such that builds are + //"repeatable" - re-running the same build results in the same output. this is nice for incremental builds. + // + //the value contains versioning information for both this module and the entire project it was built as part of. + // + //this information allows us, at runtime, to determine what libraries/versions exist on a given classpath + + String moduleVersion = String.valueOf(project.version) + String projectVersion = String.valueOf(rootProject.version) + + if ("unspecified" != projectVersion && "unspecified" != moduleVersion) { + MessageDigest md = MessageDigest.getInstance("MD5") + md.update(String.valueOf(project.group).getBytes("UTF-8")) + md.update(String.valueOf(project.name).getBytes("UTF-8")) + md.update(moduleVersion.getBytes("UTF-8")) + byte[] digest = md.digest() + String hash = DatatypeConverter.printHexBinary(digest).toUpperCase() + + String versioningFileName = "__Versioning__" + hash + String versioningPayload = "{" + + "\"format\":\"v1\"," + + "\"project\":\"" + rootProject.group + ":" + rootProject.name + ":" + projectVersion + "\"," + + "\"module\":\"" + project.group + ":" + project.name + ":" + moduleVersion + "\"" + + "}" + + String outFolder = "${project.buildDir}/resources/main/META-INF/" + mkdir outFolder + + fileTree(dir: outFolder, include:'__Versioning__*').each { + //remove any left-overs from previous builds + it.delete() + } + + file("${outFolder}/${versioningFileName}").text = versioningPayload + } + } + from "$rootDir/LICENSE" from "$rootDir/NOTICE" } @@ -401,7 +467,6 @@ subprojects { "-language:postfixOps", "-language:implicitConversions", "-language:existentials", - "-Xlint:by-name-right-associative", "-Xlint:delayedinit-select", "-Xlint:doc-detached", "-Xlint:missing-interpolator", @@ -412,8 +477,7 @@ subprojects { "-Xlint:poly-implicit-overload", "-Xlint:private-shadow", "-Xlint:stars-align", - "-Xlint:type-parameter-shadow", - "-Xlint:unsound-match" + "-Xlint:type-parameter-shadow" ] if (versions.baseScala != '2.11') { @@ -422,6 +486,15 @@ subprojects { "-Xlint:unused" ] } + + // these options are valid for Scala versions < 2.13 only + // Scala 2.13 removes them, see https://github.com/scala/scala/pull/6502 and https://github.com/scala/scala/pull/5969 + if (versions.baseScala in ['2.11','2.12']) { + scalaCompileOptions.additionalParameters += [ + "-Xlint:by-name-right-associative", + "-Xlint:unsound-match" + ] + } configure(scalaCompileOptions.forkOptions) { memoryMaximumSize = '1g' @@ -432,7 +505,7 @@ subprojects { checkstyle { configFile = new File(rootDir, "checkstyle/checkstyle.xml") configProperties = checkstyleConfigProperties("import-control.xml") - toolVersion = '8.10' + toolVersion = "$versions.checkstyle" } configure(checkstyleMain) { @@ -448,7 +521,7 @@ subprojects { test.dependsOn('checkstyleMain', 'checkstyleTest') spotbugs { - toolVersion = '3.1.8' + toolVersion = "$versions.spotbugs" excludeFilter = file("$rootDir/gradle/spotbugs-exclude.xml") ignoreFailures = false } @@ -467,7 +540,7 @@ subprojects { apply plugin: "jacoco" jacoco { - toolVersion = "0.8.2" + toolVersion = "$versions.jacoco" } // NOTE: Jacoco Gradle plugin does not support "offline instrumentation" this means that classes mocked by PowerMock @@ -511,8 +584,8 @@ def fineTuneEclipseClasspathFile(eclipse, project) { if (project.name.equals('core')) { cp.entries.findAll { it.kind == "src" && it.path.equals("src/test/scala") }*.excludes = ["integration/", "other/", "unit/"] } - /* - * Set all eclipse build output to go to 'build_eclipse' directory. This is to ensure that gradle and eclipse use different + /* + * Set all eclipse build output to go to 'build_eclipse' directory. This is to ensure that gradle and eclipse use different * build output directories, and also avoid using the eclpise default of 'bin' which clashes with some of our script directories. * https://discuss.gradle.org/t/eclipse-generated-files-should-be-put-in-the-same-place-as-the-gradle-generated-files/6986/2 */ @@ -604,8 +677,26 @@ for ( sv in availableScalaVersions ) { } } -def connectPkgs = ['connect:api', 'connect:runtime', 'connect:transforms', 'connect:json', 'connect:file', 'connect:basic-auth-extension'] -def pkgs = ['clients', 'examples', 'log4j-appender', 'tools', 'streams', 'streams:streams-scala', 'streams:test-utils', 'streams:examples'] + connectPkgs +def connectPkgs = [ + 'connect:api', + 'connect:basic-auth-extension', + 'connect:file', + 'connect:json', + 'connect:runtime', + 'connect:transforms', + 'connect:mirror', + 'connect:mirror-client' +] + +def pkgs = [ + 'clients', + 'examples', + 'log4j-appender', + 'streams', + 'streams:examples', + 'streams:test-utils', + 'tools' +] + connectPkgs /** Create one task per default Scala version */ def withDefScalaVersions(taskName) { @@ -638,22 +729,28 @@ project(':core') { dependencies { compile project(':clients') compile libs.jacksonDatabind + compile libs.jacksonModuleScala + compile libs.jacksonDataformatCsv compile libs.jacksonJDK8Datatypes compile libs.joptSimple compile libs.metrics + compile libs.scalaCollectionCompat + compile libs.scalaJava8Compat compile libs.scalaLibrary // only needed transitively, but set it explicitly to ensure it has the same version as scala-library compile libs.scalaReflect compile libs.scalaLogging compile libs.slf4jApi - compile(libs.zkclient) { - exclude module: 'zookeeper' - } compile(libs.zookeeper) { exclude module: 'slf4j-log4j12' exclude module: 'log4j' exclude module: 'netty' } + // pin the netty to 4.1.49.Final since the transitive dependent version 4.1.48.Final is vulnerable + compile libs.nettyHandler + compile libs.nettyTransportNativeEpoll + // ZooKeeperMain depends on commons-cli but declares the dependency as `provided` + compile libs.commonsCli compileOnly libs.log4j @@ -683,7 +780,7 @@ project(':core') { scoverage libs.scoveragePlugin scoverage libs.scoverageRuntime } - + scoverage { reportDir = file("${rootProject.buildDir}/scoverage") highlighting = false @@ -833,6 +930,10 @@ project(':core') { from(project(':connect:file').configurations.runtime) { into("libs/") } from(project(':connect:basic-auth-extension').jar) { into("libs/") } from(project(':connect:basic-auth-extension').configurations.runtime) { into("libs/") } + from(project(':connect:mirror').jar) { into("libs/") } + from(project(':connect:mirror').configurations.runtime) { into("libs/") } + from(project(':connect:mirror-client').jar) { into("libs/") } + from(project(':connect:mirror-client').configurations.runtime) { into("libs/") } from(project(':streams').jar) { into("libs/") } from(project(':streams').configurations.runtime) { into("libs/") } from(project(':streams:streams-scala').jar) { into("libs/") } @@ -920,8 +1021,10 @@ project(':clients') { compile libs.lz4 compile libs.snappy compile libs.slf4jApi + compileOnly libs.jacksonDatabind // for SASL/OAUTHBEARER bearer token parsing compileOnly libs.jacksonJDK8Datatypes + compile libs.conscrypt jacksonDatabindConfig libs.jacksonDatabind // to publish as provided scope dependency. @@ -984,11 +1087,23 @@ project(':clients') { task processMessages(type:JavaExec) { main = "org.apache.kafka.message.MessageGenerator" classpath = project(':generator').sourceSets.main.runtimeClasspath - args = [ "src/generated/java/org/apache/kafka/common/message", "src/main/resources/common/message" ] + args = [ "org.apache.kafka.common.message", + "src/generated/java/org/apache/kafka/common/message", + "src/main/resources/common/message" ] inputs.dir("src/main/resources/common/message") outputs.dir("src/generated/java/org/apache/kafka/common/message") } + task processTestMessages(type:JavaExec) { + main = "org.apache.kafka.message.MessageGenerator" + classpath = project(':generator').sourceSets.main.runtimeClasspath + args = [ "org.apache.kafka.common.message", + "src/generated-test/java/org/apache/kafka/common/message", + "src/test/resources/common/message" ] + inputs.dir("src/test/resources/common/message") + outputs.dir("src/generated-test/java/org/apache/kafka/common/message") + } + sourceSets { main { java { @@ -997,13 +1112,15 @@ project(':clients') { } test { java { - srcDirs = ["src/generated/java", "src/test/java"] + srcDirs = ["src/generated/java", "src/generated-test/java", "src/test/java"] } } } compileJava.dependsOn 'processMessages' + compileTestJava.dependsOn 'processTestMessages' + javadoc { include "**/org/apache/kafka/clients/admin/*" include "**/org/apache/kafka/clients/consumer/*" @@ -1022,6 +1139,7 @@ project(':clients') { include "**/org/apache/kafka/common/security/scram/*" include "**/org/apache/kafka/common/security/token/delegation/*" include "**/org/apache/kafka/common/security/oauthbearer/*" + include "**/org/apache/kafka/server/authorizer/*" include "**/org/apache/kafka/server/policy/*" include "**/org/apache/kafka/server/quota/*" } @@ -1080,6 +1198,7 @@ project(':tools') { project(':streams') { archivesBaseName = "kafka-streams" + ext.buildStreamsVersionFileName = "kafka-streams-version.properties" dependencies { compile project(':clients') @@ -1101,6 +1220,8 @@ project(':streams') { testCompile libs.log4j testCompile libs.junit testCompile libs.easymock + testCompile libs.powermockJunit4 + testCompile libs.powermockEasymock testCompile libs.bcpkix testCompile libs.hamcrest @@ -1125,7 +1246,46 @@ project(':streams') { duplicatesStrategy 'exclude' } + task determineCommitId { + def takeFromHash = 16 + if (commitId) { + commitId = commitId.take(takeFromHash) + } else if (file("$rootDir/.git/HEAD").exists()) { + def headRef = file("$rootDir/.git/HEAD").text + if (headRef.contains('ref: ')) { + headRef = headRef.replaceAll('ref: ', '').trim() + if (file("$rootDir/.git/$headRef").exists()) { + commitId = file("$rootDir/.git/$headRef").text.trim().take(takeFromHash) + } + } else { + commitId = headRef.trim().take(takeFromHash) + } + } else { + commitId = "unknown" + } + } + + task createStreamsVersionFile(dependsOn: determineCommitId) { + ext.receiptFile = file("$buildDir/kafka/$buildStreamsVersionFileName") + outputs.file receiptFile + outputs.upToDateWhen { false } + doLast { + def data = [ + commitId: commitId, + version: version, + ] + + receiptFile.parentFile.mkdirs() + def content = data.entrySet().collect { "$it.key=$it.value" }.sort().join("\n") + receiptFile.setText(content, "ISO-8859-1") + } + } + jar { + dependsOn 'createStreamsVersionFile' + from("$buildDir") { + include "kafka/$buildStreamsVersionFileName" + } dependsOn 'copyDependantLibs' } @@ -1139,6 +1299,12 @@ project(':streams') { if( !generatedDocsDir.exists() ) { generatedDocsDir.mkdirs() } standardOutput = new File(generatedDocsDir, "streams_config.html").newOutputStream() } + + test { + // The suites are for running sets of tests in IDEs. + // Gradle will run each test class, so we exclude the suites to avoid redundantly running the tests twice. + exclude '**/*Suite.class' + } } project(':streams:streams-scala') { @@ -1160,6 +1326,7 @@ project(':streams:streams-scala') { testCompile libs.junit testCompile libs.scalatest testCompile libs.easymock + testCompile libs.hamcrest testRuntime libs.slf4jlog4j } @@ -1193,6 +1360,7 @@ project(':streams:test-utils') { testCompile project(':clients').sourceSets.test.output testCompile libs.junit testCompile libs.easymock + testCompile libs.hamcrest testRuntime libs.slf4jlog4j } @@ -1342,6 +1510,30 @@ project(':streams:upgrade-system-tests-21') { } } +project(':streams:upgrade-system-tests-22') { + archivesBaseName = "kafka-streams-upgrade-system-tests-22" + + dependencies { + testCompile libs.kafkaStreams_22 + } + + systemTestLibs { + dependsOn testJar + } +} + +project(':streams:upgrade-system-tests-23') { + archivesBaseName = "kafka-streams-upgrade-system-tests-23" + + dependencies { + testCompile libs.kafkaStreams_23 + } + + systemTestLibs { + dependsOn testJar + } +} + project(':jmh-benchmarks') { apply plugin: 'com.github.johnrengelman.shadow' @@ -1353,11 +1545,17 @@ project(':jmh-benchmarks') { } dependencies { + compile project(':core') compile project(':clients') compile project(':streams') + compile project(':core') + compile project(':clients').sourceSets.test.output + compile project(':core').sourceSets.test.output compile libs.jmhCore annotationProcessor libs.jmhGeneratorAnnProcess compile libs.jmhCoreBenchmarks + compile libs.mockitoCore + compile libs.slf4jlog4j } jar { @@ -1366,6 +1564,9 @@ project(':jmh-benchmarks') { } } + checkstyle { + configProperties = checkstyleConfigProperties("import-control-jmh-benchmarks.xml") + } task jmh(type: JavaExec, dependsOn: [':jmh-benchmarks:clean', ':jmh-benchmarks:shadowJar']) { @@ -1418,7 +1619,11 @@ project(':connect:api') { javadoc { include "**/org/apache/kafka/connect/**" // needed for the `javadocAll` task - options.links "http://docs.oracle.com/javase/7/docs/api/" + // The URL structure was changed to include the locale after Java 8 + if (JavaVersion.current().isJava11Compatible()) + options.links "https://docs.oracle.com/en/java/javase/${JavaVersion.current().majorVersion}/docs/api/" + else + options.links "https://docs.oracle.com/javase/8/docs/api/" } tasks.create(name: "copyDependantLibs", type: Copy) { @@ -1546,6 +1751,7 @@ project(':connect:runtime') { testCompile libs.junit testCompile libs.powermockJunit4 testCompile libs.powermockEasymock + testCompile libs.mockitoCore testCompile libs.httpclient testCompile project(':clients').sourceSets.test.output @@ -1691,11 +1897,93 @@ project(':connect:basic-auth-extension') { } } +project(':connect:mirror') { + archivesBaseName = "connect-mirror" + + dependencies { + compile project(':connect:api') + compile project(':connect:runtime') + compile project(':connect:mirror-client') + compile project(':clients') + compile libs.argparse4j + compile libs.slf4jApi + + testCompile libs.junit + testCompile project(':clients').sourceSets.test.output + testCompile project(':connect:runtime').sourceSets.test.output + testCompile project(':core') + testCompile project(':core').sourceSets.test.output + + testRuntime project(':connect:runtime') + testRuntime libs.slf4jlog4j + } + + javadoc { + enabled = false + } + + tasks.create(name: "copyDependantLibs", type: Copy) { + from (configurations.testRuntime) { + include('slf4j-log4j12*') + include('log4j*jar') + } + from (configurations.runtime) { + exclude('kafka-clients*') + exclude('connect-*') + } + into "$buildDir/dependant-libs" + duplicatesStrategy 'exclude' + } + + jar { + dependsOn copyDependantLibs + } +} + +project(':connect:mirror-client') { + archivesBaseName = "connect-mirror-client" + + dependencies { + compile project(':clients') + compile libs.slf4jApi + + testCompile libs.junit + testCompile project(':clients').sourceSets.test.output + + testRuntime libs.slf4jlog4j + } + + javadoc { + enabled = true + } + + tasks.create(name: "copyDependantLibs", type: Copy) { + from (configurations.testRuntime) { + include('slf4j-log4j12*') + include('log4j*jar') + } + from (configurations.runtime) { + exclude('kafka-clients*') + exclude('connect-*') + } + into "$buildDir/dependant-libs" + duplicatesStrategy 'exclude' + } + + jar { + dependsOn copyDependantLibs + } +} + task aggregatedJavadoc(type: Javadoc) { def projectsWithJavadoc = subprojects.findAll { it.javadoc.enabled } source = projectsWithJavadoc.collect { it.sourceSets.main.allJava } classpath = files(projectsWithJavadoc.collect { it.sourceSets.main.compileClasspath }) includes = projectsWithJavadoc.collectMany { it.javadoc.getIncludes() } excludes = projectsWithJavadoc.collectMany { it.javadoc.getExcludes() } - options.links "http://docs.oracle.com/javase/7/docs/api/" + // The URL structure was changed to include the locale after Java 8 + if (JavaVersion.current().isJava11Compatible()) + options.links "https://docs.oracle.com/en/java/javase/${JavaVersion.current().majorVersion}/docs/api/" + else + options.links "https://docs.oracle.com/javase/8/docs/api/" } diff --git a/checkstyle/checkstyle.xml b/checkstyle/checkstyle.xml index 13cfdb82bd0a8..3cf229d51eb68 100644 --- a/checkstyle/checkstyle.xml +++ b/checkstyle/checkstyle.xml @@ -127,8 +127,15 @@ + - + + + + + diff --git a/checkstyle/import-control-jmh-benchmarks.xml b/checkstyle/import-control-jmh-benchmarks.xml new file mode 100644 index 0000000000000..270d47a926343 --- /dev/null +++ b/checkstyle/import-control-jmh-benchmarks.xml @@ -0,0 +1,51 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/checkstyle/import-control.xml b/checkstyle/import-control.xml index ffc9bf9b18822..fe567794c23ce 100644 --- a/checkstyle/import-control.xml +++ b/checkstyle/import-control.xml @@ -54,6 +54,7 @@ + @@ -118,6 +119,7 @@ + @@ -147,6 +149,8 @@ + + @@ -158,6 +162,7 @@ + @@ -238,7 +243,7 @@ - + @@ -264,10 +269,11 @@ - + + @@ -275,17 +281,6 @@ - - - - - - - - - - - @@ -322,6 +317,13 @@ + + + + + + + @@ -331,10 +333,26 @@ + + + + + + + + + + + + + + + + @@ -343,11 +361,19 @@ + + + + + + + + @@ -356,11 +382,13 @@ + + @@ -381,6 +409,8 @@ + + diff --git a/checkstyle/suppressions.xml b/checkstyle/suppressions.xml index f306aaa9c5f80..3791b7354d00e 100644 --- a/checkstyle/suppressions.xml +++ b/checkstyle/suppressions.xml @@ -1,4 +1,4 @@ - + + + + + + + + + files="(Fetcher|Sender|SenderTest|ConsumerCoordinator|KafkaConsumer|KafkaProducer|Utils|TransactionManagerTest|KafkaAdminClient|NetworkClient|Admin).java"/> + files="KerberosLogin.java|RequestResponseTest.java|ConnectMetricsRegistry.java|KafkaConsumer.java"/> - + files="(KafkaConsumer|ConsumerCoordinator|Fetcher|KafkaProducer|AbstractRequest|AbstractResponse|TransactionManager|Admin|KafkaAdminClient).java"/> + files="(Utils|Topic|KafkaLZ4BlockOutputStream|AclData|JoinGroupRequest).java"/> + files="(ConsumerCoordinator|Fetcher|Sender|KafkaProducer|BufferPool|ConfigDef|RecordAccumulator|KerberosLogin|AbstractRequest|AbstractResponse|Selector|SslFactory|SslTransportLayer|SaslClientAuthenticator|SaslClientCallbackHandler|SaslServerAuthenticator|AbstractCoordinator|TransactionManager|AbstractStickyAssignor).java"/> + files="(AbstractRequest|KerberosLogin|WorkerSinkTaskTest|TransactionManagerTest|SenderTest|KafkaAdminClient|ConsumerCoordinatorTest).java"/> + files="(ConsumerCoordinator|BufferPool|Fetcher|MetricName|Node|ConfigDef|RecordBatch|SslFactory|SslTransportLayer|MetadataResponse|KerberosLogin|Selector|Sender|Serdes|TokenInformation|Agent|Values|PluginUtils|MiniTrogdorCluster|TasksRequest|KafkaProducer|AbstractStickyAssignor).java"/> + + + + + + files="(Sender|Fetcher|KafkaConsumer|Metrics|RequestResponse|TransactionManager|KafkaAdminClient|Message)Test.java"/> + files="(ConsumerCoordinator|KafkaConsumer|RequestResponse|Fetcher|KafkaAdminClient|Message)Test.java"/> @@ -78,17 +95,19 @@ + + - + + files="(KafkaConfigBackingStore|IncrementalCooperativeAssignor|RequestResponseTest|WorkerSinkTaskTest).java"/> - + files="Worker(SinkTask|SourceTask|Coordinator).java"/> @@ -110,6 +129,8 @@ files="KafkaConfigBackingStore.java"/> + @@ -117,7 +138,7 @@ files="Values.java"/> + files="(DistributedHerder|RestClient|JsonConverter|KafkaConfigBackingStore|FileStreamSourceTask).java"/> @@ -126,12 +147,15 @@ + + + files="(KTableImpl|StreamsPartitionAssignor.java)"/> @@ -152,7 +176,10 @@ files="StreamsPartitionAssignor.java"/> + files="(ProcessorStateManager|InternalTopologyBuilder|KafkaStreams|StreamsPartitionAssignor|StreamThread).java"/> + + + + @@ -178,7 +208,7 @@ files="SmokeTestDriver.java"/> + files="KStreamKStreamJoinTest.java|KTableKTableForeignKeyJoinIntegrationTest.java"/> @@ -188,9 +218,11 @@ files="SmokeTestDriver.java"/> - + files="EosTestDriver|KStreamKStreamJoinTest.java|SmokeTestDriver.java|KStreamKStreamLeftJoinTest.java|KTableKTableForeignKeyJoinIntegrationTest.java"/> + + + - + + - - - - - diff --git a/clients/src/main/java/org/apache/kafka/clients/ApiVersion.java b/clients/src/main/java/org/apache/kafka/clients/ApiVersion.java new file mode 100644 index 0000000000000..9d606bbfa7afc --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/ApiVersion.java @@ -0,0 +1,56 @@ +/* + * 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; + +import org.apache.kafka.common.message.ApiVersionsResponseData.ApiVersionsResponseKey; +import org.apache.kafka.common.protocol.ApiKeys; + +/** + * Represents the min version and max version of an api key. + * + * NOTE: This class is intended for INTERNAL usage only within Kafka. + */ +public class ApiVersion { + public final short apiKey; + public final short minVersion; + public final short maxVersion; + + public ApiVersion(ApiKeys apiKey) { + this(apiKey.id, apiKey.oldestVersion(), apiKey.latestVersion()); + } + + public ApiVersion(short apiKey, short minVersion, short maxVersion) { + this.apiKey = apiKey; + this.minVersion = minVersion; + this.maxVersion = maxVersion; + } + + public ApiVersion(ApiVersionsResponseKey apiVersionsResponseKey) { + this.apiKey = apiVersionsResponseKey.apiKey(); + this.minVersion = apiVersionsResponseKey.minVersion(); + this.maxVersion = apiVersionsResponseKey.maxVersion(); + } + + @Override + public String toString() { + return "ApiVersion(" + + "apiKey=" + apiKey + + ", minVersion=" + minVersion + + ", maxVersion= " + maxVersion + + ")"; + } +} diff --git a/clients/src/main/java/org/apache/kafka/clients/ClientRequest.java b/clients/src/main/java/org/apache/kafka/clients/ClientRequest.java index 7b44ca3ad3c44..627615c1b3fb3 100644 --- a/clients/src/main/java/org/apache/kafka/clients/ClientRequest.java +++ b/clients/src/main/java/org/apache/kafka/clients/ClientRequest.java @@ -16,6 +16,7 @@ */ package org.apache.kafka.clients; +import org.apache.kafka.common.message.RequestHeaderData; import org.apache.kafka.common.protocol.ApiKeys; import org.apache.kafka.common.requests.AbstractRequest; import org.apache.kafka.common.requests.RequestHeader; @@ -82,7 +83,14 @@ public ApiKeys apiKey() { } public RequestHeader makeHeader(short version) { - return new RequestHeader(apiKey(), version, clientId, correlationId); + short requestApiKey = requestBuilder.apiKey().id; + return new RequestHeader( + new RequestHeaderData(). + setRequestApiKey(requestApiKey). + setRequestApiVersion(version). + setClientId(clientId). + setCorrelationId(correlationId), + ApiKeys.forId(requestApiKey).requestHeaderVersion(version)); } public AbstractRequest.Builder requestBuilder() { diff --git a/clients/src/main/java/org/apache/kafka/clients/ClientUtils.java b/clients/src/main/java/org/apache/kafka/clients/ClientUtils.java index 4d933243655ab..599baeed508dd 100644 --- a/clients/src/main/java/org/apache/kafka/clients/ClientUtils.java +++ b/clients/src/main/java/org/apache/kafka/clients/ClientUtils.java @@ -27,14 +27,12 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import java.io.Closeable; import java.net.InetAddress; import java.net.InetSocketAddress; import java.net.UnknownHostException; import java.util.ArrayList; import java.util.Collections; import java.util.List; -import java.util.concurrent.atomic.AtomicReference; import static org.apache.kafka.common.utils.Utils.getHost; import static org.apache.kafka.common.utils.Utils.getPort; @@ -49,6 +47,14 @@ public static List parseAndValidateAddresses(List url return parseAndValidateAddresses(urls, ClientDnsLookup.forConfig(clientDnsLookupConfig)); } + /** + * Kafka does not use this function directly. However, + * some third-party applications still rely on this API to parse and validate addresses. + */ + public static List parseAndValidateAddresses(List urls) { + return parseAndValidateAddresses(urls, ClientDnsLookup.DEFAULT); + } + public static List parseAndValidateAddresses(List urls, ClientDnsLookup clientDnsLookup) { List addresses = new ArrayList<>(); for (String url : urls) { @@ -65,7 +71,7 @@ public static List parseAndValidateAddresses(List url String resolvedCanonicalName = inetAddress.getCanonicalHostName(); InetSocketAddress address = new InetSocketAddress(resolvedCanonicalName, port); if (address.isUnresolved()) { - log.warn("Couldn't resolve server {} from {} as DNS resolution of the canonical hostname [} failed for {}", url, CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, resolvedCanonicalName, host); + log.warn("Couldn't resolve server {} from {} as DNS resolution of the canonical hostname {} failed for {}", url, CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, resolvedCanonicalName, host); } else { addresses.add(address); } @@ -87,21 +93,11 @@ public static List parseAndValidateAddresses(List url } } if (addresses.isEmpty()) - throw new ConfigException("No resolvable bootstrap urls given in " + CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG); + throw new ConfigException("No resolvable bootstrap server in provided urls: " + + String.join(",", urls)); return addresses; } - public static void closeQuietly(Closeable c, String name, AtomicReference firstException) { - if (c != null) { - try { - c.close(); - } catch (Throwable t) { - firstException.compareAndSet(null, t); - log.error("Failed to close " + name, t); - } - } - } - /** * @param config client configs * @return configured ChannelBuilder based on the configs. diff --git a/clients/src/main/java/org/apache/kafka/clients/ClusterConnectionStates.java b/clients/src/main/java/org/apache/kafka/clients/ClusterConnectionStates.java index e9bd9712e6127..ec3465ac4395b 100644 --- a/clients/src/main/java/org/apache/kafka/clients/ClusterConnectionStates.java +++ b/clients/src/main/java/org/apache/kafka/clients/ClusterConnectionStates.java @@ -72,11 +72,9 @@ public boolean canConnect(String id, long now) { */ public boolean isBlackedOut(String id, long now) { NodeConnectionState state = nodeState.get(id); - if (state == null) - return false; - else - return state.state.isDisconnected() && - now - state.lastConnectAttemptMs < state.reconnectBackoffMs; + return state != null + && state.state.isDisconnected() + && now - state.lastConnectAttemptMs < state.reconnectBackoffMs; } /** @@ -89,8 +87,8 @@ public boolean isBlackedOut(String id, long now) { public long connectionDelay(String id, long now) { NodeConnectionState state = nodeState.get(id); if (state == null) return 0; - long timeWaited = now - state.lastConnectAttemptMs; if (state.state.isDisconnected()) { + long timeWaited = now - state.lastConnectAttemptMs; return Math.max(state.reconnectBackoffMs - timeWaited, 0); } else { // When connecting or connected, we should be able to delay indefinitely since other events (connection or @@ -108,6 +106,17 @@ public boolean isConnecting(String id) { return state != null && state.state == ConnectionState.CONNECTING; } + /** + * Check whether a connection is either being established or awaiting API version information. + * @param id The id of the node to check + * @return true if the node is either connecting or has connected and is awaiting API versions, false otherwise + */ + public boolean isPreparingConnection(String id) { + NodeConnectionState state = nodeState.get(id); + return state != null && + (state.state == ConnectionState.CONNECTING || state.state == ConnectionState.CHECKING_API_VERSIONS); + } + /** * Enter the connecting state for the given connection, moving to a new resolved address if necessary. * @param id the id of the connection diff --git a/clients/src/main/java/org/apache/kafka/clients/CommonClientConfigs.java b/clients/src/main/java/org/apache/kafka/clients/CommonClientConfigs.java index 491b5de886f5a..03b9b31716c3a 100644 --- a/clients/src/main/java/org/apache/kafka/clients/CommonClientConfigs.java +++ b/clients/src/main/java/org/apache/kafka/clients/CommonClientConfigs.java @@ -26,7 +26,7 @@ import java.util.Map; /** - * Some configurations shared by both producer and consumer + * Configurations shared by Kafka client applications: producer, consumer, connect, etc. */ public class CommonClientConfigs { private static final Logger log = LoggerFactory.getLogger(CommonClientConfigs.class); @@ -42,13 +42,16 @@ public class CommonClientConfigs { + "servers (you may want more than one, though, in case a server is down)."; public static final String CLIENT_DNS_LOOKUP_CONFIG = "client.dns.lookup"; - public static final String CLIENT_DNS_LOOKUP_DOC = "

Controls how the client uses DNS lookups.

If set to use_all_dns_ips then, when the lookup returns multiple IP addresses for a hostname," - + " they will all be attempted to connect to before failing the connection. Applies to both bootstrap and advertised servers.

" - + "

If the value is resolve_canonical_bootstrap_servers_only each entry will be resolved and expanded into a list of canonical names.

"; + public static final String CLIENT_DNS_LOOKUP_DOC = "Controls how the client uses DNS lookups. If set to use_all_dns_ips then, when the lookup returns multiple IP addresses for a hostname," + + " they will all be attempted to connect to before failing the connection. Applies to both bootstrap and advertised servers." + + " If the value is resolve_canonical_bootstrap_servers_only each entry will be resolved and expanded into a list of canonical names."; public static final String METADATA_MAX_AGE_CONFIG = "metadata.max.age.ms"; public static final String METADATA_MAX_AGE_DOC = "The period of time in milliseconds after which we force a refresh of metadata even if we haven't seen any partition leadership changes to proactively discover any new brokers or partitions."; + public static final String METADATA_TOPIC_EXPIRY_MS_CONFIG = "metadata.topic.expiry.ms"; + public static final String METADATA_TOPIC_EXPIRY_MS_DOC = "The period of time in milliseconds before we expire topic from metadata cache"; + public static final String SEND_BUFFER_CONFIG = "send.buffer.bytes"; public static final String SEND_BUFFER_DOC = "The size of the TCP send buffer (SO_SNDBUF) to use when sending data. If the value is -1, the OS default will be used."; public static final int SEND_BUFFER_LOWER_BOUND = -1; @@ -60,6 +63,9 @@ public class CommonClientConfigs { public static final String CLIENT_ID_CONFIG = "client.id"; public static final String CLIENT_ID_DOC = "An id string to pass to the server when making requests. The purpose of this is to be able to track the source of requests beyond just ip/port by allowing a logical application name to be included in server-side request logging."; + public static final String CLIENT_RACK_CONFIG = "client.rack"; + public static final String CLIENT_RACK_DOC = "A rack identifier for this client. This can be any string value which indicates where this client is physically located. It corresponds with the broker config 'broker.rack'"; + public static final String RECONNECT_BACKOFF_MS_CONFIG = "reconnect.backoff.ms"; public static final String RECONNECT_BACKOFF_MS_DOC = "The base amount of time to wait before attempting to reconnect to a given host. This avoids repeatedly connecting to a host in a tight loop. This backoff applies to all connection attempts by the client to a broker."; @@ -81,6 +87,10 @@ public class CommonClientConfigs { public static final String METRICS_RECORDING_LEVEL_CONFIG = "metrics.recording.level"; public static final String METRICS_RECORDING_LEVEL_DOC = "The highest recording level for metrics."; + public static final String METRICS_REPLACE_ON_DUPLICATE_CONFIG = "metrics.replace.on.duplicate"; + public static final String METRICS_REPLACE_ON_DUPLICATE_DOC = "If set to true, then multiple sensors registering the same metric will not throw, but will instead log an error. This makes sensor registration errors non fatal."; + + public static final String METRIC_REPORTER_CLASSES_CONFIG = "metric.reporters"; public static final String METRIC_REPORTER_CLASSES_DOC = "A list of classes to use as metrics reporters. Implementing the org.apache.kafka.common.metrics.MetricsReporter interface allows plugging in classes that will be notified of new metric creation. The JmxReporter is always included to register JMX statistics."; @@ -98,6 +108,48 @@ public class CommonClientConfigs { + "elapses the client will resend the request if necessary or fail the request if " + "retries are exhausted."; + public static final String GROUP_ID_CONFIG = "group.id"; + public static final String GROUP_ID_DOC = "A unique string that identifies the consumer group this consumer belongs to. This property is required if the consumer uses either the group management functionality by using subscribe(topic) or the Kafka-based offset management strategy."; + + public static final String GROUP_INSTANCE_ID_CONFIG = "group.instance.id"; + public static final String GROUP_INSTANCE_ID_DOC = "A unique identifier of the consumer instance provided by the end user. " + + "Only non-empty strings are permitted. If set, the consumer is treated as a static member, " + + "which means that only one instance with this ID is allowed in the consumer group at any time. " + + "This can be used in combination with a larger session timeout to avoid group rebalances caused by transient unavailability " + + "(e.g. process restarts). If not set, the consumer will join the group as a dynamic member, which is the traditional behavior."; + + public static final String MAX_POLL_INTERVAL_MS_CONFIG = "max.poll.interval.ms"; + public static final String MAX_POLL_INTERVAL_MS_DOC = "The maximum delay between invocations of poll() when using " + + "consumer group management. This places an upper bound on the amount of time that the consumer can be idle " + + "before fetching more records. If poll() is not called before expiration of this timeout, then the consumer " + + "is considered failed and the group will rebalance in order to reassign the partitions to another member. " + + "For consumers using a non-null group.instance.id which reach this timeout, partitions will not be immediately reassigned. " + + "Instead, the consumer will stop sending heartbeats and partitions will be reassigned " + + "after expiration of session.timeout.ms. This mirrors the behavior of a static consumer which has shutdown."; + + public static final String REBALANCE_TIMEOUT_MS_CONFIG = "rebalance.timeout.ms"; + public static final String REBALANCE_TIMEOUT_MS_DOC = "The maximum allowed time for each worker to join the group " + + "once a rebalance has begun. This is basically a limit on the amount of time needed for all tasks to " + + "flush any pending data and commit offsets. If the timeout is exceeded, then the worker will be removed " + + "from the group, which will cause offset commit failures."; + + public static final String SESSION_TIMEOUT_MS_CONFIG = "session.timeout.ms"; + public static final String SESSION_TIMEOUT_MS_DOC = "The timeout used to detect client failures when using " + + "Kafka's group management facility. The client sends periodic heartbeats to indicate its liveness " + + "to the broker. If no heartbeats are received by the broker before the expiration of this session timeout, " + + "then the broker will remove this client from the group and initiate a rebalance. Note that the value " + + "must be in the allowable range as configured in the broker configuration by group.min.session.timeout.ms " + + "and group.max.session.timeout.ms."; + + public static final String HEARTBEAT_INTERVAL_MS_CONFIG = "heartbeat.interval.ms"; + public static final String HEARTBEAT_INTERVAL_MS_DOC = "The expected time between heartbeats to the consumer " + + "coordinator when using Kafka's group management facilities. Heartbeats are used to ensure that the " + + "consumer's session stays active and to facilitate rebalancing when new consumers join or leave the group. " + + "The value must be set lower than session.timeout.ms, but typically should be set no higher " + + "than 1/3 of that value. It can be adjusted even lower to control the expected time for normal rebalances."; + public static final String ENABLE_STICKY_METADATA_FETCH_CONFIG = "enable.sticky.metadata.fetch"; + public static final String ENABLE_STICKY_METADATA_FETCH_DOC = "Fetch metadata from the least loaded broker if false. Otherwise fetch metadata " + + "from the same broker until it is disconnected."; /** * Postprocess the configuration so that exponential backoff is disabled when reconnect backoff diff --git a/clients/src/main/java/org/apache/kafka/clients/FetchSessionHandler.java b/clients/src/main/java/org/apache/kafka/clients/FetchSessionHandler.java index 30ae65f2309c4..0dc8943fdcbe6 100644 --- a/clients/src/main/java/org/apache/kafka/clients/FetchSessionHandler.java +++ b/clients/src/main/java/org/apache/kafka/clients/FetchSessionHandler.java @@ -196,8 +196,10 @@ public void add(TopicPartition topicPartition, PartitionData data) { public FetchRequestData build() { if (nextMetadata.isFull()) { - log.debug("Built full fetch {} for node {} with {}.", - nextMetadata, node, partitionsToLogString(next.keySet())); + if (log.isDebugEnabled()) { + log.debug("Built full fetch {} for node {} with {}.", + nextMetadata, node, partitionsToLogString(next.keySet())); + } sessionPartitions = next; next = null; Map toSend = @@ -247,10 +249,12 @@ public FetchRequestData build() { sessionPartitions.put(topicPartition, nextData); added.add(topicPartition); } - log.debug("Built incremental fetch {} for node {}. Added {}, altered {}, removed {} " + - "out of {}", nextMetadata, node, partitionsToLogString(added), - partitionsToLogString(altered), partitionsToLogString(removed), - partitionsToLogString(sessionPartitions.keySet())); + if (log.isDebugEnabled()) { + log.debug("Built incremental fetch {} for node {}. Added {}, altered {}, removed {} " + + "out of {}", nextMetadata, node, partitionsToLogString(added), + partitionsToLogString(altered), partitionsToLogString(removed), + partitionsToLogString(sessionPartitions.keySet())); + } Map toSend = Collections.unmodifiableMap(new LinkedHashMap<>(next)); Map curSessionPartitions = @@ -393,14 +397,15 @@ public boolean handleResponse(FetchResponse response) { nextMetadata = FetchMetadata.INITIAL; return false; } else if (response.sessionId() == INVALID_SESSION_ID) { - log.debug("Node {} sent a full fetch response{}", - node, responseDataToLogString(response)); + if (log.isDebugEnabled()) + log.debug("Node {} sent a full fetch response{}", node, responseDataToLogString(response)); nextMetadata = FetchMetadata.INITIAL; return true; } else { // The server created a new incremental fetch session. - log.debug("Node {} sent a full fetch response that created a new incremental " + - "fetch session {}{}", node, response.sessionId(), responseDataToLogString(response)); + if (log.isDebugEnabled()) + log.debug("Node {} sent a full fetch response that created a new incremental " + + "fetch session {}{}", node, response.sessionId(), responseDataToLogString(response)); nextMetadata = FetchMetadata.newIncremental(response.sessionId()); return true; } @@ -412,14 +417,16 @@ public boolean handleResponse(FetchResponse response) { return false; } else if (response.sessionId() == INVALID_SESSION_ID) { // The incremental fetch session was closed by the server. - log.debug("Node {} sent an incremental fetch response closing session {}{}", - node, nextMetadata.sessionId(), responseDataToLogString(response)); + if (log.isDebugEnabled()) + log.debug("Node {} sent an incremental fetch response closing session {}{}", + node, nextMetadata.sessionId(), responseDataToLogString(response)); nextMetadata = FetchMetadata.INITIAL; return true; } else { // The incremental fetch session was continued by the server. - log.debug("Node {} sent an incremental fetch response for session {}{}", - node, response.sessionId(), responseDataToLogString(response)); + if (log.isDebugEnabled()) + log.debug("Node {} sent an incremental fetch response for session {}{}", + node, response.sessionId(), responseDataToLogString(response)); nextMetadata = nextMetadata.nextIncremental(); return true; } @@ -435,7 +442,7 @@ public boolean handleResponse(FetchResponse response) { * @param t The exception. */ public void handleError(Throwable t) { - log.info("Error sending fetch request {} to node {}: {}.", nextMetadata, node, t.toString()); + log.info("Error sending fetch request {} to node {}: {}.", nextMetadata, node, t); nextMetadata = nextMetadata.nextCloseExisting(); } } diff --git a/clients/src/main/java/org/apache/kafka/clients/GroupRebalanceConfig.java b/clients/src/main/java/org/apache/kafka/clients/GroupRebalanceConfig.java new file mode 100644 index 0000000000000..006800a3e5c4e --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/GroupRebalanceConfig.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.clients; + +import org.apache.kafka.common.config.AbstractConfig; +import org.apache.kafka.common.requests.JoinGroupRequest; + +import java.util.Locale; +import java.util.Optional; + +/** + * Class to extract group rebalance related configs. + */ +public class GroupRebalanceConfig { + + public enum ProtocolType { + CONSUMER, + CONNECT; + + @Override + public String toString() { + return super.toString().toLowerCase(Locale.ROOT); + } + } + + public final int sessionTimeoutMs; + public final int rebalanceTimeoutMs; + public final int heartbeatIntervalMs; + public final String groupId; + public final Optional groupInstanceId; + public final long retryBackoffMs; + public final boolean leaveGroupOnClose; + + public GroupRebalanceConfig(AbstractConfig config, ProtocolType protocolType) { + this.sessionTimeoutMs = config.getInt(CommonClientConfigs.SESSION_TIMEOUT_MS_CONFIG); + + // Consumer and Connect use different config names for defining rebalance timeout + if (protocolType == ProtocolType.CONSUMER) { + this.rebalanceTimeoutMs = config.getInt(CommonClientConfigs.MAX_POLL_INTERVAL_MS_CONFIG); + } else { + this.rebalanceTimeoutMs = config.getInt(CommonClientConfigs.REBALANCE_TIMEOUT_MS_CONFIG); + } + + this.heartbeatIntervalMs = config.getInt(CommonClientConfigs.HEARTBEAT_INTERVAL_MS_CONFIG); + this.groupId = config.getString(CommonClientConfigs.GROUP_ID_CONFIG); + + // Static membership is only introduced in consumer API. + if (protocolType == ProtocolType.CONSUMER) { + String groupInstanceId = config.getString(CommonClientConfigs.GROUP_INSTANCE_ID_CONFIG); + if (groupInstanceId != null) { + JoinGroupRequest.validateGroupInstanceId(groupInstanceId); + this.groupInstanceId = Optional.of(groupInstanceId); + } else { + this.groupInstanceId = Optional.empty(); + } + } else { + this.groupInstanceId = Optional.empty(); + } + + this.retryBackoffMs = config.getLong(CommonClientConfigs.RETRY_BACKOFF_MS_CONFIG); + + // Internal leave group config is only defined in Consumer. + if (protocolType == ProtocolType.CONSUMER) { + this.leaveGroupOnClose = config.getBoolean("internal.leave.group.on.close"); + } else { + this.leaveGroupOnClose = true; + } + } + + // For testing purpose. + public GroupRebalanceConfig(final int sessionTimeoutMs, + final int rebalanceTimeoutMs, + final int heartbeatIntervalMs, + String groupId, + Optional groupInstanceId, + long retryBackoffMs, + boolean leaveGroupOnClose) { + this.sessionTimeoutMs = sessionTimeoutMs; + this.rebalanceTimeoutMs = rebalanceTimeoutMs; + this.heartbeatIntervalMs = heartbeatIntervalMs; + this.groupId = groupId; + this.groupInstanceId = groupInstanceId; + this.retryBackoffMs = retryBackoffMs; + this.leaveGroupOnClose = leaveGroupOnClose; + } +} diff --git a/clients/src/main/java/org/apache/kafka/clients/ManualMetadataUpdater.java b/clients/src/main/java/org/apache/kafka/clients/ManualMetadataUpdater.java index ec007a6966879..c1c1fba4a56d0 100644 --- a/clients/src/main/java/org/apache/kafka/clients/ManualMetadataUpdater.java +++ b/clients/src/main/java/org/apache/kafka/clients/ManualMetadataUpdater.java @@ -16,15 +16,15 @@ */ package org.apache.kafka.clients; +import org.apache.kafka.common.KafkaException; import org.apache.kafka.common.Node; import org.apache.kafka.common.errors.AuthenticationException; import org.apache.kafka.common.requests.MetadataResponse; import org.apache.kafka.common.requests.RequestHeader; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; import java.util.ArrayList; import java.util.List; +import java.util.Optional; /** * A simple implementation of `MetadataUpdater` that returns the cluster nodes set via the constructor or via @@ -36,9 +36,6 @@ * This class is not thread-safe! */ public class ManualMetadataUpdater implements MetadataUpdater { - - private static final Logger log = LoggerFactory.getLogger(ManualMetadataUpdater.class); - private List nodes; public ManualMetadataUpdater() { @@ -69,24 +66,18 @@ public long maybeUpdate(long now) { } @Override - public void handleDisconnection(String destination) { - // Do nothing - } - - @Override - public void handleAuthenticationFailure(AuthenticationException exception) { - // We don't fail the broker on authentication failures, but there is sufficient information in the broker logs - // to identify the failure. - log.debug("An authentication error occurred in broker-to-broker communication.", exception); + public void handleServerDisconnect(long now, String nodeId, Optional maybeAuthException) { + // We don't fail the broker on failures. There should be sufficient information from + // the NetworkClient logs to indicate the reason for the failure. } @Override - public void handleCompletedMetadataResponse(RequestHeader requestHeader, long now, MetadataResponse response) { + public void handleFailedRequest(long now, Optional maybeFatalException) { // Do nothing } @Override - public void requestUpdate() { + public void handleSuccessfulResponse(RequestHeader requestHeader, long now, MetadataResponse response) { // Do nothing } 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 dfd461a93ea7d..47e4141c7da28 100644 --- a/clients/src/main/java/org/apache/kafka/clients/Metadata.java +++ b/clients/src/main/java/org/apache/kafka/clients/Metadata.java @@ -18,9 +18,9 @@ import org.apache.kafka.common.Cluster; import org.apache.kafka.common.KafkaException; +import org.apache.kafka.common.Node; import org.apache.kafka.common.PartitionInfo; import org.apache.kafka.common.TopicPartition; -import org.apache.kafka.common.errors.AuthenticationException; import org.apache.kafka.common.errors.InvalidMetadataException; import org.apache.kafka.common.errors.InvalidTopicException; import org.apache.kafka.common.errors.TopicAuthorizationException; @@ -35,6 +35,7 @@ import java.io.Closeable; import java.net.InetSocketAddress; import java.util.ArrayList; +import java.util.Collections; import java.util.HashMap; import java.util.HashSet; import java.util.List; @@ -44,6 +45,7 @@ import java.util.Set; import java.util.function.Consumer; import java.util.function.Predicate; +import java.util.function.Supplier; /** * A class encapsulating some of the logic around metadata. @@ -65,8 +67,9 @@ public class Metadata implements Closeable { private int requestVersion; // bumped on every new topic addition private long lastRefreshMs; private long lastSuccessfulRefreshMs; - private AuthenticationException authenticationException; - private KafkaException metadataException; + private KafkaException fatalException; + private Set invalidTopics; + private Set unauthorizedTopics; private MetadataCache cache = MetadataCache.empty(); private boolean needUpdate; private final ClusterResourceListeners clusterResourceListeners; @@ -97,6 +100,8 @@ public Metadata(long refreshBackoffMs, this.clusterResourceListeners = clusterResourceListeners; this.isClosed = false; this.lastSeenLeaderEpochs = new HashMap<>(); + this.invalidTopics = Collections.emptySet(); + this.unauthorizedTopics = Collections.emptySet(); } /** @@ -129,6 +134,10 @@ public synchronized long timeToNextUpdate(long nowMs) { return Math.max(timeToExpire, timeToAllowUpdate(nowMs)); } + public long metadataExpireMs() { + return this.metadataExpireMs; + } + /** * Request an update of the current cluster metadata info, return the current updateVersion before the update */ @@ -173,7 +182,7 @@ private synchronized boolean updateLastSeenEpoch(TopicPartition topicPartition, } return true; } else { - log.debug("Not replacing existing epoch {} with new epoch {}", oldEpoch, epoch); + log.debug("Not replacing existing epoch {} with new epoch {} for partition {}", oldEpoch, epoch, topicPartition); return false; } } @@ -200,32 +209,8 @@ public synchronized Optional partitionInfoI } } - /** - * If any non-retriable authentication exceptions were encountered during - * metadata update, clear and return the exception. - */ - public synchronized AuthenticationException getAndClearAuthenticationException() { - if (authenticationException != null) { - AuthenticationException exception = authenticationException; - authenticationException = null; - return exception; - } else - return null; - } - - synchronized KafkaException getAndClearMetadataException() { - if (this.metadataException != null) { - KafkaException metadataException = this.metadataException; - this.metadataException = null; - return metadataException; - } else - return null; - } - - public synchronized void bootstrap(List addresses, long now) { + public synchronized void bootstrap(List addresses) { this.needUpdate = true; - this.lastRefreshMs = now; - this.lastSuccessfulRefreshMs = now; this.updateVersion += 1; this.cache = MetadataCache.bootstrap(addresses); } @@ -251,6 +236,8 @@ public synchronized void update(int requestVersion, MetadataResponse response, l if (isClosed()) throw new IllegalStateException("Update requested after metadata close"); + validateCluster(response.clusterId()); + if (requestVersion == this.requestVersion) this.needUpdate = false; else @@ -278,9 +265,37 @@ public synchronized void update(int requestVersion, MetadataResponse response, l log.debug("Updated cluster metadata updateVersion {} to {}", this.updateVersion, this.cache); } + private void validateCluster(String newClusterId) { + String previousClusterId = this.cache.cluster().clusterResource().clusterId(); + + if (previousClusterId != null && newClusterId != null && !previousClusterId.equals(newClusterId)) { + // kafka cluster id is unique. + // On client side, the cluster id in Metadata is only null during bootstrap, and client is + // expected to talk to the same cluster during its life cycle. Therefore if cluster id changes + // during metadata update, meaning this metadata update response is from a different cluster, + // client should reject this response and not update cached cluster to the wrong cluster + + // Since removing brokers and adding to another cluster can be common operation, + // it might be more suitable to just throw an exception for update and fail this update operation + // instead of bringing down the client completely, so that the metadata can be updated later from + // other brokers in the same cluster. + + // this code path will only get executed when all brokers in original cluster has been removed and + // one/some brokers have been added to another. If only some of the original brokers were removed/added + // to another cluster, the client should get updated metadata with valid brokers from other hosts. + // so we can just throw an exception and close the network client + + log.error("Received metadata from a different cluster {}, current cluster {} has no valid brokers anymore," + + "please reboot the producer/consumer", newClusterId, previousClusterId); + + throw new StaleClusterMetadataException( + "Trying to access a different cluster " + newClusterId + ", previous connected cluster " + previousClusterId); + + } + } + private void maybeSetMetadataError(Cluster cluster) { - // if we encounter any invalid topics, cache the exception to later throw to the user - metadataException = null; + clearRecoverableErrors(); checkInvalidTopics(cluster); checkUnauthorizedTopics(cluster); } @@ -288,14 +303,14 @@ private void maybeSetMetadataError(Cluster cluster) { private void checkInvalidTopics(Cluster cluster) { if (!cluster.invalidTopics().isEmpty()) { log.error("Metadata response reported invalid topics {}", cluster.invalidTopics()); - metadataException = new InvalidTopicException(cluster.invalidTopics()); + invalidTopics = new HashSet<>(cluster.invalidTopics()); } } private void checkUnauthorizedTopics(Cluster cluster) { if (!cluster.unauthorizedTopics().isEmpty()) { log.error("Topic authorization failed for topics {}", cluster.unauthorizedTopics()); - metadataException = new TopicAuthorizationException(new HashSet<>(cluster.unauthorizedTopics())); + unauthorizedTopics = new HashSet<>(cluster.unauthorizedTopics()); } } @@ -306,6 +321,8 @@ private MetadataCache handleMetadataResponse(MetadataResponse metadataResponse, Predicate topicsToRetain) { Set internalTopics = new HashSet<>(); List partitions = new ArrayList<>(); + Map brokersById = metadataResponse.brokersById(); + for (MetadataResponse.TopicMetadata metadata : metadataResponse.topicMetadata()) { if (!topicsToRetain.test(metadata)) continue; @@ -313,11 +330,32 @@ private MetadataCache handleMetadataResponse(MetadataResponse metadataResponse, if (metadata.error() == Errors.NONE) { if (metadata.isInternal()) internalTopics.add(metadata.topic()); + for (MetadataResponse.PartitionMetadata partitionMetadata : metadata.partitionMetadata()) { - updatePartitionInfo(metadata.topic(), partitionMetadata, partitionInfo -> { - int epoch = partitionMetadata.leaderEpoch().orElse(RecordBatch.NO_PARTITION_LEADER_EPOCH); - partitions.add(new MetadataCache.PartitionInfoAndEpoch(partitionInfo, epoch)); - }); + // Even if the partition's metadata includes an error, we need to handle the update to catch new epochs + updatePartitionInfo(metadata.topic(), partitionMetadata, + metadataResponse.hasReliableLeaderEpochs(), partitionInfoAndEpoch -> { + Node leader = partitionInfoAndEpoch.partitionInfo().leader(); + + if (leader != null && !leader.equals(brokersById.get(leader.id()))) { + // If we are reusing metadata from a previous response (which is possible if it + // contained a larger epoch), we may not have leader information available in the + // latest response. To keep the state consistent, we override the partition metadata + // so that the leader is set consistently with the broker metadata + PartitionInfo partitionInfo = partitionInfoAndEpoch.partitionInfo(); + PartitionInfo partitionInfoWithoutLeader = new PartitionInfo( + partitionInfo.topic(), + partitionInfo.partition(), + brokersById.get(leader.id()), + partitionInfo.replicas(), + partitionInfo.inSyncReplicas(), + partitionInfo.offlineReplicas()); + partitions.add(new MetadataCache.PartitionInfoAndEpoch(partitionInfoWithoutLeader, + partitionInfoAndEpoch.epoch())); + } else { + partitions.add(partitionInfoAndEpoch); + } + }); if (partitionMetadata.error().exception() instanceof InvalidMetadataException) { log.debug("Requesting metadata update for partition {} due to error {}", @@ -331,7 +369,7 @@ private MetadataCache handleMetadataResponse(MetadataResponse metadataResponse, } } - return new MetadataCache(metadataResponse.clusterId(), new ArrayList<>(metadataResponse.brokers()), partitions, + return new MetadataCache(metadataResponse.clusterId(), brokersById.values(), partitions, metadataResponse.topicsByError(Errors.TOPIC_AUTHORIZATION_FAILED), metadataResponse.topicsByError(Errors.INVALID_TOPIC_EXCEPTION), internalTopics, metadataResponse.controller()); @@ -342,45 +380,108 @@ private MetadataCache handleMetadataResponse(MetadataResponse metadataResponse, */ private void updatePartitionInfo(String topic, MetadataResponse.PartitionMetadata partitionMetadata, - Consumer partitionInfoConsumer) { - + boolean hasReliableLeaderEpoch, + Consumer partitionInfoConsumer) { TopicPartition tp = new TopicPartition(topic, partitionMetadata.partition()); - if (partitionMetadata.leaderEpoch().isPresent()) { + + if (hasReliableLeaderEpoch && partitionMetadata.leaderEpoch().isPresent()) { int newEpoch = partitionMetadata.leaderEpoch().get(); // If the received leader epoch is at least the same as the previous one, update the metadata if (updateLastSeenEpoch(tp, newEpoch, oldEpoch -> newEpoch >= oldEpoch, false)) { - partitionInfoConsumer.accept(MetadataResponse.partitionMetaToInfo(topic, partitionMetadata)); + PartitionInfo info = MetadataResponse.partitionMetaToInfo(topic, partitionMetadata); + partitionInfoConsumer.accept(new MetadataCache.PartitionInfoAndEpoch(info, newEpoch)); } else { // Otherwise ignore the new metadata and use the previously cached info - PartitionInfo previousInfo = cache.cluster().partition(tp); - if (previousInfo != null) { - partitionInfoConsumer.accept(previousInfo); - } + cache.getPartitionInfo(tp).ifPresent(partitionInfoConsumer); } } else { - // Old cluster format (no epochs) - lastSeenLeaderEpochs.clear(); - partitionInfoConsumer.accept(MetadataResponse.partitionMetaToInfo(topic, partitionMetadata)); + // Handle old cluster formats as well as error responses where leader and epoch are missing + lastSeenLeaderEpochs.remove(tp); + PartitionInfo info = MetadataResponse.partitionMetaToInfo(topic, partitionMetadata); + partitionInfoConsumer.accept(new MetadataCache.PartitionInfoAndEpoch(info, + RecordBatch.NO_PARTITION_LEADER_EPOCH)); + } + } + + /** + * If any non-retriable exceptions were encountered during metadata update, clear and throw the exception. + * This is used by the consumer to propagate any fatal exceptions or topic exceptions for any of the topics + * in the consumer's Metadata. + */ + public synchronized void maybeThrowAnyException() { + clearErrorsAndMaybeThrowException(this::recoverableException); + } + + /** + * If any fatal exceptions were encountered during metadata update, throw the exception. This is used by + * the producer to abort waiting for metadata if there were fatal exceptions (e.g. authentication failures) + * in the last metadata update. + */ + public synchronized void maybeThrowFatalException() { + KafkaException metadataException = this.fatalException; + if (metadataException != null) { + fatalException = null; + throw metadataException; } } - public synchronized void maybeThrowException() { - AuthenticationException authenticationException = getAndClearAuthenticationException(); - if (authenticationException != null) - throw authenticationException; + /** + * If any non-retriable exceptions were encountered during metadata update, throw exception if the exception + * is fatal or related to the specified topic. All exceptions from the last metadata update are cleared. + * This is used by the producer to propagate topic metadata errors for send requests. + */ + public synchronized void maybeThrowExceptionForTopic(String topic) { + clearErrorsAndMaybeThrowException(() -> recoverableExceptionForTopic(topic)); + } - KafkaException metadataException = getAndClearMetadataException(); + private void clearErrorsAndMaybeThrowException(Supplier recoverableExceptionSupplier) { + KafkaException metadataException = Optional.ofNullable(fatalException).orElseGet(recoverableExceptionSupplier); + fatalException = null; + clearRecoverableErrors(); if (metadataException != null) throw metadataException; } + // We may be able to recover from this exception if metadata for this topic is no longer needed + private KafkaException recoverableException() { + if (!unauthorizedTopics.isEmpty()) + return new TopicAuthorizationException(unauthorizedTopics); + else if (!invalidTopics.isEmpty()) + return new InvalidTopicException(invalidTopics); + else + return null; + } + + private KafkaException recoverableExceptionForTopic(String topic) { + if (unauthorizedTopics.contains(topic)) + return new TopicAuthorizationException(Collections.singleton(topic)); + else if (invalidTopics.contains(topic)) + return new InvalidTopicException(Collections.singleton(topic)); + else + return null; + } + + private void clearRecoverableErrors() { + invalidTopics = Collections.emptySet(); + unauthorizedTopics = Collections.emptySet(); + } + /** * Record an attempt to update the metadata that failed. We need to keep track of this * to avoid retrying immediately. */ - public synchronized void failedUpdate(long now, AuthenticationException authenticationException) { + public synchronized void failedUpdate(long now) { this.lastRefreshMs = now; - this.authenticationException = authenticationException; + } + + /** + * Propagate a fatal error which affects the ability to fetch metadata for the cluster. + * Two examples are authentication and unsupported version exceptions. + * + * @param exception The fatal exception + */ + public synchronized void fatalError(KafkaException exception) { + this.fatalException = exception; } /** @@ -444,4 +545,55 @@ private MetadataRequestAndVersion(MetadataRequest.Builder requestBuilder, } } + public synchronized LeaderAndEpoch leaderAndEpoch(TopicPartition tp) { + return partitionInfoIfCurrent(tp) + .map(infoAndEpoch -> { + Node leader = infoAndEpoch.partitionInfo().leader(); + return new LeaderAndEpoch(leader == null ? Node.noNode() : leader, Optional.of(infoAndEpoch.epoch())); + }) + .orElse(new LeaderAndEpoch(Node.noNode(), lastSeenLeaderEpoch(tp))); + } + + public static class LeaderAndEpoch { + + public static final LeaderAndEpoch NO_LEADER_OR_EPOCH = new LeaderAndEpoch(Node.noNode(), Optional.empty()); + + public final Node leader; + public final Optional epoch; + + public LeaderAndEpoch(Node leader, Optional epoch) { + this.leader = Objects.requireNonNull(leader); + this.epoch = Objects.requireNonNull(epoch); + } + + public static LeaderAndEpoch noLeaderOrEpoch() { + return NO_LEADER_OR_EPOCH; + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + + LeaderAndEpoch that = (LeaderAndEpoch) o; + + if (!leader.equals(that.leader)) return false; + return epoch.equals(that.epoch); + } + + @Override + public int hashCode() { + int result = leader.hashCode(); + result = 31 * result + epoch.hashCode(); + return result; + } + + @Override + public String toString() { + return "LeaderAndEpoch{" + + "leader=" + leader + + ", epoch=" + epoch.map(Number::toString).orElse("absent") + + '}'; + } + } } diff --git a/clients/src/main/java/org/apache/kafka/clients/MetadataCache.java b/clients/src/main/java/org/apache/kafka/clients/MetadataCache.java index b928b8e0ef731..10b5cc92e7694 100644 --- a/clients/src/main/java/org/apache/kafka/clients/MetadataCache.java +++ b/clients/src/main/java/org/apache/kafka/clients/MetadataCache.java @@ -39,7 +39,7 @@ */ public class MetadataCache { private final String clusterId; - private final List nodes; + private final Collection nodes; private final Set unauthorizedTopics; private final Set invalidTopics; private final Set internalTopics; @@ -49,7 +49,7 @@ public class MetadataCache { private Cluster clusterInstance; MetadataCache(String clusterId, - List nodes, + Collection nodes, Collection partitions, Set unauthorizedTopics, Set invalidTopics, @@ -59,7 +59,7 @@ public class MetadataCache { } MetadataCache(String clusterId, - List nodes, + Collection nodes, Collection partitions, Set unauthorizedTopics, Set invalidTopics, @@ -97,13 +97,6 @@ Optional getPartitionInfo(TopicPartition topicPartition) return Optional.ofNullable(metadataByPartition.get(topicPartition)); } - synchronized void retainTopics(Collection topics) { - metadataByPartition.entrySet().removeIf(entry -> !topics.contains(entry.getKey().topic())); - unauthorizedTopics.retainAll(topics); - invalidTopics.retainAll(topics); - computeClusterView(); - } - Cluster cluster() { if (clusterInstance == null) { throw new IllegalStateException("Cached Cluster instance should not be null, but was."); @@ -137,7 +130,10 @@ static MetadataCache empty() { @Override public String toString() { return "MetadataCache{" + - "cluster=" + cluster() + + "clusterId='" + clusterId + '\'' + + ", nodes=" + nodes + + ", partitions=" + metadataByPartition.values() + + ", controller=" + controller + '}'; } diff --git a/clients/src/main/java/org/apache/kafka/clients/MetadataUpdater.java b/clients/src/main/java/org/apache/kafka/clients/MetadataUpdater.java index de765db5a8db3..77f3efadce3a0 100644 --- a/clients/src/main/java/org/apache/kafka/clients/MetadataUpdater.java +++ b/clients/src/main/java/org/apache/kafka/clients/MetadataUpdater.java @@ -16,13 +16,16 @@ */ package org.apache.kafka.clients; +import org.apache.kafka.common.KafkaException; import org.apache.kafka.common.Node; import org.apache.kafka.common.errors.AuthenticationException; +import org.apache.kafka.common.errors.UnsupportedVersionException; import org.apache.kafka.common.requests.MetadataResponse; import org.apache.kafka.common.requests.RequestHeader; import java.io.Closeable; import java.util.List; +import java.util.Optional; /** * The interface used by `NetworkClient` to request cluster metadata info to be updated and to retrieve the cluster nodes @@ -46,7 +49,7 @@ public interface MetadataUpdater extends Closeable { * Starts a cluster metadata update if needed and possible. Returns the time until the metadata update (which would * be 0 if an update has been started as a result of this call). * - * If the implementation relies on `NetworkClient` to send requests, `handleCompletedMetadataResponse` will be + * If the implementation relies on `NetworkClient` to send requests, `handleSuccessfulResponse` will be * invoked after the metadata response is received. * * The semantics of `needed` and `possible` are implementation-dependent and may take into account a number of @@ -55,20 +58,24 @@ public interface MetadataUpdater extends Closeable { long maybeUpdate(long now); /** - * Handle disconnections for metadata requests. + * Handle a server disconnect. * * This provides a mechanism for the `MetadataUpdater` implementation to use the NetworkClient instance for its own * requests with special handling for disconnections of such requests. - * @param destination + * + * @param now Current time in milliseconds + * @param nodeId The id of the node that disconnected + * @param maybeAuthException Optional authentication error */ - void handleDisconnection(String destination); + void handleServerDisconnect(long now, String nodeId, Optional maybeAuthException); /** - * Handle authentication failure. Propagate the authentication exception if awaiting metadata. + * Handle a metadata request failure. * - * @param exception authentication exception from broker + * @param now Current time in milliseconds + * @param maybeFatalException Optional fatal error (e.g. {@link UnsupportedVersionException}) */ - void handleAuthenticationFailure(AuthenticationException exception); + void handleFailedRequest(long now, Optional maybeFatalException); /** * Handle responses for metadata requests. @@ -76,13 +83,7 @@ public interface MetadataUpdater extends Closeable { * This provides a mechanism for the `MetadataUpdater` implementation to use the NetworkClient instance for its own * requests with special handling for completed receives of such requests. */ - void handleCompletedMetadataResponse(RequestHeader requestHeader, long now, MetadataResponse metadataResponse); - - /** - * Schedules an update of the current cluster metadata info. A subsequent call to `maybeUpdate` would trigger the - * start of the update if possible (see `maybeUpdate` for more information). - */ - void requestUpdate(); + void handleSuccessfulResponse(RequestHeader requestHeader, long now, MetadataResponse metadataResponse); /** * Close this updater. diff --git a/clients/src/main/java/org/apache/kafka/clients/NetworkClient.java b/clients/src/main/java/org/apache/kafka/clients/NetworkClient.java index e33429585c3b6..5d4193cb1ce26 100644 --- a/clients/src/main/java/org/apache/kafka/clients/NetworkClient.java +++ b/clients/src/main/java/org/apache/kafka/clients/NetworkClient.java @@ -17,19 +17,22 @@ package org.apache.kafka.clients; import org.apache.kafka.common.Cluster; +import org.apache.kafka.common.KafkaException; import org.apache.kafka.common.Node; import org.apache.kafka.common.TopicPartition; import org.apache.kafka.common.errors.AuthenticationException; import org.apache.kafka.common.errors.DisconnectException; import org.apache.kafka.common.errors.UnsupportedVersionException; +import org.apache.kafka.common.message.ApiVersionsResponseData.ApiVersionsResponseKey; import org.apache.kafka.common.metrics.Sensor; import org.apache.kafka.common.network.ChannelState; import org.apache.kafka.common.network.NetworkReceive; import org.apache.kafka.common.network.Selectable; 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.protocol.CommonFields; +import org.apache.kafka.common.protocol.Errors; +import org.apache.kafka.common.protocol.types.SchemaException; import org.apache.kafka.common.protocol.types.Struct; import org.apache.kafka.common.requests.AbstractRequest; import org.apache.kafka.common.requests.AbstractResponse; @@ -39,6 +42,7 @@ import org.apache.kafka.common.requests.MetadataResponse; import org.apache.kafka.common.requests.RequestHeader; import org.apache.kafka.common.requests.ResponseHeader; +import org.apache.kafka.common.security.authenticator.SaslClientAuthenticator; import org.apache.kafka.common.utils.LogContext; import org.apache.kafka.common.utils.Time; import org.apache.kafka.common.utils.Utils; @@ -49,11 +53,13 @@ import java.net.InetSocketAddress; import java.nio.ByteBuffer; import java.util.ArrayList; +import java.util.Collection; import java.util.HashMap; import java.util.Iterator; import java.util.LinkedList; import java.util.List; import java.util.Map; +import java.util.Optional; import java.util.Random; import java.util.concurrent.atomic.AtomicReference; import java.util.stream.Collectors; @@ -109,6 +115,8 @@ private enum State { private final Time time; + private boolean enableStickyMetadataFetch = true; + /** * True if we should send an ApiVersionRequest when first connecting to a broker. */ @@ -267,6 +275,10 @@ private NetworkClient(MetadataUpdater metadataUpdater, this.state = new AtomicReference<>(State.ACTIVE); } + public void setEnableStickyMetadataFetch(boolean enableStickyMetadataFetch) { + this.enableStickyMetadataFetch = enableStickyMetadataFetch; + } + /** * Begin connecting to the given node, return true if we are already connected and ready to send to that node. * @@ -306,24 +318,29 @@ public void disconnect(String nodeId) { return; selector.close(nodeId); - List requestTypes = new ArrayList<>(); long now = time.milliseconds(); - for (InFlightRequest request : inFlightRequests.clearAll(nodeId)) { - if (request.isInternalRequest) { - if (request.header.apiKey() == ApiKeys.METADATA) { - metadataUpdater.handleDisconnection(request.destination); - } - } else { - requestTypes.add(request.header.apiKey()); - abortedSends.add(new ClientResponse(request.header, - request.callback, request.destination, request.createdTimeMs, now, - true, null, null, null)); - } - } + + cancelInFlightRequests(nodeId, now, abortedSends); + connectionStates.disconnected(nodeId, now); - if (log.isDebugEnabled()) { - log.debug("Manually disconnected from {}. Removed requests: {}.", nodeId, - Utils.join(requestTypes, ", ")); + + if (log.isTraceEnabled()) { + log.trace("Manually disconnected from {}. Aborted in-flight requests: {}.", nodeId, inFlightRequests); + } + } + + private void cancelInFlightRequests(String nodeId, long now, Collection responses) { + Iterable inFlightRequests = this.inFlightRequests.clearAll(nodeId); + for (InFlightRequest request : inFlightRequests) { + log.trace("Cancelled request {} {} with correlation id {} due to node {} being disconnected", + request.header.apiKey(), request.request, request.header.correlationId(), nodeId); + + if (!request.isInternalRequest) { + if (responses != null) + responses.add(request.disconnected(now, null)); + } else if (request.header.apiKey() == ApiKeys.METADATA) { + metadataUpdater.handleFailedRequest(now, Optional.empty()); + } } } @@ -337,9 +354,8 @@ public void disconnect(String nodeId) { @Override public void close(String nodeId) { selector.close(nodeId); - for (InFlightRequest request : inFlightRequests.clearAll(nodeId)) - if (request.isInternalRequest && request.header.apiKey() == ApiKeys.METADATA) - metadataUpdater.handleDisconnection(request.destination); + long now = time.milliseconds(); + cancelInFlightRequests(nodeId, now, null); connectionStates.remove(nodeId); } @@ -433,8 +449,8 @@ public void send(ClientRequest request, long now) { doSend(request, false, now); } - private void sendInternalMetadataRequest(MetadataRequest.Builder builder, - String nodeConnectionId, long now) { + // package-private for testing + void sendInternalMetadataRequest(MetadataRequest.Builder builder, String nodeConnectionId, long now) { ClientRequest clientRequest = newClientRequest(nodeConnectionId, builder, now, true); doSend(clientRequest, true, now); } @@ -479,7 +495,11 @@ private void doSend(ClientRequest clientRequest, boolean isInternalRequest, long ClientResponse clientResponse = new ClientResponse(clientRequest.makeHeader(builder.latestAllowedVersion()), clientRequest.callback(), clientRequest.destination(), now, now, false, unsupportedVersionException, null, null); - abortedSends.add(clientResponse); + + if (!isInternalRequest) + abortedSends.add(clientResponse); + else if (clientRequest.apiKey() == ApiKeys.METADATA) + metadataUpdater.handleFailedRequest(now, Optional.of(unsupportedVersionException)); } } @@ -541,13 +561,28 @@ public List poll(long timeout, long now) { long updatedNow = this.time.milliseconds(); List responses = new ArrayList<>(); handleCompletedSends(responses, updatedNow); - handleCompletedReceives(responses, updatedNow); + + try { + handleCompletedReceives(responses, updatedNow); + } catch (StaleClusterMetadataException e) { + // upon stale metadata exception from a different cluster, close the network client + // the producer/consumer will hit closedSelector exception and close + log.error("Received stale metadata from a different cluster, close the network client now"); + this.close(); + } + handleDisconnections(responses, updatedNow); handleConnections(); handleInitiateApiVersionRequests(updatedNow); handleTimedOutRequests(responses, updatedNow); completeResponses(responses); + // We changed the metadataUpdater.maybeUpdate() such that it will keep sending MetadataRequest + // to the same broker instead choosing the least loaded node. If we don't try to send metadata here, it is possible that + // another request is sent to the broker before the next networkClient.poll(). This can cause starvation + // for the MetadataRequest and consumer's metadata may be stale for a long time. + metadataUpdater.maybeUpdate(updatedNow); + return responses; } @@ -635,7 +670,8 @@ public void close() { * Choose the node with the fewest outstanding requests which is at least eligible for connection. This method will * prefer a node with an existing connection, but will potentially choose a node for which we don't yet have a * connection if all existing connections are in use. This method will never choose a node for which there is no - * existing connection and from which we have disconnected within the reconnect backoff period. + * existing connection and from which we have disconnected within the reconnect backoff period, or an active + * connection which is being throttled. * * @return The node with the fewest in-flight requests. */ @@ -645,33 +681,51 @@ public Node leastLoadedNode(long now) { if (nodes.isEmpty()) throw new IllegalStateException("There are no nodes in the Kafka cluster"); int inflight = Integer.MAX_VALUE; - Node found = null; + + Node foundConnecting = null; + Node foundCanConnect = null; + Node foundReady = null; int offset = this.randOffset.nextInt(nodes.size()); for (int i = 0; i < nodes.size(); i++) { int idx = (offset + i) % nodes.size(); Node node = nodes.get(idx); - int currInflight = this.inFlightRequests.count(node.idString()); - if (currInflight == 0 && isReady(node, now)) { - // if we find an established connection with no in-flight requests we can stop right away - log.trace("Found least loaded node {} connected with no in-flight requests", node); - return node; - } else if (!this.connectionStates.isBlackedOut(node.idString(), now) && currInflight < inflight) { - // otherwise if this is the best we have found so far, record that - inflight = currInflight; - found = node; - } else if (log.isTraceEnabled()) { - log.trace("Removing node {} from least loaded node selection: is-blacked-out: {}, in-flight-requests: {}", - node, this.connectionStates.isBlackedOut(node.idString(), now), currInflight); + if (canSendRequest(node.idString(), now)) { + int currInflight = this.inFlightRequests.count(node.idString()); + if (currInflight == 0) { + // if we find an established connection with no in-flight requests we can stop right away + log.trace("Found least loaded node {} connected with no in-flight requests", node); + return node; + } else if (currInflight < inflight) { + // otherwise if this is the best we have found so far, record that + inflight = currInflight; + foundReady = node; + } + } else if (connectionStates.isPreparingConnection(node.idString())) { + foundConnecting = node; + } else if (canConnect(node, now)) { + foundCanConnect = node; + } else { + log.trace("Removing node {} from least loaded node selection since it is neither ready " + + "for sending or connecting", node); } } - if (found != null) - log.trace("Found least loaded node {}", found); - else + // We prefer established connections if possible. Otherwise, we will wait for connections + // which are being established before connecting to new nodes. + if (foundReady != null) { + log.trace("Found least loaded node {} with {} inflight requests", foundReady, inflight); + return foundReady; + } else if (foundConnecting != null) { + log.trace("Found least loaded connecting node {}", foundConnecting); + return foundConnecting; + } else if (foundCanConnect != null) { + log.trace("Found least loaded node {} with no active connection", foundCanConnect); + return foundCanConnect; + } else { log.trace("Least loaded node selection failed to find an available node"); - - return found; + return null; + } } public static AbstractResponse parseResponse(ByteBuffer responseBuffer, RequestHeader requestHeader) { @@ -682,7 +736,8 @@ public static AbstractResponse parseResponse(ByteBuffer responseBuffer, RequestH private static Struct parseStructMaybeUpdateThrottleTimeMetrics(ByteBuffer responseBuffer, RequestHeader requestHeader, Sensor throttleTimeSensor, long now) { - ResponseHeader responseHeader = ResponseHeader.parse(responseBuffer); + ResponseHeader responseHeader = ResponseHeader.parse(responseBuffer, + requestHeader.apiKey().responseHeaderVersion(requestHeader.apiVersion())); // Always expect the response version id to be the same as the request version id Struct responseBody = requestHeader.apiKey().parseResponse(requestHeader.apiVersion(), responseBuffer); correlate(requestHeader, responseHeader); @@ -710,7 +765,6 @@ private void processDisconnection(List responses, case AUTHENTICATION_FAILED: AuthenticationException exception = disconnectState.exception(); connectionStates.authenticationFailed(nodeId, now, exception); - metadataUpdater.handleAuthenticationFailure(exception); log.error("Connection to node {} ({}) failed authentication due to: {}", nodeId, disconnectState.remoteAddress(), exception.getMessage()); break; @@ -727,14 +781,9 @@ private void processDisconnection(List responses, default: break; // Disconnections in other states are logged at debug level in Selector } - for (InFlightRequest request : this.inFlightRequests.clearAll(nodeId)) { - log.trace("Cancelled request {} {} with correlation id {} due to node {} being disconnected", - request.header.apiKey(), request.request, request.header.correlationId(), nodeId); - if (!request.isInternalRequest) - responses.add(request.disconnected(now, disconnectState.exception())); - else if (request.header.apiKey() == ApiKeys.METADATA) - metadataUpdater.handleDisconnection(request.destination); - } + + cancelInFlightRequests(nodeId, now, responses); + metadataUpdater.handleServerDisconnect(now, nodeId, Optional.ofNullable(disconnectState.exception())); } /** @@ -752,10 +801,6 @@ private void handleTimedOutRequests(List responses, long now) { log.debug("Disconnecting from node {} due to request timeout.", nodeId); processDisconnection(responses, nodeId, now, ChannelState.LOCAL_CLOSE); } - - // we disconnected, so we should probably refresh our metadata - if (!nodeIds.isEmpty()) - metadataUpdater.requestUpdate(); } private void handleAbortedSends(List responses) { @@ -819,7 +864,7 @@ private void handleCompletedReceives(List responses, long now) { parseResponse(req.header.apiKey(), responseStruct, req.header.apiVersion()); maybeThrottle(body, req.header.apiVersion(), req.destination, now); if (req.isInternalRequest && body instanceof MetadataResponse) - metadataUpdater.handleCompletedMetadataResponse(req.header, now, (MetadataResponse) body); + metadataUpdater.handleSuccessfulResponse(req.header, now, (MetadataResponse) body); else if (req.isInternalRequest && body instanceof ApiVersionsResponse) handleApiVersionsResponse(responses, req, now, (ApiVersionsResponse) body); else @@ -830,18 +875,28 @@ else if (req.isInternalRequest && body instanceof ApiVersionsResponse) private void handleApiVersionsResponse(List responses, InFlightRequest req, long now, ApiVersionsResponse apiVersionsResponse) { final String node = req.destination; - if (apiVersionsResponse.error() != Errors.NONE) { - if (req.request.version() == 0 || apiVersionsResponse.error() != Errors.UNSUPPORTED_VERSION) { + if (apiVersionsResponse.data.errorCode() != Errors.NONE.code()) { + if (req.request.version() == 0 || apiVersionsResponse.data.errorCode() != Errors.UNSUPPORTED_VERSION.code()) { log.warn("Received error {} from node {} when making an ApiVersionsRequest with correlation id {}. Disconnecting.", - apiVersionsResponse.error(), node, req.header.correlationId()); + Errors.forCode(apiVersionsResponse.data.errorCode()), node, req.header.correlationId()); this.selector.close(node); processDisconnection(responses, node, now, ChannelState.LOCAL_CLOSE); } else { - nodesNeedingApiVersionsFetch.put(node, new ApiVersionsRequest.Builder((short) 0)); + // Starting from Apache Kafka 2.4, ApiKeys field is populated with the supported versions of + // the ApiVersionsRequest when an UNSUPPORTED_VERSION error is returned. + // If not provided, the client falls back to version 0. + short maxApiVersion = 0; + if (apiVersionsResponse.data.apiKeys().size() > 0) { + ApiVersionsResponseKey apiVersion = apiVersionsResponse.data.apiKeys().find(ApiKeys.API_VERSIONS.id); + if (apiVersion != null) { + maxApiVersion = apiVersion.maxVersion(); + } + } + nodesNeedingApiVersionsFetch.put(node, new ApiVersionsRequest.Builder(maxApiVersion)); } return; } - NodeApiVersions nodeVersionInfo = new NodeApiVersions(apiVersionsResponse.apiVersions()); + NodeApiVersions nodeVersionInfo = new NodeApiVersions(apiVersionsResponse.data.apiKeys()); apiVersions.update(node, nodeVersionInfo); this.connectionStates.ready(node); log.debug("Recorded API versions for node {}: {}", node, nodeVersionInfo); @@ -859,9 +914,6 @@ private void handleDisconnections(List responses, long now) { log.debug("Node {} disconnected.", node); processDisconnection(responses, node, now, entry.getValue()); } - // we got a disconnect so we should probably refresh our metadata and see if that broker is dead - if (this.selector.disconnected().size() > 0) - metadataUpdater.requestUpdate(); } /** @@ -869,7 +921,7 @@ private void handleDisconnections(List responses, long now) { */ private void handleConnections() { for (String node : this.selector.connected()) { - // We are now connected. Node that we might not still be able to send requests. For instance, + // We are now connected. Note that we might not still be able to send requests. For instance, // if SSL is enabled, the SSL handshake happens after the connection is established. // Therefore, it is still necessary to check isChannelReady before attempting to send on this // connection. @@ -903,9 +955,15 @@ private void handleInitiateApiVersionRequests(long now) { * Validate that the response corresponds to the request we expect or else explode */ private static void correlate(RequestHeader requestHeader, ResponseHeader responseHeader) { - if (requestHeader.correlationId() != responseHeader.correlationId()) + if (requestHeader.correlationId() != responseHeader.correlationId()) { + if (SaslClientAuthenticator.isReserved(requestHeader.correlationId()) + && !SaslClientAuthenticator.isReserved(responseHeader.correlationId())) + throw new SchemaException("the response is unrelated to Sasl request since its correlation id is " + responseHeader.correlationId() + + " and the reserved range for Sasl request is [ " + + SaslClientAuthenticator.MIN_RESERVED_CORRELATION_ID + "," + SaslClientAuthenticator.MAX_RESERVED_CORRELATION_ID + "]"); throw new IllegalStateException("Correlation id for response (" + responseHeader.correlationId() + ") does not match request (" + requestHeader.correlationId() + "), request header: " + requestHeader); + } } /** @@ -925,10 +983,10 @@ private void initiateConnect(Node node, long now) { this.socketReceiveBuffer); } catch (IOException e) { log.warn("Error connecting to node {}", node, e); - /* attempt failed, we'll try again after the backoff */ + // Attempt failed, we'll try again after the backoff connectionStates.disconnected(nodeConnectionId, now); - /* maybe the problem is our metadata, update it */ - metadataUpdater.requestUpdate(); + // Notify metadata updater of the connection failure + metadataUpdater.handleServerDisconnect(now, nodeConnectionId, Optional.empty()); } } @@ -936,6 +994,8 @@ class DefaultMetadataUpdater implements MetadataUpdater { /* the current cluster metadata */ private final Metadata metadata; + // Consumer needs to keep fetching metadata from the same node until that node goes down + private Node nodeToFetchMetadata; // Defined if there is a request in progress, null otherwise private Integer inProgressRequestVersion; @@ -943,6 +1003,7 @@ class DefaultMetadataUpdater implements MetadataUpdater { DefaultMetadataUpdater(Metadata metadata) { this.metadata = metadata; this.inProgressRequestVersion = null; + this.nodeToFetchMetadata = null; } @Override @@ -966,50 +1027,60 @@ public long maybeUpdate(long now) { long waitForMetadataFetch = hasFetchInProgress() ? defaultRequestTimeoutMs : 0; long metadataTimeout = Math.max(timeToNextMetadataUpdate, waitForMetadataFetch); - if (metadataTimeout > 0) { return metadataTimeout; } // Beware that the behavior of this method and the computation of timeouts for poll() are // highly dependent on the behavior of leastLoadedNode. - Node node = leastLoadedNode(now); - if (node == null) { + if (!enableStickyMetadataFetch || nodeToFetchMetadata == null || !connectionStates.isReady(nodeToFetchMetadata.idString(), now)) + nodeToFetchMetadata = leastLoadedNode(now); + + if (nodeToFetchMetadata == null) { log.debug("Give up sending metadata request since no node is available"); return reconnectBackoffMs; } - return maybeUpdate(now, node); + return maybeUpdate(now, nodeToFetchMetadata); } @Override - public void handleDisconnection(String destination) { + public void handleServerDisconnect(long now, String destinationId, Optional maybeFatalException) { Cluster cluster = metadata.fetch(); // 'processDisconnection' generates warnings for misconfigured bootstrap server configuration // resulting in 'Connection Refused' and misconfigured security resulting in authentication failures. - // The warning below handles the case where connection to a broker was established, but was disconnected + // The warning below handles the case where a connection to a broker was established, but was disconnected // before metadata could be obtained. if (cluster.isBootstrapConfigured()) { - int nodeId = Integer.parseInt(destination); + int nodeId = Integer.parseInt(destinationId); Node node = cluster.nodeById(nodeId); if (node != null) log.warn("Bootstrap broker {} disconnected", node); } - inProgressRequestVersion = null; + // If we have a disconnect while an update is due, we treat it as a failed update + // so that we can backoff properly + if (isUpdateDue(now)) + handleFailedRequest(now, Optional.empty()); + + maybeFatalException.ifPresent(metadata::fatalError); + + // The disconnect may be the result of stale metadata, so request an update + metadata.requestUpdate(); } @Override - public void handleAuthenticationFailure(AuthenticationException exception) { - if (metadata.updateRequested()) - metadata.failedUpdate(time.milliseconds(), exception); + public void handleFailedRequest(long now, Optional maybeFatalException) { + maybeFatalException.ifPresent(metadata::fatalError); + metadata.failedUpdate(now); inProgressRequestVersion = null; } @Override - public void handleCompletedMetadataResponse(RequestHeader requestHeader, long now, MetadataResponse response) { - // If any partition has leader with missing listeners, log a few for diagnosing broker configuration - // issues. This could be a transient issue if listeners were added dynamically to brokers. + public void handleSuccessfulResponse(RequestHeader requestHeader, long now, MetadataResponse response) { + // If any partition has leader with missing listeners, log up to ten of these partitions + // for diagnosing broker configuration issues. + // This could be a transient issue if listeners were added dynamically to brokers. List missingListenerPartitions = response.topicMetadata().stream().flatMap(topicMetadata -> topicMetadata.partitionMetadata().stream() .filter(partitionMetadata -> partitionMetadata.error() == Errors.LISTENER_NOT_FOUND) @@ -1021,16 +1092,16 @@ public void handleCompletedMetadataResponse(RequestHeader requestHeader, long no count, missingListenerPartitions.subList(0, Math.min(10, count))); } - // check if any topics metadata failed to get updated + // Check if any topic's metadata failed to get updated Map errors = response.errors(); if (!errors.isEmpty()) log.warn("Error while fetching metadata with correlation id {} : {}", requestHeader.correlationId(), errors); - // don't update the cluster if there are no valid nodes...the topic we want may still be in the process of being + // Don't update the cluster if there are no valid nodes...the topic we want may still be in the process of being // created which means we will get errors and no nodes until it exists if (response.brokers().isEmpty()) { log.trace("Ignoring empty metadata response with correlation id {}.", requestHeader.correlationId()); - this.metadata.failedUpdate(now, null); + this.metadata.failedUpdate(now); } else { this.metadata.update(inProgressRequestVersion, response, now); } @@ -1038,11 +1109,6 @@ public void handleCompletedMetadataResponse(RequestHeader requestHeader, long no inProgressRequestVersion = null; } - @Override - public void requestUpdate() { - this.metadata.requestUpdate(); - } - @Override public void close() { this.metadata.close(); @@ -1068,10 +1134,10 @@ private long maybeUpdate(long now, Node node) { if (canSendRequest(nodeConnectionId, now)) { Metadata.MetadataRequestAndVersion requestAndVersion = metadata.newMetadataRequestAndVersion(); - this.inProgressRequestVersion = requestAndVersion.requestVersion; MetadataRequest.Builder metadataRequest = requestAndVersion.requestBuilder; log.debug("Sending metadata request {} to node {}", metadataRequest, node); sendInternalMetadataRequest(metadataRequest, nodeConnectionId, now); + this.inProgressRequestVersion = requestAndVersion.requestVersion; return defaultRequestTimeoutMs; } @@ -1085,7 +1151,7 @@ private long maybeUpdate(long now, Node node) { } if (connectionStates.canConnect(nodeConnectionId, now)) { - // we don't have a connection to this node right now, make one + // We don't have a connection to this node right now, make one log.debug("Initialize connection to node {} for sending metadata request", node); initiateConnect(node, now); return reconnectBackoffMs; @@ -1107,6 +1173,15 @@ public ClientRequest newClientRequest(String nodeId, return newClientRequest(nodeId, requestBuilder, createdTimeMs, expectResponse, defaultRequestTimeoutMs, null); } + // visible for testing + int nextCorrelationId() { + if (SaslClientAuthenticator.isReserved(correlation)) { + // the numeric overflow is fine as negative values is acceptable + correlation = SaslClientAuthenticator.MAX_RESERVED_CORRELATION_ID + 1; + } + return correlation++; + } + @Override public ClientRequest newClientRequest(String nodeId, AbstractRequest.Builder requestBuilder, @@ -1114,7 +1189,7 @@ public ClientRequest newClientRequest(String nodeId, boolean expectResponse, int requestTimeoutMs, RequestCompletionHandler callback) { - return new ClientRequest(nodeId, requestBuilder, correlation++, clientId, createdTimeMs, expectResponse, + return new ClientRequest(nodeId, requestBuilder, nextCorrelationId(), clientId, createdTimeMs, expectResponse, requestTimeoutMs, callback); } diff --git a/clients/src/main/java/org/apache/kafka/clients/NodeApiVersions.java b/clients/src/main/java/org/apache/kafka/clients/NodeApiVersions.java index c8b3f44b77997..dcc5e5ffa2b2c 100644 --- a/clients/src/main/java/org/apache/kafka/clients/NodeApiVersions.java +++ b/clients/src/main/java/org/apache/kafka/clients/NodeApiVersions.java @@ -16,16 +16,17 @@ */ package org.apache.kafka.clients; +import java.util.Collection; +import java.util.Collections; +import java.util.LinkedList; import org.apache.kafka.common.errors.UnsupportedVersionException; +import org.apache.kafka.common.message.ApiVersionsResponseData.ApiVersionsResponseKey; +import org.apache.kafka.common.message.ApiVersionsResponseData.ApiVersionsResponseKeyCollection; import org.apache.kafka.common.protocol.ApiKeys; -import org.apache.kafka.common.requests.ApiVersionsResponse.ApiVersion; import org.apache.kafka.common.utils.Utils; import java.util.ArrayList; -import java.util.Collection; -import java.util.Collections; import java.util.EnumMap; -import java.util.LinkedList; import java.util.List; import java.util.Map; import java.util.TreeMap; @@ -34,6 +35,7 @@ * An internal class which represents the API versions supported by a particular node. */ public class NodeApiVersions { + // A map of the usable versions of each API, keyed by the ApiKeys instance private final Map supportedVersions = new EnumMap<>(ApiKeys.class); @@ -73,6 +75,31 @@ public static NodeApiVersions create(Collection overrides) { return new NodeApiVersions(apiVersions); } + + /** + * Create a NodeApiVersions object with a single ApiKey. It is mainly used in tests. + * + * @param apiKey ApiKey's id. + * @param minVersion ApiKey's minimum version. + * @param maxVersion ApiKey's maximum version. + * @return A new NodeApiVersions object. + */ + public static NodeApiVersions create(short apiKey, short minVersion, short maxVersion) { + return create(Collections.singleton(new ApiVersion(apiKey, minVersion, maxVersion))); + } + + public NodeApiVersions(ApiVersionsResponseKeyCollection nodeApiVersions) { + for (ApiVersionsResponseKey nodeApiVersion : nodeApiVersions) { + if (ApiKeys.hasId(nodeApiVersion.apiKey())) { + ApiKeys nodeApiKey = ApiKeys.forId(nodeApiVersion.apiKey()); + supportedVersions.put(nodeApiKey, new ApiVersion(nodeApiVersion)); + } else { + // Newer brokers may support ApiKeys we don't know about + unknownApis.add(new ApiVersion(nodeApiVersion)); + } + } + } + public NodeApiVersions(Collection nodeApiVersions) { for (ApiVersion nodeApiVersion : nodeApiVersions) { if (ApiKeys.hasId(nodeApiVersion.apiKey)) { diff --git a/clients/src/main/java/org/apache/kafka/clients/StaleClusterMetadataException.java b/clients/src/main/java/org/apache/kafka/clients/StaleClusterMetadataException.java new file mode 100644 index 0000000000000..6e012345b1738 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/StaleClusterMetadataException.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.clients; + +import org.apache.kafka.common.errors.InvalidMetadataException; + +/** + * Thrown when current metadata is from a stale cluster that has a different cluster id + * than original cluster. + * + * Note: this is not a public API. + */ +public class StaleClusterMetadataException extends InvalidMetadataException { + private static final long serialVersionUID = 1L; + + public StaleClusterMetadataException() {} + + public StaleClusterMetadataException(String message) { + super(message); + } +} 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 new file mode 100644 index 0000000000000..a0e43657f2748 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/Admin.java @@ -0,0 +1,1087 @@ +/* + * 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.admin; + +import java.time.Duration; +import java.util.Collection; +import java.util.HashSet; +import java.util.Map; +import java.util.Optional; +import java.util.Properties; +import java.util.Set; +import java.util.concurrent.TimeUnit; +import org.apache.kafka.common.ElectionType; +import org.apache.kafka.common.Metric; +import org.apache.kafka.common.MetricName; +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.TopicPartitionReplica; +import org.apache.kafka.common.acl.AclBinding; +import org.apache.kafka.common.acl.AclBindingFilter; +import org.apache.kafka.common.annotation.InterfaceStability; +import org.apache.kafka.common.config.ConfigResource; +import org.apache.kafka.common.requests.LeaveGroupResponse; + +/** + * The administrative client for Kafka, which supports managing and inspecting topics, brokers, configurations and ACLs. + *

+ * The minimum broker version required is 0.10.0.0. Methods with stricter requirements will specify the minimum broker + * version required. + *

+ * This client was introduced in 0.11.0.0 and the API is still evolving. We will try to evolve the API in a compatible + * manner, but we reserve the right to make breaking changes in minor releases, if necessary. We will update the + * {@code InterfaceStability} annotation and this notice once the API is considered stable. + */ +@InterfaceStability.Evolving +public interface Admin extends AutoCloseable { + + /** + * Create a new Admin with the given configuration. + * + * @param props The configuration. + * @return The new KafkaAdminClient. + */ + static Admin create(Properties props) { + return KafkaAdminClient.createInternal(new AdminClientConfig(props, true), null); + } + + /** + * Create a new Admin with the given configuration. + * + * @param conf The configuration. + * @return The new KafkaAdminClient. + */ + static Admin create(Map conf) { + return KafkaAdminClient.createInternal(new AdminClientConfig(conf, true), null); + } + + /** + * Close the Admin and release all associated resources. + *

+ * See {@link Admin#close(long, TimeUnit)} + */ + @Override + default void close() { + close(Long.MAX_VALUE, TimeUnit.MILLISECONDS); + } + + /** + * Close the Admin and release all associated resources. + *

+ * The close operation has a grace period during which current operations will be allowed to + * complete, specified by the given duration and time unit. + * New operations will not be accepted during the grace period. Once the grace period is over, + * all operations that have not yet been completed will be aborted with a TimeoutException. + * + * @param duration The duration to use for the wait time. + * @param unit The time unit to use for the wait time. + * @deprecated Since 2.2. Use {@link #close(Duration)} or {@link #close()}. + */ + @Deprecated + default void close(long duration, TimeUnit unit) { + close(Duration.ofMillis(unit.toMillis(duration))); + } + + /** + * Close the Admin client and release all associated resources. + *

+ * The close operation has a grace period during which current operations will be allowed to + * complete, specified by the given duration. + * New operations will not be accepted during the grace period. Once the grace period is over, + * all operations that have not yet been completed will be aborted with a TimeoutException. + * + * @param timeout The time to use for the wait time. + */ + void close(Duration timeout); + + /** + * Create a batch of new topics with the default options. + *

+ * This is a convenience method for #{@link #createTopics(Collection, CreateTopicsOptions)} with default options. + * See the overload for more details. + *

+ * This operation is supported by brokers with version 0.10.1.0 or higher. + * + * @param newTopics The new topics to create. + * @return The CreateTopicsResult. + */ + default CreateTopicsResult createTopics(Collection newTopics) { + return createTopics(newTopics, new CreateTopicsOptions()); + } + + /** + * Create a batch of new topics. + *

+ * This operation is not transactional so it may succeed for some topics while fail for others. + *

+ * It may take several seconds after {@code CreateTopicsResult} returns + * success for all the brokers to become aware that the topics have been created. + * During this time, {@link Admin#listTopics()} and {@link Admin#describeTopics(Collection)} + * may not return information about the new topics. + *

+ * This operation is supported by brokers with version 0.10.1.0 or higher. The validateOnly option is supported + * from version 0.10.2.0. + * + * @param newTopics The new topics to create. + * @param options The options to use when creating the new topics. + * @return The CreateTopicsResult. + */ + CreateTopicsResult createTopics(Collection newTopics, CreateTopicsOptions options); + + /** + * This is a convenience method for #{@link Admin#deleteTopics(Collection, DeleteTopicsOptions)} + * with default options. See the overload for more details. + *

+ * This operation is supported by brokers with version 0.10.1.0 or higher. + * + * @param topics The topic names to delete. + * @return The DeleteTopicsResult. + */ + default DeleteTopicsResult deleteTopics(Collection topics) { + return deleteTopics(topics, new DeleteTopicsOptions()); + } + + /** + * Delete a batch of topics. + *

+ * This operation is not transactional so it may succeed for some topics while fail for others. + *

+ * It may take several seconds after the {@code DeleteTopicsResult} returns + * success for all the brokers to become aware that the topics are gone. + * During this time, Admin#listTopics and Admin#describeTopics + * may continue to return information about the deleted topics. + *

+ * If delete.topic.enable is false on the brokers, deleteTopics will mark + * the topics for deletion, but not actually delete them. The futures will + * return successfully in this case. + *

+ * This operation is supported by brokers with version 0.10.1.0 or higher. + * + * @param topics The topic names to delete. + * @param options The options to use when deleting the topics. + * @return The DeleteTopicsResult. + */ + DeleteTopicsResult deleteTopics(Collection topics, DeleteTopicsOptions options); + + /** + * List the topics available in the cluster with the default options. + *

+ * This is a convenience method for #{@link Admin#listTopics(ListTopicsOptions)} with default options. + * See the overload for more details. + * + * @return The ListTopicsResult. + */ + default ListTopicsResult listTopics() { + return listTopics(new ListTopicsOptions()); + } + + /** + * List the topics available in the cluster. + * + * @param options The options to use when listing the topics. + * @return The ListTopicsResult. + */ + ListTopicsResult listTopics(ListTopicsOptions options); + + /** + * Describe some topics in the cluster, with the default options. + *

+ * This is a convenience method for #{@link Admin#describeTopics(Collection, DescribeTopicsOptions)} with + * default options. See the overload for more details. + * + * @param topicNames The names of the topics to describe. + * @return The DescribeTopicsResult. + */ + default DescribeTopicsResult describeTopics(Collection topicNames) { + return describeTopics(topicNames, new DescribeTopicsOptions()); + } + + /** + * Describe some topics in the cluster. + * + * @param topicNames The names of the topics to describe. + * @param options The options to use when describing the topic. + * @return The DescribeTopicsResult. + */ + DescribeTopicsResult describeTopics(Collection topicNames, DescribeTopicsOptions options); + + /** + * Get information about the nodes in the cluster, using the default options. + *

+ * This is a convenience method for #{@link Admin#describeCluster(DescribeClusterOptions)} with default options. + * See the overload for more details. + * + * @return The DescribeClusterResult. + */ + default DescribeClusterResult describeCluster() { + return describeCluster(new DescribeClusterOptions()); + } + + /** + * Get information about the nodes in the cluster. + * + * @param options The options to use when getting information about the cluster. + * @return The DescribeClusterResult. + */ + DescribeClusterResult describeCluster(DescribeClusterOptions options); + + /** + * This is a convenience method for #{@link Admin#describeAcls(AclBindingFilter, DescribeAclsOptions)} with + * default options. See the overload for more details. + *

+ * This operation is supported by brokers with version 0.11.0.0 or higher. + * + * @param filter The filter to use. + * @return The DeleteAclsResult. + */ + default DescribeAclsResult describeAcls(AclBindingFilter filter) { + return describeAcls(filter, new DescribeAclsOptions()); + } + + /** + * Lists access control lists (ACLs) according to the supplied filter. + *

+ * Note: it may take some time for changes made by createAcls or deleteAcls to be reflected + * in the output of describeAcls. + *

+ * This operation is supported by brokers with version 0.11.0.0 or higher. + * + * @param filter The filter to use. + * @param options The options to use when listing the ACLs. + * @return The DeleteAclsResult. + */ + DescribeAclsResult describeAcls(AclBindingFilter filter, DescribeAclsOptions options); + + /** + * This is a convenience method for #{@link Admin#createAcls(Collection, CreateAclsOptions)} with + * default options. See the overload for more details. + *

+ * This operation is supported by brokers with version 0.11.0.0 or higher. + * + * @param acls The ACLs to create + * @return The CreateAclsResult. + */ + default CreateAclsResult createAcls(Collection acls) { + return createAcls(acls, new CreateAclsOptions()); + } + + /** + * Creates access control lists (ACLs) which are bound to specific resources. + *

+ * This operation is not transactional so it may succeed for some ACLs while fail for others. + *

+ * If you attempt to add an ACL that duplicates an existing ACL, no error will be raised, but + * no changes will be made. + *

+ * This operation is supported by brokers with version 0.11.0.0 or higher. + * + * @param acls The ACLs to create + * @param options The options to use when creating the ACLs. + * @return The CreateAclsResult. + */ + CreateAclsResult createAcls(Collection acls, CreateAclsOptions options); + + /** + * This is a convenience method for #{@link Admin#deleteAcls(Collection, DeleteAclsOptions)} with default options. + * See the overload for more details. + *

+ * This operation is supported by brokers with version 0.11.0.0 or higher. + * + * @param filters The filters to use. + * @return The DeleteAclsResult. + */ + default DeleteAclsResult deleteAcls(Collection filters) { + return deleteAcls(filters, new DeleteAclsOptions()); + } + + /** + * Deletes access control lists (ACLs) according to the supplied filters. + *

+ * This operation is not transactional so it may succeed for some ACLs while fail for others. + *

+ * This operation is supported by brokers with version 0.11.0.0 or higher. + * + * @param filters The filters to use. + * @param options The options to use when deleting the ACLs. + * @return The DeleteAclsResult. + */ + DeleteAclsResult deleteAcls(Collection filters, DeleteAclsOptions options); + + + /** + * Get the configuration for the specified resources with the default options. + *

+ * This is a convenience method for #{@link Admin#describeConfigs(Collection, DescribeConfigsOptions)} with default options. + * See the overload for more details. + *

+ * This operation is supported by brokers with version 0.11.0.0 or higher. + * + * @param resources The resources (topic and broker resource types are currently supported) + * @return The DescribeConfigsResult + */ + default DescribeConfigsResult describeConfigs(Collection resources) { + return describeConfigs(resources, new DescribeConfigsOptions()); + } + + /** + * Get the configuration for the specified resources. + *

+ * The returned configuration includes default values and the isDefault() method can be used to distinguish them + * from user supplied values. + *

+ * The value of config entries where isSensitive() is true is always {@code null} so that sensitive information + * is not disclosed. + *

+ * Config entries where isReadOnly() is true cannot be updated. + *

+ * This operation is supported by brokers with version 0.11.0.0 or higher. + * + * @param resources The resources (topic and broker resource types are currently supported) + * @param options The options to use when describing configs + * @return The DescribeConfigsResult + */ + DescribeConfigsResult describeConfigs(Collection resources, DescribeConfigsOptions options); + + /** + * Update the configuration for the specified resources with the default options. + *

+ * This is a convenience method for #{@link Admin#alterConfigs(Map, AlterConfigsOptions)} with default options. + * See the overload for more details. + *

+ * This operation is supported by brokers with version 0.11.0.0 or higher. + * + * @param configs The resources with their configs (topic is the only resource type with configs that can + * be updated currently) + * @return The AlterConfigsResult + * @deprecated Since 2.3. Use {@link #incrementalAlterConfigs(Map)}. + */ + @Deprecated + default AlterConfigsResult alterConfigs(Map configs) { + return alterConfigs(configs, new AlterConfigsOptions()); + } + + /** + * Update the configuration for the specified resources with the default options. + *

+ * Updates are not transactional so they may succeed for some resources while fail for others. The configs for + * a particular resource are updated atomically. + *

+ * This operation is supported by brokers with version 0.11.0.0 or higher. + * + * @param configs The resources with their configs (topic is the only resource type with configs that can + * be updated currently) + * @param options The options to use when describing configs + * @return The AlterConfigsResult + * @deprecated Since 2.3. Use {@link #incrementalAlterConfigs(Map, AlterConfigsOptions)}. + */ + @Deprecated + AlterConfigsResult alterConfigs(Map configs, AlterConfigsOptions options); + + /** + * Incrementally updates the configuration for the specified resources with default options. + *

+ * This is a convenience method for #{@link Admin#incrementalAlterConfigs(Map, AlterConfigsOptions)} with default options. + * See the overload for more details.* + *

+ * This operation is supported by brokers with version 2.3.0 or higher. + * + * @param configs The resources with their configs + * @return The IncrementalAlterConfigsResult + */ + default AlterConfigsResult incrementalAlterConfigs(Map> configs) { + return incrementalAlterConfigs(configs, new AlterConfigsOptions()); + } + + + /** + * Incrementally update the configuration for the specified resources. + *

+ * Updates are not transactional so they may succeed for some resources while fail for others. The configs for + * a particular resource are updated atomically. + * + *

The following exceptions can be anticipated when calling {@code get()} on the futures obtained from + * the returned {@code IncrementalAlterConfigsResult}:

+ *
    + *
  • {@link org.apache.kafka.common.errors.ClusterAuthorizationException} + * if the authenticated user didn't have alter access to the cluster.
    • + *
  • {@link org.apache.kafka.common.errors.TopicAuthorizationException} + * if the authenticated user didn't have alter access to the Topic.
    • + *
  • {@link org.apache.kafka.common.errors.InvalidRequestException} + * if the request details are invalid. e.g., a configuration key was specified more than once for a resource
    • + *
* + *

+ * This operation is supported by brokers with version 2.3.0 or higher. + * + * @param configs The resources with their configs + * @param options The options to use when altering configs + * @return The IncrementalAlterConfigsResult + */ + AlterConfigsResult incrementalAlterConfigs(Map> configs, AlterConfigsOptions options); + + /** + * Change the log directory for the specified replicas. If the replica does not exist on the broker, the result + * shows REPLICA_NOT_AVAILABLE for the given replica and the replica will be created in the given log directory on the + * broker when it is created later. If the replica already exists on the broker, the replica will be moved to the given + * log directory if it is not already there. + *

+ * This operation is not transactional so it may succeed for some replicas while fail for others. + *

+ * This is a convenience method for #{@link Admin#alterReplicaLogDirs(Map, AlterReplicaLogDirsOptions)} with default options. + * See the overload for more details. + *

+ * This operation is supported by brokers with version 1.1.0 or higher. + * + * @param replicaAssignment The replicas with their log directory absolute path + * @return The AlterReplicaLogDirsResult + */ + default AlterReplicaLogDirsResult alterReplicaLogDirs(Map replicaAssignment) { + return alterReplicaLogDirs(replicaAssignment, new AlterReplicaLogDirsOptions()); + } + + /** + * Change the log directory for the specified replicas. If the replica does not exist on the broker, the result + * shows REPLICA_NOT_AVAILABLE for the given replica and the replica will be created in the given log directory on the + * broker when it is created later. If the replica already exists on the broker, the replica will be moved to the given + * log directory if it is not already there. + *

+ * This operation is not transactional so it may succeed for some replicas while fail for others. + *

+ * This operation is supported by brokers with version 1.1.0 or higher. + * + * @param replicaAssignment The replicas with their log directory absolute path + * @param options The options to use when changing replica dir + * @return The AlterReplicaLogDirsResult + */ + AlterReplicaLogDirsResult alterReplicaLogDirs(Map replicaAssignment, + AlterReplicaLogDirsOptions options); + + /** + * Query the information of all log directories on the given set of brokers + *

+ * This is a convenience method for #{@link Admin#describeLogDirs(Collection, DescribeLogDirsOptions)} with default options. + * See the overload for more details. + *

+ * This operation is supported by brokers with version 1.0.0 or higher. + * + * @param brokers A list of brokers + * @return The DescribeLogDirsResult + */ + default DescribeLogDirsResult describeLogDirs(Collection brokers) { + return describeLogDirs(brokers, new DescribeLogDirsOptions()); + } + + /** + * Query the information of all log directories on the given set of brokers + *

+ * This operation is supported by brokers with version 1.0.0 or higher. + * + * @param brokers A list of brokers + * @param options The options to use when querying log dir info + * @return The DescribeLogDirsResult + */ + DescribeLogDirsResult describeLogDirs(Collection brokers, DescribeLogDirsOptions options); + + /** + * Query the replica log directory information for the specified replicas. + *

+ * This is a convenience method for #{@link Admin#describeReplicaLogDirs(Collection, DescribeReplicaLogDirsOptions)} + * with default options. See the overload for more details. + *

+ * This operation is supported by brokers with version 1.0.0 or higher. + * + * @param replicas The replicas to query + * @return The DescribeReplicaLogDirsResult + */ + default DescribeReplicaLogDirsResult describeReplicaLogDirs(Collection replicas) { + return describeReplicaLogDirs(replicas, new DescribeReplicaLogDirsOptions()); + } + + /** + * Query the replica log directory information for the specified replicas. + *

+ * This operation is supported by brokers with version 1.0.0 or higher. + * + * @param replicas The replicas to query + * @param options The options to use when querying replica log dir info + * @return The DescribeReplicaLogDirsResult + */ + DescribeReplicaLogDirsResult describeReplicaLogDirs(Collection replicas, DescribeReplicaLogDirsOptions options); + + /** + *

Increase the number of partitions of the topics given as the keys of {@code newPartitions} + * according to the corresponding values. If partitions are increased for a topic that has a key, + * the partition logic or ordering of the messages will be affected.

+ * + *

This is a convenience method for {@link #createPartitions(Map, CreatePartitionsOptions)} with default options. + * See the overload for more details.

+ * + * @param newPartitions The topics which should have new partitions created, and corresponding parameters + * for the created partitions. + * @return The CreatePartitionsResult. + */ + default CreatePartitionsResult createPartitions(Map newPartitions) { + return createPartitions(newPartitions, new CreatePartitionsOptions()); + } + + /** + *

Increase the number of partitions of the topics given as the keys of {@code newPartitions} + * according to the corresponding values. If partitions are increased for a topic that has a key, + * the partition logic or ordering of the messages will be affected.

+ * + *

This operation is not transactional so it may succeed for some topics while fail for others.

+ * + *

It may take several seconds after this method returns + * success for all the brokers to become aware that the partitions have been created. + * During this time, {@link Admin#describeTopics(Collection)} + * may not return information about the new partitions.

+ * + *

This operation is supported by brokers with version 1.0.0 or higher.

+ * + *

The following exceptions can be anticipated when calling {@code get()} on the futures obtained from the + * {@link CreatePartitionsResult#values() values()} method of the returned {@code CreatePartitionsResult}

+ *
    + *
  • {@link org.apache.kafka.common.errors.AuthorizationException} + * if the authenticated user is not authorized to alter the topic
    • + *
  • {@link org.apache.kafka.common.errors.TimeoutException} + * if the request was not completed in within the given {@link CreatePartitionsOptions#timeoutMs()}.
    • + *
  • {@link org.apache.kafka.common.errors.ReassignmentInProgressException} + * if a partition reassignment is currently in progress
    • + *
  • {@link org.apache.kafka.common.errors.BrokerNotAvailableException} + * if the requested {@link NewPartitions#assignments()} contain a broker that is currently unavailable.
    • + *
  • {@link org.apache.kafka.common.errors.InvalidReplicationFactorException} + * if no {@link NewPartitions#assignments()} are given and it is impossible for the broker to assign + * replicas with the topics replication factor.
    • + *
  • Subclasses of {@link org.apache.kafka.common.KafkaException} + * if the request is invalid in some way.
    • + *
+ * + * @param newPartitions The topics which should have new partitions created, and corresponding parameters + * for the created partitions. + * @param options The options to use when creating the new paritions. + * @return The CreatePartitionsResult. + */ + CreatePartitionsResult createPartitions(Map newPartitions, + CreatePartitionsOptions options); + + /** + * Delete records whose offset is smaller than the given offset of the corresponding partition. + *

+ * This is a convenience method for {@link #deleteRecords(Map, DeleteRecordsOptions)} with default options. + * See the overload for more details. + *

+ * This operation is supported by brokers with version 0.11.0.0 or higher. + * + * @param recordsToDelete The topic partitions and related offsets from which records deletion starts. + * @return The DeleteRecordsResult. + */ + default DeleteRecordsResult deleteRecords(Map recordsToDelete) { + return deleteRecords(recordsToDelete, new DeleteRecordsOptions()); + } + + /** + * Delete records whose offset is smaller than the given offset of the corresponding partition. + *

+ * This operation is supported by brokers with version 0.11.0.0 or higher. + * + * @param recordsToDelete The topic partitions and related offsets from which records deletion starts. + * @param options The options to use when deleting records. + * @return The DeleteRecordsResult. + */ + DeleteRecordsResult deleteRecords(Map recordsToDelete, + DeleteRecordsOptions options); + + /** + *

Create a Delegation Token.

+ * + *

This is a convenience method for {@link #createDelegationToken(CreateDelegationTokenOptions)} with default options. + * See the overload for more details.

+ * + * @return The CreateDelegationTokenResult. + */ + default CreateDelegationTokenResult createDelegationToken() { + return createDelegationToken(new CreateDelegationTokenOptions()); + } + + + /** + *

Create a Delegation Token.

+ * + *

This operation is supported by brokers with version 1.1.0 or higher.

+ * + *

The following exceptions can be anticipated when calling {@code get()} on the futures obtained from the + * {@link CreateDelegationTokenResult#delegationToken() delegationToken()} method of the returned {@code CreateDelegationTokenResult}

+ *
    + *
  • {@link org.apache.kafka.common.errors.UnsupportedByAuthenticationException} + * If the request sent on PLAINTEXT/1-way SSL channels or delegation token authenticated channels.
    • + *
  • {@link org.apache.kafka.common.errors.InvalidPrincipalTypeException} + * if the renewers principal type is not supported.
    • + *
  • {@link org.apache.kafka.common.errors.DelegationTokenDisabledException} + * if the delegation token feature is disabled.
    • + *
  • {@link org.apache.kafka.common.errors.TimeoutException} + * if the request was not completed in within the given {@link CreateDelegationTokenOptions#timeoutMs()}.
    • + *
+ * + * @param options The options to use when creating delegation token. + * @return The DeleteRecordsResult. + */ + CreateDelegationTokenResult createDelegationToken(CreateDelegationTokenOptions options); + + + /** + *

Renew a Delegation Token.

+ * + *

This is a convenience method for {@link #renewDelegationToken(byte[], RenewDelegationTokenOptions)} with default options. + * See the overload for more details.

+ * + * @param hmac HMAC of the Delegation token + * @return The RenewDelegationTokenResult. + */ + default RenewDelegationTokenResult renewDelegationToken(byte[] hmac) { + return renewDelegationToken(hmac, new RenewDelegationTokenOptions()); + } + + /** + *

Renew a Delegation Token.

+ * + *

This operation is supported by brokers with version 1.1.0 or higher.

+ * + *

The following exceptions can be anticipated when calling {@code get()} on the futures obtained from the + * {@link RenewDelegationTokenResult#expiryTimestamp() expiryTimestamp()} method of the returned {@code RenewDelegationTokenResult}

+ *
    + *
  • {@link org.apache.kafka.common.errors.UnsupportedByAuthenticationException} + * If the request sent on PLAINTEXT/1-way SSL channels or delegation token authenticated channels.
    • + *
  • {@link org.apache.kafka.common.errors.DelegationTokenDisabledException} + * if the delegation token feature is disabled.
    • + *
  • {@link org.apache.kafka.common.errors.DelegationTokenNotFoundException} + * if the delegation token is not found on server.
    • + *
  • {@link org.apache.kafka.common.errors.DelegationTokenOwnerMismatchException} + * if the authenticated user is not owner/renewer of the token.
    • + *
  • {@link org.apache.kafka.common.errors.DelegationTokenExpiredException} + * if the delegation token is expired.
    • + *
  • {@link org.apache.kafka.common.errors.TimeoutException} + * if the request was not completed in within the given {@link RenewDelegationTokenOptions#timeoutMs()}.
    • + *
+ * + * @param hmac HMAC of the Delegation token + * @param options The options to use when renewing delegation token. + * @return The RenewDelegationTokenResult. + */ + RenewDelegationTokenResult renewDelegationToken(byte[] hmac, RenewDelegationTokenOptions options); + + /** + *

Expire a Delegation Token.

+ * + *

This is a convenience method for {@link #expireDelegationToken(byte[], ExpireDelegationTokenOptions)} with default options. + * This will expire the token immediately. See the overload for more details.

+ * + * @param hmac HMAC of the Delegation token + * @return The ExpireDelegationTokenResult. + */ + default ExpireDelegationTokenResult expireDelegationToken(byte[] hmac) { + return expireDelegationToken(hmac, new ExpireDelegationTokenOptions()); + } + + /** + *

Expire a Delegation Token.

+ * + *

This operation is supported by brokers with version 1.1.0 or higher.

+ * + *

The following exceptions can be anticipated when calling {@code get()} on the futures obtained from the + * {@link ExpireDelegationTokenResult#expiryTimestamp() expiryTimestamp()} method of the returned {@code ExpireDelegationTokenResult}

+ *
    + *
  • {@link org.apache.kafka.common.errors.UnsupportedByAuthenticationException} + * If the request sent on PLAINTEXT/1-way SSL channels or delegation token authenticated channels.
    • + *
  • {@link org.apache.kafka.common.errors.DelegationTokenDisabledException} + * if the delegation token feature is disabled.
    • + *
  • {@link org.apache.kafka.common.errors.DelegationTokenNotFoundException} + * if the delegation token is not found on server.
    • + *
  • {@link org.apache.kafka.common.errors.DelegationTokenOwnerMismatchException} + * if the authenticated user is not owner/renewer of the requested token.
    • + *
  • {@link org.apache.kafka.common.errors.DelegationTokenExpiredException} + * if the delegation token is expired.
    • + *
  • {@link org.apache.kafka.common.errors.TimeoutException} + * if the request was not completed in within the given {@link ExpireDelegationTokenOptions#timeoutMs()}.
    • + *
+ * + * @param hmac HMAC of the Delegation token + * @param options The options to use when expiring delegation token. + * @return The ExpireDelegationTokenResult. + */ + ExpireDelegationTokenResult expireDelegationToken(byte[] hmac, ExpireDelegationTokenOptions options); + + /** + *

Describe the Delegation Tokens.

+ * + *

This is a convenience method for {@link #describeDelegationToken(DescribeDelegationTokenOptions)} with default options. + * This will return all the user owned tokens and tokens where user have Describe permission. See the overload for more details.

+ * + * @return The DescribeDelegationTokenResult. + */ + default DescribeDelegationTokenResult describeDelegationToken() { + return describeDelegationToken(new DescribeDelegationTokenOptions()); + } + + /** + *

Describe the Delegation Tokens.

+ * + *

This operation is supported by brokers with version 1.1.0 or higher.

+ * + *

The following exceptions can be anticipated when calling {@code get()} on the futures obtained from the + * {@link DescribeDelegationTokenResult#delegationTokens() delegationTokens()} method of the returned {@code DescribeDelegationTokenResult}

+ *
    + *
  • {@link org.apache.kafka.common.errors.UnsupportedByAuthenticationException} + * If the request sent on PLAINTEXT/1-way SSL channels or delegation token authenticated channels.
    • + *
  • {@link org.apache.kafka.common.errors.DelegationTokenDisabledException} + * if the delegation token feature is disabled.
    • + *
  • {@link org.apache.kafka.common.errors.TimeoutException} + * if the request was not completed in within the given {@link DescribeDelegationTokenOptions#timeoutMs()}.
    • + *
+ * + * @param options The options to use when describing delegation tokens. + * @return The DescribeDelegationTokenResult. + */ + DescribeDelegationTokenResult describeDelegationToken(DescribeDelegationTokenOptions options); + + /** + * Describe some group IDs in the cluster. + * + * @param groupIds The IDs of the groups to describe. + * @param options The options to use when describing the groups. + * @return The DescribeConsumerGroupResult. + */ + DescribeConsumerGroupsResult describeConsumerGroups(Collection groupIds, + DescribeConsumerGroupsOptions options); + + /** + * Describe some group IDs in the cluster, with the default options. + *

+ * This is a convenience method for + * #{@link Admin#describeConsumerGroups(Collection, DescribeConsumerGroupsOptions)} with + * default options. See the overload for more details. + * + * @param groupIds The IDs of the groups to describe. + * @return The DescribeConsumerGroupResult. + */ + default DescribeConsumerGroupsResult describeConsumerGroups(Collection groupIds) { + return describeConsumerGroups(groupIds, new DescribeConsumerGroupsOptions()); + } + + /** + * List the consumer groups available in the cluster. + * + * @param options The options to use when listing the consumer groups. + * @return The ListGroupsResult. + */ + ListConsumerGroupsResult listConsumerGroups(ListConsumerGroupsOptions options); + + /** + * List the consumer groups available in the cluster with the default options. + *

+ * This is a convenience method for #{@link Admin#listConsumerGroups(ListConsumerGroupsOptions)} with default options. + * See the overload for more details. + * + * @return The ListGroupsResult. + */ + default ListConsumerGroupsResult listConsumerGroups() { + return listConsumerGroups(new ListConsumerGroupsOptions()); + } + + /** + * List the consumer group offsets available in the cluster. + * + * @param options The options to use when listing the consumer group offsets. + * @return The ListGroupOffsetsResult + */ + ListConsumerGroupOffsetsResult listConsumerGroupOffsets(String groupId, ListConsumerGroupOffsetsOptions options); + + /** + * List the consumer group offsets available in the cluster with the default options. + *

+ * This is a convenience method for #{@link Admin#listConsumerGroupOffsets(String, ListConsumerGroupOffsetsOptions)} with default options. + * + * @return The ListGroupOffsetsResult. + */ + default ListConsumerGroupOffsetsResult listConsumerGroupOffsets(String groupId) { + return listConsumerGroupOffsets(groupId, new ListConsumerGroupOffsetsOptions()); + } + + /** + * Delete consumer groups from the cluster. + * + * @param options The options to use when deleting a consumer group. + * @return The DeletConsumerGroupResult. + */ + DeleteConsumerGroupsResult deleteConsumerGroups(Collection groupIds, DeleteConsumerGroupsOptions options); + + /** + * Delete consumer groups from the cluster with the default options. + * + * @return The DeleteConsumerGroupResult. + */ + default DeleteConsumerGroupsResult deleteConsumerGroups(Collection groupIds) { + return deleteConsumerGroups(groupIds, new DeleteConsumerGroupsOptions()); + } + + /** + * Delete committed offsets for a set of partitions in a consumer group. This will + * succeed at the partition level only if the group is not actively subscribed + * to the corresponding topic. + * + * @param options The options to use when deleting offsets in a consumer group. + * @return The DeleteConsumerGroupOffsetsResult. + */ + DeleteConsumerGroupOffsetsResult deleteConsumerGroupOffsets(String groupId, + Set partitions, + DeleteConsumerGroupOffsetsOptions options); + + /** + * Delete committed offsets for a set of partitions in a consumer group with the default + * options. This will succeed at the partition level only if the group is not actively + * subscribed to the corresponding topic. + * + * @return The DeleteConsumerGroupOffsetsResult. + */ + default DeleteConsumerGroupOffsetsResult deleteConsumerGroupOffsets(String groupId, Set partitions) { + return deleteConsumerGroupOffsets(groupId, partitions, new DeleteConsumerGroupOffsetsOptions()); + } + + /** + * Elect the preferred replica as leader for topic partitions. + *

+ * This is a convenience method for {@link #electLeaders(ElectionType, Set, ElectLeadersOptions)} + * with preferred election type and default options. + *

+ * This operation is supported by brokers with version 2.2.0 or higher. + * + * @param partitions The partitions for which the preferred leader should be elected. + * @return The ElectPreferredLeadersResult. + * @deprecated Since 2.4.0. Use {@link #electLeaders(ElectionType, Set)}. + */ + @Deprecated + default ElectPreferredLeadersResult electPreferredLeaders(Collection partitions) { + return electPreferredLeaders(partitions, new ElectPreferredLeadersOptions()); + } + + /** + * Elect the preferred replica as leader for topic partitions. + *

+ * This is a convenience method for {@link #electLeaders(ElectionType, Set, ElectLeadersOptions)} + * with preferred election type. + *

+ * This operation is supported by brokers with version 2.2.0 or higher. + * + * @param partitions The partitions for which the preferred leader should be elected. + * @param options The options to use when electing the preferred leaders. + * @return The ElectPreferredLeadersResult. + * @deprecated Since 2.4.0. Use {@link #electLeaders(ElectionType, Set, ElectLeadersOptions)}. + */ + @Deprecated + default ElectPreferredLeadersResult electPreferredLeaders(Collection partitions, + ElectPreferredLeadersOptions options) { + final ElectLeadersOptions newOptions = new ElectLeadersOptions(); + newOptions.timeoutMs(options.timeoutMs()); + final Set topicPartitions = partitions == null ? null : new HashSet<>(partitions); + + return new ElectPreferredLeadersResult(electLeaders(ElectionType.PREFERRED, topicPartitions, newOptions)); + } + + /** + * Elect a replica as leader for topic partitions. + *

+ * This is a convenience method for {@link #electLeaders(ElectionType, Set, ElectLeadersOptions)} + * with default options. + * + * @param electionType The type of election to conduct. + * @param partitions The topics and partitions for which to conduct elections. + * @return The ElectLeadersResult. + */ + default ElectLeadersResult electLeaders(ElectionType electionType, Set partitions) { + return electLeaders(electionType, partitions, new ElectLeadersOptions()); + } + + /** + * Elect a replica as leader for the given {@code partitions}, or for all partitions if the argumentl + * to {@code partitions} is null. + *

+ * This operation is not transactional so it may succeed for some partitions while fail for others. + *

+ * It may take several seconds after this method returns success for all the brokers in the cluster + * to become aware that the partitions have new leaders. During this time, + * {@link Admin#describeTopics(Collection)} may not return information about the partitions' + * new leaders. + *

+ * This operation is supported by brokers with version 2.2.0 or later if preferred eleciton is use; + * otherwise the brokers most be 2.4.0 or higher. + * + *

The following exceptions can be anticipated when calling {@code get()} on the future obtained + * from the returned {@code ElectLeadersResult}:

+ *
    + *
  • {@link org.apache.kafka.common.errors.ClusterAuthorizationException} + * if the authenticated user didn't have alter access to the cluster.
    • + *
  • {@link org.apache.kafka.common.errors.UnknownTopicOrPartitionException} + * if the topic or partition did not exist within the cluster.
    • + *
  • {@link org.apache.kafka.common.errors.InvalidTopicException} + * if the topic was already queued for deletion.
    • + *
  • {@link org.apache.kafka.common.errors.NotControllerException} + * if the request was sent to a broker that was not the controller for the cluster.
    • + *
  • {@link org.apache.kafka.common.errors.TimeoutException} + * if the request timed out before the election was complete.
    • + *
  • {@link org.apache.kafka.common.errors.LeaderNotAvailableException} + * if the preferred leader was not alive or not in the ISR.
    • + *
+ * + * @param electionType The type of election to conduct. + * @param partitions The topics and partitions for which to conduct elections. + * @param options The options to use when electing the leaders. + * @return The ElectLeadersResult. + */ + ElectLeadersResult electLeaders( + ElectionType electionType, + Set partitions, + ElectLeadersOptions options); + + + /** + * Change the reassignments for one or more partitions. + * Providing an empty Optional (e.g via {@link Optional#empty()}) will revert the reassignment for the associated partition. + * + * This is a convenience method for {@link #alterPartitionReassignments(Map, AlterPartitionReassignmentsOptions)} + * with default options. See the overload for more details. + */ + default AlterPartitionReassignmentsResult alterPartitionReassignments( + Map> reassignments) { + return alterPartitionReassignments(reassignments, new AlterPartitionReassignmentsOptions()); + } + + /** + * Change the reassignments for one or more partitions. + * Providing an empty Optional (e.g via {@link Optional#empty()}) will revert the reassignment for the associated partition. + * + *

The following exceptions can be anticipated when calling {@code get()} on the futures obtained from + * the returned {@code AlterPartitionReassignmentsResult}:

+ *
    + *
  • {@link org.apache.kafka.common.errors.ClusterAuthorizationException} + * If the authenticated user didn't have alter access to the cluster.
    • + *
  • {@link org.apache.kafka.common.errors.UnknownTopicOrPartitionException} + * If the topic or partition does not exist within the cluster.
    • + *
  • {@link org.apache.kafka.common.errors.TimeoutException} + * if the request timed out before the controller could record the new assignments.
    • + *
  • {@link org.apache.kafka.common.errors.InvalidReplicaAssignmentException} + * If the specified assignment was not valid.
    • + *
  • {@link org.apache.kafka.common.errors.NoReassignmentInProgressException} + * If there was an attempt to cancel a reassignment for a partition which was not being reassigned.
    • + *
+ * + * @param reassignments The reassignments to add, modify, or remove. + * @param options The options to use. + * @return The result. + */ + AlterPartitionReassignmentsResult alterPartitionReassignments( + Map> reassignments, + AlterPartitionReassignmentsOptions options); + + + /** + * List all of the current partition reassignments + * + * This is a convenience method for {@link #listPartitionReassignments(ListPartitionReassignmentsOptions)} + * with default options. See the overload for more details. + */ + default ListPartitionReassignmentsResult listPartitionReassignments() { + return listPartitionReassignments(new ListPartitionReassignmentsOptions()); + } + + /** + * List the current reassignments for the given partitions + * + * This is a convenience method for {@link #listPartitionReassignments(Set, ListPartitionReassignmentsOptions)} + * with default options. See the overload for more details. + */ + default ListPartitionReassignmentsResult listPartitionReassignments(Set partitions) { + return listPartitionReassignments(partitions, new ListPartitionReassignmentsOptions()); + } + + /** + * List the current reassignments for the given partitions + * + *

The following exceptions can be anticipated when calling {@code get()} on the futures obtained from + * the returned {@code ListPartitionReassignmentsResult}:

+ *
    + *
  • {@link org.apache.kafka.common.errors.ClusterAuthorizationException} + * If the authenticated user doesn't have alter access to the cluster.
    • + *
  • {@link org.apache.kafka.common.errors.UnknownTopicOrPartitionException} + * If a given topic or partition does not exist.
    • + *
  • {@link org.apache.kafka.common.errors.TimeoutException} + * If the request timed out before the controller could list the current reassignments.
    • + *
+ * + * @param partitions The topic partitions to list reassignments for. + * @param options The options to use. + * @return The result. + */ + default ListPartitionReassignmentsResult listPartitionReassignments( + Set partitions, + ListPartitionReassignmentsOptions options) { + return listPartitionReassignments(Optional.of(partitions), options); + } + + /** + * List all of the current partition reassignments + * + *

The following exceptions can be anticipated when calling {@code get()} on the futures obtained from + * the returned {@code ListPartitionReassignmentsResult}:

+ *
    + *
  • {@link org.apache.kafka.common.errors.ClusterAuthorizationException} + * If the authenticated user doesn't have alter access to the cluster.
    • + *
  • {@link org.apache.kafka.common.errors.UnknownTopicOrPartitionException} + * If a given topic or partition does not exist.
    • + *
  • {@link org.apache.kafka.common.errors.TimeoutException} + * If the request timed out before the controller could list the current reassignments.
    • + *
+ * + * @param options The options to use. + * @return The result. + */ + default ListPartitionReassignmentsResult listPartitionReassignments(ListPartitionReassignmentsOptions options) { + return listPartitionReassignments(Optional.empty(), options); + } + + /** + * @param partitions the partitions we want to get reassignment for, or an empty optional if we want to get the reassignments for all partitions in the cluster + * @param options The options to use. + * @return The result. + */ + ListPartitionReassignmentsResult listPartitionReassignments(Optional> partitions, + ListPartitionReassignmentsOptions options); + + /** + * Remove members from the consumer group by given member identities. + *

+ * For possible error codes, refer to {@link LeaveGroupResponse}. + * + * @param groupId The ID of the group to remove member from. + * @param options The options to carry removing members' information. + * @return The MembershipChangeResult. + */ + RemoveMembersFromConsumerGroupResult removeMembersFromConsumerGroup(String groupId, RemoveMembersFromConsumerGroupOptions options); + + /** + * Get the metrics kept by the adminClient + */ + Map metrics(); +} diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/AdminClient.java b/clients/src/main/java/org/apache/kafka/clients/admin/AdminClient.java index b823cdc0509ed..75f1c5f10d29a 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/AdminClient.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/AdminClient.java @@ -17,835 +17,36 @@ package org.apache.kafka.clients.admin; -import org.apache.kafka.common.Metric; -import org.apache.kafka.common.MetricName; -import org.apache.kafka.common.TopicPartition; -import org.apache.kafka.common.TopicPartitionReplica; -import org.apache.kafka.common.acl.AclBinding; -import org.apache.kafka.common.acl.AclBindingFilter; -import org.apache.kafka.common.annotation.InterfaceStability; -import org.apache.kafka.common.config.ConfigResource; - -import java.time.Duration; -import java.util.Collection; import java.util.Map; import java.util.Properties; -import java.util.concurrent.TimeUnit; /** - * The administrative client for Kafka, which supports managing and inspecting topics, brokers, configurations and ACLs. + * The base class for in-built admin clients. * - * The minimum broker version required is 0.10.0.0. Methods with stricter requirements will specify the minimum broker - * version required. + * Client code should use the newer {@link Admin} interface in preference to this class. * - * This client was introduced in 0.11.0.0 and the API is still evolving. We will try to evolve the API in a compatible - * manner, but we reserve the right to make breaking changes in minor releases, if necessary. We will update the - * {@code InterfaceStability} annotation and this notice once the API is considered stable. + * This class may be removed in a later release, but has not be marked as deprecated to avoid unnecessary noise. */ -@InterfaceStability.Evolving -public abstract class AdminClient implements AutoCloseable { +public abstract class AdminClient implements Admin { /** - * Create a new AdminClient with the given configuration. + * Create a new Admin with the given configuration. * * @param props The configuration. * @return The new KafkaAdminClient. */ public static AdminClient create(Properties props) { - return KafkaAdminClient.createInternal(new AdminClientConfig(props, true), null); + return (AdminClient) Admin.create(props); } /** - * Create a new AdminClient with the given configuration. + * Create a new Admin with the given configuration. * * @param conf The configuration. * @return The new KafkaAdminClient. */ public static AdminClient create(Map conf) { - return KafkaAdminClient.createInternal(new AdminClientConfig(conf, true), null); - } - - /** - * Close the AdminClient and release all associated resources. - * - * See {@link AdminClient#close(long, TimeUnit)} - */ - @Override - public void close() { - close(Long.MAX_VALUE, TimeUnit.MILLISECONDS); - } - - /** - * Close the AdminClient and release all associated resources. - * - * The close operation has a grace period during which current operations will be allowed to - * complete, specified by the given duration and time unit. - * New operations will not be accepted during the grace period. Once the grace period is over, - * all operations that have not yet been completed will be aborted with a TimeoutException. - * - * @param duration The duration to use for the wait time. - * @param unit The time unit to use for the wait time. - * @deprecated Since 2.2. Use {@link #close(Duration)} or {@link #close()}. - */ - @Deprecated - public void close(long duration, TimeUnit unit) { - close(Duration.ofMillis(unit.toMillis(duration))); - } - - /** - * Close the AdminClient and release all associated resources. - * - * The close operation has a grace period during which current operations will be allowed to - * complete, specified by the given duration. - * New operations will not be accepted during the grace period. Once the grace period is over, - * all operations that have not yet been completed will be aborted with a TimeoutException. - * - * @param timeout The time to use for the wait time. - */ - public abstract void close(Duration timeout); - - /** - * Create a batch of new topics with the default options. - * - * This is a convenience method for #{@link #createTopics(Collection, CreateTopicsOptions)} with default options. - * See the overload for more details. - * - * This operation is supported by brokers with version 0.10.1.0 or higher. - * - * @param newTopics The new topics to create. - * @return The CreateTopicsResult. - */ - public CreateTopicsResult createTopics(Collection newTopics) { - return createTopics(newTopics, new CreateTopicsOptions()); - } - - /** - * Create a batch of new topics. - * - * This operation is not transactional so it may succeed for some topics while fail for others. - * - * It may take several seconds after {@code CreateTopicsResult} returns - * success for all the brokers to become aware that the topics have been created. - * During this time, {@link AdminClient#listTopics()} and {@link AdminClient#describeTopics(Collection)} - * may not return information about the new topics. - * - * This operation is supported by brokers with version 0.10.1.0 or higher. The validateOnly option is supported - * from version 0.10.2.0. - * - * @param newTopics The new topics to create. - * @param options The options to use when creating the new topics. - * @return The CreateTopicsResult. - */ - public abstract CreateTopicsResult createTopics(Collection newTopics, - CreateTopicsOptions options); - - /** - * This is a convenience method for #{@link AdminClient#deleteTopics(Collection, DeleteTopicsOptions)} - * with default options. See the overload for more details. - * - * This operation is supported by brokers with version 0.10.1.0 or higher. - * - * @param topics The topic names to delete. - * @return The DeleteTopicsResult. - */ - public DeleteTopicsResult deleteTopics(Collection topics) { - return deleteTopics(topics, new DeleteTopicsOptions()); - } - - /** - * Delete a batch of topics. - * - * This operation is not transactional so it may succeed for some topics while fail for others. - * - * It may take several seconds after the {@code DeleteTopicsResult} returns - * success for all the brokers to become aware that the topics are gone. - * During this time, AdminClient#listTopics and AdminClient#describeTopics - * may continue to return information about the deleted topics. - * - * If delete.topic.enable is false on the brokers, deleteTopics will mark - * the topics for deletion, but not actually delete them. The futures will - * return successfully in this case. - * - * This operation is supported by brokers with version 0.10.1.0 or higher. - * - * @param topics The topic names to delete. - * @param options The options to use when deleting the topics. - * @return The DeleteTopicsResult. - */ - public abstract DeleteTopicsResult deleteTopics(Collection topics, DeleteTopicsOptions options); - - /** - * List the topics available in the cluster with the default options. - * - * This is a convenience method for #{@link AdminClient#listTopics(ListTopicsOptions)} with default options. - * See the overload for more details. - * - * @return The ListTopicsResult. - */ - public ListTopicsResult listTopics() { - return listTopics(new ListTopicsOptions()); - } - - /** - * List the topics available in the cluster. - * - * @param options The options to use when listing the topics. - * @return The ListTopicsResult. - */ - public abstract ListTopicsResult listTopics(ListTopicsOptions options); - - /** - * Describe some topics in the cluster, with the default options. - * - * This is a convenience method for #{@link AdminClient#describeTopics(Collection, DescribeTopicsOptions)} with - * default options. See the overload for more details. - * - * @param topicNames The names of the topics to describe. - * - * @return The DescribeTopicsResult. - */ - public DescribeTopicsResult describeTopics(Collection topicNames) { - return describeTopics(topicNames, new DescribeTopicsOptions()); - } - - /** - * Describe some topics in the cluster. - * - * @param topicNames The names of the topics to describe. - * @param options The options to use when describing the topic. - * - * @return The DescribeTopicsResult. - */ - public abstract DescribeTopicsResult describeTopics(Collection topicNames, - DescribeTopicsOptions options); - - /** - * Get information about the nodes in the cluster, using the default options. - * - * This is a convenience method for #{@link AdminClient#describeCluster(DescribeClusterOptions)} with default options. - * See the overload for more details. - * - * @return The DescribeClusterResult. - */ - public DescribeClusterResult describeCluster() { - return describeCluster(new DescribeClusterOptions()); - } - - /** - * Get information about the nodes in the cluster. - * - * @param options The options to use when getting information about the cluster. - * @return The DescribeClusterResult. - */ - public abstract DescribeClusterResult describeCluster(DescribeClusterOptions options); - - /** - * This is a convenience method for #{@link AdminClient#describeAcls(AclBindingFilter, DescribeAclsOptions)} with - * default options. See the overload for more details. - * - * This operation is supported by brokers with version 0.11.0.0 or higher. - * - * @param filter The filter to use. - * @return The DeleteAclsResult. - */ - public DescribeAclsResult describeAcls(AclBindingFilter filter) { - return describeAcls(filter, new DescribeAclsOptions()); - } - - /** - * Lists access control lists (ACLs) according to the supplied filter. - * - * Note: it may take some time for changes made by createAcls or deleteAcls to be reflected - * in the output of describeAcls. - * - * This operation is supported by brokers with version 0.11.0.0 or higher. - * - * @param filter The filter to use. - * @param options The options to use when listing the ACLs. - * @return The DeleteAclsResult. - */ - public abstract DescribeAclsResult describeAcls(AclBindingFilter filter, DescribeAclsOptions options); - - /** - * This is a convenience method for #{@link AdminClient#createAcls(Collection, CreateAclsOptions)} with - * default options. See the overload for more details. - * - * This operation is supported by brokers with version 0.11.0.0 or higher. - * - * @param acls The ACLs to create - * @return The CreateAclsResult. - */ - public CreateAclsResult createAcls(Collection acls) { - return createAcls(acls, new CreateAclsOptions()); - } - - /** - * Creates access control lists (ACLs) which are bound to specific resources. - * - * This operation is not transactional so it may succeed for some ACLs while fail for others. - * - * If you attempt to add an ACL that duplicates an existing ACL, no error will be raised, but - * no changes will be made. - * - * This operation is supported by brokers with version 0.11.0.0 or higher. - * - * @param acls The ACLs to create - * @param options The options to use when creating the ACLs. - * @return The CreateAclsResult. - */ - public abstract CreateAclsResult createAcls(Collection acls, CreateAclsOptions options); - - /** - * This is a convenience method for #{@link AdminClient#deleteAcls(Collection, DeleteAclsOptions)} with default options. - * See the overload for more details. - * - * This operation is supported by brokers with version 0.11.0.0 or higher. - * - * @param filters The filters to use. - * @return The DeleteAclsResult. - */ - public DeleteAclsResult deleteAcls(Collection filters) { - return deleteAcls(filters, new DeleteAclsOptions()); - } - - /** - * Deletes access control lists (ACLs) according to the supplied filters. - * - * This operation is not transactional so it may succeed for some ACLs while fail for others. - * - * This operation is supported by brokers with version 0.11.0.0 or higher. - * - * @param filters The filters to use. - * @param options The options to use when deleting the ACLs. - * @return The DeleteAclsResult. - */ - public abstract DeleteAclsResult deleteAcls(Collection filters, DeleteAclsOptions options); - - - /** - * Get the configuration for the specified resources with the default options. - * - * This is a convenience method for #{@link AdminClient#describeConfigs(Collection, DescribeConfigsOptions)} with default options. - * See the overload for more details. - * - * This operation is supported by brokers with version 0.11.0.0 or higher. - * - * @param resources The resources (topic and broker resource types are currently supported) - * @return The DescribeConfigsResult - */ - public DescribeConfigsResult describeConfigs(Collection resources) { - return describeConfigs(resources, new DescribeConfigsOptions()); - } - - /** - * Get the configuration for the specified resources. - * - * The returned configuration includes default values and the isDefault() method can be used to distinguish them - * from user supplied values. - * - * The value of config entries where isSensitive() is true is always {@code null} so that sensitive information - * is not disclosed. - * - * Config entries where isReadOnly() is true cannot be updated. - * - * This operation is supported by brokers with version 0.11.0.0 or higher. - * - * @param resources The resources (topic and broker resource types are currently supported) - * @param options The options to use when describing configs - * @return The DescribeConfigsResult - */ - public abstract DescribeConfigsResult describeConfigs(Collection resources, - DescribeConfigsOptions options); - - /** - * Update the configuration for the specified resources with the default options. - * - * This is a convenience method for #{@link AdminClient#alterConfigs(Map, AlterConfigsOptions)} with default options. - * See the overload for more details. - * - * This operation is supported by brokers with version 0.11.0.0 or higher. - * - * @param configs The resources with their configs (topic is the only resource type with configs that can - * be updated currently) - * @return The AlterConfigsResult - */ - public AlterConfigsResult alterConfigs(Map configs) { - return alterConfigs(configs, new AlterConfigsOptions()); - } - - /** - * Update the configuration for the specified resources with the default options. - * - * Updates are not transactional so they may succeed for some resources while fail for others. The configs for - * a particular resource are updated atomically. - * - * This operation is supported by brokers with version 0.11.0.0 or higher. - * - * @param configs The resources with their configs (topic is the only resource type with configs that can - * be updated currently) - * @param options The options to use when describing configs - * @return The AlterConfigsResult - */ - public abstract AlterConfigsResult alterConfigs(Map configs, AlterConfigsOptions options); - - /** - * Change the log directory for the specified replicas. If the replica does not exist on the broker, the result - * shows REPLICA_NOT_AVAILABLE for the given replica and the replica will be created in the given log directory on the - * broker when it is created later. If the replica already exists on the broker, the replica will be moved to the given - * log directory if it is not already there. - * - * This operation is not transactional so it may succeed for some replicas while fail for others. - * - * This is a convenience method for #{@link AdminClient#alterReplicaLogDirs(Map, AlterReplicaLogDirsOptions)} with default options. - * See the overload for more details. - * - * This operation is supported by brokers with version 1.1.0 or higher. - * - * @param replicaAssignment The replicas with their log directory absolute path - * @return The AlterReplicaLogDirsResult - */ - public AlterReplicaLogDirsResult alterReplicaLogDirs(Map replicaAssignment) { - return alterReplicaLogDirs(replicaAssignment, new AlterReplicaLogDirsOptions()); - } - - /** - * Change the log directory for the specified replicas. If the replica does not exist on the broker, the result - * shows REPLICA_NOT_AVAILABLE for the given replica and the replica will be created in the given log directory on the - * broker when it is created later. If the replica already exists on the broker, the replica will be moved to the given - * log directory if it is not already there. - * - * This operation is not transactional so it may succeed for some replicas while fail for others. - * - * This operation is supported by brokers with version 1.1.0 or higher. - * - * @param replicaAssignment The replicas with their log directory absolute path - * @param options The options to use when changing replica dir - * @return The AlterReplicaLogDirsResult - */ - public abstract AlterReplicaLogDirsResult alterReplicaLogDirs(Map replicaAssignment, AlterReplicaLogDirsOptions options); - - /** - * Query the information of all log directories on the given set of brokers - * - * This is a convenience method for #{@link AdminClient#describeLogDirs(Collection, DescribeLogDirsOptions)} with default options. - * See the overload for more details. - * - * This operation is supported by brokers with version 1.0.0 or higher. - * - * @param brokers A list of brokers - * @return The DescribeLogDirsResult - */ - public DescribeLogDirsResult describeLogDirs(Collection brokers) { - return describeLogDirs(brokers, new DescribeLogDirsOptions()); - } - - /** - * Query the information of all log directories on the given set of brokers - * - * This operation is supported by brokers with version 1.0.0 or higher. - * - * @param brokers A list of brokers - * @param options The options to use when querying log dir info - * @return The DescribeLogDirsResult - */ - public abstract DescribeLogDirsResult describeLogDirs(Collection brokers, DescribeLogDirsOptions options); - - /** - * Query the replica log directory information for the specified replicas. - * - * This is a convenience method for #{@link AdminClient#describeReplicaLogDirs(Collection, DescribeReplicaLogDirsOptions)} - * with default options. See the overload for more details. - * - * This operation is supported by brokers with version 1.0.0 or higher. - * - * @param replicas The replicas to query - * @return The DescribeReplicaLogDirsResult - */ - public DescribeReplicaLogDirsResult describeReplicaLogDirs(Collection replicas) { - return describeReplicaLogDirs(replicas, new DescribeReplicaLogDirsOptions()); - } - - /** - * Query the replica log directory information for the specified replicas. - * - * This operation is supported by brokers with version 1.0.0 or higher. - * - * @param replicas The replicas to query - * @param options The options to use when querying replica log dir info - * @return The DescribeReplicaLogDirsResult - */ - public abstract DescribeReplicaLogDirsResult describeReplicaLogDirs(Collection replicas, DescribeReplicaLogDirsOptions options); - - /** - *

Increase the number of partitions of the topics given as the keys of {@code newPartitions} - * according to the corresponding values. If partitions are increased for a topic that has a key, - * the partition logic or ordering of the messages will be affected.

- * - *

This is a convenience method for {@link #createPartitions(Map, CreatePartitionsOptions)} with default options. - * See the overload for more details.

- * - * @param newPartitions The topics which should have new partitions created, and corresponding parameters - * for the created partitions. - * @return The CreatePartitionsResult. - */ - public CreatePartitionsResult createPartitions(Map newPartitions) { - return createPartitions(newPartitions, new CreatePartitionsOptions()); + return (AdminClient) Admin.create(conf); } - - /** - *

Increase the number of partitions of the topics given as the keys of {@code newPartitions} - * according to the corresponding values. If partitions are increased for a topic that has a key, - * the partition logic or ordering of the messages will be affected.

- * - *

This operation is not transactional so it may succeed for some topics while fail for others.

- * - *

It may take several seconds after this method returns - * success for all the brokers to become aware that the partitions have been created. - * During this time, {@link AdminClient#describeTopics(Collection)} - * may not return information about the new partitions.

- * - *

This operation is supported by brokers with version 1.0.0 or higher.

- * - *

The following exceptions can be anticipated when calling {@code get()} on the futures obtained from the - * {@link CreatePartitionsResult#values() values()} method of the returned {@code CreatePartitionsResult}

- *
    - *
  • {@link org.apache.kafka.common.errors.AuthorizationException} - * if the authenticated user is not authorized to alter the topic
    • - *
  • {@link org.apache.kafka.common.errors.TimeoutException} - * if the request was not completed in within the given {@link CreatePartitionsOptions#timeoutMs()}.
    • - *
  • {@link org.apache.kafka.common.errors.ReassignmentInProgressException} - * if a partition reassignment is currently in progress
    • - *
  • {@link org.apache.kafka.common.errors.BrokerNotAvailableException} - * if the requested {@link NewPartitions#assignments()} contain a broker that is currently unavailable.
    • - *
  • {@link org.apache.kafka.common.errors.InvalidReplicationFactorException} - * if no {@link NewPartitions#assignments()} are given and it is impossible for the broker to assign - * replicas with the topics replication factor.
    • - *
  • Subclasses of {@link org.apache.kafka.common.KafkaException} - * if the request is invalid in some way.
    • - *
- * - * @param newPartitions The topics which should have new partitions created, and corresponding parameters - * for the created partitions. - * @param options The options to use when creating the new paritions. - * @return The CreatePartitionsResult. - */ - public abstract CreatePartitionsResult createPartitions(Map newPartitions, - CreatePartitionsOptions options); - - /** - * Delete records whose offset is smaller than the given offset of the corresponding partition. - * - * This is a convenience method for {@link #deleteRecords(Map, DeleteRecordsOptions)} with default options. - * See the overload for more details. - * - * This operation is supported by brokers with version 0.11.0.0 or higher. - * - * @param recordsToDelete The topic partitions and related offsets from which records deletion starts. - * @return The DeleteRecordsResult. - */ - public DeleteRecordsResult deleteRecords(Map recordsToDelete) { - return deleteRecords(recordsToDelete, new DeleteRecordsOptions()); - } - - /** - * Delete records whose offset is smaller than the given offset of the corresponding partition. - * - * This operation is supported by brokers with version 0.11.0.0 or higher. - * - * @param recordsToDelete The topic partitions and related offsets from which records deletion starts. - * @param options The options to use when deleting records. - * @return The DeleteRecordsResult. - */ - public abstract DeleteRecordsResult deleteRecords(Map recordsToDelete, - DeleteRecordsOptions options); - - /** - *

Create a Delegation Token.

- * - *

This is a convenience method for {@link #createDelegationToken(CreateDelegationTokenOptions)} with default options. - * See the overload for more details.

- * - * @return The CreateDelegationTokenResult. - */ - public CreateDelegationTokenResult createDelegationToken() { - return createDelegationToken(new CreateDelegationTokenOptions()); - } - - - /** - *

Create a Delegation Token.

- * - *

This operation is supported by brokers with version 1.1.0 or higher.

- * - *

The following exceptions can be anticipated when calling {@code get()} on the futures obtained from the - * {@link CreateDelegationTokenResult#delegationToken() delegationToken()} method of the returned {@code CreateDelegationTokenResult}

- *
    - *
  • {@link org.apache.kafka.common.errors.UnsupportedByAuthenticationException} - * If the request sent on PLAINTEXT/1-way SSL channels or delegation token authenticated channels.
    • - *
  • {@link org.apache.kafka.common.errors.InvalidPrincipalTypeException} - * if the renewers principal type is not supported.
    • - *
  • {@link org.apache.kafka.common.errors.DelegationTokenDisabledException} - * if the delegation token feature is disabled.
    • - *
  • {@link org.apache.kafka.common.errors.TimeoutException} - * if the request was not completed in within the given {@link CreateDelegationTokenOptions#timeoutMs()}.
    • - *
- * - * @param options The options to use when creating delegation token. - * @return The DeleteRecordsResult. - */ - public abstract CreateDelegationTokenResult createDelegationToken(CreateDelegationTokenOptions options); - - - /** - *

Renew a Delegation Token.

- * - *

This is a convenience method for {@link #renewDelegationToken(byte[], RenewDelegationTokenOptions)} with default options. - * See the overload for more details.

- * - * - * @param hmac HMAC of the Delegation token - * @return The RenewDelegationTokenResult. - */ - public RenewDelegationTokenResult renewDelegationToken(byte[] hmac) { - return renewDelegationToken(hmac, new RenewDelegationTokenOptions()); - } - - /** - *

Renew a Delegation Token.

- * - *

This operation is supported by brokers with version 1.1.0 or higher.

- * - *

The following exceptions can be anticipated when calling {@code get()} on the futures obtained from the - * {@link RenewDelegationTokenResult#expiryTimestamp() expiryTimestamp()} method of the returned {@code RenewDelegationTokenResult}

- *
    - *
  • {@link org.apache.kafka.common.errors.UnsupportedByAuthenticationException} - * If the request sent on PLAINTEXT/1-way SSL channels or delegation token authenticated channels.
    • - *
  • {@link org.apache.kafka.common.errors.DelegationTokenDisabledException} - * if the delegation token feature is disabled.
    • - *
  • {@link org.apache.kafka.common.errors.DelegationTokenNotFoundException} - * if the delegation token is not found on server.
    • - *
  • {@link org.apache.kafka.common.errors.DelegationTokenOwnerMismatchException} - * if the authenticated user is not owner/renewer of the token.
    • - *
  • {@link org.apache.kafka.common.errors.DelegationTokenExpiredException} - * if the delegation token is expired.
    • - *
  • {@link org.apache.kafka.common.errors.TimeoutException} - * if the request was not completed in within the given {@link RenewDelegationTokenOptions#timeoutMs()}.
    • - *
- * - * @param hmac HMAC of the Delegation token - * @param options The options to use when renewing delegation token. - * @return The RenewDelegationTokenResult. - */ - public abstract RenewDelegationTokenResult renewDelegationToken(byte[] hmac, RenewDelegationTokenOptions options); - - /** - *

Expire a Delegation Token.

- * - *

This is a convenience method for {@link #expireDelegationToken(byte[], ExpireDelegationTokenOptions)} with default options. - * This will expire the token immediately. See the overload for more details.

- * - * @param hmac HMAC of the Delegation token - * @return The ExpireDelegationTokenResult. - */ - public ExpireDelegationTokenResult expireDelegationToken(byte[] hmac) { - return expireDelegationToken(hmac, new ExpireDelegationTokenOptions()); - } - - /** - *

Expire a Delegation Token.

- * - *

This operation is supported by brokers with version 1.1.0 or higher.

- * - *

The following exceptions can be anticipated when calling {@code get()} on the futures obtained from the - * {@link ExpireDelegationTokenResult#expiryTimestamp() expiryTimestamp()} method of the returned {@code ExpireDelegationTokenResult}

- *
    - *
  • {@link org.apache.kafka.common.errors.UnsupportedByAuthenticationException} - * If the request sent on PLAINTEXT/1-way SSL channels or delegation token authenticated channels.
    • - *
  • {@link org.apache.kafka.common.errors.DelegationTokenDisabledException} - * if the delegation token feature is disabled.
    • - *
  • {@link org.apache.kafka.common.errors.DelegationTokenNotFoundException} - * if the delegation token is not found on server.
    • - *
  • {@link org.apache.kafka.common.errors.DelegationTokenOwnerMismatchException} - * if the authenticated user is not owner/renewer of the requested token.
    • - *
  • {@link org.apache.kafka.common.errors.DelegationTokenExpiredException} - * if the delegation token is expired.
    • - *
  • {@link org.apache.kafka.common.errors.TimeoutException} - * if the request was not completed in within the given {@link ExpireDelegationTokenOptions#timeoutMs()}.
    • - *
- * - * @param hmac HMAC of the Delegation token - * @param options The options to use when expiring delegation token. - * @return The ExpireDelegationTokenResult. - */ - public abstract ExpireDelegationTokenResult expireDelegationToken(byte[] hmac, ExpireDelegationTokenOptions options); - - /** - *

Describe the Delegation Tokens.

- * - *

This is a convenience method for {@link #describeDelegationToken(DescribeDelegationTokenOptions)} with default options. - * This will return all the user owned tokens and tokens where user have Describe permission. See the overload for more details.

- * - * @return The DescribeDelegationTokenResult. - */ - public DescribeDelegationTokenResult describeDelegationToken() { - return describeDelegationToken(new DescribeDelegationTokenOptions()); - } - - /** - *

Describe the Delegation Tokens.

- * - *

This operation is supported by brokers with version 1.1.0 or higher.

- * - *

The following exceptions can be anticipated when calling {@code get()} on the futures obtained from the - * {@link DescribeDelegationTokenResult#delegationTokens() delegationTokens()} method of the returned {@code DescribeDelegationTokenResult}

- *
    - *
  • {@link org.apache.kafka.common.errors.UnsupportedByAuthenticationException} - * If the request sent on PLAINTEXT/1-way SSL channels or delegation token authenticated channels.
    • - *
  • {@link org.apache.kafka.common.errors.DelegationTokenDisabledException} - * if the delegation token feature is disabled.
    • - *
  • {@link org.apache.kafka.common.errors.TimeoutException} - * if the request was not completed in within the given {@link DescribeDelegationTokenOptions#timeoutMs()}.
    • - *
- * - * @param options The options to use when describing delegation tokens. - * @return The DescribeDelegationTokenResult. - */ - public abstract DescribeDelegationTokenResult describeDelegationToken(DescribeDelegationTokenOptions options); - - /** - * Describe some group IDs in the cluster. - * - * @param groupIds The IDs of the groups to describe. - * @param options The options to use when describing the groups. - * @return The DescribeConsumerGroupResult. - */ - public abstract DescribeConsumerGroupsResult describeConsumerGroups(Collection groupIds, - DescribeConsumerGroupsOptions options); - - /** - * Describe some group IDs in the cluster, with the default options. - *

- * This is a convenience method for - * #{@link AdminClient#describeConsumerGroups(Collection, DescribeConsumerGroupsOptions)} with - * default options. See the overload for more details. - * - * @param groupIds The IDs of the groups to describe. - * @return The DescribeConsumerGroupResult. - */ - public DescribeConsumerGroupsResult describeConsumerGroups(Collection groupIds) { - return describeConsumerGroups(groupIds, new DescribeConsumerGroupsOptions()); - } - - /** - * List the consumer groups available in the cluster. - * - * @param options The options to use when listing the consumer groups. - * @return The ListGroupsResult. - */ - public abstract ListConsumerGroupsResult listConsumerGroups(ListConsumerGroupsOptions options); - - /** - * List the consumer groups available in the cluster with the default options. - * - * This is a convenience method for #{@link AdminClient#listConsumerGroups(ListConsumerGroupsOptions)} with default options. - * See the overload for more details. - * - * @return The ListGroupsResult. - */ - public ListConsumerGroupsResult listConsumerGroups() { - return listConsumerGroups(new ListConsumerGroupsOptions()); - } - - /** - * List the consumer group offsets available in the cluster. - * - * @param options The options to use when listing the consumer group offsets. - * @return The ListGroupOffsetsResult - */ - public abstract ListConsumerGroupOffsetsResult listConsumerGroupOffsets(String groupId, ListConsumerGroupOffsetsOptions options); - - /** - * List the consumer group offsets available in the cluster with the default options. - * - * This is a convenience method for #{@link AdminClient#listConsumerGroupOffsets(String, ListConsumerGroupOffsetsOptions)} with default options. - * - * @return The ListGroupOffsetsResult. - */ - public ListConsumerGroupOffsetsResult listConsumerGroupOffsets(String groupId) { - return listConsumerGroupOffsets(groupId, new ListConsumerGroupOffsetsOptions()); - } - - /** - * Delete consumer groups from the cluster. - * - * @param options The options to use when deleting a consumer group. - * @return The DeletConsumerGroupResult. - */ - public abstract DeleteConsumerGroupsResult deleteConsumerGroups(Collection groupIds, DeleteConsumerGroupsOptions options); - - /** - * Delete consumer groups from the cluster with the default options. - * - * @return The DeleteConsumerGroupResult. - */ - public DeleteConsumerGroupsResult deleteConsumerGroups(Collection groupIds) { - return deleteConsumerGroups(groupIds, new DeleteConsumerGroupsOptions()); - } - - /** - * Elect the preferred broker of the given {@code partitions} as leader, or - * elect the preferred broker for all partitions as leader if the argument to {@code partitions} is null. - * - * This is a convenience method for {@link #electPreferredLeaders(Collection, ElectPreferredLeadersOptions)} - * with default options. - * See the overload for more details. - * - * @param partitions The partitions for which the preferred leader should be elected. - * @return The ElectPreferredLeadersResult. - */ - public ElectPreferredLeadersResult electPreferredLeaders(Collection partitions) { - return electPreferredLeaders(partitions, new ElectPreferredLeadersOptions()); - } - - /** - * Elect the preferred broker of the given {@code partitions} as leader, or - * elect the preferred broker for all partitions as leader if the argument to {@code partitions} is null. - * - * This operation is not transactional so it may succeed for some partitions while fail for others. - * - * It may take several seconds after this method returns - * success for all the brokers in the cluster to become aware that the partitions have new leaders. - * During this time, {@link AdminClient#describeTopics(Collection)} - * may not return information about the partitions' new leaders. - * - * This operation is supported by brokers with version 2.2.0 or higher. - * - *

The following exceptions can be anticipated when calling {@code get()} on the futures obtained from - * the returned {@code ElectPreferredLeadersResult}:

- *
    - *
  • {@link org.apache.kafka.common.errors.ClusterAuthorizationException} - * if the authenticated user didn't have alter access to the cluster.
    • - *
  • {@link org.apache.kafka.common.errors.UnknownTopicOrPartitionException} - * if the topic or partition did not exist within the cluster.
    • - *
  • {@link org.apache.kafka.common.errors.InvalidTopicException} - * if the topic was already queued for deletion.
    • - *
  • {@link org.apache.kafka.common.errors.NotControllerException} - * if the request was sent to a broker that was not the controller for the cluster.
    • - *
  • {@link org.apache.kafka.common.errors.TimeoutException} - * if the request timed out before the election was complete.
    • - *
  • {@link org.apache.kafka.common.errors.LeaderNotAvailableException} - * if the preferred leader was not alive or not in the ISR.
    • - *
- * - * @param partitions The partitions for which the preferred leader should be elected. - * @param options The options to use when electing the preferred leaders. - * @return The ElectPreferredLeadersResult. - */ - public abstract ElectPreferredLeadersResult electPreferredLeaders(Collection partitions, - ElectPreferredLeadersOptions options); - - /** - * Get the metrics kept by the adminClient - */ - public abstract Map metrics(); } + diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/AdminClientConfig.java b/clients/src/main/java/org/apache/kafka/clients/admin/AdminClientConfig.java index 47c76ac36af12..283d836887843 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/AdminClientConfig.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/AdminClientConfig.java @@ -23,6 +23,7 @@ import org.apache.kafka.common.config.ConfigDef; import org.apache.kafka.common.config.ConfigDef.Importance; import org.apache.kafka.common.config.ConfigDef.Type; +import org.apache.kafka.common.config.SecurityConfig; import org.apache.kafka.common.metrics.Sensor; import java.util.Map; @@ -99,14 +100,22 @@ public class AdminClientConfig extends AbstractConfig { private static final String METRICS_SAMPLE_WINDOW_MS_DOC = CommonClientConfigs.METRICS_SAMPLE_WINDOW_MS_DOC; public static final String METRICS_RECORDING_LEVEL_CONFIG = CommonClientConfigs.METRICS_RECORDING_LEVEL_CONFIG; + public static final String METRICS_REPLACE_ON_DUPLICATE_CONFIG = CommonClientConfigs.METRICS_REPLACE_ON_DUPLICATE_CONFIG; public static final String SECURITY_PROTOCOL_CONFIG = CommonClientConfigs.SECURITY_PROTOCOL_CONFIG; public static final String DEFAULT_SECURITY_PROTOCOL = CommonClientConfigs.DEFAULT_SECURITY_PROTOCOL; private static final String SECURITY_PROTOCOL_DOC = CommonClientConfigs.SECURITY_PROTOCOL_DOC; private static final String METRICS_RECORDING_LEVEL_DOC = CommonClientConfigs.METRICS_RECORDING_LEVEL_DOC; + private static final String METRICS_REPLACE_ON_DUPLICATE_DOC = CommonClientConfigs.METRICS_REPLACE_ON_DUPLICATE_DOC; public static final String RETRIES_CONFIG = CommonClientConfigs.RETRIES_CONFIG; + /** + * security.providers + */ + public static final String SECURITY_PROVIDERS_CONFIG = SecurityConfig.SECURITY_PROVIDERS_CONFIG; + private static final String SECURITY_PROVIDERS_DOC = SecurityConfig.SECURITY_PROVIDERS_DOC; + static { CONFIG = new ConfigDef().define(BOOTSTRAP_SERVERS_CONFIG, Type.LIST, @@ -173,7 +182,17 @@ public class AdminClientConfig extends AbstractConfig { ClientDnsLookup.RESOLVE_CANONICAL_BOOTSTRAP_SERVERS_ONLY.toString()), Importance.MEDIUM, CLIENT_DNS_LOOKUP_DOC) + .define(METRICS_REPLACE_ON_DUPLICATE_CONFIG, + Type.BOOLEAN, + false, + Importance.LOW, + METRICS_REPLACE_ON_DUPLICATE_DOC) // security support + .define(SECURITY_PROVIDERS_CONFIG, + Type.STRING, + null, + Importance.LOW, + SECURITY_PROVIDERS_DOC) .define(SECURITY_PROTOCOL_CONFIG, Type.STRING, DEFAULT_SECURITY_PROTOCOL, @@ -200,8 +219,12 @@ public static Set configNames() { return CONFIG.names(); } + public static ConfigDef configDef() { + return new ConfigDef(CONFIG); + } + public static void main(String[] args) { - System.out.println(CONFIG.toHtmlTable()); + System.out.println(CONFIG.toHtml()); } } diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/AlterConfigOp.java b/clients/src/main/java/org/apache/kafka/clients/admin/AlterConfigOp.java new file mode 100644 index 0000000000000..950021aabd26e --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/AlterConfigOp.java @@ -0,0 +1,96 @@ +/* + * 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.admin; + +import org.apache.kafka.common.annotation.InterfaceStability; + +import java.util.Arrays; +import java.util.Collections; +import java.util.Map; +import java.util.Objects; +import java.util.function.Function; +import java.util.stream.Collectors; + +/** + * A class representing a alter configuration entry containing name, value and operation type. + * + * The API of this class is evolving, see {@link Admin} for details. + */ +@InterfaceStability.Evolving +public class AlterConfigOp { + + public enum OpType { + SET((byte) 0), DELETE((byte) 1), APPEND((byte) 2), SUBTRACT((byte) 3); + + private static final Map OP_TYPES = Collections.unmodifiableMap( + Arrays.stream(values()).collect(Collectors.toMap(OpType::id, Function.identity())) + ); + + private final byte id; + + OpType(final byte id) { + this.id = id; + } + + public byte id() { + return id; + } + + public static OpType forId(final byte id) { + return OP_TYPES.get(id); + } + } + + private final ConfigEntry configEntry; + private final OpType opType; + + public AlterConfigOp(ConfigEntry configEntry, OpType operationType) { + this.configEntry = configEntry; + this.opType = operationType; + } + + public ConfigEntry configEntry() { + return configEntry; + }; + + public OpType opType() { + return opType; + }; + + @Override + public boolean equals(final Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + final AlterConfigOp that = (AlterConfigOp) o; + return opType == that.opType && + Objects.equals(configEntry, that.configEntry); + } + + @Override + public int hashCode() { + return Objects.hash(opType, configEntry); + } + + @Override + public String toString() { + return "AlterConfigOp{" + + "opType=" + opType + + ", configEntry=" + configEntry + + '}'; + } +} diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/AlterConfigsOptions.java b/clients/src/main/java/org/apache/kafka/clients/admin/AlterConfigsOptions.java index 2dbeba2c937f6..0b280532104e8 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/AlterConfigsOptions.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/AlterConfigsOptions.java @@ -22,9 +22,9 @@ import java.util.Map; /** - * Options for {@link AdminClient#alterConfigs(Map)}. + * Options for {@link Admin#alterConfigs(Map)}. * - * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class AlterConfigsOptions extends AbstractOptions { diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/AlterConfigsResult.java b/clients/src/main/java/org/apache/kafka/clients/admin/AlterConfigsResult.java index df6c1c29d5a63..29056ce29403d 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/AlterConfigsResult.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/AlterConfigsResult.java @@ -24,9 +24,9 @@ import java.util.Map; /** - * The result of the {@link AdminClient#alterConfigs(Map)} call. + * The result of the {@link Admin#alterConfigs(Map)} call. * - * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class AlterConfigsResult { diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/AlterPartitionReassignmentsOptions.java b/clients/src/main/java/org/apache/kafka/clients/admin/AlterPartitionReassignmentsOptions.java new file mode 100644 index 0000000000000..bee9c70ab0b19 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/AlterPartitionReassignmentsOptions.java @@ -0,0 +1,31 @@ +/* + * 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.admin; + +import org.apache.kafka.common.annotation.InterfaceStability; + +import java.util.Map; + +/** + * Options for {@link AdminClient#alterPartitionReassignments(Map, AlterPartitionReassignmentsOptions)} + * + * The API of this class is evolving. See {@link AdminClient} for details. + */ +@InterfaceStability.Evolving +public class AlterPartitionReassignmentsOptions extends AbstractOptions { +} diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/AlterPartitionReassignmentsResult.java b/clients/src/main/java/org/apache/kafka/clients/admin/AlterPartitionReassignmentsResult.java new file mode 100644 index 0000000000000..2009ab5b6b7f2 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/AlterPartitionReassignmentsResult.java @@ -0,0 +1,59 @@ +/* + * 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.admin; + +import org.apache.kafka.common.KafkaFuture; +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.annotation.InterfaceStability; + +import java.util.Map; + +/** + * The result of {@link AdminClient#alterPartitionReassignments(Map, AlterPartitionReassignmentsOptions)}. + * + * The API of this class is evolving. See {@link AdminClient} for details. + */ +@InterfaceStability.Evolving +public class AlterPartitionReassignmentsResult { + private final Map> futures; + + AlterPartitionReassignmentsResult(Map> futures) { + this.futures = futures; + } + + /** + * Return a map from partitions to futures which can be used to check the status of the reassignment. + * + * Possible error codes: + * + * INVALID_REPLICA_ASSIGNMENT (39) - if the specified replica assignment was not valid -- for example, if it included negative numbers, repeated numbers, or specified a broker ID that the controller was not aware of. + * NO_REASSIGNMENT_IN_PROGRESS (85) - if the request wants to cancel reassignments but none exist + * UNKNOWN (-1) + * + */ + public Map> values() { + return futures; + } + + /** + * Return a future which succeeds only if all the reassignments were successfully initiated. + */ + public KafkaFuture all() { + return KafkaFuture.allOf(futures.values().toArray(new KafkaFuture[0])); + } +} diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/AlterReplicaLogDirsOptions.java b/clients/src/main/java/org/apache/kafka/clients/admin/AlterReplicaLogDirsOptions.java index d6892efc9e974..76037fbb91333 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/AlterReplicaLogDirsOptions.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/AlterReplicaLogDirsOptions.java @@ -21,7 +21,7 @@ import java.util.Map; /** - * Options for {@link AdminClient#alterReplicaLogDirs(Map, AlterReplicaLogDirsOptions)}. + * Options for {@link Admin#alterReplicaLogDirs(Map, AlterReplicaLogDirsOptions)}. */ @InterfaceStability.Evolving public class AlterReplicaLogDirsOptions extends AbstractOptions { diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/AlterReplicaLogDirsResult.java b/clients/src/main/java/org/apache/kafka/clients/admin/AlterReplicaLogDirsResult.java index a3da216e32ead..2373265564141 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/AlterReplicaLogDirsResult.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/AlterReplicaLogDirsResult.java @@ -23,7 +23,7 @@ /** - * The result of {@link AdminClient#alterReplicaLogDirs(Map, AlterReplicaLogDirsOptions)}. + * The result of {@link Admin#alterReplicaLogDirs(Map, AlterReplicaLogDirsOptions)}. */ @InterfaceStability.Evolving public class AlterReplicaLogDirsResult { diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/Config.java b/clients/src/main/java/org/apache/kafka/clients/admin/Config.java index c81c0b60274b0..ae7c03a5d669e 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/Config.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/Config.java @@ -27,7 +27,7 @@ /** * A configuration object containing the configuration entries for a resource. *

- * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class Config { diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/ConfigEntry.java b/clients/src/main/java/org/apache/kafka/clients/admin/ConfigEntry.java index 9976108acb7ba..42cc6274980e7 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/ConfigEntry.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/ConfigEntry.java @@ -26,7 +26,7 @@ /** * A class representing a configuration entry containing name, value and additional metadata. * - * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class ConfigEntry { @@ -189,6 +189,7 @@ public String toString() { */ public enum ConfigSource { DYNAMIC_TOPIC_CONFIG, // dynamic topic config that is configured for a specific topic + DYNAMIC_BROKER_LOGGER_CONFIG, // dynamic broker logger config that is configured for a specific broker DYNAMIC_BROKER_CONFIG, // dynamic broker config that is configured for a specific broker DYNAMIC_DEFAULT_BROKER_CONFIG, // dynamic broker config that is configured as default for all brokers in the cluster STATIC_BROKER_CONFIG, // static broker config provided as broker properties at start up (e.g. server.properties file) diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/ConsumerGroupDescription.java b/clients/src/main/java/org/apache/kafka/clients/admin/ConsumerGroupDescription.java index 8dd6018427fba..32bd165817457 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/ConsumerGroupDescription.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/ConsumerGroupDescription.java @@ -38,15 +38,24 @@ public class ConsumerGroupDescription { private final String partitionAssignor; private final ConsumerGroupState state; private final Node coordinator; - private Set authorizedOperations; + private final Set authorizedOperations; public ConsumerGroupDescription(String groupId, boolean isSimpleConsumerGroup, Collection members, String partitionAssignor, ConsumerGroupState state, - Node coordinator, - Set authorizedOperations) { + Node coordinator) { + this(groupId, isSimpleConsumerGroup, members, partitionAssignor, state, coordinator, Collections.emptySet()); + } + + ConsumerGroupDescription(String groupId, + boolean isSimpleConsumerGroup, + Collection members, + String partitionAssignor, + ConsumerGroupState state, + Node coordinator, + Set authorizedOperations) { this.groupId = groupId == null ? "" : groupId; this.isSimpleConsumerGroup = isSimpleConsumerGroup; this.members = members == null ? Collections.emptyList() : @@ -119,7 +128,7 @@ public Node coordinator() { } /** - * authorizedOperations for this group + * authorizedOperations for this group, or null if that information is not known. */ public Set authorizedOperations() { return authorizedOperations; diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/CreateAclsOptions.java b/clients/src/main/java/org/apache/kafka/clients/admin/CreateAclsOptions.java index a7b92ba93d982..bfb8e32db157f 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/CreateAclsOptions.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/CreateAclsOptions.java @@ -22,9 +22,9 @@ import java.util.Collection; /** - * Options for {@link AdminClient#createAcls(Collection)}. + * Options for {@link Admin#createAcls(Collection)}. * - * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class CreateAclsOptions extends AbstractOptions { diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/CreateAclsResult.java b/clients/src/main/java/org/apache/kafka/clients/admin/CreateAclsResult.java index 2917f17f7560d..6e69554635efc 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/CreateAclsResult.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/CreateAclsResult.java @@ -25,9 +25,9 @@ import java.util.Map; /** - * The result of the {@link AdminClient#createAcls(Collection)} call. + * The result of the {@link Admin#createAcls(Collection)} call. * - * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class CreateAclsResult { diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/CreateDelegationTokenOptions.java b/clients/src/main/java/org/apache/kafka/clients/admin/CreateDelegationTokenOptions.java index 1b77b943800ec..6a082d499bbb4 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/CreateDelegationTokenOptions.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/CreateDelegationTokenOptions.java @@ -24,9 +24,9 @@ import org.apache.kafka.common.security.auth.KafkaPrincipal; /** - * Options for {@link AdminClient#createDelegationToken(CreateDelegationTokenOptions)}. + * Options for {@link Admin#createDelegationToken(CreateDelegationTokenOptions)}. * - * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class CreateDelegationTokenOptions extends AbstractOptions { diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/CreateDelegationTokenResult.java b/clients/src/main/java/org/apache/kafka/clients/admin/CreateDelegationTokenResult.java index 043cbe87fef5c..7aa48049226e0 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/CreateDelegationTokenResult.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/CreateDelegationTokenResult.java @@ -24,7 +24,7 @@ /** * The result of the {@link KafkaAdminClient#createDelegationToken(CreateDelegationTokenOptions)} call. * - * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class CreateDelegationTokenResult { diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/CreatePartitionsOptions.java b/clients/src/main/java/org/apache/kafka/clients/admin/CreatePartitionsOptions.java index aafc2078003ed..01eb5a0fce8e9 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/CreatePartitionsOptions.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/CreatePartitionsOptions.java @@ -22,9 +22,9 @@ import java.util.Map; /** - * Options for {@link AdminClient#createPartitions(Map)}. + * Options for {@link Admin#createPartitions(Map)}. * - * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class CreatePartitionsOptions extends AbstractOptions { diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/CreatePartitionsResult.java b/clients/src/main/java/org/apache/kafka/clients/admin/CreatePartitionsResult.java index c3a504b615c0c..8b864b697bf76 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/CreatePartitionsResult.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/CreatePartitionsResult.java @@ -23,9 +23,9 @@ import java.util.Map; /** - * The result of the {@link AdminClient#createPartitions(Map)} call. + * The result of the {@link Admin#createPartitions(Map)} call. * - * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class CreatePartitionsResult { diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/CreateTopicsOptions.java b/clients/src/main/java/org/apache/kafka/clients/admin/CreateTopicsOptions.java index 7cd0df86e8b94..a9f1009c2fc7e 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/CreateTopicsOptions.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/CreateTopicsOptions.java @@ -22,9 +22,9 @@ import java.util.Collection; /** - * Options for {@link AdminClient#createTopics(Collection)}. + * Options for {@link Admin#createTopics(Collection)}. * - * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class CreateTopicsOptions extends AbstractOptions { diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/CreateTopicsResult.java b/clients/src/main/java/org/apache/kafka/clients/admin/CreateTopicsResult.java index 404cb918b8455..85c0a905240b3 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/CreateTopicsResult.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/CreateTopicsResult.java @@ -18,20 +18,24 @@ import org.apache.kafka.common.KafkaFuture; import org.apache.kafka.common.annotation.InterfaceStability; +import org.apache.kafka.common.errors.ApiException; import java.util.Collection; import java.util.Map; +import java.util.stream.Collectors; /** - * The result of {@link AdminClient#createTopics(Collection)}. + * The result of {@link Admin#createTopics(Collection)}. * - * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class CreateTopicsResult { - private final Map> futures; + final static int UNKNOWN = -1; - CreateTopicsResult(Map> futures) { + private final Map> futures; + + CreateTopicsResult(Map> futures) { this.futures = futures; } @@ -40,7 +44,8 @@ public class CreateTopicsResult { * topic creations. */ public Map> values() { - return futures; + return futures.entrySet().stream() + .collect(Collectors.toMap(Map.Entry::getKey, e -> e.getValue().thenApply(v -> (Void) null))); } /** @@ -49,4 +54,84 @@ public Map> values() { public KafkaFuture all() { return KafkaFuture.allOf(futures.values().toArray(new KafkaFuture[0])); } + + /** + * Returns a future that provides topic configs for the topic when the request completes. + *

+ * If broker version doesn't support replication factor in the response, throw + * {@link org.apache.kafka.common.errors.UnsupportedVersionException}. + * If broker returned an error for topic configs, throw appropriate exception. For example, + * {@link org.apache.kafka.common.errors.TopicAuthorizationException} is thrown if user does not + * have permission to describe topic configs. + */ + public KafkaFuture config(String topic) { + return futures.get(topic).thenApply(TopicMetadataAndConfig::config); + } + + /** + * Returns a future that provides number of partitions in the topic when the request completes. + *

+ * If broker version doesn't support replication factor in the response, throw + * {@link org.apache.kafka.common.errors.UnsupportedVersionException}. + * If broker returned an error for topic configs, throw appropriate exception. For example, + * {@link org.apache.kafka.common.errors.TopicAuthorizationException} is thrown if user does not + * have permission to describe topic configs. + */ + public KafkaFuture numPartitions(String topic) { + return futures.get(topic).thenApply(TopicMetadataAndConfig::numPartitions); + } + + /** + * Returns a future that provides replication factor for the topic when the request completes. + *

+ * If broker version doesn't support replication factor in the response, throw + * {@link org.apache.kafka.common.errors.UnsupportedVersionException}. + * If broker returned an error for topic configs, throw appropriate exception. For example, + * {@link org.apache.kafka.common.errors.TopicAuthorizationException} is thrown if user does not + * have permission to describe topic configs. + */ + public KafkaFuture replicationFactor(String topic) { + return futures.get(topic).thenApply(TopicMetadataAndConfig::replicationFactor); + } + + static class TopicMetadataAndConfig { + private final ApiException exception; + private final int numPartitions; + private final int replicationFactor; + private final Config config; + + TopicMetadataAndConfig(int numPartitions, int replicationFactor, Config config) { + this.exception = null; + this.numPartitions = numPartitions; + this.replicationFactor = replicationFactor; + this.config = config; + } + + TopicMetadataAndConfig(ApiException exception) { + this.exception = exception; + this.numPartitions = UNKNOWN; + this.replicationFactor = UNKNOWN; + this.config = null; + } + + public int numPartitions() { + ensureSuccess(); + return numPartitions; + } + + public int replicationFactor() { + ensureSuccess(); + return replicationFactor; + } + + public Config config() { + ensureSuccess(); + return config; + } + + private void ensureSuccess() { + if (exception != null) + throw exception; + } + } } diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/DeleteAclsOptions.java b/clients/src/main/java/org/apache/kafka/clients/admin/DeleteAclsOptions.java index 56f4b78b64c70..1b67da52f386d 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/DeleteAclsOptions.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/DeleteAclsOptions.java @@ -22,9 +22,9 @@ import java.util.Collection; /** - * Options for the {@link AdminClient#deleteAcls(Collection)} call. + * Options for the {@link Admin#deleteAcls(Collection)} call. * - * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class DeleteAclsOptions extends AbstractOptions { diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/DeleteAclsResult.java b/clients/src/main/java/org/apache/kafka/clients/admin/DeleteAclsResult.java index 63310bca84ea9..391b9d1d5f4f6 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/DeleteAclsResult.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/DeleteAclsResult.java @@ -30,9 +30,9 @@ import java.util.Map; /** - * The result of the {@link AdminClient#deleteAcls(Collection)} call. + * The result of the {@link Admin#deleteAcls(Collection)} call. * - * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class DeleteAclsResult { diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/DeleteConsumerGroupOffsetsOptions.java b/clients/src/main/java/org/apache/kafka/clients/admin/DeleteConsumerGroupOffsetsOptions.java new file mode 100644 index 0000000000000..63e6b4be84bae --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/DeleteConsumerGroupOffsetsOptions.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.clients.admin; + +import java.util.Set; +import org.apache.kafka.common.annotation.InterfaceStability; + +/** + * Options for the {@link Admin#deleteConsumerGroupOffsets(String, Set)} call. + * + * The API of this class is evolving, see {@link Admin} for details. + */ +@InterfaceStability.Evolving +public class DeleteConsumerGroupOffsetsOptions extends AbstractOptions { + +} diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/DeleteConsumerGroupOffsetsResult.java b/clients/src/main/java/org/apache/kafka/clients/admin/DeleteConsumerGroupOffsetsResult.java new file mode 100644 index 0000000000000..336e9c0f3b269 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/DeleteConsumerGroupOffsetsResult.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 org.apache.kafka.clients.admin; + +import java.util.Set; +import org.apache.kafka.common.KafkaFuture; +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.annotation.InterfaceStability; + +import java.util.Map; +import org.apache.kafka.common.internals.KafkaFutureImpl; +import org.apache.kafka.common.protocol.Errors; + +/** + * The result of the {@link Admin#deleteConsumerGroupOffsets(String, Set)} call. + * + * The API of this class is evolving, see {@link Admin} for details. + */ +@InterfaceStability.Evolving +public class DeleteConsumerGroupOffsetsResult { + private final KafkaFuture> future; + private final Set partitions; + + + DeleteConsumerGroupOffsetsResult(KafkaFuture> future, Set partitions) { + this.future = future; + this.partitions = partitions; + } + + /** + * Return a future which can be used to check the result for a given partition. + */ + public KafkaFuture partitionResult(final TopicPartition partition) { + if (!partitions.contains(partition)) { + throw new IllegalArgumentException("Partition " + partition + " was not included in the original request"); + } + final KafkaFutureImpl result = new KafkaFutureImpl<>(); + + this.future.whenComplete((topicPartitions, throwable) -> { + if (throwable != null) { + result.completeExceptionally(throwable); + } else if (!maybeCompleteExceptionally(topicPartitions, partition, result)) { + result.complete(null); + } + }); + return result; + } + + /** + * Return a future which succeeds only if all the deletions succeed. + * If not, the first partition error shall be returned. + */ + public KafkaFuture all() { + final KafkaFutureImpl result = new KafkaFutureImpl<>(); + + this.future.whenComplete((topicPartitions, throwable) -> { + if (throwable != null) { + result.completeExceptionally(throwable); + } else { + for (TopicPartition partition : partitions) { + if (maybeCompleteExceptionally(topicPartitions, partition, result)) { + return; + } + } + result.complete(null); + } + }); + return result; + } + + private boolean maybeCompleteExceptionally(Map partitionLevelErrors, + TopicPartition partition, + KafkaFutureImpl result) { + Throwable exception = KafkaAdminClient.getSubLevelError(partitionLevelErrors, partition, + "Offset deletion result for partition \"" + partition + "\" was not included in the response"); + if (exception != null) { + result.completeExceptionally(exception); + return true; + } else { + return false; + } + } +} diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/DeleteConsumerGroupsOptions.java b/clients/src/main/java/org/apache/kafka/clients/admin/DeleteConsumerGroupsOptions.java index cd505f4c11138..081aeab6e750a 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/DeleteConsumerGroupsOptions.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/DeleteConsumerGroupsOptions.java @@ -21,9 +21,9 @@ import java.util.Collection; /** - * Options for the {@link AdminClient#deleteConsumerGroups(Collection)} call. + * Options for the {@link Admin#deleteConsumerGroups(Collection)} call. * - * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class DeleteConsumerGroupsOptions extends AbstractOptions { diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/DeleteConsumerGroupsResult.java b/clients/src/main/java/org/apache/kafka/clients/admin/DeleteConsumerGroupsResult.java index dd6835cf10cbc..c7d7a5ab672a5 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/DeleteConsumerGroupsResult.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/DeleteConsumerGroupsResult.java @@ -23,9 +23,9 @@ import java.util.Map; /** - * The result of the {@link AdminClient#deleteConsumerGroups(Collection)} call. + * The result of the {@link Admin#deleteConsumerGroups(Collection)} call. * - * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class DeleteConsumerGroupsResult { diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/DeleteRecordsOptions.java b/clients/src/main/java/org/apache/kafka/clients/admin/DeleteRecordsOptions.java index 2581694b4ff6a..34af759559a20 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/DeleteRecordsOptions.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/DeleteRecordsOptions.java @@ -22,9 +22,9 @@ import java.util.Map; /** - * Options for {@link AdminClient#deleteRecords(Map, DeleteRecordsOptions)}. + * Options for {@link Admin#deleteRecords(Map, DeleteRecordsOptions)}. * - * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class DeleteRecordsOptions extends AbstractOptions { diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/DeleteRecordsResult.java b/clients/src/main/java/org/apache/kafka/clients/admin/DeleteRecordsResult.java index 8a7d39ebbdf64..01966322f674d 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/DeleteRecordsResult.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/DeleteRecordsResult.java @@ -24,9 +24,9 @@ import java.util.Map; /** - * The result of the {@link AdminClient#deleteRecords(Map)} call. + * The result of the {@link Admin#deleteRecords(Map)} call. * - * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class DeleteRecordsResult { diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/DeleteTopicsOptions.java b/clients/src/main/java/org/apache/kafka/clients/admin/DeleteTopicsOptions.java index 27c9af9f17fee..91e38a196fc9a 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/DeleteTopicsOptions.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/DeleteTopicsOptions.java @@ -22,9 +22,9 @@ import java.util.Collection; /** - * Options for {@link AdminClient#deleteTopics(Collection)}. + * Options for {@link Admin#deleteTopics(Collection)}. * - * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class DeleteTopicsOptions extends AbstractOptions { diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/DeleteTopicsResult.java b/clients/src/main/java/org/apache/kafka/clients/admin/DeleteTopicsResult.java index 9148a76c50778..d48c7e00ed4a4 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/DeleteTopicsResult.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/DeleteTopicsResult.java @@ -24,9 +24,9 @@ import java.util.Map; /** - * The result of the {@link AdminClient#deleteTopics(Collection)} call. + * The result of the {@link Admin#deleteTopics(Collection)} call. * - * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class DeleteTopicsResult { diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/DescribeAclsOptions.java b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeAclsOptions.java index 89ba4bc9bece7..b17d6a7d0cb98 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/DescribeAclsOptions.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeAclsOptions.java @@ -21,9 +21,9 @@ import org.apache.kafka.common.annotation.InterfaceStability; /** - * Options for {@link AdminClient#describeAcls(AclBindingFilter)}. + * Options for {@link Admin#describeAcls(AclBindingFilter)}. * - * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class DescribeAclsOptions extends AbstractOptions { diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/DescribeAclsResult.java b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeAclsResult.java index e09bf437ab2bd..fb16222487f85 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/DescribeAclsResult.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeAclsResult.java @@ -27,7 +27,7 @@ /** * The result of the {@link KafkaAdminClient#describeAcls(AclBindingFilter)} call. * - * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class DescribeAclsResult { diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/DescribeClusterOptions.java b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeClusterOptions.java index 92640fd54ea42..670feda0d2614 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/DescribeClusterOptions.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeClusterOptions.java @@ -20,13 +20,15 @@ import org.apache.kafka.common.annotation.InterfaceStability; /** - * Options for {@link AdminClient#describeCluster()}. + * Options for {@link Admin#describeCluster()}. * - * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class DescribeClusterOptions extends AbstractOptions { + private boolean includeAuthorizedOperations; + /** * Set the request timeout in milliseconds for this operation or {@code null} if the default request timeout for the * AdminClient should be used. @@ -38,4 +40,16 @@ public DescribeClusterOptions timeoutMs(Integer timeoutMs) { return this; } + public DescribeClusterOptions includeAuthorizedOperations(boolean includeAuthorizedOperations) { + this.includeAuthorizedOperations = includeAuthorizedOperations; + return this; + } + + /** + * Specify if authorized operations should be included in the response. Note that some + * older brokers cannot not supply this information even if it is requested. + */ + public boolean includeAuthorizedOperations() { + return includeAuthorizedOperations; + } } diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/DescribeClusterResult.java b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeClusterResult.java index 7d3ffc66fb550..d307d255fe2a2 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/DescribeClusterResult.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeClusterResult.java @@ -19,27 +19,32 @@ import org.apache.kafka.common.KafkaFuture; import org.apache.kafka.common.Node; +import org.apache.kafka.common.acl.AclOperation; import org.apache.kafka.common.annotation.InterfaceStability; import java.util.Collection; +import java.util.Set; /** * The result of the {@link KafkaAdminClient#describeCluster()} call. * - * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class DescribeClusterResult { private final KafkaFuture> nodes; private final KafkaFuture controller; private final KafkaFuture clusterId; + private final KafkaFuture> authorizedOperations; DescribeClusterResult(KafkaFuture> nodes, KafkaFuture controller, - KafkaFuture clusterId) { + KafkaFuture clusterId, + KafkaFuture> authorizedOperations) { this.nodes = nodes; this.controller = controller; this.clusterId = clusterId; + this.authorizedOperations = authorizedOperations; } /** @@ -64,4 +69,12 @@ public KafkaFuture controller() { public KafkaFuture clusterId() { return clusterId; } + + /** + * Returns a future which yields authorized operations. The future value will be non-null if the + * broker supplied this information, and null otherwise. + */ + public KafkaFuture> authorizedOperations() { + return authorizedOperations; + } } diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/DescribeConfigsOptions.java b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeConfigsOptions.java index aa667af3df27e..83582b67f14aa 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/DescribeConfigsOptions.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeConfigsOptions.java @@ -22,9 +22,9 @@ import java.util.Collection; /** - * Options for {@link AdminClient#describeConfigs(Collection)}. + * Options for {@link Admin#describeConfigs(Collection)}. * - * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class DescribeConfigsOptions extends AbstractOptions { diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/DescribeConfigsResult.java b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeConfigsResult.java index e5d79e8074320..e18c3b8b81427 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/DescribeConfigsResult.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeConfigsResult.java @@ -29,7 +29,7 @@ /** * The result of the {@link KafkaAdminClient#describeConfigs(Collection)} call. * - * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class DescribeConfigsResult { diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/DescribeConsumerGroupsOptions.java b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeConsumerGroupsOptions.java index 8f05f6183d4ac..70238a8b1d5f3 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/DescribeConsumerGroupsOptions.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeConsumerGroupsOptions.java @@ -22,9 +22,9 @@ import java.util.Collection; /** - * Options for {@link AdminClient#describeConsumerGroups(Collection, DescribeConsumerGroupsOptions)}. + * Options for {@link Admin#describeConsumerGroups(Collection, DescribeConsumerGroupsOptions)}. *

- * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class DescribeConsumerGroupsOptions extends AbstractOptions { diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/DescribeConsumerGroupsResult.java b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeConsumerGroupsResult.java index 8f0ebad4f7ac6..2eddbba305db7 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/DescribeConsumerGroupsResult.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeConsumerGroupsResult.java @@ -29,7 +29,7 @@ /** * The result of the {@link KafkaAdminClient#describeConsumerGroups(Collection, DescribeConsumerGroupsOptions)}} call. * - * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class DescribeConsumerGroupsResult { diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/DescribeDelegationTokenOptions.java b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeDelegationTokenOptions.java index 60b99354e3552..ef9f105850a5f 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/DescribeDelegationTokenOptions.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeDelegationTokenOptions.java @@ -23,9 +23,9 @@ import org.apache.kafka.common.security.auth.KafkaPrincipal; /** - * Options for {@link AdminClient#describeDelegationToken(DescribeDelegationTokenOptions)}. + * Options for {@link Admin#describeDelegationToken(DescribeDelegationTokenOptions)}. * - * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class DescribeDelegationTokenOptions extends AbstractOptions { diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/DescribeDelegationTokenResult.java b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeDelegationTokenResult.java index 7a9d4b9dd9737..47b2530328199 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/DescribeDelegationTokenResult.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeDelegationTokenResult.java @@ -26,7 +26,7 @@ /** * The result of the {@link KafkaAdminClient#describeDelegationToken(DescribeDelegationTokenOptions)} call. * - * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class DescribeDelegationTokenResult { diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/DescribeLogDirsOptions.java b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeLogDirsOptions.java index 5f6352cfefc6a..17890cad3002c 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/DescribeLogDirsOptions.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeLogDirsOptions.java @@ -23,9 +23,9 @@ /** - * Options for {@link AdminClient#describeLogDirs(Collection)} + * Options for {@link Admin#describeLogDirs(Collection)} * - * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class DescribeLogDirsOptions extends AbstractOptions { diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/DescribeLogDirsResult.java b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeLogDirsResult.java index 7c7bde755e2cd..9bcc2fb29fef8 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/DescribeLogDirsResult.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeLogDirsResult.java @@ -27,9 +27,9 @@ /** - * The result of the {@link AdminClient#describeLogDirs(Collection)} call. + * The result of the {@link Admin#describeLogDirs(Collection)} call. * - * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class DescribeLogDirsResult { diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/DescribeReplicaLogDirsOptions.java b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeReplicaLogDirsOptions.java index c0924efc4b32a..589de503df5b9 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/DescribeReplicaLogDirsOptions.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeReplicaLogDirsOptions.java @@ -21,9 +21,9 @@ import java.util.Collection; /** - * Options for {@link AdminClient#describeReplicaLogDirs(Collection)}. + * Options for {@link Admin#describeReplicaLogDirs(Collection)}. * - * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class DescribeReplicaLogDirsOptions extends AbstractOptions { diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/DescribeReplicaLogDirsResult.java b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeReplicaLogDirsResult.java index d4c847946360b..54bd9c142b0b0 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/DescribeReplicaLogDirsResult.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeReplicaLogDirsResult.java @@ -28,9 +28,9 @@ /** - * The result of {@link AdminClient#describeReplicaLogDirs(Collection)}. + * The result of {@link Admin#describeReplicaLogDirs(Collection)}. * - * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class DescribeReplicaLogDirsResult { diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/DescribeTopicsOptions.java b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeTopicsOptions.java index cc3d420e134cd..196341b6b937f 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/DescribeTopicsOptions.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeTopicsOptions.java @@ -22,13 +22,15 @@ import java.util.Collection; /** - * Options for {@link AdminClient#describeTopics(Collection)}. + * Options for {@link Admin#describeTopics(Collection)}. * - * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class DescribeTopicsOptions extends AbstractOptions { + private boolean includeAuthorizedOperations; + /** * Set the request timeout in milliseconds for this operation or {@code null} if the default request timeout for the * AdminClient should be used. @@ -40,4 +42,13 @@ public DescribeTopicsOptions timeoutMs(Integer timeoutMs) { return this; } + public DescribeTopicsOptions includeAuthorizedOperations(boolean includeAuthorizedOperations) { + this.includeAuthorizedOperations = includeAuthorizedOperations; + return this; + } + + public boolean includeAuthorizedOperations() { + return includeAuthorizedOperations; + } + } diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/DescribeTopicsResult.java b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeTopicsResult.java index 6bb24d91026a4..9822b423c4d82 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/DescribeTopicsResult.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeTopicsResult.java @@ -28,7 +28,7 @@ /** * The result of the {@link KafkaAdminClient#describeTopics(Collection)} call. * - * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class DescribeTopicsResult { diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/ElectLeadersOptions.java b/clients/src/main/java/org/apache/kafka/clients/admin/ElectLeadersOptions.java new file mode 100644 index 0000000000000..db79976bcc6ac --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/ElectLeadersOptions.java @@ -0,0 +1,29 @@ +/* + * 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.admin; + +import org.apache.kafka.common.annotation.InterfaceStability; + +/** + * Options for {@link Admin#electLeaders(ElectionType, Set, ElectLeadersOptions)}. + * + * The API of this class is evolving, see {@link Admin} for details. + */ +@InterfaceStability.Evolving +final public class ElectLeadersOptions extends AbstractOptions { +} diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/ElectLeadersResult.java b/clients/src/main/java/org/apache/kafka/clients/admin/ElectLeadersResult.java new file mode 100644 index 0000000000000..92da4fcf88744 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/ElectLeadersResult.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.clients.admin; + + +import java.util.Map; +import java.util.Optional; +import org.apache.kafka.common.KafkaFuture; +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.annotation.InterfaceStability; +import org.apache.kafka.common.internals.KafkaFutureImpl; + +/** + * The result of {@link Admin#electLeaders(ElectionType, Set, ElectLeadersOptions)} + * + * The API of this class is evolving, see {@link Admin} for details. + */ +@InterfaceStability.Evolving +final public class ElectLeadersResult { + private final KafkaFutureImpl>> electionFuture; + + ElectLeadersResult(KafkaFutureImpl>> electionFuture) { + this.electionFuture = electionFuture; + } + + /** + *

Get a future for the topic partitions for which a leader election was attempted. + * If the election succeeded then the value for a topic partition will be the empty Optional. + * Otherwise the election failed and the Optional will be set with the error.

+ */ + public KafkaFuture>> partitions() { + return electionFuture; + } + + /** + * Return a future which succeeds if all the topic elections succeed. + */ + public KafkaFuture all() { + final KafkaFutureImpl result = new KafkaFutureImpl<>(); + + partitions().whenComplete( + new KafkaFuture.BiConsumer>, Throwable>() { + @Override + public void accept(Map> topicPartitions, Throwable throwable) { + if (throwable != null) { + result.completeExceptionally(throwable); + } else { + for (Optional exception : topicPartitions.values()) { + if (exception.isPresent()) { + result.completeExceptionally(exception.get()); + return; + } + } + result.complete(null); + } + } + }); + + return result; + } +} diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/ElectPreferredLeadersOptions.java b/clients/src/main/java/org/apache/kafka/clients/admin/ElectPreferredLeadersOptions.java index 80b00975f7da9..c00c92098b7b2 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/ElectPreferredLeadersOptions.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/ElectPreferredLeadersOptions.java @@ -18,14 +18,16 @@ package org.apache.kafka.clients.admin; import org.apache.kafka.common.annotation.InterfaceStability; - import java.util.Collection; /** - * Options for {@link AdminClient#electPreferredLeaders(Collection, ElectPreferredLeadersOptions)}. + * Options for {@link Admin#electPreferredLeaders(Collection, ElectPreferredLeadersOptions)}. + * + * The API of this class is evolving, see {@link Admin} for details. * - * The API of this class is evolving, see {@link AdminClient} for details. + * @deprecated Since 2.4.0. Use {@link Admin#electLeaders(ElectionType, Set, ElectLeadersOptions)}. */ @InterfaceStability.Evolving +@Deprecated public class ElectPreferredLeadersOptions extends AbstractOptions { } diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/ElectPreferredLeadersResult.java b/clients/src/main/java/org/apache/kafka/clients/admin/ElectPreferredLeadersResult.java index c76336a7cbc0a..0c851931accc7 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/ElectPreferredLeadersResult.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/ElectPreferredLeadersResult.java @@ -17,33 +17,31 @@ package org.apache.kafka.clients.admin; + +import java.util.Collection; +import java.util.Map; +import java.util.Optional; +import java.util.Set; import org.apache.kafka.common.KafkaFuture; import org.apache.kafka.common.TopicPartition; import org.apache.kafka.common.annotation.InterfaceStability; -import org.apache.kafka.common.errors.ApiException; import org.apache.kafka.common.errors.UnknownTopicOrPartitionException; import org.apache.kafka.common.internals.KafkaFutureImpl; -import org.apache.kafka.common.protocol.Errors; -import org.apache.kafka.common.requests.ApiError; - -import java.util.Collection; -import java.util.Map; -import java.util.Set; /** - * The result of {@link AdminClient#electPreferredLeaders(Collection, ElectPreferredLeadersOptions)} + * The result of {@link Admin#electPreferredLeaders(Collection, ElectPreferredLeadersOptions)} + * + * The API of this class is evolving, see {@link Admin} for details. * - * The API of this class is evolving, see {@link AdminClient} for details. + * @deprecated Since 2.4.0. Use {@link Admin#electLeaders(ElectionType, Set, ElectLeadersOptions)}. */ @InterfaceStability.Evolving +@Deprecated public class ElectPreferredLeadersResult { + private final ElectLeadersResult electionResult; - private final KafkaFutureImpl> electionFuture; - private final Set partitions; - - ElectPreferredLeadersResult(KafkaFutureImpl> electionFuture, Set partitions) { - this.electionFuture = electionFuture; - this.partitions = partitions; + ElectPreferredLeadersResult(ElectLeadersResult electionResult) { + this.electionResult = electionResult; } /** @@ -53,28 +51,28 @@ public class ElectPreferredLeadersResult { */ public KafkaFuture partitionResult(final TopicPartition partition) { final KafkaFutureImpl result = new KafkaFutureImpl<>(); - electionFuture.whenComplete(new KafkaFuture.BiConsumer, Throwable>() { - @Override - public void accept(Map topicPartitions, Throwable throwable) { - if (throwable != null) { - result.completeExceptionally(throwable); - } else if (!topicPartitions.containsKey(partition)) { - result.completeExceptionally(new UnknownTopicOrPartitionException( - "Preferred leader election for partition \"" + partition + - "\" was not attempted")); - } else { - if (partitions == null && topicPartitions.isEmpty()) { - result.completeExceptionally(Errors.CLUSTER_AUTHORIZATION_FAILED.exception()); - } - ApiException exception = topicPartitions.get(partition).exception(); - if (exception == null) { - result.complete(null); - } else { - result.completeExceptionally(exception); + + electionResult.partitions().whenComplete( + new KafkaFuture.BiConsumer>, Throwable>() { + @Override + public void accept(Map> topicPartitions, Throwable throwable) { + if (throwable != null) { + result.completeExceptionally(throwable); + } else if (!topicPartitions.containsKey(partition)) { + result.completeExceptionally(new UnknownTopicOrPartitionException( + "Preferred leader election for partition \"" + partition + + "\" was not attempted")); + } else { + Optional exception = topicPartitions.get(partition); + if (exception.isPresent()) { + result.completeExceptionally(exception.get()); + } else { + result.complete(null); + } + } } - } - } - }); + }); + return result; } @@ -84,53 +82,31 @@ public void accept(Map topicPartitions, Throwable thro * an election was attempted even if the election was not successful.

* *

This method is provided to discover the partitions attempted when - * {@link AdminClient#electPreferredLeaders(Collection)} is called + * {@link Admin#electPreferredLeaders(Collection)} is called * with a null {@code partitions} argument.

*/ public KafkaFuture> partitions() { - if (partitions != null) { - return KafkaFutureImpl.completedFuture(this.partitions); - } else { - final KafkaFutureImpl> result = new KafkaFutureImpl<>(); - electionFuture.whenComplete(new KafkaFuture.BiConsumer, Throwable>() { - @Override - public void accept(Map topicPartitions, Throwable throwable) { - if (throwable != null) { - result.completeExceptionally(throwable); - } else if (topicPartitions.isEmpty()) { - result.completeExceptionally(Errors.CLUSTER_AUTHORIZATION_FAILED.exception()); - } else { - for (ApiError apiError : topicPartitions.values()) { - if (apiError.isFailure()) { - result.completeExceptionally(apiError.exception()); - } + final KafkaFutureImpl> result = new KafkaFutureImpl<>(); + + electionResult.partitions().whenComplete( + new KafkaFuture.BiConsumer>, Throwable>() { + @Override + public void accept(Map> topicPartitions, Throwable throwable) { + if (throwable != null) { + result.completeExceptionally(throwable); + } else { + result.complete(topicPartitions.keySet()); } - result.complete(topicPartitions.keySet()); } - } - }); - return result; - } + }); + + return result; } /** * Return a future which succeeds if all the topic elections succeed. */ public KafkaFuture all() { - final KafkaFutureImpl result = new KafkaFutureImpl<>(); - electionFuture.thenApply(new KafkaFuture.Function, Void>() { - @Override - public Void apply(Map topicPartitions) { - for (ApiError apiError : topicPartitions.values()) { - if (apiError.isFailure()) { - result.completeExceptionally(apiError.exception()); - return null; - } - } - result.complete(null); - return null; - } - }); - return result; + return electionResult.all(); } } diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/ExpireDelegationTokenOptions.java b/clients/src/main/java/org/apache/kafka/clients/admin/ExpireDelegationTokenOptions.java index 138cd4e69e477..3bf9489a05e32 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/ExpireDelegationTokenOptions.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/ExpireDelegationTokenOptions.java @@ -20,9 +20,9 @@ import org.apache.kafka.common.annotation.InterfaceStability; /** - * Options for {@link AdminClient#expireDelegationToken(byte[], ExpireDelegationTokenOptions)}. + * Options for {@link Admin#expireDelegationToken(byte[], ExpireDelegationTokenOptions)}. * - * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class ExpireDelegationTokenOptions extends AbstractOptions { diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/ExpireDelegationTokenResult.java b/clients/src/main/java/org/apache/kafka/clients/admin/ExpireDelegationTokenResult.java index 41782bdcb5ce7..59b17140ae791 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/ExpireDelegationTokenResult.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/ExpireDelegationTokenResult.java @@ -23,7 +23,7 @@ /** * The result of the {@link KafkaAdminClient#expireDelegationToken(byte[], ExpireDelegationTokenOptions)} call. * - * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class ExpireDelegationTokenResult { 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 95337d093553f..26040db5acc20 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 @@ -25,15 +25,17 @@ import org.apache.kafka.clients.KafkaClient; import org.apache.kafka.clients.NetworkClient; import org.apache.kafka.clients.StaleMetadataException; +import org.apache.kafka.clients.admin.CreateTopicsResult.TopicMetadataAndConfig; import org.apache.kafka.clients.admin.DeleteAclsResult.FilterResult; import org.apache.kafka.clients.admin.DeleteAclsResult.FilterResults; import org.apache.kafka.clients.admin.DescribeReplicaLogDirsResult.ReplicaLogDirInfo; import org.apache.kafka.clients.admin.internals.AdminMetadataManager; +import org.apache.kafka.clients.consumer.ConsumerPartitionAssignor.Assignment; import org.apache.kafka.clients.consumer.OffsetAndMetadata; import org.apache.kafka.clients.consumer.internals.ConsumerProtocol; -import org.apache.kafka.clients.consumer.internals.PartitionAssignor; import org.apache.kafka.common.Cluster; import org.apache.kafka.common.ConsumerGroupState; +import org.apache.kafka.common.ElectionType; import org.apache.kafka.common.KafkaException; import org.apache.kafka.common.KafkaFuture; import org.apache.kafka.common.Metric; @@ -60,12 +62,38 @@ import org.apache.kafka.common.errors.UnknownTopicOrPartitionException; import org.apache.kafka.common.errors.UnsupportedVersionException; import org.apache.kafka.common.internals.KafkaFutureImpl; +import org.apache.kafka.common.message.AlterPartitionReassignmentsRequestData; +import org.apache.kafka.common.message.AlterPartitionReassignmentsRequestData.ReassignableTopic; +import org.apache.kafka.common.message.CreateDelegationTokenRequestData; +import org.apache.kafka.common.message.CreateDelegationTokenRequestData.CreatableRenewers; +import org.apache.kafka.common.message.CreateDelegationTokenResponseData; import org.apache.kafka.common.message.CreateTopicsRequestData; -import org.apache.kafka.common.message.CreateTopicsRequestData.CreatableTopicSet; +import org.apache.kafka.common.message.CreateTopicsRequestData.CreatableTopicCollection; +import org.apache.kafka.common.message.CreateTopicsResponseData.CreatableTopicConfigs; import org.apache.kafka.common.message.CreateTopicsResponseData.CreatableTopicResult; +import org.apache.kafka.common.message.DeleteGroupsRequestData; +import org.apache.kafka.common.message.DeleteTopicsRequestData; +import org.apache.kafka.common.message.DeleteTopicsResponseData.DeletableTopicResult; import org.apache.kafka.common.message.DescribeGroupsRequestData; import org.apache.kafka.common.message.DescribeGroupsResponseData.DescribedGroup; import org.apache.kafka.common.message.DescribeGroupsResponseData.DescribedGroupMember; +import org.apache.kafka.common.message.ExpireDelegationTokenRequestData; +import org.apache.kafka.common.message.FindCoordinatorRequestData; +import org.apache.kafka.common.message.IncrementalAlterConfigsRequestData; +import org.apache.kafka.common.message.IncrementalAlterConfigsRequestData.AlterConfigsResource; +import org.apache.kafka.common.message.IncrementalAlterConfigsRequestData.AlterableConfig; +import org.apache.kafka.common.message.IncrementalAlterConfigsRequestData.AlterableConfigCollection; +import org.apache.kafka.common.message.LeaveGroupRequestData.MemberIdentity; +import org.apache.kafka.common.message.LeaveGroupResponseData.MemberResponse; +import org.apache.kafka.common.message.ListGroupsRequestData; +import org.apache.kafka.common.message.ListGroupsResponseData; +import org.apache.kafka.common.message.ListPartitionReassignmentsRequestData; +import org.apache.kafka.common.message.MetadataRequestData; +import org.apache.kafka.common.message.OffsetDeleteRequestData; +import org.apache.kafka.common.message.OffsetDeleteRequestData.OffsetDeleteRequestPartition; +import org.apache.kafka.common.message.OffsetDeleteRequestData.OffsetDeleteRequestTopic; +import org.apache.kafka.common.message.OffsetDeleteRequestData.OffsetDeleteRequestTopicCollection; +import org.apache.kafka.common.message.RenewDelegationTokenRequestData; import org.apache.kafka.common.metrics.JmxReporter; import org.apache.kafka.common.metrics.MetricConfig; import org.apache.kafka.common.metrics.Metrics; @@ -78,6 +106,8 @@ import org.apache.kafka.common.requests.AbstractResponse; import org.apache.kafka.common.requests.AlterConfigsRequest; import org.apache.kafka.common.requests.AlterConfigsResponse; +import org.apache.kafka.common.requests.AlterPartitionReassignmentsRequest; +import org.apache.kafka.common.requests.AlterPartitionReassignmentsResponse; import org.apache.kafka.common.requests.AlterReplicaLogDirsRequest; import org.apache.kafka.common.requests.AlterReplicaLogDirsResponse; import org.apache.kafka.common.requests.ApiError; @@ -112,20 +142,31 @@ import org.apache.kafka.common.requests.DescribeGroupsResponse; import org.apache.kafka.common.requests.DescribeLogDirsRequest; import org.apache.kafka.common.requests.DescribeLogDirsResponse; -import org.apache.kafka.common.requests.ElectPreferredLeadersRequest; -import org.apache.kafka.common.requests.ElectPreferredLeadersResponse; +import org.apache.kafka.common.requests.ElectLeadersRequest; +import org.apache.kafka.common.requests.ElectLeadersResponse; import org.apache.kafka.common.requests.ExpireDelegationTokenRequest; import org.apache.kafka.common.requests.ExpireDelegationTokenResponse; import org.apache.kafka.common.requests.FindCoordinatorRequest; +import org.apache.kafka.common.requests.FindCoordinatorRequest.CoordinatorType; import org.apache.kafka.common.requests.FindCoordinatorResponse; +import org.apache.kafka.common.requests.IncrementalAlterConfigsRequest; +import org.apache.kafka.common.requests.IncrementalAlterConfigsResponse; +import org.apache.kafka.common.requests.JoinGroupRequest; +import org.apache.kafka.common.requests.LeaveGroupRequest; +import org.apache.kafka.common.requests.LeaveGroupResponse; import org.apache.kafka.common.requests.ListGroupsRequest; import org.apache.kafka.common.requests.ListGroupsResponse; +import org.apache.kafka.common.requests.ListPartitionReassignmentsRequest; +import org.apache.kafka.common.requests.ListPartitionReassignmentsResponse; import org.apache.kafka.common.requests.MetadataRequest; import org.apache.kafka.common.requests.MetadataResponse; +import org.apache.kafka.common.requests.OffsetDeleteRequest; +import org.apache.kafka.common.requests.OffsetDeleteResponse; import org.apache.kafka.common.requests.OffsetFetchRequest; import org.apache.kafka.common.requests.OffsetFetchResponse; import org.apache.kafka.common.requests.RenewDelegationTokenRequest; import org.apache.kafka.common.requests.RenewDelegationTokenResponse; +import org.apache.kafka.common.security.auth.KafkaPrincipal; import org.apache.kafka.common.security.token.delegation.DelegationToken; import org.apache.kafka.common.security.token.delegation.TokenInformation; import org.apache.kafka.common.utils.AppInfoParser; @@ -149,21 +190,32 @@ import java.util.LinkedList; import java.util.List; import java.util.Map; +import java.util.Objects; import java.util.Optional; import java.util.Set; +import java.util.TreeMap; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicLong; import java.util.function.Predicate; +import java.util.function.Supplier; import java.util.stream.Collectors; - +import java.util.stream.Stream; + +import static org.apache.kafka.common.message.AlterPartitionReassignmentsRequestData.ReassignablePartition; +import static org.apache.kafka.common.message.AlterPartitionReassignmentsResponseData.ReassignablePartitionResponse; +import static org.apache.kafka.common.message.AlterPartitionReassignmentsResponseData.ReassignableTopicResponse; +import static org.apache.kafka.common.message.ListPartitionReassignmentsRequestData.ListPartitionReassignmentsTopics; +import static org.apache.kafka.common.message.ListPartitionReassignmentsResponseData.OngoingPartitionReassignment; +import static org.apache.kafka.common.message.ListPartitionReassignmentsResponseData.OngoingTopicReassignment; +import static org.apache.kafka.common.requests.MetadataRequest.convertToMetadataRequestTopic; import static org.apache.kafka.common.utils.Utils.closeQuietly; /** - * The default implementation of {@link AdminClient}. An instance of this class is created by invoking one of the + * The default implementation of {@link Admin}. An instance of this class is created by invoking one of the * {@code create()} methods in {@code AdminClient}. Users should not refer to this class directly. * - * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class KafkaAdminClient extends AdminClient { @@ -255,12 +307,7 @@ public class KafkaAdminClient extends AdminClient { * @return The list value. */ static List getOrCreateListValue(Map> map, K key) { - List list = map.get(key); - if (list != null) - return list; - list = new LinkedList<>(); - map.put(key, list); - return list; + return map.computeIfAbsent(key, k -> new LinkedList<>()); } /** @@ -271,9 +318,18 @@ static List getOrCreateListValue(Map> map, K key) { * @param The KafkaFutureImpl result type. */ private static void completeAllExceptionally(Collection> futures, Throwable exc) { - for (KafkaFutureImpl future : futures) { - future.completeExceptionally(exc); - } + completeAllExceptionally(futures.stream(), exc); + } + + /** + * Send an exception to all futures in the provided stream + * + * @param futures The stream of KafkaFutureImpl objects. + * @param exc The exception + * @param The KafkaFutureImpl result type. + */ + private static void completeAllExceptionally(Stream> futures, Throwable exc) { + futures.forEach(future -> future.completeExceptionally(exc)); } /** @@ -366,6 +422,7 @@ static KafkaAdminClient createInternal(AdminClientConfig config, TimeoutProcesso .tags(metricTags); reporters.add(new JmxReporter(JMX_PREFIX)); metrics = new Metrics(metricConfig, reporters, time); + metrics.setReplaceOnDuplicateMetric(config.getBoolean(AdminClientConfig.METRICS_REPLACE_ON_DUPLICATE_CONFIG)); String metricGrpPrefix = "admin-client"; channelBuilder = ClientUtils.createChannelBuilder(config, time); selector = new Selector(config.getLong(AdminClientConfig.CONNECTIONS_MAX_IDLE_MS_CONFIG), @@ -441,7 +498,7 @@ private KafkaAdminClient(AdminClientConfig config, this.maxRetries = config.getInt(AdminClientConfig.RETRIES_CONFIG); this.retryBackoffMs = config.getLong(AdminClientConfig.RETRY_BACKOFF_MS_CONFIG); config.logUnused(); - AppInfoParser.registerAppInfo(JMX_PREFIX, clientId, metrics); + AppInfoParser.registerAppInfo(JMX_PREFIX, clientId, metrics, time.milliseconds()); log.debug("Kafka admin client initialized"); thread.start(); } @@ -1222,7 +1279,9 @@ public AbstractRequest.Builder createRequest(int timeoutMs) { // Since this only requests node information, it's safe to pass true // for allowAutoTopicCreation (and it simplifies communication with // older brokers) - return new MetadataRequest.Builder(Collections.emptyList(), true); + return new MetadataRequest.Builder(new MetadataRequestData() + .setTopics(Collections.emptyList()) + .setAllowAutoTopicCreation(true)); } @Override @@ -1257,14 +1316,19 @@ private static boolean groupIdIsUnrepresentable(String groupId) { return groupId == null; } + //for testing + int numPendingCalls() { + return runnable.pendingCalls.size(); + } + @Override public CreateTopicsResult createTopics(final Collection newTopics, final CreateTopicsOptions options) { - final Map> topicFutures = new HashMap<>(newTopics.size()); - final CreatableTopicSet topics = new CreatableTopicSet(); + final Map> topicFutures = new HashMap<>(newTopics.size()); + final CreatableTopicCollection topics = new CreatableTopicCollection(); for (NewTopic newTopic : newTopics) { if (topicNameIsUnrepresentable(newTopic.name())) { - KafkaFutureImpl future = new KafkaFutureImpl<>(); + KafkaFutureImpl future = new KafkaFutureImpl<>(); future.completeExceptionally(new InvalidTopicException("The given topic name '" + newTopic.name() + "' cannot be represented in a request.")); topicFutures.put(newTopic.name(), future); @@ -1299,7 +1363,7 @@ public void handleResponse(AbstractResponse abstractResponse) { } // Handle server responses for particular topics. for (CreatableTopicResult result : response.data().topics()) { - KafkaFutureImpl future = topicFutures.get(result.name()); + KafkaFutureImpl future = topicFutures.get(result.name()); if (future == null) { log.warn("Server response mentioned unknown topic {}", result.name()); } else { @@ -1309,13 +1373,33 @@ public void handleResponse(AbstractResponse abstractResponse) { if (exception != null) { future.completeExceptionally(exception); } else { - future.complete(null); + TopicMetadataAndConfig topicMetadataAndConfig; + if (result.topicConfigErrorCode() != Errors.NONE.code()) { + topicMetadataAndConfig = new TopicMetadataAndConfig(Errors.forCode(result.topicConfigErrorCode()).exception()); + } else if (result.numPartitions() == CreateTopicsResult.UNKNOWN) { + topicMetadataAndConfig = new TopicMetadataAndConfig(new UnsupportedVersionException( + "Topic metadata and configs in CreateTopics response not supported")); + } else { + List configs = result.configs(); + Config topicConfig = new Config(configs.stream() + .map(config -> new ConfigEntry(config.name(), + config.value(), + configSource(DescribeConfigsResponse.ConfigSource.forId(config.configSource())), + config.isSensitive(), + config.readOnly(), + Collections.emptyList())) + .collect(Collectors.toSet())); + topicMetadataAndConfig = new TopicMetadataAndConfig(result.numPartitions(), + result.replicationFactor(), + topicConfig); + } + future.complete(topicMetadataAndConfig); } } } // The server should send back a response for every topic. But do a sanity check anyway. - for (Map.Entry> entry : topicFutures.entrySet()) { - KafkaFutureImpl future = entry.getValue(); + for (Map.Entry> entry : topicFutures.entrySet()) { + KafkaFutureImpl future = entry.getValue(); if (!future.isDone()) { future.completeExceptionally(new ApiException("The server response did not " + "contain a reference to node " + entry.getKey())); @@ -1356,14 +1440,16 @@ public DeleteTopicsResult deleteTopics(Collection topicNames, @Override AbstractRequest.Builder createRequest(int timeoutMs) { - return new DeleteTopicsRequest.Builder(new HashSet<>(validTopicNames), timeoutMs); + return new DeleteTopicsRequest.Builder(new DeleteTopicsRequestData() + .setTopicNames(validTopicNames) + .setTimeoutMs(timeoutMs)); } @Override void handleResponse(AbstractResponse abstractResponse) { DeleteTopicsResponse response = (DeleteTopicsResponse) abstractResponse; // Check for controller change - for (Errors error : response.errors().values()) { + for (Errors error : response.errorCounts().keySet()) { if (error == Errors.NOT_CONTROLLER) { metadataManager.clearController(); metadataManager.requestUpdate(); @@ -1371,12 +1457,12 @@ void handleResponse(AbstractResponse abstractResponse) { } } // Handle server responses for particular topics. - for (Map.Entry entry : response.errors().entrySet()) { - KafkaFutureImpl future = topicFutures.get(entry.getKey()); + for (DeletableTopicResult result : response.data().responses()) { + KafkaFutureImpl future = topicFutures.get(result.name()); if (future == null) { - log.warn("Server response mentioned unknown topic {}", entry.getKey()); + log.warn("Server response mentioned unknown topic {}", result.name()); } else { - ApiException exception = entry.getValue().exception(); + ApiException exception = Errors.forCode(result.errorCode()).exception(); if (exception != null) { future.completeExceptionally(exception); } else { @@ -1455,14 +1541,17 @@ public DescribeTopicsResult describeTopics(final Collection topicNames, } final long now = time.milliseconds(); Call call = new Call("describeTopics", calcDeadlineMs(now, options.timeoutMs()), - new ControllerNodeProvider()) { + new LeastLoadedNodeProvider()) { private boolean supportsDisablingTopicCreation = true; @Override AbstractRequest.Builder createRequest(int timeoutMs) { if (supportsDisablingTopicCreation) - return new MetadataRequest.Builder(topicNamesList, false); + return new MetadataRequest.Builder(new MetadataRequestData() + .setTopics(convertToMetadataRequestTopic(topicNamesList)) + .setAllowAutoTopicCreation(false) + .setIncludeTopicAuthorizedOperations(options.includeAuthorizedOperations())); else return MetadataRequest.Builder.allTopics(); } @@ -1495,7 +1584,8 @@ void handleResponse(AbstractResponse abstractResponse) { partitions.add(topicPartitionInfo); } partitions.sort(Comparator.comparingInt(TopicPartitionInfo::partition)); - TopicDescription topicDescription = new TopicDescription(topicName, isInternal, partitions); + TopicDescription topicDescription = new TopicDescription(topicName, isInternal, partitions, + validAclOperations(response.topicAuthorizedOperations(topicName).get())); future.complete(topicDescription); } } @@ -1531,6 +1621,8 @@ public DescribeClusterResult describeCluster(DescribeClusterOptions options) { final KafkaFutureImpl> describeClusterFuture = new KafkaFutureImpl<>(); final KafkaFutureImpl controllerFuture = new KafkaFutureImpl<>(); final KafkaFutureImpl clusterIdFuture = new KafkaFutureImpl<>(); + final KafkaFutureImpl> authorizedOperationsFuture = new KafkaFutureImpl<>(); + final long now = time.milliseconds(); runnable.call(new Call("listNodes", calcDeadlineMs(now, options.timeoutMs()), new LeastLoadedNodeProvider()) { @@ -1539,7 +1631,10 @@ public DescribeClusterResult describeCluster(DescribeClusterOptions options) { AbstractRequest.Builder createRequest(int timeoutMs) { // Since this only requests node information, it's safe to pass true for allowAutoTopicCreation (and it // simplifies communication with older brokers) - return new MetadataRequest.Builder(Collections.emptyList(), true); + return new MetadataRequest.Builder(new MetadataRequestData() + .setTopics(Collections.emptyList()) + .setAllowAutoTopicCreation(true) + .setIncludeClusterAuthorizedOperations(options.includeAuthorizedOperations())); } @Override @@ -1548,6 +1643,8 @@ void handleResponse(AbstractResponse abstractResponse) { describeClusterFuture.complete(response.brokers()); controllerFuture.complete(controller(response)); clusterIdFuture.complete(response.clusterId()); + authorizedOperationsFuture.complete( + validAclOperations(response.clusterAuthorizedOperations())); } private Node controller(MetadataResponse response) { @@ -1561,10 +1658,12 @@ void handleFailure(Throwable throwable) { describeClusterFuture.completeExceptionally(throwable); controllerFuture.completeExceptionally(throwable); clusterIdFuture.completeExceptionally(throwable); + authorizedOperationsFuture.completeExceptionally(throwable); } }, now); - return new DescribeClusterResult(describeClusterFuture, controllerFuture, clusterIdFuture); + return new DescribeClusterResult(describeClusterFuture, controllerFuture, clusterIdFuture, + authorizedOperationsFuture); } @Override @@ -1724,7 +1823,7 @@ public DescribeConfigsResult describeConfigs(Collection configRe final Collection unifiedRequestResources = new ArrayList<>(configResources.size()); for (ConfigResource resource : configResources) { - if (resource.type() == ConfigResource.Type.BROKER && !resource.isDefault()) { + if (dependsOnSpecificNode(resource)) { brokerFutures.put(resource, new KafkaFutureImpl<>()); brokerResources.add(resource); } else { @@ -1849,6 +1948,9 @@ private ConfigEntry.ConfigSource configSource(DescribeConfigsResponse.ConfigSour case STATIC_BROKER_CONFIG: configSource = ConfigEntry.ConfigSource.STATIC_BROKER_CONFIG; break; + case DYNAMIC_BROKER_LOGGER_CONFIG: + configSource = ConfigEntry.ConfigSource.DYNAMIC_BROKER_LOGGER_CONFIG; + break; case DEFAULT_CONFIG: configSource = ConfigEntry.ConfigSource.DEFAULT_CONFIG; break; @@ -1859,6 +1961,7 @@ private ConfigEntry.ConfigSource configSource(DescribeConfigsResponse.ConfigSour } @Override + @Deprecated public AlterConfigsResult alterConfigs(Map configs, final AlterConfigsOptions options) { final Map> allFutures = new HashMap<>(); // We must make a separate AlterConfigs request for every BROKER resource we want to alter @@ -1867,7 +1970,7 @@ public AlterConfigsResult alterConfigs(Map configs, fina final Collection unifiedRequestResources = new ArrayList<>(); for (ConfigResource resource : configs.keySet()) { - if (resource.type() == ConfigResource.Type.BROKER && !resource.isDefault()) { + if (dependsOnSpecificNode(resource)) { NodeProvider nodeProvider = new ConstantNodeIdProvider(Integer.parseInt(resource.name())); allFutures.putAll(alterConfigs(configs, options, Collections.singleton(resource), nodeProvider)); } else @@ -1922,6 +2025,89 @@ void handleFailure(Throwable throwable) { return futures; } + @Override + public AlterConfigsResult incrementalAlterConfigs(Map> configs, + final AlterConfigsOptions options) { + final Map> allFutures = new HashMap<>(); + // We must make a separate AlterConfigs request for every BROKER resource we want to alter + // and send the request to that specific broker. Other resources are grouped together into + // a single request that may be sent to any broker. + final Collection unifiedRequestResources = new ArrayList<>(); + + for (ConfigResource resource : configs.keySet()) { + if (dependsOnSpecificNode(resource)) { + NodeProvider nodeProvider = new ConstantNodeIdProvider(Integer.parseInt(resource.name())); + allFutures.putAll(incrementalAlterConfigs(configs, options, Collections.singleton(resource), nodeProvider)); + } else + unifiedRequestResources.add(resource); + } + if (!unifiedRequestResources.isEmpty()) + allFutures.putAll(incrementalAlterConfigs(configs, options, unifiedRequestResources, new LeastLoadedNodeProvider())); + + return new AlterConfigsResult(new HashMap<>(allFutures)); + } + + private Map> incrementalAlterConfigs(Map> configs, + final AlterConfigsOptions options, + Collection resources, + NodeProvider nodeProvider) { + final Map> futures = new HashMap<>(); + for (ConfigResource resource : resources) + futures.put(resource, new KafkaFutureImpl<>()); + + final long now = time.milliseconds(); + runnable.call(new Call("incrementalAlterConfigs", calcDeadlineMs(now, options.timeoutMs()), nodeProvider) { + + @Override + public AbstractRequest.Builder createRequest(int timeoutMs) { + return new IncrementalAlterConfigsRequest.Builder( + toIncrementalAlterConfigsRequestData(resources, configs, options.shouldValidateOnly())); + } + + @Override + public void handleResponse(AbstractResponse abstractResponse) { + IncrementalAlterConfigsResponse response = (IncrementalAlterConfigsResponse) abstractResponse; + Map errors = IncrementalAlterConfigsResponse.fromResponseData(response.data()); + for (Map.Entry> entry : futures.entrySet()) { + KafkaFutureImpl future = entry.getValue(); + ApiException exception = errors.get(entry.getKey()).exception(); + if (exception != null) { + future.completeExceptionally(exception); + } else { + future.complete(null); + } + } + } + + @Override + void handleFailure(Throwable throwable) { + completeAllExceptionally(futures.values(), throwable); + } + }, now); + return futures; + } + + private IncrementalAlterConfigsRequestData toIncrementalAlterConfigsRequestData(final Collection resources, + final Map> configs, + final boolean validateOnly) { + IncrementalAlterConfigsRequestData requestData = new IncrementalAlterConfigsRequestData(); + requestData.setValidateOnly(validateOnly); + for (ConfigResource resource : resources) { + AlterableConfigCollection alterableConfigSet = new AlterableConfigCollection(); + for (AlterConfigOp configEntry : configs.get(resource)) + alterableConfigSet.add(new AlterableConfig(). + setName(configEntry.configEntry().name()). + setValue(configEntry.configEntry().value()). + setConfigOperation(configEntry.opType().id())); + + AlterConfigsResource alterConfigsResource = new AlterConfigsResource(); + alterConfigsResource.setResourceType(resource.type().id()). + setResourceName(resource.name()).setConfigs(alterableConfigSet); + requestData.resources().add(alterConfigsResource); + } + return requestData; + } + @Override public AlterReplicaLogDirsResult alterReplicaLogDirs(Map replicaAssignment, final AlterReplicaLogDirsOptions options) { final Map> futures = new HashMap<>(replicaAssignment.size()); @@ -2179,7 +2365,9 @@ public DeleteRecordsResult deleteRecords(final Map(topics), false); + return new MetadataRequest.Builder(new MetadataRequestData() + .setTopics(convertToMetadataRequestTopic(topics)) + .setAllowAutoTopicCreation(false)); } @Override @@ -2189,39 +2377,31 @@ void handleResponse(AbstractResponse abstractResponse) { Map errors = response.errors(); Cluster cluster = response.cluster(); - // completing futures for topics with errors - for (Map.Entry topicError: errors.entrySet()) { - - for (Map.Entry> future: futures.entrySet()) { - if (future.getKey().topic().equals(topicError.getKey())) { - future.getValue().completeExceptionally(topicError.getValue().exception()); - } - } - } - - // grouping topic partitions per leader + // Group topic partitions by leader Map> leaders = new HashMap<>(); for (Map.Entry entry: recordsToDelete.entrySet()) { + KafkaFutureImpl future = futures.get(entry.getKey()); - // avoiding to send deletion request for topics with errors - if (!errors.containsKey(entry.getKey().topic())) { - + // Fail partitions with topic errors + Errors topicError = errors.get(entry.getKey().topic()); + if (errors.containsKey(entry.getKey().topic())) { + future.completeExceptionally(topicError.exception()); + } else { Node node = cluster.leaderFor(entry.getKey()); if (node != null) { if (!leaders.containsKey(node)) leaders.put(node, new HashMap<>()); leaders.get(node).put(entry.getKey(), entry.getValue().beforeOffset()); } else { - KafkaFutureImpl future = futures.get(entry.getKey()); future.completeExceptionally(Errors.LEADER_NOT_AVAILABLE.exception()); } } } - for (final Map.Entry> entry: leaders.entrySet()) { - - final long nowDelete = time.milliseconds(); + final long deleteRecordsCallTimeMs = time.milliseconds(); + for (final Map.Entry> entry : leaders.entrySet()) { + final Map partitionDeleteOffsets = entry.getValue(); final int brokerId = entry.getKey().id(); runnable.call(new Call("deleteRecords", deadline, @@ -2229,7 +2409,7 @@ void handleResponse(AbstractResponse abstractResponse) { @Override AbstractRequest.Builder createRequest(int timeoutMs) { - return new DeleteRecordsRequest.Builder(timeoutMs, entry.getValue()); + return new DeleteRecordsRequest.Builder(timeoutMs, partitionDeleteOffsets); } @Override @@ -2248,9 +2428,11 @@ void handleResponse(AbstractResponse abstractResponse) { @Override void handleFailure(Throwable throwable) { - completeAllExceptionally(futures.values(), throwable); + Stream> callFutures = + partitionDeleteOffsets.keySet().stream().map(futures::get); + completeAllExceptionally(callFutures, throwable); } - }, nowDelete); + }, deleteRecordsCallTimeMs); } } @@ -2267,12 +2449,21 @@ void handleFailure(Throwable throwable) { public CreateDelegationTokenResult createDelegationToken(final CreateDelegationTokenOptions options) { final KafkaFutureImpl delegationTokenFuture = new KafkaFutureImpl<>(); final long now = time.milliseconds(); + List renewers = new ArrayList<>(); + for (KafkaPrincipal principal : options.renewers()) { + renewers.add(new CreatableRenewers() + .setPrincipalName(principal.getName()) + .setPrincipalType(principal.getPrincipalType())); + } runnable.call(new Call("createDelegationToken", calcDeadlineMs(now, options.timeoutMs()), new LeastLoadedNodeProvider()) { @Override - AbstractRequest.Builder createRequest(int timeoutMs) { - return new CreateDelegationTokenRequest.Builder(options.renewers(), options.maxlifeTimeMs()); + AbstractRequest.Builder createRequest(int timeoutMs) { + return new CreateDelegationTokenRequest.Builder( + new CreateDelegationTokenRequestData() + .setRenewers(renewers) + .setMaxLifetimeMs(options.maxlifeTimeMs())); } @Override @@ -2281,9 +2472,10 @@ void handleResponse(AbstractResponse abstractResponse) { if (response.hasError()) { delegationTokenFuture.completeExceptionally(response.error().exception()); } else { - TokenInformation tokenInfo = new TokenInformation(response.tokenId(), response.owner(), - options.renewers(), response.issueTimestamp(), response.maxTimestamp(), response.expiryTimestamp()); - DelegationToken token = new DelegationToken(tokenInfo, response.hmacBytes()); + CreateDelegationTokenResponseData data = response.data(); + TokenInformation tokenInfo = new TokenInformation(data.tokenId(), new KafkaPrincipal(data.principalType(), data.principalName()), + options.renewers(), data.issueTimestampMs(), data.maxTimestampMs(), data.expiryTimestampMs()); + DelegationToken token = new DelegationToken(tokenInfo, data.hmac()); delegationTokenFuture.complete(token); } } @@ -2305,8 +2497,11 @@ public RenewDelegationTokenResult renewDelegationToken(final byte[] hmac, final new LeastLoadedNodeProvider()) { @Override - AbstractRequest.Builder createRequest(int timeoutMs) { - return new RenewDelegationTokenRequest.Builder(hmac, options.renewTimePeriodMs()); + AbstractRequest.Builder createRequest(int timeoutMs) { + return new RenewDelegationTokenRequest.Builder( + new RenewDelegationTokenRequestData() + .setHmac(hmac) + .setRenewPeriodMs(options.renewTimePeriodMs())); } @Override @@ -2336,8 +2531,11 @@ public ExpireDelegationTokenResult expireDelegationToken(final byte[] hmac, fina new LeastLoadedNodeProvider()) { @Override - AbstractRequest.Builder createRequest(int timeoutMs) { - return new ExpireDelegationTokenRequest.Builder(hmac, options.expiryTimePeriodMs()); + AbstractRequest.Builder createRequest(int timeoutMs) { + return new ExpireDelegationTokenRequest.Builder( + new ExpireDelegationTokenRequestData() + .setHmac(hmac) + .setExpiryTimePeriodMs(options.expiryTimePeriodMs())); } @Override @@ -2390,21 +2588,93 @@ void handleFailure(Throwable throwable) { return new DescribeDelegationTokenResult(tokensFuture); } + /** + * Context class to encapsulate parameters of a call to find and use a consumer group coordinator. + * Some of the parameters are provided at construction and are immutable whereas others are provided + * as "Call" are completed and values are available, like node id of the coordinator. + * + * @param The type of return value of the KafkaFuture + * @param The type of configuration option. Different for different consumer group commands. + */ + private final static class ConsumerGroupOperationContext> { + final private String groupId; + final private O options; + final private long deadline; + final private KafkaFutureImpl future; + private Optional node; + + public ConsumerGroupOperationContext(String groupId, + O options, + long deadline, + KafkaFutureImpl future) { + this.groupId = groupId; + this.options = options; + this.deadline = deadline; + this.future = future; + this.node = Optional.empty(); + } + + public String getGroupId() { + return groupId; + } + + public O getOptions() { + return options; + } + + public long getDeadline() { + return deadline; + } + + public KafkaFutureImpl getFuture() { + return future; + } + + public Optional getNode() { + return node; + } + + public void setNode(Node node) { + this.node = Optional.ofNullable(node); + } + + public boolean hasCoordinatorMoved(AbstractResponse response) { + return response.errorCounts().keySet() + .stream() + .anyMatch(error -> error == Errors.NOT_COORDINATOR); + } + } + + private void rescheduleTask(ConsumerGroupOperationContext context, Supplier nextCall) { + log.info("Node {} is no longer the Coordinator. Retrying with new coordinator.", + context.getNode().orElse(null)); + // Requeue the task so that we can try with new coordinator + context.setNode(null); + Call findCoordinatorCall = getFindCoordinatorCall(context, nextCall); + runnable.call(findCoordinatorCall, time.milliseconds()); + } + + private static Map> createFutures(Collection groupIds) { + return new HashSet<>(groupIds).stream().collect( + Collectors.toMap(groupId -> groupId, + groupId -> { + if (groupIdIsUnrepresentable(groupId)) { + KafkaFutureImpl future = new KafkaFutureImpl<>(); + future.completeExceptionally(new InvalidGroupIdException("The given group id '" + + groupId + "' cannot be represented in a request.")); + return future; + } else { + return new KafkaFutureImpl<>(); + } + } + )); + } + @Override public DescribeConsumerGroupsResult describeConsumerGroups(final Collection groupIds, final DescribeConsumerGroupsOptions options) { - final Map> futures = new HashMap<>(groupIds.size()); - for (String groupId: groupIds) { - if (groupIdIsUnrepresentable(groupId)) { - KafkaFutureImpl future = new KafkaFutureImpl<>(); - future.completeExceptionally(new InvalidGroupIdException("The given group id '" + - groupId + "' cannot be represented in a request.")); - futures.put(groupId, future); - } else if (!futures.containsKey(groupId)) { - futures.put(groupId, new KafkaFutureImpl<>()); - } - } + final Map> futures = createFutures(groupIds); // TODO: KAFKA-6788, we should consider grouping the request per coordinator and send one request with a list of // all consumer groups this coordinator host @@ -2417,98 +2687,145 @@ public DescribeConsumerGroupsResult describeConsumerGroups(final Collection context = + new ConsumerGroupOperationContext<>(groupId, options, deadline, futures.get(groupId)); + Call findCoordinatorCall = getFindCoordinatorCall(context, + () -> getDescribeConsumerGroupsCall(context)); + runnable.call(findCoordinatorCall, startFindCoordinatorMs); + } - runnable.call(new Call("findCoordinator", deadline, new LeastLoadedNodeProvider()) { - @Override - AbstractRequest.Builder createRequest(int timeoutMs) { - return new FindCoordinatorRequest.Builder(FindCoordinatorRequest.CoordinatorType.GROUP, groupId); - } + return new DescribeConsumerGroupsResult(new HashMap<>(futures)); + } - @Override - void handleResponse(AbstractResponse abstractResponse) { - final FindCoordinatorResponse fcResponse = (FindCoordinatorResponse) abstractResponse; + /** + * Returns a {@code Call} object to fetch the coordinator for a consumer group id. Takes another Call + * parameter to schedule action that need to be taken using the coordinator. The param is a Supplier + * so that it can be lazily created, so that it can use the results of find coordinator call in its + * construction. + * + * @param The type of return value of the KafkaFuture, like ConsumerGroupDescription, Void etc. + * @param The type of configuration option, like DescribeConsumerGroupsOptions, ListConsumerGroupsOptions etc + */ + private > Call getFindCoordinatorCall(ConsumerGroupOperationContext context, + Supplier nextCall) { + return new Call("findCoordinator", context.getDeadline(), new LeastLoadedNodeProvider()) { + @Override + FindCoordinatorRequest.Builder createRequest(int timeoutMs) { + return new FindCoordinatorRequest.Builder( + new FindCoordinatorRequestData() + .setKeyType(CoordinatorType.GROUP.id()) + .setKey(context.getGroupId())); + } - if (handleFindCoordinatorError(fcResponse, futures.get(groupId))) - return; + @Override + void handleResponse(AbstractResponse abstractResponse) { + final FindCoordinatorResponse response = (FindCoordinatorResponse) abstractResponse; - final long nowDescribeConsumerGroups = time.milliseconds(); - final int nodeId = fcResponse.node().id(); - runnable.call(new Call("describeConsumerGroups", deadline, new ConstantNodeIdProvider(nodeId)) { - @Override - AbstractRequest.Builder createRequest(int timeoutMs) { - return new DescribeGroupsRequest.Builder( - new DescribeGroupsRequestData() - .setGroups(Collections.singletonList(groupId)) - .setIncludeAuthorizedOperations(options.includeAuthorizedOperations())); - } + if (handleGroupRequestError(response.error(), context.getFuture())) + return; - @Override - void handleResponse(AbstractResponse abstractResponse) { - final DescribeGroupsResponse response = (DescribeGroupsResponse) abstractResponse; - - KafkaFutureImpl future = futures.get(groupId); - final DescribedGroup describedGroup = response.data() - .groups() - .stream() - .filter(group -> groupId.equals(group.groupId())) - .findFirst().get(); - - final Errors groupError = Errors.forCode(describedGroup.errorCode()); - if (groupError != Errors.NONE) { - // TODO: KAFKA-6789, we can retry based on the error code - future.completeExceptionally(groupError.exception()); - } else { - final String protocolType = describedGroup.protocolType(); - if (protocolType.equals(ConsumerProtocol.PROTOCOL_TYPE) || protocolType.isEmpty()) { - final List members = describedGroup.members(); - final List memberDescriptions = new ArrayList<>(members.size()); - final Set authorizedOperations = validAclOperations(describedGroup.authorizedOperations()); - for (DescribedGroupMember groupMember : members) { - Set partitions = Collections.emptySet(); - if (groupMember.memberAssignment().length > 0) { - final PartitionAssignor.Assignment assignment = ConsumerProtocol. - deserializeAssignment(ByteBuffer.wrap(groupMember.memberAssignment())); - partitions = new HashSet<>(assignment.partitions()); - } - final MemberDescription memberDescription = - new MemberDescription(groupMember.memberId(), - groupMember.clientId(), - groupMember.clientHost(), - new MemberAssignment(partitions)); - memberDescriptions.add(memberDescription); - } - final ConsumerGroupDescription consumerGroupDescription = - new ConsumerGroupDescription(groupId, protocolType.isEmpty(), - memberDescriptions, - describedGroup.protocolData(), - ConsumerGroupState.parse(describedGroup.groupState()), - fcResponse.node(), - authorizedOperations); - future.complete(consumerGroupDescription); - } - } - } + context.setNode(response.node()); - @Override - void handleFailure(Throwable throwable) { - KafkaFutureImpl future = futures.get(groupId); - future.completeExceptionally(throwable); - } - }, nowDescribeConsumerGroups); + runnable.call(nextCall.get(), time.milliseconds()); + } + + @Override + void handleFailure(Throwable throwable) { + context.getFuture().completeExceptionally(throwable); + } + }; + } + + private Call getDescribeConsumerGroupsCall( + ConsumerGroupOperationContext context) { + return new Call("describeConsumerGroups", + context.getDeadline(), + new ConstantNodeIdProvider(context.getNode().get().id())) { + @Override + AbstractRequest.Builder createRequest(int timeoutMs) { + return new DescribeGroupsRequest.Builder( + new DescribeGroupsRequestData() + .setGroups(Collections.singletonList(context.getGroupId())) + .setIncludeAuthorizedOperations(context.getOptions().includeAuthorizedOperations())); + } + + @Override + void handleResponse(AbstractResponse abstractResponse) { + final DescribeGroupsResponse response = (DescribeGroupsResponse) abstractResponse; + + List describedGroups = response.data().groups(); + if (describedGroups.isEmpty()) { + context.getFuture().completeExceptionally( + new InvalidGroupIdException("No consumer group found for GroupId: " + context.getGroupId())); + return; } - @Override - void handleFailure(Throwable throwable) { - KafkaFutureImpl future = futures.get(groupId); - future.completeExceptionally(throwable); + if (describedGroups.size() > 1 || + !describedGroups.get(0).groupId().equals(context.getGroupId())) { + String ids = Arrays.toString(describedGroups.stream().map(DescribedGroup::groupId).toArray()); + context.getFuture().completeExceptionally(new InvalidGroupIdException( + "DescribeConsumerGroup request for GroupId: " + context.getGroupId() + " returned " + ids)); + return; } - }, startFindCoordinatorMs); - } - return new DescribeConsumerGroupsResult(new HashMap<>(futures)); + final DescribedGroup describedGroup = describedGroups.get(0); + + // If coordinator changed since we fetched it, retry + if (context.hasCoordinatorMoved(response)) { + rescheduleTask(context, () -> getDescribeConsumerGroupsCall(context)); + return; + } + + final Errors groupError = Errors.forCode(describedGroup.errorCode()); + if (handleGroupRequestError(groupError, context.getFuture())) + return; + + final String protocolType = describedGroup.protocolType(); + if (protocolType.equals(ConsumerProtocol.PROTOCOL_TYPE) || protocolType.isEmpty()) { + final List members = describedGroup.members(); + final List memberDescriptions = new ArrayList<>(members.size()); + final Set authorizedOperations = validAclOperations(describedGroup.authorizedOperations()); + for (DescribedGroupMember groupMember : members) { + Set partitions = Collections.emptySet(); + if (groupMember.memberAssignment().length > 0) { + final Assignment assignment = ConsumerProtocol. + deserializeAssignment(ByteBuffer.wrap(groupMember.memberAssignment())); + partitions = new HashSet<>(assignment.partitions()); + } + final MemberDescription memberDescription = new MemberDescription( + groupMember.memberId(), + Optional.ofNullable(groupMember.groupInstanceId()), + groupMember.clientId(), + groupMember.clientHost(), + new MemberAssignment(partitions)); + memberDescriptions.add(memberDescription); + } + final ConsumerGroupDescription consumerGroupDescription = + new ConsumerGroupDescription(context.getGroupId(), protocolType.isEmpty(), + memberDescriptions, + describedGroup.protocolData(), + ConsumerGroupState.parse(describedGroup.groupState()), + context.getNode().get(), + authorizedOperations); + context.getFuture().complete(consumerGroupDescription); + } else { + context.getFuture().completeExceptionally(new IllegalArgumentException( + String.format("GroupId %s is not a consumer group (%s).", + context.getGroupId(), protocolType))); + } + } + + @Override + void handleFailure(Throwable throwable) { + context.getFuture().completeExceptionally(throwable); + } + }; } private Set validAclOperations(final int authorizedOperations) { + if (authorizedOperations == MetadataResponse.AUTHORIZED_OPERATIONS_OMITTED) { + return null; + } return Utils.from32BitField(authorizedOperations) .stream() .map(AclOperation::fromCode) @@ -2518,11 +2835,10 @@ private Set validAclOperations(final int authorizedOperations) { .collect(Collectors.toSet()); } - private boolean handleFindCoordinatorError(FindCoordinatorResponse response, KafkaFutureImpl future) { - Errors error = response.error(); - if (error.exception() instanceof RetriableException) { + private boolean handleGroupRequestError(Errors error, KafkaFutureImpl future) { + if (error == Errors.COORDINATOR_LOAD_IN_PROGRESS || error == Errors.COORDINATOR_NOT_AVAILABLE) { throw error.exception(); - } else if (response.hasError()) { + } else if (error != Errors.NONE) { future.completeExceptionally(error.exception()); return true; } @@ -2583,7 +2899,9 @@ public ListConsumerGroupsResult listConsumerGroups(ListConsumerGroupsOptions opt runnable.call(new Call("findAllBrokers", deadline, new LeastLoadedNodeProvider()) { @Override AbstractRequest.Builder createRequest(int timeoutMs) { - return new MetadataRequest.Builder(Collections.emptyList(), true); + return new MetadataRequest.Builder(new MetadataRequestData() + .setTopics(Collections.emptyList()) + .setAllowAutoTopicCreation(true)); } @Override @@ -2601,10 +2919,10 @@ void handleResponse(AbstractResponse abstractResponse) { runnable.call(new Call("listConsumerGroups", deadline, new ConstantNodeIdProvider(node.id())) { @Override AbstractRequest.Builder createRequest(int timeoutMs) { - return new ListGroupsRequest.Builder(); + return new ListGroupsRequest.Builder(new ListGroupsRequestData()); } - private void maybeAddConsumerGroup(ListGroupsResponse.Group group) { + private void maybeAddConsumerGroup(ListGroupsResponseData.ListedGroup group) { String protocolType = group.protocolType(); if (protocolType.equals(ConsumerProtocol.PROTOCOL_TYPE) || protocolType.isEmpty()) { final String groupId = group.groupId(); @@ -2617,13 +2935,13 @@ private void maybeAddConsumerGroup(ListGroupsResponse.Group group) { void handleResponse(AbstractResponse abstractResponse) { final ListGroupsResponse response = (ListGroupsResponse) abstractResponse; synchronized (results) { - Errors error = response.error(); + Errors error = Errors.forCode(response.data().errorCode()); if (error == Errors.COORDINATOR_LOAD_IN_PROGRESS || error == Errors.COORDINATOR_NOT_AVAILABLE) { throw error.exception(); } else if (error != Errors.NONE) { results.addError(error.exception(), node); } else { - for (ListGroupsResponse.Group group : response.groups()) { + for (ListGroupsResponseData.ListedGroup group : response.data().groups()) { maybeAddConsumerGroup(group); } } @@ -2653,158 +2971,216 @@ void handleFailure(Throwable throwable) { } @Override - public ListConsumerGroupOffsetsResult listConsumerGroupOffsets(final String groupId, final ListConsumerGroupOffsetsOptions options) { + public ListConsumerGroupOffsetsResult listConsumerGroupOffsets(final String groupId, + final ListConsumerGroupOffsetsOptions options) { final KafkaFutureImpl> groupOffsetListingFuture = new KafkaFutureImpl<>(); - final long startFindCoordinatorMs = time.milliseconds(); final long deadline = calcDeadlineMs(startFindCoordinatorMs, options.timeoutMs()); - runnable.call(new Call("findCoordinator", deadline, new LeastLoadedNodeProvider()) { + ConsumerGroupOperationContext, ListConsumerGroupOffsetsOptions> context = + new ConsumerGroupOperationContext<>(groupId, options, deadline, groupOffsetListingFuture); + + Call findCoordinatorCall = getFindCoordinatorCall(context, + () -> getListConsumerGroupOffsetsCall(context)); + runnable.call(findCoordinatorCall, startFindCoordinatorMs); + + return new ListConsumerGroupOffsetsResult(groupOffsetListingFuture); + } + + private Call getListConsumerGroupOffsetsCall(ConsumerGroupOperationContext, + ListConsumerGroupOffsetsOptions> context) { + return new Call("listConsumerGroupOffsets", context.getDeadline(), + new ConstantNodeIdProvider(context.getNode().get().id())) { @Override AbstractRequest.Builder createRequest(int timeoutMs) { - return new FindCoordinatorRequest.Builder(FindCoordinatorRequest.CoordinatorType.GROUP, groupId); + return new OffsetFetchRequest.Builder(context.getGroupId(), context.getOptions().topicPartitions()); } @Override void handleResponse(AbstractResponse abstractResponse) { - final FindCoordinatorResponse response = (FindCoordinatorResponse) abstractResponse; + final OffsetFetchResponse response = (OffsetFetchResponse) abstractResponse; + final Map groupOffsetsListing = new HashMap<>(); - if (handleFindCoordinatorError(response, groupOffsetListingFuture)) + // If coordinator changed since we fetched it, retry + if (context.hasCoordinatorMoved(response)) { + rescheduleTask(context, () -> getListConsumerGroupOffsetsCall(context)); return; + } - final long nowListConsumerGroupOffsets = time.milliseconds(); - - final int nodeId = response.node().id(); - - runnable.call(new Call("listConsumerGroupOffsets", deadline, new ConstantNodeIdProvider(nodeId)) { - @Override - AbstractRequest.Builder createRequest(int timeoutMs) { - return new OffsetFetchRequest.Builder(groupId, options.topicPartitions()); - } - - @Override - void handleResponse(AbstractResponse abstractResponse) { - final OffsetFetchResponse response = (OffsetFetchResponse) abstractResponse; - final Map groupOffsetsListing = new HashMap<>(); + if (handleGroupRequestError(response.error(), context.getFuture())) + return; - if (response.hasError()) { - groupOffsetListingFuture.completeExceptionally(response.error().exception()); + for (Map.Entry entry : + response.responseData().entrySet()) { + final TopicPartition topicPartition = entry.getKey(); + OffsetFetchResponse.PartitionData partitionData = entry.getValue(); + final Errors error = partitionData.error; + + if (error == Errors.NONE) { + final Long offset = partitionData.offset; + final String metadata = partitionData.metadata; + final Optional leaderEpoch = partitionData.leaderEpoch; + // Negative offset indicates that the group has no committed offset for this partition + if (offset < 0) { + groupOffsetsListing.put(topicPartition, null); } else { - for (Map.Entry entry : - response.responseData().entrySet()) { - final TopicPartition topicPartition = entry.getKey(); - OffsetFetchResponse.PartitionData partitionData = entry.getValue(); - final Errors error = partitionData.error; - - if (error == Errors.NONE) { - final Long offset = partitionData.offset; - final String metadata = partitionData.metadata; - final Optional leaderEpoch = partitionData.leaderEpoch; - groupOffsetsListing.put(topicPartition, new OffsetAndMetadata(offset, leaderEpoch, metadata)); - } else { - log.warn("Skipping return offset for {} due to error {}.", topicPartition, error); - } - } - groupOffsetListingFuture.complete(groupOffsetsListing); + groupOffsetsListing.put(topicPartition, new OffsetAndMetadata(offset, leaderEpoch, metadata)); } + } else { + log.warn("Skipping return offset for {} due to error {}.", topicPartition, error); } - - @Override - void handleFailure(Throwable throwable) { - groupOffsetListingFuture.completeExceptionally(throwable); - } - }, nowListConsumerGroupOffsets); + } + context.getFuture().complete(groupOffsetsListing); } @Override void handleFailure(Throwable throwable) { - groupOffsetListingFuture.completeExceptionally(throwable); + context.getFuture().completeExceptionally(throwable); } - }, startFindCoordinatorMs); - - return new ListConsumerGroupOffsetsResult(groupOffsetListingFuture); + }; } @Override public DeleteConsumerGroupsResult deleteConsumerGroups(Collection groupIds, DeleteConsumerGroupsOptions options) { - final Map> futures = new HashMap<>(groupIds.size()); - for (String groupId: groupIds) { - if (groupIdIsUnrepresentable(groupId)) { - KafkaFutureImpl future = new KafkaFutureImpl<>(); - future.completeExceptionally(new ApiException("The given group id '" + - groupId + "' cannot be represented in a request.")); - futures.put(groupId, future); - } else if (!futures.containsKey(groupId)) { - futures.put(groupId, new KafkaFutureImpl<>()); - } - } + final Map> futures = createFutures(groupIds); // TODO: KAFKA-6788, we should consider grouping the request per coordinator and send one request with a list of // all consumer groups this coordinator host for (final String groupId : groupIds) { // skip sending request for those futures that already failed. - if (futures.get(groupId).isCompletedExceptionally()) + final KafkaFutureImpl future = futures.get(groupId); + if (future.isCompletedExceptionally()) continue; final long startFindCoordinatorMs = time.milliseconds(); final long deadline = calcDeadlineMs(startFindCoordinatorMs, options.timeoutMs()); + ConsumerGroupOperationContext context = + new ConsumerGroupOperationContext<>(groupId, options, deadline, future); + Call findCoordinatorCall = getFindCoordinatorCall(context, + () -> getDeleteConsumerGroupsCall(context)); + runnable.call(findCoordinatorCall, startFindCoordinatorMs); + } - runnable.call(new Call("findCoordinator", deadline, new LeastLoadedNodeProvider()) { - @Override - AbstractRequest.Builder createRequest(int timeoutMs) { - return new FindCoordinatorRequest.Builder(FindCoordinatorRequest.CoordinatorType.GROUP, groupId); + return new DeleteConsumerGroupsResult(new HashMap<>(futures)); + } + + private Call getDeleteConsumerGroupsCall(ConsumerGroupOperationContext context) { + return new Call("deleteConsumerGroups", context.getDeadline(), new ConstantNodeIdProvider(context.getNode().get().id())) { + + @Override + AbstractRequest.Builder createRequest(int timeoutMs) { + return new DeleteGroupsRequest.Builder( + new DeleteGroupsRequestData() + .setGroupsNames(Collections.singletonList(context.getGroupId())) + ); + } + + @Override + void handleResponse(AbstractResponse abstractResponse) { + final DeleteGroupsResponse response = (DeleteGroupsResponse) abstractResponse; + + // If coordinator changed since we fetched it, retry + if (context.hasCoordinatorMoved(response)) { + rescheduleTask(context, () -> getDeleteConsumerGroupsCall(context)); + return; } - @Override - void handleResponse(AbstractResponse abstractResponse) { - final FindCoordinatorResponse response = (FindCoordinatorResponse) abstractResponse; + final Errors groupError = response.get(context.getGroupId()); + if (handleGroupRequestError(groupError, context.getFuture())) + return; - if (handleFindCoordinatorError(response, futures.get(groupId))) - return; + context.getFuture().complete(null); + } - final long nowDeleteConsumerGroups = time.milliseconds(); + @Override + void handleFailure(Throwable throwable) { + context.getFuture().completeExceptionally(throwable); + } + }; + } - final int nodeId = response.node().id(); + @Override + public DeleteConsumerGroupOffsetsResult deleteConsumerGroupOffsets( + String groupId, + Set partitions, + DeleteConsumerGroupOffsetsOptions options) { + final KafkaFutureImpl> future = new KafkaFutureImpl<>(); - runnable.call(new Call("deleteConsumerGroups", deadline, new ConstantNodeIdProvider(nodeId)) { + if (groupIdIsUnrepresentable(groupId)) { + future.completeExceptionally(new InvalidGroupIdException("The given group id '" + + groupId + "' cannot be represented in a request.")); + return new DeleteConsumerGroupOffsetsResult(future, partitions); + } - @Override - AbstractRequest.Builder createRequest(int timeoutMs) { - return new DeleteGroupsRequest.Builder(Collections.singleton(groupId)); - } + final long startFindCoordinatorMs = time.milliseconds(); + final long deadline = calcDeadlineMs(startFindCoordinatorMs, options.timeoutMs()); + ConsumerGroupOperationContext, DeleteConsumerGroupOffsetsOptions> context = + new ConsumerGroupOperationContext<>(groupId, options, deadline, future); - @Override - void handleResponse(AbstractResponse abstractResponse) { - final DeleteGroupsResponse response = (DeleteGroupsResponse) abstractResponse; + Call findCoordinatorCall = getFindCoordinatorCall(context, + () -> getDeleteConsumerGroupOffsetsCall(context, partitions)); + runnable.call(findCoordinatorCall, startFindCoordinatorMs); - KafkaFutureImpl future = futures.get(groupId); - final Errors groupError = response.get(groupId); + return new DeleteConsumerGroupOffsetsResult(future, partitions); + } - if (groupError != Errors.NONE) { - future.completeExceptionally(groupError.exception()); - } else { - future.complete(null); - } - } + private Call getDeleteConsumerGroupOffsetsCall( + ConsumerGroupOperationContext, DeleteConsumerGroupOffsetsOptions> context, + Set partitions) { + return new Call("deleteConsumerGroupOffsets", context.getDeadline(), new ConstantNodeIdProvider(context.getNode().get().id())) { - @Override - void handleFailure(Throwable throwable) { - KafkaFutureImpl future = futures.get(groupId); - future.completeExceptionally(throwable); - } - }, nowDeleteConsumerGroups); - } + @Override + AbstractRequest.Builder createRequest(int timeoutMs) { + final OffsetDeleteRequestTopicCollection topics = new OffsetDeleteRequestTopicCollection(); + + partitions.stream().collect(Collectors.groupingBy(TopicPartition::topic)).forEach((topic, topicPartitions) -> { + topics.add( + new OffsetDeleteRequestTopic() + .setName(topic) + .setPartitions(topicPartitions.stream() + .map(tp -> new OffsetDeleteRequestPartition().setPartitionIndex(tp.partition())) + .collect(Collectors.toList()) + ) + ); + }); + + return new OffsetDeleteRequest.Builder( + new OffsetDeleteRequestData() + .setGroupId(context.groupId) + .setTopics(topics) + ); + } - @Override - void handleFailure(Throwable throwable) { - KafkaFutureImpl future = futures.get(groupId); - future.completeExceptionally(throwable); + @Override + void handleResponse(AbstractResponse abstractResponse) { + final OffsetDeleteResponse response = (OffsetDeleteResponse) abstractResponse; + + // If coordinator changed since we fetched it, retry + if (context.hasCoordinatorMoved(response)) { + rescheduleTask(context, () -> getDeleteConsumerGroupOffsetsCall(context, partitions)); + return; } - }, startFindCoordinatorMs); - } - return new DeleteConsumerGroupsResult(new HashMap<>(futures)); + // If the error is an error at the group level, the future is failed with it + final Errors groupError = Errors.forCode(response.data.errorCode()); + if (handleGroupRequestError(groupError, context.getFuture())) + return; + + final Map partitions = new HashMap<>(); + response.data.topics().forEach(topic -> topic.partitions().forEach(partition -> partitions.put( + new TopicPartition(topic.name(), partition.partitionIndex()), + Errors.forCode(partition.errorCode()))) + ); + + context.getFuture().complete(partitions); + } + + @Override + void handleFailure(Throwable throwable) { + context.getFuture().completeExceptionally(throwable); + } + }; } @Override @@ -2813,25 +3189,33 @@ void handleFailure(Throwable throwable) { } @Override - public ElectPreferredLeadersResult electPreferredLeaders(final Collection partitions, - ElectPreferredLeadersOptions options) { - final Set partitionSet = partitions != null ? new HashSet<>(partitions) : null; - final KafkaFutureImpl> electionFuture = new KafkaFutureImpl<>(); + public ElectLeadersResult electLeaders( + final ElectionType electionType, + final Set topicPartitions, + ElectLeadersOptions options) { + final KafkaFutureImpl>> electionFuture = new KafkaFutureImpl<>(); final long now = time.milliseconds(); - runnable.call(new Call("electPreferredLeaders", calcDeadlineMs(now, options.timeoutMs()), + runnable.call(new Call("electLeaders", calcDeadlineMs(now, options.timeoutMs()), new ControllerNodeProvider()) { @Override public AbstractRequest.Builder createRequest(int timeoutMs) { - return new ElectPreferredLeadersRequest.Builder( - ElectPreferredLeadersRequest.toRequestData(partitions, timeoutMs)); + return new ElectLeadersRequest.Builder(electionType, topicPartitions, timeoutMs); } @Override public void handleResponse(AbstractResponse abstractResponse) { - ElectPreferredLeadersResponse response = (ElectPreferredLeadersResponse) abstractResponse; - electionFuture.complete( - ElectPreferredLeadersRequest.fromResponseData(response.data())); + ElectLeadersResponse response = (ElectLeadersResponse) abstractResponse; + Map> result = ElectLeadersResponse.electLeadersResult(response.data()); + + // For version == 0 then errorCode would be 0 which maps to Errors.NONE + Errors error = Errors.forCode(response.data().errorCode()); + if (error != Errors.NONE) { + electionFuture.completeExceptionally(error.exception()); + return; + } + + electionFuture.complete(result); } @Override @@ -2839,7 +3223,324 @@ void handleFailure(Throwable throwable) { electionFuture.completeExceptionally(throwable); } }, now); - return new ElectPreferredLeadersResult(electionFuture, partitionSet); + + return new ElectLeadersResult(electionFuture); + } + + @Override + public AlterPartitionReassignmentsResult alterPartitionReassignments( + Map> reassignments, + AlterPartitionReassignmentsOptions options) { + final Map> futures = new HashMap<>(); + final Map>> topicsToReassignments = new TreeMap<>(); + for (Map.Entry> entry : reassignments.entrySet()) { + String topic = entry.getKey().topic(); + int partition = entry.getKey().partition(); + TopicPartition topicPartition = new TopicPartition(topic, partition); + Optional reassignment = entry.getValue(); + KafkaFutureImpl future = new KafkaFutureImpl<>(); + futures.put(topicPartition, future); + + if (topicNameIsUnrepresentable(topic)) { + future.completeExceptionally(new InvalidTopicException("The given topic name '" + + topic + "' cannot be represented in a request.")); + } else if (topicPartition.partition() < 0) { + future.completeExceptionally(new InvalidTopicException("The given partition index " + + topicPartition.partition() + " is not valid.")); + } else { + Map> partitionReassignments = + topicsToReassignments.get(topicPartition.topic()); + if (partitionReassignments == null) { + partitionReassignments = new TreeMap<>(); + topicsToReassignments.put(topic, partitionReassignments); + } + + partitionReassignments.put(partition, reassignment); + } + } + + final long now = time.milliseconds(); + Call call = new Call("alterPartitionReassignments", calcDeadlineMs(now, options.timeoutMs()), + new ControllerNodeProvider()) { + + @Override + public AbstractRequest.Builder createRequest(int timeoutMs) { + AlterPartitionReassignmentsRequestData data = + new AlterPartitionReassignmentsRequestData(); + for (Map.Entry>> entry : + topicsToReassignments.entrySet()) { + String topicName = entry.getKey(); + Map> partitionsToReassignments = entry.getValue(); + + List reassignablePartitions = new ArrayList<>(); + for (Map.Entry> partitionEntry : + partitionsToReassignments.entrySet()) { + int partitionIndex = partitionEntry.getKey(); + Optional reassignment = partitionEntry.getValue(); + + ReassignablePartition reassignablePartition = new ReassignablePartition() + .setPartitionIndex(partitionIndex) + .setReplicas(reassignment.map(NewPartitionReassignment::targetReplicas).orElse(null)); + reassignablePartitions.add(reassignablePartition); + } + + ReassignableTopic reassignableTopic = new ReassignableTopic() + .setName(topicName) + .setPartitions(reassignablePartitions); + data.topics().add(reassignableTopic); + } + data.setTimeoutMs(timeoutMs); + return new AlterPartitionReassignmentsRequest.Builder(data); + } + + @Override + public void handleResponse(AbstractResponse abstractResponse) { + AlterPartitionReassignmentsResponse response = (AlterPartitionReassignmentsResponse) abstractResponse; + Map errors = new HashMap<>(); + int receivedResponsesCount = 0; + + Errors topLevelError = Errors.forCode(response.data().errorCode()); + switch (topLevelError) { + case NONE: + receivedResponsesCount += validateTopicResponses(response.data().responses(), errors); + break; + case NOT_CONTROLLER: + handleNotControllerError(topLevelError); + break; + default: + for (ReassignableTopicResponse topicResponse : response.data().responses()) { + String topicName = topicResponse.name(); + for (ReassignablePartitionResponse partition : topicResponse.partitions()) { + errors.put( + new TopicPartition(topicName, partition.partitionIndex()), + new ApiError(topLevelError, topLevelError.message()).exception() + ); + receivedResponsesCount += 1; + } + } + break; + } + + assertResponseCountMatch(errors, receivedResponsesCount); + for (Map.Entry entry : errors.entrySet()) { + ApiException exception = entry.getValue(); + if (exception == null) + futures.get(entry.getKey()).complete(null); + else + futures.get(entry.getKey()).completeExceptionally(exception); + } + } + + private void assertResponseCountMatch(Map errors, int receivedResponsesCount) { + int expectedResponsesCount = topicsToReassignments.values().stream().mapToInt(Map::size).sum(); + if (errors.values().stream().noneMatch(Objects::nonNull) && receivedResponsesCount != expectedResponsesCount) { + String quantifier = receivedResponsesCount > expectedResponsesCount ? "many" : "less"; + throw new UnknownServerException("The server returned too " + quantifier + " results." + + "Expected " + expectedResponsesCount + " but received " + receivedResponsesCount); + } + } + + private int validateTopicResponses(List topicResponses, + Map errors) { + int receivedResponsesCount = 0; + + for (ReassignableTopicResponse topicResponse : topicResponses) { + String topicName = topicResponse.name(); + for (ReassignablePartitionResponse partResponse : topicResponse.partitions()) { + Errors partitionError = Errors.forCode(partResponse.errorCode()); + + TopicPartition tp = new TopicPartition(topicName, partResponse.partitionIndex()); + if (partitionError == Errors.NONE) { + errors.put(tp, null); + } else { + errors.put(tp, new ApiError(partitionError, partResponse.errorMessage()).exception()); + } + receivedResponsesCount += 1; + } + } + + return receivedResponsesCount; + } + + @Override + void handleFailure(Throwable throwable) { + for (KafkaFutureImpl future : futures.values()) { + future.completeExceptionally(throwable); + } + } + }; + if (!topicsToReassignments.isEmpty()) { + runnable.call(call, now); + } + return new AlterPartitionReassignmentsResult(new HashMap<>(futures)); + } + + @Override + public ListPartitionReassignmentsResult listPartitionReassignments(Optional> partitions, + ListPartitionReassignmentsOptions options) { + final KafkaFutureImpl> partitionReassignmentsFuture = new KafkaFutureImpl<>(); + if (partitions.isPresent()) { + for (TopicPartition tp : partitions.get()) { + String topic = tp.topic(); + int partition = tp.partition(); + if (topicNameIsUnrepresentable(topic)) { + partitionReassignmentsFuture.completeExceptionally(new InvalidTopicException("The given topic name '" + + topic + "' cannot be represented in a request.")); + } else if (partition < 0) { + partitionReassignmentsFuture.completeExceptionally(new InvalidTopicException("The given partition index " + + partition + " is not valid.")); + } + if (partitionReassignmentsFuture.isCompletedExceptionally()) + return new ListPartitionReassignmentsResult(partitionReassignmentsFuture); + } + } + final long now = time.milliseconds(); + runnable.call(new Call("listPartitionReassignments", calcDeadlineMs(now, options.timeoutMs()), + new ControllerNodeProvider()) { + + @Override + AbstractRequest.Builder createRequest(int timeoutMs) { + ListPartitionReassignmentsRequestData listData = new ListPartitionReassignmentsRequestData(); + listData.setTimeoutMs(timeoutMs); + + if (partitions.isPresent()) { + Map reassignmentTopicByTopicName = new HashMap<>(); + + for (TopicPartition tp : partitions.get()) { + if (!reassignmentTopicByTopicName.containsKey(tp.topic())) + reassignmentTopicByTopicName.put(tp.topic(), new ListPartitionReassignmentsTopics().setName(tp.topic())); + + reassignmentTopicByTopicName.get(tp.topic()).partitionIndexes().add(tp.partition()); + } + + listData.setTopics(new ArrayList<>(reassignmentTopicByTopicName.values())); + } + return new ListPartitionReassignmentsRequest.Builder(listData); + } + + @Override + void handleResponse(AbstractResponse abstractResponse) { + ListPartitionReassignmentsResponse response = (ListPartitionReassignmentsResponse) abstractResponse; + Errors error = Errors.forCode(response.data().errorCode()); + switch (error) { + case NONE: + break; + case NOT_CONTROLLER: + handleNotControllerError(error); + break; + default: + partitionReassignmentsFuture.completeExceptionally(new ApiError(error, response.data().errorMessage()).exception()); + break; + } + Map reassignmentMap = new HashMap<>(); + + for (OngoingTopicReassignment topicReassignment : response.data().topics()) { + String topicName = topicReassignment.name(); + for (OngoingPartitionReassignment partitionReassignment : topicReassignment.partitions()) { + reassignmentMap.put( + new TopicPartition(topicName, partitionReassignment.partitionIndex()), + new PartitionReassignment(partitionReassignment.replicas(), partitionReassignment.addingReplicas(), partitionReassignment.removingReplicas()) + ); + } + } + + partitionReassignmentsFuture.complete(reassignmentMap); + } + + @Override + void handleFailure(Throwable throwable) { + partitionReassignmentsFuture.completeExceptionally(throwable); + } + }, now); + + return new ListPartitionReassignmentsResult(partitionReassignmentsFuture); + } + + private void handleNotControllerError(Errors error) throws ApiException { + metadataManager.clearController(); + metadataManager.requestUpdate(); + throw error.exception(); + } + + /** + * Returns a boolean indicating whether the resource needs to go to a specific node + */ + private boolean dependsOnSpecificNode(ConfigResource resource) { + return (resource.type() == ConfigResource.Type.BROKER && !resource.isDefault()) + || resource.type() == ConfigResource.Type.BROKER_LOGGER; + } + + @Override + public RemoveMembersFromConsumerGroupResult removeMembersFromConsumerGroup(String groupId, + RemoveMembersFromConsumerGroupOptions options) { + final long startFindCoordinatorMs = time.milliseconds(); + final long deadline = calcDeadlineMs(startFindCoordinatorMs, options.timeoutMs()); + + KafkaFutureImpl> future = new KafkaFutureImpl<>(); + + ConsumerGroupOperationContext, RemoveMembersFromConsumerGroupOptions> context = + new ConsumerGroupOperationContext<>(groupId, options, deadline, future); + + Call findCoordinatorCall = getFindCoordinatorCall(context, + () -> getRemoveMembersFromGroupCall(context)); + runnable.call(findCoordinatorCall, startFindCoordinatorMs); + + return new RemoveMembersFromConsumerGroupResult(future, options.members()); + } + + private Call getRemoveMembersFromGroupCall(ConsumerGroupOperationContext, RemoveMembersFromConsumerGroupOptions> context) { + return new Call("leaveGroup", + context.getDeadline(), + new ConstantNodeIdProvider(context.getNode().get().id())) { + @Override + LeaveGroupRequest.Builder createRequest(int timeoutMs) { + return new LeaveGroupRequest.Builder(context.getGroupId(), + context.getOptions().members().stream().map( + MemberToRemove::toMemberIdentity).collect(Collectors.toList())); + } + + @Override + void handleResponse(AbstractResponse abstractResponse) { + final LeaveGroupResponse response = (LeaveGroupResponse) abstractResponse; + + // If coordinator changed since we fetched it, retry + if (context.hasCoordinatorMoved(response)) { + rescheduleTask(context, () -> getRemoveMembersFromGroupCall(context)); + return; + } + + if (handleGroupRequestError(response.topLevelError(), context.getFuture())) + return; + + final Map memberErrors = new HashMap<>(); + for (MemberResponse memberResponse : response.memberResponses()) { + // We set member.id to empty here explicitly, so that the lookup will succeed as user doesn't + // know the exact member.id. + memberErrors.put(new MemberIdentity() + .setMemberId(JoinGroupRequest.UNKNOWN_MEMBER_ID) + .setGroupInstanceId(memberResponse.groupInstanceId()), + Errors.forCode(memberResponse.errorCode())); + } + context.getFuture().complete(memberErrors); + } + + @Override + void handleFailure(Throwable throwable) { + context.getFuture().completeExceptionally(throwable); + } + }; + } + + /** + * Get a sub level error when the request is in batch. If given key was not found, + * return an {@link IllegalArgumentException}. + */ + static Throwable getSubLevelError(Map subLevelErrors, K subKey, String keyNotFoundMsg) { + if (!subLevelErrors.containsKey(subKey)) { + return new IllegalArgumentException(keyNotFoundMsg); + } else { + return subLevelErrors.get(subKey).exception(); + } } } diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/ListConsumerGroupOffsetsOptions.java b/clients/src/main/java/org/apache/kafka/clients/admin/ListConsumerGroupOffsetsOptions.java index c6434ebb15c13..af738ca209fb9 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/ListConsumerGroupOffsetsOptions.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/ListConsumerGroupOffsetsOptions.java @@ -23,9 +23,9 @@ import java.util.List; /** - * Options for {@link AdminClient#listConsumerGroupOffsets(String)}. + * Options for {@link Admin#listConsumerGroupOffsets(String)}. *

- * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class ListConsumerGroupOffsetsOptions extends AbstractOptions { diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/ListConsumerGroupOffsetsResult.java b/clients/src/main/java/org/apache/kafka/clients/admin/ListConsumerGroupOffsetsResult.java index 23657b5ad3cf7..48f4531418110 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/ListConsumerGroupOffsetsResult.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/ListConsumerGroupOffsetsResult.java @@ -25,9 +25,9 @@ import java.util.Map; /** - * The result of the {@link AdminClient#listConsumerGroupOffsets(String)} call. + * The result of the {@link Admin#listConsumerGroupOffsets(String)} call. *

- * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class ListConsumerGroupOffsetsResult { @@ -40,6 +40,7 @@ public class ListConsumerGroupOffsetsResult { /** * Return a future which yields a map of topic partitions to OffsetAndMetadata objects. + * If the group does not have a committed offset for this partition, the corresponding value in the returned map will be null. */ public KafkaFuture> partitionsToOffsetAndMetadata() { return future; 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 86ca171b726c0..eb27c795f2afe 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 @@ -20,9 +20,9 @@ import org.apache.kafka.common.annotation.InterfaceStability; /** - * Options for {@link AdminClient#listConsumerGroups()}. + * Options for {@link Admin#listConsumerGroups()}. * - * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class ListConsumerGroupsOptions extends AbstractOptions { diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/ListConsumerGroupsResult.java b/clients/src/main/java/org/apache/kafka/clients/admin/ListConsumerGroupsResult.java index 7de485bd2b989..7732ec98c9638 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/ListConsumerGroupsResult.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/ListConsumerGroupsResult.java @@ -25,9 +25,9 @@ import java.util.Collection; /** - * The result of the {@link AdminClient#listConsumerGroups()} call. + * The result of the {@link Admin#listConsumerGroups()} call. *

- * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class ListConsumerGroupsResult { diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/ListPartitionReassignmentsOptions.java b/clients/src/main/java/org/apache/kafka/clients/admin/ListPartitionReassignmentsOptions.java new file mode 100644 index 0000000000000..7dcc7a6c3e6d5 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/ListPartitionReassignmentsOptions.java @@ -0,0 +1,29 @@ +/* + * 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.admin; + +import org.apache.kafka.common.annotation.InterfaceStability; + +/** + * Options for {@link AdminClient#listPartitionReassignments(ListPartitionReassignmentsOptions)} + * + * The API of this class is evolving. See {@link AdminClient} for details. + */ +@InterfaceStability.Evolving +public class ListPartitionReassignmentsOptions extends AbstractOptions { +} diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/ListPartitionReassignmentsResult.java b/clients/src/main/java/org/apache/kafka/clients/admin/ListPartitionReassignmentsResult.java new file mode 100644 index 0000000000000..bc72c0683f96b --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/ListPartitionReassignmentsResult.java @@ -0,0 +1,43 @@ +/* + * 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.admin; + +import org.apache.kafka.common.KafkaFuture; +import org.apache.kafka.common.TopicPartition; + +import java.util.Map; + +/** + * The result of {@link AdminClient#listPartitionReassignments(ListPartitionReassignmentsOptions)}. + * + * The API of this class is evolving. See {@link AdminClient} for details. + */ +public class ListPartitionReassignmentsResult { + private final KafkaFuture> future; + + ListPartitionReassignmentsResult(KafkaFuture> reassignments) { + this.future = reassignments; + } + + /** + * Return a future which yields a map containing each partition's reassignments + */ + public KafkaFuture> reassignments() { + return future; + } +} diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/ListTopicsOptions.java b/clients/src/main/java/org/apache/kafka/clients/admin/ListTopicsOptions.java index 1431494646803..e288e1828fd79 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/ListTopicsOptions.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/ListTopicsOptions.java @@ -20,9 +20,9 @@ import org.apache.kafka.common.annotation.InterfaceStability; /** - * Options for {@link AdminClient#listTopics()}. + * Options for {@link Admin#listTopics()}. * - * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class ListTopicsOptions extends AbstractOptions { diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/ListTopicsResult.java b/clients/src/main/java/org/apache/kafka/clients/admin/ListTopicsResult.java index a2e17fdbd9c0a..4e7e1a2e98721 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/ListTopicsResult.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/ListTopicsResult.java @@ -25,9 +25,9 @@ import java.util.Set; /** - * The result of the {@link AdminClient#listTopics()} call. + * The result of the {@link Admin#listTopics()} call. * - * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class ListTopicsResult { diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/MemberAssignment.java b/clients/src/main/java/org/apache/kafka/clients/admin/MemberAssignment.java index 8765c17d984b2..3305de02c757b 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/MemberAssignment.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/MemberAssignment.java @@ -21,6 +21,7 @@ import java.util.Collections; import java.util.HashSet; +import java.util.Objects; import java.util.Set; /** @@ -46,7 +47,7 @@ public boolean equals(Object o) { MemberAssignment that = (MemberAssignment) o; - return topicPartitions != null ? topicPartitions.equals(that.topicPartitions) : that.topicPartitions == null; + return Objects.equals(topicPartitions, that.topicPartitions); } @Override diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/MemberDescription.java b/clients/src/main/java/org/apache/kafka/clients/admin/MemberDescription.java index a8865bed0ee2d..7bc6b14906fe2 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/MemberDescription.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/MemberDescription.java @@ -14,29 +14,40 @@ * See the License for the specific language governing permissions and * limitations under the License. */ - package org.apache.kafka.clients.admin; -import org.apache.kafka.common.TopicPartition; - import java.util.Collections; import java.util.Objects; +import java.util.Optional; /** * A detailed description of a single group instance in the cluster. */ public class MemberDescription { private final String memberId; + private final Optional groupInstanceId; private final String clientId; private final String host; private final MemberAssignment assignment; - public MemberDescription(String memberId, String clientId, String host, MemberAssignment assignment) { + public MemberDescription(String memberId, + Optional groupInstanceId, + String clientId, + String host, + MemberAssignment assignment) { this.memberId = memberId == null ? "" : memberId; + this.groupInstanceId = groupInstanceId; this.clientId = clientId == null ? "" : clientId; this.host = host == null ? "" : host; this.assignment = assignment == null ? - new MemberAssignment(Collections.emptySet()) : assignment; + new MemberAssignment(Collections.emptySet()) : assignment; + } + + public MemberDescription(String memberId, + String clientId, + String host, + MemberAssignment assignment) { + this(memberId, Optional.empty(), clientId, host, assignment); } @Override @@ -45,6 +56,7 @@ public boolean equals(Object o) { if (o == null || getClass() != o.getClass()) return false; MemberDescription that = (MemberDescription) o; return memberId.equals(that.memberId) && + groupInstanceId.equals(that.groupInstanceId) && clientId.equals(that.clientId) && host.equals(that.host) && assignment.equals(that.assignment); @@ -52,7 +64,7 @@ public boolean equals(Object o) { @Override public int hashCode() { - return Objects.hash(memberId, clientId, host, assignment); + return Objects.hash(memberId, groupInstanceId, clientId, host, assignment); } /** @@ -62,6 +74,13 @@ public String consumerId() { return memberId; } + /** + * The instance id of the group member. + */ + public Optional groupInstanceId() { + return groupInstanceId; + } + /** * The client id of the group member. */ @@ -86,6 +105,7 @@ public MemberAssignment assignment() { @Override public String toString() { return "(memberId=" + memberId + + ", groupInstanceId=" + groupInstanceId.orElse("null") + ", clientId=" + clientId + ", host=" + host + ", assignment=" + assignment + ")"; diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/MemberToRemove.java b/clients/src/main/java/org/apache/kafka/clients/admin/MemberToRemove.java new file mode 100644 index 0000000000000..4c7b16b1da650 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/MemberToRemove.java @@ -0,0 +1,58 @@ +/* + * 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.admin; + +import org.apache.kafka.common.message.LeaveGroupRequestData.MemberIdentity; +import org.apache.kafka.common.requests.JoinGroupRequest; + +import java.util.Objects; + +/** + * A struct containing information about the member to be removed. + */ +public class MemberToRemove { + private final String groupInstanceId; + + public MemberToRemove(String groupInstanceId) { + this.groupInstanceId = groupInstanceId; + } + + @Override + public boolean equals(Object o) { + if (o instanceof MemberToRemove) { + MemberToRemove otherMember = (MemberToRemove) o; + return this.groupInstanceId.equals(otherMember.groupInstanceId); + } else { + return false; + } + } + + @Override + public int hashCode() { + return Objects.hash(groupInstanceId); + } + + MemberIdentity toMemberIdentity() { + return new MemberIdentity() + .setGroupInstanceId(groupInstanceId) + .setMemberId(JoinGroupRequest.UNKNOWN_MEMBER_ID); + } + + public String groupInstanceId() { + return groupInstanceId; + } +} diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/NewPartitionReassignment.java b/clients/src/main/java/org/apache/kafka/clients/admin/NewPartitionReassignment.java new file mode 100644 index 0000000000000..f9a7008db737c --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/NewPartitionReassignment.java @@ -0,0 +1,43 @@ +/* + * 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.admin; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Map; + +/** + * A new partition reassignment, which can be applied via {@link AdminClient#alterPartitionReassignments(Map, AlterPartitionReassignmentsOptions)}. + */ +public class NewPartitionReassignment { + private final List targetReplicas; + + /** + * @throws IllegalArgumentException if no replicas are supplied + */ + public NewPartitionReassignment(List targetReplicas) { + if (targetReplicas == null || targetReplicas.size() == 0) + throw new IllegalArgumentException("Cannot create a new partition reassignment without any replicas"); + this.targetReplicas = Collections.unmodifiableList(new ArrayList<>(targetReplicas)); + } + + public List targetReplicas() { + return targetReplicas; + } +} diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/NewPartitions.java b/clients/src/main/java/org/apache/kafka/clients/admin/NewPartitions.java index 66a4d92a56705..06da256fb17d3 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/NewPartitions.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/NewPartitions.java @@ -23,9 +23,9 @@ import java.util.Map; /** - * Describes new partitions for a particular topic in a call to {@link AdminClient#createPartitions(Map)}. + * Describes new partitions for a particular topic in a call to {@link Admin#createPartitions(Map)}. * - * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class NewPartitions { diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/NewTopic.java b/clients/src/main/java/org/apache/kafka/clients/admin/NewTopic.java index 5b1bd32f89acb..088c33e274ae4 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/NewTopic.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/NewTopic.java @@ -17,6 +17,7 @@ package org.apache.kafka.clients.admin; +import java.util.Optional; import org.apache.kafka.common.message.CreateTopicsRequestData.CreatableReplicaAssignment; import org.apache.kafka.common.message.CreateTopicsRequestData.CreatableTopic; import org.apache.kafka.common.message.CreateTopicsRequestData.CreateableTopicConfig; @@ -28,12 +29,16 @@ import java.util.Map.Entry; /** - * A new topic to be created via {@link AdminClient#createTopics(Collection)}. + * A new topic to be created via {@link Admin#createTopics(Collection)}. */ public class NewTopic { + + private static final int NO_PARTITIONS = -1; + private static final short NO_REPLICATION_FACTOR = -1; + private final String name; - private final int numPartitions; - private final short replicationFactor; + private final Optional numPartitions; + private final Optional replicationFactor; private final Map> replicasAssignments; private Map configs = null; @@ -41,6 +46,15 @@ public class NewTopic { * A new topic with the specified replication factor and number of partitions. */ public NewTopic(String name, int numPartitions, short replicationFactor) { + this(name, Optional.of(numPartitions), Optional.of(replicationFactor)); + } + + /** + * A new topic that optionally defaults {@code numPartitions} and {@code replicationFactor} to + * the broker configurations for {@code num.partitions} and {@code default.replication.factor} + * respectively. + */ + public NewTopic(String name, Optional numPartitions, Optional replicationFactor) { this.name = name; this.numPartitions = numPartitions; this.replicationFactor = replicationFactor; @@ -56,8 +70,8 @@ public NewTopic(String name, int numPartitions, short replicationFactor) { */ public NewTopic(String name, Map> replicasAssignments) { this.name = name; - this.numPartitions = -1; - this.replicationFactor = -1; + this.numPartitions = Optional.empty(); + this.replicationFactor = Optional.empty(); this.replicasAssignments = Collections.unmodifiableMap(replicasAssignments); } @@ -72,14 +86,14 @@ public String name() { * The number of partitions for the new topic or -1 if a replica assignment has been specified. */ public int numPartitions() { - return numPartitions; + return numPartitions.orElse(NO_PARTITIONS); } /** * The replication factor for the new topic or -1 if a replica assignment has been specified. */ public short replicationFactor() { - return replicationFactor; + return replicationFactor.orElse(NO_REPLICATION_FACTOR); } /** @@ -111,8 +125,8 @@ public Map configs() { CreatableTopic convertToCreatableTopic() { CreatableTopic creatableTopic = new CreatableTopic(). setName(name). - setNumPartitions(numPartitions). - setReplicationFactor(replicationFactor); + setNumPartitions(numPartitions.orElse(NO_PARTITIONS)). + setReplicationFactor(replicationFactor.orElse(NO_REPLICATION_FACTOR)); if (replicasAssignments != null) { for (Entry> entry : replicasAssignments.entrySet()) { creatableTopic.assignments().add( @@ -136,8 +150,8 @@ CreatableTopic convertToCreatableTopic() { public String toString() { StringBuilder bld = new StringBuilder(); bld.append("(name=").append(name). - append(", numPartitions=").append(numPartitions). - append(", replicationFactor=").append(replicationFactor). + append(", numPartitions=").append(numPartitions.map(String::valueOf).orElse("default")). + append(", replicationFactor=").append(replicationFactor.map(String::valueOf).orElse("default")). append(", replicasAssignments=").append(replicasAssignments). append(", configs=").append(configs). append(")"); diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/PartitionReassignment.java b/clients/src/main/java/org/apache/kafka/clients/admin/PartitionReassignment.java new file mode 100644 index 0000000000000..4a9d151f1b057 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/PartitionReassignment.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.clients.admin; + +import java.util.Collections; +import java.util.List; + +/** + * A partition reassignment, which has been listed via {@link AdminClient#listPartitionReassignments()}. + */ +public class PartitionReassignment { + + private final List replicas; + private final List addingReplicas; + private final List removingReplicas; + + public PartitionReassignment(List replicas, List addingReplicas, List removingReplicas) { + this.replicas = Collections.unmodifiableList(replicas); + this.addingReplicas = Collections.unmodifiableList(addingReplicas); + this.removingReplicas = Collections.unmodifiableList(removingReplicas); + } + + /** + * The brokers which this partition currently resides on. + */ + public List replicas() { + return replicas; + } + + /** + * The brokers that we are adding this partition to as part of a reassignment. + * A subset of replicas. + */ + public List addingReplicas() { + return addingReplicas; + } + + /** + * The brokers that we are removing this partition from as part of a reassignment. + * A subset of replicas. + */ + public List removingReplicas() { + return removingReplicas; + } + + @Override + public String toString() { + return "PartitionReassignment(" + + "replicas=" + replicas + + ", addingReplicas=" + addingReplicas + + ", removingReplicas=" + removingReplicas + + ')'; + } +} diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/RecordsToDelete.java b/clients/src/main/java/org/apache/kafka/clients/admin/RecordsToDelete.java index 53a6dfb59fda6..af835c806103d 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/RecordsToDelete.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/RecordsToDelete.java @@ -22,9 +22,9 @@ import java.util.Map; /** - * Describe records to delete in a call to {@link AdminClient#deleteRecords(Map)} + * Describe records to delete in a call to {@link Admin#deleteRecords(Map)} * - * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class RecordsToDelete { diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/RemoveMembersFromConsumerGroupOptions.java b/clients/src/main/java/org/apache/kafka/clients/admin/RemoveMembersFromConsumerGroupOptions.java new file mode 100644 index 0000000000000..dc346f7c3a1be --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/RemoveMembersFromConsumerGroupOptions.java @@ -0,0 +1,43 @@ +/* + * 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.admin; + +import org.apache.kafka.common.annotation.InterfaceStability; + +import java.util.Collection; +import java.util.HashSet; +import java.util.Set; + +/** + * Options for {@link AdminClient#removeMembersFromConsumerGroup(String, RemoveMembersFromConsumerGroupOptions)}. + * It carries the members to be removed from the consumer group. + * + * The API of this class is evolving, see {@link AdminClient} for details. + */ +@InterfaceStability.Evolving +public class RemoveMembersFromConsumerGroupOptions extends AbstractOptions { + + private Set members; + + public RemoveMembersFromConsumerGroupOptions(Collection members) { + this.members = new HashSet<>(members); + } + + public Set members() { + return members; + } +} diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/RemoveMembersFromConsumerGroupResult.java b/clients/src/main/java/org/apache/kafka/clients/admin/RemoveMembersFromConsumerGroupResult.java new file mode 100644 index 0000000000000..405973b73f226 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/RemoveMembersFromConsumerGroupResult.java @@ -0,0 +1,96 @@ +/* + * 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.admin; + +import org.apache.kafka.common.KafkaFuture; +import org.apache.kafka.common.internals.KafkaFutureImpl; +import org.apache.kafka.common.message.LeaveGroupRequestData.MemberIdentity; +import org.apache.kafka.common.protocol.Errors; + +import java.util.Map; +import java.util.Set; + +/** + * The result of the {@link Admin#removeMembersFromConsumerGroup(String, RemoveMembersFromConsumerGroupOptions)} call. + * + * The API of this class is evolving, see {@link Admin} for details. + */ +public class RemoveMembersFromConsumerGroupResult { + + private final KafkaFuture> future; + private final Set memberInfos; + + RemoveMembersFromConsumerGroupResult(KafkaFuture> future, + Set memberInfos) { + this.future = future; + this.memberInfos = memberInfos; + } + + /** + * Returns a future which indicates whether the request was 100% success, i.e. no + * either top level or member level error. + * If not, the first member error shall be returned. + */ + public KafkaFuture all() { + final KafkaFutureImpl result = new KafkaFutureImpl<>(); + this.future.whenComplete((memberErrors, throwable) -> { + if (throwable != null) { + result.completeExceptionally(throwable); + } else { + for (MemberToRemove memberToRemove : memberInfos) { + if (maybeCompleteExceptionally(memberErrors, memberToRemove.toMemberIdentity(), result)) { + return; + } + } + result.complete(null); + } + }); + return result; + } + + /** + * Returns the selected member future. + */ + public KafkaFuture memberResult(MemberToRemove member) { + if (!memberInfos.contains(member)) { + throw new IllegalArgumentException("Member " + member + " was not included in the original request"); + } + + final KafkaFutureImpl result = new KafkaFutureImpl<>(); + this.future.whenComplete((memberErrors, throwable) -> { + if (throwable != null) { + result.completeExceptionally(throwable); + } else if (!maybeCompleteExceptionally(memberErrors, member.toMemberIdentity(), result)) { + result.complete(null); + } + }); + return result; + } + + private boolean maybeCompleteExceptionally(Map memberErrors, + MemberIdentity member, + KafkaFutureImpl result) { + Throwable exception = KafkaAdminClient.getSubLevelError(memberErrors, member, + "Member \"" + member + "\" was not included in the removal response"); + if (exception != null) { + result.completeExceptionally(exception); + return true; + } else { + return false; + } + } +} diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/RenewDelegationTokenOptions.java b/clients/src/main/java/org/apache/kafka/clients/admin/RenewDelegationTokenOptions.java index 238dc4a3494a8..5c2b0d1afcea1 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/RenewDelegationTokenOptions.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/RenewDelegationTokenOptions.java @@ -20,9 +20,9 @@ import org.apache.kafka.common.annotation.InterfaceStability; /** - * Options for {@link AdminClient#renewDelegationToken(byte[], RenewDelegationTokenOptions)}. + * Options for {@link Admin#renewDelegationToken(byte[], RenewDelegationTokenOptions)}. * - * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class RenewDelegationTokenOptions extends AbstractOptions { diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/RenewDelegationTokenResult.java b/clients/src/main/java/org/apache/kafka/clients/admin/RenewDelegationTokenResult.java index 38cdf1ae1b241..74725d4d633b5 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/RenewDelegationTokenResult.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/RenewDelegationTokenResult.java @@ -23,7 +23,7 @@ /** * The result of the {@link KafkaAdminClient#expireDelegationToken(byte[], ExpireDelegationTokenOptions)} call. * - * The API of this class is evolving, see {@link AdminClient} for details. + * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class RenewDelegationTokenResult { diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/TopicDescription.java b/clients/src/main/java/org/apache/kafka/clients/admin/TopicDescription.java index 4e3e59a30fd19..dc18a0eb04f0e 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/TopicDescription.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/TopicDescription.java @@ -18,9 +18,13 @@ package org.apache.kafka.clients.admin; import org.apache.kafka.common.TopicPartitionInfo; +import org.apache.kafka.common.acl.AclOperation; import org.apache.kafka.common.utils.Utils; +import java.util.Collections; import java.util.List; +import java.util.Objects; +import java.util.Set; /** * A detailed description of a single topic in the cluster. @@ -29,25 +33,22 @@ public class TopicDescription { private final String name; private final boolean internal; private final List partitions; + private final Set authorizedOperations; @Override - public boolean equals(Object o) { + public boolean equals(final Object o) { if (this == o) return true; if (o == null || getClass() != o.getClass()) return false; - - TopicDescription that = (TopicDescription) o; - - if (internal != that.internal) return false; - if (name != null ? !name.equals(that.name) : that.name != null) return false; - return partitions != null ? partitions.equals(that.partitions) : that.partitions == null; + final TopicDescription that = (TopicDescription) o; + return internal == that.internal && + Objects.equals(name, that.name) && + Objects.equals(partitions, that.partitions) && + Objects.equals(authorizedOperations, that.authorizedOperations); } @Override public int hashCode() { - int result = name != null ? name.hashCode() : 0; - result = 31 * result + (internal ? 1 : 0); - result = 31 * result + (partitions != null ? partitions.hashCode() : 0); - return result; + return Objects.hash(name, internal, partitions, authorizedOperations); } /** @@ -59,9 +60,24 @@ public int hashCode() { * leadership and replica information for that partition. */ public TopicDescription(String name, boolean internal, List partitions) { + this(name, internal, partitions, Collections.emptySet()); + } + + /** + * Create an instance with the specified parameters. + * + * @param name The topic name + * @param internal Whether the topic is internal to Kafka + * @param partitions A list of partitions where the index represents the partition id and the element contains + * leadership and replica information for that partition. + * @param authorizedOperations authorized operations for this topic, or null if this is not known. + */ + public TopicDescription(String name, boolean internal, List partitions, + Set authorizedOperations) { this.name = name; this.internal = internal; this.partitions = partitions; + this.authorizedOperations = authorizedOperations; } /** @@ -87,9 +103,16 @@ public List partitions() { return partitions; } + /** + * authorized operations for this topic, or null if this is not known. + */ + public Set authorizedOperations() { + return authorizedOperations; + } + @Override public String toString() { return "(name=" + name + ", internal=" + internal + ", partitions=" + - Utils.join(partitions, ",") + ")"; + Utils.join(partitions, ",") + ", authorizedOperations=" + authorizedOperations + ")"; } } diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/internals/AdminMetadataManager.java b/clients/src/main/java/org/apache/kafka/clients/admin/internals/AdminMetadataManager.java index 3d9e5cad037e0..6e834520c46f7 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/internals/AdminMetadataManager.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/internals/AdminMetadataManager.java @@ -19,6 +19,7 @@ import org.apache.kafka.clients.MetadataUpdater; import org.apache.kafka.common.Cluster; +import org.apache.kafka.common.KafkaException; import org.apache.kafka.common.Node; import org.apache.kafka.common.errors.AuthenticationException; import org.apache.kafka.common.requests.MetadataResponse; @@ -28,6 +29,7 @@ import java.util.Collections; import java.util.List; +import java.util.Optional; /** * Manages the metadata for KafkaAdminClient. @@ -99,23 +101,19 @@ public long maybeUpdate(long now) { } @Override - public void handleDisconnection(String destination) { - // Do nothing - } - - @Override - public void handleAuthenticationFailure(AuthenticationException e) { - updateFailed(e); + public void handleServerDisconnect(long now, String destinationId, Optional maybeFatalException) { + maybeFatalException.ifPresent(AdminMetadataManager.this::updateFailed); + AdminMetadataManager.this.requestUpdate(); } @Override - public void handleCompletedMetadataResponse(RequestHeader requestHeader, long now, MetadataResponse metadataResponse) { + public void handleFailedRequest(long now, Optional maybeFatalException) { // Do nothing } @Override - public void requestUpdate() { - AdminMetadataManager.this.requestUpdate(); + public void handleSuccessfulResponse(RequestHeader requestHeader, long now, MetadataResponse metadataResponse) { + // Do nothing } @Override diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/Consumer.java b/clients/src/main/java/org/apache/kafka/clients/consumer/Consumer.java index 7fdf4398b9add..5136d6886dd2d 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/Consumer.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/Consumer.java @@ -154,13 +154,25 @@ public interface Consumer extends Closeable { /** * @see KafkaConsumer#committed(TopicPartition) */ + @Deprecated OffsetAndMetadata committed(TopicPartition partition); /** * @see KafkaConsumer#committed(TopicPartition, Duration) */ + @Deprecated OffsetAndMetadata committed(TopicPartition partition, final Duration timeout); + /** + * @see KafkaConsumer#committed(Set) + */ + Map committed(Set partitions); + + /** + * @see KafkaConsumer#committed(Set, Duration) + */ + Map committed(Set partitions, final Duration timeout); + /** * @see KafkaConsumer#metrics() */ 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 b92cbf916c213..ddfa30c9f6121 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 @@ -22,6 +22,7 @@ import org.apache.kafka.common.config.ConfigDef; import org.apache.kafka.common.config.ConfigDef.Importance; import org.apache.kafka.common.config.ConfigDef.Type; +import org.apache.kafka.common.config.SecurityConfig; import org.apache.kafka.common.metrics.Sensor; import org.apache.kafka.common.requests.IsolationLevel; import org.apache.kafka.common.serialization.Deserializer; @@ -50,40 +51,33 @@ public class ConsumerConfig extends AbstractConfig { /** * group.id */ - public static final String GROUP_ID_CONFIG = "group.id"; - private static final String GROUP_ID_DOC = "A unique string that identifies the consumer group this consumer belongs to. This property is required if the consumer uses either the group management functionality by using subscribe(topic) or the Kafka-based offset management strategy."; + public static final String GROUP_ID_CONFIG = CommonClientConfigs.GROUP_ID_CONFIG; + private static final String GROUP_ID_DOC = CommonClientConfigs.GROUP_ID_DOC; + + /** + * group.instance.id + */ + public static final String GROUP_INSTANCE_ID_CONFIG = CommonClientConfigs.GROUP_INSTANCE_ID_CONFIG; + private static final String GROUP_INSTANCE_ID_DOC = CommonClientConfigs.GROUP_INSTANCE_ID_DOC; /** max.poll.records */ public static final String MAX_POLL_RECORDS_CONFIG = "max.poll.records"; private static final String MAX_POLL_RECORDS_DOC = "The maximum number of records returned in a single call to poll()."; /** max.poll.interval.ms */ - public static final String MAX_POLL_INTERVAL_MS_CONFIG = "max.poll.interval.ms"; - private static final String MAX_POLL_INTERVAL_MS_DOC = "The maximum delay between invocations of poll() when using " + - "consumer group management. This places an upper bound on the amount of time that the consumer can be idle " + - "before fetching more records. If poll() is not called before expiration of this timeout, then the consumer " + - "is considered failed and the group will rebalance in order to reassign the partitions to another member. "; - + public static final String MAX_POLL_INTERVAL_MS_CONFIG = CommonClientConfigs.MAX_POLL_INTERVAL_MS_CONFIG; + private static final String MAX_POLL_INTERVAL_MS_DOC = CommonClientConfigs.MAX_POLL_INTERVAL_MS_DOC; /** * session.timeout.ms */ - public static final String SESSION_TIMEOUT_MS_CONFIG = "session.timeout.ms"; - private static final String SESSION_TIMEOUT_MS_DOC = "The timeout used to detect consumer failures when using " + - "Kafka's group management facility. The consumer sends periodic heartbeats to indicate its liveness " + - "to the broker. If no heartbeats are received by the broker before the expiration of this session timeout, " + - "then the broker will remove this consumer from the group and initiate a rebalance. Note that the value " + - "must be in the allowable range as configured in the broker configuration by group.min.session.timeout.ms " + - "and group.max.session.timeout.ms."; + public static final String SESSION_TIMEOUT_MS_CONFIG = CommonClientConfigs.SESSION_TIMEOUT_MS_CONFIG; + private static final String SESSION_TIMEOUT_MS_DOC = CommonClientConfigs.SESSION_TIMEOUT_MS_DOC; /** * heartbeat.interval.ms */ - public static final String HEARTBEAT_INTERVAL_MS_CONFIG = "heartbeat.interval.ms"; - private static final String HEARTBEAT_INTERVAL_MS_DOC = "The expected time between heartbeats to the consumer " + - "coordinator when using Kafka's group management facilities. Heartbeats are used to ensure that the " + - "consumer's session stays active and to facilitate rebalancing when new consumers join or leave the group. " + - "The value must be set lower than session.timeout.ms, but typically should be set no higher " + - "than 1/3 of that value. It can be adjusted even lower to control the expected time for normal rebalances."; + public static final String HEARTBEAT_INTERVAL_MS_CONFIG = CommonClientConfigs.HEARTBEAT_INTERVAL_MS_CONFIG; + private static final String HEARTBEAT_INTERVAL_MS_DOC = CommonClientConfigs.HEARTBEAT_INTERVAL_MS_DOC; /** * bootstrap.servers @@ -109,13 +103,13 @@ public class ConsumerConfig extends AbstractConfig { * partition.assignment.strategy */ public static final String PARTITION_ASSIGNMENT_STRATEGY_CONFIG = "partition.assignment.strategy"; - private static final String PARTITION_ASSIGNMENT_STRATEGY_DOC = "The class name of the partition assignment strategy that the client will use to distribute partition ownership amongst consumer instances when group management is used"; + private static final String PARTITION_ASSIGNMENT_STRATEGY_DOC = "A list of class names or class types, ordered by preference, of supported assignors responsible for the partition assignment strategy that the client will use to distribute partition ownership amongst consumer instances when group management is used. Implementing the org.apache.kafka.clients.consumer.ConsumerPartitionAssignor interface allows you to plug in a custom assignment strategy."; /** * auto.offset.reset */ public static final String AUTO_OFFSET_RESET_CONFIG = "auto.offset.reset"; - public static final String AUTO_OFFSET_RESET_DOC = "What to do when there is no initial offset in Kafka or if the current offset does not exist any more on the server (e.g. because that data has been deleted):

  • earliest: automatically reset the offset to the earliest offset
  • latest: automatically reset the offset to the latest offset
  • none: throw exception to the consumer if no previous offset is found for the consumer's group
  • anything else: throw exception to the consumer.
"; + public static final String AUTO_OFFSET_RESET_DOC = "What to do when there is no initial offset in Kafka or if the current offset does not exist any more on the server (e.g. because that data has been deleted):
  • earliest: automatically reset the offset to the earliest offset
  • latest: automatically reset the offset to the latest offset
  • liclosest: uses 'earliest' reset strategy when there is no committed offset or when fetched offset is smaller than Log Start Offset (LSO) and uses 'latest' reset strategy when the fetched offset is greater than Log End Offset (LEO)
  • none: throw exception to the caller if no previous offset is found for the consumer's group
    • <
  • anything else: throw exception to the caller.
"; /** * fetch.min.bytes @@ -166,6 +160,11 @@ public class ConsumerConfig extends AbstractConfig { */ public static final String CLIENT_ID_CONFIG = CommonClientConfigs.CLIENT_ID_CONFIG; + /** + * client.rack + */ + public static final String CLIENT_RACK_CONFIG = CommonClientConfigs.CLIENT_RACK_CONFIG; + /** * reconnect.backoff.ms */ @@ -201,12 +200,23 @@ public class ConsumerConfig extends AbstractConfig { */ public static final String METRIC_REPORTER_CLASSES_CONFIG = CommonClientConfigs.METRIC_REPORTER_CLASSES_CONFIG; + /** + * metrics.replace.on.duplicate + */ + public static final String METRIC_REPLACE_ON_DUPLICATE_CONFIG = CommonClientConfigs.METRICS_REPLACE_ON_DUPLICATE_CONFIG; + /** * check.crcs */ public static final String CHECK_CRCS_CONFIG = "check.crcs"; private static final String CHECK_CRCS_DOC = "Automatically check the CRC32 of the records consumed. This ensures no on-the-wire or on-disk corruption to the messages occurred. This check adds some overhead, so it may be disabled in cases seeking extreme performance."; + /** + * enable.shallow.iterator + */ + public static final String ENABLE_SHALLOW_ITERATOR_CONFIG = "enable.shallow.iterator"; + private static final String ENABLE_SHALLOW_ITERATOR_DOC = "Shallow iterate message batches in the consumer, returning message batches (potentially compressed) instead of individual records"; + /** key.deserializer */ public static final String KEY_DESERIALIZER_CLASS_CONFIG = "key.deserializer"; public static final String KEY_DESERIALIZER_CLASS_DOC = "Deserializer class for key that implements the org.apache.kafka.common.serialization.Deserializer interface."; @@ -252,9 +262,9 @@ public class ConsumerConfig extends AbstractConfig { /** isolation.level */ public static final String ISOLATION_LEVEL_CONFIG = "isolation.level"; - public static final String ISOLATION_LEVEL_DOC = "

Controls how to read messages written transactionally. If set to read_committed, consumer.poll() will only return" + + public static final String ISOLATION_LEVEL_DOC = "Controls how to read messages written transactionally. If set to read_committed, consumer.poll() will only return" + " transactional messages which have been committed. If set to read_uncommitted' (the default), consumer.poll() will return all messages, even transactional messages" + - " which have been aborted. Non-transactional messages will be returned unconditionally in either mode.

Messages will always be returned in offset order. Hence, in " + + " which have been aborted. Non-transactional messages will be returned unconditionally in either mode.

Messages will always be returned in offset order. Hence, in " + " read_committed mode, consumer.poll() will only return messages up to the last stable offset (LSO), which is the one less than the offset of the first open transaction." + " In particular any messages appearing after messages belonging to ongoing transactions will be withheld until the relevant transaction has been completed. As a result, read_committed" + " consumers will not be able to read up to the high watermark when there are in flight transactions.

Further, when in read_committed the seekToEnd method will" + @@ -262,6 +272,20 @@ public class ConsumerConfig extends AbstractConfig { public static final String DEFAULT_ISOLATION_LEVEL = IsolationLevel.READ_UNCOMMITTED.toString().toLowerCase(Locale.ROOT); + /** allow.auto.create.topics */ + public static final String ALLOW_AUTO_CREATE_TOPICS_CONFIG = "allow.auto.create.topics"; + private static final String ALLOW_AUTO_CREATE_TOPICS_DOC = "Allow automatic topic creation on the broker when" + + " subscribing to or assigning a topic. A topic being subscribed to will be automatically created only if the" + + " broker allows for it using `auto.create.topics.enable` broker configuration. This configuration must" + + " be set to `false` when using brokers older than 0.11.0"; + public static final boolean DEFAULT_ALLOW_AUTO_CREATE_TOPICS = false; + + /** + * security.providers + */ + public static final String SECURITY_PROVIDERS_CONFIG = SecurityConfig.SECURITY_PROVIDERS_CONFIG; + private static final String SECURITY_PROVIDERS_DOC = SecurityConfig.SECURITY_PROVIDERS_DOC; + static { CONFIG = new ConfigDef().define(BOOTSTRAP_SERVERS_CONFIG, Type.LIST, @@ -278,6 +302,11 @@ public class ConsumerConfig extends AbstractConfig { Importance.MEDIUM, CommonClientConfigs.CLIENT_DNS_LOOKUP_DOC) .define(GROUP_ID_CONFIG, Type.STRING, null, Importance.HIGH, GROUP_ID_DOC) + .define(GROUP_INSTANCE_ID_CONFIG, + Type.STRING, + null, + Importance.MEDIUM, + GROUP_INSTANCE_ID_DOC) .define(SESSION_TIMEOUT_MS_CONFIG, Type.INT, 10000, @@ -316,6 +345,11 @@ public class ConsumerConfig extends AbstractConfig { "", Importance.LOW, CommonClientConfigs.CLIENT_ID_DOC) + .define(CLIENT_RACK_CONFIG, + Type.STRING, + "", + Importance.LOW, + CommonClientConfigs.CLIENT_RACK_DOC) .define(MAX_PARTITION_FETCH_BYTES_CONFIG, Type.INT, DEFAULT_MAX_PARTITION_FETCH_BYTES, @@ -381,6 +415,11 @@ public class ConsumerConfig extends AbstractConfig { true, Importance.LOW, CHECK_CRCS_DOC) + .define(ENABLE_SHALLOW_ITERATOR_CONFIG, + Type.BOOLEAN, + false, + Importance.LOW, + ENABLE_SHALLOW_ITERATOR_DOC) .define(METRICS_SAMPLE_WINDOW_MS_CONFIG, Type.LONG, 30000, @@ -405,6 +444,11 @@ public class ConsumerConfig extends AbstractConfig { new ConfigDef.NonNullValidator(), Importance.LOW, CommonClientConfigs.METRIC_REPORTER_CLASSES_DOC) + .define(METRIC_REPLACE_ON_DUPLICATE_CONFIG, + Type.BOOLEAN, + false, + Importance.LOW, + CommonClientConfigs.METRICS_REPLACE_ON_DUPLICATE_DOC) .define(KEY_DESERIALIZER_CLASS_CONFIG, Type.CLASS, Importance.HIGH, @@ -455,24 +499,38 @@ public class ConsumerConfig extends AbstractConfig { Importance.MEDIUM, EXCLUDE_INTERNAL_TOPICS_DOC) .defineInternal(LEAVE_GROUP_ON_CLOSE_CONFIG, - Type.BOOLEAN, - true, - Importance.LOW) + Type.BOOLEAN, + true, + Importance.LOW) .define(ISOLATION_LEVEL_CONFIG, Type.STRING, DEFAULT_ISOLATION_LEVEL, in(IsolationLevel.READ_COMMITTED.toString().toLowerCase(Locale.ROOT), IsolationLevel.READ_UNCOMMITTED.toString().toLowerCase(Locale.ROOT)), Importance.MEDIUM, ISOLATION_LEVEL_DOC) + .define(ALLOW_AUTO_CREATE_TOPICS_CONFIG, + Type.BOOLEAN, + DEFAULT_ALLOW_AUTO_CREATE_TOPICS, + Importance.MEDIUM, + ALLOW_AUTO_CREATE_TOPICS_DOC) // security support + .define(SECURITY_PROVIDERS_CONFIG, + Type.STRING, + null, + Importance.LOW, + SECURITY_PROVIDERS_DOC) .define(CommonClientConfigs.SECURITY_PROTOCOL_CONFIG, Type.STRING, CommonClientConfigs.DEFAULT_SECURITY_PROTOCOL, Importance.MEDIUM, CommonClientConfigs.SECURITY_PROTOCOL_DOC) + .define(CommonClientConfigs.ENABLE_STICKY_METADATA_FETCH_CONFIG, + Type.BOOLEAN, + true, + Importance.MEDIUM, + CommonClientConfigs.ENABLE_STICKY_METADATA_FETCH_DOC) .withClientSslSupport() .withClientSaslSupport(); - } @Override @@ -519,8 +577,12 @@ public static Set configNames() { return CONFIG.names(); } + public static ConfigDef configDef() { + return new ConfigDef(CONFIG); + } + public static void main(String[] args) { - System.out.println(CONFIG.toHtmlTable()); + System.out.println(CONFIG.toHtml()); } } diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/ConsumerGroupMetadata.java b/clients/src/main/java/org/apache/kafka/clients/consumer/ConsumerGroupMetadata.java new file mode 100644 index 0000000000000..e17894b213d0e --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/ConsumerGroupMetadata.java @@ -0,0 +1,50 @@ +/* + * 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; + +import java.util.Optional; + +public class ConsumerGroupMetadata { + private String groupId; + private int generationId; + private String memberId; + Optional groupInstanceId; + + public ConsumerGroupMetadata(String groupId, int generationId, String memberId, Optional groupInstanceId) { + this.groupId = groupId; + this.generationId = generationId; + this.memberId = memberId; + this.groupInstanceId = groupInstanceId; + } + + public String groupId() { + return groupId; + } + + public int generationId() { + return generationId; + } + + public String memberId() { + return memberId; + } + + public Optional groupInstanceId() { + return groupInstanceId; + } + +} diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/ConsumerPartitionAssignor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/ConsumerPartitionAssignor.java new file mode 100644 index 0000000000000..8708ea4f7e343 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/ConsumerPartitionAssignor.java @@ -0,0 +1,233 @@ +/* + * 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; + +import java.nio.ByteBuffer; +import java.util.Optional; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.apache.kafka.common.Cluster; +import org.apache.kafka.common.TopicPartition; + +/** + * This interface is used to define custom partition assignment for use in + * {@link org.apache.kafka.clients.consumer.KafkaConsumer}. Members of the consumer group subscribe + * to the topics they are interested in and forward their subscriptions to a Kafka broker serving + * as the group coordinator. The coordinator selects one member to perform the group assignment and + * propagates the subscriptions of all members to it. Then {@link #assign(Cluster, GroupSubscription)} is called + * to perform the assignment and the results are forwarded back to each respective members + * + * In some cases, it is useful to forward additional metadata to the assignor in order to make + * assignment decisions. For this, you can override {@link #subscriptionUserData(Set)} and provide custom + * userData in the returned Subscription. For example, to have a rack-aware assignor, an implementation + * can use this user data to forward the rackId belonging to each member. + */ +public interface ConsumerPartitionAssignor { + + /** + * Return serialized data that will be included in the {@link Subscription} sent to the leader + * and can be leveraged in {@link #assign(Cluster, GroupSubscription)} ((e.g. local host/rack information) + * + * @param topics Topics subscribed to through {@link org.apache.kafka.clients.consumer.KafkaConsumer#subscribe(java.util.Collection)} + * and variants + * @return nullable subscription user data + */ + default ByteBuffer subscriptionUserData(Set topics) { + return null; + } + + /** + * Perform the group assignment given the member subscriptions and current cluster metadata. + * @param metadata Current topic/broker metadata known by consumer + * @param groupSubscription Subscriptions from all members including metadata provided through {@link #subscriptionUserData(Set)} + * @return A map from the members to their respective assignments. This should have one entry + * for each member in the input subscription map. + */ + GroupAssignment assign(Cluster metadata, GroupSubscription groupSubscription); + + /** + * Callback which is invoked when a group member receives its assignment from the leader. + * @param assignment The local member's assignment as provided by the leader in {@link #assign(Cluster, GroupSubscription)} + * @param metadata Additional metadata on the consumer (optional) + */ + default void onAssignment(Assignment assignment, ConsumerGroupMetadata metadata) { + } + + /** + * Indicate which rebalance protocol this assignor works with; + * By default it should always work with {@link RebalanceProtocol#EAGER}. + */ + default List supportedProtocols() { + return Collections.singletonList(RebalanceProtocol.EAGER); + } + + /** + * Return the version of the assignor which indicates how the user metadata encodings + * and the assignment algorithm gets evolved. + */ + default short version() { + return (short) 0; + } + + /** + * Unique name for this assignor (e.g. "range" or "roundrobin" or "sticky"). Note, this is not required + * to be the same as the class name specified in {@link ConsumerConfig#PARTITION_ASSIGNMENT_STRATEGY_CONFIG} + * @return non-null unique name + */ + String name(); + + final class Subscription { + private final List topics; + private final ByteBuffer userData; + private final List ownedPartitions; + private Optional groupInstanceId; + + public Subscription(List topics, ByteBuffer userData, List ownedPartitions) { + this.topics = topics; + this.userData = userData; + this.ownedPartitions = ownedPartitions; + this.groupInstanceId = Optional.empty(); + } + + public Subscription(List topics, ByteBuffer userData) { + this(topics, userData, Collections.emptyList()); + } + + public Subscription(List topics) { + this(topics, null, Collections.emptyList()); + } + + public List topics() { + return topics; + } + + public ByteBuffer userData() { + return userData; + } + + public List ownedPartitions() { + return ownedPartitions; + } + + public void setGroupInstanceId(Optional groupInstanceId) { + this.groupInstanceId = groupInstanceId; + } + + public Optional groupInstanceId() { + return groupInstanceId; + } + } + + final class Assignment { + private List partitions; + private ByteBuffer userData; + + public Assignment(List partitions, ByteBuffer userData) { + this.partitions = partitions; + this.userData = userData; + } + + public Assignment(List partitions) { + this(partitions, null); + } + + public List partitions() { + return partitions; + } + + public ByteBuffer userData() { + return userData; + } + + @Override + public String toString() { + return "Assignment(" + + "partitions=" + partitions + + (userData == null ? "" : ", userDataSize=" + userData.remaining()) + + ')'; + } + } + + final class GroupSubscription { + private final Map subscriptions; + + public GroupSubscription(Map subscriptions) { + this.subscriptions = subscriptions; + } + + public Map groupSubscription() { + return subscriptions; + } + } + + final class GroupAssignment { + private final Map assignments; + + public GroupAssignment(Map assignments) { + this.assignments = assignments; + } + + public Map groupAssignment() { + return assignments; + } + } + + /** + * The rebalance protocol defines partition assignment and revocation semantics. The purpose is to establish a + * consistent set of rules that all consumers in a group follow in order to transfer ownership of a partition. + * {@link ConsumerPartitionAssignor} implementors can claim supporting one or more rebalance protocols via the + * {@link ConsumerPartitionAssignor#supportedProtocols()}, and it is their responsibility to respect the rules + * of those protocols in their {@link ConsumerPartitionAssignor#assign(Cluster, GroupSubscription)} implementations. + * Failures to follow the rules of the supported protocols would lead to runtime error or undefined behavior. + * + * The {@link RebalanceProtocol#EAGER} rebalance protocol requires a consumer to always revoke all its owned + * partitions before participating in a rebalance event. It therefore allows a complete reshuffling of the assignment. + * + * {@link RebalanceProtocol#COOPERATIVE} rebalance protocol allows a consumer to retain its currently owned + * partitions before participating in a rebalance event. The assignor should not reassign any owned partitions + * immediately, but instead may indicate consumers the need for partition revocation so that the revoked + * partitions can be reassigned to other consumers in the next rebalance event. This is designed for sticky assignment + * logic which attempts to minimize partition reassignment with cooperative adjustments. + */ + enum RebalanceProtocol { + EAGER((byte) 0), COOPERATIVE((byte) 1); + + private final byte id; + + RebalanceProtocol(byte id) { + this.id = id; + } + + public byte id() { + return id; + } + + public static RebalanceProtocol forId(byte id) { + switch (id) { + case 0: + return EAGER; + case 1: + return COOPERATIVE; + default: + throw new IllegalArgumentException("Unknown rebalance protocol id: " + id); + } + } + } + +} diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/ConsumerRebalanceListener.java b/clients/src/main/java/org/apache/kafka/clients/consumer/ConsumerRebalanceListener.java index 74e8b060c73fa..2f43b603fc8ff 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/ConsumerRebalanceListener.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/ConsumerRebalanceListener.java @@ -16,6 +16,7 @@ */ package org.apache.kafka.clients.consumer; +import java.time.Duration; import java.util.Collection; import org.apache.kafka.common.TopicPartition; @@ -29,7 +30,7 @@ *

* When Kafka is managing the group membership, a partition re-assignment will be triggered any time the members of the group change or the subscription * of the members changes. This can occur when processes die, new process instances are added or old instances come back to life after failure. - * Rebalances can also be triggered by changes affecting the subscribed topics (e.g. when the number of partitions is + * Partition re-assignments can also be triggered by changes affecting the subscribed topics (e.g. when the number of partitions is * administratively adjusted). *

* There are many uses for this functionality. One common use is saving offsets in a custom store. By saving offsets in @@ -47,11 +48,45 @@ * This callback will only execute in the user thread as part of the {@link Consumer#poll(java.time.Duration) poll(long)} call * whenever partition assignment changes. *

- * It is guaranteed that all consumer processes will invoke {@link #onPartitionsRevoked(Collection) onPartitionsRevoked} prior to - * any process invoking {@link #onPartitionsAssigned(Collection) onPartitionsAssigned}. So if offsets or other state is saved in the - * {@link #onPartitionsRevoked(Collection) onPartitionsRevoked} call it is guaranteed to be saved by the time the process taking over that - * partition has their {@link #onPartitionsAssigned(Collection) onPartitionsAssigned} callback called to load the state. + * Under normal conditions, if a partition is reassigned from one consumer to another, then the old consumer will + * always invoke {@link #onPartitionsRevoked(Collection) onPartitionsRevoked} for that partition prior to the new consumer + * invoking {@link #onPartitionsAssigned(Collection) onPartitionsAssigned} for the same partition. So if offsets or other state is saved in the + * {@link #onPartitionsRevoked(Collection) onPartitionsRevoked} call by one consumer member, it will be always accessible by the time the + * other consumer member taking over that partition and triggering its {@link #onPartitionsAssigned(Collection) onPartitionsAssigned} callback to load the state. *

+ * You can think of revocation as a graceful way to give up ownership of a partition. In some cases, the consumer may not have an opportunity to do so. + * For example, if the session times out, then the partitions may be reassigned before we have a chance to revoke them gracefully. + * For this case, we have a third callback {@link #onPartitionsLost(Collection)}. The difference between this function and + * {@link #onPartitionsRevoked(Collection)} is that upon invocation of {@link #onPartitionsLost(Collection)}, the partitions + * may already be owned by some other members in the group and therefore users would not be able to commit its consumed offsets for example. + * Users could implement these two functions differently (by default, + * {@link #onPartitionsLost(Collection)} will be calling {@link #onPartitionsRevoked(Collection)} directly); for example, in the + * {@link #onPartitionsLost(Collection)} we should not need to store the offsets since we know these partitions are no longer owned by the consumer + * at that time. + *

+ * During a rebalance event, the {@link #onPartitionsAssigned(Collection) onPartitionsAssigned} function will always be triggered exactly once when + * the rebalance completes. That is, even if there is no newly assigned partitions for a consumer member, its {@link #onPartitionsAssigned(Collection) onPartitionsAssigned} + * will still be triggered with an empty collection of partitions. As a result this function can be used also to notify when a rebalance event has happened. + * With eager rebalancing, {@link #onPartitionsRevoked(Collection)} will always be called at the start of a rebalance. On the other hand, {@link #onPartitionsLost(Collection)} + * will only be called when there were non-empty partitions that were lost. + * With cooperative rebalancing, {@link #onPartitionsRevoked(Collection)} and {@link #onPartitionsLost(Collection)} + * will only be triggered when there are non-empty partitions revoked or lost from this consumer member during a rebalance event. + *

+ * It is possible + * for a {@link org.apache.kafka.common.errors.WakeupException} or {@link org.apache.kafka.common.errors.InterruptException} + * to be raised from one of these nested invocations. In this case, the exception will be propagated to the current + * invocation of {@link KafkaConsumer#poll(java.time.Duration)} in which this callback is being executed. This means it is not + * necessary to catch these exceptions and re-attempt to wakeup or interrupt the consumer thread. + * Also if the callback function implementation itself throws an exception, this exception will be propagated to the current + * invocation of {@link KafkaConsumer#poll(java.time.Duration)} as well. + *

+ * Note that callbacks only serve as notification of an assignment change. + * They cannot be used to express acceptance of the change. + * Hence throwing an exception from a callback does not affect the assignment in any way, + * as it will be propagated all the way up to the {@link KafkaConsumer#poll(java.time.Duration)} call. + * If user captures the exception in the caller, the callback is still assumed successful and no further retries will be attempted. + *

+ * * Here is pseudo-code for a callback implementation for saving offsets: * 

  * {@code
@@ -68,6 +103,10 @@
  *              saveOffsetInExternalStore(consumer.position(partition));
  *       }
  *
+ *       public void onPartitionsLost(Collection partitions) {
+ *           // do not need to save the offsets since these partitions are probably owned by other consumers already
+ *       }
+ *
  *       public void onPartitionsAssigned(Collection partitions) {
  *           // read the offsets from an external store using some custom code not described here
  *           for(TopicPartition partition: partitions)
@@ -80,22 +119,25 @@
 public interface ConsumerRebalanceListener {
 
     /**
-     * A callback method the user can implement to provide handling of offset commits to a customized store on the start
-     * of a rebalance operation. This method will be called before a rebalance operation starts and after the consumer
-     * stops fetching data. It is recommended that offsets should be committed in this callback to either Kafka or a
+     * A callback method the user can implement to provide handling of offset commits to a customized store.
+     * This method will be called during a rebalance operation when the consumer has to give up some partitions.
+     * It can also be called when consumer is being closed ({@link KafkaConsumer#close(Duration)})
+     * or is unsubscribing ({@link KafkaConsumer#unsubscribe()}).
+     * It is recommended that offsets should be committed in this callback to either Kafka or a
      * custom offset store to prevent duplicate data.
      * 

-     * For examples on usage of this API, see Usage Examples section of {@link KafkaConsumer KafkaConsumer}
-     * 

-     * NOTE: This method is only called before rebalances. It is not called prior to {@link KafkaConsumer#close()}.
+     * In eager rebalancing, it will always be called at the start of a rebalance and after the consumer stops fetching data.
+     * In cooperative rebalancing, it will be called at the end of a rebalance on the set of partitions being revoked iff the set is non-empty.
+     * For examples on usage of this API, see Usage Examples section of {@link KafkaConsumer KafkaConsumer}.
      * 

      * It is common for the revocation callback to use the consumer instance in order to commit offsets. It is possible
      * for a {@link org.apache.kafka.common.errors.WakeupException} or {@link org.apache.kafka.common.errors.InterruptException}
-     * to be raised from one these nested invocations. In this case, the exception will be propagated to the current
+     * to be raised from one of these nested invocations. In this case, the exception will be propagated to the current
      * invocation of {@link KafkaConsumer#poll(java.time.Duration)} in which this callback is being executed. This means it is not
      * necessary to catch these exceptions and re-attempt to wakeup or interrupt the consumer thread.
      *
-     * @param partitions The list of partitions that were assigned to the consumer on the last rebalance
+     * @param partitions The list of partitions that were assigned to the consumer and now need to be revoked (may not
+     *                   include all currently assigned partitions, i.e. there may still be some partitions left)
      * @throws org.apache.kafka.common.errors.WakeupException If raised from a nested call to {@link KafkaConsumer}
      * @throws org.apache.kafka.common.errors.InterruptException If raised from a nested call to {@link KafkaConsumer}
      */
@@ -106,20 +148,53 @@ public interface ConsumerRebalanceListener {
      * partition re-assignment. This method will be called after the partition re-assignment completes and before the
      * consumer starts fetching data, and only as the result of a {@link Consumer#poll(java.time.Duration) poll(long)} call.
      * 

-     * It is guaranteed that all the processes in a consumer group will execute their
+     * It is guaranteed that under normal conditions all the processes in a consumer group will execute their
      * {@link #onPartitionsRevoked(Collection)} callback before any instance executes its
-     * {@link #onPartitionsAssigned(Collection)} callback.
+     * {@link #onPartitionsAssigned(Collection)} callback. During exceptional scenarios, partitions may be migrated
+     * without the old owner being notified (i.e. their {@link #onPartitionsRevoked(Collection)} callback not triggered),
+     * and later when the old owner consumer realized this event, the {@link #onPartitionsLost(Collection)} (Collection)} callback
+     * will be triggered by the consumer then.
      * 

      * It is common for the assignment callback to use the consumer instance in order to query offsets. It is possible
      * for a {@link org.apache.kafka.common.errors.WakeupException} or {@link org.apache.kafka.common.errors.InterruptException}
-     * to be raised from one these nested invocations. In this case, the exception will be propagated to the current
+     * to be raised from one of these nested invocations. In this case, the exception will be propagated to the current
      * invocation of {@link KafkaConsumer#poll(java.time.Duration)} in which this callback is being executed. This means it is not
      * necessary to catch these exceptions and re-attempt to wakeup or interrupt the consumer thread.
      *
-     * @param partitions The list of partitions that are now assigned to the consumer (may include partitions previously
-     *            assigned to the consumer)
+     * @param partitions The list of partitions that are now assigned to the consumer (previously owned partitions will
+     *                   NOT be included, i.e. this list will only include newly added partitions)
      * @throws org.apache.kafka.common.errors.WakeupException If raised from a nested call to {@link KafkaConsumer}
      * @throws org.apache.kafka.common.errors.InterruptException If raised from a nested call to {@link KafkaConsumer}
      */
     void onPartitionsAssigned(Collection partitions);
+
+    /**
+     * A callback method you can implement to provide handling of cleaning up resources for partitions that have already
+     * been reassigned to other consumers. This method will not be called during normal execution as the owned partitions would
+     * first be revoked by calling the {@link ConsumerRebalanceListener#onPartitionsRevoked}, before being reassigned
+     * to other consumers during a rebalance event. However, during exceptional scenarios when the consumer realized that it
+     * does not own this partition any longer, i.e. not revoked via a normal rebalance event, then this method would be invoked.
+     * 

+     * For example, this function is called if a consumer's session timeout has expired, or if a fatal error has been
+     * received indicating the consumer is no longer part of the group.
+     * 

+     * By default it will just trigger {@link ConsumerRebalanceListener#onPartitionsRevoked}; for users who want to distinguish
+     * the handling logic of revoked partitions v.s. lost partitions, they can override the default implementation.
+     * 

+     * It is possible
+     * for a {@link org.apache.kafka.common.errors.WakeupException} or {@link org.apache.kafka.common.errors.InterruptException}
+     * to be raised from one of these nested invocations. In this case, the exception will be propagated to the current
+     * invocation of {@link KafkaConsumer#poll(java.time.Duration)} in which this callback is being executed. This means it is not
+     * necessary to catch these exceptions and re-attempt to wakeup or interrupt the consumer thread.
+     *
+     * @param partitions The list of partitions that were assigned to the consumer and now have been reassigned
+     *                   to other consumers. With the current protocol this will always include all of the consumer's
+     *                   previously assigned partitions, but this may change in future protocols (ie there would still
+     *                   be some partitions left)
+     * @throws org.apache.kafka.common.errors.WakeupException If raised from a nested call to {@link KafkaConsumer}
+     * @throws org.apache.kafka.common.errors.InterruptException If raised from a nested call to {@link KafkaConsumer}
+     */
+    default void onPartitionsLost(Collection partitions) {
+        onPartitionsRevoked(partitions);
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/ConsumerRecord.java b/clients/src/main/java/org/apache/kafka/clients/consumer/ConsumerRecord.java
index 0413d5b07768e..a7dad7bc5e2a9 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/ConsumerRecord.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/ConsumerRecord.java
@@ -157,6 +157,8 @@ public ConsumerRecord(String topic,
                           Optional leaderEpoch) {
         if (topic == null)
             throw new IllegalArgumentException("Topic cannot be null");
+        if (headers == null)
+            throw new IllegalArgumentException("Headers cannot be null");
 
         this.topic = topic;
         this.partition = partition;
@@ -173,7 +175,7 @@ public ConsumerRecord(String topic,
     }
 
     /**
-     * The topic this record is received from
+     * The topic this record is received from (never null)
      */
     public String topic() {
         return this.topic;
@@ -187,7 +189,7 @@ public int partition() {
     }
 
     /**
-     * The headers
+     * The headers (never null)
      */
     public Headers headers() {
         return headers;
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/CooperativeStickyAssignor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/CooperativeStickyAssignor.java
new file mode 100644
index 0000000000000..c7c0679575a9b
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/CooperativeStickyAssignor.java
@@ -0,0 +1,109 @@
+/*
+ * 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;
+
+import java.util.Arrays;
+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 org.apache.kafka.clients.consumer.internals.AbstractStickyAssignor;
+import org.apache.kafka.common.TopicPartition;
+
+/**
+ * A cooperative version of the {@link AbstractStickyAssignor AbstractStickyAssignor}. This follows the same (sticky)
+ * assignment logic as {@link StickyAssignor StickyAssignor} but allows for cooperative rebalancing while the
+ * {@link StickyAssignor StickyAssignor} follows the eager rebalancing protocol. See
+ * {@link ConsumerPartitionAssignor.RebalanceProtocol} for an explanation of the rebalancing protocols.
+ * 

+ * Users should prefer this assignor for newer clusters.
+ * 

+ * To turn on cooperative rebalancing you must set all your consumers to use this {@code PartitionAssignor},
+ * or implement a custom one that returns {@code RebalanceProtocol.COOPERATIVE} in
+ * {@link CooperativeStickyAssignor#supportedProtocols supportedProtocols()}.
+ * 

+ * IMPORTANT: if upgrading from 2.3 or earlier, you must follow a specific upgrade path in order to safely turn on
+ * cooperative rebalancing. See the upgrade guide for details.
+ */
+public class CooperativeStickyAssignor extends AbstractStickyAssignor {
+
+    @Override
+    public String name() {
+        return "cooperative-sticky";
+    }
+
+    @Override
+    public List supportedProtocols() {
+        return Arrays.asList(RebalanceProtocol.COOPERATIVE, RebalanceProtocol.EAGER);
+    }
+
+    @Override
+    protected MemberData memberData(Subscription subscription) {
+        return new MemberData(subscription.ownedPartitions(), Optional.empty());
+    }
+
+    @Override
+    public Map> assign(Map partitionsPerTopic,
+                                                    Map subscriptions) {
+        Map> assignments = super.assign(partitionsPerTopic, subscriptions);
+
+        Map partitionsTransferringOwnership = super.partitionsTransferringOwnership == null ?
+            computePartitionsTransferringOwnership(subscriptions, assignments) :
+            super.partitionsTransferringOwnership;
+
+        adjustAssignment(assignments, partitionsTransferringOwnership);
+        return assignments;
+    }
+
+    // Following the cooperative rebalancing protocol requires removing partitions that must first be revoked from the assignment
+    private void adjustAssignment(Map> assignments,
+                                  Map partitionsTransferringOwnership) {
+        for (Map.Entry partitionEntry : partitionsTransferringOwnership.entrySet()) {
+            assignments.get(partitionEntry.getValue()).remove(partitionEntry.getKey());
+        }
+    }
+
+    private Map computePartitionsTransferringOwnership(Map subscriptions,
+                                                                               Map> assignments) {
+        Map allAddedPartitions = new HashMap<>();
+        Set allRevokedPartitions = new HashSet<>();
+
+        for (final Map.Entry> entry : assignments.entrySet()) {
+            String consumer = entry.getKey();
+
+            List ownedPartitions = subscriptions.get(consumer).ownedPartitions();
+            List assignedPartitions = entry.getValue();
+
+            Set ownedPartitionsSet = new HashSet<>(ownedPartitions);
+            for (TopicPartition tp : assignedPartitions) {
+                if (!ownedPartitionsSet.contains(tp))
+                    allAddedPartitions.put(tp, consumer);
+            }
+
+            Set assignedPartitionsSet = new HashSet<>(assignedPartitions);
+            for (TopicPartition tp : ownedPartitions) {
+                if (!assignedPartitionsSet.contains(tp))
+                    allRevokedPartitions.add(tp);
+            }
+        }
+
+        allAddedPartitions.keySet().retainAll(allRevokedPartitions);
+        return allAddedPartitions;
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/KafkaConsumer.java b/clients/src/main/java/org/apache/kafka/clients/consumer/KafkaConsumer.java
index 4cee56a23959e..7fd7d2c3bf52f 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/KafkaConsumer.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/KafkaConsumer.java
@@ -16,9 +16,14 @@
  */
 package org.apache.kafka.clients.consumer;
 
+import static org.apache.kafka.clients.consumer.internals.PartitionAssignorAdapter.getAssignorInstances;
+
 import org.apache.kafka.clients.ApiVersions;
 import org.apache.kafka.clients.ClientDnsLookup;
 import org.apache.kafka.clients.ClientUtils;
+import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.clients.GroupRebalanceConfig;
+import org.apache.kafka.clients.Metadata;
 import org.apache.kafka.clients.NetworkClient;
 import org.apache.kafka.clients.consumer.internals.ConsumerCoordinator;
 import org.apache.kafka.clients.consumer.internals.ConsumerInterceptors;
@@ -26,9 +31,8 @@
 import org.apache.kafka.clients.consumer.internals.ConsumerNetworkClient;
 import org.apache.kafka.clients.consumer.internals.Fetcher;
 import org.apache.kafka.clients.consumer.internals.FetcherMetricsRegistry;
-import org.apache.kafka.clients.consumer.internals.Heartbeat;
+import org.apache.kafka.clients.consumer.internals.KafkaConsumerMetrics;
 import org.apache.kafka.clients.consumer.internals.NoOpConsumerRebalanceListener;
-import org.apache.kafka.clients.consumer.internals.PartitionAssignor;
 import org.apache.kafka.clients.consumer.internals.SubscriptionState;
 import org.apache.kafka.common.Cluster;
 import org.apache.kafka.common.KafkaException;
@@ -69,6 +73,7 @@
 import java.util.Locale;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.Properties;
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
@@ -561,6 +566,7 @@ public class KafkaConsumer implements Consumer {
 
     // Visible for testing
     final Metrics metrics;
+    final KafkaConsumerMetrics kafkaConsumerMetrics;
 
     private final Logger log;
     private final String clientId;
@@ -579,7 +585,7 @@ public class KafkaConsumer implements Consumer {
     private final long requestTimeoutMs;
     private final int defaultApiTimeoutMs;
     private volatile boolean closed = false;
-    private List assignors;
+    private List assignors;
 
     // currentThread holds the threadId of the current thread accessing KafkaConsumer
     // and is used to prevent multi-threaded access
@@ -664,12 +670,22 @@ public KafkaConsumer(Properties properties,
     @SuppressWarnings("unchecked")
     private KafkaConsumer(ConsumerConfig config, Deserializer keyDeserializer, Deserializer valueDeserializer) {
         try {
-            String clientId = config.getString(ConsumerConfig.CLIENT_ID_CONFIG);
-            if (clientId.isEmpty())
-                clientId = "consumer-" + CONSUMER_CLIENT_ID_SEQUENCE.getAndIncrement();
-            this.clientId = clientId;
-            this.groupId = config.getString(ConsumerConfig.GROUP_ID_CONFIG);
-            LogContext logContext = new LogContext("[Consumer clientId=" + clientId + ", groupId=" + groupId + "] ");
+            GroupRebalanceConfig groupRebalanceConfig = new GroupRebalanceConfig(config,
+                    GroupRebalanceConfig.ProtocolType.CONSUMER);
+
+            this.groupId = groupRebalanceConfig.groupId;
+            this.clientId = buildClientId(config.getString(CommonClientConfigs.CLIENT_ID_CONFIG), groupRebalanceConfig);
+
+            LogContext logContext;
+
+            // If group.instance.id is set, we will append it to the log context.
+            if (groupRebalanceConfig.groupInstanceId.isPresent()) {
+                logContext = new LogContext("[Consumer instanceId=" + groupRebalanceConfig.groupInstanceId.get() +
+                        ", clientId=" + clientId + ", groupId=" + groupId + "] ");
+            } else {
+                logContext = new LogContext("[Consumer clientId=" + clientId + ", groupId=" + groupId + "] ");
+            }
+
             this.log = logContext.logger(getClass());
             boolean enableAutoCommit = config.getBoolean(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG);
             if (groupId == null) { // overwrite in case of default group id where the config is not explicitly provided
@@ -685,6 +701,7 @@ else if (enableAutoCommit)
             this.defaultApiTimeoutMs = config.getInt(ConsumerConfig.DEFAULT_API_TIMEOUT_MS_CONFIG);
             this.time = Time.SYSTEM;
             this.metrics = buildMetrics(config, time, clientId);
+            this.metrics.setReplaceOnDuplicateMetric(config.getBoolean(ConsumerConfig.METRIC_REPLACE_ON_DUPLICATE_CONFIG));
             this.retryBackoffMs = config.getLong(ConsumerConfig.RETRY_BACKOFF_MS_CONFIG);
 
             // load interceptors and make sure they get clientId
@@ -714,10 +731,11 @@ else if (enableAutoCommit)
             this.metadata = new ConsumerMetadata(retryBackoffMs,
                     config.getLong(ConsumerConfig.METADATA_MAX_AGE_CONFIG),
                     !config.getBoolean(ConsumerConfig.EXCLUDE_INTERNAL_TOPICS_CONFIG),
+                    config.getBoolean(ConsumerConfig.ALLOW_AUTO_CREATE_TOPICS_CONFIG),
                     subscriptions, logContext, clusterResourceListeners);
             List addresses = ClientUtils.parseAndValidateAddresses(
                     config.getList(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG), config.getString(ConsumerConfig.CLIENT_DNS_LOOKUP_CONFIG));
-            this.metadata.bootstrap(addresses, time.milliseconds());
+            this.metadata.bootstrap(addresses);
             String metricGrpPrefix = "consumer";
 
             FetcherMetricsRegistry metricsRegistry = new FetcherMetricsRegistry(Collections.singleton(CLIENT_ID_METRIC_TAG), metricGrpPrefix);
@@ -727,6 +745,7 @@ else if (enableAutoCommit)
             Sensor throttleTimeSensor = Fetcher.throttleTimeSensor(metrics, metricsRegistry);
             int heartbeatIntervalMs = config.getInt(ConsumerConfig.HEARTBEAT_INTERVAL_MS_CONFIG);
 
+            ApiVersions apiVersions = new ApiVersions();
             NetworkClient netClient = new NetworkClient(
                     new Selector(config.getLong(ConsumerConfig.CONNECTIONS_MAX_IDLE_MS_CONFIG), metrics, time, metricGrpPrefix, channelBuilder, logContext),
                     this.metadata,
@@ -740,9 +759,10 @@ else if (enableAutoCommit)
                     ClientDnsLookup.forConfig(config.getString(ConsumerConfig.CLIENT_DNS_LOOKUP_CONFIG)),
                     time,
                     true,
-                    new ApiVersions(),
+                    apiVersions,
                     throttleTimeSensor,
                     logContext);
+            netClient.setEnableStickyMetadataFetch(config.getBoolean(CommonClientConfigs.ENABLE_STICKY_METADATA_FETCH_CONFIG));
             this.client = new ConsumerNetworkClient(
                     logContext,
                     netClient,
@@ -751,31 +771,23 @@ else if (enableAutoCommit)
                     retryBackoffMs,
                     config.getInt(ConsumerConfig.REQUEST_TIMEOUT_MS_CONFIG),
                     heartbeatIntervalMs); //Will avoid blocking an extended period of time to prevent heartbeat thread starvation
-            this.assignors = config.getConfiguredInstances(
-                    ConsumerConfig.PARTITION_ASSIGNMENT_STRATEGY_CONFIG,
-                    PartitionAssignor.class);
 
-            int maxPollIntervalMs = config.getInt(ConsumerConfig.MAX_POLL_INTERVAL_MS_CONFIG);
-            int sessionTimeoutMs = config.getInt(ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG);
+            this.assignors = getAssignorInstances(config.getList(ConsumerConfig.PARTITION_ASSIGNMENT_STRATEGY_CONFIG), config.originals());
+
             // no coordinator will be constructed for the default (null) group id
             this.coordinator = groupId == null ? null :
-                new ConsumerCoordinator(logContext,
+                new ConsumerCoordinator(groupRebalanceConfig,
+                        logContext,
                         this.client,
-                        groupId,
-                        maxPollIntervalMs,
-                        sessionTimeoutMs,
-                        new Heartbeat(time, sessionTimeoutMs, heartbeatIntervalMs, maxPollIntervalMs, retryBackoffMs),
                         assignors,
                         this.metadata,
                         this.subscriptions,
                         metrics,
                         metricGrpPrefix,
                         this.time,
-                        retryBackoffMs,
                         enableAutoCommit,
                         config.getInt(ConsumerConfig.AUTO_COMMIT_INTERVAL_MS_CONFIG),
-                        this.interceptors,
-                        config.getBoolean(ConsumerConfig.LEAVE_GROUP_ON_CLOSE_CONFIG));
+                        this.interceptors);
             this.fetcher = new Fetcher<>(
                     logContext,
                     this.client,
@@ -785,6 +797,8 @@ else if (enableAutoCommit)
                     config.getInt(ConsumerConfig.MAX_PARTITION_FETCH_BYTES_CONFIG),
                     config.getInt(ConsumerConfig.MAX_POLL_RECORDS_CONFIG),
                     config.getBoolean(ConsumerConfig.CHECK_CRCS_CONFIG),
+                    config.getString(ConsumerConfig.CLIENT_RACK_CONFIG),
+                    config.getBoolean(ConsumerConfig.ENABLE_SHALLOW_ITERATOR_CONFIG),
                     this.keyDeserializer,
                     this.valueDeserializer,
                     this.metadata,
@@ -794,10 +808,13 @@ else if (enableAutoCommit)
                     this.time,
                     this.retryBackoffMs,
                     this.requestTimeoutMs,
-                    isolationLevel);
+                    isolationLevel,
+                    apiVersions);
+
+            this.kafkaConsumerMetrics = new KafkaConsumerMetrics(metrics, metricGrpPrefix);
 
             config.logUnused();
-            AppInfoParser.registerAppInfo(JMX_PREFIX, clientId, metrics);
+            AppInfoParser.registerAppInfo(JMX_PREFIX, clientId, metrics, time.milliseconds());
             log.debug("Kafka consumer initialized");
         } catch (Throwable t) {
             // call close methods if internal objects are already constructed; this is to prevent resource leak. see KAFKA-2121
@@ -823,7 +840,7 @@ else if (enableAutoCommit)
                   long retryBackoffMs,
                   long requestTimeoutMs,
                   int defaultApiTimeoutMs,
-                  List assignors,
+                  List assignors,
                   String groupId) {
         this.log = logContext.logger(getClass());
         this.clientId = clientId;
@@ -842,6 +859,18 @@ else if (enableAutoCommit)
         this.defaultApiTimeoutMs = defaultApiTimeoutMs;
         this.assignors = assignors;
         this.groupId = groupId;
+        this.kafkaConsumerMetrics = new KafkaConsumerMetrics(metrics, "consumer");
+    }
+
+    private static String buildClientId(String configuredClientId, GroupRebalanceConfig rebalanceConfig) {
+        if (!configuredClientId.isEmpty())
+            return configuredClientId;
+
+        if (rebalanceConfig.groupId != null && !rebalanceConfig.groupId.isEmpty())
+            return "consumer-" + rebalanceConfig.groupId + "-" + rebalanceConfig.groupInstanceId.orElseGet(() ->
+                    CONSUMER_CLIENT_ID_SEQUENCE.getAndIncrement() + "");
+
+        return "consumer-" + CONSUMER_CLIENT_ID_SEQUENCE.getAndIncrement();
     }
 
     private static Metrics buildMetrics(ConsumerConfig config, Time time, String clientId) {
@@ -867,7 +896,7 @@ private static Metrics buildMetrics(ConsumerConfig config, Time time, String cli
     public Set assignment() {
         acquireAndEnsureOpen();
         try {
-            return Collections.unmodifiableSet(new HashSet<>(this.subscriptions.assignedPartitions()));
+            return Collections.unmodifiableSet(this.subscriptions.assignedPartitions());
         } finally {
             release();
         }
@@ -1036,14 +1065,18 @@ public void subscribe(Pattern pattern) {
     /**
      * Unsubscribe from topics currently subscribed with {@link #subscribe(Collection)} or {@link #subscribe(Pattern)}.
      * This also clears any partitions directly assigned through {@link #assign(Collection)}.
+     *
+     * @throws org.apache.kafka.common.KafkaException for any other unrecoverable errors (e.g. rebalance callback errors)
      */
     public void unsubscribe() {
         acquireAndEnsureOpen();
         try {
             fetcher.clearBufferedDataForUnassignedPartitions(Collections.emptySet());
+            if (this.coordinator != null) {
+                this.coordinator.onLeavePrepare();
+                this.coordinator.maybeLeaveGroup("the consumer unsubscribed from all topics");
+            }
             this.subscriptions.unsubscribe();
-            if (this.coordinator != null)
-                this.coordinator.maybeLeaveGroup();
             log.info("Unsubscribed all topics or patterns and assigned partitions");
         } finally {
             release();
@@ -1127,7 +1160,7 @@ public void assign(Collection partitions) {
      * @throws java.lang.IllegalArgumentException if the timeout value is negative
      * @throws java.lang.IllegalStateException if the consumer is not subscribed to any topics or manually assigned any
      *             partitions to consume from
-     *
+     * @throws org.apache.kafka.common.errors.FencedInstanceIdException if this consumer instance gets fenced by broker.
      *
      * @deprecated Since 2.0. Use {@link #poll(Duration)}, which does not block beyond the timeout awaiting partition
      *             assignment. See KIP-266 for more information.
@@ -1166,22 +1199,29 @@ public ConsumerRecords poll(final long timeoutMs) {
      * @throws org.apache.kafka.common.errors.AuthorizationException if caller lacks Read access to any of the subscribed
      *             topics or to the configured groupId. See the exception for more details
      * @throws org.apache.kafka.common.KafkaException for any other unrecoverable errors (e.g. invalid groupId or
-     *             session timeout, errors deserializing key/value pairs, or any new error cases in future versions)
+     *             session timeout, errors deserializing key/value pairs, your rebalance callback thrown exceptions,
+     *             or any new error cases in future versions)
      * @throws java.lang.IllegalArgumentException if the timeout value is negative
      * @throws java.lang.IllegalStateException if the consumer is not subscribed to any topics or manually assigned any
      *             partitions to consume from
      * @throws java.lang.ArithmeticException if the timeout is greater than {@link Long#MAX_VALUE} milliseconds.
      * @throws org.apache.kafka.common.errors.InvalidTopicException if the current subscription contains any invalid
      *             topic (per {@link org.apache.kafka.common.internals.Topic#validate(String)})
+     * @throws org.apache.kafka.common.errors.FencedInstanceIdException if this consumer instance gets fenced by broker.
      */
     @Override
     public ConsumerRecords poll(final Duration timeout) {
         return poll(time.timer(timeout), true);
     }
 
+    /**
+     * @throws KafkaException if the rebalance callback throws exception
+     */
     private ConsumerRecords poll(final Timer timer, final boolean includeMetadataInTimeout) {
         acquireAndEnsureOpen();
         try {
+            this.kafkaConsumerMetrics.recordPollStart(timer.currentTimeMs());
+
             if (this.subscriptions.hasNoSubscriptionOrUserAssignment()) {
                 throw new IllegalStateException("Consumer is not subscribed to any topics or assigned any partitions");
             }
@@ -1209,7 +1249,7 @@ private ConsumerRecords poll(final Timer timer, final boolean includeMetad
                     // NOTE: since the consumed position has already been updated, we must not allow
                     // wakeups or any other errors to be triggered prior to returning the fetched records.
                     if (fetcher.sendFetches() > 0 || client.hasPendingRequests()) {
-                        client.pollNoWakeup();
+                        client.transmitSends();
                     }
 
                     return this.interceptors.onConsume(new ConsumerRecords<>(records));
@@ -1219,6 +1259,7 @@ private ConsumerRecords poll(final Timer timer, final boolean includeMetad
             return ConsumerRecords.empty();
         } finally {
             release();
+            this.kafkaConsumerMetrics.recordPollEnd(timer.currentTimeMs());
         }
     }
 
@@ -1233,6 +1274,9 @@ boolean updateAssignmentMetadataIfNeeded(final Timer timer) {
         return updateFetchPositions(timer);
     }
 
+    /**
+     * @throws KafkaException if the rebalance callback throws exception
+     */
     private Map>> pollForFetches(Timer timer) {
         long pollTimeout = coordinator == null ? timer.remainingMs() :
                 Math.min(coordinator.timeToNextPoll(timer.currentTimeMs()), timer.remainingMs());
@@ -1259,7 +1303,7 @@ private Map>> pollForFetches(Timer tim
         client.poll(pollTimer, () -> {
             // since a fetch might be completed by the background thread, we need this poll condition
             // to ensure that we do not block unnecessarily in poll()
-            return !fetcher.hasCompletedFetches();
+            return !fetcher.hasAvailableFetches();
         });
         timer.update(pollTimer.currentTimeMs());
 
@@ -1301,6 +1345,7 @@ private Map>> pollForFetches(Timer tim
      *             is too large or if the topic does not exist).
      * @throws org.apache.kafka.common.errors.TimeoutException if the timeout specified by {@code default.api.timeout.ms} expires
      *            before successful completion of the offset commit
+     * @throws org.apache.kafka.common.errors.FencedInstanceIdException if this consumer instance gets fenced by broker.
      */
     @Override
     public void commitSync() {
@@ -1335,6 +1380,7 @@ public void commitSync() {
      *             is too large or if the topic does not exist).
      * @throws org.apache.kafka.common.errors.TimeoutException if the timeout expires before successful completion
      *            of the offset commit
+     * @throws org.apache.kafka.common.errors.FencedInstanceIdException if this consumer instance gets fenced by broker.
      */
     @Override
     public void commitSync(Duration timeout) {
@@ -1381,6 +1427,7 @@ public void commitSync(Duration timeout) {
      *             is too large or if the topic does not exist).
      * @throws org.apache.kafka.common.errors.TimeoutException if the timeout expires before successful completion
      *            of the offset commit
+     * @throws org.apache.kafka.common.errors.FencedInstanceIdException if this consumer instance gets fenced by broker.
      */
     @Override
     public void commitSync(final Map offsets) {
@@ -1418,6 +1465,7 @@ public void commitSync(final Map offsets) {
     *             is too large or if the topic does not exist).
     * @throws org.apache.kafka.common.errors.TimeoutException if the timeout expires before successful completion
     *            of the offset commit
+    * @throws org.apache.kafka.common.errors.FencedInstanceIdException if this consumer instance gets fenced by broker.
     */
     @Override
     public void commitSync(final Map offsets, final Duration timeout) {
@@ -1437,6 +1485,7 @@ public void commitSync(final Map offsets, fin
     /**
      * Commit offsets returned on the last {@link #poll(Duration)} for all the subscribed list of topics and partition.
      * Same as {@link #commitAsync(OffsetCommitCallback) commitAsync(null)}
+     * @throws org.apache.kafka.common.errors.FencedInstanceIdException if this consumer instance gets fenced by broker.
      */
     @Override
     public void commitAsync() {
@@ -1459,6 +1508,7 @@ public void commitAsync() {
      * (and variants) returns.
      *
      * @param callback Callback to invoke when the commit completes
+     * @throws org.apache.kafka.common.errors.FencedInstanceIdException if this consumer instance gets fenced by broker.
      */
     @Override
     public void commitAsync(OffsetCommitCallback callback) {
@@ -1484,6 +1534,7 @@ public void commitAsync(OffsetCommitCallback callback) {
      * @param offsets A map of offsets by partition with associate metadata. This map will be copied internally, so it
      *                is safe to mutate the map after returning.
      * @param callback Callback to invoke when the commit completes
+     * @throws org.apache.kafka.common.errors.FencedInstanceIdException if this consumer instance gets fenced by broker.
      */
     @Override
     public void commitAsync(final Map offsets, OffsetCommitCallback callback) {
@@ -1508,7 +1559,20 @@ public void commitAsync(final Map offsets, Of
      */
     @Override
     public void seek(TopicPartition partition, long offset) {
-        seek(partition, new OffsetAndMetadata(offset, null));
+        if (offset < 0)
+            throw new IllegalArgumentException("seek offset must not be a negative number");
+
+        acquireAndEnsureOpen();
+        try {
+            log.debug("Seeking to offset {} for partition {}", offset, partition);
+            SubscriptionState.FetchPosition newPosition = new SubscriptionState.FetchPosition(
+                    offset,
+                    Optional.empty(), // This will ensure we skip validation
+                    this.metadata.leaderAndEpoch(partition));
+            this.subscriptions.seekUnvalidated(partition, newPosition);
+        } finally {
+            release();
+        }
     }
 
     /**
@@ -1530,13 +1594,18 @@ public void seek(TopicPartition partition, OffsetAndMetadata offsetAndMetadata)
         acquireAndEnsureOpen();
         try {
             if (offsetAndMetadata.leaderEpoch().isPresent()) {
-                log.info("Seeking to offset {} for partition {} with epoch {}",
+                log.debug("Seeking to offset {} for partition {} with epoch {}",
                         offset, partition, offsetAndMetadata.leaderEpoch().get());
             } else {
-                log.info("Seeking to offset {} for partition {}", offset, partition);
+                log.debug("Seeking to offset {} for partition {}", offset, partition);
             }
+            Metadata.LeaderAndEpoch currentLeaderAndEpoch = this.metadata.leaderAndEpoch(partition);
+            SubscriptionState.FetchPosition newPosition = new SubscriptionState.FetchPosition(
+                    offsetAndMetadata.offset(),
+                    offsetAndMetadata.leaderEpoch(),
+                    currentLeaderAndEpoch);
             this.updateLastSeenEpochIfNewer(partition, offsetAndMetadata);
-            this.subscriptions.seek(partition, offset);
+            this.subscriptions.seekUnvalidated(partition, newPosition);
         } finally {
             release();
         }
@@ -1558,10 +1627,7 @@ public void seekToBeginning(Collection partitions) {
         acquireAndEnsureOpen();
         try {
             Collection parts = partitions.size() == 0 ? this.subscriptions.assignedPartitions() : partitions;
-            for (TopicPartition tp : parts) {
-                log.info("Seeking to beginning of partition {}", tp);
-                subscriptions.requestOffsetReset(tp, OffsetResetStrategy.EARLIEST);
-            }
+            subscriptions.requestOffsetReset(parts, OffsetResetStrategy.EARLIEST);
         } finally {
             release();
         }
@@ -1586,10 +1652,7 @@ public void seekToEnd(Collection partitions) {
         acquireAndEnsureOpen();
         try {
             Collection parts = partitions.size() == 0 ? this.subscriptions.assignedPartitions() : partitions;
-            for (TopicPartition tp : parts) {
-                log.info("Seeking to end of partition {}", tp);
-                subscriptions.requestOffsetReset(tp, OffsetResetStrategy.LATEST);
-            }
+            subscriptions.requestOffsetReset(parts, OffsetResetStrategy.LATEST);
         } finally {
             release();
         }
@@ -1658,9 +1721,9 @@ public long position(TopicPartition partition, final Duration timeout) {
 
             Timer timer = time.timer(timeout);
             do {
-                Long offset = this.subscriptions.position(partition);
-                if (offset != null)
-                    return offset;
+                SubscriptionState.FetchPosition position = this.subscriptions.validPosition(partition);
+                if (position != null)
+                    return position.offset;
 
                 updateFetchPositions(timer);
                 client.poll(timer);
@@ -1694,7 +1757,10 @@ public long position(TopicPartition partition, final Duration timeout) {
      * @throws org.apache.kafka.common.KafkaException for any other unrecoverable errors
      * @throws org.apache.kafka.common.errors.TimeoutException if the committed offset cannot be found before
      *             the timeout specified by {@code default.api.timeout.ms} expires.
+     *
+     * @deprecated since 2.4 Use {@link #committed(Set)} instead
      */
+    @Deprecated
     @Override
     public OffsetAndMetadata committed(TopicPartition partition) {
         return committed(partition, Duration.ofMillis(defaultApiTimeoutMs));
@@ -1704,7 +1770,8 @@ public OffsetAndMetadata committed(TopicPartition partition) {
      * Get the last committed offset for the given partition (whether the commit happened by this process or
      * another). This offset will be used as the position for the consumer in the event of a failure.
      * 

-     * This call will block to do a remote call to get the latest committed offsets from the server.
+     * This call will block until the position can be determined, an unrecoverable error is
+     * encountered (in which case it is thrown to the caller), or the timeout expires.
      *
      * @param partition The partition to check
      * @param timeout  The maximum amount of time to await the current committed offset
@@ -1719,20 +1786,81 @@ public OffsetAndMetadata committed(TopicPartition partition) {
      * @throws org.apache.kafka.common.KafkaException for any other unrecoverable errors
      * @throws org.apache.kafka.common.errors.TimeoutException if the committed offset cannot be found before
      *             expiration of the timeout
+     *
+     * @deprecated since 2.4 Use {@link #committed(Set, Duration)} instead
      */
+    @Deprecated
     @Override
     public OffsetAndMetadata committed(TopicPartition partition, final Duration timeout) {
+        return committed(Collections.singleton(partition), timeout).get(partition);
+    }
+
+    /**
+     * Get the last committed offsets for the given partitions (whether the commit happened by this process or
+     * another). The returned offsets will be used as the position for the consumer in the event of a failure.
+     * 

+     * If any of the partitions requested do not exist, an exception would be thrown.
+     * 

+     * This call will do a remote call to get the latest committed offsets from the server, and will block until the
+     * committed offsets are gotten successfully, an unrecoverable error is encountered (in which case it is thrown to
+     * the caller), or the timeout specified by {@code default.api.timeout.ms} expires (in which case a
+     * {@link org.apache.kafka.common.errors.TimeoutException} is thrown to the caller).
+     *
+     * @param partitions The partitions to check
+     * @return The latest committed offsets for the given partitions; {@code null} will be returned for the
+     *         partition if there is no such message.
+     * @throws org.apache.kafka.common.errors.WakeupException if {@link #wakeup()} is called before or while this
+     *             function is called
+     * @throws org.apache.kafka.common.errors.InterruptException if the calling thread is interrupted before or while
+     *             this function is called
+     * @throws org.apache.kafka.common.errors.AuthenticationException if authentication fails. See the exception for more details
+     * @throws org.apache.kafka.common.errors.AuthorizationException if not authorized to the topic or to the
+     *             configured groupId. See the exception for more details
+     * @throws org.apache.kafka.common.KafkaException for any other unrecoverable errors
+     * @throws org.apache.kafka.common.errors.TimeoutException if the committed offset cannot be found before
+     *             the timeout specified by {@code default.api.timeout.ms} expires.
+     */
+    @Override
+    public Map committed(final Set partitions) {
+        return committed(partitions, Duration.ofMillis(defaultApiTimeoutMs));
+    }
+
+    /**
+     * Get the last committed offsets for the given partitions (whether the commit happened by this process or
+     * another). The returned offsets will be used as the position for the consumer in the event of a failure.
+     * 

+     * If any of the partitions requested do not exist, an exception would be thrown.
+     * 

+     * This call will block to do a remote call to get the latest committed offsets from the server.
+     *
+     * @param partitions The partitions to check
+     * @param timeout  The maximum amount of time to await the latest committed offsets
+     * @return The latest committed offsets for the given partitions; {@code null} will be returned for the
+     *         partition if there is no such message.
+     * @throws org.apache.kafka.common.errors.WakeupException if {@link #wakeup()} is called before or while this
+     *             function is called
+     * @throws org.apache.kafka.common.errors.InterruptException if the calling thread is interrupted before or while
+     *             this function is called
+     * @throws org.apache.kafka.common.errors.AuthenticationException if authentication fails. See the exception for more details
+     * @throws org.apache.kafka.common.errors.AuthorizationException if not authorized to the topic or to the
+     *             configured groupId. See the exception for more details
+     * @throws org.apache.kafka.common.KafkaException for any other unrecoverable errors
+     * @throws org.apache.kafka.common.errors.TimeoutException if the committed offset cannot be found before
+     *             expiration of the timeout
+     */
+    @Override
+    public Map committed(final Set partitions, final Duration timeout) {
         acquireAndEnsureOpen();
         try {
             maybeThrowInvalidGroupIdException();
-            Map offsets = coordinator.fetchCommittedOffsets(
-                    Collections.singleton(partition), time.timer(timeout));
+            Map offsets = coordinator.fetchCommittedOffsets(partitions, time.timer(timeout));
             if (offsets == null) {
                 throw new TimeoutException("Timeout of " + timeout.toMillis() + "ms expired before the last " +
-                        "committed offset for partition " + partition + " could be determined");
+                    "committed offset for partitions " + partitions + " could be determined. Try tuning default.api.timeout.ms " +
+                    "larger to relax the threshold.");
             } else {
                 offsets.forEach(this::updateLastSeenEpochIfNewer);
-                return offsets.get(partition);
+                return offsets;
             }
         } finally {
             release();
@@ -1799,7 +1927,7 @@ public List partitionsFor(String topic, Duration timeout) {
 
             Timer timer = time.timer(timeout);
             Map> topicMetadata = fetcher.getTopicMetadata(
-                    new MetadataRequest.Builder(Collections.singletonList(topic), true), timer);
+                    new MetadataRequest.Builder(Collections.singletonList(topic), metadata.allowAutoTopicCreation()), timer);
             return topicMetadata.get(topic);
         } finally {
             release();
@@ -2131,10 +2259,12 @@ public void close(Duration timeout) {
         acquire();
         try {
             if (!closed) {
-                closed = true;
+                // need to close before setting the flag since the close function
+                // itself may trigger rebalance callback that needs the consumer to be open still
                 close(timeout.toMillis(), false);
             }
         } finally {
+            closed = true;
             release();
         }
     }
@@ -2169,12 +2299,13 @@ private void close(long timeoutMs, boolean swallowException) {
             firstException.compareAndSet(null, t);
             log.error("Failed to close coordinator", t);
         }
-        ClientUtils.closeQuietly(fetcher, "fetcher", firstException);
-        ClientUtils.closeQuietly(interceptors, "consumer interceptors", firstException);
-        ClientUtils.closeQuietly(metrics, "consumer metrics", firstException);
-        ClientUtils.closeQuietly(client, "consumer network client", firstException);
-        ClientUtils.closeQuietly(keyDeserializer, "consumer key deserializer", firstException);
-        ClientUtils.closeQuietly(valueDeserializer, "consumer value deserializer", firstException);
+        Utils.closeQuietly(fetcher, "fetcher", firstException);
+        Utils.closeQuietly(interceptors, "consumer interceptors", firstException);
+        Utils.closeQuietly(kafkaConsumerMetrics, "kafka consumer metrics", firstException);
+        Utils.closeQuietly(metrics, "consumer metrics", firstException);
+        Utils.closeQuietly(client, "consumer network client", firstException);
+        Utils.closeQuietly(keyDeserializer, "consumer key deserializer", firstException);
+        Utils.closeQuietly(valueDeserializer, "consumer value deserializer", firstException);
         AppInfoParser.unregisterAppInfo(JMX_PREFIX, clientId, metrics);
         log.debug("Kafka consumer has been closed");
         Throwable exception = firstException.get();
@@ -2196,6 +2327,9 @@ private void close(long timeoutMs, boolean swallowException) {
      * @return true iff the operation completed without timing out
      */
     private boolean updateFetchPositions(final Timer timer) {
+        // If any partitions have been truncated due to a leader change, we need to validate the offsets
+        fetcher.validateOffsetsIfNeeded();
+
         cachedSubscriptionHashAllFetchPositions = subscriptions.hasAllFetchPositions();
         if (cachedSubscriptionHashAllFetchPositions) return true;
 
@@ -2264,7 +2398,8 @@ private void maybeThrowInvalidGroupIdException() {
     }
 
     private void updateLastSeenEpochIfNewer(TopicPartition topicPartition, OffsetAndMetadata offsetAndMetadata) {
-        offsetAndMetadata.leaderEpoch().ifPresent(epoch -> metadata.updateLastSeenEpochIfNewer(topicPartition, epoch));
+        if (offsetAndMetadata != null)
+            offsetAndMetadata.leaderEpoch().ifPresent(epoch -> metadata.updateLastSeenEpochIfNewer(topicPartition, epoch));
     }
 
     // Visible for testing
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/LogTruncationException.java b/clients/src/main/java/org/apache/kafka/clients/consumer/LogTruncationException.java
new file mode 100644
index 0000000000000..f8af50d5c902f
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/LogTruncationException.java
@@ -0,0 +1,50 @@
+/*
+ * 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;
+
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.utils.Utils;
+
+import java.util.Collections;
+import java.util.Map;
+import java.util.function.Function;
+
+/**
+ * In the even of unclean leader election, the log will be truncated,
+ * previously committed data will be lost, and new data will be written
+ * over these offsets. When this happens, the consumer will detect the
+ * truncation and raise this exception (if no automatic reset policy
+ * has been defined) with the first offset to diverge from what the
+ * consumer read.
+ */
+public class LogTruncationException extends OffsetOutOfRangeException {
+
+    private final Map divergentOffsets;
+
+    public LogTruncationException(Map divergentOffsets) {
+        super(Utils.transformMap(divergentOffsets, Function.identity(), OffsetAndMetadata::offset));
+        this.divergentOffsets = Collections.unmodifiableMap(divergentOffsets);
+    }
+
+    /**
+     * Get the offsets for the partitions which were truncated. This is the first offset which is known to diverge
+     * from what the consumer read.
+     */
+    public Map divergentOffsets() {
+        return divergentOffsets;
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/MockConsumer.java b/clients/src/main/java/org/apache/kafka/clients/consumer/MockConsumer.java
index 614ec9b4457b0..0db883ed1b413 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/MockConsumer.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/MockConsumer.java
@@ -16,11 +16,13 @@
  */
 package org.apache.kafka.clients.consumer;
 
+import org.apache.kafka.clients.Metadata;
 import org.apache.kafka.clients.consumer.internals.NoOpConsumerRebalanceListener;
 import org.apache.kafka.clients.consumer.internals.SubscriptionState;
 import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.Metric;
 import org.apache.kafka.common.MetricName;
+import org.apache.kafka.common.Node;
 import org.apache.kafka.common.PartitionInfo;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.errors.WakeupException;
@@ -40,6 +42,7 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.regex.Pattern;
+import java.util.stream.Collectors;
 
 
 /**
@@ -59,7 +62,8 @@ public class MockConsumer implements Consumer {
     private final Set paused;
 
     private Map>> records;
-    private KafkaException exception;
+    private KafkaException pollException;
+    private KafkaException offsetsException;
     private AtomicBoolean wakeup;
     private boolean closed;
 
@@ -72,7 +76,7 @@ public MockConsumer(OffsetResetStrategy offsetResetStrategy) {
         this.beginningOffsets = new HashMap<>();
         this.endOffsets = new HashMap<>();
         this.pollTasks = new LinkedList<>();
-        this.exception = null;
+        this.pollException = null;
         this.wakeup = new AtomicBoolean(false);
         this.committed = new HashMap<>();
     }
@@ -171,9 +175,9 @@ public synchronized ConsumerRecords poll(final Duration timeout) {
             throw new WakeupException();
         }
 
-        if (exception != null) {
-            RuntimeException exception = this.exception;
-            this.exception = null;
+        if (pollException != null) {
+            RuntimeException exception = this.pollException;
+            this.pollException = null;
             throw exception;
         }
 
@@ -189,13 +193,17 @@ public synchronized ConsumerRecords poll(final Duration timeout) {
             if (!subscriptions.isPaused(entry.getKey())) {
                 final List> recs = entry.getValue();
                 for (final ConsumerRecord rec : recs) {
-                    if (beginningOffsets.get(entry.getKey()) != null && beginningOffsets.get(entry.getKey()) > subscriptions.position(entry.getKey())) {
-                        throw new OffsetOutOfRangeException(Collections.singletonMap(entry.getKey(), subscriptions.position(entry.getKey())));
+                    long position = subscriptions.position(entry.getKey()).offset;
+
+                    if (beginningOffsets.get(entry.getKey()) != null && beginningOffsets.get(entry.getKey()) > position) {
+                        throw new OffsetOutOfRangeException(Collections.singletonMap(entry.getKey(), position));
                     }
 
-                    if (assignment().contains(entry.getKey()) && rec.offset() >= subscriptions.position(entry.getKey())) {
+                    if (assignment().contains(entry.getKey()) && rec.offset() >= position) {
                         results.computeIfAbsent(entry.getKey(), partition -> new ArrayList<>()).add(rec);
-                        subscriptions.position(entry.getKey(), rec.offset() + 1);
+                        SubscriptionState.FetchPosition newPosition = new SubscriptionState.FetchPosition(
+                                rec.offset() + 1, rec.leaderEpoch(), new Metadata.LeaderAndEpoch(Node.noNode(), rec.leaderEpoch()));
+                        subscriptions.position(entry.getKey(), newPosition);
                     }
                 }
             }
@@ -207,15 +215,27 @@ public synchronized ConsumerRecords poll(final Duration timeout) {
     public synchronized void addRecord(ConsumerRecord record) {
         ensureNotClosed();
         TopicPartition tp = new TopicPartition(record.topic(), record.partition());
-        Set currentAssigned = new HashSet<>(this.subscriptions.assignedPartitions());
+        Set currentAssigned = this.subscriptions.assignedPartitions();
         if (!currentAssigned.contains(tp))
             throw new IllegalStateException("Cannot add records for a partition that is not assigned to the consumer");
         List> recs = this.records.computeIfAbsent(tp, k -> new ArrayList<>());
         recs.add(record);
     }
 
+    /**
+     * @deprecated Use {@link #setPollException(KafkaException)} instead
+     */
+    @Deprecated
     public synchronized void setException(KafkaException exception) {
-        this.exception = exception;
+        setPollException(exception);
+    }
+
+    public synchronized void setPollException(KafkaException exception) {
+        this.pollException = exception;
+    }
+
+    public synchronized void setOffsetsException(KafkaException exception) {
+        this.offsetsException = exception;
     }
 
     @Override
@@ -271,31 +291,44 @@ public void seek(TopicPartition partition, OffsetAndMetadata offsetAndMetadata)
         subscriptions.seek(partition, offsetAndMetadata.offset());
     }
 
+    @Deprecated
     @Override
-    public synchronized OffsetAndMetadata committed(TopicPartition partition) {
-        ensureNotClosed();
-        if (subscriptions.isAssigned(partition)) {
-            return committed.get(partition);
-        }
-        return new OffsetAndMetadata(0);
+    public synchronized OffsetAndMetadata committed(final TopicPartition partition) {
+        return committed(Collections.singleton(partition)).get(partition);
     }
 
+    @Deprecated
     @Override
-    public OffsetAndMetadata committed(TopicPartition partition, final Duration timeout) {
+    public OffsetAndMetadata committed(final TopicPartition partition, final Duration timeout) {
         return committed(partition);
     }
 
+    @Override
+    public synchronized Map committed(final Set partitions) {
+        ensureNotClosed();
+
+        return partitions.stream()
+            .filter(committed::containsKey)
+            .collect(Collectors.toMap(tp -> tp, tp -> subscriptions.isAssigned(tp) ?
+                committed.get(tp) : new OffsetAndMetadata(0)));
+    }
+
+    @Override
+    public synchronized Map committed(final Set partitions, final Duration timeout) {
+        return committed(partitions);
+    }
+
     @Override
     public synchronized long position(TopicPartition partition) {
         ensureNotClosed();
         if (!this.subscriptions.isAssigned(partition))
             throw new IllegalArgumentException("You can only check the position for partitions assigned to this consumer.");
-        Long offset = this.subscriptions.position(partition);
-        if (offset == null) {
+        SubscriptionState.FetchPosition position = this.subscriptions.position(partition);
+        if (position == null) {
             updateFetchPosition(partition);
-            offset = this.subscriptions.position(partition);
+            position = this.subscriptions.position(partition);
         }
-        return offset;
+        return position.offset;
     }
 
     @Override
@@ -306,8 +339,7 @@ public synchronized long position(TopicPartition partition, final Duration timeo
     @Override
     public synchronized void seekToBeginning(Collection partitions) {
         ensureNotClosed();
-        for (TopicPartition tp : partitions)
-            subscriptions.requestOffsetReset(tp, OffsetResetStrategy.EARLIEST);
+        subscriptions.requestOffsetReset(partitions, OffsetResetStrategy.EARLIEST);
     }
 
     public synchronized void updateBeginningOffsets(Map newOffsets) {
@@ -317,8 +349,7 @@ public synchronized void updateBeginningOffsets(Map newOff
     @Override
     public synchronized void seekToEnd(Collection partitions) {
         ensureNotClosed();
-        for (TopicPartition tp : partitions)
-            subscriptions.requestOffsetReset(tp, OffsetResetStrategy.LATEST);
+        subscriptions.requestOffsetReset(partitions, OffsetResetStrategy.LATEST);
     }
 
     // needed for cases where you make a second call to endOffsets
@@ -389,6 +420,11 @@ public synchronized Map offsetsForTimes(Map<
 
     @Override
     public synchronized Map beginningOffsets(Collection partitions) {
+        if (offsetsException != null) {
+            RuntimeException exception = this.offsetsException;
+            this.offsetsException = null;
+            throw exception;
+        }
         Map result = new HashMap<>();
         for (TopicPartition tp : partitions) {
             Long beginningOffset = beginningOffsets.get(tp);
@@ -401,6 +437,11 @@ public synchronized Map beginningOffsets(Collection endOffsets(Collection partitions) {
+        if (offsetsException != null) {
+            RuntimeException exception = this.offsetsException;
+            this.offsetsException = null;
+            throw exception;
+        }
         Map result = new HashMap<>();
         for (TopicPartition tp : partitions) {
             Long endOffset = getEndOffset(endOffsets.get(tp));
@@ -419,7 +460,6 @@ public synchronized void close() {
     @SuppressWarnings("deprecation")
     @Override
     public synchronized void close(long timeout, TimeUnit unit) {
-        ensureNotClosed();
         this.closed = true;
     }
 
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/OffsetResetStrategy.java b/clients/src/main/java/org/apache/kafka/clients/consumer/OffsetResetStrategy.java
index 6d742b850a134..1c2ddafcb3d0c 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/OffsetResetStrategy.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/OffsetResetStrategy.java
@@ -17,5 +17,5 @@
 package org.apache.kafka.clients.consumer;
 
 public enum OffsetResetStrategy {
-    LATEST, EARLIEST, NONE
+    LATEST, EARLIEST, NONE, LICLOSEST
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/RangeAssignor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/RangeAssignor.java
index 78956abb2182d..ea642885c0623 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/RangeAssignor.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/RangeAssignor.java
@@ -40,6 +40,29 @@
  * 
  • C0: [t0p0, t0p1, t1p0, t1p1]
    • 
  * 
  • C1: [t0p2, t1p2]
    • 
  * 
+ *
+ * Since the introduction of static membership, we could leverage group.instance.id to make the assignment behavior more sticky.
+ * For the above example, after one rolling bounce, group coordinator will attempt to assign new member.id towards consumers,
+ * for example C0 -> C3 C1 -> C2.
+ *
+ * 
    The assignment could be completely shuffled to:
+ * 
      
+ * 
    • C3 (was C0): [t0p2, t1p2] (before was [t0p0, t0p1, t1p0, t1p1])
+ * 
    • C2 (was C1): [t0p0, t0p1, t1p0, t1p1] (before was [t0p2, t1p2])
+ * 
    
+ *
+ * The assignment change was caused by the change of member.id relative order, and
+ * can be avoided by setting the group.instance.id.
+ * Consumers will have individual instance ids I1, I2. As long as
+ * 1. Number of members remain the same across generation
+ * 2. Static members' identities persist across generation
+ * 3. Subscription pattern doesn't change for any member
+ *
+ * 
    The assignment will always be:
+ * 
      
+ * 
    • I0: [t0p0, t0p1, t1p0, t1p1]
+ * 
    • I1: [t0p2, t1p2]
+ * 
    
  */
 public class RangeAssignor extends AbstractPartitionAssignor {
 
@@ -48,27 +71,30 @@ public String name() {
         return "range";
     }
 
-    private Map> consumersPerTopic(Map consumerMetadata) {
-        Map> res = new HashMap<>();
+    private Map> consumersPerTopic(Map consumerMetadata) {
+        Map> topicToConsumers = new HashMap<>();
         for (Map.Entry subscriptionEntry : consumerMetadata.entrySet()) {
             String consumerId = subscriptionEntry.getKey();
-            for (String topic : subscriptionEntry.getValue().topics())
-                put(res, topic, consumerId);
+            MemberInfo memberInfo = new MemberInfo(consumerId, subscriptionEntry.getValue().groupInstanceId());
+            for (String topic : subscriptionEntry.getValue().topics()) {
+                put(topicToConsumers, topic, memberInfo);
+            }
         }
-        return res;
+        return topicToConsumers;
     }
 
     @Override
     public Map> assign(Map partitionsPerTopic,
                                                     Map subscriptions) {
-        Map> consumersPerTopic = consumersPerTopic(subscriptions);
+        Map> consumersPerTopic = consumersPerTopic(subscriptions);
+
         Map> assignment = new HashMap<>();
         for (String memberId : subscriptions.keySet())
             assignment.put(memberId, new ArrayList<>());
 
-        for (Map.Entry> topicEntry : consumersPerTopic.entrySet()) {
+        for (Map.Entry> topicEntry : consumersPerTopic.entrySet()) {
             String topic = topicEntry.getKey();
-            List consumersForTopic = topicEntry.getValue();
+            List consumersForTopic = topicEntry.getValue();
 
             Integer numPartitionsForTopic = partitionsPerTopic.get(topic);
             if (numPartitionsForTopic == null)
@@ -83,10 +109,9 @@ public Map> assign(Map partitionsP
             for (int i = 0, n = consumersForTopic.size(); i < n; i++) {
                 int start = numPartitionsPerConsumer * i + Math.min(i, consumersWithExtraPartition);
                 int length = numPartitionsPerConsumer + (i + 1 > consumersWithExtraPartition ? 0 : 1);
-                assignment.get(consumersForTopic.get(i)).addAll(partitions.subList(start, start + length));
+                assignment.get(consumersForTopic.get(i).memberId).addAll(partitions.subList(start, start + length));
             }
         }
         return assignment;
     }
-
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/RoundRobinAssignor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/RoundRobinAssignor.java
index 02cdc458fc4b4..0ac90b386834e 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/RoundRobinAssignor.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/RoundRobinAssignor.java
@@ -34,8 +34,9 @@
  * instances are identical, then the partitions will be uniformly distributed. (i.e., the partition ownership counts
  * will be within a delta of exactly one across all consumers.)
  *
- * 
    For example, suppose there are two consumers C0 and C1, two topics t0 and t1, and each topic has 3 partitions,
- * resulting in partitions t0p0, t0p1, t0p2, t1p0, t1p1, and t1p2.
+ * 
    For example, suppose there are two consumers C0 and C1, two topics t0 and t1,
+ * and each topic has 3 partitions, resulting in partitions t0p0, t0p1, t0p2,
+ * t1p0, t1p1, and t1p2.
  *
  * 
    The assignment will be:
  * 
      
@@ -46,9 +47,12 @@
  * 
      When subscriptions differ across consumer instances, the assignment process still considers each
  * consumer instance in round robin fashion but skips over an instance if it is not subscribed to
  * the topic. Unlike the case when subscriptions are identical, this can result in imbalanced
- * assignments. For example, we have three consumers C0, C1, C2, and three topics t0, t1, t2,
- * with 1, 2, and 3 partitions, respectively. Therefore, the partitions are t0p0, t1p0, t1p1, t2p0,
- * t2p1, t2p2. C0 is subscribed to t0; C1 is subscribed to t0, t1; and C2 is subscribed to t0, t1, t2.
+ * assignments. For example, we have three consumers C0, C1, C2,
+ * and three topics t0, t1, t2, with 1, 2, and 3 partitions, respectively.
+ * Therefore, the partitions are t0p0, t1p0, t1p1, t2p0, t2p1, t2p2.
+ * C0 is subscribed to t0;
+ * C1 is subscribed to t0, t1;
+ * and C2 is subscribed to t0, t1, t2.
  *
  * 
      That assignment will be:
  * 
        
@@ -56,6 +60,42 @@
  * 
      • C1: [t1p0]
  * 
      • C2: [t1p1, t2p0, t2p1, t2p2]
  * 
      
+ *
+ * Since the introduction of static membership, we could leverage group.instance.id to make the assignment behavior more sticky.
+ * For example, we have three consumers with assigned member.id C0, C1, C2,
+ * two topics t0 and t1, and each topic has 3 partitions, resulting in partitions t0p0,
+ * t0p1, t0p2, t1p0, t1p1, and t1p2. We choose to honor
+ * the sorted order based on ephemeral member.id.
+ *
+ * 
      The assignment will be:
+ * 
        
+ * 
      • C0: [t0p0, t1p0]
+ * 
      • C1: [t0p1, t1p1]
+ * 
      • C2: [t0p2, t1p2]
+ * 
      
+ *
+ * After one rolling bounce, group coordinator will attempt to assign new member.id towards consumers,
+ * for example C0 -> C5 C1 -> C3, C2 -> C4.
+ *
+ * 
      The assignment could be completely shuffled to:
+ * 
        
+ * 
      • C3 (was C1): [t0p0, t1p0] (before was [t0p1, t1p1])
+ * 
      • C4 (was C2): [t0p1, t1p1] (before was [t0p2, t1p2])
+ * 
      • C5 (was C0): [t0p2, t1p2] (before was [t0p0, t1p0])
+ * 
      
+ *
+ * This issue could be mitigated by the introduction of static membership. Consumers will have individual instance ids
+ * I1, I2, I3. As long as
+ * 1. Number of members remain the same across generation
+ * 2. Static members' identities persist across generation
+ * 3. Subscription pattern doesn't change for any member
+ *
+ * 
      The assignment will always be:
+ * 
        
+ * 
      • I0: [t0p0, t1p0]
+ * 
      • I1: [t0p1, t1p1]
+ * 
      • I2: [t0p2, t1p2]
+ * 
      
  */
 public class RoundRobinAssignor extends AbstractPartitionAssignor {
 
@@ -63,21 +103,26 @@ public class RoundRobinAssignor extends AbstractPartitionAssignor {
     public Map> assign(Map partitionsPerTopic,
                                                     Map subscriptions) {
         Map> assignment = new HashMap<>();
-        for (String memberId : subscriptions.keySet())
-            assignment.put(memberId, new ArrayList<>());
+        List memberInfoList = new ArrayList<>();
+        for (Map.Entry memberSubscription : subscriptions.entrySet()) {
+            assignment.put(memberSubscription.getKey(), new ArrayList<>());
+            memberInfoList.add(new MemberInfo(memberSubscription.getKey(),
+                                              memberSubscription.getValue().groupInstanceId()));
+        }
+
+        CircularIterator assigner = new CircularIterator<>(Utils.sorted(memberInfoList));
 
-        CircularIterator assigner = new CircularIterator<>(Utils.sorted(subscriptions.keySet()));
         for (TopicPartition partition : allPartitionsSorted(partitionsPerTopic, subscriptions)) {
             final String topic = partition.topic();
-            while (!subscriptions.get(assigner.peek()).topics().contains(topic))
+            while (!subscriptions.get(assigner.peek().memberId).topics().contains(topic))
                 assigner.next();
-            assignment.get(assigner.next()).add(partition);
+            assignment.get(assigner.next().memberId).add(partition);
         }
         return assignment;
     }
 
-    public List allPartitionsSorted(Map partitionsPerTopic,
-                                                    Map subscriptions) {
+    private List allPartitionsSorted(Map partitionsPerTopic,
+                                                     Map subscriptions) {
         SortedSet topics = new TreeSet<>();
         for (Subscription subscription : subscriptions.values())
             topics.addAll(subscription.topics());
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java
index ee537eba78802..77a61dfa2cc28 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java
@@ -16,7 +16,14 @@
  */
 package org.apache.kafka.clients.consumer;
 
-import org.apache.kafka.clients.consumer.internals.AbstractPartitionAssignor;
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import org.apache.kafka.clients.consumer.internals.AbstractStickyAssignor;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.protocol.types.ArrayOf;
 import org.apache.kafka.common.protocol.types.Field;
@@ -24,23 +31,6 @@
 import org.apache.kafka.common.protocol.types.Struct;
 import org.apache.kafka.common.protocol.types.Type;
 import org.apache.kafka.common.utils.CollectionUtils;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.io.Serializable;
-import java.nio.ByteBuffer;
-import java.util.ArrayList;
-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.Map.Entry;
-import java.util.Set;
-import java.util.TreeSet;
 
 /**
  * 
      The sticky assignor serves two purposes. First, it guarantees an assignment that is as balanced as possible, meaning either:
@@ -179,514 +169,93 @@
  * Any consumer that uses sticky assignment can leverage this listener like this:
  * consumer.subscribe(topics, new TheNewRebalanceListener());
  *
+ * Note that you can leverage the {@link CooperativeStickyAssignor} so that only partitions which are being
+ * reassigned to another consumer will be revoked. That is the preferred assignor for newer cluster. See
+ * {@link ConsumerPartitionAssignor.RebalanceProtocol} for a detailed explanation of cooperative rebalancing.
  */
-public class StickyAssignor extends AbstractPartitionAssignor {
-    private static final Logger log = LoggerFactory.getLogger(StickyAssignor.class);
+public class StickyAssignor extends AbstractStickyAssignor {
 
     // these schemas are used for preserving consumer's previously assigned partitions
     // list and sending it as user data to the leader during a rebalance
-    private static final String TOPIC_PARTITIONS_KEY_NAME = "previous_assignment";
-    private static final String TOPIC_KEY_NAME = "topic";
-    private static final String PARTITIONS_KEY_NAME = "partitions";
-    private static final Schema TOPIC_ASSIGNMENT = new Schema(
-            new Field(TOPIC_KEY_NAME, Type.STRING),
-            new Field(PARTITIONS_KEY_NAME, new ArrayOf(Type.INT32)));
-    private static final Schema STICKY_ASSIGNOR_USER_DATA = new Schema(
-            new Field(TOPIC_PARTITIONS_KEY_NAME, new ArrayOf(TOPIC_ASSIGNMENT)));
+    static final String TOPIC_PARTITIONS_KEY_NAME = "previous_assignment";
+    static final String TOPIC_KEY_NAME = "topic";
+    static final String PARTITIONS_KEY_NAME = "partitions";
+    private static final String GENERATION_KEY_NAME = "generation";
+
+    static final Schema TOPIC_ASSIGNMENT = new Schema(
+        new Field(TOPIC_KEY_NAME, Type.STRING),
+        new Field(PARTITIONS_KEY_NAME, new ArrayOf(Type.INT32)));
+    static final Schema STICKY_ASSIGNOR_USER_DATA_V0 = new Schema(
+        new Field(TOPIC_PARTITIONS_KEY_NAME, new ArrayOf(TOPIC_ASSIGNMENT)));
+    private static final Schema STICKY_ASSIGNOR_USER_DATA_V1 = new Schema(
+        new Field(TOPIC_PARTITIONS_KEY_NAME, new ArrayOf(TOPIC_ASSIGNMENT)),
+        new Field(GENERATION_KEY_NAME, Type.INT32));
 
     private List memberAssignment = null;
-    private PartitionMovements partitionMovements;
-
-    public Map> assign(Map partitionsPerTopic,
-                                                    Map subscriptions) {
-        Map> currentAssignment = new HashMap<>();
-        partitionMovements = new PartitionMovements();
-
-        prepopulateCurrentAssignments(subscriptions, currentAssignment);
-        boolean isFreshAssignment = currentAssignment.isEmpty();
-
-        // a mapping of all topic partitions to all consumers that can be assigned to them
-        final Map> partition2AllPotentialConsumers = new HashMap<>();
-        // a mapping of all consumers to all potential topic partitions that can be assigned to them
-        final Map> consumer2AllPotentialPartitions = new HashMap<>();
-
-        // initialize partition2AllPotentialConsumers and consumer2AllPotentialPartitions in the following two for loops
-        for (Entry entry: partitionsPerTopic.entrySet()) {
-            for (int i = 0; i < entry.getValue(); ++i)
-                partition2AllPotentialConsumers.put(new TopicPartition(entry.getKey(), i), new ArrayList());
-        }
-
-        for (Entry entry: subscriptions.entrySet()) {
-            String consumer = entry.getKey();
-            consumer2AllPotentialPartitions.put(consumer, new ArrayList());
-            entry.getValue().topics().stream().filter(topic -> partitionsPerTopic.get(topic) != null).forEach(topic -> {
-                for (int i = 0; i < partitionsPerTopic.get(topic); ++i) {
-                    TopicPartition topicPartition = new TopicPartition(topic, i);
-                    consumer2AllPotentialPartitions.get(consumer).add(topicPartition);
-                    partition2AllPotentialConsumers.get(topicPartition).add(consumer);
-                }
-            });
-
-            // add this consumer to currentAssignment (with an empty topic partition assignment) if it does not already exist
-            if (!currentAssignment.containsKey(consumer))
-                currentAssignment.put(consumer, new ArrayList());
-        }
-
-        // a mapping of partition to current consumer
-        Map currentPartitionConsumer = new HashMap<>();
-        for (Map.Entry> entry: currentAssignment.entrySet())
-            for (TopicPartition topicPartition: entry.getValue())
-                currentPartitionConsumer.put(topicPartition, entry.getKey());
-
-        List sortedPartitions = sortPartitions(
-                currentAssignment, isFreshAssignment, partition2AllPotentialConsumers, consumer2AllPotentialPartitions);
-
-        // all partitions that need to be assigned (initially set to all partitions but adjusted in the following loop)
-        List unassignedPartitions = new ArrayList<>(sortedPartitions);
-        for (Iterator>> it = currentAssignment.entrySet().iterator(); it.hasNext();) {
-            Map.Entry> entry = it.next();
-            if (!subscriptions.containsKey(entry.getKey())) {
-                // if a consumer that existed before (and had some partition assignments) is now removed, remove it from currentAssignment
-                for (TopicPartition topicPartition: entry.getValue())
-                    currentPartitionConsumer.remove(topicPartition);
-                it.remove();
-            } else {
-                // otherwise (the consumer still exists)
-                for (Iterator partitionIter = entry.getValue().iterator(); partitionIter.hasNext();) {
-                    TopicPartition partition = partitionIter.next();
-                    if (!partition2AllPotentialConsumers.containsKey(partition)) {
-                        // if this topic partition of this consumer no longer exists remove it from currentAssignment of the consumer
-                        partitionIter.remove();
-                        currentPartitionConsumer.remove(partition);
-                    } else if (!subscriptions.get(entry.getKey()).topics().contains(partition.topic())) {
-                        // if this partition cannot remain assigned to its current consumer because the consumer
-                        // is no longer subscribed to its topic remove it from currentAssignment of the consumer
-                        partitionIter.remove();
-                    } else
-                        // otherwise, remove the topic partition from those that need to be assigned only if
-                        // its current consumer is still subscribed to its topic (because it is already assigned
-                        // and we would want to preserve that assignment as much as possible)
-                        unassignedPartitions.remove(partition);
-                }
-            }
-        }
-        // at this point we have preserved all valid topic partition to consumer assignments and removed
-        // all invalid topic partitions and invalid consumers. Now we need to assign unassignedPartitions
-        // to consumers so that the topic partition assignments are as balanced as possible.
-
-        // an ascending sorted set of consumers based on how many topic partitions are already assigned to them
-        TreeSet sortedCurrentSubscriptions = new TreeSet<>(new SubscriptionComparator(currentAssignment));
-        sortedCurrentSubscriptions.addAll(currentAssignment.keySet());
+    private int generation = DEFAULT_GENERATION; // consumer group generation
 
-        balance(currentAssignment, sortedPartitions, unassignedPartitions, sortedCurrentSubscriptions,
-                consumer2AllPotentialPartitions, partition2AllPotentialConsumers, currentPartitionConsumer);
-        return currentAssignment;
-    }
-
-    private void prepopulateCurrentAssignments(Map subscriptions,
-                                               Map> currentAssignment) {
-        for (Map.Entry subscriptionEntry : subscriptions.entrySet()) {
-            ByteBuffer userData = subscriptionEntry.getValue().userData();
-            if (userData != null && userData.hasRemaining())
-                currentAssignment.put(subscriptionEntry.getKey(), deserializeTopicPartitionAssignment(userData));
-        }
+    @Override
+    public String name() {
+        return "sticky";
     }
 
     @Override
-    public void onAssignment(Assignment assignment) {
+    public void onAssignment(Assignment assignment, ConsumerGroupMetadata metadata) {
         memberAssignment = assignment.partitions();
+        this.generation = metadata.generationId();
     }
 
     @Override
-    public Subscription subscription(Set topics) {
+    public ByteBuffer subscriptionUserData(Set topics) {
         if (memberAssignment == null)
-            return new Subscription(new ArrayList<>(topics));
+            return null;
 
-        return new Subscription(new ArrayList<>(topics), serializeTopicPartitionAssignment(memberAssignment));
+        return serializeTopicPartitionAssignment(new MemberData(memberAssignment, Optional.of(generation)));
     }
 
     @Override
-    public String name() {
-        return "sticky";
-    }
-
-    /**
-     * determine if the current assignment is a balanced one
-     *
-     * @param currentAssignment: the assignment whose balance needs to be checked
-     * @param sortedCurrentSubscriptions: an ascending sorted set of consumers based on how many topic partitions are already assigned to them
-     * @param allSubscriptions: a mapping of all consumers to all potential topic partitions that can be assigned to them
-     * @return true if the given assignment is balanced; false otherwise
-     */
-    private boolean isBalanced(Map> currentAssignment,
-                               TreeSet sortedCurrentSubscriptions,
-                               Map> allSubscriptions) {
-        int min = currentAssignment.get(sortedCurrentSubscriptions.first()).size();
-        int max = currentAssignment.get(sortedCurrentSubscriptions.last()).size();
-        if (min >= max - 1)
-            // if minimum and maximum numbers of partitions assigned to consumers differ by at most one return true
-            return true;
-
-        // create a mapping from partitions to the consumer assigned to them
-        final Map allPartitions = new HashMap<>();
-        Set>> assignments = currentAssignment.entrySet();
-        for (Map.Entry> entry: assignments) {
-            List topicPartitions = entry.getValue();
-            for (TopicPartition topicPartition: topicPartitions) {
-                if (allPartitions.containsKey(topicPartition))
-                    log.error("{} is assigned to more than one consumer.", topicPartition);
-                allPartitions.put(topicPartition, entry.getKey());
-            }
-        }
-
-        // for each consumer that does not have all the topic partitions it can get make sure none of the topic partitions it
-        // could but did not get cannot be moved to it (because that would break the balance)
-        for (String consumer: sortedCurrentSubscriptions) {
-            List consumerPartitions = currentAssignment.get(consumer);
-            int consumerPartitionCount = consumerPartitions.size();
-
-            // skip if this consumer already has all the topic partitions it can get
-            if (consumerPartitionCount == allSubscriptions.get(consumer).size())
-                continue;
-
-            // otherwise make sure it cannot get any more
-            List potentialTopicPartitions = allSubscriptions.get(consumer);
-            for (TopicPartition topicPartition: potentialTopicPartitions) {
-                if (!currentAssignment.get(consumer).contains(topicPartition)) {
-                    String otherConsumer = allPartitions.get(topicPartition);
-                    int otherConsumerPartitionCount = currentAssignment.get(otherConsumer).size();
-                    if (consumerPartitionCount < otherConsumerPartitionCount) {
-                        log.debug("{} can be moved from consumer {} to consumer {} for a more balanced assignment.",
-                                topicPartition, otherConsumer, consumer);
-                        return false;
-                    }
-                }
-            }
-        }
-        return true;
-    }
-
-    /**
-     * @return the balance score of the given assignment, as the sum of assigned partitions size difference of all consumer pairs.
-     * A perfectly balanced assignment (with all consumers getting the same number of partitions) has a balance score of 0.
-     * Lower balance score indicates a more balanced assignment.
-     */
-    private int getBalanceScore(Map> assignment) {
-        int score = 0;
-
-        Map consumer2AssignmentSize = new HashMap<>();
-        for (Entry> entry: assignment.entrySet())
-            consumer2AssignmentSize.put(entry.getKey(), entry.getValue().size());
-
-        Iterator> it = consumer2AssignmentSize.entrySet().iterator();
-        while (it.hasNext()) {
-            Entry entry = it.next();
-            int consumerAssignmentSize = entry.getValue();
-            it.remove();
-            for (Entry otherEntry: consumer2AssignmentSize.entrySet())
-                score += Math.abs(consumerAssignmentSize - otherEntry.getValue());
-        }
-
-        return score;
-    }
-
-    /**
-     * Sort valid partitions so they are processed in the potential reassignment phase in the proper order
-     * that causes minimal partition movement among consumers (hence honoring maximal stickiness)
-     *
-     * @param currentAssignment the calculated assignment so far
-     * @param isFreshAssignment whether this is a new assignment, or a reassignment of an existing one
-     * @param partition2AllPotentialConsumers a mapping of partitions to their potential consumers
-     * @param consumer2AllPotentialPartitions a mapping of consumers to potential partitions they can consumer from
-     * @return sorted list of valid partitions
-     */
-    private List sortPartitions(Map> currentAssignment,
-                                                boolean isFreshAssignment,
-                                                Map> partition2AllPotentialConsumers,
-                                                Map> consumer2AllPotentialPartitions) {
-        List sortedPartitions = new ArrayList<>();
-
-        if (!isFreshAssignment && areSubscriptionsIdentical(partition2AllPotentialConsumers, consumer2AllPotentialPartitions)) {
-            // if this is a reassignment and the subscriptions are identical (all consumers can consumer from all topics)
-            // then we just need to simply list partitions in a round robin fashion (from consumers with
-            // most assigned partitions to those with least)
-            Map> assignments = deepCopy(currentAssignment);
-            for (Entry> entry: assignments.entrySet()) {
-                List toRemove = new ArrayList<>();
-                for (TopicPartition partition: entry.getValue())
-                    if (!partition2AllPotentialConsumers.keySet().contains(partition))
-                        toRemove.add(partition);
-                for (TopicPartition partition: toRemove)
-                    entry.getValue().remove(partition);
-            }
-            TreeSet sortedConsumers = new TreeSet<>(new SubscriptionComparator(assignments));
-            sortedConsumers.addAll(assignments.keySet());
-
-            while (!sortedConsumers.isEmpty()) {
-                String consumer = sortedConsumers.pollLast();
-                List remainingPartitions = assignments.get(consumer);
-                if (!remainingPartitions.isEmpty()) {
-                    sortedPartitions.add(remainingPartitions.remove(0));
-                    sortedConsumers.add(consumer);
-                }
-            }
-
-            for (TopicPartition partition: partition2AllPotentialConsumers.keySet()) {
-                if (!sortedPartitions.contains(partition))
-                    sortedPartitions.add(partition);
-            }
-
-        } else {
-            // an ascending sorted set of topic partitions based on how many consumers can potentially use them
-            TreeSet sortedAllPartitions = new TreeSet<>(new PartitionComparator(partition2AllPotentialConsumers));
-            sortedAllPartitions.addAll(partition2AllPotentialConsumers.keySet());
-
-            while (!sortedAllPartitions.isEmpty())
-                sortedPartitions.add(sortedAllPartitions.pollFirst());
-        }
-
-        return sortedPartitions;
-    }
-
-    /**
-     * @param partition2AllPotentialConsumers a mapping of partitions to their potential consumers
-     * @param consumer2AllPotentialPartitions a mapping of consumers to potential partitions they can consumer from
-     * @return true if potential consumers of partitions are the same, and potential partitions consumers can
-     * consumer from are the same too
-     */
-    private boolean areSubscriptionsIdentical(Map> partition2AllPotentialConsumers,
-                                              Map> consumer2AllPotentialPartitions) {
-        if (!hasIdenticalListElements(partition2AllPotentialConsumers.values()))
-            return false;
-
-        if (!hasIdenticalListElements(consumer2AllPotentialPartitions.values()))
-            return false;
-
-        return true;
-    }
-
-    /**
-     * @return the consumer to which the given partition is assigned. The assignment should improve the overall balance
-     * of the partition assignments to consumers.
-     */
-    private String assignPartition(TopicPartition partition,
-                                   TreeSet sortedCurrentSubscriptions,
-                                   Map> currentAssignment,
-                                   Map> consumer2AllPotentialPartitions,
-                                   Map currentPartitionConsumer) {
-        for (String consumer: sortedCurrentSubscriptions) {
-            if (consumer2AllPotentialPartitions.get(consumer).contains(partition)) {
-                sortedCurrentSubscriptions.remove(consumer);
-                currentAssignment.get(consumer).add(partition);
-                currentPartitionConsumer.put(partition, consumer);
-                sortedCurrentSubscriptions.add(consumer);
-                return consumer;
-            }
-        }
-        return null;
-    }
-
-    private boolean canParticipateInReassignment(TopicPartition partition,
-                                                 Map> partition2AllPotentialConsumers) {
-        // if a partition has two or more potential consumers it is subject to reassignment.
-        return partition2AllPotentialConsumers.get(partition).size() >= 2;
-    }
-
-    private boolean canParticipateInReassignment(String consumer,
-                                                 Map> currentAssignment,
-                                                 Map> consumer2AllPotentialPartitions,
-                                                 Map> partition2AllPotentialConsumers) {
-        List currentPartitions = currentAssignment.get(consumer);
-        int currentAssignmentSize = currentPartitions.size();
-        int maxAssignmentSize = consumer2AllPotentialPartitions.get(consumer).size();
-        if (currentAssignmentSize > maxAssignmentSize)
-            log.error("The consumer {} is assigned more partitions than the maximum possible.", consumer);
-
-        if (currentAssignmentSize < maxAssignmentSize)
-            // if a consumer is not assigned all its potential partitions it is subject to reassignment
-            return true;
-
-        for (TopicPartition partition: currentPartitions)
-            // if any of the partitions assigned to a consumer is subject to reassignment the consumer itself
-            // is subject to reassignment
-            if (canParticipateInReassignment(partition, partition2AllPotentialConsumers))
-                return true;
-
-        return false;
-    }
-
-    /**
-     * Balance the current assignment using the data structures created in the assign(...) method above.
-     */
-    private void balance(Map> currentAssignment,
-                         List sortedPartitions,
-                         List unassignedPartitions,
-                         TreeSet sortedCurrentSubscriptions,
-                         Map> consumer2AllPotentialPartitions,
-                         Map> partition2AllPotentialConsumers,
-                         Map currentPartitionConsumer) {
-        boolean initializing = currentAssignment.get(sortedCurrentSubscriptions.last()).isEmpty();
-        boolean reassignmentPerformed = false;
-
-        // assign all unassigned partitions
-        for (TopicPartition partition: unassignedPartitions) {
-            // skip if there is no potential consumer for the partition
-            if (partition2AllPotentialConsumers.get(partition).isEmpty())
-                continue;
-
-            assignPartition(partition, sortedCurrentSubscriptions, currentAssignment,
-                            consumer2AllPotentialPartitions, currentPartitionConsumer);
-        }
-
-        // narrow down the reassignment scope to only those partitions that can actually be reassigned
-        Set fixedPartitions = new HashSet<>();
-        for (TopicPartition partition: partition2AllPotentialConsumers.keySet())
-            if (!canParticipateInReassignment(partition, partition2AllPotentialConsumers))
-                fixedPartitions.add(partition);
-        sortedPartitions.removeAll(fixedPartitions);
-
-        // narrow down the reassignment scope to only those consumers that are subject to reassignment
-        Map> fixedAssignments = new HashMap<>();
-        for (String consumer: consumer2AllPotentialPartitions.keySet())
-            if (!canParticipateInReassignment(consumer, currentAssignment,
-                                              consumer2AllPotentialPartitions, partition2AllPotentialConsumers)) {
-                sortedCurrentSubscriptions.remove(consumer);
-                fixedAssignments.put(consumer, currentAssignment.remove(consumer));
-            }
-
-        // create a deep copy of the current assignment so we can revert to it if we do not get a more balanced assignment later
-        Map> preBalanceAssignment = deepCopy(currentAssignment);
-        Map preBalancePartitionConsumers = new HashMap<>(currentPartitionConsumer);
-
-        reassignmentPerformed = performReassignments(sortedPartitions, currentAssignment, sortedCurrentSubscriptions,
-                consumer2AllPotentialPartitions, partition2AllPotentialConsumers, currentPartitionConsumer);
-
-        // if we are not preserving existing assignments and we have made changes to the current assignment
-        // make sure we are getting a more balanced assignment; otherwise, revert to previous assignment
-        if (!initializing && reassignmentPerformed && getBalanceScore(currentAssignment) >= getBalanceScore(preBalanceAssignment)) {
-            deepCopy(preBalanceAssignment, currentAssignment);
-            currentPartitionConsumer.clear();
-            currentPartitionConsumer.putAll(preBalancePartitionConsumers);
-        }
-
-        // add the fixed assignments (those that could not change) back
-        for (Entry> entry: fixedAssignments.entrySet()) {
-            String consumer = entry.getKey();
-            currentAssignment.put(consumer, entry.getValue());
-            sortedCurrentSubscriptions.add(consumer);
+    protected MemberData memberData(Subscription subscription) {
+        ByteBuffer userData = subscription.userData();
+        if (userData == null || !userData.hasRemaining()) {
+            return new MemberData(Collections.emptyList(), Optional.empty());
         }
-
-        fixedAssignments.clear();
+        return deserializeTopicPartitionAssignment(userData);
     }
 
-    private boolean performReassignments(List reassignablePartitions,
-                                         Map> currentAssignment,
-                                         TreeSet sortedCurrentSubscriptions,
-                                         Map> consumer2AllPotentialPartitions,
-                                         Map> partition2AllPotentialConsumers,
-                                         Map currentPartitionConsumer) {
-        boolean reassignmentPerformed = false;
-        boolean modified;
-
-        // repeat reassignment until no partition can be moved to improve the balance
-        do {
-            modified = false;
-            // reassign all reassignable partitions (starting from the partition with least potential consumers and if needed)
-            // until the full list is processed or a balance is achieved
-            Iterator partitionIterator = reassignablePartitions.iterator();
-            while (partitionIterator.hasNext() && !isBalanced(currentAssignment, sortedCurrentSubscriptions, consumer2AllPotentialPartitions)) {
-                TopicPartition partition = partitionIterator.next();
-
-                // the partition must have at least two consumers
-                if (partition2AllPotentialConsumers.get(partition).size() <= 1)
-                    log.error("Expected more than one potential consumer for partition '{}'", partition);
-
-                // the partition must have a current consumer
-                String consumer = currentPartitionConsumer.get(partition);
-                if (consumer == null)
-                    log.error("Expected partition '{}' to be assigned to a consumer", partition);
-
-                // check if a better-suited consumer exist for the partition; if so, reassign it
-                for (String otherConsumer: partition2AllPotentialConsumers.get(partition)) {
-                    if (currentAssignment.get(consumer).size() > currentAssignment.get(otherConsumer).size() + 1) {
-                        reassignPartition(partition, currentAssignment, sortedCurrentSubscriptions, currentPartitionConsumer, consumer2AllPotentialPartitions);
-                        reassignmentPerformed = true;
-                        modified = true;
-                        break;
-                    }
-                }
-            }
-        } while (modified);
-
-        return reassignmentPerformed;
-    }
-
-    private void reassignPartition(TopicPartition partition,
-                                   Map> currentAssignment,
-                                   TreeSet sortedCurrentSubscriptions,
-                                   Map currentPartitionConsumer,
-                                   Map> consumer2AllPotentialPartitions) {
-        String consumer = currentPartitionConsumer.get(partition);
-
-        // find the new consumer
-        String newConsumer = null;
-        for (String anotherConsumer: sortedCurrentSubscriptions) {
-            if (consumer2AllPotentialPartitions.get(anotherConsumer).contains(partition)) {
-                newConsumer = anotherConsumer;
-                break;
-            }
-        }
-
-        assert newConsumer != null;
-
-        // find the correct partition movement considering the stickiness requirement
-        TopicPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(partition, consumer, newConsumer);
-        processPartitionMovement(partitionToBeMoved, newConsumer, currentAssignment, sortedCurrentSubscriptions, currentPartitionConsumer);
-
-        return;
-    }
-
-    private void processPartitionMovement(TopicPartition partition,
-                                          String newConsumer,
-                                          Map> currentAssignment,
-                                          TreeSet sortedCurrentSubscriptions,
-                                          Map currentPartitionConsumer) {
-        String oldConsumer = currentPartitionConsumer.get(partition);
-
-        sortedCurrentSubscriptions.remove(oldConsumer);
-        sortedCurrentSubscriptions.remove(newConsumer);
-
-        partitionMovements.movePartition(partition, oldConsumer, newConsumer);
-
-        currentAssignment.get(oldConsumer).remove(partition);
-        currentAssignment.get(newConsumer).add(partition);
-        currentPartitionConsumer.put(partition, newConsumer);
-        sortedCurrentSubscriptions.add(newConsumer);
-        sortedCurrentSubscriptions.add(oldConsumer);
-    }
-
-    boolean isSticky() {
-        return partitionMovements.isSticky();
-    }
-
-    static ByteBuffer serializeTopicPartitionAssignment(List partitions) {
-        Struct struct = new Struct(STICKY_ASSIGNOR_USER_DATA);
+    // visible for testing
+    static ByteBuffer serializeTopicPartitionAssignment(MemberData memberData) {
+        Struct struct = new Struct(STICKY_ASSIGNOR_USER_DATA_V1);
         List topicAssignments = new ArrayList<>();
-        for (Map.Entry> topicEntry : CollectionUtils.groupPartitionsByTopic(partitions).entrySet()) {
+        for (Map.Entry> topicEntry : CollectionUtils.groupPartitionsByTopic(memberData.partitions).entrySet()) {
             Struct topicAssignment = new Struct(TOPIC_ASSIGNMENT);
             topicAssignment.set(TOPIC_KEY_NAME, topicEntry.getKey());
             topicAssignment.set(PARTITIONS_KEY_NAME, topicEntry.getValue().toArray());
             topicAssignments.add(topicAssignment);
         }
         struct.set(TOPIC_PARTITIONS_KEY_NAME, topicAssignments.toArray());
-        ByteBuffer buffer = ByteBuffer.allocate(STICKY_ASSIGNOR_USER_DATA.sizeOf(struct));
-        STICKY_ASSIGNOR_USER_DATA.write(buffer, struct);
+        if (memberData.generation.isPresent())
+            struct.set(GENERATION_KEY_NAME, memberData.generation.get());
+        ByteBuffer buffer = ByteBuffer.allocate(STICKY_ASSIGNOR_USER_DATA_V1.sizeOf(struct));
+        STICKY_ASSIGNOR_USER_DATA_V1.write(buffer, struct);
         buffer.flip();
         return buffer;
     }
 
-    private static List deserializeTopicPartitionAssignment(ByteBuffer buffer) {
-        Struct struct = STICKY_ASSIGNOR_USER_DATA.read(buffer);
+    private static MemberData deserializeTopicPartitionAssignment(ByteBuffer buffer) {
+        Struct struct;
+        ByteBuffer copy = buffer.duplicate();
+        try {
+            struct = STICKY_ASSIGNOR_USER_DATA_V1.read(buffer);
+        } catch (Exception e1) {
+            try {
+                // fall back to older schema
+                struct = STICKY_ASSIGNOR_USER_DATA_V0.read(copy);
+            } catch (Exception e2) {
+                // ignore the consumer's previous assignment if it cannot be parsed
+                return new MemberData(Collections.emptyList(), Optional.of(DEFAULT_GENERATION));
+            }
+        }
+
         List partitions = new ArrayList<>();
         for (Object structObj : struct.getArray(TOPIC_PARTITIONS_KEY_NAME)) {
             Struct assignment = (Struct) structObj;
@@ -696,267 +265,8 @@ private static List deserializeTopicPartitionAssignment(ByteBuff
                 partitions.add(new TopicPartition(topic, partition));
             }
         }
-        return partitions;
-    }
-
-    /**
-     * @param col a collection of elements of type list
-     * @return true if all lists in the collection have the same members; false otherwise
-     */
-    private  boolean hasIdenticalListElements(Collection> col) {
-        Iterator> it = col.iterator();
-        if (!it.hasNext())
-            return true;
-        List cur = it.next();
-        while (it.hasNext()) {
-            List next = it.next();
-            if (!(cur.containsAll(next) && next.containsAll(cur)))
-                return false;
-            cur = next;
-        }
-        return true;
-    }
-
-    private void deepCopy(Map> source, Map> dest) {
-        dest.clear();
-        for (Entry> entry: source.entrySet())
-            dest.put(entry.getKey(), new ArrayList<>(entry.getValue()));
-    }
-
-    private Map> deepCopy(Map> assignment) {
-        Map> copy = new HashMap<>();
-        deepCopy(assignment, copy);
-        return copy;
-    }
-
-    private static class PartitionComparator implements Comparator, Serializable {
-        private static final long serialVersionUID = 1L;
-        private Map> map;
-
-        PartitionComparator(Map> map) {
-            this.map = map;
-        }
-
-        @Override
-        public int compare(TopicPartition o1, TopicPartition o2) {
-            int ret = map.get(o1).size() - map.get(o2).size();
-            if (ret == 0) {
-                ret = o1.topic().compareTo(o2.topic());
-                if (ret == 0)
-                    ret = o1.partition() - o2.partition();
-            }
-            return ret;
-        }
-    }
-
-    private static class SubscriptionComparator implements Comparator, Serializable {
-        private static final long serialVersionUID = 1L;
-        private Map> map;
-
-        SubscriptionComparator(Map> map) {
-            this.map = map;
-        }
-
-        @Override
-        public int compare(String o1, String o2) {
-            int ret = map.get(o1).size() - map.get(o2).size();
-            if (ret == 0)
-                ret = o1.compareTo(o2);
-            return ret;
-        }
-    }
-
-    /**
-     * This class maintains some data structures to simplify lookup of partition movements among consumers. At each point of
-     * time during a partition rebalance it keeps track of partition movements corresponding to each topic, and also possible
-     * movement (in form a ConsumerPair object) for each partition.
-     */
-    private static class PartitionMovements {
-        private Map>> partitionMovementsByTopic = new HashMap<>();
-        private Map partitionMovements = new HashMap<>();
-
-        private ConsumerPair removeMovementRecordOfPartition(TopicPartition partition) {
-            ConsumerPair pair = partitionMovements.remove(partition);
-
-            String topic = partition.topic();
-            Map> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
-            partitionMovementsForThisTopic.get(pair).remove(partition);
-            if (partitionMovementsForThisTopic.get(pair).isEmpty())
-                partitionMovementsForThisTopic.remove(pair);
-            if (partitionMovementsByTopic.get(topic).isEmpty())
-                partitionMovementsByTopic.remove(topic);
-
-            return pair;
-        }
-
-        private void addPartitionMovementRecord(TopicPartition partition, ConsumerPair pair) {
-            partitionMovements.put(partition, pair);
-
-            String topic = partition.topic();
-            if (!partitionMovementsByTopic.containsKey(topic))
-                partitionMovementsByTopic.put(topic, new HashMap>());
-
-            Map> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
-            if (!partitionMovementsForThisTopic.containsKey(pair))
-                partitionMovementsForThisTopic.put(pair, new HashSet());
-
-            partitionMovementsForThisTopic.get(pair).add(partition);
-        }
-
-        private void movePartition(TopicPartition partition, String oldConsumer, String newConsumer) {
-            ConsumerPair pair = new ConsumerPair(oldConsumer, newConsumer);
-
-            if (partitionMovements.containsKey(partition)) {
-                // this partition has previously moved
-                ConsumerPair existingPair = removeMovementRecordOfPartition(partition);
-                assert existingPair.dstMemberId.equals(oldConsumer);
-                if (!existingPair.srcMemberId.equals(newConsumer)) {
-                    // the partition is not moving back to its previous consumer
-                    // return new ConsumerPair2(existingPair.src, newConsumer);
-                    addPartitionMovementRecord(partition, new ConsumerPair(existingPair.srcMemberId, newConsumer));
-                }
-            } else
-                addPartitionMovementRecord(partition, pair);
-        }
-
-        private TopicPartition getTheActualPartitionToBeMoved(TopicPartition partition, String oldConsumer, String newConsumer) {
-            String topic = partition.topic();
-
-            if (!partitionMovementsByTopic.containsKey(topic))
-                return partition;
-
-            if (partitionMovements.containsKey(partition)) {
-                // this partition has previously moved
-                assert oldConsumer.equals(partitionMovements.get(partition).dstMemberId);
-                oldConsumer = partitionMovements.get(partition).srcMemberId;
-            }
-
-            Map> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
-            ConsumerPair reversePair = new ConsumerPair(newConsumer, oldConsumer);
-            if (!partitionMovementsForThisTopic.containsKey(reversePair))
-                return partition;
-
-            return partitionMovementsForThisTopic.get(reversePair).iterator().next();
-        }
-
-        private boolean isLinked(String src, String dst, Set pairs, List currentPath) {
-            if (src.equals(dst))
-                return false;
-
-            if (pairs.isEmpty())
-                return false;
-
-            if (new ConsumerPair(src, dst).in(pairs)) {
-                currentPath.add(src);
-                currentPath.add(dst);
-                return true;
-            }
-
-            for (ConsumerPair pair: pairs)
-                if (pair.srcMemberId.equals(src)) {
-                    Set reducedSet = new HashSet<>(pairs);
-                    reducedSet.remove(pair);
-                    currentPath.add(pair.srcMemberId);
-                    return isLinked(pair.dstMemberId, dst, reducedSet, currentPath);
-                }
-
-            return false;
-        }
-
-        private boolean in(List cycle, Set> cycles) {
-            List superCycle = new ArrayList<>(cycle);
-            superCycle.remove(superCycle.size() - 1);
-            superCycle.addAll(cycle);
-            for (List foundCycle: cycles) {
-                if (foundCycle.size() == cycle.size() && Collections.indexOfSubList(superCycle, foundCycle) != -1)
-                    return true;
-            }
-            return false;
-        }
-
-        private boolean hasCycles(Set pairs) {
-            Set> cycles = new HashSet<>();
-            for (ConsumerPair pair: pairs) {
-                Set reducedPairs = new HashSet<>(pairs);
-                reducedPairs.remove(pair);
-                List path = new ArrayList<>(Collections.singleton(pair.srcMemberId));
-                if (isLinked(pair.dstMemberId, pair.srcMemberId, reducedPairs, path) && !in(path, cycles)) {
-                    cycles.add(new ArrayList<>(path));
-                    log.error("A cycle of length {} was found: {}", path.size() - 1, path.toString());
-                }
-            }
-
-            // for now we want to make sure there is no partition movements of the same topic between a pair of consumers.
-            // the odds of finding a cycle among more than two consumers seem to be very low (according to various randomized
-            // tests with the given sticky algorithm) that it should not worth the added complexity of handling those cases.
-            for (List cycle: cycles)
-                if (cycle.size() == 3) // indicates a cycle of length 2
-                    return true;
-            return false;
-        }
-
-        private boolean isSticky() {
-            for (Map.Entry>> topicMovements: this.partitionMovementsByTopic.entrySet()) {
-                Set topicMovementPairs = topicMovements.getValue().keySet();
-                if (hasCycles(topicMovementPairs)) {
-                    log.error("Stickiness is violated for topic {}"
-                            + "\nPartition movements for this topic occurred among the following consumer pairs:"
-                            + "\n{}", topicMovements.getKey(), topicMovements.getValue().toString());
-                    return false;
-                }
-            }
-
-            return true;
-        }
-    }
-
-    /**
-     * ConsumerPair represents a pair of Kafka consumer ids involved in a partition reassignment. Each
-     * ConsumerPair object, which contains a source (src) and a destination (dst)
-     * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
-     * partition of the particular topic was moved from the source consumer to the destination consumer during the rebalance.
-     * This class is used, through the PartitionMovements class, by the sticky assignor and helps in determining
-     * whether a partition reassignment results in cycles among the generated graph of consumer pairs.
-     */
-    private static class ConsumerPair {
-        private final String srcMemberId;
-        private final String dstMemberId;
-
-        ConsumerPair(String srcMemberId, String dstMemberId) {
-            this.srcMemberId = srcMemberId;
-            this.dstMemberId = dstMemberId;
-        }
-
-        public String toString() {
-            return this.srcMemberId + "->" + this.dstMemberId;
-        }
-
-        @Override
-        public int hashCode() {
-            final int prime = 31;
-            int result = 1;
-            result = prime * result + ((this.srcMemberId == null) ? 0 : this.srcMemberId.hashCode());
-            result = prime * result + ((this.dstMemberId == null) ? 0 : this.dstMemberId.hashCode());
-            return result;
-        }
-
-        @Override
-        public boolean equals(Object obj) {
-            if (obj == null)
-                return false;
-
-            if (!getClass().isInstance(obj))
-                return false;
-
-            ConsumerPair otherPair = (ConsumerPair) obj;
-            return this.srcMemberId.equals(otherPair.srcMemberId) && this.dstMemberId.equals(otherPair.dstMemberId);
-        }
-
-        private boolean in(Set pairs) {
-            for (ConsumerPair pair: pairs)
-                if (this.equals(pair))
-                    return true;
-            return false;
-        }
+        // make sure this is backward compatible
+        Optional generation = struct.hasField(GENERATION_KEY_NAME) ? Optional.of(struct.getInt(GENERATION_KEY_NAME)) : Optional.empty();
+        return new MemberData(partitions, generation);
     }
-}
+}
\ No newline at end of file
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractCoordinator.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractCoordinator.java
index fd7431b01501c..efb2765ec4152 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractCoordinator.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractCoordinator.java
@@ -17,10 +17,12 @@
 package org.apache.kafka.clients.consumer.internals;
 
 import org.apache.kafka.clients.ClientResponse;
+import org.apache.kafka.clients.GroupRebalanceConfig;
 import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.Node;
 import org.apache.kafka.common.errors.AuthenticationException;
 import org.apache.kafka.common.errors.DisconnectException;
+import org.apache.kafka.common.errors.FencedInstanceIdException;
 import org.apache.kafka.common.errors.GroupAuthorizationException;
 import org.apache.kafka.common.errors.GroupMaxSizeReachedException;
 import org.apache.kafka.common.errors.IllegalGenerationException;
@@ -29,22 +31,32 @@
 import org.apache.kafka.common.errors.RebalanceInProgressException;
 import org.apache.kafka.common.errors.RetriableException;
 import org.apache.kafka.common.errors.UnknownMemberIdException;
-import org.apache.kafka.common.message.LeaveGroupRequestData;
+import org.apache.kafka.common.message.FindCoordinatorRequestData;
+import org.apache.kafka.common.message.HeartbeatRequestData;
+import org.apache.kafka.common.message.JoinGroupRequestData;
+import org.apache.kafka.common.message.JoinGroupResponseData;
+import org.apache.kafka.common.message.LeaveGroupRequestData.MemberIdentity;
+import org.apache.kafka.common.message.LeaveGroupResponseData.MemberResponse;
+import org.apache.kafka.common.message.SyncGroupRequestData;
 import org.apache.kafka.common.metrics.Measurable;
 import org.apache.kafka.common.metrics.MetricConfig;
 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.Count;
+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.Meter;
+import org.apache.kafka.common.metrics.stats.Rate;
+import org.apache.kafka.common.metrics.stats.WindowedCount;
+import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
 import org.apache.kafka.common.requests.FindCoordinatorRequest;
+import org.apache.kafka.common.requests.FindCoordinatorRequest.CoordinatorType;
 import org.apache.kafka.common.requests.FindCoordinatorResponse;
 import org.apache.kafka.common.requests.HeartbeatRequest;
 import org.apache.kafka.common.requests.HeartbeatResponse;
 import org.apache.kafka.common.requests.JoinGroupRequest;
-import org.apache.kafka.common.requests.JoinGroupRequest.ProtocolMetadata;
 import org.apache.kafka.common.requests.JoinGroupResponse;
 import org.apache.kafka.common.requests.LeaveGroupRequest;
 import org.apache.kafka.common.requests.LeaveGroupResponse;
@@ -55,10 +67,12 @@
 import org.apache.kafka.common.utils.LogContext;
 import org.apache.kafka.common.utils.Time;
 import org.apache.kafka.common.utils.Timer;
+import org.apache.kafka.common.utils.Utils;
 import org.slf4j.Logger;
 
 import java.io.Closeable;
 import java.nio.ByteBuffer;
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
@@ -86,7 +100,7 @@
  *
  * To leverage this protocol, an implementation must define the format of metadata provided by each
  * member for group registration in {@link #metadata()} and the format of the state assignment provided
- * by the leader in {@link #performAssignment(String, String, Map)} and becomes available to members in
+ * by the leader in {@link #performAssignment(String, String, List)} and becomes available to members in
  * {@link #onJoinComplete(int, String, String, ByteBuffer)}.
  *
  * Note on locking: this class shares state between the caller and a background thread which is
@@ -105,67 +119,41 @@ private enum MemberState {
     }
 
     private final Logger log;
-    private final int sessionTimeoutMs;
-    private final boolean leaveGroupOnClose;
     private final GroupCoordinatorMetrics sensors;
     private final Heartbeat heartbeat;
-    protected final int rebalanceTimeoutMs;
-    protected final String groupId;
+    private final GroupRebalanceConfig rebalanceConfig;
     protected final ConsumerNetworkClient client;
     protected final Time time;
-    protected final long retryBackoffMs;
 
-    private HeartbeatThread heartbeatThread = null;
+    private Node coordinator = null;
     private boolean rejoinNeeded = true;
     private boolean needsJoinPrepare = true;
     private MemberState state = MemberState.UNJOINED;
+    private HeartbeatThread heartbeatThread = null;
     private RequestFuture joinFuture = null;
-    private Node coordinator = null;
     private Generation generation = Generation.NO_GENERATION;
+    private long lastRebalanceStartMs = -1L;
+    private long lastRebalanceEndMs = -1L;
 
     private RequestFuture findCoordinatorFuture = null;
 
     /**
      * Initialize the coordination manager.
      */
-    public AbstractCoordinator(LogContext logContext,
+    public AbstractCoordinator(GroupRebalanceConfig rebalanceConfig,
+                               LogContext logContext,
                                ConsumerNetworkClient client,
-                               String groupId,
-                               int rebalanceTimeoutMs,
-                               int sessionTimeoutMs,
-                               Heartbeat heartbeat,
                                Metrics metrics,
                                String metricGrpPrefix,
-                               Time time,
-                               long retryBackoffMs,
-                               boolean leaveGroupOnClose) {
+                               Time time) {
+        Objects.requireNonNull(rebalanceConfig.groupId,
+                               "Expected a non-null group id for coordinator construction");
+        this.rebalanceConfig = rebalanceConfig;
         this.log = logContext.logger(AbstractCoordinator.class);
         this.client = client;
         this.time = time;
-        this.groupId = Objects.requireNonNull(groupId,
-                "Expected a non-null group id for coordinator construction");
-        this.rebalanceTimeoutMs = rebalanceTimeoutMs;
-        this.sessionTimeoutMs = sessionTimeoutMs;
-        this.leaveGroupOnClose = leaveGroupOnClose;
-        this.heartbeat = heartbeat;
+        this.heartbeat = new Heartbeat(rebalanceConfig, time);
         this.sensors = new GroupCoordinatorMetrics(metrics, metricGrpPrefix);
-        this.retryBackoffMs = retryBackoffMs;
-    }
-
-    public AbstractCoordinator(LogContext logContext,
-                               ConsumerNetworkClient client,
-                               String groupId,
-                               int rebalanceTimeoutMs,
-                               int sessionTimeoutMs,
-                               int heartbeatIntervalMs,
-                               Metrics metrics,
-                               String metricGrpPrefix,
-                               Time time,
-                               long retryBackoffMs,
-                               boolean leaveGroupOnClose) {
-        this(logContext, client, groupId, rebalanceTimeoutMs, sessionTimeoutMs,
-                new Heartbeat(time, sessionTimeoutMs, heartbeatIntervalMs, rebalanceTimeoutMs, retryBackoffMs),
-                metrics, metricGrpPrefix, time, retryBackoffMs, leaveGroupOnClose);
     }
 
     /**
@@ -183,7 +171,7 @@ public AbstractCoordinator(LogContext logContext,
      * on the preference).
      * @return Non-empty map of supported protocols and metadata
      */
-    protected abstract List metadata();
+    protected abstract JoinGroupRequestData.JoinGroupRequestProtocolCollection metadata();
 
     /**
      * Invoked prior to each group join or rejoin. This is typically used to perform any
@@ -197,12 +185,13 @@ public AbstractCoordinator(LogContext logContext,
      * Perform assignment for the group. This is used by the leader to push state to all the members
      * of the group (e.g. to push partition assignments in the case of the new consumer)
      * @param leaderId The id of the leader (which is this member)
+     * @param protocol The protocol selected by the coordinator
      * @param allMemberMetadata Metadata from all members of the group
      * @return A map from each member to their state assignment
      */
     protected abstract Map performAssignment(String leaderId,
                                                                  String protocol,
-                                                                 Map allMemberMetadata);
+                                                                 List allMemberMetadata);
 
     /**
      * Invoked when a group member has successfully joined a group. If this call fails with an exception,
@@ -218,6 +207,13 @@ protected abstract void onJoinComplete(int generation,
                                            String protocol,
                                            ByteBuffer memberAssignment);
 
+    /**
+     * Invoked prior to each leave group event. This is typically used to cleanup assigned partitions;
+     * note it is triggered by the consumer's API caller thread (i.e. background heartbeat thread would
+     * not trigger it even if it tries to force leaving group upon heartbeat session expiration)
+     */
+    protected void onLeavePrepare() {}
+
     /**
      * Visible for testing.
      *
@@ -249,7 +245,7 @@ protected synchronized boolean ensureCoordinatorReady(final Timer timer) {
                 // we found the coordinator, but the connection has failed, so mark
                 // it dead and backoff before retrying discovery
                 markCoordinatorUnknown();
-                timer.sleep(retryBackoffMs);
+                timer.sleep(rebalanceConfig.retryBackoffMs);
             }
         } while (coordinatorUnknown() && timer.notExpired());
 
@@ -331,6 +327,7 @@ public void ensureActiveGroup() {
      * Ensure the group is active (i.e., joined and synced)
      *
      * @param timer Timer bounding how long this method can block
+     * @throws KafkaException if the callback throws exception
      * @return true iff the group is active
      */
     boolean ensureActiveGroup(final Timer timer) {
@@ -379,6 +376,7 @@ private void closeHeartbeatThread() {
      * Visible for testing.
      *
      * @param timer Timer bounding how long this method can block
+     * @throws KafkaException if the callback throws exception
      * @return true iff the operation succeeded
      */
     boolean joinGroupIfNeeded(final Timer timer) {
@@ -393,8 +391,10 @@ boolean joinGroupIfNeeded(final Timer timer) {
             // refresh which changes the matched subscription set) can occur while another rebalance is
             // still in progress.
             if (needsJoinPrepare) {
-                onJoinPrepare(generation.generationId, generation.memberId);
+                // need to set the flag before calling onJoinPrepare since the user callback may throw
+                // exception, in which case upon retry we should not retry onJoinPrepare either.
                 needsJoinPrepare = false;
+                onJoinPrepare(generation.generationId, generation.memberId);
             }
 
             final RequestFuture future = initiateJoinGroup();
@@ -405,14 +405,34 @@ boolean joinGroupIfNeeded(final Timer timer) {
             }
 
             if (future.succeeded()) {
-                // Duplicate the buffer in case `onJoinComplete` does not complete and needs to be retried.
-                ByteBuffer memberAssignment = future.value().duplicate();
-                onJoinComplete(generation.generationId, generation.memberId, generation.protocol, memberAssignment);
+                Generation generationSnapshot;
 
-                // We reset the join group future only after the completion callback returns. This ensures
-                // that if the callback is woken up, we will retry it on the next joinGroupIfNeeded.
-                resetJoinGroupFuture();
-                needsJoinPrepare = true;
+                // Generation data maybe concurrently cleared by Heartbeat thread.
+                // Can't use synchronized for {@code onJoinComplete}, because it can be long enough
+                // and shouldn't block heartbeat thread.
+                // See {@link PlaintextConsumerTest#testMaxPollIntervalMsDelayInAssignment}
+                synchronized (AbstractCoordinator.this) {
+                    generationSnapshot = this.generation;
+                }
+
+                if (generationSnapshot != Generation.NO_GENERATION) {
+                    // Duplicate the buffer in case `onJoinComplete` does not complete and needs to be retried.
+                    ByteBuffer memberAssignment = future.value().duplicate();
+
+                    onJoinComplete(generationSnapshot.generationId, generationSnapshot.memberId, generationSnapshot.protocol, memberAssignment);
+
+                    // Generally speaking we should always resetJoinGroupFuture once the future is done, but here
+                    // we can only reset the join group future after the completion callback returns. This ensures
+                    // that if the callback is woken up, we will retry it on the next joinGroupIfNeeded.
+                    // And because of that we should explicitly trigger resetJoinGroupFuture in other conditions below.
+                    resetJoinGroupFuture();
+                    needsJoinPrepare = true;
+                } else {
+                    log.info("Generation data was cleared by heartbeat thread. Initiating rejoin.");
+                    resetStateAndRejoin();
+                    resetJoinGroupFuture();
+                    return false;
+                }
             } else {
                 resetJoinGroupFuture();
                 final RuntimeException exception = future.exception();
@@ -424,7 +444,7 @@ boolean joinGroupIfNeeded(final Timer timer) {
                 else if (!future.isRetriable())
                     throw exception;
 
-                timer.sleep(retryBackoffMs);
+                timer.sleep(rebalanceConfig.retryBackoffMs);
             }
         }
         return true;
@@ -434,6 +454,11 @@ private synchronized void resetJoinGroupFuture() {
         this.joinFuture = null;
     }
 
+    private synchronized void resetStateAndRejoin() {
+        rejoinNeeded = true;
+        state = MemberState.UNJOINED;
+    }
+
     private synchronized RequestFuture initiateJoinGroup() {
         // we store the join future in case we are woken up by the user after beginning the
         // rebalance in the call to poll below. This ensures that we do not mistakenly attempt
@@ -445,6 +470,10 @@ private synchronized RequestFuture initiateJoinGroup() {
             disableHeartbeatThread();
 
             state = MemberState.REBALANCING;
+            // a rebalance can be triggered consecutively if the previous one failed,
+            // in this case we would not update the start time.
+            if (lastRebalanceStartMs == -1L)
+                lastRebalanceStartMs = time.milliseconds();
             joinFuture = sendJoinGroupRequest();
             joinFuture.addListener(new RequestFutureListener() {
                 @Override
@@ -452,12 +481,21 @@ public void onSuccess(ByteBuffer value) {
                     // handle join completion in the callback so that the callback will be invoked
                     // even if the consumer is woken up before finishing the rebalance
                     synchronized (AbstractCoordinator.this) {
-                        log.info("Successfully joined group with generation {}", generation.generationId);
-                        state = MemberState.STABLE;
-                        rejoinNeeded = false;
-
-                        if (heartbeatThread != null)
-                            heartbeatThread.enable();
+                        if (generation != Generation.NO_GENERATION) {
+                            log.info("Successfully joined group with generation {}", generation.generationId);
+                            state = MemberState.STABLE;
+                            rejoinNeeded = false;
+                            // record rebalance latency
+                            lastRebalanceEndMs = time.milliseconds();
+                            sensors.successfulRebalanceSensor.record(lastRebalanceEndMs - lastRebalanceStartMs);
+                            lastRebalanceStartMs = -1L;
+
+                            if (heartbeatThread != null)
+                                heartbeatThread.enable();
+                        } else {
+                            log.info("Generation data was cleared by heartbeat thread. Rejoin failed.");
+                            recordRebalanceFailure();
+                        }
                     }
                 }
 
@@ -466,9 +504,14 @@ public void onFailure(RuntimeException e) {
                     // we handle failures below after the request finishes. if the join completes
                     // after having been woken up, the exception is ignored and we will rejoin
                     synchronized (AbstractCoordinator.this) {
-                        state = MemberState.UNJOINED;
+                        recordRebalanceFailure();
                     }
                 }
+
+                private void recordRebalanceFailure() {
+                    state = MemberState.UNJOINED;
+                    sensors.failedRebalanceSensor.record();
+                }
             });
         }
         return joinFuture;
@@ -476,7 +519,7 @@ public void onFailure(RuntimeException e) {
 
     /**
      * Join the group and return the assignment for the next generation. This function handles both
-     * JoinGroup and SyncGroup, delegating to {@link #performAssignment(String, String, Map)} if
+     * JoinGroup and SyncGroup, delegating to {@link #performAssignment(String, String, List)} if
      * elected leader by the coordinator.
      *
      * NOTE: This is visible only for testing
@@ -490,18 +533,22 @@ RequestFuture sendJoinGroupRequest() {
         // send a join group request to the coordinator
         log.info("(Re-)joining group");
         JoinGroupRequest.Builder requestBuilder = new JoinGroupRequest.Builder(
-                groupId,
-                this.sessionTimeoutMs,
-                this.generation.memberId,
-                protocolType(),
-                metadata()).setRebalanceTimeout(this.rebalanceTimeoutMs);
+                new JoinGroupRequestData()
+                        .setGroupId(rebalanceConfig.groupId)
+                        .setSessionTimeoutMs(this.rebalanceConfig.sessionTimeoutMs)
+                        .setMemberId(this.generation.memberId)
+                        .setGroupInstanceId(this.rebalanceConfig.groupInstanceId.orElse(null))
+                        .setProtocolType(protocolType())
+                        .setProtocols(metadata())
+                        .setRebalanceTimeoutMs(this.rebalanceConfig.rebalanceTimeoutMs)
+        );
 
         log.debug("Sending JoinGroup ({}) to coordinator {}", requestBuilder, this.coordinator);
 
         // Note that we override the request timeout using the rebalance timeout since that is the
         // maximum time that it may block on the coordinator. We add an extra 5 seconds for small delays.
 
-        int joinGroupTimeoutMs = Math.max(rebalanceTimeoutMs, rebalanceTimeoutMs + 5000);
+        int joinGroupTimeoutMs = Math.max(rebalanceConfig.rebalanceTimeoutMs, rebalanceConfig.rebalanceTimeoutMs + 5000);
         return client.send(coordinator, requestBuilder, joinGroupTimeoutMs)
                 .compose(new JoinGroupResponseHandler());
     }
@@ -512,7 +559,7 @@ public void handle(JoinGroupResponse joinResponse, RequestFuture fut
             Errors error = joinResponse.error();
             if (error == Errors.NONE) {
                 log.debug("Received successful JoinGroup response: {}", joinResponse);
-                sensors.joinLatency.record(response.requestLatencyMs());
+                sensors.joinSensor.record(response.requestLatencyMs());
 
                 synchronized (AbstractCoordinator.this) {
                     if (state != MemberState.REBALANCING) {
@@ -520,8 +567,8 @@ public void handle(JoinGroupResponse joinResponse, RequestFuture fut
                         // the group. In this case, we do not want to continue with the sync group.
                         future.raise(new UnjoinedGroupException());
                     } else {
-                        AbstractCoordinator.this.generation = new Generation(joinResponse.generationId(),
-                                joinResponse.memberId(), joinResponse.groupProtocol());
+                        AbstractCoordinator.this.generation = new Generation(joinResponse.data().generationId(),
+                                joinResponse.data().memberId(), joinResponse.data().protocolName());
                         if (joinResponse.isLeader()) {
                             onJoinLeader(joinResponse).chain(future);
                         } else {
@@ -535,38 +582,46 @@ public void handle(JoinGroupResponse joinResponse, RequestFuture fut
                 future.raise(error);
             } else if (error == Errors.UNKNOWN_MEMBER_ID) {
                 // reset the member id and retry immediately
-                resetGeneration();
+                resetGenerationOnResponseError(ApiKeys.JOIN_GROUP, error);
                 log.debug("Attempt to join group failed due to unknown member id.");
-                future.raise(Errors.UNKNOWN_MEMBER_ID);
+                future.raise(error);
             } else if (error == Errors.COORDINATOR_NOT_AVAILABLE
                     || error == Errors.NOT_COORDINATOR) {
                 // re-discover the coordinator and retry with backoff
                 markCoordinatorUnknown();
                 log.debug("Attempt to join group failed due to obsolete coordinator information: {}", error.message());
                 future.raise(error);
+            } else if (error == Errors.FENCED_INSTANCE_ID) {
+                log.error("Received fatal exception: group.instance.id gets fenced");
+                future.raise(error);
             } else if (error == Errors.INCONSISTENT_GROUP_PROTOCOL
                     || error == Errors.INVALID_SESSION_TIMEOUT
                     || error == Errors.INVALID_GROUP_ID
                     || error == Errors.GROUP_AUTHORIZATION_FAILED
                     || error == Errors.GROUP_MAX_SIZE_REACHED) {
+                // log the error and re-throw the exception
                 log.error("Attempt to join group failed due to fatal error: {}", error.message());
                 if (error == Errors.GROUP_MAX_SIZE_REACHED) {
-                    future.raise(new GroupMaxSizeReachedException(groupId));
+                    future.raise(new GroupMaxSizeReachedException("Consumer group " + rebalanceConfig.groupId +
+                            " already has the configured maximum number of members."));
                 } else if (error == Errors.GROUP_AUTHORIZATION_FAILED) {
-                    future.raise(new GroupAuthorizationException(groupId));
+                    future.raise(GroupAuthorizationException.forGroupId(rebalanceConfig.groupId));
                 } else {
                     future.raise(error);
                 }
+            } else if (error == Errors.UNSUPPORTED_VERSION) {
+                log.error("Attempt to join group failed due to unsupported version error. Please unset field group.instance.id and retry" +
+                        "to see if the problem resolves");
+                future.raise(error);
             } else if (error == Errors.MEMBER_ID_REQUIRED) {
                 // Broker requires a concrete member id to be allowed to join the group. Update member id
                 // and send another join group request in next cycle.
                 synchronized (AbstractCoordinator.this) {
                     AbstractCoordinator.this.generation = new Generation(OffsetCommitRequest.DEFAULT_GENERATION_ID,
-                        joinResponse.memberId(), null);
-                    AbstractCoordinator.this.rejoinNeeded = true;
-                    AbstractCoordinator.this.state = MemberState.UNJOINED;
+                            joinResponse.data().memberId(), null);
+                    AbstractCoordinator.this.resetStateAndRejoin();
                 }
-                future.raise(Errors.MEMBER_ID_REQUIRED);
+                future.raise(error);
             } else {
                 // unexpected error, throw the exception
                 log.error("Attempt to join group failed due to unexpected error: {}", error.message());
@@ -578,21 +633,42 @@ public void handle(JoinGroupResponse joinResponse, RequestFuture fut
     private RequestFuture onJoinFollower() {
         // send follower's sync group with an empty assignment
         SyncGroupRequest.Builder requestBuilder =
-                new SyncGroupRequest.Builder(groupId, generation.generationId, generation.memberId,
-                        Collections.emptyMap());
-        log.debug("Sending follower SyncGroup to coordinator {}: {}", this.coordinator, requestBuilder);
+                new SyncGroupRequest.Builder(
+                        new SyncGroupRequestData()
+                                .setGroupId(rebalanceConfig.groupId)
+                                .setMemberId(generation.memberId)
+                                .setGroupInstanceId(this.rebalanceConfig.groupInstanceId.orElse(null))
+                                .setGenerationId(generation.generationId)
+                                .setAssignments(Collections.emptyList())
+                );
+        log.debug("Sending follower SyncGroup to coordinator {} at generation {}: {}", this.coordinator, this.generation, requestBuilder);
         return sendSyncGroupRequest(requestBuilder);
     }
 
     private RequestFuture onJoinLeader(JoinGroupResponse joinResponse) {
         try {
             // perform the leader synchronization and send back the assignment for the group
-            Map groupAssignment = performAssignment(joinResponse.leaderId(), joinResponse.groupProtocol(),
-                    joinResponse.members());
+            Map groupAssignment = performAssignment(joinResponse.data().leader(), joinResponse.data().protocolName(),
+                    joinResponse.data().members());
+
+            List groupAssignmentList = new ArrayList<>();
+            for (Map.Entry assignment : groupAssignment.entrySet()) {
+                groupAssignmentList.add(new SyncGroupRequestData.SyncGroupRequestAssignment()
+                        .setMemberId(assignment.getKey())
+                        .setAssignment(Utils.toArray(assignment.getValue()))
+                );
+            }
 
             SyncGroupRequest.Builder requestBuilder =
-                    new SyncGroupRequest.Builder(groupId, generation.generationId, generation.memberId, groupAssignment);
-            log.debug("Sending leader SyncGroup to coordinator {}: {}", this.coordinator, requestBuilder);
+                    new SyncGroupRequest.Builder(
+                            new SyncGroupRequestData()
+                                    .setGroupId(rebalanceConfig.groupId)
+                                    .setMemberId(generation.memberId)
+                                    .setGroupInstanceId(this.rebalanceConfig.groupInstanceId.orElse(null))
+                                    .setGenerationId(generation.generationId)
+                                    .setAssignments(groupAssignmentList)
+                    );
+            log.debug("Sending leader SyncGroup to coordinator {} at generation {}: {}", this.coordinator, this.generation, requestBuilder);
             return sendSyncGroupRequest(requestBuilder);
         } catch (RuntimeException e) {
             return RequestFuture.failure(e);
@@ -612,20 +688,23 @@ public void handle(SyncGroupResponse syncResponse,
                            RequestFuture future) {
             Errors error = syncResponse.error();
             if (error == Errors.NONE) {
-                sensors.syncLatency.record(response.requestLatencyMs());
-                future.complete(syncResponse.memberAssignment());
+                sensors.syncSensor.record(response.requestLatencyMs());
+                future.complete(ByteBuffer.wrap(syncResponse.data.assignment()));
             } else {
                 requestRejoin();
 
                 if (error == Errors.GROUP_AUTHORIZATION_FAILED) {
-                    future.raise(new GroupAuthorizationException(groupId));
+                    future.raise(GroupAuthorizationException.forGroupId(rebalanceConfig.groupId));
                 } else if (error == Errors.REBALANCE_IN_PROGRESS) {
                     log.debug("SyncGroup failed because the group began another rebalance");
                     future.raise(error);
+                } else if (error == Errors.FENCED_INSTANCE_ID) {
+                    log.error("Received fatal exception: group.instance.id gets fenced");
+                    future.raise(error);
                 } else if (error == Errors.UNKNOWN_MEMBER_ID
                         || error == Errors.ILLEGAL_GENERATION) {
                     log.debug("SyncGroup failed: {}", error.message());
-                    resetGeneration();
+                    resetGenerationOnResponseError(ApiKeys.SYNC_GROUP, error);
                     future.raise(error);
                 } else if (error == Errors.COORDINATOR_NOT_AVAILABLE
                         || error == Errors.NOT_COORDINATOR) {
@@ -646,11 +725,14 @@ public void handle(SyncGroupResponse syncResponse,
      */
     private RequestFuture sendFindCoordinatorRequest(Node node) {
         // initiate the group metadata request
-        log.debug("Sending FindCoordinator request to broker {}", node);
+        log.info("Sending FindCoordinator request to broker {}", node);
         FindCoordinatorRequest.Builder requestBuilder =
-                new FindCoordinatorRequest.Builder(FindCoordinatorRequest.CoordinatorType.GROUP, this.groupId);
+                new FindCoordinatorRequest.Builder(
+                        new FindCoordinatorRequestData()
+                            .setKeyType(CoordinatorType.GROUP.id())
+                            .setKey(this.rebalanceConfig.groupId));
         return client.send(node, requestBuilder)
-                     .compose(new FindCoordinatorResponseHandler());
+                .compose(new FindCoordinatorResponseHandler());
     }
 
     private class FindCoordinatorResponseHandler extends RequestFutureAdapter {
@@ -666,21 +748,21 @@ public void onSuccess(ClientResponse resp, RequestFuture future) {
                 synchronized (AbstractCoordinator.this) {
                     // use MAX_VALUE - node.id as the coordinator id to allow separate connections
                     // for the coordinator in the underlying network client layer
-                    int coordinatorConnectionId = Integer.MAX_VALUE - findCoordinatorResponse.node().id();
+                    int coordinatorConnectionId = Integer.MAX_VALUE - findCoordinatorResponse.data().nodeId();
 
                     AbstractCoordinator.this.coordinator = new Node(
                             coordinatorConnectionId,
-                            findCoordinatorResponse.node().host(),
-                            findCoordinatorResponse.node().port());
+                            findCoordinatorResponse.data().host(),
+                            findCoordinatorResponse.data().port());
                     log.info("Discovered group coordinator {}", coordinator);
                     client.tryConnect(coordinator);
                     heartbeat.resetSessionTimeout();
                 }
                 future.complete(null);
             } else if (error == Errors.GROUP_AUTHORIZATION_FAILED) {
-                future.raise(new GroupAuthorizationException(groupId));
+                future.raise(GroupAuthorizationException.forGroupId(rebalanceConfig.groupId));
             } else {
-                log.debug("Group coordinator lookup failed: {}", error.message());
+                log.info("Group coordinator lookup failed: {}", findCoordinatorResponse.data().errorMessage());
                 future.raise(error);
             }
         }
@@ -740,41 +822,44 @@ protected synchronized void markCoordinatorUnknown(boolean isDisconnected) {
     }
 
     /**
-     * Get the current generation state if the group is stable.
-     * @return the current generation or null if the group is unjoined/rebalancing
+     * Get the current generation state, regardless of whether it is currently stable.
+     * Note that the generation information can be updated while we are still in the middle
+     * of a rebalance, after the join-group response is received.
+     *
+     * @return the current generation
      */
     protected synchronized Generation generation() {
-        if (this.state != MemberState.STABLE)
-            return null;
         return generation;
     }
 
     /**
-     * Check whether given generation id is matching the record within current generation.
-     * Only using in unit tests.
-     * @param generationId generation id
-     * @return true if the two ids are matching.
+     * Get the current generation state if the group is stable, otherwise return null
+     *
+     * @return the current generation or null
      */
-    final synchronized boolean hasMatchingGenerationId(int generationId) {
-        return generation != null && generation.generationId == generationId;
+    protected synchronized Generation generationIfStable() {
+        if (this.state != MemberState.STABLE)
+            return null;
+        return generation;
     }
 
-    /**
-     * @return true if the current generation's member ID is valid, false otherwise
-     */
-    // Visible for testing
-    final synchronized boolean hasValidMemberId() {
-        return generation != null && generation.hasMemberId();
+    protected synchronized String memberId() {
+        return generation.memberId;
     }
 
-
-    /**
-     * Reset the generation and memberId because we have fallen out of the group.
-     */
-    protected synchronized void resetGeneration() {
+    private synchronized void resetGeneration() {
         this.generation = Generation.NO_GENERATION;
-        this.rejoinNeeded = true;
-        this.state = MemberState.UNJOINED;
+        resetStateAndRejoin();
+    }
+
+    synchronized void resetGenerationOnResponseError(ApiKeys api, Errors error) {
+        log.debug("Resetting generation after encountering " + error + " from " + api + " response");
+        resetGeneration();
+    }
+
+    synchronized void resetGenerationOnLeaveGroup() {
+        log.debug("Resetting generation due to consumer pro-actively leaving the group");
+        resetGeneration();
     }
 
     protected synchronized void requestRejoin() {
@@ -789,6 +874,9 @@ public final void close() {
         close(time.timer(0));
     }
 
+    /**
+     * @throws KafkaException if the rebalance callback throws exception
+     */
     protected void close(Timer timer) {
         try {
             closeHeartbeatThread();
@@ -796,8 +884,9 @@ protected void close(Timer timer) {
             // Synchronize after closing the heartbeat thread since heartbeat thread
             // needs this lock to complete and terminate after close flag is set.
             synchronized (this) {
-                if (leaveGroupOnClose) {
-                    maybeLeaveGroup();
+                if (rebalanceConfig.leaveGroupOnClose) {
+                    onLeavePrepare();
+                    maybeLeaveGroup("the consumer is being closed");
                 }
 
                 // At this point, there may be pending commits (async commits or sync commits that were
@@ -813,32 +902,53 @@ protected void close(Timer timer) {
     }
 
     /**
-     * Leave the current group and reset local generation/memberId.
+     * @throws KafkaException if the rebalance callback throws exception
      */
-    public synchronized void maybeLeaveGroup() {
-        if (!coordinatorUnknown() && state != MemberState.UNJOINED && generation.hasMemberId()) {
+    public synchronized RequestFuture maybeLeaveGroup(String leaveReason) {
+        RequestFuture future = null;
+
+        // Starting from 2.3, only dynamic members will send LeaveGroupRequest to the broker,
+        // consumer with valid group.instance.id is viewed as static member that never sends LeaveGroup,
+        // and the membership expiration is only controlled by session timeout.
+        if (isDynamicMember() && !coordinatorUnknown() &&
+            state != MemberState.UNJOINED && generation.hasMemberId()) {
             // this is a minimal effort attempt to leave the group. we do not
             // attempt any resending if the request fails or times out.
-            log.info("Member {} sending LeaveGroup request to coordinator {}", generation.memberId, coordinator);
-            LeaveGroupRequest.Builder request = new LeaveGroupRequest.Builder(new LeaveGroupRequestData()
-                    .setGroupId(groupId).setMemberId(generation.memberId));
-            client.send(coordinator, request)
-                    .compose(new LeaveGroupResponseHandler());
+            log.info("Member {} sending LeaveGroup request to coordinator {} due to {}",
+                generation.memberId, coordinator, leaveReason);
+            LeaveGroupRequest.Builder request = new LeaveGroupRequest.Builder(
+                rebalanceConfig.groupId,
+                Collections.singletonList(new MemberIdentity().setMemberId(generation.memberId))
+            );
+
+            future = client.send(coordinator, request).compose(new LeaveGroupResponseHandler());
             client.pollNoWakeup();
         }
 
-        resetGeneration();
+        resetGenerationOnLeaveGroup();
+
+        return future;
+    }
+
+    protected boolean isDynamicMember() {
+        return !rebalanceConfig.groupInstanceId.isPresent();
     }
 
     private class LeaveGroupResponseHandler extends CoordinatorResponseHandler {
         @Override
         public void handle(LeaveGroupResponse leaveResponse, RequestFuture future) {
-            Errors error = leaveResponse.error();
+            final List members = leaveResponse.memberResponses();
+            if (members.size() > 1) {
+                future.raise(new IllegalStateException("The expected leave group response " +
+                                                           "should only contain no more than one member info, however get " + members));
+            }
+
+            final Errors error = leaveResponse.error();
             if (error == Errors.NONE) {
                 log.debug("LeaveGroup request returned successfully");
                 future.complete(null);
             } else {
-                log.debug("LeaveGroup request failed with error: {}", error.message());
+                log.error("LeaveGroup request failed with error: {}", error.message());
                 future.raise(error);
             }
         }
@@ -848,7 +958,11 @@ public void handle(LeaveGroupResponse leaveResponse, RequestFuture future)
     synchronized RequestFuture sendHeartbeatRequest() {
         log.debug("Sending Heartbeat request to coordinator {}", coordinator);
         HeartbeatRequest.Builder requestBuilder =
-                new HeartbeatRequest.Builder(this.groupId, this.generation.generationId, this.generation.memberId);
+                new HeartbeatRequest.Builder(new HeartbeatRequestData()
+                        .setGroupId(rebalanceConfig.groupId)
+                        .setMemberId(this.generation.memberId)
+                        .setGroupInstanceId(this.rebalanceConfig.groupInstanceId.orElse(null))
+                        .setGenerationId(this.generation.generationId));
         return client.send(coordinator, requestBuilder)
                 .compose(new HeartbeatResponseHandler());
     }
@@ -856,7 +970,7 @@ synchronized RequestFuture sendHeartbeatRequest() {
     private class HeartbeatResponseHandler extends CoordinatorResponseHandler {
         @Override
         public void handle(HeartbeatResponse heartbeatResponse, RequestFuture future) {
-            sensors.heartbeatLatency.record(response.requestLatencyMs());
+            sensors.heartbeatSensor.record(response.requestLatencyMs());
             Errors error = heartbeatResponse.error();
             if (error == Errors.NONE) {
                 log.debug("Received successful Heartbeat response");
@@ -870,17 +984,20 @@ public void handle(HeartbeatResponse heartbeatResponse, RequestFuture futu
             } else if (error == Errors.REBALANCE_IN_PROGRESS) {
                 log.info("Attempt to heartbeat failed since group is rebalancing");
                 requestRejoin();
-                future.raise(Errors.REBALANCE_IN_PROGRESS);
+                future.raise(error);
             } else if (error == Errors.ILLEGAL_GENERATION) {
                 log.info("Attempt to heartbeat failed since generation {} is not current", generation.generationId);
-                resetGeneration();
-                future.raise(Errors.ILLEGAL_GENERATION);
+                resetGenerationOnResponseError(ApiKeys.HEARTBEAT, error);
+                future.raise(error);
+            } else if (error == Errors.FENCED_INSTANCE_ID) {
+                log.error("Received fatal exception: group.instance.id gets fenced");
+                future.raise(error);
             } else if (error == Errors.UNKNOWN_MEMBER_ID) {
                 log.info("Attempt to heartbeat failed for since member id {} is not valid.", generation.memberId);
-                resetGeneration();
-                future.raise(Errors.UNKNOWN_MEMBER_ID);
+                resetGenerationOnResponseError(ApiKeys.HEARTBEAT, error);
+                future.raise(error);
             } else if (error == Errors.GROUP_AUTHORIZATION_FAILED) {
-                future.raise(new GroupAuthorizationException(groupId));
+                future.raise(GroupAuthorizationException.forGroupId(rebalanceConfig.groupId));
             } else {
                 future.raise(new KafkaException("Unexpected error in heartbeat response: " + error.message()));
             }
@@ -917,7 +1034,7 @@ public void onSuccess(ClientResponse clientResponse, RequestFuture future) {
     }
 
     protected Meter createMeter(Metrics metrics, String groupName, String baseName, String descriptiveName) {
-        return new Meter(new Count(),
+        return new Meter(new WindowedCount(),
                 metrics.metricName(baseName + "-rate", groupName,
                         String.format("The number of %s per second", descriptiveName)),
                 metrics.metricName(baseName + "-total", groupName,
@@ -927,48 +1044,121 @@ protected Meter createMeter(Metrics metrics, String groupName, String baseName,
     private class GroupCoordinatorMetrics {
         public final String metricGrpName;
 
-        public final Sensor heartbeatLatency;
-        public final Sensor joinLatency;
-        public final Sensor syncLatency;
+        public final Sensor heartbeatSensor;
+        public final Sensor joinSensor;
+        public final Sensor syncSensor;
+        public final Sensor successfulRebalanceSensor;
+        public final Sensor failedRebalanceSensor;
 
         public GroupCoordinatorMetrics(Metrics metrics, String metricGrpPrefix) {
             this.metricGrpName = metricGrpPrefix + "-coordinator-metrics";
 
-            this.heartbeatLatency = metrics.sensor("heartbeat-latency");
-            this.heartbeatLatency.add(metrics.metricName("heartbeat-response-time-max",
+            this.heartbeatSensor = metrics.sensor("heartbeat-latency");
+            this.heartbeatSensor.add(metrics.metricName("heartbeat-response-time-max",
                 this.metricGrpName,
                 "The max time taken to receive a response to a heartbeat request"), new Max());
-            this.heartbeatLatency.add(createMeter(metrics, metricGrpName, "heartbeat", "heartbeats"));
+            this.heartbeatSensor.add(createMeter(metrics, metricGrpName, "heartbeat", "heartbeats"));
 
-            this.joinLatency = metrics.sensor("join-latency");
-            this.joinLatency.add(metrics.metricName("join-time-avg",
-                    this.metricGrpName,
-                    "The average time taken for a group rejoin"), new Avg());
-            this.joinLatency.add(metrics.metricName("join-time-max",
-                    this.metricGrpName,
-                    "The max time taken for a group rejoin"), new Max());
-            this.joinLatency.add(createMeter(metrics, metricGrpName, "join", "group joins"));
+            this.joinSensor = metrics.sensor("join-latency");
+            this.joinSensor.add(metrics.metricName("join-time-avg",
+                this.metricGrpName,
+                "The average time taken for a group rejoin"), new Avg());
+            this.joinSensor.add(metrics.metricName("join-time-max",
+                this.metricGrpName,
+                "The max time taken for a group rejoin"), new Max());
+            this.joinSensor.add(createMeter(metrics, metricGrpName, "join", "group joins"));
 
+            this.syncSensor = metrics.sensor("sync-latency");
+            this.syncSensor.add(metrics.metricName("sync-time-avg",
+                this.metricGrpName,
+                "The average time taken for a group sync"), new Avg());
+            this.syncSensor.add(metrics.metricName("sync-time-max",
+                this.metricGrpName,
+                "The max time taken for a group sync"), new Max());
+            this.syncSensor.add(createMeter(metrics, metricGrpName, "sync", "group syncs"));
 
-            this.syncLatency = metrics.sensor("sync-latency");
-            this.syncLatency.add(metrics.metricName("sync-time-avg",
+            this.successfulRebalanceSensor = metrics.sensor("rebalance-latency");
+            this.successfulRebalanceSensor.add(metrics.metricName("rebalance-latency-avg",
+                this.metricGrpName,
+                "The average time taken for a group to complete a successful rebalance, which may be composed of " +
+                    "several failed re-trials until it succeeded"), new Avg());
+            this.successfulRebalanceSensor.add(metrics.metricName("rebalance-latency-max",
+                this.metricGrpName,
+                "The max time taken for a group to complete a successful rebalance, which may be composed of " +
+                    "several failed re-trials until it succeeded"), new Max());
+            this.successfulRebalanceSensor.add(metrics.metricName("rebalance-latency-total",
+                this.metricGrpName,
+                "The total number of milliseconds this consumer has spent in successful rebalances since creation"),
+                new CumulativeSum());
+            this.successfulRebalanceSensor.add(
+                metrics.metricName("rebalance-total",
                     this.metricGrpName,
-                    "The average time taken for a group sync"), new Avg());
-            this.syncLatency.add(metrics.metricName("sync-time-max",
+                    "The total number of successful rebalance events, each event is composed of " +
+                        "several failed re-trials until it succeeded"),
+                new CumulativeCount()
+            );
+            this.successfulRebalanceSensor.add(
+                metrics.metricName(
+                    "rebalance-rate-per-hour",
                     this.metricGrpName,
-                    "The max time taken for a group sync"), new Max());
-            this.syncLatency.add(createMeter(metrics, metricGrpName, "sync", "group syncs"));
+                    "The number of successful rebalance events per hour, each event is composed of " +
+                        "several failed re-trials until it succeeded"),
+                new Rate(TimeUnit.HOURS, new WindowedCount())
+            );
+
+            this.failedRebalanceSensor = metrics.sensor("failed-rebalance");
+            this.failedRebalanceSensor.add(
+                metrics.metricName("failed-rebalance-total",
+                    this.metricGrpName,
+                    "The total number of failed rebalance events"),
+                new CumulativeCount()
+            );
+            this.failedRebalanceSensor.add(
+                metrics.metricName(
+                    "failed-rebalance-rate-per-hour",
+                    this.metricGrpName,
+                    "The number of failed rebalance events per hour"),
+                new Rate(TimeUnit.HOURS, new WindowedCount())
+            );
+
+            Measurable lastRebalance = (config, now) -> {
+                if (lastRebalanceEndMs == -1L)
+                    // if no rebalance is ever triggered, we just return -1.
+                    return -1d;
+                else
+                    return TimeUnit.SECONDS.convert(now - lastRebalanceEndMs, TimeUnit.MILLISECONDS);
+            };
+            metrics.addMetric(metrics.metricName("last-rebalance-seconds-ago",
+                this.metricGrpName,
+                "The number of seconds since the last successful rebalance event"),
+                lastRebalance);
+
+            Measurable lastHeartbeat = (config, now) -> {
+                if (heartbeat.lastHeartbeatSend() == 0L)
+                    // if no heartbeat is ever triggered, just return -1.
+                    return -1d;
+                else
+                    return TimeUnit.SECONDS.convert(now - heartbeat.lastHeartbeatSend(), TimeUnit.MILLISECONDS);
+            };
+            metrics.addMetric(metrics.metricName("last-heartbeat-seconds-ago",
+                this.metricGrpName,
+                "The number of seconds since the last coordinator heartbeat was sent"),
+                lastHeartbeat);
+
+            //HOTFIX - extra liveliness-related metrics
 
-            Measurable lastHeartbeat =
+            Measurable lastHeartbeatReceived =
                 new Measurable() {
                     public double measure(MetricConfig config, long now) {
-                        return TimeUnit.SECONDS.convert(now - heartbeat.lastHeartbeatSend(), TimeUnit.MILLISECONDS);
+                        return TimeUnit.SECONDS.convert(now - heartbeat.lastHeartbeatReceive(), TimeUnit.MILLISECONDS);
                     }
                 };
-            metrics.addMetric(metrics.metricName("last-heartbeat-seconds-ago",
+            metrics.addMetric(metrics.metricName("last-heartbeat-received-seconds-ago",
                 this.metricGrpName,
-                "The number of seconds since the last coordinator heartbeat was sent"),
-                lastHeartbeat);
+                "The number of seconds since the last successful controller heartbeat was received"),
+                lastHeartbeatReceived);
+
+            //end HOTFIX
         }
     }
 
@@ -978,7 +1168,7 @@ private class HeartbeatThread extends KafkaThread implements AutoCloseable {
         private AtomicReference failed = new AtomicReference<>(null);
 
         private HeartbeatThread() {
-            super(HEARTBEAT_THREAD_PREFIX + (groupId.isEmpty() ? "" : " | " + groupId), true);
+            super(HEARTBEAT_THREAD_PREFIX + (rebalanceConfig.groupId.isEmpty() ? "" : " | " + rebalanceConfig.groupId), true);
         }
 
         public void enable() {
@@ -1040,25 +1230,24 @@ public void run() {
                             if (findCoordinatorFuture != null || lookupCoordinator().failed())
                                 // the immediate future check ensures that we backoff properly in the case that no
                                 // brokers are available to connect to.
-                                AbstractCoordinator.this.wait(retryBackoffMs);
+                                AbstractCoordinator.this.wait(rebalanceConfig.retryBackoffMs);
                         } else if (heartbeat.sessionTimeoutExpired(now)) {
                             // the session timeout has expired without seeing a successful heartbeat, so we should
                             // probably make sure the coordinator is still healthy.
                             markCoordinatorUnknown();
                         } else if (heartbeat.pollTimeoutExpired(now)) {
                             // the poll timeout has expired, which means that the foreground thread has stalled
-                            // in between calls to poll(), so we explicitly leave the group.
-                            log.warn("This member will leave the group because 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.");
-                            maybeLeaveGroup();
+                            // in between calls to poll().
+                            String leaveReason = "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.";
+                            maybeLeaveGroup(leaveReason);
                         } else if (!heartbeat.shouldHeartbeat(now)) {
                             // poll again after waiting for the retry backoff in case the heartbeat failed or the
                             // coordinator disconnected
-                            AbstractCoordinator.this.wait(retryBackoffMs);
+                            AbstractCoordinator.this.wait(rebalanceConfig.retryBackoffMs);
                         } else {
                             heartbeat.sentHeartbeat(now);
 
@@ -1079,9 +1268,12 @@ public void onFailure(RuntimeException e) {
                                             // as the duration of the rebalance timeout. If we stop sending heartbeats,
                                             // however, then the session timeout may expire before we can rejoin.
                                             heartbeat.receiveHeartbeat();
+                                        } else if (e instanceof FencedInstanceIdException) {
+                                            log.error("Caught fenced group.instance.id {} error in heartbeat thread", rebalanceConfig.groupInstanceId);
+                                            heartbeatThread.failed.set(e);
+                                            heartbeatThread.disable();
                                         } else {
                                             heartbeat.failHeartbeat();
-
                                             // wake up the thread if it's sleeping to reschedule the heartbeat
                                             AbstractCoordinator.this.notify();
                                         }
@@ -1144,8 +1336,8 @@ public boolean equals(final Object o) {
             if (o == null || getClass() != o.getClass()) return false;
             final Generation that = (Generation) o;
             return generationId == that.generationId &&
-                Objects.equals(memberId, that.memberId) &&
-                Objects.equals(protocol, that.protocol);
+                    Objects.equals(memberId, that.memberId) &&
+                    Objects.equals(protocol, that.protocol);
         }
 
         @Override
@@ -1167,4 +1359,35 @@ private static class UnjoinedGroupException extends RetriableException {
 
     }
 
+    // For testing only below
+    public Heartbeat heartbeat() {
+        return heartbeat;
+    }
+
+    final void setLastRebalanceTime(final long timestamp) {
+        lastRebalanceEndMs = timestamp;
+    }
+
+    /**
+     * Check whether given generation id is matching the record within current generation.
+     *
+     * @param generationId generation id
+     * @return true if the two ids are matching.
+     */
+    final boolean hasMatchingGenerationId(int generationId) {
+        return generation != Generation.NO_GENERATION && generation.generationId == generationId;
+    }
+
+    final boolean hasUnknownGeneration() {
+        return generation == Generation.NO_GENERATION;
+    }
+
+    /**
+     * @return true if the current generation's member ID is valid, false otherwise
+     */
+    final boolean hasValidMemberId() {
+        return generation != Generation.NO_GENERATION && generation.hasMemberId();
+    }
+
+
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractPartitionAssignor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractPartitionAssignor.java
index 35eb8eb141f93..ed0282b406299 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractPartitionAssignor.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractPartitionAssignor.java
@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.clients.consumer.internals;
 
+import org.apache.kafka.clients.consumer.ConsumerPartitionAssignor;
 import org.apache.kafka.common.Cluster;
 import org.apache.kafka.common.TopicPartition;
 import org.slf4j.Logger;
@@ -26,32 +27,29 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 
 /**
  * Abstract assignor implementation which does some common grunt work (in particular collecting
  * partition counts which are always needed in assignors).
  */
-public abstract class AbstractPartitionAssignor implements PartitionAssignor {
+public abstract class AbstractPartitionAssignor implements ConsumerPartitionAssignor {
     private static final Logger log = LoggerFactory.getLogger(AbstractPartitionAssignor.class);
 
     /**
      * Perform the group assignment given the partition counts and member subscriptions
      * @param partitionsPerTopic The number of partitions for each subscribed topic. Topics not in metadata will be excluded
      *                           from this map.
-     * @param subscriptions Map from the memberId to their respective topic subscription
+     * @param subscriptions Map from the member id to their respective topic subscription
      * @return Map from each member to the list of partitions assigned to them.
      */
     public abstract Map> assign(Map partitionsPerTopic,
                                                              Map subscriptions);
 
     @Override
-    public Subscription subscription(Set topics) {
-        return new Subscription(new ArrayList<>(topics));
-    }
-
-    @Override
-    public Map assign(Cluster metadata, Map subscriptions) {
+    public GroupAssignment assign(Cluster metadata, GroupSubscription groupSubscription) {
+        Map subscriptions = groupSubscription.groupSubscription();
         Set allSubscribedTopics = new HashSet<>();
         for (Map.Entry subscriptionEntry : subscriptions.entrySet())
             allSubscribedTopics.addAll(subscriptionEntry.getValue().topics());
@@ -71,12 +69,7 @@ public Map assign(Cluster metadata, Map assignments = new HashMap<>();
         for (Map.Entry> assignmentEntry : rawAssignments.entrySet())
             assignments.put(assignmentEntry.getKey(), new Assignment(assignmentEntry.getValue()));
-        return assignments;
-    }
-
-    @Override
-    public void onAssignment(Assignment assignment) {
-        // this assignor maintains no internal state, so nothing to do
+        return new GroupAssignment(assignments);
     }
 
     protected static  void put(Map> map, K key, V value) {
@@ -90,4 +83,50 @@ protected static List partitions(String topic, int numPartitions
             partitions.add(new TopicPartition(topic, i));
         return partitions;
     }
+
+    public static class MemberInfo implements Comparable {
+        public final String memberId;
+        public final Optional groupInstanceId;
+
+        public MemberInfo(String memberId, Optional groupInstanceId) {
+            this.memberId = memberId;
+            this.groupInstanceId = groupInstanceId;
+        }
+
+        @Override
+        public int compareTo(MemberInfo otherMemberInfo) {
+            if (this.groupInstanceId.isPresent() &&
+                    otherMemberInfo.groupInstanceId.isPresent()) {
+                return this.groupInstanceId.get()
+                        .compareTo(otherMemberInfo.groupInstanceId.get());
+            } else if (this.groupInstanceId.isPresent()) {
+                return -1;
+            } else if (otherMemberInfo.groupInstanceId.isPresent()) {
+                return 1;
+            } else {
+                return this.memberId.compareTo(otherMemberInfo.memberId);
+            }
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            return o instanceof MemberInfo && this.memberId.equals(((MemberInfo) o).memberId);
+        }
+
+        /**
+         * We could just use member.id to be the hashcode, since it's unique
+         * across the group.
+         */
+        @Override
+        public int hashCode() {
+            return memberId.hashCode();
+        }
+
+        @Override
+        public String toString() {
+            return "MemberInfo [member.id: " + memberId
+                    + ", group.instance.id: " + groupInstanceId.orElse("{}")
+                    + "]";
+        }
+    }
 }
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
new file mode 100644
index 0000000000000..a1af6d959a3f7
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignor.java
@@ -0,0 +1,969 @@
+/*
+ * 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 java.io.Serializable;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.Map;
+import java.util.Map.Entry;
+import java.util.Optional;
+import java.util.Queue;
+import java.util.Set;
+import java.util.SortedSet;
+import java.util.TreeMap;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+import org.apache.kafka.common.TopicPartition;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public abstract class AbstractStickyAssignor extends AbstractPartitionAssignor {
+    private static final Logger log = LoggerFactory.getLogger(AbstractStickyAssignor.class);
+
+    public static final int DEFAULT_GENERATION = -1;
+
+    private PartitionMovements partitionMovements;
+
+    // Keep track of the partitions being migrated from one consumer to another during assignment
+    // so the cooperative assignor can adjust the assignment
+    protected Map partitionsTransferringOwnership = new HashMap<>();
+
+    static final class ConsumerGenerationPair {
+        final String consumer;
+        final int generation;
+        ConsumerGenerationPair(String consumer, int generation) {
+            this.consumer = consumer;
+            this.generation = generation;
+        }
+    }
+
+    public static final class MemberData {
+        public final List partitions;
+        public final Optional generation;
+        public MemberData(List partitions, Optional generation) {
+            this.partitions = partitions;
+            this.generation = generation;
+        }
+    }
+
+    abstract protected MemberData memberData(Subscription subscription);
+
+    @Override
+    public Map> assign(Map partitionsPerTopic,
+                                                    Map subscriptions) {
+        partitionMovements = new PartitionMovements();
+        Map> consumerToOwnedPartitions = new HashMap<>();
+        if (allSubscriptionsEqual(partitionsPerTopic.keySet(), subscriptions, consumerToOwnedPartitions)) {
+            log.debug("Detected that all consumers were subscribed to same set of topics, invoking the "
+                          + "optimized assignment algorithm");
+            partitionsTransferringOwnership = new HashMap<>();
+            return constrainedAssign(partitionsPerTopic, consumerToOwnedPartitions);
+        } else {
+            log.debug("Detected that all not consumers were subscribed to same set of topics, falling back to the "
+                          + "general case assignment algorithm");
+            partitionsTransferringOwnership = null;
+            return generalAssign(partitionsPerTopic, subscriptions);
+        }
+    }
+
+    /**
+     * Returns true iff all consumers have an identical subscription. Also fills out the passed in
+     * {@code consumerToOwnedPartitions} with each consumer's previously owned and still-subscribed partitions
+     */
+    private boolean allSubscriptionsEqual(Set allTopics,
+                                          Map subscriptions,
+                                          Map> consumerToOwnedPartitions) {
+        Set membersWithOldGeneration = new HashSet<>();
+        Set membersOfCurrentHighestGeneration = new HashSet<>();
+        int maxGeneration = DEFAULT_GENERATION;
+
+        Set subscribedTopics = new HashSet<>();
+
+        for (Map.Entry subscriptionEntry : subscriptions.entrySet()) {
+            String consumer = subscriptionEntry.getKey();
+            Subscription subscription = subscriptionEntry.getValue();
+
+            // initialize the subscribed topics set if this is the first subscription
+            if (subscribedTopics.isEmpty()) {
+                subscribedTopics.addAll(subscription.topics());
+            } else if (!(subscription.topics().size() == subscribedTopics.size()
+                && subscribedTopics.containsAll(subscription.topics()))) {
+                return false;
+            }
+
+            MemberData memberData = memberData(subscription);
+
+            List ownedPartitions = new ArrayList<>();
+            consumerToOwnedPartitions.put(consumer, ownedPartitions);
+
+            // Only consider this consumer's owned partitions as valid if it is a member of the current highest
+            // generation, or it's generation is not present but we have not seen any known generation so far
+            if (memberData.generation.isPresent() && memberData.generation.get() >= maxGeneration
+                || !memberData.generation.isPresent() && maxGeneration == DEFAULT_GENERATION) {
+
+                // If the current member's generation is higher, all the previously owned partitions are invalid
+                if (memberData.generation.isPresent() && memberData.generation.get() > maxGeneration) {
+                    membersWithOldGeneration.addAll(membersOfCurrentHighestGeneration);
+                    membersOfCurrentHighestGeneration.clear();
+                    maxGeneration = memberData.generation.get();
+                }
+
+                membersOfCurrentHighestGeneration.add(consumer);
+                for (final TopicPartition tp : memberData.partitions) {
+                    // filter out any topics that no longer exist or aren't part of the current subscription
+                    if (allTopics.contains(tp.topic())) {
+                        ownedPartitions.add(tp);
+                    }
+                }
+            }
+        }
+
+        for (String consumer : membersWithOldGeneration) {
+            consumerToOwnedPartitions.get(consumer).clear();
+        }
+        return true;
+    }
+
+    private Map> constrainedAssign(Map partitionsPerTopic,
+                                                                Map> consumerToOwnedPartitions) {
+        SortedSet unassignedPartitions = getTopicPartitions(partitionsPerTopic);
+
+        Set allRevokedPartitions = new HashSet<>();
+
+        // Each consumer should end up in exactly one of the below
+        List unfilledMembers = new LinkedList<>();
+        Queue maxCapacityMembers = new LinkedList<>();
+        Queue minCapacityMembers = new LinkedList<>();
+
+        int numberOfConsumers = consumerToOwnedPartitions.size();
+        int minQuota = (int) Math.floor(((double) unassignedPartitions.size()) / numberOfConsumers);
+        int maxQuota = (int) Math.ceil(((double) unassignedPartitions.size()) / numberOfConsumers);
+
+        // initialize the assignment map with an empty array of size minQuota for all members
+        Map> assignment = new HashMap<>(
+            consumerToOwnedPartitions.keySet().stream().collect(Collectors.toMap(c -> c, c -> new ArrayList<>(minQuota))));
+
+        for (Map.Entry> consumerEntry : consumerToOwnedPartitions.entrySet()) {
+            String consumer = consumerEntry.getKey();
+            List ownedPartitions = consumerEntry.getValue();
+
+            List consumerAssignment = assignment.get(consumer);
+            int i = 0;
+            // assign the first N partitions up to the max quota, and mark the remaining as being revoked
+            for (TopicPartition tp : ownedPartitions) {
+                if (i < maxQuota) {
+                    consumerAssignment.add(tp);
+                    unassignedPartitions.remove(tp);
+                } else {
+                    allRevokedPartitions.add(tp);
+                }
+                ++i;
+            }
+
+            if (ownedPartitions.size() < minQuota) {
+                unfilledMembers.add(consumer);
+            } else {
+                // It's possible for a consumer to be at both min and max capacity if minQuota == maxQuota
+                if (consumerAssignment.size() == minQuota)
+                    minCapacityMembers.add(consumer);
+                if (consumerAssignment.size() == maxQuota)
+                    maxCapacityMembers.add(consumer);
+            }
+        }
+
+        Collections.sort(unfilledMembers);
+        Iterator unassignedPartitionsIter = unassignedPartitions.iterator();
+
+        while (!unfilledMembers.isEmpty() && !unassignedPartitions.isEmpty()) {
+            Iterator unfilledConsumerIter = unfilledMembers.iterator();
+
+            while (unfilledConsumerIter.hasNext()) {
+                String consumer = unfilledConsumerIter.next();
+                List consumerAssignment = assignment.get(consumer);
+
+                if (unassignedPartitionsIter.hasNext()) {
+                    TopicPartition tp = unassignedPartitionsIter.next();
+                    consumerAssignment.add(tp);
+                    unassignedPartitionsIter.remove();
+                    // We already assigned all possible ownedPartitions, so we know this must be newly to this consumer
+                    if (allRevokedPartitions.contains(tp))
+                        partitionsTransferringOwnership.put(tp, consumer);
+                } else {
+                    break;
+                }
+
+                if (consumerAssignment.size() == minQuota) {
+                    minCapacityMembers.add(consumer);
+                    unfilledConsumerIter.remove();
+                }
+            }
+        }
+
+        // If we ran out of unassigned partitions before filling all consumers, we need to start stealing partitions
+        // from the over-full consumers at max capacity
+        for (String consumer : unfilledMembers) {
+            List consumerAssignment = assignment.get(consumer);
+            int remainingCapacity = minQuota - consumerAssignment.size();
+            while (remainingCapacity > 0) {
+                String overloadedConsumer = maxCapacityMembers.poll();
+                if (overloadedConsumer == null) {
+                    throw new IllegalStateException("Some consumers are under capacity but all partitions have been assigned");
+                }
+                TopicPartition swappedPartition = assignment.get(overloadedConsumer).remove(0);
+                consumerAssignment.add(swappedPartition);
+                --remainingCapacity;
+                // This partition is by definition transferring ownership, the swapped partition must have come from
+                // the max capacity member's owned partitions since it can only reach max capacity with owned partitions
+                partitionsTransferringOwnership.put(swappedPartition, consumer);
+            }
+            minCapacityMembers.add(consumer);
+        }
+
+        // Otherwise we may have run out of unfilled consumers before assigning all partitions, in which case we
+        // should just distribute one partition each to all consumers at min capacity
+        for (TopicPartition unassignedPartition : unassignedPartitions) {
+            String underCapacityConsumer = minCapacityMembers.poll();
+            if (underCapacityConsumer == null) {
+                throw new IllegalStateException("Some partitions are unassigned but all consumers are at maximum capacity");
+            }
+            // We can skip the bookkeeping of unassignedPartitions and maxCapacityMembers here since we are at the end
+            assignment.get(underCapacityConsumer).add(unassignedPartition);
+
+            if (allRevokedPartitions.contains(unassignedPartition))
+                partitionsTransferringOwnership.put(unassignedPartition, underCapacityConsumer);
+        }
+
+        return assignment;
+    }
+
+    private SortedSet getTopicPartitions(Map partitionsPerTopic) {
+        SortedSet allPartitions =
+            new TreeSet<>(Comparator.comparing(TopicPartition::topic).thenComparing(TopicPartition::partition));
+        for (Entry entry: partitionsPerTopic.entrySet()) {
+            String topic = entry.getKey();
+            for (int i = 0; i < entry.getValue(); ++i) {
+                allPartitions.add(new TopicPartition(topic, i));
+            }
+        }
+        return allPartitions;
+    }
+
+    private Map> generalAssign(Map partitionsPerTopic,
+                                                            Map subscriptions) {
+        Map> currentAssignment = new HashMap<>();
+        Map prevAssignment = new HashMap<>();
+
+        prepopulateCurrentAssignments(subscriptions, currentAssignment, prevAssignment);
+
+        // a mapping of all topic partitions to all consumers that can be assigned to them
+        final Map> partition2AllPotentialConsumers = new HashMap<>();
+        // a mapping of all consumers to all potential topic partitions that can be assigned to them
+        final Map> consumer2AllPotentialPartitions = new HashMap<>();
+
+        // initialize partition2AllPotentialConsumers and consumer2AllPotentialPartitions in the following two for loops
+        for (Entry entry: partitionsPerTopic.entrySet()) {
+            for (int i = 0; i < entry.getValue(); ++i)
+                partition2AllPotentialConsumers.put(new TopicPartition(entry.getKey(), i), new ArrayList<>());
+        }
+
+        for (Entry entry: subscriptions.entrySet()) {
+            String consumerId = entry.getKey();
+            consumer2AllPotentialPartitions.put(consumerId, new ArrayList<>());
+            entry.getValue().topics().stream().filter(topic -> partitionsPerTopic.get(topic) != null).forEach(topic -> {
+                for (int i = 0; i < partitionsPerTopic.get(topic); ++i) {
+                    TopicPartition topicPartition = new TopicPartition(topic, i);
+                    consumer2AllPotentialPartitions.get(consumerId).add(topicPartition);
+                    partition2AllPotentialConsumers.get(topicPartition).add(consumerId);
+                }
+            });
+
+            // add this consumer to currentAssignment (with an empty topic partition assignment) if it does not already exist
+            if (!currentAssignment.containsKey(consumerId))
+                currentAssignment.put(consumerId, new ArrayList<>());
+        }
+
+        // a mapping of partition to current consumer
+        Map currentPartitionConsumer = new HashMap<>();
+        for (Map.Entry> entry: currentAssignment.entrySet())
+            for (TopicPartition topicPartition: entry.getValue())
+                currentPartitionConsumer.put(topicPartition, entry.getKey());
+
+        List sortedPartitions = sortPartitions(partition2AllPotentialConsumers);
+
+        // all partitions that need to be assigned (initially set to all partitions but adjusted in the following loop)
+        List unassignedPartitions = new ArrayList<>(sortedPartitions);
+        boolean revocationRequired = false;
+        for (Iterator>> it = currentAssignment.entrySet().iterator(); it.hasNext();) {
+            Map.Entry> entry = it.next();
+            if (!subscriptions.containsKey(entry.getKey())) {
+                // if a consumer that existed before (and had some partition assignments) is now removed, remove it from currentAssignment
+                for (TopicPartition topicPartition: entry.getValue())
+                    currentPartitionConsumer.remove(topicPartition);
+                it.remove();
+            } else {
+                // otherwise (the consumer still exists)
+                for (Iterator partitionIter = entry.getValue().iterator(); partitionIter.hasNext();) {
+                    TopicPartition partition = partitionIter.next();
+                    if (!partition2AllPotentialConsumers.containsKey(partition)) {
+                        // if this topic partition of this consumer no longer exists remove it from currentAssignment of the consumer
+                        partitionIter.remove();
+                        currentPartitionConsumer.remove(partition);
+                    } else if (!subscriptions.get(entry.getKey()).topics().contains(partition.topic())) {
+                        // if this partition cannot remain assigned to its current consumer because the consumer
+                        // is no longer subscribed to its topic remove it from currentAssignment of the consumer
+                        partitionIter.remove();
+                        revocationRequired = true;
+                    } else
+                        // otherwise, remove the topic partition from those that need to be assigned only if
+                        // its current consumer is still subscribed to its topic (because it is already assigned
+                        // and we would want to preserve that assignment as much as possible)
+                        unassignedPartitions.remove(partition);
+                }
+            }
+        }
+        // at this point we have preserved all valid topic partition to consumer assignments and removed
+        // all invalid topic partitions and invalid consumers. Now we need to assign unassignedPartitions
+        // to consumers so that the topic partition assignments are as balanced as possible.
+
+        // an ascending sorted set of consumers based on how many topic partitions are already assigned to them
+        TreeSet sortedCurrentSubscriptions = new TreeSet<>(new SubscriptionComparator(currentAssignment));
+        sortedCurrentSubscriptions.addAll(currentAssignment.keySet());
+
+        balance(currentAssignment, prevAssignment, sortedPartitions, unassignedPartitions, sortedCurrentSubscriptions,
+            consumer2AllPotentialPartitions, partition2AllPotentialConsumers, currentPartitionConsumer, revocationRequired);
+        return currentAssignment;
+    }
+
+    private void prepopulateCurrentAssignments(Map subscriptions,
+                                               Map> currentAssignment,
+                                               Map prevAssignment) {
+        // we need to process subscriptions' user data with each consumer's reported generation in mind
+        // higher generations overwrite lower generations in case of a conflict
+        // note that a conflict could exists only if user data is for different generations
+
+        // for each partition we create a sorted map of its consumers by generation
+        Map> sortedPartitionConsumersByGeneration = new HashMap<>();
+        for (Map.Entry subscriptionEntry: subscriptions.entrySet()) {
+            String consumer = subscriptionEntry.getKey();
+            MemberData memberData = memberData(subscriptionEntry.getValue());
+
+            for (TopicPartition partition: memberData.partitions) {
+                if (sortedPartitionConsumersByGeneration.containsKey(partition)) {
+                    Map consumers = sortedPartitionConsumersByGeneration.get(partition);
+                    if (memberData.generation.isPresent() && consumers.containsKey(memberData.generation.get())) {
+                        // same partition is assigned to two consumers during the same rebalance.
+                        // log a warning and skip this record
+                        log.warn("Partition '{}' is assigned to multiple consumers following sticky assignment generation {}.",
+                            partition, memberData.generation);
+                    } else
+                        consumers.put(memberData.generation.orElse(DEFAULT_GENERATION), consumer);
+                } else {
+                    TreeMap sortedConsumers = new TreeMap<>();
+                    sortedConsumers.put(memberData.generation.orElse(DEFAULT_GENERATION), consumer);
+                    sortedPartitionConsumersByGeneration.put(partition, sortedConsumers);
+                }
+            }
+        }
+
+        // prevAssignment holds the prior ConsumerGenerationPair (before current) of each partition
+        // current and previous consumers are the last two consumers of each partition in the above sorted map
+        for (Map.Entry> partitionConsumersEntry: sortedPartitionConsumersByGeneration.entrySet()) {
+            TopicPartition partition = partitionConsumersEntry.getKey();
+            TreeMap consumers = partitionConsumersEntry.getValue();
+            Iterator it = consumers.descendingKeySet().iterator();
+
+            // let's process the current (most recent) consumer first
+            String consumer = consumers.get(it.next());
+            currentAssignment.computeIfAbsent(consumer, k -> new ArrayList<>());
+            currentAssignment.get(consumer).add(partition);
+
+            // now update previous assignment if any
+            if (it.hasNext()) {
+                int generation = it.next();
+                prevAssignment.put(partition, new ConsumerGenerationPair(consumers.get(generation), generation));
+            }
+        }
+    }
+
+    /**
+     * determine if the current assignment is a balanced one
+     *
+     * @param currentAssignment: the assignment whose balance needs to be checked
+     * @param sortedCurrentSubscriptions: an ascending sorted set of consumers based on how many topic partitions are already assigned to them
+     * @param allSubscriptions: a mapping of all consumers to all potential topic partitions that can be assigned to them
+     * @return true if the given assignment is balanced; false otherwise
+     */
+    private boolean isBalanced(Map> currentAssignment,
+                               TreeSet sortedCurrentSubscriptions,
+                               Map> allSubscriptions) {
+        int min = currentAssignment.get(sortedCurrentSubscriptions.first()).size();
+        int max = currentAssignment.get(sortedCurrentSubscriptions.last()).size();
+        if (min >= max - 1)
+            // if minimum and maximum numbers of partitions assigned to consumers differ by at most one return true
+            return true;
+
+        // create a mapping from partitions to the consumer assigned to them
+        final Map allPartitions = new HashMap<>();
+        Set>> assignments = currentAssignment.entrySet();
+        for (Map.Entry> entry: assignments) {
+            List topicPartitions = entry.getValue();
+            for (TopicPartition topicPartition: topicPartitions) {
+                if (allPartitions.containsKey(topicPartition))
+                    log.error("{} is assigned to more than one consumer.", topicPartition);
+                allPartitions.put(topicPartition, entry.getKey());
+            }
+        }
+
+        // for each consumer that does not have all the topic partitions it can get make sure none of the topic partitions it
+        // could but did not get cannot be moved to it (because that would break the balance)
+        for (String consumer: sortedCurrentSubscriptions) {
+            List consumerPartitions = currentAssignment.get(consumer);
+            int consumerPartitionCount = consumerPartitions.size();
+
+            // skip if this consumer already has all the topic partitions it can get
+            if (consumerPartitionCount == allSubscriptions.get(consumer).size())
+                continue;
+
+            // otherwise make sure it cannot get any more
+            List potentialTopicPartitions = allSubscriptions.get(consumer);
+            for (TopicPartition topicPartition: potentialTopicPartitions) {
+                if (!currentAssignment.get(consumer).contains(topicPartition)) {
+                    String otherConsumer = allPartitions.get(topicPartition);
+                    int otherConsumerPartitionCount = currentAssignment.get(otherConsumer).size();
+                    if (consumerPartitionCount < otherConsumerPartitionCount) {
+                        log.debug("{} can be moved from consumer {} to consumer {} for a more balanced assignment.",
+                            topicPartition, otherConsumer, consumer);
+                        return false;
+                    }
+                }
+            }
+        }
+        return true;
+    }
+
+    /**
+     * @return the balance score of the given assignment, as the sum of assigned partitions size difference of all consumer pairs.
+     * A perfectly balanced assignment (with all consumers getting the same number of partitions) has a balance score of 0.
+     * Lower balance score indicates a more balanced assignment.
+     */
+    private int getBalanceScore(Map> assignment) {
+        int score = 0;
+
+        Map consumer2AssignmentSize = new HashMap<>();
+        for (Entry> entry: assignment.entrySet())
+            consumer2AssignmentSize.put(entry.getKey(), entry.getValue().size());
+
+        Iterator> it = consumer2AssignmentSize.entrySet().iterator();
+        while (it.hasNext()) {
+            Entry entry = it.next();
+            int consumerAssignmentSize = entry.getValue();
+            it.remove();
+            for (Entry otherEntry: consumer2AssignmentSize.entrySet())
+                score += Math.abs(consumerAssignmentSize - otherEntry.getValue());
+        }
+
+        return score;
+    }
+
+    /**
+     * Sort valid partitions so they are processed in the potential reassignment phase in the proper order
+     * that causes minimal partition movement among consumers (hence honoring maximal stickiness)
+     *
+     * @param partition2AllPotentialConsumers a mapping of partitions to their potential consumers
+     * @return  an ascending sorted list of topic partitions based on how many consumers can potentially use them
+     */
+    private List sortPartitions(Map> partition2AllPotentialConsumers) {
+        List sortedPartitions = new ArrayList<>(partition2AllPotentialConsumers.keySet());
+        Collections.sort(sortedPartitions, new PartitionComparator(partition2AllPotentialConsumers));
+        return sortedPartitions;
+    }
+
+    /**
+     * The assignment should improve the overall balance of the partition assignments to consumers.
+     */
+    private void assignPartition(TopicPartition partition,
+                                 TreeSet sortedCurrentSubscriptions,
+                                 Map> currentAssignment,
+                                 Map> consumer2AllPotentialPartitions,
+                                 Map currentPartitionConsumer) {
+        for (String consumer: sortedCurrentSubscriptions) {
+            if (consumer2AllPotentialPartitions.get(consumer).contains(partition)) {
+                sortedCurrentSubscriptions.remove(consumer);
+                currentAssignment.get(consumer).add(partition);
+                currentPartitionConsumer.put(partition, consumer);
+                sortedCurrentSubscriptions.add(consumer);
+                break;
+            }
+        }
+    }
+
+    private boolean canParticipateInReassignment(TopicPartition partition,
+                                                 Map> partition2AllPotentialConsumers) {
+        // if a partition has two or more potential consumers it is subject to reassignment.
+        return partition2AllPotentialConsumers.get(partition).size() >= 2;
+    }
+
+    private boolean canParticipateInReassignment(String consumer,
+                                                 Map> currentAssignment,
+                                                 Map> consumer2AllPotentialPartitions,
+                                                 Map> partition2AllPotentialConsumers) {
+        List currentPartitions = currentAssignment.get(consumer);
+        int currentAssignmentSize = currentPartitions.size();
+        int maxAssignmentSize = consumer2AllPotentialPartitions.get(consumer).size();
+        if (currentAssignmentSize > maxAssignmentSize)
+            log.error("The consumer {} is assigned more partitions than the maximum possible.", consumer);
+
+        if (currentAssignmentSize < maxAssignmentSize)
+            // if a consumer is not assigned all its potential partitions it is subject to reassignment
+            return true;
+
+        for (TopicPartition partition: currentPartitions)
+            // if any of the partitions assigned to a consumer is subject to reassignment the consumer itself
+            // is subject to reassignment
+            if (canParticipateInReassignment(partition, partition2AllPotentialConsumers))
+                return true;
+
+        return false;
+    }
+
+    /**
+     * Balance the current assignment using the data structures created in the assign(...) method above.
+     */
+    private void balance(Map> currentAssignment,
+                         Map prevAssignment,
+                         List sortedPartitions,
+                         List unassignedPartitions,
+                         TreeSet sortedCurrentSubscriptions,
+                         Map> consumer2AllPotentialPartitions,
+                         Map> partition2AllPotentialConsumers,
+                         Map currentPartitionConsumer,
+                         boolean revocationRequired) {
+        boolean initializing = currentAssignment.get(sortedCurrentSubscriptions.last()).isEmpty();
+        boolean reassignmentPerformed = false;
+
+        // assign all unassigned partitions
+        for (TopicPartition partition: unassignedPartitions) {
+            // skip if there is no potential consumer for the partition
+            if (partition2AllPotentialConsumers.get(partition).isEmpty())
+                continue;
+
+            assignPartition(partition, sortedCurrentSubscriptions, currentAssignment,
+                consumer2AllPotentialPartitions, currentPartitionConsumer);
+        }
+
+        // narrow down the reassignment scope to only those partitions that can actually be reassigned
+        Set fixedPartitions = new HashSet<>();
+        for (TopicPartition partition: partition2AllPotentialConsumers.keySet())
+            if (!canParticipateInReassignment(partition, partition2AllPotentialConsumers))
+                fixedPartitions.add(partition);
+        sortedPartitions.removeAll(fixedPartitions);
+        unassignedPartitions.removeAll(fixedPartitions);
+
+        // narrow down the reassignment scope to only those consumers that are subject to reassignment
+        Map> fixedAssignments = new HashMap<>();
+        for (String consumer: consumer2AllPotentialPartitions.keySet())
+            if (!canParticipateInReassignment(consumer, currentAssignment,
+                consumer2AllPotentialPartitions, partition2AllPotentialConsumers)) {
+                sortedCurrentSubscriptions.remove(consumer);
+                fixedAssignments.put(consumer, currentAssignment.remove(consumer));
+            }
+
+        // create a deep copy of the current assignment so we can revert to it if we do not get a more balanced assignment later
+        Map> preBalanceAssignment = deepCopy(currentAssignment);
+        Map preBalancePartitionConsumers = new HashMap<>(currentPartitionConsumer);
+
+        // if we don't already need to revoke something due to subscription changes, first try to balance by only moving newly added partitions
+        if (!revocationRequired) {
+            performReassignments(unassignedPartitions, currentAssignment, prevAssignment, sortedCurrentSubscriptions,
+                consumer2AllPotentialPartitions, partition2AllPotentialConsumers, currentPartitionConsumer);
+        }
+
+        reassignmentPerformed = performReassignments(sortedPartitions, currentAssignment, prevAssignment, sortedCurrentSubscriptions,
+                   consumer2AllPotentialPartitions, partition2AllPotentialConsumers, currentPartitionConsumer);
+
+        // if we are not preserving existing assignments and we have made changes to the current assignment
+        // make sure we are getting a more balanced assignment; otherwise, revert to previous assignment
+        if (!initializing && reassignmentPerformed && getBalanceScore(currentAssignment) >= getBalanceScore(preBalanceAssignment)) {
+            deepCopy(preBalanceAssignment, currentAssignment);
+            currentPartitionConsumer.clear();
+            currentPartitionConsumer.putAll(preBalancePartitionConsumers);
+        }
+
+        // add the fixed assignments (those that could not change) back
+        for (Entry> entry: fixedAssignments.entrySet()) {
+            String consumer = entry.getKey();
+            currentAssignment.put(consumer, entry.getValue());
+            sortedCurrentSubscriptions.add(consumer);
+        }
+
+        fixedAssignments.clear();
+    }
+
+    private boolean performReassignments(List reassignablePartitions,
+                                         Map> currentAssignment,
+                                         Map prevAssignment,
+                                         TreeSet sortedCurrentSubscriptions,
+                                         Map> consumer2AllPotentialPartitions,
+                                         Map> partition2AllPotentialConsumers,
+                                         Map currentPartitionConsumer) {
+        boolean reassignmentPerformed = false;
+        boolean modified;
+
+        // repeat reassignment until no partition can be moved to improve the balance
+        do {
+            modified = false;
+            // reassign all reassignable partitions (starting from the partition with least potential consumers and if needed)
+            // until the full list is processed or a balance is achieved
+            Iterator partitionIterator = reassignablePartitions.iterator();
+            while (partitionIterator.hasNext() && !isBalanced(currentAssignment, sortedCurrentSubscriptions, consumer2AllPotentialPartitions)) {
+                TopicPartition partition = partitionIterator.next();
+
+                // the partition must have at least two consumers
+                if (partition2AllPotentialConsumers.get(partition).size() <= 1)
+                    log.error("Expected more than one potential consumer for partition '{}'", partition);
+
+                // the partition must have a current consumer
+                String consumer = currentPartitionConsumer.get(partition);
+                if (consumer == null)
+                    log.error("Expected partition '{}' to be assigned to a consumer", partition);
+
+                if (prevAssignment.containsKey(partition) &&
+                    currentAssignment.get(consumer).size() > currentAssignment.get(prevAssignment.get(partition).consumer).size() + 1) {
+                    reassignPartition(partition, currentAssignment, sortedCurrentSubscriptions, currentPartitionConsumer, prevAssignment.get(partition).consumer);
+                    reassignmentPerformed = true;
+                    modified = true;
+                    continue;
+                }
+
+                // check if a better-suited consumer exist for the partition; if so, reassign it
+                for (String otherConsumer: partition2AllPotentialConsumers.get(partition)) {
+                    if (currentAssignment.get(consumer).size() > currentAssignment.get(otherConsumer).size() + 1) {
+                        reassignPartition(partition, currentAssignment, sortedCurrentSubscriptions, currentPartitionConsumer, consumer2AllPotentialPartitions);
+                        reassignmentPerformed = true;
+                        modified = true;
+                        break;
+                    }
+                }
+            }
+        } while (modified);
+
+        return reassignmentPerformed;
+    }
+
+    private void reassignPartition(TopicPartition partition,
+                                   Map> currentAssignment,
+                                   TreeSet sortedCurrentSubscriptions,
+                                   Map currentPartitionConsumer,
+                                   Map> consumer2AllPotentialPartitions) {
+        // find the new consumer
+        String newConsumer = null;
+        for (String anotherConsumer: sortedCurrentSubscriptions) {
+            if (consumer2AllPotentialPartitions.get(anotherConsumer).contains(partition)) {
+                newConsumer = anotherConsumer;
+                break;
+            }
+        }
+
+        assert newConsumer != null;
+
+        reassignPartition(partition, currentAssignment, sortedCurrentSubscriptions, currentPartitionConsumer, newConsumer);
+    }
+
+    private void reassignPartition(TopicPartition partition,
+                                   Map> currentAssignment,
+                                   TreeSet sortedCurrentSubscriptions,
+                                   Map currentPartitionConsumer,
+                                   String newConsumer) {
+        String consumer = currentPartitionConsumer.get(partition);
+        // find the correct partition movement considering the stickiness requirement
+        TopicPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(partition, consumer, newConsumer);
+        processPartitionMovement(partitionToBeMoved, newConsumer, currentAssignment, sortedCurrentSubscriptions, currentPartitionConsumer);
+    }
+
+    private void processPartitionMovement(TopicPartition partition,
+                                          String newConsumer,
+                                          Map> currentAssignment,
+                                          TreeSet sortedCurrentSubscriptions,
+                                          Map currentPartitionConsumer) {
+        String oldConsumer = currentPartitionConsumer.get(partition);
+
+        sortedCurrentSubscriptions.remove(oldConsumer);
+        sortedCurrentSubscriptions.remove(newConsumer);
+
+        partitionMovements.movePartition(partition, oldConsumer, newConsumer);
+
+        currentAssignment.get(oldConsumer).remove(partition);
+        currentAssignment.get(newConsumer).add(partition);
+        currentPartitionConsumer.put(partition, newConsumer);
+        sortedCurrentSubscriptions.add(newConsumer);
+        sortedCurrentSubscriptions.add(oldConsumer);
+    }
+
+    public boolean isSticky() {
+        return partitionMovements.isSticky();
+    }
+
+    private void deepCopy(Map> source, Map> dest) {
+        dest.clear();
+        for (Entry> entry: source.entrySet())
+            dest.put(entry.getKey(), new ArrayList<>(entry.getValue()));
+    }
+
+    private Map> deepCopy(Map> assignment) {
+        Map> copy = new HashMap<>();
+        deepCopy(assignment, copy);
+        return copy;
+    }
+
+    private static class PartitionComparator implements Comparator, Serializable {
+        private static final long serialVersionUID = 1L;
+        private Map> map;
+
+        PartitionComparator(Map> map) {
+            this.map = map;
+        }
+
+        @Override
+        public int compare(TopicPartition o1, TopicPartition o2) {
+            int ret = map.get(o1).size() - map.get(o2).size();
+            if (ret == 0) {
+                ret = o1.topic().compareTo(o2.topic());
+                if (ret == 0)
+                    ret = o1.partition() - o2.partition();
+            }
+            return ret;
+        }
+    }
+
+    private static class SubscriptionComparator implements Comparator, Serializable {
+        private static final long serialVersionUID = 1L;
+        private Map> map;
+
+        SubscriptionComparator(Map> map) {
+            this.map = map;
+        }
+
+        @Override
+        public int compare(String o1, String o2) {
+            int ret = map.get(o1).size() - map.get(o2).size();
+            if (ret == 0)
+                ret = o1.compareTo(o2);
+            return ret;
+        }
+    }
+
+    /**
+     * This class maintains some data structures to simplify lookup of partition movements among consumers. At each point of
+     * time during a partition rebalance it keeps track of partition movements corresponding to each topic, and also possible
+     * movement (in form a ConsumerPair object) for each partition.
+     */
+    private static class PartitionMovements {
+        private Map>> partitionMovementsByTopic = new HashMap<>();
+        private Map partitionMovements = new HashMap<>();
+
+        private ConsumerPair removeMovementRecordOfPartition(TopicPartition partition) {
+            ConsumerPair pair = partitionMovements.remove(partition);
+
+            String topic = partition.topic();
+            Map> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+            partitionMovementsForThisTopic.get(pair).remove(partition);
+            if (partitionMovementsForThisTopic.get(pair).isEmpty())
+                partitionMovementsForThisTopic.remove(pair);
+            if (partitionMovementsByTopic.get(topic).isEmpty())
+                partitionMovementsByTopic.remove(topic);
+
+            return pair;
+        }
+
+        private void addPartitionMovementRecord(TopicPartition partition, ConsumerPair pair) {
+            partitionMovements.put(partition, pair);
+
+            String topic = partition.topic();
+            if (!partitionMovementsByTopic.containsKey(topic))
+                partitionMovementsByTopic.put(topic, new HashMap<>());
+
+            Map> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+            if (!partitionMovementsForThisTopic.containsKey(pair))
+                partitionMovementsForThisTopic.put(pair, new HashSet<>());
+
+            partitionMovementsForThisTopic.get(pair).add(partition);
+        }
+
+        private void movePartition(TopicPartition partition, String oldConsumer, String newConsumer) {
+            ConsumerPair pair = new ConsumerPair(oldConsumer, newConsumer);
+
+            if (partitionMovements.containsKey(partition)) {
+                // this partition has previously moved
+                ConsumerPair existingPair = removeMovementRecordOfPartition(partition);
+                assert existingPair.dstMemberId.equals(oldConsumer);
+                if (!existingPair.srcMemberId.equals(newConsumer)) {
+                    // the partition is not moving back to its previous consumer
+                    // return new ConsumerPair2(existingPair.src, newConsumer);
+                    addPartitionMovementRecord(partition, new ConsumerPair(existingPair.srcMemberId, newConsumer));
+                }
+            } else
+                addPartitionMovementRecord(partition, pair);
+        }
+
+        private TopicPartition getTheActualPartitionToBeMoved(TopicPartition partition, String oldConsumer, String newConsumer) {
+            String topic = partition.topic();
+
+            if (!partitionMovementsByTopic.containsKey(topic))
+                return partition;
+
+            if (partitionMovements.containsKey(partition)) {
+                // this partition has previously moved
+                assert oldConsumer.equals(partitionMovements.get(partition).dstMemberId);
+                oldConsumer = partitionMovements.get(partition).srcMemberId;
+            }
+
+            Map> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+            ConsumerPair reversePair = new ConsumerPair(newConsumer, oldConsumer);
+            if (!partitionMovementsForThisTopic.containsKey(reversePair))
+                return partition;
+
+            return partitionMovementsForThisTopic.get(reversePair).iterator().next();
+        }
+
+        private boolean isLinked(String src, String dst, Set pairs, List currentPath) {
+            if (src.equals(dst))
+                return false;
+
+            if (pairs.isEmpty())
+                return false;
+
+            if (new ConsumerPair(src, dst).in(pairs)) {
+                currentPath.add(src);
+                currentPath.add(dst);
+                return true;
+            }
+
+            for (ConsumerPair pair: pairs)
+                if (pair.srcMemberId.equals(src)) {
+                    Set reducedSet = new HashSet<>(pairs);
+                    reducedSet.remove(pair);
+                    currentPath.add(pair.srcMemberId);
+                    return isLinked(pair.dstMemberId, dst, reducedSet, currentPath);
+                }
+
+            return false;
+        }
+
+        private boolean in(List cycle, Set> cycles) {
+            List superCycle = new ArrayList<>(cycle);
+            superCycle.remove(superCycle.size() - 1);
+            superCycle.addAll(cycle);
+            for (List foundCycle: cycles) {
+                if (foundCycle.size() == cycle.size() && Collections.indexOfSubList(superCycle, foundCycle) != -1)
+                    return true;
+            }
+            return false;
+        }
+
+        private boolean hasCycles(Set pairs) {
+            Set> cycles = new HashSet<>();
+            for (ConsumerPair pair: pairs) {
+                Set reducedPairs = new HashSet<>(pairs);
+                reducedPairs.remove(pair);
+                List path = new ArrayList<>(Collections.singleton(pair.srcMemberId));
+                if (isLinked(pair.dstMemberId, pair.srcMemberId, reducedPairs, path) && !in(path, cycles)) {
+                    cycles.add(new ArrayList<>(path));
+                    log.error("A cycle of length {} was found: {}", path.size() - 1, path.toString());
+                }
+            }
+
+            // for now we want to make sure there is no partition movements of the same topic between a pair of consumers.
+            // the odds of finding a cycle among more than two consumers seem to be very low (according to various randomized
+            // tests with the given sticky algorithm) that it should not worth the added complexity of handling those cases.
+            for (List cycle: cycles)
+                if (cycle.size() == 3) // indicates a cycle of length 2
+                    return true;
+            return false;
+        }
+
+        private boolean isSticky() {
+            for (Map.Entry>> topicMovements: this.partitionMovementsByTopic.entrySet()) {
+                Set topicMovementPairs = topicMovements.getValue().keySet();
+                if (hasCycles(topicMovementPairs)) {
+                    log.error("Stickiness is violated for topic {}"
+                        + "\nPartition movements for this topic occurred among the following consumer pairs:"
+                        + "\n{}", topicMovements.getKey(), topicMovements.getValue().toString());
+                    return false;
+                }
+            }
+
+            return true;
+        }
+    }
+
+    /**
+     * ConsumerPair represents a pair of Kafka consumer ids involved in a partition reassignment. Each
+     * ConsumerPair object, which contains a source (src) and a destination (dst)
+     * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
+     * partition of the particular topic was moved from the source consumer to the destination consumer during the rebalance.
+     * This class is used, through the PartitionMovements class, by the sticky assignor and helps in determining
+     * whether a partition reassignment results in cycles among the generated graph of consumer pairs.
+     */
+    private static class ConsumerPair {
+        private final String srcMemberId;
+        private final String dstMemberId;
+
+        ConsumerPair(String srcMemberId, String dstMemberId) {
+            this.srcMemberId = srcMemberId;
+            this.dstMemberId = dstMemberId;
+        }
+
+        public String toString() {
+            return this.srcMemberId + "->" + this.dstMemberId;
+        }
+
+        @Override
+        public int hashCode() {
+            final int prime = 31;
+            int result = 1;
+            result = prime * result + ((this.srcMemberId == null) ? 0 : this.srcMemberId.hashCode());
+            result = prime * result + ((this.dstMemberId == null) ? 0 : this.dstMemberId.hashCode());
+            return result;
+        }
+
+        @Override
+        public boolean equals(Object obj) {
+            if (obj == null)
+                return false;
+
+            if (!getClass().isInstance(obj))
+                return false;
+
+            ConsumerPair otherPair = (ConsumerPair) obj;
+            return this.srcMemberId.equals(otherPair.srcMemberId) && this.dstMemberId.equals(otherPair.dstMemberId);
+        }
+
+        private boolean in(Set pairs) {
+            for (ConsumerPair pair: pairs)
+                if (this.equals(pair))
+                    return true;
+            return false;
+        }
+    }
+}
\ No newline at end of file
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncClient.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncClient.java
new file mode 100644
index 0000000000000..8b35499a26e23
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncClient.java
@@ -0,0 +1,75 @@
+/*
+ * 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.ClientResponse;
+import org.apache.kafka.common.Node;
+import org.apache.kafka.common.requests.AbstractRequest;
+import org.apache.kafka.common.requests.AbstractResponse;
+import org.apache.kafka.common.utils.LogContext;
+import org.slf4j.Logger;
+
+public abstract class AsyncClient {
+
+    private final Logger log;
+    private final ConsumerNetworkClient client;
+
+    AsyncClient(ConsumerNetworkClient client, LogContext logContext) {
+        this.client = client;
+        this.log = logContext.logger(getClass());
+    }
+
+    public RequestFuture sendAsyncRequest(Node node, T1 requestData) {
+        AbstractRequest.Builder requestBuilder = prepareRequest(node, requestData);
+
+        return client.send(node, requestBuilder).compose(new RequestFutureAdapter() {
+            @Override
+            @SuppressWarnings("unchecked")
+            public void onSuccess(ClientResponse value, RequestFuture future) {
+                Resp resp;
+                try {
+                    resp = (Resp) value.responseBody();
+                } catch (ClassCastException cce) {
+                    log.error("Could not cast response body", cce);
+                    future.raise(cce);
+                    return;
+                }
+                log.trace("Received {} {} from broker {}", resp.getClass().getSimpleName(), resp, node);
+                try {
+                    future.complete(handleResponse(node, requestData, resp));
+                } catch (RuntimeException e) {
+                    if (!future.isDone()) {
+                        future.raise(e);
+                    }
+                }
+            }
+
+            @Override
+            public void onFailure(RuntimeException e, RequestFuture future1) {
+                future1.raise(e);
+            }
+        });
+    }
+
+    protected Logger logger() {
+        return log;
+    }
+
+    protected abstract AbstractRequest.Builder prepareRequest(Node node, T1 requestData);
+
+    protected abstract T2 handleResponse(Node node, T1 requestData, Resp response);
+}
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerCoordinator.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerCoordinator.java
index 79907074af27b..9784955e6e1ce 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerCoordinator.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerCoordinator.java
@@ -16,31 +16,44 @@
  */
 package org.apache.kafka.clients.consumer.internals;
 
+import org.apache.kafka.clients.GroupRebalanceConfig;
+import java.util.concurrent.TimeUnit;
 import org.apache.kafka.clients.consumer.CommitFailedException;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerPartitionAssignor;
+import org.apache.kafka.clients.consumer.ConsumerGroupMetadata;
+import org.apache.kafka.clients.consumer.ConsumerPartitionAssignor.GroupSubscription;
+import org.apache.kafka.clients.consumer.ConsumerPartitionAssignor.Subscription;
 import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
 import org.apache.kafka.clients.consumer.OffsetCommitCallback;
 import org.apache.kafka.clients.consumer.RetriableCommitFailedException;
-import org.apache.kafka.clients.consumer.internals.PartitionAssignor.Assignment;
-import org.apache.kafka.clients.consumer.internals.PartitionAssignor.Subscription;
+import org.apache.kafka.clients.consumer.ConsumerPartitionAssignor.Assignment;
+import org.apache.kafka.clients.consumer.ConsumerPartitionAssignor.RebalanceProtocol;
 import org.apache.kafka.common.Cluster;
 import org.apache.kafka.common.KafkaException;
 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.InterruptException;
 import org.apache.kafka.common.errors.RetriableException;
 import org.apache.kafka.common.errors.TimeoutException;
 import org.apache.kafka.common.errors.TopicAuthorizationException;
 import org.apache.kafka.common.errors.WakeupException;
+import org.apache.kafka.common.message.JoinGroupRequestData;
+import org.apache.kafka.common.message.JoinGroupResponseData;
+import org.apache.kafka.common.message.OffsetCommitRequestData;
+import org.apache.kafka.common.message.OffsetCommitResponseData;
 import org.apache.kafka.common.metrics.Measurable;
 import org.apache.kafka.common.metrics.MetricConfig;
 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.Max;
+import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.requests.JoinGroupRequest.ProtocolMetadata;
+import org.apache.kafka.common.record.RecordBatch;
 import org.apache.kafka.common.requests.OffsetCommitRequest;
 import org.apache.kafka.common.requests.OffsetCommitResponse;
 import org.apache.kafka.common.requests.OffsetFetchRequest;
@@ -59,17 +72,21 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Collectors;
 
 /**
  * This class manages the coordination process with the consumer coordinator.
  */
 public final class ConsumerCoordinator extends AbstractCoordinator {
+    private final GroupRebalanceConfig rebalanceConfig;
     private final Logger log;
-    private final List assignors;
+    private final List assignors;
     private final ConsumerMetadata metadata;
     private final ConsumerCoordinatorMetrics sensors;
     private final SubscriptionState subscriptions;
@@ -88,6 +105,10 @@ public final class ConsumerCoordinator extends AbstractCoordinator {
     private MetadataSnapshot metadataSnapshot;
     private MetadataSnapshot assignmentSnapshot;
     private Timer nextAutoCommitTimer;
+    private AtomicBoolean asyncCommitFenced;
+
+    private volatile long prevPollTime = Long.MIN_VALUE; //volatile for metrics
+
 
     // hold onto request&future for committed offset requests to enable async calls.
     private PendingCommittedOffsetRequest pendingCommittedOffsetRequest = null;
@@ -110,37 +131,30 @@ private boolean sameRequest(final Set currentRequest, final Gene
         }
     }
 
+    private final RebalanceProtocol protocol;
+
     /**
      * Initialize the coordination manager.
      */
-    public ConsumerCoordinator(LogContext logContext,
+    public ConsumerCoordinator(GroupRebalanceConfig rebalanceConfig,
+                               LogContext logContext,
                                ConsumerNetworkClient client,
-                               String groupId,
-                               int rebalanceTimeoutMs,
-                               int sessionTimeoutMs,
-                               Heartbeat heartbeat,
-                               List assignors,
+                               List assignors,
                                ConsumerMetadata metadata,
                                SubscriptionState subscriptions,
                                Metrics metrics,
                                String metricGrpPrefix,
                                Time time,
-                               long retryBackoffMs,
                                boolean autoCommitEnabled,
                                int autoCommitIntervalMs,
-                               ConsumerInterceptors interceptors,
-                               final boolean leaveGroupOnClose) {
-        super(logContext,
+                               ConsumerInterceptors interceptors) {
+        super(rebalanceConfig,
+              logContext,
               client,
-              groupId,
-              rebalanceTimeoutMs,
-              sessionTimeoutMs,
-              heartbeat,
               metrics,
               metricGrpPrefix,
-              time,
-              retryBackoffMs,
-              leaveGroupOnClose);
+              time);
+        this.rebalanceConfig = rebalanceConfig;
         this.log = logContext.logger(ConsumerCoordinator.class);
         this.metadata = metadata;
         this.metadataSnapshot = new MetadataSnapshot(subscriptions, metadata.fetch(), metadata.updateVersion());
@@ -153,10 +167,37 @@ public ConsumerCoordinator(LogContext logContext,
         this.sensors = new ConsumerCoordinatorMetrics(metrics, metricGrpPrefix);
         this.interceptors = interceptors;
         this.pendingAsyncCommits = new AtomicInteger();
+        this.asyncCommitFenced = new AtomicBoolean(false);
 
         if (autoCommitEnabled)
             this.nextAutoCommitTimer = time.timer(autoCommitIntervalMs);
 
+        // select the rebalance protocol such that:
+        //   1. only consider protocols that are supported by all the assignors. If there is no common protocols supported
+        //      across all the assignors, throw an exception.
+        //   2. if there are multiple protocols that are commonly supported, select the one with the highest id (i.e. the
+        //      id number indicates how advanced the protocol is).
+        // we know there are at least one assignor in the list, no need to double check for NPE
+        if (!assignors.isEmpty()) {
+            List supportedProtocols = new ArrayList<>(assignors.get(0).supportedProtocols());
+
+            for (ConsumerPartitionAssignor assignor : assignors) {
+                supportedProtocols.retainAll(assignor.supportedProtocols());
+            }
+
+            if (supportedProtocols.isEmpty()) {
+                throw new IllegalArgumentException("Specified assignors " +
+                    assignors.stream().map(ConsumerPartitionAssignor::name).collect(Collectors.toSet()) +
+                    " do not have commonly supported rebalance protocol");
+            }
+
+            Collections.sort(supportedProtocols);
+
+            protocol = supportedProtocols.get(supportedProtocols.size() - 1);
+        } else {
+            protocol = null;
+        }
+
         this.metadata.requestUpdate();
     }
 
@@ -166,16 +207,23 @@ public String protocolType() {
     }
 
     @Override
-    protected List metadata() {
+    protected JoinGroupRequestData.JoinGroupRequestProtocolCollection metadata() {
         log.debug("Joining group with current subscription: {}", subscriptions.subscription());
         this.joinedSubscription = subscriptions.subscription();
-        List metadataList = new ArrayList<>();
-        for (PartitionAssignor assignor : assignors) {
-            Subscription subscription = assignor.subscription(joinedSubscription);
+        JoinGroupRequestData.JoinGroupRequestProtocolCollection protocolSet = new JoinGroupRequestData.JoinGroupRequestProtocolCollection();
+
+        List topics = new ArrayList<>(joinedSubscription);
+        for (ConsumerPartitionAssignor assignor : assignors) {
+            Subscription subscription = new Subscription(topics,
+                                                         assignor.subscriptionUserData(joinedSubscription),
+                                                         subscriptions.assignedPartitionsList());
             ByteBuffer metadata = ConsumerProtocol.serializeSubscription(subscription);
-            metadataList.add(new ProtocolMetadata(assignor.name(), metadata));
+
+            protocolSet.add(new JoinGroupRequestData.JoinGroupRequestProtocol()
+                    .setName(assignor.name())
+                    .setMetadata(Utils.toArray(metadata)));
         }
-        return metadataList;
+        return protocolSet;
     }
 
     public void updatePatternSubscription(Cluster cluster) {
@@ -186,49 +234,95 @@ public void updatePatternSubscription(Cluster cluster) {
             metadata.requestUpdateForNewTopics();
     }
 
-    private PartitionAssignor lookupAssignor(String name) {
-        for (PartitionAssignor assignor : this.assignors) {
+    private ConsumerPartitionAssignor lookupAssignor(String name) {
+        for (ConsumerPartitionAssignor assignor : this.assignors) {
             if (assignor.name().equals(name))
                 return assignor;
         }
         return null;
     }
 
-    private void handleAssignmentMismatch(Assignment assignment) {
-        // We received an assignment that doesn't match our current subscription. If the subscription changed,
-        // we can ignore the assignment and rebalance. Otherwise we raise an error.
-        Set invalidAssignments = assignment.partitions().stream().filter(topicPartition ->
-                !joinedSubscription.contains(topicPartition.topic())).collect(Collectors.toSet());
-        if (invalidAssignments.size() > 0) {
-            throw new IllegalStateException("Consumer was assigned partitions " + invalidAssignments +
-                    " which didn't correspond to subscription request " + joinedSubscription);
+    private void maybeUpdateJoinedSubscription(Set assignedPartitions) {
+        if (subscriptions.hasPatternSubscription()) {
+            // Check if the assignment contains some topics that were not in the original
+            // subscription, if yes we will obey what leader has decided and add these topics
+            // into the subscriptions as long as they still match the subscribed pattern
+
+            Set addedTopics = new HashSet<>();
+            // this is a copy because its handed to listener below
+            for (TopicPartition tp : assignedPartitions) {
+                if (!joinedSubscription.contains(tp.topic()))
+                    addedTopics.add(tp.topic());
+            }
+
+            if (!addedTopics.isEmpty()) {
+                Set newSubscription = new HashSet<>(subscriptions.subscription());
+                Set newJoinedSubscription = new HashSet<>(joinedSubscription);
+                newSubscription.addAll(addedTopics);
+                newJoinedSubscription.addAll(addedTopics);
+
+                if (this.subscriptions.subscribeFromPattern(newSubscription))
+                    metadata.requestUpdateForNewTopics();
+                this.joinedSubscription = newJoinedSubscription;
+            }
+        }
+    }
+
+    private Exception invokePartitionsAssigned(final Set assignedPartitions) {
+        log.info("Adding newly assigned partitions: {}", Utils.join(assignedPartitions, ", "));
+
+        ConsumerRebalanceListener listener = subscriptions.rebalanceListener();
+        try {
+            final long startMs = time.milliseconds();
+            listener.onPartitionsAssigned(assignedPartitions);
+            sensors.assignCallbackSensor.record(time.milliseconds() - startMs);
+        } catch (WakeupException | InterruptException e) {
+            throw e;
+        } catch (Exception e) {
+            log.error("User provided listener {} failed on invocation of onPartitionsAssigned for partitions {}",
+                listener.getClass().getName(), assignedPartitions, e);
+            return e;
         }
 
-        requestRejoin();
+        return null;
     }
 
-    private void maybeUpdateJoinedSubscription(Set assignedPartitions) {
-        // Check if the assignment contains some topics that were not in the original
-        // subscription, if yes we will obey what leader has decided and add these topics
-        // into the subscriptions as long as they still match the subscribed pattern
+    private Exception invokePartitionsRevoked(final Set revokedPartitions) {
+        log.info("Revoke previously assigned partitions {}", Utils.join(revokedPartitions, ", "));
 
-        Set addedTopics = new HashSet<>();
-        //this is a copy because its handed to listener below
-        for (TopicPartition tp : assignedPartitions) {
-            if (!joinedSubscription.contains(tp.topic()))
-                addedTopics.add(tp.topic());
+        ConsumerRebalanceListener listener = subscriptions.rebalanceListener();
+        try {
+            final long startMs = time.milliseconds();
+            listener.onPartitionsRevoked(revokedPartitions);
+            sensors.revokeCallbackSensor.record(time.milliseconds() - startMs);
+        } catch (WakeupException | InterruptException e) {
+            throw e;
+        } catch (Exception e) {
+            log.error("User provided listener {} failed on invocation of onPartitionsRevoked for partitions {}",
+                listener.getClass().getName(), revokedPartitions, e);
+            return e;
         }
 
-        if (!addedTopics.isEmpty()) {
-            Set newSubscription = new HashSet<>(subscriptions.subscription());
-            Set newJoinedSubscription = new HashSet<>(joinedSubscription);
-            newSubscription.addAll(addedTopics);
-            newJoinedSubscription.addAll(addedTopics);
+        return null;
+    }
+
+    private Exception invokePartitionsLost(final Set lostPartitions) {
+        log.info("Lost previously assigned partitions {}", Utils.join(lostPartitions, ", "));
 
-            if (this.subscriptions.subscribeFromPattern(newSubscription))
-                metadata.requestUpdateForNewTopics();
-            this.joinedSubscription = newJoinedSubscription;
+        ConsumerRebalanceListener listener = subscriptions.rebalanceListener();
+        try {
+            final long startMs = time.milliseconds();
+            listener.onPartitionsLost(lostPartitions);
+            sensors.loseCallbackSensor.record(time.milliseconds() - startMs);
+        } catch (WakeupException | InterruptException e) {
+            throw e;
+        } catch (Exception e) {
+            log.error("User provided listener {} failed on invocation of onPartitionsLost for partitions {}",
+                listener.getClass().getName(), lostPartitions, e);
+            return e;
         }
+
+        return null;
     }
 
     @Override
@@ -236,43 +330,82 @@ protected void onJoinComplete(int generation,
                                   String memberId,
                                   String assignmentStrategy,
                                   ByteBuffer assignmentBuffer) {
+        log.debug("Executing onJoinComplete with generation {} and memberId {}", generation, memberId);
+
         // only the leader is responsible for monitoring for metadata changes (i.e. partition changes)
         if (!isLeader)
             assignmentSnapshot = null;
 
-        PartitionAssignor assignor = lookupAssignor(assignmentStrategy);
+        ConsumerPartitionAssignor assignor = lookupAssignor(assignmentStrategy);
         if (assignor == null)
             throw new IllegalStateException("Coordinator selected invalid assignment protocol: " + assignmentStrategy);
 
+        Set ownedPartitions = new HashSet<>(subscriptions.assignedPartitions());
+
         Assignment assignment = ConsumerProtocol.deserializeAssignment(assignmentBuffer);
-        if (!subscriptions.assignFromSubscribed(assignment.partitions())) {
-            handleAssignmentMismatch(assignment);
+
+        Set assignedPartitions = new HashSet<>(assignment.partitions());
+
+        if (!subscriptions.checkAssignmentMatchedSubscription(assignedPartitions)) {
+            log.warn("We received an assignment {} that doesn't match our current subscription {}; it is likely " +
+                "that the subscription has changed since we joined the group. Will try re-join the group with current subscription",
+                assignment.partitions(), subscriptions.prettyString());
+
+            requestRejoin();
+
             return;
         }
 
-        Set assignedPartitions = new HashSet<>(subscriptions.assignedPartitions());
+        final AtomicReference firstException = new AtomicReference<>(null);
+        Set addedPartitions = new HashSet<>(assignedPartitions);
+        addedPartitions.removeAll(ownedPartitions);
+
+        if (protocol == RebalanceProtocol.COOPERATIVE) {
+            Set revokedPartitions = new HashSet<>(ownedPartitions);
+            revokedPartitions.removeAll(assignedPartitions);
+
+            log.info("Updating assignment with\n" +
+                    "now assigned partitions: {}\n" +
+                    "compare with previously owned partitions: {}\n" +
+                    "newly added partitions: {}\n" +
+                    "revoked partitions: {}\n",
+                Utils.join(assignedPartitions, ", "),
+                Utils.join(ownedPartitions, ", "),
+                Utils.join(addedPartitions, ", "),
+                Utils.join(revokedPartitions, ", ")
+            );
+
+            if (!revokedPartitions.isEmpty()) {
+                // revoke partitions that were previously owned but no longer assigned;
+                // note that we should only change the assignment (or update the assignor's state)
+                // AFTER we've triggered  the revoke callback
+                firstException.compareAndSet(null, invokePartitionsRevoked(revokedPartitions));
+
+                // if revoked any partitions, need to re-join the group afterwards
+                log.debug("Need to revoke partitions {} and re-join the group", revokedPartitions);
+                requestRejoin();
+            }
+        }
 
         // The leader may have assigned partitions which match our subscription pattern, but which
         // were not explicitly requested, so we update the joined subscription here.
         maybeUpdateJoinedSubscription(assignedPartitions);
 
         // give the assignor a chance to update internal state based on the received assignment
-        assignor.onAssignment(assignment);
+        ConsumerGroupMetadata metadata = new ConsumerGroupMetadata(rebalanceConfig.groupId, generation, memberId, rebalanceConfig.groupInstanceId);
+        assignor.onAssignment(assignment, metadata);
 
         // reschedule the auto commit starting from now
         if (autoCommitEnabled)
             this.nextAutoCommitTimer.updateAndReset(autoCommitIntervalMs);
 
-        // execute the user's callback after rebalance
-        ConsumerRebalanceListener listener = subscriptions.rebalanceListener();
-        log.info("Setting newly assigned partitions: {}", Utils.join(assignedPartitions, ", "));
-        try {
-            listener.onPartitionsAssigned(assignedPartitions);
-        } catch (WakeupException | InterruptException e) {
-            throw e;
-        } catch (Exception e) {
-            log.error("User provided listener {} failed on partition assignment", listener.getClass().getName(), e);
-        }
+        subscriptions.assignFromSubscribed(assignedPartitions);
+
+        // add partitions that were not previously owned but are now assigned
+        firstException.compareAndSet(null, invokePartitionsAssigned(addedPartitions));
+
+        if (firstException.get() != null)
+            throw new KafkaException("User rebalance callback throws an error", firstException.get());
     }
 
     void maybeUpdateSubscriptionMetadata() {
@@ -297,14 +430,25 @@ void maybeUpdateSubscriptionMetadata() {
      * Returns early if the timeout expires
      *
      * @param timer Timer bounding how long this method can block
+     * @throws KafkaException if the rebalance callback throws an exception
      * @return true iff the operation succeeded
      */
     public boolean poll(Timer timer) {
+
+        long currentTime = time.milliseconds();
+        if (prevPollTime > Long.MIN_VALUE) {
+            sensors.pollInterval.record(currentTime - prevPollTime);
+        }
+        prevPollTime = currentTime;
+
         maybeUpdateSubscriptionMetadata();
 
         invokeCompletedOffsetCommitCallbacks();
-
         if (subscriptions.partitionsAutoAssigned()) {
+            if (protocol == null) {
+                throw new IllegalStateException("User configured " + ConsumerConfig.PARTITION_ASSIGNMENT_STRATEGY_CONFIG +
+                    " to empty while trying to subscribe for group protocol to auto assign partitions");
+            }
             // Always update the heartbeat last poll time so that the heartbeat thread does not leave the
             // group proactively due to application inactivity even if (say) the coordinator cannot be found.
             pollHeartbeat(timer.currentTimeMs());
@@ -385,17 +529,23 @@ private void updateGroupSubscription(Set topics) {
     @Override
     protected Map performAssignment(String leaderId,
                                                         String assignmentStrategy,
-                                                        Map allSubscriptions) {
-        PartitionAssignor assignor = lookupAssignor(assignmentStrategy);
+                                                        List allSubscriptions) {
+        ConsumerPartitionAssignor assignor = lookupAssignor(assignmentStrategy);
         if (assignor == null)
             throw new IllegalStateException("Coordinator selected invalid assignment protocol: " + assignmentStrategy);
 
         Set allSubscribedTopics = new HashSet<>();
         Map subscriptions = new HashMap<>();
-        for (Map.Entry subscriptionEntry : allSubscriptions.entrySet()) {
-            Subscription subscription = ConsumerProtocol.deserializeSubscription(subscriptionEntry.getValue());
-            subscriptions.put(subscriptionEntry.getKey(), subscription);
+
+        // collect all the owned partitions
+        Map> ownedPartitions = new HashMap<>();
+
+        for (JoinGroupResponseData.JoinGroupResponseMember memberSubscription : allSubscriptions) {
+            Subscription subscription = ConsumerProtocol.deserializeSubscription(ByteBuffer.wrap(memberSubscription.metadata()));
+            subscription.setGroupInstanceId(Optional.ofNullable(memberSubscription.groupInstanceId()));
+            subscriptions.put(memberSubscription.memberId(), subscription);
             allSubscribedTopics.addAll(subscription.topics());
+            ownedPartitions.put(memberSubscription.memberId(), subscription.ownedPartitions());
         }
 
         // the leader will begin watching for changes to any of the topics the group is interested in,
@@ -406,7 +556,11 @@ protected Map performAssignment(String leaderId,
 
         log.debug("Performing assignment using strategy {} with subscriptions {}", assignor.name(), subscriptions);
 
-        Map assignment = assignor.assign(metadata.fetch(), subscriptions);
+        Map assignments = assignor.assign(metadata.fetch(), new GroupSubscription(subscriptions)).groupAssignment();
+
+        if (protocol == RebalanceProtocol.COOPERATIVE) {
+            validateCooperativeAssignment(ownedPartitions, assignments);
+        }
 
         // user-customized assignor may have created some topics that are not in the subscription list
         // and assign their partitions to the members; in this case we would like to update the leader's
@@ -414,9 +568,9 @@ protected Map performAssignment(String leaderId,
         // when these topics gets updated from metadata refresh.
         //
         // TODO: this is a hack and not something we want to support long-term unless we push regex into the protocol
-        //       we may need to modify the PartitionAssignor API to better support this case.
+        //       we may need to modify the ConsumerPartitionAssignor API to better support this case.
         Set assignedTopics = new HashSet<>();
-        for (Assignment assigned : assignment.values()) {
+        for (Assignment assigned : assignments.values()) {
             for (TopicPartition tp : assigned.partitions())
                 assignedTopics.add(tp.topic());
         }
@@ -439,10 +593,10 @@ protected Map performAssignment(String leaderId,
 
         assignmentSnapshot = metadataSnapshot;
 
-        log.debug("Finished assignment for group: {}", assignment);
+        log.info("Finished assignment for group at generation {}: {}", generation().generationId, assignments);
 
         Map groupAssignment = new HashMap<>();
-        for (Map.Entry assignmentEntry : assignment.entrySet()) {
+        for (Map.Entry assignmentEntry : assignments.entrySet()) {
             ByteBuffer buffer = ConsumerProtocol.serializeAssignment(assignmentEntry.getValue());
             groupAssignment.put(assignmentEntry.getKey(), buffer);
         }
@@ -450,40 +604,150 @@ protected Map performAssignment(String leaderId,
         return groupAssignment;
     }
 
+    /**
+     * Used by COOPERATIVE rebalance protocol only.
+     *
+     * Validate the assignments returned by the assignor such that no owned partitions are going to
+     * be reassigned to a different consumer directly: if the assignor wants to reassign an owned partition,
+     * it must first remove it from the new assignment of the current owner so that it is not assigned to any
+     * member, and then in the next rebalance it can finally reassign those partitions not owned by anyone to consumers.
+     */
+    private void validateCooperativeAssignment(final Map> ownedPartitions,
+                                               final Map assignments) {
+        Set totalRevokedPartitions = new HashSet<>();
+        Set totalAddedPartitions = new HashSet<>();
+        for (final Map.Entry entry : assignments.entrySet()) {
+            final Assignment assignment = entry.getValue();
+            final Set addedPartitions = new HashSet<>(assignment.partitions());
+            addedPartitions.removeAll(ownedPartitions.get(entry.getKey()));
+            final Set revokedPartitions = new HashSet<>(ownedPartitions.get(entry.getKey()));
+            revokedPartitions.removeAll(assignment.partitions());
+
+            totalAddedPartitions.addAll(addedPartitions);
+            totalRevokedPartitions.addAll(revokedPartitions);
+        }
+
+        // if there are overlap between revoked partitions and added partitions, it means some partitions
+        // immediately gets re-assigned to another member while it is still claimed by some member
+        totalAddedPartitions.retainAll(totalRevokedPartitions);
+        if (!totalAddedPartitions.isEmpty()) {
+            log.error("With the COOPERATIVE protocol, owned partitions cannot be " +
+                "reassigned to other members; however the assignor has reassigned partitions {} which are still owned " +
+                "by some members", totalAddedPartitions);
+
+            throw new IllegalStateException("Assignor supporting the COOPERATIVE protocol violates its requirements");
+        }
+    }
+
     @Override
     protected void onJoinPrepare(int generation, String memberId) {
+        log.debug("Executing onJoinPrepare with generation {} and memberId {}", generation, memberId);
         // commit offsets prior to rebalance if auto-commit enabled
-        maybeAutoCommitOffsetsSync(time.timer(rebalanceTimeoutMs));
+        maybeAutoCommitOffsetsSync(time.timer(rebalanceConfig.rebalanceTimeoutMs));
+
+        // the generation / member-id can possibly be reset by the heartbeat thread
+        // upon getting errors or heartbeat timeouts; in this case whatever is previously
+        // owned partitions would be lost, we should trigger the callback and cleanup the assignment;
+        // otherwise we can proceed normally and revoke the partitions depending on the protocol,
+        // and in that case we should only change the assignment AFTER the revoke callback is triggered
+        // so that users can still access the previously owned partitions to commit offsets etc.
+        Exception exception = null;
+        final Set revokedPartitions;
+        if (generation == Generation.NO_GENERATION.generationId &&
+            memberId.equals(Generation.NO_GENERATION.memberId)) {
+            revokedPartitions = new HashSet<>(subscriptions.assignedPartitions());
+
+            if (!revokedPartitions.isEmpty()) {
+                log.info("Giving away all assigned partitions as lost since generation has been reset," +
+                    "indicating that consumer is no longer part of the group");
+                exception = invokePartitionsLost(revokedPartitions);
+
+                subscriptions.assignFromSubscribed(Collections.emptySet());
+            }
+        } else {
+            switch (protocol) {
+                case EAGER:
+                    // revoke all partitions
+                    revokedPartitions = new HashSet<>(subscriptions.assignedPartitions());
+                    exception = invokePartitionsRevoked(revokedPartitions);
 
-        // execute the user's callback before rebalance
-        ConsumerRebalanceListener listener = subscriptions.rebalanceListener();
-        // copy since about to be handed to user code
-        Set revoked = new HashSet<>(subscriptions.assignedPartitions());
-        log.info("Revoking previously assigned partitions {}", revoked);
-        try {
-            listener.onPartitionsRevoked(revoked);
-        } catch (WakeupException | InterruptException e) {
-            throw e;
-        } catch (Exception e) {
-            log.error("User provided listener {} failed on partition revocation", listener.getClass().getName(), e);
+                    subscriptions.assignFromSubscribed(Collections.emptySet());
+
+                    break;
+
+                case COOPERATIVE:
+                    // only revoke those partitions that are not in the subscription any more.
+                    Set ownedPartitions = new HashSet<>(subscriptions.assignedPartitions());
+                    revokedPartitions = ownedPartitions.stream()
+                        .filter(tp -> !subscriptions.subscription().contains(tp.topic()))
+                        .collect(Collectors.toSet());
+
+                    if (!revokedPartitions.isEmpty()) {
+                        exception = invokePartitionsRevoked(revokedPartitions);
+
+                        ownedPartitions.removeAll(revokedPartitions);
+                        subscriptions.assignFromSubscribed(ownedPartitions);
+                    }
+
+                    break;
+            }
         }
 
         isLeader = false;
         subscriptions.resetGroupSubscription();
+
+        if (exception != null) {
+            throw new KafkaException("User rebalance callback throws an error", exception);
+        }
+    }
+
+    @Override
+    public void onLeavePrepare() {
+        // Save the current Generation and use that to get the memberId, as the hb thread can change it at any time
+        final Generation currentGeneration = generation();
+        final String memberId = currentGeneration.memberId;
+
+        log.debug("Executing onLeavePrepare with generation {} and memberId {}", currentGeneration, memberId);
+
+        // we should reset assignment and trigger the callback before leaving group
+        Set droppedPartitions = new HashSet<>(subscriptions.assignedPartitions());
+
+        if (subscriptions.partitionsAutoAssigned() && !droppedPartitions.isEmpty()) {
+            final Exception e;
+            if (generation() != Generation.NO_GENERATION) {
+                e = invokePartitionsRevoked(droppedPartitions);
+            } else {
+                e = invokePartitionsLost(droppedPartitions);
+            }
+
+            subscriptions.assignFromSubscribed(Collections.emptySet());
+
+            if (e != null) {
+                throw new KafkaException("User rebalance callback throws an error", e);
+            }
+        }
     }
 
+    /**
+     * @throws KafkaException if the callback throws exception
+     */
     @Override
     public boolean rejoinNeededOrPending() {
         if (!subscriptions.partitionsAutoAssigned())
             return false;
 
-        // we need to rejoin if we performed the assignment and metadata has changed
-        if (assignmentSnapshot != null && !assignmentSnapshot.matches(metadataSnapshot))
+        // we need to rejoin if we performed the assignment and metadata has changed;
+        // also for those owned-but-no-longer-existed partitions we should drop them as lost
+        if (assignmentSnapshot != null && !assignmentSnapshot.matches(metadataSnapshot)) {
+            requestRejoin();
             return true;
+        }
 
         // we need to join if our subscription has changed since the last join
-        if (joinedSubscription != null && !joinedSubscription.equals(subscriptions.subscription()))
+        if (joinedSubscription != null && !joinedSubscription.equals(subscriptions.subscription())) {
+            requestRejoin();
             return true;
+        }
 
         return super.rejoinNeededOrPending();
     }
@@ -502,10 +766,17 @@ public boolean refreshCommittedOffsetsIfNeeded(Timer timer) {
 
         for (final Map.Entry entry : offsets.entrySet()) {
             final TopicPartition tp = entry.getKey();
-            final long offset = entry.getValue().offset();
-            log.info("Setting offset for partition {} to the committed offset {}", tp, offset);
-            entry.getValue().leaderEpoch().ifPresent(epoch -> this.metadata.updateLastSeenEpochIfNewer(entry.getKey(), epoch));
-            this.subscriptions.seek(tp, offset);
+            final OffsetAndMetadata offsetAndMetadata = entry.getValue();
+            if (offsetAndMetadata != null) {
+                final ConsumerMetadata.LeaderAndEpoch leaderAndEpoch = metadata.leaderAndEpoch(tp);
+                final SubscriptionState.FetchPosition position = new SubscriptionState.FetchPosition(
+                    offsetAndMetadata.offset(), offsetAndMetadata.leaderEpoch(),
+                    leaderAndEpoch);
+
+                log.info("Setting offset for partition {} to the committed offset {}", tp, position);
+                entry.getValue().leaderEpoch().ifPresent(epoch -> this.metadata.updateLastSeenEpochIfNewer(entry.getKey(), epoch));
+                this.subscriptions.seekUnvalidated(tp, position);
+            }
         }
         return true;
     }
@@ -520,8 +791,9 @@ public Map fetchCommittedOffsets(final Set fetchCommittedOffsets(final Set fetchCommittedOffsets(final Set fetchCommittedOffsets(final Set offsets,
             if (future.failed() && !future.isRetriable())
                 throw future.exception();
 
-            timer.sleep(retryBackoffMs);
+            timer.sleep(rebalanceConfig.retryBackoffMs);
         } while (timer.notExpired());
 
         return false;
@@ -698,20 +984,17 @@ private void doAutoCommitOffsetsAsync() {
         Map allConsumedOffsets = subscriptions.allConsumed();
         log.debug("Sending asynchronous auto-commit of offsets {}", allConsumedOffsets);
 
-        commitOffsetsAsync(allConsumedOffsets, new OffsetCommitCallback() {
-            @Override
-            public void onComplete(Map offsets, Exception exception) {
-                if (exception != null) {
-                    if (exception instanceof RetriableException) {
-                        log.debug("Asynchronous auto-commit of offsets {} failed due to retriable error: {}", offsets,
-                                exception);
-                        nextAutoCommitTimer.updateAndReset(retryBackoffMs);
-                    } else {
-                        log.warn("Asynchronous auto-commit of offsets {} failed: {}", offsets, exception.getMessage());
-                    }
+        commitOffsetsAsync(allConsumedOffsets, (offsets, exception) -> {
+            if (exception != null) {
+                if (exception instanceof RetriableException) {
+                    log.debug("Asynchronous auto-commit of offsets {} failed due to retriable error: {}", offsets,
+                        exception);
+                    nextAutoCommitTimer.updateAndReset(rebalanceConfig.retryBackoffMs);
                 } else {
-                    log.debug("Completed asynchronous auto-commit of offsets {}", offsets);
+                    log.warn("Asynchronous auto-commit of offsets {} failed: {}", offsets, exception.getMessage());
                 }
+            } else {
+                log.debug("Completed asynchronous auto-commit of offsets {}", offsets);
             }
         });
     }
@@ -759,30 +1042,49 @@ private RequestFuture sendOffsetCommitRequest(final Map offsetData = new HashMap<>(offsets.size());
+        Map requestTopicDataMap = new HashMap<>();
         for (Map.Entry entry : offsets.entrySet()) {
+            TopicPartition topicPartition = entry.getKey();
             OffsetAndMetadata offsetAndMetadata = entry.getValue();
             if (offsetAndMetadata.offset() < 0) {
                 return RequestFuture.failure(new IllegalArgumentException("Invalid offset: " + offsetAndMetadata.offset()));
             }
-            offsetData.put(entry.getKey(), new OffsetCommitRequest.PartitionData(offsetAndMetadata.offset(),
-                    offsetAndMetadata.leaderEpoch(), offsetAndMetadata.metadata()));
+
+            OffsetCommitRequestData.OffsetCommitRequestTopic topic = requestTopicDataMap
+                    .getOrDefault(topicPartition.topic(),
+                            new OffsetCommitRequestData.OffsetCommitRequestTopic()
+                                    .setName(topicPartition.topic())
+                    );
+
+            topic.partitions().add(new OffsetCommitRequestData.OffsetCommitRequestPartition()
+                    .setPartitionIndex(topicPartition.partition())
+                    .setCommittedOffset(offsetAndMetadata.offset())
+                    .setCommittedLeaderEpoch(offsetAndMetadata.leaderEpoch().orElse(RecordBatch.NO_PARTITION_LEADER_EPOCH))
+                    .setCommittedMetadata(offsetAndMetadata.metadata())
+            );
+            requestTopicDataMap.put(topicPartition.topic(), topic);
         }
 
         final Generation generation;
-        if (subscriptions.partitionsAutoAssigned())
-            generation = generation();
-        else
+        if (subscriptions.partitionsAutoAssigned()) {
+            generation = generationIfStable();
+            // if the generation is null, we are not part of an active group (and we expect to be).
+            // the only thing we can do is fail the commit and let the user rejoin the group in poll()
+            if (generation == null) {
+                log.info("Failing OffsetCommit request since the consumer is not part of an active group");
+                return RequestFuture.failure(new CommitFailedException());
+            }
+        } else
             generation = Generation.NO_GENERATION;
 
-        // if the generation is null, we are not part of an active group (and we expect to be).
-        // the only thing we can do is fail the commit and let the user rejoin the group in poll()
-        if (generation == null)
-            return RequestFuture.failure(new CommitFailedException());
-
-        OffsetCommitRequest.Builder builder = new OffsetCommitRequest.Builder(this.groupId, offsetData).
-                setGenerationId(generation.generationId).
-                setMemberId(generation.memberId);
+        OffsetCommitRequest.Builder builder = new OffsetCommitRequest.Builder(
+                new OffsetCommitRequestData()
+                        .setGroupId(this.rebalanceConfig.groupId)
+                        .setGenerationId(generation.generationId)
+                        .setMemberId(generation.memberId)
+                        .setGroupInstanceId(rebalanceConfig.groupInstanceId.orElse(null))
+                        .setTopics(new ArrayList<>(requestTopicDataMap.values()))
+        );
 
         log.trace("Sending OffsetCommit request with {} to coordinator {}", offsets, coordinator);
 
@@ -800,51 +1102,72 @@ private OffsetCommitResponseHandler(Map offse
 
         @Override
         public void handle(OffsetCommitResponse commitResponse, RequestFuture future) {
-            sensors.commitLatency.record(response.requestLatencyMs());
+            sensors.commitSensor.record(response.requestLatencyMs());
             Set unauthorizedTopics = new HashSet<>();
 
-            for (Map.Entry entry : commitResponse.responseData().entrySet()) {
-                TopicPartition tp = entry.getKey();
-                OffsetAndMetadata offsetAndMetadata = this.offsets.get(tp);
-                long offset = offsetAndMetadata.offset();
+            for (OffsetCommitResponseData.OffsetCommitResponseTopic topic : commitResponse.data().topics()) {
+                for (OffsetCommitResponseData.OffsetCommitResponsePartition partition : topic.partitions()) {
+                    TopicPartition tp = new TopicPartition(topic.name(), partition.partitionIndex());
+                    OffsetAndMetadata offsetAndMetadata = this.offsets.get(tp);
 
-                Errors error = entry.getValue();
-                if (error == Errors.NONE) {
-                    log.debug("Committed offset {} for partition {}", offset, tp);
-                } else {
-                    log.error("Offset commit failed on partition {} at offset {}: {}", tp, offset, error.message());
+                    long offset = offsetAndMetadata.offset();
 
-                    if (error == Errors.GROUP_AUTHORIZATION_FAILED) {
-                        future.raise(new GroupAuthorizationException(groupId));
-                        return;
-                    } else if (error == Errors.TOPIC_AUTHORIZATION_FAILED) {
-                        unauthorizedTopics.add(tp.topic());
-                    } else if (error == Errors.OFFSET_METADATA_TOO_LARGE
-                            || error == Errors.INVALID_COMMIT_OFFSET_SIZE) {
-                        // raise the error to the user
-                        future.raise(error);
-                        return;
-                    } else if (error == Errors.COORDINATOR_LOAD_IN_PROGRESS
-                            || error == Errors.UNKNOWN_TOPIC_OR_PARTITION) {
-                        // just retry
-                        future.raise(error);
-                        return;
-                    } else if (error == Errors.COORDINATOR_NOT_AVAILABLE
-                            || error == Errors.NOT_COORDINATOR
-                            || error == Errors.REQUEST_TIMED_OUT) {
-                        markCoordinatorUnknown();
-                        future.raise(error);
-                        return;
-                    } else if (error == Errors.UNKNOWN_MEMBER_ID
-                            || error == Errors.ILLEGAL_GENERATION
-                            || error == Errors.REBALANCE_IN_PROGRESS) {
-                        // need to re-join group
-                        resetGeneration();
-                        future.raise(new CommitFailedException());
-                        return;
+                    Errors error = Errors.forCode(partition.errorCode());
+                    if (error == Errors.NONE) {
+                        log.debug("Committed offset {} for partition {}", offset, tp);
                     } else {
-                        future.raise(new KafkaException("Unexpected error in commit: " + error.message()));
-                        return;
+                        if (error.exception() instanceof RetriableException) {
+                            log.warn("Offset commit failed on partition {} at offset {}: {}", tp, offset, error.message());
+                        } else {
+                            log.error("Offset commit failed on partition {} at offset {}: {}", tp, offset, error.message());
+                        }
+
+                        if (error == Errors.GROUP_AUTHORIZATION_FAILED) {
+                            future.raise(GroupAuthorizationException.forGroupId(rebalanceConfig.groupId));
+                            return;
+                        } else if (error == Errors.TOPIC_AUTHORIZATION_FAILED) {
+                            unauthorizedTopics.add(tp.topic());
+                        } else if (error == Errors.OFFSET_METADATA_TOO_LARGE
+                                || error == Errors.INVALID_COMMIT_OFFSET_SIZE) {
+                            // raise the error to the user
+                            future.raise(error);
+                            return;
+                        } else if (error == Errors.COORDINATOR_LOAD_IN_PROGRESS
+                                || error == Errors.UNKNOWN_TOPIC_OR_PARTITION) {
+                            // just retry
+                            future.raise(error);
+                            return;
+                        } else if (error == Errors.COORDINATOR_NOT_AVAILABLE
+                                || error == Errors.NOT_COORDINATOR
+                                || error == Errors.REQUEST_TIMED_OUT) {
+                            markCoordinatorUnknown();
+                            future.raise(error);
+                            return;
+                        } else if (error == Errors.FENCED_INSTANCE_ID) {
+                            log.error("Received fatal exception: group.instance.id gets fenced");
+                            future.raise(error);
+                            return;
+                        } else if (error == Errors.REBALANCE_IN_PROGRESS) {
+                            /* Consumer never tries to commit offset in between join-group and sync-group,
+                             * and hence on broker-side it is not expected to see a commit offset request
+                             * during CompletingRebalance phase; if it ever happens then broker would return
+                             * this error. In this case we should just treat as a fatal CommitFailed exception.
+                             * However, we do not need to reset generations and just request re-join, such that
+                             * if the caller decides to proceed and poll, it would still try to proceed and re-join normally.
+                             */
+                            requestRejoin();
+                            future.raise(new CommitFailedException());
+                            return;
+                        } else if (error == Errors.UNKNOWN_MEMBER_ID
+                                || error == Errors.ILLEGAL_GENERATION) {
+                            // need to reset generation and re-join group
+                            resetGenerationOnResponseError(ApiKeys.OFFSET_COMMIT, error);
+                            future.raise(new CommitFailedException());
+                            return;
+                        } else {
+                            future.raise(new KafkaException("Unexpected error in commit: " + error.message()));
+                            return;
+                        }
                     }
                 }
             }
@@ -872,7 +1195,7 @@ private RequestFuture> sendOffsetFetchReq
 
         log.debug("Fetching committed offsets for partitions: {}", partitions);
         // construct the request
-        OffsetFetchRequest.Builder requestBuilder = new OffsetFetchRequest.Builder(this.groupId,
+        OffsetFetchRequest.Builder requestBuilder = new OffsetFetchRequest.Builder(this.rebalanceConfig.groupId,
                 new ArrayList<>(partitions));
 
         // send the request with a callback
@@ -895,13 +1218,14 @@ public void handle(OffsetFetchResponse response, RequestFuture unauthorizedTopics = null;
             Map offsets = new HashMap<>(response.responseData().size());
             for (Map.Entry entry : response.responseData().entrySet()) {
                 TopicPartition tp = entry.getKey();
@@ -912,49 +1236,107 @@ public void handle(OffsetFetchResponse response, RequestFuture();
+                        }
+                        unauthorizedTopics.add(tp.topic());
                     } else {
                         future.raise(new KafkaException("Unexpected error in fetch offset response for partition " +
                             tp + ": " + error.message()));
+                        return;
                     }
-                    return;
                 } else if (data.offset >= 0) {
-                    // record the position with the offset (-1 indicates no committed offset to fetch)
+                    // record the position with the offset (-1 indicates no committed offset to fetch);
+                    // if there's no committed offset, record as null
                     offsets.put(tp, new OffsetAndMetadata(data.offset, data.leaderEpoch, data.metadata));
                 } else {
                     log.info("Found no committed offset for partition {}", tp);
+                    offsets.put(tp, null);
                 }
             }
 
-            future.complete(offsets);
+            if (unauthorizedTopics != null) {
+                future.raise(new TopicAuthorizationException(unauthorizedTopics));
+            } else {
+                future.complete(offsets);
+            }
         }
     }
 
     private class ConsumerCoordinatorMetrics {
         private final String metricGrpName;
-        private final Sensor commitLatency;
+        private final Sensor commitSensor;
+        private final Sensor revokeCallbackSensor;
+        private final Sensor assignCallbackSensor;
+        private final Sensor loseCallbackSensor;
+        private final Sensor pollInterval;
 
         private ConsumerCoordinatorMetrics(Metrics metrics, String metricGrpPrefix) {
             this.metricGrpName = metricGrpPrefix + "-coordinator-metrics";
 
-            this.commitLatency = metrics.sensor("commit-latency");
-            this.commitLatency.add(metrics.metricName("commit-latency-avg",
+            this.commitSensor = metrics.sensor("commit-latency");
+            this.commitSensor.add(metrics.metricName("commit-latency-avg",
                 this.metricGrpName,
                 "The average time taken for a commit request"), new Avg());
-            this.commitLatency.add(metrics.metricName("commit-latency-max",
+            this.commitSensor.add(metrics.metricName("commit-latency-max",
                 this.metricGrpName,
                 "The max time taken for a commit request"), new Max());
-            this.commitLatency.add(createMeter(metrics, metricGrpName, "commit", "commit calls"));
+            this.commitSensor.add(createMeter(metrics, metricGrpName, "commit", "commit calls"));
+
+            this.revokeCallbackSensor = metrics.sensor("partition-revoked-latency");
+            this.revokeCallbackSensor.add(metrics.metricName("partition-revoked-latency-avg",
+                this.metricGrpName,
+                "The average time taken for a partition-revoked rebalance listener callback"), new Avg());
+            this.revokeCallbackSensor.add(metrics.metricName("partition-revoked-latency-max",
+                this.metricGrpName,
+                "The max time taken for a partition-revoked rebalance listener callback"), new Max());
 
-            Measurable numParts =
+            this.assignCallbackSensor = metrics.sensor("partition-assigned-latency");
+            this.assignCallbackSensor.add(metrics.metricName("partition-assigned-latency-avg",
+                this.metricGrpName,
+                "The average time taken for a partition-assigned rebalance listener callback"), new Avg());
+            this.assignCallbackSensor.add(metrics.metricName("partition-assigned-latency-max",
+                this.metricGrpName,
+                "The max time taken for a partition-assigned rebalance listener callback"), new Max());
+
+            this.loseCallbackSensor = metrics.sensor("partition-lost-latency");
+            this.loseCallbackSensor.add(metrics.metricName("partition-lost-latency-avg",
+                this.metricGrpName,
+                "The average time taken for a partition-lost rebalance listener callback"), new Avg());
+            this.loseCallbackSensor.add(metrics.metricName("partition-lost-latency-max",
+                this.metricGrpName,
+                "The max time taken for a partition-lost rebalance listener callback"), new Max());
+
+            Measurable numParts = (config, now) -> subscriptions.numAssignedPartitions();
+            metrics.addMetric(metrics.metricName("assigned-partitions",
+                this.metricGrpName,
+                "The number of partitions currently assigned to this consumer"), numParts);
+
+            //HOTFIX - extra liveliness-related metrics
+
+            this.pollInterval = metrics.sensor("poll-interval");
+            this.pollInterval.add(metrics.metricName("poll-interval-avg",
+                this.metricGrpName,
+                "The average time between subsequent poll calls"), new Avg());
+            this.pollInterval.add(metrics.metricName("poll-interval-max",
+                this.metricGrpName,
+                "The max time between subsequent poll calls"), new Max());
+            this.pollInterval.add(createMeter(metrics, metricGrpName, "poll", "poll calls"));
+
+            Measurable lastHeartbeat =
                 new Measurable() {
                     public double measure(MetricConfig config, long now) {
-                        // Get the number of assigned partitions in a thread safe manner
-                        return subscriptions.numAssignedPartitions();
+                        return TimeUnit.SECONDS.convert(now - prevPollTime, TimeUnit.MILLISECONDS);
                     }
                 };
-            metrics.addMetric(metrics.metricName("assigned-partitions",
+            metrics.addMetric(metrics.metricName("last-poll-seconds-ago",
                 this.metricGrpName,
-                "The number of partitions currently assigned to this consumer"), numParts);
+                "The number of seconds since the last poll call"),
+                lastHeartbeat);
+
+            //end HOTFIX
         }
     }
 
@@ -964,8 +1346,11 @@ private static class MetadataSnapshot {
 
         private MetadataSnapshot(SubscriptionState subscription, Cluster cluster, int version) {
             Map partitionsPerTopic = new HashMap<>();
-            for (String topic : subscription.groupSubscription())
-                partitionsPerTopic.put(topic, cluster.partitionCountForTopic(topic));
+            for (String topic : subscription.metadataTopics()) {
+                Integer numPartitions = cluster.partitionCountForTopic(topic);
+                if (numPartitions != null)
+                    partitionsPerTopic.put(topic, numPartitions);
+            }
             this.partitionsPerTopic = partitionsPerTopic;
             this.version = version;
         }
@@ -973,6 +1358,7 @@ private MetadataSnapshot(SubscriptionState subscription, Cluster cluster, int ve
         boolean matches(MetadataSnapshot other) {
             return version == other.version || partitionsPerTopic.equals(other.partitionsPerTopic);
         }
+
     }
 
     private static class OffsetCommitCompletion {
@@ -992,4 +1378,8 @@ public void invoke() {
         }
     }
 
+    /* test-only classes below */
+    RebalanceProtocol getProtocol() {
+        return protocol;
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerMetadata.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerMetadata.java
index c87849dfcab2b..ef7d92417471b 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerMetadata.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerMetadata.java
@@ -28,29 +28,36 @@
 
 public class ConsumerMetadata extends Metadata {
     private final boolean includeInternalTopics;
+    private final boolean allowAutoTopicCreation;
     private final SubscriptionState subscription;
     private final Set transientTopics;
 
     public ConsumerMetadata(long refreshBackoffMs,
                             long metadataExpireMs,
                             boolean includeInternalTopics,
+                            boolean allowAutoTopicCreation,
                             SubscriptionState subscription,
                             LogContext logContext,
                             ClusterResourceListeners clusterResourceListeners) {
         super(refreshBackoffMs, metadataExpireMs, logContext, clusterResourceListeners);
         this.includeInternalTopics = includeInternalTopics;
+        this.allowAutoTopicCreation = allowAutoTopicCreation;
         this.subscription = subscription;
         this.transientTopics = new HashSet<>();
     }
 
+    public boolean allowAutoTopicCreation() {
+        return allowAutoTopicCreation;
+    }
+
     @Override
     public synchronized MetadataRequest.Builder newMetadataRequestBuilder() {
         if (subscription.hasPatternSubscription())
             return MetadataRequest.Builder.allTopics();
         List topics = new ArrayList<>();
-        topics.addAll(subscription.groupSubscription());
+        topics.addAll(subscription.metadataTopics());
         topics.addAll(transientTopics);
-        return new MetadataRequest.Builder(topics, true);
+        return new MetadataRequest.Builder(topics, allowAutoTopicCreation);
     }
 
     synchronized void addTransientTopics(Set topics) {
@@ -65,7 +72,7 @@ synchronized void clearTransientTopics() {
 
     @Override
     protected synchronized boolean retainTopic(String topic, boolean isInternal, long nowMs) {
-        if (transientTopics.contains(topic) || subscription.isGroupSubscribed(topic))
+        if (transientTopics.contains(topic) || subscription.needsMetadata(topic))
             return true;
 
         if (isInternal && !includeInternalTopics)
@@ -73,5 +80,4 @@ protected synchronized boolean retainTopic(String topic, boolean isInternal, lon
 
         return subscription.matchesSubscribedPattern(topic);
     }
-
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerNetworkClient.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerNetworkClient.java
index 753fdb0e7b1da..ffaa4835686d6 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerNetworkClient.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerNetworkClient.java
@@ -293,7 +293,7 @@ public void poll(Timer timer, PollCondition pollCondition, boolean disableWakeup
         // called without the lock to avoid deadlock potential if handlers need to acquire locks
         firePendingCompletedRequests();
 
-        metadata.maybeThrowException();
+        metadata.maybeThrowAnyException();
     }
 
     /**
@@ -303,6 +303,27 @@ public void pollNoWakeup() {
         poll(time.timer(0), null, true);
     }
 
+    /**
+     * Poll for network IO in best-effort only trying to transmit the ready-to-send request
+     * Do not check any pending requests or metadata errors so that no exception should ever
+     * be thrown, also no wakeups be triggered and no interrupted exception either.
+     */
+    public void transmitSends() {
+        Timer timer = time.timer(0);
+
+        // do not try to handle any disconnects, prev request failures, metadata exception etc;
+        // just try once and return immediately
+        lock.lock();
+        try {
+            // send all the requests we can send now
+            trySend(timer.currentTimeMs());
+
+            client.poll(0, timer.currentTimeMs());
+        } finally {
+            lock.unlock();
+        }
+    }
+
     /**
      * Block until all pending requests from the given node have finished.
      * @param node The node to await requests from
@@ -459,7 +480,8 @@ private void failUnsentRequests(Node node, RuntimeException e) {
         }
     }
 
-    private long trySend(long now) {
+    // Visible for testing
+    long trySend(long now) {
         long pollDelayMs = maxPollTimeoutMs;
 
         // send any requests that can be sent now
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerProtocol.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerProtocol.java
index 8a4aef8b33487..8b29835d62968 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerProtocol.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerProtocol.java
@@ -16,6 +16,9 @@
  */
 package org.apache.kafka.clients.consumer.internals;
 
+import java.util.Collections;
+import org.apache.kafka.clients.consumer.ConsumerPartitionAssignor.Assignment;
+import org.apache.kafka.clients.consumer.ConsumerPartitionAssignor.Subscription;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.protocol.types.ArrayOf;
 import org.apache.kafka.common.protocol.types.Field;
@@ -32,21 +35,42 @@
 
 /**
  * ConsumerProtocol contains the schemas for consumer subscriptions and assignments for use with
- * Kafka's generalized group management protocol. Below is the version 0 format:
+ * Kafka's generalized group management protocol. Below is the version 1 format:
  *
  * 
        * Subscription => Version Topics
  *   Version    => Int16
  *   Topics     => [String]
  *   UserData   => Bytes
+ *   OwnedPartitions    => [Topic Partitions]
+ *     Topic            => String
+ *     Partitions       => [int32]
  *
  * Assignment => Version TopicPartitions
- *   Version         => int16
- *   TopicPartitions => [Topic Partitions]
- *     Topic         => String
- *     Partitions    => [int32]
+ *   Version            => int16
+ *   AssignedPartitions => [Topic Partitions]
+ *     Topic            => String
+ *     Partitions       => [int32]
+ *   UserData           => Bytes
  * 
      
  *
+ * Version 0 format:
+ *
+ * 
      + * Subscription => Version Topics
+ *   Version    => Int16
+ *   Topics     => [String]
+ *   UserData   => Bytes
+ *
+ * Assignment => Version TopicPartitions
+ *   Version            => int16
+ *   AssignedPartitions => [Topic Partitions]
+ *     Topic            => String
+ *     Partitions       => [int32]
+ *   UserData           => Bytes
+ * 
      
+ *
+ *
  * The current implementation assumes that future versions will not break compatibility. When
  * it encounters a newer version, it parses it using the current format. This basically means
  * that new versions cannot remove or reorder any of the existing fields.
@@ -59,29 +83,53 @@ public class ConsumerProtocol {
     public static final String TOPICS_KEY_NAME = "topics";
     public static final String TOPIC_KEY_NAME = "topic";
     public static final String PARTITIONS_KEY_NAME = "partitions";
+    public static final String OWNED_PARTITIONS_KEY_NAME = "owned_partitions";
     public static final String TOPIC_PARTITIONS_KEY_NAME = "topic_partitions";
     public static final String USER_DATA_KEY_NAME = "user_data";
 
     public static final short CONSUMER_PROTOCOL_V0 = 0;
+    public static final short CONSUMER_PROTOCOL_V1 = 1;
+
+    public static final short CONSUMER_PROTOCOL_LATEST_VERSION = CONSUMER_PROTOCOL_V1;
+
     public static final Schema CONSUMER_PROTOCOL_HEADER_SCHEMA = new Schema(
             new Field(VERSION_KEY_NAME, Type.INT16));
     private static final Struct CONSUMER_PROTOCOL_HEADER_V0 = new Struct(CONSUMER_PROTOCOL_HEADER_SCHEMA)
             .set(VERSION_KEY_NAME, CONSUMER_PROTOCOL_V0);
+    private static final Struct CONSUMER_PROTOCOL_HEADER_V1 = new Struct(CONSUMER_PROTOCOL_HEADER_SCHEMA)
+            .set(VERSION_KEY_NAME, CONSUMER_PROTOCOL_V1);
+
+    public static final Schema TOPIC_ASSIGNMENT_V0 = new Schema(
+        new Field(TOPIC_KEY_NAME, Type.STRING),
+        new Field(PARTITIONS_KEY_NAME, new ArrayOf(Type.INT32)));
 
     public static final Schema SUBSCRIPTION_V0 = new Schema(
             new Field(TOPICS_KEY_NAME, new ArrayOf(Type.STRING)),
             new Field(USER_DATA_KEY_NAME, Type.NULLABLE_BYTES));
-    public static final Schema TOPIC_ASSIGNMENT_V0 = new Schema(
-            new Field(TOPIC_KEY_NAME, Type.STRING),
-            new Field(PARTITIONS_KEY_NAME, new ArrayOf(Type.INT32)));
+
+    public static final Schema SUBSCRIPTION_V1 = new Schema(
+        new Field(TOPICS_KEY_NAME, new ArrayOf(Type.STRING)),
+        new Field(USER_DATA_KEY_NAME, Type.NULLABLE_BYTES),
+        new Field(OWNED_PARTITIONS_KEY_NAME, new ArrayOf(TOPIC_ASSIGNMENT_V0)));
+
     public static final Schema ASSIGNMENT_V0 = new Schema(
             new Field(TOPIC_PARTITIONS_KEY_NAME, new ArrayOf(TOPIC_ASSIGNMENT_V0)),
             new Field(USER_DATA_KEY_NAME, Type.NULLABLE_BYTES));
 
-    public static ByteBuffer serializeSubscription(PartitionAssignor.Subscription subscription) {
+    public static final Schema ASSIGNMENT_V1 = new Schema(
+        new Field(TOPIC_PARTITIONS_KEY_NAME, new ArrayOf(TOPIC_ASSIGNMENT_V0)),
+        new Field(USER_DATA_KEY_NAME, Type.NULLABLE_BYTES));
+
+    public static Short deserializeVersion(ByteBuffer buffer) {
+        Struct header = CONSUMER_PROTOCOL_HEADER_SCHEMA.read(buffer);
+        return header.getShort(VERSION_KEY_NAME);
+    }
+
+    public static ByteBuffer serializeSubscriptionV0(Subscription subscription) {
         Struct struct = new Struct(SUBSCRIPTION_V0);
         struct.set(USER_DATA_KEY_NAME, subscription.userData());
         struct.set(TOPICS_KEY_NAME, subscription.topics().toArray());
+
         ByteBuffer buffer = ByteBuffer.allocate(CONSUMER_PROTOCOL_HEADER_V0.sizeOf() + SUBSCRIPTION_V0.sizeOf(struct));
         CONSUMER_PROTOCOL_HEADER_V0.writeTo(buffer);
         SUBSCRIPTION_V0.write(buffer, struct);
@@ -89,37 +137,94 @@ public static ByteBuffer serializeSubscription(PartitionAssignor.Subscription su
         return buffer;
     }
 
-    public static PartitionAssignor.Subscription deserializeSubscription(ByteBuffer buffer) {
-        Struct header = CONSUMER_PROTOCOL_HEADER_SCHEMA.read(buffer);
-        Short version = header.getShort(VERSION_KEY_NAME);
-        checkVersionCompatibility(version);
+    public static ByteBuffer serializeSubscriptionV1(Subscription subscription) {
+        Struct struct = new Struct(SUBSCRIPTION_V1);
+        struct.set(USER_DATA_KEY_NAME, subscription.userData());
+        struct.set(TOPICS_KEY_NAME, subscription.topics().toArray());
+        List topicAssignments = new ArrayList<>();
+        Map> partitionsByTopic = CollectionUtils.groupPartitionsByTopic(subscription.ownedPartitions());
+        for (Map.Entry> topicEntry : partitionsByTopic.entrySet()) {
+            Struct topicAssignment = new Struct(TOPIC_ASSIGNMENT_V0);
+            topicAssignment.set(TOPIC_KEY_NAME, topicEntry.getKey());
+            topicAssignment.set(PARTITIONS_KEY_NAME, topicEntry.getValue().toArray());
+            topicAssignments.add(topicAssignment);
+        }
+        struct.set(OWNED_PARTITIONS_KEY_NAME, topicAssignments.toArray());
+
+        ByteBuffer buffer = ByteBuffer.allocate(CONSUMER_PROTOCOL_HEADER_V1.sizeOf() + SUBSCRIPTION_V1.sizeOf(struct));
+        CONSUMER_PROTOCOL_HEADER_V1.writeTo(buffer);
+        SUBSCRIPTION_V1.write(buffer, struct);
+        buffer.flip();
+        return buffer;
+    }
+
+    public static ByteBuffer serializeSubscription(Subscription subscription) {
+        return serializeSubscription(subscription, CONSUMER_PROTOCOL_LATEST_VERSION);
+    }
+
+    public static ByteBuffer serializeSubscription(Subscription subscription, short version) {
+        switch (version) {
+            case CONSUMER_PROTOCOL_V0:
+                return serializeSubscriptionV0(subscription);
+
+            case CONSUMER_PROTOCOL_V1:
+                return serializeSubscriptionV1(subscription);
+
+            default:
+                // for any versions higher than known, try to serialize it as V1
+                return serializeSubscriptionV1(subscription);
+        }
+    }
+
+    public static Subscription deserializeSubscriptionV0(ByteBuffer buffer) {
         Struct struct = SUBSCRIPTION_V0.read(buffer);
         ByteBuffer userData = struct.getBytes(USER_DATA_KEY_NAME);
         List topics = new ArrayList<>();
         for (Object topicObj : struct.getArray(TOPICS_KEY_NAME))
             topics.add((String) topicObj);
-        return new PartitionAssignor.Subscription(topics, userData);
+
+        return new Subscription(topics, userData, Collections.emptyList());
     }
 
-    public static PartitionAssignor.Assignment deserializeAssignment(ByteBuffer buffer) {
-        Struct header = CONSUMER_PROTOCOL_HEADER_SCHEMA.read(buffer);
-        Short version = header.getShort(VERSION_KEY_NAME);
-        checkVersionCompatibility(version);
-        Struct struct = ASSIGNMENT_V0.read(buffer);
+    public static Subscription deserializeSubscriptionV1(ByteBuffer buffer) {
+        Struct struct = SUBSCRIPTION_V1.read(buffer);
         ByteBuffer userData = struct.getBytes(USER_DATA_KEY_NAME);
-        List partitions = new ArrayList<>();
-        for (Object structObj : struct.getArray(TOPIC_PARTITIONS_KEY_NAME)) {
+        List topics = new ArrayList<>();
+        for (Object topicObj : struct.getArray(TOPICS_KEY_NAME))
+            topics.add((String) topicObj);
+
+        List ownedPartitions = new ArrayList<>();
+        for (Object structObj : struct.getArray(OWNED_PARTITIONS_KEY_NAME)) {
             Struct assignment = (Struct) structObj;
             String topic = assignment.getString(TOPIC_KEY_NAME);
             for (Object partitionObj : assignment.getArray(PARTITIONS_KEY_NAME)) {
-                Integer partition = (Integer) partitionObj;
-                partitions.add(new TopicPartition(topic, partition));
+                ownedPartitions.add(new TopicPartition(topic, (Integer) partitionObj));
             }
         }
-        return new PartitionAssignor.Assignment(partitions, userData);
+
+        return new Subscription(topics, userData, ownedPartitions);
+    }
+
+    public static Subscription deserializeSubscription(ByteBuffer buffer) {
+        Short version = deserializeVersion(buffer);
+
+        if (version < CONSUMER_PROTOCOL_V0)
+            throw new SchemaException("Unsupported subscription version: " + version);
+
+        switch (version) {
+            case CONSUMER_PROTOCOL_V0:
+                return deserializeSubscriptionV0(buffer);
+
+            case CONSUMER_PROTOCOL_V1:
+                return deserializeSubscriptionV1(buffer);
+
+            // assume all higher versions can be parsed as V1
+            default:
+                return deserializeSubscriptionV1(buffer);
+        }
     }
 
-    public static ByteBuffer serializeAssignment(PartitionAssignor.Assignment assignment) {
+    public static ByteBuffer serializeAssignmentV0(Assignment assignment) {
         Struct struct = new Struct(ASSIGNMENT_V0);
         struct.set(USER_DATA_KEY_NAME, assignment.userData());
         List topicAssignments = new ArrayList<>();
@@ -131,6 +236,7 @@ public static ByteBuffer serializeAssignment(PartitionAssignor.Assignment assign
             topicAssignments.add(topicAssignment);
         }
         struct.set(TOPIC_PARTITIONS_KEY_NAME, topicAssignments.toArray());
+
         ByteBuffer buffer = ByteBuffer.allocate(CONSUMER_PROTOCOL_HEADER_V0.sizeOf() + ASSIGNMENT_V0.sizeOf(struct));
         CONSUMER_PROTOCOL_HEADER_V0.writeTo(buffer);
         ASSIGNMENT_V0.write(buffer, struct);
@@ -138,12 +244,78 @@ public static ByteBuffer serializeAssignment(PartitionAssignor.Assignment assign
         return buffer;
     }
 
-    private static void checkVersionCompatibility(short version) {
-        // check for invalid versions
-        if (version < CONSUMER_PROTOCOL_V0)
-            throw new SchemaException("Unsupported subscription version: " + version);
+    public static ByteBuffer serializeAssignmentV1(Assignment assignment) {
+        Struct struct = new Struct(ASSIGNMENT_V1);
+        struct.set(USER_DATA_KEY_NAME, assignment.userData());
+        List topicAssignments = new ArrayList<>();
+        Map> partitionsByTopic = CollectionUtils.groupPartitionsByTopic(assignment.partitions());
+        for (Map.Entry> topicEntry : partitionsByTopic.entrySet()) {
+            Struct topicAssignment = new Struct(TOPIC_ASSIGNMENT_V0);
+            topicAssignment.set(TOPIC_KEY_NAME, topicEntry.getKey());
+            topicAssignment.set(PARTITIONS_KEY_NAME, topicEntry.getValue().toArray());
+            topicAssignments.add(topicAssignment);
+        }
+        struct.set(TOPIC_PARTITIONS_KEY_NAME, topicAssignments.toArray());
 
-        // otherwise, assume versions can be parsed as V0
+        ByteBuffer buffer = ByteBuffer.allocate(CONSUMER_PROTOCOL_HEADER_V1.sizeOf() + ASSIGNMENT_V1.sizeOf(struct));
+        CONSUMER_PROTOCOL_HEADER_V1.writeTo(buffer);
+        ASSIGNMENT_V1.write(buffer, struct);
+        buffer.flip();
+        return buffer;
     }
 
+    public static ByteBuffer serializeAssignment(Assignment assignment) {
+        return serializeAssignment(assignment, CONSUMER_PROTOCOL_LATEST_VERSION);
+    }
+
+    public static ByteBuffer serializeAssignment(Assignment assignment, short version) {
+        switch (version) {
+            case CONSUMER_PROTOCOL_V0:
+                return serializeAssignmentV0(assignment);
+
+            case CONSUMER_PROTOCOL_V1:
+                return serializeAssignmentV1(assignment);
+
+            default:
+                // for any versions higher than known, try to serialize it as V1
+                return serializeAssignmentV1(assignment);
+        }
+    }
+
+    public static Assignment deserializeAssignmentV0(ByteBuffer buffer) {
+        Struct struct = ASSIGNMENT_V0.read(buffer);
+        ByteBuffer userData = struct.getBytes(USER_DATA_KEY_NAME);
+        List partitions = new ArrayList<>();
+        for (Object structObj : struct.getArray(TOPIC_PARTITIONS_KEY_NAME)) {
+            Struct assignment = (Struct) structObj;
+            String topic = assignment.getString(TOPIC_KEY_NAME);
+            for (Object partitionObj : assignment.getArray(PARTITIONS_KEY_NAME)) {
+                partitions.add(new TopicPartition(topic, (Integer) partitionObj));
+            }
+        }
+        return new Assignment(partitions, userData);
+    }
+
+    public static Assignment deserializeAssignmentV1(ByteBuffer buffer) {
+        return deserializeAssignmentV0(buffer);
+    }
+
+    public static Assignment deserializeAssignment(ByteBuffer buffer) {
+        Short version = deserializeVersion(buffer);
+
+        if (version < CONSUMER_PROTOCOL_V0)
+            throw new SchemaException("Unsupported assignment version: " + version);
+
+        switch (version) {
+            case CONSUMER_PROTOCOL_V0:
+                return deserializeAssignmentV0(buffer);
+
+            case CONSUMER_PROTOCOL_V1:
+                return deserializeAssignmentV1(buffer);
+
+            default:
+                // assume all higher versions can be parsed as V1
+                return deserializeAssignmentV1(buffer);
+        }
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/Fetcher.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/Fetcher.java
index 9009ffe092a9f..18ac4001ffdc3 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/Fetcher.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/Fetcher.java
@@ -16,20 +16,27 @@
  */
 package org.apache.kafka.clients.consumer.internals;
 
+import org.apache.kafka.clients.ApiVersions;
 import org.apache.kafka.clients.ClientResponse;
 import org.apache.kafka.clients.FetchSessionHandler;
 import org.apache.kafka.clients.MetadataCache;
+import org.apache.kafka.clients.NodeApiVersions;
+import org.apache.kafka.clients.ApiVersion;
 import org.apache.kafka.clients.StaleMetadataException;
 import org.apache.kafka.clients.consumer.ConsumerConfig;
 import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.LogTruncationException;
+import org.apache.kafka.clients.consumer.OffsetAndMetadata;
 import org.apache.kafka.clients.consumer.OffsetAndTimestamp;
 import org.apache.kafka.clients.consumer.OffsetOutOfRangeException;
 import org.apache.kafka.clients.consumer.OffsetResetStrategy;
 import org.apache.kafka.common.Cluster;
 import org.apache.kafka.common.KafkaException;
+import org.apache.kafka.common.MetricName;
 import org.apache.kafka.common.Node;
 import org.apache.kafka.common.PartitionInfo;
 import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.errors.CorruptRecordException;
 import org.apache.kafka.common.errors.InvalidTopicException;
 import org.apache.kafka.common.errors.RecordTooLargeException;
 import org.apache.kafka.common.errors.RetriableException;
@@ -38,18 +45,20 @@
 import org.apache.kafka.common.errors.TopicAuthorizationException;
 import org.apache.kafka.common.header.Headers;
 import org.apache.kafka.common.header.internals.RecordHeaders;
+import org.apache.kafka.common.metrics.Gauge;
 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.Count;
 import org.apache.kafka.common.metrics.stats.Max;
 import org.apache.kafka.common.metrics.stats.Meter;
 import org.apache.kafka.common.metrics.stats.Min;
 import org.apache.kafka.common.metrics.stats.Value;
+import org.apache.kafka.common.metrics.stats.WindowedCount;
+import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.common.record.AbstractLegacyRecordBatch;
 import org.apache.kafka.common.record.BufferSupplier;
 import org.apache.kafka.common.record.ControlRecordType;
-import org.apache.kafka.common.record.InvalidRecordException;
 import org.apache.kafka.common.record.Record;
 import org.apache.kafka.common.record.RecordBatch;
 import org.apache.kafka.common.record.Records;
@@ -61,6 +70,7 @@
 import org.apache.kafka.common.requests.ListOffsetResponse;
 import org.apache.kafka.common.requests.MetadataRequest;
 import org.apache.kafka.common.requests.MetadataResponse;
+import org.apache.kafka.common.requests.OffsetsForLeaderEpochRequest;
 import org.apache.kafka.common.serialization.Deserializer;
 import org.apache.kafka.common.utils.CloseableIterator;
 import org.apache.kafka.common.utils.LogContext;
@@ -68,9 +78,11 @@
 import org.apache.kafka.common.utils.Timer;
 import org.apache.kafka.common.utils.Utils;
 import org.slf4j.Logger;
+import org.slf4j.helpers.MessageFormatter;
 
 import java.io.Closeable;
 import java.nio.ByteBuffer;
+import java.util.ArrayDeque;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
@@ -83,6 +95,7 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.PriorityQueue;
+import java.util.Queue;
 import java.util.Set;
 import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -91,6 +104,8 @@
 import java.util.stream.Collectors;
 
 import static java.util.Collections.emptyList;
+import static org.apache.kafka.common.record.RecordBatch.MAGIC_VALUE_V2;
+
 
 /**
  * This class manages the fetching process with the brokers.
@@ -107,9 +122,13 @@
  *     the caller.
  *     
    • Responses that collate partial responses from multiple brokers (e.g. to list offsets) are
  *     synchronized on the response future.
      • 
+ *     
    • At most one request is pending for each node at any time. Nodes with pending requests are
+ *     tracked and updated after processing the response. This ensures that any state (e.g. epoch)
+ *     updated while processing responses on one thread are visible while creating the subsequent request
+ *     on a different thread.
      • 
  * 
    
  */
-public class Fetcher implements SubscriptionState.Listener, Closeable {
+public class Fetcher implements Closeable {
     private final Logger log;
     private final LogContext logContext;
     private final ConsumerNetworkClient client;
@@ -122,6 +141,8 @@ public class Fetcher implements SubscriptionState.Listener, Closeable {
     private final long requestTimeoutMs;
     private final int maxPollRecords;
     private final boolean checkCrcs;
+    private final String clientRackId;
+    private final boolean enableShallowIteration;
     private final ConsumerMetadata metadata;
     private final FetchManagerMetrics sensors;
     private final SubscriptionState subscriptions;
@@ -132,8 +153,12 @@ public class Fetcher implements SubscriptionState.Listener, Closeable {
     private final IsolationLevel isolationLevel;
     private final Map sessionHandlers;
     private final AtomicReference cachedListOffsetsException = new AtomicReference<>();
+    private final AtomicReference cachedOffsetForLeaderException = new AtomicReference<>();
+    private final OffsetsForLeaderEpochClient offsetsForLeaderEpochClient;
+    private final Set nodesWithPendingFetchRequests;
+    private final ApiVersions apiVersions;
 
-    private PartitionRecords nextInLineRecords = null;
+    private CompletedFetch nextInLineFetch = null;
 
     public Fetcher(LogContext logContext,
                    ConsumerNetworkClient client,
@@ -143,6 +168,8 @@ public Fetcher(LogContext logContext,
                    int fetchSize,
                    int maxPollRecords,
                    boolean checkCrcs,
+                   String clientRackId,
+                   boolean enableShallowIteration,
                    Deserializer keyDeserializer,
                    Deserializer valueDeserializer,
                    ConsumerMetadata metadata,
@@ -152,7 +179,8 @@ public Fetcher(LogContext logContext,
                    Time time,
                    long retryBackoffMs,
                    long requestTimeoutMs,
-                   IsolationLevel isolationLevel) {
+                   IsolationLevel isolationLevel,
+                   ApiVersions apiVersions) {
         this.log = logContext.logger(Fetcher.class);
         this.logContext = logContext;
         this.time = time;
@@ -165,27 +193,30 @@ public Fetcher(LogContext logContext,
         this.fetchSize = fetchSize;
         this.maxPollRecords = maxPollRecords;
         this.checkCrcs = checkCrcs;
+        this.clientRackId = clientRackId;
         this.keyDeserializer = keyDeserializer;
         this.valueDeserializer = valueDeserializer;
+        this.enableShallowIteration = enableShallowIteration;
         this.completedFetches = new ConcurrentLinkedQueue<>();
         this.sensors = new FetchManagerMetrics(metrics, metricsRegistry);
         this.retryBackoffMs = retryBackoffMs;
         this.requestTimeoutMs = requestTimeoutMs;
         this.isolationLevel = isolationLevel;
+        this.apiVersions = apiVersions;
         this.sessionHandlers = new HashMap<>();
-
-        subscriptions.addListener(this);
+        this.offsetsForLeaderEpochClient = new OffsetsForLeaderEpochClient(client, logContext);
+        this.nodesWithPendingFetchRequests = new HashSet<>();
     }
 
     /**
      * Represents data about an offset returned by a broker.
      */
-    private static class OffsetData {
+    private static class ListOffsetData {
         final long offset;
         final Long timestamp; //  null if the broker does not support returning timestamps
         final Optional leaderEpoch; // empty if the leader epoch is not known
 
-        OffsetData(long offset, Long timestamp, Optional leaderEpoch) {
+        ListOffsetData(long offset, Long timestamp, Optional leaderEpoch) {
             this.offset = offset;
             this.timestamp = timestamp;
             this.leaderEpoch = leaderEpoch;
@@ -193,19 +224,31 @@ private static class OffsetData {
     }
 
     /**
-     * Return whether we have any completed fetches pending return to the user. This method is thread-safe.
+     * Return whether we have any completed fetches pending return to the user. This method is thread-safe. Has
+     * visibility for testing.
      * @return true if there are completed fetches, false otherwise
      */
-    public boolean hasCompletedFetches() {
+    protected boolean hasCompletedFetches() {
         return !completedFetches.isEmpty();
     }
 
+    /**
+     * Return whether we have any completed fetches that are fetchable. This method is thread-safe.
+     * @return true if there are completed fetches that can be returned, false otherwise
+     */
+    public boolean hasAvailableFetches() {
+        return completedFetches.stream().anyMatch(fetch -> subscriptions.isFetchable(fetch.partition));
+    }
+
     /**
      * Set-up a fetch request for any node that we have assigned partitions for which doesn't already have
      * an in-flight fetch or pending fetch data.
      * @return number of fetches sent
      */
     public synchronized int sendFetches() {
+        // Update metrics in case there was an assignment change
+        sensors.maybeUpdateAssignment(subscriptions);
+
         Map fetchRequestMap = prepareFetchRequests();
         for (Map.Entry entry : fetchRequestMap.entrySet()) {
             final Node fetchTarget = entry.getKey();
@@ -215,55 +258,92 @@ public synchronized int sendFetches() {
                     .isolationLevel(isolationLevel)
                     .setMaxBytes(this.maxBytes)
                     .metadata(data.metadata())
-                    .toForget(data.toForget());
+                    .toForget(data.toForget())
+                    .rackId(clientRackId);
+
             if (log.isDebugEnabled()) {
                 log.debug("Sending {} {} to broker {}", isolationLevel, data.toString(), fetchTarget);
             }
-            client.send(fetchTarget, request)
-                    .addListener(new RequestFutureListener() {
-                        @Override
-                        public void onSuccess(ClientResponse resp) {
-                            synchronized (Fetcher.this) {
-                                @SuppressWarnings("unchecked")
-                                FetchResponse response = (FetchResponse) resp.responseBody();
-                                FetchSessionHandler handler = sessionHandler(fetchTarget.id());
-                                if (handler == null) {
-                                    log.error("Unable to find FetchSessionHandler for node {}. Ignoring fetch response.",
-                                            fetchTarget.id());
-                                    return;
-                                }
-                                if (!handler.handleResponse(response)) {
-                                    return;
-                                }
-
-                                Set partitions = new HashSet<>(response.responseData().keySet());
-                                FetchResponseMetricAggregator metricAggregator = new FetchResponseMetricAggregator(sensors, partitions);
+            RequestFuture future = client.send(fetchTarget, request);
+            // We add the node to the set of nodes with pending fetch requests before adding the
+            // listener because the future may have been fulfilled on another thread (e.g. during a
+            // disconnection being handled by the heartbeat thread) which will mean the listener
+            // will be invoked synchronously.
+            this.nodesWithPendingFetchRequests.add(entry.getKey().id());
+            future.addListener(new RequestFutureListener() {
+                @Override
+                public void onSuccess(ClientResponse resp) {
+                    synchronized (Fetcher.this) {
+                        try {
+                            @SuppressWarnings("unchecked")
+                            FetchResponse response = (FetchResponse) resp.responseBody();
+                            FetchSessionHandler handler = sessionHandler(fetchTarget.id());
+                            if (handler == null) {
+                                log.error("Unable to find FetchSessionHandler for node {}. Ignoring fetch response.",
+                                        fetchTarget.id());
+                                return;
+                            }
+                            if (!handler.handleResponse(response)) {
+                                return;
+                            }
 
-                                for (Map.Entry> entry : response.responseData().entrySet()) {
-                                    TopicPartition partition = entry.getKey();
-                                    long fetchOffset = data.sessionPartitions().get(partition).fetchOffset;
-                                    FetchResponse.PartitionData fetchData = entry.getValue();
+                            Set partitions = new HashSet<>(response.responseData().keySet());
+                            FetchResponseMetricAggregator metricAggregator = new FetchResponseMetricAggregator(sensors, partitions);
+
+                            for (Map.Entry> entry : response.responseData().entrySet()) {
+                                TopicPartition partition = entry.getKey();
+                                FetchRequest.PartitionData requestData = data.sessionPartitions().get(partition);
+                                if (requestData == null) {
+                                    String message;
+                                    if (data.metadata().isFull()) {
+                                        message = MessageFormatter.arrayFormat(
+                                                "Response for missing full request partition: partition={}; metadata={}",
+                                                new Object[]{partition, data.metadata()}).getMessage();
+                                    } else {
+                                        message = MessageFormatter.arrayFormat(
+                                                "Response for missing session request partition: partition={}; metadata={}; toSend={}; toForget={}",
+                                                new Object[]{partition, data.metadata(), data.toSend(), data.toForget()}).getMessage();
+                                    }
+
+                                    // Received fetch response for missing session partition
+                                    throw new IllegalStateException(message);
+                                } else {
+                                    long fetchOffset = requestData.fetchOffset;
+                                    FetchResponse.PartitionData partitionData = entry.getValue();
 
                                     log.debug("Fetch {} at offset {} for partition {} returned fetch data {}",
-                                            isolationLevel, fetchOffset, partition, fetchData);
-                                    completedFetches.add(new CompletedFetch(partition, fetchOffset, fetchData, metricAggregator,
-                                            resp.requestHeader().apiVersion()));
-                                }
+                                            isolationLevel, fetchOffset, partition, partitionData);
+
+                                    Iterator batches = partitionData.records.batches().iterator();
+                                    short responseVersion = resp.requestHeader().apiVersion();
 
-                                sensors.fetchLatency.record(resp.requestLatencyMs());
+                                    completedFetches.add(new CompletedFetch(partition, partitionData,
+                                            metricAggregator, batches, fetchOffset, responseVersion));
+                                }
                             }
+
+                            sensors.fetchLatency.record(resp.requestLatencyMs());
+                        } finally {
+                            nodesWithPendingFetchRequests.remove(fetchTarget.id());
                         }
+                    }
+                }
 
-                        @Override
-                        public void onFailure(RuntimeException e) {
-                            synchronized (Fetcher.this) {
-                                FetchSessionHandler handler = sessionHandler(fetchTarget.id());
-                                if (handler != null) {
-                                    handler.handleError(e);
-                                }
+                @Override
+                public void onFailure(RuntimeException e) {
+                    synchronized (Fetcher.this) {
+                        try {
+                            FetchSessionHandler handler = sessionHandler(fetchTarget.id());
+                            if (handler != null) {
+                                handler.handleError(e);
                             }
+                        } finally {
+                            nodesWithPendingFetchRequests.remove(fetchTarget.id());
                         }
-                    });
+                    }
+                }
+            });
+
         }
         return fetchRequestMap.size();
     }
@@ -286,7 +366,7 @@ public Map> getAllTopicMetadata(Timer timer) {
      */
     public Map> getTopicMetadata(MetadataRequest.Builder request, Timer timer) {
         // Save the round trip if no topics are requested.
-        if (!request.isAllTopics() && request.topics().isEmpty())
+        if (!request.isAllTopics() && request.emptyTopicList())
             return Collections.emptyMap();
 
         do {
@@ -366,6 +446,15 @@ else if (strategy == OffsetResetStrategy.LATEST)
             return null;
     }
 
+    private OffsetResetStrategy timestampToOffsetResetStrategy(long timestamp) {
+        if (timestamp == ListOffsetRequest.EARLIEST_TIMESTAMP)
+            return OffsetResetStrategy.EARLIEST;
+        else if (timestamp == ListOffsetRequest.LATEST_TIMESTAMP)
+            return OffsetResetStrategy.LATEST;
+        else
+            return null;
+    }
+
     /**
      * Reset offsets for all assigned partitions that require it.
      *
@@ -392,22 +481,45 @@ public void resetOffsetsIfNeeded() {
         resetOffsetsAsync(offsetResetTimestamps);
     }
 
+    /**
+     *  Validate offsets for all assigned partitions for which a leader change has been detected.
+     */
+    public void validateOffsetsIfNeeded() {
+        RuntimeException exception = cachedOffsetForLeaderException.getAndSet(null);
+        if (exception != null)
+            throw exception;
+
+        // Validate each partition against the current leader and epoch
+        subscriptions.assignedPartitions().forEach(topicPartition -> {
+            ConsumerMetadata.LeaderAndEpoch leaderAndEpoch = metadata.leaderAndEpoch(topicPartition);
+            subscriptions.maybeValidatePositionForCurrentLeader(topicPartition, leaderAndEpoch);
+        });
+
+        // Collect positions needing validation, with backoff
+        Map partitionsToValidate = subscriptions
+                .partitionsNeedingValidation(time.milliseconds())
+                .stream()
+                .collect(Collectors.toMap(Function.identity(), subscriptions::position));
+
+        validateOffsetsAsync(partitionsToValidate);
+    }
+
     public Map offsetsForTimes(Map timestampsToSearch,
                                                                    Timer timer) {
         metadata.addTransientTopics(topicsForPartitions(timestampsToSearch.keySet()));
 
         try {
-            Map fetchedOffsets = fetchOffsetsByTimes(timestampsToSearch,
+            Map fetchedOffsets = fetchOffsetsByTimes(timestampsToSearch,
                     timer, true).fetchedOffsets;
 
             HashMap offsetsByTimes = new HashMap<>(timestampsToSearch.size());
             for (Map.Entry entry : timestampsToSearch.entrySet())
                 offsetsByTimes.put(entry.getKey(), null);
 
-            for (Map.Entry entry : fetchedOffsets.entrySet()) {
+            for (Map.Entry entry : fetchedOffsets.entrySet()) {
                 // 'entry.getValue().timestamp' will not be null since we are guaranteed
                 // to work with a v1 (or later) ListOffset request
-                OffsetData offsetData = entry.getValue();
+                ListOffsetData offsetData = entry.getValue();
                 offsetsByTimes.put(entry.getKey(), new OffsetAndTimestamp(offsetData.offset, offsetData.timestamp,
                         offsetData.leaderEpoch));
             }
@@ -430,24 +542,21 @@ private ListOffsetResult fetchOffsetsByTimes(Map timestamp
             RequestFuture future = sendListOffsetsRequests(remainingToSearch, requireTimestamps);
             client.poll(future, timer);
 
-            if (!future.isDone())
+            if (!future.isDone()) {
                 break;
-
-            if (future.succeeded()) {
+            } else if (future.succeeded()) {
                 ListOffsetResult value = future.value();
                 result.fetchedOffsets.putAll(value.fetchedOffsets);
-                if (value.partitionsToRetry.isEmpty())
-                    return result;
-
                 remainingToSearch.keySet().retainAll(value.partitionsToRetry);
             } else if (!future.isRetriable()) {
                 throw future.exception();
             }
 
-            if (metadata.updateRequested())
+            if (remainingToSearch.isEmpty()) {
+                return result;
+            } else {
                 client.awaitMetadataUpdate(timer);
-            else
-                timer.sleep(retryBackoffMs);
+            }
         } while (timer.notExpired());
 
         throw new TimeoutException("Failed to get offsets by times in " + timer.elapsedMs() + "ms");
@@ -490,33 +599,45 @@ private Map beginningOrEndOffset(Collection>> fetchedRecords() {
         Map>> fetched = new HashMap<>();
+        Queue pausedCompletedFetches = new ArrayDeque<>();
         int recordsRemaining = maxPollRecords;
 
         try {
             while (recordsRemaining > 0) {
-                if (nextInLineRecords == null || nextInLineRecords.isFetched) {
-                    CompletedFetch completedFetch = completedFetches.peek();
-                    if (completedFetch == null) break;
-
-                    try {
-                        nextInLineRecords = parseCompletedFetch(completedFetch);
-                    } catch (Exception e) {
-                        // Remove a completedFetch upon a parse with exception if (1) it contains no records, and
-                        // (2) there are no fetched records with actual content preceding this exception.
-                        // The first condition ensures that the completedFetches is not stuck with the same completedFetch
-                        // in cases such as the TopicAuthorizationException, and the second condition ensures that no
-                        // potential data loss due to an exception in a following record.
-                        FetchResponse.PartitionData partition = completedFetch.partitionData;
-                        if (fetched.isEmpty() && (partition.records == null || partition.records.sizeInBytes() == 0)) {
-                            completedFetches.poll();
+                if (nextInLineFetch == null || nextInLineFetch.isConsumed) {
+                    CompletedFetch records = completedFetches.peek();
+                    if (records == null) break;
+
+                    if (records.notInitialized()) {
+                        try {
+                            nextInLineFetch = initializeCompletedFetch(records);
+                        } catch (Exception e) {
+                            // Remove a completedFetch upon a parse with exception if (1) it contains no records, and
+                            // (2) there are no fetched records with actual content preceding this exception.
+                            // The first condition ensures that the completedFetches is not stuck with the same completedFetch
+                            // in cases such as the TopicAuthorizationException, and the second condition ensures that no
+                            // potential data loss due to an exception in a following record.
+                            FetchResponse.PartitionData partition = records.partitionData;
+                            if (fetched.isEmpty() && (partition.records == null || partition.records.sizeInBytes() == 0)) {
+                                completedFetches.poll();
+                            }
+                            throw e;
                         }
-                        throw e;
+                    } else {
+                        nextInLineFetch = records;
                     }
                     completedFetches.poll();
+                } else if (subscriptions.isPaused(nextInLineFetch.partition)) {
+                    // when the partition is paused we add the records back to the completedFetches queue instead of draining
+                    // them so that they can be returned on a subsequent poll if the partition is resumed at that time
+                    log.debug("Skipping fetching records for assigned partition {} because it is paused", nextInLineFetch.partition);
+                    pausedCompletedFetches.add(nextInLineFetch);
+                    nextInLineFetch = null;
                 } else {
-                    List> records = fetchRecords(nextInLineRecords, recordsRemaining);
-                    TopicPartition partition = nextInLineRecords.partition;
+                    List> records = fetchRecords(nextInLineFetch, recordsRemaining);
+
                     if (!records.isEmpty()) {
+                        TopicPartition partition = nextInLineFetch.partition;
                         List> currentRecords = fetched.get(partition);
                         if (currentRecords == null) {
                             fetched.put(partition, records);
@@ -536,37 +657,47 @@ public Map>> fetchedRecords() {
         } catch (KafkaException e) {
             if (fetched.isEmpty())
                 throw e;
+        } finally {
+            // add any polled completed fetches for paused partitions back to the completed fetches queue to be
+            // re-evaluated in the next poll
+            completedFetches.addAll(pausedCompletedFetches);
         }
+
         return fetched;
     }
 
-    private List> fetchRecords(PartitionRecords partitionRecords, int maxRecords) {
-        if (!subscriptions.isAssigned(partitionRecords.partition)) {
+    private List> fetchRecords(CompletedFetch completedFetch, int maxRecords) {
+        if (!subscriptions.isAssigned(completedFetch.partition)) {
             // this can happen when a rebalance happened before fetched records are returned to the consumer's poll call
             log.debug("Not returning fetched records for partition {} since it is no longer assigned",
-                    partitionRecords.partition);
-        } else if (!subscriptions.isFetchable(partitionRecords.partition)) {
+                    completedFetch.partition);
+        } else if (!subscriptions.isFetchable(completedFetch.partition)) {
             // this can happen when a partition is paused before fetched records are returned to the consumer's
             // poll call or if the offset is being reset
             log.debug("Not returning fetched records for assigned partition {} since it is no longer fetchable",
-                    partitionRecords.partition);
+                    completedFetch.partition);
         } else {
-            long position = subscriptions.position(partitionRecords.partition);
-            if (partitionRecords.nextFetchOffset == position) {
-                List> partRecords = partitionRecords.fetchRecords(maxRecords);
-
-                long nextOffset = partitionRecords.nextFetchOffset;
-                log.trace("Returning fetched records at offset {} for assigned partition {} and update " +
-                        "position to {}", position, partitionRecords.partition, nextOffset);
-                subscriptions.position(partitionRecords.partition, nextOffset);
+            SubscriptionState.FetchPosition position = subscriptions.position(completedFetch.partition);
+            if (completedFetch.nextFetchOffset == position.offset) {
+                List> partRecords = completedFetch.fetchRecords(maxRecords);
+
+                if (completedFetch.nextFetchOffset > position.offset) {
+                    SubscriptionState.FetchPosition nextPosition = new SubscriptionState.FetchPosition(
+                            completedFetch.nextFetchOffset,
+                            completedFetch.lastEpoch,
+                            position.currentLeader);
+                    log.trace("Returning fetched records at offset {} for assigned partition {} and update " +
+                            "position to {}", position, completedFetch.partition, nextPosition);
+                    subscriptions.position(completedFetch.partition, nextPosition);
+                }
 
-                Long partitionLag = subscriptions.partitionLag(partitionRecords.partition, isolationLevel);
+                Long partitionLag = subscriptions.partitionLag(completedFetch.partition, isolationLevel);
                 if (partitionLag != null)
-                    this.sensors.recordPartitionLag(partitionRecords.partition, partitionLag);
+                    this.sensors.recordPartitionLag(completedFetch.partition, partitionLag);
 
-                Long lead = subscriptions.partitionLead(partitionRecords.partition);
+                Long lead = subscriptions.partitionLead(completedFetch.partition);
                 if (lead != null) {
-                    this.sensors.recordPartitionLead(partitionRecords.partition, lead);
+                    this.sensors.recordPartitionLead(completedFetch.partition, lead);
                 }
 
                 return partRecords;
@@ -574,28 +705,21 @@ private List> fetchRecords(PartitionRecords partitionRecord
                 // these records aren't next in line based on the last consumed position, ignore them
                 // they must be from an obsolete request
                 log.debug("Ignoring fetched records for {} at offset {} since the current position is {}",
-                        partitionRecords.partition, partitionRecords.nextFetchOffset, position);
+                        completedFetch.partition, completedFetch.nextFetchOffset, position);
             }
         }
 
-        partitionRecords.drain();
+        log.trace("Draining fetched records for partition {}", completedFetch.partition);
+        completedFetch.drain();
+
         return emptyList();
     }
 
-    private void resetOffsetIfNeeded(TopicPartition partition, Long requestedResetTimestamp, OffsetData offsetData) {
-        // we might lose the assignment while fetching the offset, or the user might seek to a different offset,
-        // so verify it is still assigned and still in need of the requested reset
-        if (!subscriptions.isAssigned(partition)) {
-            log.debug("Skipping reset of partition {} since it is no longer assigned", partition);
-        } else if (!subscriptions.isOffsetResetNeeded(partition)) {
-            log.debug("Skipping reset of partition {} since reset is no longer needed", partition);
-        } else if (!requestedResetTimestamp.equals(offsetResetStrategyTimestamp(partition))) {
-            log.debug("Skipping reset of partition {} since an alternative reset has been requested", partition);
-        } else {
-            log.info("Resetting offset for partition {} to offset {}.", partition, offsetData.offset);
-            offsetData.leaderEpoch.ifPresent(epoch -> metadata.updateLastSeenEpochIfNewer(partition, epoch));
-            subscriptions.seek(partition, offsetData.offset);
-        }
+    private void resetOffsetIfNeeded(TopicPartition partition, OffsetResetStrategy requestedResetStrategy, ListOffsetData offsetData) {
+        SubscriptionState.FetchPosition position = new SubscriptionState.FetchPosition(
+                offsetData.offset, offsetData.leaderEpoch, metadata.leaderAndEpoch(partition));
+        offsetData.leaderEpoch.ifPresent(epoch -> metadata.updateLastSeenEpochIfNewer(partition, epoch));
+        subscriptions.maybeSeekUnvalidated(partition, position.offset, requestedResetStrategy);
     }
 
     private void resetOffsetsAsync(Map partitionResetTimestamps) {
@@ -604,28 +728,28 @@ private void resetOffsetsAsync(Map partitionResetTimestamp
         for (Map.Entry> entry : timestampsToSearchByNode.entrySet()) {
             Node node = entry.getKey();
             final Map resetTimestamps = entry.getValue();
-            subscriptions.setResetPending(resetTimestamps.keySet(), time.milliseconds() + requestTimeoutMs);
+            subscriptions.setNextAllowedRetry(resetTimestamps.keySet(), time.milliseconds() + requestTimeoutMs);
 
             RequestFuture future = sendListOffsetRequest(node, resetTimestamps, false);
             future.addListener(new RequestFutureListener() {
                 @Override
                 public void onSuccess(ListOffsetResult result) {
                     if (!result.partitionsToRetry.isEmpty()) {
-                        subscriptions.resetFailed(result.partitionsToRetry, time.milliseconds() + retryBackoffMs);
+                        subscriptions.requestFailed(result.partitionsToRetry, time.milliseconds() + retryBackoffMs);
                         metadata.requestUpdate();
                     }
 
-                    for (Map.Entry fetchedOffset : result.fetchedOffsets.entrySet()) {
+                    for (Map.Entry fetchedOffset : result.fetchedOffsets.entrySet()) {
                         TopicPartition partition = fetchedOffset.getKey();
-                        OffsetData offsetData = fetchedOffset.getValue();
+                        ListOffsetData offsetData = fetchedOffset.getValue();
                         ListOffsetRequest.PartitionData requestedReset = resetTimestamps.get(partition);
-                        resetOffsetIfNeeded(partition, requestedReset.timestamp, offsetData);
+                        resetOffsetIfNeeded(partition, timestampToOffsetResetStrategy(requestedReset.timestamp), offsetData);
                     }
                 }
 
                 @Override
                 public void onFailure(RuntimeException e) {
-                    subscriptions.resetFailed(resetTimestamps.keySet(), time.milliseconds() + retryBackoffMs);
+                    subscriptions.requestFailed(resetTimestamps.keySet(), time.milliseconds() + retryBackoffMs);
                     metadata.requestUpdate();
 
                     if (!(e instanceof RetriableException) && !cachedListOffsetsException.compareAndSet(null, e))
@@ -635,6 +759,88 @@ public void onFailure(RuntimeException e) {
         }
     }
 
+    private boolean hasUsableOffsetForLeaderEpochVersion(NodeApiVersions nodeApiVersions) {
+        ApiVersion apiVersion = nodeApiVersions.apiVersion(ApiKeys.OFFSET_FOR_LEADER_EPOCH);
+        if (apiVersion == null)
+            return false;
+
+        return OffsetsForLeaderEpochRequest.supportsTopicPermission(apiVersion.maxVersion);
+    }
+
+    /**
+     * For each partition which needs validation, make an asynchronous request to get the end-offsets for the partition
+     * with the epoch less than or equal to the epoch the partition last saw.
+     *
+     * Requests are grouped by Node for efficiency.
+     */
+    private void validateOffsetsAsync(Map partitionsToValidate) {
+        final Map> regrouped =
+                regroupFetchPositionsByLeader(partitionsToValidate);
+
+        regrouped.forEach((node, fetchPostitions) -> {
+            if (node.isEmpty()) {
+                metadata.requestUpdate();
+                return;
+            }
+
+            NodeApiVersions nodeApiVersions = apiVersions.get(node.idString());
+            if (nodeApiVersions == null) {
+                client.tryConnect(node);
+                return;
+            }
+
+            if (!hasUsableOffsetForLeaderEpochVersion(nodeApiVersions)) {
+                log.debug("Skipping validation of fetch offsets for partitions {} since the broker does not " +
+                                "support the required protocol version (introduced in Kafka 2.3)",
+                        fetchPostitions.keySet());
+                for (TopicPartition partition : fetchPostitions.keySet()) {
+                    subscriptions.completeValidation(partition);
+                }
+                return;
+            }
+
+            subscriptions.setNextAllowedRetry(fetchPostitions.keySet(), time.milliseconds() + requestTimeoutMs);
+
+            RequestFuture future = offsetsForLeaderEpochClient.sendAsyncRequest(node, partitionsToValidate);
+            future.addListener(new RequestFutureListener() {
+                @Override
+                public void onSuccess(OffsetsForLeaderEpochClient.OffsetForEpochResult offsetsResult) {
+                    Map truncationWithoutResetPolicy = new HashMap<>();
+                    if (!offsetsResult.partitionsToRetry().isEmpty()) {
+                        subscriptions.setNextAllowedRetry(offsetsResult.partitionsToRetry(), time.milliseconds() + retryBackoffMs);
+                        metadata.requestUpdate();
+                    }
+
+                    // For each OffsetsForLeader response, check if the end-offset is lower than our current offset
+                    // for the partition. If so, it means we have experienced log truncation and need to reposition
+                    // that partition's offset.
+                    offsetsResult.endOffsets().forEach((respTopicPartition, respEndOffset) -> {
+                        SubscriptionState.FetchPosition requestPosition = fetchPostitions.get(respTopicPartition);
+                        Optional divergentOffsetOpt = subscriptions.maybeCompleteValidation(
+                                respTopicPartition, requestPosition, respEndOffset);
+                        divergentOffsetOpt.ifPresent(divergentOffset -> {
+                            truncationWithoutResetPolicy.put(respTopicPartition, divergentOffset);
+                        });
+                    });
+
+                    if (!truncationWithoutResetPolicy.isEmpty()) {
+                        throw new LogTruncationException(truncationWithoutResetPolicy);
+                    }
+                }
+
+                @Override
+                public void onFailure(RuntimeException e) {
+                    subscriptions.requestFailed(fetchPostitions.keySet(), time.milliseconds() + retryBackoffMs);
+                    metadata.requestUpdate();
+
+                    if (!(e instanceof RetriableException) && !cachedOffsetForLeaderException.compareAndSet(null, e)) {
+                        log.error("Discarding error in OffsetsForLeaderEpoch because another error is pending", e);
+                    }
+                }
+            });
+        });
+    }
+
     /**
      * Search the offsets by target times for the specified partitions.
      *
@@ -652,7 +858,7 @@ private RequestFuture sendListOffsetsRequests(final Map listOffsetRequestsFuture = new RequestFuture<>();
-        final Map fetchedTimestampOffsets = new HashMap<>();
+        final Map fetchedTimestampOffsets = new HashMap<>();
         final AtomicInteger remainingResponses = new AtomicInteger(timestampsToSearchByNode.size());
 
         for (Map.Entry> entry : timestampsToSearchByNode.entrySet()) {
@@ -693,17 +899,19 @@ public void onFailure(RuntimeException e) {
      *                          that need metadata update or re-connect to the leader.
      */
     private Map> groupListOffsetRequests(
-            Map timestampsToSearch, Set partitionsToRetry) {
-        final Map> timestampsToSearchByNode = new HashMap<>();
+            Map timestampsToSearch,
+            Set partitionsToRetry) {
+        final Map partitionDataMap = new HashMap<>();
         for (Map.Entry entry: timestampsToSearch.entrySet()) {
             TopicPartition tp  = entry.getKey();
+            Long offset = entry.getValue();
             Optional currentInfo = metadata.partitionInfoIfCurrent(tp);
             if (!currentInfo.isPresent()) {
-                log.debug("Leader for partition {} is unknown for fetching offset", tp);
+                log.debug("Leader for partition {} is unknown for fetching offset {}", tp, offset);
                 metadata.requestUpdate();
                 partitionsToRetry.add(tp);
             } else if (currentInfo.get().partitionInfo().leader() == null) {
-                log.debug("Leader for partition {} is unavailable for fetching offset", tp);
+                log.debug("Leader for partition {} is unavailable for fetching offset {}", tp, offset);
                 metadata.requestUpdate();
                 partitionsToRetry.add(tp);
             } else if (client.isUnavailable(currentInfo.get().partitionInfo().leader())) {
@@ -716,15 +924,12 @@ private Map> groupLis
                         currentInfo.get().partitionInfo().leader(), tp);
                 partitionsToRetry.add(tp);
             } else {
-                Node node = currentInfo.get().partitionInfo().leader();
-                Map topicData =
-                        timestampsToSearchByNode.computeIfAbsent(node, n -> new HashMap<>());
-                ListOffsetRequest.PartitionData partitionData = new ListOffsetRequest.PartitionData(
-                        entry.getValue(), Optional.of(currentInfo.get().epoch()));
-                topicData.put(entry.getKey(), partitionData);
+                int currentEpoch = currentInfo.get().epoch();
+                Optional currentEpochOpt = currentEpoch < 0 ? Optional.empty() : Optional.of(currentEpoch);
+                partitionDataMap.put(tp, new ListOffsetRequest.PartitionData(offset, currentEpochOpt));
             }
         }
-        return timestampsToSearchByNode;
+        return regroupPartitionMapByNode(partitionDataMap);
     }
 
     /**
@@ -769,7 +974,7 @@ public void onSuccess(ClientResponse response, RequestFuture f
     private void handleListOffsetResponse(Map timestampsToSearch,
                                           ListOffsetResponse listOffsetResponse,
                                           RequestFuture future) {
-        Map fetchedOffsets = new HashMap<>();
+        Map fetchedOffsets = new HashMap<>();
         Set partitionsToRetry = new HashSet<>();
         Set unauthorizedTopics = new HashSet<>();
 
@@ -777,62 +982,66 @@ private void handleListOffsetResponse(Map 1) {
-                        future.raise(new IllegalStateException("Unexpected partitionData response of length " +
-                                partitionData.offsets.size()));
-                        return;
-                    } else if (partitionData.offsets.isEmpty()) {
-                        offset = ListOffsetResponse.UNKNOWN_OFFSET;
-                    } else {
-                        offset = partitionData.offsets.get(0);
-                    }
-                    log.debug("Handling v0 ListOffsetResponse response for {}. Fetched offset {}",
+            switch (error) {
+                case NONE:
+                    if (partitionData.offsets != null) {
+                        // Handle v0 response
+                        long offset;
+                        if (partitionData.offsets.size() > 1) {
+                            future.raise(new IllegalStateException("Unexpected partitionData response of length " +
+                                                                       partitionData.offsets.size()));
+                            return;
+                        } else if (partitionData.offsets.isEmpty()) {
+                            offset = ListOffsetResponse.UNKNOWN_OFFSET;
+                        } else {
+                            offset = partitionData.offsets.get(0);
+                        }
+                        log.debug("Handling v0 ListOffsetResponse response for {}. Fetched offset {}",
                             topicPartition, offset);
-                    if (offset != ListOffsetResponse.UNKNOWN_OFFSET) {
-                        OffsetData offsetData = new OffsetData(offset, null, Optional.empty());
-                        fetchedOffsets.put(topicPartition, offsetData);
-                    }
-                } else {
-                    // Handle v1 and later response
-                    log.debug("Handling ListOffsetResponse response for {}. Fetched offset {}, timestamp {}",
+                        if (offset != ListOffsetResponse.UNKNOWN_OFFSET) {
+                            ListOffsetData offsetData = new ListOffsetData(offset, null, Optional.empty());
+                            fetchedOffsets.put(topicPartition, offsetData);
+                        }
+                    } else {
+                        // Handle v1 and later response
+                        log.debug("Handling ListOffsetResponse response for {}. Fetched offset {}, timestamp {}",
                             topicPartition, partitionData.offset, partitionData.timestamp);
-                    if (partitionData.offset != ListOffsetResponse.UNKNOWN_OFFSET) {
-                        OffsetData offsetData = new OffsetData(partitionData.offset, partitionData.timestamp,
+                        if (partitionData.offset != ListOffsetResponse.UNKNOWN_OFFSET) {
+                            ListOffsetData offsetData = new ListOffsetData(partitionData.offset, partitionData.timestamp,
                                 partitionData.leaderEpoch);
-                        fetchedOffsets.put(topicPartition, offsetData);
+                            fetchedOffsets.put(topicPartition, offsetData);
+                        }
                     }
-                }
-            } else if (error == Errors.UNSUPPORTED_FOR_MESSAGE_FORMAT) {
-                // The message format on the broker side is before 0.10.0, which means it does not
-                // support timestamps. We treat this case the same as if we weren't able to find an
-                // offset corresponding to the requested timestamp and leave it out of the result.
-                log.debug("Cannot search by timestamp for partition {} because the message format version " +
-                        "is before 0.10.0", topicPartition);
-            } else if (error == Errors.NOT_LEADER_FOR_PARTITION ||
-                       error == Errors.REPLICA_NOT_AVAILABLE ||
-                       error == Errors.KAFKA_STORAGE_ERROR ||
-                       error == Errors.OFFSET_NOT_AVAILABLE ||
-                       error == Errors.LEADER_NOT_AVAILABLE) {
-                log.debug("Attempt to fetch offsets for partition {} failed due to {}, retrying.",
-                        topicPartition, error);
-                partitionsToRetry.add(topicPartition);
-            } else if (error == Errors.FENCED_LEADER_EPOCH ||
-                       error == Errors.UNKNOWN_LEADER_EPOCH) {
-                log.debug("Attempt to fetch offsets for partition {} failed due to {}, retrying.",
+                    break;
+                case UNSUPPORTED_FOR_MESSAGE_FORMAT:
+                    // The message format on the broker side is before 0.10.0, which means it does not
+                    // support timestamps. We treat this case the same as if we weren't able to find an
+                    // offset corresponding to the requested timestamp and leave it out of the result.
+                    log.debug("Cannot search by timestamp for partition {} because the message format version " +
+                                  "is before 0.10.0", topicPartition);
+                    break;
+                case NOT_LEADER_FOR_PARTITION:
+                case REPLICA_NOT_AVAILABLE:
+                case KAFKA_STORAGE_ERROR:
+                case OFFSET_NOT_AVAILABLE:
+                case LEADER_NOT_AVAILABLE:
+                case FENCED_LEADER_EPOCH:
+                case UNKNOWN_LEADER_EPOCH:
+                    log.debug("Attempt to fetch offsets for partition {} failed due to {}, retrying.",
                         topicPartition, error);
-                partitionsToRetry.add(topicPartition);
-            } else if (error == Errors.UNKNOWN_TOPIC_OR_PARTITION) {
-                log.warn("Received unknown topic or partition error in ListOffset request for partition {}", topicPartition);
-                partitionsToRetry.add(topicPartition);
-            } else if (error == Errors.TOPIC_AUTHORIZATION_FAILED) {
-                unauthorizedTopics.add(topicPartition.topic());
-            } else {
-                log.warn("Attempt to fetch offsets for partition {} failed due to: {}, retrying.", topicPartition, error.message());
-                partitionsToRetry.add(topicPartition);
+                    partitionsToRetry.add(topicPartition);
+                    break;
+                case UNKNOWN_TOPIC_OR_PARTITION:
+                    log.warn("Received unknown topic or partition error in ListOffset request for partition {}", topicPartition);
+                    partitionsToRetry.add(topicPartition);
+                    break;
+                case TOPIC_AUTHORIZATION_FAILED:
+                    unauthorizedTopics.add(topicPartition.topic());
+                    break;
+                default:
+                    log.warn("Attempt to fetch offsets for partition {} failed due to unexpected exception: {}, retrying.",
+                        topicPartition, error.message());
+                    partitionsToRetry.add(topicPartition);
             }
         }
 
@@ -842,11 +1051,11 @@ private void handleListOffsetResponse(Map fetchedOffsets;
+    static class ListOffsetResult {
+        private final Map fetchedOffsets;
         private final Set partitionsToRetry;
 
-        public ListOffsetResult(Map fetchedOffsets, Set partitionsNeedingRetry) {
+        public ListOffsetResult(Map fetchedOffsets, Set partitionsNeedingRetry) {
             this.fetchedOffsets = fetchedOffsets;
             this.partitionsToRetry = partitionsNeedingRetry;
         }
@@ -859,15 +1068,33 @@ public ListOffsetResult() {
 
     private List fetchablePartitions() {
         Set exclude = new HashSet<>();
-        List fetchable = subscriptions.fetchablePartitions();
-        if (nextInLineRecords != null && !nextInLineRecords.isFetched) {
-            exclude.add(nextInLineRecords.partition);
+        if (nextInLineFetch != null && !nextInLineFetch.isConsumed) {
+            exclude.add(nextInLineFetch.partition);
         }
         for (CompletedFetch completedFetch : completedFetches) {
             exclude.add(completedFetch.partition);
         }
-        fetchable.removeAll(exclude);
-        return fetchable;
+        return subscriptions.fetchablePartitions(tp -> !exclude.contains(tp));
+    }
+
+    /**
+     * Determine which replica to read from.
+     */
+    Node selectReadReplica(TopicPartition partition, Node leaderReplica, long currentTimeMs) {
+        Optional nodeId = subscriptions.preferredReadReplica(partition, currentTimeMs);
+        if (nodeId.isPresent()) {
+            Optional node = nodeId.flatMap(id -> metadata.fetch().nodeIfOnline(partition, id));
+            if (node.isPresent()) {
+                return node.get();
+            } else {
+                log.trace("Not fetching from {} for partition {} since it is marked offline or is missing from our metadata," +
+                          " using the leader instead.", nodeId, partition);
+                subscriptions.clearPreferredReadReplica(partition);
+                return leaderReplica;
+            }
+        } else {
+            return leaderReplica;
+        }
     }
 
     /**
@@ -876,12 +1103,19 @@ private List fetchablePartitions() {
      */
     private Map prepareFetchRequests() {
         Map fetchable = new LinkedHashMap<>();
+
+        // Ensure the position has an up-to-date leader
+        subscriptions.assignedPartitions().forEach(
+            tp -> subscriptions.maybeValidatePositionForCurrentLeader(tp, metadata.leaderAndEpoch(tp)));
+
+        long currentTimeMs = time.milliseconds();
+
         for (TopicPartition partition : fetchablePartitions()) {
-            Node node = metadata.partitionInfoIfCurrent(partition)
-                    .map(MetadataCache.PartitionInfoAndEpoch::partitionInfo)
-                    .map(PartitionInfo::leader)
-                    .orElse(null);
-            if (node == null) {
+            // Use the preferred read replica if set, or the position's leader
+            SubscriptionState.FetchPosition position = this.subscriptions.position(partition);
+            Node node = selectReadReplica(partition, position.currentLeader.leader, currentTimeMs);
+
+            if (node == null || node.isEmpty()) {
                 metadata.requestUpdate();
             } else if (client.isUnavailable(node)) {
                 client.maybeThrowAuthFailure(node);
@@ -889,30 +1123,31 @@ private Map prepareFetchRequests() {
                 // If we try to send during the reconnect blackout window, then the request is just
                 // going to be failed anyway before being sent, so skip the send for now
                 log.trace("Skipping fetch for partition {} because node {} is awaiting reconnect backoff", partition, node);
-            } else if (client.hasPendingRequests(node)) {
-                log.trace("Skipping fetch for partition {} because there is an in-flight request to {}", partition, node);
+
+            } else if (this.nodesWithPendingFetchRequests.contains(node.id())) {
+                log.trace("Skipping fetch for partition {} because previous request to {} has not been processed", partition, node);
             } else {
                 // if there is a leader and no in-flight requests, issue a new fetch
                 FetchSessionHandler.Builder builder = fetchable.get(node);
                 if (builder == null) {
-                    FetchSessionHandler handler = sessionHandler(node.id());
+                    int id = node.id();
+                    FetchSessionHandler handler = sessionHandler(id);
                     if (handler == null) {
-                        handler = new FetchSessionHandler(logContext, node.id());
-                        sessionHandlers.put(node.id(), handler);
+                        handler = new FetchSessionHandler(logContext, id);
+                        sessionHandlers.put(id, handler);
                     }
                     builder = handler.newBuilder();
                     fetchable.put(node, builder);
                 }
 
-                long position = this.subscriptions.position(partition);
-                Optional leaderEpoch = this.metadata.lastSeenLeaderEpoch(partition);
-                builder.add(partition, new FetchRequest.PartitionData(position, FetchRequest.INVALID_LOG_START_OFFSET,
-                    this.fetchSize, leaderEpoch));
+                builder.add(partition, new FetchRequest.PartitionData(position.offset,
+                        FetchRequest.INVALID_LOG_START_OFFSET, this.fetchSize, position.currentLeader.epoch));
 
-                log.debug("Added {} fetch request for partition {} at offset {} to node {}", isolationLevel,
+                log.debug("Added {} fetch request for partition {} at position {} to node {}", isolationLevel,
                     partition, position, node);
             }
         }
+
         Map reqs = new LinkedHashMap<>();
         for (Map.Entry entry : fetchable.entrySet()) {
             reqs.put(entry.getKey(), entry.getValue().build());
@@ -920,26 +1155,40 @@ private Map prepareFetchRequests() {
         return reqs;
     }
 
+    private Map> regroupFetchPositionsByLeader(
+            Map partitionMap) {
+        return partitionMap.entrySet()
+                .stream()
+                .collect(Collectors.groupingBy(entry -> entry.getValue().currentLeader.leader,
+                        Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)));
+    }
+
+    private  Map> regroupPartitionMapByNode(Map partitionMap) {
+        return partitionMap.entrySet()
+                .stream()
+                .collect(Collectors.groupingBy(entry -> metadata.fetch().leaderFor(entry.getKey()),
+                        Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)));
+    }
+
     /**
-     * The callback for fetch completion
+     * Initialize a CompletedFetch object.
      */
-    private PartitionRecords parseCompletedFetch(CompletedFetch completedFetch) {
-        TopicPartition tp = completedFetch.partition;
-        FetchResponse.PartitionData partition = completedFetch.partitionData;
-        long fetchOffset = completedFetch.fetchedOffset;
-        PartitionRecords partitionRecords = null;
+    private CompletedFetch initializeCompletedFetch(CompletedFetch nextCompletedFetch) {
+        TopicPartition tp = nextCompletedFetch.partition;
+        FetchResponse.PartitionData partition = nextCompletedFetch.partitionData;
+        long fetchOffset = nextCompletedFetch.nextFetchOffset;
+        CompletedFetch completedFetch = null;
         Errors error = partition.error;
 
         try {
-            if (!subscriptions.isFetchable(tp)) {
-                // this can happen when a rebalance happened or a partition consumption paused
-                // while fetch is still in-flight
-                log.debug("Ignoring fetched records for partition {} since it is no longer fetchable", tp);
+            if (!subscriptions.hasValidPosition(tp)) {
+                // this can happen when a rebalance happened while fetch is still in-flight
+                log.debug("Ignoring fetched records for partition {} since it no longer has valid position", tp);
             } else if (error == Errors.NONE) {
                 // we are interested in this fetch only if the beginning offset matches the
                 // current consumed position
-                Long position = subscriptions.position(tp);
-                if (position == null || position != fetchOffset) {
+                SubscriptionState.FetchPosition position = subscriptions.position(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);
                     return null;
@@ -948,7 +1197,7 @@ private PartitionRecords parseCompletedFetch(CompletedFetch completedFetch) {
                 log.trace("Preparing to read {} bytes of data for partition {} with offset {}",
                         partition.records.sizeInBytes(), tp, position);
                 Iterator batches = partition.records.batches().iterator();
-                partitionRecords = new PartitionRecords(tp, completedFetch, batches);
+                completedFetch = nextCompletedFetch;
 
                 if (!batches.hasNext() && partition.records.sizeInBytes() > 0) {
                     if (completedFetch.responseVersion < 3) {
@@ -982,24 +1231,43 @@ private PartitionRecords parseCompletedFetch(CompletedFetch completedFetch) {
                     log.trace("Updating last stable offset for partition {} to {}", tp, partition.lastStableOffset);
                     subscriptions.updateLastStableOffset(tp, partition.lastStableOffset);
                 }
+
+                if (partition.preferredReadReplica.isPresent()) {
+                    subscriptions.updatePreferredReadReplica(completedFetch.partition, partition.preferredReadReplica.get(), () -> {
+                        long expireTimeMs = time.milliseconds() + metadata.metadataExpireMs();
+                        log.debug("Updating preferred read replica for partition {} to {}, set to expire at {}",
+                                tp, partition.preferredReadReplica.get(), expireTimeMs);
+                        return expireTimeMs;
+                    });
+                }
+
+                nextCompletedFetch.initialized = true;
             } else if (error == Errors.NOT_LEADER_FOR_PARTITION ||
                        error == Errors.REPLICA_NOT_AVAILABLE ||
                        error == Errors.KAFKA_STORAGE_ERROR ||
-                       error == Errors.FENCED_LEADER_EPOCH) {
+                       error == Errors.FENCED_LEADER_EPOCH ||
+                       error == Errors.OFFSET_NOT_AVAILABLE) {
                 log.debug("Error in fetch for partition {}: {}", tp, error.exceptionName());
                 this.metadata.requestUpdate();
             } else if (error == Errors.UNKNOWN_TOPIC_OR_PARTITION) {
                 log.warn("Received unknown topic or partition error in fetch for partition {}", tp);
                 this.metadata.requestUpdate();
             } else if (error == Errors.OFFSET_OUT_OF_RANGE) {
-                if (fetchOffset != subscriptions.position(tp)) {
-                    log.debug("Discarding stale fetch response for partition {} since the fetched offset {} " +
-                            "does not match the current offset {}", tp, fetchOffset, subscriptions.position(tp));
-                } else if (subscriptions.hasDefaultOffsetResetPolicy()) {
-                    log.info("Fetch offset {} is out of range for partition {}, resetting offset", fetchOffset, tp);
-                    subscriptions.requestOffsetReset(tp);
+                Optional clearedReplicaId = subscriptions.clearPreferredReadReplica(tp);
+                if (!clearedReplicaId.isPresent()) {
+                    // If there's no preferred replica to clear, we're fetching from the leader so handle this error normally
+                    if (fetchOffset != subscriptions.position(tp).offset) {
+                        log.debug("Discarding stale fetch response for partition {} since the fetched offset {} " +
+                                "does not match the current offset {}", tp, fetchOffset, subscriptions.position(tp));
+                    } else if (subscriptions.hasDefaultOffsetResetPolicy()) {
+                        log.info("Fetch offset {} is out of range for partition {}, resetting offset", fetchOffset, tp);
+                        subscriptions.requestOffsetReset(tp);
+                    } else {
+                        throw new OffsetOutOfRangeException(Collections.singletonMap(tp, fetchOffset));
+                    }
                 } else {
-                    throw new OffsetOutOfRangeException(Collections.singletonMap(tp, fetchOffset));
+                    log.debug("Unset the preferred read replica {} for partition {} since we got {} when fetching {}",
+                            clearedReplicaId.get(), tp, error, fetchOffset);
                 }
             } else if (error == Errors.TOPIC_AUTHORIZATION_FAILED) {
                 //we log the actual partition and not just the topic to help with ACL propagation issues in large clusters
@@ -1013,8 +1281,8 @@ private PartitionRecords parseCompletedFetch(CompletedFetch completedFetch) {
                 throw new IllegalStateException("Unexpected error code " + error.code() + " while fetching from partition " + tp);
             }
         } finally {
-            if (partitionRecords == null)
-                completedFetch.metricAggregator.record(tp, 0, 0);
+            if (completedFetch == null)
+                nextCompletedFetch.metricAggregator.record(tp, 0, 0);
 
             if (error != Errors.NONE)
                 // we move the partition to the end if there was an error. This way, it's more likely that partitions for
@@ -1022,7 +1290,7 @@ private PartitionRecords parseCompletedFetch(CompletedFetch completedFetch) {
                 subscriptions.movePartitionToEnd(tp);
         }
 
-        return partitionRecords;
+        return completedFetch;
     }
 
     /**
@@ -1040,7 +1308,8 @@ private ConsumerRecord parseRecord(TopicPartition partition,
             ByteBuffer keyBytes = record.key();
             byte[] keyByteArray = keyBytes == null ? null : Utils.toArray(keyBytes);
             K key = keyBytes == null ? null : this.keyDeserializer.deserialize(partition.topic(), headers, keyByteArray);
-            ByteBuffer valueBytes = record.value();
+            ByteBuffer valueBytes = record.hasMagic(MAGIC_VALUE_V2) ? record.value() :
+                (enableShallowIteration ? ((AbstractLegacyRecordBatch) record).outerRecord().buffer() : record.value());
             byte[] valueByteArray = valueBytes == null ? null : Utils.toArray(valueBytes);
             V value = valueBytes == null ? null : this.valueDeserializer.deserialize(partition.topic(), headers, valueByteArray);
             return new ConsumerRecord<>(partition.topic(), partition.partition(), offset,
@@ -1058,27 +1327,25 @@ private Optional maybeLeaderEpoch(int leaderEpoch) {
         return leaderEpoch == RecordBatch.NO_PARTITION_LEADER_EPOCH ? Optional.empty() : Optional.of(leaderEpoch);
     }
 
-    @Override
-    public void onAssignment(Set assignment) {
-        sensors.updatePartitionLagAndLeadSensors(assignment);
-    }
-
     /**
      * Clear the buffered data which are not a part of newly assigned partitions
      *
      * @param assignedPartitions  newly assigned {@link TopicPartition}
      */
     public void clearBufferedDataForUnassignedPartitions(Collection assignedPartitions) {
-        Iterator itr = completedFetches.iterator();
-        while (itr.hasNext()) {
-            TopicPartition tp = itr.next().partition;
+        Iterator completedFetchesItr = completedFetches.iterator();
+        while (completedFetchesItr.hasNext()) {
+            CompletedFetch records = completedFetchesItr.next();
+            TopicPartition tp = records.partition;
             if (!assignedPartitions.contains(tp)) {
-                itr.remove();
+                records.drain();
+                completedFetchesItr.remove();
             }
         }
-        if (nextInLineRecords != null && !assignedPartitions.contains(nextInLineRecords.partition)) {
-            nextInLineRecords.drain();
-            nextInLineRecords = null;
+
+        if (nextInLineFetch != null && !assignedPartitions.contains(nextInLineFetch.partition)) {
+            nextInLineFetch.drain();
+            nextInLineFetch = null;
         }
     }
 
@@ -1111,12 +1378,14 @@ public static Sensor throttleTimeSensor(Metrics metrics, FetcherMetricsRegistry
         return fetchThrottleTimeSensor;
     }
 
-    private class PartitionRecords {
+    private class CompletedFetch {
         private final TopicPartition partition;
-        private final CompletedFetch completedFetch;
         private final Iterator batches;
         private final Set abortedProducerIds;
         private final PriorityQueue abortedTransactions;
+        private final FetchResponse.PartitionData partitionData;
+        private final FetchResponseMetricAggregator metricAggregator;
+        private final short responseVersion;
 
         private int recordsRead;
         private int bytesRead;
@@ -1124,27 +1393,35 @@ private class PartitionRecords {
         private Record lastRecord;
         private CloseableIterator records;
         private long nextFetchOffset;
-        private boolean isFetched = false;
+        private Optional lastEpoch;
+        private boolean isConsumed = false;
         private Exception cachedRecordException = null;
         private boolean corruptLastRecord = false;
+        private boolean initialized = false;
 
-        private PartitionRecords(TopicPartition partition,
-                                 CompletedFetch completedFetch,
-                                 Iterator batches) {
+        private CompletedFetch(TopicPartition partition,
+                               FetchResponse.PartitionData partitionData,
+                               FetchResponseMetricAggregator metricAggregator,
+                               Iterator batches,
+                               Long fetchOffset,
+                               short responseVersion) {
             this.partition = partition;
-            this.completedFetch = completedFetch;
+            this.partitionData = partitionData;
+            this.metricAggregator = metricAggregator;
             this.batches = batches;
-            this.nextFetchOffset = completedFetch.fetchedOffset;
+            this.nextFetchOffset = fetchOffset;
+            this.responseVersion = responseVersion;
+            this.lastEpoch = Optional.empty();
             this.abortedProducerIds = new HashSet<>();
-            this.abortedTransactions = abortedTransactions(completedFetch.partitionData);
+            this.abortedTransactions = abortedTransactions(partitionData);
         }
 
         private void drain() {
-            if (!isFetched) {
+            if (!isConsumed) {
                 maybeCloseRecordStream();
                 cachedRecordException = null;
-                this.isFetched = true;
-                this.completedFetch.metricAggregator.record(partition, bytesRead, recordsRead);
+                this.isConsumed = true;
+                this.metricAggregator.record(partition, bytesRead, recordsRead);
 
                 // we move the partition to the end if we received some bytes. This way, it's more likely that partitions
                 // for the same topic can remain together (allowing for more efficient serialization).
@@ -1157,7 +1434,7 @@ private void maybeEnsureValid(RecordBatch batch) {
             if (checkCrcs && currentBatch.magic() >= RecordBatch.MAGIC_VALUE_V2) {
                 try {
                     batch.ensureValid();
-                } catch (InvalidRecordException e) {
+                } catch (CorruptRecordException e) {
                     throw new KafkaException("Record batch for partition " + partition + " at offset " +
                             batch.baseOffset() + " is invalid, cause: " + e.getMessage());
                 }
@@ -1168,7 +1445,7 @@ private void maybeEnsureValid(Record record) {
             if (checkCrcs) {
                 try {
                     record.ensureValid();
-                } catch (InvalidRecordException e) {
+                } catch (CorruptRecordException e) {
                     throw new KafkaException("Record for partition " + partition + " at offset " + record.offset()
                             + " is invalid, cause: " + e.getMessage());
                 }
@@ -1200,6 +1477,9 @@ private Record nextFetchedRecord() {
                     }
 
                     currentBatch = batches.next();
+                    lastEpoch = currentBatch.partitionLeaderEpoch() == RecordBatch.NO_PARTITION_LEADER_EPOCH ?
+                            Optional.empty() : Optional.of(currentBatch.partitionLeaderEpoch());
+
                     maybeEnsureValid(currentBatch);
 
                     if (isolationLevel == IsolationLevel.READ_COMMITTED && currentBatch.hasProducerId()) {
@@ -1220,7 +1500,11 @@ private Record nextFetchedRecord() {
                         }
                     }
 
-                    records = currentBatch.streamingIterator(decompressionBufferSupplier);
+                    if (enableShallowIteration) {
+                        records = currentBatch.shallowIterator();
+                    } else {
+                        records = currentBatch.streamingIterator(decompressionBufferSupplier);
+                    }
                 } else {
                     Record record = records.next();
                     // skip any records out of range
@@ -1247,7 +1531,7 @@ private List> fetchRecords(int maxRecords) {
                                              + ". If needed, please seek past the record to "
                                              + "continue consumption.", cachedRecordException);
 
-            if (isFetched)
+            if (isConsumed)
                 return Collections.emptyList();
 
             List> records = new ArrayList<>();
@@ -1320,25 +1604,9 @@ private boolean containsAbortMarker(RecordBatch batch) {
             Record firstRecord = batchIterator.next();
             return ControlRecordType.ABORT == ControlRecordType.parse(firstRecord.key());
         }
-    }
 
-    private static class CompletedFetch {
-        private final TopicPartition partition;
-        private final long fetchedOffset;
-        private final FetchResponse.PartitionData partitionData;
-        private final FetchResponseMetricAggregator metricAggregator;
-        private final short responseVersion;
-
-        private CompletedFetch(TopicPartition partition,
-                               long fetchedOffset,
-                               FetchResponse.PartitionData partitionData,
-                               FetchResponseMetricAggregator metricAggregator,
-                               short responseVersion) {
-            this.partition = partition;
-            this.fetchedOffset = fetchedOffset;
-            this.partitionData = partitionData;
-            this.metricAggregator = metricAggregator;
-            this.responseVersion = responseVersion;
+        private boolean notInitialized() {
+            return !this.initialized;
         }
     }
 
@@ -1410,7 +1678,8 @@ private static class FetchManagerMetrics {
         private final Sensor recordsFetchLag;
         private final Sensor recordsFetchLead;
 
-        private Set assignedPartitions;
+        private int assignmentId = 0;
+        private Set assignedPartitions = Collections.emptySet();
 
         private FetchManagerMetrics(Metrics metrics, FetcherMetricsRegistry metricsRegistry) {
             this.metrics = metrics;
@@ -1430,7 +1699,7 @@ private FetchManagerMetrics(Metrics metrics, FetcherMetricsRegistry metricsRegis
             this.fetchLatency = metrics.sensor("fetch-latency");
             this.fetchLatency.add(metrics.metricInstance(metricsRegistry.fetchLatencyAvg), new Avg());
             this.fetchLatency.add(metrics.metricInstance(metricsRegistry.fetchLatencyMax), new Max());
-            this.fetchLatency.add(new Meter(new Count(), metrics.metricInstance(metricsRegistry.fetchRequestRate),
+            this.fetchLatency.add(new Meter(new WindowedCount(), metrics.metricInstance(metricsRegistry.fetchRequestRate),
                     metrics.metricInstance(metricsRegistry.fetchRequestTotal)));
 
             this.recordsFetchLag = metrics.sensor("records-lag");
@@ -1473,16 +1742,33 @@ private void recordTopicFetchMetrics(String topic, int bytes, int records) {
             recordsFetched.record(records);
         }
 
-        private void updatePartitionLagAndLeadSensors(Set assignedPartitions) {
-            if (this.assignedPartitions != null) {
+        private void maybeUpdateAssignment(SubscriptionState subscription) {
+            int newAssignmentId = subscription.assignmentId();
+            if (this.assignmentId != newAssignmentId) {
+                Set newAssignedPartitions = subscription.assignedPartitions();
                 for (TopicPartition tp : this.assignedPartitions) {
-                    if (!assignedPartitions.contains(tp)) {
+                    if (!newAssignedPartitions.contains(tp)) {
                         metrics.removeSensor(partitionLagMetricName(tp));
                         metrics.removeSensor(partitionLeadMetricName(tp));
+                        metrics.removeMetric(partitionPreferredReadReplicaMetricName(tp));
+                    }
+                }
+
+                for (TopicPartition tp : newAssignedPartitions) {
+                    if (!this.assignedPartitions.contains(tp)) {
+                        MetricName metricName = partitionPreferredReadReplicaMetricName(tp);
+                        if (metrics.metric(metricName) == null) {
+                            metrics.addMetric(
+                                metricName,
+                                (Gauge) (config, now) -> subscription.preferredReadReplica(tp, 0L).orElse(-1)
+                            );
+                        }
                     }
                 }
+
+                this.assignedPartitions = newAssignedPartitions;
+                this.assignmentId = newAssignmentId;
             }
-            this.assignedPartitions = assignedPartitions;
         }
 
         private void recordPartitionLead(TopicPartition tp, long lead) {
@@ -1491,9 +1777,7 @@ private void recordPartitionLead(TopicPartition tp, long lead) {
             String name = partitionLeadMetricName(tp);
             Sensor recordsLead = this.metrics.getSensor(name);
             if (recordsLead == null) {
-                Map metricTags = new HashMap<>(2);
-                metricTags.put("topic", tp.topic().replace('.', '_'));
-                metricTags.put("partition", String.valueOf(tp.partition()));
+                Map metricTags = topicPartitionTags(tp);
 
                 recordsLead = this.metrics.sensor(name);
 
@@ -1510,10 +1794,7 @@ private void recordPartitionLag(TopicPartition tp, long lag) {
             String name = partitionLagMetricName(tp);
             Sensor recordsLag = this.metrics.getSensor(name);
             if (recordsLag == null) {
-                Map metricTags = new HashMap<>(2);
-                metricTags.put("topic", tp.topic().replace('.', '_'));
-                metricTags.put("partition", String.valueOf(tp.partition()));
-
+                Map metricTags = topicPartitionTags(tp);
                 recordsLag = this.metrics.sensor(name);
 
                 recordsLag.add(this.metrics.metricInstance(metricsRegistry.partitionRecordsLag, metricTags), new Value());
@@ -1531,12 +1812,23 @@ private static String partitionLeadMetricName(TopicPartition tp) {
             return tp + ".records-lead";
         }
 
+        private MetricName partitionPreferredReadReplicaMetricName(TopicPartition tp) {
+            Map metricTags = topicPartitionTags(tp);
+            return this.metrics.metricInstance(metricsRegistry.partitionPreferredReadReplica, metricTags);
+        }
+
+        private Map topicPartitionTags(TopicPartition tp) {
+            Map metricTags = new HashMap<>(2);
+            metricTags.put("topic", tp.topic().replace('.', '_'));
+            metricTags.put("partition", String.valueOf(tp.partition()));
+            return metricTags;
+        }
     }
 
     @Override
     public void close() {
-        if (nextInLineRecords != null)
-            nextInLineRecords.drain();
+        if (nextInLineFetch != null)
+            nextInLineFetch.drain();
         decompressionBufferSupplier.close();
     }
 
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/FetcherMetricsRegistry.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/FetcherMetricsRegistry.java
index f86961545cc99..501ffe9a88da8 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/FetcherMetricsRegistry.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/FetcherMetricsRegistry.java
@@ -54,6 +54,7 @@ public class FetcherMetricsRegistry {
     public MetricNameTemplate partitionRecordsLead;
     public MetricNameTemplate partitionRecordsLeadMin;
     public MetricNameTemplate partitionRecordsLeadAvg;
+    public MetricNameTemplate partitionPreferredReadReplica;
 
     public FetcherMetricsRegistry() {
         this(new HashSet(), "");
@@ -139,7 +140,9 @@ public FetcherMetricsRegistry(Set tags, String metricGrpPrefix) {
                 "The min lead of the partition", partitionTags);
         this.partitionRecordsLeadAvg = new MetricNameTemplate("records-lead-avg", groupName,
                 "The average lead of the partition", partitionTags);
-
+        this.partitionPreferredReadReplica = new MetricNameTemplate(
+                "preferred-read-replica", "consumer-fetch-manager-metrics",
+                "The current read replica for the partition, or -1 if reading from leader", partitionTags);
     }
 
     public List getAllTemplates() {
@@ -171,7 +174,8 @@ public List getAllTemplates() {
             partitionRecordsLagMax,
             partitionRecordsLead,
             partitionRecordsLeadMin,
-            partitionRecordsLeadAvg
+            partitionRecordsLeadAvg,
+            partitionPreferredReadReplica
         );
     }
 
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/Heartbeat.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/Heartbeat.java
index 8a67f310d87f0..0697d62cf6d14 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/Heartbeat.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/Heartbeat.java
@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.clients.consumer.internals;
 
+import org.apache.kafka.clients.GroupRebalanceConfig;
 import org.apache.kafka.common.utils.Time;
 import org.apache.kafka.common.utils.Timer;
 
@@ -23,36 +24,30 @@
  * A helper class for managing the heartbeat to the coordinator
  */
 public final class Heartbeat {
-    private final int sessionTimeoutMs;
-    private final int heartbeatIntervalMs;
     private final int maxPollIntervalMs;
-    private final long retryBackoffMs;
+    private final GroupRebalanceConfig rebalanceConfig;
     private final Time time;
     private final Timer heartbeatTimer;
     private final Timer sessionTimer;
     private final Timer pollTimer;
 
-    private volatile long lastHeartbeatSend;
+    private volatile long lastHeartbeatSend = 0L;
+    private volatile long lastHeartbeatReceive;
 
-    public Heartbeat(Time time,
-                     int sessionTimeoutMs,
-                     int heartbeatIntervalMs,
-                     int maxPollIntervalMs,
-                     long retryBackoffMs) {
-        if (heartbeatIntervalMs >= sessionTimeoutMs)
+    public Heartbeat(GroupRebalanceConfig config,
+                     Time time) {
+        if (config.heartbeatIntervalMs >= config.sessionTimeoutMs)
             throw new IllegalArgumentException("Heartbeat must be set lower than the session timeout");
-
+        this.rebalanceConfig = config;
         this.time = time;
-        this.sessionTimeoutMs = sessionTimeoutMs;
-        this.heartbeatIntervalMs = heartbeatIntervalMs;
-        this.maxPollIntervalMs = maxPollIntervalMs;
-        this.retryBackoffMs = retryBackoffMs;
-        this.heartbeatTimer = time.timer(heartbeatIntervalMs);
-        this.sessionTimer = time.timer(sessionTimeoutMs);
+        this.heartbeatTimer = time.timer(config.heartbeatIntervalMs);
+        this.sessionTimer = time.timer(config.sessionTimeoutMs);
+        this.maxPollIntervalMs = config.rebalanceTimeoutMs;
         this.pollTimer = time.timer(maxPollIntervalMs);
     }
 
     private void update(long now) {
+        lastHeartbeatReceive = now;
         heartbeatTimer.update(now);
         sessionTimer.update(now);
         pollTimer.update(now);
@@ -66,28 +61,32 @@ public void poll(long now) {
     public void sentHeartbeat(long now) {
         this.lastHeartbeatSend = now;
         update(now);
-        heartbeatTimer.reset(heartbeatIntervalMs);
+        heartbeatTimer.reset(rebalanceConfig.heartbeatIntervalMs);
     }
 
     public void failHeartbeat() {
         update(time.milliseconds());
-        heartbeatTimer.reset(retryBackoffMs);
+        heartbeatTimer.reset(rebalanceConfig.retryBackoffMs);
     }
 
     public void receiveHeartbeat() {
         update(time.milliseconds());
-        sessionTimer.reset(sessionTimeoutMs);
+        sessionTimer.reset(rebalanceConfig.sessionTimeoutMs);
     }
 
     public boolean shouldHeartbeat(long now) {
         update(now);
         return heartbeatTimer.isExpired();
     }
-    
+
     public long lastHeartbeatSend() {
         return this.lastHeartbeatSend;
     }
 
+    public long lastHeartbeatReceive() {
+        return lastHeartbeatReceive;
+    }
+
     public long timeToNextHeartbeat(long now) {
         update(now);
         return heartbeatTimer.remainingMs();
@@ -100,14 +99,14 @@ public boolean sessionTimeoutExpired(long now) {
 
     public void resetTimeouts() {
         update(time.milliseconds());
-        sessionTimer.reset(sessionTimeoutMs);
+        sessionTimer.reset(rebalanceConfig.sessionTimeoutMs);
         pollTimer.reset(maxPollIntervalMs);
-        heartbeatTimer.reset(heartbeatIntervalMs);
+        heartbeatTimer.reset(rebalanceConfig.heartbeatIntervalMs);
     }
 
     public void resetSessionTimeout() {
         update(time.milliseconds());
-        sessionTimer.reset(sessionTimeoutMs);
+        sessionTimer.reset(rebalanceConfig.sessionTimeoutMs);
     }
 
     public boolean pollTimeoutExpired(long now) {
@@ -119,4 +118,4 @@ public long lastPollTime() {
         return pollTimer.currentTimeMs();
     }
 
-}
\ No newline at end of file
+}
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/KafkaConsumerMetrics.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/KafkaConsumerMetrics.java
new file mode 100644
index 0000000000000..aab313320d4bd
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/KafkaConsumerMetrics.java
@@ -0,0 +1,88 @@
+/*
+ * 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.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.Max;
+
+import java.util.concurrent.TimeUnit;
+
+public class KafkaConsumerMetrics implements AutoCloseable {
+    private final Metrics metrics;
+    private final MetricName lastPollMetricName;
+    private final Sensor timeBetweenPollSensor;
+    private final Sensor pollIdleSensor;
+    private long lastPollMs;
+    private long pollStartMs;
+    private long timeSinceLastPollMs;
+
+    public KafkaConsumerMetrics(Metrics metrics, String metricGrpPrefix) {
+        this.metrics = metrics;
+
+        String metricGroupName = metricGrpPrefix + "-metrics";
+        Measurable lastPoll = (mConfig, now) -> {
+            if (lastPollMs == 0L)
+                // if no poll is ever triggered, just return -1.
+                return -1d;
+            else
+                return TimeUnit.SECONDS.convert(now - lastPollMs, TimeUnit.MILLISECONDS);
+        };
+        this.lastPollMetricName = metrics.metricName("last-poll-seconds-ago",
+                metricGroupName, "The number of seconds since the last poll() invocation.");
+        metrics.addMetric(lastPollMetricName, lastPoll);
+
+        this.timeBetweenPollSensor = metrics.sensor("time-between-poll");
+        this.timeBetweenPollSensor.add(metrics.metricName("time-between-poll-avg",
+                metricGroupName,
+                "The average delay between invocations of poll()."),
+                new Avg());
+        this.timeBetweenPollSensor.add(metrics.metricName("time-between-poll-max",
+                metricGroupName,
+                "The max delay between invocations of poll()."),
+                new Max());
+
+        this.pollIdleSensor = metrics.sensor("poll-idle-ratio-avg");
+        this.pollIdleSensor.add(metrics.metricName("poll-idle-ratio-avg",
+                metricGroupName,
+                "The average fraction of time the consumer's poll() is idle as opposed to waiting for the user code to process records."),
+                new Avg());
+    }
+
+    public void recordPollStart(long pollStartMs) {
+        this.pollStartMs = pollStartMs;
+        this.timeSinceLastPollMs = lastPollMs != 0L ? pollStartMs - lastPollMs : 0;
+        this.timeBetweenPollSensor.record(timeSinceLastPollMs);
+        this.lastPollMs = pollStartMs;
+    }
+
+    public void recordPollEnd(long pollEndMs) {
+        long pollTimeMs = pollEndMs - pollStartMs;
+        double pollIdleRatio = pollTimeMs * 1.0 / (pollTimeMs + timeSinceLastPollMs);
+        this.pollIdleSensor.record(pollIdleRatio);
+    }
+
+    @Override
+    public void close() {
+        metrics.removeMetric(lastPollMetricName);
+        metrics.removeSensor(timeBetweenPollSensor.name());
+        metrics.removeSensor(pollIdleSensor.name());
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/OffsetsForLeaderEpochClient.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/OffsetsForLeaderEpochClient.java
new file mode 100644
index 0000000000000..9556d3c64dc1b
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/OffsetsForLeaderEpochClient.java
@@ -0,0 +1,127 @@
+/*
+ * 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.common.Node;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.errors.TopicAuthorizationException;
+import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.common.requests.AbstractRequest;
+import org.apache.kafka.common.requests.EpochEndOffset;
+import org.apache.kafka.common.requests.OffsetsForLeaderEpochRequest;
+import org.apache.kafka.common.requests.OffsetsForLeaderEpochResponse;
+import org.apache.kafka.common.utils.LogContext;
+
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+
+/**
+ * Convenience class for making asynchronous requests to the OffsetsForLeaderEpoch API
+ */
+public class OffsetsForLeaderEpochClient extends AsyncClient<
+        Map,
+        OffsetsForLeaderEpochRequest,
+        OffsetsForLeaderEpochResponse,
+        OffsetsForLeaderEpochClient.OffsetForEpochResult> {
+
+    OffsetsForLeaderEpochClient(ConsumerNetworkClient client, LogContext logContext) {
+        super(client, logContext);
+    }
+
+    @Override
+    protected AbstractRequest.Builder prepareRequest(
+            Node node, Map requestData) {
+        Map partitionData = new HashMap<>(requestData.size());
+        requestData.forEach((topicPartition, fetchPosition) -> fetchPosition.offsetEpoch.ifPresent(
+            fetchEpoch -> partitionData.put(topicPartition,
+                new OffsetsForLeaderEpochRequest.PartitionData(fetchPosition.currentLeader.epoch, fetchEpoch))));
+
+        return OffsetsForLeaderEpochRequest.Builder.forConsumer(partitionData);
+    }
+
+    @Override
+    protected OffsetForEpochResult handleResponse(
+            Node node,
+            Map requestData,
+            OffsetsForLeaderEpochResponse response) {
+
+        Set partitionsToRetry = new HashSet<>();
+        Set unauthorizedTopics = new HashSet<>();
+        Map endOffsets = new HashMap<>();
+
+        for (TopicPartition topicPartition : requestData.keySet()) {
+            EpochEndOffset epochEndOffset = response.responses().get(topicPartition);
+            if (epochEndOffset == null) {
+                logger().warn("Missing partition {} from response, ignoring", topicPartition);
+                partitionsToRetry.add(topicPartition);
+                continue;
+            }
+            Errors error = epochEndOffset.error();
+            if (error == Errors.NONE) {
+                logger().debug("Handling OffsetsForLeaderEpoch response for {}. Got offset {} for epoch {}",
+                        topicPartition, epochEndOffset.endOffset(), epochEndOffset.leaderEpoch());
+                endOffsets.put(topicPartition, epochEndOffset);
+            } else if (error == Errors.NOT_LEADER_FOR_PARTITION ||
+                    error == Errors.REPLICA_NOT_AVAILABLE ||
+                    error == Errors.KAFKA_STORAGE_ERROR ||
+                    error == Errors.OFFSET_NOT_AVAILABLE ||
+                    error == Errors.LEADER_NOT_AVAILABLE) {
+                logger().debug("Attempt to fetch offsets for partition {} failed due to {}, retrying.",
+                        topicPartition, error);
+                partitionsToRetry.add(topicPartition);
+            } else if (error == Errors.FENCED_LEADER_EPOCH ||
+                    error == Errors.UNKNOWN_LEADER_EPOCH) {
+                logger().debug("Attempt to fetch offsets for partition {} failed due to {}, retrying.",
+                        topicPartition, error);
+                partitionsToRetry.add(topicPartition);
+            } else if (error == Errors.UNKNOWN_TOPIC_OR_PARTITION) {
+                logger().warn("Received unknown topic or partition error in ListOffset request for partition {}", topicPartition);
+                partitionsToRetry.add(topicPartition);
+            } else if (error == Errors.TOPIC_AUTHORIZATION_FAILED) {
+                unauthorizedTopics.add(topicPartition.topic());
+            } else {
+                logger().warn("Attempt to fetch offsets for partition {} failed due to: {}, retrying.", topicPartition, error.message());
+                partitionsToRetry.add(topicPartition);
+            }
+        }
+
+        if (!unauthorizedTopics.isEmpty())
+            throw new TopicAuthorizationException(unauthorizedTopics);
+        else
+            return new OffsetForEpochResult(endOffsets, partitionsToRetry);
+    }
+
+    public static class OffsetForEpochResult {
+        private final Map endOffsets;
+        private final Set partitionsToRetry;
+
+        OffsetForEpochResult(Map endOffsets, Set partitionsNeedingRetry) {
+            this.endOffsets = endOffsets;
+            this.partitionsToRetry = partitionsNeedingRetry;
+        }
+
+        public Map endOffsets() {
+            return endOffsets;
+        }
+
+        public Set partitionsToRetry() {
+            return partitionsToRetry;
+        }
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/PartitionAssignor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/PartitionAssignor.java
index 4a7c7a8bbd341..43bd7519169d4 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/PartitionAssignor.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/PartitionAssignor.java
@@ -36,7 +36,14 @@
  * assignment decisions. For this, you can override {@link #subscription(Set)} and provide custom
  * userData in the returned Subscription. For example, to have a rack-aware assignor, an implementation
  * can use this user data to forward the rackId belonging to each member.
+ *
+ * This interface has been deprecated in 2.4, custom assignors should now implement
+ * {@link org.apache.kafka.clients.consumer.ConsumerPartitionAssignor}. Note that maintaining compatibility
+ * for an internal interface here is a special case, as {@code PartitionAssignor} was meant to be a public API
+ * although it was placed in the internals package. Users should not expect internal interfaces or classes to
+ * not be removed or maintain compatibility in any way.
  */
+@Deprecated
 public interface PartitionAssignor {
 
     /**
@@ -58,13 +65,20 @@ public interface PartitionAssignor {
      */
     Map assign(Cluster metadata, Map subscriptions);
 
-
     /**
      * Callback which is invoked when a group member receives its assignment from the leader.
      * @param assignment The local member's assignment as provided by the leader in {@link #assign(Cluster, Map)}
      */
     void onAssignment(Assignment assignment);
 
+    /**
+     * Callback which is invoked when a group member receives its assignment from the leader.
+     * @param assignment The local member's assignment as provided by the leader in {@link #assign(Cluster, Map)}
+     * @param generation The consumer group generation associated with this partition assignment (optional)
+     */
+    default void onAssignment(Assignment assignment, int generation) {
+        onAssignment(assignment);
+    }
 
     /**
      * Unique name for this assignor (e.g. "range" or "roundrobin" or "sticky")
@@ -96,8 +110,8 @@ public ByteBuffer userData() {
         @Override
         public String toString() {
             return "Subscription(" +
-                    "topics=" + topics +
-                    ')';
+                "topics=" + topics +
+                ')';
         }
     }
 
@@ -125,8 +139,8 @@ public ByteBuffer userData() {
         @Override
         public String toString() {
             return "Assignment(" +
-                    "partitions=" + partitions +
-                    ')';
+                "partitions=" + partitions +
+                ')';
         }
     }
 
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/PartitionAssignorAdapter.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/PartitionAssignorAdapter.java
new file mode 100644
index 0000000000000..8fb791ab4faa8
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/PartitionAssignorAdapter.java
@@ -0,0 +1,140 @@
+/*
+ * 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 java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import org.apache.kafka.clients.consumer.ConsumerGroupMetadata;
+import org.apache.kafka.clients.consumer.ConsumerPartitionAssignor;
+import org.apache.kafka.common.Cluster;
+import org.apache.kafka.common.Configurable;
+import org.apache.kafka.common.KafkaException;
+import org.apache.kafka.common.utils.Utils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * This adapter class is used to ensure backwards compatibility for those who have implemented the {@link PartitionAssignor}
+ * interface, which has been deprecated in favor of the new {@link org.apache.kafka.clients.consumer.ConsumerPartitionAssignor}.
+ * 
    
+ * Note that maintaining compatibility for an internal interface here is a special case, as {@code PartitionAssignor}
+ * was meant to be a public API although it was placed in the internals package. Users should not expect internal
+ * interfaces or classes to not be removed or maintain compatibility in any way.
+ */
+@SuppressWarnings("deprecation")
+public class PartitionAssignorAdapter implements ConsumerPartitionAssignor {
+
+    private static final Logger LOG = LoggerFactory.getLogger(PartitionAssignorAdapter.class);
+    private final PartitionAssignor oldAssignor;
+
+    PartitionAssignorAdapter(PartitionAssignor oldAssignor) {
+        this.oldAssignor = oldAssignor;
+    }
+
+    @Override
+    public ByteBuffer subscriptionUserData(Set topics) {
+        return oldAssignor.subscription(topics).userData();
+    }
+
+    @Override
+    public GroupAssignment assign(Cluster metadata, GroupSubscription groupSubscription) {
+        return toNewGroupAssignment(oldAssignor.assign(metadata, toOldGroupSubscription(groupSubscription)));
+    }
+
+    @Override
+    public void onAssignment(Assignment assignment, ConsumerGroupMetadata metadata) {
+        oldAssignor.onAssignment(toOldAssignment(assignment), metadata.generationId());
+    }
+
+    @Override
+    public String name() {
+        return oldAssignor.name();
+    }
+
+    private static PartitionAssignor.Assignment toOldAssignment(Assignment newAssignment) {
+        return new PartitionAssignor.Assignment(newAssignment.partitions(), newAssignment.userData());
+    }
+
+    private static Map toOldGroupSubscription(GroupSubscription newSubscriptions) {
+        Map oldSubscriptions = new HashMap<>();
+        for (Map.Entry entry : newSubscriptions.groupSubscription().entrySet()) {
+            String member = entry.getKey();
+            Subscription newSubscription = entry.getValue();
+            oldSubscriptions.put(member, new PartitionAssignor.Subscription(
+                newSubscription.topics(), newSubscription.userData()));
+        }
+        return oldSubscriptions;
+    }
+
+    private static GroupAssignment toNewGroupAssignment(Map oldAssignments) {
+        Map newAssignments = new HashMap<>();
+        for (Map.Entry entry : oldAssignments.entrySet()) {
+            String member = entry.getKey();
+            PartitionAssignor.Assignment oldAssignment = entry.getValue();
+            newAssignments.put(member, new Assignment(oldAssignment.partitions(), oldAssignment.userData()));
+        }
+        return new GroupAssignment(newAssignments);
+    }
+
+    /**
+     * Get a list of configured instances of {@link org.apache.kafka.clients.consumer.ConsumerPartitionAssignor}
+     * based on the class names/types specified by {@link org.apache.kafka.clients.consumer.ConsumerConfig#PARTITION_ASSIGNMENT_STRATEGY_CONFIG}
+     * where any instances of the old {@link PartitionAssignor} interface are wrapped in an adapter to the new
+     * {@link org.apache.kafka.clients.consumer.ConsumerPartitionAssignor} interface
+     */
+    public static List getAssignorInstances(List assignorClasses, Map configs) {
+        List assignors = new ArrayList<>();
+
+        if (assignorClasses == null)
+            return assignors;
+
+        for (Object klass : assignorClasses) {
+            // first try to get the class if passed in as a string
+            if (klass instanceof String) {
+                try {
+                    klass = Class.forName((String) klass, true, Utils.getContextOrKafkaClassLoader());
+                } catch (ClassNotFoundException classNotFound) {
+                    throw new KafkaException(klass + " ClassNotFoundException exception occurred", classNotFound);
+                }
+            }
+
+            if (klass instanceof Class) {
+                Object assignor = Utils.newInstance((Class) klass);
+                if (assignor instanceof Configurable)
+                    ((Configurable) assignor).configure(configs);
+
+                if (assignor instanceof ConsumerPartitionAssignor) {
+                    assignors.add((ConsumerPartitionAssignor) assignor);
+                } else if (assignor instanceof PartitionAssignor) {
+                    assignors.add(new PartitionAssignorAdapter((PartitionAssignor) assignor));
+                    LOG.warn("The PartitionAssignor interface has been deprecated, "
+                        + "please implement the ConsumerPartitionAssignor interface instead.");
+                } else {
+                    throw new KafkaException(klass + " is not an instance of " + PartitionAssignor.class.getName()
+                        + " or an instance of " + ConsumerPartitionAssignor.class.getName());
+                }
+            } else {
+                throw new KafkaException("List contains element of type " + klass.getClass().getName() + ", expected String or Class");
+            }
+        }
+        return assignors;
+    }
+}
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 fe944c5642e8b..1102acafe0d64 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
@@ -16,26 +16,32 @@
  */
 package org.apache.kafka.clients.consumer.internals;
 
+import java.util.ArrayList;
+import org.apache.kafka.clients.Metadata;
 import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
 import org.apache.kafka.clients.consumer.NoOffsetForPartitionException;
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
 import org.apache.kafka.clients.consumer.OffsetResetStrategy;
+import org.apache.kafka.common.Node;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.internals.PartitionStates;
+import org.apache.kafka.common.requests.EpochEndOffset;
 import org.apache.kafka.common.requests.IsolationLevel;
 import org.apache.kafka.common.utils.LogContext;
 import org.slf4j.Logger;
 
-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.List;
 import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
 import java.util.Set;
-import java.util.concurrent.ConcurrentHashMap;
 import java.util.function.Predicate;
+import java.util.function.Supplier;
 import java.util.regex.Pattern;
 import java.util.stream.Collector;
 import java.util.stream.Collectors;
@@ -46,7 +52,7 @@
  * or with {@link #assignFromSubscribed(Collection)} (automatic assignment from subscription).
  *
  * Once assigned, the partition is not considered "fetchable" until its initial position has
- * been set with {@link #seek(TopicPartition, long)}. Fetchable partitions track a fetch
+ * been set with {@link #seekValidated(TopicPartition, FetchPosition)}. Fetchable partitions track a fetch
  * position which is used to set the offset of the next fetch, and a consumed position
  * which is the last offset that has been returned to the user. You can suspend fetching
  * from a partition through {@link #pause(TopicPartition)} without affecting the fetched/consumed
@@ -56,10 +62,7 @@
  * Note that pause state as well as fetch/consumed positions are not preserved when partition
  * assignment is changed whether directly by the user or through a group rebalance.
  *
- * Thread Safety: this class is generally not thread-safe. It should only be accessed in the
- * consumer's calling thread. The only exception is {@link ConsumerMetadata} which accesses
- * the subscription state needed to build and handle Metadata requests. The thread-safe methods
- * are documented below.
+ * Thread Safety: this class is thread-safe.
  */
 public class SubscriptionState {
     private static final String SUBSCRIPTION_EXCEPTION_MESSAGE =
@@ -72,10 +75,10 @@ private enum SubscriptionType {
     }
 
     /* the type of subscription */
-    private volatile SubscriptionType subscriptionType;
+    private SubscriptionType subscriptionType;
 
     /* the pattern user has requested */
-    private volatile Pattern subscribedPattern;
+    private Pattern subscribedPattern;
 
     /* the list of topics the user has requested */
     private Set subscription;
@@ -83,7 +86,7 @@ private enum SubscriptionType {
     /* The list of topics the group has subscribed to. This may include some topics which are not part
      * of `subscription` for the leader of a group since it is responsible for detecting metadata changes
      * which require a group rebalance. */
-    private final Set groupSubscription;
+    private Set groupSubscription;
 
     /* the partitions that are currently assigned, note that the order of partition matters (see FetchBuilder for more details) */
     private final PartitionStates assignment;
@@ -91,22 +94,57 @@ private enum SubscriptionType {
     /* Default offset reset strategy */
     private final OffsetResetStrategy defaultResetStrategy;
 
-    /* Listeners provide a hook for internal state cleanup (e.g. metrics) on assignment changes */
-    private final List listeners = new ArrayList<>();
-
     /* User-provided listener to be invoked when assignment changes */
     private ConsumerRebalanceListener rebalanceListener;
 
+    private int assignmentId = 0;
+
+    @Override
+    public synchronized String toString() {
+        return "SubscriptionState{" +
+            "type=" + subscriptionType +
+            ", subscribedPattern=" + subscribedPattern +
+            ", subscription=" + String.join(",", subscription) +
+            ", groupSubscription=" + String.join(",", groupSubscription) +
+            ", defaultResetStrategy=" + defaultResetStrategy +
+            ", assignment=" + assignment.partitionStateValues() + " (id=" + assignmentId + ")}";
+    }
+
+    public synchronized String prettyString() {
+        switch (subscriptionType) {
+            case NONE:
+                return "None";
+            case AUTO_TOPICS:
+                return "Subscribe(" + String.join(",", subscription) + ")";
+            case AUTO_PATTERN:
+                return "Subscribe(" + subscribedPattern + ")";
+            case USER_ASSIGNED:
+                return "Assign(" + assignedPartitions() + " , id=" + assignmentId + ")";
+            default:
+                throw new IllegalStateException("Unrecognized subscription type: " + subscriptionType);
+        }
+    }
+
     public SubscriptionState(LogContext logContext, OffsetResetStrategy defaultResetStrategy) {
         this.log = logContext.logger(this.getClass());
         this.defaultResetStrategy = defaultResetStrategy;
-        this.subscription = Collections.emptySet();
+        this.subscription = new HashSet<>();
         this.assignment = new PartitionStates<>();
-        this.groupSubscription = ConcurrentHashMap.newKeySet();
+        this.groupSubscription = new HashSet<>();
         this.subscribedPattern = null;
         this.subscriptionType = SubscriptionType.NONE;
     }
 
+    /**
+     * Monotonically increasing id which is incremented after every assignment change. This can
+     * be used to check when an assignment has changed.
+     *
+     * @return The current assignment Id
+     */
+    synchronized int assignmentId() {
+        return assignmentId;
+    }
+
     /**
      * This method sets the subscription type if it is not already set (i.e. when it is NONE),
      * or verifies that the subscription type is equal to the give type when it is set (i.e.
@@ -120,18 +158,19 @@ else if (this.subscriptionType != type)
             throw new IllegalStateException(SUBSCRIPTION_EXCEPTION_MESSAGE);
     }
 
-    public boolean subscribe(Set topics, ConsumerRebalanceListener listener) {
-        if (listener == null)
-            throw new IllegalArgumentException("RebalanceListener cannot be null");
-
+    public synchronized boolean subscribe(Set topics, ConsumerRebalanceListener listener) {
+        registerRebalanceListener(listener);
         setSubscriptionType(SubscriptionType.AUTO_TOPICS);
-
-        this.rebalanceListener = listener;
-
         return changeSubscription(topics);
     }
 
-    public boolean subscribeFromPattern(Set topics) {
+    public synchronized void subscribe(Pattern pattern, ConsumerRebalanceListener listener) {
+        registerRebalanceListener(listener);
+        setSubscriptionType(SubscriptionType.AUTO_PATTERN);
+        this.subscribedPattern = pattern;
+    }
+
+    public synchronized boolean subscribeFromPattern(Set topics) {
         if (subscriptionType != SubscriptionType.AUTO_PATTERN)
             throw new IllegalArgumentException("Attempt to subscribe from pattern while subscription type set to " +
                     subscriptionType);
@@ -143,27 +182,29 @@ private boolean changeSubscription(Set topicsToSubscribe) {
         if (subscription.equals(topicsToSubscribe))
             return false;
 
-        this.subscription = topicsToSubscribe;
-        this.groupSubscription.addAll(topicsToSubscribe);
+        subscription = topicsToSubscribe;
         return true;
     }
 
     /**
-     * Add topics to the current group subscription. This is used by the group leader to ensure
+     * Set the current group subscription. This is used by the group leader to ensure
      * that it receives metadata updates for all topics that the group is interested in.
-     * @param topics The topics to add to the group subscription
+     *
+     * @param topics All topics from the group subscription
+     * @return true if the group subscription contains topics which are not part of the local subscription
      */
-    public boolean groupSubscribe(Collection topics) {
+    synchronized boolean groupSubscribe(Collection topics) {
         if (!partitionsAutoAssigned())
             throw new IllegalStateException(SUBSCRIPTION_EXCEPTION_MESSAGE);
-        return this.groupSubscription.addAll(topics);
+        groupSubscription = new HashSet<>(topics);
+        return !subscription.containsAll(groupSubscription);
     }
 
     /**
      * Reset the group's subscription to only contain topics subscribed by this consumer.
      */
-    public void resetGroupSubscription() {
-        this.groupSubscription.retainAll(subscription);
+    synchronized void resetGroupSubscription() {
+        groupSubscription = Collections.emptySet();
     }
 
     /**
@@ -171,13 +212,13 @@ public void resetGroupSubscription() {
      * note this is different from {@link #assignFromSubscribed(Collection)}
      * whose input partitions are provided from the subscribed topics.
      */
-    public boolean assignFromUser(Set partitions) {
+    public synchronized boolean assignFromUser(Set partitions) {
         setSubscriptionType(SubscriptionType.USER_ASSIGNED);
 
         if (this.assignment.partitionSet().equals(partitions))
             return false;
 
-        fireOnAssignment(partitions);
+        assignmentId++;
 
         Set manualSubscribedTopics = new HashSet<>();
         Map partitionToState = new HashMap<>();
@@ -193,123 +234,126 @@ public boolean assignFromUser(Set partitions) {
     }
 
     /**
-     * Change the assignment to the specified partitions returned from the coordinator, note this is
-     * different from {@link #assignFromUser(Set)} which directly set the assignment from user inputs.
-     *
      * @return true if assignments matches subscription, otherwise false
      */
-    public boolean assignFromSubscribed(Collection assignments) {
-        if (!this.partitionsAutoAssigned())
-            throw new IllegalArgumentException("Attempt to dynamically assign partitions while manual assignment in use");
-
-        Predicate predicate = topicPartition -> {
+    public synchronized boolean checkAssignmentMatchedSubscription(Collection assignments) {
+        for (TopicPartition topicPartition : assignments) {
             if (this.subscribedPattern != null) {
-                boolean match = this.subscribedPattern.matcher(topicPartition.topic()).matches();
-                if (!match) {
+                if (!this.subscribedPattern.matcher(topicPartition.topic()).matches()) {
                     log.info("Assigned partition {} for non-subscribed topic regex pattern; subscription pattern is {}",
-                            topicPartition,
-                            this.subscribedPattern);
+                        topicPartition,
+                        this.subscribedPattern);
+
+                    return false;
                 }
-                return match;
             } else {
-                boolean match = this.subscription.contains(topicPartition.topic());
-                if (!match) {
+                if (!this.subscription.contains(topicPartition.topic())) {
                     log.info("Assigned partition {} for non-subscribed topic; subscription is {}", topicPartition, this.subscription);
+
+                    return false;
                 }
-                return match;
             }
-        };
+        }
 
-        boolean assignmentMatchedSubscription = assignments.stream().allMatch(predicate);
+        return true;
+    }
 
-        if (assignmentMatchedSubscription) {
-            Map assignedPartitionStates = partitionToStateMap(
-                    assignments);
-            fireOnAssignment(assignedPartitionStates.keySet());
+    /**
+     * Change the assignment to the specified partitions returned from the coordinator, note this is
+     * different from {@link #assignFromUser(Set)} which directly set the assignment from user inputs.
+     */
+    public synchronized void assignFromSubscribed(Collection assignments) {
+        if (!this.partitionsAutoAssigned())
+            throw new IllegalArgumentException("Attempt to dynamically assign partitions while manual assignment in use");
 
-            this.assignment.set(assignedPartitionStates);
+        Map assignedPartitionStates = new HashMap<>(assignments.size());
+        for (TopicPartition tp : assignments) {
+            TopicPartitionState state = this.assignment.stateValue(tp);
+            if (state == null)
+                state = new TopicPartitionState();
+            assignedPartitionStates.put(tp, state);
         }
 
-        return assignmentMatchedSubscription;
+        assignmentId++;
+        this.assignment.set(assignedPartitionStates);
     }
 
-    public void subscribe(Pattern pattern, ConsumerRebalanceListener listener) {
+    private void registerRebalanceListener(ConsumerRebalanceListener listener) {
         if (listener == null)
             throw new IllegalArgumentException("RebalanceListener cannot be null");
-
-        setSubscriptionType(SubscriptionType.AUTO_PATTERN);
-
         this.rebalanceListener = listener;
-        this.subscribedPattern = pattern;
     }
 
     /**
-     * Check whether pattern subscription is in use. This is thread-safe.
+     * Check whether pattern subscription is in use.
      *
      */
-    public boolean hasPatternSubscription() {
+    synchronized boolean hasPatternSubscription() {
         return this.subscriptionType == SubscriptionType.AUTO_PATTERN;
     }
 
-    public boolean hasNoSubscriptionOrUserAssignment() {
+    public synchronized boolean hasNoSubscriptionOrUserAssignment() {
         return this.subscriptionType == SubscriptionType.NONE;
     }
 
-    public void unsubscribe() {
+    public synchronized void unsubscribe() {
         this.subscription = Collections.emptySet();
-        this.groupSubscription.clear();
+        this.groupSubscription = Collections.emptySet();
         this.assignment.clear();
         this.subscribedPattern = null;
         this.subscriptionType = SubscriptionType.NONE;
-        fireOnAssignment(Collections.emptySet());
+        this.assignmentId++;
     }
 
     /**
      * Check whether a topic matches a subscribed pattern.
      *
-     * This is thread-safe, but it may not always reflect the most recent subscription pattern.
-     *
      * @return true if pattern subscription is in use and the topic matches the subscribed pattern, false otherwise
      */
-    public boolean matchesSubscribedPattern(String topic) {
+    synchronized boolean matchesSubscribedPattern(String topic) {
         Pattern pattern = this.subscribedPattern;
         if (hasPatternSubscription() && pattern != null)
             return pattern.matcher(topic).matches();
         return false;
     }
 
-    public Set subscription() {
+    public synchronized Set subscription() {
         if (partitionsAutoAssigned())
             return this.subscription;
         return Collections.emptySet();
     }
 
-    public Set pausedPartitions() {
+    public synchronized Set pausedPartitions() {
         return collectPartitions(TopicPartitionState::isPaused, Collectors.toSet());
     }
 
     /**
-     * Get the subscription for the group. For the leader, this will include the union of the
-     * subscriptions of all group members. For followers, it is just that member's subscription.
-     * This is used when querying topic metadata to detect the metadata changes which would
+     * Get the subscription topics for which metadata is required. For the leader, this will include
+     * the union of the subscriptions of all group members. For followers, it is just that member's
+     * subscription. This is used when querying topic metadata to detect the metadata changes which would
      * require rebalancing. The leader fetches metadata for all topics in the group so that it
      * can do the partition assignment (which requires at least partition counts for all topics
      * to be assigned).
      *
-     * Note this is thread-safe since the Set is backed by a ConcurrentMap.
-     *
      * @return The union of all subscribed topics in the group if this member is the leader
      *   of the current generation; otherwise it returns the same set as {@link #subscription()}
      */
-    public Set groupSubscription() {
-        return this.groupSubscription;
+    synchronized Set metadataTopics() {
+        if (groupSubscription.isEmpty())
+            return subscription;
+        else if (groupSubscription.containsAll(subscription))
+            return groupSubscription;
+        else {
+            // When subscription changes `groupSubscription` may be outdated, ensure that
+            // new subscription topics are returned.
+            Set topics = new HashSet<>(groupSubscription);
+            topics.addAll(subscription);
+            return topics;
+        }
     }
 
-    /**
-     * Note this is thread-safe since the Set is backed by a ConcurrentMap.
-     */
-    public boolean isGroupSubscribed(String topic) {
-        return groupSubscription.contains(topic);
+    synchronized boolean needsMetadata(String topic) {
+        return subscription.contains(topic) || groupSubscription.contains(topic);
     }
 
     private TopicPartitionState assignedState(TopicPartition tp) {
@@ -319,107 +363,247 @@ private TopicPartitionState assignedState(TopicPartition tp) {
         return state;
     }
 
+    private TopicPartitionState assignedStateOrNull(TopicPartition tp) {
+        return this.assignment.stateValue(tp);
+    }
+
+    public synchronized void seekValidated(TopicPartition tp, FetchPosition position) {
+        assignedState(tp).seekValidated(position);
+    }
+
     public void seek(TopicPartition tp, long offset) {
-        assignedState(tp).seek(offset);
+        seekValidated(tp, new FetchPosition(offset));
+    }
+
+    public void seekUnvalidated(TopicPartition tp, FetchPosition position) {
+        assignedState(tp).seekUnvalidated(position);
+    }
+
+    synchronized void maybeSeekUnvalidated(TopicPartition tp, long offset, OffsetResetStrategy requestedResetStrategy) {
+        TopicPartitionState state = assignedStateOrNull(tp);
+        if (state == null) {
+            log.debug("Skipping reset of partition {} since it is no longer assigned", tp);
+        } else if (!state.awaitingReset()) {
+            log.debug("Skipping reset of partition {} since reset is no longer needed", tp);
+        } else if (requestedResetStrategy != state.resetStrategy) {
+            log.debug("Skipping reset of partition {} since an alternative reset has been requested", tp);
+        } else {
+            log.info("Resetting offset for partition {} to offset {}.", tp, offset);
+            state.seekUnvalidated(new FetchPosition(offset));
+        }
     }
 
     /**
-     * @return an unmodifiable view of the currently assigned partitions
+     * @return a modifiable copy of the currently assigned partitions
      */
-    public Set assignedPartitions() {
-        return this.assignment.partitionSet();
+    public synchronized Set assignedPartitions() {
+        return new HashSet<>(this.assignment.partitionSet());
+    }
+
+    /**
+     * @return a modifiable copy of the currently assigned partitions as a list
+     */
+    public synchronized List assignedPartitionsList() {
+        return new ArrayList<>(this.assignment.partitionSet());
     }
 
     /**
      * Provides the number of assigned partitions in a thread safe manner.
      * @return the number of assigned partitions.
      */
-    public int numAssignedPartitions() {
+    synchronized int numAssignedPartitions() {
         return this.assignment.size();
     }
 
-    public List fetchablePartitions() {
-        return collectPartitions(TopicPartitionState::isFetchable, Collectors.toList());
+    synchronized List fetchablePartitions(Predicate isAvailable) {
+        return assignment.stream()
+                .filter(tpState -> isAvailable.test(tpState.topicPartition()) && tpState.value().isFetchable())
+                .map(PartitionStates.PartitionState::topicPartition)
+                .collect(Collectors.toList());
     }
 
-    public boolean partitionsAutoAssigned() {
+    synchronized boolean partitionsAutoAssigned() {
         return this.subscriptionType == SubscriptionType.AUTO_TOPICS || this.subscriptionType == SubscriptionType.AUTO_PATTERN;
     }
 
-    public void position(TopicPartition tp, long offset) {
-        assignedState(tp).position(offset);
+    public synchronized void position(TopicPartition tp, FetchPosition position) {
+        assignedState(tp).position(position);
+    }
+
+    public synchronized boolean maybeValidatePositionForCurrentLeader(TopicPartition tp, Metadata.LeaderAndEpoch leaderAndEpoch) {
+        return assignedState(tp).maybeValidatePosition(leaderAndEpoch);
+    }
+
+    /**
+     * Attempt to complete validation with the end offset returned from the OffsetForLeaderEpoch request.
+     * @return The diverging offset if truncation was detected and no reset policy is defined.
+     */
+    public synchronized Optional maybeCompleteValidation(TopicPartition tp,
+                                                                            FetchPosition requestPosition,
+                                                                            EpochEndOffset epochEndOffset) {
+        TopicPartitionState state = assignedStateOrNull(tp);
+        if (state == null) {
+            log.debug("Skipping completed validation for partition {} which is not currently assigned.", tp);
+        } else if (!state.awaitingValidation()) {
+            log.debug("Skipping completed validation for partition {} which is no longer expecting validation.", tp);
+        } else {
+            SubscriptionState.FetchPosition currentPosition = state.position;
+            if (!currentPosition.equals(requestPosition)) {
+                log.debug("Skipping completed validation for partition {} since the current position {} " +
+                                "no longer matches the position {} when the request was sent",
+                        tp, currentPosition, requestPosition);
+            } else if (epochEndOffset.endOffset() < currentPosition.offset) {
+                if (hasDefaultOffsetResetPolicy()) {
+                    SubscriptionState.FetchPosition newPosition = new SubscriptionState.FetchPosition(
+                            epochEndOffset.endOffset(), Optional.of(epochEndOffset.leaderEpoch()),
+                            currentPosition.currentLeader);
+                    log.info("Truncation detected for partition {} at offset {}, resetting offset to " +
+                            "the first offset known to diverge {}", tp, currentPosition, newPosition);
+                    state.seekValidated(newPosition);
+                } else {
+                    log.warn("Truncation detected for partition {} at offset {} (the end offset from the " +
+                                    "broker is {}), but no reset policy is set",
+                            tp, currentPosition, epochEndOffset);
+                    return Optional.of(new OffsetAndMetadata(epochEndOffset.endOffset(),
+                            Optional.of(epochEndOffset.leaderEpoch()), null));
+                }
+            } else {
+                state.completeValidation();
+            }
+        }
+
+        return Optional.empty();
+    }
+
+    public synchronized boolean awaitingValidation(TopicPartition tp) {
+        return assignedState(tp).awaitingValidation();
+    }
+
+    public synchronized void completeValidation(TopicPartition tp) {
+        assignedState(tp).completeValidation();
+    }
+
+    public synchronized FetchPosition validPosition(TopicPartition tp) {
+        return assignedState(tp).validPosition();
     }
 
-    public Long position(TopicPartition tp) {
+    public synchronized FetchPosition position(TopicPartition tp) {
         return assignedState(tp).position;
     }
 
-    public Long partitionLag(TopicPartition tp, IsolationLevel isolationLevel) {
+    synchronized Long partitionLag(TopicPartition tp, IsolationLevel isolationLevel) {
         TopicPartitionState topicPartitionState = assignedState(tp);
         if (isolationLevel == IsolationLevel.READ_COMMITTED)
-            return topicPartitionState.lastStableOffset == null ? null : topicPartitionState.lastStableOffset - topicPartitionState.position;
+            return topicPartitionState.lastStableOffset == null ? null : topicPartitionState.lastStableOffset - topicPartitionState.position.offset;
         else
-            return topicPartitionState.highWatermark == null ? null : topicPartitionState.highWatermark - topicPartitionState.position;
+            return topicPartitionState.highWatermark == null ? null : topicPartitionState.highWatermark - topicPartitionState.position.offset;
     }
 
-    public Long partitionLead(TopicPartition tp) {
+    synchronized Long partitionLead(TopicPartition tp) {
         TopicPartitionState topicPartitionState = assignedState(tp);
-        return topicPartitionState.logStartOffset == null ? null : topicPartitionState.position - topicPartitionState.logStartOffset;
+        return topicPartitionState.logStartOffset == null ? null : topicPartitionState.position.offset - topicPartitionState.logStartOffset;
+    }
+
+    synchronized void updateHighWatermark(TopicPartition tp, long highWatermark) {
+        assignedState(tp).highWatermark(highWatermark);
+    }
+
+    synchronized void updateLogStartOffset(TopicPartition tp, long logStartOffset) {
+        assignedState(tp).logStartOffset(logStartOffset);
     }
 
-    public void updateHighWatermark(TopicPartition tp, long highWatermark) {
-        assignedState(tp).highWatermark = highWatermark;
+    synchronized void updateLastStableOffset(TopicPartition tp, long lastStableOffset) {
+        assignedState(tp).lastStableOffset(lastStableOffset);
     }
 
-    public void updateLogStartOffset(TopicPartition tp, long logStartOffset) {
-        assignedState(tp).logStartOffset = logStartOffset;
+    /**
+     * 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.
+     *
+     * @param tp The topic partition
+     * @param preferredReadReplicaId The preferred read replica
+     * @param timeMs The time at which this preferred replica is no longer valid
+     */
+    public synchronized void updatePreferredReadReplica(TopicPartition tp, int preferredReadReplicaId, Supplier timeMs) {
+        assignedState(tp).updatePreferredReadReplica(preferredReadReplicaId, timeMs);
     }
 
-    public void updateLastStableOffset(TopicPartition tp, long lastStableOffset) {
-        assignedState(tp).lastStableOffset = lastStableOffset;
+    /**
+     * Get the preferred read replica
+     *
+     * @param tp The topic partition
+     * @param timeMs The current time
+     * @return Returns the current preferred read replica, if it has been set and if it has not expired.
+     */
+    public synchronized Optional preferredReadReplica(TopicPartition tp, long timeMs) {
+        final TopicPartitionState topicPartitionState = assignedStateOrNull(tp);
+        if (topicPartitionState == null) {
+            return Optional.empty();
+        } else {
+            return topicPartitionState.preferredReadReplica(timeMs);
+        }
     }
 
-    public Map allConsumed() {
+    /**
+     * Unset the preferred read replica. This causes the fetcher to go back to the leader for fetches.
+     *
+     * @param tp The topic partition
+     * @return true if the preferred read replica was set, false otherwise.
+     */
+    public synchronized Optional clearPreferredReadReplica(TopicPartition tp) {
+        return assignedState(tp).clearPreferredReadReplica();
+    }
+
+    public synchronized Map allConsumed() {
         Map allConsumed = new HashMap<>();
         assignment.stream().forEach(state -> {
-            if (state.value().hasValidPosition())
-                allConsumed.put(state.topicPartition(), new OffsetAndMetadata(state.value().position));
+            TopicPartitionState partitionState = state.value();
+            if (partitionState.hasValidPosition())
+                allConsumed.put(state.topicPartition(), new OffsetAndMetadata(partitionState.position.offset,
+                        partitionState.position.offsetEpoch, ""));
         });
         return allConsumed;
     }
 
-    public void requestOffsetReset(TopicPartition partition, OffsetResetStrategy offsetResetStrategy) {
+    public synchronized void requestOffsetReset(TopicPartition partition, OffsetResetStrategy offsetResetStrategy) {
         assignedState(partition).reset(offsetResetStrategy);
     }
 
+    public synchronized void requestOffsetReset(Collection partitions, OffsetResetStrategy offsetResetStrategy) {
+        partitions.forEach(tp -> {
+            log.info("Seeking to {} offset of partition {}", offsetResetStrategy, tp);
+            assignedState(tp).reset(offsetResetStrategy);
+        });
+    }
+
     public void requestOffsetReset(TopicPartition partition) {
         requestOffsetReset(partition, defaultResetStrategy);
     }
 
-    public void setResetPending(Set partitions, long nextAllowResetTimeMs) {
+    synchronized void setNextAllowedRetry(Set partitions, long nextAllowResetTimeMs) {
         for (TopicPartition partition : partitions) {
-            assignedState(partition).setResetPending(nextAllowResetTimeMs);
+            assignedState(partition).setNextAllowedRetry(nextAllowResetTimeMs);
         }
     }
 
-    public boolean hasDefaultOffsetResetPolicy() {
+    boolean hasDefaultOffsetResetPolicy() {
         return defaultResetStrategy != OffsetResetStrategy.NONE;
     }
 
-    public boolean isOffsetResetNeeded(TopicPartition partition) {
+    public synchronized boolean isOffsetResetNeeded(TopicPartition partition) {
         return assignedState(partition).awaitingReset();
     }
 
-    public OffsetResetStrategy resetStrategy(TopicPartition partition) {
-        return assignedState(partition).resetStrategy;
+    public synchronized OffsetResetStrategy resetStrategy(TopicPartition partition) {
+        return assignedState(partition).resetStrategy();
     }
 
-    public boolean hasAllFetchPositions() {
+    public synchronized boolean hasAllFetchPositions() {
         return assignment.stream().allMatch(state -> state.value().hasValidPosition());
     }
 
-    public Set missingFetchPositions() {
-        return collectPartitions(TopicPartitionState::isMissingPosition, Collectors.toSet());
+    public synchronized Set missingFetchPositions() {
+        return collectPartitions(state -> !state.hasPosition(), Collectors.toSet());
     }
 
     private > T collectPartitions(Predicate filter, Collector collector) {
@@ -429,16 +613,17 @@ private > T collectPartitions(Predicate partitionsWithNoOffsets = new HashSet<>();
         assignment.stream().forEach(state -> {
             TopicPartition tp = state.topicPartition();
             TopicPartitionState partitionState = state.value();
-            if (partitionState.isMissingPosition()) {
+            if (!partitionState.hasPosition()) {
                 if (defaultResetStrategy == OffsetResetStrategy.NONE)
                     partitionsWithNoOffsets.add(tp);
                 else
-                    partitionState.reset(defaultResetStrategy);
+                    requestOffsetReset(tp);
             }
         });
 
@@ -446,127 +631,230 @@ public void resetMissingPositions() {
             throw new NoOffsetForPartitionException(partitionsWithNoOffsets);
     }
 
-    public Set partitionsNeedingReset(long nowMs) {
-        return collectPartitions(state -> state.awaitingReset() && state.isResetAllowed(nowMs),
+    public synchronized Set partitionsNeedingReset(long nowMs) {
+        return collectPartitions(state -> state.awaitingReset() && !state.awaitingRetryBackoff(nowMs),
+                Collectors.toSet());
+    }
+
+    public synchronized Set partitionsNeedingValidation(long nowMs) {
+        return collectPartitions(state -> state.awaitingValidation() && !state.awaitingRetryBackoff(nowMs),
                 Collectors.toSet());
     }
 
-    public boolean isAssigned(TopicPartition tp) {
+    public synchronized boolean isAssigned(TopicPartition tp) {
         return assignment.contains(tp);
     }
 
-    public boolean isPaused(TopicPartition tp) {
-        return isAssigned(tp) && assignedState(tp).paused;
+    public synchronized boolean isPaused(TopicPartition tp) {
+        TopicPartitionState assignedOrNull = assignedStateOrNull(tp);
+        return assignedOrNull != null && assignedOrNull.isPaused();
     }
 
-    public boolean isFetchable(TopicPartition tp) {
-        return isAssigned(tp) && assignedState(tp).isFetchable();
+    synchronized boolean isFetchable(TopicPartition tp) {
+        TopicPartitionState assignedOrNull = assignedStateOrNull(tp);
+        return assignedOrNull != null && assignedOrNull.isFetchable();
     }
 
-    public boolean hasValidPosition(TopicPartition tp) {
-        return isAssigned(tp) && assignedState(tp).hasValidPosition();
+    public synchronized boolean hasValidPosition(TopicPartition tp) {
+        TopicPartitionState assignedOrNull = assignedStateOrNull(tp);
+        return assignedOrNull != null && assignedOrNull.hasValidPosition();
     }
 
-    public void pause(TopicPartition tp) {
+    public synchronized void pause(TopicPartition tp) {
         assignedState(tp).pause();
     }
 
-    public void resume(TopicPartition tp) {
+    public synchronized void resume(TopicPartition tp) {
         assignedState(tp).resume();
     }
 
-    public void resetFailed(Set partitions, long nextRetryTimeMs) {
-        for (TopicPartition partition : partitions)
-            assignedState(partition).resetFailed(nextRetryTimeMs);
+    synchronized void requestFailed(Set partitions, long nextRetryTimeMs) {
+        for (TopicPartition partition : partitions) {
+            // by the time the request failed, the assignment may no longer
+            // contain this partition any more, in which case we would just ignore.
+            final TopicPartitionState state = assignedStateOrNull(partition);
+            if (state != null)
+                state.requestFailed(nextRetryTimeMs);
+        }
     }
 
-    public void movePartitionToEnd(TopicPartition tp) {
+    synchronized void movePartitionToEnd(TopicPartition tp) {
         assignment.moveToEnd(tp);
     }
 
-    public ConsumerRebalanceListener rebalanceListener() {
+    public synchronized ConsumerRebalanceListener rebalanceListener() {
         return rebalanceListener;
     }
 
-    public void addListener(Listener listener) {
-        listeners.add(listener);
-    }
-
-    public void fireOnAssignment(Set assignment) {
-        for (Listener listener : listeners)
-            listener.onAssignment(assignment);
-    }
+    private static class TopicPartitionState {
 
-    private static Map partitionToStateMap(Collection assignments) {
-        Map map = new HashMap<>(assignments.size());
-        for (TopicPartition tp : assignments)
-            map.put(tp, new TopicPartitionState());
-        return map;
-    }
+        private FetchState fetchState;
+        private FetchPosition position; // last consumed position
 
-    private static class TopicPartitionState {
-        private Long position; // last consumed position
         private Long highWatermark; // the high watermark from last fetch
         private Long logStartOffset; // the log start offset
         private Long lastStableOffset;
         private boolean paused;  // whether this partition has been paused by the user
         private OffsetResetStrategy resetStrategy;  // the strategy to use if the offset needs resetting
-        private Long nextAllowedRetryTimeMs;
+        private Long nextRetryTimeMs;
+        private Integer preferredReadReplica;
+        private Long preferredReadReplicaExpireTimeMs;
 
         TopicPartitionState() {
             this.paused = false;
+            this.fetchState = FetchStates.INITIALIZING;
             this.position = null;
             this.highWatermark = null;
             this.logStartOffset = null;
             this.lastStableOffset = null;
             this.resetStrategy = null;
-            this.nextAllowedRetryTimeMs = null;
+            this.nextRetryTimeMs = null;
+            this.preferredReadReplica = null;
+        }
+
+        private void transitionState(FetchState newState, Runnable runIfTransitioned) {
+            FetchState nextState = this.fetchState.transitionTo(newState);
+            if (nextState.equals(newState)) {
+                this.fetchState = nextState;
+                runIfTransitioned.run();
+            }
+        }
+
+        private Optional preferredReadReplica(long timeMs) {
+            if (preferredReadReplicaExpireTimeMs != null && timeMs > preferredReadReplicaExpireTimeMs) {
+                preferredReadReplica = null;
+                return Optional.empty();
+            } else {
+                return Optional.ofNullable(preferredReadReplica);
+            }
+        }
+
+        private void updatePreferredReadReplica(int preferredReadReplica, Supplier timeMs) {
+            if (this.preferredReadReplica == null || preferredReadReplica != this.preferredReadReplica) {
+                this.preferredReadReplica = preferredReadReplica;
+                this.preferredReadReplicaExpireTimeMs = timeMs.get();
+            }
+        }
+
+        private Optional clearPreferredReadReplica() {
+            if (preferredReadReplica != null) {
+                int removedReplicaId = this.preferredReadReplica;
+                this.preferredReadReplica = null;
+                this.preferredReadReplicaExpireTimeMs = null;
+                return Optional.of(removedReplicaId);
+            } else {
+                return Optional.empty();
+            }
         }
 
         private void reset(OffsetResetStrategy strategy) {
-            this.resetStrategy = strategy;
-            this.position = null;
-            this.nextAllowedRetryTimeMs = null;
+            transitionState(FetchStates.AWAIT_RESET, () -> {
+                this.resetStrategy = strategy;
+                this.nextRetryTimeMs = null;
+            });
+        }
+
+        private boolean maybeValidatePosition(Metadata.LeaderAndEpoch currentLeaderAndEpoch) {
+            if (this.fetchState.equals(FetchStates.AWAIT_RESET)) {
+                return false;
+            }
+
+            if (currentLeaderAndEpoch.equals(Metadata.LeaderAndEpoch.noLeaderOrEpoch())) {
+                // Ignore empty LeaderAndEpochs
+                return false;
+            }
+
+            if (position != null && !position.currentLeader.equals(currentLeaderAndEpoch)) {
+                FetchPosition newPosition = new FetchPosition(position.offset, position.offsetEpoch, currentLeaderAndEpoch);
+                validatePosition(newPosition);
+                preferredReadReplica = null;
+            }
+            return this.fetchState.equals(FetchStates.AWAIT_VALIDATION);
+        }
+
+        private void validatePosition(FetchPosition position) {
+            if (position.offsetEpoch.isPresent() && position.currentLeader.epoch.isPresent()) {
+                transitionState(FetchStates.AWAIT_VALIDATION, () -> {
+                    this.position = position;
+                    this.nextRetryTimeMs = null;
+                });
+            } else {
+                // If we have no epoch information for the current position, then we can skip validation
+                transitionState(FetchStates.FETCHING, () -> {
+                    this.position = position;
+                    this.nextRetryTimeMs = null;
+                });
+            }
+        }
+
+        /**
+         * Clear the awaiting validation state and enter fetching.
+         */
+        private void completeValidation() {
+            if (hasPosition()) {
+                transitionState(FetchStates.FETCHING, () -> {
+                    this.nextRetryTimeMs = null;
+                });
+            }
+        }
+
+        private boolean awaitingValidation() {
+            return fetchState.equals(FetchStates.AWAIT_VALIDATION);
         }
 
-        private boolean isResetAllowed(long nowMs) {
-            return nextAllowedRetryTimeMs == null || nowMs >= nextAllowedRetryTimeMs;
+        private boolean awaitingRetryBackoff(long nowMs) {
+            return nextRetryTimeMs != null && nowMs < nextRetryTimeMs;
         }
 
         private boolean awaitingReset() {
-            return resetStrategy != null;
+            return fetchState.equals(FetchStates.AWAIT_RESET);
         }
 
-        private void setResetPending(long nextAllowedRetryTimeMs) {
-            this.nextAllowedRetryTimeMs = nextAllowedRetryTimeMs;
+        private void setNextAllowedRetry(long nextAllowedRetryTimeMs) {
+            this.nextRetryTimeMs = nextAllowedRetryTimeMs;
         }
 
-        private void resetFailed(long nextAllowedRetryTimeMs) {
-            this.nextAllowedRetryTimeMs = nextAllowedRetryTimeMs;
+        private void requestFailed(long nextAllowedRetryTimeMs) {
+            this.nextRetryTimeMs = nextAllowedRetryTimeMs;
         }
 
         private boolean hasValidPosition() {
-            return position != null;
+            return fetchState.hasValidPosition();
         }
 
-        private boolean isMissingPosition() {
-            return !hasValidPosition() && !awaitingReset();
+        private boolean hasPosition() {
+            return fetchState.hasPosition();
         }
 
         private boolean isPaused() {
             return paused;
         }
 
-        private void seek(long offset) {
-            this.position = offset;
-            this.resetStrategy = null;
-            this.nextAllowedRetryTimeMs = null;
+        private void seekValidated(FetchPosition position) {
+            transitionState(FetchStates.FETCHING, () -> {
+                this.position = position;
+                this.resetStrategy = null;
+                this.nextRetryTimeMs = null;
+            });
+        }
+
+        private void seekUnvalidated(FetchPosition fetchPosition) {
+            seekValidated(fetchPosition);
+            validatePosition(fetchPosition);
         }
 
-        private void position(long offset) {
+        private void position(FetchPosition position) {
             if (!hasValidPosition())
                 throw new IllegalStateException("Cannot set a new position without a valid current position");
-            this.position = offset;
+            this.position = position;
+        }
+
+        private FetchPosition validPosition() {
+            if (hasValidPosition()) {
+                return position;
+            } else {
+                return null;
+            }
         }
 
         private void pause() {
@@ -581,16 +869,162 @@ private boolean isFetchable() {
             return !paused && hasValidPosition();
         }
 
+        private void highWatermark(Long highWatermark) {
+            this.highWatermark = highWatermark;
+        }
+
+        private void logStartOffset(Long logStartOffset) {
+            this.logStartOffset = logStartOffset;
+        }
+
+        private void lastStableOffset(Long lastStableOffset) {
+            this.lastStableOffset = lastStableOffset;
+        }
+
+        private OffsetResetStrategy resetStrategy() {
+            return resetStrategy;
+        }
     }
 
-    public interface Listener {
-        /**
-         * Fired after a new assignment is received (after a group rebalance or when the user manually changes the
-         * assignment).
-         *
-         * @param assignment The topic partitions assigned to the consumer
-         */
-        void onAssignment(Set assignment);
+    /**
+     * The fetch state of a partition. This class is used to determine valid state transitions and expose the some of
+     * the behavior of the current fetch state. Actual state variables are stored in the {@link TopicPartitionState}.
+     */
+    interface FetchState {
+        default FetchState transitionTo(FetchState newState) {
+            if (validTransitions().contains(newState)) {
+                return newState;
+            } else {
+                return this;
+            }
+        }
+
+        Collection validTransitions();
+
+        boolean hasPosition();
+
+        boolean hasValidPosition();
+    }
+
+    /**
+     * An enumeration of all the possible fetch states. The state transitions are encoded in the values returned by
+     * {@link FetchState#validTransitions}.
+     */
+    enum FetchStates implements FetchState {
+        INITIALIZING() {
+            @Override
+            public Collection validTransitions() {
+                return Arrays.asList(FetchStates.FETCHING, FetchStates.AWAIT_RESET, FetchStates.AWAIT_VALIDATION);
+            }
+
+            @Override
+            public boolean hasPosition() {
+                return false;
+            }
+
+            @Override
+            public boolean hasValidPosition() {
+                return false;
+            }
+        },
+
+        FETCHING() {
+            @Override
+            public Collection validTransitions() {
+                return Arrays.asList(FetchStates.FETCHING, FetchStates.AWAIT_RESET, FetchStates.AWAIT_VALIDATION);
+            }
+
+            @Override
+            public boolean hasPosition() {
+                return true;
+            }
+
+            @Override
+            public boolean hasValidPosition() {
+                return true;
+            }
+        },
+
+        AWAIT_RESET() {
+            @Override
+            public Collection validTransitions() {
+                return Arrays.asList(FetchStates.FETCHING, FetchStates.AWAIT_RESET);
+            }
+
+            @Override
+            public boolean hasPosition() {
+                return true;
+            }
+
+            @Override
+            public boolean hasValidPosition() {
+                return false;
+            }
+        },
+
+        AWAIT_VALIDATION() {
+            @Override
+            public Collection validTransitions() {
+                return Arrays.asList(FetchStates.FETCHING, FetchStates.AWAIT_RESET, FetchStates.AWAIT_VALIDATION);
+            }
+
+            @Override
+            public boolean hasPosition() {
+                return true;
+            }
+
+            @Override
+            public boolean hasValidPosition() {
+                return false;
+            }
+        }
     }
 
+    /**
+     * Represents the position of a partition subscription.
+     *
+     * This includes the offset and epoch from the last record in
+     * the batch from a FetchResponse. It also includes the leader epoch at the time the batch was consumed.
+     *
+     * The last fetch epoch is used to
+     */
+    public static class FetchPosition {
+        public final long offset;
+        final Optional offsetEpoch;
+        final Metadata.LeaderAndEpoch currentLeader;
+
+        FetchPosition(long offset) {
+            this(offset, Optional.empty(), new Metadata.LeaderAndEpoch(Node.noNode(), Optional.empty()));
+        }
+
+        public FetchPosition(long offset, Optional offsetEpoch, Metadata.LeaderAndEpoch currentLeader) {
+            this.offset = offset;
+            this.offsetEpoch = Objects.requireNonNull(offsetEpoch);
+            this.currentLeader = Objects.requireNonNull(currentLeader);
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) return true;
+            if (o == null || getClass() != o.getClass()) return false;
+            FetchPosition that = (FetchPosition) o;
+            return offset == that.offset &&
+                    offsetEpoch.equals(that.offsetEpoch) &&
+                    currentLeader.equals(that.currentLeader);
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(offset, offsetEpoch, currentLeader);
+        }
+
+        @Override
+        public String toString() {
+            return "FetchPosition{" +
+                    "offset=" + offset +
+                    ", offsetEpoch=" + offsetEpoch +
+                    ", currentLeader=" + currentLeader +
+                    '}';
+        }
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/KafkaProducer.java b/clients/src/main/java/org/apache/kafka/clients/producer/KafkaProducer.java
index 3a0130f8447f5..1be5f06b64a96 100644
--- a/clients/src/main/java/org/apache/kafka/clients/producer/KafkaProducer.java
+++ b/clients/src/main/java/org/apache/kafka/clients/producer/KafkaProducer.java
@@ -67,6 +67,7 @@
 import org.apache.kafka.common.utils.KafkaThread;
 import org.apache.kafka.common.utils.LogContext;
 import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.common.utils.Utils;
 import org.slf4j.Logger;
 
 import java.net.InetSocketAddress;
@@ -328,13 +329,12 @@ public KafkaProducer(Properties properties, Serializer keySerializer, Seriali
             Map userProvidedConfigs = config.originals();
             this.producerConfig = config;
             this.time = time;
-            String clientId = config.getString(ProducerConfig.CLIENT_ID_CONFIG);
-            if (clientId.length() <= 0)
-                clientId = "producer-" + PRODUCER_CLIENT_ID_SEQUENCE.getAndIncrement();
-            this.clientId = clientId;
 
             String transactionalId = userProvidedConfigs.containsKey(ProducerConfig.TRANSACTIONAL_ID_CONFIG) ?
                     (String) userProvidedConfigs.get(ProducerConfig.TRANSACTIONAL_ID_CONFIG) : null;
+
+            this.clientId = buildClientId(config.getString(ProducerConfig.CLIENT_ID_CONFIG), transactionalId);
+
             LogContext logContext;
             if (transactionalId == null)
                 logContext = new LogContext(String.format("[Producer clientId=%s] ", clientId));
@@ -353,6 +353,7 @@ public KafkaProducer(Properties properties, Serializer keySerializer, Seriali
                     Collections.singletonMap(ProducerConfig.CLIENT_ID_CONFIG, clientId));
             reporters.add(new JmxReporter(JMX_PREFIX));
             this.metrics = new Metrics(metricConfig, reporters, time);
+            this.metrics.setReplaceOnDuplicateMetric(config.getBoolean(ProducerConfig.METRICS_REPLACE_ON_DUPLICATE_CONFIG));
             this.partitioner = config.getConfiguredInstance(ProducerConfig.PARTITIONER_CLASS_CONFIG, Partitioner.class);
             long retryBackoffMs = config.getLong(ProducerConfig.RETRY_BACKOFF_MS_CONFIG);
             if (keySerializer == null) {
@@ -395,7 +396,7 @@ public KafkaProducer(Properties properties, Serializer keySerializer, Seriali
             this.accumulator = new RecordAccumulator(logContext,
                     config.getInt(ProducerConfig.BATCH_SIZE_CONFIG),
                     this.compressionType,
-                    config.getInt(ProducerConfig.LINGER_MS_CONFIG),
+                    lingerMs(config),
                     retryBackoffMs,
                     deliveryTimeoutMs,
                     metrics,
@@ -414,8 +415,10 @@ public KafkaProducer(Properties properties, Serializer keySerializer, Seriali
                         config.getLong(ProducerConfig.METADATA_MAX_AGE_CONFIG),
                         logContext,
                         clusterResourceListeners,
-                        Time.SYSTEM);
-                this.metadata.bootstrap(addresses, time.milliseconds());
+                        Time.SYSTEM,
+                        config.getLong(ProducerConfig.METADATA_TOPIC_EXPIRY_MS_CONFIG),
+                        config.getBoolean(ProducerConfig.ALLOW_AUTO_CREATE_TOPICS_CONFIG));
+                this.metadata.bootstrap(addresses);
             }
             this.errors = this.metrics.sensor("errors");
             this.sender = newSender(logContext, kafkaClient, this.metadata);
@@ -423,7 +426,7 @@ public KafkaProducer(Properties properties, Serializer keySerializer, Seriali
             this.ioThread = new KafkaThread(ioThreadName, this.sender, true);
             this.ioThread.start();
             config.logUnused();
-            AppInfoParser.registerAppInfo(JMX_PREFIX, clientId, metrics);
+            AppInfoParser.registerAppInfo(JMX_PREFIX, clientId, metrics, time.milliseconds());
             log.debug("Kafka producer started");
         } catch (Throwable t) {
             // call close methods if internal objects are already constructed this is to prevent resource leak. see KAFKA-2121
@@ -433,6 +436,16 @@ public KafkaProducer(Properties properties, Serializer keySerializer, Seriali
         }
     }
 
+    private static String buildClientId(String configuredClientId, String transactionalId) {
+        if (!configuredClientId.isEmpty())
+            return configuredClientId;
+
+        if (transactionalId != null)
+            return "producer-" + transactionalId;
+
+        return "producer-" + PRODUCER_CLIENT_ID_SEQUENCE.getAndIncrement();
+    }
+
     // visible for testing
     Sender newSender(LogContext logContext, KafkaClient kafkaClient, ProducerMetadata metadata) {
         int maxInflightRequests = configureInflightRequests(producerConfig, transactionManager != null);
@@ -475,12 +488,17 @@ Sender newSender(LogContext logContext, KafkaClient kafkaClient, ProducerMetadat
                 apiVersions);
     }
 
+    private static int lingerMs(ProducerConfig config) {
+        return (int) Math.min(config.getLong(ProducerConfig.LINGER_MS_CONFIG), Integer.MAX_VALUE);
+    }
+
     private static int configureDeliveryTimeout(ProducerConfig config, Logger log) {
         int deliveryTimeoutMs = config.getInt(ProducerConfig.DELIVERY_TIMEOUT_MS_CONFIG);
-        int lingerMs = config.getInt(ProducerConfig.LINGER_MS_CONFIG);
+        int lingerMs = lingerMs(config);
         int requestTimeoutMs = config.getInt(ProducerConfig.REQUEST_TIMEOUT_MS_CONFIG);
+        int lingerAndRequestTimeoutMs = (int) Math.min((long) lingerMs + requestTimeoutMs, Integer.MAX_VALUE);
 
-        if (deliveryTimeoutMs < Integer.MAX_VALUE && deliveryTimeoutMs < lingerMs + requestTimeoutMs) {
+        if (deliveryTimeoutMs < Integer.MAX_VALUE && deliveryTimeoutMs < lingerAndRequestTimeoutMs) {
             if (config.originals().containsKey(ProducerConfig.DELIVERY_TIMEOUT_MS_CONFIG)) {
                 // throw an exception if the user explicitly set an inconsistent value
                 throw new ConfigException(ProducerConfig.DELIVERY_TIMEOUT_MS_CONFIG
@@ -488,7 +506,7 @@ private static int configureDeliveryTimeout(ProducerConfig config, Logger log) {
                     + " + " + ProducerConfig.REQUEST_TIMEOUT_MS_CONFIG);
             } else {
                 // override deliveryTimeoutMs default value to lingerMs + requestTimeoutMs for backward compatibility
-                deliveryTimeoutMs = lingerMs + requestTimeoutMs;
+                deliveryTimeoutMs = lingerAndRequestTimeoutMs;
                 log.warn("{} should be equal to or larger than {} + {}. Setting it to {}.",
                     ProducerConfig.DELIVERY_TIMEOUT_MS_CONFIG, ProducerConfig.LINGER_MS_CONFIG,
                     ProducerConfig.REQUEST_TIMEOUT_MS_CONFIG, deliveryTimeoutMs);
@@ -611,6 +629,7 @@ private static int parseAcks(String acksString) {
      */
     public void initTransactions() {
         throwIfNoTransactionManager();
+        throwIfProducerClosed();
         TransactionalRequestResult result = transactionManager.initializeTransactions();
         sender.wakeup();
         result.await(maxBlockTimeMs, TimeUnit.MILLISECONDS);
@@ -631,6 +650,7 @@ public void initTransactions() {
      */
     public void beginTransaction() throws ProducerFencedException {
         throwIfNoTransactionManager();
+        throwIfProducerClosed();
         transactionManager.beginTransaction();
     }
 
@@ -661,6 +681,7 @@ public void beginTransaction() throws ProducerFencedException {
     public void sendOffsetsToTransaction(Map offsets,
                                          String consumerGroupId) throws ProducerFencedException {
         throwIfNoTransactionManager();
+        throwIfProducerClosed();
         TransactionalRequestResult result = transactionManager.sendOffsetsToTransaction(offsets, consumerGroupId);
         sender.wakeup();
         result.await();
@@ -691,6 +712,7 @@ public void sendOffsetsToTransaction(Map offs
      */
     public void commitTransaction() throws ProducerFencedException {
         throwIfNoTransactionManager();
+        throwIfProducerClosed();
         TransactionalRequestResult result = transactionManager.beginCommit();
         sender.wakeup();
         result.await(maxBlockTimeMs, TimeUnit.MILLISECONDS);
@@ -718,6 +740,7 @@ public void commitTransaction() throws ProducerFencedException {
      */
     public void abortTransaction() throws ProducerFencedException {
         throwIfNoTransactionManager();
+        throwIfProducerClosed();
         TransactionalRequestResult result = transactionManager.beginAbort();
         sender.wakeup();
         result.await(maxBlockTimeMs, TimeUnit.MILLISECONDS);
@@ -848,7 +871,7 @@ public Future send(ProducerRecord record, Callback callbac
     // Verify that this producer instance has not been closed. This method throws IllegalStateException if the producer
     // has already been closed.
     private void throwIfProducerClosed() {
-        if (ioThread == null || !ioThread.isAlive())
+        if (sender == null || !sender.isRunning())
             throw new IllegalStateException("Cannot perform operation after producer has been closed");
     }
 
@@ -896,15 +919,36 @@ private Future doSend(ProducerRecord record, Callback call
                     compressionType, serializedKey, serializedValue, headers);
             ensureValidRecordSize(serializedSize);
             long timestamp = record.timestamp() == null ? time.milliseconds() : record.timestamp();
-            log.trace("Sending record {} with callback {} to topic {} partition {}", record, callback, record.topic(), partition);
+            if (log.isTraceEnabled()) {
+                log.trace("Attempting to append record {} with callback {} to topic {} partition {}", record, callback, record.topic(), partition);
+            }
             // producer callback will make sure to call both 'callback' and interceptor callback
             Callback interceptCallback = new InterceptorCallback<>(callback, this.interceptors, tp);
 
+            if (transactionManager != null && transactionManager.isTransactional()) {
+                transactionManager.failIfNotReadyForSend();
+            }
+            RecordAccumulator.RecordAppendResult result = accumulator.append(tp, timestamp, serializedKey,
+                    serializedValue, headers, interceptCallback, remainingWaitMs, true);
+            
+            if (result.abortForNewBatch) {
+                int prevPartition = partition;
+                partitioner.onNewBatch(record.topic(), cluster, prevPartition);
+                partition = partition(record, serializedKey, serializedValue, cluster);
+                tp = new TopicPartition(record.topic(), partition);
+                if (log.isTraceEnabled()) {
+                    log.trace("Retrying append due to new batch creation for topic {} partition {}. The old partition was {}", record.topic(), partition, prevPartition);
+                }
+                // producer callback will make sure to call both 'callback' and interceptor callback
+                interceptCallback = new InterceptorCallback<>(callback, this.interceptors, tp);
+
+                result = accumulator.append(tp, timestamp, serializedKey,
+                    serializedValue, headers, interceptCallback, remainingWaitMs, false);
+            }
+            
             if (transactionManager != null && transactionManager.isTransactional())
                 transactionManager.maybeAddPartitionToTransaction(tp);
 
-            RecordAccumulator.RecordAppendResult result = accumulator.append(tp, timestamp, serializedKey,
-                    serializedValue, headers, interceptCallback, remainingWaitMs);
             if (result.batchIsFull || result.newBatchCreated) {
                 log.trace("Waking up the sender since topic {} partition {} is either full or getting a new batch", record.topic(), partition);
                 this.sender.wakeup();
@@ -1001,7 +1045,7 @@ private ClusterAndWaitTime waitOnMetadata(String topic, Integer partition, long
                         String.format("Partition %d of topic %s with partition count %d is not present in metadata after %d ms.",
                                 partition, topic, partitionsCount, maxWaitMs));
             }
-            metadata.maybeThrowException();
+            metadata.maybeThrowExceptionForTopic(topic);
             remainingWaitMs = maxWaitMs - elapsed;
             partitionsCount = cluster.partitionCountForTopic(topic);
         } while (partitionsCount == null || (partition != null && partition >= partitionsCount));
@@ -1061,11 +1105,27 @@ private void ensureValidRecordSize(int size) {
      */
     @Override
     public void flush() {
-        log.trace("Flushing accumulated records in producer.");
-        this.accumulator.beginFlush();
-        this.sender.wakeup();
+        flush(Integer.MAX_VALUE, TimeUnit.MILLISECONDS);
+    }
+
+    /**
+     * This method waits up to timeout for the producer to send out all the buffered records.
+     * @param timeout The maximum time to wait for producer to complete. The value should be non-negative.
+     * @param unit The time unit for the timeout
+     * @throws TimeoutException If producer fail to finish in time
+     * @throws InterruptException If the thread is interrupted while blocked
+     * @throws IllegalArgumentException If the timeout is negative.
+     */
+    @Override
+    public void flush(long timeout, TimeUnit unit) {
+        if (timeout < 0)
+            throw new IllegalArgumentException("The timeout cannot be negative.");
         try {
-            this.accumulator.awaitFlushCompletion();
+            this.accumulator.beginFlush();
+            this.sender.wakeup();
+            this.accumulator.awaitFlushCompletion(unit.toMillis(timeout));
+        } catch (TimeoutException e) {
+            throw e;
         } catch (InterruptedException e) {
             throw new InterruptException("Flush interrupted.", e);
         }
@@ -1117,7 +1177,8 @@ public void close() {
      * This method waits up to timeout for the producer to complete the sending of all incomplete requests.
      * 
    
      * If the producer is unable to complete all requests before the timeout expires, this method will fail
-     * any unsent and unacknowledged records immediately.
+     * any unsent and unacknowledged records immediately. It will also abort the ongoing transaction if it's not
+     * already completing.
      * 
    
      * If invoked from within a {@link Callback} this method will not block and will be equivalent to
      * close(Duration.ofMillis(0)). This is done since no further sending will happen while
@@ -1177,13 +1238,12 @@ private void close(Duration timeout, boolean swallowException) {
             }
         }
 
-        ClientUtils.closeQuietly(interceptors, "producer interceptors", firstException);
-        ClientUtils.closeQuietly(metrics, "producer metrics", firstException);
-        ClientUtils.closeQuietly(keySerializer, "producer keySerializer", firstException);
-        ClientUtils.closeQuietly(valueSerializer, "producer valueSerializer", firstException);
-        ClientUtils.closeQuietly(partitioner, "producer partitioner", firstException);
+        Utils.closeQuietly(interceptors, "producer interceptors", firstException);
+        Utils.closeQuietly(metrics, "producer metrics", firstException);
+        Utils.closeQuietly(keySerializer, "producer keySerializer", firstException);
+        Utils.closeQuietly(valueSerializer, "producer valueSerializer", firstException);
+        Utils.closeQuietly(partitioner, "producer partitioner", firstException);
         AppInfoParser.unregisterAppInfo(JMX_PREFIX, clientId, metrics);
-        log.debug("Kafka producer has been closed");
         Throwable exception = firstException.get();
         if (exception != null && !swallowException) {
             if (exception instanceof InterruptException) {
@@ -1191,6 +1251,7 @@ private void close(Duration timeout, boolean swallowException) {
             }
             throw new KafkaException("Failed to close kafka producer", exception);
         }
+        log.debug("Kafka producer has been closed");
     }
 
     private static Map propsToMap(Properties properties) {
diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/MockProducer.java b/clients/src/main/java/org/apache/kafka/clients/producer/MockProducer.java
index c2561cfc79b69..cddcc2a463f58 100644
--- a/clients/src/main/java/org/apache/kafka/clients/producer/MockProducer.java
+++ b/clients/src/main/java/org/apache/kafka/clients/producer/MockProducer.java
@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.clients.producer;
 
+import java.util.concurrent.TimeUnit;
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
 import org.apache.kafka.clients.producer.internals.DefaultPartitioner;
 import org.apache.kafka.clients.producer.internals.FutureRecordMetadata;
@@ -229,7 +230,7 @@ private void verifyNoTransactionInFlight() {
 
     /**
      * Adds the record to the list of sent records. The {@link RecordMetadata} returned will be immediately satisfied.
-     * 
+     *
      * @see #history()
      */
     @Override
@@ -289,12 +290,16 @@ private long nextOffset(TopicPartition tp) {
         }
     }
 
-    public synchronized void flush() {
+    public synchronized void flush(long timeout, TimeUnit unit) {
         verifyProducerState();
         while (!this.completions.isEmpty())
             completeNext();
     }
 
+    public synchronized void flush() {
+        flush(Long.MAX_VALUE, null);
+    }
+
     public List partitionsFor(String topic) {
         return this.cluster.partitionsForTopic(topic);
     }
diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/Partitioner.java b/clients/src/main/java/org/apache/kafka/clients/producer/Partitioner.java
index a3a2fbea86408..1baeb0f648ba7 100644
--- a/clients/src/main/java/org/apache/kafka/clients/producer/Partitioner.java
+++ b/clients/src/main/java/org/apache/kafka/clients/producer/Partitioner.java
@@ -44,4 +44,14 @@ public interface Partitioner extends Configurable, Closeable {
      */
     public void close();
 
+
+    /**
+     * Notifies the partitioner a new batch is about to be created. When using the sticky partitioner,
+     * this method can change the chosen sticky partition for the new batch. 
+     * @param topic The topic name
+     * @param cluster The current cluster metadata
+     * @param prevPartition The partition previously selected for the record that triggered a new batch
+     */
+    default public void onNewBatch(String topic, Cluster cluster, int prevPartition) {
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/Producer.java b/clients/src/main/java/org/apache/kafka/clients/producer/Producer.java
index 96d487a79235b..0e923f09d7b4e 100644
--- a/clients/src/main/java/org/apache/kafka/clients/producer/Producer.java
+++ b/clients/src/main/java/org/apache/kafka/clients/producer/Producer.java
@@ -79,6 +79,11 @@ void sendOffsetsToTransaction(Map offsets,
     void flush();
 
     /**
+     * See {@link KafkaProducer#flush(long, TimeUnit)}
+     */
+    void flush(long timeout, TimeUnit unit);
+
+     /**
      * See {@link KafkaProducer#partitionsFor(String)}
      */
     List partitionsFor(String topic);
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 9324b9e444fe3..b8bf54012ce46 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
@@ -19,10 +19,12 @@
 import org.apache.kafka.clients.ClientDnsLookup;
 import org.apache.kafka.clients.CommonClientConfigs;
 import org.apache.kafka.clients.producer.internals.DefaultPartitioner;
+import org.apache.kafka.clients.producer.internals.ProducerMetadata;
 import org.apache.kafka.common.config.AbstractConfig;
 import org.apache.kafka.common.config.ConfigDef;
 import org.apache.kafka.common.config.ConfigDef.Importance;
 import org.apache.kafka.common.config.ConfigDef.Type;
+import org.apache.kafka.common.config.SecurityConfig;
 import org.apache.kafka.common.metrics.Sensor;
 import org.apache.kafka.common.serialization.Serializer;
 
@@ -52,9 +54,13 @@ public class ProducerConfig extends AbstractConfig {
     /** bootstrap.servers */
     public static final String BOOTSTRAP_SERVERS_CONFIG = CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG;
 
+
     /** client.dns.lookup */
     public static final String CLIENT_DNS_LOOKUP_CONFIG = CommonClientConfigs.CLIENT_DNS_LOOKUP_CONFIG;
 
+    public static final String METADATA_TOPIC_EXPIRY_MS_CONFIG = CommonClientConfigs.METADATA_TOPIC_EXPIRY_MS_CONFIG;
+    private static final String METADATA_TOPIC_EXPIRY_MS_DOC = CommonClientConfigs.METADATA_TOPIC_EXPIRY_MS_DOC;
+
     /** metadata.max.age.ms */
     public static final String METADATA_MAX_AGE_CONFIG = CommonClientConfigs.METADATA_MAX_AGE_CONFIG;
     private static final String METADATA_MAX_AGE_DOC = CommonClientConfigs.METADATA_MAX_AGE_DOC;
@@ -82,13 +88,14 @@ public class ProducerConfig extends AbstractConfig {
                                            + " server at all. The record will be immediately added to the socket buffer and considered sent. No guarantee can be"
                                            + " made that the server has received the record in this case, and the retries configuration will not"
                                            + " take effect (as the client won't generally know of any failures). The offset given back for each record will"
-                                           + " always be set to -1."
+                                           + " always be set to -1."
                                            + " 
  • acks=1 This will mean the leader will write the record to its local log but will respond"
                                            + " without awaiting full acknowledgement from all followers. In this case should the leader fail immediately after"
                                            + " acknowledging the record but before the followers have replicated it then the record will be lost."
                                            + " 
  • acks=all This means the leader will wait for the full set of in-sync replicas to"
                                            + " acknowledge the record. This guarantees that the record will not be lost as long as at least one in-sync replica"
-                                           + " remains alive. This is the strongest available guarantee. This is equivalent to the acks=-1 setting.";
+                                           + " remains alive. This is the strongest available guarantee. This is equivalent to the acks=-1 setting."
+                                           + "";
 
     /** linger.ms */
     public static final String LINGER_MS_CONFIG = "linger.ms";
@@ -177,6 +184,11 @@ public class ProducerConfig extends AbstractConfig {
      */
     public static final String METRICS_RECORDING_LEVEL_CONFIG = CommonClientConfigs.METRICS_RECORDING_LEVEL_CONFIG;
 
+    /**
+     * metrics.replace.on.duplicate
+     */
+    public static final String METRICS_REPLACE_ON_DUPLICATE_CONFIG = CommonClientConfigs.METRICS_REPLACE_ON_DUPLICATE_CONFIG;
+
     /** metric.reporters */
     public static final String METRIC_REPORTER_CLASSES_CONFIG = CommonClientConfigs.METRIC_REPORTER_CLASSES_CONFIG;
 
@@ -240,6 +252,16 @@ public class ProducerConfig extends AbstractConfig {
             "The default is null, which means transactions cannot be used. " +
             "Note that, by default, transactions require a cluster of at least three brokers which is the recommended setting for production; for development you can change this, by adjusting broker setting transaction.state.log.replication.factor.";
 
+    /**
+     * security.providers
+     */
+    public static final String SECURITY_PROVIDERS_CONFIG = SecurityConfig.SECURITY_PROVIDERS_CONFIG;
+    private static final String SECURITY_PROVIDERS_DOC = SecurityConfig.SECURITY_PROVIDERS_DOC;
+    /** allow.auto.create.topics */
+    public static final String ALLOW_AUTO_CREATE_TOPICS_CONFIG = "allow.auto.create.topics";
+    private static final String ALLOW_AUTO_CREATE_TOPICS_DOC = "The client-side (producer) permission to allow auto-topic creation. Both the client-side and the broker-side should enable auto-topic creation in order for a topic to be automatically created";
+    public static final boolean DEFAULT_ALLOW_AUTO_CREATE_TOPICS = false;
+
     static {
         CONFIG = new ConfigDef().define(BOOTSTRAP_SERVERS_CONFIG, Type.LIST, Collections.emptyList(), new ConfigDef.NonNullValidator(), Importance.HIGH, CommonClientConfigs.BOOTSTRAP_SERVERS_DOC)
                                 .define(CLIENT_DNS_LOOKUP_CONFIG,
@@ -260,8 +282,8 @@ public class ProducerConfig extends AbstractConfig {
                                         ACKS_DOC)
                                 .define(COMPRESSION_TYPE_CONFIG, Type.STRING, "none", Importance.HIGH, COMPRESSION_TYPE_DOC)
                                 .define(BATCH_SIZE_CONFIG, Type.INT, 16384, atLeast(0), Importance.MEDIUM, BATCH_SIZE_DOC)
-                                .define(LINGER_MS_CONFIG, Type.INT, 0, atLeast(0), Importance.MEDIUM, LINGER_MS_DOC)
-                                .define(DELIVERY_TIMEOUT_MS_CONFIG, Type.INT, 120 * 1000, atLeast(0), Importance.MEDIUM, DELIVERY_TIMEOUT_MS_DOC)
+                                .define(LINGER_MS_CONFIG, Type.LONG, 0, atLeast(0), Importance.MEDIUM, LINGER_MS_DOC)
+                                .define(DELIVERY_TIMEOUT_MS_CONFIG, Type.INT, Integer.MAX_VALUE, atLeast(0), Importance.MEDIUM, DELIVERY_TIMEOUT_MS_DOC)
                                 .define(CLIENT_ID_CONFIG, Type.STRING, "", Importance.MEDIUM, CommonClientConfigs.CLIENT_ID_DOC)
                                 .define(SEND_BUFFER_CONFIG, Type.INT, 128 * 1024, atLeast(CommonClientConfigs.SEND_BUFFER_LOWER_BOUND), Importance.MEDIUM, CommonClientConfigs.SEND_BUFFER_DOC)
                                 .define(RECEIVE_BUFFER_CONFIG, Type.INT, 32 * 1024, atLeast(CommonClientConfigs.RECEIVE_BUFFER_LOWER_BOUND), Importance.MEDIUM, CommonClientConfigs.RECEIVE_BUFFER_DOC)
@@ -306,6 +328,11 @@ public class ProducerConfig extends AbstractConfig {
                                         new ConfigDef.NonNullValidator(),
                                         Importance.LOW,
                                         CommonClientConfigs.METRIC_REPORTER_CLASSES_DOC)
+                                .define(METRICS_REPLACE_ON_DUPLICATE_CONFIG,
+                                        Type.BOOLEAN,
+                                        false,
+                                        Importance.LOW,
+                                        CommonClientConfigs.METRICS_REPLACE_ON_DUPLICATE_DOC)
                                 .define(MAX_IN_FLIGHT_REQUESTS_PER_CONNECTION,
                                         Type.INT,
                                         5,
@@ -341,6 +368,16 @@ public class ProducerConfig extends AbstractConfig {
                                         CommonClientConfigs.DEFAULT_SECURITY_PROTOCOL,
                                         Importance.MEDIUM,
                                         CommonClientConfigs.SECURITY_PROTOCOL_DOC)
+                                .define(SECURITY_PROVIDERS_CONFIG,
+                                        Type.STRING,
+                                        null,
+                                        Importance.LOW,
+                                        SECURITY_PROVIDERS_DOC)
+                                .define(CommonClientConfigs.ENABLE_STICKY_METADATA_FETCH_CONFIG,
+                                        Type.BOOLEAN,
+                                        false,
+                                        Importance.MEDIUM,
+                                        CommonClientConfigs.ENABLE_STICKY_METADATA_FETCH_DOC)
                                 .withClientSslSupport()
                                 .withClientSaslSupport()
                                 .define(ENABLE_IDEMPOTENCE_CONFIG,
@@ -353,12 +390,22 @@ public class ProducerConfig extends AbstractConfig {
                                         60000,
                                         Importance.LOW,
                                         TRANSACTION_TIMEOUT_DOC)
+                                .define(METADATA_TOPIC_EXPIRY_MS_CONFIG,
+                                        Type.LONG,
+                                        ProducerMetadata.TOPIC_EXPIRY_MS,
+                                        Importance.LOW,
+                                        METADATA_TOPIC_EXPIRY_MS_DOC)
                                 .define(TRANSACTIONAL_ID_CONFIG,
                                         Type.STRING,
                                         null,
                                         new ConfigDef.NonEmptyString(),
                                         Importance.LOW,
-                                        TRANSACTIONAL_ID_DOC);
+                                        TRANSACTIONAL_ID_DOC)
+                                .define(ALLOW_AUTO_CREATE_TOPICS_CONFIG,
+                                        Type.BOOLEAN,
+                                        DEFAULT_ALLOW_AUTO_CREATE_TOPICS,
+                                        Importance.MEDIUM,
+                                        ALLOW_AUTO_CREATE_TOPICS_DOC);
     }
 
     @Override
@@ -404,8 +451,12 @@ public static Set configNames() {
         return CONFIG.names();
     }
 
+    public static ConfigDef configDef() {
+        return new ConfigDef(CONFIG);
+    }
+
     public static void main(String[] args) {
-        System.out.println(CONFIG.toHtmlTable());
+        System.out.println(CONFIG.toHtml());
     }
 
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/ProducerRecord.java b/clients/src/main/java/org/apache/kafka/clients/producer/ProducerRecord.java
index 5d6bb3ff82e54..0fa37dc15d068 100644
--- a/clients/src/main/java/org/apache/kafka/clients/producer/ProducerRecord.java
+++ b/clients/src/main/java/org/apache/kafka/clients/producer/ProducerRecord.java
@@ -20,6 +20,8 @@
 import org.apache.kafka.common.header.Headers;
 import org.apache.kafka.common.header.internals.RecordHeaders;
 
+import java.util.Objects;
+
 /**
  * A key/value pair to be sent to Kafka. This consists of a topic name to which the record is being sent, an optional
  * partition number, and an optional key and value.
@@ -202,20 +204,12 @@ else if (!(o instanceof ProducerRecord))
 
         ProducerRecord that = (ProducerRecord) o;
 
-        if (key != null ? !key.equals(that.key) : that.key != null) 
-            return false;
-        else if (partition != null ? !partition.equals(that.partition) : that.partition != null) 
-            return false;
-        else if (topic != null ? !topic.equals(that.topic) : that.topic != null) 
-            return false;
-        else if (headers != null ? !headers.equals(that.headers) : that.headers != null)
-            return false;
-        else if (value != null ? !value.equals(that.value) : that.value != null) 
-            return false;
-        else if (timestamp != null ? !timestamp.equals(that.timestamp) : that.timestamp != null)
-            return false;
-
-        return true;
+        return Objects.equals(key, that.key) &&
+            Objects.equals(partition, that.partition) &&
+            Objects.equals(topic, that.topic) &&
+            Objects.equals(headers, that.headers) &&
+            Objects.equals(value, that.value) &&
+            Objects.equals(timestamp, that.timestamp);
     }
 
     @Override
diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/RoundRobinPartitioner.java b/clients/src/main/java/org/apache/kafka/clients/producer/RoundRobinPartitioner.java
new file mode 100644
index 0000000000000..80c47252b13f2
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/clients/producer/RoundRobinPartitioner.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.clients.producer;
+
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import org.apache.kafka.common.Cluster;
+import org.apache.kafka.common.PartitionInfo;
+import org.apache.kafka.common.utils.Utils;
+
+/**
+ * The "Round-Robin" partitioner
+ * 
+ * This partitioning strategy can be used when user wants 
+ * to distribute the writes to all partitions equally. This
+ * is the behaviour regardless of record key hash. 
+ *
+ */
+public class RoundRobinPartitioner implements Partitioner {
+    private final ConcurrentMap topicCounterMap = new ConcurrentHashMap<>();
+
+    public void configure(Map configs) {}
+
+    /**
+     * Compute the partition for the given record.
+     *
+     * @param topic The topic name
+     * @param key The key to partition on (or null if no key)
+     * @param keyBytes serialized key to partition on (or null if no key)
+     * @param value The value to partition on or null
+     * @param valueBytes serialized value to partition on or null
+     * @param cluster The current cluster metadata
+     */
+    @Override
+    public int partition(String topic, Object key, byte[] keyBytes, Object value, byte[] valueBytes, Cluster cluster) {
+        List partitions = cluster.partitionsForTopic(topic);
+        int numPartitions = partitions.size();
+        int nextValue = nextValue(topic);
+        List availablePartitions = cluster.availablePartitionsForTopic(topic);
+        if (!availablePartitions.isEmpty()) {
+            int part = Utils.toPositive(nextValue) % availablePartitions.size();
+            return availablePartitions.get(part).partition();
+        } else {
+            // no partitions are available, give a non-available partition
+            return Utils.toPositive(nextValue) % numPartitions;
+        }
+    }
+
+    private int nextValue(String topic) {
+        AtomicInteger counter = topicCounterMap.computeIfAbsent(topic, k -> {
+            return new AtomicInteger(0);
+        });
+        return counter.getAndIncrement();
+    }
+
+    public void close() {}
+
+}
diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/UniformStickyPartitioner.java b/clients/src/main/java/org/apache/kafka/clients/producer/UniformStickyPartitioner.java
new file mode 100644
index 0000000000000..3b06be47b146f
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/clients/producer/UniformStickyPartitioner.java
@@ -0,0 +1,65 @@
+/*
+ * 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.producer;
+
+import java.util.Map;
+
+import org.apache.kafka.clients.producer.internals.StickyPartitionCache;
+import org.apache.kafka.common.Cluster;
+
+
+/**
+ * The partitioning strategy:
+ * 
      
+ * 
    • If a partition is specified in the record, use it
+ * 
    • Otherwise choose the sticky partition that changes when the batch is full.
+ * 
+ * NOTE: In constrast to the DefaultPartitioner, the record key is NOT used as part of the partitioning strategy in this 
+ *       partitioner. Records with the same key are not guaranteed to be sent to the same partition.
+ * 
+ * See KIP-480 for details about sticky partitioning.
+ */
+public class UniformStickyPartitioner implements Partitioner {
+
+    private final StickyPartitionCache stickyPartitionCache = new StickyPartitionCache();
+
+    public void configure(Map configs) {}
+
+    /**
+     * Compute the partition for the given record.
+     *
+     * @param topic The topic name
+     * @param key The key to partition on (or null if no key)
+     * @param keyBytes serialized key to partition on (or null if no key)
+     * @param value The value to partition on or null
+     * @param valueBytes serialized value to partition on or null
+     * @param cluster The current cluster metadata
+     */
+    public int partition(String topic, Object key, byte[] keyBytes, Object value, byte[] valueBytes, Cluster cluster) {
+        return stickyPartitionCache.partition(topic, cluster);
+    }
+
+    public void close() {}
+    
+    /**
+     * If a batch completed for the current sticky partition, change the sticky partition. 
+     * Alternately, if no sticky partition has been determined, set one.
+     */
+    public void onNewBatch(String topic, Cluster cluster, int prevPartition) {
+        stickyPartitionCache.nextPartition(topic, cluster, prevPartition);
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/BufferPool.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/BufferPool.java
index 22d747265d618..70fc3a62605b2 100644
--- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/BufferPool.java
+++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/BufferPool.java
@@ -19,16 +19,22 @@
 import java.nio.ByteBuffer;
 import java.util.ArrayDeque;
 import java.util.Deque;
+import java.util.Iterator;
+import java.util.LinkedHashSet;
+import java.util.Set;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.locks.Condition;
 import java.util.concurrent.locks.ReentrantLock;
 
+import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.MetricName;
 import org.apache.kafka.common.errors.TimeoutException;
 import org.apache.kafka.common.metrics.Metrics;
 import org.apache.kafka.common.metrics.Sensor;
 import org.apache.kafka.common.metrics.stats.Meter;
 import org.apache.kafka.common.utils.Time;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 
 /**
@@ -43,18 +49,22 @@
  */
 public class BufferPool {
 
+    private static final Logger log = LoggerFactory.getLogger(BufferPool.class);
+
     static final String WAIT_TIME_SENSOR_NAME = "bufferpool-wait-time";
 
     private final long totalMemory;
     private final int poolableSize;
     private final ReentrantLock lock;
-    private final Deque free;
+    private final Set free;
     private final Deque waiters;
     /** Total available memory is the sum of nonPooledAvailableMemory and the number of byte buffers in free * poolableSize.  */
     private long nonPooledAvailableMemory;
     private final Metrics metrics;
     private final Time time;
     private final Sensor waitTime;
+    private boolean closed;
+    private long nextOvermemoryWarn;
 
     /**
      * Create a new buffer pool
@@ -68,7 +78,7 @@ public class BufferPool {
     public BufferPool(long memory, int poolableSize, Metrics metrics, Time time, String metricGrpName) {
         this.poolableSize = poolableSize;
         this.lock = new ReentrantLock();
-        this.free = new ArrayDeque<>();
+        this.free = new LinkedHashSet<>();
         this.waiters = new ArrayDeque<>();
         this.totalMemory = memory;
         this.nonPooledAvailableMemory = memory;
@@ -82,6 +92,8 @@ public BufferPool(long memory, int poolableSize, Metrics metrics, Time time, Str
                                                    metricGrpName,
                                                    "The total time an appender waits for space allocation.");
         this.waitTime.add(new Meter(TimeUnit.NANOSECONDS, rateMetricName, totalMetricName));
+        this.closed = false;
+        this.nextOvermemoryWarn = 0;
     }
 
     /**
@@ -104,10 +116,16 @@ public ByteBuffer allocate(int size, long maxTimeToBlockMs) throws InterruptedEx
 
         ByteBuffer buffer = null;
         this.lock.lock();
+
+        if (this.closed) {
+            this.lock.unlock();
+            throw new KafkaException("Producer closed while allocating memory");
+        }
+
         try {
             // check if we have a free buffer of the right size pooled
             if (size == poolableSize && !this.free.isEmpty())
-                return this.free.pollFirst();
+                return popBuffer();
 
             // now check if the request is immediately satisfiable with the
             // memory on hand or if we need to block
@@ -138,6 +156,9 @@ public ByteBuffer allocate(int size, long maxTimeToBlockMs) throws InterruptedEx
                             recordWaitTime(timeNs);
                         }
 
+                        if (this.closed)
+                            throw new KafkaException("Producer closed while allocating memory");
+
                         if (waitingTimeElapsed) {
                             throw new TimeoutException("Failed to allocate memory within the configured max blocking time " + maxTimeToBlockMs + " ms.");
                         }
@@ -148,7 +169,7 @@ public ByteBuffer allocate(int size, long maxTimeToBlockMs) throws InterruptedEx
                         // otherwise allocate memory
                         if (accumulated == 0 && size == this.poolableSize && !this.free.isEmpty()) {
                             // just grab a buffer from the free list
-                            buffer = this.free.pollFirst();
+                            buffer = popBuffer();
                             accumulated = size;
                         } else {
                             // we'll need to allocate memory, but we may only get
@@ -225,25 +246,36 @@ protected ByteBuffer allocateByteBuffer(int size) {
      */
     private void freeUp(int size) {
         while (!this.free.isEmpty() && this.nonPooledAvailableMemory < size)
-            this.nonPooledAvailableMemory += this.free.pollLast().capacity();
+            this.nonPooledAvailableMemory += popBuffer().capacity();
     }
 
     /**
      * Return buffers to the pool. If they are of the poolable size add them to the free list, otherwise just mark the
      * memory as free.
      *
-     * @param buffer The buffer to return
+     * @param buffer The buffer to return to the pool.
      * @param size The size of the buffer to mark as deallocated, note that this may be smaller than buffer.capacity
-     *             since the buffer may re-allocate itself during in-place compression
+     *          since the buffer may re-allocate itself during in-place compression
      */
     public void deallocate(ByteBuffer buffer, int size) {
         lock.lock();
         try {
-            if (size == this.poolableSize && size == buffer.capacity()) {
-                buffer.clear();
-                this.free.add(buffer);
+            long availableMemory = availableMemoryUnlocked();
+            if (availableMemory + size > this.totalMemory && nextOvermemoryWarn < this.time.milliseconds()) {
+                //Don't flood logs.
+                log.error("Detected an attempt to bring available memory " + availableMemory + " to " +
+                    (availableMemory + size) + " which is higher than totalMemory " + this.totalMemory + ".", new Exception());
+                this.nextOvermemoryWarn = this.time.milliseconds() + TimeUnit.HOURS.toMillis(1);
+            }
+            if (buffer.hasArray() && buffer.array().length == this.poolableSize &&
+                this.poolableSize + availableMemory <= this.totalMemory) {
+                if (!this.free.add(buffer.array()) && this.nextOvermemoryWarn < time.milliseconds()) {
+                    log.error("Detected an attempt to double deallocate the same buffer.", new Exception());
+                    this.nextOvermemoryWarn = this.time.milliseconds() + TimeUnit.HOURS.toMillis(1);
+                }
             } else {
-                this.nonPooledAvailableMemory += size;
+                long freeMem = Math.min(size, this.totalMemory - availableMemory);
+                this.nonPooledAvailableMemory += freeMem;
             }
             Condition moreMem = this.waiters.peekFirst();
             if (moreMem != null)
@@ -263,12 +295,16 @@ public void deallocate(ByteBuffer buffer) {
     public long availableMemory() {
         lock.lock();
         try {
-            return this.nonPooledAvailableMemory + freeSize() * (long) this.poolableSize;
+            return availableMemoryUnlocked();
         } finally {
             lock.unlock();
         }
     }
 
+    private long availableMemoryUnlocked() {
+        return this.nonPooledAvailableMemory + freeSize() * (long) this.poolableSize;
+    }
+
     // Protected for testing.
     protected int freeSize() {
         return this.free.size();
@@ -316,4 +352,26 @@ public long totalMemory() {
     Deque waiters() {
         return this.waiters;
     }
+
+    /**
+     * Closes the buffer pool. Memory will be prevented from being allocated, but may be deallocated. All allocations
+     * awaiting available memory will be notified to abort.
+     */
+    public void close() {
+        this.lock.lock();
+        this.closed = true;
+        try {
+            for (Condition waiter : this.waiters)
+                waiter.signal();
+        } finally {
+            this.lock.unlock();
+        }
+    }
+
+    private ByteBuffer popBuffer() {
+        Iterator it = free.iterator();
+        byte[] array = it.next();
+        it.remove();
+        return ByteBuffer.wrap(array);
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/DefaultPartitioner.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/DefaultPartitioner.java
index 9d4ecbf151372..85c5e4eb9f216 100644
--- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/DefaultPartitioner.java
+++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/DefaultPartitioner.java
@@ -18,10 +18,6 @@
 
 import java.util.List;
 import java.util.Map;
-import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.ConcurrentMap;
-import java.util.concurrent.ThreadLocalRandom;
-import java.util.concurrent.atomic.AtomicInteger;
 
 import org.apache.kafka.clients.producer.Partitioner;
 import org.apache.kafka.common.Cluster;
@@ -33,11 +29,13 @@
  * 
        
  * 
      • If a partition is specified in the record, use it
  * 
      • If no partition is specified but a key is present choose a partition based on a hash of the key
- * 
      • If no partition or key is present choose a partition in a round-robin fashion
+ * 
      • If no partition or key is present choose the sticky partition that changes when the batch is full.
+ * 
+ * See KIP-480 for details about sticky partitioning.
  */
 public class DefaultPartitioner implements Partitioner {
 
-    private final ConcurrentMap topicCounterMap = new ConcurrentHashMap<>();
+    private final StickyPartitionCache stickyPartitionCache = new StickyPartitionCache();
 
     public void configure(Map configs) {}
 
@@ -52,36 +50,22 @@ public void configure(Map configs) {}
      * @param cluster The current cluster metadata
      */
     public int partition(String topic, Object key, byte[] keyBytes, Object value, byte[] valueBytes, Cluster cluster) {
+        if (keyBytes == null) {
+            return stickyPartitionCache.partition(topic, cluster);
+        } 
         List partitions = cluster.partitionsForTopic(topic);
         int numPartitions = partitions.size();
-        if (keyBytes == null) {
-            int nextValue = nextValue(topic);
-            List availablePartitions = cluster.availablePartitionsForTopic(topic);
-            if (availablePartitions.size() > 0) {
-                int part = Utils.toPositive(nextValue) % availablePartitions.size();
-                return availablePartitions.get(part).partition();
-            } else {
-                // no partitions are available, give a non-available partition
-                return Utils.toPositive(nextValue) % numPartitions;
-            }
-        } else {
-            // hash the keyBytes to choose a partition
-            return Utils.toPositive(Utils.murmur2(keyBytes)) % numPartitions;
-        }
-    }
-
-    private int nextValue(String topic) {
-        AtomicInteger counter = topicCounterMap.get(topic);
-        if (null == counter) {
-            counter = new AtomicInteger(ThreadLocalRandom.current().nextInt());
-            AtomicInteger currentCounter = topicCounterMap.putIfAbsent(topic, counter);
-            if (currentCounter != null) {
-                counter = currentCounter;
-            }
-        }
-        return counter.getAndIncrement();
+        // hash the keyBytes to choose a partition
+        return Utils.toPositive(Utils.murmur2(keyBytes)) % numPartitions;
     }
 
     public void close() {}
-
+    
+    /**
+     * If a batch completed for the current sticky partition, change the sticky partition. 
+     * Alternately, if no sticky partition has been determined, set one.
+     */
+    public void onNewBatch(String topic, Cluster cluster, int prevPartition) {
+        stickyPartitionCache.nextPartition(topic, cluster, prevPartition);
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ProducerMetadata.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ProducerMetadata.java
index 90e7970492c07..8759be649e842 100644
--- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ProducerMetadata.java
+++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ProducerMetadata.java
@@ -18,7 +18,6 @@
 
 import org.apache.kafka.clients.Metadata;
 import org.apache.kafka.common.KafkaException;
-import org.apache.kafka.common.errors.AuthenticationException;
 import org.apache.kafka.common.internals.ClusterResourceListeners;
 import org.apache.kafka.common.requests.MetadataRequest;
 import org.apache.kafka.common.requests.MetadataResponse;
@@ -34,26 +33,44 @@
 
 public class ProducerMetadata extends Metadata {
     private static final long TOPIC_EXPIRY_NEEDS_UPDATE = -1L;
-    static final long TOPIC_EXPIRY_MS = 5 * 60 * 1000;
+    public static final long TOPIC_EXPIRY_MS = 5 * 60 * 1000;
 
+    private final boolean allowAutoTopicCreation;
     /* Topics with expiry time */
     private final Map topics = new HashMap<>();
     private final Logger log;
     private final Time time;
+    private final long topicExpiryMs;
+
+    // LI-HOTFIX: this constructor should only be used for unit tests
+    // after the following hotfix changes:
+    // 1) add metadata.topic.expiry.ms config to KafkaProducer
+    // 2) Make client-side auto.topic.creation configurable and default to be false
+    public ProducerMetadata(long refreshBackoffMs,
+        long metadataExpireMs,
+        LogContext logContext,
+        ClusterResourceListeners clusterResourceListeners,
+        Time time) {
+        this(refreshBackoffMs, metadataExpireMs, logContext, clusterResourceListeners, time, TOPIC_EXPIRY_MS, true);
+    }
 
     public ProducerMetadata(long refreshBackoffMs,
                             long metadataExpireMs,
                             LogContext logContext,
                             ClusterResourceListeners clusterResourceListeners,
-                            Time time) {
+                            Time time,
+                            long topicExpiryMs,
+                            boolean allowAutoTopicCreation) {
         super(refreshBackoffMs, metadataExpireMs, logContext, clusterResourceListeners);
         this.log = logContext.logger(ProducerMetadata.class);
         this.time = time;
+        this.topicExpiryMs = topicExpiryMs;
+        this.allowAutoTopicCreation = allowAutoTopicCreation;
     }
 
     @Override
     public synchronized MetadataRequest.Builder newMetadataRequestBuilder() {
-        return new MetadataRequest.Builder(new ArrayList<>(topics.keySet()), true);
+        return new MetadataRequest.Builder(new ArrayList<>(topics.keySet()), allowAutoTopicCreation);
     }
 
     public synchronized void add(String topic) {
@@ -78,7 +95,7 @@ public synchronized boolean retainTopic(String topic, boolean isInternal, long n
         if (expireMs == null) {
             return false;
         } else if (expireMs == TOPIC_EXPIRY_NEEDS_UPDATE) {
-            topics.put(topic, nowMs + TOPIC_EXPIRY_MS);
+            topics.put(topic, Long.MAX_VALUE - nowMs > topicExpiryMs ? nowMs + topicExpiryMs : Long.MAX_VALUE);
             return true;
         } else if (expireMs <= nowMs) {
             log.debug("Removing unused topic {} from the metadata list, expiryMs {} now {}", topic, expireMs, nowMs);
@@ -96,7 +113,8 @@ public synchronized void awaitUpdate(final int lastVersion, final long timeoutMs
         long currentTimeMs = time.milliseconds();
         long deadlineMs = currentTimeMs + timeoutMs < 0 ? Long.MAX_VALUE : currentTimeMs + timeoutMs;
         time.waitObject(this, () -> {
-            maybeThrowException();
+            // Throw fatal exceptions, if there are any. Recoverable topic errors will be handled by the caller.
+            maybeThrowFatalException();
             return updateVersion() > lastVersion || isClosed();
         }, deadlineMs);
 
@@ -111,10 +129,9 @@ public synchronized void update(int requestVersion, MetadataResponse response, l
     }
 
     @Override
-    public synchronized void failedUpdate(long now, AuthenticationException authenticationException) {
-        super.failedUpdate(now, authenticationException);
-        if (authenticationException != null)
-            notifyAll();
+    public synchronized void fatalError(KafkaException fatalException) {
+        super.fatalError(fatalException);
+        notifyAll();
     }
 
     /**
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 80a9d0c2bd870..49a59f5be5e0c 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
@@ -25,9 +25,11 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.TimeUnit;
 import org.apache.kafka.clients.ApiVersions;
 import org.apache.kafka.clients.producer.Callback;
 import org.apache.kafka.common.Cluster;
@@ -36,6 +38,8 @@
 import org.apache.kafka.common.Node;
 import org.apache.kafka.common.PartitionInfo;
 import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.errors.RecordBatchTooLargeException;
+import org.apache.kafka.common.errors.TimeoutException;
 import org.apache.kafka.common.errors.UnsupportedVersionException;
 import org.apache.kafka.common.header.Header;
 import org.apache.kafka.common.metrics.Measurable;
@@ -54,7 +58,6 @@
 import org.apache.kafka.common.utils.CopyOnWriteMap;
 import org.apache.kafka.common.utils.LogContext;
 import org.apache.kafka.common.utils.Time;
-import org.apache.kafka.common.utils.Utils;
 import org.slf4j.Logger;
 
 /**
@@ -72,9 +75,9 @@ public final class RecordAccumulator {
     private final AtomicInteger appendsInProgress;
     private final int batchSize;
     private final CompressionType compression;
-    private final long lingerMs;
+    private final int lingerMs;
     private final long retryBackoffMs;
-    private final long deliveryTimeoutMs;
+    private final int deliveryTimeoutMs;
     private final BufferPool free;
     private final Time time;
     private final ApiVersions apiVersions;
@@ -106,9 +109,9 @@ public final class RecordAccumulator {
     public RecordAccumulator(LogContext logContext,
                              int batchSize,
                              CompressionType compression,
-                             long lingerMs,
+                             int lingerMs,
                              long retryBackoffMs,
-                             long deliveryTimeoutMs,
+                             int deliveryTimeoutMs,
                              Metrics metrics,
                              String metricGrpName,
                              Time time,
@@ -179,6 +182,8 @@ public double measure(MetricConfig config, long now) {
      * @param headers the Headers for the record
      * @param callback The user-supplied callback to execute when the request is complete
      * @param maxTimeToBlock The maximum time in milliseconds to block for buffer memory to be available
+     * @param abortOnNewBatch A boolean that indicates returning before a new batch is created and 
+     *                        running the the partitioner's onNewBatch method before trying to append again
      */
     public RecordAppendResult append(TopicPartition tp,
                                      long timestamp,
@@ -186,7 +191,8 @@ public RecordAppendResult append(TopicPartition tp,
                                      byte[] value,
                                      Header[] headers,
                                      Callback callback,
-                                     long maxTimeToBlock) throws InterruptedException {
+                                     long maxTimeToBlock,
+                                     boolean abortOnNewBatch) throws InterruptedException {
         // We keep track of the number of appending thread to make sure we do not miss batches in
         // abortIncompleteBatches().
         appendsInProgress.incrementAndGet();
@@ -204,6 +210,11 @@ public RecordAppendResult append(TopicPartition tp,
             }
 
             // we don't have an in-progress record batch try to allocate a new batch
+            if (abortOnNewBatch) {
+                // Return a result that will cause another call to append.
+                return new RecordAppendResult(null, false, false, true);
+            }
+            
             byte maxUsableMagic = apiVersions.maxUsableProduceMagic();
             int size = Math.max(this.batchSize, AbstractRecords.estimateSizeInBytesUpperBound(maxUsableMagic, compression, key, value, headers));
             log.trace("Allocating a new {} byte message buffer for topic {} partition {}", size, tp.topic(), tp.partition());
@@ -221,14 +232,15 @@ public RecordAppendResult append(TopicPartition tp,
 
                 MemoryRecordsBuilder recordsBuilder = recordsBuilder(buffer, maxUsableMagic);
                 ProducerBatch batch = new ProducerBatch(tp, recordsBuilder, time.milliseconds());
-                FutureRecordMetadata future = Utils.notNull(batch.tryAppend(timestamp, key, value, headers, callback, time.milliseconds()));
+                FutureRecordMetadata future = Objects.requireNonNull(batch.tryAppend(timestamp, key, value, headers,
+                        callback, time.milliseconds()));
 
                 dq.addLast(batch);
                 incomplete.add(batch);
 
                 // Don't deallocate this buffer in the finally block as it's being used in the record batch
                 buffer = null;
-                return new RecordAppendResult(future, dq.size() > 1 || batch.isFull(), true);
+                return new RecordAppendResult(future, dq.size() > 1 || batch.isFull(), true, false);
             }
         } finally {
             if (buffer != null)
@@ -261,16 +273,24 @@ private RecordAppendResult tryAppend(long timestamp, byte[] key, byte[] value, H
             if (future == null)
                 last.closeForRecordAppends();
             else
-                return new RecordAppendResult(future, deque.size() > 1 || last.isFull(), false);
+                return new RecordAppendResult(future, deque.size() > 1 || last.isFull(), false, false);
         }
         return null;
     }
 
     private boolean isMuted(TopicPartition tp, long now) {
-        boolean result = muted.containsKey(tp) && muted.get(tp) > now;
-        if (!result)
+        // Take care to avoid unnecessary map look-ups because this method is a hotspot if producing to a
+        // large number of partitions
+        Long throttleUntilTime = muted.get(tp);
+        if (throttleUntilTime == null)
+            return false;
+
+        if (now >= throttleUntilTime) {
             muted.remove(tp);
-        return result;
+            return false;
+        }
+        
+        return true;
     }
 
     public void resetNextBatchExpiryTime() {
@@ -442,18 +462,19 @@ public ReadyCheckResult ready(Cluster cluster, long nowMs) {
 
         boolean exhausted = this.free.queued() > 0;
         for (Map.Entry> entry : this.batches.entrySet()) {
-            TopicPartition part = entry.getKey();
             Deque deque = entry.getValue();
-
-            Node leader = cluster.leaderFor(part);
             synchronized (deque) {
-                if (leader == null && !deque.isEmpty()) {
-                    // This is a partition for which leader is not known, but messages are available to send.
-                    // Note that entries are currently not removed from batches when deque is empty.
-                    unknownLeaderTopics.add(part.topic());
-                } else if (!readyNodes.contains(leader) && !isMuted(part, nowMs)) {
-                    ProducerBatch batch = deque.peekFirst();
-                    if (batch != null) {
+                // When producing to a large number of partitions, this path is hot and deques are often empty.
+                // We check whether a batch exists first to avoid the more expensive checks whenever possible.
+                ProducerBatch batch = deque.peekFirst();
+                if (batch != null) {
+                    TopicPartition part = entry.getKey();
+                    Node leader = cluster.leaderFor(part);
+                    if (leader == null) {
+                        // This is a partition for which leader is not known, but messages are available to send.
+                        // Note that entries are currently not removed from batches when deque is empty.
+                        unknownLeaderTopics.add(part.topic());
+                    } else if (!readyNodes.contains(leader) && !isMuted(part, nowMs)) {
                         long waitedTimeMs = batch.waitedTimeMs(nowMs);
                         boolean backingOff = batch.attempts() > 0 && waitedTimeMs < retryBackoffMs;
                         long timeToWaitMs = backingOff ? retryBackoffMs : lingerMs;
@@ -558,7 +579,7 @@ private List drainBatchesForOneNode(Cluster cluster, Node node, i
                     if (shouldStopDrainBatchesForPartition(first, tp))
                         break;
 
-                    boolean isTransactional = transactionManager != null ? transactionManager.isTransactional() : false;
+                    boolean isTransactional = transactionManager != null && transactionManager.isTransactional();
                     ProducerIdAndEpoch producerIdAndEpoch =
                         transactionManager != null ? transactionManager.producerIdAndEpoch() : null;
                     ProducerBatch batch = deque.pollFirst();
@@ -665,6 +686,13 @@ boolean flushInProgress() {
         return flushesInProgress.get() > 0;
     }
 
+    /**
+     * This method should be used only for testing.
+     */
+    IncompleteBatches incompleteBatches() {
+        return incomplete;
+    }
+
     /* Visible for testing */
     Map> batches() {
         return Collections.unmodifiableMap(batches);
@@ -685,12 +713,33 @@ private boolean appendsInProgress() {
     }
 
     /**
-     * Mark all partitions as ready to send and block until the send is complete
+     * Mark all partitions as ready to send and block until the send is complete or time expires
      */
-    public void awaitFlushCompletion() throws InterruptedException {
+    public void awaitFlushCompletion(long timeoutMs) throws InterruptedException {
         try {
-            for (ProducerBatch batch : this.incomplete.copyAll())
-                batch.produceFuture.await();
+            boolean retry;
+            Long expireMs = System.currentTimeMillis() + timeoutMs;
+            do {
+                retry = false;
+                for (ProducerBatch batch : this.incomplete.copyAll()) {
+                    Long currentMs = System.currentTimeMillis();
+                    if (currentMs > expireMs) {
+                        throw new TimeoutException("Failed to flush accumulated records within" + timeoutMs + "milliseconds.");
+                    }
+                    boolean completed = batch.produceFuture.await(Math.max(expireMs - currentMs, 0), TimeUnit.MILLISECONDS);
+                    if (!completed) {
+                        throw new TimeoutException("Failed to flush accumulated records within" + timeoutMs + "milliseconds.");
+                    }
+                    // If the produceFuture failed with RecordBatchTooLargeException, it means that the
+                    // batch was split into smaller batches and re-enqueued into the RecordAccumulator by Sender thread.
+                    // This if condition will make sure to retry and send all the split batches.
+                    // Note that, More records get sent to the broker than necessary because the retry mechanism
+                    // will also include all the newly added records via kafkaProducer.send() api.
+                    if (batch.produceFuture.error() instanceof RecordBatchTooLargeException) {
+                        retry = true;
+                    }
+                }
+            } while (retry);
         } finally {
             this.flushesInProgress.decrementAndGet();
         }
@@ -778,6 +827,7 @@ public void unmutePartition(TopicPartition tp, long throttleUntilTimeMs) {
      */
     public void close() {
         this.closed = true;
+        this.free.close();
     }
 
     /*
@@ -787,11 +837,13 @@ public final static class RecordAppendResult {
         public final FutureRecordMetadata future;
         public final boolean batchIsFull;
         public final boolean newBatchCreated;
+        public final boolean abortForNewBatch;
 
-        public RecordAppendResult(FutureRecordMetadata future, boolean batchIsFull, boolean newBatchCreated) {
+        public RecordAppendResult(FutureRecordMetadata future, boolean batchIsFull, boolean newBatchCreated, boolean abortForNewBatch) {
             this.future = future;
             this.batchIsFull = batchIsFull;
             this.newBatchCreated = newBatchCreated;
+            this.abortForNewBatch = abortForNewBatch;
         }
     }
 
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 6189aae173643..ad2e1eaaa60cf 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
@@ -16,7 +16,6 @@
  */
 package org.apache.kafka.clients.producer.internals;
 
-import java.util.ArrayList;
 import org.apache.kafka.clients.ApiVersions;
 import org.apache.kafka.clients.ClientRequest;
 import org.apache.kafka.clients.ClientResponse;
@@ -33,13 +32,12 @@
 import org.apache.kafka.common.errors.ClusterAuthorizationException;
 import org.apache.kafka.common.errors.InvalidMetadataException;
 import org.apache.kafka.common.errors.OutOfOrderSequenceException;
-import org.apache.kafka.common.errors.ProducerFencedException;
 import org.apache.kafka.common.errors.RetriableException;
 import org.apache.kafka.common.errors.TimeoutException;
 import org.apache.kafka.common.errors.TopicAuthorizationException;
-import org.apache.kafka.common.errors.TransactionalIdAuthorizationException;
 import org.apache.kafka.common.errors.UnknownTopicOrPartitionException;
 import org.apache.kafka.common.errors.UnsupportedVersionException;
+import org.apache.kafka.common.message.InitProducerIdRequestData;
 import org.apache.kafka.common.metrics.Sensor;
 import org.apache.kafka.common.metrics.stats.Avg;
 import org.apache.kafka.common.metrics.stats.Max;
@@ -48,6 +46,7 @@
 import org.apache.kafka.common.record.MemoryRecords;
 import org.apache.kafka.common.record.RecordBatch;
 import org.apache.kafka.common.requests.AbstractRequest;
+import org.apache.kafka.common.requests.FindCoordinatorRequest;
 import org.apache.kafka.common.requests.InitProducerIdRequest;
 import org.apache.kafka.common.requests.InitProducerIdResponse;
 import org.apache.kafka.common.requests.ProduceRequest;
@@ -55,15 +54,16 @@
 import org.apache.kafka.common.requests.RequestHeader;
 import org.apache.kafka.common.utils.LogContext;
 import org.apache.kafka.common.utils.Time;
-import org.apache.kafka.common.utils.Utils;
 import org.slf4j.Logger;
 
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 
 import static org.apache.kafka.common.record.RecordBatch.NO_TIMESTAMP;
 
@@ -159,7 +159,7 @@ public List inFlightBatches(TopicPartition tp) {
         return inFlightBatches.containsKey(tp) ? inFlightBatches.get(tp) : new ArrayList<>();
     }
 
-    public void maybeRemoveFromInflightBatches(ProducerBatch batch) {
+    private void maybeRemoveFromInflightBatches(ProducerBatch batch) {
         List batches = inFlightBatches.get(batch.topicPartition);
         if (batches != null) {
             batches.remove(batch);
@@ -169,6 +169,11 @@ public void maybeRemoveFromInflightBatches(ProducerBatch batch) {
         }
     }
 
+    private void maybeRemoveAndDeallocateBatch(ProducerBatch batch) {
+        maybeRemoveFromInflightBatches(batch);
+        this.accumulator.deallocate(batch);
+    }
+
     /**
      *  Get the in-flight batches that has reached delivery timeout.
      */
@@ -223,6 +228,10 @@ public void addToInflightBatches(Map> batches) {
         }
     }
 
+    private boolean hasPendingTransactionalRequests() {
+        return transactionManager != null && transactionManager.hasPendingRequests() && transactionManager.hasOngoingTransaction();
+    }
+
     /**
      * The main run loop for the sender thread
      */
@@ -232,7 +241,7 @@ public void run() {
         // main loop, runs until close is called
         while (running) {
             try {
-                run(time.milliseconds());
+                runOnce();
             } catch (Exception e) {
                 log.error("Uncaught error in kafka producer I/O thread: ", e);
             }
@@ -241,18 +250,36 @@ public void run() {
         log.debug("Beginning shutdown of Kafka producer I/O thread, sending remaining records.");
 
         // okay we stopped accepting requests but there may still be
-        // requests in the accumulator or waiting for acknowledgment,
+        // requests in the transaction manager, accumulator or waiting for acknowledgment,
         // wait until these are completed.
-        while (!forceClose && (this.accumulator.hasUndrained() || this.client.inFlightRequestCount() > 0)) {
+        while (!forceClose && ((this.accumulator.hasUndrained() || this.client.inFlightRequestCount() > 0) || hasPendingTransactionalRequests())) {
             try {
-                run(time.milliseconds());
+                runOnce();
             } catch (Exception e) {
                 log.error("Uncaught error in kafka producer I/O thread: ", e);
             }
         }
+
+        // Abort the transaction if any commit or abort didn't go through the transaction manager's queue
+        while (!forceClose && transactionManager != null && transactionManager.hasOngoingTransaction()) {
+            if (!transactionManager.isCompleting()) {
+                log.info("Aborting incomplete transaction due to shutdown");
+                transactionManager.beginAbort();
+            }
+            try {
+                runOnce();
+            } catch (Exception e) {
+                log.error("Uncaught error in kafka producer I/O thread: ", e);
+            }
+        }
+
         if (forceClose) {
-            // We need to fail all the incomplete batches and wake up the threads waiting on
+            // We need to fail all the incomplete transactional requests and batches and wake up the threads waiting on
             // the futures.
+            if (transactionManager != null) {
+                log.debug("Aborting incomplete transactional requests due to forced shutdown");
+                transactionManager.close();
+            }
             log.debug("Aborting incomplete batches due to forced shutdown");
             this.accumulator.abortIncompleteBatches();
         }
@@ -268,14 +295,12 @@ public void run() {
     /**
      * Run a single iteration of sending
      *
-     * @param now The current POSIX time in milliseconds
      */
-    void run(long now) {
+    void runOnce() {
         if (transactionManager != null) {
             try {
-                if (transactionManager.shouldResetProducerStateAfterResolvingSequences())
-                    // Check if the previous run expired batches which requires a reset of the producer state.
-                    transactionManager.resetProducerId();
+                transactionManager.resetProducerIdIfNeeded();
+
                 if (!transactionManager.isTransactional()) {
                     // this is an idempotent producer, so make sure we have a producer id
                     maybeWaitForProducerId();
@@ -283,9 +308,7 @@ void run(long now) {
                     transactionManager.transitionToFatalError(
                         new KafkaException("The client hasn't received acknowledgment for " +
                             "some previously sent messages and can no longer retry them. It isn't safe to continue."));
-                } else if (transactionManager.hasInFlightTransactionalRequest() || maybeSendTransactionalRequest(now)) {
-                    // as long as there are outstanding transactional requests, we simply wait for them to return
-                    client.poll(retryBackoffMs, now);
+                } else if (maybeSendAndPollTransactionalRequest()) {
                     return;
                 }
 
@@ -295,7 +318,7 @@ void run(long now) {
                     RuntimeException lastError = transactionManager.lastError();
                     if (lastError != null)
                         maybeAbortBatches(lastError);
-                    client.poll(retryBackoffMs, now);
+                    client.poll(retryBackoffMs, time.milliseconds());
                     return;
                 } else if (transactionManager.hasAbortableError()) {
                     accumulator.abortUndrainedBatches(transactionManager.lastError());
@@ -307,8 +330,9 @@ void run(long now) {
             }
         }
 
-        long pollTimeout = sendProducerData(now);
-        client.poll(pollTimeout, now);
+        long currentTimeMs = time.milliseconds();
+        long pollTimeout = sendProducerData(currentTimeMs);
+        client.poll(pollTimeout, currentTimeMs);
     }
 
     private long sendProducerData(long now) {
@@ -392,7 +416,16 @@ private long sendProducerData(long now) {
         return pollTimeout;
     }
 
-    private boolean maybeSendTransactionalRequest(long now) {
+    /**
+     * Returns true if a transactional request is sent or polled, or if a FindCoordinator request is enqueued
+     */
+    private boolean maybeSendAndPollTransactionalRequest() {
+        if (transactionManager.hasInFlightTransactionalRequest()) {
+            // as long as there are outstanding transactional requests, we simply wait for them to return
+            client.poll(retryBackoffMs, time.milliseconds());
+            return true;
+        }
+
         if (transactionManager.isCompleting() && accumulator.hasIncomplete()) {
             if (transactionManager.isAborting())
                 accumulator.abortUndrainedBatches(new KafkaException("Failing batch since transaction was aborted"));
@@ -409,47 +442,43 @@ private boolean maybeSendTransactionalRequest(long now) {
             return false;
 
         AbstractRequest.Builder requestBuilder = nextRequestHandler.requestBuilder();
-        while (!forceClose) {
-            Node targetNode = null;
-            try {
-                if (nextRequestHandler.needsCoordinator()) {
-                    targetNode = transactionManager.coordinator(nextRequestHandler.coordinatorType());
-                    if (targetNode == null) {
-                        transactionManager.lookupCoordinator(nextRequestHandler);
-                        break;
-                    }
-                    if (!NetworkClientUtils.awaitReady(client, targetNode, time, requestTimeoutMs)) {
-                        transactionManager.lookupCoordinator(nextRequestHandler);
-                        break;
-                    }
-                } else {
-                    targetNode = awaitLeastLoadedNodeReady(requestTimeoutMs);
-                }
-
-                if (targetNode != null) {
-                    if (nextRequestHandler.isRetry())
-                        time.sleep(nextRequestHandler.retryBackoffMs());
-                    ClientRequest clientRequest = client.newClientRequest(
-                        targetNode.idString(), requestBuilder, now, true, requestTimeoutMs, nextRequestHandler);
-                    transactionManager.setInFlightTransactionalRequestCorrelationId(clientRequest.correlationId());
-                    log.debug("Sending transactional request {} to node {}", requestBuilder, targetNode);
-                    client.send(clientRequest, now);
-                    return true;
-                }
-            } catch (IOException e) {
-                log.debug("Disconnect from {} while trying to send request {}. Going " +
-                        "to back off and retry.", targetNode, requestBuilder, e);
-                if (nextRequestHandler.needsCoordinator()) {
-                    // We break here so that we pick up the FindCoordinator request immediately.
-                    transactionManager.lookupCoordinator(nextRequestHandler);
-                    break;
-                }
+        Node targetNode = null;
+        try {
+            targetNode = awaitNodeReady(nextRequestHandler.coordinatorType());
+            if (targetNode == null) {
+                lookupCoordinatorAndRetry(nextRequestHandler);
+                return true;
             }
+
+            if (nextRequestHandler.isRetry())
+                time.sleep(nextRequestHandler.retryBackoffMs());
+            long currentTimeMs = time.milliseconds();
+            ClientRequest clientRequest = client.newClientRequest(
+                targetNode.idString(), requestBuilder, currentTimeMs, true, requestTimeoutMs, nextRequestHandler);
+            log.debug("Sending transactional request {} to node {}", requestBuilder, targetNode);
+            client.send(clientRequest, currentTimeMs);
+            transactionManager.setInFlightCorrelationId(clientRequest.correlationId());
+            client.poll(retryBackoffMs, time.milliseconds());
+            return true;
+        } catch (IOException e) {
+            log.debug("Disconnect from {} while trying to send request {}. Going " +
+                    "to back off and retry.", targetNode, requestBuilder, e);
+            // We break here so that we pick up the FindCoordinator request immediately.
+            lookupCoordinatorAndRetry(nextRequestHandler);
+            return true;
+        }
+    }
+
+    private void lookupCoordinatorAndRetry(TransactionManager.TxnRequestHandler nextRequestHandler) {
+        if (nextRequestHandler.needsCoordinator()) {
+            transactionManager.lookupCoordinator(nextRequestHandler);
+        } else {
+            // For non-coordinator requests, sleep here to prevent a tight loop when no node is available
             time.sleep(retryBackoffMs);
             metadata.requestUpdate();
         }
+
         transactionManager.retry(nextRequestHandler);
-        return true;
     }
 
     private void maybeAbortBatches(RuntimeException exception) {
@@ -478,16 +507,26 @@ public void forceClose() {
         initiateClose();
     }
 
+    public boolean isRunning() {
+        return running;
+    }
+
     private ClientResponse sendAndAwaitInitProducerIdRequest(Node node) throws IOException {
         String nodeId = node.idString();
-        InitProducerIdRequest.Builder builder = new InitProducerIdRequest.Builder(null);
+        InitProducerIdRequestData requestData = new InitProducerIdRequestData()
+                .setTransactionalId(null)
+                .setTransactionTimeoutMs(Integer.MAX_VALUE);
+        InitProducerIdRequest.Builder builder = new InitProducerIdRequest.Builder(requestData);
         ClientRequest request = client.newClientRequest(nodeId, builder, time.milliseconds(), true, requestTimeoutMs, null);
         return NetworkClientUtils.sendAndReceive(client, request, time);
     }
 
-    private Node awaitLeastLoadedNodeReady(long remainingTimeMs) throws IOException {
-        Node node = client.leastLoadedNode(time.milliseconds());
-        if (node != null && NetworkClientUtils.awaitReady(client, node, time, remainingTimeMs)) {
+    private Node awaitNodeReady(FindCoordinatorRequest.CoordinatorType coordinatorType) throws IOException {
+        Node node = coordinatorType != null ?
+                transactionManager.coordinator(coordinatorType) :
+                client.leastLoadedNode(time.milliseconds());
+
+        if (node != null && NetworkClientUtils.awaitReady(client, node, time, requestTimeoutMs)) {
             return node;
         }
         return null;
@@ -497,14 +536,14 @@ private void maybeWaitForProducerId() {
         while (!forceClose && !transactionManager.hasProducerId() && !transactionManager.hasError()) {
             Node node = null;
             try {
-                node = awaitLeastLoadedNodeReady(requestTimeoutMs);
+                node = awaitNodeReady(null);
                 if (node != null) {
                     ClientResponse response = sendAndAwaitInitProducerIdRequest(node);
                     InitProducerIdResponse initProducerIdResponse = (InitProducerIdResponse) response.responseBody();
                     Errors error = initProducerIdResponse.error();
                     if (error == Errors.NONE) {
                         ProducerIdAndEpoch producerIdAndEpoch = new ProducerIdAndEpoch(
-                                initProducerIdResponse.producerId(), initProducerIdResponse.epoch());
+                                initProducerIdResponse.data.producerId(), initProducerIdResponse.data.producerEpoch());
                         transactionManager.setProducerIdAndEpoch(producerIdAndEpoch);
                         return;
                     } else if (error.exception() instanceof RetriableException) {
@@ -591,7 +630,7 @@ private void completeBatch(ProducerBatch batch, ProduceResponse.PartitionRespons
             if (transactionManager != null)
                 transactionManager.removeInFlightBatch(batch);
             this.accumulator.splitAndReenqueue(batch);
-            this.accumulator.deallocate(batch);
+            maybeRemoveAndDeallocateBatch(batch);
             this.sensors.recordBatchSplit();
         } else if (error != Errors.NONE) {
             if (canRetry(batch, response, now)) {
@@ -624,7 +663,7 @@ private void completeBatch(ProducerBatch batch, ProduceResponse.PartitionRespons
             } else {
                 final RuntimeException exception;
                 if (error == Errors.TOPIC_AUTHORIZATION_FAILED)
-                    exception = new TopicAuthorizationException(batch.topicPartition.topic());
+                    exception = new TopicAuthorizationException(Collections.singleton(batch.topicPartition.topic()));
                 else if (error == Errors.CLUSTER_AUTHORIZATION_FAILED)
                     exception = new ClusterAuthorizationException("The producer is not authorized to do idempotent sends");
                 else
@@ -662,61 +701,34 @@ private void reenqueueBatch(ProducerBatch batch, long currentTimeMs) {
 
     private void completeBatch(ProducerBatch batch, ProduceResponse.PartitionResponse response) {
         if (transactionManager != null) {
-            if (transactionManager.hasProducerIdAndEpoch(batch.producerId(), batch.producerEpoch())) {
-                transactionManager
-                    .maybeUpdateLastAckedSequence(batch.topicPartition, batch.baseSequence() + batch.recordCount - 1);
-                log.debug("ProducerId: {}; Set last ack'd sequence number for topic-partition {} to {}",
-                    batch.producerId(),
-                    batch.topicPartition,
-                    transactionManager.lastAckedSequence(batch.topicPartition).orElse(-1));
-            }
-            transactionManager.updateLastAckedOffset(response, batch);
-            transactionManager.removeInFlightBatch(batch);
+            transactionManager.handleCompletedBatch(batch, response);
         }
 
         if (batch.done(response.baseOffset, response.logAppendTime, null)) {
-            maybeRemoveFromInflightBatches(batch);
-            this.accumulator.deallocate(batch);
+            maybeRemoveAndDeallocateBatch(batch);
         }
     }
 
-    private void failBatch(ProducerBatch batch, ProduceResponse.PartitionResponse response, RuntimeException exception,
+    private void failBatch(ProducerBatch batch,
+                           ProduceResponse.PartitionResponse response,
+                           RuntimeException exception,
                            boolean adjustSequenceNumbers) {
         failBatch(batch, response.baseOffset, response.logAppendTime, exception, adjustSequenceNumbers);
     }
 
-    private void failBatch(ProducerBatch batch, long baseOffset, long logAppendTime, RuntimeException exception,
-        boolean adjustSequenceNumbers) {
+    private void failBatch(ProducerBatch batch,
+                           long baseOffset,
+                           long logAppendTime,
+                           RuntimeException exception,
+                           boolean adjustSequenceNumbers) {
         if (transactionManager != null) {
-            if (exception instanceof OutOfOrderSequenceException
-                    && !transactionManager.isTransactional()
-                    && transactionManager.hasProducerId(batch.producerId())) {
-                log.error("The broker returned {} for topic-partition " +
-                            "{} at offset {}. This indicates data loss on the broker, and should be investigated.",
-                        exception, batch.topicPartition, baseOffset);
-
-                // Reset the transaction state since we have hit an irrecoverable exception and cannot make any guarantees
-                // about the previously committed message. Note that this will discard the producer id and sequence
-                // numbers for all existing partitions.
-                transactionManager.resetProducerId();
-            } else if (exception instanceof ClusterAuthorizationException
-                    || exception instanceof TransactionalIdAuthorizationException
-                    || exception instanceof ProducerFencedException
-                    || exception instanceof UnsupportedVersionException) {
-                transactionManager.transitionToFatalError(exception);
-            } else if (transactionManager.isTransactional()) {
-                transactionManager.transitionToAbortableError(exception);
-            }
-            transactionManager.removeInFlightBatch(batch);
-            if (adjustSequenceNumbers)
-                transactionManager.adjustSequencesDueToFailedBatch(batch);
+            transactionManager.handleFailedBatch(batch, exception, adjustSequenceNumbers);
         }
 
         this.sensors.recordErrors(batch.topicPartition.topic(), batch.recordCount);
 
         if (batch.done(baseOffset, logAppendTime, exception)) {
-            maybeRemoveFromInflightBatches(batch);
-            this.accumulator.deallocate(batch);
+            maybeRemoveAndDeallocateBatch(batch);
         }
     }
 
@@ -914,17 +926,17 @@ public void updateProduceRequestMetrics(Map> batche
 
                     // per-topic record send rate
                     String topicRecordsCountName = "topic." + topic + ".records-per-batch";
-                    Sensor topicRecordCount = Utils.notNull(this.metrics.getSensor(topicRecordsCountName));
+                    Sensor topicRecordCount = Objects.requireNonNull(this.metrics.getSensor(topicRecordsCountName));
                     topicRecordCount.record(batch.recordCount);
 
                     // per-topic bytes send rate
                     String topicByteRateName = "topic." + topic + ".bytes";
-                    Sensor topicByteRate = Utils.notNull(this.metrics.getSensor(topicByteRateName));
+                    Sensor topicByteRate = Objects.requireNonNull(this.metrics.getSensor(topicByteRateName));
                     topicByteRate.record(batch.estimatedSizeInBytes());
 
                     // per-topic compression rate
                     String topicCompressionRateName = "topic." + topic + ".compression-rate";
-                    Sensor topicCompressionRate = Utils.notNull(this.metrics.getSensor(topicCompressionRateName));
+                    Sensor topicCompressionRate = Objects.requireNonNull(this.metrics.getSensor(topicCompressionRateName));
                     topicCompressionRate.record(batch.compressionRatio());
 
                     // global metrics
diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/StickyPartitionCache.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/StickyPartitionCache.java
new file mode 100644
index 0000000000000..bfe1ac6c6fc3d
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/StickyPartitionCache.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.clients.producer.internals;
+
+import org.apache.kafka.common.Cluster;
+import org.apache.kafka.common.PartitionInfo;
+
+import java.util.List;
+import java.util.concurrent.ThreadLocalRandom;
+import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.ConcurrentHashMap;
+import org.apache.kafka.common.utils.Utils;
+
+/**
+ * An internal class that implements a cache used for sticky partitioning behavior. The cache tracks the current sticky
+ * partition for any given topic. This class should not be used externally. 
+ */
+public class StickyPartitionCache {
+    private final ConcurrentMap indexCache;
+    public StickyPartitionCache() {
+        this.indexCache = new ConcurrentHashMap<>();
+    }
+
+    public int partition(String topic, Cluster cluster) {
+        Integer part = indexCache.get(topic);
+        if (part == null) {
+            return nextPartition(topic, cluster, -1);
+        }
+        return part;
+    }
+
+    public int nextPartition(String topic, Cluster cluster, int prevPartition) {
+        List partitions = cluster.partitionsForTopic(topic);
+        Integer oldPart = indexCache.get(topic);
+        Integer newPart = oldPart;
+        // Check that the current sticky partition for the topic is either not set or that the partition that 
+        // triggered the new batch matches the sticky partition that needs to be changed.
+        if (oldPart == null || oldPart == prevPartition) {
+            List availablePartitions = cluster.availablePartitionsForTopic(topic);
+            if (availablePartitions.size() < 1) {
+                Integer random = Utils.toPositive(ThreadLocalRandom.current().nextInt());
+                newPart = random % partitions.size();
+            } else if (availablePartitions.size() == 1) {
+                newPart = availablePartitions.get(0).partition();
+            } else {
+                while (newPart == null || newPart.equals(oldPart)) {
+                    Integer random = Utils.toPositive(ThreadLocalRandom.current().nextInt());
+                    newPart = availablePartitions.get(random % availablePartitions.size()).partition();
+                }
+            }
+            // Only change the sticky partition if it is null or prevPartition matches the current sticky partition.
+            if (oldPart == null) {
+                indexCache.putIfAbsent(topic, newPart);
+            } else {
+                indexCache.replace(topic, prevPartition, newPart);
+            }
+            return indexCache.get(topic);
+        }
+        return indexCache.get(topic);
+    }
+
+}
\ No newline at end of file
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 b6190932661e1..944eb8c5c381c 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
@@ -23,8 +23,16 @@
 import org.apache.kafka.common.Node;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.errors.AuthenticationException;
+import org.apache.kafka.common.errors.ClusterAuthorizationException;
 import org.apache.kafka.common.errors.GroupAuthorizationException;
+import org.apache.kafka.common.errors.OutOfOrderSequenceException;
+import org.apache.kafka.common.errors.ProducerFencedException;
 import org.apache.kafka.common.errors.TopicAuthorizationException;
+import org.apache.kafka.common.errors.TransactionalIdAuthorizationException;
+import org.apache.kafka.common.errors.UnsupportedVersionException;
+import org.apache.kafka.common.message.FindCoordinatorRequestData;
+import org.apache.kafka.common.message.InitProducerIdRequestData;
+import org.apache.kafka.common.message.TxnOffsetCommitRequestData;
 import org.apache.kafka.common.protocol.Errors;
 import org.apache.kafka.common.record.DefaultRecordBatch;
 import org.apache.kafka.common.record.RecordBatch;
@@ -37,6 +45,7 @@
 import org.apache.kafka.common.requests.EndTxnRequest;
 import org.apache.kafka.common.requests.EndTxnResponse;
 import org.apache.kafka.common.requests.FindCoordinatorRequest;
+import org.apache.kafka.common.requests.FindCoordinatorRequest.CoordinatorType;
 import org.apache.kafka.common.requests.FindCoordinatorResponse;
 import org.apache.kafka.common.requests.InitProducerIdRequest;
 import org.apache.kafka.common.requests.InitProducerIdResponse;
@@ -46,6 +55,7 @@
 import org.apache.kafka.common.requests.TxnOffsetCommitRequest.CommittedOffset;
 import org.apache.kafka.common.requests.TxnOffsetCommitResponse;
 import org.apache.kafka.common.utils.LogContext;
+import org.apache.kafka.common.utils.PrimitiveRef;
 import org.slf4j.Logger;
 
 import java.util.ArrayList;
@@ -58,6 +68,9 @@
 import java.util.OptionalLong;
 import java.util.PriorityQueue;
 import java.util.Set;
+import java.util.SortedSet;
+import java.util.TreeSet;
+import java.util.function.Consumer;
 import java.util.function.Supplier;
 
 import static org.apache.kafka.common.record.RecordBatch.NO_PRODUCER_EPOCH;
@@ -129,7 +142,7 @@ private static class TopicPartitionEntry {
         // we continue to order batches by the sequence numbers even when the responses come back out of order during
         // leader failover. We add a batch to the queue when it is drained, and remove it when the batch completes
         // (either successfully or through a fatal failure).
-        private PriorityQueue inflightBatchesBySequence;
+        private SortedSet inflightBatchesBySequence;
 
         // We keep track of the last acknowledged offset on a per partition basis in order to disambiguate UnknownProducer
         // responses which are due to the retention period elapsing, and those which are due to actual lost data.
@@ -139,8 +152,18 @@ private static class TopicPartitionEntry {
             this.nextSequence = 0;
             this.lastAckedSequence = NO_LAST_ACKED_SEQUENCE_NUMBER;
             this.lastAckedOffset = ProduceResponse.INVALID_OFFSET;
-            this.inflightBatchesBySequence = new PriorityQueue<>(5, Comparator.comparingInt(ProducerBatch::baseSequence));
+            this.inflightBatchesBySequence = new TreeSet<>(Comparator.comparingInt(ProducerBatch::baseSequence));
         }
+
+        public void resetSequenceNumbers(Consumer resetSequence) {
+            TreeSet newInflights = new TreeSet<>(Comparator.comparingInt(ProducerBatch::baseSequence));
+            for (ProducerBatch inflightBatch : inflightBatchesBySequence) {
+                resetSequence.accept(inflightBatch);
+                newInflights.add(inflightBatch);
+            }
+            inflightBatchesBySequence = newInflights;
+        }
+
     }
 
     private final TopicPartitionBookkeeper topicPartitionBookkeeper;
@@ -254,8 +277,10 @@ public synchronized TransactionalRequestResult initializeTransactions() {
         return handleCachedTransactionRequestResult(() -> {
             transitionTo(State.INITIALIZING);
             setProducerIdAndEpoch(ProducerIdAndEpoch.NONE);
-            InitProducerIdRequest.Builder builder = new InitProducerIdRequest.Builder(transactionalId, transactionTimeoutMs);
-            InitProducerIdHandler handler = new InitProducerIdHandler(builder);
+            InitProducerIdRequestData requestData = new InitProducerIdRequestData()
+                    .setTransactionalId(transactionalId)
+                    .setTransactionTimeoutMs(transactionTimeoutMs);
+            InitProducerIdHandler handler = new InitProducerIdHandler(new InitProducerIdRequest.Builder(requestData));
             enqueueRequest(handler);
             return handler.result;
         }, State.INITIALIZING);
@@ -307,15 +332,13 @@ public synchronized TransactionalRequestResult sendOffsetsToTransaction(Map inflightBatches = topicPartitionBookkeeper.getPartition(topicPartition).inflightBatchesBySequence;
+        if (inflightBatches.isEmpty())
             return RecordBatch.NO_SEQUENCE;
-
-        return first.baseSequence();
+        else
+            return inflightBatches.first().baseSequence();
     }
 
     synchronized ProducerBatch nextBatchBySequence(TopicPartition topicPartition) {
-        PriorityQueue queue = topicPartitionBookkeeper.getPartition(topicPartition).inflightBatchesBySequence;
-        return queue.peek();
+        SortedSet queue = topicPartitionBookkeeper.getPartition(topicPartition).inflightBatchesBySequence;
+        return queue.isEmpty() ? null : queue.first();
     }
 
     synchronized void removeInFlightBatch(ProducerBatch batch) {
         if (hasInflightBatches(batch.topicPartition)) {
-            PriorityQueue queue = topicPartitionBookkeeper.getPartition(batch.topicPartition).inflightBatchesBySequence;
-            queue.remove(batch);
+            topicPartitionBookkeeper.getPartition(batch.topicPartition).inflightBatchesBySequence.remove(batch);
         }
     }
 
-    synchronized void maybeUpdateLastAckedSequence(TopicPartition topicPartition, int sequence) {
+    private void maybeUpdateLastAckedSequence(TopicPartition topicPartition, int sequence) {
         if (sequence > lastAckedSequence(topicPartition).orElse(NO_LAST_ACKED_SEQUENCE_NUMBER))
             topicPartitionBookkeeper.getPartition(topicPartition).lastAckedSequence = sequence;
     }
@@ -519,7 +547,7 @@ synchronized OptionalLong lastAckedOffset(TopicPartition topicPartition) {
         return topicPartitionBookkeeper.lastAckedOffset(topicPartition);
     }
 
-    synchronized void updateLastAckedOffset(ProduceResponse.PartitionResponse response, ProducerBatch batch) {
+    private void updateLastAckedOffset(ProduceResponse.PartitionResponse response, ProducerBatch batch) {
         if (response.baseOffset == ProduceResponse.INVALID_OFFSET)
             return;
         long lastOffset = response.baseOffset + batch.recordCount - 1;
@@ -537,12 +565,66 @@ synchronized void updateLastAckedOffset(ProduceResponse.PartitionResponse respon
         }
     }
 
+    public synchronized void handleCompletedBatch(ProducerBatch batch, ProduceResponse.PartitionResponse response) {
+        if (!hasProducerIdAndEpoch(batch.producerId(), batch.producerEpoch())) {
+            log.debug("Ignoring completed batch {} with producer id {}, epoch {}, and sequence number {} " +
+                            "since the producerId has been reset internally", batch, batch.producerId(),
+                    batch.producerEpoch(), batch.baseSequence());
+            return;
+        }
+
+        maybeUpdateLastAckedSequence(batch.topicPartition, batch.baseSequence() + batch.recordCount - 1);
+        log.debug("ProducerId: {}; Set last ack'd sequence number for topic-partition {} to {}",
+                batch.producerId(),
+                batch.topicPartition,
+                lastAckedSequence(batch.topicPartition).orElse(-1));
+
+        updateLastAckedOffset(response, batch);
+        removeInFlightBatch(batch);
+    }
+
+    private void maybeTransitionToErrorState(RuntimeException exception) {
+        if (exception instanceof ClusterAuthorizationException
+                || exception instanceof TransactionalIdAuthorizationException
+                || exception instanceof ProducerFencedException
+                || exception instanceof UnsupportedVersionException) {
+            transitionToFatalError(exception);
+        } else if (isTransactional()) {
+            transitionToAbortableError(exception);
+        }
+    }
+
+    public synchronized void handleFailedBatch(ProducerBatch batch, RuntimeException exception, boolean adjustSequenceNumbers) {
+        maybeTransitionToErrorState(exception);
+
+        if (!hasProducerIdAndEpoch(batch.producerId(), batch.producerEpoch())) {
+            log.debug("Ignoring failed batch {} with producer id {}, epoch {}, and sequence number {} " +
+                    "since the producerId has been reset internally", batch, batch.producerId(),
+                    batch.producerEpoch(), batch.baseSequence(), exception);
+            return;
+        }
+
+        if (exception instanceof OutOfOrderSequenceException && !isTransactional()) {
+            log.error("The broker returned {} for topic-partition {} with producerId {}, epoch {}, and sequence number {}",
+                    exception, batch.topicPartition, batch.producerId(), batch.producerEpoch(), batch.baseSequence());
+
+            // Reset the producerId since we have hit an irrecoverable exception and cannot make any guarantees
+            // about the previously committed message. Note that this will discard the producer id and sequence
+            // numbers for all existing partitions.
+            resetProducerId();
+        } else {
+            removeInFlightBatch(batch);
+            if (adjustSequenceNumbers)
+                adjustSequencesDueToFailedBatch(batch);
+        }
+    }
+
     // If a batch is failed fatally, the sequence numbers for future batches bound for the partition must be adjusted
     // so that they don't fail with the OutOfOrderSequenceException.
     //
     // This method must only be called when we know that the batch is question has been unequivocally failed by the broker,
     // ie. it has received a confirmed fatal status code like 'Message Too Large' or something similar.
-    synchronized void adjustSequencesDueToFailedBatch(ProducerBatch batch) {
+    private void adjustSequencesDueToFailedBatch(ProducerBatch batch) {
         if (!topicPartitionBookkeeper.contains(batch.topicPartition))
             // Sequence numbers are not being tracked for this partition. This could happen if the producer id was just
             // reset due to a previous OutOfOrderSequenceException.
@@ -552,38 +634,39 @@ synchronized void adjustSequencesDueToFailedBatch(ProducerBatch batch) {
         int currentSequence = sequenceNumber(batch.topicPartition);
         currentSequence -= batch.recordCount;
         if (currentSequence < 0)
-            throw new IllegalStateException("Sequence number for partition " + batch.topicPartition + " is going to become negative : " + currentSequence);
+            throw new IllegalStateException("Sequence number for partition " + batch.topicPartition + " is going to become negative: " + currentSequence);
 
         setNextSequence(batch.topicPartition, currentSequence);
 
-        for (ProducerBatch inFlightBatch : topicPartitionBookkeeper.getPartition(batch.topicPartition).inflightBatchesBySequence) {
+        topicPartitionBookkeeper.getPartition(batch.topicPartition).resetSequenceNumbers(inFlightBatch -> {
             if (inFlightBatch.baseSequence() < batch.baseSequence())
-                continue;
+                return;
+
             int newSequence = inFlightBatch.baseSequence() - batch.recordCount;
             if (newSequence < 0)
                 throw new IllegalStateException("Sequence number for batch with sequence " + inFlightBatch.baseSequence()
-                        + " for partition " + batch.topicPartition + " is going to become negative :" + newSequence);
+                        + " for partition " + batch.topicPartition + " is going to become negative: " + newSequence);
 
             log.info("Resetting sequence number of batch with current sequence {} for partition {} to {}", inFlightBatch.baseSequence(), batch.topicPartition, newSequence);
             inFlightBatch.resetProducerState(new ProducerIdAndEpoch(inFlightBatch.producerId(), inFlightBatch.producerEpoch()), newSequence, inFlightBatch.isTransactional());
-        }
+
+        });
     }
 
-    private synchronized void startSequencesAtBeginning(TopicPartition topicPartition) {
-        int sequence = 0;
-        for (ProducerBatch inFlightBatch : topicPartitionBookkeeper.getPartition(topicPartition).inflightBatchesBySequence) {
+    private void startSequencesAtBeginning(TopicPartition topicPartition) {
+        final PrimitiveRef.IntRef sequence = PrimitiveRef.ofInt(0);
+        topicPartitionBookkeeper.getPartition(topicPartition).resetSequenceNumbers(inFlightBatch -> {
             log.info("Resetting sequence number of batch with current sequence {} for partition {} to {}",
-                    inFlightBatch.baseSequence(), inFlightBatch.topicPartition, sequence);
+                    inFlightBatch.baseSequence(), inFlightBatch.topicPartition, sequence.value);
             inFlightBatch.resetProducerState(new ProducerIdAndEpoch(inFlightBatch.producerId(),
-                    inFlightBatch.producerEpoch()), sequence, inFlightBatch.isTransactional());
-
-            sequence += inFlightBatch.recordCount;
-        }
-        setNextSequence(topicPartition, sequence);
+                    inFlightBatch.producerEpoch()), sequence.value, inFlightBatch.isTransactional());
+            sequence.value += inFlightBatch.recordCount;
+        });
+        setNextSequence(topicPartition, sequence.value);
         topicPartitionBookkeeper.getPartition(topicPartition).lastAckedSequence = NO_LAST_ACKED_SEQUENCE_NUMBER;
     }
 
-    private synchronized boolean hasInflightBatches(TopicPartition topicPartition) {
+    private boolean hasInflightBatches(TopicPartition topicPartition) {
         return topicPartitionBookkeeper.contains(topicPartition)
                 && !topicPartitionBookkeeper.getPartition(topicPartition).inflightBatchesBySequence.isEmpty();
     }
@@ -603,7 +686,7 @@ synchronized void markSequenceUnresolved(TopicPartition topicPartition) {
 
     // Checks if there are any partitions with unresolved partitions which may now be resolved. Returns true if
     // the producer id needs a reset, false otherwise.
-    synchronized boolean shouldResetProducerStateAfterResolvingSequences() {
+    private boolean shouldResetProducerStateAfterResolvingSequences() {
         if (isTransactional())
             // We should not reset producer state if we are transactional. We will transition to a fatal error instead.
             return false;
@@ -628,11 +711,11 @@ synchronized boolean shouldResetProducerStateAfterResolvingSequences() {
         return false;
     }
 
-    private synchronized boolean isNextSequence(TopicPartition topicPartition, int sequence) {
+    private boolean isNextSequence(TopicPartition topicPartition, int sequence) {
         return sequence - lastAckedSequence(topicPartition).orElse(NO_LAST_ACKED_SEQUENCE_NUMBER) == 1;
     }
 
-    private synchronized void setNextSequence(TopicPartition topicPartition, int sequence) {
+    private void setNextSequence(TopicPartition topicPartition, int sequence) {
         topicPartitionBookkeeper.getPartition(topicPartition).nextSequence = sequence;
     }
 
@@ -681,6 +764,16 @@ synchronized void authenticationFailed(AuthenticationException e) {
             request.fatalError(e);
     }
 
+    synchronized void close() {
+        KafkaException shutdownException = new KafkaException("The producer closed forcefully");
+        pendingRequests.forEach(handler ->
+                handler.fatalError(shutdownException));
+        if (pendingResult != null) {
+            pendingResult.setError(shutdownException);
+            pendingResult.done();
+        }
+    }
+
     Node coordinator(FindCoordinatorRequest.CoordinatorType type) {
         switch (type) {
             case GROUP:
@@ -696,7 +789,7 @@ void lookupCoordinator(TxnRequestHandler request) {
         lookupCoordinator(request.coordinatorType(), request.coordinatorKey());
     }
 
-    void setInFlightTransactionalRequestCorrelationId(int correlationId) {
+    void setInFlightCorrelationId(int correlationId) {
         inFlightRequestCorrelationId = correlationId;
     }
 
@@ -728,6 +821,10 @@ synchronized boolean hasPendingOffsetCommits() {
         return !pendingTxnOffsetCommits.isEmpty();
     }
 
+    synchronized boolean hasPendingRequests() {
+        return !pendingRequests.isEmpty();
+    }
+
     // visible for testing
     synchronized boolean hasOngoingTransaction() {
         // transactions are considered ongoing once started until completion or a fatal error
@@ -735,7 +832,7 @@ synchronized boolean hasOngoingTransaction() {
     }
 
     synchronized boolean canRetry(ProduceResponse.PartitionResponse response, ProducerBatch batch) {
-        if (!hasProducerId(batch.producerId()))
+        if (!hasProducerIdAndEpoch(batch.producerId(), batch.producerEpoch()))
             return false;
 
         Errors error = response.error;
@@ -787,7 +884,7 @@ private void transitionTo(State target) {
         transitionTo(target, null);
     }
 
-    private synchronized void transitionTo(State target, RuntimeException error) {
+    private void transitionTo(State target, RuntimeException error) {
         if (!currentState.isTransitionValid(currentState, target)) {
             String idString = transactionalId == null ?  "" : "TransactionalId " + transactionalId + ": ";
             throw new KafkaException(idString + "Invalid transition attempted from state "
@@ -796,7 +893,7 @@ private synchronized void transitionTo(State target, RuntimeException error) {
 
         if (target == State.FATAL_ERROR || target == State.ABORTABLE_ERROR) {
             if (error == null)
-                throw new IllegalArgumentException("Cannot transition to " + target + " with an null exception");
+                throw new IllegalArgumentException("Cannot transition to " + target + " with a null exception");
             lastError = error;
         } else {
             lastError = null;
@@ -816,8 +913,16 @@ private void ensureTransactional() {
     }
 
     private void maybeFailWithError() {
-        if (hasError())
-            throw new KafkaException("Cannot execute transactional method because we are in an error state", lastError);
+        if (hasError()) {
+            // for ProducerFencedException, do not wrap it as a KafkaException
+            // but create a new instance without the call trace since it was not thrown because of the current call
+            if (lastError instanceof ProducerFencedException) {
+                throw new ProducerFencedException("The producer has been rejected from the broker because " +
+                    "it tried to use an old epoch with the transactionalId");
+            } else {
+                throw new KafkaException("Cannot execute transactional method because we are in an error state", lastError);
+            }
+        }
     }
 
     private boolean maybeTerminateRequestWithError(TxnRequestHandler requestHandler) {
@@ -837,7 +942,7 @@ private void enqueueRequest(TxnRequestHandler requestHandler) {
         pendingRequests.add(requestHandler);
     }
 
-    private synchronized void lookupCoordinator(FindCoordinatorRequest.CoordinatorType type, String coordinatorKey) {
+    private void lookupCoordinator(FindCoordinatorRequest.CoordinatorType type, String coordinatorKey) {
         switch (type) {
             case GROUP:
                 consumerGroupCoordinator = null;
@@ -849,11 +954,14 @@ private synchronized void lookupCoordinator(FindCoordinatorRequest.CoordinatorTy
                 throw new IllegalStateException("Invalid coordinator type: " + type);
         }
 
-        FindCoordinatorRequest.Builder builder = new FindCoordinatorRequest.Builder(type, coordinatorKey);
+        FindCoordinatorRequest.Builder builder = new FindCoordinatorRequest.Builder(
+                new FindCoordinatorRequestData()
+                    .setKeyType(type.id())
+                    .setKey(coordinatorKey));
         enqueueRequest(new FindCoordinatorHandler(builder));
     }
 
-    private synchronized void completeTransaction() {
+    private void completeTransaction() {
         transitionTo(State.READY);
         lastError = null;
         transactionStarted = false;
@@ -862,7 +970,7 @@ private synchronized void completeTransaction() {
         partitionsInTransaction.clear();
     }
 
-    private synchronized TxnRequestHandler addPartitionsToTransactionHandler() {
+    private TxnRequestHandler addPartitionsToTransactionHandler() {
         pendingPartitionsInTransaction.addAll(newPartitionsInTransaction);
         newPartitionsInTransaction.clear();
         AddPartitionsToTxnRequest.Builder builder = new AddPartitionsToTxnRequest.Builder(transactionalId,
@@ -879,8 +987,14 @@ private TxnOffsetCommitHandler txnOffsetCommitHandler(TransactionalRequestResult
                     offsetAndMetadata.metadata(), offsetAndMetadata.leaderEpoch());
             pendingTxnOffsetCommits.put(entry.getKey(), committedOffset);
         }
-        TxnOffsetCommitRequest.Builder builder = new TxnOffsetCommitRequest.Builder(transactionalId, consumerGroupId,
-                producerIdAndEpoch.producerId, producerIdAndEpoch.epoch, pendingTxnOffsetCommits);
+        TxnOffsetCommitRequest.Builder builder = new TxnOffsetCommitRequest.Builder(
+            new TxnOffsetCommitRequestData()
+                .setTransactionalId(transactionalId)
+                .setGroupId(consumerGroupId)
+                .setProducerId(producerIdAndEpoch.producerId)
+                .setProducerEpoch(producerIdAndEpoch.epoch)
+                .setTopics(TxnOffsetCommitRequest.getTopics(pendingTxnOffsetCommits))
+        );
         return new TxnOffsetCommitHandler(result, builder);
     }
 
@@ -1020,7 +1134,8 @@ public void handleResponse(AbstractResponse response) {
             Errors error = initProducerIdResponse.error();
 
             if (error == Errors.NONE) {
-                ProducerIdAndEpoch producerIdAndEpoch = new ProducerIdAndEpoch(initProducerIdResponse.producerId(), initProducerIdResponse.epoch());
+                ProducerIdAndEpoch producerIdAndEpoch = new ProducerIdAndEpoch(initProducerIdResponse.data.producerId(),
+                        initProducerIdResponse.data.producerEpoch());
                 setProducerIdAndEpoch(producerIdAndEpoch);
                 transitionTo(State.READY);
                 lastError = null;
@@ -1175,10 +1290,11 @@ String coordinatorKey() {
         public void handleResponse(AbstractResponse response) {
             FindCoordinatorResponse findCoordinatorResponse = (FindCoordinatorResponse) response;
             Errors error = findCoordinatorResponse.error();
+            CoordinatorType coordinatorType = CoordinatorType.forId(builder.data().keyType());
 
             if (error == Errors.NONE) {
                 Node node = findCoordinatorResponse.node();
-                switch (builder.coordinatorType()) {
+                switch (coordinatorType) {
                     case GROUP:
                         consumerGroupCoordinator = node;
                         break;
@@ -1191,11 +1307,11 @@ public void handleResponse(AbstractResponse response) {
             } else if (error == Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED) {
                 fatalError(error.exception());
             } else if (findCoordinatorResponse.error() == Errors.GROUP_AUTHORIZATION_FAILED) {
-                abortableError(new GroupAuthorizationException(builder.coordinatorKey()));
+                abortableError(GroupAuthorizationException.forGroupId(builder.data().key()));
             } else {
                 fatalError(new KafkaException(String.format("Could not find a coordinator with type %s with key %s due to" +
-                        "unexpected error: %s", builder.coordinatorType(), builder.coordinatorKey(),
-                        findCoordinatorResponse.error().message())));
+                        "unexpected error: %s", coordinatorType, builder.data().key(),
+                        findCoordinatorResponse.data().errorMessage())));
             }
         }
     }
@@ -1290,7 +1406,7 @@ public void handleResponse(AbstractResponse response) {
             } else if (error == Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED) {
                 fatalError(error.exception());
             } else if (error == Errors.GROUP_AUTHORIZATION_FAILED) {
-                abortableError(new GroupAuthorizationException(builder.consumerGroupId()));
+                abortableError(GroupAuthorizationException.forGroupId(builder.consumerGroupId()));
             } else {
                 fatalError(new KafkaException("Unexpected error in AddOffsetsToTxnResponse: " + error.message()));
             }
@@ -1323,7 +1439,7 @@ FindCoordinatorRequest.CoordinatorType coordinatorType() {
 
         @Override
         String coordinatorKey() {
-            return builder.consumerGroupId();
+            return builder.data.groupId();
         }
 
         @Override
@@ -1332,7 +1448,7 @@ public void handleResponse(AbstractResponse response) {
             boolean coordinatorReloaded = false;
             Map errors = txnOffsetCommitResponse.errors();
 
-            log.debug("Received TxnOffsetCommit response for consumer group {}: {}", builder.consumerGroupId(),
+            log.debug("Received TxnOffsetCommit response for consumer group {}: {}", builder.data.groupId(),
                     errors);
 
             for (Map.Entry entry : errors.entrySet()) {
@@ -1345,14 +1461,14 @@ public void handleResponse(AbstractResponse response) {
                         || error == Errors.REQUEST_TIMED_OUT) {
                     if (!coordinatorReloaded) {
                         coordinatorReloaded = true;
-                        lookupCoordinator(FindCoordinatorRequest.CoordinatorType.GROUP, builder.consumerGroupId());
+                        lookupCoordinator(FindCoordinatorRequest.CoordinatorType.GROUP, builder.data.groupId());
                     }
                 } else if (error == Errors.UNKNOWN_TOPIC_OR_PARTITION
                         || error == Errors.COORDINATOR_LOAD_IN_PROGRESS) {
                     // If the topic is unknown or the coordinator is loading, retry with the current coordinator
                     continue;
                 } else if (error == Errors.GROUP_AUTHORIZATION_FAILED) {
-                    abortableError(new GroupAuthorizationException(builder.consumerGroupId()));
+                    abortableError(GroupAuthorizationException.forGroupId(builder.data.groupId()));
                     break;
                 } else if (error == Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED
                         || error == Errors.INVALID_PRODUCER_EPOCH
diff --git a/clients/src/main/java/org/apache/kafka/common/Cluster.java b/clients/src/main/java/org/apache/kafka/common/Cluster.java
index 753b7f914f3c0..c765cdc0b0c46 100644
--- a/clients/src/main/java/org/apache/kafka/common/Cluster.java
+++ b/clients/src/main/java/org/apache/kafka/common/Cluster.java
@@ -16,10 +16,9 @@
  */
 package org.apache.kafka.common;
 
-import org.apache.kafka.common.utils.Utils;
-
 import java.net.InetSocketAddress;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
@@ -27,6 +26,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.Set;
 
 /**
@@ -106,23 +106,64 @@ private Cluster(String clusterId,
 
         // Index the nodes for quick lookup
         Map tmpNodesById = new HashMap<>();
-        for (Node node : nodes)
+        Map> tmpPartitionsByNode = new HashMap<>(nodes.size());
+        for (Node node : nodes) {
             tmpNodesById.put(node.id(), node);
+            // Populate the map here to make it easy to add the partitions per node efficiently when iterating over
+            // the partitions
+            tmpPartitionsByNode.put(node.id(), new ArrayList<>());
+        }
         this.nodesById = Collections.unmodifiableMap(tmpNodesById);
 
         // index the partition infos by topic, topic+partition, and node
+        // note that this code is performance sensitive if there are a large number of partitions so we are careful
+        // to avoid unnecessary work
         Map tmpPartitionsByTopicPartition = new HashMap<>(partitions.size());
         Map> tmpPartitionsByTopic = new HashMap<>();
-        Map> tmpAvailablePartitionsByTopic = new HashMap<>();
-        Map> tmpPartitionsByNode = new HashMap<>();
         for (PartitionInfo p : partitions) {
             tmpPartitionsByTopicPartition.put(new TopicPartition(p.topic(), p.partition()), p);
-            tmpPartitionsByTopic.merge(p.topic(), Collections.singletonList(p), Utils::concatListsUnmodifiable);
+            List partitionsForTopic = tmpPartitionsByTopic.get(p.topic());
+            if (partitionsForTopic == null) {
+                partitionsForTopic = new ArrayList<>();
+                tmpPartitionsByTopic.put(p.topic(), partitionsForTopic);
+            }
+            partitionsForTopic.add(p);
             if (p.leader() != null) {
-                tmpAvailablePartitionsByTopic.merge(p.topic(), Collections.singletonList(p), Utils::concatListsUnmodifiable);
-                tmpPartitionsByNode.merge(p.leader().id(), Collections.singletonList(p), Utils::concatListsUnmodifiable);
+                // The broker guarantees that if a partition has a non-null leader, it is one of the brokers returned
+                // in the metadata response
+                List partitionsForNode = Objects.requireNonNull(tmpPartitionsByNode.get(p.leader().id()));
+                partitionsForNode.add(p);
+            }
+        }
+
+        // Update the values of `tmpPartitionsByNode` to contain unmodifiable lists
+        for (Map.Entry> entry : tmpPartitionsByNode.entrySet()) {
+            tmpPartitionsByNode.put(entry.getKey(), Collections.unmodifiableList(entry.getValue()));
+        }
+
+        // Populate `tmpAvailablePartitionsByTopic` and update the values of `tmpPartitionsByTopic` to contain
+        // unmodifiable lists
+        Map> tmpAvailablePartitionsByTopic = new HashMap<>(tmpPartitionsByTopic.size());
+        for (Map.Entry> entry : tmpPartitionsByTopic.entrySet()) {
+            String topic = entry.getKey();
+            List partitionsForTopic = Collections.unmodifiableList(entry.getValue());
+            tmpPartitionsByTopic.put(topic, partitionsForTopic);
+            // Optimise for the common case where all partitions are available
+            boolean foundUnavailablePartition = partitionsForTopic.stream().anyMatch(p -> p.leader() == null);
+            List availablePartitionsForTopic;
+            if (foundUnavailablePartition) {
+                availablePartitionsForTopic = new ArrayList<>(partitionsForTopic.size());
+                for (PartitionInfo p : partitionsForTopic) {
+                    if (p.leader() != null)
+                        availablePartitionsForTopic.add(p);
+                }
+                availablePartitionsForTopic = Collections.unmodifiableList(availablePartitionsForTopic);
+            } else {
+                availablePartitionsForTopic = partitionsForTopic;
             }
+            tmpAvailablePartitionsByTopic.put(topic, availablePartitionsForTopic);
         }
+
         this.partitionsByTopicPartition = Collections.unmodifiableMap(tmpPartitionsByTopicPartition);
         this.partitionsByTopic = Collections.unmodifiableMap(tmpPartitionsByTopic);
         this.availablePartitionsByTopic = Collections.unmodifiableMap(tmpAvailablePartitionsByTopic);
@@ -183,6 +224,21 @@ public Node nodeById(int id) {
         return this.nodesById.get(id);
     }
 
+    /**
+     * Get the node by node id if the replica for the given partition is online
+     * @param partition
+     * @param id
+     * @return the node
+     */
+    public Optional nodeIfOnline(TopicPartition partition, int id) {
+        Node node = nodeById(id);
+        if (node != null && !Arrays.asList(partition(partition).offlineReplicas()).contains(node)) {
+            return Optional.of(node);
+        } else {
+            return Optional.empty();
+        }
+    }
+
     /**
      * Get the current leader for the given topic-partition
      * @param topicPartition The topic and partition we want to know the leader for
diff --git a/clients/src/main/java/org/apache/kafka/common/ElectionType.java b/clients/src/main/java/org/apache/kafka/common/ElectionType.java
new file mode 100644
index 0000000000000..4e29853550677
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/ElectionType.java
@@ -0,0 +1,48 @@
+/*
+ * 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;
+
+import java.util.Arrays;
+import org.apache.kafka.common.annotation.InterfaceStability;
+
+/**
+ * Options for {@link org.apache.kafka.clients.admin.Admin#electLeaders(ElectionType, Set, ElectLeadersOptions)}.
+ *
+ * The API of this class is evolving, see {@link org.apache.kafka.clients.admin.Admin} for details.
+ */
+@InterfaceStability.Evolving
+public enum ElectionType {
+    PREFERRED((byte) 0), UNCLEAN((byte) 1);
+
+    public final byte value;
+
+    ElectionType(byte value) {
+        this.value = value;
+    }
+
+    public static ElectionType valueOf(byte value) {
+        if (value == PREFERRED.value) {
+            return PREFERRED;
+        } else if (value == UNCLEAN.value) {
+            return UNCLEAN;
+        } else {
+            throw new IllegalArgumentException(
+                    String.format("Value %s must be one of %s", value, Arrays.asList(ElectionType.values())));
+        }
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/Endpoint.java b/clients/src/main/java/org/apache/kafka/common/Endpoint.java
new file mode 100644
index 0000000000000..2353de26ec4df
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/Endpoint.java
@@ -0,0 +1,104 @@
+/*
+ * 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;
+
+import java.util.Objects;
+import java.util.Optional;
+
+import org.apache.kafka.common.annotation.InterfaceStability;
+import org.apache.kafka.common.security.auth.SecurityProtocol;
+
+/**
+ * Represents a broker endpoint.
+ */
+
+@InterfaceStability.Evolving
+public class Endpoint {
+
+    private final String listenerName;
+    private final SecurityProtocol securityProtocol;
+    private final String host;
+    private final int port;
+
+    public Endpoint(String listenerName, SecurityProtocol securityProtocol, String host, int port) {
+        this.listenerName = listenerName;
+        this.securityProtocol = securityProtocol;
+        this.host = host;
+        this.port = port;
+    }
+
+    /**
+     * Returns the listener name of this endpoint. This is non-empty for endpoints provided
+     * to broker plugins, but may be empty when used in clients.
+     */
+    public Optional listenerName() {
+        return Optional.ofNullable(listenerName);
+    }
+
+    /**
+     * Returns the security protocol of this endpoint.
+     */
+    public SecurityProtocol securityProtocol() {
+        return securityProtocol;
+    }
+
+    /**
+     * Returns advertised host name of this endpoint.
+     */
+    public String host() {
+        return host;
+    }
+
+    /**
+     * Returns the port to which the listener is bound.
+     */
+    public int port() {
+        return port;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+        if (!(o instanceof Endpoint)) {
+            return false;
+        }
+
+        Endpoint that = (Endpoint) o;
+        return Objects.equals(this.listenerName, that.listenerName) &&
+            Objects.equals(this.securityProtocol, that.securityProtocol) &&
+            Objects.equals(this.host, that.host) &&
+            this.port == that.port;
+
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(listenerName, securityProtocol, host, port);
+    }
+
+    @Override
+    public String toString() {
+        return "Endpoint(" +
+            "listenerName='" + listenerName + '\'' +
+            ", securityProtocol=" + securityProtocol +
+            ", host='" + host + '\'' +
+            ", port=" + port +
+            ')';
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/record/InvalidRecordException.java b/clients/src/main/java/org/apache/kafka/common/InvalidRecordException.java
similarity index 85%
rename from clients/src/main/java/org/apache/kafka/common/record/InvalidRecordException.java
rename to clients/src/main/java/org/apache/kafka/common/InvalidRecordException.java
index 49f616611589e..4c2815bb3bda5 100644
--- a/clients/src/main/java/org/apache/kafka/common/record/InvalidRecordException.java
+++ b/clients/src/main/java/org/apache/kafka/common/InvalidRecordException.java
@@ -14,11 +14,11 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.kafka.common.record;
+package org.apache.kafka.common;
 
-import org.apache.kafka.common.errors.CorruptRecordException;
+import org.apache.kafka.common.errors.ApiException;
 
-public class InvalidRecordException extends CorruptRecordException {
+public class InvalidRecordException extends ApiException {
 
     private static final long serialVersionUID = 1;
 
diff --git a/clients/src/main/java/org/apache/kafka/common/MetricName.java b/clients/src/main/java/org/apache/kafka/common/MetricName.java
index 0af77a038b4b4..9d695b5cf5efd 100644
--- a/clients/src/main/java/org/apache/kafka/common/MetricName.java
+++ b/clients/src/main/java/org/apache/kafka/common/MetricName.java
@@ -17,8 +17,7 @@
 package org.apache.kafka.common;
 
 import java.util.Map;
-
-import org.apache.kafka.common.utils.Utils;
+import java.util.Objects;
 
 /**
  * The MetricName class encapsulates a metric's name, logical group and its related attributes. It should be constructed using metrics.MetricName(...).
@@ -78,10 +77,10 @@ public final class MetricName {
      * @param tags        additional key/value attributes of the metric
      */
     public MetricName(String name, String group, String description, Map tags) {
-        this.name = Utils.notNull(name);
-        this.group = Utils.notNull(group);
-        this.description = Utils.notNull(description);
-        this.tags = Utils.notNull(tags);
+        this.name = Objects.requireNonNull(name);
+        this.group = Objects.requireNonNull(group);
+        this.description = Objects.requireNonNull(description);
+        this.tags = Objects.requireNonNull(tags);
     }
 
     public String name() {
diff --git a/clients/src/main/java/org/apache/kafka/common/MetricNameTemplate.java b/clients/src/main/java/org/apache/kafka/common/MetricNameTemplate.java
index abe121f7ac972..f9c4ef5f730e6 100644
--- a/clients/src/main/java/org/apache/kafka/common/MetricNameTemplate.java
+++ b/clients/src/main/java/org/apache/kafka/common/MetricNameTemplate.java
@@ -21,8 +21,6 @@
 import java.util.Objects;
 import java.util.Set;
 
-import org.apache.kafka.common.utils.Utils;
-
 /**
  * A template for a MetricName. It contains a name, group, and description, as
  * well as all the tags that will be used to create the mBean name. Tag values
@@ -46,10 +44,10 @@ public class MetricNameTemplate {
      * @param tagsNames the set of metric tag names, which can/should be a set that maintains order; may not be null
      */
     public MetricNameTemplate(String name, String group, String description, Set tagsNames) {
-        this.name = Utils.notNull(name);
-        this.group = Utils.notNull(group);
-        this.description = Utils.notNull(description);
-        this.tags = new LinkedHashSet<>(Utils.notNull(tagsNames));
+        this.name = Objects.requireNonNull(name);
+        this.group = Objects.requireNonNull(group);
+        this.description = Objects.requireNonNull(description);
+        this.tags = new LinkedHashSet<>(Objects.requireNonNull(tagsNames));
     }
 
     /**
diff --git a/clients/src/main/java/org/apache/kafka/common/Node.java b/clients/src/main/java/org/apache/kafka/common/Node.java
index d51fa13946b22..020d2bcaf3355 100644
--- a/clients/src/main/java/org/apache/kafka/common/Node.java
+++ b/clients/src/main/java/org/apache/kafka/common/Node.java
@@ -16,6 +16,8 @@
  */
 package org.apache.kafka.common;
 
+import java.util.Objects;
+
 /**
  * Information about a Kafka node
  */
@@ -122,10 +124,10 @@ public boolean equals(Object obj) {
         if (obj == null || getClass() != obj.getClass())
             return false;
         Node other = (Node) obj;
-        return (host == null ? other.host == null : host.equals(other.host)) &&
-            id == other.id &&
+        return id == other.id &&
             port == other.port &&
-            (rack == null ? other.rack == null : rack.equals(other.rack));
+            Objects.equals(host, other.host) &&
+            Objects.equals(rack, other.rack);
     }
 
     @Override
diff --git a/clients/src/main/java/org/apache/kafka/common/TopicPartitionInfo.java b/clients/src/main/java/org/apache/kafka/common/TopicPartitionInfo.java
index 7edf71408f815..60b5d37b04ad0 100644
--- a/clients/src/main/java/org/apache/kafka/common/TopicPartitionInfo.java
+++ b/clients/src/main/java/org/apache/kafka/common/TopicPartitionInfo.java
@@ -21,6 +21,7 @@
 
 import java.util.Collections;
 import java.util.List;
+import java.util.Objects;
 
 /**
  * A class containing leadership, replicas and ISR information for a topic partition.
@@ -90,10 +91,10 @@ public boolean equals(Object o) {
 
         TopicPartitionInfo that = (TopicPartitionInfo) o;
 
-        if (partition != that.partition) return false;
-        if (leader != null ? !leader.equals(that.leader) : that.leader != null) return false;
-        if (replicas != null ? !replicas.equals(that.replicas) : that.replicas != null) return false;
-        return isr != null ? isr.equals(that.isr) : that.isr == null;
+        return partition == that.partition &&
+            Objects.equals(leader, that.leader) &&
+            Objects.equals(replicas, that.replicas) &&
+            Objects.equals(isr, that.isr);
     }
 
     @Override
diff --git a/clients/src/main/java/org/apache/kafka/common/TopicPartitionReplica.java b/clients/src/main/java/org/apache/kafka/common/TopicPartitionReplica.java
index 244260295a108..0a7c419c933bd 100644
--- a/clients/src/main/java/org/apache/kafka/common/TopicPartitionReplica.java
+++ b/clients/src/main/java/org/apache/kafka/common/TopicPartitionReplica.java
@@ -16,9 +16,8 @@
  */
 package org.apache.kafka.common;
 
-import org.apache.kafka.common.utils.Utils;
-
 import java.io.Serializable;
+import java.util.Objects;
 
 
 /**
@@ -32,7 +31,7 @@ public final class TopicPartitionReplica implements Serializable {
     private final String topic;
 
     public TopicPartitionReplica(String topic, int partition, int brokerId) {
-        this.topic = Utils.notNull(topic);
+        this.topic = Objects.requireNonNull(topic);
         this.partition = partition;
         this.brokerId = brokerId;
     }
diff --git a/clients/src/main/java/org/apache/kafka/common/acl/AclOperation.java b/clients/src/main/java/org/apache/kafka/common/acl/AclOperation.java
index 7befd6e9d36dd..671069775ca0e 100644
--- a/clients/src/main/java/org/apache/kafka/common/acl/AclOperation.java
+++ b/clients/src/main/java/org/apache/kafka/common/acl/AclOperation.java
@@ -108,6 +108,9 @@ public enum AclOperation {
      */
     IDEMPOTENT_WRITE((byte) 12);
 
+    // Note: we cannot have more than 30 ACL operations without modifying the format used
+    // to describe ACL operations in MetadataResponse.
+
     private final static HashMap CODE_TO_VALUE = new HashMap<>();
 
     static {
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 9cf13dd35df02..e91a9a5f836f9 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
@@ -22,6 +22,7 @@
 import org.apache.kafka.common.utils.Utils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import org.apache.kafka.common.config.provider.ConfigProvider;
 
 import java.util.ArrayList;
 import java.util.Collections;
@@ -52,13 +53,58 @@ public class AbstractConfig {
 
     private final ConfigDef definition;
 
+    public static final String CONFIG_PROVIDERS_CONFIG = "config.providers";
+
+    private static final String CONFIG_PROVIDERS_PARAM = ".param.";
+
+    /**
+     * Construct a configuration with a ConfigDef and the configuration properties, which can include properties
+     * for zero or more {@link ConfigProvider} that will be used to resolve variables in configuration property
+     * values.
+     *
+     * The originals is a name-value pair configuration properties and optional config provider configs. The
+     * value of the configuration can be a variable as defined below or the actual value. This constructor will
+     * first instantiate the ConfigProviders using the config provider configs, then it will find all the
+     * variables in the values of the originals configurations, attempt to resolve the variables using the named
+     * ConfigProviders, and then parse and validate the configurations.
+     *
+     * ConfigProvider configs can be passed either as configs in the originals map or in the separate
+     * configProviderProps map. If config providers properties are passed in the configProviderProps any config
+     * provider properties in originals map will be ignored. If ConfigProvider properties are not provided, the
+     * constructor will skip the variable substitution step and will simply validate and parse the supplied
+     * configuration.
+     *
+     * The "{@code config.providers}" configuration property and all configuration properties that begin with the
+     * "{@code config.providers.}" prefix are reserved. The "{@code config.providers}" configuration property
+     * specifies the names of the config providers, and properties that begin with the "{@code config.providers..}"
+     * prefix correspond to the properties for that named provider. For example, the "{@code config.providers..class}"
+     * property specifies the name of the {@link ConfigProvider} implementation class that should be used for
+     * the provider.
+     *
+     * The keys for ConfigProvider configs in both originals and configProviderProps will start with the above
+     * mentioned "{@code config.providers.}" prefix.
+     *
+     * Variables have the form "${providerName:[path:]key}", where "providerName" is the name of a ConfigProvider,
+     * "path" is an optional string, and "key" is a required string. This variable is resolved by passing the "key"
+     * and optional "path" to a ConfigProvider with the specified name, and the result from the ConfigProvider is
+     * then used in place of the variable. Variables that cannot be resolved by the AbstractConfig constructor will
+     * be left unchanged in the configuration.
+     *
+     *
+     * @param definition the definition of the configurations; may not be null
+     * @param originals the configuration properties plus any optional config provider properties;
+     * @param configProviderProps the map of properties of config providers which will be instantiated by
+     *        the constructor to resolve any variables in {@code originals}; may be null or empty
+     * @param doLog whether the configurations should be logged
+     */
     @SuppressWarnings("unchecked")
-    public AbstractConfig(ConfigDef definition, Map originals, boolean doLog) {
+    public AbstractConfig(ConfigDef definition, Map originals,  Map 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.");
-        this.originals = (Map) originals;
+
+        this.originals = resolveConfigVariables(configProviderProps, (Map) originals);
         this.values = definition.parse(this.originals);
         Map configUpdates = postProcessParsedConfig(Collections.unmodifiableMap(this.values));
         for (Map.Entry update : configUpdates.entrySet()) {
@@ -71,8 +117,30 @@ public AbstractConfig(ConfigDef definition, Map originals, boolean doLog)
             logAll();
     }
 
+    /**
+     * Construct a configuration with a ConfigDef and the configuration properties,
+     * which can include properties for zero or more {@link ConfigProvider}
+     * that will be used to resolve variables in configuration property values.
+     *
+     * @param definition the definition of the configurations; may not be null
+     * @param originals the configuration properties plus any optional config provider properties; may not be null
+     */
     public AbstractConfig(ConfigDef definition, Map originals) {
-        this(definition, originals, true);
+        this(definition, originals, Collections.emptyMap(), true);
+    }
+
+    /**
+     * Construct a configuration with a ConfigDef and the configuration properties,
+     * which can include properties for zero or more {@link ConfigProvider}
+     * that will be used to resolve variables in configuration property values.
+     *
+     * @param definition the definition of the configurations; may not be null
+     * @param originals the configuration properties plus any optional config provider properties; may not be null
+     * @param doLog whether the configurations should be logged
+     */
+    public AbstractConfig(ConfigDef definition, Map originals, boolean doLog) {
+        this(definition, originals, Collections.emptyMap(), doLog);
+
     }
 
     /**
@@ -287,6 +355,29 @@ public void logUnused() {
             log.warn("The configuration '{}' was supplied but isn't a known config.", key);
     }
 
+    private  T getConfiguredInstance(Object klass, Class t, Map configPairs) {
+        if (klass == null)
+            return null;
+
+        Object o;
+        if (klass instanceof String) {
+            try {
+                o = Utils.newInstance((String) klass, t);
+            } catch (ClassNotFoundException e) {
+                throw new KafkaException("Class " + klass + " cannot be found", e);
+            }
+        } else if (klass instanceof Class) {
+            o = Utils.newInstance((Class) klass);
+        } else
+            throw new KafkaException("Unexpected element of type " + klass.getClass().getName() + ", expected String or Class");
+        if (!t.isInstance(o))
+            throw new KafkaException(klass + " is not an instance of " + t.getName());
+        if (o instanceof Configurable)
+            ((Configurable) o).configure(configPairs);
+
+        return t.cast(o);
+    }
+
     /**
      * Get a configured instance of the give class specified by the given configuration key. If the object implements
      * Configurable configure it using the configuration.
@@ -297,14 +388,8 @@ public void logUnused() {
      */
     public  T getConfiguredInstance(String key, Class t) {
         Class c = getClass(key);
-        if (c == null)
-            return null;
-        Object o = Utils.newInstance(c);
-        if (!t.isInstance(o))
-            throw new KafkaException(c.getName() + " is not an instance of " + t.getName());
-        if (o instanceof Configurable)
-            ((Configurable) o).configure(originals());
-        return t.cast(o);
+
+        return getConfiguredInstance(c, t, originals());
     }
 
     /**
@@ -332,7 +417,6 @@ public  List getConfiguredInstances(String key, Class t, Map List getConfiguredInstances(List classNames, Class t, M
         Map configPairs = originals();
         configPairs.putAll(configOverrides);
         for (Object klass : classNames) {
-            Object o;
-            if (klass instanceof String) {
-                try {
-                    o = Utils.newInstance((String) klass, t);
-                } catch (ClassNotFoundException e) {
-                    throw new KafkaException(klass + " ClassNotFoundException exception occurred", e);
-                }
-            } else if (klass instanceof Class) {
-                o = Utils.newInstance((Class) klass);
-            } else
-                throw new KafkaException("List contains element of type " + klass.getClass().getName() + ", expected String or Class");
-            if (!t.isInstance(o))
-                throw new KafkaException(klass + " is not an instance of " + t.getName());
-            if (o instanceof Configurable)
-                ((Configurable) o).configure(configPairs);
+            Object o = getConfiguredInstance(klass, t, configPairs);
             objects.add(t.cast(o));
         }
         return objects;
     }
 
+    private Map extractPotentialVariables(Map  configMap) {
+        // Variables are tuples of the form "${providerName:[path:]key}". From the configMap we extract the subset of configs with string
+        // values as potential variables.
+        Map configMapAsString = new HashMap<>();
+        for (Map.Entry entry : configMap.entrySet()) {
+            if (entry.getValue() instanceof String)
+                configMapAsString.put((String) entry.getKey(), (String) entry.getValue());
+        }
+
+        return configMapAsString;
+    }
+
+    /**
+     * Instantiates given list of config providers and fetches the actual values of config variables from the config providers.
+     * returns a map of config key and resolved values.
+     * @param configProviderProps The map of config provider configs
+     * @param originals The map of raw configs.
+     * @return map of resolved config variable.
+     */
+    @SuppressWarnings("unchecked")
+    private  Map resolveConfigVariables(Map configProviderProps, Map originals) {
+        Map providerConfigString;
+        Map configProperties;
+        Map resolvedOriginals = new HashMap<>();
+        // As variable configs are strings, parse the originals and obtain the potential variable configs.
+        Map indirectVariables = extractPotentialVariables(originals);
+
+        resolvedOriginals.putAll(originals);
+        if (configProviderProps == null || configProviderProps.isEmpty()) {
+            providerConfigString = indirectVariables;
+            configProperties = originals;
+        } else {
+            providerConfigString = extractPotentialVariables(configProviderProps);
+            configProperties = configProviderProps;
+        }
+        Map providers = instantiateConfigProviders(providerConfigString, configProperties);
+
+        if (!providers.isEmpty()) {
+            ConfigTransformer configTransformer = new ConfigTransformer(providers);
+            ConfigTransformerResult result = configTransformer.transform(indirectVariables);
+            if (!result.data().isEmpty()) {
+                resolvedOriginals.putAll(result.data());
+            }
+        }
+        providers.values().forEach(x -> Utils.closeQuietly(x, "config provider"));
+
+        return new ResolvingMap<>(resolvedOriginals, originals);
+    }
+
+    private Map configProviderProperties(String configProviderPrefix, Map providerConfigProperties) {
+        Map result = new HashMap<>();
+        for (Map.Entry entry : providerConfigProperties.entrySet()) {
+            String key = entry.getKey();
+            if (key.startsWith(configProviderPrefix) && key.length() > configProviderPrefix.length()) {
+                result.put(key.substring(configProviderPrefix.length()), entry.getValue());
+            }
+        }
+        return result;
+    }
+
+    /**
+     * Instantiates and configures the ConfigProviders. The config providers configs are defined as follows:
+     * config.providers : A comma-separated list of names for providers.
+     * config.providers.{name}.class : The Java class name for a provider.
+     * config.providers.{name}.param.{param-name} : A parameter to be passed to the above Java class on initialization.
+     * returns a map of config provider name and its instance.
+     * @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.
+     */
+    private Map instantiateConfigProviders(Map indirectConfigs, Map providerConfigProperties) {
+        final String configProviders = indirectConfigs.get(CONFIG_PROVIDERS_CONFIG);
+
+        if (configProviders == null || configProviders.isEmpty()) {
+            return Collections.emptyMap();
+        }
+
+        Map providerMap = new HashMap<>();
+
+        for (String provider: configProviders.split(",")) {
+            String providerClass = CONFIG_PROVIDERS_CONFIG + "." + provider + ".class";
+            if (indirectConfigs.containsKey(providerClass))
+                providerMap.put(provider, indirectConfigs.get(providerClass));
+
+        }
+        // Instantiate Config Providers
+        Map configProviderInstances = new HashMap<>();
+        for (Map.Entry entry : providerMap.entrySet()) {
+            try {
+                String prefix = CONFIG_PROVIDERS_CONFIG + "." + entry.getKey() + CONFIG_PROVIDERS_PARAM;
+                Map configProperties = configProviderProperties(prefix, providerConfigProperties);
+                ConfigProvider provider = Utils.newInstance(entry.getValue(), ConfigProvider.class);
+                provider.configure(configProperties);
+                configProviderInstances.put(entry.getKey(), provider);
+            } catch (ClassNotFoundException e) {
+                log.error("ClassNotFoundException exception occurred: " + entry.getValue());
+                throw new ConfigException("Invalid config:" + entry.getValue() + " ClassNotFoundException exception occurred", e);
+            }
+        }
+
+        return configProviderInstances;
+    }
+
     @Override
     public boolean equals(Object o) {
         if (this == o) return true;
@@ -429,4 +602,31 @@ public V get(Object key) {
             return super.get(key);
         }
     }
+
+    /**
+     * ResolvingMap keeps a track of the original map instance and the resolved configs.
+     * The originals are tracked in a separate nested map and may be a `RecordingMap`; thus
+     * any access to a value for a key needs to be recorded on the originals map.
+     * The resolved configs are kept in the inherited map and are therefore mutable, though any
+     * mutations are not applied to the originals.
+     */
+    private static class ResolvingMap extends HashMap {
+
+        private final Map originals;
+
+        ResolvingMap(Map resolved, Map originals) {
+            super(resolved);
+            this.originals = Collections.unmodifiableMap(originals);
+        }
+
+        @Override
+        public V get(Object key) {
+            if (key instanceof String && originals.containsKey(key)) {
+                // Intentionally ignore the result; call just to mark the original entry as used
+                originals.get(key);
+            }
+            // But always use the resolved entry
+            return super.get(key);
+        }
+    }
 }
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 7669f1dd99942..43ab70d9b7a4f 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
@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.common.config;
 
+import java.util.stream.Collectors;
 import org.apache.kafka.common.config.types.Password;
 import org.apache.kafka.common.utils.Utils;
 
@@ -30,6 +31,8 @@
 import java.util.Locale;
 import java.util.Map;
 import java.util.Set;
+import java.util.function.BiConsumer;
+import java.util.function.Supplier;
 import java.util.regex.Pattern;
 
 /**
@@ -53,7 +56,7 @@
  * defs.define("config_with_validator", Type.INT, 42, Range.atLeast(0), "Configuration with user provided validator.");
  * defs.define("config_with_dependents", Type.INT, "Configuration with dependents.", "group", 1, "Config With Dependents", Arrays.asList("config_with_default","config_with_validator"));
  *
- * Map<String, String> props = new HashMap<>();
+ * Map<String, String> props = new HashMap<>();
  * props.put("config_with_default", "some value");
  * props.put("config_with_dependents", "some other value");
  *
@@ -705,9 +708,16 @@ else if (value instanceof String)
                 case CLASS:
                     if (value instanceof Class)
                         return value;
-                    else if (value instanceof String)
-                        return Class.forName(trimmed, true, Utils.getContextOrKafkaClassLoader());
-                    else
+                    else if (value instanceof String) {
+                        ClassLoader contextOrKafkaClassLoader = Utils.getContextOrKafkaClassLoader();
+                        // Use loadClass here instead of Class.forName because the name we use here may be an alias
+                        // and not match the name of the class that gets loaded. If that happens, Class.forName can
+                        // throw an exception.
+                        Class klass = contextOrKafkaClassLoader.loadClass(trimmed);
+                        // Invoke forName here with the true name of the requested class to cause class
+                        // initialization to take place.
+                        return Class.forName(klass.getName(), true, contextOrKafkaClassLoader);
+                    } else
                         throw new ConfigException(name, value, "Expected a Class instance or class name.");
                 default:
                     throw new IllegalStateException("Unknown type.");
@@ -938,6 +948,33 @@ public String toString() {
         }
     }
 
+    public static class CaseInsensitiveValidString implements Validator {
+
+        final Set validStrings;
+
+        private CaseInsensitiveValidString(List validStrings) {
+            this.validStrings = validStrings.stream()
+                .map(s -> s.toUpperCase(Locale.ROOT))
+                .collect(Collectors.toSet());
+        }
+
+        public static CaseInsensitiveValidString in(String... validStrings) {
+            return new CaseInsensitiveValidString(Arrays.asList(validStrings));
+        }
+
+        @Override
+        public void ensureValid(String name, Object o) {
+            String s = (String) o;
+            if (s == null || !validStrings.contains(s.toUpperCase(Locale.ROOT))) {
+                throw new ConfigException(name, o, "String must be one of (case insensitive): " + Utils.join(validStrings, ", "));
+            }
+        }
+
+        public String toString() {
+            return "(case insensitive) [" + Utils.join(validStrings, ", ") + "]";
+        }
+    }
+
     public static class NonNullValidator implements Validator {
         @Override
         public void ensureValid(String name, Object value) {
@@ -952,6 +989,32 @@ public String toString() {
         }
     }
 
+    public static class LambdaValidator implements Validator {
+        BiConsumer ensureValid;
+        Supplier toStringFunction;
+
+        private LambdaValidator(BiConsumer ensureValid,
+                                Supplier toStringFunction) {
+            this.ensureValid = ensureValid;
+            this.toStringFunction = toStringFunction;
+        }
+
+        public static LambdaValidator with(BiConsumer ensureValid,
+                                           Supplier toStringFunction) {
+            return new LambdaValidator(ensureValid, toStringFunction);
+        }
+
+        @Override
+        public void ensureValid(String name, Object value) {
+            ensureValid.accept(name, value);
+        }
+
+        @Override
+        public String toString() {
+            return toStringFunction.get();
+        }
+    }
+
     public static class CompositeValidator implements Validator {
         private final List validators;
 
@@ -1078,6 +1141,10 @@ public ConfigKey(String name, Type type, Object defaultValue, Validator validato
         public boolean hasDefault() {
             return !NO_DEFAULT_VALUE.equals(this.defaultValue);
         }
+
+        public Type type() {
+            return type;
+        }
     }
 
     protected List headers() {
@@ -1133,7 +1200,7 @@ private void addColumnValue(StringBuilder builder, String value) {
      * If dynamicUpdateModes is non-empty, a "Dynamic Update Mode" column
      * will be included n the table with the value of the update mode. Default
      * mode is "read-only".
-     * @param dynamicUpdateModes Config name -> update mode mapping
+     * @param dynamicUpdateModes Config name -> update mode mapping
      */
     public String toHtmlTable(Map dynamicUpdateModes) {
         boolean hasUpdateModes = !dynamicUpdateModes.isEmpty();
@@ -1314,7 +1381,16 @@ public void embed(final String keyPrefix, final String groupPrefix, final int st
      */
     private static Validator embeddedValidator(final String keyPrefix, final Validator base) {
         if (base == null) return null;
-        return (name, value) -> base.ensureValid(name.substring(keyPrefix.length()), value);
+        return new Validator() {
+            public void ensureValid(String name, Object value) {
+                base.ensureValid(name.substring(keyPrefix.length()), value);
+            }
+
+            @Override
+            public String toString() {
+                return base.toString();
+            }
+        };
     }
 
     /**
@@ -1361,4 +1437,54 @@ public boolean visible(String name, Map parsedConfig) {
         };
     }
 
+    public String toHtml() {
+        return toHtml(Collections.emptyMap());
+    }
+
+    /**
+     * Converts this config into an HTML list that can be embedded into docs.
+     * If dynamicUpdateModes is non-empty, a "Dynamic Update Mode" label
+     * will be included in the config details with the value of the update mode. Default
+     * mode is "read-only".
+     * @param dynamicUpdateModes Config name -> update mode mapping
+     */
+    public String toHtml(Map dynamicUpdateModes) {
+        boolean hasUpdateModes = !dynamicUpdateModes.isEmpty();
+        List configs = sortedConfigs();
+        StringBuilder b = new StringBuilder();
+        b.append("
          \n");
+
+        for (ConfigKey key : configs) {
+            if (key.internalConfig) {
+                continue;
+            }
+            b.append("
        • ");
+            b.append("");
+            b.append(key.name);
+            b.append(": ");
+            b.append(key.documentation);
+            b.append("
          ");
+            // details
+            b.append("
            ");
+            for (String detail : headers()) {
+                if (detail.equals("Name") || detail.equals("Description")) continue;
+                addConfigDetail(b, detail, getConfigValue(key, detail));
+            }
+            if (hasUpdateModes) {
+                String updateMode = dynamicUpdateModes.get(key.name);
+                if (updateMode == null)
+                    updateMode = "read-only";
+                addConfigDetail(b, "Update Mode", updateMode);
+            }
+            b.append("
          ");
+            b.append("
          • \n");
+        }
+        b.append("
        \n");
+        return b.toString();
+    }
+
+    private static void addConfigDetail(StringBuilder builder, String name, String value) {
+        builder.append("
      • " + name + ": " + value + "
        • ");
+    }
+
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/config/ConfigResource.java b/clients/src/main/java/org/apache/kafka/common/config/ConfigResource.java
index 5343a6bcdfd3d..8870238f638a1 100644
--- a/clients/src/main/java/org/apache/kafka/common/config/ConfigResource.java
+++ b/clients/src/main/java/org/apache/kafka/common/config/ConfigResource.java
@@ -33,7 +33,7 @@ public final class ConfigResource {
      * Type of resource.
      */
     public enum Type {
-        BROKER((byte) 4), TOPIC((byte) 2), UNKNOWN((byte) 0);
+        BROKER_LOGGER((byte) 8), BROKER((byte) 4), TOPIC((byte) 2), UNKNOWN((byte) 0);
 
         private static final Map TYPES = Collections.unmodifiableMap(
             Arrays.stream(values()).collect(Collectors.toMap(Type::id, Function.identity()))
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 a830f9fe379a1..4f078b15a32f3 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
@@ -53,7 +53,7 @@
  * {@link ConfigProvider#unsubscribe(String, Set, ConfigChangeCallback)} methods.
  */
 public class ConfigTransformer {
-    public static final Pattern DEFAULT_PATTERN = Pattern.compile("\\$\\{(.*?):((.*?):)?(.*?)\\}");
+    public static final Pattern DEFAULT_PATTERN = Pattern.compile("\\$\\{([^}]*?):(([^}]*?):)?([^}]*?)\\}");
     private static final String EMPTY_PATH = "";
 
     private final Map configProviders;
diff --git a/clients/src/main/java/org/apache/kafka/common/config/LogLevelConfig.java b/clients/src/main/java/org/apache/kafka/common/config/LogLevelConfig.java
new file mode 100644
index 0000000000000..fe7e2eb6669e7
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/config/LogLevelConfig.java
@@ -0,0 +1,71 @@
+/*
+ * 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;
+
+import java.util.Arrays;
+import java.util.HashSet;
+import java.util.Set;
+
+/**
+ * This class holds definitions for log level configurations related to Kafka's application logging. See KIP-412 for additional information
+ */
+public class LogLevelConfig {
+    /*
+     * NOTE: DO NOT CHANGE EITHER CONFIG NAMES AS THESE ARE PART OF THE PUBLIC API AND CHANGE WILL BREAK USER CODE.
+     */
+
+    /**
+     * The FATAL level designates a very severe error
+     * that will lead the Kafka broker to abort.
+     */
+    public static final String FATAL_LOG_LEVEL = "FATAL";
+
+    /**
+     * The ERROR level designates error events that
+     * might still allow the broker to continue running.
+     */
+    public static final String ERROR_LOG_LEVEL = "ERROR";
+
+    /**
+     * The WARN level designates potentially harmful situations.
+     */
+    public static final String WARN_LOG_LEVEL = "WARN";
+
+    /**
+     * The INFO level designates informational messages
+     * that highlight normal Kafka events at a coarse-grained level
+     */
+    public static final String INFO_LOG_LEVEL = "INFO";
+
+    /**
+     * The DEBUG level designates fine-grained
+     * informational events that are most useful to debug Kafka
+     */
+    public static final String DEBUG_LOG_LEVEL = "DEBUG";
+
+    /**
+     * The TRACE level designates finer-grained
+     * informational events than the DEBUG level.
+     */
+    public static final String TRACE_LOG_LEVEL = "TRACE";
+
+    public static final Set VALID_LOG_LEVELS = new HashSet<>(Arrays.asList(
+            FATAL_LOG_LEVEL, ERROR_LOG_LEVEL, WARN_LOG_LEVEL,
+            INFO_LOG_LEVEL, DEBUG_LOG_LEVEL, TRACE_LOG_LEVEL
+    ));
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/config/SecurityConfig.java b/clients/src/main/java/org/apache/kafka/common/config/SecurityConfig.java
new file mode 100644
index 0000000000000..b4dc26c7ea5cb
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/config/SecurityConfig.java
@@ -0,0 +1,29 @@
+/*
+ * 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;
+
+/**
+ * Contains the common security config for SSL and SASL
+ */
+public class SecurityConfig {
+
+    public static final String SECURITY_PROVIDERS_CONFIG = "security.providers";
+    public static final String SECURITY_PROVIDERS_DOC = "A list of configurable creator classes each returning a provider" +
+        " implementing security algorithms. These classes should implement the" +
+        " org.apache.kafka.common.security.auth.SecurityProviderCreator interface.";
+
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/config/SslClientAuth.java b/clients/src/main/java/org/apache/kafka/common/config/SslClientAuth.java
new file mode 100644
index 0000000000000..9d85b184ab9ab
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/config/SslClientAuth.java
@@ -0,0 +1,48 @@
+/*
+ * 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;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Locale;
+
+/**
+ * Describes whether the server should require or request client authentication.
+ */
+public enum SslClientAuth {
+    REQUIRED,
+    REQUESTED,
+    NONE;
+
+    public static final List VALUES =
+        Collections.unmodifiableList(Arrays.asList(SslClientAuth.values()));
+
+    public static SslClientAuth forConfig(String key) {
+        if (key == null) {
+            return SslClientAuth.NONE;
+        }
+        String upperCaseKey = key.toUpperCase(Locale.ROOT);
+        for (SslClientAuth auth : VALUES) {
+            if (auth.name().equals(upperCaseKey)) {
+                return auth;
+            }
+        }
+        return null;
+    }
+}
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 2ed177b6fbc55..47b7cf7643bff 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
@@ -17,6 +17,7 @@
 package org.apache.kafka.common.config;
 
 import org.apache.kafka.common.config.internals.BrokerSecurityConfigs;
+import org.apache.kafka.common.security.ssl.SimpleSslContextProvider;
 import org.apache.kafka.common.utils.Utils;
 
 import javax.net.ssl.KeyManagerFactory;
@@ -47,6 +48,11 @@ public class SslConfigs {
     public static final String DEFAULT_PRINCIPAL_BUILDER_CLASS =
             org.apache.kafka.common.security.auth.DefaultPrincipalBuilder.class.getName();
 
+    public static final String SSL_CONTEXT_PROVIDER_CLASS_CONFIG = "ssl.context.provider.class";
+    public static final String SSL_CONTEXT_PROVIDER_CLASS_DOC = "The fully qualified name of a class that implements the "
+            + "SslContextProvider interface, which is used to build SSLContext insides encrypted channels.";
+    public static final String DEFAULT_SSL_CONTEXT_PROVIDER_CLASS = SimpleSslContextProvider.class.getName();
+
     public static final String SSL_PROTOCOL_CONFIG = "ssl.protocol";
     public static final String SSL_PROTOCOL_DOC = "The SSL protocol used to generate the SSLContext. "
             + "Default setting is TLS, which is fine for most cases. "
@@ -123,6 +129,7 @@ public class SslConfigs {
 
     public static void addClientSslSupport(ConfigDef config) {
         config.define(SslConfigs.SSL_PROTOCOL_CONFIG, ConfigDef.Type.STRING, SslConfigs.DEFAULT_SSL_PROTOCOL, ConfigDef.Importance.MEDIUM, SslConfigs.SSL_PROTOCOL_DOC)
+                .define(SslConfigs.SSL_CONTEXT_PROVIDER_CLASS_CONFIG, ConfigDef.Type.STRING, SslConfigs.DEFAULT_SSL_CONTEXT_PROVIDER_CLASS, ConfigDef.Importance.MEDIUM, SslConfigs.SSL_CONTEXT_PROVIDER_CLASS_DOC)
                 .define(SslConfigs.SSL_PROVIDER_CONFIG, ConfigDef.Type.STRING, null, ConfigDef.Importance.MEDIUM, SslConfigs.SSL_PROVIDER_DOC)
                 .define(SslConfigs.SSL_CIPHER_SUITES_CONFIG, ConfigDef.Type.LIST, null, ConfigDef.Importance.LOW, SslConfigs.SSL_CIPHER_SUITES_DOC)
                 .define(SslConfigs.SSL_ENABLED_PROTOCOLS_CONFIG, ConfigDef.Type.LIST, SslConfigs.DEFAULT_SSL_ENABLED_PROTOCOLS, ConfigDef.Importance.MEDIUM, SslConfigs.SSL_ENABLED_PROTOCOLS_DOC)
@@ -147,4 +154,16 @@ public static void addClientSslSupport(ConfigDef config) {
             SslConfigs.SSL_TRUSTSTORE_TYPE_CONFIG,
             SslConfigs.SSL_TRUSTSTORE_LOCATION_CONFIG,
             SslConfigs.SSL_TRUSTSTORE_PASSWORD_CONFIG);
+
+    public static final Set NON_RECONFIGURABLE_CONFIGS = Utils.mkSet(
+            BrokerSecurityConfigs.SSL_CLIENT_AUTH_CONFIG,
+            SslConfigs.SSL_PROTOCOL_CONFIG,
+            SslConfigs.SSL_PROVIDER_CONFIG,
+            SslConfigs.SSL_CONTEXT_PROVIDER_CLASS_CONFIG,
+            SslConfigs.SSL_CIPHER_SUITES_CONFIG,
+            SslConfigs.SSL_ENABLED_PROTOCOLS_CONFIG,
+            SslConfigs.SSL_KEYMANAGER_ALGORITHM_CONFIG,
+            SslConfigs.SSL_TRUSTMANAGER_ALGORITHM_CONFIG,
+            SslConfigs.SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_CONFIG,
+            SslConfigs.SSL_SECURE_RANDOM_IMPLEMENTATION_CONFIG);
 }
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 57662d5a8663c..9b36736524f34 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
@@ -76,12 +76,12 @@ public class TopicConfig {
         "their data. If set to -1, no time limit is applied.";
 
     public static final String MAX_MESSAGE_BYTES_CONFIG = "max.message.bytes";
-    public static final String MAX_MESSAGE_BYTES_DOC = "
        The largest record batch size allowed by Kafka. If this " +
+    public static final String MAX_MESSAGE_BYTES_DOC = "The largest record batch size allowed by Kafka. If this " +
         "is increased and there are consumers older than 0.10.2, the consumers' fetch size must also be increased so that " +
-        "the they can fetch record batches this large.
        " +
-        "
        In the latest message format version, records are always grouped into batches for efficiency. In previous " +
+        "the they can fetch record batches this large. " +
+        "In the latest message format version, records are always grouped into batches for efficiency. In previous " +
         "message format versions, uncompressed records are not grouped into batches and this limit only applies to a " +
-        "single record in that case.
        ";
+        "single record in that case.";
 
     public static final String INDEX_INTERVAL_BYTES_CONFIG = "index.interval.bytes";
     public static final String INDEX_INTERVAL_BYTES_DOCS = "This setting controls how frequently " +
@@ -104,6 +104,10 @@ public class TopicConfig {
     public static final String 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 MAX_COMPACTION_LAG_MS_CONFIG = "max.compaction.lag.ms";
+    public static final String 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 static final String MIN_CLEANABLE_DIRTY_RATIO_CONFIG = "min.cleanable.dirty.ratio";
     public static final String MIN_CLEANABLE_DIRTY_RATIO_DOC = "This configuration controls how frequently " +
         "the log compactor will attempt to clean the log (assuming log " +
@@ -111,7 +115,11 @@ public class TopicConfig {
         "50% of the log has been compacted. This ratio bounds the maximum space wasted in " +
         "the log by duplicates (at 50% at most 50% of the log could be duplicates). A " +
         "higher ratio will mean fewer, more efficient cleanings but will mean more wasted " +
-        "space in the log.";
+        "space in the log. If the " + MAX_COMPACTION_LAG_MS_CONFIG + " or the " + MIN_COMPACTION_LAG_MS_CONFIG +
+        " configurations are also specified, then the log compactor considers the log to be 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 " + MIN_COMPACTION_LAG_MS_CONFIG + " duration, or (ii) if the log has had " +
+        "dirty (uncompacted) records for at most the " + MAX_COMPACTION_LAG_MS_CONFIG + " period.";
 
     public static final String CLEANUP_POLICY_CONFIG = "cleanup.policy";
     public static final String CLEANUP_POLICY_COMPACT = "compact";
diff --git a/clients/src/main/java/org/apache/kafka/common/config/internals/BrokerSecurityConfigs.java b/clients/src/main/java/org/apache/kafka/common/config/internals/BrokerSecurityConfigs.java
index 98f6467a48429..8a4a51738dedd 100644
--- a/clients/src/main/java/org/apache/kafka/common/config/internals/BrokerSecurityConfigs.java
+++ b/clients/src/main/java/org/apache/kafka/common/config/internals/BrokerSecurityConfigs.java
@@ -55,7 +55,7 @@ public class BrokerSecurityConfigs {
             " see  security authorization and acls. Note that this configuration is ignored" +
             " if an extension of KafkaPrincipalBuilder is provided by the " + PRINCIPAL_BUILDER_CLASS_CONFIG + "" +
            " configuration.";
-    public static final List DEFAULT_SSL_PRINCIPAL_MAPPING_RULES = Collections.singletonList("DEFAULT");
+    public static final String DEFAULT_SSL_PRINCIPAL_MAPPING_RULES = "DEFAULT";
 
     public static final String SASL_KERBEROS_PRINCIPAL_TO_LOCAL_RULES_DOC = "A list of rules for mapping from principal " +
             "names to short names (typically operating system usernames). The rules are evaluated in order and the " +
@@ -73,7 +73,8 @@ public class BrokerSecurityConfigs {
             + " client authentication is required."
             + " 
      • ssl.client.auth=requested This means client authentication is optional."
             + " unlike requested , if this option is set client can choose not to provide authentication information about itself"
-            + " 
      • ssl.client.auth=none This means client authentication is not needed.";
+            + " 
      • ssl.client.auth=none This means client authentication is not needed."
+            + "
      ";
 
     public static final String SASL_ENABLED_MECHANISMS_DOC = "The list of SASL mechanisms enabled in the Kafka server. "
             + "The list may contain any mechanism for which a security provider is available. "
diff --git a/clients/src/main/java/org/apache/kafka/common/config/provider/ConfigProvider.java b/clients/src/main/java/org/apache/kafka/common/config/provider/ConfigProvider.java
index 8561511e64e30..fe65ddfa30656 100644
--- a/clients/src/main/java/org/apache/kafka/common/config/provider/ConfigProvider.java
+++ b/clients/src/main/java/org/apache/kafka/common/config/provider/ConfigProvider.java
@@ -25,6 +25,7 @@
 
 /**
  * A provider of configuration data, which may optionally support subscriptions to configuration changes.
+ * Implementations are required to safely support concurrent calls to any of the methods in this interface.
  */
 public interface ConfigProvider extends Configurable, Closeable {
 
diff --git a/clients/src/main/java/org/apache/kafka/common/errors/ElectionNotNeededException.java b/clients/src/main/java/org/apache/kafka/common/errors/ElectionNotNeededException.java
new file mode 100644
index 0000000000000..74fc7d670158b
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/errors/ElectionNotNeededException.java
@@ -0,0 +1,28 @@
+/*
+ * 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 ElectionNotNeededException extends InvalidMetadataException {
+
+    public ElectionNotNeededException(String message) {
+        super(message);
+    }
+
+    public ElectionNotNeededException(String message, Throwable cause) {
+        super(message, cause);
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/errors/EligibleLeadersNotAvailableException.java b/clients/src/main/java/org/apache/kafka/common/errors/EligibleLeadersNotAvailableException.java
new file mode 100644
index 0000000000000..87679652e55f1
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/errors/EligibleLeadersNotAvailableException.java
@@ -0,0 +1,28 @@
+/*
+ * 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 EligibleLeadersNotAvailableException extends InvalidMetadataException {
+
+    public EligibleLeadersNotAvailableException(String message) {
+        super(message);
+    }
+
+    public EligibleLeadersNotAvailableException(String message, Throwable cause) {
+        super(message, cause);
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/errors/FencedInstanceIdException.java b/clients/src/main/java/org/apache/kafka/common/errors/FencedInstanceIdException.java
new file mode 100644
index 0000000000000..78e4034a24a2f
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/errors/FencedInstanceIdException.java
@@ -0,0 +1,29 @@
+/*
+ * 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 FencedInstanceIdException extends ApiException {
+    private static final long serialVersionUID = 1L;
+
+    public FencedInstanceIdException(String message) {
+        super(message);
+    }
+
+    public FencedInstanceIdException(String message, Throwable cause) {
+        super(message, cause);
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/errors/GroupAuthorizationException.java b/clients/src/main/java/org/apache/kafka/common/errors/GroupAuthorizationException.java
index c3f0795de43e5..22eae3b57b4c8 100644
--- a/clients/src/main/java/org/apache/kafka/common/errors/GroupAuthorizationException.java
+++ b/clients/src/main/java/org/apache/kafka/common/errors/GroupAuthorizationException.java
@@ -19,13 +19,27 @@
 public class GroupAuthorizationException extends AuthorizationException {
     private final String groupId;
 
-    public GroupAuthorizationException(String groupId) {
-        super("Not authorized to access group: " + groupId);
+    public GroupAuthorizationException(String message, String groupId) {
+        super(message);
         this.groupId = groupId;
     }
 
+    public GroupAuthorizationException(String message) {
+        this(message, null);
+    }
+
+    /**
+     * Return the group ID that failed authorization. May be null if it is not known
+     * in the context the exception was raised in.
+     *
+     * @return nullable groupId
+     */
     public String groupId() {
         return groupId;
     }
 
+    public static GroupAuthorizationException forGroupId(String groupId) {
+        return new GroupAuthorizationException("Not authorized to access group: " + groupId, groupId);
+    }
+
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/errors/GroupMaxSizeReachedException.java b/clients/src/main/java/org/apache/kafka/common/errors/GroupMaxSizeReachedException.java
index 086b79ff0284e..85d0c7d25ce85 100644
--- a/clients/src/main/java/org/apache/kafka/common/errors/GroupMaxSizeReachedException.java
+++ b/clients/src/main/java/org/apache/kafka/common/errors/GroupMaxSizeReachedException.java
@@ -22,7 +22,7 @@
 public class GroupMaxSizeReachedException extends ApiException {
     private static final long serialVersionUID = 1L;
 
-    public GroupMaxSizeReachedException(String groupId) {
-        super("Consumer group " + groupId + " already has the configured maximum number of members.");
+    public GroupMaxSizeReachedException(String message) {
+        super(message);
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/errors/GroupSubscribedToTopicException.java b/clients/src/main/java/org/apache/kafka/common/errors/GroupSubscribedToTopicException.java
new file mode 100644
index 0000000000000..a62fe325d8136
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/errors/GroupSubscribedToTopicException.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 GroupSubscribedToTopicException extends ApiException {
+    public GroupSubscribedToTopicException(String message) {
+        super(message);
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/errors/NoReassignmentInProgressException.java b/clients/src/main/java/org/apache/kafka/common/errors/NoReassignmentInProgressException.java
new file mode 100644
index 0000000000000..9fd8a73c80916
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/errors/NoReassignmentInProgressException.java
@@ -0,0 +1,31 @@
+/*
+ * 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;
+
+/**
+ * Thrown if a reassignment cannot be cancelled because none is in progress.
+ */
+public class NoReassignmentInProgressException extends ApiException {
+    public NoReassignmentInProgressException(String message) {
+        super(message);
+    }
+
+    public NoReassignmentInProgressException(String message, Throwable cause) {
+        super(message, cause);
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/errors/TopicAuthorizationException.java b/clients/src/main/java/org/apache/kafka/common/errors/TopicAuthorizationException.java
index 6bf260d5cde79..e2235f804e170 100644
--- a/clients/src/main/java/org/apache/kafka/common/errors/TopicAuthorizationException.java
+++ b/clients/src/main/java/org/apache/kafka/common/errors/TopicAuthorizationException.java
@@ -22,15 +22,25 @@
 public class TopicAuthorizationException extends AuthorizationException {
     private final Set unauthorizedTopics;
 
-    public TopicAuthorizationException(Set unauthorizedTopics) {
-        super("Not authorized to access topics: " + unauthorizedTopics);
+    public TopicAuthorizationException(String message, Set unauthorizedTopics) {
+        super(message);
         this.unauthorizedTopics = unauthorizedTopics;
     }
 
-    public TopicAuthorizationException(String unauthorizedTopic) {
-        this(Collections.singleton(unauthorizedTopic));
+    public TopicAuthorizationException(Set unauthorizedTopics) {
+        this("Not authorized to access topics: " + unauthorizedTopics, unauthorizedTopics);
+    }
+
+    public TopicAuthorizationException(String message) {
+        this(message, Collections.emptySet());
     }
 
+    /**
+     * Get the set of topics which failed authorization. May be empty if the set is not known
+     * in the context the exception was raised in.
+     *
+     * @return possibly empty set of unauthorized topics
+     */
     public Set unauthorizedTopics() {
         return unauthorizedTopics;
     }
diff --git a/clients/src/main/java/org/apache/kafka/common/header/internals/RecordHeader.java b/clients/src/main/java/org/apache/kafka/common/header/internals/RecordHeader.java
index a6c5375dbb220..3b73c9337ac97 100644
--- a/clients/src/main/java/org/apache/kafka/common/header/internals/RecordHeader.java
+++ b/clients/src/main/java/org/apache/kafka/common/header/internals/RecordHeader.java
@@ -60,7 +60,7 @@ public boolean equals(Object o) {
             return false;
 
         RecordHeader header = (RecordHeader) o;
-        return (key == null ? header.key == null : key.equals(header.key)) && 
+        return Objects.equals(key, header.key) &&
                Arrays.equals(value(), header.value());
     }
 
diff --git a/clients/src/main/java/org/apache/kafka/common/header/internals/RecordHeaders.java b/clients/src/main/java/org/apache/kafka/common/header/internals/RecordHeaders.java
index 577e758348d7a..1277408270800 100644
--- a/clients/src/main/java/org/apache/kafka/common/header/internals/RecordHeaders.java
+++ b/clients/src/main/java/org/apache/kafka/common/header/internals/RecordHeaders.java
@@ -16,16 +16,17 @@
  */
 package org.apache.kafka.common.header.internals;
 
+import org.apache.kafka.common.header.Header;
+import org.apache.kafka.common.header.Headers;
+import org.apache.kafka.common.record.Record;
+import org.apache.kafka.common.utils.AbstractIterator;
+
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Iterator;
 import java.util.List;
-
-import org.apache.kafka.common.header.Header;
-import org.apache.kafka.common.header.Headers;
-import org.apache.kafka.common.record.Record;
-import org.apache.kafka.common.utils.AbstractIterator;
+import java.util.Objects;
 
 public class RecordHeaders implements Headers {
 
@@ -61,6 +62,7 @@ public RecordHeaders(Iterable
       headers) {
 
     @Override
     public Headers add(Header header) throws IllegalStateException {
+        Objects.requireNonNull(header, "Header cannot be null.");
         canWrite();
         headers.add(header);
         return this;
@@ -155,7 +157,7 @@ public boolean equals(Object o) {
 
         RecordHeaders headers1 = (RecordHeaders) o;
 
-        return headers != null ? headers.equals(headers1.headers) : headers1.headers == null;
+        return Objects.equals(headers, headers1.headers);
     }
 
     @Override
diff --git a/clients/src/main/java/org/apache/kafka/common/internals/PartitionStates.java b/clients/src/main/java/org/apache/kafka/common/internals/PartitionStates.java
index 22e183ab23865..daad3550738b0 100644
--- a/clients/src/main/java/org/apache/kafka/common/internals/PartitionStates.java
+++ b/clients/src/main/java/org/apache/kafka/common/internals/PartitionStates.java
@@ -89,7 +89,7 @@ public boolean contains(TopicPartition topicPartition) {
      * Returns the partition states in order.
      */
     public List> partitionStates() {
-        List> result = new ArrayList<>();
+        List> result = new ArrayList<>(map.size());
         for (Map.Entry entry : map.entrySet()) {
             result.add(new PartitionState<>(entry.getKey(), entry.getValue()));
         }
diff --git a/clients/src/main/java/org/apache/kafka/common/internals/Topic.java b/clients/src/main/java/org/apache/kafka/common/internals/Topic.java
index 619477d326063..a5ef3357dd135 100644
--- a/clients/src/main/java/org/apache/kafka/common/internals/Topic.java
+++ b/clients/src/main/java/org/apache/kafka/common/internals/Topic.java
@@ -28,7 +28,7 @@ public class Topic {
     public static final String TRANSACTION_STATE_TOPIC_NAME = "__transaction_state";
     public static final String LEGAL_CHARS = "[a-zA-Z0-9._-]";
 
-    public static final Set INTERNAL_TOPICS = Collections.unmodifiableSet(
+    private static final Set INTERNAL_TOPICS = Collections.unmodifiableSet(
             Utils.mkSet(GROUP_METADATA_TOPIC_NAME, TRANSACTION_STATE_TOPIC_NAME));
 
     private static final int MAX_NAME_LENGTH = 249;
diff --git a/clients/src/main/java/org/apache/kafka/common/memory/GarbageCollectedMemoryPool.java b/clients/src/main/java/org/apache/kafka/common/memory/GarbageCollectedMemoryPool.java
index 18f8ffe91c66e..bda419eddca5a 100644
--- a/clients/src/main/java/org/apache/kafka/common/memory/GarbageCollectedMemoryPool.java
+++ b/clients/src/main/java/org/apache/kafka/common/memory/GarbageCollectedMemoryPool.java
@@ -42,7 +42,7 @@ public class GarbageCollectedMemoryPool extends SimpleMemoryPool implements Auto
     private volatile boolean alive = true;
 
     public GarbageCollectedMemoryPool(long sizeBytes, int maxSingleAllocationSize, boolean strict, Sensor oomPeriodSensor) {
-        super(sizeBytes, maxSingleAllocationSize, strict, oomPeriodSensor);
+        super(sizeBytes, maxSingleAllocationSize, strict, oomPeriodSensor, null);
         this.alive = true;
         this.gcListenerThread = new Thread(gcListener, "memory pool GC listener");
         this.gcListenerThread.setDaemon(true); //so we dont need to worry about shutdown
diff --git a/clients/src/main/java/org/apache/kafka/common/memory/RecyclingMemoryPool.java b/clients/src/main/java/org/apache/kafka/common/memory/RecyclingMemoryPool.java
new file mode 100644
index 0000000000000..971f0be72972c
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/memory/RecyclingMemoryPool.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.common.memory;
+
+import java.nio.ByteBuffer;
+import java.util.concurrent.LinkedBlockingQueue;
+import org.apache.kafka.common.metrics.Sensor;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * An implementation of memory pool which recycles buffers of commonly used size.
+ * This memory pool is useful if most of the requested buffers' size are within close size range.
+ * In this case, instead of deallocate and reallocate the buffer, the memory pool will recycle the buffer for future use.
+ */
+public class RecyclingMemoryPool implements MemoryPool {
+    protected static final Logger log = LoggerFactory.getLogger(RecyclingMemoryPool.class);
+    protected final int cacheableBufferSizeUpperThreshold;
+    protected final int cacheableBufferSizeLowerThreshold;
+    protected final LinkedBlockingQueue bufferCache;
+    protected final Sensor requestSensor;
+
+    public RecyclingMemoryPool(int cacheableBufferSize, int bufferCacheCapacity, Sensor requestSensor) {
+        if (bufferCacheCapacity <= 0 || cacheableBufferSize <= 0) {
+            throw new IllegalArgumentException(String.format("Must provide a positive cacheable buffer size and buffer cache " +
+                    "capacity, provided %d and %d respectively.", cacheableBufferSize, bufferCacheCapacity));
+        }
+        this.bufferCache = new LinkedBlockingQueue<>(bufferCacheCapacity);
+        for (int i = 0; i < bufferCacheCapacity; i++) {
+            bufferCache.offer(ByteBuffer.allocate(cacheableBufferSize));
+        }
+        this.cacheableBufferSizeUpperThreshold = cacheableBufferSize;
+        this.cacheableBufferSizeLowerThreshold = cacheableBufferSize / 2;
+        this.requestSensor = requestSensor;
+    }
+
+    @Override
+    public ByteBuffer tryAllocate(int sizeBytes) {
+        if (sizeBytes < 1) {
+            throw new IllegalArgumentException("requested size " + sizeBytes + "<=0");
+        }
+
+        ByteBuffer allocated = null;
+        if (sizeBytes > cacheableBufferSizeLowerThreshold  && sizeBytes <= cacheableBufferSizeUpperThreshold) {
+            allocated = bufferCache.poll();
+        }
+        if (allocated != null) {
+            allocated.limit(sizeBytes);
+        } else {
+            allocated = ByteBuffer.allocate(sizeBytes);
+        }
+        bufferToBeAllocated(allocated);
+        return allocated;
+    }
+
+    @Override
+    public void release(ByteBuffer previouslyAllocated) {
+        if (previouslyAllocated == null) {
+            throw new IllegalArgumentException("provided null buffer");
+        }
+        if (previouslyAllocated.capacity() == cacheableBufferSizeUpperThreshold) {
+            previouslyAllocated.clear();
+            bufferCache.offer(previouslyAllocated);
+        } else {
+            bufferToBeReleased(previouslyAllocated);
+        }
+    }
+
+    //allows subclasses to do their own bookkeeping (and validation) _before_ memory is returned to client code.
+    protected void bufferToBeAllocated(ByteBuffer justAllocated) {
+        try {
+            this.requestSensor.record(justAllocated.limit());
+        } catch (Exception e) {
+            log.debug("failed to record size of allocated buffer");
+        }
+        log.trace("allocated buffer of size {}", justAllocated.capacity());
+    }
+
+    //allows subclasses to do their own bookkeeping (and validation) _before_ memory is marked as reclaimed.
+    protected void bufferToBeReleased(ByteBuffer justReleased) {
+        log.trace("released buffer of size {}", justReleased.capacity());
+    }
+
+    @Override
+    public long size() {
+        return Long.MAX_VALUE;
+    }
+
+    @Override
+    public long availableMemory() {
+        return Long.MAX_VALUE;
+    }
+
+    @Override
+    public boolean isOutOfMemory() {
+        return false;
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/memory/SimpleMemoryPool.java b/clients/src/main/java/org/apache/kafka/common/memory/SimpleMemoryPool.java
index f1ab8f71854e7..f6f20808cfffb 100644
--- a/clients/src/main/java/org/apache/kafka/common/memory/SimpleMemoryPool.java
+++ b/clients/src/main/java/org/apache/kafka/common/memory/SimpleMemoryPool.java
@@ -38,8 +38,9 @@ public class SimpleMemoryPool implements MemoryPool {
     protected final int maxSingleAllocationSize;
     protected final AtomicLong startOfNoMemPeriod = new AtomicLong(); //nanoseconds
     protected volatile Sensor oomTimeSensor;
+    protected volatile Sensor allocateSensor;
 
-    public SimpleMemoryPool(long sizeInBytes, int maxSingleAllocationBytes, boolean strict, Sensor oomPeriodSensor) {
+    public SimpleMemoryPool(long sizeInBytes, int maxSingleAllocationBytes, boolean strict, Sensor oomPeriodSensor, Sensor allocateSensor) {
         if (sizeInBytes <= 0 || maxSingleAllocationBytes <= 0 || maxSingleAllocationBytes > sizeInBytes)
             throw new IllegalArgumentException("must provide a positive size and max single allocation size smaller than size."
                 + "provided " + sizeInBytes + " and " + maxSingleAllocationBytes + " respectively");
@@ -48,6 +49,7 @@ public SimpleMemoryPool(long sizeInBytes, int maxSingleAllocationBytes, boolean
         this.availableMemory = new AtomicLong(sizeInBytes);
         this.maxSingleAllocationSize = maxSingleAllocationBytes;
         this.oomTimeSensor = oomPeriodSensor;
+        this.allocateSensor = allocateSensor;
     }
 
     @Override
@@ -111,6 +113,7 @@ public boolean isOutOfMemory() {
 
     //allows subclasses to do their own bookkeeping (and validation) _before_ memory is returned to client code.
     protected void bufferToBeReturned(ByteBuffer justAllocated) {
+        this.allocateSensor.record(justAllocated.capacity());
         log.trace("allocated buffer of size {} ", justAllocated.capacity());
     }
 
diff --git a/clients/src/main/java/org/apache/kafka/common/memory/WeakMemoryPool.java b/clients/src/main/java/org/apache/kafka/common/memory/WeakMemoryPool.java
new file mode 100644
index 0000000000000..48fb1b2194dd0
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/memory/WeakMemoryPool.java
@@ -0,0 +1,138 @@
+/*
+ * 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.memory;
+
+import java.lang.ref.WeakReference;
+import java.nio.ByteBuffer;
+import java.util.LinkedList;
+import java.util.Map;
+import java.util.NavigableMap;
+import java.util.TreeMap;
+
+
+/**
+ * An implementation of memory pool with weak references.
+ *
+ * Note: This will be a no-op for all the other clients.
+ *
+ * If and when the GC runs, the buffers will be reclaimed and no side effects will be seen.
+ */
+public class WeakMemoryPool implements MemoryPool {
+    private final NavigableMap>> cache;
+    private final Object lock = new Object();
+
+    public WeakMemoryPool() {
+        cache = new TreeMap<>();
+    }
+
+    /**
+     * Tries to acquire a ByteBuffer of the specified size
+     * @param sizeBytes size required
+     * @return a ByteBuffer (which later needs to be release()ed), or null if no memory available.
+     *         the buffer will be of the exact size requested, even if backed by a larger chunk of memory
+     */
+    @Override
+    public ByteBuffer tryAllocate(int sizeBytes) {
+        if (sizeBytes < 0) {
+            throw new IllegalArgumentException("The buffer size cannot be less than 0");
+        }
+        ByteBuffer buffer = null;
+        synchronized (lock) {
+            NavigableMap>> tailMap = cache.tailMap(sizeBytes, true);
+            LinkedList keysToDelete = new LinkedList<>();
+            for (Map.Entry>> entry : tailMap.entrySet()) {
+
+                LinkedList> queue = entry.getValue();
+                while (!queue.isEmpty() && buffer == null) {
+                    buffer = queue.pop().get();
+                }
+
+                if (queue.isEmpty()) {
+                    keysToDelete.add(entry.getKey());
+                }
+
+                // If buffer is non-null; break out of the loop!
+                if (buffer != null) {
+                    break;
+                }
+            }
+
+            // Remove all keys from the map which have zero queue lengths
+            for (Integer key : keysToDelete) {
+                cache.remove(key);
+            }
+        }
+
+        // If buffer is non-null, clear it before returning back to the user!
+        if (buffer != null) {
+            buffer.clear();
+            return buffer;
+        }
+
+        return ByteBuffer.allocate(sizeBytes);
+    }
+
+    /**
+     * Returns a previously allocated buffer to the pool.
+     * @param previouslyAllocated a buffer previously returned from tryAllocate()
+     */
+    @Override
+    public void release(ByteBuffer previouslyAllocated) {
+        if (previouslyAllocated == null) {
+            throw new IllegalArgumentException("the buffer to be released cannot be null!");
+        }
+
+        synchronized (lock) {
+            LinkedList> queue =
+                cache.computeIfAbsent(previouslyAllocated.capacity(), k -> new LinkedList<>());
+            queue.add(new WeakReference<>(previouslyAllocated));
+        }
+    }
+
+    /**
+     * Returns the total size of this pool
+     * @return total size, in bytes
+     */
+    @Override
+    public long size() {
+        return Long.MAX_VALUE;
+    }
+
+    /**
+     * Returns the amount of memory available for allocation by this pool.
+     * NOTE: result may be negative (pools may over allocate to avoid starvation issues)
+     * @return bytes available
+     */
+    @Override
+    public long availableMemory() {
+        return Long.MAX_VALUE;
+    }
+
+    /**
+     * Returns true if the pool cannot currently allocate any more buffers
+     * - meaning total outstanding buffers meets or exceeds pool size and
+     * some would need to be released before further allocations are possible.
+     *
+     * This is equivalent to availableMemory() <= 0
+     * @return true if out of memory
+     */
+    @Override
+    public boolean isOutOfMemory() {
+        return false;
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/metrics/CompoundStat.java b/clients/src/main/java/org/apache/kafka/common/metrics/CompoundStat.java
index 69f3efafe2fc0..f2a7ac63bae3e 100644
--- a/clients/src/main/java/org/apache/kafka/common/metrics/CompoundStat.java
+++ b/clients/src/main/java/org/apache/kafka/common/metrics/CompoundStat.java
@@ -26,9 +26,9 @@
  */
 public interface CompoundStat extends Stat {
 
-    public List stats();
+    List stats();
 
-    public static class NamedMeasurable {
+    class NamedMeasurable {
 
         private final MetricName name;
         private final Measurable stat;
diff --git a/clients/src/main/java/org/apache/kafka/common/metrics/JmxReporter.java b/clients/src/main/java/org/apache/kafka/common/metrics/JmxReporter.java
index 8d7bdbe85a036..3f852cf1a5387 100644
--- a/clients/src/main/java/org/apache/kafka/common/metrics/JmxReporter.java
+++ b/clients/src/main/java/org/apache/kafka/common/metrics/JmxReporter.java
@@ -78,6 +78,7 @@ public void init(List metrics) {
     public boolean containsMbean(String mbeanName) {
         return mbeans.containsKey(mbeanName);
     }
+
     @Override
     public void metricChange(KafkaMetric metric) {
         synchronized (LOCK) {
diff --git a/clients/src/main/java/org/apache/kafka/common/metrics/MeasurableStat.java b/clients/src/main/java/org/apache/kafka/common/metrics/MeasurableStat.java
index aedac9a764c82..035449e7363a2 100644
--- a/clients/src/main/java/org/apache/kafka/common/metrics/MeasurableStat.java
+++ b/clients/src/main/java/org/apache/kafka/common/metrics/MeasurableStat.java
@@ -19,7 +19,7 @@
 /**
  * A MeasurableStat is a {@link Stat} that is also {@link Measurable} (i.e. can produce a single floating point value).
  * This is the interface used for most of the simple statistics such as {@link org.apache.kafka.common.metrics.stats.Avg},
- * {@link org.apache.kafka.common.metrics.stats.Max}, {@link org.apache.kafka.common.metrics.stats.Count}, etc.
+ * {@link org.apache.kafka.common.metrics.stats.Max}, {@link org.apache.kafka.common.metrics.stats.CumulativeCount}, etc.
  */
 public interface MeasurableStat extends Stat, Measurable {
 
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 91c245dd5a186..dd776cdc708e9 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
@@ -20,7 +20,6 @@
 import org.apache.kafka.common.MetricNameTemplate;
 import org.apache.kafka.common.utils.KafkaThread;
 import org.apache.kafka.common.utils.Time;
-import org.apache.kafka.common.utils.Utils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -33,6 +32,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
+import java.util.Objects;
 import java.util.Set;
 import java.util.TreeMap;
 import java.util.concurrent.ConcurrentHashMap;
@@ -77,6 +77,8 @@ public class Metrics implements Closeable {
     private final ScheduledThreadPoolExecutor metricsScheduler;
     private static final Logger log = LoggerFactory.getLogger(Metrics.class);
 
+    private volatile boolean replaceOnDuplicate = false;
+
     /**
      * Create a metrics repository with no metric reporters and default configuration.
      * Expiration of Sensors is disabled.
@@ -90,7 +92,7 @@ public Metrics() {
      * Expiration of Sensors is disabled.
      */
     public Metrics(Time time) {
-        this(new MetricConfig(), new ArrayList(0), time);
+        this(new MetricConfig(), new ArrayList<>(0), time);
     }
 
     /**
@@ -98,7 +100,7 @@ public Metrics(Time time) {
      * Expiration of Sensors is disabled.
      */
     public Metrics(MetricConfig defaultConfig, Time time) {
-        this(defaultConfig, new ArrayList(0), time);
+        this(defaultConfig, new ArrayList<>(0), time);
     }
 
 
@@ -108,7 +110,7 @@ public Metrics(MetricConfig defaultConfig, Time time) {
      * @param defaultConfig The default config to use for all metrics that don't override their config
      */
     public Metrics(MetricConfig defaultConfig) {
-        this(defaultConfig, new ArrayList(0), Time.SYSTEM);
+        this(defaultConfig, new ArrayList<>(0), Time.SYSTEM);
     }
 
     /**
@@ -134,7 +136,7 @@ public Metrics(MetricConfig defaultConfig, List reporters, Time
         this.sensors = new ConcurrentHashMap<>();
         this.metrics = new ConcurrentHashMap<>();
         this.childrenSensors = new ConcurrentHashMap<>();
-        this.reporters = Utils.notNull(reporters);
+        this.reporters = Objects.requireNonNull(reporters);
         this.time = time;
         for (MetricsReporter reporter : reporters)
             reporter.init(new ArrayList());
@@ -313,7 +315,7 @@ public MetricConfig config() {
      * @return Return the sensor or null if no such sensor exists
      */
     public Sensor getSensor(String name) {
-        return this.sensors.get(Utils.notNull(name));
+        return this.sensors.get(Objects.requireNonNull(name));
     }
 
     /**
@@ -413,7 +415,7 @@ public synchronized Sensor sensor(String name, MetricConfig config, long inactiv
                     children.add(s);
                 }
             }
-            log.debug("Added sensor with name {}", name);
+            log.trace("Added sensor with name {}", name);
         }
         return s;
     }
@@ -446,7 +448,7 @@ public void removeSensor(String name) {
                     if (sensors.remove(name, sensor)) {
                         for (KafkaMetric metric : sensor.metrics())
                             removeMetric(metric.metricName());
-                        log.debug("Removed sensor with name {}", name);
+                        log.trace("Removed sensor with name {}", name);
                         childSensors = childrenSensors.remove(sensor);
                         for (final Sensor parent : sensor.parents()) {
                             childrenSensors.getOrDefault(parent, emptyList()).remove(sensor);
@@ -500,8 +502,8 @@ public void addMetric(MetricName metricName, MetricConfig config, Measurable mea
      */
     public void addMetric(MetricName metricName, MetricConfig config, MetricValueProvider metricValueProvider) {
         KafkaMetric m = new KafkaMetric(new Object(),
-                                        Utils.notNull(metricName),
-                                        Utils.notNull(metricValueProvider),
+                                        Objects.requireNonNull(metricName),
+                                        Objects.requireNonNull(metricValueProvider),
                                         config == null ? this.config : config,
                                         time);
         registerMetric(m);
@@ -533,7 +535,7 @@ public synchronized KafkaMetric removeMetric(MetricName metricName) {
                 try {
                     reporter.metricRemoval(metric);
                 } catch (Exception e) {
-                    log.error("Error when removing metric from " + reporter.getClass().getName(), e);
+                    log.error("Error when removing metric " + metricName + " from " + reporter.getClass().getName(), e);
                 }
             }
             log.trace("Removed metric named {}", metricName);
@@ -545,7 +547,7 @@ public synchronized KafkaMetric removeMetric(MetricName metricName) {
      * Add a MetricReporter
      */
     public synchronized void addReporter(MetricsReporter reporter) {
-        Utils.notNull(reporter).init(new ArrayList<>(metrics.values()));
+        Objects.requireNonNull(reporter).init(new ArrayList<>(metrics.values()));
         this.reporters.add(reporter);
     }
 
@@ -560,8 +562,16 @@ public synchronized void removeReporter(MetricsReporter reporter) {
 
     synchronized void registerMetric(KafkaMetric metric) {
         MetricName metricName = metric.metricName();
-        if (this.metrics.containsKey(metricName))
-            throw new IllegalArgumentException("A metric named '" + metricName + "' already exists, can't register another one.");
+
+        if (this.metrics.containsKey(metricName)) {
+            if (replaceOnDuplicate) {
+                log.error("The metric " + metricName + " is being replaced since it had already been registered. Please file a bug report.");
+                removeMetric(metricName);
+            } else {
+                throw new IllegalArgumentException("A metric named '" + metricName + "' already exists, can't register another one.");
+            }
+        }
+
         this.metrics.put(metricName, metric);
         for (MetricsReporter reporter : reporters) {
             try {
@@ -636,6 +646,10 @@ public MetricName metricInstance(MetricNameTemplate template, Map metrics);
+    void init(List metrics);
 
     /**
      * This is called whenever a metric is updated or added
      * @param metric
      */
-    public void metricChange(KafkaMetric metric);
+    void metricChange(KafkaMetric metric);
 
     /**
      * This is called whenever a metric is removed
      * @param metric
      */
-    public void metricRemoval(KafkaMetric metric);
+    void metricRemoval(KafkaMetric metric);
 
     /**
      * Called when the metrics repository is closed.
      */
-    public void close();
+    void close();
 
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/metrics/Sensor.java b/clients/src/main/java/org/apache/kafka/common/metrics/Sensor.java
index 1af9419bc757a..7d52d01f20640 100644
--- a/clients/src/main/java/org/apache/kafka/common/metrics/Sensor.java
+++ b/clients/src/main/java/org/apache/kafka/common/metrics/Sensor.java
@@ -19,7 +19,6 @@
 import org.apache.kafka.common.MetricName;
 import org.apache.kafka.common.metrics.CompoundStat.NamedMeasurable;
 import org.apache.kafka.common.utils.Time;
-import org.apache.kafka.common.utils.Utils;
 
 import java.util.ArrayList;
 import java.util.HashSet;
@@ -28,6 +27,7 @@
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
 
@@ -107,7 +107,7 @@ public boolean shouldRecord(final int configId) {
            long inactiveSensorExpirationTimeSeconds, RecordingLevel recordingLevel) {
         super();
         this.registry = registry;
-        this.name = Utils.notNull(name);
+        this.name = Objects.requireNonNull(name);
         this.parents = parents == null ? new Sensor[0] : parents;
         this.metrics = new LinkedHashMap<>();
         this.stats = new ArrayList<>();
@@ -238,7 +238,7 @@ public synchronized boolean add(CompoundStat stat, MetricConfig config) {
         if (hasExpired())
             return false;
 
-        this.stats.add(Utils.notNull(stat));
+        this.stats.add(Objects.requireNonNull(stat));
         Object lock = metricLock();
         for (NamedMeasurable m : stat.stats()) {
             final KafkaMetric metric = new KafkaMetric(lock, m.name(), m.stat(), config == null ? this.config : config, time);
@@ -276,8 +276,8 @@ public synchronized boolean add(final MetricName metricName, final MeasurableSta
         } else {
             final KafkaMetric metric = new KafkaMetric(
                 metricLock(),
-                Utils.notNull(metricName),
-                Utils.notNull(stat),
+                Objects.requireNonNull(metricName),
+                Objects.requireNonNull(stat),
                 config == null ? this.config : config,
                 time
             );
diff --git a/clients/src/main/java/org/apache/kafka/common/metrics/Stat.java b/clients/src/main/java/org/apache/kafka/common/metrics/Stat.java
index 1ddf5be2ec514..fa5aa1aec3bed 100644
--- a/clients/src/main/java/org/apache/kafka/common/metrics/Stat.java
+++ b/clients/src/main/java/org/apache/kafka/common/metrics/Stat.java
@@ -27,6 +27,6 @@ public interface Stat {
      * @param value The value to record
      * @param timeMs The POSIX time in milliseconds this value occurred
      */
-    public void record(MetricConfig config, double value, long timeMs);
+    void record(MetricConfig config, double value, long timeMs);
 
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/metrics/stats/Count.java b/clients/src/main/java/org/apache/kafka/common/metrics/stats/Count.java
index 3da91c4e908c6..2bef7cf93d8bb 100644
--- a/clients/src/main/java/org/apache/kafka/common/metrics/stats/Count.java
+++ b/clients/src/main/java/org/apache/kafka/common/metrics/stats/Count.java
@@ -16,30 +16,14 @@
  */
 package org.apache.kafka.common.metrics.stats;
 
-import java.util.List;
-
-import org.apache.kafka.common.metrics.MetricConfig;
-
 /**
  * A {@link SampledStat} that maintains a simple count of what it has seen.
+ * This is a special kind of {@link WindowedSum} that always records a value of {@code 1} instead of the provided value.
+ *
+ * See also {@link CumulativeCount} for a non-sampled version of this metric.
+ *
+ * @deprecated since 2.4 . Use {@link WindowedCount} instead
  */
-public class Count extends SampledStat {
-
-    public Count() {
-        super(0);
-    }
-
-    @Override
-    protected void update(Sample sample, MetricConfig config, double value, long now) {
-        sample.value += 1.0;
-    }
-
-    @Override
-    public double combine(List samples, MetricConfig config, long now) {
-        double total = 0.0;
-        for (Sample sample : samples)
-            total += sample.value;
-        return total;
-    }
-
+@Deprecated
+public class Count extends WindowedCount {
 }
diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/metrics/CumulativeCount.java b/clients/src/main/java/org/apache/kafka/common/metrics/stats/CumulativeCount.java
similarity index 66%
rename from streams/src/main/java/org/apache/kafka/streams/processor/internals/metrics/CumulativeCount.java
rename to clients/src/main/java/org/apache/kafka/common/metrics/stats/CumulativeCount.java
index 2c12c2b6e9d41..85591b5cf726f 100644
--- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/metrics/CumulativeCount.java
+++ b/clients/src/main/java/org/apache/kafka/common/metrics/stats/CumulativeCount.java
@@ -14,25 +14,21 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.kafka.streams.processor.internals.metrics;
+package org.apache.kafka.common.metrics.stats;
 
-import org.apache.kafka.common.metrics.MeasurableStat;
 import org.apache.kafka.common.metrics.MetricConfig;
 
 /**
- * A non-SampledStat version of Count for measuring -total metrics in streams
+ * A non-sampled version of {@link WindowedCount} maintained over all time.
+ *
+ * This is a special kind of {@link CumulativeSum} that always records {@code 1} instead of the provided value.
+ * In other words, it counts the number of
+ * {@link CumulativeCount#record(MetricConfig, double, long)} invocations,
+ * instead of summing the recorded values.
  */
-public class CumulativeCount implements MeasurableStat {
-
-    private double count = 0.0;
-
+public class CumulativeCount extends CumulativeSum {
     @Override
     public void record(final MetricConfig config, final double value, final long timeMs) {
-        count += 1;
-    }
-
-    @Override
-    public double measure(final MetricConfig config, final long now) {
-        return count;
+        super.record(config, 1, timeMs);
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/metrics/stats/CumulativeSum.java b/clients/src/main/java/org/apache/kafka/common/metrics/stats/CumulativeSum.java
new file mode 100644
index 0000000000000..13f12a1bb09d4
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/metrics/stats/CumulativeSum.java
@@ -0,0 +1,50 @@
+/*
+ * 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.metrics.stats;
+
+import org.apache.kafka.common.metrics.MeasurableStat;
+import org.apache.kafka.common.metrics.MetricConfig;
+
+/**
+ * An non-sampled cumulative total maintained over all time.
+ * This is a non-sampled version of {@link WindowedSum}.
+ *
+ * See also {@link CumulativeCount} if you just want to increment the value by 1 on each recording.
+ */
+public class CumulativeSum implements MeasurableStat {
+
+    private double total;
+
+    public CumulativeSum() {
+        total = 0.0;
+    }
+
+    public CumulativeSum(double value) {
+        total = value;
+    }
+
+    @Override
+    public void record(MetricConfig config, double value, long now) {
+        total += value;
+    }
+
+    @Override
+    public double measure(MetricConfig config, long now) {
+        return total;
+    }
+
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/metrics/stats/Meter.java b/clients/src/main/java/org/apache/kafka/common/metrics/stats/Meter.java
index 91d4461d2b52f..a6bdc9f3c108f 100644
--- a/clients/src/main/java/org/apache/kafka/common/metrics/stats/Meter.java
+++ b/clients/src/main/java/org/apache/kafka/common/metrics/stats/Meter.java
@@ -23,48 +23,46 @@
 import org.apache.kafka.common.MetricName;
 import org.apache.kafka.common.metrics.CompoundStat;
 import org.apache.kafka.common.metrics.MetricConfig;
-import org.apache.kafka.common.metrics.stats.Rate.SampledTotal;
 
 
 /**
  * A compound stat that includes a rate metric and a cumulative total metric.
  */
 public class Meter implements CompoundStat {
-
     private final MetricName rateMetricName;
     private final MetricName totalMetricName;
     private final Rate rate;
-    private final Total total;
+    private final CumulativeSum total;
 
     /**
-     * Construct a Meter with seconds as time unit and {@link SampledTotal} stats for Rate
+     * Construct a Meter with seconds as time unit
      */
     public Meter(MetricName rateMetricName, MetricName totalMetricName) {
-        this(TimeUnit.SECONDS, new SampledTotal(), rateMetricName, totalMetricName);
+        this(TimeUnit.SECONDS, new WindowedSum(), rateMetricName, totalMetricName);
     }
 
     /**
-     * Construct a Meter with provided time unit and {@link SampledTotal} stats for Rate
+     * Construct a Meter with provided time unit
      */
     public Meter(TimeUnit unit, MetricName rateMetricName, MetricName totalMetricName) {
-        this(unit, new SampledTotal(), rateMetricName, totalMetricName);
+        this(unit, new WindowedSum(), rateMetricName, totalMetricName);
     }
 
     /**
-     * Construct a Meter with seconds as time unit and provided {@link SampledStat} stats for Rate
+     * Construct a Meter with seconds as time unit
      */
     public Meter(SampledStat rateStat, MetricName rateMetricName, MetricName totalMetricName) {
         this(TimeUnit.SECONDS, rateStat, rateMetricName, totalMetricName);
     }
 
     /**
-     * Construct a Meter with provided time unit and provided {@link SampledStat} stats for Rate
+     * Construct a Meter with provided time unit
      */
     public Meter(TimeUnit unit, SampledStat rateStat, MetricName rateMetricName, MetricName totalMetricName) {
-        if (!(rateStat instanceof SampledTotal) && !(rateStat instanceof Count)) {
-            throw new IllegalArgumentException("Meter is supported only for SampledTotal and Count");
+        if (!(rateStat instanceof WindowedSum)) {
+            throw new IllegalArgumentException("Meter is supported only for WindowedCount or WindowedSum.");
         }
-        this.total = new Total();
+        this.total = new CumulativeSum();
         this.rate = new Rate(unit, rateStat);
         this.rateMetricName = rateMetricName;
         this.totalMetricName = totalMetricName;
@@ -81,7 +79,7 @@ public List stats() {
     public void record(MetricConfig config, double value, long timeMs) {
         rate.record(config, value, timeMs);
         // Total metrics with Count stat should record 1.0 (as recorded in the count)
-        double totalValue = (rate.stat instanceof Count) ? 1.0 : value;
+        double totalValue = (rate.stat instanceof WindowedCount) ? 1.0 : value;
         total.record(config, totalValue, timeMs);
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/metrics/stats/Rate.java b/clients/src/main/java/org/apache/kafka/common/metrics/stats/Rate.java
index a56734cd86dfb..604f860ae4f2d 100644
--- a/clients/src/main/java/org/apache/kafka/common/metrics/stats/Rate.java
+++ b/clients/src/main/java/org/apache/kafka/common/metrics/stats/Rate.java
@@ -16,7 +16,6 @@
  */
 package org.apache.kafka.common.metrics.stats;
 
-import java.util.List;
 import java.util.Locale;
 import java.util.concurrent.TimeUnit;
 
@@ -40,7 +39,7 @@ public Rate() {
     }
 
     public Rate(TimeUnit unit) {
-        this(unit, new SampledTotal());
+        this(unit, new WindowedSum());
     }
 
     public Rate(SampledStat stat) {
@@ -115,24 +114,10 @@ private double convert(long timeMs) {
         }
     }
 
-    public static class SampledTotal extends SampledStat {
-
-        public SampledTotal() {
-            super(0.0d);
-        }
-
-        @Override
-        protected void update(Sample sample, MetricConfig config, double value, long timeMs) {
-            sample.value += value;
-        }
-
-        @Override
-        public double combine(List samples, MetricConfig config, long now) {
-            double total = 0.0;
-            for (Sample sample : samples)
-                total += sample.value;
-            return total;
-        }
-
+    /**
+     * @deprecated since 2.4 Use {@link WindowedSum} instead.
+     */
+    @Deprecated
+    public static class SampledTotal extends WindowedSum {
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/metrics/stats/Sum.java b/clients/src/main/java/org/apache/kafka/common/metrics/stats/Sum.java
index b40e9cdb649a0..17188b835c434 100644
--- a/clients/src/main/java/org/apache/kafka/common/metrics/stats/Sum.java
+++ b/clients/src/main/java/org/apache/kafka/common/metrics/stats/Sum.java
@@ -16,30 +16,14 @@
  */
 package org.apache.kafka.common.metrics.stats;
 
-import java.util.List;
-
-import org.apache.kafka.common.metrics.MetricConfig;
-
 /**
  * A {@link SampledStat} that maintains the sum of what it has seen.
+ * This is a sampled version of {@link CumulativeSum}.
+ *
+ * See also {@link WindowedCount} if you want to increment the value by 1 on each recording.
+ *
+ * @deprecated since 2.4 . Use {@link WindowedSum} instead
  */
-public class Sum extends SampledStat {
-
-    public Sum() {
-        super(0);
-    }
-
-    @Override
-    protected void update(Sample sample, MetricConfig config, double value, long now) {
-        sample.value += value;
-    }
-
-    @Override
-    public double combine(List samples, MetricConfig config, long now) {
-        double total = 0.0;
-        for (Sample sample : samples)
-            total += sample.value;
-        return total;
-    }
-
+@Deprecated
+public class Sum extends WindowedSum {
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/metrics/stats/Total.java b/clients/src/main/java/org/apache/kafka/common/metrics/stats/Total.java
index b8a83f5004bd6..23f7d040f8f90 100644
--- a/clients/src/main/java/org/apache/kafka/common/metrics/stats/Total.java
+++ b/clients/src/main/java/org/apache/kafka/common/metrics/stats/Total.java
@@ -16,32 +16,14 @@
  */
 package org.apache.kafka.common.metrics.stats;
 
-import org.apache.kafka.common.metrics.MeasurableStat;
-import org.apache.kafka.common.metrics.MetricConfig;
-
 /**
- * An un-windowed cumulative total maintained over all time.
+ * An non-sampled cumulative total maintained over all time.
+ * This is a non-sampled version of {@link WindowedSum}.
+ *
+ * See also {@link CumulativeCount} if you just want to increment the value by 1 on each recording.
+ *
+ * @deprecated since 2.4 . Use {@link CumulativeSum} instead.
  */
-public class Total implements MeasurableStat {
-
-    private double total;
-
-    public Total() {
-        this.total = 0.0;
-    }
-
-    public Total(double value) {
-        this.total = value;
-    }
-
-    @Override
-    public void record(MetricConfig config, double value, long now) {
-        this.total += value;
-    }
-
-    @Override
-    public double measure(MetricConfig config, long now) {
-        return this.total;
-    }
-
-}
+@Deprecated
+public class Total extends CumulativeSum {
+}
\ No newline at end of file
diff --git a/clients/src/main/java/org/apache/kafka/common/metrics/stats/WindowedCount.java b/clients/src/main/java/org/apache/kafka/common/metrics/stats/WindowedCount.java
new file mode 100644
index 0000000000000..825f404ef02d1
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/metrics/stats/WindowedCount.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.common.metrics.stats;
+
+import org.apache.kafka.common.metrics.MetricConfig;
+
+/**
+ * A {@link SampledStat} that maintains a simple count of what it has seen.
+ * This is a special kind of {@link WindowedSum} that always records a value of {@code 1} instead of the provided value.
+ * In other words, it counts the number of
+ * {@link WindowedCount#record(MetricConfig, double, long)} invocations,
+ * instead of summing the recorded values.
+ *
+ * See also {@link CumulativeCount} for a non-sampled version of this metric.
+ */
+public class WindowedCount extends WindowedSum {
+    @Override
+    protected void update(Sample sample, MetricConfig config, double value, long now) {
+        super.update(sample, config, 1.0, now);
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/metrics/stats/WindowedSum.java b/clients/src/main/java/org/apache/kafka/common/metrics/stats/WindowedSum.java
new file mode 100644
index 0000000000000..14aa5620fcc44
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/metrics/stats/WindowedSum.java
@@ -0,0 +1,48 @@
+/*
+ * 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.metrics.stats;
+
+import org.apache.kafka.common.metrics.MetricConfig;
+
+import java.util.List;
+
+/**
+ * A {@link SampledStat} that maintains the sum of what it has seen.
+ * This is a sampled version of {@link CumulativeSum}.
+ *
+ * See also {@link WindowedCount} if you want to increment the value by 1 on each recording.
+ */
+public class WindowedSum extends SampledStat {
+
+    public WindowedSum() {
+        super(0);
+    }
+
+    @Override
+    protected void update(Sample sample, MetricConfig config, double value, long now) {
+        sample.value += value;
+    }
+
+    @Override
+    public double combine(List samples, MetricConfig config, long now) {
+        double total = 0.0;
+        for (Sample sample : samples)
+            total += sample.value;
+        return total;
+    }
+
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/network/KafkaChannel.java b/clients/src/main/java/org/apache/kafka/common/network/KafkaChannel.java
index 3bca276bc7356..df85e97cea768 100644
--- a/clients/src/main/java/org/apache/kafka/common/network/KafkaChannel.java
+++ b/clients/src/main/java/org/apache/kafka/common/network/KafkaChannel.java
@@ -571,8 +571,8 @@ public boolean maybeBeginClientReauthentication(Supplier nowNanosSupplier)
          * We've delayed getting the time as long as possible in case we don't need it,
          * but at this point we need it -- so get it now.
          */
-        long nowNanos = nowNanosSupplier.get().longValue();
-        if (nowNanos < authenticator.clientSessionReauthenticationTimeNanos().longValue())
+        long nowNanos = nowNanosSupplier.get();
+        if (nowNanos < authenticator.clientSessionReauthenticationTimeNanos())
             return false;
         swapAuthenticatorsAndBeginReauthentication(new ReauthenticationContext(authenticator, receive, nowNanos));
         receive = null;
@@ -605,7 +605,7 @@ public Long reauthenticationLatencyMs() {
      */
     public boolean serverAuthenticationSessionExpired(long nowNanos) {
         Long serverSessionExpirationTimeNanos = authenticator.serverSessionExpirationTimeNanos();
-        return serverSessionExpirationTimeNanos != null && nowNanos - serverSessionExpirationTimeNanos.longValue() > 0;
+        return serverSessionExpirationTimeNanos != null && nowNanos - serverSessionExpirationTimeNanos > 0;
     }
     
     /**
diff --git a/clients/src/main/java/org/apache/kafka/common/network/NetworkSend.java b/clients/src/main/java/org/apache/kafka/common/network/NetworkSend.java
index 8820059aa7985..791d5f5cec13c 100644
--- a/clients/src/main/java/org/apache/kafka/common/network/NetworkSend.java
+++ b/clients/src/main/java/org/apache/kafka/common/network/NetworkSend.java
@@ -24,11 +24,23 @@
 public class NetworkSend extends ByteBufferSend {
 
     public NetworkSend(String destination, ByteBuffer buffer) {
-        super(destination, sizeDelimit(buffer));
+        super(destination, sizeBuffer(buffer.remaining()), buffer);
     }
 
-    private static ByteBuffer[] sizeDelimit(ByteBuffer buffer) {
-        return new ByteBuffer[] {sizeBuffer(buffer.remaining()), buffer};
+    public NetworkSend(String destination, ByteBuffer[] buffers) {
+        super(destination, sizeDelimit(buffers));
+    }
+
+    private static ByteBuffer[] sizeDelimit(ByteBuffer[] buffers) {
+        int totalRemaining = 0;
+        for (ByteBuffer byteBuffer: buffers) {
+            totalRemaining += byteBuffer.remaining();
+        }
+
+        ByteBuffer[] result = new ByteBuffer[buffers.length + 1];
+        result[0] = sizeBuffer(totalRemaining);
+        System.arraycopy(buffers, 0, result, 1, buffers.length);
+        return result;
     }
 
     private static ByteBuffer sizeBuffer(int size) {
diff --git a/clients/src/main/java/org/apache/kafka/common/network/SaslChannelBuilder.java b/clients/src/main/java/org/apache/kafka/common/network/SaslChannelBuilder.java
index 6b7187b99a555..a2c843108be10 100644
--- a/clients/src/main/java/org/apache/kafka/common/network/SaslChannelBuilder.java
+++ b/clients/src/main/java/org/apache/kafka/common/network/SaslChannelBuilder.java
@@ -46,14 +46,11 @@
 import org.apache.kafka.common.security.scram.internals.ScramServerCallbackHandler;
 import org.apache.kafka.common.security.ssl.SslFactory;
 import org.apache.kafka.common.security.token.delegation.internals.DelegationTokenCache;
-import org.apache.kafka.common.utils.Java;
 import org.apache.kafka.common.utils.Time;
 import org.apache.kafka.common.utils.Utils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.lang.reflect.InvocationTargetException;
-import java.lang.reflect.Method;
 import java.io.IOException;
 import java.net.Socket;
 import java.nio.channels.SelectionKey;
@@ -66,6 +63,7 @@
 import java.util.function.Supplier;
 
 import javax.security.auth.Subject;
+import javax.security.auth.kerberos.KerberosPrincipal;
 
 public class SaslChannelBuilder implements ChannelBuilder, ListenerReconfigurable {
     private static final Logger log = LoggerFactory.getLogger(SaslChannelBuilder.class);
@@ -260,25 +258,9 @@ Map loginManagers() {
         return loginManagers;
     }
 
-    private static String defaultKerberosRealm() throws ClassNotFoundException, NoSuchMethodException,
-            IllegalArgumentException, IllegalAccessException, InvocationTargetException {
-
-        //TODO Find a way to avoid using these proprietary classes as access to Java 9 will block access by default
-        //due to the Jigsaw module system
-
-        Object kerbConf;
-        Class classRef;
-        Method getInstanceMethod;
-        Method getDefaultRealmMethod;
-        if (Java.isIbmJdk()) {
-            classRef = Class.forName("com.ibm.security.krb5.internal.Config");
-        } else {
-            classRef = Class.forName("sun.security.krb5.Config");
-        }
-        getInstanceMethod = classRef.getMethod("getInstance", new Class[0]);
-        kerbConf = getInstanceMethod.invoke(classRef, new Object[0]);
-        getDefaultRealmMethod = classRef.getDeclaredMethod("getDefaultRealm", new Class[0]);
-        return (String) getDefaultRealmMethod.invoke(kerbConf, new Object[0]);
+    private static String defaultKerberosRealm() {
+        // see https://issues.apache.org/jira/browse/HADOOP-10848 for details
+        return new KerberosPrincipal("tmp", 1).getRealm();
     }
 
     private void createClientCallbackHandler(Map configs) {
@@ -290,7 +272,7 @@ private void createClientCallbackHandler(Map configs) {
         saslCallbackHandlers.put(clientSaslMechanism, callbackHandler);
     }
 
-    private void createServerCallbackHandlers(Map configs) throws ClassNotFoundException {
+    private void createServerCallbackHandlers(Map configs) {
         for (String mechanism : jaasContexts.keySet()) {
             AuthenticateCallbackHandler callbackHandler;
             String prefix = ListenerName.saslMechanismPrefix(mechanism);
diff --git a/clients/src/main/java/org/apache/kafka/common/network/Selector.java b/clients/src/main/java/org/apache/kafka/common/network/Selector.java
index e431e2788d858..27c28f15a47d6 100644
--- a/clients/src/main/java/org/apache/kafka/common/network/Selector.java
+++ b/clients/src/main/java/org/apache/kafka/common/network/Selector.java
@@ -23,13 +23,14 @@
 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.Count;
+import org.apache.kafka.common.metrics.stats.CumulativeSum;
 import org.apache.kafka.common.metrics.stats.Max;
 import org.apache.kafka.common.metrics.stats.Meter;
+import org.apache.kafka.common.metrics.stats.WindowedCount;
 import org.apache.kafka.common.metrics.stats.SampledStat;
-import org.apache.kafka.common.metrics.stats.Total;
 import org.apache.kafka.common.utils.LogContext;
 import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.common.utils.Utils;
 import org.slf4j.Logger;
 
 import java.io.IOException;
@@ -52,6 +53,7 @@
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
 
 /**
  * A nioSelector interface for doing non-blocking multi-connection network I/O.
@@ -358,15 +360,24 @@ public void wakeup() {
     @Override
     public void close() {
         List connections = new ArrayList<>(channels.keySet());
-        for (String id : connections)
-            close(id);
         try {
-            this.nioSelector.close();
-        } catch (IOException | SecurityException e) {
-            log.error("Exception closing nioSelector:", e);
+            for (String id : connections)
+                close(id);
+        } finally {
+            // If there is any exception thrown in close(id), we should still be able
+            // to close the remaining objects, especially the sensors because keeping
+            // the sensors may lead to failure to start up the ReplicaFetcherThread if
+            // the old sensors with the same names has not yet been cleaned up.
+            AtomicReference firstException = new AtomicReference<>();
+            Utils.closeQuietly(nioSelector, "nioSelector", firstException);
+            Utils.closeQuietly(sensors, "sensors", firstException);
+            Utils.closeQuietly(channelBuilder, "channelBuilder", firstException);
+            Throwable exception = firstException.get();
+            if (exception instanceof RuntimeException && !(exception instanceof SecurityException)) {
+                throw (RuntimeException) exception;
+            }
+
         }
-        sensors.close();
-        channelBuilder.close();
     }
 
     /**
@@ -919,6 +930,28 @@ public KafkaChannel closingChannel(String id) {
         return closingChannels.get(id);
     }
 
+    /**
+     * Returns the lowest priority channel chosen using the following sequence:
+     *   1) If one or more channels are in closing state, return any one of them
+     *   2) If idle expiry manager is enabled, return the least recently updated channel
+     *   3) Otherwise return any of the channels
+     *
+     * This method is used to close a channel to accommodate a new channel on the inter-broker listener
+     * when broker-wide `max.connections` limit is enabled.
+     */
+    public KafkaChannel lowestPriorityChannel() {
+        KafkaChannel channel = null;
+        if (!closingChannels.isEmpty()) {
+            channel = closingChannels.values().iterator().next();
+        } else if (idleExpiryManager != null && !idleExpiryManager.lruConnections.isEmpty()) {
+            String channelId = idleExpiryManager.lruConnections.keySet().iterator().next();
+            channel = channel(channelId);
+        } else if (!channels.isEmpty()) {
+            channel = channels.values().iterator().next();
+        }
+        return channel;
+    }
+
     /**
      * Get the channel associated with selectionKey
      */
@@ -997,6 +1030,8 @@ private class SelectorMetrics implements AutoCloseable {
         private final String metricGrpPrefix;
         private final Map metricTags;
         private final boolean metricsPerConnection;
+        private final String metricGrpName;
+        private final String perConnectionMetricGrpName;
 
         public final Sensor connectionClosed;
         public final Sensor connectionCreated;
@@ -1021,7 +1056,8 @@ public SelectorMetrics(Metrics metrics, String metricGrpPrefix, Map tag: metricTags.entrySet()) {
@@ -1051,7 +1087,7 @@ public SelectorMetrics(Metrics metrics, String metricGrpPrefix, Map metricTags,
                 SampledStat stat, String baseName, String descriptiveName) {
-            MetricName rateMetricName = metrics.metricName(baseName + "-rate", groupName,
-                            String.format("The number of %s per second", descriptiveName), metricTags);
-            MetricName totalMetricName = metrics.metricName(baseName + "-total", groupName,
-                            String.format("The total number of %s", descriptiveName), metricTags);
+            // Reason for using intern here: These strings are very repetitive and can lead to 10K+ identical objects in
+            // memory. They also low QPS (access). Hence, it is ok to intern() them.
+            MetricName rateMetricName = metrics.metricName((baseName + "-rate").intern(), groupName,
+                            String.format("The number of %s per second", descriptiveName).intern(), metricTags);
+            MetricName totalMetricName = metrics.metricName((baseName + "-total").intern(), groupName,
+                            String.format("The total number of %s", descriptiveName).intern(), metricTags);
             if (stat == null)
                 return new Meter(rateMetricName, totalMetricName);
             else
@@ -1147,29 +1185,27 @@ public void maybeRegisterConnectionMetrics(String connectionId) {
                 String nodeRequestName = "node-" + connectionId + ".bytes-sent";
                 Sensor nodeRequest = this.metrics.getSensor(nodeRequestName);
                 if (nodeRequest == null) {
-                    String metricGrpName = metricGrpPrefix + "-node-metrics";
-
                     Map tags = new LinkedHashMap<>(metricTags);
                     tags.put("node-id", "node-" + connectionId);
 
                     nodeRequest = sensor(nodeRequestName);
-                    nodeRequest.add(createMeter(metrics, metricGrpName, tags, "outgoing-byte", "outgoing bytes"));
-                    nodeRequest.add(createMeter(metrics, metricGrpName, tags, new Count(), "request", "requests sent"));
-                    MetricName metricName = metrics.metricName("request-size-avg", metricGrpName, "The average size of requests sent.", tags);
+                    nodeRequest.add(createMeter(metrics, perConnectionMetricGrpName, tags, "outgoing-byte", "outgoing bytes"));
+                    nodeRequest.add(createMeter(metrics, perConnectionMetricGrpName, tags, new WindowedCount(), "request", "requests sent"));
+                    MetricName metricName = metrics.metricName("request-size-avg", perConnectionMetricGrpName, "The average size of requests sent.", tags);
                     nodeRequest.add(metricName, new Avg());
-                    metricName = metrics.metricName("request-size-max", metricGrpName, "The maximum size of any request sent.", tags);
+                    metricName = metrics.metricName("request-size-max", perConnectionMetricGrpName, "The maximum size of any request sent.", tags);
                     nodeRequest.add(metricName, new Max());
 
                     String nodeResponseName = "node-" + connectionId + ".bytes-received";
                     Sensor nodeResponse = sensor(nodeResponseName);
-                    nodeResponse.add(createMeter(metrics, metricGrpName, tags, "incoming-byte", "incoming bytes"));
-                    nodeResponse.add(createMeter(metrics, metricGrpName, tags, new Count(), "response", "responses received"));
+                    nodeResponse.add(createMeter(metrics, perConnectionMetricGrpName, tags, "incoming-byte", "incoming bytes"));
+                    nodeResponse.add(createMeter(metrics, perConnectionMetricGrpName, tags, new WindowedCount(), "response", "responses received"));
 
                     String nodeTimeName = "node-" + connectionId + ".latency";
                     Sensor nodeRequestTime = sensor(nodeTimeName);
-                    metricName = metrics.metricName("request-latency-avg", metricGrpName, tags);
+                    metricName = metrics.metricName("request-latency-avg", perConnectionMetricGrpName, tags);
                     nodeRequestTime.add(metricName, new Avg());
-                    metricName = metrics.metricName("request-latency-max", metricGrpName, tags);
+                    metricName = metrics.metricName("request-latency-max", perConnectionMetricGrpName, tags);
                     nodeRequestTime.add(metricName, new Max());
                 }
             }
diff --git a/clients/src/main/java/org/apache/kafka/common/network/SslChannelBuilder.java b/clients/src/main/java/org/apache/kafka/common/network/SslChannelBuilder.java
index 50081999a27cb..9c6816124132e 100644
--- a/clients/src/main/java/org/apache/kafka/common/network/SslChannelBuilder.java
+++ b/clients/src/main/java/org/apache/kafka/common/network/SslChannelBuilder.java
@@ -35,7 +35,6 @@
 import java.net.InetSocketAddress;
 import java.nio.channels.SelectionKey;
 import java.nio.channels.SocketChannel;
-import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.function.Supplier;
@@ -51,7 +50,7 @@ public class SslChannelBuilder implements ChannelBuilder, ListenerReconfigurable
     private SslPrincipalMapper sslPrincipalMapper;
 
     /**
-     * Constructs a SSL channel builder. ListenerName is provided only
+     * Constructs an SSL channel builder. ListenerName is provided only
      * for server channel builder and will be null for client channel builder.
      */
     public SslChannelBuilder(Mode mode, ListenerName listenerName, boolean isInterBrokerListener) {
@@ -63,8 +62,7 @@ public SslChannelBuilder(Mode mode, ListenerName listenerName, boolean isInterBr
     public void configure(Map configs) throws KafkaException {
         try {
             this.configs = configs;
-            @SuppressWarnings("unchecked")
-            List sslPrincipalMappingRules = (List) configs.get(BrokerSecurityConfigs.SSL_PRINCIPAL_MAPPING_RULES_CONFIG);
+            String sslPrincipalMappingRules = (String) configs.get(BrokerSecurityConfigs.SSL_PRINCIPAL_MAPPING_RULES_CONFIG);
             if (sslPrincipalMappingRules != null)
                 sslPrincipalMapper = SslPrincipalMapper.fromRules(sslPrincipalMappingRules);
             this.sslFactory = new SslFactory(mode, null, isInterBrokerListener);
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 9410fddef5080..adcbcb22fc688 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
@@ -152,7 +152,7 @@ public boolean isConnected() {
 
 
     /**
-    * Sends a SSL close message and closes socketChannel.
+    * Sends an SSL close message and closes socketChannel.
     */
     @Override
     public void close() throws IOException {
@@ -316,9 +316,9 @@ private void doHandshake() throws IOException {
                     netWriteBuffer.compact();
                     netWriteBuffer = Utils.ensureCapacity(netWriteBuffer, currentNetWriteBufferSize);
                     netWriteBuffer.flip();
-                    if (netWriteBuffer.limit() >= currentNetWriteBufferSize) {
+                    if (netWriteBuffer.limit() > currentNetWriteBufferSize) {
                         throw new IllegalStateException("Buffer overflow when available data size (" + netWriteBuffer.limit() +
-                                                        ") >= network buffer size (" + currentNetWriteBufferSize + ")");
+                                                        ") > network buffer size (" + currentNetWriteBufferSize + ")");
                     }
                 } else if (handshakeResult.getStatus() == Status.BUFFER_UNDERFLOW) {
                     throw new IllegalStateException("Should not have received BUFFER_UNDERFLOW during handshake WRAP.");
@@ -350,7 +350,7 @@ private void doHandshake() throws IOException {
                 if (handshakeResult.getStatus() == Status.BUFFER_UNDERFLOW) {
                     int currentNetReadBufferSize = netReadBufferSize();
                     netReadBuffer = Utils.ensureCapacity(netReadBuffer, currentNetReadBufferSize);
-                    if (netReadBuffer.position() >= currentNetReadBufferSize) {
+                    if (netReadBuffer.position() > currentNetReadBufferSize) {
                         throw new IllegalStateException("Buffer underflow when there is available data");
                     }
                 } else if (handshakeResult.getStatus() == Status.CLOSED) {
@@ -422,7 +422,7 @@ private void handshakeFinished() throws IOException {
                 key.interestOps(key.interestOps() & ~SelectionKey.OP_WRITE);
                 SSLSession session = sslEngine.getSession();
                 log.debug("SSL handshake completed successfully with peerHost '{}' peerPort {} peerPrincipal '{}' cipherSuite '{}'",
-                        session.getPeerHost(), session.getPeerPort(), peerPrincipal(), session.getCipherSuite());
+                            session.getPeerHost(), session.getPeerPort(), peerPrincipal(), session.getCipherSuite());
             }
 
             log.trace("SSLHandshake FINISHED channelId {}, appReadBuffer pos {}, netReadBuffer pos {}, netWriteBuffer pos {} ",
@@ -547,9 +547,9 @@ public int read(ByteBuffer dst) throws IOException {
                 } else if (unwrapResult.getStatus() == Status.BUFFER_OVERFLOW) {
                     int currentApplicationBufferSize = applicationBufferSize();
                     appReadBuffer = Utils.ensureCapacity(appReadBuffer, currentApplicationBufferSize);
-                    if (appReadBuffer.position() >= currentApplicationBufferSize) {
+                    if (appReadBuffer.position() > currentApplicationBufferSize) {
                         throw new IllegalStateException("Buffer overflow when available data size (" + appReadBuffer.position() +
-                                                        ") >= application buffer size (" + currentApplicationBufferSize + ")");
+                                                        ") > application buffer size (" + currentApplicationBufferSize + ")");
                     }
 
                     // appReadBuffer will extended upto currentApplicationBufferSize
@@ -562,7 +562,7 @@ public int read(ByteBuffer dst) throws IOException {
                 } else if (unwrapResult.getStatus() == Status.BUFFER_UNDERFLOW) {
                     int currentNetReadBufferSize = netReadBufferSize();
                     netReadBuffer = Utils.ensureCapacity(netReadBuffer, currentNetReadBufferSize);
-                    if (netReadBuffer.position() >= currentNetReadBufferSize) {
+                    if (netReadBuffer.position() > currentNetReadBufferSize) {
                         throw new IllegalStateException("Buffer underflow when available data size (" + netReadBuffer.position() +
                                                         ") > packet buffer size (" + currentNetReadBufferSize + ")");
                     }
@@ -667,8 +667,8 @@ public int write(ByteBuffer src) throws IOException {
             netWriteBuffer.compact();
             netWriteBuffer = Utils.ensureCapacity(netWriteBuffer, currentNetWriteBufferSize);
             netWriteBuffer.flip();
-            if (netWriteBuffer.limit() >= currentNetWriteBufferSize)
-                throw new IllegalStateException("SSL BUFFER_OVERFLOW when available data size (" + netWriteBuffer.limit() + ") >= network buffer size (" + currentNetWriteBufferSize + ")");
+            if (netWriteBuffer.limit() > currentNetWriteBufferSize)
+                throw new IllegalStateException("SSL BUFFER_OVERFLOW when available data size (" + netWriteBuffer.limit() + ") > network buffer size (" + currentNetWriteBufferSize + ")");
         } else if (wrapResult.getStatus() == Status.BUFFER_UNDERFLOW) {
             throw new IllegalStateException("SSL BUFFER_UNDERFLOW during write");
         } else if (wrapResult.getStatus() == Status.CLOSED) {
@@ -737,7 +737,7 @@ public Principal peerPrincipal() {
     }
 
     /**
-     * returns a SSL Session after the handshake is established
+     * returns an SSL Session after the handshake is established
      * throws IllegalStateException if the handshake is not established
      */
     public SSLSession sslSession() throws IllegalStateException {
@@ -863,7 +863,8 @@ private void handshakeFailure(SSLException sslException, boolean flush) throws I
     private void maybeProcessHandshakeFailure(SSLException sslException, boolean flush, IOException ioException) throws IOException {
         if (sslException instanceof SSLHandshakeException || sslException instanceof SSLProtocolException ||
                 sslException instanceof SSLPeerUnverifiedException || sslException instanceof SSLKeyException ||
-                sslException.getMessage().contains("Unrecognized SSL message"))
+                sslException.getMessage().contains("Unrecognized SSL message") ||
+                sslException.getMessage().contains("Received fatal alert: "))
             handshakeFailure(sslException, flush);
         else if (ioException == null)
             throw sslException;
diff --git a/clients/src/main/java/org/apache/kafka/common/network/TransportLayer.java b/clients/src/main/java/org/apache/kafka/common/network/TransportLayer.java
index a8a4b8730283c..b196c5be96c65 100644
--- a/clients/src/main/java/org/apache/kafka/common/network/TransportLayer.java
+++ b/clients/src/main/java/org/apache/kafka/common/network/TransportLayer.java
@@ -81,7 +81,7 @@ public interface TransportLayer extends ScatteringByteChannel, GatheringByteChan
     boolean hasPendingWrites();
 
     /**
-     * Returns `SSLSession.getPeerPrincipal()` if this is a SslTransportLayer and there is an authenticated peer,
+     * Returns `SSLSession.getPeerPrincipal()` if this is an SslTransportLayer and there is an authenticated peer,
      * `KafkaPrincipal.ANONYMOUS` is returned otherwise.
      */
     Principal peerPrincipal() throws IOException;
diff --git a/clients/src/main/java/org/apache/kafka/common/protocol/ApiKeys.java b/clients/src/main/java/org/apache/kafka/common/protocol/ApiKeys.java
index 0a19939164703..e2f7f99a082e7 100644
--- a/clients/src/main/java/org/apache/kafka/common/protocol/ApiKeys.java
+++ b/clients/src/main/java/org/apache/kafka/common/protocol/ApiKeys.java
@@ -16,18 +16,69 @@
  */
 package org.apache.kafka.common.protocol;
 
+import org.apache.kafka.common.message.ApiMessageType;
+import org.apache.kafka.common.message.ApiVersionsRequestData;
+import org.apache.kafka.common.message.ApiVersionsResponseData;
+import org.apache.kafka.common.message.ControlledShutdownRequestData;
+import org.apache.kafka.common.message.ControlledShutdownResponseData;
+import org.apache.kafka.common.message.CreateDelegationTokenRequestData;
+import org.apache.kafka.common.message.CreateDelegationTokenResponseData;
 import org.apache.kafka.common.message.CreateTopicsRequestData;
 import org.apache.kafka.common.message.CreateTopicsResponseData;
+import org.apache.kafka.common.message.DeleteGroupsRequestData;
+import org.apache.kafka.common.message.DeleteGroupsResponseData;
+import org.apache.kafka.common.message.DeleteTopicsRequestData;
+import org.apache.kafka.common.message.DeleteTopicsResponseData;
+import org.apache.kafka.common.message.DescribeDelegationTokenRequestData;
+import org.apache.kafka.common.message.DescribeDelegationTokenResponseData;
 import org.apache.kafka.common.message.DescribeGroupsRequestData;
 import org.apache.kafka.common.message.DescribeGroupsResponseData;
-import org.apache.kafka.common.message.ElectPreferredLeadersRequestData;
-import org.apache.kafka.common.message.ElectPreferredLeadersResponseData;
+import org.apache.kafka.common.message.ElectLeadersRequestData;
+import org.apache.kafka.common.message.ElectLeadersResponseData;
+import org.apache.kafka.common.message.ExpireDelegationTokenRequestData;
+import org.apache.kafka.common.message.ExpireDelegationTokenResponseData;
+import org.apache.kafka.common.message.FindCoordinatorRequestData;
+import org.apache.kafka.common.message.FindCoordinatorResponseData;
+import org.apache.kafka.common.message.HeartbeatRequestData;
+import org.apache.kafka.common.message.HeartbeatResponseData;
+import org.apache.kafka.common.message.IncrementalAlterConfigsRequestData;
+import org.apache.kafka.common.message.IncrementalAlterConfigsResponseData;
+import org.apache.kafka.common.message.InitProducerIdRequestData;
+import org.apache.kafka.common.message.InitProducerIdResponseData;
+import org.apache.kafka.common.message.JoinGroupRequestData;
+import org.apache.kafka.common.message.JoinGroupResponseData;
+import org.apache.kafka.common.message.LeaderAndIsrRequestData;
+import org.apache.kafka.common.message.LeaderAndIsrResponseData;
 import org.apache.kafka.common.message.LeaveGroupRequestData;
 import org.apache.kafka.common.message.LeaveGroupResponseData;
+import org.apache.kafka.common.message.ListGroupsRequestData;
+import org.apache.kafka.common.message.ListGroupsResponseData;
+import org.apache.kafka.common.message.AlterPartitionReassignmentsRequestData;
+import org.apache.kafka.common.message.AlterPartitionReassignmentsResponseData;
+import org.apache.kafka.common.message.ListPartitionReassignmentsRequestData;
+import org.apache.kafka.common.message.ListPartitionReassignmentsResponseData;
+import org.apache.kafka.common.message.MetadataRequestData;
+import org.apache.kafka.common.message.MetadataResponseData;
+import org.apache.kafka.common.message.OffsetCommitRequestData;
+import org.apache.kafka.common.message.OffsetCommitResponseData;
+import org.apache.kafka.common.message.OffsetDeleteRequestData;
+import org.apache.kafka.common.message.OffsetDeleteResponseData;
+import org.apache.kafka.common.message.OffsetFetchRequestData;
+import org.apache.kafka.common.message.OffsetFetchResponseData;
+import org.apache.kafka.common.message.RenewDelegationTokenRequestData;
+import org.apache.kafka.common.message.RenewDelegationTokenResponseData;
 import org.apache.kafka.common.message.SaslAuthenticateRequestData;
 import org.apache.kafka.common.message.SaslAuthenticateResponseData;
 import org.apache.kafka.common.message.SaslHandshakeRequestData;
 import org.apache.kafka.common.message.SaslHandshakeResponseData;
+import org.apache.kafka.common.message.StopReplicaRequestData;
+import org.apache.kafka.common.message.StopReplicaResponseData;
+import org.apache.kafka.common.message.SyncGroupRequestData;
+import org.apache.kafka.common.message.SyncGroupResponseData;
+import org.apache.kafka.common.message.UpdateMetadataRequestData;
+import org.apache.kafka.common.message.UpdateMetadataResponseData;
+import org.apache.kafka.common.message.TxnOffsetCommitRequestData;
+import org.apache.kafka.common.message.TxnOffsetCommitResponseData;
 import org.apache.kafka.common.protocol.types.Schema;
 import org.apache.kafka.common.protocol.types.SchemaException;
 import org.apache.kafka.common.protocol.types.Struct;
@@ -41,72 +92,30 @@
 import org.apache.kafka.common.requests.AlterConfigsResponse;
 import org.apache.kafka.common.requests.AlterReplicaLogDirsRequest;
 import org.apache.kafka.common.requests.AlterReplicaLogDirsResponse;
-import org.apache.kafka.common.requests.ApiVersionsRequest;
-import org.apache.kafka.common.requests.ApiVersionsResponse;
-import org.apache.kafka.common.requests.ControlledShutdownRequest;
-import org.apache.kafka.common.requests.ControlledShutdownResponse;
 import org.apache.kafka.common.requests.CreateAclsRequest;
 import org.apache.kafka.common.requests.CreateAclsResponse;
-import org.apache.kafka.common.requests.CreateDelegationTokenRequest;
-import org.apache.kafka.common.requests.CreateDelegationTokenResponse;
 import org.apache.kafka.common.requests.CreatePartitionsRequest;
 import org.apache.kafka.common.requests.CreatePartitionsResponse;
 import org.apache.kafka.common.requests.DeleteAclsRequest;
 import org.apache.kafka.common.requests.DeleteAclsResponse;
-import org.apache.kafka.common.requests.DeleteGroupsRequest;
-import org.apache.kafka.common.requests.DeleteGroupsResponse;
 import org.apache.kafka.common.requests.DeleteRecordsRequest;
 import org.apache.kafka.common.requests.DeleteRecordsResponse;
-import org.apache.kafka.common.requests.DeleteTopicsRequest;
-import org.apache.kafka.common.requests.DeleteTopicsResponse;
 import org.apache.kafka.common.requests.DescribeAclsRequest;
 import org.apache.kafka.common.requests.DescribeAclsResponse;
 import org.apache.kafka.common.requests.DescribeConfigsRequest;
 import org.apache.kafka.common.requests.DescribeConfigsResponse;
-import org.apache.kafka.common.requests.DescribeDelegationTokenRequest;
-import org.apache.kafka.common.requests.DescribeDelegationTokenResponse;
 import org.apache.kafka.common.requests.DescribeLogDirsRequest;
 import org.apache.kafka.common.requests.DescribeLogDirsResponse;
 import org.apache.kafka.common.requests.EndTxnRequest;
 import org.apache.kafka.common.requests.EndTxnResponse;
-import org.apache.kafka.common.requests.ExpireDelegationTokenRequest;
-import org.apache.kafka.common.requests.ExpireDelegationTokenResponse;
 import org.apache.kafka.common.requests.FetchRequest;
 import org.apache.kafka.common.requests.FetchResponse;
-import org.apache.kafka.common.requests.FindCoordinatorRequest;
-import org.apache.kafka.common.requests.FindCoordinatorResponse;
-import org.apache.kafka.common.requests.HeartbeatRequest;
-import org.apache.kafka.common.requests.HeartbeatResponse;
-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.JoinGroupResponse;
-import org.apache.kafka.common.requests.LeaderAndIsrRequest;
-import org.apache.kafka.common.requests.LeaderAndIsrResponse;
-import org.apache.kafka.common.requests.ListGroupsRequest;
-import org.apache.kafka.common.requests.ListGroupsResponse;
 import org.apache.kafka.common.requests.ListOffsetRequest;
 import org.apache.kafka.common.requests.ListOffsetResponse;
-import org.apache.kafka.common.requests.MetadataRequest;
-import org.apache.kafka.common.requests.MetadataResponse;
-import org.apache.kafka.common.requests.OffsetCommitRequest;
-import org.apache.kafka.common.requests.OffsetCommitResponse;
-import org.apache.kafka.common.requests.OffsetFetchRequest;
-import org.apache.kafka.common.requests.OffsetFetchResponse;
 import org.apache.kafka.common.requests.OffsetsForLeaderEpochRequest;
 import org.apache.kafka.common.requests.OffsetsForLeaderEpochResponse;
 import org.apache.kafka.common.requests.ProduceRequest;
 import org.apache.kafka.common.requests.ProduceResponse;
-import org.apache.kafka.common.requests.RenewDelegationTokenRequest;
-import org.apache.kafka.common.requests.RenewDelegationTokenResponse;
-import org.apache.kafka.common.requests.StopReplicaRequest;
-import org.apache.kafka.common.requests.StopReplicaResponse;
-import org.apache.kafka.common.requests.SyncGroupRequest;
-import org.apache.kafka.common.requests.SyncGroupResponse;
-import org.apache.kafka.common.requests.TxnOffsetCommitRequest;
-import org.apache.kafka.common.requests.TxnOffsetCommitResponse;
-import org.apache.kafka.common.requests.UpdateMetadataRequest;
-import org.apache.kafka.common.requests.UpdateMetadataResponse;
 import org.apache.kafka.common.requests.WriteTxnMarkersRequest;
 import org.apache.kafka.common.requests.WriteTxnMarkersResponse;
 
@@ -114,6 +123,8 @@
 import java.util.concurrent.atomic.AtomicBoolean;
 
 import static org.apache.kafka.common.protocol.types.Type.BYTES;
+import static org.apache.kafka.common.protocol.types.Type.COMPACT_BYTES;
+import static org.apache.kafka.common.protocol.types.Type.COMPACT_NULLABLE_BYTES;
 import static org.apache.kafka.common.protocol.types.Type.NULLABLE_BYTES;
 import static org.apache.kafka.common.protocol.types.Type.RECORDS;
 
@@ -124,26 +135,25 @@ public enum ApiKeys {
     PRODUCE(0, "Produce", ProduceRequest.schemaVersions(), ProduceResponse.schemaVersions()),
     FETCH(1, "Fetch", FetchRequest.schemaVersions(), FetchResponse.schemaVersions()),
     LIST_OFFSETS(2, "ListOffsets", ListOffsetRequest.schemaVersions(), ListOffsetResponse.schemaVersions()),
-    METADATA(3, "Metadata", MetadataRequest.schemaVersions(), MetadataResponse.schemaVersions()),
-    LEADER_AND_ISR(4, "LeaderAndIsr", true, LeaderAndIsrRequest.schemaVersions(), LeaderAndIsrResponse.schemaVersions()),
-    STOP_REPLICA(5, "StopReplica", true, StopReplicaRequest.schemaVersions(), StopReplicaResponse.schemaVersions()),
-    UPDATE_METADATA(6, "UpdateMetadata", true, UpdateMetadataRequest.schemaVersions(),
-            UpdateMetadataResponse.schemaVersions()),
-    CONTROLLED_SHUTDOWN(7, "ControlledShutdown", true, ControlledShutdownRequest.schemaVersions(),
-            ControlledShutdownResponse.schemaVersions()),
-    OFFSET_COMMIT(8, "OffsetCommit", OffsetCommitRequest.schemaVersions(), OffsetCommitResponse.schemaVersions()),
-    OFFSET_FETCH(9, "OffsetFetch", OffsetFetchRequest.schemaVersions(), OffsetFetchResponse.schemaVersions()),
-    FIND_COORDINATOR(10, "FindCoordinator", FindCoordinatorRequest.schemaVersions(),
-            FindCoordinatorResponse.schemaVersions()),
-    JOIN_GROUP(11, "JoinGroup", JoinGroupRequest.schemaVersions(), JoinGroupResponse.schemaVersions()),
-    HEARTBEAT(12, "Heartbeat", HeartbeatRequest.schemaVersions(), HeartbeatResponse.schemaVersions()),
+    METADATA(3, "Metadata", MetadataRequestData.SCHEMAS, MetadataResponseData.SCHEMAS),
+    LEADER_AND_ISR(4, "LeaderAndIsr", true, LeaderAndIsrRequestData.SCHEMAS, LeaderAndIsrResponseData.SCHEMAS),
+    STOP_REPLICA(5, "StopReplica", true, StopReplicaRequestData.SCHEMAS, StopReplicaResponseData.SCHEMAS),
+    UPDATE_METADATA(6, "UpdateMetadata", true, UpdateMetadataRequestData.SCHEMAS, UpdateMetadataResponseData.SCHEMAS),
+    CONTROLLED_SHUTDOWN(7, "ControlledShutdown", true, ControlledShutdownRequestData.SCHEMAS,
+            ControlledShutdownResponseData.SCHEMAS),
+    OFFSET_COMMIT(8, "OffsetCommit", OffsetCommitRequestData.SCHEMAS, OffsetCommitResponseData.SCHEMAS),
+    OFFSET_FETCH(9, "OffsetFetch", OffsetFetchRequestData.SCHEMAS, OffsetFetchResponseData.SCHEMAS),
+    FIND_COORDINATOR(10, "FindCoordinator", FindCoordinatorRequestData.SCHEMAS,
+        FindCoordinatorResponseData.SCHEMAS),
+    JOIN_GROUP(11, "JoinGroup", JoinGroupRequestData.SCHEMAS, JoinGroupResponseData.SCHEMAS),
+    HEARTBEAT(12, "Heartbeat", HeartbeatRequestData.SCHEMAS, HeartbeatResponseData.SCHEMAS),
     LEAVE_GROUP(13, "LeaveGroup", LeaveGroupRequestData.SCHEMAS, LeaveGroupResponseData.SCHEMAS),
-    SYNC_GROUP(14, "SyncGroup", SyncGroupRequest.schemaVersions(), SyncGroupResponse.schemaVersions()),
+    SYNC_GROUP(14, "SyncGroup", SyncGroupRequestData.SCHEMAS, SyncGroupResponseData.SCHEMAS),
     DESCRIBE_GROUPS(15, "DescribeGroups", DescribeGroupsRequestData.SCHEMAS,
             DescribeGroupsResponseData.SCHEMAS),
-    LIST_GROUPS(16, "ListGroups", ListGroupsRequest.schemaVersions(), ListGroupsResponse.schemaVersions()),
+    LIST_GROUPS(16, "ListGroups", ListGroupsRequestData.SCHEMAS, ListGroupsResponseData.SCHEMAS),
     SASL_HANDSHAKE(17, "SaslHandshake", SaslHandshakeRequestData.SCHEMAS, SaslHandshakeResponseData.SCHEMAS),
-    API_VERSIONS(18, "ApiVersions", ApiVersionsRequest.schemaVersions(), ApiVersionsResponse.schemaVersions()) {
+    API_VERSIONS(18, "ApiVersions", ApiVersionsRequestData.SCHEMAS, ApiVersionsResponseData.SCHEMAS) {
         @Override
         public Struct parseResponse(short version, ByteBuffer buffer) {
             // Fallback to version 0 for ApiVersions response. If a client sends an ApiVersionsRequest
@@ -153,10 +163,9 @@ public Struct parseResponse(short version, ByteBuffer buffer) {
         }
     },
     CREATE_TOPICS(19, "CreateTopics", CreateTopicsRequestData.SCHEMAS, CreateTopicsResponseData.SCHEMAS),
-    DELETE_TOPICS(20, "DeleteTopics", DeleteTopicsRequest.schemaVersions(), DeleteTopicsResponse.schemaVersions()),
+    DELETE_TOPICS(20, "DeleteTopics", DeleteTopicsRequestData.SCHEMAS, DeleteTopicsResponseData.SCHEMAS),
     DELETE_RECORDS(21, "DeleteRecords", DeleteRecordsRequest.schemaVersions(), DeleteRecordsResponse.schemaVersions()),
-    INIT_PRODUCER_ID(22, "InitProducerId", InitProducerIdRequest.schemaVersions(),
-            InitProducerIdResponse.schemaVersions()),
+    INIT_PRODUCER_ID(22, "InitProducerId", InitProducerIdRequestData.SCHEMAS, InitProducerIdResponseData.SCHEMAS),
     OFFSET_FOR_LEADER_EPOCH(23, "OffsetForLeaderEpoch", false, OffsetsForLeaderEpochRequest.schemaVersions(),
             OffsetsForLeaderEpochResponse.schemaVersions()),
     ADD_PARTITIONS_TO_TXN(24, "AddPartitionsToTxn", false, RecordBatch.MAGIC_VALUE_V2,
@@ -167,8 +176,8 @@ public Struct parseResponse(short version, ByteBuffer buffer) {
             EndTxnResponse.schemaVersions()),
     WRITE_TXN_MARKERS(27, "WriteTxnMarkers", true, RecordBatch.MAGIC_VALUE_V2, WriteTxnMarkersRequest.schemaVersions(),
             WriteTxnMarkersResponse.schemaVersions()),
-    TXN_OFFSET_COMMIT(28, "TxnOffsetCommit", false, RecordBatch.MAGIC_VALUE_V2, TxnOffsetCommitRequest.schemaVersions(),
-            TxnOffsetCommitResponse.schemaVersions()),
+    TXN_OFFSET_COMMIT(28, "TxnOffsetCommit", false, RecordBatch.MAGIC_VALUE_V2, TxnOffsetCommitRequestData.SCHEMAS,
+                      TxnOffsetCommitResponseData.SCHEMAS),
     DESCRIBE_ACLS(29, "DescribeAcls", DescribeAclsRequest.schemaVersions(), DescribeAclsResponse.schemaVersions()),
     CREATE_ACLS(30, "CreateAcls", CreateAclsRequest.schemaVersions(), CreateAclsResponse.schemaVersions()),
     DELETE_ACLS(31, "DeleteAcls", DeleteAclsRequest.schemaVersions(), DeleteAclsResponse.schemaVersions()),
@@ -184,13 +193,20 @@ public Struct parseResponse(short version, ByteBuffer buffer) {
             SaslAuthenticateResponseData.SCHEMAS),
     CREATE_PARTITIONS(37, "CreatePartitions", CreatePartitionsRequest.schemaVersions(),
             CreatePartitionsResponse.schemaVersions()),
-    CREATE_DELEGATION_TOKEN(38, "CreateDelegationToken", CreateDelegationTokenRequest.schemaVersions(), CreateDelegationTokenResponse.schemaVersions()),
-    RENEW_DELEGATION_TOKEN(39, "RenewDelegationToken", RenewDelegationTokenRequest.schemaVersions(), RenewDelegationTokenResponse.schemaVersions()),
-    EXPIRE_DELEGATION_TOKEN(40, "ExpireDelegationToken", ExpireDelegationTokenRequest.schemaVersions(), ExpireDelegationTokenResponse.schemaVersions()),
-    DESCRIBE_DELEGATION_TOKEN(41, "DescribeDelegationToken", DescribeDelegationTokenRequest.schemaVersions(), DescribeDelegationTokenResponse.schemaVersions()),
-    DELETE_GROUPS(42, "DeleteGroups", DeleteGroupsRequest.schemaVersions(), DeleteGroupsResponse.schemaVersions()),
-    ELECT_PREFERRED_LEADERS(43, "ElectPreferredLeaders", ElectPreferredLeadersRequestData.SCHEMAS,
-            ElectPreferredLeadersResponseData.SCHEMAS);
+    CREATE_DELEGATION_TOKEN(38, "CreateDelegationToken", CreateDelegationTokenRequestData.SCHEMAS, CreateDelegationTokenResponseData.SCHEMAS),
+    RENEW_DELEGATION_TOKEN(39, "RenewDelegationToken", RenewDelegationTokenRequestData.SCHEMAS, RenewDelegationTokenResponseData.SCHEMAS),
+    EXPIRE_DELEGATION_TOKEN(40, "ExpireDelegationToken", ExpireDelegationTokenRequestData.SCHEMAS, ExpireDelegationTokenResponseData.SCHEMAS),
+    DESCRIBE_DELEGATION_TOKEN(41, "DescribeDelegationToken", DescribeDelegationTokenRequestData.SCHEMAS, DescribeDelegationTokenResponseData.SCHEMAS),
+    DELETE_GROUPS(42, "DeleteGroups", DeleteGroupsRequestData.SCHEMAS, DeleteGroupsResponseData.SCHEMAS),
+    ELECT_LEADERS(43, "ElectLeaders", ElectLeadersRequestData.SCHEMAS,
+            ElectLeadersResponseData.SCHEMAS),
+    INCREMENTAL_ALTER_CONFIGS(44, "IncrementalAlterConfigs", IncrementalAlterConfigsRequestData.SCHEMAS,
+                              IncrementalAlterConfigsResponseData.SCHEMAS),
+    ALTER_PARTITION_REASSIGNMENTS(45, "AlterPartitionReassignments", AlterPartitionReassignmentsRequestData.SCHEMAS,
+                                  AlterPartitionReassignmentsResponseData.SCHEMAS),
+    LIST_PARTITION_REASSIGNMENTS(46, "ListPartitionReassignments", ListPartitionReassignmentsRequestData.SCHEMAS,
+                                 ListPartitionReassignmentsResponseData.SCHEMAS),
+    OFFSET_DELETE(47, "OffsetDelete", OffsetDeleteRequestData.SCHEMAS, OffsetDeleteResponseData.SCHEMAS);
 
     private static final ApiKeys[] ID_TO_TYPE;
     private static final int MIN_API_KEY = 0;
@@ -321,6 +337,14 @@ public boolean isVersionSupported(short apiVersion) {
         return apiVersion >= oldestVersion() && apiVersion <= latestVersion();
     }
 
+    public short requestHeaderVersion(short apiVersion) {
+        return ApiMessageType.fromApiKey(id).requestHeaderVersion(apiVersion);
+    }
+
+    public short responseHeaderVersion(short apiVersion) {
+        return ApiMessageType.fromApiKey(id).responseHeaderVersion(apiVersion);
+    }
+
     private static String toHtml() {
         final StringBuilder b = new StringBuilder();
         b.append("\n");
@@ -351,7 +375,8 @@ private static boolean retainsBufferReference(Schema schema) {
         Schema.Visitor detector = new Schema.Visitor() {
             @Override
             public void visit(Type field) {
-                if (field == BYTES || field == NULLABLE_BYTES || field == RECORDS)
+                if (field == BYTES || field == NULLABLE_BYTES || field == RECORDS ||
+                    field == COMPACT_BYTES || field == COMPACT_NULLABLE_BYTES)
                     hasBuffer.set(true);
             }
         };
diff --git a/clients/src/main/java/org/apache/kafka/common/protocol/ByteBufferAccessor.java b/clients/src/main/java/org/apache/kafka/common/protocol/ByteBufferAccessor.java
index 60a923a70362d..f3dbb87f45cd6 100644
--- a/clients/src/main/java/org/apache/kafka/common/protocol/ByteBufferAccessor.java
+++ b/clients/src/main/java/org/apache/kafka/common/protocol/ByteBufferAccessor.java
@@ -17,6 +17,8 @@
 
 package org.apache.kafka.common.protocol;
 
+import org.apache.kafka.common.utils.ByteUtils;
+
 import java.nio.ByteBuffer;
 
 public class ByteBufferAccessor implements Readable, Writable {
@@ -51,6 +53,11 @@ public void readArray(byte[] arr) {
         buf.get(arr);
     }
 
+    @Override
+    public int readUnsignedVarint() {
+        return ByteUtils.readUnsignedVarint(buf);
+    }
+
     @Override
     public void writeByte(byte val) {
         buf.put(val);
@@ -72,7 +79,12 @@ public void writeLong(long val) {
     }
 
     @Override
-    public void writeArray(byte[] arr) {
+    public void writeByteArray(byte[] arr) {
         buf.put(arr);
     }
+
+    @Override
+    public void writeUnsignedVarint(int i) {
+        ByteUtils.writeUnsignedVarint(i, buf);
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/protocol/CommonFields.java b/clients/src/main/java/org/apache/kafka/common/protocol/CommonFields.java
index 708500c65e6ce..5da38aa13ed53 100644
--- a/clients/src/main/java/org/apache/kafka/common/protocol/CommonFields.java
+++ b/clients/src/main/java/org/apache/kafka/common/protocol/CommonFields.java
@@ -36,9 +36,6 @@ public class CommonFields {
 
     // Group APIs
     public static final Field.Str GROUP_ID = new Field.Str("group_id", "The unique group identifier");
-    public static final Field.Int32 GENERATION_ID = new Field.Int32("generation_id", "The generation of the group.");
-    public static final Field.Str MEMBER_ID = new Field.Str("member_id", "The member id assigned by the group " +
-            "coordinator or null if joining for the first time.");
 
     // Transactional APIs
     public static final Field.Str TRANSACTIONAL_ID = new Field.Str("transactional_id", "The transactional id corresponding to the transaction.");
@@ -51,7 +48,7 @@ public class CommonFields {
     public static final Field.Int8 RESOURCE_TYPE = new Field.Int8("resource_type", "The resource type");
     public static final Field.Str RESOURCE_NAME = new Field.Str("resource_name", "The resource name");
     public static final Field.NullableStr RESOURCE_NAME_FILTER = new Field.NullableStr("resource_name", "The resource name filter");
-    public static final Field.Int8 RESOURCE_PATTERN_TYPE = new Field.Int8("resource_pattten_type", "The resource pattern type", PatternType.LITERAL.code());
+    public static final Field.Int8 RESOURCE_PATTERN_TYPE = new Field.Int8("resource_pattern_type", "The resource pattern type", PatternType.LITERAL.code());
     public static final Field.Int8 RESOURCE_PATTERN_TYPE_FILTER = new Field.Int8("resource_pattern_type_filter", "The resource pattern type filter", PatternType.LITERAL.code());
     public static final Field.Str PRINCIPAL = new Field.Str("principal", "The ACL principal");
     public static final Field.NullableStr PRINCIPAL_FILTER = new Field.NullableStr("principal", "The ACL principal filter");
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 cc493c22f199a..13ce6684b1c16 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
@@ -16,10 +16,12 @@
  */
 package org.apache.kafka.common.protocol;
 
+import org.apache.kafka.common.InvalidRecordException;
 import org.apache.kafka.common.errors.ApiException;
 import org.apache.kafka.common.errors.BrokerNotAvailableException;
 import org.apache.kafka.common.errors.ClusterAuthorizationException;
 import org.apache.kafka.common.errors.ConcurrentTransactionsException;
+import org.apache.kafka.common.errors.GroupSubscribedToTopicException;
 import org.apache.kafka.common.errors.ControllerMovedException;
 import org.apache.kafka.common.errors.CoordinatorLoadInProgressException;
 import org.apache.kafka.common.errors.CoordinatorNotAvailableException;
@@ -60,8 +62,12 @@
 import org.apache.kafka.common.errors.KafkaStorageException;
 import org.apache.kafka.common.errors.LeaderNotAvailableException;
 import org.apache.kafka.common.errors.LogDirNotFoundException;
+import org.apache.kafka.common.errors.FencedInstanceIdException;
 import org.apache.kafka.common.errors.MemberIdRequiredException;
+import org.apache.kafka.common.errors.ElectionNotNeededException;
+import org.apache.kafka.common.errors.EligibleLeadersNotAvailableException;
 import org.apache.kafka.common.errors.NetworkException;
+import org.apache.kafka.common.errors.NoReassignmentInProgressException;
 import org.apache.kafka.common.errors.NotControllerException;
 import org.apache.kafka.common.errors.NotCoordinatorException;
 import org.apache.kafka.common.errors.NotEnoughReplicasAfterAppendException;
@@ -182,10 +188,8 @@ public enum Errors {
             RebalanceInProgressException::new),
     INVALID_COMMIT_OFFSET_SIZE(28, "The committing offset data size is not valid.",
             InvalidCommitOffsetSizeException::new),
-    TOPIC_AUTHORIZATION_FAILED(29, "Topic authorization failed.",
-            TopicAuthorizationException::new),
-    GROUP_AUTHORIZATION_FAILED(30, "Group authorization failed.",
-            GroupAuthorizationException::new),
+    TOPIC_AUTHORIZATION_FAILED(29, "Topic authorization failed.", TopicAuthorizationException::new),
+    GROUP_AUTHORIZATION_FAILED(30, "Group authorization failed.", GroupAuthorizationException::new),
     CLUSTER_AUTHORIZATION_FAILED(31, "Cluster authorization failed.",
             ClusterAuthorizationException::new),
     INVALID_TIMESTAMP(32, "The timestamp of the message is out of acceptable range.",
@@ -302,7 +306,18 @@ public enum Errors {
             MemberIdRequiredException::new),
     PREFERRED_LEADER_NOT_AVAILABLE(80, "The preferred leader was not available",
             PreferredLeaderNotAvailableException::new),
-    GROUP_MAX_SIZE_REACHED(81, "The consumer group has reached its max size.", GroupMaxSizeReachedException::new);
+    GROUP_MAX_SIZE_REACHED(81, "The consumer group has reached its max size.", GroupMaxSizeReachedException::new),
+    FENCED_INSTANCE_ID(82, "The broker rejected this static consumer since " +
+            "another consumer with the same group.instance.id has registered with a different member.id.",
+            FencedInstanceIdException::new),
+    ELIGIBLE_LEADERS_NOT_AVAILABLE(83, "Eligible topic partition leaders are not available",
+            EligibleLeadersNotAvailableException::new),
+    ELECTION_NOT_NEEDED(84, "Leader election not needed for topic partition", ElectionNotNeededException::new),
+    NO_REASSIGNMENT_IN_PROGRESS(85, "No partition reassignment is in progress.",
+            NoReassignmentInProgressException::new),
+    GROUP_SUBSCRIBED_TO_TOPIC(86, "Deleting offsets of a topic is forbidden while the consumer group is actively subscribed to it.",
+        GroupSubscribedToTopicException::new),
+    INVALID_RECORD(87, "This record has failed the validation on broker and hence be rejected.", InvalidRecordException::new);
 
     private static final Logger log = LoggerFactory.getLogger(Errors.class);
 
diff --git a/clients/src/main/java/org/apache/kafka/common/protocol/Message.java b/clients/src/main/java/org/apache/kafka/common/protocol/Message.java
index 885c69239c570..b31593f1e6f0a 100644
--- a/clients/src/main/java/org/apache/kafka/common/protocol/Message.java
+++ b/clients/src/main/java/org/apache/kafka/common/protocol/Message.java
@@ -17,8 +17,11 @@
 
 package org.apache.kafka.common.protocol;
 
+import org.apache.kafka.common.protocol.types.RawTaggedField;
 import org.apache.kafka.common.protocol.types.Struct;
 
+import java.util.List;
+
 /**
  * An object that can serialize itself.  The serialization protocol is versioned.
  * Messages also implement toString, equals, and hashCode.
@@ -37,28 +40,31 @@ public interface Message {
     /**
      * Returns the number of bytes it would take to write out this message.
      *
+     * @param cache         The serialization size cache to populate.
      * @param version       The version to use.
      *
      * @throws {@see org.apache.kafka.common.errors.UnsupportedVersionException}
      *                      If the specified version is too new to be supported
      *                      by this software.
      */
-    int size(short version);
+    int size(ObjectSerializationCache cache, short version);
 
     /**
-     * Writes out this message to the given ByteBuffer.
+     * Writes out this message to the given Writable.
      *
      * @param writable      The destination writable.
+     * @param cache         The object serialization cache to use.  You must have
+     *                      previously populated the size cache using #{Message#size()}.
      * @param version       The version to use.
      *
      * @throws {@see org.apache.kafka.common.errors.UnsupportedVersionException}
      *                      If the specified version is too new to be supported
      *                      by this software.
      */
-    void write(Writable writable, short version);
+    void write(Writable writable, ObjectSerializationCache cache, short version);
 
     /**
-     * Reads this message from the given ByteBuffer.  This will overwrite all
+     * Reads this message from the given Readable.  This will overwrite all
      * relevant fields with information from the byte buffer.
      *
      * @param readable      The source readable.
@@ -89,4 +95,11 @@ public interface Message {
      *                      by this software.
      */
     Struct toStruct(short version);
+
+    /**
+     * Returns a list of tagged fields which this software can't understand.
+     *
+     * @return              The raw tagged fields.
+     */
+    List unknownTaggedFields();
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/protocol/MessageUtil.java b/clients/src/main/java/org/apache/kafka/common/protocol/MessageUtil.java
index e3fdea7feae3b..6e473e780f3c0 100644
--- a/clients/src/main/java/org/apache/kafka/common/protocol/MessageUtil.java
+++ b/clients/src/main/java/org/apache/kafka/common/protocol/MessageUtil.java
@@ -17,21 +17,26 @@
 
 package org.apache.kafka.common.protocol;
 
-import org.apache.kafka.common.utils.Utils;
-
+import java.nio.ByteBuffer;
 import java.util.Iterator;
+import java.util.UUID;
 
 public final class MessageUtil {
+    public static final UUID ZERO_UUID = new UUID(0L, 0L);
+
     /**
-     * Get the length of the UTF8 representation of a string, without allocating
-     * a byte buffer for the string.
+     * Copy a byte buffer into an array.  This will not affect the buffer's
+     * position or mark.
      */
-    public static short serializedUtf8Length(CharSequence input) {
-        int count = Utils.utf8Length(input);
-        if (count > Short.MAX_VALUE) {
-            throw new RuntimeException("String " + input + " is too long to serialize.");
+    public static byte[] byteBufferToArray(ByteBuffer buf) {
+        byte[] arr = new byte[buf.remaining()];
+        int prevPosition = buf.position();
+        try {
+            buf.get(arr);
+        } finally {
+            buf.position(prevPosition);
         }
-        return (short) count;
+        return arr;
     }
 
     public static String deepToString(Iterator iter) {
diff --git a/clients/src/main/java/org/apache/kafka/common/protocol/ObjectSerializationCache.java b/clients/src/main/java/org/apache/kafka/common/protocol/ObjectSerializationCache.java
new file mode 100644
index 0000000000000..98addf443ff0f
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/protocol/ObjectSerializationCache.java
@@ -0,0 +1,57 @@
+/*
+ * 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.protocol;
+
+import java.util.IdentityHashMap;
+
+/**
+ * The ObjectSerializationCache stores sizes and values computed during the
+ * first serialization pass.  This avoids recalculating and recomputing the same
+ * values during the second pass.
+ *
+ * It is intended to be used as part of a two-pass serialization process like:
+ * ObjectSerializationCache cache = new ObjectSerializationCache();
+ * message.size(version, cache);
+ * message.write(version, cache);
+ */
+public final class ObjectSerializationCache {
+    private final IdentityHashMap map;
+
+    public ObjectSerializationCache() {
+        this.map = new IdentityHashMap<>();
+    }
+
+    public void setArraySizeInBytes(Object o, int size) {
+        map.put(o, Integer.valueOf(size));
+    }
+
+    public int getArraySizeInBytes(Object o) {
+        Object value = map.get(o);
+        Integer sizeInBytes = (Integer) value;
+        return sizeInBytes;
+    }
+
+    public void cacheSerializedValue(Object o, byte[] val) {
+        map.put(o, val);
+    }
+
+    public byte[] getSerializedValue(Object o) {
+        Object value = map.get(o);
+        return (byte[]) value;
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/protocol/Protocol.java b/clients/src/main/java/org/apache/kafka/common/protocol/Protocol.java
index b5042c3f92a8d..706ed31ce043b 100644
--- a/clients/src/main/java/org/apache/kafka/common/protocol/Protocol.java
+++ b/clients/src/main/java/org/apache/kafka/common/protocol/Protocol.java
@@ -16,12 +16,12 @@
  */
 package org.apache.kafka.common.protocol;
 
-import org.apache.kafka.common.protocol.types.ArrayOf;
+import org.apache.kafka.common.message.RequestHeaderData;
+import org.apache.kafka.common.message.ResponseHeaderData;
 import org.apache.kafka.common.protocol.types.BoundField;
 import org.apache.kafka.common.protocol.types.Schema;
+import org.apache.kafka.common.protocol.types.TaggedFields;
 import org.apache.kafka.common.protocol.types.Type;
-import org.apache.kafka.common.requests.RequestHeader;
-import org.apache.kafka.common.requests.ResponseHeader;
 
 import java.util.LinkedHashMap;
 import java.util.LinkedHashSet;
@@ -43,18 +43,21 @@ private static void schemaToBnfHtml(Schema schema, StringBuilder b, int indentSi
 
         // Top level fields
         for (BoundField field: schema.fields()) {
-            if (field.def.type instanceof ArrayOf) {
+            Type type = field.def.type;
+            if (type.isArray()) {
                 b.append("[");
                 b.append(field.def.name);
                 b.append("] ");
-                Type innerType = ((ArrayOf) field.def.type).type();
-                if (!subTypes.containsKey(field.def.name))
-                    subTypes.put(field.def.name, innerType);
+                if (!subTypes.containsKey(field.def.name)) {
+                    subTypes.put(field.def.name, type.arrayElementType().get());
+                }
+            } else if (type instanceof TaggedFields) {
+                b.append("TAG_BUFFER ");
             } else {
                 b.append(field.def.name);
                 b.append(" ");
                 if (!subTypes.containsKey(field.def.name))
-                    subTypes.put(field.def.name, field.def.type);
+                    subTypes.put(field.def.name, type);
             }
         }
         b.append("\n");
@@ -81,8 +84,8 @@ private static void schemaToBnfHtml(Schema schema, StringBuilder b, int indentSi
     private static void populateSchemaFields(Schema schema, Set fields) {
         for (BoundField field: schema.fields()) {
             fields.add(field);
-            if (field.def.type instanceof ArrayOf) {
-                Type innerType = ((ArrayOf) field.def.type).type();
+            if (field.def.type.isArray()) {
+                Type innerType = field.def.type.arrayElementType().get();
                 if (innerType instanceof Schema)
                     populateSchemaFields((Schema) innerType, fields);
             } else if (field.def.type instanceof Schema)
@@ -116,18 +119,20 @@ public static String toHtml() {
         final StringBuilder b = new StringBuilder();
         b.append("
      Headers:
      \n");
 
-        b.append("
      ");
-        b.append("Request Header => ");
-        schemaToBnfHtml(RequestHeader.SCHEMA, b, 2);
-        b.append("
      \n");
-        schemaToFieldTableHtml(RequestHeader.SCHEMA, b);
-
-        b.append("
      ");
-        b.append("Response Header => ");
-        schemaToBnfHtml(ResponseHeader.SCHEMA, b, 2);
-        b.append("
      \n");
-        schemaToFieldTableHtml(ResponseHeader.SCHEMA, b);
-
+        for (int i = 0; i < RequestHeaderData.SCHEMAS.length; i++) {
+            b.append("
      ");
+            b.append("Request Header v").append(i).append(" => ");
+            schemaToBnfHtml(RequestHeaderData.SCHEMAS[i], b, 2);
+            b.append("
      \n");
+            schemaToFieldTableHtml(RequestHeaderData.SCHEMAS[i], b);
+        }
+        for (int i = 0; i < ResponseHeaderData.SCHEMAS.length; i++) {
+            b.append("
      ");
+            b.append("Response Header v").append(i).append(" => ");
+            schemaToBnfHtml(ResponseHeaderData.SCHEMAS[i], b, 2);
+            b.append("
      \n");
+            schemaToFieldTableHtml(ResponseHeaderData.SCHEMAS[i], b);
+        }
         for (ApiKeys key : ApiKeys.values()) {
             // Key
             b.append("
      ");
diff --git a/clients/src/main/java/org/apache/kafka/common/protocol/Readable.java b/clients/src/main/java/org/apache/kafka/common/protocol/Readable.java
index a527239ca486c..b3e2defd2186c 100644
--- a/clients/src/main/java/org/apache/kafka/common/protocol/Readable.java
+++ b/clients/src/main/java/org/apache/kafka/common/protocol/Readable.java
@@ -17,7 +17,12 @@
 
 package org.apache.kafka.common.protocol;
 
+import org.apache.kafka.common.protocol.types.RawTaggedField;
+
 import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.UUID;
 
 public interface Readable {
     byte readByte();
@@ -25,33 +30,28 @@ public interface Readable {
     int readInt();
     long readLong();
     void readArray(byte[] arr);
+    int readUnsignedVarint();
 
-    /**
-     * Read a Kafka-delimited string from a byte buffer.  The UTF-8 string
-     * length is stored in a two-byte short.  If the length is negative, the
-     * string is null.
-     */
-    default String readNullableString() {
-        int length = readShort();
-        if (length < 0) {
-            return null;
-        }
+    default String readString(int length) {
         byte[] arr = new byte[length];
         readArray(arr);
         return new String(arr, StandardCharsets.UTF_8);
     }
 
+    default List readUnknownTaggedField(List unknowns, int tag, int size) {
+        if (unknowns == null) {
+            unknowns = new ArrayList<>();
+        }
+        byte[] data = new byte[size];
+        readArray(data);
+        unknowns.add(new RawTaggedField(tag, data));
+        return unknowns;
+    }
+
     /**
-     * Read a Kafka-delimited array from a byte buffer.  The array length is
-     * stored in a four-byte short.
+     * Read a UUID with the most significant digits first.
      */
-    default byte[] readNullableBytes() {
-        int length = readInt();
-        if (length < 0) {
-            return null;
-        }
-        byte[] arr = new byte[length];
-        readArray(arr);
-        return arr;
+    default UUID readUUID() {
+        return new UUID(readLong(), readLong());
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/protocol/SecurityProtocol.java b/clients/src/main/java/org/apache/kafka/common/protocol/SecurityProtocol.java
new file mode 100644
index 0000000000000..9739ad50078b8
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/protocol/SecurityProtocol.java
@@ -0,0 +1,78 @@
+/*
+ * 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.protocol;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Locale;
+import java.util.Map;
+
+// A duplicate of org.apache.kafka.common.security.auth.SecurityProtocol. Temporarily added to break circular dependency
+// among kafka-server, container, and linkedin-kafka-clients.
+public enum SecurityProtocol {
+    /** Un-authenticated, non-encrypted channel */
+    PLAINTEXT(0, "PLAINTEXT"),
+    /** SSL channel */
+    SSL(1, "SSL"),
+    /** SASL authenticated, non-encrypted channel */
+    SASL_PLAINTEXT(2, "SASL_PLAINTEXT"),
+    /** SASL authenticated, SSL channel */
+    SASL_SSL(3, "SASL_SSL");
+
+    private static final Map CODE_TO_SECURITY_PROTOCOL;
+    private static final List NAMES;
+
+    static {
+        SecurityProtocol[] protocols = SecurityProtocol.values();
+        List names = new ArrayList<>(protocols.length);
+        Map codeToSecurityProtocol = new HashMap<>(protocols.length);
+        for (SecurityProtocol proto : protocols) {
+            codeToSecurityProtocol.put(proto.id, proto);
+            names.add(proto.name);
+        }
+        CODE_TO_SECURITY_PROTOCOL = Collections.unmodifiableMap(codeToSecurityProtocol);
+        NAMES = Collections.unmodifiableList(names);
+    }
+
+    /** The permanent and immutable id of a security protocol -- this can't change, and must match kafka.cluster.SecurityProtocol  */
+    public final short id;
+
+    /** Name of the security protocol. This may be used by client configuration. */
+    public final String name;
+
+    SecurityProtocol(int id, String name) {
+        this.id = (short) id;
+        this.name = name;
+    }
+
+    public static List names() {
+        return NAMES;
+    }
+
+    public static SecurityProtocol forId(short id) {
+        return CODE_TO_SECURITY_PROTOCOL.get(id);
+    }
+
+    /** Case insensitive lookup by protocol name */
+    public static SecurityProtocol forName(String name) {
+        return SecurityProtocol.valueOf(name.toUpperCase(Locale.ROOT));
+    }
+
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/protocol/Writable.java b/clients/src/main/java/org/apache/kafka/common/protocol/Writable.java
index 9478ed3ec6225..34f6db34bd68c 100644
--- a/clients/src/main/java/org/apache/kafka/common/protocol/Writable.java
+++ b/clients/src/main/java/org/apache/kafka/common/protocol/Writable.java
@@ -17,55 +17,18 @@
 
 package org.apache.kafka.common.protocol;
 
-import java.nio.charset.StandardCharsets;
+import java.util.UUID;
 
 public interface Writable {
     void writeByte(byte val);
     void writeShort(short val);
     void writeInt(int val);
     void writeLong(long val);
-    void writeArray(byte[] arr);
+    void writeByteArray(byte[] arr);
+    void writeUnsignedVarint(int i);
 
-    /**
-     * Write a nullable byte array delimited by a four-byte length prefix.
-     */
-    default void writeNullableBytes(byte[] arr) {
-        if (arr == null) {
-            writeInt(-1);
-        } else {
-            writeBytes(arr);
-        }
-    }
-
-    /**
-     * Write a byte array delimited by a four-byte length prefix.
-     */
-    default void writeBytes(byte[] arr) {
-        writeInt(arr.length);
-        writeArray(arr);
-    }
-
-    /**
-     * Write a nullable string delimited by a two-byte length prefix.
-     */
-    default void writeNullableString(String string) {
-        if (string == null) {
-            writeShort((short) -1);
-        } else {
-            writeString(string);
-        }
-    }
-
-    /**
-     * Write a string delimited by a two-byte length prefix.
-     */
-    default void writeString(String string) {
-        byte[] arr = string.getBytes(StandardCharsets.UTF_8);
-        if (arr.length > Short.MAX_VALUE) {
-            throw new RuntimeException("Can't store string longer than " +
-                Short.MAX_VALUE);
-        }
-        writeShort((short) arr.length);
-        writeArray(arr);
+    default void writeUUID(UUID uuid) {
+        writeLong(uuid.getMostSignificantBits());
+        writeLong(uuid.getLeastSignificantBits());
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/protocol/types/ArrayOf.java b/clients/src/main/java/org/apache/kafka/common/protocol/types/ArrayOf.java
index 6609dfd519619..3333084ef663a 100644
--- a/clients/src/main/java/org/apache/kafka/common/protocol/types/ArrayOf.java
+++ b/clients/src/main/java/org/apache/kafka/common/protocol/types/ArrayOf.java
@@ -19,6 +19,7 @@
 import org.apache.kafka.common.protocol.types.Type.DocumentedType;
 
 import java.nio.ByteBuffer;
+import java.util.Optional;
 
 /**
  * Represents a type for an array of a particular type
@@ -91,8 +92,9 @@ public int sizeOf(Object o) {
         return size;
     }
 
-    public Type type() {
-        return type;
+    @Override
+    public Optional arrayElementType() {
+        return Optional.of(type);
     }
 
     @Override
diff --git a/clients/src/main/java/org/apache/kafka/common/protocol/types/CompactArrayOf.java b/clients/src/main/java/org/apache/kafka/common/protocol/types/CompactArrayOf.java
new file mode 100644
index 0000000000000..4e9f8f8a98153
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/protocol/types/CompactArrayOf.java
@@ -0,0 +1,139 @@
+/*
+ * 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.protocol.types;
+
+import org.apache.kafka.common.protocol.types.Type.DocumentedType;
+import org.apache.kafka.common.utils.ByteUtils;
+
+import java.nio.ByteBuffer;
+import java.util.Optional;
+
+/**
+ * Represents a type for a compact array of a particular type.
+ * A compact array represents its length with a varint rather than a
+ * fixed-length field.
+ */
+public class CompactArrayOf extends DocumentedType {
+    private static final String COMPACT_ARRAY_TYPE_NAME = "COMPACT_ARRAY";
+
+    private final Type type;
+    private final boolean nullable;
+
+
+    public CompactArrayOf(Type type) {
+        this(type, false);
+    }
+
+    public static CompactArrayOf nullable(Type type) {
+        return new CompactArrayOf(type, true);
+    }
+
+    private CompactArrayOf(Type type, boolean nullable) {
+        this.type = type;
+        this.nullable = nullable;
+    }
+
+    @Override
+    public boolean isNullable() {
+        return nullable;
+    }
+
+    @Override
+    public void write(ByteBuffer buffer, Object o) {
+        if (o == null) {
+            ByteUtils.writeUnsignedVarint(0, buffer);
+            return;
+        }
+        Object[] objs = (Object[]) o;
+        int size = objs.length;
+        ByteUtils.writeUnsignedVarint(size + 1, buffer);
+
+        for (Object obj : objs)
+            type.write(buffer, obj);
+    }
+
+    @Override
+    public Object read(ByteBuffer buffer) {
+        int n = ByteUtils.readUnsignedVarint(buffer);
+        if (n == 0) {
+            if (isNullable()) {
+                return null;
+            } else {
+                throw new SchemaException("This array is not nullable.");
+            }
+        }
+        int size = n - 1;
+        if (size > buffer.remaining())
+            throw new SchemaException("Error reading array of size " + size + ", only " + buffer.remaining() + " bytes available");
+        Object[] objs = new Object[size];
+        for (int i = 0; i < size; i++)
+            objs[i] = type.read(buffer);
+        return objs;
+    }
+
+    @Override
+    public int sizeOf(Object o) {
+        if (o == null) {
+            return 1;
+        }
+        Object[] objs = (Object[]) o;
+        int size = ByteUtils.sizeOfUnsignedVarint(objs.length + 1);
+        for (Object obj : objs) {
+            size += type.sizeOf(obj);
+        }
+        return size;
+    }
+
+    @Override
+    public Optional arrayElementType() {
+        return Optional.of(type);
+    }
+
+    @Override
+    public String toString() {
+        return COMPACT_ARRAY_TYPE_NAME + "(" + type + ")";
+    }
+
+    @Override
+    public Object[] validate(Object item) {
+        try {
+            if (isNullable() && item == null)
+                return null;
+
+            Object[] array = (Object[]) item;
+            for (Object obj : array)
+                type.validate(obj);
+            return array;
+        } catch (ClassCastException e) {
+            throw new SchemaException("Not an Object[].");
+        }
+    }
+
+    @Override
+    public String typeName() {
+        return COMPACT_ARRAY_TYPE_NAME;
+    }
+
+    @Override
+    public String documentation() {
+        return "Represents a sequence of objects of a given type T. " +
+                "Type T can be either a primitive type (e.g. " + STRING + ") or a structure. " +
+                "First, the length N + 1 is given as an UNSIGNED_VARINT. Then N instances of type T follow. " +
+                "A null array is represented with a length of 0. " +
+                "In protocol documentation an array of T instances is referred to as [T].";
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/protocol/types/Field.java b/clients/src/main/java/org/apache/kafka/common/protocol/types/Field.java
index 72e051c179e73..749b7f3dc961f 100644
--- a/clients/src/main/java/org/apache/kafka/common/protocol/types/Field.java
+++ b/clients/src/main/java/org/apache/kafka/common/protocol/types/Field.java
@@ -75,6 +75,16 @@ public Int64(String name, String docString, long defaultValue) {
         }
     }
 
+    public static class UUID extends Field {
+        public UUID(String name, String docString) {
+            super(name, Type.UUID, docString, false, null);
+        }
+
+        public UUID(String name, String docString, UUID defaultValue) {
+            super(name, Type.UUID, docString, true, defaultValue);
+        }
+    }
+
     public static class Int16 extends Field {
         public Int16(String name, String docString) {
             super(name, Type.INT16, docString, false, null);
@@ -87,12 +97,24 @@ public Str(String name, String docString) {
         }
     }
 
+    public static class CompactStr extends Field {
+        public CompactStr(String name, String docString) {
+            super(name, Type.COMPACT_STRING, docString, false, null);
+        }
+    }
+
     public static class NullableStr extends Field {
         public NullableStr(String name, String docString) {
             super(name, Type.NULLABLE_STRING, docString, false, null);
         }
     }
 
+    public static class CompactNullableStr extends Field {
+        public CompactNullableStr(String name, String docString) {
+            super(name, Type.COMPACT_NULLABLE_STRING, docString, false, null);
+        }
+    }
+
     public static class Bool extends Field {
         public Bool(String name, String docString) {
             super(name, Type.BOOLEAN, docString, false, null);
@@ -105,6 +127,32 @@ public Array(String name, Type elementType, String docString) {
         }
     }
 
+    public static class CompactArray extends Field {
+        public CompactArray(String name, Type elementType, String docString) {
+            super(name, new CompactArrayOf(elementType), docString, false, null);
+        }
+    }
+
+    public static class TaggedFieldsSection extends Field {
+        private static final String NAME = "_tagged_fields";
+        private static final String DOC_STRING = "The tagged fields";
+
+        /**
+         * Create a new TaggedFieldsSection with the given tags and fields.
+         *
+         * @param fields    This is an array containing Integer tags followed
+         *                  by associated Field objects.
+         * @return          The new {@link TaggedFieldsSection}
+         */
+        public static TaggedFieldsSection of(Object... fields) {
+            return new TaggedFieldsSection(TaggedFields.of(fields));
+        }
+
+        public TaggedFieldsSection(Type type) {
+            super(NAME, type, DOC_STRING, false, null);
+        }
+    }
+
     public static class ComplexArray {
         public final String name;
         public final String docString;
diff --git a/clients/src/main/java/org/apache/kafka/common/protocol/types/RawTaggedField.java b/clients/src/main/java/org/apache/kafka/common/protocol/types/RawTaggedField.java
new file mode 100644
index 0000000000000..60deb5ed5fc26
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/protocol/types/RawTaggedField.java
@@ -0,0 +1,56 @@
+/*
+ * 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.protocol.types;
+
+import java.util.Arrays;
+
+public class RawTaggedField {
+    private final int tag;
+    private final byte[] data;
+
+    public RawTaggedField(int tag, byte[] data) {
+        this.tag = tag;
+        this.data = data;
+    }
+
+    public int tag() {
+        return tag;
+    }
+
+    public byte[] data() {
+        return data;
+    }
+
+    public int size() {
+        return data.length;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if ((o == null) || (!o.getClass().equals(getClass()))) {
+            return false;
+        }
+        RawTaggedField other = (RawTaggedField) o;
+        return tag == other.tag && Arrays.equals(data, other.data);
+    }
+
+    @Override
+    public int hashCode() {
+        return tag ^ Arrays.hashCode(data);
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/protocol/types/RawTaggedFieldWriter.java b/clients/src/main/java/org/apache/kafka/common/protocol/types/RawTaggedFieldWriter.java
new file mode 100644
index 0000000000000..7218d34032788
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/protocol/types/RawTaggedFieldWriter.java
@@ -0,0 +1,80 @@
+/*
+ * 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.protocol.types;
+
+import org.apache.kafka.common.protocol.Writable;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.ListIterator;
+
+/**
+ * The RawTaggedFieldWriter is used by Message subclasses to serialize their
+ * lists of raw tags.
+ */
+public class RawTaggedFieldWriter {
+    private static final RawTaggedFieldWriter EMPTY_WRITER =
+        new RawTaggedFieldWriter(new ArrayList<>(0));
+
+    private final List fields;
+    private final ListIterator iter;
+    private int prevTag;
+
+    public static RawTaggedFieldWriter forFields(List fields) {
+        if (fields == null) {
+            return EMPTY_WRITER;
+        }
+        return new RawTaggedFieldWriter(fields);
+    }
+
+    private RawTaggedFieldWriter(List fields) {
+        this.fields = fields;
+        this.iter = this.fields.listIterator();
+        this.prevTag = -1;
+    }
+
+    public int numFields() {
+        return fields.size();
+    }
+
+    public void writeRawTags(Writable writable, int nextDefinedTag) {
+        while (iter.hasNext()) {
+            RawTaggedField field = iter.next();
+            int tag = field.tag();
+            if (tag >= nextDefinedTag) {
+                if (tag == nextDefinedTag) {
+                    // We must not have a raw tag field that duplicates the tag of another field.
+                    throw new RuntimeException("Attempted to use tag " + tag + " as an " +
+                        "undefined tag.");
+                }
+                iter.previous();
+                return;
+            }
+            if (tag <= prevTag) {
+                // The raw tag field list must be sorted by tag, and there must not be
+                // any duplicate tags.
+                throw new RuntimeException("Invalid raw tag field list: tag " + tag +
+                    " comes after tag " + prevTag + ", but is not higher than it.");
+            }
+            writable.writeUnsignedVarint(field.tag());
+            writable.writeUnsignedVarint(field.data().length);
+            writable.writeByteArray(field.data());
+            prevTag = tag;
+        }
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/protocol/types/Schema.java b/clients/src/main/java/org/apache/kafka/common/protocol/types/Schema.java
index cbcd4491d37fb..8aa54fc3689c8 100644
--- a/clients/src/main/java/org/apache/kafka/common/protocol/types/Schema.java
+++ b/clients/src/main/java/org/apache/kafka/common/protocol/types/Schema.java
@@ -25,18 +25,38 @@
  * The schema for a compound record definition
  */
 public class Schema extends Type {
+    private final static Object[] NO_VALUES = new Object[0];
 
     private final BoundField[] fields;
     private final Map fieldsByName;
+    private final boolean tolerateMissingFieldsWithDefaults;
+    private final Struct cachedStruct;
 
     /**
      * Construct the schema with a given list of its field values
      *
+     * @param fs the fields of this schema
+     *
      * @throws SchemaException If the given list have duplicate fields
      */
     public Schema(Field... fs) {
+        this(false, fs);
+    }
+
+    /**
+     * Construct the schema with a given list of its field values and the ability to tolerate
+     * missing optional fields with defaults at the end of the schema definition.
+     *
+     * @param tolerateMissingFieldsWithDefaults whether to accept records with missing optional
+     * fields the end of the schema
+     * @param fs the fields of this schema
+     *
+     * @throws SchemaException If the given list have duplicate fields
+     */
+    public Schema(boolean tolerateMissingFieldsWithDefaults, Field... fs) {
         this.fields = new BoundField[fs.length];
-        this.fieldsByName = new HashMap<>();
+        this.fieldsByName = new HashMap<>(fs.length);
+        this.tolerateMissingFieldsWithDefaults = tolerateMissingFieldsWithDefaults;
         for (int i = 0; i < this.fields.length; i++) {
             Field def = fs[i];
             if (fieldsByName.containsKey(def.name))
@@ -44,6 +64,9 @@ public Schema(Field... fs) {
             this.fields[i] = new BoundField(def, this, i);
             this.fieldsByName.put(def.name, this.fields[i]);
         }
+        //6 schemas have no fields at the time of this writing (3 versions each of list_groups and api_versions)
+        //for such schemas there's no point in even creating a unique Struct object when deserializing.
+        this.cachedStruct = this.fields.length > 0 ? null : new Struct(this, NO_VALUES);
     }
 
     /**
@@ -64,14 +87,32 @@ public void write(ByteBuffer buffer, Object o) {
     }
 
     /**
-     * Read a struct from the buffer
+     * Read a struct from the buffer. If this schema is configured to tolerate missing
+     * optional fields at the end of the buffer, these fields are replaced with their default
+     * values; otherwise, if the schema does not tolerate missing fields, or if missing fields
+     * don't have a default value, a {@code SchemaException} is thrown to signify that mandatory
+     * fields are missing.
      */
     @Override
     public Struct read(ByteBuffer buffer) {
+        if (cachedStruct != null) {
+            return cachedStruct;
+        }
         Object[] objects = new Object[fields.length];
         for (int i = 0; i < fields.length; i++) {
             try {
-                objects[i] = fields[i].def.type.read(buffer);
+                if (tolerateMissingFieldsWithDefaults) {
+                    if (buffer.hasRemaining()) {
+                        objects[i] = fields[i].def.type.read(buffer);
+                    } else if (fields[i].def.hasDefaultValue) {
+                        objects[i] = fields[i].def.defaultValue;
+                    } else {
+                        throw new SchemaException("Missing value for field '" + fields[i].def.name +
+                                "' which has no default value.");
+                    }
+                } else {
+                    objects[i] = fields[i].def.type.read(buffer);
+                }
             } catch (Exception e) {
                 throw new SchemaException("Error reading field '" + fields[i].def.name + "': " +
                                           (e.getMessage() == null ? e.getClass().getName() : e.getMessage()));
@@ -107,7 +148,7 @@ public int numFields() {
 
     /**
      * Get a field by its slot in the record array
-     * 
+     *
      * @param slot The slot at which this field sits
      * @return The field
      */
@@ -117,7 +158,7 @@ public BoundField get(int slot) {
 
     /**
      * Get a field by its name
-     * 
+     *
      * @param name The name of the field
      * @return The field
      */
@@ -176,10 +217,9 @@ private static void handleNode(Type node, Visitor visitor) {
             visitor.visit(schema);
             for (BoundField f : schema.fields())
                 handleNode(f.def.type, visitor);
-        } else if (node instanceof ArrayOf) {
-            ArrayOf array = (ArrayOf) node;
-            visitor.visit(array);
-            handleNode(array.type(), visitor);
+        } else if (node.isArray()) {
+            visitor.visit(node);
+            handleNode(node.arrayElementType().get(), visitor);
         } else {
             visitor.visit(node);
         }
@@ -190,7 +230,6 @@ private static void handleNode(Type node, Visitor visitor) {
      */
     public static abstract class Visitor {
         public void visit(Schema schema) {}
-        public void visit(ArrayOf array) {}
         public void visit(Type field) {}
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/protocol/types/Struct.java b/clients/src/main/java/org/apache/kafka/common/protocol/types/Struct.java
index b93c0cf7cbcda..0b69851f16441 100644
--- a/clients/src/main/java/org/apache/kafka/common/protocol/types/Struct.java
+++ b/clients/src/main/java/org/apache/kafka/common/protocol/types/Struct.java
@@ -20,6 +20,8 @@
 
 import java.nio.ByteBuffer;
 import java.util.Arrays;
+import java.util.Objects;
+import java.util.UUID;
 
 /**
  * A record that can be serialized and deserialized according to a pre-defined schema
@@ -87,6 +89,10 @@ public Long get(Field.Int64 field) {
         return getLong(field.name);
     }
 
+    public UUID get(Field.UUID field) {
+        return getUUID(field.name);
+    }
+
     public Short get(Field.Int16 field) {
         return getShort(field.name);
     }
@@ -117,6 +123,12 @@ public Long getOrElse(Field.Int64 field, long alternative) {
         return alternative;
     }
 
+    public UUID getOrElse(Field.UUID field, UUID alternative) {
+        if (hasField(field.name))
+            return getUUID(field.name);
+        return alternative;
+    }
+
     public Short getOrElse(Field.Int16 field, short alternative) {
         if (hasField(field.name))
             return getShort(field.name);
@@ -244,6 +256,14 @@ public Long getLong(String name) {
         return (Long) get(name);
     }
 
+    public UUID getUUID(BoundField field) {
+        return (UUID) get(field);
+    }
+
+    public UUID getUUID(String name) {
+        return (UUID) get(name);
+    }
+
     public Object[] getArray(BoundField field) {
         return (Object[]) get(field);
     }
@@ -289,6 +309,7 @@ public byte[] getByteArray(String name) {
         ByteBuffer buf = (ByteBuffer) result;
         byte[] arr = new byte[buf.remaining()];
         buf.get(arr);
+        buf.flip();
         return arr;
     }
 
@@ -340,6 +361,10 @@ public Struct set(Field.Int64 def, long value) {
         return set(def.name, value);
     }
 
+    public Struct set(Field.UUID def, UUID value) {
+        return set(def.name, value);
+    }
+
     public Struct set(Field.Int16 def, short value) {
         return set(def.name, value);
     }
@@ -393,9 +418,8 @@ public Struct instance(BoundField field) {
         validateField(field);
         if (field.def.type instanceof Schema) {
             return new Struct((Schema) field.def.type);
-        } else if (field.def.type instanceof ArrayOf) {
-            ArrayOf array = (ArrayOf) field.def.type;
-            return new Struct((Schema) array.type());
+        } else if (field.def.type.isArray()) {
+            return new Struct((Schema) field.def.type.arrayElementType().get());
         } else {
             throw new SchemaException("Field '" + field.def.name + "' is not a container type, it is of type " + field.def.type);
         }
@@ -447,6 +471,7 @@ public void writeTo(ByteBuffer buffer) {
      * @throws SchemaException If validation fails
      */
     private void validateField(BoundField field) {
+        Objects.requireNonNull(field, "`field` must be non-null");
         if (this.schema != field.schema)
             throw new SchemaException("Attempt to access field '" + field.def.name + "' from a different schema instance.");
         if (field.index > values.length)
@@ -470,7 +495,7 @@ public String toString() {
             BoundField f = this.schema.get(i);
             b.append(f.def.name);
             b.append('=');
-            if (f.def.type instanceof ArrayOf && this.values[i] != null) {
+            if (f.def.type.isArray() && this.values[i] != null) {
                 Object[] arrayValue = (Object[]) this.values[i];
                 b.append('[');
                 for (int j = 0; j < arrayValue.length; j++) {
@@ -494,7 +519,7 @@ public int hashCode() {
         int result = 1;
         for (int i = 0; i < this.values.length; i++) {
             BoundField f = this.schema.get(i);
-            if (f.def.type instanceof ArrayOf) {
+            if (f.def.type.isArray()) {
                 if (this.get(f) != null) {
                     Object[] arrayObject = (Object[]) this.get(f);
                     for (Object arrayItem: arrayObject)
@@ -524,12 +549,12 @@ public boolean equals(Object obj) {
         for (int i = 0; i < this.values.length; i++) {
             BoundField f = this.schema.get(i);
             boolean result;
-            if (f.def.type instanceof ArrayOf) {
+            if (f.def.type.isArray()) {
                 result = Arrays.equals((Object[]) this.get(f), (Object[]) other.get(f));
             } else {
                 Object thisField = this.get(f);
                 Object otherField = other.get(f);
-                return (thisField == null) ? (otherField == null) : thisField.equals(otherField);
+                result = Objects.equals(thisField, otherField);
             }
             if (!result)
                 return false;
diff --git a/clients/src/main/java/org/apache/kafka/common/protocol/types/TaggedFields.java b/clients/src/main/java/org/apache/kafka/common/protocol/types/TaggedFields.java
new file mode 100644
index 0000000000000..4e1ab0d4d5add
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/protocol/types/TaggedFields.java
@@ -0,0 +1,181 @@
+/*
+ * 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.protocol.types;
+
+import org.apache.kafka.common.protocol.types.Type.DocumentedType;
+import org.apache.kafka.common.utils.ByteUtils;
+
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.Map;
+import java.util.NavigableMap;
+import java.util.TreeMap;
+
+/**
+ * Represents a tagged fields section.
+ */
+public class TaggedFields extends DocumentedType {
+    private static final String TAGGED_FIELDS_TYPE_NAME = "TAGGED_FIELDS";
+
+    private final Map fields;
+
+    /**
+     * Create a new TaggedFields object with the given tags and fields.
+     *
+     * @param fields    This is an array containing Integer tags followed
+     *                  by associated Field objects.
+     * @return          The new {@link TaggedFields}
+     */
+    @SuppressWarnings("unchecked")
+    public static TaggedFields of(Object... fields) {
+        if (fields.length % 2 != 0) {
+            throw new RuntimeException("TaggedFields#of takes an even " +
+                "number of parameters.");
+        }
+        TreeMap newFields = new TreeMap<>();
+        for (int i = 0; i < fields.length; i += 2) {
+            Integer tag = (Integer) fields[i];
+            Field field = (Field) fields[i + 1];
+            newFields.put(tag, field);
+        }
+        return new TaggedFields(newFields);
+    }
+
+    public TaggedFields(Map fields) {
+        this.fields = fields;
+    }
+
+    @Override
+    public boolean isNullable() {
+        return false;
+    }
+
+    @SuppressWarnings("unchecked")
+    @Override
+    public void write(ByteBuffer buffer, Object o) {
+        NavigableMap objects = (NavigableMap) o;
+        ByteUtils.writeUnsignedVarint(objects.size(), buffer);
+        for (Map.Entry entry : objects.entrySet()) {
+            Integer tag = entry.getKey();
+            Field field = fields.get(tag);
+            ByteUtils.writeUnsignedVarint(tag, buffer);
+            if (field == null) {
+                RawTaggedField value = (RawTaggedField) entry.getValue();
+                ByteUtils.writeUnsignedVarint(value.data().length, buffer);
+                buffer.put(value.data());
+            } else {
+                ByteUtils.writeUnsignedVarint(field.type.sizeOf(entry.getValue()), buffer);
+                field.type.write(buffer, entry.getValue());
+            }
+        }
+    }
+
+    @SuppressWarnings("unchecked")
+    @Override
+    public NavigableMap read(ByteBuffer buffer) {
+        int numTaggedFields = ByteUtils.readUnsignedVarint(buffer);
+        if (numTaggedFields == 0) {
+            return Collections.emptyNavigableMap();
+        }
+        NavigableMap objects = new TreeMap<>();
+        int prevTag = -1;
+        for (int i = 0; i < numTaggedFields; i++) {
+            int tag = ByteUtils.readUnsignedVarint(buffer);
+            if (tag <= prevTag) {
+                throw new RuntimeException("Invalid or out-of-order tag " + tag);
+            }
+            prevTag = tag;
+            int size = ByteUtils.readUnsignedVarint(buffer);
+            Field field = fields.get(tag);
+            if (field == null) {
+                byte[] bytes = new byte[size];
+                buffer.get(bytes);
+                objects.put(tag, new RawTaggedField(tag, bytes));
+            } else {
+                objects.put(tag, field.type.read(buffer));
+            }
+        }
+        return objects;
+    }
+
+    @SuppressWarnings("unchecked")
+    @Override
+    public int sizeOf(Object o) {
+        int size = 0;
+        NavigableMap objects = (NavigableMap) o;
+        size += ByteUtils.sizeOfUnsignedVarint(objects.size());
+        for (Map.Entry entry : objects.entrySet()) {
+            Integer tag = entry.getKey();
+            size += ByteUtils.sizeOfUnsignedVarint(tag);
+            Field field = fields.get(tag);
+            if (field == null) {
+                RawTaggedField value = (RawTaggedField) entry.getValue();
+                size += value.data().length + ByteUtils.sizeOfUnsignedVarint(value.data().length);
+            } else {
+                int valueSize = field.type.sizeOf(entry.getValue());
+                size += valueSize + ByteUtils.sizeOfUnsignedVarint(valueSize);
+            }
+        }
+        return size;
+    }
+
+    @Override
+    public String toString() {
+        StringBuilder bld = new StringBuilder("TAGGED_FIELDS_TYPE_NAME(");
+        String prefix = "";
+        for (Map.Entry field : fields.entrySet()) {
+            bld.append(prefix);
+            prefix = ", ";
+            bld.append(field.getKey()).append(" -> ").append(field.getValue().toString());
+        }
+        bld.append(")");
+        return bld.toString();
+    }
+
+    @SuppressWarnings("unchecked")
+    @Override
+    public Map validate(Object item) {
+        try {
+            NavigableMap objects = (NavigableMap) item;
+            for (Map.Entry entry : objects.entrySet()) {
+                Integer tag = entry.getKey();
+                Field field = fields.get(tag);
+                if (field == null) {
+                    if (!(entry.getValue() instanceof RawTaggedField)) {
+                        throw new SchemaException("The value associated with tag " + tag +
+                            " must be a RawTaggedField in this version of the software.");
+                    }
+                } else {
+                    field.type.validate(entry.getValue());
+                }
+            }
+            return objects;
+        } catch (ClassCastException e) {
+            throw new SchemaException("Not a NavigableMap.");
+        }
+    }
+
+    @Override
+    public String typeName() {
+        return TAGGED_FIELDS_TYPE_NAME;
+    }
+
+    @Override
+    public String documentation() {
+        return "Represents a series of tagged fields.";
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/protocol/types/Type.java b/clients/src/main/java/org/apache/kafka/common/protocol/types/Type.java
index 4bd508b77989f..6183625a785f8 100644
--- a/clients/src/main/java/org/apache/kafka/common/protocol/types/Type.java
+++ b/clients/src/main/java/org/apache/kafka/common/protocol/types/Type.java
@@ -23,6 +23,8 @@
 import org.apache.kafka.common.utils.Utils;
 
 import java.nio.ByteBuffer;
+import java.util.Optional;
+import java.util.UUID;
 
 /**
  * A serializable type
@@ -63,6 +65,20 @@ public boolean isNullable() {
         return false;
     }
 
+    /**
+     * If the type is an array, return the type of the array elements.  Otherwise, return empty.
+     */
+    public Optional arrayElementType() {
+        return Optional.empty();
+    }
+
+    /**
+     * Returns true if the type is an array.
+     */
+    public final boolean isArray() {
+        return arrayElementType().isPresent();
+    }
+
     /**
      * A Type that can return its description for documentation purposes.
      */
@@ -313,6 +329,44 @@ public String documentation() {
         }
     };
 
+    public static final DocumentedType UUID = new DocumentedType() {
+        @Override
+        public void write(ByteBuffer buffer, Object o) {
+            final java.util.UUID uuid = (java.util.UUID) o;
+            buffer.putLong(uuid.getMostSignificantBits());
+            buffer.putLong(uuid.getLeastSignificantBits());
+        }
+
+        @Override
+        public Object read(ByteBuffer buffer) {
+            return new java.util.UUID(buffer.getLong(), buffer.getLong());
+        }
+
+        @Override
+        public int sizeOf(Object o) {
+            return 16;
+        }
+
+        @Override
+        public String typeName() {
+            return "UUID";
+        }
+
+        @Override
+        public UUID validate(Object item) {
+            if (item instanceof UUID)
+                return (UUID) item;
+            else
+                throw new SchemaException(item + " is not a UUID.");
+        }
+
+        @Override
+        public String documentation() {
+            return "Represents a java.util.UUID. " +
+                    "The values are encoded using sixteen bytes in network byte order (big-endian).";
+        }
+    };
+
     public static final DocumentedType STRING = new DocumentedType() {
         @Override
         public void write(ByteBuffer buffer, Object o) {
@@ -361,6 +415,56 @@ public String documentation() {
         }
     };
 
+    public static final DocumentedType COMPACT_STRING = new DocumentedType() {
+        @Override
+        public void write(ByteBuffer buffer, Object o) {
+            byte[] bytes = Utils.utf8((String) o);
+            if (bytes.length > Short.MAX_VALUE)
+                throw new SchemaException("String length " + bytes.length + " is larger than the maximum string length.");
+            ByteUtils.writeUnsignedVarint(bytes.length + 1, buffer);
+            buffer.put(bytes);
+        }
+
+        @Override
+        public String read(ByteBuffer buffer) {
+            int length = ByteUtils.readUnsignedVarint(buffer) - 1;
+            if (length < 0)
+                throw new SchemaException("String length " + length + " cannot be negative");
+            if (length > Short.MAX_VALUE)
+                throw new SchemaException("String length " + length + " is larger than the maximum string length.");
+            if (length > buffer.remaining())
+                throw new SchemaException("Error reading string of length " + length + ", only " + buffer.remaining() + " bytes available");
+            String result = Utils.utf8(buffer, length);
+            buffer.position(buffer.position() + length);
+            return result;
+        }
+
+        @Override
+        public int sizeOf(Object o) {
+            int length = Utils.utf8Length((String) o);
+            return ByteUtils.sizeOfUnsignedVarint(length + 1) + length;
+        }
+
+        @Override
+        public String typeName() {
+            return "COMPACT_STRING";
+        }
+
+        @Override
+        public String validate(Object item) {
+            if (item instanceof String)
+                return (String) item;
+            else
+                throw new SchemaException(item + " is not a String.");
+        }
+
+        @Override
+        public String documentation() {
+            return "Represents a sequence of characters. First the length N + 1 is given as an UNSIGNED_VARINT " +
+                    ". Then N bytes follow which are the UTF-8 encoding of the character sequence.";
+        }
+    };
+
     public static final DocumentedType NULLABLE_STRING = new DocumentedType() {
         @Override
         public boolean isNullable() {
@@ -425,6 +529,74 @@ public String documentation() {
         }
     };
 
+    public static final DocumentedType COMPACT_NULLABLE_STRING = new DocumentedType() {
+        @Override
+        public boolean isNullable() {
+            return true;
+        }
+
+        @Override
+        public void write(ByteBuffer buffer, Object o) {
+            if (o == null) {
+                ByteUtils.writeUnsignedVarint(0, buffer);
+            } else {
+                byte[] bytes = Utils.utf8((String) o);
+                if (bytes.length > Short.MAX_VALUE)
+                    throw new SchemaException("String length " + bytes.length + " is larger than the maximum string length.");
+                ByteUtils.writeUnsignedVarint(bytes.length + 1, buffer);
+                buffer.put(bytes);
+            }
+        }
+
+        @Override
+        public String read(ByteBuffer buffer) {
+            int length = ByteUtils.readUnsignedVarint(buffer) - 1;
+            if (length < 0) {
+                return null;
+            } else if (length > Short.MAX_VALUE) {
+                throw new SchemaException("String length " + length + " is larger than the maximum string length.");
+            } else if (length > buffer.remaining()) {
+                throw new SchemaException("Error reading string of length " + length + ", only " + buffer.remaining() + " bytes available");
+            } else {
+                String result = Utils.utf8(buffer, length);
+                buffer.position(buffer.position() + length);
+                return result;
+            }
+        }
+
+        @Override
+        public int sizeOf(Object o) {
+            if (o == null) {
+                return 1;
+            }
+            int length = Utils.utf8Length((String) o);
+            return ByteUtils.sizeOfUnsignedVarint(length + 1) + length;
+        }
+
+        @Override
+        public String typeName() {
+            return "COMPACT_NULLABLE_STRING";
+        }
+
+        @Override
+        public String validate(Object item) {
+            if (item == null) {
+                return null;
+            } else if (item instanceof String) {
+                return (String) item;
+            } else {
+                throw new SchemaException(item + " is not a String.");
+            }
+        }
+
+        @Override
+        public String documentation() {
+            return "Represents a sequence of characters. First the length N + 1 is given as an UNSIGNED_VARINT " +
+                    ". Then N bytes follow which are the UTF-8 encoding of the character sequence. " +
+                    "A null string is represented with a length of 0.";
+        }
+    };
+
     public static final DocumentedType BYTES = new DocumentedType() {
         @Override
         public void write(ByteBuffer buffer, Object o) {
@@ -475,6 +647,57 @@ public String documentation() {
         }
     };
 
+    public static final DocumentedType COMPACT_BYTES = new DocumentedType() {
+        @Override
+        public void write(ByteBuffer buffer, Object o) {
+            ByteBuffer arg = (ByteBuffer) o;
+            int pos = arg.position();
+            ByteUtils.writeUnsignedVarint(arg.remaining() + 1, buffer);
+            buffer.put(arg);
+            arg.position(pos);
+        }
+
+        @Override
+        public Object read(ByteBuffer buffer) {
+            int size = ByteUtils.readUnsignedVarint(buffer) - 1;
+            if (size < 0)
+                throw new SchemaException("Bytes size " + size + " cannot be negative");
+            if (size > buffer.remaining())
+                throw new SchemaException("Error reading bytes of size " + size + ", only " + buffer.remaining() + " bytes available");
+
+            ByteBuffer val = buffer.slice();
+            val.limit(size);
+            buffer.position(buffer.position() + size);
+            return val;
+        }
+
+        @Override
+        public int sizeOf(Object o) {
+            ByteBuffer buffer = (ByteBuffer) o;
+            int remaining = buffer.remaining();
+            return ByteUtils.sizeOfUnsignedVarint(remaining + 1) + remaining;
+        }
+
+        @Override
+        public String typeName() {
+            return "COMPACT_BYTES";
+        }
+
+        @Override
+        public ByteBuffer validate(Object item) {
+            if (item instanceof ByteBuffer)
+                return (ByteBuffer) item;
+            else
+                throw new SchemaException(item + " is not a java.nio.ByteBuffer.");
+        }
+
+        @Override
+        public String documentation() {
+            return "Represents a raw sequence of bytes. First the length N+1 is given as an UNSIGNED_VARINT." +
+                    "Then N bytes follow.";
+        }
+    };
+
     public static final DocumentedType NULLABLE_BYTES = new DocumentedType() {
         @Override
         public boolean isNullable() {
@@ -541,6 +764,72 @@ public String documentation() {
         }
     };
 
+    public static final DocumentedType COMPACT_NULLABLE_BYTES = new DocumentedType() {
+        @Override
+        public boolean isNullable() {
+            return true;
+        }
+
+        @Override
+        public void write(ByteBuffer buffer, Object o) {
+            if (o == null) {
+                ByteUtils.writeUnsignedVarint(0, buffer);
+            } else {
+                ByteBuffer arg = (ByteBuffer) o;
+                int pos = arg.position();
+                ByteUtils.writeUnsignedVarint(arg.remaining() + 1, buffer);
+                buffer.put(arg);
+                arg.position(pos);
+            }
+        }
+
+        @Override
+        public Object read(ByteBuffer buffer) {
+            int size = ByteUtils.readUnsignedVarint(buffer) - 1;
+            if (size < 0)
+                return null;
+            if (size > buffer.remaining())
+                throw new SchemaException("Error reading bytes of size " + size + ", only " + buffer.remaining() + " bytes available");
+
+            ByteBuffer val = buffer.slice();
+            val.limit(size);
+            buffer.position(buffer.position() + size);
+            return val;
+        }
+
+        @Override
+        public int sizeOf(Object o) {
+            if (o == null) {
+                return 1;
+            }
+            ByteBuffer buffer = (ByteBuffer) o;
+            int remaining = buffer.remaining();
+            return ByteUtils.sizeOfUnsignedVarint(remaining + 1) + remaining;
+        }
+
+        @Override
+        public String typeName() {
+            return "COMPACT_NULLABLE_BYTES";
+        }
+
+        @Override
+        public ByteBuffer validate(Object item) {
+            if (item == null)
+                return null;
+
+            if (item instanceof ByteBuffer)
+                return (ByteBuffer) item;
+
+            throw new SchemaException(item + " is not a java.nio.ByteBuffer.");
+        }
+
+        @Override
+        public String documentation() {
+            return "Represents a raw sequence of bytes. First the length N+1 is given as an UNSIGNED_VARINT." +
+                    "Then N bytes follow. A null object is represented with a length of 0.";
+        }
+    };
+
     public static final DocumentedType RECORDS = new DocumentedType() {
         @Override
         public boolean isNullable() {
@@ -665,12 +954,12 @@ public String documentation() {
     };
 
     private static String toHtml() {
-
         DocumentedType[] types = {
             BOOLEAN, INT8, INT16, INT32, INT64,
-            UNSIGNED_INT32, VARINT, VARLONG,
-            STRING, NULLABLE_STRING, BYTES, NULLABLE_BYTES,
-            RECORDS, new ArrayOf(STRING)};
+            UNSIGNED_INT32, VARINT, VARLONG, UUID,
+            STRING, COMPACT_STRING, NULLABLE_STRING, COMPACT_NULLABLE_STRING,
+            BYTES, COMPACT_BYTES, NULLABLE_BYTES, COMPACT_NULLABLE_BYTES,
+            RECORDS, new ArrayOf(STRING), new CompactArrayOf(COMPACT_STRING)};
         final StringBuilder b = new StringBuilder();
         b.append("
      \n");
         b.append("");
diff --git a/clients/src/main/java/org/apache/kafka/common/record/AbstractLegacyRecordBatch.java b/clients/src/main/java/org/apache/kafka/common/record/AbstractLegacyRecordBatch.java
index 1d50a15a04938..93901da15c74c 100644
--- a/clients/src/main/java/org/apache/kafka/common/record/AbstractLegacyRecordBatch.java
+++ b/clients/src/main/java/org/apache/kafka/common/record/AbstractLegacyRecordBatch.java
@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.common.record;
 
+import org.apache.kafka.common.InvalidRecordException;
 import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.errors.CorruptRecordException;
 import org.apache.kafka.common.header.Header;
@@ -32,6 +33,7 @@
 import java.util.ArrayDeque;
 import java.util.Iterator;
 import java.util.NoSuchElementException;
+import java.util.Objects;
 
 import static org.apache.kafka.common.record.Records.LOG_OVERHEAD;
 import static org.apache.kafka.common.record.Records.OFFSET_OFFSET;
@@ -226,10 +228,15 @@ public Iterator iterator() {
         return iterator(BufferSupplier.NO_CACHING);
     }
 
-    private CloseableIterator iterator(BufferSupplier bufferSupplier) {
+    CloseableIterator iterator(BufferSupplier bufferSupplier) {
         if (isCompressed())
             return new DeepRecordsIterator(this, false, Integer.MAX_VALUE, bufferSupplier);
 
+        return shallowIterator();
+    }
+
+    @Override
+    public CloseableIterator shallowIterator() {
         return new CloseableIterator() {
             private boolean hasNext = true;
 
@@ -439,13 +446,13 @@ public boolean equals(Object o) {
             BasicLegacyRecordBatch that = (BasicLegacyRecordBatch) o;
 
             return offset == that.offset &&
-                    (record != null ? record.equals(that.record) : that.record == null);
+                Objects.equals(record, that.record);
         }
 
         @Override
         public int hashCode() {
             int result = record != null ? record.hashCode() : 0;
-            result = 31 * result + (int) (offset ^ (offset >>> 32));
+            result = 31 * result + Long.hashCode(offset);
             return result;
         }
     }
@@ -502,6 +509,16 @@ private void setTimestampAndUpdateCrc(TimestampType timestampType, long timestam
             ByteUtils.writeUnsignedInt(buffer, LOG_OVERHEAD + LegacyRecord.CRC_OFFSET, crc);
         }
 
+        /**
+         * LegacyRecordBatch does not implement this iterator and would hence fallback to the normal iterator.
+         *
+         * @return An iterator over the records contained within this batch
+         */
+        @Override
+        public CloseableIterator skipKeyValueIterator(BufferSupplier bufferSupplier) {
+            return CloseableIterator.wrap(iterator(bufferSupplier));
+        }
+
         @Override
         public void writeTo(ByteBufferOutputStream outputStream) {
             outputStream.write(buffer.duplicate());
@@ -516,7 +533,7 @@ public boolean equals(Object o) {
 
             ByteBufferLegacyRecordBatch that = (ByteBufferLegacyRecordBatch) o;
 
-            return buffer != null ? buffer.equals(that.buffer) : that.buffer == null;
+            return Objects.equals(buffer, that.buffer);
         }
 
         @Override
diff --git a/clients/src/main/java/org/apache/kafka/common/record/AbstractRecordBatch.java b/clients/src/main/java/org/apache/kafka/common/record/AbstractRecordBatch.java
index 78ad05046d257..d104fcde68dba 100644
--- a/clients/src/main/java/org/apache/kafka/common/record/AbstractRecordBatch.java
+++ b/clients/src/main/java/org/apache/kafka/common/record/AbstractRecordBatch.java
@@ -16,8 +16,8 @@
  */
 package org.apache.kafka.common.record;
 
-abstract class AbstractRecordBatch implements RecordBatch {
 
+abstract class AbstractRecordBatch implements RecordBatch {
     @Override
     public boolean hasProducerId() {
         return RecordBatch.NO_PRODUCER_ID < producerId();
diff --git a/clients/src/main/java/org/apache/kafka/common/record/AbstractRecords.java b/clients/src/main/java/org/apache/kafka/common/record/AbstractRecords.java
index 1994a71968a1c..09bd2e501f751 100644
--- a/clients/src/main/java/org/apache/kafka/common/record/AbstractRecords.java
+++ b/clients/src/main/java/org/apache/kafka/common/record/AbstractRecords.java
@@ -43,6 +43,15 @@ public boolean hasCompatibleMagic(byte magic) {
         return true;
     }
 
+    public RecordBatch firstBatch() {
+        Iterator iterator = batches().iterator();
+
+        if (!iterator.hasNext())
+            return null;
+
+        return iterator.next();
+    }
+
     /**
      * Get an iterator over the deep records.
      * @return An iterator over the records
diff --git a/clients/src/main/java/org/apache/kafka/common/record/CompressionType.java b/clients/src/main/java/org/apache/kafka/common/record/CompressionType.java
index 352d12d834977..454dd17bf8240 100644
--- a/clients/src/main/java/org/apache/kafka/common/record/CompressionType.java
+++ b/clients/src/main/java/org/apache/kafka/common/record/CompressionType.java
@@ -133,6 +133,18 @@ public InputStream wrapForInput(ByteBuffer buffer, byte messageVersion, BufferSu
                 throw new KafkaException(e);
             }
         }
+    },
+
+    PASSTHROUGH(5, "passthrough", 1.0f) {
+        @Override
+        public OutputStream wrapForOutput(ByteBufferOutputStream buffer, byte messageVersion) {
+            return buffer;
+        }
+
+        @Override
+        public InputStream wrapForInput(ByteBuffer buffer, byte messageVersion, BufferSupplier decompressionBufferSupplier) {
+            return new ByteBufferInputStream(buffer);
+        }
     };
 
     public final int id;
@@ -178,6 +190,8 @@ public static CompressionType forId(int id) {
                 return LZ4;
             case 4:
                 return ZSTD;
+            case 5:
+                return PASSTHROUGH;
             default:
                 throw new IllegalArgumentException("Unknown compression type id: " + id);
         }
@@ -194,6 +208,8 @@ else if (LZ4.name.equals(name))
             return LZ4;
         else if (ZSTD.name.equals(name))
             return ZSTD;
+        else if (PASSTHROUGH.name.equals(name))
+            return PASSTHROUGH;
         else
             throw new IllegalArgumentException("Unknown compression name: " + name);
     }
diff --git a/clients/src/main/java/org/apache/kafka/common/record/ControlRecordType.java b/clients/src/main/java/org/apache/kafka/common/record/ControlRecordType.java
index d5ead14df9126..ad41f1d9cd8dd 100644
--- a/clients/src/main/java/org/apache/kafka/common/record/ControlRecordType.java
+++ b/clients/src/main/java/org/apache/kafka/common/record/ControlRecordType.java
@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.common.record;
 
+import org.apache.kafka.common.InvalidRecordException;
 import org.apache.kafka.common.protocol.types.Field;
 import org.apache.kafka.common.protocol.types.Schema;
 import org.apache.kafka.common.protocol.types.Struct;
diff --git a/clients/src/main/java/org/apache/kafka/common/record/DefaultRecord.java b/clients/src/main/java/org/apache/kafka/common/record/DefaultRecord.java
index 109528fa1da18..637b5a9a522e1 100644
--- a/clients/src/main/java/org/apache/kafka/common/record/DefaultRecord.java
+++ b/clients/src/main/java/org/apache/kafka/common/record/DefaultRecord.java
@@ -16,11 +16,14 @@
  */
 package org.apache.kafka.common.record;
 
+import org.apache.kafka.common.InvalidRecordException;
 import org.apache.kafka.common.header.Header;
 import org.apache.kafka.common.header.internals.RecordHeader;
 import org.apache.kafka.common.utils.ByteUtils;
 import org.apache.kafka.common.utils.Checksums;
 import org.apache.kafka.common.utils.Crc32C;
+import org.apache.kafka.common.utils.PrimitiveRef;
+import org.apache.kafka.common.utils.PrimitiveRef.IntRef;
 import org.apache.kafka.common.utils.Utils;
 
 import java.io.DataInput;
@@ -29,6 +32,7 @@
 import java.nio.BufferUnderflowException;
 import java.nio.ByteBuffer;
 import java.util.Arrays;
+import java.util.Objects;
 import java.util.zip.Checksum;
 
 import static org.apache.kafka.common.record.RecordBatch.MAGIC_VALUE_V2;
@@ -77,14 +81,14 @@ public class DefaultRecord implements Record {
     private final ByteBuffer value;
     private final Header[] headers;
 
-    private DefaultRecord(int sizeInBytes,
-                          byte attributes,
-                          long offset,
-                          long timestamp,
-                          int sequence,
-                          ByteBuffer key,
-                          ByteBuffer value,
-                          Header[] headers) {
+    DefaultRecord(int sizeInBytes,
+                  byte attributes,
+                  long offset,
+                  long timestamp,
+                  int sequence,
+                  ByteBuffer key,
+                  ByteBuffer value,
+                  Header[] headers) {
         this.sizeInBytes = sizeInBytes;
         this.attributes = attributes;
         this.offset = offset;
@@ -266,8 +270,8 @@ public boolean equals(Object o) {
                 offset == that.offset &&
                 timestamp == that.timestamp &&
                 sequence == that.sequence &&
-                (key == null ? that.key == null : key.equals(that.key)) &&
-                (value == null ? that.value == null : value.equals(that.value)) &&
+                Objects.equals(key, that.key) &&
+                Objects.equals(value, that.value) &&
                 Arrays.equals(headers, that.headers);
     }
 
@@ -275,8 +279,8 @@ public boolean equals(Object o) {
     public int hashCode() {
         int result = sizeInBytes;
         result = 31 * result + (int) attributes;
-        result = 31 * result + (int) (offset ^ (offset >>> 32));
-        result = 31 * result + (int) (timestamp ^ (timestamp >>> 32));
+        result = 31 * result + Long.hashCode(offset);
+        result = 31 * result + Long.hashCode(timestamp);
         result = 31 * result + sequence;
         result = 31 * result + (key != null ? key.hashCode() : 0);
         result = 31 * result + (value != null ? value.hashCode() : 0);
@@ -311,6 +315,16 @@ public static DefaultRecord readFrom(ByteBuffer buffer,
                 baseSequence, logAppendTime);
     }
 
+    public static DefaultRecord readAllRecords(int sizeInBytes,
+        byte attributes,
+        long offset,
+        long timestamp,
+        int sequence,
+        ByteBuffer buffer) {
+        return new DefaultRecord(sizeInBytes, attributes, offset,
+            timestamp, sequence, null, buffer, Record.EMPTY_HEADERS);
+    }
+
     private static DefaultRecord readFrom(ByteBuffer buffer,
                                           int sizeInBytes,
                                           int sizeOfBodyInBytes,
@@ -369,6 +383,161 @@ private static DefaultRecord readFrom(ByteBuffer buffer,
         }
     }
 
+    public static PartialDefaultRecord readPartiallyFrom(DataInput input,
+                                                         byte[] skipArray,
+                                                         long baseOffset,
+                                                         long baseTimestamp,
+                                                         int baseSequence,
+                                                         Long logAppendTime) throws IOException {
+        int sizeOfBodyInBytes = ByteUtils.readVarint(input);
+        int totalSizeInBytes = ByteUtils.sizeOfVarint(sizeOfBodyInBytes) + sizeOfBodyInBytes;
+
+        return readPartiallyFrom(input, skipArray, totalSizeInBytes, sizeOfBodyInBytes, baseOffset, baseTimestamp,
+            baseSequence, logAppendTime);
+    }
+
+    private static PartialDefaultRecord readPartiallyFrom(DataInput input,
+                                                          byte[] skipArray,
+                                                          int sizeInBytes,
+                                                          int sizeOfBodyInBytes,
+                                                          long baseOffset,
+                                                          long baseTimestamp,
+                                                          int baseSequence,
+                                                          Long logAppendTime) throws IOException {
+        ByteBuffer skipBuffer = ByteBuffer.wrap(skipArray);
+        // set its limit to 0 to indicate no bytes readable yet
+        skipBuffer.limit(0);
+
+        try {
+            // reading the attributes / timestamp / offset and key-size does not require
+            // any byte array allocation and therefore we can just read them straight-forwardly
+            IntRef bytesRemaining = PrimitiveRef.ofInt(sizeOfBodyInBytes);
+
+            byte attributes = readByte(skipBuffer, input, bytesRemaining);
+            long timestampDelta = readVarLong(skipBuffer, input, bytesRemaining);
+            long timestamp = baseTimestamp + timestampDelta;
+            if (logAppendTime != null)
+                timestamp = logAppendTime;
+
+            int offsetDelta = readVarInt(skipBuffer, input, bytesRemaining);
+            long offset = baseOffset + offsetDelta;
+            int sequence = baseSequence >= 0 ?
+                DefaultRecordBatch.incrementSequence(baseSequence, offsetDelta) :
+                RecordBatch.NO_SEQUENCE;
+
+            // first skip key
+            int keySize = skipLengthDelimitedField(skipBuffer, input, bytesRemaining);
+
+            // then skip value
+            int valueSize = skipLengthDelimitedField(skipBuffer, input, bytesRemaining);
+
+            // then skip header
+            int numHeaders = readVarInt(skipBuffer, input, bytesRemaining);
+            if (numHeaders < 0)
+                throw new InvalidRecordException("Found invalid number of record headers " + numHeaders);
+            for (int i = 0; i < numHeaders; i++) {
+                int headerKeySize = skipLengthDelimitedField(skipBuffer, input, bytesRemaining);
+                if (headerKeySize < 0)
+                    throw new InvalidRecordException("Invalid negative header key size " + headerKeySize);
+
+                // headerValueSize
+                skipLengthDelimitedField(skipBuffer, input, bytesRemaining);
+            }
+
+            if (bytesRemaining.value > 0 || skipBuffer.remaining() > 0)
+                throw new InvalidRecordException("Invalid record size: expected to read " + sizeOfBodyInBytes +
+                    " bytes in record payload, but there are still bytes remaining");
+
+            return new PartialDefaultRecord(sizeInBytes, attributes, offset, timestamp, sequence, keySize, valueSize);
+        } catch (BufferUnderflowException | IllegalArgumentException e) {
+            throw new InvalidRecordException("Found invalid record structure", e);
+        }
+    }
+
+    private static byte readByte(ByteBuffer buffer, DataInput input, IntRef bytesRemaining) throws IOException {
+        if (buffer.remaining() < 1 && bytesRemaining.value > 0) {
+            readMore(buffer, input, bytesRemaining);
+        }
+
+        return buffer.get();
+    }
+
+    private static long readVarLong(ByteBuffer buffer, DataInput input, IntRef bytesRemaining) throws IOException {
+        if (buffer.remaining() < 10 && bytesRemaining.value > 0) {
+            readMore(buffer, input, bytesRemaining);
+        }
+
+        return ByteUtils.readVarlong(buffer);
+    }
+
+    private static int readVarInt(ByteBuffer buffer, DataInput input, IntRef bytesRemaining) throws IOException {
+        if (buffer.remaining() < 5 && bytesRemaining.value > 0) {
+            readMore(buffer, input, bytesRemaining);
+        }
+
+        return ByteUtils.readVarint(buffer);
+    }
+
+    private static int skipLengthDelimitedField(ByteBuffer buffer, DataInput input, IntRef bytesRemaining) throws IOException {
+        boolean needMore = false;
+        int sizeInBytes = -1;
+        int bytesToSkip = -1;
+
+        while (true) {
+            if (needMore) {
+                readMore(buffer, input, bytesRemaining);
+                needMore = false;
+            }
+
+            if (bytesToSkip < 0) {
+                if (buffer.remaining() < 5 && bytesRemaining.value > 0) {
+                    needMore = true;
+                } else {
+                    sizeInBytes = ByteUtils.readVarint(buffer);
+                    if (sizeInBytes <= 0)
+                        return sizeInBytes;
+                    else
+                        bytesToSkip = sizeInBytes;
+
+                }
+            } else {
+                if (bytesToSkip > buffer.remaining()) {
+                    bytesToSkip -= buffer.remaining();
+                    buffer.position(buffer.limit());
+                    needMore = true;
+                } else {
+                    buffer.position(buffer.position() + bytesToSkip);
+                    return sizeInBytes;
+                }
+            }
+        }
+    }
+
+    private static void readMore(ByteBuffer buffer, DataInput input, IntRef bytesRemaining) throws IOException {
+        if (bytesRemaining.value > 0) {
+            byte[] array = buffer.array();
+
+            // first copy the remaining bytes to the beginning of the array;
+            // at most 4 bytes would be shifted here
+            int stepsToLeftShift = buffer.position();
+            int bytesToLeftShift = buffer.remaining();
+            for (int i = 0; i < bytesToLeftShift; i++) {
+                array[i] = array[i + stepsToLeftShift];
+            }
+
+            // then try to read more bytes to the remaining of the array
+            int bytesRead = Math.min(bytesRemaining.value, array.length - bytesToLeftShift);
+            input.readFully(array, bytesToLeftShift, bytesRead);
+            buffer.rewind();
+            // only those many bytes are readable
+            buffer.limit(bytesToLeftShift + bytesRead);
+
+            bytesRemaining.value -= bytesRead;
+        } else {
+            throw new InvalidRecordException("Invalid record size: expected to read more bytes in record payload");
+        }
+    }
+
     private static Header[] readHeaders(ByteBuffer buffer, int numHeaders) {
         Header[] headers = new Header[numHeaders];
         for (int i = 0; i < numHeaders; i++) {
diff --git a/clients/src/main/java/org/apache/kafka/common/record/DefaultRecordBatch.java b/clients/src/main/java/org/apache/kafka/common/record/DefaultRecordBatch.java
index 19ddb0ef3fa1a..09e5abd4d9d66 100644
--- a/clients/src/main/java/org/apache/kafka/common/record/DefaultRecordBatch.java
+++ b/clients/src/main/java/org/apache/kafka/common/record/DefaultRecordBatch.java
@@ -16,7 +16,9 @@
  */
 package org.apache.kafka.common.record;
 
+import org.apache.kafka.common.InvalidRecordException;
 import org.apache.kafka.common.KafkaException;
+import org.apache.kafka.common.errors.CorruptRecordException;
 import org.apache.kafka.common.header.Header;
 import org.apache.kafka.common.utils.ByteBufferOutputStream;
 import org.apache.kafka.common.utils.ByteUtils;
@@ -33,6 +35,7 @@
 import java.util.Iterator;
 import java.util.List;
 import java.util.NoSuchElementException;
+import java.util.Objects;
 
 import static org.apache.kafka.common.record.Records.LOG_OVERHEAD;
 
@@ -127,6 +130,8 @@ public class DefaultRecordBatch extends AbstractRecordBatch implements MutableRe
     private static final int CONTROL_FLAG_MASK = 0x20;
     private static final byte TIMESTAMP_TYPE_MASK = 0x08;
 
+    private static final int MAX_SKIP_BUFFER_SIZE = 2048;
+
     private final ByteBuffer buffer;
 
     DefaultRecordBatch(ByteBuffer buffer) {
@@ -141,11 +146,11 @@ public byte magic() {
     @Override
     public void ensureValid() {
         if (sizeInBytes() < RECORD_BATCH_OVERHEAD)
-            throw new InvalidRecordException("Record batch is corrupt (the size " + sizeInBytes() +
+            throw new CorruptRecordException("Record batch is corrupt (the size " + sizeInBytes() +
                     " is smaller than the minimum allowed overhead " + RECORD_BATCH_OVERHEAD + ")");
 
         if (!isValid())
-            throw new InvalidRecordException("Record is corrupt (stored crc = " + checksum()
+            throw new CorruptRecordException("Record is corrupt (stored crc = " + checksum()
                     + ", computed crc = " + computeChecksum() + ")");
     }
 
@@ -250,42 +255,30 @@ public int partitionLeaderEpoch() {
         return buffer.getInt(PARTITION_LEADER_EPOCH_OFFSET);
     }
 
-    private CloseableIterator compressedIterator(BufferSupplier bufferSupplier) {
+    private CloseableIterator compressedIterator(BufferSupplier bufferSupplier, boolean skipKeyValue) {
         final ByteBuffer buffer = this.buffer.duplicate();
         buffer.position(RECORDS_OFFSET);
         final DataInputStream inputStream = new DataInputStream(compressionType().wrapForInput(buffer, magic(),
-                bufferSupplier));
+            bufferSupplier));
 
-        return new RecordIterator() {
-            @Override
-            protected Record readNext(long baseOffset, long firstTimestamp, int baseSequence, Long logAppendTime) {
-                try {
-                    return DefaultRecord.readFrom(inputStream, baseOffset, firstTimestamp, baseSequence, logAppendTime);
-                } catch (EOFException e) {
-                    throw new InvalidRecordException("Incorrect declared batch size, premature EOF reached");
-                } catch (IOException e) {
-                    throw new KafkaException("Failed to decompress record stream", e);
-                }
-            }
+        if (skipKeyValue) {
+            // this buffer is used to skip length delimited fields like key, value, headers
+            byte[] skipArray = new byte[MAX_SKIP_BUFFER_SIZE];
 
-            @Override
-            protected boolean ensureNoneRemaining() {
-                try {
-                    return inputStream.read() == -1;
-                } catch (IOException e) {
-                    throw new KafkaException("Error checking for remaining bytes after reading batch", e);
+            return new StreamRecordIterator(inputStream) {
+                @Override
+                protected Record doReadRecord(long baseOffset, long firstTimestamp, int baseSequence, Long logAppendTime) throws IOException {
+                    return DefaultRecord.readPartiallyFrom(inputStream, skipArray, baseOffset, firstTimestamp, baseSequence, logAppendTime);
                 }
-            }
-
-            @Override
-            public void close() {
-                try {
-                    inputStream.close();
-                } catch (IOException e) {
-                    throw new KafkaException("Failed to close record stream", e);
+            };
+        } else {
+            return new StreamRecordIterator(inputStream) {
+                @Override
+                protected Record doReadRecord(long baseOffset, long firstTimestamp, int baseSequence, Long logAppendTime) throws IOException {
+                    return DefaultRecord.readFrom(inputStream, baseOffset, firstTimestamp, baseSequence, logAppendTime);
                 }
-            }
-        };
+            };
+        }
     }
 
     private CloseableIterator uncompressedIterator() {
@@ -320,7 +313,7 @@ public Iterator iterator() {
         // for a normal iterator, we cannot ensure that the underlying compression stream is closed,
         // so we decompress the full record set here. Use cases which call for a lower memory footprint
         // can use `streamingIterator` at the cost of additional complexity
-        try (CloseableIterator iterator = compressedIterator(BufferSupplier.NO_CACHING)) {
+        try (CloseableIterator iterator = compressedIterator(BufferSupplier.NO_CACHING, false)) {
             List records = new ArrayList<>(count());
             while (iterator.hasNext())
                 records.add(iterator.next());
@@ -328,10 +321,59 @@ public Iterator iterator() {
         }
     }
 
+    @Override
+    public CloseableIterator skipKeyValueIterator(BufferSupplier bufferSupplier) {
+        if (count() == 0) {
+            return CloseableIterator.wrap(Collections.emptyIterator());
+        }
+
+        /*
+         * For uncompressed iterator, it is actually not worth skipping key / value / headers at all since
+         * its ByteBufferInputStream's skip() function is less efficient compared with just reading it actually
+         * as it will allocate new byte array.
+         */
+        if (!isCompressed())
+            return uncompressedIterator();
+
+        // we define this to be a closable iterator so that caller (i.e. the log validator) needs to close it
+        // while we can save memory footprint of not decompressing the full record set ahead of time
+        return compressedIterator(bufferSupplier, true);
+    }
+
+    @Override
+    public CloseableIterator shallowIterator() {
+        return new CloseableIterator() {
+            private boolean hasNext = true;
+            @Override
+            public void close() {
+
+            }
+
+            @Override
+            public boolean hasNext() {
+                return hasNext;
+            }
+
+            @Override
+            public Record next() {
+                if (!hasNext)
+                    throw new NoSuchElementException();
+                hasNext = false;
+                return DefaultRecord.readAllRecords(sizeInBytes(), attributes(), baseOffset(),
+                    firstTimestamp(), baseSequence(), buffer);
+            }
+
+            @Override
+            public void remove() {
+                throw new UnsupportedOperationException();
+            }
+        };
+    }
+
     @Override
     public CloseableIterator streamingIterator(BufferSupplier bufferSupplier) {
         if (isCompressed())
-            return compressedIterator(bufferSupplier);
+            return compressedIterator(bufferSupplier, false);
         else
             return uncompressedIterator();
     }
@@ -386,7 +428,7 @@ public boolean equals(Object o) {
             return false;
 
         DefaultRecordBatch that = (DefaultRecordBatch) o;
-        return buffer != null ? buffer.equals(that.buffer) : that.buffer == null;
+        return Objects.equals(buffer, that.buffer);
     }
 
     @Override
@@ -522,16 +564,16 @@ static int estimateBatchSizeUpperBound(ByteBuffer key, ByteBuffer value, Header[
         return RECORD_BATCH_OVERHEAD + DefaultRecord.recordSizeUpperBound(key, value, headers);
     }
 
-    public static int incrementSequence(int baseSequence, int increment) {
-        if (baseSequence > Integer.MAX_VALUE - increment)
-            return increment - (Integer.MAX_VALUE - baseSequence) - 1;
-        return baseSequence + increment;
+    public static int incrementSequence(int sequence, int increment) {
+        if (sequence > Integer.MAX_VALUE - increment)
+            return increment - (Integer.MAX_VALUE - sequence) - 1;
+        return sequence + increment;
     }
 
-    public static int decrementSequence(int baseSequence, int decrement) {
-        if (baseSequence < decrement)
-            return Integer.MAX_VALUE - (decrement - baseSequence) + 1;
-        return baseSequence - decrement;
+    public static int decrementSequence(int sequence, int decrement) {
+        if (sequence < decrement)
+            return Integer.MAX_VALUE - (decrement - sequence) + 1;
+        return sequence - decrement;
     }
 
     private abstract class RecordIterator implements CloseableIterator {
@@ -542,7 +584,7 @@ private abstract class RecordIterator implements CloseableIterator {
         private final int numRecords;
         private int readRecords = 0;
 
-        public RecordIterator() {
+        RecordIterator() {
             this.logAppendTime = timestampType() == TimestampType.LOG_APPEND_TIME ? maxTimestamp() : null;
             this.baseOffset = baseOffset();
             this.firstTimestamp = firstTimestamp();
@@ -587,6 +629,46 @@ public void remove() {
 
     }
 
+    private abstract class StreamRecordIterator extends RecordIterator {
+        private final DataInputStream inputStream;
+
+        StreamRecordIterator(DataInputStream inputStream) {
+            super();
+            this.inputStream = inputStream;
+        }
+
+        abstract Record doReadRecord(long baseOffset, long firstTimestamp, int baseSequence, Long logAppendTime) throws IOException;
+
+        @Override
+        protected Record readNext(long baseOffset, long firstTimestamp, int baseSequence, Long logAppendTime) {
+            try {
+                return doReadRecord(baseOffset, firstTimestamp, baseSequence, logAppendTime);
+            } catch (EOFException e) {
+                throw new InvalidRecordException("Incorrect declared batch size, premature EOF reached");
+            } catch (IOException e) {
+                throw new KafkaException("Failed to decompress record stream", e);
+            }
+        }
+
+        @Override
+        protected boolean ensureNoneRemaining() {
+            try {
+                return inputStream.read() == -1;
+            } catch (IOException e) {
+                throw new KafkaException("Error checking for remaining bytes after reading batch", e);
+            }
+        }
+
+        @Override
+        public void close() {
+            try {
+                inputStream.close();
+            } catch (IOException e) {
+                throw new KafkaException("Failed to close record stream", e);
+            }
+        }
+    }
+
     static class DefaultFileChannelRecordBatch extends FileLogInputStream.FileChannelRecordBatch {
 
         DefaultFileChannelRecordBatch(long offset,
diff --git a/clients/src/main/java/org/apache/kafka/common/record/EndTransactionMarker.java b/clients/src/main/java/org/apache/kafka/common/record/EndTransactionMarker.java
index 726b52a197378..4bf1ebf94a098 100644
--- a/clients/src/main/java/org/apache/kafka/common/record/EndTransactionMarker.java
+++ b/clients/src/main/java/org/apache/kafka/common/record/EndTransactionMarker.java
@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.common.record;
 
+import org.apache.kafka.common.InvalidRecordException;
 import org.apache.kafka.common.protocol.types.Field;
 import org.apache.kafka.common.protocol.types.Schema;
 import org.apache.kafka.common.protocol.types.Struct;
diff --git a/clients/src/main/java/org/apache/kafka/common/record/FileLogInputStream.java b/clients/src/main/java/org/apache/kafka/common/record/FileLogInputStream.java
index 472c7a7ac3e13..eb04705af2e48 100644
--- a/clients/src/main/java/org/apache/kafka/common/record/FileLogInputStream.java
+++ b/clients/src/main/java/org/apache/kafka/common/record/FileLogInputStream.java
@@ -27,6 +27,7 @@
 import java.nio.ByteBuffer;
 import java.nio.channels.FileChannel;
 import java.util.Iterator;
+import java.util.Objects;
 
 import static org.apache.kafka.common.record.Records.LOG_OVERHEAD;
 import static org.apache.kafka.common.record.Records.HEADER_SIZE_UP_TO_MAGIC;
@@ -156,6 +157,12 @@ public CloseableIterator streamingIterator(BufferSupplier bufferSupplier
             return loadFullBatch().streamingIterator(bufferSupplier);
         }
 
+        @Override
+        public CloseableIterator shallowIterator() {
+            // Not implemented
+            throw new UnsupportedOperationException();
+        }
+
         @Override
         public boolean isValid() {
             return loadFullBatch().isValid();
@@ -233,14 +240,14 @@ public boolean equals(Object o) {
             return offset == that.offset &&
                     position == that.position &&
                     batchSize == that.batchSize &&
-                    (channel == null ? thatChannel == null : channel.equals(thatChannel));
+                    Objects.equals(channel, thatChannel);
         }
 
         @Override
         public int hashCode() {
             FileChannel channel = fileRecords == null ? null : fileRecords.channel();
 
-            int result = (int) (offset ^ (offset >>> 32));
+            int result = Long.hashCode(offset);
             result = 31 * result + (channel != null ? channel.hashCode() : 0);
             result = 31 * result + position;
             result = 31 * result + batchSize;
diff --git a/clients/src/main/java/org/apache/kafka/common/record/FileRecords.java b/clients/src/main/java/org/apache/kafka/common/record/FileRecords.java
index d723ba091b72e..46ac481e77388 100644
--- a/clients/src/main/java/org/apache/kafka/common/record/FileRecords.java
+++ b/clients/src/main/java/org/apache/kafka/common/record/FileRecords.java
@@ -135,17 +135,20 @@ public void readInto(ByteBuffer buffer, int position) throws IOException {
      * @return A sliced wrapper on this message set limited based on the given position and size
      */
     public FileRecords slice(int position, int size) throws IOException {
+        // Cache current size in case concurrent write changes it
+        int currentSizeInBytes = sizeInBytes();
+
         if (position < 0)
             throw new IllegalArgumentException("Invalid position: " + position + " in read from " + this);
-        if (position > sizeInBytes() - start)
+        if (position > currentSizeInBytes - start)
             throw new IllegalArgumentException("Slice from position " + position + " exceeds end position of " + this);
         if (size < 0)
             throw new IllegalArgumentException("Invalid size: " + size + " in read from " + this);
 
         int end = this.start + position + size;
         // handle integer overflow or if end is beyond the end of the file
-        if (end < 0 || end >= start + sizeInBytes())
-            end = start + sizeInBytes();
+        if (end < 0 || end > start + currentSizeInBytes)
+            end = start + currentSizeInBytes;
         return new FileRecords(file, channel, this.start + position, end, true);
     }
 
@@ -208,11 +211,11 @@ public void trim() throws IOException {
     }
 
     /**
-     * Update the file reference (to be used with caution since this does not reopen the file channel)
-     * @param file The new file to use
+     * Update the parent directory (to be used with caution since this does not reopen the file channel)
+     * @param parentDir The new parent directory
      */
-    public void setFile(File file) {
-        this.file = file;
+    public void updateParentDir(File parentDir) {
+        this.file = new File(parentDir, file.getName());
     }
 
     /**
@@ -370,7 +373,8 @@ public Iterable batches() {
 
     @Override
     public String toString() {
-        return "FileRecords(file= " + file +
+        return "FileRecords(size=" + sizeInBytes() +
+                ", file=" + file +
                 ", start=" + start +
                 ", end=" + end +
                 ")";
@@ -484,7 +488,7 @@ public boolean equals(Object o) {
 
         @Override
         public int hashCode() {
-            int result = (int) (offset ^ (offset >>> 32));
+            int result = Long.hashCode(offset);
             result = 31 * result + position;
             result = 31 * result + size;
             return result;
diff --git a/clients/src/main/java/org/apache/kafka/common/record/LazyDownConversionRecords.java b/clients/src/main/java/org/apache/kafka/common/record/LazyDownConversionRecords.java
index 217870d5faa0d..56c6809b45463 100644
--- a/clients/src/main/java/org/apache/kafka/common/record/LazyDownConversionRecords.java
+++ b/clients/src/main/java/org/apache/kafka/common/record/LazyDownConversionRecords.java
@@ -100,12 +100,21 @@ public boolean equals(Object o) {
     @Override
     public int hashCode() {
         int result = toMagic;
-        result = 31 * result + (int) (firstOffset ^ (firstOffset >>> 32));
+        result = 31 * result + Long.hashCode(firstOffset);
         result = 31 * result + topicPartition.hashCode();
         result = 31 * result + records.hashCode();
         return result;
     }
 
+    @Override
+    public String toString() {
+        return "LazyDownConversionRecords(size=" + sizeInBytes +
+                ", underlying=" + records +
+                ", toMagic=" + toMagic +
+                ", firstOffset=" + firstOffset +
+                ")";
+    }
+
     public java.util.Iterator> iterator(long maximumReadSize) {
         // We typically expect only one iterator instance to be created, so null out the first converted batch after
         // first use to make it available for GC.
diff --git a/clients/src/main/java/org/apache/kafka/common/record/LegacyRecord.java b/clients/src/main/java/org/apache/kafka/common/record/LegacyRecord.java
index 482c4a65efc1e..32c5aa81d7530 100644
--- a/clients/src/main/java/org/apache/kafka/common/record/LegacyRecord.java
+++ b/clients/src/main/java/org/apache/kafka/common/record/LegacyRecord.java
@@ -17,6 +17,7 @@
 package org.apache.kafka.common.record;
 
 import org.apache.kafka.common.KafkaException;
+import org.apache.kafka.common.errors.CorruptRecordException;
 import org.apache.kafka.common.utils.ByteBufferOutputStream;
 import org.apache.kafka.common.utils.ByteUtils;
 import org.apache.kafka.common.utils.Checksums;
@@ -135,11 +136,11 @@ public TimestampType wrapperRecordTimestampType() {
      */
     public void ensureValid() {
         if (sizeInBytes() < RECORD_OVERHEAD_V0)
-            throw new InvalidRecordException("Record is corrupt (crc could not be retrieved as the record is too "
+            throw new CorruptRecordException("Record is corrupt (crc could not be retrieved as the record is too "
                     + "small, size = " + sizeInBytes() + ")");
 
         if (!isValid())
-            throw new InvalidRecordException("Record is corrupt (stored crc = " + checksum()
+            throw new CorruptRecordException("Record is corrupt (stored crc = " + checksum()
                     + ", computed crc = " + computeChecksum() + ")");
     }
 
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 c6db18d419162..8f73565d1b40f 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
@@ -16,7 +16,6 @@
  */
 package org.apache.kafka.common.record;
 
-import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.errors.CorruptRecordException;
 import org.apache.kafka.common.record.MemoryRecords.RecordFilter.BatchRetention;
@@ -32,7 +31,6 @@
 import java.nio.channels.GatheringByteChannel;
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Iterator;
 import java.util.List;
 import java.util.Objects;
 
@@ -279,34 +277,9 @@ public Iterable batches() {
 
     @Override
     public String toString() {
-        StringBuilder builder = new StringBuilder();
-        builder.append('[');
-
-        Iterator batchIterator = batches.iterator();
-        while (batchIterator.hasNext()) {
-            RecordBatch batch = batchIterator.next();
-            try (CloseableIterator recordsIterator = batch.streamingIterator(BufferSupplier.create())) {
-                while (recordsIterator.hasNext()) {
-                    Record record = recordsIterator.next();
-                    appendRecordToStringBuilder(builder, record.toString());
-                    if (recordsIterator.hasNext())
-                        builder.append(", ");
-                }
-            } catch (KafkaException e) {
-                appendRecordToStringBuilder(builder, "CORRUPTED");
-            }
-            if (batchIterator.hasNext())
-                builder.append(", ");
-        }
-        builder.append(']');
-        return builder.toString();
-    }
-
-    private void appendRecordToStringBuilder(StringBuilder builder, String recordAsString) {
-        builder.append('(')
-            .append("record=")
-            .append(recordAsString)
-            .append(")");
+        return "MemoryRecords(size=" + sizeInBytes() +
+                ", buffer=" + buffer +
+                ")";
     }
 
     @Override
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 7512c8273dab4..4319d199028e4 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
@@ -20,6 +20,7 @@
 import org.apache.kafka.common.header.Header;
 import org.apache.kafka.common.protocol.types.Struct;
 import org.apache.kafka.common.utils.ByteBufferOutputStream;
+import org.apache.kafka.common.utils.Utils;
 
 import java.io.DataOutputStream;
 import java.io.IOException;
@@ -38,6 +39,9 @@
  * This will release resources like compression buffers that can be relatively large (64 KB for LZ4).
  */
 public class MemoryRecordsBuilder implements AutoCloseable {
+    // TODO(viswamy): Revert this change once all brokers are on Message_v2
+    private static final String INTERNAL_HEADER_PREFIX = "_";
+
     private static final float COMPRESSION_RATE_ESTIMATION_FACTOR = 1.05f;
     private static final DataOutputStream CLOSED_STREAM = new DataOutputStream(new OutputStream() {
         @Override
@@ -68,6 +72,7 @@ public void write(int b) {
     // Used to append records, may compress data on the fly
     private DataOutputStream appendStream;
     private boolean isTransactional;
+    private boolean usePassthrough = false;
     private long producerId;
     private short producerEpoch;
     private int baseSequence;
@@ -108,7 +113,6 @@ public MemoryRecordsBuilder(ByteBufferOutputStream bufferStream,
 
         this.magic = magic;
         this.timestampType = timestampType;
-        this.compressionType = compressionType;
         this.baseOffset = baseOffset;
         this.logAppendTime = logAppendTime;
         this.numRecords = 0;
@@ -123,7 +127,15 @@ public MemoryRecordsBuilder(ByteBufferOutputStream bufferStream,
         this.partitionLeaderEpoch = partitionLeaderEpoch;
         this.writeLimit = writeLimit;
         this.initialPosition = bufferStream.position();
-        this.batchHeaderSizeInBytes = AbstractRecords.recordBatchHeaderSizeInBytes(magic, compressionType);
+
+        if (compressionType == CompressionType.PASSTHROUGH) {
+            usePassthrough = true;
+            this.compressionType = CompressionType.NONE;
+        } else {
+            this.compressionType = compressionType;
+        }
+
+        this.batchHeaderSizeInBytes = usePassthrough ? 0 : AbstractRecords.recordBatchHeaderSizeInBytes(magic, this.compressionType);
 
         bufferStream.position(initialPosition + batchHeaderSizeInBytes);
         this.bufferStream = bufferStream;
@@ -316,9 +328,11 @@ public void close() {
             buffer().position(initialPosition);
             builtRecords = MemoryRecords.EMPTY;
         } else {
-            if (magic > RecordBatch.MAGIC_VALUE_V1)
-                this.actualCompressionRatio = (float) writeDefaultBatchHeader() / this.uncompressedRecordsSizeInBytes;
-            else if (compressionType != CompressionType.NONE)
+            if (magic > RecordBatch.MAGIC_VALUE_V1) {
+                if (!usePassthrough) {
+                    this.actualCompressionRatio = (float) writeDefaultBatchHeader() / this.uncompressedRecordsSizeInBytes;
+                }
+            } else if (compressionType != CompressionType.NONE)
                 this.actualCompressionRatio = (float) writeLegacyCompressedWrapperHeader() / this.uncompressedRecordsSizeInBytes;
 
             ByteBuffer buffer = buffer().duplicate();
@@ -392,6 +406,18 @@ private int writeLegacyCompressedWrapperHeader() {
         return writtenCompressed;
     }
 
+    // Validates and throws an exception if the headers contain non-internal record headers
+    private void validateHeaders(Header[] headers) {
+        if (magic < RecordBatch.MAGIC_VALUE_V2 && headers != null && headers.length > 0) {
+            for (Header header : headers) {
+                if (!header.key().startsWith(INTERNAL_HEADER_PREFIX)) {
+                    throw new IllegalArgumentException("Magic v" + magic + " does not support record headers. [Key = "
+                        + header.key() + "]");
+                }
+            }
+        }
+    }
+
     /**
      * Append a record and return its checksum for message format v0 and v1, or null for v2 and above.
      */
@@ -408,8 +434,7 @@ private Long appendWithOffset(long offset, boolean isControlRecord, long timesta
             if (timestamp < 0 && timestamp != RecordBatch.NO_TIMESTAMP)
                 throw new IllegalArgumentException("Invalid negative timestamp " + timestamp);
 
-            if (magic < RecordBatch.MAGIC_VALUE_V2 && headers != null && headers.length > 0)
-                throw new IllegalArgumentException("Magic v" + magic + " does not support record headers");
+            validateHeaders(headers);
 
             if (firstTimestamp == null)
                 firstTimestamp = timestamp;
@@ -418,7 +443,11 @@ private Long appendWithOffset(long offset, boolean isControlRecord, long timesta
                 appendDefaultRecord(offset, timestamp, key, value, headers);
                 return null;
             } else {
-                return appendLegacyRecord(offset, timestamp, key, value);
+                if (usePassthrough) {
+                    return appendPassthroughLegacyBatch(timestamp, value);
+                } else {
+                    return appendLegacyRecord(offset, timestamp, key, value, magic);
+                }
             }
         } catch (IOException e) {
             throw new KafkaException("I/O exception when writing to the append stream, closing", e);
@@ -587,6 +616,60 @@ public void appendUncheckedWithOffset(long offset, LegacyRecord record) {
         }
     }
 
+    /**
+     * Append a record without doing offset/magic validation (this should only be used in testing).
+     *
+     * @param offset The offset of the record
+     * @param record The record to add
+     */
+    public void appendUncheckedWithOffset(long offset, SimpleRecord record) throws IOException {
+        if (magic >= RecordBatch.MAGIC_VALUE_V2) {
+            int offsetDelta = (int) (offset - baseOffset);
+            long timestamp = record.timestamp();
+            if (firstTimestamp == null)
+                firstTimestamp = timestamp;
+
+            int sizeInBytes = DefaultRecord.writeTo(appendStream,
+                offsetDelta,
+                timestamp - firstTimestamp,
+                record.key(),
+                record.value(),
+                record.headers());
+            recordWritten(offset, timestamp, sizeInBytes);
+        } else {
+            LegacyRecord legacyRecord = LegacyRecord.create(magic,
+                record.timestamp(),
+                Utils.toNullableArray(record.key()),
+                Utils.toNullableArray(record.value()));
+            appendUncheckedWithOffset(offset, legacyRecord);
+        }
+    }
+
+    /**
+     * Append an already assembled (and compressed) batch. The offset is always recorded as 0 as the value is
+     * a single complete batch. If multiple batches are appended, they will be handled with request aggregation on
+     * the broker
+     * @param timestamp
+     * @param value
+     * @return crc of the record
+     */
+    public long appendPassthroughLegacyBatch(long timestamp, ByteBuffer value) {
+        try {
+            ensureOpenForRecordAppend();
+            // For passthrough compression (from shallow iterator), the value is a wrapper message
+            LegacyRecord wrapperMessage = new LegacyRecord(value);
+
+            int size = LegacyRecord.recordSize(wrapperMessage.magic(), wrapperMessage.key(), wrapperMessage.value());
+            AbstractLegacyRecordBatch.writeHeader(appendStream, toInnerOffset(0L), size);
+
+            long crc = LegacyRecord.write(appendStream, wrapperMessage.magic(), timestamp, wrapperMessage.key(), wrapperMessage.value(), wrapperMessage.compressionType(), timestampType);
+            recordWritten(0L, timestamp, size + Records.LOG_OVERHEAD);
+            return crc;
+        } catch (IOException e) {
+            throw new KafkaException("I/O exception when writing to the legacy append stream, closing", e);
+        }
+    }
+
     /**
      * Append a record at the next sequential offset.
      * @param record the record to add
@@ -623,16 +706,27 @@ public void append(LegacyRecord record) {
         appendWithOffset(nextSequentialOffset(), record);
     }
 
+    /**
+     * Write the entire buffer to `out` as-it-is and return its size
+     */
+    public int writePassthrough(DataOutputStream out,
+        ByteBuffer buffer) throws IOException {
+        int bufferSize = buffer.remaining();
+        Utils.writeTo(out, buffer, bufferSize);
+        return bufferSize;
+    }
+
     private void appendDefaultRecord(long offset, long timestamp, ByteBuffer key, ByteBuffer value,
                                      Header[] headers) throws IOException {
         ensureOpenForRecordAppend();
         int offsetDelta = (int) (offset - baseOffset);
         long timestampDelta = timestamp - firstTimestamp;
-        int sizeInBytes = DefaultRecord.writeTo(appendStream, offsetDelta, timestampDelta, key, value, headers);
+        int sizeInBytes = usePassthrough ? writePassthrough(appendStream, value) :
+            DefaultRecord.writeTo(appendStream, offsetDelta, timestampDelta, key, value, headers);
         recordWritten(offset, timestamp, sizeInBytes);
     }
 
-    private long appendLegacyRecord(long offset, long timestamp, ByteBuffer key, ByteBuffer value) throws IOException {
+    private long appendLegacyRecord(long offset, long timestamp, ByteBuffer key, ByteBuffer value, byte magic) throws IOException {
         ensureOpenForRecordAppend();
         if (compressionType == CompressionType.NONE && timestampType == TimestampType.LOG_APPEND_TIME)
             timestamp = logAppendTime;
@@ -727,6 +821,11 @@ public boolean hasRoomFor(long timestamp, ByteBuffer key, ByteBuffer value, Head
         if (numRecords == 0)
             return true;
 
+        // For passthrough V2, ensure one producerBatch only has one DefaultRecordBatch, and since
+        // in this case, DefaultRecordBatch is the value part of a DefaultRecord, so we only allow one record
+        if (magic >= RecordBatch.MAGIC_VALUE_V2 && usePassthrough)
+            return false;
+
         final int recordSize;
         if (magic < RecordBatch.MAGIC_VALUE_V2) {
             recordSize = Records.LOG_OVERHEAD + LegacyRecord.recordSize(magic, key, value);
diff --git a/clients/src/main/java/org/apache/kafka/common/record/MutableRecordBatch.java b/clients/src/main/java/org/apache/kafka/common/record/MutableRecordBatch.java
index c13bb5a2880e5..8c0dc2363e948 100644
--- a/clients/src/main/java/org/apache/kafka/common/record/MutableRecordBatch.java
+++ b/clients/src/main/java/org/apache/kafka/common/record/MutableRecordBatch.java
@@ -17,6 +17,7 @@
 package org.apache.kafka.common.record;
 
 import org.apache.kafka.common.utils.ByteBufferOutputStream;
+import org.apache.kafka.common.utils.CloseableIterator;
 
 /**
  * A mutable record batch is one that can be modified in place (without copying). This is used by the broker
@@ -55,4 +56,12 @@ public interface MutableRecordBatch extends RecordBatch {
      */
     void writeTo(ByteBufferOutputStream outputStream);
 
+    /**
+     * Return an iterator which skips parsing key, value and headers from the record stream, and therefore the resulted
+     * {@code org.apache.kafka.common.record.Record}'s key and value fields would be empty. This iterator is used
+     * when the read record's key and value are not needed and hence can save some byte buffer allocating / GC overhead.
+     *
+     * @return The closeable iterator
+     */
+    CloseableIterator skipKeyValueIterator(BufferSupplier bufferSupplier);
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/record/PartialDefaultRecord.java b/clients/src/main/java/org/apache/kafka/common/record/PartialDefaultRecord.java
new file mode 100644
index 0000000000000..67ca1abafbc9f
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/record/PartialDefaultRecord.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.common.record;
+
+import org.apache.kafka.common.header.Header;
+
+import java.nio.ByteBuffer;
+
+public class PartialDefaultRecord extends DefaultRecord {
+
+    private final int keySize;
+    private final int valueSize;
+
+    PartialDefaultRecord(int sizeInBytes,
+                         byte attributes,
+                         long offset,
+                         long timestamp,
+                         int sequence,
+                         int keySize,
+                         int valueSize) {
+        super(sizeInBytes, attributes, offset, timestamp, sequence, null, null, null);
+
+        this.keySize = keySize;
+        this.valueSize = valueSize;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        return super.equals(o) &&
+            this.keySize == ((PartialDefaultRecord) o).keySize &&
+            this.valueSize == ((PartialDefaultRecord) o).valueSize;
+    }
+
+    @Override
+    public int hashCode() {
+        int result = super.hashCode();
+        result = 31 * result + keySize;
+        result = 31 * result + valueSize;
+        return result;
+    }
+
+    @Override
+    public String toString() {
+        return String.format("PartialDefaultRecord(offset=%d, timestamp=%d, key=%d bytes, value=%d bytes)",
+            offset(),
+            timestamp(),
+            keySize,
+            valueSize);
+    }
+
+    @Override
+    public int keySize() {
+        return keySize;
+    }
+
+    @Override
+    public boolean hasKey() {
+        return keySize >= 0;
+    }
+
+    @Override
+    public ByteBuffer key() {
+        throw new UnsupportedOperationException("key is skipped in PartialDefaultRecord");
+    }
+
+    @Override
+    public int valueSize() {
+        return valueSize;
+    }
+
+    @Override
+    public boolean hasValue() {
+        return valueSize >= 0;
+    }
+
+    @Override
+    public ByteBuffer value() {
+        throw new UnsupportedOperationException("value is skipped in PartialDefaultRecord");
+    }
+
+    @Override
+    public Header[] headers() {
+        throw new UnsupportedOperationException("headers is skipped in PartialDefaultRecord");
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/record/RecordBatch.java b/clients/src/main/java/org/apache/kafka/common/record/RecordBatch.java
index 65a6a95fbe41f..2a8b547f09df1 100644
--- a/clients/src/main/java/org/apache/kafka/common/record/RecordBatch.java
+++ b/clients/src/main/java/org/apache/kafka/common/record/RecordBatch.java
@@ -230,6 +230,16 @@ public interface RecordBatch extends Iterable {
      */
     CloseableIterator streamingIterator(BufferSupplier decompressionBufferSupplier);
 
+    /**
+     * Return a iterator which does shallow iteration over the record batch, i.e. no decompression.
+     * If the record batch is not compressed, it will also return the entire recordBatch.
+     * 1. For V0/V1, since the record and recordBatch shares the same schema, the behavior is not changed, i.e. it will return the entire
+     *    recordBatch and send it to the producer.
+     * 2. For V2, it will also return the entire recordBatch and later wrap the recordBatch as a single record and send it to producer.
+     * @return The closeable iterator
+     */
+    CloseableIterator shallowIterator();
+
     /**
      * Check whether this is a control batch (i.e. whether the control bit is set in the batch attributes).
      * For magic versions prior to 2, this is always false.
diff --git a/clients/src/main/java/org/apache/kafka/common/record/RecordsUtil.java b/clients/src/main/java/org/apache/kafka/common/record/RecordsUtil.java
index 3b0c59a24a521..423d1e1656f86 100644
--- a/clients/src/main/java/org/apache/kafka/common/record/RecordsUtil.java
+++ b/clients/src/main/java/org/apache/kafka/common/record/RecordsUtil.java
@@ -18,6 +18,7 @@
 
 import org.apache.kafka.common.errors.UnsupportedCompressionTypeException;
 import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.common.utils.Utils;
 
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
@@ -80,9 +81,11 @@ protected static ConvertedRecords downConvert(Iterable>> 32));
+        result = 31 * result + Long.hashCode(timestamp);
         result = 31 * result + Arrays.hashCode(headers);
         return result;
     }
diff --git a/clients/src/main/java/org/apache/kafka/common/replica/ClientMetadata.java b/clients/src/main/java/org/apache/kafka/common/replica/ClientMetadata.java
new file mode 100644
index 0000000000000..b328733dc7b3e
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/replica/ClientMetadata.java
@@ -0,0 +1,124 @@
+/*
+ * 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.replica;
+
+import org.apache.kafka.common.security.auth.KafkaPrincipal;
+
+import java.net.InetAddress;
+import java.util.Objects;
+
+/**
+ * Holder for all the client metadata required to determine a preferred replica.
+ */
+public interface ClientMetadata {
+
+    /**
+     * Rack ID sent by the client
+     */
+    String rackId();
+
+    /**
+     * Client ID sent by the client
+     */
+    String clientId();
+
+    /**
+     * Incoming address of the client
+     */
+    InetAddress clientAddress();
+
+    /**
+     * Security principal of the client
+     */
+    KafkaPrincipal principal();
+
+    /**
+     * Listener name for the client
+     */
+    String listenerName();
+
+
+    class DefaultClientMetadata implements ClientMetadata {
+        private final String rackId;
+        private final String clientId;
+        private final InetAddress clientAddress;
+        private final KafkaPrincipal principal;
+        private final String listenerName;
+
+        public DefaultClientMetadata(String rackId, String clientId, InetAddress clientAddress,
+                                     KafkaPrincipal principal, String listenerName) {
+            this.rackId = rackId;
+            this.clientId = clientId;
+            this.clientAddress = clientAddress;
+            this.principal = principal;
+            this.listenerName = listenerName;
+        }
+
+        @Override
+        public String rackId() {
+            return rackId;
+        }
+
+        @Override
+        public String clientId() {
+            return clientId;
+        }
+
+        @Override
+        public InetAddress clientAddress() {
+            return clientAddress;
+        }
+
+        @Override
+        public KafkaPrincipal principal() {
+            return principal;
+        }
+
+        @Override
+        public String listenerName() {
+            return listenerName;
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) return true;
+            if (o == null || getClass() != o.getClass()) return false;
+            DefaultClientMetadata that = (DefaultClientMetadata) o;
+            return Objects.equals(rackId, that.rackId) &&
+                    Objects.equals(clientId, that.clientId) &&
+                    Objects.equals(clientAddress, that.clientAddress) &&
+                    Objects.equals(principal, that.principal) &&
+                    Objects.equals(listenerName, that.listenerName);
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(rackId, clientId, clientAddress, principal, listenerName);
+        }
+
+        @Override
+        public String toString() {
+            return "DefaultClientMetadata{" +
+                    "rackId='" + rackId + '\'' +
+                    ", clientId='" + clientId + '\'' +
+                    ", clientAddress=" + clientAddress +
+                    ", principal=" + principal +
+                    ", listenerName='" + listenerName + '\'' +
+                    '}';
+        }
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/replica/PartitionView.java b/clients/src/main/java/org/apache/kafka/common/replica/PartitionView.java
new file mode 100644
index 0000000000000..8174e631305ab
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/replica/PartitionView.java
@@ -0,0 +1,72 @@
+/*
+ * 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.replica;
+
+import java.util.Collections;
+import java.util.Objects;
+import java.util.Set;
+
+/**
+ * View of a partition used by {@link ReplicaSelector} to determine a preferred replica.
+ */
+public interface PartitionView {
+    Set replicas();
+
+    ReplicaView leader();
+
+    class DefaultPartitionView implements PartitionView {
+        private final Set replicas;
+        private final ReplicaView leader;
+
+        public DefaultPartitionView(Set replicas, ReplicaView leader) {
+            this.replicas = Collections.unmodifiableSet(replicas);
+            this.leader = leader;
+        }
+
+        @Override
+        public Set replicas() {
+            return replicas;
+        }
+
+        @Override
+        public ReplicaView leader() {
+            return leader;
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) return true;
+            if (o == null || getClass() != o.getClass()) return false;
+            DefaultPartitionView that = (DefaultPartitionView) o;
+            return Objects.equals(replicas, that.replicas) &&
+                    Objects.equals(leader, that.leader);
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(replicas, leader);
+        }
+
+        @Override
+        public String toString() {
+            return "DefaultPartitionView{" +
+                    "replicas=" + replicas +
+                    ", leader=" + leader +
+                    '}';
+        }
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/replica/RackAwareReplicaSelector.java b/clients/src/main/java/org/apache/kafka/common/replica/RackAwareReplicaSelector.java
new file mode 100644
index 0000000000000..8ae68723af1d5
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/replica/RackAwareReplicaSelector.java
@@ -0,0 +1,54 @@
+/*
+ * 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.replica;
+
+import org.apache.kafka.common.TopicPartition;
+
+import java.util.Optional;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+/**
+ * Returns a replica whose rack id is equal to the rack id specified in the client request metadata. If no such replica
+ * is found, returns the leader.
+ */
+public class RackAwareReplicaSelector implements ReplicaSelector {
+
+    @Override
+    public Optional select(TopicPartition topicPartition,
+                                        ClientMetadata clientMetadata,
+                                        PartitionView partitionView) {
+        if (clientMetadata.rackId() != null && !clientMetadata.rackId().isEmpty()) {
+            Set sameRackReplicas = partitionView.replicas().stream()
+                    .filter(replicaInfo -> clientMetadata.rackId().equals(replicaInfo.endpoint().rack()))
+                    .collect(Collectors.toSet());
+            if (sameRackReplicas.isEmpty()) {
+                return Optional.of(partitionView.leader());
+            } else {
+                if (sameRackReplicas.contains(partitionView.leader())) {
+                    // Use the leader if it's in this rack
+                    return Optional.of(partitionView.leader());
+                } else {
+                    // Otherwise, get the most caught-up replica
+                    return sameRackReplicas.stream().max(ReplicaView.comparator());
+                }
+            }
+        } else {
+            return Optional.of(partitionView.leader());
+        }
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/replica/ReplicaSelector.java b/clients/src/main/java/org/apache/kafka/common/replica/ReplicaSelector.java
new file mode 100644
index 0000000000000..301fc9fdc4b36
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/replica/ReplicaSelector.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.common.replica;
+
+import org.apache.kafka.common.Configurable;
+import org.apache.kafka.common.TopicPartition;
+
+import java.io.Closeable;
+import java.io.IOException;
+import java.util.Map;
+import java.util.Optional;
+
+/**
+ * Plug-able interface for selecting a preferred read replica given the current set of replicas for a partition
+ * and metadata from the client.
+ */
+public interface ReplicaSelector extends Configurable, Closeable {
+
+    /**
+     * Select the preferred replica a client should use for fetching. If no replica is available, this will return an
+     * empty optional.
+     */
+    Optional select(TopicPartition topicPartition,
+                                 ClientMetadata clientMetadata,
+                                 PartitionView partitionView);
+    @Override
+    default void close() throws IOException {
+        // No-op by default
+    }
+
+    @Override
+    default void configure(Map configs) {
+        // No-op by default
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/replica/ReplicaView.java b/clients/src/main/java/org/apache/kafka/common/replica/ReplicaView.java
new file mode 100644
index 0000000000000..69c6cf7fef970
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/replica/ReplicaView.java
@@ -0,0 +1,105 @@
+/*
+ * 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.replica;
+
+import org.apache.kafka.common.Node;
+
+import java.util.Comparator;
+import java.util.Objects;
+
+/**
+ * View of a replica used by {@link ReplicaSelector} to determine a preferred replica.
+ */
+public interface ReplicaView {
+
+    /**
+     * The endpoint information for this replica (hostname, port, rack, etc)
+     */
+    Node endpoint();
+
+    /**
+     * The log end offset for this replica
+     */
+    long logEndOffset();
+
+    /**
+     * The number of milliseconds (if any) since the last time this replica was caught up to the high watermark.
+     * For a leader replica, this is always zero.
+     */
+    long timeSinceLastCaughtUpMs();
+
+    /**
+     * Comparator for ReplicaView that returns in the order of "most caught up". This is used for deterministic
+     * selection of a replica when there is a tie from a selector.
+     */
+    static Comparator comparator() {
+        return Comparator.comparingLong(ReplicaView::logEndOffset)
+            .thenComparing(Comparator.comparingLong(ReplicaView::timeSinceLastCaughtUpMs).reversed())
+            .thenComparing(replicaInfo -> replicaInfo.endpoint().id());
+    }
+
+    class DefaultReplicaView implements ReplicaView {
+        private final Node endpoint;
+        private final long logEndOffset;
+        private final long timeSinceLastCaughtUpMs;
+
+        public DefaultReplicaView(Node endpoint, long logEndOffset, long timeSinceLastCaughtUpMs) {
+            this.endpoint = endpoint;
+            this.logEndOffset = logEndOffset;
+            this.timeSinceLastCaughtUpMs = timeSinceLastCaughtUpMs;
+        }
+
+        @Override
+        public Node endpoint() {
+            return endpoint;
+        }
+
+        @Override
+        public long logEndOffset() {
+            return logEndOffset;
+        }
+
+        @Override
+        public long timeSinceLastCaughtUpMs() {
+            return timeSinceLastCaughtUpMs;
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) return true;
+            if (o == null || getClass() != o.getClass()) return false;
+            DefaultReplicaView that = (DefaultReplicaView) o;
+            return logEndOffset == that.logEndOffset &&
+                    Objects.equals(endpoint, that.endpoint) &&
+                    Objects.equals(timeSinceLastCaughtUpMs, that.timeSinceLastCaughtUpMs);
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(endpoint, logEndOffset, timeSinceLastCaughtUpMs);
+        }
+
+        @Override
+        public String toString() {
+            return "DefaultReplicaView{" +
+                    "endpoint=" + endpoint +
+                    ", logEndOffset=" + logEndOffset +
+                    ", timeSinceLastCaughtUpMs=" + timeSinceLastCaughtUpMs +
+                    '}';
+        }
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/AbstractControlRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/AbstractControlRequest.java
index d1fae0502c438..cf22160aad480 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/AbstractControlRequest.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/AbstractControlRequest.java
@@ -17,64 +17,38 @@
 package org.apache.kafka.common.requests;
 
 import org.apache.kafka.common.protocol.ApiKeys;
-import org.apache.kafka.common.protocol.types.Field;
-import org.apache.kafka.common.protocol.types.Struct;
 
 // Abstract class for all control requests including UpdateMetadataRequest, LeaderAndIsrRequest and StopReplicaRequest
 public abstract class AbstractControlRequest extends AbstractRequest {
-    public static final long UNKNOWN_BROKER_EPOCH = -1L;
-
-    protected static final Field.Int32 CONTROLLER_ID = new Field.Int32("controller_id", "The controller id");
-    protected static final Field.Int32 CONTROLLER_EPOCH = new Field.Int32("controller_epoch", "The controller epoch");
-    protected static final Field.Int64 BROKER_EPOCH = new Field.Int64("broker_epoch", "The broker epoch");
 
-    protected final int controllerId;
-    protected final int controllerEpoch;
-    protected final long brokerEpoch;
+    public static final long UNKNOWN_BROKER_EPOCH = -1L;
 
     public static abstract class Builder extends AbstractRequest.Builder {
         protected final int controllerId;
         protected final int controllerEpoch;
         protected final long brokerEpoch;
+        protected final long maxBrokerEpoch;
 
-        protected Builder(ApiKeys api, short version, int controllerId, int controllerEpoch, long brokerEpoch) {
+        protected Builder(ApiKeys api, short version, int controllerId, int controllerEpoch, long brokerEpoch, long maxBrokerEpoch) {
             super(api, version);
             this.controllerId = controllerId;
             this.controllerEpoch = controllerEpoch;
             this.brokerEpoch = brokerEpoch;
+            this.maxBrokerEpoch = maxBrokerEpoch;
         }
 
     }
 
-    public int controllerId() {
-        return controllerId;
-    }
-
-    public int controllerEpoch() {
-        return controllerEpoch;
+    protected AbstractControlRequest(ApiKeys api, short version) {
+        super(api, version);
     }
 
-    public long brokerEpoch() {
-        return brokerEpoch;
-    }
+    public abstract int controllerId();
 
-    protected AbstractControlRequest(ApiKeys api, short version, int controllerId, int controllerEpoch, long brokerEpoch) {
-        super(api, version);
-        this.controllerId = controllerId;
-        this.controllerEpoch = controllerEpoch;
-        this.brokerEpoch = brokerEpoch;
-    }
+    public abstract int controllerEpoch();
 
-    protected AbstractControlRequest(ApiKeys api, Struct struct, short version) {
-        super(api, version);
-        this.controllerId = struct.get(CONTROLLER_ID);
-        this.controllerEpoch = struct.get(CONTROLLER_EPOCH);
-        this.brokerEpoch = struct.getOrElse(BROKER_EPOCH, UNKNOWN_BROKER_EPOCH);
-    }
+    public abstract long brokerEpoch();
 
-    // Used for test
-    long size() {
-        return toStruct().sizeOf();
-    }
+    public abstract long maxBrokerEpoch();
 
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/AbstractRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/AbstractRequest.java
index 239024f86327c..97bc72852e39f 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/AbstractRequest.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/AbstractRequest.java
@@ -26,7 +26,7 @@
 import java.nio.ByteBuffer;
 import java.util.Map;
 
-public abstract class AbstractRequest extends AbstractRequestResponse {
+public abstract class AbstractRequest implements AbstractRequestResponse {
 
     public static abstract class Builder {
         private final ApiKeys apiKey;
@@ -76,11 +76,13 @@ public T build() {
     }
 
     private final short version;
+    public final ApiKeys api;
 
     public AbstractRequest(ApiKeys api, short version) {
         if (!api.isVersionSupported(version))
             throw new UnsupportedVersionException("The " + api + " protocol does not support version " + version);
         this.version = version;
+        this.api = api;
     }
 
     /**
@@ -98,7 +100,7 @@ public Send toSend(String destination, RequestHeader header) {
      * Use with care, typically {@link #toSend(String, RequestHeader)} should be used instead.
      */
     public ByteBuffer serialize(RequestHeader header) {
-        return serialize(header.toStruct(), toStruct());
+        return RequestUtils.serialize(header.toStruct(), toStruct());
     }
 
     protected abstract Struct toStruct();
@@ -227,8 +229,16 @@ public static AbstractRequest parseRequest(ApiKeys apiKey, short apiVersion, Str
                 return new DescribeDelegationTokenRequest(struct, apiVersion);
             case DELETE_GROUPS:
                 return new DeleteGroupsRequest(struct, apiVersion);
-            case ELECT_PREFERRED_LEADERS:
-                return new ElectPreferredLeadersRequest(struct, apiVersion);
+            case ELECT_LEADERS:
+                return new ElectLeadersRequest(struct, apiVersion);
+            case INCREMENTAL_ALTER_CONFIGS:
+                return new IncrementalAlterConfigsRequest(struct, apiVersion);
+            case ALTER_PARTITION_REASSIGNMENTS:
+                return new AlterPartitionReassignmentsRequest(struct, apiVersion);
+            case LIST_PARTITION_REASSIGNMENTS:
+                return new ListPartitionReassignmentsRequest(struct, apiVersion);
+            case OFFSET_DELETE:
+                return new OffsetDeleteRequest(struct, apiVersion);
             default:
                 throw new AssertionError(String.format("ApiKey %s is not currently handled in `parseRequest`, the " +
                         "code should be updated to do so.", apiKey));
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/AbstractRequestResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/AbstractRequestResponse.java
index 0ba373d6fea7a..b02659d4f354c 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/AbstractRequestResponse.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/AbstractRequestResponse.java
@@ -16,19 +16,5 @@
  */
 package org.apache.kafka.common.requests;
 
-import org.apache.kafka.common.protocol.types.Struct;
-
-import java.nio.ByteBuffer;
-
-public abstract class AbstractRequestResponse {
-    /**
-     * Visible for testing.
-     */
-    public static ByteBuffer serialize(Struct headerStruct, Struct bodyStruct) {
-        ByteBuffer buffer = ByteBuffer.allocate(headerStruct.sizeOf() + bodyStruct.sizeOf());
-        headerStruct.writeTo(buffer);
-        bodyStruct.writeTo(buffer);
-        buffer.rewind();
-        return buffer;
-    }
+public interface AbstractRequestResponse {
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/AbstractResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/AbstractResponse.java
index 712d732de77b7..64701529403de 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/AbstractResponse.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/AbstractResponse.java
@@ -23,22 +23,32 @@
 import org.apache.kafka.common.protocol.types.Struct;
 
 import java.nio.ByteBuffer;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
 
-public abstract class AbstractResponse extends AbstractRequestResponse {
+public abstract class AbstractResponse implements AbstractRequestResponse {
     public static final int DEFAULT_THROTTLE_TIME = 0;
 
     protected Send toSend(String destination, ResponseHeader header, short apiVersion) {
-        return new NetworkSend(destination, serialize(apiVersion, header));
+        return new NetworkSend(destination, RequestUtils.serialize(header.toStruct(), toStruct(apiVersion)));
     }
 
     /**
      * Visible for testing, typically {@link #toSend(String, ResponseHeader, short)} should be used instead.
      */
     public ByteBuffer serialize(short version, ResponseHeader responseHeader) {
-        return serialize(responseHeader.toStruct(), toStruct(version));
+        return RequestUtils.serialize(responseHeader.toStruct(), toStruct(version));
+    }
+
+    /**
+     * Visible for testing, typically {@link #toSend(String, ResponseHeader, short)} should be used instead.
+     */
+    public ByteBuffer serialize(ApiKeys apiKey, short version, int correlationId) {
+        ResponseHeader header =
+            new ResponseHeader(correlationId, apiKey.responseHeaderVersion(version));
+        return RequestUtils.serialize(header.toStruct(), toStruct(version));
     }
 
     public abstract Map errorCounts();
@@ -47,9 +57,9 @@ protected Map errorCounts(Errors error) {
         return Collections.singletonMap(error, 1);
     }
 
-    protected Map errorCounts(Map errors) {
+    protected Map errorCounts(Collection errors) {
         Map errorCounts = new HashMap<>();
-        for (Errors error : errors.values())
+        for (Errors error : errors)
             updateErrorCounts(errorCounts, error);
         return errorCounts;
     }
@@ -77,45 +87,45 @@ public static AbstractResponse parseResponse(ApiKeys apiKey, Struct struct, shor
             case LIST_OFFSETS:
                 return new ListOffsetResponse(struct);
             case METADATA:
-                return new MetadataResponse(struct);
+                return new MetadataResponse(struct, version);
             case OFFSET_COMMIT:
-                return new OffsetCommitResponse(struct);
+                return new OffsetCommitResponse(struct, version);
             case OFFSET_FETCH:
-                return new OffsetFetchResponse(struct);
+                return new OffsetFetchResponse(struct, version);
             case FIND_COORDINATOR:
-                return new FindCoordinatorResponse(struct);
+                return new FindCoordinatorResponse(struct, version);
             case JOIN_GROUP:
-                return new JoinGroupResponse(struct);
+                return new JoinGroupResponse(struct, version);
             case HEARTBEAT:
-                return new HeartbeatResponse(struct);
+                return new HeartbeatResponse(struct, version);
             case LEAVE_GROUP:
                 return new LeaveGroupResponse(struct, version);
             case SYNC_GROUP:
-                return new SyncGroupResponse(struct);
+                return new SyncGroupResponse(struct, version);
             case STOP_REPLICA:
-                return new StopReplicaResponse(struct);
+                return new StopReplicaResponse(struct, version);
             case CONTROLLED_SHUTDOWN:
-                return new ControlledShutdownResponse(struct);
+                return new ControlledShutdownResponse(struct, version);
             case UPDATE_METADATA:
-                return new UpdateMetadataResponse(struct);
+                return new UpdateMetadataResponse(struct, version);
             case LEADER_AND_ISR:
-                return new LeaderAndIsrResponse(struct);
+                return new LeaderAndIsrResponse(struct, version);
             case DESCRIBE_GROUPS:
                 return new DescribeGroupsResponse(struct, version);
             case LIST_GROUPS:
-                return new ListGroupsResponse(struct);
+                return new ListGroupsResponse(struct, version);
             case SASL_HANDSHAKE:
                 return new SaslHandshakeResponse(struct, version);
             case API_VERSIONS:
-                return new ApiVersionsResponse(struct);
+                return ApiVersionsResponse.fromStruct(struct, version);
             case CREATE_TOPICS:
                 return new CreateTopicsResponse(struct, version);
             case DELETE_TOPICS:
-                return new DeleteTopicsResponse(struct);
+                return new DeleteTopicsResponse(struct, version);
             case DELETE_RECORDS:
                 return new DeleteRecordsResponse(struct);
             case INIT_PRODUCER_ID:
-                return new InitProducerIdResponse(struct);
+                return new InitProducerIdResponse(struct, version);
             case OFFSET_FOR_LEADER_EPOCH:
                 return new OffsetsForLeaderEpochResponse(struct);
             case ADD_PARTITIONS_TO_TXN:
@@ -127,7 +137,7 @@ public static AbstractResponse parseResponse(ApiKeys apiKey, Struct struct, shor
             case WRITE_TXN_MARKERS:
                 return new WriteTxnMarkersResponse(struct);
             case TXN_OFFSET_COMMIT:
-                return new TxnOffsetCommitResponse(struct);
+                return new TxnOffsetCommitResponse(struct, version);
             case DESCRIBE_ACLS:
                 return new DescribeAclsResponse(struct);
             case CREATE_ACLS:
@@ -147,17 +157,25 @@ public static AbstractResponse parseResponse(ApiKeys apiKey, Struct struct, shor
             case CREATE_PARTITIONS:
                 return new CreatePartitionsResponse(struct);
             case CREATE_DELEGATION_TOKEN:
-                return new CreateDelegationTokenResponse(struct);
+                return new CreateDelegationTokenResponse(struct, version);
             case RENEW_DELEGATION_TOKEN:
-                return new RenewDelegationTokenResponse(struct);
+                return new RenewDelegationTokenResponse(struct, version);
             case EXPIRE_DELEGATION_TOKEN:
-                return new ExpireDelegationTokenResponse(struct);
+                return new ExpireDelegationTokenResponse(struct, version);
             case DESCRIBE_DELEGATION_TOKEN:
-                return new DescribeDelegationTokenResponse(struct);
+                return new DescribeDelegationTokenResponse(struct, version);
             case DELETE_GROUPS:
-                return new DeleteGroupsResponse(struct);
-            case ELECT_PREFERRED_LEADERS:
-                return new ElectPreferredLeadersResponse(struct, version);
+                return new DeleteGroupsResponse(struct, version);
+            case ELECT_LEADERS:
+                return new ElectLeadersResponse(struct, version);
+            case INCREMENTAL_ALTER_CONFIGS:
+                return new IncrementalAlterConfigsResponse(struct, version);
+            case ALTER_PARTITION_REASSIGNMENTS:
+                return new AlterPartitionReassignmentsResponse(struct, version);
+            case LIST_PARTITION_REASSIGNMENTS:
+                return new ListPartitionReassignmentsResponse(struct, version);
+            case OFFSET_DELETE:
+                return new OffsetDeleteResponse(struct, version);
             default:
                 throw new AssertionError(String.format("ApiKey %s is not currently handled in `parseResponse`, the " +
                         "code should be updated to do so.", apiKey));
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/AddPartitionsToTxnResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/AddPartitionsToTxnResponse.java
index ea8a073e01755..7f20f07b6a117 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/AddPartitionsToTxnResponse.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/AddPartitionsToTxnResponse.java
@@ -102,7 +102,7 @@ public Map errors() {
 
     @Override
     public Map errorCounts() {
-        return errorCounts(errors);
+        return errorCounts(errors.values());
     }
 
     @Override
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/AlterPartitionReassignmentsRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/AlterPartitionReassignmentsRequest.java
new file mode 100644
index 0000000000000..7b2f848f614f9
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/AlterPartitionReassignmentsRequest.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.common.requests;
+
+import org.apache.kafka.common.message.AlterPartitionReassignmentsRequestData;
+import org.apache.kafka.common.message.AlterPartitionReassignmentsRequestData.ReassignableTopic;
+import org.apache.kafka.common.message.AlterPartitionReassignmentsResponseData;
+import org.apache.kafka.common.message.AlterPartitionReassignmentsResponseData.ReassignablePartitionResponse;
+import org.apache.kafka.common.message.AlterPartitionReassignmentsResponseData.ReassignableTopicResponse;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.types.Struct;
+
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.stream.Collectors;
+
+public class AlterPartitionReassignmentsRequest extends AbstractRequest {
+
+    public static class Builder extends AbstractRequest.Builder {
+        private final AlterPartitionReassignmentsRequestData data;
+
+        public Builder(AlterPartitionReassignmentsRequestData data) {
+            super(ApiKeys.ALTER_PARTITION_REASSIGNMENTS);
+            this.data = data;
+        }
+
+        @Override
+        public AlterPartitionReassignmentsRequest build(short version) {
+            return new AlterPartitionReassignmentsRequest(data, version);
+        }
+
+        @Override
+        public String toString() {
+            return data.toString();
+        }
+    }
+
+    private final AlterPartitionReassignmentsRequestData data;
+    private final short version;
+
+    private AlterPartitionReassignmentsRequest(AlterPartitionReassignmentsRequestData data, short version) {
+        super(ApiKeys.ALTER_PARTITION_REASSIGNMENTS, version);
+        this.data = data;
+        this.version = version;
+    }
+
+    AlterPartitionReassignmentsRequest(Struct struct, short version) {
+        super(ApiKeys.ALTER_PARTITION_REASSIGNMENTS, version);
+        this.data = new AlterPartitionReassignmentsRequestData(struct, version);
+        this.version = version;
+    }
+
+    public static AlterPartitionReassignmentsRequest parse(ByteBuffer buffer, short version) {
+        return new AlterPartitionReassignmentsRequest(
+                ApiKeys.ALTER_PARTITION_REASSIGNMENTS.parseRequest(version, buffer),
+                version
+        );
+    }
+
+    public AlterPartitionReassignmentsRequestData data() {
+        return data;
+    }
+
+    /**
+     * Visible for testing.
+     */
+    @Override
+    public Struct toStruct() {
+        return data.toStruct(version);
+    }
+
+    @Override
+    public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
+        ApiError apiError = ApiError.fromThrowable(e);
+        List topicResponses = new ArrayList<>();
+
+        for (ReassignableTopic topic : data.topics()) {
+            List partitionResponses = topic.partitions().stream().map(partition ->
+                    new ReassignablePartitionResponse()
+                            .setPartitionIndex(partition.partitionIndex())
+                            .setErrorCode(apiError.error().code())
+                            .setErrorMessage(apiError.message())
+            ).collect(Collectors.toList());
+            topicResponses.add(
+                    new ReassignableTopicResponse()
+                            .setName(topic.name())
+                            .setPartitions(partitionResponses)
+            );
+        }
+
+        AlterPartitionReassignmentsResponseData responseData = new AlterPartitionReassignmentsResponseData()
+                .setResponses(topicResponses)
+                .setErrorCode(apiError.error().code())
+                .setErrorMessage(apiError.message())
+                .setThrottleTimeMs(throttleTimeMs);
+        return new AlterPartitionReassignmentsResponse(responseData);
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/ElectPreferredLeadersResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/AlterPartitionReassignmentsResponse.java
similarity index 50%
rename from clients/src/main/java/org/apache/kafka/common/requests/ElectPreferredLeadersResponse.java
rename to clients/src/main/java/org/apache/kafka/common/requests/AlterPartitionReassignmentsResponse.java
index c168c67c6b399..ef235cda34f78 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/ElectPreferredLeadersResponse.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/AlterPartitionReassignmentsResponse.java
@@ -17,9 +17,9 @@
 
 package org.apache.kafka.common.requests;
 
-import org.apache.kafka.common.message.ElectPreferredLeadersResponseData;
-import org.apache.kafka.common.message.ElectPreferredLeadersResponseData.PartitionResult;
-import org.apache.kafka.common.message.ElectPreferredLeadersResponseData.ReplicaElectionResult;
+import org.apache.kafka.common.message.AlterPartitionReassignmentsResponseData;
+import org.apache.kafka.common.message.AlterPartitionReassignmentsResponseData.ReassignableTopicResponse;
+import org.apache.kafka.common.message.AlterPartitionReassignmentsResponseData.ReassignablePartitionResponse;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
 import org.apache.kafka.common.protocol.types.Struct;
@@ -28,30 +28,33 @@
 import java.util.HashMap;
 import java.util.Map;
 
-public class ElectPreferredLeadersResponse extends AbstractResponse {
+public class AlterPartitionReassignmentsResponse extends AbstractResponse {
 
-    private final ElectPreferredLeadersResponseData data;
+    private final AlterPartitionReassignmentsResponseData data;
 
-    public ElectPreferredLeadersResponse(ElectPreferredLeadersResponseData data) {
+    public AlterPartitionReassignmentsResponse(Struct struct) {
+        this(struct, ApiKeys.ALTER_PARTITION_REASSIGNMENTS.latestVersion());
+    }
+
+    public AlterPartitionReassignmentsResponse(AlterPartitionReassignmentsResponseData data) {
         this.data = data;
     }
 
-    public ElectPreferredLeadersResponse(Struct struct, short version) {
-        this.data = new ElectPreferredLeadersResponseData(struct, version);
+    AlterPartitionReassignmentsResponse(Struct struct, short version) {
+        this.data = new AlterPartitionReassignmentsResponseData(struct, version);
     }
 
-    public ElectPreferredLeadersResponse(Struct struct) {
-        short latestVersion = (short) (ElectPreferredLeadersResponseData.SCHEMAS.length - 1);
-        this.data = new ElectPreferredLeadersResponseData(struct, latestVersion);
+    public static AlterPartitionReassignmentsResponse parse(ByteBuffer buffer, short version) {
+        return new AlterPartitionReassignmentsResponse(ApiKeys.ALTER_PARTITION_REASSIGNMENTS.responseSchema(version).read(buffer), version);
     }
 
-    public ElectPreferredLeadersResponseData data() {
+    public AlterPartitionReassignmentsResponseData data() {
         return data;
     }
 
     @Override
-    protected Struct toStruct(short version) {
-        return data.toStruct(version);
+    public boolean shouldClientThrottle(short version) {
+        return true;
     }
 
     @Override
@@ -61,23 +64,21 @@ public int throttleTimeMs() {
 
     @Override
     public Map errorCounts() {
-        HashMap counts = new HashMap<>();
-        for (ReplicaElectionResult result : data.replicaElectionResults()) {
-            for (PartitionResult partitionResult : result.partitionResult()) {
-                Errors error = Errors.forCode(partitionResult.errorCode());
+        Map counts = new HashMap<>();
+        Errors topLevelErr = Errors.forCode(data.errorCode());
+        counts.put(topLevelErr, counts.getOrDefault(topLevelErr, 0) + 1);
+
+        for (ReassignableTopicResponse topicResponse : data.responses()) {
+            for (ReassignablePartitionResponse partitionResponse : topicResponse.partitions()) {
+                Errors error = Errors.forCode(partitionResponse.errorCode());
                 counts.put(error, counts.getOrDefault(error, 0) + 1);
             }
         }
         return counts;
     }
 
-    public static ElectPreferredLeadersResponse parse(ByteBuffer buffer, short version) {
-        return new ElectPreferredLeadersResponse(
-                ApiKeys.ELECT_PREFERRED_LEADERS.responseSchema(version).read(buffer), version);
-    }
-
     @Override
-    public boolean shouldClientThrottle(short version) {
-        return version >= 3;
+    protected Struct toStruct(short version) {
+        return data.toStruct(version);
     }
-}
\ No newline at end of file
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/AlterReplicaLogDirsResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/AlterReplicaLogDirsResponse.java
index 7a73275d74577..2bbdce2122bc2 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/AlterReplicaLogDirsResponse.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/AlterReplicaLogDirsResponse.java
@@ -132,7 +132,7 @@ public Map responses() {
 
     @Override
     public Map errorCounts() {
-        return errorCounts(responses);
+        return errorCounts(responses.values());
     }
 
     public static AlterReplicaLogDirsResponse parse(ByteBuffer buffer, short version) {
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/ApiError.java b/clients/src/main/java/org/apache/kafka/common/requests/ApiError.java
index dad21b398d9a9..6cb09f078ecca 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/ApiError.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/ApiError.java
@@ -50,6 +50,10 @@ public ApiError(Struct struct) {
         message = struct.getOrElse(ERROR_MESSAGE, null);
     }
 
+    public ApiError(Errors error) {
+        this(error, error.message());
+    }
+
     public ApiError(Errors error, String message) {
         this.error = error;
         this.message = message;
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/ApiVersionsRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/ApiVersionsRequest.java
index 03f385b9ba586..87ffceaa26e1c 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/ApiVersionsRequest.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/ApiVersionsRequest.java
@@ -16,30 +16,26 @@
  */
 package org.apache.kafka.common.requests;
 
+import java.util.regex.Pattern;
+import org.apache.kafka.common.message.ApiVersionsRequestData;
+import org.apache.kafka.common.message.ApiVersionsResponseData;
+import org.apache.kafka.common.message.ApiVersionsResponseData.ApiVersionsResponseKey;
+import org.apache.kafka.common.message.ApiVersionsResponseData.ApiVersionsResponseKeyCollection;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.protocol.types.Schema;
 import org.apache.kafka.common.protocol.types.Struct;
+import org.apache.kafka.common.utils.AppInfoParser;
 
 import java.nio.ByteBuffer;
-import java.util.Collections;
 
 public class ApiVersionsRequest extends AbstractRequest {
-    private static final Schema API_VERSIONS_REQUEST_V0 = new Schema();
-
-    /* v1 request is the same as v0. Throttle time has been added to response */
-    private static final Schema API_VERSIONS_REQUEST_V1 = API_VERSIONS_REQUEST_V0;
-
-    /**
-     * The version number is bumped to indicate that on quota violation brokers send out responses before throttling.
-     */
-    private static final Schema API_VERSIONS_REQUEST_V2 = API_VERSIONS_REQUEST_V1;
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{API_VERSIONS_REQUEST_V0, API_VERSIONS_REQUEST_V1, API_VERSIONS_REQUEST_V2};
-    }
 
     public static class Builder extends AbstractRequest.Builder {
+        private static final String DEFAULT_CLIENT_SOFTWARE_NAME = "apache-kafka-java";
+
+        private static final ApiVersionsRequestData DATA = new ApiVersionsRequestData()
+            .setClientSoftwareName(DEFAULT_CLIENT_SOFTWARE_NAME)
+            .setClientSoftwareVersion(AppInfoParser.getVersion());
 
         public Builder() {
             super(ApiKeys.API_VERSIONS);
@@ -51,23 +47,28 @@ public Builder(short version) {
 
         @Override
         public ApiVersionsRequest build(short version) {
-            return new ApiVersionsRequest(version);
+            return new ApiVersionsRequest(DATA, version);
         }
 
         @Override
         public String toString() {
-            return "(type=ApiVersionsRequest)";
+            return DATA.toString();
         }
     }
 
+    private static final Pattern SOFTWARE_NAME_VERSION_PATTERN = Pattern.compile("[a-zA-Z0-9](?:[a-zA-Z0-9\\-.]*[a-zA-Z0-9])?");
+
     private final Short unsupportedRequestVersion;
 
-    public ApiVersionsRequest(short version) {
-        this(version, null);
+    public final ApiVersionsRequestData data;
+
+    public ApiVersionsRequest(ApiVersionsRequestData data, short version) {
+        this(data, version, null);
     }
 
-    public ApiVersionsRequest(short version, Short unsupportedRequestVersion) {
+    public ApiVersionsRequest(ApiVersionsRequestData data, short version, Short unsupportedRequestVersion) {
         super(ApiKeys.API_VERSIONS, version);
+        this.data = data;
 
         // Unlike other request types, the broker handles ApiVersion requests with higher versions than
         // supported. It does so by treating the request as if it were v0 and returns a response using
@@ -78,31 +79,48 @@ public ApiVersionsRequest(short version, Short unsupportedRequestVersion) {
     }
 
     public ApiVersionsRequest(Struct struct, short version) {
-        this(version, null);
+        this(new ApiVersionsRequestData(struct, version), version);
     }
 
     public boolean hasUnsupportedRequestVersion() {
         return unsupportedRequestVersion != null;
     }
 
+    public boolean isValid() {
+        if (version() >= 3) {
+            return SOFTWARE_NAME_VERSION_PATTERN.matcher(data.clientSoftwareName()).matches() &&
+                SOFTWARE_NAME_VERSION_PATTERN.matcher(data.clientSoftwareVersion()).matches();
+        } else {
+            return true;
+        }
+    }
+
     @Override
     protected Struct toStruct() {
-        return new Struct(ApiKeys.API_VERSIONS.requestSchema(version()));
+        return data.toStruct(version());
     }
 
     @Override
     public ApiVersionsResponse getErrorResponse(int throttleTimeMs, Throwable e) {
-        short version = version();
-        switch (version) {
-            case 0:
-                return new ApiVersionsResponse(Errors.forException(e), Collections.emptyList());
-            case 1:
-            case 2:
-                return new ApiVersionsResponse(throttleTimeMs, Errors.forException(e), Collections.emptyList());
-            default:
-                throw new IllegalArgumentException(String.format("Version %d is not valid. Valid versions for %s are 0 to %d",
-                        version, this.getClass().getSimpleName(), ApiKeys.API_VERSIONS.latestVersion()));
+        ApiVersionsResponseData data = new ApiVersionsResponseData()
+            .setErrorCode(Errors.forException(e).code());
+
+        if (version() >= 1) {
+            data.setThrottleTimeMs(throttleTimeMs);
         }
+
+        // Starting from Apache Kafka 2.4 (KIP-511), ApiKeys field is populated with the supported
+        // versions of the ApiVersionsRequest when an UNSUPPORTED_VERSION error is returned.
+        if (Errors.forException(e) == Errors.UNSUPPORTED_VERSION) {
+            ApiVersionsResponseKeyCollection apiKeys = new ApiVersionsResponseKeyCollection();
+            apiKeys.add(new ApiVersionsResponseKey()
+                .setApiKey(ApiKeys.API_VERSIONS.id)
+                .setMinVersion(ApiKeys.API_VERSIONS.oldestVersion())
+                .setMaxVersion(ApiKeys.API_VERSIONS.latestVersion()));
+            data.setApiKeys(apiKeys);
+        }
+
+        return new ApiVersionsResponse(data);
     }
 
     public static ApiVersionsRequest parse(ByteBuffer buffer, short version) {
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 a37b6da0b7273..4604fedd0b0e3 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
@@ -16,190 +16,126 @@
  */
 package org.apache.kafka.common.requests;
 
+import org.apache.kafka.common.message.ApiVersionsResponseData;
+import org.apache.kafka.common.message.ApiVersionsResponseData.ApiVersionsResponseKey;
+import org.apache.kafka.common.message.ApiVersionsResponseData.ApiVersionsResponseKeyCollection;
 import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
 import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.protocol.types.ArrayOf;
-import org.apache.kafka.common.protocol.types.Field;
-import org.apache.kafka.common.protocol.types.Schema;
+import org.apache.kafka.common.protocol.types.SchemaException;
 import org.apache.kafka.common.protocol.types.Struct;
 import org.apache.kafka.common.record.RecordBatch;
 
 import java.nio.ByteBuffer;
-import java.util.ArrayList;
-import java.util.Collection;
-import java.util.HashMap;
-import java.util.List;
 import java.util.Map;
 
-import static org.apache.kafka.common.protocol.CommonFields.ERROR_CODE;
-import static org.apache.kafka.common.protocol.CommonFields.THROTTLE_TIME_MS;
-import static org.apache.kafka.common.protocol.types.Type.INT16;
-
+/**
+ * Possible error codes:
+ * - {@link Errors#UNSUPPORTED_VERSION}
+ * - {@link Errors#INVALID_REQUEST}
+ */
 public class ApiVersionsResponse extends AbstractResponse {
-    private static final String API_VERSIONS_KEY_NAME = "api_versions";
-    private static final String API_KEY_NAME = "api_key";
-    private static final String MIN_VERSION_KEY_NAME = "min_version";
-    private static final String MAX_VERSION_KEY_NAME = "max_version";
-
-    private static final Schema API_VERSIONS_V0 = new Schema(
-            new Field(API_KEY_NAME, INT16, "API key."),
-            new Field(MIN_VERSION_KEY_NAME, INT16, "Minimum supported version."),
-            new Field(MAX_VERSION_KEY_NAME, INT16, "Maximum supported version."));
-
-    private static final Schema API_VERSIONS_RESPONSE_V0 = new Schema(
-            ERROR_CODE,
-            new Field(API_VERSIONS_KEY_NAME, new ArrayOf(API_VERSIONS_V0), "API versions supported by the broker."));
-    private static final Schema API_VERSIONS_RESPONSE_V1 = new Schema(
-            ERROR_CODE,
-            new Field(API_VERSIONS_KEY_NAME, new ArrayOf(API_VERSIONS_V0), "API versions supported by the broker."),
-            THROTTLE_TIME_MS);
-
-    /**
-     * The version number is bumped to indicate that on quota violation brokers send out responses before throttling.
-     */
-    private static final Schema API_VERSIONS_RESPONSE_V2 = API_VERSIONS_RESPONSE_V1;
-
-    // initialized lazily to avoid circular initialization dependence with ApiKeys
-    private static volatile ApiVersionsResponse defaultApiVersionsResponse;
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{API_VERSIONS_RESPONSE_V0, API_VERSIONS_RESPONSE_V1, API_VERSIONS_RESPONSE_V2};
-    }
-
-    /**
-     * Possible error codes:
-     *
-     * UNSUPPORTED_VERSION (33)
-     */
-    private final Errors error;
-    private final int throttleTimeMs;
-    private final Map apiKeyToApiVersion;
-
-    public static final class ApiVersion {
-        public final short apiKey;
-        public final short minVersion;
-        public final short maxVersion;
-
-        public ApiVersion(ApiKeys apiKey) {
-            this(apiKey.id, apiKey.oldestVersion(), apiKey.latestVersion());
-        }
 
-        public ApiVersion(short apiKey, short minVersion, short maxVersion) {
-            this.apiKey = apiKey;
-            this.minVersion = minVersion;
-            this.maxVersion = maxVersion;
-        }
+    public static final ApiVersionsResponse DEFAULT_API_VERSIONS_RESPONSE =
+        createApiVersionsResponse(DEFAULT_THROTTLE_TIME, RecordBatch.CURRENT_MAGIC_VALUE);
 
-        @Override
-        public String toString() {
-            return "ApiVersion(" +
-                    "apiKey=" + apiKey +
-                    ", minVersion=" + minVersion +
-                    ", maxVersion= " + maxVersion +
-                    ")";
-        }
-    }
+    public final ApiVersionsResponseData data;
 
-    public ApiVersionsResponse(Errors error, List apiVersions) {
-        this(DEFAULT_THROTTLE_TIME, error, apiVersions);
+    public ApiVersionsResponse(ApiVersionsResponseData data) {
+        this.data = data;
     }
 
-    public ApiVersionsResponse(int throttleTimeMs, Errors error, List apiVersions) {
-        this.throttleTimeMs = throttleTimeMs;
-        this.error = error;
-        this.apiKeyToApiVersion = buildApiKeyToApiVersion(apiVersions);
+    public ApiVersionsResponse(Struct struct) {
+        this(new ApiVersionsResponseData(struct, (short) (ApiVersionsResponseData.SCHEMAS.length - 1)));
     }
 
-    public ApiVersionsResponse(Struct struct) {
-        this.throttleTimeMs = struct.getOrElse(THROTTLE_TIME_MS, DEFAULT_THROTTLE_TIME);
-        this.error = Errors.forCode(struct.get(ERROR_CODE));
-        List tempApiVersions = new ArrayList<>();
-        for (Object apiVersionsObj : struct.getArray(API_VERSIONS_KEY_NAME)) {
-            Struct apiVersionStruct = (Struct) apiVersionsObj;
-            short apiKey = apiVersionStruct.getShort(API_KEY_NAME);
-            short minVersion = apiVersionStruct.getShort(MIN_VERSION_KEY_NAME);
-            short maxVersion = apiVersionStruct.getShort(MAX_VERSION_KEY_NAME);
-            tempApiVersions.add(new ApiVersion(apiKey, minVersion, maxVersion));
-        }
-        this.apiKeyToApiVersion = buildApiKeyToApiVersion(tempApiVersions);
+    public ApiVersionsResponse(Struct struct, short version) {
+        this(new ApiVersionsResponseData(struct, version));
     }
 
     @Override
     protected Struct toStruct(short version) {
-        Struct struct = new Struct(ApiKeys.API_VERSIONS.responseSchema(version));
-        struct.setIfExists(THROTTLE_TIME_MS, throttleTimeMs);
-        struct.set(ERROR_CODE, error.code());
-        List apiVersionList = new ArrayList<>();
-        for (ApiVersion apiVersion : apiKeyToApiVersion.values()) {
-            Struct apiVersionStruct = struct.instance(API_VERSIONS_KEY_NAME);
-            apiVersionStruct.set(API_KEY_NAME, apiVersion.apiKey);
-            apiVersionStruct.set(MIN_VERSION_KEY_NAME, apiVersion.minVersion);
-            apiVersionStruct.set(MAX_VERSION_KEY_NAME, apiVersion.maxVersion);
-            apiVersionList.add(apiVersionStruct);
-        }
-        struct.set(API_VERSIONS_KEY_NAME, apiVersionList.toArray());
-        return struct;
+        return this.data.toStruct(version);
     }
 
-    public static ApiVersionsResponse apiVersionsResponse(int throttleTimeMs, byte maxMagic) {
-        if (maxMagic == RecordBatch.CURRENT_MAGIC_VALUE && throttleTimeMs == DEFAULT_THROTTLE_TIME) {
-            return defaultApiVersionsResponse();
-        }
-        return createApiVersionsResponse(throttleTimeMs, maxMagic);
+    public ApiVersionsResponseKey apiVersion(short apiKey) {
+        return data.apiKeys().find(apiKey);
     }
 
     @Override
-    public int throttleTimeMs() {
-        return throttleTimeMs;
-    }
-
-    public Collection apiVersions() {
-        return apiKeyToApiVersion.values();
-    }
-
-    public ApiVersion apiVersion(short apiKey) {
-        return apiKeyToApiVersion.get(apiKey);
+    public Map errorCounts() {
+        return errorCounts(Errors.forCode(this.data.errorCode()));
     }
 
-    public Errors error() {
-        return error;
+    @Override
+    public int throttleTimeMs() {
+        return this.data.throttleTimeMs();
     }
 
     @Override
-    public Map errorCounts() {
-        return errorCounts(error);
+    public boolean shouldClientThrottle(short version) {
+        return version >= 2;
     }
 
     public static ApiVersionsResponse parse(ByteBuffer buffer, short version) {
-        return new ApiVersionsResponse(ApiKeys.API_VERSIONS.parseResponse(version, buffer));
+        // Fallback to version 0 for ApiVersions response. If a client sends an ApiVersionsRequest
+        // using a version higher than that supported by the broker, a version 0 response is sent
+        // to the client indicating UNSUPPORTED_VERSION. When the client receives the response, it
+        // falls back while parsing it into a Struct which means that the version received by this
+        // method is not necessary the real one. It may be version 0 as well.
+        int prev = buffer.position();
+        try {
+            return new ApiVersionsResponse(
+                new ApiVersionsResponseData(new ByteBufferAccessor(buffer), version));
+        } catch (RuntimeException e) {
+            buffer.position(prev);
+            if (version != 0)
+                return new ApiVersionsResponse(
+                    new ApiVersionsResponseData(new ByteBufferAccessor(buffer), (short) 0));
+            else
+                throw e;
+        }
     }
 
-    private Map buildApiKeyToApiVersion(List apiVersions) {
-        Map tempApiIdToApiVersion = new HashMap<>();
-        for (ApiVersion apiVersion : apiVersions) {
-            tempApiIdToApiVersion.put(apiVersion.apiKey, apiVersion);
+    public static ApiVersionsResponse fromStruct(Struct struct, short version) {
+        // Fallback to version 0 for ApiVersions response. If a client sends an ApiVersionsRequest
+        // using a version higher than that supported by the broker, a version 0 response is sent
+        // to the client indicating UNSUPPORTED_VERSION. When the client receives the response, it
+        // falls back while parsing it into a Struct which means that the version received by this
+        // method is not necessary the real one. It may be version 0 as well.
+        try {
+            return new ApiVersionsResponse(struct, version);
+        } catch (SchemaException e) {
+            if (version != 0)
+                return new ApiVersionsResponse(struct, (short) 0);
+            else
+                throw e;
         }
-        return tempApiIdToApiVersion;
+    }
+
+    public static ApiVersionsResponse apiVersionsResponse(int throttleTimeMs, byte maxMagic) {
+        if (maxMagic == RecordBatch.CURRENT_MAGIC_VALUE && throttleTimeMs == DEFAULT_THROTTLE_TIME) {
+            return DEFAULT_API_VERSIONS_RESPONSE;
+        }
+        return createApiVersionsResponse(throttleTimeMs, maxMagic);
     }
 
     public static ApiVersionsResponse createApiVersionsResponse(int throttleTimeMs, final byte minMagic) {
-        List versionList = new ArrayList<>();
+        ApiVersionsResponseKeyCollection apiKeys = new ApiVersionsResponseKeyCollection();
         for (ApiKeys apiKey : ApiKeys.values()) {
             if (apiKey.minRequiredInterBrokerMagic <= minMagic) {
-                versionList.add(new ApiVersionsResponse.ApiVersion(apiKey));
+                apiKeys.add(new ApiVersionsResponseKey()
+                    .setApiKey(apiKey.id)
+                    .setMinVersion(apiKey.oldestVersion())
+                    .setMaxVersion(apiKey.latestVersion()));
             }
         }
-        return new ApiVersionsResponse(throttleTimeMs, Errors.NONE, versionList);
-    }
 
-    public static ApiVersionsResponse defaultApiVersionsResponse() {
-        if (defaultApiVersionsResponse == null)
-            defaultApiVersionsResponse = createApiVersionsResponse(DEFAULT_THROTTLE_TIME, RecordBatch.CURRENT_MAGIC_VALUE);
-        return defaultApiVersionsResponse;
-    }
+        ApiVersionsResponseData data = new ApiVersionsResponseData();
+        data.setThrottleTimeMs(throttleTimeMs);
+        data.setErrorCode(Errors.NONE.code());
+        data.setApiKeys(apiKeys);
 
-    @Override
-    public boolean shouldClientThrottle(short version) {
-        return version >= 2;
+        return new ApiVersionsResponse(data);
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/ControlledShutdownRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/ControlledShutdownRequest.java
index ea8ff5f6da41d..71238452cf782 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/ControlledShutdownRequest.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/ControlledShutdownRequest.java
@@ -16,95 +16,56 @@
  */
 package org.apache.kafka.common.requests;
 
+import org.apache.kafka.common.message.ControlledShutdownRequestData;
+import org.apache.kafka.common.message.ControlledShutdownResponseData;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.protocol.types.Field;
-import org.apache.kafka.common.protocol.types.Schema;
 import org.apache.kafka.common.protocol.types.Struct;
 
 import java.nio.ByteBuffer;
-import java.util.Collections;
-
-import static org.apache.kafka.common.protocol.types.Type.INT32;
-import static org.apache.kafka.common.protocol.types.Type.INT64;
 
 public class ControlledShutdownRequest extends AbstractRequest {
-    private static final String BROKER_ID_KEY_NAME = "broker_id";
-    private static final String BROKER_EPOCH_KEY_NAME = "broker_epoch";
-
-    private static final Schema CONTROLLED_SHUTDOWN_REQUEST_V0 = new Schema(
-            new Field(BROKER_ID_KEY_NAME, INT32, "The id of the broker for which controlled shutdown has been requested."));
-    private static final Schema CONTROLLED_SHUTDOWN_REQUEST_V1 = CONTROLLED_SHUTDOWN_REQUEST_V0;
-    // Introduce broker_epoch to allow controller to reject stale ControlledShutdownRequest
-    private static final Schema CONTROLLED_SHUTDOWN_REQUEST_V2 = new Schema(
-            new Field(BROKER_ID_KEY_NAME, INT32, "The id of the broker for which controlled shutdown has been requested."),
-            new Field(BROKER_EPOCH_KEY_NAME, INT64, "The broker epoch"));
-
-    public static Schema[] schemaVersions() {
-        return new Schema[] {CONTROLLED_SHUTDOWN_REQUEST_V0, CONTROLLED_SHUTDOWN_REQUEST_V1, CONTROLLED_SHUTDOWN_REQUEST_V2};
-    }
 
     public static class Builder extends AbstractRequest.Builder {
-        private final int brokerId;
-        private final long brokerEpoch;
 
-        public Builder(int brokerId, long brokerEpoch, short desiredVersion) {
+        private final ControlledShutdownRequestData data;
+
+        public Builder(ControlledShutdownRequestData data, short desiredVersion) {
             super(ApiKeys.CONTROLLED_SHUTDOWN, desiredVersion);
-            this.brokerId = brokerId;
-            this.brokerEpoch = brokerEpoch;
+            this.data = data;
         }
 
         @Override
         public ControlledShutdownRequest build(short version) {
-            return new ControlledShutdownRequest(brokerId, brokerEpoch, version);
+            return new ControlledShutdownRequest(data, version);
         }
 
         @Override
         public String toString() {
-            StringBuilder bld = new StringBuilder();
-            bld.append("(type=ControlledShutdownRequest").
-                append(", brokerId=").append(brokerId).
-                append(", brokerEpoch=").append(brokerEpoch).
-                append(")");
-            return bld.toString();
+            return data.toString();
         }
     }
-    private final int brokerId;
-    private final long brokerEpoch;
 
-    private ControlledShutdownRequest(int brokerId, long brokerEpoch, short version) {
+    private final ControlledShutdownRequestData data;
+    private final short version;
+
+    private ControlledShutdownRequest(ControlledShutdownRequestData data, short version) {
         super(ApiKeys.CONTROLLED_SHUTDOWN, version);
-        this.brokerId = brokerId;
-        this.brokerEpoch = brokerEpoch;
+        this.data = data;
+        this.version = version;
     }
 
     public ControlledShutdownRequest(Struct struct, short version) {
         super(ApiKeys.CONTROLLED_SHUTDOWN, version);
-        brokerId = struct.getInt(BROKER_ID_KEY_NAME);
-        brokerEpoch = struct.hasField(BROKER_EPOCH_KEY_NAME) ? struct.getLong(BROKER_EPOCH_KEY_NAME) :
-                AbstractControlRequest.UNKNOWN_BROKER_EPOCH;
+        this.data = new ControlledShutdownRequestData(struct, version);
+        this.version = version;
     }
 
     @Override
-    public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
-        short versionId = version();
-        switch (versionId) {
-            case 0:
-            case 1:
-            case 2:
-                return new ControlledShutdownResponse(Errors.forException(e), Collections.emptySet());
-            default:
-                throw new IllegalArgumentException(String.format("Version %d is not valid. Valid versions for %s are 0 to %d",
-                    versionId, this.getClass().getSimpleName(), ApiKeys.CONTROLLED_SHUTDOWN.latestVersion()));
-        }
-    }
-
-    public int brokerId() {
-        return brokerId;
-    }
-
-    public long brokerEpoch() {
-        return brokerEpoch;
+    public ControlledShutdownResponse getErrorResponse(int throttleTimeMs, Throwable e) {
+        ControlledShutdownResponseData data = new ControlledShutdownResponseData()
+                .setErrorCode(Errors.forException(e).code());
+        return new ControlledShutdownResponse(data);
     }
 
     public static ControlledShutdownRequest parse(ByteBuffer buffer, short version) {
@@ -114,9 +75,10 @@ public static ControlledShutdownRequest parse(ByteBuffer buffer, short version)
 
     @Override
     protected Struct toStruct() {
-        Struct struct = new Struct(ApiKeys.CONTROLLED_SHUTDOWN.requestSchema(version()));
-        struct.set(BROKER_ID_KEY_NAME, brokerId);
-        struct.setIfExists(BROKER_EPOCH_KEY_NAME, brokerEpoch);
-        return struct;
+        return data.toStruct(version);
+    }
+
+    public ControlledShutdownRequestData data() {
+        return data;
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/ControlledShutdownResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/ControlledShutdownResponse.java
index 1d62b2d2954ba..a49b84a851f41 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/ControlledShutdownResponse.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/ControlledShutdownResponse.java
@@ -17,44 +17,20 @@
 package org.apache.kafka.common.requests;
 
 import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.message.ControlledShutdownResponseData;
+import org.apache.kafka.common.message.ControlledShutdownResponseData.RemainingPartition;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.protocol.types.ArrayOf;
-import org.apache.kafka.common.protocol.types.Field;
-import org.apache.kafka.common.protocol.types.Schema;
 import org.apache.kafka.common.protocol.types.Struct;
 
 import java.nio.ByteBuffer;
-import java.util.ArrayList;
-import java.util.HashSet;
-import java.util.List;
+import java.util.Collections;
 import java.util.Map;
 import java.util.Set;
 
-import static org.apache.kafka.common.protocol.CommonFields.ERROR_CODE;
-import static org.apache.kafka.common.protocol.CommonFields.PARTITION_ID;
-import static org.apache.kafka.common.protocol.CommonFields.TOPIC_NAME;
 
 public class ControlledShutdownResponse extends AbstractResponse {
 
-    private static final String PARTITIONS_REMAINING_KEY_NAME = "partitions_remaining";
-
-    private static final Schema CONTROLLED_SHUTDOWN_PARTITION_V0 = new Schema(
-            TOPIC_NAME,
-            PARTITION_ID);
-
-    private static final Schema CONTROLLED_SHUTDOWN_RESPONSE_V0 = new Schema(
-            ERROR_CODE,
-            new Field(PARTITIONS_REMAINING_KEY_NAME, new ArrayOf(CONTROLLED_SHUTDOWN_PARTITION_V0), "The partitions " +
-                    "that the broker still leads."));
-
-    private static final Schema CONTROLLED_SHUTDOWN_RESPONSE_V1 = CONTROLLED_SHUTDOWN_RESPONSE_V0;
-    private static final Schema CONTROLLED_SHUTDOWN_RESPONSE_V2 = CONTROLLED_SHUTDOWN_RESPONSE_V1;
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{CONTROLLED_SHUTDOWN_RESPONSE_V0, CONTROLLED_SHUTDOWN_RESPONSE_V1, CONTROLLED_SHUTDOWN_RESPONSE_V2};
-    }
-
     /**
      * Possible error codes:
      *
@@ -62,58 +38,49 @@ public static Schema[] schemaVersions() {
      * BROKER_NOT_AVAILABLE(8)
      * STALE_CONTROLLER_EPOCH(11)
      */
-    private final Errors error;
+    private final ControlledShutdownResponseData data;
 
-    private final Set partitionsRemaining;
-
-    public ControlledShutdownResponse(Errors error, Set partitionsRemaining) {
-        this.error = error;
-        this.partitionsRemaining = partitionsRemaining;
+    public ControlledShutdownResponse(ControlledShutdownResponseData data) {
+        this.data = data;
     }
 
-    public ControlledShutdownResponse(Struct struct) {
-        error = Errors.forCode(struct.get(ERROR_CODE));
-        Set partitions = new HashSet<>();
-        for (Object topicPartitionObj : struct.getArray(PARTITIONS_REMAINING_KEY_NAME)) {
-            Struct topicPartition = (Struct) topicPartitionObj;
-            String topic = topicPartition.get(TOPIC_NAME);
-            int partition = topicPartition.get(PARTITION_ID);
-            partitions.add(new TopicPartition(topic, partition));
-        }
-        partitionsRemaining = partitions;
+    public ControlledShutdownResponse(Struct struct, short version) {
+        this.data = new ControlledShutdownResponseData(struct, version);
     }
 
     public Errors error() {
-        return error;
+        return Errors.forCode(data.errorCode());
     }
 
     @Override
     public Map errorCounts() {
-        return errorCounts(error);
-    }
-
-    public Set partitionsRemaining() {
-        return partitionsRemaining;
+        return Collections.singletonMap(error(), 1);
     }
 
     public static ControlledShutdownResponse parse(ByteBuffer buffer, short version) {
-        return new ControlledShutdownResponse(ApiKeys.CONTROLLED_SHUTDOWN.parseResponse(version, buffer));
+        return new ControlledShutdownResponse(ApiKeys.CONTROLLED_SHUTDOWN.parseResponse(version, buffer), version);
     }
 
     @Override
     protected Struct toStruct(short version) {
-        Struct struct = new Struct(ApiKeys.CONTROLLED_SHUTDOWN.responseSchema(version));
-        struct.set(ERROR_CODE, error.code());
+        return data.toStruct(version);
+    }
 
-        List partitionsRemainingList = new ArrayList<>(partitionsRemaining.size());
-        for (TopicPartition topicPartition : partitionsRemaining) {
-            Struct topicPartitionStruct = struct.instance(PARTITIONS_REMAINING_KEY_NAME);
-            topicPartitionStruct.set(TOPIC_NAME, topicPartition.topic());
-            topicPartitionStruct.set(PARTITION_ID, topicPartition.partition());
-            partitionsRemainingList.add(topicPartitionStruct);
-        }
-        struct.set(PARTITIONS_REMAINING_KEY_NAME, partitionsRemainingList.toArray());
+    public ControlledShutdownResponseData data() {
+        return data;
+    }
 
-        return struct;
+    public static ControlledShutdownResponse prepareResponse(Errors error, Set tps) {
+        ControlledShutdownResponseData data = new ControlledShutdownResponseData();
+        data.setErrorCode(error.code());
+        ControlledShutdownResponseData.RemainingPartitionCollection pSet = new ControlledShutdownResponseData.RemainingPartitionCollection();
+        tps.forEach(tp -> {
+            pSet.add(new RemainingPartition()
+                    .setTopicName(tp.topic())
+                    .setPartitionIndex(tp.partition()));
+        });
+        data.setRemainingPartitions(pSet);
+        return new ControlledShutdownResponse(data);
     }
+
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/CreateDelegationTokenRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/CreateDelegationTokenRequest.java
index 3277f100382e8..61b404a7e0623 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/CreateDelegationTokenRequest.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/CreateDelegationTokenRequest.java
@@ -16,126 +16,62 @@
  */
 package org.apache.kafka.common.requests;
 
+import org.apache.kafka.common.message.CreateDelegationTokenRequestData;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.protocol.types.ArrayOf;
-import org.apache.kafka.common.protocol.types.Field;
-import org.apache.kafka.common.protocol.types.Schema;
 import org.apache.kafka.common.protocol.types.Struct;
 import org.apache.kafka.common.security.auth.KafkaPrincipal;
 
 import java.nio.ByteBuffer;
-import java.util.ArrayList;
-import java.util.List;
-
-import static org.apache.kafka.common.protocol.types.Type.INT64;
-import static org.apache.kafka.common.protocol.CommonFields.PRINCIPAL_TYPE;
-import static org.apache.kafka.common.protocol.CommonFields.PRINCIPAL_NAME;
 
 public class CreateDelegationTokenRequest extends AbstractRequest {
-    private static final String RENEWERS_KEY_NAME = "renewers";
-    private static final String MAX_LIFE_TIME_KEY_NAME = "max_life_time";
-
-    private static final Schema TOKEN_CREATE_REQUEST_V0 = new Schema(
-            new Field(RENEWERS_KEY_NAME, new ArrayOf(new Schema(PRINCIPAL_TYPE, PRINCIPAL_NAME)),
-                    "An array of token renewers. Renewer is an Kafka PrincipalType and name string," +
-                        " who is allowed to renew this token before the max lifetime expires."),
-            new Field(MAX_LIFE_TIME_KEY_NAME, INT64,
-                    "Max lifetime period for token in milli seconds. if value is -1, then max lifetime" +
-                        "  will default to a server side config value."));
 
-    /**
-     * The version number is bumped to indicate that on quota violation brokers send out responses before throttling.
-     */
-    private static final Schema TOKEN_CREATE_REQUEST_V1 = TOKEN_CREATE_REQUEST_V0;
+    private final CreateDelegationTokenRequestData data;
 
-    private final List renewers;
-    private final long maxLifeTime;
-
-    private CreateDelegationTokenRequest(short version, List renewers, long maxLifeTime) {
+    private CreateDelegationTokenRequest(CreateDelegationTokenRequestData data, short version) {
         super(ApiKeys.CREATE_DELEGATION_TOKEN, version);
-        this.maxLifeTime = maxLifeTime;
-        this.renewers = renewers;
+        this.data = data;
     }
 
     public CreateDelegationTokenRequest(Struct struct, short version) {
         super(ApiKeys.CREATE_DELEGATION_TOKEN, version);
-        maxLifeTime = struct.getLong(MAX_LIFE_TIME_KEY_NAME);
-        Object[] renewerArray = struct.getArray(RENEWERS_KEY_NAME);
-        renewers = new ArrayList<>();
-        if (renewerArray != null) {
-            for (Object renewerObj : renewerArray) {
-                Struct renewerObjStruct = (Struct) renewerObj;
-                String principalType = renewerObjStruct.get(PRINCIPAL_TYPE);
-                String principalName = renewerObjStruct.get(PRINCIPAL_NAME);
-                renewers.add(new KafkaPrincipal(principalType, principalName));
-            }
-        }
+        this.data = new CreateDelegationTokenRequestData(struct, version);
     }
 
     public static CreateDelegationTokenRequest parse(ByteBuffer buffer, short version) {
         return new CreateDelegationTokenRequest(ApiKeys.CREATE_DELEGATION_TOKEN.parseRequest(version, buffer), version);
     }
 
-    public static Schema[] schemaVersions() {
-        return new Schema[]{TOKEN_CREATE_REQUEST_V0, TOKEN_CREATE_REQUEST_V1};
-    }
-
     @Override
     protected Struct toStruct() {
-        short version = version();
-        Struct struct = new Struct(ApiKeys.CREATE_DELEGATION_TOKEN.requestSchema(version));
-        Object[] renewersArray = new Object[renewers.size()];
-
-        int i = 0;
-        for (KafkaPrincipal principal: renewers) {
-            Struct renewerStruct = struct.instance(RENEWERS_KEY_NAME);
-            renewerStruct.set(PRINCIPAL_TYPE, principal.getPrincipalType());
-            renewerStruct.set(PRINCIPAL_NAME, principal.getName());
-            renewersArray[i++] = renewerStruct;
-        }
+        return data.toStruct(version());
+    }
 
-        struct.set(RENEWERS_KEY_NAME, renewersArray);
-        struct.set(MAX_LIFE_TIME_KEY_NAME, maxLifeTime);
-        return struct;
+    public CreateDelegationTokenRequestData data() {
+        return data;
     }
 
     @Override
     public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
-        return new CreateDelegationTokenResponse(throttleTimeMs, Errors.forException(e), KafkaPrincipal.ANONYMOUS);
-    }
-
-    public List renewers() {
-        return renewers;
-    }
-
-    public long maxLifeTime() {
-        return maxLifeTime;
+        return CreateDelegationTokenResponse.prepareResponse(throttleTimeMs, Errors.forException(e), KafkaPrincipal.ANONYMOUS);
     }
 
     public static class Builder extends AbstractRequest.Builder {
-        private final List renewers;
-        private final long maxLifeTime;
+        private final CreateDelegationTokenRequestData data;
 
-        public Builder(List renewers, long maxLifeTime) {
+        public Builder(CreateDelegationTokenRequestData data) {
             super(ApiKeys.CREATE_DELEGATION_TOKEN);
-            this.renewers = renewers;
-            this.maxLifeTime = maxLifeTime;
+            this.data = data;
         }
 
         @Override
         public CreateDelegationTokenRequest build(short version) {
-            return new CreateDelegationTokenRequest(version, renewers, maxLifeTime);
+            return new CreateDelegationTokenRequest(data, version);
         }
 
         @Override
         public String toString() {
-            StringBuilder bld = new StringBuilder();
-            bld.append("(type: CreateDelegationTokenRequest").
-                append(", renewers=").append(renewers).
-                append(", maxLifeTime=").append(maxLifeTime).
-                append(")");
-            return bld.toString();
+            return data.toString();
         }
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/CreateDelegationTokenResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/CreateDelegationTokenResponse.java
index 8bb6631140b54..04db17ac7a2f3 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/CreateDelegationTokenResponse.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/CreateDelegationTokenResponse.java
@@ -16,160 +16,82 @@
  */
 package org.apache.kafka.common.requests;
 
+import org.apache.kafka.common.message.CreateDelegationTokenResponseData;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.protocol.types.Field;
-import org.apache.kafka.common.protocol.types.Schema;
 import org.apache.kafka.common.protocol.types.Struct;
 import org.apache.kafka.common.security.auth.KafkaPrincipal;
 
 import java.nio.ByteBuffer;
+import java.util.Collections;
 import java.util.Map;
 
-import static org.apache.kafka.common.protocol.CommonFields.ERROR_CODE;
-import static org.apache.kafka.common.protocol.CommonFields.PRINCIPAL_NAME;
-import static org.apache.kafka.common.protocol.CommonFields.PRINCIPAL_TYPE;
-import static org.apache.kafka.common.protocol.CommonFields.THROTTLE_TIME_MS;
-import static org.apache.kafka.common.protocol.types.Type.BYTES;
-import static org.apache.kafka.common.protocol.types.Type.INT64;
-import static org.apache.kafka.common.protocol.types.Type.STRING;
-
 public class CreateDelegationTokenResponse extends AbstractResponse {
 
-    private static final String OWNER_KEY_NAME = "owner";
-    private static final String ISSUE_TIMESTAMP_KEY_NAME = "issue_timestamp";
-    private static final String EXPIRY_TIMESTAMP_NAME = "expiry_timestamp";
-    private static final String MAX_TIMESTAMP_NAME = "max_timestamp";
-    private static final String TOKEN_ID_KEY_NAME = "token_id";
-    private static final String HMAC_KEY_NAME = "hmac";
-
-    private final Errors error;
-    private final long issueTimestamp;
-    private final long expiryTimestamp;
-    private final long maxTimestamp;
-    private final String tokenId;
-    private final ByteBuffer hmac;
-    private final int throttleTimeMs;
-    private KafkaPrincipal owner;
-
-    private static final Schema TOKEN_CREATE_RESPONSE_V0 = new Schema(
-        ERROR_CODE,
-        new Field(OWNER_KEY_NAME, new Schema(PRINCIPAL_TYPE, PRINCIPAL_NAME), "token owner."),
-        new Field(ISSUE_TIMESTAMP_KEY_NAME, INT64, "timestamp (in msec) when this token was generated."),
-        new Field(EXPIRY_TIMESTAMP_NAME, INT64, "timestamp (in msec) at which this token expires."),
-        new Field(MAX_TIMESTAMP_NAME, INT64, "max life time of this token."),
-        new Field(TOKEN_ID_KEY_NAME, STRING, "UUID to ensure uniqueness."),
-        new Field(HMAC_KEY_NAME, BYTES, "HMAC of the delegation token."),
-        THROTTLE_TIME_MS);
-
-    /**
-     * The version number is bumped to indicate that on quota violation brokers send out responses before throttling.
-     */
-    private static final Schema TOKEN_CREATE_RESPONSE_V1 = TOKEN_CREATE_RESPONSE_V0;
-
-    public CreateDelegationTokenResponse(int throttleTimeMs,
-                                         Errors error,
-                                         KafkaPrincipal owner,
-                                         long issueTimestamp,
-                                         long expiryTimestamp,
-                                         long maxTimestamp,
-                                         String tokenId,
-                                         ByteBuffer hmac) {
-        this.throttleTimeMs = throttleTimeMs;
-        this.error = error;
-        this.owner = owner;
-        this.issueTimestamp = issueTimestamp;
-        this.expiryTimestamp = expiryTimestamp;
-        this.maxTimestamp = maxTimestamp;
-        this.tokenId = tokenId;
-        this.hmac = hmac;
-    }
+    private final CreateDelegationTokenResponseData data;
 
-    public CreateDelegationTokenResponse(int throttleTimeMs, Errors error, KafkaPrincipal owner) {
-        this(throttleTimeMs, error, owner, -1, -1, -1, "", ByteBuffer.wrap(new byte[] {}));
+    public CreateDelegationTokenResponse(CreateDelegationTokenResponseData data) {
+        this.data = data;
     }
 
-    public CreateDelegationTokenResponse(Struct struct) {
-        error = Errors.forCode(struct.get(ERROR_CODE));
-        Struct ownerStruct = (Struct) struct.get(OWNER_KEY_NAME);
-        String principalType = ownerStruct.get(PRINCIPAL_TYPE);
-        String principalName = ownerStruct.get(PRINCIPAL_NAME);
-        owner = new KafkaPrincipal(principalType, principalName);
-        issueTimestamp = struct.getLong(ISSUE_TIMESTAMP_KEY_NAME);
-        expiryTimestamp = struct.getLong(EXPIRY_TIMESTAMP_NAME);
-        maxTimestamp = struct.getLong(MAX_TIMESTAMP_NAME);
-        tokenId = struct.getString(TOKEN_ID_KEY_NAME);
-        hmac = struct.getBytes(HMAC_KEY_NAME);
-        this.throttleTimeMs = struct.getOrElse(THROTTLE_TIME_MS, DEFAULT_THROTTLE_TIME);
+    public CreateDelegationTokenResponse(Struct struct, short version) {
+        this.data = new CreateDelegationTokenResponseData(struct, version);
     }
 
     public static CreateDelegationTokenResponse parse(ByteBuffer buffer, short version) {
-        return new CreateDelegationTokenResponse(ApiKeys.CREATE_DELEGATION_TOKEN.responseSchema(version).read(buffer));
-    }
-
-    public static Schema[] schemaVersions() {
-        return new Schema[] {TOKEN_CREATE_RESPONSE_V0, TOKEN_CREATE_RESPONSE_V1};
-    }
-
-    @Override
-    public Map errorCounts() {
-        return errorCounts(error);
-    }
-
-    @Override
-    protected Struct toStruct(short version) {
-        Struct struct = new Struct(ApiKeys.CREATE_DELEGATION_TOKEN.responseSchema(version));
-        struct.set(ERROR_CODE, error.code());
-        Struct ownerStruct = struct.instance(OWNER_KEY_NAME);
-        ownerStruct.set(PRINCIPAL_TYPE, owner.getPrincipalType());
-        ownerStruct.set(PRINCIPAL_NAME, owner.getName());
-        struct.set(OWNER_KEY_NAME, ownerStruct);
-        struct.set(ISSUE_TIMESTAMP_KEY_NAME, issueTimestamp);
-        struct.set(EXPIRY_TIMESTAMP_NAME, expiryTimestamp);
-        struct.set(MAX_TIMESTAMP_NAME, maxTimestamp);
-        struct.set(TOKEN_ID_KEY_NAME, tokenId);
-        struct.set(HMAC_KEY_NAME, hmac);
-        struct.setIfExists(THROTTLE_TIME_MS, throttleTimeMs);
-        return struct;
+        return new CreateDelegationTokenResponse(ApiKeys.CREATE_DELEGATION_TOKEN.responseSchema(version).read(buffer), version);
     }
 
-    public Errors error() {
-        return error;
-    }
-
-    public KafkaPrincipal owner() {
-        return owner;
+    public static CreateDelegationTokenResponse prepareResponse(int throttleTimeMs,
+            Errors error,
+            KafkaPrincipal owner,
+            long issueTimestamp,
+            long expiryTimestamp,
+            long maxTimestamp,
+            String tokenId,
+            ByteBuffer hmac) {
+        CreateDelegationTokenResponseData data = new CreateDelegationTokenResponseData()
+                .setThrottleTimeMs(throttleTimeMs)
+                .setErrorCode(error.code())
+                .setPrincipalType(owner.getPrincipalType())
+                .setPrincipalName(owner.getName())
+                .setIssueTimestampMs(issueTimestamp)
+                .setExpiryTimestampMs(expiryTimestamp)
+                .setMaxTimestampMs(maxTimestamp)
+                .setTokenId(tokenId)
+                .setHmac(hmac.array());
+        return new CreateDelegationTokenResponse(data);
     }
 
-    public long issueTimestamp() {
-        return issueTimestamp;
+    public static CreateDelegationTokenResponse prepareResponse(int throttleTimeMs, Errors error, KafkaPrincipal owner) {
+        return prepareResponse(throttleTimeMs, error, owner, -1, -1, -1, "", ByteBuffer.wrap(new byte[] {}));
     }
 
-    public long expiryTimestamp() {
-        return expiryTimestamp;
+    public CreateDelegationTokenResponseData data() {
+        return data;
     }
 
-    public long maxTimestamp() {
-        return maxTimestamp;
-    }
-
-    public String tokenId() {
-        return tokenId;
+    @Override
+    public Map errorCounts() {
+        return Collections.singletonMap(error(), 1);
     }
 
-    public byte[] hmacBytes() {
-        byte[] byteArray = new byte[hmac.remaining()];
-        hmac.get(byteArray);
-        return byteArray;
+    @Override
+    protected Struct toStruct(short version) {
+        return data.toStruct(version);
     }
 
     @Override
     public int throttleTimeMs() {
-        return throttleTimeMs;
+        return data.throttleTimeMs();
+    }
+
+    public Errors error() {
+        return Errors.forCode(data.errorCode());
     }
 
     public boolean hasError() {
-        return this.error != Errors.NONE;
+        return error() != Errors.NONE;
     }
 
     @Override
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/CreateTopicsRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/CreateTopicsRequest.java
index 93f7ab269a4ba..dd26e5642632e 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/CreateTopicsRequest.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/CreateTopicsRequest.java
@@ -16,15 +16,16 @@
  */
 package org.apache.kafka.common.requests;
 
+import java.nio.ByteBuffer;
+import java.util.List;
+import java.util.stream.Collectors;
 import org.apache.kafka.common.errors.UnsupportedVersionException;
-import org.apache.kafka.common.protocol.ApiKeys;
-import org.apache.kafka.common.protocol.types.Struct;
 import org.apache.kafka.common.message.CreateTopicsRequestData;
 import org.apache.kafka.common.message.CreateTopicsRequestData.CreatableTopic;
 import org.apache.kafka.common.message.CreateTopicsResponseData;
 import org.apache.kafka.common.message.CreateTopicsResponseData.CreatableTopicResult;
-
-import java.nio.ByteBuffer;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.types.Struct;
 
 public class CreateTopicsRequest extends AbstractRequest {
     public static class Builder extends AbstractRequest.Builder {
@@ -40,6 +41,23 @@ public CreateTopicsRequest build(short version) {
             if (data.validateOnly() && version == 0)
                 throw new UnsupportedVersionException("validateOnly is not supported in version 0 of " +
                         "CreateTopicsRequest");
+
+            final List topicsWithDefaults = data.topics()
+                .stream()
+                .filter(topic -> topic.assignments().isEmpty())
+                .filter(topic ->
+                    topic.numPartitions() == CreateTopicsRequest.NO_NUM_PARTITIONS
+                        || topic.replicationFactor() == CreateTopicsRequest.NO_REPLICATION_FACTOR)
+                .map(CreatableTopic::name)
+                .collect(Collectors.toList());
+
+            if (!topicsWithDefaults.isEmpty() && version < 4) {
+                throw new UnsupportedVersionException("Creating topics with default "
+                    + "partitions/replication factor are only supported in CreateTopicRequest "
+                    + "version 4+. The following topics need values for partitions and replicas: "
+                    + topicsWithDefaults);
+            }
+
             return new CreateTopicsRequest(data, version);
         }
 
@@ -50,7 +68,6 @@ public String toString() {
     }
 
     private final CreateTopicsRequestData data;
-    private final short version;
 
     public static final int NO_NUM_PARTITIONS = -1;
     public static final short NO_REPLICATION_FACTOR = -1;
@@ -58,13 +75,11 @@ public String toString() {
     private CreateTopicsRequest(CreateTopicsRequestData data, short version) {
         super(ApiKeys.CREATE_TOPICS, version);
         this.data = data;
-        this.version = version;
     }
 
     public CreateTopicsRequest(Struct struct, short version) {
         super(ApiKeys.CREATE_TOPICS, version);
         this.data = new CreateTopicsRequestData(struct, version);
-        this.version = version;
     }
 
     public CreateTopicsRequestData data() {
@@ -96,6 +111,6 @@ public static CreateTopicsRequest parse(ByteBuffer buffer, short version) {
      */
     @Override
     public Struct toStruct() {
-        return data.toStruct(version);
+        return data.toStruct(version());
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/DeleteGroupsRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/DeleteGroupsRequest.java
index f2a5d92e9882d..2ad947c5904be 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/DeleteGroupsRequest.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/DeleteGroupsRequest.java
@@ -16,108 +16,71 @@
  */
 package org.apache.kafka.common.requests;
 
+import org.apache.kafka.common.message.DeleteGroupsRequestData;
+import org.apache.kafka.common.message.DeleteGroupsResponseData;
+import org.apache.kafka.common.message.DeleteGroupsResponseData.DeletableGroupResult;
+import org.apache.kafka.common.message.DeleteGroupsResponseData.DeletableGroupResultCollection;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.protocol.types.ArrayOf;
-import org.apache.kafka.common.protocol.types.Field;
-import org.apache.kafka.common.protocol.types.Schema;
 import org.apache.kafka.common.protocol.types.Struct;
-import org.apache.kafka.common.utils.Utils;
 
 import java.nio.ByteBuffer;
-import java.util.HashMap;
-import java.util.HashSet;
-import java.util.Map;
-import java.util.Set;
-
-import static org.apache.kafka.common.protocol.types.Type.STRING;
 
 public class DeleteGroupsRequest extends AbstractRequest {
-    private static final String GROUPS_KEY_NAME = "groups";
-
-    /* DeleteGroups api */
-    private static final Schema DELETE_GROUPS_REQUEST_V0 = new Schema(
-            new Field(GROUPS_KEY_NAME, new ArrayOf(STRING), "An array of groups to be deleted."));
-
-    /**
-     * The version number is bumped to indicate that on quota violation brokers send out responses before throttling.
-     */
-    private static final Schema DELETE_GROUPS_REQUEST_V1 = DELETE_GROUPS_REQUEST_V0;
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{DELETE_GROUPS_REQUEST_V0, DELETE_GROUPS_REQUEST_V1};
-    }
-
-    private final Set groups;
-
     public static class Builder extends AbstractRequest.Builder {
-        private final Set groups;
+        private final DeleteGroupsRequestData data;
 
-        public Builder(Set groups) {
+        public Builder(DeleteGroupsRequestData data) {
             super(ApiKeys.DELETE_GROUPS);
-            this.groups = groups;
+            this.data = data;
         }
 
         @Override
         public DeleteGroupsRequest build(short version) {
-            return new DeleteGroupsRequest(groups, version);
+            return new DeleteGroupsRequest(data, version);
         }
 
         @Override
         public String toString() {
-            StringBuilder bld = new StringBuilder();
-            bld.append("(type=DeleteGroupsRequest").
-                append(", groups=(").append(Utils.join(groups, ", ")).append(")").
-                append(")");
-            return bld.toString();
+            return data.toString();
         }
     }
 
-    private DeleteGroupsRequest(Set groups, short version) {
+    public final DeleteGroupsRequestData data;
+
+    public DeleteGroupsRequest(DeleteGroupsRequestData data, short version) {
         super(ApiKeys.DELETE_GROUPS, version);
-        this.groups = groups;
+        this.data = data;
     }
 
     public DeleteGroupsRequest(Struct struct, short version) {
         super(ApiKeys.DELETE_GROUPS, version);
-        Object[] groupsArray = struct.getArray(GROUPS_KEY_NAME);
-        Set groups = new HashSet<>(groupsArray.length);
-        for (Object group : groupsArray)
-            groups.add((String) group);
-
-        this.groups = groups;
-    }
-
-    @Override
-    protected Struct toStruct() {
-        Struct struct = new Struct(ApiKeys.DELETE_GROUPS.requestSchema(version()));
-        struct.set(GROUPS_KEY_NAME, groups.toArray());
-        return struct;
+        this.data = new DeleteGroupsRequestData(struct, version);
     }
 
     @Override
     public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
         Errors error = Errors.forException(e);
-        Map groupErrors = new HashMap<>(groups.size());
-        for (String group : groups)
-            groupErrors.put(group, error);
-
-        switch (version()) {
-            case 0:
-            case 1:
-                return new DeleteGroupsResponse(throttleTimeMs, groupErrors);
-            default:
-                throw new IllegalArgumentException(String.format("Version %d is not valid. Valid versions for %s are 0 to %d",
-                    version(), ApiKeys.DELETE_GROUPS.name, ApiKeys.DELETE_GROUPS.latestVersion()));
+        DeletableGroupResultCollection groupResults = new DeletableGroupResultCollection();
+        for (String groupId : data.groupsNames()) {
+            groupResults.add(new DeletableGroupResult()
+                                 .setGroupId(groupId)
+                                 .setErrorCode(error.code()));
         }
-    }
 
-    public Set groups() {
-        return groups;
+        return new DeleteGroupsResponse(
+            new DeleteGroupsResponseData()
+                .setResults(groupResults)
+                .setThrottleTimeMs(throttleTimeMs)
+        );
     }
 
     public static DeleteGroupsRequest parse(ByteBuffer buffer, short version) {
         return new DeleteGroupsRequest(ApiKeys.DELETE_GROUPS.parseRequest(version, buffer), version);
     }
 
+    @Override
+    protected Struct toStruct() {
+        return data.toStruct(version());
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/DeleteGroupsResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/DeleteGroupsResponse.java
index 032f15de37ae2..a98a7c881d96e 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/DeleteGroupsResponse.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/DeleteGroupsResponse.java
@@ -16,121 +16,82 @@
  */
 package org.apache.kafka.common.requests;
 
+import org.apache.kafka.common.message.DeleteGroupsResponseData;
+import org.apache.kafka.common.message.DeleteGroupsResponseData.DeletableGroupResult;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.protocol.types.ArrayOf;
-import org.apache.kafka.common.protocol.types.Field;
-import org.apache.kafka.common.protocol.types.Schema;
 import org.apache.kafka.common.protocol.types.Struct;
 
 import java.nio.ByteBuffer;
-import java.util.ArrayList;
 import java.util.HashMap;
-import java.util.List;
 import java.util.Map;
 
-import static org.apache.kafka.common.protocol.CommonFields.ERROR_CODE;
-import static org.apache.kafka.common.protocol.CommonFields.GROUP_ID;
-import static org.apache.kafka.common.protocol.CommonFields.THROTTLE_TIME_MS;
-
+/**
+ * Possible error codes:
+ *
+ * COORDINATOR_LOAD_IN_PROGRESS (14)
+ * COORDINATOR_NOT_AVAILABLE(15)
+ * NOT_COORDINATOR (16)
+ * INVALID_GROUP_ID(24)
+ * GROUP_AUTHORIZATION_FAILED(30)
+ * NON_EMPTY_GROUP(68)
+ * GROUP_ID_NOT_FOUND(69)
+ */
 public class DeleteGroupsResponse extends AbstractResponse {
-    private static final String GROUP_ERROR_CODES_KEY_NAME = "group_error_codes";
-
-    private static final Schema GROUP_ERROR_CODE = new Schema(
-            GROUP_ID,
-            ERROR_CODE);
-
-    private static final Schema DELETE_GROUPS_RESPONSE_V0 = new Schema(
-            THROTTLE_TIME_MS,
-            new Field(GROUP_ERROR_CODES_KEY_NAME, new ArrayOf(GROUP_ERROR_CODE), "An array of per group error codes."));
-
-    /**
-     * The version number is bumped to indicate that on quota violation brokers send out responses before throttling.
-     */
-    private static final Schema DELETE_GROUPS_RESPONSE_V1 = DELETE_GROUPS_RESPONSE_V0;
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{DELETE_GROUPS_RESPONSE_V0, DELETE_GROUPS_RESPONSE_V1};
-    }
-
 
-    /**
-     * Possible error codes:
-     *
-     * COORDINATOR_LOAD_IN_PROGRESS (14)
-     * COORDINATOR_NOT_AVAILABLE(15)
-     * NOT_COORDINATOR (16)
-     * INVALID_GROUP_ID(24)
-     * GROUP_AUTHORIZATION_FAILED(30)
-     * NON_EMPTY_GROUP(68)
-     * GROUP_ID_NOT_FOUND(69)
-     */
+    public final DeleteGroupsResponseData data;
 
-    private final Map errors;
-    private final int throttleTimeMs;
-
-    public DeleteGroupsResponse(Map errors) {
-        this(DEFAULT_THROTTLE_TIME, errors);
-    }
-
-    public DeleteGroupsResponse(int throttleTimeMs, Map errors) {
-        this.throttleTimeMs = throttleTimeMs;
-        this.errors = errors;
+    public DeleteGroupsResponse(DeleteGroupsResponseData data) {
+        this.data = data;
     }
 
     public DeleteGroupsResponse(Struct struct) {
-        this.throttleTimeMs = struct.getOrElse(THROTTLE_TIME_MS, DEFAULT_THROTTLE_TIME);
-        Object[] groupErrorCodesStructs = struct.getArray(GROUP_ERROR_CODES_KEY_NAME);
-        Map errors = new HashMap<>();
-        for (Object groupErrorCodeStructObj : groupErrorCodesStructs) {
-            Struct groupErrorCodeStruct = (Struct) groupErrorCodeStructObj;
-            String group = groupErrorCodeStruct.get(GROUP_ID);
-            Errors error = Errors.forCode(groupErrorCodeStruct.get(ERROR_CODE));
-            errors.put(group, error);
-        }
-
-        this.errors = errors;
+        short latestVersion = (short) (DeleteGroupsResponseData.SCHEMAS.length - 1);
+        this.data = new DeleteGroupsResponseData(struct, latestVersion);
     }
 
-    @Override
-    protected Struct toStruct(short version) {
-        Struct struct = new Struct(ApiKeys.DELETE_GROUPS.responseSchema(version));
-        struct.setIfExists(THROTTLE_TIME_MS, throttleTimeMs);
-        List groupErrorCodeStructs = new ArrayList<>(errors.size());
-        for (Map.Entry groupError : errors.entrySet()) {
-            Struct groupErrorCodeStruct = struct.instance(GROUP_ERROR_CODES_KEY_NAME);
-            groupErrorCodeStruct.set(GROUP_ID, groupError.getKey());
-            groupErrorCodeStruct.set(ERROR_CODE, groupError.getValue().code());
-            groupErrorCodeStructs.add(groupErrorCodeStruct);
-        }
-        struct.set(GROUP_ERROR_CODES_KEY_NAME, groupErrorCodeStructs.toArray());
-        return struct;
+    public DeleteGroupsResponse(Struct struct, short version) {
+        this.data = new DeleteGroupsResponseData(struct, version);
     }
 
     @Override
-    public int throttleTimeMs() {
-        return throttleTimeMs;
+    protected Struct toStruct(short version) {
+        return data.toStruct(version);
     }
 
     public Map errors() {
-        return errors;
-    }
-
-    public boolean hasError(String group) {
-        return errors.containsKey(group) && errors.get(group) != Errors.NONE;
+        Map errorMap = new HashMap<>();
+        for (DeletableGroupResult result : data.results()) {
+            errorMap.put(result.groupId(), Errors.forCode(result.errorCode()));
+        }
+        return errorMap;
     }
 
-    public Errors get(String group) {
-        return errors.get(group);
+    public Errors get(String group) throws IllegalArgumentException {
+        DeletableGroupResult result = data.results().find(group);
+        if (result == null) {
+            throw new IllegalArgumentException("could not find group " + group + " in the delete group response");
+        }
+        return Errors.forCode(result.errorCode());
     }
 
     @Override
     public Map errorCounts() {
-        return errorCounts(errors);
+        Map counts = new HashMap<>();
+        for (DeletableGroupResult result : data.results()) {
+            Errors error = Errors.forCode(result.errorCode());
+            counts.put(error, counts.getOrDefault(error, 0) + 1);
+        }
+        return counts;
     }
 
     public static DeleteGroupsResponse parse(ByteBuffer buffer, short version) {
-        return new DeleteGroupsResponse(ApiKeys.DELETE_GROUPS.responseSchema(version).read(buffer));
+        return new DeleteGroupsResponse(ApiKeys.DELETE_GROUPS.parseResponse(version, buffer), version);
+    }
+
+    @Override
+    public int throttleTimeMs() {
+        return data.throttleTimeMs();
     }
 
     @Override
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/DeleteTopicsRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/DeleteTopicsRequest.java
index facb55e69f2a7..440acfd2b838a 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/DeleteTopicsRequest.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/DeleteTopicsRequest.java
@@ -16,131 +16,72 @@
  */
 package org.apache.kafka.common.requests;
 
+import org.apache.kafka.common.message.DeleteTopicsRequestData;
+import org.apache.kafka.common.message.DeleteTopicsResponseData;
+import org.apache.kafka.common.message.DeleteTopicsResponseData.DeletableTopicResult;
 import org.apache.kafka.common.protocol.ApiKeys;
-import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.protocol.types.ArrayOf;
-import org.apache.kafka.common.protocol.types.Field;
-import org.apache.kafka.common.protocol.types.Schema;
 import org.apache.kafka.common.protocol.types.Struct;
-import org.apache.kafka.common.utils.Utils;
 
 import java.nio.ByteBuffer;
-import java.util.HashMap;
-import java.util.HashSet;
-import java.util.Map;
-import java.util.Set;
-
-import static org.apache.kafka.common.protocol.types.Type.INT32;
-import static org.apache.kafka.common.protocol.types.Type.STRING;
 
 public class DeleteTopicsRequest extends AbstractRequest {
-    private static final String TOPICS_KEY_NAME = "topics";
-    private static final String TIMEOUT_KEY_NAME = "timeout";
-
-    /* DeleteTopic api */
-    private static final Schema DELETE_TOPICS_REQUEST_V0 = new Schema(
-            new Field(TOPICS_KEY_NAME, new ArrayOf(STRING), "An array of topics to be deleted."),
-            new Field(TIMEOUT_KEY_NAME, INT32, "The time in ms to wait for a topic to be completely deleted on the " +
-                    "controller node. Values <= 0 will trigger topic deletion and return immediately"));
-
-    /* v1 request is the same as v0. Throttle time has been added to the response */
-    private static final Schema DELETE_TOPICS_REQUEST_V1 = DELETE_TOPICS_REQUEST_V0;
-
-    /**
-     * The version number is bumped to indicate that on quota violation brokers send out responses before throttling.
-     */
-    private static final Schema DELETE_TOPICS_REQUEST_V2 = DELETE_TOPICS_REQUEST_V1;
 
-    /**
-     * v3 request is the same that as v2. The response is different based on the request version.
-     * In v3 version a TopicDeletionDisabledException is returned
-     */
-    private static final Schema DELETE_TOPICS_REQUEST_V3 = DELETE_TOPICS_REQUEST_V2;
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{DELETE_TOPICS_REQUEST_V0, DELETE_TOPICS_REQUEST_V1,
-            DELETE_TOPICS_REQUEST_V2, DELETE_TOPICS_REQUEST_V3};
-    }
-
-    private final Set topics;
-    private final Integer timeout;
+    private DeleteTopicsRequestData data;
+    private final short version;
 
     public static class Builder extends AbstractRequest.Builder {
-        private final Set topics;
-        private final Integer timeout;
+        private DeleteTopicsRequestData data;
 
-        public Builder(Set topics, Integer timeout) {
+        public Builder(DeleteTopicsRequestData data) {
             super(ApiKeys.DELETE_TOPICS);
-            this.topics = topics;
-            this.timeout = timeout;
+            this.data = data;
         }
 
         @Override
         public DeleteTopicsRequest build(short version) {
-            return new DeleteTopicsRequest(topics, timeout, version);
+            return new DeleteTopicsRequest(data, version);
         }
 
         @Override
         public String toString() {
-            StringBuilder bld = new StringBuilder();
-            bld.append("(type=DeleteTopicsRequest").
-                append(", topics=(").append(Utils.join(topics, ", ")).append(")").
-                append(", timeout=").append(timeout).
-                append(")");
-            return bld.toString();
+            return data.toString();
         }
     }
 
-    private DeleteTopicsRequest(Set topics, Integer timeout, short version) {
+    private DeleteTopicsRequest(DeleteTopicsRequestData data, short version) {
         super(ApiKeys.DELETE_TOPICS, version);
-        this.topics = topics;
-        this.timeout = timeout;
+        this.data = data;
+        this.version = version;
     }
 
     public DeleteTopicsRequest(Struct struct, short version) {
         super(ApiKeys.DELETE_TOPICS, version);
-        Object[] topicsArray = struct.getArray(TOPICS_KEY_NAME);
-        Set topics = new HashSet<>(topicsArray.length);
-        for (Object topic : topicsArray)
-            topics.add((String) topic);
-
-        this.topics = topics;
-        this.timeout = struct.getInt(TIMEOUT_KEY_NAME);
+        this.data = new DeleteTopicsRequestData(struct, version);
+        this.version = version;
     }
 
     @Override
     protected Struct toStruct() {
-        Struct struct = new Struct(ApiKeys.DELETE_TOPICS.requestSchema(version()));
-        struct.set(TOPICS_KEY_NAME, topics.toArray());
-        struct.set(TIMEOUT_KEY_NAME, timeout);
-        return struct;
+        return data.toStruct(version);
+    }
+
+    public DeleteTopicsRequestData data() {
+        return data;
     }
 
     @Override
     public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
-        Map topicErrors = new HashMap<>();
-        for (String topic : topics)
-            topicErrors.put(topic, Errors.forException(e));
-
-        switch (version()) {
-            case 0:
-                return new DeleteTopicsResponse(topicErrors);
-            case 1:
-            case 2:
-            case 3:
-                return new DeleteTopicsResponse(throttleTimeMs, topicErrors);
-            default:
-                throw new IllegalArgumentException(String.format("Version %d is not valid. Valid versions for %s are 0 to %d",
-                    version(), this.getClass().getSimpleName(), ApiKeys.DELETE_TOPICS.latestVersion()));
+        DeleteTopicsResponseData response = new DeleteTopicsResponseData();
+        if (version >= 1) {
+            response.setThrottleTimeMs(throttleTimeMs);
         }
-    }
-
-    public Set topics() {
-        return topics;
-    }
-
-    public Integer timeout() {
-        return this.timeout;
+        ApiError apiError = ApiError.fromThrowable(e);
+        for (String topic : data.topicNames()) {
+            response.responses().add(new DeletableTopicResult()
+                    .setName(topic)
+                    .setErrorCode(apiError.error().code()));
+        }
+        return new DeleteTopicsResponse(response);
     }
 
     public static DeleteTopicsRequest parse(ByteBuffer buffer, short version) {
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/DeleteTopicsResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/DeleteTopicsResponse.java
index 650caa829ab38..aa8e552d206df 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/DeleteTopicsResponse.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/DeleteTopicsResponse.java
@@ -16,54 +16,18 @@
  */
 package org.apache.kafka.common.requests;
 
+import org.apache.kafka.common.message.DeleteTopicsResponseData;
+import org.apache.kafka.common.message.DeleteTopicsResponseData.DeletableTopicResult;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.protocol.types.ArrayOf;
-import org.apache.kafka.common.protocol.types.Field;
-import org.apache.kafka.common.protocol.types.Schema;
 import org.apache.kafka.common.protocol.types.Struct;
 
 import java.nio.ByteBuffer;
-import java.util.ArrayList;
 import java.util.HashMap;
-import java.util.List;
 import java.util.Map;
 
-import static org.apache.kafka.common.protocol.CommonFields.ERROR_CODE;
-import static org.apache.kafka.common.protocol.CommonFields.THROTTLE_TIME_MS;
-import static org.apache.kafka.common.protocol.CommonFields.TOPIC_NAME;
 
 public class DeleteTopicsResponse extends AbstractResponse {
-    private static final String TOPIC_ERROR_CODES_KEY_NAME = "topic_error_codes";
-
-    private static final Schema TOPIC_ERROR_CODE = new Schema(
-            TOPIC_NAME,
-            ERROR_CODE);
-
-    private static final Schema DELETE_TOPICS_RESPONSE_V0 = new Schema(
-            new Field(TOPIC_ERROR_CODES_KEY_NAME,
-                    new ArrayOf(TOPIC_ERROR_CODE), "An array of per topic error codes."));
-
-    private static final Schema DELETE_TOPICS_RESPONSE_V1 = new Schema(
-            THROTTLE_TIME_MS,
-            new Field(TOPIC_ERROR_CODES_KEY_NAME, new ArrayOf(TOPIC_ERROR_CODE), "An array of per topic error codes."));
-
-    /**
-     * The version number is bumped to indicate that on quota violation brokers send out responses before throttling.
-     */
-    private static final Schema DELETE_TOPICS_RESPONSE_V2 = DELETE_TOPICS_RESPONSE_V1;
-
-    /**
-     * v3 request is the same that as v2. The response is different based on the request version.
-     * In v3 version a TopicDeletionDisabledException is returned
-     */
-    private static final Schema DELETE_TOPICS_RESPONSE_V3 = DELETE_TOPICS_RESPONSE_V2;
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{DELETE_TOPICS_RESPONSE_V0, DELETE_TOPICS_RESPONSE_V1,
-            DELETE_TOPICS_RESPONSE_V2, DELETE_TOPICS_RESPONSE_V3};
-    }
-
 
     /**
      * Possible error codes:
@@ -75,63 +39,42 @@ public static Schema[] schemaVersions() {
      * INVALID_REQUEST(42)
      * TOPIC_DELETION_DISABLED(73)
      */
-    private final Map errors;
-    private final int throttleTimeMs;
+    private DeleteTopicsResponseData data;
 
-    public DeleteTopicsResponse(Map errors) {
-        this(DEFAULT_THROTTLE_TIME, errors);
+    public DeleteTopicsResponse(DeleteTopicsResponseData data) {
+        this.data = data;
     }
 
-    public DeleteTopicsResponse(int throttleTimeMs, Map errors) {
-        this.throttleTimeMs = throttleTimeMs;
-        this.errors = errors;
-    }
-
-    public DeleteTopicsResponse(Struct struct) {
-        this.throttleTimeMs = struct.getOrElse(THROTTLE_TIME_MS, DEFAULT_THROTTLE_TIME);
-        Object[] topicErrorCodesStructs = struct.getArray(TOPIC_ERROR_CODES_KEY_NAME);
-        Map errors = new HashMap<>();
-        for (Object topicErrorCodeStructObj : topicErrorCodesStructs) {
-            Struct topicErrorCodeStruct = (Struct) topicErrorCodeStructObj;
-            String topic = topicErrorCodeStruct.get(TOPIC_NAME);
-            Errors error = Errors.forCode(topicErrorCodeStruct.get(ERROR_CODE));
-            errors.put(topic, error);
-        }
-
-        this.errors = errors;
+    public DeleteTopicsResponse(Struct struct, short version) {
+        this.data = new DeleteTopicsResponseData(struct, version);
     }
 
     @Override
     protected Struct toStruct(short version) {
-        Struct struct = new Struct(ApiKeys.DELETE_TOPICS.responseSchema(version));
-        struct.setIfExists(THROTTLE_TIME_MS, throttleTimeMs);
-        List topicErrorCodeStructs = new ArrayList<>(errors.size());
-        for (Map.Entry topicError : errors.entrySet()) {
-            Struct topicErrorCodeStruct = struct.instance(TOPIC_ERROR_CODES_KEY_NAME);
-            topicErrorCodeStruct.set(TOPIC_NAME, topicError.getKey());
-            topicErrorCodeStruct.set(ERROR_CODE, topicError.getValue().code());
-            topicErrorCodeStructs.add(topicErrorCodeStruct);
-        }
-        struct.set(TOPIC_ERROR_CODES_KEY_NAME, topicErrorCodeStructs.toArray());
-        return struct;
+        return data.toStruct(version);
     }
 
     @Override
     public int throttleTimeMs() {
-        return throttleTimeMs;
+        return data.throttleTimeMs();
     }
 
-    public Map errors() {
-        return errors;
+    public DeleteTopicsResponseData data() {
+        return data;
     }
 
     @Override
     public Map errorCounts() {
-        return errorCounts(errors);
+        HashMap counts = new HashMap<>();
+        for (DeletableTopicResult result : data.responses()) {
+            Errors error = Errors.forCode(result.errorCode());
+            counts.put(error, counts.getOrDefault(error, 0) + 1);
+        }
+        return counts;
     }
 
     public static DeleteTopicsResponse parse(ByteBuffer buffer, short version) {
-        return new DeleteTopicsResponse(ApiKeys.DELETE_TOPICS.responseSchema(version).read(buffer));
+        return new DeleteTopicsResponse(ApiKeys.DELETE_TOPICS.parseResponse(version, buffer), version);
     }
 
     @Override
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/DescribeAclsRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/DescribeAclsRequest.java
index 04bfee8d74057..ed417c577309c 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/DescribeAclsRequest.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/DescribeAclsRequest.java
@@ -86,14 +86,14 @@ public String toString() {
     private final AclBindingFilter filter;
 
     DescribeAclsRequest(AclBindingFilter filter, short version) {
-        super(ApiKeys.DELETE_ACLS, version);
+        super(ApiKeys.DESCRIBE_ACLS, version);
         this.filter = filter;
 
         validate(filter, version);
     }
 
     public DescribeAclsRequest(Struct struct, short version) {
-        super(ApiKeys.DELETE_ACLS, version);
+        super(ApiKeys.DESCRIBE_ACLS, version);
         ResourcePatternFilter resourceFilter = RequestUtils.resourcePatternFilterFromStructFields(struct);
         AccessControlEntryFilter entryFilter = RequestUtils.aceFilterFromStructFields(struct);
         this.filter = new AclBindingFilter(resourceFilter, entryFilter);
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/DescribeConfigsResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/DescribeConfigsResponse.java
index 51c35d56f22dc..6d424f2f3e378 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/DescribeConfigsResponse.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/DescribeConfigsResponse.java
@@ -179,7 +179,8 @@ public enum ConfigSource {
         DYNAMIC_BROKER_CONFIG((byte) 2),
         DYNAMIC_DEFAULT_BROKER_CONFIG((byte) 3),
         STATIC_BROKER_CONFIG((byte) 4),
-        DEFAULT_CONFIG((byte) 5);
+        DEFAULT_CONFIG((byte) 5),
+        DYNAMIC_BROKER_LOGGER_CONFIG((byte) 6);
 
         final byte id;
         private static final ConfigSource[] VALUES = values();
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/DescribeDelegationTokenRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/DescribeDelegationTokenRequest.java
index 21e14608c3691..7db3209d621fb 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/DescribeDelegationTokenRequest.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/DescribeDelegationTokenRequest.java
@@ -16,122 +16,69 @@
  */
 package org.apache.kafka.common.requests;
 
+import org.apache.kafka.common.message.DescribeDelegationTokenRequestData;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.protocol.types.ArrayOf;
-import org.apache.kafka.common.protocol.types.Field;
-import org.apache.kafka.common.protocol.types.Schema;
 import org.apache.kafka.common.protocol.types.Struct;
 import org.apache.kafka.common.security.auth.KafkaPrincipal;
 
-import java.nio.ByteBuffer;
-import java.util.ArrayList;
 import java.util.List;
-
-import static org.apache.kafka.common.protocol.CommonFields.PRINCIPAL_NAME;
-import static org.apache.kafka.common.protocol.CommonFields.PRINCIPAL_TYPE;
+import java.util.stream.Collectors;
 
 public class DescribeDelegationTokenRequest extends AbstractRequest {
-    private static final String OWNER_KEY_NAME = "owners";
-
-    private final List owners;
-
-    public static final Schema TOKEN_DESCRIBE_REQUEST_V0 = new Schema(
-        new Field(OWNER_KEY_NAME, ArrayOf.nullable(new Schema(PRINCIPAL_TYPE, PRINCIPAL_NAME)), "An array of token owners."));
 
-    /**
-     * The version number is bumped to indicate that on quota violation brokers send out responses before throttling.
-     */
-    public static final Schema TOKEN_DESCRIBE_REQUEST_V1 = TOKEN_DESCRIBE_REQUEST_V0;
+    private final DescribeDelegationTokenRequestData data;
 
     public static class Builder extends AbstractRequest.Builder {
-        // describe token for the given list of owners, or null if we want to describe all tokens.
-        private final List owners;
+        private final DescribeDelegationTokenRequestData data;
 
         public Builder(List owners) {
             super(ApiKeys.DESCRIBE_DELEGATION_TOKEN);
-            this.owners = owners;
+            this.data = new DescribeDelegationTokenRequestData()
+                .setOwners(owners == null ? null : owners
+                    .stream()
+                    .map(owner -> new DescribeDelegationTokenRequestData.DescribeDelegationTokenOwner()
+                            .setPrincipalName(owner.getName())
+                            .setPrincipalType(owner.getPrincipalType()))
+                    .collect(Collectors.toList()));
         }
 
         @Override
         public DescribeDelegationTokenRequest build(short version) {
-            return new DescribeDelegationTokenRequest(version, owners);
+            return new DescribeDelegationTokenRequest(data, version);
         }
 
         @Override
         public String toString() {
-            StringBuilder bld = new StringBuilder();
-            bld.append("(type: DescribeDelegationTokenRequest").
-                append(", owners=").append(owners).
-                append(")");
-            return bld.toString();
+            return data.toString();
         }
     }
 
-    private DescribeDelegationTokenRequest(short version, List owners) {
+    public DescribeDelegationTokenRequest(Struct struct, short version) {
         super(ApiKeys.DESCRIBE_DELEGATION_TOKEN, version);
-        this.owners = owners;
+        this.data = new DescribeDelegationTokenRequestData(struct, version);
     }
 
-    public DescribeDelegationTokenRequest(Struct struct, short versionId) {
-        super(ApiKeys.DESCRIBE_DELEGATION_TOKEN, versionId);
-
-        Object[] ownerArray = struct.getArray(OWNER_KEY_NAME);
-
-        if (ownerArray != null) {
-            owners = new ArrayList<>();
-            for (Object ownerObj : ownerArray) {
-                Struct ownerObjStruct = (Struct) ownerObj;
-                String principalType = ownerObjStruct.get(PRINCIPAL_TYPE);
-                String principalName = ownerObjStruct.get(PRINCIPAL_NAME);
-                owners.add(new KafkaPrincipal(principalType, principalName));
-            }
-        } else
-            owners = null;
-    }
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{TOKEN_DESCRIBE_REQUEST_V0, TOKEN_DESCRIBE_REQUEST_V1};
+    public DescribeDelegationTokenRequest(DescribeDelegationTokenRequestData data, short version) {
+        super(ApiKeys.DESCRIBE_DELEGATION_TOKEN, version);
+        this.data = data;
     }
 
     @Override
     protected Struct toStruct() {
-        short version = version();
-        Struct struct = new Struct(ApiKeys.DESCRIBE_DELEGATION_TOKEN.requestSchema(version));
-
-        if (owners == null) {
-            struct.set(OWNER_KEY_NAME, null);
-        } else {
-            Object[] ownersArray = new Object[owners.size()];
-
-            int i = 0;
-            for (KafkaPrincipal principal: owners) {
-                Struct ownerStruct = struct.instance(OWNER_KEY_NAME);
-                ownerStruct.set(PRINCIPAL_TYPE, principal.getPrincipalType());
-                ownerStruct.set(PRINCIPAL_NAME, principal.getName());
-                ownersArray[i++] = ownerStruct;
-            }
-
-            struct.set(OWNER_KEY_NAME, ownersArray);
-        }
-
-        return struct;
-    }
-
-    @Override
-    public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
-        return new DescribeDelegationTokenResponse(throttleTimeMs, Errors.forException(e));
+        return data.toStruct(version());
     }
 
-    public List owners() {
-        return owners;
+    public DescribeDelegationTokenRequestData data() {
+        return data;
     }
 
     public boolean ownersListEmpty() {
-        return owners != null && owners.isEmpty();
+        return data.owners() != null && data.owners().isEmpty();
     }
 
-    public static DescribeDelegationTokenRequest parse(ByteBuffer buffer, short version) {
-        return new DescribeDelegationTokenRequest(ApiKeys.DESCRIBE_DELEGATION_TOKEN.parseRequest(version, buffer), version);
+    @Override
+    public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
+        return new DescribeDelegationTokenResponse(throttleTimeMs, Errors.forException(e));
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/DescribeDelegationTokenResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/DescribeDelegationTokenResponse.java
index 64ed27bfd0003..91e7db0fccb1e 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/DescribeDelegationTokenResponse.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/DescribeDelegationTokenResponse.java
@@ -16,11 +16,11 @@
  */
 package org.apache.kafka.common.requests;
 
+import org.apache.kafka.common.message.DescribeDelegationTokenResponseData;
+import org.apache.kafka.common.message.DescribeDelegationTokenResponseData.DescribedDelegationToken;
+import org.apache.kafka.common.message.DescribeDelegationTokenResponseData.DescribedDelegationTokenRenewer;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.protocol.types.ArrayOf;
-import org.apache.kafka.common.protocol.types.Field;
-import org.apache.kafka.common.protocol.types.Schema;
 import org.apache.kafka.common.protocol.types.Struct;
 import org.apache.kafka.common.security.auth.KafkaPrincipal;
 import org.apache.kafka.common.security.token.delegation.DelegationToken;
@@ -30,169 +30,82 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
-
-import static org.apache.kafka.common.protocol.CommonFields.ERROR_CODE;
-import static org.apache.kafka.common.protocol.CommonFields.PRINCIPAL_NAME;
-import static org.apache.kafka.common.protocol.CommonFields.PRINCIPAL_TYPE;
-import static org.apache.kafka.common.protocol.CommonFields.THROTTLE_TIME_MS;
-import static org.apache.kafka.common.protocol.types.Type.BYTES;
-import static org.apache.kafka.common.protocol.types.Type.INT64;
-import static org.apache.kafka.common.protocol.types.Type.STRING;
+import java.util.stream.Collectors;
 
 public class DescribeDelegationTokenResponse extends AbstractResponse {
 
-    private static final String TOKEN_DETAILS_KEY_NAME = "token_details";
-    private static final String ISSUE_TIMESTAMP_KEY_NAME = "issue_timestamp";
-    private static final String EXPIRY_TIMESTAMP_NAME = "expiry_timestamp";
-    private static final String MAX_TIMESTAMP_NAME = "max_timestamp";
-    private static final String TOKEN_ID_KEY_NAME = "token_id";
-    private static final String HMAC_KEY_NAME = "hmac";
-    private static final String OWNER_KEY_NAME = "owner";
-    private static final String RENEWERS_KEY_NAME = "renewers";
-
-    private final Errors error;
-    private final List tokens;
-    private final int throttleTimeMs;
-
-    private static final Schema TOKEN_DETAILS_V0 = new Schema(
-        new Field(OWNER_KEY_NAME, new Schema(PRINCIPAL_TYPE, PRINCIPAL_NAME), "token owner."),
-        new Field(ISSUE_TIMESTAMP_KEY_NAME, INT64, "timestamp (in msec) when this token was generated."),
-        new Field(EXPIRY_TIMESTAMP_NAME, INT64, "timestamp (in msec) at which this token expires."),
-        new Field(MAX_TIMESTAMP_NAME, INT64, "max life time of this token."),
-        new Field(TOKEN_ID_KEY_NAME, STRING, "UUID to ensure uniqueness."),
-        new Field(HMAC_KEY_NAME, BYTES, "HMAC of the delegation token to be expired."),
-        new Field(RENEWERS_KEY_NAME, new ArrayOf(new Schema(PRINCIPAL_TYPE, PRINCIPAL_NAME)),
-            "An array of token renewers. Renewer is an Kafka PrincipalType and name string," +
-                " who is allowed to renew this token before the max lifetime expires."));
-
-    private static final Schema TOKEN_DESCRIBE_RESPONSE_V0 = new Schema(
-        ERROR_CODE,
-        new Field(TOKEN_DETAILS_KEY_NAME, new ArrayOf(TOKEN_DETAILS_V0)),
-        THROTTLE_TIME_MS);
-
-    /**
-     * The version number is bumped to indicate that on quota violation brokers send out responses before throttling.
-     */
-    private static final Schema TOKEN_DESCRIBE_RESPONSE_V1 = TOKEN_DESCRIBE_RESPONSE_V0;
+    private final DescribeDelegationTokenResponseData data;
 
     public DescribeDelegationTokenResponse(int throttleTimeMs, Errors error, List tokens) {
-        this.throttleTimeMs = throttleTimeMs;
-        this.error = error;
-        this.tokens = tokens;
+        List describedDelegationTokenList = tokens
+            .stream()
+            .map(dt -> new DescribedDelegationToken()
+                .setTokenId(dt.tokenInfo().tokenId())
+                .setPrincipalType(dt.tokenInfo().owner().getPrincipalType())
+                .setPrincipalName(dt.tokenInfo().owner().getName())
+                .setIssueTimestamp(dt.tokenInfo().issueTimestamp())
+                .setMaxTimestamp(dt.tokenInfo().maxTimestamp())
+                .setExpiryTimestamp(dt.tokenInfo().expiryTimestamp())
+                .setHmac(dt.hmac())
+                .setRenewers(dt.tokenInfo().renewers()
+                    .stream()
+                    .map(r -> new DescribedDelegationTokenRenewer().setPrincipalName(r.getName()).setPrincipalType(r.getPrincipalType()))
+                    .collect(Collectors.toList())))
+            .collect(Collectors.toList());
+
+        this.data = new DescribeDelegationTokenResponseData()
+            .setThrottleTimeMs(throttleTimeMs)
+            .setErrorCode(error.code())
+            .setTokens(describedDelegationTokenList);
     }
 
     public DescribeDelegationTokenResponse(int throttleTimeMs, Errors error) {
         this(throttleTimeMs, error, new ArrayList<>());
     }
 
-    public DescribeDelegationTokenResponse(Struct struct) {
-        Object[] requestStructs = struct.getArray(TOKEN_DETAILS_KEY_NAME);
-        List tokens = new ArrayList<>();
-
-        for (Object requestStructObj : requestStructs) {
-            Struct singleRequestStruct = (Struct) requestStructObj;
-
-            Struct ownerStruct = (Struct) singleRequestStruct.get(OWNER_KEY_NAME);
-            KafkaPrincipal owner = new KafkaPrincipal(ownerStruct.get(PRINCIPAL_TYPE), ownerStruct.get(PRINCIPAL_NAME));
-            long issueTimestamp = singleRequestStruct.getLong(ISSUE_TIMESTAMP_KEY_NAME);
-            long expiryTimestamp = singleRequestStruct.getLong(EXPIRY_TIMESTAMP_NAME);
-            long maxTimestamp = singleRequestStruct.getLong(MAX_TIMESTAMP_NAME);
-            String tokenId = singleRequestStruct.getString(TOKEN_ID_KEY_NAME);
-            ByteBuffer hmac = singleRequestStruct.getBytes(HMAC_KEY_NAME);
-
-            Object[] renewerArray = singleRequestStruct.getArray(RENEWERS_KEY_NAME);
-            List  renewers = new ArrayList<>();
-            if (renewerArray != null) {
-                for (Object renewerObj : renewerArray) {
-                    Struct renewerObjStruct = (Struct) renewerObj;
-                    String principalType = renewerObjStruct.get(PRINCIPAL_TYPE);
-                    String principalName = renewerObjStruct.get(PRINCIPAL_NAME);
-                    renewers.add(new KafkaPrincipal(principalType, principalName));
-                }
-            }
-
-            TokenInformation tokenInfo = new TokenInformation(tokenId, owner, renewers, issueTimestamp, maxTimestamp, expiryTimestamp);
-
-            byte[] hmacBytes = new byte[hmac.remaining()];
-            hmac.get(hmacBytes);
-
-            DelegationToken tokenDetails = new DelegationToken(tokenInfo, hmacBytes);
-            tokens.add(tokenDetails);
-        }
-
-        this.tokens = tokens;
-        error = Errors.forCode(struct.get(ERROR_CODE));
-        this.throttleTimeMs = struct.getOrElse(THROTTLE_TIME_MS, DEFAULT_THROTTLE_TIME);
+    public DescribeDelegationTokenResponse(Struct struct, short version) {
+        this.data = new DescribeDelegationTokenResponseData(struct, version);
     }
 
     public static DescribeDelegationTokenResponse parse(ByteBuffer buffer, short version) {
-        return new DescribeDelegationTokenResponse(ApiKeys.DESCRIBE_DELEGATION_TOKEN.responseSchema(version).read(buffer));
+        return new DescribeDelegationTokenResponse(ApiKeys.DESCRIBE_DELEGATION_TOKEN.responseSchema(version).read(buffer), version);
     }
 
     @Override
     public Map errorCounts() {
-        return errorCounts(error);
+        return errorCounts(error());
     }
 
     @Override
     protected Struct toStruct(short version) {
-        Struct struct = new Struct(ApiKeys.DESCRIBE_DELEGATION_TOKEN.responseSchema(version));
-        List tokenDetailsStructs = new ArrayList<>(tokens.size());
-
-        struct.set(ERROR_CODE, error.code());
-
-        for (DelegationToken token : tokens) {
-            TokenInformation tokenInfo = token.tokenInfo();
-            Struct singleRequestStruct = struct.instance(TOKEN_DETAILS_KEY_NAME);
-            Struct ownerStruct = singleRequestStruct.instance(OWNER_KEY_NAME);
-            ownerStruct.set(PRINCIPAL_TYPE, tokenInfo.owner().getPrincipalType());
-            ownerStruct.set(PRINCIPAL_NAME, tokenInfo.owner().getName());
-            singleRequestStruct.set(OWNER_KEY_NAME, ownerStruct);
-            singleRequestStruct.set(ISSUE_TIMESTAMP_KEY_NAME, tokenInfo.issueTimestamp());
-            singleRequestStruct.set(EXPIRY_TIMESTAMP_NAME, tokenInfo.expiryTimestamp());
-            singleRequestStruct.set(MAX_TIMESTAMP_NAME, tokenInfo.maxTimestamp());
-            singleRequestStruct.set(TOKEN_ID_KEY_NAME, tokenInfo.tokenId());
-            singleRequestStruct.set(HMAC_KEY_NAME, ByteBuffer.wrap(token.hmac()));
-
-            Object[] renewersArray = new Object[tokenInfo.renewers().size()];
-
-            int i = 0;
-            for (KafkaPrincipal principal: tokenInfo.renewers()) {
-                Struct renewerStruct = singleRequestStruct.instance(RENEWERS_KEY_NAME);
-                renewerStruct.set(PRINCIPAL_TYPE, principal.getPrincipalType());
-                renewerStruct.set(PRINCIPAL_NAME, principal.getName());
-                renewersArray[i++] = renewerStruct;
-            }
-
-            singleRequestStruct.set(RENEWERS_KEY_NAME, renewersArray);
-
-            tokenDetailsStructs.add(singleRequestStruct);
-        }
-        struct.set(TOKEN_DETAILS_KEY_NAME, tokenDetailsStructs.toArray());
-        struct.setIfExists(THROTTLE_TIME_MS, throttleTimeMs);
-
-        return struct;
-    }
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{TOKEN_DESCRIBE_RESPONSE_V0, TOKEN_DESCRIBE_RESPONSE_V1};
+        return data.toStruct(version);
     }
 
     @Override
     public int throttleTimeMs() {
-        return throttleTimeMs;
+        return data.throttleTimeMs();
     }
 
     public Errors error() {
-        return error;
+        return Errors.forCode(data.errorCode());
     }
 
     public List tokens() {
-        return tokens;
+        return data.tokens()
+            .stream()
+            .map(ddt -> new DelegationToken(new TokenInformation(
+                ddt.tokenId(),
+                new KafkaPrincipal(ddt.principalType(), ddt.principalName()),
+                ddt.renewers()
+                    .stream()
+                    .map(ddtr -> new KafkaPrincipal(ddtr.principalType(), ddtr.principalName()))
+                    .collect(Collectors.toList()), ddt.issueTimestamp(), ddt.maxTimestamp(), ddt.expiryTimestamp()),
+                ddt.hmac()))
+            .collect(Collectors.toList());
     }
 
     public boolean hasError() {
-        return this.error != Errors.NONE;
+        return error() != Errors.NONE;
     }
 
     @Override
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/DescribeGroupsResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/DescribeGroupsResponse.java
index 61b9ea2e1da87..cb369a0980867 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/DescribeGroupsResponse.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/DescribeGroupsResponse.java
@@ -54,12 +54,14 @@ public DescribeGroupsResponse(Struct struct, short version) {
 
     public static DescribedGroupMember groupMember(
         final String memberId,
+        final String groupInstanceId,
         final String clientId,
         final String clientHost,
         final byte[] assignment,
         final byte[] metadata) {
         return new DescribedGroupMember()
             .setMemberId(memberId)
+            .setGroupInstanceId(groupInstanceId)
             .setClientId(clientId)
             .setClientHost(clientHost)
             .setMemberAssignment(assignment)
@@ -85,6 +87,25 @@ public static DescribedGroup groupMetadata(
         return  groupMetada;
     }
 
+    public static DescribedGroup groupMetadata(
+        final String groupId,
+        final Errors error,
+        final String state,
+        final String protocolType,
+        final String protocol,
+        final List members,
+        final int authorizedOperations) {
+        DescribedGroup groupMetada = new DescribedGroup();
+        groupMetada.setGroupId(groupId)
+            .setErrorCode(error.code())
+            .setGroupState(state)
+            .setProtocolType(protocolType)
+            .setProtocolData(protocol)
+            .setMembers(members)
+            .setAuthorizedOperations(authorizedOperations);
+        return  groupMetada;
+    }
+
     public DescribeGroupsResponseData data() {
         return data;
     }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/ElectLeadersRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/ElectLeadersRequest.java
new file mode 100644
index 0000000000000..025733e39c0bc
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/ElectLeadersRequest.java
@@ -0,0 +1,134 @@
+/*
+ * 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.requests;
+
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+import java.util.Map;
+import org.apache.kafka.common.ElectionType;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.errors.UnsupportedVersionException;
+import org.apache.kafka.common.message.ElectLeadersRequestData.TopicPartitions;
+import org.apache.kafka.common.message.ElectLeadersRequestData;
+import org.apache.kafka.common.message.ElectLeadersResponseData.PartitionResult;
+import org.apache.kafka.common.message.ElectLeadersResponseData.ReplicaElectionResult;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.MessageUtil;
+import org.apache.kafka.common.protocol.types.Struct;
+import org.apache.kafka.common.utils.CollectionUtils;
+
+public class ElectLeadersRequest extends AbstractRequest {
+    public static class Builder extends AbstractRequest.Builder {
+        private final ElectionType electionType;
+        private final Collection topicPartitions;
+        private final int timeoutMs;
+
+        public Builder(ElectionType electionType, Collection topicPartitions, int timeoutMs) {
+            super(ApiKeys.ELECT_LEADERS);
+            this.electionType = electionType;
+            this.topicPartitions = topicPartitions;
+            this.timeoutMs = timeoutMs;
+        }
+
+        @Override
+        public ElectLeadersRequest build(short version) {
+            return new ElectLeadersRequest(toRequestData(version), version);
+        }
+
+        @Override
+        public String toString() {
+            return "ElectLeadersRequest("
+                + "electionType=" + electionType
+                + ", topicPartitions=" + ((topicPartitions == null) ? "null" : MessageUtil.deepToString(topicPartitions.iterator()))
+                + ", timeoutMs=" + timeoutMs
+                + ")";
+        }
+
+        private ElectLeadersRequestData toRequestData(short version) {
+            if (electionType != ElectionType.PREFERRED && version == 0) {
+                throw new UnsupportedVersionException("API Version 0 only supports PREFERRED election type");
+            }
+
+            ElectLeadersRequestData data = new ElectLeadersRequestData()
+                .setTimeoutMs(timeoutMs);
+
+            if (topicPartitions != null) {
+                for (Map.Entry> tp : CollectionUtils.groupPartitionsByTopic(topicPartitions).entrySet()) {
+                    data.topicPartitions().add(new ElectLeadersRequestData.TopicPartitions().setTopic(tp.getKey()).setPartitionId(tp.getValue()));
+                }
+            } else {
+                data.setTopicPartitions(null);
+            }
+
+            data.setElectionType(electionType.value);
+
+            return data;
+        }
+    }
+
+    private final ElectLeadersRequestData data;
+
+    private ElectLeadersRequest(ElectLeadersRequestData data, short version) {
+        super(ApiKeys.ELECT_LEADERS, version);
+        this.data = data;
+    }
+
+    public ElectLeadersRequest(Struct struct, short version) {
+        super(ApiKeys.ELECT_LEADERS, version);
+        this.data = new ElectLeadersRequestData(struct, version);
+    }
+
+    public ElectLeadersRequestData data() {
+        return data;
+    }
+
+    @Override
+    public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
+        ApiError apiError = ApiError.fromThrowable(e);
+        List electionResults = new ArrayList<>();
+
+        for (TopicPartitions topic : data.topicPartitions()) {
+            ReplicaElectionResult electionResult = new ReplicaElectionResult();
+
+            electionResult.setTopic(topic.topic());
+            for (Integer partitionId : topic.partitionId()) {
+                PartitionResult partitionResult = new PartitionResult();
+                partitionResult.setPartitionId(partitionId);
+                partitionResult.setErrorCode(apiError.error().code());
+                partitionResult.setErrorMessage(apiError.message());
+
+                electionResult.partitionResult().add(partitionResult);
+            }
+
+            electionResults.add(electionResult);
+        }
+
+        return new ElectLeadersResponse(throttleTimeMs, apiError.error().code(), electionResults, version());
+    }
+
+    public static ElectLeadersRequest parse(ByteBuffer buffer, short version) {
+        return new ElectLeadersRequest(ApiKeys.ELECT_LEADERS.parseRequest(version, buffer), version);
+    }
+
+    @Override
+    protected Struct toStruct() {
+        return data.toStruct(version());
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/ElectLeadersResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/ElectLeadersResponse.java
new file mode 100644
index 0000000000000..b0852040612f0
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/ElectLeadersResponse.java
@@ -0,0 +1,129 @@
+/*
+ * 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.requests;
+
+import java.nio.ByteBuffer;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.message.ElectLeadersResponseData.PartitionResult;
+import org.apache.kafka.common.message.ElectLeadersResponseData.ReplicaElectionResult;
+import org.apache.kafka.common.message.ElectLeadersResponseData;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.common.protocol.types.Struct;
+
+public class ElectLeadersResponse extends AbstractResponse {
+
+    private final short version;
+    private final ElectLeadersResponseData data;
+
+    public ElectLeadersResponse(Struct struct) {
+        this(struct, ApiKeys.ELECT_LEADERS.latestVersion());
+    }
+
+    public ElectLeadersResponse(Struct struct, short version) {
+        this.version = version;
+        this.data = new ElectLeadersResponseData(struct, version);
+    }
+
+    public ElectLeadersResponse(
+            int throttleTimeMs,
+            short errorCode,
+            List electionResults) {
+        this(throttleTimeMs, errorCode, electionResults, ApiKeys.ELECT_LEADERS.latestVersion());
+    }
+
+    public ElectLeadersResponse(
+            int throttleTimeMs,
+            short errorCode,
+            List electionResults,
+            short version) {
+
+        this.version = version;
+        this.data = new ElectLeadersResponseData();
+
+        data.setThrottleTimeMs(throttleTimeMs);
+
+        if (version >= 1) {
+            data.setErrorCode(errorCode);
+        }
+
+        data.setReplicaElectionResults(electionResults);
+    }
+
+    public ElectLeadersResponseData data() {
+        return data;
+    }
+
+    public short version() {
+        return version;
+    }
+
+    @Override
+    protected Struct toStruct(short version) {
+        return data.toStruct(version);
+    }
+
+    @Override
+    public int throttleTimeMs() {
+        return data.throttleTimeMs();
+    }
+
+    @Override
+    public Map errorCounts() {
+        HashMap counts = new HashMap<>();
+        for (ReplicaElectionResult result : data.replicaElectionResults()) {
+            for (PartitionResult partitionResult : result.partitionResult()) {
+                Errors error = Errors.forCode(partitionResult.errorCode());
+                counts.put(error, counts.getOrDefault(error, 0) + 1);
+            }
+        }
+        return counts;
+    }
+
+    public static ElectLeadersResponse parse(ByteBuffer buffer, short version) {
+        return new ElectLeadersResponse(ApiKeys.ELECT_LEADERS.responseSchema(version).read(buffer), version);
+    }
+
+    @Override
+    public boolean shouldClientThrottle(short version) {
+        return true;
+    }
+
+    public static Map> electLeadersResult(ElectLeadersResponseData data) {
+        Map> map = new HashMap<>();
+
+        for (ElectLeadersResponseData.ReplicaElectionResult topicResults : data.replicaElectionResults()) {
+            for (ElectLeadersResponseData.PartitionResult partitionResult : topicResults.partitionResult()) {
+                Optional value = Optional.empty();
+                Errors error = Errors.forCode(partitionResult.errorCode());
+                if (error != Errors.NONE) {
+                    value = Optional.of(error.exception(partitionResult.errorMessage()));
+                }
+
+                map.put(new TopicPartition(topicResults.topic(), partitionResult.partitionId()),
+                        value);
+            }
+        }
+
+        return map;
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/ElectPreferredLeadersRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/ElectPreferredLeadersRequest.java
deleted file mode 100644
index ab96e3ba9e98c..0000000000000
--- a/clients/src/main/java/org/apache/kafka/common/requests/ElectPreferredLeadersRequest.java
+++ /dev/null
@@ -1,129 +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.common.requests;
-
-import org.apache.kafka.common.TopicPartition;
-import org.apache.kafka.common.message.ElectPreferredLeadersRequestData;
-import org.apache.kafka.common.message.ElectPreferredLeadersRequestData.TopicPartitions;
-import org.apache.kafka.common.message.ElectPreferredLeadersResponseData;
-import org.apache.kafka.common.message.ElectPreferredLeadersResponseData.ReplicaElectionResult;
-import org.apache.kafka.common.protocol.ApiKeys;
-import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.protocol.types.Struct;
-import org.apache.kafka.common.utils.CollectionUtils;
-
-import java.nio.ByteBuffer;
-import java.util.Collection;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-
-public class ElectPreferredLeadersRequest extends AbstractRequest {
-    public static class Builder extends AbstractRequest.Builder {
-        private final ElectPreferredLeadersRequestData data;
-
-        public Builder(ElectPreferredLeadersRequestData data) {
-            super(ApiKeys.ELECT_PREFERRED_LEADERS);
-            this.data = data;
-        }
-
-        @Override
-        public ElectPreferredLeadersRequest build(short version) {
-            return new ElectPreferredLeadersRequest(data, version);
-        }
-
-        @Override
-        public String toString() {
-            return data.toString();
-        }
-    }
-
-    public static ElectPreferredLeadersRequestData toRequestData(Collection partitions, int timeoutMs) {
-        ElectPreferredLeadersRequestData d = new ElectPreferredLeadersRequestData()
-                .setTimeoutMs(timeoutMs);
-        if (partitions != null) {
-            for (Map.Entry> tp : CollectionUtils.groupPartitionsByTopic(partitions).entrySet()) {
-                d.topicPartitions().add(new ElectPreferredLeadersRequestData.TopicPartitions().setTopic(tp.getKey()).setPartitionId(tp.getValue()));
-            }
-        } else {
-            d.setTopicPartitions(null);
-        }
-        return d;
-    }
-
-    public static Map fromResponseData(ElectPreferredLeadersResponseData data) {
-        Map map = new HashMap<>();
-        for (ElectPreferredLeadersResponseData.ReplicaElectionResult topicResults : data.replicaElectionResults()) {
-            for (ElectPreferredLeadersResponseData.PartitionResult partitionResult : topicResults.partitionResult()) {
-                map.put(new TopicPartition(topicResults.topic(), partitionResult.partitionId()),
-                        new ApiError(Errors.forCode(partitionResult.errorCode()),
-                                partitionResult.errorMessage()));
-            }
-        }
-        return map;
-    }
-
-    private final ElectPreferredLeadersRequestData data;
-    private final short version;
-
-    private ElectPreferredLeadersRequest(ElectPreferredLeadersRequestData data, short version) {
-        super(ApiKeys.ELECT_PREFERRED_LEADERS, version);
-        this.data = data;
-        this.version = version;
-    }
-
-    public ElectPreferredLeadersRequest(Struct struct, short version) {
-        super(ApiKeys.ELECT_PREFERRED_LEADERS, version);
-        this.data = new ElectPreferredLeadersRequestData(struct, version);
-        this.version = version;
-    }
-
-    public ElectPreferredLeadersRequestData data() {
-        return data;
-    }
-
-    @Override
-    public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
-        ElectPreferredLeadersResponseData response = new ElectPreferredLeadersResponseData();
-        response.setThrottleTimeMs(throttleTimeMs);
-        ApiError apiError = ApiError.fromThrowable(e);
-        for (TopicPartitions topic : data.topicPartitions()) {
-            ReplicaElectionResult electionResult = new ReplicaElectionResult().setTopic(topic.topic());
-            for (Integer partitionId : topic.partitionId()) {
-                electionResult.partitionResult().add(new ElectPreferredLeadersResponseData.PartitionResult()
-                        .setPartitionId(partitionId)
-                        .setErrorCode(apiError.error().code())
-                        .setErrorMessage(apiError.message()));
-            }
-            response.replicaElectionResults().add(electionResult);
-        }
-        return new ElectPreferredLeadersResponse(response);
-    }
-
-    public static ElectPreferredLeadersRequest parse(ByteBuffer buffer, short version) {
-        return new ElectPreferredLeadersRequest(ApiKeys.ELECT_PREFERRED_LEADERS.parseRequest(version, buffer), version);
-    }
-
-    /**
-     * Visible for testing.
-     */
-    @Override
-    public Struct toStruct() {
-        return data.toStruct(version);
-    }
-}
\ No newline at end of file
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/ExpireDelegationTokenRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/ExpireDelegationTokenRequest.java
index 5b996763cca9e..ca6d2d67d1f79 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/ExpireDelegationTokenRequest.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/ExpireDelegationTokenRequest.java
@@ -16,102 +16,69 @@
  */
 package org.apache.kafka.common.requests;
 
+import java.nio.ByteBuffer;
+
+import org.apache.kafka.common.message.ExpireDelegationTokenRequestData;
+import org.apache.kafka.common.message.ExpireDelegationTokenResponseData;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.protocol.types.Field;
-import org.apache.kafka.common.protocol.types.Schema;
 import org.apache.kafka.common.protocol.types.Struct;
 
-import java.nio.ByteBuffer;
-
-import static org.apache.kafka.common.protocol.types.Type.BYTES;
-import static org.apache.kafka.common.protocol.types.Type.INT64;
-
 public class ExpireDelegationTokenRequest extends AbstractRequest {
 
-    private static final String HMAC_KEY_NAME = "hmac";
-    private static final String EXPIRY_TIME_PERIOD_KEY_NAME = "expiry_time_period";
-    private final ByteBuffer hmac;
-    private final long expiryTimePeriod;
-
-    private static final Schema TOKEN_EXPIRE_REQUEST_V0 = new Schema(
-        new Field(HMAC_KEY_NAME, BYTES, "HMAC of the delegation token to be expired."),
-        new Field(EXPIRY_TIME_PERIOD_KEY_NAME, INT64, "expiry time period in milli seconds."));
-
-    /**
-     * The version number is bumped to indicate that on quota violation brokers send out responses before throttling.
-     */
-    private static final Schema TOKEN_EXPIRE_REQUEST_V1 = TOKEN_EXPIRE_REQUEST_V0;
+    private final ExpireDelegationTokenRequestData data;
 
-    private ExpireDelegationTokenRequest(short version, ByteBuffer hmac, long renewTimePeriod) {
+    private ExpireDelegationTokenRequest(ExpireDelegationTokenRequestData data, short version) {
         super(ApiKeys.EXPIRE_DELEGATION_TOKEN, version);
-
-        this.hmac = hmac;
-        this.expiryTimePeriod = renewTimePeriod;
+        this.data = data;
     }
 
-    public ExpireDelegationTokenRequest(Struct struct, short versionId) {
-        super(ApiKeys.EXPIRE_DELEGATION_TOKEN, versionId);
-
-        hmac = struct.getBytes(HMAC_KEY_NAME);
-        expiryTimePeriod = struct.getLong(EXPIRY_TIME_PERIOD_KEY_NAME);
+    public ExpireDelegationTokenRequest(Struct struct, short version) {
+        super(ApiKeys.EXPIRE_DELEGATION_TOKEN, version);
+        this.data = new ExpireDelegationTokenRequestData(struct, version);
     }
 
     public static ExpireDelegationTokenRequest parse(ByteBuffer buffer, short version) {
         return new ExpireDelegationTokenRequest(ApiKeys.EXPIRE_DELEGATION_TOKEN.parseRequest(version, buffer), version);
     }
 
-    public static Schema[] schemaVersions() {
-        return new Schema[] {TOKEN_EXPIRE_REQUEST_V0, TOKEN_EXPIRE_REQUEST_V1};
-    }
-
     @Override
     protected Struct toStruct() {
-        short version = version();
-        Struct struct = new Struct(ApiKeys.EXPIRE_DELEGATION_TOKEN.requestSchema(version));
-
-        struct.set(HMAC_KEY_NAME, hmac);
-        struct.set(EXPIRY_TIME_PERIOD_KEY_NAME, expiryTimePeriod);
-
-        return struct;
+        return data.toStruct(version());
     }
 
     @Override
     public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
-        return new ExpireDelegationTokenResponse(throttleTimeMs, Errors.forException(e));
+        return new ExpireDelegationTokenResponse(
+                new ExpireDelegationTokenResponseData()
+                    .setErrorCode(Errors.forException(e).code())
+                    .setThrottleTimeMs(throttleTimeMs));
     }
 
     public ByteBuffer hmac() {
-        return hmac;
+        return ByteBuffer.wrap(data.hmac());
     }
 
     public long expiryTimePeriod() {
-        return expiryTimePeriod;
+        return data.expiryTimePeriodMs();
     }
 
     public static class Builder extends AbstractRequest.Builder {
-        private final ByteBuffer hmac;
-        private final long expiryTimePeriod;
+        private final ExpireDelegationTokenRequestData data;
 
-        public Builder(byte[] hmac, long expiryTimePeriod) {
+        public Builder(ExpireDelegationTokenRequestData data) {
             super(ApiKeys.EXPIRE_DELEGATION_TOKEN);
-            this.hmac = ByteBuffer.wrap(hmac);
-            this.expiryTimePeriod = expiryTimePeriod;
+            this.data = data;
         }
 
         @Override
         public ExpireDelegationTokenRequest build(short version) {
-            return new ExpireDelegationTokenRequest(version, hmac, expiryTimePeriod);
+            return new ExpireDelegationTokenRequest(data, version);
         }
 
         @Override
         public String toString() {
-            StringBuilder bld = new StringBuilder();
-            bld.append("(type: ExpireDelegationTokenRequest").
-                append(", hmac=").append(hmac).
-                append(", expiryTimePeriod=").append(expiryTimePeriod).
-                append(")");
-            return bld.toString();
+            return data.toString();
         }
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/ExpireDelegationTokenResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/ExpireDelegationTokenResponse.java
index 9491a3546a1b9..16a6e8c5b40f5 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/ExpireDelegationTokenResponse.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/ExpireDelegationTokenResponse.java
@@ -16,92 +16,56 @@
  */
 package org.apache.kafka.common.requests;
 
-import org.apache.kafka.common.protocol.ApiKeys;
-import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.protocol.types.Field;
-import org.apache.kafka.common.protocol.types.Schema;
-import org.apache.kafka.common.protocol.types.Struct;
-
 import java.nio.ByteBuffer;
+import java.util.Collections;
 import java.util.Map;
 
-import static org.apache.kafka.common.protocol.CommonFields.ERROR_CODE;
-import static org.apache.kafka.common.protocol.CommonFields.THROTTLE_TIME_MS;
-import static org.apache.kafka.common.protocol.types.Type.INT64;
+import org.apache.kafka.common.message.ExpireDelegationTokenResponseData;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.common.protocol.types.Struct;
 
 public class ExpireDelegationTokenResponse extends AbstractResponse {
 
-    private static final String EXPIRY_TIMESTAMP_KEY_NAME = "expiry_timestamp";
-
-    private final Errors error;
-    private final long expiryTimestamp;
-    private final int throttleTimeMs;
+    private final ExpireDelegationTokenResponseData data;
 
-    private  static final Schema TOKEN_EXPIRE_RESPONSE_V0 = new Schema(
-        ERROR_CODE,
-        new Field(EXPIRY_TIMESTAMP_KEY_NAME, INT64, "timestamp (in msec) at which this token expires.."),
-        THROTTLE_TIME_MS);
-
-    /**
-     * The version number is bumped to indicate that on quota violation brokers send out responses before throttling.
-     */
-    private static final Schema TOKEN_EXPIRE_RESPONSE_V1 = TOKEN_EXPIRE_RESPONSE_V0;
-
-    public ExpireDelegationTokenResponse(int throttleTimeMs, Errors error, long expiryTimestamp) {
-        this.throttleTimeMs = throttleTimeMs;
-        this.error = error;
-        this.expiryTimestamp = expiryTimestamp;
-    }
-
-    public ExpireDelegationTokenResponse(int throttleTimeMs, Errors error) {
-        this(throttleTimeMs, error, -1);
+    public ExpireDelegationTokenResponse(ExpireDelegationTokenResponseData data) {
+        this.data = data;
     }
 
-    public ExpireDelegationTokenResponse(Struct struct) {
-        error = Errors.forCode(struct.get(ERROR_CODE));
-        this.expiryTimestamp = struct.getLong(EXPIRY_TIMESTAMP_KEY_NAME);
-        this.throttleTimeMs = struct.getOrElse(THROTTLE_TIME_MS, DEFAULT_THROTTLE_TIME);
+    public ExpireDelegationTokenResponse(Struct struct, short version) {
+        this.data = new ExpireDelegationTokenResponseData(struct, version);
     }
 
     public static ExpireDelegationTokenResponse parse(ByteBuffer buffer, short version) {
-        return new ExpireDelegationTokenResponse(ApiKeys.EXPIRE_DELEGATION_TOKEN.responseSchema(version).read(buffer));
-    }
-
-    public static Schema[] schemaVersions() {
-        return new Schema[] {TOKEN_EXPIRE_RESPONSE_V0, TOKEN_EXPIRE_RESPONSE_V1};
+        return new ExpireDelegationTokenResponse(ApiKeys.EXPIRE_DELEGATION_TOKEN.responseSchema(version).read(buffer), version);
     }
 
     public Errors error() {
-        return error;
+        return Errors.forCode(data.errorCode());
     }
 
     public long expiryTimestamp() {
-        return expiryTimestamp;
+        return data.expiryTimestampMs();
     }
 
     @Override
     public Map errorCounts() {
-        return errorCounts(error);
+        return Collections.singletonMap(error(), 1);
     }
 
     @Override
     protected Struct toStruct(short version) {
-        Struct struct = new Struct(ApiKeys.EXPIRE_DELEGATION_TOKEN.responseSchema(version));
-
-        struct.set(ERROR_CODE, error.code());
-        struct.set(EXPIRY_TIMESTAMP_KEY_NAME, expiryTimestamp);
-        struct.setIfExists(THROTTLE_TIME_MS, throttleTimeMs);
-
-        return struct;
+        return data.toStruct(version);
     }
 
     @Override
     public int throttleTimeMs() {
-        return throttleTimeMs;
+        return data.throttleTimeMs();
     }
 
     public boolean hasError() {
-        return this.error != Errors.NONE;
+        return error() != Errors.NONE;
     }
 
     @Override
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/FetchRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/FetchRequest.java
index b3443a1f70447..2c0455ae4ae1e 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/FetchRequest.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/FetchRequest.java
@@ -63,6 +63,7 @@ public class FetchRequest extends AbstractRequest {
                     "consumers to discard ABORTED transactional records");
     private static final Field.Int32 SESSION_ID = new Field.Int32("session_id", "The fetch session ID");
     private static final Field.Int32 SESSION_EPOCH = new Field.Int32("session_epoch", "The fetch session epoch");
+    private static final Field.Str RACK_ID = new Field.Str("rack_id", "The consumer's rack id");
 
     // topic level fields
     private static final Field.ComplexArray PARTITIONS = new Field.ComplexArray("partitions",
@@ -194,10 +195,23 @@ public class FetchRequest extends AbstractRequest {
     // V10 bumped up to indicate ZStandard capability. (see KIP-110)
     private static final Schema FETCH_REQUEST_V10 = FETCH_REQUEST_V9;
 
+    // V11 added rack ID to support read from followers (KIP-392)
+    private static final Schema FETCH_REQUEST_V11 = new Schema(
+            REPLICA_ID,
+            MAX_WAIT_TIME,
+            MIN_BYTES,
+            MAX_BYTES,
+            ISOLATION_LEVEL,
+            SESSION_ID,
+            SESSION_EPOCH,
+            FETCH_REQUEST_TOPIC_V9,
+            FORGOTTEN_TOPIC_DATA_V7,
+            RACK_ID);
+
     public static Schema[] schemaVersions() {
         return new Schema[]{FETCH_REQUEST_V0, FETCH_REQUEST_V1, FETCH_REQUEST_V2, FETCH_REQUEST_V3, FETCH_REQUEST_V4,
             FETCH_REQUEST_V5, FETCH_REQUEST_V6, FETCH_REQUEST_V7, FETCH_REQUEST_V8, FETCH_REQUEST_V9,
-            FETCH_REQUEST_V10};
+            FETCH_REQUEST_V10, FETCH_REQUEST_V11};
     }
 
     // default values for older versions where a request level limit did not exist
@@ -217,6 +231,7 @@ public static Schema[] schemaVersions() {
 
     private final List toForget;
     private final FetchMetadata metadata;
+    private final String rackId;
 
     public static final class PartitionData {
         public final long fetchOffset;
@@ -290,6 +305,7 @@ public static class Builder extends AbstractRequest.Builder {
         private int maxBytes = DEFAULT_RESPONSE_MAX_BYTES;
         private FetchMetadata metadata = FetchMetadata.LEGACY;
         private List toForget = Collections.emptyList();
+        private String rackId = "";
 
         public static Builder forConsumer(int maxWait, int minBytes, Map fetchData) {
             return new Builder(ApiKeys.FETCH.oldestVersion(), ApiKeys.FETCH.latestVersion(),
@@ -320,6 +336,11 @@ public Builder metadata(FetchMetadata metadata) {
             return this;
         }
 
+        public Builder rackId(String rackId) {
+            this.rackId = rackId;
+            return this;
+        }
+
         public Map fetchData() {
             return this.fetchData;
         }
@@ -345,7 +366,7 @@ public FetchRequest build(short version) {
             }
 
             return new FetchRequest(version, replicaId, maxWait, minBytes, maxBytes, fetchData,
-                isolationLevel, toForget, metadata);
+                isolationLevel, toForget, metadata, rackId);
         }
 
         @Override
@@ -360,6 +381,7 @@ public String toString() {
                     append(", isolationLevel=").append(isolationLevel).
                     append(", toForget=").append(Utils.join(toForget, ", ")).
                     append(", metadata=").append(metadata).
+                    append(", rackId=").append(rackId).
                     append(")");
             return bld.toString();
         }
@@ -367,7 +389,7 @@ public String toString() {
 
     private FetchRequest(short version, int replicaId, int maxWait, int minBytes, int maxBytes,
                          Map fetchData, IsolationLevel isolationLevel,
-                         List toForget, FetchMetadata metadata) {
+                         List toForget, FetchMetadata metadata, String rackId) {
         super(ApiKeys.FETCH, version);
         this.replicaId = replicaId;
         this.maxWait = maxWait;
@@ -377,6 +399,7 @@ private FetchRequest(short version, int replicaId, int maxWait, int minBytes, in
         this.isolationLevel = isolationLevel;
         this.toForget = toForget;
         this.metadata = metadata;
+        this.rackId = rackId;
     }
 
     public FetchRequest(Struct struct, short version) {
@@ -421,6 +444,7 @@ public FetchRequest(Struct struct, short version) {
                 fetchData.put(new TopicPartition(topic, partition), partitionData);
             }
         }
+        rackId = struct.getOrElse(RACK_ID, "");
     }
 
     @Override
@@ -436,7 +460,7 @@ public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
         for (Map.Entry entry : fetchData.entrySet()) {
             FetchResponse.PartitionData partitionResponse = new FetchResponse.PartitionData<>(error,
                 FetchResponse.INVALID_HIGHWATERMARK, FetchResponse.INVALID_LAST_STABLE_OFFSET,
-                FetchResponse.INVALID_LOG_START_OFFSET, null, MemoryRecords.EMPTY);
+                FetchResponse.INVALID_LOG_START_OFFSET, Optional.empty(), null, MemoryRecords.EMPTY);
             responseData.put(entry.getKey(), partitionResponse);
         }
         return new FetchResponse<>(error, responseData, throttleTimeMs, metadata.sessionId());
@@ -478,6 +502,10 @@ public FetchMetadata metadata() {
         return metadata;
     }
 
+    public String rackId() {
+        return rackId;
+    }
+
     public static FetchRequest parse(ByteBuffer buffer, short version) {
         return new FetchRequest(ApiKeys.FETCH.parseRequest(version, buffer), version);
     }
@@ -530,6 +558,7 @@ protected Struct toStruct() {
             }
             struct.set(FORGOTTEN_TOPICS, toForgetStructs.toArray());
         }
+        struct.setIfExists(RACK_ID, rackId);
         return struct;
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/FetchResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/FetchResponse.java
index 9c29d375ccba0..d387e8199966c 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/FetchResponse.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/FetchResponse.java
@@ -37,7 +37,10 @@
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
 import java.util.Queue;
+import java.util.function.Predicate;
 
 import static org.apache.kafka.common.protocol.CommonFields.ERROR_CODE;
 import static org.apache.kafka.common.protocol.CommonFields.PARTITION_ID;
@@ -77,6 +80,8 @@ public class FetchResponse extends AbstractResponse {
             "Last committed offset.");
     private static final Field.Int64 LOG_START_OFFSET = new Field.Int64("log_start_offset",
             "Earliest available offset.");
+    private static final Field.Int32 PREFERRED_READ_REPLICA = new Field.Int32("preferred_read_replica",
+            "The ID of the replica that the consumer should prefer.");
 
     private static final String PARTITION_HEADER_KEY_NAME = "partition_header";
     private static final String ABORTED_TRANSACTIONS_KEY_NAME = "aborted_transactions";
@@ -139,6 +144,16 @@ public class FetchResponse extends AbstractResponse {
             LOG_START_OFFSET,
             new Field(ABORTED_TRANSACTIONS_KEY_NAME, ArrayOf.nullable(FETCH_RESPONSE_ABORTED_TRANSACTION_V4)));
 
+    // Introduced in V11 to support read from followers (KIP-392)
+    private static final Schema FETCH_RESPONSE_PARTITION_HEADER_V6 = new Schema(
+            PARTITION_ID,
+            ERROR_CODE,
+            HIGH_WATERMARK,
+            LAST_STABLE_OFFSET,
+            LOG_START_OFFSET,
+            new Field(ABORTED_TRANSACTIONS_KEY_NAME, ArrayOf.nullable(FETCH_RESPONSE_ABORTED_TRANSACTION_V4)),
+            PREFERRED_READ_REPLICA);
+
     private static final Schema FETCH_RESPONSE_PARTITION_V4 = new Schema(
             new Field(PARTITION_HEADER_KEY_NAME, FETCH_RESPONSE_PARTITION_HEADER_V4),
             new Field(RECORD_SET_KEY_NAME, RECORDS));
@@ -147,6 +162,10 @@ public class FetchResponse extends AbstractResponse {
             new Field(PARTITION_HEADER_KEY_NAME, FETCH_RESPONSE_PARTITION_HEADER_V5),
             new Field(RECORD_SET_KEY_NAME, RECORDS));
 
+    private static final Schema FETCH_RESPONSE_PARTITION_V6 = new Schema(
+            new Field(PARTITION_HEADER_KEY_NAME, FETCH_RESPONSE_PARTITION_HEADER_V6),
+            new Field(RECORD_SET_KEY_NAME, RECORDS));
+
     private static final Schema FETCH_RESPONSE_TOPIC_V4 = new Schema(
             TOPIC_NAME,
             new Field(PARTITIONS_KEY_NAME, new ArrayOf(FETCH_RESPONSE_PARTITION_V4)));
@@ -155,6 +174,10 @@ public class FetchResponse extends AbstractResponse {
             TOPIC_NAME,
             new Field(PARTITIONS_KEY_NAME, new ArrayOf(FETCH_RESPONSE_PARTITION_V5)));
 
+    private static final Schema FETCH_RESPONSE_TOPIC_V6 = new Schema(
+            TOPIC_NAME,
+            new Field(PARTITIONS_KEY_NAME, new ArrayOf(FETCH_RESPONSE_PARTITION_V6)));
+
     private static final Schema FETCH_RESPONSE_V4 = new Schema(
             THROTTLE_TIME_MS,
             new Field(RESPONSES_KEY_NAME, new ArrayOf(FETCH_RESPONSE_TOPIC_V4)));
@@ -185,15 +208,25 @@ public class FetchResponse extends AbstractResponse {
     // V10 bumped up to indicate ZStandard capability. (see KIP-110)
     private static final Schema FETCH_RESPONSE_V10 = FETCH_RESPONSE_V9;
 
+    // V11 added preferred read replica for each partition response to support read from followers (KIP-392)
+    private static final Schema FETCH_RESPONSE_V11 = new Schema(
+            THROTTLE_TIME_MS,
+            ERROR_CODE,
+            SESSION_ID,
+            new Field(RESPONSES_KEY_NAME, new ArrayOf(FETCH_RESPONSE_TOPIC_V6)));
+
+
     public static Schema[] schemaVersions() {
         return new Schema[] {FETCH_RESPONSE_V0, FETCH_RESPONSE_V1, FETCH_RESPONSE_V2,
             FETCH_RESPONSE_V3, FETCH_RESPONSE_V4, FETCH_RESPONSE_V5, FETCH_RESPONSE_V6,
-            FETCH_RESPONSE_V7, FETCH_RESPONSE_V8, FETCH_RESPONSE_V9, FETCH_RESPONSE_V10};
+            FETCH_RESPONSE_V7, FETCH_RESPONSE_V8, FETCH_RESPONSE_V9, FETCH_RESPONSE_V10,
+            FETCH_RESPONSE_V11};
     }
 
     public static final long INVALID_HIGHWATERMARK = -1L;
     public static final long INVALID_LAST_STABLE_OFFSET = -1L;
     public static final long INVALID_LOG_START_OFFSET = -1L;
+    public static final int INVALID_PREFERRED_REPLICA_ID = -1;
 
     private final int throttleTimeMs;
     private final Errors error;
@@ -223,8 +256,8 @@ public boolean equals(Object o) {
 
         @Override
         public int hashCode() {
-            int result = (int) (producerId ^ (producerId >>> 32));
-            result = 31 * result + (int) (firstOffset ^ (firstOffset >>> 32));
+            int result = Long.hashCode(producerId);
+            result = 31 * result + Long.hashCode(firstOffset);
             return result;
         }
 
@@ -239,9 +272,26 @@ public static final class PartitionData {
         public final long highWatermark;
         public final long lastStableOffset;
         public final long logStartOffset;
+        public final Optional preferredReadReplica;
         public final List abortedTransactions;
         public final T records;
 
+        public PartitionData(Errors error,
+                             long highWatermark,
+                             long lastStableOffset,
+                             long logStartOffset,
+                             Optional preferredReadReplica,
+                             List abortedTransactions,
+                             T records) {
+            this.error = error;
+            this.highWatermark = highWatermark;
+            this.lastStableOffset = lastStableOffset;
+            this.logStartOffset = logStartOffset;
+            this.preferredReadReplica = preferredReadReplica;
+            this.abortedTransactions = abortedTransactions;
+            this.records = records;
+        }
+
         public PartitionData(Errors error,
                              long highWatermark,
                              long lastStableOffset,
@@ -252,6 +302,7 @@ public PartitionData(Errors error,
             this.highWatermark = highWatermark;
             this.lastStableOffset = lastStableOffset;
             this.logStartOffset = logStartOffset;
+            this.preferredReadReplica = Optional.empty();
             this.abortedTransactions = abortedTransactions;
             this.records = records;
         }
@@ -269,16 +320,18 @@ public boolean equals(Object o) {
                     highWatermark == that.highWatermark &&
                     lastStableOffset == that.lastStableOffset &&
                     logStartOffset == that.logStartOffset &&
-                    (abortedTransactions == null ? that.abortedTransactions == null : abortedTransactions.equals(that.abortedTransactions)) &&
-                    (records == null ? that.records == null : records.equals(that.records));
+                    Objects.equals(preferredReadReplica, that.preferredReadReplica) &&
+                    Objects.equals(abortedTransactions, that.abortedTransactions) &&
+                    Objects.equals(records, that.records);
         }
 
         @Override
         public int hashCode() {
             int result = error != null ? error.hashCode() : 0;
-            result = 31 * result + (int) (highWatermark ^ (highWatermark >>> 32));
-            result = 31 * result + (int) (lastStableOffset ^ (lastStableOffset >>> 32));
-            result = 31 * result + (int) (logStartOffset ^ (logStartOffset >>> 32));
+            result = 31 * result + Long.hashCode(highWatermark);
+            result = 31 * result + Long.hashCode(lastStableOffset);
+            result = 31 * result + Long.hashCode(logStartOffset);
+            result = 31 * result + Objects.hashCode(preferredReadReplica);
             result = 31 * result + (abortedTransactions != null ? abortedTransactions.hashCode() : 0);
             result = 31 * result + (records != null ? records.hashCode() : 0);
             return result;
@@ -290,6 +343,7 @@ public String toString() {
                     ", highWaterMark=" + highWatermark +
                     ", lastStableOffset = " + lastStableOffset +
                     ", logStartOffset = " + logStartOffset +
+                    ", preferredReadReplica = " + preferredReadReplica.map(Object::toString).orElse("absent") +
                     ", abortedTransactions = " + abortedTransactions +
                     ", recordsSizeInBytes=" + records.sizeInBytes() + ")";
         }
@@ -328,6 +382,9 @@ public static FetchResponse parse(Struct struct) {
                 long highWatermark = partitionResponseHeader.get(HIGH_WATERMARK);
                 long lastStableOffset = partitionResponseHeader.getOrElse(LAST_STABLE_OFFSET, INVALID_LAST_STABLE_OFFSET);
                 long logStartOffset = partitionResponseHeader.getOrElse(LOG_START_OFFSET, INVALID_LOG_START_OFFSET);
+                Optional preferredReadReplica = Optional.of(
+                    partitionResponseHeader.getOrElse(PREFERRED_READ_REPLICA, INVALID_PREFERRED_REPLICA_ID)
+                ).filter(Predicate.isEqual(INVALID_PREFERRED_REPLICA_ID).negate());
 
                 BaseRecords baseRecords = partitionResponse.getRecords(RECORD_SET_KEY_NAME);
                 if (!(baseRecords instanceof MemoryRecords))
@@ -349,7 +406,7 @@ public static FetchResponse parse(Struct struct) {
                 }
 
                 PartitionData partitionData = new PartitionData<>(error, highWatermark, lastStableOffset,
-                        logStartOffset, abortedTransactions, records);
+                        logStartOffset, preferredReadReplica, abortedTransactions, records);
                 responseData.put(new TopicPartition(topic, partition), partitionData);
             }
         }
@@ -512,6 +569,7 @@ private static  Struct toStruct(short version, int thrott
                     }
                 }
                 partitionDataHeader.setIfExists(LOG_START_OFFSET, fetchPartitionData.logStartOffset);
+                partitionDataHeader.setIfExists(PREFERRED_READ_REPLICA, fetchPartitionData.preferredReadReplica.orElse(-1));
                 partitionData.set(PARTITION_HEADER_KEY_NAME, partitionDataHeader);
                 partitionData.set(RECORD_SET_KEY_NAME, fetchPartitionData.records);
                 partitionArray.add(partitionData);
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/FindCoordinatorRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/FindCoordinatorRequest.java
index 2d44ab3015bae..0e728439169d4 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/FindCoordinatorRequest.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/FindCoordinatorRequest.java
@@ -18,122 +18,64 @@
 
 import org.apache.kafka.common.Node;
 import org.apache.kafka.common.errors.UnsupportedVersionException;
+import org.apache.kafka.common.message.FindCoordinatorRequestData;
+import org.apache.kafka.common.message.FindCoordinatorResponseData;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.protocol.types.Field;
-import org.apache.kafka.common.protocol.types.Schema;
 import org.apache.kafka.common.protocol.types.Struct;
 
 import java.nio.ByteBuffer;
 
-import static org.apache.kafka.common.protocol.CommonFields.GROUP_ID;
-import static org.apache.kafka.common.protocol.types.Type.INT8;
-import static org.apache.kafka.common.protocol.types.Type.STRING;
 
 public class FindCoordinatorRequest extends AbstractRequest {
-    private static final String COORDINATOR_KEY_KEY_NAME = "coordinator_key";
-    private static final String COORDINATOR_TYPE_KEY_NAME = "coordinator_type";
-
-    private static final Schema FIND_COORDINATOR_REQUEST_V0 = new Schema(GROUP_ID);
-
-    private static final Schema FIND_COORDINATOR_REQUEST_V1 = new Schema(
-            new Field("coordinator_key", STRING, "Id to use for finding the coordinator (for groups, this is the groupId, " +
-                            "for transactional producers, this is the transactional id)"),
-            new Field("coordinator_type", INT8, "The type of coordinator to find (0 = group, 1 = transaction)"));
-
-    /**
-     * The version number is bumped to indicate that on quota violation brokers send out responses before throttling.
-     */
-    private static final Schema FIND_COORDINATOR_REQUEST_V2 = FIND_COORDINATOR_REQUEST_V1;
-
-    public static Schema[] schemaVersions() {
-        return new Schema[] {FIND_COORDINATOR_REQUEST_V0, FIND_COORDINATOR_REQUEST_V1, FIND_COORDINATOR_REQUEST_V2};
-    }
 
     public static class Builder extends AbstractRequest.Builder {
-        private final String coordinatorKey;
-        private final CoordinatorType coordinatorType;
-        private final short minVersion;
+        private final FindCoordinatorRequestData data;
 
-        public Builder(CoordinatorType coordinatorType, String coordinatorKey) {
+        public Builder(FindCoordinatorRequestData data) {
             super(ApiKeys.FIND_COORDINATOR);
-            this.coordinatorType = coordinatorType;
-            this.coordinatorKey = coordinatorKey;
-            this.minVersion = coordinatorType == CoordinatorType.TRANSACTION ? (short) 1 : (short) 0;
+            this.data = data;
         }
 
         @Override
         public FindCoordinatorRequest build(short version) {
-            if (version < minVersion)
+            if (version < 1 && data.keyType() == CoordinatorType.TRANSACTION.id()) {
                 throw new UnsupportedVersionException("Cannot create a v" + version + " FindCoordinator request " +
-                        "because we require features supported only in " + minVersion + " or later.");
-            return new FindCoordinatorRequest(coordinatorType, coordinatorKey, version);
-        }
-
-        public String coordinatorKey() {
-            return coordinatorKey;
-        }
-
-        public CoordinatorType coordinatorType() {
-            return coordinatorType;
+                        "because we require features supported only in 2 or later.");
+            }
+            return new FindCoordinatorRequest(data, version);
         }
 
         @Override
         public String toString() {
-            StringBuilder bld = new StringBuilder();
-            bld.append("(type=FindCoordinatorRequest, coordinatorKey=");
-            bld.append(coordinatorKey);
-            bld.append(", coordinatorType=");
-            bld.append(coordinatorType);
-            bld.append(")");
-            return bld.toString();
+            return data.toString();
+        }
+
+        public FindCoordinatorRequestData data() {
+            return data;
         }
     }
 
-    private final String coordinatorKey;
-    private final CoordinatorType coordinatorType;
+    private final FindCoordinatorRequestData data;
 
-    private FindCoordinatorRequest(CoordinatorType coordinatorType, String coordinatorKey, short version) {
+    private FindCoordinatorRequest(FindCoordinatorRequestData data, short version) {
         super(ApiKeys.FIND_COORDINATOR, version);
-        this.coordinatorType = coordinatorType;
-        this.coordinatorKey = coordinatorKey;
+        this.data = data;
     }
 
     public FindCoordinatorRequest(Struct struct, short version) {
         super(ApiKeys.FIND_COORDINATOR, version);
-
-        if (struct.hasField(COORDINATOR_TYPE_KEY_NAME))
-            this.coordinatorType = CoordinatorType.forId(struct.getByte(COORDINATOR_TYPE_KEY_NAME));
-        else
-            this.coordinatorType = CoordinatorType.GROUP;
-        if (struct.hasField(GROUP_ID))
-            this.coordinatorKey = struct.get(GROUP_ID);
-        else
-            this.coordinatorKey = struct.getString(COORDINATOR_KEY_KEY_NAME);
+        this.data = new FindCoordinatorRequestData(struct, version);
     }
 
     @Override
     public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
-        short versionId = version();
-        switch (versionId) {
-            case 0:
-                return new FindCoordinatorResponse(Errors.forException(e), Node.noNode());
-            case 1:
-            case 2:
-                return new FindCoordinatorResponse(throttleTimeMs, Errors.forException(e), Node.noNode());
-
-            default:
-                throw new IllegalArgumentException(String.format("Version %d is not valid. Valid versions for %s are 0 to %d",
-                        versionId, this.getClass().getSimpleName(), ApiKeys.FIND_COORDINATOR.latestVersion()));
+        FindCoordinatorResponseData response = new FindCoordinatorResponseData();
+        if (version() >= 2) {
+            response.setThrottleTimeMs(throttleTimeMs);
         }
-    }
-
-    public String coordinatorKey() {
-        return coordinatorKey;
-    }
-
-    public CoordinatorType coordinatorType() {
-        return coordinatorType;
+        Errors error = Errors.forException(e);
+        return FindCoordinatorResponse.prepareResponse(error, Node.noNode());
     }
 
     public static FindCoordinatorRequest parse(ByteBuffer buffer, short version) {
@@ -142,14 +84,11 @@ public static FindCoordinatorRequest parse(ByteBuffer buffer, short version) {
 
     @Override
     protected Struct toStruct() {
-        Struct struct = new Struct(ApiKeys.FIND_COORDINATOR.requestSchema(version()));
-        if (struct.hasField(GROUP_ID))
-            struct.set(GROUP_ID, coordinatorKey);
-        else
-            struct.set(COORDINATOR_KEY_KEY_NAME, coordinatorKey);
-        if (struct.hasField(COORDINATOR_TYPE_KEY_NAME))
-            struct.set(COORDINATOR_TYPE_KEY_NAME, coordinatorType.id);
-        return struct;
+        return data.toStruct(version());
+    }
+
+    public FindCoordinatorRequestData data() {
+        return data;
     }
 
     public enum CoordinatorType {
@@ -161,6 +100,10 @@ public enum CoordinatorType {
             this.id = id;
         }
 
+        public byte id() {
+            return id;
+        }
+
         public static CoordinatorType forId(byte id) {
             switch (id) {
                 case 0:
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/FindCoordinatorResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/FindCoordinatorResponse.java
index bc7f654c0bb6e..c880408a8ed30 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/FindCoordinatorResponse.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/FindCoordinatorResponse.java
@@ -17,53 +17,17 @@
 package org.apache.kafka.common.requests;
 
 import org.apache.kafka.common.Node;
+import org.apache.kafka.common.message.FindCoordinatorResponseData;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.protocol.types.Field;
-import org.apache.kafka.common.protocol.types.Schema;
 import org.apache.kafka.common.protocol.types.Struct;
 
 import java.nio.ByteBuffer;
+import java.util.Collections;
 import java.util.Map;
 
-import static org.apache.kafka.common.protocol.CommonFields.ERROR_CODE;
-import static org.apache.kafka.common.protocol.CommonFields.ERROR_MESSAGE;
-import static org.apache.kafka.common.protocol.CommonFields.THROTTLE_TIME_MS;
-import static org.apache.kafka.common.protocol.types.Type.INT32;
-import static org.apache.kafka.common.protocol.types.Type.STRING;
 
 public class FindCoordinatorResponse extends AbstractResponse {
-    private static final String COORDINATOR_KEY_NAME = "coordinator";
-
-    // coordinator level field names
-    private static final String NODE_ID_KEY_NAME = "node_id";
-    private static final String HOST_KEY_NAME = "host";
-    private static final String PORT_KEY_NAME = "port";
-
-    private static final Schema FIND_COORDINATOR_BROKER_V0 = new Schema(
-            new Field(NODE_ID_KEY_NAME, INT32, "The broker id."),
-            new Field(HOST_KEY_NAME, STRING, "The hostname of the broker."),
-            new Field(PORT_KEY_NAME, INT32, "The port on which the broker accepts requests."));
-
-    private static final Schema FIND_COORDINATOR_RESPONSE_V0 = new Schema(
-            ERROR_CODE,
-            new Field(COORDINATOR_KEY_NAME, FIND_COORDINATOR_BROKER_V0, "Host and port information for the coordinator " +
-                    "for a consumer group."));
-
-    private static final Schema FIND_COORDINATOR_RESPONSE_V1 = new Schema(
-            THROTTLE_TIME_MS,
-            ERROR_CODE,
-            ERROR_MESSAGE,
-            new Field(COORDINATOR_KEY_NAME, FIND_COORDINATOR_BROKER_V0, "Host and port information for the coordinator"));
-
-    /**
-     * The version number is bumped to indicate that on quota violation brokers send out responses before throttling.
-     */
-    private static final Schema FIND_COORDINATOR_RESPONSE_V2 = FIND_COORDINATOR_RESPONSE_V1;
-
-    public static Schema[] schemaVersions() {
-        return new Schema[] {FIND_COORDINATOR_RESPONSE_V0, FIND_COORDINATOR_RESPONSE_V1, FIND_COORDINATOR_RESPONSE_V2};
-    }
 
     /**
      * Possible error codes:
@@ -75,88 +39,68 @@ public static Schema[] schemaVersions() {
      * TRANSACTIONAL_ID_AUTHORIZATION_FAILED (53)
      */
 
+    private final FindCoordinatorResponseData data;
 
-    private final int throttleTimeMs;
-    private final String errorMessage;
-    private final Errors error;
-    private final Node node;
-
-    public FindCoordinatorResponse(Errors error, Node node) {
-        this(DEFAULT_THROTTLE_TIME, error, node);
+    public FindCoordinatorResponse(FindCoordinatorResponseData data) {
+        this.data = data;
     }
 
-    public FindCoordinatorResponse(int throttleTimeMs, Errors error, Node node) {
-        this.throttleTimeMs = throttleTimeMs;
-        this.error = error;
-        this.node = node;
-        this.errorMessage = null;
+    public FindCoordinatorResponse(Struct struct, short version) {
+        this.data = new FindCoordinatorResponseData(struct, version);
     }
 
-    public FindCoordinatorResponse(Struct struct) {
-        this.throttleTimeMs = struct.getOrElse(THROTTLE_TIME_MS, DEFAULT_THROTTLE_TIME);
-        error = Errors.forCode(struct.get(ERROR_CODE));
-        errorMessage = struct.getOrElse(ERROR_MESSAGE, null);
+    public FindCoordinatorResponseData data() {
+        return data;
+    }
 
-        Struct broker = (Struct) struct.get(COORDINATOR_KEY_NAME);
-        int nodeId = broker.getInt(NODE_ID_KEY_NAME);
-        String host = broker.getString(HOST_KEY_NAME);
-        int port = broker.getInt(PORT_KEY_NAME);
-        node = new Node(nodeId, host, port);
+    public Node node() {
+        return new Node(data.nodeId(), data.host(), data.port());
     }
 
     @Override
     public int throttleTimeMs() {
-        return throttleTimeMs;
+        return data.throttleTimeMs();
     }
 
     public boolean hasError() {
-        return this.error != Errors.NONE;
+        return error() != Errors.NONE;
     }
 
     public Errors error() {
-        return error;
+        return Errors.forCode(data.errorCode());
     }
 
     @Override
     public Map errorCounts() {
-        return errorCounts(error);
-    }
-
-    public Node node() {
-        return node;
+        return Collections.singletonMap(error(), 1);
     }
 
     @Override
     protected Struct toStruct(short version) {
-        Struct struct = new Struct(ApiKeys.FIND_COORDINATOR.responseSchema(version));
-        struct.setIfExists(THROTTLE_TIME_MS, throttleTimeMs);
-        struct.set(ERROR_CODE, error.code());
-        struct.setIfExists(ERROR_MESSAGE, errorMessage);
-
-        Struct coordinator = struct.instance(COORDINATOR_KEY_NAME);
-        coordinator.set(NODE_ID_KEY_NAME, node.id());
-        coordinator.set(HOST_KEY_NAME, node.host());
-        coordinator.set(PORT_KEY_NAME, node.port());
-        struct.set(COORDINATOR_KEY_NAME, coordinator);
-        return struct;
+        return data.toStruct(version);
     }
 
     public static FindCoordinatorResponse parse(ByteBuffer buffer, short version) {
-        return new FindCoordinatorResponse(ApiKeys.FIND_COORDINATOR.responseSchema(version).read(buffer));
+        return new FindCoordinatorResponse(ApiKeys.FIND_COORDINATOR.responseSchema(version).read(buffer), version);
     }
 
     @Override
     public String toString() {
-        return "FindCoordinatorResponse(" +
-                "throttleTimeMs=" + throttleTimeMs +
-                ", errorMessage='" + errorMessage + '\'' +
-                ", error=" + error +
-                ", node=" + node +
-                ')';
+        return data.toString();
     }
 
     @Override
     public boolean shouldClientThrottle(short version) {
         return version >= 2;
     }
+
+    public static FindCoordinatorResponse prepareResponse(Errors error, Node node) {
+        FindCoordinatorResponseData data = new FindCoordinatorResponseData();
+        data.setErrorCode(error.code())
+            .setErrorMessage(error.message())
+            .setNodeId(node.id())
+            .setHost(node.host())
+            .setPort(node.port());
+        return new FindCoordinatorResponse(data);
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/HeartbeatRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/HeartbeatRequest.java
index 9a131478df723..e5cc6c8860a58 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/HeartbeatRequest.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/HeartbeatRequest.java
@@ -16,108 +16,61 @@
  */
 package org.apache.kafka.common.requests;
 
+import org.apache.kafka.common.errors.UnsupportedVersionException;
+import org.apache.kafka.common.message.HeartbeatRequestData;
+import org.apache.kafka.common.message.HeartbeatResponseData;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.protocol.types.Schema;
 import org.apache.kafka.common.protocol.types.Struct;
 
 import java.nio.ByteBuffer;
 
-import static org.apache.kafka.common.protocol.CommonFields.GENERATION_ID;
-import static org.apache.kafka.common.protocol.CommonFields.GROUP_ID;
-import static org.apache.kafka.common.protocol.CommonFields.MEMBER_ID;
 
 public class HeartbeatRequest extends AbstractRequest {
-    private static final Schema HEARTBEAT_REQUEST_V0 = new Schema(
-            GROUP_ID,
-            GENERATION_ID,
-            MEMBER_ID);
-
-    /* v1 request is the same as v0. Throttle time has been added to response */
-    private static final Schema HEARTBEAT_REQUEST_V1 = HEARTBEAT_REQUEST_V0;
-
-    /**
-     * The version number is bumped to indicate that on quota violation brokers send out responses before throttling.
-     */
-    private static final Schema HEARTBEAT_REQUEST_V2 = HEARTBEAT_REQUEST_V1;
-
-    public static Schema[] schemaVersions() {
-        return new Schema[] {HEARTBEAT_REQUEST_V0, HEARTBEAT_REQUEST_V1,
-            HEARTBEAT_REQUEST_V2};
-    }
 
     public static class Builder extends AbstractRequest.Builder {
-        private final String groupId;
-        private final int groupGenerationId;
-        private final String memberId;
+        private final HeartbeatRequestData data;
 
-        public Builder(String groupId, int groupGenerationId, String memberId) {
+        public Builder(HeartbeatRequestData data) {
             super(ApiKeys.HEARTBEAT);
-            this.groupId = groupId;
-            this.groupGenerationId = groupGenerationId;
-            this.memberId = memberId;
+            this.data = data;
         }
 
         @Override
         public HeartbeatRequest build(short version) {
-            return new HeartbeatRequest(groupId, groupGenerationId, memberId, version);
+            if (data.groupInstanceId() != null && version < 3) {
+                throw new UnsupportedVersionException("The broker heartbeat protocol version " +
+                        version + " does not support usage of config group.instance.id.");
+            }
+            return new HeartbeatRequest(data, version);
         }
 
         @Override
         public String toString() {
-            StringBuilder bld = new StringBuilder();
-            bld.append("(type=HeartbeatRequest").
-                append(", groupId=").append(groupId).
-                append(", groupGenerationId=").append(groupGenerationId).
-                append(", memberId=").append(memberId).
-                append(")");
-            return bld.toString();
+            return data.toString();
         }
     }
 
-    private final String groupId;
-    private final int groupGenerationId;
-    private final String memberId;
+    public final HeartbeatRequestData data;
 
-    private HeartbeatRequest(String groupId, int groupGenerationId, String memberId, short version) {
+    private HeartbeatRequest(HeartbeatRequestData data, short version) {
         super(ApiKeys.HEARTBEAT, version);
-        this.groupId = groupId;
-        this.groupGenerationId = groupGenerationId;
-        this.memberId = memberId;
+        this.data = data;
     }
 
     public HeartbeatRequest(Struct struct, short version) {
         super(ApiKeys.HEARTBEAT, version);
-        groupId = struct.get(GROUP_ID);
-        groupGenerationId = struct.get(GENERATION_ID);
-        memberId = struct.get(MEMBER_ID);
+        this.data = new HeartbeatRequestData(struct, version);
     }
 
     @Override
     public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
-        short versionId = version();
-        switch (versionId) {
-            case 0:
-                return new HeartbeatResponse(Errors.forException(e));
-            case 1:
-            case 2:
-                return new HeartbeatResponse(throttleTimeMs, Errors.forException(e));
-            default:
-                throw new IllegalArgumentException(String.format("Version %d is not valid. Valid versions for %s are 0 to %d",
-                        versionId, this.getClass().getSimpleName(), ApiKeys.HEARTBEAT.latestVersion()));
+        HeartbeatResponseData responseData = new HeartbeatResponseData().
+            setErrorCode(Errors.forException(e).code());
+        if (version() >= 1) {
+            responseData.setThrottleTimeMs(throttleTimeMs);
         }
-    }
-
-    public String groupId() {
-        return groupId;
-    }
-
-    public int groupGenerationId() {
-        return groupGenerationId;
-    }
-
-    public String memberId() {
-        return memberId;
+        return new HeartbeatResponse(responseData);
     }
 
     public static HeartbeatRequest parse(ByteBuffer buffer, short version) {
@@ -126,10 +79,6 @@ public static HeartbeatRequest parse(ByteBuffer buffer, short version) {
 
     @Override
     protected Struct toStruct() {
-        Struct struct = new Struct(ApiKeys.HEARTBEAT.requestSchema(version()));
-        struct.set(GROUP_ID, groupId);
-        struct.set(GENERATION_ID, groupGenerationId);
-        struct.set(MEMBER_ID, memberId);
-        return struct;
+        return data.toStruct(version());
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/HeartbeatResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/HeartbeatResponse.java
index efe2ed89f2929..cc36a2037ae08 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/HeartbeatResponse.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/HeartbeatResponse.java
@@ -16,35 +16,18 @@
  */
 package org.apache.kafka.common.requests;
 
+import org.apache.kafka.common.message.HeartbeatResponseData;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.protocol.types.Schema;
 import org.apache.kafka.common.protocol.types.Struct;
 
 import java.nio.ByteBuffer;
+import java.util.Collections;
 import java.util.Map;
 
-import static org.apache.kafka.common.protocol.CommonFields.ERROR_CODE;
-import static org.apache.kafka.common.protocol.CommonFields.THROTTLE_TIME_MS;
 
 public class HeartbeatResponse extends AbstractResponse {
 
-    private static final Schema HEARTBEAT_RESPONSE_V0 = new Schema(
-            ERROR_CODE);
-    private static final Schema HEARTBEAT_RESPONSE_V1 = new Schema(
-            THROTTLE_TIME_MS,
-            ERROR_CODE);
-
-    /**
-     * The version number is bumped to indicate that on quota violation brokers send out responses before throttling.
-     */
-    private static final Schema HEARTBEAT_RESPONSE_V2 = HEARTBEAT_RESPONSE_V1;
-
-    public static Schema[] schemaVersions() {
-        return new Schema[] {HEARTBEAT_RESPONSE_V0, HEARTBEAT_RESPONSE_V1,
-            HEARTBEAT_RESPONSE_V2};
-    }
-
     /**
      * Possible error codes:
      *
@@ -55,47 +38,37 @@ public static Schema[] schemaVersions() {
      * REBALANCE_IN_PROGRESS (27)
      * GROUP_AUTHORIZATION_FAILED (30)
      */
-    private final Errors error;
-    private final int throttleTimeMs;
-
-    public HeartbeatResponse(Errors error) {
-        this(DEFAULT_THROTTLE_TIME, error);
-    }
+    private final HeartbeatResponseData data;
 
-    public HeartbeatResponse(int throttleTimeMs, Errors error) {
-        this.throttleTimeMs = throttleTimeMs;
-        this.error = error;
+    public HeartbeatResponse(HeartbeatResponseData data) {
+        this.data = data;
     }
 
-    public HeartbeatResponse(Struct struct) {
-        this.throttleTimeMs = struct.getOrElse(THROTTLE_TIME_MS, DEFAULT_THROTTLE_TIME);
-        error = Errors.forCode(struct.get(ERROR_CODE));
+    public HeartbeatResponse(Struct struct, short version) {
+        this.data = new HeartbeatResponseData(struct, version);
     }
 
     @Override
     public int throttleTimeMs() {
-        return throttleTimeMs;
+        return data.throttleTimeMs();
     }
 
     public Errors error() {
-        return error;
+        return Errors.forCode(data.errorCode());
     }
 
     @Override
     public Map errorCounts() {
-        return errorCounts(error);
+        return Collections.singletonMap(error(), 1);
     }
 
     @Override
     protected Struct toStruct(short version) {
-        Struct struct = new Struct(ApiKeys.HEARTBEAT.responseSchema(version));
-        struct.setIfExists(THROTTLE_TIME_MS, throttleTimeMs);
-        struct.set(ERROR_CODE, error.code());
-        return struct;
+        return data.toStruct(version);
     }
 
     public static HeartbeatResponse parse(ByteBuffer buffer, short version) {
-        return new HeartbeatResponse(ApiKeys.HEARTBEAT.parseResponse(version, buffer));
+        return new HeartbeatResponse(ApiKeys.HEARTBEAT.parseResponse(version, buffer), version);
     }
 
     @Override
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/IncrementalAlterConfigsRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/IncrementalAlterConfigsRequest.java
new file mode 100644
index 0000000000000..0c6c0b2814730
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/IncrementalAlterConfigsRequest.java
@@ -0,0 +1,91 @@
+/*
+ * 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.requests;
+
+import org.apache.kafka.common.message.IncrementalAlterConfigsRequestData;
+import org.apache.kafka.common.message.IncrementalAlterConfigsRequestData.AlterConfigsResource;
+import org.apache.kafka.common.message.IncrementalAlterConfigsResponseData;
+import org.apache.kafka.common.message.IncrementalAlterConfigsResponseData.AlterConfigsResourceResponse;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.types.Struct;
+
+import java.nio.ByteBuffer;
+
+public class IncrementalAlterConfigsRequest extends AbstractRequest {
+
+    public static class Builder extends AbstractRequest.Builder {
+        private final IncrementalAlterConfigsRequestData data;
+
+        public Builder(IncrementalAlterConfigsRequestData data) {
+            super(ApiKeys.INCREMENTAL_ALTER_CONFIGS);
+            this.data = data;
+        }
+
+        @Override
+        public IncrementalAlterConfigsRequest build(short version) {
+            return new IncrementalAlterConfigsRequest(data, version);
+        }
+
+        @Override
+        public String toString() {
+            return data.toString();
+        }
+    }
+
+    private final IncrementalAlterConfigsRequestData data;
+    private final short version;
+
+    private IncrementalAlterConfigsRequest(IncrementalAlterConfigsRequestData data, short version) {
+        super(ApiKeys.INCREMENTAL_ALTER_CONFIGS, version);
+        this.data = data;
+        this.version = version;
+    }
+
+    IncrementalAlterConfigsRequest(final Struct struct, final short version) {
+        super(ApiKeys.INCREMENTAL_ALTER_CONFIGS, version);
+        this.data = new IncrementalAlterConfigsRequestData(struct, version);
+        this.version = version;
+    }
+
+    public static IncrementalAlterConfigsRequest parse(ByteBuffer buffer, short version) {
+        return new IncrementalAlterConfigsRequest(ApiKeys.INCREMENTAL_ALTER_CONFIGS.parseRequest(version, buffer), version);
+    }
+
+    public IncrementalAlterConfigsRequestData data() {
+        return data;
+    }
+
+    @Override
+    protected Struct toStruct() {
+        return data.toStruct(version);
+    }
+
+    @Override
+    public AbstractResponse getErrorResponse(final int throttleTimeMs, final Throwable e) {
+        IncrementalAlterConfigsResponseData response = new IncrementalAlterConfigsResponseData();
+        ApiError apiError = ApiError.fromThrowable(e);
+        for (AlterConfigsResource resource : data.resources()) {
+            response.responses().add(new AlterConfigsResourceResponse()
+                    .setResourceName(resource.resourceName())
+                    .setResourceType(resource.resourceType())
+                    .setErrorCode(apiError.error().code())
+                    .setErrorMessage(apiError.message()));
+        }
+        return new IncrementalAlterConfigsResponse(response);
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/IncrementalAlterConfigsResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/IncrementalAlterConfigsResponse.java
new file mode 100644
index 0000000000000..46b1d530f5635
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/IncrementalAlterConfigsResponse.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.common.requests;
+
+import org.apache.kafka.common.config.ConfigResource;
+import org.apache.kafka.common.message.IncrementalAlterConfigsResponseData;
+import org.apache.kafka.common.message.IncrementalAlterConfigsResponseData.AlterConfigsResourceResponse;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.common.protocol.types.Struct;
+
+import java.nio.ByteBuffer;
+import java.util.HashMap;
+import java.util.Map;
+
+public class IncrementalAlterConfigsResponse extends AbstractResponse {
+
+    public static IncrementalAlterConfigsResponseData toResponseData(final int requestThrottleMs,
+                                                                     final Map results) {
+        IncrementalAlterConfigsResponseData responseData = new IncrementalAlterConfigsResponseData();
+        responseData.setThrottleTimeMs(requestThrottleMs);
+        for (Map.Entry entry : results.entrySet()) {
+            responseData.responses().add(new AlterConfigsResourceResponse().
+                    setResourceName(entry.getKey().name()).
+                    setResourceType(entry.getKey().type().id()).
+                    setErrorCode(entry.getValue().error().code()).
+                    setErrorMessage(entry.getValue().message()));
+        }
+        return responseData;
+    }
+
+    public static Map fromResponseData(final IncrementalAlterConfigsResponseData data) {
+        Map map = new HashMap<>();
+        for (AlterConfigsResourceResponse response : data.responses()) {
+            map.put(new ConfigResource(ConfigResource.Type.forId(response.resourceType()), response.resourceName()),
+                    new ApiError(Errors.forCode(response.errorCode()), response.errorMessage()));
+        }
+        return map;
+    }
+
+    private final IncrementalAlterConfigsResponseData data;
+
+    public IncrementalAlterConfigsResponse(IncrementalAlterConfigsResponseData data) {
+        this.data = data;
+    }
+
+    public IncrementalAlterConfigsResponse(final Struct struct, final short version) {
+        this.data = new IncrementalAlterConfigsResponseData(struct, version);
+    }
+
+    public IncrementalAlterConfigsResponseData data() {
+        return data;
+    }
+
+    @Override
+    public Map errorCounts() {
+        HashMap counts = new HashMap<>();
+        for (AlterConfigsResourceResponse response : data.responses()) {
+            Errors error = Errors.forCode(response.errorCode());
+            counts.put(error, counts.getOrDefault(error, 0) + 1);
+        }
+        return counts;
+    }
+
+    @Override
+    protected Struct toStruct(final short version) {
+        return data.toStruct(version);
+    }
+
+    @Override
+    public boolean shouldClientThrottle(short version) {
+        return version >= 0;
+    }
+
+    @Override
+    public int throttleTimeMs() {
+        return data.throttleTimeMs();
+    }
+
+    public static IncrementalAlterConfigsResponse parse(ByteBuffer buffer, short version) {
+        return new IncrementalAlterConfigsResponse(
+                ApiKeys.INCREMENTAL_ALTER_CONFIGS.responseSchema(version).read(buffer), version);
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/InitProducerIdRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/InitProducerIdRequest.java
index aab7c72cf430e..8351c2159836a 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/InitProducerIdRequest.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/InitProducerIdRequest.java
@@ -16,106 +16,71 @@
  */
 package org.apache.kafka.common.requests;
 
+import org.apache.kafka.common.message.InitProducerIdRequestData;
+import org.apache.kafka.common.message.InitProducerIdResponseData;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.protocol.types.Field;
-import org.apache.kafka.common.protocol.types.Schema;
 import org.apache.kafka.common.protocol.types.Struct;
+import org.apache.kafka.common.record.RecordBatch;
 
 import java.nio.ByteBuffer;
 
-import static org.apache.kafka.common.protocol.CommonFields.NULLABLE_TRANSACTIONAL_ID;
-import static org.apache.kafka.common.protocol.types.Type.INT32;
-
 public class InitProducerIdRequest extends AbstractRequest {
-    public static final int NO_TRANSACTION_TIMEOUT_MS = Integer.MAX_VALUE;
-
-    private static final String TRANSACTION_TIMEOUT_KEY_NAME = "transaction_timeout_ms";
-
-    private static final Schema INIT_PRODUCER_ID_REQUEST_V0 = new Schema(
-            NULLABLE_TRANSACTIONAL_ID,
-            new Field(TRANSACTION_TIMEOUT_KEY_NAME, INT32, "The time in ms to wait for before aborting idle transactions sent by this producer."));
-
-    /**
-     * The version number is bumped to indicate that on quota violation brokers send out responses before throttling.
-     */
-    private static final Schema INIT_PRODUCER_ID_REQUEST_V1 = INIT_PRODUCER_ID_REQUEST_V0;
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{INIT_PRODUCER_ID_REQUEST_V0, INIT_PRODUCER_ID_REQUEST_V1};
-    }
-
-    private final String transactionalId;
-    private final int transactionTimeoutMs;
-
     public static class Builder extends AbstractRequest.Builder {
-        private final String transactionalId;
-        private final int transactionTimeoutMs;
-
-        public Builder(String transactionalId) {
-            this(transactionalId, NO_TRANSACTION_TIMEOUT_MS);
-        }
+        private final InitProducerIdRequestData data;
 
-        public Builder(String transactionalId, int transactionTimeoutMs) {
+        public Builder(InitProducerIdRequestData data) {
             super(ApiKeys.INIT_PRODUCER_ID);
-
-            if (transactionTimeoutMs <= 0)
-                throw new IllegalArgumentException("transaction timeout value is not positive: " + transactionTimeoutMs);
-
-            if (transactionalId != null && transactionalId.isEmpty())
-                throw new IllegalArgumentException("Must set either a null or a non-empty transactional id.");
-
-            this.transactionalId = transactionalId;
-            this.transactionTimeoutMs = transactionTimeoutMs;
+            this.data = data;
         }
 
         @Override
         public InitProducerIdRequest build(short version) {
-            return new InitProducerIdRequest(version, transactionalId, transactionTimeoutMs);
+            if (data.transactionTimeoutMs() <= 0)
+                throw new IllegalArgumentException("transaction timeout value is not positive: " + data.transactionTimeoutMs());
+
+            if (data.transactionalId() != null && data.transactionalId().isEmpty())
+                throw new IllegalArgumentException("Must set either a null or a non-empty transactional id.");
+
+            return new InitProducerIdRequest(data, version);
         }
 
         @Override
         public String toString() {
-            return "(type=InitProducerIdRequest, transactionalId=" + transactionalId + ", transactionTimeoutMs=" +
-                    transactionTimeoutMs + ")";
+            return data.toString();
         }
     }
 
-    public InitProducerIdRequest(Struct struct, short version) {
+    public final InitProducerIdRequestData data;
+
+    private InitProducerIdRequest(InitProducerIdRequestData data, short version) {
         super(ApiKeys.INIT_PRODUCER_ID, version);
-        this.transactionalId = struct.get(NULLABLE_TRANSACTIONAL_ID);
-        this.transactionTimeoutMs = struct.getInt(TRANSACTION_TIMEOUT_KEY_NAME);
+        this.data = data;
     }
 
-    private InitProducerIdRequest(short version, String transactionalId, int transactionTimeoutMs) {
+    public InitProducerIdRequest(Struct struct, short version) {
         super(ApiKeys.INIT_PRODUCER_ID, version);
-        this.transactionalId = transactionalId;
-        this.transactionTimeoutMs = transactionTimeoutMs;
+        this.data = new InitProducerIdRequestData(struct, version);
     }
 
+
     @Override
     public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
-        return new InitProducerIdResponse(throttleTimeMs, Errors.forException(e));
+        InitProducerIdResponseData response = new InitProducerIdResponseData()
+                .setErrorCode(Errors.forException(e).code())
+                .setProducerId(RecordBatch.NO_PRODUCER_ID)
+                .setProducerEpoch(RecordBatch.NO_PRODUCER_EPOCH)
+                .setThrottleTimeMs(0);
+        return new InitProducerIdResponse(response);
     }
 
     public static InitProducerIdRequest parse(ByteBuffer buffer, short version) {
         return new InitProducerIdRequest(ApiKeys.INIT_PRODUCER_ID.parseRequest(version, buffer), version);
     }
 
-    public String transactionalId() {
-        return transactionalId;
-    }
-
-    public int transactionTimeoutMs() {
-        return transactionTimeoutMs;
-    }
-
     @Override
     protected Struct toStruct() {
-        Struct struct = new Struct(ApiKeys.INIT_PRODUCER_ID.requestSchema(version()));
-        struct.set(NULLABLE_TRANSACTIONAL_ID, transactionalId);
-        struct.set(TRANSACTION_TIMEOUT_KEY_NAME, transactionTimeoutMs);
-        return struct;
+        return data.toStruct(version());
     }
 
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/InitProducerIdResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/InitProducerIdResponse.java
index 9a1e0f744cd64..a33daf3e2d220 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/InitProducerIdResponse.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/InitProducerIdResponse.java
@@ -16,110 +16,59 @@
  */
 package org.apache.kafka.common.requests;
 
+import org.apache.kafka.common.message.InitProducerIdResponseData;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.protocol.types.Schema;
 import org.apache.kafka.common.protocol.types.Struct;
-import org.apache.kafka.common.record.RecordBatch;
 
 import java.nio.ByteBuffer;
 import java.util.Map;
 
-import static org.apache.kafka.common.protocol.CommonFields.ERROR_CODE;
-import static org.apache.kafka.common.protocol.CommonFields.PRODUCER_EPOCH;
-import static org.apache.kafka.common.protocol.CommonFields.PRODUCER_ID;
-import static org.apache.kafka.common.protocol.CommonFields.THROTTLE_TIME_MS;
-
+/**
+ * Possible error codes:
+ * - {@link Errors#NOT_COORDINATOR}
+ * - {@link Errors#COORDINATOR_NOT_AVAILABLE}
+ * - {@link Errors#COORDINATOR_LOAD_IN_PROGRESS}
+ * - {@link Errors#TRANSACTIONAL_ID_AUTHORIZATION_FAILED}
+ * - {@link Errors#CLUSTER_AUTHORIZATION_FAILED}
+ */
 public class InitProducerIdResponse extends AbstractResponse {
-    // Possible error codes:
-    //   NotCoordinator
-    //   CoordinatorNotAvailable
-    //   CoordinatorLoadInProgress
-    //   TransactionalIdAuthorizationFailed
-    //   ClusterAuthorizationFailed
-
-    private static final Schema INIT_PRODUCER_ID_RESPONSE_V0 = new Schema(
-            THROTTLE_TIME_MS,
-            ERROR_CODE,
-            PRODUCER_ID,
-            PRODUCER_EPOCH);
-
-    /**
-     * The version number is bumped to indicate that on quota violation brokers send out responses before throttling.
-     */
-    private static final Schema INIT_PRODUCER_ID_RESPONSE_V1 = INIT_PRODUCER_ID_RESPONSE_V0;
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{INIT_PRODUCER_ID_RESPONSE_V0, INIT_PRODUCER_ID_RESPONSE_V1};
-    }
-
-    private final int throttleTimeMs;
-    private final Errors error;
-    private final long producerId;
-    private final short epoch;
+    public final InitProducerIdResponseData data;
 
-    public InitProducerIdResponse(int throttleTimeMs, Errors error, long producerId, short epoch) {
-        this.throttleTimeMs = throttleTimeMs;
-        this.error = error;
-        this.producerId = producerId;
-        this.epoch = epoch;
+    public InitProducerIdResponse(InitProducerIdResponseData data) {
+        this.data = data;
     }
 
-    public InitProducerIdResponse(Struct struct) {
-        this.throttleTimeMs = struct.get(THROTTLE_TIME_MS);
-        this.error = Errors.forCode(struct.get(ERROR_CODE));
-        this.producerId = struct.get(PRODUCER_ID);
-        this.epoch = struct.get(PRODUCER_EPOCH);
-    }
-
-    public InitProducerIdResponse(int throttleTimeMs, Errors errors) {
-        this(throttleTimeMs, errors, RecordBatch.NO_PRODUCER_ID, (short) 0);
+    public InitProducerIdResponse(Struct struct, short version) {
+        this.data = new InitProducerIdResponseData(struct, version);
     }
 
     @Override
     public int throttleTimeMs() {
-        return throttleTimeMs;
-    }
-
-    public long producerId() {
-        return producerId;
-    }
-
-    public Errors error() {
-        return error;
+        return data.throttleTimeMs();
     }
 
     @Override
     public Map errorCounts() {
-        return errorCounts(error);
-    }
-
-    public short epoch() {
-        return epoch;
+        return errorCounts(Errors.forCode(data.errorCode()));
     }
 
     @Override
     protected Struct toStruct(short version) {
-        Struct struct = new Struct(ApiKeys.INIT_PRODUCER_ID.responseSchema(version));
-        struct.set(THROTTLE_TIME_MS, throttleTimeMs);
-        struct.set(PRODUCER_ID, producerId);
-        struct.set(PRODUCER_EPOCH, epoch);
-        struct.set(ERROR_CODE, error.code());
-        return struct;
+        return data.toStruct(version);
     }
 
     public static InitProducerIdResponse parse(ByteBuffer buffer, short version) {
-        return new InitProducerIdResponse(ApiKeys.INIT_PRODUCER_ID.parseResponse(version, buffer));
+        return new InitProducerIdResponse(ApiKeys.INIT_PRODUCER_ID.parseResponse(version, buffer), version);
     }
 
     @Override
     public String toString() {
-        return "InitProducerIdResponse(" +
-                "error=" + error +
-                ", producerId=" + producerId +
-                ", producerEpoch=" + epoch +
-                ", throttleTimeMs=" + throttleTimeMs +
-                ')';
+        return data.toString();
+    }
+
+    public Errors error() {
+        return Errors.forCode(data.errorCode());
     }
 
     @Override
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/JoinGroupRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/JoinGroupRequest.java
index 236d6848f53aa..95d125d312a17 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/JoinGroupRequest.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/JoinGroupRequest.java
@@ -16,190 +16,105 @@
  */
 package org.apache.kafka.common.requests;
 
+import org.apache.kafka.common.errors.InvalidConfigurationException;
+import org.apache.kafka.common.errors.UnsupportedVersionException;
+import org.apache.kafka.common.message.JoinGroupRequestData;
+import org.apache.kafka.common.message.JoinGroupResponseData;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.protocol.types.ArrayOf;
-import org.apache.kafka.common.protocol.types.Field;
-import org.apache.kafka.common.protocol.types.Schema;
 import org.apache.kafka.common.protocol.types.Struct;
-import org.apache.kafka.common.utils.Utils;
 
 import java.nio.ByteBuffer;
-import java.util.ArrayList;
 import java.util.Collections;
-import java.util.List;
-
-import static org.apache.kafka.common.protocol.CommonFields.GROUP_ID;
-import static org.apache.kafka.common.protocol.CommonFields.MEMBER_ID;
-import static org.apache.kafka.common.protocol.types.Type.BYTES;
-import static org.apache.kafka.common.protocol.types.Type.INT32;
-import static org.apache.kafka.common.protocol.types.Type.STRING;
 
 public class JoinGroupRequest extends AbstractRequest {
-    private static final String SESSION_TIMEOUT_KEY_NAME = "session_timeout";
-    private static final String REBALANCE_TIMEOUT_KEY_NAME = "rebalance_timeout";
-    private static final String PROTOCOL_TYPE_KEY_NAME = "protocol_type";
-    private static final String GROUP_PROTOCOLS_KEY_NAME = "group_protocols";
-    private static final String PROTOCOL_NAME_KEY_NAME = "protocol_name";
-    private static final String PROTOCOL_METADATA_KEY_NAME = "protocol_metadata";
-
-    /* Join group api */
-    private static final Schema JOIN_GROUP_REQUEST_PROTOCOL_V0 = new Schema(
-            new Field(PROTOCOL_NAME_KEY_NAME, STRING),
-            new Field(PROTOCOL_METADATA_KEY_NAME, BYTES));
-
-    private static final Schema JOIN_GROUP_REQUEST_V0 = new Schema(
-            GROUP_ID,
-            new Field(SESSION_TIMEOUT_KEY_NAME, INT32, "The coordinator considers the consumer dead if it receives " +
-                    "no heartbeat after this timeout in ms."),
-            MEMBER_ID,
-            new Field(PROTOCOL_TYPE_KEY_NAME, STRING, "Unique name for class of protocols implemented by group"),
-            new Field(GROUP_PROTOCOLS_KEY_NAME, new ArrayOf(JOIN_GROUP_REQUEST_PROTOCOL_V0), "List of protocols " +
-                    "that the member supports"));
-
-    private static final Schema JOIN_GROUP_REQUEST_V1 = new Schema(
-            GROUP_ID,
-            new Field(SESSION_TIMEOUT_KEY_NAME, INT32, "The coordinator considers the consumer dead if it receives no " +
-                    "heartbeat after this timeout in ms."),
-            new Field(REBALANCE_TIMEOUT_KEY_NAME, INT32, "The maximum time that the coordinator will wait for each " +
-                    "member to rejoin when rebalancing the group"),
-            MEMBER_ID,
-            new Field(PROTOCOL_TYPE_KEY_NAME, STRING, "Unique name for class of protocols implemented by group"),
-            new Field(GROUP_PROTOCOLS_KEY_NAME, new ArrayOf(JOIN_GROUP_REQUEST_PROTOCOL_V0), "List of protocols " +
-                    "that the member supports"));
-
-    /* v2 request is the same as v1. Throttle time has been added to response */
-    private static final Schema JOIN_GROUP_REQUEST_V2 = JOIN_GROUP_REQUEST_V1;
-
-    /**
-     * The version number is bumped to indicate that on quota violation brokers send out responses before throttling.
-     */
-    private static final Schema JOIN_GROUP_REQUEST_V3 = JOIN_GROUP_REQUEST_V2;
-
-    /**
-     * The version number is bumped to indicate that client needs to issue a second join group request under first try
-     * with UNKNOWN_MEMBER_ID.
-     */
-    private static final Schema JOIN_GROUP_REQUEST_V4 = JOIN_GROUP_REQUEST_V3;
-
-    public static Schema[] schemaVersions() {
-        return new Schema[] {JOIN_GROUP_REQUEST_V0, JOIN_GROUP_REQUEST_V1, JOIN_GROUP_REQUEST_V2,
-            JOIN_GROUP_REQUEST_V3, JOIN_GROUP_REQUEST_V4};
-    }
-
-    public static final String UNKNOWN_MEMBER_ID = "";
-
-    private final String groupId;
-    private final int sessionTimeout;
-    private final int rebalanceTimeout;
-    private final String memberId;
-    private final String protocolType;
-    private final List groupProtocols;
-
-    public static class ProtocolMetadata {
-        private final String name;
-        private final ByteBuffer metadata;
-
-        public ProtocolMetadata(String name, ByteBuffer metadata) {
-            this.name = name;
-            this.metadata = metadata;
-        }
-
-        public String name() {
-            return name;
-        }
-
-        public ByteBuffer metadata() {
-            return metadata;
-        }
-    }
 
     public static class Builder extends AbstractRequest.Builder {
-        private final String groupId;
-        private final int sessionTimeout;
-        private final String memberId;
-        private final String protocolType;
-        private final List groupProtocols;
-        private int rebalanceTimeout = 0;
 
-        public Builder(String groupId, int sessionTimeout, String memberId,
-                       String protocolType, List groupProtocols) {
-            super(ApiKeys.JOIN_GROUP);
-            this.groupId = groupId;
-            this.sessionTimeout = sessionTimeout;
-            this.rebalanceTimeout = sessionTimeout;
-            this.memberId = memberId;
-            this.protocolType = protocolType;
-            this.groupProtocols = groupProtocols;
-        }
+        private final JoinGroupRequestData data;
 
-        public Builder setRebalanceTimeout(int rebalanceTimeout) {
-            this.rebalanceTimeout = rebalanceTimeout;
-            return this;
+        public Builder(JoinGroupRequestData data) {
+            super(ApiKeys.JOIN_GROUP);
+            this.data = data;
         }
 
         @Override
         public JoinGroupRequest build(short version) {
-            if (version < 1) {
-                // v0 had no rebalance timeout but used session timeout implicitly
-                rebalanceTimeout = sessionTimeout;
+            if (data.groupInstanceId() != null && version < 5) {
+                throw new UnsupportedVersionException("The broker join group protocol version " +
+                        version + " does not support usage of config group.instance.id.");
             }
-            return new JoinGroupRequest(version, groupId, sessionTimeout,
-                    rebalanceTimeout, memberId, protocolType, groupProtocols);
+            return new JoinGroupRequest(data, version);
         }
 
         @Override
         public String toString() {
-            StringBuilder bld = new StringBuilder();
-            bld.append("(type: JoinGroupRequest").
-                append(", groupId=").append(groupId).
-                append(", sessionTimeout=").append(sessionTimeout).
-                append(", rebalanceTimeout=").append(rebalanceTimeout).
-                append(", memberId=").append(memberId).
-                append(", protocolType=").append(protocolType).
-                append(", groupProtocols=").append(Utils.join(groupProtocols, ", ")).
-                append(")");
-            return bld.toString();
+            return data.toString();
         }
     }
 
-    private JoinGroupRequest(short version, String groupId, int sessionTimeout,
-            int rebalanceTimeout, String memberId, String protocolType,
-            List groupProtocols) {
-        super(ApiKeys.JOIN_GROUP, version);
-        this.groupId = groupId;
-        this.sessionTimeout = sessionTimeout;
-        this.rebalanceTimeout = rebalanceTimeout;
-        this.memberId = memberId;
-        this.protocolType = protocolType;
-        this.groupProtocols = groupProtocols;
-    }
+    private final JoinGroupRequestData data;
+
+    public static final String UNKNOWN_MEMBER_ID = "";
 
-    public JoinGroupRequest(Struct struct, short versionId) {
-        super(ApiKeys.JOIN_GROUP, versionId);
+    private static final int MAX_GROUP_INSTANCE_ID_LENGTH = 249;
 
-        groupId = struct.get(GROUP_ID);
-        sessionTimeout = struct.getInt(SESSION_TIMEOUT_KEY_NAME);
+    /**
+     * Ported from class Topic in {@link org.apache.kafka.common.internals} to restrict the charset for
+     * static member id.
+     */
+    public static void validateGroupInstanceId(String id) {
+        if (id.equals(""))
+            throw new InvalidConfigurationException("Group instance id must be non-empty string");
+        if (id.equals(".") || id.equals(".."))
+            throw new InvalidConfigurationException("Group instance id cannot be \".\" or \"..\"");
+        if (id.length() > MAX_GROUP_INSTANCE_ID_LENGTH)
+            throw new InvalidConfigurationException("Group instance id can't be longer than " + MAX_GROUP_INSTANCE_ID_LENGTH +
+                    " characters: " + id);
+        if (!containsValidPattern(id))
+            throw new InvalidConfigurationException("Group instance id \"" + id + "\" is illegal, it contains a character other than " +
+                    "ASCII alphanumerics, '.', '_' and '-'");
+    }
 
-        if (struct.hasField(REBALANCE_TIMEOUT_KEY_NAME))
-            // rebalance timeout is added in v1
-            rebalanceTimeout = struct.getInt(REBALANCE_TIMEOUT_KEY_NAME);
-        else
-            // v0 had no rebalance timeout but used session timeout implicitly
-            rebalanceTimeout = sessionTimeout;
+    /**
+     * Valid characters for Consumer group.instance.id are the ASCII alphanumerics, '.', '_', and '-'
+     */
+    static boolean containsValidPattern(String topic) {
+        for (int i = 0; i < topic.length(); ++i) {
+            char c = topic.charAt(i);
+
+            boolean validChar = (c >= 'a' && c <= 'z') || (c >= '0' && c <= '9') || (c >= 'A' && c <= 'Z') || c == '.' ||
+                    c == '_' || c == '-';
+            if (!validChar)
+                return false;
+        }
+        return true;
+    }
 
-        memberId = struct.get(MEMBER_ID);
-        protocolType = struct.getString(PROTOCOL_TYPE_KEY_NAME);
+    public JoinGroupRequest(JoinGroupRequestData data, short version) {
+        super(ApiKeys.JOIN_GROUP, version);
+        this.data = data;
+        maybeOverrideRebalanceTimeout(version);
+    }
 
-        groupProtocols = new ArrayList<>();
-        for (Object groupProtocolObj : struct.getArray(GROUP_PROTOCOLS_KEY_NAME)) {
-            Struct groupProtocolStruct = (Struct) groupProtocolObj;
-            String name = groupProtocolStruct.getString(PROTOCOL_NAME_KEY_NAME);
-            ByteBuffer metadata = groupProtocolStruct.getBytes(PROTOCOL_METADATA_KEY_NAME);
-            groupProtocols.add(new ProtocolMetadata(name, metadata));
+    public JoinGroupRequest(Struct struct, short version) {
+        super(ApiKeys.JOIN_GROUP, version);
+        this.data = new JoinGroupRequestData(struct, version);
+        maybeOverrideRebalanceTimeout(version);
+    }
+
+    private void maybeOverrideRebalanceTimeout(short version) {
+        if (version == 0) {
+            // Version 0 has no rebalance timeout, so we use the session timeout
+            // to be consistent with the original behavior of the API.
+            data.setRebalanceTimeoutMs(data.sessionTimeoutMs());
         }
     }
 
+    public JoinGroupRequestData data() {
+        return data;
+    }
+
     @Override
     public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
         short versionId = version();
@@ -207,77 +122,40 @@ public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
             case 0:
             case 1:
                 return new JoinGroupResponse(
-                    Errors.forException(e),
-                    JoinGroupResponse.UNKNOWN_GENERATION_ID,
-                    JoinGroupResponse.UNKNOWN_PROTOCOL,
-                    JoinGroupResponse.UNKNOWN_MEMBER_ID, // memberId
-                    JoinGroupResponse.UNKNOWN_MEMBER_ID, // leaderId
-                    Collections.emptyMap());
+                        new JoinGroupResponseData()
+                                .setErrorCode(Errors.forException(e).code())
+                                .setGenerationId(JoinGroupResponse.UNKNOWN_GENERATION_ID)
+                                .setProtocolName(JoinGroupResponse.UNKNOWN_PROTOCOL)
+                                .setLeader(JoinGroupResponse.UNKNOWN_MEMBER_ID)
+                                .setMemberId(JoinGroupResponse.UNKNOWN_MEMBER_ID)
+                                .setMembers(Collections.emptyList())
+                );
             case 2:
             case 3:
             case 4:
+            case 5:
                 return new JoinGroupResponse(
-                    throttleTimeMs,
-                    Errors.forException(e),
-                    JoinGroupResponse.UNKNOWN_GENERATION_ID,
-                    JoinGroupResponse.UNKNOWN_PROTOCOL,
-                    JoinGroupResponse.UNKNOWN_MEMBER_ID, // memberId
-                    JoinGroupResponse.UNKNOWN_MEMBER_ID, // leaderId
-                    Collections.emptyMap());
-
+                        new JoinGroupResponseData()
+                                .setThrottleTimeMs(throttleTimeMs)
+                                .setErrorCode(Errors.forException(e).code())
+                                .setGenerationId(JoinGroupResponse.UNKNOWN_GENERATION_ID)
+                                .setProtocolName(JoinGroupResponse.UNKNOWN_PROTOCOL)
+                                .setLeader(JoinGroupResponse.UNKNOWN_MEMBER_ID)
+                                .setMemberId(JoinGroupResponse.UNKNOWN_MEMBER_ID)
+                                .setMembers(Collections.emptyList())
+                );
             default:
                 throw new IllegalArgumentException(String.format("Version %d is not valid. Valid versions for %s are 0 to %d",
                         versionId, this.getClass().getSimpleName(), ApiKeys.JOIN_GROUP.latestVersion()));
         }
     }
 
-    public String groupId() {
-        return groupId;
-    }
-
-    public int sessionTimeout() {
-        return sessionTimeout;
-    }
-
-    public int rebalanceTimeout() {
-        return rebalanceTimeout;
-    }
-
-    public String memberId() {
-        return memberId;
-    }
-
-    public List groupProtocols() {
-        return groupProtocols;
-    }
-
-    public String protocolType() {
-        return protocolType;
-    }
-
     public static JoinGroupRequest parse(ByteBuffer buffer, short version) {
         return new JoinGroupRequest(ApiKeys.JOIN_GROUP.parseRequest(version, buffer), version);
     }
 
     @Override
     protected Struct toStruct() {
-        short version = version();
-        Struct struct = new Struct(ApiKeys.JOIN_GROUP.requestSchema(version));
-        struct.set(GROUP_ID, groupId);
-        struct.set(SESSION_TIMEOUT_KEY_NAME, sessionTimeout);
-        if (version >= 1) {
-            struct.set(REBALANCE_TIMEOUT_KEY_NAME, rebalanceTimeout);
-        }
-        struct.set(MEMBER_ID, memberId);
-        struct.set(PROTOCOL_TYPE_KEY_NAME, protocolType);
-        List groupProtocolsList = new ArrayList<>(groupProtocols.size());
-        for (ProtocolMetadata protocol : groupProtocols) {
-            Struct protocolStruct = struct.instance(GROUP_PROTOCOLS_KEY_NAME);
-            protocolStruct.set(PROTOCOL_NAME_KEY_NAME, protocol.name);
-            protocolStruct.set(PROTOCOL_METADATA_KEY_NAME, protocol.metadata);
-            groupProtocolsList.add(protocolStruct);
-        }
-        struct.set(GROUP_PROTOCOLS_KEY_NAME, groupProtocolsList.toArray());
-        return struct;
+        return data.toStruct(version());
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/JoinGroupResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/JoinGroupResponse.java
index 55cb97f01c512..0e1644e465923 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/JoinGroupResponse.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/JoinGroupResponse.java
@@ -16,217 +16,70 @@
  */
 package org.apache.kafka.common.requests;
 
+import org.apache.kafka.common.message.JoinGroupResponseData;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.protocol.types.ArrayOf;
-import org.apache.kafka.common.protocol.types.Field;
-import org.apache.kafka.common.protocol.types.Schema;
 import org.apache.kafka.common.protocol.types.Struct;
-import org.apache.kafka.common.utils.Utils;
 
 import java.nio.ByteBuffer;
-import java.util.ArrayList;
-import java.util.HashMap;
-import java.util.List;
+import java.util.Collections;
 import java.util.Map;
 
-import static org.apache.kafka.common.protocol.CommonFields.ERROR_CODE;
-import static org.apache.kafka.common.protocol.CommonFields.GENERATION_ID;
-import static org.apache.kafka.common.protocol.CommonFields.MEMBER_ID;
-import static org.apache.kafka.common.protocol.CommonFields.THROTTLE_TIME_MS;
-import static org.apache.kafka.common.protocol.types.Type.BYTES;
-import static org.apache.kafka.common.protocol.types.Type.STRING;
-
 public class JoinGroupResponse extends AbstractResponse {
 
-    private static final String GROUP_PROTOCOL_KEY_NAME = "group_protocol";
-    private static final String LEADER_ID_KEY_NAME = "leader_id";
-    private static final String MEMBERS_KEY_NAME = "members";
-
-    private static final String MEMBER_METADATA_KEY_NAME = "member_metadata";
-
-    private static final Schema JOIN_GROUP_RESPONSE_MEMBER_V0 = new Schema(
-            MEMBER_ID,
-            new Field(MEMBER_METADATA_KEY_NAME, BYTES));
-
-    private static final Schema JOIN_GROUP_RESPONSE_V0 = new Schema(
-            ERROR_CODE,
-            GENERATION_ID,
-            new Field(GROUP_PROTOCOL_KEY_NAME, STRING, "The group protocol selected by the coordinator"),
-            new Field(LEADER_ID_KEY_NAME, STRING, "The leader of the group"),
-            MEMBER_ID,
-            new Field(MEMBERS_KEY_NAME, new ArrayOf(JOIN_GROUP_RESPONSE_MEMBER_V0)));
-
-    private static final Schema JOIN_GROUP_RESPONSE_V1 = JOIN_GROUP_RESPONSE_V0;
-
-    private static final Schema JOIN_GROUP_RESPONSE_V2 = new Schema(
-            THROTTLE_TIME_MS,
-            ERROR_CODE,
-            GENERATION_ID,
-            new Field(GROUP_PROTOCOL_KEY_NAME, STRING, "The group protocol selected by the coordinator"),
-            new Field(LEADER_ID_KEY_NAME, STRING, "The leader of the group"),
-            MEMBER_ID,
-            new Field(MEMBERS_KEY_NAME, new ArrayOf(JOIN_GROUP_RESPONSE_MEMBER_V0)));
-
-    /**
-     * The version number is bumped to indicate that on quota violation brokers send out responses before throttling.
-     */
-    private static final Schema JOIN_GROUP_RESPONSE_V3 = JOIN_GROUP_RESPONSE_V2;
-
-    /**
-     * The version number is bumped to indicate that client needs to issue a second join group request under first try
-     * with UNKNOWN_MEMBER_ID.
-     */
-    private static final Schema JOIN_GROUP_RESPONSE_V4 = JOIN_GROUP_RESPONSE_V3;
-
-    public static Schema[] schemaVersions() {
-        return new Schema[] {JOIN_GROUP_RESPONSE_V0, JOIN_GROUP_RESPONSE_V1, JOIN_GROUP_RESPONSE_V2,
-            JOIN_GROUP_RESPONSE_V3, JOIN_GROUP_RESPONSE_V4};
-    }
+    private final JoinGroupResponseData data;
 
     public static final String UNKNOWN_PROTOCOL = "";
     public static final int UNKNOWN_GENERATION_ID = -1;
     public static final String UNKNOWN_MEMBER_ID = "";
 
-    /**
-     * Possible error codes:
-     *
-     * COORDINATOR_LOAD_IN_PROGRESS (14)
-     * GROUP_COORDINATOR_NOT_AVAILABLE (15)
-     * NOT_COORDINATOR (16)
-     * INCONSISTENT_GROUP_PROTOCOL (23)
-     * UNKNOWN_MEMBER_ID (25)
-     * INVALID_SESSION_TIMEOUT (26)
-     * GROUP_AUTHORIZATION_FAILED (30)
-     * MEMBER_ID_REQUIRED (79)
-     */
-
-    private final int throttleTimeMs;
-    private final Errors error;
-    private final int generationId;
-    private final String groupProtocol;
-    private final String memberId;
-    private final String leaderId;
-    private final Map members;
-
-    public JoinGroupResponse(Errors error,
-                             int generationId,
-                             String groupProtocol,
-                             String memberId,
-                             String leaderId,
-                             Map groupMembers) {
-        this(DEFAULT_THROTTLE_TIME, error, generationId, groupProtocol, memberId, leaderId, groupMembers);
-    }
-
-    public JoinGroupResponse(int throttleTimeMs,
-            Errors error,
-            int generationId,
-            String groupProtocol,
-            String memberId,
-            String leaderId,
-            Map groupMembers) {
-        this.throttleTimeMs = throttleTimeMs;
-        this.error = error;
-        this.generationId = generationId;
-        this.groupProtocol = groupProtocol;
-        this.memberId = memberId;
-        this.leaderId = leaderId;
-        this.members = groupMembers;
+    public JoinGroupResponse(JoinGroupResponseData data) {
+        this.data = data;
     }
 
     public JoinGroupResponse(Struct struct) {
-        this.throttleTimeMs = struct.getOrElse(THROTTLE_TIME_MS, DEFAULT_THROTTLE_TIME);
-        members = new HashMap<>();
-
-        for (Object memberDataObj : struct.getArray(MEMBERS_KEY_NAME)) {
-            Struct memberData = (Struct) memberDataObj;
-            String memberId = memberData.get(MEMBER_ID);
-            ByteBuffer memberMetadata = memberData.getBytes(MEMBER_METADATA_KEY_NAME);
-            members.put(memberId, memberMetadata);
-        }
-        error = Errors.forCode(struct.get(ERROR_CODE));
-        generationId = struct.get(GENERATION_ID);
-        groupProtocol = struct.getString(GROUP_PROTOCOL_KEY_NAME);
-        memberId = struct.get(MEMBER_ID);
-        leaderId = struct.getString(LEADER_ID_KEY_NAME);
-    }
-
-    @Override
-    public int throttleTimeMs() {
-        return throttleTimeMs;
-    }
-
-    public Errors error() {
-        return error;
+        short latestVersion = (short) (JoinGroupResponseData.SCHEMAS.length - 1);
+        this.data = new JoinGroupResponseData(struct, latestVersion);
     }
 
-    @Override
-    public Map errorCounts() {
-        return errorCounts(error);
-    }
-
-    public int generationId() {
-        return generationId;
+    public JoinGroupResponse(Struct struct, short version) {
+        this.data = new JoinGroupResponseData(struct, version);
     }
 
-    public String groupProtocol() {
-        return groupProtocol;
+    public JoinGroupResponseData data() {
+        return data;
     }
 
-    public String memberId() {
-        return memberId;
+    public boolean isLeader() {
+        return data.memberId().equals(data.leader());
     }
 
-    public String leaderId() {
-        return leaderId;
+    @Override
+    public int throttleTimeMs() {
+        return data.throttleTimeMs();
     }
 
-    public boolean isLeader() {
-        return memberId.equals(leaderId);
+    public Errors error() {
+        return Errors.forCode(data.errorCode());
     }
 
-    public Map members() {
-        return members;
+    @Override
+    public Map errorCounts() {
+        return Collections.singletonMap(Errors.forCode(data.errorCode()), 1);
     }
 
-    public static JoinGroupResponse parse(ByteBuffer buffer, short version) {
-        return new JoinGroupResponse(ApiKeys.JOIN_GROUP.parseResponse(version, buffer));
+    public static JoinGroupResponse parse(ByteBuffer buffer, short versionId) {
+        return new JoinGroupResponse(ApiKeys.JOIN_GROUP.parseResponse(versionId, buffer), versionId);
     }
 
     @Override
     protected Struct toStruct(short version) {
-        Struct struct = new Struct(ApiKeys.JOIN_GROUP.responseSchema(version));
-        struct.setIfExists(THROTTLE_TIME_MS, throttleTimeMs);
-
-        struct.set(ERROR_CODE, error.code());
-        struct.set(GENERATION_ID, generationId);
-        struct.set(GROUP_PROTOCOL_KEY_NAME, groupProtocol);
-        struct.set(MEMBER_ID, memberId);
-        struct.set(LEADER_ID_KEY_NAME, leaderId);
-
-        List memberArray = new ArrayList<>();
-        for (Map.Entry entries : members.entrySet()) {
-            Struct memberData = struct.instance(MEMBERS_KEY_NAME);
-            memberData.set(MEMBER_ID, entries.getKey());
-            memberData.set(MEMBER_METADATA_KEY_NAME, entries.getValue());
-            memberArray.add(memberData);
-        }
-        struct.set(MEMBERS_KEY_NAME, memberArray.toArray());
-
-        return struct;
+        return data.toStruct(version);
     }
 
     @Override
     public String toString() {
-        return "JoinGroupResponse" +
-            "(throttleTimeMs=" + throttleTimeMs +
-            ", error=" + error +
-            ", generationId=" + generationId +
-            ", groupProtocol=" + groupProtocol +
-            ", memberId=" + memberId +
-            ", leaderId=" + leaderId +
-            ", members=" + ((members == null) ? "null" :
-                Utils.join(members.keySet(), ",")) + ")";
+        return data.toString();
     }
 
     @Override
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/LeaderAndIsrRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/LeaderAndIsrRequest.java
index 230c80734408a..54d0e9762f3e2 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/LeaderAndIsrRequest.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/LeaderAndIsrRequest.java
@@ -17,129 +17,76 @@
 package org.apache.kafka.common.requests;
 
 import org.apache.kafka.common.Node;
-import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.message.LeaderAndIsrRequestData;
+import org.apache.kafka.common.message.LeaderAndIsrRequestData.LeaderAndIsrLiveLeader;
+import org.apache.kafka.common.message.LeaderAndIsrRequestData.LeaderAndIsrTopicState;
+import org.apache.kafka.common.message.LeaderAndIsrRequestData.LeaderAndIsrPartitionState;
+import org.apache.kafka.common.message.LeaderAndIsrResponseData;
+import org.apache.kafka.common.message.LeaderAndIsrResponseData.LeaderAndIsrPartitionError;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.protocol.types.Field;
-import org.apache.kafka.common.protocol.types.Schema;
 import org.apache.kafka.common.protocol.types.Struct;
-import org.apache.kafka.common.utils.CollectionUtils;
+import org.apache.kafka.common.utils.FlattenedIterator;
 import org.apache.kafka.common.utils.Utils;
 
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
-import java.util.Set;
-
-import static org.apache.kafka.common.protocol.CommonFields.PARTITION_ID;
-import static org.apache.kafka.common.protocol.CommonFields.TOPIC_NAME;
-import static org.apache.kafka.common.protocol.types.Type.INT32;
+import java.util.stream.Collectors;
 
 public class LeaderAndIsrRequest extends AbstractControlRequest {
-    private static final Field.ComplexArray TOPIC_STATES = new Field.ComplexArray("topic_states", "Topic states");
-    private static final Field.ComplexArray PARTITION_STATES = new Field.ComplexArray("partition_states", "Partition states");
-    private static final Field.ComplexArray LIVE_LEADERS = new Field.ComplexArray("live_leaders", "Live leaders");
-
-    // PartitionState fields
-    private static final Field.Int32 LEADER = new Field.Int32("leader", "The broker id for the leader.");
-    private static final Field.Int32 LEADER_EPOCH = new Field.Int32("leader_epoch", "The leader epoch.");
-    private static final Field.Array ISR = new Field.Array("isr", INT32, "The in sync replica ids.");
-    private static final Field.Int32 ZK_VERSION = new Field.Int32("zk_version", "The ZK version.");
-    private static final Field.Array REPLICAS = new Field.Array("replicas", INT32, "The replica ids.");
-    private static final Field.Bool IS_NEW = new Field.Bool("is_new", "Whether the replica should have existed on the broker or not");
-
-    // live_leaders fields
-    private static final Field.Int32 END_POINT_ID = new Field.Int32("id", "The broker id");
-    private static final Field.Str HOST = new Field.Str("host", "The hostname of the broker.");
-    private static final Field.Int32  PORT = new Field.Int32("port", "The port on which the broker accepts requests.");
-
-    private static final Field PARTITION_STATES_V0  = PARTITION_STATES.withFields(
-            TOPIC_NAME,
-            PARTITION_ID,
-            CONTROLLER_EPOCH,
-            LEADER,
-            LEADER_EPOCH,
-            ISR,
-            ZK_VERSION,
-            REPLICAS);
-
-    // PARTITION_STATES_V1 added a per-partition is_new Field.
-    // This field specifies whether the replica should have existed on the broker or not.
-    private static final Field PARTITION_STATES_V1  = PARTITION_STATES.withFields(
-            TOPIC_NAME,
-            PARTITION_ID,
-            CONTROLLER_EPOCH,
-            LEADER,
-            LEADER_EPOCH,
-            ISR,
-            ZK_VERSION,
-            REPLICAS,
-            IS_NEW);
-
-    private static final Field PARTITION_STATES_V2  = PARTITION_STATES.withFields(
-            PARTITION_ID,
-            CONTROLLER_EPOCH,
-            LEADER,
-            LEADER_EPOCH,
-            ISR,
-            ZK_VERSION,
-            REPLICAS,
-            IS_NEW);
-
-    // TOPIC_STATES_V2 normalizes TOPIC_STATES_V1 to make it more memory efficient
-    private static final Field TOPIC_STATES_V2 = TOPIC_STATES.withFields(
-            TOPIC_NAME,
-            PARTITION_STATES_V2);
-
-    private static final Field LIVE_LEADERS_V0 = LIVE_LEADERS.withFields(
-            END_POINT_ID,
-            HOST,
-            PORT);
-
-    private static final Schema LEADER_AND_ISR_REQUEST_V0 = new Schema(
-            CONTROLLER_ID,
-            CONTROLLER_EPOCH,
-            PARTITION_STATES_V0,
-            LIVE_LEADERS_V0);
-
-    // LEADER_AND_ISR_REQUEST_V1 added a per-partition is_new Field. This field specifies whether the replica should
-    // have existed on the broker or not.
-    private static final Schema LEADER_AND_ISR_REQUEST_V1 = new Schema(
-            CONTROLLER_ID,
-            CONTROLLER_EPOCH,
-            PARTITION_STATES_V1,
-            LIVE_LEADERS_V0);
-
-    // LEADER_AND_ISR_REQUEST_V2 added a broker_epoch Field. This field specifies the generation of the broker across
-    // bounces. It also normalizes partitions under each topic.
-    private static final Schema LEADER_AND_ISR_REQUEST_V2 = new Schema(
-            CONTROLLER_ID,
-            CONTROLLER_EPOCH,
-            BROKER_EPOCH,
-            TOPIC_STATES_V2,
-            LIVE_LEADERS_V0);
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{LEADER_AND_ISR_REQUEST_V0, LEADER_AND_ISR_REQUEST_V1, LEADER_AND_ISR_REQUEST_V2};
-    }
 
     public static class Builder extends AbstractControlRequest.Builder {
-        private final Map partitionStates;
-        private final Set liveLeaders;
 
-        public Builder(short version, int controllerId, int controllerEpoch, long brokerEpoch,
-                       Map partitionStates, Set liveLeaders) {
-            super(ApiKeys.LEADER_AND_ISR, version, controllerId, controllerEpoch, brokerEpoch);
+        private final List partitionStates;
+        private final Collection liveLeaders;
+
+        public Builder(short version, int controllerId, int controllerEpoch, long brokerEpoch, long maxBrokerEpoch,
+                       List partitionStates, Collection liveLeaders) {
+            super(ApiKeys.LEADER_AND_ISR, version, controllerId, controllerEpoch, brokerEpoch, maxBrokerEpoch);
             this.partitionStates = partitionStates;
             this.liveLeaders = liveLeaders;
         }
 
         @Override
         public LeaderAndIsrRequest build(short version) {
-            return new LeaderAndIsrRequest(controllerId, controllerEpoch, brokerEpoch, partitionStates, liveLeaders, version);
+            List leaders = liveLeaders.stream().map(n -> new LeaderAndIsrLiveLeader()
+                .setBrokerId(n.id())
+                .setHostName(n.host())
+                .setPort(n.port())
+            ).collect(Collectors.toList());
+
+            LeaderAndIsrRequestData data = new LeaderAndIsrRequestData()
+                .setControllerId(controllerId)
+                .setControllerEpoch(controllerEpoch)
+                .setBrokerEpoch(brokerEpoch)
+                .setMaxBrokerEpoch(maxBrokerEpoch)
+                .setLiveLeaders(leaders);
+
+            if (version >= 2) {
+                Map topicStatesMap = groupByTopic(partitionStates);
+                data.setTopicStates(new ArrayList<>(topicStatesMap.values()));
+            } else {
+                data.setUngroupedPartitionStates(partitionStates);
+            }
+
+            return new LeaderAndIsrRequest(data, version);
+        }
+
+        private static Map groupByTopic(List partitionStates) {
+            Map topicStates = new HashMap<>();
+            // We don't null out the topic name in LeaderAndIsrRequestPartition since it's ignored by
+            // the generated code if version >= 2
+            for (LeaderAndIsrPartitionState partition : partitionStates) {
+                LeaderAndIsrTopicState topicState = topicStates.computeIfAbsent(partition.topicName(),
+                    t -> new LeaderAndIsrTopicState().setTopicName(partition.topicName()));
+                topicState.partitionStates().add(partition);
+            }
+            return topicStates;
         }
 
         @Override
@@ -149,209 +96,98 @@ public String toString() {
                 .append(", controllerId=").append(controllerId)
                 .append(", controllerEpoch=").append(controllerEpoch)
                 .append(", brokerEpoch=").append(brokerEpoch)
+                .append(", maxBrokerEpoch=").append(maxBrokerEpoch)
                 .append(", partitionStates=").append(partitionStates)
                 .append(", liveLeaders=(").append(Utils.join(liveLeaders, ", ")).append(")")
                 .append(")");
             return bld.toString();
+
         }
     }
 
-    private final Map partitionStates;
-    private final Set liveLeaders;
+    private final LeaderAndIsrRequestData data;
 
-    private LeaderAndIsrRequest(int controllerId, int controllerEpoch, long brokerEpoch, Map partitionStates,
-                                Set liveLeaders, short version) {
-        super(ApiKeys.LEADER_AND_ISR, version, controllerId, controllerEpoch, brokerEpoch);
-        this.partitionStates = partitionStates;
-        this.liveLeaders = liveLeaders;
+    LeaderAndIsrRequest(LeaderAndIsrRequestData data, short version) {
+        super(ApiKeys.LEADER_AND_ISR, version);
+        this.data = data;
+        // Do this from the constructor to make it thread-safe (even though it's only needed when some methods are called)
+        normalize();
     }
 
-    public LeaderAndIsrRequest(Struct struct, short version) {
-        super(ApiKeys.LEADER_AND_ISR, struct, version);
-
-        Map partitionStates = new HashMap<>();
-        if (struct.hasField(TOPIC_STATES)) {
-            for (Object topicStatesDataObj : struct.get(TOPIC_STATES)) {
-                Struct topicStatesData = (Struct) topicStatesDataObj;
-                String topic = topicStatesData.get(TOPIC_NAME);
-                for (Object partitionStateDataObj : topicStatesData.get(PARTITION_STATES)) {
-                    Struct partitionStateData = (Struct) partitionStateDataObj;
-                    int partition = partitionStateData.get(PARTITION_ID);
-                    PartitionState partitionState = new PartitionState(partitionStateData);
-                    partitionStates.put(new TopicPartition(topic, partition), partitionState);
+    private void normalize() {
+        if (version() >= 2) {
+            for (LeaderAndIsrTopicState topicState : data.topicStates()) {
+                for (LeaderAndIsrPartitionState partitionState : topicState.partitionStates()) {
+                    // Set the topic name so that we can always present the ungrouped view to callers
+                    partitionState.setTopicName(topicState.topicName());
                 }
             }
-        } else {
-            for (Object partitionStateDataObj : struct.get(PARTITION_STATES)) {
-                Struct partitionStateData = (Struct) partitionStateDataObj;
-                String topic = partitionStateData.get(TOPIC_NAME);
-                int partition = partitionStateData.get(PARTITION_ID);
-                PartitionState partitionState = new PartitionState(partitionStateData);
-                partitionStates.put(new TopicPartition(topic, partition), partitionState);
-            }
-        }
-
-        Set leaders = new HashSet<>();
-        for (Object leadersDataObj : struct.get(LIVE_LEADERS)) {
-            Struct leadersData = (Struct) leadersDataObj;
-            int id = leadersData.get(END_POINT_ID);
-            String host = leadersData.get(HOST);
-            int port = leadersData.get(PORT);
-            leaders.add(new Node(id, host, port));
         }
+    }
 
-        this.partitionStates = partitionStates;
-        this.liveLeaders = leaders;
+    public LeaderAndIsrRequest(Struct struct, short version) {
+        this(new LeaderAndIsrRequestData(struct, version), version);
     }
 
     @Override
     protected Struct toStruct() {
-        short version = version();
-        Struct struct = new Struct(ApiKeys.LEADER_AND_ISR.requestSchema(version));
-        struct.set(CONTROLLER_ID, controllerId);
-        struct.set(CONTROLLER_EPOCH, controllerEpoch);
-        struct.setIfExists(BROKER_EPOCH, brokerEpoch);
-
-        if (struct.hasField(TOPIC_STATES)) {
-            Map> topicStates = CollectionUtils.groupPartitionDataByTopic(partitionStates);
-            List topicStatesData = new ArrayList<>(topicStates.size());
-            for (Map.Entry> entry : topicStates.entrySet()) {
-                Struct topicStateData = struct.instance(TOPIC_STATES);
-                topicStateData.set(TOPIC_NAME, entry.getKey());
-                Map partitionMap = entry.getValue();
-                List partitionStatesData = new ArrayList<>(partitionMap.size());
-                for (Map.Entry partitionEntry : partitionMap.entrySet()) {
-                    Struct partitionStateData = topicStateData.instance(PARTITION_STATES);
-                    partitionStateData.set(PARTITION_ID, partitionEntry.getKey());
-                    partitionEntry.getValue().setStruct(partitionStateData);
-                    partitionStatesData.add(partitionStateData);
-                }
-                topicStateData.set(PARTITION_STATES, partitionStatesData.toArray());
-                topicStatesData.add(topicStateData);
-            }
-            struct.set(TOPIC_STATES, topicStatesData.toArray());
-        } else {
-            List partitionStatesData = new ArrayList<>(partitionStates.size());
-            for (Map.Entry entry : partitionStates.entrySet()) {
-                Struct partitionStateData = struct.instance(PARTITION_STATES);
-                TopicPartition topicPartition = entry.getKey();
-                partitionStateData.set(TOPIC_NAME, topicPartition.topic());
-                partitionStateData.set(PARTITION_ID, topicPartition.partition());
-                entry.getValue().setStruct(partitionStateData);
-                partitionStatesData.add(partitionStateData);
-            }
-            struct.set(PARTITION_STATES, partitionStatesData.toArray());
-        }
-
-        List leadersData = new ArrayList<>(liveLeaders.size());
-        for (Node leader : liveLeaders) {
-            Struct leaderData = struct.instance(LIVE_LEADERS);
-            leaderData.set(END_POINT_ID, leader.id());
-            leaderData.set(HOST, leader.host());
-            leaderData.set(PORT, leader.port());
-            leadersData.add(leaderData);
-        }
-        struct.set(LIVE_LEADERS, leadersData.toArray());
-        return struct;
+        return data.toStruct(version());
     }
 
     @Override
     public LeaderAndIsrResponse getErrorResponse(int throttleTimeMs, Throwable e) {
+        LeaderAndIsrResponseData responseData = new LeaderAndIsrResponseData();
         Errors error = Errors.forException(e);
-
-        Map responses = new HashMap<>(partitionStates.size());
-        for (TopicPartition partition : partitionStates.keySet()) {
-            responses.put(partition, error);
-        }
-
-        short versionId = version();
-        switch (versionId) {
-            case 0:
-            case 1:
-            case 2:
-                return new LeaderAndIsrResponse(error, responses);
-            default:
-                throw new IllegalArgumentException(String.format("Version %d is not valid. Valid versions for %s are 0 to %d",
-                        versionId, this.getClass().getSimpleName(), ApiKeys.LEADER_AND_ISR.latestVersion()));
+        responseData.setErrorCode(error.code());
+
+        List partitions = new ArrayList<>();
+        for (LeaderAndIsrPartitionState partition : partitionStates()) {
+            partitions.add(new LeaderAndIsrPartitionError()
+                .setTopicName(partition.topicName())
+                .setPartitionIndex(partition.partitionIndex())
+                .setErrorCode(error.code()));
         }
+        responseData.setPartitionErrors(partitions);
+        return new LeaderAndIsrResponse(responseData);
     }
 
+    @Override
     public int controllerId() {
-        return controllerId;
+        return data.controllerId();
     }
 
+    @Override
     public int controllerEpoch() {
-        return controllerEpoch;
+        return data.controllerEpoch();
     }
 
-    public Map partitionStates() {
-        return partitionStates;
+    @Override
+    public long brokerEpoch() {
+        return data.brokerEpoch();
     }
 
-    public Set liveLeaders() {
-        return liveLeaders;
+    @Override
+    public long maxBrokerEpoch() {
+        return data.maxBrokerEpoch();
     }
 
-    public static LeaderAndIsrRequest parse(ByteBuffer buffer, short version) {
-        return new LeaderAndIsrRequest(ApiKeys.LEADER_AND_ISR.parseRequest(version, buffer), version);
+    public Iterable partitionStates() {
+        if (version() >= 2)
+            return () -> new FlattenedIterator<>(data.topicStates().iterator(),
+                topicState -> topicState.partitionStates().iterator());
+        return data.ungroupedPartitionStates();
     }
 
-    public static final class PartitionState {
-        public final BasePartitionState basePartitionState;
-        public final boolean isNew;
-
-        public PartitionState(int controllerEpoch,
-                              int leader,
-                              int leaderEpoch,
-                              List isr,
-                              int zkVersion,
-                              List replicas,
-                              boolean isNew) {
-            this.basePartitionState = new BasePartitionState(controllerEpoch, leader, leaderEpoch, isr, zkVersion, replicas);
-            this.isNew = isNew;
-        }
-
-        private PartitionState(Struct struct) {
-            int controllerEpoch = struct.get(CONTROLLER_EPOCH);
-            int leader = struct.get(LEADER);
-            int leaderEpoch = struct.get(LEADER_EPOCH);
-
-            Object[] isrArray = struct.get(ISR);
-            List isr = new ArrayList<>(isrArray.length);
-            for (Object r : isrArray)
-                isr.add((Integer) r);
-
-            int zkVersion = struct.get(ZK_VERSION);
-
-            Object[] replicasArray = struct.get(REPLICAS);
-            List replicas = new ArrayList<>(replicasArray.length);
-            for (Object r : replicasArray)
-                replicas.add((Integer) r);
-
-            this.basePartitionState = new BasePartitionState(controllerEpoch, leader, leaderEpoch, isr, zkVersion, replicas);
-            this.isNew = struct.getOrElse(IS_NEW, false);
-        }
-
-        @Override
-        public String toString() {
-            return "PartitionState(controllerEpoch=" + basePartitionState.controllerEpoch +
-                ", leader=" + basePartitionState.leader +
-                ", leaderEpoch=" + basePartitionState.leaderEpoch +
-                ", isr=" + Utils.join(basePartitionState.isr, ",") +
-                ", zkVersion=" + basePartitionState.zkVersion +
-                ", replicas=" + Utils.join(basePartitionState.replicas, ",") +
-                ", isNew=" + isNew + ")";
-        }
+    public List liveLeaders() {
+        return Collections.unmodifiableList(data.liveLeaders());
+    }
 
-        private void setStruct(Struct struct) {
-            struct.set(CONTROLLER_EPOCH, basePartitionState.controllerEpoch);
-            struct.set(LEADER, basePartitionState.leader);
-            struct.set(LEADER_EPOCH, basePartitionState.leaderEpoch);
-            struct.set(ISR, basePartitionState.isr.toArray());
-            struct.set(ZK_VERSION, basePartitionState.zkVersion);
-            struct.set(REPLICAS, basePartitionState.replicas.toArray());
-            struct.setIfExists(IS_NEW, isNew);
-        }
+    // Visible for testing
+    LeaderAndIsrRequestData data() {
+        return data;
     }
 
+    public static LeaderAndIsrRequest parse(ByteBuffer buffer, short version) {
+        return new LeaderAndIsrRequest(ApiKeys.LEADER_AND_ISR.parseRequest(version, buffer), version);
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/LeaderAndIsrResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/LeaderAndIsrResponse.java
index 3ab9bf79de4a2..0329307e66693 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/LeaderAndIsrResponse.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/LeaderAndIsrResponse.java
@@ -16,43 +16,19 @@
  */
 package org.apache.kafka.common.requests;
 
-import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.message.LeaderAndIsrResponseData;
+import org.apache.kafka.common.message.LeaderAndIsrResponseData.LeaderAndIsrPartitionError;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.protocol.types.Field;
-import org.apache.kafka.common.protocol.types.Schema;
 import org.apache.kafka.common.protocol.types.Struct;
 
 import java.nio.ByteBuffer;
-import java.util.ArrayList;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-
-import static org.apache.kafka.common.protocol.CommonFields.ERROR_CODE;
-import static org.apache.kafka.common.protocol.CommonFields.PARTITION_ID;
-import static org.apache.kafka.common.protocol.CommonFields.TOPIC_NAME;
+import java.util.stream.Collectors;
 
 public class LeaderAndIsrResponse extends AbstractResponse {
-    private static final Field.ComplexArray PARTITIONS = new Field.ComplexArray("partitions", "Response for the requests partitions");
-
-    private static final Field PARTITIONS_V0 = PARTITIONS.withFields(
-            TOPIC_NAME,
-            PARTITION_ID,
-            ERROR_CODE);
-    private static final Schema LEADER_AND_ISR_RESPONSE_V0 = new Schema(
-            ERROR_CODE,
-            PARTITIONS_V0);
-
-    // LeaderAndIsrResponse V1 may receive KAFKA_STORAGE_ERROR in the response
-    private static final Schema LEADER_AND_ISR_RESPONSE_V1 = LEADER_AND_ISR_RESPONSE_V0;
-
-    private static final Schema LEADER_AND_ISR_RESPONSE_V2 = LEADER_AND_ISR_RESPONSE_V1;
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{LEADER_AND_ISR_RESPONSE_V0, LEADER_AND_ISR_RESPONSE_V1, LEADER_AND_ISR_RESPONSE_V2};
-    }
 
     /**
      * Possible error code:
@@ -60,74 +36,45 @@ public static Schema[] schemaVersions() {
      * STALE_CONTROLLER_EPOCH (11)
      * STALE_BROKER_EPOCH (77)
      */
-    private final Errors error;
+    private final LeaderAndIsrResponseData data;
 
-    private final Map responses;
-
-    public LeaderAndIsrResponse(Errors error, Map responses) {
-        this.responses = responses;
-        this.error = error;
+    public LeaderAndIsrResponse(LeaderAndIsrResponseData data) {
+        this.data = data;
     }
 
-    public LeaderAndIsrResponse(Struct struct) {
-        responses = new HashMap<>();
-        for (Object responseDataObj : struct.get(PARTITIONS)) {
-            Struct responseData = (Struct) responseDataObj;
-            String topic = responseData.get(TOPIC_NAME);
-            int partition = responseData.get(PARTITION_ID);
-            Errors error = Errors.forCode(responseData.get(ERROR_CODE));
-            responses.put(new TopicPartition(topic, partition), error);
-        }
-
-        error = Errors.forCode(struct.get(ERROR_CODE));
+    public LeaderAndIsrResponse(Struct struct, short version) {
+        this.data = new LeaderAndIsrResponseData(struct, version);
     }
 
-    public Map responses() {
-        return responses;
+    public List partitions() {
+        return data.partitionErrors();
     }
 
     public Errors error() {
-        return error;
+        return Errors.forCode(data.errorCode());
     }
 
     @Override
     public Map errorCounts() {
+        Errors error = error();
         if (error != Errors.NONE)
             // Minor optimization since the top-level error applies to all partitions
-            return Collections.singletonMap(error, responses.size());
-        return errorCounts(responses);
+            return Collections.singletonMap(error, data.partitionErrors().size());
+        return errorCounts(data.partitionErrors().stream().map(l -> Errors.forCode(l.errorCode())).collect(Collectors.toList()));
     }
 
     public static LeaderAndIsrResponse parse(ByteBuffer buffer, short version) {
-        return new LeaderAndIsrResponse(ApiKeys.LEADER_AND_ISR.parseResponse(version, buffer));
+        return new LeaderAndIsrResponse(ApiKeys.LEADER_AND_ISR.parseResponse(version, buffer), version);
     }
 
     @Override
     protected Struct toStruct(short version) {
-        Struct struct = new Struct(ApiKeys.LEADER_AND_ISR.responseSchema(version));
-
-        List responseDatas = new ArrayList<>(responses.size());
-        for (Map.Entry response : responses.entrySet()) {
-            Struct partitionData = struct.instance(PARTITIONS);
-            TopicPartition partition = response.getKey();
-            partitionData.set(TOPIC_NAME, partition.topic());
-            partitionData.set(PARTITION_ID, partition.partition());
-            partitionData.set(ERROR_CODE, response.getValue().code());
-            responseDatas.add(partitionData);
-        }
-
-        struct.set(PARTITIONS, responseDatas.toArray());
-        struct.set(ERROR_CODE, error.code());
-
-        return struct;
+        return data.toStruct(version);
     }
 
     @Override
     public String toString() {
-        return "LeaderAndIsrResponse(" +
-                "responses=" + responses +
-                ", error=" + error +
-                ")";
+        return data.toString();
     }
 
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/LeaveGroupRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/LeaveGroupRequest.java
index 20fdf4ffac78e..4cd72a6cf44ca 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/LeaveGroupRequest.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/LeaveGroupRequest.java
@@ -16,32 +16,68 @@
  */
 package org.apache.kafka.common.requests;
 
+import org.apache.kafka.common.errors.UnsupportedVersionException;
 import org.apache.kafka.common.message.LeaveGroupRequestData;
+import org.apache.kafka.common.message.LeaveGroupRequestData.MemberIdentity;
 import org.apache.kafka.common.message.LeaveGroupResponseData;
 import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.common.protocol.MessageUtil;
 import org.apache.kafka.common.protocol.types.Struct;
 
 import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.List;
 
 public class LeaveGroupRequest extends AbstractRequest {
 
     public static class Builder extends AbstractRequest.Builder {
+        private final String groupId;
+        private final List members;
 
-        private final LeaveGroupRequestData data;
+        public Builder(String groupId, List members) {
+            this(groupId, members, ApiKeys.LEAVE_GROUP.oldestVersion(), ApiKeys.LEAVE_GROUP.latestVersion());
+        }
 
-        public Builder(LeaveGroupRequestData data) {
-            super(ApiKeys.LEAVE_GROUP);
-            this.data = data;
+        Builder(String groupId, List members, short oldestVersion, short latestVersion) {
+            super(ApiKeys.LEAVE_GROUP, oldestVersion, latestVersion);
+            this.groupId = groupId;
+            this.members = members;
+            if (members.isEmpty()) {
+                throw new IllegalArgumentException("leaving members should not be empty");
+            }
         }
 
+        /**
+         * Based on the request version to choose fields.
+         */
         @Override
         public LeaveGroupRequest build(short version) {
+            final LeaveGroupRequestData data;
+            // Starting from version 3, all the leave group request will be in batch.
+            if (version >= 3) {
+                data = new LeaveGroupRequestData()
+                           .setGroupId(groupId)
+                           .setMembers(members);
+            } else {
+                if (members.size() != 1) {
+                    throw new UnsupportedVersionException("Version " + version + " leave group request only " +
+                                                              "supports single member instance than " + members.size() + " members");
+                }
+
+                data = new LeaveGroupRequestData()
+                           .setGroupId(groupId)
+                           .setMemberId(members.get(0).memberId());
+            }
             return new LeaveGroupRequest(data, version);
         }
 
         @Override
         public String toString() {
-            return data.toString();
+            return "(type=LeaveGroupRequest" +
+                       ", groupId=" + groupId +
+                       ", members=" + MessageUtil.deepToString(members.iterator()) +
+                       ")";
         }
     }
     private final LeaveGroupRequestData data;
@@ -63,13 +99,22 @@ public LeaveGroupRequestData data() {
         return data;
     }
 
+    public List members() {
+        // Before version 3, leave group request is still in single mode
+        return version <= 2 ? Collections.singletonList(
+            new MemberIdentity()
+                .setMemberId(data.memberId())) : data.members();
+    }
+
     @Override
     public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
-        LeaveGroupResponseData response = new LeaveGroupResponseData();
-        if (version() >= 2) {
-            response.setThrottleTimeMs(throttleTimeMs);
+        LeaveGroupResponseData responseData = new LeaveGroupResponseData()
+                                                  .setErrorCode(Errors.forException(e).code());
+
+        if (version() >= 1) {
+            responseData.setThrottleTimeMs(throttleTimeMs);
         }
-        return new LeaveGroupResponse(response);
+        return new LeaveGroupResponse(responseData);
     }
 
     public static LeaveGroupRequest parse(ByteBuffer buffer, short version) {
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/LeaveGroupResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/LeaveGroupResponse.java
index 3f67bb0b1ad30..e64fde168980b 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/LeaveGroupResponse.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/LeaveGroupResponse.java
@@ -17,46 +17,121 @@
 package org.apache.kafka.common.requests;
 
 import org.apache.kafka.common.message.LeaveGroupResponseData;
+import org.apache.kafka.common.message.LeaveGroupResponseData.MemberResponse;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
 import org.apache.kafka.common.protocol.types.Struct;
 
 import java.nio.ByteBuffer;
-import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 
+/**
+ * Possible error codes.
+ *
+ * Top level errors:
+ * - {@link Errors#COORDINATOR_LOAD_IN_PROGRESS}
+ * - {@link Errors#COORDINATOR_NOT_AVAILABLE}
+ * - {@link Errors#NOT_COORDINATOR}
+ * - {@link Errors#GROUP_AUTHORIZATION_FAILED}
+ *
+ * Member level errors:
+ * - {@link Errors#FENCED_INSTANCE_ID}
+ * - {@link Errors#UNKNOWN_MEMBER_ID}
+ *
+ * If the top level error code is set, normally this indicates that broker early stops the request
+ * handling due to some severe global error, so it is expected to see the member level errors to be empty.
+ * For older version response, we may populate member level error towards top level because older client
+ * couldn't parse member level.
+ */
 public class LeaveGroupResponse extends AbstractResponse {
 
-    private final LeaveGroupResponseData data;
+    public final LeaveGroupResponseData data;
 
     public LeaveGroupResponse(LeaveGroupResponseData data) {
         this.data = data;
     }
 
+    public LeaveGroupResponse(List memberResponses,
+                              Errors topLevelError,
+                              final int throttleTimeMs,
+                              final short version) {
+        if (version <= 2) {
+            // Populate member level error.
+            final short errorCode = getError(topLevelError, memberResponses).code();
+
+            this.data = new LeaveGroupResponseData()
+                            .setErrorCode(errorCode);
+        } else {
+            this.data = new LeaveGroupResponseData()
+                            .setErrorCode(topLevelError.code())
+                            .setMembers(memberResponses);
+        }
+
+        if (version >= 1) {
+            this.data.setThrottleTimeMs(throttleTimeMs);
+        }
+    }
+
     public LeaveGroupResponse(Struct struct) {
         short latestVersion = (short) (LeaveGroupResponseData.SCHEMAS.length - 1);
         this.data = new LeaveGroupResponseData(struct, latestVersion);
     }
+
     public LeaveGroupResponse(Struct struct, short version) {
         this.data = new LeaveGroupResponseData(struct, version);
     }
 
-    public LeaveGroupResponseData data() {
-        return data;
-    }
-
     @Override
     public int throttleTimeMs() {
         return data.throttleTimeMs();
     }
 
+    public List memberResponses() {
+        return data.members();
+    }
+
     public Errors error() {
+        return getError(Errors.forCode(data.errorCode()), data.members());
+    }
+
+    public Errors topLevelError() {
         return Errors.forCode(data.errorCode());
     }
 
+    private static Errors getError(Errors topLevelError, List memberResponses) {
+        if (topLevelError != Errors.NONE) {
+            return topLevelError;
+        } else {
+            for (MemberResponse memberResponse : memberResponses) {
+                Errors memberError = Errors.forCode(memberResponse.errorCode());
+                if (memberError != Errors.NONE) {
+                    return memberError;
+                }
+            }
+            return Errors.NONE;
+        }
+    }
+
     @Override
     public Map errorCounts() {
-        return Collections.singletonMap(Errors.forCode(data.errorCode()), 1);
+        Map combinedErrorCounts = new HashMap<>();
+        // Top level error.
+        Errors topLevelError = Errors.forCode(data.errorCode());
+        if (topLevelError != Errors.NONE) {
+            updateErrorCounts(combinedErrorCounts, topLevelError);
+        }
+
+        // Member level error.
+        for (MemberResponse memberResponse : data.members()) {
+            Errors memberError = Errors.forCode(memberResponse.errorCode());
+            if (memberError != Errors.NONE) {
+                updateErrorCounts(combinedErrorCounts, memberError);
+            }
+        }
+        return combinedErrorCounts;
     }
 
     @Override
@@ -72,4 +147,15 @@ public static LeaveGroupResponse parse(ByteBuffer buffer, short versionId) {
     public boolean shouldClientThrottle(short version) {
         return version >= 2;
     }
+
+    @Override
+    public boolean equals(Object other) {
+        return other instanceof LeaveGroupResponse &&
+                   ((LeaveGroupResponse) other).data.equals(this.data);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hashCode(data);
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/ListGroupsRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/ListGroupsRequest.java
index baab1e1c56ca4..caa6f47906bfd 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/ListGroupsRequest.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/ListGroupsRequest.java
@@ -16,69 +16,65 @@
  */
 package org.apache.kafka.common.requests;
 
+import org.apache.kafka.common.message.ListGroupsRequestData;
+import org.apache.kafka.common.message.ListGroupsResponseData;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.protocol.types.Schema;
 import org.apache.kafka.common.protocol.types.Struct;
 
 import java.nio.ByteBuffer;
 import java.util.Collections;
 
+/**
+ * Possible error codes:
+ *
+ * COORDINATOR_LOAD_IN_PROGRESS (14)
+ * COORDINATOR_NOT_AVAILABLE (15)
+ * AUTHORIZATION_FAILED (29)
+ */
 public class ListGroupsRequest extends AbstractRequest {
 
-    /* List groups api */
-    private static final Schema LIST_GROUPS_REQUEST_V0 = new Schema();
-
-    /* v1 request is the same as v0. Throttle time has been added to response */
-    private static final Schema LIST_GROUPS_REQUEST_V1 = LIST_GROUPS_REQUEST_V0;
+    public static class Builder extends AbstractRequest.Builder {
 
-    /**
-     * The version number is bumped to indicate that on quota violation brokers send out responses before throttling.
-     */
-    private static final Schema LIST_GROUPS_REQUEST_V2 = LIST_GROUPS_REQUEST_V1;
+        private final ListGroupsRequestData data;
 
-    public static Schema[] schemaVersions() {
-        return new Schema[] {LIST_GROUPS_REQUEST_V0, LIST_GROUPS_REQUEST_V1,
-            LIST_GROUPS_REQUEST_V2};
-    }
-
-    public static class Builder extends AbstractRequest.Builder {
-        public Builder() {
+        public Builder(ListGroupsRequestData data) {
             super(ApiKeys.LIST_GROUPS);
+            this.data = data;
         }
 
         @Override
         public ListGroupsRequest build(short version) {
-            return new ListGroupsRequest(version);
+            return new ListGroupsRequest(data, version);
         }
 
         @Override
         public String toString() {
-            return "(type=ListGroupsRequest)";
+            return data.toString();
         }
     }
 
-    public ListGroupsRequest(short version) {
+    private final ListGroupsRequestData data;
+
+    public ListGroupsRequest(ListGroupsRequestData data, short version) {
         super(ApiKeys.LIST_GROUPS, version);
+        this.data = data;
     }
 
-    public ListGroupsRequest(Struct struct, short versionId) {
-        super(ApiKeys.LIST_GROUPS, versionId);
+    public ListGroupsRequest(Struct struct, short version) {
+        super(ApiKeys.LIST_GROUPS, version);
+        this.data = new ListGroupsRequestData(struct, version);
     }
 
     @Override
     public ListGroupsResponse getErrorResponse(int throttleTimeMs, Throwable e) {
-        short versionId = version();
-        switch (versionId) {
-            case 0:
-                return new ListGroupsResponse(Errors.forException(e), Collections.emptyList());
-            case 1:
-            case 2:
-                return new ListGroupsResponse(throttleTimeMs, Errors.forException(e), Collections.emptyList());
-            default:
-                throw new IllegalArgumentException(String.format("Version %d is not valid. Valid versions for %s are 0 to %d",
-                        versionId, this.getClass().getSimpleName(), ApiKeys.LIST_GROUPS.latestVersion()));
+        ListGroupsResponseData listGroupsResponseData = new ListGroupsResponseData().
+            setGroups(Collections.emptyList()).
+            setErrorCode(Errors.forException(e).code());
+        if (version() >= 1) {
+            listGroupsResponseData.setThrottleTimeMs(throttleTimeMs);
         }
+        return new ListGroupsResponse(listGroupsResponseData);
     }
 
     public static ListGroupsRequest parse(ByteBuffer buffer, short version) {
@@ -87,6 +83,6 @@ public static ListGroupsRequest parse(ByteBuffer buffer, short version) {
 
     @Override
     protected Struct toStruct() {
-        return new Struct(ApiKeys.LIST_GROUPS.requestSchema(version()));
+        return data.toStruct(version());
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/ListGroupsResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/ListGroupsResponse.java
index af6f7212e8cd5..f3a67ced7e9ac 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/ListGroupsResponse.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/ListGroupsResponse.java
@@ -16,138 +16,48 @@
  */
 package org.apache.kafka.common.requests;
 
+import org.apache.kafka.common.message.ListGroupsResponseData;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.protocol.types.ArrayOf;
-import org.apache.kafka.common.protocol.types.Field;
-import org.apache.kafka.common.protocol.types.Schema;
 import org.apache.kafka.common.protocol.types.Struct;
 
 import java.nio.ByteBuffer;
-import java.util.ArrayList;
-import java.util.List;
+import java.util.Collections;
 import java.util.Map;
 
-import static org.apache.kafka.common.protocol.CommonFields.ERROR_CODE;
-import static org.apache.kafka.common.protocol.CommonFields.GROUP_ID;
-import static org.apache.kafka.common.protocol.CommonFields.THROTTLE_TIME_MS;
-import static org.apache.kafka.common.protocol.types.Type.STRING;
-
 public class ListGroupsResponse extends AbstractResponse {
 
-    private static final String GROUPS_KEY_NAME = "groups";
-    private static final String PROTOCOL_TYPE_KEY_NAME = "protocol_type";
-
-    private static final Schema LIST_GROUPS_RESPONSE_GROUP_V0 = new Schema(
-            GROUP_ID,
-            new Field(PROTOCOL_TYPE_KEY_NAME, STRING));
-    private static final Schema LIST_GROUPS_RESPONSE_V0 = new Schema(
-            ERROR_CODE,
-            new Field(GROUPS_KEY_NAME, new ArrayOf(LIST_GROUPS_RESPONSE_GROUP_V0)));
-    private static final Schema LIST_GROUPS_RESPONSE_V1 = new Schema(
-            THROTTLE_TIME_MS,
-            ERROR_CODE,
-            new Field(GROUPS_KEY_NAME, new ArrayOf(LIST_GROUPS_RESPONSE_GROUP_V0)));
-
-    /**
-     * The version number is bumped to indicate that on quota violation brokers send out responses before throttling.
-     */
-    private static final Schema LIST_GROUPS_RESPONSE_V2 = LIST_GROUPS_RESPONSE_V1;
+    private final ListGroupsResponseData data;
 
-    public static Schema[] schemaVersions() {
-        return new Schema[] {LIST_GROUPS_RESPONSE_V0, LIST_GROUPS_RESPONSE_V1,
-            LIST_GROUPS_RESPONSE_V2};
+    public ListGroupsResponse(ListGroupsResponseData data) {
+        this.data = data;
     }
 
-    /**
-     * Possible error codes:
-     *
-     * COORDINATOR_LOADING_IN_PROGRESS (14)
-     * COORDINATOR_NOT_AVAILABLE (15)
-     * AUTHORIZATION_FAILED (29)
-     */
-
-    private final Errors error;
-    private final int throttleTimeMs;
-    private final List groups;
-
-    public ListGroupsResponse(Errors error, List groups) {
-        this(DEFAULT_THROTTLE_TIME, error, groups);
+    public ListGroupsResponse(Struct struct, short version) {
+        this.data = new ListGroupsResponseData(struct, version);
     }
 
-    public ListGroupsResponse(int throttleTimeMs, Errors error, List groups) {
-        this.throttleTimeMs = throttleTimeMs;
-        this.error = error;
-        this.groups = groups;
-    }
-
-    public ListGroupsResponse(Struct struct) {
-        this.throttleTimeMs = struct.getOrElse(THROTTLE_TIME_MS, DEFAULT_THROTTLE_TIME);
-        this.error = Errors.forCode(struct.get(ERROR_CODE));
-        this.groups = new ArrayList<>();
-        for (Object groupObj : struct.getArray(GROUPS_KEY_NAME)) {
-            Struct groupStruct = (Struct) groupObj;
-            String groupId = groupStruct.get(GROUP_ID);
-            String protocolType = groupStruct.getString(PROTOCOL_TYPE_KEY_NAME);
-            this.groups.add(new Group(groupId, protocolType));
-        }
+    public ListGroupsResponseData data() {
+        return data;
     }
 
     @Override
     public int throttleTimeMs() {
-        return throttleTimeMs;
-    }
-
-    public List groups() {
-        return groups;
-    }
-
-    public Errors error() {
-        return error;
+        return data.throttleTimeMs();
     }
 
     @Override
     public Map errorCounts() {
-        return errorCounts(error);
-    }
-
-    public static class Group {
-        private final String groupId;
-        private final String protocolType;
-
-        public Group(String groupId, String protocolType) {
-            this.groupId = groupId;
-            this.protocolType = protocolType;
-        }
-
-        public String groupId() {
-            return groupId;
-        }
-
-        public String protocolType() {
-            return protocolType;
-        }
-
+        return Collections.singletonMap(Errors.forCode(data.errorCode()), 1);
     }
 
     @Override
     protected Struct toStruct(short version) {
-        Struct struct = new Struct(ApiKeys.LIST_GROUPS.responseSchema(version));
-        struct.setIfExists(THROTTLE_TIME_MS, throttleTimeMs);
-        struct.set(ERROR_CODE, error.code());
-        List groupList = new ArrayList<>();
-        for (Group group : groups) {
-            Struct groupStruct = struct.instance(GROUPS_KEY_NAME);
-            groupStruct.set(GROUP_ID, group.groupId);
-            groupStruct.set(PROTOCOL_TYPE_KEY_NAME, group.protocolType);
-            groupList.add(groupStruct);
-        }
-        struct.set(GROUPS_KEY_NAME, groupList.toArray());
-        return struct;
+        return data.toStruct(version);
     }
 
     public static ListGroupsResponse parse(ByteBuffer buffer, short version) {
-        return new ListGroupsResponse(ApiKeys.LIST_GROUPS.parseResponse(version, buffer));
+        return new ListGroupsResponse(ApiKeys.LIST_GROUPS.responseSchema(version).read(buffer), version);
     }
 
     @Override
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/ListOffsetRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/ListOffsetRequest.java
index e9fe942e8581a..01418c86e5832 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/ListOffsetRequest.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/ListOffsetRequest.java
@@ -31,6 +31,7 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
 
@@ -203,6 +204,19 @@ public PartitionData(long timestamp, Optional currentLeaderEpoch) {
             this(timestamp, 1, currentLeaderEpoch);
         }
 
+        @Override
+        public boolean equals(Object obj) {
+            if (!(obj instanceof PartitionData)) return false;
+            PartitionData other = (PartitionData) obj;
+            return this.timestamp == other.timestamp &&
+                this.currentLeaderEpoch.equals(other.currentLeaderEpoch);
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(timestamp, currentLeaderEpoch);
+        }
+
         @Override
         public String toString() {
             StringBuilder bld = new StringBuilder();
@@ -269,12 +283,13 @@ public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
             responseData.put(partition, partitionError);
         }
 
-        switch (version()) {
+        switch (versionId) {
             case 0:
             case 1:
             case 2:
             case 3:
             case 4:
+            case 5:
                 return new ListOffsetResponse(throttleTimeMs, responseData);
             default:
                 throw new IllegalArgumentException(String.format("Version %d is not valid. Valid versions for %s are 0 to %d",
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/ListOffsetResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/ListOffsetResponse.java
index 769c850e22c2f..5cf35e1d7900e 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/ListOffsetResponse.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/ListOffsetResponse.java
@@ -70,7 +70,7 @@ public class ListOffsetResponse extends AbstractResponse {
     // partition level fields
     // This key is only used by ListOffsetResponse v0
     @Deprecated
-    private static final Field.Array OFFSETS = new Field.Array("offsets'", INT64, "A list of offsets.");
+    private static final Field.Array OFFSETS = new Field.Array("offsets", INT64, "A list of offsets.");
     private static final Field.Int64 TIMESTAMP = new Field.Int64("timestamp",
             "The timestamp associated with the returned offset");
     private static final Field.Int64 OFFSET = new Field.Int64("offset",
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/ListPartitionReassignmentsRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/ListPartitionReassignmentsRequest.java
new file mode 100644
index 0000000000000..0aa5e557556ea
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/ListPartitionReassignmentsRequest.java
@@ -0,0 +1,109 @@
+/*
+ * 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.requests;
+
+import org.apache.kafka.common.message.ListPartitionReassignmentsRequestData;
+import org.apache.kafka.common.message.ListPartitionReassignmentsResponseData;
+import org.apache.kafka.common.message.ListPartitionReassignmentsResponseData.OngoingPartitionReassignment;
+import org.apache.kafka.common.message.ListPartitionReassignmentsResponseData.OngoingTopicReassignment;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.types.Struct;
+
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.stream.Collectors;
+
+import static org.apache.kafka.common.message.ListPartitionReassignmentsRequestData.ListPartitionReassignmentsTopics;
+
+public class ListPartitionReassignmentsRequest extends AbstractRequest {
+
+    public static class Builder extends AbstractRequest.Builder {
+        private final ListPartitionReassignmentsRequestData data;
+
+        public Builder(ListPartitionReassignmentsRequestData data) {
+            super(ApiKeys.LIST_PARTITION_REASSIGNMENTS);
+            this.data = data;
+        }
+
+        @Override
+        public ListPartitionReassignmentsRequest build(short version) {
+            return new ListPartitionReassignmentsRequest(data, version);
+        }
+
+        @Override
+        public String toString() {
+            return data.toString();
+        }
+    }
+
+    private ListPartitionReassignmentsRequestData data;
+    private final short version;
+
+    private ListPartitionReassignmentsRequest(ListPartitionReassignmentsRequestData data, short version) {
+        super(ApiKeys.LIST_PARTITION_REASSIGNMENTS, version);
+        this.data = data;
+        this.version = version;
+    }
+
+    ListPartitionReassignmentsRequest(Struct struct, short version) {
+        super(ApiKeys.LIST_PARTITION_REASSIGNMENTS, version);
+        this.data = new ListPartitionReassignmentsRequestData(struct, version);
+        this.version = version;
+    }
+
+    public static ListPartitionReassignmentsRequest parse(ByteBuffer buffer, short version) {
+        return new ListPartitionReassignmentsRequest(
+                ApiKeys.LIST_PARTITION_REASSIGNMENTS.parseRequest(version, buffer), version
+        );
+    }
+
+    public ListPartitionReassignmentsRequestData data() {
+        return data;
+    }
+
+    /**
+     * Visible for testing.
+     */
+    @Override
+    public Struct toStruct() {
+        return data.toStruct(version);
+    }
+
+    @Override
+    public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
+        ApiError apiError = ApiError.fromThrowable(e);
+
+        List ongoingTopicReassignments = new ArrayList<>();
+        if (data.topics() != null) {
+            for (ListPartitionReassignmentsTopics topic : data.topics()) {
+                ongoingTopicReassignments.add(
+                        new OngoingTopicReassignment()
+                                .setName(topic.name())
+                                .setPartitions(topic.partitionIndexes().stream().map(partitionIndex ->
+                                        new OngoingPartitionReassignment().setPartitionIndex(partitionIndex)).collect(Collectors.toList()))
+                );
+            }
+        }
+        ListPartitionReassignmentsResponseData responseData = new ListPartitionReassignmentsResponseData()
+                .setTopics(ongoingTopicReassignments)
+                .setErrorCode(apiError.error().code())
+                .setErrorMessage(apiError.message())
+                .setThrottleTimeMs(throttleTimeMs);
+        return new ListPartitionReassignmentsResponse(responseData);
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/ListPartitionReassignmentsResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/ListPartitionReassignmentsResponse.java
new file mode 100644
index 0000000000000..0d53e4dd7caee
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/ListPartitionReassignmentsResponse.java
@@ -0,0 +1,75 @@
+/*
+ * 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.requests;
+
+import org.apache.kafka.common.message.ListPartitionReassignmentsResponseData;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.common.protocol.types.Struct;
+
+import java.nio.ByteBuffer;
+import java.util.HashMap;
+import java.util.Map;
+
+public class ListPartitionReassignmentsResponse extends AbstractResponse {
+
+    private final ListPartitionReassignmentsResponseData data;
+
+    public ListPartitionReassignmentsResponse(Struct struct) {
+        this(struct, ApiKeys.LIST_PARTITION_REASSIGNMENTS.latestVersion());
+    }
+
+    public ListPartitionReassignmentsResponse(ListPartitionReassignmentsResponseData responseData) {
+        this.data = responseData;
+    }
+
+    ListPartitionReassignmentsResponse(Struct struct, short version) {
+        this.data = new ListPartitionReassignmentsResponseData(struct, version);
+    }
+
+    public static ListPartitionReassignmentsResponse parse(ByteBuffer buffer, short version) {
+        return new ListPartitionReassignmentsResponse(ApiKeys.LIST_PARTITION_REASSIGNMENTS.responseSchema(version).read(buffer), version);
+    }
+
+    public ListPartitionReassignmentsResponseData data() {
+        return data;
+    }
+
+    @Override
+    public boolean shouldClientThrottle(short version) {
+        return true;
+    }
+
+    @Override
+    public int throttleTimeMs() {
+        return data.throttleTimeMs();
+    }
+
+    @Override
+    public Map errorCounts() {
+        Map counts = new HashMap<>();
+        Errors topLevelErr = Errors.forCode(data.errorCode());
+        counts.put(topLevelErr, 1);
+
+        return counts;
+    }
+
+    @Override
+    protected Struct toStruct(short version) {
+        return data.toStruct(version);
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/MetadataRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/MetadataRequest.java
index 3f12f1de9a87a..cbfca4bbdf6f2 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/MetadataRequest.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/MetadataRequest.java
@@ -17,159 +17,120 @@
 package org.apache.kafka.common.requests;
 
 import org.apache.kafka.common.errors.UnsupportedVersionException;
+import org.apache.kafka.common.message.MetadataRequestData;
+import org.apache.kafka.common.message.MetadataRequestData.MetadataRequestTopic;
+import org.apache.kafka.common.message.MetadataResponseData;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.protocol.types.ArrayOf;
-import org.apache.kafka.common.protocol.types.Field;
-import org.apache.kafka.common.protocol.types.Schema;
 import org.apache.kafka.common.protocol.types.Struct;
-import org.apache.kafka.common.utils.Utils;
 
 import java.nio.ByteBuffer;
-import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
-
-import static org.apache.kafka.common.protocol.types.Type.STRING;
+import java.util.stream.Collectors;
 
 public class MetadataRequest extends AbstractRequest {
 
-    private static final String TOPICS_KEY_NAME = "topics";
-
-    private static final Schema METADATA_REQUEST_V0 = new Schema(
-            new Field(TOPICS_KEY_NAME, new ArrayOf(STRING), "An array of topics to fetch metadata for. If no topics are specified fetch metadata for all topics."));
-
-    private static final Schema METADATA_REQUEST_V1 = new Schema(
-            new Field(TOPICS_KEY_NAME, ArrayOf.nullable(STRING), "An array of topics to fetch metadata for. If the topics array is null fetch metadata for all topics."));
-
-    /* The v2 metadata request is the same as v1. An additional field for cluster id has been added to the v2 metadata response */
-    private static final Schema METADATA_REQUEST_V2 = METADATA_REQUEST_V1;
-
-    /* The v3 metadata request is the same as v1 and v2. An additional field for throttle time has been added to the v3 metadata response */
-    private static final Schema METADATA_REQUEST_V3 = METADATA_REQUEST_V2;
-
-    /* The v4 metadata request has an additional field for allowing auto topic creation. The response is the same as v3. */
-    private static final Field.Bool ALLOW_AUTO_TOPIC_CREATION = new Field.Bool("allow_auto_topic_creation",
-            "If this and the broker config auto.create.topics.enable are true, topics that " +
-                    "don't exist will be created by the broker. Otherwise, no topics will be created by the broker.");
-
-    private static final Schema METADATA_REQUEST_V4 = new Schema(
-            new Field(TOPICS_KEY_NAME, ArrayOf.nullable(STRING), "An array of topics to fetch metadata for. " +
-                    "If the topics array is null fetch metadata for all topics."),
-            ALLOW_AUTO_TOPIC_CREATION);
+    public static class Builder extends AbstractRequest.Builder {
+        private static final MetadataRequestData ALL_TOPICS_REQUEST_DATA = new MetadataRequestData().
+            setTopics(null).setAllowAutoTopicCreation(true);
 
-    /* The v5 metadata request is the same as v4. An additional field for offline_replicas has been added to the v5 metadata response */
-    private static final Schema METADATA_REQUEST_V5 = METADATA_REQUEST_V4;
+        private final MetadataRequestData data;
 
-    /**
-     * The version number is bumped to indicate that on quota violation brokers send out responses before throttling.
-     */
-    private static final Schema METADATA_REQUEST_V6 = METADATA_REQUEST_V5;
+        public Builder(MetadataRequestData data) {
+            super(ApiKeys.METADATA);
+            this.data = data;
+        }
 
-    /**
-     * Bumped for the addition of the current leader epoch in the metadata response.
-     */
-    private static final Schema METADATA_REQUEST_V7 = METADATA_REQUEST_V6;
+        public Builder(List topics, boolean allowAutoTopicCreation, short allowedVersion) {
+            this(topics, allowAutoTopicCreation, allowedVersion, allowedVersion);
+        }
 
-    public static Schema[] schemaVersions() {
-        return new Schema[] {METADATA_REQUEST_V0, METADATA_REQUEST_V1, METADATA_REQUEST_V2, METADATA_REQUEST_V3,
-            METADATA_REQUEST_V4, METADATA_REQUEST_V5, METADATA_REQUEST_V6, METADATA_REQUEST_V7};
-    }
+        public Builder(List topics, boolean allowAutoTopicCreation, short minVersion, short maxVersion) {
+            super(ApiKeys.METADATA, minVersion, maxVersion);
+            MetadataRequestData data = new MetadataRequestData();
+            if (topics == null)
+                data.setTopics(null);
+            else {
+                topics.forEach(topic -> data.topics().add(new MetadataRequestTopic().setName(topic)));
+            }
 
-    public static class Builder extends AbstractRequest.Builder {
-        private static final List ALL_TOPICS = null;
+            data.setAllowAutoTopicCreation(allowAutoTopicCreation);
+            this.data = data;
+        }
 
-        // The list of topics, or null if we want to request metadata about all topics.
-        private final List topics;
-        private final boolean allowAutoTopicCreation;
+        public Builder(List topics, boolean allowAutoTopicCreation) {
+            this(topics, allowAutoTopicCreation, ApiKeys.METADATA.oldestVersion(),  ApiKeys.METADATA.latestVersion());
+        }
 
         public static Builder allTopics() {
             // This never causes auto-creation, but we set the boolean to true because that is the default value when
             // deserializing V2 and older. This way, the value is consistent after serialization and deserialization.
-            return new Builder(ALL_TOPICS, true);
+            return new Builder(ALL_TOPICS_REQUEST_DATA);
         }
 
-        public Builder(List topics, boolean allowAutoTopicCreation) {
-            super(ApiKeys.METADATA);
-            this.topics = topics;
-            this.allowAutoTopicCreation = allowAutoTopicCreation;
+        public boolean emptyTopicList() {
+            return data.topics().isEmpty();
         }
 
-        public List topics() {
-            return this.topics;
+        public boolean isAllTopics() {
+            return data.topics() == null;
         }
 
-        public boolean isAllTopics() {
-            return this.topics == ALL_TOPICS;
+        public List topics() {
+            return data.topics()
+                .stream()
+                .map(MetadataRequestTopic::name)
+                .collect(Collectors.toList());
         }
 
         @Override
         public MetadataRequest build(short version) {
             if (version < 1)
                 throw new UnsupportedVersionException("MetadataRequest versions older than 1 are not supported.");
-            if (!allowAutoTopicCreation && version < 4)
+            if (!data.allowAutoTopicCreation() && version < 4)
                 throw new UnsupportedVersionException("MetadataRequest versions older than 4 don't support the " +
                         "allowAutoTopicCreation field");
-            return new MetadataRequest(this.topics, allowAutoTopicCreation, version);
+            return new MetadataRequest(data, version);
         }
 
         @Override
         public String toString() {
-            StringBuilder bld = new StringBuilder();
-            bld.append("(type=MetadataRequest").
-                append(", topics=");
-            if (topics == null) {
-                bld.append("");
-            } else {
-                bld.append(Utils.join(topics, ","));
-            }
-            bld.append(")");
-            return bld.toString();
+            return data.toString();
         }
     }
 
-    private final List topics;
-    private final boolean allowAutoTopicCreation;
+    private final MetadataRequestData data;
+    private final short version;
 
-    /**
-     * In v0 null is not allowed and an empty list indicates requesting all topics.
-     * Note: modern clients do not support sending v0 requests.
-     * In v1 null indicates requesting all topics, and an empty list indicates requesting no topics.
-     */
-    public MetadataRequest(List topics, boolean allowAutoTopicCreation, short version) {
+    public MetadataRequest(MetadataRequestData data, short version) {
         super(ApiKeys.METADATA, version);
-        this.topics = topics;
-        this.allowAutoTopicCreation = allowAutoTopicCreation;
+        this.data = data;
+        this.version = version;
     }
 
     public MetadataRequest(Struct struct, short version) {
         super(ApiKeys.METADATA, version);
-        Object[] topicArray = struct.getArray(TOPICS_KEY_NAME);
-        if (topicArray != null) {
-            if (topicArray.length == 0 && version == 0) {
-                topics = null;
-            } else {
-                topics = new ArrayList<>();
-                for (Object topicObj: topicArray) {
-                    topics.add((String) topicObj);
-                }
-            }
-        } else {
-            topics = null;
-        }
+        this.data = new MetadataRequestData(struct, version);
+        this.version = version;
+    }
 
-        allowAutoTopicCreation = struct.getOrElse(ALLOW_AUTO_TOPIC_CREATION, true);
+    public MetadataRequestData data() {
+        return data;
     }
 
     @Override
     public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
-        List topicMetadatas = new ArrayList<>();
         Errors error = Errors.forException(e);
-        List partitions = Collections.emptyList();
-
-        if (topics != null) {
-            for (String topic : topics)
-                topicMetadatas.add(new MetadataResponse.TopicMetadata(error, topic, false, partitions));
+        MetadataResponseData responseData = new MetadataResponseData();
+        if (topics() != null) {
+            for (String topic :topics())
+                responseData.topics().add(new MetadataResponseData.MetadataResponseTopic()
+                    .setName(topic)
+                    .setErrorCode(error.code())
+                    .setIsInternal(false)
+                    .setPartitions(Collections.emptyList()));
         }
 
         short versionId = version();
@@ -177,13 +138,15 @@ public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
             case 0:
             case 1:
             case 2:
-                return new MetadataResponse(Collections.emptyList(), null, MetadataResponse.NO_CONTROLLER_ID, topicMetadatas);
+                return new MetadataResponse(responseData);
             case 3:
             case 4:
             case 5:
             case 6:
             case 7:
-                return new MetadataResponse(throttleTimeMs, Collections.emptyList(), null, MetadataResponse.NO_CONTROLLER_ID, topicMetadatas);
+            case 8:
+                responseData.setThrottleTimeMs(throttleTimeMs);
+                return new MetadataResponse(responseData);
             default:
                 throw new IllegalArgumentException(String.format("Version %d is not valid. Valid versions for %s are 0 to %d",
                         versionId, this.getClass().getSimpleName(), ApiKeys.METADATA.latestVersion()));
@@ -191,29 +154,37 @@ public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
     }
 
     public boolean isAllTopics() {
-        return topics == null;
+        return (data.topics() == null) ||
+            (data.topics().isEmpty() && version == 0); //In version 0, an empty topic list indicates
+                                                         // "request metadata for all topics."
     }
 
     public List topics() {
-        return topics;
+        if (isAllTopics()) //In version 0, we return null for empty topic list
+            return null;
+        else
+            return data.topics()
+                .stream()
+                .map(MetadataRequestTopic::name)
+                .collect(Collectors.toList());
     }
 
     public boolean allowAutoTopicCreation() {
-        return allowAutoTopicCreation;
+        return data.allowAutoTopicCreation();
     }
 
     public static MetadataRequest parse(ByteBuffer buffer, short version) {
         return new MetadataRequest(ApiKeys.METADATA.parseRequest(version, buffer), version);
     }
 
+    public static List convertToMetadataRequestTopic(final Collection topics) {
+        return topics.stream().map(topic -> new MetadataRequestTopic()
+            .setName(topic))
+            .collect(Collectors.toList());
+    }
+
     @Override
     protected Struct toStruct() {
-        Struct struct = new Struct(ApiKeys.METADATA.requestSchema(version()));
-        if (topics == null)
-            struct.set(TOPICS_KEY_NAME, null);
-        else
-            struct.set(TOPICS_KEY_NAME, topics.toArray());
-        struct.setIfExists(ALLOW_AUTO_TOPIC_CREATION, allowAutoTopicCreation);
-        return struct;
+        return data.toStruct(version);
     }
 }
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 f90876fd261bb..fb8ae8f46223f 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
@@ -19,11 +19,14 @@
 import org.apache.kafka.common.Cluster;
 import org.apache.kafka.common.Node;
 import org.apache.kafka.common.PartitionInfo;
+import org.apache.kafka.common.message.MetadataResponseData;
+import org.apache.kafka.common.message.MetadataResponseData.MetadataResponseBroker;
+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.protocol.types.Field;
-import org.apache.kafka.common.protocol.types.Schema;
 import org.apache.kafka.common.protocol.types.Struct;
+import org.apache.kafka.common.record.RecordBatch;
 import org.apache.kafka.common.utils.Utils;
 
 import java.nio.ByteBuffer;
@@ -33,15 +36,10 @@
 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 static org.apache.kafka.common.protocol.CommonFields.ERROR_CODE;
-import static org.apache.kafka.common.protocol.CommonFields.LEADER_EPOCH;
-import static org.apache.kafka.common.protocol.CommonFields.PARTITION_ID;
-import static org.apache.kafka.common.protocol.CommonFields.THROTTLE_TIME_MS;
-import static org.apache.kafka.common.protocol.CommonFields.TOPIC_NAME;
-import static org.apache.kafka.common.protocol.types.Type.INT32;
+import java.util.stream.Collectors;
 
 /**
  * Possible topic-level error codes:
@@ -57,253 +55,39 @@
 public class MetadataResponse extends AbstractResponse {
     public static final int NO_CONTROLLER_ID = -1;
 
-    private static final Field.ComplexArray BROKERS = new Field.ComplexArray("brokers",
-            "Host and port information for all brokers.");
-    private static final Field.ComplexArray TOPIC_METADATA = new Field.ComplexArray("topic_metadata",
-            "Metadata for requested topics");
-
-    // cluster level fields
-    private static final Field.NullableStr CLUSTER_ID = new Field.NullableStr("cluster_id",
-            "The cluster id that this broker belongs to.");
-    private static final Field.Int32 CONTROLLER_ID = new Field.Int32("controller_id",
-            "The broker id of the controller broker.");
-
-    // broker level fields
-    private static final Field.Int32 NODE_ID = new Field.Int32("node_id", "The broker id.");
-    private static final Field.Str HOST = new Field.Str("host", "The hostname of the broker.");
-    private static final Field.Int32 PORT = new Field.Int32("port", "The port on which the broker accepts requests.");
-    private static final Field.NullableStr RACK = new Field.NullableStr("rack", "The rack of the broker.");
-
-    // topic level fields
-    private static final Field.ComplexArray PARTITION_METADATA = new Field.ComplexArray("partition_metadata",
-            "Metadata for each partition of the topic.");
-    private static final Field.Bool IS_INTERNAL = new Field.Bool("is_internal",
-            "Indicates if the topic is considered a Kafka internal topic");
-
-    // partition level fields
-    private static final Field.Int32 LEADER = new Field.Int32("leader",
-            "The id of the broker acting as leader for this partition.");
-    private static final Field.Array REPLICAS = new Field.Array("replicas", INT32,
-            "The set of all nodes that host this partition.");
-    private static final Field.Array ISR = new Field.Array("isr", INT32,
-            "The set of nodes that are in sync with the leader for this partition.");
-    private static final Field.Array OFFLINE_REPLICAS = new Field.Array("offline_replicas", INT32,
-            "The set of offline replicas of this partition.");
-
-    private static final Field METADATA_BROKER_V0 = BROKERS.withFields(
-            NODE_ID,
-            HOST,
-            PORT);
-
-    private static final Field PARTITION_METADATA_V0 = PARTITION_METADATA.withFields(
-            ERROR_CODE,
-            PARTITION_ID,
-            LEADER,
-            REPLICAS,
-            ISR);
-
-    private static final Field TOPIC_METADATA_V0 = TOPIC_METADATA.withFields(
-            ERROR_CODE,
-            TOPIC_NAME,
-            PARTITION_METADATA_V0);
-
-    private static final Schema METADATA_RESPONSE_V0 = new Schema(
-            METADATA_BROKER_V0,
-            TOPIC_METADATA_V0);
-
-    // V1 adds fields for the rack of each broker, the controller id, and whether or not the topic is internal
-    private static final Field METADATA_BROKER_V1 = BROKERS.withFields(
-            NODE_ID,
-            HOST,
-            PORT,
-            RACK);
-
-    private static final Field TOPIC_METADATA_V1 = TOPIC_METADATA.withFields(
-            ERROR_CODE,
-            TOPIC_NAME,
-            IS_INTERNAL,
-            PARTITION_METADATA_V0);
-
-    private static final Schema METADATA_RESPONSE_V1 = new Schema(
-            METADATA_BROKER_V1,
-            CONTROLLER_ID,
-            TOPIC_METADATA_V1);
-
-    // V2 added a field for the cluster id
-    private static final Schema METADATA_RESPONSE_V2 = new Schema(
-            METADATA_BROKER_V1,
-            CLUSTER_ID,
-            CONTROLLER_ID,
-            TOPIC_METADATA_V1);
-
-    // V3 adds the throttle time to the response
-    private static final Schema METADATA_RESPONSE_V3 = new Schema(
-            THROTTLE_TIME_MS,
-            METADATA_BROKER_V1,
-            CLUSTER_ID,
-            CONTROLLER_ID,
-            TOPIC_METADATA_V1);
-
-    private static final Schema METADATA_RESPONSE_V4 = METADATA_RESPONSE_V3;
-
-    // V5 added a per-partition offline_replicas field. This field specifies the list of replicas that are offline.
-    private static final Field PARTITION_METADATA_V5 = PARTITION_METADATA.withFields(
-            ERROR_CODE,
-            PARTITION_ID,
-            LEADER,
-            REPLICAS,
-            ISR,
-            OFFLINE_REPLICAS);
-
-    private static final Field TOPIC_METADATA_V5 = TOPIC_METADATA.withFields(
-            ERROR_CODE,
-            TOPIC_NAME,
-            IS_INTERNAL,
-            PARTITION_METADATA_V5);
-
-    private static final Schema METADATA_RESPONSE_V5 = new Schema(
-            THROTTLE_TIME_MS,
-            METADATA_BROKER_V1,
-            CLUSTER_ID,
-            CONTROLLER_ID,
-            TOPIC_METADATA_V5);
-
-    // V6 bump used to indicate that on quota violation brokers send out responses before throttling.
-    private static final Schema METADATA_RESPONSE_V6 = METADATA_RESPONSE_V5;
-
-    // V7 adds the leader epoch to the partition metadata
-    private static final Field PARTITION_METADATA_V7 = PARTITION_METADATA.withFields(
-            ERROR_CODE,
-            PARTITION_ID,
-            LEADER,
-            LEADER_EPOCH,
-            REPLICAS,
-            ISR,
-            OFFLINE_REPLICAS);
-
-    private static final Field TOPIC_METADATA_V7 = TOPIC_METADATA.withFields(
-            ERROR_CODE,
-            TOPIC_NAME,
-            IS_INTERNAL,
-            PARTITION_METADATA_V7);
-
-    private static final Schema METADATA_RESPONSE_V7 = new Schema(
-            THROTTLE_TIME_MS,
-            METADATA_BROKER_V1,
-            CLUSTER_ID,
-            CONTROLLER_ID,
-            TOPIC_METADATA_V7);
-
-    public static Schema[] schemaVersions() {
-        return new Schema[] {METADATA_RESPONSE_V0, METADATA_RESPONSE_V1, METADATA_RESPONSE_V2, METADATA_RESPONSE_V3,
-            METADATA_RESPONSE_V4, METADATA_RESPONSE_V5, METADATA_RESPONSE_V6, METADATA_RESPONSE_V7};
-    }
+    public static final int AUTHORIZED_OPERATIONS_OMITTED = Integer.MIN_VALUE;
 
-    private final int throttleTimeMs;
-    private final Collection brokers;
-    private final Node controller;
-    private final List topicMetadata;
-    private final String clusterId;
+    private final MetadataResponseData data;
+    private volatile Holder holder;
+    private final boolean hasReliableLeaderEpochs;
 
-    /**
-     * Constructor for all versions.
-     */
-    public MetadataResponse(List brokers, String clusterId, int controllerId, List topicMetadata) {
-        this(DEFAULT_THROTTLE_TIME, brokers, clusterId, controllerId, topicMetadata);
+    public MetadataResponse(MetadataResponseData data) {
+        this(data, true);
     }
 
-    public MetadataResponse(int throttleTimeMs, List brokers, String clusterId, int controllerId, List topicMetadata) {
-        this.throttleTimeMs = throttleTimeMs;
-        this.brokers = brokers;
-        this.controller = getControllerNode(controllerId, brokers);
-        this.topicMetadata = topicMetadata;
-        this.clusterId = clusterId;
-    }
-
-    public MetadataResponse(Struct struct) {
-        this.throttleTimeMs = struct.getOrElse(THROTTLE_TIME_MS, DEFAULT_THROTTLE_TIME);
-        Map brokers = new HashMap<>();
-        Object[] brokerStructs = struct.get(BROKERS);
-        for (Object brokerStruct : brokerStructs) {
-            Struct broker = (Struct) brokerStruct;
-            int nodeId = broker.get(NODE_ID);
-            String host = broker.get(HOST);
-            int port = broker.get(PORT);
-            // This field only exists in v1+
-            // When we can't know if a rack exists in a v0 response we default to null
-            String rack = broker.getOrElse(RACK, null);
-            brokers.put(nodeId, new Node(nodeId, host, port, rack));
-        }
-
-        // This field only exists in v1+
-        // When we can't know the controller id in a v0 response we default to NO_CONTROLLER_ID
-        int controllerId = struct.getOrElse(CONTROLLER_ID, NO_CONTROLLER_ID);
-
-        // This field only exists in v2+
-        this.clusterId = struct.getOrElse(CLUSTER_ID, null);
-
-        List topicMetadata = new ArrayList<>();
-        Object[] topicInfos = struct.get(TOPIC_METADATA);
-        for (Object topicInfoObj : topicInfos) {
-            Struct topicInfo = (Struct) topicInfoObj;
-            Errors topicError = Errors.forCode(topicInfo.get(ERROR_CODE));
-            String topic = topicInfo.get(TOPIC_NAME);
-            // This field only exists in v1+
-            // When we can't know if a topic is internal or not in a v0 response we default to false
-            boolean isInternal = topicInfo.getOrElse(IS_INTERNAL, false);
-            List partitionMetadata = new ArrayList<>();
-
-            Object[] partitionInfos = topicInfo.get(PARTITION_METADATA);
-            for (Object partitionInfoObj : partitionInfos) {
-                Struct partitionInfo = (Struct) partitionInfoObj;
-                Errors partitionError = Errors.forCode(partitionInfo.get(ERROR_CODE));
-                int partition = partitionInfo.get(PARTITION_ID);
-                int leader = partitionInfo.get(LEADER);
-                Optional leaderEpoch = RequestUtils.getLeaderEpoch(partitionInfo, LEADER_EPOCH);
-                Node leaderNode = leader == -1 ? null : brokers.get(leader);
-
-                Object[] replicas = partitionInfo.get(REPLICAS);
-                List replicaNodes = convertToNodes(brokers, replicas);
-
-                Object[] isr = partitionInfo.get(ISR);
-                List isrNodes = convertToNodes(brokers, isr);
-
-                Object[] offlineReplicas = partitionInfo.getOrEmpty(OFFLINE_REPLICAS);
-                List offlineNodes = convertToNodes(brokers, offlineReplicas);
-
-                partitionMetadata.add(new PartitionMetadata(partitionError, partition, leaderNode, leaderEpoch,
-                        replicaNodes, isrNodes, offlineNodes));
-            }
-
-            topicMetadata.add(new TopicMetadata(topicError, topic, isInternal, partitionMetadata));
-        }
-
-        this.brokers = brokers.values();
-        this.controller = getControllerNode(controllerId, brokers.values());
-        this.topicMetadata = topicMetadata;
+    public MetadataResponse(Struct struct, short version) {
+        // Prior to Kafka version 2.4 (which coincides with Metadata version 9), the broker
+        // does not propagate leader epoch information accurately while a reassignment is in
+        // progress. Relying on a stale epoch can lead to FENCED_LEADER_EPOCH errors which
+        // can prevent consumption throughout the course of a reassignment. It is safer in
+        // this case to revert to the behavior in previous protocol versions which checks
+        // leader status only.
+        this(new MetadataResponseData(struct, version), version >= 9);
     }
 
-    private List convertToNodes(Map brokers, Object[] brokerIds) {
-        List nodes = new ArrayList<>(brokerIds.length);
-        for (Object brokerId : brokerIds)
-            if (brokers.containsKey(brokerId))
-                nodes.add(brokers.get(brokerId));
-            else
-                nodes.add(new Node((int) brokerId, "", -1));
-        return nodes;
+    private MetadataResponse(MetadataResponseData data, boolean hasReliableLeaderEpochs) {
+        this.data = data;
+        this.hasReliableLeaderEpochs = hasReliableLeaderEpochs;
     }
 
-    private Node getControllerNode(int controllerId, Collection brokers) {
-        for (Node broker : brokers) {
-            if (broker.id() == controllerId)
-                return broker;
-        }
-        return null;
+    @Override
+    protected Struct toStruct(short version) {
+        return data.toStruct(version);
     }
 
     @Override
     public int throttleTimeMs() {
-        return throttleTimeMs;
+        return data.throttleTimeMs();
     }
 
     /**
@@ -312,9 +96,9 @@ public int throttleTimeMs() {
      */
     public Map errors() {
         Map errors = new HashMap<>();
-        for (TopicMetadata metadata : topicMetadata) {
-            if (metadata.error != Errors.NONE)
-                errors.put(metadata.topic(), metadata.error);
+        for (MetadataResponseTopic metadata : data.topics()) {
+            if (metadata.errorCode() != Errors.NONE.code())
+                errors.put(metadata.name(), Errors.forCode(metadata.errorCode()));
         }
         return errors;
     }
@@ -322,8 +106,8 @@ public Map errors() {
     @Override
     public Map errorCounts() {
         Map errorCounts = new HashMap<>();
-        for (TopicMetadata metadata : topicMetadata)
-            updateErrorCounts(errorCounts, metadata.error);
+        for (MetadataResponseTopic metadata : data.topics())
+            updateErrorCounts(errorCounts, Errors.forCode(metadata.errorCode()));
         return errorCounts;
     }
 
@@ -332,9 +116,9 @@ public Map errorCounts() {
      */
     public Set topicsByError(Errors error) {
         Set errorTopics = new HashSet<>();
-        for (TopicMetadata metadata : topicMetadata) {
-            if (metadata.error == error)
-                errorTopics.add(metadata.topic());
+        for (MetadataResponseTopic metadata : data.topics()) {
+            if (metadata.errorCode() == error.code())
+                errorTopics.add(metadata.name());
         }
         return errorTopics;
     }
@@ -346,8 +130,7 @@ public Set topicsByError(Errors error) {
     public Cluster cluster() {
         Set internalTopics = new HashSet<>();
         List partitions = new ArrayList<>();
-        for (TopicMetadata metadata : topicMetadata) {
-
+        for (TopicMetadata metadata : topicMetadata()) {
             if (metadata.error == Errors.NONE) {
                 if (metadata.isInternal)
                     internalTopics.add(metadata.topic);
@@ -356,13 +139,30 @@ public Cluster cluster() {
                 }
             }
         }
-        return new Cluster(this.clusterId, this.brokers, partitions, topicsByError(Errors.TOPIC_AUTHORIZATION_FAILED),
-                topicsByError(Errors.INVALID_TOPIC_EXCEPTION), internalTopics, this.controller);
+        return new Cluster(data.clusterId(), brokers(), partitions, topicsByError(Errors.TOPIC_AUTHORIZATION_FAILED),
+                topicsByError(Errors.INVALID_TOPIC_EXCEPTION), internalTopics, controller());
+    }
+
+    /**
+     * Returns a 32-bit bitfield to represent authorized operations for this topic.
+     */
+    public Optional topicAuthorizedOperations(String topicName) {
+        MetadataResponseTopic topic = data.topics().find(topicName);
+        if (topic == null)
+            return Optional.empty();
+        else
+            return Optional.of(topic.topicAuthorizedOperations());
     }
 
     /**
-     * Transform a topic and PartitionMetadata into PartitionInfo
-     * @return
+     * Returns a 32-bit bitfield to represent authorized operations for this cluster.
+     */
+    public int clusterAuthorizedOperations() {
+        return data.clusterAuthorizedOperations();
+    }
+
+    /**
+     * Transform a topic and PartitionMetadata into PartitionInfo.
      */
     public static PartitionInfo partitionMetaToInfo(String topic, PartitionMetadata partitionMetadata) {
         return new PartitionInfo(
@@ -374,12 +174,30 @@ public static PartitionInfo partitionMetaToInfo(String topic, PartitionMetadata
                 partitionMetadata.offlineReplicas().toArray(new Node[0]));
     }
 
+    private Holder holder() {
+        if (holder == null) {
+            synchronized (data) {
+                if (holder == null)
+                    holder = new Holder(data);
+            }
+        }
+        return holder;
+    }
+
     /**
      * Get all brokers returned in metadata response
      * @return the brokers
      */
     public Collection brokers() {
-        return brokers;
+        return holder().brokers.values();
+    }
+
+    /**
+     * Get a map of all brokers keyed by the brokerId.
+     * @return A map from the brokerId to the broker `Node` information
+     */
+    public Map brokersById() {
+        return holder().brokers;
     }
 
     /**
@@ -387,7 +205,7 @@ public Collection brokers() {
      * @return the topicMetadata
      */
     public Collection topicMetadata() {
-        return topicMetadata;
+        return holder().topicMetadata;
     }
 
     /**
@@ -395,7 +213,7 @@ public Collection topicMetadata() {
      * @return the controller node or null if it doesn't exist
      */
     public Node controller() {
-        return controller;
+        return holder().controller;
     }
 
     /**
@@ -403,11 +221,23 @@ public Node controller() {
      * @return cluster identifier if it is present in the response, null otherwise.
      */
     public String clusterId() {
-        return this.clusterId;
+        return this.data.clusterId();
+    }
+
+    /**
+     * Check whether the leader epochs returned from the response can be relied on
+     * for epoch validation in Fetch, ListOffsets, and OffsetsForLeaderEpoch requests.
+     * If not, then the client will not retain the leader epochs and hence will not
+     * forward them in requests.
+     *
+     * @return true if the epoch can be used for validation
+     */
+    public boolean hasReliableLeaderEpochs() {
+        return hasReliableLeaderEpochs;
     }
 
     public static MetadataResponse parse(ByteBuffer buffer, short version) {
-        return new MetadataResponse(ApiKeys.METADATA.parseResponse(version, buffer));
+        return new MetadataResponse(ApiKeys.METADATA.responseSchema(version).read(buffer), version);
     }
 
     public static class TopicMetadata {
@@ -415,15 +245,25 @@ public static class TopicMetadata {
         private final String topic;
         private final boolean isInternal;
         private final List partitionMetadata;
+        private int authorizedOperations;
 
         public TopicMetadata(Errors error,
                              String topic,
                              boolean isInternal,
-                             List partitionMetadata) {
+                             List partitionMetadata,
+                             int authorizedOperations) {
             this.error = error;
             this.topic = topic;
             this.isInternal = isInternal;
             this.partitionMetadata = partitionMetadata;
+            this.authorizedOperations = authorizedOperations;
+        }
+
+        public TopicMetadata(Errors error,
+                             String topic,
+                             boolean isInternal,
+                             List partitionMetadata) {
+            this(error, topic, isInternal, partitionMetadata, 0);
         }
 
         public Errors error() {
@@ -442,13 +282,40 @@ public List partitionMetadata() {
             return partitionMetadata;
         }
 
+        public void authorizedOperations(int authorizedOperations) {
+            this.authorizedOperations = authorizedOperations;
+        }
+
+        public int authorizedOperations() {
+            return authorizedOperations;
+        }
+
+        @Override
+        public boolean equals(final Object o) {
+            if (this == o) return true;
+            if (o == null || getClass() != o.getClass()) return false;
+            final TopicMetadata that = (TopicMetadata) o;
+            return isInternal == that.isInternal &&
+                error == that.error &&
+                Objects.equals(topic, that.topic) &&
+                Objects.equals(partitionMetadata, that.partitionMetadata) &&
+                Objects.equals(authorizedOperations, that.authorizedOperations);
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(error, topic, isInternal, partitionMetadata, authorizedOperations);
+        }
+
         @Override
         public String toString() {
-            return "(type=TopicMetadata" +
-                    ", error=" + error +
-                    ", topic=" + topic +
-                    ", isInternal=" + isInternal +
-                    ", partitionMetadata=" + partitionMetadata + ')';
+            return "TopicMetadata{" +
+                "error=" + error +
+                ", topic='" + topic + '\'' +
+                ", isInternal=" + isInternal +
+                ", partitionMetadata=" + partitionMetadata +
+                ", authorizedOperations=" + authorizedOperations +
+                '}';
         }
     }
 
@@ -523,68 +390,112 @@ public String toString() {
         }
     }
 
-    @Override
-    protected Struct toStruct(short version) {
-        Struct struct = new Struct(ApiKeys.METADATA.responseSchema(version));
-        struct.setIfExists(THROTTLE_TIME_MS, throttleTimeMs);
-        List brokerArray = new ArrayList<>();
-        for (Node node : brokers) {
-            Struct broker = struct.instance(BROKERS);
-            broker.set(NODE_ID, node.id());
-            broker.set(HOST, node.host());
-            broker.set(PORT, node.port());
-            // This field only exists in v1+
-            broker.setIfExists(RACK, node.rack());
-            brokerArray.add(broker);
+    private static class Holder {
+        private final Map brokers;
+        private final Node controller;
+        private final Collection topicMetadata;
+
+        Holder(MetadataResponseData data) {
+            this.brokers = createBrokers(data);
+            this.topicMetadata = createTopicMetadata(data);
+            this.controller = brokers.get(data.controllerId());
+        }
+
+        private Map createBrokers(MetadataResponseData data) {
+            return data.brokers().valuesList().stream().map(b -> new Node(b.nodeId(), b.host(), b.port(), b.rack()))
+                    .collect(Collectors.toMap(Node::id, b -> b));
         }
-        struct.set(BROKERS, brokerArray.toArray());
-
-        // This field only exists in v1+
-        struct.setIfExists(CONTROLLER_ID, controller == null ? NO_CONTROLLER_ID : controller.id());
-
-        // This field only exists in v2+
-        struct.setIfExists(CLUSTER_ID, clusterId);
-
-        List topicMetadataArray = new ArrayList<>(topicMetadata.size());
-        for (TopicMetadata metadata : topicMetadata) {
-            Struct topicData = struct.instance(TOPIC_METADATA);
-            topicData.set(TOPIC_NAME, metadata.topic);
-            topicData.set(ERROR_CODE, metadata.error.code());
-            // This field only exists in v1+
-            topicData.setIfExists(IS_INTERNAL, metadata.isInternal());
-
-            List partitionMetadataArray = new ArrayList<>(metadata.partitionMetadata.size());
-            for (PartitionMetadata partitionMetadata : metadata.partitionMetadata()) {
-                Struct partitionData = topicData.instance(PARTITION_METADATA);
-                partitionData.set(ERROR_CODE, partitionMetadata.error.code());
-                partitionData.set(PARTITION_ID, partitionMetadata.partition);
-                partitionData.set(LEADER, partitionMetadata.leaderId());
-
-                // Leader epoch exists in v7 forward
-                RequestUtils.setLeaderEpochIfExists(partitionData, LEADER_EPOCH, partitionMetadata.leaderEpoch);
-
-                ArrayList replicas = new ArrayList<>(partitionMetadata.replicas.size());
-                for (Node node : partitionMetadata.replicas)
-                    replicas.add(node.id());
-                partitionData.set(REPLICAS, replicas.toArray());
-                ArrayList isr = new ArrayList<>(partitionMetadata.isr.size());
-                for (Node node : partitionMetadata.isr)
-                    isr.add(node.id());
-                partitionData.set(ISR, isr.toArray());
-                if (partitionData.hasField(OFFLINE_REPLICAS)) {
-                    ArrayList offlineReplicas = new ArrayList<>(partitionMetadata.offlineReplicas.size());
-                    for (Node node : partitionMetadata.offlineReplicas)
-                        offlineReplicas.add(node.id());
-                    partitionData.set(OFFLINE_REPLICAS, offlineReplicas.toArray());
+
+        private Collection createTopicMetadata(MetadataResponseData data) {
+            List topicMetadataList = new ArrayList<>();
+            for (MetadataResponseTopic topicMetadata : data.topics()) {
+                Errors topicError = Errors.forCode(topicMetadata.errorCode());
+                String topic = topicMetadata.name();
+                boolean isInternal = topicMetadata.isInternal();
+                List partitionMetadataList = new ArrayList<>();
+
+                for (MetadataResponsePartition partitionMetadata : topicMetadata.partitions()) {
+                    Errors partitionError = Errors.forCode(partitionMetadata.errorCode());
+                    int partitionIndex = partitionMetadata.partitionIndex();
+                    int leader = partitionMetadata.leaderId();
+                    Optional leaderEpoch = RequestUtils.getLeaderEpoch(partitionMetadata.leaderEpoch());
+                    Node leaderNode = leader == -1 ? null : brokers.get(leader);
+                    List replicaNodes = convertToNodes(partitionMetadata.replicaNodes());
+                    List isrNodes = convertToNodes(partitionMetadata.isrNodes());
+                    List offlineNodes = convertToNodes(partitionMetadata.offlineReplicas());
+                    partitionMetadataList.add(new PartitionMetadata(partitionError, partitionIndex, leaderNode, leaderEpoch,
+                            replicaNodes, isrNodes, offlineNodes));
                 }
-                partitionMetadataArray.add(partitionData);
 
+                topicMetadataList.add(new TopicMetadata(topicError, topic, isInternal, partitionMetadataList,
+                        topicMetadata.topicAuthorizedOperations()));
+            }
+            return topicMetadataList;
+        }
+
+        private List convertToNodes(List brokerIds) {
+            List nodes = new ArrayList<>(brokerIds.size());
+            for (Integer brokerId : brokerIds) {
+                Node node = brokers.get(brokerId);
+                if (node == null)
+                    nodes.add(new Node(brokerId, "", -1));
+                else
+                    nodes.add(node);
             }
-            topicData.set(PARTITION_METADATA, partitionMetadataArray.toArray());
-            topicMetadataArray.add(topicData);
+            return nodes;
         }
-        struct.set(TOPIC_METADATA, topicMetadataArray.toArray());
-        return struct;
+
+    }
+
+    public static MetadataResponse prepareResponse(int throttleTimeMs, Collection brokers, String clusterId,
+                                                   int controllerId, List topicMetadataList,
+                                                   int clusterAuthorizedOperations) {
+        MetadataResponseData responseData = new MetadataResponseData();
+        responseData.setThrottleTimeMs(throttleTimeMs);
+        brokers.forEach(broker ->
+            responseData.brokers().add(new MetadataResponseBroker()
+                .setNodeId(broker.id())
+                .setHost(broker.host())
+                .setPort(broker.port())
+                .setRack(broker.rack()))
+        );
+
+        responseData.setClusterId(clusterId);
+        responseData.setControllerId(controllerId);
+        responseData.setClusterAuthorizedOperations(clusterAuthorizedOperations);
+
+        topicMetadataList.forEach(topicMetadata -> {
+            MetadataResponseTopic metadataResponseTopic = new MetadataResponseTopic();
+            metadataResponseTopic
+                .setErrorCode(topicMetadata.error.code())
+                .setName(topicMetadata.topic)
+                .setIsInternal(topicMetadata.isInternal)
+                .setTopicAuthorizedOperations(topicMetadata.authorizedOperations);
+
+            for (PartitionMetadata partitionMetadata : topicMetadata.partitionMetadata) {
+                metadataResponseTopic.partitions().add(new MetadataResponsePartition()
+                    .setErrorCode(partitionMetadata.error.code())
+                    .setPartitionIndex(partitionMetadata.partition)
+                    .setLeaderId(partitionMetadata.leader == null ? -1 : partitionMetadata.leader.id())
+                    .setLeaderEpoch(partitionMetadata.leaderEpoch().orElse(RecordBatch.NO_PARTITION_LEADER_EPOCH))
+                    .setReplicaNodes(partitionMetadata.replicas.stream().map(Node::id).collect(Collectors.toList()))
+                    .setIsrNodes(partitionMetadata.isr.stream().map(Node::id).collect(Collectors.toList()))
+                    .setOfflineReplicas(partitionMetadata.offlineReplicas.stream().map(Node::id).collect(Collectors.toList())));
+            }
+            responseData.topics().add(metadataResponseTopic);
+        });
+        return new MetadataResponse(responseData);
+    }
+
+    public static MetadataResponse prepareResponse(int throttleTimeMs, Collection brokers, String clusterId,
+                                                   int controllerId, List topicMetadataList) {
+        return prepareResponse(throttleTimeMs, brokers, clusterId, controllerId, topicMetadataList,
+                MetadataResponse.AUTHORIZED_OPERATIONS_OMITTED);
+    }
+
+    public static MetadataResponse prepareResponse(Collection brokers, String clusterId, int controllerId,
+                                                   List topicMetadata) {
+        return prepareResponse(AbstractResponse.DEFAULT_THROTTLE_TIME, brokers, clusterId, controllerId, topicMetadata);
     }
 
     @Override
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/OffsetCommitRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/OffsetCommitRequest.java
index fdb5b1081996d..72ebef261f4df 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/OffsetCommitRequest.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/OffsetCommitRequest.java
@@ -17,126 +17,23 @@
 package org.apache.kafka.common.requests;
 
 import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.errors.UnsupportedVersionException;
+import org.apache.kafka.common.message.OffsetCommitRequestData;
+import org.apache.kafka.common.message.OffsetCommitRequestData.OffsetCommitRequestTopic;
+import org.apache.kafka.common.message.OffsetCommitResponseData;
+import org.apache.kafka.common.message.OffsetCommitResponseData.OffsetCommitResponsePartition;
+import org.apache.kafka.common.message.OffsetCommitResponseData.OffsetCommitResponseTopic;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.protocol.types.Field;
-import org.apache.kafka.common.protocol.types.Schema;
 import org.apache.kafka.common.protocol.types.Struct;
-import org.apache.kafka.common.utils.CollectionUtils;
 
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.Optional;
-
-import static org.apache.kafka.common.protocol.CommonFields.COMMITTED_LEADER_EPOCH;
-import static org.apache.kafka.common.protocol.CommonFields.COMMITTED_METADATA;
-import static org.apache.kafka.common.protocol.CommonFields.COMMITTED_OFFSET;
-import static org.apache.kafka.common.protocol.CommonFields.GENERATION_ID;
-import static org.apache.kafka.common.protocol.CommonFields.GROUP_ID;
-import static org.apache.kafka.common.protocol.CommonFields.MEMBER_ID;
-import static org.apache.kafka.common.protocol.CommonFields.PARTITION_ID;
-import static org.apache.kafka.common.protocol.CommonFields.TOPIC_NAME;
 
 public class OffsetCommitRequest extends AbstractRequest {
-    // top level fields
-    private static final Field.ComplexArray TOPICS = new Field.ComplexArray("topics",
-            "Topics to commit offsets");
-
-    // topic level fields
-    private static final Field.ComplexArray PARTITIONS = new Field.ComplexArray("partitions",
-            "Partitions to commit offsets");
-
-    // partition level fields
-    @Deprecated
-    private static final Field.Int64 COMMIT_TIMESTAMP = new Field.Int64("timestamp", "Timestamp of the commit");
-    private static final Field.Int64 RETENTION_TIME = new Field.Int64("retention_time",
-            "Time period in ms to retain the offset.");
-
-    private static final Field PARTITIONS_V0 = PARTITIONS.withFields(
-            PARTITION_ID,
-            COMMITTED_OFFSET,
-            COMMITTED_METADATA);
-
-    private static final Field TOPICS_V0 = TOPICS.withFields(
-            TOPIC_NAME,
-            PARTITIONS_V0);
-
-    private static final Schema OFFSET_COMMIT_REQUEST_V0 = new Schema(
-            GROUP_ID,
-            TOPICS_V0);
-
-    // V1 adds timestamp and group membership information (generation and memberId)
-    private static final Field PARTITIONS_V1 = PARTITIONS.withFields(
-            PARTITION_ID,
-            COMMITTED_OFFSET,
-            COMMIT_TIMESTAMP,
-            COMMITTED_METADATA);
-
-    private static final Field TOPICS_V1 = TOPICS.withFields(
-            TOPIC_NAME,
-            PARTITIONS_V1);
-
-    private static final Schema OFFSET_COMMIT_REQUEST_V1 = new Schema(
-            GROUP_ID,
-            GENERATION_ID,
-            MEMBER_ID,
-            TOPICS_V1);
-
-    // V2 adds retention time
-    private static final Field PARTITIONS_V2 = PARTITIONS.withFields(
-            PARTITION_ID,
-            COMMITTED_OFFSET,
-            COMMITTED_METADATA);
-
-    private static final Field TOPICS_V2 = TOPICS.withFields(
-            TOPIC_NAME,
-            PARTITIONS_V2);
-
-    private static final Schema OFFSET_COMMIT_REQUEST_V2 = new Schema(
-            GROUP_ID,
-            GENERATION_ID,
-            MEMBER_ID,
-            RETENTION_TIME,
-            TOPICS_V2);
-
-    // V3 adds throttle time
-    private static final Schema OFFSET_COMMIT_REQUEST_V3 = OFFSET_COMMIT_REQUEST_V2;
-
-    // V4 bump used to indicate that on quota violation brokers send out responses before throttling.
-    private static final Schema OFFSET_COMMIT_REQUEST_V4 = OFFSET_COMMIT_REQUEST_V3;
-
-    // V5 removes the retention time which is now controlled only by a broker configuration
-    private static final Schema OFFSET_COMMIT_REQUEST_V5 = new Schema(
-            GROUP_ID,
-            GENERATION_ID,
-            MEMBER_ID,
-            TOPICS_V2);
-
-    // V6 adds the leader epoch to the partition data
-    private static final Field PARTITIONS_V6 = PARTITIONS.withFields(
-            PARTITION_ID,
-            COMMITTED_OFFSET,
-            COMMITTED_LEADER_EPOCH,
-            COMMITTED_METADATA);
-
-    private static final Field TOPICS_V6 = TOPICS.withFields(
-            TOPIC_NAME,
-            PARTITIONS_V6);
-
-    private static final Schema OFFSET_COMMIT_REQUEST_V6 = new Schema(
-            GROUP_ID,
-            GENERATION_ID,
-            MEMBER_ID,
-            TOPICS_V6);
-
-    public static Schema[] schemaVersions() {
-        return new Schema[] {OFFSET_COMMIT_REQUEST_V0, OFFSET_COMMIT_REQUEST_V1, OFFSET_COMMIT_REQUEST_V2,
-            OFFSET_COMMIT_REQUEST_V3, OFFSET_COMMIT_REQUEST_V4, OFFSET_COMMIT_REQUEST_V5, OFFSET_COMMIT_REQUEST_V6};
-    }
-
     // default values for the current version
     public static final int DEFAULT_GENERATION_ID = -1;
     public static final String DEFAULT_MEMBER_ID = "";
@@ -147,221 +44,119 @@ public static Schema[] schemaVersions() {
     @Deprecated
     public static final long DEFAULT_TIMESTAMP = -1L;            // for V0, V1
 
-    private final String groupId;
-    private final String memberId;
-    private final int generationId;
-    private final long retentionTime;
-    private final Map offsetData;
-
-    public static final class PartitionData {
-        @Deprecated
-        public final long timestamp;                // for V1
-
-        public final long offset;
-        public final String metadata;
-        public final Optional leaderEpoch;
-
-        private PartitionData(long offset, Optional leaderEpoch, long timestamp, String metadata) {
-            this.offset = offset;
-            this.leaderEpoch = leaderEpoch;
-            this.timestamp = timestamp;
-            this.metadata = metadata;
-        }
-
-        @Deprecated
-        public PartitionData(long offset, long timestamp, String metadata) {
-            this(offset, Optional.empty(), timestamp, metadata);
-        }
-
-        public PartitionData(long offset, Optional leaderEpoch, String metadata) {
-            this(offset, leaderEpoch, DEFAULT_TIMESTAMP, metadata);
-        }
-
-        @Override
-        public String toString() {
-            StringBuilder bld = new StringBuilder();
-            bld.append("(timestamp=").append(timestamp).
-                    append(", offset=").append(offset).
-                    append(", leaderEpoch=").append(leaderEpoch).
-                    append(", metadata=").append(metadata).
-                    append(")");
-            return bld.toString();
-        }
-    }
+    private final OffsetCommitRequestData data;
 
     public static class Builder extends AbstractRequest.Builder {
-        private final String groupId;
-        private final Map offsetData;
-        private String memberId = DEFAULT_MEMBER_ID;
-        private int generationId = DEFAULT_GENERATION_ID;
 
-        public Builder(String groupId, Map offsetData) {
-            super(ApiKeys.OFFSET_COMMIT);
-            this.groupId = groupId;
-            this.offsetData = offsetData;
-        }
-
-        public Builder setMemberId(String memberId) {
-            this.memberId = memberId;
-            return this;
-        }
+        private final OffsetCommitRequestData data;
 
-        public Builder setGenerationId(int generationId) {
-            this.generationId = generationId;
-            return this;
+        public Builder(OffsetCommitRequestData data) {
+            super(ApiKeys.OFFSET_COMMIT);
+            this.data = data;
         }
 
         @Override
         public OffsetCommitRequest build(short version) {
-            if (version == 0) {
-                return new OffsetCommitRequest(groupId, DEFAULT_GENERATION_ID, DEFAULT_MEMBER_ID,
-                        DEFAULT_RETENTION_TIME, offsetData, version);
-            } else {
-                return new OffsetCommitRequest(groupId, generationId, memberId, DEFAULT_RETENTION_TIME,
-                        offsetData, version);
+            if (data.groupInstanceId() != null && version < 7) {
+                throw new UnsupportedVersionException("The broker offset commit protocol version " +
+                        version + " does not support usage of config group.instance.id.");
             }
+            return new OffsetCommitRequest(data, version);
         }
 
         @Override
         public String toString() {
-            StringBuilder bld = new StringBuilder();
-            bld.append("(type=OffsetCommitRequest").
-                append(", groupId=").append(groupId).
-                append(", memberId=").append(memberId).
-                append(", generationId=").append(generationId).
-                append(", offsetData=").append(offsetData).
-                append(")");
-            return bld.toString();
+            return data.toString();
         }
     }
 
-    private OffsetCommitRequest(String groupId, int generationId, String memberId, long retentionTime,
-                                Map offsetData, short version) {
+    private final short version;
+
+    public OffsetCommitRequest(OffsetCommitRequestData data, short version) {
         super(ApiKeys.OFFSET_COMMIT, version);
-        this.groupId = groupId;
-        this.generationId = generationId;
-        this.memberId = memberId;
-        this.retentionTime = retentionTime;
-        this.offsetData = offsetData;
+        this.data = data;
+        this.version = version;
     }
 
-    public OffsetCommitRequest(Struct struct, short versionId) {
-        super(ApiKeys.OFFSET_COMMIT, versionId);
 
-        groupId = struct.get(GROUP_ID);
-
-        // These fields only exists in v1.
-        generationId = struct.getOrElse(GENERATION_ID, DEFAULT_GENERATION_ID);
-        memberId = struct.getOrElse(MEMBER_ID, DEFAULT_MEMBER_ID);
+    public OffsetCommitRequest(Struct struct, short version) {
+        super(ApiKeys.OFFSET_COMMIT, version);
+        this.data = new OffsetCommitRequestData(struct, version);
+        this.version = version;
+    }
 
-        // This field only exists in v2
-        retentionTime = struct.getOrElse(RETENTION_TIME, DEFAULT_RETENTION_TIME);
+    public OffsetCommitRequestData data() {
+        return data;
+    }
 
-        offsetData = new HashMap<>();
-        for (Object topicDataObj : struct.get(TOPICS)) {
-            Struct topicData = (Struct) topicDataObj;
-            String topic = topicData.get(TOPIC_NAME);
-            for (Object partitionDataObj : topicData.get(PARTITIONS)) {
-                Struct partitionDataStruct = (Struct) partitionDataObj;
-                int partition = partitionDataStruct.get(PARTITION_ID);
-                long offset = partitionDataStruct.get(COMMITTED_OFFSET);
-                String metadata = partitionDataStruct.get(COMMITTED_METADATA);
-                PartitionData partitionOffset;
-                // This field only exists in v1
-                if (partitionDataStruct.hasField(COMMIT_TIMESTAMP)) {
-                    long timestamp = partitionDataStruct.get(COMMIT_TIMESTAMP);
-                    partitionOffset = new PartitionData(offset, timestamp, metadata);
-                } else {
-                    Optional leaderEpochOpt = RequestUtils.getLeaderEpoch(partitionDataStruct,
-                            COMMITTED_LEADER_EPOCH);
-                    partitionOffset = new PartitionData(offset, leaderEpochOpt, metadata);
-                }
-                offsetData.put(new TopicPartition(topic, partition), partitionOffset);
+    public Map offsets() {
+        Map offsets = new HashMap<>();
+        for (OffsetCommitRequestTopic topic : data.topics()) {
+            for (OffsetCommitRequestData.OffsetCommitRequestPartition partition : topic.partitions()) {
+                offsets.put(new TopicPartition(topic.name(), partition.partitionIndex()),
+                        partition.committedOffset());
             }
         }
+        return offsets;
     }
 
-    @Override
-    public Struct toStruct() {
-        short version = version();
-        Struct struct = new Struct(ApiKeys.OFFSET_COMMIT.requestSchema(version));
-        struct.set(GROUP_ID, groupId);
-
-        Map> topicsData = CollectionUtils.groupPartitionDataByTopic(offsetData);
-        List topicArray = new ArrayList<>();
-        for (Map.Entry> topicEntry: topicsData.entrySet()) {
-            Struct topicData = struct.instance(TOPICS);
-            topicData.set(TOPIC_NAME, topicEntry.getKey());
-            List partitionArray = new ArrayList<>();
-            for (Map.Entry partitionEntry : topicEntry.getValue().entrySet()) {
-                PartitionData fetchPartitionData = partitionEntry.getValue();
-                Struct partitionData = topicData.instance(PARTITIONS);
-                partitionData.set(PARTITION_ID, partitionEntry.getKey());
-                partitionData.set(COMMITTED_OFFSET, fetchPartitionData.offset);
-                // Only for v1
-                partitionData.setIfExists(COMMIT_TIMESTAMP, fetchPartitionData.timestamp);
-                // Only for v6
-                RequestUtils.setLeaderEpochIfExists(partitionData, COMMITTED_LEADER_EPOCH, fetchPartitionData.leaderEpoch);
-                partitionData.set(COMMITTED_METADATA, fetchPartitionData.metadata);
-                partitionArray.add(partitionData);
+    public static List getErrorResponseTopics(
+            List requestTopics,
+            Errors e) {
+        List responseTopicData = new ArrayList<>();
+        for (OffsetCommitRequestTopic entry : requestTopics) {
+            List responsePartitions =
+                    new ArrayList<>();
+            for (OffsetCommitRequestData.OffsetCommitRequestPartition requestPartition : entry.partitions()) {
+                responsePartitions.add(new OffsetCommitResponsePartition()
+                                           .setPartitionIndex(requestPartition.partitionIndex())
+                                           .setErrorCode(e.code()));
             }
-            topicData.set(PARTITIONS, partitionArray.toArray());
-            topicArray.add(topicData);
+            responseTopicData.add(new OffsetCommitResponseTopic()
+                    .setName(entry.name())
+                    .setPartitions(responsePartitions)
+            );
         }
-        struct.set(TOPICS, topicArray.toArray());
-        struct.setIfExists(GENERATION_ID, generationId);
-        struct.setIfExists(MEMBER_ID, memberId);
-        struct.setIfExists(RETENTION_TIME, retentionTime);
-        return struct;
+        return responseTopicData;
     }
 
     @Override
-    public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
-        Map responseData = new HashMap<>();
-        for (Map.Entry entry: offsetData.entrySet()) {
-            responseData.put(entry.getKey(), Errors.forException(e));
-        }
+    public OffsetCommitResponse getErrorResponse(int throttleTimeMs, Throwable e) {
+        List
+                responseTopicData = getErrorResponseTopics(data.topics(), Errors.forException(e));
 
         short versionId = version();
         switch (versionId) {
             case 0:
             case 1:
             case 2:
-                return new OffsetCommitResponse(responseData);
+                return new OffsetCommitResponse(
+                        new OffsetCommitResponseData()
+                                .setTopics(responseTopicData)
+                );
             case 3:
             case 4:
             case 5:
             case 6:
-                return new OffsetCommitResponse(throttleTimeMs, responseData);
+            case 7:
+                return new OffsetCommitResponse(
+                        new OffsetCommitResponseData()
+                                .setTopics(responseTopicData)
+                                .setThrottleTimeMs(throttleTimeMs)
+
+                );
             default:
                 throw new IllegalArgumentException(String.format("Version %d is not valid. Valid versions for %s are 0 to %d",
                         versionId, this.getClass().getSimpleName(), ApiKeys.OFFSET_COMMIT.latestVersion()));
         }
     }
 
-    public String groupId() {
-        return groupId;
-    }
-
-    public int generationId() {
-        return generationId;
-    }
-
-    public String memberId() {
-        return memberId;
-    }
-
-    @Deprecated
-    public long retentionTime() {
-        return retentionTime;
-    }
-
-    public Map offsetData() {
-        return offsetData;
+    public static OffsetCommitRequest parse(ByteBuffer buffer, short version) {
+        return new OffsetCommitRequest(ApiKeys.OFFSET_COMMIT.parseRequest(version, buffer), version);
     }
 
-    public static OffsetCommitRequest parse(ByteBuffer buffer, short version) {
-        Schema schema = ApiKeys.OFFSET_COMMIT.requestSchema(version);
-        return new OffsetCommitRequest(schema.read(buffer), version);
+    @Override
+    protected Struct toStruct() {
+        return data.toStruct(version);
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/OffsetCommitResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/OffsetCommitResponse.java
index 4d724a32813bc..3adeb157c188c 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/OffsetCommitResponse.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/OffsetCommitResponse.java
@@ -17,12 +17,12 @@
 package org.apache.kafka.common.requests;
 
 import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.message.OffsetCommitResponseData;
+import org.apache.kafka.common.message.OffsetCommitResponseData.OffsetCommitResponsePartition;
+import org.apache.kafka.common.message.OffsetCommitResponseData.OffsetCommitResponseTopic;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.protocol.types.Field;
-import org.apache.kafka.common.protocol.types.Schema;
 import org.apache.kafka.common.protocol.types.Struct;
-import org.apache.kafka.common.utils.CollectionUtils;
 
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
@@ -30,139 +30,97 @@
 import java.util.List;
 import java.util.Map;
 
-import static org.apache.kafka.common.protocol.CommonFields.ERROR_CODE;
-import static org.apache.kafka.common.protocol.CommonFields.PARTITION_ID;
-import static org.apache.kafka.common.protocol.CommonFields.THROTTLE_TIME_MS;
-import static org.apache.kafka.common.protocol.CommonFields.TOPIC_NAME;
-
 /**
  * Possible error codes:
  *
- * UNKNOWN_TOPIC_OR_PARTITION (3)
- * REQUEST_TIMED_OUT (7)
- * OFFSET_METADATA_TOO_LARGE (12)
- * COORDINATOR_LOAD_IN_PROGRESS (14)
- * GROUP_COORDINATOR_NOT_AVAILABLE (15)
- * NOT_COORDINATOR (16)
- * ILLEGAL_GENERATION (22)
- * UNKNOWN_MEMBER_ID (25)
- * REBALANCE_IN_PROGRESS (27)
- * INVALID_COMMIT_OFFSET_SIZE (28)
- * TOPIC_AUTHORIZATION_FAILED (29)
- * GROUP_AUTHORIZATION_FAILED (30)
+ *   - {@link Errors#UNKNOWN_TOPIC_OR_PARTITION}
+ *   - {@link Errors#REQUEST_TIMED_OUT}
+ *   - {@link Errors#OFFSET_METADATA_TOO_LARGE}
+ *   - {@link Errors#COORDINATOR_LOAD_IN_PROGRESS}
+ *   - {@link Errors#COORDINATOR_NOT_AVAILABLE}
+ *   - {@link Errors#NOT_COORDINATOR}
+ *   - {@link Errors#ILLEGAL_GENERATION}
+ *   - {@link Errors#UNKNOWN_MEMBER_ID}
+ *   - {@link Errors#REBALANCE_IN_PROGRESS}
+ *   - {@link Errors#INVALID_COMMIT_OFFSET_SIZE}
+ *   - {@link Errors#TOPIC_AUTHORIZATION_FAILED}
+ *   - {@link Errors#GROUP_AUTHORIZATION_FAILED}
  */
 public class OffsetCommitResponse extends AbstractResponse {
-    private static final Field.ComplexArray TOPICS = new Field.ComplexArray("responses",
-            "Responses by topic for committed partitions");
-
-    // topic level fields
-    private static final Field.ComplexArray PARTITIONS = new Field.ComplexArray("partition_responses",
-            "Responses for committed partitions");
-
-    private static final Field PARTITIONS_V0 = PARTITIONS.withFields(
-            PARTITION_ID,
-            ERROR_CODE);
-
-    private static final Field TOPICS_V0 = TOPICS.withFields(
-            TOPIC_NAME,
-            PARTITIONS_V0);
 
-    private static final Schema OFFSET_COMMIT_RESPONSE_V0 = new Schema(
-            TOPICS_V0);
+    private final OffsetCommitResponseData data;
 
-    // V1 adds timestamp and group membership information (generation and memberId) to the request
-    private static final Schema OFFSET_COMMIT_RESPONSE_V1 = OFFSET_COMMIT_RESPONSE_V0;
-
-    // V2 adds retention time to the request
-    private static final Schema OFFSET_COMMIT_RESPONSE_V2 = OFFSET_COMMIT_RESPONSE_V0;
+    public OffsetCommitResponse(OffsetCommitResponseData data) {
+        this.data = data;
+    }
 
-    // V3 adds throttle time
-    private static final Schema OFFSET_COMMIT_RESPONSE_V3 = new Schema(
-            THROTTLE_TIME_MS,
-            TOPICS_V0);
+    public OffsetCommitResponse(int requestThrottleMs, Map responseData) {
+        Map
+                responseTopicDataMap = new HashMap<>();
 
-    // V4 bump used to indicate that on quota violation brokers send out responses before throttling.
-    private static final Schema OFFSET_COMMIT_RESPONSE_V4 = OFFSET_COMMIT_RESPONSE_V3;
+        for (Map.Entry entry : responseData.entrySet()) {
+            TopicPartition topicPartition = entry.getKey();
+            String topicName = topicPartition.topic();
 
-    // V5 removes retention time from the request
-    private static final Schema OFFSET_COMMIT_RESPONSE_V5 = OFFSET_COMMIT_RESPONSE_V4;
+            OffsetCommitResponseTopic topic = responseTopicDataMap.getOrDefault(
+                topicName, new OffsetCommitResponseTopic().setName(topicName));
 
-    // V6 adds leader epoch to the request
-    private static final Schema OFFSET_COMMIT_RESPONSE_V6 = OFFSET_COMMIT_RESPONSE_V5;
+            topic.partitions().add(new OffsetCommitResponsePartition()
+                                       .setErrorCode(entry.getValue().code())
+                                       .setPartitionIndex(topicPartition.partition()));
+            responseTopicDataMap.put(topicName, topic);
+        }
 
-    public static Schema[] schemaVersions() {
-        return new Schema[] {OFFSET_COMMIT_RESPONSE_V0, OFFSET_COMMIT_RESPONSE_V1, OFFSET_COMMIT_RESPONSE_V2,
-            OFFSET_COMMIT_RESPONSE_V3, OFFSET_COMMIT_RESPONSE_V4, OFFSET_COMMIT_RESPONSE_V5, OFFSET_COMMIT_RESPONSE_V6};
+        data = new OffsetCommitResponseData()
+                .setTopics(new ArrayList<>(responseTopicDataMap.values()))
+                .setThrottleTimeMs(requestThrottleMs);
     }
 
-    private final Map responseData;
-    private final int throttleTimeMs;
-
     public OffsetCommitResponse(Map responseData) {
         this(DEFAULT_THROTTLE_TIME, responseData);
     }
 
-    public OffsetCommitResponse(int throttleTimeMs, Map responseData) {
-        this.throttleTimeMs = throttleTimeMs;
-        this.responseData = responseData;
+    public OffsetCommitResponse(Struct struct) {
+        short latestVersion = (short) (OffsetCommitResponseData.SCHEMAS.length - 1);
+        this.data = new OffsetCommitResponseData(struct, latestVersion);
     }
 
-    public OffsetCommitResponse(Struct struct) {
-        this.throttleTimeMs = struct.getOrElse(THROTTLE_TIME_MS, DEFAULT_THROTTLE_TIME);
-        responseData = new HashMap<>();
-        for (Object topicResponseObj : struct.get(TOPICS)) {
-            Struct topicResponse = (Struct) topicResponseObj;
-            String topic = topicResponse.get(TOPIC_NAME);
-            for (Object partitionResponseObj : topicResponse.get(PARTITIONS)) {
-                Struct partitionResponse = (Struct) partitionResponseObj;
-                int partition = partitionResponse.get(PARTITION_ID);
-                Errors error = Errors.forCode(partitionResponse.get(ERROR_CODE));
-                responseData.put(new TopicPartition(topic, partition), error);
-            }
-        }
+    public OffsetCommitResponse(Struct struct, short version) {
+        this.data = new OffsetCommitResponseData(struct, version);
+    }
+
+    public OffsetCommitResponseData data() {
+        return data;
     }
 
     @Override
-    public Struct toStruct(short version) {
-        Struct struct = new Struct(ApiKeys.OFFSET_COMMIT.responseSchema(version));
-        struct.setIfExists(THROTTLE_TIME_MS, throttleTimeMs);
-
-        Map> topicsData = CollectionUtils.groupPartitionDataByTopic(responseData);
-        List topicArray = new ArrayList<>();
-        for (Map.Entry> entries: topicsData.entrySet()) {
-            Struct topicData = struct.instance(TOPICS);
-            topicData.set(TOPIC_NAME, entries.getKey());
-            List partitionArray = new ArrayList<>();
-            for (Map.Entry partitionEntry : entries.getValue().entrySet()) {
-                Struct partitionData = topicData.instance(PARTITIONS);
-                partitionData.set(PARTITION_ID, partitionEntry.getKey());
-                partitionData.set(ERROR_CODE, partitionEntry.getValue().code());
-                partitionArray.add(partitionData);
+    public Map errorCounts() {
+        List errors = new ArrayList<>();
+        for (OffsetCommitResponseTopic topic : data.topics()) {
+            for (OffsetCommitResponsePartition partition : topic.partitions()) {
+                errors.add(Errors.forCode(partition.errorCode()));
             }
-            topicData.set(PARTITIONS, partitionArray.toArray());
-            topicArray.add(topicData);
         }
-        struct.set(TOPICS, topicArray.toArray());
-
-        return struct;
+        return errorCounts(errors);
     }
 
-    @Override
-    public int throttleTimeMs() {
-        return throttleTimeMs;
+    public static OffsetCommitResponse parse(ByteBuffer buffer, short version) {
+        return new OffsetCommitResponse(ApiKeys.OFFSET_COMMIT.parseResponse(version, buffer), version);
     }
 
-    public Map responseData() {
-        return responseData;
+    @Override
+    public Struct toStruct(short version) {
+        return data.toStruct(version);
     }
 
     @Override
-    public Map errorCounts() {
-        return errorCounts(responseData);
+    public String toString() {
+        return data.toString();
     }
 
-    public static OffsetCommitResponse parse(ByteBuffer buffer, short version) {
-        return new OffsetCommitResponse(ApiKeys.OFFSET_COMMIT.parseResponse(version, buffer));
+    @Override
+    public int throttleTimeMs() {
+        return data.throttleTimeMs();
     }
 
     @Override
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/OffsetDeleteRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/OffsetDeleteRequest.java
new file mode 100644
index 0000000000000..293efc2f6f2db
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/OffsetDeleteRequest.java
@@ -0,0 +1,83 @@
+/*
+ * 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.requests;
+
+import org.apache.kafka.common.message.OffsetDeleteRequestData;
+import org.apache.kafka.common.message.OffsetDeleteResponseData;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.common.protocol.types.Struct;
+
+import java.nio.ByteBuffer;
+
+public class OffsetDeleteRequest extends AbstractRequest {
+
+    public static class Builder extends AbstractRequest.Builder {
+
+        private final OffsetDeleteRequestData data;
+
+        public Builder(OffsetDeleteRequestData data) {
+            super(ApiKeys.OFFSET_DELETE);
+            this.data = data;
+        }
+
+        @Override
+        public OffsetDeleteRequest build(short version) {
+            return new OffsetDeleteRequest(data, version);
+        }
+
+        @Override
+        public String toString() {
+            return data.toString();
+        }
+    }
+
+    public final OffsetDeleteRequestData data;
+
+    public OffsetDeleteRequest(OffsetDeleteRequestData data, short version) {
+        super(ApiKeys.OFFSET_DELETE, version);
+        this.data = data;
+    }
+
+    public OffsetDeleteRequest(Struct struct, short version) {
+        super(ApiKeys.OFFSET_DELETE, version);
+        this.data = new OffsetDeleteRequestData(struct, version);
+    }
+
+    public AbstractResponse getErrorResponse(int throttleTimeMs, Errors error) {
+        return new OffsetDeleteResponse(
+            new OffsetDeleteResponseData()
+                .setThrottleTimeMs(throttleTimeMs)
+                .setErrorCode(error.code())
+        );
+    }
+
+    @Override
+    public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
+        return getErrorResponse(throttleTimeMs, Errors.forException(e));
+    }
+
+    public static OffsetDeleteRequest parse(ByteBuffer buffer, short version) {
+        return new OffsetDeleteRequest(ApiKeys.OFFSET_DELETE.parseRequest(version, buffer),
+            version);
+    }
+
+    @Override
+    protected Struct toStruct() {
+        return data.toStruct(version());
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/OffsetDeleteResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/OffsetDeleteResponse.java
new file mode 100644
index 0000000000000..ff0c2967be890
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/OffsetDeleteResponse.java
@@ -0,0 +1,95 @@
+/*
+ * 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.requests;
+
+import org.apache.kafka.common.message.OffsetDeleteResponseData;
+import org.apache.kafka.common.message.OffsetDeleteResponseData.OffsetDeleteResponsePartition;
+import org.apache.kafka.common.message.OffsetDeleteResponseData.OffsetDeleteResponseTopic;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.common.protocol.types.Struct;
+
+import java.nio.ByteBuffer;
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * Possible error codes:
+ *
+ * - Partition errors:
+ *   - {@link Errors#GROUP_SUBSCRIBED_TO_TOPIC}
+ *   - {@link Errors#TOPIC_AUTHORIZATION_FAILED}
+ *   - {@link Errors#UNKNOWN_TOPIC_OR_PARTITION}
+ *
+ * - Group or coordinator errors:
+ *   - {@link Errors#COORDINATOR_LOAD_IN_PROGRESS}
+ *   - {@link Errors#COORDINATOR_NOT_AVAILABLE}
+ *   - {@link Errors#NOT_COORDINATOR}
+ *   - {@link Errors#GROUP_AUTHORIZATION_FAILED}
+ *   - {@link Errors#INVALID_GROUP_ID}
+ *   - {@link Errors#GROUP_ID_NOT_FOUND}
+ *   - {@link Errors#NON_EMPTY_GROUP}
+ */
+public class OffsetDeleteResponse extends AbstractResponse {
+
+    public final OffsetDeleteResponseData data;
+
+    public OffsetDeleteResponse(OffsetDeleteResponseData data) {
+        this.data = data;
+    }
+
+    public OffsetDeleteResponse(Struct struct) {
+        short latestVersion = (short) (OffsetDeleteResponseData.SCHEMAS.length - 1);
+        this.data = new OffsetDeleteResponseData(struct, latestVersion);
+    }
+
+    public OffsetDeleteResponse(Struct struct, short version) {
+        this.data = new OffsetDeleteResponseData(struct, version);
+    }
+
+    @Override
+    protected Struct toStruct(short version) {
+        return data.toStruct(version);
+    }
+
+    @Override
+    public Map errorCounts() {
+        Map counts = new HashMap<>();
+        counts.put(Errors.forCode(data.errorCode()), 1);
+        for (OffsetDeleteResponseTopic topic : data.topics()) {
+            for (OffsetDeleteResponsePartition partition : topic.partitions()) {
+                Errors error = Errors.forCode(partition.errorCode());
+                counts.put(error, counts.getOrDefault(error, 0) + 1);
+            }
+        }
+        return counts;
+    }
+
+    public static OffsetDeleteResponse parse(ByteBuffer buffer, short version) {
+        return new OffsetDeleteResponse(ApiKeys.OFFSET_DELETE.parseResponse(version, buffer));
+    }
+
+    @Override
+    public int throttleTimeMs() {
+        return data.throttleTimeMs();
+    }
+
+    @Override
+    public boolean shouldClientThrottle(short version) {
+        return version >= 0;
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/OffsetFetchRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/OffsetFetchRequest.java
index d2f5c888cec19..2303f89f91984 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/OffsetFetchRequest.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/OffsetFetchRequest.java
@@ -18,13 +18,11 @@
 
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.errors.UnsupportedVersionException;
+import org.apache.kafka.common.message.OffsetFetchRequestData;
+import org.apache.kafka.common.message.OffsetFetchRequestData.OffsetFetchRequestTopic;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.protocol.types.Field;
-import org.apache.kafka.common.protocol.types.Schema;
 import org.apache.kafka.common.protocol.types.Struct;
-import org.apache.kafka.common.utils.CollectionUtils;
-import org.apache.kafka.common.utils.Utils;
 
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
@@ -33,84 +31,45 @@
 import java.util.Map;
 import java.util.Optional;
 
-import static org.apache.kafka.common.protocol.CommonFields.GROUP_ID;
-import static org.apache.kafka.common.protocol.CommonFields.PARTITION_ID;
-import static org.apache.kafka.common.protocol.CommonFields.TOPIC_NAME;
-
 public class OffsetFetchRequest extends AbstractRequest {
-    // top level fields
-    private static final Field.ComplexArray TOPICS = new Field.ComplexArray("topics",
-            "Topics to fetch offsets. If the topic array is null fetch offsets for all topics.");
-
-    // topic level fields
-    private static final Field.ComplexArray PARTITIONS = new Field.ComplexArray("partitions",
-            "Partitions to fetch offsets.");
-
-    /*
-     * Wire formats of version 0 and 1 are the same, but with different functionality.
-     * Wire format of version 2 is similar to version 1, with the exception of
-     * - accepting 'null' as list of topics
-     * - returning a top level error code
-     * Version 0 will read the offsets from ZK.
-     * Version 1 will read the offsets from Kafka.
-     * Version 2 will read the offsets from Kafka, and returns all associated topic partition offsets if
-     * a 'null' is passed instead of a list of specific topic partitions. It also returns a top level error code
-     * for group or coordinator level errors.
-     */
-    private static final Field PARTITIONS_V0 = PARTITIONS.withFields(
-            PARTITION_ID);
-
-    private static final Field TOPICS_V0 = TOPICS.withFields("Topics to fetch offsets.",
-            TOPIC_NAME,
-            PARTITIONS_V0);
-
-    private static final Schema OFFSET_FETCH_REQUEST_V0 = new Schema(
-            GROUP_ID,
-            TOPICS_V0);
-
-    // V1 begins support for fetching offsets from the internal __consumer_offsets topic
-    private static final Schema OFFSET_FETCH_REQUEST_V1 = OFFSET_FETCH_REQUEST_V0;
-
-    // V2 adds top-level error code to the response as well as allowing a null offset array to indicate fetch
-    // of all committed offsets for a group
-    private static final Field TOPICS_V2 = TOPICS.nullableWithFields(
-            TOPIC_NAME,
-            PARTITIONS_V0);
-    private static final Schema OFFSET_FETCH_REQUEST_V2 = new Schema(
-            GROUP_ID,
-            TOPICS_V2);
-
-    // V3 request is the same as v2. Throttle time has been added to v3 response
-    private static final Schema OFFSET_FETCH_REQUEST_V3 = OFFSET_FETCH_REQUEST_V2;
-
-    // V4 bump used to indicate that on quota violation brokers send out responses before throttling.
-    private static final Schema OFFSET_FETCH_REQUEST_V4 = OFFSET_FETCH_REQUEST_V3;
-
-    // V5 adds the leader epoch of the committed offset in the response
-    private static final Schema OFFSET_FETCH_REQUEST_V5 = OFFSET_FETCH_REQUEST_V4;
-
-    public static Schema[] schemaVersions() {
-        return new Schema[] {OFFSET_FETCH_REQUEST_V0, OFFSET_FETCH_REQUEST_V1, OFFSET_FETCH_REQUEST_V2,
-            OFFSET_FETCH_REQUEST_V3, OFFSET_FETCH_REQUEST_V4, OFFSET_FETCH_REQUEST_V5};
-    }
+
+    private static final List ALL_TOPIC_PARTITIONS = null;
+    public final OffsetFetchRequestData data;
 
     public static class Builder extends AbstractRequest.Builder {
-        private static final List ALL_TOPIC_PARTITIONS = null;
-        private final String groupId;
-        private final List partitions;
+
+        public final OffsetFetchRequestData data;
 
         public Builder(String groupId, List partitions) {
             super(ApiKeys.OFFSET_FETCH);
-            this.groupId = groupId;
-            this.partitions = partitions;
+
+            final List topics;
+            if (partitions != null) {
+                Map offsetFetchRequestTopicMap = new HashMap<>();
+                for (TopicPartition topicPartition : partitions) {
+                    String topicName = topicPartition.topic();
+                    OffsetFetchRequestTopic topic = offsetFetchRequestTopicMap.getOrDefault(
+                        topicName, new OffsetFetchRequestTopic().setName(topicName));
+                    topic.partitionIndexes().add(topicPartition.partition());
+                    offsetFetchRequestTopicMap.put(topicName, topic);
+                }
+                topics = new ArrayList<>(offsetFetchRequestTopicMap.values());
+            } else {
+                // If passed in partition list is null, it is requesting offsets for all topic partitions.
+                topics = ALL_TOPIC_PARTITIONS;
+            }
+
+            this.data = new OffsetFetchRequestData()
+                            .setGroupId(groupId)
+                            .setTopics(topics);
         }
 
         public static Builder allTopicPartitions(String groupId) {
-            return new Builder(groupId, ALL_TOPIC_PARTITIONS);
+            return new Builder(groupId, null);
         }
 
         public boolean isAllTopicPartitions() {
-            return this.partitions == ALL_TOPIC_PARTITIONS;
+            return this.data.topics() == ALL_TOPIC_PARTITIONS;
         }
 
         @Override
@@ -118,54 +77,40 @@ public OffsetFetchRequest build(short version) {
             if (isAllTopicPartitions() && version < 2)
                 throw new UnsupportedVersionException("The broker only supports OffsetFetchRequest " +
                         "v" + version + ", but we need v2 or newer to request all topic partitions.");
-            return new OffsetFetchRequest(groupId, partitions, version);
+            return new OffsetFetchRequest(data, version);
         }
 
         @Override
         public String toString() {
-            StringBuilder bld = new StringBuilder();
-            String partitionsString = partitions == null ? "" : Utils.join(partitions, ",");
-            bld.append("(type=OffsetFetchRequest, ").
-                    append("groupId=").append(groupId).
-                    append(", partitions=").append(partitionsString).
-                    append(")");
-            return bld.toString();
+            return data.toString();
         }
     }
 
-    private final String groupId;
-    private final List partitions;
+    public List partitions() {
+        if (isAllPartitions()) {
+            return null;
+        }
+        List partitions = new ArrayList<>();
+        for (OffsetFetchRequestTopic topic : data.topics()) {
+            for (Integer partitionIndex : topic.partitionIndexes()) {
+                partitions.add(new TopicPartition(topic.name(), partitionIndex));
+            }
+        }
+        return partitions;
+    }
 
-    public static OffsetFetchRequest forAllPartitions(String groupId) {
-        return new OffsetFetchRequest.Builder(groupId, null).build((short) 2);
+    public String groupId() {
+        return data.groupId();
     }
 
-    private OffsetFetchRequest(String groupId, List partitions, short version) {
+    private OffsetFetchRequest(OffsetFetchRequestData data, short version) {
         super(ApiKeys.OFFSET_FETCH, version);
-        this.groupId = groupId;
-        this.partitions = partitions;
+        this.data = data;
     }
 
     public OffsetFetchRequest(Struct struct, short version) {
         super(ApiKeys.OFFSET_FETCH, version);
-
-        Object[] topicArray = struct.get(TOPICS);
-        if (topicArray != null) {
-            partitions = new ArrayList<>();
-            for (Object topicResponseObj : topicArray) {
-                Struct topicResponse = (Struct) topicResponseObj;
-                String topic = topicResponse.get(TOPIC_NAME);
-                for (Object partitionResponseObj : topicResponse.get(PARTITIONS)) {
-                    Struct partitionResponse = (Struct) partitionResponseObj;
-                    int partition = partitionResponse.get(PARTITION_ID);
-                    partitions.add(new TopicPartition(topic, partition));
-                }
-            }
-        } else
-            partitions = null;
-
-
-        groupId = struct.get(GROUP_ID);
+        this.data = new OffsetFetchRequestData(struct, version);
     }
 
     public OffsetFetchResponse getErrorResponse(Errors error) {
@@ -173,32 +118,26 @@ public OffsetFetchResponse getErrorResponse(Errors error) {
     }
 
     public OffsetFetchResponse getErrorResponse(int throttleTimeMs, Errors error) {
-        short versionId = version();
-
         Map responsePartitions = new HashMap<>();
-        if (versionId < 2) {
+        if (version() < 2) {
             OffsetFetchResponse.PartitionData partitionError = new OffsetFetchResponse.PartitionData(
                     OffsetFetchResponse.INVALID_OFFSET,
                     Optional.empty(),
                     OffsetFetchResponse.NO_METADATA,
                     error);
 
-            for (TopicPartition partition : this.partitions)
-                responsePartitions.put(partition, partitionError);
+            for (OffsetFetchRequestTopic topic : this.data.topics()) {
+                for (int partitionIndex : topic.partitionIndexes()) {
+                    responsePartitions.put(
+                        new TopicPartition(topic.name(), partitionIndex), partitionError);
+                }
+            }
         }
 
-        switch (versionId) {
-            case 0:
-            case 1:
-            case 2:
-                return new OffsetFetchResponse(error, responsePartitions);
-            case 3:
-            case 4:
-            case 5:
-                return new OffsetFetchResponse(throttleTimeMs, error, responsePartitions);
-            default:
-                throw new IllegalArgumentException(String.format("Version %d is not valid. Valid versions for %s are 0 to %d",
-                        versionId, this.getClass().getSimpleName(), ApiKeys.OFFSET_FETCH.latestVersion()));
+        if (version() >= 3) {
+            return new OffsetFetchResponse(throttleTimeMs, error, responsePartitions);
+        } else {
+            return new OffsetFetchResponse(error, responsePartitions);
         }
     }
 
@@ -207,47 +146,16 @@ public OffsetFetchResponse getErrorResponse(int throttleTimeMs, Throwable e) {
         return getErrorResponse(throttleTimeMs, Errors.forException(e));
     }
 
-    public String groupId() {
-        return groupId;
-    }
-
-    public List partitions() {
-        return partitions;
-    }
-
     public static OffsetFetchRequest parse(ByteBuffer buffer, short version) {
         return new OffsetFetchRequest(ApiKeys.OFFSET_FETCH.parseRequest(version, buffer), version);
     }
 
     public boolean isAllPartitions() {
-        return partitions == null;
+        return data.topics() == ALL_TOPIC_PARTITIONS;
     }
 
     @Override
     protected Struct toStruct() {
-        Struct struct = new Struct(ApiKeys.OFFSET_FETCH.requestSchema(version()));
-        struct.set(GROUP_ID, groupId);
-        if (partitions != null) {
-            Map> topicsData = CollectionUtils.groupPartitionsByTopic(partitions);
-
-            List topicArray = new ArrayList<>();
-            for (Map.Entry> entries : topicsData.entrySet()) {
-                Struct topicData = struct.instance(TOPICS);
-                topicData.set(TOPIC_NAME, entries.getKey());
-                List partitionArray = new ArrayList<>();
-                for (Integer partitionId : entries.getValue()) {
-                    Struct partitionData = topicData.instance(PARTITIONS);
-                    partitionData.set(PARTITION_ID, partitionId);
-                    partitionArray.add(partitionData);
-                }
-                topicData.set(PARTITIONS, partitionArray.toArray());
-                topicArray.add(topicData);
-            }
-            struct.set(TOPICS, topicArray.toArray());
-        } else
-            struct.set(TOPICS, null);
-
-        return struct;
+        return data.toStruct(version());
     }
-
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/OffsetFetchResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/OffsetFetchResponse.java
index 2022bb7791089..722ef4e468a7a 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/OffsetFetchResponse.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/OffsetFetchResponse.java
@@ -17,114 +17,53 @@
 package org.apache.kafka.common.requests;
 
 import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.message.OffsetFetchResponseData;
+import org.apache.kafka.common.message.OffsetFetchResponseData.OffsetFetchResponsePartition;
+import org.apache.kafka.common.message.OffsetFetchResponseData.OffsetFetchResponseTopic;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.protocol.types.Field;
-import org.apache.kafka.common.protocol.types.Schema;
 import org.apache.kafka.common.protocol.types.Struct;
-import org.apache.kafka.common.utils.CollectionUtils;
 
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
-import java.util.Collection;
-import java.util.Collections;
+import java.util.Arrays;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Optional;
 
-import static org.apache.kafka.common.protocol.CommonFields.COMMITTED_LEADER_EPOCH;
-import static org.apache.kafka.common.protocol.CommonFields.COMMITTED_METADATA;
-import static org.apache.kafka.common.protocol.CommonFields.COMMITTED_OFFSET;
-import static org.apache.kafka.common.protocol.CommonFields.ERROR_CODE;
-import static org.apache.kafka.common.protocol.CommonFields.PARTITION_ID;
-import static org.apache.kafka.common.protocol.CommonFields.THROTTLE_TIME_MS;
-import static org.apache.kafka.common.protocol.CommonFields.TOPIC_NAME;
+import static org.apache.kafka.common.record.RecordBatch.NO_PARTITION_LEADER_EPOCH;
 
 /**
  * Possible error codes:
  *
  * - Partition errors:
- *   - UNKNOWN_TOPIC_OR_PARTITION (3)
+ *   - {@link Errors#UNKNOWN_TOPIC_OR_PARTITION}
+ *   - {@link Errors#TOPIC_AUTHORIZATION_FAILED}
  *
  * - Group or coordinator errors:
- *   - COORDINATOR_LOAD_IN_PROGRESS (14)
- *   - COORDINATOR_NOT_AVAILABLE (15)
- *   - NOT_COORDINATOR (16)
- *   - GROUP_AUTHORIZATION_FAILED (30)
+ *   - {@link Errors#COORDINATOR_LOAD_IN_PROGRESS}
+ *   - {@link Errors#COORDINATOR_NOT_AVAILABLE}
+ *   - {@link Errors#NOT_COORDINATOR}
+ *   - {@link Errors#GROUP_AUTHORIZATION_FAILED}
  */
 public class OffsetFetchResponse extends AbstractResponse {
-    private static final Field.ComplexArray TOPICS = new Field.ComplexArray("responses",
-            "Responses by topic for fetched offsets");
-
-    // topic level fields
-    private static final Field.ComplexArray PARTITIONS = new Field.ComplexArray("partition_responses",
-            "Responses by partition for fetched offsets");
-
-    private static final Field PARTITIONS_V0 = PARTITIONS.withFields(
-            PARTITION_ID,
-            COMMITTED_OFFSET,
-            COMMITTED_METADATA,
-            ERROR_CODE);
-
-    private static final Field TOPICS_V0 = TOPICS.withFields(
-            TOPIC_NAME,
-            PARTITIONS_V0);
-
-    private static final Schema OFFSET_FETCH_RESPONSE_V0 = new Schema(
-            TOPICS_V0);
-
-    // V1 begins support for fetching offsets from the internal __consumer_offsets topic
-    private static final Schema OFFSET_FETCH_RESPONSE_V1 = OFFSET_FETCH_RESPONSE_V0;
-
-    // V2 adds top-level error code
-    private static final Schema OFFSET_FETCH_RESPONSE_V2 = new Schema(
-            TOPICS_V0,
-            ERROR_CODE);
-
-    // V3 request includes throttle time
-    private static final Schema OFFSET_FETCH_RESPONSE_V3 = new Schema(
-            THROTTLE_TIME_MS,
-            TOPICS_V0,
-            ERROR_CODE);
-
-    // V4 bump used to indicate that on quota violation brokers send out responses before throttling.
-    private static final Schema OFFSET_FETCH_RESPONSE_V4 = OFFSET_FETCH_RESPONSE_V3;
-
-    // V5 adds the leader epoch to the committed offset
-    private static final Field PARTITIONS_V5 = PARTITIONS.withFields(
-            PARTITION_ID,
-            COMMITTED_OFFSET,
-            COMMITTED_LEADER_EPOCH,
-            COMMITTED_METADATA,
-            ERROR_CODE);
-
-    private static final Field TOPICS_V5 = TOPICS.withFields(
-            TOPIC_NAME,
-            PARTITIONS_V5);
-
-    private static final Schema OFFSET_FETCH_RESPONSE_V5 = new Schema(
-            THROTTLE_TIME_MS,
-            TOPICS_V5,
-            ERROR_CODE);
-
-    public static Schema[] schemaVersions() {
-        return new Schema[] {OFFSET_FETCH_RESPONSE_V0, OFFSET_FETCH_RESPONSE_V1, OFFSET_FETCH_RESPONSE_V2,
-            OFFSET_FETCH_RESPONSE_V3, OFFSET_FETCH_RESPONSE_V4, OFFSET_FETCH_RESPONSE_V5};
-    }
-
     public static final long INVALID_OFFSET = -1L;
     public static final String NO_METADATA = "";
     public static final PartitionData UNKNOWN_PARTITION = new PartitionData(INVALID_OFFSET,
-            Optional.empty(), NO_METADATA, Errors.UNKNOWN_TOPIC_OR_PARTITION);
+                                                                            Optional.empty(),
+                                                                            NO_METADATA,
+                                                                            Errors.UNKNOWN_TOPIC_OR_PARTITION);
     public static final PartitionData UNAUTHORIZED_PARTITION = new PartitionData(INVALID_OFFSET,
-            Optional.empty(), NO_METADATA, Errors.TOPIC_AUTHORIZATION_FAILED);
+                                                                                 Optional.empty(),
+                                                                                 NO_METADATA,
+                                                                                 Errors.TOPIC_AUTHORIZATION_FAILED);
+    private static final List PARTITION_ERRORS = Arrays.asList(
+        Errors.UNKNOWN_TOPIC_OR_PARTITION, Errors.TOPIC_AUTHORIZATION_FAILED);
 
-    private static final List PARTITION_ERRORS = Collections.singletonList(Errors.UNKNOWN_TOPIC_OR_PARTITION);
-
-    private final Map responseData;
+    public final OffsetFetchResponseData data;
     private final Errors error;
-    private final int throttleTimeMs;
 
     public static final class PartitionData {
         public final long offset;
@@ -145,6 +84,32 @@ public PartitionData(long offset,
         public boolean hasError() {
             return this.error != Errors.NONE;
         }
+
+        @Override
+        public boolean equals(Object other) {
+            if (!(other instanceof PartitionData))
+                return false;
+            PartitionData otherPartition = (PartitionData) other;
+            return this.offset == otherPartition.offset
+                       && this.leaderEpoch.equals(otherPartition.leaderEpoch)
+                       && this.metadata.equals(otherPartition.metadata)
+                       && this.error.equals(otherPartition.error);
+        }
+
+        @Override
+        public String toString() {
+            return "PartitionData("
+                       + "offset=" + offset
+                       + ", leaderEpoch=" + leaderEpoch.orElse(NO_PARTITION_LEADER_EPOCH)
+                       + ", metadata=" + metadata
+                       + ", error='" + error.toString()
+                       + ")";
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(offset, leaderEpoch, metadata, error);
+        }
     }
 
     /**
@@ -163,31 +128,41 @@ public OffsetFetchResponse(Errors error, Map resp
      * @param responseData Fetched offset information grouped by topic-partition
      */
     public OffsetFetchResponse(int throttleTimeMs, Errors error, Map responseData) {
-        this.throttleTimeMs = throttleTimeMs;
-        this.responseData = responseData;
+        Map offsetFetchResponseTopicMap = new HashMap<>();
+        for (Map.Entry entry : responseData.entrySet()) {
+            String topicName = entry.getKey().topic();
+            OffsetFetchResponseTopic topic = offsetFetchResponseTopicMap.getOrDefault(
+                topicName, new OffsetFetchResponseTopic().setName(topicName));
+            PartitionData partitionData = entry.getValue();
+            topic.partitions().add(new OffsetFetchResponsePartition()
+                                       .setPartitionIndex(entry.getKey().partition())
+                                       .setErrorCode(partitionData.error.code())
+                                       .setCommittedOffset(partitionData.offset)
+                                       .setCommittedLeaderEpoch(
+                                           partitionData.leaderEpoch.orElse(NO_PARTITION_LEADER_EPOCH))
+                                       .setMetadata(partitionData.metadata)
+            );
+            offsetFetchResponseTopicMap.put(topicName, topic);
+        }
+
+        this.data = new OffsetFetchResponseData()
+            .setTopics(new ArrayList<>(offsetFetchResponseTopicMap.values()))
+            .setErrorCode(error.code())
+            .setThrottleTimeMs(throttleTimeMs);
         this.error = error;
     }
 
-    public OffsetFetchResponse(Struct struct) {
-        this.throttleTimeMs = struct.getOrElse(THROTTLE_TIME_MS, DEFAULT_THROTTLE_TIME);
-        Errors topLevelError = Errors.NONE;
-        this.responseData = new HashMap<>();
-        for (Object topicResponseObj : struct.get(TOPICS)) {
-            Struct topicResponse = (Struct) topicResponseObj;
-            String topic = topicResponse.get(TOPIC_NAME);
-            for (Object partitionResponseObj : topicResponse.get(PARTITIONS)) {
-                Struct partitionResponse = (Struct) partitionResponseObj;
-                int partition = partitionResponse.get(PARTITION_ID);
-                long offset = partitionResponse.get(COMMITTED_OFFSET);
-                String metadata = partitionResponse.get(COMMITTED_METADATA);
-                Optional leaderEpochOpt = RequestUtils.getLeaderEpoch(partitionResponse, COMMITTED_LEADER_EPOCH);
-
-                Errors error = Errors.forCode(partitionResponse.get(ERROR_CODE));
-                if (error != Errors.NONE && !PARTITION_ERRORS.contains(error))
-                    topLevelError = error;
+    public OffsetFetchResponse(Struct struct, short version) {
+        this.data = new OffsetFetchResponseData(struct, version);
 
-                PartitionData partitionData = new PartitionData(offset, leaderEpochOpt, metadata, error);
-                this.responseData.put(new TopicPartition(topic, partition), partitionData);
+        Errors topLevelError = Errors.NONE;
+        for (OffsetFetchResponseTopic topic : data.topics()) {
+            for (OffsetFetchResponsePartition partition : topic.partitions()) {
+                Errors partitionError = Errors.forCode(partition.errorCode());
+                if (partitionError != Errors.NONE && !PARTITION_ERRORS.contains(partitionError)) {
+                    topLevelError = partitionError;
+                    break;
+                }
             }
         }
 
@@ -195,28 +170,20 @@ public OffsetFetchResponse(Struct struct) {
         // for older versions there is no top-level error in the response and all errors are partition errors,
         // so if there is a group or coordinator error at the partition level use that as the top-level error.
         // this way clients can depend on the top-level error regardless of the offset fetch version.
-        this.error = struct.hasField(ERROR_CODE) ? Errors.forCode(struct.get(ERROR_CODE)) : topLevelError;
-    }
-
-    public void maybeThrowFirstPartitionError() {
-        Collection partitionsData = this.responseData.values();
-        for (PartitionData data : partitionsData) {
-            if (data.hasError())
-                throw data.error.exception();
-        }
+        this.error = version >= 2 ? Errors.forCode(data.errorCode()) : topLevelError;
     }
 
     @Override
     public int throttleTimeMs() {
-        return throttleTimeMs;
+        return data.throttleTimeMs();
     }
 
     public boolean hasError() {
-        return this.error != Errors.NONE;
+        return error != Errors.NONE;
     }
 
     public Errors error() {
-        return this.error;
+        return error;
     }
 
     @Override
@@ -225,43 +192,27 @@ public Map errorCounts() {
     }
 
     public Map responseData() {
+        Map responseData = new HashMap<>();
+        for (OffsetFetchResponseTopic topic : data.topics()) {
+            for (OffsetFetchResponsePartition partition : topic.partitions()) {
+                responseData.put(new TopicPartition(topic.name(), partition.partitionIndex()),
+                                 new PartitionData(partition.committedOffset(),
+                                                   RequestUtils.getLeaderEpoch(partition.committedLeaderEpoch()),
+                                                   partition.metadata(),
+                                                   Errors.forCode(partition.errorCode()))
+                );
+            }
+        }
         return responseData;
     }
 
     public static OffsetFetchResponse parse(ByteBuffer buffer, short version) {
-        return new OffsetFetchResponse(ApiKeys.OFFSET_FETCH.parseResponse(version, buffer));
+        return new OffsetFetchResponse(ApiKeys.OFFSET_FETCH.parseResponse(version, buffer), version);
     }
 
     @Override
     protected Struct toStruct(short version) {
-        Struct struct = new Struct(ApiKeys.OFFSET_FETCH.responseSchema(version));
-        struct.setIfExists(THROTTLE_TIME_MS, throttleTimeMs);
-
-        Map> topicsData = CollectionUtils.groupPartitionDataByTopic(responseData);
-        List topicArray = new ArrayList<>();
-        for (Map.Entry> entries : topicsData.entrySet()) {
-            Struct topicData = struct.instance(TOPICS);
-            topicData.set(TOPIC_NAME, entries.getKey());
-            List partitionArray = new ArrayList<>();
-            for (Map.Entry partitionEntry : entries.getValue().entrySet()) {
-                PartitionData fetchPartitionData = partitionEntry.getValue();
-                Struct partitionData = topicData.instance(PARTITIONS);
-                partitionData.set(PARTITION_ID, partitionEntry.getKey());
-                partitionData.set(COMMITTED_OFFSET, fetchPartitionData.offset);
-                RequestUtils.setLeaderEpochIfExists(partitionData, COMMITTED_LEADER_EPOCH, fetchPartitionData.leaderEpoch);
-                partitionData.set(COMMITTED_METADATA, fetchPartitionData.metadata);
-                partitionData.set(ERROR_CODE, fetchPartitionData.error.code());
-                partitionArray.add(partitionData);
-            }
-            topicData.set(PARTITIONS, partitionArray.toArray());
-            topicArray.add(topicData);
-        }
-        struct.set(TOPICS, topicArray.toArray());
-
-        if (version > 1)
-            struct.set(ERROR_CODE, this.error.code());
-
-        return struct;
+        return data.toStruct(version);
     }
 
     @Override
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/OffsetsForLeaderEpochRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/OffsetsForLeaderEpochRequest.java
index 75788a06ba0df..d5c78da5e281f 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/OffsetsForLeaderEpochRequest.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/OffsetsForLeaderEpochRequest.java
@@ -17,6 +17,7 @@
 package org.apache.kafka.common.requests;
 
 import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.errors.UnsupportedVersionException;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
 import org.apache.kafka.common.protocol.types.Field;
@@ -36,6 +37,20 @@
 import static org.apache.kafka.common.protocol.CommonFields.TOPIC_NAME;
 
 public class OffsetsForLeaderEpochRequest extends AbstractRequest {
+    private static final Field.Int32 REPLICA_ID = new Field.Int32("replica_id",
+            "Broker id of the follower. For normal consumers, use -1.");
+
+    /**
+     * Sentinel replica_id value to indicate a regular consumer rather than another broker
+     */
+    public static final int CONSUMER_REPLICA_ID = -1;
+
+    /**
+     * Sentinel replica_id which indicates either a debug consumer or a replica which is using
+     * an old version of the protocol.
+     */
+    public static final int DEBUGGING_REPLICA_ID = -2;
+
     private static final Field.ComplexArray TOPICS = new Field.ComplexArray("topics",
             "An array of topics to get epochs for");
     private static final Field.ComplexArray PARTITIONS = new Field.ComplexArray("partitions",
@@ -67,28 +82,55 @@ public class OffsetsForLeaderEpochRequest extends AbstractRequest {
     private static final Schema OFFSET_FOR_LEADER_EPOCH_REQUEST_V2 = new Schema(
             TOPICS_V2);
 
+    private static final Schema OFFSET_FOR_LEADER_EPOCH_REQUEST_V3 = new Schema(
+            REPLICA_ID,
+            TOPICS_V2);
+
     public static Schema[] schemaVersions() {
         return new Schema[]{OFFSET_FOR_LEADER_EPOCH_REQUEST_V0, OFFSET_FOR_LEADER_EPOCH_REQUEST_V1,
-            OFFSET_FOR_LEADER_EPOCH_REQUEST_V2};
+            OFFSET_FOR_LEADER_EPOCH_REQUEST_V2, OFFSET_FOR_LEADER_EPOCH_REQUEST_V3};
     }
 
-    private Map epochsByPartition;
+    private final Map epochsByPartition;
+
+    private final int replicaId;
 
     public Map epochsByTopicPartition() {
         return epochsByPartition;
     }
 
+    public int replicaId() {
+        return replicaId;
+    }
+
     public static class Builder extends AbstractRequest.Builder {
         private final Map epochsByPartition;
+        private final int replicaId;
 
-        public Builder(short version, Map epochsByPartition) {
-            super(ApiKeys.OFFSET_FOR_LEADER_EPOCH, version);
+        Builder(short oldestAllowedVersion, short latestAllowedVersion, Map epochsByPartition, int replicaId) {
+            super(ApiKeys.OFFSET_FOR_LEADER_EPOCH, oldestAllowedVersion, latestAllowedVersion);
             this.epochsByPartition = epochsByPartition;
+            this.replicaId = replicaId;
+        }
+
+        public static Builder forConsumer(Map epochsByPartition) {
+            // Old versions of this API require CLUSTER permission which is not typically granted
+            // to clients. Beginning with version 3, the broker requires only TOPIC Describe
+            // permission for the topic of each requested partition. In order to ensure client
+            // compatibility, we only send this request when we can guarantee the relaxed permissions.
+            return new Builder((short) 3, ApiKeys.OFFSET_FOR_LEADER_EPOCH.latestVersion(),
+                    epochsByPartition, CONSUMER_REPLICA_ID);
+        }
+
+        public static Builder forFollower(short version, Map epochsByPartition, int replicaId) {
+            return new Builder(version, version, epochsByPartition, replicaId);
         }
 
         @Override
         public OffsetsForLeaderEpochRequest build(short version) {
-            return new OffsetsForLeaderEpochRequest(epochsByPartition, version);
+            if (version < oldestAllowedVersion() || version > latestAllowedVersion())
+                throw new UnsupportedVersionException("Cannot build " + this + " with version " + version);
+            return new OffsetsForLeaderEpochRequest(epochsByPartition, replicaId, version);
         }
 
         public static OffsetsForLeaderEpochRequest parse(ByteBuffer buffer, short version) {
@@ -105,13 +147,15 @@ public String toString() {
         }
     }
 
-    public OffsetsForLeaderEpochRequest(Map epochsByPartition, short version) {
+    public OffsetsForLeaderEpochRequest(Map epochsByPartition, int replicaId, short version) {
         super(ApiKeys.OFFSET_FOR_LEADER_EPOCH, version);
         this.epochsByPartition = epochsByPartition;
+        this.replicaId = replicaId;
     }
 
     public OffsetsForLeaderEpochRequest(Struct struct, short version) {
         super(ApiKeys.OFFSET_FOR_LEADER_EPOCH, version);
+        replicaId = struct.getOrElse(REPLICA_ID, DEBUGGING_REPLICA_ID);
         epochsByPartition = new HashMap<>();
         for (Object topicAndEpochsObj : struct.get(TOPICS)) {
             Struct topicAndEpochs = (Struct) topicAndEpochsObj;
@@ -134,6 +178,7 @@ public static OffsetsForLeaderEpochRequest parse(ByteBuffer buffer, short versio
     @Override
     protected Struct toStruct() {
         Struct requestStruct = new Struct(ApiKeys.OFFSET_FOR_LEADER_EPOCH.requestSchema(version()));
+        requestStruct.setIfExists(REPLICA_ID, replicaId);
 
         Map> topicsToPartitionEpochs = CollectionUtils.groupPartitionDataByTopic(epochsByPartition);
 
@@ -189,4 +234,13 @@ public String toString() {
             return bld.toString();
         }
     }
+
+    /**
+     * Check whether a broker allows Topic-level permissions in order to use the
+     * OffsetForLeaderEpoch API. Old versions require Cluster permission.
+     */
+    public static boolean supportsTopicPermission(short latestUsableVersion) {
+        return latestUsableVersion >= 3;
+    }
+
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/OffsetsForLeaderEpochResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/OffsetsForLeaderEpochResponse.java
index d5d12653c3da0..3fe3cdc0da6b5 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/OffsetsForLeaderEpochResponse.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/OffsetsForLeaderEpochResponse.java
@@ -85,9 +85,12 @@ public class OffsetsForLeaderEpochResponse extends AbstractResponse {
             THROTTLE_TIME_MS,
             TOPICS_V1);
 
+    private static final Schema OFFSET_FOR_LEADER_EPOCH_RESPONSE_V3 = OFFSET_FOR_LEADER_EPOCH_RESPONSE_V2;
+
+
     public static Schema[] schemaVersions() {
         return new Schema[]{OFFSET_FOR_LEADER_EPOCH_RESPONSE_V0, OFFSET_FOR_LEADER_EPOCH_RESPONSE_V1,
-            OFFSET_FOR_LEADER_EPOCH_RESPONSE_V2};
+            OFFSET_FOR_LEADER_EPOCH_RESPONSE_V2, OFFSET_FOR_LEADER_EPOCH_RESPONSE_V3};
     }
 
     private final int throttleTimeMs;
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/ProduceRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/ProduceRequest.java
index ad23f3f61d83f..7b3ae1cf2afe0 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/ProduceRequest.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/ProduceRequest.java
@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.common.requests;
 
+import org.apache.kafka.common.InvalidRecordException;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.errors.UnsupportedCompressionTypeException;
 import org.apache.kafka.common.protocol.ApiKeys;
@@ -26,7 +27,6 @@
 import org.apache.kafka.common.protocol.types.Schema;
 import org.apache.kafka.common.protocol.types.Struct;
 import org.apache.kafka.common.record.CompressionType;
-import org.apache.kafka.common.record.InvalidRecordException;
 import org.apache.kafka.common.record.MemoryRecords;
 import org.apache.kafka.common.record.MutableRecordBatch;
 import org.apache.kafka.common.record.RecordBatch;
@@ -120,9 +120,15 @@ public class ProduceRequest extends AbstractRequest {
      */
     private static final Schema PRODUCE_REQUEST_V7 = PRODUCE_REQUEST_V6;
 
+    /**
+     * V8 bumped up to add two new fields record_errors offset list and error_message to {@link org.apache.kafka.common.requests.ProduceResponse.PartitionResponse}
+     * (See KIP-467)
+     */
+    private static final Schema PRODUCE_REQUEST_V8 = PRODUCE_REQUEST_V7;
+
     public static Schema[] schemaVersions() {
         return new Schema[] {PRODUCE_REQUEST_V0, PRODUCE_REQUEST_V1, PRODUCE_REQUEST_V2, PRODUCE_REQUEST_V3,
-            PRODUCE_REQUEST_V4, PRODUCE_REQUEST_V5, PRODUCE_REQUEST_V6, PRODUCE_REQUEST_V7};
+            PRODUCE_REQUEST_V4, PRODUCE_REQUEST_V5, PRODUCE_REQUEST_V6, PRODUCE_REQUEST_V7, PRODUCE_REQUEST_V8};
     }
 
     public static class Builder extends AbstractRequest.Builder {
@@ -337,6 +343,7 @@ public ProduceResponse getErrorResponse(int throttleTimeMs, Throwable e) {
             case 5:
             case 6:
             case 7:
+            case 8:
                 return new ProduceResponse(responseMap, throttleTimeMs);
             default:
                 throw new IllegalArgumentException(String.format("Version %d is not valid. Valid versions for %s are 0 to %d",
@@ -434,6 +441,7 @@ public static byte requiredMagicForVersion(short produceRequestVersion) {
             case 5:
             case 6:
             case 7:
+            case 8:
                 return RecordBatch.MAGIC_VALUE_V2;
 
             default:
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/ProduceResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/ProduceResponse.java
index 7d3e4fed36280..7bbab0879bf42 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/ProduceResponse.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/ProduceResponse.java
@@ -28,6 +28,7 @@
 
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -53,28 +54,37 @@ public class ProduceResponse extends AbstractResponse {
     /**
      * Possible error code:
      *
-     * CORRUPT_MESSAGE (2)
-     * UNKNOWN_TOPIC_OR_PARTITION (3)
-     * NOT_LEADER_FOR_PARTITION (6)
-     * MESSAGE_TOO_LARGE (10)
-     * INVALID_TOPIC (17)
-     * RECORD_LIST_TOO_LARGE (18)
-     * NOT_ENOUGH_REPLICAS (19)
-     * NOT_ENOUGH_REPLICAS_AFTER_APPEND (20)
-     * INVALID_REQUIRED_ACKS (21)
-     * TOPIC_AUTHORIZATION_FAILED (29)
-     * UNSUPPORTED_FOR_MESSAGE_FORMAT (43)
-     * INVALID_PRODUCER_EPOCH (47)
-     * CLUSTER_AUTHORIZATION_FAILED (31)
-     * TRANSACTIONAL_ID_AUTHORIZATION_FAILED (53)
+     * {@link Errors#CORRUPT_MESSAGE}
+     * {@link Errors#UNKNOWN_TOPIC_OR_PARTITION}
+     * {@link Errors#NOT_LEADER_FOR_PARTITION}
+     * {@link Errors#MESSAGE_TOO_LARGE}
+     * {@link Errors#INVALID_TOPIC_EXCEPTION}
+     * {@link Errors#RECORD_LIST_TOO_LARGE}
+     * {@link Errors#NOT_ENOUGH_REPLICAS}
+     * {@link Errors#NOT_ENOUGH_REPLICAS_AFTER_APPEND}
+     * {@link Errors#INVALID_REQUIRED_ACKS}
+     * {@link Errors#TOPIC_AUTHORIZATION_FAILED}
+     * {@link Errors#UNSUPPORTED_FOR_MESSAGE_FORMAT}
+     * {@link Errors#INVALID_PRODUCER_EPOCH}
+     * {@link Errors#CLUSTER_AUTHORIZATION_FAILED}
+     * {@link Errors#TRANSACTIONAL_ID_AUTHORIZATION_FAILED}
+     * {@link Errors#INVALID_RECORD}
      */
 
     private static final String BASE_OFFSET_KEY_NAME = "base_offset";
     private static final String LOG_APPEND_TIME_KEY_NAME = "log_append_time";
     private static final String LOG_START_OFFSET_KEY_NAME = "log_start_offset";
+    private static final String RECORD_ERRORS_KEY_NAME = "record_errors";
+    private static final String BATCH_INDEX_KEY_NAME = "batch_index";
+    private static final String BATCH_INDEX_ERROR_MESSAGE_KEY_NAME = "batch_index_error_message";
+    private static final String ERROR_MESSAGE_KEY_NAME = "error_message";
 
     private static final Field.Int64 LOG_START_OFFSET_FIELD = new Field.Int64(LOG_START_OFFSET_KEY_NAME,
             "The start offset of the log at the time this produce response was created", INVALID_OFFSET);
+    private static final Field.NullableStr BATCH_INDEX_ERROR_MESSAGE_FIELD = new Field.NullableStr(BATCH_INDEX_ERROR_MESSAGE_KEY_NAME,
+            "The error message of the record that caused the batch to be dropped");
+    private static final Field.NullableStr ERROR_MESSAGE_FIELD = new Field.NullableStr(ERROR_MESSAGE_KEY_NAME,
+            "The global error message summarizing the common root cause of the records that caused the batch to be dropped");
 
     private static final Schema PRODUCE_RESPONSE_V0 = new Schema(
             new Field(RESPONSES_KEY_NAME, new ArrayOf(new Schema(
@@ -149,9 +159,32 @@ public class ProduceResponse extends AbstractResponse {
      */
     private static final Schema PRODUCE_RESPONSE_V7 = PRODUCE_RESPONSE_V6;
 
+    /**
+     * V8 adds record_errors and error_message. (see KIP-467)
+     */
+    public static final Schema PRODUCE_RESPONSE_V8 = new Schema(
+            new Field(RESPONSES_KEY_NAME, new ArrayOf(new Schema(
+                    TOPIC_NAME,
+                    new Field(PARTITION_RESPONSES_KEY_NAME, new ArrayOf(new Schema(
+                            PARTITION_ID,
+                            ERROR_CODE,
+                            new Field(BASE_OFFSET_KEY_NAME, INT64),
+                            new Field(LOG_APPEND_TIME_KEY_NAME, INT64, "The timestamp returned by broker after appending " +
+                                    "the messages. If CreateTime is used for the topic, the timestamp will be -1. " +
+                                    "If LogAppendTime is used for the topic, the timestamp will be the broker local " +
+                                    "time when the messages are appended."),
+                            LOG_START_OFFSET_FIELD,
+                            new Field(RECORD_ERRORS_KEY_NAME, new ArrayOf(new Schema(
+                                    new Field.Int32(BATCH_INDEX_KEY_NAME, "The batch index of the record " +
+                                            "that caused the batch to be dropped"),
+                                    BATCH_INDEX_ERROR_MESSAGE_FIELD
+                            )), "The batch indices of records that caused the batch to be dropped"),
+                            ERROR_MESSAGE_FIELD)))))),
+            THROTTLE_TIME_MS);
+
     public static Schema[] schemaVersions() {
         return new Schema[]{PRODUCE_RESPONSE_V0, PRODUCE_RESPONSE_V1, PRODUCE_RESPONSE_V2, PRODUCE_RESPONSE_V3,
-            PRODUCE_RESPONSE_V4, PRODUCE_RESPONSE_V5, PRODUCE_RESPONSE_V6, PRODUCE_RESPONSE_V7};
+            PRODUCE_RESPONSE_V4, PRODUCE_RESPONSE_V5, PRODUCE_RESPONSE_V6, PRODUCE_RESPONSE_V7, PRODUCE_RESPONSE_V8};
     }
 
     private final Map responses;
@@ -183,15 +216,34 @@ public ProduceResponse(Struct struct) {
         for (Object topicResponse : struct.getArray(RESPONSES_KEY_NAME)) {
             Struct topicRespStruct = (Struct) topicResponse;
             String topic = topicRespStruct.get(TOPIC_NAME);
+
             for (Object partResponse : topicRespStruct.getArray(PARTITION_RESPONSES_KEY_NAME)) {
                 Struct partRespStruct = (Struct) partResponse;
                 int partition = partRespStruct.get(PARTITION_ID);
                 Errors error = Errors.forCode(partRespStruct.get(ERROR_CODE));
                 long offset = partRespStruct.getLong(BASE_OFFSET_KEY_NAME);
-                long logAppendTime = partRespStruct.getLong(LOG_APPEND_TIME_KEY_NAME);
+                long logAppendTime = partRespStruct.hasField(LOG_APPEND_TIME_KEY_NAME) ?
+                        partRespStruct.getLong(LOG_APPEND_TIME_KEY_NAME) : RecordBatch.NO_TIMESTAMP;
                 long logStartOffset = partRespStruct.getOrElse(LOG_START_OFFSET_FIELD, INVALID_OFFSET);
+
+                List recordErrors = Collections.emptyList();
+                if (partRespStruct.hasField(RECORD_ERRORS_KEY_NAME)) {
+                    Object[] recordErrorsArray = partRespStruct.getArray(RECORD_ERRORS_KEY_NAME);
+                    if (recordErrorsArray.length > 0) {
+                        recordErrors = new ArrayList<>(recordErrorsArray.length);
+                        for (Object indexAndMessage : recordErrorsArray) {
+                            Struct indexAndMessageStruct = (Struct) indexAndMessage;
+                            recordErrors.add(new RecordError(
+                                    indexAndMessageStruct.getInt(BATCH_INDEX_KEY_NAME),
+                                    indexAndMessageStruct.get(BATCH_INDEX_ERROR_MESSAGE_FIELD)
+                            ));
+                        }
+                    }
+                }
+
+                String errorMessage = partRespStruct.getOrElse(ERROR_MESSAGE_FIELD, null);
                 TopicPartition tp = new TopicPartition(topic, partition);
-                responses.put(tp, new PartitionResponse(error, offset, logAppendTime, logStartOffset));
+                responses.put(tp, new PartitionResponse(error, offset, logAppendTime, logStartOffset, recordErrors, errorMessage));
             }
         }
         this.throttleTimeMs = struct.getOrElse(THROTTLE_TIME_MS, DEFAULT_THROTTLE_TIME);
@@ -220,9 +272,24 @@ protected Struct toStruct(short version) {
                         .set(PARTITION_ID, partitionEntry.getKey())
                         .set(ERROR_CODE, errorCode)
                         .set(BASE_OFFSET_KEY_NAME, part.baseOffset);
-                if (partStruct.hasField(LOG_APPEND_TIME_KEY_NAME))
-                    partStruct.set(LOG_APPEND_TIME_KEY_NAME, part.logAppendTime);
+                partStruct.setIfExists(LOG_APPEND_TIME_KEY_NAME, part.logAppendTime);
                 partStruct.setIfExists(LOG_START_OFFSET_FIELD, part.logStartOffset);
+
+                if (partStruct.hasField(RECORD_ERRORS_KEY_NAME)) {
+                    List recordErrors = Collections.emptyList();
+                    if (!part.recordErrors.isEmpty()) {
+                        recordErrors = new ArrayList<>();
+                        for (RecordError indexAndMessage : part.recordErrors) {
+                            Struct indexAndMessageStruct = partStruct.instance(RECORD_ERRORS_KEY_NAME)
+                                    .set(BATCH_INDEX_KEY_NAME, indexAndMessage.batchIndex)
+                                    .set(BATCH_INDEX_ERROR_MESSAGE_FIELD, indexAndMessage.message);
+                            recordErrors.add(indexAndMessageStruct);
+                        }
+                    }
+                    partStruct.set(RECORD_ERRORS_KEY_NAME, recordErrors.toArray());
+                }
+
+                partStruct.setIfExists(ERROR_MESSAGE_FIELD, part.errorMessage);
                 partitionArray.add(partStruct);
             }
             topicData.set(PARTITION_RESPONSES_KEY_NAME, partitionArray.toArray());
@@ -256,16 +323,28 @@ public static final class PartitionResponse {
         public long baseOffset;
         public long logAppendTime;
         public long logStartOffset;
+        public List recordErrors;
+        public String errorMessage;
 
         public PartitionResponse(Errors error) {
             this(error, INVALID_OFFSET, RecordBatch.NO_TIMESTAMP, INVALID_OFFSET);
         }
 
         public PartitionResponse(Errors error, long baseOffset, long logAppendTime, long logStartOffset) {
+            this(error, baseOffset, logAppendTime, logStartOffset, Collections.emptyList(), null);
+        }
+
+        public PartitionResponse(Errors error, long baseOffset, long logAppendTime, long logStartOffset, List recordErrors) {
+            this(error, baseOffset, logAppendTime, logStartOffset, recordErrors, null);
+        }
+
+        public PartitionResponse(Errors error, long baseOffset, long logAppendTime, long logStartOffset, List recordErrors, String errorMessage) {
             this.error = error;
             this.baseOffset = baseOffset;
             this.logAppendTime = logAppendTime;
             this.logStartOffset = logStartOffset;
+            this.recordErrors = recordErrors;
+            this.errorMessage = errorMessage;
         }
 
         @Override
@@ -280,11 +359,35 @@ public String toString() {
             b.append(logAppendTime);
             b.append(", logStartOffset: ");
             b.append(logStartOffset);
+            b.append(", recordErrors: ");
+            b.append(recordErrors);
+            b.append(", errorMessage: ");
+            if (errorMessage != null) {
+                b.append(errorMessage);
+            } else {
+                b.append("null");
+            }
             b.append('}');
             return b.toString();
         }
     }
 
+    public static final class RecordError {
+        public final int batchIndex;
+        public final String message;
+
+        public RecordError(int batchIndex, String message) {
+            this.batchIndex = batchIndex;
+            this.message = message;
+        }
+
+        public RecordError(int batchIndex) {
+            this.batchIndex = batchIndex;
+            this.message = null;
+        }
+
+    }
+
     public static ProduceResponse parse(ByteBuffer buffer, short version) {
         return new ProduceResponse(ApiKeys.PRODUCE.responseSchema(version).read(buffer));
     }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/RenewDelegationTokenRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/RenewDelegationTokenRequest.java
index d73561abbb88e..2c83f06adda62 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/RenewDelegationTokenRequest.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/RenewDelegationTokenRequest.java
@@ -16,102 +16,65 @@
  */
 package org.apache.kafka.common.requests;
 
+import java.nio.ByteBuffer;
+
+import org.apache.kafka.common.message.RenewDelegationTokenRequestData;
+import org.apache.kafka.common.message.RenewDelegationTokenResponseData;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.protocol.types.Field;
-import org.apache.kafka.common.protocol.types.Schema;
 import org.apache.kafka.common.protocol.types.Struct;
 
-import java.nio.ByteBuffer;
-
-import static org.apache.kafka.common.protocol.types.Type.BYTES;
-import static org.apache.kafka.common.protocol.types.Type.INT64;
-
 public class RenewDelegationTokenRequest extends AbstractRequest {
 
-    private static final String HMAC_KEY_NAME = "hmac";
-    private static final String RENEW_TIME_PERIOD_KEY_NAME = "renew_time_period";
-    private final ByteBuffer hmac;
-    private final long renewTimePeriod;
+    private final RenewDelegationTokenRequestData data;
 
-    public static final Schema TOKEN_RENEW_REQUEST_V0 = new Schema(
-        new Field(HMAC_KEY_NAME, BYTES, "HMAC of the delegation token to be renewed."),
-        new Field(RENEW_TIME_PERIOD_KEY_NAME, INT64, "Renew time period in milli seconds."));
-
-    /**
-     * The version number is bumped to indicate that on quota violation brokers send out responses before throttling.
-     */
-    public static final Schema TOKEN_RENEW_REQUEST_V1 = TOKEN_RENEW_REQUEST_V0;
-
-    private RenewDelegationTokenRequest(short version, ByteBuffer hmac, long renewTimePeriod) {
+    public RenewDelegationTokenRequest(RenewDelegationTokenRequestData data, short version) {
         super(ApiKeys.RENEW_DELEGATION_TOKEN, version);
-
-        this.hmac = hmac;
-        this.renewTimePeriod = renewTimePeriod;
+        this.data = data;
     }
 
-    public RenewDelegationTokenRequest(Struct struct, short versionId) {
-        super(ApiKeys.RENEW_DELEGATION_TOKEN, versionId);
-
-        hmac = struct.getBytes(HMAC_KEY_NAME);
-        renewTimePeriod = struct.getLong(RENEW_TIME_PERIOD_KEY_NAME);
+    public RenewDelegationTokenRequest(Struct struct, short version) {
+        super(ApiKeys.RENEW_DELEGATION_TOKEN, version);
+        this.data = new RenewDelegationTokenRequestData(struct, version);
     }
 
     public static RenewDelegationTokenRequest parse(ByteBuffer buffer, short version) {
         return new RenewDelegationTokenRequest(ApiKeys.RENEW_DELEGATION_TOKEN.parseRequest(version, buffer), version);
     }
 
-    public static Schema[] schemaVersions() {
-        return new Schema[] {TOKEN_RENEW_REQUEST_V0, TOKEN_RENEW_REQUEST_V1};
-    }
-
     @Override
     protected Struct toStruct() {
-        short version = version();
-        Struct struct = new Struct(ApiKeys.RENEW_DELEGATION_TOKEN.requestSchema(version));
-
-        struct.set(HMAC_KEY_NAME, hmac);
-        struct.set(RENEW_TIME_PERIOD_KEY_NAME, renewTimePeriod);
+        return data.toStruct(version());
+    }
 
-        return struct;
+    public RenewDelegationTokenRequestData data() {
+        return data;
     }
 
     @Override
     public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
-        return new RenewDelegationTokenResponse(throttleTimeMs, Errors.forException(e));
-    }
-
-    public ByteBuffer hmac() {
-        return hmac;
-    }
-
-    public long renewTimePeriod() {
-        return renewTimePeriod;
+        return new RenewDelegationTokenResponse(
+                new RenewDelegationTokenResponseData()
+                    .setThrottleTimeMs(throttleTimeMs)
+                    .setErrorCode(Errors.forException(e).code()));
     }
 
     public static class Builder extends AbstractRequest.Builder {
-        private final ByteBuffer hmac;
-        private final long renewTimePeriod;
+        private final RenewDelegationTokenRequestData data;
 
-        public Builder(byte[] hmac, long renewTimePeriod) {
+        public Builder(RenewDelegationTokenRequestData data) {
             super(ApiKeys.RENEW_DELEGATION_TOKEN);
-            this.hmac = ByteBuffer.wrap(hmac);
-            this.renewTimePeriod = renewTimePeriod;
+            this.data = data;
         }
 
         @Override
         public RenewDelegationTokenRequest build(short version) {
-            return new RenewDelegationTokenRequest(version, hmac, renewTimePeriod);
+            return new RenewDelegationTokenRequest(data, version);
         }
 
         @Override
         public String toString() {
-            StringBuilder bld = new StringBuilder();
-            bld.append("(type: RenewDelegationTokenRequest").
-                append(", hmac=").append(hmac).
-                append(", renewTimePeriod=").append(renewTimePeriod).
-                append(")");
-            return bld.toString();
+            return data.toString();
         }
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/RenewDelegationTokenResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/RenewDelegationTokenResponse.java
index dc961e176df3a..35cd6152632b6 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/RenewDelegationTokenResponse.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/RenewDelegationTokenResponse.java
@@ -16,92 +16,56 @@
  */
 package org.apache.kafka.common.requests;
 
-import org.apache.kafka.common.protocol.ApiKeys;
-import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.protocol.types.Field;
-import org.apache.kafka.common.protocol.types.Schema;
-import org.apache.kafka.common.protocol.types.Struct;
-
 import java.nio.ByteBuffer;
+import java.util.Collections;
 import java.util.Map;
 
-import static org.apache.kafka.common.protocol.CommonFields.ERROR_CODE;
-import static org.apache.kafka.common.protocol.CommonFields.THROTTLE_TIME_MS;
-import static org.apache.kafka.common.protocol.types.Type.INT64;
+import org.apache.kafka.common.message.RenewDelegationTokenResponseData;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.common.protocol.types.Struct;
 
 public class RenewDelegationTokenResponse extends AbstractResponse {
 
-    private static final String EXPIRY_TIMESTAMP_KEY_NAME = "expiry_timestamp";
-
-    private final Errors error;
-    private final long expiryTimestamp;
-    private final int throttleTimeMs;
+    private final RenewDelegationTokenResponseData data;
 
-    private static final Schema TOKEN_RENEW_RESPONSE_V0 = new Schema(
-            ERROR_CODE,
-            new Field(EXPIRY_TIMESTAMP_KEY_NAME, INT64, "timestamp (in msec) at which this token expires.."),
-            THROTTLE_TIME_MS);
-
-    /**
-     * The version number is bumped to indicate that on quota violation brokers send out responses before throttling.
-     */
-    private static final Schema TOKEN_RENEW_RESPONSE_V1 = TOKEN_RENEW_RESPONSE_V0;
-
-    public RenewDelegationTokenResponse(int throttleTimeMs, Errors error, long expiryTimestamp) {
-        this.throttleTimeMs = throttleTimeMs;
-        this.error = error;
-        this.expiryTimestamp = expiryTimestamp;
-    }
-
-    public RenewDelegationTokenResponse(int throttleTimeMs, Errors error) {
-        this(throttleTimeMs, error, -1);
+    public RenewDelegationTokenResponse(RenewDelegationTokenResponseData data) {
+        this.data = data;
     }
 
-    public RenewDelegationTokenResponse(Struct struct) {
-        error = Errors.forCode(struct.get(ERROR_CODE));
-        expiryTimestamp = struct.getLong(EXPIRY_TIMESTAMP_KEY_NAME);
-        this.throttleTimeMs = struct.getOrElse(THROTTLE_TIME_MS, DEFAULT_THROTTLE_TIME);
+    public RenewDelegationTokenResponse(Struct struct, short version) {
+        data = new RenewDelegationTokenResponseData(struct, version);
     }
 
     public static RenewDelegationTokenResponse parse(ByteBuffer buffer, short version) {
-        return new RenewDelegationTokenResponse(ApiKeys.RENEW_DELEGATION_TOKEN.responseSchema(version).read(buffer));
-    }
-
-    public static Schema[] schemaVersions() {
-        return new Schema[] {TOKEN_RENEW_RESPONSE_V0, TOKEN_RENEW_RESPONSE_V1};
+        return new RenewDelegationTokenResponse(ApiKeys.RENEW_DELEGATION_TOKEN.responseSchema(version).read(buffer), version);
     }
 
     @Override
     public Map errorCounts() {
-        return errorCounts(error);
+        return Collections.singletonMap(error(), 1);
     }
 
     @Override
     protected Struct toStruct(short version) {
-        Struct struct = new Struct(ApiKeys.RENEW_DELEGATION_TOKEN.responseSchema(version));
-
-        struct.set(ERROR_CODE, error.code());
-        struct.set(EXPIRY_TIMESTAMP_KEY_NAME, expiryTimestamp);
-        struct.setIfExists(THROTTLE_TIME_MS, throttleTimeMs);
-
-        return struct;
+        return data.toStruct(version);
     }
 
     @Override
     public int throttleTimeMs() {
-        return throttleTimeMs;
+        return data.throttleTimeMs();
     }
 
     public Errors error() {
-        return error;
+        return Errors.forCode(data.errorCode());
     }
 
     public long expiryTimestamp() {
-        return expiryTimestamp;
+        return data.expiryTimestampMs();
     }
 
     public boolean hasError() {
-        return this.error != Errors.NONE;
+        return error() != Errors.NONE;
     }
 
     @Override
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/RequestContext.java b/clients/src/main/java/org/apache/kafka/common/requests/RequestContext.java
index 663d7460db93e..8e21f8e40be5e 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/RequestContext.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/RequestContext.java
@@ -17,6 +17,7 @@
 package org.apache.kafka.common.requests;
 
 import org.apache.kafka.common.errors.InvalidRequestException;
+import org.apache.kafka.common.message.ApiVersionsRequestData;
 import org.apache.kafka.common.network.ListenerName;
 import org.apache.kafka.common.network.Send;
 import org.apache.kafka.common.protocol.ApiKeys;
@@ -26,10 +27,11 @@
 
 import java.net.InetAddress;
 import java.nio.ByteBuffer;
+import org.apache.kafka.server.authorizer.AuthorizableRequestContext;
 
 import static org.apache.kafka.common.protocol.ApiKeys.API_VERSIONS;
 
-public class RequestContext {
+public class RequestContext implements AuthorizableRequestContext {
     public final RequestHeader header;
     public final String connectionId;
     public final InetAddress clientAddress;
@@ -54,7 +56,7 @@ public RequestContext(RequestHeader header,
     public RequestAndSize parseRequest(ByteBuffer buffer) {
         if (isUnsupportedApiVersionsRequest()) {
             // Unsupported ApiVersion requests are treated as v0 requests and are not parsed
-            ApiVersionsRequest apiVersionsRequest = new ApiVersionsRequest((short) 0, header.apiVersion());
+            ApiVersionsRequest apiVersionsRequest = new ApiVersionsRequest(new ApiVersionsRequestData(), (short) 0, header.apiVersion());
             return new RequestAndSize(apiVersionsRequest, 0);
         } else {
             ApiKeys apiKey = header.apiKey();
@@ -89,4 +91,43 @@ public short apiVersion() {
         return header.apiVersion();
     }
 
+    @Override
+    public String listenerName() {
+        return listenerName.value();
+    }
+
+    @Override
+    public SecurityProtocol securityProtocol() {
+        return securityProtocol;
+    }
+
+    @Override
+    public KafkaPrincipal principal() {
+        return principal;
+    }
+
+    @Override
+    public InetAddress clientAddress() {
+        return clientAddress;
+    }
+
+    @Override
+    public int requestType() {
+        return header.apiKey().id;
+    }
+
+    @Override
+    public int requestVersion() {
+        return header.apiVersion();
+    }
+
+    @Override
+    public String clientId() {
+        return header.clientId();
+    }
+
+    @Override
+    public int correlationId() {
+        return header.correlationId();
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/RequestHeader.java b/clients/src/main/java/org/apache/kafka/common/requests/RequestHeader.java
index 956d81335701a..3d80c4e7563d5 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/RequestHeader.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/RequestHeader.java
@@ -17,122 +17,95 @@
 package org.apache.kafka.common.requests;
 
 import org.apache.kafka.common.errors.InvalidRequestException;
+import org.apache.kafka.common.errors.UnsupportedVersionException;
+import org.apache.kafka.common.message.RequestHeaderData;
 import org.apache.kafka.common.protocol.ApiKeys;
-import org.apache.kafka.common.protocol.types.Field;
-import org.apache.kafka.common.protocol.types.Schema;
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
 import org.apache.kafka.common.protocol.types.Struct;
 
 import java.nio.ByteBuffer;
 
-import static java.util.Objects.requireNonNull;
-import static org.apache.kafka.common.protocol.types.Type.INT16;
-import static org.apache.kafka.common.protocol.types.Type.INT32;
-import static org.apache.kafka.common.protocol.types.Type.NULLABLE_STRING;
-
 /**
  * The header for a request in the Kafka protocol
  */
-public class RequestHeader extends AbstractRequestResponse {
-    private static final String API_KEY_FIELD_NAME = "api_key";
-    private static final String API_VERSION_FIELD_NAME = "api_version";
-    private static final String CLIENT_ID_FIELD_NAME = "client_id";
-    private static final String CORRELATION_ID_FIELD_NAME = "correlation_id";
-
-    public static final Schema SCHEMA = new Schema(
-            new Field(API_KEY_FIELD_NAME, INT16, "The id of the request type."),
-            new Field(API_VERSION_FIELD_NAME, INT16, "The version of the API."),
-            new Field(CORRELATION_ID_FIELD_NAME, INT32, "A user-supplied integer value that will be passed back with the response"),
-            new Field(CLIENT_ID_FIELD_NAME, NULLABLE_STRING, "A user specified identifier for the client making the request.", ""));
-
-    // Version 0 of the controlled shutdown API used a non-standard request header (the clientId is missing).
-    // This can be removed once we drop support for that version.
-    private static final Schema CONTROLLED_SHUTDOWN_V0_SCHEMA = new Schema(
-            new Field(API_KEY_FIELD_NAME, INT16, "The id of the request type."),
-            new Field(API_VERSION_FIELD_NAME, INT16, "The version of the API."),
-            new Field(CORRELATION_ID_FIELD_NAME, INT32, "A user-supplied integer value that will be passed back with the response"));
-
-    private final ApiKeys apiKey;
-    private final short apiVersion;
-    private final String clientId;
-    private final int correlationId;
-
-    public RequestHeader(Struct struct) {
-        short apiKey = struct.getShort(API_KEY_FIELD_NAME);
-        if (!ApiKeys.hasId(apiKey))
-            throw new InvalidRequestException("Unknown API key " + apiKey);
-
-        this.apiKey = ApiKeys.forId(apiKey);
-        apiVersion = struct.getShort(API_VERSION_FIELD_NAME);
-
-        // only v0 of the controlled shutdown request is missing the clientId
-        if (struct.hasField(CLIENT_ID_FIELD_NAME))
-            clientId = struct.getString(CLIENT_ID_FIELD_NAME);
-        else
-            clientId = "";
-        correlationId = struct.getInt(CORRELATION_ID_FIELD_NAME);
+public class RequestHeader implements AbstractRequestResponse {
+    private final RequestHeaderData data;
+    private final short headerVersion;
+
+    public RequestHeader(Struct struct, short headerVersion) {
+        this(new RequestHeaderData(struct, headerVersion), headerVersion);
+    }
+
+    public RequestHeader(ApiKeys requestApiKey, short requestVersion, String clientId, int correlationId) {
+        this(new RequestHeaderData().
+                setRequestApiKey(requestApiKey.id).
+                setRequestApiVersion(requestVersion).
+                setClientId(clientId).
+                setCorrelationId(correlationId),
+            ApiKeys.forId(requestApiKey.id).requestHeaderVersion(requestVersion));
     }
 
-    public RequestHeader(ApiKeys apiKey, short version, String clientId, int correlation) {
-        this.apiKey = requireNonNull(apiKey);
-        this.apiVersion = version;
-        this.clientId = clientId;
-        this.correlationId = correlation;
+    public RequestHeader(RequestHeaderData data, short headerVersion) {
+        this.data = data;
+        this.headerVersion = headerVersion;
     }
 
     public Struct toStruct() {
-        Schema schema = schema(apiKey.id, apiVersion);
-        Struct struct = new Struct(schema);
-        struct.set(API_KEY_FIELD_NAME, apiKey.id);
-        struct.set(API_VERSION_FIELD_NAME, apiVersion);
-
-        // only v0 of the controlled shutdown request is missing the clientId
-        if (struct.hasField(CLIENT_ID_FIELD_NAME))
-            struct.set(CLIENT_ID_FIELD_NAME, clientId);
-        struct.set(CORRELATION_ID_FIELD_NAME, correlationId);
-        return struct;
+        return this.data.toStruct(headerVersion);
     }
 
     public ApiKeys apiKey() {
-        return apiKey;
+        return ApiKeys.forId(data.requestApiKey());
     }
 
     public short apiVersion() {
-        return apiVersion;
+        return data.requestApiVersion();
+    }
+
+    public short headerVersion() {
+        return headerVersion;
     }
 
     public String clientId() {
-        return clientId;
+        return data.clientId();
     }
 
     public int correlationId() {
-        return correlationId;
+        return data.correlationId();
+    }
+
+    public RequestHeaderData data() {
+        return data;
     }
 
     public ResponseHeader toResponseHeader() {
-        return new ResponseHeader(correlationId);
+        return new ResponseHeader(data.correlationId(),
+            apiKey().responseHeaderVersion(apiVersion()));
     }
 
     public static RequestHeader parse(ByteBuffer buffer) {
+        short apiKey = -1;
         try {
-            short apiKey = buffer.getShort();
+            apiKey = buffer.getShort();
             short apiVersion = buffer.getShort();
-            Schema schema = schema(apiKey, apiVersion);
+            short headerVersion = ApiKeys.forId(apiKey).requestHeaderVersion(apiVersion);
             buffer.rewind();
-            return new RequestHeader(schema.read(buffer));
-        } catch (InvalidRequestException e) {
-            throw e;
-        } catch (Throwable  ex) {
+            return new RequestHeader(new RequestHeaderData(
+                new ByteBufferAccessor(buffer), headerVersion), headerVersion);
+        } catch (UnsupportedVersionException e) {
+            throw new InvalidRequestException("Unknown API key " + apiKey, e);
+        } catch (Throwable ex) {
             throw new InvalidRequestException("Error parsing request header. Our best guess of the apiKey is: " +
-                    buffer.getShort(0), ex);
+                    apiKey, ex);
         }
     }
 
     @Override
     public String toString() {
-        return "RequestHeader(apiKey=" + apiKey +
-                ", apiVersion=" + apiVersion +
-                ", clientId=" + clientId +
-                ", correlationId=" + correlationId +
+        return "RequestHeader(apiKey=" + apiKey() +
+                ", apiVersion=" + apiVersion() +
+                ", clientId=" + clientId() +
+                ", correlationId=" + correlationId() +
                 ")";
     }
 
@@ -140,28 +113,12 @@ public String toString() {
     public boolean equals(Object o) {
         if (this == o) return true;
         if (o == null || getClass() != o.getClass()) return false;
-
         RequestHeader that = (RequestHeader) o;
-        return apiKey == that.apiKey &&
-                apiVersion == that.apiVersion &&
-                correlationId == that.correlationId &&
-                (clientId == null ? that.clientId == null : clientId.equals(that.clientId));
+        return this.data.equals(that.data);
     }
 
     @Override
     public int hashCode() {
-        int result = apiKey.hashCode();
-        result = 31 * result + (int) apiVersion;
-        result = 31 * result + (clientId != null ? clientId.hashCode() : 0);
-        result = 31 * result + correlationId;
-        return result;
-    }
-
-    private static Schema schema(short apiKey, short version) {
-        if (apiKey == ApiKeys.CONTROLLED_SHUTDOWN.id && version == 0)
-            // This will be removed once we remove support for v0 of ControlledShutdownRequest, which
-            // depends on a non-standard request header (it does not have a clientId)
-            return CONTROLLED_SHUTDOWN_V0_SCHEMA;
-        return SCHEMA;
+        return this.data.hashCode();
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/RequestUtils.java b/clients/src/main/java/org/apache/kafka/common/requests/RequestUtils.java
index 24c2fbe441661..aae7a98fd42cd 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/RequestUtils.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/RequestUtils.java
@@ -28,6 +28,7 @@
 import org.apache.kafka.common.protocol.types.Struct;
 import org.apache.kafka.common.resource.ResourceType;
 
+import java.nio.ByteBuffer;
 import java.util.Optional;
 
 import static org.apache.kafka.common.protocol.CommonFields.HOST;
@@ -42,7 +43,7 @@
 import static org.apache.kafka.common.protocol.CommonFields.RESOURCE_PATTERN_TYPE_FILTER;
 import static org.apache.kafka.common.protocol.CommonFields.RESOURCE_TYPE;
 
-final class RequestUtils {
+public final class RequestUtils {
 
     private RequestUtils() {}
 
@@ -117,4 +118,24 @@ static Optional getLeaderEpoch(Struct struct, Field.Int32 leaderEpochFi
         return leaderEpochOpt;
     }
 
+    static Optional getLeaderEpoch(int leaderEpoch) {
+        Optional leaderEpochOpt = leaderEpoch == RecordBatch.NO_PARTITION_LEADER_EPOCH ?
+            Optional.empty() : Optional.of(leaderEpoch);
+        return leaderEpochOpt;
+    }
+
+    public static ByteBuffer serialize(Struct headerStruct, Struct bodyStruct) {
+        ByteBuffer buffer = ByteBuffer.allocate(headerStruct.sizeOf() + bodyStruct.sizeOf());
+        headerStruct.writeTo(buffer);
+        bodyStruct.writeTo(buffer);
+        buffer.rewind();
+        return buffer;
+    }
+
+    public static ByteBuffer serialize(Struct struct) {
+        ByteBuffer buffer = ByteBuffer.allocate(struct.sizeOf());
+        struct.writeTo(buffer);
+        buffer.rewind();
+        return buffer;
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/ResponseHeader.java b/clients/src/main/java/org/apache/kafka/common/requests/ResponseHeader.java
index fe452a2dc0808..118e5d3506d72 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/ResponseHeader.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/ResponseHeader.java
@@ -16,31 +16,30 @@
  */
 package org.apache.kafka.common.requests;
 
-import org.apache.kafka.common.protocol.types.BoundField;
-import org.apache.kafka.common.protocol.types.Field;
-import org.apache.kafka.common.protocol.types.Schema;
+import org.apache.kafka.common.message.ResponseHeaderData;
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
 import org.apache.kafka.common.protocol.types.Struct;
 
 import java.nio.ByteBuffer;
 
-import static org.apache.kafka.common.protocol.types.Type.INT32;
-
 /**
  * A response header in the kafka protocol.
  */
-public class ResponseHeader extends AbstractRequestResponse {
-    public static final Schema SCHEMA = new Schema(
-            new Field("correlation_id", INT32, "The user-supplied value passed in with the request"));
-    private static final BoundField CORRELATION_KEY_FIELD = SCHEMA.get("correlation_id");
+public class ResponseHeader implements AbstractRequestResponse {
+    private final ResponseHeaderData data;
+    private final short headerVersion;
 
-    private final int correlationId;
+    public ResponseHeader(Struct struct, short headerVersion) {
+        this(new ResponseHeaderData(struct, headerVersion), headerVersion);
+    }
 
-    public ResponseHeader(Struct struct) {
-        correlationId = struct.getInt(CORRELATION_KEY_FIELD);
+    public ResponseHeader(int correlationId, short headerVersion) {
+        this(new ResponseHeaderData().setCorrelationId(correlationId), headerVersion);
     }
 
-    public ResponseHeader(int correlationId) {
-        this.correlationId = correlationId;
+    public ResponseHeader(ResponseHeaderData data, short headerVersion) {
+        this.data = data;
+        this.headerVersion = headerVersion;
     }
 
     public int sizeOf() {
@@ -48,17 +47,24 @@ public int sizeOf() {
     }
 
     public Struct toStruct() {
-        Struct struct = new Struct(SCHEMA);
-        struct.set(CORRELATION_KEY_FIELD, correlationId);
-        return struct;
+        return data.toStruct(headerVersion);
     }
 
     public int correlationId() {
-        return correlationId;
+        return this.data.correlationId();
     }
 
-    public static ResponseHeader parse(ByteBuffer buffer) {
-        return new ResponseHeader(SCHEMA.read(buffer));
+    public short headerVersion() {
+        return headerVersion;
     }
 
+    public ResponseHeaderData data() {
+        return data;
+    }
+
+    public static ResponseHeader parse(ByteBuffer buffer, short headerVersion) {
+        return new ResponseHeader(
+            new ResponseHeaderData(new ByteBufferAccessor(buffer), headerVersion),
+                headerVersion);
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/StopReplicaRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/StopReplicaRequest.java
index 7c03c799e7543..ea8046875b6bf 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/StopReplicaRequest.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/StopReplicaRequest.java
@@ -17,73 +17,65 @@
 package org.apache.kafka.common.requests;
 
 import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.message.StopReplicaRequestData;
+import org.apache.kafka.common.message.StopReplicaRequestData.StopReplicaPartitionV0;
+import org.apache.kafka.common.message.StopReplicaRequestData.StopReplicaTopic;
+import org.apache.kafka.common.message.StopReplicaResponseData;
+import org.apache.kafka.common.message.StopReplicaResponseData.StopReplicaPartitionError;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.protocol.types.Field;
-import org.apache.kafka.common.protocol.types.Schema;
 import org.apache.kafka.common.protocol.types.Struct;
 import org.apache.kafka.common.utils.CollectionUtils;
+import org.apache.kafka.common.utils.FlattenedIterator;
+import org.apache.kafka.common.utils.MappedIterator;
 import org.apache.kafka.common.utils.Utils;
 
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
-import java.util.HashMap;
-import java.util.HashSet;
+import java.util.Collection;
 import java.util.List;
 import java.util.Map;
-import java.util.Set;
-
-import static org.apache.kafka.common.protocol.CommonFields.PARTITION_ID;
-import static org.apache.kafka.common.protocol.CommonFields.TOPIC_NAME;
-import static org.apache.kafka.common.protocol.types.Type.INT32;
+import java.util.stream.Collectors;
 
 public class StopReplicaRequest extends AbstractControlRequest {
-    private static final Field.Bool DELETE_PARTITIONS = new Field.Bool("delete_partitions", "Boolean which indicates if replica's partitions must be deleted.");
-    private static final Field.ComplexArray PARTITIONS = new Field.ComplexArray("partitions", "The partitions");
-    private static final Field.Array PARTITION_IDS = new Field.Array("partition_ids", INT32, "The partition ids of a topic");
-
-    private static final Field PARTITIONS_V0 = PARTITIONS.withFields(
-            TOPIC_NAME,
-            PARTITION_ID);
-    private static final Field PARTITIONS_V1 = PARTITIONS.withFields(
-            TOPIC_NAME,
-            PARTITION_IDS);
-
-    private static final Schema STOP_REPLICA_REQUEST_V0 = new Schema(
-            CONTROLLER_ID,
-            CONTROLLER_EPOCH,
-            DELETE_PARTITIONS,
-            PARTITIONS_V0);
-
-    // STOP_REPLICA_REQUEST_V1 added a broker_epoch Field. This field specifies the generation of the broker across
-    // bounces. It also normalizes partitions under each topic.
-    private static final Schema STOP_REPLICA_REQUEST_V1 = new Schema(
-            CONTROLLER_ID,
-            CONTROLLER_EPOCH,
-            BROKER_EPOCH,
-            DELETE_PARTITIONS,
-            PARTITIONS_V1);
-
-
-    public static Schema[] schemaVersions() {
-        return new Schema[] {STOP_REPLICA_REQUEST_V0, STOP_REPLICA_REQUEST_V1};
-    }
 
     public static class Builder extends AbstractControlRequest.Builder {
         private final boolean deletePartitions;
-        private final Set partitions;
+        private final Collection partitions;
 
-        public Builder(short version, int controllerId, int controllerEpoch, long brokerEpoch, boolean deletePartitions,
-                       Set partitions) {
-            super(ApiKeys.STOP_REPLICA, version, controllerId, controllerEpoch, brokerEpoch);
+        public Builder(short version, int controllerId, int controllerEpoch, long brokerEpoch, long maxBrokerEpoch, boolean deletePartitions,
+                       Collection partitions) {
+            super(ApiKeys.STOP_REPLICA, version, controllerId, controllerEpoch, brokerEpoch, maxBrokerEpoch);
             this.deletePartitions = deletePartitions;
             this.partitions = partitions;
         }
 
-        @Override
         public StopReplicaRequest build(short version) {
-            return new StopReplicaRequest(controllerId, controllerEpoch, brokerEpoch,
-                    deletePartitions, partitions, version);
+            StopReplicaRequestData data = new StopReplicaRequestData()
+                .setControllerId(controllerId)
+                .setControllerEpoch(controllerEpoch)
+                .setBrokerEpoch(brokerEpoch)
+                .setMaxBrokerEpoch(maxBrokerEpoch)
+                .setDeletePartitions(deletePartitions);
+
+            if (version >= 1) {
+                Map> topicPartitionsMap = CollectionUtils.groupPartitionsByTopic(partitions);
+                List topics = topicPartitionsMap.entrySet().stream().map(entry ->
+                    new StopReplicaTopic()
+                        .setName(entry.getKey())
+                        .setPartitionIndexes(entry.getValue())
+                ).collect(Collectors.toList());
+                data.setTopics(topics);
+            } else {
+                List requestPartitions = partitions.stream().map(tp ->
+                    new StopReplicaPartitionV0()
+                        .setTopicName(tp.topic())
+                        .setPartitionIndex(tp.partition())
+                ).collect(Collectors.toList());
+                data.setUngroupedPartitions(requestPartitions);
+            }
+
+            return new StopReplicaRequest(data, version);
         }
 
         @Override
@@ -94,108 +86,94 @@ public String toString() {
                 append(", controllerEpoch=").append(controllerEpoch).
                 append(", deletePartitions=").append(deletePartitions).
                 append(", brokerEpoch=").append(brokerEpoch).
+                append(", maxBrokerEpoch=").append(maxBrokerEpoch).
                 append(", partitions=").append(Utils.join(partitions, ",")).
                 append(")");
             return bld.toString();
         }
     }
 
-    private final boolean deletePartitions;
-    private final Set partitions;
+    private final StopReplicaRequestData data;
 
-    private StopReplicaRequest(int controllerId, int controllerEpoch, long brokerEpoch, boolean deletePartitions,
-                               Set partitions, short version) {
-        super(ApiKeys.STOP_REPLICA, version, controllerId, controllerEpoch, brokerEpoch);
-        this.deletePartitions = deletePartitions;
-        this.partitions = partitions;
+    private StopReplicaRequest(StopReplicaRequestData data, short version) {
+        super(ApiKeys.STOP_REPLICA, version);
+        this.data = data;
     }
 
     public StopReplicaRequest(Struct struct, short version) {
-        super(ApiKeys.STOP_REPLICA, struct, version);
-
-        partitions = new HashSet<>();
-        if (version > 0) { // V1
-            for (Object topicObj : struct.get(PARTITIONS)) {
-                Struct topicData = (Struct) topicObj;
-                String topic = topicData.get(TOPIC_NAME);
-                for (Object partitionObj : topicData.get(PARTITION_IDS)) {
-                    int partition = (Integer) partitionObj;
-                    partitions.add(new TopicPartition(topic, partition));
-                }
-            }
-        } else { // V0
-            for (Object partitionDataObj : struct.get(PARTITIONS)) {
-                Struct partitionData = (Struct) partitionDataObj;
-                String topic = partitionData.get(TOPIC_NAME);
-                int partition = partitionData.get(PARTITION_ID);
-                partitions.add(new TopicPartition(topic, partition));
-            }
-        }
-        deletePartitions = struct.get(DELETE_PARTITIONS);
+        this(new StopReplicaRequestData(struct, version), version);
     }
 
     @Override
     public StopReplicaResponse getErrorResponse(int throttleTimeMs, Throwable e) {
         Errors error = Errors.forException(e);
 
-        Map responses = new HashMap<>(partitions.size());
-        for (TopicPartition partition : partitions) {
-            responses.put(partition, error);
+        StopReplicaResponseData data = new StopReplicaResponseData();
+        data.setErrorCode(error.code());
+        List partitions = new ArrayList<>();
+        for (TopicPartition tp : partitions()) {
+            partitions.add(new StopReplicaPartitionError()
+                .setTopicName(tp.topic())
+                .setPartitionIndex(tp.partition())
+                .setErrorCode(error.code()));
         }
+        data.setPartitionErrors(partitions);
+        return new StopReplicaResponse(data);
+    }
+
+    public boolean deletePartitions() {
+        return data.deletePartitions();
+    }
 
-        short versionId = version();
-        switch (versionId) {
-            case 0:
-            case 1:
-                return new StopReplicaResponse(error, responses);
-            default:
-                throw new IllegalArgumentException(String.format("Version %d is not valid. Valid versions for %s are 0 to %d",
-                        versionId, this.getClass().getSimpleName(), ApiKeys.STOP_REPLICA.latestVersion()));
+    /**
+     * Note that this method has allocation overhead per iterated element, so callers should copy the result into
+     * another collection if they need to iterate more than once.
+     *
+     * Implementation note: we should strive to avoid allocation overhead per element, see
+     * `UpdateMetadataRequest.partitionStates()` for the preferred approach. That's not possible in this case and
+     * StopReplicaRequest should be relatively rare in comparison to other request types.
+     */
+    public Iterable partitions() {
+        if (version() >= 1) {
+            return () -> new FlattenedIterator<>(data.topics().iterator(), topic ->
+                new MappedIterator<>(topic.partitionIndexes().iterator(), partition ->
+                    new TopicPartition(topic.name(), partition)));
         }
+        return () -> new MappedIterator<>(data.ungroupedPartitions().iterator(),
+            partition -> new TopicPartition(partition.topicName(), partition.partitionIndex()));
     }
 
-    public boolean deletePartitions() {
-        return deletePartitions;
+    @Override
+    public int controllerId() {
+        return data.controllerId();
     }
 
-    public Set partitions() {
-        return partitions;
+    @Override
+    public int controllerEpoch() {
+        return data.controllerEpoch();
+    }
+
+    @Override
+    public long brokerEpoch() {
+        return data.brokerEpoch();
+    }
+
+    @Override
+    public long maxBrokerEpoch() {
+        return data.maxBrokerEpoch();
     }
 
     public static StopReplicaRequest parse(ByteBuffer buffer, short version) {
         return new StopReplicaRequest(ApiKeys.STOP_REPLICA.parseRequest(version, buffer), version);
     }
 
+    // Visible for testing
+    StopReplicaRequestData data() {
+        return data;
+    }
+
     @Override
     protected Struct toStruct() {
-        Struct struct = new Struct(ApiKeys.STOP_REPLICA.requestSchema(version()));
-
-        struct.set(CONTROLLER_ID, controllerId);
-        struct.set(CONTROLLER_EPOCH, controllerEpoch);
-        struct.setIfExists(BROKER_EPOCH, brokerEpoch);
-        struct.set(DELETE_PARTITIONS, deletePartitions);
-
-        if (version() > 0) { // V1
-            Map> topicPartitionsMap = CollectionUtils.groupPartitionsByTopic(partitions);
-            List topicsData = new ArrayList<>(topicPartitionsMap.size());
-            for (Map.Entry> entry : topicPartitionsMap.entrySet()) {
-                Struct topicData = struct.instance(PARTITIONS);
-                topicData.set(TOPIC_NAME, entry.getKey());
-                topicData.set(PARTITION_IDS, entry.getValue().toArray());
-                topicsData.add(topicData);
-            }
-            struct.set(PARTITIONS, topicsData.toArray());
-
-        } else { // V0
-            List partitionDatas = new ArrayList<>(partitions.size());
-            for (TopicPartition partition : partitions) {
-                Struct partitionData = struct.instance(PARTITIONS);
-                partitionData.set(TOPIC_NAME, partition.topic());
-                partitionData.set(PARTITION_ID, partition.partition());
-                partitionDatas.add(partitionData);
-            }
-            struct.set(PARTITIONS, partitionDatas.toArray());
-        }
-        return struct;
+        return data.toStruct(version());
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/StopReplicaResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/StopReplicaResponse.java
index 6089bea8d4f0a..7d6c7a0742899 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/StopReplicaResponse.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/StopReplicaResponse.java
@@ -16,43 +16,19 @@
  */
 package org.apache.kafka.common.requests;
 
-import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.message.StopReplicaResponseData;
+import org.apache.kafka.common.message.StopReplicaResponseData.StopReplicaPartitionError;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.protocol.types.Field;
-import org.apache.kafka.common.protocol.types.Schema;
 import org.apache.kafka.common.protocol.types.Struct;
 
 import java.nio.ByteBuffer;
-import java.util.ArrayList;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-
-import static org.apache.kafka.common.protocol.CommonFields.ERROR_CODE;
-import static org.apache.kafka.common.protocol.CommonFields.PARTITION_ID;
-import static org.apache.kafka.common.protocol.CommonFields.TOPIC_NAME;
+import java.util.stream.Collectors;
 
 public class StopReplicaResponse extends AbstractResponse {
-    private static final Field.ComplexArray PARTITIONS = new Field.ComplexArray("partitions", "Response for the requests partitions");
-
-    private static final Field PARTITIONS_V0 = PARTITIONS.withFields(
-            TOPIC_NAME,
-            PARTITION_ID,
-            ERROR_CODE);
-    private static final Schema STOP_REPLICA_RESPONSE_V0 = new Schema(
-            ERROR_CODE,
-            PARTITIONS_V0);
-
-    private static final Schema STOP_REPLICA_RESPONSE_V1 = STOP_REPLICA_RESPONSE_V0;
-
-
-    public static Schema[] schemaVersions() {
-        return new Schema[] {STOP_REPLICA_RESPONSE_V0, STOP_REPLICA_RESPONSE_V1};
-    }
-
-    private final Map responses;
 
     /**
      * Possible error code:
@@ -60,71 +36,44 @@ public static Schema[] schemaVersions() {
      * STALE_CONTROLLER_EPOCH (11)
      * STALE_BROKER_EPOCH (77)
      */
-    private final Errors error;
+    private final StopReplicaResponseData data;
 
-    public StopReplicaResponse(Errors error, Map responses) {
-        this.responses = responses;
-        this.error = error;
+    public StopReplicaResponse(StopReplicaResponseData data) {
+        this.data = data;
     }
 
-    public StopReplicaResponse(Struct struct) {
-        responses = new HashMap<>();
-        for (Object responseDataObj : struct.get(PARTITIONS)) {
-            Struct responseData = (Struct) responseDataObj;
-            String topic = responseData.get(TOPIC_NAME);
-            int partition = responseData.get(PARTITION_ID);
-            Errors error = Errors.forCode(responseData.get(ERROR_CODE));
-            responses.put(new TopicPartition(topic, partition), error);
-        }
-
-        error = Errors.forCode(struct.get(ERROR_CODE));
+    public StopReplicaResponse(Struct struct, short version) {
+        data = new StopReplicaResponseData(struct, version);
     }
 
-    public Map responses() {
-        return responses;
+    public List partitionErrors() {
+        return data.partitionErrors();
     }
 
     public Errors error() {
-        return error;
+        return Errors.forCode(data.errorCode());
     }
 
     @Override
     public Map errorCounts() {
-        if (error != Errors.NONE)
+        if (data.errorCode() != Errors.NONE.code())
             // Minor optimization since the top-level error applies to all partitions
-            return Collections.singletonMap(error, responses.size());
-        return errorCounts(responses);
+            return Collections.singletonMap(error(), data.partitionErrors().size());
+        return errorCounts(data.partitionErrors().stream().map(p -> Errors.forCode(p.errorCode())).collect(Collectors.toList()));
     }
 
     public static StopReplicaResponse parse(ByteBuffer buffer, short version) {
-        return new StopReplicaResponse(ApiKeys.STOP_REPLICA.parseResponse(version, buffer));
+        return new StopReplicaResponse(ApiKeys.STOP_REPLICA.parseResponse(version, buffer), version);
     }
 
     @Override
     protected Struct toStruct(short version) {
-        Struct struct = new Struct(ApiKeys.STOP_REPLICA.responseSchema(version));
-
-        List responseDatas = new ArrayList<>(responses.size());
-        for (Map.Entry response : responses.entrySet()) {
-            Struct partitionData = struct.instance(PARTITIONS);
-            TopicPartition partition = response.getKey();
-            partitionData.set(TOPIC_NAME, partition.topic());
-            partitionData.set(PARTITION_ID, partition.partition());
-            partitionData.set(ERROR_CODE, response.getValue().code());
-            responseDatas.add(partitionData);
-        }
-
-        struct.set(PARTITIONS, responseDatas.toArray());
-        struct.set(ERROR_CODE, error.code());
-        return struct;
+        return data.toStruct(version);
     }
 
     @Override
     public String toString() {
-        return "StopReplicaResponse(" +
-                "responses=" + responses +
-                ", error=" + error +
-                ")";
+        return data.toString();
     }
 
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/SyncGroupRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/SyncGroupRequest.java
index 237320f9cd25e..48319d55cde5f 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/SyncGroupRequest.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/SyncGroupRequest.java
@@ -16,112 +16,53 @@
  */
 package org.apache.kafka.common.requests;
 
+import org.apache.kafka.common.errors.UnsupportedVersionException;
+import org.apache.kafka.common.message.SyncGroupRequestData;
+import org.apache.kafka.common.message.SyncGroupResponseData;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.protocol.types.ArrayOf;
-import org.apache.kafka.common.protocol.types.Field;
-import org.apache.kafka.common.protocol.types.Schema;
 import org.apache.kafka.common.protocol.types.Struct;
-import org.apache.kafka.common.utils.Utils;
 
 import java.nio.ByteBuffer;
-import java.util.ArrayList;
 import java.util.HashMap;
-import java.util.List;
 import java.util.Map;
 
-import static org.apache.kafka.common.protocol.CommonFields.GENERATION_ID;
-import static org.apache.kafka.common.protocol.CommonFields.GROUP_ID;
-import static org.apache.kafka.common.protocol.CommonFields.MEMBER_ID;
-import static org.apache.kafka.common.protocol.types.Type.BYTES;
-
 public class SyncGroupRequest extends AbstractRequest {
-    private static final String MEMBER_ASSIGNMENT_KEY_NAME = "member_assignment";
-    private static final String GROUP_ASSIGNMENT_KEY_NAME = "group_assignment";
-
-    private static final Schema SYNC_GROUP_REQUEST_MEMBER_V0 = new Schema(
-            MEMBER_ID,
-            new Field(MEMBER_ASSIGNMENT_KEY_NAME, BYTES));
-    private static final Schema SYNC_GROUP_REQUEST_V0 = new Schema(
-            GROUP_ID,
-            GENERATION_ID,
-            MEMBER_ID,
-            new Field(GROUP_ASSIGNMENT_KEY_NAME, new ArrayOf(SYNC_GROUP_REQUEST_MEMBER_V0)));
-
-    /* v1 request is the same as v0. Throttle time has been added to response */
-    private static final Schema SYNC_GROUP_REQUEST_V1 = SYNC_GROUP_REQUEST_V0;
-
-    /**
-     * The version number is bumped to indicate that on quota violation brokers send out responses before throttling.
-     */
-    private static final Schema SYNC_GROUP_REQUEST_V2 = SYNC_GROUP_REQUEST_V1;
-
-    public static Schema[] schemaVersions() {
-        return new Schema[] {SYNC_GROUP_REQUEST_V0, SYNC_GROUP_REQUEST_V1,
-            SYNC_GROUP_REQUEST_V2};
-    }
 
     public static class Builder extends AbstractRequest.Builder {
-        private final String groupId;
-        private final int generationId;
-        private final String memberId;
-        private final Map groupAssignment;
 
-        public Builder(String groupId, int generationId, String memberId,
-                       Map groupAssignment) {
+        private final SyncGroupRequestData data;
+
+        public Builder(SyncGroupRequestData data) {
             super(ApiKeys.SYNC_GROUP);
-            this.groupId = groupId;
-            this.generationId = generationId;
-            this.memberId = memberId;
-            this.groupAssignment = groupAssignment;
+            this.data = data;
         }
 
         @Override
         public SyncGroupRequest build(short version) {
-            return new SyncGroupRequest(groupId, generationId, memberId, groupAssignment, version);
+            if (data.groupInstanceId() != null && version < 3) {
+                throw new UnsupportedVersionException("The broker sync group protocol version " +
+                        version + " does not support usage of config group.instance.id.");
+            }
+            return new SyncGroupRequest(data, version);
         }
 
         @Override
         public String toString() {
-            StringBuilder bld = new StringBuilder();
-            bld.append("(type=SyncGroupRequest").
-                    append(", groupId=").append(groupId).
-                    append(", generationId=").append(generationId).
-                    append(", memberId=").append(memberId).
-                    append(", groupAssignment=").
-                    append(Utils.join(groupAssignment.keySet(), ",")).
-                    append(")");
-            return bld.toString();
+            return data.toString();
         }
     }
-    private final String groupId;
-    private final int generationId;
-    private final String memberId;
-    private final Map groupAssignment;
 
-    private SyncGroupRequest(String groupId, int generationId, String memberId,
-                             Map groupAssignment, short version) {
+    public final SyncGroupRequestData data;
+
+    public SyncGroupRequest(SyncGroupRequestData data, short version) {
         super(ApiKeys.SYNC_GROUP, version);
-        this.groupId = groupId;
-        this.generationId = generationId;
-        this.memberId = memberId;
-        this.groupAssignment = groupAssignment;
+        this.data = data;
     }
 
     public SyncGroupRequest(Struct struct, short version) {
         super(ApiKeys.SYNC_GROUP, version);
-        this.groupId = struct.get(GROUP_ID);
-        this.generationId = struct.get(GENERATION_ID);
-        this.memberId = struct.get(MEMBER_ID);
-
-        groupAssignment = new HashMap<>();
-
-        for (Object memberDataObj : struct.getArray(GROUP_ASSIGNMENT_KEY_NAME)) {
-            Struct memberData = (Struct) memberDataObj;
-            String memberId = memberData.get(MEMBER_ID);
-            ByteBuffer memberMetadata = memberData.getBytes(MEMBER_ASSIGNMENT_KEY_NAME);
-            groupAssignment.put(memberId, memberMetadata);
-        }
+        this.data = new SyncGroupRequestData(struct, version);
     }
 
     @Override
@@ -130,34 +71,31 @@ public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
         switch (versionId) {
             case 0:
                 return new SyncGroupResponse(
-                        Errors.forException(e),
-                        ByteBuffer.wrap(new byte[]{}));
+                        new SyncGroupResponseData()
+                            .setErrorCode(Errors.forException(e).code())
+                            .setAssignment(new byte[0])
+                       );
             case 1:
             case 2:
+            case 3:
                 return new SyncGroupResponse(
-                        throttleTimeMs,
-                        Errors.forException(e),
-                        ByteBuffer.wrap(new byte[]{}));
+                        new SyncGroupResponseData()
+                            .setErrorCode(Errors.forException(e).code())
+                            .setAssignment(new byte[0])
+                            .setThrottleTimeMs(throttleTimeMs)
+                );
             default:
                 throw new IllegalArgumentException(String.format("Version %d is not valid. Valid versions for %s are 0 to %d",
                         versionId, this.getClass().getSimpleName(), ApiKeys.SYNC_GROUP.latestVersion()));
         }
     }
 
-    public String groupId() {
-        return groupId;
-    }
-
-    public int generationId() {
-        return generationId;
-    }
-
-    public Map groupAssignment() {
-        return groupAssignment;
-    }
-
-    public String memberId() {
-        return memberId;
+    public Map groupAssignments() {
+        Map groupAssignments = new HashMap<>();
+        for (SyncGroupRequestData.SyncGroupRequestAssignment assignment : data.assignments()) {
+            groupAssignments.put(assignment.memberId(), ByteBuffer.wrap(assignment.assignment()));
+        }
+        return groupAssignments;
     }
 
     public static SyncGroupRequest parse(ByteBuffer buffer, short version) {
@@ -166,19 +104,6 @@ public static SyncGroupRequest parse(ByteBuffer buffer, short version) {
 
     @Override
     protected Struct toStruct() {
-        Struct struct = new Struct(ApiKeys.SYNC_GROUP.requestSchema(version()));
-        struct.set(GROUP_ID, groupId);
-        struct.set(GENERATION_ID, generationId);
-        struct.set(MEMBER_ID, memberId);
-
-        List memberArray = new ArrayList<>();
-        for (Map.Entry entries: groupAssignment.entrySet()) {
-            Struct memberData = struct.instance(GROUP_ASSIGNMENT_KEY_NAME);
-            memberData.set(MEMBER_ID, entries.getKey());
-            memberData.set(MEMBER_ASSIGNMENT_KEY_NAME, entries.getValue());
-            memberArray.add(memberData);
-        }
-        struct.set(GROUP_ASSIGNMENT_KEY_NAME, memberArray.toArray());
-        return struct;
+        return data.toStruct(version());
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/SyncGroupResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/SyncGroupResponse.java
index 2b2fc6fbc3390..8c9bad13d9d24 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/SyncGroupResponse.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/SyncGroupResponse.java
@@ -16,100 +16,49 @@
  */
 package org.apache.kafka.common.requests;
 
+import org.apache.kafka.common.message.SyncGroupResponseData;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.protocol.types.Field;
-import org.apache.kafka.common.protocol.types.Schema;
 import org.apache.kafka.common.protocol.types.Struct;
 
 import java.nio.ByteBuffer;
+import java.util.Collections;
 import java.util.Map;
 
-import static org.apache.kafka.common.protocol.CommonFields.ERROR_CODE;
-import static org.apache.kafka.common.protocol.CommonFields.THROTTLE_TIME_MS;
-import static org.apache.kafka.common.protocol.types.Type.BYTES;
-
 public class SyncGroupResponse extends AbstractResponse {
-    private static final String MEMBER_ASSIGNMENT_KEY_NAME = "member_assignment";
-
-    private static final Schema SYNC_GROUP_RESPONSE_V0 = new Schema(
-            ERROR_CODE,
-            new Field(MEMBER_ASSIGNMENT_KEY_NAME, BYTES));
-    private static final Schema SYNC_GROUP_RESPONSE_V1 = new Schema(
-            THROTTLE_TIME_MS,
-            ERROR_CODE,
-            new Field(MEMBER_ASSIGNMENT_KEY_NAME, BYTES));
-
-    /**
-     * The version number is bumped to indicate that on quota violation brokers send out responses before throttling.
-     */
-    private static final Schema SYNC_GROUP_RESPONSE_V2 = SYNC_GROUP_RESPONSE_V1;
-
-    public static Schema[] schemaVersions() {
-        return new Schema[] {SYNC_GROUP_RESPONSE_V0, SYNC_GROUP_RESPONSE_V1,
-            SYNC_GROUP_RESPONSE_V2};
-    }
 
-    /**
-     * Possible error codes:
-     *
-     * COORDINATOR_NOT_AVAILABLE (15)
-     * NOT_COORDINATOR (16)
-     * ILLEGAL_GENERATION (22)
-     * UNKNOWN_MEMBER_ID (25)
-     * REBALANCE_IN_PROGRESS (27)
-     * GROUP_AUTHORIZATION_FAILED (30)
-     *
-     * NOTE: Currently the coordinator returns REBALANCE_IN_PROGRESS while the coordinator is
-     * loading. On the next protocol bump, we should consider using COORDINATOR_LOAD_IN_PROGRESS
-     * to be consistent with the other APIs.
-     */
+    public final SyncGroupResponseData data;
 
-    private final Errors error;
-    private final int throttleTimeMs;
-    private final ByteBuffer memberState;
-
-    public SyncGroupResponse(Errors error, ByteBuffer memberState) {
-        this(DEFAULT_THROTTLE_TIME, error, memberState);
+    public SyncGroupResponse(SyncGroupResponseData data) {
+        this.data = data;
     }
 
-    public SyncGroupResponse(int throttleTimeMs, Errors error, ByteBuffer memberState) {
-        this.throttleTimeMs = throttleTimeMs;
-        this.error = error;
-        this.memberState = memberState;
+    public SyncGroupResponse(Struct struct) {
+        short latestVersion = (short) (SyncGroupResponseData.SCHEMAS.length - 1);
+        this.data = new SyncGroupResponseData(struct, latestVersion);
     }
 
-    public SyncGroupResponse(Struct struct) {
-        this.throttleTimeMs = struct.getOrElse(THROTTLE_TIME_MS, DEFAULT_THROTTLE_TIME);
-        this.error = Errors.forCode(struct.get(ERROR_CODE));
-        this.memberState = struct.getBytes(MEMBER_ASSIGNMENT_KEY_NAME);
+    public SyncGroupResponse(Struct struct, short version) {
+        this.data = new SyncGroupResponseData(struct, version);
     }
 
     @Override
     public int throttleTimeMs() {
-        return throttleTimeMs;
+        return data.throttleTimeMs();
     }
 
     public Errors error() {
-        return error;
+        return Errors.forCode(data.errorCode());
     }
 
     @Override
     public Map errorCounts() {
-        return errorCounts(error);
-    }
-
-    public ByteBuffer memberAssignment() {
-        return memberState;
+        return Collections.singletonMap(Errors.forCode(data.errorCode()), 1);
     }
 
     @Override
     protected Struct toStruct(short version) {
-        Struct struct = new Struct(ApiKeys.SYNC_GROUP.responseSchema(version));
-        struct.setIfExists(THROTTLE_TIME_MS, throttleTimeMs);
-        struct.set(ERROR_CODE, error.code());
-        struct.set(MEMBER_ASSIGNMENT_KEY_NAME, memberState);
-        return struct;
+        return data.toStruct(version);
     }
 
     public static SyncGroupResponse parse(ByteBuffer buffer, short version) {
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/TxnOffsetCommitRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/TxnOffsetCommitRequest.java
index 1c922e1dd525c..279fc9cb070a7 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/TxnOffsetCommitRequest.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/TxnOffsetCommitRequest.java
@@ -17,225 +17,129 @@
 package org.apache.kafka.common.requests;
 
 import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.message.TxnOffsetCommitRequestData;
+import org.apache.kafka.common.message.TxnOffsetCommitRequestData.TxnOffsetCommitRequestPartition;
+import org.apache.kafka.common.message.TxnOffsetCommitRequestData.TxnOffsetCommitRequestTopic;
+import org.apache.kafka.common.message.TxnOffsetCommitResponseData;
+import org.apache.kafka.common.message.TxnOffsetCommitResponseData.TxnOffsetCommitResponsePartition;
+import org.apache.kafka.common.message.TxnOffsetCommitResponseData.TxnOffsetCommitResponseTopic;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.protocol.types.Field;
-import org.apache.kafka.common.protocol.types.Schema;
 import org.apache.kafka.common.protocol.types.Struct;
-import org.apache.kafka.common.utils.CollectionUtils;
+import org.apache.kafka.common.record.RecordBatch;
 
 import java.nio.ByteBuffer;
+import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Optional;
-
-import static org.apache.kafka.common.protocol.CommonFields.COMMITTED_LEADER_EPOCH;
-import static org.apache.kafka.common.protocol.CommonFields.COMMITTED_METADATA;
-import static org.apache.kafka.common.protocol.CommonFields.COMMITTED_OFFSET;
-import static org.apache.kafka.common.protocol.CommonFields.GROUP_ID;
-import static org.apache.kafka.common.protocol.CommonFields.PARTITION_ID;
-import static org.apache.kafka.common.protocol.CommonFields.PRODUCER_EPOCH;
-import static org.apache.kafka.common.protocol.CommonFields.PRODUCER_ID;
-import static org.apache.kafka.common.protocol.CommonFields.TOPIC_NAME;
-import static org.apache.kafka.common.protocol.CommonFields.TRANSACTIONAL_ID;
+import java.util.stream.Collectors;
 
 public class TxnOffsetCommitRequest extends AbstractRequest {
-    // top level fields
-    private static final Field.ComplexArray TOPICS = new Field.ComplexArray("topics",
-            "Topics to commit offsets");
-
-    // topic level fields
-    private static final Field.ComplexArray PARTITIONS = new Field.ComplexArray("partitions",
-            "Partitions to commit offsets");
-
-    private static final Field PARTITIONS_V0 = PARTITIONS.withFields(
-            PARTITION_ID,
-            COMMITTED_OFFSET,
-            COMMITTED_METADATA);
-
-    private static final Field TOPICS_V0 = TOPICS.withFields(
-            TOPIC_NAME,
-            PARTITIONS_V0);
-
-    private static final Schema TXN_OFFSET_COMMIT_REQUEST_V0 = new Schema(
-            TRANSACTIONAL_ID,
-            GROUP_ID,
-            PRODUCER_ID,
-            PRODUCER_EPOCH,
-            TOPICS_V0);
-
-    // V1 bump used to indicate that on quota violation brokers send out responses before throttling.
-    private static final Schema TXN_OFFSET_COMMIT_REQUEST_V1 = TXN_OFFSET_COMMIT_REQUEST_V0;
 
-    // V2 adds the leader epoch to the partition data
-    private static final Field PARTITIONS_V2 = PARTITIONS.withFields(
-            PARTITION_ID,
-            COMMITTED_OFFSET,
-            COMMITTED_LEADER_EPOCH,
-            COMMITTED_METADATA);
-
-    private static final Field TOPICS_V2 = TOPICS.withFields(
-            TOPIC_NAME,
-            PARTITIONS_V2);
-
-    private static final Schema TXN_OFFSET_COMMIT_REQUEST_V2 = new Schema(
-            TRANSACTIONAL_ID,
-            GROUP_ID,
-            PRODUCER_ID,
-            PRODUCER_EPOCH,
-            TOPICS_V2);
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{TXN_OFFSET_COMMIT_REQUEST_V0, TXN_OFFSET_COMMIT_REQUEST_V1, TXN_OFFSET_COMMIT_REQUEST_V2};
-    }
+    public final TxnOffsetCommitRequestData data;
 
     public static class Builder extends AbstractRequest.Builder {
-        private final String transactionalId;
-        private final String consumerGroupId;
-        private final long producerId;
-        private final short producerEpoch;
-        private final Map offsets;
 
-        public Builder(String transactionalId, String consumerGroupId, long producerId, short producerEpoch,
-                       Map offsets) {
-            super(ApiKeys.TXN_OFFSET_COMMIT);
-            this.transactionalId = transactionalId;
-            this.consumerGroupId = consumerGroupId;
-            this.producerId = producerId;
-            this.producerEpoch = producerEpoch;
-            this.offsets = offsets;
-        }
-
-        public String consumerGroupId() {
-            return consumerGroupId;
-        }
+        public final TxnOffsetCommitRequestData data;
 
-        public Map offsets() {
-            return offsets;
+        public Builder(TxnOffsetCommitRequestData data) {
+            super(ApiKeys.TXN_OFFSET_COMMIT);
+            this.data = data;
         }
 
         @Override
         public TxnOffsetCommitRequest build(short version) {
-            return new TxnOffsetCommitRequest(version, transactionalId, consumerGroupId, producerId, producerEpoch, offsets);
+            return new TxnOffsetCommitRequest(data, version);
         }
 
         @Override
         public String toString() {
-            StringBuilder bld = new StringBuilder();
-            bld.append("(type=TxnOffsetCommitRequest").
-                    append(", transactionalId=").append(transactionalId).
-                    append(", producerId=").append(producerId).
-                    append(", producerEpoch=").append(producerEpoch).
-                    append(", consumerGroupId=").append(consumerGroupId).
-                    append(", offsets=").append(offsets).
-                    append(")");
-            return bld.toString();
+            return data.toString();
         }
     }
 
-    private final String transactionalId;
-    private final String consumerGroupId;
-    private final long producerId;
-    private final short producerEpoch;
-    private final Map offsets;
-
-    public TxnOffsetCommitRequest(short version, String transactionalId, String consumerGroupId, long producerId,
-                                  short producerEpoch, Map offsets) {
+    public TxnOffsetCommitRequest(TxnOffsetCommitRequestData data, short version) {
         super(ApiKeys.TXN_OFFSET_COMMIT, version);
-        this.transactionalId = transactionalId;
-        this.consumerGroupId = consumerGroupId;
-        this.producerId = producerId;
-        this.producerEpoch = producerEpoch;
-        this.offsets = offsets;
+        this.data = data;
     }
 
     public TxnOffsetCommitRequest(Struct struct, short version) {
         super(ApiKeys.TXN_OFFSET_COMMIT, version);
-        this.transactionalId = struct.get(TRANSACTIONAL_ID);
-        this.consumerGroupId = struct.get(GROUP_ID);
-        this.producerId = struct.get(PRODUCER_ID);
-        this.producerEpoch = struct.get(PRODUCER_EPOCH);
+        this.data = new TxnOffsetCommitRequestData(struct, version);
+    }
 
-        Map offsets = new HashMap<>();
-        Object[] topicPartitionsArray = struct.get(TOPICS);
-        for (Object topicPartitionObj : topicPartitionsArray) {
-            Struct topicPartitionStruct = (Struct) topicPartitionObj;
-            String topic = topicPartitionStruct.get(TOPIC_NAME);
-            for (Object partitionObj : topicPartitionStruct.get(PARTITIONS)) {
-                Struct partitionStruct = (Struct) partitionObj;
-                TopicPartition partition = new TopicPartition(topic, partitionStruct.get(PARTITION_ID));
-                long offset = partitionStruct.get(COMMITTED_OFFSET);
-                String metadata = partitionStruct.get(COMMITTED_METADATA);
-                Optional leaderEpoch = RequestUtils.getLeaderEpoch(partitionStruct, COMMITTED_LEADER_EPOCH);
-                offsets.put(partition, new CommittedOffset(offset, metadata, leaderEpoch));
+    public Map offsets() {
+        List topics = data.topics();
+        Map offsetMap = new HashMap<>();
+        for (TxnOffsetCommitRequestTopic topic : topics) {
+            for (TxnOffsetCommitRequestPartition partition : topic.partitions()) {
+                offsetMap.put(new TopicPartition(topic.name(), partition.partitionIndex()),
+                              new CommittedOffset(partition.committedOffset(),
+                                                  partition.committedMetadata(),
+                                                  RequestUtils.getLeaderEpoch(partition.committedLeaderEpoch()))
+                );
             }
         }
-        this.offsets = offsets;
+        return offsetMap;
     }
 
-    public String transactionalId() {
-        return transactionalId;
-    }
-
-    public String consumerGroupId() {
-        return consumerGroupId;
-    }
-
-    public long producerId() {
-        return producerId;
-    }
-
-    public short producerEpoch() {
-        return producerEpoch;
-    }
-
-    public Map offsets() {
-        return offsets;
+    public static List getTopics(Map pendingTxnOffsetCommits) {
+        Map> topicPartitionMap = new HashMap<>();
+        for (Map.Entry entry : pendingTxnOffsetCommits.entrySet()) {
+            TopicPartition topicPartition = entry.getKey();
+            CommittedOffset offset = entry.getValue();
+
+            List partitions =
+                topicPartitionMap.getOrDefault(topicPartition.topic(), new ArrayList<>());
+            partitions.add(new TxnOffsetCommitRequestPartition()
+                               .setPartitionIndex(topicPartition.partition())
+                               .setCommittedOffset(offset.offset)
+                               .setCommittedLeaderEpoch(offset.leaderEpoch.orElse(RecordBatch.NO_PARTITION_LEADER_EPOCH))
+                               .setCommittedMetadata(offset.metadata)
+            );
+            topicPartitionMap.put(topicPartition.topic(), partitions);
+        }
+        return topicPartitionMap.entrySet().stream()
+                   .map(entry -> new TxnOffsetCommitRequestTopic()
+                                     .setName(entry.getKey())
+                                     .setPartitions(entry.getValue()))
+                   .collect(Collectors.toList());
     }
 
     @Override
     protected Struct toStruct() {
-        Struct struct = new Struct(ApiKeys.TXN_OFFSET_COMMIT.requestSchema(version()));
-        struct.set(TRANSACTIONAL_ID, transactionalId);
-        struct.set(GROUP_ID, consumerGroupId);
-        struct.set(PRODUCER_ID, producerId);
-        struct.set(PRODUCER_EPOCH, producerEpoch);
-
-        Map> mappedPartitionOffsets = CollectionUtils.groupPartitionDataByTopic(offsets);
-        Object[] partitionsArray = new Object[mappedPartitionOffsets.size()];
-        int i = 0;
-        for (Map.Entry> topicAndPartitions : mappedPartitionOffsets.entrySet()) {
-            Struct topicPartitionsStruct = struct.instance(TOPICS);
-            topicPartitionsStruct.set(TOPIC_NAME, topicAndPartitions.getKey());
+        return data.toStruct(version());
+    }
 
-            Map partitionOffsets = topicAndPartitions.getValue();
-            Object[] partitionOffsetsArray = new Object[partitionOffsets.size()];
-            int j = 0;
-            for (Map.Entry partitionOffset : partitionOffsets.entrySet()) {
-                Struct partitionOffsetStruct = topicPartitionsStruct.instance(PARTITIONS);
-                partitionOffsetStruct.set(PARTITION_ID, partitionOffset.getKey());
-                CommittedOffset committedOffset = partitionOffset.getValue();
-                partitionOffsetStruct.set(COMMITTED_OFFSET, committedOffset.offset);
-                partitionOffsetStruct.set(COMMITTED_METADATA, committedOffset.metadata);
-                RequestUtils.setLeaderEpochIfExists(partitionOffsetStruct, COMMITTED_LEADER_EPOCH,
-                        committedOffset.leaderEpoch);
-                partitionOffsetsArray[j++] = partitionOffsetStruct;
+    static List getErrorResponseTopics(List requestTopics,
+                                                                     Errors e) {
+        List responseTopicData = new ArrayList<>();
+        for (TxnOffsetCommitRequestTopic entry : requestTopics) {
+            List responsePartitions = new ArrayList<>();
+            for (TxnOffsetCommitRequestPartition requestPartition : entry.partitions()) {
+                responsePartitions.add(new TxnOffsetCommitResponsePartition()
+                                           .setPartitionIndex(requestPartition.partitionIndex())
+                                           .setErrorCode(e.code()));
             }
-            topicPartitionsStruct.set(PARTITIONS, partitionOffsetsArray);
-            partitionsArray[i++] = topicPartitionsStruct;
+            responseTopicData.add(new TxnOffsetCommitResponseTopic()
+                                      .setName(entry.name())
+                                      .setPartitions(responsePartitions)
+            );
         }
-
-        struct.set(TOPICS, partitionsArray);
-        return struct;
+        return responseTopicData;
     }
 
     @Override
     public TxnOffsetCommitResponse getErrorResponse(int throttleTimeMs, Throwable e) {
-        Errors error = Errors.forException(e);
-        Map errors = new HashMap<>(offsets.size());
-        for (TopicPartition partition : offsets.keySet())
-            errors.put(partition, error);
-        return new TxnOffsetCommitResponse(throttleTimeMs, errors);
+        List responseTopicData =
+            getErrorResponseTopics(data.topics(), Errors.forException(e));
+
+        return new TxnOffsetCommitResponse(new TxnOffsetCommitResponseData()
+                                               .setThrottleTimeMs(throttleTimeMs)
+                                               .setTopics(responseTopicData));
     }
 
     public static TxnOffsetCommitRequest parse(ByteBuffer buffer, short version) {
@@ -260,6 +164,22 @@ public String toString() {
                     ", leaderEpoch=" + leaderEpoch +
                     ", metadata='" + metadata + "')";
         }
-    }
 
+        @Override
+        public boolean equals(Object other) {
+            if (!(other instanceof CommittedOffset)) {
+                return false;
+            }
+            CommittedOffset otherOffset = (CommittedOffset) other;
+
+            return this.offset == otherOffset.offset
+                       && this.leaderEpoch.equals(otherOffset.leaderEpoch)
+                       && Objects.equals(this.metadata, otherOffset.metadata);
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(offset, leaderEpoch, metadata);
+        }
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/TxnOffsetCommitResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/TxnOffsetCommitResponse.java
index ba8f7d6be2eaf..256b9983671e0 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/TxnOffsetCommitResponse.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/TxnOffsetCommitResponse.java
@@ -17,142 +17,98 @@
 package org.apache.kafka.common.requests;
 
 import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.message.TxnOffsetCommitResponseData;
+import org.apache.kafka.common.message.TxnOffsetCommitResponseData.TxnOffsetCommitResponsePartition;
+import org.apache.kafka.common.message.TxnOffsetCommitResponseData.TxnOffsetCommitResponseTopic;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.protocol.types.Field;
-import org.apache.kafka.common.protocol.types.Schema;
 import org.apache.kafka.common.protocol.types.Struct;
-import org.apache.kafka.common.utils.CollectionUtils;
 
 import java.nio.ByteBuffer;
+import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.Map;
 
-import static org.apache.kafka.common.protocol.CommonFields.ERROR_CODE;
-import static org.apache.kafka.common.protocol.CommonFields.PARTITION_ID;
-import static org.apache.kafka.common.protocol.CommonFields.THROTTLE_TIME_MS;
-import static org.apache.kafka.common.protocol.CommonFields.TOPIC_NAME;
-
 /**
- *
  * Possible error codes:
- *   InvalidProducerEpoch
- *   NotCoordinator
- *   CoordinatorNotAvailable
- *   CoordinatorLoadInProgress
- *   OffsetMetadataTooLarge
- *   GroupAuthorizationFailed
- *   InvalidCommitOffsetSize
- *   TransactionalIdAuthorizationFailed
- *   RequestTimedOut
+ *
+ *   - {@link Errors#INVALID_PRODUCER_EPOCH}
+ *   - {@link Errors#NOT_COORDINATOR}
+ *   - {@link Errors#COORDINATOR_NOT_AVAILABLE}
+ *   - {@link Errors#COORDINATOR_LOAD_IN_PROGRESS}
+ *   - {@link Errors#OFFSET_METADATA_TOO_LARGE}
+ *   - {@link Errors#GROUP_AUTHORIZATION_FAILED}
+ *   - {@link Errors#INVALID_COMMIT_OFFSET_SIZE}
+ *   - {@link Errors#TRANSACTIONAL_ID_AUTHORIZATION_FAILED}
+ *   - {@link Errors#REQUEST_TIMED_OUT}
  */
 public class TxnOffsetCommitResponse extends AbstractResponse {
-    private static final Field.ComplexArray TOPICS = new Field.ComplexArray("topics",
-            "Responses by topic for committed offsets");
-
-    // topic level fields
-    private static final Field.ComplexArray PARTITIONS = new Field.ComplexArray("partitions",
-            "Responses by partition for committed offsets");
-
-    private static final Field PARTITIONS_V0 = PARTITIONS.withFields(
-            PARTITION_ID,
-            ERROR_CODE);
-
-    private static final Field TOPICS_V0 = TOPICS.withFields(
-            TOPIC_NAME,
-            PARTITIONS_V0);
 
-    private static final Schema TXN_OFFSET_COMMIT_RESPONSE_V0 = new Schema(
-            THROTTLE_TIME_MS,
-            TOPICS_V0);
+    public final TxnOffsetCommitResponseData data;
 
-    // V1 bump used to indicate that on quota violation brokers send out responses before throttling.
-    private static final Schema TXN_OFFSET_COMMIT_RESPONSE_V1 = TXN_OFFSET_COMMIT_RESPONSE_V0;
-
-    // V2 adds the leader epoch to the partition data
-    private static final Schema TXN_OFFSET_COMMIT_RESPONSE_V2 = TXN_OFFSET_COMMIT_RESPONSE_V1;
+    public TxnOffsetCommitResponse(TxnOffsetCommitResponseData data) {
+        this.data = data;
+    }
 
-    public static Schema[] schemaVersions() {
-        return new Schema[]{TXN_OFFSET_COMMIT_RESPONSE_V0, TXN_OFFSET_COMMIT_RESPONSE_V1, TXN_OFFSET_COMMIT_RESPONSE_V2};
+    public TxnOffsetCommitResponse(Struct struct, short version) {
+        this.data = new TxnOffsetCommitResponseData(struct, version);
     }
 
-    private final Map errors;
-    private final int throttleTimeMs;
+    public TxnOffsetCommitResponse(int requestThrottleMs, Map responseData) {
+        Map responseTopicDataMap = new HashMap<>();
 
-    public TxnOffsetCommitResponse(int throttleTimeMs, Map errors) {
-        this.throttleTimeMs = throttleTimeMs;
-        this.errors = errors;
-    }
+        for (Map.Entry entry : responseData.entrySet()) {
+            TopicPartition topicPartition = entry.getKey();
+            String topicName = topicPartition.topic();
 
-    public TxnOffsetCommitResponse(Struct struct) {
-        this.throttleTimeMs = struct.get(THROTTLE_TIME_MS);
-        Map errors = new HashMap<>();
-        Object[] topicPartitionsArray = struct.get(TOPICS);
-        for (Object topicPartitionObj : topicPartitionsArray) {
-            Struct topicPartitionStruct = (Struct) topicPartitionObj;
-            String topic = topicPartitionStruct.get(TOPIC_NAME);
-            for (Object partitionObj : topicPartitionStruct.get(PARTITIONS)) {
-                Struct partitionStruct = (Struct) partitionObj;
-                Integer partition = partitionStruct.get(PARTITION_ID);
-                Errors error = Errors.forCode(partitionStruct.get(ERROR_CODE));
-                errors.put(new TopicPartition(topic, partition), error);
-            }
+            TxnOffsetCommitResponseTopic topic = responseTopicDataMap.getOrDefault(
+                topicName, new TxnOffsetCommitResponseTopic().setName(topicName));
+
+            topic.partitions().add(new TxnOffsetCommitResponsePartition()
+                                       .setErrorCode(entry.getValue().code())
+                                       .setPartitionIndex(topicPartition.partition())
+            );
+            responseTopicDataMap.put(topicName, topic);
         }
-        this.errors = errors;
+
+        data = new TxnOffsetCommitResponseData()
+                   .setTopics(new ArrayList<>(responseTopicDataMap.values()))
+                   .setThrottleTimeMs(requestThrottleMs);
     }
 
     @Override
     protected Struct toStruct(short version) {
-        Struct struct = new Struct(ApiKeys.TXN_OFFSET_COMMIT.responseSchema(version));
-        struct.set(THROTTLE_TIME_MS, throttleTimeMs);
-        Map> mappedPartitions = CollectionUtils.groupPartitionDataByTopic(errors);
-        Object[] partitionsArray = new Object[mappedPartitions.size()];
-        int i = 0;
-        for (Map.Entry> topicAndPartitions : mappedPartitions.entrySet()) {
-            Struct topicPartitionsStruct = struct.instance(TOPICS);
-            topicPartitionsStruct.set(TOPIC_NAME, topicAndPartitions.getKey());
-            Map partitionAndErrors = topicAndPartitions.getValue();
-
-            Object[] partitionAndErrorsArray = new Object[partitionAndErrors.size()];
-            int j = 0;
-            for (Map.Entry partitionAndError : partitionAndErrors.entrySet()) {
-                Struct partitionAndErrorStruct = topicPartitionsStruct.instance(PARTITIONS);
-                partitionAndErrorStruct.set(PARTITION_ID, partitionAndError.getKey());
-                partitionAndErrorStruct.set(ERROR_CODE, partitionAndError.getValue().code());
-                partitionAndErrorsArray[j++] = partitionAndErrorStruct;
-            }
-            topicPartitionsStruct.set(PARTITIONS, partitionAndErrorsArray);
-            partitionsArray[i++] = topicPartitionsStruct;
-        }
-
-        struct.set(TOPICS, partitionsArray);
-        return struct;
+        return data.toStruct(version);
     }
 
     @Override
     public int throttleTimeMs() {
-        return throttleTimeMs;
-    }
-
-    public Map errors() {
-        return errors;
+        return data.throttleTimeMs();
     }
 
     @Override
     public Map errorCounts() {
-        return errorCounts(errors);
+        return errorCounts(errors().values());
+    }
+
+    public Map errors() {
+        Map errorMap = new HashMap<>();
+        for (TxnOffsetCommitResponseTopic topic : data.topics()) {
+            for (TxnOffsetCommitResponsePartition partition : topic.partitions()) {
+                errorMap.put(new TopicPartition(topic.name(), partition.partitionIndex()),
+                             Errors.forCode(partition.errorCode()));
+            }
+        }
+        return errorMap;
     }
 
     public static TxnOffsetCommitResponse parse(ByteBuffer buffer, short version) {
-        return new TxnOffsetCommitResponse(ApiKeys.TXN_OFFSET_COMMIT.parseResponse(version, buffer));
+        return new TxnOffsetCommitResponse(ApiKeys.TXN_OFFSET_COMMIT.parseResponse(version, buffer), version);
     }
 
     @Override
     public String toString() {
-        return "TxnOffsetCommitResponse(" +
-                "errors=" + errors +
-                ", throttleTimeMs=" + throttleTimeMs +
-                ')';
+        return data.toString();
     }
 
     @Override
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/UpdateMetadataRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/UpdateMetadataRequest.java
index b715c2aa65b6a..6ab4d25f61f90 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/UpdateMetadataRequest.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/UpdateMetadataRequest.java
@@ -16,481 +16,271 @@
  */
 package org.apache.kafka.common.requests;
 
-import org.apache.kafka.common.TopicPartition;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
 import org.apache.kafka.common.errors.UnsupportedVersionException;
+import org.apache.kafka.common.message.UpdateMetadataRequestData;
+import org.apache.kafka.common.message.UpdateMetadataRequestData.UpdateMetadataBroker;
+import org.apache.kafka.common.message.UpdateMetadataRequestData.UpdateMetadataEndpoint;
+import org.apache.kafka.common.message.UpdateMetadataRequestData.UpdateMetadataPartitionState;
+import org.apache.kafka.common.message.UpdateMetadataRequestData.UpdateMetadataTopicState;
+import org.apache.kafka.common.message.UpdateMetadataResponseData;
 import org.apache.kafka.common.network.ListenerName;
+import org.apache.kafka.common.network.NetworkSend;
+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.protocol.types.Field;
-import org.apache.kafka.common.protocol.types.Schema;
 import org.apache.kafka.common.protocol.types.Struct;
 import org.apache.kafka.common.security.auth.SecurityProtocol;
-import org.apache.kafka.common.utils.CollectionUtils;
+import org.apache.kafka.common.utils.FlattenedIterator;
 import org.apache.kafka.common.utils.Utils;
 
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
-import java.util.Set;
 
-import static org.apache.kafka.common.protocol.CommonFields.PARTITION_ID;
-import static org.apache.kafka.common.protocol.CommonFields.TOPIC_NAME;
-import static org.apache.kafka.common.protocol.types.Type.INT32;
+import static java.util.Collections.singletonList;
 
 public class UpdateMetadataRequest extends AbstractControlRequest {
-    private static final Field.ComplexArray TOPIC_STATES = new Field.ComplexArray("topic_states", "Topic states");
-    private static final Field.ComplexArray PARTITION_STATES = new Field.ComplexArray("partition_states", "Partition states");
-    private static final Field.ComplexArray LIVE_BROKERS = new Field.ComplexArray("live_brokers", "Live broekrs");
-
-    // PartitionState fields
-    private static final Field.Int32 LEADER = new Field.Int32("leader", "The broker id for the leader.");
-    private static final Field.Int32 LEADER_EPOCH = new Field.Int32("leader_epoch", "The leader epoch.");
-    private static final Field.Array ISR = new Field.Array("isr", INT32, "The in sync replica ids.");
-    private static final Field.Int32 ZK_VERSION = new Field.Int32("zk_version", "The ZK version.");
-    private static final Field.Array REPLICAS = new Field.Array("replicas", INT32, "The replica ids.");
-    private static final Field.Array OFFLINE_REPLICAS = new Field.Array("offline_replicas", INT32, "The offline replica ids");
-
-    // Live brokers fields
-    private static final Field.Int32 BROKER_ID = new Field.Int32("id", "The broker id");
-    private static final Field.ComplexArray ENDPOINTS = new Field.ComplexArray("end_points", "The endpoints");
-    private static final Field.NullableStr RACK = new Field.NullableStr("rack", "The rack");
-
-    // EndPoint fields
-    private static final Field.Str HOST = new Field.Str("host", "The hostname of the broker.");
-    private static final Field.Int32  PORT = new Field.Int32("port", "The port on which the broker accepts requests.");
-    private static final Field.Str LISTENER_NAME = new Field.Str("listener_name", "The listener name.");
-    private static final Field.Int16  SECURITY_PROTOCOL_TYPE = new Field.Int16("security_protocol_type", "The security protocol type.");
-
-    private static final Field PARTITION_STATES_V0 = PARTITION_STATES.withFields(
-            TOPIC_NAME,
-            PARTITION_ID,
-            CONTROLLER_EPOCH,
-            LEADER,
-            LEADER_EPOCH,
-            ISR,
-            ZK_VERSION,
-            REPLICAS);
-
-    // PARTITION_STATES_V4 added a per-partition offline_replicas field. This field specifies
-    // the list of replicas that are offline.
-    private static final Field PARTITION_STATES_V4 = PARTITION_STATES.withFields(
-            TOPIC_NAME,
-            PARTITION_ID,
-            CONTROLLER_EPOCH,
-            LEADER,
-            LEADER_EPOCH,
-            ISR,
-            ZK_VERSION,
-            REPLICAS,
-            OFFLINE_REPLICAS);
-
-    private static final Field PARTITION_STATES_V5 = PARTITION_STATES.withFields(
-            PARTITION_ID,
-            CONTROLLER_EPOCH,
-            LEADER,
-            LEADER_EPOCH,
-            ISR,
-            ZK_VERSION,
-            REPLICAS,
-            OFFLINE_REPLICAS);
-
-    // TOPIC_STATES_V5 normalizes TOPIC_STATES_V4 to
-    // make it more memory efficient
-    private static final Field TOPIC_STATES_V5 = TOPIC_STATES.withFields(
-            TOPIC_NAME,
-            PARTITION_STATES_V5);
-
-    // for some reason, V1 sends `port` before `host` while V0 sends `host` before `port
-    private static final Field ENDPOINTS_V1 = ENDPOINTS.withFields(
-            PORT,
-            HOST,
-            SECURITY_PROTOCOL_TYPE);
-
-    private static final Field ENDPOINTS_V3 = ENDPOINTS.withFields(
-            PORT,
-            HOST,
-            LISTENER_NAME,
-            SECURITY_PROTOCOL_TYPE);
-
-    private static final Field LIVE_BROKERS_V0 = LIVE_BROKERS.withFields(
-            BROKER_ID,
-            HOST,
-            PORT);
-
-    private static final Field LIVE_BROKERS_V1 = LIVE_BROKERS.withFields(
-            BROKER_ID,
-            ENDPOINTS_V1);
-
-    private static final Field LIVE_BROKERS_V2 = LIVE_BROKERS.withFields(
-            BROKER_ID,
-            ENDPOINTS_V1,
-            RACK);
-
-    private static final Field LIVE_BROKERS_V3 = LIVE_BROKERS.withFields(
-            BROKER_ID,
-            ENDPOINTS_V3,
-            RACK);
-
-    private static final Schema UPDATE_METADATA_REQUEST_V0 = new Schema(
-            CONTROLLER_ID,
-            CONTROLLER_EPOCH,
-            PARTITION_STATES_V0,
-            LIVE_BROKERS_V0);
-
-    private static final Schema UPDATE_METADATA_REQUEST_V1 = new Schema(
-            CONTROLLER_ID,
-            CONTROLLER_EPOCH,
-            PARTITION_STATES_V0,
-            LIVE_BROKERS_V1);
-
-    private static final Schema UPDATE_METADATA_REQUEST_V2 = new Schema(
-            CONTROLLER_ID,
-            CONTROLLER_EPOCH,
-            PARTITION_STATES_V0,
-            LIVE_BROKERS_V2);
-
-
-    private static final Schema UPDATE_METADATA_REQUEST_V3 = new Schema(
-            CONTROLLER_ID,
-            CONTROLLER_EPOCH,
-            PARTITION_STATES_V0,
-            LIVE_BROKERS_V3);
-
-    // UPDATE_METADATA_REQUEST_V4 added a per-partition offline_replicas field. This field specifies the list of replicas that are offline.
-    private static final Schema UPDATE_METADATA_REQUEST_V4 = new Schema(
-            CONTROLLER_ID,
-            CONTROLLER_EPOCH,
-            PARTITION_STATES_V4,
-            LIVE_BROKERS_V3);
-
-    // UPDATE_METADATA_REQUEST_V5 added a broker_epoch Field. This field specifies the generation of the broker across
-    // bounces. It also normalizes partitions under each topic.
-    private static final Schema UPDATE_METADATA_REQUEST_V5 = new Schema(
-            CONTROLLER_ID,
-            CONTROLLER_EPOCH,
-            BROKER_EPOCH,
-            TOPIC_STATES_V5,
-            LIVE_BROKERS_V3);
-
-    public static Schema[] schemaVersions() {
-        return new Schema[] {UPDATE_METADATA_REQUEST_V0, UPDATE_METADATA_REQUEST_V1, UPDATE_METADATA_REQUEST_V2,
-            UPDATE_METADATA_REQUEST_V3, UPDATE_METADATA_REQUEST_V4, UPDATE_METADATA_REQUEST_V5};
-    }
+    private final ReentrantLock bodyBufferLock = new ReentrantLock();
+    private byte[] bodyBuffer;
 
     public static class Builder extends AbstractControlRequest.Builder {
-        private final Map partitionStates;
-        private final Set liveBrokers;
+        private final List partitionStates;
+        private final List liveBrokers;
+        private Lock buildLock = new ReentrantLock();
 
-        public Builder(short version, int controllerId, int controllerEpoch, long brokerEpoch,
-                       Map partitionStates, Set liveBrokers) {
-            super(ApiKeys.UPDATE_METADATA, version, controllerId, controllerEpoch, brokerEpoch);
+        // LIKAFKA-18349 - Cache the UpdateMetadataRequest Objects to reduce memory usage
+        private final Map requestCache = new HashMap<>();
+
+        public Builder(short version, int controllerId, int controllerEpoch, long brokerEpoch, long maxBrokerEpoch,
+                       List partitionStates, List liveBrokers) {
+            super(ApiKeys.UPDATE_METADATA, version, controllerId, controllerEpoch, brokerEpoch, maxBrokerEpoch);
             this.partitionStates = partitionStates;
             this.liveBrokers = liveBrokers;
         }
 
         @Override
         public UpdateMetadataRequest build(short version) {
-            if (version == 0) {
-                for (Broker broker : liveBrokers) {
-                    if (broker.endPoints.size() != 1 || broker.endPoints.get(0).securityProtocol != SecurityProtocol.PLAINTEXT) {
-                        throw new UnsupportedVersionException("UpdateMetadataRequest v0 only handles PLAINTEXT endpoints");
+            // the following inner blocks are not indented in order to make hotfix cherry-picking easier
+            buildLock.lock();
+            try {
+            UpdateMetadataRequest updateMetadataRequest = requestCache.get(version);
+            if (updateMetadataRequest == null) {
+            if (version < 3) {
+                for (UpdateMetadataBroker broker : liveBrokers) {
+                    if (version == 0) {
+                        if (broker.endpoints().size() != 1)
+                            throw new UnsupportedVersionException("UpdateMetadataRequest v0 requires a single endpoint");
+                        if (broker.endpoints().get(0).securityProtocol() != SecurityProtocol.PLAINTEXT.id)
+                            throw new UnsupportedVersionException("UpdateMetadataRequest v0 only handles PLAINTEXT endpoints");
+                        // Don't null out `endpoints` since it's ignored by the generated code if version >= 1
+                        UpdateMetadataEndpoint endpoint = broker.endpoints().get(0);
+                        broker.setV0Host(endpoint.host());
+                        broker.setV0Port(endpoint.port());
+                    } else {
+                        if (broker.endpoints().stream().anyMatch(endpoint -> !endpoint.listener().isEmpty() &&
+                                !endpoint.listener().equals(listenerNameFromSecurityProtocol(endpoint)))) {
+                            throw new UnsupportedVersionException("UpdateMetadataRequest v0-v3 does not support custom " +
+                                "listeners, request version: " + version + ", endpoints: " + broker.endpoints());
+                        }
                     }
                 }
             }
-            return new UpdateMetadataRequest(version, controllerId, controllerEpoch, brokerEpoch, partitionStates, liveBrokers);
+
+            UpdateMetadataRequestData data = new UpdateMetadataRequestData()
+                    .setControllerId(controllerId)
+                    .setControllerEpoch(controllerEpoch)
+                    .setBrokerEpoch(brokerEpoch)
+                    .setMaxBrokerEpoch(maxBrokerEpoch)
+                    .setLiveBrokers(liveBrokers);
+
+            if (version >= 5) {
+                Map topicStatesMap = groupByTopic(partitionStates);
+                data.setTopicStates(new ArrayList<>(topicStatesMap.values()));
+            } else {
+                data.setUngroupedPartitionStates(partitionStates);
+            }
+
+            updateMetadataRequest = new UpdateMetadataRequest(data, version);
+            requestCache.put(version, updateMetadataRequest);
+            }
+            return updateMetadataRequest;
+            } finally {
+                buildLock.unlock();
+            }
+        }
+
+        private static Map groupByTopic(List partitionStates) {
+            Map topicStates = new HashMap<>();
+            for (UpdateMetadataPartitionState partition : partitionStates) {
+                // We don't null out the topic name in UpdateMetadataTopicState since it's ignored by the generated
+                // code if version >= 5
+                UpdateMetadataTopicState topicState = topicStates.computeIfAbsent(partition.topicName(),
+                    t -> new UpdateMetadataTopicState().setTopicName(partition.topicName()));
+                topicState.partitionStates().add(partition);
+            }
+            return topicStates;
         }
 
         @Override
         public String toString() {
             StringBuilder bld = new StringBuilder();
+            // HOTFIX: LIKAFKA-24478
+            // large cluster with large metadata can create really large string
+            // potentially causing OOM
             bld.append("(type: UpdateMetadataRequest=").
                 append(", controllerId=").append(controllerId).
                 append(", controllerEpoch=").append(controllerEpoch).
                 append(", brokerEpoch=").append(brokerEpoch).
-                append(", partitionStates=").append(partitionStates).
+                append(", maxBrokerEpoch=").append(maxBrokerEpoch).
                 append(", liveBrokers=").append(Utils.join(liveBrokers, ", ")).
                 append(")");
+
+            // bld.append("(type: UpdateMetadataRequest=").
+            //   append(", controllerId=").append(controllerId).
+            //   append(", controllerEpoch=").append(controllerEpoch).
+            //   append(", brokerEpoch=").append(brokerEpoch).
+            //   append(", partitionStates=").append(partitionStates).
+            //   append(", liveBrokers=").append(Utils.join(liveBrokers, ", ")).
+            //   append(")");
             return bld.toString();
         }
     }
 
-    public static final class PartitionState {
-        public final BasePartitionState basePartitionState;
-        public final List offlineReplicas;
-
-        public PartitionState(int controllerEpoch,
-                              int leader,
-                              int leaderEpoch,
-                              List isr,
-                              int zkVersion,
-                              List replicas,
-                              List offlineReplicas) {
-            this.basePartitionState = new BasePartitionState(controllerEpoch, leader, leaderEpoch, isr, zkVersion, replicas);
-            this.offlineReplicas = offlineReplicas;
-        }
-
-        private PartitionState(Struct struct) {
-            int controllerEpoch = struct.get(CONTROLLER_EPOCH);
-            int leader = struct.get(LEADER);
-            int leaderEpoch = struct.get(LEADER_EPOCH);
-
-            Object[] isrArray = struct.get(ISR);
-            List isr = new ArrayList<>(isrArray.length);
-            for (Object r : isrArray)
-                isr.add((Integer) r);
+    private final UpdateMetadataRequestData data;
+    // LIKAFKA-18349 - Cache the UpdateMetadataRequest struct to reduce memory usage
+    private Struct struct = null;
+    private Lock structLock = new ReentrantLock();
 
-            int zkVersion = struct.get(ZK_VERSION);
-
-            Object[] replicasArray = struct.get(REPLICAS);
-            List replicas = new ArrayList<>(replicasArray.length);
-            for (Object r : replicasArray)
-                replicas.add((Integer) r);
-
-            this.basePartitionState = new BasePartitionState(controllerEpoch, leader, leaderEpoch, isr, zkVersion, replicas);
+    UpdateMetadataRequest(UpdateMetadataRequestData data, short version) {
+        super(ApiKeys.UPDATE_METADATA, version);
+        this.data = data;
+        // Do this from the constructor to make it thread-safe (even though it's only needed when some methods are called)
+        normalize();
+    }
 
-            this.offlineReplicas = new ArrayList<>();
-            if (struct.hasField(OFFLINE_REPLICAS)) {
-                Object[] offlineReplicasArray = struct.get(OFFLINE_REPLICAS);
-                for (Object r : offlineReplicasArray)
-                    offlineReplicas.add((Integer) r);
+    private void normalize() {
+        // Version 0 only supported a single host and port and the protocol was always plaintext
+        // Version 1 added support for multiple endpoints, each with its own security protocol
+        // Version 2 added support for rack
+        // Version 3 added support for listener name, which we can infer from the security protocol for older versions
+        if (version() < 3) {
+            for (UpdateMetadataBroker liveBroker : data.liveBrokers()) {
+                // Set endpoints so that callers can rely on it always being present
+                if (version() == 0 && liveBroker.endpoints().isEmpty()) {
+                    SecurityProtocol securityProtocol = SecurityProtocol.PLAINTEXT;
+                    liveBroker.setEndpoints(singletonList(new UpdateMetadataEndpoint()
+                        .setHost(liveBroker.v0Host())
+                        .setPort(liveBroker.v0Port())
+                        .setSecurityProtocol(securityProtocol.id)
+                        .setListener(ListenerName.forSecurityProtocol(securityProtocol).value())));
+                } else {
+                    for (UpdateMetadataEndpoint endpoint : liveBroker.endpoints()) {
+                        // Set listener so that callers can rely on it always being present
+                        if (endpoint.listener().isEmpty())
+                            endpoint.setListener(listenerNameFromSecurityProtocol(endpoint));
+                    }
+                }
             }
         }
 
-        @Override
-        public String toString() {
-            return "PartitionState(controllerEpoch=" + basePartitionState.controllerEpoch +
-                ", leader=" + basePartitionState.leader +
-                ", leaderEpoch=" + basePartitionState.leaderEpoch +
-                ", isr=" + Arrays.toString(basePartitionState.isr.toArray()) +
-                ", zkVersion=" + basePartitionState.zkVersion +
-                ", replicas=" + Arrays.toString(basePartitionState.replicas.toArray()) +
-                ", offlineReplicas=" + Arrays.toString(offlineReplicas.toArray()) + ")";
-        }
-
-        private void setStruct(Struct struct) {
-            struct.set(CONTROLLER_EPOCH, basePartitionState.controllerEpoch);
-            struct.set(LEADER, basePartitionState.leader);
-            struct.set(LEADER_EPOCH, basePartitionState.leaderEpoch);
-            struct.set(ISR, basePartitionState.isr.toArray());
-            struct.set(ZK_VERSION, basePartitionState.zkVersion);
-            struct.set(REPLICAS, basePartitionState.replicas.toArray());
-            struct.setIfExists(OFFLINE_REPLICAS, offlineReplicas.toArray());
+        if (version() >= 5) {
+            for (UpdateMetadataTopicState topicState : data.topicStates()) {
+                for (UpdateMetadataPartitionState partitionState : topicState.partitionStates()) {
+                    // Set the topic name so that we can always present the ungrouped view to callers
+                    partitionState.setTopicName(topicState.topicName());
+                }
+            }
         }
     }
 
-    public static final class Broker {
-        public final int id;
-        public final List endPoints;
-        public final String rack; // introduced in V2
-
-        public Broker(int id, List endPoints, String rack) {
-            this.id = id;
-            this.endPoints = endPoints;
-            this.rack = rack;
-        }
+    private static String listenerNameFromSecurityProtocol(UpdateMetadataEndpoint endpoint) {
+        SecurityProtocol securityProtocol = SecurityProtocol.forId(endpoint.securityProtocol());
+        return ListenerName.forSecurityProtocol(securityProtocol).value();
+    }
 
-        @Override
-        public String toString() {
-            StringBuilder bld = new StringBuilder();
-            bld.append("(id=").append(id);
-            bld.append(", endPoints=").append(Utils.join(endPoints, ","));
-            bld.append(", rack=").append(rack);
-            bld.append(")");
-            return bld.toString();
-        }
+    public UpdateMetadataRequest(Struct struct, short version) {
+        this(new UpdateMetadataRequestData(struct, version), version);
     }
 
-    public static final class EndPoint {
-        public final String host;
-        public final int port;
-        public final SecurityProtocol securityProtocol;
-        public final ListenerName listenerName; // introduced in V3
+    @Override
+    public int controllerId() {
+        return data.controllerId();
+    }
 
-        public EndPoint(String host, int port, SecurityProtocol securityProtocol, ListenerName listenerName) {
-            this.host = host;
-            this.port = port;
-            this.securityProtocol = securityProtocol;
-            this.listenerName = listenerName;
-        }
+    @Override
+    public int controllerEpoch() {
+        return data.controllerEpoch();
+    }
 
-        @Override
-        public String toString() {
-            return "(host=" + host + ", port=" + port + ", listenerName=" + listenerName +
-                    ", securityProtocol=" + securityProtocol + ")";
-        }
+    @Override
+    public long brokerEpoch() {
+        return data.brokerEpoch();
     }
 
-    private final Map partitionStates;
-    private final Set liveBrokers;
+    @Override
+    public long maxBrokerEpoch() {
+        return data.maxBrokerEpoch();
+    }
 
-    private UpdateMetadataRequest(short version, int controllerId, int controllerEpoch, long brokerEpoch,
-                                  Map partitionStates, Set liveBrokers) {
-        super(ApiKeys.UPDATE_METADATA, version, controllerId, controllerEpoch, brokerEpoch);
-        this.partitionStates = partitionStates;
-        this.liveBrokers = liveBrokers;
+    @Override
+    public UpdateMetadataResponse getErrorResponse(int throttleTimeMs, Throwable e) {
+        UpdateMetadataResponseData data = new UpdateMetadataResponseData()
+                .setErrorCode(Errors.forException(e).code());
+        return new UpdateMetadataResponse(data);
     }
 
-    public UpdateMetadataRequest(Struct struct, short versionId) {
-        super(ApiKeys.UPDATE_METADATA, struct, versionId);
-        Map partitionStates = new HashMap<>();
-        if (struct.hasField(TOPIC_STATES)) {
-            for (Object topicStatesDataObj : struct.get(TOPIC_STATES)) {
-                Struct topicStatesData = (Struct) topicStatesDataObj;
-                String topic = topicStatesData.get(TOPIC_NAME);
-                for (Object partitionStateDataObj : topicStatesData.get(PARTITION_STATES)) {
-                    Struct partitionStateData = (Struct) partitionStateDataObj;
-                    int partition = partitionStateData.get(PARTITION_ID);
-                    PartitionState partitionState = new PartitionState(partitionStateData);
-                    partitionStates.put(new TopicPartition(topic, partition), partitionState);
-                }
-            }
-        } else {
-            for (Object partitionStateDataObj : struct.get(PARTITION_STATES)) {
-                Struct partitionStateData = (Struct) partitionStateDataObj;
-                String topic = partitionStateData.get(TOPIC_NAME);
-                int partition = partitionStateData.get(PARTITION_ID);
-                PartitionState partitionState = new PartitionState(partitionStateData);
-                partitionStates.put(new TopicPartition(topic, partition), partitionState);
-            }
+    public Iterable partitionStates() {
+        if (version() >= 5) {
+            return () -> new FlattenedIterator<>(data.topicStates().iterator(),
+                topicState -> topicState.partitionStates().iterator());
         }
+        return data.ungroupedPartitionStates();
+    }
 
-        Set liveBrokers = new HashSet<>();
-
-        for (Object brokerDataObj : struct.get(LIVE_BROKERS)) {
-            Struct brokerData = (Struct) brokerDataObj;
-            int brokerId = brokerData.get(BROKER_ID);
+    public List liveBrokers() {
+        return data.liveBrokers();
+    }
 
-            // V0
-            if (brokerData.hasField(HOST)) {
-                String host = brokerData.get(HOST);
-                int port = brokerData.get(PORT);
-                List endPoints = new ArrayList<>(1);
-                SecurityProtocol securityProtocol = SecurityProtocol.PLAINTEXT;
-                endPoints.add(new EndPoint(host, port, securityProtocol, ListenerName.forSecurityProtocol(securityProtocol)));
-                liveBrokers.add(new Broker(brokerId, endPoints, null));
-            } else { // V1, V2 or V3
-                List endPoints = new ArrayList<>();
-                for (Object endPointDataObj : brokerData.get(ENDPOINTS)) {
-                    Struct endPointData = (Struct) endPointDataObj;
-                    int port = endPointData.get(PORT);
-                    String host = endPointData.get(HOST);
-                    short protocolTypeId = endPointData.get(SECURITY_PROTOCOL_TYPE);
-                    SecurityProtocol securityProtocol = SecurityProtocol.forId(protocolTypeId);
-                    String listenerName;
-                    if (endPointData.hasField(LISTENER_NAME)) // V3
-                        listenerName = endPointData.get(LISTENER_NAME);
-                    else
-                        listenerName = securityProtocol.name;
-                    endPoints.add(new EndPoint(host, port, securityProtocol, new ListenerName(listenerName)));
-                }
-                String rack = null;
-                if (brokerData.hasField(RACK)) { // V2
-                    rack = brokerData.get(RACK);
-                }
-                liveBrokers.add(new Broker(brokerId, endPoints, rack));
+    @Override
+    public Send toSend(String destination, RequestHeader header) {
+        // For UpdateMetadataRequest, the toSend method on the same object will be called many times, each time with a different destination
+        // value and a header containing a different correlation id.
+        ByteBuffer headerBuffer = serialize(header);
+        bodyBufferLock.lock();
+        try {
+            if (bodyBuffer == null) {
+                bodyBuffer = RequestUtils.serialize(toStruct()).array();
             }
+        } finally {
+            bodyBufferLock.unlock();
         }
-        this.partitionStates = partitionStates;
-        this.liveBrokers = liveBrokers;
+        return new NetworkSend(destination, new ByteBuffer[]{headerBuffer, ByteBuffer.wrap(bodyBuffer)});
     }
 
     @Override
     protected Struct toStruct() {
-        short version = version();
-        Struct struct = new Struct(ApiKeys.UPDATE_METADATA.requestSchema(version));
-        struct.set(CONTROLLER_ID, controllerId);
-        struct.set(CONTROLLER_EPOCH, controllerEpoch);
-        struct.setIfExists(BROKER_EPOCH, brokerEpoch);
-
-        if (struct.hasField(TOPIC_STATES)) {
-            Map> topicStates = CollectionUtils.groupPartitionDataByTopic(partitionStates);
-            List topicStatesData = new ArrayList<>(topicStates.size());
-            for (Map.Entry> entry : topicStates.entrySet()) {
-                Struct topicStateData = struct.instance(TOPIC_STATES);
-                topicStateData.set(TOPIC_NAME, entry.getKey());
-                Map partitionMap = entry.getValue();
-                List partitionStatesData = new ArrayList<>(partitionMap.size());
-                for (Map.Entry partitionEntry : partitionMap.entrySet()) {
-                    Struct partitionStateData = topicStateData.instance(PARTITION_STATES);
-                    partitionStateData.set(PARTITION_ID, partitionEntry.getKey());
-                    partitionEntry.getValue().setStruct(partitionStateData);
-                    partitionStatesData.add(partitionStateData);
-                }
-                topicStateData.set(PARTITION_STATES, partitionStatesData.toArray());
-                topicStatesData.add(topicStateData);
-            }
-            struct.set(TOPIC_STATES, topicStatesData.toArray());
-        } else {
-            List partitionStatesData = new ArrayList<>(partitionStates.size());
-            for (Map.Entry entry : partitionStates.entrySet()) {
-                Struct partitionStateData = struct.instance(PARTITION_STATES);
-                TopicPartition topicPartition = entry.getKey();
-                partitionStateData.set(TOPIC_NAME, topicPartition.topic());
-                partitionStateData.set(PARTITION_ID, topicPartition.partition());
-                entry.getValue().setStruct(partitionStateData);
-                partitionStatesData.add(partitionStateData);
-            }
-            struct.set(PARTITION_STATES, partitionStatesData.toArray());
-        }
-
-        List brokersData = new ArrayList<>(liveBrokers.size());
-        for (Broker broker : liveBrokers) {
-            Struct brokerData = struct.instance(LIVE_BROKERS);
-            brokerData.set(BROKER_ID, broker.id);
-
-            if (version == 0) {
-                EndPoint endPoint = broker.endPoints.get(0);
-                brokerData.set(HOST, endPoint.host);
-                brokerData.set(PORT, endPoint.port);
-            } else {
-                List endPointsData = new ArrayList<>(broker.endPoints.size());
-                for (EndPoint endPoint : broker.endPoints) {
-                    Struct endPointData = brokerData.instance(ENDPOINTS);
-                    endPointData.set(PORT, endPoint.port);
-                    endPointData.set(HOST, endPoint.host);
-                    endPointData.set(SECURITY_PROTOCOL_TYPE, endPoint.securityProtocol.id);
-                    if (version >= 3)
-                        endPointData.set(LISTENER_NAME, endPoint.listenerName.value());
-                    endPointsData.add(endPointData);
-
-                }
-                brokerData.set(ENDPOINTS, endPointsData.toArray());
-                if (version >= 2) {
-                    brokerData.set(RACK, broker.rack);
-                }
-            }
-
-            brokersData.add(brokerData);
+        // the following inner blocks are not indented in order to make hotfix cherry-picking easier
+        structLock.lock();
+        try {
+        if (struct == null) {
+        struct = data.toStruct(version());
         }
-        struct.set(LIVE_BROKERS, brokersData.toArray());
-
         return struct;
+        } finally {
+            structLock.unlock();
+        }
     }
 
-    @Override
-    public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
-        short versionId = version();
-        if (versionId <= 5)
-            return new UpdateMetadataResponse(Errors.forException(e));
-        else
-            throw new IllegalArgumentException(String.format("Version %d is not valid. Valid versions for %s are 0 to %d",
-                versionId, this.getClass().getSimpleName(), ApiKeys.UPDATE_METADATA.latestVersion()));
-    }
-
-    public Map partitionStates() {
-        return partitionStates;
-    }
-
-    public Set liveBrokers() {
-        return liveBrokers;
+    // Visible for testing
+    UpdateMetadataRequestData data() {
+        return data;
     }
 
     public static UpdateMetadataRequest parse(ByteBuffer buffer, short version) {
         return new UpdateMetadataRequest(ApiKeys.UPDATE_METADATA.parseRequest(version, buffer), version);
     }
-
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/UpdateMetadataResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/UpdateMetadataResponse.java
index 478265b395863..c7d803f560eb1 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/UpdateMetadataResponse.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/UpdateMetadataResponse.java
@@ -16,62 +16,41 @@
  */
 package org.apache.kafka.common.requests;
 
+import org.apache.kafka.common.message.UpdateMetadataResponseData;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.protocol.types.Schema;
 import org.apache.kafka.common.protocol.types.Struct;
 
 import java.nio.ByteBuffer;
 import java.util.Map;
 
-import static org.apache.kafka.common.protocol.CommonFields.ERROR_CODE;
-
 public class UpdateMetadataResponse extends AbstractResponse {
-    private static final Schema UPDATE_METADATA_RESPONSE_V0 = new Schema(ERROR_CODE);
-    private static final Schema UPDATE_METADATA_RESPONSE_V1 = UPDATE_METADATA_RESPONSE_V0;
-    private static final Schema UPDATE_METADATA_RESPONSE_V2 = UPDATE_METADATA_RESPONSE_V1;
-    private static final Schema UPDATE_METADATA_RESPONSE_V3 = UPDATE_METADATA_RESPONSE_V2;
-    private static final Schema UPDATE_METADATA_RESPONSE_V4 = UPDATE_METADATA_RESPONSE_V3;
-    private static final Schema UPDATE_METADATA_RESPONSE_V5 = UPDATE_METADATA_RESPONSE_V4;
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{UPDATE_METADATA_RESPONSE_V0, UPDATE_METADATA_RESPONSE_V1, UPDATE_METADATA_RESPONSE_V2,
-            UPDATE_METADATA_RESPONSE_V3, UPDATE_METADATA_RESPONSE_V4, UPDATE_METADATA_RESPONSE_V5};
-    }
 
-    /**
-     * Possible error code:
-     *
-     * STALE_CONTROLLER_EPOCH (11)
-     * STALE_BROKER_EPOCH (77)
-     */
-    private final Errors error;
+    private final UpdateMetadataResponseData data;
 
-    public UpdateMetadataResponse(Errors error) {
-        this.error = error;
+    public UpdateMetadataResponse(UpdateMetadataResponseData data) {
+        this.data = data;
     }
 
-    public UpdateMetadataResponse(Struct struct) {
-        error = Errors.forCode(struct.get(ERROR_CODE));
+    public UpdateMetadataResponse(Struct struct, short version) {
+        this(new UpdateMetadataResponseData(struct, version));
     }
 
     public Errors error() {
-        return error;
+        return Errors.forCode(data.errorCode());
     }
 
     @Override
     public Map errorCounts() {
-        return errorCounts(error);
+        return errorCounts(error());
     }
 
     public static UpdateMetadataResponse parse(ByteBuffer buffer, short version) {
-        return new UpdateMetadataResponse(ApiKeys.UPDATE_METADATA.parseResponse(version, buffer));
+        return new UpdateMetadataResponse(ApiKeys.UPDATE_METADATA.parseResponse(version, buffer), version);
     }
 
     @Override
     protected Struct toStruct(short version) {
-        Struct struct = new Struct(ApiKeys.UPDATE_METADATA.responseSchema(version));
-        struct.set(ERROR_CODE, error.code());
-        return struct;
+        return data.toStruct(version);
     }
-}
\ No newline at end of file
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/JaasUtils.java b/clients/src/main/java/org/apache/kafka/common/security/JaasUtils.java
index c2e7fbdc74a2a..593f104e300fa 100644
--- a/clients/src/main/java/org/apache/kafka/common/security/JaasUtils.java
+++ b/clients/src/main/java/org/apache/kafka/common/security/JaasUtils.java
@@ -16,12 +16,12 @@
  */
 package org.apache.kafka.common.security;
 
-import javax.security.auth.login.Configuration;
-
 import org.apache.kafka.common.KafkaException;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import javax.security.auth.login.Configuration;
+
 public final class JaasUtils {
     private static final Logger LOG = LoggerFactory.getLogger(JaasUtils.class);
     public static final String JAVA_LOGIN_CONFIG_PARAM = "java.security.auth.login.config";
@@ -31,27 +31,48 @@ public final class JaasUtils {
     public static final String ZK_SASL_CLIENT = "zookeeper.sasl.client";
     public static final String ZK_LOGIN_CONTEXT_NAME_KEY = "zookeeper.sasl.clientconfig";
 
+    private static final String DEFAULT_ZK_LOGIN_CONTEXT_NAME = "Client";
+    private static final String DEFAULT_ZK_SASL_CLIENT = "true";
+
     private JaasUtils() {}
 
+    public static String zkSecuritySysConfigString() {
+        String loginConfig = System.getProperty(JAVA_LOGIN_CONFIG_PARAM);
+        String clientEnabled = System.getProperty(ZK_SASL_CLIENT, "default:" + DEFAULT_ZK_SASL_CLIENT);
+        String contextName = System.getProperty(ZK_LOGIN_CONTEXT_NAME_KEY, "default:" + DEFAULT_ZK_LOGIN_CONTEXT_NAME);
+        return "[" +
+                JAVA_LOGIN_CONFIG_PARAM + "=" + loginConfig +
+                ", " +
+                ZK_SASL_CLIENT + "=" + clientEnabled +
+                ", " +
+                ZK_LOGIN_CONTEXT_NAME_KEY + "=" + contextName +
+                "]";
+    }
+
     public static boolean isZkSecurityEnabled() {
-        boolean zkSaslEnabled = Boolean.parseBoolean(System.getProperty(ZK_SASL_CLIENT, "true"));
-        String zkLoginContextName = System.getProperty(ZK_LOGIN_CONTEXT_NAME_KEY, "Client");
+        boolean zkSaslEnabled = Boolean.parseBoolean(System.getProperty(ZK_SASL_CLIENT, DEFAULT_ZK_SASL_CLIENT));
+        String zkLoginContextName = System.getProperty(ZK_LOGIN_CONTEXT_NAME_KEY, DEFAULT_ZK_LOGIN_CONTEXT_NAME);
 
-        boolean isSecurityEnabled;
+        LOG.debug("Checking login config for Zookeeper JAAS context {}", zkSecuritySysConfigString());
+
+        boolean foundLoginConfigEntry;
         try {
             Configuration loginConf = Configuration.getConfiguration();
-            isSecurityEnabled = loginConf.getAppConfigurationEntry(zkLoginContextName) != null;
+            foundLoginConfigEntry = loginConf.getAppConfigurationEntry(zkLoginContextName) != null;
         } catch (Exception e) {
-            throw new KafkaException("Exception while loading Zookeeper JAAS login context '" + zkLoginContextName + "'", e);
+            throw new KafkaException("Exception while loading Zookeeper JAAS login context " +
+                    zkSecuritySysConfigString(), e);
         }
-        if (isSecurityEnabled && !zkSaslEnabled) {
+
+        if (foundLoginConfigEntry && !zkSaslEnabled) {
             LOG.error("JAAS configuration is present, but system property " +
                         ZK_SASL_CLIENT + " is set to false, which disables " +
                         "SASL in the ZooKeeper client");
-            throw new KafkaException("Exception while determining if ZooKeeper is secure");
+            throw new KafkaException("Exception while determining if ZooKeeper is secure " +
+                    zkSecuritySysConfigString());
         }
 
-        return isSecurityEnabled;
+        return foundLoginConfigEntry;
     }
 }
 
diff --git a/clients/src/main/java/org/apache/kafka/common/security/auth/SecurityProviderCreator.java b/clients/src/main/java/org/apache/kafka/common/security/auth/SecurityProviderCreator.java
new file mode 100644
index 0000000000000..ae56f9a3520a3
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/auth/SecurityProviderCreator.java
@@ -0,0 +1,43 @@
+/*
+ * 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.security.auth;
+
+import org.apache.kafka.common.Configurable;
+import org.apache.kafka.common.annotation.InterfaceStability;
+
+import java.security.Provider;
+import java.util.Map;
+
+/**
+ * An interface for generating security providers.
+ */
+@InterfaceStability.Evolving
+public interface SecurityProviderCreator extends Configurable {
+
+    /**
+     * Configure method is used to configure the generator to create the Security Provider
+     * @param config configuration parameters for initialising security provider
+     */
+    default void configure(Map config) {
+
+    }
+
+    /**
+     * Generate the security provider configured
+     */
+    Provider getProvider();
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/authenticator/LoginManager.java b/clients/src/main/java/org/apache/kafka/common/security/authenticator/LoginManager.java
index 3d9481c497d90..6613fd147f89a 100644
--- a/clients/src/main/java/org/apache/kafka/common/security/authenticator/LoginManager.java
+++ b/clients/src/main/java/org/apache/kafka/common/security/authenticator/LoginManager.java
@@ -25,6 +25,7 @@
 import org.apache.kafka.common.security.auth.Login;
 import org.apache.kafka.common.security.oauthbearer.OAuthBearerLoginModule;
 import org.apache.kafka.common.security.oauthbearer.internals.unsecured.OAuthBearerUnsecuredLoginCallbackHandler;
+import org.apache.kafka.common.utils.SecurityUtils;
 import org.apache.kafka.common.utils.Utils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -112,6 +113,7 @@ public static LoginManager acquireLoginManager(JaasContext jaasContext, String s
                     STATIC_INSTANCES.put(loginMetadata, loginManager);
                 }
             }
+            SecurityUtils.addConfiguredSecurityProviders(configs);
             return loginManager.acquire();
         }
     }
diff --git a/clients/src/main/java/org/apache/kafka/common/security/authenticator/SaslClientAuthenticator.java b/clients/src/main/java/org/apache/kafka/common/security/authenticator/SaslClientAuthenticator.java
index 3133f44a88f3f..c75bb82265b03 100644
--- a/clients/src/main/java/org/apache/kafka/common/security/authenticator/SaslClientAuthenticator.java
+++ b/clients/src/main/java/org/apache/kafka/common/security/authenticator/SaslClientAuthenticator.java
@@ -23,6 +23,8 @@
 import org.apache.kafka.common.errors.IllegalSaslStateException;
 import org.apache.kafka.common.errors.SaslAuthenticationException;
 import org.apache.kafka.common.errors.UnsupportedSaslMechanismException;
+import org.apache.kafka.common.message.ApiVersionsResponseData.ApiVersionsResponseKey;
+import org.apache.kafka.common.message.RequestHeaderData;
 import org.apache.kafka.common.message.SaslAuthenticateRequestData;
 import org.apache.kafka.common.message.SaslHandshakeRequestData;
 import org.apache.kafka.common.network.Authenticator;
@@ -37,7 +39,6 @@
 import org.apache.kafka.common.requests.AbstractResponse;
 import org.apache.kafka.common.requests.ApiVersionsRequest;
 import org.apache.kafka.common.requests.ApiVersionsResponse;
-import org.apache.kafka.common.requests.ApiVersionsResponse.ApiVersion;
 import org.apache.kafka.common.requests.RequestHeader;
 import org.apache.kafka.common.requests.SaslAuthenticateRequest;
 import org.apache.kafka.common.requests.SaslAuthenticateResponse;
@@ -105,6 +106,35 @@ public enum SaslState {
     private static final short DISABLE_KAFKA_SASL_AUTHENTICATE_HEADER = -1;
     private static final Random RNG = new Random();
 
+    /**
+     * the reserved range of correlation id for Sasl requests.
+     *
+     * Noted: there is a story about reserved range. The response of LIST_OFFSET is compatible to response of SASL_HANDSHAKE.
+     * Hence, we could miss the schema error when using schema of SASL_HANDSHAKE to parse response of LIST_OFFSET.
+     * For example: the IllegalStateException caused by mismatched correlation id is thrown if following steps happens.
+     * 1) sent LIST_OFFSET
+     * 2) sent SASL_HANDSHAKE
+     * 3) receive response of LIST_OFFSET
+     * 4) succeed to use schema of SASL_HANDSHAKE to parse response of LIST_OFFSET
+     * 5) throw IllegalStateException due to mismatched correlation id
+     * As a simple approach, we force Sasl requests to use a reserved correlation id which is separated from those
+     * used in NetworkClient for Kafka requests. Hence, we can guarantee that every SASL request will throw
+     * SchemaException due to correlation id mismatch during reauthentication
+     */
+    public static final int MAX_RESERVED_CORRELATION_ID = Integer.MAX_VALUE;
+
+    /**
+     * We only expect one request in-flight a time during authentication so the small range is fine.
+     */
+    public static final int MIN_RESERVED_CORRELATION_ID = MAX_RESERVED_CORRELATION_ID - 7;
+
+    /**
+     * @return true if the correlation id is reserved for SASL request. otherwise, false
+     */
+    public static boolean isReserved(int correlationId) {
+        return correlationId >= MIN_RESERVED_CORRELATION_ID;
+    }
+
     private final Subject subject;
     private final String servicePrincipal;
     private final String host;
@@ -149,7 +179,7 @@ public SaslClientAuthenticator(Map configs,
         this.host = host;
         this.servicePrincipal = servicePrincipal;
         this.mechanism = mechanism;
-        this.correlationId = -1;
+        this.correlationId = 0;
         this.transportLayer = transportLayer;
         this.configs = configs;
         this.saslAuthenticateVersion = DISABLE_KAFKA_SASL_AUTHENTICATE_HEADER;
@@ -173,7 +203,8 @@ public SaslClientAuthenticator(Map configs,
         }
     }
 
-    private SaslClient createSaslClient() {
+    // visible for testing
+    SaslClient createSaslClient() {
         try {
             return Subject.doAs(subject, (PrivilegedExceptionAction) () -> {
                 String[] mechs = {mechanism};
@@ -201,7 +232,7 @@ public void authenticate() throws IOException {
         switch (saslState) {
             case SEND_APIVERSIONS_REQUEST:
                 // Always use version 0 request since brokers treat requests with schema exceptions as GSSAPI tokens
-                ApiVersionsRequest apiVersionsRequest = new ApiVersionsRequest((short) 0);
+                ApiVersionsRequest apiVersionsRequest = new ApiVersionsRequest.Builder().build((short) 0);
                 send(apiVersionsRequest.toSend(node, nextRequestHeader(ApiKeys.API_VERSIONS, apiVersionsRequest.version())));
                 setSaslState(SaslState.RECEIVE_APIVERSIONS_RESPONSE);
                 break;
@@ -284,7 +315,7 @@ public void authenticate() throws IOException {
 
     private void sendHandshakeRequest(ApiVersionsResponse apiVersionsResponse) throws IOException {
         SaslHandshakeRequest handshakeRequest = createSaslHandshakeRequest(
-                apiVersionsResponse.apiVersion(ApiKeys.SASL_HANDSHAKE.id).maxVersion);
+                apiVersionsResponse.apiVersion(ApiKeys.SASL_HANDSHAKE.id).maxVersion());
         send(handshakeRequest.toSend(node, nextRequestHeader(ApiKeys.SASL_HANDSHAKE, handshakeRequest.version())));
     }
 
@@ -322,9 +353,23 @@ public Long reauthenticationLatencyMs() {
         return reauthInfo.reauthenticationLatencyMs();
     }
 
+    // visible for testing
+    int nextCorrelationId() {
+        if (!isReserved(correlationId))
+            correlationId = MIN_RESERVED_CORRELATION_ID;
+        return correlationId++;
+    }
+
     private RequestHeader nextRequestHeader(ApiKeys apiKey, short version) {
         String clientId = (String) configs.get(CommonClientConfigs.CLIENT_ID_CONFIG);
-        currentRequestHeader = new RequestHeader(apiKey, version, clientId, correlationId++);
+        short requestApiKey = apiKey.id;
+        currentRequestHeader = new RequestHeader(
+            new RequestHeaderData().
+                setRequestApiKey(requestApiKey).
+                setRequestApiVersion(version).
+                setClientId(clientId).
+                setCorrelationId(nextCorrelationId()),
+            apiKey.requestHeaderVersion(version));
         return currentRequestHeader;
     }
 
@@ -336,9 +381,9 @@ protected SaslHandshakeRequest createSaslHandshakeRequest(short version) {
 
     // Visible to override for testing
     protected void saslAuthenticateVersion(ApiVersionsResponse apiVersionsResponse) {
-        ApiVersion authenticateVersion = apiVersionsResponse.apiVersion(ApiKeys.SASL_AUTHENTICATE.id);
+        ApiVersionsResponseKey authenticateVersion = apiVersionsResponse.apiVersion(ApiKeys.SASL_AUTHENTICATE.id);
         if (authenticateVersion != null)
-            this.saslAuthenticateVersion = (short) Math.min(authenticateVersion.maxVersion,
+            this.saslAuthenticateVersion = (short) Math.min(authenticateVersion.maxVersion(),
                     ApiKeys.SASL_AUTHENTICATE.latestVersion());
     }
 
diff --git a/clients/src/main/java/org/apache/kafka/common/security/authenticator/SaslServerAuthenticator.java b/clients/src/main/java/org/apache/kafka/common/security/authenticator/SaslServerAuthenticator.java
index 7aca17716bb41..afaf6b728a112 100644
--- a/clients/src/main/java/org/apache/kafka/common/security/authenticator/SaslServerAuthenticator.java
+++ b/clients/src/main/java/org/apache/kafka/common/security/authenticator/SaslServerAuthenticator.java
@@ -568,7 +568,7 @@ private String handleHandshakeRequest(RequestContext context, SaslHandshakeReque
 
     // Visible to override for testing
     protected ApiVersionsResponse apiVersionsResponse() {
-        return ApiVersionsResponse.defaultApiVersionsResponse();
+        return ApiVersionsResponse.DEFAULT_API_VERSIONS_RESPONSE;
     }
 
     // Visible to override for testing
@@ -582,6 +582,8 @@ private void handleApiVersionsRequest(RequestContext context, ApiVersionsRequest
 
         if (apiVersionsRequest.hasUnsupportedRequestVersion())
             sendKafkaResponse(context, apiVersionsRequest.getErrorResponse(0, Errors.UNSUPPORTED_VERSION.exception()));
+        else if (!apiVersionsRequest.isValid())
+            sendKafkaResponse(context, apiVersionsRequest.getErrorResponse(0, Errors.INVALID_REQUEST.exception()));
         else {
             sendKafkaResponse(context, apiVersionsResponse());
             setSaslState(SaslState.HANDSHAKE_REQUEST);
@@ -676,30 +678,29 @@ private long calcCompletionTimesAndReturnSessionLifetimeMs() {
             Long connectionsMaxReauthMs = connectionsMaxReauthMsByMechanism.get(saslMechanism);
             if (credentialExpirationMs != null || connectionsMaxReauthMs != null) {
                 if (credentialExpirationMs == null)
-                    retvalSessionLifetimeMs = zeroIfNegative(connectionsMaxReauthMs.longValue());
+                    retvalSessionLifetimeMs = zeroIfNegative(connectionsMaxReauthMs);
                 else if (connectionsMaxReauthMs == null)
-                    retvalSessionLifetimeMs = zeroIfNegative(credentialExpirationMs.longValue() - authenticationEndMs);
+                    retvalSessionLifetimeMs = zeroIfNegative(credentialExpirationMs - authenticationEndMs);
                 else
                     retvalSessionLifetimeMs = zeroIfNegative(
-                            Math.min(credentialExpirationMs.longValue() - authenticationEndMs,
-                                    connectionsMaxReauthMs.longValue()));
+                            Math.min(credentialExpirationMs - authenticationEndMs,
+                                    connectionsMaxReauthMs));
                 if (retvalSessionLifetimeMs > 0L)
-                    sessionExpirationTimeNanos = Long
-                            .valueOf(authenticationEndNanos + 1000 * 1000 * retvalSessionLifetimeMs);
+                    sessionExpirationTimeNanos = authenticationEndNanos + 1000 * 1000 * retvalSessionLifetimeMs;
             }
             if (credentialExpirationMs != null) {
                 if (sessionExpirationTimeNanos != null)
                     LOG.debug(
                             "Authentication complete; session max lifetime from broker config={} ms, credential expiration={} ({} ms); session expiration = {} ({} ms), sending {} ms to client",
                             connectionsMaxReauthMs, new Date(credentialExpirationMs),
-                            Long.valueOf(credentialExpirationMs.longValue() - authenticationEndMs),
+                            credentialExpirationMs - authenticationEndMs,
                             new Date(authenticationEndMs + retvalSessionLifetimeMs), retvalSessionLifetimeMs,
                             retvalSessionLifetimeMs);
                 else
                     LOG.debug(
                             "Authentication complete; session max lifetime from broker config={} ms, credential expiration={} ({} ms); no session expiration, sending 0 ms to client",
                             connectionsMaxReauthMs, new Date(credentialExpirationMs),
-                            Long.valueOf(credentialExpirationMs.longValue() - authenticationEndMs));
+                            credentialExpirationMs - authenticationEndMs);
             } else {
                 if (sessionExpirationTimeNanos != null)
                     LOG.debug(
@@ -719,7 +720,7 @@ public Long reauthenticationLatencyMs() {
                 return null;
             // record at least 1 ms if there is some latency
             long latencyNanos = authenticationEndNanos - reauthenticationBeginNanos;
-            return latencyNanos == 0L ? 0L : Math.max(1L, Long.valueOf(Math.round(latencyNanos / 1000.0 / 1000.0)));
+            return latencyNanos == 0L ? 0L : Math.max(1L, Math.round(latencyNanos / 1000.0 / 1000.0));
         }
 
         private long zeroIfNegative(long value) {
diff --git a/clients/src/main/java/org/apache/kafka/common/security/kerberos/KerberosRule.java b/clients/src/main/java/org/apache/kafka/common/security/kerberos/KerberosRule.java
index 92a70b1f19456..91280ca65a347 100644
--- a/clients/src/main/java/org/apache/kafka/common/security/kerberos/KerberosRule.java
+++ b/clients/src/main/java/org/apache/kafka/common/security/kerberos/KerberosRule.java
@@ -46,6 +46,7 @@ class KerberosRule {
     private final String toPattern;
     private final boolean repeat;
     private final boolean toLowerCase;
+    private final boolean toUpperCase;
 
     KerberosRule(String defaultRealm) {
         this.defaultRealm = defaultRealm;
@@ -57,10 +58,11 @@ class KerberosRule {
         toPattern = null;
         repeat = false;
         toLowerCase = false;
+        toUpperCase = false;
     }
 
     KerberosRule(String defaultRealm, int numOfComponents, String format, String match, String fromPattern,
-                 String toPattern, boolean repeat, boolean toLowerCase) {
+                 String toPattern, boolean repeat, boolean toLowerCase, boolean toUpperCase) {
         this.defaultRealm = defaultRealm;
         isDefault = false;
         this.numOfComponents = numOfComponents;
@@ -71,6 +73,7 @@ class KerberosRule {
         this.toPattern = toPattern;
         this.repeat = repeat;
         this.toLowerCase = toLowerCase;
+        this.toUpperCase = toUpperCase;
     }
 
     @Override
@@ -102,6 +105,9 @@ public String toString() {
             if (toLowerCase) {
                 buf.append("/L");
             }
+            if (toUpperCase) {
+                buf.append("/U");
+            }
         }
         return buf.toString();
     }
@@ -191,7 +197,10 @@ String apply(String[] params) throws IOException {
         }
         if (toLowerCase && result != null) {
             result = result.toLowerCase(Locale.ENGLISH);
+        } else if (toUpperCase && result != null) {
+            result = result.toUpperCase(Locale.ENGLISH);
         }
+
         return result;
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/security/kerberos/KerberosShortNamer.java b/clients/src/main/java/org/apache/kafka/common/security/kerberos/KerberosShortNamer.java
index 69b4689009baa..96e01f17942b4 100644
--- a/clients/src/main/java/org/apache/kafka/common/security/kerberos/KerberosShortNamer.java
+++ b/clients/src/main/java/org/apache/kafka/common/security/kerberos/KerberosShortNamer.java
@@ -33,7 +33,7 @@ public class KerberosShortNamer {
     /**
      * A pattern for parsing a auth_to_local rule.
      */
-    private static final Pattern RULE_PARSER = Pattern.compile("((DEFAULT)|((RULE:\\[(\\d*):([^\\]]*)](\\(([^)]*)\\))?(s/([^/]*)/([^/]*)/(g)?)?/?(L)?)))");
+    private static final Pattern RULE_PARSER = Pattern.compile("((DEFAULT)|((RULE:\\[(\\d*):([^\\]]*)](\\(([^)]*)\\))?(s/([^/]*)/([^/]*)/(g)?)?/?(L|U)?)))");
 
     /* Rules for the translation of the principal name into an operating system name */
     private final List principalToLocalRules;
@@ -66,7 +66,8 @@ private static List parseRules(String defaultRealm, List r
                         matcher.group(10),
                         matcher.group(11),
                         "g".equals(matcher.group(12)),
-                        "L".equals(matcher.group(13))));
+                        "L".equals(matcher.group(13)),
+                        "U".equals(matcher.group(13))));
 
             }
         }
diff --git a/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/OAuthBearerExtensionsValidatorCallback.java b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/OAuthBearerExtensionsValidatorCallback.java
index 97ac4d96838f5..eab208baa6f2b 100644
--- a/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/OAuthBearerExtensionsValidatorCallback.java
+++ b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/OAuthBearerExtensionsValidatorCallback.java
@@ -76,7 +76,7 @@ public Map validatedExtensions() {
     }
 
     /**
-     * @return An immutable {@link Map} consisting of the name->error messages of extensions which failed validation
+     * @return An immutable {@link Map} consisting of the name->error messages of extensions which failed validation
      */
     public Map invalidExtensions() {
         return Collections.unmodifiableMap(invalidExtensions);
diff --git a/clients/src/main/java/org/apache/kafka/common/security/ssl/BoringSslContextProvider.java b/clients/src/main/java/org/apache/kafka/common/security/ssl/BoringSslContextProvider.java
new file mode 100644
index 0000000000000..0a41be81efeb0
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/ssl/BoringSslContextProvider.java
@@ -0,0 +1,43 @@
+/*
+ * 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.security.ssl;
+
+import java.security.NoSuchAlgorithmException;
+import java.security.Provider;
+import java.util.Map;
+import javax.net.ssl.SSLContext;
+import org.apache.kafka.common.config.SslConfigs;
+import org.conscrypt.Conscrypt;
+
+/**
+ * A BoringSslContext provider based on Conscrypt library.
+ */
+public class BoringSslContextProvider implements SslContextProvider {
+    private String protocol;
+    private Provider provider;
+
+    @Override
+    public void configure(Map configs) {
+        protocol = (String) configs.get(SslConfigs.SSL_PROTOCOL_CONFIG);
+        provider = Conscrypt.newProvider();
+    }
+
+    @Override
+    public SSLContext getSSLContext()  throws NoSuchAlgorithmException {
+        return SSLContext.getInstance(protocol, provider);
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/ssl/SimpleSslContextProvider.java b/clients/src/main/java/org/apache/kafka/common/security/ssl/SimpleSslContextProvider.java
new file mode 100644
index 0000000000000..90ba337e88e1f
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/ssl/SimpleSslContextProvider.java
@@ -0,0 +1,43 @@
+/*
+ * 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.security.ssl;
+
+import java.security.NoSuchAlgorithmException;
+import java.security.NoSuchProviderException;
+import java.util.Map;
+import javax.net.ssl.SSLContext;
+import org.apache.kafka.common.config.SslConfigs;
+
+
+public class SimpleSslContextProvider implements SslContextProvider {
+    private String protocol;
+    private String provider;
+
+    @Override
+    public void configure(Map configs) {
+        this.protocol =  (String) configs.get(SslConfigs.SSL_PROTOCOL_CONFIG);
+        this.provider = (String) configs.get(SslConfigs.SSL_PROVIDER_CONFIG);
+    }
+
+    @Override
+    public SSLContext getSSLContext() throws NoSuchAlgorithmException, NoSuchProviderException {
+        if (provider == null) {
+            return SSLContext.getInstance(protocol);
+        }
+        return SSLContext.getInstance(protocol, provider);
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/ssl/SslContextProvider.java b/clients/src/main/java/org/apache/kafka/common/security/ssl/SslContextProvider.java
new file mode 100644
index 0000000000000..ff2e8fc954348
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/ssl/SslContextProvider.java
@@ -0,0 +1,37 @@
+/*
+ * 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.security.ssl;
+
+import java.security.NoSuchAlgorithmException;
+import java.security.NoSuchProviderException;
+import javax.net.ssl.SSLContext;
+import org.apache.kafka.common.Configurable;
+
+/**
+ * A SSLContext provider interface to create SSLContext based on configs
+ */
+public interface SslContextProvider extends Configurable {
+
+    /**
+    * Create a SSLContext based on protocol and provider config.
+    * @return SSLContext
+    * @throws NoSuchAlgorithmException if protocol config is invalid.
+    * @throws NoSuchProviderException if provider config is invalid.
+    */
+    SSLContext getSSLContext() throws NoSuchAlgorithmException, NoSuchProviderException;
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/ssl/SslEngineBuilder.java b/clients/src/main/java/org/apache/kafka/common/security/ssl/SslEngineBuilder.java
new file mode 100644
index 0000000000000..b44298e25b990
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/ssl/SslEngineBuilder.java
@@ -0,0 +1,317 @@
+/*
+ * 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.security.ssl;
+
+import org.apache.kafka.common.KafkaException;
+import org.apache.kafka.common.config.SslClientAuth;
+import org.apache.kafka.common.config.SslConfigs;
+import org.apache.kafka.common.config.internals.BrokerSecurityConfigs;
+import org.apache.kafka.common.network.Mode;
+import org.apache.kafka.common.config.types.Password;
+import org.apache.kafka.common.utils.SecurityUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.net.ssl.KeyManager;
+import javax.net.ssl.KeyManagerFactory;
+import javax.net.ssl.SSLContext;
+import javax.net.ssl.SSLEngine;
+import javax.net.ssl.SSLParameters;
+import javax.net.ssl.TrustManagerFactory;
+import java.io.IOException;
+import java.io.InputStream;
+import java.nio.file.Files;
+import java.nio.file.Paths;
+import java.security.SecureRandom;
+import java.security.GeneralSecurityException;
+import java.security.KeyStore;
+import java.util.Collections;
+import java.util.Date;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.stream.Collectors;
+
+public class SslEngineBuilder {
+    private static final Logger log = LoggerFactory.getLogger(SslEngineBuilder.class);
+
+    private final Map configs;
+    private SslContextProvider sslContextProvider;
+    private final String kmfAlgorithm;
+    private final String tmfAlgorithm;
+    private final SecurityStore keystore;
+    private final SecurityStore truststore;
+    private final String[] cipherSuites;
+    private final String[] enabledProtocols;
+    private final SecureRandom secureRandomImplementation;
+    private final SSLContext sslContext;
+    private final SslClientAuth sslClientAuth;
+
+    @SuppressWarnings("unchecked")
+    SslEngineBuilder(Map configs) {
+        this.configs = Collections.unmodifiableMap(configs);
+        try {
+            String sslContextProvider = (String) configs.get(SslConfigs.SSL_CONTEXT_PROVIDER_CLASS_CONFIG);
+            this.sslContextProvider = (SslContextProvider) Class.forName(sslContextProvider).newInstance();
+        } catch (Exception e) {
+            throw new KafkaException(e);
+        }
+        this.sslContextProvider.configure(configs);
+        SecurityUtils.addConfiguredSecurityProviders(this.configs);
+
+        List cipherSuitesList = (List) configs.get(SslConfigs.SSL_CIPHER_SUITES_CONFIG);
+        if (cipherSuitesList != null && !cipherSuitesList.isEmpty()) {
+            this.cipherSuites = cipherSuitesList.toArray(new String[cipherSuitesList.size()]);
+        } else {
+            this.cipherSuites = null;
+        }
+
+        List enabledProtocolsList = (List) configs.get(SslConfigs.SSL_ENABLED_PROTOCOLS_CONFIG);
+        if (enabledProtocolsList != null && !enabledProtocolsList.isEmpty()) {
+            this.enabledProtocols = enabledProtocolsList.toArray(new String[enabledProtocolsList.size()]);
+        } else {
+            this.enabledProtocols = null;
+        }
+
+        this.secureRandomImplementation = createSecureRandom((String)
+                configs.get(SslConfigs.SSL_SECURE_RANDOM_IMPLEMENTATION_CONFIG));
+
+        this.sslClientAuth = createSslClientAuth((String) configs.get(
+                BrokerSecurityConfigs.SSL_CLIENT_AUTH_CONFIG));
+
+        this.kmfAlgorithm = (String) configs.get(SslConfigs.SSL_KEYMANAGER_ALGORITHM_CONFIG);
+        this.tmfAlgorithm = (String) configs.get(SslConfigs.SSL_TRUSTMANAGER_ALGORITHM_CONFIG);
+
+        this.keystore = createKeystore((String) configs.get(SslConfigs.SSL_KEYSTORE_TYPE_CONFIG),
+                (String) configs.get(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG),
+                (Password) configs.get(SslConfigs.SSL_KEYSTORE_PASSWORD_CONFIG),
+                (Password) configs.get(SslConfigs.SSL_KEY_PASSWORD_CONFIG));
+
+        this.truststore = createTruststore((String) configs.get(SslConfigs.SSL_TRUSTSTORE_TYPE_CONFIG),
+                (String) configs.get(SslConfigs.SSL_TRUSTSTORE_LOCATION_CONFIG),
+                (Password) configs.get(SslConfigs.SSL_TRUSTSTORE_PASSWORD_CONFIG));
+
+        this.sslContext = createSSLContext();
+    }
+
+    private static SslClientAuth createSslClientAuth(String key) {
+        SslClientAuth auth = SslClientAuth.forConfig(key);
+        if (auth != null) {
+            return auth;
+        }
+        log.warn("Unrecognized client authentication configuration {}.  Falling " +
+                "back to NONE.  Recognized client authentication configurations are {}.",
+                key, String.join(", ", SslClientAuth.VALUES.stream().
+                        map(a -> a.name()).collect(Collectors.toList())));
+        return SslClientAuth.NONE;
+    }
+
+    private static SecureRandom createSecureRandom(String key) {
+        if (key == null) {
+            return null;
+        }
+        try {
+            return SecureRandom.getInstance(key);
+        } catch (GeneralSecurityException e) {
+            throw new KafkaException(e);
+        }
+    }
+
+    private SSLContext createSSLContext() {
+        try {
+            SSLContext sslContext = sslContextProvider.getSSLContext();
+
+            KeyManager[] keyManagers = null;
+            if (keystore != null || kmfAlgorithm != null) {
+                String kmfAlgorithm = this.kmfAlgorithm != null ?
+                        this.kmfAlgorithm : KeyManagerFactory.getDefaultAlgorithm();
+                KeyManagerFactory kmf = KeyManagerFactory.getInstance(kmfAlgorithm);
+                if (keystore != null) {
+                    KeyStore ks = keystore.load();
+                    Password keyPassword = keystore.keyPassword != null ? keystore.keyPassword : keystore.password;
+                    kmf.init(ks, keyPassword.value().toCharArray());
+                } else {
+                    kmf.init(null, null);
+                }
+                keyManagers = kmf.getKeyManagers();
+            }
+
+            String tmfAlgorithm = this.tmfAlgorithm != null ? this.tmfAlgorithm : TrustManagerFactory.getDefaultAlgorithm();
+            TrustManagerFactory tmf = TrustManagerFactory.getInstance(tmfAlgorithm);
+            KeyStore ts = truststore == null ? null : truststore.load();
+            tmf.init(ts);
+
+            sslContext.init(keyManagers, tmf.getTrustManagers(), this.secureRandomImplementation);
+            log.debug("Created SSL context with keystore {}, truststore {}, provider {}.",
+                    keystore, truststore, sslContext.getProvider().getName());
+            return sslContext;
+        } catch (Exception e) {
+            throw new KafkaException(e);
+        }
+    }
+
+    private static SecurityStore createKeystore(String type, String path, Password password, Password keyPassword) {
+        if (path == null && password != null) {
+            throw new KafkaException("SSL key store is not specified, but key store password is specified.");
+        } else if (path != null && password == null) {
+            throw new KafkaException("SSL key store is specified, but key store password is not specified.");
+        } else if (path != null && password != null) {
+            return new SecurityStore(type, path, password, keyPassword);
+        } else
+            return null; // path == null, clients may use this path with brokers that don't require client auth
+    }
+
+    private static SecurityStore createTruststore(String type, String path, Password password) {
+        if (path == null && password != null) {
+            throw new KafkaException("SSL trust store is not specified, but trust store password is specified.");
+        } else if (path != null) {
+            return new SecurityStore(type, path, password, null);
+        } else
+            return null;
+    }
+
+    @SuppressWarnings("unchecked")
+    Map configs() {
+        return (Map) configs;
+    }
+
+    public SecurityStore keystore() {
+        return keystore;
+    }
+
+    public SecurityStore truststore() {
+        return truststore;
+    }
+
+    /**
+     * Create a new SSLEngine object.
+     *
+     * @param mode      Whether to use client or server mode.
+     * @param peerHost  The peer host to use. This is used in client mode if endpoint validation is enabled.
+     * @param peerPort  The peer port to use. This is a hint and not used for validation.
+     * @param endpointIdentification Endpoint identification algorithm for client mode.
+     * @return          The new SSLEngine.
+     */
+    public SSLEngine createSslEngine(Mode mode, String peerHost, int peerPort, String endpointIdentification) {
+        SSLEngine sslEngine = sslContext.createSSLEngine(peerHost, peerPort);
+        if (cipherSuites != null) sslEngine.setEnabledCipherSuites(cipherSuites);
+        if (enabledProtocols != null) sslEngine.setEnabledProtocols(enabledProtocols);
+
+        if (mode == Mode.SERVER) {
+            sslEngine.setUseClientMode(false);
+            switch (sslClientAuth) {
+                case REQUIRED:
+                    sslEngine.setNeedClientAuth(true);
+                    break;
+                case REQUESTED:
+                    sslEngine.setWantClientAuth(true);
+                    break;
+                case NONE:
+                    break;
+            }
+            sslEngine.setUseClientMode(false);
+        } else {
+            sslEngine.setUseClientMode(true);
+            SSLParameters sslParams = sslEngine.getSSLParameters();
+            // SSLParameters#setEndpointIdentificationAlgorithm enables endpoint validation
+            // only in client mode. Hence, validation is enabled only for clients.
+            sslParams.setEndpointIdentificationAlgorithm(endpointIdentification);
+            sslEngine.setSSLParameters(sslParams);
+        }
+        return sslEngine;
+    }
+
+    public SSLContext sslContext() {
+        return sslContext;
+    }
+
+    /**
+     * Returns true if this SslEngineBuilder needs to be rebuilt.
+     *
+     * @param nextConfigs       The configuration we want to use.
+     * @return                  True only if this builder should be rebuilt.
+     */
+    public boolean shouldBeRebuilt(Map nextConfigs) {
+        if (!nextConfigs.equals(configs)) {
+            return true;
+        }
+        if (truststore != null && truststore.modified()) {
+            return true;
+        }
+        if (keystore != null && keystore.modified()) {
+            return true;
+        }
+        return false;
+    }
+
+    // package access for testing
+    static class SecurityStore {
+        private final String type;
+        private final String path;
+        private final Password password;
+        private final Password keyPassword;
+        private final Long fileLastModifiedMs;
+
+        SecurityStore(String type, String path, Password password, Password keyPassword) {
+            Objects.requireNonNull(type, "type must not be null");
+            this.type = type;
+            this.path = path;
+            this.password = password;
+            this.keyPassword = keyPassword;
+            fileLastModifiedMs = lastModifiedMs(path);
+        }
+
+        /**
+         * Loads this keystore
+         * @return the keystore
+         * @throws KafkaException if the file could not be read or if the keystore could not be loaded
+         *   using the specified configs (e.g. if the password or keystore type is invalid)
+         */
+        KeyStore load() {
+            try (InputStream in = Files.newInputStream(Paths.get(path))) {
+                KeyStore ks = KeyStore.getInstance(type);
+                // If a password is not set access to the truststore is still available, but integrity checking is disabled.
+                char[] passwordChars = password != null ? password.value().toCharArray() : null;
+                ks.load(in, passwordChars);
+                return ks;
+            } catch (GeneralSecurityException | IOException e) {
+                throw new KafkaException("Failed to load SSL keystore " + path + " of type " + type, e);
+            }
+        }
+
+        private Long lastModifiedMs(String path) {
+            try {
+                return Files.getLastModifiedTime(Paths.get(path)).toMillis();
+            } catch (IOException e) {
+                log.error("Modification time of key store could not be obtained: " + path, e);
+                return null;
+            }
+        }
+
+        boolean modified() {
+            Long modifiedMs = lastModifiedMs(path);
+            return modifiedMs != null && !Objects.equals(modifiedMs, this.fileLastModifiedMs);
+        }
+
+        @Override
+        public String toString() {
+            return "SecurityStore(" +
+                    "path=" + path +
+                    ", modificationTime=" + (fileLastModifiedMs == null ? null : new Date(fileLastModifiedMs)) + ")";
+        }
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/ssl/SslFactory.java b/clients/src/main/java/org/apache/kafka/common/security/ssl/SslFactory.java
index b9b52037c52a0..a338138eb51be 100644
--- a/clients/src/main/java/org/apache/kafka/common/security/ssl/SslFactory.java
+++ b/clients/src/main/java/org/apache/kafka/common/security/ssl/SslFactory.java
@@ -19,38 +19,29 @@
 import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.Reconfigurable;
 import org.apache.kafka.common.config.ConfigException;
+import org.apache.kafka.common.config.SecurityConfig;
 import org.apache.kafka.common.config.SslConfigs;
 import org.apache.kafka.common.config.internals.BrokerSecurityConfigs;
 import org.apache.kafka.common.network.Mode;
-import org.apache.kafka.common.config.types.Password;
 import org.apache.kafka.common.utils.Utils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import javax.net.ssl.KeyManager;
-import javax.net.ssl.KeyManagerFactory;
 import javax.net.ssl.SSLContext;
 import javax.net.ssl.SSLEngine;
 import javax.net.ssl.SSLEngineResult;
 import javax.net.ssl.SSLException;
-import javax.net.ssl.SSLParameters;
-import javax.net.ssl.TrustManagerFactory;
-import java.io.IOException;
-import java.io.InputStream;
 import java.nio.ByteBuffer;
-import java.nio.file.Files;
-import java.nio.file.Paths;
+import java.security.cert.Certificate;
+import java.security.cert.X509Certificate;
 import java.security.GeneralSecurityException;
 import java.security.KeyStore;
 import java.security.Principal;
-import java.security.SecureRandom;
-import java.security.cert.Certificate;
-import java.security.cert.X509Certificate;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
-import java.util.Date;
 import java.util.Enumeration;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -63,26 +54,25 @@ public class SslFactory implements Reconfigurable {
     private final Mode mode;
     private final String clientAuthConfigOverride;
     private final boolean keystoreVerifiableUsingTruststore;
-
-    private String protocol;
-    private String provider;
-    private String kmfAlgorithm;
-    private String tmfAlgorithm;
-    private SecurityStore keystore = null;
-    private SecurityStore truststore;
-    private String[] cipherSuites;
-    private String[] enabledProtocols;
     private String endpointIdentification;
-    private SecureRandom secureRandomImplementation;
-    private SSLContext sslContext;
-    private boolean needClientAuth;
-    private boolean wantClientAuth;
+    private SslEngineBuilder sslEngineBuilder;
 
     public SslFactory(Mode mode) {
         this(mode, null, false);
     }
 
-    public SslFactory(Mode mode, String clientAuthConfigOverride, boolean keystoreVerifiableUsingTruststore) {
+    /**
+     * Create an SslFactory.
+     *
+     * @param mode                                  Whether to use client or server mode.
+     * @param clientAuthConfigOverride              The value to override ssl.client.auth with, or null
+     *                                              if we don't want to override it.
+     * @param keystoreVerifiableUsingTruststore     True if we should require the keystore to be verifiable
+     *                                              using the truststore.
+     */
+    public SslFactory(Mode mode,
+                      String clientAuthConfigOverride,
+                      boolean keystoreVerifiableUsingTruststore) {
         this.mode = mode;
         this.clientAuthConfigOverride = clientAuthConfigOverride;
         this.keystoreVerifiableUsingTruststore = keystoreVerifiableUsingTruststore;
@@ -90,58 +80,28 @@ public SslFactory(Mode mode, String clientAuthConfigOverride, boolean keystoreVe
 
     @Override
     public void configure(Map configs) throws KafkaException {
-        this.protocol =  (String) configs.get(SslConfigs.SSL_PROTOCOL_CONFIG);
-        this.provider = (String) configs.get(SslConfigs.SSL_PROVIDER_CONFIG);
-
-        @SuppressWarnings("unchecked")
-        List cipherSuitesList = (List) configs.get(SslConfigs.SSL_CIPHER_SUITES_CONFIG);
-        if (cipherSuitesList != null && !cipherSuitesList.isEmpty())
-            this.cipherSuites = cipherSuitesList.toArray(new String[cipherSuitesList.size()]);
-
-        @SuppressWarnings("unchecked")
-        List enabledProtocolsList = (List) configs.get(SslConfigs.SSL_ENABLED_PROTOCOLS_CONFIG);
-        if (enabledProtocolsList != null && !enabledProtocolsList.isEmpty())
-            this.enabledProtocols = enabledProtocolsList.toArray(new String[enabledProtocolsList.size()]);
-
-        String endpointIdentification = (String) configs.get(SslConfigs.SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_CONFIG);
-        if (endpointIdentification != null)
-            this.endpointIdentification = endpointIdentification;
-
-        String secureRandomImplementation = (String) configs.get(SslConfigs.SSL_SECURE_RANDOM_IMPLEMENTATION_CONFIG);
-        if (secureRandomImplementation != null) {
-            try {
-                this.secureRandomImplementation = SecureRandom.getInstance(secureRandomImplementation);
-            } catch (GeneralSecurityException e) {
-                throw new KafkaException(e);
-            }
+        if (sslEngineBuilder != null) {
+            throw new IllegalStateException("SslFactory was already configured.");
         }
-
-        String clientAuthConfig = clientAuthConfigOverride;
-        if (clientAuthConfig == null)
-            clientAuthConfig = (String) configs.get(BrokerSecurityConfigs.SSL_CLIENT_AUTH_CONFIG);
-        if (clientAuthConfig != null) {
-            if (clientAuthConfig.equals("required"))
-                this.needClientAuth = true;
-            else if (clientAuthConfig.equals("requested"))
-                this.wantClientAuth = true;
+        this.endpointIdentification = (String) configs.get(SslConfigs.SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_CONFIG);
+
+        Map nextConfigs = new HashMap<>();
+        copyMapEntries(nextConfigs, configs, SslConfigs.NON_RECONFIGURABLE_CONFIGS);
+        copyMapEntries(nextConfigs, configs, SslConfigs.RECONFIGURABLE_CONFIGS);
+        copyMapEntry(nextConfigs, configs, SecurityConfig.SECURITY_PROVIDERS_CONFIG);
+        if (clientAuthConfigOverride != null) {
+            nextConfigs.put(BrokerSecurityConfigs.SSL_CLIENT_AUTH_CONFIG, clientAuthConfigOverride);
         }
-
-        this.kmfAlgorithm = (String) configs.get(SslConfigs.SSL_KEYMANAGER_ALGORITHM_CONFIG);
-        this.tmfAlgorithm = (String) configs.get(SslConfigs.SSL_TRUSTMANAGER_ALGORITHM_CONFIG);
-
-        this.keystore = createKeystore((String) configs.get(SslConfigs.SSL_KEYSTORE_TYPE_CONFIG),
-                       (String) configs.get(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG),
-                       (Password) configs.get(SslConfigs.SSL_KEYSTORE_PASSWORD_CONFIG),
-                       (Password) configs.get(SslConfigs.SSL_KEY_PASSWORD_CONFIG));
-
-        this.truststore = createTruststore((String) configs.get(SslConfigs.SSL_TRUSTSTORE_TYPE_CONFIG),
-                         (String) configs.get(SslConfigs.SSL_TRUSTSTORE_LOCATION_CONFIG),
-                         (Password) configs.get(SslConfigs.SSL_TRUSTSTORE_PASSWORD_CONFIG));
-        try {
-            this.sslContext = createSSLContext(keystore, truststore);
-        } catch (Exception e) {
-            throw new KafkaException(e);
+        SslEngineBuilder builder = new SslEngineBuilder(nextConfigs);
+        if (keystoreVerifiableUsingTruststore) {
+            try {
+                SslEngineValidator.validate(builder, builder);
+            } catch (Exception e) {
+                throw new ConfigException("A client SSLEngine created with the provided settings " +
+                        "can't connect to a server SSLEngine created with those settings.", e);
+            }
         }
+        this.sslEngineBuilder = builder;
     }
 
     @Override
@@ -150,216 +110,156 @@ public Set reconfigurableConfigs() {
     }
 
     @Override
-    public void validateReconfiguration(Map configs) {
-        try {
-            SecurityStore newKeystore = maybeCreateNewKeystore(configs);
-            SecurityStore newTruststore = maybeCreateNewTruststore(configs);
-            if (newKeystore != null || newTruststore != null) {
-                SecurityStore keystore = newKeystore != null ? newKeystore : this.keystore;
-                SecurityStore truststore = newTruststore != null ? newTruststore : this.truststore;
-                createSSLContext(keystore, truststore);
-            }
-        } catch (Exception e) {
-            throw new ConfigException("Validation of dynamic config update failed", e);
-        }
+    public void validateReconfiguration(Map newConfigs) {
+        createNewSslEngineBuilder(newConfigs);
     }
 
     @Override
-    public void reconfigure(Map configs) throws KafkaException {
-        SecurityStore newKeystore = maybeCreateNewKeystore(configs);
-        SecurityStore newTruststore = maybeCreateNewTruststore(configs);
-        if (newKeystore != null || newTruststore != null) {
-            try {
-                SecurityStore keystore = newKeystore != null ? newKeystore : this.keystore;
-                SecurityStore truststore = newTruststore != null ? newTruststore : this.truststore;
-                this.sslContext = createSSLContext(keystore, truststore);
-                this.keystore = keystore;
-                this.truststore = truststore;
-            } catch (Exception e) {
-                throw new ConfigException("Reconfiguration of SSL keystore/truststore failed", e);
-            }
+    public void reconfigure(Map newConfigs) throws KafkaException {
+        SslEngineBuilder newSslEngineBuilder = createNewSslEngineBuilder(newConfigs);
+        if (newSslEngineBuilder != this.sslEngineBuilder) {
+            this.sslEngineBuilder = newSslEngineBuilder;
+            log.info("Created new {} SSL engine builder with keystore {} truststore {}", mode,
+                    newSslEngineBuilder.keystore(), newSslEngineBuilder.truststore());
         }
     }
 
-    private SecurityStore maybeCreateNewKeystore(Map configs) {
-        boolean keystoreChanged = !Objects.equals(configs.get(SslConfigs.SSL_KEYSTORE_TYPE_CONFIG), keystore.type) ||
-                !Objects.equals(configs.get(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG), keystore.path) ||
-                !Objects.equals(configs.get(SslConfigs.SSL_KEYSTORE_PASSWORD_CONFIG), keystore.password) ||
-                !Objects.equals(configs.get(SslConfigs.SSL_KEY_PASSWORD_CONFIG), keystore.keyPassword);
-
-        if (!keystoreChanged) {
-            keystoreChanged = keystore.modified();
+    private SslEngineBuilder createNewSslEngineBuilder(Map newConfigs) {
+        if (sslEngineBuilder == null) {
+            throw new IllegalStateException("SslFactory has not been configured.");
         }
-        if (keystoreChanged) {
-            return createKeystore((String) configs.get(SslConfigs.SSL_KEYSTORE_TYPE_CONFIG),
-                    (String) configs.get(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG),
-                    (Password) configs.get(SslConfigs.SSL_KEYSTORE_PASSWORD_CONFIG),
-                    (Password) configs.get(SslConfigs.SSL_KEY_PASSWORD_CONFIG));
-        } else
-            return null;
-    }
-
-    private SecurityStore maybeCreateNewTruststore(Map configs) {
-        boolean truststoreChanged = !Objects.equals(configs.get(SslConfigs.SSL_TRUSTSTORE_TYPE_CONFIG), truststore.type) ||
-                !Objects.equals(configs.get(SslConfigs.SSL_TRUSTSTORE_LOCATION_CONFIG), truststore.path) ||
-                !Objects.equals(configs.get(SslConfigs.SSL_TRUSTSTORE_PASSWORD_CONFIG), truststore.password);
-
-        if (!truststoreChanged) {
-            truststoreChanged = truststore.modified();
+        Map nextConfigs = new HashMap<>(sslEngineBuilder.configs());
+        copyMapEntries(nextConfigs, newConfigs, SslConfigs.RECONFIGURABLE_CONFIGS);
+        if (clientAuthConfigOverride != null) {
+            nextConfigs.put(BrokerSecurityConfigs.SSL_CLIENT_AUTH_CONFIG, clientAuthConfigOverride);
         }
-        if (truststoreChanged) {
-            return createTruststore((String) configs.get(SslConfigs.SSL_TRUSTSTORE_TYPE_CONFIG),
-                    (String) configs.get(SslConfigs.SSL_TRUSTSTORE_LOCATION_CONFIG),
-                    (Password) configs.get(SslConfigs.SSL_TRUSTSTORE_PASSWORD_CONFIG));
-        } else
-            return null;
-    }
-
-    // package access for testing
-    SSLContext createSSLContext(SecurityStore keystore, SecurityStore truststore) throws GeneralSecurityException, IOException  {
-        SSLContext sslContext;
-        if (provider != null)
-            sslContext = SSLContext.getInstance(protocol, provider);
-        else
-            sslContext = SSLContext.getInstance(protocol);
-
-        KeyManager[] keyManagers = null;
-        if (keystore != null) {
-            String kmfAlgorithm = this.kmfAlgorithm != null ? this.kmfAlgorithm : KeyManagerFactory.getDefaultAlgorithm();
-            KeyManagerFactory kmf = KeyManagerFactory.getInstance(kmfAlgorithm);
-            KeyStore ks = keystore.load();
-            Password keyPassword = keystore.keyPassword != null ? keystore.keyPassword : keystore.password;
-            kmf.init(ks, keyPassword.value().toCharArray());
-            keyManagers = kmf.getKeyManagers();
+        if (!sslEngineBuilder.shouldBeRebuilt(nextConfigs)) {
+            return sslEngineBuilder;
         }
-
-        String tmfAlgorithm = this.tmfAlgorithm != null ? this.tmfAlgorithm : TrustManagerFactory.getDefaultAlgorithm();
-        TrustManagerFactory tmf = TrustManagerFactory.getInstance(tmfAlgorithm);
-        KeyStore ts = truststore == null ? null : truststore.load();
-        tmf.init(ts);
-
-        sslContext.init(keyManagers, tmf.getTrustManagers(), this.secureRandomImplementation);
-        boolean verifyKeystore = keystore != null && keystore != this.keystore;
-        boolean verifyTruststore = truststore != null && truststore != this.truststore;
-        if (verifyKeystore || verifyTruststore) {
-            if (this.keystore == null)
-                throw new ConfigException("Cannot add SSL keystore to an existing listener for which no keystore was configured.");
-            if (keystoreVerifiableUsingTruststore) {
-                SSLConfigValidatorEngine.validate(this, sslContext, this.sslContext);
-                SSLConfigValidatorEngine.validate(this, this.sslContext, sslContext);
+        try {
+            SslEngineBuilder newSslEngineBuilder = new SslEngineBuilder(nextConfigs);
+            if (sslEngineBuilder.keystore() == null) {
+                if (newSslEngineBuilder.keystore() != null) {
+                    throw new ConfigException("Cannot add SSL keystore to an existing listener for " +
+                            "which no keystore was configured.");
+                }
+            } else {
+                if (newSslEngineBuilder.keystore() == null) {
+                    throw new ConfigException("Cannot remove the SSL keystore from an existing listener for " +
+                            "which a keystore was configured.");
+                }
+                if (!CertificateEntries.create(sslEngineBuilder.keystore().load()).equals(
+                        CertificateEntries.create(newSslEngineBuilder.keystore().load()))) {
+                    throw new ConfigException("Keystore DistinguishedName or SubjectAltNames do not match");
+                }
             }
-            if (verifyKeystore &&
-                    !CertificateEntries.create(this.keystore.load()).equals(CertificateEntries.create(keystore.load()))) {
-                throw new ConfigException("Keystore DistinguishedName or SubjectAltNames do not match");
+            if (sslEngineBuilder.truststore() == null && newSslEngineBuilder.truststore() != null) {
+                throw new ConfigException("Cannot add SSL truststore to an existing listener for which no " +
+                        "truststore was configured.");
+            }
+            if (keystoreVerifiableUsingTruststore) {
+                if (sslEngineBuilder.truststore() != null || sslEngineBuilder.keystore() != null) {
+                    SslEngineValidator.validate(sslEngineBuilder, newSslEngineBuilder);
+                }
             }
+            return newSslEngineBuilder;
+        } catch (Exception e) {
+            log.debug("Validation of dynamic config update of SSLFactory failed.", e);
+            throw new ConfigException("Validation of dynamic config update of SSLFactory failed: " + e);
         }
-        return sslContext;
     }
 
     public SSLEngine createSslEngine(String peerHost, int peerPort) {
-        return createSslEngine(sslContext, peerHost, peerPort);
+        if (sslEngineBuilder == null) {
+            throw new IllegalStateException("SslFactory has not been configured.");
+        }
+        return sslEngineBuilder.createSslEngine(mode, peerHost, peerPort, endpointIdentification);
     }
 
-    private SSLEngine createSslEngine(SSLContext sslContext, String peerHost, int peerPort) {
-        SSLEngine sslEngine = sslContext.createSSLEngine(peerHost, peerPort);
-        if (cipherSuites != null) sslEngine.setEnabledCipherSuites(cipherSuites);
-        if (enabledProtocols != null) sslEngine.setEnabledProtocols(enabledProtocols);
+    @Deprecated
+    public SSLContext sslContext() {
+        return sslEngineBuilder.sslContext();
+    }
 
-        // SSLParameters#setEndpointIdentificationAlgorithm enables endpoint validation
-        // only in client mode. Hence, validation is enabled only for clients.
-        if (mode == Mode.SERVER) {
-            sslEngine.setUseClientMode(false);
-            if (needClientAuth)
-                sslEngine.setNeedClientAuth(needClientAuth);
-            else
-                sslEngine.setWantClientAuth(wantClientAuth);
-        } else {
-            sslEngine.setUseClientMode(true);
-            SSLParameters sslParams = sslEngine.getSSLParameters();
-            sslParams.setEndpointIdentificationAlgorithm(endpointIdentification);
-            sslEngine.setSSLParameters(sslParams);
-        }
-        return sslEngine;
+    public SslEngineBuilder sslEngineBuilder() {
+        return sslEngineBuilder;
     }
 
     /**
-     * Returns a configured SSLContext.
-     * @return SSLContext.
+     * Copy entries from one map into another.
+     *
+     * @param destMap   The map to copy entries into.
+     * @param srcMap    The map to copy entries from.
+     * @param keySet    Only entries with these keys will be copied.
+     * @param        The map key type.
+     * @param        The map value type.
      */
-    public SSLContext sslContext() {
-        return sslContext;
-    }
-
-    private SecurityStore createKeystore(String type, String path, Password password, Password keyPassword) {
-        if (path == null && password != null) {
-            throw new KafkaException("SSL key store is not specified, but key store password is specified.");
-        } else if (path != null && password == null) {
-            throw new KafkaException("SSL key store is specified, but key store password is not specified.");
-        } else if (path != null && password != null) {
-            return new SecurityStore(type, path, password, keyPassword);
-        } else
-            return null; // path == null, clients may use this path with brokers that don't require client auth
+    private static  void copyMapEntries(Map destMap,
+                                              Map srcMap,
+                                              Set keySet) {
+        for (K k : keySet) {
+            copyMapEntry(destMap, srcMap, k);
+        }
     }
 
-    private SecurityStore createTruststore(String type, String path, Password password) {
-        if (path == null && password != null) {
-            throw new KafkaException("SSL trust store is not specified, but trust store password is specified.");
-        } else if (path != null) {
-            return new SecurityStore(type, path, password, null);
-        } else
-            return null;
+    /**
+     * Copy entry from one map into another.
+     *
+     * @param destMap   The map to copy entries into.
+     * @param srcMap    The map to copy entries from.
+     * @param key       The entry with this key will be copied
+     * @param        The map key type.
+     * @param        The map value type.
+     */
+    private static  void copyMapEntry(Map destMap,
+                                            Map srcMap,
+                                            K key) {
+        if (srcMap.containsKey(key)) {
+            destMap.put(key, srcMap.get(key));
+        }
     }
 
-    // package access for testing
-    static class SecurityStore {
-        private final String type;
-        private final String path;
-        private final Password password;
-        private final Password keyPassword;
-        private Long fileLastModifiedMs;
+    static class CertificateEntries {
+        private final Principal subjectPrincipal;
+        private final Set> subjectAltNames;
 
-        SecurityStore(String type, String path, Password password, Password keyPassword) {
-            Objects.requireNonNull(type, "type must not be null");
-            this.type = type;
-            this.path = path;
-            this.password = password;
-            this.keyPassword = keyPassword;
+        static List create(KeyStore keystore) throws GeneralSecurityException {
+            Enumeration aliases = keystore.aliases();
+            List entries = new ArrayList<>();
+            while (aliases.hasMoreElements()) {
+                String alias = aliases.nextElement();
+                Certificate cert  = keystore.getCertificate(alias);
+                if (cert instanceof X509Certificate)
+                    entries.add(new CertificateEntries((X509Certificate) cert));
+            }
+            return entries;
         }
 
-        /**
-         * Loads this keystore
-         * @return the keystore
-         * @throws KafkaException if the file could not be read or if the keystore could not be loaded
-         *   using the specified configs (e.g. if the password or keystore type is invalid)
-         */
-        KeyStore load() {
-            try (InputStream in = Files.newInputStream(Paths.get(path))) {
-                KeyStore ks = KeyStore.getInstance(type);
-                // If a password is not set access to the truststore is still available, but integrity checking is disabled.
-                char[] passwordChars = password != null ? password.value().toCharArray() : null;
-                ks.load(in, passwordChars);
-                fileLastModifiedMs = lastModifiedMs(path);
+        CertificateEntries(X509Certificate cert) throws GeneralSecurityException {
+            this.subjectPrincipal = cert.getSubjectX500Principal();
+            Collection> altNames = cert.getSubjectAlternativeNames();
+            // use a set for comparison
+            this.subjectAltNames = altNames != null ? new HashSet<>(altNames) : Collections.emptySet();
+        }
 
-                log.debug("Loaded key store with path {} modification time {}", path,
-                        fileLastModifiedMs == null ? null : new Date(fileLastModifiedMs));
-                return ks;
-            } catch (GeneralSecurityException | IOException e) {
-                throw new KafkaException("Failed to load SSL keystore " + path + " of type " + type, e);
-            }
+        @Override
+        public int hashCode() {
+            return Objects.hash(subjectPrincipal, subjectAltNames);
         }
 
-        private Long lastModifiedMs(String path) {
-            try {
-                return Files.getLastModifiedTime(Paths.get(path)).toMillis();
-            } catch (IOException e) {
-                log.error("Modification time of key store could not be obtained: " + path, e);
-                return null;
-            }
+        @Override
+        public boolean equals(Object obj) {
+            if (!(obj instanceof CertificateEntries))
+                return false;
+            CertificateEntries other = (CertificateEntries) obj;
+            return Objects.equals(subjectPrincipal, other.subjectPrincipal) &&
+                    Objects.equals(subjectAltNames, other.subjectAltNames);
         }
 
-        boolean modified() {
-            Long modifiedMs = lastModifiedMs(path);
-            return modifiedMs != null && !Objects.equals(modifiedMs, this.fileLastModifiedMs);
+        @Override
+        public String toString() {
+            return "subjectPrincipal=" + subjectPrincipal +
+                    ", subjectAltNames=" + subjectAltNames;
         }
     }
 
@@ -368,32 +268,44 @@ boolean modified() {
      * The validator checks that a successful handshake can be performed using the keystore and
      * truststore configured on this SslFactory.
      */
-    private static class SSLConfigValidatorEngine {
+    private static class SslEngineValidator {
         private static final ByteBuffer EMPTY_BUF = ByteBuffer.allocate(0);
         private final SSLEngine sslEngine;
         private SSLEngineResult handshakeResult;
         private ByteBuffer appBuffer;
         private ByteBuffer netBuffer;
 
-        static void validate(SslFactory sslFactory, SSLContext clientSslContext, SSLContext serverSslContext) throws SSLException {
-            SSLConfigValidatorEngine clientEngine = new SSLConfigValidatorEngine(sslFactory, clientSslContext, Mode.CLIENT);
-            SSLConfigValidatorEngine serverEngine = new SSLConfigValidatorEngine(sslFactory, serverSslContext, Mode.SERVER);
+        static void validate(SslEngineBuilder oldEngineBuilder,
+                             SslEngineBuilder newEngineBuilder) throws SSLException {
+            validate(createSslEngineForValidation(oldEngineBuilder, Mode.SERVER),
+                    createSslEngineForValidation(newEngineBuilder, Mode.CLIENT));
+            validate(createSslEngineForValidation(newEngineBuilder, Mode.SERVER),
+                    createSslEngineForValidation(oldEngineBuilder, Mode.CLIENT));
+        }
+
+        private static SSLEngine createSslEngineForValidation(SslEngineBuilder sslEngineBuilder, Mode mode) {
+            // Use empty hostname, disable hostname verification
+            return sslEngineBuilder.createSslEngine(mode, "", 0, "");
+        }
+
+        static void validate(SSLEngine clientEngine, SSLEngine serverEngine) throws SSLException {
+            SslEngineValidator clientValidator = new SslEngineValidator(clientEngine);
+            SslEngineValidator serverValidator = new SslEngineValidator(serverEngine);
             try {
-                clientEngine.beginHandshake();
-                serverEngine.beginHandshake();
-                while (!serverEngine.complete() || !clientEngine.complete()) {
-                    clientEngine.handshake(serverEngine);
-                    serverEngine.handshake(clientEngine);
+                clientValidator.beginHandshake();
+                serverValidator.beginHandshake();
+                while (!serverValidator.complete() || !clientValidator.complete()) {
+                    clientValidator.handshake(serverValidator);
+                    serverValidator.handshake(clientValidator);
                 }
             } finally {
-                clientEngine.close();
-                serverEngine.close();
+                clientValidator.close();
+                serverValidator.close();
             }
         }
 
-        private SSLConfigValidatorEngine(SslFactory sslFactory, SSLContext sslContext, Mode mode) {
-            this.sslEngine = sslFactory.createSslEngine(sslContext, "localhost", 0); // these hints are not used for validation
-            sslEngine.setUseClientMode(mode == Mode.CLIENT);
+        private SslEngineValidator(SSLEngine engine) {
+            this.sslEngine = engine;
             appBuffer = ByteBuffer.allocate(sslEngine.getSession().getApplicationBufferSize());
             netBuffer = ByteBuffer.allocate(sslEngine.getSession().getPacketBufferSize());
         }
@@ -401,8 +313,7 @@ private SSLConfigValidatorEngine(SslFactory sslFactory, SSLContext sslContext, M
         void beginHandshake() throws SSLException {
             sslEngine.beginHandshake();
         }
-
-        void handshake(SSLConfigValidatorEngine peerEngine) throws SSLException {
+        void handshake(SslEngineValidator peerValidator) throws SSLException {
             SSLEngineResult.HandshakeStatus handshakeStatus = sslEngine.getHandshakeStatus();
             while (true) {
                 switch (handshakeStatus) {
@@ -422,11 +333,11 @@ void handshake(SSLConfigValidatorEngine peerEngine) throws SSLException {
                         }
                         return;
                     case NEED_UNWRAP:
-                        if (peerEngine.netBuffer.position() == 0) // no data to unwrap, return to process peer
+                        if (peerValidator.netBuffer.position() == 0) // no data to unwrap, return to process peer
                             return;
-                        peerEngine.netBuffer.flip(); // unwrap the data from peer
-                        handshakeResult = sslEngine.unwrap(peerEngine.netBuffer, appBuffer);
-                        peerEngine.netBuffer.compact();
+                        peerValidator.netBuffer.flip(); // unwrap the data from peer
+                        handshakeResult = sslEngine.unwrap(peerValidator.netBuffer, appBuffer);
+                        peerValidator.netBuffer.compact();
                         handshakeStatus = handshakeResult.getHandshakeStatus();
                         switch (handshakeResult.getStatus()) {
                             case OK: break;
@@ -471,48 +382,4 @@ void close() {
             }
         }
     }
-
-    static class CertificateEntries {
-        private final Principal subjectPrincipal;
-        private final Set> subjectAltNames;
-
-        static List create(KeyStore keystore) throws GeneralSecurityException {
-            Enumeration aliases = keystore.aliases();
-            List entries = new ArrayList<>();
-            while (aliases.hasMoreElements()) {
-                String alias = aliases.nextElement();
-                Certificate cert  = keystore.getCertificate(alias);
-                if (cert instanceof X509Certificate)
-                    entries.add(new CertificateEntries((X509Certificate) cert));
-            }
-            return entries;
-        }
-
-        CertificateEntries(X509Certificate cert) throws GeneralSecurityException {
-            this.subjectPrincipal = cert.getSubjectX500Principal();
-            Collection> altNames = cert.getSubjectAlternativeNames();
-            // use a set for comparison
-            this.subjectAltNames = altNames != null ? new HashSet<>(altNames) : Collections.emptySet();
-        }
-
-        @Override
-        public int hashCode() {
-            return Objects.hash(subjectPrincipal, subjectAltNames);
-        }
-
-        @Override
-        public boolean equals(Object obj) {
-            if (!(obj instanceof CertificateEntries))
-                return false;
-            CertificateEntries other = (CertificateEntries) obj;
-            return Objects.equals(subjectPrincipal, other.subjectPrincipal) &&
-                    Objects.equals(subjectAltNames, other.subjectAltNames);
-        }
-
-        @Override
-        public String toString() {
-            return "subjectPrincipal=" + subjectPrincipal +
-                    ", subjectAltNames=" + subjectAltNames;
-        }
-    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/security/ssl/SslPrincipalMapper.java b/clients/src/main/java/org/apache/kafka/common/security/ssl/SslPrincipalMapper.java
index 3b95e1a5f4d41..977fdc332cbfa 100644
--- a/clients/src/main/java/org/apache/kafka/common/security/ssl/SslPrincipalMapper.java
+++ b/clients/src/main/java/org/apache/kafka/common/security/ssl/SslPrincipalMapper.java
@@ -18,53 +18,44 @@
 
 import java.io.IOException;
 import java.util.List;
-import java.util.Collections;
 import java.util.ArrayList;
 import java.util.Locale;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
+import static org.apache.kafka.common.config.internals.BrokerSecurityConfigs.DEFAULT_SSL_PRINCIPAL_MAPPING_RULES;
+
 public class SslPrincipalMapper {
 
-    private static final Pattern RULE_PARSER = Pattern.compile("((DEFAULT)|(RULE:(([^/]*)/([^/]*))/([LU])?))");
+    private static final String RULE_PATTERN = "(DEFAULT)|RULE:((\\\\.|[^\\\\/])*)/((\\\\.|[^\\\\/])*)/([LU]?).*?|(.*?)";
+    private static final Pattern RULE_SPLITTER = Pattern.compile("\\s*(" + RULE_PATTERN + ")\\s*(,\\s*|$)");
+    private static final Pattern RULE_PARSER = Pattern.compile(RULE_PATTERN);
 
     private final List rules;
 
-    public SslPrincipalMapper(List sslPrincipalMappingRules) {
-        this.rules = sslPrincipalMappingRules;
+    public SslPrincipalMapper(String sslPrincipalMappingRules) {
+        this.rules = parseRules(splitRules(sslPrincipalMappingRules));
     }
 
-    public static SslPrincipalMapper fromRules(List sslPrincipalMappingRules) {
-        List rules = sslPrincipalMappingRules == null ? Collections.singletonList("DEFAULT") : sslPrincipalMappingRules;
-        return new SslPrincipalMapper(parseRules(rules));
+    public static SslPrincipalMapper fromRules(String sslPrincipalMappingRules) {
+        return new SslPrincipalMapper(sslPrincipalMappingRules);
     }
 
-    private static List joinSplitRules(List rules) {
-        String rule = "RULE:";
-        String defaultRule = "DEFAULT";
-        List retVal = new ArrayList<>();
-        StringBuilder currentRule = new StringBuilder();
-        for (String r : rules) {
-            if (currentRule.length() > 0) {
-                if (r.startsWith(rule) || r.equals(defaultRule)) {
-                    retVal.add(currentRule.toString());
-                    currentRule.setLength(0);
-                    currentRule.append(r);
-                } else {
-                    currentRule.append(String.format(",%s", r));
-                }
-            } else {
-                currentRule.append(r);
-            }
+    private static List splitRules(String sslPrincipalMappingRules) {
+        if (sslPrincipalMappingRules == null) {
+            sslPrincipalMappingRules = DEFAULT_SSL_PRINCIPAL_MAPPING_RULES;
         }
-        if (currentRule.length() > 0) {
-            retVal.add(currentRule.toString());
+
+        List result = new ArrayList<>();
+        Matcher matcher = RULE_SPLITTER.matcher(sslPrincipalMappingRules.trim());
+        while (matcher.find()) {
+            result.add(matcher.group(1));
         }
-        return retVal;
+
+        return result;
     }
 
     private static List parseRules(List rules) {
-        rules = joinSplitRules(rules);
         List result = new ArrayList<>();
         for (String rule : rules) {
             Matcher matcher = RULE_PARSER.matcher(rule);
@@ -74,15 +65,18 @@ private static List parseRules(List rules) {
             if (rule.length() != matcher.end()) {
                 throw new IllegalArgumentException("Invalid rule: `" + rule + "`, unmatched substring: `" + rule.substring(matcher.end()) + "`");
             }
-            if (matcher.group(2) != null) {
+
+            // empty rules are ignored
+            if (matcher.group(1) != null) {
                 result.add(new Rule());
-            } else {
-                result.add(new Rule(matcher.group(5),
-                                    matcher.group(6),
-                                    "L".equals(matcher.group(7)),
-                                    "U".equals(matcher.group(7))));
+            } else if (matcher.group(2) != null) {
+                result.add(new Rule(matcher.group(2),
+                                    matcher.group(4),
+                                    "L".equals(matcher.group(6)),
+                                    "U".equals(matcher.group(6))));
             }
         }
+
         return result;
     }
 
diff --git a/clients/src/main/java/org/apache/kafka/common/security/token/delegation/DelegationToken.java b/clients/src/main/java/org/apache/kafka/common/security/token/delegation/DelegationToken.java
index a1e2372d7a028..b389a199a9201 100644
--- a/clients/src/main/java/org/apache/kafka/common/security/token/delegation/DelegationToken.java
+++ b/clients/src/main/java/org/apache/kafka/common/security/token/delegation/DelegationToken.java
@@ -20,6 +20,7 @@
 
 import java.util.Arrays;
 import java.util.Base64;
+import java.util.Objects;
 
 /**
  * A class representing a delegation token.
@@ -58,10 +59,7 @@ public boolean equals(Object o) {
 
         DelegationToken token = (DelegationToken) o;
 
-        if (tokenInformation != null ? !tokenInformation.equals(token.tokenInformation) : token.tokenInformation != null) {
-            return false;
-        }
-        return Arrays.equals(hmac, token.hmac);
+        return Objects.equals(tokenInformation, token.tokenInformation) && Arrays.equals(hmac, token.hmac);
     }
 
     @Override
diff --git a/clients/src/main/java/org/apache/kafka/common/security/token/delegation/TokenInformation.java b/clients/src/main/java/org/apache/kafka/common/security/token/delegation/TokenInformation.java
index 8f360ca153baf..9903eb51b235f 100644
--- a/clients/src/main/java/org/apache/kafka/common/security/token/delegation/TokenInformation.java
+++ b/clients/src/main/java/org/apache/kafka/common/security/token/delegation/TokenInformation.java
@@ -21,6 +21,7 @@
 
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.Objects;
 
 /**
  * A class representing a delegation token details.
@@ -113,27 +114,19 @@ public boolean equals(Object o) {
 
         TokenInformation that = (TokenInformation) o;
 
-        if (issueTimestamp != that.issueTimestamp) {
-            return false;
-        }
-        if (maxTimestamp != that.maxTimestamp) {
-            return false;
-        }
-        if (owner != null ? !owner.equals(that.owner) : that.owner != null) {
-            return false;
-        }
-        if (renewers != null ? !renewers.equals(that.renewers) : that.renewers != null) {
-            return false;
-        }
-        return tokenId != null ? tokenId.equals(that.tokenId) : that.tokenId == null;
+        return issueTimestamp == that.issueTimestamp &&
+            maxTimestamp == that.maxTimestamp &&
+            Objects.equals(owner, that.owner) &&
+            Objects.equals(renewers, that.renewers) &&
+            Objects.equals(tokenId, that.tokenId);
     }
 
     @Override
     public int hashCode() {
         int result = owner != null ? owner.hashCode() : 0;
         result = 31 * result + (renewers != null ? renewers.hashCode() : 0);
-        result = 31 * result + (int) (issueTimestamp ^ (issueTimestamp >>> 32));
-        result = 31 * result + (int) (maxTimestamp ^ (maxTimestamp >>> 32));
+        result = 31 * result + Long.hashCode(issueTimestamp);
+        result = 31 * result + Long.hashCode(maxTimestamp);
         result = 31 * result + (tokenId != null ? tokenId.hashCode() : 0);
         return result;
     }
diff --git a/clients/src/main/java/org/apache/kafka/common/utils/AppInfoParser.java b/clients/src/main/java/org/apache/kafka/common/utils/AppInfoParser.java
index 8a12fbc94af35..997ba7610340b 100644
--- a/clients/src/main/java/org/apache/kafka/common/utils/AppInfoParser.java
+++ b/clients/src/main/java/org/apache/kafka/common/utils/AppInfoParser.java
@@ -36,6 +36,8 @@ public class AppInfoParser {
     private static final String VERSION;
     private static final String COMMIT_ID;
 
+    protected static final String DEFAULT_VALUE = "unknown";
+
     static {
         Properties props = new Properties();
         try (InputStream resourceStream = AppInfoParser.class.getResourceAsStream("/kafka/kafka-version.properties")) {
@@ -43,8 +45,8 @@ public class AppInfoParser {
         } catch (Exception e) {
             log.warn("Error while loading kafka-version.properties: {}", e.getMessage());
         }
-        VERSION = props.getProperty("version", "unknown").trim();
-        COMMIT_ID = props.getProperty("commitId", "unknown").trim();
+        VERSION = props.getProperty("version", DEFAULT_VALUE).trim();
+        COMMIT_ID = props.getProperty("commitId", DEFAULT_VALUE).trim();
     }
 
     public static String getVersion() {
@@ -55,13 +57,14 @@ public static String getCommitId() {
         return COMMIT_ID;
     }
 
-    public static synchronized void registerAppInfo(String prefix, String id, Metrics metrics) {
+    public static synchronized void registerAppInfo(String prefix, String id, Metrics metrics, long nowMs) {
         try {
             ObjectName name = new ObjectName(prefix + ":type=app-info,id=" + Sanitizer.jmxSanitize(id));
-            AppInfo mBean = new AppInfo();
-            ManagementFactory.getPlatformMBeanServer().registerMBean(mBean, name);
-
-            registerMetrics(metrics); // prefix will be added later by JmxReporter
+            AppInfo mBean = new AppInfo(nowMs);
+            if (!ManagementFactory.getPlatformMBeanServer().isRegistered(name)) {
+                ManagementFactory.getPlatformMBeanServer().registerMBean(mBean, name);
+            }
+            registerMetrics(metrics, mBean); // prefix will be added later by JmxReporter
         } catch (JMException e) {
             log.warn("Error registering AppInfo mbean", e);
         }
@@ -84,10 +87,11 @@ private static MetricName metricName(Metrics metrics, String name) {
         return metrics.metricName(name, "app-info", "Metric indicating " + name);
     }
 
-    private static void registerMetrics(Metrics metrics) {
+    private static void registerMetrics(Metrics metrics, AppInfo appInfo) {
         if (metrics != null) {
-            metrics.addMetric(metricName(metrics, "version"), new ImmutableValue<>(VERSION));
-            metrics.addMetric(metricName(metrics, "commit-id"), new ImmutableValue<>(COMMIT_ID));
+            metrics.addMetric(metricName(metrics, "version"), new ImmutableValue<>(appInfo.getVersion()));
+            metrics.addMetric(metricName(metrics, "commit-id"), new ImmutableValue<>(appInfo.getCommitId()));
+            metrics.addMetric(metricName(metrics, "start-time-ms"), new ImmutableValue<>(appInfo.getStartTimeMs()));
         }
     }
 
@@ -95,19 +99,25 @@ private static void unregisterMetrics(Metrics metrics) {
         if (metrics != null) {
             metrics.removeMetric(metricName(metrics, "version"));
             metrics.removeMetric(metricName(metrics, "commit-id"));
+            metrics.removeMetric(metricName(metrics, "start-time-ms"));
         }
     }
 
     public interface AppInfoMBean {
         String getVersion();
         String getCommitId();
+        Long getStartTimeMs();
     }
 
     public static class AppInfo implements AppInfoMBean {
 
-        public AppInfo() {
+        private final Long startTimeMs;
+
+        public AppInfo(long startTimeMs) {
+            this.startTimeMs = startTimeMs;
             log.info("Kafka version: {}", AppInfoParser.getVersion());
             log.info("Kafka commitId: {}", AppInfoParser.getCommitId());
+            log.info("Kafka startTimeMs: {}", startTimeMs);
         }
 
         @Override
@@ -120,6 +130,11 @@ public String getCommitId() {
             return AppInfoParser.getCommitId();
         }
 
+        @Override
+        public Long getStartTimeMs() {
+            return startTimeMs;
+        }
+
     }
 
     static class ImmutableValue implements Gauge {
diff --git a/clients/src/main/java/org/apache/kafka/common/utils/ByteUtils.java b/clients/src/main/java/org/apache/kafka/common/utils/ByteUtils.java
index 6efe3112cdb43..ae93075065710 100644
--- a/clients/src/main/java/org/apache/kafka/common/utils/ByteUtils.java
+++ b/clients/src/main/java/org/apache/kafka/common/utils/ByteUtils.java
@@ -129,7 +129,7 @@ public static void writeUnsignedIntLE(byte[] buffer, int offset, int value) {
     }
 
     /**
-     * Read an integer stored in variable-length format using zig-zag decoding from
+     * Read an integer stored in variable-length format using unsigned decoding from
      *  Google Protocol Buffers.
      *
      * @param buffer The buffer to read from
@@ -137,7 +137,7 @@ public static void writeUnsignedIntLE(byte[] buffer, int offset, int value) {
      *
      * @throws IllegalArgumentException if variable-length value does not terminate after 5 bytes have been read
      */
-    public static int readVarint(ByteBuffer buffer) {
+    public static int readUnsignedVarint(ByteBuffer buffer) {
         int value = 0;
         int i = 0;
         int b;
@@ -148,11 +148,11 @@ public static int readVarint(ByteBuffer buffer) {
                 throw illegalVarintException(value);
         }
         value |= b << i;
-        return (value >>> 1) ^ -(value & 1);
+        return value;
     }
 
     /**
-     * Read an integer stored in variable-length format using zig-zag decoding from
+     * Read an integer stored in variable-length format using unsigned decoding from
      *  Google Protocol Buffers.
      *
      * @param in The input to read from
@@ -161,7 +161,7 @@ public static int readVarint(ByteBuffer buffer) {
      * @throws IllegalArgumentException if variable-length value does not terminate after 5 bytes have been read
      * @throws IOException              if {@link DataInput} throws {@link IOException}
      */
-    public static int readVarint(DataInput in) throws IOException {
+    public static int readUnsignedVarint(DataInput in) throws IOException {
         int value = 0;
         int i = 0;
         int b;
@@ -172,6 +172,35 @@ public static int readVarint(DataInput in) throws IOException {
                 throw illegalVarintException(value);
         }
         value |= b << i;
+        return value;
+    }
+
+    /**
+     * Read an integer stored in variable-length format using zig-zag decoding from
+     *  Google Protocol Buffers.
+     *
+     * @param buffer The buffer to read from
+     * @return The integer read
+     *
+     * @throws IllegalArgumentException if variable-length value does not terminate after 5 bytes have been read
+     */
+    public static int readVarint(ByteBuffer buffer) {
+        int value = readUnsignedVarint(buffer);
+        return (value >>> 1) ^ -(value & 1);
+    }
+
+    /**
+     * Read an integer stored in variable-length format using zig-zag decoding from
+     *  Google Protocol Buffers.
+     *
+     * @param in The input to read from
+     * @return The integer read
+     *
+     * @throws IllegalArgumentException if variable-length value does not terminate after 5 bytes have been read
+     * @throws IOException              if {@link DataInput} throws {@link IOException}
+     */
+    public static int readVarint(DataInput in) throws IOException {
+        int value = readUnsignedVarint(in);
         return (value >>> 1) ^ -(value & 1);
     }
 
@@ -222,6 +251,40 @@ public static long readVarlong(ByteBuffer buffer)  {
         return (value >>> 1) ^ -(value & 1);
     }
 
+    /**
+     * Write the given integer following the variable-length unsigned encoding from
+     *  Google Protocol Buffers
+     * into the buffer.
+     *
+     * @param value The value to write
+     * @param buffer The output to write to
+     */
+    public static void writeUnsignedVarint(int value, ByteBuffer buffer) {
+        while ((value & 0xffffff80) != 0L) {
+            byte b = (byte) ((value & 0x7f) | 0x80);
+            buffer.put(b);
+            value >>>= 7;
+        }
+        buffer.put((byte) value);
+    }
+
+    /**
+     * Write the given integer following the variable-length unsigned encoding from
+     *  Google Protocol Buffers
+     * into the buffer.
+     *
+     * @param value The value to write
+     * @param out The output to write to
+     */
+    public static void writeUnsignedVarint(int value, DataOutput out) throws IOException {
+        while ((value & 0xffffff80) != 0L) {
+            byte b = (byte) ((value & 0x7f) | 0x80);
+            out.writeByte(b);
+            value >>>= 7;
+        }
+        out.writeByte((byte) value);
+    }
+
     /**
      * Write the given integer following the variable-length zig-zag encoding from
      *  Google Protocol Buffers
@@ -231,12 +294,7 @@ public static long readVarlong(ByteBuffer buffer)  {
      * @param out The output to write to
      */
     public static void writeVarint(int value, DataOutput out) throws IOException {
-        int v = (value << 1) ^ (value >> 31);
-        while ((v & 0xffffff80) != 0L) {
-            out.writeByte((v & 0x7f) | 0x80);
-            v >>>= 7;
-        }
-        out.writeByte((byte) v);
+        writeUnsignedVarint((value << 1) ^ (value >> 31), out);
     }
 
     /**
@@ -248,13 +306,7 @@ public static void writeVarint(int value, DataOutput out) throws IOException {
      * @param buffer The output to write to
      */
     public static void writeVarint(int value, ByteBuffer buffer) {
-        int v = (value << 1) ^ (value >> 31);
-        while ((v & 0xffffff80) != 0L) {
-            byte b = (byte) ((v & 0x7f) | 0x80);
-            buffer.put(b);
-            v >>>= 7;
-        }
-        buffer.put((byte) v);
+        writeUnsignedVarint((value << 1) ^ (value >> 31), buffer);
     }
 
     /**
@@ -293,20 +345,28 @@ public static void writeVarlong(long value, ByteBuffer buffer) {
     }
 
     /**
-     * Number of bytes needed to encode an integer in variable-length format.
+     * Number of bytes needed to encode an integer in unsigned variable-length format.
      *
      * @param value The signed value
      */
-    public static int sizeOfVarint(int value) {
-        int v = (value << 1) ^ (value >> 31);
+    public static int sizeOfUnsignedVarint(int value) {
         int bytes = 1;
-        while ((v & 0xffffff80) != 0L) {
+        while ((value & 0xffffff80) != 0L) {
             bytes += 1;
-            v >>>= 7;
+            value >>>= 7;
         }
         return bytes;
     }
 
+    /**
+     * Number of bytes needed to encode an integer in variable-length format.
+     *
+     * @param value The signed value
+     */
+    public static int sizeOfVarint(int value) {
+        return sizeOfUnsignedVarint((value << 1) ^ (value >> 31));
+    }
+
     /**
      * Number of bytes needed to encode a long in variable-length format.
      *
diff --git a/clients/src/main/java/org/apache/kafka/common/utils/Bytes.java b/clients/src/main/java/org/apache/kafka/common/utils/Bytes.java
index 19cd711928fc8..df754596a1940 100644
--- a/clients/src/main/java/org/apache/kafka/common/utils/Bytes.java
+++ b/clients/src/main/java/org/apache/kafka/common/utils/Bytes.java
@@ -140,6 +140,32 @@ private static String toString(final byte[] b, int off, int len) {
         return result.toString();
     }
 
+    /**
+     * Increment the underlying byte array by adding 1. Throws an IndexOutOfBoundsException if incrementing would cause
+     * the underlying input byte array to overflow.
+     *
+     * @param input - The byte array to increment
+     * @return A new copy of the incremented byte array.
+     */
+    public static Bytes increment(Bytes input) throws IndexOutOfBoundsException {
+        byte[] inputArr = input.get();
+        byte[] ret = new byte[inputArr.length];
+        int carry = 1;
+        for (int i = inputArr.length - 1; i >= 0; i--) {
+            if (inputArr[i] == (byte) 0xFF && carry == 1) {
+                ret[i] = (byte) 0x00;
+            } else {
+                ret[i] = (byte) (inputArr[i] + carry);
+                carry = 0;
+            }
+        }
+        if (carry == 0) {
+            return wrap(ret);
+        } else {
+            throw new IndexOutOfBoundsException();
+        }
+    }
+
     /**
      * A byte array comparator based on lexicograpic ordering.
      */
diff --git a/clients/src/main/java/org/apache/kafka/common/utils/CircularIterator.java b/clients/src/main/java/org/apache/kafka/common/utils/CircularIterator.java
index 3e7d5eb29fbcc..925f4adf9e9f7 100644
--- a/clients/src/main/java/org/apache/kafka/common/utils/CircularIterator.java
+++ b/clients/src/main/java/org/apache/kafka/common/utils/CircularIterator.java
@@ -14,22 +14,54 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+
 package org.apache.kafka.common.utils;
 
+import java.util.Collection;
+import java.util.ConcurrentModificationException;
 import java.util.Iterator;
-import java.util.List;
+import java.util.Objects;
 
+/**
+ * An iterator that cycles through the {@code Iterator} of a {@code Collection}
+ * indefinitely. Useful for tasks such as round-robin load balancing. This class
+ * does not provide thread-safe access. This {@code Iterator} supports
+ * {@code null} elements in the underlying {@code Collection}. This
+ * {@code Iterator} does not support any modification to the underlying
+ * {@code Collection} after it has been wrapped by this class. Changing the
+ * underlying {@code Collection} may cause a
+ * {@link ConcurrentModificationException} or some other undefined behavior.
+ */
 public class CircularIterator implements Iterator {
-    int i = 0;
-    private List list;
 
-    public CircularIterator(List list) {
-        if (list.isEmpty()) {
+    private final Iterable iterable;
+    private Iterator iterator;
+    private T nextValue;
+
+    /**
+     * Create a new instance of a CircularIterator. The ordering of this
+     * Iterator will be dictated by the Iterator returned by Collection itself.
+     *
+     * @param col The collection to iterate indefinitely
+     *
+     * @throws NullPointerException if col is {@code null}
+     * @throws IllegalArgumentException if col is empty.
+     */
+    public CircularIterator(final Collection col) {
+        this.iterable = Objects.requireNonNull(col);
+        this.iterator = col.iterator();
+        if (col.isEmpty()) {
             throw new IllegalArgumentException("CircularIterator can only be used on non-empty lists");
         }
-        this.list = list;
+        this.nextValue = advance();
     }
 
+    /**
+     * Returns true since the iteration will forever cycle through the provided
+     * {@code Collection}.
+     *
+     * @return Always true
+     */
     @Override
     public boolean hasNext() {
         return true;
@@ -37,13 +69,34 @@ public boolean hasNext() {
 
     @Override
     public T next() {
-        T next = list.get(i);
-        i = (i + 1) % list.size();
+        final T next = nextValue;
+        nextValue = advance();
         return next;
     }
 
+    /**
+     * Return the next value in the {@code Iterator}, restarting the
+     * {@code Iterator} if necessary.
+     *
+     * @return The next value in the iterator
+     */
+    private T advance() {
+        if (!iterator.hasNext()) {
+            iterator = iterable.iterator();
+        }
+        return iterator.next();
+    }
+
+    /**
+     * Peek at the next value in the Iterator. Calling this method multiple
+     * times will return the same element without advancing this Iterator. The
+     * value returned by this method will be the next item returned by
+     * {@code next()}.
+     *
+     * @return The next value in this {@code Iterator}
+     */
     public T peek() {
-        return list.get(i);
+        return nextValue;
     }
 
     @Override
diff --git a/clients/src/main/java/org/apache/kafka/common/utils/CloseableIterator.java b/clients/src/main/java/org/apache/kafka/common/utils/CloseableIterator.java
index 38fba8ec9d4e1..50b06369b471e 100644
--- a/clients/src/main/java/org/apache/kafka/common/utils/CloseableIterator.java
+++ b/clients/src/main/java/org/apache/kafka/common/utils/CloseableIterator.java
@@ -27,4 +27,26 @@
  */
 public interface CloseableIterator extends Iterator, Closeable {
     void close();
+
+    static  CloseableIterator wrap(Iterator inner) {
+        return new CloseableIterator() {
+            @Override
+            public void close() {}
+
+            @Override
+            public boolean hasNext() {
+                return inner.hasNext();
+            }
+
+            @Override
+            public R next() {
+                return inner.next();
+            }
+
+            @Override
+            public void remove() {
+                inner.remove();
+            }
+        };
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/utils/FixedOrderMap.java b/clients/src/main/java/org/apache/kafka/common/utils/FixedOrderMap.java
new file mode 100644
index 0000000000000..6518878877028
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/utils/FixedOrderMap.java
@@ -0,0 +1,63 @@
+/*
+ * 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.utils;
+
+import java.util.LinkedHashMap;
+import java.util.Map;
+
+/**
+ * An ordered map (LinkedHashMap) implementation for which the order is immutable.
+ * To accomplish this, all methods of removing mappings are disabled (they are marked
+ * deprecated and throw an exception).
+ *
+ * This class is final to prevent subclasses from violating the desired property.
+ *
+ * @param  The key type
+ * @param  The value type
+ */
+public final class FixedOrderMap extends LinkedHashMap {
+    private static final long serialVersionUID = -6504110858733236170L;
+
+    @Deprecated
+    @Override
+    protected boolean removeEldestEntry(final Map.Entry eldest) {
+        return false;
+    }
+
+    @Deprecated
+    @Override
+    public V remove(final Object key) {
+        throw new UnsupportedOperationException("Removing from registeredStores is not allowed");
+    }
+
+    @Deprecated
+    @Override
+    public boolean remove(final Object key, final Object value) {
+        throw new UnsupportedOperationException("Removing from registeredStores is not allowed");
+    }
+
+    @Deprecated
+    @Override
+    public void clear() {
+        throw new UnsupportedOperationException("Removing from registeredStores is not allowed");
+    }
+
+    @Override
+    public FixedOrderMap clone() {
+        throw new UnsupportedOperationException();
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/utils/FlattenedIterator.java b/clients/src/main/java/org/apache/kafka/common/utils/FlattenedIterator.java
new file mode 100644
index 0000000000000..48bf3b7199e14
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/utils/FlattenedIterator.java
@@ -0,0 +1,45 @@
+/*
+ * 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.utils;
+
+import java.util.Iterator;
+import java.util.function.Function;
+
+/**
+ * Provides a flattened iterator over the inner elements of an outer iterator.
+ */
+public final class FlattenedIterator extends AbstractIterator {
+    private final Iterator outerIterator;
+    private final Function> innerIteratorFunction;
+    private Iterator innerIterator;
+
+    public FlattenedIterator(Iterator outerIterator, Function> innerIteratorFunction) {
+        this.outerIterator = outerIterator;
+        this.innerIteratorFunction = innerIteratorFunction;
+    }
+
+    @Override
+    public I makeNext() {
+        while (innerIterator == null || !innerIterator.hasNext()) {
+            if (outerIterator.hasNext())
+                innerIterator = innerIteratorFunction.apply(outerIterator.next());
+            else
+                return allDone();
+        }
+        return innerIterator.next();
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/utils/ImplicitLinkedHashSet.java b/clients/src/main/java/org/apache/kafka/common/utils/ImplicitLinkedHashCollection.java
similarity index 66%
rename from clients/src/main/java/org/apache/kafka/common/utils/ImplicitLinkedHashSet.java
rename to clients/src/main/java/org/apache/kafka/common/utils/ImplicitLinkedHashCollection.java
index 75fc9ee8273c5..fba7d7ac584ec 100644
--- a/clients/src/main/java/org/apache/kafka/common/utils/ImplicitLinkedHashSet.java
+++ b/clients/src/main/java/org/apache/kafka/common/utils/ImplicitLinkedHashCollection.java
@@ -17,9 +17,14 @@
 
 package org.apache.kafka.common.utils;
 
+import java.util.AbstractCollection;
+import java.util.AbstractSequentialList;
 import java.util.AbstractSet;
 import java.util.Iterator;
+import java.util.List;
+import java.util.ListIterator;
 import java.util.NoSuchElementException;
+import java.util.Set;
 
 /**
  * A memory-efficient hash set which tracks the order of insertion of elements.
@@ -40,7 +45,7 @@
  *
  * This set does not allow null elements.  It does not have internal synchronization.
  */
-public class ImplicitLinkedHashSet extends AbstractSet {
+public class ImplicitLinkedHashCollection extends AbstractCollection {
     public interface Element {
         int prev();
         void setPrev(int prev);
@@ -127,35 +132,137 @@ private static void removeFromList(Element head, Element[] elements, int element
         element.setPrev(INVALID_INDEX);
     }
 
-    private class ImplicitLinkedHashSetIterator implements Iterator {
+    private class ImplicitLinkedHashCollectionIterator implements ListIterator {
+        private int cursor = 0;
         private Element cur = head;
+        private int lastReturnedSlot = INVALID_INDEX;
 
-        private Element next = indexToElement(head, elements, head.next());
+        ImplicitLinkedHashCollectionIterator(int index) {
+            for (int i = 0; i < index; ++i) {
+                cur = indexToElement(head, elements, cur.next());
+                cursor++;
+            }
+        }
 
         @Override
         public boolean hasNext() {
-            return next != head;
+            return cursor != size;
+        }
+
+        @Override
+        public boolean hasPrevious() {
+            return cursor != 0;
         }
 
         @Override
         public E next() {
-            if (next == head) {
+            if (cursor == size) {
+                throw new NoSuchElementException();
+            }
+            lastReturnedSlot = cur.next();
+            cur = indexToElement(head, elements, cur.next());
+            ++cursor;
+            @SuppressWarnings("unchecked")
+            E returnValue = (E) cur;
+            return returnValue;
+        }
+
+        @Override
+        public E previous() {
+            if (cursor == 0) {
                 throw new NoSuchElementException();
             }
-            cur = next;
-            next = indexToElement(head, elements, cur.next());
             @SuppressWarnings("unchecked")
             E returnValue = (E) cur;
+            cur = indexToElement(head, elements, cur.prev());
+            lastReturnedSlot = cur.next();
+            --cursor;
             return returnValue;
         }
 
+        @Override
+        public int nextIndex() {
+            return cursor;
+        }
+
+        @Override
+        public int previousIndex() {
+            return cursor - 1;
+        }
+
         @Override
         public void remove() {
-            if (cur == head) {
+            if (lastReturnedSlot == INVALID_INDEX) {
                 throw new IllegalStateException();
             }
-            ImplicitLinkedHashSet.this.remove(cur);
-            cur = head;
+
+            if (cur == indexToElement(head, elements, lastReturnedSlot)) {
+                cursor--;
+                cur = indexToElement(head, elements, cur.prev());
+            }
+            ImplicitLinkedHashCollection.this.removeElementAtSlot(lastReturnedSlot);
+
+            lastReturnedSlot = INVALID_INDEX;
+        }
+
+        @Override
+        public void set(E e) {
+            throw new UnsupportedOperationException();
+        }
+
+        @Override
+        public void add(E e) {
+            throw new UnsupportedOperationException();
+        }
+    }
+
+    private class ImplicitLinkedHashCollectionListView extends AbstractSequentialList {
+
+        @Override
+        public ListIterator listIterator(int index) {
+            if (index < 0 || index > size) {
+                throw new IndexOutOfBoundsException();
+            }
+
+            return ImplicitLinkedHashCollection.this.listIterator(index);
+        }
+
+        @Override
+        public int size() {
+            return size;
+        }
+    }
+
+    private class ImplicitLinkedHashCollectionSetView extends AbstractSet {
+
+        @Override
+        public Iterator iterator() {
+            return ImplicitLinkedHashCollection.this.iterator();
+        }
+
+        @Override
+        public int size() {
+            return size;
+        }
+
+        @Override
+        public boolean add(E newElement) {
+            return ImplicitLinkedHashCollection.this.add(newElement);
+        }
+
+        @Override
+        public boolean remove(Object key) {
+            return ImplicitLinkedHashCollection.this.remove(key);
+        }
+
+        @Override
+        public boolean contains(Object key) {
+            return ImplicitLinkedHashCollection.this.contains(key);
+        }
+
+        @Override
+        public void clear() {
+            ImplicitLinkedHashCollection.this.clear();
         }
     }
 
@@ -174,7 +281,11 @@ public void remove() {
      */
     @Override
     final public Iterator iterator() {
-        return new ImplicitLinkedHashSetIterator();
+        return listIterator(0);
+    }
+
+    private ListIterator listIterator(int index) {
+        return new ImplicitLinkedHashCollectionIterator(index);
     }
 
     final int slot(Element[] curElements, Object e) {
@@ -193,7 +304,7 @@ final int slot(Element[] curElements, Object e) {
      * @return                  The match index, or INVALID_INDEX if no match was found.
      */
     final private int findIndexOfEqualElement(Object key) {
-        if (key == null) {
+        if (key == null || size == 0) {
             return INVALID_INDEX;
         }
         int slot = slot(elements, key);
@@ -402,30 +513,30 @@ private void reseat(int prevSlot) {
     }
 
     /**
-     * Create a new ImplicitLinkedHashSet.
+     * Create a new ImplicitLinkedHashCollection.
      */
-    public ImplicitLinkedHashSet() {
+    public ImplicitLinkedHashCollection() {
         this(0);
     }
 
     /**
-     * Create a new ImplicitLinkedHashSet.
+     * Create a new ImplicitLinkedHashCollection.
      *
      * @param expectedNumElements   The number of elements we expect to have in this set.
      *                              This is used to optimize by setting the capacity ahead
      *                              of time rather than growing incrementally.
      */
-    public ImplicitLinkedHashSet(int expectedNumElements) {
+    public ImplicitLinkedHashCollection(int expectedNumElements) {
         clear(expectedNumElements);
     }
 
     /**
-     * Create a new ImplicitLinkedHashSet.
+     * Create a new ImplicitLinkedHashCollection.
      *
      * @param iter                  We will add all the elements accessible through this iterator
      *                              to the set.
      */
-    public ImplicitLinkedHashSet(Iterator iter) {
+    public ImplicitLinkedHashCollection(Iterator iter) {
         clear(0);
         while (iter.hasNext()) {
             mustAdd(iter.next());
@@ -457,8 +568,81 @@ final public void clear(int expectedNumElements) {
         }
     }
 
+    /**
+     * Compares the specified object with this collection for equality. Two
+     * {@code ImplicitLinkedHashCollection} objects are equal if they contain the
+     * same elements (as determined by the element's {@code equals} method), and
+     * those elements were inserted in the same order. Because
+     * {@code ImplicitLinkedHashCollectionListIterator} iterates over the elements
+     * in insertion order, it is sufficient to call {@code valuesList.equals}.
+     *
+     * Note that {@link ImplicitLinkedHashMultiCollection} does not override
+     * {@code equals} and uses this method as well. This means that two
+     * {@code ImplicitLinkedHashMultiCollection} objects will be considered equal even
+     * if they each contain two elements A and B such that A.equals(B) but A != B and
+     * A and B have switched insertion positions between the two collections. This
+     * is an acceptable definition of equality, because the collections are still
+     * equal in terms of the order and value of each element.
+     *
+     * @param o object to be compared for equality with this collection
+     * @return true is the specified object is equal to this collection
+     */
+    @Override
+    public boolean equals(Object o) {
+        if (o == this)
+            return true;
+
+        if (!(o instanceof ImplicitLinkedHashCollection))
+            return false;
+
+        ImplicitLinkedHashCollection ilhs = (ImplicitLinkedHashCollection) o;
+        return this.valuesList().equals(ilhs.valuesList());
+    }
+
+    /**
+     * Returns the hash code value for this collection. Because
+     * {@code ImplicitLinkedHashCollection.equals} compares the {@code valuesList}
+     * of two {@code ImplicitLinkedHashCollection} objects to determine equality,
+     * this method uses the @{code valuesList} to compute the has code value as well.
+     *
+     * @return the hash code value for this collection
+     */
+    @Override
+    public int hashCode() {
+        return this.valuesList().hashCode();
+    }
+
     // Visible for testing
     final int numSlots() {
         return elements.length;
     }
+
+    /**
+     * Returns a {@link List} view of the elements contained in the collection,
+     * ordered by order of insertion into the collection. The list is backed by the
+     * collection, so changes to the collection are reflected in the list and
+     * vice-versa. The list supports element removal, which removes the corresponding
+     * element from the collection, but does not support the {@code add} or
+     * {@code set} operations.
+     *
+     * The list is implemented as a circular linked list, so all index-based
+     * operations, such as {@code List.get}, run in O(n) time.
+     *
+     * @return a list view of the elements contained in this collection
+     */
+    public List valuesList() {
+        return new ImplicitLinkedHashCollectionListView();
+    }
+
+    /**
+     * Returns a {@link Set} view of the elements contained in the collection. The
+     * set is backed by the collection, so changes to the collection are reflected in
+     * the set, and vice versa. The set supports element removal and addition, which
+     * removes from or adds to the collection, respectively.
+     *
+     * @return a set view of the elements contained in this collection
+     */
+    public Set valuesSet() {
+        return new ImplicitLinkedHashCollectionSetView();
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/utils/ImplicitLinkedHashMultiSet.java b/clients/src/main/java/org/apache/kafka/common/utils/ImplicitLinkedHashMultiCollection.java
similarity index 90%
rename from clients/src/main/java/org/apache/kafka/common/utils/ImplicitLinkedHashMultiSet.java
rename to clients/src/main/java/org/apache/kafka/common/utils/ImplicitLinkedHashMultiCollection.java
index 2eb53f6f1d892..b5ae8f9793a46 100644
--- a/clients/src/main/java/org/apache/kafka/common/utils/ImplicitLinkedHashMultiSet.java
+++ b/clients/src/main/java/org/apache/kafka/common/utils/ImplicitLinkedHashMultiCollection.java
@@ -24,7 +24,7 @@
 
 /**
  * A memory-efficient hash multiset which tracks the order of insertion of elements.
- * See org.apache.kafka.common.utils.ImplicitLinkedHashSet for implementation details.
+ * See org.apache.kafka.common.utils.ImplicitLinkedHashCollection for implementation details.
  *
  * This class is a multi-set because it allows multiple elements to be inserted that are
  * equal to each other.
@@ -42,17 +42,17 @@
  *
  * This multiset does not allow null elements.  It does not have internal synchronization.
  */
-public class ImplicitLinkedHashMultiSet
-        extends ImplicitLinkedHashSet {
-    public ImplicitLinkedHashMultiSet() {
+public class ImplicitLinkedHashMultiCollection
+        extends ImplicitLinkedHashCollection {
+    public ImplicitLinkedHashMultiCollection() {
         super(0);
     }
 
-    public ImplicitLinkedHashMultiSet(int expectedNumElements) {
+    public ImplicitLinkedHashMultiCollection(int expectedNumElements) {
         super(expectedNumElements);
     }
 
-    public ImplicitLinkedHashMultiSet(Iterator iter) {
+    public ImplicitLinkedHashMultiCollection(Iterator iter) {
         super(iter);
     }
 
@@ -91,7 +91,7 @@ int addInternal(Element newElement, Element[] addElements) {
      */
     @Override
     int findElementToRemove(Object key) {
-        if (key == null) {
+        if (key == null || size() == 0) {
             return INVALID_INDEX;
         }
         int slot = slot(elements, key);
@@ -120,7 +120,7 @@ int findElementToRemove(Object key) {
      * @return          All of the matching elements.
      */
     final public List findAll(E key) {
-        if (key == null) {
+        if (key == null || size() == 0) {
             return Collections.emptyList();
         }
         ArrayList results = new ArrayList<>();
diff --git a/clients/src/main/java/org/apache/kafka/common/utils/LoggingSignalHandler.java b/clients/src/main/java/org/apache/kafka/common/utils/LoggingSignalHandler.java
index fbe67366ed35e..112d7fdb3edaa 100644
--- a/clients/src/main/java/org/apache/kafka/common/utils/LoggingSignalHandler.java
+++ b/clients/src/main/java/org/apache/kafka/common/utils/LoggingSignalHandler.java
@@ -23,6 +23,8 @@
 import java.lang.reflect.InvocationHandler;
 import java.lang.reflect.Method;
 import java.lang.reflect.Proxy;
+import java.util.Arrays;
+import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 
@@ -30,6 +32,8 @@ public class LoggingSignalHandler {
 
     private static final Logger log = LoggerFactory.getLogger(LoggingSignalHandler.class);
 
+    private static final List SIGNALS = Arrays.asList("TERM", "INT", "HUP");
+
     private final Constructor signalConstructor;
     private final Class signalHandlerClass;
     private final Class signalClass;
@@ -61,9 +65,11 @@ public LoggingSignalHandler() throws ReflectiveOperationException {
      */
     public void register() throws ReflectiveOperationException {
         Map jvmSignalHandlers = new ConcurrentHashMap<>();
-        register("TERM", jvmSignalHandlers);
-        register("INT", jvmSignalHandlers);
-        register("HUP", jvmSignalHandlers);
+
+        for (String signal : SIGNALS) {
+            register(signal, jvmSignalHandlers);
+        }
+        log.info("Registered signal handlers for " + String.join(", ", SIGNALS));
     }
 
     private Object createSignalHandler(final Map jvmSignalHandlers) {
diff --git a/clients/src/main/java/org/apache/kafka/common/utils/MappedIterator.java b/clients/src/main/java/org/apache/kafka/common/utils/MappedIterator.java
new file mode 100644
index 0000000000000..f6eb270c56a2d
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/utils/MappedIterator.java
@@ -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.
+ */
+package org.apache.kafka.common.utils;
+
+import java.util.Iterator;
+import java.util.function.Function;
+
+/**
+ * An iterator that maps another iterator's elements from type `F` to type `T`.
+ */
+public final class MappedIterator implements Iterator {
+    private final Iterator underlyingIterator;
+    private final Function mapper;
+
+    public MappedIterator(Iterator underlyingIterator, Function mapper) {
+        this.underlyingIterator = underlyingIterator;
+        this.mapper = mapper;
+    }
+
+    @Override
+    public final boolean hasNext() {
+        return underlyingIterator.hasNext();
+    }
+
+    @Override
+    public final T next() {
+        return mapper.apply(underlyingIterator.next());
+    }
+
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/utils/PoisonPill.java b/clients/src/main/java/org/apache/kafka/common/utils/PoisonPill.java
new file mode 100644
index 0000000000000..383820d5cc989
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/utils/PoisonPill.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.utils;
+
+import com.sun.management.HotSpotDiagnosticMXBean;
+import java.io.File;
+import java.lang.management.ManagementFactory;
+import java.nio.file.Files;
+import java.nio.file.StandardCopyOption;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import javax.management.MBeanServer;
+
+
+public class PoisonPill {
+
+    public static void die() {
+        die(null, -1, 1);
+    }
+
+    public static void die(File heapDumpFolder, final long maxWaitForDump) {
+        die(heapDumpFolder, maxWaitForDump, 1);
+    }
+
+    public static void die(File heapDumpFolder, final long maxWaitForDump, int haltStatusCode) {
+        try {
+            if (maxWaitForDump > 0 && heapDumpFolder != null) {
+                grabHeapDump(heapDumpFolder, maxWaitForDump, haltStatusCode);
+            }
+        } catch (Exception e) {
+            System.err.println("unable to complete heap dump");
+            e.printStackTrace(System.err);
+            System.err.flush();
+        } finally {
+            Runtime.getRuntime().halt(haltStatusCode);
+        }
+    }
+
+    private static void grabHeapDump(File heapDumpFolder, final long maxWait, final int haltStatusCode) throws Exception {
+
+        //set up a watchdog background thread that will halt in ~maxWait regardless of whether or not
+        //we succeed in taking a heap dump (since we dont know when it'll ever complete)
+        final CountDownLatch latch = new CountDownLatch(1);
+        Thread watchdog = new Thread(new Runnable() {
+            @Override
+            public void run() {
+                try {
+                    latch.countDown();
+                    Thread.sleep(maxWait);
+                    //at this point ~maxWait has passed since the call to die().
+                    //if the heap dump process completed successfully die() would
+                    //have called halt() and we wouldnt be here (99.99%)
+                    System.err.println("heap dump (probably) did not complete within timeout. halting.");
+                    System.err.flush();
+                } catch (Exception e) {
+                    System.err.println("watchdog caught exception");
+                    e.printStackTrace(System.err);
+                    System.err.flush();
+                } finally {
+                    Runtime.getRuntime().halt(haltStatusCode);
+                }
+            }
+        });
+        watchdog.setDaemon(true);
+        watchdog.setName("clark the death watchdog");
+        watchdog.start();
+
+        //make sure the watchdog is up and running before we go off attempting to dump
+        if (!latch.await(maxWait, TimeUnit.MILLISECONDS)) {
+            System.err.println("unable to start watchdog within timeout. will not proceed with dump");
+            System.err.flush();
+            return;
+        }
+
+        System.err.println("dumping heap to " + heapDumpFolder.getCanonicalPath());
+        System.err.flush();
+
+        //we dump into dump.inprogress and atomically rename it to be dump.complete
+        //(overwriting any previous such file). this attempts to guarantee there are
+        //at most 2 (potentially large) dump files at any point in time.
+
+        File inProgress = new File(heapDumpFolder, "dump.inprogress.hprof");
+        File complete = new File(heapDumpFolder, "dump.complete.hprof");
+        if (inProgress.exists() && !inProgress.delete()) {
+            System.err.println("unable to delete existing dump file. will not proceed with dump");
+            System.err.flush();
+            return;
+        }
+
+        MBeanServer server = ManagementFactory.getPlatformMBeanServer();
+        HotSpotDiagnosticMXBean diagnosticMBean =
+            ManagementFactory.newPlatformMXBeanProxy(server, "com.sun.management:type=HotSpotDiagnostic",
+                HotSpotDiagnosticMXBean.class);
+        diagnosticMBean.dumpHeap(inProgress.getCanonicalPath(), false /* disable only live - dump all objects */);
+        Files.move(inProgress.toPath(), complete.toPath(), StandardCopyOption.ATOMIC_MOVE,
+            StandardCopyOption.REPLACE_EXISTING);
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/utils/PrimitiveRef.java b/clients/src/main/java/org/apache/kafka/common/utils/PrimitiveRef.java
new file mode 100644
index 0000000000000..e1bbfe373161c
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/utils/PrimitiveRef.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.common.utils;
+
+/**
+ * Primitive reference used to pass primitive typed values as parameter-by-reference.
+ *
+ * This is cheaper than using Atomic references.
+ */
+public class PrimitiveRef {
+    public static IntRef ofInt(int value) {
+        return new IntRef(value);
+    }
+
+    public static class IntRef {
+        public int value;
+
+        IntRef(int value) {
+            this.value = value;
+        }
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/utils/SecurityUtils.java b/clients/src/main/java/org/apache/kafka/common/utils/SecurityUtils.java
index 9c9bd44bfcfde..09c97be15a31c 100644
--- a/clients/src/main/java/org/apache/kafka/common/utils/SecurityUtils.java
+++ b/clients/src/main/java/org/apache/kafka/common/utils/SecurityUtils.java
@@ -16,10 +16,39 @@
  */
 package org.apache.kafka.common.utils;
 
+import org.apache.kafka.common.acl.AclOperation;
+import org.apache.kafka.common.config.SecurityConfig;
+import org.apache.kafka.common.resource.ResourceType;
+import org.apache.kafka.common.security.auth.SecurityProviderCreator;
 import org.apache.kafka.common.security.auth.KafkaPrincipal;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.security.Security;
+import java.util.HashMap;
+import java.util.Map;
 
 public class SecurityUtils {
 
+    private static final Logger LOGGER = LoggerFactory.getLogger(SecurityConfig.class);
+
+    private static final Map NAME_TO_RESOURCE_TYPES;
+    private static final Map NAME_TO_OPERATIONS;
+
+    static {
+        NAME_TO_RESOURCE_TYPES = new HashMap<>(ResourceType.values().length);
+        NAME_TO_OPERATIONS = new HashMap<>(AclOperation.values().length);
+
+        for (ResourceType resourceType : ResourceType.values()) {
+            String resourceTypeName = toPascalCase(resourceType.name());
+            NAME_TO_RESOURCE_TYPES.put(resourceTypeName, resourceType);
+        }
+        for (AclOperation operation : AclOperation.values()) {
+            String operationName = toPascalCase(operation.name());
+            NAME_TO_OPERATIONS.put(operationName, operation);
+        }
+    }
+
     public static KafkaPrincipal parseKafkaPrincipal(String str) {
         if (str == null || str.isEmpty()) {
             throw new IllegalArgumentException("expected a string in format principalType:principalName but got " + str);
@@ -34,4 +63,50 @@ public static KafkaPrincipal parseKafkaPrincipal(String str) {
         return new KafkaPrincipal(split[0], split[1]);
     }
 
+    public static void addConfiguredSecurityProviders(Map configs) {
+        String securityProviderClassesStr = (String) configs.get(SecurityConfig.SECURITY_PROVIDERS_CONFIG);
+        if (securityProviderClassesStr == null || securityProviderClassesStr.equals("")) {
+            return;
+        }
+        try {
+            String[] securityProviderClasses = securityProviderClassesStr.replaceAll("\\s+", "").split(",");
+            for (int index = 0; index < securityProviderClasses.length; index++) {
+                SecurityProviderCreator securityProviderCreator = (SecurityProviderCreator) Class.forName(securityProviderClasses[index]).newInstance();
+                securityProviderCreator.configure(configs);
+                Security.insertProviderAt(securityProviderCreator.getProvider(), index + 1);
+            }
+        } catch (ClassCastException e) {
+            LOGGER.error("Creators provided through " + SecurityConfig.SECURITY_PROVIDERS_CONFIG +
+                    " are expected to be sub-classes of SecurityProviderCreator");
+        } catch (ClassNotFoundException cnfe) {
+            LOGGER.error("Unrecognized security provider creator class", cnfe);
+        } catch (IllegalAccessException | InstantiationException e) {
+            LOGGER.error("Unexpected implementation of security provider creator class", e);
+        }
+    }
+
+    public static ResourceType resourceType(String name) {
+        ResourceType resourceType = NAME_TO_RESOURCE_TYPES.get(name);
+        return resourceType == null ? ResourceType.UNKNOWN : resourceType;
+    }
+
+    public static AclOperation operation(String name) {
+        AclOperation operation = NAME_TO_OPERATIONS.get(name);
+        return operation == null ? AclOperation.UNKNOWN : operation;
+    }
+
+    private static String toPascalCase(String name) {
+        StringBuilder builder = new StringBuilder();
+        boolean capitalizeNext = true;
+        for (char c : name.toCharArray()) {
+            if (c == '_')
+                capitalizeNext = true;
+            else if (capitalizeNext) {
+                builder.append(Character.toUpperCase(c));
+                capitalizeNext = false;
+            } else
+                builder.append(Character.toLowerCase(c));
+        }
+        return builder.toString();
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/utils/SystemTime.java b/clients/src/main/java/org/apache/kafka/common/utils/SystemTime.java
index 9ef096f27d227..31919a22155b6 100644
--- a/clients/src/main/java/org/apache/kafka/common/utils/SystemTime.java
+++ b/clients/src/main/java/org/apache/kafka/common/utils/SystemTime.java
@@ -38,12 +38,7 @@ public long nanoseconds() {
 
     @Override
     public void sleep(long ms) {
-        try {
-            Thread.sleep(ms);
-        } catch (InterruptedException e) {
-            // just wake up early
-            Thread.currentThread().interrupt();
-        }
+        Utils.sleep(ms);
     }
 
     @Override
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 5d2a5cf63b751..5a69db0c021a4 100755
--- a/clients/src/main/java/org/apache/kafka/common/utils/Utils.java
+++ b/clients/src/main/java/org/apache/kafka/common/utils/Utils.java
@@ -39,11 +39,12 @@
 import java.nio.file.Files;
 import java.nio.file.NoSuchFileException;
 import java.nio.file.Path;
-import java.nio.file.SimpleFileVisitor;
 import java.nio.file.Paths;
+import java.nio.file.SimpleFileVisitor;
 import java.nio.file.StandardCopyOption;
 import java.nio.file.attribute.BasicFileAttributes;
 import java.text.DecimalFormat;
+import java.text.DecimalFormatSymbols;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
@@ -53,10 +54,12 @@
 import java.util.Iterator;
 import java.util.LinkedHashMap;
 import java.util.List;
+import java.util.Locale;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Properties;
 import java.util.Set;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Function;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
@@ -74,7 +77,8 @@ private Utils() {}
     private static final Pattern VALID_HOST_CHARACTERS = Pattern.compile("([0-9a-zA-Z\\-%._:]*)");
 
     // Prints up to 2 decimal digits. Used for human readable printing
-    private static final DecimalFormat TWO_DIGIT_FORMAT = new DecimalFormat("0.##");
+    private static final DecimalFormat TWO_DIGIT_FORMAT = new DecimalFormat("0.##",
+        DecimalFormatSymbols.getInstance(Locale.ENGLISH));
 
     private static final String[] BYTE_SCALE_SUFFIXES = new String[] {"B", "KB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB"};
 
@@ -283,20 +287,6 @@ public static byte[] copyArray(byte[] src) {
         return Arrays.copyOf(src, src.length);
     }
 
-    /**
-     * Check that the parameter t is not null
-     *
-     * @param t The object to check
-     * @return t if it isn't null
-     * @throws NullPointerException if t is null.
-     */
-    public static  T notNull(T t) {
-        if (t == null)
-            throw new NullPointerException();
-        else
-            return t;
-    }
-
     /**
      * Sleep for a bit
      * @param ms The duration of the sleep
@@ -601,15 +591,6 @@ public static String stackTrace(Throwable e) {
         return sw.toString();
     }
 
-    /**
-     * Print an error message and shutdown the JVM
-     * @param message The error message
-     */
-    public static void croak(String message) {
-        System.err.println(message);
-        Exit.exit(1);
-    }
-
     /**
      * Read a buffer into a Byte array for the given offset and length
      */
@@ -844,6 +825,18 @@ public static void closeAll(Closeable... closeables) throws IOException {
             throw exception;
     }
 
+    /**
+     * An {@link AutoCloseable} interface without a throws clause in the signature
+     *
+     * This is used with lambda expressions in try-with-resources clauses
+     * to avoid casting un-checked exceptions to checked exceptions unnecessarily.
+     */
+    @FunctionalInterface
+    public interface UncheckedCloseable extends AutoCloseable {
+        @Override
+        void close();
+    }
+
     /**
      * Closes {@code closeable} and if an exception is thrown, it is logged at the WARN level.
      */
@@ -857,6 +850,17 @@ public static void closeQuietly(AutoCloseable closeable, String name) {
         }
     }
 
+    public static void closeQuietly(AutoCloseable closeable, String name, AtomicReference firstException) {
+        if (closeable != null) {
+            try {
+                closeable.close();
+            } catch (Throwable t) {
+                firstException.compareAndSet(null, t);
+                log.error("Failed to close {} with type {}", name, closeable.getClass().getName(), t);
+            }
+        }
+    }
+
     /**
      * A cheap way to deterministically convert a number to a positive value. When the input is
      * positive, the original value is returned. When the input number is negative, the returned
@@ -874,10 +878,6 @@ public static int toPositive(int number) {
         return number & 0x7fffffff;
     }
 
-    public static int longHashcode(long value) {
-        return (int) (value ^ (value >>> 32));
-    }
-
     /**
      * Read a size-delimited byte buffer starting at the given offset.
      * @param buffer Buffer containing the size and data
@@ -1038,4 +1038,16 @@ public static Set from32BitField(final int intValue) {
         }
         return result;
     }
+
+    public static  Map transformMap(
+            Map map,
+            Function keyMapper,
+            Function valueMapper) {
+        return map.entrySet().stream().collect(
+            Collectors.toMap(
+                entry -> keyMapper.apply(entry.getKey()),
+                entry -> valueMapper.apply(entry.getValue())
+            )
+        );
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/server/authorizer/AclCreateResult.java b/clients/src/main/java/org/apache/kafka/server/authorizer/AclCreateResult.java
new file mode 100644
index 0000000000000..70b9c00b3646c
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/server/authorizer/AclCreateResult.java
@@ -0,0 +1,45 @@
+/*
+ * 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.authorizer;
+
+import org.apache.kafka.common.annotation.InterfaceStability;
+import org.apache.kafka.common.errors.ApiException;
+
+import java.util.Optional;
+
+@InterfaceStability.Evolving
+public class AclCreateResult {
+    public static final AclCreateResult SUCCESS = new AclCreateResult();
+
+    private final ApiException exception;
+
+    private AclCreateResult() {
+        this(null);
+    }
+
+    public AclCreateResult(ApiException exception) {
+        this.exception = exception;
+    }
+
+    /**
+     * Returns any exception during create. If exception is empty, the request has succeeded.
+     */
+    public Optional exception() {
+        return exception == null ? Optional.empty() : Optional.of(exception);
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/server/authorizer/AclDeleteResult.java b/clients/src/main/java/org/apache/kafka/server/authorizer/AclDeleteResult.java
new file mode 100644
index 0000000000000..994d6fde74f1c
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/server/authorizer/AclDeleteResult.java
@@ -0,0 +1,95 @@
+/*
+ * 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.authorizer;
+
+import java.util.Collections;
+import java.util.Collection;
+import java.util.Optional;
+
+import org.apache.kafka.common.acl.AclBinding;
+import org.apache.kafka.common.annotation.InterfaceStability;
+import org.apache.kafka.common.errors.ApiException;
+
+@InterfaceStability.Evolving
+public class AclDeleteResult {
+    private final ApiException exception;
+    private final Collection aclBindingDeleteResults;
+
+    public AclDeleteResult(ApiException exception) {
+        this(Collections.emptySet(), exception);
+    }
+
+    public AclDeleteResult(Collection deleteResults) {
+        this(deleteResults, null);
+    }
+
+    private AclDeleteResult(Collection deleteResults, ApiException exception) {
+        this.aclBindingDeleteResults = deleteResults;
+        this.exception = exception;
+    }
+
+    /**
+     * Returns any exception while attempting to match ACL filter to delete ACLs.
+     * If exception is empty, filtering has succeeded. See {@link #aclBindingDeleteResults()}
+     * for deletion results for each filter.
+     */
+    public Optional exception() {
+        return exception == null ? Optional.empty() : Optional.of(exception);
+    }
+
+    /**
+     * Returns delete result for each matching ACL binding.
+     */
+    public Collection aclBindingDeleteResults() {
+        return aclBindingDeleteResults;
+    }
+
+
+    /**
+     * Delete result for each ACL binding that matched a delete filter.
+     */
+    public static class AclBindingDeleteResult {
+        private final AclBinding aclBinding;
+        private final ApiException exception;
+
+        public AclBindingDeleteResult(AclBinding aclBinding) {
+            this(aclBinding, null);
+        }
+
+        public AclBindingDeleteResult(AclBinding aclBinding, ApiException exception) {
+            this.aclBinding = aclBinding;
+            this.exception = exception;
+        }
+
+        /**
+         * Returns ACL binding that matched the delete filter. If {@link #exception()} is
+         * empty, the ACL binding was successfully deleted.
+         */
+        public AclBinding aclBinding() {
+            return aclBinding;
+        }
+
+        /**
+         * Returns any exception that resulted in failure to delete ACL binding.
+         * If exception is empty, the ACL binding was successfully deleted.
+         */
+        public Optional exception() {
+            return exception == null ? Optional.empty() : Optional.of(exception);
+        }
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/server/authorizer/Action.java b/clients/src/main/java/org/apache/kafka/server/authorizer/Action.java
new file mode 100644
index 0000000000000..a62b7f0965b2e
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/server/authorizer/Action.java
@@ -0,0 +1,124 @@
+/*
+ * 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.authorizer;
+
+import java.util.Objects;
+import org.apache.kafka.common.acl.AclOperation;
+import org.apache.kafka.common.annotation.InterfaceStability;
+import org.apache.kafka.common.resource.ResourcePattern;
+
+@InterfaceStability.Evolving
+public class Action {
+
+    private final ResourcePattern resourcePattern;
+    private final AclOperation operation;
+    private final int resourceReferenceCount;
+    private final boolean logIfAllowed;
+    private final boolean logIfDenied;
+
+    public Action(AclOperation operation,
+                  ResourcePattern resourcePattern,
+                  int resourceReferenceCount,
+                  boolean logIfAllowed,
+                  boolean logIfDenied) {
+        this.operation = operation;
+        this.resourcePattern = resourcePattern;
+        this.logIfAllowed = logIfAllowed;
+        this.logIfDenied = logIfDenied;
+        this.resourceReferenceCount = resourceReferenceCount;
+    }
+
+    /**
+     * Resource on which action is being performed.
+     */
+    public ResourcePattern resourcePattern() {
+        return resourcePattern;
+    }
+
+    /**
+     * Operation being performed.
+     */
+    public AclOperation operation() {
+        return operation;
+    }
+
+    /**
+     * Indicates if audit logs tracking ALLOWED access should include this action if result is
+     * ALLOWED. The flag is true if access to a resource is granted while processing the request as a
+     * result of this authorization. The flag is false only for requests used to describe access where
+     * no operation on the resource is actually performed based on the authorization result.
+     */
+    public boolean logIfAllowed() {
+        return logIfAllowed;
+    }
+
+    /**
+     * Indicates if audit logs tracking DENIED access should include this action if result is
+     * DENIED. The flag is true if access to a resource was explicitly requested and request
+     * is denied as a result of this authorization request. The flag is false if request was
+     * filtering out authorized resources (e.g. to subscribe to regex pattern). The flag is also
+     * false if this is an optional authorization where an alternative resource authorization is
+     * applied if this fails (e.g. Cluster:Create which is subsequently overridden by Topic:Create).
+     */
+    public boolean logIfDenied() {
+        return logIfDenied;
+    }
+
+    /**
+     * Number of times the resource being authorized is referenced within the request. For example, a single
+     * request may reference `n` topic partitions of the same topic. Brokers will authorize the topic once
+     * with `resourceReferenceCount=n`. Authorizers may include the count in audit logs.
+     */
+    public int resourceReferenceCount() {
+        return resourceReferenceCount;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+        if (!(o instanceof Action)) {
+            return false;
+        }
+
+        Action that = (Action) o;
+        return Objects.equals(this.resourcePattern, that.resourcePattern) &&
+            Objects.equals(this.operation, that.operation) &&
+            this.resourceReferenceCount == that.resourceReferenceCount &&
+            this.logIfAllowed == that.logIfAllowed &&
+            this.logIfDenied == that.logIfDenied;
+
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(resourcePattern, operation, resourceReferenceCount, logIfAllowed, logIfDenied);
+    }
+
+    @Override
+    public String toString() {
+        return "Action(" +
+            ", resourcePattern='" + resourcePattern + '\'' +
+            ", operation='" + operation + '\'' +
+            ", resourceReferenceCount='" + resourceReferenceCount + '\'' +
+            ", logIfAllowed='" + logIfAllowed + '\'' +
+            ", logIfDenied='" + logIfDenied + '\'' +
+            ')';
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/server/authorizer/AuthorizableRequestContext.java b/clients/src/main/java/org/apache/kafka/server/authorizer/AuthorizableRequestContext.java
new file mode 100644
index 0000000000000..f68b9381ce1d9
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/server/authorizer/AuthorizableRequestContext.java
@@ -0,0 +1,72 @@
+/*
+ * 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.authorizer;
+
+import java.net.InetAddress;
+import org.apache.kafka.common.annotation.InterfaceStability;
+import org.apache.kafka.common.security.auth.KafkaPrincipal;
+import org.apache.kafka.common.security.auth.SecurityProtocol;
+
+/**
+ * Request context interface that provides data from request header as well as connection
+ * and authentication information to plugins.
+ */
+@InterfaceStability.Evolving
+public interface AuthorizableRequestContext {
+
+    /**
+     * Returns name of listener on which request was received.
+     */
+    String listenerName();
+
+    /**
+     * Returns the security protocol for the listener on which request was received.
+     */
+    SecurityProtocol securityProtocol();
+
+    /**
+     * Returns authenticated principal for the connection on which request was received.
+     */
+    KafkaPrincipal principal();
+
+    /**
+     * Returns client IP address from which request was sent.
+     */
+    InetAddress clientAddress();
+
+    /**
+     * 16-bit API key of the request from the request header. See
+     * https://kafka.apache.org/protocol#protocol_api_keys for request types.
+     */
+    int requestType();
+
+    /**
+     * Returns the request version from the request header.
+     */
+    int requestVersion();
+
+    /**
+     * Returns the client id from the request header.
+     */
+    String clientId();
+
+    /**
+     * Returns the correlation id from the request header.
+     */
+    int correlationId();
+}
diff --git a/clients/src/main/java/org/apache/kafka/server/authorizer/AuthorizationResult.java b/clients/src/main/java/org/apache/kafka/server/authorizer/AuthorizationResult.java
new file mode 100644
index 0000000000000..d4ad15d6b0322
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/server/authorizer/AuthorizationResult.java
@@ -0,0 +1,26 @@
+/*
+ * 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.authorizer;
+
+import org.apache.kafka.common.annotation.InterfaceStability;
+
+@InterfaceStability.Evolving
+public enum AuthorizationResult {
+    ALLOWED,
+    DENIED
+}
diff --git a/clients/src/main/java/org/apache/kafka/server/authorizer/Authorizer.java b/clients/src/main/java/org/apache/kafka/server/authorizer/Authorizer.java
new file mode 100644
index 0000000000000..45bd6d939928d
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/server/authorizer/Authorizer.java
@@ -0,0 +1,140 @@
+/*
+ * 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.authorizer;
+
+import java.io.Closeable;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.CompletionStage;
+
+import org.apache.kafka.common.Configurable;
+import org.apache.kafka.common.Endpoint;
+import org.apache.kafka.common.acl.AclBinding;
+import org.apache.kafka.common.acl.AclBindingFilter;
+import org.apache.kafka.common.annotation.InterfaceStability;
+
+/**
+ *
+ * Pluggable authorizer interface for Kafka brokers.
+ *
+ * Startup sequence in brokers:
+ * 
        
+ *   
      1. Broker creates authorizer instance if configured in `authorizer.class.name`.
        2. 
+ *   
      2. Broker configures and starts authorizer instance. Authorizer implementation starts loading its metadata.
        3. 
+ *   
      3. Broker starts SocketServer to accept connections and process requests.
        4. 
+ *   
      4. For each listener, SocketServer waits for authorization metadata to be available in the
+ *       authorizer before accepting connections. The future returned by {@link #start(AuthorizerServerInfo)}
+ *       for each listener must return only when authorizer is ready to authorize requests on the listener.
        5. 
+ *   
      5. Broker accepts connections. For each connection, broker performs authentication and then accepts Kafka requests.
+ *       For each request, broker invokes {@link #authorize(AuthorizableRequestContext, List)} to authorize
+ *       actions performed by the request.
        6. 
+ * 
      
+ *
+ * Authorizer implementation class may optionally implement @{@link org.apache.kafka.common.Reconfigurable}
+ * to enable dynamic reconfiguration without restarting the broker.
+ * 
      
+ * Threading model:
+ * 
        
+ *   
      • All authorizer operations including authorization and ACL updates must be thread-safe.
        • 
+ *   
      • ACL update methods are asynchronous. Implementations with low update latency may return a
+ *       completed future using {@link java.util.concurrent.CompletableFuture#completedFuture(Object)}.
+ *       This ensures that the request will be handled synchronously by the caller without using a
+ *       purgatory to wait for the result. If ACL updates require remote communication which may block,
+ *       return a future that is completed asynchronously when the remote operation completes. This enables
+ *       the caller to process other requests on the request threads without blocking.
        • 
+ *   
      • Any threads or thread pools used for processing remote operations asynchronously can be started during
+ *       {@link #start(AuthorizerServerInfo)}. These threads must be shutdown during {@link Authorizer#close()}.
        • 
+ * 
      
+ * 
      
+ */
+@InterfaceStability.Evolving
+public interface Authorizer extends Configurable, Closeable {
+
+    /**
+     * Starts loading authorization metadata and returns futures that can be used to wait until
+     * metadata for authorizing requests on each listener is available. Each listener will be
+     * started only after its metadata is available and authorizer is ready to start authorizing
+     * requests on that listener.
+     *
+     * @param serverInfo Metadata for the broker including broker id and listener endpoints
+     * @return CompletionStage for each endpoint that completes when authorizer is ready to
+     *         start authorizing requests on that listener.
+     */
+    Map> start(AuthorizerServerInfo serverInfo);
+
+    /**
+     * Authorizes the specified action. Additional metadata for the action is specified
+     * in `requestContext`.
+     * 
      
+     * This is a synchronous API designed for use with locally cached ACLs. Since this method is invoked on the
+     * request thread while processing each request, implementations of this method should avoid time-consuming
+     * remote communication that may block request threads.
+     *
+     * @param requestContext Request context including request type, security protocol and listener name
+     * @param actions Actions being authorized including resource and operation for each action
+     * @return List of authorization results for each action in the same order as the provided actions
+     */
+    List authorize(AuthorizableRequestContext requestContext, List actions);
+
+    /**
+     * Creates new ACL bindings.
+     * 
      
+     * This is an asynchronous API that enables the caller to avoid blocking during the update. Implementations of this
+     * API can return completed futures using {@link java.util.concurrent.CompletableFuture#completedFuture(Object)}
+     * to process the update synchronously on the request thread.
+     *
+     * @param requestContext Request context if the ACL is being created by a broker to handle
+     *        a client request to create ACLs. This may be null if ACLs are created directly in ZooKeeper
+     *        using AclCommand.
+     * @param aclBindings ACL bindings to create
+     *
+     * @return Create result for each ACL binding in the same order as in the input list. Each result
+     *         is returned as a CompletionStage that completes when the result is available.
+     */
+    List> createAcls(AuthorizableRequestContext requestContext, List aclBindings);
+
+    /**
+     * Deletes all ACL bindings that match the provided filters.
+     * 
      
+     * This is an asynchronous API that enables the caller to avoid blocking during the update. Implementations of this
+     * API can return completed futures using {@link java.util.concurrent.CompletableFuture#completedFuture(Object)}
+     * to process the update synchronously on the request thread.
+     *
+     * @param requestContext Request context if the ACL is being deleted by a broker to handle
+     *        a client request to delete ACLs. This may be null if ACLs are deleted directly in ZooKeeper
+     *        using AclCommand.
+     * @param aclBindingFilters Filters to match ACL bindings that are to be deleted
+     *
+     * @return Delete result for each filter in the same order as in the input list.
+     *         Each result indicates which ACL bindings were actually deleted as well as any
+     *         bindings that matched but could not be deleted. Each result is returned as a
+     *         CompletionStage that completes when the result is available.
+     */
+    List> deleteAcls(AuthorizableRequestContext requestContext, List aclBindingFilters);
+
+    /**
+     * Returns ACL bindings which match the provided filter.
+     * 
      
+     * This is a synchronous API designed for use with locally cached ACLs. This method is invoked on the request
+     * thread while processing DescribeAcls requests and should avoid time-consuming remote communication that may
+     * block request threads.
+     *
+     * @return Iterator for ACL bindings, which may be populated lazily.
+     */
+    Iterable acls(AclBindingFilter filter);
+}
diff --git a/clients/src/main/java/org/apache/kafka/server/authorizer/AuthorizerServerInfo.java b/clients/src/main/java/org/apache/kafka/server/authorizer/AuthorizerServerInfo.java
new file mode 100644
index 0000000000000..51e23fba57fad
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/server/authorizer/AuthorizerServerInfo.java
@@ -0,0 +1,51 @@
+/*
+ * 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.authorizer;
+
+import java.util.Collection;
+import org.apache.kafka.common.ClusterResource;
+import org.apache.kafka.common.Endpoint;
+import org.apache.kafka.common.annotation.InterfaceStability;
+
+/**
+ * Runtime broker configuration metadata provided to authorizers during start up.
+ */
+@InterfaceStability.Evolving
+public interface AuthorizerServerInfo {
+
+    /**
+     * Returns cluster metadata for the broker running this authorizer including cluster id.
+     */
+    ClusterResource clusterResource();
+
+    /**
+     * Returns broker id. This may be a generated broker id if `broker.id` was not configured.
+     */
+    int brokerId();
+
+    /**
+     * Returns endpoints for all listeners including the advertised host and port to which
+     * the listener is bound.
+     */
+    Collection endpoints();
+
+    /**
+     * Returns the inter-broker endpoint. This is one of the endpoints returned by {@link #endpoints()}.
+     */
+    Endpoint interBrokerEndpoint();
+}
diff --git a/clients/src/main/resources/common/message/AddOffsetsToTxnRequest.json b/clients/src/main/resources/common/message/AddOffsetsToTxnRequest.json
index 981650f6f4c0b..acf7301664aca 100644
--- a/clients/src/main/resources/common/message/AddOffsetsToTxnRequest.json
+++ b/clients/src/main/resources/common/message/AddOffsetsToTxnRequest.json
@@ -19,14 +19,15 @@
   "name": "AddOffsetsToTxnRequest",
   // Version 1 is the same as version 0.
   "validVersions": "0-1",
+  "flexibleVersions": "none",
   "fields": [
-    { "name": "TransactionalId", "type": "string", "versions": "0+",
+    { "name": "TransactionalId", "type": "string", "versions": "0+", "entityType": "transactionalId",
       "about": "The transactional id corresponding to the transaction."},
-    { "name": "ProducerId", "type": "int64", "versions": "0+",
+    { "name": "ProducerId", "type": "int64", "versions": "0+", "entityType": "producerId",
       "about": "Current producer id in use by the transactional id." },
     { "name": "ProducerEpoch", "type": "int16", "versions": "0+",
       "about": "Current epoch associated with the producer id." },
-    { "name": "GroupId", "type": "string", "versions": "0+",
+    { "name": "GroupId", "type": "string", "versions": "0+", "entityType": "groupId",
       "about": "The unique group identifier." }
   ]
 }
diff --git a/clients/src/main/resources/common/message/AddOffsetsToTxnResponse.json b/clients/src/main/resources/common/message/AddOffsetsToTxnResponse.json
index 080974226fa7a..4e44006aa9324 100644
--- a/clients/src/main/resources/common/message/AddOffsetsToTxnResponse.json
+++ b/clients/src/main/resources/common/message/AddOffsetsToTxnResponse.json
@@ -19,6 +19,7 @@
   "name": "AddOffsetsToTxnResponse",
   // Starting in version 1, on quota violation brokers send out responses before throttling.
   "validVersions": "0-1",
+  "flexibleVersions": "none",
   "fields": [
     { "name": "ThrottleTimeMs", "type": "int32", "versions": "0+",
       "about": "Duration in milliseconds for which the request was throttled due to a quota violation, or zero if the request did not violate any quota." },
diff --git a/clients/src/main/resources/common/message/AddPartitionsToTxnRequest.json b/clients/src/main/resources/common/message/AddPartitionsToTxnRequest.json
index 1c71fa7855c4d..13a5a6e7e9b47 100644
--- a/clients/src/main/resources/common/message/AddPartitionsToTxnRequest.json
+++ b/clients/src/main/resources/common/message/AddPartitionsToTxnRequest.json
@@ -19,16 +19,17 @@
   "name": "AddPartitionsToTxnRequest",
   // Version 1 is the same as version 0.
   "validVersions": "0-1",
+  "flexibleVersions": "none",
   "fields": [
-    { "name": "TransactionalId", "type": "string", "versions": "0+",
+    { "name": "TransactionalId", "type": "string", "versions": "0+", "entityType": "transactionalId",
       "about": "The transactional id corresponding to the transaction."},
-    { "name": "ProducerId", "type": "int64", "versions": "0+",
+    { "name": "ProducerId", "type": "int64", "versions": "0+", "entityType": "producerId",
       "about": "Current producer id in use by the transactional id." },
     { "name": "ProducerEpoch", "type": "int16", "versions": "0+",
       "about": "Current epoch associated with the producer id." },
     { "name": "Topics", "type": "[]AddPartitionsToTxnTopic", "versions": "0+",
       "about": "The partitions to add to the transation.", "fields": [
-      { "name": "Name", "type": "string", "versions": "0+", "mapKey": true,
+      { "name": "Name", "type": "string", "versions": "0+", "mapKey": true, "entityType": "topicName",
         "about": "The name of the topic." },
       { "name": "Partitions", "type": "[]int32", "versions": "0+",
         "about": "The partition indexes to add to the transaction" }
diff --git a/clients/src/main/resources/common/message/AddPartitionsToTxnResponse.json b/clients/src/main/resources/common/message/AddPartitionsToTxnResponse.json
index 50ae5cd6a4a2d..e143a3f1772c5 100644
--- a/clients/src/main/resources/common/message/AddPartitionsToTxnResponse.json
+++ b/clients/src/main/resources/common/message/AddPartitionsToTxnResponse.json
@@ -19,12 +19,13 @@
   "name": "AddPartitionsToTxnResponse",
   // Starting in version 1, on quota violation brokers send out responses before throttling.
   "validVersions": "0-1",
+  "flexibleVersions": "none",
   "fields": [
     { "name": "ThrottleTimeMs", "type": "int32", "versions": "0+",
       "about": "Duration in milliseconds for which the request was throttled due to a quota violation, or zero if the request did not violate any quota." },
     { "name": "Results", "type": "[]AddPartitionsToTxnTopicResult", "versions": "0+",
       "about": "The results for each topic.", "fields": [
-      { "name": "Name", "type": "string", "versions": "0+", "mapKey": true,
+      { "name": "Name", "type": "string", "versions": "0+", "mapKey": true, "entityType": "topicName",
         "about": "The topic name." },
       { "name": "Results", "type": "[]AddPartitionsToTxnPartitionResult", "versions": "0+", 
         "about": "The results for each partition", "fields": [
diff --git a/clients/src/main/resources/common/message/AlterConfigsRequest.json b/clients/src/main/resources/common/message/AlterConfigsRequest.json
index 6d037654d0c50..6f8d72c95d13b 100644
--- a/clients/src/main/resources/common/message/AlterConfigsRequest.json
+++ b/clients/src/main/resources/common/message/AlterConfigsRequest.json
@@ -19,6 +19,7 @@
   "name": "AlterConfigsRequest",
   // Version 1 is the same as version 0.
   "validVersions": "0-1",
+  "flexibleVersions": "none",
   "fields": [
     { "name": "Resources", "type": "[]AlterConfigsResource", "versions": "0+", 
       "about": "The updates for each resource.", "fields": [
diff --git a/clients/src/main/resources/common/message/AlterConfigsResponse.json b/clients/src/main/resources/common/message/AlterConfigsResponse.json
index 135a46741366d..288103df46237 100644
--- a/clients/src/main/resources/common/message/AlterConfigsResponse.json
+++ b/clients/src/main/resources/common/message/AlterConfigsResponse.json
@@ -19,10 +19,11 @@
   "name": "AlterConfigsResponse",
   // Starting in version 1, on quota violation brokers send out responses before throttling.
   "validVersions": "0-1",
+  "flexibleVersions": "none",
   "fields": [
     { "name": "ThrottleTimeMs", "type": "int32", "versions": "0+",
       "about": "Duration in milliseconds for which the request was throttled due to a quota violation, or zero if the request did not violate any quota." },
-    { "name": "Resources", "type": "[]AlterConfigsResourceResponse", "versions": "0+",
+    { "name": "Responses", "type": "[]AlterConfigsResourceResponse", "versions": "0+",
       "about": "The responses for each resource.", "fields": [
       { "name": "ErrorCode", "type": "int16", "versions": "0+",
         "about": "The resource error code." },
diff --git a/clients/src/main/resources/common/message/AlterPartitionReassignmentsRequest.json b/clients/src/main/resources/common/message/AlterPartitionReassignmentsRequest.json
new file mode 100644
index 0000000000000..ecf7eca6eae43
--- /dev/null
+++ b/clients/src/main/resources/common/message/AlterPartitionReassignmentsRequest.json
@@ -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.
+
+{
+  "apiKey": 45,
+  "type": "request",
+  "name": "AlterPartitionReassignmentsRequest",
+  "validVersions": "0",
+  "flexibleVersions": "0+",
+  "fields": [
+    { "name": "TimeoutMs", "type": "int32", "versions": "0+", "default": "60000",
+      "about": "The time in ms to wait for the request to complete." },
+    { "name": "Topics", "type": "[]ReassignableTopic", "versions": "0+",
+      "about": "The topics to reassign.", "fields": [
+      { "name": "Name", "type": "string", "versions": "0+", "entityType": "topicName",
+        "about": "The topic name." },
+      { "name": "Partitions", "type": "[]ReassignablePartition", "versions": "0+",
+        "about": "The partitions to reassign.", "fields": [
+        { "name": "PartitionIndex", "type": "int32", "versions": "0+",
+          "about": "The partition index." },
+        { "name": "Replicas", "type": "[]int32", "versions": "0+", "nullableVersions": "0+", "default": "null",
+          "about": "The replicas to place the partitions on, or null to cancel a pending reassignment for this partition." }
+      ]}
+    ]}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/AlterPartitionReassignmentsResponse.json b/clients/src/main/resources/common/message/AlterPartitionReassignmentsResponse.json
new file mode 100644
index 0000000000000..3fa08883d1bd5
--- /dev/null
+++ b/clients/src/main/resources/common/message/AlterPartitionReassignmentsResponse.json
@@ -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.
+
+{
+  "apiKey": 45,
+  "type": "response",
+  "name": "AlterPartitionReassignmentsResponse",
+  "validVersions": "0",
+  "flexibleVersions": "0+",
+  "fields": [
+    { "name": "ThrottleTimeMs", "type": "int32", "versions": "0+",
+      "about": "The duration in milliseconds for which the request was throttled due to a quota violation, or zero if the request did not violate any quota." },
+    { "name": "ErrorCode", "type": "int16", "versions": "0+",
+      "about": "The top-level error code, or 0 if there was no error." },
+    { "name": "ErrorMessage", "type": "string", "versions": "0+", "nullableVersions": "0+",
+      "about": "The top-level error message, or null if there was no error." },
+    { "name": "Responses", "type": "[]ReassignableTopicResponse", "versions": "0+",
+      "about": "The responses to topics to reassign.", "fields": [
+      { "name": "Name", "type": "string", "versions": "0+", "entityType": "topicName",
+        "about": "The topic name" },
+      { "name": "Partitions", "type": "[]ReassignablePartitionResponse", "versions": "0+",
+        "about": "The responses to partitions to reassign", "fields": [
+        { "name": "PartitionIndex", "type": "int32", "versions": "0+",
+          "about": "The partition index." },
+        { "name": "ErrorCode", "type": "int16", "versions": "0+",
+          "about": "The error code for this partition, or 0 if there was no error." },
+        { "name": "ErrorMessage", "type": "string", "versions": "0+", "nullableVersions": "0+",
+          "about": "The error message for this partition, or null if there was no error." }
+      ]}
+    ]}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/AlterReplicaLogDirsRequest.json b/clients/src/main/resources/common/message/AlterReplicaLogDirsRequest.json
index 4a00249afeef0..d82ee3c975951 100644
--- a/clients/src/main/resources/common/message/AlterReplicaLogDirsRequest.json
+++ b/clients/src/main/resources/common/message/AlterReplicaLogDirsRequest.json
@@ -19,6 +19,7 @@
   "name": "AlterReplicaLogDirsRequest",
   // Version 1 is the same as version 0.
   "validVersions": "0-1",
+  "flexibleVersions": "none",
   "fields": [
     { "name": "Dirs", "type": "[]AlterReplicaLogDir", "versions": "0+", 
       "about": "The alterations to make for each directory.", "fields": [
@@ -26,7 +27,7 @@
         "about": "The absolute directory path." },
       { "name": "Topics", "type": "[]AlterReplicaLogDirTopic", "versions": "0+",
         "about": "The topics to add to the directory.",  "fields": [
-        { "name": "Name", "type": "string", "versions": "0+", "mapKey": true,
+        { "name": "Name", "type": "string", "versions": "0+", "mapKey": true, "entityType": "topicName",
           "about": "The topic name." },
         { "name": "Partitions", "type": "[]int32", "versions": "0+",
           "about": "The partition indexes." }
diff --git a/clients/src/main/resources/common/message/AlterReplicaLogDirsResponse.json b/clients/src/main/resources/common/message/AlterReplicaLogDirsResponse.json
index 2551a15985c9b..75e6f607822a8 100644
--- a/clients/src/main/resources/common/message/AlterReplicaLogDirsResponse.json
+++ b/clients/src/main/resources/common/message/AlterReplicaLogDirsResponse.json
@@ -19,12 +19,13 @@
   "name": "AlterReplicaLogDirsResponse",
   // Starting in version 1, on quota violation brokers send out responses before throttling.
   "validVersions": "0-1",
+  "flexibleVersions": "none",
   "fields": [
     { "name": "ThrottleTimeMs", "type": "int32", "versions": "0+",
       "about": "Duration in milliseconds for which the request was throttled due to a quota violation, or zero if the request did not violate any quota." },
     { "name": "Results", "type": "[]AlterReplicaLogDirTopicResult", "versions": "0+",
       "about": "The results for each topic.", "fields": [
-      { "name": "TopicName", "type": "string", "versions": "0+",
+      { "name": "TopicName", "type": "string", "versions": "0+", "entityType": "topicName",
         "about": "The name of the topic." },
       { "name": "Partitions", "type": "[]AlterReplicaLogDirPartitionResult", "versions": "0+",
         "about": "The results for each partition.", "fields": [
diff --git a/clients/src/main/resources/common/message/ApiVersionsRequest.json b/clients/src/main/resources/common/message/ApiVersionsRequest.json
index d6ddb5475c7b8..79fe7a735274f 100644
--- a/clients/src/main/resources/common/message/ApiVersionsRequest.json
+++ b/clients/src/main/resources/common/message/ApiVersionsRequest.json
@@ -18,7 +18,14 @@
   "type": "request",
   "name": "ApiVersionsRequest",
   // Versions 0 through 2 of ApiVersionsRequest are the same.
-  "validVersions": "0-2",
+  //
+  // Version 3 is the first flexible version and adds ClientSoftwareName and ClientSoftwareVersion.
+  "validVersions": "0-3",
+  "flexibleVersions": "3+",
   "fields": [
+    { "name": "ClientSoftwareName", "type": "string", "versions": "3+",
+      "about": "The name of the client." },
+    { "name": "ClientSoftwareVersion", "type": "string", "versions": "3+",
+      "about": "The version of the client." }
   ]
 }
diff --git a/clients/src/main/resources/common/message/ApiVersionsResponse.json b/clients/src/main/resources/common/message/ApiVersionsResponse.json
index 21fc84a9798c2..ecbc448504918 100644
--- a/clients/src/main/resources/common/message/ApiVersionsResponse.json
+++ b/clients/src/main/resources/common/message/ApiVersionsResponse.json
@@ -18,14 +18,23 @@
   "type": "response",
   "name": "ApiVersionsResponse",
   // Version 1 adds throttle time to the response.
+  //
   // Starting in version 2, on quota violation, brokers send out responses before throttling.
-  "validVersions": "0-2",
+  //
+  // Version 3 is the first flexible version. Tagged fields are only supported in the body but
+  // not in the header. The length of the header must not change in order to guarantee the
+  // backward compatibility.
+  //
+  // Starting from Apache Kafka 2.4 (KIP-511), ApiKeys field is populated with the supported
+  // versions of the ApiVersionsRequest when an UNSUPPORTED_VERSION error is returned.
+  "validVersions": "0-3",
+  "flexibleVersions": "3+",
   "fields": [
     { "name": "ErrorCode", "type": "int16", "versions": "0+",
       "about": "The top-level error code." },
     { "name": "ApiKeys", "type": "[]ApiVersionsResponseKey", "versions": "0+",
       "about": "The APIs supported by the broker.", "fields": [
-      { "name": "Index", "type": "int16", "versions": "0+", "mapKey": true,
+      { "name": "ApiKey", "type": "int16", "versions": "0+", "mapKey": true,
         "about": "The API index." },
       { "name": "MinVersion", "type": "int16", "versions": "0+",
         "about": "The minimum supported version, inclusive." },
diff --git a/clients/src/main/resources/common/message/ControlledShutdownRequest.json b/clients/src/main/resources/common/message/ControlledShutdownRequest.json
index 60ceaa59cdcff..5756d1c16b5df 100644
--- a/clients/src/main/resources/common/message/ControlledShutdownRequest.json
+++ b/clients/src/main/resources/common/message/ControlledShutdownRequest.json
@@ -24,9 +24,10 @@
   // Version 1 is the same as version 0.
   //
   // Version 2 adds BrokerEpoch.
-  "validVersions": "0-2",
+  "validVersions": "0-3",
+  "flexibleVersions": "3+",
   "fields": [
-    { "name": "BrokerId", "type": "int32", "versions": "0+",
+    { "name": "BrokerId", "type": "int32", "versions": "0+", "entityType": "brokerId",
       "about": "The id of the broker for which controlled shutdown has been requested." },
     { "name": "BrokerEpoch", "type": "int64", "versions": "2+", "default": "-1", "ignorable": true,
       "about": "The broker epoch." }
diff --git a/clients/src/main/resources/common/message/ControlledShutdownResponse.json b/clients/src/main/resources/common/message/ControlledShutdownResponse.json
index d0fbcf2fbc87e..27feb1b69b087 100644
--- a/clients/src/main/resources/common/message/ControlledShutdownResponse.json
+++ b/clients/src/main/resources/common/message/ControlledShutdownResponse.json
@@ -18,13 +18,14 @@
   "type": "response",
   "name": "ControlledShutdownResponse",
   // Versions 1 and 2 are the same as version 0.
-  "validVersions": "0-2",
+  "validVersions": "0-3",
+  "flexibleVersions": "3+",
   "fields": [
     { "name": "ErrorCode", "type": "int16", "versions": "0+",
       "about": "The top-level error code." },
     { "name": "RemainingPartitions", "type": "[]RemainingPartition", "versions": "0+",
       "about": "The partitions that the broker still leads.", "fields": [
-      { "name": "TopicName", "type": "string", "versions": "0+", "mapKey": true,
+      { "name": "TopicName", "type": "string", "versions": "0+", "mapKey": true, "entityType": "topicName",
         "about": "The name of the topic." },
       { "name": "PartitionIndex", "type": "int32", "versions": "0+", "mapKey": true,
         "about": "The index of the partition." }
diff --git a/clients/src/main/resources/common/message/CreateAclsRequest.json b/clients/src/main/resources/common/message/CreateAclsRequest.json
index 0e022a8954dc2..72d4d112c0d55 100644
--- a/clients/src/main/resources/common/message/CreateAclsRequest.json
+++ b/clients/src/main/resources/common/message/CreateAclsRequest.json
@@ -19,6 +19,7 @@
   "name": "CreateAclsRequest",
   // Version 1 adds resource pattern type.
   "validVersions": "0-1",
+  "flexibleVersions": "none",
   "fields": [
     { "name": "Creations", "type": "[]CreatableAcl", "versions": "0+", 
       "about": "The ACLs that we want to create.", "fields": [
diff --git a/clients/src/main/resources/common/message/CreateAclsResponse.json b/clients/src/main/resources/common/message/CreateAclsResponse.json
index ede3fef9058c4..d84b7235c3889 100644
--- a/clients/src/main/resources/common/message/CreateAclsResponse.json
+++ b/clients/src/main/resources/common/message/CreateAclsResponse.json
@@ -19,6 +19,7 @@
   "name": "CreateAclsResponse",
   // Starting in version 1, on quota violation, brokers send out responses before throttling.
   "validVersions": "0-1",
+  "flexibleVersions": "none",
   "fields": [
     { "name": "ThrottleTimeMs", "type": "int32", "versions": "0+",
       "about": "The duration in milliseconds for which the request was throttled due to a quota violation, or zero if the request did not violate any quota." },
diff --git a/clients/src/main/resources/common/message/CreateDelegationTokenRequest.json b/clients/src/main/resources/common/message/CreateDelegationTokenRequest.json
index 2de43b9654261..8d881355aa9d2 100644
--- a/clients/src/main/resources/common/message/CreateDelegationTokenRequest.json
+++ b/clients/src/main/resources/common/message/CreateDelegationTokenRequest.json
@@ -18,7 +18,10 @@
   "type": "request",
   "name": "CreateDelegationTokenRequest",
   // Version 1 is the same as version 0.
-  "validVersions": "0-1",
+  //
+  // Version 2 is the first flexible version.
+  "validVersions": "0-2",
+  "flexibleVersions": "2+",
   "fields": [
     { "name": "Renewers", "type": "[]CreatableRenewers", "versions": "0+",
       "about": "A list of those who are allowed to renew this token before it expires.", "fields": [
diff --git a/clients/src/main/resources/common/message/CreateDelegationTokenResponse.json b/clients/src/main/resources/common/message/CreateDelegationTokenResponse.json
index 61367e4991c60..74ad905b94b26 100644
--- a/clients/src/main/resources/common/message/CreateDelegationTokenResponse.json
+++ b/clients/src/main/resources/common/message/CreateDelegationTokenResponse.json
@@ -18,7 +18,10 @@
   "type": "response",
   "name": "CreateDelegationTokenResponse",
   // Starting in version 1, on quota violation, brokers send out responses before throttling.
-  "validVersions": "0-1",
+  //
+  // Version 2 is the first flexible version.
+  "validVersions": "0-2",
+  "flexibleVersions": "2+",
   "fields": [
     { "name": "ErrorCode", "type": "int16", "versions": "0+",
       "about": "The top-level error, or zero if there was no error."},
diff --git a/clients/src/main/resources/common/message/CreatePartitionsRequest.json b/clients/src/main/resources/common/message/CreatePartitionsRequest.json
index 2dc75c7cca445..6db7049899876 100644
--- a/clients/src/main/resources/common/message/CreatePartitionsRequest.json
+++ b/clients/src/main/resources/common/message/CreatePartitionsRequest.json
@@ -19,16 +19,17 @@
   "name": "CreatePartitionsRequest",
   // Version 1 is the same as version 0.
   "validVersions": "0-1",
+  "flexibleVersions": "none",
   "fields": [
     { "name": "Topics", "type": "[]CreatePartitionsTopic", "versions": "0+",
       "about": "Each topic that we want to create new partitions inside.",  "fields": [
-      { "name": "Name", "type": "string", "versions": "0+",
+      { "name": "Name", "type": "string", "versions": "0+", "entityType": "topicName",
         "about": "The topic name." },
       { "name": "Count", "type": "int32", "versions": "0+",
         "about": "The new partition count." },
       { "name": "Assignments", "type": "[]CreatePartitionsAssignment", "versions": "0+", "nullableVersions": "0+", 
         "about": "The new partition assignments.", "fields": [
-        { "name": "BrokerIds", "type": "[]int32", "versions": "0+",
+        { "name": "BrokerIds", "type": "[]int32", "versions": "0+", "entityType": "brokerId",
           "about": "The assigned broker IDs." }
       ]}
     ]},
diff --git a/clients/src/main/resources/common/message/CreatePartitionsResponse.json b/clients/src/main/resources/common/message/CreatePartitionsResponse.json
index 2a0c01ed935fb..f6527266a3436 100644
--- a/clients/src/main/resources/common/message/CreatePartitionsResponse.json
+++ b/clients/src/main/resources/common/message/CreatePartitionsResponse.json
@@ -19,12 +19,13 @@
   "name": "CreatePartitionsResponse",
   // Starting in version 1, on quota violation, brokers send out responses before throttling.
   "validVersions": "0-1",
+  "flexibleVersions": "none",
   "fields": [
     { "name": "ThrottleTimeMs", "type": "int32", "versions": "0+",
       "about": "The duration in milliseconds for which the request was throttled due to a quota violation, or zero if the request did not violate any quota." },
     { "name": "Results", "type": "[]CreatePartitionsTopicResult", "versions": "0+",
       "about": "The partition creation results for each topic.", "fields": [
-      { "name": "Name", "type": "string", "versions": "0+",
+      { "name": "Name", "type": "string", "versions": "0+", "entityType": "topicName",
         "about": "The topic name." },
       { "name": "ErrorCode", "type": "int16", "versions": "0+",
         "about": "The result error, or zero if there was no error."},
diff --git a/clients/src/main/resources/common/message/CreateTopicsRequest.json b/clients/src/main/resources/common/message/CreateTopicsRequest.json
index da6e6c206bcc5..2fc27717a4dee 100644
--- a/clients/src/main/resources/common/message/CreateTopicsRequest.json
+++ b/clients/src/main/resources/common/message/CreateTopicsRequest.json
@@ -18,21 +18,27 @@
   "type": "request",
   "name": "CreateTopicsRequest",
   // Version 1 adds validateOnly.
-  "validVersions": "0-3",
+  //
+  // Version 4 makes partitions/replicationFactor optional even when assignments are not present (KIP-464)
+  //
+  // Version 5 is the first flexible version.
+  // Version 5 also returns topic configs in the response (KIP-525).
+  "validVersions": "0-5",
+  "flexibleVersions": "5+",
   "fields": [
     { "name": "Topics", "type": "[]CreatableTopic", "versions": "0+",
       "about": "The topics to create.", "fields": [
-      { "name": "Name", "type": "string", "versions": "0+", "mapKey": true,
+      { "name": "Name", "type": "string", "versions": "0+", "mapKey": true, "entityType": "topicName",
         "about": "The topic name." },
       { "name": "NumPartitions", "type": "int32", "versions": "0+",
-        "about": "The number of partitions to create in the topic, or -1 if we are specifying a manual partition assignment." },
+        "about": "The number of partitions to create in the topic, or -1 if we are either specifying a manual partition assignment or using the default partitions." },
       { "name": "ReplicationFactor", "type": "int16", "versions": "0+",
-        "about": "The number of replicas to create for each partition in the topic, or -1 if we are specifying a manual partition assignment." },
+        "about": "The number of replicas to create for each partition in the topic, or -1 if we are either specifying a manual partition assignment or using the default replication factor." },
       { "name": "Assignments", "type": "[]CreatableReplicaAssignment", "versions": "0+",
         "about": "The manual partition assignment, or the empty array if we are using automatic assignment.", "fields": [
         { "name": "PartitionIndex", "type": "int32", "versions": "0+", "mapKey": true,
           "about": "The partition index." },
-        { "name": "BrokerIds", "type": "[]int32", "versions": "0+",
+        { "name": "BrokerIds", "type": "[]int32", "versions": "0+", "entityType": "brokerId",
           "about": "The brokers to place the partition on." }
       ]},
       { "name": "Configs", "type": "[]CreateableTopicConfig", "versions": "0+",
diff --git a/clients/src/main/resources/common/message/CreateTopicsResponse.json b/clients/src/main/resources/common/message/CreateTopicsResponse.json
index 49e4d7b888bf9..9d777ab6dd4f2 100644
--- a/clients/src/main/resources/common/message/CreateTopicsResponse.json
+++ b/clients/src/main/resources/common/message/CreateTopicsResponse.json
@@ -18,20 +18,47 @@
   "type": "response",
   "name": "CreateTopicsResponse",
   // Version 1 adds a per-topic error message string.
+  //
   // Version 2 adds the throttle time.
+  //
   // Starting in version 3, on quota violation, brokers send out responses before throttling.
-  "validVersions": "0-3",
+  //
+  // Version 4 makes partitions/replicationFactor optional even when assignments are not present (KIP-464).
+  //
+  // Version 5 is the first flexible version.
+  // Version 5 also returns topic configs in the response (KIP-525).
+  "validVersions": "0-5",
+  "flexibleVersions": "5+",
   "fields": [
     { "name": "ThrottleTimeMs", "type": "int32", "versions": "2+", "ignorable": true,
       "about": "The duration in milliseconds for which the request was throttled due to a quota violation, or zero if the request did not violate any quota." },
     { "name": "Topics", "type": "[]CreatableTopicResult", "versions": "0+",
       "about": "Results for each topic we tried to create.", "fields": [
-      { "name": "Name", "type": "string", "versions": "0+", "mapKey": true,
+      { "name": "Name", "type": "string", "versions": "0+", "mapKey": true, "entityType": "topicName",
         "about": "The topic name." },
       { "name": "ErrorCode", "type": "int16", "versions": "0+",
         "about": "The error code, or 0 if there was no error." },
       { "name": "ErrorMessage", "type": "string", "versions": "1+", "nullableVersions": "0+", "ignorable": true,
-        "about": "The error message, or null if there was no error." }
+        "about": "The error message, or null if there was no error." },
+      { "name": "TopicConfigErrorCode", "type": "int16", "versions": "5+", "tag": 0, "taggedVersions": "5+",
+        "about": "Optional topic config error returned if configs are not returned in the response." },
+      { "name": "NumPartitions", "type": "int32", "versions": "5+", "default": "-1",
+        "about": "Number of partitions of the topic." },
+      { "name": "ReplicationFactor", "type": "int16", "versions": "5+", "default": "-1",
+        "about": "Replicator factor of the topic." },
+      { "name": "Configs", "type": "[]CreatableTopicConfigs", "versions": "5+", "nullableVersions": "5+",
+        "about": "Configuration of the topic.", "fields": [
+        { "name": "Name", "type": "string", "versions": "5+",
+          "about": "The configuration name." },
+        { "name": "Value", "type": "string", "versions": "5+", "nullableVersions": "5+",
+          "about": "The configuration value." },
+        { "name": "ReadOnly", "type": "bool", "versions": "5+",
+          "about": "True if the configuration is read-only." },
+        { "name": "ConfigSource", "type": "int8", "versions": "5+", "default": "-1", "ignorable": true,
+          "about": "The configuration source." },
+        { "name": "IsSensitive", "type": "bool", "versions": "5+",
+          "about": "True if this configuration is sensitive." }
+      ]}
     ]}
   ]
 }
diff --git a/clients/src/main/resources/common/message/DeleteAclsRequest.json b/clients/src/main/resources/common/message/DeleteAclsRequest.json
index 3a4aed173349f..c9d6d4aaff21d 100644
--- a/clients/src/main/resources/common/message/DeleteAclsRequest.json
+++ b/clients/src/main/resources/common/message/DeleteAclsRequest.json
@@ -19,6 +19,7 @@
   "name": "DeleteAclsRequest",
   // Version 1 adds the pattern type.
   "validVersions": "0-1",
+  "flexibleVersions": "none",
   "fields": [
     { "name": "Filters", "type": "[]DeleteAclsFilter", "versions": "0+",
       "about": "The filters to use when deleting ACLs.", "fields": [
diff --git a/clients/src/main/resources/common/message/DeleteAclsResponse.json b/clients/src/main/resources/common/message/DeleteAclsResponse.json
index d752888866990..303fa2bc94dac 100644
--- a/clients/src/main/resources/common/message/DeleteAclsResponse.json
+++ b/clients/src/main/resources/common/message/DeleteAclsResponse.json
@@ -20,6 +20,7 @@
   // Version 1 adds the resource pattern type.
   // Starting in version 1, on quota violation, brokers send out responses before throttling.
   "validVersions": "0-1",
+  "flexibleVersions": "none",
   "fields": [
     { "name": "ThrottleTimeMs", "type": "int32", "versions": "0+",
       "about": "The duration in milliseconds for which the request was throttled due to a quota violation, or zero if the request did not violate any quota." },
diff --git a/clients/src/main/resources/common/message/DeleteGroupsRequest.json b/clients/src/main/resources/common/message/DeleteGroupsRequest.json
index 8dd8172de30b7..833ed7a4dc6d7 100644
--- a/clients/src/main/resources/common/message/DeleteGroupsRequest.json
+++ b/clients/src/main/resources/common/message/DeleteGroupsRequest.json
@@ -18,9 +18,12 @@
   "type": "request",
   "name": "DeleteGroupsRequest",
   // Version 1 is the same as version 0.
-  "validVersions": "0-1",
+  //
+  // Version 2 is the first flexible version.
+  "validVersions": "0-2",
+  "flexibleVersions": "2+",
   "fields": [
-    { "name": "GroupsNames", "type": "[]string", "versions": "0+",
+    { "name": "GroupsNames", "type": "[]string", "versions": "0+", "entityType": "groupId",
       "about": "The group names to delete." }
   ]
 }
diff --git a/clients/src/main/resources/common/message/DeleteGroupsResponse.json b/clients/src/main/resources/common/message/DeleteGroupsResponse.json
index 818331bf61b93..37e06a55b9913 100644
--- a/clients/src/main/resources/common/message/DeleteGroupsResponse.json
+++ b/clients/src/main/resources/common/message/DeleteGroupsResponse.json
@@ -18,13 +18,16 @@
   "type": "response",
   "name": "DeleteGroupsResponse",
   // Starting in version 1, on quota violation, brokers send out responses before throttling.
-  "validVersions": "0-1",
+  //
+  // Version 2 is the first flexible version.
+  "validVersions": "0-2",
+  "flexibleVersions": "2+",
   "fields": [
     { "name": "ThrottleTimeMs", "type": "int32", "versions": "0+",
       "about": "The duration in milliseconds for which the request was throttled due to a quota violation, or zero if the request did not violate any quota." },
     { "name": "Results", "type": "[]DeletableGroupResult", "versions": "0+",
       "about": "The deletion results", "fields": [
-      { "name": "GroupId", "type": "string", "versions": "0+", "mapKey": true,
+      { "name": "GroupId", "type": "string", "versions": "0+", "mapKey": true, "entityType": "groupId",
         "about": "The group id" },
       { "name": "ErrorCode", "type": "int16", "versions": "0+",
         "about": "The deletion error, or 0 if the deletion succeeded." }
diff --git a/clients/src/main/resources/common/message/DeleteRecordsRequest.json b/clients/src/main/resources/common/message/DeleteRecordsRequest.json
index be6c5e755f954..bfaac07900931 100644
--- a/clients/src/main/resources/common/message/DeleteRecordsRequest.json
+++ b/clients/src/main/resources/common/message/DeleteRecordsRequest.json
@@ -19,10 +19,11 @@
   "name": "DeleteRecordsRequest",
   // Version 1 is the same as version 0.
   "validVersions": "0-1",
+  "flexibleVersions": "none",
   "fields": [
     { "name": "Topics", "type": "[]DeleteRecordsTopic", "versions": "0+",
       "about": "Each topic that we want to delete records from.", "fields": [
-      { "name": "Name", "type": "string", "versions": "0+",
+      { "name": "Name", "type": "string", "versions": "0+", "entityType": "topicName",
         "about": "The topic name." },
       { "name": "Partitions", "type": "[]DeleteRecordsPartition", "versions": "0+",
         "about": "Each partition that we want to delete records from.", "fields": [
diff --git a/clients/src/main/resources/common/message/DeleteRecordsResponse.json b/clients/src/main/resources/common/message/DeleteRecordsResponse.json
index 88ac4ab874cfd..140d3e75e4dfc 100644
--- a/clients/src/main/resources/common/message/DeleteRecordsResponse.json
+++ b/clients/src/main/resources/common/message/DeleteRecordsResponse.json
@@ -19,12 +19,13 @@
   "name": "DeleteRecordsResponse",
   // Starting in version 1, on quota violation, brokers send out responses before throttling.
   "validVersions": "0-1",
+  "flexibleVersions": "none",
   "fields": [
     { "name": "ThrottleTimeMs", "type": "int32", "versions": "0+",
       "about": "The duration in milliseconds for which the request was throttled due to a quota violation, or zero if the request did not violate any quota." },
     { "name": "Topics", "type": "[]DeleteRecordsTopicResult", "versions": "0+",
       "about": "Each topic that we wanted to delete records from.", "fields": [
-      { "name": "Name", "type": "string", "versions": "0+",
+      { "name": "Name", "type": "string", "versions": "0+", "entityType": "topicName",
         "about": "The topic name." },
       { "name": "Partitions", "type": "[]DeleteRecordsPartitionResult", "versions": "0+",
         "about": "Each partition that we wanted to delete records from.", "fields": [
diff --git a/clients/src/main/resources/common/message/DeleteTopicsRequest.json b/clients/src/main/resources/common/message/DeleteTopicsRequest.json
index 269a3c01734d6..6e86fc26746bb 100644
--- a/clients/src/main/resources/common/message/DeleteTopicsRequest.json
+++ b/clients/src/main/resources/common/message/DeleteTopicsRequest.json
@@ -18,9 +18,12 @@
   "type": "request",
   "name": "DeleteTopicsRequest",
   // Versions 0, 1, 2, and 3 are the same.
-  "validVersions": "0-3",
+  //
+  // Version 4 is the first flexible version.
+  "validVersions": "0-4",
+  "flexibleVersions": "4+",
   "fields": [
-    { "name": "TopicNames", "type": "[]string", "versions": "0+",
+    { "name": "TopicNames", "type": "[]string", "versions": "0+", "entityType": "topicName",
       "about": "The names of the topics to delete" },
     { "name": "TimeoutMs", "type": "int32", "versions": "0+",
       "about": "The length of time in milliseconds to wait for the deletions to complete." }
diff --git a/clients/src/main/resources/common/message/DeleteTopicsResponse.json b/clients/src/main/resources/common/message/DeleteTopicsResponse.json
index cf0837bbf2065..814909a48d247 100644
--- a/clients/src/main/resources/common/message/DeleteTopicsResponse.json
+++ b/clients/src/main/resources/common/message/DeleteTopicsResponse.json
@@ -18,15 +18,20 @@
   "type": "response",
   "name": "DeleteTopicsResponse",
   // Version 1 adds the throttle time.
+  //
   // Starting in version 2, on quota violation, brokers send out responses before throttling.
+  //
   // Starting in version 3, a TOPIC_DELETION_DISABLED error code may be returned.
-  "validVersions": "0-3",
+  //
+  // Version 4 is the first flexible version.
+  "validVersions": "0-4",
+  "flexibleVersions": "4+",
   "fields": [
-    { "name": "throttleTimeMs", "type": "int32", "versions": "1+",
+    { "name": "ThrottleTimeMs", "type": "int32", "versions": "1+",
       "about": "The duration in milliseconds for which the request was throttled due to a quota violation, or zero if the request did not violate any quota." },
     { "name": "Responses", "type": "[]DeletableTopicResult", "versions": "0+",
-      "about": "The results for each topic.", "fields": [
-      { "name": "Name", "type": "string", "versions": "0+",
+      "about": "The results for each topic we tried to delete.", "fields": [
+      { "name": "Name", "type": "string", "versions": "0+", "mapKey": true, "entityType": "topicName",
         "about": "The topic name" },
       { "name": "ErrorCode", "type": "int16", "versions": "0+",
         "about": "The deletion error, or 0 if the deletion succeeded." }
diff --git a/clients/src/main/resources/common/message/DescribeAclsRequest.json b/clients/src/main/resources/common/message/DescribeAclsRequest.json
index 9be6bbe24617f..258dddc1f51ef 100644
--- a/clients/src/main/resources/common/message/DescribeAclsRequest.json
+++ b/clients/src/main/resources/common/message/DescribeAclsRequest.json
@@ -19,6 +19,7 @@
   "name": "DescribeAclsRequest",
   // Version 1 adds resource pattern type.
   "validVersions": "0-1",
+  "flexibleVersions": "none",
   "fields": [
     { "name": "ResourceType", "type": "int8", "versions": "0+",
       "about": "The resource type." },
diff --git a/clients/src/main/resources/common/message/DescribeAclsResponse.json b/clients/src/main/resources/common/message/DescribeAclsResponse.json
index 0fdf0c0335864..9f04de01b8757 100644
--- a/clients/src/main/resources/common/message/DescribeAclsResponse.json
+++ b/clients/src/main/resources/common/message/DescribeAclsResponse.json
@@ -20,6 +20,7 @@
   // Version 1 adds PatternType.
   // Starting in version 1, on quota violation, brokers send out responses before throttling.
   "validVersions": "0-1",
+  "flexibleVersions": "none",
   "fields": [
     { "name": "ThrottleTimeMs", "type": "int32", "versions": "0+",
       "about": "The duration in milliseconds for which the request was throttled due to a quota violation, or zero if the request did not violate any quota." },
diff --git a/clients/src/main/resources/common/message/DescribeConfigsRequest.json b/clients/src/main/resources/common/message/DescribeConfigsRequest.json
index 5113d4d8dc4b9..5d4d13b3fdb8b 100644
--- a/clients/src/main/resources/common/message/DescribeConfigsRequest.json
+++ b/clients/src/main/resources/common/message/DescribeConfigsRequest.json
@@ -20,6 +20,7 @@
   // Version 1 adds IncludeSynoyms.
   // Version 2 is the same as version 1.
   "validVersions": "0-2",
+  "flexibleVersions": "none",
   "fields": [
     { "name": "Resources", "type": "[]DescribeConfigsResource", "versions": "0+",
       "about": "The resources whose configurations we want to describe.", "fields": [
diff --git a/clients/src/main/resources/common/message/DescribeConfigsResponse.json b/clients/src/main/resources/common/message/DescribeConfigsResponse.json
index 89cb145b1d678..82e44f7856ce8 100644
--- a/clients/src/main/resources/common/message/DescribeConfigsResponse.json
+++ b/clients/src/main/resources/common/message/DescribeConfigsResponse.json
@@ -20,6 +20,7 @@
   // Version 1 adds ConfigSource and the synonyms.
   // Starting in version 2, on quota violation, brokers send out responses before throttling.
   "validVersions": "0-2",
+  "flexibleVersions": "none",
   "fields": [
     { "name": "ThrottleTimeMs", "type": "int32", "versions": "0+",
       "about": "The duration in milliseconds for which the request was throttled due to a quota violation, or zero if the request did not violate any quota." },
@@ -43,7 +44,7 @@
           "about": "True if the configuration is read-only." },
         { "name": "IsDefault", "type": "bool", "versions": "0",
           "about": "True if the configuration is not set." },
-        // Note: the v0 default for this field that shouldd be exposed to callers is
+        // Note: the v0 default for this field that should be exposed to callers is
         // context-dependent. For example, if the resource is a broker, this should default to 4.
         // -1 is just a placeholder value.
         { "name": "ConfigSource", "type": "int8", "versions": "1+", "default": "-1", "ignorable": true,
diff --git a/clients/src/main/resources/common/message/DescribeDelegationTokenRequest.json b/clients/src/main/resources/common/message/DescribeDelegationTokenRequest.json
index e9ec3baffc5c2..2bc2eba721492 100644
--- a/clients/src/main/resources/common/message/DescribeDelegationTokenRequest.json
+++ b/clients/src/main/resources/common/message/DescribeDelegationTokenRequest.json
@@ -19,6 +19,7 @@
   "name": "DescribeDelegationTokenRequest",
   // Version 1 is the same as version 0.
   "validVersions": "0-1",
+  "flexibleVersions": "none",
   "fields": [
     { "name": "Owners", "type": "[]DescribeDelegationTokenOwner", "versions": "0+", "nullableVersions": "0+",
       "about": "Each owner that we want to describe delegation tokens for, or null to describe all tokens.", "fields": [
diff --git a/clients/src/main/resources/common/message/DescribeDelegationTokenResponse.json b/clients/src/main/resources/common/message/DescribeDelegationTokenResponse.json
index 627246d254edf..302ff24a5cc88 100644
--- a/clients/src/main/resources/common/message/DescribeDelegationTokenResponse.json
+++ b/clients/src/main/resources/common/message/DescribeDelegationTokenResponse.json
@@ -19,6 +19,7 @@
   "name": "DescribeDelegationTokenResponse",
   // Starting in version 1, on quota violation, brokers send out responses before throttling.
   "validVersions": "0-1",
+  "flexibleVersions": "none",
   "fields": [
     { "name": "ErrorCode", "type": "int16", "versions": "0+",
       "about": "The error code, or 0 if there was no error." },
diff --git a/clients/src/main/resources/common/message/DescribeGroupsRequest.json b/clients/src/main/resources/common/message/DescribeGroupsRequest.json
index 3557bae3fb237..8a5887a5680a8 100644
--- a/clients/src/main/resources/common/message/DescribeGroupsRequest.json
+++ b/clients/src/main/resources/common/message/DescribeGroupsRequest.json
@@ -18,10 +18,16 @@
   "type": "request",
   "name": "DescribeGroupsRequest",
   // Versions 1 and 2 are the same as version 0.
+  //
   // Starting in version 3, authorized operations can be requested.
-  "validVersions": "0-3",
+  //
+  // Starting in version 4, the response will include group.instance.id info for members.
+  //
+  // Version 5 is the first flexible version.
+  "validVersions": "0-5",
+  "flexibleVersions": "5+",
   "fields": [
-    { "name": "Groups", "type": "[]string", "versions": "0+",
+    { "name": "Groups", "type": "[]string", "versions": "0+", "entityType": "groupId",
       "about": "The names of the groups to describe" },
     { "name": "IncludeAuthorizedOperations", "type": "bool", "versions": "3+",
       "about": "Whether to include authorized operations." }
diff --git a/clients/src/main/resources/common/message/DescribeGroupsResponse.json b/clients/src/main/resources/common/message/DescribeGroupsResponse.json
index f4677a7e4140b..7e8a4d98b0727 100644
--- a/clients/src/main/resources/common/message/DescribeGroupsResponse.json
+++ b/clients/src/main/resources/common/message/DescribeGroupsResponse.json
@@ -18,9 +18,16 @@
   "type": "response",
   "name": "DescribeGroupsResponse",
   // Version 1 added throttle time.
+  //
   // Starting in version 2, on quota violation, brokers send out responses before throttling.
+  //
   // Starting in version 3, brokers can send authorized operations.
-  "validVersions": "0-3",
+  //
+  // Starting in version 4, the response will include group.instance.id info for members.
+  //
+  // Version 5 is the first flexible version.
+  "validVersions": "0-5",
+  "flexibleVersions": "5+",
   "fields": [
     { "name": "ThrottleTimeMs", "type": "int32", "versions": "1+", "ignorable": true,
       "about": "The duration in milliseconds for which the request was throttled due to a quota violation, or zero if the request did not violate any quota." },
@@ -28,7 +35,7 @@
       "about": "Each described group.", "fields": [
       { "name": "ErrorCode", "type": "int16", "versions": "0+",
         "about": "The describe error, or 0 if there was no error." },
-      { "name": "GroupId", "type": "string", "versions": "0+",
+      { "name": "GroupId", "type": "string", "versions": "0+", "entityType": "groupId",
         "about": "The group ID string." },
       { "name": "GroupState", "type": "string", "versions": "0+",
         "about": "The group state string, or the empty string." },
@@ -42,6 +49,9 @@
         "about": "The group members.", "fields": [
         { "name": "MemberId", "type": "string", "versions": "0+",
           "about": "The member ID assigned by the group coordinator." },
+        { "name": "GroupInstanceId", "type": "string", "versions": "4+",
+          "nullableVersions": "4+", "default": "null",
+          "about": "The unique identifier of the consumer instance provided by end user." },
         { "name": "ClientId", "type": "string", "versions": "0+",
           "about": "The client ID used in the member's latest join group request." },
         { "name": "ClientHost", "type": "string", "versions": "0+",
@@ -53,7 +63,7 @@
         { "name": "MemberAssignment", "type": "bytes", "versions": "0+",
           "about": "The current assignment provided by the group leader." }
       ]},
-      { "name": "AuthorizedOperations", "type": "int32", "versions": "3+",
+      { "name": "AuthorizedOperations", "type": "int32", "versions": "3+",  "default": "-2147483648",
         "about": "32-bit bitfield to represent authorized operations for this group." }
     ]}
   ]
diff --git a/clients/src/main/resources/common/message/DescribeLogDirsRequest.json b/clients/src/main/resources/common/message/DescribeLogDirsRequest.json
index b17cd4b3eb3b4..333fa682b1dcb 100644
--- a/clients/src/main/resources/common/message/DescribeLogDirsRequest.json
+++ b/clients/src/main/resources/common/message/DescribeLogDirsRequest.json
@@ -19,10 +19,11 @@
   "name": "DescribeLogDirsRequest",
   // Version 1 is the same as version 0.
   "validVersions": "0-1",
+  "flexibleVersions": "none",
   "fields": [
     { "name": "Topics", "type": "[]DescribableLogDirTopic", "versions": "0+", "nullableVersions": "0+",
       "about": "Each topic that we want to describe log directories for, or null for all topics.", "fields": [
-      { "name": "Topic", "type": "string", "versions": "0+",
+      { "name": "Topic", "type": "string", "versions": "0+", "entityType": "topicName",
         "about": "The topic name" },
       { "name": "PartitionIndex", "type": "[]int32", "versions": "0+",
         "about": "The partition indxes." }
diff --git a/clients/src/main/resources/common/message/DescribeLogDirsResponse.json b/clients/src/main/resources/common/message/DescribeLogDirsResponse.json
index 85355d3a4eaa8..d75b348b298e2 100644
--- a/clients/src/main/resources/common/message/DescribeLogDirsResponse.json
+++ b/clients/src/main/resources/common/message/DescribeLogDirsResponse.json
@@ -19,6 +19,7 @@
   "name": "DescribeLogDirsResponse",
   // Starting in version 1, on quota violation, brokers send out responses before throttling.
   "validVersions": "0-1",
+  "flexibleVersions": "none",
   "fields": [
     { "name": "ThrottleTimeMs", "type": "int32", "versions": "0+",
       "about": "The duration in milliseconds for which the request was throttled due to a quota violation, or zero if the request did not violate any quota." },
@@ -30,7 +31,7 @@
         "about": "The absolute log directory path." },
       { "name": "Topics", "type": "[]DescribeLogDirsTopic", "versions": "0+",
         "about": "Each topic.", "fields": [
-        { "name": "Name", "type": "string", "versions": "0+",
+        { "name": "Name", "type": "string", "versions": "0+", "entityType": "topicName",
           "about": "The topic name." },
         { "name": "Partitions", "type": "[]DescribeLogDirsPartition", "versions": "0+", "fields": [
           { "name": "PartitionIndex", "type": "int32", "versions": "0+",
diff --git a/clients/src/main/resources/common/message/ElectPreferredLeadersRequest.json b/clients/src/main/resources/common/message/ElectLeadersRequest.json
similarity index 63%
rename from clients/src/main/resources/common/message/ElectPreferredLeadersRequest.json
rename to clients/src/main/resources/common/message/ElectLeadersRequest.json
index f566cdf3f6416..5b86c96b04d60 100644
--- a/clients/src/main/resources/common/message/ElectPreferredLeadersRequest.json
+++ b/clients/src/main/resources/common/message/ElectLeadersRequest.json
@@ -16,18 +16,25 @@
 {
   "apiKey": 43,
   "type": "request",
-  "name": "ElectPreferredLeadersRequest",
-  "validVersions": "0",
+  "name": "ElectLeadersRequest",
+  // Version 1 implements multiple leader election types, as described by KIP-460.
+  //
+  // Version 2 is the first flexible version.
+  "validVersions": "0-2",
+  "flexibleVersions": "2+",
   "fields": [
+    { "name": "ElectionType", "type": "int8", "versions": "1+",
+      "about": "Type of elections to conduct for the partition. A value of '0' elects the preferred replica. A value of '1' elects the first live replica if there are no in-sync replica." },
     { "name": "TopicPartitions", "type": "[]TopicPartitions", "versions": "0+", "nullableVersions": "0+",
-      "about": "The topic partitions to elect the preferred leader of.",
+      "about": "The topic partitions to elect leaders.",
       "fields": [
-        { "name": "Topic", "type": "string", "versions": "0+",
+        { "name": "Topic", "type": "string", "versions": "0+", "entityType": "topicName",
           "about": "The name of a topic." },
         { "name": "PartitionId", "type": "[]int32", "versions": "0+",
-          "about": "The partitions of this topic whose preferred leader should be elected" }
-      ]},
+          "about": "The partitions of this topic whose leader should be elected." }
+      ]
+    },
     { "name": "TimeoutMs", "type": "int32", "versions": "0+", "default": "60000",
       "about": "The time in ms to wait for the election to complete." }
   ]
-}
\ No newline at end of file
+}
diff --git a/clients/src/main/resources/common/message/ElectPreferredLeadersResponse.json b/clients/src/main/resources/common/message/ElectLeadersResponse.json
similarity index 76%
rename from clients/src/main/resources/common/message/ElectPreferredLeadersResponse.json
rename to clients/src/main/resources/common/message/ElectLeadersResponse.json
index f34599cf03f39..15468c78d1f70 100644
--- a/clients/src/main/resources/common/message/ElectPreferredLeadersResponse.json
+++ b/clients/src/main/resources/common/message/ElectLeadersResponse.json
@@ -16,14 +16,20 @@
 {
   "apiKey": 43,
   "type": "response",
-  "name": "ElectPreferredLeadersResponse",
-  "validVersions": "0",
+  "name": "ElectLeadersResponse",
+  // Version 1 adds a top-level error code.
+  //
+  // Version 2 is the first flexible version.
+  "validVersions": "0-2",
+  "flexibleVersions": "2+",
   "fields": [
     { "name": "ThrottleTimeMs", "type": "int32", "versions": "0+",
       "about": "The duration in milliseconds for which the request was throttled due to a quota violation, or zero if the request did not violate any quota." },
+    { "name": "ErrorCode", "type": "int16", "versions": "1+", "ignorable": false,
+      "about": "The top level response error code." },
     { "name": "ReplicaElectionResults", "type": "[]ReplicaElectionResult", "versions": "0+",
-      "about": "The error code, or 0 if there was no error.", "fields": [
-      { "name": "Topic", "type": "string", "versions": "0+",
+      "about": "The election results, or an empty array if the requester did not have permission and the request asks for all partitions.", "fields": [
+      { "name": "Topic", "type": "string", "versions": "0+", "entityType": "topicName",
         "about": "The topic name" },
       { "name": "PartitionResult", "type": "[]PartitionResult", "versions": "0+",
         "about": "The results for each partition", "fields": [
@@ -36,4 +42,4 @@
       ]}
     ]}
   ]
-}
\ No newline at end of file
+}
diff --git a/clients/src/main/resources/common/message/EndTxnRequest.json b/clients/src/main/resources/common/message/EndTxnRequest.json
index ebf122439b910..4105479b69f29 100644
--- a/clients/src/main/resources/common/message/EndTxnRequest.json
+++ b/clients/src/main/resources/common/message/EndTxnRequest.json
@@ -19,10 +19,11 @@
   "name": "EndTxnRequest",
   // Version 1 is the same as version 0.
   "validVersions": "0-1",
+  "flexibleVersions": "none",
   "fields": [
-    { "name": "TransactionalId", "type": "string", "versions": "0+",
+    { "name": "TransactionalId", "type": "string", "versions": "0+", "entityType": "transactionalId",
       "about": "The ID of the transaction to end." },
-    { "name": "ProducerId", "type": "int64", "versions": "0+",
+    { "name": "ProducerId", "type": "int64", "versions": "0+", "entityType": "producerId",
       "about": "The producer ID." },
     { "name": "ProducerEpoch", "type": "int16", "versions": "0+",
       "about": "The current epoch associated with the producer." },
diff --git a/clients/src/main/resources/common/message/EndTxnResponse.json b/clients/src/main/resources/common/message/EndTxnResponse.json
index b0d4010a7e408..247e9b974b963 100644
--- a/clients/src/main/resources/common/message/EndTxnResponse.json
+++ b/clients/src/main/resources/common/message/EndTxnResponse.json
@@ -19,6 +19,7 @@
   "name": "EndTxnResponse",
   // Starting in version 1, on quota violation, brokers send out responses before throttling.
   "validVersions": "0-1",
+  "flexibleVersions": "none",
   "fields": [
     { "name": "ThrottleTimeMs", "type": "int32", "versions": "0+",
       "about": "The duration in milliseconds for which the request was throttled due to a quota violation, or zero if the request did not violate any quota." },
diff --git a/clients/src/main/resources/common/message/ExpireDelegationTokenRequest.json b/clients/src/main/resources/common/message/ExpireDelegationTokenRequest.json
index fdc50d0559f43..5192b3f1173cc 100644
--- a/clients/src/main/resources/common/message/ExpireDelegationTokenRequest.json
+++ b/clients/src/main/resources/common/message/ExpireDelegationTokenRequest.json
@@ -19,6 +19,7 @@
   "name": "ExpireDelegationTokenRequest",
   // Version 1 is the same as version 0.
   "validVersions": "0-1",
+  "flexibleVersions": "none",
   "fields": [
     { "name": "Hmac", "type": "bytes", "versions": "0+",
       "about": "The HMAC of the delegation token to be expired." },
diff --git a/clients/src/main/resources/common/message/ExpireDelegationTokenResponse.json b/clients/src/main/resources/common/message/ExpireDelegationTokenResponse.json
index 53fd5a0a7ea1c..4182ec0ca06d9 100644
--- a/clients/src/main/resources/common/message/ExpireDelegationTokenResponse.json
+++ b/clients/src/main/resources/common/message/ExpireDelegationTokenResponse.json
@@ -19,6 +19,7 @@
   "name": "ExpireDelegationTokenResponse",
   // Starting in version 1, on quota violation, brokers send out responses before throttling.
   "validVersions": "0-1",
+  "flexibleVersions": "none",
   "fields": [
     { "name": "ErrorCode", "type": "int16", "versions": "0+",
       "about": "The error code, or 0 if there was no error." },
diff --git a/clients/src/main/resources/common/message/FetchRequest.json b/clients/src/main/resources/common/message/FetchRequest.json
index ee1e88b65e562..71fe3790e4ab7 100644
--- a/clients/src/main/resources/common/message/FetchRequest.json
+++ b/clients/src/main/resources/common/message/FetchRequest.json
@@ -44,7 +44,8 @@
   // Version 10 indicates that we can use the ZStd compression algorithm, as
   // described in KIP-110.
   //
-  "validVersions": "0-10",
+  "validVersions": "0-11",
+  "flexibleVersions": "none",
   "fields": [
     { "name": "ReplicaId", "type": "int32", "versions": "0+",
       "about": "The broker ID of the follower, of -1 if this request is from a consumer." },
@@ -59,10 +60,10 @@
     { "name": "SessionId", "type": "int32", "versions": "7+", "default": "0", "ignorable": false,
       "about": "The fetch session ID." },
     { "name": "Epoch", "type": "int32", "versions": "7+", "default": "-1", "ignorable": false,
-      "about": "The fetch session ID." },
+      "about": "The epoch of the partition leader as known to the follower replica or a consumer." },
     { "name": "Topics", "type": "[]FetchableTopic", "versions": "0+",
       "about": "The topics to fetch.", "fields": [
-      { "name": "Name", "type": "string", "versions": "0+",
+      { "name": "Name", "type": "string", "versions": "0+", "entityType": "topicName",
         "about": "The name of the topic to fetch." },
       { "name": "FetchPartitions", "type": "[]FetchPartition", "versions": "0+",
         "about": "The partitions to fetch.", "fields": [
@@ -80,10 +81,12 @@
     ]},
     { "name": "Forgotten", "type": "[]ForgottenTopic", "versions": "7+", "ignorable": false,
       "about": "In an incremental fetch request, the partitions to remove.", "fields": [
-      { "name": "Name", "type": "string", "versions": "7+",
+      { "name": "Name", "type": "string", "versions": "7+", "entityType": "topicName",
         "about": "The partition name." },
       { "name": "ForgottenPartitionIndexes", "type": "[]int32", "versions": "7+",
         "about": "The partitions indexes to forget." }
-    ]}
+    ]},
+    { "name": "RackId", "type":  "string", "versions": "11+", "default": "", "ignorable": true,
+      "about": "Rack ID of the consumer making this request"}
   ]
 }
diff --git a/clients/src/main/resources/common/message/FetchResponse.json b/clients/src/main/resources/common/message/FetchResponse.json
index afee391c71065..0fe98ade2f5f7 100644
--- a/clients/src/main/resources/common/message/FetchResponse.json
+++ b/clients/src/main/resources/common/message/FetchResponse.json
@@ -38,7 +38,8 @@
   // Version 10 indicates that the response data can use the ZStd compression
   // algorithm, as described in KIP-110.
   //
-  "validVersions": "0-10",
+  "validVersions": "0-11",
+  "flexibleVersions": "none",
   "fields": [
     { "name": "ThrottleTimeMs", "type": "int32", "versions": "1+", "ignorable": true,
       "about": "The duration in milliseconds for which the request was throttled due to a quota violation, or zero if the request did not violate any quota." },
@@ -48,7 +49,7 @@
       "about": "The fetch session ID, or 0 if this is not part of a fetch session." },
     { "name": "Topics", "type": "[]FetchableTopicResponse", "versions": "0+",
       "about": "The response topics.", "fields": [
-      { "name": "Name", "type": "string", "versions": "0+",
+      { "name": "Name", "type": "string", "versions": "0+", "entityType": "topicName",
         "about": "The topic name." },
       { "name": "Partitions", "type": "[]FetchablePartitionResponse", "versions": "0+",
         "about": "The topic partitions.", "fields": [
@@ -64,11 +65,13 @@
           "about": "The current log start offset." },
         { "name": "Aborted", "type": "[]AbortedTransaction", "versions": "4+", "nullableVersions": "4+", "ignorable": false,
           "about": "The aborted transactions.",  "fields": [
-          { "name": "ProducerId", "type": "int64", "versions": "4+",
+          { "name": "ProducerId", "type": "int64", "versions": "4+", "entityType": "producerId",
             "about": "The producer id associated with the aborted transaction." },
           { "name": "FirstOffset", "type": "int64", "versions": "4+",
             "about": "The first offset in the aborted transaction." }
         ]},
+        { "name": "PreferredReadReplica", "type": "int32", "versions": "11+", "ignorable": true,
+          "about": "The preferred read replica for the consumer to use on its next fetch request"},
         { "name": "Records", "type": "bytes", "versions": "0+", "nullableVersions": "0+",
           "about": "The record data." }
       ]}
diff --git a/clients/src/main/resources/common/message/FindCoordinatorRequest.json b/clients/src/main/resources/common/message/FindCoordinatorRequest.json
index 8ef3f7420e2f3..6a90887b1aa01 100644
--- a/clients/src/main/resources/common/message/FindCoordinatorRequest.json
+++ b/clients/src/main/resources/common/message/FindCoordinatorRequest.json
@@ -18,8 +18,12 @@
   "type": "request",
   "name": "FindCoordinatorRequest",
   // Version 1 adds KeyType.
+  //
   // Version 2 is the same as version 1.
-  "validVersions": "0-2",
+  //
+  // Version 3 is the first flexible version.
+  "validVersions": "0-3",
+  "flexibleVersions": "3+",
   "fields": [
     { "name": "Key", "type": "string", "versions": "0+",
       "about": "The coordinator key." },
diff --git a/clients/src/main/resources/common/message/FindCoordinatorResponse.json b/clients/src/main/resources/common/message/FindCoordinatorResponse.json
index aed8f1a8ae323..996e8fb5fba5c 100644
--- a/clients/src/main/resources/common/message/FindCoordinatorResponse.json
+++ b/clients/src/main/resources/common/message/FindCoordinatorResponse.json
@@ -18,8 +18,12 @@
   "type": "response",
   "name": "FindCoordinatorResponse",
   // Version 1 adds throttle time and error messages.
+  //
   // Starting in version 2, on quota violation, brokers send out responses before throttling.
-  "validVersions": "0-2",
+  //
+  // Version 3 is the first flexible version.
+  "validVersions": "0-3",
+  "flexibleVersions": "3+",
   "fields": [
     { "name": "ThrottleTimeMs", "type": "int32", "versions": "1+", "ignorable": true,
       "about": "The duration in milliseconds for which the request was throttled due to a quota violation, or zero if the request did not violate any quota." },
@@ -27,7 +31,7 @@
       "about": "The error code, or 0 if there was no error." },
     { "name": "ErrorMessage", "type": "string", "versions": "1+", "nullableVersions": "1+", "ignorable": true,
       "about": "The error message, or null if there was no error." },
-    { "name": "NodeId", "type": "int32", "versions": "0+",
+    { "name": "NodeId", "type": "int32", "versions": "0+", "entityType": "brokerId",
       "about": "The node id." },
     { "name": "Host", "type": "string", "versions": "0+",
       "about": "The host name." },
diff --git a/clients/src/main/resources/common/message/HeartbeatRequest.json b/clients/src/main/resources/common/message/HeartbeatRequest.json
index 61cb20e034369..4d799aa3c14c1 100644
--- a/clients/src/main/resources/common/message/HeartbeatRequest.json
+++ b/clients/src/main/resources/common/message/HeartbeatRequest.json
@@ -18,13 +18,21 @@
   "type": "request",
   "name": "HeartbeatRequest",
   // Version 1 and version 2 are the same as version 0.
-  "validVersions": "0-2",
+  //
+  // Starting from version 3, we add a new field called groupInstanceId to indicate member identity across restarts.
+  //
+  // Version 4 is the first flexible version.
+  "validVersions": "0-4",
+  "flexibleVersions": "4+",
   "fields": [
-    { "name": "GroupId", "type": "string", "versions": "0+",
+    { "name": "GroupId", "type": "string", "versions": "0+", "entityType": "groupId",
       "about": "The group id." },
-    { "name": "Generationid", "type": "int32", "versions": "0+",
+    { "name": "GenerationId", "type": "int32", "versions": "0+",
       "about": "The generation of the group." },
     { "name": "MemberId", "type": "string", "versions": "0+",
-      "about": "The member ID." }
+      "about": "The member ID." },
+    { "name": "GroupInstanceId", "type": "string", "versions": "3+",
+      "nullableVersions": "3+", "default": "null",
+      "about": "The unique identifier of the consumer instance provided by end user." }
   ]
 }
diff --git a/clients/src/main/resources/common/message/HeartbeatResponse.json b/clients/src/main/resources/common/message/HeartbeatResponse.json
index c19ba37ad1774..280ba1103b437 100644
--- a/clients/src/main/resources/common/message/HeartbeatResponse.json
+++ b/clients/src/main/resources/common/message/HeartbeatResponse.json
@@ -18,8 +18,14 @@
   "type": "response",
   "name": "HeartbeatResponse",
   // Version 1 adds throttle time.
+  //
   // Starting in version 2, on quota violation, brokers send out responses before throttling.
-  "validVersions": "0-2",
+  //
+  // Starting from version 3, heartbeatRequest supports a new field called groupInstanceId to indicate member identity across restarts.
+  //
+  // Version 4 is the first flexible version.
+  "validVersions": "0-4",
+  "flexibleVersions": "4+",
   "fields": [
     { "name": "ThrottleTimeMs", "type": "int32", "versions": "1+", "ignorable": true,
       "about": "The duration in milliseconds for which the request was throttled due to a quota violation, or zero if the request did not violate any quota." },
diff --git a/clients/src/main/resources/common/message/IncrementalAlterConfigsRequest.json b/clients/src/main/resources/common/message/IncrementalAlterConfigsRequest.json
new file mode 100644
index 0000000000000..b1fb1e9481ac2
--- /dev/null
+++ b/clients/src/main/resources/common/message/IncrementalAlterConfigsRequest.json
@@ -0,0 +1,43 @@
+// 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.
+
+{
+  "apiKey": 44,
+  "type": "request",
+  "name": "IncrementalAlterConfigsRequest",
+  // Version 1 is the first flexible version.
+  "validVersions": "0-1",
+  "flexibleVersions": "1+",
+  "fields": [
+    { "name": "Resources", "type": "[]AlterConfigsResource", "versions": "0+",
+      "about": "The incremental updates for each resource.", "fields": [
+      { "name": "ResourceType", "type": "int8", "versions": "0+", "mapKey": true,
+        "about": "The resource type." },
+      { "name": "ResourceName", "type": "string", "versions": "0+", "mapKey": true,
+        "about": "The resource name." },
+      { "name": "Configs", "type": "[]AlterableConfig", "versions": "0+",
+        "about": "The configurations.",  "fields": [
+        { "name": "Name", "type": "string", "versions": "0+", "mapKey": true,
+          "about": "The configuration key name." },
+        { "name": "ConfigOperation", "type": "int8", "versions": "0+", "mapKey": true,
+          "about": "The type (Set, Delete, Append, Subtract) of operation." },
+        { "name": "Value", "type": "string", "versions": "0+", "nullableVersions": "0+",
+          "about": "The value to set for the configuration key."}
+      ]}
+    ]},
+    { "name": "ValidateOnly", "type": "bool", "versions": "0+",
+      "about": "True if we should validate the request, but not change the configurations."}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/IncrementalAlterConfigsResponse.json b/clients/src/main/resources/common/message/IncrementalAlterConfigsResponse.json
new file mode 100644
index 0000000000000..d4dad294f5b7c
--- /dev/null
+++ b/clients/src/main/resources/common/message/IncrementalAlterConfigsResponse.json
@@ -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.
+
+{
+  "apiKey": 44,
+  "type": "response",
+  "name": "IncrementalAlterConfigsResponse",
+  // Version 1 is the first flexible version.
+  "validVersions": "0-1",
+  "flexibleVersions": "1+",
+  "fields": [
+    { "name": "ThrottleTimeMs", "type": "int32", "versions": "0+",
+      "about": "Duration in milliseconds for which the request was throttled due to a quota violation, or zero if the request did not violate any quota." },
+    { "name": "Responses", "type": "[]AlterConfigsResourceResponse", "versions": "0+",
+      "about": "The responses for each resource.", "fields": [
+      { "name": "ErrorCode", "type": "int16", "versions": "0+",
+        "about": "The resource error code." },
+      { "name": "ErrorMessage", "type": "string", "nullableVersions": "0+", "versions": "0+",
+        "about": "The resource error message, or null if there was no error." },
+      { "name": "ResourceType", "type": "int8", "versions": "0+",
+        "about": "The resource type." },
+      { "name": "ResourceName", "type": "string", "versions": "0+",
+        "about": "The resource name." }
+    ]}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/InitProducerIdRequest.json b/clients/src/main/resources/common/message/InitProducerIdRequest.json
index 8bf2ce35748d4..6c0cbb420ca85 100644
--- a/clients/src/main/resources/common/message/InitProducerIdRequest.json
+++ b/clients/src/main/resources/common/message/InitProducerIdRequest.json
@@ -18,11 +18,14 @@
   "type": "request",
   "name": "InitProducerIdRequest",
   // Version 1 is the same as version 0.
-  "validVersions": "0-1",
+  //
+  // Version 2 is the first flexible version.
+  "validVersions": "0-2",
+  "flexibleVersions": "2+",
   "fields": [
-    { "name": "TransactionalId", "type": "string", "versions": "0+", "nullableVersions": "0+",
+    { "name": "TransactionalId", "type": "string", "versions": "0+", "nullableVersions": "0+", "entityType": "transactionalId",
       "about": "The transactional id, or null if the producer is not transactional." },
     { "name": "TransactionTimeoutMs", "type": "int32", "versions": "0+",
-      "about": "The time in ms to wait for before aborting idle transactions sent by this producer." }
+      "about": "The time in ms to wait for before aborting idle transactions sent by this producer. This is only relevant if a TransactionalId has been defined." }
   ]
 }
diff --git a/clients/src/main/resources/common/message/InitProducerIdResponse.json b/clients/src/main/resources/common/message/InitProducerIdResponse.json
index b251051ec9022..0e8ad2e007e01 100644
--- a/clients/src/main/resources/common/message/InitProducerIdResponse.json
+++ b/clients/src/main/resources/common/message/InitProducerIdResponse.json
@@ -18,14 +18,17 @@
   "type": "response",
   "name": "InitProducerIdResponse",
   // Starting in version 1, on quota violation, brokers send out responses before throttling.
-  "validVersions": "0-1",
+  //
+  // Version 2 is the first flexible version.
+  "validVersions": "0-2",
+  "flexibleVersions": "2+",
   "fields": [
     { "name": "ThrottleTimeMs", "type": "int32", "versions": "0+", "ignorable": true,
       "about": "The duration in milliseconds for which the request was throttled due to a quota violation, or zero if the request did not violate any quota." },
     { "name": "ErrorCode", "type": "int16", "versions": "0+",
       "about": "The error code, or 0 if there was no error." },
-    { "name": "ProducerId", "type": "int64", "versions": "0+",
-      "about": "The current producer id." },
+    { "name": "ProducerId", "type": "int64", "versions": "0+", "entityType": "producerId",
+      "default": -1, "about": "The current producer id." },
     { "name": "ProducerEpoch", "type": "int16", "versions": "0+",
       "about": "The current epoch associated with the producer id." }
   ]
diff --git a/clients/src/main/resources/common/message/JoinGroupRequest.json b/clients/src/main/resources/common/message/JoinGroupRequest.json
index 2a3000cdcf0be..707436f6e9ad4 100644
--- a/clients/src/main/resources/common/message/JoinGroupRequest.json
+++ b/clients/src/main/resources/common/message/JoinGroupRequest.json
@@ -18,12 +18,19 @@
   "type": "request",
   "name": "JoinGroupRequest",
   // Version 1 adds RebalanceTimeoutMs.
+  //
   // Version 2 and 3 are the same as version 1.
+  //
   // Starting from version 4, the client needs to issue a second request to join group
+  //
+  // Starting from version 5, we add a new field called groupInstanceId to indicate member identity across restarts.
   // with assigned id.
-  "validVersions": "0-4",
+  //
+  // Version 6 is the first flexible version.
+  "validVersions": "0-6",
+  "flexibleVersions": "6+",
   "fields": [
-    { "name": "GroupId", "type": "string", "versions": "0+",
+    { "name": "GroupId", "type": "string", "versions": "0+", "entityType": "groupId",
       "about": "The group identifier." },
     { "name": "SessionTimeoutMs", "type": "int32", "versions": "0+",
       "about": "The coordinator considers the consumer dead if it receives no heartbeat after this timeout in milliseconds." },
@@ -33,6 +40,9 @@
       "about": "The maximum time in milliseconds that the coordinator will wait for each member to rejoin when rebalancing the group." },
     { "name": "MemberId", "type": "string", "versions": "0+",
       "about": "The member id assigned by the group coordinator." },
+    { "name": "GroupInstanceId", "type": "string", "versions": "5+", 
+      "nullableVersions": "5+", "default": "null",
+      "about": "The unique identifier of the consumer instance provided by end user." },
     { "name": "ProtocolType", "type": "string", "versions": "0+",
       "about": "The unique name the for class of protocols implemented by the group we want to join." },
     { "name": "Protocols", "type": "[]JoinGroupRequestProtocol", "versions": "0+",
diff --git a/clients/src/main/resources/common/message/JoinGroupResponse.json b/clients/src/main/resources/common/message/JoinGroupResponse.json
index fb36bb076afe4..2ab1e41cd8efc 100644
--- a/clients/src/main/resources/common/message/JoinGroupResponse.json
+++ b/clients/src/main/resources/common/message/JoinGroupResponse.json
@@ -18,17 +18,25 @@
   "type": "response",
   "name": "JoinGroupResponse",
   // Version 1 is the same as version 0.
+  //
   // Version 2 adds throttle time.
+  //
   // Starting in version 3, on quota violation, brokers send out responses before throttling.
+  //
   // Starting in version 4, the client needs to issue a second request to join group
   // with assigned id.
-  "validVersions": "0-4",
+  //
+  // Version 5 is bumped to apply group.instance.id to identify member across restarts.
+  //
+  // Version 6 is the first flexible version.
+  "validVersions": "0-6",
+  "flexibleVersions": "6+",
   "fields": [
     { "name": "ThrottleTimeMs", "type": "int32", "versions": "2+", "ignorable": true,
       "about": "The duration in milliseconds for which the request was throttled due to a quota violation, or zero if the request did not violate any quota." },
     { "name": "ErrorCode", "type": "int16", "versions": "0+",
       "about": "The error code, or 0 if there was no error." },
-    { "name": "GenerationId", "type": "int32", "versions": "0+",
+    { "name": "GenerationId", "type": "int32", "versions": "0+", "default": "-1",
       "about": "The generation ID of the group." },
     { "name": "ProtocolName", "type": "string", "versions": "0+",
       "about": "The group protocol selected by the coordinator." },
@@ -39,6 +47,9 @@
     { "name": "Members", "type": "[]JoinGroupResponseMember", "versions": "0+", "fields": [
       { "name": "MemberId", "type": "string", "versions": "0+",
         "about": "The group member ID." },
+      { "name": "GroupInstanceId", "type": "string", "versions": "5+",
+        "nullableVersions": "5+", "default": "null",
+        "about": "The unique identifier of the consumer instance provided by end user." },
       { "name": "Metadata", "type": "bytes", "versions": "0+",
         "about": "The group member metadata." }
     ]}
diff --git a/clients/src/main/resources/common/message/LeaderAndIsrRequest.json b/clients/src/main/resources/common/message/LeaderAndIsrRequest.json
index b8988351c8da0..3a1f01e4790b0 100644
--- a/clients/src/main/resources/common/message/LeaderAndIsrRequest.json
+++ b/clients/src/main/resources/common/message/LeaderAndIsrRequest.json
@@ -20,67 +20,69 @@
   // Version 1 adds IsNew.
   //
   // Version 2 adds broker epoch and reorganizes the partitions by topic.
-  "validVersions": "0-2",
+  //
+  // Version 3 adds the max broker epoch to make the UpdateMetadataRequest cacheable
+  //
+  // Version 4 adds AddingReplicas and RemovingReplicas
+  //
+  // Version 5 is the first flexible version.
+  "validVersions": "0-5",
+  "flexibleVersions": "5+",
   "fields": [
-    { "name": "ControllerId", "type": "int32", "versions": "0+",
+    { "name": "ControllerId", "type": "int32", "versions": "0+", "entityType": "brokerId",
       "about": "The current controller ID." },
     { "name": "ControllerEpoch", "type": "int32", "versions": "0+",
       "about": "The current controller epoch." },
-    { "name": "BrokerEpoch", "type": "int64", "versions": "2+", "ignorable": true, "default": "-1",
+    { "name": "BrokerEpoch", "type": "int64", "versions": "2", "ignorable": true, "default": "-1",
       "about": "The current broker epoch." },
-    { "name": "TopicStates", "type": "[]LeaderAndIsrRequestTopicState", "versions": "2+",
+    { "name": "UngroupedPartitionStates", "type": "[]LeaderAndIsrPartitionState", "versions": "0-1",
+      "about": "The state of each partition, in a v0 or v1 message." },
+    { "name": "MaxBrokerEpoch", "type": "int64", "versions": "3+", "ignorable": true, "default": "-1",
+      "about": "The max broker epoch." },
+    // In v0 or v1 requests, each partition is listed alongside its topic name.
+    // In v2+ requests, partitions are organized by topic, so that each topic name
+    // only needs to be listed once.
+    { "name": "TopicStates", "type": "[]LeaderAndIsrTopicState", "versions": "2+",
       "about": "Each topic.", "fields": [
-      { "name": "Name", "type": "string", "versions": "2+",
+      { "name": "TopicName", "type": "string", "versions": "2+", "entityType": "topicName",
         "about": "The topic name." },
-      { "name": "PartitionStates", "type": "[]LeaderAndIsrRequestPartitionState", "versions": "0+",
-        "about": "The state of each partition", "fields": [
-        { "name": "PartitionIndex", "type": "int32", "versions": "0+",
-          "about": "The partition index." },
-        { "name": "ControllerEpoch", "type": "int32", "versions": "0+",
-          "about": "The controller epoch." },
-        { "name": "LeaderKey", "type": "int32", "versions": "0+",
-          "about": "The broker ID of the leader." },
-        { "name": "LeaderEpoch", "type": "int32", "versions": "0+",
-          "about": "The leader epoch." },
-        { "name": "IsrReplicas", "type": "[]int32", "versions": "0+",
-          "about": "The in-sync replica IDs." },
-        { "name": "ZkVersion", "type": "int32", "versions": "0+",
-          "about": "The ZooKeeper version." },
-        { "name": "Replicas", "type": "[]int32", "versions": "0+",
-          "about": "The replica IDs." },
-        { "name": "IsNew", "type": "bool", "versions": "1+", "default": "false", "ignorable": true, 
-          "about": "Whether the replica should have existed on the broker or not." }
-      ]}
+      { "name": "PartitionStates", "type": "[]LeaderAndIsrPartitionState", "versions": "2+",
+        "about": "The state of each partition" }
     ]},
-    { "name": "PartitionStatesV0", "type": "[]LeaderAndIsrRequestPartitionStateV0", "versions": "0-1",
-      "about": "The state of each partition", "fields": [
-      { "name": "TopicName", "type": "string", "versions": "0-1",
-        "about": "The topic name." },
-      { "name": "PartitionIndex", "type": "int32", "versions": "0-1",
+    { "name": "LiveLeaders", "type": "[]LeaderAndIsrLiveLeader", "versions": "0+",
+      "about": "The current live leaders.", "fields": [
+      { "name": "BrokerId", "type": "int32", "versions": "0+", "entityType": "brokerId",
+        "about": "The leader's broker ID." },
+      { "name": "HostName", "type": "string", "versions": "0+",
+        "about": "The leader's hostname." },
+      { "name": "Port", "type": "int32", "versions": "0+",
+        "about": "The leader's port." }
+    ]}
+  ],
+  "commonStructs": [
+    { "name": "LeaderAndIsrPartitionState", "versions": "0+", "fields": [
+      { "name": "TopicName", "type": "string", "versions": "0-1", "entityType": "topicName", "ignorable": true,
+        "about": "The topic name.  This is only present in v0 or v1." },
+      { "name": "PartitionIndex", "type": "int32", "versions": "0+",
         "about": "The partition index." },
-      { "name": "ControllerEpoch", "type": "int32", "versions": "0-1",
+      { "name": "ControllerEpoch", "type": "int32", "versions": "0+",
         "about": "The controller epoch." },
-      { "name": "LeaderKey", "type": "int32", "versions": "0-1",
+      { "name": "Leader", "type": "int32", "versions": "0+", "entityType": "brokerId",
         "about": "The broker ID of the leader." },
-      { "name": "LeaderEpoch", "type": "int32", "versions": "0-1",
+      { "name": "LeaderEpoch", "type": "int32", "versions": "0+",
         "about": "The leader epoch." },
-      { "name": "IsrReplicas", "type": "[]int32", "versions": "0-1",
+      { "name": "Isr", "type": "[]int32", "versions": "0+",
         "about": "The in-sync replica IDs." },
-      { "name": "ZkVersion", "type": "int32", "versions": "0-1",
+      { "name": "ZkVersion", "type": "int32", "versions": "0+",
         "about": "The ZooKeeper version." },
-      { "name": "Replicas", "type": "[]int32", "versions": "0-1",
+      { "name": "Replicas", "type": "[]int32", "versions": "0+",
         "about": "The replica IDs." },
-      { "name": "IsNew", "type": "bool", "versions": "1", "default": "false", "ignorable": true, 
+      { "name": "AddingReplicas", "type": "[]int32", "versions": "4+", "ignorable": true,
+        "about": "The replica IDs that we are adding this partition to, or null if no replicas are being added." },
+      { "name": "RemovingReplicas", "type": "[]int32", "versions": "4+", "ignorable": true,
+        "about": "The replica IDs that we are removing this partition from, or null if no replicas are being removed." },
+      { "name": "IsNew", "type": "bool", "versions": "1+", "default": "false", "ignorable": true,
         "about": "Whether the replica should have existed on the broker or not." }
-    ]},
-    { "name": "LiveLeaders", "type": "[]LeaderAndIsrLiveLeader", "versions": "0+",
-      "about": "The current live leaders.", "fields": [
-      { "name": "BrokerId", "type": "int32", "versions": "0+",
-        "about": "The leader's broker ID." },
-      { "name": "HostName", "type": "string", "versions": "0+",
-        "about": "The leader's hostname." },
-      { "name": "Port", "type": "int32", "versions": "0+",
-        "about": "The leader's port." }
     ]}
   ]
 }
diff --git a/clients/src/main/resources/common/message/LeaderAndIsrResponse.json b/clients/src/main/resources/common/message/LeaderAndIsrResponse.json
index e4e1e09098c46..073d89b190ee2 100644
--- a/clients/src/main/resources/common/message/LeaderAndIsrResponse.json
+++ b/clients/src/main/resources/common/message/LeaderAndIsrResponse.json
@@ -20,13 +20,20 @@
   // Version 1 adds KAFKA_STORAGE_ERROR as a valid error code.
   //
   // Version 2 is the same as version 1.
-  "validVersions": "0-2",
+  //
+  // Version 3 is the same as version 2.
+  //
+  // Version 4 is the same as version 3
+  //
+  // Version 5 is is the first flexible version
+  "validVersions": "0-5",
+  "flexibleVersions": "5+",
   "fields": [
     { "name": "ErrorCode", "type": "int16", "versions": "0+",
       "about": "The error code, or 0 if there was no error." },
-    { "name": "Partitions", "type": "[]LeaderAndIsrResponsePartition", "versions": "0+",
+    { "name": "PartitionErrors", "type": "[]LeaderAndIsrPartitionError", "versions": "0+",
       "about": "Each partition.", "fields": [
-      { "name": "TopicName", "type": "string", "versions": "0+",
+      { "name": "TopicName", "type": "string", "versions": "0+", "entityType": "topicName",
         "about": "The topic name." },
       { "name": "PartitionIndex", "type": "int32", "versions": "0+",
         "about": "The partition index." },
diff --git a/clients/src/main/resources/common/message/LeaveGroupRequest.json b/clients/src/main/resources/common/message/LeaveGroupRequest.json
index 9448705fd66b4..acc7938c387b1 100644
--- a/clients/src/main/resources/common/message/LeaveGroupRequest.json
+++ b/clients/src/main/resources/common/message/LeaveGroupRequest.json
@@ -18,11 +18,24 @@
   "type": "request",
   "name": "LeaveGroupRequest",
   // Version 1 and 2 are the same as version 0.
-  "validVersions": "0-2",
+  //
+  // Version 3 defines batch processing scheme with group.instance.id + member.id for identity
+  //
+  // Version 4 is the first flexible version.
+  "validVersions": "0-4",
+  "flexibleVersions": "4+",
   "fields": [
-    { "name": "GroupId", "type": "string", "versions": "0+",
+    { "name": "GroupId", "type": "string", "versions": "0+", "entityType": "groupId",
       "about": "The ID of the group to leave." },
-    { "name": "MemberId", "type": "string", "versions": "0+",
-      "about": "The member ID to remove from the group." }
+    { "name": "MemberId", "type": "string", "versions": "0-2",
+      "about": "The member ID to remove from the group." },
+    { "name": "Members", "type": "[]MemberIdentity", "versions": "3+",
+      "about": "List of leaving member identities.", "fields": [
+      { "name": "MemberId", "type": "string", "versions": "3+",
+        "about": "The member ID to remove from the group." },
+      { "name": "GroupInstanceId", "type": "string",
+        "versions": "3+", "nullableVersions": "3+", "default": "null",
+        "about": "The group instance ID to remove from the group." }
+    ]}
   ]
 }
diff --git a/clients/src/main/resources/common/message/LeaveGroupResponse.json b/clients/src/main/resources/common/message/LeaveGroupResponse.json
index 0d887cdcb0289..0ddb4c6eb5323 100644
--- a/clients/src/main/resources/common/message/LeaveGroupResponse.json
+++ b/clients/src/main/resources/common/message/LeaveGroupResponse.json
@@ -18,12 +18,28 @@
   "type": "response",
   "name": "LeaveGroupResponse",
   // Version 1 adds the throttle time.
+  //
   // Starting in version 2, on quota violation, brokers send out responses before throttling.
-  "validVersions": "0-2",
+  //
+  // Starting in version 3, we will make leave group request into batch mode and add group.instance.id.
+  //
+  // Version 4 is the first flexible version.
+  "validVersions": "0-4",
+  "flexibleVersions": "4+",
   "fields": [
     { "name": "ThrottleTimeMs", "type": "int32", "versions": "1+", "ignorable": true,
       "about": "The duration in milliseconds for which the request was throttled due to a quota violation, or zero if the request did not violate any quota." },
     { "name": "ErrorCode", "type": "int16", "versions": "0+",
-      "about": "The error code, or 0 if there was no error." }
+      "about": "The error code, or 0 if there was no error." },
+
+    { "name": "Members", "type": "[]MemberResponse", "versions": "3+",
+      "about": "List of leaving member responses.", "fields": [
+      { "name": "MemberId", "type": "string", "versions": "3+",
+        "about": "The member ID to remove from the group." },
+      { "name": "GroupInstanceId", "type": "string", "versions": "3+", "nullableVersions": "3+",
+        "about": "The group instance ID to remove from the group." },
+      { "name": "ErrorCode", "type": "int16", "versions": "3+",
+        "about": "The error code, or 0 if there was no error." }
+    ]}
   ]
 }
diff --git a/clients/src/main/resources/common/message/ListGroupsRequest.json b/clients/src/main/resources/common/message/ListGroupsRequest.json
index b0c5aa0e3e81d..f0130e264902b 100644
--- a/clients/src/main/resources/common/message/ListGroupsRequest.json
+++ b/clients/src/main/resources/common/message/ListGroupsRequest.json
@@ -18,7 +18,10 @@
   "type": "request",
   "name": "ListGroupsRequest",
   // Version 1 and 2 are the same as version 0.
-  "validVersions": "0-2",
+  //
+  // Version 3 is the first flexible version.
+  "validVersions": "0-3",
+  "flexibleVersions": "3+",
   "fields": [
   ]
 }
diff --git a/clients/src/main/resources/common/message/ListGroupsResponse.json b/clients/src/main/resources/common/message/ListGroupsResponse.json
index 2dc83fa0121d8..aa8bba6ebc738 100644
--- a/clients/src/main/resources/common/message/ListGroupsResponse.json
+++ b/clients/src/main/resources/common/message/ListGroupsResponse.json
@@ -18,8 +18,12 @@
   "type": "response",
   "name": "ListGroupsResponse",
   // Version 1 adds the throttle time.
+  //
   // Starting in version 2, on quota violation, brokers send out responses before throttling.
-  "validVersions": "0-2",
+  //
+  // Version 3 is the first flexible version.
+  "validVersions": "0-3",
+  "flexibleVersions": "3+",
   "fields": [
     { "name": "ThrottleTimeMs", "type": "int32", "versions": "1+", "ignorable": true,
       "about": "The duration in milliseconds for which the request was throttled due to a quota violation, or zero if the request did not violate any quota." },
@@ -27,7 +31,7 @@
       "about": "The error code, or 0 if there was no error." },
     { "name": "Groups", "type": "[]ListedGroup", "versions": "0+",
       "about": "Each group in the response.", "fields": [
-      { "name": "GroupId", "type": "string", "versions": "0+",
+      { "name": "GroupId", "type": "string", "versions": "0+", "entityType": "groupId",
         "about": "The group ID." },
       { "name": "ProtocolType", "type": "string", "versions": "0+",
         "about": "The group protocol type." }
diff --git a/clients/src/main/resources/common/message/ListOffsetRequest.json b/clients/src/main/resources/common/message/ListOffsetRequest.json
index 8f2c7383e698c..2e12b8b2b7d60 100644
--- a/clients/src/main/resources/common/message/ListOffsetRequest.json
+++ b/clients/src/main/resources/common/message/ListOffsetRequest.json
@@ -19,19 +19,24 @@
   "name": "ListOffsetRequest",
   // Version 1 removes MaxNumOffsets.  From this version forward, only a single
   // offset can be returned.
+  //
   // Version 2 adds the isolation level, which is used for transactional reads.
+  //
   // Version 3 is the same as version 2.
+  //
   // Version 4 adds the current leader epoch, which is used for fencing.
+  //
   // Version 5 is the same as version 5.
   "validVersions": "0-5",
+  "flexibleVersions": "none",
   "fields": [
-    { "name": "ReplicaId", "type": "int32", "versions": "0+",
+    { "name": "ReplicaId", "type": "int32", "versions": "0+", "entityType": "brokerId",
       "about": "The broker ID of the requestor, or -1 if this request is being made by a normal consumer." },
     { "name": "IsolationLevel", "type": "int8", "versions": "2+",
       "about": "This setting controls the visibility of transactional records. Using READ_UNCOMMITTED (isolation_level = 0) makes all records visible. With READ_COMMITTED (isolation_level = 1), non-transactional and COMMITTED transactional records are visible. To be more concrete, READ_COMMITTED returns all data from offsets smaller than the current LSO (last stable offset), and enables the inclusion of the list of aborted transactions in the result, which allows consumers to discard ABORTED transactional records" },
     { "name": "Topics", "type": "[]ListOffsetTopic", "versions": "0+", 
       "about": "Each topic in the request.", "fields": [
-      { "name": "Name", "type": "string", "versions": "0+",
+      { "name": "Name", "type": "string", "versions": "0+", "entityType": "topicName",
         "about": "The topic name." },
       { "name": "Partitions", "type": "[]ListOffsetPartition", "versions": "0+",
         "about": "Each partition in the request.", "fields": [
diff --git a/clients/src/main/resources/common/message/ListOffsetResponse.json b/clients/src/main/resources/common/message/ListOffsetResponse.json
index 9476a192b7f01..4037e1af80bfb 100644
--- a/clients/src/main/resources/common/message/ListOffsetResponse.json
+++ b/clients/src/main/resources/common/message/ListOffsetResponse.json
@@ -19,17 +19,22 @@
   "name": "ListOffsetResponse",
   // Version 1 removes the offsets array in favor of returning a single offset.
   // Version 1 also adds the timestamp associated with the returned offset.
+  //
   // Version 2 adds the throttle time.
+  //
   // Starting in version 3, on quota violation, brokers send out responses before throttling.
+  //
   // Version 4 adds the leader epoch, which is used for fencing.
+  //
   // Version 5 adds a new error code, OFFSET_NOT_AVAILABLE.
   "validVersions": "0-5",
+  "flexibleVersions": "none",
   "fields": [
     { "name": "ThrottleTimeMs", "type": "int32", "versions": "2+", "ignorable": true,
       "about": "The duration in milliseconds for which the request was throttled due to a quota violation, or zero if the request did not violate any quota." },
     { "name": "Topics", "type": "[]ListOffsetTopicResponse", "versions": "0+",
       "about": "Each topic in the response.", "fields": [
-      { "name": "Name", "type": "string", "versions": "0+",
+      { "name": "Name", "type": "string", "versions": "0+", "entityType": "topicName",
         "about": "The topic name" },
       { "name": "Partitions", "type": "[]ListOffsetPartitionResponse", "versions": "0+",
         "about": "Each partition in the response.", "fields": [
diff --git a/clients/src/main/resources/common/message/ListPartitionReassignmentsRequest.json b/clients/src/main/resources/common/message/ListPartitionReassignmentsRequest.json
new file mode 100644
index 0000000000000..7322f25bc456e
--- /dev/null
+++ b/clients/src/main/resources/common/message/ListPartitionReassignmentsRequest.json
@@ -0,0 +1,33 @@
+// 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.
+
+{
+  "apiKey": 46,
+  "type": "request",
+  "name": "ListPartitionReassignmentsRequest",
+  "validVersions": "0",
+  "flexibleVersions": "0+",
+  "fields": [
+    { "name": "TimeoutMs", "type": "int32", "versions": "0+", "default": "60000",
+      "about": "The time in ms to wait for the request to complete." },
+    { "name": "Topics", "type": "[]ListPartitionReassignmentsTopics", "versions": "0+", "nullableVersions": "0+", "default": "null",
+      "about": "The topics to list partition reassignments for, or null to list everything.", "fields": [
+      { "name": "Name", "type": "string", "versions": "0+", "entityType": "topicName",
+        "about": "The topic name" },
+      { "name": "PartitionIndexes", "type": "[]int32", "versions": "0+",
+        "about": "The partitions to list partition reassignments for." }
+    ]}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/ListPartitionReassignmentsResponse.json b/clients/src/main/resources/common/message/ListPartitionReassignmentsResponse.json
new file mode 100644
index 0000000000000..36cf70cab80c9
--- /dev/null
+++ b/clients/src/main/resources/common/message/ListPartitionReassignmentsResponse.json
@@ -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.
+
+{
+  "apiKey": 46,
+  "type": "response",
+  "name": "ListPartitionReassignmentsResponse",
+  "validVersions": "0",
+  "flexibleVersions": "0+",
+  "fields": [
+    { "name": "ThrottleTimeMs", "type": "int32", "versions": "0+",
+      "about": "The duration in milliseconds for which the request was throttled due to a quota violation, or zero if the request did not violate any quota." },
+    { "name": "ErrorCode", "type": "int16", "versions": "0+",
+      "about": "The top-level error code, or 0 if there was no error" },
+    { "name": "ErrorMessage", "type": "string", "versions": "0+", "nullableVersions": "0+",
+      "about": "The top-level error message, or null if there was no error." },
+    { "name": "Topics", "type": "[]OngoingTopicReassignment", "versions": "0+",
+      "about": "The ongoing reassignments for each topic.", "fields": [
+      { "name": "Name", "type": "string", "versions": "0+", "entityType": "topicName",
+        "about": "The topic name." },
+      { "name": "Partitions", "type": "[]OngoingPartitionReassignment", "versions": "0+",
+        "about": "The ongoing reassignments for each partition.", "fields": [
+        { "name": "PartitionIndex", "type": "int32", "versions": "0+",
+          "about": "The index of the partition." },
+        { "name": "Replicas", "type": "[]int32", "versions": "0+",
+          "about": "The current replica set." },
+        { "name": "AddingReplicas", "type": "[]int32", "versions": "0+",
+          "about": "The set of replicas we are currently adding." },
+        { "name": "RemovingReplicas", "type": "[]int32", "versions": "0+",
+          "about": "The set of replicas we are currently removing." }
+      ]}
+    ]}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/MetadataRequest.json b/clients/src/main/resources/common/message/MetadataRequest.json
index 74f3fab5cd585..6129d427b7eef 100644
--- a/clients/src/main/resources/common/message/MetadataRequest.json
+++ b/clients/src/main/resources/common/message/MetadataRequest.json
@@ -17,7 +17,8 @@
   "apiKey": 3,
   "type": "request",
   "name": "MetadataRequest",
-  "validVersions": "0-7",
+  "validVersions": "0-9",
+  "flexibleVersions": "9+",
   "fields": [
     // In version 0, an empty array indicates "request metadata for all topics."  In version 1 and
     // higher, an empty array indicates "request metadata for no topics," and a null array is used to
@@ -26,12 +27,20 @@
     // Version 2 and 3 are the same as version 1.
     //
     // Version 4 adds AllowAutoTopicCreation.
+    //
+    // Starting in version 8, authorized operations can be requested for cluster and topic resource.
+    //
+    // Version 9 is the first flexible version.
     { "name": "Topics", "type": "[]MetadataRequestTopic", "versions": "0+", "nullableVersions": "1+",
       "about": "The topics to fetch metadata for.", "fields": [
-      { "name": "Name", "type": "string", "versions": "0+",
+      { "name": "Name", "type": "string", "versions": "0+", "entityType": "topicName",
         "about": "The topic name." }
     ]},
     { "name": "AllowAutoTopicCreation", "type": "bool", "versions": "4+", "default": "true", "ignorable": false,
-      "about": "If this is true, the broker may auto-create topics that we requested which do not already exist, if it is configured to do so." }
+      "about": "If this is true, the broker may auto-create topics that we requested which do not already exist, if it is configured to do so." },
+    { "name": "IncludeClusterAuthorizedOperations", "type": "bool", "versions": "8+",
+      "about": "Whether to include cluster authorized operations." },
+    { "name": "IncludeTopicAuthorizedOperations", "type": "bool", "versions": "8+",
+      "about": "Whether to include topic authorized operations." }
   ]
 }
diff --git a/clients/src/main/resources/common/message/MetadataResponse.json b/clients/src/main/resources/common/message/MetadataResponse.json
index e58a720c23e42..ce232b1233110 100644
--- a/clients/src/main/resources/common/message/MetadataResponse.json
+++ b/clients/src/main/resources/common/message/MetadataResponse.json
@@ -32,13 +32,18 @@
   // Starting in version 6, on quota violation, brokers send out responses before throttling.
   //
   // Version 7 adds the leader epoch to the partition metadata.
-  "validVersions": "0-7",
+  //
+  // Starting in version 8, brokers can send authorized operations for topic and cluster.
+  //
+  // Version 9 is the first flexible version.
+  "validVersions": "0-9",
+  "flexibleVersions": "9+",
   "fields": [
     { "name": "ThrottleTimeMs", "type": "int32", "versions": "3+",
       "about": "The duration in milliseconds for which the request was throttled due to a quota violation, or zero if the request did not violate any quota." },
     { "name": "Brokers", "type": "[]MetadataResponseBroker", "versions": "0+",
       "about": "Each broker in the response.", "fields": [
-      { "name": "NodeId", "type": "int32", "versions": "0+", "mapKey": true,
+      { "name": "NodeId", "type": "int32", "versions": "0+", "mapKey": true, "entityType": "brokerId",
         "about": "The broker ID." },
       { "name": "Host", "type": "string", "versions": "0+",
         "about": "The broker hostname." },
@@ -47,15 +52,15 @@
       { "name": "Rack", "type": "string", "versions": "1+", "nullableVersions": "1+", "ignorable": true, "default": "null",
         "about": "The rack of the broker, or null if it has not been assigned to a rack." }
     ]},
-    { "name": "ClusterId", "type": "string", "nullableVersions": "2+", "versions": "2+", "ignorable": true,
+    { "name": "ClusterId", "type": "string", "nullableVersions": "2+", "versions": "2+", "ignorable": true, "default": "null",
       "about": "The cluster ID that responding broker belongs to." },
-    { "name": "ControllerId", "type": "int32", "versions": "1+", "default": "-1", "ignorable": true,
+    { "name": "ControllerId", "type": "int32", "versions": "1+", "default": "-1", "ignorable": true, "entityType": "brokerId",
       "about": "The ID of the controller broker." },
     { "name": "Topics", "type": "[]MetadataResponseTopic", "versions": "0+",
       "about": "Each topic in the response.", "fields": [
       { "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,
+      { "name": "Name", "type": "string", "versions": "0+", "mapKey": true, "entityType": "topicName",
         "about": "The topic name." },
       { "name": "IsInternal", "type": "bool", "versions": "1+", "default": "false", "ignorable": true,
         "about": "True if the topic is internal." },
@@ -65,17 +70,21 @@
           "about": "The partition error, or 0 if there was no error." },
         { "name": "PartitionIndex", "type": "int32", "versions": "0+",
           "about": "The partition index." },
-        { "name": "LeaderId", "type": "int32", "versions": "0+",
+        { "name": "LeaderId", "type": "int32", "versions": "0+", "entityType": "brokerId",
           "about": "The ID of the leader broker." },
         { "name": "LeaderEpoch", "type": "int32", "versions": "7+", "default": "-1", "ignorable": true,
           "about": "The leader epoch of this partition." },
-        { "name": "ReplicaNodes", "type": "[]int32", "versions": "0+",
+        { "name": "ReplicaNodes", "type": "[]int32", "versions": "0+", "entityType": "brokerId",
           "about": "The set of all nodes that host this partition." },
         { "name": "IsrNodes", "type": "[]int32", "versions": "0+",
           "about": "The set of nodes that are in sync with the leader for this partition." },
         { "name": "OfflineReplicas", "type": "[]int32", "versions": "5+", "ignorable": true,
           "about": "The set of offline replicas of this partition." }
-      ]}
-    ]}
+      ]},
+      { "name": "TopicAuthorizedOperations", "type": "int32", "versions": "8+", "default": "-2147483648",
+        "about": "32-bit bitfield to represent authorized operations for this topic." }
+    ]},
+    { "name": "ClusterAuthorizedOperations", "type": "int32", "versions": "8+", "default": "-2147483648",
+      "about": "32-bit bitfield to represent authorized operations for this cluster." }
   ]
 }
diff --git a/clients/src/main/resources/common/message/OffsetCommitRequest.json b/clients/src/main/resources/common/message/OffsetCommitRequest.json
index ebccedef58772..096b61917ae0f 100644
--- a/clients/src/main/resources/common/message/OffsetCommitRequest.json
+++ b/clients/src/main/resources/common/message/OffsetCommitRequest.json
@@ -26,19 +26,27 @@
   // Version 5 removes the retention time, which is now controlled only by a broker configuration.
   //
   // Version 6 adds the leader epoch for fencing.
-  "validVersions": "0-6",
+  //
+  // version 7 adds a new field called groupInstanceId to indicate member identity across restarts.
+  //
+  // Version 8 is the first flexible version.
+  "validVersions": "0-8",
+  "flexibleVersions": "8+",
   "fields": [
-    { "name": "GroupId", "type": "string", "versions": "0+",
+    { "name": "GroupId", "type": "string", "versions": "0+", "entityType": "groupId",
       "about": "The unique group identifier." },
     { "name": "GenerationId", "type": "int32", "versions": "1+", "default": "-1", "ignorable": true,
       "about": "The generation of the group." },
     { "name": "MemberId", "type": "string", "versions": "1+", "ignorable": true,
       "about": "The member ID assigned by the group coordinator." },
+    { "name": "GroupInstanceId", "type": "string", "versions": "7+",
+      "nullableVersions": "7+", "default": "null",
+      "about": "The unique identifier of the consumer instance provided by end user." },
     { "name": "RetentionTimeMs", "type": "int64", "versions": "2-4", "default": "-1", "ignorable": true,
       "about": "The time period in ms to retain the offset." },
     { "name": "Topics", "type": "[]OffsetCommitRequestTopic", "versions": "0+",
       "about": "The topics to commit offsets for.",  "fields": [
-      { "name": "Name", "type": "string", "versions": "0+",
+      { "name": "Name", "type": "string", "versions": "0+", "entityType": "topicName",
         "about": "The topic name." },
       { "name": "Partitions", "type": "[]OffsetCommitRequestPartition", "versions": "0+",
         "about": "Each partition to commit offsets for.", "fields": [
diff --git a/clients/src/main/resources/common/message/OffsetCommitResponse.json b/clients/src/main/resources/common/message/OffsetCommitResponse.json
index 39daa56562f50..3d547794cfb42 100644
--- a/clients/src/main/resources/common/message/OffsetCommitResponse.json
+++ b/clients/src/main/resources/common/message/OffsetCommitResponse.json
@@ -24,13 +24,18 @@
   // Starting in version 4, on quota violation, brokers send out responses before throttling.
   //
   // Versions 5 and 6 are the same as version 4.
-  "validVersions": "0-6",
+  //
+  // Version 7 offsetCommitRequest supports a new field called groupInstanceId to indicate member identity across restarts.
+  //
+  // Version 8 is the first flexible version.
+  "validVersions": "0-8",
+  "flexibleVersions": "8+",
   "fields": [
     { "name": "ThrottleTimeMs", "type": "int32", "versions": "3+", "ignorable": true,
       "about": "The duration in milliseconds for which the request was throttled due to a quota violation, or zero if the request did not violate any quota." },
     { "name": "Topics", "type": "[]OffsetCommitResponseTopic", "versions": "0+",
       "about": "The responses for each topic.", "fields": [
-      { "name": "Name", "type": "string", "versions": "0+",
+      { "name": "Name", "type": "string", "versions": "0+", "entityType": "topicName",
         "about": "The topic name." },
       { "name": "Partitions", "type": "[]OffsetCommitResponsePartition", "versions": "0+",
         "about": "The responses for each partition in the topic.",  "fields": [
diff --git a/clients/src/main/resources/common/message/OffsetDeleteRequest.json b/clients/src/main/resources/common/message/OffsetDeleteRequest.json
new file mode 100644
index 0000000000000..563594997097c
--- /dev/null
+++ b/clients/src/main/resources/common/message/OffsetDeleteRequest.json
@@ -0,0 +1,37 @@
+// 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.
+
+{
+  "apiKey": 47,
+  "type": "request",
+  "name": "OffsetDeleteRequest",
+  "validVersions": "0",
+  "fields": [
+    { "name": "GroupId", "type": "string", "versions": "0+",
+      "about": "The unique group identifier." },
+    { "name": "Topics", "type": "[]OffsetDeleteRequestTopic", "versions": "0+",
+      "about": "The topics to delete offsets for", "fields": [
+        { "name": "Name",  "type": "string",  "versions": "0+", "mapKey": true,
+          "about": "The topic name." },
+        { "name": "Partitions", "type": "[]OffsetDeleteRequestPartition", "versions": "0+",
+          "about": "Each partition to delete offsets for.", "fields": [
+            { "name": "PartitionIndex", "type": "int32", "versions": "0+",
+              "about": "The partition index." }
+          ]
+        }
+      ]
+    }
+  ]
+}
diff --git a/clients/src/main/resources/common/message/OffsetDeleteResponse.json b/clients/src/main/resources/common/message/OffsetDeleteResponse.json
new file mode 100644
index 0000000000000..f2be321f5d8fd
--- /dev/null
+++ b/clients/src/main/resources/common/message/OffsetDeleteResponse.json
@@ -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.
+
+{
+  "apiKey": 47,
+  "type": "response",
+  "name": "OffsetDeleteResponse",
+  "validVersions": "0",
+  "fields": [
+    { "name": "ErrorCode", "type": "int16", "versions": "0+",
+      "about": "The top-level error code, or 0 if there was no error." },
+    { "name": "ThrottleTimeMs",  "type": "int32",  "versions": "0+", "ignorable": true,
+      "about": "The duration in milliseconds for which the request was throttled due to a quota violation, or zero if the request did not violate any quota." },
+    { "name": "Topics", "type": "[]OffsetDeleteResponseTopic", "versions": "0+",
+      "about": "The responses for each topic.", "fields": [
+        { "name": "Name", "type": "string", "versions": "0+", "mapKey": true,
+          "about": "The topic name." },
+        { "name": "Partitions", "type": "[]OffsetDeleteResponsePartition", "versions": "0+",
+          "about": "The responses for each partition in the topic.", "fields": [
+            { "name": "PartitionIndex", "type": "int32", "versions": "0+", "mapKey": true,
+              "about": "The partition index." },
+            { "name": "ErrorCode", "type": "int16", "versions": "0+",
+              "about": "The error code, or 0 if there was no error." }
+          ]
+        }
+      ]
+    }
+  ]
+}
diff --git a/clients/src/main/resources/common/message/OffsetFetchRequest.json b/clients/src/main/resources/common/message/OffsetFetchRequest.json
index e634f7c1a279b..57306880df415 100644
--- a/clients/src/main/resources/common/message/OffsetFetchRequest.json
+++ b/clients/src/main/resources/common/message/OffsetFetchRequest.json
@@ -17,19 +17,26 @@
   "apiKey": 9,
   "type": "request",
   "name": "OffsetFetchRequest",
+  // In version 0, the request read offsets from ZK.
+  //
   // Starting in version 1, the broker supports fetching offsets from the internal __consumer_offsets topic.
   //
   // Starting in version 2, the request can contain a null topics array to indicate that offsets
-  // for all topics should be fetched.
+  // for all topics should be fetched. It also returns a top level error code
+  // for group or coordinator level errors.
   //
   // Version 3, 4, and 5 are the same as version 2.
-  "validVersions": "0-5",
+  //
+  // Version 6 is the first flexible version.
+  "validVersions": "0-6",
+  "flexibleVersions": "6+",
   "fields": [
-    { "name": "GroupId", "type": "string", "versions": "0+",
+    { "name": "GroupId", "type": "string", "versions": "0+", "entityType": "groupId",
       "about": "The group to fetch offsets for." },
     { "name": "Topics", "type": "[]OffsetFetchRequestTopic", "versions": "0+", "nullableVersions": "2+",
       "about": "Each topic we would like to fetch offsets for, or null to fetch offsets for all topics.", "fields": [
-      { "name": "Name", "type": "string", "versions": "0+" },
+      { "name": "Name", "type": "string", "versions": "0+", "entityType": "topicName",
+        "about": "The topic name."},
       { "name": "PartitionIndexes", "type": "[]int32", "versions": "0+",
         "about": "The partition indexes we would like to fetch offsets for." }
     ]}
diff --git a/clients/src/main/resources/common/message/OffsetFetchResponse.json b/clients/src/main/resources/common/message/OffsetFetchResponse.json
index 70fd27712791f..e0d5e358a9173 100644
--- a/clients/src/main/resources/common/message/OffsetFetchResponse.json
+++ b/clients/src/main/resources/common/message/OffsetFetchResponse.json
@@ -26,13 +26,16 @@
   // Starting in version 4, on quota violation, brokers send out responses before throttling.
   //
   // Version 5 adds the leader epoch to the committed offset.
-  "validVersions": "0-5",
+  //
+  // Version 6 is the first flexible version.
+  "validVersions": "0-6",
+  "flexibleVersions": "6+",
   "fields": [
     { "name": "ThrottleTimeMs", "type": "int32", "versions": "3+", "ignorable": true,
       "about": "The duration in milliseconds for which the request was throttled due to a quota violation, or zero if the request did not violate any quota." },
     { "name": "Topics", "type": "[]OffsetFetchResponseTopic", "versions": "0+", 
       "about": "The responses per topic.", "fields": [
-      { "name": "Name", "type": "string", "versions": "0+",
+      { "name": "Name", "type": "string", "versions": "0+", "entityType": "topicName",
         "about": "The topic name." },
       { "name": "Partitions", "type": "[]OffsetFetchResponsePartition", "versions": "0+",
         "about": "The responses per partition", "fields": [
@@ -40,7 +43,7 @@
           "about": "The partition index." },
         { "name": "CommittedOffset", "type": "int64", "versions": "0+",
           "about": "The committed message offset." },
-        { "name": "CommittedLeaderEpoch", "type": "int32", "versions": "5+",
+        { "name": "CommittedLeaderEpoch", "type": "int32", "versions": "5+", "default": "-1",
           "about": "The leader epoch." },
         { "name": "Metadata", "type": "string", "versions": "0+", "nullableVersions": "0+",
           "about": "The partition metadata." },
diff --git a/clients/src/main/resources/common/message/OffsetForLeaderEpochRequest.json b/clients/src/main/resources/common/message/OffsetForLeaderEpochRequest.json
index 40227edae5a66..0d1b063b90cba 100644
--- a/clients/src/main/resources/common/message/OffsetForLeaderEpochRequest.json
+++ b/clients/src/main/resources/common/message/OffsetForLeaderEpochRequest.json
@@ -20,11 +20,18 @@
   // Version 1 is the same as version 0.
   //
   // Version 2 adds the current leader epoch to support fencing.
-  "validVersions": "0-2",
+  //
+  // Version 3 adds ReplicaId (the default is -2 which conventionally represents a
+  //    "debug" consumer which is allowed to see offsets beyond the high watermark).
+  //    Followers will use this replicaId when using an older version of the protocol.
+  "validVersions": "0-3",
+  "flexibleVersions": "none",
   "fields": [
+    { "name": "ReplicaId", "type": "int32", "versions": "3+", "default": -2, "ignorable": true,
+      "about": "The broker ID of the follower, of -1 if this request is from a consumer." },
     { "name": "Topics", "type": "[]OffsetForLeaderTopic", "versions": "0+",
       "about": "Each topic to get offsets for.", "fields": [
-      { "name": "Name", "type": "string", "versions": "0+",
+      { "name": "Name", "type": "string", "versions": "0+", "entityType": "topicName",
         "about": "The topic name." },
       { "name": "Partitions", "type": "[]OffsetForLeaderPartition", "versions": "0+",
         "about": "Each partition to get offsets for.", "fields": [
diff --git a/clients/src/main/resources/common/message/OffsetForLeaderEpochResponse.json b/clients/src/main/resources/common/message/OffsetForLeaderEpochResponse.json
index 26bd490fdaf91..867b763a4214b 100644
--- a/clients/src/main/resources/common/message/OffsetForLeaderEpochResponse.json
+++ b/clients/src/main/resources/common/message/OffsetForLeaderEpochResponse.json
@@ -19,13 +19,14 @@
   "name": "OffsetForLeaderEpochResponse",
   // Version 1 added the leader epoch to the response.
   // Version 2 added the throttle time.
-  "validVersions": "0-2",
+  "validVersions": "0-3",
+  "flexibleVersions": "none",
   "fields": [
     { "name": "ThrottleTimeMs", "type": "int32", "versions": "2+", "ignorable": true,
       "about": "The duration in milliseconds for which the request was throttled due to a quota violation, or zero if the request did not violate any quota." },
     { "name": "Topics", "type": "[]OffsetForLeaderTopicResult", "versions": "0+",
       "about": "Each topic we fetched offsets for.", "fields": [
-      { "name": "Name", "type": "string", "versions": "0+",
+      { "name": "Name", "type": "string", "versions": "0+", "entityType": "topicName",
         "about": "The topic name." },
       { "name": "Partitions", "type": "[]OffsetForLeaderPartitionResult", "versions": "0+",
         "about": "Each partition in the topic we fetched offsets for.", "fields": [
diff --git a/clients/src/main/resources/common/message/ProduceRequest.json b/clients/src/main/resources/common/message/ProduceRequest.json
index 4f35db2e712ca..a872da75c7e9b 100644
--- a/clients/src/main/resources/common/message/ProduceRequest.json
+++ b/clients/src/main/resources/common/message/ProduceRequest.json
@@ -28,9 +28,12 @@
   // Version 5 and 6 are the same as version 3.
   //
   // Starting in version 7, records can be produced using ZStandard compression.  See KIP-110.
-  "validVersions": "0-7",
+  //
+  // Starting in Version 8, response has RecordErrors and ErrorMEssage. See KIP-467.
+  "validVersions": "0-8",
+  "flexibleVersions": "none",
   "fields": [
-    { "name": "TransactionalId", "type": "string", "versions": "3+", "nullableVersions": "0+",
+    { "name": "TransactionalId", "type": "string", "versions": "3+", "nullableVersions": "0+", "entityType": "transactionalId",
       "about": "The transactional ID, or null if the producer is not transactional." },
     { "name": "Acks", "type": "int16", "versions": "0+",
       "about": "The number of acknowledgments the producer requires the leader to have received before considering a request complete. Allowed values: 0 for no acknowledgments, 1 for only the leader and -1 for the full ISR." },
@@ -38,7 +41,7 @@
       "about": "The timeout to await a response in miliseconds." },
     { "name": "Topics", "type": "[]TopicProduceData", "versions": "0+", 
       "about": "Each topic to produce to.", "fields": [
-      { "name": "Name", "type": "string", "versions": "0+",
+      { "name": "Name", "type": "string", "versions": "0+", "entityType": "topicName",
         "about": "The topic name." },
       { "name": "Partitions", "type": "[]PartitionProduceData", "versions": "0+",
         "about": "Each partition to produce to.", "fields": [
diff --git a/clients/src/main/resources/common/message/ProduceResponse.json b/clients/src/main/resources/common/message/ProduceResponse.json
index 14d38aa6917db..77ab0655be57f 100644
--- a/clients/src/main/resources/common/message/ProduceResponse.json
+++ b/clients/src/main/resources/common/message/ProduceResponse.json
@@ -27,11 +27,15 @@
   //
   // Version 5 added LogStartOffset to filter out spurious
   // OutOfOrderSequenceExceptions on the client.
-  "validVersions": "0-7",
+  //
+  // Version 8 added RecordErrors and ErrorMessage to include information about
+  // records that cause the whole batch to be dropped.  See KIP-467 for details.
+  "validVersions": "0-8",
+  "flexibleVersions": "none",
   "fields": [
     { "name": "Responses", "type": "[]TopicProduceResponse", "versions": "0+",
       "about": "Each produce response", "fields": [
-      { "name": "Name", "type": "string", "versions": "0+",
+      { "name": "Name", "type": "string", "versions": "0+", "entityType": "topicName",
         "about": "The topic name" },
       { "name": "Partitions", "type": "[]PartitionProduceResponse", "versions": "0+",
         "about": "Each partition that we produced to within the topic.", "fields": [
@@ -44,7 +48,16 @@
         { "name": "LogAppendTimeMs", "type": "int64", "versions": "2+", "default": "-1", "ignorable": true,
           "about": "The timestamp returned by broker after appending the messages. If CreateTime is used for the topic, the timestamp will be -1.  If LogAppendTime is used for the topic, the timestamp will be the broker local time when the messages are appended." },
         { "name": "LogStartOffset", "type": "int64", "versions": "5+", "default": "-1", "ignorable": true,
-          "about": "The log start offset." }
+          "about": "The log start offset." },
+        { "name": "RecordErrors", "type": "[]BatchIndexAndErrorMessage", "versions": "8+", "ignorable": true,
+          "about": "The batch indices of records that caused the batch to be dropped", "fields": [
+          { "name":  "BatchIndex", "type": "int32", "versions":  "8+",
+            "about":  "The batch index of the record that cause the batch to be dropped" },
+          { "name": "BatchIndexErrorMessage", "type": "string", "default": "null", "versions": "8+", "nullableVersions": "8+",
+            "about":  "The error message of the record that caused the batch to be dropped"}
+        ]},
+        { "name":  "ErrorMessage", "type": "string", "default": "null", "versions": "8+", "nullableVersions": "8+", "ignorable":  true,
+          "about":  "The global error message summarizing the common root cause of the records that caused the batch to be dropped"}
       ]}
     ]},
     { "name": "ThrottleTimeMs", "type": "int32", "versions": "1+", "ignorable": true,
diff --git a/clients/src/main/resources/common/message/README.md b/clients/src/main/resources/common/message/README.md
index 2de09bccfe229..64e17b16b5075 100644
--- a/clients/src/main/resources/common/message/README.md
+++ b/clients/src/main/resources/common/message/README.md
@@ -69,25 +69,28 @@ Field Types
 -----------
 There are several primitive field types available.
 
-* "boolean": either true or false.  This takes up 1 byte on the wire.
+* "boolean": either true or false.
 
-* "int8": an 8-bit integer.  This also takes up 1 byte on the wire.
+* "int8": an 8-bit integer.
 
-* "int16": a 16-bit integer.  This takes up 2 bytes on the wire.
+* "int16": a 16-bit integer.
 
-* "int32": a 32-bit integer.  This takes up 4 bytes on the wire.
+* "int32": a 32-bit integer.
 
-* "int64": a 64-bit integer.  This takes up 8 bytes on the wire.
+* "int64": a 64-bit integer.
 
-* "string": a string.  This must be less than 64kb in size when serialized as UTF-8.  This takes up 2 bytes on the wire, plus the length of the string when serialized to UTF-8.
+* "string": a UTF-8 string.
 
-* "bytes": binary data.  This takes up 4 bytes on the wire, plus the length of the bytes.
+* "bytes": binary data.
 
 In addition to these primitive field types, there is also an array type.  Array
 types start with a "[]" and end with the name of the element type.  For
 example, []Foo declares an array of "Foo" objects.  Array fields have their own
 array of fields, which specifies what is in the contained objects.
 
+For information about how fields are serialized, see the [Kafka Protocol
+Guide](https://kafka.apache.org/protocol.html).
+
 Nullable Fields
 ---------------
 Booleans and ints can never be null.  However, fields that are strings, bytes,
@@ -104,6 +107,39 @@ If a field is declared as non-nullable, and it is present in the message
 version you are using, you should set it to a non-null value before serializing
 the message.  Otherwise, you will get a runtime error.
 
+Tagged Fields
+-------------
+Tagged fields are an extension to the Kafka protocol which allows optional data
+to be attached to messages.  Tagged fields can appear at the root level of
+messages, or within any structure in the message.
+
+Unlike mandatory fields, tagged fields can be added to message versions that
+already exists.  Older servers will ignore new tagged fields which they do not
+understand.
+
+In order to make a field tagged, set a "tag" for the field, and also set up
+tagged versions for the field.  The taggedVersions you specify should be
+open-ended-- that is, they should specify a start version, but not an end
+version.
+
+You can remove support for a tagged field from a specific version of a message,
+but you can't reuse a tag once it has been used for something else.  Once tags
+have been used for something, they can't be used for anything else, without
+breaking compatibilty.
+
+Note that tagged fields can only be added to "flexible" message versions.
+
+Flexible Versions
+-----------------
+Kafka serialization has been improved over time to be more flexible and
+efficient.  Message versions that contain these improvements are referred to as
+"flexible versions."
+
+In flexible verisons, variable-length fields such as strings, arrays, and bytes
+fields are serialized in a more efficient way that saves space.  The new
+serialization types start with compact.  For example COMPACT_STRING is a more
+efficient form of STRING.
+
 Serializing Messages
 --------------------
 The Message#write method writes out a message to a buffer.  The fields that are
diff --git a/clients/src/main/resources/common/message/RenewDelegationTokenRequest.json b/clients/src/main/resources/common/message/RenewDelegationTokenRequest.json
index ba1db7ede5eb8..83fde2ee837e6 100644
--- a/clients/src/main/resources/common/message/RenewDelegationTokenRequest.json
+++ b/clients/src/main/resources/common/message/RenewDelegationTokenRequest.json
@@ -19,6 +19,7 @@
   "name": "RenewDelegationTokenRequest",
   // Version 1 is the same as version 0.
   "validVersions": "0-1",
+  "flexibleVersions": "none",
   "fields": [
     { "name": "Hmac", "type": "bytes", "versions": "0+",
       "about": "The HMAC of the delegation token to be renewed." },
diff --git a/clients/src/main/resources/common/message/RenewDelegationTokenResponse.json b/clients/src/main/resources/common/message/RenewDelegationTokenResponse.json
index aed734a879645..5c70521164563 100644
--- a/clients/src/main/resources/common/message/RenewDelegationTokenResponse.json
+++ b/clients/src/main/resources/common/message/RenewDelegationTokenResponse.json
@@ -19,6 +19,7 @@
   "name": "RenewDelegationTokenResponse",
   // Starting in version 1, on quota violation, brokers send out responses before throttling.
   "validVersions": "0-1",
+  "flexibleVersions": "none",
   "fields": [
     { "name": "ErrorCode", "type": "int16", "versions": "0+",
       "about": "The error code, or 0 if there was no error." },
diff --git a/clients/src/main/resources/common/message/RequestHeader.json b/clients/src/main/resources/common/message/RequestHeader.json
index d24dcf0e98755..fbf4e2ca9f454 100644
--- a/clients/src/main/resources/common/message/RequestHeader.json
+++ b/clients/src/main/resources/common/message/RequestHeader.json
@@ -16,7 +16,13 @@
 {
   "type": "header",
   "name": "RequestHeader",
-  "validVersions": "0",
+  // Version 0 of the RequestHeader is only used by v0 of ControlledShutdownRequest.
+  //
+  // Version 1 is the first version with ClientId.
+  //
+  // Version 2 is the first flexible version.
+  "validVersions": "0-2",
+  "flexibleVersions": "2+",
   "fields": [
     { "name": "RequestApiKey", "type": "int16", "versions": "0+",
       "about": "The API key of this request." },
@@ -24,7 +30,13 @@
       "about": "The API version of this request." },
     { "name": "CorrelationId", "type": "int32", "versions": "0+",
       "about": "The correlation ID of this request." },
-    { "name": "ClientId", "type": "string", "versions": "0+",
-      "about": "The client ID string." }
+
+    // The ClientId string must be serialized with the old-style two-byte length prefix.
+    // The reason is that older brokers must be able to read the request header for any
+    // ApiVersionsRequest, even if it is from a newer version.
+    // Since the client is sending the ApiVersionsRequest in order to discover what
+    // versions are supported, the client does not know the best version to use.
+    { "name": "ClientId", "type": "string", "versions": "1+", "nullableVersions": "1+", "ignorable": true,
+      "flexibleVersions": "none", "about": "The client ID string." }
   ]
 }
diff --git a/clients/src/main/resources/common/message/ResponseHeader.json b/clients/src/main/resources/common/message/ResponseHeader.json
index d44825977da7d..773673601a886 100644
--- a/clients/src/main/resources/common/message/ResponseHeader.json
+++ b/clients/src/main/resources/common/message/ResponseHeader.json
@@ -16,7 +16,9 @@
 {
   "type": "header",
   "name": "ResponseHeader",
-  "validVersions": "0",
+  // Version 1 is the first flexible version.
+  "validVersions": "0-1",
+  "flexibleVersions": "1+",
   "fields": [
     { "name": "CorrelationId", "type": "int32", "versions": "0+",
       "about": "The correlation ID of this response." }
diff --git a/clients/src/main/resources/common/message/SaslAuthenticateRequest.json b/clients/src/main/resources/common/message/SaslAuthenticateRequest.json
index 96f6f3aa4efb4..27e97700304f4 100644
--- a/clients/src/main/resources/common/message/SaslAuthenticateRequest.json
+++ b/clients/src/main/resources/common/message/SaslAuthenticateRequest.json
@@ -19,6 +19,7 @@
   "name": "SaslAuthenticateRequest",
   // Version 1 is the same as version 0.
   "validVersions": "0-1",
+  "flexibleVersions": "none",
   "fields": [
     { "name": "AuthBytes", "type": "bytes", "versions": "0+",
       "about": "The SASL authentication bytes from the client, as defined by the SASL mechanism." }
diff --git a/clients/src/main/resources/common/message/SaslAuthenticateResponse.json b/clients/src/main/resources/common/message/SaslAuthenticateResponse.json
index f6644f6d3459d..66dfe74651fec 100644
--- a/clients/src/main/resources/common/message/SaslAuthenticateResponse.json
+++ b/clients/src/main/resources/common/message/SaslAuthenticateResponse.json
@@ -19,6 +19,7 @@
   "name": "SaslAuthenticateResponse",
   // Version 1 adds the session lifetime.
   "validVersions": "0-1",
+  "flexibleVersions": "none",
   "fields": [
     { "name": "ErrorCode", "type": "int16", "versions": "0+",
       "about": "The error code, or 0 if there was no error." },
diff --git a/clients/src/main/resources/common/message/SaslHandshakeRequest.json b/clients/src/main/resources/common/message/SaslHandshakeRequest.json
index 7ad0a4716ef74..162448238dfe8 100644
--- a/clients/src/main/resources/common/message/SaslHandshakeRequest.json
+++ b/clients/src/main/resources/common/message/SaslHandshakeRequest.json
@@ -19,6 +19,7 @@
   "name": "SaslHandshakeRequest",
   // Version 1 supports SASL_AUTHENTICATE.
   "validVersions": "0-1",
+  "flexibleVersions": "none",
   "fields": [
     { "name": "Mechanism", "type": "string", "versions": "0+",
       "about": "The SASL mechanism chosen by the client." }
diff --git a/clients/src/main/resources/common/message/SaslHandshakeResponse.json b/clients/src/main/resources/common/message/SaslHandshakeResponse.json
index 821f3290833cf..2d26945f433f5 100644
--- a/clients/src/main/resources/common/message/SaslHandshakeResponse.json
+++ b/clients/src/main/resources/common/message/SaslHandshakeResponse.json
@@ -19,6 +19,7 @@
   "name": "SaslHandshakeResponse",
   // Version 1 is the same as version 0.
   "validVersions": "0-1",
+  "flexibleVersions": "none",
   "fields": [
     { "name": "ErrorCode", "type": "int16", "versions": "0+",
       "about": "The error code, or 0 if there was no error." },
diff --git a/clients/src/main/resources/common/message/StopReplicaRequest.json b/clients/src/main/resources/common/message/StopReplicaRequest.json
index dffa11da9ee4c..3d016b0fe181c 100644
--- a/clients/src/main/resources/common/message/StopReplicaRequest.json
+++ b/clients/src/main/resources/common/message/StopReplicaRequest.json
@@ -19,26 +19,33 @@
   "name": "StopReplicaRequest",
   // Version 1 adds the broker epoch and reorganizes the partitions to be stored
   // per topic.
-  "validVersions": "0-1",
+  //
+  // Version 2 adds the max broker epoch to make the UpdateMetadataRequest cacheable
+  //
+  // Version 3 is the first flexible version.
+  "validVersions": "0-3",
+  "flexibleVersions": "3+",
   "fields": [
-    { "name": "ControllerId", "type": "int32", "versions": "0+",
+    { "name": "ControllerId", "type": "int32", "versions": "0+", "entityType": "brokerId",
       "about": "The controller id." },
     { "name": "ControllerEpoch", "type": "int32", "versions": "0+",
       "about": "The controller epoch." },
-    { "name": "BrokerEpoch", "type": "int64", "versions": "1+", "default": "-1", "ignorable": true,
+    { "name": "BrokerEpoch", "type": "int64", "versions": "1", "default": "-1", "ignorable": true,
       "about": "The broker epoch." },
+    { "name": "MaxBrokerEpoch", "type": "int64", "versions": "2+", "default": "-1", "ignorable": true,
+      "about": "The max broker epoch." },
     { "name": "DeletePartitions", "type": "bool", "versions": "0+",
       "about": "Whether these partitions should be deleted." },
-    { "name": "PartitionsV0", "type": "[]StopReplicaRequestPartitionV0", "versions": "0",
+    { "name": "UngroupedPartitions", "type": "[]StopReplicaPartitionV0", "versions": "0",
       "about": "The partitions to stop.", "fields": [
-      { "name": "TopicName", "type": "string", "versions": "0",
+      { "name": "TopicName", "type": "string", "versions": "0", "entityType": "topicName",
         "about": "The topic name." },
       { "name": "PartitionIndex", "type": "int32", "versions": "0",
         "about": "The partition index." }
     ]},
-    { "name": "Topics", "type": "[]StopReplicaRequestTopic", "versions": "1+",
+    { "name": "Topics", "type": "[]StopReplicaTopic", "versions": "1+",
       "about": "The topics to stop.", "fields": [
-      { "name": "Name", "type": "string", "versions": "1+",
+      { "name": "Name", "type": "string", "versions": "1+", "entityType": "topicName",
         "about": "The topic name." },
       { "name": "PartitionIndexes", "type": "[]int32", "versions": "1+",
         "about": "The partition indexes." }
diff --git a/clients/src/main/resources/common/message/StopReplicaResponse.json b/clients/src/main/resources/common/message/StopReplicaResponse.json
index 55daac5e9c391..6fb5128bab9e1 100644
--- a/clients/src/main/resources/common/message/StopReplicaResponse.json
+++ b/clients/src/main/resources/common/message/StopReplicaResponse.json
@@ -18,13 +18,16 @@
   "type": "response",
   "name": "StopReplicaResponse",
   // Version 1 is the same as version 0.
-  "validVersions": "0-1",
+  // Version 2 is the same as version 1.
+  // Version 3 is the first flexible version.
+  "validVersions": "0-3",
+  "flexibleVersions": "3+",
   "fields": [
     { "name": "ErrorCode", "type": "int16", "versions": "0+",
       "about": "The top-level error code, or 0 if there was no top-level error." },
-    { "name": "Partitions", "type": "[]StopReplicaResponsePartition", "versions": "0+",
+    { "name": "PartitionErrors", "type": "[]StopReplicaPartitionError", "versions": "0+",
       "about": "The responses for each partition.", "fields": [
-      { "name": "TopicName", "type": "string", "versions": "0+",
+      { "name": "TopicName", "type": "string", "versions": "0+", "entityType": "topicName",
         "about": "The topic name." },
       { "name": "PartitionIndex", "type": "int32", "versions": "0+",
         "about": "The partition index." },
diff --git a/clients/src/main/resources/common/message/SyncGroupRequest.json b/clients/src/main/resources/common/message/SyncGroupRequest.json
index ec910a03e60a7..0e65d87f8db4a 100644
--- a/clients/src/main/resources/common/message/SyncGroupRequest.json
+++ b/clients/src/main/resources/common/message/SyncGroupRequest.json
@@ -18,14 +18,22 @@
   "type": "request",
   "name": "SyncGroupRequest",
   // Versions 1 and 2 are the same as version 0.
-  "validVersions": "0-2",
+  //
+  // Starting from version 3, we add a new field called groupInstanceId to indicate member identity across restarts.
+  //
+  // Version 4 is the first flexible version.
+  "validVersions": "0-4",
+  "flexibleVersions": "4+",
   "fields": [
-    { "name": "GroupId", "type": "string", "versions": "0+",
+    { "name": "GroupId", "type": "string", "versions": "0+", "entityType": "groupId",
       "about": "The unique group identifier." },
     { "name": "GenerationId", "type": "int32", "versions": "0+",
       "about": "The generation of the group." },
     { "name": "MemberId", "type": "string", "versions": "0+",
       "about": "The member ID assigned by the group." },
+    { "name": "GroupInstanceId", "type": "string", "versions": "3+", 
+      "nullableVersions": "3+", "default": "null",
+      "about": "The unique identifier of the consumer instance provided by end user." },
     { "name": "Assignments", "type": "[]SyncGroupRequestAssignment", "versions": "0+",
       "about": "Each assignment.", "fields": [
       { "name": "MemberId", "type": "string", "versions": "0+",
diff --git a/clients/src/main/resources/common/message/SyncGroupResponse.json b/clients/src/main/resources/common/message/SyncGroupResponse.json
index 0faa158223bd3..ce60e3f562fa9 100644
--- a/clients/src/main/resources/common/message/SyncGroupResponse.json
+++ b/clients/src/main/resources/common/message/SyncGroupResponse.json
@@ -18,8 +18,14 @@
   "type": "response",
   "name": "SyncGroupResponse",
   // Version 1 adds throttle time.
+  //
   // Starting in version 2, on quota violation, brokers send out responses before throttling.
-  "validVersions": "0-2",
+  //
+  // Starting from version 3, syncGroupRequest supports a new field called groupInstanceId to indicate member identity across restarts.
+  //
+  // Version 4 is the first flexible version.
+  "validVersions": "0-4",
+  "flexibleVersions": "4+",
   "fields": [
     { "name": "ThrottleTimeMs", "type": "int32", "versions": "1+", "ignorable": true,
       "about": "The duration in milliseconds for which the request was throttled due to a quota violation, or zero if the request did not violate any quota." },
diff --git a/clients/src/main/resources/common/message/TxnOffsetCommitRequest.json b/clients/src/main/resources/common/message/TxnOffsetCommitRequest.json
index 357d1d0ad5c52..6ad3ceb113d84 100644
--- a/clients/src/main/resources/common/message/TxnOffsetCommitRequest.json
+++ b/clients/src/main/resources/common/message/TxnOffsetCommitRequest.json
@@ -21,18 +21,19 @@
   //
   // Version 2 adds the committed leader epoch.
   "validVersions": "0-2",
+  "flexibleVersions": "none",
   "fields": [
     { "name": "TransactionalId", "type": "string", "versions": "0+",
       "about": "The ID of the transaction." },
-    { "name": "GroupId", "type": "string", "versions": "0+",
+    { "name": "GroupId", "type": "string", "versions": "0+", "entityType": "groupId",
       "about": "The ID of the group." },
-    { "name": "ProducerId", "type": "int64", "versions": "0+",
+    { "name": "ProducerId", "type": "int64", "versions": "0+", "entityType": "producerId",
       "about": "The current producer ID in use by the transactional ID." },
     { "name": "ProducerEpoch", "type": "int16", "versions": "0+",
       "about": "The current epoch associated with the producer ID." },
     { "name": "Topics", "type" : "[]TxnOffsetCommitRequestTopic", "versions": "0+",
       "about": "Each topic that we want to committ offsets for.", "fields": [
-      { "name": "Name", "type": "string", "versions": "0+",
+      { "name": "Name", "type": "string", "versions": "0+", "entityType": "topicName",
         "about": "The topic name." },
       { "name": "Partitions", "type": "[]TxnOffsetCommitRequestPartition", "versions": "0+",
         "about": "The partitions inside the topic that we want to committ offsets for.", "fields": [
diff --git a/clients/src/main/resources/common/message/TxnOffsetCommitResponse.json b/clients/src/main/resources/common/message/TxnOffsetCommitResponse.json
index 58667cb503661..7cfe0c3e75301 100644
--- a/clients/src/main/resources/common/message/TxnOffsetCommitResponse.json
+++ b/clients/src/main/resources/common/message/TxnOffsetCommitResponse.json
@@ -20,12 +20,13 @@
   // Starting in version 1, on quota violation, brokers send out responses before throttling.
   // Version 2 is the same as version 1.
   "validVersions": "0-2",
+  "flexibleVersions": "none",
   "fields": [
     { "name": "ThrottleTimeMs", "type": "int32", "versions": "0+",
       "about": "The duration in milliseconds for which the request was throttled due to a quota violation, or zero if the request did not violate any quota." },
     { "name": "Topics", "type": "[]TxnOffsetCommitResponseTopic", "versions": "0+",
       "about": "The responses for each topic.", "fields": [
-      { "name": "Name", "type": "string", "versions": "0+",
+      { "name": "Name", "type": "string", "versions": "0+", "entityType": "topicName",
         "about": "The topic name." },
       { "name": "Partitions", "type": "[]TxnOffsetCommitResponsePartition", "versions": "0+",
         "about": "The responses for each partition in the topic.", "fields": [
diff --git a/clients/src/main/resources/common/message/UpdateMetadataRequest.json b/clients/src/main/resources/common/message/UpdateMetadataRequest.json
index 07e0f0346eec8..19ed06184b98d 100644
--- a/clients/src/main/resources/common/message/UpdateMetadataRequest.json
+++ b/clients/src/main/resources/common/message/UpdateMetadataRequest.json
@@ -26,74 +26,46 @@
   // Version 4 adds the offline replica list.
   //
   // Version 5 adds the broker epoch field and normalizes partitions by topic.
-  "validVersions": "0-5",
+  //
+  // Version 6 adds the max broker epoch to make the UpdateMetadataRequest cacheable
+  //
+  // Version 7 is the first flexible version.
+  "validVersions": "0-7",
+  "flexibleVersions": "7+",
   "fields": [
-    { "name": "ControllerId", "type": "int32", "versions": "0+",
+    { "name": "ControllerId", "type": "int32", "versions": "0+", "entityType": "brokerId",
       "about": "The controller id." },
     { "name": "ControllerEpoch", "type": "int32", "versions": "0+",
       "about": "The controller epoch." },
-    { "name": "BrokerEpoch", "type": "int64", "versions": "5+", "ignorable": true, "default": "-1",
+    { "name": "BrokerEpoch", "type": "int64", "versions": "5", "ignorable": true, "default": "-1",
       "about": "The broker epoch." },
-    { "name": "TopicStates", "type": "[]UpdateMetadataRequestTopicState", "versions": "5+",
-      "about": "Each topic that we would like to update.", "fields": [
-      { "name": "TopicName", "type": "string", "versions": "0+",
+    { "name": "UngroupedPartitionStates", "type": "[]UpdateMetadataPartitionState", "versions": "0-4",
+      "about": "In older versions of this RPC, each partition that we would like to update." },
+    { "name": "MaxBrokerEpoch", "type": "int64", "versions": "6+", "ignorable": true, "default": "-1",
+      "about": "The max broker epoch." },
+    { "name": "TopicStates", "type": "[]UpdateMetadataTopicState", "versions": "5+",
+      "about": "In newer versions of this RPC, each topic that we would like to update.", "fields": [
+      { "name": "TopicName", "type": "string", "versions": "5+", "entityType": "topicName",
         "about": "The topic name." },
       { "name": "PartitionStates", "type": "[]UpdateMetadataPartitionState", "versions": "5+",
-        "about": "The partition that we would like to update.", "fields": [
-        { "name": "PartitionIndex", "type": "int32", "versions": "5+",
-          "about": "The partition index." },
-        { "name": "ControllerEpoch", "type": "int32", "versions": "5+",
-          "about": "The controller epoch." },
-        { "name": "Leader", "type": "int32", "versions": "5+",
-          "about": "The ID of the broker which is the current partition leader." },
-        { "name": "LeaderEpoch", "type": "int32", "versions": "5+",
-          "about": "The leader epoch of this partition." },
-        { "name": "Isr", "type": "[]int32", "versions": "5+",
-          "about": "The brokers which are in the ISR for this partition." },
-        { "name": "ZkVersion", "type": "int32", "versions": "5+",
-          "about": "The Zookeeper version." },
-        { "name": "Replicas", "type": "[]int32", "versions": "5+",
-          "about": "All the replicas of this partition." },
-        { "name": "OfflineReplicas", "type": "[]int32", "versions": "5+",
-          "about": "The replicas of this partition which are offline." }
-      ]}
+        "about": "The partition that we would like to update." }
     ]},
-    { "name": "PartitionStatesV0", "type": "[]UpdateMetadataRequestPartitionStateV0", "versions": "0-4",
-      "about": "Each partition that we would like to update.", "fields": [
-      { "name": "TopicName", "type": "string", "versions": "0-4",
-        "about": "The topic name." },
-      { "name": "PartitionIndex", "type": "int32", "versions": "0-4",
-        "about": "The partition index." },
-      { "name": "ControllerEpoch", "type": "int32", "versions": "0-4",
-        "about": "The controller epoch." },
-      { "name": "Leader", "type": "int32", "versions": "0-4",
-        "about": "The ID of the broker which is the current partition leader." },
-      { "name": "LeaderEpoch", "type": "int32", "versions": "0-4",
-        "about": "The leader epoch of this partition." },
-      { "name": "Isr", "type": "[]int32", "versions": "0-4",
-        "about": "The brokers which are in the ISR for this partition." },
-      { "name": "ZkVersion", "type": "int32", "versions": "0-4",
-        "about": "The Zookeeper version." },
-      { "name": "Replicas", "type": "[]int32", "versions": "0-4",
-        "about": "All the replicas of this partition." },
-      { "name": "OfflineReplicas", "type": "[]int32", "versions": "4",
-        "about": "The replicas of this partition which are offline." }
-    ]},
-    { "name": "Brokers", "type": "[]UpdateMetadataRequestBroker", "versions": "0+", "fields": [
-        { "name": "Id", "type": "int32", "versions": "0+" },
+    { "name": "LiveBrokers", "type": "[]UpdateMetadataBroker", "versions": "0+", "fields": [
+        { "name": "Id", "type": "int32", "versions": "0+", "entityType": "brokerId",
+          "about": "The broker id." },
         // Version 0 of the protocol only allowed specifying a single host and
         // port per broker, rather than an array of endpoints.
         { "name": "V0Host", "type": "string", "versions": "0", "ignorable": true,
           "about": "The broker hostname." },
         { "name": "V0Port", "type": "int32", "versions": "0", "ignorable": true,
           "about": "The broker port." },
-        { "name": "Endpoints", "type": "[]UpdateMetadataRequestEndpoint", "versions": "1+",
+        { "name": "Endpoints", "type": "[]UpdateMetadataEndpoint", "versions": "1+", "ignorable": true,
           "about": "The broker endpoints.", "fields": [
           { "name": "Port", "type": "int32", "versions": "1+",
             "about": "The port of this endpoint" },
           { "name": "Host", "type": "string", "versions": "1+",
             "about": "The hostname of this endpoint" },
-          { "name": "Listener", "type": "string", "versions": "3+",
+          { "name": "Listener", "type": "string", "versions": "3+", "ignorable": true,
             "about": "The listener name." },
           { "name": "SecurityProtocol", "type": "int16", "versions": "1+",
             "about": "The security protocol type." }
@@ -101,5 +73,27 @@
         { "name": "Rack", "type": "string", "versions": "2+", "nullableVersions": "0+", "ignorable": true,
           "about": "The rack which this broker belongs to." }
     ]}
+  ],
+  "commonStructs": [
+    { "name": "UpdateMetadataPartitionState", "versions": "0+", "fields": [
+      { "name": "TopicName", "type": "string", "versions": "0-4", "entityType": "topicName", "ignorable": true,
+        "about": "In older versions of this RPC, the topic name." },
+      { "name": "PartitionIndex", "type": "int32", "versions": "0+",
+        "about": "The partition index." },
+      { "name": "ControllerEpoch", "type": "int32", "versions": "0+",
+        "about": "The controller epoch." },
+      { "name": "Leader", "type": "int32", "versions": "0+", "entityType": "brokerId",
+        "about": "The ID of the broker which is the current partition leader." },
+      { "name": "LeaderEpoch", "type": "int32", "versions": "0+",
+        "about": "The leader epoch of this partition." },
+      { "name": "Isr", "type": "[]int32", "versions": "0+", "entityType": "brokerId",
+        "about": "The brokers which are in the ISR for this partition." },
+      { "name": "ZkVersion", "type": "int32", "versions": "0+",
+        "about": "The Zookeeper version." },
+      { "name": "Replicas", "type": "[]int32", "versions": "0+", "entityType": "brokerId",
+        "about": "All the replicas of this partition." },
+      { "name": "OfflineReplicas", "type": "[]int32", "versions": "4+", "entityType": "brokerId", "ignorable": true,
+        "about": "The replicas of this partition which are offline." }
+    ]}
   ]
 }
diff --git a/clients/src/main/resources/common/message/UpdateMetadataResponse.json b/clients/src/main/resources/common/message/UpdateMetadataResponse.json
index 5069a63d77345..68d093f687951 100644
--- a/clients/src/main/resources/common/message/UpdateMetadataResponse.json
+++ b/clients/src/main/resources/common/message/UpdateMetadataResponse.json
@@ -17,8 +17,11 @@
   "apiKey": 6,
   "type": "response",
   "name": "UpdateMetadataResponse",
-  // Versions 1, 2, 3, 4, and 5 are the same as version 0
-  "validVersions": "0-5",
+  // Versions 1, 2, 3, 4, 5 and 6 are the same as version 0
+  //
+  // Version 7 is the first flexible version.
+  "validVersions": "0-7",
+  "flexibleVersions": "7+",
   "fields": [
       { "name": "ErrorCode", "type": "int16", "versions": "0+",
         "about": "The error code, or 0 if there was no error." }
diff --git a/clients/src/main/resources/common/message/WriteTxnMarkersRequest.json b/clients/src/main/resources/common/message/WriteTxnMarkersRequest.json
index 89868fc9710b4..845a4ca3735f1 100644
--- a/clients/src/main/resources/common/message/WriteTxnMarkersRequest.json
+++ b/clients/src/main/resources/common/message/WriteTxnMarkersRequest.json
@@ -18,10 +18,11 @@
   "type": "request",
   "name": "WriteTxnMarkersRequest",
   "validVersions": "0",
+  "flexibleVersions": "none",
   "fields": [
     { "name": "Markers", "type": "[]WritableTxnMarker", "versions": "0+",
       "about": "The transaction markers to be written.", "fields": [
-      { "name": "ProducerId", "type": "int64", "versions": "0+",
+      { "name": "ProducerId", "type": "int64", "versions": "0+", "entityType": "producerId",
         "about": "The current producer ID."},
       { "name": "ProducerEpoch", "type": "int16", "versions": "0+",
         "about": "The current epoch associated with the producer ID." },
@@ -29,7 +30,7 @@
         "about": "The result of the transaction to write to the partitions (false = ABORT, true = COMMIT)." },
       { "name": "Topics", "type": "[]WritableTxnMarkerTopic", "versions": "0+",
         "about": "Each topic that we want to write transaction marker(s) for.", "fields": [
-        { "name": "Name", "type": "string", "versions": "0+",
+        { "name": "Name", "type": "string", "versions": "0+", "entityType": "topicName",
           "about": "The topic name." },
         { "name": "PartitionIndexes", "type": "[]int32", "versions": "0+",
           "about": "The indexes of the partitions to write transaction markers for." }
diff --git a/clients/src/main/resources/common/message/WriteTxnMarkersResponse.json b/clients/src/main/resources/common/message/WriteTxnMarkersResponse.json
index ca080548f7594..67bb791541b3b 100644
--- a/clients/src/main/resources/common/message/WriteTxnMarkersResponse.json
+++ b/clients/src/main/resources/common/message/WriteTxnMarkersResponse.json
@@ -18,14 +18,15 @@
   "type": "response",
   "name": "WriteTxnMarkersResponse",
   "validVersions": "0",
+  "flexibleVersions": "none",
   "fields": [
     { "name": "Markers", "type": "[]WritableTxnMarkerResult", "versions": "0+",
       "about": "The results for writing makers.", "fields": [
-      { "name": "ProducerId", "type": "int64", "versions": "0+",
+      { "name": "ProducerId", "type": "int64", "versions": "0+", "entityType": "producerId",
         "about": "The current producer ID in use by the transactional ID." },
       { "name": "Topics", "type": "[]WritableTxnMarkerTopicResult", "versions": "0+",
         "about": "The results by topic.", "fields": [
-        { "name": "Name", "type": "string", "versions": "0+",
+        { "name": "Name", "type": "string", "versions": "0+", "entityType": "topicName",
           "about": "The topic name." },
         { "name": "Partitions", "type": "[]WritableTxnMarkerPartitionResult", "versions": "0+",
           "about": "The results by partition.", "fields": [
diff --git a/clients/src/test/java/org/apache/kafka/clients/ApiVersionsTest.java b/clients/src/test/java/org/apache/kafka/clients/ApiVersionsTest.java
index 654a6064d5401..9458b04ec89e0 100644
--- a/clients/src/test/java/org/apache/kafka/clients/ApiVersionsTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/ApiVersionsTest.java
@@ -18,11 +18,8 @@
 
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.record.RecordBatch;
-import org.apache.kafka.common.requests.ApiVersionsResponse;
 import org.junit.Test;
 
-import java.util.Collections;
-
 import static org.junit.Assert.assertEquals;
 
 public class ApiVersionsTest {
@@ -35,8 +32,7 @@ public void testMaxUsableProduceMagic() {
         apiVersions.update("0", NodeApiVersions.create());
         assertEquals(RecordBatch.CURRENT_MAGIC_VALUE, apiVersions.maxUsableProduceMagic());
 
-        apiVersions.update("1", NodeApiVersions.create(Collections.singleton(
-                new ApiVersionsResponse.ApiVersion(ApiKeys.PRODUCE.id, (short) 0, (short) 2))));
+        apiVersions.update("1", NodeApiVersions.create(ApiKeys.PRODUCE.id, (short) 0, (short) 2));
         assertEquals(RecordBatch.MAGIC_VALUE_V1, apiVersions.maxUsableProduceMagic());
 
         apiVersions.remove("1");
diff --git a/clients/src/test/java/org/apache/kafka/clients/ClientUtilsTest.java b/clients/src/test/java/org/apache/kafka/clients/ClientUtilsTest.java
index afe5a5d1e7d37..572896f56bd11 100644
--- a/clients/src/test/java/org/apache/kafka/clients/ClientUtilsTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/ClientUtilsTest.java
@@ -16,18 +16,17 @@
  */
 package org.apache.kafka.clients;
 
-import org.apache.kafka.common.config.ConfigException;
-import org.junit.Test;
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertTrue;
-
 import java.net.InetAddress;
 import java.net.InetSocketAddress;
 import java.net.UnknownHostException;
 import java.util.Arrays;
 import java.util.List;
 import java.util.stream.Collectors;
+import org.apache.kafka.common.config.ConfigException;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import org.junit.Test;
 
 public class ClientUtilsTest {
 
@@ -108,7 +107,7 @@ public void testResolveDnsLookup() throws UnknownHostException {
 
     @Test
     public void testResolveDnsLookupAllIps() throws UnknownHostException {
-        assertEquals(2, ClientUtils.resolve("kafka.apache.org", ClientDnsLookup.USE_ALL_DNS_IPS).size());
+        assertTrue(ClientUtils.resolve("kafka.apache.org", ClientDnsLookup.USE_ALL_DNS_IPS).size() > 1);
     }
 
     private List checkWithoutLookup(String... url) {
diff --git a/clients/src/test/java/org/apache/kafka/clients/ClusterConnectionStatesTest.java b/clients/src/test/java/org/apache/kafka/clients/ClusterConnectionStatesTest.java
index 19b701df6f3ad..eb9d22fc1643f 100644
--- a/clients/src/test/java/org/apache/kafka/clients/ClusterConnectionStatesTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/ClusterConnectionStatesTest.java
@@ -256,7 +256,7 @@ public void testSingleIPWithUseAll() throws UnknownHostException {
 
     @Test
     public void testMultipleIPsWithDefault() throws UnknownHostException {
-        assertEquals(2, ClientUtils.resolve(hostTwoIps, ClientDnsLookup.USE_ALL_DNS_IPS).size());
+        assertTrue(ClientUtils.resolve(hostTwoIps, ClientDnsLookup.USE_ALL_DNS_IPS).size() > 1);
 
         connectionStates.connecting(nodeId1, time.milliseconds(), hostTwoIps, ClientDnsLookup.DEFAULT);
         InetAddress currAddress = connectionStates.currentAddress(nodeId1);
@@ -266,7 +266,7 @@ public void testMultipleIPsWithDefault() throws UnknownHostException {
 
     @Test
     public void testMultipleIPsWithUseAll() throws UnknownHostException {
-        assertEquals(2, ClientUtils.resolve(hostTwoIps, ClientDnsLookup.USE_ALL_DNS_IPS).size());
+        assertTrue(ClientUtils.resolve(hostTwoIps, ClientDnsLookup.USE_ALL_DNS_IPS).size() > 1);
 
         connectionStates.connecting(nodeId1, time.milliseconds(), hostTwoIps, ClientDnsLookup.USE_ALL_DNS_IPS);
         InetAddress addr1 = connectionStates.currentAddress(nodeId1);
@@ -276,12 +276,12 @@ public void testMultipleIPsWithUseAll() throws UnknownHostException {
 
         connectionStates.connecting(nodeId1, time.milliseconds(), hostTwoIps, ClientDnsLookup.USE_ALL_DNS_IPS);
         InetAddress addr3 = connectionStates.currentAddress(nodeId1);
-        assertSame(addr1, addr3);
+        assertNotSame(addr1, addr3);
     }
 
     @Test
     public void testHostResolveChange() throws UnknownHostException, ReflectiveOperationException {
-        assertEquals(2, ClientUtils.resolve(hostTwoIps, ClientDnsLookup.USE_ALL_DNS_IPS).size());
+        assertTrue(ClientUtils.resolve(hostTwoIps, ClientDnsLookup.USE_ALL_DNS_IPS).size() > 1);
 
         connectionStates.connecting(nodeId1, time.milliseconds(), hostTwoIps, ClientDnsLookup.DEFAULT);
         InetAddress addr1 = connectionStates.currentAddress(nodeId1);
@@ -310,4 +310,15 @@ public void testNodeWithNewHostname() throws UnknownHostException {
 
         assertNotSame(addr1, addr2);
     }
+
+    @Test
+    public void testIsPreparingConnection() {
+        assertFalse(connectionStates.isPreparingConnection(nodeId1));
+        connectionStates.connecting(nodeId1, time.milliseconds(), "localhost", ClientDnsLookup.DEFAULT);
+        assertTrue(connectionStates.isPreparingConnection(nodeId1));
+        connectionStates.checkingApiVersions(nodeId1);
+        assertTrue(connectionStates.isPreparingConnection(nodeId1));
+        connectionStates.disconnected(nodeId1, time.milliseconds());
+        assertFalse(connectionStates.isPreparingConnection(nodeId1));
+    }
 }
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 63a9312846001..d65e6d1cbdc50 100644
--- a/clients/src/test/java/org/apache/kafka/clients/CommonClientConfigsTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/CommonClientConfigsTest.java
@@ -58,7 +58,7 @@ public TestConfig(Map props) {
     }
 
     @Test
-    public void testExponentialBackoffDefaults() throws Exception {
+    public void testExponentialBackoffDefaults() {
         TestConfig defaultConf = new TestConfig(Collections.emptyMap());
         assertEquals(Long.valueOf(50L),
                 defaultConf.getLong(CommonClientConfigs.RECONNECT_BACKOFF_MS_CONFIG));
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 3d282971ed1b2..f690dc685b27b 100644
--- a/clients/src/test/java/org/apache/kafka/clients/MetadataTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/MetadataTest.java
@@ -17,12 +17,21 @@
 package org.apache.kafka.clients;
 
 import org.apache.kafka.common.Cluster;
+import org.apache.kafka.common.Node;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.errors.InvalidTopicException;
 import org.apache.kafka.common.errors.TopicAuthorizationException;
 import org.apache.kafka.common.internals.ClusterResourceListeners;
 import org.apache.kafka.common.internals.Topic;
+import org.apache.kafka.common.message.MetadataResponseData;
+import org.apache.kafka.common.message.MetadataResponseData.MetadataResponseBroker;
+import org.apache.kafka.common.message.MetadataResponseData.MetadataResponseBrokerCollection;
+import org.apache.kafka.common.message.MetadataResponseData.MetadataResponsePartition;
+import org.apache.kafka.common.message.MetadataResponseData.MetadataResponseTopic;
+import org.apache.kafka.common.message.MetadataResponseData.MetadataResponseTopicCollection;
+import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.common.protocol.types.Struct;
 import org.apache.kafka.common.requests.MetadataResponse;
 import org.apache.kafka.common.utils.LogContext;
 import org.apache.kafka.common.utils.MockTime;
@@ -32,11 +41,15 @@
 import org.junit.Test;
 
 import java.net.InetSocketAddress;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.stream.Collectors;
 
+import static org.apache.kafka.test.TestUtils.assertOptional;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
@@ -52,7 +65,7 @@ public class MetadataTest {
             new ClusterResourceListeners());
 
     private static MetadataResponse emptyMetadataResponse() {
-        return new MetadataResponse(
+        return MetadataResponse.prepareResponse(
                 Collections.emptyList(),
                 null,
                 -1,
@@ -104,6 +117,18 @@ private static void checkTimeToNextUpdate(long refreshBackoffMs, long metadataEx
         assertEquals(0, metadata.timeToNextUpdate(now + 1));
     }
 
+    @Test
+    public void testUpdateMetadataAllowedImmediatelyAfterBootstrap() {
+        MockTime time = new MockTime();
+
+        Metadata metadata = new Metadata(refreshBackoffMs, metadataExpireMs, new LogContext(),
+                new ClusterResourceListeners());
+        metadata.bootstrap(Collections.singletonList(new InetSocketAddress("localhost", 9002)));
+
+        assertEquals(0, metadata.timeToAllowUpdate(time.milliseconds()));
+        assertEquals(0, metadata.timeToNextUpdate(time.milliseconds()));
+    }
+
     @Test
     public void testTimeToNextUpdate() {
         checkTimeToNextUpdate(100, 1000);
@@ -118,7 +143,7 @@ public void testTimeToNextUpdate_RetryBackoff() {
         long now = 10000;
 
         // lastRefreshMs updated to now.
-        metadata.failedUpdate(now, null);
+        metadata.failedUpdate(now);
 
         // Backing off. Remaining time until next try should be returned.
         assertEquals(refreshBackoffMs, metadata.timeToNextUpdate(now));
@@ -134,13 +159,119 @@ public void testTimeToNextUpdate_RetryBackoff() {
         assertEquals(0, metadata.timeToNextUpdate(now + 1));
     }
 
+    /**
+     * Prior to Kafka version 2.4 (which coincides with Metadata version 9), the broker does not propagate leader epoch
+     * information accurately while a reassignment is in progress, so we cannot rely on it. This is explained in more
+     * detail in MetadataResponse's constructor.
+     */
+    @Test
+    public void testIgnoreLeaderEpochInOlderMetadataResponse() {
+        TopicPartition tp = new TopicPartition("topic", 0);
+
+        MetadataResponsePartition partitionMetadata = new MetadataResponsePartition()
+                .setPartitionIndex(tp.partition())
+                .setLeaderId(5)
+                .setLeaderEpoch(10)
+                .setReplicaNodes(Arrays.asList(1, 2, 3))
+                .setIsrNodes(Arrays.asList(1, 2, 3))
+                .setOfflineReplicas(Collections.emptyList())
+                .setErrorCode(Errors.NONE.code());
+
+        MetadataResponseTopic topicMetadata = new MetadataResponseTopic()
+                .setName(tp.topic())
+                .setErrorCode(Errors.NONE.code())
+                .setPartitions(Collections.singletonList(partitionMetadata))
+                .setIsInternal(false);
+
+        MetadataResponseTopicCollection topics = new MetadataResponseTopicCollection();
+        topics.add(topicMetadata);
+
+        MetadataResponseData data = new MetadataResponseData()
+                .setClusterId("clusterId")
+                .setControllerId(0)
+                .setTopics(topics)
+                .setBrokers(new MetadataResponseBrokerCollection());
+
+        for (short version = ApiKeys.METADATA.oldestVersion(); version < 9; version++) {
+            Struct struct = data.toStruct(version);
+            MetadataResponse response = new MetadataResponse(struct, version);
+            assertFalse(response.hasReliableLeaderEpochs());
+            metadata.update(response, 100);
+            assertTrue(metadata.partitionInfoIfCurrent(tp).isPresent());
+            MetadataCache.PartitionInfoAndEpoch info = metadata.partitionInfoIfCurrent(tp).get();
+            assertEquals(-1, info.epoch());
+        }
+
+        for (short version = 9; version <= ApiKeys.METADATA.latestVersion(); version++) {
+            Struct struct = data.toStruct(version);
+            MetadataResponse response = new MetadataResponse(struct, version);
+            assertTrue(response.hasReliableLeaderEpochs());
+            metadata.update(response, 100);
+            assertTrue(metadata.partitionInfoIfCurrent(tp).isPresent());
+            MetadataCache.PartitionInfoAndEpoch info = metadata.partitionInfoIfCurrent(tp).get();
+            assertEquals(10, info.epoch());
+        }
+    }
+
+    @Test
+    public void testStaleMetadata() {
+        TopicPartition tp = new TopicPartition("topic", 0);
+
+        MetadataResponsePartition partitionMetadata = new MetadataResponsePartition()
+                .setPartitionIndex(tp.partition())
+                .setLeaderId(1)
+                .setLeaderEpoch(10)
+                .setReplicaNodes(Arrays.asList(1, 2, 3))
+                .setIsrNodes(Arrays.asList(1, 2, 3))
+                .setOfflineReplicas(Collections.emptyList())
+                .setErrorCode(Errors.NONE.code());
+
+        MetadataResponseTopic topicMetadata = new MetadataResponseTopic()
+                .setName(tp.topic())
+                .setErrorCode(Errors.NONE.code())
+                .setPartitions(Collections.singletonList(partitionMetadata))
+                .setIsInternal(false);
+
+        MetadataResponseTopicCollection topics = new MetadataResponseTopicCollection();
+        topics.add(topicMetadata);
+
+        MetadataResponseData data = new MetadataResponseData()
+                .setClusterId("clusterId")
+                .setControllerId(0)
+                .setTopics(topics)
+                .setBrokers(new MetadataResponseBrokerCollection());
+
+        metadata.update(new MetadataResponse(data), 100);
+
+        // Older epoch with changed ISR should be ignored
+        partitionMetadata
+                .setPartitionIndex(tp.partition())
+                .setLeaderId(1)
+                .setLeaderEpoch(9)
+                .setReplicaNodes(Arrays.asList(1, 2, 3))
+                .setIsrNodes(Arrays.asList(1, 2))
+                .setOfflineReplicas(Collections.emptyList())
+                .setErrorCode(Errors.NONE.code());
+
+        metadata.update(new MetadataResponse(data), 101);
+        assertEquals(Optional.of(10), metadata.lastSeenLeaderEpoch(tp));
+
+        assertTrue(metadata.partitionInfoIfCurrent(tp).isPresent());
+        MetadataCache.PartitionInfoAndEpoch info = metadata.partitionInfoIfCurrent(tp).get();
+
+        List cachedIsr = Arrays.stream(info.partitionInfo().inSyncReplicas())
+                .map(Node::id).collect(Collectors.toList());
+        assertEquals(Arrays.asList(1, 2, 3), cachedIsr);
+        assertEquals(10, info.epoch());
+    }
+
     @Test
     public void testFailedUpdate() {
         long time = 100;
         metadata.update(emptyMetadataResponse(), time);
 
         assertEquals(100, metadata.timeToNextUpdate(1000));
-        metadata.failedUpdate(1100, null);
+        metadata.failedUpdate(1100);
 
         assertEquals(100, metadata.timeToNextUpdate(1100));
         assertEquals(100, metadata.lastSuccessfulUpdate());
@@ -151,14 +282,13 @@ public void testFailedUpdate() {
 
     @Test
     public void testClusterListenerGetsNotifiedOfUpdate() {
-        long time = 0;
         MockClusterResourceListener mockClusterListener = new MockClusterResourceListener();
         ClusterResourceListeners listeners = new ClusterResourceListeners();
         listeners.maybeAdd(mockClusterListener);
         metadata = new Metadata(refreshBackoffMs, metadataExpireMs, new LogContext(), listeners);
 
         String hostName = "www.example.com";
-        metadata.bootstrap(Collections.singletonList(new InetSocketAddress(hostName, 9002)), time);
+        metadata.bootstrap(Collections.singletonList(new InetSocketAddress(hostName, 9002)));
         assertFalse("ClusterResourceListener should not called when metadata is updated with bootstrap Cluster",
                 MockClusterResourceListener.IS_ON_UPDATE_CALLED.get());
 
@@ -205,9 +335,7 @@ public void testRejectOldMetadata() {
 
         // First epoch seen, accept it
         {
-            MetadataResponse metadataResponse = TestUtils.metadataUpdateWith("dummy", 1, Collections.emptyMap(), partitionCounts,
-                (error, partition, leader, leaderEpoch, replicas, isr, offlineReplicas) ->
-                        new MetadataResponse.PartitionMetadata(error, partition, leader, Optional.of(100), replicas, isr, offlineReplicas));
+            MetadataResponse metadataResponse = TestUtils.metadataUpdateWith("dummy", 1, Collections.emptyMap(), partitionCounts, _tp -> 100);
             metadata.update(metadataResponse, 10L);
             assertNotNull(metadata.fetch().partition(tp));
             assertEquals(metadata.lastSeenLeaderEpoch(tp).get().longValue(), 100);
@@ -215,9 +343,9 @@ public void testRejectOldMetadata() {
 
         // Fake an empty ISR, but with an older epoch, should reject it
         {
-            MetadataResponse metadataResponse = TestUtils.metadataUpdateWith("dummy", 1, Collections.emptyMap(), partitionCounts,
+            MetadataResponse metadataResponse = TestUtils.metadataUpdateWith("dummy", 1, Collections.emptyMap(), partitionCounts, _tp -> 99,
                 (error, partition, leader, leaderEpoch, replicas, isr, offlineReplicas) ->
-                        new MetadataResponse.PartitionMetadata(error, partition, leader, Optional.of(99), replicas, Collections.emptyList(), offlineReplicas));
+                        new MetadataResponse.PartitionMetadata(error, partition, leader, leaderEpoch, replicas, Collections.emptyList(), offlineReplicas));
             metadata.update(metadataResponse, 20L);
             assertEquals(metadata.fetch().partition(tp).inSyncReplicas().length, 1);
             assertEquals(metadata.lastSeenLeaderEpoch(tp).get().longValue(), 100);
@@ -225,9 +353,9 @@ public void testRejectOldMetadata() {
 
         // Fake an empty ISR, with same epoch, accept it
         {
-            MetadataResponse metadataResponse = TestUtils.metadataUpdateWith("dummy", 1, Collections.emptyMap(), partitionCounts,
+            MetadataResponse metadataResponse = TestUtils.metadataUpdateWith("dummy", 1, Collections.emptyMap(), partitionCounts, _tp -> 100,
                 (error, partition, leader, leaderEpoch, replicas, isr, offlineReplicas) ->
-                        new MetadataResponse.PartitionMetadata(error, partition, leader, Optional.of(100), replicas, Collections.emptyList(), offlineReplicas));
+                        new MetadataResponse.PartitionMetadata(error, partition, leader, leaderEpoch, replicas, Collections.emptyList(), offlineReplicas));
             metadata.update(metadataResponse, 20L);
             assertEquals(metadata.fetch().partition(tp).inSyncReplicas().length, 0);
             assertEquals(metadata.lastSeenLeaderEpoch(tp).get().longValue(), 100);
@@ -235,8 +363,7 @@ public void testRejectOldMetadata() {
 
         // Empty metadata response, should not keep old partition but should keep the last-seen epoch
         {
-            MetadataResponse metadataResponse = TestUtils.metadataUpdateWith(
-                    "dummy", 1, Collections.emptyMap(), Collections.emptyMap(), MetadataResponse.PartitionMetadata::new);
+            MetadataResponse metadataResponse = TestUtils.metadataUpdateWith("dummy", 1, Collections.emptyMap(), Collections.emptyMap());
             metadata.update(metadataResponse, 20L);
             assertNull(metadata.fetch().partition(tp));
             assertEquals(metadata.lastSeenLeaderEpoch(tp).get().longValue(), 100);
@@ -244,9 +371,7 @@ public void testRejectOldMetadata() {
 
         // Back in the metadata, with old epoch, should not get added
         {
-            MetadataResponse metadataResponse = TestUtils.metadataUpdateWith("dummy", 1, Collections.emptyMap(), partitionCounts,
-                (error, partition, leader, leaderEpoch, replicas, isr, offlineReplicas) ->
-                        new MetadataResponse.PartitionMetadata(error, partition, leader, Optional.of(99), replicas, isr, offlineReplicas));
+            MetadataResponse metadataResponse = TestUtils.metadataUpdateWith("dummy", 1, Collections.emptyMap(), partitionCounts, _tp -> 99);
             metadata.update(metadataResponse, 10L);
             assertNull(metadata.fetch().partition(tp));
             assertEquals(metadata.lastSeenLeaderEpoch(tp).get().longValue(), 100);
@@ -284,9 +409,7 @@ public void testOutOfBandEpochUpdate() {
         assertTrue(metadata.updateLastSeenEpochIfNewer(tp, 99));
 
         // Update epoch to 100
-        MetadataResponse metadataResponse = TestUtils.metadataUpdateWith("dummy", 1, Collections.emptyMap(), partitionCounts,
-            (error, partition, leader, leaderEpoch, replicas, isr, offlineReplicas) ->
-                    new MetadataResponse.PartitionMetadata(error, partition, leader, Optional.of(100), replicas, isr, offlineReplicas));
+        MetadataResponse metadataResponse = TestUtils.metadataUpdateWith("dummy", 1, Collections.emptyMap(), partitionCounts, _tp -> 100);
         metadata.update(metadataResponse, 10L);
         assertNotNull(metadata.fetch().partition(tp));
         assertEquals(metadata.lastSeenLeaderEpoch(tp).get().longValue(), 100);
@@ -308,9 +431,7 @@ public void testOutOfBandEpochUpdate() {
         assertEquals(metadata.lastSeenLeaderEpoch(tp).get().longValue(), 101);
 
         // Metadata with equal or newer epoch is accepted
-        metadataResponse = TestUtils.metadataUpdateWith("dummy", 1, Collections.emptyMap(), partitionCounts,
-            (error, partition, leader, leaderEpoch, replicas, isr, offlineReplicas) ->
-                    new MetadataResponse.PartitionMetadata(error, partition, leader, Optional.of(101), replicas, isr, offlineReplicas));
+        metadataResponse = TestUtils.metadataUpdateWith("dummy", 1, Collections.emptyMap(), partitionCounts, _tp -> 101);
         metadata.update(metadataResponse, 30L);
         assertNotNull(metadata.fetch().partition(tp));
         assertEquals(metadata.fetch().partitionCountForTopic("topic-1").longValue(), 5);
@@ -321,9 +442,7 @@ public void testOutOfBandEpochUpdate() {
     @Test
     public void testNoEpoch() {
         metadata.update(emptyMetadataResponse(), 0L);
-        MetadataResponse metadataResponse = TestUtils.metadataUpdateWith("dummy", 1, Collections.emptyMap(), Collections.singletonMap("topic-1", 1),
-            (error, partition, leader, leaderEpoch, replicas, isr, offlineReplicas) ->
-                new MetadataResponse.PartitionMetadata(error, partition, leader, Optional.empty(), replicas, isr, offlineReplicas));
+        MetadataResponse metadataResponse = TestUtils.metadataUpdateWith("dummy", 1, Collections.emptyMap(), Collections.singletonMap("topic-1", 1));
         metadata.update(metadataResponse, 10L);
 
         TopicPartition tp = new TopicPartition("topic-1", 0);
@@ -413,18 +532,18 @@ public void testInvalidTopicError() {
                 Collections.singletonMap(invalidTopic, Errors.INVALID_TOPIC_EXCEPTION), Collections.emptyMap());
         metadata.update(invalidTopicResponse, time.milliseconds());
 
-        InvalidTopicException e = assertThrows(InvalidTopicException.class, () -> metadata.maybeThrowException());
+        InvalidTopicException e = assertThrows(InvalidTopicException.class, () -> metadata.maybeThrowAnyException());
 
         assertEquals(Collections.singleton(invalidTopic), e.invalidTopics());
         // We clear the exception once it has been raised to the user
-        assertNull(metadata.getAndClearMetadataException());
+        metadata.maybeThrowAnyException();
 
         // Reset the invalid topic error
         metadata.update(invalidTopicResponse, time.milliseconds());
 
         // If we get a good update, the error should clear even if we haven't had a chance to raise it to the user
         metadata.update(emptyMetadataResponse(), time.milliseconds());
-        assertNull(metadata.getAndClearMetadataException());
+        metadata.maybeThrowAnyException();
     }
 
     @Test
@@ -436,17 +555,147 @@ public void testTopicAuthorizationError() {
                 Collections.singletonMap(invalidTopic, Errors.TOPIC_AUTHORIZATION_FAILED), Collections.emptyMap());
         metadata.update(unauthorizedTopicResponse, time.milliseconds());
 
-        TopicAuthorizationException e = assertThrows(TopicAuthorizationException.class, () -> metadata.maybeThrowException());
+        TopicAuthorizationException e = assertThrows(TopicAuthorizationException.class, () -> metadata.maybeThrowAnyException());
         assertEquals(Collections.singleton(invalidTopic), e.unauthorizedTopics());
         // We clear the exception once it has been raised to the user
-        assertNull(metadata.getAndClearMetadataException());
+        metadata.maybeThrowAnyException();
 
         // Reset the unauthorized topic error
         metadata.update(unauthorizedTopicResponse, time.milliseconds());
 
         // If we get a good update, the error should clear even if we haven't had a chance to raise it to the user
         metadata.update(emptyMetadataResponse(), time.milliseconds());
-        assertNull(metadata.getAndClearMetadataException());
+        metadata.maybeThrowAnyException();
+    }
+
+    @Test
+    public void testMetadataTopicErrors() {
+        Time time = new MockTime();
+
+        Map topicErrors = new HashMap<>(3);
+        topicErrors.put("invalidTopic", Errors.INVALID_TOPIC_EXCEPTION);
+        topicErrors.put("sensitiveTopic1", Errors.TOPIC_AUTHORIZATION_FAILED);
+        topicErrors.put("sensitiveTopic2", Errors.TOPIC_AUTHORIZATION_FAILED);
+        MetadataResponse metadataResponse = TestUtils.metadataUpdateWith("clusterId", 1, topicErrors, Collections.emptyMap());
+
+        metadata.update(metadataResponse, time.milliseconds());
+        TopicAuthorizationException e1 = assertThrows(TopicAuthorizationException.class,
+            () -> metadata.maybeThrowExceptionForTopic("sensitiveTopic1"));
+        assertEquals(Collections.singleton("sensitiveTopic1"), e1.unauthorizedTopics());
+        // We clear the exception once it has been raised to the user
+        metadata.maybeThrowAnyException();
+
+        metadata.update(metadataResponse, time.milliseconds());
+        TopicAuthorizationException e2 = assertThrows(TopicAuthorizationException.class,
+            () -> metadata.maybeThrowExceptionForTopic("sensitiveTopic2"));
+        assertEquals(Collections.singleton("sensitiveTopic2"), e2.unauthorizedTopics());
+        metadata.maybeThrowAnyException();
+
+        metadata.update(metadataResponse, time.milliseconds());
+        InvalidTopicException e3 = assertThrows(InvalidTopicException.class,
+            () -> metadata.maybeThrowExceptionForTopic("invalidTopic"));
+        assertEquals(Collections.singleton("invalidTopic"), e3.invalidTopics());
+        metadata.maybeThrowAnyException();
+
+        // Other topics should not throw exception, but they should clear existing exception
+        metadata.update(metadataResponse, time.milliseconds());
+        metadata.maybeThrowExceptionForTopic("anotherTopic");
+        metadata.maybeThrowAnyException();
+    }
+
+    @Test
+    public void testNodeIfOffline() {
+        Map partitionCounts = new HashMap<>();
+        partitionCounts.put("topic-1", 1);
+        Node node0 = new Node(0, "localhost", 9092);
+        Node node1 = new Node(1, "localhost", 9093);
+
+        MetadataResponse metadataResponse = TestUtils.metadataUpdateWith("dummy", 2, Collections.emptyMap(), partitionCounts, _tp -> 99,
+            (error, partition, leader, leaderEpoch, replicas, isr, offlineReplicas) ->
+                new MetadataResponse.PartitionMetadata(error, partition, node0, leaderEpoch,
+                    Collections.singletonList(node0), Collections.emptyList(), Collections.singletonList(node1)));
+        metadata.update(emptyMetadataResponse(), 0L);
+        metadata.update(metadataResponse, 10L);
+
+        TopicPartition tp = new TopicPartition("topic-1", 0);
+
+        assertOptional(metadata.fetch().nodeIfOnline(tp, 0), node -> assertEquals(node.id(), 0));
+        assertFalse(metadata.fetch().nodeIfOnline(tp, 1).isPresent());
+        assertEquals(metadata.fetch().nodeById(0).id(), 0);
+        assertEquals(metadata.fetch().nodeById(1).id(), 1);
+    }
+
+    @Test
+    public void testLeaderMetadataInconsistentWithBrokerMetadata() {
+        // Tests a reordering scenario which can lead to inconsistent leader state.
+        // A partition initially has one broker offline. That broker comes online and
+        // is elected leader. The client sees these two events in the opposite order.
+
+        TopicPartition tp = new TopicPartition("topic", 0);
+
+        Node node0 = new Node(0, "localhost", 9092);
+        Node node1 = new Node(1, "localhost", 9093);
+        Node node2 = new Node(2, "localhost", 9094);
+
+        // The first metadata received by broker (epoch=10)
+        MetadataResponsePartition firstPartitionMetadata = new MetadataResponsePartition()
+                .setPartitionIndex(tp.partition())
+                .setErrorCode(Errors.NONE.code())
+                .setLeaderEpoch(10)
+                .setLeaderId(0)
+                .setReplicaNodes(Arrays.asList(0, 1, 2))
+                .setIsrNodes(Arrays.asList(0, 1, 2))
+                .setOfflineReplicas(Collections.emptyList());
+
+        // The second metadata received has stale metadata (epoch=8)
+        MetadataResponsePartition secondPartitionMetadata = new MetadataResponsePartition()
+                .setPartitionIndex(tp.partition())
+                .setErrorCode(Errors.NONE.code())
+                .setLeaderEpoch(8)
+                .setLeaderId(1)
+                .setReplicaNodes(Arrays.asList(0, 1, 2))
+                .setIsrNodes(Arrays.asList(1, 2))
+                .setOfflineReplicas(Collections.singletonList(0));
+
+        metadata.update(new MetadataResponse(new MetadataResponseData()
+                        .setTopics(buildTopicCollection(tp.topic(), firstPartitionMetadata))
+                        .setBrokers(buildBrokerCollection(Arrays.asList(node0, node1, node2)))),
+                10L);
+
+        metadata.update(new MetadataResponse(new MetadataResponseData()
+                        .setTopics(buildTopicCollection(tp.topic(), secondPartitionMetadata))
+                        .setBrokers(buildBrokerCollection(Arrays.asList(node1, node2)))),
+                20L);
+
+        assertNull(metadata.fetch().leaderFor(tp));
+        assertEquals(Optional.of(10), metadata.lastSeenLeaderEpoch(tp));
+        assertTrue(metadata.leaderAndEpoch(tp).leader.isEmpty());
+    }
+
+    private MetadataResponseTopicCollection buildTopicCollection(String topic, MetadataResponsePartition partitionMetadata) {
+        MetadataResponseTopic topicMetadata = new MetadataResponseTopic()
+                .setErrorCode(Errors.NONE.code())
+                .setName(topic)
+                .setIsInternal(false);
+
+        topicMetadata.setPartitions(Collections.singletonList(partitionMetadata));
+
+        MetadataResponseTopicCollection topics = new MetadataResponseTopicCollection();
+        topics.add(topicMetadata);
+        return topics;
+    }
+
+    private MetadataResponseBrokerCollection buildBrokerCollection(List nodes) {
+        MetadataResponseBrokerCollection brokers = new MetadataResponseBrokerCollection();
+        for (Node node : nodes) {
+            MetadataResponseBroker broker = new MetadataResponseBroker()
+                    .setNodeId(node.id())
+                    .setHost(node.host())
+                    .setPort(node.port())
+                    .setRack(node.rack());
+            brokers.add(broker);
+        }
+        return brokers;
     }
 
 }
diff --git a/clients/src/test/java/org/apache/kafka/clients/MockClient.java b/clients/src/test/java/org/apache/kafka/clients/MockClient.java
index 7a1febd1575dc..f9fb96ce0c0a8 100644
--- a/clients/src/test/java/org/apache/kafka/clients/MockClient.java
+++ b/clients/src/test/java/org/apache/kafka/clients/MockClient.java
@@ -73,17 +73,10 @@ public FutureResponse(Node node,
     }
 
     private int correlation;
+    private Runnable wakeupHook;
     private final Time time;
     private final MockMetadataUpdater metadataUpdater;
-    private final Set ready = new HashSet<>();
-
-    // Nodes awaiting reconnect backoff, will not be chosen by leastLoadedNode
-    private final TransientSet blackedOut;
-    // Nodes which will always fail to connect, but can be chosen by leastLoadedNode
-    private final TransientSet unreachable;
-    // Nodes which have a delay before ultimately succeeding to connect
-    private final TransientSet delayedReady;
-
+    private final Map connections = new HashMap<>();
     private final Map pendingAuthenticationErrors = new HashMap<>();
     private final Map authenticationErrors = new HashMap<>();
     // Use concurrent queue for requests so that requests may be queried from a different thread
@@ -103,36 +96,34 @@ public MockClient(Time time, Metadata metadata) {
     public MockClient(Time time, MockMetadataUpdater metadataUpdater) {
         this.time = time;
         this.metadataUpdater = metadataUpdater;
-        this.blackedOut = new TransientSet<>(time);
-        this.unreachable = new TransientSet<>(time);
-        this.delayedReady = new TransientSet<>(time);
+    }
+
+    public boolean isConnected(String idString) {
+        return connectionState(idString).state == ConnectionState.State.CONNECTED;
+    }
+
+    private ConnectionState connectionState(String idString) {
+        ConnectionState connectionState = connections.get(idString);
+        if (connectionState == null) {
+            connectionState = new ConnectionState();
+            connections.put(idString, connectionState);
+        }
+        return connectionState;
     }
 
     @Override
     public boolean isReady(Node node, long now) {
-        return ready.contains(node.idString());
+        return connectionState(node.idString()).isReady(now);
     }
 
     @Override
     public boolean ready(Node node, long now) {
-        if (blackedOut.contains(node, now))
-            return false;
-
-        if (unreachable.contains(node, now)) {
-            blackout(node, 100);
-            return false;
-        }
-
-        if (delayedReady.contains(node, now))
-            return false;
-
-        ready.add(node.idString());
-        return true;
+        return connectionState(node.idString()).ready(now);
     }
 
     @Override
     public long connectionDelay(Node node, long now) {
-        return blackedOut.expirationDelayMs(node, now);
+        return connectionState(node.idString()).connectionDelay(now);
     }
 
     @Override
@@ -141,16 +132,20 @@ public long pollDelayMs(Node node, long now) {
     }
 
     public void blackout(Node node, long durationMs) {
-        blackedOut.add(node, durationMs);
+        connectionState(node.idString()).backoff(time.milliseconds() + durationMs);
     }
 
     public void setUnreachable(Node node, long durationMs) {
         disconnect(node.idString());
-        unreachable.add(node, durationMs);
+        connectionState(node.idString()).setUnreachable(time.milliseconds() + durationMs);
+    }
+
+    public void throttle(Node node, long durationMs) {
+        connectionState(node.idString()).throttle(time.milliseconds() + durationMs);
     }
 
     public void delayReady(Node node, long durationMs) {
-        delayedReady.add(node, durationMs);
+        connectionState(node.idString()).setReadyDelayed(time.milliseconds() + durationMs);
     }
 
     public void authenticationFailed(Node node, long blackoutMs) {
@@ -166,7 +161,7 @@ public void createPendingAuthenticationError(Node node, long blackoutMs) {
 
     @Override
     public boolean connectionFailed(Node node) {
-        return blackedOut.contains(node);
+        return connectionState(node.idString()).isBackingOff(time.milliseconds());
     }
 
     @Override
@@ -187,11 +182,14 @@ public void disconnect(String node) {
                 iter.remove();
             }
         }
-        ready.remove(node);
+        connectionState(node).disconnect();
     }
 
     @Override
     public void send(ClientRequest request, long now) {
+        if (!connectionState(request.destination()).isReady(now))
+            throw new IllegalStateException("Cannot send " + request + " since the destination is not ready");
+
         // Check if the request is directed to a node with a pending authentication error.
         for (Iterator> authErrorIter =
              pendingAuthenticationErrors.entrySet().iterator(); authErrorIter.hasNext(); ) {
@@ -224,7 +222,7 @@ public void send(ClientRequest request, long now) {
                     builder.latestAllowedVersion());
             AbstractRequest abstractRequest = request.requestBuilder().build(version);
             if (!futureResp.requestMatcher.matches(abstractRequest))
-                throw new IllegalStateException("Request matcher did not match next-in-line request " + abstractRequest);
+                throw new IllegalStateException("Request matcher did not match next-in-line request " + abstractRequest + " with prepared response " + futureResp.responseBody);
 
             UnsupportedVersionException unsupportedVersionException = null;
             if (futureResp.isUnsupportedRequest)
@@ -257,6 +255,9 @@ public synchronized void wakeup() {
             numBlockingWakeups--;
             notify();
         }
+        if (wakeupHook != null) {
+            wakeupHook.run();
+        }
     }
 
     private synchronized void maybeAwaitWakeup() {
@@ -437,9 +438,7 @@ public boolean conditionMet() {
     }
 
     public void reset() {
-        ready.clear();
-        blackedOut.clear();
-        unreachable.clear();
+        connections.clear();
         requests.clear();
         responses.clear();
         futureResponses.clear();
@@ -499,7 +498,7 @@ public boolean hasInFlightRequests(String node) {
 
     @Override
     public boolean hasReadyNodes(long now) {
-        return !ready.isEmpty();
+        return connections.values().stream().anyMatch(cxn -> cxn.isReady(now));
     }
 
     @Override
@@ -537,19 +536,23 @@ public void close() {
 
     @Override
     public void close(String node) {
-        ready.remove(node);
+        connections.remove(node);
     }
 
     @Override
     public Node leastLoadedNode(long now) {
         // Consistent with NetworkClient, we do not return nodes awaiting reconnect backoff
         for (Node node : metadataUpdater.fetchNodes()) {
-            if (!blackedOut.contains(node, now))
+            if (!connectionState(node.idString()).isBackingOff(now))
                 return node;
         }
         return null;
     }
 
+    public void setWakeupHook(Runnable wakeupHook) {
+        this.wakeupHook = wakeupHook;
+    }
+
     /**
      * The RequestMatcher provides a way to match a particular request to a response prepared
      * through {@link #prepareResponse(RequestMatcher, AbstractResponse)}. Basically this allows testers
@@ -580,45 +583,6 @@ private Set topics() {
         }
     }
 
-    private static class TransientSet {
-        // The elements in the set mapped to their expiration timestamps
-        private final Map elements = new HashMap<>();
-        private final Time time;
-
-        private TransientSet(Time time) {
-            this.time = time;
-        }
-
-        boolean contains(T element) {
-            return contains(element, time.milliseconds());
-        }
-
-        boolean contains(T element, long now) {
-            return expirationDelayMs(element, now) > 0;
-        }
-
-        void add(T element, long durationMs) {
-            elements.put(element, time.milliseconds() + durationMs);
-        }
-
-        long expirationDelayMs(T element, long now) {
-            Long expirationTimeMs = elements.get(element);
-            if (expirationTimeMs == null) {
-                return 0;
-            } else if (now > expirationTimeMs) {
-                elements.remove(element);
-                return 0;
-            } else {
-                return expirationTimeMs - now;
-            }
-        }
-
-        void clear() {
-            elements.clear();
-        }
-
-    }
-
     /**
      * This is a dumbed down version of {@link MetadataUpdater} which is used to facilitate
      * metadata tracking primarily in order to serve {@link KafkaClient#leastLoadedNode(long)}
@@ -664,7 +628,7 @@ public void updateWithCurrentMetadata(Time time) {
 
         private void maybeCheckExpectedTopics(MetadataUpdate update, MetadataRequest.Builder builder) {
             if (update.expectMatchRefreshTopics) {
-                if (builder.topics() == null)
+                if (builder.isAllTopics())
                     throw new IllegalStateException("The metadata topics does not match expectation. "
                             + "Expected topics: " + update.topics()
                             + ", asked topics: ALL");
@@ -692,4 +656,92 @@ public void close() {
         }
     }
 
+    private static class ConnectionState {
+        enum State { CONNECTING, CONNECTED, DISCONNECTED }
+
+        private long throttledUntilMs = 0L;
+        private long readyDelayedUntilMs = 0L;
+        private long backingOffUntilMs = 0L;
+        private long unreachableUntilMs = 0L;
+        private State state = State.DISCONNECTED;
+
+        void backoff(long untilMs) {
+            backingOffUntilMs = untilMs;
+        }
+
+        void throttle(long untilMs) {
+            throttledUntilMs = untilMs;
+        }
+
+        void setUnreachable(long untilMs) {
+            unreachableUntilMs = untilMs;
+        }
+
+        void setReadyDelayed(long untilMs) {
+            readyDelayedUntilMs = untilMs;
+        }
+
+        boolean isReady(long now) {
+            return state == State.CONNECTED && notThrottled(now);
+        }
+
+        boolean isReadyDelayed(long now) {
+            return now < readyDelayedUntilMs;
+        }
+
+        boolean notThrottled(long now) {
+            return now > throttledUntilMs;
+        }
+
+        boolean isBackingOff(long now) {
+            return now < backingOffUntilMs;
+        }
+
+        boolean isUnreachable(long now) {
+            return now < unreachableUntilMs;
+        }
+
+        void disconnect() {
+            state = State.DISCONNECTED;
+        }
+
+        long connectionDelay(long now) {
+            if (state != State.DISCONNECTED)
+                return Long.MAX_VALUE;
+
+            if (backingOffUntilMs > now)
+                return backingOffUntilMs - now;
+
+            return 0;
+        }
+
+        boolean ready(long now) {
+            switch (state) {
+                case CONNECTED:
+                    return notThrottled(now);
+
+                case CONNECTING:
+                    if (isReadyDelayed(now))
+                        return false;
+                    state = State.CONNECTED;
+                    return ready(now);
+
+                case DISCONNECTED:
+                    if (isBackingOff(now)) {
+                        return false;
+                    } else if (isUnreachable(now)) {
+                        backingOffUntilMs = now + 100;
+                        return false;
+                    }
+
+                    state = State.CONNECTING;
+                    return ready(now);
+
+                default:
+                    throw new IllegalArgumentException("Invalid state: " + state);
+            }
+        }
+
+    }
+
 }
diff --git a/clients/src/test/java/org/apache/kafka/clients/NetworkClientTest.java b/clients/src/test/java/org/apache/kafka/clients/NetworkClientTest.java
index b40f690e788e8..e445c55413675 100644
--- a/clients/src/test/java/org/apache/kafka/clients/NetworkClientTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/NetworkClientTest.java
@@ -16,8 +16,16 @@
  */
 package org.apache.kafka.clients;
 
+import org.apache.kafka.common.Cluster;
+import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.Node;
 import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.errors.AuthenticationException;
+import org.apache.kafka.common.errors.UnsupportedVersionException;
+import org.apache.kafka.common.internals.ClusterResourceListeners;
+import org.apache.kafka.common.message.ApiVersionsResponseData;
+import org.apache.kafka.common.message.ApiVersionsResponseData.ApiVersionsResponseKey;
+import org.apache.kafka.common.message.ApiVersionsResponseData.ApiVersionsResponseKeyCollection;
 import org.apache.kafka.common.network.NetworkReceive;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.CommonFields;
@@ -26,8 +34,11 @@
 import org.apache.kafka.common.record.MemoryRecords;
 import org.apache.kafka.common.requests.ApiVersionsResponse;
 import org.apache.kafka.common.requests.MetadataRequest;
+import org.apache.kafka.common.requests.MetadataResponse;
 import org.apache.kafka.common.requests.ProduceRequest;
+import org.apache.kafka.common.requests.RequestHeader;
 import org.apache.kafka.common.requests.ResponseHeader;
+import org.apache.kafka.common.security.authenticator.SaslClientAuthenticator;
 import org.apache.kafka.common.utils.LogContext;
 import org.apache.kafka.common.utils.MockTime;
 import org.apache.kafka.test.DelayedReceive;
@@ -42,10 +53,16 @@
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
+import java.util.Optional;
+import java.util.Set;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
 
+import static org.apache.kafka.common.protocol.ApiKeys.PRODUCE;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotEquals;
+import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 
@@ -58,25 +75,32 @@ public class NetworkClientTest {
     protected final long reconnectBackoffMsTest = 10 * 1000;
     protected final long reconnectBackoffMaxMsTest = 10 * 10000;
 
+    private final TestMetadataUpdater metadataUpdater = new TestMetadataUpdater(Collections.singletonList(node));
     private final NetworkClient client = createNetworkClient(reconnectBackoffMaxMsTest);
     private final NetworkClient clientWithNoExponentialBackoff = createNetworkClient(reconnectBackoffMsTest);
     private final NetworkClient clientWithStaticNodes = createNetworkClientWithStaticNodes();
     private final NetworkClient clientWithNoVersionDiscovery = createNetworkClientWithNoVersionDiscovery();
 
     private NetworkClient createNetworkClient(long reconnectBackoffMaxMs) {
-        return new NetworkClient(selector, new ManualMetadataUpdater(Collections.singletonList(node)), "mock", Integer.MAX_VALUE,
+        return new NetworkClient(selector, metadataUpdater, "mock", Integer.MAX_VALUE,
                 reconnectBackoffMsTest, reconnectBackoffMaxMs, 64 * 1024, 64 * 1024,
                 defaultRequestTimeoutMs, ClientDnsLookup.DEFAULT, time, true, new ApiVersions(), new LogContext());
     }
 
     private NetworkClient createNetworkClientWithStaticNodes() {
-        return new NetworkClient(selector, new ManualMetadataUpdater(Collections.singletonList(node)),
+        return new NetworkClient(selector, metadataUpdater,
                 "mock-static", Integer.MAX_VALUE, 0, 0, 64 * 1024, 64 * 1024, defaultRequestTimeoutMs,
                 ClientDnsLookup.DEFAULT, time, true, new ApiVersions(), new LogContext());
     }
 
+    private NetworkClient createNetworkClientWithNoVersionDiscovery(Metadata metadata) {
+        return new NetworkClient(selector, metadata, "mock", Integer.MAX_VALUE,
+                reconnectBackoffMsTest, 0, 64 * 1024, 64 * 1024,
+                defaultRequestTimeoutMs, ClientDnsLookup.DEFAULT, time, false, new ApiVersions(), new LogContext());
+    }
+
     private NetworkClient createNetworkClientWithNoVersionDiscovery() {
-        return new NetworkClient(selector, new ManualMetadataUpdater(Collections.singletonList(node)), "mock", Integer.MAX_VALUE,
+        return new NetworkClient(selector, metadataUpdater, "mock", Integer.MAX_VALUE,
                 reconnectBackoffMsTest, reconnectBackoffMaxMsTest,
                 64 * 1024, 64 * 1024, defaultRequestTimeoutMs,
                 ClientDnsLookup.DEFAULT, time, false, new ApiVersions(), new LogContext());
@@ -140,18 +164,35 @@ public void testClose() {
         assertFalse("Connection should not be ready after close", client.isReady(node, 0));
     }
 
+    @Test
+    public void testUnsupportedVersionDuringInternalMetadataRequest() {
+        List topics = Arrays.asList("topic_1");
+
+        // disabling auto topic creation for versions less than 4 is not supported
+        MetadataRequest.Builder builder = new MetadataRequest.Builder(topics, false, (short) 3);
+        client.sendInternalMetadataRequest(builder, node.idString(), time.milliseconds());
+        assertEquals(UnsupportedVersionException.class, metadataUpdater.getAndClearFailure().getClass());
+    }
+
     private void checkSimpleRequestResponse(NetworkClient networkClient) {
         awaitReady(networkClient, node); // has to be before creating any request, as it may send ApiVersionsRequest and its response is mocked with correlation id 0
-        ProduceRequest.Builder builder = ProduceRequest.Builder.forCurrentMagic((short) 1, 1000,
-                        Collections.emptyMap());
+        ProduceRequest.Builder builder = new ProduceRequest.Builder(
+                PRODUCE.latestVersion(),
+                PRODUCE.latestVersion(),
+                (short) 1,
+                1000,
+                Collections.emptyMap(),
+                null);
         TestCallbackHandler handler = new TestCallbackHandler();
         ClientRequest request = networkClient.newClientRequest(
                 node.idString(), builder, time.milliseconds(), true, defaultRequestTimeoutMs, handler);
         networkClient.send(request, time.milliseconds());
         networkClient.poll(1, time.milliseconds());
         assertEquals(1, networkClient.inFlightRequestCount());
-        ResponseHeader respHeader = new ResponseHeader(request.correlationId());
-        Struct resp = new Struct(ApiKeys.PRODUCE.responseSchema(ApiKeys.PRODUCE.latestVersion()));
+        ResponseHeader respHeader =
+            new ResponseHeader(request.correlationId(),
+                request.apiKey().responseHeaderVersion(PRODUCE.latestVersion()));
+        Struct resp = new Struct(PRODUCE.responseSchema(PRODUCE.latestVersion()));
         resp.set("responses", new Object[0]);
         Struct responseHeaderStruct = respHeader.toStruct();
         int size = responseHeaderStruct.sizeOf() + resp.sizeOf();
@@ -168,21 +209,219 @@ private void checkSimpleRequestResponse(NetworkClient networkClient) {
                 request.correlationId(), handler.response.requestHeader().correlationId());
     }
 
-    private void setExpectedApiVersionsResponse(ApiVersionsResponse response) {
-        short apiVersionsResponseVersion = response.apiVersion(ApiKeys.API_VERSIONS.id).maxVersion;
-        ByteBuffer buffer = response.serialize(apiVersionsResponseVersion, new ResponseHeader(0));
+    private void delayedApiVersionsResponse(int correlationId, short version, ApiVersionsResponse response) {
+        ByteBuffer buffer = response.serialize(ApiKeys.API_VERSIONS, version, correlationId);
         selector.delayedReceive(new DelayedReceive(node.idString(), new NetworkReceive(node.idString(), buffer)));
     }
 
+    private void setExpectedApiVersionsResponse(ApiVersionsResponse response) {
+        short apiVersionsResponseVersion = response.apiVersion(ApiKeys.API_VERSIONS.id).maxVersion();
+        delayedApiVersionsResponse(0, apiVersionsResponseVersion, response);
+    }
+
     private void awaitReady(NetworkClient client, Node node) {
         if (client.discoverBrokerVersions()) {
-            setExpectedApiVersionsResponse(ApiVersionsResponse.defaultApiVersionsResponse());
+            setExpectedApiVersionsResponse(ApiVersionsResponse.DEFAULT_API_VERSIONS_RESPONSE);
         }
         while (!client.ready(node, time.milliseconds()))
             client.poll(1, time.milliseconds());
         selector.clear();
     }
 
+    @Test
+    public void testInvalidApiVersionsRequest() {
+        // initiate the connection
+        client.ready(node, time.milliseconds());
+
+        // handle the connection, send the ApiVersionsRequest
+        client.poll(0, time.milliseconds());
+
+        // check that the ApiVersionsRequest has been initiated
+        assertTrue(client.hasInFlightRequests(node.idString()));
+
+        // prepare response
+        delayedApiVersionsResponse(0, ApiKeys.API_VERSIONS.latestVersion(),
+            new ApiVersionsResponse(
+                new ApiVersionsResponseData()
+                    .setErrorCode(Errors.INVALID_REQUEST.code())
+                    .setThrottleTimeMs(0)
+            ));
+
+        // handle completed receives
+        client.poll(0, time.milliseconds());
+
+        // the ApiVersionsRequest is gone
+        assertFalse(client.hasInFlightRequests(node.idString()));
+
+        // various assertions
+        assertFalse(client.isReady(node, time.milliseconds()));
+    }
+
+    @Test
+    public void testApiVersionsRequest() {
+        // initiate the connection
+        client.ready(node, time.milliseconds());
+
+        // handle the connection, send the ApiVersionsRequest
+        client.poll(0, time.milliseconds());
+
+        // check that the ApiVersionsRequest has been initiated
+        assertTrue(client.hasInFlightRequests(node.idString()));
+
+        // prepare response
+        delayedApiVersionsResponse(0, ApiKeys.API_VERSIONS.latestVersion(),
+            ApiVersionsResponse.DEFAULT_API_VERSIONS_RESPONSE);
+
+        // handle completed receives
+        client.poll(0, time.milliseconds());
+
+        // the ApiVersionsRequest is gone
+        assertFalse(client.hasInFlightRequests(node.idString()));
+
+        // various assertions
+        assertTrue(client.isReady(node, time.milliseconds()));
+    }
+
+    @Test
+    public void testUnsupportedApiVersionsRequestWithVersionProvidedByTheBroker() {
+        // initiate the connection
+        client.ready(node, time.milliseconds());
+
+        // handle the connection, initiate first ApiVersionsRequest
+        client.poll(0, time.milliseconds());
+
+        // ApiVersionsRequest is in flight but not sent yet
+        assertTrue(client.hasInFlightRequests(node.idString()));
+
+        // completes initiated sends
+        client.poll(0, time.milliseconds());
+        assertEquals(1, selector.completedSends().size());
+
+        ByteBuffer buffer = selector.completedSendBuffers().get(0).buffer();
+        RequestHeader header = parseHeader(buffer);
+        assertEquals(ApiKeys.API_VERSIONS, header.apiKey());
+        assertEquals(3, header.apiVersion());
+
+        // prepare response
+        ApiVersionsResponseKeyCollection apiKeys = new ApiVersionsResponseKeyCollection();
+        apiKeys.add(new ApiVersionsResponseKey()
+            .setApiKey(ApiKeys.API_VERSIONS.id)
+            .setMinVersion((short) 0)
+            .setMaxVersion((short) 2));
+        delayedApiVersionsResponse(0, (short) 0,
+            new ApiVersionsResponse(
+                new ApiVersionsResponseData()
+                    .setErrorCode(Errors.UNSUPPORTED_VERSION.code())
+                    .setApiKeys(apiKeys)
+            ));
+
+        // handle ApiVersionResponse, initiate second ApiVersionRequest
+        client.poll(0, time.milliseconds());
+
+        // ApiVersionsRequest is in flight but not sent yet
+        assertTrue(client.hasInFlightRequests(node.idString()));
+
+        // ApiVersionsResponse has been received
+        assertEquals(1, selector.completedReceives().size());
+
+        // clean up the buffers
+        selector.completedSends().clear();
+        selector.completedSendBuffers().clear();
+        selector.completedReceives().clear();
+
+        // completes initiated sends
+        client.poll(0, time.milliseconds());
+
+        // ApiVersionsRequest has been sent
+        assertEquals(1, selector.completedSends().size());
+
+        buffer = selector.completedSendBuffers().get(0).buffer();
+        header = parseHeader(buffer);
+        assertEquals(ApiKeys.API_VERSIONS, header.apiKey());
+        assertEquals(2, header.apiVersion());
+
+        // prepare response
+        delayedApiVersionsResponse(1, (short) 0,
+            ApiVersionsResponse.DEFAULT_API_VERSIONS_RESPONSE);
+
+        // handle completed receives
+        client.poll(0, time.milliseconds());
+
+        // the ApiVersionsRequest is gone
+        assertFalse(client.hasInFlightRequests(node.idString()));
+        assertEquals(1, selector.completedReceives().size());
+
+        // the client is ready
+        assertTrue(client.isReady(node, time.milliseconds()));
+    }
+
+    @Test
+    public void testUnsupportedApiVersionsRequestWithoutVersionProvidedByTheBroker() {
+        // initiate the connection
+        client.ready(node, time.milliseconds());
+
+        // handle the connection, initiate first ApiVersionsRequest
+        client.poll(0, time.milliseconds());
+
+        // ApiVersionsRequest is in flight but not sent yet
+        assertTrue(client.hasInFlightRequests(node.idString()));
+
+        // completes initiated sends
+        client.poll(0, time.milliseconds());
+        assertEquals(1, selector.completedSends().size());
+
+        ByteBuffer buffer = selector.completedSendBuffers().get(0).buffer();
+        RequestHeader header = parseHeader(buffer);
+        assertEquals(ApiKeys.API_VERSIONS, header.apiKey());
+        assertEquals(3, header.apiVersion());
+
+        // prepare response
+        delayedApiVersionsResponse(0, (short) 0,
+            new ApiVersionsResponse(
+                new ApiVersionsResponseData()
+                    .setErrorCode(Errors.UNSUPPORTED_VERSION.code())
+            ));
+
+        // handle ApiVersionResponse, initiate second ApiVersionRequest
+        client.poll(0, time.milliseconds());
+
+        // ApiVersionsRequest is in flight but not sent yet
+        assertTrue(client.hasInFlightRequests(node.idString()));
+
+        // ApiVersionsResponse has been received
+        assertEquals(1, selector.completedReceives().size());
+
+        // clean up the buffers
+        selector.completedSends().clear();
+        selector.completedSendBuffers().clear();
+        selector.completedReceives().clear();
+
+        // completes initiated sends
+        client.poll(0, time.milliseconds());
+
+        // ApiVersionsRequest has been sent
+        assertEquals(1, selector.completedSends().size());
+        
+        buffer = selector.completedSendBuffers().get(0).buffer();
+        header = parseHeader(buffer);
+        assertEquals(ApiKeys.API_VERSIONS, header.apiKey());
+        assertEquals(0, header.apiVersion());
+
+        // prepare response
+        delayedApiVersionsResponse(1, (short) 0,
+            ApiVersionsResponse.DEFAULT_API_VERSIONS_RESPONSE);
+
+        // handle completed receives
+        client.poll(0, time.milliseconds());
+
+        // the ApiVersionsRequest is gone
+        assertFalse(client.hasInFlightRequests(node.idString()));
+        assertEquals(1, selector.completedReceives().size());
+
+        // the client is ready
+        assertTrue(client.isReady(node, time.milliseconds()));
+    }
+
     @Test
     public void testRequestTimeout() {
         awaitReady(client, node); // has to be before creating any request, as it may send ApiVersionsRequest and its response is mocked with correlation id 0
@@ -222,15 +461,22 @@ private void testRequestTimeout(ClientRequest request) {
     public void testConnectionThrottling() {
         // Instrument the test to return a response with a 100ms throttle delay.
         awaitReady(client, node);
-        ProduceRequest.Builder builder = ProduceRequest.Builder.forCurrentMagic((short) 1, 1000,
-            Collections.emptyMap());
+        ProduceRequest.Builder builder = new ProduceRequest.Builder(
+            PRODUCE.latestVersion(),
+            PRODUCE.latestVersion(),
+            (short) 1,
+            1000,
+            Collections.emptyMap(),
+            null);
         TestCallbackHandler handler = new TestCallbackHandler();
         ClientRequest request = client.newClientRequest(node.idString(), builder, time.milliseconds(), true,
                 defaultRequestTimeoutMs, handler);
         client.send(request, time.milliseconds());
         client.poll(1, time.milliseconds());
-        ResponseHeader respHeader = new ResponseHeader(request.correlationId());
-        Struct resp = new Struct(ApiKeys.PRODUCE.responseSchema(ApiKeys.PRODUCE.latestVersion()));
+        ResponseHeader respHeader =
+            new ResponseHeader(request.correlationId(),
+                request.apiKey().responseHeaderVersion(PRODUCE.latestVersion()));
+        Struct resp = new Struct(PRODUCE.responseSchema(PRODUCE.latestVersion()));
         resp.set("responses", new Object[0]);
         resp.set(CommonFields.THROTTLE_TIME_MS, 100);
         Struct responseHeaderStruct = respHeader.toStruct();
@@ -259,57 +505,87 @@ public void testConnectionThrottling() {
 
     // Creates expected ApiVersionsResponse from the specified node, where the max protocol version for the specified
     // key is set to the specified version.
-    private ApiVersionsResponse createExpectedApiVersionsResponse(Node node, ApiKeys key,
-        short apiVersionsMaxProtocolVersion) {
-        List versionList = new ArrayList<>();
+    private ApiVersionsResponse createExpectedApiVersionsResponse(ApiKeys key, short maxVersion) {
+        ApiVersionsResponseKeyCollection versionList = new ApiVersionsResponseKeyCollection();
         for (ApiKeys apiKey : ApiKeys.values()) {
             if (apiKey == key) {
-                versionList.add(new ApiVersionsResponse.ApiVersion(apiKey.id, (short) 0, apiVersionsMaxProtocolVersion));
+                versionList.add(new ApiVersionsResponseKey()
+                    .setApiKey(apiKey.id)
+                    .setMinVersion((short) 0)
+                    .setMaxVersion(maxVersion));
             } else {
-                versionList.add(new ApiVersionsResponse.ApiVersion(apiKey));
+                versionList.add(new ApiVersionsResponseKey()
+                    .setApiKey(apiKey.id)
+                    .setMinVersion(apiKey.oldestVersion())
+                    .setMaxVersion(apiKey.latestVersion()));
             }
         }
-        return new ApiVersionsResponse(0, Errors.NONE, versionList);
+        return new ApiVersionsResponse(new ApiVersionsResponseData()
+            .setErrorCode(Errors.NONE.code())
+            .setThrottleTimeMs(0)
+            .setApiKeys(versionList));
     }
 
     @Test
     public void testThrottlingNotEnabledForConnectionToOlderBroker() {
         // Instrument the test so that the max protocol version for PRODUCE returned from the node is 5 and thus
         // client-side throttling is not enabled. Also, return a response with a 100ms throttle delay.
-        setExpectedApiVersionsResponse(createExpectedApiVersionsResponse(node, ApiKeys.PRODUCE, (short) 5));
+        setExpectedApiVersionsResponse(createExpectedApiVersionsResponse(PRODUCE, (short) 5));
         while (!client.ready(node, time.milliseconds()))
             client.poll(1, time.milliseconds());
         selector.clear();
 
+        int correlationId = sendEmptyProduceRequest();
+        client.poll(1, time.milliseconds());
+
+        sendThrottledProduceResponse(correlationId, 100);
+        client.poll(1, time.milliseconds());
+
+        // Since client-side throttling is disabled, the connection is ready even though the response indicated a
+        // throttle delay.
+        assertTrue(client.ready(node, time.milliseconds()));
+        assertEquals(0, client.throttleDelayMs(node, time.milliseconds()));
+    }
+
+    private int sendEmptyProduceRequest() {
+        return sendEmptyProduceRequest(node.idString());
+    }
+
+    private int sendEmptyProduceRequest(String nodeId) {
         ProduceRequest.Builder builder = ProduceRequest.Builder.forCurrentMagic((short) 1, 1000,
-            Collections.emptyMap());
+                Collections.emptyMap());
         TestCallbackHandler handler = new TestCallbackHandler();
-        ClientRequest request = client.newClientRequest(node.idString(), builder, time.milliseconds(), true,
+        ClientRequest request = client.newClientRequest(nodeId, builder, time.milliseconds(), true,
                 defaultRequestTimeoutMs, handler);
         client.send(request, time.milliseconds());
-        client.poll(1, time.milliseconds());
-        ResponseHeader respHeader = new ResponseHeader(request.correlationId());
-        Struct resp = new Struct(ApiKeys.PRODUCE.responseSchema(ApiKeys.PRODUCE.latestVersion()));
-        resp.set("responses", new Object[0]);
-        resp.set(CommonFields.THROTTLE_TIME_MS, 100);
+        return request.correlationId();
+    }
+
+    private void sendResponse(ResponseHeader respHeader, Struct response) {
         Struct responseHeaderStruct = respHeader.toStruct();
-        int size = responseHeaderStruct.sizeOf() + resp.sizeOf();
+        int size = responseHeaderStruct.sizeOf() + response.sizeOf();
         ByteBuffer buffer = ByteBuffer.allocate(size);
         responseHeaderStruct.writeTo(buffer);
-        resp.writeTo(buffer);
+        response.writeTo(buffer);
         buffer.flip();
         selector.completeReceive(new NetworkReceive(node.idString(), buffer));
-        client.poll(1, time.milliseconds());
+    }
 
-        // Since client-side throttling is disabled, the connection is ready even though the response indicated a
-        // throttle delay.
-        assertTrue(client.ready(node, time.milliseconds()));
-        assertEquals(0, client.throttleDelayMs(node, time.milliseconds()));
+    private void sendThrottledProduceResponse(int correlationId, int throttleMs) {
+        Struct resp = new Struct(PRODUCE.responseSchema(PRODUCE.latestVersion()));
+        resp.set("responses", new Object[0]);
+        resp.set(CommonFields.THROTTLE_TIME_MS, throttleMs);
+        sendResponse(new ResponseHeader(correlationId,
+            PRODUCE.responseHeaderVersion(PRODUCE.latestVersion())),
+            resp);
     }
 
     @Test
     public void testLeastLoadedNode() {
         client.ready(node, time.milliseconds());
+        assertFalse(client.isReady(node, time.milliseconds()));
+        assertEquals(node, client.leastLoadedNode(time.milliseconds()));
+
         awaitReady(client, node);
         client.poll(1, time.milliseconds());
         assertTrue("The client should be ready", client.isReady(node, time.milliseconds()));
@@ -330,6 +606,66 @@ public void testLeastLoadedNode() {
         assertNull("There should be NO leastloadednode", leastNode);
     }
 
+    @Test
+    public void testAuthenticationFailureWithInFlightMetadataRequest() {
+        int refreshBackoffMs = 50;
+
+        MetadataResponse metadataResponse = TestUtils.metadataUpdateWith(2, Collections.emptyMap());
+        Metadata metadata = new Metadata(refreshBackoffMs, 5000, new LogContext(), new ClusterResourceListeners());
+        metadata.update(metadataResponse, time.milliseconds());
+
+        Cluster cluster = metadata.fetch();
+        Node node1 = cluster.nodes().get(0);
+        Node node2 = cluster.nodes().get(1);
+
+        NetworkClient client = createNetworkClientWithNoVersionDiscovery(metadata);
+
+        awaitReady(client, node1);
+
+        metadata.requestUpdate();
+        time.sleep(refreshBackoffMs);
+
+        client.poll(0, time.milliseconds());
+
+        Optional nodeWithPendingMetadataOpt = cluster.nodes().stream()
+                .filter(node -> client.hasInFlightRequests(node.idString()))
+                .findFirst();
+        assertEquals(Optional.of(node1), nodeWithPendingMetadataOpt);
+
+        assertFalse(client.ready(node2, time.milliseconds()));
+        selector.serverAuthenticationFailed(node2.idString());
+        client.poll(0, time.milliseconds());
+        assertNotNull(client.authenticationException(node2));
+
+        ByteBuffer requestBuffer = selector.completedSendBuffers().get(0).buffer();
+        RequestHeader header = parseHeader(requestBuffer);
+        assertEquals(ApiKeys.METADATA, header.apiKey());
+
+        ByteBuffer responseBuffer = metadataResponse.serialize(ApiKeys.METADATA, header.apiVersion(), header.correlationId());
+        selector.delayedReceive(new DelayedReceive(node1.idString(), new NetworkReceive(node1.idString(), responseBuffer)));
+
+        int initialUpdateVersion = metadata.updateVersion();
+        client.poll(0, time.milliseconds());
+        assertEquals(initialUpdateVersion + 1, metadata.updateVersion());
+    }
+
+    @Test
+    public void testLeastLoadedNodeConsidersThrottledConnections() {
+        client.ready(node, time.milliseconds());
+        awaitReady(client, node);
+        client.poll(1, time.milliseconds());
+        assertTrue("The client should be ready", client.isReady(node, time.milliseconds()));
+
+        int correlationId = sendEmptyProduceRequest();
+        client.poll(1, time.milliseconds());
+
+        sendThrottledProduceResponse(correlationId, 100);
+        client.poll(1, time.milliseconds());
+
+        // leastloadednode should return null since the node is throttled
+        assertNull(client.leastLoadedNode(time.milliseconds()));
+    }
+
     @Test
     public void testConnectionDelayWithNoExponentialBackoff() {
         long now = time.milliseconds();
@@ -530,6 +866,21 @@ public void testCallDisconnect() throws Exception {
         assertTrue(client.canConnect(node, time.milliseconds()));
     }
 
+    @Test
+    public void testCorrelationId() {
+        int count = 100;
+        Set ids = IntStream.range(0, count)
+            .mapToObj(i -> client.nextCorrelationId())
+            .collect(Collectors.toSet());
+        assertEquals(count, ids.size());
+        ids.forEach(id -> assertTrue(id < SaslClientAuthenticator.MIN_RESERVED_CORRELATION_ID));
+    }
+
+    private RequestHeader parseHeader(ByteBuffer buffer) {
+        buffer.getInt(); // skip size
+        return RequestHeader.parse(buffer.slice());
+    }
+
     private void awaitInFlightApiVersionRequest() throws Exception {
         client.ready(node, time.milliseconds());
         TestUtils.waitForCondition(new TestCondition() {
@@ -551,4 +902,34 @@ public void onComplete(ClientResponse response) {
             this.response = response;
         }
     }
+
+    // ManualMetadataUpdater with ability to keep track of failures
+    private static class TestMetadataUpdater extends ManualMetadataUpdater {
+        KafkaException failure;
+
+        public TestMetadataUpdater(List nodes) {
+            super(nodes);
+        }
+
+        @Override
+        public void handleServerDisconnect(long now, String destinationId, Optional maybeAuthException) {
+            maybeAuthException.ifPresent(exception -> {
+                failure = exception;
+            });
+            super.handleServerDisconnect(now, destinationId, maybeAuthException);
+        }
+
+        @Override
+        public void handleFailedRequest(long now, Optional maybeFatalException) {
+            maybeFatalException.ifPresent(exception -> {
+                failure = exception;
+            });
+        }
+
+        public KafkaException getAndClearFailure() {
+            KafkaException failure = this.failure;
+            this.failure = null;
+            return failure;
+        }
+    }
 }
diff --git a/clients/src/test/java/org/apache/kafka/clients/NodeApiVersionsTest.java b/clients/src/test/java/org/apache/kafka/clients/NodeApiVersionsTest.java
index 88d0c2e719e2a..99af49ef73551 100644
--- a/clients/src/test/java/org/apache/kafka/clients/NodeApiVersionsTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/NodeApiVersionsTest.java
@@ -16,17 +16,16 @@
  */
 package org.apache.kafka.clients;
 
+import java.util.ArrayList;
+import java.util.LinkedList;
+import java.util.List;
 import org.apache.kafka.common.errors.UnsupportedVersionException;
+import org.apache.kafka.common.message.ApiVersionsResponseData.ApiVersionsResponseKey;
+import org.apache.kafka.common.message.ApiVersionsResponseData.ApiVersionsResponseKeyCollection;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.requests.ApiVersionsResponse;
-import org.apache.kafka.common.requests.ApiVersionsResponse.ApiVersion;
 import org.junit.Test;
 
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.LinkedList;
-import java.util.List;
-
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 
@@ -34,7 +33,7 @@ public class NodeApiVersionsTest {
 
     @Test
     public void testUnsupportedVersionsToString() {
-        NodeApiVersions versions = new NodeApiVersions(Collections.emptyList());
+        NodeApiVersions versions = new NodeApiVersions(new ApiVersionsResponseKeyCollection());
         StringBuilder bld = new StringBuilder();
         String prefix = "(";
         for (ApiKeys apiKey : ApiKeys.values()) {
@@ -48,8 +47,7 @@ public void testUnsupportedVersionsToString() {
 
     @Test
     public void testUnknownApiVersionsToString() {
-        ApiVersion unknownApiVersion = new ApiVersion((short) 337, (short) 0, (short) 1);
-        NodeApiVersions versions = new NodeApiVersions(Collections.singleton(unknownApiVersion));
+        NodeApiVersions versions = NodeApiVersions.create((short) 337, (short) 0, (short) 1);
         assertTrue(versions.toString().endsWith("UNKNOWN(337): 0 to 1)"));
     }
 
@@ -92,8 +90,7 @@ public void testVersionsToString() {
 
     @Test
     public void testLatestUsableVersion() {
-        NodeApiVersions apiVersions = NodeApiVersions.create(Collections.singleton(
-                new ApiVersion(ApiKeys.PRODUCE.id, (short) 1, (short) 3)));
+        NodeApiVersions apiVersions = NodeApiVersions.create(ApiKeys.PRODUCE.id, (short) 1, (short) 3);
         assertEquals(3, apiVersions.latestUsableVersion(ApiKeys.PRODUCE));
         assertEquals(1, apiVersions.latestUsableVersion(ApiKeys.PRODUCE, (short) 0, (short) 1));
         assertEquals(1, apiVersions.latestUsableVersion(ApiKeys.PRODUCE, (short) 1, (short) 1));
@@ -107,43 +104,52 @@ public void testLatestUsableVersion() {
 
     @Test(expected = UnsupportedVersionException.class)
     public void testLatestUsableVersionOutOfRangeLow() {
-        NodeApiVersions apiVersions = NodeApiVersions.create(Collections.singleton(
-                new ApiVersion(ApiKeys.PRODUCE.id, (short) 1, (short) 2)));
+        NodeApiVersions apiVersions = NodeApiVersions.create(ApiKeys.PRODUCE.id, (short) 1, (short) 2);
         apiVersions.latestUsableVersion(ApiKeys.PRODUCE, (short) 3, (short) 4);
     }
 
     @Test(expected = UnsupportedVersionException.class)
     public void testLatestUsableVersionOutOfRangeHigh() {
-        NodeApiVersions apiVersions = NodeApiVersions.create(Collections.singleton(
-                new ApiVersion(ApiKeys.PRODUCE.id, (short) 2, (short) 3)));
+        NodeApiVersions apiVersions = NodeApiVersions.create(ApiKeys.PRODUCE.id, (short) 2, (short) 3);
         apiVersions.latestUsableVersion(ApiKeys.PRODUCE, (short) 0, (short) 1);
     }
 
     @Test(expected = UnsupportedVersionException.class)
     public void testUsableVersionCalculationNoKnownVersions() {
-        List versionList = new ArrayList<>();
-        NodeApiVersions versions =  new NodeApiVersions(versionList);
+        NodeApiVersions versions = new NodeApiVersions(new ApiVersionsResponseKeyCollection());
         versions.latestUsableVersion(ApiKeys.FETCH);
     }
 
     @Test(expected = UnsupportedVersionException.class)
     public void testLatestUsableVersionOutOfRange() {
-        NodeApiVersions apiVersions = NodeApiVersions.create(Collections.singleton(
-                new ApiVersion(ApiKeys.PRODUCE.id, (short) 300, (short) 300)));
+        NodeApiVersions apiVersions = NodeApiVersions.create(ApiKeys.PRODUCE.id, (short) 300, (short) 300);
         apiVersions.latestUsableVersion(ApiKeys.PRODUCE);
     }
 
     @Test
     public void testUsableVersionLatestVersions() {
         List versionList = new LinkedList<>();
-        for (ApiVersion apiVersion: ApiVersionsResponse.defaultApiVersionsResponse().apiVersions()) {
-            versionList.add(apiVersion);
+        for (ApiVersionsResponseKey apiVersion: ApiVersionsResponse.DEFAULT_API_VERSIONS_RESPONSE.data.apiKeys()) {
+            versionList.add(new ApiVersion(apiVersion));
         }
         // Add an API key that we don't know about.
         versionList.add(new ApiVersion((short) 100, (short) 0, (short) 1));
-        NodeApiVersions versions =  new NodeApiVersions(versionList);
+        NodeApiVersions versions = new NodeApiVersions(versionList);
         for (ApiKeys apiKey: ApiKeys.values()) {
             assertEquals(apiKey.latestVersion(), versions.latestUsableVersion(apiKey));
         }
     }
+
+    @Test
+    public void testConstructionFromApiVersionsResponse() {
+        ApiVersionsResponse apiVersionsResponse = ApiVersionsResponse.DEFAULT_API_VERSIONS_RESPONSE;
+        NodeApiVersions versions = new NodeApiVersions(apiVersionsResponse.data.apiKeys());
+
+        for (ApiVersionsResponseKey apiVersionKey : apiVersionsResponse.data.apiKeys()) {
+            ApiVersion apiVersion = versions.apiVersion(ApiKeys.forId(apiVersionKey.apiKey()));
+            assertEquals(apiVersionKey.apiKey(), apiVersion.apiKey);
+            assertEquals(apiVersionKey.minVersion(), apiVersion.minVersion);
+            assertEquals(apiVersionKey.maxVersion(), apiVersion.maxVersion);
+        }
+    }
 }
diff --git a/clients/src/test/java/org/apache/kafka/clients/admin/AdminClientUnitTestEnv.java b/clients/src/test/java/org/apache/kafka/clients/admin/AdminClientUnitTestEnv.java
index 6023c63022b29..f6a808f74727e 100644
--- a/clients/src/test/java/org/apache/kafka/clients/admin/AdminClientUnitTestEnv.java
+++ b/clients/src/test/java/org/apache/kafka/clients/admin/AdminClientUnitTestEnv.java
@@ -23,6 +23,7 @@
 import org.apache.kafka.common.utils.LogContext;
 import org.apache.kafka.common.utils.Time;
 
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -33,7 +34,7 @@
  * easily create a simple cluster.
  * 
      
  * To use in a test, create an instance and prepare its {@link #kafkaClient() MockClient} with the expected responses
- * for the {@link AdminClient}. Then, use the {@link #adminClient() AdminClient} in the test, which will then use the MockClient
+ * for the {@link Admin}. Then, use the {@link #adminClient() AdminClient} in the test, which will then use the MockClient
  * and receive the responses you provided.
  *
  * Since {@link #kafkaClient() MockClient} is not thread-safe,
@@ -53,14 +54,18 @@ public AdminClientUnitTestEnv(Cluster cluster, String... vals) {
     }
 
     public AdminClientUnitTestEnv(Time time, Cluster cluster, String... vals) {
-        this(time, cluster, newStrMap(vals));
+        this(time, cluster, clientConfigs(vals));
     }
 
     public AdminClientUnitTestEnv(Time time, Cluster cluster) {
-        this(time, cluster, newStrMap());
+        this(time, cluster, clientConfigs());
     }
 
     public AdminClientUnitTestEnv(Time time, Cluster cluster, Map config) {
+        this(time, cluster, config, Collections.emptyMap());
+    }
+
+    public AdminClientUnitTestEnv(Time time, Cluster cluster, Map config, Map unreachableNodes) {
         this.time = time;
         this.cluster = cluster;
         AdminClientConfig adminClientConfig = new AdminClientConfig(config);
@@ -86,6 +91,7 @@ public void update(Time time, MockClient.MetadataUpdate update) {
         });
 
         metadataManager.update(cluster, time.milliseconds());
+        unreachableNodes.forEach(mockClient::setUnreachable);
         this.adminClient = KafkaAdminClient.createInternal(adminClientConfig, metadataManager, mockClient, time);
     }
 
@@ -97,7 +103,7 @@ public Cluster cluster() {
         return cluster;
     }
 
-    public AdminClient adminClient() {
+    public Admin adminClient() {
         return adminClient;
     }
 
@@ -110,15 +116,15 @@ public void close() {
         this.adminClient.close();
     }
 
-    private static Map newStrMap(String... vals) {
+    static Map clientConfigs(String... overrides) {
         Map map = new HashMap<>();
         map.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:8121");
         map.put(AdminClientConfig.REQUEST_TIMEOUT_MS_CONFIG, "1000");
-        if (vals.length % 2 != 0) {
+        if (overrides.length % 2 != 0) {
             throw new IllegalStateException();
         }
-        for (int i = 0; i < vals.length; i += 2) {
-            map.put(vals[i], vals[i + 1]);
+        for (int i = 0; i < overrides.length; i += 2) {
+            map.put(overrides[i], overrides[i + 1]);
         }
         return map;
     }
diff --git a/clients/src/test/java/org/apache/kafka/clients/admin/DeleteConsumerGroupOffsetsResultTest.java b/clients/src/test/java/org/apache/kafka/clients/admin/DeleteConsumerGroupOffsetsResultTest.java
new file mode 100644
index 0000000000000..19ce76da7dc64
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/clients/admin/DeleteConsumerGroupOffsetsResultTest.java
@@ -0,0 +1,118 @@
+/*
+ * 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.admin;
+
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.errors.GroupAuthorizationException;
+import org.apache.kafka.common.errors.UnknownTopicOrPartitionException;
+import org.apache.kafka.common.internals.KafkaFutureImpl;
+
+import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.test.TestUtils;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.ExecutionException;
+
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertThrows;
+
+public class DeleteConsumerGroupOffsetsResultTest {
+
+    private final String topic = "topic";
+    private final TopicPartition tpZero = new TopicPartition(topic, 0);
+    private final TopicPartition tpOne = new TopicPartition(topic, 1);
+    private Set partitions;
+    private Map errorsMap;
+
+    private KafkaFutureImpl> partitionFutures;
+
+    @Before
+    public void setUp() {
+        partitionFutures = new KafkaFutureImpl<>();
+        partitions = new HashSet<>();
+        partitions.add(tpZero);
+        partitions.add(tpOne);
+
+        errorsMap = new HashMap<>();
+        errorsMap.put(tpZero, Errors.NONE);
+        errorsMap.put(tpOne, Errors.UNKNOWN_TOPIC_OR_PARTITION);
+    }
+
+    @Test
+    public void testTopLevelErrorConstructor() throws InterruptedException {
+        partitionFutures.completeExceptionally(Errors.GROUP_AUTHORIZATION_FAILED.exception());
+        DeleteConsumerGroupOffsetsResult topLevelErrorResult =
+            new DeleteConsumerGroupOffsetsResult(partitionFutures, partitions);
+        TestUtils.assertFutureError(topLevelErrorResult.all(), GroupAuthorizationException.class);
+    }
+
+    @Test
+    public void testPartitionLevelErrorConstructor() throws ExecutionException, InterruptedException {
+        createAndVerifyPartitionLevelErrror();
+    }
+
+    @Test
+    public void testPartitionMissingInResponseErrorConstructor() throws InterruptedException, ExecutionException {
+        errorsMap.remove(tpOne);
+        partitionFutures.complete(errorsMap);
+        assertFalse(partitionFutures.isCompletedExceptionally());
+        DeleteConsumerGroupOffsetsResult missingPartitionResult =
+            new DeleteConsumerGroupOffsetsResult(partitionFutures, partitions);
+
+        TestUtils.assertFutureError(missingPartitionResult.all(), IllegalArgumentException.class);
+        assertNull(missingPartitionResult.partitionResult(tpZero).get());
+        TestUtils.assertFutureError(missingPartitionResult.partitionResult(tpOne), IllegalArgumentException.class);
+    }
+
+    @Test
+    public void testPartitionMissingInRequestErrorConstructor() throws InterruptedException, ExecutionException {
+        DeleteConsumerGroupOffsetsResult partitionLevelErrorResult = createAndVerifyPartitionLevelErrror();
+        assertThrows(IllegalArgumentException.class, () -> partitionLevelErrorResult.partitionResult(new TopicPartition("invalid-topic", 0)));
+    }
+
+    @Test
+    public void testNoErrorConstructor() throws ExecutionException, InterruptedException {
+        Map errorsMap = new HashMap<>();
+        errorsMap.put(tpZero, Errors.NONE);
+        errorsMap.put(tpOne, Errors.NONE);
+        DeleteConsumerGroupOffsetsResult noErrorResult =
+            new DeleteConsumerGroupOffsetsResult(partitionFutures, partitions);
+        partitionFutures.complete(errorsMap);
+
+        assertNull(noErrorResult.all().get());
+        assertNull(noErrorResult.partitionResult(tpZero).get());
+        assertNull(noErrorResult.partitionResult(tpOne).get());
+    }
+
+    private DeleteConsumerGroupOffsetsResult createAndVerifyPartitionLevelErrror() throws InterruptedException, ExecutionException {
+        partitionFutures.complete(errorsMap);
+        assertFalse(partitionFutures.isCompletedExceptionally());
+        DeleteConsumerGroupOffsetsResult partitionLevelErrorResult =
+            new DeleteConsumerGroupOffsetsResult(partitionFutures, partitions);
+
+        TestUtils.assertFutureError(partitionLevelErrorResult.all(), UnknownTopicOrPartitionException.class);
+        assertNull(partitionLevelErrorResult.partitionResult(tpZero).get());
+        TestUtils.assertFutureError(partitionLevelErrorResult.partitionResult(tpOne), UnknownTopicOrPartitionException.class);
+        return partitionLevelErrorResult;
+    }
+}
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 782dc1660f71f..4105f802c50a0 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
@@ -21,10 +21,11 @@
 import org.apache.kafka.clients.MockClient;
 import org.apache.kafka.clients.NodeApiVersions;
 import org.apache.kafka.clients.admin.DeleteAclsResult.FilterResults;
+import org.apache.kafka.clients.consumer.ConsumerPartitionAssignor;
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
 import org.apache.kafka.clients.consumer.internals.ConsumerProtocol;
-import org.apache.kafka.clients.consumer.internals.PartitionAssignor;
 import org.apache.kafka.common.Cluster;
+import org.apache.kafka.common.ElectionType;
 import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.KafkaFuture;
 import org.apache.kafka.common.Node;
@@ -39,7 +40,9 @@
 import org.apache.kafka.common.config.ConfigResource;
 import org.apache.kafka.common.errors.AuthenticationException;
 import org.apache.kafka.common.errors.ClusterAuthorizationException;
+import org.apache.kafka.common.errors.FencedInstanceIdException;
 import org.apache.kafka.common.errors.GroupAuthorizationException;
+import org.apache.kafka.common.errors.GroupSubscribedToTopicException;
 import org.apache.kafka.common.errors.InvalidRequestException;
 import org.apache.kafka.common.errors.InvalidTopicException;
 import org.apache.kafka.common.errors.LeaderNotAvailableException;
@@ -48,16 +51,37 @@
 import org.apache.kafka.common.errors.SaslAuthenticationException;
 import org.apache.kafka.common.errors.SecurityDisabledException;
 import org.apache.kafka.common.errors.TimeoutException;
+import org.apache.kafka.common.errors.TopicAuthorizationException;
 import org.apache.kafka.common.errors.TopicDeletionDisabledException;
+import org.apache.kafka.common.errors.UnknownMemberIdException;
 import org.apache.kafka.common.errors.UnknownServerException;
 import org.apache.kafka.common.errors.UnknownTopicOrPartitionException;
+import org.apache.kafka.common.message.AlterPartitionReassignmentsResponseData;
 import org.apache.kafka.common.message.CreateTopicsResponseData;
 import org.apache.kafka.common.message.CreateTopicsResponseData.CreatableTopicResult;
+import org.apache.kafka.common.message.DeleteGroupsResponseData;
+import org.apache.kafka.common.message.DeleteGroupsResponseData.DeletableGroupResult;
+import org.apache.kafka.common.message.DeleteGroupsResponseData.DeletableGroupResultCollection;
+import org.apache.kafka.common.message.DeleteTopicsResponseData;
+import org.apache.kafka.common.message.DeleteTopicsResponseData.DeletableTopicResult;
 import org.apache.kafka.common.message.DescribeGroupsResponseData;
-import org.apache.kafka.common.message.ElectPreferredLeadersResponseData;
-import org.apache.kafka.common.message.ElectPreferredLeadersResponseData.PartitionResult;
-import org.apache.kafka.common.message.ElectPreferredLeadersResponseData.ReplicaElectionResult;
+import org.apache.kafka.common.message.DescribeGroupsResponseData.DescribedGroupMember;
+import org.apache.kafka.common.message.ElectLeadersResponseData.PartitionResult;
+import org.apache.kafka.common.message.ElectLeadersResponseData.ReplicaElectionResult;
+import org.apache.kafka.common.message.IncrementalAlterConfigsResponseData;
+import org.apache.kafka.common.message.IncrementalAlterConfigsResponseData.AlterConfigsResourceResponse;
+import org.apache.kafka.common.message.LeaveGroupRequestData.MemberIdentity;
+import org.apache.kafka.common.message.LeaveGroupResponseData;
+import org.apache.kafka.common.message.LeaveGroupResponseData.MemberResponse;
+import org.apache.kafka.common.message.ListGroupsResponseData;
+import org.apache.kafka.common.message.ListPartitionReassignmentsResponseData;
+import org.apache.kafka.common.message.OffsetDeleteResponseData;
+import org.apache.kafka.common.message.OffsetDeleteResponseData.OffsetDeleteResponsePartition;
+import org.apache.kafka.common.message.OffsetDeleteResponseData.OffsetDeleteResponsePartitionCollection;
+import org.apache.kafka.common.message.OffsetDeleteResponseData.OffsetDeleteResponseTopic;
+import org.apache.kafka.common.message.OffsetDeleteResponseData.OffsetDeleteResponseTopicCollection;
 import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.common.requests.AlterPartitionReassignmentsResponse;
 import org.apache.kafka.common.requests.ApiError;
 import org.apache.kafka.common.requests.CreateAclsResponse;
 import org.apache.kafka.common.requests.CreateAclsResponse.AclCreationResponse;
@@ -74,11 +98,15 @@
 import org.apache.kafka.common.requests.DescribeAclsResponse;
 import org.apache.kafka.common.requests.DescribeConfigsResponse;
 import org.apache.kafka.common.requests.DescribeGroupsResponse;
-import org.apache.kafka.common.requests.ElectPreferredLeadersResponse;
+import org.apache.kafka.common.requests.ElectLeadersResponse;
 import org.apache.kafka.common.requests.FindCoordinatorResponse;
+import org.apache.kafka.common.requests.IncrementalAlterConfigsResponse;
+import org.apache.kafka.common.requests.LeaveGroupResponse;
 import org.apache.kafka.common.requests.ListGroupsResponse;
+import org.apache.kafka.common.requests.ListPartitionReassignmentsResponse;
 import org.apache.kafka.common.requests.MetadataRequest;
 import org.apache.kafka.common.requests.MetadataResponse;
+import org.apache.kafka.common.requests.OffsetDeleteResponse;
 import org.apache.kafka.common.requests.OffsetFetchResponse;
 import org.apache.kafka.common.resource.PatternType;
 import org.apache.kafka.common.resource.ResourcePattern;
@@ -99,6 +127,7 @@
 import java.net.InetSocketAddress;
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
@@ -108,14 +137,23 @@
 import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 import static java.util.Arrays.asList;
 import static java.util.Collections.singletonList;
+import static org.apache.kafka.common.message.AlterPartitionReassignmentsResponseData.ReassignablePartitionResponse;
+import static org.apache.kafka.common.message.AlterPartitionReassignmentsResponseData.ReassignableTopicResponse;
+import static org.apache.kafka.common.message.ListPartitionReassignmentsResponseData.OngoingPartitionReassignment;
+import static org.apache.kafka.common.message.ListPartitionReassignmentsResponseData.OngoingTopicReassignment;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
@@ -203,7 +241,7 @@ private static Cluster mockCluster(int controllerIndex) {
 
     private static Cluster mockBootstrapCluster() {
         return Cluster.bootstrap(ClientUtils.parseAndValidateAddresses(
-                Collections.singletonList("localhost:8121"), ClientDnsLookup.DEFAULT));
+                singletonList("localhost:8121"), ClientDnsLookup.DEFAULT));
     }
 
     private static AdminClientUnitTestEnv mockClientEnv(String... configVals) {
@@ -216,6 +254,30 @@ public void testCloseAdminClient() {
         }
     }
 
+    private static OffsetDeleteResponse prepareOffsetDeleteResponse(Errors error) {
+        return new OffsetDeleteResponse(
+            new OffsetDeleteResponseData()
+                .setErrorCode(error.code())
+                .setTopics(new OffsetDeleteResponseTopicCollection())
+        );
+    }
+
+    private static OffsetDeleteResponse prepareOffsetDeleteResponse(String topic, int partition, Errors error) {
+        return new OffsetDeleteResponse(
+            new OffsetDeleteResponseData()
+                .setErrorCode(Errors.NONE.code())
+                .setTopics(new OffsetDeleteResponseTopicCollection(Stream.of(
+                    new OffsetDeleteResponseTopic()
+                        .setName(topic)
+                        .setPartitions(new OffsetDeleteResponsePartitionCollection(Collections.singletonList(
+                            new OffsetDeleteResponsePartition()
+                                .setPartitionIndex(partition)
+                                .setErrorCode(error.code())
+                        ).iterator()))
+                ).collect(Collectors.toList()).iterator()))
+        );
+    }
+
     private static CreateTopicsResponse prepareCreateTopicsResponse(String topicName, Errors error) {
         CreateTopicsResponseData data = new CreateTopicsResponseData();
         data.topics().add(new CreatableTopicResult().
@@ -223,6 +285,18 @@ private static CreateTopicsResponse prepareCreateTopicsResponse(String topicName
         return new CreateTopicsResponse(data);
     }
 
+    private static DeleteTopicsResponse prepareDeleteTopicsResponse(String topicName, Errors error) {
+        DeleteTopicsResponseData data = new DeleteTopicsResponseData();
+        data.responses().add(new DeletableTopicResult()
+                .setName(topicName)
+                .setErrorCode(error.code()));
+        return new DeleteTopicsResponse(data);
+    }
+
+    private static FindCoordinatorResponse prepareFindCoordinatorResponse(Errors error, Node node) {
+        return FindCoordinatorResponse.prepareResponse(error, node);
+    }
+
     /**
      * Test that the client properly times out when we don't receive any metadata.
      */
@@ -250,7 +324,7 @@ public void testConnectionFailureOnMetadataUpdate() throws Exception {
             env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
             env.kafkaClient().prepareResponse(request -> request instanceof MetadataRequest, null, true);
             env.kafkaClient().prepareResponse(request -> request instanceof MetadataRequest,
-                    new MetadataResponse(discoveredCluster.nodes(), discoveredCluster.clusterResource().clusterId(),
+                    MetadataResponse.prepareResponse(discoveredCluster.nodes(), discoveredCluster.clusterResource().clusterId(),
                             1, Collections.emptyList()));
             env.kafkaClient().prepareResponse(body -> body instanceof CreateTopicsRequest,
                     prepareCreateTopicsResponse("myTopic", Errors.NONE));
@@ -268,13 +342,14 @@ public void testUnreachableBootstrapServer() throws Exception {
         // This tests the scenario in which the bootstrap server is unreachable for a short while,
         // which prevents AdminClient from being able to send the initial metadata request
 
-        Cluster cluster = Cluster.bootstrap(Collections.singletonList(new InetSocketAddress("localhost", 8121)));
-        try (final AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(Time.SYSTEM, cluster)) {
+        Cluster cluster = Cluster.bootstrap(singletonList(new InetSocketAddress("localhost", 8121)));
+        Map unreachableNodes = Collections.singletonMap(cluster.nodes().get(0), 200L);
+        try (final AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(Time.SYSTEM, cluster,
+                AdminClientUnitTestEnv.clientConfigs(), unreachableNodes)) {
             Cluster discoveredCluster = mockCluster(0);
             env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
-            env.kafkaClient().setUnreachable(cluster.nodes().get(0), 200);
             env.kafkaClient().prepareResponse(body -> body instanceof MetadataRequest,
-                    new  MetadataResponse(discoveredCluster.nodes(), discoveredCluster.clusterResource().clusterId(),
+                    MetadataResponse.prepareResponse(discoveredCluster.nodes(), discoveredCluster.clusterResource().clusterId(),
                             1, Collections.emptyList()));
             env.kafkaClient().prepareResponse(body -> body instanceof CreateTopicsRequest,
                 prepareCreateTopicsResponse("myTopic", Errors.NONE));
@@ -352,6 +427,11 @@ public void testCreateTopicsRetryBackoff() throws Exception {
             // Wait until the first attempt has failed, then advance the time
             TestUtils.waitForCondition(() -> mockClient.numAwaitingResponses() == 1,
                     "Failed awaiting CreateTopics first request failure");
+
+            // Wait until the retry call added to the queue in AdminClient
+            TestUtils.waitForCondition(() -> ((KafkaAdminClient) env.adminClient()).numPendingCalls() == 1,
+                "Failed to add retry CreateTopics call");
+
             time.sleep(retryBackoff);
 
             future.get();
@@ -369,7 +449,7 @@ public void testCreateTopicsHandleNotControllerException() throws Exception {
             env.kafkaClient().prepareResponseFrom(
                 prepareCreateTopicsResponse("myTopic", Errors.NOT_CONTROLLER),
                 env.cluster().nodeById(0));
-            env.kafkaClient().prepareResponse(new MetadataResponse(env.cluster().nodes(),
+            env.kafkaClient().prepareResponse(MetadataResponse.prepareResponse(env.cluster().nodes(),
                 env.cluster().clusterResource().clusterId(),
                 1,
                 Collections.emptyList()));
@@ -389,20 +469,20 @@ public void testDeleteTopics() throws Exception {
             env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
 
             env.kafkaClient().prepareResponse(body -> body instanceof DeleteTopicsRequest,
-                    new DeleteTopicsResponse(Collections.singletonMap("myTopic", Errors.NONE)));
-            KafkaFuture future = env.adminClient().deleteTopics(Collections.singletonList("myTopic"),
+                    prepareDeleteTopicsResponse("myTopic", Errors.NONE));
+            KafkaFuture future = env.adminClient().deleteTopics(singletonList("myTopic"),
                     new DeleteTopicsOptions()).all();
             future.get();
 
             env.kafkaClient().prepareResponse(body -> body instanceof DeleteTopicsRequest,
-                    new DeleteTopicsResponse(Collections.singletonMap("myTopic", Errors.TOPIC_DELETION_DISABLED)));
-            future = env.adminClient().deleteTopics(Collections.singletonList("myTopic"),
+                    prepareDeleteTopicsResponse("myTopic", Errors.TOPIC_DELETION_DISABLED));
+            future = env.adminClient().deleteTopics(singletonList("myTopic"),
                     new DeleteTopicsOptions()).all();
             TestUtils.assertFutureError(future, TopicDeletionDisabledException.class);
 
             env.kafkaClient().prepareResponse(body -> body instanceof DeleteTopicsRequest,
-                    new DeleteTopicsResponse(Collections.singletonMap("myTopic", Errors.UNKNOWN_TOPIC_OR_PARTITION)));
-            future = env.adminClient().deleteTopics(Collections.singletonList("myTopic"),
+                    prepareDeleteTopicsResponse("myTopic", Errors.UNKNOWN_TOPIC_OR_PARTITION));
+            future = env.adminClient().deleteTopics(singletonList("myTopic"),
                     new DeleteTopicsOptions()).all();
             TestUtils.assertFutureError(future, UnknownTopicOrPartitionException.class);
         }
@@ -457,7 +537,7 @@ public void testMetadataRetries() throws Exception {
             env.kafkaClient().prepareResponse(null, true);
 
             // The next one succeeds and gives us the controller id
-            env.kafkaClient().prepareResponse(new MetadataResponse(initializedCluster.nodes(),
+            env.kafkaClient().prepareResponse(MetadataResponse.prepareResponse(initializedCluster.nodes(),
                     initializedCluster.clusterResource().clusterId(),
                     initializedCluster.controller().id(),
                     Collections.emptyList()));
@@ -467,14 +547,15 @@ public void testMetadataRetries() throws Exception {
             MetadataResponse.PartitionMetadata partitionMetadata = new MetadataResponse.PartitionMetadata(
                     Errors.NONE, 0, leader, Optional.of(10), singletonList(leader),
                     singletonList(leader), singletonList(leader));
-            env.kafkaClient().prepareResponse(new MetadataResponse(initializedCluster.nodes(),
+            env.kafkaClient().prepareResponse(MetadataResponse.prepareResponse(initializedCluster.nodes(),
                     initializedCluster.clusterResource().clusterId(), 1,
                     singletonList(new MetadataResponse.TopicMetadata(Errors.NONE, topic, false,
-                            singletonList(partitionMetadata)))));
+                            singletonList(partitionMetadata), MetadataResponse.AUTHORIZED_OPERATIONS_OMITTED))));
 
             DescribeTopicsResult result = env.adminClient().describeTopics(Collections.singleton(topic));
             Map topicDescriptions = result.all().get();
             assertEquals(leader, topicDescriptions.get(topic).partitions().get(0).leader());
+            assertEquals(null, topicDescriptions.get(topic).authorizedOperations());
         }
     }
 
@@ -650,51 +731,59 @@ public void testDeleteAcls() throws Exception {
     }
 
     @Test
-    public void testElectPreferredLeaders()  throws Exception {
+    public void testElectLeaders()  throws Exception {
         TopicPartition topic1 = new TopicPartition("topic", 0);
         TopicPartition topic2 = new TopicPartition("topic", 2);
         try (AdminClientUnitTestEnv env = mockClientEnv()) {
-            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
-
-            // Test a call where one partition has an error.
-            ApiError value = ApiError.fromThrowable(new ClusterAuthorizationException(null));
-            ElectPreferredLeadersResponseData responseData = new ElectPreferredLeadersResponseData();
-            ReplicaElectionResult r = new ReplicaElectionResult().setTopic(topic1.topic());
-            r.partitionResult().add(new PartitionResult()
-                    .setPartitionId(topic1.partition())
-                    .setErrorCode(ApiError.NONE.error().code())
-                    .setErrorMessage(ApiError.NONE.message()));
-            r.partitionResult().add(new PartitionResult()
-                    .setPartitionId(topic2.partition())
-                    .setErrorCode(value.error().code())
-                    .setErrorMessage(value.message()));
-            responseData.replicaElectionResults().add(r);
-            env.kafkaClient().prepareResponse(new ElectPreferredLeadersResponse(responseData));
-            ElectPreferredLeadersResult results = env.adminClient().electPreferredLeaders(asList(topic1, topic2));
-            results.partitionResult(topic1).get();
-            TestUtils.assertFutureError(results.partitionResult(topic2), ClusterAuthorizationException.class);
-            TestUtils.assertFutureError(results.all(), ClusterAuthorizationException.class);
-
-            // Test a call where there are no errors.
-            r.partitionResult().clear();
-            r.partitionResult().add(new PartitionResult()
-                    .setPartitionId(topic1.partition())
-                    .setErrorCode(ApiError.NONE.error().code())
-                    .setErrorMessage(ApiError.NONE.message()));
-            r.partitionResult().add(new PartitionResult()
-                    .setPartitionId(topic2.partition())
-                    .setErrorCode(ApiError.NONE.error().code())
-                    .setErrorMessage(ApiError.NONE.message()));
-            env.kafkaClient().prepareResponse(new ElectPreferredLeadersResponse(responseData));
-
-            results = env.adminClient().electPreferredLeaders(asList(topic1, topic2));
-            results.partitionResult(topic1).get();
-            results.partitionResult(topic2).get();
-
-            // Now try a timeout
-            results = env.adminClient().electPreferredLeaders(asList(topic1, topic2), new ElectPreferredLeadersOptions().timeoutMs(100));
-            TestUtils.assertFutureError(results.partitionResult(topic1), TimeoutException.class);
-            TestUtils.assertFutureError(results.partitionResult(topic2), TimeoutException.class);
+            for (ElectionType electionType : ElectionType.values()) {
+                env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+                // Test a call where one partition has an error.
+                ApiError value = ApiError.fromThrowable(new ClusterAuthorizationException(null));
+                List electionResults = new ArrayList<>();
+                ReplicaElectionResult electionResult = new ReplicaElectionResult();
+                electionResult.setTopic(topic1.topic());
+                // Add partition 1 result
+                PartitionResult partition1Result = new PartitionResult();
+                partition1Result.setPartitionId(topic1.partition());
+                partition1Result.setErrorCode(value.error().code());
+                partition1Result.setErrorMessage(value.message());
+                electionResult.partitionResult().add(partition1Result);
+
+                // Add partition 2 result
+                PartitionResult partition2Result = new PartitionResult();
+                partition2Result.setPartitionId(topic2.partition());
+                partition2Result.setErrorCode(value.error().code());
+                partition2Result.setErrorMessage(value.message());
+                electionResult.partitionResult().add(partition2Result);
+
+                electionResults.add(electionResult);
+
+                env.kafkaClient().prepareResponse(new ElectLeadersResponse(0, Errors.NONE.code(), electionResults));
+                ElectLeadersResult results = env.adminClient().electLeaders(
+                        electionType,
+                        new HashSet<>(asList(topic1, topic2)));
+                assertEquals(results.partitions().get().get(topic2).get().getClass(), ClusterAuthorizationException.class);
+
+                // Test a call where there are no errors. By mutating the internal of election results
+                partition1Result.setErrorCode(ApiError.NONE.error().code());
+                partition1Result.setErrorMessage(ApiError.NONE.message());
+
+                partition2Result.setErrorCode(ApiError.NONE.error().code());
+                partition2Result.setErrorMessage(ApiError.NONE.message());
+
+                env.kafkaClient().prepareResponse(new ElectLeadersResponse(0, Errors.NONE.code(), electionResults));
+                results = env.adminClient().electLeaders(electionType, new HashSet<>(asList(topic1, topic2)));
+                assertFalse(results.partitions().get().get(topic1).isPresent());
+                assertFalse(results.partitions().get().get(topic2).isPresent());
+
+                // Now try a timeout
+                results = env.adminClient().electLeaders(
+                        electionType,
+                        new HashSet<>(asList(topic1, topic2)),
+                        new ElectLeadersOptions().timeoutMs(100));
+                TestUtils.assertFutureError(results.partitions(), TimeoutException.class);
+            }
         }
     }
 
@@ -795,6 +884,69 @@ public void testCreatePartitions() throws Exception {
         }
     }
 
+    @Test
+    public void testDeleteRecordsTopicAuthorizationError() {
+        String topic = "foo";
+        TopicPartition partition = new TopicPartition(topic, 0);
+
+        try (AdminClientUnitTestEnv env = mockClientEnv()) {
+            List topics = new ArrayList<>();
+            topics.add(new MetadataResponse.TopicMetadata(Errors.TOPIC_AUTHORIZATION_FAILED, topic, false,
+                    Collections.emptyList()));
+
+            env.kafkaClient().prepareResponse(MetadataResponse.prepareResponse(env.cluster().nodes(),
+                    env.cluster().clusterResource().clusterId(), env.cluster().controller().id(), topics));
+
+            Map recordsToDelete = new HashMap<>();
+            recordsToDelete.put(partition, RecordsToDelete.beforeOffset(10L));
+            DeleteRecordsResult results = env.adminClient().deleteRecords(recordsToDelete);
+
+            TestUtils.assertFutureThrows(results.lowWatermarks().get(partition), TopicAuthorizationException.class);
+        }
+    }
+
+    @Test
+    public void testDeleteRecordsMultipleSends() throws Exception {
+        String topic = "foo";
+        TopicPartition tp0 = new TopicPartition(topic, 0);
+        TopicPartition tp1 = new TopicPartition(topic, 1);
+
+        Cluster cluster = mockCluster(0);
+        MockTime time = new MockTime();
+
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(time, cluster)) {
+            List nodes = cluster.nodes();
+
+            List partitionMetadata = new ArrayList<>();
+            partitionMetadata.add(new MetadataResponse.PartitionMetadata(Errors.NONE, tp0.partition(), nodes.get(0),
+                    Optional.of(5), singletonList(nodes.get(0)), singletonList(nodes.get(0)),
+                    Collections.emptyList()));
+            partitionMetadata.add(new MetadataResponse.PartitionMetadata(Errors.NONE, tp1.partition(), nodes.get(1),
+                    Optional.of(5), singletonList(nodes.get(1)), singletonList(nodes.get(1)), Collections.emptyList()));
+
+            List topicMetadata = new ArrayList<>();
+            topicMetadata.add(new MetadataResponse.TopicMetadata(Errors.NONE, topic, false, partitionMetadata));
+
+            env.kafkaClient().prepareResponse(MetadataResponse.prepareResponse(cluster.nodes(),
+                    cluster.clusterResource().clusterId(), cluster.controller().id(), topicMetadata));
+
+            Map deletedPartitions = new HashMap<>();
+            deletedPartitions.put(tp0, new DeleteRecordsResponse.PartitionResponse(3, Errors.NONE));
+            env.kafkaClient().prepareResponseFrom(new DeleteRecordsResponse(0, deletedPartitions), nodes.get(0));
+
+            env.kafkaClient().disconnect(nodes.get(1).idString());
+            env.kafkaClient().createPendingAuthenticationError(nodes.get(1), 100);
+
+            Map recordsToDelete = new HashMap<>();
+            recordsToDelete.put(tp0, RecordsToDelete.beforeOffset(10L));
+            recordsToDelete.put(tp1, RecordsToDelete.beforeOffset(10L));
+            DeleteRecordsResult results = env.adminClient().deleteRecords(recordsToDelete);
+
+            assertEquals(3L, results.lowWatermarks().get(tp0).get().lowWatermark());
+            TestUtils.assertFutureThrows(results.lowWatermarks().get(tp1), AuthenticationException.class);
+        }
+    }
+
     @Test
     public void testDeleteRecords() throws Exception {
 
@@ -845,7 +997,7 @@ public void testDeleteRecords() throws Exception {
 
             t.add(new MetadataResponse.TopicMetadata(Errors.NONE, "my_topic", false, p));
 
-            env.kafkaClient().prepareResponse(new MetadataResponse(cluster.nodes(), cluster.clusterResource().clusterId(), cluster.controller().id(), t));
+            env.kafkaClient().prepareResponse(MetadataResponse.prepareResponse(cluster.nodes(), cluster.clusterResource().clusterId(), cluster.controller().id(), t));
             env.kafkaClient().prepareResponse(new DeleteRecordsResponse(0, m));
 
             Map recordsToDelete = new HashMap<>();
@@ -901,6 +1053,61 @@ public void testDeleteRecords() throws Exception {
         }
     }
 
+    @Test
+    public void testDescribeCluster() throws Exception {
+        final HashMap nodes = new HashMap<>();
+        Node node0 = new Node(0, "localhost", 8121);
+        Node node1 = new Node(1, "localhost", 8122);
+        Node node2 = new Node(2, "localhost", 8123);
+        Node node3 = new Node(3, "localhost", 8124);
+        nodes.put(0, node0);
+        nodes.put(1, node1);
+        nodes.put(2, node2);
+        nodes.put(3, node3);
+
+        final Cluster cluster = new Cluster(
+                "mockClusterId",
+                nodes.values(),
+                Collections.emptyList(),
+                Collections.emptySet(),
+                Collections.emptySet(), nodes.get(0));
+
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(cluster, AdminClientConfig.RETRIES_CONFIG, "2")) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            // Prepare the metadata response used for the first describe cluster
+            MetadataResponse response = MetadataResponse.prepareResponse(0,
+                    new ArrayList<>(nodes.values()),
+                    env.cluster().clusterResource().clusterId(),
+                    2,
+                    Collections.emptyList(),
+                    MetadataResponse.AUTHORIZED_OPERATIONS_OMITTED);
+            env.kafkaClient().prepareResponse(response);
+
+            // Prepare the metadata response used for the second describe cluster
+            MetadataResponse response2 = MetadataResponse.prepareResponse(0,
+                    new ArrayList<>(nodes.values()),
+                    env.cluster().clusterResource().clusterId(),
+                    3,
+                    Collections.emptyList(),
+                    1 << AclOperation.DESCRIBE.code() | 1 << AclOperation.ALTER.code());
+            env.kafkaClient().prepareResponse(response2);
+
+            // Test DescribeCluster with the authorized operations omitted.
+            final DescribeClusterResult result = env.adminClient().describeCluster();
+            assertEquals(env.cluster().clusterResource().clusterId(), result.clusterId().get());
+            assertEquals(2, result.controller().get().id());
+            assertEquals(null, result.authorizedOperations().get());
+
+            // Test DescribeCluster with the authorized operations included.
+            final DescribeClusterResult result2 = env.adminClient().describeCluster();
+            assertEquals(env.cluster().clusterResource().clusterId(), result2.clusterId().get());
+            assertEquals(3, result2.controller().get().id());
+            assertEquals(new HashSet<>(Arrays.asList(AclOperation.DESCRIBE, AclOperation.ALTER)),
+                    result2.authorizedOperations().get());
+        }
+    }
+
     @Test
     public void testListConsumerGroups() throws Exception {
         final HashMap nodes = new HashMap<>();
@@ -920,19 +1127,19 @@ public void testListConsumerGroups() throws Exception {
                 Collections.emptySet(),
                 Collections.emptySet(), nodes.get(0));
 
-        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(cluster)) {
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(cluster, AdminClientConfig.RETRIES_CONFIG, "2")) {
             env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
 
             // Empty metadata response should be retried
             env.kafkaClient().prepareResponse(
-                    new MetadataResponse(
+                     MetadataResponse.prepareResponse(
                             Collections.emptyList(),
                             env.cluster().clusterResource().clusterId(),
                             -1,
                             Collections.emptyList()));
 
             env.kafkaClient().prepareResponse(
-                    new MetadataResponse(
+                     MetadataResponse.prepareResponse(
                             env.cluster().nodes(),
                             env.cluster().clusterResource().clusterId(),
                             env.cluster().controller().id(),
@@ -940,52 +1147,69 @@ public void testListConsumerGroups() throws Exception {
 
             env.kafkaClient().prepareResponseFrom(
                     new ListGroupsResponse(
-                            Errors.NONE,
-                            asList(
-                                    new ListGroupsResponse.Group("group-1", ConsumerProtocol.PROTOCOL_TYPE),
-                                    new ListGroupsResponse.Group("group-connect-1", "connector")
-                            )),
+                            new ListGroupsResponseData()
+                            .setErrorCode(Errors.NONE.code())
+                            .setGroups(Arrays.asList(
+                                    new ListGroupsResponseData.ListedGroup()
+                                            .setGroupId("group-1")
+                                            .setProtocolType(ConsumerProtocol.PROTOCOL_TYPE),
+                                    new ListGroupsResponseData.ListedGroup()
+                                            .setGroupId("group-connect-1")
+                                            .setProtocolType("connector")
+                            ))),
                     node0);
 
             // handle retriable errors
             env.kafkaClient().prepareResponseFrom(
                     new ListGroupsResponse(
-                            Errors.COORDINATOR_NOT_AVAILABLE,
-                            Collections.emptyList()
+                            new ListGroupsResponseData()
+                                    .setErrorCode(Errors.COORDINATOR_NOT_AVAILABLE.code())
+                                    .setGroups(Collections.emptyList())
                     ),
                     node1);
             env.kafkaClient().prepareResponseFrom(
                     new ListGroupsResponse(
-                            Errors.COORDINATOR_LOAD_IN_PROGRESS,
-                            Collections.emptyList()
+                            new ListGroupsResponseData()
+                                    .setErrorCode(Errors.COORDINATOR_LOAD_IN_PROGRESS.code())
+                                    .setGroups(Collections.emptyList())
                     ),
                     node1);
             env.kafkaClient().prepareResponseFrom(
                     new ListGroupsResponse(
-                            Errors.NONE,
-                            asList(
-                                    new ListGroupsResponse.Group("group-2", ConsumerProtocol.PROTOCOL_TYPE),
-                                    new ListGroupsResponse.Group("group-connect-2", "connector")
-                            )),
+                            new ListGroupsResponseData()
+                                    .setErrorCode(Errors.NONE.code())
+                                    .setGroups(Arrays.asList(
+                                            new ListGroupsResponseData.ListedGroup()
+                                                    .setGroupId("group-2")
+                                                    .setProtocolType(ConsumerProtocol.PROTOCOL_TYPE),
+                                            new ListGroupsResponseData.ListedGroup()
+                                                    .setGroupId("group-connect-2")
+                                                    .setProtocolType("connector")
+                            ))),
                     node1);
 
             env.kafkaClient().prepareResponseFrom(
                     new ListGroupsResponse(
-                            Errors.NONE,
-                            asList(
-                                    new ListGroupsResponse.Group("group-3", ConsumerProtocol.PROTOCOL_TYPE),
-                                    new ListGroupsResponse.Group("group-connect-3", "connector")
-                            )),
+                            new ListGroupsResponseData()
+                                    .setErrorCode(Errors.NONE.code())
+                                    .setGroups(Arrays.asList(
+                                            new ListGroupsResponseData.ListedGroup()
+                                                    .setGroupId("group-3")
+                                                    .setProtocolType(ConsumerProtocol.PROTOCOL_TYPE),
+                                            new ListGroupsResponseData.ListedGroup()
+                                                    .setGroupId("group-connect-3")
+                                                    .setProtocolType("connector")
+                                    ))),
                     node2);
 
             // fatal error
             env.kafkaClient().prepareResponseFrom(
                     new ListGroupsResponse(
-                            Errors.UNKNOWN_SERVER_ERROR,
-                            Collections.emptyList()),
+                            new ListGroupsResponseData()
+                                    .setErrorCode(Errors.UNKNOWN_SERVER_ERROR.code())
+                                    .setGroups(Collections.emptyList())),
                     node3);
 
-
             final ListConsumerGroupsResult result = env.adminClient().listConsumerGroups();
             TestUtils.assertFutureError(result.all(), UnknownServerException.class);
 
@@ -1027,7 +1251,7 @@ public void testListConsumerGroupsMetadataFailure() throws Exception {
             // Empty metadata causes the request to fail since we have no list of brokers
             // to send the ListGroups requests to
             env.kafkaClient().prepareResponse(
-                    new MetadataResponse(
+                     MetadataResponse.prepareResponse(
                             Collections.emptyList(),
                             env.cluster().clusterResource().clusterId(),
                             -1,
@@ -1054,9 +1278,54 @@ public void testDescribeConsumerGroups() throws Exception {
         try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(cluster)) {
             env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
 
-            env.kafkaClient().prepareResponse(new FindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
+            //Retriable FindCoordinatorResponse errors should be retried
+            env.kafkaClient().prepareResponse(prepareFindCoordinatorResponse(Errors.COORDINATOR_NOT_AVAILABLE,  Node.noNode()));
+            env.kafkaClient().prepareResponse(prepareFindCoordinatorResponse(Errors.COORDINATOR_LOAD_IN_PROGRESS,  Node.noNode()));
+
+            env.kafkaClient().prepareResponse(prepareFindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
 
             DescribeGroupsResponseData data = new DescribeGroupsResponseData();
+
+            //Retriable errors should be retried
+            data.groups().add(DescribeGroupsResponse.groupMetadata(
+                "group-0",
+                Errors.COORDINATOR_LOAD_IN_PROGRESS,
+                "",
+                "",
+                "",
+                Collections.emptyList(),
+                Collections.emptySet()));
+            env.kafkaClient().prepareResponse(new DescribeGroupsResponse(data));
+
+            data = new DescribeGroupsResponseData();
+            data.groups().add(DescribeGroupsResponse.groupMetadata(
+                "group-0",
+                Errors.COORDINATOR_NOT_AVAILABLE,
+                "",
+                "",
+                "",
+                Collections.emptyList(),
+                Collections.emptySet()));
+            env.kafkaClient().prepareResponse(new DescribeGroupsResponse(data));
+
+            /*
+             * We need to return two responses here, one with NOT_COORDINATOR error when calling describe consumer group
+             * api using coordinator that has moved. This will retry whole operation. So we need to again respond with a
+             * FindCoordinatorResponse.
+             */
+            data = new DescribeGroupsResponseData();
+            data.groups().add(DescribeGroupsResponse.groupMetadata(
+                    "group-0",
+                    Errors.NOT_COORDINATOR,
+                    "",
+                    "",
+                    "",
+                    Collections.emptyList(),
+                    Collections.emptySet()));
+            env.kafkaClient().prepareResponse(new DescribeGroupsResponse(data));
+            env.kafkaClient().prepareResponse(prepareFindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
+
+            data = new DescribeGroupsResponseData();
             TopicPartition myTopicPartition0 = new TopicPartition("my_topic", 0);
             TopicPartition myTopicPartition1 = new TopicPartition("my_topic", 1);
             TopicPartition myTopicPartition2 = new TopicPartition("my_topic", 2);
@@ -1066,42 +1335,179 @@ public void testDescribeConsumerGroups() throws Exception {
             topicPartitions.add(1, myTopicPartition1);
             topicPartitions.add(2, myTopicPartition2);
 
-            final ByteBuffer memberAssignment = ConsumerProtocol.serializeAssignment(new PartitionAssignor.Assignment(topicPartitions));
+            final ByteBuffer memberAssignment = ConsumerProtocol.serializeAssignment(new ConsumerPartitionAssignor.Assignment(topicPartitions));
             byte[] memberAssignmentBytes = new byte[memberAssignment.remaining()];
             memberAssignment.get(memberAssignmentBytes);
 
+            DescribedGroupMember memberOne = DescribeGroupsResponse.groupMember("0", "instance1", "clientId0", "clientHost", memberAssignmentBytes, null);
+            DescribedGroupMember memberTwo = DescribeGroupsResponse.groupMember("1", "instance2", "clientId1", "clientHost", memberAssignmentBytes, null);
+
+            List expectedMemberDescriptions = new ArrayList<>();
+            expectedMemberDescriptions.add(convertToMemberDescriptions(memberOne,
+                                                                       new MemberAssignment(new HashSet<>(topicPartitions))));
+            expectedMemberDescriptions.add(convertToMemberDescriptions(memberTwo,
+                                                                       new MemberAssignment(new HashSet<>(topicPartitions))));
             data.groups().add(DescribeGroupsResponse.groupMetadata(
+                    "group-0",
+                    Errors.NONE,
+                    "",
+                    ConsumerProtocol.PROTOCOL_TYPE,
+                    "",
+                    asList(memberOne, memberTwo),
+                    Collections.emptySet()));
+
+            env.kafkaClient().prepareResponse(new DescribeGroupsResponse(data));
+
+            final DescribeConsumerGroupsResult result = env.adminClient().describeConsumerGroups(singletonList("group-0"));
+            final ConsumerGroupDescription groupDescription = result.describedGroups().get("group-0").get();
+
+            assertEquals(1, result.describedGroups().size());
+            assertEquals("group-0", groupDescription.groupId());
+            assertEquals(2, groupDescription.members().size());
+            assertEquals(expectedMemberDescriptions, groupDescription.members());
+        }
+    }
+
+    @Test
+    public void testDescribeMultipleConsumerGroups() throws Exception {
+        final HashMap nodes = new HashMap<>();
+        nodes.put(0, new Node(0, "localhost", 8121));
+
+        final Cluster cluster =
+                new Cluster(
+                        "mockClusterId",
+                        nodes.values(),
+                        Collections.emptyList(),
+                        Collections.emptySet(),
+                        Collections.emptySet(), nodes.get(0));
+
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(cluster)) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            env.kafkaClient().prepareResponse(prepareFindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
+
+            TopicPartition myTopicPartition0 = new TopicPartition("my_topic", 0);
+            TopicPartition myTopicPartition1 = new TopicPartition("my_topic", 1);
+            TopicPartition myTopicPartition2 = new TopicPartition("my_topic", 2);
+
+            final List topicPartitions = new ArrayList<>();
+            topicPartitions.add(0, myTopicPartition0);
+            topicPartitions.add(1, myTopicPartition1);
+            topicPartitions.add(2, myTopicPartition2);
+
+            final ByteBuffer memberAssignment = ConsumerProtocol.serializeAssignment(new ConsumerPartitionAssignor.Assignment(topicPartitions));
+            byte[] memberAssignmentBytes = new byte[memberAssignment.remaining()];
+            memberAssignment.get(memberAssignmentBytes);
+
+            DescribeGroupsResponseData group0Data = new DescribeGroupsResponseData();
+            group0Data.groups().add(DescribeGroupsResponse.groupMetadata(
                     "group-0",
                     Errors.NONE,
                     "",
                     ConsumerProtocol.PROTOCOL_TYPE,
                     "",
                     asList(
-                        DescribeGroupsResponse.groupMember("0", "clientId0", "clientHost", memberAssignmentBytes, null),
-                        DescribeGroupsResponse.groupMember("1", "clientId1", "clientHost", memberAssignmentBytes, null)
+                            DescribeGroupsResponse.groupMember("0", null, "clientId0", "clientHost", memberAssignmentBytes, null),
+                            DescribeGroupsResponse.groupMember("1", null, "clientId1", "clientHost", memberAssignmentBytes, null)
                     ),
                     Collections.emptySet()));
 
-            data.groups().add(DescribeGroupsResponse.groupMetadata(
+            DescribeGroupsResponseData groupConnectData = new DescribeGroupsResponseData();
+            group0Data.groups().add(DescribeGroupsResponse.groupMetadata(
                     "group-connect-0",
                     Errors.NONE,
                     "",
                     "connect",
                     "",
                     asList(
-                        DescribeGroupsResponse.groupMember("0", "clientId0", "clientHost", memberAssignmentBytes, null),
-                        DescribeGroupsResponse.groupMember("1", "clientId1", "clientHost", memberAssignmentBytes, null)
+                            DescribeGroupsResponse.groupMember("0", null, "clientId0", "clientHost", memberAssignmentBytes, null),
+                            DescribeGroupsResponse.groupMember("1", null, "clientId1", "clientHost", memberAssignmentBytes, null)
                     ),
                     Collections.emptySet()));
 
+            env.kafkaClient().prepareResponse(new DescribeGroupsResponse(group0Data));
+            env.kafkaClient().prepareResponse(new DescribeGroupsResponse(groupConnectData));
+
+            Collection groups = new HashSet<>();
+            groups.add("group-0");
+            groups.add("group-connect-0");
+            final DescribeConsumerGroupsResult result = env.adminClient().describeConsumerGroups(groups);
+            assertEquals(2, result.describedGroups().size());
+            assertEquals(groups, result.describedGroups().keySet());
+        }
+    }
+
+    @Test
+    public void testDescribeConsumerGroupsWithAuthorizedOperationsOmitted() throws Exception {
+        final HashMap nodes = new HashMap<>();
+        nodes.put(0, new Node(0, "localhost", 8121));
+
+        final Cluster cluster =
+            new Cluster(
+                "mockClusterId",
+                nodes.values(),
+                Collections.emptyList(),
+                Collections.emptySet(),
+                Collections.emptySet(), nodes.get(0));
+
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(cluster)) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            env.kafkaClient().prepareResponse(FindCoordinatorResponse.prepareResponse(Errors.NONE, env.cluster().controller()));
+
+            DescribeGroupsResponseData data = new DescribeGroupsResponseData();
+            data.groups().add(DescribeGroupsResponse.groupMetadata(
+                "group-0",
+                Errors.NONE,
+                "",
+                ConsumerProtocol.PROTOCOL_TYPE,
+                "",
+                Collections.emptyList(),
+                MetadataResponse.AUTHORIZED_OPERATIONS_OMITTED));
+
             env.kafkaClient().prepareResponse(new DescribeGroupsResponse(data));
 
             final DescribeConsumerGroupsResult result = env.adminClient().describeConsumerGroups(singletonList("group-0"));
             final ConsumerGroupDescription groupDescription = result.describedGroups().get("group-0").get();
 
-            assertEquals(1, result.describedGroups().size());
-            assertEquals("group-0", groupDescription.groupId());
-            assertEquals(2, groupDescription.members().size());
+            assertNull(groupDescription.authorizedOperations());
+        }
+    }
+
+    @Test
+    public void testDescribeNonConsumerGroups() throws Exception {
+        final HashMap nodes = new HashMap<>();
+        nodes.put(0, new Node(0, "localhost", 8121));
+
+        final Cluster cluster =
+            new Cluster(
+                "mockClusterId",
+                nodes.values(),
+                Collections.emptyList(),
+                Collections.emptySet(),
+                Collections.emptySet(), nodes.get(0));
+
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(cluster)) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            env.kafkaClient().prepareResponse(prepareFindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
+
+            DescribeGroupsResponseData data = new DescribeGroupsResponseData();
+
+            data.groups().add(DescribeGroupsResponse.groupMetadata(
+                "group-0",
+                Errors.NONE,
+                "",
+                "non-consumer",
+                "",
+                asList(),
+                Collections.emptySet()));
+
+            env.kafkaClient().prepareResponse(new DescribeGroupsResponse(data));
+
+            final DescribeConsumerGroupsResult result = env.adminClient().describeConsumerGroups(singletonList("group-0"));
+
+            TestUtils.assertFutureError(result.describedGroups().get("group-0"), IllegalArgumentException.class);
         }
     }
 
@@ -1121,11 +1527,27 @@ public void testDescribeConsumerGroupOffsets() throws Exception {
         try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(cluster)) {
             env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
 
-            env.kafkaClient().prepareResponse(new FindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
+            // Retriable FindCoordinatorResponse errors should be retried
+            env.kafkaClient().prepareResponse(prepareFindCoordinatorResponse(Errors.COORDINATOR_NOT_AVAILABLE,  Node.noNode()));
+
+            env.kafkaClient().prepareResponse(prepareFindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
+
+            // Retriable errors should be retried
+            env.kafkaClient().prepareResponse(new OffsetFetchResponse(Errors.COORDINATOR_NOT_AVAILABLE, Collections.emptyMap()));
+            env.kafkaClient().prepareResponse(new OffsetFetchResponse(Errors.COORDINATOR_LOAD_IN_PROGRESS, Collections.emptyMap()));
+
+            /*
+             * We need to return two responses here, one for NOT_COORDINATOR error when calling list consumer group offsets
+             * api using coordinator that has moved. This will retry whole operation. So we need to again respond with a
+             * FindCoordinatorResponse.
+             */
+            env.kafkaClient().prepareResponse(new OffsetFetchResponse(Errors.NOT_COORDINATOR, Collections.emptyMap()));
+            env.kafkaClient().prepareResponse(prepareFindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
 
             TopicPartition myTopicPartition0 = new TopicPartition("my_topic", 0);
             TopicPartition myTopicPartition1 = new TopicPartition("my_topic", 1);
             TopicPartition myTopicPartition2 = new TopicPartition("my_topic", 2);
+            TopicPartition myTopicPartition3 = new TopicPartition("my_topic", 3);
 
             final Map responseData = new HashMap<>();
             responseData.put(myTopicPartition0, new OffsetFetchResponse.PartitionData(10,
@@ -1134,21 +1556,25 @@ public void testDescribeConsumerGroupOffsets() throws Exception {
                     Optional.empty(), "", Errors.NONE));
             responseData.put(myTopicPartition2, new OffsetFetchResponse.PartitionData(20,
                     Optional.empty(), "", Errors.NONE));
+            responseData.put(myTopicPartition3, new OffsetFetchResponse.PartitionData(OffsetFetchResponse.INVALID_OFFSET,
+                    Optional.empty(), "", Errors.NONE));
             env.kafkaClient().prepareResponse(new OffsetFetchResponse(Errors.NONE, responseData));
 
             final ListConsumerGroupOffsetsResult result = env.adminClient().listConsumerGroupOffsets("group-0");
             final Map partitionToOffsetAndMetadata = result.partitionsToOffsetAndMetadata().get();
 
-            assertEquals(3, partitionToOffsetAndMetadata.size());
+            assertEquals(4, partitionToOffsetAndMetadata.size());
             assertEquals(10, partitionToOffsetAndMetadata.get(myTopicPartition0).offset());
             assertEquals(0, partitionToOffsetAndMetadata.get(myTopicPartition1).offset());
             assertEquals(20, partitionToOffsetAndMetadata.get(myTopicPartition2).offset());
+            assertTrue(partitionToOffsetAndMetadata.containsKey(myTopicPartition3));
+            assertNull(partitionToOffsetAndMetadata.get(myTopicPartition3));
         }
     }
 
     @Test
     public void testDeleteConsumerGroups() throws Exception {
-        final HashMap nodes = new HashMap<>();
+        final Map nodes = new HashMap<>();
         nodes.put(0, new Node(0, "localhost", 8121));
 
         final Cluster cluster =
@@ -1165,14 +1591,19 @@ public void testDeleteConsumerGroups() throws Exception {
             env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
 
             //Retriable FindCoordinatorResponse errors should be retried
-            env.kafkaClient().prepareResponse(new FindCoordinatorResponse(Errors.COORDINATOR_NOT_AVAILABLE,  Node.noNode()));
-            env.kafkaClient().prepareResponse(new FindCoordinatorResponse(Errors.COORDINATOR_LOAD_IN_PROGRESS, Node.noNode()));
-
-            env.kafkaClient().prepareResponse(new FindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
-
-            final Map response = new HashMap<>();
-            response.put("group-0", Errors.NONE);
-            env.kafkaClient().prepareResponse(new DeleteGroupsResponse(response));
+            env.kafkaClient().prepareResponse(prepareFindCoordinatorResponse(Errors.COORDINATOR_NOT_AVAILABLE,  Node.noNode()));
+            env.kafkaClient().prepareResponse(prepareFindCoordinatorResponse(Errors.COORDINATOR_LOAD_IN_PROGRESS, Node.noNode()));
+
+            env.kafkaClient().prepareResponse(prepareFindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
+
+            final DeletableGroupResultCollection validResponse = new DeletableGroupResultCollection();
+            validResponse.add(new DeletableGroupResult()
+                                  .setGroupId("group-0")
+                                  .setErrorCode(Errors.NONE.code()));
+            env.kafkaClient().prepareResponse(new DeleteGroupsResponse(
+                new DeleteGroupsResponseData()
+                    .setResults(validResponse)
+            ));
 
             final DeleteConsumerGroupsResult result = env.adminClient().deleteConsumerGroups(groupIds);
 
@@ -1180,14 +1611,658 @@ public void testDeleteConsumerGroups() throws Exception {
             assertNull(results.get());
 
             //should throw error for non-retriable errors
-            env.kafkaClient().prepareResponse(new FindCoordinatorResponse(Errors.GROUP_AUTHORIZATION_FAILED,  Node.noNode()));
+            env.kafkaClient().prepareResponse(prepareFindCoordinatorResponse(Errors.GROUP_AUTHORIZATION_FAILED,  Node.noNode()));
 
             final DeleteConsumerGroupsResult errorResult = env.adminClient().deleteConsumerGroups(groupIds);
             TestUtils.assertFutureError(errorResult.deletedGroups().get("group-0"), GroupAuthorizationException.class);
 
+            //Retriable errors should be retried
+            env.kafkaClient().prepareResponse(FindCoordinatorResponse.prepareResponse(Errors.NONE, env.cluster().controller()));
+
+            final DeletableGroupResultCollection errorResponse1 = new DeletableGroupResultCollection();
+            errorResponse1.add(new DeletableGroupResult()
+                                   .setGroupId("group-0")
+                                   .setErrorCode(Errors.COORDINATOR_NOT_AVAILABLE.code())
+            );
+            env.kafkaClient().prepareResponse(new DeleteGroupsResponse(
+                new DeleteGroupsResponseData()
+                    .setResults(errorResponse1)));
+
+            final DeletableGroupResultCollection errorResponse2 = new DeletableGroupResultCollection();
+            errorResponse2.add(new DeletableGroupResult()
+                                   .setGroupId("group-0")
+                                   .setErrorCode(Errors.COORDINATOR_LOAD_IN_PROGRESS.code())
+            );
+            env.kafkaClient().prepareResponse(new DeleteGroupsResponse(
+                new DeleteGroupsResponseData()
+                    .setResults(errorResponse2)));
+
+            /*
+             * We need to return two responses here, one for NOT_COORDINATOR call when calling delete a consumer group
+             * api using coordinator that has moved. This will retry whole operation. So we need to again respond with a
+             * FindCoordinatorResponse.
+             */
+            final DeletableGroupResultCollection coordinatorMoved = new DeletableGroupResultCollection();
+            coordinatorMoved.add(new DeletableGroupResult()
+                                     .setGroupId("UnitTestError")
+                                     .setErrorCode(Errors.NOT_COORDINATOR.code())
+            );
+            env.kafkaClient().prepareResponse(new DeleteGroupsResponse(
+                new DeleteGroupsResponseData()
+                    .setResults(coordinatorMoved)));
+            env.kafkaClient().prepareResponse(prepareFindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
+
+            env.kafkaClient().prepareResponse(new DeleteGroupsResponse(
+                new DeleteGroupsResponseData()
+                    .setResults(validResponse)));
+
+            final DeleteConsumerGroupsResult errorResult1 = env.adminClient().deleteConsumerGroups(groupIds);
+
+            final KafkaFuture errorResults = errorResult1.deletedGroups().get("group-0");
+            assertNull(errorResults.get());
+        }
+    }
+
+    @Test
+    public void testDeleteConsumerGroupOffsets() throws Exception {
+        // Happy path
+
+        final Map nodes = new HashMap<>();
+        nodes.put(0, new Node(0, "localhost", 8121));
+
+        final Cluster cluster =
+            new Cluster(
+                "mockClusterId",
+                nodes.values(),
+                Collections.emptyList(),
+                Collections.emptySet(),
+                Collections.emptySet(), nodes.get(0));
+
+        final String groupId = "group-0";
+        final TopicPartition tp1 = new TopicPartition("foo", 0);
+        final TopicPartition tp2 = new TopicPartition("bar", 0);
+        final TopicPartition tp3 = new TopicPartition("foobar", 0);
+
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(cluster)) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            env.kafkaClient().prepareResponse(
+                prepareFindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
+
+            env.kafkaClient().prepareResponse(new OffsetDeleteResponse(
+                new OffsetDeleteResponseData()
+                    .setTopics(new OffsetDeleteResponseTopicCollection(Stream.of(
+                        new OffsetDeleteResponseTopic()
+                            .setName("foo")
+                            .setPartitions(new OffsetDeleteResponsePartitionCollection(Collections.singletonList(
+                                new OffsetDeleteResponsePartition()
+                                    .setPartitionIndex(0)
+                                    .setErrorCode(Errors.NONE.code())
+                            ).iterator())),
+                        new OffsetDeleteResponseTopic()
+                            .setName("bar")
+                            .setPartitions(new OffsetDeleteResponsePartitionCollection(Collections.singletonList(
+                                new OffsetDeleteResponsePartition()
+                                    .setPartitionIndex(0)
+                                    .setErrorCode(Errors.GROUP_SUBSCRIBED_TO_TOPIC.code())
+                            ).iterator()))
+                    ).collect(Collectors.toList()).iterator()))
+                )
+            );
+
+            final DeleteConsumerGroupOffsetsResult errorResult = env.adminClient().deleteConsumerGroupOffsets(
+                groupId, Stream.of(tp1, tp2).collect(Collectors.toSet()));
+
+            assertNull(errorResult.partitionResult(tp1).get());
+            TestUtils.assertFutureError(errorResult.all(), GroupSubscribedToTopicException.class);
+            TestUtils.assertFutureError(errorResult.partitionResult(tp2), GroupSubscribedToTopicException.class);
+            assertThrows(IllegalArgumentException.class, () -> errorResult.partitionResult(tp3));
+        }
+    }
+
+    @Test
+    public void testDeleteConsumerGroupOffsetsRetriableErrors() throws Exception {
+        // Retriable errors should be retried
+
+        final Map nodes = new HashMap<>();
+        nodes.put(0, new Node(0, "localhost", 8121));
+
+        final Cluster cluster =
+            new Cluster(
+                "mockClusterId",
+                nodes.values(),
+                Collections.emptyList(),
+                Collections.emptySet(),
+                Collections.emptySet(), nodes.get(0));
+
+        final String groupId = "group-0";
+        final TopicPartition tp1 = new TopicPartition("foo", 0);
+
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(cluster)) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            env.kafkaClient().prepareResponse(
+                FindCoordinatorResponse.prepareResponse(Errors.NONE, env.cluster().controller()));
+
+            env.kafkaClient().prepareResponse(
+                prepareOffsetDeleteResponse(Errors.COORDINATOR_NOT_AVAILABLE));
+
+            env.kafkaClient().prepareResponse(
+                prepareOffsetDeleteResponse(Errors.COORDINATOR_LOAD_IN_PROGRESS));
+
+            /*
+             * We need to return two responses here, one for NOT_COORDINATOR call when calling delete a consumer group
+             * api using coordinator that has moved. This will retry whole operation. So we need to again respond with a
+             * FindCoordinatorResponse.
+             */
+            env.kafkaClient().prepareResponse(
+                prepareOffsetDeleteResponse(Errors.NOT_COORDINATOR));
+
+            env.kafkaClient().prepareResponse(
+                prepareFindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
+
+            env.kafkaClient().prepareResponse(
+                prepareOffsetDeleteResponse("foo", 0, Errors.NONE));
+
+            final DeleteConsumerGroupOffsetsResult errorResult1 = env.adminClient()
+                .deleteConsumerGroupOffsets(groupId, Stream.of(tp1).collect(Collectors.toSet()));
+
+            assertNull(errorResult1.all().get());
+            assertNull(errorResult1.partitionResult(tp1).get());
         }
     }
 
+    @Test
+    public void testDeleteConsumerGroupOffsetsNonRetriableErrors() throws Exception {
+        // Non-retriable errors throw an exception
+
+        final Map nodes = new HashMap<>();
+        nodes.put(0, new Node(0, "localhost", 8121));
+
+        final Cluster cluster =
+            new Cluster(
+                "mockClusterId",
+                nodes.values(),
+                Collections.emptyList(),
+                Collections.emptySet(),
+                Collections.emptySet(), nodes.get(0));
+
+        final String groupId = "group-0";
+        final TopicPartition tp1 = new TopicPartition("foo", 0);
+        final List nonRetriableErrors = Arrays.asList(
+            Errors.GROUP_AUTHORIZATION_FAILED, Errors.INVALID_GROUP_ID, Errors.GROUP_ID_NOT_FOUND);
+
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(cluster)) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            for (Errors error : nonRetriableErrors) {
+                env.kafkaClient().prepareResponse(
+                    prepareFindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
+
+                env.kafkaClient().prepareResponse(
+                    prepareOffsetDeleteResponse(error));
+
+                DeleteConsumerGroupOffsetsResult errorResult = env.adminClient()
+                    .deleteConsumerGroupOffsets(groupId, Stream.of(tp1).collect(Collectors.toSet()));
+
+                TestUtils.assertFutureError(errorResult.all(), error.exception().getClass());
+                TestUtils.assertFutureError(errorResult.partitionResult(tp1), error.exception().getClass());
+            }
+        }
+    }
+
+    @Test
+    public void testDeleteConsumerGroupOffsetsFindCoordinatorRetriableErrors() throws Exception {
+        // Retriable FindCoordinatorResponse errors should be retried
+
+        final Map nodes = new HashMap<>();
+        nodes.put(0, new Node(0, "localhost", 8121));
+
+        final Cluster cluster =
+            new Cluster(
+                "mockClusterId",
+                nodes.values(),
+                Collections.emptyList(),
+                Collections.emptySet(),
+                Collections.emptySet(), nodes.get(0));
+
+        final String groupId = "group-0";
+        final TopicPartition tp1 = new TopicPartition("foo", 0);
+
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(cluster)) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            env.kafkaClient().prepareResponse(
+                prepareFindCoordinatorResponse(Errors.COORDINATOR_NOT_AVAILABLE,  Node.noNode()));
+            env.kafkaClient().prepareResponse(
+                prepareFindCoordinatorResponse(Errors.COORDINATOR_LOAD_IN_PROGRESS, Node.noNode()));
+
+            env.kafkaClient().prepareResponse(
+                prepareFindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
+
+            env.kafkaClient().prepareResponse(
+                prepareOffsetDeleteResponse("foo", 0, Errors.NONE));
+
+            final DeleteConsumerGroupOffsetsResult result = env.adminClient()
+                .deleteConsumerGroupOffsets(groupId, Stream.of(tp1).collect(Collectors.toSet()));
+
+            assertNull(result.all().get());
+            assertNull(result.partitionResult(tp1).get());
+        }
+    }
+
+    @Test
+    public void testDeleteConsumerGroupOffsetsFindCoordinatorNonRetriableErrors() throws Exception {
+        // Non-retriable FindCoordinatorResponse errors throw an exception
+
+        final Map nodes = new HashMap<>();
+        nodes.put(0, new Node(0, "localhost", 8121));
+
+        final Cluster cluster =
+            new Cluster(
+                "mockClusterId",
+                nodes.values(),
+                Collections.emptyList(),
+                Collections.emptySet(),
+                Collections.emptySet(), nodes.get(0));
+
+        final String groupId = "group-0";
+        final TopicPartition tp1 = new TopicPartition("foo", 0);
+
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(cluster)) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            env.kafkaClient().prepareResponse(
+                prepareFindCoordinatorResponse(Errors.GROUP_AUTHORIZATION_FAILED,  Node.noNode()));
+
+            final DeleteConsumerGroupOffsetsResult errorResult = env.adminClient()
+                .deleteConsumerGroupOffsets(groupId, Stream.of(tp1).collect(Collectors.toSet()));
+
+            TestUtils.assertFutureError(errorResult.all(), GroupAuthorizationException.class);
+            TestUtils.assertFutureError(errorResult.partitionResult(tp1), GroupAuthorizationException.class);
+        }
+    }
+
+    @Test
+    public void testIncrementalAlterConfigs()  throws Exception {
+        try (AdminClientUnitTestEnv env = mockClientEnv()) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            //test error scenarios
+            IncrementalAlterConfigsResponseData responseData =  new IncrementalAlterConfigsResponseData();
+            responseData.responses().add(new AlterConfigsResourceResponse()
+                    .setResourceName("")
+                    .setResourceType(ConfigResource.Type.BROKER.id())
+                    .setErrorCode(Errors.CLUSTER_AUTHORIZATION_FAILED.code())
+                    .setErrorMessage("authorization error"));
+
+            responseData.responses().add(new AlterConfigsResourceResponse()
+                    .setResourceName("topic1")
+                    .setResourceType(ConfigResource.Type.TOPIC.id())
+                    .setErrorCode(Errors.INVALID_REQUEST.code())
+                    .setErrorMessage("Config value append is not allowed for config"));
+
+            env.kafkaClient().prepareResponse(new IncrementalAlterConfigsResponse(responseData));
+
+            ConfigResource brokerResource = new ConfigResource(ConfigResource.Type.BROKER, "");
+            ConfigResource topicResource = new ConfigResource(ConfigResource.Type.TOPIC, "topic1");
+
+            AlterConfigOp alterConfigOp1 = new AlterConfigOp(
+                    new ConfigEntry("log.segment.bytes", "1073741"),
+                    AlterConfigOp.OpType.SET);
+
+            AlterConfigOp alterConfigOp2 = new AlterConfigOp(
+                    new ConfigEntry("compression.type", "gzip"),
+                    AlterConfigOp.OpType.APPEND);
+
+            final Map> configs = new HashMap<>();
+            configs.put(brokerResource, singletonList(alterConfigOp1));
+            configs.put(topicResource, singletonList(alterConfigOp2));
+
+            AlterConfigsResult result = env.adminClient().incrementalAlterConfigs(configs);
+            TestUtils.assertFutureError(result.values().get(brokerResource), ClusterAuthorizationException.class);
+            TestUtils.assertFutureError(result.values().get(topicResource), InvalidRequestException.class);
+
+            // Test a call where there are no errors.
+            responseData =  new IncrementalAlterConfigsResponseData();
+            responseData.responses().add(new AlterConfigsResourceResponse()
+                    .setResourceName("")
+                    .setResourceType(ConfigResource.Type.BROKER.id())
+                    .setErrorCode(Errors.NONE.code())
+                    .setErrorMessage(ApiError.NONE.message()));
+
+            env.kafkaClient().prepareResponse(new IncrementalAlterConfigsResponse(responseData));
+            env.adminClient().incrementalAlterConfigs(Collections.singletonMap(brokerResource, singletonList(alterConfigOp1))).all().get();
+        }
+    }
+
+    @Test
+    public void testRemoveMembersFromGroup() throws Exception {
+        try (AdminClientUnitTestEnv env = mockClientEnv()) {
+            final String instanceOne = "instance-1";
+            final String instanceTwo = "instance-2";
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            // Retriable FindCoordinatorResponse errors should be retried
+            env.kafkaClient().prepareResponse(prepareFindCoordinatorResponse(Errors.COORDINATOR_NOT_AVAILABLE,  Node.noNode()));
+            env.kafkaClient().prepareResponse(prepareFindCoordinatorResponse(Errors.COORDINATOR_LOAD_IN_PROGRESS,  Node.noNode()));
+            env.kafkaClient().prepareResponse(prepareFindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
+
+            // Retriable errors should be retried
+            env.kafkaClient().prepareResponse(null, true);
+            env.kafkaClient().prepareResponse(new LeaveGroupResponse(new LeaveGroupResponseData()
+                                                                         .setErrorCode(Errors.COORDINATOR_NOT_AVAILABLE.code())));
+            env.kafkaClient().prepareResponse(new LeaveGroupResponse(new LeaveGroupResponseData()
+                                                                         .setErrorCode(Errors.COORDINATOR_LOAD_IN_PROGRESS.code())));
+
+            // Inject a top-level non-retriable error
+            env.kafkaClient().prepareResponse(new LeaveGroupResponse(new LeaveGroupResponseData()
+                                                                         .setErrorCode(Errors.UNKNOWN_SERVER_ERROR.code())));
+
+            String groupId = "groupId";
+            Collection membersToRemove = Arrays.asList(new MemberToRemove(instanceOne),
+                                                                       new MemberToRemove(instanceTwo));
+            final RemoveMembersFromConsumerGroupResult unknownErrorResult = env.adminClient().removeMembersFromConsumerGroup(
+                groupId,
+                new RemoveMembersFromConsumerGroupOptions(membersToRemove)
+            );
+
+            MemberToRemove memberOne = new MemberToRemove(instanceOne);
+            MemberToRemove memberTwo = new MemberToRemove(instanceTwo);
+
+            TestUtils.assertFutureError(unknownErrorResult.all(), UnknownServerException.class);
+            TestUtils.assertFutureError(unknownErrorResult.memberResult(memberOne), UnknownServerException.class);
+            TestUtils.assertFutureError(unknownErrorResult.memberResult(memberTwo), UnknownServerException.class);
+
+            MemberResponse responseOne = new MemberResponse()
+                                             .setGroupInstanceId(instanceOne)
+                                             .setErrorCode(Errors.UNKNOWN_MEMBER_ID.code());
+
+            MemberResponse responseTwo = new MemberResponse()
+                                             .setGroupInstanceId(instanceTwo)
+                                             .setErrorCode(Errors.NONE.code());
+
+            // Inject one member level error.
+            env.kafkaClient().prepareResponse(prepareFindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
+            env.kafkaClient().prepareResponse(new LeaveGroupResponse(new LeaveGroupResponseData()
+                                                                         .setErrorCode(Errors.NONE.code())
+                                                                         .setMembers(Arrays.asList(responseOne, responseTwo))));
+
+            final RemoveMembersFromConsumerGroupResult memberLevelErrorResult = env.adminClient().removeMembersFromConsumerGroup(
+                groupId,
+                new RemoveMembersFromConsumerGroupOptions(membersToRemove)
+            );
+
+            TestUtils.assertFutureError(memberLevelErrorResult.all(), UnknownMemberIdException.class);
+            TestUtils.assertFutureError(memberLevelErrorResult.memberResult(memberOne), UnknownMemberIdException.class);
+            assertNull(memberLevelErrorResult.memberResult(memberTwo).get());
+
+            // Return with missing member.
+            env.kafkaClient().prepareResponse(prepareFindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
+            env.kafkaClient().prepareResponse(new LeaveGroupResponse(new LeaveGroupResponseData()
+                                                                         .setErrorCode(Errors.NONE.code())
+                                                                         .setMembers(Collections.singletonList(responseTwo))));
+
+            final RemoveMembersFromConsumerGroupResult missingMemberResult = env.adminClient().removeMembersFromConsumerGroup(
+                groupId,
+                new RemoveMembersFromConsumerGroupOptions(membersToRemove)
+            );
+
+            TestUtils.assertFutureError(missingMemberResult.all(), IllegalArgumentException.class);
+            // The memberOne was not included in the response.
+            TestUtils.assertFutureError(missingMemberResult.memberResult(memberOne), IllegalArgumentException.class);
+            assertNull(missingMemberResult.memberResult(memberTwo).get());
+
+
+            // Return with success.
+            env.kafkaClient().prepareResponse(prepareFindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
+            env.kafkaClient().prepareResponse(new LeaveGroupResponse(
+                    new LeaveGroupResponseData().setErrorCode(Errors.NONE.code()).setMembers(
+                        Arrays.asList(responseTwo,
+                                      new MemberResponse().setGroupInstanceId(instanceOne).setErrorCode(Errors.NONE.code())
+                        ))
+            ));
+
+            final RemoveMembersFromConsumerGroupResult noErrorResult = env.adminClient().removeMembersFromConsumerGroup(
+                groupId,
+                new RemoveMembersFromConsumerGroupOptions(membersToRemove)
+            );
+            assertNull(noErrorResult.all().get());
+            assertNull(noErrorResult.memberResult(memberOne).get());
+            assertNull(noErrorResult.memberResult(memberTwo).get());
+        }
+    }
+
+    @Test
+    public void testAlterPartitionReassignments() throws Exception {
+        try (AdminClientUnitTestEnv env = mockClientEnv()) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            TopicPartition tp1 = new TopicPartition("A", 0);
+            TopicPartition tp2 = new TopicPartition("B", 0);
+            Map> reassignments = new HashMap<>();
+            reassignments.put(tp1, Optional.empty());
+            reassignments.put(tp2, Optional.of(new NewPartitionReassignment(Arrays.asList(1, 2, 3))));
+
+            // 1. server returns less responses than number of partitions we sent
+            AlterPartitionReassignmentsResponseData responseData1 = new AlterPartitionReassignmentsResponseData();
+            ReassignablePartitionResponse normalPartitionResponse = new ReassignablePartitionResponse().setPartitionIndex(0);
+            responseData1.setResponses(Collections.singletonList(
+                    new ReassignableTopicResponse()
+                            .setName("A")
+                            .setPartitions(Collections.singletonList(normalPartitionResponse))));
+            env.kafkaClient().prepareResponse(new AlterPartitionReassignmentsResponse(responseData1));
+            AlterPartitionReassignmentsResult result1 = env.adminClient().alterPartitionReassignments(reassignments);
+            Future future1 = result1.all();
+            Future future2 = result1.values().get(tp1);
+            TestUtils.assertFutureError(future1, UnknownServerException.class);
+            TestUtils.assertFutureError(future2, UnknownServerException.class);
+
+            // 2. NOT_CONTROLLER error handling
+            AlterPartitionReassignmentsResponseData controllerErrResponseData =
+                    new AlterPartitionReassignmentsResponseData()
+                            .setErrorCode(Errors.NOT_CONTROLLER.code())
+                            .setErrorMessage(Errors.NOT_CONTROLLER.message())
+                            .setResponses(Arrays.asList(
+                                new ReassignableTopicResponse()
+                                        .setName("A")
+                                        .setPartitions(Collections.singletonList(normalPartitionResponse)),
+                                new ReassignableTopicResponse()
+                                        .setName("B")
+                                        .setPartitions(Collections.singletonList(normalPartitionResponse)))
+                            );
+            MetadataResponse controllerNodeResponse = MetadataResponse.prepareResponse(env.cluster().nodes(),
+                    env.cluster().clusterResource().clusterId(), 1, Collections.emptyList());
+            AlterPartitionReassignmentsResponseData normalResponse =
+                    new AlterPartitionReassignmentsResponseData()
+                            .setResponses(Arrays.asList(
+                                    new ReassignableTopicResponse()
+                                            .setName("A")
+                                            .setPartitions(Collections.singletonList(normalPartitionResponse)),
+                                    new ReassignableTopicResponse()
+                                            .setName("B")
+                                            .setPartitions(Collections.singletonList(normalPartitionResponse)))
+                            );
+            env.kafkaClient().prepareResponse(new AlterPartitionReassignmentsResponse(controllerErrResponseData));
+            env.kafkaClient().prepareResponse(controllerNodeResponse);
+            env.kafkaClient().prepareResponse(new AlterPartitionReassignmentsResponse(normalResponse));
+            AlterPartitionReassignmentsResult controllerErrResult = env.adminClient().alterPartitionReassignments(reassignments);
+            controllerErrResult.all().get();
+            controllerErrResult.values().get(tp1).get();
+            controllerErrResult.values().get(tp2).get();
+
+            // 3. partition-level error
+            AlterPartitionReassignmentsResponseData partitionLevelErrData =
+                    new AlterPartitionReassignmentsResponseData()
+                            .setResponses(Arrays.asList(
+                                    new ReassignableTopicResponse()
+                                            .setName("A")
+                                            .setPartitions(Collections.singletonList(new ReassignablePartitionResponse()
+                                                .setPartitionIndex(0).setErrorMessage(Errors.INVALID_REPLICA_ASSIGNMENT.message())
+                                                .setErrorCode(Errors.INVALID_REPLICA_ASSIGNMENT.code())
+                                            )),
+                                    new ReassignableTopicResponse()
+                                            .setName("B")
+                                            .setPartitions(Collections.singletonList(normalPartitionResponse)))
+                            );
+            env.kafkaClient().prepareResponse(new AlterPartitionReassignmentsResponse(partitionLevelErrData));
+            AlterPartitionReassignmentsResult partitionLevelErrResult = env.adminClient().alterPartitionReassignments(reassignments);
+            TestUtils.assertFutureError(partitionLevelErrResult.values().get(tp1), Errors.INVALID_REPLICA_ASSIGNMENT.exception().getClass());
+            partitionLevelErrResult.values().get(tp2).get();
+
+            // 4. top-level error
+            AlterPartitionReassignmentsResponseData topLevelErrResponseData =
+                    new AlterPartitionReassignmentsResponseData()
+                            .setErrorCode(Errors.CLUSTER_AUTHORIZATION_FAILED.code())
+                            .setErrorMessage(Errors.CLUSTER_AUTHORIZATION_FAILED.message())
+                            .setResponses(Arrays.asList(
+                                    new ReassignableTopicResponse()
+                                            .setName("A")
+                                            .setPartitions(Collections.singletonList(normalPartitionResponse)),
+                                    new ReassignableTopicResponse()
+                                            .setName("B")
+                                            .setPartitions(Collections.singletonList(normalPartitionResponse)))
+                            );
+            env.kafkaClient().prepareResponse(new AlterPartitionReassignmentsResponse(topLevelErrResponseData));
+            AlterPartitionReassignmentsResult topLevelErrResult = env.adminClient().alterPartitionReassignments(reassignments);
+            TestUtils.assertFutureError(topLevelErrResult.all(), Errors.CLUSTER_AUTHORIZATION_FAILED.exception().getClass());
+            TestUtils.assertFutureError(topLevelErrResult.values().get(tp1), Errors.CLUSTER_AUTHORIZATION_FAILED.exception().getClass());
+            TestUtils.assertFutureError(topLevelErrResult.values().get(tp2), Errors.CLUSTER_AUTHORIZATION_FAILED.exception().getClass());
+
+            // 5. unrepresentable topic name error
+            TopicPartition invalidTopicTP = new TopicPartition("", 0);
+            TopicPartition invalidPartitionTP = new TopicPartition("ABC", -1);
+            Map> invalidTopicReassignments = new HashMap<>();
+            invalidTopicReassignments.put(invalidPartitionTP, Optional.of(new NewPartitionReassignment(Arrays.asList(1, 2, 3))));
+            invalidTopicReassignments.put(invalidTopicTP, Optional.of(new NewPartitionReassignment(Arrays.asList(1, 2, 3))));
+            invalidTopicReassignments.put(tp1, Optional.of(new NewPartitionReassignment(Arrays.asList(1, 2, 3))));
+
+            AlterPartitionReassignmentsResponseData singlePartResponseData =
+                    new AlterPartitionReassignmentsResponseData()
+                            .setResponses(Collections.singletonList(
+                                    new ReassignableTopicResponse()
+                                            .setName("A")
+                                            .setPartitions(Collections.singletonList(normalPartitionResponse)))
+                            );
+            env.kafkaClient().prepareResponse(new AlterPartitionReassignmentsResponse(singlePartResponseData));
+            AlterPartitionReassignmentsResult unrepresentableTopicResult = env.adminClient().alterPartitionReassignments(invalidTopicReassignments);
+            TestUtils.assertFutureError(unrepresentableTopicResult.values().get(invalidTopicTP), InvalidTopicException.class);
+            TestUtils.assertFutureError(unrepresentableTopicResult.values().get(invalidPartitionTP), InvalidTopicException.class);
+            unrepresentableTopicResult.values().get(tp1).get();
+
+            // Test success scenario
+            AlterPartitionReassignmentsResponseData noErrResponseData =
+                    new AlterPartitionReassignmentsResponseData()
+                            .setErrorCode(Errors.NONE.code())
+                            .setErrorMessage(Errors.NONE.message())
+                            .setResponses(Arrays.asList(
+                                    new ReassignableTopicResponse()
+                                            .setName("A")
+                                            .setPartitions(Collections.singletonList(normalPartitionResponse)),
+                                    new ReassignableTopicResponse()
+                                            .setName("B")
+                                            .setPartitions(Collections.singletonList(normalPartitionResponse)))
+                            );
+            env.kafkaClient().prepareResponse(new AlterPartitionReassignmentsResponse(noErrResponseData));
+            AlterPartitionReassignmentsResult noErrResult = env.adminClient().alterPartitionReassignments(reassignments);
+            noErrResult.all().get();
+            noErrResult.values().get(tp1).get();
+            noErrResult.values().get(tp2).get();
+        }
+    }
+
+    @Test
+    public void testListPartitionReassignments() throws Exception {
+        try (AdminClientUnitTestEnv env = mockClientEnv()) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            TopicPartition tp1 = new TopicPartition("A", 0);
+            OngoingPartitionReassignment tp1PartitionReassignment = new OngoingPartitionReassignment()
+                    .setPartitionIndex(0)
+                    .setRemovingReplicas(Arrays.asList(1, 2, 3))
+                    .setAddingReplicas(Arrays.asList(4, 5, 6))
+                    .setReplicas(Arrays.asList(1, 2, 3, 4, 5, 6));
+            OngoingTopicReassignment tp1Reassignment = new OngoingTopicReassignment().setName("A")
+                    .setPartitions(Collections.singletonList(tp1PartitionReassignment));
+
+            TopicPartition tp2 = new TopicPartition("B", 0);
+            OngoingPartitionReassignment tp2PartitionReassignment = new OngoingPartitionReassignment()
+                    .setPartitionIndex(0)
+                    .setRemovingReplicas(Arrays.asList(1, 2, 3))
+                    .setAddingReplicas(Arrays.asList(4, 5, 6))
+                    .setReplicas(Arrays.asList(1, 2, 3, 4, 5, 6));
+            OngoingTopicReassignment tp2Reassignment = new OngoingTopicReassignment().setName("B")
+                    .setPartitions(Collections.singletonList(tp2PartitionReassignment));
+
+            // 1. NOT_CONTROLLER error handling
+            ListPartitionReassignmentsResponseData notControllerData = new ListPartitionReassignmentsResponseData()
+                    .setErrorCode(Errors.NOT_CONTROLLER.code())
+                    .setErrorMessage(Errors.NOT_CONTROLLER.message());
+            MetadataResponse controllerNodeResponse = MetadataResponse.prepareResponse(env.cluster().nodes(),
+                    env.cluster().clusterResource().clusterId(), 1, Collections.emptyList());
+            ListPartitionReassignmentsResponseData reassignmentsData = new ListPartitionReassignmentsResponseData()
+                    .setTopics(Arrays.asList(tp1Reassignment, tp2Reassignment));
+            env.kafkaClient().prepareResponse(new ListPartitionReassignmentsResponse(notControllerData));
+            env.kafkaClient().prepareResponse(controllerNodeResponse);
+            env.kafkaClient().prepareResponse(new ListPartitionReassignmentsResponse(reassignmentsData));
+
+            ListPartitionReassignmentsResult noControllerResult = env.adminClient().listPartitionReassignments();
+            noControllerResult.reassignments().get(); // no error
+
+            // 2. UNKNOWN_TOPIC_OR_EXCEPTION_ERROR
+            ListPartitionReassignmentsResponseData unknownTpData = new ListPartitionReassignmentsResponseData()
+                    .setErrorCode(Errors.UNKNOWN_TOPIC_OR_PARTITION.code())
+                    .setErrorMessage(Errors.UNKNOWN_TOPIC_OR_PARTITION.message());
+            env.kafkaClient().prepareResponse(new ListPartitionReassignmentsResponse(unknownTpData));
+
+            ListPartitionReassignmentsResult unknownTpResult = env.adminClient().listPartitionReassignments(new HashSet<>(Arrays.asList(tp1, tp2)));
+            TestUtils.assertFutureError(unknownTpResult.reassignments(), UnknownTopicOrPartitionException.class);
+
+            // 3. Success
+            ListPartitionReassignmentsResponseData responseData = new ListPartitionReassignmentsResponseData()
+                    .setTopics(Arrays.asList(tp1Reassignment, tp2Reassignment));
+            env.kafkaClient().prepareResponse(new ListPartitionReassignmentsResponse(responseData));
+            ListPartitionReassignmentsResult responseResult = env.adminClient().listPartitionReassignments();
+
+            Map reassignments = responseResult.reassignments().get();
+
+            PartitionReassignment tp1Result = reassignments.get(tp1);
+            assertEquals(tp1PartitionReassignment.addingReplicas(), tp1Result.addingReplicas());
+            assertEquals(tp1PartitionReassignment.removingReplicas(), tp1Result.removingReplicas());
+            assertEquals(tp1PartitionReassignment.replicas(), tp1Result.replicas());
+            assertEquals(tp1PartitionReassignment.replicas(), tp1Result.replicas());
+            PartitionReassignment tp2Result = reassignments.get(tp2);
+            assertEquals(tp2PartitionReassignment.addingReplicas(), tp2Result.addingReplicas());
+            assertEquals(tp2PartitionReassignment.removingReplicas(), tp2Result.removingReplicas());
+            assertEquals(tp2PartitionReassignment.replicas(), tp2Result.replicas());
+            assertEquals(tp2PartitionReassignment.replicas(), tp2Result.replicas());
+        }
+    }
+
+    @Test
+    public void testGetSubLevelError() {
+        List memberIdentities = Arrays.asList(
+            new MemberIdentity().setGroupInstanceId("instance-0"),
+            new MemberIdentity().setGroupInstanceId("instance-1"));
+        Map errorsMap = new HashMap<>();
+        errorsMap.put(memberIdentities.get(0), Errors.NONE);
+        errorsMap.put(memberIdentities.get(1), Errors.FENCED_INSTANCE_ID);
+        assertEquals(IllegalArgumentException.class, KafkaAdminClient.getSubLevelError(errorsMap,
+                                                                                       new MemberIdentity().setGroupInstanceId("non-exist-id"), "For unit test").getClass());
+        assertNull(KafkaAdminClient.getSubLevelError(errorsMap, memberIdentities.get(0), "For unit test"));
+        assertEquals(FencedInstanceIdException.class, KafkaAdminClient.getSubLevelError(
+            errorsMap, memberIdentities.get(1), "For unit test").getClass());
+    }
+
+    private static MemberDescription convertToMemberDescriptions(DescribedGroupMember member,
+                                                                 MemberAssignment assignment) {
+        return new MemberDescription(member.memberId(),
+                                     Optional.ofNullable(member.groupInstanceId()),
+                                     member.clientId(),
+                                     member.clientHost(),
+                                     assignment);
+    }
+
     @SafeVarargs
     private static  void assertCollectionIs(Collection collection, T... elements) {
         for (T element : elements) {
@@ -1241,7 +2316,5 @@ boolean callHasExpired(KafkaAdminClient.Call call) {
                 }
             }
         }
-
     }
-
 }
diff --git a/clients/src/test/java/org/apache/kafka/clients/admin/MemberDescriptionTest.java b/clients/src/test/java/org/apache/kafka/clients/admin/MemberDescriptionTest.java
new file mode 100644
index 0000000000000..be309758829b1
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/clients/admin/MemberDescriptionTest.java
@@ -0,0 +1,102 @@
+/*
+ * 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.admin;
+
+import org.apache.kafka.common.TopicPartition;
+import org.junit.Test;
+
+import java.util.Collections;
+import java.util.Optional;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotEquals;
+
+public class MemberDescriptionTest {
+
+    private static final String MEMBER_ID = "member_id";
+    private static final Optional INSTANCE_ID = Optional.of("instanceId");
+    private static final String CLIENT_ID = "client_id";
+    private static final String HOST = "host";
+    private static final MemberAssignment ASSIGNMENT;
+    private static final MemberDescription STATIC_MEMBER_DESCRIPTION;
+
+    static {
+        ASSIGNMENT = new MemberAssignment(Collections.singleton(new TopicPartition("topic", 1)));
+        STATIC_MEMBER_DESCRIPTION = new MemberDescription(MEMBER_ID,
+                                                          INSTANCE_ID,
+                                                          CLIENT_ID,
+                                                          HOST,
+                                                          ASSIGNMENT);
+    }
+
+    @Test
+    public void testEqualsWithoutGroupInstanceId() {
+        MemberDescription dynamicMemberDescription = new MemberDescription(MEMBER_ID,
+                                                                           CLIENT_ID,
+                                                                           HOST,
+                                                                           ASSIGNMENT);
+
+        MemberDescription identityDescription = new MemberDescription(MEMBER_ID,
+                                                                      CLIENT_ID,
+                                                                      HOST,
+                                                                      ASSIGNMENT);
+
+        assertNotEquals(STATIC_MEMBER_DESCRIPTION, dynamicMemberDescription);
+        assertNotEquals(STATIC_MEMBER_DESCRIPTION.hashCode(), dynamicMemberDescription.hashCode());
+
+        // Check self equality.
+        assertEquals(dynamicMemberDescription, dynamicMemberDescription);
+        assertEquals(dynamicMemberDescription, identityDescription);
+        assertEquals(dynamicMemberDescription.hashCode(), identityDescription.hashCode());
+    }
+
+    @Test
+    public void testEqualsWithGroupInstanceId() {
+        // Check self equality.
+        assertEquals(STATIC_MEMBER_DESCRIPTION, STATIC_MEMBER_DESCRIPTION);
+
+        MemberDescription identityDescription = new MemberDescription(MEMBER_ID,
+                                                                      INSTANCE_ID,
+                                                                      CLIENT_ID,
+                                                                      HOST,
+                                                                      ASSIGNMENT);
+
+        assertEquals(STATIC_MEMBER_DESCRIPTION, identityDescription);
+        assertEquals(STATIC_MEMBER_DESCRIPTION.hashCode(), identityDescription.hashCode());
+    }
+
+    @Test
+    public void testNonEqual() {
+        MemberDescription newMemberDescription = new MemberDescription("new_member",
+                                                                       INSTANCE_ID,
+                                                                       CLIENT_ID,
+                                                                       HOST,
+                                                                       ASSIGNMENT);
+
+        assertNotEquals(STATIC_MEMBER_DESCRIPTION, newMemberDescription);
+        assertNotEquals(STATIC_MEMBER_DESCRIPTION.hashCode(), newMemberDescription.hashCode());
+
+        MemberDescription newInstanceDescription = new MemberDescription(MEMBER_ID,
+                                                                         Optional.of("new_instance"),
+                                                                         CLIENT_ID,
+                                                                         HOST,
+                                                                         ASSIGNMENT);
+
+        assertNotEquals(STATIC_MEMBER_DESCRIPTION, newInstanceDescription);
+        assertNotEquals(STATIC_MEMBER_DESCRIPTION.hashCode(), newInstanceDescription.hashCode());
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/clients/admin/MockAdminClient.java b/clients/src/test/java/org/apache/kafka/clients/admin/MockAdminClient.java
index d721245be854e..240071abc4d54 100644
--- a/clients/src/test/java/org/apache/kafka/clients/admin/MockAdminClient.java
+++ b/clients/src/test/java/org/apache/kafka/clients/admin/MockAdminClient.java
@@ -21,10 +21,12 @@
 import org.apache.kafka.common.MetricName;
 import org.apache.kafka.common.Node;
 import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.ElectionType;
 import org.apache.kafka.common.TopicPartitionInfo;
 import org.apache.kafka.common.TopicPartitionReplica;
 import org.apache.kafka.common.acl.AclBinding;
 import org.apache.kafka.common.acl.AclBindingFilter;
+import org.apache.kafka.common.acl.AclOperation;
 import org.apache.kafka.common.config.ConfigResource;
 import org.apache.kafka.common.errors.TimeoutException;
 import org.apache.kafka.common.errors.TopicExistsException;
@@ -38,6 +40,8 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
 
 public class MockAdminClient extends AdminClient {
     public static final String DEFAULT_CLUSTER_ID = "I4ZmrWqfT2e-upky_4fdPA";
@@ -125,30 +129,33 @@ public DescribeClusterResult describeCluster(DescribeClusterOptions options) {
         KafkaFutureImpl> nodesFuture = new KafkaFutureImpl<>();
         KafkaFutureImpl controllerFuture = new KafkaFutureImpl<>();
         KafkaFutureImpl brokerIdFuture = new KafkaFutureImpl<>();
+        KafkaFutureImpl> authorizedOperationsFuture = new KafkaFutureImpl<>();
 
         if (timeoutNextRequests > 0) {
             nodesFuture.completeExceptionally(new TimeoutException());
             controllerFuture.completeExceptionally(new TimeoutException());
             brokerIdFuture.completeExceptionally(new TimeoutException());
+            authorizedOperationsFuture.completeExceptionally(new TimeoutException());
             --timeoutNextRequests;
         } else {
             nodesFuture.complete(brokers);
             controllerFuture.complete(controller);
             brokerIdFuture.complete(clusterId);
+            authorizedOperationsFuture.complete(Collections.emptySet());
         }
 
-        return new DescribeClusterResult(nodesFuture, controllerFuture, brokerIdFuture);
+        return new DescribeClusterResult(nodesFuture, controllerFuture, brokerIdFuture, authorizedOperationsFuture);
     }
 
     @Override
     public CreateTopicsResult createTopics(Collection newTopics, CreateTopicsOptions options) {
-        Map> createTopicResult = new HashMap<>();
+        Map> createTopicResult = new HashMap<>();
 
         if (timeoutNextRequests > 0) {
             for (final NewTopic newTopic : newTopics) {
                 String topicName = newTopic.name();
 
-                KafkaFutureImpl future = new KafkaFutureImpl<>();
+                KafkaFutureImpl future = new KafkaFutureImpl<>();
                 future.completeExceptionally(new TimeoutException());
                 createTopicResult.put(topicName, future);
             }
@@ -158,7 +165,7 @@ public CreateTopicsResult createTopics(Collection newTopics, CreateTop
         }
 
         for (final NewTopic newTopic : newTopics) {
-            KafkaFutureImpl future = new KafkaFutureImpl<>();
+            KafkaFutureImpl future = new KafkaFutureImpl<>();
 
             String topicName = newTopic.name();
             if (allTopics.containsKey(topicName)) {
@@ -199,7 +206,11 @@ public ListTopicsResult listTopics(ListTopicsOptions options) {
 
         for (Map.Entry topicDescription : allTopics.entrySet()) {
             String topicName = topicDescription.getKey();
-            topicListings.put(topicName, new TopicListing(topicName, topicDescription.getValue().isInternalTopic));
+            if (topicDescription.getValue().fetchesRemainingUntilVisible > 0) {
+                topicDescription.getValue().fetchesRemainingUntilVisible--;
+            } else {
+                topicListings.put(topicName, new TopicListing(topicName, topicDescription.getValue().isInternalTopic));
+            }
         }
 
         KafkaFutureImpl> future = new KafkaFutureImpl<>();
@@ -226,11 +237,16 @@ public DescribeTopicsResult describeTopics(Collection topicNames, Descri
             for (Map.Entry topicDescription : allTopics.entrySet()) {
                 String topicName = topicDescription.getKey();
                 if (topicName.equals(requestedTopic) && !topicDescription.getValue().markedForDeletion) {
-                    TopicMetadata topicMetadata = topicDescription.getValue();
-                    KafkaFutureImpl future = new KafkaFutureImpl<>();
-                    future.complete(new TopicDescription(topicName, topicMetadata.isInternalTopic, topicMetadata.partitions));
-                    topicDescriptions.put(topicName, future);
-                    break;
+                    if (topicDescription.getValue().fetchesRemainingUntilVisible > 0) {
+                        topicDescription.getValue().fetchesRemainingUntilVisible--;
+                    } else {
+                        TopicMetadata topicMetadata = topicDescription.getValue();
+                        KafkaFutureImpl future = new KafkaFutureImpl<>();
+                        future.complete(new TopicDescription(topicName, topicMetadata.isInternalTopic, topicMetadata.partitions,
+                                Collections.emptySet()));
+                        topicDescriptions.put(topicName, future);
+                        break;
+                    }
                 }
             }
             if (!topicDescriptions.containsKey(requestedTopic)) {
@@ -327,10 +343,30 @@ public DeleteConsumerGroupsResult deleteConsumerGroups(Collection groupI
         throw new UnsupportedOperationException("Not implemented yet");
     }
 
+    @Override
+    public DeleteConsumerGroupOffsetsResult deleteConsumerGroupOffsets(String groupId, Set partitions, DeleteConsumerGroupOffsetsOptions options) {
+        throw new UnsupportedOperationException("Not implemented yet");
+    }
+
+    @Deprecated
+    @Override
     public ElectPreferredLeadersResult electPreferredLeaders(Collection partitions, ElectPreferredLeadersOptions options) {
         throw new UnsupportedOperationException("Not implemented yet");
     }
 
+    @Override
+    public ElectLeadersResult electLeaders(
+            ElectionType electionType,
+            Set partitions,
+            ElectLeadersOptions options) {
+        throw new UnsupportedOperationException("Not implemented yet");
+    }
+
+    @Override
+    public RemoveMembersFromConsumerGroupResult removeMembersFromConsumerGroup(String groupId, RemoveMembersFromConsumerGroupOptions options) {
+        throw new UnsupportedOperationException("Not implemented yet");
+    }
+
     @Override
     public CreateAclsResult createAcls(Collection acls, CreateAclsOptions options) {
         throw new UnsupportedOperationException("Not implemented yet");
@@ -369,10 +405,17 @@ public DescribeConfigsResult describeConfigs(Collection resource
     }
 
     @Override
+    @Deprecated
     public AlterConfigsResult alterConfigs(Map configs, AlterConfigsOptions options) {
         throw new UnsupportedOperationException("Not implemented yet");
     }
 
+    @Override
+    public AlterConfigsResult incrementalAlterConfigs(Map> configs,
+                                                      AlterConfigsOptions options) {
+        throw new UnsupportedOperationException("Not implemented yet");
+    }
+
     @Override
     public AlterReplicaLogDirsResult alterReplicaLogDirs(Map replicaAssignment, AlterReplicaLogDirsOptions options) {
         throw new UnsupportedOperationException("Not implemented yet");
@@ -388,6 +431,17 @@ public DescribeReplicaLogDirsResult describeReplicaLogDirs(Collection> reassignments,
+                                                                         AlterPartitionReassignmentsOptions options) {
+        throw new UnsupportedOperationException("Not implemented yet");
+    }
+
+    @Override
+    public ListPartitionReassignmentsResult listPartitionReassignments(Optional> partitions, ListPartitionReassignmentsOptions options) {
+        throw new UnsupportedOperationException("Not implemented yet");
+    }
+
     @Override
     public void close(Duration timeout) {}
 
@@ -396,6 +450,7 @@ private final static class TopicMetadata {
         final boolean isInternalTopic;
         final List partitions;
         final Map configs;
+        int fetchesRemainingUntilVisible;
 
         public boolean markedForDeletion;
 
@@ -406,6 +461,7 @@ private final static class TopicMetadata {
             this.partitions = partitions;
             this.configs = configs != null ? configs : Collections.emptyMap();
             this.markedForDeletion = false;
+            this.fetchesRemainingUntilVisible = 0;
         }
     }
 
@@ -417,4 +473,12 @@ public void setMockMetrics(MetricName name, Metric metric) {
     public Map metrics() {
         return mockMetrics;
     }
+
+    public void setFetchesRemainingUntilVisible(String topicName, int fetchesRemainingUntilVisible) {
+        TopicMetadata metadata = allTopics.get(topicName);
+        if (metadata == null) {
+            throw new RuntimeException("No such topic as " + topicName);
+        }
+        metadata.fetchesRemainingUntilVisible = fetchesRemainingUntilVisible;
+    }
 }
diff --git a/clients/src/test/java/org/apache/kafka/common/record/SimpleMemoryRecordsTest.java b/clients/src/test/java/org/apache/kafka/clients/admin/RemoveMembersFromConsumerGroupOptionsTest.java
similarity index 52%
rename from clients/src/test/java/org/apache/kafka/common/record/SimpleMemoryRecordsTest.java
rename to clients/src/test/java/org/apache/kafka/clients/admin/RemoveMembersFromConsumerGroupOptionsTest.java
index 2909566d90efd..41aa386a0590a 100644
--- a/clients/src/test/java/org/apache/kafka/common/record/SimpleMemoryRecordsTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/admin/RemoveMembersFromConsumerGroupOptionsTest.java
@@ -14,28 +14,22 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-
-package org.apache.kafka.common.record;
+package org.apache.kafka.clients.admin;
 
 import org.junit.Test;
 
+import java.util.Collections;
+
 import static org.junit.Assert.assertEquals;
 
-/**
- * Non-parameterized MemoryRecords tests.
- */
-public class SimpleMemoryRecordsTest {
+public class RemoveMembersFromConsumerGroupOptionsTest {
 
     @Test
-    public void testToStringIfLz4ChecksumIsCorrupted() {
-        long timestamp = 1000000;
-        MemoryRecords memoryRecords = MemoryRecords.withRecords(CompressionType.LZ4,
-                new SimpleRecord(timestamp, "key1".getBytes(), "value1".getBytes()),
-                new SimpleRecord(timestamp + 1, "key2".getBytes(), "value2".getBytes()));
-        // Change the lz4 checksum value (not the kafka record crc) so that it doesn't match the contents
-        int lz4ChecksumOffset = 6;
-        memoryRecords.buffer().array()[DefaultRecordBatch.RECORD_BATCH_OVERHEAD + lz4ChecksumOffset] = 0;
-        assertEquals("[(record=CORRUPTED)]", memoryRecords.toString());
-    }
+    public void testConstructor() {
+        RemoveMembersFromConsumerGroupOptions options = new RemoveMembersFromConsumerGroupOptions(
+            Collections.singleton(new MemberToRemove("instance-1")));
 
+        assertEquals(Collections.singleton(
+            new MemberToRemove("instance-1")), options.members());
+    }
 }
diff --git a/clients/src/test/java/org/apache/kafka/clients/admin/RemoveMembersFromConsumerGroupResultTest.java b/clients/src/test/java/org/apache/kafka/clients/admin/RemoveMembersFromConsumerGroupResultTest.java
new file mode 100644
index 0000000000000..e2da23b768e3d
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/clients/admin/RemoveMembersFromConsumerGroupResultTest.java
@@ -0,0 +1,119 @@
+/*
+ * 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.admin;
+
+import org.apache.kafka.common.errors.FencedInstanceIdException;
+import org.apache.kafka.common.errors.GroupAuthorizationException;
+import org.apache.kafka.common.internals.KafkaFutureImpl;
+import org.apache.kafka.common.message.LeaveGroupRequestData.MemberIdentity;
+
+import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.test.TestUtils;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.ExecutionException;
+
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertThrows;
+
+public class RemoveMembersFromConsumerGroupResultTest {
+
+    private final MemberToRemove instanceOne = new MemberToRemove("instance-1");
+    private final MemberToRemove instanceTwo = new MemberToRemove("instance-2");
+    private Set membersToRemove;
+    private Map errorsMap;
+
+    private KafkaFutureImpl> memberFutures;
+
+    @Before
+    public void setUp() {
+        memberFutures = new KafkaFutureImpl<>();
+        membersToRemove = new HashSet<>();
+        membersToRemove.add(instanceOne);
+        membersToRemove.add(instanceTwo);
+
+        errorsMap = new HashMap<>();
+        errorsMap.put(instanceOne.toMemberIdentity(), Errors.NONE);
+        errorsMap.put(instanceTwo.toMemberIdentity(), Errors.FENCED_INSTANCE_ID);
+    }
+
+    @Test
+    public void testTopLevelErrorConstructor() throws InterruptedException {
+        memberFutures.completeExceptionally(Errors.GROUP_AUTHORIZATION_FAILED.exception());
+        RemoveMembersFromConsumerGroupResult topLevelErrorResult =
+            new RemoveMembersFromConsumerGroupResult(memberFutures, membersToRemove);
+        TestUtils.assertFutureError(topLevelErrorResult.all(), GroupAuthorizationException.class);
+    }
+
+    @Test
+    public void testMemberLevelErrorConstructor() throws InterruptedException, ExecutionException {
+        createAndVerifyMemberLevelError();
+    }
+
+    @Test
+    public void testMemberMissingErrorInRequestConstructor() throws InterruptedException, ExecutionException {
+        errorsMap.remove(instanceTwo.toMemberIdentity());
+        memberFutures.complete(errorsMap);
+        assertFalse(memberFutures.isCompletedExceptionally());
+        RemoveMembersFromConsumerGroupResult missingMemberResult =
+            new RemoveMembersFromConsumerGroupResult(memberFutures, membersToRemove);
+
+        TestUtils.assertFutureError(missingMemberResult.all(), IllegalArgumentException.class);
+        assertNull(missingMemberResult.memberResult(instanceOne).get());
+        TestUtils.assertFutureError(missingMemberResult.memberResult(instanceTwo), IllegalArgumentException.class);
+    }
+
+    @Test
+    public void testMemberLevelErrorInResponseConstructor() throws InterruptedException, ExecutionException {
+        RemoveMembersFromConsumerGroupResult memberLevelErrorResult = createAndVerifyMemberLevelError();
+        assertThrows(IllegalArgumentException.class, () -> memberLevelErrorResult.memberResult(
+            new MemberToRemove("invalid-instance-id"))
+        );
+    }
+
+    @Test
+    public void testNoErrorConstructor() throws ExecutionException, InterruptedException {
+        Map errorsMap = new HashMap<>();
+        errorsMap.put(instanceOne.toMemberIdentity(), Errors.NONE);
+        errorsMap.put(instanceTwo.toMemberIdentity(), Errors.NONE);
+        RemoveMembersFromConsumerGroupResult noErrorResult =
+            new RemoveMembersFromConsumerGroupResult(memberFutures, membersToRemove);
+        memberFutures.complete(errorsMap);
+
+        assertNull(noErrorResult.all().get());
+        assertNull(noErrorResult.memberResult(instanceOne).get());
+        assertNull(noErrorResult.memberResult(instanceTwo).get());
+    }
+
+    private RemoveMembersFromConsumerGroupResult createAndVerifyMemberLevelError() throws InterruptedException, ExecutionException {
+        memberFutures.complete(errorsMap);
+        assertFalse(memberFutures.isCompletedExceptionally());
+        RemoveMembersFromConsumerGroupResult memberLevelErrorResult =
+            new RemoveMembersFromConsumerGroupResult(memberFutures, membersToRemove);
+
+        TestUtils.assertFutureError(memberLevelErrorResult.all(), FencedInstanceIdException.class);
+        assertNull(memberLevelErrorResult.memberResult(instanceOne).get());
+        TestUtils.assertFutureError(memberLevelErrorResult.memberResult(instanceTwo), FencedInstanceIdException.class);
+        return memberLevelErrorResult;
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/CooperativeStickyAssignorTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/CooperativeStickyAssignorTest.java
new file mode 100644
index 0000000000000..aed8c09537068
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/CooperativeStickyAssignorTest.java
@@ -0,0 +1,105 @@
+/*
+ * 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;
+
+import static org.junit.Assert.assertTrue;
+
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import org.apache.kafka.clients.consumer.ConsumerPartitionAssignor.Subscription;
+import org.apache.kafka.clients.consumer.internals.AbstractStickyAssignor;
+import org.apache.kafka.clients.consumer.internals.AbstractStickyAssignorTest;
+import org.apache.kafka.common.TopicPartition;
+
+public class CooperativeStickyAssignorTest extends AbstractStickyAssignorTest {
+
+    @Override
+    public AbstractStickyAssignor createAssignor() {
+        return new CooperativeStickyAssignor();
+    }
+
+    @Override
+    public Subscription buildSubscription(List topics, List partitions) {
+        return new Subscription(topics, assignor.subscriptionUserData(new HashSet<>(topics)), partitions);
+    }
+
+    /**
+     * The cooperative assignor must do some additional work and verification of some assignments relative to the eager
+     * assignor, since it may or may not need to trigger a second follow-up rebalance.
+     * 
      
+     * In addition to the validity requirements described in
+     * {@link org.apache.kafka.clients.consumer.internals.AbstractStickyAssignorTest#verifyValidityAndBalance(Map, Map, Map)},
+     * we must verify that no partition is being revoked and reassigned during the same rebalance. This means the initial
+     * assignment may be unbalanced, so if we do detect partitions being revoked we should trigger a second "rebalance"
+     * to get the final assignment and then verify that it is both valid and balanced.
+     */
+    @Override
+    public void verifyValidityAndBalance(Map subscriptions,
+                                         Map> assignments,
+                                         Map partitionsPerTopic) {
+        int rebalances = 0;
+        // partitions are being revoked, we must go through another assignment to get the final state
+        while (verifyCooperativeValidity(subscriptions, assignments)) {
+
+            // update the subscriptions with the now owned partitions
+            for (Map.Entry> entry : assignments.entrySet()) {
+                String consumer = entry.getKey();
+                Subscription oldSubscription = subscriptions.get(consumer);
+                subscriptions.put(consumer, buildSubscription(oldSubscription.topics(), entry.getValue()));
+            }
+
+            assignments.clear();
+            assignments.putAll(assignor.assign(partitionsPerTopic, subscriptions));
+            ++rebalances;
+
+            assertTrue(rebalances <= 4);
+        }
+
+        // Check the validity and balance of the final assignment
+        super.verifyValidityAndBalance(subscriptions, assignments, partitionsPerTopic);
+    }
+
+    // Returns true if partitions are being revoked, indicating a second rebalance will be triggered
+    private boolean verifyCooperativeValidity(Map subscriptions, Map> assignments) {
+        Set allAddedPartitions = new HashSet<>();
+        Set allRevokedPartitions = new HashSet<>();
+        for (Map.Entry> entry : assignments.entrySet()) {
+            List ownedPartitions = subscriptions.get(entry.getKey()).ownedPartitions();
+            List assignedPartitions = entry.getValue();
+
+            Set revokedPartitions = new HashSet<>(ownedPartitions);
+            revokedPartitions.removeAll(assignedPartitions);
+
+            Set addedPartitions = new HashSet<>(assignedPartitions);
+            addedPartitions.removeAll(ownedPartitions);
+
+            allAddedPartitions.addAll(addedPartitions);
+            allRevokedPartitions.addAll(revokedPartitions);
+        }
+
+        Set intersection = new HashSet<>(allAddedPartitions);
+        intersection.retainAll(allRevokedPartitions);
+        assertTrue("Error: Some partitions were assigned to a new consumer during the same rebalance they are being "
+            + "revoked from their previous owner."
+            + "Partitions: " + intersection.toString(),
+            intersection.isEmpty());
+
+        return !allRevokedPartitions.isEmpty();
+    }
+}
\ No newline at end of file
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 a5161b44a8d2b..10027194bed9c 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
@@ -16,7 +16,9 @@
  */
 package org.apache.kafka.clients.consumer;
 
+import org.apache.kafka.clients.ApiVersions;
 import org.apache.kafka.clients.ClientRequest;
+import org.apache.kafka.clients.GroupRebalanceConfig;
 import org.apache.kafka.clients.KafkaClient;
 import org.apache.kafka.clients.MockClient;
 import org.apache.kafka.clients.consumer.internals.ConsumerCoordinator;
@@ -26,11 +28,10 @@
 import org.apache.kafka.clients.consumer.internals.ConsumerNetworkClient;
 import org.apache.kafka.clients.consumer.internals.ConsumerProtocol;
 import org.apache.kafka.clients.consumer.internals.Fetcher;
-import org.apache.kafka.clients.consumer.internals.Heartbeat;
-import org.apache.kafka.clients.consumer.internals.PartitionAssignor;
 import org.apache.kafka.clients.consumer.internals.SubscriptionState;
 import org.apache.kafka.common.Cluster;
 import org.apache.kafka.common.KafkaException;
+import org.apache.kafka.common.MetricName;
 import org.apache.kafka.common.Node;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.errors.AuthenticationException;
@@ -39,8 +40,12 @@
 import org.apache.kafka.common.errors.InvalidGroupIdException;
 import org.apache.kafka.common.errors.InvalidTopicException;
 import org.apache.kafka.common.errors.WakeupException;
+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.LeaveGroupResponseData;
 import org.apache.kafka.common.internals.ClusterResourceListeners;
+import org.apache.kafka.common.message.SyncGroupResponseData;
 import org.apache.kafka.common.metrics.Metrics;
 import org.apache.kafka.common.metrics.Sensor;
 import org.apache.kafka.common.network.Selectable;
@@ -89,6 +94,7 @@
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.Iterator;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
@@ -142,6 +148,13 @@ public class KafkaConsumerTest {
     private final int autoCommitIntervalMs = 500;
 
     private final String groupId = "mock-group";
+    private final Optional groupInstanceId = Optional.of("mock-instance");
+
+    private final String partitionRevoked = "Hit partition revoke ";
+    private final String partitionAssigned = "Hit partition assign ";
+    private final String partitionLost = "Hit partition lost ";
+
+    private final Collection singleTopicPartition = Collections.singleton(new TopicPartition(topic, 0));
 
     @Test
     public void testMetricsReporterAutoGeneratedClientId() {
@@ -203,6 +216,16 @@ public void testInvalidSocketReceiveBufferSize() {
         new KafkaConsumer<>(config, new ByteArrayDeserializer(), new ByteArrayDeserializer());
     }
 
+    @Test
+    public void shouldIgnoreGroupInstanceIdForEmptyGroupId() {
+        Map config = new HashMap<>();
+        config.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
+        config.put(ConsumerConfig.GROUP_INSTANCE_ID_CONFIG, "instance_id");
+        KafkaConsumer consumer = new KafkaConsumer<>(
+                config, new ByteArrayDeserializer(), new ByteArrayDeserializer());
+        consumer.close();
+    }
+
     @Test
     public void testSubscription() {
         KafkaConsumer consumer = newConsumer(groupId);
@@ -379,9 +402,10 @@ public void verifyHeartbeatSent() throws Exception {
         initMetadata(client, Collections.singletonMap(topic, 1));
         Node node = metadata.fetch().nodes().get(0);
 
-        PartitionAssignor assignor = new RoundRobinAssignor();
+        ConsumerPartitionAssignor assignor = new RoundRobinAssignor();
+
+        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, true, groupInstanceId);
 
-        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, true);
         consumer.subscribe(singleton(topic), getConsumerRebalanceListener(consumer));
         Node coordinator = prepareRebalance(client, node, assignor, singletonList(tp0), null);
 
@@ -391,7 +415,7 @@ public void verifyHeartbeatSent() throws Exception {
 
         assertEquals(singleton(tp0), consumer.assignment());
 
-        AtomicBoolean heartbeatReceived = prepareHeartbeatResponse(client, coordinator);
+        AtomicBoolean heartbeatReceived = prepareHeartbeatResponse(client, coordinator, Errors.NONE);
 
         // heartbeat interval is 2 seconds
         time.sleep(heartbeatIntervalMs);
@@ -412,9 +436,9 @@ public void verifyHeartbeatSentWhenFetchedDataReady() throws Exception {
 
         initMetadata(client, Collections.singletonMap(topic, 1));
         Node node = metadata.fetch().nodes().get(0);
-        PartitionAssignor assignor = new RoundRobinAssignor();
+        ConsumerPartitionAssignor assignor = new RoundRobinAssignor();
 
-        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, true);
+        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, true, groupInstanceId);
         consumer.subscribe(singleton(topic), getConsumerRebalanceListener(consumer));
         Node coordinator = prepareRebalance(client, node, assignor, singletonList(tp0), null);
 
@@ -426,7 +450,7 @@ public void verifyHeartbeatSentWhenFetchedDataReady() throws Exception {
         client.poll(0, time.milliseconds());
 
         client.prepareResponseFrom(fetchResponse(tp0, 5, 0), node);
-        AtomicBoolean heartbeatReceived = prepareHeartbeatResponse(client, coordinator);
+        AtomicBoolean heartbeatReceived = prepareHeartbeatResponse(client, coordinator, Errors.NONE);
 
         time.sleep(heartbeatIntervalMs);
         Thread.sleep(heartbeatIntervalMs);
@@ -447,9 +471,9 @@ public void verifyPollTimesOutDuringMetadataUpdate() {
         initMetadata(client, Collections.singletonMap(topic, 1));
         Node node = metadata.fetch().nodes().get(0);
 
-        final PartitionAssignor assignor = new RoundRobinAssignor();
+        final ConsumerPartitionAssignor assignor = new RoundRobinAssignor();
 
-        final KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, true);
+        final KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, true, groupInstanceId);
         consumer.subscribe(singleton(topic), getConsumerRebalanceListener(consumer));
         prepareRebalance(client, node, assignor, singletonList(tp0), null);
 
@@ -471,9 +495,9 @@ public void verifyDeprecatedPollDoesNotTimeOutDuringMetadataUpdate() {
         initMetadata(client, Collections.singletonMap(topic, 1));
         Node node = metadata.fetch().nodes().get(0);
 
-        final PartitionAssignor assignor = new RoundRobinAssignor();
+        final ConsumerPartitionAssignor assignor = new RoundRobinAssignor();
 
-        final KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, true);
+        final KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, true, groupInstanceId);
         consumer.subscribe(singleton(topic), getConsumerRebalanceListener(consumer));
         prepareRebalance(client, node, assignor, singletonList(tp0), null);
 
@@ -494,9 +518,9 @@ public void verifyNoCoordinatorLookupForManualAssignmentWithSeek() {
         MockClient client = new MockClient(time, metadata);
 
         initMetadata(client, Collections.singletonMap(topic, 1));
-        PartitionAssignor assignor = new RoundRobinAssignor();
+        ConsumerPartitionAssignor assignor = new RoundRobinAssignor();
 
-        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, true);
+        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, true, groupInstanceId);
         consumer.assign(singleton(tp0));
         consumer.seekToBeginning(singleton(tp0));
 
@@ -569,13 +593,13 @@ public void testMissingOffsetNoResetPolicy() {
         initMetadata(client, Collections.singletonMap(topic, 1));
         Node node = metadata.fetch().nodes().get(0);
 
-        PartitionAssignor assignor = new RoundRobinAssignor();
+        ConsumerPartitionAssignor assignor = new RoundRobinAssignor();
 
         KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor,
-                true, groupId);
+                true, groupId, groupInstanceId);
         consumer.assign(singletonList(tp0));
 
-        client.prepareResponseFrom(new FindCoordinatorResponse(Errors.NONE, node), node);
+        client.prepareResponseFrom(FindCoordinatorResponse.prepareResponse(Errors.NONE, node), node);
         Node coordinator = new Node(Integer.MAX_VALUE - node.id(), node.host(), node.port());
 
         // lookup committed offset and find nothing
@@ -593,13 +617,13 @@ public void testResetToCommittedOffset() {
         initMetadata(client, Collections.singletonMap(topic, 1));
         Node node = metadata.fetch().nodes().get(0);
 
-        PartitionAssignor assignor = new RoundRobinAssignor();
+        ConsumerPartitionAssignor assignor = new RoundRobinAssignor();
 
         KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor,
-                true, groupId);
+                true, groupId, groupInstanceId);
         consumer.assign(singletonList(tp0));
 
-        client.prepareResponseFrom(new FindCoordinatorResponse(Errors.NONE, node), node);
+        client.prepareResponseFrom(FindCoordinatorResponse.prepareResponse(Errors.NONE, node), node);
         Node coordinator = new Node(Integer.MAX_VALUE - node.id(), node.host(), node.port());
 
         client.prepareResponseFrom(offsetResponse(Collections.singletonMap(tp0, 539L), Errors.NONE), coordinator);
@@ -618,13 +642,13 @@ public void testResetUsingAutoResetPolicy() {
         initMetadata(client, Collections.singletonMap(topic, 1));
         Node node = metadata.fetch().nodes().get(0);
 
-        PartitionAssignor assignor = new RoundRobinAssignor();
+        ConsumerPartitionAssignor assignor = new RoundRobinAssignor();
 
         KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor,
-                true, groupId);
+                true, groupId, groupInstanceId);
         consumer.assign(singletonList(tp0));
 
-        client.prepareResponseFrom(new FindCoordinatorResponse(Errors.NONE, node), node);
+        client.prepareResponseFrom(FindCoordinatorResponse.prepareResponse(Errors.NONE, node), node);
         Node coordinator = new Node(Integer.MAX_VALUE - node.id(), node.host(), node.port());
 
         client.prepareResponseFrom(offsetResponse(Collections.singletonMap(tp0, -1L), Errors.NONE), coordinator);
@@ -635,6 +659,25 @@ public void testResetUsingAutoResetPolicy() {
         assertEquals(50L, consumer.position(tp0));
     }
 
+    @Test
+    public void testOffsetIsValidAfterSeek() {
+        Time time = new MockTime();
+        SubscriptionState subscription = new SubscriptionState(new LogContext(), OffsetResetStrategy.LATEST);
+        ConsumerMetadata metadata = createMetadata(subscription);
+        MockClient client = new MockClient(time, metadata);
+
+        initMetadata(client, Collections.singletonMap(topic, 1));
+
+        ConsumerPartitionAssignor assignor = new RoundRobinAssignor();
+
+        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor,
+                true, groupId, Optional.empty());
+        consumer.assign(singletonList(tp0));
+        consumer.seek(tp0, 20L);
+        consumer.poll(Duration.ZERO);
+        assertEquals(subscription.validPosition(tp0).offset, 20L);
+    }
+
     @Test
     public void testCommitsFetchedDuringAssign() {
         long offset1 = 10000;
@@ -648,18 +691,18 @@ public void testCommitsFetchedDuringAssign() {
         initMetadata(client, Collections.singletonMap(topic, 2));
         Node node = metadata.fetch().nodes().get(0);
 
-        PartitionAssignor assignor = new RoundRobinAssignor();
+        ConsumerPartitionAssignor assignor = new RoundRobinAssignor();
 
-        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, true);
+        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, true, groupInstanceId);
         consumer.assign(singletonList(tp0));
 
         // lookup coordinator
-        client.prepareResponseFrom(new FindCoordinatorResponse(Errors.NONE, node), node);
+        client.prepareResponseFrom(FindCoordinatorResponse.prepareResponse(Errors.NONE, node), node);
         Node coordinator = new Node(Integer.MAX_VALUE - node.id(), node.host(), node.port());
 
         // fetch offset for one topic
         client.prepareResponseFrom(offsetResponse(Collections.singletonMap(tp0, offset1), Errors.NONE), coordinator);
-        assertEquals(offset1, consumer.committed(tp0).offset());
+        assertEquals(offset1, consumer.committed(Collections.singleton(tp0)).get(tp0).offset());
 
         consumer.assign(Arrays.asList(tp0, tp1));
 
@@ -667,12 +710,43 @@ public void testCommitsFetchedDuringAssign() {
         Map offsets = new HashMap<>();
         offsets.put(tp0, offset1);
         client.prepareResponseFrom(offsetResponse(offsets, Errors.NONE), coordinator);
-        assertEquals(offset1, consumer.committed(tp0).offset());
+        assertEquals(offset1, consumer.committed(Collections.singleton(tp0)).get(tp0).offset());
 
         offsets.remove(tp0);
         offsets.put(tp1, offset2);
         client.prepareResponseFrom(offsetResponse(offsets, Errors.NONE), coordinator);
-        assertEquals(offset2, consumer.committed(tp1).offset());
+        assertEquals(offset2, consumer.committed(Collections.singleton(tp1)).get(tp1).offset());
+        consumer.close(Duration.ofMillis(0));
+    }
+
+    @Test
+    public void testNoCommittedOffsets() {
+        long offset1 = 10000;
+
+        Time time = new MockTime();
+        SubscriptionState subscription = new SubscriptionState(new LogContext(), OffsetResetStrategy.EARLIEST);
+        ConsumerMetadata metadata = createMetadata(subscription);
+        MockClient client = new MockClient(time, metadata);
+
+        initMetadata(client, Collections.singletonMap(topic, 2));
+        Node node = metadata.fetch().nodes().get(0);
+
+        ConsumerPartitionAssignor assignor = new RoundRobinAssignor();
+
+        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, true, groupInstanceId);
+        consumer.assign(Arrays.asList(tp0, tp1));
+
+        // lookup coordinator
+        client.prepareResponseFrom(FindCoordinatorResponse.prepareResponse(Errors.NONE, node), node);
+        Node coordinator = new Node(Integer.MAX_VALUE - node.id(), node.host(), node.port());
+
+        // fetch offset for one topic
+        client.prepareResponseFrom(offsetResponse(Utils.mkMap(Utils.mkEntry(tp0, offset1), Utils.mkEntry(tp1, -1L)), Errors.NONE), coordinator);
+        final Map committed = consumer.committed(Utils.mkSet(tp0, tp1));
+        assertEquals(2, committed.size());
+        assertEquals(offset1, committed.get(tp0).offset());
+        assertNull(committed.get(tp1));
+
         consumer.close(Duration.ofMillis(0));
     }
 
@@ -686,9 +760,9 @@ public void testAutoCommitSentBeforePositionUpdate() {
         initMetadata(client, Collections.singletonMap(topic, 1));
         Node node = metadata.fetch().nodes().get(0);
 
-        PartitionAssignor assignor = new RoundRobinAssignor();
+        ConsumerPartitionAssignor assignor = new RoundRobinAssignor();
 
-        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, true);
+        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, true, groupInstanceId);
         consumer.subscribe(singleton(topic), getConsumerRebalanceListener(consumer));
         Node coordinator = prepareRebalance(client, node, assignor, singletonList(tp0), null);
 
@@ -726,9 +800,9 @@ public void testRegexSubscription() {
         initMetadata(client, partitionCounts);
         Node node = metadata.fetch().nodes().get(0);
 
-        PartitionAssignor assignor = new RoundRobinAssignor();
+        ConsumerPartitionAssignor assignor = new RoundRobinAssignor();
 
-        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, true);
+        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, true, groupInstanceId);
         prepareRebalance(client, node, singleton(topic), assignor, singletonList(tp0), null);
 
         consumer.subscribe(Pattern.compile(topic), getConsumerRebalanceListener(consumer));
@@ -744,7 +818,7 @@ public void testRegexSubscription() {
 
     @Test
     public void testChangingRegexSubscription() {
-        PartitionAssignor assignor = new RoundRobinAssignor();
+        ConsumerPartitionAssignor assignor = new RoundRobinAssignor();
 
         String otherTopic = "other";
         TopicPartition otherTopicPartition = new TopicPartition(otherTopic, 0);
@@ -760,7 +834,7 @@ public void testChangingRegexSubscription() {
         initMetadata(client, partitionCounts);
         Node node = metadata.fetch().nodes().get(0);
 
-        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, false);
+        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, false, groupInstanceId);
 
         Node coordinator = prepareRebalance(client, node, singleton(topic), assignor, singletonList(tp0), null);
         consumer.subscribe(Pattern.compile(topic), getConsumerRebalanceListener(consumer));
@@ -790,9 +864,9 @@ public void testWakeupWithFetchDataAvailable() throws Exception {
         initMetadata(client, Collections.singletonMap(topic, 1));
         Node node = metadata.fetch().nodes().get(0);
 
-        PartitionAssignor assignor = new RoundRobinAssignor();
+        ConsumerPartitionAssignor assignor = new RoundRobinAssignor();
 
-        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, true);
+        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, true, groupInstanceId);
         consumer.subscribe(singleton(topic), getConsumerRebalanceListener(consumer));
         prepareRebalance(client, node, assignor, singletonList(tp0), null);
 
@@ -840,9 +914,9 @@ public void testPollThrowsInterruptExceptionIfInterrupted() {
         initMetadata(client, Collections.singletonMap(topic, 1));
         Node node = metadata.fetch().nodes().get(0);
 
-        final PartitionAssignor assignor = new RoundRobinAssignor();
+        final ConsumerPartitionAssignor assignor = new RoundRobinAssignor();
 
-        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, false);
+        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, false, groupInstanceId);
         consumer.subscribe(singleton(topic), getConsumerRebalanceListener(consumer));
         prepareRebalance(client, node, assignor, singletonList(tp0), null);
 
@@ -870,9 +944,9 @@ public void fetchResponseWithUnexpectedPartitionIsIgnored() {
         initMetadata(client, Collections.singletonMap(topic, 1));
         Node node = metadata.fetch().nodes().get(0);
 
-        PartitionAssignor assignor = new RangeAssignor();
+        ConsumerPartitionAssignor assignor = new RangeAssignor();
 
-        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, true);
+        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, true, groupInstanceId);
         consumer.subscribe(singletonList(topic), getConsumerRebalanceListener(consumer));
 
         prepareRebalance(client, node, assignor, singletonList(tp0), null);
@@ -910,9 +984,9 @@ public void testSubscriptionChangesWithAutoCommitEnabled() {
         initMetadata(client, tpCounts);
         Node node = metadata.fetch().nodes().get(0);
 
-        PartitionAssignor assignor = new RangeAssignor();
+        ConsumerPartitionAssignor assignor = new RangeAssignor();
 
-        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, true);
+        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, true, groupInstanceId);
 
         // initial subscription
         consumer.subscribe(Arrays.asList(topic, topic2), getConsumerRebalanceListener(consumer));
@@ -1024,16 +1098,11 @@ public void testSubscriptionChangesWithAutoCommitDisabled() {
         initMetadata(client, tpCounts);
         Node node = metadata.fetch().nodes().get(0);
 
-        PartitionAssignor assignor = new RangeAssignor();
-
-        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, false);
+        ConsumerPartitionAssignor assignor = new RangeAssignor();
 
-        // initial subscription
-        consumer.subscribe(singleton(topic), getConsumerRebalanceListener(consumer));
+        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, false, groupInstanceId);
 
-        // verify that subscription has changed but assignment is still unchanged
-        assertEquals(singleton(topic), consumer.subscription());
-        assertEquals(Collections.emptySet(), consumer.assignment());
+        initializeSubscriptionWithSingleTopic(consumer, getConsumerRebalanceListener(consumer));
 
         // mock rebalance responses
         prepareRebalance(client, node, assignor, singletonList(tp0), null);
@@ -1073,6 +1142,71 @@ public void testSubscriptionChangesWithAutoCommitDisabled() {
         consumer.close();
     }
 
+    @Test
+    public void testUnsubscribeShouldTriggerPartitionsRevokedWithValidGeneration() {
+        Time time = new MockTime();
+        SubscriptionState subscription = new SubscriptionState(new LogContext(), OffsetResetStrategy.EARLIEST);
+        ConsumerMetadata metadata = createMetadata(subscription);
+        MockClient client = new MockClient(time, metadata);
+
+        initMetadata(client, Collections.singletonMap(topic, 1));
+        Node node = metadata.fetch().nodes().get(0);
+
+        CooperativeStickyAssignor assignor = new CooperativeStickyAssignor();
+        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, false, groupInstanceId);
+
+        initializeSubscriptionWithSingleTopic(consumer, getExceptionConsumerRebalanceListener());
+
+        prepareRebalance(client, node, assignor, singletonList(tp0), null);
+
+        RuntimeException assignmentException = assertThrows(RuntimeException.class,
+            () -> consumer.updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE)));
+        assertEquals(partitionAssigned + singleTopicPartition, assignmentException.getCause().getMessage());
+
+        RuntimeException unsubscribeException = assertThrows(RuntimeException.class, consumer::unsubscribe);
+        assertEquals(partitionRevoked + singleTopicPartition, unsubscribeException.getCause().getMessage());
+    }
+
+    @Test
+    public void testUnsubscribeShouldTriggerPartitionsLostWithNoGeneration() throws Exception {
+        Time time = new MockTime();
+        SubscriptionState subscription = new SubscriptionState(new LogContext(), OffsetResetStrategy.EARLIEST);
+        ConsumerMetadata metadata = createMetadata(subscription);
+        MockClient client = new MockClient(time, metadata);
+
+        initMetadata(client, Collections.singletonMap(topic, 1));
+        Node node = metadata.fetch().nodes().get(0);
+
+        CooperativeStickyAssignor assignor = new CooperativeStickyAssignor();
+        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, false, groupInstanceId);
+
+        initializeSubscriptionWithSingleTopic(consumer, getExceptionConsumerRebalanceListener());
+        Node coordinator = prepareRebalance(client, node, assignor, singletonList(tp0), null);
+
+        RuntimeException assignException = assertThrows(RuntimeException.class,
+            () -> consumer.updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE)));
+        assertEquals(partitionAssigned + singleTopicPartition, assignException.getCause().getMessage());
+
+        AtomicBoolean heartbeatReceived = prepareHeartbeatResponse(client, coordinator, Errors.UNKNOWN_MEMBER_ID);
+
+        time.sleep(heartbeatIntervalMs);
+        TestUtils.waitForCondition(heartbeatReceived::get, "Heartbeat response did not occur within timeout.");
+
+        consumer.updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE));
+        assertTrue(heartbeatReceived.get());
+
+        RuntimeException unsubscribeException = assertThrows(RuntimeException.class, consumer::unsubscribe);
+        assertEquals(partitionLost + singleTopicPartition, unsubscribeException.getCause().getMessage());
+    }
+
+    private void initializeSubscriptionWithSingleTopic(KafkaConsumer consumer,
+                                                       ConsumerRebalanceListener consumerRebalanceListener) {
+        consumer.subscribe(singleton(topic), consumerRebalanceListener);
+        // verify that subscription has changed but assignment is still unchanged
+        assertEquals(singleton(topic), consumer.subscription());
+        assertEquals(Collections.emptySet(), consumer.assignment());
+    }
+
     @Test
     public void testManualAssignmentChangeWithAutoCommitEnabled() {
         Time time = new MockTime();
@@ -1086,12 +1220,12 @@ public void testManualAssignmentChangeWithAutoCommitEnabled() {
         initMetadata(client, tpCounts);
         Node node = metadata.fetch().nodes().get(0);
 
-        PartitionAssignor assignor = new RangeAssignor();
+        ConsumerPartitionAssignor assignor = new RangeAssignor();
 
-        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, true);
+        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, true, groupInstanceId);
 
         // lookup coordinator
-        client.prepareResponseFrom(new FindCoordinatorResponse(Errors.NONE, node), node);
+        client.prepareResponseFrom(FindCoordinatorResponse.prepareResponse(Errors.NONE, node), node);
         Node coordinator = new Node(Integer.MAX_VALUE - node.id(), node.host(), node.port());
 
         // manual assignment
@@ -1100,7 +1234,7 @@ public void testManualAssignmentChangeWithAutoCommitEnabled() {
 
         // fetch offset for one topic
         client.prepareResponseFrom(offsetResponse(Collections.singletonMap(tp0, 0L), Errors.NONE), coordinator);
-        assertEquals(0, consumer.committed(tp0).offset());
+        assertEquals(0, consumer.committed(Collections.singleton(tp0)).get(tp0).offset());
 
         // verify that assignment immediately changes
         assertTrue(consumer.assignment().equals(singleton(tp0)));
@@ -1142,12 +1276,12 @@ public void testManualAssignmentChangeWithAutoCommitDisabled() {
         initMetadata(client, tpCounts);
         Node node = metadata.fetch().nodes().get(0);
 
-        PartitionAssignor assignor = new RangeAssignor();
+        ConsumerPartitionAssignor assignor = new RangeAssignor();
 
-        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, false);
+        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, false, groupInstanceId);
 
         // lookup coordinator
-        client.prepareResponseFrom(new FindCoordinatorResponse(Errors.NONE, node), node);
+        client.prepareResponseFrom(FindCoordinatorResponse.prepareResponse(Errors.NONE, node), node);
         Node coordinator = new Node(Integer.MAX_VALUE - node.id(), node.host(), node.port());
 
         // manual assignment
@@ -1158,7 +1292,7 @@ public void testManualAssignmentChangeWithAutoCommitDisabled() {
         client.prepareResponseFrom(
                 offsetResponse(Collections.singletonMap(tp0, 0L), Errors.NONE),
                 coordinator);
-        assertEquals(0, consumer.committed(tp0).offset());
+        assertEquals(0, consumer.committed(Collections.singleton(tp0)).get(tp0).offset());
 
         // verify that assignment immediately changes
         assertTrue(consumer.assignment().equals(singleton(tp0)));
@@ -1196,12 +1330,12 @@ public void testOffsetOfPausedPartitions() {
         initMetadata(client, Collections.singletonMap(topic, 2));
         Node node = metadata.fetch().nodes().get(0);
 
-        PartitionAssignor assignor = new RangeAssignor();
+        ConsumerPartitionAssignor assignor = new RangeAssignor();
 
-        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, true);
+        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, true, groupInstanceId);
 
         // lookup coordinator
-        client.prepareResponseFrom(new FindCoordinatorResponse(Errors.NONE, node), node);
+        client.prepareResponseFrom(FindCoordinatorResponse.prepareResponse(Errors.NONE, node), node);
         Node coordinator = new Node(Integer.MAX_VALUE - node.id(), node.host(), node.port());
 
         // manual assignment
@@ -1219,12 +1353,12 @@ public void testOffsetOfPausedPartitions() {
         offsets.put(tp1, 0L);
 
         client.prepareResponseFrom(offsetResponse(offsets, Errors.NONE), coordinator);
-        assertEquals(0, consumer.committed(tp0).offset());
+        assertEquals(0, consumer.committed(Collections.singleton(tp0)).get(tp0).offset());
 
         offsets.remove(tp0);
         offsets.put(tp1, 0L);
         client.prepareResponseFrom(offsetResponse(offsets, Errors.NONE), coordinator);
-        assertEquals(0, consumer.committed(tp1).offset());
+        assertEquals(0, consumer.committed(Collections.singleton(tp1)).get(tp1).offset());
 
         // fetch and verify consumer's position in the two partitions
         final Map offsetResponse = new HashMap<>();
@@ -1295,7 +1429,7 @@ public void testCloseInterrupt() throws Exception {
     }
 
     @Test
-    public void closeShouldBeIdempotent() {
+    public void testCloseShouldBeIdempotent() {
         KafkaConsumer consumer = newConsumer((String) null);
         consumer.close();
         consumer.close();
@@ -1319,7 +1453,7 @@ public void testOperationsBySubscribingConsumerWithDefaultGroupId() {
         }
 
         try {
-            newConsumer((String) null).committed(tp0);
+            newConsumer((String) null).committed(Collections.singleton(tp0)).get(tp0);
             fail("Expected an InvalidGroupIdException");
         } catch (InvalidGroupIdException e) {
             // OK, expected
@@ -1346,7 +1480,7 @@ public void testOperationsByAssigningConsumerWithDefaultGroupId() {
         consumer.assign(singleton(tp0));
 
         try {
-            consumer.committed(tp0);
+            consumer.committed(Collections.singleton(tp0)).get(tp0);
             fail("Expected an InvalidGroupIdException");
         } catch (InvalidGroupIdException e) {
             // OK, expected
@@ -1382,7 +1516,7 @@ public void testMetricConfigRecordingLevel() {
     }
 
     @Test
-    public void shouldAttemptToRejoinGroupAfterSyncGroupFailed() throws Exception {
+    public void testShouldAttemptToRejoinGroupAfterSyncGroupFailed() throws Exception {
         Time time = new MockTime();
         SubscriptionState subscription = new SubscriptionState(new LogContext(), OffsetResetStrategy.EARLIEST);
         ConsumerMetadata metadata = createMetadata(subscription);
@@ -1391,11 +1525,11 @@ public void shouldAttemptToRejoinGroupAfterSyncGroupFailed() throws Exception {
         initMetadata(client, Collections.singletonMap(topic, 1));
         Node node = metadata.fetch().nodes().get(0);
 
-        PartitionAssignor assignor = new RoundRobinAssignor();
+        ConsumerPartitionAssignor assignor = new RoundRobinAssignor();
 
-        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, false);
+        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, false, groupInstanceId);
         consumer.subscribe(singleton(topic), getConsumerRebalanceListener(consumer));
-        client.prepareResponseFrom(new FindCoordinatorResponse(Errors.NONE, node), node);
+        client.prepareResponseFrom(FindCoordinatorResponse.prepareResponse(Errors.NONE, node), node);
         Node coordinator = new Node(Integer.MAX_VALUE - node.id(), node.host(), node.port());
 
 
@@ -1414,21 +1548,35 @@ public void shouldAttemptToRejoinGroupAfterSyncGroupFailed() throws Exception {
             public boolean matches(AbstractRequest body) {
                 return true;
             }
-        }, new HeartbeatResponse(Errors.REBALANCE_IN_PROGRESS), coordinator);
+        }, new HeartbeatResponse(
+                new HeartbeatResponseData().setErrorCode(Errors.REBALANCE_IN_PROGRESS.code())),
+                coordinator);
 
         // join group
-        final ByteBuffer byteBuffer = ConsumerProtocol.serializeSubscription(new PartitionAssignor.Subscription(singletonList(topic)));
+        final ByteBuffer byteBuffer = ConsumerProtocol.serializeSubscription(new ConsumerPartitionAssignor.Subscription(singletonList(topic)));
 
         // This member becomes the leader
-        final JoinGroupResponse leaderResponse = new JoinGroupResponse(Errors.NONE, 1, assignor.name(), "memberId", "memberId",
-                Collections.singletonMap("memberId", byteBuffer));
+        String memberId = "memberId";
+        final JoinGroupResponse leaderResponse = new JoinGroupResponse(
+                new JoinGroupResponseData()
+                        .setErrorCode(Errors.NONE.code())
+                        .setGenerationId(1).setProtocolName(assignor.name())
+                        .setLeader(memberId).setMemberId(memberId)
+                        .setMembers(Collections.singletonList(
+                                new JoinGroupResponseData.JoinGroupResponseMember()
+                                        .setMemberId(memberId)
+                                        .setMetadata(byteBuffer.array())
+                                )
+                        )
+        );
+
         client.prepareResponseFrom(leaderResponse, coordinator);
 
         // sync group fails due to disconnect
         client.prepareResponseFrom(syncGroupResponse(singletonList(tp0), Errors.NONE), coordinator, true);
 
         // should try and find the new coordinator
-        client.prepareResponseFrom(new FindCoordinatorResponse(Errors.NONE, node), node);
+        client.prepareResponseFrom(FindCoordinatorResponse.prepareResponse(Errors.NONE, node), node);
 
         // rejoin group
         client.prepareResponseFrom(joinGroupFollowerResponse(assignor, 1, "memberId", "leaderId", Errors.NONE), coordinator);
@@ -1460,9 +1608,9 @@ private void consumerCloseTest(final long closeTimeoutMs,
         initMetadata(client, Collections.singletonMap(topic, 1));
         Node node = metadata.fetch().nodes().get(0);
 
-        PartitionAssignor assignor = new RoundRobinAssignor();
+        ConsumerPartitionAssignor assignor = new RoundRobinAssignor();
 
-        final KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, false);
+        final KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, false, Optional.empty());
         consumer.subscribe(singleton(topic), getConsumerRebalanceListener(consumer));
         Node coordinator = prepareRebalance(client, node, assignor, singletonList(tp0), null);
 
@@ -1585,7 +1733,53 @@ public void testCommitSyncAuthenticationFailure() {
     @Test(expected = AuthenticationException.class)
     public void testCommittedAuthenticationFaiure() {
         final KafkaConsumer consumer = consumerWithPendingAuthenticationError();
-        consumer.committed(tp0);
+        consumer.committed(Collections.singleton(tp0)).get(tp0);
+    }
+
+    @Test
+    public void testRebalanceException() {
+        Time time = new MockTime();
+        SubscriptionState subscription = new SubscriptionState(new LogContext(), OffsetResetStrategy.EARLIEST);
+        ConsumerMetadata metadata = createMetadata(subscription);
+        MockClient client = new MockClient(time, metadata);
+
+        initMetadata(client, Collections.singletonMap(topic, 1));
+        Node node = metadata.fetch().nodes().get(0);
+
+        ConsumerPartitionAssignor assignor = new RoundRobinAssignor();
+
+        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, true, groupInstanceId);
+
+        consumer.subscribe(singleton(topic), getExceptionConsumerRebalanceListener());
+        Node coordinator = new Node(Integer.MAX_VALUE - node.id(), node.host(), node.port());
+
+        client.prepareResponseFrom(FindCoordinatorResponse.prepareResponse(Errors.NONE, node), node);
+        client.prepareResponseFrom(joinGroupFollowerResponse(assignor, 1, "memberId", "leaderId", Errors.NONE), coordinator);
+        client.prepareResponseFrom(syncGroupResponse(singletonList(tp0), Errors.NONE), coordinator);
+
+        // assign throws
+        try {
+            consumer.updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE));
+            fail("Should throw exception");
+        } catch (Throwable e) {
+            assertEquals(partitionAssigned + singleTopicPartition, e.getCause().getMessage());
+        }
+
+        // the assignment is still updated regardless of the exception
+        assertEquals(singleton(tp0), subscription.assignedPartitions());
+
+        // close's revoke throws
+        try {
+            consumer.close(Duration.ofMillis(0));
+            fail("Should throw exception");
+        } catch (Throwable e) {
+            assertEquals(partitionRevoked + singleTopicPartition, e.getCause().getCause().getMessage());
+        }
+
+        consumer.close(Duration.ofMillis(0));
+
+        // the assignment is still updated regardless of the exception
+        assertTrue(subscription.assignedPartitions().isEmpty());
     }
 
     private KafkaConsumer consumerWithPendingAuthenticationError() {
@@ -1597,10 +1791,10 @@ private KafkaConsumer consumerWithPendingAuthenticationError() {
         initMetadata(client, singletonMap(topic, 1));
         Node node = metadata.fetch().nodes().get(0);
 
-        PartitionAssignor assignor = new RangeAssignor();
+        ConsumerPartitionAssignor assignor = new RangeAssignor();
 
         client.createPendingAuthenticationError(node, 0);
-        return newConsumer(time, client, subscription, metadata, assignor, false);
+        return newConsumer(time, client, subscription, metadata, assignor, false, groupInstanceId);
     }
 
     private ConsumerRebalanceListener getConsumerRebalanceListener(final KafkaConsumer consumer) {
@@ -1618,15 +1812,34 @@ public void onPartitionsAssigned(Collection partitions) {
         };
     }
 
+    private ConsumerRebalanceListener getExceptionConsumerRebalanceListener() {
+        return new ConsumerRebalanceListener() {
+            @Override
+            public void onPartitionsRevoked(Collection partitions) {
+                throw new RuntimeException(partitionRevoked + partitions);
+            }
+
+            @Override
+            public void onPartitionsAssigned(Collection partitions) {
+                throw new RuntimeException(partitionAssigned + partitions);
+            }
+
+            @Override
+            public void onPartitionsLost(Collection partitions) {
+                throw new RuntimeException(partitionLost + partitions);
+            }
+        };
+    }
+
     private ConsumerMetadata createMetadata(SubscriptionState subscription) {
-        return new ConsumerMetadata(0, Long.MAX_VALUE, false, subscription,
-                new LogContext(), new ClusterResourceListeners());
+        return new ConsumerMetadata(0, Long.MAX_VALUE, false, false,
+                                    subscription, new LogContext(), new ClusterResourceListeners());
     }
 
-    private Node prepareRebalance(MockClient client, Node node, final Set subscribedTopics, PartitionAssignor assignor, List partitions, Node coordinator) {
+    private Node prepareRebalance(MockClient client, Node node, final Set subscribedTopics, ConsumerPartitionAssignor assignor, List partitions, Node coordinator) {
         if (coordinator == null) {
             // lookup coordinator
-            client.prepareResponseFrom(new FindCoordinatorResponse(Errors.NONE, node), node);
+            client.prepareResponseFrom(FindCoordinatorResponse.prepareResponse(Errors.NONE, node), node);
             coordinator = new Node(Integer.MAX_VALUE - node.id(), node.host(), node.port());
         }
 
@@ -1635,7 +1848,12 @@ private Node prepareRebalance(MockClient client, Node node, final Set su
             @Override
             public boolean matches(AbstractRequest body) {
                 JoinGroupRequest joinGroupRequest = (JoinGroupRequest) body;
-                PartitionAssignor.Subscription subscription = ConsumerProtocol.deserializeSubscription(joinGroupRequest.groupProtocols().get(0).metadata());
+                Iterator protocolIterator =
+                        joinGroupRequest.data().protocols().iterator();
+                assertTrue(protocolIterator.hasNext());
+
+                ByteBuffer protocolMetadata = ByteBuffer.wrap(protocolIterator.next().metadata());
+                ConsumerPartitionAssignor.Subscription subscription = ConsumerProtocol.deserializeSubscription(protocolMetadata);
                 return subscribedTopics.equals(new HashSet<>(subscription.topics()));
             }
         }, joinGroupFollowerResponse(assignor, 1, "memberId", "leaderId", Errors.NONE), coordinator);
@@ -1646,10 +1864,10 @@ public boolean matches(AbstractRequest body) {
         return coordinator;
     }
 
-    private Node prepareRebalance(MockClient client, Node node, PartitionAssignor assignor, List partitions, Node coordinator) {
+    private Node prepareRebalance(MockClient client, Node node, ConsumerPartitionAssignor assignor, List partitions, Node coordinator) {
         if (coordinator == null) {
             // lookup coordinator
-            client.prepareResponseFrom(new FindCoordinatorResponse(Errors.NONE, node), node);
+            client.prepareResponseFrom(FindCoordinatorResponse.prepareResponse(Errors.NONE, node), node);
             coordinator = new Node(Integer.MAX_VALUE - node.id(), node.host(), node.port());
         }
 
@@ -1662,7 +1880,7 @@ private Node prepareRebalance(MockClient client, Node node, PartitionAssignor as
         return coordinator;
     }
 
-    private AtomicBoolean prepareHeartbeatResponse(MockClient client, Node coordinator) {
+    private AtomicBoolean prepareHeartbeatResponse(MockClient client, Node coordinator, Errors error) {
         final AtomicBoolean heartbeatReceived = new AtomicBoolean(false);
         client.prepareResponseFrom(new MockClient.RequestMatcher() {
             @Override
@@ -1670,7 +1888,7 @@ public boolean matches(AbstractRequest body) {
                 heartbeatReceived.set(true);
                 return true;
             }
-        }, new HeartbeatResponse(Errors.NONE), coordinator);
+        }, new HeartbeatResponse(new HeartbeatResponseData().setErrorCode(error.code())), coordinator);
         return heartbeatReceived;
     }
 
@@ -1684,10 +1902,11 @@ private AtomicBoolean prepareOffsetCommitResponse(MockClient client, Node coordi
             @Override
             public boolean matches(AbstractRequest body) {
                 OffsetCommitRequest commitRequest = (OffsetCommitRequest) body;
+                Map commitErrors = commitRequest.offsets();
+
                 for (Map.Entry partitionOffset : partitionOffsets.entrySet()) {
-                    OffsetCommitRequest.PartitionData partitionData = commitRequest.offsetData().get(partitionOffset.getKey());
                     // verify that the expected offset has been committed
-                    if (partitionData.offset != partitionOffset.getValue()) {
+                    if (!commitErrors.get(partitionOffset.getKey()).equals(partitionOffset.getValue())) {
                         commitReceived.set(false);
                         return false;
                     }
@@ -1706,14 +1925,25 @@ private OffsetCommitResponse offsetCommitResponse(Map re
         return new OffsetCommitResponse(responseData);
     }
 
-    private JoinGroupResponse joinGroupFollowerResponse(PartitionAssignor assignor, int generationId, String memberId, String leaderId, Errors error) {
-        return new JoinGroupResponse(error, generationId, assignor.name(), memberId, leaderId,
-                Collections.emptyMap());
+    private JoinGroupResponse joinGroupFollowerResponse(ConsumerPartitionAssignor assignor, int generationId, String memberId, String leaderId, Errors error) {
+        return new JoinGroupResponse(
+                new JoinGroupResponseData()
+                        .setErrorCode(error.code())
+                        .setGenerationId(generationId)
+                        .setProtocolName(assignor.name())
+                        .setLeader(leaderId)
+                        .setMemberId(memberId)
+                        .setMembers(Collections.emptyList())
+        );
     }
 
     private SyncGroupResponse syncGroupResponse(List partitions, Errors error) {
-        ByteBuffer buf = ConsumerProtocol.serializeAssignment(new PartitionAssignor.Assignment(partitions));
-        return new SyncGroupResponse(error, buf);
+        ByteBuffer buf = ConsumerProtocol.serializeAssignment(new ConsumerPartitionAssignor.Assignment(partitions));
+        return new SyncGroupResponse(
+                new SyncGroupResponseData()
+                        .setErrorCode(error.code())
+                        .setAssignment(Utils.toArray(buf))
+        );
     }
 
     private OffsetFetchResponse offsetResponse(Map offsets, Errors error) {
@@ -1747,7 +1977,6 @@ private ListOffsetResponse listOffsetsResponse(Map partiti
         return new ListOffsetResponse(partitionData);
     }
 
-
     private FetchResponse fetchResponse(Map fetches) {
         LinkedHashMap> tpResponses = new LinkedHashMap<>();
         for (Map.Entry fetchEntry : fetches.entrySet()) {
@@ -1780,25 +2009,27 @@ private KafkaConsumer newConsumer(Time time,
                                                       KafkaClient client,
                                                       SubscriptionState subscription,
                                                       ConsumerMetadata metadata,
-                                                      PartitionAssignor assignor,
-                                                      boolean autoCommitEnabled) {
-        return newConsumer(time, client, subscription, metadata, assignor, autoCommitEnabled, groupId);
+                                                      ConsumerPartitionAssignor assignor,
+                                                      boolean autoCommitEnabled,
+                                                      Optional groupInstanceId) {
+        return newConsumer(time, client, subscription, metadata, assignor, autoCommitEnabled, groupId, groupInstanceId);
     }
 
     private KafkaConsumer newConsumerNoAutoCommit(Time time,
                                                                   KafkaClient client,
                                                                   SubscriptionState subscription,
                                                                   ConsumerMetadata metadata) {
-        return newConsumer(time, client, subscription, metadata, new RangeAssignor(), false, groupId);
+        return newConsumer(time, client, subscription, metadata, new RangeAssignor(), false, groupId, groupInstanceId);
     }
 
     private KafkaConsumer newConsumer(Time time,
                                                       KafkaClient client,
                                                       SubscriptionState subscription,
                                                       ConsumerMetadata metadata,
-                                                      PartitionAssignor assignor,
+                                                      ConsumerPartitionAssignor assignor,
                                                       boolean autoCommitEnabled,
-                                                      String groupId) {
+                                                      String groupId,
+                                                      Optional groupInstanceId) {
         String clientId = "mock-consumer";
         String metricGroupPrefix = "consumer";
         long retryBackoffMs = 100;
@@ -1810,41 +2041,41 @@ private KafkaConsumer newConsumer(Time time,
         int fetchSize = 1024 * 1024;
         int maxPollRecords = Integer.MAX_VALUE;
         boolean checkCrcs = true;
+        boolean usePassthrough = false;
         int rebalanceTimeoutMs = 60000;
 
         Deserializer keyDeserializer = new StringDeserializer();
         Deserializer valueDeserializer = new StringDeserializer();
 
-        List assignors = singletonList(assignor);
+        List assignors = singletonList(assignor);
         ConsumerInterceptors interceptors = new ConsumerInterceptors<>(Collections.emptyList());
 
-        Metrics metrics = new Metrics();
+        Metrics metrics = new Metrics(time);
         ConsumerMetrics metricsRegistry = new ConsumerMetrics(metricGroupPrefix);
 
         LogContext loggerFactory = new LogContext();
         ConsumerNetworkClient consumerClient = new ConsumerNetworkClient(loggerFactory, client, metadata, time,
                 retryBackoffMs, requestTimeoutMs, heartbeatIntervalMs);
 
-        Heartbeat heartbeat = new Heartbeat(time, sessionTimeoutMs, heartbeatIntervalMs, rebalanceTimeoutMs, retryBackoffMs);
-        ConsumerCoordinator consumerCoordinator = new ConsumerCoordinator(
-                loggerFactory,
-                consumerClient,
-                groupId,
-                rebalanceTimeoutMs,
-                sessionTimeoutMs,
-                heartbeat,
-                assignors,
-                metadata,
-                subscription,
-                metrics,
-                metricGroupPrefix,
-                time,
-                retryBackoffMs,
-                autoCommitEnabled,
-                autoCommitIntervalMs,
-                interceptors,
-                true);
-
+        GroupRebalanceConfig rebalanceConfig = new GroupRebalanceConfig(sessionTimeoutMs,
+                                                                        rebalanceTimeoutMs,
+                                                                        heartbeatIntervalMs,
+                                                                        groupId,
+                                                                        groupInstanceId,
+                                                                        retryBackoffMs,
+                                                                        true);
+        ConsumerCoordinator consumerCoordinator = new ConsumerCoordinator(rebalanceConfig,
+                                                                          loggerFactory,
+                                                                          consumerClient,
+                                                                          assignors,
+                                                                          metadata,
+                                                                          subscription,
+                                                                          metrics,
+                                                                          metricGroupPrefix,
+                                                                          time,
+                                                                          autoCommitEnabled,
+                                                                          autoCommitIntervalMs,
+                                                                          interceptors);
         Fetcher fetcher = new Fetcher<>(
                 loggerFactory,
                 consumerClient,
@@ -1854,6 +2085,8 @@ private KafkaConsumer newConsumer(Time time,
                 fetchSize,
                 maxPollRecords,
                 checkCrcs,
+                "",
+                usePassthrough,
                 keyDeserializer,
                 valueDeserializer,
                 metadata,
@@ -1863,7 +2096,8 @@ private KafkaConsumer newConsumer(Time time,
                 time,
                 retryBackoffMs,
                 requestTimeoutMs,
-                IsolationLevel.READ_UNCOMMITTED);
+                IsolationLevel.READ_UNCOMMITTED,
+                new ApiVersions());
 
         return new KafkaConsumer<>(
                 loggerFactory,
@@ -1914,22 +2148,140 @@ public void testSubscriptionOnInvalidTopic() {
         initMetadata(client, Collections.singletonMap(topic, 1));
         Cluster cluster = metadata.fetch();
 
-        PartitionAssignor assignor = new RoundRobinAssignor();
+        ConsumerPartitionAssignor assignor = new RoundRobinAssignor();
 
         String invalidTopicName = "topic abc";  // Invalid topic name due to space
 
         List topicMetadata = new ArrayList<>();
         topicMetadata.add(new MetadataResponse.TopicMetadata(Errors.INVALID_TOPIC_EXCEPTION,
                 invalidTopicName, false, Collections.emptyList()));
-        MetadataResponse updateResponse = new MetadataResponse(cluster.nodes(),
+        MetadataResponse updateResponse = MetadataResponse.prepareResponse(cluster.nodes(),
                 cluster.clusterResource().clusterId(),
                 cluster.controller().id(),
                 topicMetadata);
         client.prepareMetadataUpdate(updateResponse);
 
-        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, true);
+        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, true, groupInstanceId);
         consumer.subscribe(singleton(invalidTopicName), getConsumerRebalanceListener(consumer));
 
         consumer.poll(Duration.ZERO);
     }
+
+    @Test
+    public void testPollTimeMetrics() {
+        Time time = new MockTime();
+        SubscriptionState subscription = new SubscriptionState(new LogContext(), OffsetResetStrategy.EARLIEST);
+        ConsumerMetadata metadata = createMetadata(subscription);
+        MockClient client = new MockClient(time, metadata);
+        initMetadata(client, Collections.singletonMap(topic, 1));
+
+        ConsumerPartitionAssignor assignor = new RoundRobinAssignor();
+
+        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, true, groupInstanceId);
+        consumer.subscribe(singletonList(topic));
+        // MetricName objects to check
+        Metrics metrics = consumer.metrics;
+        MetricName lastPollSecondsAgoName = metrics.metricName("last-poll-seconds-ago", "consumer-metrics");
+        MetricName timeBetweenPollAvgName = metrics.metricName("time-between-poll-avg", "consumer-metrics");
+        MetricName timeBetweenPollMaxName = metrics.metricName("time-between-poll-max", "consumer-metrics");
+        // Test default values
+        assertEquals(-1.0d, consumer.metrics().get(lastPollSecondsAgoName).metricValue());
+        assertEquals(Double.NaN, consumer.metrics().get(timeBetweenPollAvgName).metricValue());
+        assertEquals(Double.NaN, consumer.metrics().get(timeBetweenPollMaxName).metricValue());
+        // Call first poll
+        consumer.poll(Duration.ZERO);
+        assertEquals(0.0d, consumer.metrics().get(lastPollSecondsAgoName).metricValue());
+        assertEquals(0.0d, consumer.metrics().get(timeBetweenPollAvgName).metricValue());
+        assertEquals(0.0d, consumer.metrics().get(timeBetweenPollMaxName).metricValue());
+        // Advance time by 5,000 (total time = 5,000)
+        time.sleep(5 * 1000L);
+        assertEquals(5.0d, consumer.metrics().get(lastPollSecondsAgoName).metricValue());
+        // Call second poll
+        consumer.poll(Duration.ZERO);
+        assertEquals(2.5 * 1000d, consumer.metrics().get(timeBetweenPollAvgName).metricValue());
+        assertEquals(5 * 1000d, consumer.metrics().get(timeBetweenPollMaxName).metricValue());
+        // Advance time by 10,000 (total time = 15,000)
+        time.sleep(10 * 1000L);
+        assertEquals(10.0d, consumer.metrics().get(lastPollSecondsAgoName).metricValue());
+        // Call third poll
+        consumer.poll(Duration.ZERO);
+        assertEquals(5 * 1000d, consumer.metrics().get(timeBetweenPollAvgName).metricValue());
+        assertEquals(10 * 1000d, consumer.metrics().get(timeBetweenPollMaxName).metricValue());
+        // Advance time by 5,000 (total time = 20,000)
+        time.sleep(5 * 1000L);
+        assertEquals(5.0d, consumer.metrics().get(lastPollSecondsAgoName).metricValue());
+        // Call fourth poll
+        consumer.poll(Duration.ZERO);
+        assertEquals(5 * 1000d, consumer.metrics().get(timeBetweenPollAvgName).metricValue());
+        assertEquals(10 * 1000d, consumer.metrics().get(timeBetweenPollMaxName).metricValue());
+    }
+
+    @Test
+    public void testPollIdleRatio() {
+        Time time = new MockTime();
+        SubscriptionState subscription = new SubscriptionState(new LogContext(), OffsetResetStrategy.EARLIEST);
+        ConsumerMetadata metadata = createMetadata(subscription);
+        MockClient client = new MockClient(time, metadata);
+        initMetadata(client, Collections.singletonMap(topic, 1));
+
+        ConsumerPartitionAssignor assignor = new RoundRobinAssignor();
+
+        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, true, groupInstanceId);
+        // MetricName object to check
+        Metrics metrics = consumer.metrics;
+        MetricName pollIdleRatio = metrics.metricName("poll-idle-ratio-avg", "consumer-metrics");
+        // Test default value
+        assertEquals(Double.NaN, consumer.metrics().get(pollIdleRatio).metricValue());
+
+        // 1st poll
+        // Spend 50ms in poll so value = 1.0
+        consumer.kafkaConsumerMetrics.recordPollStart(time.milliseconds());
+        time.sleep(50);
+        consumer.kafkaConsumerMetrics.recordPollEnd(time.milliseconds());
+
+        assertEquals(1.0d, consumer.metrics().get(pollIdleRatio).metricValue());
+
+        // 2nd poll
+        // Spend 50m outside poll and 0ms in poll so value = 0.0
+        time.sleep(50);
+        consumer.kafkaConsumerMetrics.recordPollStart(time.milliseconds());
+        consumer.kafkaConsumerMetrics.recordPollEnd(time.milliseconds());
+
+        // Avg of first two data points
+        assertEquals((1.0d + 0.0d) / 2, consumer.metrics().get(pollIdleRatio).metricValue());
+
+        // 3rd poll
+        // Spend 25ms outside poll and 25ms in poll so value = 0.5
+        time.sleep(25);
+        consumer.kafkaConsumerMetrics.recordPollStart(time.milliseconds());
+        time.sleep(25);
+        consumer.kafkaConsumerMetrics.recordPollEnd(time.milliseconds());
+
+        // Avg of three data points
+        assertEquals((1.0d + 0.0d + 0.5d) / 3, consumer.metrics().get(pollIdleRatio).metricValue());
+    }
+
+    private static boolean consumerMetricPresent(KafkaConsumer consumer, String name) {
+        MetricName metricName = new MetricName(name, "consumer-metrics", "", Collections.emptyMap());
+        return consumer.metrics.metrics().containsKey(metricName);
+    }
+
+    @Test
+    public void testClosingConsumerUnregistersConsumerMetrics() {
+        Time time = new MockTime();
+        SubscriptionState subscription = new SubscriptionState(new LogContext(), OffsetResetStrategy.EARLIEST);
+        ConsumerMetadata metadata = createMetadata(subscription);
+        MockClient client = new MockClient(time, metadata);
+        initMetadata(client, Collections.singletonMap(topic, 1));
+        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata,
+            new RoundRobinAssignor(), true, groupInstanceId);
+        consumer.subscribe(singletonList(topic));
+        assertTrue(consumerMetricPresent(consumer, "last-poll-seconds-ago"));
+        assertTrue(consumerMetricPresent(consumer, "time-between-poll-avg"));
+        assertTrue(consumerMetricPresent(consumer, "time-between-poll-max"));
+        consumer.close();
+        assertFalse(consumerMetricPresent(consumer, "last-poll-seconds-ago"));
+        assertFalse(consumerMetricPresent(consumer, "time-between-poll-avg"));
+        assertFalse(consumerMetricPresent(consumer, "time-between-poll-max"));
+    }
 }
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/MockConsumerTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/MockConsumerTest.java
index aad4d2973a032..5a012b2cf67a9 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/MockConsumerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/MockConsumerTest.java
@@ -55,9 +55,10 @@ public void testSimpleMock() {
         assertEquals(rec1, iter.next());
         assertEquals(rec2, iter.next());
         assertFalse(iter.hasNext());
-        assertEquals(2L, consumer.position(new TopicPartition("test", 0)));
+        final TopicPartition tp = new TopicPartition("test", 0);
+        assertEquals(2L, consumer.position(tp));
         consumer.commitSync();
-        assertEquals(2L, consumer.committed(new TopicPartition("test", 0)).offset());
+        assertEquals(2L, consumer.committed(Collections.singleton(tp)).get(tp).offset());
     }
 
     @SuppressWarnings("deprecation")
@@ -81,9 +82,10 @@ public void testSimpleMockDeprecated() {
         assertEquals(rec1, iter.next());
         assertEquals(rec2, iter.next());
         assertFalse(iter.hasNext());
-        assertEquals(2L, consumer.position(new TopicPartition("test", 0)));
+        final TopicPartition tp = new TopicPartition("test", 0);
+        assertEquals(2L, consumer.position(tp));
         consumer.commitSync();
-        assertEquals(2L, consumer.committed(new TopicPartition("test", 0)).offset());
+        assertEquals(2L, consumer.committed(Collections.singleton(tp)).get(tp).offset());
     }
 
     @Test
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/RangeAssignorTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/RangeAssignorTest.java
index 8158f54a10e4e..e5c5073afee00 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/RangeAssignorTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/RangeAssignorTest.java
@@ -16,16 +16,21 @@
  */
 package org.apache.kafka.clients.consumer;
 
-import org.apache.kafka.clients.consumer.internals.PartitionAssignor.Subscription;
+import org.apache.kafka.clients.consumer.ConsumerPartitionAssignor.Subscription;
+import org.apache.kafka.clients.consumer.internals.AbstractPartitionAssignor;
+import org.apache.kafka.clients.consumer.internals.AbstractPartitionAssignor.MemberInfo;
 import org.apache.kafka.common.TopicPartition;
+import org.junit.Before;
 import org.junit.Test;
 
+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 static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
@@ -34,128 +39,115 @@ public class RangeAssignorTest {
 
     private RangeAssignor assignor = new RangeAssignor();
 
+    // For plural tests
+    private String topic1 = "topic1";
+    private String topic2 = "topic2";
+    private final String consumer1 = "consumer1";
+    private final String instance1 = "instance1";
+    private final String consumer2 = "consumer2";
+    private final String instance2 = "instance2";
+    private final String consumer3 = "consumer3";
+    private final String instance3 = "instance3";
+
+    private List staticMemberInfos;
+
+    @Before
+    public void setUp() {
+        staticMemberInfos = new ArrayList<>();
+        staticMemberInfos.add(new MemberInfo(consumer1, Optional.of(instance1)));
+        staticMemberInfos.add(new MemberInfo(consumer2, Optional.of(instance2)));
+        staticMemberInfos.add(new MemberInfo(consumer3, Optional.of(instance3)));
+    }
 
     @Test
     public void testOneConsumerNoTopic() {
-        String consumerId = "consumer";
-
         Map partitionsPerTopic = new HashMap<>();
 
         Map> assignment = assignor.assign(partitionsPerTopic,
-                Collections.singletonMap(consumerId, new Subscription(Collections.emptyList())));
+                Collections.singletonMap(consumer1, new Subscription(Collections.emptyList())));
 
-        assertEquals(Collections.singleton(consumerId), assignment.keySet());
-        assertTrue(assignment.get(consumerId).isEmpty());
+        assertEquals(Collections.singleton(consumer1), assignment.keySet());
+        assertTrue(assignment.get(consumer1).isEmpty());
     }
 
     @Test
     public void testOneConsumerNonexistentTopic() {
-        String topic = "topic";
-        String consumerId = "consumer";
-
         Map partitionsPerTopic = new HashMap<>();
         Map> assignment = assignor.assign(partitionsPerTopic,
-                Collections.singletonMap(consumerId, new Subscription(topics(topic))));
-        assertEquals(Collections.singleton(consumerId), assignment.keySet());
-        assertTrue(assignment.get(consumerId).isEmpty());
+                Collections.singletonMap(consumer1, new Subscription(topics(topic1))));
+        assertEquals(Collections.singleton(consumer1), assignment.keySet());
+        assertTrue(assignment.get(consumer1).isEmpty());
     }
 
     @Test
     public void testOneConsumerOneTopic() {
-        String topic = "topic";
-        String consumerId = "consumer";
-
         Map partitionsPerTopic = new HashMap<>();
-        partitionsPerTopic.put(topic, 3);
+        partitionsPerTopic.put(topic1, 3);
 
         Map> assignment = assignor.assign(partitionsPerTopic,
-                Collections.singletonMap(consumerId, new Subscription(topics(topic))));
+                Collections.singletonMap(consumer1, new Subscription(topics(topic1))));
 
-        assertEquals(Collections.singleton(consumerId), assignment.keySet());
-        assertAssignment(partitions(tp(topic, 0), tp(topic, 1), tp(topic, 2)), assignment.get(consumerId));
+        assertEquals(Collections.singleton(consumer1), assignment.keySet());
+        assertAssignment(partitions(tp(topic1, 0), tp(topic1, 1), tp(topic1, 2)), assignment.get(consumer1));
     }
 
     @Test
     public void testOnlyAssignsPartitionsFromSubscribedTopics() {
-        String topic = "topic";
         String otherTopic = "other";
-        String consumerId = "consumer";
 
         Map partitionsPerTopic = new HashMap<>();
-        partitionsPerTopic.put(topic, 3);
+        partitionsPerTopic.put(topic1, 3);
         partitionsPerTopic.put(otherTopic, 3);
 
         Map> assignment = assignor.assign(partitionsPerTopic,
-                Collections.singletonMap(consumerId, new Subscription(topics(topic))));
-        assertEquals(Collections.singleton(consumerId), assignment.keySet());
-        assertAssignment(partitions(tp(topic, 0), tp(topic, 1), tp(topic, 2)), assignment.get(consumerId));
+                Collections.singletonMap(consumer1, new Subscription(topics(topic1))));
+        assertEquals(Collections.singleton(consumer1), assignment.keySet());
+        assertAssignment(partitions(tp(topic1, 0), tp(topic1, 1), tp(topic1, 2)), assignment.get(consumer1));
     }
 
     @Test
     public void testOneConsumerMultipleTopics() {
-        String topic1 = "topic1";
-        String topic2 = "topic2";
-        String consumerId = "consumer";
-
-        Map partitionsPerTopic = new HashMap<>();
-        partitionsPerTopic.put(topic1, 1);
-        partitionsPerTopic.put(topic2, 2);
+        Map partitionsPerTopic = setupPartitionsPerTopicWithTwoTopics(1, 2);
 
         Map> assignment = assignor.assign(partitionsPerTopic,
-                Collections.singletonMap(consumerId, new Subscription(topics(topic1, topic2))));
+                Collections.singletonMap(consumer1, new Subscription(topics(topic1, topic2))));
 
-        assertEquals(Collections.singleton(consumerId), assignment.keySet());
-        assertAssignment(partitions(tp(topic1, 0), tp(topic2, 0), tp(topic2, 1)), assignment.get(consumerId));
+        assertEquals(Collections.singleton(consumer1), assignment.keySet());
+        assertAssignment(partitions(tp(topic1, 0), tp(topic2, 0), tp(topic2, 1)), assignment.get(consumer1));
     }
 
     @Test
     public void testTwoConsumersOneTopicOnePartition() {
-        String topic = "topic";
-        String consumer1 = "consumer1";
-        String consumer2 = "consumer2";
-
         Map partitionsPerTopic = new HashMap<>();
-        partitionsPerTopic.put(topic, 1);
+        partitionsPerTopic.put(topic1, 1);
 
         Map consumers = new HashMap<>();
-        consumers.put(consumer1, new Subscription(topics(topic)));
-        consumers.put(consumer2, new Subscription(topics(topic)));
+        consumers.put(consumer1, new Subscription(topics(topic1)));
+        consumers.put(consumer2, new Subscription(topics(topic1)));
 
         Map> assignment = assignor.assign(partitionsPerTopic, consumers);
-        assertAssignment(partitions(tp(topic, 0)), assignment.get(consumer1));
-        assertAssignment(Collections.emptyList(), assignment.get(consumer2));
+        assertAssignment(partitions(tp(topic1, 0)), assignment.get(consumer1));
+        assertAssignment(Collections.emptyList(), assignment.get(consumer2));
     }
 
 
     @Test
     public void testTwoConsumersOneTopicTwoPartitions() {
-        String topic = "topic";
-        String consumer1 = "consumer1";
-        String consumer2 = "consumer2";
-
         Map partitionsPerTopic = new HashMap<>();
-        partitionsPerTopic.put(topic, 2);
+        partitionsPerTopic.put(topic1, 2);
 
         Map consumers = new HashMap<>();
-        consumers.put(consumer1, new Subscription(topics(topic)));
-        consumers.put(consumer2, new Subscription(topics(topic)));
+        consumers.put(consumer1, new Subscription(topics(topic1)));
+        consumers.put(consumer2, new Subscription(topics(topic1)));
 
         Map> assignment = assignor.assign(partitionsPerTopic, consumers);
-        assertAssignment(partitions(tp(topic, 0)), assignment.get(consumer1));
-        assertAssignment(partitions(tp(topic, 1)), assignment.get(consumer2));
+        assertAssignment(partitions(tp(topic1, 0)), assignment.get(consumer1));
+        assertAssignment(partitions(tp(topic1, 1)), assignment.get(consumer2));
     }
 
     @Test
     public void testMultipleConsumersMixedTopics() {
-        String topic1 = "topic1";
-        String topic2 = "topic2";
-        String consumer1 = "consumer1";
-        String consumer2 = "consumer2";
-        String consumer3 = "consumer3";
-
-        Map partitionsPerTopic = new HashMap<>();
-        partitionsPerTopic.put(topic1, 3);
-        partitionsPerTopic.put(topic2, 2);
+        Map partitionsPerTopic = setupPartitionsPerTopicWithTwoTopics(3, 2);
 
         Map consumers = new HashMap<>();
         consumers.put(consumer1, new Subscription(topics(topic1)));
@@ -175,9 +167,7 @@ public void testTwoConsumersTwoTopicsSixPartitions() {
         String consumer1 = "consumer1";
         String consumer2 = "consumer2";
 
-        Map partitionsPerTopic = new HashMap<>();
-        partitionsPerTopic.put(topic1, 3);
-        partitionsPerTopic.put(topic2, 3);
+        Map partitionsPerTopic = setupPartitionsPerTopicWithTwoTopics(3, 3);
 
         Map consumers = new HashMap<>();
         consumers.put(consumer1, new Subscription(topics(topic1, topic2)));
@@ -188,11 +178,155 @@ public void testTwoConsumersTwoTopicsSixPartitions() {
         assertAssignment(partitions(tp(topic1, 2), tp(topic2, 2)), assignment.get(consumer2));
     }
 
+    @Test
+    public void testTwoStaticConsumersTwoTopicsSixPartitions() {
+        // although consumer high has a higher rank than consumer low, the comparison happens on
+        // instance id level.
+        String consumerIdLow = "consumer-b";
+        String consumerIdHigh = "consumer-a";
+
+        Map partitionsPerTopic = setupPartitionsPerTopicWithTwoTopics(3, 3);
+
+        Map consumers = new HashMap<>();
+        Subscription consumerLowSubscription = new Subscription(topics(topic1, topic2),
+                                                              null,
+                                                              Collections.emptyList());
+        consumerLowSubscription.setGroupInstanceId(Optional.of(instance1));
+        consumers.put(consumerIdLow, consumerLowSubscription);
+        Subscription consumerHighSubscription = new Subscription(topics(topic1, topic2),
+                                                              null,
+                                                              Collections.emptyList());
+        consumerHighSubscription.setGroupInstanceId(Optional.of(instance2));
+        consumers.put(consumerIdHigh, consumerHighSubscription);
+        Map> assignment = assignor.assign(partitionsPerTopic, consumers);
+        assertAssignment(partitions(tp(topic1, 0), tp(topic1, 1), tp(topic2, 0), tp(topic2, 1)), assignment.get(consumerIdLow));
+        assertAssignment(partitions(tp(topic1, 2), tp(topic2, 2)), assignment.get(consumerIdHigh));
+    }
+
+    @Test
+    public void testOneStaticConsumerAndOneDynamicConsumerTwoTopicsSixPartitions() {
+        // although consumer high has a higher rank than low, consumer low will win the comparison
+        // because it has instance id while consumer 2 doesn't.
+        String consumerIdLow = "consumer-b";
+        String consumerIdHigh = "consumer-a";
+
+        Map partitionsPerTopic = setupPartitionsPerTopicWithTwoTopics(3, 3);
+
+        Map consumers = new HashMap<>();
+
+        Subscription consumerLowSubscription = new Subscription(topics(topic1, topic2),
+                                                              null,
+                                                              Collections.emptyList());
+        consumerLowSubscription.setGroupInstanceId(Optional.of(instance1));
+        consumers.put(consumerIdLow, consumerLowSubscription);
+        consumers.put(consumerIdHigh, new Subscription(topics(topic1, topic2)));
+
+        Map> assignment = assignor.assign(partitionsPerTopic, consumers);
+        assertAssignment(partitions(tp(topic1, 0), tp(topic1, 1), tp(topic2, 0), tp(topic2, 1)), assignment.get(consumerIdLow));
+        assertAssignment(partitions(tp(topic1, 2), tp(topic2, 2)), assignment.get(consumerIdHigh));
+    }
+
+    @Test
+    public void testStaticMemberRangeAssignmentPersistent() {
+        Map partitionsPerTopic = setupPartitionsPerTopicWithTwoTopics(5, 4);
+
+        Map consumers = new HashMap<>();
+        for (MemberInfo m : staticMemberInfos) {
+            Subscription subscription = new Subscription(topics(topic1, topic2),
+                                                         null,
+                                                         Collections.emptyList());
+            subscription.setGroupInstanceId(m.groupInstanceId);
+            consumers.put(m.memberId, subscription);
+        }
+        // Consumer 4 is a dynamic member.
+        String consumer4 = "consumer4";
+        consumers.put(consumer4, new Subscription(topics(topic1, topic2)));
+
+        Map> expectedAssignment = new HashMap<>();
+        // Have 3 static members instance1, instance2, instance3 to be persistent
+        // across generations. Their assignment shall be the same.
+        expectedAssignment.put(consumer1, partitions(tp(topic1, 0), tp(topic1, 1), tp(topic2, 0)));
+        expectedAssignment.put(consumer2, partitions(tp(topic1, 2), tp(topic2, 1)));
+        expectedAssignment.put(consumer3, partitions(tp(topic1, 3), tp(topic2, 2)));
+        expectedAssignment.put(consumer4, partitions(tp(topic1, 4), tp(topic2, 3)));
+
+        Map> assignment = assignor.assign(partitionsPerTopic, consumers);
+        assertEquals(expectedAssignment, assignment);
+
+        // Replace dynamic member 4 with a new dynamic member 5.
+        consumers.remove(consumer4);
+        String consumer5 = "consumer5";
+        consumers.put(consumer5, new Subscription(topics(topic1, topic2)));
+
+        expectedAssignment.remove(consumer4);
+        expectedAssignment.put(consumer5, partitions(tp(topic1, 4), tp(topic2, 3)));
+        assignment = assignor.assign(partitionsPerTopic, consumers);
+        assertEquals(expectedAssignment, assignment);
+    }
+
+    @Test
+    public void testStaticMemberRangeAssignmentPersistentAfterMemberIdChanges() {
+        Map partitionsPerTopic = setupPartitionsPerTopicWithTwoTopics(5, 5);
+
+        Map consumers = new HashMap<>();
+        for (MemberInfo m : staticMemberInfos) {
+            Subscription subscription = new Subscription(topics(topic1, topic2),
+                                                         null,
+                                                         Collections.emptyList());
+            subscription.setGroupInstanceId(m.groupInstanceId);
+            consumers.put(m.memberId, subscription);
+        }
+        Map> expectedInstanceAssignment = new HashMap<>();
+        expectedInstanceAssignment.put(instance1,
+                                       partitions(tp(topic1, 0), tp(topic1, 1), tp(topic2, 0), tp(topic2, 1)));
+        expectedInstanceAssignment.put(instance2,
+                                       partitions(tp(topic1, 2), tp(topic1, 3), tp(topic2, 2), tp(topic2, 3)));
+        expectedInstanceAssignment.put(instance3,
+                                       partitions(tp(topic1, 4), tp(topic2, 4)));
+
+        Map> staticAssignment =
+            checkStaticAssignment(assignor, partitionsPerTopic, consumers);
+        assertEquals(expectedInstanceAssignment, staticAssignment);
+
+        // Now switch the member.id fields for each member info, the assignment should
+        // stay the same as last time.
+        String consumer4 = "consumer4";
+        String consumer5 = "consumer5";
+        consumers.put(consumer4, consumers.get(consumer3));
+        consumers.remove(consumer3);
+        consumers.put(consumer5, consumers.get(consumer2));
+        consumers.remove(consumer2);
+
+        Map> newStaticAssignment =
+            checkStaticAssignment(assignor, partitionsPerTopic, consumers);
+        assertEquals(staticAssignment, newStaticAssignment);
+    }
+
+    static Map> checkStaticAssignment(AbstractPartitionAssignor assignor,
+                                                                   Map partitionsPerTopic,
+                                                                   Map consumers) {
+        Map> assignmentByMemberId = assignor.assign(partitionsPerTopic, consumers);
+        Map> assignmentByInstanceId = new HashMap<>();
+        for (Map.Entry entry : consumers.entrySet()) {
+            String memberId = entry.getKey();
+            Optional instanceId = entry.getValue().groupInstanceId();
+            instanceId.ifPresent(id -> assignmentByInstanceId.put(id, assignmentByMemberId.get(memberId)));
+        }
+        return assignmentByInstanceId;
+    }
+
     private void assertAssignment(List expected, List actual) {
         // order doesn't matter for assignment, so convert to a set
         assertEquals(new HashSet<>(expected), new HashSet<>(actual));
     }
 
+    private Map setupPartitionsPerTopicWithTwoTopics(int numberOfPartitions1, int numberOfPartitions2) {
+        Map partitionsPerTopic = new HashMap<>();
+        partitionsPerTopic.put(topic1, numberOfPartitions1);
+        partitionsPerTopic.put(topic2, numberOfPartitions2);
+        return partitionsPerTopic;
+    }
+
     private static List topics(String... topics) {
         return Arrays.asList(topics);
     }
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/RoundRobinAssignorTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/RoundRobinAssignorTest.java
index 799a58af43e81..5358a814d4305 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/RoundRobinAssignorTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/RoundRobinAssignorTest.java
@@ -16,41 +16,45 @@
  */
 package org.apache.kafka.clients.consumer;
 
-import org.apache.kafka.clients.consumer.internals.PartitionAssignor.Subscription;
+import org.apache.kafka.clients.consumer.internals.AbstractPartitionAssignor;
+import org.apache.kafka.clients.consumer.internals.AbstractPartitionAssignor.MemberInfo;
+import org.apache.kafka.clients.consumer.ConsumerPartitionAssignor.Subscription;
 import org.apache.kafka.common.TopicPartition;
 import org.junit.Test;
 
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 
+import static org.apache.kafka.clients.consumer.RangeAssignorTest.checkStaticAssignment;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 
 public class RoundRobinAssignorTest {
 
     private RoundRobinAssignor assignor = new RoundRobinAssignor();
+    private String topic = "topic";
+    private String consumerId = "consumer";
 
+    private String topic1 = "topic1";
+    private String topic2 = "topic2";
 
     @Test
     public void testOneConsumerNoTopic() {
-        String consumerId = "consumer";
-
         Map partitionsPerTopic = new HashMap<>();
 
         Map> assignment = assignor.assign(partitionsPerTopic,
-                Collections.singletonMap(consumerId, new Subscription(Collections.emptyList())));
+                Collections.singletonMap(consumerId, new Subscription(Collections.emptyList())));
         assertEquals(Collections.singleton(consumerId), assignment.keySet());
         assertTrue(assignment.get(consumerId).isEmpty());
     }
 
     @Test
     public void testOneConsumerNonexistentTopic() {
-        String topic = "topic";
-        String consumerId = "consumer";
-
         Map partitionsPerTopic = new HashMap<>();
         Map> assignment = assignor.assign(partitionsPerTopic,
                 Collections.singletonMap(consumerId, new Subscription(topics(topic))));
@@ -61,9 +65,6 @@ public void testOneConsumerNonexistentTopic() {
 
     @Test
     public void testOneConsumerOneTopic() {
-        String topic = "topic";
-        String consumerId = "consumer";
-
         Map partitionsPerTopic = new HashMap<>();
         partitionsPerTopic.put(topic, 3);
 
@@ -74,9 +75,7 @@ public void testOneConsumerOneTopic() {
 
     @Test
     public void testOnlyAssignsPartitionsFromSubscribedTopics() {
-        String topic = "topic";
         String otherTopic = "other";
-        String consumerId = "consumer";
 
         Map partitionsPerTopic = new HashMap<>();
         partitionsPerTopic.put(topic, 3);
@@ -89,13 +88,7 @@ public void testOnlyAssignsPartitionsFromSubscribedTopics() {
 
     @Test
     public void testOneConsumerMultipleTopics() {
-        String topic1 = "topic1";
-        String topic2 = "topic2";
-        String consumerId = "consumer";
-
-        Map partitionsPerTopic = new HashMap<>();
-        partitionsPerTopic.put(topic1, 1);
-        partitionsPerTopic.put(topic2, 2);
+        Map partitionsPerTopic = setupPartitionsPerTopicWithTwoTopics(1, 2);
 
         Map> assignment = assignor.assign(partitionsPerTopic,
                 Collections.singletonMap(consumerId, new Subscription(topics(topic1, topic2))));
@@ -104,7 +97,6 @@ public void testOneConsumerMultipleTopics() {
 
     @Test
     public void testTwoConsumersOneTopicOnePartition() {
-        String topic = "topic";
         String consumer1 = "consumer1";
         String consumer2 = "consumer2";
 
@@ -122,7 +114,6 @@ public void testTwoConsumersOneTopicOnePartition() {
 
     @Test
     public void testTwoConsumersOneTopicTwoPartitions() {
-        String topic = "topic";
         String consumer1 = "consumer1";
         String consumer2 = "consumer2";
 
@@ -146,9 +137,7 @@ public void testMultipleConsumersMixedTopics() {
         String consumer2 = "consumer2";
         String consumer3 = "consumer3";
 
-        Map partitionsPerTopic = new HashMap<>();
-        partitionsPerTopic.put(topic1, 3);
-        partitionsPerTopic.put(topic2, 2);
+        Map partitionsPerTopic = setupPartitionsPerTopicWithTwoTopics(3, 2);
 
         Map consumers = new HashMap<>();
         consumers.put(consumer1, new Subscription(topics(topic1)));
@@ -162,15 +151,13 @@ public void testMultipleConsumersMixedTopics() {
     }
 
     @Test
-    public void testTwoConsumersTwoTopicsSixPartitions() {
+    public void testTwoDynamicConsumersTwoTopicsSixPartitions() {
         String topic1 = "topic1";
         String topic2 = "topic2";
         String consumer1 = "consumer1";
         String consumer2 = "consumer2";
 
-        Map partitionsPerTopic = new HashMap<>();
-        partitionsPerTopic.put(topic1, 3);
-        partitionsPerTopic.put(topic2, 3);
+        Map partitionsPerTopic = setupPartitionsPerTopicWithTwoTopics(3, 3);
 
         Map consumers = new HashMap<>();
         consumers.put(consumer1, new Subscription(topics(topic1, topic2)));
@@ -181,6 +168,155 @@ public void testTwoConsumersTwoTopicsSixPartitions() {
         assertEquals(partitions(tp(topic1, 1), tp(topic2, 0), tp(topic2, 2)), assignment.get(consumer2));
     }
 
+    @Test
+    public void testTwoStaticConsumersTwoTopicsSixPartitions() {
+        // although consumer 2 has a higher rank than 1, the comparison happens on
+        // instance id level.
+        String topic1 = "topic1";
+        String topic2 = "topic2";
+        String consumer1 = "consumer-b";
+        String instance1 = "instance1";
+        String consumer2 = "consumer-a";
+        String instance2 = "instance2";
+
+        Map partitionsPerTopic = setupPartitionsPerTopicWithTwoTopics(3, 3);
+
+        Map consumers = new HashMap<>();
+        Subscription consumer1Subscription = new Subscription(topics(topic1, topic2), null);
+        consumer1Subscription.setGroupInstanceId(Optional.of(instance1));
+        consumers.put(consumer1, consumer1Subscription);
+        Subscription consumer2Subscription = new Subscription(topics(topic1, topic2), null);
+        consumer2Subscription.setGroupInstanceId(Optional.of(instance2));
+        consumers.put(consumer2, consumer2Subscription);
+        Map> assignment = assignor.assign(partitionsPerTopic, consumers);
+        assertEquals(partitions(tp(topic1, 0), tp(topic1, 2), tp(topic2, 1)), assignment.get(consumer1));
+        assertEquals(partitions(tp(topic1, 1), tp(topic2, 0), tp(topic2, 2)), assignment.get(consumer2));
+    }
+
+    @Test
+    public void testOneStaticConsumerAndOneDynamicConsumerTwoTopicsSixPartitions() {
+        // although consumer 2 has a higher rank than 1, consumer 1 will win the comparison
+        // because it has instance id while consumer 2 doesn't.
+        String consumer1 = "consumer-b";
+        String instance1 = "instance1";
+        String consumer2 = "consumer-a";
+
+        Map partitionsPerTopic = setupPartitionsPerTopicWithTwoTopics(3, 3);
+
+        Map consumers = new HashMap<>();
+
+        Subscription consumer1Subscription = new Subscription(topics(topic1, topic2), null);
+        consumer1Subscription.setGroupInstanceId(Optional.of(instance1));
+        consumers.put(consumer1, consumer1Subscription);
+        consumers.put(consumer2, new Subscription(topics(topic1, topic2)));
+
+        Map> assignment = assignor.assign(partitionsPerTopic, consumers);
+        assertEquals(partitions(tp(topic1, 0), tp(topic1, 2), tp(topic2, 1)), assignment.get(consumer1));
+        assertEquals(partitions(tp(topic1, 1), tp(topic2, 0), tp(topic2, 2)), assignment.get(consumer2));
+    }
+
+    @Test
+    public void testStaticMemberRoundRobinAssignmentPersistent() {
+        // Have 3 static members instance1, instance2, instance3 to be persistent
+        // across generations. Their assignment shall be the same.
+        String consumer1 = "consumer1";
+        String instance1 = "instance1";
+        String consumer2 = "consumer2";
+        String instance2 = "instance2";
+        String consumer3 = "consumer3";
+        String instance3 = "instance3";
+
+        List staticMemberInfos = new ArrayList<>();
+        staticMemberInfos.add(new MemberInfo(consumer1, Optional.of(instance1)));
+        staticMemberInfos.add(new MemberInfo(consumer2, Optional.of(instance2)));
+        staticMemberInfos.add(new MemberInfo(consumer3, Optional.of(instance3)));
+
+        // Consumer 4 is a dynamic member.
+        String consumer4 = "consumer4";
+
+        Map partitionsPerTopic = setupPartitionsPerTopicWithTwoTopics(3, 3);
+
+        Map consumers = new HashMap<>();
+        for (MemberInfo m : staticMemberInfos) {
+            Subscription subscription = new Subscription(topics(topic1, topic2), null);
+            subscription.setGroupInstanceId(m.groupInstanceId);
+            consumers.put(m.memberId, subscription);
+        }
+        consumers.put(consumer4, new Subscription(topics(topic1, topic2)));
+
+        Map> expectedAssignment = new HashMap<>();
+        expectedAssignment.put(consumer1, partitions(tp(topic1, 0), tp(topic2, 1)));
+        expectedAssignment.put(consumer2, partitions(tp(topic1, 1), tp(topic2, 2)));
+        expectedAssignment.put(consumer3, partitions(tp(topic1, 2)));
+        expectedAssignment.put(consumer4, partitions(tp(topic2, 0)));
+
+        Map> assignment = assignor.assign(partitionsPerTopic, consumers);
+        assertEquals(expectedAssignment, assignment);
+
+        // Replace dynamic member 4 with a new dynamic member 5.
+        consumers.remove(consumer4);
+        String consumer5 = "consumer5";
+        consumers.put(consumer5, new Subscription(topics(topic1, topic2)));
+
+        expectedAssignment.remove(consumer4);
+        expectedAssignment.put(consumer5, partitions(tp(topic2, 0)));
+        assignment = assignor.assign(partitionsPerTopic, consumers);
+        assertEquals(expectedAssignment, assignment);
+    }
+
+    @Test
+    public void testStaticMemberRoundRobinAssignmentPersistentAfterMemberIdChanges() {
+        String consumer1 = "consumer1";
+        String instance1 = "instance1";
+        String consumer2 = "consumer2";
+        String instance2 = "instance2";
+        String consumer3 = "consumer3";
+        String instance3 = "instance3";
+        Map memberIdToInstanceId = new HashMap<>();
+        memberIdToInstanceId.put(consumer1, instance1);
+        memberIdToInstanceId.put(consumer2, instance2);
+        memberIdToInstanceId.put(consumer3, instance3);
+
+        Map partitionsPerTopic = setupPartitionsPerTopicWithTwoTopics(5, 5);
+
+        Map> expectedInstanceAssignment = new HashMap<>();
+        expectedInstanceAssignment.put(instance1,
+                                       partitions(tp(topic1, 0), tp(topic1, 3), tp(topic2, 1), tp(topic2, 4)));
+        expectedInstanceAssignment.put(instance2,
+                                       partitions(tp(topic1, 1), tp(topic1, 4), tp(topic2, 2)));
+        expectedInstanceAssignment.put(instance3,
+                                       partitions(tp(topic1, 2), tp(topic2, 0), tp(topic2, 3)));
+
+        List staticMemberInfos = new ArrayList<>();
+        for (Map.Entry entry : memberIdToInstanceId.entrySet()) {
+            staticMemberInfos.add(new AbstractPartitionAssignor.MemberInfo(entry.getKey(), Optional.of(entry.getValue())));
+        }
+        Map consumers = new HashMap<>();
+        for (MemberInfo m : staticMemberInfos) {
+            Subscription subscription = new Subscription(topics(topic1, topic2), null);
+            subscription.setGroupInstanceId(m.groupInstanceId);
+            consumers.put(m.memberId, subscription);
+        }
+
+        Map> staticAssignment =
+            checkStaticAssignment(assignor, partitionsPerTopic, consumers);
+        assertEquals(expectedInstanceAssignment, staticAssignment);
+
+        memberIdToInstanceId.clear();
+
+        // Now switch the member.id fields for each member info, the assignment should
+        // stay the same as last time.
+        String consumer4 = "consumer4";
+        String consumer5 = "consumer5";
+        consumers.put(consumer4, consumers.get(consumer3));
+        consumers.remove(consumer3);
+        consumers.put(consumer5, consumers.get(consumer2));
+        consumers.remove(consumer2);
+        Map> newStaticAssignment =
+            checkStaticAssignment(assignor, partitionsPerTopic, consumers);
+        assertEquals(staticAssignment, newStaticAssignment);
+    }
+
     private static List topics(String... topics) {
         return Arrays.asList(topics);
     }
@@ -192,4 +328,11 @@ private static List partitions(TopicPartition... partitions) {
     private static TopicPartition tp(String topic, int partition) {
         return new TopicPartition(topic, partition);
     }
+
+    private Map setupPartitionsPerTopicWithTwoTopics(int numberOfPartitions1, int numberOfPartitions2) {
+        Map partitionsPerTopic = new HashMap<>();
+        partitionsPerTopic.put(topic1, numberOfPartitions1);
+        partitionsPerTopic.put(topic2, numberOfPartitions2);
+        return partitionsPerTopic;
+    }
 }
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/StickyAssignorTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/StickyAssignorTest.java
index 32ba16a482023..fb89944903739 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/StickyAssignorTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/StickyAssignorTest.java
@@ -16,752 +16,230 @@
  */
 package org.apache.kafka.clients.consumer;
 
+import static org.apache.kafka.clients.consumer.StickyAssignor.serializeTopicPartitionAssignment;
+import static org.apache.kafka.clients.consumer.internals.AbstractStickyAssignor.DEFAULT_GENERATION;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 
+import java.nio.ByteBuffer;
 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.Map.Entry;
-import java.util.Random;
-import java.util.Set;
-
-import org.apache.kafka.clients.consumer.internals.PartitionAssignor.Subscription;
+import java.util.Optional;
+import org.apache.kafka.clients.consumer.ConsumerPartitionAssignor.Subscription;
+import org.apache.kafka.clients.consumer.internals.AbstractStickyAssignor;
+import org.apache.kafka.clients.consumer.internals.AbstractStickyAssignor.MemberData;
+import org.apache.kafka.clients.consumer.internals.AbstractStickyAssignorTest;
 import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.protocol.types.Struct;
 import org.apache.kafka.common.utils.CollectionUtils;
-import org.apache.kafka.common.utils.Utils;
 import org.junit.Test;
 
-public class StickyAssignorTest {
-
-    private StickyAssignor assignor = new StickyAssignor();
-
-    @Test
-    public void testOneConsumerNoTopic() {
-        String consumerId = "consumer";
-
-        Map partitionsPerTopic = new HashMap<>();
-        Map subscriptions =
-                Collections.singletonMap(consumerId, new Subscription(Collections.emptyList()));
-
-        Map> assignment = assignor.assign(partitionsPerTopic, subscriptions);
-        assertEquals(Collections.singleton(consumerId), assignment.keySet());
-        assertTrue(assignment.get(consumerId).isEmpty());
+public class StickyAssignorTest extends AbstractStickyAssignorTest {
 
-        verifyValidityAndBalance(subscriptions, assignment);
-        assertTrue(isFullyBalanced(assignment));
+    @Override
+    public AbstractStickyAssignor createAssignor() {
+        return new StickyAssignor();
     }
 
-    @Test
-    public void testOneConsumerNonexistentTopic() {
-        String topic = "topic";
-        String consumerId = "consumer";
-
-        Map partitionsPerTopic = new HashMap<>();
-        partitionsPerTopic.put(topic, 0);
-        Map subscriptions = Collections.singletonMap(consumerId, new Subscription(topics(topic)));
-
-        Map> assignment = assignor.assign(partitionsPerTopic, subscriptions);
-
-        assertEquals(Collections.singleton(consumerId), assignment.keySet());
-        assertTrue(assignment.get(consumerId).isEmpty());
-
-        verifyValidityAndBalance(subscriptions, assignment);
-        assertTrue(isFullyBalanced(assignment));
-    }
-
-    @Test
-    public void testOneConsumerOneTopic() {
-        String topic = "topic";
-        String consumerId = "consumer";
-
-        Map partitionsPerTopic = new HashMap<>();
-        partitionsPerTopic.put(topic, 3);
-        Map subscriptions = Collections.singletonMap(consumerId, new Subscription(topics(topic)));
-
-        Map> assignment = assignor.assign(partitionsPerTopic, subscriptions);
-        assertEquals(partitions(tp(topic, 0), tp(topic, 1), tp(topic, 2)), assignment.get(consumerId));
-
-        verifyValidityAndBalance(subscriptions, assignment);
-        assertTrue(isFullyBalanced(assignment));
+    @Override
+    public Subscription buildSubscription(List topics, List partitions) {
+        return new Subscription(topics,
+            serializeTopicPartitionAssignment(new MemberData(partitions, Optional.of(DEFAULT_GENERATION))));
     }
 
     @Test
-    public void testOnlyAssignsPartitionsFromSubscribedTopics() {
-        String topic = "topic";
-        String otherTopic = "other";
-        String consumerId = "consumer";
-
-        Map partitionsPerTopic = new HashMap<>();
-        partitionsPerTopic.put(topic, 3);
-        partitionsPerTopic.put(otherTopic, 3);
-        Map subscriptions = Collections.singletonMap(consumerId, new Subscription(topics(topic)));
-
-        Map> assignment = assignor.assign(partitionsPerTopic, subscriptions);
-        assertEquals(partitions(tp(topic, 0), tp(topic, 1), tp(topic, 2)), assignment.get(consumerId));
-
-        verifyValidityAndBalance(subscriptions, assignment);
-        assertTrue(isFullyBalanced(assignment));
-    }
-
-    @Test
-    public void testOneConsumerMultipleTopics() {
-        String topic1 = "topic1";
-        String topic2 = "topic2";
-        String consumerId = "consumer";
-
-        Map partitionsPerTopic = new HashMap<>();
-        partitionsPerTopic.put(topic1, 1);
-        partitionsPerTopic.put(topic2, 2);
-        Map subscriptions = Collections.singletonMap(consumerId, new Subscription(topics(topic1, topic2)));
-
-        Map> assignment = assignor.assign(partitionsPerTopic, subscriptions);
-        assertEquals(partitions(tp(topic1, 0), tp(topic2, 0), tp(topic2, 1)), assignment.get(consumerId));
-
-        verifyValidityAndBalance(subscriptions, assignment);
-        assertTrue(isFullyBalanced(assignment));
-    }
-
-    @Test
-    public void testTwoConsumersOneTopicOnePartition() {
-        String topic = "topic";
-        String consumer1 = "consumer1";
-        String consumer2 = "consumer2";
-
-        Map partitionsPerTopic = new HashMap<>();
-        partitionsPerTopic.put(topic, 1);
-
-        Map subscriptions = new HashMap<>();
-        subscriptions.put(consumer1, new Subscription(topics(topic)));
-        subscriptions.put(consumer2, new Subscription(topics(topic)));
-
-        Map> assignment = assignor.assign(partitionsPerTopic, subscriptions);
-        assertEquals(partitions(tp(topic, 0)), assignment.get(consumer1));
-        assertEquals(Collections.emptyList(), assignment.get(consumer2));
-
-        verifyValidityAndBalance(subscriptions, assignment);
-        assertTrue(isFullyBalanced(assignment));
-    }
-
-    @Test
-    public void testTwoConsumersOneTopicTwoPartitions() {
-        String topic = "topic";
+    public void testAssignmentWithMultipleGenerations1() {
         String consumer1 = "consumer1";
         String consumer2 = "consumer2";
+        String consumer3 = "consumer3";
 
         Map partitionsPerTopic = new HashMap<>();
-        partitionsPerTopic.put(topic, 2);
-
-        Map subscriptions = new HashMap<>();
+        partitionsPerTopic.put(topic, 6);
         subscriptions.put(consumer1, new Subscription(topics(topic)));
         subscriptions.put(consumer2, new Subscription(topics(topic)));
+        subscriptions.put(consumer3, new Subscription(topics(topic)));
 
         Map> assignment = assignor.assign(partitionsPerTopic, subscriptions);
-        assertEquals(partitions(tp(topic, 0)), assignment.get(consumer1));
-        assertEquals(partitions(tp(topic, 1)), assignment.get(consumer2));
-
-        verifyValidityAndBalance(subscriptions, assignment);
-        assertTrue(isFullyBalanced(assignment));
-    }
-
-    @Test
-    public void testMultipleConsumersMixedTopicSubscriptions() {
-        String topic1 = "topic1";
-        String topic2 = "topic2";
-        String consumer1 = "consumer1";
-        String consumer2 = "consumer2";
-        String consumer3 = "consumer3";
-
-        Map partitionsPerTopic = new HashMap<>();
-        partitionsPerTopic.put(topic1, 3);
-        partitionsPerTopic.put(topic2, 2);
-
-        Map subscriptions = new HashMap<>();
-        subscriptions.put(consumer1, new Subscription(topics(topic1)));
-        subscriptions.put(consumer2, new Subscription(topics(topic1, topic2)));
-        subscriptions.put(consumer3, new Subscription(topics(topic1)));
-
-        Map> assignment = assignor.assign(partitionsPerTopic, subscriptions);
-        assertEquals(partitions(tp(topic1, 0), tp(topic1, 2)), assignment.get(consumer1));
-        assertEquals(partitions(tp(topic2, 0), tp(topic2, 1)), assignment.get(consumer2));
-        assertEquals(partitions(tp(topic1, 1)), assignment.get(consumer3));
-
-        verifyValidityAndBalance(subscriptions, assignment);
-        assertTrue(isFullyBalanced(assignment));
-    }
-
-    @Test
-    public void testTwoConsumersTwoTopicsSixPartitions() {
-        String topic1 = "topic1";
-        String topic2 = "topic2";
-        String consumer1 = "consumer1";
-        String consumer2 = "consumer2";
-
-        Map partitionsPerTopic = new HashMap<>();
-        partitionsPerTopic.put(topic1, 3);
-        partitionsPerTopic.put(topic2, 3);
-
-        Map subscriptions = new HashMap<>();
-        subscriptions.put(consumer1, new Subscription(topics(topic1, topic2)));
-        subscriptions.put(consumer2, new Subscription(topics(topic1, topic2)));
-
-        Map> assignment = assignor.assign(partitionsPerTopic, subscriptions);
-        assertEquals(partitions(tp(topic1, 0), tp(topic1, 2), tp(topic2, 1)), assignment.get(consumer1));
-        assertEquals(partitions(tp(topic1, 1), tp(topic2, 0), tp(topic2, 2)), assignment.get(consumer2));
-
-        verifyValidityAndBalance(subscriptions, assignment);
+        List r1partitions1 = assignment.get(consumer1);
+        List r1partitions2 = assignment.get(consumer2);
+        List r1partitions3 = assignment.get(consumer3);
+        assertTrue(r1partitions1.size() == 2 && r1partitions2.size() == 2 && r1partitions3.size() == 2);
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
         assertTrue(isFullyBalanced(assignment));
-    }
-
-    @Test
-    public void testAddRemoveConsumerOneTopic() {
-        String topic = "topic";
-        String consumer1 = "consumer";
-
-        Map partitionsPerTopic = new HashMap<>();
-        partitionsPerTopic.put(topic, 3);
-        Map subscriptions = new HashMap<>();
-        subscriptions.put(consumer1, new Subscription(topics(topic)));
 
-        Map> assignment = assignor.assign(partitionsPerTopic, subscriptions);
-        assertEquals(partitions(tp(topic, 0), tp(topic, 1), tp(topic, 2)), assignment.get(consumer1));
+        subscriptions.put(consumer1, buildSubscription(topics(topic), r1partitions1));
+        subscriptions.put(consumer2, buildSubscription(topics(topic), r1partitions2));
+        subscriptions.remove(consumer3);
 
-        verifyValidityAndBalance(subscriptions, assignment);
-        assertTrue(isFullyBalanced(assignment));
-
-        String consumer2 = "consumer2";
-        subscriptions.put(consumer1,
-                new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment(assignment.get(consumer1))));
-        subscriptions.put(consumer2, new Subscription(topics(topic)));
         assignment = assignor.assign(partitionsPerTopic, subscriptions);
-        assertEquals(partitions(tp(topic, 1), tp(topic, 2)), assignment.get(consumer1));
-        assertEquals(partitions(tp(topic, 0)), assignment.get(consumer2));
-
-        verifyValidityAndBalance(subscriptions, assignment);
+        List r2partitions1 = assignment.get(consumer1);
+        List r2partitions2 = assignment.get(consumer2);
+        assertTrue(r2partitions1.size() == 3 && r2partitions2.size() == 3);
+        assertTrue(r2partitions1.containsAll(r1partitions1));
+        assertTrue(r2partitions2.containsAll(r1partitions2));
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
         assertTrue(isFullyBalanced(assignment));
         assertTrue(assignor.isSticky());
+        assertFalse(Collections.disjoint(r2partitions2, r1partitions3));
 
         subscriptions.remove(consumer1);
-        subscriptions.put(consumer2,
-                new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment(assignment.get(consumer2))));
-        assignment = assignor.assign(partitionsPerTopic, subscriptions);
-        assertTrue(assignment.get(consumer2).contains(tp(topic, 0)));
-        assertTrue(assignment.get(consumer2).contains(tp(topic, 1)));
-        assertTrue(assignment.get(consumer2).contains(tp(topic, 2)));
+        subscriptions.put(consumer2, buildSubscriptionWithGeneration(topics(topic), r2partitions2, 2));
+        subscriptions.put(consumer3, buildSubscriptionWithGeneration(topics(topic), r1partitions3, 1));
 
-        verifyValidityAndBalance(subscriptions, assignment);
+        assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        List r3partitions2 = assignment.get(consumer2);
+        List r3partitions3 = assignment.get(consumer3);
+        assertTrue(r3partitions2.size() == 3 && r3partitions3.size() == 3);
+        assertTrue(Collections.disjoint(r3partitions2, r3partitions3));
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
         assertTrue(isFullyBalanced(assignment));
         assertTrue(assignor.isSticky());
     }
 
-    /**
-     * This unit test performs sticky assignment for a scenario that round robin assignor handles poorly.
-     * Topics (partitions per topic): topic1 (2), topic2 (1), topic3 (2), topic4 (1), topic5 (2)
-     * Subscriptions:
-     *  - consumer1: topic1, topic2, topic3, topic4, topic5
-     *  - consumer2: topic1, topic3, topic5
-     *  - consumer3: topic1, topic3, topic5
-     *  - consumer4: topic1, topic2, topic3, topic4, topic5
-     * Round Robin Assignment Result:
-     *  - consumer1: topic1-0, topic3-0, topic5-0
-     *  - consumer2: topic1-1, topic3-1, topic5-1
-     *  - consumer3:
-     *  - consumer4: topic2-0, topic4-0
-     * Sticky Assignment Result:
-     *  - consumer1: topic2-0, topic3-0
-     *  - consumer2: topic1-0, topic3-1
-     *  - consumer3: topic1-1, topic5-0
-     *  - consumer4: topic4-0, topic5-1
-     */
-    @Test
-    public void testPoorRoundRobinAssignmentScenario() {
-        Map partitionsPerTopic = new HashMap<>();
-        for (int i = 1; i <= 5; i++)
-            partitionsPerTopic.put(String.format("topic%d", i), (i % 2) + 1);
-
-        Map subscriptions = new HashMap<>();
-        subscriptions.put("consumer1", new Subscription(topics("topic1", "topic2", "topic3", "topic4", "topic5")));
-        subscriptions.put("consumer2", new Subscription(topics("topic1", "topic3", "topic5")));
-        subscriptions.put("consumer3", new Subscription(topics("topic1", "topic3", "topic5")));
-        subscriptions.put("consumer4", new Subscription(topics("topic1", "topic2", "topic3", "topic4", "topic5")));
-
-        Map> assignment = assignor.assign(partitionsPerTopic, subscriptions);
-        verifyValidityAndBalance(subscriptions, assignment);
-    }
-
     @Test
-    public void testAddRemoveTopicTwoConsumers() {
-        String topic = "topic";
-        String consumer1 = "consumer";
+    public void testAssignmentWithMultipleGenerations2() {
+        String consumer1 = "consumer1";
         String consumer2 = "consumer2";
+        String consumer3 = "consumer3";
 
         Map partitionsPerTopic = new HashMap<>();
-        partitionsPerTopic.put(topic, 3);
-        Map subscriptions = new HashMap<>();
+        partitionsPerTopic.put(topic, 6);
         subscriptions.put(consumer1, new Subscription(topics(topic)));
         subscriptions.put(consumer2, new Subscription(topics(topic)));
+        subscriptions.put(consumer3, new Subscription(topics(topic)));
 
         Map> assignment = assignor.assign(partitionsPerTopic, subscriptions);
-        // verify balance
+        List r1partitions1 = assignment.get(consumer1);
+        List r1partitions2 = assignment.get(consumer2);
+        List r1partitions3 = assignment.get(consumer3);
+        assertTrue(r1partitions1.size() == 2 && r1partitions2.size() == 2 && r1partitions3.size() == 2);
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
         assertTrue(isFullyBalanced(assignment));
-        verifyValidityAndBalance(subscriptions, assignment);
-        // verify stickiness
-        List consumer1Assignment1 = assignment.get(consumer1);
-        List consumer2Assignment1 = assignment.get(consumer2);
-        assertTrue((consumer1Assignment1.size() == 1 && consumer2Assignment1.size() == 2) ||
-                   (consumer1Assignment1.size() == 2 && consumer2Assignment1.size() == 1));
-
-        String topic2 = "topic2";
-        partitionsPerTopic.put(topic2, 3);
-        subscriptions.put(consumer1,
-                new Subscription(topics(topic, topic2), StickyAssignor.serializeTopicPartitionAssignment(assignment.get(consumer1))));
-        subscriptions.put(consumer2,
-                new Subscription(topics(topic, topic2), StickyAssignor.serializeTopicPartitionAssignment(assignment.get(consumer2))));
-        assignment = assignor.assign(partitionsPerTopic, subscriptions);
-        // verify balance
-        verifyValidityAndBalance(subscriptions, assignment);
-        assertTrue(isFullyBalanced(assignment));
-        // verify stickiness
-        List consumer1assignment = assignment.get(consumer1);
-        List consumer2assignment = assignment.get(consumer2);
-        assertTrue(consumer1assignment.size() == 3 && consumer2assignment.size() == 3);
-        assertTrue(consumer1assignment.containsAll(consumer1Assignment1));
-        assertTrue(consumer2assignment.containsAll(consumer2Assignment1));
-        assertTrue(assignor.isSticky());
-
-        partitionsPerTopic.remove(topic);
-        subscriptions.put(consumer1,
-                new Subscription(topics(topic2), StickyAssignor.serializeTopicPartitionAssignment(assignment.get(consumer1))));
-        subscriptions.put(consumer2,
-                new Subscription(topics(topic2), StickyAssignor.serializeTopicPartitionAssignment(assignment.get(consumer2))));
-        assignment = assignor.assign(partitionsPerTopic, subscriptions);
-        // verify balance
-        verifyValidityAndBalance(subscriptions, assignment);
-        assertTrue(isFullyBalanced(assignment));
-        // verify stickiness
-        List consumer1Assignment3 = assignment.get(consumer1);
-        List consumer2Assignment3 = assignment.get(consumer2);
-        assertTrue((consumer1Assignment3.size() == 1 && consumer2Assignment3.size() == 2) ||
-                   (consumer1Assignment3.size() == 2 && consumer2Assignment3.size() == 1));
-        assertTrue(consumer1assignment.containsAll(consumer1Assignment3));
-        assertTrue(consumer2assignment.containsAll(consumer2Assignment3));
-        assertTrue(assignor.isSticky());
-    }
-
-    @Test
-    public void testReassignmentAfterOneConsumerLeaves() {
-        Map partitionsPerTopic = new HashMap<>();
-        for (int i = 1; i < 20; i++)
-            partitionsPerTopic.put(getTopicName(i, 20), i);
-
-        Map subscriptions = new HashMap<>();
-        for (int i = 1; i < 20; i++) {
-            List topics = new ArrayList();
-            for (int j = 1; j <= i; j++)
-                topics.add(getTopicName(j, 20));
-            subscriptions.put(getConsumerName(i, 20), new Subscription(topics));
-        }
-
-        Map> assignment = assignor.assign(partitionsPerTopic, subscriptions);
-        verifyValidityAndBalance(subscriptions, assignment);
 
-        for (int i = 1; i < 20; i++) {
-            String consumer = getConsumerName(i, 20);
-            subscriptions.put(consumer,
-                    new Subscription(subscriptions.get(consumer).topics(), StickyAssignor.serializeTopicPartitionAssignment(assignment.get(consumer))));
-        }
-        subscriptions.remove("consumer10");
-
-        assignment = assignor.assign(partitionsPerTopic, subscriptions);
-        verifyValidityAndBalance(subscriptions, assignment);
-        assertTrue(assignor.isSticky());
-    }
-
-    @Test
-    public void testReassignmentAfterOneConsumerAdded() {
-        Map partitionsPerTopic = new HashMap<>();
-        partitionsPerTopic.put("topic", 20);
-
-        Map subscriptions = new HashMap<>();
-        for (int i = 1; i < 10; i++)
-            subscriptions.put(getConsumerName(i, 10), new Subscription(topics("topic")));
-
-        Map> assignment = assignor.assign(partitionsPerTopic, subscriptions);
-        verifyValidityAndBalance(subscriptions, assignment);
-
-        // add a new consumer
-        subscriptions.put(getConsumerName(10, 10), new Subscription(topics("topic")));
+        subscriptions.remove(consumer1);
+        subscriptions.put(consumer2, buildSubscriptionWithGeneration(topics(topic), r1partitions2, 1));
+        subscriptions.remove(consumer3);
 
         assignment = assignor.assign(partitionsPerTopic, subscriptions);
-        verifyValidityAndBalance(subscriptions, assignment);
+        List r2partitions2 = assignment.get(consumer2);
+        assertEquals(6, r2partitions2.size());
+        assertTrue(r2partitions2.containsAll(r1partitions2));
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
+        assertTrue(isFullyBalanced(assignment));
         assertTrue(assignor.isSticky());
-    }
-
-    @Test
-    public void testSameSubscriptions() {
-        Map partitionsPerTopic = new HashMap<>();
-        for (int i = 1; i < 15; i++)
-            partitionsPerTopic.put(getTopicName(i, 15), i);
-
-        Map subscriptions = new HashMap<>();
-        for (int i = 1; i < 9; i++) {
-            List topics = new ArrayList();
-            for (int j = 1; j <= partitionsPerTopic.size(); j++)
-                topics.add(getTopicName(j, 15));
-            subscriptions.put(getConsumerName(i, 9), new Subscription(topics));
-        }
-
-        Map> assignment = assignor.assign(partitionsPerTopic, subscriptions);
-        verifyValidityAndBalance(subscriptions, assignment);
 
-        for (int i = 1; i < 9; i++) {
-            String consumer = getConsumerName(i, 9);
-            subscriptions.put(consumer,
-                    new Subscription(subscriptions.get(consumer).topics(), StickyAssignor.serializeTopicPartitionAssignment(assignment.get(consumer))));
-        }
-        subscriptions.remove(getConsumerName(5, 9));
+        subscriptions.put(consumer1, buildSubscriptionWithGeneration(topics(topic), r1partitions1, 1));
+        subscriptions.put(consumer2, buildSubscriptionWithGeneration(topics(topic), r2partitions2, 2));
+        subscriptions.put(consumer3, buildSubscriptionWithGeneration(topics(topic), r1partitions3, 1));
 
         assignment = assignor.assign(partitionsPerTopic, subscriptions);
-        verifyValidityAndBalance(subscriptions, assignment);
+        List r3partitions1 = assignment.get(consumer1);
+        List r3partitions2 = assignment.get(consumer2);
+        List r3partitions3 = assignment.get(consumer3);
+        assertTrue(r3partitions1.size() == 2 && r3partitions2.size() == 2 && r3partitions3.size() == 2);
+        assertEquals(r1partitions1, r3partitions1);
+        assertEquals(r1partitions2, r3partitions2);
+        assertEquals(r1partitions3, r3partitions3);
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
+        assertTrue(isFullyBalanced(assignment));
         assertTrue(assignor.isSticky());
     }
 
     @Test
-    public void testLargeAssignmentWithMultipleConsumersLeaving() {
-        Random rand = new Random();
-        int topicCount = 40;
-        int consumerCount = 200;
+    public void testAssignmentWithConflictingPreviousGenerations() {
+        String consumer1 = "consumer1";
+        String consumer2 = "consumer2";
+        String consumer3 = "consumer3";
 
         Map partitionsPerTopic = new HashMap<>();
-        for (int i = 0; i < topicCount; i++)
-            partitionsPerTopic.put(getTopicName(i, topicCount), rand.nextInt(10) + 1);
-
-        Map subscriptions = new HashMap<>();
-        for (int i = 0; i < consumerCount; i++) {
-            List topics = new ArrayList();
-            for (int j = 0; j < rand.nextInt(20); j++)
-                topics.add(getTopicName(rand.nextInt(topicCount), topicCount));
-            subscriptions.put(getConsumerName(i, consumerCount), new Subscription(topics));
-        }
-
-        Map> assignment = assignor.assign(partitionsPerTopic, subscriptions);
-        verifyValidityAndBalance(subscriptions, assignment);
-
-        for (int i = 1; i < consumerCount; i++) {
-            String consumer = getConsumerName(i, consumerCount);
-            subscriptions.put(consumer,
-                    new Subscription(subscriptions.get(consumer).topics(), StickyAssignor.serializeTopicPartitionAssignment(assignment.get(consumer))));
-        }
-        for (int i = 0; i < 50; ++i) {
-            String c = getConsumerName(rand.nextInt(consumerCount), consumerCount);
-            subscriptions.remove(c);
-        }
+        partitionsPerTopic.put(topic, 6);
+        subscriptions.put(consumer1, new Subscription(topics(topic)));
+        subscriptions.put(consumer2, new Subscription(topics(topic)));
+        subscriptions.put(consumer3, new Subscription(topics(topic)));
 
-        assignment = assignor.assign(partitionsPerTopic, subscriptions);
-        verifyValidityAndBalance(subscriptions, assignment);
-        assertTrue(assignor.isSticky());
-    }
+        TopicPartition tp0 = new TopicPartition(topic, 0);
+        TopicPartition tp1 = new TopicPartition(topic, 1);
+        TopicPartition tp2 = new TopicPartition(topic, 2);
+        TopicPartition tp3 = new TopicPartition(topic, 3);
+        TopicPartition tp4 = new TopicPartition(topic, 4);
+        TopicPartition tp5 = new TopicPartition(topic, 5);
 
-    @Test
-    public void testNewSubscription() {
-        Map partitionsPerTopic = new HashMap<>();
-        for (int i = 1; i < 5; i++)
-            partitionsPerTopic.put(getTopicName(i, 5), 1);
-
-        Map subscriptions = new HashMap<>();
-        for (int i = 0; i < 3; i++) {
-            List topics = new ArrayList();
-            for (int j = i; j <= 3 * i - 2; j++)
-                topics.add(getTopicName(j, 5));
-            subscriptions.put(getConsumerName(i, 3), new Subscription(topics));
-        }
+        List c1partitions0 = partitions(tp0, tp1, tp4);
+        List c2partitions0 = partitions(tp0, tp1, tp2);
+        List c3partitions0 = partitions(tp3, tp4, tp5);
+        subscriptions.put(consumer1, buildSubscriptionWithGeneration(topics(topic), c1partitions0, 1));
+        subscriptions.put(consumer2, buildSubscriptionWithGeneration(topics(topic), c2partitions0, 2));
+        subscriptions.put(consumer3, buildSubscriptionWithGeneration(topics(topic), c3partitions0, 2));
 
         Map> assignment = assignor.assign(partitionsPerTopic, subscriptions);
-        verifyValidityAndBalance(subscriptions, assignment);
-
-        subscriptions.get(getConsumerName(0, 3)).topics().add(getTopicName(1, 5));
+        List c1partitions = assignment.get(consumer1);
+        List c2partitions = assignment.get(consumer2);
+        List c3partitions = assignment.get(consumer3);
 
-        assignment = assignor.assign(partitionsPerTopic, subscriptions);
-        verifyValidityAndBalance(subscriptions, assignment);
+        assertTrue(c1partitions.size() == 2 && c2partitions.size() == 2 && c3partitions.size() == 2);
+        assertTrue(c2partitions0.containsAll(c2partitions));
+        assertTrue(c3partitions0.containsAll(c3partitions));
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
+        assertTrue(isFullyBalanced(assignment));
         assertTrue(assignor.isSticky());
     }
 
     @Test
-    public void testReassignmentWithRandomSubscriptionsAndChanges() {
-        final int minNumConsumers = 20;
-        final int maxNumConsumers = 40;
-        final int minNumTopics = 10;
-        final int maxNumTopics = 20;
-
-        for (int round = 1; round <= 100; ++round) {
-            int numTopics = minNumTopics + new Random().nextInt(maxNumTopics - minNumTopics);
-
-            ArrayList topics = new ArrayList<>();
-            for (int i = 0; i < numTopics; ++i)
-                topics.add(getTopicName(i, maxNumTopics));
-
-            Map partitionsPerTopic = new HashMap<>();
-            for (int i = 0; i < numTopics; ++i)
-                partitionsPerTopic.put(getTopicName(i, maxNumTopics), i + 1);
-
-            int numConsumers = minNumConsumers + new Random().nextInt(maxNumConsumers - minNumConsumers);
-
-            Map subscriptions = new HashMap<>();
-            for (int i = 0; i < numConsumers; ++i) {
-                List sub = Utils.sorted(getRandomSublist(topics));
-                subscriptions.put(getConsumerName(i, maxNumConsumers), new Subscription(sub));
-            }
-
-            StickyAssignor assignor = new StickyAssignor();
-
-            Map> assignment = assignor.assign(partitionsPerTopic, subscriptions);
-            verifyValidityAndBalance(subscriptions, assignment);
-
-            subscriptions.clear();
-            for (int i = 0; i < numConsumers; ++i) {
-                List sub = Utils.sorted(getRandomSublist(topics));
-                String consumer = getConsumerName(i, maxNumConsumers);
-                subscriptions.put(consumer,
-                        new Subscription(sub, StickyAssignor.serializeTopicPartitionAssignment(assignment.get(consumer))));
-            }
-
-            assignment = assignor.assign(partitionsPerTopic, subscriptions);
-            verifyValidityAndBalance(subscriptions, assignment);
-            assertTrue(assignor.isSticky());
-        }
-    }
+    public void testSchemaBackwardCompatibility() {
+        String consumer1 = "consumer1";
+        String consumer2 = "consumer2";
+        String consumer3 = "consumer3";
 
-    @Test
-    public void testMoveExistingAssignments() {
         Map partitionsPerTopic = new HashMap<>();
-        for (int i = 1; i <= 6; i++)
-            partitionsPerTopic.put(String.format("topic%02d", i), 1);
-
-        Map subscriptions = new HashMap<>();
-        subscriptions.put("consumer01",
-                new Subscription(topics("topic01", "topic02"),
-                        StickyAssignor.serializeTopicPartitionAssignment(partitions(tp("topic01", 0)))));
-        subscriptions.put("consumer02",
-                new Subscription(topics("topic01", "topic02", "topic03", "topic04"),
-                        StickyAssignor.serializeTopicPartitionAssignment(partitions(tp("topic02", 0), tp("topic03", 0)))));
-        subscriptions.put("consumer03",
-                new Subscription(topics("topic02", "topic03", "topic04", "topic05", "topic06"),
-                        StickyAssignor.serializeTopicPartitionAssignment(partitions(tp("topic04", 0), tp("topic05", 0), tp("topic06", 0)))));
-
-        Map> assignment = assignor.assign(partitionsPerTopic, subscriptions);
-        verifyValidityAndBalance(subscriptions, assignment);
-    }
+        partitionsPerTopic.put(topic, 3);
+        subscriptions.put(consumer1, new Subscription(topics(topic)));
+        subscriptions.put(consumer2, new Subscription(topics(topic)));
+        subscriptions.put(consumer3, new Subscription(topics(topic)));
 
-    @Test
-    public void testStickiness() {
-        Map partitionsPerTopic = new HashMap<>();
-        partitionsPerTopic.put("topic01", 3);
-        Map subscriptions = new HashMap<>();
-        subscriptions.put("consumer01", new Subscription(topics("topic01")));
-        subscriptions.put("consumer02", new Subscription(topics("topic01")));
-        subscriptions.put("consumer03", new Subscription(topics("topic01")));
-        subscriptions.put("consumer04", new Subscription(topics("topic01")));
+        TopicPartition tp0 = new TopicPartition(topic, 0);
+        TopicPartition tp1 = new TopicPartition(topic, 1);
+        TopicPartition tp2 = new TopicPartition(topic, 2);
 
+        List c1partitions0 = partitions(tp0, tp2);
+        List c2partitions0 = partitions(tp1);
+        subscriptions.put(consumer1, buildSubscriptionWithGeneration(topics(topic), c1partitions0, 1));
+        subscriptions.put(consumer2, buildSubscriptionWithOldSchema(topics(topic), c2partitions0));
         Map> assignment = assignor.assign(partitionsPerTopic, subscriptions);
-        verifyValidityAndBalance(subscriptions, assignment);
-        Map partitionsAssigned = new HashMap<>();
-
-        Set>> assignments = assignment.entrySet();
-        for (Map.Entry> entry: assignments) {
-            String consumer = entry.getKey();
-            List topicPartitions = entry.getValue();
-            int size = topicPartitions.size();
-            assertTrue("Consumer " + consumer + " is assigned more topic partitions than expected.", size <= 1);
-            if (size == 1)
-                partitionsAssigned.put(consumer, topicPartitions.get(0));
-        }
-
-        // removing the potential group leader
-        subscriptions.remove("consumer01");
-        subscriptions.put("consumer02",
-                new Subscription(topics("topic01"),
-                        StickyAssignor.serializeTopicPartitionAssignment(assignment.get("consumer02"))));
-        subscriptions.put("consumer03",
-                new Subscription(topics("topic01"),
-                        StickyAssignor.serializeTopicPartitionAssignment(assignment.get("consumer03"))));
-        subscriptions.put("consumer04",
-                new Subscription(topics("topic01"),
-                        StickyAssignor.serializeTopicPartitionAssignment(assignment.get("consumer04"))));
-
-        assignment = assignor.assign(partitionsPerTopic, subscriptions);
-        verifyValidityAndBalance(subscriptions, assignment);
-        assertTrue(assignor.isSticky());
-
-        assignments = assignment.entrySet();
-        for (Map.Entry> entry: assignments) {
-            String consumer = entry.getKey();
-            List topicPartitions = entry.getValue();
-            assertEquals("Consumer " + consumer + " is assigned more topic partitions than expected.", 1, topicPartitions.size());
-            assertTrue("Stickiness was not honored for consumer " + consumer,
-                (!partitionsAssigned.containsKey(consumer)) || (assignment.get(consumer).contains(partitionsAssigned.get(consumer))));
-        }
-    }
-
-    @Test
-    public void testAssignmentUpdatedForDeletedTopic() {
-        String consumerId = "consumer";
+        List c1partitions = assignment.get(consumer1);
+        List c2partitions = assignment.get(consumer2);
+        List c3partitions = assignment.get(consumer3);
 
-        Map partitionsPerTopic = new HashMap<>();
-        partitionsPerTopic.put("topic01", 1);
-        partitionsPerTopic.put("topic03", 100);
-        Map subscriptions =
-                Collections.singletonMap(consumerId, new Subscription(topics("topic01", "topic02", "topic03")));
-
-        Map> assignment = assignor.assign(partitionsPerTopic, subscriptions);
-        assertEquals(assignment.values().stream().mapToInt(topicPartitions -> topicPartitions.size()).sum(), 1 + 100);
-        assertEquals(Collections.singleton(consumerId), assignment.keySet());
+        assertTrue(c1partitions.size() == 1 && c2partitions.size() == 1 && c3partitions.size() == 1);
+        assertTrue(c1partitions0.containsAll(c1partitions));
+        assertTrue(c2partitions0.containsAll(c2partitions));
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
         assertTrue(isFullyBalanced(assignment));
+        assertTrue(assignor.isSticky());
     }
 
-    @Test
-    public void testNoExceptionThrownWhenOnlySubscribedTopicDeleted() {
-        String topic = "topic01";
-        String consumer = "consumer01";
-        Map partitionsPerTopic = new HashMap<>();
-        partitionsPerTopic.put(topic, 3);
-        Map subscriptions = new HashMap<>();
-        subscriptions.put(consumer, new Subscription(topics(topic)));
-        Map> assignment = assignor.assign(partitionsPerTopic, subscriptions);
-        subscriptions.put(consumer, new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment(assignment.get(consumer))));
-
-        assignment = assignor.assign(Collections.emptyMap(), subscriptions);
-        assertEquals(assignment.size(), 1);
-        assertTrue(assignment.get(consumer).isEmpty());
-    }
-
-    private String getTopicName(int i, int maxNum) {
-        return getCanonicalName("t", i, maxNum);
-    }
-
-    private String getConsumerName(int i, int maxNum) {
-        return getCanonicalName("c", i, maxNum);
-    }
-
-    private String getCanonicalName(String str, int i, int maxNum) {
-        return str + pad(i, Integer.toString(maxNum).length());
-    }
-
-    private String pad(int num, int digits) {
-        StringBuilder sb = new StringBuilder();
-        int iDigits = Integer.toString(num).length();
-
-        for (int i = 1; i <= digits - iDigits; ++i)
-            sb.append("0");
-
-        sb.append(num);
-        return sb.toString();
-    }
-
-    private static List topics(String... topics) {
-        return Arrays.asList(topics);
-    }
-
-    private static List partitions(TopicPartition... partitions) {
-        return Arrays.asList(partitions);
-    }
-
-    private static TopicPartition tp(String topic, int partition) {
-        return new TopicPartition(topic, partition);
+    private Subscription buildSubscriptionWithGeneration(List topics, List partitions, int generation) {
+        return new Subscription(topics,
+            serializeTopicPartitionAssignment(new MemberData(partitions, Optional.of(generation))));
     }
 
-    private static boolean isFullyBalanced(Map> assignment) {
-        int min = Integer.MAX_VALUE;
-        int max = Integer.MIN_VALUE;
-        for (List topicPartitions: assignment.values()) {
-            int size = topicPartitions.size();
-            if (size < min)
-                min = size;
-            if (size > max)
-                max = size;
+    private static Subscription buildSubscriptionWithOldSchema(List topics, List partitions) {
+        Struct struct = new Struct(StickyAssignor.STICKY_ASSIGNOR_USER_DATA_V0);
+        List topicAssignments = new ArrayList<>();
+        for (Map.Entry> topicEntry : CollectionUtils.groupPartitionsByTopic(partitions).entrySet()) {
+            Struct topicAssignment = new Struct(StickyAssignor.TOPIC_ASSIGNMENT);
+            topicAssignment.set(StickyAssignor.TOPIC_KEY_NAME, topicEntry.getKey());
+            topicAssignment.set(StickyAssignor.PARTITIONS_KEY_NAME, topicEntry.getValue().toArray());
+            topicAssignments.add(topicAssignment);
         }
-        return max - min <= 1;
-    }
-
-    private static List getRandomSublist(ArrayList list) {
-        List selectedItems = new ArrayList<>(list);
-        int len = list.size();
-        Random random = new Random();
-        int howManyToRemove = random.nextInt(len);
-
-        for (int i = 1; i <= howManyToRemove; ++i)
-            selectedItems.remove(random.nextInt(selectedItems.size()));
+        struct.set(StickyAssignor.TOPIC_PARTITIONS_KEY_NAME, topicAssignments.toArray());
+        ByteBuffer buffer = ByteBuffer.allocate(StickyAssignor.STICKY_ASSIGNOR_USER_DATA_V0.sizeOf(struct));
+        StickyAssignor.STICKY_ASSIGNOR_USER_DATA_V0.write(buffer, struct);
+        buffer.flip();
 
-        return selectedItems;
-    }
-
-    /**
-     * Verifies that the given assignment is valid and balanced with respect to the given subscriptions
-     * Validity requirements:
-     * - each consumer is subscribed to topics of all partitions assigned to it, and
-     * - each partition is assigned to no more than one consumer
-     * Balance requirements:
-     * - the assignment is fully balanced (the numbers of topic partitions assigned to consumers differ by at most one), or
-     * - there is no topic partition that can be moved from one consumer to another with 2+ fewer topic partitions
-     *
-     * @param subscriptions: topic subscriptions of each consumer
-     * @param assignments: given assignment for balance check
-     */
-    private static void verifyValidityAndBalance(Map subscriptions, Map> assignments) {
-        int size = subscriptions.size();
-        assert size == assignments.size();
-
-        List consumers = Utils.sorted(assignments.keySet());
-
-        for (int i = 0; i < size; ++i) {
-            String consumer = consumers.get(i);
-            List partitions = assignments.get(consumer);
-            for (TopicPartition partition: partitions)
-                assertTrue("Error: Partition " + partition + "is assigned to c" + i + ", but it is not subscribed to Topic t" + partition.topic()
-                        + "\nSubscriptions: " + subscriptions.toString() + "\nAssignments: " + assignments.toString(),
-                        subscriptions.get(consumer).topics().contains(partition.topic()));
-
-            if (i == size - 1)
-                continue;
-
-            for (int j = i + 1; j < size; ++j) {
-                String otherConsumer = consumers.get(j);
-                List otherPartitions = assignments.get(otherConsumer);
-
-                Set intersection = new HashSet<>(partitions);
-                intersection.retainAll(otherPartitions);
-                assertTrue("Error: Consumers c" + i + " and c" + j + " have common partitions assigned to them: " + intersection.toString()
-                        + "\nSubscriptions: " + subscriptions.toString() + "\nAssignments: " + assignments.toString(),
-                        intersection.isEmpty());
-
-                int len = partitions.size();
-                int otherLen = otherPartitions.size();
-
-                if (Math.abs(len - otherLen) <= 1)
-                    continue;
-
-                Map> map = CollectionUtils.groupPartitionsByTopic(partitions);
-                Map> otherMap = CollectionUtils.groupPartitionsByTopic(otherPartitions);
-
-                if (len > otherLen) {
-                    for (String topic: map.keySet())
-                        assertTrue("Error: Some partitions can be moved from c" + i + " to c" + j + " to achieve a better balance"
-                                + "\nc" + i + " has " + len + " partitions, and c" + j + " has " + otherLen + " partitions."
-                                + "\nSubscriptions: " + subscriptions.toString() + "\nAssignments: " + assignments.toString(),
-                                !otherMap.containsKey(topic));
-                }
-
-                if (otherLen > len) {
-                    for (String topic: otherMap.keySet())
-                        assertTrue("Error: Some partitions can be moved from c" + j + " to c" + i + " to achieve a better balance"
-                                + "\nc" + i + " has " + len + " partitions, and c" + j + " has " + otherLen + " partitions."
-                                + "\nSubscriptions: " + subscriptions.toString() + "\nAssignments: " + assignments.toString(),
-                                !map.containsKey(topic));
-                }
-            }
-        }
+        return new Subscription(topics, buffer);
     }
-}
+}
\ No newline at end of file
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 9a68db706b4df..6a5b1f0bffbc2 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
@@ -16,12 +16,22 @@
  */
 package org.apache.kafka.clients.consumer.internals;
 
+import org.apache.kafka.clients.GroupRebalanceConfig;
 import org.apache.kafka.clients.MockClient;
 import org.apache.kafka.clients.consumer.OffsetResetStrategy;
 import org.apache.kafka.common.Node;
 import org.apache.kafka.common.errors.AuthenticationException;
+import org.apache.kafka.common.errors.FencedInstanceIdException;
+import org.apache.kafka.common.errors.UnknownMemberIdException;
 import org.apache.kafka.common.errors.WakeupException;
 import org.apache.kafka.common.internals.ClusterResourceListeners;
+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.LeaveGroupResponseData;
+import org.apache.kafka.common.message.LeaveGroupResponseData.MemberResponse;
+import org.apache.kafka.common.message.SyncGroupResponseData;
+import org.apache.kafka.common.metrics.KafkaMetric;
 import org.apache.kafka.common.metrics.Metrics;
 import org.apache.kafka.common.protocol.Errors;
 import org.apache.kafka.common.requests.AbstractRequest;
@@ -30,6 +40,8 @@
 import org.apache.kafka.common.requests.HeartbeatResponse;
 import org.apache.kafka.common.requests.JoinGroupRequest;
 import org.apache.kafka.common.requests.JoinGroupResponse;
+import org.apache.kafka.common.requests.LeaveGroupRequest;
+import org.apache.kafka.common.requests.LeaveGroupResponse;
 import org.apache.kafka.common.requests.SyncGroupRequest;
 import org.apache.kafka.common.requests.SyncGroupResponse;
 import org.apache.kafka.common.utils.LogContext;
@@ -41,10 +53,12 @@
 import org.junit.Test;
 
 import java.nio.ByteBuffer;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
@@ -55,8 +69,10 @@
 import static java.util.Collections.emptyMap;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNotSame;
 import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
@@ -66,42 +82,149 @@ public class AbstractCoordinatorTest {
     private static final int SESSION_TIMEOUT_MS = 10000;
     private static final int HEARTBEAT_INTERVAL_MS = 3000;
     private static final int RETRY_BACKOFF_MS = 100;
-    private static final int LONG_RETRY_BACKOFF_MS = 10000;
     private static final int REQUEST_TIMEOUT_MS = 40000;
     private static final String GROUP_ID = "dummy-group";
     private static final String METRIC_GROUP_PREFIX = "consumer";
 
-    private MockClient mockClient;
-    private MockTime mockTime;
     private Node node;
+    private Metrics metrics;
+    private MockTime mockTime;
     private Node coordinatorNode;
-    private ConsumerNetworkClient consumerClient;
+    private MockClient mockClient;
     private DummyCoordinator coordinator;
+    private ConsumerNetworkClient consumerClient;
+
+    private final String memberId = "memberId";
+    private final String leaderId = "leaderId";
+    private final int defaultGeneration = -1;
 
     private void setupCoordinator() {
-        setupCoordinator(RETRY_BACKOFF_MS, REBALANCE_TIMEOUT_MS);
+        setupCoordinator(RETRY_BACKOFF_MS, REBALANCE_TIMEOUT_MS,
+            Optional.empty());
     }
 
     private void setupCoordinator(int retryBackoffMs) {
-        setupCoordinator(retryBackoffMs, REBALANCE_TIMEOUT_MS);
+        setupCoordinator(retryBackoffMs, REBALANCE_TIMEOUT_MS,
+            Optional.empty());
     }
 
-    private void setupCoordinator(int retryBackoffMs, int rebalanceTimeoutMs) {
+    private void setupCoordinator(int retryBackoffMs, int rebalanceTimeoutMs, Optional groupInstanceId) {
         LogContext logContext = new LogContext();
         this.mockTime = new MockTime();
         ConsumerMetadata metadata = new ConsumerMetadata(retryBackoffMs, 60 * 60 * 1000L,
-                false, new SubscriptionState(logContext, OffsetResetStrategy.EARLIEST),
+                false, false, new SubscriptionState(logContext, OffsetResetStrategy.EARLIEST),
                 logContext, new ClusterResourceListeners());
 
         this.mockClient = new MockClient(mockTime, metadata);
-        this.consumerClient = new ConsumerNetworkClient(logContext, mockClient, metadata, mockTime,
-                retryBackoffMs, REQUEST_TIMEOUT_MS, HEARTBEAT_INTERVAL_MS);
-        Metrics metrics = new Metrics();
+        this.consumerClient = new ConsumerNetworkClient(logContext,
+                                                        mockClient,
+                                                        metadata,
+                                                        mockTime,
+                                                        retryBackoffMs,
+                                                        REQUEST_TIMEOUT_MS,
+                                                        HEARTBEAT_INTERVAL_MS);
+        metrics = new Metrics(mockTime);
 
         mockClient.updateMetadata(TestUtils.metadataUpdateWith(1, emptyMap()));
         this.node = metadata.fetch().nodes().get(0);
         this.coordinatorNode = new Node(Integer.MAX_VALUE - node.id(), node.host(), node.port());
-        this.coordinator = new DummyCoordinator(consumerClient, metrics, mockTime, rebalanceTimeoutMs, retryBackoffMs);
+
+        GroupRebalanceConfig rebalanceConfig = new GroupRebalanceConfig(SESSION_TIMEOUT_MS,
+                                                                        rebalanceTimeoutMs,
+                                                                        HEARTBEAT_INTERVAL_MS,
+                                                                        GROUP_ID,
+                                                                        groupInstanceId,
+                                                                        retryBackoffMs,
+                                                                        !groupInstanceId.isPresent());
+        this.coordinator = new DummyCoordinator(rebalanceConfig,
+                                                consumerClient,
+                                                metrics,
+                                                mockTime);
+    }
+
+    @Test
+    public void testMetrics() {
+        setupCoordinator();
+
+        assertNotNull(getMetric("heartbeat-response-time-max"));
+        assertNotNull(getMetric("heartbeat-rate"));
+        assertNotNull(getMetric("heartbeat-total"));
+        assertNotNull(getMetric("last-heartbeat-seconds-ago"));
+        assertNotNull(getMetric("join-time-avg"));
+        assertNotNull(getMetric("join-time-max"));
+        assertNotNull(getMetric("join-rate"));
+        assertNotNull(getMetric("join-total"));
+        assertNotNull(getMetric("sync-time-avg"));
+        assertNotNull(getMetric("sync-time-max"));
+        assertNotNull(getMetric("sync-rate"));
+        assertNotNull(getMetric("sync-total"));
+        assertNotNull(getMetric("rebalance-latency-avg"));
+        assertNotNull(getMetric("rebalance-latency-max"));
+        assertNotNull(getMetric("rebalance-latency-total"));
+        assertNotNull(getMetric("rebalance-rate-per-hour"));
+        assertNotNull(getMetric("rebalance-total"));
+        assertNotNull(getMetric("last-rebalance-seconds-ago"));
+        assertNotNull(getMetric("failed-rebalance-rate-per-hour"));
+        assertNotNull(getMetric("failed-rebalance-total"));
+
+        metrics.sensor("heartbeat-latency").record(1.0d);
+        metrics.sensor("heartbeat-latency").record(6.0d);
+        metrics.sensor("heartbeat-latency").record(2.0d);
+
+        assertEquals(6.0d, getMetric("heartbeat-response-time-max").metricValue());
+        assertEquals(0.1d, getMetric("heartbeat-rate").metricValue());
+        assertEquals(3.0d, getMetric("heartbeat-total").metricValue());
+
+        assertEquals(-1.0d, getMetric("last-heartbeat-seconds-ago").metricValue());
+        coordinator.heartbeat().sentHeartbeat(mockTime.milliseconds());
+        assertEquals(0.0d, getMetric("last-heartbeat-seconds-ago").metricValue());
+        mockTime.sleep(10 * 1000L);
+        assertEquals(10.0d, getMetric("last-heartbeat-seconds-ago").metricValue());
+
+        metrics.sensor("join-latency").record(1.0d);
+        metrics.sensor("join-latency").record(6.0d);
+        metrics.sensor("join-latency").record(2.0d);
+
+        assertEquals(3.0d, getMetric("join-time-avg").metricValue());
+        assertEquals(6.0d, getMetric("join-time-max").metricValue());
+        assertEquals(0.1d, getMetric("join-rate").metricValue());
+        assertEquals(3.0d, getMetric("join-total").metricValue());
+
+        metrics.sensor("sync-latency").record(1.0d);
+        metrics.sensor("sync-latency").record(6.0d);
+        metrics.sensor("sync-latency").record(2.0d);
+
+        assertEquals(3.0d, getMetric("sync-time-avg").metricValue());
+        assertEquals(6.0d, getMetric("sync-time-max").metricValue());
+        assertEquals(0.1d, getMetric("sync-rate").metricValue());
+        assertEquals(3.0d, getMetric("sync-total").metricValue());
+
+        metrics.sensor("rebalance-latency").record(1.0d);
+        metrics.sensor("rebalance-latency").record(6.0d);
+        metrics.sensor("rebalance-latency").record(2.0d);
+
+        assertEquals(3.0d, getMetric("rebalance-latency-avg").metricValue());
+        assertEquals(6.0d, getMetric("rebalance-latency-max").metricValue());
+        assertEquals(9.0d, getMetric("rebalance-latency-total").metricValue());
+        assertEquals(360.0d, getMetric("rebalance-rate-per-hour").metricValue());
+        assertEquals(3.0d, getMetric("rebalance-total").metricValue());
+
+        metrics.sensor("failed-rebalance").record(1.0d);
+        metrics.sensor("failed-rebalance").record(6.0d);
+        metrics.sensor("failed-rebalance").record(2.0d);
+
+        assertEquals(360.0d, getMetric("failed-rebalance-rate-per-hour").metricValue());
+        assertEquals(3.0d, getMetric("failed-rebalance-total").metricValue());
+
+        assertEquals(-1.0d, getMetric("last-rebalance-seconds-ago").metricValue());
+        coordinator.setLastRebalanceTime(mockTime.milliseconds());
+        assertEquals(0.0d, getMetric("last-rebalance-seconds-ago").metricValue());
+        mockTime.sleep(10 * 1000L);
+        assertEquals(10.0d, getMetric("last-rebalance-seconds-ago").metricValue());
+    }
+
+    private KafkaMetric getMetric(final String name) {
+        return metrics.metrics().get(metrics.metricName(name, "consumer-coordinator-metrics"));
     }
 
     @Test
@@ -137,7 +260,7 @@ public void testTimeoutAndRetryJoinGroupIfNeeded() throws Exception {
             assertFalse(firstAttempt.get());
             assertTrue(consumerClient.hasPendingRequests(coordinatorNode));
 
-            mockClient.respond(joinGroupFollowerResponse(1, "memberId", "leaderId", Errors.NONE));
+            mockClient.respond(joinGroupFollowerResponse(1, memberId, leaderId, Errors.NONE));
             mockClient.prepareResponse(syncGroupResponse(Errors.NONE));
 
             Timer secondAttemptTimer = mockTime.timer(REQUEST_TIMEOUT_MS);
@@ -156,10 +279,7 @@ public void testGroupMaxSizeExceptionIsFatal() {
         mockClient.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
         coordinator.ensureCoordinatorReady(mockTime.timer(0));
 
-        final String memberId = "memberId";
-        final int generation = -1;
-
-        mockClient.prepareResponse(joinGroupFollowerResponse(generation, memberId, JoinGroupResponse.UNKNOWN_MEMBER_ID, Errors.GROUP_MAX_SIZE_REACHED));
+        mockClient.prepareResponse(joinGroupFollowerResponse(defaultGeneration, memberId, JoinGroupResponse.UNKNOWN_MEMBER_ID, Errors.GROUP_MAX_SIZE_REACHED));
 
         RequestFuture future = coordinator.sendJoinGroupRequest();
         assertTrue(consumerClient.poll(future, mockTime.timer(REQUEST_TIMEOUT_MS)));
@@ -169,7 +289,8 @@ public void testGroupMaxSizeExceptionIsFatal() {
 
     @Test
     public void testJoinGroupRequestTimeout() {
-        setupCoordinator(RETRY_BACKOFF_MS, REBALANCE_TIMEOUT_MS);
+        setupCoordinator(RETRY_BACKOFF_MS, REBALANCE_TIMEOUT_MS,
+            Optional.empty());
         mockClient.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
         coordinator.ensureCoordinatorReady(mockTime.timer(0));
 
@@ -186,7 +307,8 @@ public void testJoinGroupRequestTimeout() {
     public void testJoinGroupRequestMaxTimeout() {
         // Ensure we can handle the maximum allowed rebalance timeout
 
-        setupCoordinator(RETRY_BACKOFF_MS, Integer.MAX_VALUE);
+        setupCoordinator(RETRY_BACKOFF_MS, Integer.MAX_VALUE,
+            Optional.empty());
         mockClient.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
         coordinator.ensureCoordinatorReady(mockTime.timer(0));
 
@@ -203,23 +325,14 @@ public void testJoinGroupRequestWithMemberIdRequired() {
         mockClient.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
         coordinator.ensureCoordinatorReady(mockTime.timer(0));
 
-        final String memberId = "memberId";
-        final int generation = -1;
+        mockClient.prepareResponse(joinGroupFollowerResponse(defaultGeneration, memberId, JoinGroupResponse.UNKNOWN_MEMBER_ID, Errors.MEMBER_ID_REQUIRED));
 
-        mockClient.prepareResponse(joinGroupFollowerResponse(generation, memberId, JoinGroupResponse.UNKNOWN_MEMBER_ID, Errors.MEMBER_ID_REQUIRED));
-
-        mockClient.prepareResponse(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                if (!(body instanceof JoinGroupRequest)) {
-                    return false;
-                }
-                JoinGroupRequest joinGroupRequest = (JoinGroupRequest) body;
-                if (!joinGroupRequest.memberId().equals(memberId)) {
-                    return false;
-                }
-                return true;
+        mockClient.prepareResponse(body -> {
+            if (!(body instanceof JoinGroupRequest)) {
+                return false;
             }
+            JoinGroupRequest joinGroupRequest = (JoinGroupRequest) body;
+            return joinGroupRequest.data().memberId().equals(memberId);
         }, joinGroupResponse(Errors.UNKNOWN_MEMBER_ID));
 
         RequestFuture future = coordinator.sendJoinGroupRequest();
@@ -227,29 +340,189 @@ public boolean matches(AbstractRequest body) {
         assertEquals(Errors.MEMBER_ID_REQUIRED.message(), future.exception().getMessage());
         assertTrue(coordinator.rejoinNeededOrPending());
         assertTrue(coordinator.hasValidMemberId());
-        assertTrue(coordinator.hasMatchingGenerationId(generation));
+        assertTrue(coordinator.hasMatchingGenerationId(defaultGeneration));
         future = coordinator.sendJoinGroupRequest();
         assertTrue(consumerClient.poll(future, mockTime.timer(REBALANCE_TIMEOUT_MS)));
     }
 
+    @Test
+    public void testJoinGroupRequestWithFencedInstanceIdException() {
+        setupCoordinator();
+        mockClient.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+        coordinator.ensureCoordinatorReady(mockTime.timer(0));
+
+        mockClient.prepareResponse(joinGroupFollowerResponse(defaultGeneration, memberId, JoinGroupResponse.UNKNOWN_MEMBER_ID, Errors.FENCED_INSTANCE_ID));
+
+        RequestFuture future = coordinator.sendJoinGroupRequest();
+        assertTrue(consumerClient.poll(future, mockTime.timer(REQUEST_TIMEOUT_MS)));
+        assertEquals(Errors.FENCED_INSTANCE_ID.message(), future.exception().getMessage());
+        // Make sure the exception is fatal.
+        assertFalse(future.isRetriable());
+    }
+
+    @Test
+    public void testSyncGroupRequestWithFencedInstanceIdException() {
+        setupCoordinator();
+        mockClient.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+
+        final int generation = -1;
+
+        mockClient.prepareResponse(joinGroupFollowerResponse(generation, memberId, JoinGroupResponse.UNKNOWN_MEMBER_ID, Errors.NONE));
+        mockClient.prepareResponse(syncGroupResponse(Errors.FENCED_INSTANCE_ID));
+
+        assertThrows(FencedInstanceIdException.class, () -> coordinator.ensureActiveGroup());
+    }
+
+    @Test
+    public void testHeartbeatRequestWithFencedInstanceIdException() throws InterruptedException {
+        setupCoordinator();
+        mockClient.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+
+        final int generation = -1;
+
+        mockClient.prepareResponse(joinGroupFollowerResponse(generation, memberId, JoinGroupResponse.UNKNOWN_MEMBER_ID, Errors.NONE));
+        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);
+        }
+    }
+
+    @Test
+    public void testJoinGroupRequestWithGroupInstanceIdNotFound() {
+        setupCoordinator();
+        mockClient.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+        coordinator.ensureCoordinatorReady(mockTime.timer(0));
+
+        mockClient.prepareResponse(joinGroupFollowerResponse(defaultGeneration, memberId, JoinGroupResponse.UNKNOWN_MEMBER_ID, Errors.UNKNOWN_MEMBER_ID));
+
+        RequestFuture future = coordinator.sendJoinGroupRequest();
+
+        assertTrue(consumerClient.poll(future, mockTime.timer(REQUEST_TIMEOUT_MS)));
+        assertEquals(Errors.UNKNOWN_MEMBER_ID.message(), future.exception().getMessage());
+        assertTrue(coordinator.rejoinNeededOrPending());
+        assertTrue(coordinator.hasUnknownGeneration());
+    }
+
+    @Test
+    public void testLeaveGroupSentWithGroupInstanceIdUnSet() {
+        checkLeaveGroupRequestSent(Optional.empty());
+        checkLeaveGroupRequestSent(Optional.of("groupInstanceId"));
+    }
+
+    private void checkLeaveGroupRequestSent(Optional groupInstanceId) {
+        setupCoordinator(RETRY_BACKOFF_MS, Integer.MAX_VALUE, groupInstanceId);
+
+        mockClient.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+        mockClient.prepareResponse(joinGroupFollowerResponse(1, memberId, leaderId, Errors.NONE));
+        mockClient.prepareResponse(syncGroupResponse(Errors.NONE));
+
+        final RuntimeException e = new RuntimeException();
+
+        // raise the error when the coordinator tries to send leave group request.
+        mockClient.prepareResponse(body -> {
+            if (body instanceof LeaveGroupRequest)
+                throw e;
+            return false;
+        }, heartbeatResponse(Errors.UNKNOWN_SERVER_ERROR));
+
+        try {
+            coordinator.ensureActiveGroup();
+            coordinator.close();
+            if (coordinator.isDynamicMember()) {
+                fail("Expected leavegroup to raise an error.");
+            }
+        } catch (RuntimeException exception) {
+            if (coordinator.isDynamicMember()) {
+                assertEquals(exception, e);
+            } else {
+                fail("Coordinator with group.instance.id set shouldn't send leave group request.");
+            }
+        }
+    }
+
+    @Test
+    public void testHandleNormalLeaveGroupResponse() {
+        MemberResponse memberResponse = new MemberResponse()
+                                            .setMemberId(memberId)
+                                            .setErrorCode(Errors.NONE.code());
+        LeaveGroupResponse response =
+            leaveGroupResponse(Collections.singletonList(memberResponse));
+        RequestFuture leaveGroupFuture = setupLeaveGroup(response);
+        assertNotNull(leaveGroupFuture);
+        assertTrue(leaveGroupFuture.succeeded());
+    }
+
+    @Test
+    public void testHandleMultipleMembersLeaveGroupResponse() {
+        MemberResponse memberResponse = new MemberResponse()
+                                            .setMemberId(memberId)
+                                            .setErrorCode(Errors.NONE.code());
+        LeaveGroupResponse response =
+            leaveGroupResponse(Arrays.asList(memberResponse, memberResponse));
+        RequestFuture leaveGroupFuture = setupLeaveGroup(response);
+        assertNotNull(leaveGroupFuture);
+        assertTrue(leaveGroupFuture.exception() instanceof IllegalStateException);
+    }
+
+    @Test
+    public void testHandleLeaveGroupResponseWithEmptyMemberResponse() {
+        LeaveGroupResponse response =
+            leaveGroupResponse(Collections.emptyList());
+        RequestFuture leaveGroupFuture = setupLeaveGroup(response);
+        assertNotNull(leaveGroupFuture);
+        assertTrue(leaveGroupFuture.succeeded());
+    }
+
+    @Test
+    public void testHandleLeaveGroupResponseWithException() {
+        MemberResponse memberResponse = new MemberResponse()
+                                            .setMemberId(memberId)
+                                            .setErrorCode(Errors.UNKNOWN_MEMBER_ID.code());
+        LeaveGroupResponse response =
+            leaveGroupResponse(Collections.singletonList(memberResponse));
+        RequestFuture leaveGroupFuture = setupLeaveGroup(response);
+        assertNotNull(leaveGroupFuture);
+        assertTrue(leaveGroupFuture.exception() instanceof UnknownMemberIdException);
+    }
+
+    private RequestFuture setupLeaveGroup(LeaveGroupResponse leaveGroupResponse) {
+        setupCoordinator(RETRY_BACKOFF_MS, Integer.MAX_VALUE, Optional.empty());
+
+        mockClient.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+        mockClient.prepareResponse(joinGroupFollowerResponse(1, memberId, leaderId, Errors.NONE));
+        mockClient.prepareResponse(syncGroupResponse(Errors.NONE));
+        mockClient.prepareResponse(leaveGroupResponse);
+
+        coordinator.ensureActiveGroup();
+        return coordinator.maybeLeaveGroup("test maybe leave group");
+    }
+
     @Test
     public void testUncaughtExceptionInHeartbeatThread() throws Exception {
         setupCoordinator();
 
         mockClient.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        mockClient.prepareResponse(joinGroupFollowerResponse(1, "memberId", "leaderId", Errors.NONE));
+        mockClient.prepareResponse(joinGroupFollowerResponse(1, memberId, leaderId, Errors.NONE));
         mockClient.prepareResponse(syncGroupResponse(Errors.NONE));
 
         final RuntimeException e = new RuntimeException();
 
         // raise the error when the background thread tries to send a heartbeat
-        mockClient.prepareResponse(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                if (body instanceof HeartbeatRequest)
-                    throw e;
-                return false;
-            }
+        mockClient.prepareResponse(body -> {
+            if (body instanceof HeartbeatRequest)
+                throw e;
+            return false;
         }, heartbeatResponse(Errors.UNKNOWN_SERVER_ERROR));
 
         try {
@@ -268,21 +541,19 @@ public boolean matches(AbstractRequest body) {
 
     @Test
     public void testPollHeartbeatAwakesHeartbeatThread() throws Exception {
-        setupCoordinator(LONG_RETRY_BACKOFF_MS);
+        final int longRetryBackoffMs = 10000;
+        setupCoordinator(longRetryBackoffMs);
 
         mockClient.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        mockClient.prepareResponse(joinGroupFollowerResponse(1, "memberId", "leaderId", Errors.NONE));
+        mockClient.prepareResponse(joinGroupFollowerResponse(1, memberId, leaderId, Errors.NONE));
         mockClient.prepareResponse(syncGroupResponse(Errors.NONE));
 
         coordinator.ensureActiveGroup();
 
         final CountDownLatch heartbeatDone = new CountDownLatch(1);
-        mockClient.prepareResponse(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                heartbeatDone.countDown();
-                return body instanceof HeartbeatRequest;
-            }
+        mockClient.prepareResponse(body -> {
+            heartbeatDone.countDown();
+            return body instanceof HeartbeatRequest;
         }, heartbeatResponse(Errors.NONE));
 
         mockTime.sleep(HEARTBEAT_INTERVAL_MS);
@@ -327,7 +598,7 @@ public boolean matches(AbstractRequest body) {
                     throw new WakeupException();
                 return isJoinGroupRequest;
             }
-        }, joinGroupFollowerResponse(1, "memberId", "leaderId", Errors.NONE));
+        }, joinGroupFollowerResponse(1, memberId, leaderId, Errors.NONE));
         mockClient.prepareResponse(syncGroupResponse(Errors.NONE));
         AtomicBoolean heartbeatReceived = prepareFirstHeartbeat();
 
@@ -365,7 +636,7 @@ public boolean matches(AbstractRequest body) {
                     throw new WakeupException();
                 return isJoinGroupRequest;
             }
-        }, joinGroupFollowerResponse(1, "memberId", "leaderId", Errors.NONE));
+        }, joinGroupFollowerResponse(1, memberId, leaderId, Errors.NONE));
         mockClient.prepareResponse(syncGroupResponse(Errors.NONE));
         AtomicBoolean heartbeatReceived = prepareFirstHeartbeat();
 
@@ -394,16 +665,13 @@ public void testWakeupAfterJoinGroupReceived() throws Exception {
         setupCoordinator();
 
         mockClient.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        mockClient.prepareResponse(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                boolean isJoinGroupRequest = body instanceof JoinGroupRequest;
-                if (isJoinGroupRequest)
-                    // wakeup after the request returns
-                    consumerClient.wakeup();
-                return isJoinGroupRequest;
-            }
-        }, joinGroupFollowerResponse(1, "memberId", "leaderId", Errors.NONE));
+        mockClient.prepareResponse(body -> {
+            boolean isJoinGroupRequest = body instanceof JoinGroupRequest;
+            if (isJoinGroupRequest)
+                // wakeup after the request returns
+                consumerClient.wakeup();
+            return isJoinGroupRequest;
+        }, joinGroupFollowerResponse(1, memberId, leaderId, Errors.NONE));
         mockClient.prepareResponse(syncGroupResponse(Errors.NONE));
         AtomicBoolean heartbeatReceived = prepareFirstHeartbeat();
 
@@ -430,16 +698,13 @@ public void testWakeupAfterJoinGroupReceivedExternalCompletion() throws Exceptio
         setupCoordinator();
 
         mockClient.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        mockClient.prepareResponse(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                boolean isJoinGroupRequest = body instanceof JoinGroupRequest;
-                if (isJoinGroupRequest)
-                    // wakeup after the request returns
-                    consumerClient.wakeup();
-                return isJoinGroupRequest;
-            }
-        }, joinGroupFollowerResponse(1, "memberId", "leaderId", Errors.NONE));
+        mockClient.prepareResponse(body -> {
+            boolean isJoinGroupRequest = body instanceof JoinGroupRequest;
+            if (isJoinGroupRequest)
+                // wakeup after the request returns
+                consumerClient.wakeup();
+            return isJoinGroupRequest;
+        }, joinGroupFollowerResponse(1, memberId, leaderId, Errors.NONE));
         mockClient.prepareResponse(syncGroupResponse(Errors.NONE));
         AtomicBoolean heartbeatReceived = prepareFirstHeartbeat();
 
@@ -468,7 +733,7 @@ public void testWakeupAfterSyncGroupSent() throws Exception {
         setupCoordinator();
 
         mockClient.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        mockClient.prepareResponse(joinGroupFollowerResponse(1, "memberId", "leaderId", Errors.NONE));
+        mockClient.prepareResponse(joinGroupFollowerResponse(1, memberId, leaderId, Errors.NONE));
         mockClient.prepareResponse(new MockClient.RequestMatcher() {
             private int invocations = 0;
             @Override
@@ -506,7 +771,7 @@ public void testWakeupAfterSyncGroupSentExternalCompletion() throws Exception {
         setupCoordinator();
 
         mockClient.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        mockClient.prepareResponse(joinGroupFollowerResponse(1, "memberId", "leaderId", Errors.NONE));
+        mockClient.prepareResponse(joinGroupFollowerResponse(1, memberId, leaderId, Errors.NONE));
         mockClient.prepareResponse(new MockClient.RequestMatcher() {
             private int invocations = 0;
             @Override
@@ -546,16 +811,13 @@ public void testWakeupAfterSyncGroupReceived() throws Exception {
         setupCoordinator();
 
         mockClient.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        mockClient.prepareResponse(joinGroupFollowerResponse(1, "memberId", "leaderId", Errors.NONE));
-        mockClient.prepareResponse(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                boolean isSyncGroupRequest = body instanceof SyncGroupRequest;
-                if (isSyncGroupRequest)
-                    // wakeup after the request returns
-                    consumerClient.wakeup();
-                return isSyncGroupRequest;
-            }
+        mockClient.prepareResponse(joinGroupFollowerResponse(1, memberId, leaderId, Errors.NONE));
+        mockClient.prepareResponse(body -> {
+            boolean isSyncGroupRequest = body instanceof SyncGroupRequest;
+            if (isSyncGroupRequest)
+                // wakeup after the request returns
+                consumerClient.wakeup();
+            return isSyncGroupRequest;
         }, syncGroupResponse(Errors.NONE));
         AtomicBoolean heartbeatReceived = prepareFirstHeartbeat();
 
@@ -582,16 +844,13 @@ public void testWakeupAfterSyncGroupReceivedExternalCompletion() throws Exceptio
         setupCoordinator();
 
         mockClient.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        mockClient.prepareResponse(joinGroupFollowerResponse(1, "memberId", "leaderId", Errors.NONE));
-        mockClient.prepareResponse(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                boolean isSyncGroupRequest = body instanceof SyncGroupRequest;
-                if (isSyncGroupRequest)
-                    // wakeup after the request returns
-                    consumerClient.wakeup();
-                return isSyncGroupRequest;
-            }
+        mockClient.prepareResponse(joinGroupFollowerResponse(1, memberId, leaderId, Errors.NONE));
+        mockClient.prepareResponse(body -> {
+            boolean isSyncGroupRequest = body instanceof SyncGroupRequest;
+            if (isSyncGroupRequest)
+                // wakeup after the request returns
+                consumerClient.wakeup();
+            return isSyncGroupRequest;
         }, syncGroupResponse(Errors.NONE));
         AtomicBoolean heartbeatReceived = prepareFirstHeartbeat();
 
@@ -619,7 +878,7 @@ public void testWakeupInOnJoinComplete() throws Exception {
 
         coordinator.wakeupOnJoinComplete = true;
         mockClient.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        mockClient.prepareResponse(joinGroupFollowerResponse(1, "memberId", "leaderId", Errors.NONE));
+        mockClient.prepareResponse(joinGroupFollowerResponse(1, memberId, leaderId, Errors.NONE));
         mockClient.prepareResponse(syncGroupResponse(Errors.NONE));
         AtomicBoolean heartbeatReceived = prepareFirstHeartbeat();
 
@@ -660,14 +919,11 @@ public void testAuthenticationErrorInEnsureCoordinatorReady() {
 
     private AtomicBoolean prepareFirstHeartbeat() {
         final AtomicBoolean heartbeatReceived = new AtomicBoolean(false);
-        mockClient.prepareResponse(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                boolean isHeartbeatRequest = body instanceof HeartbeatRequest;
-                if (isHeartbeatRequest)
-                    heartbeatReceived.set(true);
-                return isHeartbeatRequest;
-            }
+        mockClient.prepareResponse(body -> {
+            boolean isHeartbeatRequest = body instanceof HeartbeatRequest;
+            if (isHeartbeatRequest)
+                heartbeatReceived.set(true);
+            return isHeartbeatRequest;
         }, heartbeatResponse(Errors.UNKNOWN_SERVER_ERROR));
         return heartbeatReceived;
     }
@@ -683,16 +939,23 @@ public boolean conditionMet() {
     }
 
     private FindCoordinatorResponse groupCoordinatorResponse(Node node, Errors error) {
-        return new FindCoordinatorResponse(error, node);
+        return FindCoordinatorResponse.prepareResponse(error, node);
     }
 
     private HeartbeatResponse heartbeatResponse(Errors error) {
-        return new HeartbeatResponse(error);
+        return new HeartbeatResponse(new HeartbeatResponseData().setErrorCode(error.code()));
     }
 
     private JoinGroupResponse joinGroupFollowerResponse(int generationId, String memberId, String leaderId, Errors error) {
-        return new JoinGroupResponse(error, generationId, "dummy-subprotocol", memberId, leaderId,
-                Collections.emptyMap());
+        return new JoinGroupResponse(
+                new JoinGroupResponseData()
+                        .setErrorCode(error.code())
+                        .setGenerationId(generationId)
+                        .setProtocolName("dummy-subprotocol")
+                        .setMemberId(memberId)
+                        .setLeader(leaderId)
+                        .setMembers(Collections.emptyList())
+        );
     }
 
     private JoinGroupResponse joinGroupResponse(Errors error) {
@@ -701,7 +964,17 @@ private JoinGroupResponse joinGroupResponse(Errors error) {
     }
 
     private SyncGroupResponse syncGroupResponse(Errors error) {
-        return new SyncGroupResponse(error, ByteBuffer.allocate(0));
+        return new SyncGroupResponse(
+                new SyncGroupResponseData()
+                        .setErrorCode(error.code())
+                        .setAssignment(new byte[0])
+        );
+    }
+
+    private LeaveGroupResponse leaveGroupResponse(List members) {
+        return new LeaveGroupResponse(new LeaveGroupResponseData()
+                .setErrorCode(Errors.NONE.code())
+                .setMembers(members));
     }
 
     public static class DummyCoordinator extends AbstractCoordinator {
@@ -710,13 +983,11 @@ public static class DummyCoordinator extends AbstractCoordinator {
         private int onJoinCompleteInvokes = 0;
         private boolean wakeupOnJoinComplete = false;
 
-        public DummyCoordinator(ConsumerNetworkClient client,
+        public DummyCoordinator(GroupRebalanceConfig rebalanceConfig,
+                                ConsumerNetworkClient client,
                                 Metrics metrics,
-                                Time time,
-                                int rebalanceTimeoutMs,
-                                int retryBackoffMs) {
-            super(new LogContext(), client, GROUP_ID, rebalanceTimeoutMs, SESSION_TIMEOUT_MS,
-                    HEARTBEAT_INTERVAL_MS, metrics, METRIC_GROUP_PREFIX, time, retryBackoffMs, false);
+                                Time time) {
+            super(rebalanceConfig, new LogContext(), client, metrics, METRIC_GROUP_PREFIX, time);
         }
 
         @Override
@@ -725,15 +996,22 @@ protected String protocolType() {
         }
 
         @Override
-        protected List metadata() {
-            return Collections.singletonList(new JoinGroupRequest.ProtocolMetadata("dummy-subprotocol", EMPTY_DATA));
+        protected JoinGroupRequestData.JoinGroupRequestProtocolCollection metadata() {
+            return new JoinGroupRequestData.JoinGroupRequestProtocolCollection(
+                    Collections.singleton(new JoinGroupRequestData.JoinGroupRequestProtocol()
+                            .setName("dummy-subprotocol")
+                            .setMetadata(EMPTY_DATA.array())).iterator()
+            );
         }
 
         @Override
-        protected Map performAssignment(String leaderId, String protocol, Map allMemberMetadata) {
+        protected Map performAssignment(String leaderId,
+                                                            String protocol,
+                                                            List allMemberMetadata) {
             Map assignment = new HashMap<>();
-            for (Map.Entry metadata : allMemberMetadata.entrySet())
-                assignment.put(metadata.getKey(), EMPTY_DATA);
+            for (JoinGroupResponseData.JoinGroupResponseMember member : allMemberMetadata) {
+                assignment.put(member.memberId(), EMPTY_DATA);
+            }
             return assignment;
         }
 
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AbstractPartitionAssignorTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AbstractPartitionAssignorTest.java
new file mode 100644
index 0000000000000..e5f73295b5f78
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AbstractPartitionAssignorTest.java
@@ -0,0 +1,89 @@
+/*
+ * 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.consumer.internals.AbstractPartitionAssignor.MemberInfo;
+import org.apache.kafka.common.utils.Utils;
+import org.junit.Test;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Optional;
+import java.util.Random;
+
+import static org.junit.Assert.assertEquals;
+
+public class AbstractPartitionAssignorTest {
+
+    @Test
+    public void testMemberInfoSortingWithoutGroupInstanceId() {
+        MemberInfo m1 = new MemberInfo("a", Optional.empty());
+        MemberInfo m2 = new MemberInfo("b", Optional.empty());
+        MemberInfo m3 = new MemberInfo("c", Optional.empty());
+
+        List memberInfoList = Arrays.asList(m1, m2, m3);
+        assertEquals(memberInfoList, Utils.sorted(memberInfoList));
+    }
+
+    @Test
+    public void testMemberInfoSortingWithAllGroupInstanceId() {
+        MemberInfo m1 = new MemberInfo("a", Optional.of("y"));
+        MemberInfo m2 = new MemberInfo("b", Optional.of("z"));
+        MemberInfo m3 = new MemberInfo("c", Optional.of("x"));
+
+        List memberInfoList = Arrays.asList(m1, m2, m3);
+        assertEquals(Arrays.asList(m3, m1, m2), Utils.sorted(memberInfoList));
+    }
+
+    @Test
+    public void testMemberInfoSortingSomeGroupInstanceId() {
+        MemberInfo m1 = new MemberInfo("a", Optional.empty());
+        MemberInfo m2 = new MemberInfo("b", Optional.of("y"));
+        MemberInfo m3 = new MemberInfo("c", Optional.of("x"));
+
+        List memberInfoList = Arrays.asList(m1, m2, m3);
+        assertEquals(Arrays.asList(m3, m2, m1), Utils.sorted(memberInfoList));
+    }
+
+    @Test
+    public void testMergeSortManyMemberInfo() {
+        Random rand = new Random();
+        int bound = 2;
+        List memberInfoList = new ArrayList<>();
+        List staticMemberList = new ArrayList<>();
+        List dynamicMemberList = new ArrayList<>();
+        for (int i = 0; i < 100; i++) {
+            // Need to make sure all the ids are defined as 3-digits otherwise
+            // the comparison result will break.
+            String id = Integer.toString(i + 100);
+            Optional groupInstanceId = rand.nextInt(bound) < bound / 2 ?
+                                                       Optional.of(id) : Optional.empty();
+            MemberInfo m = new MemberInfo(id, groupInstanceId);
+            memberInfoList.add(m);
+            if (m.groupInstanceId.isPresent()) {
+                staticMemberList.add(m);
+            } else {
+                dynamicMemberList.add(m);
+            }
+        }
+        staticMemberList.addAll(dynamicMemberList);
+        Collections.shuffle(memberInfoList);
+        assertEquals(staticMemberList, Utils.sorted(memberInfoList));
+    }
+}
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
new file mode 100644
index 0000000000000..c7b45233ca1e9
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignorTest.java
@@ -0,0 +1,801 @@
+/*
+ * 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 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.Random;
+import java.util.Set;
+
+import org.apache.kafka.clients.consumer.ConsumerPartitionAssignor.Subscription;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.utils.CollectionUtils;
+import org.apache.kafka.common.utils.Utils;
+import org.junit.Before;
+import org.junit.Test;
+
+import static org.apache.kafka.common.utils.Utils.mkEntry;
+import static org.apache.kafka.common.utils.Utils.mkMap;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+public abstract class AbstractStickyAssignorTest {
+    protected AbstractStickyAssignor assignor;
+    protected String consumerId = "consumer";
+    protected Map subscriptions;
+    protected String topic = "topic";
+
+    protected abstract AbstractStickyAssignor createAssignor();
+
+    protected abstract Subscription buildSubscription(List topics, List partitions);
+
+    @Before
+    public void setUp() {
+        assignor = createAssignor();
+
+        if (subscriptions != null) {
+            subscriptions.clear();
+        } else {
+            subscriptions = new HashMap<>();
+        }
+    }
+
+    @Test
+    public void testOneConsumerNoTopic() {
+        Map partitionsPerTopic = new HashMap<>();
+        subscriptions = Collections.singletonMap(consumerId, new Subscription(Collections.emptyList()));
+
+        Map> assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        assertEquals(Collections.singleton(consumerId), assignment.keySet());
+        assertTrue(assignment.get(consumerId).isEmpty());
+
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
+        assertTrue(isFullyBalanced(assignment));
+    }
+
+    @Test
+    public void testOneConsumerNonexistentTopic() {
+        Map partitionsPerTopic = new HashMap<>();
+        partitionsPerTopic.put(topic, 0);
+        subscriptions = Collections.singletonMap(consumerId, new Subscription(topics(topic)));
+
+        Map> assignment = assignor.assign(partitionsPerTopic, subscriptions);
+
+        assertEquals(Collections.singleton(consumerId), assignment.keySet());
+        assertTrue(assignment.get(consumerId).isEmpty());
+
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
+        assertTrue(isFullyBalanced(assignment));
+    }
+
+    @Test
+    public void testOneConsumerOneTopic() {
+        Map partitionsPerTopic = new HashMap<>();
+        partitionsPerTopic.put(topic, 3);
+        subscriptions = Collections.singletonMap(consumerId, new Subscription(topics(topic)));
+
+        Map> assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        assertEquals(partitions(tp(topic, 0), tp(topic, 1), tp(topic, 2)), assignment.get(consumerId));
+
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
+        assertTrue(isFullyBalanced(assignment));
+    }
+
+    @Test
+    public void testOnlyAssignsPartitionsFromSubscribedTopics() {
+        String otherTopic = "other";
+
+        Map partitionsPerTopic = new HashMap<>();
+        partitionsPerTopic.put(topic, 2);
+        subscriptions = mkMap(
+                mkEntry(consumerId, buildSubscription(
+                        topics(topic),
+                        Arrays.asList(tp(topic, 0), tp(topic, 1), tp(otherTopic, 0), tp(otherTopic, 1)))
+                )
+        );
+
+        Map> assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        assertEquals(partitions(tp(topic, 0), tp(topic, 1)), assignment.get(consumerId));
+
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
+        assertTrue(isFullyBalanced(assignment));
+    }
+
+    @Test
+    public void testOneConsumerMultipleTopics() {
+        String topic1 = "topic1";
+        String topic2 = "topic2";
+
+        Map partitionsPerTopic = new HashMap<>();
+        partitionsPerTopic.put(topic1, 1);
+        partitionsPerTopic.put(topic2, 2);
+        subscriptions = Collections.singletonMap(consumerId, new Subscription(topics(topic1, topic2)));
+
+        Map> assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        assertEquals(partitions(tp(topic1, 0), tp(topic2, 0), tp(topic2, 1)), assignment.get(consumerId));
+
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
+        assertTrue(isFullyBalanced(assignment));
+    }
+
+    @Test
+    public void testTwoConsumersOneTopicOnePartition() {
+        String consumer1 = "consumer1";
+        String consumer2 = "consumer2";
+
+        Map partitionsPerTopic = new HashMap<>();
+        partitionsPerTopic.put(topic, 1);
+
+        subscriptions.put(consumer1, new Subscription(topics(topic)));
+        subscriptions.put(consumer2, new Subscription(topics(topic)));
+
+        Map> assignment = assignor.assign(partitionsPerTopic, subscriptions);
+
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
+        assertTrue(isFullyBalanced(assignment));
+    }
+
+    @Test
+    public void testTwoConsumersOneTopicTwoPartitions() {
+        String consumer1 = "consumer1";
+        String consumer2 = "consumer2";
+
+        Map partitionsPerTopic = new HashMap<>();
+        partitionsPerTopic.put(topic, 2);
+
+        subscriptions.put(consumer1, new Subscription(topics(topic)));
+        subscriptions.put(consumer2, new Subscription(topics(topic)));
+
+        Map> assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        assertEquals(partitions(tp(topic, 0)), assignment.get(consumer1));
+        assertEquals(partitions(tp(topic, 1)), assignment.get(consumer2));
+
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
+        assertTrue(isFullyBalanced(assignment));
+    }
+
+    @Test
+    public void testMultipleConsumersMixedTopicSubscriptions() {
+        String topic1 = "topic1";
+        String topic2 = "topic2";
+        String consumer1 = "consumer1";
+        String consumer2 = "consumer2";
+        String consumer3 = "consumer3";
+
+        Map partitionsPerTopic = new HashMap<>();
+        partitionsPerTopic.put(topic1, 3);
+        partitionsPerTopic.put(topic2, 2);
+
+        subscriptions.put(consumer1, new Subscription(topics(topic1)));
+        subscriptions.put(consumer2, new Subscription(topics(topic1, topic2)));
+        subscriptions.put(consumer3, new Subscription(topics(topic1)));
+
+        Map> assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        assertEquals(partitions(tp(topic1, 0), tp(topic1, 2)), assignment.get(consumer1));
+        assertEquals(partitions(tp(topic2, 0), tp(topic2, 1)), assignment.get(consumer2));
+        assertEquals(partitions(tp(topic1, 1)), assignment.get(consumer3));
+
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
+        assertTrue(isFullyBalanced(assignment));
+    }
+
+    @Test
+    public void testTwoConsumersTwoTopicsSixPartitions() {
+        String topic1 = "topic1";
+        String topic2 = "topic2";
+        String consumer1 = "consumer1";
+        String consumer2 = "consumer2";
+
+        Map partitionsPerTopic = new HashMap<>();
+        partitionsPerTopic.put(topic1, 3);
+        partitionsPerTopic.put(topic2, 3);
+
+        subscriptions.put(consumer1, new Subscription(topics(topic1, topic2)));
+        subscriptions.put(consumer2, new Subscription(topics(topic1, topic2)));
+
+        Map> assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        assertEquals(partitions(tp(topic1, 0), tp(topic1, 2), tp(topic2, 1)), assignment.get(consumer1));
+        assertEquals(partitions(tp(topic1, 1), tp(topic2, 0), tp(topic2, 2)), assignment.get(consumer2));
+
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
+        assertTrue(isFullyBalanced(assignment));
+    }
+
+    @Test
+    public void testAddRemoveConsumerOneTopic() {
+        String consumer1 = "consumer1";
+
+        Map partitionsPerTopic = new HashMap<>();
+        partitionsPerTopic.put(topic, 3);
+        subscriptions.put(consumer1, new Subscription(topics(topic)));
+
+        Map> assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        assertEquals(partitions(tp(topic, 0), tp(topic, 1), tp(topic, 2)), assignment.get(consumer1));
+
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
+        assertTrue(isFullyBalanced(assignment));
+
+        String consumer2 = "consumer2";
+        subscriptions.put(consumer1, buildSubscription(topics(topic), assignment.get(consumer1)));
+        subscriptions.put(consumer2, buildSubscription(topics(topic), Collections.emptyList()));
+        assignment = assignor.assign(partitionsPerTopic, subscriptions);
+
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
+        assertEquals(partitions(tp(topic, 0), tp(topic, 1)), assignment.get(consumer1));
+        assertEquals(partitions(tp(topic, 2)), assignment.get(consumer2));
+        assertTrue(isFullyBalanced(assignment));
+        assertTrue(assignor.isSticky());
+
+        subscriptions.remove(consumer1);
+        subscriptions.put(consumer2, buildSubscription(topics(topic), assignment.get(consumer2)));
+        assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        assertEquals(new HashSet<>(partitions(tp(topic, 2), tp(topic, 1), tp(topic, 0))),
+            new HashSet<>(assignment.get(consumer2)));
+
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
+        assertTrue(isFullyBalanced(assignment));
+        assertTrue(assignor.isSticky());
+    }
+
+    /**
+     * This unit test performs sticky assignment for a scenario that round robin assignor handles poorly.
+     * Topics (partitions per topic): topic1 (2), topic2 (1), topic3 (2), topic4 (1), topic5 (2)
+     * Subscriptions:
+     *  - consumer1: topic1, topic2, topic3, topic4, topic5
+     *  - consumer2: topic1, topic3, topic5
+     *  - consumer3: topic1, topic3, topic5
+     *  - consumer4: topic1, topic2, topic3, topic4, topic5
+     * Round Robin Assignment Result:
+     *  - consumer1: topic1-0, topic3-0, topic5-0
+     *  - consumer2: topic1-1, topic3-1, topic5-1
+     *  - consumer3:
+     *  - consumer4: topic2-0, topic4-0
+     * Sticky Assignment Result:
+     *  - consumer1: topic2-0, topic3-0
+     *  - consumer2: topic1-0, topic3-1
+     *  - consumer3: topic1-1, topic5-0
+     *  - consumer4: topic4-0, topic5-1
+     */
+    @Test
+    public void testPoorRoundRobinAssignmentScenario() {
+        Map partitionsPerTopic = new HashMap<>();
+        for (int i = 1; i <= 5; i++)
+            partitionsPerTopic.put(String.format("topic%d", i), (i % 2) + 1);
+
+        subscriptions.put("consumer1",
+            new Subscription(topics("topic1", "topic2", "topic3", "topic4", "topic5")));
+        subscriptions.put("consumer2",
+            new Subscription(topics("topic1", "topic3", "topic5")));
+        subscriptions.put("consumer3",
+            new Subscription(topics("topic1", "topic3", "topic5")));
+        subscriptions.put("consumer4",
+            new Subscription(topics("topic1", "topic2", "topic3", "topic4", "topic5")));
+
+        Map> assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
+    }
+
+    @Test
+    public void testAddRemoveTopicTwoConsumers() {
+        String consumer1 = "consumer";
+        String consumer2 = "consumer2";
+
+        Map partitionsPerTopic = new HashMap<>();
+        partitionsPerTopic.put(topic, 3);
+        subscriptions.put(consumer1, new Subscription(topics(topic)));
+        subscriptions.put(consumer2, new Subscription(topics(topic)));
+
+        Map> assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        // verify balance
+        assertTrue(isFullyBalanced(assignment));
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
+        // verify stickiness
+        List consumer1Assignment1 = assignment.get(consumer1);
+        List consumer2Assignment1 = assignment.get(consumer2);
+        assertTrue((consumer1Assignment1.size() == 1 && consumer2Assignment1.size() == 2) ||
+            (consumer1Assignment1.size() == 2 && consumer2Assignment1.size() == 1));
+
+        String topic2 = "topic2";
+        partitionsPerTopic.put(topic2, 3);
+        subscriptions.put(consumer1, buildSubscription(topics(topic, topic2), assignment.get(consumer1)));
+        subscriptions.put(consumer2, buildSubscription(topics(topic, topic2), assignment.get(consumer2)));
+
+        assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        // verify balance
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
+        assertTrue(isFullyBalanced(assignment));
+        // verify stickiness
+        List consumer1assignment = assignment.get(consumer1);
+        List consumer2assignment = assignment.get(consumer2);
+        assertTrue(consumer1assignment.size() == 3 && consumer2assignment.size() == 3);
+        assertTrue(consumer1assignment.containsAll(consumer1Assignment1));
+        assertTrue(consumer2assignment.containsAll(consumer2Assignment1));
+        assertTrue(assignor.isSticky());
+
+        partitionsPerTopic.remove(topic);
+        subscriptions.put(consumer1, buildSubscription(topics(topic2), assignment.get(consumer1)));
+        subscriptions.put(consumer2, buildSubscription(topics(topic2), assignment.get(consumer2)));
+
+        assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        // verify balance
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
+        assertTrue(isFullyBalanced(assignment));
+        // verify stickiness
+        List consumer1Assignment3 = assignment.get(consumer1);
+        List consumer2Assignment3 = assignment.get(consumer2);
+        assertTrue((consumer1Assignment3.size() == 1 && consumer2Assignment3.size() == 2) ||
+            (consumer1Assignment3.size() == 2 && consumer2Assignment3.size() == 1));
+        assertTrue(consumer1assignment.containsAll(consumer1Assignment3));
+        assertTrue(consumer2assignment.containsAll(consumer2Assignment3));
+        assertTrue(assignor.isSticky());
+    }
+
+
+    @Test
+    public void testReassignmentAfterOneConsumerLeaves() {
+        Map partitionsPerTopic = new HashMap<>();
+        for (int i = 1; i < 20; i++)
+            partitionsPerTopic.put(getTopicName(i, 20), i);
+
+        for (int i = 1; i < 20; i++) {
+            List topics = new ArrayList<>();
+            for (int j = 1; j <= i; j++)
+                topics.add(getTopicName(j, 20));
+            subscriptions.put(getConsumerName(i, 20), new Subscription(topics));
+        }
+
+        Map> assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
+
+        for (int i = 1; i < 20; i++) {
+            String consumer = getConsumerName(i, 20);
+            subscriptions.put(consumer,
+                buildSubscription(subscriptions.get(consumer).topics(), assignment.get(consumer)));
+        }
+        subscriptions.remove("consumer10");
+
+        assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
+        assertTrue(assignor.isSticky());
+    }
+
+
+    @Test
+    public void testReassignmentAfterOneConsumerAdded() {
+        Map partitionsPerTopic = new HashMap<>();
+        partitionsPerTopic.put("topic", 20);
+
+        for (int i = 1; i < 10; i++)
+            subscriptions.put(getConsumerName(i, 10),
+                new Subscription(topics("topic")));
+
+        Map> assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
+
+        // add a new consumer
+        subscriptions.put(getConsumerName(10, 10), new Subscription(topics("topic")));
+
+        assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
+        assertTrue(assignor.isSticky());
+    }
+
+    @Test
+    public void testSameSubscriptions() {
+        Map partitionsPerTopic = new HashMap<>();
+        for (int i = 1; i < 15; i++)
+            partitionsPerTopic.put(getTopicName(i, 15), i);
+
+        for (int i = 1; i < 9; i++) {
+            List topics = new ArrayList<>();
+            for (int j = 1; j <= partitionsPerTopic.size(); j++)
+                topics.add(getTopicName(j, 15));
+            subscriptions.put(getConsumerName(i, 9), new Subscription(topics));
+        }
+
+        Map> assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
+
+        for (int i = 1; i < 9; i++) {
+            String consumer = getConsumerName(i, 9);
+            subscriptions.put(consumer,
+                buildSubscription(subscriptions.get(consumer).topics(), assignment.get(consumer)));
+        }
+        subscriptions.remove(getConsumerName(5, 9));
+
+        assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
+        assertTrue(assignor.isSticky());
+    }
+
+    @Test(timeout = 30 * 1000)
+    public void testLargeAssignmentAndGroupWithUniformSubscription() {
+        // 1 million partitions!
+        int topicCount = 500;
+        int partitionCount = 2_000;
+        int consumerCount = 2_000;
+
+        List topics = new ArrayList<>();
+        Map partitionsPerTopic = new HashMap<>();
+        for (int i = 0; i < topicCount; i++) {
+            String topicName = getTopicName(i, topicCount);
+            topics.add(topicName);
+            partitionsPerTopic.put(topicName, partitionCount);
+        }
+
+        for (int i = 0; i < consumerCount; i++) {
+            subscriptions.put(getConsumerName(i, consumerCount), new Subscription(topics));
+        }
+
+        Map> assignment = assignor.assign(partitionsPerTopic, subscriptions);
+
+        for (int i = 1; i < consumerCount; i++) {
+            String consumer = getConsumerName(i, consumerCount);
+            subscriptions.put(consumer, buildSubscription(topics, assignment.get(consumer)));
+        }
+
+        assignor.assign(partitionsPerTopic, subscriptions);
+    }
+
+    @Test
+    public void testLargeAssignmentWithMultipleConsumersLeavingAndRandomSubscription() {
+        Random rand = new Random();
+        int topicCount = 40;
+        int consumerCount = 200;
+
+        Map partitionsPerTopic = new HashMap<>();
+        for (int i = 0; i < topicCount; i++)
+            partitionsPerTopic.put(getTopicName(i, topicCount), rand.nextInt(10) + 1);
+
+        for (int i = 0; i < consumerCount; i++) {
+            List topics = new ArrayList<>();
+            for (int j = 0; j < rand.nextInt(20); j++)
+                topics.add(getTopicName(rand.nextInt(topicCount), topicCount));
+            subscriptions.put(getConsumerName(i, consumerCount), new Subscription(topics));
+        }
+
+        Map> assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
+
+        for (int i = 1; i < consumerCount; i++) {
+            String consumer = getConsumerName(i, consumerCount);
+            subscriptions.put(consumer,
+                buildSubscription(subscriptions.get(consumer).topics(), assignment.get(consumer)));
+        }
+        for (int i = 0; i < 50; ++i) {
+            String c = getConsumerName(rand.nextInt(consumerCount), consumerCount);
+            subscriptions.remove(c);
+        }
+
+        assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
+        assertTrue(assignor.isSticky());
+    }
+
+    @Test
+    public void testNewSubscription() {
+        Map partitionsPerTopic = new HashMap<>();
+        for (int i = 1; i < 5; i++)
+            partitionsPerTopic.put(getTopicName(i, 5), 1);
+
+        for (int i = 0; i < 3; i++) {
+            List topics = new ArrayList<>();
+            for (int j = i; j <= 3 * i - 2; j++)
+                topics.add(getTopicName(j, 5));
+            subscriptions.put(getConsumerName(i, 3), new Subscription(topics));
+        }
+
+        Map> assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
+
+        subscriptions.get(getConsumerName(0, 3)).topics().add(getTopicName(1, 5));
+
+        assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
+        assertTrue(assignor.isSticky());
+    }
+
+    @Test
+    public void testMoveExistingAssignments() {
+        Map partitionsPerTopic = new HashMap<>();
+        for (int i = 1; i <= 6; i++)
+            partitionsPerTopic.put(String.format("topic%02d", i), 1);
+
+        subscriptions.put("consumer01",
+            buildSubscription(topics("topic01", "topic02"),
+                partitions(tp("topic01", 0))));
+        subscriptions.put("consumer02",
+            buildSubscription(topics("topic01", "topic02", "topic03", "topic04"),
+                partitions(tp("topic02", 0), tp("topic03", 0))));
+        subscriptions.put("consumer03",
+            buildSubscription(topics("topic02", "topic03", "topic04", "topic05", "topic06"),
+                partitions(tp("topic04", 0), tp("topic05", 0), tp("topic06", 0))));
+
+        Map> assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
+    }
+
+    @Test
+    public void testStickiness() {
+        Map partitionsPerTopic = new HashMap<>();
+        partitionsPerTopic.put("topic01", 3);
+        String consumer1 = "consumer01";
+        String consumer2 = "consumer02";
+        String consumer3 = "consumer03";
+        String consumer4 = "consumer04";
+
+        subscriptions.put(consumer1, new Subscription(topics("topic01")));
+        subscriptions.put(consumer2, new Subscription(topics("topic01")));
+        subscriptions.put(consumer3, new Subscription(topics("topic01")));
+        subscriptions.put(consumer4, new Subscription(topics("topic01")));
+
+        Map> assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
+        Map partitionsAssigned = new HashMap<>();
+
+        Set>> assignments = assignment.entrySet();
+        for (Map.Entry> entry: assignments) {
+            String consumer = entry.getKey();
+            List topicPartitions = entry.getValue();
+            int size = topicPartitions.size();
+            assertTrue("Consumer " + consumer + " is assigned more topic partitions than expected.", size <= 1);
+            if (size == 1)
+                partitionsAssigned.put(consumer, topicPartitions.get(0));
+        }
+
+        // removing the potential group leader
+        subscriptions.remove(consumer1);
+        subscriptions.put(consumer2,
+            buildSubscription(topics("topic01"), assignment.get(consumer2)));
+        subscriptions.put(consumer3,
+            buildSubscription(topics("topic01"), assignment.get(consumer3)));
+        subscriptions.put(consumer4,
+            buildSubscription(topics("topic01"), assignment.get(consumer4)));
+
+
+        assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
+        assertTrue(assignor.isSticky());
+
+        assignments = assignment.entrySet();
+        for (Map.Entry> entry: assignments) {
+            String consumer = entry.getKey();
+            List topicPartitions = entry.getValue();
+            assertEquals("Consumer " + consumer + " is assigned more topic partitions than expected.", 1, topicPartitions.size());
+            assertTrue("Stickiness was not honored for consumer " + consumer,
+                (!partitionsAssigned.containsKey(consumer)) || (assignment.get(consumer).contains(partitionsAssigned.get(consumer))));
+        }
+    }
+
+    @Test
+    public void testAssignmentUpdatedForDeletedTopic() {
+        Map partitionsPerTopic = new HashMap<>();
+        partitionsPerTopic.put("topic01", 1);
+        partitionsPerTopic.put("topic03", 100);
+        subscriptions = Collections.singletonMap(consumerId, new Subscription(topics("topic01", "topic02", "topic03")));
+
+        Map> assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        assertEquals(assignment.values().stream().mapToInt(topicPartitions -> topicPartitions.size()).sum(), 1 + 100);
+        assertEquals(Collections.singleton(consumerId), assignment.keySet());
+        assertTrue(isFullyBalanced(assignment));
+    }
+
+    @Test
+    public void testNoExceptionThrownWhenOnlySubscribedTopicDeleted() {
+        Map partitionsPerTopic = new HashMap<>();
+        partitionsPerTopic.put(topic, 3);
+        subscriptions.put(consumerId, new Subscription(topics(topic)));
+        Map> assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        subscriptions.put(consumerId, buildSubscription(topics(topic), assignment.get(consumerId)));
+
+        assignment = assignor.assign(Collections.emptyMap(), subscriptions);
+        assertEquals(assignment.size(), 1);
+        assertTrue(assignment.get(consumerId).isEmpty());
+    }
+
+    @Test
+    public void testReassignmentWithRandomSubscriptionsAndChanges() {
+        final int minNumConsumers = 20;
+        final int maxNumConsumers = 40;
+        final int minNumTopics = 10;
+        final int maxNumTopics = 20;
+
+        for (int round = 1; round <= 100; ++round) {
+            int numTopics = minNumTopics + new Random().nextInt(maxNumTopics - minNumTopics);
+
+            ArrayList topics = new ArrayList<>();
+            for (int i = 0; i < numTopics; ++i)
+                topics.add(getTopicName(i, maxNumTopics));
+
+            Map partitionsPerTopic = new HashMap<>();
+            for (int i = 0; i < numTopics; ++i)
+                partitionsPerTopic.put(getTopicName(i, maxNumTopics), i + 1);
+
+            int numConsumers = minNumConsumers + new Random().nextInt(maxNumConsumers - minNumConsumers);
+
+            for (int i = 0; i < numConsumers; ++i) {
+                List sub = Utils.sorted(getRandomSublist(topics));
+                subscriptions.put(getConsumerName(i, maxNumConsumers), new Subscription(sub));
+            }
+
+            assignor = createAssignor();
+
+            Map> assignment = assignor.assign(partitionsPerTopic, subscriptions);
+            verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
+
+            subscriptions.clear();
+            for (int i = 0; i < numConsumers; ++i) {
+                List sub = Utils.sorted(getRandomSublist(topics));
+                String consumer = getConsumerName(i, maxNumConsumers);
+                subscriptions.put(consumer, buildSubscription(sub, assignment.get(consumer)));
+            }
+
+            assignment = assignor.assign(partitionsPerTopic, subscriptions);
+            verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
+            assertTrue(assignor.isSticky());
+        }
+    }
+
+    private String getTopicName(int i, int maxNum) {
+        return getCanonicalName("t", i, maxNum);
+    }
+
+    private String getConsumerName(int i, int maxNum) {
+        return getCanonicalName("c", i, maxNum);
+    }
+
+    private String getCanonicalName(String str, int i, int maxNum) {
+        return str + pad(i, Integer.toString(maxNum).length());
+    }
+
+    private String pad(int num, int digits) {
+        StringBuilder sb = new StringBuilder();
+        int iDigits = Integer.toString(num).length();
+
+        for (int i = 1; i <= digits - iDigits; ++i)
+            sb.append("0");
+
+        sb.append(num);
+        return sb.toString();
+    }
+
+    protected static List topics(String... topics) {
+        return Arrays.asList(topics);
+    }
+
+    protected static List partitions(TopicPartition... partitions) {
+        return Arrays.asList(partitions);
+    }
+
+    protected static TopicPartition tp(String topic, int partition) {
+        return new TopicPartition(topic, partition);
+    }
+
+    protected static boolean isFullyBalanced(Map> assignment) {
+        int min = Integer.MAX_VALUE;
+        int max = Integer.MIN_VALUE;
+        for (List topicPartitions: assignment.values()) {
+            int size = topicPartitions.size();
+            if (size < min)
+                min = size;
+            if (size > max)
+                max = size;
+        }
+        return max - min <= 1;
+    }
+
+    protected static List getRandomSublist(ArrayList list) {
+        List selectedItems = new ArrayList<>(list);
+        int len = list.size();
+        Random random = new Random();
+        int howManyToRemove = random.nextInt(len);
+
+        for (int i = 1; i <= howManyToRemove; ++i)
+            selectedItems.remove(random.nextInt(selectedItems.size()));
+
+        return selectedItems;
+    }
+
+    /**
+     * Verifies that the given assignment is valid with respect to the given subscriptions
+     * Validity requirements:
+     * - each consumer is subscribed to topics of all partitions assigned to it, and
+     * - each partition is assigned to no more than one consumer
+     * Balance requirements:
+     * - the assignment is fully balanced (the numbers of topic partitions assigned to consumers differ by at most one), or
+     * - there is no topic partition that can be moved from one consumer to another with 2+ fewer topic partitions
+     *
+     * @param subscriptions: topic subscriptions of each consumer
+     * @param assignments: given assignment for balance check
+     * @param partitionsPerTopic: number of partitions per topic
+     */
+    protected void verifyValidityAndBalance(Map subscriptions,
+                                            Map> assignments,
+                                            Map partitionsPerTopic) {
+        int size = subscriptions.size();
+        assert size == assignments.size();
+
+        List consumers = Utils.sorted(assignments.keySet());
+
+        for (int i = 0; i < size; ++i) {
+            String consumer = consumers.get(i);
+            List partitions = assignments.get(consumer);
+            for (TopicPartition partition: partitions)
+                assertTrue("Error: Partition " + partition + "is assigned to c" + i + ", but it is not subscribed to Topic t" + partition.topic()
+                        + "\nSubscriptions: " + subscriptions.toString() + "\nAssignments: " + assignments.toString(),
+                    subscriptions.get(consumer).topics().contains(partition.topic()));
+
+            if (i == size - 1)
+                continue;
+
+            for (int j = i + 1; j < size; ++j) {
+                String otherConsumer = consumers.get(j);
+                List otherPartitions = assignments.get(otherConsumer);
+
+                Set intersection = new HashSet<>(partitions);
+                intersection.retainAll(otherPartitions);
+                assertTrue("Error: Consumers c" + i + " and c" + j + " have common partitions assigned to them: " + intersection.toString()
+                        + "\nSubscriptions: " + subscriptions.toString() + "\nAssignments: " + assignments.toString(),
+                    intersection.isEmpty());
+
+                int len = partitions.size();
+                int otherLen = otherPartitions.size();
+
+                if (Math.abs(len - otherLen) <= 1)
+                    continue;
+
+                Map> map = CollectionUtils.groupPartitionsByTopic(partitions);
+                Map> otherMap = CollectionUtils.groupPartitionsByTopic(otherPartitions);
+
+                if (len > otherLen) {
+                    for (String topic: map.keySet())
+                        if (otherMap.containsKey(topic))
+                            //assertTrue(true);
+                        assertFalse("Error: Some partitions can be moved from c" + i + " to c" + j
+                                + " to achieve a better balance"
+                                + "\nc" + i + " has " + len + " partitions, and c" + j + " has " + otherLen
+                                + " partitions."
+                                + "\nSubscriptions: " + subscriptions.toString() + "\nAssignments: " + assignments
+                                .toString(),
+                            otherMap.containsKey(topic));
+
+
+
+                }
+
+                if (otherLen > len) {
+                    for (String topic: otherMap.keySet())
+                        if (otherMap.containsKey(topic))
+                            //assertTrue(true);
+                        assertFalse("Error: Some partitions can be moved from c" + j + " to c" + i + " to achieve a better balance"
+                                + "\nc" + i + " has " + len + " partitions, and c" + j + " has " + otherLen + " partitions."
+                                + "\nSubscriptions: " + subscriptions.toString() + "\nAssignments: " + assignments.toString(),
+                            map.containsKey(topic));
+
+
+                }
+            }
+        }
+    }
+
+}
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 885b3574034e3..6d818df83d9b1 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
@@ -17,15 +17,15 @@
 package org.apache.kafka.clients.consumer.internals;
 
 import org.apache.kafka.clients.ClientResponse;
+import org.apache.kafka.clients.GroupRebalanceConfig;
 import org.apache.kafka.clients.MockClient;
 import org.apache.kafka.clients.consumer.CommitFailedException;
+import org.apache.kafka.clients.consumer.ConsumerPartitionAssignor;
 import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
 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.clients.consumer.RangeAssignor;
 import org.apache.kafka.clients.consumer.RetriableCommitFailedException;
-import org.apache.kafka.clients.consumer.RoundRobinAssignor;
 import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.Metric;
 import org.apache.kafka.common.MetricName;
@@ -34,19 +34,30 @@
 import org.apache.kafka.common.errors.ApiException;
 import org.apache.kafka.common.errors.AuthenticationException;
 import org.apache.kafka.common.errors.DisconnectException;
+import org.apache.kafka.common.errors.FencedInstanceIdException;
 import org.apache.kafka.common.errors.GroupAuthorizationException;
 import org.apache.kafka.common.errors.OffsetMetadataTooLarge;
+import org.apache.kafka.common.errors.TopicAuthorizationException;
 import org.apache.kafka.common.errors.WakeupException;
 import org.apache.kafka.common.internals.ClusterResourceListeners;
 import org.apache.kafka.common.internals.Topic;
+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.MemberIdentity;
 import org.apache.kafka.common.message.LeaveGroupResponseData;
+import org.apache.kafka.common.message.OffsetCommitRequestData;
+import org.apache.kafka.common.message.OffsetCommitResponseData;
+import org.apache.kafka.common.message.SyncGroupResponseData;
+import org.apache.kafka.common.metrics.KafkaMetric;
 import org.apache.kafka.common.metrics.Metrics;
+import org.apache.kafka.common.metrics.Sensor;
 import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.common.record.RecordBatch;
 import org.apache.kafka.common.requests.AbstractRequest;
 import org.apache.kafka.common.requests.FindCoordinatorResponse;
 import org.apache.kafka.common.requests.HeartbeatResponse;
 import org.apache.kafka.common.requests.JoinGroupRequest;
-import org.apache.kafka.common.requests.JoinGroupRequest.ProtocolMetadata;
 import org.apache.kafka.common.requests.JoinGroupResponse;
 import org.apache.kafka.common.requests.LeaveGroupRequest;
 import org.apache.kafka.common.requests.LeaveGroupResponse;
@@ -58,11 +69,14 @@
 import org.apache.kafka.common.requests.SyncGroupResponse;
 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.test.TestCondition;
 import org.apache.kafka.test.TestUtils;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
 
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
@@ -71,6 +85,7 @@
 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,6 +103,9 @@
 import static java.util.Collections.singleton;
 import static java.util.Collections.singletonList;
 import static java.util.Collections.singletonMap;
+import static org.apache.kafka.clients.consumer.ConsumerPartitionAssignor.RebalanceProtocol.COOPERATIVE;
+import static org.apache.kafka.clients.consumer.ConsumerPartitionAssignor.RebalanceProtocol.EAGER;
+import static org.apache.kafka.test.TestUtils.toSet;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
@@ -96,12 +114,14 @@
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
+@RunWith(value = Parameterized.class)
 public class ConsumerCoordinatorTest {
     private final String topic1 = "test1";
     private final String topic2 = "test2";
     private final TopicPartition t1p = new TopicPartition(topic1, 0);
     private final TopicPartition t2p = new TopicPartition(topic2, 0);
     private final String groupId = "test-group";
+    private final Optional groupInstanceId = Optional.of("test-instance");
     private final int rebalanceTimeoutMs = 60000;
     private final int sessionTimeoutMs = 10000;
     private final int heartbeatIntervalMs = 5000;
@@ -109,11 +129,11 @@ public class ConsumerCoordinatorTest {
     private final int autoCommitIntervalMs = 2000;
     private final int requestTimeoutMs = 30000;
     private final MockTime time = new MockTime();
-    private final Heartbeat heartbeat = new Heartbeat(time, sessionTimeoutMs, heartbeatIntervalMs,
-            rebalanceTimeoutMs, retryBackoffMs);
+    private GroupRebalanceConfig rebalanceConfig;
 
-    private MockPartitionAssignor partitionAssignor = new MockPartitionAssignor();
-    private List assignors = Collections.singletonList(partitionAssignor);
+    private final ConsumerPartitionAssignor.RebalanceProtocol protocol;
+    private final MockPartitionAssignor partitionAssignor;
+    private final List assignors;
     private MockClient client;
     private MetadataResponse metadataResponse = TestUtils.metadataUpdateWith(1, new HashMap() {
         {
@@ -130,12 +150,27 @@ public class ConsumerCoordinatorTest {
     private MockCommitCallback mockOffsetCommitCallback;
     private ConsumerCoordinator coordinator;
 
+    public ConsumerCoordinatorTest(final ConsumerPartitionAssignor.RebalanceProtocol protocol) {
+        this.protocol = protocol;
+        this.partitionAssignor = new MockPartitionAssignor(Collections.singletonList(protocol));
+        this.assignors = Collections.singletonList(partitionAssignor);
+    }
+
+    @Parameterized.Parameters(name = "rebalance protocol = {0}")
+    public static Collection data() {
+        final List values = new ArrayList<>();
+        for (final ConsumerPartitionAssignor.RebalanceProtocol protocol: ConsumerPartitionAssignor.RebalanceProtocol.values()) {
+            values.add(new Object[]{protocol});
+        }
+        return values;
+    }
+
     @Before
     public void setup() {
         LogContext logContext = new LogContext();
         this.subscriptions = new SubscriptionState(logContext, OffsetResetStrategy.EARLIEST);
         this.metadata = new ConsumerMetadata(0, Long.MAX_VALUE, false,
-                subscriptions, logContext, new ClusterResourceListeners());
+                false, subscriptions, logContext, new ClusterResourceListeners());
         this.client = new MockClient(time, metadata);
         this.client.updateMetadata(metadataResponse);
         this.consumerClient = new ConsumerNetworkClient(logContext, client, metadata, time, 100,
@@ -144,8 +179,21 @@ public void setup() {
         this.rebalanceListener = new MockRebalanceListener();
         this.mockOffsetCommitCallback = new MockCommitCallback();
         this.partitionAssignor.clear();
+        this.rebalanceConfig = buildRebalanceConfig(Optional.empty());
+        this.coordinator = buildCoordinator(rebalanceConfig,
+                                            metrics,
+                                            assignors,
+                                            false);
+    }
 
-        this.coordinator = buildCoordinator(metrics, assignors, false, true);
+    private GroupRebalanceConfig buildRebalanceConfig(Optional groupInstanceId) {
+        return new GroupRebalanceConfig(sessionTimeoutMs,
+                                        rebalanceTimeoutMs,
+                                        heartbeatIntervalMs,
+                                        groupId,
+                                        groupInstanceId,
+                                        retryBackoffMs,
+                                        !groupInstanceId.isPresent());
     }
 
     @After
@@ -154,6 +202,77 @@ public void teardown() {
         this.coordinator.close(time.timer(0));
     }
 
+    @Test
+    public void testMetrics() {
+        assertNotNull(getMetric("commit-latency-avg"));
+        assertNotNull(getMetric("commit-latency-max"));
+        assertNotNull(getMetric("commit-rate"));
+        assertNotNull(getMetric("commit-total"));
+        assertNotNull(getMetric("partition-revoked-latency-avg"));
+        assertNotNull(getMetric("partition-revoked-latency-max"));
+        assertNotNull(getMetric("partition-assigned-latency-avg"));
+        assertNotNull(getMetric("partition-assigned-latency-max"));
+        assertNotNull(getMetric("partition-lost-latency-avg"));
+        assertNotNull(getMetric("partition-lost-latency-max"));
+        assertNotNull(getMetric("assigned-partitions"));
+
+        metrics.sensor("commit-latency").record(1.0d);
+        metrics.sensor("commit-latency").record(6.0d);
+        metrics.sensor("commit-latency").record(2.0d);
+
+        assertEquals(3.0d, getMetric("commit-latency-avg").metricValue());
+        assertEquals(6.0d, getMetric("commit-latency-max").metricValue());
+        assertEquals(0.1d, getMetric("commit-rate").metricValue());
+        assertEquals(3.0d, getMetric("commit-total").metricValue());
+
+        metrics.sensor("partition-revoked-latency").record(1.0d);
+        metrics.sensor("partition-revoked-latency").record(2.0d);
+        metrics.sensor("partition-assigned-latency").record(1.0d);
+        metrics.sensor("partition-assigned-latency").record(2.0d);
+        metrics.sensor("partition-lost-latency").record(1.0d);
+        metrics.sensor("partition-lost-latency").record(2.0d);
+
+        assertEquals(1.5d, getMetric("partition-revoked-latency-avg").metricValue());
+        assertEquals(2.0d, getMetric("partition-revoked-latency-max").metricValue());
+        assertEquals(1.5d, getMetric("partition-assigned-latency-avg").metricValue());
+        assertEquals(2.0d, getMetric("partition-assigned-latency-max").metricValue());
+        assertEquals(1.5d, getMetric("partition-lost-latency-avg").metricValue());
+        assertEquals(2.0d, getMetric("partition-lost-latency-max").metricValue());
+
+        assertEquals(0.0d, getMetric("assigned-partitions").metricValue());
+        subscriptions.assignFromUser(Collections.singleton(t1p));
+        assertEquals(1.0d, getMetric("assigned-partitions").metricValue());
+        subscriptions.assignFromUser(Utils.mkSet(t1p, t2p));
+        assertEquals(2.0d, getMetric("assigned-partitions").metricValue());
+    }
+
+    private KafkaMetric getMetric(final String name) {
+        return metrics.metrics().get(metrics.metricName(name, "consumer" + groupId + "-coordinator-metrics"));
+    }
+
+    private Sensor getSensor(final String name) {
+        return metrics.sensor(name);
+    }
+
+    @Test
+    public void testSelectRebalanceProtcol() {
+        List assignors = new ArrayList<>();
+        assignors.add(new MockPartitionAssignor(Collections.singletonList(ConsumerPartitionAssignor.RebalanceProtocol.EAGER)));
+        assignors.add(new MockPartitionAssignor(Collections.singletonList(COOPERATIVE)));
+
+        // no commonly supported protocols
+        assertThrows(IllegalArgumentException.class, () -> buildCoordinator(rebalanceConfig, new Metrics(), assignors, false));
+
+        assignors.clear();
+        assignors.add(new MockPartitionAssignor(Arrays.asList(ConsumerPartitionAssignor.RebalanceProtocol.EAGER, COOPERATIVE)));
+        assignors.add(new MockPartitionAssignor(Arrays.asList(ConsumerPartitionAssignor.RebalanceProtocol.EAGER, COOPERATIVE)));
+
+        // select higher indexed (more advanced) protocols
+        try (ConsumerCoordinator coordinator = buildCoordinator(rebalanceConfig, new Metrics(), assignors, false)) {
+            assertEquals(COOPERATIVE, coordinator.getProtocol());
+        }
+    }
+
     @Test
     public void testNormalHeartbeat() {
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
@@ -185,7 +304,7 @@ public void testGroupReadUnauthorized() {
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
         coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
-        client.prepareResponse(joinGroupLeaderResponse(0, "memberId", Collections.>emptyMap(),
+        client.prepareResponse(joinGroupLeaderResponse(0, "memberId", Collections.emptyMap(),
                 Errors.GROUP_AUTHORIZATION_FAILED));
         coordinator.poll(time.timer(Long.MAX_VALUE));
     }
@@ -212,7 +331,7 @@ public void testCoordinatorNotAvailable() {
     }
 
     @Test
-    public void testManyInFlightAsyncCommitsWithCoordinatorDisconnect() throws Exception {
+    public void testManyInFlightAsyncCommitsWithCoordinatorDisconnect() {
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
         coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
@@ -240,7 +359,7 @@ public void onComplete(Map offsets, Exception
     }
 
     @Test
-    public void testCoordinatorUnknownInUnsentCallbacksAfterCoordinatorDead() throws Exception {
+    public void testCoordinatorUnknownInUnsentCallbacksAfterCoordinatorDead() {
         // When the coordinator is marked dead, all unsent or in-flight requests are cancelled
         // with a disconnect error. This test case ensures that the corresponding callbacks see
         // the coordinator as unknown which prevents additional retries to the same coordinator.
@@ -249,10 +368,24 @@ public void testCoordinatorUnknownInUnsentCallbacksAfterCoordinatorDead() throws
         coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
         final AtomicBoolean asyncCallbackInvoked = new AtomicBoolean(false);
-        Map offsets = singletonMap(
-                new TopicPartition("foo", 0), new OffsetCommitRequest.PartitionData(13L,
-                        Optional.empty(), ""));
-        consumerClient.send(coordinator.checkAndGetCoordinator(), new OffsetCommitRequest.Builder(groupId, offsets))
+
+        OffsetCommitRequestData offsetCommitRequestData = new OffsetCommitRequestData()
+                .setGroupId(groupId)
+                .setTopics(Collections.singletonList(new
+                        OffsetCommitRequestData.OffsetCommitRequestTopic()
+                                .setName("foo")
+                                .setPartitions(Collections.singletonList(
+                                        new OffsetCommitRequestData.OffsetCommitRequestPartition()
+                                                .setPartitionIndex(0)
+                                                .setCommittedLeaderEpoch(RecordBatch.NO_PARTITION_LEADER_EPOCH)
+                                                .setCommittedMetadata("")
+                                                .setCommittedOffset(13L)
+                                                .setCommitTimestamp(0)
+                                ))
+                        )
+                );
+
+        consumerClient.send(coordinator.checkAndGetCoordinator(), new OffsetCommitRequest.Builder(offsetCommitRequestData))
                 .compose(new RequestFutureAdapter() {
                     @Override
                     public void onSuccess(ClientResponse value, RequestFuture future) {}
@@ -313,10 +446,43 @@ public void testIllegalGeneration() {
         assertTrue(future.failed());
         assertEquals(Errors.ILLEGAL_GENERATION.exception(), future.exception());
         assertTrue(coordinator.rejoinNeededOrPending());
+
+        coordinator.poll(time.timer(0));
+
+        assertEquals(1, rebalanceListener.lostCount);
+        assertEquals(Collections.singleton(t1p), rebalanceListener.lost);
     }
 
     @Test
-    public void testUnknownConsumerId() {
+    public void testUnsubscribeWithValidGeneration() {
+        client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
+
+        subscriptions.subscribe(singleton(topic1), rebalanceListener);
+        ByteBuffer buffer = ConsumerProtocol.serializeAssignment(
+            new ConsumerPartitionAssignor.Assignment(Collections.singletonList(t1p), ByteBuffer.wrap(new byte[0])));
+        coordinator.onJoinComplete(1, "memberId", partitionAssignor.name(), buffer);
+
+        coordinator.onLeavePrepare();
+        assertEquals(1, rebalanceListener.lostCount);
+        assertEquals(0, rebalanceListener.revokedCount);
+    }
+
+    @Test
+    public void testUnsubscribeWithInvalidGeneration() {
+        client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
+
+        subscriptions.subscribe(singleton(topic1), rebalanceListener);
+        subscriptions.assignFromSubscribed(Collections.singletonList(t1p));
+
+        coordinator.onLeavePrepare();
+        assertEquals(1, rebalanceListener.lostCount);
+        assertEquals(0, rebalanceListener.revokedCount);
+    }
+
+    @Test
+    public void testUnknownMemberId() {
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
         coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
@@ -337,6 +503,11 @@ public void testUnknownConsumerId() {
         assertTrue(future.failed());
         assertEquals(Errors.UNKNOWN_MEMBER_ID.exception(), future.exception());
         assertTrue(coordinator.rejoinNeededOrPending());
+
+        coordinator.poll(time.timer(0));
+
+        assertEquals(1, rebalanceListener.lostCount);
+        assertEquals(Collections.singleton(t1p), rebalanceListener.lost);
     }
 
     @Test
@@ -372,7 +543,7 @@ public void testJoinGroupInvalidGroupId() {
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
         coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
-        client.prepareResponse(joinGroupLeaderResponse(0, consumerId, Collections.>emptyMap(),
+        client.prepareResponse(joinGroupLeaderResponse(0, consumerId, Collections.emptyMap(),
                 Errors.INVALID_GROUP_ID));
         coordinator.poll(time.timer(Long.MAX_VALUE));
     }
@@ -380,6 +551,9 @@ public void testJoinGroupInvalidGroupId() {
     @Test
     public void testNormalJoinGroupLeader() {
         final String consumerId = "leader";
+        final Set subscription = singleton(topic1);
+        final List owned = Collections.emptyList();
+        final List assigned = Arrays.asList(t1p);
 
         subscriptions.subscribe(singleton(topic1), rebalanceListener);
 
@@ -391,116 +565,118 @@ public void testNormalJoinGroupLeader() {
 
         // normal join group
         Map> memberSubscriptions = singletonMap(consumerId, singletonList(topic1));
-        partitionAssignor.prepare(singletonMap(consumerId, singletonList(t1p)));
+        partitionAssignor.prepare(singletonMap(consumerId, assigned));
 
         client.prepareResponse(joinGroupLeaderResponse(1, consumerId, memberSubscriptions, Errors.NONE));
-        client.prepareResponse(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                SyncGroupRequest sync = (SyncGroupRequest) body;
-                return sync.memberId().equals(consumerId) &&
-                        sync.generationId() == 1 &&
-                        sync.groupAssignment().containsKey(consumerId);
-            }
-        }, syncGroupResponse(singletonList(t1p), Errors.NONE));
+        client.prepareResponse(body -> {
+            SyncGroupRequest sync = (SyncGroupRequest) body;
+            return sync.data.memberId().equals(consumerId) &&
+                    sync.data.generationId() == 1 &&
+                    sync.groupAssignments().containsKey(consumerId);
+        }, syncGroupResponse(assigned, Errors.NONE));
         coordinator.poll(time.timer(Long.MAX_VALUE));
 
         assertFalse(coordinator.rejoinNeededOrPending());
-        assertEquals(singleton(t1p), subscriptions.assignedPartitions());
-        assertEquals(singleton(topic1), subscriptions.groupSubscription());
-        assertEquals(1, rebalanceListener.revokedCount);
-        assertEquals(Collections.emptySet(), rebalanceListener.revoked);
+        assertEquals(toSet(assigned), subscriptions.assignedPartitions());
+        assertEquals(subscription, subscriptions.metadataTopics());
+        assertEquals(0, rebalanceListener.revokedCount);
+        assertNull(rebalanceListener.revoked);
         assertEquals(1, rebalanceListener.assignedCount);
-        assertEquals(singleton(t1p), rebalanceListener.assigned);
+        assertEquals(getAdded(owned, assigned), rebalanceListener.assigned);
     }
 
     @Test
     public void testOutdatedCoordinatorAssignment() {
         final String consumerId = "outdated_assignment";
+        final List owned = Collections.emptyList();
+        final List oldSubscription = singletonList(topic2);
+        final List oldAssignment = Arrays.asList(t2p);
+        final List newSubscription = singletonList(topic1);
+        final List newAssignment = Arrays.asList(t1p);
 
-        subscriptions.subscribe(singleton(topic2), rebalanceListener);
+        subscriptions.subscribe(toSet(oldSubscription), rebalanceListener);
 
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
         coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
         // Test coordinator returning unsubscribed partitions
-        partitionAssignor.prepare(singletonMap(consumerId, singletonList(t1p)));
+        partitionAssignor.prepare(singletonMap(consumerId, newAssignment));
 
         // First incorrect assignment for subscription
         client.prepareResponse(
                 joinGroupLeaderResponse(
-                    1, consumerId, singletonMap(consumerId, singletonList(topic2)), Errors.NONE));
-        client.prepareResponse(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                SyncGroupRequest sync = (SyncGroupRequest) body;
-                return sync.memberId().equals(consumerId) &&
-                        sync.generationId() == 1 &&
-                        sync.groupAssignment().containsKey(consumerId);
-            }
-        }, syncGroupResponse(Arrays.asList(t2p), Errors.NONE));
+                    1, consumerId, singletonMap(consumerId, oldSubscription), Errors.NONE));
+        client.prepareResponse(body -> {
+            SyncGroupRequest sync = (SyncGroupRequest) body;
+            return sync.data.memberId().equals(consumerId) &&
+                    sync.data.generationId() == 1 &&
+                    sync.groupAssignments().containsKey(consumerId);
+        }, syncGroupResponse(oldAssignment, Errors.NONE));
 
         // Second correct assignment for subscription
         client.prepareResponse(
                 joinGroupLeaderResponse(
-                    1, consumerId, singletonMap(consumerId, singletonList(topic1)), Errors.NONE));
-        client.prepareResponse(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                SyncGroupRequest sync = (SyncGroupRequest) body;
-                return sync.memberId().equals(consumerId) &&
-                        sync.generationId() == 1 &&
-                        sync.groupAssignment().containsKey(consumerId);
-            }
-        }, syncGroupResponse(singletonList(t1p), Errors.NONE));
+                    1, consumerId, singletonMap(consumerId, newSubscription), Errors.NONE));
+        client.prepareResponse(body -> {
+            SyncGroupRequest sync = (SyncGroupRequest) body;
+            return sync.data.memberId().equals(consumerId) &&
+                    sync.data.generationId() == 1 &&
+                    sync.groupAssignments().containsKey(consumerId);
+        }, syncGroupResponse(newAssignment, Errors.NONE));
 
         // Poll once so that the join group future gets created and complete
         coordinator.poll(time.timer(0));
 
         // Before the sync group response gets completed change the subscription
-        subscriptions.subscribe(singleton(topic1), rebalanceListener);
+        subscriptions.subscribe(toSet(newSubscription), rebalanceListener);
         coordinator.poll(time.timer(0));
 
         coordinator.poll(time.timer(Long.MAX_VALUE));
 
+        final Collection assigned = getAdded(owned, newAssignment);
+
         assertFalse(coordinator.rejoinNeededOrPending());
-        assertEquals(singleton(t1p), subscriptions.assignedPartitions());
-        assertEquals(singleton(topic1), subscriptions.groupSubscription());
-        assertEquals(2, rebalanceListener.revokedCount);
-        assertEquals(Collections.emptySet(), rebalanceListener.revoked);
+        assertEquals(toSet(newAssignment), subscriptions.assignedPartitions());
+        assertEquals(toSet(newSubscription), subscriptions.metadataTopics());
+        assertEquals(protocol == EAGER ? 1 : 0, rebalanceListener.revokedCount);
         assertEquals(1, rebalanceListener.assignedCount);
-        assertEquals(singleton(t1p), rebalanceListener.assigned);
+        assertEquals(assigned, rebalanceListener.assigned);
     }
 
     @Test
-    public void testInvalidCoordinatorAssignment() {
-        final String consumerId = "invalid_assignment";
+    public void testMetadataTopicsDuringSubscriptionChange() {
+        final String consumerId = "subscription_change";
+        final List oldSubscription = singletonList(topic1);
+        final List oldAssignment = Collections.singletonList(t1p);
+        final List newSubscription = singletonList(topic2);
+        final List newAssignment = Collections.singletonList(t2p);
 
-        subscriptions.subscribe(singleton(topic1), rebalanceListener);
+        subscriptions.subscribe(toSet(oldSubscription), rebalanceListener);
+        assertEquals(toSet(oldSubscription), subscriptions.metadataTopics());
 
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
         coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
-        // normal join group
-        Map> memberSubscriptions = singletonMap(consumerId, singletonList(topic2));
-        partitionAssignor.prepare(singletonMap(consumerId, singletonList(t2p)));
+        prepareJoinAndSyncResponse(consumerId, 1, oldSubscription, oldAssignment);
 
-        client.prepareResponse(joinGroupLeaderResponse(1, consumerId, memberSubscriptions, Errors.NONE));
-        client.prepareResponse(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                SyncGroupRequest sync = (SyncGroupRequest) body;
-                return sync.memberId().equals(consumerId) &&
-                        sync.generationId() == 1 &&
-                        sync.groupAssignment().containsKey(consumerId);
-            }
-        }, syncGroupResponse(singletonList(t2p), Errors.NONE));
-        assertThrows(IllegalStateException.class, () -> coordinator.poll(time.timer(Long.MAX_VALUE)));
+        coordinator.poll(time.timer(0));
+        assertEquals(toSet(oldSubscription), subscriptions.metadataTopics());
+
+        subscriptions.subscribe(toSet(newSubscription), rebalanceListener);
+        assertEquals(Utils.mkSet(topic1, topic2), subscriptions.metadataTopics());
+
+        prepareJoinAndSyncResponse(consumerId, 2, newSubscription, newAssignment);
+        coordinator.poll(time.timer(Long.MAX_VALUE));
+        assertFalse(coordinator.rejoinNeededOrPending());
+        assertEquals(toSet(newAssignment), subscriptions.assignedPartitions());
+        assertEquals(toSet(newSubscription), subscriptions.metadataTopics());
     }
 
     @Test
     public void testPatternJoinGroupLeader() {
         final String consumerId = "leader";
+        final List assigned = Arrays.asList(t1p, t2p);
+        final List owned = Collections.emptyList();
 
         subscriptions.subscribe(Pattern.compile("test.*"), rebalanceListener);
 
@@ -513,18 +689,15 @@ public void testPatternJoinGroupLeader() {
 
         // normal join group
         Map> memberSubscriptions = singletonMap(consumerId, singletonList(topic1));
-        partitionAssignor.prepare(singletonMap(consumerId, Arrays.asList(t1p, t2p)));
+        partitionAssignor.prepare(singletonMap(consumerId, assigned));
 
         client.prepareResponse(joinGroupLeaderResponse(1, consumerId, memberSubscriptions, Errors.NONE));
-        client.prepareResponse(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                SyncGroupRequest sync = (SyncGroupRequest) body;
-                return sync.memberId().equals(consumerId) &&
-                        sync.generationId() == 1 &&
-                        sync.groupAssignment().containsKey(consumerId);
-            }
-        }, syncGroupResponse(Arrays.asList(t1p, t2p), Errors.NONE));
+        client.prepareResponse(body -> {
+            SyncGroupRequest sync = (SyncGroupRequest) body;
+            return sync.data.memberId().equals(consumerId) &&
+                    sync.data.generationId() == 1 &&
+                    sync.groupAssignments().containsKey(consumerId);
+        }, syncGroupResponse(assigned, Errors.NONE));
         // expect client to force updating the metadata, if yes gives it both topics
         client.prepareMetadataUpdate(metadataResponse);
 
@@ -532,18 +705,20 @@ public boolean matches(AbstractRequest body) {
 
         assertFalse(coordinator.rejoinNeededOrPending());
         assertEquals(2, subscriptions.numAssignedPartitions());
-        assertEquals(2, subscriptions.groupSubscription().size());
+        assertEquals(2, subscriptions.metadataTopics().size());
         assertEquals(2, subscriptions.subscription().size());
-        assertEquals(1, rebalanceListener.revokedCount);
-        assertEquals(Collections.emptySet(), rebalanceListener.revoked);
+        // callback not triggered at all since there's nothing to be revoked
+        assertEquals(0, rebalanceListener.revokedCount);
+        assertNull(rebalanceListener.revoked);
         assertEquals(1, rebalanceListener.assignedCount);
-        assertEquals(2, rebalanceListener.assigned.size());
+        assertEquals(getAdded(owned, assigned), rebalanceListener.assigned);
     }
 
     @Test
     public void testMetadataRefreshDuringRebalance() {
         final String consumerId = "leader";
-
+        final List owned = Collections.emptyList();
+        final List oldAssigned = Arrays.asList(t1p);
         subscriptions.subscribe(Pattern.compile(".*"), rebalanceListener);
         client.updateMetadata(TestUtils.metadataUpdateWith(1, singletonMap(topic1, 1)));
         coordinator.maybeUpdateSubscriptionMetadata();
@@ -554,53 +729,106 @@ public void testMetadataRefreshDuringRebalance() {
         coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
         Map> initialSubscription = singletonMap(consumerId, singletonList(topic1));
-        partitionAssignor.prepare(singletonMap(consumerId, singletonList(t1p)));
+        partitionAssignor.prepare(singletonMap(consumerId, oldAssigned));
 
         // the metadata will be updated in flight with a new topic added
         final List updatedSubscription = Arrays.asList(topic1, topic2);
-        final Set updatedSubscriptionSet = new HashSet<>(updatedSubscription);
 
         client.prepareResponse(joinGroupLeaderResponse(1, consumerId, initialSubscription, Errors.NONE));
+        client.prepareResponse(body -> {
+            final Map updatedPartitions = new HashMap<>();
+            for (String topic : updatedSubscription)
+                updatedPartitions.put(topic, 1);
+            client.updateMetadata(TestUtils.metadataUpdateWith(1, updatedPartitions));
+            return true;
+        }, syncGroupResponse(oldAssigned, Errors.NONE));
+        coordinator.poll(time.timer(Long.MAX_VALUE));
+
+        // rejoin will only be set in the next poll call
+        assertFalse(coordinator.rejoinNeededOrPending());
+        assertEquals(singleton(topic1), subscriptions.subscription());
+        assertEquals(toSet(oldAssigned), subscriptions.assignedPartitions());
+        // nothing to be revoked and hence no callback triggered
+        assertEquals(0, rebalanceListener.revokedCount);
+        assertNull(rebalanceListener.revoked);
+        assertEquals(1, rebalanceListener.assignedCount);
+        assertEquals(getAdded(owned, oldAssigned), rebalanceListener.assigned);
+
+        List newAssigned = Arrays.asList(t1p, t2p);
+
+        final Map> updatedSubscriptions = singletonMap(consumerId, Arrays.asList(topic1, topic2));
+        partitionAssignor.prepare(singletonMap(consumerId, newAssigned));
+
+        // we expect to see a second rebalance with the new-found topics
+        client.prepareResponse(body -> {
+            JoinGroupRequest join = (JoinGroupRequest) body;
+            Iterator protocolIterator =
+                    join.data().protocols().iterator();
+            assertTrue(protocolIterator.hasNext());
+            JoinGroupRequestData.JoinGroupRequestProtocol protocolMetadata = protocolIterator.next();
+
+            ByteBuffer metadata = ByteBuffer.wrap(protocolMetadata.metadata());
+            ConsumerPartitionAssignor.Subscription subscription = ConsumerProtocol.deserializeSubscription(metadata);
+            metadata.rewind();
+            return subscription.topics().containsAll(updatedSubscription);
+        }, joinGroupLeaderResponse(2, consumerId, updatedSubscriptions, Errors.NONE));
         client.prepareResponse(new MockClient.RequestMatcher() {
+            // update the metadata again back to topic1
             @Override
             public boolean matches(AbstractRequest body) {
-                final Map updatedPartitions = new HashMap<>();
-                for (String topic : updatedSubscription)
-                    updatedPartitions.put(topic, 1);
-                client.updateMetadata(TestUtils.metadataUpdateWith(1, updatedPartitions));
+                client.updateMetadata(TestUtils.metadataUpdateWith(1, singletonMap(topic1, 1)));
                 return true;
             }
-        }, syncGroupResponse(singletonList(t1p), Errors.NONE));
+        }, syncGroupResponse(newAssigned, Errors.NONE));
+
         coordinator.poll(time.timer(Long.MAX_VALUE));
 
-        List newAssignment = Arrays.asList(t1p, t2p);
-        Set newAssignmentSet = new HashSet<>(newAssignment);
+        Collection revoked = getRevoked(oldAssigned, newAssigned);
+        int revokedCount = revoked.isEmpty() ? 0 : 1;
 
-        Map> updatedSubscriptions = singletonMap(consumerId, Arrays.asList(topic1, topic2));
-        partitionAssignor.prepare(singletonMap(consumerId, newAssignment));
+        assertFalse(coordinator.rejoinNeededOrPending());
+        assertEquals(toSet(updatedSubscription), subscriptions.subscription());
+        assertEquals(toSet(newAssigned), subscriptions.assignedPartitions());
+        assertEquals(revokedCount, rebalanceListener.revokedCount);
+        assertEquals(revoked.isEmpty() ? null : revoked, rebalanceListener.revoked);
+        assertEquals(2, rebalanceListener.assignedCount);
+        assertEquals(getAdded(oldAssigned, newAssigned), rebalanceListener.assigned);
+
+        // we expect to see a third rebalance with the new-found topics
+        partitionAssignor.prepare(singletonMap(consumerId, oldAssigned));
 
-        // we expect to see a second rebalance with the new-found topics
         client.prepareResponse(new MockClient.RequestMatcher() {
             @Override
             public boolean matches(AbstractRequest body) {
                 JoinGroupRequest join = (JoinGroupRequest) body;
-                ProtocolMetadata protocolMetadata = join.groupProtocols().iterator().next();
-                PartitionAssignor.Subscription subscription = ConsumerProtocol.deserializeSubscription(protocolMetadata.metadata());
-                protocolMetadata.metadata().rewind();
-                return subscription.topics().containsAll(updatedSubscriptionSet);
+                Iterator protocolIterator =
+                    join.data().protocols().iterator();
+                assertTrue(protocolIterator.hasNext());
+                JoinGroupRequestData.JoinGroupRequestProtocol protocolMetadata = protocolIterator.next();
+
+                ByteBuffer metadata = ByteBuffer.wrap(protocolMetadata.metadata());
+                ConsumerPartitionAssignor.Subscription subscription = ConsumerProtocol.deserializeSubscription(metadata);
+                metadata.rewind();
+                return subscription.topics().contains(topic1);
             }
-        }, joinGroupLeaderResponse(2, consumerId, updatedSubscriptions, Errors.NONE));
-        client.prepareResponse(syncGroupResponse(newAssignment, Errors.NONE));
+        }, joinGroupLeaderResponse(3, consumerId, initialSubscription, Errors.NONE));
+        client.prepareResponse(syncGroupResponse(oldAssigned, Errors.NONE));
 
         coordinator.poll(time.timer(Long.MAX_VALUE));
 
+        revoked = getRevoked(newAssigned, oldAssigned);
+        assertFalse(revoked.isEmpty());
+        revokedCount += 1;
+        Collection added = getAdded(newAssigned, oldAssigned);
+
         assertFalse(coordinator.rejoinNeededOrPending());
-        assertEquals(updatedSubscriptionSet, subscriptions.subscription());
-        assertEquals(newAssignmentSet, subscriptions.assignedPartitions());
-        assertEquals(2, rebalanceListener.revokedCount);
-        assertEquals(singleton(t1p), rebalanceListener.revoked);
-        assertEquals(2, rebalanceListener.assignedCount);
-        assertEquals(newAssignmentSet, rebalanceListener.assigned);
+        assertEquals(singleton(topic1), subscriptions.subscription());
+        assertEquals(toSet(oldAssigned), subscriptions.assignedPartitions());
+        assertEquals(revokedCount, rebalanceListener.revokedCount);
+        assertEquals(revoked.isEmpty() ? null : revoked, rebalanceListener.revoked);
+        assertEquals(3, rebalanceListener.assignedCount);
+        assertEquals(added, rebalanceListener.assigned);
+        assertEquals(0, rebalanceListener.lostCount);
     }
 
     @Test
@@ -621,14 +849,11 @@ public void testForceMetadataRefreshForPatternSubscriptionDuringRebalance() {
         client.prepareMetadataUpdate(metadataResponse);
 
         client.prepareResponse(joinGroupFollowerResponse(1, consumerId, "leader", Errors.NONE));
-        client.prepareResponse(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                SyncGroupRequest sync = (SyncGroupRequest) body;
-                return sync.memberId().equals(consumerId) &&
-                    sync.generationId() == 1 &&
-                    sync.groupAssignment().isEmpty();
-            }
+        client.prepareResponse(body -> {
+            SyncGroupRequest sync = (SyncGroupRequest) body;
+            return sync.data.memberId().equals(consumerId) &&
+                sync.data.generationId() == 1 &&
+                sync.groupAssignments().isEmpty();
         }, syncGroupResponse(singletonList(t1p), Errors.NONE));
 
         partitionAssignor.prepare(singletonMap(consumerId, singletonList(t1p)));
@@ -647,9 +872,69 @@ public boolean matches(AbstractRequest body) {
         assertFalse(coordinator.rejoinNeededOrPending());
     }
 
+    /**
+     * Verifies that the consumer re-joins after a metadata change. If JoinGroup fails
+     * and metadata reverts to its original value, the consumer should still retry JoinGroup.
+     */
+    @Test
+    public void testRebalanceWithMetadataChange() {
+        final String consumerId = "leader";
+        final List topics = Arrays.asList(topic1, topic2);
+        final List partitions = Arrays.asList(t1p, t2p);
+        subscriptions.subscribe(toSet(topics), rebalanceListener);
+        client.updateMetadata(TestUtils.metadataUpdateWith(1,
+                Utils.mkMap(Utils.mkEntry(topic1, 1), Utils.mkEntry(topic2, 1))));
+        coordinator.maybeUpdateSubscriptionMetadata();
+
+        client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
+
+        Map> initialSubscription = singletonMap(consumerId, topics);
+        partitionAssignor.prepare(singletonMap(consumerId, partitions));
+
+        client.prepareResponse(joinGroupLeaderResponse(1, consumerId, initialSubscription, Errors.NONE));
+        client.prepareResponse(syncGroupResponse(partitions, Errors.NONE));
+        coordinator.poll(time.timer(Long.MAX_VALUE));
+
+        // rejoin will only be set in the next poll call
+        assertFalse(coordinator.rejoinNeededOrPending());
+        assertEquals(toSet(topics), subscriptions.subscription());
+        assertEquals(toSet(partitions), subscriptions.assignedPartitions());
+        assertEquals(0, rebalanceListener.revokedCount);
+        assertNull(rebalanceListener.revoked);
+        assertEquals(1, rebalanceListener.assignedCount);
+
+        // Change metadata to trigger rebalance.
+        client.updateMetadata(TestUtils.metadataUpdateWith(1, singletonMap(topic1, 1)));
+        coordinator.poll(time.timer(0));
+
+        // Revert metadata to original value. Fail pending JoinGroup. Another
+        // JoinGroup should be sent, which will be completed successfully.
+        client.updateMetadata(TestUtils.metadataUpdateWith(1,
+                Utils.mkMap(Utils.mkEntry(topic1, 1), Utils.mkEntry(topic2, 1))));
+        client.respond(joinGroupFollowerResponse(1, consumerId, "leader", Errors.NOT_COORDINATOR));
+        client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+        coordinator.poll(time.timer(0));
+        assertTrue(coordinator.rejoinNeededOrPending());
+
+        client.respond(joinGroupLeaderResponse(2, consumerId, initialSubscription, Errors.NONE));
+        client.prepareResponse(syncGroupResponse(partitions, Errors.NONE));
+        coordinator.poll(time.timer(Long.MAX_VALUE));
+
+        assertFalse(coordinator.rejoinNeededOrPending());
+        Collection revoked = getRevoked(partitions, partitions);
+        assertEquals(revoked.isEmpty() ? 0 : 1, rebalanceListener.revokedCount);
+        assertEquals(revoked.isEmpty() ? null : revoked, rebalanceListener.revoked);
+        assertEquals(2, rebalanceListener.assignedCount);
+        assertEquals(getAdded(partitions, partitions), rebalanceListener.assigned);
+        assertEquals(toSet(partitions), subscriptions.assignedPartitions());
+    }
+
     @Test
     public void testWakeupDuringJoin() {
         final String consumerId = "leader";
+        final List owned = Collections.emptyList();
+        final List assigned = Arrays.asList(t1p);
 
         subscriptions.subscribe(singleton(topic1), rebalanceListener);
 
@@ -660,7 +945,7 @@ public void testWakeupDuringJoin() {
         coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
         Map> memberSubscriptions = singletonMap(consumerId, singletonList(topic1));
-        partitionAssignor.prepare(singletonMap(consumerId, singletonList(t1p)));
+        partitionAssignor.prepare(singletonMap(consumerId, assigned));
 
         // prepare only the first half of the join and then trigger the wakeup
         client.prepareResponse(joinGroupLeaderResponse(1, consumerId, memberSubscriptions, Errors.NONE));
@@ -673,47 +958,47 @@ public void testWakeupDuringJoin() {
         }
 
         // now complete the second half
-        client.prepareResponse(syncGroupResponse(singletonList(t1p), Errors.NONE));
+        client.prepareResponse(syncGroupResponse(assigned, Errors.NONE));
         coordinator.poll(time.timer(Long.MAX_VALUE));
 
         assertFalse(coordinator.rejoinNeededOrPending());
-        assertEquals(singleton(t1p), subscriptions.assignedPartitions());
-        assertEquals(1, rebalanceListener.revokedCount);
-        assertEquals(Collections.emptySet(), rebalanceListener.revoked);
+        assertEquals(toSet(assigned), subscriptions.assignedPartitions());
+        assertEquals(0, rebalanceListener.revokedCount);
+        assertNull(rebalanceListener.revoked);
         assertEquals(1, rebalanceListener.assignedCount);
-        assertEquals(singleton(t1p), rebalanceListener.assigned);
+        assertEquals(getAdded(owned, assigned), rebalanceListener.assigned);
     }
 
     @Test
     public void testNormalJoinGroupFollower() {
         final String consumerId = "consumer";
+        final Set subscription = singleton(topic1);
+        final List owned = Collections.emptyList();
+        final List assigned = Arrays.asList(t1p);
 
-        subscriptions.subscribe(singleton(topic1), rebalanceListener);
+        subscriptions.subscribe(subscription, rebalanceListener);
 
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
         coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
         // normal join group
         client.prepareResponse(joinGroupFollowerResponse(1, consumerId, "leader", Errors.NONE));
-        client.prepareResponse(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                SyncGroupRequest sync = (SyncGroupRequest) body;
-                return sync.memberId().equals(consumerId) &&
-                        sync.generationId() == 1 &&
-                        sync.groupAssignment().isEmpty();
-            }
-        }, syncGroupResponse(singletonList(t1p), Errors.NONE));
+        client.prepareResponse(body -> {
+            SyncGroupRequest sync = (SyncGroupRequest) body;
+            return sync.data.memberId().equals(consumerId) &&
+                    sync.data.generationId() == 1 &&
+                    sync.groupAssignments().isEmpty();
+        }, syncGroupResponse(assigned, Errors.NONE));
 
         coordinator.joinGroupIfNeeded(time.timer(Long.MAX_VALUE));
 
         assertFalse(coordinator.rejoinNeededOrPending());
-        assertEquals(singleton(t1p), subscriptions.assignedPartitions());
-        assertEquals(singleton(topic1), subscriptions.groupSubscription());
-        assertEquals(1, rebalanceListener.revokedCount);
-        assertEquals(Collections.emptySet(), rebalanceListener.revoked);
+        assertEquals(toSet(assigned), subscriptions.assignedPartitions());
+        assertEquals(subscription, subscriptions.metadataTopics());
+        assertEquals(0, rebalanceListener.revokedCount);
+        assertNull(rebalanceListener.revoked);
         assertEquals(1, rebalanceListener.assignedCount);
-        assertEquals(singleton(t1p), rebalanceListener.assigned);
+        assertEquals(getAdded(owned, assigned), rebalanceListener.assigned);
     }
 
     @Test
@@ -739,15 +1024,18 @@ public void testUpdateLastHeartbeatPollWhenCoordinatorUnknown() throws Exception
 
         assertTrue(coordinator.coordinatorUnknown());
         assertFalse(coordinator.poll(time.timer(0)));
-        assertEquals(time.milliseconds(), heartbeat.lastPollTime());
+        assertEquals(time.milliseconds(), coordinator.heartbeat().lastPollTime());
 
         time.sleep(rebalanceTimeoutMs - 1);
-        assertFalse(heartbeat.pollTimeoutExpired(time.milliseconds()));
+        assertFalse(coordinator.heartbeat().pollTimeoutExpired(time.milliseconds()));
     }
 
     @Test
     public void testPatternJoinGroupFollower() {
         final String consumerId = "consumer";
+        final Set subscription = Utils.mkSet(topic1, topic2);
+        final List owned = Collections.emptyList();
+        final List assigned = Arrays.asList(t1p, t2p);
 
         subscriptions.subscribe(Pattern.compile("test.*"), rebalanceListener);
 
@@ -760,26 +1048,24 @@ public void testPatternJoinGroupFollower() {
 
         // normal join group
         client.prepareResponse(joinGroupFollowerResponse(1, consumerId, "leader", Errors.NONE));
-        client.prepareResponse(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                SyncGroupRequest sync = (SyncGroupRequest) body;
-                return sync.memberId().equals(consumerId) &&
-                        sync.generationId() == 1 &&
-                        sync.groupAssignment().isEmpty();
-            }
-        }, syncGroupResponse(Arrays.asList(t1p, t2p), Errors.NONE));
+        client.prepareResponse(body -> {
+            SyncGroupRequest sync = (SyncGroupRequest) body;
+            return sync.data.memberId().equals(consumerId) &&
+                    sync.data.generationId() == 1 &&
+                    sync.groupAssignments().isEmpty();
+        }, syncGroupResponse(assigned, Errors.NONE));
         // expect client to force updating the metadata, if yes gives it both topics
         client.prepareMetadataUpdate(metadataResponse);
 
         coordinator.joinGroupIfNeeded(time.timer(Long.MAX_VALUE));
 
         assertFalse(coordinator.rejoinNeededOrPending());
-        assertEquals(2, subscriptions.numAssignedPartitions());
-        assertEquals(2, subscriptions.subscription().size());
-        assertEquals(1, rebalanceListener.revokedCount);
+        assertEquals(assigned.size(), subscriptions.numAssignedPartitions());
+        assertEquals(subscription, subscriptions.subscription());
+        assertEquals(0, rebalanceListener.revokedCount);
+        assertNull(rebalanceListener.revoked);
         assertEquals(1, rebalanceListener.assignedCount);
-        assertEquals(2, rebalanceListener.assigned.size());
+        assertEquals(getAdded(owned, assigned), rebalanceListener.assigned);
     }
 
     @Test
@@ -790,15 +1076,12 @@ public void testLeaveGroupOnClose() {
         joinAsFollowerAndReceiveAssignment(consumerId, coordinator, singletonList(t1p));
 
         final AtomicBoolean received = new AtomicBoolean(false);
-        client.prepareResponse(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                received.set(true);
-                LeaveGroupRequest leaveRequest = (LeaveGroupRequest) body;
-                return leaveRequest.data().memberId().equals(consumerId) &&
-                        leaveRequest.data().groupId().equals(groupId);
-            }
-        }, new LeaveGroupResponse(new LeaveGroupResponseData().setErrorCode(Errors.NONE.code())));
+        client.prepareResponse(body -> {
+            received.set(true);
+            LeaveGroupRequest leaveRequest = (LeaveGroupRequest) body;
+            return validateLeaveGroup(groupId, consumerId, leaveRequest);
+        }, new LeaveGroupResponse(
+            new LeaveGroupResponseData().setErrorCode(Errors.NONE.code())));
         coordinator.close(time.timer(0));
         assertTrue(received.get());
     }
@@ -811,22 +1094,27 @@ public void testMaybeLeaveGroup() {
         joinAsFollowerAndReceiveAssignment(consumerId, coordinator, singletonList(t1p));
 
         final AtomicBoolean received = new AtomicBoolean(false);
-        client.prepareResponse(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                received.set(true);
-                LeaveGroupRequest leaveRequest = (LeaveGroupRequest) body;
-                return leaveRequest.data().memberId().equals(consumerId) &&
-                        leaveRequest.data().groupId().equals(groupId);
-            }
+        client.prepareResponse(body -> {
+            received.set(true);
+            LeaveGroupRequest leaveRequest = (LeaveGroupRequest) body;
+            return validateLeaveGroup(groupId, consumerId, leaveRequest);
         }, new LeaveGroupResponse(new LeaveGroupResponseData().setErrorCode(Errors.NONE.code())));
-        coordinator.maybeLeaveGroup();
+        coordinator.maybeLeaveGroup("test maybe leave group");
         assertTrue(received.get());
 
-        AbstractCoordinator.Generation generation = coordinator.generation();
+        AbstractCoordinator.Generation generation = coordinator.generationIfStable();
         assertNull(generation);
     }
 
+    private boolean validateLeaveGroup(String groupId,
+                                       String consumerId,
+                                       LeaveGroupRequest leaveRequest) {
+        List members = leaveRequest.data().members();
+        return leaveRequest.data().groupId().equals(groupId) &&
+                   members.size() == 1 &&
+                   members.get(0).memberId().equals(consumerId);
+    }
+
     /**
      * This test checks if a consumer that has a valid member ID but an invalid generation
      * ({@link org.apache.kafka.clients.consumer.internals.AbstractCoordinator.Generation#NO_GENERATION})
@@ -849,16 +1137,13 @@ public void testPendingMemberShouldLeaveGroup() {
         coordinator.joinGroupIfNeeded(time.timer(0));
 
         final AtomicBoolean received = new AtomicBoolean(false);
-        client.prepareResponse(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                received.set(true);
-                LeaveGroupRequest leaveRequest = (LeaveGroupRequest) body;
-                return leaveRequest.data().memberId().equals(consumerId);
-            }
+        client.prepareResponse(body -> {
+            received.set(true);
+            LeaveGroupRequest leaveRequest = (LeaveGroupRequest) body;
+            return validateLeaveGroup(groupId, consumerId, leaveRequest);
         }, new LeaveGroupResponse(new LeaveGroupResponseData().setErrorCode(Errors.NONE.code())));
 
-        coordinator.maybeLeaveGroup();
+        coordinator.maybeLeaveGroup("pending member leaves");
         assertTrue(received.get());
     }
 
@@ -888,15 +1173,12 @@ public void testUnknownMemberIdOnSyncGroup() {
 
         // join initially, but let coordinator returns unknown member id
         client.prepareResponse(joinGroupFollowerResponse(1, consumerId, "leader", Errors.NONE));
-        client.prepareResponse(syncGroupResponse(Collections.emptyList(), Errors.UNKNOWN_MEMBER_ID));
+        client.prepareResponse(syncGroupResponse(Collections.emptyList(), Errors.UNKNOWN_MEMBER_ID));
 
         // now we should see a new join with the empty UNKNOWN_MEMBER_ID
-        client.prepareResponse(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                JoinGroupRequest joinRequest = (JoinGroupRequest) body;
-                return joinRequest.memberId().equals(JoinGroupRequest.UNKNOWN_MEMBER_ID);
-            }
+        client.prepareResponse(body -> {
+            JoinGroupRequest joinRequest = (JoinGroupRequest) body;
+            return joinRequest.data().memberId().equals(JoinGroupRequest.UNKNOWN_MEMBER_ID);
         }, joinGroupFollowerResponse(2, consumerId, "leader", Errors.NONE));
         client.prepareResponse(syncGroupResponse(singletonList(t1p), Errors.NONE));
 
@@ -917,7 +1199,7 @@ public void testRebalanceInProgressOnSyncGroup() {
 
         // join initially, but let coordinator rebalance on sync
         client.prepareResponse(joinGroupFollowerResponse(1, consumerId, "leader", Errors.NONE));
-        client.prepareResponse(syncGroupResponse(Collections.emptyList(), Errors.REBALANCE_IN_PROGRESS));
+        client.prepareResponse(syncGroupResponse(Collections.emptyList(), Errors.REBALANCE_IN_PROGRESS));
 
         // then let the full join/sync finish successfully
         client.prepareResponse(joinGroupFollowerResponse(2, consumerId, "leader", Errors.NONE));
@@ -940,15 +1222,12 @@ public void testIllegalGenerationOnSyncGroup() {
 
         // join initially, but let coordinator rebalance on sync
         client.prepareResponse(joinGroupFollowerResponse(1, consumerId, "leader", Errors.NONE));
-        client.prepareResponse(syncGroupResponse(Collections.emptyList(), Errors.ILLEGAL_GENERATION));
+        client.prepareResponse(syncGroupResponse(Collections.emptyList(), Errors.ILLEGAL_GENERATION));
 
         // then let the full join/sync finish successfully
-        client.prepareResponse(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                JoinGroupRequest joinRequest = (JoinGroupRequest) body;
-                return joinRequest.memberId().equals(JoinGroupRequest.UNKNOWN_MEMBER_ID);
-            }
+        client.prepareResponse(body -> {
+            JoinGroupRequest joinRequest = (JoinGroupRequest) body;
+            return joinRequest.data().memberId().equals(JoinGroupRequest.UNKNOWN_MEMBER_ID);
         }, joinGroupFollowerResponse(2, consumerId, "leader", Errors.NONE));
         client.prepareResponse(syncGroupResponse(singletonList(t1p), Errors.NONE));
 
@@ -970,7 +1249,7 @@ public void testMetadataChangeTriggersRebalance() {
         coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
         Map> memberSubscriptions = singletonMap(consumerId, singletonList(topic1));
-        partitionAssignor.prepare(singletonMap(consumerId, singletonList(t1p)));
+        partitionAssignor.prepare(singletonMap(consumerId, Arrays.asList(t1p)));
 
         // the leader is responsible for picking up metadata changes and forcing a group rebalance
         client.prepareResponse(joinGroupLeaderResponse(1, consumerId, memberSubscriptions, Errors.NONE));
@@ -1008,25 +1287,22 @@ public void testUpdateMetadataDuringRebalance() {
 
         // prepare initial rebalance
         Map> memberSubscriptions = singletonMap(consumerId, topics);
-        partitionAssignor.prepare(singletonMap(consumerId, Collections.singletonList(tp1)));
+        partitionAssignor.prepare(singletonMap(consumerId, Arrays.asList(tp1)));
 
         client.prepareResponse(joinGroupLeaderResponse(1, consumerId, memberSubscriptions, Errors.NONE));
-        client.prepareResponse(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                SyncGroupRequest sync = (SyncGroupRequest) body;
-                if (sync.memberId().equals(consumerId) &&
-                        sync.generationId() == 1 &&
-                        sync.groupAssignment().containsKey(consumerId)) {
-                    // trigger the metadata update including both topics after the sync group request has been sent
-                    Map topicPartitionCounts = new HashMap<>();
-                    topicPartitionCounts.put(topic1, 1);
-                    topicPartitionCounts.put(topic2, 1);
-                    client.updateMetadata(TestUtils.metadataUpdateWith(1, topicPartitionCounts));
-                    return true;
-                }
-                return false;
+        client.prepareResponse(body -> {
+            SyncGroupRequest sync = (SyncGroupRequest) body;
+            if (sync.data.memberId().equals(consumerId) &&
+                    sync.data.generationId() == 1 &&
+                    sync.groupAssignments().containsKey(consumerId)) {
+                // trigger the metadata update including both topics after the sync group request has been sent
+                Map topicPartitionCounts = new HashMap<>();
+                topicPartitionCounts.put(topic1, 1);
+                topicPartitionCounts.put(topic2, 1);
+                client.updateMetadata(TestUtils.metadataUpdateWith(1, topicPartitionCounts));
+                return true;
             }
+            return false;
         }, syncGroupResponse(Collections.singletonList(tp1), Errors.NONE));
         coordinator.poll(time.timer(Long.MAX_VALUE));
 
@@ -1040,11 +1316,47 @@ public boolean matches(AbstractRequest body) {
         assertEquals(new HashSet<>(Arrays.asList(tp1, tp2)), subscriptions.assignedPartitions());
     }
 
+    /**
+     * Verifies that subscription change updates SubscriptionState correctly even after JoinGroup failures
+     * that don't re-invoke onJoinPrepare.
+     */
     @Test
-    public void testWakeupFromAssignmentCallback() {
-        ConsumerCoordinator coordinator = buildCoordinator(new Metrics(), assignors,
-                false, true);
+    public void testSubscriptionChangeWithAuthorizationFailure() {
+        final String consumerId = "consumer";
+
+        // Subscribe to two topics of which only one is authorized and verify that metadata failure is propagated.
+        subscriptions.subscribe(Utils.mkSet(topic1, topic2), rebalanceListener);
+        client.prepareMetadataUpdate(TestUtils.metadataUpdateWith("kafka-cluster", 1,
+                Collections.singletonMap(topic2, Errors.TOPIC_AUTHORIZATION_FAILED), singletonMap(topic1, 1)));
+        assertThrows(TopicAuthorizationException.class, () -> coordinator.poll(time.timer(Long.MAX_VALUE)));
+
+        client.respond(groupCoordinatorResponse(node, Errors.NONE));
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
+
+        // Fail the first JoinGroup request
+        client.prepareResponse(joinGroupLeaderResponse(0, consumerId, Collections.emptyMap(),
+                Errors.GROUP_AUTHORIZATION_FAILED));
+        assertThrows(GroupAuthorizationException.class, () -> coordinator.poll(time.timer(Long.MAX_VALUE)));
+
+        // Change subscription to include only the authorized topic. Complete rebalance and check that
+        // references to topic2 have been removed from SubscriptionState.
+        subscriptions.subscribe(Utils.mkSet(topic1), rebalanceListener);
+        assertEquals(Collections.singleton(topic1), subscriptions.metadataTopics());
+        client.prepareMetadataUpdate(TestUtils.metadataUpdateWith("kafka-cluster", 1,
+                Collections.emptyMap(), singletonMap(topic1, 1)));
+
+        Map> memberSubscriptions = singletonMap(consumerId, singletonList(topic1));
+        partitionAssignor.prepare(singletonMap(consumerId, Arrays.asList(t1p)));
+        client.prepareResponse(joinGroupLeaderResponse(1, consumerId, memberSubscriptions, Errors.NONE));
+        client.prepareResponse(syncGroupResponse(singletonList(t1p), Errors.NONE));
+        coordinator.poll(time.timer(Long.MAX_VALUE));
 
+        assertEquals(singleton(topic1), subscriptions.subscription());
+        assertEquals(singleton(topic1), subscriptions.metadataTopics());
+    }
+
+    @Test
+    public void testWakeupFromAssignmentCallback() {
         final String topic = "topic1";
         TopicPartition partition = new TopicPartition(topic, 0);
         final String consumerId = "follower";
@@ -1074,7 +1386,6 @@ public void onPartitionsAssigned(Collection partitions) {
         client.prepareResponse(joinGroupFollowerResponse(1, consumerId, "leader", Errors.NONE));
         client.prepareResponse(syncGroupResponse(Collections.singletonList(partition), Errors.NONE));
 
-
         // The first call to poll should raise the exception from the rebalance listener
         try {
             coordinator.poll(time.timer(Long.MAX_VALUE));
@@ -1086,7 +1397,7 @@ public void onPartitionsAssigned(Collection partitions) {
         coordinator.poll(time.timer(Long.MAX_VALUE));
 
         assertFalse(coordinator.rejoinNeededOrPending());
-        assertEquals(1, rebalanceListener.revokedCount);
+        assertEquals(0, rebalanceListener.revokedCount);
         assertEquals(2, rebalanceListener.assignedCount);
     }
 
@@ -1126,6 +1437,7 @@ private void unavailableTopicTest(boolean patternSubscribe, Set unavaila
         client.prepareResponse(syncGroupResponse(Collections.emptyList(), Errors.NONE));
         coordinator.poll(time.timer(Long.MAX_VALUE));
         assertFalse(coordinator.rejoinNeededOrPending());
+        // callback not triggered since there's nothing to be assigned
         assertEquals(Collections.emptySet(), rebalanceListener.assigned);
         assertTrue("Metadata refresh not requested for unavailable partitions", metadata.updateRequested());
 
@@ -1158,55 +1470,62 @@ public void testIncludeInternalTopicsConfigOption() {
 
     private void testInternalTopicInclusion(boolean includeInternalTopics) {
         metadata = new ConsumerMetadata(0, Long.MAX_VALUE, includeInternalTopics,
-                subscriptions, new LogContext(), new ClusterResourceListeners());
+                false, subscriptions, new LogContext(), new ClusterResourceListeners());
         client = new MockClient(time, metadata);
-        coordinator = buildCoordinator(new Metrics(), assignors, false, true);
-
-        subscriptions.subscribe(Pattern.compile(".*"), rebalanceListener);
-
-        Node node = new Node(0, "localhost", 9999);
-        MetadataResponse.PartitionMetadata partitionMetadata =
+        try (ConsumerCoordinator coordinator = buildCoordinator(rebalanceConfig, new Metrics(), assignors, false)) {
+            subscriptions.subscribe(Pattern.compile(".*"), rebalanceListener);
+            Node node = new Node(0, "localhost", 9999);
+            MetadataResponse.PartitionMetadata partitionMetadata =
                 new MetadataResponse.PartitionMetadata(Errors.NONE, 0, node, Optional.empty(),
-                        singletonList(node), singletonList(node), singletonList(node));
-        MetadataResponse.TopicMetadata topicMetadata = new MetadataResponse.TopicMetadata(Errors.NONE,
+                    singletonList(node), singletonList(node), singletonList(node));
+            MetadataResponse.TopicMetadata topicMetadata = new MetadataResponse.TopicMetadata(Errors.NONE,
                 Topic.GROUP_METADATA_TOPIC_NAME, true, singletonList(partitionMetadata));
 
-        client.updateMetadata(new MetadataResponse(singletonList(node), "clusterId", node.id(),
+            client.updateMetadata(MetadataResponse.prepareResponse(singletonList(node), "clusterId", node.id(),
                 singletonList(topicMetadata)));
-        coordinator.maybeUpdateSubscriptionMetadata();
+            coordinator.maybeUpdateSubscriptionMetadata();
 
-        assertEquals(includeInternalTopics, subscriptions.subscription().contains(Topic.GROUP_METADATA_TOPIC_NAME));
+            assertEquals(includeInternalTopics, subscriptions.subscription().contains(Topic.GROUP_METADATA_TOPIC_NAME));
+        }
     }
 
     @Test
     public void testRejoinGroup() {
         String otherTopic = "otherTopic";
+        final List owned = Collections.emptyList();
+        final List assigned = Arrays.asList(t1p);
 
         subscriptions.subscribe(singleton(topic1), rebalanceListener);
 
         // join the group once
-        joinAsFollowerAndReceiveAssignment("consumer", coordinator, singletonList(t1p));
+        joinAsFollowerAndReceiveAssignment("consumer", coordinator, assigned);
 
-        assertEquals(1, rebalanceListener.revokedCount);
-        assertTrue(rebalanceListener.revoked.isEmpty());
+        assertEquals(0, rebalanceListener.revokedCount);
+        assertNull(rebalanceListener.revoked);
         assertEquals(1, rebalanceListener.assignedCount);
-        assertEquals(singleton(t1p), rebalanceListener.assigned);
+        assertEquals(getAdded(owned, assigned), rebalanceListener.assigned);
 
         // and join the group again
+        rebalanceListener.revoked = null;
+        rebalanceListener.assigned = null;
         subscriptions.subscribe(new HashSet<>(Arrays.asList(topic1, otherTopic)), rebalanceListener);
         client.prepareResponse(joinGroupFollowerResponse(2, "consumer", "leader", Errors.NONE));
-        client.prepareResponse(syncGroupResponse(singletonList(t1p), Errors.NONE));
+        client.prepareResponse(syncGroupResponse(assigned, Errors.NONE));
         coordinator.joinGroupIfNeeded(time.timer(Long.MAX_VALUE));
 
-        assertEquals(2, rebalanceListener.revokedCount);
-        assertEquals(singleton(t1p), rebalanceListener.revoked);
+        Collection revoked = getRevoked(assigned, assigned);
+        Collection added = getAdded(assigned, assigned);
+        assertEquals(revoked.isEmpty() ? 0 : 1, rebalanceListener.revokedCount);
+        assertEquals(revoked.isEmpty() ? null : revoked, rebalanceListener.revoked);
         assertEquals(2, rebalanceListener.assignedCount);
-        assertEquals(singleton(t1p), rebalanceListener.assigned);
+        assertEquals(added, rebalanceListener.assigned);
     }
 
     @Test
     public void testDisconnectInJoin() {
         subscriptions.subscribe(singleton(topic1), rebalanceListener);
+        final List owned = Collections.emptyList();
+        final List assigned = Arrays.asList(t1p);
 
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
         coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
@@ -1215,14 +1534,16 @@ public void testDisconnectInJoin() {
         client.prepareResponse(joinGroupFollowerResponse(1, "consumer", "leader", Errors.NONE), true);
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
         client.prepareResponse(joinGroupFollowerResponse(1, "consumer", "leader", Errors.NONE));
-        client.prepareResponse(syncGroupResponse(singletonList(t1p), Errors.NONE));
+        client.prepareResponse(syncGroupResponse(assigned, Errors.NONE));
         coordinator.joinGroupIfNeeded(time.timer(Long.MAX_VALUE));
 
         assertFalse(coordinator.rejoinNeededOrPending());
-        assertEquals(singleton(t1p), subscriptions.assignedPartitions());
-        assertEquals(1, rebalanceListener.revokedCount);
+        assertEquals(toSet(assigned), subscriptions.assignedPartitions());
+        // nothing to be revoked hence callback not triggered
+        assertEquals(0, rebalanceListener.revokedCount);
+        assertNull(rebalanceListener.revoked);
         assertEquals(1, rebalanceListener.assignedCount);
-        assertEquals(singleton(t1p), rebalanceListener.assigned);
+        assertEquals(getAdded(owned, assigned), rebalanceListener.assigned);
     }
 
     @Test(expected = ApiException.class)
@@ -1288,158 +1609,153 @@ private void testInFlightRequestsFailedAfterCoordinatorMarkedDead(Errors error)
     public void testAutoCommitDynamicAssignment() {
         final String consumerId = "consumer";
 
-        ConsumerCoordinator coordinator = buildCoordinator(new Metrics(), assignors,
-                true, true);
-
-        subscriptions.subscribe(singleton(topic1), rebalanceListener);
-        joinAsFollowerAndReceiveAssignment(consumerId, coordinator, singletonList(t1p));
-        subscriptions.seek(t1p, 100);
-
-        prepareOffsetCommitRequest(singletonMap(t1p, 100L), Errors.NONE);
-        time.sleep(autoCommitIntervalMs);
-        coordinator.poll(time.timer(Long.MAX_VALUE));
-        assertFalse(client.hasPendingResponses());
+        try (ConsumerCoordinator coordinator = buildCoordinator(rebalanceConfig, new Metrics(), assignors, true)
+        ) {
+            subscriptions.subscribe(singleton(topic1), rebalanceListener);
+            joinAsFollowerAndReceiveAssignment(consumerId, coordinator, singletonList(t1p));
+            subscriptions.seek(t1p, 100);
+            prepareOffsetCommitRequest(singletonMap(t1p, 100L), Errors.NONE);
+            time.sleep(autoCommitIntervalMs);
+            coordinator.poll(time.timer(Long.MAX_VALUE));
+            assertFalse(client.hasPendingResponses());
+        }
     }
 
     @Test
     public void testAutoCommitRetryBackoff() {
         final String consumerId = "consumer";
-        ConsumerCoordinator coordinator = buildCoordinator(new Metrics(), assignors,
-                true, true);
-        subscriptions.subscribe(singleton(topic1), rebalanceListener);
-        joinAsFollowerAndReceiveAssignment(consumerId, coordinator, singletonList(t1p));
 
-        subscriptions.seek(t1p, 100);
-        time.sleep(autoCommitIntervalMs);
+        try (ConsumerCoordinator coordinator = buildCoordinator(rebalanceConfig, new Metrics(), assignors, true)) {
+            subscriptions.subscribe(singleton(topic1), rebalanceListener);
+            joinAsFollowerAndReceiveAssignment(consumerId, coordinator, singletonList(t1p));
 
-        // Send an offset commit, but let it fail with a retriable error
-        prepareOffsetCommitRequest(singletonMap(t1p, 100L), Errors.NOT_COORDINATOR);
-        coordinator.poll(time.timer(Long.MAX_VALUE));
-        assertTrue(coordinator.coordinatorUnknown());
+            subscriptions.seek(t1p, 100);
+            time.sleep(autoCommitIntervalMs);
 
-        // After the disconnect, we should rediscover the coordinator
-        client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.poll(time.timer(Long.MAX_VALUE));
+            // Send an offset commit, but let it fail with a retriable error
+            prepareOffsetCommitRequest(singletonMap(t1p, 100L), Errors.NOT_COORDINATOR);
+            coordinator.poll(time.timer(Long.MAX_VALUE));
+            assertTrue(coordinator.coordinatorUnknown());
 
-        subscriptions.seek(t1p, 200);
+            // After the disconnect, we should rediscover the coordinator
+            client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+            coordinator.poll(time.timer(Long.MAX_VALUE));
 
-        // Until the retry backoff has expired, we should not retry the offset commit
-        time.sleep(retryBackoffMs / 2);
-        coordinator.poll(time.timer(Long.MAX_VALUE));
-        assertEquals(0, client.inFlightRequestCount());
+            subscriptions.seek(t1p, 200);
 
-        // Once the backoff expires, we should retry
-        time.sleep(retryBackoffMs / 2);
-        coordinator.poll(time.timer(Long.MAX_VALUE));
-        assertEquals(1, client.inFlightRequestCount());
-        respondToOffsetCommitRequest(singletonMap(t1p, 200L), Errors.NONE);
+            // Until the retry backoff has expired, we should not retry the offset commit
+            time.sleep(retryBackoffMs / 2);
+            coordinator.poll(time.timer(Long.MAX_VALUE));
+            assertEquals(0, client.inFlightRequestCount());
+
+            // Once the backoff expires, we should retry
+            time.sleep(retryBackoffMs / 2);
+            coordinator.poll(time.timer(Long.MAX_VALUE));
+            assertEquals(1, client.inFlightRequestCount());
+            respondToOffsetCommitRequest(singletonMap(t1p, 200L), Errors.NONE);
+        }
     }
 
     @Test
     public void testAutoCommitAwaitsInterval() {
         final String consumerId = "consumer";
-        ConsumerCoordinator coordinator = buildCoordinator(new Metrics(), assignors,
-                true, true);
-        subscriptions.subscribe(singleton(topic1), rebalanceListener);
-        joinAsFollowerAndReceiveAssignment(consumerId, coordinator, singletonList(t1p));
+        try (ConsumerCoordinator coordinator = buildCoordinator(rebalanceConfig, new Metrics(), assignors, true)) {
+            subscriptions.subscribe(singleton(topic1), rebalanceListener);
+            joinAsFollowerAndReceiveAssignment(consumerId, coordinator, singletonList(t1p));
 
-        subscriptions.seek(t1p, 100);
-        time.sleep(autoCommitIntervalMs);
+            subscriptions.seek(t1p, 100);
+            time.sleep(autoCommitIntervalMs);
 
-        // Send the offset commit request, but do not respond
-        coordinator.poll(time.timer(Long.MAX_VALUE));
-        assertEquals(1, client.inFlightRequestCount());
+            // Send the offset commit request, but do not respond
+            coordinator.poll(time.timer(Long.MAX_VALUE));
+            assertEquals(1, client.inFlightRequestCount());
 
-        time.sleep(autoCommitIntervalMs / 2);
+            time.sleep(autoCommitIntervalMs / 2);
 
-        // Ensure that no additional offset commit is sent
-        coordinator.poll(time.timer(Long.MAX_VALUE));
-        assertEquals(1, client.inFlightRequestCount());
+            // Ensure that no additional offset commit is sent
+            coordinator.poll(time.timer(Long.MAX_VALUE));
+            assertEquals(1, client.inFlightRequestCount());
 
-        respondToOffsetCommitRequest(singletonMap(t1p, 100L), Errors.NONE);
-        coordinator.poll(time.timer(Long.MAX_VALUE));
-        assertEquals(0, client.inFlightRequestCount());
+            respondToOffsetCommitRequest(singletonMap(t1p, 100L), Errors.NONE);
+            coordinator.poll(time.timer(Long.MAX_VALUE));
+            assertEquals(0, client.inFlightRequestCount());
 
-        subscriptions.seek(t1p, 200);
+            subscriptions.seek(t1p, 200);
 
-        // If we poll again before the auto-commit interval, there should be no new sends
-        coordinator.poll(time.timer(Long.MAX_VALUE));
-        assertEquals(0, client.inFlightRequestCount());
+            // If we poll again before the auto-commit interval, there should be no new sends
+            coordinator.poll(time.timer(Long.MAX_VALUE));
+            assertEquals(0, client.inFlightRequestCount());
 
-        // After the remainder of the interval passes, we send a new offset commit
-        time.sleep(autoCommitIntervalMs / 2);
-        coordinator.poll(time.timer(Long.MAX_VALUE));
-        assertEquals(1, client.inFlightRequestCount());
-        respondToOffsetCommitRequest(singletonMap(t1p, 200L), Errors.NONE);
+            // After the remainder of the interval passes, we send a new offset commit
+            time.sleep(autoCommitIntervalMs / 2);
+            coordinator.poll(time.timer(Long.MAX_VALUE));
+            assertEquals(1, client.inFlightRequestCount());
+            respondToOffsetCommitRequest(singletonMap(t1p, 200L), Errors.NONE);
+        }
     }
 
     @Test
     public void testAutoCommitDynamicAssignmentRebalance() {
         final String consumerId = "consumer";
 
-        ConsumerCoordinator coordinator = buildCoordinator(new Metrics(), assignors,
-                true, true);
-
-        subscriptions.subscribe(singleton(topic1), rebalanceListener);
-
-        client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
+        try (ConsumerCoordinator coordinator = buildCoordinator(rebalanceConfig, new Metrics(), assignors, true)) {
+            subscriptions.subscribe(singleton(topic1), rebalanceListener);
+            client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+            coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
-        // haven't joined, so should not cause a commit
-        time.sleep(autoCommitIntervalMs);
-        consumerClient.poll(time.timer(0));
+            // haven't joined, so should not cause a commit
+            time.sleep(autoCommitIntervalMs);
+            consumerClient.poll(time.timer(0));
 
-        client.prepareResponse(joinGroupFollowerResponse(1, consumerId, "leader", Errors.NONE));
-        client.prepareResponse(syncGroupResponse(singletonList(t1p), Errors.NONE));
-        coordinator.joinGroupIfNeeded(time.timer(Long.MAX_VALUE));
+            client.prepareResponse(joinGroupFollowerResponse(1, consumerId, "leader", Errors.NONE));
+            client.prepareResponse(syncGroupResponse(singletonList(t1p), Errors.NONE));
+            coordinator.joinGroupIfNeeded(time.timer(Long.MAX_VALUE));
 
-        subscriptions.seek(t1p, 100);
+            subscriptions.seek(t1p, 100);
 
-        prepareOffsetCommitRequest(singletonMap(t1p, 100L), Errors.NONE);
-        time.sleep(autoCommitIntervalMs);
-        coordinator.poll(time.timer(Long.MAX_VALUE));
-        assertFalse(client.hasPendingResponses());
+            prepareOffsetCommitRequest(singletonMap(t1p, 100L), Errors.NONE);
+            time.sleep(autoCommitIntervalMs);
+            coordinator.poll(time.timer(Long.MAX_VALUE));
+            assertFalse(client.hasPendingResponses());
+        }
     }
 
     @Test
     public void testAutoCommitManualAssignment() {
-        ConsumerCoordinator coordinator = buildCoordinator(new Metrics(), assignors,
-                true, true);
-
-        subscriptions.assignFromUser(singleton(t1p));
-        subscriptions.seek(t1p, 100);
-
-        client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
+        try (ConsumerCoordinator coordinator = buildCoordinator(rebalanceConfig, new Metrics(), assignors, true)) {
+            subscriptions.assignFromUser(singleton(t1p));
+            subscriptions.seek(t1p, 100);
+            client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+            coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
-        prepareOffsetCommitRequest(singletonMap(t1p, 100L), Errors.NONE);
-        time.sleep(autoCommitIntervalMs);
-        coordinator.poll(time.timer(Long.MAX_VALUE));
-        assertFalse(client.hasPendingResponses());
+            prepareOffsetCommitRequest(singletonMap(t1p, 100L), Errors.NONE);
+            time.sleep(autoCommitIntervalMs);
+            coordinator.poll(time.timer(Long.MAX_VALUE));
+            assertFalse(client.hasPendingResponses());
+        }
     }
 
     @Test
     public void testAutoCommitManualAssignmentCoordinatorUnknown() {
-        ConsumerCoordinator coordinator = buildCoordinator(new Metrics(), assignors,
-                true, true);
+        try (ConsumerCoordinator coordinator = buildCoordinator(rebalanceConfig, new Metrics(), assignors, true)) {
+            subscriptions.assignFromUser(singleton(t1p));
+            subscriptions.seek(t1p, 100);
 
-        subscriptions.assignFromUser(singleton(t1p));
-        subscriptions.seek(t1p, 100);
+            // no commit initially since coordinator is unknown
+            consumerClient.poll(time.timer(0));
+            time.sleep(autoCommitIntervalMs);
+            consumerClient.poll(time.timer(0));
 
-        // no commit initially since coordinator is unknown
-        consumerClient.poll(time.timer(0));
-        time.sleep(autoCommitIntervalMs);
-        consumerClient.poll(time.timer(0));
+            // now find the coordinator
+            client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+            coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
-        // now find the coordinator
-        client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
-
-        // sleep only for the retry backoff
-        time.sleep(retryBackoffMs);
-        prepareOffsetCommitRequest(singletonMap(t1p, 100L), Errors.NONE);
-        coordinator.poll(time.timer(Long.MAX_VALUE));
-        assertFalse(client.hasPendingResponses());
+            // sleep only for the retry backoff
+            time.sleep(retryBackoffMs);
+            prepareOffsetCommitRequest(singletonMap(t1p, 100L), Errors.NONE);
+            coordinator.poll(time.timer(Long.MAX_VALUE));
+            assertFalse(client.hasPendingResponses());
+        }
     }
 
     @Test
@@ -1479,19 +1795,17 @@ public void testCommitAfterLeaveGroup() {
         joinAsFollowerAndReceiveAssignment("consumer", coordinator, singletonList(t1p));
 
         // now switch to manual assignment
-        client.prepareResponse(new LeaveGroupResponse(new LeaveGroupResponseData().setErrorCode(Errors.NONE.code())));
+        client.prepareResponse(new LeaveGroupResponse(new LeaveGroupResponseData()
+                .setErrorCode(Errors.NONE.code())));
         subscriptions.unsubscribe();
-        coordinator.maybeLeaveGroup();
+        coordinator.maybeLeaveGroup("test commit after leave");
         subscriptions.assignFromUser(singleton(t1p));
 
         // the client should not reuse generation/memberId from auto-subscribed generation
-        client.prepareResponse(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                OffsetCommitRequest commitRequest = (OffsetCommitRequest) body;
-                return commitRequest.memberId().equals(OffsetCommitRequest.DEFAULT_MEMBER_ID) &&
-                        commitRequest.generationId() == OffsetCommitRequest.DEFAULT_GENERATION_ID;
-            }
+        client.prepareResponse(body -> {
+            OffsetCommitRequest commitRequest = (OffsetCommitRequest) body;
+            return commitRequest.data().memberId().equals(OffsetCommitRequest.DEFAULT_MEMBER_ID) &&
+                    commitRequest.data().generationId() == OffsetCommitRequest.DEFAULT_GENERATION_ID;
         }, offsetCommitResponse(singletonMap(t1p, Errors.NONE)));
 
         AtomicBoolean success = new AtomicBoolean(false);
@@ -1601,7 +1915,7 @@ public void testAsyncCommitCallbacksInvokedPriorToSyncCommitCompletion() throws
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
         coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
-        final List committedOffsets = Collections.synchronizedList(new ArrayList());
+        final List committedOffsets = Collections.synchronizedList(new ArrayList<>());
         final OffsetAndMetadata firstOffset = new OffsetAndMetadata(0L);
         final OffsetAndMetadata secondOffset = new OffsetAndMetadata(1L);
 
@@ -1677,15 +1991,43 @@ public void testCommitOffsetUnknownMemberId() {
                 new OffsetAndMetadata(100L, "metadata")), time.timer(Long.MAX_VALUE));
     }
 
-    @Test(expected = CommitFailedException.class)
+    @Test
     public void testCommitOffsetRebalanceInProgress() {
         // we cannot retry if a rebalance occurs before the commit completed
+        final String consumerId = "leader";
+
+        subscriptions.subscribe(singleton(topic1), rebalanceListener);
+
+        // ensure metadata is up-to-date for leader
+        client.updateMetadata(metadataResponse);
+
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
         coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
+        // normal join group
+        Map> memberSubscriptions = singletonMap(consumerId, singletonList(topic1));
+        partitionAssignor.prepare(singletonMap(consumerId, singletonList(t1p)));
+
+        client.prepareResponse(joinGroupLeaderResponse(1, consumerId, memberSubscriptions, Errors.NONE));
+        client.prepareResponse(body -> {
+            SyncGroupRequest sync = (SyncGroupRequest) body;
+            return sync.data.memberId().equals(consumerId) &&
+                sync.data.generationId() == 1 &&
+                sync.groupAssignments().containsKey(consumerId);
+        }, syncGroupResponse(singletonList(t1p), Errors.NONE));
+        coordinator.poll(time.timer(Long.MAX_VALUE));
+
+        AbstractCoordinator.Generation expectedGeneration = new AbstractCoordinator.Generation(1, consumerId, partitionAssignor.name());
+        assertFalse(coordinator.rejoinNeededOrPending());
+        assertEquals(expectedGeneration, coordinator.generationIfStable());
+
         prepareOffsetCommitRequest(singletonMap(t1p, 100L), Errors.REBALANCE_IN_PROGRESS);
-        coordinator.commitOffsetsSync(singletonMap(t1p,
-                new OffsetAndMetadata(100L, "metadata")), time.timer(Long.MAX_VALUE));
+
+        assertThrows(CommitFailedException.class, () -> coordinator.commitOffsetsSync(singletonMap(t1p,
+            new OffsetAndMetadata(100L, "metadata")), time.timer(Long.MAX_VALUE)));
+
+        assertTrue(coordinator.rejoinNeededOrPending());
+        assertEquals(expectedGeneration, coordinator.generationIfStable());
     }
 
     @Test(expected = KafkaException.class)
@@ -1716,7 +2058,31 @@ public void testRefreshOffset() {
 
         assertEquals(Collections.emptySet(), subscriptions.missingFetchPositions());
         assertTrue(subscriptions.hasAllFetchPositions());
-        assertEquals(100L, subscriptions.position(t1p).longValue());
+        assertEquals(100L, subscriptions.position(t1p).offset);
+    }
+
+    @Test
+    public void testRefreshOffsetWithValidation() {
+        client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
+
+        subscriptions.assignFromUser(singleton(t1p));
+
+        // Initial leader epoch of 4
+        MetadataResponse metadataResponse = TestUtils.metadataUpdateWith("kafka-cluster", 1,
+                Collections.emptyMap(), singletonMap(topic1, 1), tp -> 4);
+        client.updateMetadata(metadataResponse);
+
+        // Load offsets from previous epoch
+        client.prepareResponse(offsetFetchResponse(t1p, Errors.NONE, "", 100L, 3));
+        coordinator.refreshCommittedOffsetsIfNeeded(time.timer(Long.MAX_VALUE));
+
+        // Offset gets loaded, but requires validation
+        assertEquals(Collections.emptySet(), subscriptions.missingFetchPositions());
+        assertFalse(subscriptions.hasAllFetchPositions());
+        assertTrue(subscriptions.awaitingValidation(t1p));
+        assertEquals(subscriptions.position(t1p).offset, 100L);
+        assertNull(subscriptions.validPosition(t1p));
     }
 
     @Test
@@ -1738,6 +2104,21 @@ public void testFetchCommittedOffsets() {
         assertEquals(new OffsetAndMetadata(offset, leaderEpoch, metadata), fetchedOffsets.get(t1p));
     }
 
+    @Test
+    public void testTopicAuthorizationFailedInOffsetFetch() {
+        client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
+
+        OffsetFetchResponse.PartitionData data = new OffsetFetchResponse.PartitionData(-1, Optional.empty(),
+                "", Errors.TOPIC_AUTHORIZATION_FAILED);
+
+        client.prepareResponse(new OffsetFetchResponse(Errors.NONE, singletonMap(t1p, data)));
+        TopicAuthorizationException exception = assertThrows(TopicAuthorizationException.class, () ->
+                coordinator.fetchCommittedOffsets(singleton(t1p), time.timer(Long.MAX_VALUE)));
+
+        assertEquals(singleton(topic1), exception.unauthorizedTopics());
+    }
+
     @Test
     public void testRefreshOffsetLoadInProgress() {
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
@@ -1750,7 +2131,7 @@ public void testRefreshOffsetLoadInProgress() {
 
         assertEquals(Collections.emptySet(), subscriptions.missingFetchPositions());
         assertTrue(subscriptions.hasAllFetchPositions());
-        assertEquals(100L, subscriptions.position(t1p).longValue());
+        assertEquals(100L, subscriptions.position(t1p).offset);
     }
 
     @Test
@@ -1791,7 +2172,7 @@ public void testRefreshOffsetNotCoordinatorForConsumer() {
 
         assertEquals(Collections.emptySet(), subscriptions.missingFetchPositions());
         assertTrue(subscriptions.hasAllFetchPositions());
-        assertEquals(100L, subscriptions.position(t1p).longValue());
+        assertEquals(100L, subscriptions.position(t1p).offset);
     }
 
     @Test
@@ -1806,7 +2187,7 @@ public void testRefreshOffsetWithNoFetchableOffsets() {
         assertEquals(Collections.singleton(t1p), subscriptions.missingFetchPositions());
         assertEquals(Collections.emptySet(), subscriptions.partitionsNeedingReset(time.milliseconds()));
         assertFalse(subscriptions.hasAllFetchPositions());
-        assertEquals(null, subscriptions.position(t1p));
+        assertNull(subscriptions.position(t1p));
     }
 
     @Test
@@ -1819,7 +2200,7 @@ public void testNoCoordinatorDiscoveryIfPositionsKnown() {
 
         assertEquals(Collections.emptySet(), subscriptions.missingFetchPositions());
         assertTrue(subscriptions.hasAllFetchPositions());
-        assertEquals(500L, subscriptions.position(t1p).longValue());
+        assertEquals(500L, subscriptions.position(t1p).offset);
         assertTrue(coordinator.coordinatorUnknown());
     }
 
@@ -1850,35 +2231,11 @@ public void testAuthenticationFailureInEnsureActiveGroup() {
         }
     }
 
-    @Test
-    public void testProtocolMetadataOrder() {
-        RoundRobinAssignor roundRobin = new RoundRobinAssignor();
-        RangeAssignor range = new RangeAssignor();
-
-        try (Metrics metrics = new Metrics(time)) {
-            ConsumerCoordinator coordinator = buildCoordinator(metrics, Arrays.asList(roundRobin, range),
-                    false, true);
-            List metadata = coordinator.metadata();
-            assertEquals(2, metadata.size());
-            assertEquals(roundRobin.name(), metadata.get(0).name());
-            assertEquals(range.name(), metadata.get(1).name());
-        }
-
-        try (Metrics metrics = new Metrics(time)) {
-            ConsumerCoordinator coordinator = buildCoordinator(metrics, Arrays.asList(range, roundRobin),
-                    false, true);
-            List metadata = coordinator.metadata();
-            assertEquals(2, metadata.size());
-            assertEquals(range.name(), metadata.get(0).name());
-            assertEquals(roundRobin.name(), metadata.get(1).name());
-        }
-    }
-
     @Test
     public void testThreadSafeAssignedPartitionsMetric() throws Exception {
         // Get the assigned-partitions metric
         final Metric metric = metrics.metric(new MetricName("assigned-partitions", "consumer" + groupId + "-coordinator-metrics",
-                "", Collections.emptyMap()));
+                "", Collections.emptyMap()));
 
         // Start polling the metric in the background
         final AtomicBoolean doStop = new AtomicBoolean();
@@ -1930,126 +2287,230 @@ public boolean conditionMet() {
 
     @Test
     public void testCloseDynamicAssignment() throws Exception {
-        ConsumerCoordinator coordinator = prepareCoordinatorForCloseTest(true, true, true);
-        gracefulCloseTest(coordinator, true);
+        try (ConsumerCoordinator coordinator = prepareCoordinatorForCloseTest(true, true, Optional.empty())) {
+            gracefulCloseTest(coordinator, true);
+        }
     }
 
     @Test
     public void testCloseManualAssignment() throws Exception {
-        ConsumerCoordinator coordinator = prepareCoordinatorForCloseTest(false, true, true);
-        gracefulCloseTest(coordinator, false);
-    }
-
-    @Test
-    public void shouldNotLeaveGroupWhenLeaveGroupFlagIsFalse() throws Exception {
-        final ConsumerCoordinator coordinator = prepareCoordinatorForCloseTest(true, true, false);
-        gracefulCloseTest(coordinator, false);
+        try (ConsumerCoordinator coordinator = prepareCoordinatorForCloseTest(false, true, Optional.empty())) {
+            gracefulCloseTest(coordinator, false);
+        }
     }
 
     @Test
     public void testCloseCoordinatorNotKnownManualAssignment() throws Exception {
-        ConsumerCoordinator coordinator = prepareCoordinatorForCloseTest(false, true, true);
-        makeCoordinatorUnknown(coordinator, Errors.NOT_COORDINATOR);
-        time.sleep(autoCommitIntervalMs);
-        closeVerifyTimeout(coordinator, 1000, 1000, 1000);
+        try (ConsumerCoordinator coordinator = prepareCoordinatorForCloseTest(false, true, Optional.empty())) {
+            makeCoordinatorUnknown(coordinator, Errors.NOT_COORDINATOR);
+            time.sleep(autoCommitIntervalMs);
+            closeVerifyTimeout(coordinator, 1000, 1000, 1000);
+        }
     }
 
     @Test
     public void testCloseCoordinatorNotKnownNoCommits() throws Exception {
-        ConsumerCoordinator coordinator = prepareCoordinatorForCloseTest(true, false, true);
-        makeCoordinatorUnknown(coordinator, Errors.NOT_COORDINATOR);
-        closeVerifyTimeout(coordinator, 1000, 0, 0);
+        try (ConsumerCoordinator coordinator = prepareCoordinatorForCloseTest(true, false, Optional.empty())) {
+            makeCoordinatorUnknown(coordinator, Errors.NOT_COORDINATOR);
+            closeVerifyTimeout(coordinator, 1000, 0, 0);
+        }
     }
 
     @Test
     public void testCloseCoordinatorNotKnownWithCommits() throws Exception {
-        ConsumerCoordinator coordinator = prepareCoordinatorForCloseTest(true, true, true);
-        makeCoordinatorUnknown(coordinator, Errors.NOT_COORDINATOR);
-        time.sleep(autoCommitIntervalMs);
-        closeVerifyTimeout(coordinator, 1000, 1000, 1000);
+        try (ConsumerCoordinator coordinator = prepareCoordinatorForCloseTest(true, true, Optional.empty())) {
+            makeCoordinatorUnknown(coordinator, Errors.NOT_COORDINATOR);
+            time.sleep(autoCommitIntervalMs);
+            closeVerifyTimeout(coordinator, 1000, 1000, 1000);
+        }
     }
 
     @Test
     public void testCloseCoordinatorUnavailableNoCommits() throws Exception {
-        ConsumerCoordinator coordinator = prepareCoordinatorForCloseTest(true, false, true);
-        makeCoordinatorUnknown(coordinator, Errors.COORDINATOR_NOT_AVAILABLE);
-        closeVerifyTimeout(coordinator, 1000, 0, 0);
+        try (ConsumerCoordinator coordinator = prepareCoordinatorForCloseTest(true, false, Optional.empty())) {
+            makeCoordinatorUnknown(coordinator, Errors.COORDINATOR_NOT_AVAILABLE);
+            closeVerifyTimeout(coordinator, 1000, 0, 0);
+        }
     }
 
     @Test
     public void testCloseTimeoutCoordinatorUnavailableForCommit() throws Exception {
-        ConsumerCoordinator coordinator = prepareCoordinatorForCloseTest(true, true, true);
-        makeCoordinatorUnknown(coordinator, Errors.COORDINATOR_NOT_AVAILABLE);
-        time.sleep(autoCommitIntervalMs);
-        closeVerifyTimeout(coordinator, 1000, 1000, 1000);
+        try (ConsumerCoordinator coordinator = prepareCoordinatorForCloseTest(true, true, groupInstanceId)) {
+            makeCoordinatorUnknown(coordinator, Errors.COORDINATOR_NOT_AVAILABLE);
+            time.sleep(autoCommitIntervalMs);
+            closeVerifyTimeout(coordinator, 1000, 1000, 1000);
+        }
     }
 
     @Test
     public void testCloseMaxWaitCoordinatorUnavailableForCommit() throws Exception {
-        ConsumerCoordinator coordinator = prepareCoordinatorForCloseTest(true, true, true);
-        makeCoordinatorUnknown(coordinator, Errors.COORDINATOR_NOT_AVAILABLE);
-        time.sleep(autoCommitIntervalMs);
-        closeVerifyTimeout(coordinator, Long.MAX_VALUE, requestTimeoutMs, requestTimeoutMs);
+        try (ConsumerCoordinator coordinator = prepareCoordinatorForCloseTest(true, true, groupInstanceId)) {
+            makeCoordinatorUnknown(coordinator, Errors.COORDINATOR_NOT_AVAILABLE);
+            time.sleep(autoCommitIntervalMs);
+            closeVerifyTimeout(coordinator, Long.MAX_VALUE, requestTimeoutMs, requestTimeoutMs);
+        }
     }
 
     @Test
     public void testCloseNoResponseForCommit() throws Exception {
-        ConsumerCoordinator coordinator = prepareCoordinatorForCloseTest(true, true, true);
-        time.sleep(autoCommitIntervalMs);
-        closeVerifyTimeout(coordinator, Long.MAX_VALUE, requestTimeoutMs, requestTimeoutMs);
+        try (ConsumerCoordinator coordinator = prepareCoordinatorForCloseTest(true, true, groupInstanceId)) {
+            time.sleep(autoCommitIntervalMs);
+            closeVerifyTimeout(coordinator, Long.MAX_VALUE, requestTimeoutMs, requestTimeoutMs);
+        }
     }
 
     @Test
     public void testCloseNoResponseForLeaveGroup() throws Exception {
-        ConsumerCoordinator coordinator = prepareCoordinatorForCloseTest(true, false, true);
-        closeVerifyTimeout(coordinator, Long.MAX_VALUE, requestTimeoutMs, requestTimeoutMs);
+        try (ConsumerCoordinator coordinator = prepareCoordinatorForCloseTest(true, false, Optional.empty())) {
+            closeVerifyTimeout(coordinator, Long.MAX_VALUE, requestTimeoutMs, requestTimeoutMs);
+        }
     }
 
     @Test
     public void testCloseNoWait() throws Exception {
-        ConsumerCoordinator coordinator = prepareCoordinatorForCloseTest(true, true, true);
-        time.sleep(autoCommitIntervalMs);
-        closeVerifyTimeout(coordinator, 0, 0, 0);
+        try (ConsumerCoordinator coordinator = prepareCoordinatorForCloseTest(true, true, groupInstanceId)) {
+            time.sleep(autoCommitIntervalMs);
+            closeVerifyTimeout(coordinator, 0, 0, 0);
+        }
     }
 
     @Test
     public void testHeartbeatThreadClose() throws Exception {
-        ConsumerCoordinator coordinator = prepareCoordinatorForCloseTest(true, true, true);
-        coordinator.ensureActiveGroup();
-        time.sleep(heartbeatIntervalMs + 100);
-        Thread.yield(); // Give heartbeat thread a chance to attempt heartbeat
-        closeVerifyTimeout(coordinator, Long.MAX_VALUE, requestTimeoutMs, requestTimeoutMs);
-        Thread[] threads = new Thread[Thread.activeCount()];
-        int threadCount = Thread.enumerate(threads);
-        for (int i = 0; i < threadCount; i++)
-            assertFalse("Heartbeat thread active after close", threads[i].getName().contains(groupId));
+        try (ConsumerCoordinator coordinator = prepareCoordinatorForCloseTest(true, true, groupInstanceId)) {
+            coordinator.ensureActiveGroup();
+            time.sleep(heartbeatIntervalMs + 100);
+            Thread.yield(); // Give heartbeat thread a chance to attempt heartbeat
+            closeVerifyTimeout(coordinator, Long.MAX_VALUE, requestTimeoutMs, requestTimeoutMs);
+            Thread[] threads = new Thread[Thread.activeCount()];
+            int threadCount = Thread.enumerate(threads);
+            for (int i = 0; i < threadCount; i++) {
+                assertFalse("Heartbeat thread active after close", threads[i].getName().contains(groupId));
+            }
+        }
     }
 
     @Test
     public void testAutoCommitAfterCoordinatorBackToService() {
-        ConsumerCoordinator coordinator = buildCoordinator(new Metrics(), assignors,
-                true, true);
-        subscriptions.assignFromUser(Collections.singleton(t1p));
-        subscriptions.seek(t1p, 100L);
+        try (ConsumerCoordinator coordinator = buildCoordinator(rebalanceConfig, new Metrics(), assignors, true)) {
+            subscriptions.assignFromUser(Collections.singleton(t1p));
+            subscriptions.seek(t1p, 100L);
+
+            coordinator.markCoordinatorUnknown();
+            assertTrue(coordinator.coordinatorUnknown());
+            client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+            prepareOffsetCommitRequest(singletonMap(t1p, 100L), Errors.NONE);
+
+            // async commit offset should find coordinator
+            time.sleep(autoCommitIntervalMs); // sleep for a while to ensure auto commit does happen
+            coordinator.maybeAutoCommitOffsetsAsync(time.milliseconds());
+            assertFalse(coordinator.coordinatorUnknown());
+            assertEquals(100L, subscriptions.position(t1p).offset);
+        }
+    }
 
-        coordinator.markCoordinatorUnknown();
-        assertTrue(coordinator.coordinatorUnknown());
+    @Test(expected = FencedInstanceIdException.class)
+    public void testCommitOffsetRequestSyncWithFencedInstanceIdException() {
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        prepareOffsetCommitRequest(singletonMap(t1p, 100L), Errors.NONE);
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
-        // async commit offset should find coordinator
-        time.sleep(autoCommitIntervalMs); // sleep for a while to ensure auto commit does happen
-        coordinator.maybeAutoCommitOffsetsAsync(time.milliseconds());
-        assertFalse(coordinator.coordinatorUnknown());
-        assertEquals(100L, subscriptions.position(t1p).longValue());
+        // sync commit with invalid partitions should throw if we have no callback
+        prepareOffsetCommitRequest(singletonMap(t1p, 100L), Errors.FENCED_INSTANCE_ID);
+        coordinator.commitOffsetsSync(singletonMap(t1p, new OffsetAndMetadata(100L)), time.timer(Long.MAX_VALUE));
+    }
+
+    @Test(expected = FencedInstanceIdException.class)
+    public void testCommitOffsetRequestAsyncWithFencedInstanceIdException() {
+        receiveFencedInstanceIdException();
+    }
+
+    @Test
+    public void testCommitOffsetRequestAsyncAlwaysReceiveFencedException() {
+        // Once we get fenced exception once, we should always hit fencing case.
+        assertThrows(FencedInstanceIdException.class, this::receiveFencedInstanceIdException);
+        assertThrows(FencedInstanceIdException.class, () ->
+                coordinator.commitOffsetsAsync(singletonMap(t1p, new OffsetAndMetadata(100L)), new MockCommitCallback()));
+        assertThrows(FencedInstanceIdException.class, () ->
+                coordinator.commitOffsetsSync(singletonMap(t1p, new OffsetAndMetadata(100L)), time.timer(Long.MAX_VALUE)));
+    }
+
+    @Test
+    public void testConsumerRejoinAfterRebalance() throws Exception {
+        try (ConsumerCoordinator coordinator = prepareCoordinatorForCloseTest(true, false, Optional.of("group-id"))) {
+            coordinator.ensureActiveGroup();
+
+            prepareOffsetCommitRequest(singletonMap(t1p, 100L), Errors.REBALANCE_IN_PROGRESS);
+
+            assertThrows(CommitFailedException.class, () -> coordinator.commitOffsetsSync(
+                singletonMap(t1p, new OffsetAndMetadata(100L)),
+                time.timer(Long.MAX_VALUE)));
+
+            assertFalse(client.hasPendingResponses());
+            assertFalse(client.hasInFlightRequests());
+
+            int generationId = 42;
+            String memberId = "consumer-42";
+
+            client.prepareResponse(joinGroupFollowerResponse(generationId, memberId, "leader", Errors.NONE));
+
+            MockTime time = new MockTime(1);
+
+            // onJoinPrepare will be executed and onJoinComplete will not.
+            boolean res = coordinator.joinGroupIfNeeded(time.timer(2));
+
+            assertFalse(res);
+            assertFalse(client.hasPendingResponses());
+            // SynGroupRequest not responded.
+            assertEquals(1, client.inFlightRequestCount());
+            assertEquals(generationId, coordinator.generation().generationId);
+            assertEquals(memberId, coordinator.generation().memberId);
+
+            // Imitating heartbeat thread that clears generation data.
+            coordinator.maybeLeaveGroup("Clear generation data.");
+
+            assertEquals(AbstractCoordinator.Generation.NO_GENERATION, coordinator.generation());
+
+            client.respond(syncGroupResponse(singletonList(t1p), Errors.NONE));
+
+            // Join future should succeed but generation already cleared so result of join is false.
+            res = coordinator.joinGroupIfNeeded(time.timer(1));
+
+            assertFalse(res);
+            assertFalse(client.hasPendingResponses());
+            assertFalse(client.hasInFlightRequests());
+
+            // Retry join should then succeed
+            client.prepareResponse(joinGroupFollowerResponse(generationId, memberId, "leader", Errors.NONE));
+            client.prepareResponse(syncGroupResponse(singletonList(t1p), Errors.NONE));
+
+            res = coordinator.joinGroupIfNeeded(time.timer(2));
+
+            assertTrue(res);
+            assertFalse(client.hasPendingResponses());
+            assertFalse(client.hasInFlightRequests());
+        }
+    }
+
+    private void receiveFencedInstanceIdException() {
+        subscriptions.assignFromUser(singleton(t1p));
+
+        client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
+
+        prepareOffsetCommitRequest(singletonMap(t1p, 100L), Errors.FENCED_INSTANCE_ID);
+
+        coordinator.commitOffsetsAsync(singletonMap(t1p, new OffsetAndMetadata(100L)), new MockCommitCallback());
+        coordinator.invokeCompletedOffsetCommitCallbacks();
     }
 
     private ConsumerCoordinator prepareCoordinatorForCloseTest(final boolean useGroupManagement,
                                                                final boolean autoCommit,
-                                                               final boolean leaveGroup) {
+                                                               final Optional groupInstanceId) {
         final String consumerId = "consumer";
-        ConsumerCoordinator coordinator = buildCoordinator(new Metrics(), assignors,
-                autoCommit, leaveGroup);
+        rebalanceConfig = buildRebalanceConfig(groupInstanceId);
+        ConsumerCoordinator coordinator = buildCoordinator(rebalanceConfig,
+                                                           new Metrics(),
+                                                           assignors,
+                                                           autoCommit);
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
         coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
         if (useGroupManagement) {
@@ -2057,8 +2518,9 @@ private ConsumerCoordinator prepareCoordinatorForCloseTest(final boolean useGrou
             client.prepareResponse(joinGroupFollowerResponse(1, consumerId, "leader", Errors.NONE));
             client.prepareResponse(syncGroupResponse(singletonList(t1p), Errors.NONE));
             coordinator.joinGroupIfNeeded(time.timer(Long.MAX_VALUE));
-        } else
+        } else {
             subscriptions.assignFromUser(singleton(t1p));
+        }
 
         subscriptions.seek(t1p, 100);
         coordinator.poll(time.timer(Long.MAX_VALUE));
@@ -2116,81 +2578,126 @@ public void run() {
     private void gracefulCloseTest(ConsumerCoordinator coordinator, boolean shouldLeaveGroup) throws Exception {
         final AtomicBoolean commitRequested = new AtomicBoolean();
         final AtomicBoolean leaveGroupRequested = new AtomicBoolean();
-        client.prepareResponse(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                commitRequested.set(true);
-                OffsetCommitRequest commitRequest = (OffsetCommitRequest) body;
-                return commitRequest.groupId().equals(groupId);
-            }
-        }, new OffsetCommitResponse(new HashMap()));
-        client.prepareResponse(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                leaveGroupRequested.set(true);
-                LeaveGroupRequest leaveRequest = (LeaveGroupRequest) body;
-                return leaveRequest.data().groupId().equals(groupId);
-            }
-        }, new LeaveGroupResponse(new LeaveGroupResponseData().setErrorCode(Errors.NONE.code())));
+        client.prepareResponse(body -> {
+            commitRequested.set(true);
+            OffsetCommitRequest commitRequest = (OffsetCommitRequest) body;
+            return commitRequest.data().groupId().equals(groupId);
+        }, new OffsetCommitResponse(new OffsetCommitResponseData()));
+        client.prepareResponse(body -> {
+            leaveGroupRequested.set(true);
+            LeaveGroupRequest leaveRequest = (LeaveGroupRequest) body;
+            return leaveRequest.data().groupId().equals(groupId);
+        }, new LeaveGroupResponse(new LeaveGroupResponseData()
+                .setErrorCode(Errors.NONE.code())));
 
         coordinator.close();
         assertTrue("Commit not requested", commitRequested.get());
         assertEquals("leaveGroupRequested should be " + shouldLeaveGroup, shouldLeaveGroup, leaveGroupRequested.get());
+
+        if (shouldLeaveGroup) {
+            assertEquals(1, rebalanceListener.revokedCount);
+            assertEquals(singleton(t1p), rebalanceListener.revoked);
+        }
     }
 
-    private ConsumerCoordinator buildCoordinator(final Metrics metrics,
-                                                 final List assignors,
-                                                 final boolean autoCommitEnabled,
-                                                 final boolean leaveGroup) {
+    private ConsumerCoordinator buildCoordinator(final GroupRebalanceConfig rebalanceConfig,
+                                                 final Metrics metrics,
+                                                 final List assignors,
+                                                 final boolean autoCommitEnabled) {
         return new ConsumerCoordinator(
+                rebalanceConfig,
                 new LogContext(),
                 consumerClient,
-                groupId,
-                rebalanceTimeoutMs,
-                sessionTimeoutMs,
-                heartbeat,
                 assignors,
                 metadata,
                 subscriptions,
                 metrics,
                 "consumer" + groupId,
                 time,
-                retryBackoffMs,
                 autoCommitEnabled,
                 autoCommitIntervalMs,
-                null,
-                leaveGroup);
+                null);
+    }
+
+    private Collection getRevoked(final List owned,
+                                                  final List assigned) {
+        switch (protocol) {
+            case EAGER:
+                return toSet(owned);
+            case COOPERATIVE:
+                final List revoked = new ArrayList<>(owned);
+                revoked.removeAll(assigned);
+                return toSet(revoked);
+            default:
+                throw new IllegalStateException("This should not happen");
+        }
+    }
+
+    private Collection getAdded(final List owned,
+                                                final List assigned) {
+        switch (protocol) {
+            case EAGER:
+                return toSet(assigned);
+            case COOPERATIVE:
+                final List added = new ArrayList<>(assigned);
+                added.removeAll(owned);
+                return toSet(added);
+            default:
+                throw new IllegalStateException("This should not happen");
+        }
     }
 
     private FindCoordinatorResponse groupCoordinatorResponse(Node node, Errors error) {
-        return new FindCoordinatorResponse(error, node);
+        return FindCoordinatorResponse.prepareResponse(error, node);
     }
 
     private HeartbeatResponse heartbeatResponse(Errors error) {
-        return new HeartbeatResponse(error);
+        return new HeartbeatResponse(new HeartbeatResponseData().setErrorCode(error.code()));
     }
 
     private JoinGroupResponse joinGroupLeaderResponse(int generationId,
                                                       String memberId,
                                                       Map> subscriptions,
                                                       Errors error) {
-        Map metadata = new HashMap<>();
+        List metadata = new ArrayList<>();
         for (Map.Entry> subscriptionEntry : subscriptions.entrySet()) {
-            PartitionAssignor.Subscription subscription = new PartitionAssignor.Subscription(subscriptionEntry.getValue());
+            ConsumerPartitionAssignor.Subscription subscription = new ConsumerPartitionAssignor.Subscription(subscriptionEntry.getValue());
             ByteBuffer buf = ConsumerProtocol.serializeSubscription(subscription);
-            metadata.put(subscriptionEntry.getKey(), buf);
+            metadata.add(new JoinGroupResponseData.JoinGroupResponseMember()
+                    .setMemberId(subscriptionEntry.getKey())
+                    .setMetadata(buf.array()));
         }
-        return new JoinGroupResponse(error, generationId, partitionAssignor.name(), memberId, memberId, metadata);
+
+        return new JoinGroupResponse(
+                new JoinGroupResponseData()
+                        .setErrorCode(error.code())
+                        .setGenerationId(generationId)
+                        .setProtocolName(partitionAssignor.name())
+                        .setLeader(memberId)
+                        .setMemberId(memberId)
+                        .setMembers(metadata)
+        );
     }
 
     private JoinGroupResponse joinGroupFollowerResponse(int generationId, String memberId, String leaderId, Errors error) {
-        return new JoinGroupResponse(error, generationId, partitionAssignor.name(), memberId, leaderId,
-                Collections.emptyMap());
+        return new JoinGroupResponse(
+                new JoinGroupResponseData()
+                        .setErrorCode(error.code())
+                        .setGenerationId(generationId)
+                        .setProtocolName(partitionAssignor.name())
+                        .setLeader(leaderId)
+                        .setMemberId(memberId)
+                        .setMembers(Collections.emptyList())
+        );
     }
 
     private SyncGroupResponse syncGroupResponse(List partitions, Errors error) {
-        ByteBuffer buf = ConsumerProtocol.serializeAssignment(new PartitionAssignor.Assignment(partitions));
-        return new SyncGroupResponse(error, buf);
+        ByteBuffer buf = ConsumerProtocol.serializeAssignment(new ConsumerPartitionAssignor.Assignment(partitions));
+        return new SyncGroupResponse(
+                new SyncGroupResponseData()
+                        .setErrorCode(error.code())
+                        .setAssignment(Utils.toArray(buf))
+        );
     }
 
     private OffsetCommitResponse offsetCommitResponse(Map responseData) {
@@ -2207,6 +2714,12 @@ private OffsetFetchResponse offsetFetchResponse(TopicPartition tp, Errors partit
         return new OffsetFetchResponse(Errors.NONE, singletonMap(tp, data));
     }
 
+    private OffsetFetchResponse offsetFetchResponse(TopicPartition tp, Errors partitionLevelError, String metadata, long offset, int epoch) {
+        OffsetFetchResponse.PartitionData data = new OffsetFetchResponse.PartitionData(offset,
+                Optional.of(epoch), metadata, partitionLevelError);
+        return new OffsetFetchResponse(Errors.NONE, singletonMap(tp, data));
+    }
+
     private OffsetCommitCallback callback(final AtomicBoolean success) {
         return new OffsetCommitCallback() {
             @Override
@@ -2242,6 +2755,19 @@ private void prepareOffsetCommitRequest(final Map expected
         client.prepareResponse(offsetCommitRequestMatcher(expectedOffsets), offsetCommitResponse(errors), disconnected);
     }
 
+    private void prepareJoinAndSyncResponse(String consumerId, int generation, List subscription, List assignment) {
+        partitionAssignor.prepare(singletonMap(consumerId, assignment));
+        client.prepareResponse(
+                joinGroupLeaderResponse(
+                        generation, consumerId, singletonMap(consumerId, subscription), Errors.NONE));
+        client.prepareResponse(body -> {
+            SyncGroupRequest sync = (SyncGroupRequest) body;
+            return sync.data.memberId().equals(consumerId) &&
+                    sync.data.generationId() == generation &&
+                    sync.groupAssignments().containsKey(consumerId);
+        }, syncGroupResponse(assignment, Errors.NONE));
+    }
+
     private Map partitionErrors(Collection partitions, Errors error) {
         final Map errors = new HashMap<>();
         for (TopicPartition partition : partitions) {
@@ -2256,36 +2782,31 @@ private void respondToOffsetCommitRequest(final Map expect
     }
 
     private MockClient.RequestMatcher offsetCommitRequestMatcher(final Map expectedOffsets) {
-        return new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                OffsetCommitRequest req = (OffsetCommitRequest) body;
-                Map offsets = req.offsetData();
-                if (offsets.size() != expectedOffsets.size())
-                    return false;
-
-                for (Map.Entry expectedOffset : expectedOffsets.entrySet()) {
-                    if (!offsets.containsKey(expectedOffset.getKey()))
-                        return false;
+        return body -> {
+            OffsetCommitRequest req = (OffsetCommitRequest) body;
+            Map offsets = req.offsets();
+            if (offsets.size() != expectedOffsets.size())
+                return false;
 
-                    OffsetCommitRequest.PartitionData offsetCommitData = offsets.get(expectedOffset.getKey());
-                    if (offsetCommitData.offset != expectedOffset.getValue())
+            for (Map.Entry expectedOffset : expectedOffsets.entrySet()) {
+                if (!offsets.containsKey(expectedOffset.getKey())) {
+                    return false;
+                } else {
+                    Long actualOffset = offsets.get(expectedOffset.getKey());
+                    if (!actualOffset.equals(expectedOffset.getValue())) {
                         return false;
+                    }
                 }
-
-                return true;
             }
+            return true;
         };
     }
 
     private OffsetCommitCallback callback(final Map expectedOffsets,
                                           final AtomicBoolean success) {
-        return new OffsetCommitCallback() {
-            @Override
-            public void onComplete(Map offsets, Exception exception) {
-                if (expectedOffsets.equals(offsets) && exception == null)
-                    success.set(true);
-            }
+        return (offsets, exception) -> {
+            if (expectedOffsets.equals(offsets) && exception == null)
+                success.set(true);
         };
     }
 
@@ -2301,12 +2822,13 @@ public void onComplete(Map offsets, Exception
     }
 
     private static class MockRebalanceListener implements ConsumerRebalanceListener {
+        public Collection lost;
         public Collection revoked;
         public Collection assigned;
+        public int lostCount = 0;
         public int revokedCount = 0;
         public int assignedCount = 0;
 
-
         @Override
         public void onPartitionsAssigned(Collection partitions) {
             this.assigned = partitions;
@@ -2319,5 +2841,10 @@ public void onPartitionsRevoked(Collection partitions) {
             revokedCount++;
         }
 
+        @Override
+        public void onPartitionsLost(Collection partitions) {
+            this.lost = partitions;
+            lostCount++;
+        }
     }
 }
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerMetadataTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerMetadataTest.java
index 871ef30c4ff27..b373192ca7066 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerMetadataTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerMetadataTest.java
@@ -75,7 +75,8 @@ private void testPatternSubscription(boolean includeInternalTopics) {
         topics.add(topicMetadata("__matching_topic", false));
         topics.add(topicMetadata("non_matching_topic", false));
 
-        MetadataResponse response = new MetadataResponse(singletonList(node), "clusterId", node.id(), topics);
+        MetadataResponse response = MetadataResponse.prepareResponse(singletonList(node),
+            "clusterId", node.id(), topics);
         metadata.update(response, time.milliseconds());
 
         if (includeInternalTopics)
@@ -91,13 +92,21 @@ public void testUserAssignment() {
                 new TopicPartition("bar", 0),
                 new TopicPartition("__consumer_offsets", 0)));
         testBasicSubscription(Utils.mkSet("foo", "bar"), Utils.mkSet("__consumer_offsets"));
+
+        subscription.assignFromUser(Utils.mkSet(
+                new TopicPartition("baz", 0),
+                new TopicPartition("__consumer_offsets", 0)));
+        testBasicSubscription(Utils.mkSet("baz"), Utils.mkSet("__consumer_offsets"));
     }
 
     @Test
     public void testNormalSubscription() {
         subscription.subscribe(Utils.mkSet("foo", "bar", "__consumer_offsets"), new NoOpConsumerRebalanceListener());
-        subscription.groupSubscribe(Utils.mkSet("baz"));
+        subscription.groupSubscribe(Utils.mkSet("baz", "foo", "bar", "__consumer_offsets"));
         testBasicSubscription(Utils.mkSet("foo", "bar", "baz"), Utils.mkSet("__consumer_offsets"));
+
+        subscription.resetGroupSubscription();
+        testBasicSubscription(Utils.mkSet("foo", "bar"), Utils.mkSet("__consumer_offsets"));
     }
 
     @Test
@@ -142,7 +151,8 @@ private void testBasicSubscription(Set expectedTopics, Set expec
         for (String expectedInternalTopic : expectedInternalTopics)
             topics.add(topicMetadata(expectedInternalTopic, true));
 
-        MetadataResponse response = new MetadataResponse(singletonList(node), "clusterId", node.id(), topics);
+        MetadataResponse response = MetadataResponse.prepareResponse(singletonList(node),
+            "clusterId", node.id(), topics);
         metadata.update(response, time.milliseconds());
 
         assertEquals(allTopics, metadata.fetch().topics());
@@ -157,8 +167,8 @@ private MetadataResponse.TopicMetadata topicMetadata(String topic, boolean isInt
     private ConsumerMetadata newConsumerMetadata(boolean includeInternalTopics) {
         long refreshBackoffMs = 50;
         long expireMs = 50000;
-        return new ConsumerMetadata(refreshBackoffMs, expireMs, includeInternalTopics, subscription, new LogContext(),
-                new ClusterResourceListeners());
+        return new ConsumerMetadata(refreshBackoffMs, expireMs, includeInternalTopics, false,
+                subscription, new LogContext(), new ClusterResourceListeners());
     }
 
 }
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 14c2cba72a985..c50bce45dcfe0 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
@@ -21,6 +21,7 @@
 import org.apache.kafka.clients.MockClient;
 import org.apache.kafka.clients.NetworkClient;
 import org.apache.kafka.common.Cluster;
+import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.Node;
 import org.apache.kafka.common.errors.AuthenticationException;
 import org.apache.kafka.common.errors.DisconnectException;
@@ -29,6 +30,8 @@
 import org.apache.kafka.common.errors.TopicAuthorizationException;
 import org.apache.kafka.common.errors.WakeupException;
 import org.apache.kafka.common.internals.ClusterResourceListeners;
+import org.apache.kafka.common.message.HeartbeatRequestData;
+import org.apache.kafka.common.message.HeartbeatResponseData;
 import org.apache.kafka.common.protocol.Errors;
 import org.apache.kafka.common.requests.HeartbeatRequest;
 import org.apache.kafka.common.requests.HeartbeatResponse;
@@ -41,10 +44,10 @@
 import java.time.Duration;
 import java.util.Collections;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 import static org.mockito.ArgumentMatchers.anyLong;
@@ -232,13 +235,13 @@ public void run() {
 
     @Test
     public void testAuthenticationExceptionPropagatedFromMetadata() {
-        metadata.failedUpdate(time.milliseconds(), new AuthenticationException("Authentication failed"));
+        metadata.fatalError(new AuthenticationException("Authentication failed"));
         try {
             consumerClient.poll(time.timer(Duration.ZERO));
             fail("Expected authentication error thrown");
         } catch (AuthenticationException e) {
             // After the exception is raised, it should have been cleared
-            assertNull(metadata.getAndClearAuthenticationException());
+            metadata.maybeThrowAnyException();
         }
     }
 
@@ -258,6 +261,18 @@ public void testTopicAuthorizationExceptionPropagatedFromMetadata() {
         consumerClient.poll(time.timer(Duration.ZERO));
     }
 
+    @Test
+    public void testMetadataFailurePropagated() {
+        KafkaException metadataException = new KafkaException();
+        metadata.fatalError(metadataException);
+        try {
+            consumerClient.poll(time.timer(Duration.ZERO));
+            fail("Expected poll to throw exception");
+        } catch (Exception e) {
+            assertEquals(metadataException, e);
+        }
+    }
+
     @Test
     public void testFutureCompletionOutsidePoll() throws Exception {
         // Tests the scenario in which the request that is being awaited in one thread
@@ -364,12 +379,49 @@ public boolean connectionFailed(Node node) {
         assertEquals(0, consumerClient.pendingRequestCount(node));
     }
 
+    @Test
+    public void testTrySend() {
+        final AtomicBoolean isReady = new AtomicBoolean();
+        final AtomicInteger checkCount = new AtomicInteger();
+        client = new MockClient(time, metadata) {
+            @Override
+            public boolean ready(Node node, long now) {
+                checkCount.incrementAndGet();
+                if (isReady.get())
+                    return super.ready(node, now);
+                else
+                    return false;
+            }
+        };
+        consumerClient = new ConsumerNetworkClient(new LogContext(), client, metadata, time, 100, 10, Integer.MAX_VALUE);
+        consumerClient.send(node, heartbeat());
+        consumerClient.send(node, heartbeat());
+        assertEquals(2, consumerClient.pendingRequestCount(node));
+        assertEquals(0, client.inFlightRequestCount(node.idString()));
+
+        consumerClient.trySend(time.milliseconds());
+        // only check one time when the node doesn't ready
+        assertEquals(1, checkCount.getAndSet(0));
+        assertEquals(2, consumerClient.pendingRequestCount(node));
+        assertEquals(0, client.inFlightRequestCount(node.idString()));
+
+        isReady.set(true);
+        consumerClient.trySend(time.milliseconds());
+        // check node ready or not for every request
+        assertEquals(2, checkCount.getAndSet(0));
+        assertEquals(2, consumerClient.pendingRequestCount(node));
+        assertEquals(2, client.inFlightRequestCount(node.idString()));
+    }
+
     private HeartbeatRequest.Builder heartbeat() {
-        return new HeartbeatRequest.Builder("group", 1, "memberId");
+        return new HeartbeatRequest.Builder(new HeartbeatRequestData()
+                .setGroupId("group")
+                .setGenerationId(1)
+                .setMemberId("memberId"));
     }
 
     private HeartbeatResponse heartbeatResponse(Errors error) {
-        return new HeartbeatResponse(error);
+        return new HeartbeatResponse(new HeartbeatResponseData().setErrorCode(error.code()));
     }
 
 }
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerProtocolTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerProtocolTest.java
index 37d105cf1cc3a..3cf7be51665c4 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerProtocolTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerProtocolTest.java
@@ -16,7 +16,8 @@
  */
 package org.apache.kafka.clients.consumer.internals;
 
-import org.apache.kafka.clients.consumer.internals.PartitionAssignor.Subscription;
+import org.apache.kafka.clients.consumer.ConsumerPartitionAssignor.Assignment;
+import org.apache.kafka.clients.consumer.ConsumerPartitionAssignor.Subscription;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.protocol.types.ArrayOf;
 import org.apache.kafka.common.protocol.types.Field;
@@ -27,22 +28,49 @@
 
 import java.nio.ByteBuffer;
 import java.util.Arrays;
-import java.util.Collection;
-import java.util.HashSet;
+import java.util.Collections;
 import java.util.List;
-import java.util.Set;
-
+import java.util.Optional;
+
+import static org.apache.kafka.clients.consumer.internals.ConsumerProtocol.CONSUMER_PROTOCOL_HEADER_SCHEMA;
+import static org.apache.kafka.clients.consumer.internals.ConsumerProtocol.OWNED_PARTITIONS_KEY_NAME;
+import static org.apache.kafka.clients.consumer.internals.ConsumerProtocol.TOPICS_KEY_NAME;
+import static org.apache.kafka.clients.consumer.internals.ConsumerProtocol.TOPIC_ASSIGNMENT_V0;
+import static org.apache.kafka.clients.consumer.internals.ConsumerProtocol.TOPIC_PARTITIONS_KEY_NAME;
+import static org.apache.kafka.clients.consumer.internals.ConsumerProtocol.USER_DATA_KEY_NAME;
+import static org.apache.kafka.clients.consumer.internals.ConsumerProtocol.VERSION_KEY_NAME;
+import static org.apache.kafka.test.TestUtils.toSet;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
 
 public class ConsumerProtocolTest {
 
+    private final TopicPartition tp1 = new TopicPartition("foo", 1);
+    private final TopicPartition tp2 = new TopicPartition("bar", 2);
+    private final Optional groupInstanceId = Optional.of("instance.id");
+
     @Test
     public void serializeDeserializeMetadata() {
-        Subscription subscription = new Subscription(Arrays.asList("foo", "bar"));
+        Subscription subscription = new Subscription(Arrays.asList("foo", "bar"), ByteBuffer.wrap(new byte[0]));
         ByteBuffer buffer = ConsumerProtocol.serializeSubscription(subscription);
         Subscription parsedSubscription = ConsumerProtocol.deserializeSubscription(buffer);
         assertEquals(subscription.topics(), parsedSubscription.topics());
+        assertEquals(0, parsedSubscription.userData().limit());
+        assertFalse(parsedSubscription.groupInstanceId().isPresent());
+    }
+
+    @Test
+    public void serializeDeserializeMetadataAndGroupInstanceId() {
+        Subscription subscription = new Subscription(Arrays.asList("foo", "bar"), ByteBuffer.wrap(new byte[0]));
+        ByteBuffer buffer = ConsumerProtocol.serializeSubscription(subscription);
+
+        Subscription parsedSubscription = ConsumerProtocol.deserializeSubscription(buffer);
+        parsedSubscription.setGroupInstanceId(groupInstanceId);
+        assertEquals(subscription.topics(), parsedSubscription.topics());
+        assertEquals(0, parsedSubscription.userData().limit());
+        assertEquals(groupInstanceId, parsedSubscription.groupInstanceId());
     }
 
     @Test
@@ -51,26 +79,54 @@ public void serializeDeserializeNullSubscriptionUserData() {
         ByteBuffer buffer = ConsumerProtocol.serializeSubscription(subscription);
         Subscription parsedSubscription = ConsumerProtocol.deserializeSubscription(buffer);
         assertEquals(subscription.topics(), parsedSubscription.topics());
-        assertNull(subscription.userData());
+        assertNull(parsedSubscription.userData());
     }
 
     @Test
-    public void deserializeNewSubscriptionVersion() {
+    public void deserializeOldSubscriptionVersion() {
+        Subscription subscription = new Subscription(Arrays.asList("foo", "bar"), null);
+        ByteBuffer buffer = ConsumerProtocol.serializeSubscriptionV0(subscription);
+        Subscription parsedSubscription = ConsumerProtocol.deserializeSubscription(buffer);
+        assertEquals(parsedSubscription.topics(), parsedSubscription.topics());
+        assertNull(parsedSubscription.userData());
+        assertTrue(parsedSubscription.ownedPartitions().isEmpty());
+    }
+
+    @Test
+    public void deserializeNewSubscriptionWithOldVersion() {
+        Subscription subscription = new Subscription(Arrays.asList("foo", "bar"), null, Collections.singletonList(tp2));
+        ByteBuffer buffer = ConsumerProtocol.serializeSubscription(subscription);
+        // ignore the version assuming it is the old byte code, as it will blindly deserialize as V0
+        Struct header = CONSUMER_PROTOCOL_HEADER_SCHEMA.read(buffer);
+        header.getShort(VERSION_KEY_NAME);
+        Subscription parsedSubscription = ConsumerProtocol.deserializeSubscriptionV0(buffer);
+        assertEquals(subscription.topics(), parsedSubscription.topics());
+        assertNull(parsedSubscription.userData());
+        assertTrue(parsedSubscription.ownedPartitions().isEmpty());
+        assertFalse(parsedSubscription.groupInstanceId().isPresent());
+    }
+
+    @Test
+    public void deserializeFutureSubscriptionVersion() {
         // verify that a new version which adds a field is still parseable
         short version = 100;
 
         Schema subscriptionSchemaV100 = new Schema(
-                new Field(ConsumerProtocol.TOPICS_KEY_NAME, new ArrayOf(Type.STRING)),
-                new Field(ConsumerProtocol.USER_DATA_KEY_NAME, Type.BYTES),
+                new Field(TOPICS_KEY_NAME, new ArrayOf(Type.STRING)),
+                new Field(USER_DATA_KEY_NAME, Type.BYTES),
+                new Field(OWNED_PARTITIONS_KEY_NAME, new ArrayOf(TOPIC_ASSIGNMENT_V0)),
                 new Field("foo", Type.STRING));
 
         Struct subscriptionV100 = new Struct(subscriptionSchemaV100);
-        subscriptionV100.set(ConsumerProtocol.TOPICS_KEY_NAME, new Object[]{"topic"});
-        subscriptionV100.set(ConsumerProtocol.USER_DATA_KEY_NAME, ByteBuffer.wrap(new byte[0]));
+        subscriptionV100.set(TOPICS_KEY_NAME, new Object[]{"topic"});
+        subscriptionV100.set(USER_DATA_KEY_NAME, ByteBuffer.wrap(new byte[0]));
+        subscriptionV100.set(OWNED_PARTITIONS_KEY_NAME, new Object[]{new Struct(TOPIC_ASSIGNMENT_V0)
+            .set(ConsumerProtocol.TOPIC_KEY_NAME, tp2.topic())
+            .set(ConsumerProtocol.PARTITIONS_KEY_NAME, new Object[]{tp2.partition()})});
         subscriptionV100.set("foo", "bar");
 
-        Struct headerV100 = new Struct(ConsumerProtocol.CONSUMER_PROTOCOL_HEADER_SCHEMA);
-        headerV100.set(ConsumerProtocol.VERSION_KEY_NAME, version);
+        Struct headerV100 = new Struct(CONSUMER_PROTOCOL_HEADER_SCHEMA);
+        headerV100.set(VERSION_KEY_NAME, version);
 
         ByteBuffer buffer = ByteBuffer.allocate(subscriptionV100.sizeOf() + headerV100.sizeOf());
         headerV100.writeTo(buffer);
@@ -79,46 +135,50 @@ public void deserializeNewSubscriptionVersion() {
         buffer.flip();
 
         Subscription subscription = ConsumerProtocol.deserializeSubscription(buffer);
-        assertEquals(Arrays.asList("topic"), subscription.topics());
+        subscription.setGroupInstanceId(groupInstanceId);
+        assertEquals(Collections.singletonList("topic"), subscription.topics());
+        assertEquals(Collections.singletonList(tp2), subscription.ownedPartitions());
+        assertEquals(groupInstanceId, subscription.groupInstanceId());
     }
 
     @Test
     public void serializeDeserializeAssignment() {
-        List partitions = Arrays.asList(new TopicPartition("foo", 0), new TopicPartition("bar", 2));
-        ByteBuffer buffer = ConsumerProtocol.serializeAssignment(new PartitionAssignor.Assignment(partitions));
-        PartitionAssignor.Assignment parsedAssignment = ConsumerProtocol.deserializeAssignment(buffer);
+        List partitions = Arrays.asList(tp1, tp2);
+        ByteBuffer buffer = ConsumerProtocol.serializeAssignment(new Assignment(partitions, ByteBuffer.wrap(new byte[0])));
+        Assignment parsedAssignment = ConsumerProtocol.deserializeAssignment(buffer);
         assertEquals(toSet(partitions), toSet(parsedAssignment.partitions()));
+        assertEquals(0, parsedAssignment.userData().limit());
     }
 
     @Test
     public void deserializeNullAssignmentUserData() {
-        List partitions = Arrays.asList(new TopicPartition("foo", 0), new TopicPartition("bar", 2));
-        ByteBuffer buffer = ConsumerProtocol.serializeAssignment(new PartitionAssignor.Assignment(partitions, null));
-        PartitionAssignor.Assignment parsedAssignment = ConsumerProtocol.deserializeAssignment(buffer);
+        List partitions = Arrays.asList(tp1, tp2);
+        ByteBuffer buffer = ConsumerProtocol.serializeAssignment(new Assignment(partitions, null));
+        Assignment parsedAssignment = ConsumerProtocol.deserializeAssignment(buffer);
         assertEquals(toSet(partitions), toSet(parsedAssignment.partitions()));
         assertNull(parsedAssignment.userData());
     }
 
     @Test
-    public void deserializeNewAssignmentVersion() {
+    public void deserializeFutureAssignmentVersion() {
         // verify that a new version which adds a field is still parseable
         short version = 100;
 
         Schema assignmentSchemaV100 = new Schema(
-                new Field(ConsumerProtocol.TOPIC_PARTITIONS_KEY_NAME, new ArrayOf(ConsumerProtocol.TOPIC_ASSIGNMENT_V0)),
-                new Field(ConsumerProtocol.USER_DATA_KEY_NAME, Type.BYTES),
+                new Field(TOPIC_PARTITIONS_KEY_NAME, new ArrayOf(TOPIC_ASSIGNMENT_V0)),
+                new Field(USER_DATA_KEY_NAME, Type.BYTES),
                 new Field("foo", Type.STRING));
 
         Struct assignmentV100 = new Struct(assignmentSchemaV100);
-        assignmentV100.set(ConsumerProtocol.TOPIC_PARTITIONS_KEY_NAME,
-                new Object[]{new Struct(ConsumerProtocol.TOPIC_ASSIGNMENT_V0)
-                        .set(ConsumerProtocol.TOPIC_KEY_NAME, "foo")
-                        .set(ConsumerProtocol.PARTITIONS_KEY_NAME, new Object[]{1})});
-        assignmentV100.set(ConsumerProtocol.USER_DATA_KEY_NAME, ByteBuffer.wrap(new byte[0]));
+        assignmentV100.set(TOPIC_PARTITIONS_KEY_NAME,
+                new Object[]{new Struct(TOPIC_ASSIGNMENT_V0)
+                        .set(ConsumerProtocol.TOPIC_KEY_NAME, tp1.topic())
+                        .set(ConsumerProtocol.PARTITIONS_KEY_NAME, new Object[]{tp1.partition()})});
+        assignmentV100.set(USER_DATA_KEY_NAME, ByteBuffer.wrap(new byte[0]));
         assignmentV100.set("foo", "bar");
 
-        Struct headerV100 = new Struct(ConsumerProtocol.CONSUMER_PROTOCOL_HEADER_SCHEMA);
-        headerV100.set(ConsumerProtocol.VERSION_KEY_NAME, version);
+        Struct headerV100 = new Struct(CONSUMER_PROTOCOL_HEADER_SCHEMA);
+        headerV100.set(VERSION_KEY_NAME, version);
 
         ByteBuffer buffer = ByteBuffer.allocate(assignmentV100.sizeOf() + headerV100.sizeOf());
         headerV100.writeTo(buffer);
@@ -126,11 +186,7 @@ public void deserializeNewAssignmentVersion() {
 
         buffer.flip();
 
-        PartitionAssignor.Assignment assignment = ConsumerProtocol.deserializeAssignment(buffer);
-        assertEquals(toSet(Arrays.asList(new TopicPartition("foo", 1))), toSet(assignment.partitions()));
-    }
-
-    private static  Set toSet(Collection collection) {
-        return new HashSet<>(collection);
+        Assignment assignment = ConsumerProtocol.deserializeAssignment(buffer);
+        assertEquals(toSet(Collections.singletonList(tp1)), toSet(assignment.partitions()));
     }
 }
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/FetcherTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/FetcherTest.java
index 3fe7ca05c0c67..27c198f705ef7 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/FetcherTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/FetcherTest.java
@@ -21,6 +21,7 @@
 import org.apache.kafka.clients.ClientRequest;
 import org.apache.kafka.clients.ClientUtils;
 import org.apache.kafka.clients.FetchSessionHandler;
+import org.apache.kafka.clients.Metadata;
 import org.apache.kafka.clients.MockClient;
 import org.apache.kafka.clients.NetworkClient;
 import org.apache.kafka.clients.NodeApiVersions;
@@ -64,8 +65,8 @@
 import org.apache.kafka.common.record.Records;
 import org.apache.kafka.common.record.SimpleRecord;
 import org.apache.kafka.common.record.TimestampType;
-import org.apache.kafka.common.requests.AbstractRequest;
 import org.apache.kafka.common.requests.ApiVersionsResponse;
+import org.apache.kafka.common.requests.EpochEndOffset;
 import org.apache.kafka.common.requests.FetchRequest;
 import org.apache.kafka.common.requests.FetchResponse;
 import org.apache.kafka.common.requests.IsolationLevel;
@@ -73,8 +74,10 @@
 import org.apache.kafka.common.requests.ListOffsetResponse;
 import org.apache.kafka.common.requests.MetadataRequest;
 import org.apache.kafka.common.requests.MetadataResponse;
-import org.apache.kafka.common.requests.ResponseHeader;
+import org.apache.kafka.common.requests.OffsetsForLeaderEpochRequest;
+import org.apache.kafka.common.requests.OffsetsForLeaderEpochResponse;
 import org.apache.kafka.common.serialization.ByteArrayDeserializer;
+import org.apache.kafka.common.serialization.BytesDeserializer;
 import org.apache.kafka.common.serialization.Deserializer;
 import org.apache.kafka.common.serialization.StringDeserializer;
 import org.apache.kafka.common.utils.ByteBufferOutputStream;
@@ -93,6 +96,7 @@
 import java.lang.reflect.Field;
 import java.nio.ByteBuffer;
 import java.nio.charset.StandardCharsets;
+import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
@@ -104,17 +108,21 @@
 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.Future;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 
+import static java.util.Arrays.asList;
 import static java.util.Collections.singleton;
 import static java.util.Collections.singletonMap;
 import static org.apache.kafka.common.requests.FetchMetadata.INVALID_SESSION_ID;
+import static org.apache.kafka.test.TestUtils.assertOptional;
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
@@ -124,6 +132,8 @@
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.spy;
 
 @SuppressWarnings("deprecation")
 public class FetcherTest {
@@ -151,9 +161,12 @@ public class FetcherTest {
     private FetcherMetricsRegistry metricsRegistry;
     private MockClient client;
     private Metrics metrics;
+    private ApiVersions apiVersions = new ApiVersions();
     private ConsumerNetworkClient consumerClient;
     private Fetcher fetcher;
 
+    private boolean testPassthrough = false;
+
     private MemoryRecords records;
     private MemoryRecords nextRecords;
     private MemoryRecords emptyRecords;
@@ -206,7 +219,7 @@ public void testFetchNormal() {
 
         List> records = partitionRecords.get(tp0);
         assertEquals(3, records.size());
-        assertEquals(4L, subscriptions.position(tp0).longValue()); // this is the next fetching position
+        assertEquals(4L, subscriptions.position(tp0).offset); // this is the next fetching position
         long offset = 1;
         for (ConsumerRecord record : records) {
             assertEquals(offset, record.offset());
@@ -370,7 +383,7 @@ public void testFetcherIgnoresControlRecords() {
 
         List> records = partitionRecords.get(tp0);
         assertEquals(1, records.size());
-        assertEquals(2L, subscriptions.position(tp0).longValue());
+        assertEquals(2L, subscriptions.position(tp0).offset);
 
         ConsumerRecord record = records.get(0);
         assertArrayEquals("key".getBytes(), record.key());
@@ -395,13 +408,10 @@ public void testFetchError() {
     }
 
     private MockClient.RequestMatcher matchesOffset(final TopicPartition tp, final long offset) {
-        return new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                FetchRequest fetch = (FetchRequest) body;
-                return fetch.fetchData().containsKey(tp) &&
-                        fetch.fetchData().get(tp).fetchOffset == offset;
-            }
+        return body -> {
+            FetchRequest fetch = (FetchRequest) body;
+            return fetch.fetchData().containsKey(tp) &&
+                    fetch.fetchData().get(tp).fetchOffset == offset;
         };
     }
 
@@ -438,7 +448,7 @@ public byte[] deserialize(String topic, byte[] data) {
                 fail("fetchedRecords should have raised");
             } catch (SerializationException e) {
                 // the position should not advance since no data has been returned
-                assertEquals(1, subscriptions.position(tp0).longValue());
+                assertEquals(1, subscriptions.position(tp0).offset);
             }
         }
     }
@@ -496,7 +506,7 @@ public void testParseCorruptedRecord() throws Exception {
 
         // the first fetchedRecords() should return the first valid message
         assertEquals(1, fetcher.fetchedRecords().get(tp0).size());
-        assertEquals(1, subscriptions.position(tp0).longValue());
+        assertEquals(1, subscriptions.position(tp0).offset);
 
         ensureBlockOnRecord(1L);
         seekAndConsumeRecord(buffer, 2L);
@@ -518,7 +528,7 @@ private void ensureBlockOnRecord(long blockedOffset) {
                 fetcher.fetchedRecords();
                 fail("fetchedRecords should have raised KafkaException");
             } catch (KafkaException e) {
-                assertEquals(blockedOffset, subscriptions.position(tp0).longValue());
+                assertEquals(blockedOffset, subscriptions.position(tp0).offset);
             }
         }
     }
@@ -536,7 +546,7 @@ private void seekAndConsumeRecord(ByteBuffer responseBuffer, long toOffset) {
         List> records = recordsByPartition.get(tp0);
         assertEquals(1, records.size());
         assertEquals(toOffset, records.get(0).offset());
-        assertEquals(toOffset + 1, subscriptions.position(tp0).longValue());
+        assertEquals(toOffset + 1, subscriptions.position(tp0).offset);
     }
 
     @Test
@@ -574,7 +584,7 @@ public void testInvalidDefaultRecordBatch() {
                 fetcher.fetchedRecords();
                 fail("fetchedRecords should have raised KafkaException");
             } catch (KafkaException e) {
-                assertEquals(0, subscriptions.position(tp0).longValue());
+                assertEquals(0, subscriptions.position(tp0).offset);
             }
         }
     }
@@ -604,10 +614,47 @@ public void testParseInvalidRecordBatch() {
             fail("fetchedRecords should have raised");
         } catch (KafkaException e) {
             // the position should not advance since no data has been returned
-            assertEquals(0, subscriptions.position(tp0).longValue());
+            assertEquals(0, subscriptions.position(tp0).offset);
+        }
+    }
+
+    @Test
+    public void testFetchEntireBatchWithShallowIteratorEnabled() {
+        for (byte magic : asList(RecordBatch.MAGIC_VALUE_V0, RecordBatch.MAGIC_VALUE_V1, RecordBatch.MAGIC_VALUE_V2)) {
+            ByteBuffer buffer = ByteBuffer.allocate(1024);
+            // create compressed batches
+            MemoryRecordsBuilder builder =
+                new MemoryRecordsBuilder(buffer, magic, CompressionType.GZIP, TimestampType.CREATE_TIME,
+                    0L, 10L, RecordBatch.NO_PRODUCER_ID, RecordBatch.NO_PRODUCER_EPOCH, RecordBatch.NO_SEQUENCE, false, false, 0, 1024);
+
+            builder.append(10L, "key".getBytes(), "value".getBytes());
+            builder.append(11L, "key1".getBytes(), "value1".getBytes());
+            builder.append(12L, "key2".getBytes(), "value2".getBytes());
+            builder.append(13L, "key3".getBytes(), "value3".getBytes());
+
+            builder.close();
+            buffer.flip();
+
+            // create new fetcher with shallow iterator enabled and maxPollRecords set to 6
+            testPassthrough = true;
+            buildFetcher(6);
+
+            assignFromUser(singleton(tp0));
+            subscriptions.seek(tp0, 0);
+
+            // normal fetch
+            assertEquals(1, fetcher.sendFetches());
+            client.prepareResponse(fullFetchResponse(tp0, MemoryRecords.readableRecords(buffer), Errors.NONE, 100L, 0));
+            consumerClient.poll(time.timer(0));
+
+            Map>> partitionRecords = fetchedRecords();
+            List> records = partitionRecords.get(tp0);
+            assertEquals(1, records.size());
         }
     }
 
+
+
     @Test
     public void testHeaders() {
         buildFetcher();
@@ -669,7 +716,7 @@ public void testFetchMaxPollRecords() {
         Map>> recordsByPartition = fetchedRecords();
         records = recordsByPartition.get(tp0);
         assertEquals(2, records.size());
-        assertEquals(3L, subscriptions.position(tp0).longValue());
+        assertEquals(3L, subscriptions.position(tp0).offset);
         assertEquals(1, records.get(0).offset());
         assertEquals(2, records.get(1).offset());
 
@@ -678,7 +725,7 @@ public void testFetchMaxPollRecords() {
         recordsByPartition = fetchedRecords();
         records = recordsByPartition.get(tp0);
         assertEquals(1, records.size());
-        assertEquals(4L, subscriptions.position(tp0).longValue());
+        assertEquals(4L, subscriptions.position(tp0).offset);
         assertEquals(3, records.get(0).offset());
 
         assertTrue(fetcher.sendFetches() > 0);
@@ -686,7 +733,7 @@ public void testFetchMaxPollRecords() {
         recordsByPartition = fetchedRecords();
         records = recordsByPartition.get(tp0);
         assertEquals(2, records.size());
-        assertEquals(6L, subscriptions.position(tp0).longValue());
+        assertEquals(6L, subscriptions.position(tp0).offset);
         assertEquals(4, records.get(0).offset());
         assertEquals(5, records.get(1).offset());
     }
@@ -712,7 +759,7 @@ public void testFetchAfterPartitionWithFetchedRecordsIsUnassigned() {
         Map>> recordsByPartition = fetchedRecords();
         records = recordsByPartition.get(tp0);
         assertEquals(2, records.size());
-        assertEquals(3L, subscriptions.position(tp0).longValue());
+        assertEquals(3L, subscriptions.position(tp0).offset);
         assertEquals(1, records.get(0).offset());
         assertEquals(2, records.get(1).offset());
 
@@ -726,7 +773,7 @@ public void testFetchAfterPartitionWithFetchedRecordsIsUnassigned() {
         assertNull(fetchedRecords.get(tp0));
         records = fetchedRecords.get(tp1);
         assertEquals(2, records.size());
-        assertEquals(6L, subscriptions.position(tp1).longValue());
+        assertEquals(6L, subscriptions.position(tp1).offset);
         assertEquals(4, records.get(0).offset());
         assertEquals(5, records.get(1).offset());
     }
@@ -755,7 +802,7 @@ public void testFetchNonContinuousRecords() {
         Map>> recordsByPartition = fetchedRecords();
         consumerRecords = recordsByPartition.get(tp0);
         assertEquals(3, consumerRecords.size());
-        assertEquals(31L, subscriptions.position(tp0).longValue()); // this is the next fetching position
+        assertEquals(31L, subscriptions.position(tp0).offset); // this is the next fetching position
 
         assertEquals(15L, consumerRecords.get(0).offset());
         assertEquals(20L, consumerRecords.get(1).offset());
@@ -771,8 +818,7 @@ public void testFetchRequestWhenRecordTooLarge() {
         try {
             buildFetcher();
 
-            client.setNodeApiVersions(NodeApiVersions.create(Collections.singletonList(
-                new ApiVersionsResponse.ApiVersion(ApiKeys.FETCH.id, (short) 2, (short) 2))));
+            client.setNodeApiVersions(NodeApiVersions.create(ApiKeys.FETCH.id, (short) 2, (short) 2));
             makeFetchRequestWithIncompleteRecord();
             try {
                 fetcher.fetchedRecords();
@@ -780,7 +826,7 @@ public void testFetchRequestWhenRecordTooLarge() {
             } catch (RecordTooLargeException e) {
                 assertTrue(e.getMessage().startsWith("There are some messages at [Partition=Offset]: "));
                 // the position should not advance since no data has been returned
-                assertEquals(0, subscriptions.position(tp0).longValue());
+                assertEquals(0, subscriptions.position(tp0).offset);
             }
         } finally {
             client.setNodeApiVersions(NodeApiVersions.create());
@@ -803,7 +849,7 @@ public void testFetchRequestInternalError() {
         } catch (KafkaException e) {
             assertTrue(e.getMessage().startsWith("Failed to make progress reading messages"));
             // the position should not advance since no data has been returned
-            assertEquals(0, subscriptions.position(tp0).longValue());
+            assertEquals(0, subscriptions.position(tp0).offset);
         }
     }
 
@@ -839,7 +885,7 @@ public void testUnauthorizedTopic() {
     }
 
     @Test
-    public void testFetchDuringRebalance() {
+    public void testFetchDuringEagerRebalance() {
         buildFetcher();
 
         subscriptions.subscribe(singleton(topicName), listener);
@@ -850,7 +896,9 @@ public void testFetchDuringRebalance() {
 
         assertEquals(1, fetcher.sendFetches());
 
-        // Now the rebalance happens and fetch positions are cleared
+        // Now the eager rebalance happens and fetch positions are cleared
+        subscriptions.assignFromSubscribed(Collections.emptyList());
+
         subscriptions.assignFromSubscribed(singleton(tp0));
         client.prepareResponse(fullFetchResponse(tp0, this.records, Errors.NONE, 100L, 0));
         consumerClient.poll(time.timer(0));
@@ -859,6 +907,31 @@ public void testFetchDuringRebalance() {
         assertTrue(fetcher.fetchedRecords().isEmpty());
     }
 
+    @Test
+    public void testFetchDuringCooperativeRebalance() {
+        buildFetcher();
+
+        subscriptions.subscribe(singleton(topicName), listener);
+        subscriptions.assignFromSubscribed(singleton(tp0));
+        subscriptions.seek(tp0, 0);
+
+        client.updateMetadata(initialUpdateResponse);
+
+        assertEquals(1, fetcher.sendFetches());
+
+        // Now the cooperative rebalance happens and fetch positions are NOT cleared for unrevoked partitions
+        subscriptions.assignFromSubscribed(singleton(tp0));
+
+        client.prepareResponse(fullFetchResponse(tp0, this.records, Errors.NONE, 100L, 0));
+        consumerClient.poll(time.timer(0));
+
+        Map>> fetchedRecords = fetchedRecords();
+
+        // The active fetch should NOT be ignored since the position for tp0 is still valid
+        assertEquals(1, fetchedRecords.size());
+        assertEquals(3, fetchedRecords.get(tp0).size());
+    }
+
     @Test
     public void testInFlightFetchOnPausedPartition() {
         buildFetcher();
@@ -886,6 +959,171 @@ public void testFetchOnPausedPartition() {
         assertTrue(client.requests().isEmpty());
     }
 
+    @Test
+    public void testFetchOnCompletedFetchesForPausedAndResumedPartitions() {
+        buildFetcher();
+
+        assignFromUser(singleton(tp0));
+        subscriptions.seek(tp0, 0);
+
+        assertEquals(1, fetcher.sendFetches());
+
+        subscriptions.pause(tp0);
+
+        client.prepareResponse(fullFetchResponse(tp0, this.records, Errors.NONE, 100L, 0));
+        consumerClient.poll(time.timer(0));
+
+        Map>> fetchedRecords = fetchedRecords();
+        assertEquals("Should not return any records when partition is paused", 0, fetchedRecords.size());
+        assertTrue("Should still contain completed fetches", fetcher.hasCompletedFetches());
+        assertFalse("Should not have any available (non-paused) completed fetches", fetcher.hasAvailableFetches());
+        assertNull(fetchedRecords.get(tp0));
+        assertEquals(0, fetcher.sendFetches());
+
+        subscriptions.resume(tp0);
+
+        assertTrue("Should have available (non-paused) completed fetches", fetcher.hasAvailableFetches());
+
+        consumerClient.poll(time.timer(0));
+        fetchedRecords = fetchedRecords();
+        assertEquals("Should return records when partition is resumed", 1, fetchedRecords.size());
+        assertNotNull(fetchedRecords.get(tp0));
+        assertEquals(3, fetchedRecords.get(tp0).size());
+
+        consumerClient.poll(time.timer(0));
+        fetchedRecords = fetchedRecords();
+        assertEquals("Should not return records after previously paused partitions are fetched", 0, fetchedRecords.size());
+        assertFalse("Should no longer contain completed fetches", fetcher.hasCompletedFetches());
+    }
+
+    @Test
+    public void testFetchOnCompletedFetchesForSomePausedPartitions() {
+        buildFetcher();
+
+        Map>> fetchedRecords;
+
+        assignFromUser(Utils.mkSet(tp0, tp1));
+
+        // seek to tp0 and tp1 in two polls to generate 2 complete requests and responses
+
+        // #1 seek, request, poll, response
+        subscriptions.seek(tp0, 1);
+        assertEquals(1, fetcher.sendFetches());
+        client.prepareResponse(fullFetchResponse(tp0, this.records, Errors.NONE, 100L, 0));
+        consumerClient.poll(time.timer(0));
+
+        // #2 seek, request, poll, response
+        subscriptions.seek(tp1, 1);
+        assertEquals(1, fetcher.sendFetches());
+        client.prepareResponse(fullFetchResponse(tp1, this.nextRecords, Errors.NONE, 100L, 0));
+
+        subscriptions.pause(tp0);
+        consumerClient.poll(time.timer(0));
+
+        fetchedRecords = fetchedRecords();
+        assertEquals("Should return completed fetch for unpaused partitions", 1, fetchedRecords.size());
+        assertTrue("Should still contain completed fetches", fetcher.hasCompletedFetches());
+        assertNotNull(fetchedRecords.get(tp1));
+        assertNull(fetchedRecords.get(tp0));
+
+        fetchedRecords = fetchedRecords();
+        assertEquals("Should return no records for remaining paused partition", 0, fetchedRecords.size());
+        assertTrue("Should still contain completed fetches", fetcher.hasCompletedFetches());
+    }
+
+    @Test
+    public void testFetchOnCompletedFetchesForAllPausedPartitions() {
+        buildFetcher();
+
+        Map>> fetchedRecords;
+
+        assignFromUser(Utils.mkSet(tp0, tp1));
+
+        // seek to tp0 and tp1 in two polls to generate 2 complete requests and responses
+
+        // #1 seek, request, poll, response
+        subscriptions.seek(tp0, 1);
+        assertEquals(1, fetcher.sendFetches());
+        client.prepareResponse(fullFetchResponse(tp0, this.records, Errors.NONE, 100L, 0));
+        consumerClient.poll(time.timer(0));
+
+        // #2 seek, request, poll, response
+        subscriptions.seek(tp1, 1);
+        assertEquals(1, fetcher.sendFetches());
+        client.prepareResponse(fullFetchResponse(tp1, this.nextRecords, Errors.NONE, 100L, 0));
+
+        subscriptions.pause(tp0);
+        subscriptions.pause(tp1);
+
+        consumerClient.poll(time.timer(0));
+
+        fetchedRecords = fetchedRecords();
+        assertEquals("Should return no records for all paused partitions", 0, fetchedRecords.size());
+        assertTrue("Should still contain completed fetches", fetcher.hasCompletedFetches());
+        assertFalse("Should not have any available (non-paused) completed fetches", fetcher.hasAvailableFetches());
+    }
+
+    @Test
+    public void testPartialFetchWithPausedPartitions() {
+        // this test sends creates a completed fetch with 3 records and a max poll of 2 records to assert
+        // that a fetch that must be returned over at least 2 polls can be cached successfully when its partition is
+        // paused, then returned successfully after its been resumed again later
+        buildFetcher(2);
+
+        Map>> fetchedRecords;
+
+        assignFromUser(Utils.mkSet(tp0, tp1));
+
+        subscriptions.seek(tp0, 1);
+        assertEquals(1, fetcher.sendFetches());
+        client.prepareResponse(fullFetchResponse(tp0, this.records, Errors.NONE, 100L, 0));
+        consumerClient.poll(time.timer(0));
+
+        fetchedRecords = fetchedRecords();
+
+        assertEquals("Should return 2 records from fetch with 3 records", 2, fetchedRecords.get(tp0).size());
+        assertFalse("Should have no completed fetches", fetcher.hasCompletedFetches());
+
+        subscriptions.pause(tp0);
+        consumerClient.poll(time.timer(0));
+
+        fetchedRecords = fetchedRecords();
+
+        assertEquals("Should return no records for paused partitions", 0, fetchedRecords.size());
+        assertTrue("Should have 1 entry in completed fetches", fetcher.hasCompletedFetches());
+        assertFalse("Should not have any available (non-paused) completed fetches", fetcher.hasAvailableFetches());
+
+        subscriptions.resume(tp0);
+
+        consumerClient.poll(time.timer(0));
+
+        fetchedRecords = fetchedRecords();
+
+        assertEquals("Should return last remaining record", 1, fetchedRecords.get(tp0).size());
+        assertFalse("Should have no completed fetches", fetcher.hasCompletedFetches());
+    }
+
+    @Test
+    public void testFetchDiscardedAfterPausedPartitionResumedAndSeekedToNewOffset() {
+        buildFetcher();
+        assignFromUser(singleton(tp0));
+        subscriptions.seek(tp0, 0);
+
+        assertEquals(1, fetcher.sendFetches());
+        subscriptions.pause(tp0);
+        client.prepareResponse(fullFetchResponse(tp0, this.records, Errors.NONE, 100L, 0));
+
+        subscriptions.seek(tp0, 3);
+        subscriptions.resume(tp0);
+        consumerClient.poll(time.timer(0));
+
+        assertTrue("Should have 1 entry in completed fetches", fetcher.hasCompletedFetches());
+        Map>> fetchedRecords = fetchedRecords();
+        assertEquals("Should not return any records because we seeked to a new offset", 0, fetchedRecords.size());
+        assertNull(fetchedRecords.get(tp0));
+        assertFalse("Should have no completed fetches", fetcher.hasCompletedFetches());
+    }
+
     @Test
     public void testFetchNotLeaderForPartition() {
         buildFetcher();
@@ -944,15 +1182,11 @@ public void testFetchUnknownLeaderEpoch() {
     public void testEpochSetInFetchRequest() {
         buildFetcher();
         subscriptions.assignFromUser(singleton(tp0));
-        client.updateMetadata(initialUpdateResponse);
-
-        // Metadata update with leader epochs
-        MetadataResponse metadataResponse = TestUtils.metadataUpdateWith("dummy", 1, Collections.emptyMap(), Collections.singletonMap(topicName, 4),
-            (error, partition, leader, leaderEpoch, replicas, isr, offlineReplicas) ->
-                    new MetadataResponse.PartitionMetadata(error, partition, leader, Optional.of(99), replicas, Collections.emptyList(), offlineReplicas));
+        MetadataResponse metadataResponse = TestUtils.metadataUpdateWith(null, 1,
+                Collections.emptyMap(), Collections.singletonMap(topicName, 4), tp -> 99);
         client.updateMetadata(metadataResponse);
 
-        subscriptions.seek(tp0, 0);
+        subscriptions.seek(tp0, 10);
         assertEquals(1, fetcher.sendFetches());
 
         // Check for epoch in outgoing request
@@ -961,7 +1195,7 @@ public void testEpochSetInFetchRequest() {
                 FetchRequest fetchRequest = (FetchRequest) body;
                 fetchRequest.fetchData().values().forEach(partitionData -> {
                     assertTrue("Expected Fetcher to set leader epoch in request", partitionData.currentLeaderEpoch.isPresent());
-                    assertEquals("Expected leader epoch to match epoch from metadata update", partitionData.currentLeaderEpoch.get().longValue(), 99);
+                    assertEquals("Expected leader epoch to match epoch from metadata update", 99, partitionData.currentLeaderEpoch.get().longValue());
                 });
                 return true;
             } else {
@@ -984,7 +1218,8 @@ public void testFetchOffsetOutOfRange() {
         consumerClient.poll(time.timer(0));
         assertEquals(0, fetcher.fetchedRecords().size());
         assertTrue(subscriptions.isOffsetResetNeeded(tp0));
-        assertEquals(null, subscriptions.position(tp0));
+        assertNull(subscriptions.validPosition(tp0));
+        assertNotNull(subscriptions.position(tp0));
     }
 
     @Test
@@ -1001,7 +1236,7 @@ public void testStaleOutOfRangeError() {
         consumerClient.poll(time.timer(0));
         assertEquals(0, fetcher.fetchedRecords().size());
         assertFalse(subscriptions.isOffsetResetNeeded(tp0));
-        assertEquals(1, subscriptions.position(tp0).longValue());
+        assertEquals(1, subscriptions.position(tp0).offset);
     }
 
     @Test
@@ -1065,8 +1300,8 @@ public void testFetchPositionAfterException() {
         List> allFetchedRecords = new ArrayList<>();
         fetchRecordsInto(allFetchedRecords);
 
-        assertEquals(1, subscriptions.position(tp0).longValue());
-        assertEquals(4, subscriptions.position(tp1).longValue());
+        assertEquals(1, subscriptions.position(tp0).offset);
+        assertEquals(4, subscriptions.position(tp1).offset);
         assertEquals(3, allFetchedRecords.size());
 
         OffsetOutOfRangeException e = assertThrows(OffsetOutOfRangeException.class, () ->
@@ -1075,8 +1310,8 @@ public void testFetchPositionAfterException() {
         assertEquals(singleton(tp0), e.offsetOutOfRangePartitions().keySet());
         assertEquals(1L, e.offsetOutOfRangePartitions().get(tp0).longValue());
 
-        assertEquals(1, subscriptions.position(tp0).longValue());
-        assertEquals(4, subscriptions.position(tp1).longValue());
+        assertEquals(1, subscriptions.position(tp0).offset);
+        assertEquals(4, subscriptions.position(tp1).offset);
         assertEquals(3, allFetchedRecords.size());
     }
 
@@ -1118,8 +1353,8 @@ public void testCompletedFetchRemoval() {
         for (List> records : recordsByPartition.values())
             fetchedRecords.addAll(records);
 
-        assertEquals(fetchedRecords.size(), subscriptions.position(tp1) - 1);
-        assertEquals(4, subscriptions.position(tp1).longValue());
+        assertEquals(fetchedRecords.size(), subscriptions.position(tp1).offset - 1);
+        assertEquals(4, subscriptions.position(tp1).offset);
         assertEquals(3, fetchedRecords.size());
 
         List oorExceptions = new ArrayList<>();
@@ -1142,7 +1377,7 @@ public void testCompletedFetchRemoval() {
             fetchedRecords.addAll(records);
 
         // Should not have received an Exception for tp2.
-        assertEquals(6, subscriptions.position(tp2).longValue());
+        assertEquals(6, subscriptions.position(tp2).offset);
         assertEquals(5, fetchedRecords.size());
 
         int numExceptionsExpected = 3;
@@ -1170,7 +1405,7 @@ public void testSeekBeforeException() {
         assertEquals(1, fetcher.sendFetches());
         Map> partitions = new HashMap<>();
         partitions.put(tp0, new FetchResponse.PartitionData<>(Errors.NONE, 100,
-                FetchResponse.INVALID_LAST_STABLE_OFFSET, FetchResponse.INVALID_LOG_START_OFFSET, null, records));
+                FetchResponse.INVALID_LAST_STABLE_OFFSET, FetchResponse.INVALID_LOG_START_OFFSET, Optional.empty(), null, records));
         client.prepareResponse(fullFetchResponse(tp0, this.records, Errors.NONE, 100L, 0));
         consumerClient.poll(time.timer(0));
 
@@ -1181,7 +1416,7 @@ public void testSeekBeforeException() {
         assertEquals(1, fetcher.sendFetches());
         partitions = new HashMap<>();
         partitions.put(tp1, new FetchResponse.PartitionData<>(Errors.OFFSET_OUT_OF_RANGE, 100,
-                FetchResponse.INVALID_LAST_STABLE_OFFSET, FetchResponse.INVALID_LOG_START_OFFSET, null, MemoryRecords.EMPTY));
+                FetchResponse.INVALID_LAST_STABLE_OFFSET, FetchResponse.INVALID_LOG_START_OFFSET, Optional.empty(), null, MemoryRecords.EMPTY));
         client.prepareResponse(new FetchResponse<>(Errors.NONE, new LinkedHashMap<>(partitions), 0, INVALID_SESSION_ID));
         consumerClient.poll(time.timer(0));
         assertEquals(1, fetcher.fetchedRecords().get(tp0).size());
@@ -1206,7 +1441,7 @@ public void testFetchDisconnected() {
         // disconnects should have no affect on subscription state
         assertFalse(subscriptions.isOffsetResetNeeded(tp0));
         assertTrue(subscriptions.isFetchable(tp0));
-        assertEquals(0, subscriptions.position(tp0).longValue());
+        assertEquals(0, subscriptions.position(tp0).offset);
     }
 
     @Test
@@ -1218,7 +1453,7 @@ public void testUpdateFetchPositionNoOpWithPositionSet() {
         fetcher.resetOffsetsIfNeeded();
         assertFalse(client.hasInFlightRequests());
         assertTrue(subscriptions.isFetchable(tp0));
-        assertEquals(5, subscriptions.position(tp0).longValue());
+        assertEquals(5, subscriptions.position(tp0).offset);
     }
 
     @Test
@@ -1233,7 +1468,7 @@ public void testUpdateFetchPositionResetToDefaultOffset() {
         consumerClient.pollNoWakeup();
         assertFalse(subscriptions.isOffsetResetNeeded(tp0));
         assertTrue(subscriptions.isFetchable(tp0));
-        assertEquals(5, subscriptions.position(tp0).longValue());
+        assertEquals(5, subscriptions.position(tp0).offset);
     }
 
     @Test
@@ -1250,7 +1485,7 @@ public void testUpdateFetchPositionResetToLatestOffset() {
         consumerClient.pollNoWakeup();
         assertFalse(subscriptions.isOffsetResetNeeded(tp0));
         assertTrue(subscriptions.isFetchable(tp0));
-        assertEquals(5, subscriptions.position(tp0).longValue());
+        assertEquals(5, subscriptions.position(tp0).offset);
     }
 
     /**
@@ -1290,7 +1525,7 @@ public void testFetchOffsetErrors() {
         assertTrue(subscriptions.hasValidPosition(tp0));
         assertFalse(subscriptions.isOffsetResetNeeded(tp0));
         assertTrue(subscriptions.isFetchable(tp0));
-        assertEquals(subscriptions.position(tp0).longValue(), 5L);
+        assertEquals(subscriptions.position(tp0).offset, 5L);
     }
 
     @Test
@@ -1319,7 +1554,7 @@ private void testListOffsetsSendsIsolationLevel(IsolationLevel isolationLevel) {
 
         assertFalse(subscriptions.isOffsetResetNeeded(tp0));
         assertTrue(subscriptions.isFetchable(tp0));
-        assertEquals(5, subscriptions.position(tp0).longValue());
+        assertEquals(5, subscriptions.position(tp0).offset);
     }
 
     @Test
@@ -1346,7 +1581,7 @@ public void testResetOffsetsSkipsBlackedOutConnections() {
 
         assertFalse(subscriptions.isOffsetResetNeeded(tp0));
         assertTrue(subscriptions.isFetchable(tp0));
-        assertEquals(5, subscriptions.position(tp0).longValue());
+        assertEquals(5, subscriptions.position(tp0).offset);
     }
 
     @Test
@@ -1362,7 +1597,7 @@ public void testUpdateFetchPositionResetToEarliestOffset() {
 
         assertFalse(subscriptions.isOffsetResetNeeded(tp0));
         assertTrue(subscriptions.isFetchable(tp0));
-        assertEquals(5, subscriptions.position(tp0).longValue());
+        assertEquals(5, subscriptions.position(tp0).offset);
     }
 
     @Test
@@ -1392,7 +1627,7 @@ public void testResetOffsetsMetadataRefresh() {
 
         assertFalse(subscriptions.isOffsetResetNeeded(tp0));
         assertTrue(subscriptions.isFetchable(tp0));
-        assertEquals(5, subscriptions.position(tp0).longValue());
+        assertEquals(5, subscriptions.position(tp0).offset);
     }
 
     @Test
@@ -1428,7 +1663,7 @@ public void testUpdateFetchPositionDisconnect() {
 
         assertFalse(subscriptions.isOffsetResetNeeded(tp0));
         assertTrue(subscriptions.isFetchable(tp0));
-        assertEquals(5, subscriptions.position(tp0).longValue());
+        assertEquals(5, subscriptions.position(tp0).offset);
     }
 
     @Test
@@ -1476,7 +1711,52 @@ public void testSeekWithInFlightReset() {
 
         assertFalse(client.hasPendingResponses());
         assertFalse(client.hasInFlightRequests());
-        assertEquals(237L, subscriptions.position(tp0).longValue());
+        assertEquals(237L, subscriptions.position(tp0).offset);
+    }
+
+    @Test(timeout = 10000)
+    public void testEarlierOffsetResetArrivesLate() throws InterruptedException {
+        LogContext lc = new LogContext();
+        buildFetcher(spy(new SubscriptionState(lc, OffsetResetStrategy.EARLIEST)), lc);
+        assignFromUser(singleton(tp0));
+
+        ExecutorService es = Executors.newSingleThreadExecutor();
+        CountDownLatch latchLatestStart = new CountDownLatch(1);
+        CountDownLatch latchEarliestStart = new CountDownLatch(1);
+        CountDownLatch latchEarliestDone = new CountDownLatch(1);
+        CountDownLatch latchEarliestFinish = new CountDownLatch(1);
+        try {
+            doAnswer(invocation -> {
+                latchLatestStart.countDown();
+                latchEarliestStart.await();
+                Object result = invocation.callRealMethod();
+                latchEarliestDone.countDown();
+                return result;
+            }).when(subscriptions).maybeSeekUnvalidated(tp0, 0L, OffsetResetStrategy.EARLIEST);
+
+            es.submit(() -> {
+                subscriptions.requestOffsetReset(tp0, OffsetResetStrategy.EARLIEST);
+                fetcher.resetOffsetsIfNeeded();
+                consumerClient.pollNoWakeup();
+                client.respond(listOffsetResponse(Errors.NONE, 1L, 0L));
+                consumerClient.pollNoWakeup();
+                latchEarliestFinish.countDown();
+            }, Void.class);
+
+            latchLatestStart.await();
+            subscriptions.requestOffsetReset(tp0, OffsetResetStrategy.LATEST);
+            fetcher.resetOffsetsIfNeeded();
+            consumerClient.pollNoWakeup();
+            client.respond(listOffsetResponse(Errors.NONE, 1L, 10L));
+            latchEarliestStart.countDown();
+            latchEarliestDone.await();
+            consumerClient.pollNoWakeup();
+            latchEarliestFinish.await();
+            assertEquals(10, subscriptions.position(tp0).offset);
+        } finally {
+            es.shutdown();
+            es.awaitTermination(10000, TimeUnit.MILLISECONDS);
+        }
     }
 
     @Test
@@ -1524,7 +1804,7 @@ public void testIdempotentResetWithInFlightReset() {
 
         assertFalse(client.hasInFlightRequests());
         assertFalse(subscriptions.isOffsetResetNeeded(tp0));
-        assertEquals(5L, subscriptions.position(tp0).longValue());
+        assertEquals(5L, subscriptions.position(tp0).offset);
     }
 
     @Test
@@ -1562,7 +1842,7 @@ public void testRestOffsetsAuthorizationFailure() {
 
         assertFalse(subscriptions.isOffsetResetNeeded(tp0));
         assertTrue(subscriptions.isFetchable(tp0));
-        assertEquals(5, subscriptions.position(tp0).longValue());
+        assertEquals(5, subscriptions.position(tp0).offset);
     }
 
     @Test
@@ -1580,7 +1860,7 @@ public void testUpdateFetchPositionOfPausedPartitionsRequiringOffsetReset() {
         assertFalse(subscriptions.isOffsetResetNeeded(tp0));
         assertFalse(subscriptions.isFetchable(tp0)); // because tp is paused
         assertTrue(subscriptions.hasValidPosition(tp0));
-        assertEquals(10, subscriptions.position(tp0).longValue());
+        assertEquals(10, subscriptions.position(tp0).offset);
     }
 
     @Test
@@ -1610,7 +1890,7 @@ public void testUpdateFetchPositionOfPausedPartitionsWithAValidPosition() {
         assertFalse(subscriptions.isOffsetResetNeeded(tp0));
         assertFalse(subscriptions.isFetchable(tp0)); // because tp is paused
         assertTrue(subscriptions.hasValidPosition(tp0));
-        assertEquals(10, subscriptions.position(tp0).longValue());
+        assertEquals(10, subscriptions.position(tp0).offset);
     }
 
     @Test
@@ -1720,8 +2000,8 @@ public void testGetTopicMetadataOfflinePartitions() {
             altTopics.add(alteredTopic);
         }
         Node controller = originalResponse.controller();
-        MetadataResponse altered = new MetadataResponse(
-            (List) originalResponse.brokers(),
+        MetadataResponse altered = MetadataResponse.prepareResponse(
+            originalResponse.brokers(),
             originalResponse.clusterId(),
             controller != null ? controller.id() : MetadataResponse.NO_CONTROLLER_ID,
             altTopics);
@@ -1753,8 +2033,9 @@ public void testQuotaMetrics() {
                 1000, 1000, 64 * 1024, 64 * 1024, 1000,  ClientDnsLookup.DEFAULT,
                 time, true, new ApiVersions(), throttleTimeSensor, new LogContext());
 
-        short apiVersionsResponseVersion = ApiKeys.API_VERSIONS.latestVersion();
-        ByteBuffer buffer = ApiVersionsResponse.createApiVersionsResponse(400, RecordBatch.CURRENT_MAGIC_VALUE).serialize(apiVersionsResponseVersion, new ResponseHeader(0));
+        ByteBuffer buffer = ApiVersionsResponse.
+            createApiVersionsResponse(400, RecordBatch.CURRENT_MAGIC_VALUE).
+                serialize(ApiKeys.API_VERSIONS, ApiKeys.API_VERSIONS.latestVersion(), 0);
         selector.delayedReceive(new DelayedReceive(node.idString(), new NetworkReceive(node.idString(), buffer)));
         while (!client.ready(node, time.milliseconds())) {
             client.poll(1, time.milliseconds());
@@ -1766,11 +2047,14 @@ public void testQuotaMetrics() {
         for (int i = 1; i <= 3; i++) {
             int throttleTimeMs = 100 * i;
             FetchRequest.Builder builder = FetchRequest.Builder.forConsumer(100, 100, new LinkedHashMap<>());
+            builder.rackId("");
             ClientRequest request = client.newClientRequest(node.idString(), builder, time.milliseconds(), true);
             client.send(request, time.milliseconds());
             client.poll(1, time.milliseconds());
             FetchResponse response = fullFetchResponse(tp0, nextRecords, Errors.NONE, i, throttleTimeMs);
-            buffer = response.serialize(ApiKeys.FETCH.latestVersion(), new ResponseHeader(request.correlationId()));
+            buffer = response.serialize(ApiKeys.FETCH,
+                    ApiKeys.FETCH.latestVersion(),
+                    request.correlationId());
             selector.completeReceive(new NetworkReceive(node.idString(), buffer));
             client.poll(1, time.milliseconds());
             // If a throttled response is received, advance the time to ensure progress.
@@ -1825,6 +2109,7 @@ public void testFetcherMetrics() {
 
         // verify de-registration of partition lag
         subscriptions.unsubscribe();
+        fetcher.sendFetches();
         assertFalse(allMetrics.containsKey(partitionLagMetric));
     }
 
@@ -1866,6 +2151,7 @@ public void testFetcherLeadMetric() {
 
         // verify de-registration of partition lag
         subscriptions.unsubscribe();
+        fetcher.sendFetches();
         assertFalse(allMetrics.containsKey(partitionLeadMetric));
     }
 
@@ -1908,6 +2194,7 @@ public void testReadCommittedLagMetric() {
 
         // verify de-registration of partition lag
         subscriptions.unsubscribe();
+        fetcher.sendFetches();
         assertFalse(allMetrics.containsKey(partitionLagMetric));
     }
 
@@ -2120,13 +2407,9 @@ private Map>> fetchRecords(
 
     @Test
     public void testGetOffsetsForTimesTimeout() {
-        try {
-            buildFetcher();
-            fetcher.offsetsForTimes(Collections.singletonMap(new TopicPartition(topicName, 2), 1000L), time.timer(100L));
-            fail("Should throw timeout exception.");
-        } catch (TimeoutException e) {
-            // let it go.
-        }
+        buildFetcher();
+        assertThrows(TimeoutException.class, () -> fetcher.offsetsForTimes(
+            Collections.singletonMap(new TopicPartition(topicName, 2), 1000L), time.timer(100L)));
     }
 
     @Test
@@ -2134,7 +2417,7 @@ public void testGetOffsetsForTimes() {
         buildFetcher();
 
         // Empty map
-        assertTrue(fetcher.offsetsForTimes(new HashMap(), time.timer(100L)).isEmpty());
+        assertTrue(fetcher.offsetsForTimes(new HashMap<>(), time.timer(100L)).isEmpty());
         // Unknown Offset
         testGetOffsetsForTimesWithUnknownOffset();
         // Error code none with unknown offset
@@ -2170,6 +2453,89 @@ public void testGetOffsetsFencedLeaderEpoch() {
         assertEquals(0L, metadata.timeToNextUpdate(time.milliseconds()));
     }
 
+    @Test
+    public void testGetOffsetByTimeWithPartitionsRetryCouldTriggerMetadataUpdate() {
+        List retriableErrors = Arrays.asList(Errors.NOT_LEADER_FOR_PARTITION,
+            Errors.REPLICA_NOT_AVAILABLE, Errors.KAFKA_STORAGE_ERROR, Errors.OFFSET_NOT_AVAILABLE,
+            Errors.LEADER_NOT_AVAILABLE, Errors.FENCED_LEADER_EPOCH, Errors.UNKNOWN_LEADER_EPOCH);
+
+        final int newLeaderEpoch = 3;
+        MetadataResponse updatedMetadata = TestUtils.metadataUpdateWith(null, 3,
+            singletonMap(topicName, Errors.NONE), singletonMap(topicName, 4), tp -> newLeaderEpoch);
+
+        Node originalLeader = initialUpdateResponse.cluster().leaderFor(tp1);
+        Node newLeader = updatedMetadata.cluster().leaderFor(tp1);
+        assertNotEquals(originalLeader, newLeader);
+
+        for (Errors retriableError : retriableErrors) {
+            buildFetcher();
+
+            subscriptions.assignFromUser(Utils.mkSet(tp0, tp1));
+            client.updateMetadata(initialUpdateResponse);
+
+            final long fetchTimestamp = 10L;
+            Map allPartitionData = new HashMap<>();
+            allPartitionData.put(tp0, new ListOffsetResponse.PartitionData(
+                Errors.NONE, fetchTimestamp, 4L, Optional.empty()));
+            allPartitionData.put(tp1, new ListOffsetResponse.PartitionData(
+                retriableError, ListOffsetRequest.LATEST_TIMESTAMP, -1L, Optional.empty()));
+
+            client.prepareResponseFrom(body -> {
+                boolean isListOffsetRequest = body instanceof ListOffsetRequest;
+                if (isListOffsetRequest) {
+                    ListOffsetRequest request = (ListOffsetRequest) body;
+                    Map expectedTopicPartitions = new HashMap<>();
+                    expectedTopicPartitions.put(tp0, new ListOffsetRequest.PartitionData(
+                        fetchTimestamp, Optional.empty()));
+                    expectedTopicPartitions.put(tp1, new ListOffsetRequest.PartitionData(
+                        fetchTimestamp, Optional.empty()));
+
+                    return request.partitionTimestamps().equals(expectedTopicPartitions);
+                } else {
+                    return false;
+                }
+            }, new ListOffsetResponse(allPartitionData), originalLeader);
+
+            client.prepareMetadataUpdate(updatedMetadata);
+
+            // If the metadata wasn't updated before retrying, the fetcher would consult the original leader and hit a NOT_LEADER exception.
+            // We will count the answered future response in the end to verify if this is the case.
+            Map paritionDataWithFatalError = new HashMap<>(allPartitionData);
+            paritionDataWithFatalError.put(tp1, new ListOffsetResponse.PartitionData(
+                Errors.NOT_LEADER_FOR_PARTITION, ListOffsetRequest.LATEST_TIMESTAMP, -1L, Optional.empty()));
+            client.prepareResponseFrom(new ListOffsetResponse(paritionDataWithFatalError), originalLeader);
+
+            // The request to new leader must only contain one partition tp1 with error.
+            client.prepareResponseFrom(body -> {
+                boolean isListOffsetRequest = body instanceof ListOffsetRequest;
+                if (isListOffsetRequest) {
+                    ListOffsetRequest request = (ListOffsetRequest) body;
+
+                    return request.partitionTimestamps().equals(
+                        Collections.singletonMap(tp1, new ListOffsetRequest.PartitionData(
+                            fetchTimestamp, Optional.of(newLeaderEpoch))));
+                } else {
+                    return false;
+                }
+            }, listOffsetResponse(tp1, Errors.NONE, fetchTimestamp, 5L), newLeader);
+
+            Map offsetAndTimestampMap =
+                fetcher.offsetsForTimes(
+                    Utils.mkMap(Utils.mkEntry(tp0, fetchTimestamp),
+                    Utils.mkEntry(tp1, fetchTimestamp)), time.timer(Integer.MAX_VALUE));
+
+            assertEquals(Utils.mkMap(
+                Utils.mkEntry(tp0, new OffsetAndTimestamp(4L, fetchTimestamp)),
+                Utils.mkEntry(tp1, new OffsetAndTimestamp(5L, fetchTimestamp))), offsetAndTimestampMap);
+
+            // The NOT_LEADER exception future should not be cleared as we already refreshed the metadata before
+            // first retry, thus never hitting.
+            assertEquals(1, client.numAwaitingResponses());
+
+            fetcher.close();
+        }
+    }
+
     @Test
     public void testGetOffsetsUnknownLeaderEpoch() {
         buildFetcher();
@@ -2194,9 +2560,8 @@ public void testGetOffsetsIncludesLeaderEpoch() {
         client.updateMetadata(initialUpdateResponse);
 
         // Metadata update with leader epochs
-        MetadataResponse metadataResponse = TestUtils.metadataUpdateWith("dummy", 1, Collections.emptyMap(), Collections.singletonMap(topicName, 4),
-            (error, partition, leader, leaderEpoch, replicas, isr, offlineReplicas) ->
-                new MetadataResponse.PartitionMetadata(error, partition, leader, Optional.of(99), replicas, Collections.emptyList(), offlineReplicas));
+        MetadataResponse metadataResponse = TestUtils.metadataUpdateWith(null, 1,
+                Collections.emptyMap(), Collections.singletonMap(topicName, 4), tp -> 99);
         client.updateMetadata(metadataResponse);
 
         // Request latest offset
@@ -2225,6 +2590,7 @@ public void testGetOffsetsIncludesLeaderEpoch() {
     public void testGetOffsetsForTimesWhenSomeTopicPartitionLeadersNotKnownInitially() {
         buildFetcher();
 
+        subscriptions.assignFromUser(Utils.mkSet(tp0, tp1));
         final String anotherTopic = "another-topic";
         final TopicPartition t2p0 = new TopicPartition(anotherTopic, 0);
 
@@ -2333,7 +2699,7 @@ public void testReturnCommittedTransactions() {
                 new SimpleRecord(time.milliseconds(), "key".getBytes(), "value".getBytes()),
                 new SimpleRecord(time.milliseconds(), "key".getBytes(), "value".getBytes()));
 
-        currentOffset += commitTransaction(buffer, 1L, currentOffset);
+        commitTransaction(buffer, 1L, currentOffset);
         buffer.flip();
 
         List abortedTransactions = new ArrayList<>();
@@ -2345,13 +2711,10 @@ public void testReturnCommittedTransactions() {
         // normal fetch
         assertEquals(1, fetcher.sendFetches());
         assertFalse(fetcher.hasCompletedFetches());
-        client.prepareResponse(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                FetchRequest request = (FetchRequest) body;
-                assertEquals(IsolationLevel.READ_COMMITTED, request.isolationLevel());
-                return true;
-            }
+        client.prepareResponse(body -> {
+            FetchRequest request = (FetchRequest) body;
+            assertEquals(IsolationLevel.READ_COMMITTED, request.isolationLevel());
+            return true;
         }, fullFetchResponseWithAbortedTransactions(records, abortedTransactions, Errors.NONE, 100L, 100L, 0));
 
         consumerClient.poll(time.timer(0));
@@ -2483,7 +2846,7 @@ public void testMultipleAbortMarkers() {
         for (ConsumerRecord consumerRecord : fetchedConsumerRecords) {
             actuallyCommittedKeys.add(new String(consumerRecord.key(), StandardCharsets.UTF_8));
         }
-        assertTrue(actuallyCommittedKeys.equals(committedKeys));
+        assertEquals(actuallyCommittedKeys, committedKeys);
     }
 
     @Test
@@ -2568,7 +2931,7 @@ protected boolean shouldRetainRecord(RecordBatch recordBatch, Record record) {
         }
 
         // The next offset should point to the next batch
-        assertEquals(4L, subscriptions.position(tp0).longValue());
+        assertEquals(4L, subscriptions.position(tp0).offset);
     }
 
     @Test
@@ -2599,7 +2962,7 @@ public void testUpdatePositionOnEmptyBatch() {
         assertTrue(allFetchedRecords.isEmpty());
 
         // The next offset should point to the next batch
-        assertEquals(lastOffset + 1, subscriptions.position(tp0).longValue());
+        assertEquals(lastOffset + 1, subscriptions.position(tp0).offset);
     }
 
     @Test
@@ -2731,7 +3094,7 @@ public void testConsumerPositionUpdatedWhenSkippingAbortedTransactions() {
 
         // Ensure that we don't return any of the aborted records, but yet advance the consumer position.
         assertFalse(fetchedRecords.containsKey(tp0));
-        assertEquals(currentOffset, (long) subscriptions.position(tp0));
+        assertEquals(currentOffset, subscriptions.position(tp0).offset);
     }
 
     @Test
@@ -2740,8 +3103,8 @@ public void testConsumingViaIncrementalFetchRequests() {
 
         List> records;
         assignFromUser(new HashSet<>(Arrays.asList(tp0, tp1)));
-        subscriptions.seek(tp0, 0);
-        subscriptions.seek(tp1, 1);
+        subscriptions.seekValidated(tp0, new SubscriptionState.FetchPosition(0, Optional.empty(), metadata.leaderAndEpoch(tp0)));
+        subscriptions.seekValidated(tp1, new SubscriptionState.FetchPosition(1, Optional.empty(), metadata.leaderAndEpoch(tp1)));
 
         // Fetch some records and establish an incremental fetch session.
         LinkedHashMap> partitions1 = new LinkedHashMap<>();
@@ -2759,8 +3122,8 @@ public void testConsumingViaIncrementalFetchRequests() {
         assertFalse(fetchedRecords.containsKey(tp1));
         records = fetchedRecords.get(tp0);
         assertEquals(2, records.size());
-        assertEquals(3L, subscriptions.position(tp0).longValue());
-        assertEquals(1L, subscriptions.position(tp1).longValue());
+        assertEquals(3L, subscriptions.position(tp0).offset);
+        assertEquals(1L, subscriptions.position(tp1).offset);
         assertEquals(1, records.get(0).offset());
         assertEquals(2, records.get(1).offset());
 
@@ -2771,7 +3134,7 @@ public void testConsumingViaIncrementalFetchRequests() {
         records = fetchedRecords.get(tp0);
         assertEquals(1, records.size());
         assertEquals(3, records.get(0).offset());
-        assertEquals(4L, subscriptions.position(tp0).longValue());
+        assertEquals(4L, subscriptions.position(tp0).offset);
 
         // The second response contains no new records.
         LinkedHashMap> partitions2 = new LinkedHashMap<>();
@@ -2781,8 +3144,8 @@ public void testConsumingViaIncrementalFetchRequests() {
         consumerClient.poll(time.timer(0));
         fetchedRecords = fetchedRecords();
         assertTrue(fetchedRecords.isEmpty());
-        assertEquals(4L, subscriptions.position(tp0).longValue());
-        assertEquals(1L, subscriptions.position(tp1).longValue());
+        assertEquals(4L, subscriptions.position(tp0).offset);
+        assertEquals(1L, subscriptions.position(tp1).offset);
 
         // The third response contains some new records for tp0.
         LinkedHashMap> partitions3 = new LinkedHashMap<>();
@@ -2796,8 +3159,8 @@ public void testConsumingViaIncrementalFetchRequests() {
         assertFalse(fetchedRecords.containsKey(tp1));
         records = fetchedRecords.get(tp0);
         assertEquals(2, records.size());
-        assertEquals(6L, subscriptions.position(tp0).longValue());
-        assertEquals(1L, subscriptions.position(tp1).longValue());
+        assertEquals(6L, subscriptions.position(tp0).offset);
+        assertEquals(1L, subscriptions.position(tp1).offset);
         assertEquals(4, records.get(0).offset());
         assertEquals(5, records.get(1).offset());
     }
@@ -2809,7 +3172,8 @@ public void testFetcherConcurrency() throws Exception {
         for (int i = 0; i < numPartitions; i++)
             topicPartitions.add(new TopicPartition(topicName, i));
 
-        buildDependencies(new MetricConfig(), OffsetResetStrategy.EARLIEST);
+        LogContext logContext = new LogContext();
+        buildDependencies(new MetricConfig(), Long.MAX_VALUE, new SubscriptionState(logContext, OffsetResetStrategy.EARLIEST), logContext);
 
         fetcher = new Fetcher(
                 new LogContext(),
@@ -2820,6 +3184,8 @@ public void testFetcherConcurrency() throws Exception {
                 fetchSize,
                 2 * numPartitions,
                 true,
+                "",
+                testPassthrough,
                 new ByteArrayDeserializer(),
                 new ByteArrayDeserializer(),
                 metadata,
@@ -2829,7 +3195,8 @@ public void testFetcherConcurrency() throws Exception {
                 time,
                 retryBackoffMs,
                 requestTimeoutMs,
-                IsolationLevel.READ_UNCOMMITTED) {
+                IsolationLevel.READ_UNCOMMITTED,
+                apiVersions) {
             @Override
             protected FetchSessionHandler sessionHandler(int id) {
                 final FetchSessionHandler handler = super.sessionHandler(id);
@@ -2933,6 +3300,62 @@ private void verifySessionPartitions() {
         assertEquals(0, future.get());
     }
 
+    @Test
+    public void testFetcherSessionEpochUpdate() throws Exception {
+        buildFetcher(2);
+
+        MetadataResponse initialMetadataResponse = TestUtils.metadataUpdateWith(1, singletonMap(topicName, 1));
+        client.updateMetadata(initialMetadataResponse);
+        assignFromUser(Collections.singleton(tp0));
+        subscriptions.seek(tp0, 0L);
+
+        AtomicInteger fetchesRemaining = new AtomicInteger(1000);
+        executorService = Executors.newSingleThreadExecutor();
+        Future future = executorService.submit(() -> {
+            long nextOffset = 0;
+            long nextEpoch = 0;
+            while (fetchesRemaining.get() > 0) {
+                synchronized (consumerClient) {
+                    if (!client.requests().isEmpty()) {
+                        ClientRequest request = client.requests().peek();
+                        FetchRequest fetchRequest = (FetchRequest) request.requestBuilder().build();
+                        int epoch = fetchRequest.metadata().epoch();
+                        assertTrue(String.format("Unexpected epoch expected %d got %d", nextEpoch, epoch), epoch == 0 || epoch == nextEpoch);
+                        nextEpoch++;
+                        LinkedHashMap> responseMap = new LinkedHashMap<>();
+                        responseMap.put(tp0, new FetchResponse.PartitionData<>(Errors.NONE, nextOffset + 2L, nextOffset + 2,
+                                0L, null, buildRecords(nextOffset, 2, nextOffset)));
+                        nextOffset += 2;
+                        client.respondToRequest(request, new FetchResponse<>(Errors.NONE, responseMap, 0, 123));
+                        consumerClient.poll(time.timer(0));
+                    }
+                }
+            }
+            return fetchesRemaining.get();
+        });
+        long nextFetchOffset = 0;
+        while (fetchesRemaining.get() > 0 && !future.isDone()) {
+            if (fetcher.sendFetches() == 1) {
+                synchronized (consumerClient) {
+                    consumerClient.poll(time.timer(0));
+                }
+            }
+            if (fetcher.hasCompletedFetches()) {
+                Map>> fetchedRecords = fetchedRecords();
+                if (!fetchedRecords.isEmpty()) {
+                    fetchesRemaining.decrementAndGet();
+                    List> records = fetchedRecords.get(tp0);
+                    assertEquals(2, records.size());
+                    assertEquals(nextFetchOffset, records.get(0).offset());
+                    assertEquals(nextFetchOffset + 1, records.get(1).offset());
+                    nextFetchOffset += 2;
+                }
+                assertTrue(fetchedRecords().isEmpty());
+            }
+        }
+        assertEquals(0, future.get());
+    }
+
     @Test
     public void testEmptyControlBatch() {
         buildFetcher(OffsetResetStrategy.EARLIEST, new ByteArrayDeserializer(),
@@ -2962,13 +3385,10 @@ public void testEmptyControlBatch() {
         // normal fetch
         assertEquals(1, fetcher.sendFetches());
         assertFalse(fetcher.hasCompletedFetches());
-        client.prepareResponse(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                FetchRequest request = (FetchRequest) body;
-                assertEquals(IsolationLevel.READ_COMMITTED, request.isolationLevel());
-                return true;
-            }
+        client.prepareResponse(body -> {
+            FetchRequest request = (FetchRequest) body;
+            assertEquals(IsolationLevel.READ_COMMITTED, request.isolationLevel());
+            return true;
         }, fullFetchResponseWithAbortedTransactions(records, abortedTransactions, Errors.NONE, 100L, 100L, 0));
 
         consumerClient.poll(time.timer(0));
@@ -3002,12 +3422,11 @@ private int appendTransactionalRecords(ByteBuffer buffer, long pid, long baseOff
         return appendTransactionalRecords(buffer, pid, baseOffset, (int) baseOffset, records);
     }
 
-    private int commitTransaction(ByteBuffer buffer, long producerId, long baseOffset) {
+    private void commitTransaction(ByteBuffer buffer, long producerId, long baseOffset) {
         short producerEpoch = 0;
         int partitionLeaderEpoch = 0;
         MemoryRecords.writeEndTransactionalMarker(buffer, baseOffset, time.milliseconds(), partitionLeaderEpoch, producerId, producerEpoch,
                 new EndTransactionMarker(ControlRecordType.COMMIT, 0));
-        return 1;
     }
 
     private int abortTransaction(ByteBuffer buffer, long producerId, long baseOffset) {
@@ -3029,7 +3448,7 @@ private void testGetOffsetsForTimesWithError(Errors errorForP0,
         TopicPartition t2p0 = new TopicPartition(topicName2, 0);
         // Expect a metadata refresh.
         metadata.bootstrap(ClientUtils.parseAndValidateAddresses(Collections.singletonList("1.1.1.1:1111"),
-                ClientDnsLookup.DEFAULT), time.milliseconds());
+                ClientDnsLookup.DEFAULT));
 
         Map partitionNumByTopic = new HashMap<>();
         partitionNumByTopic.put(topicName, 2);
@@ -3095,14 +3514,410 @@ private void testGetOffsetsForTimesWithUnknownOffset() {
         assertNull(offsetAndTimestampMap.get(tp0));
     }
 
+    @Test
+    public void testSubscriptionPositionUpdatedWithEpoch() {
+        // Create some records that include a leader epoch (1)
+        MemoryRecordsBuilder builder = MemoryRecords.builder(
+                ByteBuffer.allocate(1024),
+                RecordBatch.CURRENT_MAGIC_VALUE,
+                CompressionType.NONE,
+                TimestampType.CREATE_TIME,
+                0L,
+                RecordBatch.NO_TIMESTAMP,
+                RecordBatch.NO_PRODUCER_ID,
+                RecordBatch.NO_PRODUCER_EPOCH,
+                RecordBatch.NO_SEQUENCE,
+                false,
+                1
+        );
+        builder.appendWithOffset(0L, 0L, "key".getBytes(), "value-1".getBytes());
+        builder.appendWithOffset(1L, 0L, "key".getBytes(), "value-2".getBytes());
+        builder.appendWithOffset(2L, 0L, "key".getBytes(), "value-3".getBytes());
+        MemoryRecords records = builder.build();
+
+        buildFetcher();
+        assignFromUser(singleton(tp0));
+
+        // Initialize the epoch=1
+        Map partitionCounts = new HashMap<>();
+        partitionCounts.put(tp0.topic(), 4);
+        MetadataResponse metadataResponse = TestUtils.metadataUpdateWith(null, 1, Collections.emptyMap(), partitionCounts, tp -> 1);
+        metadata.update(metadataResponse, 0L);
+
+        // Seek
+        subscriptions.seek(tp0, 0);
+
+        // Do a normal fetch
+        assertEquals(1, fetcher.sendFetches());
+        assertFalse(fetcher.hasCompletedFetches());
+
+        client.prepareResponse(fullFetchResponse(tp0, records, Errors.NONE, 100L, 0));
+        consumerClient.pollNoWakeup();
+        assertTrue(fetcher.hasCompletedFetches());
+
+        Map>> partitionRecords = fetchedRecords();
+        assertTrue(partitionRecords.containsKey(tp0));
+
+        assertEquals(subscriptions.position(tp0).offset, 3L);
+        assertOptional(subscriptions.position(tp0).offsetEpoch, value -> assertEquals(value.intValue(), 1));
+    }
+
+    @Test
+    public void testOffsetValidationAwaitsNodeApiVersion() {
+        buildFetcher();
+        assignFromUser(singleton(tp0));
+
+        Map partitionCounts = new HashMap<>();
+        partitionCounts.put(tp0.topic(), 4);
+
+        final int epochOne = 1;
+
+        metadata.update(TestUtils.metadataUpdateWith(null, 1,
+                Collections.emptyMap(), partitionCounts, tp -> epochOne), 0L);
+
+        Node node = metadata.fetch().nodes().get(0);
+        assertFalse(client.isConnected(node.idString()));
+
+        // Seek with a position and leader+epoch
+        Metadata.LeaderAndEpoch leaderAndEpoch = new Metadata.LeaderAndEpoch(
+                metadata.leaderAndEpoch(tp0).leader, Optional.of(epochOne));
+        subscriptions.seekUnvalidated(tp0, new SubscriptionState.FetchPosition(20L, Optional.of(epochOne), leaderAndEpoch));
+        assertFalse(client.isConnected(node.idString()));
+        assertTrue(subscriptions.awaitingValidation(tp0));
+
+        // No version information is initially available, but the node is now connected
+        fetcher.validateOffsetsIfNeeded();
+        assertTrue(subscriptions.awaitingValidation(tp0));
+        assertTrue(client.isConnected(node.idString()));
+        apiVersions.update(node.idString(), NodeApiVersions.create());
+
+        // On the next call, the OffsetForLeaderEpoch request is sent and validation completes
+        Map endOffsetMap = new HashMap<>();
+        endOffsetMap.put(tp0, new EpochEndOffset(Errors.NONE, epochOne, 30L));
+        OffsetsForLeaderEpochResponse resp = new OffsetsForLeaderEpochResponse(endOffsetMap);
+        client.prepareResponseFrom(resp, node);
+
+        fetcher.validateOffsetsIfNeeded();
+        consumerClient.pollNoWakeup();
+
+        assertFalse(subscriptions.awaitingValidation(tp0));
+        assertEquals(20L, subscriptions.position(tp0).offset);
+    }
+
+    @Test
+    public void testOffsetValidationSkippedForOldBroker() {
+        // Old brokers may require CLUSTER permission to use the OffsetForLeaderEpoch API,
+        // so we should skip offset validation and not send the request.
+
+        buildFetcher();
+        assignFromUser(singleton(tp0));
+
+        Map partitionCounts = new HashMap<>();
+        partitionCounts.put(tp0.topic(), 4);
+
+        final int epochOne = 1;
+        final int epochTwo = 2;
+
+        // Start with metadata, epoch=1
+        metadata.update(TestUtils.metadataUpdateWith(null, 1,
+                Collections.emptyMap(), partitionCounts, tp -> epochOne), 0L);
+
+        // Offset validation requires OffsetForLeaderEpoch request v3 or higher
+        Node node = metadata.fetch().nodes().get(0);
+        apiVersions.update(node.idString(), NodeApiVersions.create(
+            ApiKeys.OFFSET_FOR_LEADER_EPOCH.id, (short) 0, (short) 2));
+
+        // Seek with a position and leader+epoch
+        Metadata.LeaderAndEpoch leaderAndEpoch = new Metadata.LeaderAndEpoch(
+                metadata.leaderAndEpoch(tp0).leader, Optional.of(epochOne));
+        subscriptions.seekUnvalidated(tp0, new SubscriptionState.FetchPosition(0, Optional.of(epochOne), leaderAndEpoch));
+
+        // Update metadata to epoch=2, enter validation
+        metadata.update(TestUtils.metadataUpdateWith(null, 1,
+                Collections.emptyMap(), partitionCounts, tp -> epochTwo), 0L);
+        fetcher.validateOffsetsIfNeeded();
+
+        // Offset validation is skipped
+        assertFalse(subscriptions.awaitingValidation(tp0));
+    }
+
+    @Test
+    public void testOffsetValidationHandlesSeekWithInflightOffsetForLeaderRequest() {
+        buildFetcher();
+        assignFromUser(singleton(tp0));
+
+        Map partitionCounts = new HashMap<>();
+        partitionCounts.put(tp0.topic(), 4);
+
+        final int epochOne = 1;
+
+        metadata.update(TestUtils.metadataUpdateWith(null, 1, Collections.emptyMap(), partitionCounts, tp -> epochOne), 0L);
+
+        // Offset validation requires OffsetForLeaderEpoch request v3 or higher
+        Node node = metadata.fetch().nodes().get(0);
+        apiVersions.update(node.idString(), NodeApiVersions.create());
+
+        Metadata.LeaderAndEpoch leaderAndEpoch = new Metadata.LeaderAndEpoch(metadata.leaderAndEpoch(tp0).leader, Optional.of(epochOne));
+        subscriptions.seekUnvalidated(tp0, new SubscriptionState.FetchPosition(0, Optional.of(epochOne), leaderAndEpoch));
+
+        fetcher.validateOffsetsIfNeeded();
+        consumerClient.poll(time.timer(Duration.ZERO));
+        assertTrue(subscriptions.awaitingValidation(tp0));
+        assertTrue(client.hasInFlightRequests());
+
+        // While the OffsetForLeaderEpoch request is in-flight, we seek to a different offset.
+        subscriptions.seekUnvalidated(tp0, new SubscriptionState.FetchPosition(5, Optional.of(epochOne), leaderAndEpoch));
+        assertTrue(subscriptions.awaitingValidation(tp0));
+
+        client.respond(request -> {
+            OffsetsForLeaderEpochRequest epochRequest = (OffsetsForLeaderEpochRequest) request;
+            OffsetsForLeaderEpochRequest.PartitionData partitionData = epochRequest.epochsByTopicPartition().get(tp0);
+            return partitionData.currentLeaderEpoch.equals(Optional.of(epochOne)) && partitionData.leaderEpoch == epochOne;
+        }, new OffsetsForLeaderEpochResponse(singletonMap(tp0, new EpochEndOffset(0, 0L))));
+        consumerClient.poll(time.timer(Duration.ZERO));
+
+        // The response should be ignored since we were validating a different position.
+        assertTrue(subscriptions.awaitingValidation(tp0));
+    }
+
+    @Test
+    public void testOffsetValidationFencing() {
+        buildFetcher();
+        assignFromUser(singleton(tp0));
+
+        Map partitionCounts = new HashMap<>();
+        partitionCounts.put(tp0.topic(), 4);
+
+        final int epochOne = 1;
+        final int epochTwo = 2;
+        final int epochThree = 3;
+
+        // Start with metadata, epoch=1
+        metadata.update(TestUtils.metadataUpdateWith(null, 1, Collections.emptyMap(), partitionCounts, tp -> epochOne), 0L);
+
+        // Offset validation requires OffsetForLeaderEpoch request v3 or higher
+        Node node = metadata.fetch().nodes().get(0);
+        apiVersions.update(node.idString(), NodeApiVersions.create());
+
+        // Seek with a position and leader+epoch
+        Metadata.LeaderAndEpoch leaderAndEpoch = new Metadata.LeaderAndEpoch(metadata.leaderAndEpoch(tp0).leader, Optional.of(epochOne));
+        subscriptions.seekValidated(tp0, new SubscriptionState.FetchPosition(0, Optional.of(epochOne), leaderAndEpoch));
+
+        // Update metadata to epoch=2, enter validation
+        metadata.update(TestUtils.metadataUpdateWith(null, 1, Collections.emptyMap(), partitionCounts, tp -> epochTwo), 0L);
+        fetcher.validateOffsetsIfNeeded();
+        assertTrue(subscriptions.awaitingValidation(tp0));
+
+        // Update the position to epoch=3, as we would from a fetch
+        subscriptions.completeValidation(tp0);
+        SubscriptionState.FetchPosition nextPosition = new SubscriptionState.FetchPosition(
+                10,
+                Optional.of(epochTwo),
+                new Metadata.LeaderAndEpoch(leaderAndEpoch.leader, Optional.of(epochTwo)));
+        subscriptions.position(tp0, nextPosition);
+        subscriptions.maybeValidatePositionForCurrentLeader(tp0, new Metadata.LeaderAndEpoch(leaderAndEpoch.leader, Optional.of(epochThree)));
+
+        // Prepare offset list response from async validation with epoch=2
+        Map endOffsetMap = new HashMap<>();
+        endOffsetMap.put(tp0, new EpochEndOffset(Errors.NONE, epochTwo, 10L));
+        OffsetsForLeaderEpochResponse resp = new OffsetsForLeaderEpochResponse(endOffsetMap);
+        client.prepareResponse(resp);
+        consumerClient.pollNoWakeup();
+        assertTrue("Expected validation to fail since leader epoch changed", subscriptions.awaitingValidation(tp0));
+
+        // Next round of validation, should succeed in validating the position
+        fetcher.validateOffsetsIfNeeded();
+        endOffsetMap.clear();
+        endOffsetMap.put(tp0, new EpochEndOffset(Errors.NONE, epochThree, 10L));
+        resp = new OffsetsForLeaderEpochResponse(endOffsetMap);
+        client.prepareResponse(resp);
+        consumerClient.pollNoWakeup();
+        assertFalse("Expected validation to succeed with latest epoch", subscriptions.awaitingValidation(tp0));
+    }
+
+    @Test
+    public void testTruncationDetected() {
+        // Create some records that include a leader epoch (1)
+        MemoryRecordsBuilder builder = MemoryRecords.builder(
+                ByteBuffer.allocate(1024),
+                RecordBatch.CURRENT_MAGIC_VALUE,
+                CompressionType.NONE,
+                TimestampType.CREATE_TIME,
+                0L,
+                RecordBatch.NO_TIMESTAMP,
+                RecordBatch.NO_PRODUCER_ID,
+                RecordBatch.NO_PRODUCER_EPOCH,
+                RecordBatch.NO_SEQUENCE,
+                false,
+                1 // record epoch is earlier than the leader epoch on the client
+        );
+        builder.appendWithOffset(0L, 0L, "key".getBytes(), "value-1".getBytes());
+        builder.appendWithOffset(1L, 0L, "key".getBytes(), "value-2".getBytes());
+        builder.appendWithOffset(2L, 0L, "key".getBytes(), "value-3".getBytes());
+        MemoryRecords records = builder.build();
+
+        buildFetcher();
+        assignFromUser(singleton(tp0));
+
+        // Initialize the epoch=2
+        Map partitionCounts = new HashMap<>();
+        partitionCounts.put(tp0.topic(), 4);
+        MetadataResponse metadataResponse = TestUtils.metadataUpdateWith(null, 1, Collections.emptyMap(), partitionCounts, tp -> 2);
+        metadata.update(metadataResponse, 0L);
+
+        // Offset validation requires OffsetForLeaderEpoch request v3 or higher
+        Node node = metadata.fetch().nodes().get(0);
+        apiVersions.update(node.idString(), NodeApiVersions.create());
+
+        // Seek
+        Metadata.LeaderAndEpoch leaderAndEpoch = new Metadata.LeaderAndEpoch(metadata.leaderAndEpoch(tp0).leader, Optional.of(1));
+        subscriptions.seekValidated(tp0, new SubscriptionState.FetchPosition(0, Optional.of(1), leaderAndEpoch));
+
+        // Check for truncation, this should cause tp0 to go into validation
+        fetcher.validateOffsetsIfNeeded();
+
+        // No fetches sent since we entered validation
+        assertEquals(0, fetcher.sendFetches());
+        assertFalse(fetcher.hasCompletedFetches());
+        assertTrue(subscriptions.awaitingValidation(tp0));
+
+        // Prepare OffsetForEpoch response then check that we update the subscription position correctly.
+        Map endOffsetMap = new HashMap<>();
+        endOffsetMap.put(tp0, new EpochEndOffset(Errors.NONE, 1, 10L));
+        OffsetsForLeaderEpochResponse resp = new OffsetsForLeaderEpochResponse(endOffsetMap);
+        client.prepareResponse(resp);
+        consumerClient.pollNoWakeup();
+
+        assertFalse(subscriptions.awaitingValidation(tp0));
+
+        // Fetch again, now it works
+        assertEquals(1, fetcher.sendFetches());
+        assertFalse(fetcher.hasCompletedFetches());
+
+        client.prepareResponse(fullFetchResponse(tp0, records, Errors.NONE, 100L, 0));
+        consumerClient.pollNoWakeup();
+        assertTrue(fetcher.hasCompletedFetches());
+
+        Map>> partitionRecords = fetchedRecords();
+        assertTrue(partitionRecords.containsKey(tp0));
+
+        assertEquals(subscriptions.position(tp0).offset, 3L);
+        assertOptional(subscriptions.position(tp0).offsetEpoch, value -> assertEquals(value.intValue(), 1));
+    }
+
+    @Test
+    public void testPreferredReadReplica() {
+        buildFetcher(new MetricConfig(), OffsetResetStrategy.EARLIEST, new BytesDeserializer(), new BytesDeserializer(),
+                Integer.MAX_VALUE, IsolationLevel.READ_COMMITTED, Duration.ofMinutes(5).toMillis());
+
+        subscriptions.assignFromUser(singleton(tp0));
+        client.updateMetadata(TestUtils.metadataUpdateWith(2, singletonMap(topicName, 4)));
+        subscriptions.seek(tp0, 0);
+
+        // Node preferred replica before first fetch response
+        Node selected = fetcher.selectReadReplica(tp0, Node.noNode(), time.milliseconds());
+        assertEquals(selected.id(), -1);
+
+        assertEquals(1, fetcher.sendFetches());
+        assertFalse(fetcher.hasCompletedFetches());
+
+        // Set preferred read replica to node=1
+        client.prepareResponse(fullFetchResponse(tp0, this.records, Errors.NONE, 100L,
+                FetchResponse.INVALID_LAST_STABLE_OFFSET, 0, Optional.of(1)));
+        consumerClient.poll(time.timer(0));
+        assertTrue(fetcher.hasCompletedFetches());
+
+        Map>> partitionRecords = fetchedRecords();
+        assertTrue(partitionRecords.containsKey(tp0));
+
+        // verify
+        selected = fetcher.selectReadReplica(tp0, Node.noNode(), time.milliseconds());
+        assertEquals(selected.id(), 1);
+
+
+        assertEquals(1, fetcher.sendFetches());
+        assertFalse(fetcher.hasCompletedFetches());
+
+        // Set preferred read replica to node=2, which isn't in our metadata, should revert to leader
+        client.prepareResponse(fullFetchResponse(tp0, this.records, Errors.NONE, 100L,
+                FetchResponse.INVALID_LAST_STABLE_OFFSET, 0, Optional.of(2)));
+        consumerClient.poll(time.timer(0));
+        assertTrue(fetcher.hasCompletedFetches());
+        fetchedRecords();
+        selected = fetcher.selectReadReplica(tp0, Node.noNode(), time.milliseconds());
+        assertEquals(selected.id(), -1);
+    }
+
+    @Test
+    public void testPreferredReadReplicaOffsetError() {
+        buildFetcher(new MetricConfig(), OffsetResetStrategy.EARLIEST, new BytesDeserializer(), new BytesDeserializer(),
+                Integer.MAX_VALUE, IsolationLevel.READ_COMMITTED, Duration.ofMinutes(5).toMillis());
+
+        subscriptions.assignFromUser(singleton(tp0));
+        client.updateMetadata(TestUtils.metadataUpdateWith(2, singletonMap(topicName, 4)));
+        subscriptions.seek(tp0, 0);
+
+        assertEquals(1, fetcher.sendFetches());
+        assertFalse(fetcher.hasCompletedFetches());
+
+        client.prepareResponse(fullFetchResponse(tp0, this.records, Errors.NONE, 100L,
+                FetchResponse.INVALID_LAST_STABLE_OFFSET, 0, Optional.of(1)));
+        consumerClient.poll(time.timer(0));
+        assertTrue(fetcher.hasCompletedFetches());
+
+        fetchedRecords();
+
+        Node selected = fetcher.selectReadReplica(tp0, Node.noNode(), time.milliseconds());
+        assertEquals(selected.id(), 1);
+
+        // Return an error, should unset the preferred read replica
+        assertEquals(1, fetcher.sendFetches());
+        assertFalse(fetcher.hasCompletedFetches());
+
+        client.prepareResponse(fullFetchResponse(tp0, this.records, Errors.OFFSET_OUT_OF_RANGE, 100L,
+                FetchResponse.INVALID_LAST_STABLE_OFFSET, 0, Optional.empty()));
+        consumerClient.poll(time.timer(0));
+        assertTrue(fetcher.hasCompletedFetches());
+
+        fetchedRecords();
+
+        selected = fetcher.selectReadReplica(tp0, Node.noNode(), time.milliseconds());
+        assertEquals(selected.id(), -1);
+    }
+
+    @Test
+    public void testFetchCompletedBeforeHandlerAdded() {
+        buildFetcher();
+        assignFromUser(singleton(tp0));
+        subscriptions.seek(tp0, 0);
+        fetcher.sendFetches();
+        client.prepareResponse(fullFetchResponse(tp0, buildRecords(1L, 1, 1), Errors.NONE, 100L, 0));
+        consumerClient.poll(time.timer(0));
+        fetchedRecords();
+        Node node = fetcher.selectReadReplica(tp0, subscriptions.position(tp0).currentLeader.leader, time.milliseconds());
+
+        AtomicBoolean wokenUp = new AtomicBoolean(false);
+        client.setWakeupHook(() -> {
+            if (!wokenUp.getAndSet(true)) {
+                consumerClient.disconnectAsync(node);
+                consumerClient.poll(time.timer(0));
+            }
+        });
+
+        assertEquals(1, fetcher.sendFetches());
+
+        consumerClient.disconnectAsync(node);
+        consumerClient.poll(time.timer(0));
+
+        assertEquals(1, fetcher.sendFetches());
+    }
+
     private MockClient.RequestMatcher listOffsetRequestMatcher(final long timestamp) {
         // matches any list offset request with the provided timestamp
-        return new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                ListOffsetRequest req = (ListOffsetRequest) body;
-                return timestamp == req.partitionTimestamps().get(tp0).timestamp;
-            }
+        return body -> {
+            ListOffsetRequest req = (ListOffsetRequest) body;
+            return req.partitionTimestamps().equals(Collections.singletonMap(
+                tp0, new ListOffsetRequest.PartitionData(timestamp, Optional.empty())));
         };
     }
 
@@ -3140,6 +3955,14 @@ private FetchResponse fullFetchResponse(TopicPartition tp, Memory
         return new FetchResponse<>(Errors.NONE, new LinkedHashMap<>(partitions), throttleTime, INVALID_SESSION_ID);
     }
 
+    private FetchResponse fullFetchResponse(TopicPartition tp, MemoryRecords records, Errors error, long hw,
+                                                           long lastStableOffset, int throttleTime, Optional preferredReplicaId) {
+        Map> partitions = Collections.singletonMap(tp,
+                new FetchResponse.PartitionData<>(error, hw, lastStableOffset, 0L,
+                        preferredReplicaId, null, records));
+        return new FetchResponse<>(Errors.NONE, new LinkedHashMap<>(partitions), throttleTime, INVALID_SESSION_ID);
+    }
+
     private FetchResponse fetchResponse(TopicPartition tp, MemoryRecords records, Errors error, long hw,
                                         long lastStableOffset, long logStartOffset, int throttleTime) {
         Map> partitions = Collections.singletonMap(tp,
@@ -3162,7 +3985,7 @@ private MetadataResponse newMetadataResponse(String topic, Errors error) {
         MetadataResponse.TopicMetadata topicMetadata = new MetadataResponse.TopicMetadata(error, topic, false,
                 partitionsMetadata);
         List brokers = new ArrayList<>(initialUpdateResponse.brokers());
-        return new MetadataResponse(brokers, initialUpdateResponse.clusterId(),
+        return MetadataResponse.prepareResponse(brokers, initialUpdateResponse.clusterId(),
                 initialUpdateResponse.controller().id(), Collections.singletonList(topicMetadata));
     }
 
@@ -3201,7 +4024,36 @@ private  void buildFetcher(MetricConfig metricConfig,
                                      Deserializer valueDeserializer,
                                      int maxPollRecords,
                                      IsolationLevel isolationLevel) {
-        buildDependencies(metricConfig, offsetResetStrategy);
+        buildFetcher(metricConfig, offsetResetStrategy, keyDeserializer, valueDeserializer, maxPollRecords, isolationLevel, Long.MAX_VALUE);
+    }
+
+    private  void buildFetcher(MetricConfig metricConfig,
+                                     OffsetResetStrategy offsetResetStrategy,
+                                     Deserializer keyDeserializer,
+                                     Deserializer valueDeserializer,
+                                     int maxPollRecords,
+                                     IsolationLevel isolationLevel,
+                                     long metadataExpireMs) {
+        LogContext logContext = new LogContext();
+        SubscriptionState subscriptionState = new SubscriptionState(logContext, offsetResetStrategy);
+        buildFetcher(metricConfig, keyDeserializer, valueDeserializer, maxPollRecords, isolationLevel, metadataExpireMs,
+                subscriptionState, logContext);
+    }
+
+    private void buildFetcher(SubscriptionState subscriptionState, LogContext logContext) {
+        buildFetcher(new MetricConfig(), new ByteArrayDeserializer(), new ByteArrayDeserializer(), Integer.MAX_VALUE,
+                IsolationLevel.READ_UNCOMMITTED, Long.MAX_VALUE, subscriptionState, logContext);
+    }
+
+    private  void buildFetcher(MetricConfig metricConfig,
+                                     Deserializer keyDeserializer,
+                                     Deserializer valueDeserializer,
+                                     int maxPollRecords,
+                                     IsolationLevel isolationLevel,
+                                     long metadataExpireMs,
+                                     SubscriptionState subscriptionState,
+                                     LogContext logContext) {
+        buildDependencies(metricConfig, metadataExpireMs, subscriptionState, logContext);
         fetcher = new Fetcher<>(
                 new LogContext(),
                 consumerClient,
@@ -3211,6 +4063,8 @@ private  void buildFetcher(MetricConfig metricConfig,
                 fetchSize,
                 maxPollRecords,
                 true, // check crc
+                "",
+                testPassthrough,
                 keyDeserializer,
                 valueDeserializer,
                 metadata,
@@ -3220,14 +4074,17 @@ private  void buildFetcher(MetricConfig metricConfig,
                 time,
                 retryBackoffMs,
                 requestTimeoutMs,
-                isolationLevel);
+                isolationLevel,
+                apiVersions);
     }
 
-    private void buildDependencies(MetricConfig metricConfig, OffsetResetStrategy offsetResetStrategy) {
-        LogContext logContext = new LogContext();
+    private void buildDependencies(MetricConfig metricConfig,
+                                   long metadataExpireMs,
+                                   SubscriptionState subscriptionState,
+                                   LogContext logContext) {
         time = new MockTime(1);
-        subscriptions = new SubscriptionState(logContext, offsetResetStrategy);
-        metadata = new ConsumerMetadata(0, Long.MAX_VALUE, false,
+        subscriptions = subscriptionState;
+        metadata = new ConsumerMetadata(0, metadataExpireMs, false, false,
                 subscriptions, logContext, new ClusterResourceListeners());
         client = new MockClient(time, metadata);
         metrics = new Metrics(metricConfig, time);
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/HeartbeatTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/HeartbeatTest.java
index c382de610dcb9..b014bec637c1c 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/HeartbeatTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/HeartbeatTest.java
@@ -16,23 +16,38 @@
  */
 package org.apache.kafka.clients.consumer.internals;
 
+import org.apache.kafka.clients.GroupRebalanceConfig;
 import org.apache.kafka.common.utils.MockTime;
 
+import org.junit.Before;
 import org.junit.Test;
 
+import java.util.Optional;
+
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 
 public class HeartbeatTest {
-
     private int sessionTimeoutMs = 300;
     private int heartbeatIntervalMs = 100;
     private int maxPollIntervalMs = 900;
     private long retryBackoffMs = 10L;
     private MockTime time = new MockTime();
-    private Heartbeat heartbeat = new Heartbeat(time, sessionTimeoutMs, heartbeatIntervalMs,
-            maxPollIntervalMs, retryBackoffMs);
+
+    private Heartbeat heartbeat;
+
+    @Before
+    public void setUp() {
+        GroupRebalanceConfig rebalanceConfig = new GroupRebalanceConfig(sessionTimeoutMs,
+                                                                        maxPollIntervalMs,
+                                                                        heartbeatIntervalMs,
+                                                                        "group_id",
+                                                                        Optional.empty(),
+                                                                        retryBackoffMs,
+                                                                        true);
+        heartbeat = new Heartbeat(rebalanceConfig, time);
+    }
 
     @Test
     public void testShouldHeartbeat() {
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MockPartitionAssignor.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MockPartitionAssignor.java
index 609c773ff6ef7..ca7cb34ddd414 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MockPartitionAssignor.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MockPartitionAssignor.java
@@ -23,8 +23,14 @@
 
 public class MockPartitionAssignor extends AbstractPartitionAssignor {
 
+    private final List supportedProtocols;
+
     private Map> result = null;
 
+    MockPartitionAssignor(final List supportedProtocols) {
+        this.supportedProtocols = supportedProtocols;
+    }
+
     @Override
     public Map> assign(Map partitionsPerTopic,
                                                     Map subscriptions) {
@@ -38,6 +44,11 @@ public String name() {
         return "consumer-mock-assignor";
     }
 
+    @Override
+    public List supportedProtocols() {
+        return supportedProtocols;
+    }
+
     public void clear() {
         this.result = null;
     }
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/OffsetForLeaderEpochClientTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/OffsetForLeaderEpochClientTest.java
new file mode 100644
index 0000000000000..55b8754ac9584
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/OffsetForLeaderEpochClientTest.java
@@ -0,0 +1,166 @@
+/*
+ * 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.Metadata;
+import org.apache.kafka.clients.MockClient;
+import org.apache.kafka.clients.consumer.OffsetResetStrategy;
+import org.apache.kafka.common.Node;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.errors.TopicAuthorizationException;
+import org.apache.kafka.common.internals.ClusterResourceListeners;
+import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.common.requests.EpochEndOffset;
+import org.apache.kafka.common.requests.OffsetsForLeaderEpochResponse;
+import org.apache.kafka.common.utils.LogContext;
+import org.apache.kafka.common.utils.MockTime;
+import org.apache.kafka.common.utils.Time;
+import org.junit.Test;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+public class OffsetForLeaderEpochClientTest {
+
+    private ConsumerNetworkClient consumerClient;
+    private SubscriptionState subscriptions;
+    private Metadata metadata;
+    private MockClient client;
+    private Time time;
+
+    private TopicPartition tp0 = new TopicPartition("topic", 0);
+
+    @Test
+    public void testEmptyResponse() {
+        OffsetsForLeaderEpochClient offsetClient = newOffsetClient();
+        RequestFuture future =
+                offsetClient.sendAsyncRequest(Node.noNode(), Collections.emptyMap());
+
+        OffsetsForLeaderEpochResponse resp = new OffsetsForLeaderEpochResponse(Collections.emptyMap());
+        client.prepareResponse(resp);
+        consumerClient.pollNoWakeup();
+
+        OffsetsForLeaderEpochClient.OffsetForEpochResult result = future.value();
+        assertTrue(result.partitionsToRetry().isEmpty());
+        assertTrue(result.endOffsets().isEmpty());
+    }
+
+    @Test
+    public void testUnexpectedEmptyResponse() {
+        Map positionMap = new HashMap<>();
+        positionMap.put(tp0, new SubscriptionState.FetchPosition(0, Optional.of(1),
+                new Metadata.LeaderAndEpoch(Node.noNode(), Optional.of(1))));
+
+        OffsetsForLeaderEpochClient offsetClient = newOffsetClient();
+        RequestFuture future =
+                offsetClient.sendAsyncRequest(Node.noNode(), positionMap);
+
+        OffsetsForLeaderEpochResponse resp = new OffsetsForLeaderEpochResponse(Collections.emptyMap());
+        client.prepareResponse(resp);
+        consumerClient.pollNoWakeup();
+
+        OffsetsForLeaderEpochClient.OffsetForEpochResult result = future.value();
+        assertFalse(result.partitionsToRetry().isEmpty());
+        assertTrue(result.endOffsets().isEmpty());
+    }
+
+    @Test
+    public void testOkResponse() {
+        Map positionMap = new HashMap<>();
+        positionMap.put(tp0, new SubscriptionState.FetchPosition(0, Optional.of(1),
+                new Metadata.LeaderAndEpoch(Node.noNode(), Optional.of(1))));
+
+        OffsetsForLeaderEpochClient offsetClient = newOffsetClient();
+        RequestFuture future =
+                offsetClient.sendAsyncRequest(Node.noNode(), positionMap);
+
+        Map endOffsetMap = new HashMap<>();
+        endOffsetMap.put(tp0, new EpochEndOffset(Errors.NONE, 1, 10L));
+        client.prepareResponse(new OffsetsForLeaderEpochResponse(endOffsetMap));
+        consumerClient.pollNoWakeup();
+
+        OffsetsForLeaderEpochClient.OffsetForEpochResult result = future.value();
+        assertTrue(result.partitionsToRetry().isEmpty());
+        assertTrue(result.endOffsets().containsKey(tp0));
+        assertEquals(result.endOffsets().get(tp0).error(), Errors.NONE);
+        assertEquals(result.endOffsets().get(tp0).leaderEpoch(), 1);
+        assertEquals(result.endOffsets().get(tp0).endOffset(), 10L);
+    }
+
+    @Test
+    public void testUnauthorizedTopic() {
+        Map positionMap = new HashMap<>();
+        positionMap.put(tp0, new SubscriptionState.FetchPosition(0, Optional.of(1),
+                new Metadata.LeaderAndEpoch(Node.noNode(), Optional.of(1))));
+
+        OffsetsForLeaderEpochClient offsetClient = newOffsetClient();
+        RequestFuture future =
+                offsetClient.sendAsyncRequest(Node.noNode(), positionMap);
+
+        Map endOffsetMap = new HashMap<>();
+        endOffsetMap.put(tp0, new EpochEndOffset(Errors.TOPIC_AUTHORIZATION_FAILED, -1, -1));
+        client.prepareResponse(new OffsetsForLeaderEpochResponse(endOffsetMap));
+        consumerClient.pollNoWakeup();
+
+        assertTrue(future.failed());
+        assertEquals(future.exception().getClass(), TopicAuthorizationException.class);
+        assertTrue(((TopicAuthorizationException) future.exception()).unauthorizedTopics().contains(tp0.topic()));
+    }
+
+    @Test
+    public void testRetriableError() {
+        Map positionMap = new HashMap<>();
+        positionMap.put(tp0, new SubscriptionState.FetchPosition(0, Optional.of(1),
+                new Metadata.LeaderAndEpoch(Node.noNode(), Optional.of(1))));
+
+        OffsetsForLeaderEpochClient offsetClient = newOffsetClient();
+        RequestFuture future =
+                offsetClient.sendAsyncRequest(Node.noNode(), positionMap);
+
+        Map endOffsetMap = new HashMap<>();
+        endOffsetMap.put(tp0, new EpochEndOffset(Errors.LEADER_NOT_AVAILABLE, -1, -1));
+        client.prepareResponse(new OffsetsForLeaderEpochResponse(endOffsetMap));
+        consumerClient.pollNoWakeup();
+
+        assertFalse(future.failed());
+        OffsetsForLeaderEpochClient.OffsetForEpochResult result = future.value();
+        assertTrue(result.partitionsToRetry().contains(tp0));
+        assertFalse(result.endOffsets().containsKey(tp0));
+    }
+
+    private OffsetsForLeaderEpochClient newOffsetClient() {
+        buildDependencies(OffsetResetStrategy.EARLIEST);
+        return new OffsetsForLeaderEpochClient(consumerClient, new LogContext());
+    }
+
+    private void buildDependencies(OffsetResetStrategy offsetResetStrategy) {
+        LogContext logContext = new LogContext();
+        time = new MockTime(1);
+        subscriptions = new SubscriptionState(logContext, offsetResetStrategy);
+        metadata = new ConsumerMetadata(0, Long.MAX_VALUE, false, false,
+                subscriptions, logContext, new ClusterResourceListeners());
+        client = new MockClient(time, metadata);
+        consumerClient = new ConsumerNetworkClient(logContext, client, metadata, time,
+                100, 1000, Integer.MAX_VALUE);
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/PartitionAssignorAdapterTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/PartitionAssignorAdapterTest.java
new file mode 100644
index 0000000000000..52950cba4cc6c
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/PartitionAssignorAdapterTest.java
@@ -0,0 +1,173 @@
+/*
+ * 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 static org.apache.kafka.clients.consumer.internals.PartitionAssignorAdapter.getAssignorInstances;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+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 org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerGroupMetadata;
+import org.apache.kafka.clients.consumer.ConsumerPartitionAssignor;
+import org.apache.kafka.clients.consumer.ConsumerPartitionAssignor.Assignment;
+import org.apache.kafka.clients.consumer.ConsumerPartitionAssignor.GroupSubscription;
+import org.apache.kafka.clients.consumer.ConsumerPartitionAssignor.Subscription;
+import org.apache.kafka.clients.consumer.KafkaConsumer;
+import org.apache.kafka.clients.consumer.StickyAssignor;
+import org.apache.kafka.common.Cluster;
+import org.apache.kafka.common.KafkaException;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.StringDeserializer;
+import org.junit.Test;
+
+public class PartitionAssignorAdapterTest {
+
+    private List classNames;
+    private List classTypes;
+
+    @Test
+    public void shouldInstantiateNewAssignors() {
+        classNames = Arrays.asList(StickyAssignor.class.getName());
+        List assignors = getAssignorInstances(classNames, Collections.emptyMap());
+        assertTrue(StickyAssignor.class.isInstance(assignors.get(0)));
+    }
+
+    @Test
+    public void shouldAdaptOldAssignors() {
+        classNames = Arrays.asList(OldPartitionAssignor.class.getName());
+        List assignors = getAssignorInstances(classNames, Collections.emptyMap());
+        assertTrue(PartitionAssignorAdapter.class.isInstance(assignors.get(0)));
+    }
+
+    @Test
+    public void shouldThrowKafkaExceptionOnNonAssignor() {
+        classNames = Arrays.asList(String.class.getName());
+        assertThrows(KafkaException.class, () -> getAssignorInstances(classNames, Collections.emptyMap()));
+    }
+
+    @Test
+    public void shouldThrowKafkaExceptionOnAssignorNotFound() {
+        classNames = Arrays.asList("Non-existent assignor");
+        assertThrows(KafkaException.class, () -> getAssignorInstances(classNames, Collections.emptyMap()));
+    }
+
+    @Test
+    public void shouldInstantiateFromListOfOldAndNewClassTypes() {
+        Properties props = new Properties();
+        props.setProperty(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
+
+        classTypes = Arrays.asList(StickyAssignor.class, OldPartitionAssignor.class);
+
+        props.put(ConsumerConfig.PARTITION_ASSIGNMENT_STRATEGY_CONFIG, classTypes);
+        KafkaConsumer consumer = new KafkaConsumer<>(
+            props, new StringDeserializer(), new StringDeserializer());
+
+        consumer.close();
+    }
+
+    @Test
+    public void shouldThrowKafkaExceptionOnListWithNonAssignorClassType() {
+        Properties props = new Properties();
+        props.setProperty(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
+
+        classTypes = Arrays.asList(StickyAssignor.class, OldPartitionAssignor.class, String.class);
+
+        props.put(ConsumerConfig.PARTITION_ASSIGNMENT_STRATEGY_CONFIG, classTypes);
+        assertThrows(KafkaException.class, () -> new KafkaConsumer<>(
+            props, new StringDeserializer(), new StringDeserializer()));
+    }
+
+    @Test
+    @SuppressWarnings("deprecation")
+    public void testOnAssignment() {
+        OldPartitionAssignor oldAssignor = new OldPartitionAssignor();
+        ConsumerPartitionAssignor adaptedAssignor = new PartitionAssignorAdapter(oldAssignor);
+
+        TopicPartition tp1 = new TopicPartition("tp1", 1);
+        TopicPartition tp2 = new TopicPartition("tp2", 2);
+        List partitions = Arrays.asList(tp1, tp2);
+
+        adaptedAssignor.onAssignment(new Assignment(partitions), new ConsumerGroupMetadata("", 1, "", Optional.empty()));
+
+        assertEquals(oldAssignor.partitions, partitions);
+    }
+
+    @Test
+    public void testAssign() {
+        ConsumerPartitionAssignor adaptedAssignor = new PartitionAssignorAdapter(new OldPartitionAssignor());
+
+        Map subscriptions = new HashMap<>();
+        subscriptions.put("C1", new Subscription(Arrays.asList("topic1")));
+        subscriptions.put("C2", new Subscription(Arrays.asList("topic1", "topic2")));
+        subscriptions.put("C3", new Subscription(Arrays.asList("topic2", "topic3")));
+        GroupSubscription groupSubscription = new GroupSubscription(subscriptions);
+
+        Map assignments = adaptedAssignor.assign(null, groupSubscription).groupAssignment();
+
+        assertEquals(assignments.get("C1").partitions(), Arrays.asList(new TopicPartition("topic1", 1)));
+        assertEquals(assignments.get("C2").partitions(), Arrays.asList(new TopicPartition("topic1", 1), new TopicPartition("topic2", 1)));
+        assertEquals(assignments.get("C3").partitions(), Arrays.asList(new TopicPartition("topic2", 1), new TopicPartition("topic3", 1)));
+    }
+
+    /*
+     * Dummy assignor just gives each consumer partition 1 of each topic it's subscribed to
+     */
+    @SuppressWarnings("deprecation")
+    public static class OldPartitionAssignor implements PartitionAssignor {
+
+        List partitions = null;
+
+        @Override
+        public Subscription subscription(Set topics) {
+            return new Subscription(new ArrayList<>(topics), null);
+        }
+
+        @Override
+        public Map assign(Cluster metadata, Map subscriptions) {
+            Map assignments = new HashMap<>();
+            for (Map.Entry entry : subscriptions.entrySet()) {
+                List partitions = new ArrayList<>();
+                for (String topic : entry.getValue().topics()) {
+                    partitions.add(new TopicPartition(topic, 1));
+                }
+                assignments.put(entry.getKey(), new Assignment(partitions, null));
+            }
+            return assignments;
+        }
+
+        @Override
+        public void onAssignment(Assignment assignment) {
+            partitions = assignment.partitions();
+        }
+
+        @Override
+        public String name() {
+            return "old-assignor";
+        }
+    }
+
+}
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 8f8e96068dc65..58d8b888822ec 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
@@ -16,37 +16,42 @@
  */
 package org.apache.kafka.clients.consumer.internals;
 
+import org.apache.kafka.clients.Metadata;
 import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
+import org.apache.kafka.clients.consumer.OffsetAndMetadata;
 import org.apache.kafka.clients.consumer.OffsetResetStrategy;
+import org.apache.kafka.common.Node;
 import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.requests.EpochEndOffset;
 import org.apache.kafka.common.utils.LogContext;
 import org.apache.kafka.common.utils.Utils;
+import org.apache.kafka.test.TestUtils;
 import org.junit.Test;
 
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashSet;
+import java.util.Optional;
 import java.util.Set;
-import java.util.concurrent.atomic.AtomicReference;
 import java.util.regex.Pattern;
 
 import static java.util.Collections.singleton;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
 
 public class SubscriptionStateTest {
 
-    private final SubscriptionState state = new SubscriptionState(
-            new LogContext(),
-            OffsetResetStrategy.EARLIEST);
+    private SubscriptionState state = new SubscriptionState(new LogContext(), OffsetResetStrategy.EARLIEST);
     private final String topic = "test";
     private final String topic1 = "test1";
     private final TopicPartition tp0 = new TopicPartition(topic, 0);
     private final TopicPartition tp1 = new TopicPartition(topic, 1);
     private final TopicPartition t1p0 = new TopicPartition(topic1, 0);
     private final MockRebalanceListener rebalanceListener = new MockRebalanceListener();
+    private final Metadata.LeaderAndEpoch leaderAndEpoch = new Metadata.LeaderAndEpoch(Node.noNode(), Optional.empty());
 
     @Test
     public void partitionAssignment() {
@@ -56,8 +61,8 @@ public void partitionAssignment() {
         assertFalse(state.hasAllFetchPositions());
         state.seek(tp0, 1);
         assertTrue(state.isFetchable(tp0));
-        assertEquals(1L, state.position(tp0).longValue());
-        state.assignFromUser(Collections.emptySet());
+        assertEquals(1L, state.position(tp0).offset);
+        state.assignFromUser(Collections.emptySet());
         assertTrue(state.assignedPartitions().isEmpty());
         assertEquals(0, state.numAssignedPartitions());
         assertFalse(state.isAssigned(tp0));
@@ -83,7 +88,8 @@ public void partitionAssignmentChangeOnTopicSubscription() {
         assertTrue(state.assignedPartitions().isEmpty());
         assertEquals(0, state.numAssignedPartitions());
 
-        assertTrue(state.assignFromSubscribed(singleton(t1p0)));
+        assertTrue(state.checkAssignmentMatchedSubscription(singleton(t1p0)));
+        state.assignFromSubscribed(singleton(t1p0));
         // assigned partitions should immediately change
         assertEquals(singleton(t1p0), state.assignedPartitions());
         assertEquals(1, state.numAssignedPartitions());
@@ -99,6 +105,28 @@ public void partitionAssignmentChangeOnTopicSubscription() {
         assertEquals(0, state.numAssignedPartitions());
     }
 
+    @Test
+    public void testGroupSubscribe() {
+        state.subscribe(singleton(topic1), rebalanceListener);
+        assertEquals(singleton(topic1), state.metadataTopics());
+
+        assertFalse(state.groupSubscribe(singleton(topic1)));
+        assertEquals(singleton(topic1), state.metadataTopics());
+
+        assertTrue(state.groupSubscribe(Utils.mkSet(topic, topic1)));
+        assertEquals(Utils.mkSet(topic, topic1), state.metadataTopics());
+
+        // `groupSubscribe` does not accumulate
+        assertFalse(state.groupSubscribe(singleton(topic1)));
+        assertEquals(singleton(topic1), state.metadataTopics());
+
+        state.subscribe(singleton("anotherTopic"), rebalanceListener);
+        assertEquals(Utils.mkSet(topic1, "anotherTopic"), state.metadataTopics());
+
+        assertFalse(state.groupSubscribe(singleton("anotherTopic")));
+        assertEquals(singleton("anotherTopic"), state.metadataTopics());
+    }
+
     @Test
     public void partitionAssignmentChangeOnPatternSubscription() {
         state.subscribe(Pattern.compile(".*"), rebalanceListener);
@@ -106,18 +134,22 @@ public void partitionAssignmentChangeOnPatternSubscription() {
         assertTrue(state.assignedPartitions().isEmpty());
         assertEquals(0, state.numAssignedPartitions());
 
-        state.subscribeFromPattern(new HashSet<>(Collections.singletonList(topic)));
+        state.subscribeFromPattern(Collections.singleton(topic));
         // assigned partitions should remain unchanged
         assertTrue(state.assignedPartitions().isEmpty());
         assertEquals(0, state.numAssignedPartitions());
 
-        assertTrue(state.assignFromSubscribed(singleton(tp1)));
+        assertTrue(state.checkAssignmentMatchedSubscription(singleton(tp1)));
+        state.assignFromSubscribed(singleton(tp1));
+
         // assigned partitions should immediately change
         assertEquals(singleton(tp1), state.assignedPartitions());
         assertEquals(1, state.numAssignedPartitions());
         assertEquals(singleton(topic), state.subscription());
 
-        assertTrue(state.assignFromSubscribed(Collections.singletonList(t1p0)));
+        assertTrue(state.checkAssignmentMatchedSubscription(singleton(t1p0)));
+        state.assignFromSubscribed(singleton(t1p0));
+
         // assigned partitions should immediately change
         assertEquals(singleton(t1p0), state.assignedPartitions());
         assertEquals(1, state.numAssignedPartitions());
@@ -133,7 +165,9 @@ public void partitionAssignmentChangeOnPatternSubscription() {
         assertEquals(singleton(t1p0), state.assignedPartitions());
         assertEquals(1, state.numAssignedPartitions());
 
-        assertTrue(state.assignFromSubscribed(Collections.singletonList(tp0)));
+        assertTrue(state.checkAssignmentMatchedSubscription(singleton(tp0)));
+        state.assignFromSubscribed(singleton(tp0));
+
         // assigned partitions should immediately change
         assertEquals(singleton(tp0), state.assignedPartitions());
         assertEquals(1, state.numAssignedPartitions());
@@ -146,36 +180,34 @@ public void partitionAssignmentChangeOnPatternSubscription() {
     }
 
     @Test
-    public void verifyAssignmentListener() {
-        final AtomicReference> assignmentRef = new AtomicReference<>();
-        state.addListener(new SubscriptionState.Listener() {
-            @Override
-            public void onAssignment(Set assignment) {
-                assignmentRef.set(assignment);
-            }
-        });
+    public void verifyAssignmentId() {
+        assertEquals(0, state.assignmentId());
         Set userAssignment = Utils.mkSet(tp0, tp1);
         state.assignFromUser(userAssignment);
-        assertEquals(userAssignment, assignmentRef.get());
+        assertEquals(1, state.assignmentId());
+        assertEquals(userAssignment, state.assignedPartitions());
 
         state.unsubscribe();
-        assertEquals(Collections.emptySet(), assignmentRef.get());
+        assertEquals(2, state.assignmentId());
+        assertEquals(Collections.emptySet(), state.assignedPartitions());
 
         Set autoAssignment = Utils.mkSet(t1p0);
         state.subscribe(singleton(topic1), rebalanceListener);
-        assertTrue(state.assignFromSubscribed(autoAssignment));
-        assertEquals(autoAssignment, assignmentRef.get());
+        assertTrue(state.checkAssignmentMatchedSubscription(autoAssignment));
+        state.assignFromSubscribed(autoAssignment);
+        assertEquals(3, state.assignmentId());
+        assertEquals(autoAssignment, state.assignedPartitions());
     }
 
     @Test
     public void partitionReset() {
         state.assignFromUser(singleton(tp0));
         state.seek(tp0, 5);
-        assertEquals(5L, (long) state.position(tp0));
+        assertEquals(5L, state.position(tp0).offset);
         state.requestOffsetReset(tp0);
         assertFalse(state.isFetchable(tp0));
         assertTrue(state.isOffsetResetNeeded(tp0));
-        assertEquals(null, state.position(tp0));
+        assertNotNull(state.position(tp0));
 
         // seek should clear the reset and make the partition fetchable
         state.seek(tp0, 0);
@@ -190,10 +222,14 @@ public void topicSubscription() {
         assertTrue(state.assignedPartitions().isEmpty());
         assertEquals(0, state.numAssignedPartitions());
         assertTrue(state.partitionsAutoAssigned());
-        assertTrue(state.assignFromSubscribed(singleton(tp0)));
+        assertTrue(state.checkAssignmentMatchedSubscription(singleton(tp0)));
+        state.assignFromSubscribed(singleton(tp0));
+
         state.seek(tp0, 1);
-        assertEquals(1L, state.position(tp0).longValue());
-        assertTrue(state.assignFromSubscribed(singleton(tp1)));
+        assertEquals(1L, state.position(tp0).offset);
+        assertTrue(state.checkAssignmentMatchedSubscription(singleton(tp1)));
+        state.assignFromSubscribed(singleton(tp1));
+
         assertTrue(state.isAssigned(tp1));
         assertFalse(state.isAssigned(tp0));
         assertFalse(state.isFetchable(tp1));
@@ -215,26 +251,28 @@ public void partitionPause() {
     @Test(expected = IllegalStateException.class)
     public void invalidPositionUpdate() {
         state.subscribe(singleton(topic), rebalanceListener);
-        assertTrue(state.assignFromSubscribed(singleton(tp0)));
-        state.position(tp0, 0);
+        assertTrue(state.checkAssignmentMatchedSubscription(singleton(tp0)));
+        state.assignFromSubscribed(singleton(tp0));
+
+        state.position(tp0, new SubscriptionState.FetchPosition(0, Optional.empty(), leaderAndEpoch));
     }
 
     @Test
     public void cantAssignPartitionForUnsubscribedTopics() {
         state.subscribe(singleton(topic), rebalanceListener);
-        assertFalse(state.assignFromSubscribed(Collections.singletonList(t1p0)));
+        assertFalse(state.checkAssignmentMatchedSubscription(Collections.singletonList(t1p0)));
     }
 
     @Test
     public void cantAssignPartitionForUnmatchedPattern() {
         state.subscribe(Pattern.compile(".*t"), rebalanceListener);
-        state.subscribeFromPattern(new HashSet<>(Collections.singletonList(topic)));
-        assertFalse(state.assignFromSubscribed(Collections.singletonList(t1p0)));
+        state.subscribeFromPattern(Collections.singleton(topic));
+        assertFalse(state.checkAssignmentMatchedSubscription(Collections.singletonList(t1p0)));
     }
 
     @Test(expected = IllegalStateException.class)
     public void cantChangePositionForNonAssignedPartition() {
-        state.position(tp0, 1);
+        state.position(tp0, new SubscriptionState.FetchPosition(1, Optional.empty(), leaderAndEpoch));
     }
 
     @Test(expected = IllegalStateException.class)
@@ -289,7 +327,9 @@ public void unsubscribeUserSubscribe() {
     public void unsubscription() {
         state.subscribe(Pattern.compile(".*"), rebalanceListener);
         state.subscribeFromPattern(new HashSet<>(Arrays.asList(topic, topic1)));
-        assertTrue(state.assignFromSubscribed(singleton(tp1)));
+        assertTrue(state.checkAssignmentMatchedSubscription(singleton(tp1)));
+        state.assignFromSubscribed(singleton(tp1));
+
         assertEquals(singleton(tp1), state.assignedPartitions());
         assertEquals(1, state.numAssignedPartitions());
 
@@ -308,13 +348,275 @@ public void unsubscription() {
         assertEquals(0, state.numAssignedPartitions());
     }
 
+    @Test
+    public void testPreferredReadReplicaLease() {
+        state.assignFromUser(Collections.singleton(tp0));
+
+        // Default state
+        assertFalse(state.preferredReadReplica(tp0, 0L).isPresent());
+
+        // Set the preferred replica with lease
+        state.updatePreferredReadReplica(tp0, 42, () -> 10L);
+        TestUtils.assertOptional(state.preferredReadReplica(tp0, 9L),  value -> assertEquals(value.intValue(), 42));
+        TestUtils.assertOptional(state.preferredReadReplica(tp0, 10L),  value -> assertEquals(value.intValue(), 42));
+        assertFalse(state.preferredReadReplica(tp0, 11L).isPresent());
+
+        // Unset the preferred replica
+        state.clearPreferredReadReplica(tp0);
+        assertFalse(state.preferredReadReplica(tp0, 9L).isPresent());
+        assertFalse(state.preferredReadReplica(tp0, 11L).isPresent());
+
+        // Set to new preferred replica with lease
+        state.updatePreferredReadReplica(tp0, 43, () -> 20L);
+        TestUtils.assertOptional(state.preferredReadReplica(tp0, 11L),  value -> assertEquals(value.intValue(), 43));
+        TestUtils.assertOptional(state.preferredReadReplica(tp0, 20L),  value -> assertEquals(value.intValue(), 43));
+        assertFalse(state.preferredReadReplica(tp0, 21L).isPresent());
+
+        // Set to new preferred replica without clearing first
+        state.updatePreferredReadReplica(tp0, 44, () -> 30L);
+        TestUtils.assertOptional(state.preferredReadReplica(tp0, 30L),  value -> assertEquals(value.intValue(), 44));
+        assertFalse(state.preferredReadReplica(tp0, 31L).isPresent());
+    }
+
+    @Test
+    public void testSeekUnvalidatedWithNoOffsetEpoch() {
+        Node broker1 = new Node(1, "localhost", 9092);
+        state.assignFromUser(Collections.singleton(tp0));
+
+        // Seek with no offset epoch requires no validation no matter what the current leader is
+        state.seekUnvalidated(tp0, new SubscriptionState.FetchPosition(0L, Optional.empty(),
+                new Metadata.LeaderAndEpoch(broker1, Optional.of(5))));
+        assertTrue(state.hasValidPosition(tp0));
+        assertFalse(state.awaitingValidation(tp0));
+
+        assertFalse(state.maybeValidatePositionForCurrentLeader(tp0, new Metadata.LeaderAndEpoch(broker1, Optional.empty())));
+        assertTrue(state.hasValidPosition(tp0));
+        assertFalse(state.awaitingValidation(tp0));
+
+        assertFalse(state.maybeValidatePositionForCurrentLeader(tp0, new Metadata.LeaderAndEpoch(broker1, Optional.of(10))));
+        assertTrue(state.hasValidPosition(tp0));
+        assertFalse(state.awaitingValidation(tp0));
+    }
+
+    @Test
+    public void testSeekUnvalidatedWithNoEpochClearsAwaitingValidation() {
+        Node broker1 = new Node(1, "localhost", 9092);
+        state.assignFromUser(Collections.singleton(tp0));
+
+        // Seek with no offset epoch requires no validation no matter what the current leader is
+        state.seekUnvalidated(tp0, new SubscriptionState.FetchPosition(0L, Optional.of(2),
+                new Metadata.LeaderAndEpoch(broker1, Optional.of(5))));
+        assertFalse(state.hasValidPosition(tp0));
+        assertTrue(state.awaitingValidation(tp0));
+
+        state.seekUnvalidated(tp0, new SubscriptionState.FetchPosition(0L, Optional.empty(),
+                new Metadata.LeaderAndEpoch(broker1, Optional.of(5))));
+        assertTrue(state.hasValidPosition(tp0));
+        assertFalse(state.awaitingValidation(tp0));
+    }
+
+    @Test
+    public void testSeekUnvalidatedWithOffsetEpoch() {
+        Node broker1 = new Node(1, "localhost", 9092);
+        state.assignFromUser(Collections.singleton(tp0));
+
+        state.seekUnvalidated(tp0, new SubscriptionState.FetchPosition(0L, Optional.of(2),
+                new Metadata.LeaderAndEpoch(broker1, Optional.of(5))));
+        assertFalse(state.hasValidPosition(tp0));
+        assertTrue(state.awaitingValidation(tp0));
+
+        // Update using the current leader and epoch
+        assertTrue(state.maybeValidatePositionForCurrentLeader(tp0, new Metadata.LeaderAndEpoch(broker1, Optional.of(5))));
+        assertFalse(state.hasValidPosition(tp0));
+        assertTrue(state.awaitingValidation(tp0));
+
+        // Update with a newer leader and epoch
+        assertTrue(state.maybeValidatePositionForCurrentLeader(tp0, new Metadata.LeaderAndEpoch(broker1, Optional.of(15))));
+        assertFalse(state.hasValidPosition(tp0));
+        assertTrue(state.awaitingValidation(tp0));
+
+        // If the updated leader has no epoch information, then skip validation and begin fetching
+        assertFalse(state.maybeValidatePositionForCurrentLeader(tp0, new Metadata.LeaderAndEpoch(broker1, Optional.empty())));
+        assertTrue(state.hasValidPosition(tp0));
+        assertFalse(state.awaitingValidation(tp0));
+    }
+
+    @Test
+    public void testSeekValidatedShouldClearAwaitingValidation() {
+        Node broker1 = new Node(1, "localhost", 9092);
+        state.assignFromUser(Collections.singleton(tp0));
+
+        state.seekUnvalidated(tp0, new SubscriptionState.FetchPosition(10L, Optional.of(5),
+                new Metadata.LeaderAndEpoch(broker1, Optional.of(10))));
+        assertFalse(state.hasValidPosition(tp0));
+        assertTrue(state.awaitingValidation(tp0));
+        assertEquals(10L, state.position(tp0).offset);
+
+        state.seekValidated(tp0, new SubscriptionState.FetchPosition(8L, Optional.of(4),
+                new Metadata.LeaderAndEpoch(broker1, Optional.of(10))));
+        assertTrue(state.hasValidPosition(tp0));
+        assertFalse(state.awaitingValidation(tp0));
+        assertEquals(8L, state.position(tp0).offset);
+    }
+
+    @Test
+    public void testCompleteValidationShouldClearAwaitingValidation() {
+        Node broker1 = new Node(1, "localhost", 9092);
+        state.assignFromUser(Collections.singleton(tp0));
+
+        state.seekUnvalidated(tp0, new SubscriptionState.FetchPosition(10L, Optional.of(5),
+                new Metadata.LeaderAndEpoch(broker1, Optional.of(10))));
+        assertFalse(state.hasValidPosition(tp0));
+        assertTrue(state.awaitingValidation(tp0));
+        assertEquals(10L, state.position(tp0).offset);
+
+        state.completeValidation(tp0);
+        assertTrue(state.hasValidPosition(tp0));
+        assertFalse(state.awaitingValidation(tp0));
+        assertEquals(10L, state.position(tp0).offset);
+    }
+
+    @Test
+    public void testOffsetResetWhileAwaitingValidation() {
+        Node broker1 = new Node(1, "localhost", 9092);
+        state.assignFromUser(Collections.singleton(tp0));
+
+        state.seekUnvalidated(tp0, new SubscriptionState.FetchPosition(10L, Optional.of(5),
+                new Metadata.LeaderAndEpoch(broker1, Optional.of(10))));
+        assertTrue(state.awaitingValidation(tp0));
+
+        state.requestOffsetReset(tp0, OffsetResetStrategy.EARLIEST);
+        assertFalse(state.awaitingValidation(tp0));
+        assertTrue(state.isOffsetResetNeeded(tp0));
+    }
+
+    @Test
+    public void testMaybeCompleteValidation() {
+        Node broker1 = new Node(1, "localhost", 9092);
+        state.assignFromUser(Collections.singleton(tp0));
+
+        int currentEpoch = 10;
+        long initialOffset = 10L;
+        int initialOffsetEpoch = 5;
+
+        SubscriptionState.FetchPosition initialPosition = new SubscriptionState.FetchPosition(initialOffset,
+                Optional.of(initialOffsetEpoch), new Metadata.LeaderAndEpoch(broker1, Optional.of(currentEpoch)));
+        state.seekUnvalidated(tp0, initialPosition);
+        assertTrue(state.awaitingValidation(tp0));
+
+        Optional divergentOffsetMetadataOpt = state.maybeCompleteValidation(tp0, initialPosition,
+                new EpochEndOffset(initialOffsetEpoch, initialOffset + 5));
+        assertEquals(Optional.empty(), divergentOffsetMetadataOpt);
+        assertFalse(state.awaitingValidation(tp0));
+        assertEquals(initialPosition, state.position(tp0));
+    }
+
+    @Test
+    public void testMaybeCompleteValidationAfterPositionChange() {
+        Node broker1 = new Node(1, "localhost", 9092);
+        state.assignFromUser(Collections.singleton(tp0));
+
+        int currentEpoch = 10;
+        long initialOffset = 10L;
+        int initialOffsetEpoch = 5;
+        long updateOffset = 20L;
+        int updateOffsetEpoch = 8;
+
+        SubscriptionState.FetchPosition initialPosition = new SubscriptionState.FetchPosition(initialOffset,
+                Optional.of(initialOffsetEpoch), new Metadata.LeaderAndEpoch(broker1, Optional.of(currentEpoch)));
+        state.seekUnvalidated(tp0, initialPosition);
+        assertTrue(state.awaitingValidation(tp0));
+
+        SubscriptionState.FetchPosition updatePosition = new SubscriptionState.FetchPosition(updateOffset,
+                Optional.of(updateOffsetEpoch), new Metadata.LeaderAndEpoch(broker1, Optional.of(currentEpoch)));
+        state.seekUnvalidated(tp0, updatePosition);
+
+        Optional divergentOffsetMetadataOpt = state.maybeCompleteValidation(tp0, initialPosition,
+                new EpochEndOffset(initialOffsetEpoch, initialOffset + 5));
+        assertEquals(Optional.empty(), divergentOffsetMetadataOpt);
+        assertTrue(state.awaitingValidation(tp0));
+        assertEquals(updatePosition, state.position(tp0));
+    }
+
+    @Test
+    public void testMaybeCompleteValidationAfterOffsetReset() {
+        Node broker1 = new Node(1, "localhost", 9092);
+        state.assignFromUser(Collections.singleton(tp0));
+
+        int currentEpoch = 10;
+        long initialOffset = 10L;
+        int initialOffsetEpoch = 5;
+
+        SubscriptionState.FetchPosition initialPosition = new SubscriptionState.FetchPosition(initialOffset,
+                Optional.of(initialOffsetEpoch), new Metadata.LeaderAndEpoch(broker1, Optional.of(currentEpoch)));
+        state.seekUnvalidated(tp0, initialPosition);
+        assertTrue(state.awaitingValidation(tp0));
+
+        state.requestOffsetReset(tp0);
+
+        Optional divergentOffsetMetadataOpt = state.maybeCompleteValidation(tp0, initialPosition,
+                new EpochEndOffset(initialOffsetEpoch, initialOffset + 5));
+        assertEquals(Optional.empty(), divergentOffsetMetadataOpt);
+        assertFalse(state.awaitingValidation(tp0));
+        assertTrue(state.isOffsetResetNeeded(tp0));
+    }
+
+    @Test
+    public void testTruncationDetectionWithResetPolicy() {
+        Node broker1 = new Node(1, "localhost", 9092);
+        state.assignFromUser(Collections.singleton(tp0));
+
+        int currentEpoch = 10;
+        long initialOffset = 10L;
+        int initialOffsetEpoch = 5;
+        long divergentOffset = 5L;
+        int divergentOffsetEpoch = 7;
+
+        SubscriptionState.FetchPosition initialPosition = new SubscriptionState.FetchPosition(initialOffset,
+                Optional.of(initialOffsetEpoch), new Metadata.LeaderAndEpoch(broker1, Optional.of(currentEpoch)));
+        state.seekUnvalidated(tp0, initialPosition);
+        assertTrue(state.awaitingValidation(tp0));
+
+        Optional divergentOffsetMetadata = state.maybeCompleteValidation(tp0, initialPosition,
+                new EpochEndOffset(divergentOffsetEpoch, divergentOffset));
+        assertEquals(Optional.empty(), divergentOffsetMetadata);
+        assertFalse(state.awaitingValidation(tp0));
+
+        SubscriptionState.FetchPosition updatedPosition = new SubscriptionState.FetchPosition(divergentOffset,
+                Optional.of(divergentOffsetEpoch), new Metadata.LeaderAndEpoch(broker1, Optional.of(currentEpoch)));
+        assertEquals(updatedPosition, state.position(tp0));
+    }
+
+    @Test
+    public void testTruncationDetectionWithoutResetPolicy() {
+        Node broker1 = new Node(1, "localhost", 9092);
+        state = new SubscriptionState(new LogContext(), OffsetResetStrategy.NONE);
+        state.assignFromUser(Collections.singleton(tp0));
+
+        int currentEpoch = 10;
+        long initialOffset = 10L;
+        int initialOffsetEpoch = 5;
+        long divergentOffset = 5L;
+        int divergentOffsetEpoch = 7;
+
+        SubscriptionState.FetchPosition initialPosition = new SubscriptionState.FetchPosition(initialOffset,
+                Optional.of(initialOffsetEpoch), new Metadata.LeaderAndEpoch(broker1, Optional.of(currentEpoch)));
+        state.seekUnvalidated(tp0, initialPosition);
+        assertTrue(state.awaitingValidation(tp0));
+
+        Optional divergentOffsetMetadata = state.maybeCompleteValidation(tp0, initialPosition,
+                new EpochEndOffset(divergentOffsetEpoch, divergentOffset));
+        assertEquals(Optional.of(new OffsetAndMetadata(divergentOffset, Optional.of(divergentOffsetEpoch), "")),
+                divergentOffsetMetadata);
+        assertTrue(state.awaitingValidation(tp0));
+    }
+
     private static class MockRebalanceListener implements ConsumerRebalanceListener {
         public Collection revoked;
         public Collection assigned;
         public int revokedCount = 0;
         public int assignedCount = 0;
 
-
         @Override
         public void onPartitionsAssigned(Collection partitions) {
             this.assigned = partitions;
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 638cb7b130702..a209a923b36c0 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
@@ -32,9 +32,12 @@
 import org.apache.kafka.common.errors.TimeoutException;
 import org.apache.kafka.common.header.internals.RecordHeader;
 import org.apache.kafka.common.internals.ClusterResourceListeners;
+import org.apache.kafka.common.message.InitProducerIdResponseData;
 import org.apache.kafka.common.metrics.Sensor;
 import org.apache.kafka.common.network.Selectable;
 import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.common.requests.FindCoordinatorResponse;
+import org.apache.kafka.common.requests.InitProducerIdResponse;
 import org.apache.kafka.common.requests.MetadataResponse;
 import org.apache.kafka.common.serialization.ByteArraySerializer;
 import org.apache.kafka.common.serialization.Serializer;
@@ -59,6 +62,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Properties;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
@@ -72,6 +76,7 @@
 import static java.util.Collections.singletonMap;
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 import static org.mockito.ArgumentMatchers.any;
@@ -86,20 +91,21 @@
 
 public class KafkaProducerTest {
     private String topic = "topic";
-    private Collection nodes = Collections.singletonList(new Node(0, "host1", 1000));
+    private Node host1 = new Node(0, "host1", 1000);
+    private Collection nodes = Collections.singletonList(host1);
     private final Cluster emptyCluster = new Cluster(null, nodes,
             Collections.emptySet(),
             Collections.emptySet(),
             Collections.emptySet());
     private final Cluster onePartitionCluster = new Cluster(
             "dummy",
-            Collections.singletonList(new Node(0, "host1", 1000)),
+            Collections.singletonList(host1),
             Collections.singletonList(new PartitionInfo(topic, 0, null, null, null)),
             Collections.emptySet(),
             Collections.emptySet());
     private final Cluster threePartitionCluster = new Cluster(
             "dummy",
-            Collections.singletonList(new Node(0, "host1", 1000)),
+            Collections.singletonList(host1),
             Arrays.asList(
                     new PartitionInfo(topic, 0, null, null, null),
                     new PartitionInfo(topic, 1, null, null, null),
@@ -497,12 +503,7 @@ public void testTopicRefreshInMetadata() throws InterruptedException {
                 }
             });
             t.start();
-            try {
-                producer.partitionsFor(topic);
-                fail("Expect TimeoutException");
-            } catch (TimeoutException e) {
-                // skip
-            }
+            assertThrows(TimeoutException.class, () -> producer.partitionsFor(topic));
             running.set(false);
             t.join();
         }
@@ -553,12 +554,7 @@ private > void doTestHeaders(Class serializerCla
         producer.send(record, null);
 
         //ensure headers are closed and cannot be mutated post send
-        try {
-            record.headers().add(new RecordHeader("test", "test".getBytes()));
-            fail("Expected IllegalStateException to be raised");
-        } catch (IllegalStateException ise) {
-            //expected
-        }
+        assertThrows(IllegalStateException.class, () -> record.headers().add(new RecordHeader("test", "test".getBytes())));
 
         //ensure existing headers are not changed, and last header for key is still original value
         assertArrayEquals(record.headers().lastHeader("test").value(), "header2".getBytes());
@@ -625,14 +621,11 @@ public void testPartitionsForWithNullTopic() {
         Properties props = new Properties();
         props.setProperty(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9000");
         try (KafkaProducer producer = new KafkaProducer<>(props, new ByteArraySerializer(), new ByteArraySerializer())) {
-            producer.partitionsFor(null);
-            fail("Expected NullPointerException to be raised");
-        } catch (NullPointerException e) {
-            // expected
+            assertThrows(NullPointerException.class, () -> producer.partitionsFor(null));
         }
     }
 
-    @Test(expected = TimeoutException.class)
+    @Test
     public void testInitTransactionTimeout() {
         Map configs = new HashMap<>();
         configs.put(ProducerConfig.TRANSACTIONAL_ID_CONFIG, "bad-transaction");
@@ -646,14 +639,49 @@ public void testInitTransactionTimeout() {
 
         MockClient client = new MockClient(time, metadata);
 
+        try (Producer producer = new KafkaProducer<>(configs, new StringSerializer(),
+                new StringSerializer(), metadata, client, null, time)) {
+            assertThrows(TimeoutException.class, producer::initTransactions);
+        }
+    }
+
+    @Test
+    public void testInitTransactionWhileThrottled() {
+        Map configs = new HashMap<>();
+        configs.put(ProducerConfig.TRANSACTIONAL_ID_CONFIG, "some.id");
+        configs.put(ProducerConfig.MAX_BLOCK_MS_CONFIG, 10000);
+        configs.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9000");
+
+        Time time = new MockTime(1);
+        MetadataResponse initialUpdateResponse = TestUtils.metadataUpdateWith(1, singletonMap("topic", 1));
+        ProducerMetadata metadata = newMetadata(0, Long.MAX_VALUE);
+
+        MockClient client = new MockClient(time, metadata);
+        client.updateMetadata(initialUpdateResponse);
+
+        Node node = metadata.fetch().nodes().get(0);
+        client.throttle(node, 5000);
+
+        client.prepareResponse(FindCoordinatorResponse.prepareResponse(Errors.NONE, host1));
+        client.prepareResponse(initProducerIdResponse(1L, (short) 5, Errors.NONE));
+
         try (Producer producer = new KafkaProducer<>(configs, new StringSerializer(),
                 new StringSerializer(), metadata, client, null, time)) {
             producer.initTransactions();
-            fail("initTransactions() should have raised TimeoutException");
         }
     }
 
-    @Test(expected = KafkaException.class)
+    private InitProducerIdResponse initProducerIdResponse(long producerId, short producerEpoch, Errors error) {
+        InitProducerIdResponseData responseData = new InitProducerIdResponseData()
+                .setErrorCode(error.code())
+                .setProducerEpoch(producerEpoch)
+                .setProducerId(producerId)
+                .setThrottleTimeMs(0);
+        return new InitProducerIdResponse(responseData);
+    }
+
+
+    @Test
     public void testOnlyCanExecuteCloseAfterInitTransactionsTimeout() {
         Map configs = new HashMap<>();
         configs.put(ProducerConfig.TRANSACTIONAL_ID_CONFIG, "bad-transaction");
@@ -669,14 +697,10 @@ public void testOnlyCanExecuteCloseAfterInitTransactionsTimeout() {
 
         Producer producer = new KafkaProducer<>(configs, new StringSerializer(), new StringSerializer(),
                 metadata, client, null, time);
-        try {
-            producer.initTransactions();
-        } catch (TimeoutException e) {
-            // expected
-        }
+        assertThrows(TimeoutException.class, producer::initTransactions);
         // other transactional operations should not be allowed if we catch the error after initTransactions failed
         try {
-            producer.beginTransaction();
+            assertThrows(KafkaException.class, producer::beginTransaction);
         } finally {
             producer.close(Duration.ofMillis(0));
         }
@@ -704,7 +728,7 @@ public void testSendToInvalidTopic() throws Exception {
         List topicMetadata = new ArrayList<>();
         topicMetadata.add(new MetadataResponse.TopicMetadata(Errors.INVALID_TOPIC_EXCEPTION,
                 invalidTopicName, false, Collections.emptyList()));
-        MetadataResponse updateResponse = new MetadataResponse(
+        MetadataResponse updateResponse =  MetadataResponse.prepareResponse(
                 new ArrayList<>(initialUpdateResponse.brokers()),
                 initialUpdateResponse.clusterId(),
                 initialUpdateResponse.controller().id(),
@@ -766,6 +790,111 @@ public void testCloseWhenWaitingForMetadataUpdate() throws InterruptedException
         }
     }
 
+    @Test
+    public void testTransactionalMethodThrowsWhenSenderClosed() {
+        Map configs = new HashMap<>();
+        configs.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9000");
+        configs.put(ProducerConfig.TRANSACTIONAL_ID_CONFIG, "this-is-a-transactional-id");
+
+        Time time = new MockTime();
+        MetadataResponse initialUpdateResponse = TestUtils.metadataUpdateWith(1, emptyMap());
+        ProducerMetadata metadata = newMetadata(0, Long.MAX_VALUE);
+        metadata.update(initialUpdateResponse, time.milliseconds());
+
+        MockClient client = new MockClient(time, metadata);
+
+        Producer producer = new KafkaProducer<>(configs, new StringSerializer(), new StringSerializer(),
+                metadata, client, null, time);
+        producer.close();
+        assertThrows(IllegalStateException.class, producer::initTransactions);
+    }
+
+    @Test
+    public void testCloseIsForcedOnPendingFindCoordinator() throws InterruptedException {
+        Map configs = new HashMap<>();
+        configs.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9000");
+        configs.put(ProducerConfig.TRANSACTIONAL_ID_CONFIG, "this-is-a-transactional-id");
+
+        Time time = new MockTime();
+        MetadataResponse initialUpdateResponse = TestUtils.metadataUpdateWith(1, singletonMap("testTopic", 1));
+        ProducerMetadata metadata = newMetadata(0, Long.MAX_VALUE);
+        metadata.update(initialUpdateResponse, time.milliseconds());
+
+        MockClient client = new MockClient(time, metadata);
+
+        Producer producer = new KafkaProducer<>(configs, new StringSerializer(), new StringSerializer(),
+                metadata, client, null, time);
+
+        ExecutorService executorService = Executors.newSingleThreadExecutor();
+        CountDownLatch assertionDoneLatch = new CountDownLatch(1);
+        executorService.submit(() -> {
+            assertThrows(KafkaException.class, producer::initTransactions);
+            assertionDoneLatch.countDown();
+        });
+
+        client.waitForRequests(1, 2000);
+        producer.close(Duration.ofMillis(1000));
+        assertionDoneLatch.await(5000, TimeUnit.MILLISECONDS);
+    }
+
+    @Test
+    public void testCloseIsForcedOnPendingInitProducerId() throws InterruptedException {
+        Map configs = new HashMap<>();
+        configs.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9000");
+        configs.put(ProducerConfig.TRANSACTIONAL_ID_CONFIG, "this-is-a-transactional-id");
+
+        Time time = new MockTime();
+        MetadataResponse initialUpdateResponse = TestUtils.metadataUpdateWith(1, singletonMap("testTopic", 1));
+        ProducerMetadata metadata = newMetadata(0, Long.MAX_VALUE);
+        metadata.update(initialUpdateResponse, time.milliseconds());
+
+        MockClient client = new MockClient(time, metadata);
+
+        Producer producer = new KafkaProducer<>(configs, new StringSerializer(), new StringSerializer(),
+                metadata, client, null, time);
+
+        ExecutorService executorService = Executors.newSingleThreadExecutor();
+        CountDownLatch assertionDoneLatch = new CountDownLatch(1);
+        client.prepareResponse(FindCoordinatorResponse.prepareResponse(Errors.NONE, host1));
+        executorService.submit(() -> {
+            assertThrows(KafkaException.class, producer::initTransactions);
+            assertionDoneLatch.countDown();
+        });
+
+        client.waitForRequests(1, 2000);
+        producer.close(Duration.ofMillis(1000));
+        assertionDoneLatch.await(5000, TimeUnit.MILLISECONDS);
+    }
+
+    @Test
+    public void testCloseIsForcedOnPendingAddOffsetRequest() throws InterruptedException {
+        Map configs = new HashMap<>();
+        configs.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9000");
+        configs.put(ProducerConfig.TRANSACTIONAL_ID_CONFIG, "this-is-a-transactional-id");
+
+        Time time = new MockTime();
+        MetadataResponse initialUpdateResponse = TestUtils.metadataUpdateWith(1, singletonMap("testTopic", 1));
+        ProducerMetadata metadata = newMetadata(0, Long.MAX_VALUE);
+        metadata.update(initialUpdateResponse, time.milliseconds());
+
+        MockClient client = new MockClient(time, metadata);
+
+        Producer producer = new KafkaProducer<>(configs, new StringSerializer(), new StringSerializer(),
+                metadata, client, null, time);
+
+        ExecutorService executorService = Executors.newSingleThreadExecutor();
+        CountDownLatch assertionDoneLatch = new CountDownLatch(1);
+        client.prepareResponse(FindCoordinatorResponse.prepareResponse(Errors.NONE, host1));
+        executorService.submit(() -> {
+            assertThrows(KafkaException.class, producer::initTransactions);
+            assertionDoneLatch.countDown();
+        });
+
+        client.waitForRequests(1, 2000);
+        producer.close(Duration.ofMillis(1000));
+        assertionDoneLatch.await(5000, TimeUnit.MILLISECONDS);
+    }
+
     private ProducerMetadata newMetadata(long refreshBackoffMs, long expirationMs) {
         return new ProducerMetadata(refreshBackoffMs, expirationMs,
                 new LogContext(), new ClusterResourceListeners(), Time.SYSTEM);
diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/RoundRobinPartitionerTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/RoundRobinPartitionerTest.java
new file mode 100644
index 0000000000000..c6bb616047061
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/clients/producer/RoundRobinPartitionerTest.java
@@ -0,0 +1,128 @@
+/*
+ * 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.producer;
+
+import org.apache.kafka.common.Cluster;
+import org.apache.kafka.common.Node;
+import org.apache.kafka.common.PartitionInfo;
+import org.junit.Test;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import static java.util.Arrays.asList;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+public class RoundRobinPartitionerTest {
+    private final static Node[] NODES = new Node[] {
+        new Node(0, "localhost", 99),
+        new Node(1, "localhost", 100),
+        new Node(2, "localhost", 101)
+    };
+
+    @Test
+    public void testRoundRobinWithUnavailablePartitions() {
+        // Intentionally make the partition list not in partition order to test the edge
+        // cases.
+        List partitions = asList(
+                new PartitionInfo("test", 1, null, NODES, NODES),
+                new PartitionInfo("test", 2, NODES[1], NODES, NODES),
+                new PartitionInfo("test", 0, NODES[0], NODES, NODES));
+        // When there are some unavailable partitions, we want to make sure that (1) we
+        // always pick an available partition,
+        // and (2) the available partitions are selected in a round robin way.
+        int countForPart0 = 0;
+        int countForPart2 = 0;
+        Partitioner partitioner = new RoundRobinPartitioner();
+        Cluster cluster = new Cluster("clusterId", asList(NODES[0], NODES[1], NODES[2]), partitions,
+            Collections.emptySet(), Collections.emptySet());
+        for (int i = 1; i <= 100; i++) {
+            int part = partitioner.partition("test", null, null, null, null, cluster);
+            assertTrue("We should never choose a leader-less node in round robin", part == 0 || part == 2);
+            if (part == 0)
+                countForPart0++;
+            else
+                countForPart2++;
+        }
+        assertEquals("The distribution between two available partitions should be even", countForPart0, countForPart2);
+    }
+
+    @Test
+    public void testRoundRobinWithKeyBytes() throws InterruptedException {
+        final String topicA = "topicA";
+        final String topicB = "topicB";
+
+        List allPartitions = asList(new PartitionInfo(topicA, 0, NODES[0], NODES, NODES),
+                new PartitionInfo(topicA, 1, NODES[1], NODES, NODES), new PartitionInfo(topicA, 2, NODES[2], NODES, NODES),
+                new PartitionInfo(topicB, 0, NODES[0], NODES, NODES));
+        Cluster testCluster = new Cluster("clusterId", asList(NODES[0], NODES[1], NODES[2]), allPartitions,
+                Collections.emptySet(), Collections.emptySet());
+
+        final Map partitionCount = new HashMap<>();
+
+        final byte[] keyBytes = "key".getBytes();
+        Partitioner partitioner = new RoundRobinPartitioner();
+        for (int i = 0; i < 30; ++i) {
+            int partition = partitioner.partition(topicA, null, keyBytes, null, null, testCluster);
+            Integer count = partitionCount.get(partition);
+            if (null == count)
+                count = 0;
+            partitionCount.put(partition, count + 1);
+
+            if (i % 5 == 0) {
+                partitioner.partition(topicB, null, keyBytes, null, null, testCluster);
+            }
+        }
+
+        assertEquals(10, partitionCount.get(0).intValue());
+        assertEquals(10, partitionCount.get(1).intValue());
+        assertEquals(10, partitionCount.get(2).intValue());
+    }
+    
+    @Test
+    public void testRoundRobinWithNullKeyBytes() throws InterruptedException {
+        final String topicA = "topicA";
+        final String topicB = "topicB";
+
+        List allPartitions = asList(new PartitionInfo(topicA, 0, NODES[0], NODES, NODES),
+                new PartitionInfo(topicA, 1, NODES[1], NODES, NODES), new PartitionInfo(topicA, 2, NODES[2], NODES, NODES),
+                new PartitionInfo(topicB, 0, NODES[0], NODES, NODES));
+        Cluster testCluster = new Cluster("clusterId", asList(NODES[0], NODES[1], NODES[2]), allPartitions,
+                Collections.emptySet(), Collections.emptySet());
+
+        final Map partitionCount = new HashMap<>();
+
+        Partitioner partitioner = new RoundRobinPartitioner();
+        for (int i = 0; i < 30; ++i) {
+            int partition = partitioner.partition(topicA, null, null, null, null, testCluster);
+            Integer count = partitionCount.get(partition);
+            if (null == count)
+                count = 0;
+            partitionCount.put(partition, count + 1);
+
+            if (i % 5 == 0) {
+                partitioner.partition(topicB, null, null, null, null, testCluster);
+            }
+        }
+
+        assertEquals(10, partitionCount.get(0).intValue());
+        assertEquals(10, partitionCount.get(1).intValue());
+        assertEquals(10, partitionCount.get(2).intValue());
+    }    
+}
diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/UniformStickyPartitionerTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/UniformStickyPartitionerTest.java
new file mode 100644
index 0000000000000..918390786f8b0
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/clients/producer/UniformStickyPartitionerTest.java
@@ -0,0 +1,207 @@
+/*
+ * 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.producer;
+
+import org.apache.kafka.common.Cluster;
+import org.apache.kafka.common.Node;
+import org.apache.kafka.common.PartitionInfo;
+import org.junit.Test;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import static java.util.Arrays.asList;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+public class UniformStickyPartitionerTest {
+    private final static Node[] NODES = new Node[] {
+        new Node(0, "localhost", 99),
+        new Node(1, "localhost", 100),
+        new Node(2, "localhost", 101)
+    };
+
+    private final static String TOPIC_A = "TOPIC_A";
+    private final static String TOPIC_B = "TOPIC_B";
+
+    @Test
+    public void testRoundRobinWithUnavailablePartitions() {
+        // Intentionally make the partition list not in partition order to test the edge
+        // cases.
+        List partitions = asList(
+                new PartitionInfo("test", 1, null, NODES, NODES),
+                new PartitionInfo("test", 2, NODES[1], NODES, NODES),
+                new PartitionInfo("test", 0, NODES[0], NODES, NODES));
+        // When there are some unavailable partitions, we want to make sure that (1) we
+        // always pick an available partition,
+        // and (2) the available partitions are selected in a sticky way.
+        int countForPart0 = 0;
+        int countForPart2 = 0;
+        int part = 0;
+        Partitioner partitioner = new UniformStickyPartitioner();
+        Cluster cluster = new Cluster("clusterId", asList(NODES[0], NODES[1], NODES[2]), partitions,
+            Collections.emptySet(), Collections.emptySet());
+        for (int i = 0; i < 50; i++) {
+            part = partitioner.partition("test", null, null, null, null, cluster);
+            assertTrue("We should never choose a leader-less node in round robin", part == 0 || part == 2);
+            if (part == 0)
+                countForPart0++;
+            else
+                countForPart2++;
+        }
+        // Simulates switching the sticky partition on a new batch.
+        partitioner.onNewBatch("test", cluster, part);       
+        for (int i = 1; i <= 50; i++) {
+            part = partitioner.partition("test", null, null, null, null, cluster);
+            assertTrue("We should never choose a leader-less node in round robin", part == 0 || part == 2);
+            if (part == 0)
+                countForPart0++;
+            else
+                countForPart2++;
+        }
+        assertEquals("The distribution between two available partitions should be even", countForPart0, countForPart2);
+    }
+
+    @Test
+    public void testRoundRobinWithKeyBytes() throws InterruptedException {
+        List allPartitions = asList(new PartitionInfo(TOPIC_A, 0, NODES[0], NODES, NODES),
+                new PartitionInfo(TOPIC_A, 1, NODES[1], NODES, NODES), new PartitionInfo(TOPIC_A, 2, NODES[1], NODES, NODES),
+                new PartitionInfo(TOPIC_B, 0, NODES[0], NODES, NODES));
+        Cluster testCluster = new Cluster("clusterId", asList(NODES[0], NODES[1], NODES[2]), allPartitions,
+                Collections.emptySet(), Collections.emptySet());
+
+        final Map partitionCount = new HashMap<>();
+
+        final byte[] keyBytes = "key".getBytes();
+        int partition = 0;
+        Partitioner partitioner = new UniformStickyPartitioner();
+        for (int i = 0; i < 30; ++i) {
+            partition = partitioner.partition(TOPIC_A, null, keyBytes, null, null, testCluster);
+            Integer count = partitionCount.get(partition);
+            if (null == count)
+                count = 0;
+            partitionCount.put(partition, count + 1);
+
+            if (i % 5 == 0) {
+                partitioner.partition(TOPIC_B, null, keyBytes, null, null, testCluster);
+            }
+        }
+        // Simulate a batch filling up and switching the sticky partition.
+        partitioner.onNewBatch(TOPIC_A, testCluster, partition);
+        partitioner.onNewBatch(TOPIC_B, testCluster, 0);
+        
+        // Save old partition to ensure that the wrong partition does not trigger a new batch.
+        int oldPart = partition; 
+        
+        for (int i = 0; i < 30; ++i) {
+            partition = partitioner.partition(TOPIC_A, null, keyBytes, null, null, testCluster);
+            Integer count = partitionCount.get(partition);
+            if (null == count)
+                count = 0;
+            partitionCount.put(partition, count + 1);
+
+            if (i % 5 == 0) {
+                partitioner.partition(TOPIC_B, null, keyBytes, null, null, testCluster);
+            }
+        }
+        
+        int newPart = partition;
+        
+        // Attempt to switch the partition with the wrong previous partition. Sticky partition should not change.
+        partitioner.onNewBatch(TOPIC_A, testCluster, oldPart);
+        
+        for (int i = 0; i < 30; ++i) {
+            partition = partitioner.partition(TOPIC_A, null, keyBytes, null, null, testCluster);
+            Integer count = partitionCount.get(partition);
+            if (null == count)
+                count = 0;
+            partitionCount.put(partition, count + 1);
+
+            if (i % 5 == 0) {
+                partitioner.partition(TOPIC_B, null, keyBytes, null, null, testCluster);
+            }
+        } 
+               
+        assertEquals(30, partitionCount.get(oldPart).intValue());
+        assertEquals(60, partitionCount.get(newPart).intValue());
+    }
+    
+    @Test
+    public void testRoundRobinWithNullKeyBytes() throws InterruptedException {
+        List allPartitions = asList(new PartitionInfo(TOPIC_A, 0, NODES[0], NODES, NODES),
+                new PartitionInfo(TOPIC_A, 1, NODES[1], NODES, NODES), new PartitionInfo(TOPIC_A, 2, NODES[1], NODES, NODES),
+                new PartitionInfo(TOPIC_B, 0, NODES[0], NODES, NODES));
+        Cluster testCluster = new Cluster("clusterId", asList(NODES[0], NODES[1], NODES[2]), allPartitions,
+                Collections.emptySet(), Collections.emptySet());
+
+        final Map partitionCount = new HashMap<>();
+
+        int partition = 0;
+        Partitioner partitioner = new UniformStickyPartitioner();
+        for (int i = 0; i < 30; ++i) {
+            partition = partitioner.partition(TOPIC_A, null, null, null, null, testCluster);
+            Integer count = partitionCount.get(partition);
+            if (null == count)
+                count = 0;
+            partitionCount.put(partition, count + 1);
+
+            if (i % 5 == 0) {
+                partitioner.partition(TOPIC_B, null, null, null, null, testCluster);
+            }
+        }
+        // Simulate a batch filling up and switching the sticky partition.
+        partitioner.onNewBatch(TOPIC_A, testCluster, partition);
+        partitioner.onNewBatch(TOPIC_B, testCluster, 0);
+
+        // Save old partition to ensure that the wrong partition does not trigger a new batch.
+        int oldPart = partition;
+
+        for (int i = 0; i < 30; ++i) {
+            partition = partitioner.partition(TOPIC_A, null, null, null, null, testCluster);
+            Integer count = partitionCount.get(partition);
+            if (null == count)
+                count = 0;
+            partitionCount.put(partition, count + 1);
+
+            if (i % 5 == 0) {
+                partitioner.partition(TOPIC_B, null, null, null, null, testCluster);
+            }
+        }
+
+        int newPart = partition;
+
+        // Attempt to switch the partition with the wrong previous partition. Sticky partition should not change.
+        partitioner.onNewBatch(TOPIC_A, testCluster, oldPart);
+
+        for (int i = 0; i < 30; ++i) {
+            partition = partitioner.partition(TOPIC_A, null, null, null, null, testCluster);
+            Integer count = partitionCount.get(partition);
+            if (null == count)
+                count = 0;
+            partitionCount.put(partition, count + 1);
+
+            if (i % 5 == 0) {
+                partitioner.partition(TOPIC_B, null, null, null, null, testCluster);
+            }
+        }
+
+        assertEquals(30, partitionCount.get(oldPart).intValue());
+        assertEquals(60, partitionCount.get(newPart).intValue());
+    }
+}
+
diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/BufferPoolTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/BufferPoolTest.java
index 8e44fa18146d4..2c10284adf260 100644
--- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/BufferPoolTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/BufferPoolTest.java
@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.clients.producer.internals;
 
+import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.errors.TimeoutException;
 import org.apache.kafka.common.metrics.Metrics;
 import org.apache.kafka.common.utils.MockTime;
@@ -28,7 +29,10 @@
 import java.util.ArrayList;
 import java.util.Deque;
 import java.util.List;
+import java.util.concurrent.Callable;
 import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ExecutorService;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -36,6 +40,7 @@
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotEquals;
+import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 import static org.mockito.ArgumentMatchers.anyLong;
@@ -232,7 +237,7 @@ public void testCleanupMemoryAvailabilityWaiterOnInterruption() throws Exception
         // both the allocate() called by threads t1 and t2 should have been interrupted and the waiters queue should be empty
         assertEquals(pool.queued(), 0);
     }
-    
+
     @Test
     public void testCleanupMemoryAvailabilityOnMetricsException() throws Exception {
         BufferPool bufferPool = spy(new BufferPool(2, 1, new Metrics(), time, metricGroup));
@@ -346,6 +351,44 @@ protected ByteBuffer allocateByteBuffer(int size) {
         assertEquals(bufferPool.availableMemory(), 1024);
     }
 
+    @Test
+    public void overDeallocate() {
+        BufferPool bufferPool = new BufferPool(1024, 512, metrics, time, metricGroup);
+        bufferPool.deallocate(ByteBuffer.allocate(512));
+        assertEquals(bufferPool.availableMemory(), bufferPool.totalMemory());
+    }
+
+    @Test
+    public void dedupeDeallocations() throws Exception {
+        BufferPool bufferPool = new BufferPool(1024, 512, metrics, time, metricGroup);
+        ByteBuffer bbuf = bufferPool.allocate(512, 1);
+        ByteBuffer shallowCopy = bbuf.slice();
+        bufferPool.deallocate(bbuf);
+        bufferPool.deallocate(shallowCopy);
+        assertEquals(bufferPool.totalMemory(), bufferPool.availableMemory());
+    }
+
+    @Test
+    public void testExpandBufferDeallocations() throws Exception {
+
+        int batchSize = 512;
+        BufferPool bufferPool = new BufferPool(2048, batchSize, metrics, time, metricGroup);
+
+        ByteBuffer bbuf1 = bufferPool.allocate(batchSize, 1);
+        ByteBuffer bbuf2 = bufferPool.allocate(batchSize, 1);
+
+        int expandSize = batchSize * 2;
+        // expand bbuf1 to 2x of the original size
+        bbuf1 = ByteBuffer.allocate(expandSize);
+
+        // return initially allocated size of bbuf1 to bufferpool
+        bufferPool.deallocate(bbuf1, batchSize);
+        assertEquals(bufferPool.totalMemory(), bufferPool.availableMemory() + batchSize);
+
+        bufferPool.deallocate(bbuf2, batchSize);
+        assertEquals(bufferPool.totalMemory(), bufferPool.availableMemory());
+    }
+
     public static class StressTestThread extends Thread {
         private final int iterations;
         private final BufferPool pool;
@@ -377,4 +420,48 @@ public void run() {
         }
     }
 
+    @Test
+    public void testCloseAllocations() throws Exception {
+        BufferPool pool = new BufferPool(10, 1, metrics, Time.SYSTEM, metricGroup);
+        ByteBuffer buffer = pool.allocate(1, maxBlockTimeMs);
+
+        // Close the buffer pool. This should prevent any further allocations.
+        pool.close();
+
+        assertThrows(KafkaException.class, () -> pool.allocate(1, maxBlockTimeMs));
+
+        // Ensure deallocation still works.
+        pool.deallocate(buffer);
+    }
+
+    @Test
+    public void testCloseNotifyWaiters() throws Exception {
+        final int numWorkers = 2;
+
+        BufferPool pool = new BufferPool(1, 1, metrics, Time.SYSTEM, metricGroup);
+        ByteBuffer buffer = pool.allocate(1, Long.MAX_VALUE);
+
+        CountDownLatch completed = new CountDownLatch(numWorkers);
+        ExecutorService executor = Executors.newFixedThreadPool(numWorkers);
+        Callable work = new Callable() {
+                public Void call() throws Exception {
+                    assertThrows(KafkaException.class, () -> pool.allocate(1, Long.MAX_VALUE));
+                    completed.countDown();
+                    return null;
+                }
+            };
+        for (int i = 0; i < numWorkers; ++i) {
+            executor.submit(work);
+        }
+
+        assertEquals("Allocation shouldn't have happened yet, waiting on memory", numWorkers, completed.getCount());
+
+        // Close the buffer pool. This should notify all waiters.
+        pool.close();
+
+        completed.await(15, TimeUnit.SECONDS);
+
+        pool.deallocate(buffer);
+    }
+
 }
diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/DefaultPartitionerTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/DefaultPartitionerTest.java
index f3fdd656c919e..d88b0dad4237d 100644
--- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/DefaultPartitionerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/DefaultPartitionerTest.java
@@ -23,79 +23,29 @@
 import org.junit.Test;
 
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.List;
-import java.util.Map;
 import static java.util.Arrays.asList;
 import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertTrue;
 
 public class DefaultPartitionerTest {
-    private byte[] keyBytes = "key".getBytes();
-    private Partitioner partitioner = new DefaultPartitioner();
-    private Node node0 = new Node(0, "localhost", 99);
-    private Node node1 = new Node(1, "localhost", 100);
-    private Node node2 = new Node(2, "localhost", 101);
-    private Node[] nodes = new Node[] {node0, node1, node2};
-    private String topic = "test";
+    private final static byte[] KEY_BYTES = "key".getBytes();
+    private final static Node[] NODES = new Node[] {
+        new Node(0, "localhost", 99),
+        new Node(1, "localhost", 100),
+        new Node(12, "localhost", 101)
+    };
+    private final static String TOPIC = "test";
     // Intentionally make the partition list not in partition order to test the edge cases.
-    private List partitions = asList(new PartitionInfo(topic, 1, null, nodes, nodes),
-                                                    new PartitionInfo(topic, 2, node1, nodes, nodes),
-                                                    new PartitionInfo(topic, 0, node0, nodes, nodes));
-    private Cluster cluster = new Cluster("clusterId", asList(node0, node1, node2), partitions,
-            Collections.emptySet(), Collections.emptySet());
+    private final static List PARTITIONS = asList(new PartitionInfo(TOPIC, 1, null, NODES, NODES),
+                                                    new PartitionInfo(TOPIC, 2, NODES[1], NODES, NODES),
+                                                    new PartitionInfo(TOPIC, 0, NODES[0], NODES, NODES));
 
     @Test
     public void testKeyPartitionIsStable() {
-        int partition = partitioner.partition("test",  null, keyBytes, null, null, cluster);
-        assertEquals("Same key should yield same partition", partition, partitioner.partition("test", null, keyBytes, null, null, cluster));
-    }
-
-    @Test
-    public void testRoundRobinWithUnavailablePartitions() {
-        // When there are some unavailable partitions, we want to make sure that (1) we always pick an available partition,
-        // and (2) the available partitions are selected in a round robin way.
-        int countForPart0 = 0;
-        int countForPart2 = 0;
-        for (int i = 1; i <= 100; i++) {
-            int part = partitioner.partition("test", null, null, null, null, cluster);
-            assertTrue("We should never choose a leader-less node in round robin", part == 0 || part == 2);
-            if (part == 0)
-                countForPart0++;
-            else
-                countForPart2++;
-        }
-        assertEquals("The distribution between two available partitions should be even", countForPart0, countForPart2);
-    }
-
-    @Test
-    public void testRoundRobin() throws InterruptedException {
-        final String topicA = "topicA";
-        final String topicB = "topicB";
-
-        List allPartitions = asList(new PartitionInfo(topicA, 0, node0, nodes, nodes),
-                new PartitionInfo(topicA, 1, node1, nodes, nodes),
-                new PartitionInfo(topicA, 2, node2, nodes, nodes),
-                new PartitionInfo(topicB, 0, node0, nodes, nodes)
-                );
-        Cluster testCluster = new Cluster("clusterId", asList(node0, node1, node2), allPartitions,
-                Collections.emptySet(), Collections.emptySet());
-
-        final Map partitionCount = new HashMap<>();
-
-        for (int i = 0; i < 30; ++i) {
-            int partition = partitioner.partition(topicA, null, null, null, null, testCluster);
-            Integer count = partitionCount.get(partition);
-            if (null == count) count = 0;
-            partitionCount.put(partition, count + 1);
-
-            if (i % 5 == 0) {
-                partitioner.partition(topicB, null, null, null, null, testCluster);
-            }
-        }
-
-        assertEquals(10, (int) partitionCount.get(0));
-        assertEquals(10, (int) partitionCount.get(1));
-        assertEquals(10, (int) partitionCount.get(2));
+        final Partitioner partitioner = new DefaultPartitioner();
+        final Cluster cluster = new Cluster("clusterId", asList(NODES), PARTITIONS,
+            Collections.emptySet(), Collections.emptySet());
+        int partition = partitioner.partition("test",  null, KEY_BYTES, null, null, cluster);
+        assertEquals("Same key should yield same partition", partition, partitioner.partition("test", null, KEY_BYTES, null, null, cluster));
     }
 }
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 f9887f9033e81..fd25c844afd95 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
@@ -158,6 +158,9 @@ public void testAppendedChecksumMagicV0AndV1() {
     @Test
     public void testSplitPreservesHeaders() {
         for (CompressionType compressionType : CompressionType.values()) {
+            if (compressionType == CompressionType.PASSTHROUGH)
+                continue;
+
             MemoryRecordsBuilder builder = MemoryRecords.builder(
                     ByteBuffer.allocate(1024),
                     MAGIC_VALUE_V2,
@@ -194,7 +197,8 @@ public void testSplitPreservesHeaders() {
     public void testSplitPreservesMagicAndCompressionType() {
         for (byte magic : Arrays.asList(MAGIC_VALUE_V0, MAGIC_VALUE_V1, MAGIC_VALUE_V2)) {
             for (CompressionType compressionType : CompressionType.values()) {
-                if (compressionType == CompressionType.NONE && magic < MAGIC_VALUE_V2)
+                if ((compressionType == CompressionType.NONE && magic < MAGIC_VALUE_V2) ||
+                    compressionType == CompressionType.PASSTHROUGH)
                     continue;
 
                 if (compressionType == CompressionType.ZSTD && magic < MAGIC_VALUE_V2)
diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ProducerMetadataTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ProducerMetadataTest.java
index aaf38576a0252..58fbd445b9f10 100644
--- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ProducerMetadataTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ProducerMetadataTest.java
@@ -17,6 +17,7 @@
 package org.apache.kafka.clients.producer.internals;
 
 import org.apache.kafka.common.KafkaException;
+import org.apache.kafka.common.errors.AuthenticationException;
 import org.apache.kafka.common.errors.TimeoutException;
 import org.apache.kafka.common.internals.ClusterResourceListeners;
 import org.apache.kafka.common.requests.MetadataResponse;
@@ -35,6 +36,7 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
@@ -174,6 +176,12 @@ public void testTopicExpiry() {
         }
     }
 
+    @Test
+    public void testMetadataWaitAbortedOnFatalException() {
+        metadata.fatalError(new AuthenticationException("Fatal exception from test"));
+        assertThrows(AuthenticationException.class, () -> metadata.awaitUpdate(0, 1000));
+    }
+
     private MetadataResponse responseWithCurrentTopics() {
         return responseWithTopics(metadata.topics());
     }
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 13b0d1ba1c5de..9b01d6bc7cdaa 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
@@ -16,16 +16,22 @@
  */
 package org.apache.kafka.clients.producer.internals;
 
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.atomic.AtomicBoolean;
 import org.apache.kafka.clients.ApiVersions;
 import org.apache.kafka.clients.NodeApiVersions;
 import org.apache.kafka.clients.producer.Callback;
+import org.apache.kafka.clients.producer.Partitioner;
 import org.apache.kafka.clients.producer.RecordMetadata;
 import org.apache.kafka.common.Cluster;
 import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.Node;
 import org.apache.kafka.common.PartitionInfo;
 import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.errors.TimeoutException;
 import org.apache.kafka.common.errors.UnsupportedVersionException;
+import org.apache.kafka.common.header.Header;
+import org.apache.kafka.common.header.internals.RecordHeader;
 import org.apache.kafka.common.metrics.Metrics;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.record.CompressionRatioEstimator;
@@ -37,12 +43,12 @@
 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.ApiVersionsResponse;
 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.test.TestUtils;
 import org.junit.After;
+import org.junit.Assert;
 import org.junit.Test;
 
 import java.nio.ByteBuffer;
@@ -103,11 +109,11 @@ public void testFull() throws Exception {
         int batchSize = 1025;
 
         RecordAccumulator accum = createTestRecordAccumulator(
-                batchSize + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 10L * batchSize, CompressionType.NONE, 10L);
+                batchSize + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 10L * batchSize, CompressionType.NONE, 10);
         int appends = expectedNumAppends(batchSize);
         for (int i = 0; i < appends; i++) {
             // append to the first batch
-            accum.append(tp1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs);
+            accum.append(tp1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false);
             Deque partitionBatches = accum.batches().get(tp1);
             assertEquals(1, partitionBatches.size());
 
@@ -118,7 +124,7 @@ public void testFull() throws Exception {
 
         // this append doesn't fit in the first batch, so a new batch is created and the first batch is closed
 
-        accum.append(tp1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs);
+        accum.append(tp1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false);
         Deque partitionBatches = accum.batches().get(tp1);
         assertEquals(2, partitionBatches.size());
         Iterator partitionBatchesIterator = partitionBatches.iterator();
@@ -152,8 +158,8 @@ private void testAppendLarge(CompressionType compressionType) throws Exception {
         int batchSize = 512;
         byte[] value = new byte[2 * batchSize];
         RecordAccumulator accum = createTestRecordAccumulator(
-                batchSize + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 10 * 1024, compressionType, 0L);
-        accum.append(tp1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs);
+                batchSize + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 10 * 1024, compressionType, 0);
+        accum.append(tp1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false);
         assertEquals("Our partition's leader should be ready", Collections.singleton(node1), accum.ready(cluster, time.milliseconds()).readyNodes);
 
         Deque batches = accum.batches().get(tp1);
@@ -187,12 +193,11 @@ private void testAppendLargeOldMessageFormat(CompressionType compressionType) th
         byte[] value = new byte[2 * batchSize];
 
         ApiVersions apiVersions = new ApiVersions();
-        apiVersions.update(node1.idString(), NodeApiVersions.create(Collections.singleton(
-                new ApiVersionsResponse.ApiVersion(ApiKeys.PRODUCE.id, (short) 0, (short) 2))));
+        apiVersions.update(node1.idString(), NodeApiVersions.create(ApiKeys.PRODUCE.id, (short) 0, (short) 2));
 
         RecordAccumulator accum = createTestRecordAccumulator(
-                batchSize + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 10 * 1024, compressionType, 0L);
-        accum.append(tp1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs);
+                batchSize + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 10 * 1024, compressionType, 0);
+        accum.append(tp1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false);
         assertEquals("Our partition's leader should be ready", Collections.singleton(node1), accum.ready(cluster, time.milliseconds()).readyNodes);
 
         Deque batches = accum.batches().get(tp1);
@@ -213,10 +218,10 @@ private void testAppendLargeOldMessageFormat(CompressionType compressionType) th
 
     @Test
     public void testLinger() throws Exception {
-        long lingerMs = 10L;
+        int lingerMs = 10;
         RecordAccumulator accum = createTestRecordAccumulator(
                 1024 + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 10 * 1024, CompressionType.NONE, lingerMs);
-        accum.append(tp1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs);
+        accum.append(tp1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false);
         assertEquals("No partitions should be ready", 0, accum.ready(cluster, time.milliseconds()).readyNodes.size());
         time.sleep(10);
         assertEquals("Our partition's leader should be ready", Collections.singleton(node1), accum.ready(cluster, time.milliseconds()).readyNodes);
@@ -234,12 +239,12 @@ public void testLinger() throws Exception {
     @Test
     public void testPartialDrain() throws Exception {
         RecordAccumulator accum = createTestRecordAccumulator(
-                1024 + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 10 * 1024, CompressionType.NONE, 10L);
+                1024 + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 10 * 1024, CompressionType.NONE, 10);
         int appends = 1024 / msgSize + 1;
         List partitions = asList(tp1, tp2);
         for (TopicPartition tp : partitions) {
             for (int i = 0; i < appends; i++)
-                accum.append(tp, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs);
+                accum.append(tp, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false);
         }
         assertEquals("Partition's leader should be ready", Collections.singleton(node1), accum.ready(cluster, time.milliseconds()).readyNodes);
 
@@ -254,14 +259,14 @@ public void testStressfulSituation() throws Exception {
         final int msgs = 10000;
         final int numParts = 2;
         final RecordAccumulator accum = createTestRecordAccumulator(
-            1024 + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 10 * 1024, CompressionType.NONE, 0L);
+            1024 + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 10 * 1024, CompressionType.NONE, 0);
         List threads = new ArrayList<>();
         for (int i = 0; i < numThreads; i++) {
             threads.add(new Thread() {
                 public void run() {
                     for (int i = 0; i < msgs; i++) {
                         try {
-                            accum.append(new TopicPartition(topic, i % numParts), 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs);
+                            accum.append(new TopicPartition(topic, i % numParts), 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false);
                         } catch (Exception e) {
                             e.printStackTrace();
                         }
@@ -293,7 +298,7 @@ public void run() {
     @Test
     public void testNextReadyCheckDelay() throws Exception {
         // Next check time will use lingerMs since this test won't trigger any retries/backoff
-        long lingerMs = 10L;
+        int lingerMs = 10;
 
         // test case assumes that the records do not fill the batch completely
         int batchSize = 1025;
@@ -305,7 +310,7 @@ public void testNextReadyCheckDelay() throws Exception {
 
         // Partition on node1 only
         for (int i = 0; i < appends; i++)
-            accum.append(tp1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs);
+            accum.append(tp1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false);
         RecordAccumulator.ReadyCheckResult result = accum.ready(cluster, time.milliseconds());
         assertEquals("No nodes should be ready.", 0, result.readyNodes.size());
         assertEquals("Next check time should be the linger time", lingerMs, result.nextReadyCheckDelayMs);
@@ -314,14 +319,14 @@ public void testNextReadyCheckDelay() throws Exception {
 
         // Add partition on node2 only
         for (int i = 0; i < appends; i++)
-            accum.append(tp3, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs);
+            accum.append(tp3, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false);
         result = accum.ready(cluster, time.milliseconds());
         assertEquals("No nodes should be ready.", 0, result.readyNodes.size());
         assertEquals("Next check time should be defined by node1, half remaining linger time", lingerMs / 2, result.nextReadyCheckDelayMs);
 
         // Add data for another partition on node1, enough to make data sendable immediately
         for (int i = 0; i < appends + 1; i++)
-            accum.append(tp2, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs);
+            accum.append(tp2, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false);
         result = accum.ready(cluster, time.milliseconds());
         assertEquals("Node1 should be ready", Collections.singleton(node1), result.readyNodes);
         // Note this can actually be < linger time because it may use delays from partitions that aren't sendable
@@ -331,10 +336,9 @@ public void testNextReadyCheckDelay() throws Exception {
 
     @Test
     public void testRetryBackoff() throws Exception {
-        long lingerMs = Integer.MAX_VALUE / 16;
+        int lingerMs = Integer.MAX_VALUE / 16;
         long retryBackoffMs = Integer.MAX_VALUE / 8;
-        int requestTimeoutMs = Integer.MAX_VALUE / 4;
-        long deliveryTimeoutMs = Integer.MAX_VALUE;
+        int deliveryTimeoutMs = Integer.MAX_VALUE;
         long totalSize = 10 * 1024;
         int batchSize = 1024 + DefaultRecordBatch.RECORD_BATCH_OVERHEAD;
         String metricGrpName = "producer-metrics";
@@ -344,7 +348,7 @@ CompressionType.NONE, lingerMs, retryBackoffMs, deliveryTimeoutMs, metrics, metr
             new BufferPool(totalSize, batchSize, metrics, time, metricGrpName));
 
         long now = time.milliseconds();
-        accum.append(tp1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs);
+        accum.append(tp1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false);
         RecordAccumulator.ReadyCheckResult result = accum.ready(cluster, now + lingerMs + 1);
         assertEquals("Node1 should be ready", Collections.singleton(node1), result.readyNodes);
         Map> batches = accum.drain(cluster, result.readyNodes, Integer.MAX_VALUE, now + lingerMs + 1);
@@ -356,7 +360,7 @@ CompressionType.NONE, lingerMs, retryBackoffMs, deliveryTimeoutMs, metrics, metr
         accum.reenqueue(batches.get(0).get(0), now);
 
         // Put message for partition 1 into accumulator
-        accum.append(tp2, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs);
+        accum.append(tp2, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false);
         result = accum.ready(cluster, now + lingerMs + 1);
         assertEquals("Node1 should be ready", Collections.singleton(node1), result.readyNodes);
 
@@ -377,12 +381,12 @@ CompressionType.NONE, lingerMs, retryBackoffMs, deliveryTimeoutMs, metrics, metr
 
     @Test
     public void testFlush() throws Exception {
-        long lingerMs = Integer.MAX_VALUE;
+        int lingerMs = Integer.MAX_VALUE;
         final RecordAccumulator accum = createTestRecordAccumulator(
                 4 * 1024 + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 64 * 1024, CompressionType.NONE, lingerMs);
 
         for (int i = 0; i < 100; i++) {
-            accum.append(new TopicPartition(topic, i % 3), 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs);
+            accum.append(new TopicPartition(topic, i % 3), 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false);
             assertTrue(accum.hasIncomplete());
         }
         RecordAccumulator.ReadyCheckResult result = accum.ready(cluster, time.milliseconds());
@@ -400,7 +404,7 @@ public void testFlush() throws Exception {
                 accum.deallocate(batch);
 
         // should be complete with no unsent records.
-        accum.awaitFlushCompletion();
+        accum.awaitFlushCompletion(Integer.MAX_VALUE);
         assertFalse(accum.hasUndrained());
         assertFalse(accum.hasIncomplete());
     }
@@ -416,23 +420,88 @@ public void run() {
         t.start();
     }
 
+    @Test
+    public void testSplitAwaitFlushComplete() throws Exception {
+        RecordAccumulator accum = createTestRecordAccumulator(1024, 10 * 1024, CompressionType.GZIP, 10);
+
+        // Create a big batch
+        byte[] value = new byte[256];
+        // Create a batch such that it fails with RecordBatchTooLargeException
+        accum.append(new TopicPartition(topic, 0), 0L, null, value, null, null, maxBlockTimeMs, false);
+        accum.append(new TopicPartition(topic, 0), 0L, null, value, null, null, maxBlockTimeMs, false);
+
+        CountDownLatch flushInProgress = new CountDownLatch(1);
+        Iterator incompleteBatches = accum.incompleteBatches().copyAll().iterator();
+
+        // Assert that there is only one batch
+        Assert.assertTrue(incompleteBatches.hasNext());
+        ProducerBatch producerBatch = incompleteBatches.next();
+        Assert.assertFalse(incompleteBatches.hasNext());
+
+        AtomicBoolean timedOut = new AtomicBoolean(false);
+        Thread thread = new Thread(() -> {
+            Assert.assertTrue(accum.hasIncomplete());
+            accum.beginFlush();
+            Assert.assertTrue(accum.flushInProgress());
+            try {
+                flushInProgress.countDown();
+                accum.awaitFlushCompletion(2000);
+            } catch (TimeoutException timeoutException) {
+              // Catch it and set the timedout variable
+              // This is the only valid path for this thread.
+                timedOut.set(true);
+            } catch (InterruptedException e) {
+            }
+        });
+        thread.start();
+        flushInProgress.await();
+        // Wait for 100ms to make sure that the flush is actually in progress
+        Thread.sleep(100);
+
+        // Split the big batch and re-enqueue
+        accum.splitAndReenqueue(producerBatch);
+        accum.deallocate(producerBatch);
+
+        thread.join();
+        // The thread would have failed with timeout exception because the child batches
+        // are not evaluated and it would have waited for 2seconds before the timeout.
+        Assert.assertTrue("The thread should have timed out", timedOut.get());
+    }
+
     @Test
     public void testAwaitFlushComplete() throws Exception {
         RecordAccumulator accum = createTestRecordAccumulator(
-            4 * 1024 + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 64 * 1024, CompressionType.NONE, Long.MAX_VALUE);
-        accum.append(new TopicPartition(topic, 0), 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs);
+            4 * 1024 + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 64 * 1024, CompressionType.NONE, Integer.MAX_VALUE);
+        accum.append(new TopicPartition(topic, 0), 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false);
 
         accum.beginFlush();
         assertTrue(accum.flushInProgress());
         delayedInterrupt(Thread.currentThread(), 1000L);
         try {
-            accum.awaitFlushCompletion();
+            accum.awaitFlushCompletion(Integer.MAX_VALUE);
             fail("awaitFlushCompletion should throw InterruptException");
         } catch (InterruptedException e) {
             assertFalse("flushInProgress count should be decremented even if thread is interrupted", accum.flushInProgress());
         }
     }
 
+    @Test
+    public void testAwaitFlushTimeout() throws Exception {
+        RecordAccumulator accum = createTestRecordAccumulator(
+             4 * 1024 + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 64 * 1024, CompressionType.NONE, Integer.MAX_VALUE);
+        accum.append(new TopicPartition(topic, 0), 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false);
+
+        accum.beginFlush();
+        assertTrue(accum.flushInProgress());
+        try {
+            accum.awaitFlushCompletion(100);
+            fail("testAwaitFlushTimeout should throw TimeoutException");
+        } catch (TimeoutException e) {
+            assertFalse("flushInProgress count should be decremented even if flush timeout expires", accum.flushInProgress());
+        }
+    }
+
+
     @Test
     public void testAbortIncompleteBatches() throws Exception {
         int lingerMs = Integer.MAX_VALUE;
@@ -449,7 +518,7 @@ public void onCompletion(RecordMetadata metadata, Exception exception) {
             }
         }
         for (int i = 0; i < numRecords; i++)
-            accum.append(new TopicPartition(topic, i % 3), 0L, key, value, null, new TestCallback(), maxBlockTimeMs);
+            accum.append(new TopicPartition(topic, i % 3), 0L, key, value, null, new TestCallback(), maxBlockTimeMs, false);
         RecordAccumulator.ReadyCheckResult result = accum.ready(cluster, time.milliseconds());
         assertFalse(result.readyNodes.isEmpty());
         Map> drained = accum.drain(cluster, result.readyNodes, Integer.MAX_VALUE, time.milliseconds());
@@ -490,7 +559,7 @@ public void onCompletion(RecordMetadata metadata, Exception exception) {
             }
         }
         for (int i = 0; i < numRecords; i++)
-            accum.append(new TopicPartition(topic, i % 3), 0L, key, value, null, new TestCallback(), maxBlockTimeMs);
+            accum.append(new TopicPartition(topic, i % 3), 0L, key, value, null, new TestCallback(), maxBlockTimeMs, false);
         RecordAccumulator.ReadyCheckResult result = accum.ready(cluster, time.milliseconds());
         assertFalse(result.readyNodes.isEmpty());
         Map> drained = accum.drain(cluster, result.readyNodes, Integer.MAX_VALUE,
@@ -515,8 +584,8 @@ public void onCompletion(RecordMetadata metadata, Exception exception) {
         assertTrue(accum.hasIncomplete());
     }
 
-    private void doExpireBatchSingle(long deliveryTimeoutMs) throws InterruptedException {
-        long lingerMs = 300L;
+    private void doExpireBatchSingle(int deliveryTimeoutMs) throws InterruptedException {
+        int lingerMs = 300;
         List muteStates = Arrays.asList(false, true);
         Set readyNodes = null;
         List expiredBatches = new ArrayList<>();
@@ -529,7 +598,7 @@ private void doExpireBatchSingle(long deliveryTimeoutMs) throws InterruptedExcep
         for (Boolean mute: muteStates) {
             if (time.milliseconds() < System.currentTimeMillis())
                 time.setCurrentTimeMs(System.currentTimeMillis());
-            accum.append(tp1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs);
+            accum.append(tp1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false);
             assertEquals("No partition should be ready.", 0, accum.ready(cluster, time.milliseconds()).readyNodes.size());
 
             time.sleep(lingerMs);
@@ -554,20 +623,20 @@ private void doExpireBatchSingle(long deliveryTimeoutMs) throws InterruptedExcep
 
     @Test
     public void testExpiredBatchSingle() throws InterruptedException {
-        doExpireBatchSingle(3200L);
+        doExpireBatchSingle(3200);
     }
 
     @Test
     public void testExpiredBatchSingleMaxValue() throws InterruptedException {
-        doExpireBatchSingle(Long.MAX_VALUE);
+        doExpireBatchSingle(Integer.MAX_VALUE);
     }
 
     @Test
     public void testExpiredBatches() throws InterruptedException {
         long retryBackoffMs = 100L;
-        long lingerMs = 30L;
+        int lingerMs = 30;
         int requestTimeout = 60;
-        long deliveryTimeoutMs = 3200L;
+        int deliveryTimeoutMs = 3200;
 
         // test case assumes that the records do not fill the batch completely
         int batchSize = 1025;
@@ -578,11 +647,11 @@ public void testExpiredBatches() throws InterruptedException {
 
         // Test batches not in retry
         for (int i = 0; i < appends; i++) {
-            accum.append(tp1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs);
+            accum.append(tp1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false);
             assertEquals("No partitions should be ready.", 0, accum.ready(cluster, time.milliseconds()).readyNodes.size());
         }
         // Make the batches ready due to batch full
-        accum.append(tp1, 0L, key, value, Record.EMPTY_HEADERS, null, 0);
+        accum.append(tp1, 0L, key, value, Record.EMPTY_HEADERS, null, 0, false);
         Set readyNodes = accum.ready(cluster, time.milliseconds()).readyNodes;
         assertEquals("Our partition's leader should be ready", Collections.singleton(node1), readyNodes);
         // Advance the clock to expire the batch.
@@ -612,7 +681,7 @@ public void testExpiredBatches() throws InterruptedException {
 
         // Test batches in retry.
         // Create a retried batch
-        accum.append(tp1, 0L, key, value, Record.EMPTY_HEADERS, null, 0);
+        accum.append(tp1, 0L, key, value, Record.EMPTY_HEADERS, null, 0, false);
         time.sleep(lingerMs);
         readyNodes = accum.ready(cluster, time.milliseconds()).readyNodes;
         assertEquals("Our partition's leader should be ready", Collections.singleton(node1), readyNodes);
@@ -636,7 +705,7 @@ public void testExpiredBatches() throws InterruptedException {
         assertEquals("All batches should have been expired.", 0, expiredBatches.size());
 
         // Test that when being throttled muted batches are expired before the throttle time is over.
-        accum.append(tp1, 0L, key, value, Record.EMPTY_HEADERS, null, 0);
+        accum.append(tp1, 0L, key, value, Record.EMPTY_HEADERS, null, 0, false);
         time.sleep(lingerMs);
         readyNodes = accum.ready(cluster, time.milliseconds()).readyNodes;
         assertEquals("Our partition's leader should be ready", Collections.singleton(node1), readyNodes);
@@ -669,7 +738,7 @@ public void testMutedPartitions() throws InterruptedException {
                 batchSize + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 10 * batchSize, CompressionType.NONE, 10);
         int appends = expectedNumAppends(batchSize);
         for (int i = 0; i < appends; i++) {
-            accum.append(tp1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs);
+            accum.append(tp1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false);
             assertEquals("No partitions should be ready.", 0, accum.ready(cluster, now).readyNodes.size());
         }
         time.sleep(2000);
@@ -700,19 +769,52 @@ public void testIdempotenceWithOldMagic() throws InterruptedException {
         // Simulate talking to an older broker, ie. one which supports a lower magic.
         ApiVersions apiVersions = new ApiVersions();
         int batchSize = 1025;
-        int requestTimeoutMs = 1600;
-        long deliveryTimeoutMs = 3200L;
-        long lingerMs = 10L;
+        int deliveryTimeoutMs = 3200;
+        int lingerMs = 10;
         long retryBackoffMs = 100L;
         long totalSize = 10 * batchSize;
         String metricGrpName = "producer-metrics";
 
-        apiVersions.update("foobar", NodeApiVersions.create(Arrays.asList(new ApiVersionsResponse.ApiVersion(ApiKeys.PRODUCE.id,
-                (short) 0, (short) 2))));
+        apiVersions.update("foobar", NodeApiVersions.create(ApiKeys.PRODUCE.id, (short) 0, (short) 2));
         RecordAccumulator accum = new RecordAccumulator(logContext, batchSize + DefaultRecordBatch.RECORD_BATCH_OVERHEAD,
             CompressionType.NONE, lingerMs, retryBackoffMs, deliveryTimeoutMs, metrics, metricGrpName, time, apiVersions, new TransactionManager(),
             new BufferPool(totalSize, batchSize, metrics, time, metricGrpName));
-        accum.append(tp1, 0L, key, value, Record.EMPTY_HEADERS, null, 0);
+        accum.append(tp1, 0L, key, value, Record.EMPTY_HEADERS, null, 0, false);
+    }
+
+    @Test
+    public void testRecordHeadersWithMagicV1() {
+        ByteBuffer buffer = ByteBuffer.allocate(4096);
+        MemoryRecordsBuilder builder = MemoryRecords.builder(
+            buffer,
+            (byte) 1, // Enforces that Message format v1 is used on the client
+            CompressionType.NONE,
+            TimestampType.CREATE_TIME,
+            0L
+        );
+        long now = System.currentTimeMillis();
+        ProducerBatch batch = new ProducerBatch(tp1, builder, now, true);
+        try {
+            batch.tryAppend(
+                1L,
+                new byte[10],
+                new byte[20],
+                new Header[] {new RecordHeader("key", new byte[20])},
+                null,
+                now
+            );
+            throw new IllegalStateException("The append of a record with non internal headers should have failed");
+        } catch (IllegalArgumentException exception) {
+            assertEquals(exception.getMessage(), "Magic v1 does not support record headers. [Key = key]");
+        }
+        batch.tryAppend(
+                1L,
+                new byte[10],
+                new byte[20],
+                new Header[] {new RecordHeader("_key", new byte[20])},
+                null,
+                now
+        );
     }
 
     @Test
@@ -777,7 +879,7 @@ public void testSplitBatchOffAccumulator() throws InterruptedException {
 
         // First set the compression ratio estimation to be good.
         CompressionRatioEstimator.setEstimation(tp1.topic(), CompressionType.GZIP, 0.1f);
-        RecordAccumulator accum = createTestRecordAccumulator(batchSize, bufferCapacity, CompressionType.GZIP, 0L);
+        RecordAccumulator accum = createTestRecordAccumulator(batchSize, bufferCapacity, CompressionType.GZIP, 0);
         int numSplitBatches = prepareSplitBatches(accum, seed, 100, 20);
         assertTrue("There should be some split batches", numSplitBatches > 0);
         // Drain all the split batches.
@@ -812,7 +914,7 @@ public void testSplitFrequency() throws InterruptedException {
                 int dice = random.nextInt(100);
                 byte[] value = (dice < goodCompRatioPercentage) ?
                         bytesWithGoodCompression(random) : bytesWithPoorCompression(random, 100);
-                accum.append(tp1, 0L, null, value, Record.EMPTY_HEADERS, null, 0);
+                accum.append(tp1, 0L, null, value, Record.EMPTY_HEADERS, null, 0, false);
                 BatchDrainedResult result = completeOrSplitBatches(accum, batchSize);
                 numSplit += result.numSplit;
                 numBatches += result.numBatches;
@@ -829,13 +931,13 @@ public void testSplitFrequency() throws InterruptedException {
 
     @Test
     public void testSoonToExpireBatchesArePickedUpForExpiry() throws InterruptedException {
-        long lingerMs = 500L;
+        int lingerMs = 500;
         int batchSize = 1025;
 
         RecordAccumulator accum = createTestRecordAccumulator(
             batchSize + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 10 * batchSize, CompressionType.NONE, lingerMs);
 
-        accum.append(tp1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs);
+        accum.append(tp1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false);
         Set readyNodes = accum.ready(cluster, time.milliseconds()).readyNodes;
         Map> drained = accum.drain(cluster, readyNodes, Integer.MAX_VALUE, time.milliseconds());
         assertTrue(drained.isEmpty());
@@ -850,7 +952,7 @@ public void testSoonToExpireBatchesArePickedUpForExpiry() throws InterruptedExce
         //assertTrue(accum.soonToExpireInFlightBatches().isEmpty());
 
         // Queue another batch and advance clock such that batch expiry time is earlier than request timeout.
-        accum.append(tp2, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs);
+        accum.append(tp2, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false);
         time.sleep(lingerMs * 4);
 
         // Now drain and check that accumulator picked up the drained batch because its expiry is soon.
@@ -874,8 +976,8 @@ public void testExpiredBatchesRetry() throws InterruptedException {
             batchSize + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 10 * batchSize, CompressionType.NONE, lingerMs);
 
         // Test batches in retry.
-        for (Boolean mute: muteStates) {
-            accum.append(tp1, 0L, key, value, Record.EMPTY_HEADERS, null, 0);
+        for (Boolean mute : muteStates) {
+            accum.append(tp1, 0L, key, value, Record.EMPTY_HEADERS, null, 0, false);
             time.sleep(lingerMs);
             readyNodes = accum.ready(cluster, time.milliseconds()).readyNodes;
             assertEquals("Our partition's leader should be ready", Collections.singleton(node1), readyNodes);
@@ -893,8 +995,82 @@ public void testExpiredBatchesRetry() throws InterruptedException {
             time.sleep(deliveryTimeoutMs - rtt);
             accum.drain(cluster, Collections.singleton(node1), Integer.MAX_VALUE, time.milliseconds());
             expiredBatches = accum.expiredBatches(time.milliseconds());
-            assertEquals("RecordAccumulator has expired batches if the partition is not muted", mute  ? 1 : 0, expiredBatches.size());
+            assertEquals("RecordAccumulator has expired batches if the partition is not muted", mute ? 1 : 0, expiredBatches.size());
+        }
+    }
+        
+    @Test
+    public void testStickyBatches() throws Exception {
+        long now = time.milliseconds();
+
+        // Test case assumes that the records do not fill the batch completely
+        int batchSize = 1025;
+
+        Partitioner partitioner = new DefaultPartitioner();
+        RecordAccumulator accum = createTestRecordAccumulator(3200,
+                batchSize + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 10L * batchSize, CompressionType.NONE, 10);
+        int expectedAppends = expectedNumAppendsNoKey(batchSize);
+        
+        // Create first batch
+        int partition = partitioner.partition(topic, null, null, "value", value, cluster);
+        TopicPartition tp = new TopicPartition(topic, partition);
+        accum.append(tp, 0L, null, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false);
+        int appends = 1;
+        
+        boolean switchPartition = false;
+        while (!switchPartition) {
+            // Append to the first batch
+            partition = partitioner.partition(topic, null, null, "value", value, cluster);
+            tp = new TopicPartition(topic, partition);
+            RecordAccumulator.RecordAppendResult result = accum.append(tp, 0L, null, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, true);
+            Deque partitionBatches1 = accum.batches().get(tp1);
+            Deque partitionBatches2 = accum.batches().get(tp2);
+            Deque partitionBatches3 = accum.batches().get(tp3);
+            int numBatches = (partitionBatches1 == null ? 0 : partitionBatches1.size()) + (partitionBatches2 == null ? 0 : partitionBatches2.size()) + (partitionBatches3 == null ? 0 : partitionBatches3.size());
+            // Only one batch is created because the partition is sticky.
+            assertEquals(1, numBatches);
+
+            switchPartition = result.abortForNewBatch;
+            // We only appended if we do not retry.
+            if (!switchPartition) {
+                appends++;
+                assertEquals("No partitions should be ready.", 0, accum.ready(cluster, now).readyNodes.size());
+            }
+        }
+        
+        // Batch should be full.
+        assertEquals(1, accum.ready(cluster, time.milliseconds()).readyNodes.size());
+        assertEquals(appends, expectedAppends);
+        switchPartition = false;
+        
+        // KafkaProducer would call this method in this case, make second batch
+        partitioner.onNewBatch(topic, cluster, partition);
+        partition = partitioner.partition(topic, null, null, "value", value, cluster);
+        tp = new TopicPartition(topic, partition);
+        accum.append(tp, 0L, null, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false);
+        appends++;
+        
+        // These appends all go into the second batch
+        while (!switchPartition) {
+            partition = partitioner.partition(topic, null, null, "value", value, cluster);
+            tp = new TopicPartition(topic, partition);
+            RecordAccumulator.RecordAppendResult result = accum.append(tp, 0L, null, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, true);
+            Deque partitionBatches1 = accum.batches().get(tp1);
+            Deque partitionBatches2 = accum.batches().get(tp2);
+            Deque partitionBatches3 = accum.batches().get(tp3);
+            int numBatches = (partitionBatches1 == null ? 0 : partitionBatches1.size()) + (partitionBatches2 == null ? 0 : partitionBatches2.size()) + (partitionBatches3 == null ? 0 : partitionBatches3.size());
+            // Only two batches because the new partition is also sticky.
+            assertEquals(2, numBatches);
+            
+            switchPartition = result.abortForNewBatch;
+            // We only appended if we do not retry.
+            if (!switchPartition) {
+                appends++;
+            }
         }
+        
+        // There should be two full batches now.
+        assertEquals(appends, 2 * expectedAppends);
     }
 
     private int prepareSplitBatches(RecordAccumulator accum, long seed, int recordSize, int numRecords)
@@ -906,7 +1082,7 @@ private int prepareSplitBatches(RecordAccumulator accum, long seed, int recordSi
         CompressionRatioEstimator.setEstimation(tp1.topic(), CompressionType.GZIP, 0.1f);
         // Append 20 records of 100 bytes size with poor compression ratio should make the batch too big.
         for (int i = 0; i < numRecords; i++) {
-            accum.append(tp1, 0L, null, bytesWithPoorCompression(random, recordSize), Record.EMPTY_HEADERS, null, 0);
+            accum.append(tp1, 0L, null, bytesWithPoorCompression(random, recordSize), Record.EMPTY_HEADERS, null, 0, false);
         }
 
         RecordAccumulator.ReadyCheckResult result = accum.ready(cluster, time.milliseconds());
@@ -991,19 +1167,34 @@ private int expectedNumAppends(int batchSize) {
             size += recordSize;
         }
     }
+    
+     /**
+     * Return the offset delta when there is no key.
+     */
+    private int expectedNumAppendsNoKey(int batchSize) {
+        int size = 0;
+        int offsetDelta = 0;
+        while (true) {
+            int recordSize = DefaultRecord.sizeInBytes(offsetDelta, 0, 0, value.length,
+                Record.EMPTY_HEADERS);
+            if (size + recordSize > batchSize)
+                return offsetDelta;
+            offsetDelta += 1;
+            size += recordSize;
+        }
+    }
 
 
-    private RecordAccumulator createTestRecordAccumulator(int batchSize, long totalSize, CompressionType type, long lingerMs) {
-        long deliveryTimeoutMs = 3200L;
+    private RecordAccumulator createTestRecordAccumulator(int batchSize, long totalSize, CompressionType type, int lingerMs) {
+        int deliveryTimeoutMs = 3200;
         return createTestRecordAccumulator(deliveryTimeoutMs, batchSize, totalSize, type, lingerMs);
     }
 
     /**
      * Return a test RecordAccumulator instance
      */
-    private RecordAccumulator createTestRecordAccumulator(long deliveryTimeoutMs, int batchSize, long totalSize, CompressionType type, long lingerMs) {
+    private RecordAccumulator createTestRecordAccumulator(int deliveryTimeoutMs, int batchSize, long totalSize, CompressionType type, int lingerMs) {
         long retryBackoffMs = 100L;
-        int requestTimeoutMs = 1600;
         String metricGrpName = "producer-metrics";
 
         return new RecordAccumulator(
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 4cbdaa25fa0b7..1b35e0b43116e 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
@@ -25,6 +25,7 @@
 import org.apache.kafka.clients.producer.Callback;
 import org.apache.kafka.clients.producer.RecordMetadata;
 import org.apache.kafka.common.Cluster;
+import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.MetricName;
 import org.apache.kafka.common.MetricNameTemplate;
 import org.apache.kafka.common.Node;
@@ -38,6 +39,7 @@
 import org.apache.kafka.common.errors.UnsupportedForMessageFormatException;
 import org.apache.kafka.common.errors.UnsupportedVersionException;
 import org.apache.kafka.common.internals.ClusterResourceListeners;
+import org.apache.kafka.common.message.InitProducerIdResponseData;
 import org.apache.kafka.common.metrics.KafkaMetric;
 import org.apache.kafka.common.metrics.MetricConfig;
 import org.apache.kafka.common.metrics.Metrics;
@@ -54,13 +56,15 @@
 import org.apache.kafka.common.requests.AbstractRequest;
 import org.apache.kafka.common.requests.AddPartitionsToTxnResponse;
 import org.apache.kafka.common.requests.ApiVersionsResponse;
+import org.apache.kafka.common.requests.EndTxnRequest;
+import org.apache.kafka.common.requests.EndTxnResponse;
 import org.apache.kafka.common.requests.FindCoordinatorResponse;
 import org.apache.kafka.common.requests.InitProducerIdRequest;
 import org.apache.kafka.common.requests.InitProducerIdResponse;
 import org.apache.kafka.common.requests.MetadataResponse;
 import org.apache.kafka.common.requests.ProduceRequest;
 import org.apache.kafka.common.requests.ProduceResponse;
-import org.apache.kafka.common.requests.ResponseHeader;
+import org.apache.kafka.common.requests.TransactionResult;
 import org.apache.kafka.common.utils.LogContext;
 import org.apache.kafka.common.utils.MockTime;
 import org.apache.kafka.common.utils.Time;
@@ -92,9 +96,11 @@
 import java.util.concurrent.atomic.AtomicReference;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNull;
-import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 import static org.mockito.AdditionalMatchers.geq;
 import static org.mockito.ArgumentMatchers.any;
@@ -106,15 +112,17 @@
 import static org.mockito.Mockito.atLeastOnce;
 import static org.mockito.Mockito.inOrder;
 import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
 
 public class SenderTest {
-
     private static final int MAX_REQUEST_SIZE = 1024 * 1024;
     private static final short ACKS_ALL = -1;
     private static final String CLIENT_ID = "clientId";
     private static final double EPS = 0.0001;
     private static final int MAX_BLOCK_TIMEOUT = 1000;
     private static final int REQUEST_TIMEOUT = 1000;
+    private static final long RETRY_BACKOFF_MS = 50;
 
     private TopicPartition tp0 = new TopicPartition("test", 0);
     private TopicPartition tp1 = new TopicPartition("test", 1);
@@ -144,18 +152,18 @@ public void tearDown() {
     public void testSimple() throws Exception {
         long offset = 0;
         Future future = accumulator.append(tp0, 0L, "key".getBytes(), "value".getBytes(),
-                null, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds()); // connect
-        sender.run(time.milliseconds()); // send produce request
+                null, null, MAX_BLOCK_TIMEOUT, false).future;
+        sender.runOnce(); // connect
+        sender.runOnce(); // send produce request
         assertEquals("We should have a single produce request in flight.", 1, client.inFlightRequestCount());
         assertEquals(1, sender.inFlightBatches(tp0).size());
         assertTrue(client.hasInFlightRequests());
         client.respond(produceResponse(tp0, offset, Errors.NONE, 0));
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertEquals("All requests completed.", 0, client.inFlightRequestCount());
         assertEquals(0, sender.inFlightBatches(tp0).size());
         assertFalse(client.hasInFlightRequests());
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertTrue("Request should be completed", future.isDone());
         assertEquals(offset, future.get().offset());
     }
@@ -171,11 +179,10 @@ public void testMessageFormatDownConversion() throws Exception {
         apiVersions.update("0", NodeApiVersions.create());
 
         Future future = accumulator.append(tp0, 0L, "key".getBytes(), "value".getBytes(),
-                null, null, MAX_BLOCK_TIMEOUT).future;
+                null, null, MAX_BLOCK_TIMEOUT, false).future;
 
         // now the partition leader supports only v2
-        apiVersions.update("0", NodeApiVersions.create(Collections.singleton(
-                new ApiVersionsResponse.ApiVersion(ApiKeys.PRODUCE.id, (short) 0, (short) 2))));
+        apiVersions.update("0", NodeApiVersions.create(ApiKeys.PRODUCE.id, (short) 0, (short) 2));
 
         client.prepareResponse(new MockClient.RequestMatcher() {
             @Override
@@ -191,8 +198,8 @@ public boolean matches(AbstractRequest body) {
             }
         }, produceResponse(tp0, offset, Errors.NONE, 0));
 
-        sender.run(time.milliseconds()); // connect
-        sender.run(time.milliseconds()); // send produce request
+        sender.runOnce(); // connect
+        sender.runOnce(); // send produce request
 
         assertTrue("Request should be completed", future.isDone());
         assertEquals(offset, future.get().offset());
@@ -211,14 +218,13 @@ public void testDownConversionForMismatchedMagicValues() throws Exception {
         apiVersions.update("0", NodeApiVersions.create());
 
         Future future1 = accumulator.append(tp0, 0L, "key".getBytes(), "value".getBytes(),
-                null, null, MAX_BLOCK_TIMEOUT).future;
+                null, null, MAX_BLOCK_TIMEOUT, false).future;
 
         // now the partition leader supports only v2
-        apiVersions.update("0", NodeApiVersions.create(Collections.singleton(
-                new ApiVersionsResponse.ApiVersion(ApiKeys.PRODUCE.id, (short) 0, (short) 2))));
+        apiVersions.update("0", NodeApiVersions.create(ApiKeys.PRODUCE.id, (short) 0, (short) 2));
 
         Future future2 = accumulator.append(tp1, 0L, "key".getBytes(), "value".getBytes(),
-                null, null, MAX_BLOCK_TIMEOUT).future;
+                null, null, MAX_BLOCK_TIMEOUT, false).future;
 
         // start off support produce request v3
         apiVersions.update("0", NodeApiVersions.create());
@@ -248,8 +254,8 @@ public boolean matches(AbstractRequest body) {
             }
         }, produceResponse);
 
-        sender.run(time.milliseconds()); // connect
-        sender.run(time.milliseconds()); // send produce request
+        sender.runOnce(); // connect
+        sender.runOnce(); // send produce request
 
         assertTrue("Request should be completed", future1.isDone());
         assertTrue("Request should be completed", future2.isDone());
@@ -259,7 +265,6 @@ public boolean matches(AbstractRequest body) {
      * Send multiple requests. Verify that the client side quota metrics have the right values
      */
     @Test
-    @SuppressWarnings("deprecation")
     public void testQuotaMetrics() throws Exception {
         MockSelector selector = new MockSelector(time);
         Sensor throttleTimeSensor = Sender.throttleTimeSensor(this.senderMetricsRegistry);
@@ -269,8 +274,8 @@ public void testQuotaMetrics() throws Exception {
                 1000, 1000, 64 * 1024, 64 * 1024, 1000,  ClientDnsLookup.DEFAULT,
                 time, true, new ApiVersions(), throttleTimeSensor, logContext);
 
-        short apiVersionsResponseVersion = ApiKeys.API_VERSIONS.latestVersion();
-        ByteBuffer buffer = ApiVersionsResponse.createApiVersionsResponse(400, RecordBatch.CURRENT_MAGIC_VALUE).serialize(apiVersionsResponseVersion, new ResponseHeader(0));
+        ByteBuffer buffer = ApiVersionsResponse.createApiVersionsResponse(400, RecordBatch.CURRENT_MAGIC_VALUE).
+            serialize(ApiKeys.API_VERSIONS, ApiKeys.API_VERSIONS.latestVersion(), 0);
         selector.delayedReceive(new DelayedReceive(node.idString(), new NetworkReceive(node.idString(), buffer)));
         while (!client.ready(node, time.milliseconds())) {
             client.poll(1, time.milliseconds());
@@ -287,7 +292,8 @@ public void testQuotaMetrics() throws Exception {
             client.send(request, time.milliseconds());
             client.poll(1, time.milliseconds());
             ProduceResponse response = produceResponse(tp0, i, Errors.NONE, throttleTimeMs);
-            buffer = response.serialize(ApiKeys.PRODUCE.latestVersion(), new ResponseHeader(request.correlationId()));
+            buffer = response.
+                serialize(ApiKeys.PRODUCE, ApiKeys.PRODUCE.latestVersion(), request.correlationId());
             selector.completeReceive(new NetworkReceive(node.idString(), buffer));
             client.poll(1, time.milliseconds());
             // If a throttled response is received, advance the time to ensure progress.
@@ -310,14 +316,14 @@ public void testSenderMetricsTemplates() throws Exception {
         metrics = new Metrics(new MetricConfig().tags(clientTags));
         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, 50, null, apiVersions);
+                1, metricsRegistry, time, REQUEST_TIMEOUT, RETRY_BACKOFF_MS, null, apiVersions);
 
         // Append a message so that topic metrics are created
-        accumulator.append(tp0, 0L, "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT);
-        sender.run(time.milliseconds()); // connect
-        sender.run(time.milliseconds()); // send produce request
+        accumulator.append(tp0, 0L, "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false);
+        sender.runOnce(); // connect
+        sender.runOnce(); // send produce request
         client.respond(produceResponse(tp0, 0, Errors.NONE, 0));
-        sender.run(time.milliseconds());
+        sender.runOnce();
         // Create throttle time metrics
         Sender.throttleTimeSensor(metricsRegistry);
 
@@ -338,49 +344,49 @@ public void testRetries() throws Exception {
         SenderMetricsRegistry senderMetrics = new SenderMetricsRegistry(m);
         try {
             Sender sender = new Sender(logContext, client, metadata, this.accumulator, false, MAX_REQUEST_SIZE, ACKS_ALL,
-                    maxRetries, senderMetrics, time, REQUEST_TIMEOUT, 50, null, apiVersions);
+                    maxRetries, senderMetrics, time, REQUEST_TIMEOUT, RETRY_BACKOFF_MS, null, apiVersions);
             // do a successful retry
-            Future future = accumulator.append(tp0, 0L, "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
-            sender.run(time.milliseconds()); // connect
-            sender.run(time.milliseconds()); // send produce request
+            Future future = accumulator.append(tp0, 0L, "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
+            sender.runOnce(); // connect
+            sender.runOnce(); // send produce request
             String id = client.requests().peek().destination();
             Node node = new Node(Integer.parseInt(id), "localhost", 0);
             assertEquals(1, client.inFlightRequestCount());
             assertTrue(client.hasInFlightRequests());
             assertEquals(1, sender.inFlightBatches(tp0).size());
-            assertTrue("Client ready status should be true", client.isReady(node, 0L));
+            assertTrue("Client ready status should be true", client.isReady(node, time.milliseconds()));
             client.disconnect(id);
             assertEquals(0, client.inFlightRequestCount());
             assertFalse(client.hasInFlightRequests());
-            assertFalse("Client ready status should be false", client.isReady(node, 0L));
+            assertFalse("Client ready status should be false", client.isReady(node, time.milliseconds()));
             // the batch is in accumulator.inFlightBatches until it expires
             assertEquals(1, sender.inFlightBatches(tp0).size());
-            sender.run(time.milliseconds()); // receive error
-            sender.run(time.milliseconds()); // reconnect
-            sender.run(time.milliseconds()); // resend
+            sender.runOnce(); // receive error
+            sender.runOnce(); // reconnect
+            sender.runOnce(); // resend
             assertEquals(1, client.inFlightRequestCount());
             assertTrue(client.hasInFlightRequests());
             assertEquals(1, sender.inFlightBatches(tp0).size());
             long offset = 0;
             client.respond(produceResponse(tp0, offset, Errors.NONE, 0));
-            sender.run(time.milliseconds());
+            sender.runOnce();
             assertTrue("Request should have retried and completed", future.isDone());
             assertEquals(offset, future.get().offset());
             assertEquals(0, sender.inFlightBatches(tp0).size());
 
             // do an unsuccessful retry
-            future = accumulator.append(tp0, 0L, "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
-            sender.run(time.milliseconds()); // send produce request
+            future = accumulator.append(tp0, 0L, "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
+            sender.runOnce(); // send produce request
             assertEquals(1, sender.inFlightBatches(tp0).size());
             for (int i = 0; i < maxRetries + 1; i++) {
                 client.disconnect(client.requests().peek().destination());
-                sender.run(time.milliseconds()); // receive error
+                sender.runOnce(); // receive error
                 assertEquals(0, sender.inFlightBatches(tp0).size());
-                sender.run(time.milliseconds()); // reconnect
-                sender.run(time.milliseconds()); // resend
+                sender.runOnce(); // reconnect
+                sender.runOnce(); // resend
                 assertEquals(i > 0 ? 0 : 1, sender.inFlightBatches(tp0).size());
             }
-            sender.run(time.milliseconds());
+            sender.runOnce();
             assertFutureFailure(future, NetworkException.class);
             assertEquals(0, sender.inFlightBatches(tp0).size());
         } finally {
@@ -396,34 +402,34 @@ public void testSendInOrder() throws Exception {
 
         try {
             Sender sender = new Sender(logContext, client, metadata, this.accumulator, true, MAX_REQUEST_SIZE, ACKS_ALL, maxRetries,
-                    senderMetrics, time, REQUEST_TIMEOUT, 50, null, apiVersions);
+                    senderMetrics, time, REQUEST_TIMEOUT, RETRY_BACKOFF_MS, null, apiVersions);
             // Create a two broker cluster, with partition 0 on broker 0 and partition 1 on broker 1
             MetadataResponse metadataUpdate1 = TestUtils.metadataUpdateWith(2, Collections.singletonMap("test", 2));
             client.prepareMetadataUpdate(metadataUpdate1);
 
             // Send the first message.
             TopicPartition tp2 = new TopicPartition("test", 1);
-            accumulator.append(tp2, 0L, "key1".getBytes(), "value1".getBytes(), null, null, MAX_BLOCK_TIMEOUT);
-            sender.run(time.milliseconds()); // connect
-            sender.run(time.milliseconds()); // send produce request
+            accumulator.append(tp2, 0L, "key1".getBytes(), "value1".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false);
+            sender.runOnce(); // connect
+            sender.runOnce(); // send produce request
             String id = client.requests().peek().destination();
             assertEquals(ApiKeys.PRODUCE, client.requests().peek().requestBuilder().apiKey());
             Node node = new Node(Integer.parseInt(id), "localhost", 0);
             assertEquals(1, client.inFlightRequestCount());
             assertTrue(client.hasInFlightRequests());
-            assertTrue("Client ready status should be true", client.isReady(node, 0L));
+            assertTrue("Client ready status should be true", client.isReady(node, time.milliseconds()));
             assertEquals(1, sender.inFlightBatches(tp2).size());
 
             time.sleep(900);
             // Now send another message to tp2
-            accumulator.append(tp2, 0L, "key2".getBytes(), "value2".getBytes(), null, null, MAX_BLOCK_TIMEOUT);
+            accumulator.append(tp2, 0L, "key2".getBytes(), "value2".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false);
 
             // Update metadata before sender receives response from broker 0. Now partition 2 moves to broker 0
             MetadataResponse metadataUpdate2 = TestUtils.metadataUpdateWith(1, Collections.singletonMap("test", 2));
             client.prepareMetadataUpdate(metadataUpdate2);
             // Sender should not send the second message to node 0.
             assertEquals(1, sender.inFlightBatches(tp2).size());
-            sender.run(time.milliseconds());  // receive the response for the previous send, and send the new batch
+            sender.runOnce();  // receive the response for the previous send, and send the new batch
             assertEquals(1, client.inFlightRequestCount());
             assertTrue(client.hasInFlightRequests());
             assertEquals(1, sender.inFlightBatches(tp2).size());
@@ -446,7 +452,7 @@ public void onCompletion(RecordMetadata metadata, Exception exception) {
                 if (exception instanceof TimeoutException) {
                     expiryCallbackCount.incrementAndGet();
                     try {
-                        accumulator.append(tp1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs);
+                        accumulator.append(tp1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false);
                     } catch (InterruptedException e) {
                         throw new RuntimeException("Unexpected interruption", e);
                     }
@@ -456,7 +462,7 @@ public void onCompletion(RecordMetadata metadata, Exception exception) {
         };
 
         for (int i = 0; i < messagesPerBatch; i++)
-            accumulator.append(tp1, 0L, key, value, null, callback, maxBlockTimeMs);
+            accumulator.append(tp1, 0L, key, value, null, callback, maxBlockTimeMs, false);
 
         // Advance the clock to expire the first batch.
         time.sleep(10000);
@@ -471,7 +477,7 @@ public void onCompletion(RecordMetadata metadata, Exception exception) {
         client.disconnect(clusterNode.idString());
         client.blackout(clusterNode, 100);
 
-        sender.run(time.milliseconds());  // We should try to flush the batch, but we expire it instead without sending anything.
+        sender.runOnce();  // We should try to flush the batch, but we expire it instead without sending anything.
         assertEquals("Callbacks not invoked for expiry", messagesPerBatch, expiryCallbackCount.get());
         assertNull("Unexpected exception", unexpectedException.get());
         // Make sure that the reconds were appended back to the batch.
@@ -489,34 +495,34 @@ public void testMetadataTopicExpiry() throws Exception {
         long offset = 0;
         client.updateMetadata(TestUtils.metadataUpdateWith(1, Collections.singletonMap("test", 2)));
 
-        Future future = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds());
+        Future future = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
+        sender.runOnce();
         assertTrue("Topic not added to metadata", metadata.containsTopic(tp0.topic()));
         client.updateMetadata(TestUtils.metadataUpdateWith(1, Collections.singletonMap("test", 2)));
-        sender.run(time.milliseconds());  // send produce request
+        sender.runOnce();  // send produce request
         client.respond(produceResponse(tp0, offset++, Errors.NONE, 0));
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertEquals("Request completed.", 0, client.inFlightRequestCount());
         assertFalse(client.hasInFlightRequests());
         assertEquals(0, sender.inFlightBatches(tp0).size());
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertTrue("Request should be completed", future.isDone());
 
         assertTrue("Topic not retained in metadata list", metadata.containsTopic(tp0.topic()));
         time.sleep(ProducerMetadata.TOPIC_EXPIRY_MS);
         client.updateMetadata(TestUtils.metadataUpdateWith(1, Collections.singletonMap("test", 2)));
         assertFalse("Unused topic has not been expired", metadata.containsTopic(tp0.topic()));
-        future = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds());
+        future = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
+        sender.runOnce();
         assertTrue("Topic not added to metadata", metadata.containsTopic(tp0.topic()));
         client.updateMetadata(TestUtils.metadataUpdateWith(1, Collections.singletonMap("test", 2)));
-        sender.run(time.milliseconds());  // send produce request
+        sender.runOnce();  // send produce request
         client.respond(produceResponse(tp0, offset++, Errors.NONE, 0));
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertEquals("Request completed.", 0, client.inFlightRequestCount());
         assertFalse(client.hasInFlightRequests());
         assertEquals(0, sender.inFlightBatches(tp0).size());
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertTrue("Request should be completed", future.isDone());
     }
 
@@ -548,15 +554,15 @@ public void testClusterAuthorizationExceptionInInitProducerIdRequest() throws Ex
     @Test
     public void testCanRetryWithoutIdempotence() throws Exception {
         // do a successful retry
-        Future future = accumulator.append(tp0, 0L, "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds()); // connect
-        sender.run(time.milliseconds()); // send produce request
+        Future future = accumulator.append(tp0, 0L, "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
+        sender.runOnce(); // connect
+        sender.runOnce(); // send produce request
         String id = client.requests().peek().destination();
         Node node = new Node(Integer.parseInt(id), "localhost", 0);
         assertEquals(1, client.inFlightRequestCount());
         assertTrue(client.hasInFlightRequests());
         assertEquals(1, sender.inFlightBatches(tp0).size());
-        assertTrue("Client ready status should be true", client.isReady(node, 0L));
+        assertTrue("Client ready status should be true", client.isReady(node, time.milliseconds()));
         assertFalse(future.isDone());
 
         client.respond(new MockClient.RequestMatcher() {
@@ -567,7 +573,7 @@ public boolean matches(AbstractRequest body) {
                 return true;
             }
         }, produceResponse(tp0, -1L, Errors.TOPIC_AUTHORIZATION_FAILED, 0));
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertTrue(future.isDone());
         try {
             future.get();
@@ -587,8 +593,8 @@ public void testIdempotenceWithMultipleInflights() throws Exception {
         assertEquals(0, transactionManager.sequenceNumber(tp0).longValue());
 
         // Send first ProduceRequest
-        Future request1 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds());
+        Future request1 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
+        sender.runOnce();
         String nodeId = client.requests().peek().destination();
         Node node = new Node(Integer.valueOf(nodeId), "localhost", 0);
         assertEquals(1, client.inFlightRequestCount());
@@ -596,8 +602,8 @@ public void testIdempotenceWithMultipleInflights() throws Exception {
         assertEquals(OptionalInt.empty(), transactionManager.lastAckedSequence(tp0));
 
         // Send second ProduceRequest
-        Future request2 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds());
+        Future request2 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
+        sender.runOnce();
         assertEquals(2, client.inFlightRequestCount());
         assertEquals(2, transactionManager.sequenceNumber(tp0).longValue());
         assertEquals(OptionalInt.empty(), transactionManager.lastAckedSequence(tp0));
@@ -607,7 +613,7 @@ public void testIdempotenceWithMultipleInflights() throws Exception {
 
         sendIdempotentProducerResponse(0, tp0, Errors.NONE, 0L);
 
-        sender.run(time.milliseconds()); // receive response 0
+        sender.runOnce(); // receive response 0
 
         assertEquals(1, client.inFlightRequestCount());
         assertEquals(OptionalInt.of(0), transactionManager.lastAckedSequence(tp0));
@@ -616,7 +622,7 @@ public void testIdempotenceWithMultipleInflights() throws Exception {
         assertFalse(request2.isDone());
 
         sendIdempotentProducerResponse(1, tp0, Errors.NONE, 1L);
-        sender.run(time.milliseconds()); // receive response 1
+        sender.runOnce(); // receive response 1
         assertEquals(OptionalInt.of(1), transactionManager.lastAckedSequence(tp0));
         assertFalse(client.hasInFlightRequests());
         assertEquals(0, sender.inFlightBatches(tp0).size());
@@ -637,8 +643,8 @@ public void testIdempotenceWithMultipleInflightsRetriedInOrder() throws Exceptio
         assertEquals(0, transactionManager.sequenceNumber(tp0).longValue());
 
         // Send first ProduceRequest
-        Future request1 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds());
+        Future request1 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
+        sender.runOnce();
         String nodeId = client.requests().peek().destination();
         Node node = new Node(Integer.valueOf(nodeId), "localhost", 0);
         assertEquals(1, client.inFlightRequestCount());
@@ -646,12 +652,12 @@ public void testIdempotenceWithMultipleInflightsRetriedInOrder() throws Exceptio
         assertEquals(OptionalInt.empty(), transactionManager.lastAckedSequence(tp0));
 
         // Send second ProduceRequest
-        Future request2 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds());
+        Future request2 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
+        sender.runOnce();
 
          // Send third ProduceRequest
-        Future request3 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds());
+        Future request3 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
+        sender.runOnce();
 
         assertEquals(3, client.inFlightRequestCount());
         assertEquals(3, transactionManager.sequenceNumber(tp0).longValue());
@@ -662,24 +668,24 @@ public void testIdempotenceWithMultipleInflightsRetriedInOrder() throws Exceptio
         assertTrue(client.isReady(node, time.milliseconds()));
 
         sendIdempotentProducerResponse(0, tp0, Errors.LEADER_NOT_AVAILABLE, -1L);
-        sender.run(time.milliseconds()); // receive response 0
+        sender.runOnce(); // receive response 0
 
         // Queue the fourth request, it shouldn't be sent until the first 3 complete.
-        Future request4 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
+        Future request4 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
 
         assertEquals(2, client.inFlightRequestCount());
         assertEquals(OptionalInt.empty(), transactionManager.lastAckedSequence(tp0));
 
         sendIdempotentProducerResponse(1, tp0, Errors.OUT_OF_ORDER_SEQUENCE_NUMBER, -1L);
-        sender.run(time.milliseconds()); // re send request 1, receive response 2
+        sender.runOnce(); // re send request 1, receive response 2
 
         sendIdempotentProducerResponse(2, tp0, Errors.OUT_OF_ORDER_SEQUENCE_NUMBER, -1L);
-        sender.run(time.milliseconds()); // receive response 3
+        sender.runOnce(); // receive response 3
 
         assertEquals(OptionalInt.empty(), transactionManager.lastAckedSequence(tp0));
         assertEquals(1, client.inFlightRequestCount());
 
-        sender.run(time.milliseconds()); // Do nothing, we are reduced to one in flight request during retries.
+        sender.runOnce(); // Do nothing, we are reduced to one in flight request during retries.
 
         assertEquals(3, transactionManager.sequenceNumber(tp0).longValue());  // the batch for request 4 shouldn't have been drained, and hence the sequence should not have been incremented.
         assertEquals(1, client.inFlightRequestCount());
@@ -687,19 +693,19 @@ public void testIdempotenceWithMultipleInflightsRetriedInOrder() throws Exceptio
         assertEquals(OptionalInt.empty(), transactionManager.lastAckedSequence(tp0));
 
         sendIdempotentProducerResponse(0, tp0, Errors.NONE, 0L);
-        sender.run(time.milliseconds());  // receive response 1
+        sender.runOnce();  // receive response 1
         assertEquals(OptionalInt.of(0), transactionManager.lastAckedSequence(tp0));
         assertTrue(request1.isDone());
         assertEquals(0, request1.get().offset());
         assertFalse(client.hasInFlightRequests());
         assertEquals(0, sender.inFlightBatches(tp0).size());
 
-        sender.run(time.milliseconds()); // send request 2;
+        sender.runOnce(); // send request 2;
         assertEquals(1, client.inFlightRequestCount());
         assertEquals(1, sender.inFlightBatches(tp0).size());
 
         sendIdempotentProducerResponse(1, tp0, Errors.NONE, 1L);
-        sender.run(time.milliseconds());  // receive response 2
+        sender.runOnce();  // receive response 2
         assertEquals(OptionalInt.of(1), transactionManager.lastAckedSequence(tp0));
         assertTrue(request2.isDone());
         assertEquals(1, request2.get().offset());
@@ -707,12 +713,12 @@ public void testIdempotenceWithMultipleInflightsRetriedInOrder() throws Exceptio
         assertFalse(client.hasInFlightRequests());
         assertEquals(0, sender.inFlightBatches(tp0).size());
 
-        sender.run(time.milliseconds()); // send request 3
+        sender.runOnce(); // send request 3
         assertEquals(1, client.inFlightRequestCount());
         assertEquals(1, sender.inFlightBatches(tp0).size());
 
         sendIdempotentProducerResponse(2, tp0, Errors.NONE, 2L);
-        sender.run(time.milliseconds());  // receive response 3, send request 4 since we are out of 'retry' mode.
+        sender.runOnce();  // receive response 3, send request 4 since we are out of 'retry' mode.
         assertEquals(OptionalInt.of(2), transactionManager.lastAckedSequence(tp0));
         assertTrue(request3.isDone());
         assertEquals(2, request3.get().offset());
@@ -720,7 +726,7 @@ public void testIdempotenceWithMultipleInflightsRetriedInOrder() throws Exceptio
         assertEquals(1, sender.inFlightBatches(tp0).size());
 
         sendIdempotentProducerResponse(3, tp0, Errors.NONE, 3L);
-        sender.run(time.milliseconds());  // receive response 4
+        sender.runOnce();  // receive response 4
         assertEquals(OptionalInt.of(3), transactionManager.lastAckedSequence(tp0));
         assertTrue(request4.isDone());
         assertEquals(3, request4.get().offset());
@@ -737,8 +743,8 @@ public void testIdempotenceWithMultipleInflightsWhereFirstFailsFatallyAndSequenc
         assertEquals(0, transactionManager.sequenceNumber(tp0).longValue());
 
         // Send first ProduceRequest
-        Future request1 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds());
+        Future request1 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
+        sender.runOnce();
         String nodeId = client.requests().peek().destination();
         Node node = new Node(Integer.valueOf(nodeId), "localhost", 0);
         assertEquals(1, client.inFlightRequestCount());
@@ -746,8 +752,8 @@ public void testIdempotenceWithMultipleInflightsWhereFirstFailsFatallyAndSequenc
         assertEquals(OptionalInt.empty(), transactionManager.lastAckedSequence(tp0));
 
         // Send second ProduceRequest
-        Future request2 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds());
+        Future request2 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
+        sender.runOnce();
         assertEquals(2, client.inFlightRequestCount());
         assertEquals(2, transactionManager.sequenceNumber(tp0).longValue());
         assertEquals(OptionalInt.empty(), transactionManager.lastAckedSequence(tp0));
@@ -757,7 +763,7 @@ public void testIdempotenceWithMultipleInflightsWhereFirstFailsFatallyAndSequenc
 
         sendIdempotentProducerResponse(0, tp0, Errors.MESSAGE_TOO_LARGE, -1L);
 
-        sender.run(time.milliseconds()); // receive response 0, should adjust sequences of future batches.
+        sender.runOnce(); // receive response 0, should adjust sequences of future batches.
         assertFutureFailure(request1, RecordTooLargeException.class);
 
         assertEquals(1, client.inFlightRequestCount());
@@ -765,19 +771,19 @@ public void testIdempotenceWithMultipleInflightsWhereFirstFailsFatallyAndSequenc
 
         sendIdempotentProducerResponse(1, tp0, Errors.OUT_OF_ORDER_SEQUENCE_NUMBER, -1L);
 
-        sender.run(time.milliseconds()); // receive response 1
+        sender.runOnce(); // receive response 1
 
         assertEquals(OptionalInt.empty(), transactionManager.lastAckedSequence(tp0));
         assertEquals(0, client.inFlightRequestCount());
 
-        sender.run(time.milliseconds()); // resend request 1
+        sender.runOnce(); // resend request 1
 
         assertEquals(1, client.inFlightRequestCount());
 
         assertEquals(OptionalInt.empty(), transactionManager.lastAckedSequence(tp0));
 
         sendIdempotentProducerResponse(0, tp0, Errors.NONE, 0L);
-        sender.run(time.milliseconds());  // receive response 1
+        sender.runOnce();  // receive response 1
         assertEquals(OptionalInt.of(0), transactionManager.lastAckedSequence(tp0));
         assertEquals(0, client.inFlightRequestCount());
 
@@ -796,9 +802,9 @@ public void testMustNotRetryOutOfOrderSequenceForNextBatch() throws Exception {
         assertEquals(0, transactionManager.sequenceNumber(tp0).longValue());
 
         // Send first ProduceRequest with multiple messages.
-        Future request1 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
-        accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT);
-        sender.run(time.milliseconds());
+        Future request1 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
+        accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false);
+        sender.runOnce();
         String nodeId = client.requests().peek().destination();
         Node node = new Node(Integer.valueOf(nodeId), "localhost", 0);
         assertEquals(1, client.inFlightRequestCount());
@@ -808,11 +814,11 @@ public void testMustNotRetryOutOfOrderSequenceForNextBatch() throws Exception {
         assertEquals(OptionalInt.empty(), transactionManager.lastAckedSequence(tp0));
         sendIdempotentProducerResponse(0, tp0, Errors.NONE, 0);
 
-        sender.run(time.milliseconds());
+        sender.runOnce();
 
         // Send second ProduceRequest
-        Future request2 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds());
+        Future request2 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
+        sender.runOnce();
         assertEquals(1, client.inFlightRequestCount());
         assertEquals(3, transactionManager.sequenceNumber(tp0).longValue());
         assertEquals(OptionalInt.of(1), transactionManager.lastAckedSequence(tp0));
@@ -824,7 +830,7 @@ public void testMustNotRetryOutOfOrderSequenceForNextBatch() throws Exception {
         // This OutOfOrderSequence is fatal since it is returned for the batch succeeding the last acknowledged batch.
         sendIdempotentProducerResponse(2, tp0, Errors.OUT_OF_ORDER_SEQUENCE_NUMBER, -1L);
 
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertFutureFailure(request2, OutOfOrderSequenceException.class);
     }
 
@@ -838,8 +844,8 @@ public void testCorrectHandlingOfOutOfOrderResponses() throws Exception {
         assertEquals(0, transactionManager.sequenceNumber(tp0).longValue());
 
         // Send first ProduceRequest
-        Future request1 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds());
+        Future request1 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
+        sender.runOnce();
         String nodeId = client.requests().peek().destination();
         Node node = new Node(Integer.valueOf(nodeId), "localhost", 0);
         assertEquals(1, client.inFlightRequestCount());
@@ -847,8 +853,8 @@ public void testCorrectHandlingOfOutOfOrderResponses() throws Exception {
         assertEquals(OptionalInt.empty(), transactionManager.lastAckedSequence(tp0));
 
         // Send second ProduceRequest
-        Future request2 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds());
+        Future request2 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
+        sender.runOnce();
         assertEquals(2, client.inFlightRequestCount());
         assertEquals(2, transactionManager.sequenceNumber(tp0).longValue());
         assertEquals(OptionalInt.empty(), transactionManager.lastAckedSequence(tp0));
@@ -861,7 +867,7 @@ public void testCorrectHandlingOfOutOfOrderResponses() throws Exception {
 
         client.respondToRequest(secondClientRequest, produceResponse(tp0, -1, Errors.OUT_OF_ORDER_SEQUENCE_NUMBER, -1));
 
-        sender.run(time.milliseconds()); // receive response 1
+        sender.runOnce(); // receive response 1
         Deque queuedBatches = accumulator.batches().get(tp0);
 
         // Make sure that we are queueing the second batch first.
@@ -872,7 +878,7 @@ public void testCorrectHandlingOfOutOfOrderResponses() throws Exception {
 
         client.respondToRequest(firstClientRequest, produceResponse(tp0, -1, Errors.NOT_LEADER_FOR_PARTITION, -1));
 
-        sender.run(time.milliseconds()); // receive response 0
+        sender.runOnce(); // receive response 0
 
         // Make sure we requeued both batches in the correct order.
         assertEquals(2, queuedBatches.size());
@@ -883,25 +889,25 @@ public void testCorrectHandlingOfOutOfOrderResponses() throws Exception {
         assertFalse(request1.isDone());
         assertFalse(request2.isDone());
 
-        sender.run(time.milliseconds()); // send request 0
+        sender.runOnce(); // send request 0
         assertEquals(1, client.inFlightRequestCount());
-        sender.run(time.milliseconds()); // don't do anything, only one inflight allowed once we are retrying.
+        sender.runOnce(); // don't do anything, only one inflight allowed once we are retrying.
 
         assertEquals(1, client.inFlightRequestCount());
         assertEquals(OptionalInt.empty(), transactionManager.lastAckedSequence(tp0));
 
         // Make sure that the requests are sent in order, even though the previous responses were not in order.
         sendIdempotentProducerResponse(0, tp0, Errors.NONE, 0L);
-        sender.run(time.milliseconds());  // receive response 0
+        sender.runOnce();  // receive response 0
         assertEquals(OptionalInt.of(0), transactionManager.lastAckedSequence(tp0));
         assertEquals(0, client.inFlightRequestCount());
         assertTrue(request1.isDone());
         assertEquals(0, request1.get().offset());
 
-        sender.run(time.milliseconds()); // send request 1
+        sender.runOnce(); // send request 1
         assertEquals(1, client.inFlightRequestCount());
         sendIdempotentProducerResponse(1, tp0, Errors.NONE, 1L);
-        sender.run(time.milliseconds());  // receive response 1
+        sender.runOnce();  // receive response 1
 
         assertFalse(client.hasInFlightRequests());
         assertEquals(OptionalInt.of(1), transactionManager.lastAckedSequence(tp0));
@@ -920,15 +926,15 @@ public void testCorrectHandlingOfOutOfOrderResponsesWhenSecondSucceeds() throws
         assertEquals(0, transactionManager.sequenceNumber(tp0).longValue());
 
         // Send first ProduceRequest
-        Future request1 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds());
+        Future request1 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
+        sender.runOnce();
         String nodeId = client.requests().peek().destination();
         Node node = new Node(Integer.valueOf(nodeId), "localhost", 0);
         assertEquals(1, client.inFlightRequestCount());
 
         // Send second ProduceRequest
-        Future request2 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds());
+        Future request2 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
+        sender.runOnce();
         assertEquals(2, client.inFlightRequestCount());
         assertFalse(request1.isDone());
         assertFalse(request2.isDone());
@@ -939,7 +945,7 @@ public void testCorrectHandlingOfOutOfOrderResponsesWhenSecondSucceeds() throws
 
         client.respondToRequest(secondClientRequest, produceResponse(tp0, 1, Errors.NONE, 1));
 
-        sender.run(time.milliseconds()); // receive response 1
+        sender.runOnce(); // receive response 1
         assertTrue(request2.isDone());
         assertEquals(1, request2.get().offset());
         assertFalse(request1.isDone());
@@ -951,7 +957,7 @@ public void testCorrectHandlingOfOutOfOrderResponsesWhenSecondSucceeds() throws
 
         client.respondToRequest(firstClientRequest, produceResponse(tp0, -1, Errors.REQUEST_TIMED_OUT, -1));
 
-        sender.run(time.milliseconds()); // receive response 0
+        sender.runOnce(); // receive response 0
 
         // Make sure we requeued both batches in the correct order.
         assertEquals(1, queuedBatches.size());
@@ -959,7 +965,7 @@ public void testCorrectHandlingOfOutOfOrderResponsesWhenSecondSucceeds() throws
         assertEquals(OptionalInt.of(1), transactionManager.lastAckedSequence(tp0));
         assertEquals(0, client.inFlightRequestCount());
 
-        sender.run(time.milliseconds()); // resend request 0
+        sender.runOnce(); // resend request 0
         assertEquals(1, client.inFlightRequestCount());
 
         assertEquals(1, client.inFlightRequestCount());
@@ -967,7 +973,7 @@ public void testCorrectHandlingOfOutOfOrderResponsesWhenSecondSucceeds() throws
 
         // Make sure we handle the out of order successful responses correctly.
         sendIdempotentProducerResponse(0, tp0, Errors.NONE, 0L);
-        sender.run(time.milliseconds());  // receive response 0
+        sender.runOnce();  // receive response 0
         assertEquals(0, queuedBatches.size());
         assertEquals(OptionalInt.of(1), transactionManager.lastAckedSequence(tp0));
         assertEquals(0, client.inFlightRequestCount());
@@ -988,13 +994,13 @@ public void testExpiryOfUnsentBatchesShouldNotCauseUnresolvedSequences() throws
         assertEquals(0, transactionManager.sequenceNumber(tp0).longValue());
 
         // Send first ProduceRequest
-        Future request1 = accumulator.append(tp0, 0L, "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
+        Future request1 = accumulator.append(tp0, 0L, "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
         Node node = metadata.fetch().nodes().get(0);
         time.sleep(10000L);
         client.disconnect(node.idString());
         client.blackout(node, 10);
 
-        sender.run(time.milliseconds());
+        sender.runOnce();
 
         assertFutureFailure(request1, TimeoutException.class);
         assertFalse(transactionManager.hasUnresolvedSequence(tp0));
@@ -1010,19 +1016,19 @@ public void testExpiryOfFirstBatchShouldNotCauseUnresolvedSequencesIfFutureBatch
         assertEquals(0, transactionManager.sequenceNumber(tp0).longValue());
 
         // Send first ProduceRequest
-        Future request1 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds());  // send request
+        Future request1 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
+        sender.runOnce();  // send request
         // We separate the two appends by 1 second so that the two batches
         // don't expire at the same time.
         time.sleep(1000L);
 
-        Future request2 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds());  // send request
+        Future request2 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
+        sender.runOnce();  // send request
         assertEquals(2, client.inFlightRequestCount());
         assertEquals(2, sender.inFlightBatches(tp0).size());
 
         sendIdempotentProducerResponse(0, tp0, Errors.REQUEST_TIMED_OUT, -1);
-        sender.run(time.milliseconds());  // receive first response
+        sender.runOnce();  // receive first response
         assertEquals(1, sender.inFlightBatches(tp0).size());
 
         Node node = metadata.fetch().nodes().get(0);
@@ -1032,21 +1038,21 @@ public void testExpiryOfFirstBatchShouldNotCauseUnresolvedSequencesIfFutureBatch
         client.disconnect(node.idString());
         client.blackout(node, 10);
 
-        sender.run(time.milliseconds()); // now expire the first batch.
+        sender.runOnce(); // now expire the first batch.
         assertFutureFailure(request1, TimeoutException.class);
         assertTrue(transactionManager.hasUnresolvedSequence(tp0));
         assertEquals(0, sender.inFlightBatches(tp0).size());
 
         // let's enqueue another batch, which should not be dequeued until the unresolved state is clear.
-        Future request3 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
+        Future request3 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
         time.sleep(20);
         assertFalse(request2.isDone());
 
-        sender.run(time.milliseconds());  // send second request
+        sender.runOnce();  // send second request
         sendIdempotentProducerResponse(1, tp0, Errors.NONE, 1);
         assertEquals(1, sender.inFlightBatches(tp0).size());
 
-        sender.run(time.milliseconds()); // receive second response, the third request shouldn't be sent since we are in an unresolved state.
+        sender.runOnce(); // receive second response, the third request shouldn't be sent since we are in an unresolved state.
         assertTrue(request2.isDone());
         assertEquals(1, request2.get().offset());
         assertEquals(0, sender.inFlightBatches(tp0).size());
@@ -1058,7 +1064,7 @@ public void testExpiryOfFirstBatchShouldNotCauseUnresolvedSequencesIfFutureBatch
         assertEquals(2L, transactionManager.sequenceNumber(tp0).longValue());
         assertTrue(transactionManager.hasUnresolvedSequence(tp0));
 
-        sender.run(time.milliseconds());  // clear the unresolved state, send the pending request.
+        sender.runOnce();  // clear the unresolved state, send the pending request.
         assertFalse(transactionManager.hasUnresolvedSequence(tp0));
         assertTrue(transactionManager.hasProducerId());
         assertEquals(0, batches.size());
@@ -1077,34 +1083,34 @@ public void testExpiryOfFirstBatchShouldCauseResetIfFutureBatchesFail() throws E
         assertEquals(0, transactionManager.sequenceNumber(tp0).longValue());
 
         // Send first ProduceRequest
-        Future request1 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds());  // send request
+        Future request1 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
+        sender.runOnce();  // send request
 
         time.sleep(1000L);
-        Future request2 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds());  // send request
+        Future request2 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
+        sender.runOnce();  // send request
 
         assertEquals(2, client.inFlightRequestCount());
 
         sendIdempotentProducerResponse(0, tp0, Errors.NOT_LEADER_FOR_PARTITION, -1);
-        sender.run(time.milliseconds());  // receive first response
+        sender.runOnce();  // receive first response
 
         Node node = metadata.fetch().nodes().get(0);
         time.sleep(1000L);
         client.disconnect(node.idString());
         client.blackout(node, 10);
 
-        sender.run(time.milliseconds()); // now expire the first batch.
+        sender.runOnce(); // now expire the first batch.
         assertFutureFailure(request1, TimeoutException.class);
         assertTrue(transactionManager.hasUnresolvedSequence(tp0));
         // let's enqueue another batch, which should not be dequeued until the unresolved state is clear.
-        Future request3 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
+        Future request3 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
 
         time.sleep(20);
         assertFalse(request2.isDone());
-        sender.run(time.milliseconds());  // send second request
+        sender.runOnce();  // send second request
         sendIdempotentProducerResponse(1, tp0, Errors.OUT_OF_ORDER_SEQUENCE_NUMBER, 1);
-        sender.run(time.milliseconds()); // receive second response, the third request shouldn't be sent since we are in an unresolved state.
+        sender.runOnce(); // receive second response, the third request shouldn't be sent since we are in an unresolved state.
         assertFutureFailure(request2, OutOfOrderSequenceException.class);
 
         Deque batches = accumulator.batches().get(tp0);
@@ -1130,11 +1136,11 @@ public void testExpiryOfAllSentBatchesShouldCauseUnresolvedSequences() throws Ex
         assertEquals(0, transactionManager.sequenceNumber(tp0).longValue());
 
         // Send first ProduceRequest
-        Future request1 = accumulator.append(tp0, 0L, "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds());  // send request
+        Future request1 = accumulator.append(tp0, 0L, "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
+        sender.runOnce();  // send request
         sendIdempotentProducerResponse(0, tp0, Errors.NOT_LEADER_FOR_PARTITION, -1);
 
-        sender.run(time.milliseconds());  // receive response
+        sender.runOnce();  // receive response
         assertEquals(1L, transactionManager.sequenceNumber(tp0).longValue());
 
         Node node = metadata.fetch().nodes().get(0);
@@ -1142,7 +1148,7 @@ public void testExpiryOfAllSentBatchesShouldCauseUnresolvedSequences() throws Ex
         client.disconnect(node.idString());
         client.blackout(node, 10);
 
-        sender.run(time.milliseconds()); // now expire the batch.
+        sender.runOnce(); // now expire the batch.
 
         assertFutureFailure(request1, TimeoutException.class);
         assertTrue(transactionManager.hasUnresolvedSequence(tp0));
@@ -1167,13 +1173,13 @@ public void testResetOfProducerStateShouldAllowQueuedBatchesToDrain() throws Exc
         SenderMetricsRegistry senderMetrics = new SenderMetricsRegistry(m);
 
         Sender sender = new Sender(logContext, client, metadata, this.accumulator, true, MAX_REQUEST_SIZE, ACKS_ALL, maxRetries,
-                senderMetrics, time, REQUEST_TIMEOUT, 50, transactionManager, apiVersions);
+                senderMetrics, time, REQUEST_TIMEOUT, RETRY_BACKOFF_MS, transactionManager, apiVersions);
 
         Future failedResponse = accumulator.append(tp0, time.milliseconds(), "key".getBytes(),
-                "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
+                "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
         Future successfulResponse = accumulator.append(tp1, time.milliseconds(), "key".getBytes(),
-                "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds());  // connect and send.
+                "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
+        sender.runOnce();  // connect and send.
 
         assertEquals(1, client.inFlightRequestCount());
 
@@ -1181,7 +1187,7 @@ public void testResetOfProducerStateShouldAllowQueuedBatchesToDrain() throws Exc
         responses.put(tp1, new OffsetAndError(-1, Errors.NOT_LEADER_FOR_PARTITION));
         responses.put(tp0, new OffsetAndError(-1, Errors.OUT_OF_ORDER_SEQUENCE_NUMBER));
         client.respond(produceResponse(responses));
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertTrue(failedResponse.isDone());
         assertFalse("Expected transaction state to be reset upon receiving an OutOfOrderSequenceException", transactionManager.hasProducerId());
         prepareAndReceiveInitProducerId(producerId + 1, Errors.NONE); // also send request to tp1
@@ -1189,7 +1195,7 @@ public void testResetOfProducerStateShouldAllowQueuedBatchesToDrain() throws Exc
 
         assertFalse(successfulResponse.isDone());
         client.respond(produceResponse(tp1, 10, Errors.NONE, -1));
-        sender.run(time.milliseconds());
+        sender.runOnce();
 
         assertTrue(successfulResponse.isDone());
         assertEquals(10, successfulResponse.get().offset());
@@ -1209,13 +1215,13 @@ public void testCloseWithProducerIdReset() throws Exception {
         SenderMetricsRegistry senderMetrics = new SenderMetricsRegistry(m);
 
         Sender sender = new Sender(logContext, client, metadata, this.accumulator, true, MAX_REQUEST_SIZE, ACKS_ALL, 10,
-            senderMetrics, time, REQUEST_TIMEOUT, 50, transactionManager, apiVersions);
+            senderMetrics, time, REQUEST_TIMEOUT, RETRY_BACKOFF_MS, transactionManager, apiVersions);
 
         Future failedResponse = accumulator.append(tp0, time.milliseconds(), "key".getBytes(),
-            "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
+            "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
         Future successfulResponse = accumulator.append(tp1, time.milliseconds(), "key".getBytes(),
-            "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds());  // connect and send.
+            "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
+        sender.runOnce();  // connect and send.
 
         assertEquals(1, client.inFlightRequestCount());
 
@@ -1224,15 +1230,15 @@ public void testCloseWithProducerIdReset() throws Exception {
         responses.put(tp0, new OffsetAndError(-1, Errors.OUT_OF_ORDER_SEQUENCE_NUMBER));
         client.respond(produceResponse(responses));
         sender.initiateClose(); // initiate close
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertTrue(failedResponse.isDone());
         assertFalse("Expected transaction state to be reset upon receiving an OutOfOrderSequenceException", transactionManager.hasProducerId());
 
         TestUtils.waitForCondition(new TestCondition() {
             @Override
             public boolean conditionMet() {
-                prepareInitPidResponse(Errors.NONE, producerId + 1, (short) 1);
-                sender.run(time.milliseconds());
+                prepareInitProducerResponse(Errors.NONE, producerId + 1, (short) 1);
+                sender.runOnce();
                 return !accumulator.hasUndrained();
             }
         }, 5000, "Failed to drain batches");
@@ -1248,13 +1254,13 @@ public void testForceCloseWithProducerIdReset() throws Exception {
         SenderMetricsRegistry senderMetrics = new SenderMetricsRegistry(m);
 
         Sender sender = new Sender(logContext, client, metadata, this.accumulator, true, MAX_REQUEST_SIZE, ACKS_ALL, 10,
-            senderMetrics, time, REQUEST_TIMEOUT, 50, transactionManager, apiVersions);
+            senderMetrics, time, REQUEST_TIMEOUT, RETRY_BACKOFF_MS, transactionManager, apiVersions);
 
         Future failedResponse = accumulator.append(tp0, time.milliseconds(), "key".getBytes(),
-            "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
+            "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
         Future successfulResponse = accumulator.append(tp1, time.milliseconds(), "key".getBytes(),
-            "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds());  // connect and send.
+            "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
+        sender.runOnce();  // connect and send.
 
         assertEquals(1, client.inFlightRequestCount());
 
@@ -1262,11 +1268,11 @@ public void testForceCloseWithProducerIdReset() throws Exception {
         responses.put(tp1, new OffsetAndError(-1, Errors.NOT_LEADER_FOR_PARTITION));
         responses.put(tp0, new OffsetAndError(-1, Errors.OUT_OF_ORDER_SEQUENCE_NUMBER));
         client.respond(produceResponse(responses));
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertTrue(failedResponse.isDone());
         assertFalse("Expected transaction state to be reset upon receiving an OutOfOrderSequenceException", transactionManager.hasProducerId());
         sender.forceClose(); // initiate force close
-        sender.run(time.milliseconds()); // this should not block
+        sender.runOnce(); // this should not block
         sender.run(); // run main loop to test forceClose flag
         assertTrue("Pending batches are not aborted.", !accumulator.hasUndrained());
         assertTrue(successfulResponse.isDone());
@@ -1284,14 +1290,14 @@ public void testBatchesDrainedWithOldProducerIdShouldFailWithOutOfOrderSequenceO
         SenderMetricsRegistry senderMetrics = new SenderMetricsRegistry(m);
 
         Sender sender = new Sender(logContext, client, metadata, this.accumulator, true, MAX_REQUEST_SIZE, ACKS_ALL, maxRetries,
-                senderMetrics, time, REQUEST_TIMEOUT, 50, transactionManager, apiVersions);
+                senderMetrics, time, REQUEST_TIMEOUT, RETRY_BACKOFF_MS, transactionManager, apiVersions);
 
         Future failedResponse = accumulator.append(tp0, time.milliseconds(), "key".getBytes(),
-                "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
+                "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
         Future successfulResponse = accumulator.append(tp1, time.milliseconds(), "key".getBytes(),
-                "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds());  // connect.
-        sender.run(time.milliseconds());  // send.
+                "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
+        sender.runOnce();  // connect.
+        sender.runOnce();  // send.
 
         assertEquals(1, client.inFlightRequestCount());
 
@@ -1299,17 +1305,17 @@ public void testBatchesDrainedWithOldProducerIdShouldFailWithOutOfOrderSequenceO
         responses.put(tp1, new OffsetAndError(-1, Errors.NOT_LEADER_FOR_PARTITION));
         responses.put(tp0, new OffsetAndError(-1, Errors.OUT_OF_ORDER_SEQUENCE_NUMBER));
         client.respond(produceResponse(responses));
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertTrue(failedResponse.isDone());
         assertFalse("Expected transaction state to be reset upon receiving an OutOfOrderSequenceException", transactionManager.hasProducerId());
         prepareAndReceiveInitProducerId(producerId + 1, Errors.NONE);
         assertEquals(producerId + 1, transactionManager.producerIdAndEpoch().producerId);
-        sender.run(time.milliseconds());  // send request to tp1 with the old producerId
+        sender.runOnce();  // send request to tp1 with the old producerId
 
         assertFalse(successfulResponse.isDone());
         // The response comes back with a retriable error.
         client.respond(produceResponse(tp1, 0, Errors.NOT_LEADER_FOR_PARTITION, -1));
-        sender.run(time.milliseconds());
+        sender.runOnce();
 
         assertTrue(successfulResponse.isDone());
         // Since the batch has an old producerId, it will not be retried yet again, but will be failed with a Fatal
@@ -1333,8 +1339,8 @@ public void testCorrectHandlingOfDuplicateSequenceError() throws Exception {
         assertEquals(0, transactionManager.sequenceNumber(tp0).longValue());
 
         // Send first ProduceRequest
-        Future request1 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds());
+        Future request1 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
+        sender.runOnce();
         String nodeId = client.requests().peek().destination();
         Node node = new Node(Integer.valueOf(nodeId), "localhost", 0);
         assertEquals(1, client.inFlightRequestCount());
@@ -1342,8 +1348,8 @@ public void testCorrectHandlingOfDuplicateSequenceError() throws Exception {
         assertEquals(OptionalInt.empty(), transactionManager.lastAckedSequence(tp0));
 
         // Send second ProduceRequest
-        Future request2 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds());
+        Future request2 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
+        sender.runOnce();
         assertEquals(2, client.inFlightRequestCount());
         assertEquals(2, transactionManager.sequenceNumber(tp0).longValue());
         assertEquals(OptionalInt.empty(), transactionManager.lastAckedSequence(tp0));
@@ -1356,14 +1362,14 @@ public void testCorrectHandlingOfDuplicateSequenceError() throws Exception {
 
         client.respondToRequest(secondClientRequest, produceResponse(tp0, 1000, Errors.NONE, 0));
 
-        sender.run(time.milliseconds()); // receive response 1
+        sender.runOnce(); // receive response 1
 
         assertEquals(OptionalLong.of(1000), transactionManager.lastAckedOffset(tp0));
         assertEquals(OptionalInt.of(1), transactionManager.lastAckedSequence(tp0));
 
         client.respondToRequest(firstClientRequest, produceResponse(tp0, ProduceResponse.INVALID_OFFSET, Errors.DUPLICATE_SEQUENCE_NUMBER, 0));
 
-        sender.run(time.milliseconds()); // receive response 0
+        sender.runOnce(); // receive response 0
 
         // Make sure that the last ack'd sequence doesn't change.
         assertEquals(OptionalInt.of(1), transactionManager.lastAckedSequence(tp0));
@@ -1386,8 +1392,8 @@ public void testUnknownProducerHandlingWhenRetentionLimitReached() throws Except
         assertEquals(0, transactionManager.sequenceNumber(tp0).longValue());
 
         // Send first ProduceRequest
-        Future request1 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds());
+        Future request1 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
+        sender.runOnce();
 
         assertEquals(1, client.inFlightRequestCount());
         assertEquals(1, transactionManager.sequenceNumber(tp0).longValue());
@@ -1395,7 +1401,7 @@ public void testUnknownProducerHandlingWhenRetentionLimitReached() throws Except
 
         sendIdempotentProducerResponse(0, tp0, Errors.NONE, 1000L, 10L);
 
-        sender.run(time.milliseconds());  // receive the response.
+        sender.runOnce();  // receive the response.
 
         assertTrue(request1.isDone());
         assertEquals(1000L, request1.get().offset());
@@ -1403,16 +1409,16 @@ public void testUnknownProducerHandlingWhenRetentionLimitReached() throws Except
         assertEquals(OptionalLong.of(1000L), transactionManager.lastAckedOffset(tp0));
 
         // Send second ProduceRequest, a single batch with 2 records.
-        accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT);
-        Future request2 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds());
+        accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false);
+        Future request2 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
+        sender.runOnce();
         assertEquals(3, transactionManager.sequenceNumber(tp0).longValue());
         assertEquals(OptionalInt.of(0), transactionManager.lastAckedSequence(tp0));
 
         assertFalse(request2.isDone());
 
         sendIdempotentProducerResponse(1, tp0, Errors.UNKNOWN_PRODUCER_ID, -1L, 1010L);
-        sender.run(time.milliseconds()); // receive response 0, should be retried since the logStartOffset > lastAckedOffset.
+        sender.runOnce(); // receive response 0, should be retried since the logStartOffset > lastAckedOffset.
 
         // We should have reset the sequence number state of the partition because the state was lost on the broker.
         assertEquals(OptionalInt.empty(), transactionManager.lastAckedSequence(tp0));
@@ -1420,11 +1426,11 @@ public void testUnknownProducerHandlingWhenRetentionLimitReached() throws Except
         assertFalse(request2.isDone());
         assertFalse(client.hasInFlightRequests());
 
-        sender.run(time.milliseconds()); // should retry request 1
+        sender.runOnce(); // should retry request 1
 
         // resend the request. Note that the expected sequence is 0, since we have lost producer state on the broker.
         sendIdempotentProducerResponse(0, tp0, Errors.NONE, 1011L, 1010L);
-        sender.run(time.milliseconds()); // receive response 1
+        sender.runOnce(); // receive response 1
         assertEquals(OptionalInt.of(1), transactionManager.lastAckedSequence(tp0));
         assertEquals(2, transactionManager.sequenceNumber(tp0).longValue());
         assertFalse(client.hasInFlightRequests());
@@ -1444,8 +1450,8 @@ public void testUnknownProducerErrorShouldBeRetriedWhenLogStartOffsetIsUnknown()
         assertEquals(0, transactionManager.sequenceNumber(tp0).longValue());
 
         // Send first ProduceRequest
-        Future request1 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds());
+        Future request1 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
+        sender.runOnce();
 
         assertEquals(1, client.inFlightRequestCount());
         assertEquals(1, transactionManager.sequenceNumber(tp0).longValue());
@@ -1453,7 +1459,7 @@ public void testUnknownProducerErrorShouldBeRetriedWhenLogStartOffsetIsUnknown()
 
         sendIdempotentProducerResponse(0, tp0, Errors.NONE, 1000L, 10L);
 
-        sender.run(time.milliseconds());  // receive the response.
+        sender.runOnce();  // receive the response.
 
         assertTrue(request1.isDone());
         assertEquals(1000L, request1.get().offset());
@@ -1461,15 +1467,15 @@ public void testUnknownProducerErrorShouldBeRetriedWhenLogStartOffsetIsUnknown()
         assertEquals(OptionalLong.of(1000L), transactionManager.lastAckedOffset(tp0));
 
         // Send second ProduceRequest
-        Future request2 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds());
+        Future request2 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
+        sender.runOnce();
         assertEquals(2, transactionManager.sequenceNumber(tp0).longValue());
         assertEquals(OptionalInt.of(0), transactionManager.lastAckedSequence(tp0));
 
         assertFalse(request2.isDone());
 
         sendIdempotentProducerResponse(1, tp0, Errors.UNKNOWN_PRODUCER_ID, -1L, -1L);
-        sender.run(time.milliseconds()); // receive response 0, should be retried without resetting the sequence numbers since the log start offset is unknown.
+        sender.runOnce(); // receive response 0, should be retried without resetting the sequence numbers since the log start offset is unknown.
 
         // We should have reset the sequence number state of the partition because the state was lost on the broker.
         assertEquals(OptionalInt.of(0), transactionManager.lastAckedSequence(tp0));
@@ -1477,12 +1483,12 @@ public void testUnknownProducerErrorShouldBeRetriedWhenLogStartOffsetIsUnknown()
         assertFalse(request2.isDone());
         assertFalse(client.hasInFlightRequests());
 
-        sender.run(time.milliseconds()); // should retry request 1
+        sender.runOnce(); // should retry request 1
 
         // resend the request. Note that the expected sequence is 1, since we never got the logStartOffset in the previous
         // response and hence we didn't reset the sequence numbers.
         sendIdempotentProducerResponse(1, tp0, Errors.NONE, 1011L, 1010L);
-        sender.run(time.milliseconds()); // receive response 1
+        sender.runOnce(); // receive response 1
         assertEquals(OptionalInt.of(1), transactionManager.lastAckedSequence(tp0));
         assertEquals(2, transactionManager.sequenceNumber(tp0).longValue());
         assertFalse(client.hasInFlightRequests());
@@ -1502,8 +1508,8 @@ public void testUnknownProducerErrorShouldBeRetriedForFutureBatchesWhenFirstFail
         assertEquals(0, transactionManager.sequenceNumber(tp0).longValue());
 
         // Send first ProduceRequest
-        Future request1 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds());
+        Future request1 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
+        sender.runOnce();
 
         assertEquals(1, client.inFlightRequestCount());
         assertEquals(1, transactionManager.sequenceNumber(tp0).longValue());
@@ -1511,7 +1517,7 @@ public void testUnknownProducerErrorShouldBeRetriedForFutureBatchesWhenFirstFail
 
         sendIdempotentProducerResponse(0, tp0, Errors.NONE, 1000L, 10L);
 
-        sender.run(time.milliseconds());  // receive the response.
+        sender.runOnce();  // receive the response.
 
         assertTrue(request1.isDone());
         assertEquals(1000L, request1.get().offset());
@@ -1519,15 +1525,15 @@ public void testUnknownProducerErrorShouldBeRetriedForFutureBatchesWhenFirstFail
         assertEquals(OptionalLong.of(1000L), transactionManager.lastAckedOffset(tp0));
 
         // Send second ProduceRequest
-        Future request2 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds());
+        Future request2 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
+        sender.runOnce();
         assertEquals(2, transactionManager.sequenceNumber(tp0).longValue());
         assertEquals(OptionalInt.of(0), transactionManager.lastAckedSequence(tp0));
 
         // Send the third ProduceRequest, in parallel with the second. It should be retried even though the
         // lastAckedOffset > logStartOffset when its UnknownProducerResponse comes back.
-        Future request3 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds());
+        Future request3 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
+        sender.runOnce();
         assertEquals(3, transactionManager.sequenceNumber(tp0).longValue());
         assertEquals(OptionalInt.of(0), transactionManager.lastAckedSequence(tp0));
 
@@ -1537,7 +1543,7 @@ public void testUnknownProducerErrorShouldBeRetriedForFutureBatchesWhenFirstFail
 
 
         sendIdempotentProducerResponse(1, tp0, Errors.UNKNOWN_PRODUCER_ID, -1L, 1010L);
-        sender.run(time.milliseconds()); // receive response 2, should reset the sequence numbers and be retried.
+        sender.runOnce(); // receive response 2, should reset the sequence numbers and be retried.
 
         // We should have reset the sequence number state of the partition because the state was lost on the broker.
         assertEquals(OptionalInt.empty(), transactionManager.lastAckedSequence(tp0));
@@ -1546,20 +1552,20 @@ public void testUnknownProducerErrorShouldBeRetriedForFutureBatchesWhenFirstFail
         assertFalse(request3.isDone());
         assertEquals(1, client.inFlightRequestCount());
 
-        sender.run(time.milliseconds()); // resend request 2.
+        sender.runOnce(); // resend request 2.
 
         assertEquals(2, client.inFlightRequestCount());
 
         // receive the original response 3. note the expected sequence is still the originally assigned sequence.
         sendIdempotentProducerResponse(2, tp0, Errors.UNKNOWN_PRODUCER_ID, -1, 1010L);
-        sender.run(time.milliseconds()); // receive response 3
+        sender.runOnce(); // receive response 3
 
         assertEquals(1, client.inFlightRequestCount());
         assertEquals(OptionalInt.empty(), transactionManager.lastAckedSequence(tp0));
         assertEquals(2, transactionManager.sequenceNumber(tp0).longValue());
 
         sendIdempotentProducerResponse(0, tp0, Errors.NONE, 1011L, 1010L);
-        sender.run(time.milliseconds());  // receive response 2, don't send request 3 since we can have at most 1 in flight when retrying
+        sender.runOnce();  // receive response 2, don't send request 3 since we can have at most 1 in flight when retrying
 
         assertTrue(request2.isDone());
         assertFalse(request3.isDone());
@@ -1568,11 +1574,11 @@ public void testUnknownProducerErrorShouldBeRetriedForFutureBatchesWhenFirstFail
         assertEquals(1011L, request2.get().offset());
         assertEquals(OptionalLong.of(1011L), transactionManager.lastAckedOffset(tp0));
 
-        sender.run(time.milliseconds());  // resend request 3.
+        sender.runOnce();  // resend request 3.
         assertEquals(1, client.inFlightRequestCount());
 
         sendIdempotentProducerResponse(1, tp0, Errors.NONE, 1012L, 1010L);
-        sender.run(time.milliseconds());  // receive response 3.
+        sender.runOnce();  // receive response 3.
 
         assertFalse(client.hasInFlightRequests());
         assertTrue(request3.isDone());
@@ -1591,8 +1597,8 @@ public void testShouldRaiseOutOfOrderSequenceExceptionToUserIfLogWasNotTruncated
         assertEquals(0, transactionManager.sequenceNumber(tp0).longValue());
 
         // Send first ProduceRequest
-        Future request1 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds());
+        Future request1 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
+        sender.runOnce();
 
         assertEquals(1, client.inFlightRequestCount());
         assertEquals(1, transactionManager.sequenceNumber(tp0).longValue());
@@ -1600,7 +1606,7 @@ public void testShouldRaiseOutOfOrderSequenceExceptionToUserIfLogWasNotTruncated
 
         sendIdempotentProducerResponse(0, tp0, Errors.NONE, 1000L, 10L);
 
-        sender.run(time.milliseconds());  // receive the response.
+        sender.runOnce();  // receive the response.
 
         assertTrue(request1.isDone());
         assertEquals(1000L, request1.get().offset());
@@ -1608,15 +1614,15 @@ public void testShouldRaiseOutOfOrderSequenceExceptionToUserIfLogWasNotTruncated
         assertEquals(OptionalLong.of(1000L), transactionManager.lastAckedOffset(tp0));
 
         // Send second ProduceRequest,
-        Future request2 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds());
+        Future request2 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
+        sender.runOnce();
         assertEquals(2, transactionManager.sequenceNumber(tp0).longValue());
         assertEquals(OptionalInt.of(0), transactionManager.lastAckedSequence(tp0));
 
         assertFalse(request2.isDone());
 
         sendIdempotentProducerResponse(1, tp0, Errors.UNKNOWN_PRODUCER_ID, -1L, 10L);
-        sender.run(time.milliseconds()); // receive response 0, should cause a producerId reset since the logStartOffset < lastAckedOffset
+        sender.runOnce(); // receive response 0, should cause a producerId reset since the logStartOffset < lastAckedOffset
         assertFutureFailure(request2, OutOfOrderSequenceException.class);
 
     }
@@ -1652,7 +1658,7 @@ public void testClusterAuthorizationExceptionInProduceRequest() throws Exception
 
         // cluster authorization is a fatal error for the producer
         Future future = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(),
-                null, null, MAX_BLOCK_TIMEOUT).future;
+                null, null, MAX_BLOCK_TIMEOUT, false).future;
         client.prepareResponse(new MockClient.RequestMatcher() {
             @Override
             public boolean matches(AbstractRequest body) {
@@ -1660,7 +1666,7 @@ public boolean matches(AbstractRequest body) {
             }
         }, produceResponse(tp0, -1, Errors.CLUSTER_AUTHORIZATION_FAILED, 0));
 
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertFutureFailure(future, ClusterAuthorizationException.class);
 
         // cluster authorization errors are fatal, so we should continue seeing it on future sends
@@ -1679,12 +1685,12 @@ public void testCancelInFlightRequestAfterFatalError() throws Exception {
 
         // cluster authorization is a fatal error for the producer
         Future future1 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(),
-                null, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds());
+                null, null, MAX_BLOCK_TIMEOUT, false).future;
+        sender.runOnce();
 
         Future future2 = accumulator.append(tp1, time.milliseconds(), "key".getBytes(), "value".getBytes(),
-                null, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds());
+                null, null, MAX_BLOCK_TIMEOUT, false).future;
+        sender.runOnce();
 
         client.respond(new MockClient.RequestMatcher() {
             @Override
@@ -1693,11 +1699,11 @@ public boolean matches(AbstractRequest body) {
             }
         }, produceResponse(tp0, -1, Errors.CLUSTER_AUTHORIZATION_FAILED, 0));
 
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertTrue(transactionManager.hasFatalError());
         assertFutureFailure(future1, ClusterAuthorizationException.class);
 
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertFutureFailure(future2, ClusterAuthorizationException.class);
 
         // Should be fine if the second response eventually returns
@@ -1707,7 +1713,7 @@ public boolean matches(AbstractRequest body) {
                 return body instanceof ProduceRequest && ((ProduceRequest) body).hasIdempotentRecords();
             }
         }, produceResponse(tp1, 0, Errors.NONE, 0));
-        sender.run(time.milliseconds());
+        sender.runOnce();
     }
 
     @Test
@@ -1720,7 +1726,7 @@ public void testUnsupportedForMessageFormatInProduceRequest() throws Exception {
         assertTrue(transactionManager.hasProducerId());
 
         Future future = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(),
-                null, null, MAX_BLOCK_TIMEOUT).future;
+                null, null, MAX_BLOCK_TIMEOUT, false).future;
         client.prepareResponse(new MockClient.RequestMatcher() {
             @Override
             public boolean matches(AbstractRequest body) {
@@ -1728,7 +1734,7 @@ public boolean matches(AbstractRequest body) {
             }
         }, produceResponse(tp0, -1, Errors.UNSUPPORTED_FOR_MESSAGE_FORMAT, 0));
 
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertFutureFailure(future, UnsupportedForMessageFormatException.class);
 
         // unsupported for message format is not a fatal error
@@ -1745,7 +1751,7 @@ public void testUnsupportedVersionInProduceRequest() throws Exception {
         assertTrue(transactionManager.hasProducerId());
 
         Future future = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(),
-                null, null, MAX_BLOCK_TIMEOUT).future;
+                null, null, MAX_BLOCK_TIMEOUT, false).future;
         client.prepareUnsupportedVersionResponse(new MockClient.RequestMatcher() {
             @Override
             public boolean matches(AbstractRequest body) {
@@ -1753,7 +1759,7 @@ public boolean matches(AbstractRequest body) {
             }
         });
 
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertFutureFailure(future, UnsupportedVersionException.class);
 
         // unsupported version errors are fatal, so we should continue seeing it on future sends
@@ -1773,9 +1779,9 @@ public void testSequenceNumberIncrement() throws InterruptedException {
         SenderMetricsRegistry senderMetrics = new SenderMetricsRegistry(m);
 
         Sender sender = new Sender(logContext, client, metadata, this.accumulator, true, MAX_REQUEST_SIZE, ACKS_ALL, maxRetries,
-                senderMetrics, time, REQUEST_TIMEOUT, 50, transactionManager, apiVersions);
+                senderMetrics, time, REQUEST_TIMEOUT, RETRY_BACKOFF_MS, transactionManager, apiVersions);
 
-        Future responseFuture = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
+        Future responseFuture = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
         client.prepareResponse(new MockClient.RequestMatcher() {
             @Override
             public boolean matches(AbstractRequest body) {
@@ -1795,17 +1801,16 @@ public boolean matches(AbstractRequest body) {
             }
         }, produceResponse(tp0, 0, Errors.NONE, 0));
 
-        sender.run(time.milliseconds());  // connect.
-        sender.run(time.milliseconds());  // send.
+        sender.runOnce();  // connect.
+        sender.runOnce();  // send.
 
-        sender.run(time.milliseconds());  // receive response
+        sender.runOnce();  // receive response
         assertTrue(responseFuture.isDone());
         assertEquals(OptionalInt.of(0), transactionManager.lastAckedSequence(tp0));
         assertEquals(1L, (long) transactionManager.sequenceNumber(tp0));
     }
 
     @Test
-    @SuppressWarnings("deprecation")
     public void testAbortRetryWhenProducerIdChanges() throws InterruptedException {
         final long producerId = 343434L;
         TransactionManager transactionManager = new TransactionManager();
@@ -1816,24 +1821,24 @@ public void testAbortRetryWhenProducerIdChanges() throws InterruptedException {
         Metrics m = new Metrics();
         SenderMetricsRegistry senderMetrics = new SenderMetricsRegistry(m);
         Sender sender = new Sender(logContext, client, metadata, this.accumulator, true, MAX_REQUEST_SIZE, ACKS_ALL, maxRetries,
-                senderMetrics, time, REQUEST_TIMEOUT, 50, transactionManager, apiVersions);
+                senderMetrics, time, REQUEST_TIMEOUT, RETRY_BACKOFF_MS, transactionManager, apiVersions);
 
-        Future responseFuture = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds());  // connect.
-        sender.run(time.milliseconds());  // send.
+        Future responseFuture = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
+        sender.runOnce();  // connect.
+        sender.runOnce();  // send.
         String id = client.requests().peek().destination();
         Node node = new Node(Integer.valueOf(id), "localhost", 0);
         assertEquals(1, client.inFlightRequestCount());
-        assertTrue("Client ready status should be true", client.isReady(node, 0L));
+        assertTrue("Client ready status should be true", client.isReady(node, time.milliseconds()));
         client.disconnect(id);
         assertEquals(0, client.inFlightRequestCount());
-        assertFalse("Client ready status should be false", client.isReady(node, 0L));
+        assertFalse("Client ready status should be false", client.isReady(node, time.milliseconds()));
 
         transactionManager.resetProducerId();
         transactionManager.setProducerIdAndEpoch(new ProducerIdAndEpoch(producerId + 1, (short) 0));
-        sender.run(time.milliseconds()); // receive error
-        sender.run(time.milliseconds()); // reconnect
-        sender.run(time.milliseconds()); // nothing to do, since the pid has changed. We should check the metrics for errors.
+        sender.runOnce(); // receive error
+        sender.runOnce(); // reconnect
+        sender.runOnce(); // nothing to do, since the pid has changed. We should check the metrics for errors.
         assertEquals("Expected requests to be aborted after pid change", 0, client.inFlightRequestCount());
 
         KafkaMetric recordErrors = m.metrics().get(senderMetrics.recordErrorRate);
@@ -1855,18 +1860,18 @@ public void testResetWhenOutOfOrderSequenceReceived() throws InterruptedExceptio
         SenderMetricsRegistry senderMetrics = new SenderMetricsRegistry(m);
 
         Sender sender = new Sender(logContext, client, metadata, this.accumulator, true, MAX_REQUEST_SIZE, ACKS_ALL, maxRetries,
-                senderMetrics, time, REQUEST_TIMEOUT, 50, transactionManager, apiVersions);
+                senderMetrics, time, REQUEST_TIMEOUT, RETRY_BACKOFF_MS, transactionManager, apiVersions);
 
-        Future responseFuture = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds());  // connect.
-        sender.run(time.milliseconds());  // send.
+        Future responseFuture = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
+        sender.runOnce();  // connect.
+        sender.runOnce();  // send.
 
         assertEquals(1, client.inFlightRequestCount());
         assertEquals(1, sender.inFlightBatches(tp0).size());
 
         client.respond(produceResponse(tp0, 0, Errors.OUT_OF_ORDER_SEQUENCE_NUMBER, 0));
 
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertTrue(responseFuture.isDone());
         assertEquals(0, sender.inFlightBatches(tp0).size());
         assertFalse("Expected transaction state to be reset upon receiving an OutOfOrderSequenceException", transactionManager.hasProducerId());
@@ -1891,27 +1896,27 @@ public void testTransactionalSplitBatchAndSend() throws Exception {
         doInitTransactions(txnManager, producerIdAndEpoch);
 
         txnManager.beginTransaction();
+        txnManager.failIfNotReadyForSend();
         txnManager.maybeAddPartitionToTransaction(tp);
         client.prepareResponse(new AddPartitionsToTxnResponse(0, Collections.singletonMap(tp, Errors.NONE)));
-        sender.run(time.milliseconds());
+        sender.runOnce();
 
         testSplitBatchAndSend(txnManager, producerIdAndEpoch, tp);
     }
 
-    @SuppressWarnings("deprecation")
     private void testSplitBatchAndSend(TransactionManager txnManager,
                                        ProducerIdAndEpoch producerIdAndEpoch,
                                        TopicPartition tp) throws Exception {
         int maxRetries = 1;
         String topic = tp.topic();
-        long deliveryTimeoutMs = 3000L;
+        int deliveryTimeoutMs = 3000;
         long totalSize = 1024 * 1024;
         String metricGrpName = "producer-metrics";
         // Set a good compression ratio.
         CompressionRatioEstimator.setEstimation(topic, CompressionType.GZIP, 0.2f);
         try (Metrics m = new Metrics()) {
             accumulator = new RecordAccumulator(logContext, batchSize, CompressionType.GZIP,
-                0L, 0L, deliveryTimeoutMs, m, metricGrpName, time, new ApiVersions(), txnManager,
+                0, 0L, deliveryTimeoutMs, m, metricGrpName, time, new ApiVersions(), txnManager,
                 new BufferPool(totalSize, batchSize, metrics, time, "producer-internal-metrics"));
             SenderMetricsRegistry senderMetrics = new SenderMetricsRegistry(m);
             Sender sender = new Sender(logContext, client, metadata, this.accumulator, true, MAX_REQUEST_SIZE, ACKS_ALL, maxRetries,
@@ -1921,28 +1926,28 @@ private void testSplitBatchAndSend(TransactionManager txnManager,
             client.prepareMetadataUpdate(metadataUpdate1);
             // Send the first message.
             Future f1 =
-                    accumulator.append(tp, 0L, "key1".getBytes(), new byte[batchSize / 2], null, null, MAX_BLOCK_TIMEOUT).future;
+                    accumulator.append(tp, 0L, "key1".getBytes(), new byte[batchSize / 2], null, null, MAX_BLOCK_TIMEOUT, false).future;
             Future f2 =
-                    accumulator.append(tp, 0L, "key2".getBytes(), new byte[batchSize / 2], null, null, MAX_BLOCK_TIMEOUT).future;
-            sender.run(time.milliseconds()); // connect
-            sender.run(time.milliseconds()); // send produce request
+                    accumulator.append(tp, 0L, "key2".getBytes(), new byte[batchSize / 2], null, null, MAX_BLOCK_TIMEOUT, false).future;
+            sender.runOnce(); // connect
+            sender.runOnce(); // send produce request
 
             assertEquals("The next sequence should be 2", 2, txnManager.sequenceNumber(tp).longValue());
             String id = client.requests().peek().destination();
             assertEquals(ApiKeys.PRODUCE, client.requests().peek().requestBuilder().apiKey());
             Node node = new Node(Integer.valueOf(id), "localhost", 0);
             assertEquals(1, client.inFlightRequestCount());
-            assertTrue("Client ready status should be true", client.isReady(node, 0L));
+            assertTrue("Client ready status should be true", client.isReady(node, time.milliseconds()));
 
             Map responseMap = new HashMap<>();
             responseMap.put(tp, new ProduceResponse.PartitionResponse(Errors.MESSAGE_TOO_LARGE));
             client.respond(new ProduceResponse(responseMap));
-            sender.run(time.milliseconds()); // split and reenqueue
+            sender.runOnce(); // split and reenqueue
             assertEquals("The next sequence should be 2", 2, txnManager.sequenceNumber(tp).longValue());
             // The compression ratio should have been improved once.
             assertEquals(CompressionType.GZIP.rate - CompressionRatioEstimator.COMPRESSION_RATIO_IMPROVING_STEP,
                     CompressionRatioEstimator.estimation(topic, CompressionType.GZIP), 0.01);
-            sender.run(time.milliseconds()); // send the first produce request
+            sender.runOnce(); // send the first produce request
             assertEquals("The next sequence number should be 2", 2, txnManager.sequenceNumber(tp).longValue());
             assertFalse("The future shouldn't have been done.", f1.isDone());
             assertFalse("The future shouldn't have been done.", f2.isDone());
@@ -1950,30 +1955,30 @@ private void testSplitBatchAndSend(TransactionManager txnManager,
             assertEquals(ApiKeys.PRODUCE, client.requests().peek().requestBuilder().apiKey());
             node = new Node(Integer.valueOf(id), "localhost", 0);
             assertEquals(1, client.inFlightRequestCount());
-            assertTrue("Client ready status should be true", client.isReady(node, 0L));
+            assertTrue("Client ready status should be true", client.isReady(node, time.milliseconds()));
 
             responseMap.put(tp, new ProduceResponse.PartitionResponse(Errors.NONE, 0L, 0L, 0L));
             client.respond(produceRequestMatcher(tp, producerIdAndEpoch, 0, txnManager.isTransactional()),
                     new ProduceResponse(responseMap));
 
-            sender.run(time.milliseconds()); // receive
+            sender.runOnce(); // receive
             assertTrue("The future should have been done.", f1.isDone());
             assertEquals("The next sequence number should still be 2", 2, txnManager.sequenceNumber(tp).longValue());
             assertEquals("The last ack'd sequence number should be 0", OptionalInt.of(0), txnManager.lastAckedSequence(tp));
             assertFalse("The future shouldn't have been done.", f2.isDone());
             assertEquals("Offset of the first message should be 0", 0L, f1.get().offset());
-            sender.run(time.milliseconds()); // send the seconcd produce request
+            sender.runOnce(); // send the seconcd produce request
             id = client.requests().peek().destination();
             assertEquals(ApiKeys.PRODUCE, client.requests().peek().requestBuilder().apiKey());
             node = new Node(Integer.valueOf(id), "localhost", 0);
             assertEquals(1, client.inFlightRequestCount());
-            assertTrue("Client ready status should be true", client.isReady(node, 0L));
+            assertTrue("Client ready status should be true", client.isReady(node, time.milliseconds()));
 
             responseMap.put(tp, new ProduceResponse.PartitionResponse(Errors.NONE, 1L, 0L, 0L));
             client.respond(produceRequestMatcher(tp, producerIdAndEpoch, 1, txnManager.isTransactional()),
                     new ProduceResponse(responseMap));
 
-            sender.run(time.milliseconds()); // receive
+            sender.runOnce(); // receive
             assertTrue("The future should have been done.", f2.isDone());
             assertEquals("The next sequence number should be 2", 2, txnManager.sequenceNumber(tp).longValue());
             assertEquals("The last ack'd sequence number should be 1", OptionalInt.of(1), txnManager.lastAckedSequence(tp));
@@ -1993,20 +1998,20 @@ public void testNoDoubleDeallocation() throws Exception {
 
         // Send first ProduceRequest
         Future request1 =
-            accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds());  // send request
+            accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
+        sender.runOnce();  // send request
         assertEquals(1, client.inFlightRequestCount());
         assertEquals(1, sender.inFlightBatches(tp0).size());
 
         time.sleep(deliverTimeoutMs);
         assertFalse(pool.allMatch());
 
-        sender.run(time.milliseconds());  // expire the batch
+        sender.runOnce();  // expire the batch
         assertTrue(request1.isDone());
         assertTrue("The batch should have been de-allocated", pool.allMatch());
         assertTrue(pool.allMatch());
 
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertTrue("The batch should have been de-allocated", pool.allMatch());
         assertEquals(0, client.inFlightRequestCount());
         assertEquals(0, sender.inFlightBatches(tp0).size());
@@ -2018,8 +2023,8 @@ public void testInflightBatchesExpireOnDeliveryTimeout() throws InterruptedExcep
         setupWithTransactionState(null, true, null);
 
         // Send first ProduceRequest
-        Future request = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds());  // send request
+        Future request = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
+        sender.runOnce();  // send request
         assertEquals(1, client.inFlightRequestCount());
         assertEquals("Expect one in-flight batch in accumulator", 1, sender.inFlightBatches(tp0).size());
 
@@ -2028,7 +2033,7 @@ public void testInflightBatchesExpireOnDeliveryTimeout() throws InterruptedExcep
         client.respond(new ProduceResponse(responseMap));
 
         time.sleep(deliveryTimeoutMs);
-        sender.run(time.milliseconds());  // receive first response
+        sender.runOnce();  // receive first response
         assertEquals("Expect zero in-flight batch in accumulator", 0, sender.inFlightBatches(tp0).size());
         try {
             request.get();
@@ -2044,27 +2049,27 @@ public void testWhenFirstBatchExpireNoSendSecondBatchIfGuaranteeOrder() throws I
         setupWithTransactionState(null, true, null);
 
         // Send first ProduceRequest
-        accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT);
-        sender.run(time.milliseconds());  // send request
+        accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false);
+        sender.runOnce();  // send request
         assertEquals(1, client.inFlightRequestCount());
         assertEquals(1, sender.inFlightBatches(tp0).size());
 
         time.sleep(deliveryTimeoutMs / 2);
 
         // Send second ProduceRequest
-        accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT);
-        sender.run(time.milliseconds());  // must not send request because the partition is muted
+        accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false);
+        sender.runOnce();  // must not send request because the partition is muted
         assertEquals(1, client.inFlightRequestCount());
         assertEquals(1, sender.inFlightBatches(tp0).size());
 
         time.sleep(deliveryTimeoutMs / 2); // expire the first batch only
 
         client.respond(produceResponse(tp0, 0L, Errors.NONE, 0, 0L));
-        sender.run(time.milliseconds());  // receive response (offset=0)
+        sender.runOnce();  // receive response (offset=0)
         assertEquals(0, client.inFlightRequestCount());
         assertEquals(0, sender.inFlightBatches(tp0).size());
 
-        sender.run(time.milliseconds());  // Drain the second request only this time
+        sender.runOnce();  // Drain the second request only this time
         assertEquals(1, client.inFlightRequestCount());
         assertEquals(1, sender.inFlightBatches(tp0).size());
     }
@@ -2077,8 +2082,8 @@ public void testExpiredBatchDoesNotRetry() throws Exception {
         // Send first ProduceRequest
         Future request1 =
             accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null,
-                MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds());  // send request
+                MAX_BLOCK_TIMEOUT, false).future;
+        sender.runOnce();  // send request
         assertEquals(1, client.inFlightRequestCount());
         time.sleep(deliverTimeoutMs);
 
@@ -2086,16 +2091,16 @@ public void testExpiredBatchDoesNotRetry() throws Exception {
         responseMap.put(tp0, new ProduceResponse.PartitionResponse(Errors.NONE, 0L, 0L, 0L));
         client.respond(produceResponse(tp0, -1, Errors.NOT_LEADER_FOR_PARTITION, -1)); // return a retriable error
 
-        sender.run(time.milliseconds());  // expire the batch
+        sender.runOnce();  // expire the batch
         assertTrue(request1.isDone());
         assertEquals(0, client.inFlightRequestCount());
         assertEquals(0, sender.inFlightBatches(tp0).size());
 
-        sender.run(time.milliseconds()); // receive first response and do not reenqueue.
+        sender.runOnce(); // receive first response and do not reenqueue.
         assertEquals(0, client.inFlightRequestCount());
         assertEquals(0, sender.inFlightBatches(tp0).size());
 
-        sender.run(time.milliseconds()); // run again and must not send anything.
+        sender.runOnce(); // run again and must not send anything.
         assertEquals(0, client.inFlightRequestCount());
         assertEquals(0, sender.inFlightBatches(tp0).size());
     }
@@ -2103,29 +2108,30 @@ public void testExpiredBatchDoesNotRetry() throws Exception {
     @Test
     public void testExpiredBatchDoesNotSplitOnMessageTooLargeError() throws Exception {
         long deliverTimeoutMs = 1500L;
-        // create a producer batch with more than one record so it is eligible to split
+        // create a producer batch with more than one record so it is eligible for splitting
         Future request1 =
             accumulator.append(tp0, time.milliseconds(), "key1".getBytes(), "value1".getBytes(), null, null,
-                MAX_BLOCK_TIMEOUT).future;
+                MAX_BLOCK_TIMEOUT, false).future;
         Future request2 =
             accumulator.append(tp0, time.milliseconds(), "key2".getBytes(), "value2".getBytes(), null, null,
-                MAX_BLOCK_TIMEOUT).future;
+                MAX_BLOCK_TIMEOUT, false).future;
 
-        sender.run(time.milliseconds());  // send request
+        // send request
+        sender.runOnce();
         assertEquals(1, client.inFlightRequestCount());
         // return a MESSAGE_TOO_LARGE error
         client.respond(produceResponse(tp0, -1, Errors.MESSAGE_TOO_LARGE, -1));
 
         time.sleep(deliverTimeoutMs);
         // expire the batch and process the response
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertTrue(request1.isDone());
         assertTrue(request2.isDone());
         assertEquals(0, client.inFlightRequestCount());
         assertEquals(0, sender.inFlightBatches(tp0).size());
 
         // run again and must not split big batch and resend anything.
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertEquals(0, client.inFlightRequestCount());
         assertEquals(0, sender.inFlightBatches(tp0).size());
     }
@@ -2137,12 +2143,12 @@ public void testResetNextBatchExpiry() throws Exception {
         setupWithTransactionState(null);
 
         accumulator.append(tp0, 0L, "key".getBytes(), "value".getBytes(), null, null,
-                MAX_BLOCK_TIMEOUT);
+                MAX_BLOCK_TIMEOUT, false);
 
-        sender.run(time.milliseconds());
-        sender.run(time.milliseconds());
+        sender.runOnce();
+        sender.runOnce();
         time.setCurrentTimeMs(time.milliseconds() + accumulator.getDeliveryTimeoutMs() + 1);
-        sender.run(time.milliseconds());
+        sender.runOnce();
 
         InOrder inOrder = inOrder(client);
         inOrder.verify(client, atLeastOnce()).ready(any(), anyLong());
@@ -2161,11 +2167,11 @@ public void testExpiredBatchesInMultiplePartitions() throws Exception {
         setupWithTransactionState(null, true, null);
 
         // Send multiple ProduceRequest across multiple partitions.
-        Future request1 = accumulator.append(tp0, time.milliseconds(), "k1".getBytes(), "v1".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
-        Future request2 = accumulator.append(tp1, time.milliseconds(), "k2".getBytes(), "v2".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
+        Future request1 = accumulator.append(tp0, time.milliseconds(), "k1".getBytes(), "v1".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
+        Future request2 = accumulator.append(tp1, time.milliseconds(), "k2".getBytes(), "v2".getBytes(), null, null, MAX_BLOCK_TIMEOUT, false).future;
 
         // Send request.
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertEquals(1, client.inFlightRequestCount());
         assertEquals("Expect one in-flight batch in accumulator", 1, sender.inFlightBatches(tp0).size());
 
@@ -2175,7 +2181,7 @@ public void testExpiredBatchesInMultiplePartitions() throws Exception {
 
         // Successfully expire both batches.
         time.sleep(deliveryTimeoutMs);
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertEquals("Expect zero in-flight batch in accumulator", 0, sender.inFlightBatches(tp0).size());
 
         try {
@@ -2193,6 +2199,179 @@ public void testExpiredBatchesInMultiplePartitions() throws Exception {
         }
     }
 
+    @Test
+    public void testTransactionalRequestsSentOnShutdown() {
+        // create a sender with retries = 1
+        int maxRetries = 1;
+        Metrics m = new Metrics();
+        SenderMetricsRegistry senderMetrics = new SenderMetricsRegistry(m);
+        try {
+            TransactionManager txnManager = new TransactionManager(logContext, "testTransactionalRequestsSentOnShutdown", 6000, 100);
+            Sender sender = new Sender(logContext, client, metadata, this.accumulator, false, MAX_REQUEST_SIZE, ACKS_ALL,
+                    maxRetries, senderMetrics, time, REQUEST_TIMEOUT, RETRY_BACKOFF_MS, txnManager, apiVersions);
+
+            ProducerIdAndEpoch producerIdAndEpoch = new ProducerIdAndEpoch(123456L, (short) 0);
+            TopicPartition tp = new TopicPartition("testTransactionalRequestsSentOnShutdown", 1);
+
+            setupWithTransactionState(txnManager);
+            doInitTransactions(txnManager, producerIdAndEpoch);
+
+            txnManager.beginTransaction();
+            txnManager.failIfNotReadyForSend();
+            txnManager.maybeAddPartitionToTransaction(tp);
+            client.prepareResponse(new AddPartitionsToTxnResponse(0, Collections.singletonMap(tp, Errors.NONE)));
+            sender.runOnce();
+            sender.initiateClose();
+            txnManager.beginCommit();
+            AssertEndTxnRequestMatcher endTxnMatcher = new AssertEndTxnRequestMatcher(TransactionResult.COMMIT);
+            client.prepareResponse(endTxnMatcher, new EndTxnResponse(0, Errors.NONE));
+            sender.run();
+            assertTrue("Response didn't match in test", endTxnMatcher.matched);
+        } finally {
+            m.close();
+        }
+    }
+
+    @Test
+    public void testIncompleteTransactionAbortOnShutdown() {
+        // create a sender with retries = 1
+        int maxRetries = 1;
+        Metrics m = new Metrics();
+        SenderMetricsRegistry senderMetrics = new SenderMetricsRegistry(m);
+        try {
+            TransactionManager txnManager = new TransactionManager(logContext, "testIncompleteTransactionAbortOnShutdown", 6000, 100);
+            Sender sender = new Sender(logContext, client, metadata, this.accumulator, false, MAX_REQUEST_SIZE, ACKS_ALL,
+                    maxRetries, senderMetrics, time, REQUEST_TIMEOUT, RETRY_BACKOFF_MS, txnManager, apiVersions);
+
+            ProducerIdAndEpoch producerIdAndEpoch = new ProducerIdAndEpoch(123456L, (short) 0);
+            TopicPartition tp = new TopicPartition("testIncompleteTransactionAbortOnShutdown", 1);
+
+            setupWithTransactionState(txnManager);
+            doInitTransactions(txnManager, producerIdAndEpoch);
+
+            txnManager.beginTransaction();
+            txnManager.failIfNotReadyForSend();
+            txnManager.maybeAddPartitionToTransaction(tp);
+            client.prepareResponse(new AddPartitionsToTxnResponse(0, Collections.singletonMap(tp, Errors.NONE)));
+            sender.runOnce();
+            sender.initiateClose();
+            AssertEndTxnRequestMatcher endTxnMatcher = new AssertEndTxnRequestMatcher(TransactionResult.ABORT);
+            client.prepareResponse(endTxnMatcher, new EndTxnResponse(0, Errors.NONE));
+            sender.run();
+            assertTrue("Response didn't match in test", endTxnMatcher.matched);
+        } finally {
+            m.close();
+        }
+    }
+
+    @Test(timeout = 10000L)
+    public void testForceShutdownWithIncompleteTransaction() {
+        // create a sender with retries = 1
+        int maxRetries = 1;
+        Metrics m = new Metrics();
+        SenderMetricsRegistry senderMetrics = new SenderMetricsRegistry(m);
+        try {
+            TransactionManager txnManager = new TransactionManager(logContext, "testForceShutdownWithIncompleteTransaction", 6000, 100);
+            Sender sender = new Sender(logContext, client, metadata, this.accumulator, false, MAX_REQUEST_SIZE, ACKS_ALL,
+                    maxRetries, senderMetrics, time, REQUEST_TIMEOUT, RETRY_BACKOFF_MS, txnManager, apiVersions);
+
+            ProducerIdAndEpoch producerIdAndEpoch = new ProducerIdAndEpoch(123456L, (short) 0);
+            TopicPartition tp = new TopicPartition("testForceShutdownWithIncompleteTransaction", 1);
+
+            setupWithTransactionState(txnManager);
+            doInitTransactions(txnManager, producerIdAndEpoch);
+
+            txnManager.beginTransaction();
+            txnManager.failIfNotReadyForSend();
+            txnManager.maybeAddPartitionToTransaction(tp);
+            client.prepareResponse(new AddPartitionsToTxnResponse(0, Collections.singletonMap(tp, Errors.NONE)));
+            sender.runOnce();
+
+            // Try to commit the transaction but it won't happen as we'll forcefully close the sender
+            TransactionalRequestResult commitResult = txnManager.beginCommit();
+
+            sender.forceClose();
+            sender.run();
+            assertThrows("The test expected to throw a KafkaException for forcefully closing the sender",
+                    KafkaException.class, commitResult::await);
+        } finally {
+            m.close();
+        }
+    }
+
+    @Test
+    public void testDoNotPollWhenNoRequestSent() {
+        client = spy(new MockClient(time, metadata));
+
+        TransactionManager txnManager = new TransactionManager(logContext, "testDoNotPollWhenNoRequestSent", 6000, 100);
+        ProducerIdAndEpoch producerIdAndEpoch = new ProducerIdAndEpoch(123456L, (short) 0);
+        setupWithTransactionState(txnManager);
+        doInitTransactions(txnManager, producerIdAndEpoch);
+
+        // doInitTransactions calls sender.doOnce three times, only two requests are sent, so we should only poll twice
+        verify(client, times(2)).poll(eq(RETRY_BACKOFF_MS), anyLong());
+    }
+
+    @Test
+    public void testTooLargeBatchesAreSafelyRemoved() throws InterruptedException {
+        ProducerIdAndEpoch producerIdAndEpoch = new ProducerIdAndEpoch(123456L, (short) 0);
+        TransactionManager txnManager = new TransactionManager(logContext, "testSplitBatchAndSend", 60000, 100);
+
+        setupWithTransactionState(txnManager, false, null);
+        doInitTransactions(txnManager, producerIdAndEpoch);
+
+        txnManager.beginTransaction();
+        txnManager.maybeAddPartitionToTransaction(tp0);
+        client.prepareResponse(new AddPartitionsToTxnResponse(0, Collections.singletonMap(tp0, Errors.NONE)));
+        sender.runOnce();
+
+        // create a producer batch with more than one record so it is eligible for splitting
+        Future request1 =
+                accumulator.append(tp0, time.milliseconds(), "key1".getBytes(), "value1".getBytes(), null, null,
+                        MAX_BLOCK_TIMEOUT, false).future;
+        Future request2 =
+                accumulator.append(tp0, time.milliseconds(), "key2".getBytes(), "value2".getBytes(), null, null,
+                        MAX_BLOCK_TIMEOUT, false).future;
+
+        // send request
+        sender.runOnce();
+        assertEquals(1, sender.inFlightBatches(tp0).size());
+        // return a MESSAGE_TOO_LARGE error
+        client.respond(produceResponse(tp0, -1, Errors.MESSAGE_TOO_LARGE, -1));
+        sender.runOnce();
+
+        // process retried response
+        sender.runOnce();
+        client.respond(produceResponse(tp0, 0, Errors.NONE, 0));
+        sender.runOnce();
+
+        // In-flight batches should be empty. Sleep past the expiration time of the batch and run once, no error should be thrown
+        assertEquals(0, sender.inFlightBatches(tp0).size());
+        time.sleep(2000);
+        sender.runOnce();
+    }
+
+    class AssertEndTxnRequestMatcher implements MockClient.RequestMatcher {
+
+        private TransactionResult requiredResult;
+        private boolean matched = false;
+
+        AssertEndTxnRequestMatcher(TransactionResult requiredResult) {
+            this.requiredResult = requiredResult;
+        }
+
+        @Override
+        public boolean matches(AbstractRequest body) {
+            if (body instanceof EndTxnRequest) {
+                assertSame(requiredResult, ((EndTxnRequest) body).command());
+                matched = true;
+                return true;
+            } else {
+                return false;
+            }
+        }
+    }
+
     private class MatchingBufferPool extends BufferPool {
         IdentityHashMap allocatedBuffers;
 
@@ -2286,18 +2465,18 @@ private void setupWithTransactionState(TransactionManager transactionManager) {
     }
 
     private void setupWithTransactionState(TransactionManager transactionManager, boolean guaranteeOrder, BufferPool customPool) {
-        long deliveryTimeoutMs = 1500L;
+        int deliveryTimeoutMs = 1500;
         long totalSize = 1024 * 1024;
         String metricGrpName = "producer-metrics";
         MetricConfig metricConfig = new MetricConfig().tags(Collections.singletonMap("client-id", CLIENT_ID));
         this.metrics = new Metrics(metricConfig, time);
         BufferPool pool = (customPool == null) ? new BufferPool(totalSize, batchSize, metrics, time, metricGrpName) : customPool;
 
-        this.accumulator = new RecordAccumulator(logContext, batchSize, CompressionType.NONE, 0L, 0L,
+        this.accumulator = new RecordAccumulator(logContext, batchSize, CompressionType.NONE, 0, 0L,
                 deliveryTimeoutMs, metrics, metricGrpName, time, apiVersions, transactionManager, pool);
         this.senderMetricsRegistry = new SenderMetricsRegistry(this.metrics);
         this.sender = new Sender(logContext, this.client, this.metadata, this.accumulator, guaranteeOrder, MAX_REQUEST_SIZE, ACKS_ALL,
-                Integer.MAX_VALUE, this.senderMetricsRegistry, this.time, REQUEST_TIMEOUT, 50, transactionManager, apiVersions);
+                Integer.MAX_VALUE, this.senderMetricsRegistry, this.time, REQUEST_TIMEOUT, RETRY_BACKOFF_MS, transactionManager, apiVersions);
 
         metadata.add("test");
         this.client.updateMetadata(TestUtils.metadataUpdateWith(1, Collections.singletonMap("test", 2)));
@@ -2305,8 +2484,8 @@ private void setupWithTransactionState(TransactionManager transactionManager, bo
 
     private void assertSendFailure(Class expectedError) throws Exception {
         Future future = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(),
-                null, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds());
+                null, null, MAX_BLOCK_TIMEOUT, false).future;
+        sender.runOnce();
         assertTrue(future.isDone());
         try {
             future.get();
@@ -2321,32 +2500,40 @@ private void prepareAndReceiveInitProducerId(long producerId, Errors error) {
         if (error != Errors.NONE)
             producerEpoch = RecordBatch.NO_PRODUCER_EPOCH;
 
-        client.prepareResponse(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                return body instanceof InitProducerIdRequest && ((InitProducerIdRequest) body).transactionalId() == null;
-            }
-        }, new InitProducerIdResponse(0, error, producerId, producerEpoch));
-        sender.run(time.milliseconds());
+        client.prepareResponse(body -> {
+            return body instanceof InitProducerIdRequest &&
+                    ((InitProducerIdRequest) body).data.transactionalId() == null;
+        }, initProducerIdResponse(producerId, producerEpoch, error));
+        sender.runOnce();
+    }
+
+    private InitProducerIdResponse initProducerIdResponse(long producerId, short producerEpoch, Errors error) {
+        InitProducerIdResponseData responseData = new InitProducerIdResponseData()
+                .setErrorCode(error.code())
+                .setProducerEpoch(producerEpoch)
+                .setProducerId(producerId)
+                .setThrottleTimeMs(0);
+        return new InitProducerIdResponse(responseData);
     }
 
     private void doInitTransactions(TransactionManager transactionManager, ProducerIdAndEpoch producerIdAndEpoch) {
         transactionManager.initializeTransactions();
         prepareFindCoordinatorResponse(Errors.NONE);
-        sender.run(time.milliseconds());
-        sender.run(time.milliseconds());
+        sender.runOnce();
+        sender.runOnce();
 
-        prepareInitPidResponse(Errors.NONE, producerIdAndEpoch.producerId, producerIdAndEpoch.epoch);
-        sender.run(time.milliseconds());
+        prepareInitProducerResponse(Errors.NONE, producerIdAndEpoch.producerId, producerIdAndEpoch.epoch);
+        sender.runOnce();
         assertTrue(transactionManager.hasProducerId());
     }
 
     private void prepareFindCoordinatorResponse(Errors error) {
-        client.prepareResponse(new FindCoordinatorResponse(error, metadata.fetch().nodes().get(0)));
+        Node node = metadata.fetch().nodes().get(0);
+        client.prepareResponse(FindCoordinatorResponse.prepareResponse(error, node));
     }
 
-    private void prepareInitPidResponse(Errors error, long pid, short epoch) {
-        client.prepareResponse(new InitProducerIdResponse(0, error, pid, epoch));
+    private void prepareInitProducerResponse(Errors error, long producerId, short producerEpoch) {
+        client.prepareResponse(initProducerIdResponse(producerId, producerEpoch, error));
     }
 
     private void assertFutureFailure(Future future, Class expectedExceptionType)
diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/StickyPartitionCacheTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/StickyPartitionCacheTest.java
new file mode 100644
index 0000000000000..9f4a481c9a071
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/StickyPartitionCacheTest.java
@@ -0,0 +1,111 @@
+/*
+ * 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.producer.internals;
+
+import org.apache.kafka.common.Cluster;
+import org.apache.kafka.common.Node;
+import org.apache.kafka.common.PartitionInfo;
+import org.junit.Test;
+
+import java.util.Collections;
+import java.util.List;
+
+import static java.util.Arrays.asList;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotEquals;
+
+public class StickyPartitionCacheTest {
+    private final static Node[] NODES = new Node[] {
+        new Node(0, "localhost", 99),
+        new Node(1, "localhost", 100),
+        new Node(2, "localhost", 101),
+        new Node(11, "localhost", 102)
+    };
+    final static String TOPIC_A = "topicA";
+    final static String TOPIC_B = "topicB";
+    final static String TOPIC_C = "topicC";
+
+    @Test
+    public void testStickyPartitionCache() {
+        List allPartitions = asList(new PartitionInfo(TOPIC_A, 0, NODES[0], NODES, NODES),
+            new PartitionInfo(TOPIC_A, 1, NODES[1], NODES, NODES),
+            new PartitionInfo(TOPIC_A, 2, NODES[2], NODES, NODES),
+            new PartitionInfo(TOPIC_B, 0, NODES[0], NODES, NODES)
+        );
+        Cluster testCluster = new Cluster("clusterId", asList(NODES), allPartitions,
+            Collections.emptySet(), Collections.emptySet());
+        StickyPartitionCache stickyPartitionCache = new StickyPartitionCache();
+
+        int partA = stickyPartitionCache.partition(TOPIC_A, testCluster);
+        assertEquals(partA, stickyPartitionCache.partition(TOPIC_A, testCluster));
+
+        int partB = stickyPartitionCache.partition(TOPIC_B, testCluster);
+        assertEquals(partB, stickyPartitionCache.partition(TOPIC_B, testCluster));
+
+        int changedPartA = stickyPartitionCache.nextPartition(TOPIC_A, testCluster, partA);
+        assertEquals(changedPartA, stickyPartitionCache.partition(TOPIC_A, testCluster));
+        assertNotEquals(partA, changedPartA);
+        int changedPartA2 = stickyPartitionCache.partition(TOPIC_A, testCluster);
+        assertEquals(changedPartA2, changedPartA);
+
+        // We do not want to change partitions because the previous partition does not match the current sticky one.
+        int changedPartA3 = stickyPartitionCache.nextPartition(TOPIC_A, testCluster, partA);
+        assertEquals(changedPartA3, changedPartA2);
+
+        // Check that the we can still use the partitioner when there is only one partition
+        int changedPartB = stickyPartitionCache.nextPartition(TOPIC_B, testCluster, partB);
+        assertEquals(changedPartB, stickyPartitionCache.partition(TOPIC_B, testCluster));
+    }
+    
+    @Test
+    public void unavailablePartitionsTest() {
+        // Partition 1 in topic A and partition 0 in topic B are unavailable partitions.
+        List allPartitions = asList(new PartitionInfo(TOPIC_A, 0, NODES[0], NODES, NODES),
+            new PartitionInfo(TOPIC_A, 1, null, NODES, NODES),
+            new PartitionInfo(TOPIC_A, 2, NODES[2], NODES, NODES),
+            new PartitionInfo(TOPIC_B, 0, null, NODES, NODES),
+            new PartitionInfo(TOPIC_B, 1, NODES[0], NODES, NODES),
+            new PartitionInfo(TOPIC_C, 0, null, NODES, NODES)
+        );
+        
+        Cluster testCluster = new Cluster("clusterId", asList(NODES[0], NODES[1], NODES[2]), allPartitions,
+            Collections.emptySet(), Collections.emptySet());
+        StickyPartitionCache stickyPartitionCache = new StickyPartitionCache();
+        
+        // Assure we never choose partition 1 because it is unavailable.
+        int partA = stickyPartitionCache.partition(TOPIC_A, testCluster);
+        assertNotEquals(1, partA);
+        for (int aPartitions = 0; aPartitions < 100; aPartitions++) {
+            partA = stickyPartitionCache.nextPartition(TOPIC_A, testCluster, partA);
+            assertNotEquals(1, stickyPartitionCache.partition(TOPIC_A, testCluster));
+        }
+        
+        // Assure we always choose partition 1 for topic B.
+        int partB = stickyPartitionCache.partition(TOPIC_B, testCluster);
+        assertEquals(1, partB);
+        for (int bPartitions = 0; bPartitions < 100; bPartitions++) {
+            partB = stickyPartitionCache.nextPartition(TOPIC_B, testCluster, partB);
+            assertEquals(1, stickyPartitionCache.partition(TOPIC_B, testCluster));
+        }
+        
+        // Assure that we still choose the partition when there are no partitions available.
+        int partC = stickyPartitionCache.partition(TOPIC_C, testCluster);
+        assertEquals(0, partC);
+        partC = stickyPartitionCache.nextPartition(TOPIC_C, testCluster, partC);
+        assertEquals(0, partC);
+    }
+}
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 97f7f5d23b621..3d33a43e60d56 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
@@ -27,20 +27,25 @@
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.errors.GroupAuthorizationException;
 import org.apache.kafka.common.errors.OutOfOrderSequenceException;
+import org.apache.kafka.common.errors.ProducerFencedException;
 import org.apache.kafka.common.errors.TimeoutException;
 import org.apache.kafka.common.errors.TopicAuthorizationException;
 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.header.Header;
 import org.apache.kafka.common.internals.ClusterResourceListeners;
+import org.apache.kafka.common.message.InitProducerIdResponseData;
 import org.apache.kafka.common.metrics.MetricConfig;
 import org.apache.kafka.common.metrics.Metrics;
 import org.apache.kafka.common.protocol.Errors;
 import org.apache.kafka.common.record.CompressionType;
 import org.apache.kafka.common.record.MemoryRecords;
+import org.apache.kafka.common.record.MemoryRecordsBuilder;
 import org.apache.kafka.common.record.MutableRecordBatch;
 import org.apache.kafka.common.record.Record;
 import org.apache.kafka.common.record.RecordBatch;
+import org.apache.kafka.common.record.TimestampType;
 import org.apache.kafka.common.requests.AddOffsetsToTxnRequest;
 import org.apache.kafka.common.requests.AddOffsetsToTxnResponse;
 import org.apache.kafka.common.requests.AddPartitionsToTxnRequest;
@@ -63,6 +68,7 @@
 import org.junit.Before;
 import org.junit.Test;
 
+import java.nio.ByteBuffer;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
@@ -71,6 +77,7 @@
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.OptionalInt;
 import java.util.Set;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Future;
@@ -83,6 +90,7 @@
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
@@ -117,7 +125,7 @@ public void setup() {
         Map metricTags = new LinkedHashMap<>();
         metricTags.put("client-id", CLIENT_ID);
         int batchSize = 16 * 1024;
-        long deliveryTimeoutMs = 3000L;
+        int deliveryTimeoutMs = 3000;
         long totalSize = 1024 * 1024;
         String metricGrpName = "producer-metrics";
         MetricConfig metricConfig = new MetricConfig().tags(metricTags);
@@ -127,7 +135,7 @@ public void setup() {
         Metrics metrics = new Metrics(metricConfig, time);
         SenderMetricsRegistry senderMetrics = new SenderMetricsRegistry(metrics);
 
-        this.accumulator = new RecordAccumulator(logContext, batchSize, CompressionType.NONE, 0L, 0L,
+        this.accumulator = new RecordAccumulator(logContext, batchSize, CompressionType.NONE, 0, 0L,
                 deliveryTimeoutMs, metrics, metricGrpName, time, apiVersions, transactionManager,
                 new BufferPool(totalSize, batchSize, metrics, time, metricGrpName));
         this.sender = new Sender(logContext, this.client, this.metadata, this.accumulator, true, MAX_REQUEST_SIZE, ACKS_ALL,
@@ -137,20 +145,27 @@ public void setup() {
     }
 
     @Test
-    public void testSenderShutdownWithPendingAddPartitions() throws Exception {
+    public void testSenderShutdownWithPendingTransactions() throws Exception {
         long pid = 13131L;
         short epoch = 1;
         doInitTransactions(pid, epoch);
         transactionManager.beginTransaction();
 
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
         FutureRecordMetadata sendFuture = accumulator.append(tp0, time.milliseconds(), "key".getBytes(),
-                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT).future;
+                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT, false).future;
 
         prepareAddPartitionsToTxn(tp0, Errors.NONE);
         prepareProduceResponse(Errors.NONE, pid, epoch);
 
         sender.initiateClose();
+        sender.runOnce();
+        TransactionalRequestResult result = transactionManager.beginCommit();
+        sender.runOnce();
+        prepareEndTxnResponse(Errors.NONE, TransactionResult.COMMIT, pid, epoch);
+        sender.runOnce();
+        assertTrue(result.isCompleted());
         sender.run();
 
         assertTrue(sendFuture.isDone());
@@ -163,9 +178,10 @@ public void testEndTxnNotSentIfIncompleteBatches() {
         doInitTransactions(pid, epoch);
         transactionManager.beginTransaction();
 
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
         prepareAddPartitionsToTxn(tp0, Errors.NONE);
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertTrue(transactionManager.isPartitionAdded(tp0));
 
         transactionManager.beginCommit();
@@ -231,17 +247,18 @@ public void testHasOngoingTransactionSuccessfulAbort() {
         transactionManager.beginTransaction();
         assertTrue(transactionManager.hasOngoingTransaction());
 
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(partition);
         assertTrue(transactionManager.hasOngoingTransaction());
 
         prepareAddPartitionsToTxn(partition, Errors.NONE);
-        sender.run(time.milliseconds());
+        sender.runOnce();
 
         transactionManager.beginAbort();
         assertTrue(transactionManager.hasOngoingTransaction());
 
         prepareEndTxnResponse(Errors.NONE, TransactionResult.ABORT, pid, epoch);
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertFalse(transactionManager.hasOngoingTransaction());
     }
 
@@ -258,17 +275,18 @@ public void testHasOngoingTransactionSuccessfulCommit() {
         transactionManager.beginTransaction();
         assertTrue(transactionManager.hasOngoingTransaction());
 
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(partition);
         assertTrue(transactionManager.hasOngoingTransaction());
 
         prepareAddPartitionsToTxn(partition, Errors.NONE);
-        sender.run(time.milliseconds());
+        sender.runOnce();
 
         transactionManager.beginCommit();
         assertTrue(transactionManager.hasOngoingTransaction());
 
         prepareEndTxnResponse(Errors.NONE, TransactionResult.COMMIT, pid, epoch);
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertFalse(transactionManager.hasOngoingTransaction());
     }
 
@@ -285,11 +303,12 @@ public void testHasOngoingTransactionAbortableError() {
         transactionManager.beginTransaction();
         assertTrue(transactionManager.hasOngoingTransaction());
 
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(partition);
         assertTrue(transactionManager.hasOngoingTransaction());
 
         prepareAddPartitionsToTxn(partition, Errors.NONE);
-        sender.run(time.milliseconds());
+        sender.runOnce();
 
         transactionManager.transitionToAbortableError(new KafkaException());
         assertTrue(transactionManager.hasOngoingTransaction());
@@ -298,7 +317,7 @@ public void testHasOngoingTransactionAbortableError() {
         assertTrue(transactionManager.hasOngoingTransaction());
 
         prepareEndTxnResponse(Errors.NONE, TransactionResult.ABORT, pid, epoch);
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertFalse(transactionManager.hasOngoingTransaction());
     }
 
@@ -315,11 +334,12 @@ public void testHasOngoingTransactionFatalError() {
         transactionManager.beginTransaction();
         assertTrue(transactionManager.hasOngoingTransaction());
 
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(partition);
         assertTrue(transactionManager.hasOngoingTransaction());
 
         prepareAddPartitionsToTxn(partition, Errors.NONE);
-        sender.run(time.milliseconds());
+        sender.runOnce();
 
         transactionManager.transitionToFatalError(new KafkaException());
         assertFalse(transactionManager.hasOngoingTransaction());
@@ -333,19 +353,21 @@ public void testMaybeAddPartitionToTransaction() {
         doInitTransactions(pid, epoch);
         transactionManager.beginTransaction();
 
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(partition);
         assertTrue(transactionManager.hasPartitionsToAdd());
         assertFalse(transactionManager.isPartitionAdded(partition));
         assertTrue(transactionManager.isPartitionPendingAdd(partition));
 
         prepareAddPartitionsToTxn(partition, Errors.NONE);
-        sender.run(time.milliseconds());
+        sender.runOnce();
 
         assertFalse(transactionManager.hasPartitionsToAdd());
         assertTrue(transactionManager.isPartitionAdded(partition));
         assertFalse(transactionManager.isPartitionPendingAdd(partition));
 
         // adding the partition again should not have any effect
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(partition);
         assertFalse(transactionManager.hasPartitionsToAdd());
         assertTrue(transactionManager.isPartitionAdded(partition));
@@ -360,13 +382,14 @@ public void testAddPartitionToTransactionOverridesRetryBackoffForConcurrentTrans
         doInitTransactions(pid, epoch);
         transactionManager.beginTransaction();
 
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(partition);
         assertTrue(transactionManager.hasPartitionsToAdd());
         assertFalse(transactionManager.isPartitionAdded(partition));
         assertTrue(transactionManager.isPartitionPendingAdd(partition));
 
         prepareAddPartitionsToTxn(partition, Errors.CONCURRENT_TRANSACTIONS);
-        sender.run(time.milliseconds());
+        sender.runOnce();
 
         TransactionManager.TxnRequestHandler handler = transactionManager.nextRequestHandler(false);
         assertNotNull(handler);
@@ -381,13 +404,14 @@ public void testAddPartitionToTransactionRetainsRetryBackoffForRegularRetriableE
         doInitTransactions(pid, epoch);
         transactionManager.beginTransaction();
 
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(partition);
         assertTrue(transactionManager.hasPartitionsToAdd());
         assertFalse(transactionManager.isPartitionAdded(partition));
         assertTrue(transactionManager.isPartitionPendingAdd(partition));
 
         prepareAddPartitionsToTxn(partition, Errors.COORDINATOR_NOT_AVAILABLE);
-        sender.run(time.milliseconds());
+        sender.runOnce();
 
         TransactionManager.TxnRequestHandler handler = transactionManager.nextRequestHandler(false);
         assertNotNull(handler);
@@ -402,16 +426,18 @@ public void testAddPartitionToTransactionRetainsRetryBackoffWhenPartitionsAlread
         doInitTransactions(pid, epoch);
         transactionManager.beginTransaction();
 
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(partition);
         assertTrue(transactionManager.hasPartitionsToAdd());
         assertFalse(transactionManager.isPartitionAdded(partition));
         assertTrue(transactionManager.isPartitionPendingAdd(partition));
 
         prepareAddPartitionsToTxn(partition, Errors.NONE);
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertTrue(transactionManager.isPartitionAdded(partition));
 
         TopicPartition otherPartition = new TopicPartition("foo", 1);
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(otherPartition);
 
         prepareAddPartitionsToTxn(otherPartition, Errors.CONCURRENT_TRANSACTIONS);
@@ -422,6 +448,7 @@ public void testAddPartitionToTransactionRetainsRetryBackoffWhenPartitionsAlread
 
     @Test(expected = IllegalStateException.class)
     public void testMaybeAddPartitionToTransactionBeforeInitTransactions() {
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(new TopicPartition("foo", 0));
     }
 
@@ -430,6 +457,7 @@ public void testMaybeAddPartitionToTransactionBeforeBeginTransaction() {
         long pid = 13131L;
         short epoch = 1;
         doInitTransactions(pid, epoch);
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(new TopicPartition("foo", 0));
     }
 
@@ -440,6 +468,7 @@ public void testMaybeAddPartitionToTransactionAfterAbortableError() {
         doInitTransactions(pid, epoch);
         transactionManager.beginTransaction();
         transactionManager.transitionToAbortableError(new KafkaException());
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(new TopicPartition("foo", 0));
     }
 
@@ -449,6 +478,7 @@ public void testMaybeAddPartitionToTransactionAfterFatalError() {
         short epoch = 1;
         doInitTransactions(pid, epoch);
         transactionManager.transitionToFatalError(new KafkaException());
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(new TopicPartition("foo", 0));
     }
 
@@ -460,6 +490,7 @@ public void testIsSendToPartitionAllowedWithPendingPartitionAfterAbortableError(
         doInitTransactions(pid, epoch);
 
         transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
         transactionManager.transitionToAbortableError(new KafkaException());
 
@@ -475,10 +506,11 @@ public void testIsSendToPartitionAllowedWithInFlightPartitionAddAfterAbortableEr
         doInitTransactions(pid, epoch);
 
         transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
 
         // Send the AddPartitionsToTxn request and leave it in-flight
-        sender.run(time.milliseconds());
+        sender.runOnce();
         transactionManager.transitionToAbortableError(new KafkaException());
 
         assertFalse(transactionManager.isSendToPartitionAllowed(tp0));
@@ -493,6 +525,7 @@ public void testIsSendToPartitionAllowedWithPendingPartitionAfterFatalError() {
         doInitTransactions(pid, epoch);
 
         transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
         transactionManager.transitionToFatalError(new KafkaException());
 
@@ -508,10 +541,11 @@ public void testIsSendToPartitionAllowedWithInFlightPartitionAddAfterFatalError(
         doInitTransactions(pid, epoch);
 
         transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
 
         // Send the AddPartitionsToTxn request and leave it in-flight
-        sender.run(time.milliseconds());
+        sender.runOnce();
         transactionManager.transitionToFatalError(new KafkaException());
 
         assertFalse(transactionManager.isSendToPartitionAllowed(tp0));
@@ -527,9 +561,10 @@ public void testIsSendToPartitionAllowedWithAddedPartitionAfterAbortableError()
 
         transactionManager.beginTransaction();
 
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
         prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, pid);
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertFalse(transactionManager.hasPartitionsToAdd());
         transactionManager.transitionToAbortableError(new KafkaException());
 
@@ -545,9 +580,10 @@ public void testIsSendToPartitionAllowedWithAddedPartitionAfterFatalError() {
         doInitTransactions(pid, epoch);
 
         transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
         prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, pid);
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertFalse(transactionManager.hasPartitionsToAdd());
         transactionManager.transitionToFatalError(new KafkaException());
 
@@ -572,6 +608,164 @@ public void testDefaultSequenceNumber() {
         assertEquals((int) transactionManager.sequenceNumber(tp0), 3);
     }
 
+    @Test
+    public void testResetSequenceNumbersAfterUnknownProducerId() {
+        final long producerId = 13131L;
+        final short epoch = 1;
+        ProducerIdAndEpoch producerIdAndEpoch = new ProducerIdAndEpoch(producerId, epoch);
+
+        TransactionManager transactionManager = new TransactionManager();
+        transactionManager.setProducerIdAndEpoch(producerIdAndEpoch);
+
+        ProducerBatch b1 = writeIdempotentBatchWithValue(transactionManager, tp0, "1");
+        ProducerBatch b2 = writeIdempotentBatchWithValue(transactionManager, tp0, "2");
+        ProducerBatch b3 = writeIdempotentBatchWithValue(transactionManager, tp0, "3");
+        ProducerBatch b4 = writeIdempotentBatchWithValue(transactionManager, tp0, "4");
+        ProducerBatch b5 = writeIdempotentBatchWithValue(transactionManager, tp0, "5");
+        assertEquals(5, transactionManager.sequenceNumber(tp0).intValue());
+
+        // First batch succeeds
+        long b1AppendTime = time.milliseconds();
+        ProduceResponse.PartitionResponse b1Response = new ProduceResponse.PartitionResponse(
+                Errors.NONE, 500L, b1AppendTime, 0L);
+        b1.done(500L, b1AppendTime, null);
+        transactionManager.handleCompletedBatch(b1, b1Response);
+
+        // Retention caused log start offset to jump forward. We set sequence numbers back to 0
+        ProduceResponse.PartitionResponse b2Response = new ProduceResponse.PartitionResponse(
+                Errors.UNKNOWN_PRODUCER_ID, -1, -1, 600L);
+        assertTrue(transactionManager.canRetry(b2Response, b2));
+        assertEquals(4, transactionManager.sequenceNumber(tp0).intValue());
+        assertEquals(0, b2.baseSequence());
+        assertEquals(1, b3.baseSequence());
+        assertEquals(2, b4.baseSequence());
+        assertEquals(3, b5.baseSequence());
+    }
+
+    @Test
+    public void testAdjustSequenceNumbersAfterFatalError() {
+        final long producerId = 13131L;
+        final short epoch = 1;
+        ProducerIdAndEpoch producerIdAndEpoch = new ProducerIdAndEpoch(producerId, epoch);
+
+        TransactionManager transactionManager = new TransactionManager();
+        transactionManager.setProducerIdAndEpoch(producerIdAndEpoch);
+
+        ProducerBatch b1 = writeIdempotentBatchWithValue(transactionManager, tp0, "1");
+        ProducerBatch b2 = writeIdempotentBatchWithValue(transactionManager, tp0, "2");
+        ProducerBatch b3 = writeIdempotentBatchWithValue(transactionManager, tp0, "3");
+        ProducerBatch b4 = writeIdempotentBatchWithValue(transactionManager, tp0, "4");
+        ProducerBatch b5 = writeIdempotentBatchWithValue(transactionManager, tp0, "5");
+        assertEquals(5, transactionManager.sequenceNumber(tp0).intValue());
+
+        // First batch succeeds
+        long b1AppendTime = time.milliseconds();
+        ProduceResponse.PartitionResponse b1Response = new ProduceResponse.PartitionResponse(
+                Errors.NONE, 500L, b1AppendTime, 0L);
+        b1.done(500L, b1AppendTime, null);
+        transactionManager.handleCompletedBatch(b1, b1Response);
+
+        // Second batch fails with a fatal error. Sequence numbers are adjusted by one for remaining
+        // inflight batches.
+        ProduceResponse.PartitionResponse b2Response = new ProduceResponse.PartitionResponse(
+                Errors.MESSAGE_TOO_LARGE, -1, -1, 0L);
+        assertFalse(transactionManager.canRetry(b2Response, b2));
+
+        b2.done(-1L, -1L, Errors.MESSAGE_TOO_LARGE.exception());
+        transactionManager.handleFailedBatch(b2, Errors.MESSAGE_TOO_LARGE.exception(), true);
+        assertEquals(4, transactionManager.sequenceNumber(tp0).intValue());
+        assertEquals(1, b3.baseSequence());
+        assertEquals(2, b4.baseSequence());
+        assertEquals(3, b5.baseSequence());
+
+        // The remaining batches are doomed to fail, but they can be retried. Expected
+        // sequence numbers should remain the same.
+        ProduceResponse.PartitionResponse b3Response = new ProduceResponse.PartitionResponse(
+                Errors.OUT_OF_ORDER_SEQUENCE_NUMBER, -1, -1, 0L);
+        assertTrue(transactionManager.canRetry(b3Response, b3));
+        assertEquals(4, transactionManager.sequenceNumber(tp0).intValue());
+        assertEquals(1, b3.baseSequence());
+        assertEquals(2, b4.baseSequence());
+        assertEquals(3, b5.baseSequence());
+    }
+
+    @Test
+    public void testBatchFailureAfterProducerReset() {
+        // This tests a scenario where the producerId is reset while pending requests are still inflight.
+        // The returned responses should not update internal state.
+
+        final long producerId = 13131L;
+        final short epoch = 1;
+        ProducerIdAndEpoch producerIdAndEpoch = new ProducerIdAndEpoch(producerId, epoch);
+        TransactionManager transactionManager = new TransactionManager();
+        transactionManager.setProducerIdAndEpoch(producerIdAndEpoch);
+
+        ProducerBatch b1 = writeIdempotentBatchWithValue(transactionManager, tp0, "1");
+
+        ProducerIdAndEpoch updatedProducerIdAndEpoch = new ProducerIdAndEpoch(producerId + 1, epoch);
+        transactionManager.resetProducerId();
+        transactionManager.setProducerIdAndEpoch(updatedProducerIdAndEpoch);
+
+        ProducerBatch b2 = writeIdempotentBatchWithValue(transactionManager, tp0, "2");
+        assertEquals(1, transactionManager.sequenceNumber(tp0).intValue());
+
+        ProduceResponse.PartitionResponse b1Response = new ProduceResponse.PartitionResponse(
+                Errors.UNKNOWN_PRODUCER_ID, -1, -1, 400L);
+        assertFalse(transactionManager.canRetry(b1Response, b1));
+        transactionManager.handleFailedBatch(b1, Errors.UNKNOWN_PRODUCER_ID.exception(), true);
+
+        assertEquals(1, transactionManager.sequenceNumber(tp0).intValue());
+        assertEquals(b2, transactionManager.nextBatchBySequence(tp0));
+    }
+
+    @Test
+    public void testBatchCompletedAfterProducerReset() {
+        final long producerId = 13131L;
+        final short epoch = 1;
+        ProducerIdAndEpoch producerIdAndEpoch = new ProducerIdAndEpoch(producerId, epoch);
+        TransactionManager transactionManager = new TransactionManager();
+        transactionManager.setProducerIdAndEpoch(producerIdAndEpoch);
+
+        ProducerBatch b1 = writeIdempotentBatchWithValue(transactionManager, tp0, "1");
+
+        // The producerId might be reset due to a failure on another partition
+        ProducerIdAndEpoch updatedProducerIdAndEpoch = new ProducerIdAndEpoch(producerId + 1, epoch);
+        transactionManager.resetProducerId();
+        transactionManager.setProducerIdAndEpoch(updatedProducerIdAndEpoch);
+
+        ProducerBatch b2 = writeIdempotentBatchWithValue(transactionManager, tp0, "2");
+        assertEquals(1, transactionManager.sequenceNumber(tp0).intValue());
+
+        // If the request returns successfully, we should ignore the response and not update any state
+        ProduceResponse.PartitionResponse b1Response = new ProduceResponse.PartitionResponse(
+                Errors.NONE, 500L, time.milliseconds(), 0L);
+        transactionManager.handleCompletedBatch(b1, b1Response);
+
+        assertEquals(1, transactionManager.sequenceNumber(tp0).intValue());
+        assertEquals(b2, transactionManager.nextBatchBySequence(tp0));
+    }
+
+    private ProducerBatch writeIdempotentBatchWithValue(TransactionManager manager,
+                                                        TopicPartition tp,
+                                                        String value) {
+        int seq = manager.sequenceNumber(tp);
+        manager.incrementSequenceNumber(tp, 1);
+        ProducerBatch batch = batchWithValue(tp, value);
+        batch.setProducerState(manager.producerIdAndEpoch(), seq, false);
+        manager.addInFlightBatch(batch);
+        batch.close();
+        return batch;
+    }
+
+    private ProducerBatch batchWithValue(TopicPartition tp, String value) {
+        MemoryRecordsBuilder builder = MemoryRecords.builder(ByteBuffer.allocate(64),
+                CompressionType.NONE, TimestampType.CREATE_TIME, 0L);
+        long currentTimeMs = time.milliseconds();
+        ProducerBatch batch = new ProducerBatch(tp, builder, currentTimeMs);
+        batch.tryAppend(currentTimeMs, new byte[0], value.getBytes(), new Header[0], null, currentTimeMs);
+        return batch;
+    }
+
     @Test
     public void testSequenceNumberOverflow() {
         TransactionManager transactionManager = new TransactionManager();
@@ -602,10 +796,11 @@ public void testBasicTransaction() throws InterruptedException {
         doInitTransactions(pid, epoch);
 
         transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
 
         Future responseFuture = accumulator.append(tp0, time.milliseconds(), "key".getBytes(),
-                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT).future;
+                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT, false).future;
 
         assertFalse(responseFuture.isDone());
         prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, pid);
@@ -613,13 +808,13 @@ public void testBasicTransaction() throws InterruptedException {
         prepareProduceResponse(Errors.NONE, pid, epoch);
         assertFalse(transactionManager.transactionContainsPartition(tp0));
         assertFalse(transactionManager.isSendToPartitionAllowed(tp0));
-        sender.run(time.milliseconds());  // send addPartitions.
+        sender.runOnce();  // send addPartitions.
         // Check that only addPartitions was sent.
         assertTrue(transactionManager.transactionContainsPartition(tp0));
         assertTrue(transactionManager.isSendToPartitionAllowed(tp0));
         assertFalse(responseFuture.isDone());
 
-        sender.run(time.milliseconds());  // send produce request.
+        sender.runOnce();  // send produce request.
         assertTrue(responseFuture.isDone());
 
         Map offsets = new HashMap<>();
@@ -631,7 +826,7 @@ public void testBasicTransaction() throws InterruptedException {
 
         prepareAddOffsetsToTxnResponse(Errors.NONE, consumerGroupId, pid, epoch);
 
-        sender.run(time.milliseconds());  // Send AddOffsetsRequest
+        sender.runOnce();  // Send AddOffsetsRequest
         assertTrue(transactionManager.hasPendingOffsetCommits());  // We should now have created and queued the offset commit request.
         assertFalse(addOffsetsResult.isCompleted()); // the result doesn't complete until TxnOffsetCommit returns
 
@@ -642,19 +837,19 @@ public void testBasicTransaction() throws InterruptedException {
         prepareTxnOffsetCommitResponse(consumerGroupId, pid, epoch, txnOffsetCommitResponse);
 
         assertNull(transactionManager.coordinator(CoordinatorType.GROUP));
-        sender.run(time.milliseconds());  // try to send TxnOffsetCommitRequest, but find we don't have a group coordinator.
-        sender.run(time.milliseconds());  // send find coordinator for group request
+        sender.runOnce();  // try to send TxnOffsetCommitRequest, but find we don't have a group coordinator.
+        sender.runOnce();  // send find coordinator for group request
         assertNotNull(transactionManager.coordinator(CoordinatorType.GROUP));
         assertTrue(transactionManager.hasPendingOffsetCommits());
 
-        sender.run(time.milliseconds());  // send TxnOffsetCommitRequest commit.
+        sender.runOnce();  // send TxnOffsetCommitRequest commit.
 
         assertFalse(transactionManager.hasPendingOffsetCommits());
         assertTrue(addOffsetsResult.isCompleted());  // We should only be done after both RPCs complete.
 
         transactionManager.beginCommit();
         prepareEndTxnResponse(Errors.NONE, TransactionResult.COMMIT, pid, epoch);
-        sender.run(time.milliseconds());  // commit.
+        sender.runOnce();  // commit.
 
         assertFalse(transactionManager.hasOngoingTransaction());
         assertFalse(transactionManager.isCompleting());
@@ -667,11 +862,11 @@ public void testDisconnectAndRetry() {
         // It finds the coordinator and then gets a PID.
         transactionManager.initializeTransactions();
         prepareFindCoordinatorResponse(Errors.NONE, true, CoordinatorType.TRANSACTION, transactionalId);
-        sender.run(time.milliseconds());  // find coordinator, connection lost.
+        sender.runOnce();  // find coordinator, connection lost.
 
         prepareFindCoordinatorResponse(Errors.NONE, false, CoordinatorType.TRANSACTION, transactionalId);
-        sender.run(time.milliseconds());  // find coordinator
-        sender.run(time.milliseconds());
+        sender.runOnce();  // find coordinator
+        sender.runOnce();
         assertEquals(brokerNode, transactionManager.coordinator(CoordinatorType.TRANSACTION));
     }
 
@@ -680,13 +875,13 @@ public void testUnsupportedFindCoordinator() {
         transactionManager.initializeTransactions();
         client.prepareUnsupportedVersionResponse(body -> {
             FindCoordinatorRequest findCoordinatorRequest = (FindCoordinatorRequest) body;
-            assertEquals(findCoordinatorRequest.coordinatorType(), CoordinatorType.TRANSACTION);
-            assertEquals(findCoordinatorRequest.coordinatorKey(), transactionalId);
+            assertEquals(CoordinatorType.forId(findCoordinatorRequest.data().keyType()), CoordinatorType.TRANSACTION);
+            assertEquals(findCoordinatorRequest.data().key(), transactionalId);
             return true;
         });
 
-        sender.run(time.milliseconds()); // InitProducerRequest is queued
-        sender.run(time.milliseconds()); // FindCoordinator is queued after peeking InitProducerRequest
+        sender.runOnce(); // InitProducerRequest is queued
+        sender.runOnce(); // FindCoordinator is queued after peeking InitProducerRequest
         assertTrue(transactionManager.hasFatalError());
         assertTrue(transactionManager.lastError() instanceof UnsupportedVersionException);
     }
@@ -695,20 +890,20 @@ public void testUnsupportedFindCoordinator() {
     public void testUnsupportedInitTransactions() {
         transactionManager.initializeTransactions();
         prepareFindCoordinatorResponse(Errors.NONE, false, CoordinatorType.TRANSACTION, transactionalId);
-        sender.run(time.milliseconds()); // InitProducerRequest is queued
-        sender.run(time.milliseconds()); // FindCoordinator is queued after peeking InitProducerRequest
+        sender.runOnce(); // InitProducerRequest is queued
+        sender.runOnce(); // FindCoordinator is queued after peeking InitProducerRequest
 
         assertFalse(transactionManager.hasError());
         assertNotNull(transactionManager.coordinator(CoordinatorType.TRANSACTION));
 
         client.prepareUnsupportedVersionResponse(body -> {
             InitProducerIdRequest initProducerIdRequest = (InitProducerIdRequest) body;
-            assertEquals(initProducerIdRequest.transactionalId(), transactionalId);
-            assertEquals(initProducerIdRequest.transactionTimeoutMs(), transactionTimeoutMs);
+            assertEquals(initProducerIdRequest.data.transactionalId(), transactionalId);
+            assertEquals(initProducerIdRequest.data.transactionTimeoutMs(), transactionTimeoutMs);
             return true;
         });
 
-        sender.run(time.milliseconds()); // InitProducerRequest is dequeued
+        sender.runOnce(); // InitProducerRequest is dequeued
         assertTrue(transactionManager.hasFatalError());
         assertTrue(transactionManager.lastError() instanceof UnsupportedVersionException);
     }
@@ -727,14 +922,14 @@ public void testUnsupportedForMessageFormatInTxnOffsetCommit() {
                 singletonMap(tp, new OffsetAndMetadata(39L)), consumerGroupId);
 
         prepareAddOffsetsToTxnResponse(Errors.NONE, consumerGroupId, pid, epoch);
-        sender.run(time.milliseconds());  // AddOffsetsToTxn Handled, TxnOffsetCommit Enqueued
-        sender.run(time.milliseconds());  // FindCoordinator Enqueued
+        sender.runOnce();  // AddOffsetsToTxn Handled, TxnOffsetCommit Enqueued
+        sender.runOnce();  // FindCoordinator Enqueued
 
         prepareFindCoordinatorResponse(Errors.NONE, false, CoordinatorType.GROUP, consumerGroupId);
-        sender.run(time.milliseconds());  // FindCoordinator Returned
+        sender.runOnce();  // FindCoordinator Returned
 
         prepareTxnOffsetCommitResponse(consumerGroupId, pid, epoch, singletonMap(tp, Errors.UNSUPPORTED_FOR_MESSAGE_FORMAT));
-        sender.run(time.milliseconds());  // TxnOffsetCommit Handled
+        sender.runOnce();  // TxnOffsetCommit Handled
 
         assertTrue(transactionManager.hasError());
         assertTrue(transactionManager.lastError() instanceof UnsupportedForMessageFormatException);
@@ -752,24 +947,24 @@ public void testLookupCoordinatorOnDisconnectAfterSend() {
         final short epoch = 1;
         TransactionalRequestResult initPidResult = transactionManager.initializeTransactions();
         prepareFindCoordinatorResponse(Errors.NONE, false, CoordinatorType.TRANSACTION, transactionalId);
-        sender.run(time.milliseconds());  // find coordinator
-        sender.run(time.milliseconds());
+        sender.runOnce();  // find coordinator
+        sender.runOnce();
         assertEquals(brokerNode, transactionManager.coordinator(CoordinatorType.TRANSACTION));
         prepareInitPidResponse(Errors.NONE, true, pid, epoch);
         // send pid to coordinator, should get disconnected before receiving the response, and resend the
         // FindCoordinator and InitPid requests.
-        sender.run(time.milliseconds());
+        sender.runOnce();
 
         assertNull(transactionManager.coordinator(CoordinatorType.TRANSACTION));
         assertFalse(initPidResult.isCompleted());
         assertFalse(transactionManager.hasProducerId());
 
         prepareFindCoordinatorResponse(Errors.NONE, false, CoordinatorType.TRANSACTION, transactionalId);
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertEquals(brokerNode, transactionManager.coordinator(CoordinatorType.TRANSACTION));
         assertFalse(initPidResult.isCompleted());
         prepareInitPidResponse(Errors.NONE, false, pid, epoch);
-        sender.run(time.milliseconds());  // get pid and epoch
+        sender.runOnce();  // get pid and epoch
 
         assertTrue(initPidResult.isCompleted()); // The future should only return after the second round of retries succeed.
         assertTrue(transactionManager.hasProducerId());
@@ -785,15 +980,15 @@ public void testLookupCoordinatorOnDisconnectBeforeSend() {
         final short epoch = 1;
         TransactionalRequestResult initPidResult = transactionManager.initializeTransactions();
         prepareFindCoordinatorResponse(Errors.NONE, false, CoordinatorType.TRANSACTION, transactionalId);
-        sender.run(time.milliseconds());  // one loop to realize we need a coordinator.
-        sender.run(time.milliseconds());  // next loop to find coordintor.
+        sender.runOnce();  // one loop to realize we need a coordinator.
+        sender.runOnce();  // next loop to find coordintor.
         assertEquals(brokerNode, transactionManager.coordinator(CoordinatorType.TRANSACTION));
 
         client.disconnect(brokerNode.idString());
         client.blackout(brokerNode, 100);
         // send pid to coordinator. Should get disconnected before the send and resend the FindCoordinator
         // and InitPid requests.
-        sender.run(time.milliseconds());
+        sender.runOnce();
         time.sleep(110);  // waiting for the blackout period for the node to expire.
 
         assertNull(transactionManager.coordinator(CoordinatorType.TRANSACTION));
@@ -801,11 +996,11 @@ public void testLookupCoordinatorOnDisconnectBeforeSend() {
         assertFalse(transactionManager.hasProducerId());
 
         prepareFindCoordinatorResponse(Errors.NONE, false, CoordinatorType.TRANSACTION, transactionalId);
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertEquals(brokerNode, transactionManager.coordinator(CoordinatorType.TRANSACTION));
         assertFalse(initPidResult.isCompleted());
         prepareInitPidResponse(Errors.NONE, false, pid, epoch);
-        sender.run(time.milliseconds());  // get pid and epoch
+        sender.runOnce();  // get pid and epoch
 
         assertTrue(initPidResult.isCompleted()); // The future should only return after the second round of retries succeed.
         assertTrue(transactionManager.hasProducerId());
@@ -821,23 +1016,23 @@ public void testLookupCoordinatorOnNotCoordinatorError() {
         final short epoch = 1;
         TransactionalRequestResult initPidResult = transactionManager.initializeTransactions();
         prepareFindCoordinatorResponse(Errors.NONE, false, CoordinatorType.TRANSACTION, transactionalId);
-        sender.run(time.milliseconds());  // find coordinator
-        sender.run(time.milliseconds());
+        sender.runOnce();  // find coordinator
+        sender.runOnce();
         assertEquals(brokerNode, transactionManager.coordinator(CoordinatorType.TRANSACTION));
 
         prepareInitPidResponse(Errors.NOT_COORDINATOR, false, pid, epoch);
-        sender.run(time.milliseconds());  // send pid, get not coordinator. Should resend the FindCoordinator and InitPid requests
+        sender.runOnce();  // send pid, get not coordinator. Should resend the FindCoordinator and InitPid requests
 
         assertNull(transactionManager.coordinator(CoordinatorType.TRANSACTION));
         assertFalse(initPidResult.isCompleted());
         assertFalse(transactionManager.hasProducerId());
 
         prepareFindCoordinatorResponse(Errors.NONE, false, CoordinatorType.TRANSACTION, transactionalId);
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertEquals(brokerNode, transactionManager.coordinator(CoordinatorType.TRANSACTION));
         assertFalse(initPidResult.isCompleted());
         prepareInitPidResponse(Errors.NONE, false, pid, epoch);
-        sender.run(time.milliseconds());  // get pid and epoch
+        sender.runOnce();  // get pid and epoch
 
         assertTrue(initPidResult.isCompleted()); // The future should only return after the second round of retries succeed.
         assertTrue(transactionManager.hasProducerId());
@@ -850,13 +1045,13 @@ public void testTransactionalIdAuthorizationFailureInFindCoordinator() {
         TransactionalRequestResult initPidResult = transactionManager.initializeTransactions();
         prepareFindCoordinatorResponse(Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED, false,
                 CoordinatorType.TRANSACTION, transactionalId);
-        sender.run(time.milliseconds());  // find coordinator
-        sender.run(time.milliseconds());
+        sender.runOnce();  // find coordinator
+        sender.runOnce();
 
         assertTrue(transactionManager.hasError());
         assertTrue(transactionManager.lastError() instanceof TransactionalIdAuthorizationException);
 
-        sender.run(time.milliseconds()); // one more run to fail the InitProducerId future
+        sender.runOnce(); // one more run to fail the InitProducerId future
         assertTrue(initPidResult.isCompleted());
         assertFalse(initPidResult.isSuccessful());
         assertTrue(initPidResult.error() instanceof TransactionalIdAuthorizationException);
@@ -869,12 +1064,12 @@ public void testTransactionalIdAuthorizationFailureInInitProducerId() {
         final long pid = 13131L;
         TransactionalRequestResult initPidResult = transactionManager.initializeTransactions();
         prepareFindCoordinatorResponse(Errors.NONE, false, CoordinatorType.TRANSACTION, transactionalId);
-        sender.run(time.milliseconds());  // find coordinator
-        sender.run(time.milliseconds());
+        sender.runOnce();  // find coordinator
+        sender.runOnce();
         assertEquals(brokerNode, transactionManager.coordinator(CoordinatorType.TRANSACTION));
 
         prepareInitPidResponse(Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED, false, pid, RecordBatch.NO_PRODUCER_EPOCH);
-        sender.run(time.milliseconds());
+        sender.runOnce();
 
         assertTrue(transactionManager.hasError());
         assertTrue(initPidResult.isCompleted());
@@ -897,12 +1092,12 @@ public void testGroupAuthorizationFailureInFindCoordinator() {
                 singletonMap(new TopicPartition("foo", 0), new OffsetAndMetadata(39L)), consumerGroupId);
 
         prepareAddOffsetsToTxnResponse(Errors.NONE, consumerGroupId, pid, epoch);
-        sender.run(time.milliseconds());  // AddOffsetsToTxn Handled, TxnOffsetCommit Enqueued
-        sender.run(time.milliseconds());  // FindCoordinator Enqueued
+        sender.runOnce();  // AddOffsetsToTxn Handled, TxnOffsetCommit Enqueued
+        sender.runOnce();  // FindCoordinator Enqueued
 
         prepareFindCoordinatorResponse(Errors.GROUP_AUTHORIZATION_FAILED, false, CoordinatorType.GROUP, consumerGroupId);
-        sender.run(time.milliseconds());  // FindCoordinator Failed
-        sender.run(time.milliseconds());  // TxnOffsetCommit Aborted
+        sender.runOnce();  // FindCoordinator Failed
+        sender.runOnce();  // TxnOffsetCommit Aborted
         assertTrue(transactionManager.hasError());
         assertTrue(transactionManager.lastError() instanceof GroupAuthorizationException);
         assertTrue(sendOffsetsResult.isCompleted());
@@ -929,14 +1124,14 @@ public void testGroupAuthorizationFailureInTxnOffsetCommit() {
                 singletonMap(tp1, new OffsetAndMetadata(39L)), consumerGroupId);
 
         prepareAddOffsetsToTxnResponse(Errors.NONE, consumerGroupId, pid, epoch);
-        sender.run(time.milliseconds());  // AddOffsetsToTxn Handled, TxnOffsetCommit Enqueued
-        sender.run(time.milliseconds());  // FindCoordinator Enqueued
+        sender.runOnce();  // AddOffsetsToTxn Handled, TxnOffsetCommit Enqueued
+        sender.runOnce();  // FindCoordinator Enqueued
 
         prepareFindCoordinatorResponse(Errors.NONE, false, CoordinatorType.GROUP, consumerGroupId);
-        sender.run(time.milliseconds());  // FindCoordinator Returned
+        sender.runOnce();  // FindCoordinator Returned
 
         prepareTxnOffsetCommitResponse(consumerGroupId, pid, epoch, singletonMap(tp1, Errors.GROUP_AUTHORIZATION_FAILED));
-        sender.run(time.milliseconds());  // TxnOffsetCommit Handled
+        sender.runOnce();  // TxnOffsetCommit Handled
 
         assertTrue(transactionManager.hasError());
         assertTrue(transactionManager.lastError() instanceof GroupAuthorizationException);
@@ -965,7 +1160,7 @@ public void testTransactionalIdAuthorizationFailureInAddOffsetsToTxn() {
                 singletonMap(tp, new OffsetAndMetadata(39L)), consumerGroupId);
 
         prepareAddOffsetsToTxnResponse(Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED, consumerGroupId, pid, epoch);
-        sender.run(time.milliseconds());  // AddOffsetsToTxn Handled
+        sender.runOnce();  // AddOffsetsToTxn Handled
 
         assertTrue(transactionManager.hasError());
         assertTrue(transactionManager.lastError() instanceof TransactionalIdAuthorizationException);
@@ -990,14 +1185,14 @@ public void testTransactionalIdAuthorizationFailureInTxnOffsetCommit() {
                 singletonMap(tp, new OffsetAndMetadata(39L)), consumerGroupId);
 
         prepareAddOffsetsToTxnResponse(Errors.NONE, consumerGroupId, pid, epoch);
-        sender.run(time.milliseconds());  // AddOffsetsToTxn Handled, TxnOffsetCommit Enqueued
-        sender.run(time.milliseconds());  // FindCoordinator Enqueued
+        sender.runOnce();  // AddOffsetsToTxn Handled, TxnOffsetCommit Enqueued
+        sender.runOnce();  // FindCoordinator Enqueued
 
         prepareFindCoordinatorResponse(Errors.NONE, false, CoordinatorType.GROUP, consumerGroupId);
-        sender.run(time.milliseconds());  // FindCoordinator Returned
+        sender.runOnce();  // FindCoordinator Returned
 
         prepareTxnOffsetCommitResponse(consumerGroupId, pid, epoch, singletonMap(tp, Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED));
-        sender.run(time.milliseconds());  // TxnOffsetCommit Handled
+        sender.runOnce();  // TxnOffsetCommit Handled
 
         assertTrue(transactionManager.hasError());
         assertTrue(transactionManager.lastError() instanceof TransactionalIdAuthorizationException);
@@ -1018,14 +1213,16 @@ public void testTopicAuthorizationFailureInAddPartitions() {
         doInitTransactions(pid, epoch);
 
         transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp1);
         Map errors = new HashMap<>();
         errors.put(tp0, Errors.TOPIC_AUTHORIZATION_FAILED);
         errors.put(tp1, Errors.OPERATION_NOT_ATTEMPTED);
 
         prepareAddPartitionsToTxn(errors);
-        sender.run(time.milliseconds());
+        sender.runOnce();
 
         assertTrue(transactionManager.hasError());
         assertTrue(transactionManager.lastError() instanceof TopicAuthorizationException);
@@ -1050,22 +1247,23 @@ public void testRecoveryFromAbortableErrorTransactionNotStarted() throws Excepti
         doInitTransactions(pid, epoch);
 
         transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(unauthorizedPartition);
 
         Future responseFuture = accumulator.append(unauthorizedPartition, time.milliseconds(), "key".getBytes(),
-                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT).future;
+                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT, false).future;
 
         prepareAddPartitionsToTxn(singletonMap(unauthorizedPartition, Errors.TOPIC_AUTHORIZATION_FAILED));
-        sender.run(time.milliseconds());
+        sender.runOnce();
 
         assertTrue(transactionManager.hasAbortableError());
         transactionManager.beginAbort();
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertTrue(responseFuture.isDone());
         assertFutureFailed(responseFuture);
 
         // No partitions added, so no need to prepare EndTxn response
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertTrue(transactionManager.isReady());
         assertFalse(transactionManager.hasPartitionsToAdd());
         assertFalse(accumulator.hasIncomplete());
@@ -1073,25 +1271,26 @@ public void testRecoveryFromAbortableErrorTransactionNotStarted() throws Excepti
         // ensure we can now start a new transaction
 
         transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
 
         responseFuture = accumulator.append(tp0, time.milliseconds(), "key".getBytes(),
-                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT).future;
+                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT, false).future;
 
         prepareAddPartitionsToTxn(singletonMap(tp0, Errors.NONE));
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertTrue(transactionManager.isPartitionAdded(tp0));
         assertFalse(transactionManager.hasPartitionsToAdd());
 
         transactionManager.beginCommit();
         prepareProduceResponse(Errors.NONE, pid, epoch);
-        sender.run(time.milliseconds());
+        sender.runOnce();
 
         assertTrue(responseFuture.isDone());
         assertNotNull(responseFuture.get());
 
         prepareEndTxnResponse(Errors.NONE, TransactionResult.COMMIT, pid, epoch);
-        sender.run(time.milliseconds());
+        sender.runOnce();
 
         assertTrue(transactionManager.isReady());
     }
@@ -1105,19 +1304,21 @@ public void testRecoveryFromAbortableErrorTransactionStarted() throws Exception
         doInitTransactions(pid, epoch);
 
         transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
         prepareAddPartitionsToTxn(tp0, Errors.NONE);
 
         Future authorizedTopicProduceFuture = accumulator.append(unauthorizedPartition, time.milliseconds(),
-                "key".getBytes(), "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds());
+                "key".getBytes(), "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT, false).future;
+        sender.runOnce();
         assertTrue(transactionManager.isPartitionAdded(tp0));
 
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(unauthorizedPartition);
         Future unauthorizedTopicProduceFuture = accumulator.append(unauthorizedPartition, time.milliseconds(),
-                "key".getBytes(), "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT).future;
+                "key".getBytes(), "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT, false).future;
         prepareAddPartitionsToTxn(singletonMap(unauthorizedPartition, Errors.TOPIC_AUTHORIZATION_FAILED));
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertTrue(transactionManager.hasAbortableError());
         assertTrue(transactionManager.isPartitionAdded(tp0));
         assertFalse(transactionManager.isPartitionAdded(unauthorizedPartition));
@@ -1126,7 +1327,7 @@ public void testRecoveryFromAbortableErrorTransactionStarted() throws Exception
 
         prepareEndTxnResponse(Errors.NONE, TransactionResult.ABORT, pid, epoch);
         transactionManager.beginAbort();
-        sender.run(time.milliseconds());
+        sender.runOnce();
         // neither produce request has been sent, so they should both be failed immediately
         assertFutureFailed(authorizedTopicProduceFuture);
         assertFutureFailed(unauthorizedTopicProduceFuture);
@@ -1137,25 +1338,26 @@ public void testRecoveryFromAbortableErrorTransactionStarted() throws Exception
         // ensure we can now start a new transaction
 
         transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
 
         FutureRecordMetadata nextTransactionFuture = accumulator.append(tp0, time.milliseconds(), "key".getBytes(),
-                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT).future;
+                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT, false).future;
 
         prepareAddPartitionsToTxn(singletonMap(tp0, Errors.NONE));
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertTrue(transactionManager.isPartitionAdded(tp0));
         assertFalse(transactionManager.hasPartitionsToAdd());
 
         transactionManager.beginCommit();
         prepareProduceResponse(Errors.NONE, pid, epoch);
-        sender.run(time.milliseconds());
+        sender.runOnce();
 
         assertTrue(nextTransactionFuture.isDone());
         assertNotNull(nextTransactionFuture.get());
 
         prepareEndTxnResponse(Errors.NONE, TransactionResult.COMMIT, pid, epoch);
-        sender.run(time.milliseconds());
+        sender.runOnce();
 
         assertTrue(transactionManager.isReady());
     }
@@ -1169,32 +1371,34 @@ public void testRecoveryFromAbortableErrorProduceRequestInRetry() throws Excepti
         doInitTransactions(pid, epoch);
 
         transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
         prepareAddPartitionsToTxn(tp0, Errors.NONE);
 
         Future authorizedTopicProduceFuture = accumulator.append(tp0, time.milliseconds(),
-                "key".getBytes(), "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds());
+                "key".getBytes(), "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT, false).future;
+        sender.runOnce();
         assertTrue(transactionManager.isPartitionAdded(tp0));
 
         accumulator.beginFlush();
         prepareProduceResponse(Errors.REQUEST_TIMED_OUT, pid, epoch);
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertFalse(authorizedTopicProduceFuture.isDone());
         assertTrue(accumulator.hasIncomplete());
 
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(unauthorizedPartition);
         Future unauthorizedTopicProduceFuture = accumulator.append(unauthorizedPartition, time.milliseconds(),
-                "key".getBytes(), "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT).future;
+                "key".getBytes(), "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT, false).future;
         prepareAddPartitionsToTxn(singletonMap(unauthorizedPartition, Errors.TOPIC_AUTHORIZATION_FAILED));
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertTrue(transactionManager.hasAbortableError());
         assertTrue(transactionManager.isPartitionAdded(tp0));
         assertFalse(transactionManager.isPartitionAdded(unauthorizedPartition));
         assertFalse(authorizedTopicProduceFuture.isDone());
 
         prepareProduceResponse(Errors.NONE, pid, epoch);
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertFutureFailed(unauthorizedTopicProduceFuture);
         assertTrue(authorizedTopicProduceFuture.isDone());
         assertNotNull(authorizedTopicProduceFuture.get());
@@ -1202,7 +1406,7 @@ public void testRecoveryFromAbortableErrorProduceRequestInRetry() throws Excepti
 
         prepareEndTxnResponse(Errors.NONE, TransactionResult.ABORT, pid, epoch);
         transactionManager.beginAbort();
-        sender.run(time.milliseconds());
+        sender.runOnce();
         // neither produce request has been sent, so they should both be failed immediately
         assertTrue(transactionManager.isReady());
         assertFalse(transactionManager.hasPartitionsToAdd());
@@ -1211,25 +1415,26 @@ public void testRecoveryFromAbortableErrorProduceRequestInRetry() throws Excepti
         // ensure we can now start a new transaction
 
         transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
 
         FutureRecordMetadata nextTransactionFuture = accumulator.append(tp0, time.milliseconds(), "key".getBytes(),
-                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT).future;
+                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT, false).future;
 
         prepareAddPartitionsToTxn(singletonMap(tp0, Errors.NONE));
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertTrue(transactionManager.isPartitionAdded(tp0));
         assertFalse(transactionManager.hasPartitionsToAdd());
 
         transactionManager.beginCommit();
         prepareProduceResponse(Errors.NONE, pid, epoch);
-        sender.run(time.milliseconds());
+        sender.runOnce();
 
         assertTrue(nextTransactionFuture.isDone());
         assertNotNull(nextTransactionFuture.get());
 
         prepareEndTxnResponse(Errors.NONE, TransactionResult.COMMIT, pid, epoch);
-        sender.run(time.milliseconds());
+        sender.runOnce();
 
         assertTrue(transactionManager.isReady());
     }
@@ -1243,10 +1448,11 @@ public void testTransactionalIdAuthorizationFailureInAddPartitions() {
         doInitTransactions(pid, epoch);
 
         transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp);
 
         prepareAddPartitionsToTxn(tp, Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED);
-        sender.run(time.milliseconds());
+        sender.runOnce();
 
         assertTrue(transactionManager.hasError());
         assertTrue(transactionManager.lastError() instanceof TransactionalIdAuthorizationException);
@@ -1262,10 +1468,11 @@ public void testFlushPendingPartitionsOnCommit() throws InterruptedException {
         doInitTransactions(pid, epoch);
 
         transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
 
         Future responseFuture = accumulator.append(tp0, time.milliseconds(), "key".getBytes(),
-                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT).future;
+                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT, false).future;
 
         assertFalse(responseFuture.isDone());
 
@@ -1279,13 +1486,13 @@ public void testFlushPendingPartitionsOnCommit() throws InterruptedException {
         assertFalse(transactionManager.transactionContainsPartition(tp0));
         prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, pid);
 
-        sender.run(time.milliseconds());  // AddPartitions.
+        sender.runOnce();  // AddPartitions.
         assertTrue(transactionManager.transactionContainsPartition(tp0));
         assertFalse(responseFuture.isDone());
         assertFalse(commitResult.isCompleted());
 
         prepareProduceResponse(Errors.NONE, pid, epoch);
-        sender.run(time.milliseconds());  // Produce.
+        sender.runOnce();  // Produce.
         assertTrue(responseFuture.isDone());
 
         prepareEndTxnResponse(Errors.NONE, TransactionResult.COMMIT, pid, epoch);
@@ -1293,7 +1500,7 @@ public void testFlushPendingPartitionsOnCommit() throws InterruptedException {
         assertTrue(transactionManager.hasOngoingTransaction());
         assertTrue(transactionManager.isCompleting());
 
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertTrue(commitResult.isCompleted());
         assertFalse(transactionManager.hasOngoingTransaction());
     }
@@ -1306,11 +1513,12 @@ public void testMultipleAddPartitionsPerForOneProduce() throws InterruptedExcept
         doInitTransactions(pid, epoch);
 
         transactionManager.beginTransaction();
-        // User does one producer.sed
+        // User does one producer.send
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
 
         Future responseFuture = accumulator.append(tp0, time.milliseconds(), "key".getBytes(),
-                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT).future;
+                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT, false).future;
 
         assertFalse(responseFuture.isDone());
         prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, pid);
@@ -1318,14 +1526,15 @@ public void testMultipleAddPartitionsPerForOneProduce() throws InterruptedExcept
         assertFalse(transactionManager.transactionContainsPartition(tp0));
 
         // Sender flushes one add partitions. The produce goes next.
-        sender.run(time.milliseconds());  // send addPartitions.
+        sender.runOnce();  // send addPartitions.
         // Check that only addPartitions was sent.
         assertTrue(transactionManager.transactionContainsPartition(tp0));
 
         // In the mean time, the user does a second produce to a different partition
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp1);
         Future secondResponseFuture = accumulator.append(tp0, time.milliseconds(), "key".getBytes(),
-                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT).future;
+                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT, false).future;
 
         prepareAddPartitionsToTxnResponse(Errors.NONE, tp1, epoch, pid);
         prepareProduceResponse(Errors.NONE, pid, epoch);
@@ -1336,42 +1545,56 @@ public void testMultipleAddPartitionsPerForOneProduce() throws InterruptedExcept
         assertFalse(secondResponseFuture.isDone());
 
         // The second add partitions should go out here.
-        sender.run(time.milliseconds());  // send second add partitions request
+        sender.runOnce();  // send second add partitions request
         assertTrue(transactionManager.transactionContainsPartition(tp1));
 
         assertFalse(responseFuture.isDone());
         assertFalse(secondResponseFuture.isDone());
 
         // Finally we get to the produce.
-        sender.run(time.milliseconds());  // send produce request
+        sender.runOnce();  // send produce request
 
         assertTrue(responseFuture.isDone());
         assertTrue(secondResponseFuture.isDone());
     }
 
-    @Test(expected = ExecutionException.class)
-    public void testProducerFencedException() throws InterruptedException, ExecutionException {
+    @Test
+    public void testProducerFencedException() throws InterruptedException {
         final long pid = 13131L;
         final short epoch = 1;
 
         doInitTransactions(pid, epoch);
 
         transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
 
         Future responseFuture = accumulator.append(tp0, time.milliseconds(), "key".getBytes(),
-                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT).future;
+                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT, false).future;
 
         assertFalse(responseFuture.isDone());
         prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, pid);
         prepareProduceResponse(Errors.INVALID_PRODUCER_EPOCH, pid, epoch);
-        sender.run(time.milliseconds()); // Add partitions.
+        sender.runOnce(); // Add partitions.
 
-        sender.run(time.milliseconds());  // send produce.
+        sender.runOnce();  // send produce.
 
         assertTrue(responseFuture.isDone());
         assertTrue(transactionManager.hasError());
-        responseFuture.get();
+
+        try {
+            // make sure the produce was expired.
+            responseFuture.get();
+            fail("Expected to get a ExecutionException from the response");
+        } catch (ExecutionException e) {
+            assertTrue(e.getCause() instanceof ProducerFencedException);
+        }
+
+        // make sure the exception was thrown directly from the follow-up calls.
+        assertThrows(ProducerFencedException.class, () -> transactionManager.beginTransaction());
+        assertThrows(ProducerFencedException.class, () -> transactionManager.beginCommit());
+        assertThrows(ProducerFencedException.class, () -> transactionManager.beginAbort());
+        assertThrows(ProducerFencedException.class, () -> transactionManager.sendOffsetsToTransaction(Collections.emptyMap(), "dummyId"));
     }
 
     @Test
@@ -1382,21 +1605,22 @@ public void testDisallowCommitOnProduceFailure() throws InterruptedException {
         doInitTransactions(pid, epoch);
 
         transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
 
         Future responseFuture = accumulator.append(tp0, time.milliseconds(), "key".getBytes(),
-                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT).future;
+                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT, false).future;
 
         TransactionalRequestResult commitResult = transactionManager.beginCommit();
         assertFalse(responseFuture.isDone());
         prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, pid);
         prepareProduceResponse(Errors.OUT_OF_ORDER_SEQUENCE_NUMBER, pid, epoch);
 
-        sender.run(time.milliseconds());  // Send AddPartitionsRequest
+        sender.runOnce();  // Send AddPartitionsRequest
         assertFalse(commitResult.isCompleted());
-        sender.run(time.milliseconds());  // Send Produce Request, returns OutOfOrderSequenceException.
+        sender.runOnce();  // Send Produce Request, returns OutOfOrderSequenceException.
 
-        sender.run(time.milliseconds());  // try to commit.
+        sender.runOnce();  // try to commit.
         assertTrue(commitResult.isCompleted());  // commit should be cancelled with exception without being sent.
 
         try {
@@ -1416,7 +1640,7 @@ public void testDisallowCommitOnProduceFailure() throws InterruptedException {
         // Commit is not allowed, so let's abort and try again.
         TransactionalRequestResult abortResult = transactionManager.beginAbort();
         prepareEndTxnResponse(Errors.NONE, TransactionResult.ABORT, pid, epoch);
-        sender.run(time.milliseconds());  // Send abort request. It is valid to transition from ERROR to ABORT
+        sender.runOnce();  // Send abort request. It is valid to transition from ERROR to ABORT
 
         assertTrue(abortResult.isCompleted());
         assertTrue(abortResult.isSuccessful());
@@ -1431,21 +1655,22 @@ public void testAllowAbortOnProduceFailure() throws InterruptedException {
         doInitTransactions(pid, epoch);
 
         transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
 
         Future responseFuture = accumulator.append(tp0, time.milliseconds(), "key".getBytes(),
-                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT).future;
+                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT, false).future;
 
         assertFalse(responseFuture.isDone());
         prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, pid);
         prepareProduceResponse(Errors.OUT_OF_ORDER_SEQUENCE_NUMBER, pid, epoch);
         prepareEndTxnResponse(Errors.NONE, TransactionResult.ABORT, pid, epoch);
 
-        sender.run(time.milliseconds());  // Send AddPartitionsRequest
-        sender.run(time.milliseconds());  // Send Produce Request, returns OutOfOrderSequenceException.
+        sender.runOnce();  // Send AddPartitionsRequest
+        sender.runOnce();  // Send Produce Request, returns OutOfOrderSequenceException.
 
         TransactionalRequestResult abortResult = transactionManager.beginAbort();
-        sender.run(time.milliseconds());  // try to abort
+        sender.runOnce();  // try to abort
         assertTrue(abortResult.isCompleted());
         assertTrue(abortResult.isSuccessful());
         assertTrue(transactionManager.isReady());  // make sure we are ready for a transaction now.
@@ -1459,16 +1684,17 @@ public void testAbortableErrorWhileAbortInProgress() throws InterruptedException
         doInitTransactions(pid, epoch);
 
         transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
 
         Future responseFuture = accumulator.append(tp0, time.milliseconds(), "key".getBytes(),
-                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT).future;
+                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT, false).future;
 
         assertFalse(responseFuture.isDone());
         prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, pid);
 
-        sender.run(time.milliseconds());  // Send AddPartitionsRequest
-        sender.run(time.milliseconds());  // Send Produce Request
+        sender.runOnce();  // Send AddPartitionsRequest
+        sender.runOnce();  // Send Produce Request
 
         TransactionalRequestResult abortResult = transactionManager.beginAbort();
         assertTrue(transactionManager.isAborting());
@@ -1476,13 +1702,13 @@ public void testAbortableErrorWhileAbortInProgress() throws InterruptedException
 
         sendProduceResponse(Errors.OUT_OF_ORDER_SEQUENCE_NUMBER, pid, epoch);
         prepareEndTxnResponse(Errors.NONE, TransactionResult.ABORT, pid, epoch);
-        sender.run(time.milliseconds());  // receive the produce response
+        sender.runOnce();  // receive the produce response
 
         // we do not transition to ABORTABLE_ERROR since we were already aborting
         assertTrue(transactionManager.isAborting());
         assertFalse(transactionManager.hasError());
 
-        sender.run(time.milliseconds());  // handle the abort
+        sender.runOnce();  // handle the abort
         assertTrue(abortResult.isCompleted());
         assertTrue(abortResult.isSuccessful());
         assertTrue(transactionManager.isReady());  // make sure we are ready for a transaction now.
@@ -1496,25 +1722,26 @@ public void testCommitTransactionWithUnsentProduceRequest() throws Exception {
         doInitTransactions(pid, epoch);
 
         transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
 
         Future responseFuture = accumulator.append(tp0, time.milliseconds(), "key".getBytes(),
-                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT).future;
+                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT, false).future;
 
         prepareAddPartitionsToTxn(tp0, Errors.NONE);
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertTrue(accumulator.hasUndrained());
 
         // committing the transaction should cause the unsent batch to be flushed
         transactionManager.beginCommit();
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertFalse(accumulator.hasUndrained());
         assertTrue(accumulator.hasIncomplete());
         assertFalse(transactionManager.hasInFlightTransactionalRequest());
         assertFalse(responseFuture.isDone());
 
         // until the produce future returns, we will not send EndTxn
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertFalse(accumulator.hasUndrained());
         assertTrue(accumulator.hasIncomplete());
         assertFalse(transactionManager.hasInFlightTransactionalRequest());
@@ -1522,17 +1749,17 @@ public void testCommitTransactionWithUnsentProduceRequest() throws Exception {
 
         // now the produce response returns
         sendProduceResponse(Errors.NONE, pid, epoch);
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertTrue(responseFuture.isDone());
         assertFalse(accumulator.hasUndrained());
         assertFalse(accumulator.hasIncomplete());
         assertFalse(transactionManager.hasInFlightTransactionalRequest());
 
         // now we send EndTxn
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertTrue(transactionManager.hasInFlightTransactionalRequest());
         sendEndTxnResponse(Errors.NONE, TransactionResult.COMMIT, pid, epoch);
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertFalse(transactionManager.hasInFlightTransactionalRequest());
         assertTrue(transactionManager.isReady());
     }
@@ -1545,31 +1772,32 @@ public void testCommitTransactionWithInFlightProduceRequest() throws Exception {
         doInitTransactions(pid, epoch);
 
         transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
 
         Future responseFuture = accumulator.append(tp0, time.milliseconds(), "key".getBytes(),
-                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT).future;
+                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT, false).future;
 
         prepareAddPartitionsToTxn(tp0, Errors.NONE);
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertTrue(accumulator.hasUndrained());
 
         accumulator.beginFlush();
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertFalse(accumulator.hasUndrained());
         assertTrue(accumulator.hasIncomplete());
         assertFalse(transactionManager.hasInFlightTransactionalRequest());
 
         // now we begin the commit with the produce request still pending
         transactionManager.beginCommit();
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertFalse(accumulator.hasUndrained());
         assertTrue(accumulator.hasIncomplete());
         assertFalse(transactionManager.hasInFlightTransactionalRequest());
         assertFalse(responseFuture.isDone());
 
         // until the produce future returns, we will not send EndTxn
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertFalse(accumulator.hasUndrained());
         assertTrue(accumulator.hasIncomplete());
         assertFalse(transactionManager.hasInFlightTransactionalRequest());
@@ -1577,17 +1805,17 @@ public void testCommitTransactionWithInFlightProduceRequest() throws Exception {
 
         // now the produce response returns
         sendProduceResponse(Errors.NONE, pid, epoch);
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertTrue(responseFuture.isDone());
         assertFalse(accumulator.hasUndrained());
         assertFalse(accumulator.hasIncomplete());
         assertFalse(transactionManager.hasInFlightTransactionalRequest());
 
         // now we send EndTxn
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertTrue(transactionManager.hasInFlightTransactionalRequest());
         sendEndTxnResponse(Errors.NONE, TransactionResult.COMMIT, pid, epoch);
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertFalse(transactionManager.hasInFlightTransactionalRequest());
         assertTrue(transactionManager.isReady());
     }
@@ -1600,22 +1828,23 @@ public void testFindCoordinatorAllowedInAbortableErrorState() throws Interrupted
         doInitTransactions(pid, epoch);
 
         transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
 
         Future responseFuture = accumulator.append(tp0, time.milliseconds(), "key".getBytes(),
-                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT).future;
+                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT, false).future;
 
         assertFalse(responseFuture.isDone());
-        sender.run(time.milliseconds());  // Send AddPartitionsRequest
+        sender.runOnce();  // Send AddPartitionsRequest
 
         transactionManager.transitionToAbortableError(new KafkaException());
         sendAddPartitionsToTxnResponse(Errors.NOT_COORDINATOR, tp0, epoch, pid);
-        sender.run(time.milliseconds()); // AddPartitions returns
+        sender.runOnce(); // AddPartitions returns
         assertTrue(transactionManager.hasAbortableError());
 
         assertNull(transactionManager.coordinator(CoordinatorType.TRANSACTION));
         prepareFindCoordinatorResponse(Errors.NONE, false, CoordinatorType.TRANSACTION, transactionalId);
-        sender.run(time.milliseconds()); // FindCoordinator handled
+        sender.runOnce(); // FindCoordinator handled
         assertEquals(brokerNode, transactionManager.coordinator(CoordinatorType.TRANSACTION));
         assertTrue(transactionManager.hasAbortableError());
     }
@@ -1628,17 +1857,18 @@ public void testCancelUnsentAddPartitionsAndProduceOnAbort() throws InterruptedE
         doInitTransactions(pid, epoch);
 
         transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
 
         Future responseFuture = accumulator.append(tp0, time.milliseconds(), "key".getBytes(),
-                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT).future;
+                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT, false).future;
 
         assertFalse(responseFuture.isDone());
 
         TransactionalRequestResult abortResult = transactionManager.beginAbort();
         // note since no partitions were added to the transaction, no EndTxn will be sent
 
-        sender.run(time.milliseconds());  // try to abort
+        sender.runOnce();  // try to abort
         assertTrue(abortResult.isCompleted());
         assertTrue(abortResult.isSuccessful());
         assertTrue(transactionManager.isReady());  // make sure we are ready for a transaction now.
@@ -1659,13 +1889,14 @@ public void testAbortResendsAddPartitionErrorIfRetried() throws InterruptedExcep
         doInitTransactions(producerId, producerEpoch);
 
         transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
         prepareAddPartitionsToTxnResponse(Errors.UNKNOWN_TOPIC_OR_PARTITION, tp0, producerEpoch, producerId);
 
         Future responseFuture = accumulator.append(tp0, time.milliseconds(), "key".getBytes(),
-                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT).future;
+                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT, false).future;
 
-        sender.run(time.milliseconds());  // Send AddPartitions and let it fail
+        sender.runOnce();  // Send AddPartitions and let it fail
         assertFalse(responseFuture.isDone());
 
         TransactionalRequestResult abortResult = transactionManager.beginAbort();
@@ -1674,8 +1905,8 @@ public void testAbortResendsAddPartitionErrorIfRetried() throws InterruptedExcep
         prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, producerEpoch, producerId);
         prepareEndTxnResponse(Errors.NONE, TransactionResult.ABORT, producerId, producerEpoch);
 
-        sender.run(time.milliseconds());  // Resend AddPartitions
-        sender.run(time.milliseconds());  // Send EndTxn
+        sender.runOnce();  // Resend AddPartitions
+        sender.runOnce();  // Send EndTxn
 
         assertTrue(abortResult.isCompleted());
         assertTrue(abortResult.isSuccessful());
@@ -1697,15 +1928,16 @@ public void testAbortResendsProduceRequestIfRetried() throws Exception {
         doInitTransactions(producerId, producerEpoch);
 
         transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
         prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, producerEpoch, producerId);
         prepareProduceResponse(Errors.REQUEST_TIMED_OUT, producerId, producerEpoch);
 
         Future responseFuture = accumulator.append(tp0, time.milliseconds(), "key".getBytes(),
-                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT).future;
+                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT, false).future;
 
-        sender.run(time.milliseconds());  // Send AddPartitions
-        sender.run(time.milliseconds());  // Send ProduceRequest and let it fail
+        sender.runOnce();  // Send AddPartitions
+        sender.runOnce();  // Send ProduceRequest and let it fail
 
         assertFalse(responseFuture.isDone());
 
@@ -1715,8 +1947,8 @@ public void testAbortResendsProduceRequestIfRetried() throws Exception {
         prepareProduceResponse(Errors.NONE, producerId, producerEpoch);
         prepareEndTxnResponse(Errors.NONE, TransactionResult.ABORT, producerId, producerEpoch);
 
-        sender.run(time.milliseconds());  // Resend ProduceRequest
-        sender.run(time.milliseconds());  // Send EndTxn
+        sender.runOnce();  // Resend ProduceRequest
+        sender.runOnce();  // Send EndTxn
 
         assertTrue(abortResult.isCompleted());
         assertTrue(abortResult.isSuccessful());
@@ -1734,23 +1966,24 @@ public void testHandlingOfUnknownTopicPartitionErrorOnAddPartitions() throws Int
         doInitTransactions(pid, epoch);
 
         transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
 
         Future responseFuture = accumulator.append(tp0, time.milliseconds(), "key".getBytes(),
-                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT).future;
+                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT, false).future;
 
         assertFalse(responseFuture.isDone());
         prepareAddPartitionsToTxnResponse(Errors.UNKNOWN_TOPIC_OR_PARTITION, tp0, epoch, pid);
 
-        sender.run(time.milliseconds());  // Send AddPartitionsRequest
+        sender.runOnce();  // Send AddPartitionsRequest
         assertFalse(transactionManager.transactionContainsPartition(tp0));  // The partition should not yet be added.
 
         prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, pid);
         prepareProduceResponse(Errors.NONE, pid, epoch);
-        sender.run(time.milliseconds());  // Send AddPartitionsRequest successfully.
+        sender.runOnce();  // Send AddPartitionsRequest successfully.
         assertTrue(transactionManager.transactionContainsPartition(tp0));
 
-        sender.run(time.milliseconds());  // Send ProduceRequest.
+        sender.runOnce();  // Send ProduceRequest.
         assertTrue(responseFuture.isDone());
     }
 
@@ -1780,7 +2013,7 @@ private void testRetriableErrorInTxnOffsetCommit(Errors error) {
         TransactionalRequestResult addOffsetsResult = transactionManager.sendOffsetsToTransaction(offsets, consumerGroupId);
         prepareAddOffsetsToTxnResponse(Errors.NONE, consumerGroupId, pid, epoch);
 
-        sender.run(time.milliseconds());  // send AddOffsetsToTxnResult
+        sender.runOnce();  // send AddOffsetsToTxnResult
 
         assertFalse(addOffsetsResult.isCompleted());  // The request should complete only after the TxnOffsetCommit completes.
 
@@ -1792,19 +2025,19 @@ private void testRetriableErrorInTxnOffsetCommit(Errors error) {
         prepareTxnOffsetCommitResponse(consumerGroupId, pid, epoch, txnOffsetCommitResponse);
 
         assertNull(transactionManager.coordinator(CoordinatorType.GROUP));
-        sender.run(time.milliseconds());  // try to send TxnOffsetCommitRequest, but find we don't have a group coordinator.
-        sender.run(time.milliseconds());  // send find coordinator for group request
+        sender.runOnce();  // try to send TxnOffsetCommitRequest, but find we don't have a group coordinator.
+        sender.runOnce();  // send find coordinator for group request
         assertNotNull(transactionManager.coordinator(CoordinatorType.GROUP));
         assertTrue(transactionManager.hasPendingOffsetCommits());
 
-        sender.run(time.milliseconds());  // send TxnOffsetCommitRequest request.
+        sender.runOnce();  // send TxnOffsetCommitRequest request.
 
         assertTrue(transactionManager.hasPendingOffsetCommits());  // The TxnOffsetCommit failed.
         assertFalse(addOffsetsResult.isCompleted());  // We should only be done after both RPCs complete successfully.
 
         txnOffsetCommitResponse.put(tp1, Errors.NONE);
         prepareTxnOffsetCommitResponse(consumerGroupId, pid, epoch, txnOffsetCommitResponse);
-        sender.run(time.milliseconds());  // Send TxnOffsetCommitRequest again.
+        sender.runOnce();  // Send TxnOffsetCommitRequest again.
 
         assertTrue(addOffsetsResult.isCompleted());
         assertTrue(addOffsetsResult.isSuccessful());
@@ -1823,15 +2056,16 @@ public void shouldNotSendAbortTxnRequestWhenOnlyAddPartitionsRequestFailed() {
         doInitTransactions(pid, epoch);
 
         transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
 
         prepareAddPartitionsToTxnResponse(Errors.TOPIC_AUTHORIZATION_FAILED, tp0, epoch, pid);
-        sender.run(time.milliseconds());  // Send AddPartitionsRequest
+        sender.runOnce();  // Send AddPartitionsRequest
 
         TransactionalRequestResult abortResult = transactionManager.beginAbort();
         assertFalse(abortResult.isCompleted());
 
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertTrue(abortResult.isCompleted());
         assertTrue(abortResult.isSuccessful());
     }
@@ -1853,10 +2087,10 @@ public void shouldNotSendAbortTxnRequestWhenOnlyAddOffsetsRequestFailed() {
         TransactionalRequestResult abortResult = transactionManager.beginAbort();
 
         prepareAddOffsetsToTxnResponse(Errors.GROUP_AUTHORIZATION_FAILED, consumerGroupId, pid, epoch);
-        sender.run(time.milliseconds());  // Send AddOffsetsToTxnRequest
+        sender.runOnce();  // Send AddOffsetsToTxnRequest
         assertFalse(abortResult.isCompleted());
 
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertTrue(transactionManager.isReady());
         assertTrue(abortResult.isCompleted());
         assertTrue(abortResult.isSuccessful());
@@ -1879,10 +2113,10 @@ public void shouldFailAbortIfAddOffsetsFailsWithFatalError() {
         TransactionalRequestResult abortResult = transactionManager.beginAbort();
 
         prepareAddOffsetsToTxnResponse(Errors.UNKNOWN_SERVER_ERROR, consumerGroupId, pid, epoch);
-        sender.run(time.milliseconds());  // Send AddOffsetsToTxnRequest
+        sender.runOnce();  // Send AddOffsetsToTxnRequest
         assertFalse(abortResult.isCompleted());
 
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertTrue(abortResult.isCompleted());
         assertFalse(abortResult.isSuccessful());
         assertTrue(transactionManager.hasFatalError());
@@ -1894,12 +2128,14 @@ public void testNoDrainWhenPartitionsPending() throws InterruptedException {
         final short epoch = 1;
         doInitTransactions(pid, epoch);
         transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
         accumulator.append(tp0, time.milliseconds(), "key".getBytes(),
-                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT);
+                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT, false);
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp1);
         accumulator.append(tp1, time.milliseconds(), "key".getBytes(),
-                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT);
+                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT, false);
 
         assertFalse(transactionManager.isSendToPartitionAllowed(tp0));
         assertFalse(transactionManager.isSendToPartitionAllowed(tp1));
@@ -1931,15 +2167,17 @@ public void testAllowDrainInAbortableErrorState() throws InterruptedException {
         final short epoch = 1;
         doInitTransactions(pid, epoch);
         transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp1);
         prepareAddPartitionsToTxn(tp1, Errors.NONE);
-        sender.run(time.milliseconds());  // Send AddPartitions, tp1 should be in the transaction now.
+        sender.runOnce();  // Send AddPartitions, tp1 should be in the transaction now.
 
         assertTrue(transactionManager.transactionContainsPartition(tp1));
 
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
         prepareAddPartitionsToTxn(tp0, Errors.TOPIC_AUTHORIZATION_FAILED);
-        sender.run(time.milliseconds());  // Send AddPartitions, should be in abortable state.
+        sender.runOnce();  // Send AddPartitions, should be in abortable state.
 
         assertTrue(transactionManager.hasAbortableError());
         assertTrue(transactionManager.isSendToPartitionAllowed(tp1));
@@ -1950,7 +2188,7 @@ public void testAllowDrainInAbortableErrorState() throws InterruptedException {
         Cluster cluster = new Cluster(null, Collections.singletonList(node1), Collections.singletonList(part1),
                 Collections.emptySet(), Collections.emptySet());
         accumulator.append(tp1, time.milliseconds(), "key".getBytes(),
-                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT);
+                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT, false);
         Map> drainedBatches = accumulator.drain(cluster, Collections.singleton(node1),
                 Integer.MAX_VALUE,
                 time.milliseconds());
@@ -1970,7 +2208,7 @@ public void testRaiseErrorWhenNoPartitionsPendingOnDrain() throws InterruptedExc
         transactionManager.beginTransaction();
         // Don't execute transactionManager.maybeAddPartitionToTransaction(tp0). This should result in an error on drain.
         accumulator.append(tp0, time.milliseconds(), "key".getBytes(),
-                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT);
+                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT, false);
         Node node1 = new Node(0, "localhost", 1111);
         PartitionInfo part1 = new PartitionInfo(topic, 0, node1, null, null);
 
@@ -1993,21 +2231,22 @@ public void resendFailedProduceRequestAfterAbortableError() throws Exception {
         doInitTransactions(pid, epoch);
         transactionManager.beginTransaction();
 
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
 
         Future responseFuture = accumulator.append(tp0, time.milliseconds(), "key".getBytes(),
-                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT).future;
+                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT, false).future;
 
         prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, pid);
         prepareProduceResponse(Errors.NOT_LEADER_FOR_PARTITION, pid, epoch);
-        sender.run(time.milliseconds()); // Add partitions
-        sender.run(time.milliseconds()); // Produce
+        sender.runOnce(); // Add partitions
+        sender.runOnce(); // Produce
 
         assertFalse(responseFuture.isDone());
 
         transactionManager.transitionToAbortableError(new KafkaException());
         prepareProduceResponse(Errors.NONE, pid, epoch);
-        sender.run(time.milliseconds());
+        sender.runOnce();
 
         assertTrue(responseFuture.isDone());
         assertNotNull(responseFuture.get()); // should throw the exception which caused the transaction to be aborted.
@@ -2021,10 +2260,11 @@ public void testTransitionToAbortableErrorOnBatchExpiry() throws InterruptedExce
         doInitTransactions(pid, epoch);
 
         transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
 
         Future responseFuture = accumulator.append(tp0, time.milliseconds(), "key".getBytes(),
-                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT).future;
+                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT, false).future;
 
         assertFalse(responseFuture.isDone());
 
@@ -2032,7 +2272,7 @@ public void testTransitionToAbortableErrorOnBatchExpiry() throws InterruptedExce
 
         assertFalse(transactionManager.transactionContainsPartition(tp0));
         assertFalse(transactionManager.isSendToPartitionAllowed(tp0));
-        sender.run(time.milliseconds());  // send addPartitions.
+        sender.runOnce();  // send addPartitions.
         // Check that only addPartitions was sent.
         assertTrue(transactionManager.transactionContainsPartition(tp0));
         assertTrue(transactionManager.isSendToPartitionAllowed(tp0));
@@ -2046,7 +2286,7 @@ public void testTransitionToAbortableErrorOnBatchExpiry() throws InterruptedExce
         client.disconnect(clusterNode.idString());
         client.blackout(clusterNode, 100);
 
-        sender.run(time.milliseconds());  // We should try to flush the produce, but expire it instead without sending anything.
+        sender.runOnce();  // We should try to flush the produce, but expire it instead without sending anything.
         assertTrue(responseFuture.isDone());
 
         try {
@@ -2067,13 +2307,15 @@ public void testTransitionToAbortableErrorOnMultipleBatchExpiry() throws Interru
         doInitTransactions(pid, epoch);
 
         transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp1);
 
         Future firstBatchResponse = accumulator.append(tp0, time.milliseconds(), "key".getBytes(),
-                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT).future;
+                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT, false).future;
         Future secondBatchResponse = accumulator.append(tp1, time.milliseconds(), "key".getBytes(),
-               "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT).future;
+               "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT, false).future;
 
         assertFalse(firstBatchResponse.isDone());
         assertFalse(secondBatchResponse.isDone());
@@ -2085,7 +2327,7 @@ public void testTransitionToAbortableErrorOnMultipleBatchExpiry() throws Interru
 
         assertFalse(transactionManager.transactionContainsPartition(tp0));
         assertFalse(transactionManager.isSendToPartitionAllowed(tp0));
-        sender.run(time.milliseconds());  // send addPartitions.
+        sender.runOnce();  // send addPartitions.
         // Check that only addPartitions was sent.
         assertTrue(transactionManager.transactionContainsPartition(tp0));
         assertTrue(transactionManager.transactionContainsPartition(tp1));
@@ -2102,7 +2344,7 @@ public void testTransitionToAbortableErrorOnMultipleBatchExpiry() throws Interru
         client.disconnect(clusterNode.idString());
         client.blackout(clusterNode, 100);
 
-        sender.run(time.milliseconds());  // We should try to flush the produce, but expire it instead without sending anything.
+        sender.runOnce();  // We should try to flush the produce, but expire it instead without sending anything.
         assertTrue(firstBatchResponse.isDone());
         assertTrue(secondBatchResponse.isDone());
 
@@ -2132,10 +2374,11 @@ public void testDropCommitOnBatchExpiry() throws InterruptedException {
         doInitTransactions(pid, epoch);
 
         transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
 
         Future responseFuture = accumulator.append(tp0, time.milliseconds(), "key".getBytes(),
-                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT).future;
+                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT, false).future;
 
         assertFalse(responseFuture.isDone());
 
@@ -2143,7 +2386,7 @@ public void testDropCommitOnBatchExpiry() throws InterruptedException {
 
         assertFalse(transactionManager.transactionContainsPartition(tp0));
         assertFalse(transactionManager.isSendToPartitionAllowed(tp0));
-        sender.run(time.milliseconds());  // send addPartitions.
+        sender.runOnce();  // send addPartitions.
         // Check that only addPartitions was sent.
         assertTrue(transactionManager.transactionContainsPartition(tp0));
         assertTrue(transactionManager.isSendToPartitionAllowed(tp0));
@@ -2157,9 +2400,8 @@ public void testDropCommitOnBatchExpiry() throws InterruptedException {
         // expire the batch.
         Node clusterNode = metadata.fetch().nodes().get(0);
         client.disconnect(clusterNode.idString());
-        client.blackout(clusterNode, 100);
 
-        sender.run(time.milliseconds());  // We should try to flush the produce, but expire it instead without sending anything.
+        sender.runOnce();  // We should try to flush the produce, but expire it instead without sending anything.
         assertTrue(responseFuture.isDone());
 
         try {
@@ -2169,7 +2411,7 @@ public void testDropCommitOnBatchExpiry() throws InterruptedException {
         } catch (ExecutionException e) {
             assertTrue(e.getCause() instanceof  TimeoutException);
         }
-        sender.run(time.milliseconds());  // the commit shouldn't be completed without being sent since the produce request failed.
+        sender.runOnce();  // the commit shouldn't be completed without being sent since the produce request failed.
 
         assertTrue(commitResult.isCompleted());
         assertFalse(commitResult.isSuccessful());  // the commit shouldn't succeed since the produce request failed.
@@ -2183,7 +2425,7 @@ public void testDropCommitOnBatchExpiry() throws InterruptedException {
 
         prepareEndTxnResponse(Errors.NONE, TransactionResult.ABORT, pid, epoch);
 
-        sender.run(time.milliseconds());  // send the abort.
+        sender.runOnce();  // send the abort.
 
         assertTrue(abortResult.isCompleted());
         assertTrue(abortResult.isSuccessful());
@@ -2199,10 +2441,11 @@ public void testTransitionToFatalErrorWhenRetriedBatchIsExpired() throws Interru
         doInitTransactions(pid, epoch);
 
         transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
 
         Future responseFuture = accumulator.append(tp0, time.milliseconds(), "key".getBytes(),
-                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT).future;
+                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT, false).future;
 
         assertFalse(responseFuture.isDone());
 
@@ -2210,13 +2453,13 @@ public void testTransitionToFatalErrorWhenRetriedBatchIsExpired() throws Interru
 
         assertFalse(transactionManager.transactionContainsPartition(tp0));
         assertFalse(transactionManager.isSendToPartitionAllowed(tp0));
-        sender.run(time.milliseconds());  // send addPartitions.
+        sender.runOnce();  // send addPartitions.
         // Check that only addPartitions was sent.
         assertTrue(transactionManager.transactionContainsPartition(tp0));
         assertTrue(transactionManager.isSendToPartitionAllowed(tp0));
 
         prepareProduceResponse(Errors.NOT_LEADER_FOR_PARTITION, pid, epoch);
-        sender.run(time.milliseconds());  // send the produce request.
+        sender.runOnce();  // send the produce request.
 
         assertFalse(responseFuture.isDone());
 
@@ -2230,7 +2473,7 @@ public void testTransitionToFatalErrorWhenRetriedBatchIsExpired() throws Interru
         client.disconnect(clusterNode.idString());
         client.blackout(clusterNode, 100);
 
-        sender.run(time.milliseconds());  // We should try to flush the produce, but expire it instead without sending anything.
+        sender.runOnce();  // We should try to flush the produce, but expire it instead without sending anything.
         assertTrue(responseFuture.isDone());
 
         try {
@@ -2240,8 +2483,8 @@ public void testTransitionToFatalErrorWhenRetriedBatchIsExpired() throws Interru
         } catch (ExecutionException e) {
             assertTrue(e.getCause() instanceof  TimeoutException);
         }
-        sender.run(time.milliseconds());  // Transition to fatal error since we have unresolved batches.
-        sender.run(time.milliseconds());  // Fail the queued transactional requests
+        sender.runOnce();  // Transition to fatal error since we have unresolved batches.
+        sender.runOnce();  // Fail the queued transactional requests
 
         assertTrue(commitResult.isCompleted());
         assertFalse(commitResult.isSuccessful());  // the commit should have been dropped.
@@ -2251,30 +2494,121 @@ public void testTransitionToFatalErrorWhenRetriedBatchIsExpired() throws Interru
     }
 
     @Test
-    public void testShouldResetProducerStateAfterResolvingSequences() {
-        // Create a TransactionManager without a transactionalId to test
-        // shouldResetProducerStateAfterResolvingSequences.
+    public void testResetProducerIdAfterWithoutPendingInflightRequests() {
         TransactionManager manager = new TransactionManager(logContext, null, transactionTimeoutMs,
             DEFAULT_RETRY_BACKOFF_MS);
-        assertFalse(manager.shouldResetProducerStateAfterResolvingSequences());
+        long producerId = 15L;
+        short epoch = 5;
+        ProducerIdAndEpoch producerIdAndEpoch = new ProducerIdAndEpoch(producerId, epoch);
+        manager.setProducerIdAndEpoch(producerIdAndEpoch);
+
+        // Nothing to resolve, so no reset is needed
+        manager.resetProducerIdIfNeeded();
+        assertEquals(producerIdAndEpoch, manager.producerIdAndEpoch());
+
         TopicPartition tp0 = new TopicPartition("foo", 0);
-        TopicPartition tp1 = new TopicPartition("foo", 1);
         assertEquals(Integer.valueOf(0), manager.sequenceNumber(tp0));
-        assertEquals(Integer.valueOf(0), manager.sequenceNumber(tp1));
 
-        manager.incrementSequenceNumber(tp0, 1);
-        manager.incrementSequenceNumber(tp1, 1);
-        manager.maybeUpdateLastAckedSequence(tp0, 0);
-        manager.maybeUpdateLastAckedSequence(tp1, 0);
+        ProducerBatch b1 = writeIdempotentBatchWithValue(manager, tp0, "1");
+        assertEquals(Integer.valueOf(1), manager.sequenceNumber(tp0));
+        manager.handleCompletedBatch(b1, new ProduceResponse.PartitionResponse(
+                Errors.NONE, 500L, time.milliseconds(), 0L));
+        assertEquals(OptionalInt.of(0), manager.lastAckedSequence(tp0));
+
+        // Marking sequence numbers unresolved without inflight requests is basically a no-op.
         manager.markSequenceUnresolved(tp0);
-        manager.markSequenceUnresolved(tp1);
-        assertFalse(manager.shouldResetProducerStateAfterResolvingSequences());
+        manager.resetProducerIdIfNeeded();
+        assertEquals(producerIdAndEpoch, manager.producerIdAndEpoch());
+        assertFalse(manager.hasUnresolvedSequences());
 
-        manager.maybeUpdateLastAckedSequence(tp0, 5);
-        manager.incrementSequenceNumber(tp0, 1);
+        // We have a new batch which fails with a timeout
+        ProducerBatch b2 = writeIdempotentBatchWithValue(manager, tp0, "2");
+        assertEquals(Integer.valueOf(2), manager.sequenceNumber(tp0));
         manager.markSequenceUnresolved(tp0);
-        manager.markSequenceUnresolved(tp1);
-        assertTrue(manager.shouldResetProducerStateAfterResolvingSequences());
+        manager.handleFailedBatch(b2, new TimeoutException(), false);
+        assertTrue(manager.hasUnresolvedSequences());
+
+        // We only had one inflight batch, so we should be able to clear the unresolved status
+        // and reset the producerId
+        manager.resetProducerIdIfNeeded();
+        assertFalse(manager.hasUnresolvedSequences());
+        assertFalse(manager.hasProducerId());
+    }
+
+    @Test
+    public void testNoProducerIdResetAfterLastInFlightBatchSucceeds() {
+        TransactionManager manager = new TransactionManager(logContext, null, transactionTimeoutMs,
+                DEFAULT_RETRY_BACKOFF_MS);
+        long producerId = 15L;
+        short epoch = 5;
+        ProducerIdAndEpoch producerIdAndEpoch = new ProducerIdAndEpoch(producerId, epoch);
+        manager.setProducerIdAndEpoch(producerIdAndEpoch);
+
+        TopicPartition tp0 = new TopicPartition("foo", 0);
+        ProducerBatch b1 = writeIdempotentBatchWithValue(manager, tp0, "1");
+        ProducerBatch b2 = writeIdempotentBatchWithValue(manager, tp0, "2");
+        ProducerBatch b3 = writeIdempotentBatchWithValue(manager, tp0, "3");
+        assertEquals(3, manager.sequenceNumber(tp0).intValue());
+
+        // The first batch fails with a timeout
+        manager.markSequenceUnresolved(tp0);
+        manager.handleFailedBatch(b1, new TimeoutException(), false);
+        assertTrue(manager.hasUnresolvedSequences());
+
+        // The reset should not occur until sequence numbers have been resolved
+        manager.resetProducerIdIfNeeded();
+        assertEquals(producerIdAndEpoch, manager.producerIdAndEpoch());
+        assertTrue(manager.hasUnresolvedSequences());
+
+        // The second batch fails as well with a timeout
+        manager.handleFailedBatch(b2, new TimeoutException(), false);
+        manager.resetProducerIdIfNeeded();
+        assertEquals(producerIdAndEpoch, manager.producerIdAndEpoch());
+        assertTrue(manager.hasUnresolvedSequences());
+
+        // The third batch succeeds, which should resolve the sequence number without
+        // requiring a producerId reset.
+        manager.handleCompletedBatch(b3, new ProduceResponse.PartitionResponse(
+                Errors.NONE, 500L, time.milliseconds(), 0L));
+        manager.resetProducerIdIfNeeded();
+        assertEquals(producerIdAndEpoch, manager.producerIdAndEpoch());
+        assertFalse(manager.hasUnresolvedSequences());
+        assertEquals(3, manager.sequenceNumber(tp0).intValue());
+    }
+
+    @Test
+    public void testProducerIdResetAfterLastInFlightBatchFails() {
+        TransactionManager manager = new TransactionManager(logContext, null, transactionTimeoutMs,
+                DEFAULT_RETRY_BACKOFF_MS);
+        long producerId = 15L;
+        short epoch = 5;
+        ProducerIdAndEpoch producerIdAndEpoch = new ProducerIdAndEpoch(producerId, epoch);
+        manager.setProducerIdAndEpoch(producerIdAndEpoch);
+
+        TopicPartition tp0 = new TopicPartition("foo", 0);
+        ProducerBatch b1 = writeIdempotentBatchWithValue(manager, tp0, "1");
+        ProducerBatch b2 = writeIdempotentBatchWithValue(manager, tp0, "2");
+        ProducerBatch b3 = writeIdempotentBatchWithValue(manager, tp0, "3");
+        assertEquals(Integer.valueOf(3), manager.sequenceNumber(tp0));
+
+        // The first batch fails with a timeout
+        manager.markSequenceUnresolved(tp0);
+        manager.handleFailedBatch(b1, new TimeoutException(), false);
+        assertTrue(manager.hasUnresolvedSequences());
+
+        // The second batch succeeds, but sequence numbers are still not resolved
+        manager.handleCompletedBatch(b2, new ProduceResponse.PartitionResponse(
+                Errors.NONE, 500L, time.milliseconds(), 0L));
+        manager.resetProducerIdIfNeeded();
+        assertEquals(producerIdAndEpoch, manager.producerIdAndEpoch());
+        assertTrue(manager.hasUnresolvedSequences());
+
+        // When the last inflight batch fails, we have to reset the producerId
+        manager.handleFailedBatch(b3, new TimeoutException(), false);
+        manager.resetProducerIdIfNeeded();
+        assertFalse(manager.hasProducerId());
+        assertFalse(manager.hasUnresolvedSequences());
+        assertEquals(0, manager.sequenceNumber(tp0).intValue());
     }
 
     @Test
@@ -2306,21 +2640,22 @@ private void verifyCommitOrAbortTranscationRetriable(TransactionResult firstTran
         doInitTransactions(pid, epoch);
 
         transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
 
         accumulator.append(tp0, time.milliseconds(), "key".getBytes(),
-                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT);
+                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT, false);
 
         prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, pid);
 
         prepareProduceResponse(Errors.NONE, pid, epoch);
-        sender.run(time.milliseconds());  // send addPartitions.
-        sender.run(time.milliseconds());  // send produce request.
+        sender.runOnce();  // send addPartitions.
+        sender.runOnce();  // send produce request.
 
         TransactionalRequestResult result = firstTransactionResult == TransactionResult.COMMIT ?
                 transactionManager.beginCommit() : transactionManager.beginAbort();
         prepareEndTxnResponse(Errors.NONE, firstTransactionResult, pid, epoch, true);
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertFalse(result.isCompleted());
 
         try {
@@ -2330,12 +2665,12 @@ private void verifyCommitOrAbortTranscationRetriable(TransactionResult firstTran
         }
 
         prepareFindCoordinatorResponse(Errors.NONE, false, CoordinatorType.TRANSACTION, transactionalId);
-        sender.run(time.milliseconds());
+        sender.runOnce();
         TransactionalRequestResult retryResult = retryTransactionResult == TransactionResult.COMMIT ?
                 transactionManager.beginCommit() : transactionManager.beginAbort();
         assertEquals(retryResult, result); // check if cached result is reused.
         prepareEndTxnResponse(Errors.NONE, retryTransactionResult, pid, epoch, false);
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertTrue(retryResult.isCompleted());
         assertFalse(transactionManager.hasOngoingTransaction());
     }
@@ -2347,13 +2682,14 @@ private void verifyAddPartitionsFailsWithPartitionLevelError(final Errors error)
         doInitTransactions(pid, epoch);
 
         transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
 
         Future responseFuture = accumulator.append(tp0, time.milliseconds(), "key".getBytes(),
-                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT).future;
+                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT, false).future;
         assertFalse(responseFuture.isDone());
         prepareAddPartitionsToTxn(tp0, error);
-        sender.run(time.milliseconds());  // attempt send addPartitions.
+        sender.runOnce();  // attempt send addPartitions.
         assertTrue(transactionManager.hasError());
         assertFalse(transactionManager.transactionContainsPartition(tp0));
     }
@@ -2375,27 +2711,32 @@ private void prepareFindCoordinatorResponse(Errors error, boolean shouldDisconne
                                                 final String coordinatorKey) {
         client.prepareResponse(body -> {
             FindCoordinatorRequest findCoordinatorRequest = (FindCoordinatorRequest) body;
-            assertEquals(findCoordinatorRequest.coordinatorType(), coordinatorType);
-            assertEquals(findCoordinatorRequest.coordinatorKey(), coordinatorKey);
+            assertEquals(CoordinatorType.forId(findCoordinatorRequest.data().keyType()), coordinatorType);
+            assertEquals(findCoordinatorRequest.data().key(), coordinatorKey);
             return true;
-        }, new FindCoordinatorResponse(error, brokerNode), shouldDisconnect);
+        }, FindCoordinatorResponse.prepareResponse(error, brokerNode), shouldDisconnect);
     }
 
-    private void prepareInitPidResponse(Errors error, boolean shouldDisconnect, long pid, short epoch) {
+    private void prepareInitPidResponse(Errors error, boolean shouldDisconnect, long producerId, short producerEpoch) {
+        InitProducerIdResponseData responseData = new InitProducerIdResponseData()
+                .setErrorCode(error.code())
+                .setProducerEpoch(producerEpoch)
+                .setProducerId(producerId)
+                .setThrottleTimeMs(0);
         client.prepareResponse(body -> {
             InitProducerIdRequest initProducerIdRequest = (InitProducerIdRequest) body;
-            assertEquals(initProducerIdRequest.transactionalId(), transactionalId);
-            assertEquals(initProducerIdRequest.transactionTimeoutMs(), transactionTimeoutMs);
+            assertEquals(initProducerIdRequest.data.transactionalId(), transactionalId);
+            assertEquals(initProducerIdRequest.data.transactionTimeoutMs(), transactionTimeoutMs);
             return true;
-        }, new InitProducerIdResponse(0, error, pid, epoch), shouldDisconnect);
+        }, new InitProducerIdResponse(responseData), shouldDisconnect);
     }
 
-    private void sendProduceResponse(Errors error, final long pid, final short epoch) {
-        client.respond(produceRequestMatcher(pid, epoch), produceResponse(tp0, 0, error, 0));
+    private void sendProduceResponse(Errors error, final long producerId, final short producerEpoch) {
+        client.respond(produceRequestMatcher(producerId, producerEpoch), produceResponse(tp0, 0, error, 0));
     }
 
-    private void prepareProduceResponse(Errors error, final long pid, final short epoch) {
-        client.prepareResponse(produceRequestMatcher(pid, epoch), produceResponse(tp0, 0, error, 0));
+    private void prepareProduceResponse(Errors error, final long producerId, final short producerEpoch) {
+        client.prepareResponse(produceRequestMatcher(producerId, producerEpoch), produceResponse(tp0, 0, error, 0));
     }
     private MockClient.RequestMatcher produceRequestMatcher(final long pid, final short epoch) {
         return body -> {
@@ -2485,9 +2826,9 @@ private void prepareTxnOffsetCommitResponse(final String consumerGroupId,
                                                 Map txnOffsetCommitResponse) {
         client.prepareResponse(request -> {
             TxnOffsetCommitRequest txnOffsetCommitRequest = (TxnOffsetCommitRequest) request;
-            assertEquals(consumerGroupId, txnOffsetCommitRequest.consumerGroupId());
-            assertEquals(producerId, txnOffsetCommitRequest.producerId());
-            assertEquals(producerEpoch, txnOffsetCommitRequest.producerEpoch());
+            assertEquals(consumerGroupId, txnOffsetCommitRequest.data.groupId());
+            assertEquals(producerId, txnOffsetCommitRequest.data.producerId());
+            assertEquals(producerEpoch, txnOffsetCommitRequest.data.producerEpoch());
             return true;
         }, new TxnOffsetCommitResponse(0, txnOffsetCommitResponse));
     }
@@ -2501,12 +2842,12 @@ private ProduceResponse produceResponse(TopicPartition tp, long offset, Errors e
     private void doInitTransactions(long pid, short epoch) {
         transactionManager.initializeTransactions();
         prepareFindCoordinatorResponse(Errors.NONE, false, CoordinatorType.TRANSACTION, transactionalId);
-        sender.run(time.milliseconds());  // find coordinator
-        sender.run(time.milliseconds());
+        sender.runOnce();  // find coordinator
+        sender.runOnce();
         assertEquals(brokerNode, transactionManager.coordinator(CoordinatorType.TRANSACTION));
 
         prepareInitPidResponse(Errors.NONE, false, pid, epoch);
-        sender.run(time.milliseconds());  // get pid.
+        sender.runOnce();  // get pid.
         assertTrue(transactionManager.hasProducerId());
     }
 
diff --git a/clients/src/test/java/org/apache/kafka/common/ClusterTest.java b/clients/src/test/java/org/apache/kafka/common/ClusterTest.java
index 0a7049bcc8899..2c80d08a3747f 100644
--- a/clients/src/test/java/org/apache/kafka/common/ClusterTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/ClusterTest.java
@@ -22,19 +22,35 @@
 import java.net.InetSocketAddress;
 import java.util.Arrays;
 import java.util.HashSet;
+import java.util.List;
 import java.util.Set;
 
+import static java.util.Arrays.asList;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
 
 public class ClusterTest {
 
+    private final static Node[] NODES = new Node[] {
+        new Node(0, "localhost", 99),
+        new Node(1, "localhost", 100),
+        new Node(2, "localhost", 101),
+        new Node(11, "localhost", 102)
+    };
+
+    private final static String TOPIC_A = "topicA";
+    private final static String TOPIC_B = "topicB";
+    private final static String TOPIC_C = "topicC";
+    private final static String TOPIC_D = "topicD";
+    private final static String TOPIC_E = "topicE";
+
     @Test
     public void testBootstrap() {
         String ipAddress = "140.211.11.105";
         String hostName = "www.example.com";
         Cluster cluster = Cluster.bootstrap(Arrays.asList(
-                new InetSocketAddress(ipAddress, 9002),
-                new InetSocketAddress(hostName, 9002)
+            new InetSocketAddress(ipAddress, 9002),
+            new InetSocketAddress(hostName, 9002)
         ));
         Set expectedHosts = Utils.mkSet(ipAddress, hostName);
         Set actualHosts = new HashSet<>();
@@ -43,4 +59,34 @@ public void testBootstrap() {
         assertEquals(expectedHosts, actualHosts);
     }
 
+    @Test
+    public void testReturnUnmodifiableCollections() {
+        List allPartitions = asList(new PartitionInfo(TOPIC_A, 0, NODES[0], NODES, NODES),
+            new PartitionInfo(TOPIC_A, 1, null, NODES, NODES),
+            new PartitionInfo(TOPIC_A, 2, NODES[2], NODES, NODES),
+            new PartitionInfo(TOPIC_B, 0, null, NODES, NODES),
+            new PartitionInfo(TOPIC_B, 1, NODES[0], NODES, NODES),
+            new PartitionInfo(TOPIC_C, 0, null, NODES, NODES),
+            new PartitionInfo(TOPIC_D, 0, NODES[1], NODES, NODES),
+            new PartitionInfo(TOPIC_E, 0, NODES[0], NODES, NODES)
+        );
+        Set unauthorizedTopics = Utils.mkSet(TOPIC_C);
+        Set invalidTopics = Utils.mkSet(TOPIC_D);
+        Set internalTopics = Utils.mkSet(TOPIC_E);
+        Cluster cluster = new Cluster("clusterId", asList(NODES), allPartitions, unauthorizedTopics,
+            invalidTopics, internalTopics, NODES[1]);
+
+        assertThrows(UnsupportedOperationException.class, () -> cluster.invalidTopics().add("foo"));
+        assertThrows(UnsupportedOperationException.class, () -> cluster.internalTopics().add("foo"));
+        assertThrows(UnsupportedOperationException.class, () -> cluster.unauthorizedTopics().add("foo"));
+        assertThrows(UnsupportedOperationException.class, () -> cluster.topics().add("foo"));
+        assertThrows(UnsupportedOperationException.class, () -> cluster.nodes().add(NODES[3]));
+        assertThrows(UnsupportedOperationException.class, () -> cluster.partitionsForTopic(TOPIC_A).add(
+            new PartitionInfo(TOPIC_A, 3, NODES[0], NODES, NODES)));
+        assertThrows(UnsupportedOperationException.class, () -> cluster.availablePartitionsForTopic(TOPIC_B).add(
+            new PartitionInfo(TOPIC_B, 2, NODES[0], NODES, NODES)));
+        assertThrows(UnsupportedOperationException.class, () -> cluster.partitionsForNode(NODES[1].id()).add(
+            new PartitionInfo(TOPIC_B, 2, NODES[1], NODES, NODES)));
+    }
+
 }
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 071deed47e422..834278b18d5f4 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
@@ -24,6 +24,8 @@
 import org.apache.kafka.common.metrics.JmxReporter;
 import org.apache.kafka.common.metrics.MetricsReporter;
 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.junit.Test;
 
 import java.util.Arrays;
@@ -32,6 +34,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Properties;
+import java.util.UUID;
 
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNull;
@@ -263,7 +266,7 @@ public RestrictedClassLoader() {
             @Override
             protected Class findClass(String name) throws ClassNotFoundException {
                 if (name.equals(ClassTestConfig.DEFAULT_CLASS.getName()) || name.equals(ClassTestConfig.RESTRICTED_CLASS.getName()))
-                    return null;
+                    throw new ClassNotFoundException();
                 else
                     return ClassTestConfig.class.getClassLoader().loadClass(name);
             }
@@ -321,6 +324,188 @@ protected Class findClass(String name) throws ClassNotFoundException {
         }
     }
 
+    @SuppressWarnings("unchecked")
+    public Map 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) props;
+    }
+
+    @Test
+    public void testOriginalsWithConfigProvidersProps() {
+        Properties props = new Properties();
+
+        // Test Case: Valid Test Case for ConfigProviders as part of config.properties
+        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);
+        assertEquals(config.originals().get("sasl.kerberos.key"), "testKey");
+        assertEquals(config.originals().get("sasl.kerberos.password"), "randomPassword");
+        assertEquals(config.originals().get("prefix.ssl.truststore.location.number"), 5);
+        assertEquals(config.originals().get("sasl.kerberos.service.name"), "service name");
+        MockFileConfigProvider.assertClosed(id);
+    }
+
+    @Test
+    public void testConfigProvidersPropsAsParam() {
+        // Test Case: Valid Test Case for ConfigProviders as a separate variable
+        Properties providers = new Properties();
+        providers.put("config.providers", "file");
+        providers.put("config.providers.file.class", MockFileConfigProvider.class.getName());
+        String id = UUID.randomUUID().toString();
+        providers.put("config.providers.file.param.testId", id);
+        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));
+        assertEquals(config.originals().get("sasl.kerberos.key"), "testKey");
+        assertEquals(config.originals().get("sasl.kerberos.password"), "randomPassword");
+        MockFileConfigProvider.assertClosed(id);
+    }
+
+    @Test
+    public void testImmutableOriginalsWithConfigProvidersProps() {
+        // Test Case: Valid Test Case for ConfigProviders as a separate variable
+        Properties providers = new Properties();
+        providers.put("config.providers", "file");
+        providers.put("config.providers.file.class", MockFileConfigProvider.class.getName());
+        String id = UUID.randomUUID().toString();
+        providers.put("config.providers.file.param.testId", id);
+        Properties props = new Properties();
+        props.put("sasl.kerberos.key", "${file:/usr/kerberos:key}");
+        Map immutableMap = Collections.unmodifiableMap(props);
+        Map provMap = convertPropertiesToMap(providers);
+        TestIndirectConfigResolution config = new TestIndirectConfigResolution(immutableMap, provMap);
+        assertEquals(config.originals().get("sasl.kerberos.key"), "testKey");
+        MockFileConfigProvider.assertClosed(id);
+    }
+
+    @Test
+    public void testAutoConfigResolutionWithMultipleConfigProviders() {
+        // Test Case: Valid Test Case With Multiple ConfigProviders as a separate variable
+        Properties providers = new Properties();
+        providers.put("config.providers", "file,vault");
+        providers.put("config.providers.file.class", MockFileConfigProvider.class.getName());
+        String id = UUID.randomUUID().toString();
+        providers.put("config.providers.file.param.testId", id);
+        providers.put("config.providers.vault.class", MockVaultConfigProvider.class.getName());
+        Properties props = new Properties();
+        props.put("sasl.kerberos.key", "${file:/usr/kerberos:key}");
+        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));
+        assertEquals(config.originals().get("sasl.kerberos.key"), "testKey");
+        assertEquals(config.originals().get("sasl.kerberos.password"), "randomPassword");
+        assertEquals(config.originals().get("sasl.truststore.key"), "testTruststoreKey");
+        assertEquals(config.originals().get("sasl.truststore.password"), "randomtruststorePassword");
+        MockFileConfigProvider.assertClosed(id);
+    }
+
+    @Test
+    public void testAutoConfigResolutionWithInvalidConfigProviderClass() {
+        // Test Case: Invalid class for Config Provider
+        Properties props = new Properties();
+        props.put("config.providers", "file");
+        props.put("config.providers.file.class",
+            "org.apache.kafka.common.config.provider.InvalidConfigProvider");
+        props.put("testKey", "${test:/foo/bar/testpath:testKey}");
+        try {
+            TestIndirectConfigResolution config = new TestIndirectConfigResolution(props);
+            fail("Expected a config exception due to invalid props :" + props);
+        } catch (KafkaException e) {
+            // this is good
+        }
+    }
+
+    @Test
+    public void testAutoConfigResolutionWithMissingConfigProvider() {
+        // Test Case: Config Provider for a variable missing in config file.
+        Properties props = new Properties();
+        props.put("testKey", "${test:/foo/bar/testpath:testKey}");
+        TestIndirectConfigResolution config = new TestIndirectConfigResolution(props);
+        assertEquals(config.originals().get("testKey"), "${test:/foo/bar/testpath:testKey}");
+    }
+
+    @Test
+    public void testAutoConfigResolutionWithMissingConfigKey() {
+        // Test Case: Config Provider fails to resolve the config (key not present)
+        Properties props = new Properties();
+        props.put("config.providers", "test");
+        props.put("config.providers.test.class", MockFileConfigProvider.class.getName());
+        String id = UUID.randomUUID().toString();
+        props.put("config.providers.test.param.testId", id);
+        props.put("random", "${test:/foo/bar/testpath:random}");
+        TestIndirectConfigResolution config = new TestIndirectConfigResolution(props);
+        assertEquals(config.originals().get("random"), "${test:/foo/bar/testpath:random}");
+        MockFileConfigProvider.assertClosed(id);
+    }
+
+    @Test
+    public void testAutoConfigResolutionWithDuplicateConfigProvider() {
+        // Test Case: If ConfigProvider is provided in both originals and provider. Only the ones in provider should be used.
+        Properties providers = new Properties();
+        providers.put("config.providers", "test");
+        providers.put("config.providers.test.class", MockVaultConfigProvider.class.getName());
+
+        Properties props = new Properties();
+        props.put("sasl.kerberos.key", "${file:/usr/kerberos:key}");
+        props.put("config.providers", "file");
+        props.put("config.providers.file.class", MockVaultConfigProvider.class.getName());
+
+        TestIndirectConfigResolution config = new TestIndirectConfigResolution(props, convertPropertiesToMap(providers));
+        assertEquals(config.originals().get("sasl.kerberos.key"), "${file:/usr/kerberos:key}");
+    }
+
+    @Test
+    public void testConfigProviderConfigurationWithConfigParams() {
+        // Test Case: Valid Test Case With Multiple ConfigProviders as a separate variable
+        Properties providers = new Properties();
+        providers.put("config.providers", "vault");
+        providers.put("config.providers.vault.class", MockVaultConfigProvider.class.getName());
+        providers.put("config.providers.vault.param.key", "randomKey");
+        providers.put("config.providers.vault.param.location", "/usr/vault");
+        Properties props = new Properties();
+        props.put("sasl.truststore.location", "${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));
+        assertEquals(config.originals().get("sasl.truststore.location"), "/usr/vault");
+    }
+
+    private static class TestIndirectConfigResolution extends AbstractConfig {
+
+        private static final ConfigDef CONFIG;
+
+        public static final String INDIRECT_CONFIGS = "indirect.variables";
+        private static final String INDIRECT_CONFIGS_DOC = "Variables whose values can be obtained from ConfigProviders";
+
+        static {
+            CONFIG = new ConfigDef().define(INDIRECT_CONFIGS,
+                    Type.LIST,
+                    "",
+                    Importance.LOW,
+                    INDIRECT_CONFIGS_DOC);
+        }
+
+        public TestIndirectConfigResolution(Map props) {
+            super(CONFIG, props, true);
+        }
+
+        public TestIndirectConfigResolution(Map props, Map providers) {
+            super(CONFIG, props, providers, true);
+        }
+    }
+
     private static class ClassTestConfig extends AbstractConfig {
         static final Class DEFAULT_CLASS = FakeMetricsReporter.class;
         static final Class VISIBLE_CLASS = JmxReporter.class;
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 974d39f934d24..9751545ca4de1 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
@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.common.config;
 
+import org.apache.kafka.common.config.ConfigDef.CaseInsensitiveValidString;
 import org.apache.kafka.common.config.ConfigDef.Importance;
 import org.apache.kafka.common.config.ConfigDef.Range;
 import org.apache.kafka.common.config.ConfigDef.Type;
@@ -157,7 +158,9 @@ public void testNestedClass() {
     public void testValidators() {
         testValidators(Type.INT, Range.between(0, 10), 5, new Object[]{1, 5, 9}, new Object[]{-1, 11, null});
         testValidators(Type.STRING, ValidString.in("good", "values", "default"), "default",
-                new Object[]{"good", "values", "default"}, new Object[]{"bad", "inputs", null});
+                new Object[]{"good", "values", "default"}, new Object[]{"bad", "inputs", "DEFAULT", null});
+        testValidators(Type.STRING, CaseInsensitiveValidString.in("good", "values", "default"), "default",
+            new Object[]{"gOOd", "VALUES", "default"}, new Object[]{"Bad", "iNPUts", null});
         testValidators(Type.LIST, ConfigDef.ValidList.in("1", "2", "3"), "1", new Object[]{"1", "2", "3"}, new Object[]{"4", "5", "6"});
         testValidators(Type.STRING, new ConfigDef.NonNullValidator(), "a", new Object[]{"abb"}, new Object[] {null});
         testValidators(Type.STRING, ConfigDef.CompositeValidator.of(new ConfigDef.NonNullValidator(), ValidString.in("a", "b")), "a", new Object[]{"a", "b"}, new Object[] {null, -1, "c"});
@@ -639,6 +642,29 @@ public void testConvertValueToStringNestedClass() throws ClassNotFoundException
         assertEquals(NestedClass.class, Class.forName(actual));
     }
 
+    @Test
+    public void testClassWithAlias() {
+        final String alias = "PluginAlias";
+        ClassLoader originalClassLoader = Thread.currentThread().getContextClassLoader();
+        try {
+            // Could try to use the Plugins class from Connect here, but this should simulate enough
+            // of the aliasing logic to suffice for this test.
+            Thread.currentThread().setContextClassLoader(new ClassLoader(originalClassLoader) {
+                @Override
+                public Class loadClass(String name, boolean resolve) throws ClassNotFoundException {
+                    if (alias.equals(name)) {
+                        return NestedClass.class;
+                    } else {
+                        return super.loadClass(name, resolve);
+                    }
+                }
+            });
+            ConfigDef.parseType("Test config", alias, Type.CLASS);
+        } finally {
+            Thread.currentThread().setContextClassLoader(originalClassLoader);
+        }
+    }
+
     private class NestedClass {
     }
 }
diff --git a/clients/src/test/java/org/apache/kafka/common/config/ConfigTransformerTest.java b/clients/src/test/java/org/apache/kafka/common/config/ConfigTransformerTest.java
index e2b9f6b001cee..12c6b1f4a2785 100644
--- a/clients/src/test/java/org/apache/kafka/common/config/ConfigTransformerTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/config/ConfigTransformerTest.java
@@ -95,6 +95,13 @@ public void testReplaceVariableNoPath() throws Exception {
         assertTrue(ttls.isEmpty());
     }
 
+    @Test
+    public void testReplaceMultipleVariablesWithoutPathInValue() throws Exception {
+        ConfigTransformerResult result = configTransformer.transform(Collections.singletonMap(MY_KEY, "first ${test:testKey}; second ${test:testKey}"));
+        Map data = result.data();
+        assertEquals("first testResultNoPath; second testResultNoPath", data.get(MY_KEY));
+    }
+
     @Test
     public void testNullConfigValue() throws Exception {
         ConfigTransformerResult result = configTransformer.transform(Collections.singletonMap(MY_KEY, null));
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
new file mode 100644
index 0000000000000..3409096446895
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/config/provider/MockFileConfigProvider.java
@@ -0,0 +1,64 @@
+/*
+ * 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 java.io.IOException;
+import java.io.Reader;
+import java.io.StringReader;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
+
+public class MockFileConfigProvider extends FileConfigProvider {
+
+    private static final Map INSTANCES = Collections.synchronizedMap(new HashMap<>());
+    private String id;
+    private boolean closed = false;
+
+    public void configure(Map configs) {
+        Object id = configs.get("testId");
+        if (id == null) {
+            throw new RuntimeException(getClass().getName() + " missing 'testId' config");
+        }
+        if (this.id != null) {
+            throw new RuntimeException(getClass().getName() + " instance was configured twice");
+        }
+        this.id = id.toString();
+        INSTANCES.put(id.toString(), this);
+    }
+
+    @Override
+    protected Reader reader(String path) throws IOException {
+        return new StringReader("key=testKey\npassword=randomPassword");
+    }
+
+    @Override
+    public synchronized void close() {
+        closed = true;
+    }
+
+    public static void assertClosed(String id) {
+        MockFileConfigProvider instance = INSTANCES.remove(id);
+        assertNotNull(instance);
+        synchronized (instance) {
+            assertTrue(instance.closed);
+        }
+    }
+}
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
new file mode 100644
index 0000000000000..c741798a72972
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/config/provider/MockVaultConfigProvider.java
@@ -0,0 +1,45 @@
+/*
+ * 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 java.io.IOException;
+import java.io.Reader;
+import java.io.StringReader;
+import java.util.Map;
+
+public class MockVaultConfigProvider extends FileConfigProvider {
+
+    Map vaultConfigs;
+    private boolean configured = false;
+    private static final String LOCATION = "location";
+
+    @Override
+    protected Reader reader(String 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 configs) {
+        this.vaultConfigs = configs;
+        configured = true;
+    }
+
+    public boolean configured() {
+        return configured;
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/header/internals/RecordHeadersTest.java b/clients/src/test/java/org/apache/kafka/common/header/internals/RecordHeadersTest.java
index 39c1c9c9a7bce..5b9f95ea91f18 100644
--- a/clients/src/test/java/org/apache/kafka/common/header/internals/RecordHeadersTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/header/internals/RecordHeadersTest.java
@@ -16,19 +16,19 @@
  */
 package org.apache.kafka.common.header.internals;
 
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertNull;
-import static org.junit.Assert.assertTrue;
-import static org.junit.Assert.fail;
+import org.apache.kafka.common.header.Header;
+import org.apache.kafka.common.header.Headers;
+import org.junit.Test;
 
 import java.io.IOException;
 import java.util.Arrays;
 import java.util.Iterator;
 
-import org.apache.kafka.common.header.Header;
-import org.apache.kafka.common.header.Headers;
-import org.junit.Test;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
 
 public class RecordHeadersTest {
 
@@ -206,6 +206,11 @@ public void testNew() throws IOException {
         assertEquals(2, getCount(newHeaders));
     }
 
+    @Test(expected = NullPointerException.class)
+    public void shouldThrowNpeWhenAddingNullHeader() {
+        new RecordHeaders().add(null);
+    }
+
     private int getCount(Headers headers) {
         int count = 0;
         Iterator
       headerIterator = headers.iterator();
diff --git a/clients/src/test/java/org/apache/kafka/common/memory/RecyclingMemoryPoolTest.java b/clients/src/test/java/org/apache/kafka/common/memory/RecyclingMemoryPoolTest.java
new file mode 100644
index 0000000000000..eec9e7e41668c
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/memory/RecyclingMemoryPoolTest.java
@@ -0,0 +1,110 @@
+/*
+ * 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.memory;
+
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.atomic.AtomicReference;
+import org.apache.kafka.common.metrics.Metrics;
+import org.apache.kafka.common.metrics.Sensor;
+import org.junit.Assert;
+import org.junit.Test;
+
+
+public class RecyclingMemoryPoolTest {
+    private static final int TWO_KILOBYTES = 2048;
+    private static final int CACHEABLE_BUFFER_SIZE = 1024;
+    private static final int BUFFER_CACHE_CAPACITY = 2;
+    private static final Sensor ALLOCATE_SENSOR = new Metrics().sensor("allocate_sensor");
+
+    @Test(expected = IllegalArgumentException.class)
+    public void testNegativeAllocation() {
+        RecyclingMemoryPool memoryPool = new RecyclingMemoryPool(CACHEABLE_BUFFER_SIZE, BUFFER_CACHE_CAPACITY, ALLOCATE_SENSOR);
+        memoryPool.tryAllocate(-1);
+    }
+
+    @Test(expected = IllegalArgumentException.class)
+    public void testZeroAllocation() {
+        RecyclingMemoryPool memoryPool = new RecyclingMemoryPool(CACHEABLE_BUFFER_SIZE, BUFFER_CACHE_CAPACITY, ALLOCATE_SENSOR);
+        memoryPool.tryAllocate(0);
+    }
+
+    @Test(expected = IllegalArgumentException.class)
+    public void testNullRelease() {
+        RecyclingMemoryPool memoryPool = new RecyclingMemoryPool(CACHEABLE_BUFFER_SIZE, BUFFER_CACHE_CAPACITY, ALLOCATE_SENSOR);
+        memoryPool.release(null);
+    }
+
+    @Test
+    public void testAllocation() {
+        RecyclingMemoryPool memoryPool = new RecyclingMemoryPool(CACHEABLE_BUFFER_SIZE, BUFFER_CACHE_CAPACITY, ALLOCATE_SENSOR);
+        ByteBuffer buffer1 = memoryPool.tryAllocate(TWO_KILOBYTES);
+        ByteBuffer buffer2 = memoryPool.tryAllocate(CACHEABLE_BUFFER_SIZE);
+        ByteBuffer buffer3 = memoryPool.tryAllocate(CACHEABLE_BUFFER_SIZE * 2 / 3);
+        ByteBuffer buffer4 = memoryPool.tryAllocate(CACHEABLE_BUFFER_SIZE);
+
+        memoryPool.release(buffer1);
+        ByteBuffer reuse1 = memoryPool.tryAllocate(TWO_KILOBYTES);
+        // Compare the references
+        Assert.assertNotEquals(System.identityHashCode(reuse1), System.identityHashCode(buffer1));
+
+        memoryPool.release(buffer2);
+        memoryPool.release(buffer3);
+        memoryPool.release(buffer4);
+        ByteBuffer reuse2 = memoryPool.tryAllocate(CACHEABLE_BUFFER_SIZE);
+        ByteBuffer reuse3 = memoryPool.tryAllocate(CACHEABLE_BUFFER_SIZE * 2 / 3);
+        ByteBuffer reuse4 = memoryPool.tryAllocate(CACHEABLE_BUFFER_SIZE);
+
+        Assert.assertEquals(System.identityHashCode(reuse2), System.identityHashCode(buffer2));
+        Assert.assertEquals(System.identityHashCode(reuse3), System.identityHashCode(buffer3));
+        Assert.assertNotEquals(System.identityHashCode(reuse4), System.identityHashCode(buffer4));
+    }
+
+    @Test
+    public void testMultiThreadAllocation() {
+        RecyclingMemoryPool memoryPool = new RecyclingMemoryPool(CACHEABLE_BUFFER_SIZE, BUFFER_CACHE_CAPACITY, ALLOCATE_SENSOR);
+        AtomicReference error = new AtomicReference<>();
+        List processorThreads = new ArrayList<>(3);
+        for (int i = 0; i < 3; i++) {
+            processorThreads.add(new Thread(() -> {
+                try {
+                    ByteBuffer buffer = memoryPool.tryAllocate(CACHEABLE_BUFFER_SIZE);
+                    Thread.sleep(1000);
+                    memoryPool.release(buffer);
+                } catch (InterruptedException e) {
+                    error.compareAndSet(null, e);
+                }
+            }));
+        }
+        processorThreads.forEach(t -> {
+            t.setDaemon(true);
+            t.start();
+        });
+        processorThreads.forEach(t -> {
+            try {
+                t.join(30000);
+            } catch (InterruptedException e) {
+                error.compareAndSet(null, e);
+            }
+        });
+
+        Assert.assertNull(error.get());
+        Assert.assertEquals(memoryPool.bufferCache.size(), BUFFER_CACHE_CAPACITY);
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/memory/WeakMemoryPoolTest.java b/clients/src/test/java/org/apache/kafka/common/memory/WeakMemoryPoolTest.java
new file mode 100644
index 0000000000000..e1f89ea8f024e
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/memory/WeakMemoryPoolTest.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.common.memory;
+
+import java.nio.ByteBuffer;
+import org.junit.Assert;
+import org.junit.Ignore;
+import org.junit.Test;
+
+
+public class WeakMemoryPoolTest {
+    private static final int FORTY_MEGABYTES = 40 * 1024 * 1024;
+
+    @Test(expected = IllegalArgumentException.class)
+    public void testNegativeAllocation() {
+        WeakMemoryPool memoryPool = new WeakMemoryPool();
+        memoryPool.tryAllocate(-1);
+    }
+
+    @Test
+    public void testZeroAllocation() {
+        WeakMemoryPool memoryPool = new WeakMemoryPool();
+        memoryPool.tryAllocate(0);
+    }
+
+    @Test(expected = IllegalArgumentException.class)
+    public void testNullRelease() {
+        WeakMemoryPool memoryPool = new WeakMemoryPool();
+        memoryPool.release(null);
+    }
+
+    @Ignore("flaky test")
+    @Test
+    public void testAllocationMemorySize() {
+        WeakMemoryPool pool = new WeakMemoryPool();
+        long freeMemory = Runtime.getRuntime().freeMemory();
+        ByteBuffer buffer1 = pool.tryAllocate(FORTY_MEGABYTES + 1);
+        ByteBuffer buffer2 = pool.tryAllocate(FORTY_MEGABYTES + 2);
+        ByteBuffer buffer3 = pool.tryAllocate(FORTY_MEGABYTES + 3);
+        Assert.assertTrue(Runtime.getRuntime().freeMemory()
+            <= freeMemory - buffer1.capacity() - buffer2.capacity() - buffer3.capacity());
+
+        pool.release(buffer1);
+        ByteBuffer reuse1 = pool.tryAllocate(FORTY_MEGABYTES);
+        // Compare the references
+        Assert.assertTrue(reuse1 == buffer1);
+
+        pool.release(buffer2);
+        pool.release(buffer3);
+        ByteBuffer reuse3 = pool.tryAllocate(FORTY_MEGABYTES + 3);
+        ByteBuffer reuse2 = pool.tryAllocate(FORTY_MEGABYTES + 2);
+
+        Assert.assertTrue(reuse3 == buffer3);
+        Assert.assertTrue(reuse2 == buffer2);
+    }
+
+    @Test
+    public void testAllocation() {
+        WeakMemoryPool pool = new WeakMemoryPool();
+        ByteBuffer buffer1 = pool.tryAllocate(FORTY_MEGABYTES + 1);
+        ByteBuffer buffer2 = pool.tryAllocate(FORTY_MEGABYTES + 2);
+        ByteBuffer buffer3 = pool.tryAllocate(FORTY_MEGABYTES + 3);
+
+        pool.release(buffer1);
+        ByteBuffer reuse1 = pool.tryAllocate(FORTY_MEGABYTES);
+        // Compare the references
+        Assert.assertEquals(System.identityHashCode(reuse1), System.identityHashCode(buffer1));
+
+        pool.release(buffer2);
+        pool.release(buffer3);
+        ByteBuffer reuse3 = pool.tryAllocate(FORTY_MEGABYTES + 3);
+        ByteBuffer reuse2 = pool.tryAllocate(FORTY_MEGABYTES + 2);
+
+        Assert.assertEquals(System.identityHashCode(reuse3), System.identityHashCode(buffer3));
+        Assert.assertEquals(System.identityHashCode(reuse2), System.identityHashCode(buffer2));
+    }
+
+    @Test
+    public void testAllocationGC() {
+        // Clean all garbage before we begin!
+        System.gc();
+
+        WeakMemoryPool pool = new WeakMemoryPool();
+
+        ByteBuffer buffer1 = ByteBuffer.allocate(FORTY_MEGABYTES + 1);
+        ByteBuffer buffer2 = ByteBuffer.allocate(FORTY_MEGABYTES + 5);
+        ByteBuffer buffer3 = ByteBuffer.allocate(FORTY_MEGABYTES + 9);
+
+        // The byte buffers are not reachable from gc-roots
+        int identifier1 = System.identityHashCode(buffer1);
+        int identifier2 = System.identityHashCode(buffer2);
+        int identifier3 = System.identityHashCode(buffer3);
+
+        pool.release(buffer1);
+        pool.release(buffer2);
+        pool.release(buffer3);
+
+        Assert.assertEquals(identifier2, System.identityHashCode(pool.tryAllocate(FORTY_MEGABYTES + 3)));
+        Assert.assertEquals(identifier3, System.identityHashCode(pool.tryAllocate(FORTY_MEGABYTES + 7)));
+        Assert.assertEquals(identifier1, System.identityHashCode(pool.tryAllocate(FORTY_MEGABYTES + 0)));
+
+        pool.release(buffer1);
+        pool.release(buffer2);
+        pool.release(buffer3);
+
+        buffer1 = null;
+        buffer2 = null;
+        buffer3 = null;
+
+        // Reclaim all the objects
+        System.gc();
+
+        // Assert that the object is a different one!
+        Assert.assertNotEquals(identifier2, System.identityHashCode(pool.tryAllocate(FORTY_MEGABYTES + 3)));
+        Assert.assertNotEquals(identifier3, System.identityHashCode(pool.tryAllocate(FORTY_MEGABYTES + 7)));
+        Assert.assertNotEquals(identifier1, System.identityHashCode(pool.tryAllocate(FORTY_MEGABYTES + 0)));
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/message/ApiMessageTypeTest.java b/clients/src/test/java/org/apache/kafka/common/message/ApiMessageTypeTest.java
new file mode 100644
index 0000000000000..8fdb6a090ad96
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/message/ApiMessageTypeTest.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 org.apache.kafka.common.message;
+
+import org.apache.kafka.common.errors.UnsupportedVersionException;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.Timeout;
+
+import java.util.HashSet;
+import java.util.Set;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.fail;
+
+public class ApiMessageTypeTest {
+    @Rule
+    final public Timeout globalTimeout = Timeout.millis(120000);
+
+    @Test
+    public void testFromApiKey() {
+        for (ApiMessageType type : ApiMessageType.values()) {
+            ApiMessageType type2 = ApiMessageType.fromApiKey(type.apiKey());
+            assertEquals(type2, type);
+        }
+    }
+
+    @Test
+    public void testInvalidFromApiKey() {
+        try {
+            ApiMessageType.fromApiKey((short) -1);
+            fail("expected to get an UnsupportedVersionException");
+        } catch (UnsupportedVersionException uve) {
+            // expected
+        }
+    }
+
+    @Test
+    public void testUniqueness() {
+        Set ids = new HashSet<>();
+        Set requestNames = new HashSet<>();
+        Set responseNames = new HashSet<>();
+        for (ApiMessageType type : ApiMessageType.values()) {
+            assertFalse("found two ApiMessageType objects with id " + type.apiKey(),
+                ids.contains(type.apiKey()));
+            ids.add(type.apiKey());
+            String requestName = type.newRequest().getClass().getSimpleName();
+            assertFalse("found two ApiMessageType objects with requestName " + requestName,
+                requestNames.contains(requestName));
+            requestNames.add(requestName);
+            String responseName = type.newResponse().getClass().getSimpleName();
+            assertFalse("found two ApiMessageType objects with responseName " + responseName,
+                responseNames.contains(responseName));
+            responseNames.add(responseName);
+        }
+        assertEquals(ApiMessageType.values().length, ids.size());
+        assertEquals(ApiMessageType.values().length, requestNames.size());
+        assertEquals(ApiMessageType.values().length, responseNames.size());
+    }
+
+    @Test
+    public void testHeaderVersion() {
+        assertEquals((short) 1, ApiMessageType.PRODUCE.requestHeaderVersion((short) 0));
+        assertEquals((short) 0, ApiMessageType.PRODUCE.responseHeaderVersion((short) 0));
+
+        assertEquals((short) 1, ApiMessageType.PRODUCE.requestHeaderVersion((short) 1));
+        assertEquals((short) 0, ApiMessageType.PRODUCE.responseHeaderVersion((short) 1));
+
+        assertEquals((short) 0, ApiMessageType.CONTROLLED_SHUTDOWN.requestHeaderVersion((short) 0));
+        assertEquals((short) 0, ApiMessageType.CONTROLLED_SHUTDOWN.responseHeaderVersion((short) 0));
+
+        assertEquals((short) 1, ApiMessageType.CONTROLLED_SHUTDOWN.requestHeaderVersion((short) 1));
+        assertEquals((short) 0, ApiMessageType.CONTROLLED_SHUTDOWN.responseHeaderVersion((short) 1));
+
+        assertEquals((short) 1, ApiMessageType.CREATE_TOPICS.requestHeaderVersion((short) 4));
+        assertEquals((short) 0, ApiMessageType.CREATE_TOPICS.responseHeaderVersion((short) 4));
+
+        assertEquals((short) 2, ApiMessageType.CREATE_TOPICS.requestHeaderVersion((short) 5));
+        assertEquals((short) 1, ApiMessageType.CREATE_TOPICS.responseHeaderVersion((short) 5));
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/message/MessageTest.java b/clients/src/test/java/org/apache/kafka/common/message/MessageTest.java
index b725e70fb48ae..d5662ada18e07 100644
--- a/clients/src/test/java/org/apache/kafka/common/message/MessageTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/message/MessageTest.java
@@ -18,106 +18,627 @@
 package org.apache.kafka.common.message;
 
 import org.apache.kafka.common.errors.UnsupportedVersionException;
-import org.apache.kafka.common.message.CreateTopicsRequestData.CreatableTopicSet;
+import org.apache.kafka.common.message.AddPartitionsToTxnRequestData.AddPartitionsToTxnTopic;
+import org.apache.kafka.common.message.AddPartitionsToTxnRequestData.AddPartitionsToTxnTopicCollection;
+import org.apache.kafka.common.message.JoinGroupResponseData.JoinGroupResponseMember;
+import org.apache.kafka.common.message.LeaveGroupResponseData.MemberResponse;
+import org.apache.kafka.common.message.OffsetCommitRequestData.OffsetCommitRequestPartition;
+import org.apache.kafka.common.message.OffsetCommitRequestData.OffsetCommitRequestTopic;
+import org.apache.kafka.common.message.OffsetCommitResponseData.OffsetCommitResponsePartition;
+import org.apache.kafka.common.message.OffsetCommitResponseData.OffsetCommitResponseTopic;
+import org.apache.kafka.common.message.OffsetFetchRequestData.OffsetFetchRequestTopic;
+import org.apache.kafka.common.message.OffsetFetchResponseData.OffsetFetchResponsePartition;
+import org.apache.kafka.common.message.OffsetFetchResponseData.OffsetFetchResponseTopic;
+import org.apache.kafka.common.message.TxnOffsetCommitRequestData.TxnOffsetCommitRequestPartition;
+import org.apache.kafka.common.message.TxnOffsetCommitRequestData.TxnOffsetCommitRequestTopic;
+import org.apache.kafka.common.message.TxnOffsetCommitResponseData.TxnOffsetCommitResponsePartition;
+import org.apache.kafka.common.message.TxnOffsetCommitResponseData.TxnOffsetCommitResponseTopic;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.ByteBufferAccessor;
+import org.apache.kafka.common.protocol.Errors;
 import org.apache.kafka.common.protocol.Message;
-import org.apache.kafka.common.protocol.types.ArrayOf;
+import org.apache.kafka.common.protocol.ObjectSerializationCache;
 import org.apache.kafka.common.protocol.types.BoundField;
+import org.apache.kafka.common.protocol.types.RawTaggedField;
 import org.apache.kafka.common.protocol.types.Schema;
+import org.apache.kafka.common.protocol.types.SchemaException;
 import org.apache.kafka.common.protocol.types.Struct;
-
-import java.nio.ByteBuffer;
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.Collections;
-import java.util.List;
-
 import org.apache.kafka.common.protocol.types.Type;
 import org.apache.kafka.common.utils.Utils;
-import org.apache.kafka.common.message.AddPartitionsToTxnRequestData.AddPartitionsToTxnTopic;
-import org.apache.kafka.common.message.AddPartitionsToTxnRequestData.AddPartitionsToTxnTopicSet;
 import org.junit.Assert;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.Timeout;
 
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.function.Supplier;
+
+import static java.util.Collections.singletonList;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
 public final class MessageTest {
+
+    private final String memberId = "memberId";
+    private final String instanceId = "instanceId";
+
     @Rule
     final public Timeout globalTimeout = Timeout.millis(120000);
 
+    @Test
+    public void testAddOffsetsToTxnVersions() throws Exception {
+        testAllMessageRoundTrips(new AddOffsetsToTxnRequestData().
+                setTransactionalId("foobar").
+                setProducerId(0xbadcafebadcafeL).
+                setProducerEpoch((short) 123).
+                setGroupId("baaz"));
+        testAllMessageRoundTrips(new AddOffsetsToTxnResponseData().
+                setThrottleTimeMs(42).
+                setErrorCode((short) 0));
+    }
+
+    @Test
+    public void testAddPartitionsToTxnVersions() throws Exception {
+        testAllMessageRoundTrips(new AddPartitionsToTxnRequestData().
+                setTransactionalId("blah").
+                setProducerId(0xbadcafebadcafeL).
+                setProducerEpoch((short) 30000).
+                setTopics(new AddPartitionsToTxnTopicCollection(singletonList(
+                        new AddPartitionsToTxnTopic().
+                                setName("Topic").
+                                setPartitions(singletonList(1))).iterator())));
+    }
+
+    @Test
+    public void testCreateTopicsVersions() throws Exception {
+        testAllMessageRoundTrips(new CreateTopicsRequestData().
+                setTimeoutMs(1000).setTopics(new CreateTopicsRequestData.CreatableTopicCollection()));
+    }
+
+    @Test
+    public void testDescribeAclsRequest() throws Exception {
+        testAllMessageRoundTrips(new DescribeAclsRequestData().
+                setResourceType((byte) 42).
+                setResourceNameFilter(null).
+                setResourcePatternType((byte) 3).
+                setPrincipalFilter("abc").
+                setHostFilter(null).
+                setOperation((byte) 0).
+                setPermissionType((byte) 0));
+    }
+
+    @Test
+    public void testMetadataVersions() throws Exception {
+        testAllMessageRoundTrips(new MetadataRequestData().setTopics(
+                Arrays.asList(new MetadataRequestData.MetadataRequestTopic().setName("foo"),
+                        new MetadataRequestData.MetadataRequestTopic().setName("bar")
+                )));
+        testAllMessageRoundTripsFromVersion((short) 1, new MetadataRequestData().
+                setTopics(null).
+                setAllowAutoTopicCreation(true).
+                setIncludeClusterAuthorizedOperations(false).
+                setIncludeTopicAuthorizedOperations(false));
+        testAllMessageRoundTripsFromVersion((short) 4, new MetadataRequestData().
+                setTopics(null).
+                setAllowAutoTopicCreation(false).
+                setIncludeClusterAuthorizedOperations(false).
+                setIncludeTopicAuthorizedOperations(false));
+    }
+
+    @Test
+    public void testHeartbeatVersions() throws Exception {
+        Supplier newRequest = () -> new HeartbeatRequestData()
+                .setGroupId("groupId")
+                .setMemberId(memberId)
+                .setGenerationId(15);
+        testAllMessageRoundTrips(newRequest.get());
+        testAllMessageRoundTrips(newRequest.get().setGroupInstanceId(null));
+        testAllMessageRoundTripsFromVersion((short) 3, newRequest.get().setGroupInstanceId("instanceId"));
+    }
+
+    @Test
+    public void testJoinGroupRequestVersions() throws Exception {
+        Supplier newRequest = () -> new JoinGroupRequestData()
+                .setGroupId("groupId")
+                .setMemberId(memberId)
+                .setProtocolType("consumer")
+                .setProtocols(new JoinGroupRequestData.JoinGroupRequestProtocolCollection())
+                .setSessionTimeoutMs(10000);
+        testAllMessageRoundTrips(newRequest.get());
+        testAllMessageRoundTripsFromVersion((short) 1, newRequest.get().setRebalanceTimeoutMs(20000));
+        testAllMessageRoundTrips(newRequest.get().setGroupInstanceId(null));
+        testAllMessageRoundTripsFromVersion((short) 5, newRequest.get().setGroupInstanceId("instanceId"));
+    }
+
+    @Test
+    public void testJoinGroupResponseVersions() throws Exception {
+        Supplier newResponse = () -> new JoinGroupResponseData()
+                .setMemberId(memberId)
+                .setLeader(memberId)
+                .setGenerationId(1)
+                .setMembers(Collections.singletonList(
+                        new JoinGroupResponseMember()
+                                .setMemberId(memberId)
+                ));
+        testAllMessageRoundTrips(newResponse.get());
+        testAllMessageRoundTripsFromVersion((short) 2, newResponse.get().setThrottleTimeMs(1000));
+        testAllMessageRoundTrips(newResponse.get().members().get(0).setGroupInstanceId(null));
+        testAllMessageRoundTripsFromVersion((short) 5, newResponse.get().members().get(0).setGroupInstanceId("instanceId"));
+    }
+
+    @Test
+    public void testLeaveGroupResponseVersions() throws Exception {
+        Supplier newResponse = () -> new LeaveGroupResponseData()
+                                                                 .setErrorCode(Errors.NOT_COORDINATOR.code());
+
+        testAllMessageRoundTrips(newResponse.get());
+        testAllMessageRoundTripsFromVersion((short) 1, newResponse.get().setThrottleTimeMs(1000));
+
+        testAllMessageRoundTripsFromVersion((short) 3, newResponse.get().setMembers(
+            Collections.singletonList(new MemberResponse()
+            .setMemberId(memberId)
+            .setGroupInstanceId(instanceId))
+        ));
+    }
+
+    @Test
+    public void testSyncGroupDefaultGroupInstanceId() throws Exception {
+        Supplier request = () -> new SyncGroupRequestData()
+                .setGroupId("groupId")
+                .setMemberId(memberId)
+                .setGenerationId(15)
+                .setAssignments(new ArrayList<>());
+        testAllMessageRoundTrips(request.get());
+        testAllMessageRoundTrips(request.get().setGroupInstanceId(null));
+        testAllMessageRoundTripsFromVersion((short) 3, request.get().setGroupInstanceId(instanceId));
+    }
+
+    @Test
+    public void testOffsetCommitDefaultGroupInstanceId() throws Exception {
+        testAllMessageRoundTrips(new OffsetCommitRequestData()
+                .setTopics(new ArrayList<>())
+                .setGroupId("groupId"));
+
+        Supplier request = () -> new OffsetCommitRequestData()
+                .setGroupId("groupId")
+                .setMemberId(memberId)
+                .setTopics(new ArrayList<>())
+                .setGenerationId(15);
+        testAllMessageRoundTripsFromVersion((short) 1, request.get());
+        testAllMessageRoundTripsFromVersion((short) 1, request.get().setGroupInstanceId(null));
+        testAllMessageRoundTripsFromVersion((short) 7, request.get().setGroupInstanceId(instanceId));
+    }
+
+    @Test
+    public void testOffsetForLeaderEpochVersions() throws Exception {
+        // Version 2 adds optional current leader epoch
+        OffsetForLeaderEpochRequestData.OffsetForLeaderPartition partitionDataNoCurrentEpoch =
+                new OffsetForLeaderEpochRequestData.OffsetForLeaderPartition()
+                        .setPartitionIndex(0)
+                        .setLeaderEpoch(3);
+        OffsetForLeaderEpochRequestData.OffsetForLeaderPartition partitionDataWithCurrentEpoch =
+                new OffsetForLeaderEpochRequestData.OffsetForLeaderPartition()
+                        .setPartitionIndex(0)
+                        .setLeaderEpoch(3)
+                        .setCurrentLeaderEpoch(5);
+
+        testAllMessageRoundTrips(new OffsetForLeaderEpochRequestData().setTopics(singletonList(
+                new OffsetForLeaderEpochRequestData.OffsetForLeaderTopic()
+                        .setName("foo")
+                        .setPartitions(singletonList(partitionDataNoCurrentEpoch)))
+        ));
+        testAllMessageRoundTripsBeforeVersion((short) 2, partitionDataWithCurrentEpoch, partitionDataNoCurrentEpoch);
+        testAllMessageRoundTripsFromVersion((short) 2, partitionDataWithCurrentEpoch);
+
+        // Version 3 adds the optional replica Id field
+        testAllMessageRoundTripsFromVersion((short) 3, new OffsetForLeaderEpochRequestData().setReplicaId(5));
+        testAllMessageRoundTripsBeforeVersion((short) 3,
+                new OffsetForLeaderEpochRequestData().setReplicaId(5),
+                new OffsetForLeaderEpochRequestData());
+        testAllMessageRoundTripsBeforeVersion((short) 3,
+                new OffsetForLeaderEpochRequestData().setReplicaId(5),
+                new OffsetForLeaderEpochRequestData().setReplicaId(-2));
+
+    }
+
+    @Test
+    public void testLeaderAndIsrVersions() throws Exception {
+        // Version 3 adds two new fields - AddingReplicas and RemovingReplicas
+        LeaderAndIsrRequestData.LeaderAndIsrTopicState partitionStateNoAddingRemovingReplicas =
+            new LeaderAndIsrRequestData.LeaderAndIsrTopicState()
+                .setTopicName("topic")
+                .setPartitionStates(Collections.singletonList(
+                    new LeaderAndIsrRequestData.LeaderAndIsrPartitionState()
+                        .setPartitionIndex(0)
+                        .setReplicas(Collections.singletonList(0))
+                ));
+        LeaderAndIsrRequestData.LeaderAndIsrTopicState partitionStateWithAddingRemovingReplicas =
+            new LeaderAndIsrRequestData.LeaderAndIsrTopicState()
+                .setTopicName("topic")
+                .setPartitionStates(Collections.singletonList(
+                    new LeaderAndIsrRequestData.LeaderAndIsrPartitionState()
+                        .setPartitionIndex(0)
+                        .setReplicas(Collections.singletonList(0))
+                        .setAddingReplicas(Collections.singletonList(1))
+                        .setRemovingReplicas(Collections.singletonList(1))
+                ));
+        testAllMessageRoundTripsBetweenVersions(
+            (short) 2,
+            (short) 3,
+            new LeaderAndIsrRequestData().setTopicStates(Collections.singletonList(partitionStateWithAddingRemovingReplicas)),
+            new LeaderAndIsrRequestData().setTopicStates(Collections.singletonList(partitionStateNoAddingRemovingReplicas)));
+        testAllMessageRoundTripsFromVersion((short) 4, new LeaderAndIsrRequestData().setTopicStates(Collections.singletonList(partitionStateWithAddingRemovingReplicas)));
+    }
+
+    @Test
+    public void testOffsetCommitRequestVersions() throws Exception {
+        String groupId = "groupId";
+        String topicName = "topic";
+        String metadata = "metadata";
+        int partition = 2;
+        int offset = 100;
+
+        testAllMessageRoundTrips(new OffsetCommitRequestData()
+                                     .setGroupId(groupId)
+                                     .setTopics(Collections.singletonList(
+                                         new OffsetCommitRequestTopic()
+                                             .setName(topicName)
+                                             .setPartitions(Collections.singletonList(
+                                                 new OffsetCommitRequestPartition()
+                                                     .setPartitionIndex(partition)
+                                                     .setCommittedMetadata(metadata)
+                                                     .setCommittedOffset(offset)
+                                             )))));
+
+        Supplier request =
+            () -> new OffsetCommitRequestData()
+                      .setGroupId(groupId)
+                      .setMemberId("memberId")
+                      .setGroupInstanceId("instanceId")
+                      .setTopics(Collections.singletonList(
+                          new OffsetCommitRequestTopic()
+                              .setName(topicName)
+                              .setPartitions(Collections.singletonList(
+                                  new OffsetCommitRequestPartition()
+                                      .setPartitionIndex(partition)
+                                      .setCommittedLeaderEpoch(10)
+                                      .setCommittedMetadata(metadata)
+                                      .setCommittedOffset(offset)
+                                      .setCommitTimestamp(20)
+                            ))))
+                    .setRetentionTimeMs(20);
+
+        for (short version = 0; version <= ApiKeys.OFFSET_COMMIT.latestVersion(); version++) {
+            OffsetCommitRequestData requestData = request.get();
+            if (version < 1) {
+                requestData.setMemberId("");
+                requestData.setGenerationId(-1);
+            }
+
+            if (version != 1) {
+                requestData.topics().get(0).partitions().get(0).setCommitTimestamp(-1);
+            }
+
+            if (version < 2 || version > 4) {
+                requestData.setRetentionTimeMs(-1);
+            }
+
+            if (version < 6) {
+                requestData.topics().get(0).partitions().get(0).setCommittedLeaderEpoch(-1);
+
+            }
+
+            if (version < 7) {
+                requestData.setGroupInstanceId(null);
+            }
+
+            if (version == 1) {
+                testEquivalentMessageRoundTrip(version, requestData);
+            } else if (version >= 2 && version <= 4) {
+                testAllMessageRoundTripsBetweenVersions(version, (short) 4, requestData, requestData);
+            } else {
+                testAllMessageRoundTripsFromVersion(version, requestData);
+            }
+        }
+    }
+
+    @Test
+    public void testOffsetCommitResponseVersions() throws Exception {
+        Supplier response =
+            () -> new OffsetCommitResponseData()
+                      .setTopics(
+                          singletonList(
+                              new OffsetCommitResponseTopic()
+                                  .setName("topic")
+                                  .setPartitions(singletonList(
+                                      new OffsetCommitResponsePartition()
+                                          .setPartitionIndex(1)
+                                          .setErrorCode(Errors.UNKNOWN_MEMBER_ID.code())
+                                  ))
+                          )
+                      )
+                      .setThrottleTimeMs(20);
+
+        for (short version = 0; version <= ApiKeys.OFFSET_COMMIT.latestVersion(); version++) {
+            OffsetCommitResponseData responseData = response.get();
+            if (version < 3) {
+                responseData.setThrottleTimeMs(0);
+            }
+            testAllMessageRoundTripsFromVersion(version, responseData);
+        }
+    }
+
+    @Test
+    public void testTxnOffsetCommitRequestVersions() throws Exception {
+        String groupId = "groupId";
+        String topicName = "topic";
+        String metadata = "metadata";
+        String txnId = "transactionalId";
+        int producerId = 25;
+        short producerEpoch = 10;
+
+        int partition = 2;
+        int offset = 100;
+
+        testAllMessageRoundTrips(new TxnOffsetCommitRequestData()
+                                     .setGroupId(groupId)
+                                     .setTransactionalId(txnId)
+                                     .setProducerId(producerId)
+                                     .setProducerEpoch(producerEpoch)
+                                     .setTopics(Collections.singletonList(
+                                         new TxnOffsetCommitRequestTopic()
+                                             .setName(topicName)
+                                             .setPartitions(Collections.singletonList(
+                                                 new TxnOffsetCommitRequestPartition()
+                                                     .setPartitionIndex(partition)
+                                                     .setCommittedMetadata(metadata)
+                                                     .setCommittedOffset(offset)
+                                             )))));
+
+        Supplier request =
+            () -> new TxnOffsetCommitRequestData()
+                      .setGroupId(groupId)
+                      .setTransactionalId(txnId)
+                      .setProducerId(producerId)
+                      .setProducerEpoch(producerEpoch)
+                      .setTopics(Collections.singletonList(
+                          new TxnOffsetCommitRequestTopic()
+                              .setName(topicName)
+                              .setPartitions(Collections.singletonList(
+                                  new TxnOffsetCommitRequestPartition()
+                                      .setPartitionIndex(partition)
+                                      .setCommittedLeaderEpoch(10)
+                                      .setCommittedMetadata(metadata)
+                                      .setCommittedOffset(offset)
+                              ))));
+
+        for (short version = 0; version <= ApiKeys.TXN_OFFSET_COMMIT.latestVersion(); version++) {
+            TxnOffsetCommitRequestData requestData = request.get();
+            if (version < 6) {
+                requestData.topics().get(0).partitions().get(0).setCommittedLeaderEpoch(-1);
+            }
+
+            testAllMessageRoundTripsFromVersion(version, requestData);
+        }
+    }
+
+    @Test
+    public void testTxnOffsetCommitResponseVersions() throws Exception {
+        testAllMessageRoundTrips(
+            new TxnOffsetCommitResponseData()
+                .setTopics(
+                   singletonList(
+                       new TxnOffsetCommitResponseTopic()
+                           .setName("topic")
+                           .setPartitions(singletonList(
+                               new TxnOffsetCommitResponsePartition()
+                                   .setPartitionIndex(1)
+                                   .setErrorCode(Errors.UNKNOWN_MEMBER_ID.code())
+                           ))
+                   )
+               )
+               .setThrottleTimeMs(20));
+    }
+
+    @Test
+    public void testOffsetFetchVersions() throws Exception {
+        String groupId = "groupId";
+        String topicName = "topic";
+
+        testAllMessageRoundTrips(new OffsetFetchRequestData()
+                                     .setTopics(new ArrayList<>())
+                                     .setGroupId(groupId));
+
+        testAllMessageRoundTrips(new OffsetFetchRequestData()
+                                     .setGroupId(groupId)
+                                     .setTopics(Collections.singletonList(
+                                         new OffsetFetchRequestTopic()
+                                             .setName(topicName)
+                                             .setPartitionIndexes(Collections.singletonList(5))))
+        );
+
+        OffsetFetchRequestData allPartitionData = new OffsetFetchRequestData()
+                                                      .setGroupId(groupId)
+                                                      .setTopics(null);
+        for (short version = 0; version <= ApiKeys.OFFSET_FETCH.latestVersion(); version++) {
+            if (version < 2) {
+                final short finalVersion = version;
+                assertThrows(SchemaException.class, () -> testAllMessageRoundTripsFromVersion(finalVersion, allPartitionData));
+            } else {
+                testAllMessageRoundTripsFromVersion(version, allPartitionData);
+            }
+        }
+
+        Supplier response =
+            () -> new OffsetFetchResponseData()
+                      .setTopics(Collections.singletonList(
+                          new OffsetFetchResponseTopic()
+                              .setName(topicName)
+                              .setPartitions(Collections.singletonList(
+                                  new OffsetFetchResponsePartition()
+                                      .setPartitionIndex(5)
+                                      .setMetadata(null)
+                                      .setCommittedOffset(100)
+                                      .setCommittedLeaderEpoch(3)
+                                      .setErrorCode(Errors.UNKNOWN_TOPIC_OR_PARTITION.code())))))
+                      .setErrorCode(Errors.NOT_COORDINATOR.code())
+                      .setThrottleTimeMs(10);
+        for (short version = 0; version <= ApiKeys.OFFSET_FETCH.latestVersion(); version++) {
+            OffsetFetchResponseData responseData = response.get();
+            if (version <= 1) {
+                responseData.setErrorCode(Errors.NONE.code());
+            }
+
+            if (version <= 2) {
+                responseData.setThrottleTimeMs(0);
+            }
+
+            if (version <= 4) {
+                responseData.topics().get(0).partitions().get(0).setCommittedLeaderEpoch(-1);
+            }
+
+            testAllMessageRoundTripsFromVersion(version, responseData);
+        }
+    }
+
+    @Test
+    public void testProduceResponseVersions() throws Exception {
+        String topicName = "topic";
+        int partitionIndex = 0;
+        short errorCode = Errors.INVALID_TOPIC_EXCEPTION.code();
+        long baseOffset = 12L;
+        int throttleTimeMs = 1234;
+        long logAppendTimeMs = 1234L;
+        long logStartOffset = 1234L;
+        int batchIndex = 0;
+        String batchIndexErrorMessage = "error message";
+        String errorMessage = "global error message";
+
+        testAllMessageRoundTrips(new ProduceResponseData()
+                                     .setResponses(singletonList(
+                                         new ProduceResponseData.TopicProduceResponse()
+                                             .setName(topicName)
+                                             .setPartitions(singletonList(
+                                                 new ProduceResponseData.PartitionProduceResponse()
+                                                     .setPartitionIndex(partitionIndex)
+                                                     .setErrorCode(errorCode)
+                                                     .setBaseOffset(baseOffset))))));
+
+        Supplier response =
+            () -> new ProduceResponseData()
+                      .setResponses(singletonList(
+                            new ProduceResponseData.TopicProduceResponse()
+                                .setName(topicName)
+                                .setPartitions(singletonList(
+                                     new ProduceResponseData.PartitionProduceResponse()
+                                         .setPartitionIndex(partitionIndex)
+                                         .setErrorCode(errorCode)
+                                         .setBaseOffset(baseOffset)
+                                         .setLogAppendTimeMs(logAppendTimeMs)
+                                         .setLogStartOffset(logStartOffset)
+                                         .setRecordErrors(singletonList(
+                                             new ProduceResponseData.BatchIndexAndErrorMessage()
+                                                 .setBatchIndex(batchIndex)
+                                                 .setBatchIndexErrorMessage(batchIndexErrorMessage)))
+                                         .setErrorMessage(errorMessage)))))
+                      .setThrottleTimeMs(throttleTimeMs);
+
+        for (short version = 0; version <= ApiKeys.PRODUCE.latestVersion(); version++) {
+            ProduceResponseData responseData = response.get();
+
+            if (version < 8) {
+                responseData.responses().get(0).partitions().get(0).setRecordErrors(Collections.emptyList());
+                responseData.responses().get(0).partitions().get(0).setErrorMessage(null);
+            }
+
+            if (version < 5) {
+                responseData.responses().get(0).partitions().get(0).setLogStartOffset(-1);
+            }
+
+            if (version < 2) {
+                responseData.responses().get(0).partitions().get(0).setLogAppendTimeMs(-1);
+            }
+
+            if (version < 1) {
+                responseData.setThrottleTimeMs(0);
+            }
+
+            if (version >= 3 && version <= 4) {
+                testAllMessageRoundTripsBetweenVersions(version, (short) 4, responseData, responseData);
+            } else if (version >= 6 && version <= 7) {
+                testAllMessageRoundTripsBetweenVersions(version, (short) 7, responseData, responseData);
+            } else {
+                testEquivalentMessageRoundTrip(version, responseData);
+            }
+        }
+    }
+
+    private void testAllMessageRoundTrips(Message message) throws Exception {
+        testAllMessageRoundTripsFromVersion(message.lowestSupportedVersion(), message);
+    }
+
+    private void testAllMessageRoundTripsBeforeVersion(short beforeVersion, Message message, Message expected) throws Exception {
+        testAllMessageRoundTripsBetweenVersions((short) 0, beforeVersion, message, expected);
+    }
+
     /**
-     * Test serializing and deserializing some messages.
+     * @param startVersion - the version we want to start at, inclusive
+     * @param endVersion - the version we want to end at, exclusive
      */
-    @Test
-    public void testRoundTrips() throws Exception {
-        testMessageRoundTrips(new MetadataRequestData().setTopics(
-            Arrays.asList(new MetadataRequestData.MetadataRequestTopic().setName("foo"),
-                new MetadataRequestData.MetadataRequestTopic().setName("bar")
-            )), (short) 6);
-        testMessageRoundTrips(new AddOffsetsToTxnRequestData().
-            setTransactionalId("foobar").
-            setProducerId(0xbadcafebadcafeL).
-            setProducerEpoch((short) 123).
-            setGroupId("baaz"), (short) 1);
-        testMessageRoundTrips(new AddOffsetsToTxnResponseData().
-            setThrottleTimeMs(42).
-            setErrorCode((short) 0), (short) 0);
-        testMessageRoundTrips(new AddPartitionsToTxnRequestData().
-            setTransactionalId("blah").
-            setProducerId(0xbadcafebadcafeL).
-            setProducerEpoch((short) 30000).
-            setTopics(new AddPartitionsToTxnTopicSet(Collections.singletonList(
-                new AddPartitionsToTxnTopic().
-                    setName("Topic").
-                    setPartitions(Collections.singletonList(1))).iterator())));
-        testMessageRoundTrips(new CreateTopicsRequestData().
-            setTimeoutMs(1000).setTopics(new CreatableTopicSet()));
-        testMessageRoundTrips(new DescribeAclsRequestData().
-            setResourceType((byte) 42).
-            setResourceNameFilter(null).
-            setResourcePatternType((byte) 3).
-            setPrincipalFilter("abc").
-            setHostFilter(null).
-            setOperation((byte) 0).
-            setPermissionType((byte) 0), (short) 0);
-    }
-
-    private void testMessageRoundTrips(Message message) throws Exception {
-        testMessageRoundTrips(message, message.highestSupportedVersion());
-    }
-
-    private void testMessageRoundTrips(Message message, short version) throws Exception {
-        testStructRoundTrip(message, version);
-        testByteBufferRoundTrip(message, version);
-    }
-
-    private void testByteBufferRoundTrip(Message message, short version) throws Exception {
-        int size = message.size(version);
+    private void testAllMessageRoundTripsBetweenVersions(short startVersion, short endVersion, Message message, Message expected) throws Exception {
+        for (short version = startVersion; version < endVersion; version++) {
+            testMessageRoundTrip(version, message, expected);
+        }
+    }
+
+    private void testAllMessageRoundTripsFromVersion(short fromVersion, Message message) throws Exception {
+        for (short version = fromVersion; version < message.highestSupportedVersion(); version++) {
+            testEquivalentMessageRoundTrip(version, message);
+        }
+    }
+
+    private void testMessageRoundTrip(short version, Message message, Message expected) throws Exception {
+        testByteBufferRoundTrip(version, message, expected);
+        testStructRoundTrip(version, message, expected);
+    }
+
+    private void testEquivalentMessageRoundTrip(short version, Message message) throws Exception {
+        testStructRoundTrip(version, message, message);
+        testByteBufferRoundTrip(version, message, message);
+    }
+
+    private void testByteBufferRoundTrip(short version, Message message, Message expected) throws Exception {
+        ObjectSerializationCache cache = new ObjectSerializationCache();
+        int size = message.size(cache, version);
         ByteBuffer buf = ByteBuffer.allocate(size);
         ByteBufferAccessor byteBufferAccessor = new ByteBufferAccessor(buf);
-        message.write(byteBufferAccessor, version);
-        assertEquals(size, buf.position());
+        message.write(byteBufferAccessor, cache, version);
+        assertEquals("The result of the size function does not match the number of bytes " +
+            "written for version " + version, size, buf.position());
         Message message2 = message.getClass().newInstance();
         buf.flip();
         message2.read(byteBufferAccessor, version);
-        assertEquals(size, buf.position());
-        assertEquals(message, message2);
-        assertEquals(message.hashCode(), message2.hashCode());
-        assertEquals(message.toString(), message2.toString());
+        assertEquals("The result of the size function does not match the number of bytes " +
+            "read back in for version " + version, size, buf.position());
+        assertEquals("The message object created after a round trip did not match for " +
+            "version " + version, expected, message2);
+        assertEquals(expected.hashCode(), message2.hashCode());
+        assertEquals(expected.toString(), message2.toString());
     }
 
-    private void testStructRoundTrip(Message message, short version) throws Exception {
+    private void testStructRoundTrip(short version, Message message, Message expected) throws Exception {
         Struct struct = message.toStruct(version);
         Message message2 = message.getClass().newInstance();
         message2.fromStruct(struct, version);
-        assertEquals(message, message2);
-        assertEquals(message.hashCode(), message2.hashCode());
-        assertEquals(message.toString(), message2.toString());
+        assertEquals(expected, message2);
+        assertEquals(expected.hashCode(), message2.hashCode());
+        assertEquals(expected.toString(), message2.toString());
     }
 
     /**
@@ -129,7 +650,7 @@ public void testMessageVersions() throws Exception {
         for (ApiKeys apiKey : ApiKeys.values()) {
             Message message = null;
             try {
-                message = ApiMessageFactory.newRequest(apiKey.id);
+                message = ApiMessageType.fromApiKey(apiKey.id).newRequest();
             } catch (UnsupportedVersionException e) {
                 fail("No request message spec found for API " + apiKey);
             }
@@ -137,7 +658,7 @@ public void testMessageVersions() throws Exception {
                     "supports versions up to " + message.highestSupportedVersion(),
                 apiKey.latestVersion() <= message.highestSupportedVersion());
             try {
-                message = ApiMessageFactory.newResponse(apiKey.id);
+                message = ApiMessageType.fromApiKey(apiKey.id).newResponse();
             } catch (UnsupportedVersionException e) {
                 fail("No response message spec found for API " + apiKey);
             }
@@ -154,7 +675,7 @@ public void testMessageVersions() throws Exception {
     public void testRequestSchemas() throws Exception {
         for (ApiKeys apiKey : ApiKeys.values()) {
             Schema[] manualSchemas = apiKey.requestSchemas;
-            Schema[] generatedSchemas = ApiMessageFactory.requestSchemas(apiKey.id);
+            Schema[] generatedSchemas = ApiMessageType.fromApiKey(apiKey.id).requestSchemas();
             Assert.assertEquals("Mismatching request SCHEMAS lengths " +
                 "for api key " + apiKey, manualSchemas.length, generatedSchemas.length);
             for (int v = 0; v < manualSchemas.length; v++) {
@@ -174,10 +695,10 @@ public void testRequestSchemas() throws Exception {
      * Test that the JSON response files match the schemas accessible through the ApiKey class.
      */
     @Test
-    public void testResponseSchemas() throws Exception {
+    public void testResponseSchemas() {
         for (ApiKeys apiKey : ApiKeys.values()) {
             Schema[] manualSchemas = apiKey.responseSchemas;
-            Schema[] generatedSchemas = ApiMessageFactory.responseSchemas(apiKey.id);
+            Schema[] generatedSchemas = ApiMessageType.fromApiKey(apiKey.id).responseSchemas();
             Assert.assertEquals("Mismatching response SCHEMAS lengths " +
                 "for api key " + apiKey, manualSchemas.length, generatedSchemas.length);
             for (int v = 0; v < manualSchemas.length; v++) {
@@ -251,9 +772,9 @@ private static void compareTypes(NamedType typeA, NamedType typeB) {
                     entryA, entryA.type.isNullable() ? "nullable" : "non-nullable",
                     entryB, entryB.type.isNullable() ? "nullable" : "non-nullable"));
             }
-            if (entryA.type instanceof ArrayOf) {
-                compareTypes(new NamedType(entryA.name, ((ArrayOf) entryA.type).type()),
-                             new NamedType(entryB.name, ((ArrayOf) entryB.type).type()));
+            if (entryA.type.isArray()) {
+                compareTypes(new NamedType(entryA.name, entryA.type.arrayElementType().get()),
+                             new NamedType(entryB.name, entryB.type.arrayElementType().get()));
             }
         }
     }
@@ -267,7 +788,7 @@ private static void compareTypes(NamedType typeA, NamedType typeB) {
      */
     private static List flatten(NamedType type) {
         if (!(type.type instanceof Schema)) {
-            return Collections.singletonList(type);
+            return singletonList(type);
         }
         Schema schema = (Schema) type.type;
         ArrayList results = new ArrayList<>();
@@ -279,30 +800,95 @@ private static List flatten(NamedType type) {
 
     @Test
     public void testDefaultValues() throws Exception {
-        verifySizeRaisesUve((short) 0, "validateOnly",
+        verifyWriteRaisesUve((short) 0, "validateOnly",
             new CreateTopicsRequestData().setValidateOnly(true));
-        verifySizeSucceeds((short) 0,
+        verifyWriteSucceeds((short) 0,
             new CreateTopicsRequestData().setValidateOnly(false));
-        verifySizeSucceeds((short) 0,
+        verifyWriteSucceeds((short) 0,
             new OffsetCommitRequestData().setRetentionTimeMs(123));
-        verifySizeRaisesUve((short) 5, "forgotten",
-            new FetchRequestData().setForgotten(Collections.singletonList(
+        verifyWriteRaisesUve((short) 5, "forgotten",
+            new FetchRequestData().setForgotten(singletonList(
                 new FetchRequestData.ForgottenTopic().setName("foo"))));
     }
 
-    private void verifySizeRaisesUve(short version, String problemFieldName,
-                                     Message message) throws Exception {
-        try {
-            message.size(version);
-            fail("Expected to see an UnsupportedVersionException when writing " +
-                message + " at version " + version);
-        } catch (UnsupportedVersionException e) {
-            assertTrue("Expected to get an error message about " + problemFieldName,
-                e.getMessage().contains(problemFieldName));
+    @Test
+    public void testNonIgnorableFieldWithDefaultNull() throws Exception {
+        // Test non-ignorable string field `groupInstanceId` with default null
+        verifyWriteRaisesUve((short) 0, "groupInstanceId", new HeartbeatRequestData()
+                .setGroupId("groupId")
+                .setGenerationId(15)
+                .setMemberId(memberId)
+                .setGroupInstanceId(instanceId));
+        verifyWriteSucceeds((short) 0, new HeartbeatRequestData()
+                .setGroupId("groupId")
+                .setGenerationId(15)
+                .setMemberId(memberId)
+                .setGroupInstanceId(null));
+        verifyWriteSucceeds((short) 0, new HeartbeatRequestData()
+                .setGroupId("groupId")
+                .setGenerationId(15)
+                .setMemberId(memberId));
+    }
+
+    @Test
+    public void testWriteNullForNonNullableFieldRaisesException() throws Exception {
+        CreateTopicsRequestData createTopics = new CreateTopicsRequestData().setTopics(null);
+        for (short i = (short) 0; i <= createTopics.highestSupportedVersion(); i++) {
+            verifyWriteRaisesNpe(i, createTopics);
         }
+        MetadataRequestData metadata = new MetadataRequestData().setTopics(null);
+        verifyWriteRaisesNpe((short) 0, metadata);
     }
 
-    private void verifySizeSucceeds(short version, Message message) throws Exception {
-        message.size(version);
+    @Test
+    public void testUnknownTaggedFields() throws Exception {
+        CreateTopicsRequestData createTopics = new CreateTopicsRequestData();
+        verifyWriteSucceeds((short) 6, createTopics);
+        RawTaggedField field1000 = new RawTaggedField(1000, new byte[] {0x1, 0x2, 0x3});
+        createTopics.unknownTaggedFields().add(field1000);
+        verifyWriteRaisesUve((short) 0, "Tagged fields were set", createTopics);
+        verifyWriteSucceeds((short) 6, createTopics);
+    }
+
+    private void verifyWriteRaisesNpe(short version, Message message) throws Exception {
+        ObjectSerializationCache cache = new ObjectSerializationCache();
+        assertThrows(NullPointerException.class, () -> {
+            int size = message.size(cache, version);
+            ByteBuffer buf = ByteBuffer.allocate(size);
+            ByteBufferAccessor byteBufferAccessor = new ByteBufferAccessor(buf);
+            message.write(byteBufferAccessor, cache, version);
+        });
+    }
+
+    private void verifyWriteRaisesUve(short version,
+                                      String problemText,
+                                     Message message) throws Exception {
+        ObjectSerializationCache cache = new ObjectSerializationCache();
+        UnsupportedVersionException e =
+            assertThrows(UnsupportedVersionException.class, () -> {
+                int size = message.size(cache, version);
+                ByteBuffer buf = ByteBuffer.allocate(size);
+                ByteBufferAccessor byteBufferAccessor = new ByteBufferAccessor(buf);
+                message.write(byteBufferAccessor, cache, version);
+            });
+        assertTrue("Expected to get an error message about " + problemText +
+                ", but got: " + e.getMessage(),
+                e.getMessage().contains(problemText));
+    }
+
+    private void verifyWriteSucceeds(short version, Message message) throws Exception {
+        ObjectSerializationCache cache = new ObjectSerializationCache();
+        int size = message.size(cache, version);
+        ByteBuffer buf = ByteBuffer.allocate(size * 2);
+        ByteBufferAccessor byteBufferAccessor = new ByteBufferAccessor(buf);
+        message.write(byteBufferAccessor, cache, version);
+        ByteBuffer alt = buf.duplicate();
+        alt.flip();
+        StringBuilder bld = new StringBuilder();
+        while (alt.hasRemaining()) {
+            bld.append(String.format(" %02x", alt.get()));
+        }
+        assertEquals("Expected the serialized size to be " + size +
+            ", but it was " + buf.position(), size, buf.position());
     }
 }
diff --git a/clients/src/test/java/org/apache/kafka/common/message/SimpleExampleMessageTest.java b/clients/src/test/java/org/apache/kafka/common/message/SimpleExampleMessageTest.java
new file mode 100644
index 0000000000000..5a3026f1873c3
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/message/SimpleExampleMessageTest.java
@@ -0,0 +1,104 @@
+/*
+ * 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.message;
+
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
+import org.apache.kafka.common.protocol.MessageTestUtil;
+import org.apache.kafka.common.protocol.ObjectSerializationCache;
+import org.apache.kafka.common.protocol.types.Struct;
+import org.junit.Test;
+
+import java.nio.ByteBuffer;
+import java.util.Arrays;
+import java.util.UUID;
+
+import static org.junit.Assert.assertEquals;
+
+public class SimpleExampleMessageTest {
+
+    @Test
+    public void shouldStoreField() {
+        final UUID uuid = UUID.randomUUID();
+        final SimpleExampleMessageData out = new SimpleExampleMessageData();
+        out.setProcessId(uuid);
+        assertEquals(uuid, out.processId());
+    }
+
+    @Test
+    public void shouldDefaultField() {
+        final SimpleExampleMessageData out = new SimpleExampleMessageData();
+        assertEquals(UUID.fromString("00000000-0000-0000-0000-000000000000"), out.processId());
+    }
+
+    @Test
+    public void shouldRoundTripFieldThroughStruct() {
+        final UUID uuid = UUID.randomUUID();
+        final SimpleExampleMessageData out = new SimpleExampleMessageData();
+        out.setProcessId(uuid);
+
+        final Struct struct = out.toStruct((short) 1);
+        final SimpleExampleMessageData in = new SimpleExampleMessageData();
+        in.fromStruct(struct, (short) 1);
+
+        assertEquals(uuid, in.processId());
+    }
+
+    @Test
+    public void shouldRoundTripFieldThroughBuffer() {
+        final UUID uuid = UUID.randomUUID();
+        final SimpleExampleMessageData out = new SimpleExampleMessageData();
+        out.setProcessId(uuid);
+
+        ObjectSerializationCache cache = new ObjectSerializationCache();
+        final ByteBuffer buffer = ByteBuffer.allocate(out.size(cache, (short) 1));
+        out.write(new ByteBufferAccessor(buffer), cache, (short) 1);
+        buffer.rewind();
+
+        final SimpleExampleMessageData in = new SimpleExampleMessageData();
+        in.read(new ByteBufferAccessor(buffer), (short) 1);
+
+        assertEquals(uuid, in.processId());
+    }
+
+    @Test
+    public void shouldImplementJVMMethods() {
+        final UUID uuid = UUID.randomUUID();
+        final SimpleExampleMessageData a = new SimpleExampleMessageData();
+        a.setProcessId(uuid);
+
+        final SimpleExampleMessageData b = new SimpleExampleMessageData();
+        b.setProcessId(uuid);
+
+        assertEquals(a, b);
+        assertEquals(a.hashCode(), b.hashCode());
+        // just tagging this on here
+        assertEquals(a.toString(), b.toString());
+    }
+
+    @Test
+    public void testMyTaggedIntArray() {
+        final SimpleExampleMessageData data = new SimpleExampleMessageData();
+        data.setMyTaggedIntArray(Arrays.asList(1, 2, 3));
+        short version = 1;
+        ByteBuffer buf = MessageTestUtil.messageToByteBuffer(data, version);
+        final SimpleExampleMessageData data2 = new SimpleExampleMessageData();
+        data2.read(new ByteBufferAccessor(buf.duplicate()), version);
+        assertEquals(Arrays.asList(1, 2, 3), data.myTaggedIntArray());
+        assertEquals(Arrays.asList(1, 2, 3), data2.myTaggedIntArray());
+        assertEquals(data, data2);
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/metrics/JmxReporterTest.java b/clients/src/test/java/org/apache/kafka/common/metrics/JmxReporterTest.java
index c6e112adf52ec..37d1182262667 100644
--- a/clients/src/test/java/org/apache/kafka/common/metrics/JmxReporterTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/metrics/JmxReporterTest.java
@@ -18,7 +18,7 @@
 
 import org.apache.kafka.common.MetricName;
 import org.apache.kafka.common.metrics.stats.Avg;
-import org.apache.kafka.common.metrics.stats.Total;
+import org.apache.kafka.common.metrics.stats.CumulativeSum;
 import org.junit.Test;
 
 import javax.management.MBeanServer;
@@ -43,7 +43,7 @@ public void testJmxRegistration() throws Exception {
 
             Sensor sensor = metrics.sensor("kafka.requests");
             sensor.add(metrics.metricName("pack.bean1.avg", "grp1"), new Avg());
-            sensor.add(metrics.metricName("pack.bean2.total", "grp2"), new Total());
+            sensor.add(metrics.metricName("pack.bean2.total", "grp2"), new CumulativeSum());
 
             assertTrue(server.isRegistered(new ObjectName(":type=grp1")));
             assertEquals(Double.NaN, server.getAttribute(new ObjectName(":type=grp1"), "pack.bean1.avg"));
@@ -79,11 +79,11 @@ public void testJmxRegistrationSanitization() throws Exception {
             metrics.addReporter(new JmxReporter());
 
             Sensor sensor = metrics.sensor("kafka.requests");
-            sensor.add(metrics.metricName("name", "group", "desc", "id", "foo*"), new Total());
-            sensor.add(metrics.metricName("name", "group", "desc", "id", "foo+"), new Total());
-            sensor.add(metrics.metricName("name", "group", "desc", "id", "foo?"), new Total());
-            sensor.add(metrics.metricName("name", "group", "desc", "id", "foo:"), new Total());
-            sensor.add(metrics.metricName("name", "group", "desc", "id", "foo%"), new Total());
+            sensor.add(metrics.metricName("name", "group", "desc", "id", "foo*"), new CumulativeSum());
+            sensor.add(metrics.metricName("name", "group", "desc", "id", "foo+"), new CumulativeSum());
+            sensor.add(metrics.metricName("name", "group", "desc", "id", "foo?"), new CumulativeSum());
+            sensor.add(metrics.metricName("name", "group", "desc", "id", "foo:"), new CumulativeSum());
+            sensor.add(metrics.metricName("name", "group", "desc", "id", "foo%"), new CumulativeSum());
 
             assertTrue(server.isRegistered(new ObjectName(":type=group,id=\"foo\\*\"")));
             assertEquals(0.0, server.getAttribute(new ObjectName(":type=group,id=\"foo\\*\""), "name"));
diff --git a/clients/src/test/java/org/apache/kafka/common/metrics/KafkaMbeanTest.java b/clients/src/test/java/org/apache/kafka/common/metrics/KafkaMbeanTest.java
index dd494c88778d9..ea1178e0c9788 100644
--- a/clients/src/test/java/org/apache/kafka/common/metrics/KafkaMbeanTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/metrics/KafkaMbeanTest.java
@@ -17,8 +17,8 @@
 package org.apache.kafka.common.metrics;
 
 import org.apache.kafka.common.MetricName;
-import org.apache.kafka.common.metrics.stats.Count;
-import org.apache.kafka.common.metrics.stats.Sum;
+import org.apache.kafka.common.metrics.stats.WindowedCount;
+import org.apache.kafka.common.metrics.stats.WindowedSum;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
@@ -51,9 +51,9 @@ public void setup() throws Exception {
         metrics.addReporter(new JmxReporter());
         sensor = metrics.sensor("kafka.requests");
         countMetricName = metrics.metricName("pack.bean1.count", "grp1");
-        sensor.add(countMetricName, new Count());
+        sensor.add(countMetricName, new WindowedCount());
         sumMetricName = metrics.metricName("pack.bean1.sum", "grp1");
-        sensor.add(sumMetricName, new Sum());
+        sensor.add(sumMetricName, new WindowedSum());
     }
 
     @After
diff --git a/clients/src/test/java/org/apache/kafka/common/metrics/MetricsTest.java b/clients/src/test/java/org/apache/kafka/common/metrics/MetricsTest.java
index 98b468afc3d1c..a221d4cdfd875 100644
--- a/clients/src/test/java/org/apache/kafka/common/metrics/MetricsTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/metrics/MetricsTest.java
@@ -44,7 +44,7 @@
 import org.apache.kafka.common.Metric;
 import org.apache.kafka.common.MetricName;
 import org.apache.kafka.common.metrics.stats.Avg;
-import org.apache.kafka.common.metrics.stats.Count;
+import org.apache.kafka.common.metrics.stats.CumulativeSum;
 import org.apache.kafka.common.metrics.stats.Max;
 import org.apache.kafka.common.metrics.stats.Meter;
 import org.apache.kafka.common.metrics.stats.Min;
@@ -52,9 +52,9 @@
 import org.apache.kafka.common.metrics.stats.Percentiles;
 import org.apache.kafka.common.metrics.stats.Percentiles.BucketSizing;
 import org.apache.kafka.common.metrics.stats.Rate;
+import org.apache.kafka.common.metrics.stats.WindowedCount;
+import org.apache.kafka.common.metrics.stats.WindowedSum;
 import org.apache.kafka.common.metrics.stats.SimpleRate;
-import org.apache.kafka.common.metrics.stats.Sum;
-import org.apache.kafka.common.metrics.stats.Total;
 import org.apache.kafka.common.metrics.stats.Value;
 import org.apache.kafka.common.utils.MockTime;
 import org.junit.After;
@@ -63,7 +63,6 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-@SuppressWarnings("deprecation")
 public class MetricsTest {
     private static final Logger log = LoggerFactory.getLogger(MetricsTest.class);
 
@@ -75,7 +74,7 @@ public class MetricsTest {
 
     @Before
     public void setup() {
-        this.metrics = new Metrics(config, Arrays.asList((MetricsReporter) new JmxReporter()), time, true);
+        this.metrics = new Metrics(config, Arrays.asList(new JmxReporter()), time, true);
     }
 
     @After
@@ -119,15 +118,15 @@ private void verifyStats(Function metricValueFunc) {
         s.add(metrics.metricName("test.min", "grp1"), new Min());
         s.add(new Meter(TimeUnit.SECONDS, metrics.metricName("test.rate", "grp1"),
                 metrics.metricName("test.total", "grp1")));
-        s.add(new Meter(TimeUnit.SECONDS, new Count(), metrics.metricName("test.occurences", "grp1"),
+        s.add(new Meter(TimeUnit.SECONDS, new WindowedCount(), metrics.metricName("test.occurences", "grp1"),
                 metrics.metricName("test.occurences.total", "grp1")));
-        s.add(metrics.metricName("test.count", "grp1"), new Count());
+        s.add(metrics.metricName("test.count", "grp1"), new WindowedCount());
         s.add(new Percentiles(100, -100, 100, BucketSizing.CONSTANT,
                              new Percentile(metrics.metricName("test.median", "grp1"), 50.0),
                              new Percentile(metrics.metricName("test.perc99_9", "grp1"), 99.9)));
 
         Sensor s2 = metrics.sensor("test.sensor2");
-        s2.add(metrics.metricName("s2.total", "grp1"), new Total());
+        s2.add(metrics.metricName("s2.total", "grp1"), new CumulativeSum());
         s2.record(5.0);
 
         int sum = 0;
@@ -162,15 +161,15 @@ private void verifyStats(Function metricValueFunc) {
     @Test
     public void testHierarchicalSensors() {
         Sensor parent1 = metrics.sensor("test.parent1");
-        parent1.add(metrics.metricName("test.parent1.count", "grp1"), new Count());
+        parent1.add(metrics.metricName("test.parent1.count", "grp1"), new WindowedCount());
         Sensor parent2 = metrics.sensor("test.parent2");
-        parent2.add(metrics.metricName("test.parent2.count", "grp1"), new Count());
+        parent2.add(metrics.metricName("test.parent2.count", "grp1"), new WindowedCount());
         Sensor child1 = metrics.sensor("test.child1", parent1, parent2);
-        child1.add(metrics.metricName("test.child1.count", "grp1"), new Count());
+        child1.add(metrics.metricName("test.child1.count", "grp1"), new WindowedCount());
         Sensor child2 = metrics.sensor("test.child2", parent1);
-        child2.add(metrics.metricName("test.child2.count", "grp1"), new Count());
+        child2.add(metrics.metricName("test.child2.count", "grp1"), new WindowedCount());
         Sensor grandchild = metrics.sensor("test.grandchild", child1);
-        grandchild.add(metrics.metricName("test.grandchild.count", "grp1"), new Count());
+        grandchild.add(metrics.metricName("test.grandchild.count", "grp1"), new WindowedCount());
 
         /* increment each sensor one time */
         parent1.record();
@@ -222,15 +221,15 @@ public void testRemoveChildSensor() {
     public void testRemoveSensor() {
         int size = metrics.metrics().size();
         Sensor parent1 = metrics.sensor("test.parent1");
-        parent1.add(metrics.metricName("test.parent1.count", "grp1"), new Count());
+        parent1.add(metrics.metricName("test.parent1.count", "grp1"), new WindowedCount());
         Sensor parent2 = metrics.sensor("test.parent2");
-        parent2.add(metrics.metricName("test.parent2.count", "grp1"), new Count());
+        parent2.add(metrics.metricName("test.parent2.count", "grp1"), new WindowedCount());
         Sensor child1 = metrics.sensor("test.child1", parent1, parent2);
-        child1.add(metrics.metricName("test.child1.count", "grp1"), new Count());
+        child1.add(metrics.metricName("test.child1.count", "grp1"), new WindowedCount());
         Sensor child2 = metrics.sensor("test.child2", parent2);
-        child2.add(metrics.metricName("test.child2.count", "grp1"), new Count());
+        child2.add(metrics.metricName("test.child2.count", "grp1"), new WindowedCount());
         Sensor grandChild1 = metrics.sensor("test.gchild2", child2);
-        grandChild1.add(metrics.metricName("test.gchild2.count", "grp1"), new Count());
+        grandChild1.add(metrics.metricName("test.gchild2.count", "grp1"), new WindowedCount());
 
         Sensor sensor = metrics.getSensor("test.parent1");
         assertNotNull(sensor);
@@ -268,10 +267,10 @@ public void testRemoveSensor() {
     @Test
     public void testRemoveInactiveMetrics() {
         Sensor s1 = metrics.sensor("test.s1", null, 1);
-        s1.add(metrics.metricName("test.s1.count", "grp1"), new Count());
+        s1.add(metrics.metricName("test.s1.count", "grp1"), new WindowedCount());
 
         Sensor s2 = metrics.sensor("test.s2", null, 3);
-        s2.add(metrics.metricName("test.s2.count", "grp1"), new Count());
+        s2.add(metrics.metricName("test.s2.count", "grp1"), new WindowedCount());
 
         Metrics.ExpireSensorTask purger = metrics.new ExpireSensorTask();
         purger.run();
@@ -309,7 +308,7 @@ public void testRemoveInactiveMetrics() {
 
         // After purging, it should be possible to recreate a metric
         s1 = metrics.sensor("test.s1", null, 1);
-        s1.add(metrics.metricName("test.s1.count", "grp1"), new Count());
+        s1.add(metrics.metricName("test.s1.count", "grp1"), new WindowedCount());
         assertNotNull("Sensor test.s1 must be present", metrics.getSensor("test.s1"));
         assertNotNull("MetricName test.s1.count must be present",
                 metrics.metrics().get(metrics.metricName("test.s1.count", "grp1")));
@@ -318,8 +317,8 @@ public void testRemoveInactiveMetrics() {
     @Test
     public void testRemoveMetric() {
         int size = metrics.metrics().size();
-        metrics.addMetric(metrics.metricName("test1", "grp1"), new Count());
-        metrics.addMetric(metrics.metricName("test2", "grp1"), new Count());
+        metrics.addMetric(metrics.metricName("test1", "grp1"), new WindowedCount());
+        metrics.addMetric(metrics.metricName("test2", "grp1"), new WindowedCount());
 
         assertNotNull(metrics.removeMetric(metrics.metricName("test1", "grp1")));
         assertNull(metrics.metrics().get(metrics.metricName("test1", "grp1")));
@@ -333,7 +332,7 @@ public void testRemoveMetric() {
 
     @Test
     public void testEventWindowing() {
-        Count count = new Count();
+        WindowedCount count = new WindowedCount();
         MetricConfig config = new MetricConfig().eventWindow(1).samples(2);
         count.record(config, 1.0, time.milliseconds());
         count.record(config, 1.0, time.milliseconds());
@@ -344,7 +343,7 @@ public void testEventWindowing() {
 
     @Test
     public void testTimeWindowing() {
-        Count count = new Count();
+        WindowedCount count = new WindowedCount();
         MetricConfig config = new MetricConfig().timeWindow(1, TimeUnit.MILLISECONDS).samples(2);
         count.record(config, 1.0, time.milliseconds());
         time.sleep(1);
@@ -397,8 +396,8 @@ public void testSampledStatReturnsNaNWhenNoValuesExist() {
      */
     @Test
     public void testSampledStatReturnsInitialValueWhenNoValuesExist() {
-        Count count = new Count();
-        Rate.SampledTotal sampledTotal = new Rate.SampledTotal();
+        WindowedCount count = new WindowedCount();
+        WindowedSum sampledTotal = new WindowedSum();
         long windowMs = 100;
         int samples = 2;
         MetricConfig config = new MetricConfig().timeWindow(windowMs, TimeUnit.MILLISECONDS).samples(samples);
@@ -415,14 +414,32 @@ public void testSampledStatReturnsInitialValueWhenNoValuesExist() {
     @Test(expected = IllegalArgumentException.class)
     public void testDuplicateMetricName() {
         metrics.sensor("test").add(metrics.metricName("test", "grp1"), new Avg());
-        metrics.sensor("test2").add(metrics.metricName("test", "grp1"), new Total());
+        metrics.sensor("test2").add(metrics.metricName("test", "grp1"), new CumulativeSum());
+    }
+
+    @Test
+    public void testDuplicateMetricNameOptionallyReplace() {
+        metrics.setReplaceOnDuplicateMetric(true);
+
+        int initialSize = metrics.metrics().size();
+        MetricName metricName = metrics.metricName("test1", "grp1");
+        metrics.addMetric(metricName, new WindowedCount());
+        assertEquals(initialSize + 1, metrics.metrics().size());
+
+        metrics.addMetric(metricName, new WindowedCount());
+        assertEquals(initialSize + 1, metrics.metrics().size());
+
+        assertNotNull(metrics.removeMetric(metricName));
+        assertNull(metrics.metrics().get(metricName));
+
+        assertEquals(initialSize, metrics.metrics().size());
     }
 
     @Test
     public void testQuotas() {
         Sensor sensor = metrics.sensor("test");
-        sensor.add(metrics.metricName("test1.total", "grp1"), new Total(), new MetricConfig().quota(Quota.upperBound(5.0)));
-        sensor.add(metrics.metricName("test2.total", "grp1"), new Total(), new MetricConfig().quota(Quota.lowerBound(0.0)));
+        sensor.add(metrics.metricName("test1.total", "grp1"), new CumulativeSum(), new MetricConfig().quota(Quota.upperBound(5.0)));
+        sensor.add(metrics.metricName("test2.total", "grp1"), new CumulativeSum(), new MetricConfig().quota(Quota.lowerBound(0.0)));
         sensor.record(5.0);
         try {
             sensor.record(1.0);
@@ -503,7 +520,7 @@ public void testRateWindowing() throws Exception {
         MetricName countRateMetricName = metrics.metricName("test.count.rate", "grp1");
         MetricName countTotalMetricName = metrics.metricName("test.count.total", "grp1");
         s.add(new Meter(TimeUnit.SECONDS, rateMetricName, totalMetricName));
-        s.add(new Meter(TimeUnit.SECONDS, new Count(), countRateMetricName, countTotalMetricName));
+        s.add(new Meter(TimeUnit.SECONDS, new WindowedCount(), countRateMetricName, countTotalMetricName));
         KafkaMetric totalMetric = metrics.metrics().get(totalMetricName);
         KafkaMetric countTotalMetric = metrics.metrics().get(countTotalMetricName);
 
@@ -631,7 +648,7 @@ public void testMetricInstances() {
         Map childTagsWithValues = new HashMap<>();
         childTagsWithValues.put("child-tag", "child-tag-value");
 
-        try (Metrics inherited = new Metrics(new MetricConfig().tags(parentTagsWithValues), Arrays.asList((MetricsReporter) new JmxReporter()), time, true)) {
+        try (Metrics inherited = new Metrics(new MetricConfig().tags(parentTagsWithValues), Arrays.asList(new JmxReporter()), time, true)) {
             MetricName inheritedMetric = inherited.metricInstance(SampleMetrics.METRIC_WITH_INHERITED_TAGS, childTagsWithValues);
 
             Map filledOutTags = inheritedMetric.tags();
@@ -731,7 +748,7 @@ synchronized void processMetrics() {
 
         final LockingReporter reporter = new LockingReporter();
         this.metrics.close();
-        this.metrics = new Metrics(config, Arrays.asList((MetricsReporter) reporter), new MockTime(10), true);
+        this.metrics = new Metrics(config, Arrays.asList(reporter), new MockTime(10), true);
         final Deque sensors = new ConcurrentLinkedDeque<>();
         SensorCreator sensorCreator = new SensorCreator(metrics);
 
@@ -825,10 +842,10 @@ private Sensor createSensor(StatType statType, int index) {
                     sensor.add(metrics.metricName("test.metric.avg", "avg", tags), new Avg());
                     break;
                 case TOTAL:
-                    sensor.add(metrics.metricName("test.metric.total", "total", tags), new Total());
+                    sensor.add(metrics.metricName("test.metric.total", "total", tags), new CumulativeSum());
                     break;
                 case COUNT:
-                    sensor.add(metrics.metricName("test.metric.count", "count", tags), new Count());
+                    sensor.add(metrics.metricName("test.metric.count", "count", tags), new WindowedCount());
                     break;
                 case MAX:
                     sensor.add(metrics.metricName("test.metric.max", "max", tags), new Max());
@@ -843,7 +860,7 @@ private Sensor createSensor(StatType statType, int index) {
                     sensor.add(metrics.metricName("test.metric.simpleRate", "simpleRate", tags), new SimpleRate());
                     break;
                 case SUM:
-                    sensor.add(metrics.metricName("test.metric.sum", "sum", tags), new Sum());
+                    sensor.add(metrics.metricName("test.metric.sum", "sum", tags), new WindowedSum());
                     break;
                 case VALUE:
                     sensor.add(metrics.metricName("test.metric.value", "value", tags), new Value());
diff --git a/clients/src/test/java/org/apache/kafka/common/metrics/SensorTest.java b/clients/src/test/java/org/apache/kafka/common/metrics/SensorTest.java
index 8e3dfeb3b9fad..1dc59a8aad543 100644
--- a/clients/src/test/java/org/apache/kafka/common/metrics/SensorTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/metrics/SensorTest.java
@@ -20,7 +20,7 @@
 import org.apache.kafka.common.metrics.stats.Avg;
 import org.apache.kafka.common.metrics.stats.Meter;
 import org.apache.kafka.common.metrics.stats.Rate;
-import org.apache.kafka.common.metrics.stats.Sum;
+import org.apache.kafka.common.metrics.stats.WindowedSum;
 import org.apache.kafka.common.utils.MockTime;
 import org.apache.kafka.common.utils.SystemTime;
 import org.apache.kafka.common.utils.Time;
@@ -86,25 +86,23 @@ public void testShouldRecord() {
     public void testExpiredSensor() {
         MetricConfig config = new MetricConfig();
         Time mockTime = new MockTime();
-        Metrics metrics =  new Metrics(config, Arrays.asList((MetricsReporter) new JmxReporter()), mockTime, true);
-
-        long inactiveSensorExpirationTimeSeconds = 60L;
-        Sensor sensor = new Sensor(metrics, "sensor", null, config, mockTime,
-            inactiveSensorExpirationTimeSeconds, Sensor.RecordingLevel.INFO);
-
-        assertTrue(sensor.add(metrics.metricName("test1", "grp1"), new Avg()));
-
-        Map emptyTags = Collections.emptyMap();
-        MetricName rateMetricName = new MetricName("rate", "test", "", emptyTags);
-        MetricName totalMetricName = new MetricName("total", "test", "", emptyTags);
-        Meter meter = new Meter(rateMetricName, totalMetricName);
-        assertTrue(sensor.add(meter));
-
-        mockTime.sleep(TimeUnit.SECONDS.toMillis(inactiveSensorExpirationTimeSeconds + 1));
-        assertFalse(sensor.add(metrics.metricName("test3", "grp1"), new Avg()));
-        assertFalse(sensor.add(meter));
-
-        metrics.close();
+        try (Metrics metrics = new Metrics(config, Arrays.asList(new JmxReporter()), mockTime, true)) {
+            long inactiveSensorExpirationTimeSeconds = 60L;
+            Sensor sensor = new Sensor(metrics, "sensor", null, config, mockTime,
+                    inactiveSensorExpirationTimeSeconds, Sensor.RecordingLevel.INFO);
+
+            assertTrue(sensor.add(metrics.metricName("test1", "grp1"), new Avg()));
+
+            Map emptyTags = Collections.emptyMap();
+            MetricName rateMetricName = new MetricName("rate", "test", "", emptyTags);
+            MetricName totalMetricName = new MetricName("total", "test", "", emptyTags);
+            Meter meter = new Meter(rateMetricName, totalMetricName);
+            assertTrue(sensor.add(meter));
+
+            mockTime.sleep(TimeUnit.SECONDS.toMillis(inactiveSensorExpirationTimeSeconds + 1));
+            assertFalse(sensor.add(metrics.metricName("test3", "grp1"), new Avg()));
+            assertFalse(sensor.add(meter));
+        }
     }
 
     @Test
@@ -127,8 +125,13 @@ public void testIdempotentAdd() {
             // pass
         }
 
+        // Unless duplicate sensors are allowed, then the sensor will be replaced with an error in the logs
+        metrics.setReplaceOnDuplicateMetric(true);
+        final Sensor anotherAnotherSensor = metrics.sensor("another-sensor");
+        anotherAnotherSensor.add(metrics.metricName("test-metric", "test-group"), new Avg());
+
         // note that adding a different metric with the same name is also a no-op
-        assertTrue(sensor.add(metrics.metricName("test-metric", "test-group"), new Sum()));
+        assertTrue(sensor.add(metrics.metricName("test-metric", "test-group"), new WindowedSum()));
 
         // so after all this, we still just have the original metric registered
         assertEquals(1, sensor.metrics().size());
diff --git a/clients/src/test/java/org/apache/kafka/common/metrics/stats/FrequenciesTest.java b/clients/src/test/java/org/apache/kafka/common/metrics/stats/FrequenciesTest.java
index e3e623fae078e..bf7647a9afb40 100644
--- a/clients/src/test/java/org/apache/kafka/common/metrics/stats/FrequenciesTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/metrics/stats/FrequenciesTest.java
@@ -22,7 +22,6 @@
 import org.apache.kafka.common.metrics.JmxReporter;
 import org.apache.kafka.common.metrics.MetricConfig;
 import org.apache.kafka.common.metrics.Metrics;
-import org.apache.kafka.common.metrics.MetricsReporter;
 import org.apache.kafka.common.metrics.Sensor;
 import org.apache.kafka.common.utils.MockTime;
 import org.apache.kafka.common.utils.Time;
@@ -46,7 +45,7 @@ public class FrequenciesTest {
     public void setup() {
         config = new MetricConfig().eventWindow(50).samples(2);
         time = new MockTime();
-        metrics = new Metrics(config, Arrays.asList((MetricsReporter) new JmxReporter()), time, true);
+        metrics = new Metrics(config, Arrays.asList(new JmxReporter()), time, true);
     }
 
     @After
diff --git a/clients/src/test/java/org/apache/kafka/common/metrics/stats/MeterTest.java b/clients/src/test/java/org/apache/kafka/common/metrics/stats/MeterTest.java
index 8204771d794db..27198ea1e79d0 100644
--- a/clients/src/test/java/org/apache/kafka/common/metrics/stats/MeterTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/metrics/stats/MeterTest.java
@@ -44,7 +44,7 @@ public void testMeter() {
         assertEquals(rateMetricName, rate.name());
         assertEquals(totalMetricName, total.name());
         Rate rateStat = (Rate) rate.stat();
-        Total totalStat = (Total) total.stat();
+        CumulativeSum totalStat = (CumulativeSum) total.stat();
 
         MetricConfig config = new MetricConfig();
         double nextValue = 0.0;
diff --git a/clients/src/test/java/org/apache/kafka/common/network/CertStores.java b/clients/src/test/java/org/apache/kafka/common/network/CertStores.java
index 045e280880aea..d02bdfd91e582 100644
--- a/clients/src/test/java/org/apache/kafka/common/network/CertStores.java
+++ b/clients/src/test/java/org/apache/kafka/common/network/CertStores.java
@@ -41,23 +41,23 @@ public class CertStores {
 
     private final Map sslConfig;
 
-    public CertStores(boolean server, String hostName) throws Exception {
-        this(server, hostName, new TestSslUtils.CertificateBuilder());
+    public CertStores(boolean server, String hostName, TestSslUtils.SSLProvider provider) throws Exception {
+        this(server, hostName, new TestSslUtils.CertificateBuilder(), provider);
     }
 
-    public CertStores(boolean server, String commonName, String sanHostName) throws Exception {
-        this(server, commonName, new TestSslUtils.CertificateBuilder().sanDnsName(sanHostName));
+    public CertStores(boolean server, String commonName, String sanHostName, TestSslUtils.SSLProvider provider) throws Exception {
+        this(server, commonName, new TestSslUtils.CertificateBuilder().sanDnsName(sanHostName), provider);
     }
 
-    public CertStores(boolean server, String commonName, InetAddress hostAddress) throws Exception {
-        this(server, commonName, new TestSslUtils.CertificateBuilder().sanIpAddress(hostAddress));
+    public CertStores(boolean server, String commonName, InetAddress hostAddress, TestSslUtils.SSLProvider provider) throws Exception {
+        this(server, commonName, new TestSslUtils.CertificateBuilder().sanIpAddress(hostAddress), provider);
     }
 
-    private CertStores(boolean server, String commonName, TestSslUtils.CertificateBuilder certBuilder) throws Exception {
+    private CertStores(boolean server, String commonName, TestSslUtils.CertificateBuilder certBuilder, TestSslUtils.SSLProvider provider) throws Exception {
         String name = server ? "server" : "client";
         Mode mode = server ? Mode.SERVER : Mode.CLIENT;
         File truststoreFile = File.createTempFile(name + "TS", ".jks");
-        sslConfig = TestSslUtils.createSslConfig(!server, true, mode, truststoreFile, name, commonName, certBuilder);
+        sslConfig = TestSslUtils.createSslConfig(!server, true, mode, truststoreFile, name, commonName, certBuilder, provider);
     }
 
     public Map getTrustingConfig(CertStores truststoreConfig) {
diff --git a/clients/src/test/java/org/apache/kafka/common/network/EchoServer.java b/clients/src/test/java/org/apache/kafka/common/network/EchoServer.java
index e986598f4ebaf..faf0a9f1e5b48 100644
--- a/clients/src/test/java/org/apache/kafka/common/network/EchoServer.java
+++ b/clients/src/test/java/org/apache/kafka/common/network/EchoServer.java
@@ -50,7 +50,7 @@ public EchoServer(SecurityProtocol securityProtocol, Map configs) thr
             case SSL:
                 this.sslFactory = new SslFactory(Mode.SERVER);
                 this.sslFactory.configure(configs);
-                SSLContext sslContext = this.sslFactory.sslContext();
+                SSLContext sslContext = this.sslFactory.sslEngineBuilder().sslContext();
                 this.serverSocket = sslContext.getServerSocketFactory().createServerSocket(0);
                 break;
             case PLAINTEXT:
diff --git a/clients/src/test/java/org/apache/kafka/common/network/SelectorTest.java b/clients/src/test/java/org/apache/kafka/common/network/SelectorTest.java
index 9d48ab5264967..952d1ceebd9af 100644
--- a/clients/src/test/java/org/apache/kafka/common/network/SelectorTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/network/SelectorTest.java
@@ -21,7 +21,9 @@
 import org.apache.kafka.common.memory.MemoryPool;
 import org.apache.kafka.common.memory.SimpleMemoryPool;
 import org.apache.kafka.common.metrics.KafkaMetric;
+import org.apache.kafka.common.metrics.MetricConfig;
 import org.apache.kafka.common.metrics.Metrics;
+import org.apache.kafka.common.metrics.Sensor;
 import org.apache.kafka.common.security.auth.SecurityProtocol;
 import org.apache.kafka.common.utils.LogContext;
 import org.apache.kafka.common.utils.MockTime;
@@ -59,6 +61,7 @@
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
+import static org.junit.Assert.assertThrows;
 import static org.mockito.Mockito.atLeastOnce;
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.mock;
@@ -78,6 +81,7 @@ public class SelectorTest {
     protected Selector selector;
     protected ChannelBuilder channelBuilder;
     protected Metrics metrics;
+    protected Sensor sensor;
 
     @Before
     public void setUp() throws Exception {
@@ -89,6 +93,7 @@ public void setUp() throws Exception {
         this.channelBuilder.configure(configs);
         this.metrics = new Metrics();
         this.selector = new Selector(5000, this.metrics, time, "MetricGroup", channelBuilder, new LogContext());
+        this.sensor = metrics.sensor("infoSensor", new MetricConfig().recordLevel(Sensor.RecordingLevel.INFO), Sensor.RecordingLevel.INFO);
     }
 
     @After
@@ -493,7 +498,7 @@ private KafkaChannel createConnectionWithStagedReceives(int maxStagedReceives) t
             do {
                 selector.poll(1000);
             } while (selector.completedReceives().isEmpty());
-        } while (selector.numStagedReceives(channel) == 0 && !channel.hasBytesBuffered() && --retries > 0);
+        } while ((selector.numStagedReceives(channel) == 0 || channel.hasBytesBuffered()) && --retries > 0);
         assertTrue("No staged receives after 100 attempts", selector.numStagedReceives(channel) > 0);
         // We want to return without any bytes buffered to ensure that channel will be closed after idle time
         assertFalse("Channel has bytes buffered", channel.hasBytesBuffered());
@@ -533,7 +538,7 @@ private void verifyCloseOldestConnectionWithStagedReceives(int maxStagedReceives
     public void testMuteOnOOM() throws Exception {
         //clean up default selector, replace it with one that uses a finite mem pool
         selector.close();
-        MemoryPool pool = new SimpleMemoryPool(900, 900, false, null);
+        MemoryPool pool = new SimpleMemoryPool(900, 900, false, null, sensor);
         selector = new Selector(NetworkReceive.UNLIMITED, 5000, metrics, time, "MetricGroup",
             new HashMap(), true, false, channelBuilder, pool, new LogContext());
 
@@ -688,6 +693,56 @@ public void testInboundConnectionsCountInConnectionCreationMetric() throws Excep
         assertEquals((double) conns, getMetric("connection-count").metricValue());
     }
 
+    @SuppressWarnings("unchecked")
+    @Test
+    public void testLowestPriorityChannel() throws Exception {
+        int conns = 5;
+        InetSocketAddress addr = new InetSocketAddress("localhost", server.port);
+        for (int i = 0; i < conns; i++) {
+            connect(String.valueOf(i), addr);
+        }
+        assertNotNull(selector.lowestPriorityChannel());
+        for (int i = conns - 1; i >= 0; i--) {
+            if (i != 2)
+              assertEquals("", blockingRequest(String.valueOf(i), ""));
+            time.sleep(10);
+        }
+        assertEquals("2", selector.lowestPriorityChannel().id());
+
+        Field field = Selector.class.getDeclaredField("closingChannels");
+        field.setAccessible(true);
+        Map closingChannels = (Map) field.get(selector);
+        closingChannels.put("3", selector.channel("3"));
+        assertEquals("3", selector.lowestPriorityChannel().id());
+        closingChannels.remove("3");
+
+        for (int i = 0; i < conns; i++) {
+            selector.close(String.valueOf(i));
+        }
+        assertNull(selector.lowestPriorityChannel());
+    }
+
+    @Test
+    public void testMetricsCleanupOnSelectorClose() throws Exception {
+        Metrics metrics = new Metrics();
+        Selector selector = new ImmediatelyConnectingSelector(5000, metrics, time, "MetricGroup", channelBuilder, new LogContext()) {
+            @Override
+            public void close(String id) {
+                throw new RuntimeException();
+            }
+        };
+        assertTrue(metrics.metrics().size() > 1);
+        String id = "0";
+        selector.connect(id, new InetSocketAddress("localhost", server.port), BUFFER_SIZE, BUFFER_SIZE);
+
+        // Close the selector and ensure a RuntimeException has been throw
+        assertThrows(RuntimeException.class, selector::close);
+
+        // We should only have one remaining metric for kafka-metrics-count, which is a global metric
+        assertEquals(1, metrics.metrics().size());
+    }
+
+
     private String blockingRequest(String node, String s) throws IOException {
         selector.send(createSend(node, s));
         selector.poll(1000L);
@@ -762,7 +817,7 @@ protected void verifySelectorEmpty() throws Exception {
         verifySelectorEmpty(this.selector);
     }
 
-    private void verifySelectorEmpty(Selector selector) throws Exception {
+    public void verifySelectorEmpty(Selector selector) throws Exception {
         for (KafkaChannel channel : selector.channels()) {
             selector.close(channel.id());
             assertNull(channel.selectionKey().attachment());
diff --git a/clients/src/test/java/org/apache/kafka/common/network/SslSelectorTest.java b/clients/src/test/java/org/apache/kafka/common/network/SslSelectorTest.java
index 1f9739bf762d0..273d7a531baa4 100644
--- a/clients/src/test/java/org/apache/kafka/common/network/SslSelectorTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/network/SslSelectorTest.java
@@ -18,13 +18,21 @@
 
 import java.nio.channels.SelectionKey;
 import javax.net.ssl.SSLEngine;
+
+import org.apache.kafka.common.config.SecurityConfig;
 import org.apache.kafka.common.memory.MemoryPool;
 import org.apache.kafka.common.memory.SimpleMemoryPool;
+import org.apache.kafka.common.metrics.MetricConfig;
 import org.apache.kafka.common.metrics.Metrics;
+import org.apache.kafka.common.metrics.Sensor;
 import org.apache.kafka.common.security.auth.SecurityProtocol;
 import org.apache.kafka.common.security.ssl.SslFactory;
+import org.apache.kafka.common.security.ssl.mock.TestKeyManagerFactory;
+import org.apache.kafka.common.security.ssl.mock.TestProviderCreator;
+import org.apache.kafka.common.security.ssl.mock.TestTrustManagerFactory;
 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.test.TestCondition;
 import org.apache.kafka.test.TestSslUtils;
 import org.apache.kafka.test.TestUtils;
@@ -37,6 +45,7 @@
 import java.net.InetSocketAddress;
 import java.nio.channels.ServerSocketChannel;
 import java.nio.channels.SocketChannel;
+import java.security.Security;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
@@ -71,6 +80,7 @@ public void setUp() throws Exception {
         this.channelBuilder.configure(sslClientConfigs);
         this.metrics = new Metrics();
         this.selector = new Selector(5000, metrics, time, "MetricGroup", channelBuilder, new LogContext());
+        this.sensor = metrics.sensor("infoSensor", new MetricConfig().recordLevel(Sensor.RecordingLevel.INFO), Sensor.RecordingLevel.INFO);
     }
 
     @After
@@ -85,6 +95,50 @@ public SecurityProtocol securityProtocol() {
         return SecurityProtocol.PLAINTEXT;
     }
 
+    @Test
+    public void testConnectionWithCustomKeyManager() throws Exception {
+
+        TestProviderCreator testProviderCreator = new TestProviderCreator();
+
+        int requestSize = 100 * 1024;
+        final String node = "0";
+        String request = TestUtils.randomString(requestSize);
+
+        Map sslServerConfigs = TestSslUtils.createSslConfig(
+                TestKeyManagerFactory.ALGORITHM,
+                TestTrustManagerFactory.ALGORITHM
+        );
+        sslServerConfigs.put(SecurityConfig.SECURITY_PROVIDERS_CONFIG, testProviderCreator.getClass().getName());
+        EchoServer server = new EchoServer(SecurityProtocol.SSL, sslServerConfigs);
+        server.start();
+        Time time = new MockTime();
+        File trustStoreFile = new File(TestKeyManagerFactory.TestKeyManager.mockTrustStoreFile);
+        Map sslClientConfigs = TestSslUtils.createSslConfig(true, true, Mode.CLIENT, trustStoreFile, "client");
+
+        ChannelBuilder channelBuilder = new TestSslChannelBuilder(Mode.CLIENT);
+        channelBuilder.configure(sslClientConfigs);
+        Metrics metrics = new Metrics();
+        Selector selector = new Selector(5000, metrics, time, "MetricGroup", channelBuilder, new LogContext());
+
+        selector.connect(node, new InetSocketAddress("localhost", server.port), BUFFER_SIZE, BUFFER_SIZE);
+        while (!selector.connected().contains(node))
+            selector.poll(10000L);
+        while (!selector.isChannelReady(node))
+            selector.poll(10000L);
+
+        selector.send(createSend(node, request));
+
+        waitForBytesBuffered(selector, node);
+
+        selector.close(node);
+        super.verifySelectorEmpty(selector);
+
+        Security.removeProvider(testProviderCreator.getProvider().getName());
+        selector.close();
+        server.close();
+        metrics.close();
+    }
+
     @Test
     public void testDisconnectWithIntermediateBufferedBytes() throws Exception {
         int requestSize = 100 * 1024;
@@ -219,7 +273,7 @@ public void testRenegotiationFails() throws Exception {
     public void testMuteOnOOM() throws Exception {
         //clean up default selector, replace it with one that uses a finite mem pool
         selector.close();
-        MemoryPool pool = new SimpleMemoryPool(900, 900, false, null);
+        MemoryPool pool = new SimpleMemoryPool(900, 900, false, null, sensor);
         //the initial channel builder is for clients, we need a server one
         File trustStoreFile = File.createTempFile("truststore", ".jks");
         Map sslServerConfigs = TestSslUtils.createSslConfig(false, true, Mode.SERVER, trustStoreFile, "server");
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 118baecd61c61..e7dd0d3956102 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
@@ -16,6 +16,8 @@
  */
 package org.apache.kafka.common.network;
 
+import java.util.ArrayList;
+import java.util.Collection;
 import java.util.List;
 import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.config.SslConfigs;
@@ -30,12 +32,12 @@
 import org.apache.kafka.common.utils.Time;
 import org.apache.kafka.common.utils.Utils;
 import org.apache.kafka.test.TestCondition;
+import org.apache.kafka.test.TestSslUtils.SSLProvider;
 import org.apache.kafka.test.TestUtils;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 
-import javax.net.ssl.SSLContext;
 import javax.net.ssl.SSLEngine;
 import javax.net.ssl.SSLParameters;
 import java.io.ByteArrayOutputStream;
@@ -51,6 +53,8 @@
 import java.util.Map;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicLong;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
@@ -60,7 +64,13 @@
 /**
  * Tests for the SSL transport layer. These use a test harness that runs a simple socket server that echos back responses.
  */
+@RunWith(Parameterized.class)
 public class SslTransportLayerTest {
+    private final SSLProvider provider;
+
+    public SslTransportLayerTest(SSLProvider provider) {
+        this.provider = provider;
+    }
 
     private static final int BUFFER_SIZE = 4 * 1024;
     private static Time time = Time.SYSTEM;
@@ -76,8 +86,8 @@ public class SslTransportLayerTest {
     @Before
     public void setup() throws Exception {
         // Create certificates for use by client and server. Add server cert to client truststore and vice versa.
-        serverCertStores = new CertStores(true, "server",  "localhost");
-        clientCertStores = new CertStores(false, "client", "localhost");
+        serverCertStores = new CertStores(true, "server",  "localhost", provider);
+        clientCertStores = new CertStores(false, "client", "localhost", provider);
         sslServerConfigs = serverCertStores.getTrustingConfig(clientCertStores);
         sslClientConfigs = clientCertStores.getTrustingConfig(serverCertStores);
         this.channelBuilder = new SslChannelBuilder(Mode.CLIENT, null, false);
@@ -117,8 +127,8 @@ public void testValidEndpointIdentificationSanDns() throws Exception {
     @Test
     public void testValidEndpointIdentificationSanIp() throws Exception {
         String node = "0";
-        serverCertStores = new CertStores(true, "server", InetAddress.getByName("127.0.0.1"));
-        clientCertStores = new CertStores(false, "client", InetAddress.getByName("127.0.0.1"));
+        serverCertStores = new CertStores(true, "server", InetAddress.getByName("127.0.0.1"), provider);
+        clientCertStores = new CertStores(false, "client", InetAddress.getByName("127.0.0.1"), provider);
         sslServerConfigs = serverCertStores.getTrustingConfig(clientCertStores);
         sslClientConfigs = clientCertStores.getTrustingConfig(serverCertStores);
         server = createEchoServer(SecurityProtocol.SSL);
@@ -137,8 +147,8 @@ public void testValidEndpointIdentificationSanIp() throws Exception {
     @Test
     public void testValidEndpointIdentificationCN() throws Exception {
         String node = "0";
-        serverCertStores = new CertStores(true, "localhost");
-        clientCertStores = new CertStores(false, "localhost");
+        serverCertStores = new CertStores(true, "localhost", provider);
+        clientCertStores = new CertStores(false, "localhost", provider);
         sslServerConfigs = serverCertStores.getTrustingConfig(clientCertStores);
         sslClientConfigs = clientCertStores.getTrustingConfig(serverCertStores);
         server = createEchoServer(SecurityProtocol.SSL);
@@ -187,8 +197,8 @@ public void testClientEndpointNotValidated() throws Exception {
         String node = "0";
 
         // Create client certificate with an invalid hostname
-        clientCertStores = new CertStores(false, "non-existent.com");
-        serverCertStores = new CertStores(true, "localhost");
+        clientCertStores = new CertStores(false, "non-existent.com", provider);
+        serverCertStores = new CertStores(true, "localhost", provider);
         sslServerConfigs = serverCertStores.getTrustingConfig(clientCertStores);
         sslClientConfigs = clientCertStores.getTrustingConfig(serverCertStores);
 
@@ -222,8 +232,8 @@ protected TestSslTransportLayer newTransportLayer(String id, SelectionKey key, S
     @Test
     public void testInvalidEndpointIdentification() throws Exception {
         String node = "0";
-        serverCertStores = new CertStores(true, "server", "notahost");
-        clientCertStores = new CertStores(false, "client", "localhost");
+        serverCertStores = new CertStores(true, "server", "notahost", provider);
+        clientCertStores = new CertStores(false, "client", "localhost", provider);
         sslServerConfigs = serverCertStores.getTrustingConfig(clientCertStores);
         sslClientConfigs = clientCertStores.getTrustingConfig(serverCertStores);
         sslClientConfigs.put(SslConfigs.SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_CONFIG, "HTTPS");
@@ -242,8 +252,8 @@ public void testInvalidEndpointIdentification() throws Exception {
      */
     @Test
     public void testEndpointIdentificationDisabled() throws Exception {
-        serverCertStores = new CertStores(true, "server", "notahost");
-        clientCertStores = new CertStores(false, "client", "localhost");
+        serverCertStores = new CertStores(true, "server", "notahost", provider);
+        clientCertStores = new CertStores(false, "client", "localhost", provider);
         sslServerConfigs = serverCertStores.getTrustingConfig(clientCertStores);
         sslClientConfigs = clientCertStores.getTrustingConfig(serverCertStores);
 
@@ -543,13 +553,12 @@ public void testUnsupportedTLSVersion() throws Exception {
     @Test
     public void testUnsupportedCiphers() throws Exception {
         String node = "0";
-        SSLContext context = SSLContext.getInstance("TLSv1.2");
-        context.init(null, null, null);
-        String[] cipherSuites = context.getDefaultSSLParameters().getCipherSuites();
-        sslServerConfigs.put(SslConfigs.SSL_CIPHER_SUITES_CONFIG, Arrays.asList(cipherSuites[0]));
+        // Since Conscrypt does not support all the default JDK cipher suites, specifically pick two suites which are supported
+        // by both JDK and Conscrypt.
+        sslServerConfigs.put(SslConfigs.SSL_CIPHER_SUITES_CONFIG, Arrays.asList("TLS_ECDHE_ECDSA_WITH_AES_128_CBC_SHA"));
         server = createEchoServer(SecurityProtocol.SSL);
 
-        sslClientConfigs.put(SslConfigs.SSL_CIPHER_SUITES_CONFIG, Arrays.asList(cipherSuites[1]));
+        sslClientConfigs.put(SslConfigs.SSL_CIPHER_SUITES_CONFIG, Arrays.asList("TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256"));
         createSelector(sslClientConfigs);
         InetSocketAddress addr = new InetSocketAddress("localhost", server.port());
         selector.connect(node, addr, BUFFER_SIZE, BUFFER_SIZE);
@@ -888,6 +897,46 @@ public boolean conditionMet() {
         }, 5000, "All requests sent were not processed");
     }
 
+    /**
+     * Verifies that inter-broker listener with validation of truststore against keystore works
+     * with configs including mutual authentication and hostname verification.
+     */
+    @Test
+    public void testInterBrokerSslConfigValidation() throws Exception {
+        SecurityProtocol securityProtocol = SecurityProtocol.SSL;
+        sslServerConfigs.put(BrokerSecurityConfigs.SSL_CLIENT_AUTH_CONFIG, "required");
+        sslServerConfigs.put(SslConfigs.SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_CONFIG, "HTTPS");
+        sslServerConfigs.putAll(serverCertStores.keyStoreProps());
+        sslServerConfigs.putAll(serverCertStores.trustStoreProps());
+        sslClientConfigs.putAll(serverCertStores.keyStoreProps());
+        sslClientConfigs.putAll(serverCertStores.trustStoreProps());
+        TestSecurityConfig config = new TestSecurityConfig(sslServerConfigs);
+        ListenerName listenerName = ListenerName.forSecurityProtocol(securityProtocol);
+        ChannelBuilder serverChannelBuilder = ChannelBuilders.serverChannelBuilder(listenerName,
+                true, securityProtocol, config, null, null, time);
+        server = new NioEchoServer(listenerName, securityProtocol, config,
+                "localhost", serverChannelBuilder, null, time);
+        server.start();
+
+        this.selector = createSelector(sslClientConfigs, null, null, null);
+        InetSocketAddress addr = new InetSocketAddress("localhost", server.port());
+        selector.connect("0", addr, BUFFER_SIZE, BUFFER_SIZE);
+        NetworkTestUtils.checkClientConnection(selector, "0", 100, 10);
+    }
+
+    /**
+     * Verifies that inter-broker listener with validation of truststore against keystore
+     * fails if certs from keystore are not trusted.
+     */
+    @Test(expected = KafkaException.class)
+    public void testInterBrokerSslConfigValidationFailure() throws Exception {
+        SecurityProtocol securityProtocol = SecurityProtocol.SSL;
+        sslServerConfigs.put(BrokerSecurityConfigs.SSL_CLIENT_AUTH_CONFIG, "required");
+        TestSecurityConfig config = new TestSecurityConfig(sslServerConfigs);
+        ListenerName listenerName = ListenerName.forSecurityProtocol(securityProtocol);
+        ChannelBuilders.serverChannelBuilder(listenerName, true, securityProtocol, config, null, null, time);
+    }
+
     /**
      * Tests reconfiguration of server keystore. Verifies that existing connections continue
      * to work with old keystore and new connections work with new keystore.
@@ -910,7 +959,7 @@ public void testServerKeystoreDynamicUpdate() throws Exception {
         oldClientSelector.connect(oldNode, addr, BUFFER_SIZE, BUFFER_SIZE);
         NetworkTestUtils.checkClientConnection(selector, oldNode, 100, 10);
 
-        CertStores newServerCertStores = new CertStores(true, "server", "localhost");
+        CertStores newServerCertStores = new CertStores(true, "server", "localhost", provider);
         Map newKeystoreConfigs = newServerCertStores.keyStoreProps();
         assertTrue("SslChannelBuilder not reconfigurable", serverChannelBuilder instanceof ListenerReconfigurable);
         ListenerReconfigurable reconfigurableBuilder = (ListenerReconfigurable) serverChannelBuilder;
@@ -931,7 +980,7 @@ public void testServerKeystoreDynamicUpdate() throws Exception {
         // Verify that old client continues to work
         NetworkTestUtils.checkClientConnection(oldClientSelector, oldNode, 100, 10);
 
-        CertStores invalidCertStores = new CertStores(true, "server", "127.0.0.1");
+        CertStores invalidCertStores = new CertStores(true, "server", "127.0.0.1", provider);
         Map  invalidConfigs = invalidCertStores.getTrustingConfig(clientCertStores);
         verifyInvalidReconfigure(reconfigurableBuilder, invalidConfigs, "keystore with different SubjectAltName");
 
@@ -970,7 +1019,7 @@ public void testServerTruststoreDynamicUpdate() throws Exception {
         oldClientSelector.connect(oldNode, addr, BUFFER_SIZE, BUFFER_SIZE);
         NetworkTestUtils.checkClientConnection(selector, oldNode, 100, 10);
 
-        CertStores newClientCertStores = new CertStores(true, "client", "localhost");
+        CertStores newClientCertStores = new CertStores(true, "client", "localhost", provider);
         sslClientConfigs = newClientCertStores.getTrustingConfig(serverCertStores);
         Map newTruststoreConfigs = newClientCertStores.trustStoreProps();
         assertTrue("SslChannelBuilder not reconfigurable", serverChannelBuilder instanceof ListenerReconfigurable);
@@ -1179,4 +1228,11 @@ int updateAndGet(int actualSize, boolean update) {
             }
         }
     }
+
+    @Parameterized.Parameters(name = "SSLProvider={0}")
+    public static Collection data() {
+        Collection p = new ArrayList<>();
+        p.add(new Object[]{SSLProvider.DEFAULT});
+        return p;
+    }
 }
diff --git a/clients/src/test/java/org/apache/kafka/common/protocol/MessageTestUtil.java b/clients/src/test/java/org/apache/kafka/common/protocol/MessageTestUtil.java
new file mode 100644
index 0000000000000..a175ccd7c14e1
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/protocol/MessageTestUtil.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.common.protocol;
+
+import java.nio.ByteBuffer;
+
+public final class MessageTestUtil {
+    public static int messageSize(Message message, short version) {
+        return message.size(new ObjectSerializationCache(), version);
+    }
+
+    public static ByteBuffer messageToByteBuffer(Message message, short version) {
+        ObjectSerializationCache cache = new ObjectSerializationCache();
+        int size = message.size(cache, version);
+        ByteBuffer bytes = ByteBuffer.allocate(size);
+        message.write(new ByteBufferAccessor(bytes), cache, version);
+        bytes.flip();
+        return bytes;
+    }
+
+    public static void messageFromByteBuffer(ByteBuffer bytes, Message message, short version) {
+        message.read(new ByteBufferAccessor(bytes.duplicate()), version);
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/protocol/MessageUtilTest.java b/clients/src/test/java/org/apache/kafka/common/protocol/MessageUtilTest.java
index 8620a1b03dc8d..7494d66268934 100755
--- a/clients/src/test/java/org/apache/kafka/common/protocol/MessageUtilTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/protocol/MessageUtilTest.java
@@ -21,34 +21,16 @@
 import org.junit.Test;
 import org.junit.rules.Timeout;
 
-import java.nio.charset.StandardCharsets;
+import java.nio.ByteBuffer;
 import java.util.Arrays;
 
+import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
 
 public final class MessageUtilTest {
     @Rule
     final public Timeout globalTimeout = Timeout.millis(120000);
 
-    @Test
-    public void testSimpleUtf8Lengths() {
-        validateUtf8Length("");
-        validateUtf8Length("abc");
-        validateUtf8Length("This is a test string.");
-    }
-
-    @Test
-    public void testMultibyteUtf8Lengths() {
-        validateUtf8Length("A\u00ea\u00f1\u00fcC");
-        validateUtf8Length("\ud801\udc00");
-        validateUtf8Length("M\u00fcO");
-    }
-
-    private void validateUtf8Length(String string) {
-        byte[] arr = string.getBytes(StandardCharsets.UTF_8);
-        assertEquals(arr.length, MessageUtil.serializedUtf8Length(string));
-    }
-
     @Test
     public void testDeepToString() {
         assertEquals("[1, 2, 3]",
@@ -56,4 +38,12 @@ public void testDeepToString() {
         assertEquals("[foo]",
             MessageUtil.deepToString(Arrays.asList("foo").iterator()));
     }
+
+    @Test
+    public void testByteBufferToArray() {
+        assertArrayEquals(new byte[] {1, 2, 3},
+                MessageUtil.byteBufferToArray(ByteBuffer.wrap(new byte[] {1, 2, 3})));
+        assertArrayEquals(new byte[] {},
+                MessageUtil.byteBufferToArray(ByteBuffer.wrap(new byte[] {})));
+    }
 }
diff --git a/clients/src/test/java/org/apache/kafka/common/protocol/types/ProtocolSerializationTest.java b/clients/src/test/java/org/apache/kafka/common/protocol/types/ProtocolSerializationTest.java
index 6e9341affcdfa..e1c85671c51f5 100644
--- a/clients/src/test/java/org/apache/kafka/common/protocol/types/ProtocolSerializationTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/protocol/types/ProtocolSerializationTest.java
@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.common.protocol.types;
 
+import org.apache.kafka.common.utils.ByteUtils;
 import org.junit.Before;
 import org.junit.Test;
 
@@ -26,6 +27,7 @@
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.fail;
 
 public class ProtocolSerializationTest {
@@ -43,11 +45,17 @@ public void setup() {
                                  new Field("varint", Type.VARINT),
                                  new Field("varlong", Type.VARLONG),
                                  new Field("string", Type.STRING),
+                                 new Field("compact_string", Type.COMPACT_STRING),
                                  new Field("nullable_string", Type.NULLABLE_STRING),
+                                 new Field("compact_nullable_string", Type.COMPACT_NULLABLE_STRING),
                                  new Field("bytes", Type.BYTES),
+                                 new Field("compact_bytes", Type.COMPACT_BYTES),
                                  new Field("nullable_bytes", Type.NULLABLE_BYTES),
+                                 new Field("compact_nullable_bytes", Type.COMPACT_NULLABLE_BYTES),
                                  new Field("array", new ArrayOf(Type.INT32)),
+                                 new Field("compact_array", new CompactArrayOf(Type.INT32)),
                                  new Field("null_array", ArrayOf.nullable(Type.INT32)),
+                                 new Field("compact_null_array", CompactArrayOf.nullable(Type.INT32)),
                                  new Field("struct", new Schema(new Field("field", new ArrayOf(Type.INT32)))));
         this.struct = new Struct(this.schema).set("boolean", true)
                                              .set("int8", (byte) 1)
@@ -57,41 +65,69 @@ public void setup() {
                                              .set("varint", 300)
                                              .set("varlong", 500L)
                                              .set("string", "1")
+                                             .set("compact_string", "1")
                                              .set("nullable_string", null)
+                                             .set("compact_nullable_string", null)
                                              .set("bytes", ByteBuffer.wrap("1".getBytes()))
+                                             .set("compact_bytes", ByteBuffer.wrap("1".getBytes()))
                                              .set("nullable_bytes", null)
+                                             .set("compact_nullable_bytes", null)
                                              .set("array", new Object[] {1})
-                                             .set("null_array", null);
+                                             .set("compact_array", new Object[] {1})
+                                             .set("null_array", null)
+                                             .set("compact_null_array", null);
         this.struct.set("struct", this.struct.instance("struct").set("field", new Object[] {1, 2, 3}));
     }
 
     @Test
     public void testSimple() {
-        check(Type.BOOLEAN, false);
-        check(Type.BOOLEAN, true);
-        check(Type.INT8, (byte) -111);
-        check(Type.INT16, (short) -11111);
-        check(Type.INT32, -11111111);
-        check(Type.INT64, -11111111111L);
-        check(Type.STRING, "");
-        check(Type.STRING, "hello");
-        check(Type.STRING, "A\u00ea\u00f1\u00fcC");
-        check(Type.NULLABLE_STRING, null);
-        check(Type.NULLABLE_STRING, "");
-        check(Type.NULLABLE_STRING, "hello");
-        check(Type.BYTES, ByteBuffer.allocate(0));
-        check(Type.BYTES, ByteBuffer.wrap("abcd".getBytes()));
-        check(Type.NULLABLE_BYTES, null);
-        check(Type.NULLABLE_BYTES, ByteBuffer.allocate(0));
-        check(Type.NULLABLE_BYTES, ByteBuffer.wrap("abcd".getBytes()));
-        check(Type.VARINT, Integer.MAX_VALUE);
-        check(Type.VARINT, Integer.MIN_VALUE);
-        check(Type.VARLONG, Long.MAX_VALUE);
-        check(Type.VARLONG, Long.MIN_VALUE);
-        check(new ArrayOf(Type.INT32), new Object[] {1, 2, 3, 4});
-        check(new ArrayOf(Type.STRING), new Object[] {});
-        check(new ArrayOf(Type.STRING), new Object[] {"hello", "there", "beautiful"});
-        check(ArrayOf.nullable(Type.STRING), null);
+        check(Type.BOOLEAN, false, "BOOLEAN");
+        check(Type.BOOLEAN, true, "BOOLEAN");
+        check(Type.INT8, (byte) -111, "INT8");
+        check(Type.INT16, (short) -11111, "INT16");
+        check(Type.INT32, -11111111, "INT32");
+        check(Type.INT64, -11111111111L, "INT64");
+        check(Type.STRING, "", "STRING");
+        check(Type.STRING, "hello", "STRING");
+        check(Type.STRING, "A\u00ea\u00f1\u00fcC", "STRING");
+        check(Type.COMPACT_STRING, "", "COMPACT_STRING");
+        check(Type.COMPACT_STRING, "hello", "COMPACT_STRING");
+        check(Type.COMPACT_STRING, "A\u00ea\u00f1\u00fcC", "COMPACT_STRING");
+        check(Type.NULLABLE_STRING, null, "NULLABLE_STRING");
+        check(Type.NULLABLE_STRING, "", "NULLABLE_STRING");
+        check(Type.NULLABLE_STRING, "hello", "NULLABLE_STRING");
+        check(Type.COMPACT_NULLABLE_STRING, null, "COMPACT_NULLABLE_STRING");
+        check(Type.COMPACT_NULLABLE_STRING, "", "COMPACT_NULLABLE_STRING");
+        check(Type.COMPACT_NULLABLE_STRING, "hello", "COMPACT_NULLABLE_STRING");
+        check(Type.BYTES, ByteBuffer.allocate(0), "BYTES");
+        check(Type.BYTES, ByteBuffer.wrap("abcd".getBytes()), "BYTES");
+        check(Type.COMPACT_BYTES, ByteBuffer.allocate(0), "COMPACT_BYTES");
+        check(Type.COMPACT_BYTES, ByteBuffer.wrap("abcd".getBytes()), "COMPACT_BYTES");
+        check(Type.NULLABLE_BYTES, null, "NULLABLE_BYTES");
+        check(Type.NULLABLE_BYTES, ByteBuffer.allocate(0), "NULLABLE_BYTES");
+        check(Type.NULLABLE_BYTES, ByteBuffer.wrap("abcd".getBytes()), "NULLABLE_BYTES");
+        check(Type.COMPACT_NULLABLE_BYTES, null, "COMPACT_NULLABLE_BYTES");
+        check(Type.COMPACT_NULLABLE_BYTES, ByteBuffer.allocate(0), "COMPACT_NULLABLE_BYTES");
+        check(Type.COMPACT_NULLABLE_BYTES, ByteBuffer.wrap("abcd".getBytes()),
+                "COMPACT_NULLABLE_BYTES");
+        check(Type.VARINT, Integer.MAX_VALUE, "VARINT");
+        check(Type.VARINT, Integer.MIN_VALUE, "VARINT");
+        check(Type.VARLONG, Long.MAX_VALUE, "VARLONG");
+        check(Type.VARLONG, Long.MIN_VALUE, "VARLONG");
+        check(new ArrayOf(Type.INT32), new Object[] {1, 2, 3, 4}, "ARRAY(INT32)");
+        check(new ArrayOf(Type.STRING), new Object[] {}, "ARRAY(STRING)");
+        check(new ArrayOf(Type.STRING), new Object[] {"hello", "there", "beautiful"},
+                "ARRAY(STRING)");
+        check(new CompactArrayOf(Type.INT32), new Object[] {1, 2, 3, 4},
+                "COMPACT_ARRAY(INT32)");
+        check(new CompactArrayOf(Type.COMPACT_STRING), new Object[] {},
+                "COMPACT_ARRAY(COMPACT_STRING)");
+        check(new CompactArrayOf(Type.COMPACT_STRING),
+                new Object[] {"hello", "there", "beautiful"},
+                "COMPACT_ARRAY(COMPACT_STRING)");
+        check(ArrayOf.nullable(Type.STRING), null, "ARRAY(STRING)");
+        check(CompactArrayOf.nullable(Type.COMPACT_STRING), null,
+                "COMPACT_ARRAY(COMPACT_STRING)");
     }
 
     @Test
@@ -104,7 +140,7 @@ public void testNulls() {
                 if (!f.def.type.isNullable())
                     fail("Should not allow serialization of null value.");
             } catch (SchemaException e) {
-                assertFalse(f.def.type.isNullable());
+                assertFalse(f.toString() + " should not be nullable", f.def.type.isNullable());
             } finally {
                 this.struct.set(f, o);
             }
@@ -122,7 +158,9 @@ public void testDefault() {
     @Test
     public void testNullableDefault() {
         checkNullableDefault(Type.NULLABLE_BYTES, ByteBuffer.allocate(0));
+        checkNullableDefault(Type.COMPACT_NULLABLE_BYTES, ByteBuffer.allocate(0));
         checkNullableDefault(Type.NULLABLE_STRING, "default");
+        checkNullableDefault(Type.COMPACT_NULLABLE_STRING, "default");
     }
 
     private void checkNullableDefault(Type type, Object defaultValue) {
@@ -150,6 +188,24 @@ public void testReadArraySizeTooLarge() {
         }
     }
 
+    @Test
+    public void testReadCompactArraySizeTooLarge() {
+        Type type = new CompactArrayOf(Type.INT8);
+        int size = 10;
+        ByteBuffer invalidBuffer = ByteBuffer.allocate(
+            ByteUtils.sizeOfUnsignedVarint(Integer.MAX_VALUE) + size);
+        ByteUtils.writeUnsignedVarint(Integer.MAX_VALUE, invalidBuffer);
+        for (int i = 0; i < size; i++)
+            invalidBuffer.put((byte) i);
+        invalidBuffer.rewind();
+        try {
+            type.read(invalidBuffer);
+            fail("Array size not validated");
+        } catch (SchemaException e) {
+            // Expected exception
+        }
+    }
+
     @Test
     public void testReadNegativeArraySize() {
         Type type = new ArrayOf(Type.INT8);
@@ -167,6 +223,24 @@ public void testReadNegativeArraySize() {
         }
     }
 
+    @Test
+    public void testReadZeroCompactArraySize() {
+        Type type = new CompactArrayOf(Type.INT8);
+        int size = 10;
+        ByteBuffer invalidBuffer = ByteBuffer.allocate(
+            ByteUtils.sizeOfUnsignedVarint(0) + size);
+        ByteUtils.writeUnsignedVarint(0, invalidBuffer);
+        for (int i = 0; i < size; i++)
+            invalidBuffer.put((byte) i);
+        invalidBuffer.rewind();
+        try {
+            type.read(invalidBuffer);
+            fail("Array size not validated");
+        } catch (SchemaException e) {
+            // Expected exception
+        }
+    }
+
     @Test
     public void testReadStringSizeTooLarge() {
         byte[] stringBytes = "foo".getBytes();
@@ -258,12 +332,13 @@ private Object roundtrip(Type type, Object obj) {
         return read;
     }
 
-    private void check(Type type, Object obj) {
+    private void check(Type type, Object obj, String expectedTypeName) {
         Object result = roundtrip(type, obj);
         if (obj instanceof Object[]) {
             obj = Arrays.asList((Object[]) obj);
             result = Arrays.asList((Object[]) result);
         }
+        assertEquals(expectedTypeName, type.toString());
         assertEquals("The object read back should be the same as what was written.", obj, result);
     }
 
@@ -278,4 +353,71 @@ public void testStructEquals() {
         assertNotEquals(emptyStruct1, mostlyEmptyStruct);
         assertNotEquals(mostlyEmptyStruct, emptyStruct1);
     }
+
+    @Test
+    public void testReadIgnoringExtraDataAtTheEnd() {
+        Schema oldSchema = new Schema(new Field("field1", Type.NULLABLE_STRING), new Field("field2", Type.NULLABLE_STRING));
+        Schema newSchema = new Schema(new Field("field1", Type.NULLABLE_STRING));
+        String value = "foo bar baz";
+        Struct oldFormat = new Struct(oldSchema).set("field1", value).set("field2", "fine to ignore");
+        ByteBuffer buffer = ByteBuffer.allocate(oldSchema.sizeOf(oldFormat));
+        oldFormat.writeTo(buffer);
+        buffer.flip();
+        Struct newFormat = newSchema.read(buffer);
+        assertEquals(value, newFormat.get("field1"));
+    }
+
+    @Test
+    public void testReadWhenOptionalDataMissingAtTheEndIsTolerated() {
+        Schema oldSchema = new Schema(new Field("field1", Type.NULLABLE_STRING));
+        Schema newSchema = new Schema(
+                true,
+                new Field("field1", Type.NULLABLE_STRING),
+                new Field("field2", Type.NULLABLE_STRING, "", true, "default"),
+                new Field("field3", Type.NULLABLE_STRING, "", true, null),
+                new Field("field4", Type.NULLABLE_BYTES, "", true, ByteBuffer.allocate(0)),
+                new Field("field5", Type.INT64, "doc", true, Long.MAX_VALUE));
+        String value = "foo bar baz";
+        Struct oldFormat = new Struct(oldSchema).set("field1", value);
+        ByteBuffer buffer = ByteBuffer.allocate(oldSchema.sizeOf(oldFormat));
+        oldFormat.writeTo(buffer);
+        buffer.flip();
+        Struct newFormat = newSchema.read(buffer);
+        assertEquals(value, newFormat.get("field1"));
+        assertEquals("default", newFormat.get("field2"));
+        assertEquals(null, newFormat.get("field3"));
+        assertEquals(ByteBuffer.allocate(0), newFormat.get("field4"));
+        assertEquals(Long.MAX_VALUE, newFormat.get("field5"));
+    }
+
+    @Test
+    public void testReadWhenOptionalDataMissingAtTheEndIsNotTolerated() {
+        Schema oldSchema = new Schema(new Field("field1", Type.NULLABLE_STRING));
+        Schema newSchema = new Schema(
+                new Field("field1", Type.NULLABLE_STRING),
+                new Field("field2", Type.NULLABLE_STRING, "", true, "default"));
+        String value = "foo bar baz";
+        Struct oldFormat = new Struct(oldSchema).set("field1", value);
+        ByteBuffer buffer = ByteBuffer.allocate(oldSchema.sizeOf(oldFormat));
+        oldFormat.writeTo(buffer);
+        buffer.flip();
+        SchemaException e = assertThrows(SchemaException.class, () -> newSchema.read(buffer));
+        e.getMessage().contains("Error reading field 'field2': java.nio.BufferUnderflowException");
+    }
+
+    @Test
+    public void testReadWithMissingNonOptionalExtraDataAtTheEnd() {
+        Schema oldSchema = new Schema(new Field("field1", Type.NULLABLE_STRING));
+        Schema newSchema = new Schema(
+                true,
+                new Field("field1", Type.NULLABLE_STRING),
+                new Field("field2", Type.NULLABLE_STRING));
+        String value = "foo bar baz";
+        Struct oldFormat = new Struct(oldSchema).set("field1", value);
+        ByteBuffer buffer = ByteBuffer.allocate(oldSchema.sizeOf(oldFormat));
+        oldFormat.writeTo(buffer);
+        buffer.flip();
+        SchemaException e = assertThrows(SchemaException.class, () -> newSchema.read(buffer));
+        e.getMessage().contains("Missing value for field 'field2' which has no default value");
+    }
 }
diff --git a/clients/src/test/java/org/apache/kafka/common/protocol/types/RawTaggedFieldWriterTest.java b/clients/src/test/java/org/apache/kafka/common/protocol/types/RawTaggedFieldWriterTest.java
new file mode 100644
index 0000000000000..9ea38f9b84805
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/protocol/types/RawTaggedFieldWriterTest.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.common.protocol.types;
+
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.Timeout;
+
+import java.nio.ByteBuffer;
+import java.util.Arrays;
+import java.util.List;
+
+import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.fail;
+
+public class RawTaggedFieldWriterTest {
+    @Rule
+    final public Timeout globalTimeout = Timeout.millis(120000);
+
+    @Test
+    public void testWritingZeroRawTaggedFields() {
+        RawTaggedFieldWriter writer = RawTaggedFieldWriter.forFields(null);
+        assertEquals(0, writer.numFields());
+        ByteBufferAccessor accessor = new ByteBufferAccessor(ByteBuffer.allocate(0));
+        writer.writeRawTags(accessor, Integer.MAX_VALUE);
+    }
+
+    @Test
+    public void testWritingSeveralRawTaggedFields() {
+        List tags = Arrays.asList(
+            new RawTaggedField(2, new byte[] {0x1, 0x2, 0x3}),
+            new RawTaggedField(5, new byte[] {0x4, 0x5})
+        );
+        RawTaggedFieldWriter writer = RawTaggedFieldWriter.forFields(tags);
+        assertEquals(2, writer.numFields());
+        byte[] arr = new byte[9];
+        ByteBufferAccessor accessor = new ByteBufferAccessor(ByteBuffer.wrap(arr));
+        writer.writeRawTags(accessor, 1);
+        assertArrayEquals(new byte[] {0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0}, arr);
+        writer.writeRawTags(accessor, 3);
+        assertArrayEquals(new byte[] {0x2, 0x3, 0x1, 0x2, 0x3, 0x0, 0x0, 0x0, 0x0}, arr);
+        writer.writeRawTags(accessor, 7);
+        assertArrayEquals(new byte[] {0x2, 0x3, 0x1, 0x2, 0x3, 0x5, 0x2, 0x4, 0x5}, arr);
+        writer.writeRawTags(accessor, Integer.MAX_VALUE);
+        assertArrayEquals(new byte[] {0x2, 0x3, 0x1, 0x2, 0x3, 0x5, 0x2, 0x4, 0x5}, arr);
+    }
+
+    @Test
+    public void testInvalidNextDefinedTag() {
+        List tags = Arrays.asList(
+            new RawTaggedField(2, new byte[] {0x1, 0x2, 0x3}),
+            new RawTaggedField(5, new byte[] {0x4, 0x5, 0x6}),
+            new RawTaggedField(7, new byte[] {0x0})
+        );
+        RawTaggedFieldWriter writer = RawTaggedFieldWriter.forFields(tags);
+        assertEquals(3, writer.numFields());
+        try {
+            writer.writeRawTags(new ByteBufferAccessor(ByteBuffer.allocate(1024)), 2);
+            fail("expected to get RuntimeException");
+        } catch (RuntimeException e) {
+            assertEquals("Attempted to use tag 2 as an undefined tag.", e.getMessage());
+        }
+    }
+
+    @Test
+    public void testOutOfOrderTags() {
+        List tags = Arrays.asList(
+            new RawTaggedField(5, new byte[] {0x4, 0x5, 0x6}),
+            new RawTaggedField(2, new byte[] {0x1, 0x2, 0x3}),
+            new RawTaggedField(7, new byte[] {0x0 })
+        );
+        RawTaggedFieldWriter writer = RawTaggedFieldWriter.forFields(tags);
+        assertEquals(3, writer.numFields());
+        try {
+            writer.writeRawTags(new ByteBufferAccessor(ByteBuffer.allocate(1024)), 8);
+            fail("expected to get RuntimeException");
+        } catch (RuntimeException e) {
+            assertEquals("Invalid raw tag field list: tag 2 comes after tag 5, but is " +
+                "not higher than it.", e.getMessage());
+        }
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/protocol/types/StructTest.java b/clients/src/test/java/org/apache/kafka/common/protocol/types/StructTest.java
new file mode 100644
index 0000000000000..44cd09bc19d01
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/protocol/types/StructTest.java
@@ -0,0 +1,84 @@
+/*
+ * 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.protocol.types;
+
+import org.junit.Test;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotEquals;
+
+public class StructTest {
+    private static final Schema FLAT_STRUCT_SCHEMA = new Schema(
+        new Field.Int8("int8", ""),
+        new Field.Int16("int16", ""),
+        new Field.Int32("int32", ""),
+        new Field.Int64("int64", ""),
+        new Field.Bool("boolean", ""),
+        new Field.Str("string", ""));
+
+    private static final Schema ARRAY_SCHEMA = new Schema(new Field.Array("array", new ArrayOf(Type.INT8), ""));
+    private static final Schema NESTED_CHILD_SCHEMA = new Schema(
+            new Field.Int8("int8", ""));
+    private static final Schema NESTED_SCHEMA = new Schema(
+            new Field.Array("array", ARRAY_SCHEMA, ""),
+            new Field("nested", NESTED_CHILD_SCHEMA, ""));
+
+    @Test
+    public void testEquals() {
+        Struct struct1 = new Struct(FLAT_STRUCT_SCHEMA)
+                .set("int8", (byte) 12)
+                .set("int16", (short) 12)
+                .set("int32", 12)
+                .set("int64", (long) 12)
+                .set("boolean", true)
+                .set("string", "foobar");
+        Struct struct2 = new Struct(FLAT_STRUCT_SCHEMA)
+                .set("int8", (byte) 12)
+                .set("int16", (short) 12)
+                .set("int32", 12)
+                .set("int64", (long) 12)
+                .set("boolean", true)
+                .set("string", "foobar");
+        Struct struct3 = new Struct(FLAT_STRUCT_SCHEMA)
+                .set("int8", (byte) 12)
+                .set("int16", (short) 12)
+                .set("int32", 12)
+                .set("int64", (long) 12)
+                .set("boolean", true)
+                .set("string", "mismatching string");
+
+        assertEquals(struct1, struct2);
+        assertNotEquals(struct1, struct3);
+
+        Object[] array = {(byte) 1, (byte) 2};
+        struct1 = new Struct(NESTED_SCHEMA)
+                .set("array", array)
+                .set("nested", new Struct(NESTED_CHILD_SCHEMA).set("int8", (byte) 12));
+        Object[] array2 = {(byte) 1, (byte) 2};
+        struct2 = new Struct(NESTED_SCHEMA)
+                .set("array", array2)
+                .set("nested", new Struct(NESTED_CHILD_SCHEMA).set("int8", (byte) 12));
+        Object[] array3 = {(byte) 1, (byte) 2, (byte) 3};
+        struct3 = new Struct(NESTED_SCHEMA)
+                .set("array", array3)
+                .set("nested", new Struct(NESTED_CHILD_SCHEMA).set("int8", (byte) 13));
+
+        assertEquals(struct1, struct2);
+        assertNotEquals(struct1, struct3);
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/record/AbstractLegacyRecordBatchTest.java b/clients/src/test/java/org/apache/kafka/common/record/AbstractLegacyRecordBatchTest.java
index fe6ffabaf61eb..87811b8bb0d80 100644
--- a/clients/src/test/java/org/apache/kafka/common/record/AbstractLegacyRecordBatchTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/record/AbstractLegacyRecordBatchTest.java
@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.common.record;
 
+import org.apache.kafka.common.InvalidRecordException;
 import org.apache.kafka.common.record.AbstractLegacyRecordBatch.ByteBufferLegacyRecordBatch;
 import org.apache.kafka.common.utils.Utils;
 import org.junit.Test;
diff --git a/clients/src/test/java/org/apache/kafka/common/record/ByteBufferLogInputStreamTest.java b/clients/src/test/java/org/apache/kafka/common/record/ByteBufferLogInputStreamTest.java
index 3745006a58304..063e188dfd353 100644
--- a/clients/src/test/java/org/apache/kafka/common/record/ByteBufferLogInputStreamTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/record/ByteBufferLogInputStreamTest.java
@@ -19,7 +19,6 @@
 import org.apache.kafka.common.errors.CorruptRecordException;
 import org.junit.Test;
 
-import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.util.Iterator;
 
@@ -56,7 +55,7 @@ public void iteratorIgnoresIncompleteEntries() {
     }
 
     @Test(expected = CorruptRecordException.class)
-    public void iteratorRaisesOnTooSmallRecords() throws IOException {
+    public void iteratorRaisesOnTooSmallRecords() {
         ByteBuffer buffer = ByteBuffer.allocate(1024);
         MemoryRecordsBuilder builder = MemoryRecords.builder(buffer, CompressionType.NONE, TimestampType.CREATE_TIME, 0L);
         builder.append(15L, "a".getBytes(), "1".getBytes());
@@ -79,7 +78,7 @@ public void iteratorRaisesOnTooSmallRecords() throws IOException {
     }
 
     @Test(expected = CorruptRecordException.class)
-    public void iteratorRaisesOnInvalidMagic() throws IOException {
+    public void iteratorRaisesOnInvalidMagic() {
         ByteBuffer buffer = ByteBuffer.allocate(1024);
         MemoryRecordsBuilder builder = MemoryRecords.builder(buffer, CompressionType.NONE, TimestampType.CREATE_TIME, 0L);
         builder.append(15L, "a".getBytes(), "1".getBytes());
@@ -102,7 +101,7 @@ public void iteratorRaisesOnInvalidMagic() throws IOException {
     }
 
     @Test(expected = CorruptRecordException.class)
-    public void iteratorRaisesOnTooLargeRecords() throws IOException {
+    public void iteratorRaisesOnTooLargeRecords() {
         ByteBuffer buffer = ByteBuffer.allocate(1024);
         MemoryRecordsBuilder builder = MemoryRecords.builder(buffer, CompressionType.NONE, TimestampType.CREATE_TIME, 0L);
         builder.append(15L, "a".getBytes(), "1".getBytes());
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 ab8cbb7ee2cd7..2d42c03a1a276 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
@@ -16,6 +16,8 @@
  */
 package org.apache.kafka.common.record;
 
+import org.apache.kafka.common.InvalidRecordException;
+import org.apache.kafka.common.errors.CorruptRecordException;
 import org.apache.kafka.common.header.Header;
 import org.apache.kafka.common.header.internals.RecordHeader;
 import org.apache.kafka.common.utils.CloseableIterator;
@@ -173,7 +175,7 @@ public void testSizeInBytes() {
         assertEquals(actualSize, DefaultRecordBatch.sizeInBytes(Arrays.asList(records)));
     }
 
-    @Test(expected = InvalidRecordException.class)
+    @Test(expected = CorruptRecordException.class)
     public void testInvalidRecordSize() {
         MemoryRecords records = MemoryRecords.withRecords(RecordBatch.MAGIC_VALUE_V2, 0L,
                 CompressionType.NONE, TimestampType.CREATE_TIME,
@@ -233,7 +235,7 @@ public void testInvalidRecordCountTooLittleCompressedV2() {
         }
     }
 
-    @Test(expected = InvalidRecordException.class)
+    @Test(expected = CorruptRecordException.class)
     public void testInvalidCrc() {
         MemoryRecords records = MemoryRecords.withRecords(RecordBatch.MAGIC_VALUE_V2, 0L,
                 CompressionType.NONE, TimestampType.CREATE_TIME,
@@ -375,6 +377,32 @@ public void testStreamingIteratorConsistency() {
         }
     }
 
+    @Test
+    public void testSkipKeyValueIteratorCorrectness() {
+        Header[] headers = {new RecordHeader("k1", "v1".getBytes()), new RecordHeader("k2", "v2".getBytes())};
+
+        MemoryRecords records = MemoryRecords.withRecords(RecordBatch.MAGIC_VALUE_V2, 0L,
+            CompressionType.LZ4, TimestampType.CREATE_TIME,
+            new SimpleRecord(1L, "a".getBytes(), "1".getBytes()),
+            new SimpleRecord(2L, "b".getBytes(), "2".getBytes()),
+            new SimpleRecord(3L, "c".getBytes(), "3".getBytes()),
+            new SimpleRecord(1000L, "abc".getBytes(), "0".getBytes()),
+            new SimpleRecord(9999L, "abc".getBytes(), "0".getBytes(), headers)
+            );
+        DefaultRecordBatch batch = new DefaultRecordBatch(records.buffer());
+        try (CloseableIterator streamingIterator = batch.skipKeyValueIterator(BufferSupplier.NO_CACHING)) {
+            assertEquals(Arrays.asList(
+                new PartialDefaultRecord(9, (byte) 0, 0L, 1L, -1, 1, 1),
+                new PartialDefaultRecord(9, (byte) 0, 1L, 2L, -1, 1, 1),
+                new PartialDefaultRecord(9, (byte) 0, 2L, 3L, -1, 1, 1),
+                new PartialDefaultRecord(12, (byte) 0, 3L, 1000L, -1, 3, 1),
+                new PartialDefaultRecord(25, (byte) 0, 4L, 9999L, -1, 3, 1)
+                ),
+                Utils.toList(streamingIterator)
+            );
+        }
+    }
+
     @Test
     public void testIncrementSequence() {
         assertEquals(10, DefaultRecordBatch.incrementSequence(5, 5));
@@ -382,6 +410,12 @@ public void testIncrementSequence() {
         assertEquals(4, DefaultRecordBatch.incrementSequence(Integer.MAX_VALUE - 5, 10));
     }
 
+    @Test
+    public void testDecrementSequence() {
+        assertEquals(0, DefaultRecordBatch.decrementSequence(5, 5));
+        assertEquals(Integer.MAX_VALUE, DefaultRecordBatch.decrementSequence(0, 1));
+    }
+
     private static DefaultRecordBatch recordsWithInvalidRecordCount(Byte magicValue, long timestamp,
                                               CompressionType codec, int invalidCount) {
         ByteBuffer buf = ByteBuffer.allocate(512);
diff --git a/clients/src/test/java/org/apache/kafka/common/record/DefaultRecordTest.java b/clients/src/test/java/org/apache/kafka/common/record/DefaultRecordTest.java
index 3ff73c9c67bd9..822b3b9c6aff0 100644
--- a/clients/src/test/java/org/apache/kafka/common/record/DefaultRecordTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/record/DefaultRecordTest.java
@@ -16,12 +16,16 @@
  */
 package org.apache.kafka.common.record;
 
+import org.apache.kafka.common.InvalidRecordException;
 import org.apache.kafka.common.header.Header;
 import org.apache.kafka.common.header.internals.RecordHeader;
+import org.apache.kafka.common.utils.ByteBufferInputStream;
 import org.apache.kafka.common.utils.ByteBufferOutputStream;
 import org.apache.kafka.common.utils.ByteUtils;
+import org.junit.Before;
 import org.junit.Test;
 
+import java.io.DataInputStream;
 import java.io.DataOutputStream;
 import java.io.IOException;
 import java.nio.ByteBuffer;
@@ -32,6 +36,13 @@
 
 public class DefaultRecordTest {
 
+    private byte[] skipArray;
+
+    @Before
+    public void setUp() {
+        skipArray = new byte[64];
+    }
+
     @Test
     public void testBasicSerde() throws IOException {
         Header[] headers = new Header[] {
@@ -152,6 +163,27 @@ public void testInvalidKeySize() {
         DefaultRecord.readFrom(buf, 0L, 0L, RecordBatch.NO_SEQUENCE, null);
     }
 
+    @Test(expected = InvalidRecordException.class)
+    public void testInvalidKeySizePartial() throws IOException {
+        byte attributes = 0;
+        long timestampDelta = 2;
+        int offsetDelta = 1;
+        int sizeOfBodyInBytes = 100;
+        int keySize = 105; // use a key size larger than the full message
+
+        ByteBuffer buf = ByteBuffer.allocate(sizeOfBodyInBytes + ByteUtils.sizeOfVarint(sizeOfBodyInBytes));
+        ByteUtils.writeVarint(sizeOfBodyInBytes, buf);
+        buf.put(attributes);
+        ByteUtils.writeVarlong(timestampDelta, buf);
+        ByteUtils.writeVarint(offsetDelta, buf);
+        ByteUtils.writeVarint(keySize, buf);
+        buf.position(buf.limit());
+
+        buf.flip();
+        DataInputStream inputStream = new DataInputStream(new ByteBufferInputStream(buf));
+        DefaultRecord.readPartiallyFrom(inputStream, skipArray, 0L, 0L, RecordBatch.NO_SEQUENCE, null);
+    }
+
     @Test(expected = InvalidRecordException.class)
     public void testInvalidValueSize() throws IOException {
         byte attributes = 0;
@@ -173,6 +205,210 @@ public void testInvalidValueSize() throws IOException {
         DefaultRecord.readFrom(buf, 0L, 0L, RecordBatch.NO_SEQUENCE, null);
     }
 
+    @Test(expected = InvalidRecordException.class)
+    public void testInvalidValueSizePartial() throws IOException {
+        byte attributes = 0;
+        long timestampDelta = 2;
+        int offsetDelta = 1;
+        int sizeOfBodyInBytes = 100;
+        int valueSize = 105; // use a value size larger than the full message
+
+        ByteBuffer buf = ByteBuffer.allocate(sizeOfBodyInBytes + ByteUtils.sizeOfVarint(sizeOfBodyInBytes));
+        ByteUtils.writeVarint(sizeOfBodyInBytes, buf);
+        buf.put(attributes);
+        ByteUtils.writeVarlong(timestampDelta, buf);
+        ByteUtils.writeVarint(offsetDelta, buf);
+        ByteUtils.writeVarint(-1, buf); // null key
+        ByteUtils.writeVarint(valueSize, buf);
+        buf.position(buf.limit());
+
+        buf.flip();
+        DataInputStream inputStream = new DataInputStream(new ByteBufferInputStream(buf));
+        DefaultRecord.readPartiallyFrom(inputStream, skipArray, 0L, 0L, RecordBatch.NO_SEQUENCE, null);
+    }
+
+    @Test(expected = InvalidRecordException.class)
+    public void testInvalidNumHeaders() throws IOException {
+        byte attributes = 0;
+        long timestampDelta = 2;
+        int offsetDelta = 1;
+        int sizeOfBodyInBytes = 100;
+
+        ByteBuffer buf = ByteBuffer.allocate(sizeOfBodyInBytes + ByteUtils.sizeOfVarint(sizeOfBodyInBytes));
+        ByteUtils.writeVarint(sizeOfBodyInBytes, buf);
+        buf.put(attributes);
+        ByteUtils.writeVarlong(timestampDelta, buf);
+        ByteUtils.writeVarint(offsetDelta, buf);
+        ByteUtils.writeVarint(-1, buf); // null key
+        ByteUtils.writeVarint(-1, buf); // null value
+        ByteUtils.writeVarint(-1, buf); // -1 num.headers, not allowed
+        buf.position(buf.limit());
+
+        buf.flip();
+        DefaultRecord.readFrom(buf, 0L, 0L, RecordBatch.NO_SEQUENCE, null);
+    }
+
+    @Test(expected = InvalidRecordException.class)
+    public void testInvalidNumHeadersPartial() throws IOException {
+        byte attributes = 0;
+        long timestampDelta = 2;
+        int offsetDelta = 1;
+        int sizeOfBodyInBytes = 100;
+
+        ByteBuffer buf = ByteBuffer.allocate(sizeOfBodyInBytes + ByteUtils.sizeOfVarint(sizeOfBodyInBytes));
+        ByteUtils.writeVarint(sizeOfBodyInBytes, buf);
+        buf.put(attributes);
+        ByteUtils.writeVarlong(timestampDelta, buf);
+        ByteUtils.writeVarint(offsetDelta, buf);
+        ByteUtils.writeVarint(-1, buf); // null key
+        ByteUtils.writeVarint(-1, buf); // null value
+        ByteUtils.writeVarint(-1, buf); // -1 num.headers, not allowed
+        buf.position(buf.limit());
+
+        buf.flip();
+        DataInputStream inputStream = new DataInputStream(new ByteBufferInputStream(buf));
+        DefaultRecord.readPartiallyFrom(inputStream, skipArray, 0L, 0L, RecordBatch.NO_SEQUENCE, null);
+    }
+
+    @Test(expected = StringIndexOutOfBoundsException.class)
+    public void testInvalidHeaderKey() {
+        byte attributes = 0;
+        long timestampDelta = 2;
+        int offsetDelta = 1;
+        int sizeOfBodyInBytes = 100;
+
+        ByteBuffer buf = ByteBuffer.allocate(sizeOfBodyInBytes + ByteUtils.sizeOfVarint(sizeOfBodyInBytes));
+        ByteUtils.writeVarint(sizeOfBodyInBytes, buf);
+        buf.put(attributes);
+        ByteUtils.writeVarlong(timestampDelta, buf);
+        ByteUtils.writeVarint(offsetDelta, buf);
+        ByteUtils.writeVarint(-1, buf); // null key
+        ByteUtils.writeVarint(-1, buf); // null value
+        ByteUtils.writeVarint(1, buf);
+        ByteUtils.writeVarint(105, buf); // header key too long
+        buf.position(buf.limit());
+
+        buf.flip();
+        DefaultRecord.readFrom(buf, 0L, 0L, RecordBatch.NO_SEQUENCE, null);
+    }
+
+    @Test(expected = InvalidRecordException.class)
+    public void testInvalidHeaderKeyPartial() throws IOException {
+        byte attributes = 0;
+        long timestampDelta = 2;
+        int offsetDelta = 1;
+        int sizeOfBodyInBytes = 100;
+
+        ByteBuffer buf = ByteBuffer.allocate(sizeOfBodyInBytes + ByteUtils.sizeOfVarint(sizeOfBodyInBytes));
+        ByteUtils.writeVarint(sizeOfBodyInBytes, buf);
+        buf.put(attributes);
+        ByteUtils.writeVarlong(timestampDelta, buf);
+        ByteUtils.writeVarint(offsetDelta, buf);
+        ByteUtils.writeVarint(-1, buf); // null key
+        ByteUtils.writeVarint(-1, buf); // null value
+        ByteUtils.writeVarint(1, buf);
+        ByteUtils.writeVarint(105, buf); // header key too long
+        buf.position(buf.limit());
+
+        buf.flip();
+        DataInputStream inputStream = new DataInputStream(new ByteBufferInputStream(buf));
+        DefaultRecord.readPartiallyFrom(inputStream, skipArray, 0L, 0L, RecordBatch.NO_SEQUENCE, null);
+    }
+
+    @Test(expected = InvalidRecordException.class)
+    public void testNullHeaderKey() {
+        byte attributes = 0;
+        long timestampDelta = 2;
+        int offsetDelta = 1;
+        int sizeOfBodyInBytes = 100;
+
+        ByteBuffer buf = ByteBuffer.allocate(sizeOfBodyInBytes + ByteUtils.sizeOfVarint(sizeOfBodyInBytes));
+        ByteUtils.writeVarint(sizeOfBodyInBytes, buf);
+        buf.put(attributes);
+        ByteUtils.writeVarlong(timestampDelta, buf);
+        ByteUtils.writeVarint(offsetDelta, buf);
+        ByteUtils.writeVarint(-1, buf); // null key
+        ByteUtils.writeVarint(-1, buf); // null value
+        ByteUtils.writeVarint(1, buf);
+        ByteUtils.writeVarint(-1, buf); // null header key not allowed
+        buf.position(buf.limit());
+
+        buf.flip();
+        DefaultRecord.readFrom(buf, 0L, 0L, RecordBatch.NO_SEQUENCE, null);
+    }
+
+    @Test(expected = InvalidRecordException.class)
+    public void testNullHeaderKeyPartial() throws IOException {
+        byte attributes = 0;
+        long timestampDelta = 2;
+        int offsetDelta = 1;
+        int sizeOfBodyInBytes = 100;
+
+        ByteBuffer buf = ByteBuffer.allocate(sizeOfBodyInBytes + ByteUtils.sizeOfVarint(sizeOfBodyInBytes));
+        ByteUtils.writeVarint(sizeOfBodyInBytes, buf);
+        buf.put(attributes);
+        ByteUtils.writeVarlong(timestampDelta, buf);
+        ByteUtils.writeVarint(offsetDelta, buf);
+        ByteUtils.writeVarint(-1, buf); // null key
+        ByteUtils.writeVarint(-1, buf); // null value
+        ByteUtils.writeVarint(1, buf);
+        ByteUtils.writeVarint(-1, buf); // null header key not allowed
+        buf.position(buf.limit());
+
+        buf.flip();
+        DataInputStream inputStream = new DataInputStream(new ByteBufferInputStream(buf));
+        DefaultRecord.readPartiallyFrom(inputStream, skipArray, 0L, 0L, RecordBatch.NO_SEQUENCE, null);
+    }
+
+    @Test(expected = InvalidRecordException.class)
+    public void testInvalidHeaderValue() {
+        byte attributes = 0;
+        long timestampDelta = 2;
+        int offsetDelta = 1;
+        int sizeOfBodyInBytes = 100;
+
+        ByteBuffer buf = ByteBuffer.allocate(sizeOfBodyInBytes + ByteUtils.sizeOfVarint(sizeOfBodyInBytes));
+        ByteUtils.writeVarint(sizeOfBodyInBytes, buf);
+        buf.put(attributes);
+        ByteUtils.writeVarlong(timestampDelta, buf);
+        ByteUtils.writeVarint(offsetDelta, buf);
+        ByteUtils.writeVarint(-1, buf); // null key
+        ByteUtils.writeVarint(-1, buf); // null value
+        ByteUtils.writeVarint(1, buf);
+        ByteUtils.writeVarint(1, buf);
+        buf.put((byte) 1);
+        ByteUtils.writeVarint(105, buf); // header value too long
+        buf.position(buf.limit());
+
+        buf.flip();
+        DefaultRecord.readFrom(buf, 0L, 0L, RecordBatch.NO_SEQUENCE, null);
+    }
+
+    @Test(expected = InvalidRecordException.class)
+    public void testInvalidHeaderValuePartial() throws IOException {
+        byte attributes = 0;
+        long timestampDelta = 2;
+        int offsetDelta = 1;
+        int sizeOfBodyInBytes = 100;
+
+        ByteBuffer buf = ByteBuffer.allocate(sizeOfBodyInBytes + ByteUtils.sizeOfVarint(sizeOfBodyInBytes));
+        ByteUtils.writeVarint(sizeOfBodyInBytes, buf);
+        buf.put(attributes);
+        ByteUtils.writeVarlong(timestampDelta, buf);
+        ByteUtils.writeVarint(offsetDelta, buf);
+        ByteUtils.writeVarint(-1, buf); // null key
+        ByteUtils.writeVarint(-1, buf); // null value
+        ByteUtils.writeVarint(1, buf);
+        ByteUtils.writeVarint(1, buf);
+        buf.put((byte) 1);
+        ByteUtils.writeVarint(105, buf); // header value too long
+        buf.position(buf.limit());
+
+        buf.flip();
+        DataInputStream inputStream = new DataInputStream(new ByteBufferInputStream(buf));
+        DefaultRecord.readPartiallyFrom(inputStream, skipArray, 0L, 0L, RecordBatch.NO_SEQUENCE, null);
+    }
+
     @Test(expected = InvalidRecordException.class)
     public void testUnderflowReadingTimestamp() {
         byte attributes = 0;
diff --git a/clients/src/test/java/org/apache/kafka/common/record/EndTransactionMarkerTest.java b/clients/src/test/java/org/apache/kafka/common/record/EndTransactionMarkerTest.java
index 903f674ed1965..8698c7cfca86b 100644
--- a/clients/src/test/java/org/apache/kafka/common/record/EndTransactionMarkerTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/record/EndTransactionMarkerTest.java
@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.common.record;
 
+import org.apache.kafka.common.InvalidRecordException;
 import org.junit.Test;
 
 import java.nio.ByteBuffer;
diff --git a/clients/src/test/java/org/apache/kafka/common/record/FileLogInputStreamTest.java b/clients/src/test/java/org/apache/kafka/common/record/FileLogInputStreamTest.java
index 783a5b531ef27..eba1a9cf48ea5 100644
--- a/clients/src/test/java/org/apache/kafka/common/record/FileLogInputStreamTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/record/FileLogInputStreamTest.java
@@ -283,8 +283,11 @@ private void assertGenericRecordBatchData(RecordBatch batch, long baseOffset, lo
     public static Collection data() {
         List values = new ArrayList<>();
         for (byte magic : asList(MAGIC_VALUE_V0, MAGIC_VALUE_V1, MAGIC_VALUE_V2))
-            for (CompressionType type: CompressionType.values())
-                values.add(new Object[] {magic, type});
+            for (CompressionType type: CompressionType.values()) {
+                if (type == CompressionType.PASSTHROUGH)
+                    continue;
+                values.add(new Object[]{magic, type});
+            }
         return values;
     }
 }
diff --git a/clients/src/test/java/org/apache/kafka/common/record/FileRecordsTest.java b/clients/src/test/java/org/apache/kafka/common/record/FileRecordsTest.java
index 20ecba1d18816..08dbc57d7f29d 100644
--- a/clients/src/test/java/org/apache/kafka/common/record/FileRecordsTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/record/FileRecordsTest.java
@@ -35,6 +35,10 @@
 import java.util.Iterator;
 import java.util.List;
 import java.util.Optional;
+import java.util.Random;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
 
 import static java.util.Arrays.asList;
 import static org.apache.kafka.common.utils.Utils.utf8;
@@ -117,6 +121,36 @@ public void testIterationOverPartialAndTruncation() throws IOException {
         testPartialWrite(6, fileRecords);
     }
 
+    @Test
+    public void testSliceSizeLimitWithConcurrentWrite() throws Exception {
+        FileRecords log = FileRecords.open(tempFile());
+        ExecutorService executor = Executors.newFixedThreadPool(2);
+        int maxSizeInBytes = 16384;
+
+        try {
+            Future readerCompletion = executor.submit(() -> {
+                while (log.sizeInBytes() < maxSizeInBytes) {
+                    int currentSize = log.sizeInBytes();
+                    FileRecords slice = log.slice(0, currentSize);
+                    assertEquals(currentSize, slice.sizeInBytes());
+                }
+                return null;
+            });
+
+            Future writerCompletion = executor.submit(() -> {
+                while (log.sizeInBytes() < maxSizeInBytes) {
+                    append(log, values);
+                }
+                return null;
+            });
+
+            writerCompletion.get();
+            readerCompletion.get();
+        } finally {
+            executor.shutdownNow();
+        }
+    }
+
     private void testPartialWrite(int size, FileRecords fileRecords) throws IOException {
         ByteBuffer buffer = ByteBuffer.allocate(size);
         for (int i = 0; i < size; i++)
@@ -434,6 +468,39 @@ private void appendWithOffsetAndTimestamp(FileRecords fileRecords,
         fileRecords.append(builder.build());
     }
 
+    @Test
+    public void testDownconversionAfterMessageFormatDowngrade() throws IOException {
+        // random bytes
+        Random random = new Random();
+        byte[] bytes = new byte[3000];
+        random.nextBytes(bytes);
+
+        // records
+        CompressionType compressionType = CompressionType.GZIP;
+        List offsets = asList(0L, 1L);
+        List magic = asList(RecordBatch.MAGIC_VALUE_V2, RecordBatch.MAGIC_VALUE_V1);  // downgrade message format from v2 to v1
+        List records = asList(
+                new SimpleRecord(1L, "k1".getBytes(), bytes),
+                new SimpleRecord(2L, "k2".getBytes(), bytes));
+        byte toMagic = 1;
+
+        // create MemoryRecords
+        ByteBuffer buffer = ByteBuffer.allocate(8000);
+        for (int i = 0; i < records.size(); i++) {
+            MemoryRecordsBuilder builder = MemoryRecords.builder(buffer, magic.get(i), compressionType, TimestampType.CREATE_TIME, 0L);
+            builder.appendWithOffset(offsets.get(i), records.get(i));
+            builder.close();
+        }
+        buffer.flip();
+
+        // create FileRecords, down-convert and verify
+        try (FileRecords fileRecords = FileRecords.open(tempFile())) {
+            fileRecords.append(MemoryRecords.readableRecords(buffer));
+            fileRecords.flush();
+            downConvertAndVerifyRecords(records, offsets, fileRecords, compressionType, toMagic, 0L, time);
+        }
+    }
+
     @Test
     public void testConversion() throws IOException {
         doTestConversion(CompressionType.NONE, RecordBatch.MAGIC_VALUE_V0);
diff --git a/clients/src/test/java/org/apache/kafka/common/record/LegacyRecordTest.java b/clients/src/test/java/org/apache/kafka/common/record/LegacyRecordTest.java
index 9480c60ca940b..848f0a3c5f28e 100644
--- a/clients/src/test/java/org/apache/kafka/common/record/LegacyRecordTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/record/LegacyRecordTest.java
@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.common.record;
 
+import org.apache.kafka.common.errors.CorruptRecordException;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.Parameterized;
@@ -93,7 +94,7 @@ public void testChecksum() {
             try {
                 copy.ensureValid();
                 fail("Should fail the above test.");
-            } catch (InvalidRecordException e) {
+            } catch (CorruptRecordException e) {
                 // this is good
             }
         }
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 522915f064034..ccd62b3ba8bcb 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
@@ -686,8 +686,11 @@ public void shouldThrowIllegalStateExceptionOnAppendWhenAborted() {
     public static Collection data() {
         List values = new ArrayList<>();
         for (int bufferOffset : Arrays.asList(0, 15))
-            for (CompressionType compressionType : CompressionType.values())
-                values.add(new Object[] {bufferOffset, compressionType});
+            for (CompressionType compressionType : CompressionType.values()) {
+                if (compressionType == CompressionType.PASSTHROUGH)
+                    continue;
+                values.add(new Object[]{bufferOffset, compressionType});
+            }
         return values;
     }
 
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 3d5a4f1a498b3..23f4dca9fd117 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
@@ -20,6 +20,7 @@
 import org.apache.kafka.common.errors.CorruptRecordException;
 import org.apache.kafka.common.header.internals.RecordHeaders;
 import org.apache.kafka.common.record.MemoryRecords.RecordFilter.BatchRetention;
+import org.apache.kafka.common.utils.CloseableIterator;
 import org.apache.kafka.common.utils.Utils;
 import org.apache.kafka.test.TestUtils;
 import org.junit.Test;
@@ -30,8 +31,8 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
-import java.util.function.Supplier;
 import java.util.List;
+import java.util.function.Supplier;
 
 import static java.util.Arrays.asList;
 import static org.apache.kafka.common.record.RecordBatch.MAGIC_VALUE_V2;
@@ -41,7 +42,6 @@
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
-import static org.junit.Assert.fail;
 import static org.junit.Assume.assumeTrue;
 
 @RunWith(value = Parameterized.class)
@@ -155,6 +155,125 @@ public void testIterator() {
         }
     }
 
+    @Test
+    public void testShallowIteration() {
+        assumeAtLeastV2OrNotZstd();
+        // Create the record batch
+        ByteBuffer buffer = ByteBuffer.allocate(1024);
+
+        MemoryRecordsBuilder builder1 = new MemoryRecordsBuilder(buffer, magic, compression,
+            TimestampType.CREATE_TIME, firstOffset, logAppendTime, pid, epoch, firstSequence, false, false,
+            partitionLeaderEpoch, buffer.limit());
+
+        SimpleRecord[] records = new SimpleRecord[] {
+            new SimpleRecord(0L, "a".getBytes(), "1".getBytes()),
+            new SimpleRecord(0L, "b".getBytes(), "2".getBytes()),
+            new SimpleRecord(0L, "c".getBytes(), "3".getBytes()),
+            new SimpleRecord(0L, "d".getBytes(), "4".getBytes())
+        };
+
+        for (SimpleRecord record : records)
+            builder1.append(record);
+
+        // Build MemoryRecords with normal appender
+        MemoryRecords records1 = builder1.build();
+        List recordBatches = Utils.toList(records1.batches().iterator());
+
+        // Shallow iterate over the records, if compression is NONE, there should be 4 entries during shallow iteration;
+        // otherwise, there should be only 1 entry
+        int recordCount = 0;
+
+        for (int i = 0; i < recordBatches.size(); i++) {
+            CloseableIterator iter = recordBatches.get(i).shallowIterator();
+
+            while (iter.hasNext()) {
+                Record curRecord = iter.next();
+                curRecord.ensureValid();
+                recordCount++;
+            }
+        }
+
+        int expectedRecordCount = compression == CompressionType.NONE ?
+            (magic == RecordBatch.MAGIC_VALUE_V2 ? 1 : 4) : 1;
+        assertEquals(expectedRecordCount, recordCount);
+    }
+
+    @Test
+    public void testIteratorWithProducerCompression() {
+        assumeAtLeastV2OrNotZstd();
+        ByteBuffer buffer = ByteBuffer.allocate(1024);
+
+        MemoryRecordsBuilder builder1 = new MemoryRecordsBuilder(buffer, magic, compression,
+            TimestampType.CREATE_TIME, firstOffset, logAppendTime, pid, epoch, firstSequence, false, false,
+            partitionLeaderEpoch, buffer.limit());
+
+        SimpleRecord[] records = new SimpleRecord[] {
+            new SimpleRecord(0L, "a".getBytes(), "1".getBytes()),
+            new SimpleRecord(0L, "b".getBytes(), "2".getBytes()),
+            new SimpleRecord(0L, "c".getBytes(), "3".getBytes()),
+            new SimpleRecord(0L, "d".getBytes(), "4".getBytes())
+        };
+
+        for (SimpleRecord record : records)
+            builder1.append(record);
+
+        // Build MemoryRecords with normal appender
+        MemoryRecords records1 = builder1.build();
+
+        List recordBatches = Utils.toList(records1.batches().iterator());
+
+        ByteBuffer buffer2 = ByteBuffer.allocate(1024);
+
+        MemoryRecordsBuilder builder2 = new MemoryRecordsBuilder(buffer2, magic, CompressionType.PASSTHROUGH,
+            TimestampType.CREATE_TIME, firstOffset, logAppendTime, pid, epoch, firstSequence, false, false,
+            partitionLeaderEpoch, buffer2.limit());
+
+        List recordList = new ArrayList<>();
+        // Part 1: Shallow iterate over the original record batch and use passthrough to append to new buffer
+        for (RecordBatch currentBatch : recordBatches) {
+            CloseableIterator currentRecords = currentBatch.shallowIterator();
+            while (currentRecords.hasNext()) {
+                Record curRecord = currentRecords.next();
+                recordList.add(curRecord);
+
+                if (magic >= RecordBatch.MAGIC_VALUE_V2) {
+                    builder2.append(0, null, curRecord.value());
+                } else {
+                    builder2.append(0, null, ((AbstractLegacyRecordBatch) curRecord).outerRecord().buffer());
+                }
+            }
+        }
+
+        MemoryRecords records2 = builder2.build();
+
+        List list2 = Utils.toList(records2.batches().iterator());
+
+        // Part 2: the newly produced records should be the same as the original ones
+        for (int i = 0; i < list2.size(); i++) {
+            CloseableIterator origRecIter = recordBatches.get(i).shallowIterator();
+            CloseableIterator curRecIter = list2.get(i).shallowIterator();
+
+            while (curRecIter.hasNext()) {
+                Record curRecord = curRecIter.next(), origRecord = origRecIter.next();
+
+                assertEquals(magic == RecordBatch.MAGIC_VALUE_V2 ? firstOffset : 0, curRecord.offset());
+
+                if (magic == RecordBatch.MAGIC_VALUE_V2) {
+                    assertEquals(curRecord, origRecord);
+                } else {
+                    assertEquals(((AbstractLegacyRecordBatch) curRecord).outerRecord().buffer(),
+                                 ((AbstractLegacyRecordBatch) origRecord).outerRecord().buffer());
+                }
+
+                curRecord.ensureValid();
+            }
+
+            assertFalse(origRecIter.hasNext());
+        }
+    }
+
+
+
     @Test
     public void testHasRoomForMethod() {
         assumeAtLeastV2OrNotZstd();
@@ -688,38 +807,6 @@ public void testFilterToWithUndersizedBuffer() {
             assertNotNull(record.key());
     }
 
-    @Test
-    public void testToString() {
-        assumeAtLeastV2OrNotZstd();
-
-        long timestamp = 1000000;
-        MemoryRecords memoryRecords = MemoryRecords.withRecords(magic, compression,
-                new SimpleRecord(timestamp, "key1".getBytes(), "value1".getBytes()),
-                new SimpleRecord(timestamp + 1, "key2".getBytes(), "value2".getBytes()));
-        switch (magic) {
-            case RecordBatch.MAGIC_VALUE_V0:
-                assertEquals("[(record=LegacyRecordBatch(offset=0, Record(magic=0, attributes=0, compression=NONE, " +
-                                "crc=1978725405, key=4 bytes, value=6 bytes))), (record=LegacyRecordBatch(offset=1, Record(magic=0, " +
-                                "attributes=0, compression=NONE, crc=1964753830, key=4 bytes, value=6 bytes)))]",
-                        memoryRecords.toString());
-                break;
-            case RecordBatch.MAGIC_VALUE_V1:
-                assertEquals("[(record=LegacyRecordBatch(offset=0, Record(magic=1, attributes=0, compression=NONE, " +
-                                "crc=97210616, CreateTime=1000000, key=4 bytes, value=6 bytes))), (record=LegacyRecordBatch(offset=1, " +
-                                "Record(magic=1, attributes=0, compression=NONE, crc=3535988507, CreateTime=1000001, key=4 bytes, " +
-                                "value=6 bytes)))]",
-                        memoryRecords.toString());
-                break;
-            case RecordBatch.MAGIC_VALUE_V2:
-                assertEquals("[(record=DefaultRecord(offset=0, timestamp=1000000, key=4 bytes, value=6 bytes)), " +
-                                "(record=DefaultRecord(offset=1, timestamp=1000001, key=4 bytes, value=6 bytes))]",
-                        memoryRecords.toString());
-                break;
-            default:
-                fail("Unexpected magic " + magic);
-        }
-    }
-
     @Test
     public void testFilterTo() {
         assumeAtLeastV2OrNotZstd();
@@ -936,8 +1023,11 @@ public static Collection data() {
         List values = new ArrayList<>();
         for (long firstOffset : asList(0L, 57L))
             for (byte magic : asList(RecordBatch.MAGIC_VALUE_V0, RecordBatch.MAGIC_VALUE_V1, RecordBatch.MAGIC_VALUE_V2))
-                for (CompressionType type: CompressionType.values())
-                    values.add(new Object[] {magic, firstOffset, type});
+                for (CompressionType type: CompressionType.values()) {
+                    if (type == CompressionType.PASSTHROUGH)
+                        continue;
+                    values.add(new Object[]{magic, firstOffset, type});
+                }
         return values;
     }
 
diff --git a/clients/src/test/java/org/apache/kafka/common/record/SimpleLegacyRecordTest.java b/clients/src/test/java/org/apache/kafka/common/record/SimpleLegacyRecordTest.java
index 5f578a873d0dc..cd287bbbf1c5b 100644
--- a/clients/src/test/java/org/apache/kafka/common/record/SimpleLegacyRecordTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/record/SimpleLegacyRecordTest.java
@@ -16,6 +16,8 @@
  */
 package org.apache.kafka.common.record;
 
+import org.apache.kafka.common.InvalidRecordException;
+import org.apache.kafka.common.errors.CorruptRecordException;
 import org.apache.kafka.common.utils.ByteBufferOutputStream;
 import org.apache.kafka.common.utils.Utils;
 import org.junit.Test;
@@ -66,7 +68,7 @@ public void testCompressedIterationWithEmptyRecords() throws Exception {
     }
 
     /* This scenario can happen if the record size field is corrupt and we end up allocating a buffer that is too small */
-    @Test(expected = InvalidRecordException.class)
+    @Test(expected = CorruptRecordException.class)
     public void testIsValidWithTooSmallBuffer() {
         ByteBuffer buffer = ByteBuffer.allocate(2);
         LegacyRecord record = new LegacyRecord(buffer);
@@ -74,7 +76,7 @@ public void testIsValidWithTooSmallBuffer() {
         record.ensureValid();
     }
 
-    @Test(expected = InvalidRecordException.class)
+    @Test(expected = CorruptRecordException.class)
     public void testIsValidWithChecksumMismatch() {
         ByteBuffer buffer = ByteBuffer.allocate(4);
         // set checksum
diff --git a/clients/src/test/java/org/apache/kafka/common/replica/ReplicaSelectorTest.java b/clients/src/test/java/org/apache/kafka/common/replica/ReplicaSelectorTest.java
new file mode 100644
index 0000000000000..da03e5787ddd3
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/replica/ReplicaSelectorTest.java
@@ -0,0 +1,88 @@
+/*
+ * 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.replica;
+
+import org.apache.kafka.common.Node;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.security.auth.KafkaPrincipal;
+import org.junit.Test;
+
+import java.net.InetAddress;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Optional;
+import java.util.Set;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+import static org.apache.kafka.test.TestUtils.assertOptional;
+import static org.junit.Assert.assertEquals;
+
+public class ReplicaSelectorTest {
+
+    @Test
+    public void testSameRackSelector() {
+        TopicPartition tp = new TopicPartition("test", 0);
+
+        List replicaViewSet = replicaInfoSet();
+        ReplicaView leader = replicaViewSet.get(0);
+        PartitionView partitionView = partitionInfo(new HashSet<>(replicaViewSet), leader);
+
+        ReplicaSelector selector = new RackAwareReplicaSelector();
+        Optional selected = selector.select(tp, metadata("rack-b"), partitionView);
+        assertOptional(selected, replicaInfo -> {
+            assertEquals("Expect replica to be in rack-b", replicaInfo.endpoint().rack(), "rack-b");
+            assertEquals("Expected replica 3 since it is more caught-up", replicaInfo.endpoint().id(), 3);
+        });
+
+        selected = selector.select(tp, metadata("not-a-rack"), partitionView);
+        assertOptional(selected, replicaInfo -> {
+            assertEquals("Expect leader when we can't find any nodes in given rack", replicaInfo, leader);
+        });
+
+        selected = selector.select(tp, metadata("rack-a"), partitionView);
+        assertOptional(selected, replicaInfo -> {
+            assertEquals("Expect replica to be in rack-a", replicaInfo.endpoint().rack(), "rack-a");
+            assertEquals("Expect the leader since it's in rack-a", replicaInfo, leader);
+        });
+
+
+    }
+
+    static List replicaInfoSet() {
+        return Stream.of(
+                replicaInfo(new Node(0, "host0", 1234, "rack-a"), 4, 0),
+                replicaInfo(new Node(1, "host1", 1234, "rack-a"), 2, 5),
+                replicaInfo(new Node(2, "host2", 1234, "rack-b"), 3, 3),
+                replicaInfo(new Node(3, "host3", 1234, "rack-b"), 4, 2)
+
+        ).collect(Collectors.toList());
+    }
+
+    static ReplicaView replicaInfo(Node node, long logOffset, long timeSinceLastCaughtUpMs) {
+        return new ReplicaView.DefaultReplicaView(node, logOffset, timeSinceLastCaughtUpMs);
+    }
+
+    static PartitionView partitionInfo(Set replicaViewSet, ReplicaView leader) {
+        return new PartitionView.DefaultPartitionView(replicaViewSet, leader);
+    }
+
+    static ClientMetadata metadata(String rack) {
+        return new ClientMetadata.DefaultClientMetadata(rack, "test-client",
+                InetAddress.getLoopbackAddress(), KafkaPrincipal.ANONYMOUS, "TEST");
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/requests/ApiVersionsResponseTest.java b/clients/src/test/java/org/apache/kafka/common/requests/ApiVersionsResponseTest.java
index 2b526d12fe248..c7f1ecc001d94 100644
--- a/clients/src/test/java/org/apache/kafka/common/requests/ApiVersionsResponseTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/requests/ApiVersionsResponseTest.java
@@ -17,6 +17,7 @@
 
 package org.apache.kafka.common.requests;
 
+import org.apache.kafka.common.message.ApiVersionsResponseData.ApiVersionsResponseKey;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.record.RecordBatch;
 import org.apache.kafka.common.utils.Utils;
@@ -43,7 +44,7 @@ public void shouldCreateApiResponseOnlyWithKeysSupportedByMagicValue() {
 
     @Test
     public void shouldCreateApiResponseThatHasAllApiKeysSupportedByBroker() {
-        assertEquals(apiKeysInResponse(ApiVersionsResponse.defaultApiVersionsResponse()), Utils.mkSet(ApiKeys.values()));
+        assertEquals(apiKeysInResponse(ApiVersionsResponse.DEFAULT_API_VERSIONS_RESPONSE), Utils.mkSet(ApiKeys.values()));
     }
 
     @Test
@@ -55,39 +56,39 @@ public void shouldReturnAllKeysWhenMagicIsCurrentValueAndThrottleMsIsDefaultThro
 
     @Test
     public void shouldHaveCorrectDefaultApiVersionsResponse() {
-        Collection apiVersions = ApiVersionsResponse.defaultApiVersionsResponse().apiVersions();
+        Collection apiVersions = ApiVersionsResponse.DEFAULT_API_VERSIONS_RESPONSE.data.apiKeys();
         assertEquals("API versions for all API keys must be maintained.", apiVersions.size(), ApiKeys.values().length);
 
         for (ApiKeys key : ApiKeys.values()) {
-            ApiVersionsResponse.ApiVersion version = ApiVersionsResponse.defaultApiVersionsResponse().apiVersion(key.id);
+            ApiVersionsResponseKey version = ApiVersionsResponse.DEFAULT_API_VERSIONS_RESPONSE.apiVersion(key.id);
             assertNotNull("Could not find ApiVersion for API " + key.name, version);
-            assertEquals("Incorrect min version for Api " + key.name, version.minVersion, key.oldestVersion());
-            assertEquals("Incorrect max version for Api " + key.name, version.maxVersion, key.latestVersion());
+            assertEquals("Incorrect min version for Api " + key.name, version.minVersion(), key.oldestVersion());
+            assertEquals("Incorrect max version for Api " + key.name, version.maxVersion(), key.latestVersion());
 
             // Check if versions less than min version are indeed set as null, i.e., deprecated.
-            for (int i = 0; i < version.minVersion; ++i) {
-                assertNull("Request version " + i + " for API " + version.apiKey + " must be null", key.requestSchemas[i]);
-                assertNull("Response version " + i + " for API " + version.apiKey + " must be null", key.responseSchemas[i]);
+            for (int i = 0; i < version.minVersion(); ++i) {
+                assertNull("Request version " + i + " for API " + version.apiKey() + " must be null", key.requestSchemas[i]);
+                assertNull("Response version " + i + " for API " + version.apiKey() + " must be null", key.responseSchemas[i]);
             }
 
             // Check if versions between min and max versions are non null, i.e., valid.
-            for (int i = version.minVersion; i <= version.maxVersion; ++i) {
-                assertNotNull("Request version " + i + " for API " + version.apiKey + " must not be null", key.requestSchemas[i]);
-                assertNotNull("Response version " + i + " for API " + version.apiKey + " must not be null", key.responseSchemas[i]);
+            for (int i = version.minVersion(); i <= version.maxVersion(); ++i) {
+                assertNotNull("Request version " + i + " for API " + version.apiKey() + " must not be null", key.requestSchemas[i]);
+                assertNotNull("Response version " + i + " for API " + version.apiKey() + " must not be null", key.responseSchemas[i]);
             }
         }
     }
     
     private void verifyApiKeysForMagic(final ApiVersionsResponse response, final byte maxMagic) {
-        for (final ApiVersionsResponse.ApiVersion version : response.apiVersions()) {
-            assertTrue(ApiKeys.forId(version.apiKey).minRequiredInterBrokerMagic <= maxMagic);
+        for (final ApiVersionsResponseKey version : response.data.apiKeys()) {
+            assertTrue(ApiKeys.forId(version.apiKey()).minRequiredInterBrokerMagic <= maxMagic);
         }
     }
 
     private Set apiKeysInResponse(final ApiVersionsResponse apiVersions) {
         final Set apiKeys = new HashSet<>();
-        for (final ApiVersionsResponse.ApiVersion version : apiVersions.apiVersions()) {
-            apiKeys.add(ApiKeys.forId(version.apiKey));
+        for (final ApiVersionsResponseKey version : apiVersions.data.apiKeys()) {
+            apiKeys.add(ApiKeys.forId(version.apiKey()));
         }
         return apiKeys;
     }
diff --git a/clients/src/test/java/org/apache/kafka/common/requests/ControlRequestTest.java b/clients/src/test/java/org/apache/kafka/common/requests/ControlRequestTest.java
deleted file mode 100644
index 2f4689ed0a09c..0000000000000
--- a/clients/src/test/java/org/apache/kafka/common/requests/ControlRequestTest.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.common.requests;
-
-import org.apache.kafka.common.TopicPartition;
-import org.junit.Assert;
-import org.junit.Test;
-
-import java.util.Collections;
-import java.util.Set;
-import java.util.Map;
-import java.util.HashMap;
-import java.util.HashSet;
-import java.util.Random;
-
-
-public class ControlRequestTest {
-    @Test
-    public void testLeaderAndIsrRequestNormalization() {
-        Set tps = generateRandomTopicPartitions(10, 10);
-        Map partitionStates = new HashMap<>();
-        for (TopicPartition tp: tps) {
-            partitionStates.put(tp, new LeaderAndIsrRequest.PartitionState(0, 0, 0,
-                    Collections.emptyList(), 0, Collections.emptyList(), false));
-        }
-        LeaderAndIsrRequest.Builder builder = new LeaderAndIsrRequest.Builder((short) 2, 0, 0, 0,
-                partitionStates, Collections.emptySet());
-
-        Assert.assertTrue(builder.build((short) 2).size() <  builder.build((short) 1).size());
-    }
-
-    @Test
-    public void testUpdateMetadataRequestNormalization() {
-        Set tps = generateRandomTopicPartitions(10, 10);
-        Map partitionStates = new HashMap<>();
-        for (TopicPartition tp: tps) {
-            partitionStates.put(tp, new UpdateMetadataRequest.PartitionState(0, 0, 0,
-                    Collections.emptyList(), 0, Collections.emptyList(), Collections.emptyList()));
-        }
-        UpdateMetadataRequest.Builder builder = new UpdateMetadataRequest.Builder((short) 5, 0, 0, 0,
-                partitionStates, Collections.emptySet());
-
-        Assert.assertTrue(builder.build((short) 5).size() <  builder.build((short) 4).size());
-    }
-
-    @Test
-    public void testStopReplicaRequestNormalization() {
-        Set tps = generateRandomTopicPartitions(10, 10);
-        Map partitionStates = new HashMap<>();
-        for (TopicPartition tp: tps) {
-            partitionStates.put(tp, new UpdateMetadataRequest.PartitionState(0, 0, 0,
-                    Collections.emptyList(), 0, Collections.emptyList(), Collections.emptyList()));
-        }
-        StopReplicaRequest.Builder builder = new StopReplicaRequest.Builder((short) 5, 0, 0, 0, false, tps);
-
-        Assert.assertTrue(builder.build((short) 1).size() <  builder.build((short) 0).size());
-    }
-
-    private Set generateRandomTopicPartitions(int numTopic, int numPartitionPerTopic) {
-        Set tps = new HashSet<>();
-        Random r = new Random();
-        for (int i = 0; i < numTopic; i++) {
-            byte[] array = new byte[32];
-            r.nextBytes(array);
-            String topic = new String(array);
-            for (int j = 0; j < numPartitionPerTopic; j++) {
-                tps.add(new TopicPartition(topic, j));
-            }
-        }
-        return tps;
-    }
-
-}
diff --git a/clients/src/test/java/org/apache/kafka/common/requests/ControlledShutdownRequestTest.java b/clients/src/test/java/org/apache/kafka/common/requests/ControlledShutdownRequestTest.java
new file mode 100644
index 0000000000000..b7ec3da9d1b01
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/ControlledShutdownRequestTest.java
@@ -0,0 +1,51 @@
+/*
+ * 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.requests;
+
+import org.apache.kafka.common.errors.ClusterAuthorizationException;
+import org.apache.kafka.common.errors.UnsupportedVersionException;
+import org.apache.kafka.common.message.ControlledShutdownRequestData;
+import org.apache.kafka.common.protocol.Errors;
+import org.junit.Test;
+
+import static org.apache.kafka.common.protocol.ApiKeys.CONTROLLED_SHUTDOWN;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
+
+public class ControlledShutdownRequestTest {
+
+    @Test
+    public void testUnsupportedVersion() {
+        ControlledShutdownRequest.Builder builder = new ControlledShutdownRequest.Builder(
+                new ControlledShutdownRequestData().setBrokerId(1),
+                (short) (CONTROLLED_SHUTDOWN.latestVersion() + 1));
+        assertThrows(UnsupportedVersionException.class, builder::build);
+    }
+
+    @Test
+    public void testGetErrorResponse() {
+        for (short version = CONTROLLED_SHUTDOWN.oldestVersion(); version < CONTROLLED_SHUTDOWN.latestVersion(); version++) {
+            ControlledShutdownRequest.Builder builder = new ControlledShutdownRequest.Builder(
+                    new ControlledShutdownRequestData().setBrokerId(1), version);
+            ControlledShutdownRequest request = builder.build();
+            ControlledShutdownResponse response = request.getErrorResponse(0,
+                    new ClusterAuthorizationException("Not authorized"));
+            assertEquals(Errors.CLUSTER_AUTHORIZATION_FAILED, response.error());
+        }
+    }
+
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/requests/DeleteGroupsResponseTest.java b/clients/src/test/java/org/apache/kafka/common/requests/DeleteGroupsResponseTest.java
new file mode 100644
index 0000000000000..37e8a5202ee75
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/DeleteGroupsResponseTest.java
@@ -0,0 +1,79 @@
+/*
+ * 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.requests;
+
+import org.apache.kafka.common.message.DeleteGroupsResponseData;
+import org.apache.kafka.common.message.DeleteGroupsResponseData.DeletableGroupResult;
+import org.apache.kafka.common.message.DeleteGroupsResponseData.DeletableGroupResultCollection;
+import org.apache.kafka.common.protocol.Errors;
+import org.junit.Test;
+
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.Map;
+
+import static org.junit.Assert.assertEquals;
+
+public class DeleteGroupsResponseTest {
+
+    private static final String GROUP_ID_1 = "groupId1";
+    private static final String GROUP_ID_2 = "groupId2";
+    private static final int THROTTLE_TIME_MS = 10;
+    private static DeleteGroupsResponse deleteGroupsResponse;
+
+    static {
+        deleteGroupsResponse = new DeleteGroupsResponse(
+            new DeleteGroupsResponseData()
+                .setResults(
+                    new DeletableGroupResultCollection(Arrays.asList(
+                        new DeletableGroupResult()
+                            .setGroupId(GROUP_ID_1)
+                            .setErrorCode(Errors.NONE.code()),
+                        new DeletableGroupResult()
+                            .setGroupId(GROUP_ID_2)
+                            .setErrorCode(Errors.GROUP_AUTHORIZATION_FAILED.code())).iterator()
+                    )
+                )
+                .setThrottleTimeMs(THROTTLE_TIME_MS));
+    }
+
+    @Test
+    public void testGetErrorWithExistingGroupIds() {
+        assertEquals(Errors.NONE, deleteGroupsResponse.get(GROUP_ID_1));
+        assertEquals(Errors.GROUP_AUTHORIZATION_FAILED, deleteGroupsResponse.get(GROUP_ID_2));
+
+        Map expectedErrors = new HashMap<>();
+        expectedErrors.put(GROUP_ID_1, Errors.NONE);
+        expectedErrors.put(GROUP_ID_2, Errors.GROUP_AUTHORIZATION_FAILED);
+        assertEquals(expectedErrors, deleteGroupsResponse.errors());
+
+        Map expectedErrorCounts = new HashMap<>();
+        expectedErrorCounts.put(Errors.NONE, 1);
+        expectedErrorCounts.put(Errors.GROUP_AUTHORIZATION_FAILED, 1);
+        assertEquals(expectedErrorCounts, deleteGroupsResponse.errorCounts());
+    }
+
+    @Test(expected = IllegalArgumentException.class)
+    public void testGetErrorWithInvalidGroupId() {
+        deleteGroupsResponse.get("invalid-group-id");
+    }
+
+    @Test
+    public void testGetThrottleTimeMs() {
+        assertEquals(THROTTLE_TIME_MS, deleteGroupsResponse.throttleTimeMs());
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/BasePartitionState.java b/clients/src/test/java/org/apache/kafka/common/requests/HeartbeatRequestTest.java
similarity index 53%
rename from clients/src/main/java/org/apache/kafka/common/requests/BasePartitionState.java
rename to clients/src/test/java/org/apache/kafka/common/requests/HeartbeatRequestTest.java
index c9a461093caaf..2532213041d93 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/BasePartitionState.java
+++ b/clients/src/test/java/org/apache/kafka/common/requests/HeartbeatRequestTest.java
@@ -16,25 +16,19 @@
  */
 package org.apache.kafka.common.requests;
 
-import java.util.List;
+import org.apache.kafka.common.errors.UnsupportedVersionException;
+import org.apache.kafka.common.message.HeartbeatRequestData;
+import org.junit.Test;
 
-// This class contains the common fields shared between LeaderAndIsrRequest.PartitionState and UpdateMetadataRequest.PartitionState
-public class BasePartitionState {
+public class HeartbeatRequestTest {
 
-    public final int controllerEpoch;
-    public final int leader;
-    public final int leaderEpoch;
-    public final List isr;
-    public final int zkVersion;
-    public final List replicas;
-
-    BasePartitionState(int controllerEpoch, int leader, int leaderEpoch, List isr, int zkVersion, List replicas) {
-        this.controllerEpoch = controllerEpoch;
-        this.leader = leader;
-        this.leaderEpoch = leaderEpoch;
-        this.isr = isr;
-        this.zkVersion = zkVersion;
-        this.replicas = replicas;
+    @Test(expected = UnsupportedVersionException.class)
+    public void testRequestVersionCompatibilityFailBuild() {
+        new HeartbeatRequest.Builder(
+                new HeartbeatRequestData()
+                        .setGroupId("groupId")
+                        .setMemberId("consumerId")
+                        .setGroupInstanceId("groupInstanceId")
+        ).build((short) 2);
     }
-
 }
diff --git a/clients/src/test/java/org/apache/kafka/common/requests/JoinGroupRequestTest.java b/clients/src/test/java/org/apache/kafka/common/requests/JoinGroupRequestTest.java
new file mode 100644
index 0000000000000..9d8031cd30074
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/JoinGroupRequestTest.java
@@ -0,0 +1,102 @@
+/*
+ * 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.requests;
+
+import org.apache.kafka.common.errors.InvalidConfigurationException;
+import org.apache.kafka.common.errors.UnsupportedVersionException;
+import org.apache.kafka.common.message.JoinGroupRequestData;
+import org.apache.kafka.common.protocol.types.Struct;
+import org.apache.kafka.test.TestUtils;
+import org.junit.Test;
+
+import java.nio.ByteBuffer;
+import java.util.Arrays;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.fail;
+
+public class JoinGroupRequestTest {
+
+    @Test
+    public void shouldAcceptValidGroupInstanceIds() {
+        String maxLengthString = TestUtils.randomString(249);
+        String[] validGroupInstanceIds = {"valid", "INSTANCE", "gRoUp", "ar6", "VaL1d", "_0-9_.", "...", maxLengthString};
+
+        for (String instanceId : validGroupInstanceIds) {
+            JoinGroupRequest.validateGroupInstanceId(instanceId);
+        }
+    }
+
+    @Test
+    public void shouldThrowOnInvalidGroupInstanceIds() {
+        char[] longString = new char[250];
+        Arrays.fill(longString, 'a');
+        String[] invalidGroupInstanceIds = {"", "foo bar", "..", "foo:bar", "foo=bar", ".", new String(longString)};
+
+        for (String instanceId : invalidGroupInstanceIds) {
+            try {
+                JoinGroupRequest.validateGroupInstanceId(instanceId);
+                fail("No exception was thrown for invalid instance id: " + instanceId);
+            } catch (InvalidConfigurationException e) {
+                // Good
+            }
+        }
+    }
+
+    @Test
+    public void shouldRecognizeInvalidCharactersInGroupInstanceIds() {
+        char[] invalidChars = {'/', '\\', ',', '\u0000', ':', '"', '\'', ';', '*', '?', ' ', '\t', '\r', '\n', '='};
+
+        for (char c : invalidChars) {
+            String instanceId = "Is " + c + "illegal";
+            assertFalse(JoinGroupRequest.containsValidPattern(instanceId));
+        }
+    }
+
+    @Test(expected = UnsupportedVersionException.class)
+    public void testRequestVersionCompatibilityFailBuild() {
+        new JoinGroupRequest.Builder(
+                new JoinGroupRequestData()
+                        .setGroupId("groupId")
+                        .setMemberId("consumerId")
+                        .setGroupInstanceId("groupInstanceId")
+                        .setProtocolType("consumer")
+        ).build((short) 4);
+    }
+
+    @Test
+    public void testRebalanceTimeoutDefaultsToSessionTimeoutV0() {
+        int sessionTimeoutMs = 30000;
+
+        Struct struct = new JoinGroupRequestData()
+                .setGroupId("groupId")
+                .setMemberId("consumerId")
+                .setGroupInstanceId("groupInstanceId")
+                .setProtocolType("consumer")
+                .setSessionTimeoutMs(sessionTimeoutMs)
+                .toStruct((short) 0);
+
+        ByteBuffer buffer = ByteBuffer.allocate(struct.sizeOf());
+        struct.writeTo(buffer);
+        buffer.flip();
+
+        JoinGroupRequest request = JoinGroupRequest.parse(buffer, (short) 0);
+        assertEquals(sessionTimeoutMs, request.data().sessionTimeoutMs());
+        assertEquals(sessionTimeoutMs, request.data().rebalanceTimeoutMs());
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/requests/LeaderAndIsrRequestTest.java b/clients/src/test/java/org/apache/kafka/common/requests/LeaderAndIsrRequestTest.java
new file mode 100644
index 0000000000000..7a7475cc4790c
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/LeaderAndIsrRequestTest.java
@@ -0,0 +1,175 @@
+/*
+ * 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.requests;
+
+import org.apache.kafka.common.Node;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.errors.ClusterAuthorizationException;
+import org.apache.kafka.common.errors.UnsupportedVersionException;
+import org.apache.kafka.common.message.LeaderAndIsrRequestData;
+import org.apache.kafka.common.message.LeaderAndIsrRequestData.LeaderAndIsrLiveLeader;
+import org.apache.kafka.common.message.LeaderAndIsrRequestData.LeaderAndIsrPartitionState;
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
+import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.common.protocol.MessageTestUtil;
+import org.apache.kafka.test.TestUtils;
+import org.junit.Test;
+
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+import java.util.stream.Collectors;
+import java.util.stream.StreamSupport;
+
+import static java.util.Arrays.asList;
+import static java.util.Collections.emptyList;
+import static org.apache.kafka.common.protocol.ApiKeys.LEADER_AND_ISR;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
+
+public class LeaderAndIsrRequestTest {
+
+    @Test
+    public void testUnsupportedVersion() {
+        LeaderAndIsrRequest.Builder builder = new LeaderAndIsrRequest.Builder(
+                (short) (LEADER_AND_ISR.latestVersion() + 1), 0, 0, 0L, 0L,
+                Collections.emptyList(), Collections.emptySet());
+        assertThrows(UnsupportedVersionException.class, builder::build);
+    }
+
+    @Test
+    public void testGetErrorResponse() {
+        for (short version = LEADER_AND_ISR.oldestVersion(); version < LEADER_AND_ISR.latestVersion(); version++) {
+            LeaderAndIsrRequest.Builder builder = new LeaderAndIsrRequest.Builder(version, 0, 0, 0, 0,
+                    Collections.emptyList(), Collections.emptySet());
+            LeaderAndIsrRequest request = builder.build();
+            LeaderAndIsrResponse response = request.getErrorResponse(0,
+                    new ClusterAuthorizationException("Not authorized"));
+            assertEquals(Errors.CLUSTER_AUTHORIZATION_FAILED, response.error());
+        }
+    }
+
+    /**
+     * Verifies the logic we have in LeaderAndIsrRequest to present a unified interface across the various versions
+     * works correctly. For example, `LeaderAndIsrPartitionState.topicName` is not serialiazed/deserialized in
+     * recent versions, but we set it manually so that we can always present the ungrouped partition states
+     * independently of the version.
+     */
+    @Test
+    public void testVersionLogic() {
+        for (short version = LEADER_AND_ISR.oldestVersion(); version <= LEADER_AND_ISR.latestVersion(); version++) {
+            List partitionStates = asList(
+                new LeaderAndIsrPartitionState()
+                    .setTopicName("topic0")
+                    .setPartitionIndex(0)
+                    .setControllerEpoch(2)
+                    .setLeader(0)
+                    .setLeaderEpoch(10)
+                    .setIsr(asList(0, 1))
+                    .setZkVersion(10)
+                    .setReplicas(asList(0, 1, 2))
+                    .setAddingReplicas(asList(3))
+                    .setRemovingReplicas(asList(2)),
+                new LeaderAndIsrPartitionState()
+                    .setTopicName("topic0")
+                    .setPartitionIndex(1)
+                    .setControllerEpoch(2)
+                    .setLeader(1)
+                    .setLeaderEpoch(11)
+                    .setIsr(asList(1, 2, 3))
+                    .setZkVersion(11)
+                    .setReplicas(asList(1, 2, 3))
+                    .setAddingReplicas(emptyList())
+                    .setRemovingReplicas(emptyList()),
+                new LeaderAndIsrPartitionState()
+                    .setTopicName("topic1")
+                    .setPartitionIndex(0)
+                    .setControllerEpoch(2)
+                    .setLeader(2)
+                    .setLeaderEpoch(11)
+                    .setIsr(asList(2, 3, 4))
+                    .setZkVersion(11)
+                    .setReplicas(asList(2, 3, 4))
+                    .setAddingReplicas(emptyList())
+                    .setRemovingReplicas(emptyList())
+            );
+
+            List liveNodes = asList(
+                new Node(0, "host0", 9090),
+                new Node(1, "host1", 9091)
+            );
+            LeaderAndIsrRequest request = new LeaderAndIsrRequest.Builder(version, 1, 2, 3, 3, partitionStates,
+                liveNodes).build();
+
+            List liveLeaders = liveNodes.stream().map(n -> new LeaderAndIsrLiveLeader()
+                .setBrokerId(n.id())
+                .setHostName(n.host())
+                .setPort(n.port())).collect(Collectors.toList());
+            assertEquals(new HashSet<>(partitionStates), iterableToSet(request.partitionStates()));
+            assertEquals(liveLeaders, request.liveLeaders());
+            assertEquals(1, request.controllerId());
+            assertEquals(2, request.controllerEpoch());
+            assertEquals(3, request.brokerEpoch());
+
+            ByteBuffer byteBuffer = MessageTestUtil.messageToByteBuffer(request.data(), request.version());
+            LeaderAndIsrRequest deserializedRequest = new LeaderAndIsrRequest(new LeaderAndIsrRequestData(
+                new ByteBufferAccessor(byteBuffer), version), version);
+
+            // Adding/removing replicas is only supported from version 4, so the deserialized request won't have
+            // them for earlier versions.
+            if (version < 4) {
+                partitionStates.get(0)
+                    .setAddingReplicas(emptyList())
+                    .setRemovingReplicas(emptyList());
+            }
+
+            assertEquals(new HashSet<>(partitionStates), iterableToSet(deserializedRequest.partitionStates()));
+            assertEquals(liveLeaders, deserializedRequest.liveLeaders());
+            assertEquals(1, request.controllerId());
+            assertEquals(2, request.controllerEpoch());
+            assertEquals(3, request.brokerEpoch());
+        }
+    }
+
+    @Test
+    public void testTopicPartitionGroupingSizeReduction() {
+        Set tps = TestUtils.generateRandomTopicPartitions(10, 10);
+        List partitionStates = new ArrayList<>();
+        for (TopicPartition tp : tps) {
+            partitionStates.add(new LeaderAndIsrPartitionState()
+                .setTopicName(tp.topic())
+                .setPartitionIndex(tp.partition()));
+        }
+        LeaderAndIsrRequest.Builder builder = new LeaderAndIsrRequest.Builder((short) 2, 0, 0, 0, 0,
+            partitionStates, Collections.emptySet());
+
+        LeaderAndIsrRequest v2 = builder.build((short) 2);
+        LeaderAndIsrRequest v1 = builder.build((short) 1);
+        int size2 = MessageTestUtil.messageSize(v2.data(), v2.version());
+        int size1 = MessageTestUtil.messageSize(v1.data(), v1.version());
+
+        assertTrue("Expected v2 < v1: v2=" + size2 + ", v1=" + size1, size2 < size1);
+    }
+
+    private  Set iterableToSet(Iterable iterable) {
+        return StreamSupport.stream(iterable.spliterator(), false).collect(Collectors.toSet());
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/requests/LeaderAndIsrResponseTest.java b/clients/src/test/java/org/apache/kafka/common/requests/LeaderAndIsrResponseTest.java
index 74e097191c118..4782a1a408741 100644
--- a/clients/src/test/java/org/apache/kafka/common/requests/LeaderAndIsrResponseTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/requests/LeaderAndIsrResponseTest.java
@@ -16,16 +16,19 @@
  */
 package org.apache.kafka.common.requests;
 
-import org.apache.kafka.common.Node;
-import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.message.LeaderAndIsrRequestData.LeaderAndIsrPartitionState;
+import org.apache.kafka.common.message.LeaderAndIsrResponseData;
+import org.apache.kafka.common.message.LeaderAndIsrResponseData.LeaderAndIsrPartitionError;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
 import org.junit.Test;
 
+import java.util.ArrayList;
 import java.util.Collections;
-import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 
+import static java.util.Arrays.asList;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 
@@ -33,32 +36,50 @@ public class LeaderAndIsrResponseTest {
 
     @Test
     public void testErrorCountsFromGetErrorResponse() {
-        HashMap partitionStates = new HashMap<>();
-        partitionStates.put(new TopicPartition("foo", 0), new LeaderAndIsrRequest.PartitionState(15, 1, 10,
-                Collections.singletonList(10), 20, Collections.singletonList(10), false));
-        partitionStates.put(new TopicPartition("foo", 1), new LeaderAndIsrRequest.PartitionState(15, 1, 10,
-                Collections.singletonList(10), 20, Collections.singletonList(10), false));
+        List partitionStates = new ArrayList<>();
+        partitionStates.add(new LeaderAndIsrPartitionState()
+            .setTopicName("foo")
+            .setPartitionIndex(0)
+            .setControllerEpoch(15)
+            .setLeader(1)
+            .setLeaderEpoch(10)
+            .setIsr(Collections.singletonList(10))
+            .setZkVersion(20)
+            .setReplicas(Collections.singletonList(10))
+            .setIsNew(false));
+        partitionStates.add(new LeaderAndIsrPartitionState()
+            .setTopicName("foo")
+            .setPartitionIndex(1)
+            .setControllerEpoch(15)
+            .setLeader(1)
+            .setLeaderEpoch(10)
+            .setIsr(Collections.singletonList(10))
+            .setZkVersion(20)
+            .setReplicas(Collections.singletonList(10))
+            .setIsNew(false));
         LeaderAndIsrRequest request = new LeaderAndIsrRequest.Builder(ApiKeys.LEADER_AND_ISR.latestVersion(),
-                15, 20, 0, partitionStates, Collections.emptySet()).build();
+                15, 20, 0, 0, partitionStates, Collections.emptySet()).build();
         LeaderAndIsrResponse response = request.getErrorResponse(0, Errors.CLUSTER_AUTHORIZATION_FAILED.exception());
         assertEquals(Collections.singletonMap(Errors.CLUSTER_AUTHORIZATION_FAILED, 2), response.errorCounts());
     }
 
     @Test
     public void testErrorCountsWithTopLevelError() {
-        Map errors = new HashMap<>();
-        errors.put(new TopicPartition("foo", 0), Errors.NONE);
-        errors.put(new TopicPartition("foo", 1), Errors.NOT_LEADER_FOR_PARTITION);
-        LeaderAndIsrResponse response = new LeaderAndIsrResponse(Errors.UNKNOWN_SERVER_ERROR, errors);
+        List partitions = createPartitions("foo",
+            asList(Errors.NONE, Errors.NOT_LEADER_FOR_PARTITION));
+        LeaderAndIsrResponse response = new LeaderAndIsrResponse(new LeaderAndIsrResponseData()
+            .setErrorCode(Errors.UNKNOWN_SERVER_ERROR.code())
+            .setPartitionErrors(partitions));
         assertEquals(Collections.singletonMap(Errors.UNKNOWN_SERVER_ERROR, 2), response.errorCounts());
     }
 
     @Test
     public void testErrorCountsNoTopLevelError() {
-        Map errors = new HashMap<>();
-        errors.put(new TopicPartition("foo", 0), Errors.NONE);
-        errors.put(new TopicPartition("foo", 1), Errors.CLUSTER_AUTHORIZATION_FAILED);
-        LeaderAndIsrResponse response = new LeaderAndIsrResponse(Errors.NONE, errors);
+        List partitions = createPartitions("foo",
+            asList(Errors.NONE, Errors.CLUSTER_AUTHORIZATION_FAILED));
+        LeaderAndIsrResponse response = new LeaderAndIsrResponse(new LeaderAndIsrResponseData()
+            .setErrorCode(Errors.NONE.code())
+            .setPartitionErrors(partitions));
         Map errorCounts = response.errorCounts();
         assertEquals(2, errorCounts.size());
         assertEquals(1, errorCounts.get(Errors.NONE).intValue());
@@ -67,14 +88,27 @@ public void testErrorCountsNoTopLevelError() {
 
     @Test
     public void testToString() {
-        Map errors = new HashMap<>();
-        errors.put(new TopicPartition("foo", 0), Errors.NONE);
-        errors.put(new TopicPartition("foo", 1), Errors.CLUSTER_AUTHORIZATION_FAILED);
-        LeaderAndIsrResponse response = new LeaderAndIsrResponse(Errors.NONE, errors);
+        List partitions = createPartitions("foo",
+            asList(Errors.NONE, Errors.CLUSTER_AUTHORIZATION_FAILED));
+        LeaderAndIsrResponse response = new LeaderAndIsrResponse(new LeaderAndIsrResponseData()
+            .setErrorCode(Errors.NONE.code())
+            .setPartitionErrors(partitions));
         String responseStr = response.toString();
         assertTrue(responseStr.contains(LeaderAndIsrResponse.class.getSimpleName()));
-        assertTrue(responseStr.contains(errors.toString()));
-        assertTrue(responseStr.contains(Errors.NONE.name()));
+        assertTrue(responseStr.contains(partitions.toString()));
+        assertTrue(responseStr.contains("errorCode=" + Errors.NONE.code()));
+    }
+
+    private List createPartitions(String topicName, List errors) {
+        List partitions = new ArrayList<>();
+        int partitionIndex = 0;
+        for (Errors error : errors) {
+            partitions.add(new LeaderAndIsrPartitionError()
+                .setTopicName(topicName)
+                .setPartitionIndex(partitionIndex++)
+                .setErrorCode(error.code()));
+        }
+        return partitions;
     }
 
 }
diff --git a/clients/src/test/java/org/apache/kafka/common/requests/LeaveGroupRequestTest.java b/clients/src/test/java/org/apache/kafka/common/requests/LeaveGroupRequestTest.java
new file mode 100644
index 0000000000000..2ff928bea5c0c
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/LeaveGroupRequestTest.java
@@ -0,0 +1,128 @@
+/*
+ * 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.requests;
+
+import org.apache.kafka.common.errors.UnsupportedVersionException;
+import org.apache.kafka.common.message.LeaveGroupRequestData;
+import org.apache.kafka.common.message.LeaveGroupRequestData.MemberIdentity;
+import org.apache.kafka.common.message.LeaveGroupResponseData;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.Errors;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
+
+public class LeaveGroupRequestTest {
+
+    private final String groupId = "group_id";
+    private final String memberIdOne = "member_1";
+    private final String instanceIdOne = "instance_1";
+    private final String memberIdTwo = "member_2";
+    private final String instanceIdTwo = "instance_2";
+
+    private final int throttleTimeMs = 10;
+
+    private LeaveGroupRequest.Builder builder;
+    private List members;
+
+    @Before
+    public void setUp() {
+        members = Arrays.asList(new MemberIdentity()
+                                         .setMemberId(memberIdOne)
+                                         .setGroupInstanceId(instanceIdOne),
+                                new MemberIdentity()
+                                         .setMemberId(memberIdTwo)
+                                         .setGroupInstanceId(instanceIdTwo));
+        builder = new LeaveGroupRequest.Builder(
+            groupId,
+            members
+        );
+    }
+
+    @Test
+    public void testMultiLeaveConstructor() {
+        final LeaveGroupRequestData expectedData = new LeaveGroupRequestData()
+                                                       .setGroupId(groupId)
+                                                       .setMembers(members);
+
+        for (short version = 0; version <= ApiKeys.LEAVE_GROUP.latestVersion(); version++) {
+            try {
+                LeaveGroupRequest request = builder.build(version);
+                if (version <= 2) {
+                    fail("Older version " + version +
+                             " request data should not be created due to non-single members");
+                }
+                assertEquals(expectedData, request.data());
+                assertEquals(members, request.members());
+
+                LeaveGroupResponse expectedResponse = new LeaveGroupResponse(
+                    Collections.emptyList(),
+                    Errors.COORDINATOR_LOAD_IN_PROGRESS,
+                    throttleTimeMs,
+                    version
+                );
+
+                assertEquals(expectedResponse, request.getErrorResponse(throttleTimeMs,
+                                                                        Errors.COORDINATOR_LOAD_IN_PROGRESS.exception()));
+            } catch (UnsupportedVersionException e) {
+                assertTrue(e.getMessage().contains("leave group request only supports single member instance"));
+            }
+        }
+
+    }
+
+    @Test
+    public void testSingleLeaveConstructor() {
+        final LeaveGroupRequestData expectedData = new LeaveGroupRequestData()
+                                                       .setGroupId(groupId)
+                                                       .setMemberId(memberIdOne);
+        List singleMember = Collections.singletonList(
+            new MemberIdentity()
+                .setMemberId(memberIdOne));
+
+        builder = new LeaveGroupRequest.Builder(groupId, singleMember);
+
+        for (short version = 0; version <= 2; version++) {
+            LeaveGroupRequest request = builder.build(version);
+            assertEquals(expectedData, request.data());
+            assertEquals(singleMember, request.members());
+
+            int expectedThrottleTime = version >= 1 ? throttleTimeMs
+                                           : AbstractResponse.DEFAULT_THROTTLE_TIME;
+            LeaveGroupResponse expectedResponse = new LeaveGroupResponse(
+                new LeaveGroupResponseData()
+                    .setErrorCode(Errors.NOT_CONTROLLER.code())
+                    .setThrottleTimeMs(expectedThrottleTime)
+            );
+
+            assertEquals(expectedResponse, request.getErrorResponse(throttleTimeMs,
+                                                                    Errors.NOT_CONTROLLER.exception()));
+        }
+    }
+
+    @Test(expected = IllegalArgumentException.class)
+    public void testBuildEmptyMembers() {
+        new LeaveGroupRequest.Builder(groupId, Collections.emptyList());
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/requests/LeaveGroupResponseTest.java b/clients/src/test/java/org/apache/kafka/common/requests/LeaveGroupResponseTest.java
new file mode 100644
index 0000000000000..187640bfdeb71
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/LeaveGroupResponseTest.java
@@ -0,0 +1,165 @@
+/*
+ * 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.requests;
+
+import org.apache.kafka.common.message.LeaveGroupResponseData;
+import org.apache.kafka.common.message.LeaveGroupResponseData.MemberResponse;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.Errors;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import static org.apache.kafka.common.requests.AbstractResponse.DEFAULT_THROTTLE_TIME;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+public class LeaveGroupResponseTest {
+
+    private final String memberIdOne = "member_1";
+    private final String instanceIdOne = "instance_1";
+    private final String memberIdTwo = "member_2";
+    private final String instanceIdTwo = "instance_2";
+
+    private final int throttleTimeMs = 10;
+
+    private List memberResponses;
+
+    @Before
+    public void setUp() {
+        memberResponses = Arrays.asList(new MemberResponse()
+                                            .setMemberId(memberIdOne)
+                                            .setGroupInstanceId(instanceIdOne)
+                                            .setErrorCode(Errors.UNKNOWN_MEMBER_ID.code()),
+                                        new MemberResponse()
+                                            .setMemberId(memberIdTwo)
+                                            .setGroupInstanceId(instanceIdTwo)
+                                            .setErrorCode(Errors.FENCED_INSTANCE_ID.code())
+        );
+    }
+
+    @Test
+    public void testConstructorWithStruct() {
+        Map expectedErrorCounts = Collections.singletonMap(Errors.NOT_COORDINATOR, 1);
+
+        LeaveGroupResponseData responseData = new LeaveGroupResponseData()
+                                                  .setErrorCode(Errors.NOT_COORDINATOR.code())
+                                                  .setThrottleTimeMs(throttleTimeMs);
+        for (short version = 0; version <= ApiKeys.LEAVE_GROUP.latestVersion(); version++) {
+            LeaveGroupResponse leaveGroupResponse = new LeaveGroupResponse(responseData.toStruct(version), version);
+
+            assertEquals(expectedErrorCounts, leaveGroupResponse.errorCounts());
+
+            if (version >= 1) {
+                assertEquals(throttleTimeMs, leaveGroupResponse.throttleTimeMs());
+            } else {
+                assertEquals(DEFAULT_THROTTLE_TIME, leaveGroupResponse.throttleTimeMs());
+            }
+
+            assertEquals(Errors.NOT_COORDINATOR, leaveGroupResponse.error());
+        }
+    }
+
+
+    @Test
+    public void testConstructorWithMemberResponses() {
+        Map expectedErrorCounts = new HashMap<>();
+        expectedErrorCounts.put(Errors.UNKNOWN_MEMBER_ID, 1);
+        expectedErrorCounts.put(Errors.FENCED_INSTANCE_ID, 1);
+
+        for (short version = 0; version <= ApiKeys.LEAVE_GROUP.latestVersion(); version++) {
+            LeaveGroupResponse leaveGroupResponse = new LeaveGroupResponse(memberResponses,
+                                                                           Errors.NONE,
+                                                                           throttleTimeMs,
+                                                                           version);
+
+            if (version >= 3) {
+                assertEquals(expectedErrorCounts, leaveGroupResponse.errorCounts());
+                assertEquals(memberResponses, leaveGroupResponse.memberResponses());
+            } else {
+                assertEquals(Collections.singletonMap(Errors.UNKNOWN_MEMBER_ID, 1),
+                             leaveGroupResponse.errorCounts());
+                assertEquals(Collections.emptyList(), leaveGroupResponse.memberResponses());
+            }
+
+            if (version >= 1) {
+                assertEquals(throttleTimeMs, leaveGroupResponse.throttleTimeMs());
+            } else {
+                assertEquals(DEFAULT_THROTTLE_TIME, leaveGroupResponse.throttleTimeMs());
+            }
+
+            assertEquals(Errors.UNKNOWN_MEMBER_ID, leaveGroupResponse.error());
+        }
+    }
+
+    @Test
+    public void testShouldThrottle() {
+        LeaveGroupResponse response = new LeaveGroupResponse(new LeaveGroupResponseData());
+        for (short version = 0; version <= ApiKeys.LEAVE_GROUP.latestVersion(); version++) {
+            if (version >= 2) {
+                assertTrue(response.shouldClientThrottle(version));
+            } else {
+                assertFalse(response.shouldClientThrottle(version));
+            }
+        }
+    }
+
+    @Test
+    public void testEqualityWithStruct() {
+        LeaveGroupResponseData responseData = new LeaveGroupResponseData()
+            .setErrorCode(Errors.NONE.code())
+            .setThrottleTimeMs(throttleTimeMs);
+        for (short version = 0; version <= ApiKeys.LEAVE_GROUP.latestVersion(); version++) {
+            LeaveGroupResponse primaryResponse = new LeaveGroupResponse(responseData.toStruct(version), version);
+
+            LeaveGroupResponse secondaryResponse = new LeaveGroupResponse(responseData.toStruct(version), version);
+
+            assertEquals(primaryResponse, primaryResponse);
+            assertEquals(primaryResponse, secondaryResponse);
+            assertEquals(primaryResponse.hashCode(), secondaryResponse.hashCode());
+
+        }
+    }
+
+    @Test
+    public void testEqualityWithMemberResponses() {
+        for (short version = 0; version <= ApiKeys.LEAVE_GROUP.latestVersion(); version++) {
+            List localResponses = version > 2 ? memberResponses : memberResponses.subList(0, 1);
+            LeaveGroupResponse primaryResponse = new LeaveGroupResponse(localResponses,
+                                                                        Errors.NONE,
+                                                                        throttleTimeMs,
+                                                                        version);
+
+            // The order of members should not alter result data.
+            Collections.reverse(localResponses);
+            LeaveGroupResponse reversedResponse = new LeaveGroupResponse(localResponses,
+                                                                         Errors.NONE,
+                                                                         throttleTimeMs,
+                                                                         version);
+
+            assertEquals(primaryResponse, primaryResponse);
+            assertEquals(primaryResponse, reversedResponse);
+            assertEquals(primaryResponse.hashCode(), reversedResponse.hashCode());
+        }
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/requests/MetadataRequestTest.java b/clients/src/test/java/org/apache/kafka/common/requests/MetadataRequestTest.java
index 207cac7670fc0..31e22625ce8d1 100644
--- a/clients/src/test/java/org/apache/kafka/common/requests/MetadataRequestTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/requests/MetadataRequestTest.java
@@ -16,8 +16,8 @@
  */
 package org.apache.kafka.common.requests;
 
-import org.apache.kafka.common.protocol.types.Schema;
-import org.apache.kafka.common.protocol.types.Struct;
+import org.apache.kafka.common.message.MetadataRequestData;
+import org.apache.kafka.common.protocol.ApiKeys;
 import org.junit.Test;
 
 import java.util.Collections;
@@ -31,25 +31,38 @@ public class MetadataRequestTest {
 
     @Test
     public void testEmptyMeansAllTopicsV0() {
-        Struct rawRequest = new Struct(MetadataRequest.schemaVersions()[0]);
-        rawRequest.set("topics", new Object[0]);
-        MetadataRequest parsedRequest = new MetadataRequest(rawRequest, (short) 0);
+        MetadataRequestData data = new MetadataRequestData();
+        MetadataRequest parsedRequest = new MetadataRequest(data, (short) 0);
         assertTrue(parsedRequest.isAllTopics());
         assertNull(parsedRequest.topics());
     }
 
     @Test
     public void testEmptyMeansEmptyForVersionsAboveV0() {
-        for (int i = 1; i < MetadataRequest.schemaVersions().length; i++) {
-            Schema schema = MetadataRequest.schemaVersions()[i];
-            Struct rawRequest = new Struct(schema);
-            rawRequest.set("topics", new Object[0]);
-            if (rawRequest.hasField("allow_auto_topic_creation"))
-                rawRequest.set("allow_auto_topic_creation", true);
-            MetadataRequest parsedRequest = new MetadataRequest(rawRequest, (short) i);
+        for (int i = 1; i < MetadataRequestData.SCHEMAS.length; i++) {
+            MetadataRequestData data = new MetadataRequestData();
+            data.setAllowAutoTopicCreation(true);
+            MetadataRequest parsedRequest = new MetadataRequest(data, (short) i);
             assertFalse(parsedRequest.isAllTopics());
             assertEquals(Collections.emptyList(), parsedRequest.topics());
         }
     }
 
+    @Test
+    public void testMetadataRequestVersion() {
+        MetadataRequest.Builder builder = new MetadataRequest.Builder(Collections.singletonList("topic"), false);
+        assertEquals(ApiKeys.METADATA.oldestVersion(), builder.oldestAllowedVersion());
+        assertEquals(ApiKeys.METADATA.latestVersion(), builder.latestAllowedVersion());
+
+        short version = 5;
+        MetadataRequest.Builder builder2 = new MetadataRequest.Builder(Collections.singletonList("topic"), false, version);
+        assertEquals(version, builder2.oldestAllowedVersion());
+        assertEquals(version, builder2.latestAllowedVersion());
+
+        short minVersion = 1;
+        short maxVersion = 6;
+        MetadataRequest.Builder builder3 = new MetadataRequest.Builder(Collections.singletonList("topic"), false, minVersion, maxVersion);
+        assertEquals(minVersion, builder3.oldestAllowedVersion());
+        assertEquals(maxVersion, builder3.latestAllowedVersion());
+    }
 }
diff --git a/clients/src/test/java/org/apache/kafka/common/requests/OffsetCommitRequestTest.java b/clients/src/test/java/org/apache/kafka/common/requests/OffsetCommitRequestTest.java
new file mode 100644
index 0000000000000..c6c5a7bea24c4
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/OffsetCommitRequestTest.java
@@ -0,0 +1,146 @@
+/*
+ * 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.requests;
+
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.errors.UnsupportedVersionException;
+import org.apache.kafka.common.message.OffsetCommitRequestData;
+import org.apache.kafka.common.message.OffsetCommitRequestData.OffsetCommitRequestPartition;
+import org.apache.kafka.common.message.OffsetCommitRequestData.OffsetCommitRequestTopic;
+import org.apache.kafka.common.message.OffsetCommitResponseData.OffsetCommitResponsePartition;
+import org.apache.kafka.common.message.OffsetCommitResponseData.OffsetCommitResponseTopic;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.Errors;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import static org.apache.kafka.common.requests.AbstractResponse.DEFAULT_THROTTLE_TIME;
+import static org.apache.kafka.common.requests.OffsetCommitRequest.getErrorResponseTopics;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
+
+public class OffsetCommitRequestTest {
+
+    protected static String groupId = "groupId";
+    protected static String topicOne = "topicOne";
+    protected static String topicTwo = "topicTwo";
+    protected static int partitionOne = 1;
+    protected static int partitionTwo = 2;
+    protected static long offset = 100L;
+    protected static short leaderEpoch = 20;
+    protected static String metadata = "metadata";
+
+    protected static int throttleTimeMs = 10;
+
+    private static OffsetCommitRequestData data;
+    private static List topics;
+
+    @Before
+    public void setUp() {
+        topics = Arrays.asList(
+            new OffsetCommitRequestTopic()
+                .setName(topicOne)
+                .setPartitions(Collections.singletonList(
+                    new OffsetCommitRequestPartition()
+                        .setPartitionIndex(partitionOne)
+                        .setCommittedOffset(offset)
+                        .setCommittedLeaderEpoch(leaderEpoch)
+                        .setCommittedMetadata(metadata)
+                )),
+            new OffsetCommitRequestTopic()
+                .setName(topicTwo)
+                .setPartitions(Collections.singletonList(
+                    new OffsetCommitRequestPartition()
+                        .setPartitionIndex(partitionTwo)
+                        .setCommittedOffset(offset)
+                        .setCommittedLeaderEpoch(leaderEpoch)
+                        .setCommittedMetadata(metadata)
+                ))
+        );
+        data = new OffsetCommitRequestData()
+                   .setGroupId(groupId)
+                   .setTopics(topics);
+    }
+
+    @Test
+    public void testConstructor() {
+        Map expectedOffsets = new HashMap<>();
+        expectedOffsets.put(new TopicPartition(topicOne, partitionOne), offset);
+        expectedOffsets.put(new TopicPartition(topicTwo, partitionTwo), offset);
+
+        OffsetCommitRequest.Builder builder = new OffsetCommitRequest.Builder(data);
+
+        for (short version = 0; version <= ApiKeys.TXN_OFFSET_COMMIT.latestVersion(); version++) {
+            OffsetCommitRequest request = builder.build(version);
+            assertEquals(expectedOffsets, request.offsets());
+
+            OffsetCommitResponse response = request.getErrorResponse(throttleTimeMs, Errors.NOT_COORDINATOR.exception());
+
+            assertEquals(Collections.singletonMap(Errors.NOT_COORDINATOR, 2), response.errorCounts());
+
+            if (version >= 3) {
+                assertEquals(throttleTimeMs, response.throttleTimeMs());
+            } else {
+                assertEquals(DEFAULT_THROTTLE_TIME, response.throttleTimeMs());
+            }
+        }
+    }
+
+    @Test
+    public void testGetErrorResponseTopics() {
+        List expectedTopics = Arrays.asList(
+            new OffsetCommitResponseTopic()
+                .setName(topicOne)
+                .setPartitions(Collections.singletonList(
+                    new OffsetCommitResponsePartition()
+                        .setErrorCode(Errors.UNKNOWN_MEMBER_ID.code())
+                        .setPartitionIndex(partitionOne))),
+            new OffsetCommitResponseTopic()
+                .setName(topicTwo)
+                .setPartitions(Collections.singletonList(
+                    new OffsetCommitResponsePartition()
+                        .setErrorCode(Errors.UNKNOWN_MEMBER_ID.code())
+                        .setPartitionIndex(partitionTwo)))
+        );
+        assertEquals(expectedTopics, getErrorResponseTopics(topics, Errors.UNKNOWN_MEMBER_ID));
+    }
+
+    @Test
+    public void testVersionSupportForGroupInstanceId() {
+        OffsetCommitRequest.Builder builder = new OffsetCommitRequest.Builder(
+            new OffsetCommitRequestData()
+                .setGroupId("groupId")
+                .setMemberId("consumerId")
+                .setGroupInstanceId("groupInstanceId")
+        );
+
+        for (short version = 0; version <= ApiKeys.OFFSET_COMMIT.latestVersion(); version++) {
+            if (version >= 7) {
+                builder.build(version);
+            } else {
+                final short finalVersion = version;
+                assertThrows(UnsupportedVersionException.class, () -> builder.build(finalVersion));
+            }
+        }
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/requests/OffsetCommitResponseTest.java b/clients/src/test/java/org/apache/kafka/common/requests/OffsetCommitResponseTest.java
new file mode 100644
index 0000000000000..62229168b2405
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/OffsetCommitResponseTest.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.common.requests;
+
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.message.OffsetCommitResponseData;
+import org.apache.kafka.common.message.OffsetCommitResponseData.OffsetCommitResponsePartition;
+import org.apache.kafka.common.message.OffsetCommitResponseData.OffsetCommitResponseTopic;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.Errors;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
+import static org.apache.kafka.common.requests.AbstractResponse.DEFAULT_THROTTLE_TIME;
+import static org.junit.Assert.assertEquals;
+
+public class OffsetCommitResponseTest {
+
+    protected final int throttleTimeMs = 10;
+
+    protected final String topicOne = "topic1";
+    protected final int partitionOne = 1;
+    protected final Errors errorOne = Errors.COORDINATOR_NOT_AVAILABLE;
+    protected final Errors errorTwo = Errors.NOT_COORDINATOR;
+    protected final String topicTwo = "topic2";
+    protected final int partitionTwo = 2;
+
+    protected TopicPartition tp1 = new TopicPartition(topicOne, partitionOne);
+    protected TopicPartition tp2 = new TopicPartition(topicTwo, partitionTwo);
+    protected Map expectedErrorCounts;
+    protected Map errorsMap;
+
+    @Before
+    public void setUp() {
+        expectedErrorCounts = new HashMap<>();
+        expectedErrorCounts.put(errorOne, 1);
+        expectedErrorCounts.put(errorTwo, 1);
+
+        errorsMap = new HashMap<>();
+        errorsMap.put(tp1, errorOne);
+        errorsMap.put(tp2, errorTwo);
+    }
+
+    @Test
+    public void testConstructorWithErrorResponse() {
+        OffsetCommitResponse response = new OffsetCommitResponse(throttleTimeMs, errorsMap);
+
+        assertEquals(expectedErrorCounts, response.errorCounts());
+        assertEquals(throttleTimeMs, response.throttleTimeMs());
+    }
+
+    @Test
+    public void testConstructorWithStruct() {
+        OffsetCommitResponseData data = new OffsetCommitResponseData()
+                                            .setTopics(Arrays.asList(
+                                                new OffsetCommitResponseTopic().setPartitions(
+                                                    Collections.singletonList(new OffsetCommitResponsePartition()
+                                                                                  .setPartitionIndex(partitionOne)
+                                                                                  .setErrorCode(errorOne.code()))),
+                                                new OffsetCommitResponseTopic().setPartitions(
+                                                    Collections.singletonList(new OffsetCommitResponsePartition()
+                                                                                  .setPartitionIndex(partitionTwo)
+                                                                                  .setErrorCode(errorTwo.code())))
+                                            ))
+                                            .setThrottleTimeMs(throttleTimeMs);
+
+        for (short version = 0; version <= ApiKeys.OFFSET_COMMIT.latestVersion(); version++) {
+            OffsetCommitResponse response = new OffsetCommitResponse(data.toStruct(version), version);
+            assertEquals(expectedErrorCounts, response.errorCounts());
+
+            if (version >= 3) {
+                assertEquals(throttleTimeMs, response.throttleTimeMs());
+            } else {
+                assertEquals(DEFAULT_THROTTLE_TIME, response.throttleTimeMs());
+            }
+
+            assertEquals(version >= 4, response.shouldClientThrottle(version));
+        }
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/requests/OffsetFetchRequestTest.java b/clients/src/test/java/org/apache/kafka/common/requests/OffsetFetchRequestTest.java
new file mode 100644
index 0000000000000..3bafb76561ae8
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/OffsetFetchRequestTest.java
@@ -0,0 +1,115 @@
+/*
+ * 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.requests;
+
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.errors.UnsupportedVersionException;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.common.requests.OffsetFetchResponse.PartitionData;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+
+import static org.apache.kafka.common.requests.AbstractResponse.DEFAULT_THROTTLE_TIME;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
+
+public class OffsetFetchRequestTest {
+
+    private final String topicOne = "topic1";
+    private final int partitionOne = 1;
+    private final String topicTwo = "topic2";
+    private final int partitionTwo = 2;
+    private final String groupId = "groupId";
+
+    private OffsetFetchRequest.Builder builder;
+    private List partitions;
+
+    @Before
+    public void setUp() {
+        partitions = Arrays.asList(new TopicPartition(topicOne, partitionOne),
+                                   new TopicPartition(topicTwo, partitionTwo));
+        builder = new OffsetFetchRequest.Builder(
+            groupId,
+            partitions
+        );
+    }
+
+    @Test
+    public void testConstructor() {
+        assertFalse(builder.isAllTopicPartitions());
+        int throttleTimeMs = 10;
+
+        Map expectedData = new HashMap<>();
+        for (TopicPartition partition : partitions) {
+            expectedData.put(partition, new PartitionData(
+                OffsetFetchResponse.INVALID_OFFSET,
+                Optional.empty(),
+                OffsetFetchResponse.NO_METADATA,
+                Errors.NONE
+            ));
+        }
+
+        for (short version = 0; version <= ApiKeys.OFFSET_FETCH.latestVersion(); version++) {
+            OffsetFetchRequest request = builder.build(version);
+            assertFalse(request.isAllPartitions());
+            assertEquals(groupId, request.groupId());
+            assertEquals(partitions, request.partitions());
+
+            OffsetFetchResponse response = request.getErrorResponse(throttleTimeMs, Errors.NONE);
+            assertEquals(Errors.NONE, response.error());
+            assertFalse(response.hasError());
+            assertEquals(Collections.singletonMap(Errors.NONE, 1), response.errorCounts());
+
+            if (version <= 1) {
+                assertEquals(expectedData, response.responseData());
+            }
+
+            if (version >= 3) {
+                assertEquals(throttleTimeMs, response.throttleTimeMs());
+            } else {
+                assertEquals(DEFAULT_THROTTLE_TIME, response.throttleTimeMs());
+            }
+        }
+    }
+
+    @Test
+    public void testConstructorFailForUnsupportedAllPartition() {
+        builder = OffsetFetchRequest.Builder.allTopicPartitions(groupId);
+        for (short version = 0; version <= ApiKeys.OFFSET_FETCH.latestVersion(); version++) {
+            short finalVersion = version;
+            if (version <= 1) {
+                assertThrows(UnsupportedVersionException.class, () -> builder.build(finalVersion));
+            } else {
+                OffsetFetchRequest request = builder.build(finalVersion);
+                assertEquals(groupId, request.groupId());
+                assertNull(request.partitions());
+                assertTrue(request.isAllPartitions());
+            }
+        }
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/requests/OffsetFetchResponseTest.java b/clients/src/test/java/org/apache/kafka/common/requests/OffsetFetchResponseTest.java
new file mode 100644
index 0000000000000..d3ff161ac0b07
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/OffsetFetchResponseTest.java
@@ -0,0 +1,223 @@
+/*
+ * 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.requests;
+
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.message.OffsetFetchResponseData;
+import org.apache.kafka.common.message.OffsetFetchResponseData.OffsetFetchResponsePartition;
+import org.apache.kafka.common.message.OffsetFetchResponseData.OffsetFetchResponseTopic;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.common.protocol.types.Struct;
+import org.apache.kafka.common.record.RecordBatch;
+import org.apache.kafka.common.requests.OffsetFetchResponse.PartitionData;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+
+import static org.apache.kafka.common.protocol.CommonFields.ERROR_CODE;
+import static org.apache.kafka.common.requests.AbstractResponse.DEFAULT_THROTTLE_TIME;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+public class OffsetFetchResponseTest {
+
+    private final int throttleTimeMs = 10;
+    private final int offset = 100;
+    private final String metadata = "metadata";
+
+    private final String topicOne = "topic1";
+    private final int partitionOne = 1;
+    private final Optional leaderEpochOne = Optional.of(1);
+    private final String topicTwo = "topic2";
+    private final int partitionTwo = 2;
+    private final Optional leaderEpochTwo = Optional.of(2);
+
+    private Map partitionDataMap;
+
+    @Before
+    public void setUp() {
+        partitionDataMap = new HashMap<>();
+        partitionDataMap.put(new TopicPartition(topicOne, partitionOne), new PartitionData(
+            offset,
+            leaderEpochOne,
+            metadata,
+            Errors.TOPIC_AUTHORIZATION_FAILED
+        ));
+        partitionDataMap.put(new TopicPartition(topicTwo, partitionTwo), new PartitionData(
+            offset,
+            leaderEpochTwo,
+            metadata,
+            Errors.UNKNOWN_TOPIC_OR_PARTITION
+        ));
+    }
+
+    @Test
+    public void testConstructor() {
+        OffsetFetchResponse response = new OffsetFetchResponse(throttleTimeMs, Errors.NOT_COORDINATOR, partitionDataMap);
+        assertEquals(Errors.NOT_COORDINATOR, response.error());
+        assertEquals(Collections.singletonMap(Errors.NOT_COORDINATOR, 1), response.errorCounts());
+
+        assertEquals(throttleTimeMs, response.throttleTimeMs());
+
+        Map responseData = response.responseData();
+        assertEquals(partitionDataMap, responseData);
+        responseData.forEach(
+            (tp, data) -> assertTrue(data.hasError())
+        );
+    }
+
+    /**
+     * Test behavior changes over the versions. Refer to resources.common.messages.OffsetFetchResponse.json
+     */
+    @Test
+    public void testStructBuild() {
+        partitionDataMap.put(new TopicPartition(topicTwo, partitionTwo), new PartitionData(
+            offset,
+            leaderEpochTwo,
+            metadata,
+            Errors.GROUP_AUTHORIZATION_FAILED
+        ));
+
+        OffsetFetchResponse latestResponse = new OffsetFetchResponse(throttleTimeMs, Errors.NONE, partitionDataMap);
+
+        for (short version = 0; version <= ApiKeys.OFFSET_FETCH.latestVersion(); version++) {
+            Struct struct = latestResponse.data.toStruct(version);
+
+            OffsetFetchResponse oldResponse =  new OffsetFetchResponse(struct, version);
+
+            if (version <= 1) {
+                assertFalse(struct.hasField(ERROR_CODE));
+
+                // Partition level error populated in older versions.
+                assertEquals(Errors.GROUP_AUTHORIZATION_FAILED, oldResponse.error());
+                assertEquals(Collections.singletonMap(Errors.GROUP_AUTHORIZATION_FAILED, 1), oldResponse.errorCounts());
+
+            } else {
+                assertTrue(struct.hasField(ERROR_CODE));
+
+                assertEquals(Errors.NONE, oldResponse.error());
+                assertEquals(Collections.singletonMap(Errors.NONE, 1), oldResponse.errorCounts());
+            }
+
+            if (version <= 2) {
+                assertEquals(DEFAULT_THROTTLE_TIME, oldResponse.throttleTimeMs());
+            } else {
+                assertEquals(throttleTimeMs, oldResponse.throttleTimeMs());
+            }
+
+            Map expectedDataMap = new HashMap<>();
+            for (Map.Entry entry : partitionDataMap.entrySet()) {
+                PartitionData partitionData = entry.getValue();
+                expectedDataMap.put(entry.getKey(), new PartitionData(
+                    partitionData.offset,
+                    version <= 4 ? Optional.empty() : partitionData.leaderEpoch,
+                    partitionData.metadata,
+                    partitionData.error
+                ));
+            }
+
+            Map responseData = oldResponse.responseData();
+            assertEquals(expectedDataMap, responseData);
+
+            responseData.forEach(
+                (tp, data) -> assertTrue(data.hasError())
+            );
+        }
+    }
+
+    @Test
+    public void testShouldThrottle() {
+        OffsetFetchResponse response = new OffsetFetchResponse(throttleTimeMs, Errors.NONE, partitionDataMap);
+        for (short version = 0; version <= ApiKeys.OFFSET_FETCH.latestVersion(); version++) {
+            if (version >= 4) {
+                assertTrue(response.shouldClientThrottle(version));
+            } else {
+                assertFalse(response.shouldClientThrottle(version));
+            }
+        }
+    }
+
+    @Test
+    public void testNullableMetadata() {
+        partitionDataMap.clear();
+        partitionDataMap.put(new TopicPartition(topicOne, partitionOne),
+                             new PartitionData(
+                                 offset,
+                                 leaderEpochOne,
+                                 null,
+                                 Errors.UNKNOWN_TOPIC_OR_PARTITION)
+        );
+
+        OffsetFetchResponse response = new OffsetFetchResponse(throttleTimeMs, Errors.GROUP_AUTHORIZATION_FAILED, partitionDataMap);
+        OffsetFetchResponseData expectedData =
+            new OffsetFetchResponseData()
+                .setErrorCode(Errors.GROUP_AUTHORIZATION_FAILED.code())
+                .setThrottleTimeMs(throttleTimeMs)
+                .setTopics(Collections.singletonList(
+                    new OffsetFetchResponseTopic()
+                        .setName(topicOne)
+                        .setPartitions(Collections.singletonList(
+                            new OffsetFetchResponsePartition()
+                                .setPartitionIndex(partitionOne)
+                                .setCommittedOffset(offset)
+                                .setCommittedLeaderEpoch(leaderEpochOne.orElse(-1))
+                                .setErrorCode(Errors.UNKNOWN_TOPIC_OR_PARTITION.code())
+                                .setMetadata(null))
+                        ))
+                );
+        assertEquals(expectedData, response.data);
+    }
+
+    @Test
+    public void testUseDefaultLeaderEpoch() {
+        final Optional emptyLeaderEpoch = Optional.empty();
+        partitionDataMap.clear();
+
+        partitionDataMap.put(new TopicPartition(topicOne, partitionOne),
+                             new PartitionData(
+                                 offset,
+                                 emptyLeaderEpoch,
+                                 metadata,
+                                 Errors.UNKNOWN_TOPIC_OR_PARTITION)
+        );
+
+        OffsetFetchResponse response = new OffsetFetchResponse(throttleTimeMs, Errors.NOT_COORDINATOR, partitionDataMap);
+        OffsetFetchResponseData expectedData =
+            new OffsetFetchResponseData()
+                .setErrorCode(Errors.NOT_COORDINATOR.code())
+                .setThrottleTimeMs(throttleTimeMs)
+                .setTopics(Collections.singletonList(
+                new OffsetFetchResponseTopic()
+                    .setName(topicOne)
+                    .setPartitions(Collections.singletonList(
+                        new OffsetFetchResponsePartition()
+                            .setPartitionIndex(partitionOne)
+                            .setCommittedOffset(offset)
+                            .setCommittedLeaderEpoch(RecordBatch.NO_PARTITION_LEADER_EPOCH)
+                            .setErrorCode(Errors.UNKNOWN_TOPIC_OR_PARTITION.code())
+                            .setMetadata(metadata))
+                    ))
+                );
+        assertEquals(expectedData, response.data);
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/requests/OffsetsForLeaderEpochRequestTest.java b/clients/src/test/java/org/apache/kafka/common/requests/OffsetsForLeaderEpochRequestTest.java
new file mode 100644
index 0000000000000..604f88f04aeaa
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/OffsetsForLeaderEpochRequestTest.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.common.requests;
+
+import org.apache.kafka.common.errors.UnsupportedVersionException;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.junit.Test;
+
+import java.util.Collections;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
+
+public class OffsetsForLeaderEpochRequestTest {
+
+    @Test
+    public void testForConsumerRequiresVersion3() {
+        OffsetsForLeaderEpochRequest.Builder builder = OffsetsForLeaderEpochRequest.Builder.forConsumer(Collections.emptyMap());
+        for (short version = 0; version < 3; version++) {
+            final short v = version;
+            assertThrows(UnsupportedVersionException.class, () -> builder.build(v));
+        }
+
+        for (short version = 3; version < ApiKeys.OFFSET_FOR_LEADER_EPOCH.latestVersion(); version++) {
+            OffsetsForLeaderEpochRequest request = builder.build((short) 3);
+            assertEquals(OffsetsForLeaderEpochRequest.CONSUMER_REPLICA_ID, request.replicaId());
+        }
+    }
+
+    @Test
+    public void testDefaultReplicaId() {
+        for (short version = 0; version < ApiKeys.OFFSET_FOR_LEADER_EPOCH.latestVersion(); version++) {
+            int replicaId = 1;
+            OffsetsForLeaderEpochRequest.Builder builder = OffsetsForLeaderEpochRequest.Builder.forFollower(
+                    version, Collections.emptyMap(), replicaId);
+            OffsetsForLeaderEpochRequest request = builder.build();
+            OffsetsForLeaderEpochRequest parsed = (OffsetsForLeaderEpochRequest) AbstractRequest.parseRequest(
+                    ApiKeys.OFFSET_FOR_LEADER_EPOCH, version, request.toStruct());
+            if (version < 3)
+                assertEquals(OffsetsForLeaderEpochRequest.DEBUGGING_REPLICA_ID, parsed.replicaId());
+            else
+                assertEquals(replicaId, parsed.replicaId());
+        }
+    }
+
+}
\ No newline at end of file
diff --git a/clients/src/test/java/org/apache/kafka/common/requests/ProduceRequestTest.java b/clients/src/test/java/org/apache/kafka/common/requests/ProduceRequestTest.java
index 809d64f1e327b..95d719adb3bb7 100644
--- a/clients/src/test/java/org/apache/kafka/common/requests/ProduceRequestTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/requests/ProduceRequestTest.java
@@ -17,10 +17,10 @@
 
 package org.apache.kafka.common.requests;
 
+import org.apache.kafka.common.InvalidRecordException;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.record.CompressionType;
-import org.apache.kafka.common.record.InvalidRecordException;
 import org.apache.kafka.common.record.MemoryRecords;
 import org.apache.kafka.common.record.MemoryRecordsBuilder;
 import org.apache.kafka.common.record.RecordBatch;
diff --git a/clients/src/test/java/org/apache/kafka/common/requests/ProduceResponseTest.java b/clients/src/test/java/org/apache/kafka/common/requests/ProduceResponseTest.java
new file mode 100644
index 0000000000000..ea6e99882d200
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/ProduceResponseTest.java
@@ -0,0 +1,115 @@
+/*
+ * 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.requests;
+
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.common.protocol.types.Struct;
+import org.apache.kafka.common.record.RecordBatch;
+import org.junit.Test;
+
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
+import static org.apache.kafka.common.protocol.ApiKeys.PRODUCE;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+public class ProduceResponseTest {
+
+    @Test
+    public void produceResponseV5Test() {
+        Map responseData = new HashMap<>();
+        TopicPartition tp0 = new TopicPartition("test", 0);
+        responseData.put(tp0, new ProduceResponse.PartitionResponse(Errors.NONE,
+                10000, RecordBatch.NO_TIMESTAMP, 100));
+
+        ProduceResponse v5Response = new ProduceResponse(responseData, 10);
+        short version = 5;
+
+        ByteBuffer buffer = v5Response.serialize(ApiKeys.PRODUCE, version, 0);
+        buffer.rewind();
+
+        ResponseHeader.parse(buffer, ApiKeys.PRODUCE.responseHeaderVersion(version)); // throw away.
+
+        Struct deserializedStruct = ApiKeys.PRODUCE.parseResponse(version, buffer);
+
+        ProduceResponse v5FromBytes = (ProduceResponse) AbstractResponse.parseResponse(ApiKeys.PRODUCE,
+                deserializedStruct, version);
+
+        assertEquals(1, v5FromBytes.responses().size());
+        assertTrue(v5FromBytes.responses().containsKey(tp0));
+        ProduceResponse.PartitionResponse partitionResponse = v5FromBytes.responses().get(tp0);
+        assertEquals(100, partitionResponse.logStartOffset);
+        assertEquals(10000, partitionResponse.baseOffset);
+        assertEquals(10, v5FromBytes.throttleTimeMs());
+        assertEquals(responseData, v5Response.responses());
+    }
+
+    @Test
+    public void produceResponseVersionTest() {
+        Map responseData = new HashMap<>();
+        responseData.put(new TopicPartition("test", 0), new ProduceResponse.PartitionResponse(Errors.NONE,
+                10000, RecordBatch.NO_TIMESTAMP, 100));
+        ProduceResponse v0Response = new ProduceResponse(responseData);
+        ProduceResponse v1Response = new ProduceResponse(responseData, 10);
+        ProduceResponse v2Response = new ProduceResponse(responseData, 10);
+        assertEquals("Throttle time must be zero", 0, v0Response.throttleTimeMs());
+        assertEquals("Throttle time must be 10", 10, v1Response.throttleTimeMs());
+        assertEquals("Throttle time must be 10", 10, v2Response.throttleTimeMs());
+        assertEquals("Should use schema version 0", ApiKeys.PRODUCE.responseSchema((short) 0),
+                v0Response.toStruct((short) 0).schema());
+        assertEquals("Should use schema version 1", ApiKeys.PRODUCE.responseSchema((short) 1),
+                v1Response.toStruct((short) 1).schema());
+        assertEquals("Should use schema version 2", ApiKeys.PRODUCE.responseSchema((short) 2),
+                v2Response.toStruct((short) 2).schema());
+        assertEquals("Response data does not match", responseData, v0Response.responses());
+        assertEquals("Response data does not match", responseData, v1Response.responses());
+        assertEquals("Response data does not match", responseData, v2Response.responses());
+    }
+
+    @Test
+    public void produceResponseRecordErrorsTest() {
+        Map responseData = new HashMap<>();
+        TopicPartition tp = new TopicPartition("test", 0);
+        ProduceResponse.PartitionResponse partResponse = new ProduceResponse.PartitionResponse(Errors.NONE,
+                10000, RecordBatch.NO_TIMESTAMP, 100,
+                Collections.singletonList(new ProduceResponse.RecordError(3, "Record error")),
+                "Produce failed");
+        responseData.put(tp, partResponse);
+
+        for (short ver = 0; ver <= PRODUCE.latestVersion(); ver++) {
+            ProduceResponse response = new ProduceResponse(responseData);
+            Struct struct = response.toStruct(ver);
+            assertEquals("Should use schema version " + ver, ApiKeys.PRODUCE.responseSchema(ver), struct.schema());
+            ProduceResponse.PartitionResponse deserialized = new ProduceResponse(struct).responses().get(tp);
+            if (ver >= 8) {
+                assertEquals(1, deserialized.recordErrors.size());
+                assertEquals(3, deserialized.recordErrors.get(0).batchIndex);
+                assertEquals("Record error", deserialized.recordErrors.get(0).message);
+                assertEquals("Produce failed", deserialized.errorMessage);
+            } else {
+                assertEquals(0, deserialized.recordErrors.size());
+                assertEquals(null, deserialized.errorMessage);
+            }
+        }
+    }
+}
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 7c24d1f2e1d04..7d066250d2f2b 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
@@ -16,6 +16,8 @@
  */
 package org.apache.kafka.common.requests;
 
+import org.apache.kafka.common.message.ApiVersionsResponseData;
+import org.apache.kafka.common.message.ApiVersionsResponseData.ApiVersionsResponseKeyCollection;
 import org.apache.kafka.common.network.ListenerName;
 import org.apache.kafka.common.network.Send;
 import org.apache.kafka.common.protocol.ApiKeys;
@@ -27,7 +29,6 @@
 
 import java.net.InetAddress;
 import java.nio.ByteBuffer;
-import java.util.Collections;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
@@ -55,8 +56,10 @@ public void testSerdeUnsupportedApiVersionRequest() throws Exception {
         ApiVersionsRequest request = (ApiVersionsRequest) requestAndSize.request;
         assertTrue(request.hasUnsupportedRequestVersion());
 
-        Send send = context.buildResponse(new ApiVersionsResponse(0, Errors.UNSUPPORTED_VERSION,
-                Collections.emptyList()));
+        Send send = context.buildResponse(new ApiVersionsResponse(new ApiVersionsResponseData()
+            .setThrottleTimeMs(0)
+            .setErrorCode(Errors.UNSUPPORTED_VERSION.code())
+            .setApiKeys(new ApiVersionsResponseKeyCollection())));
         ByteBufferChannel channel = new ByteBufferChannel(256);
         send.writeTo(channel);
 
@@ -64,14 +67,15 @@ public void testSerdeUnsupportedApiVersionRequest() throws Exception {
         responseBuffer.flip();
         responseBuffer.getInt(); // strip off the size
 
-        ResponseHeader responseHeader = ResponseHeader.parse(responseBuffer);
+        ResponseHeader responseHeader = ResponseHeader.parse(responseBuffer,
+            ApiKeys.API_VERSIONS.responseHeaderVersion(header.apiVersion()));
         assertEquals(correlationId, responseHeader.correlationId());
 
         Struct struct = ApiKeys.API_VERSIONS.parseResponse((short) 0, responseBuffer);
         ApiVersionsResponse response = (ApiVersionsResponse)
             AbstractResponse.parseResponse(ApiKeys.API_VERSIONS, struct, (short) 0);
-        assertEquals(Errors.UNSUPPORTED_VERSION, response.error());
-        assertTrue(response.apiVersions().isEmpty());
+        assertEquals(Errors.UNSUPPORTED_VERSION.code(), response.data.errorCode());
+        assertTrue(response.data.apiKeys().isEmpty());
     }
 
 }
diff --git a/clients/src/test/java/org/apache/kafka/common/requests/RequestHeaderTest.java b/clients/src/test/java/org/apache/kafka/common/requests/RequestHeaderTest.java
index f73ee2f3823b1..2842ae8e019ce 100644
--- a/clients/src/test/java/org/apache/kafka/common/requests/RequestHeaderTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/requests/RequestHeaderTest.java
@@ -43,6 +43,7 @@ public void testSerdeControlledShutdownV0() {
         assertEquals(0, deserialized.apiVersion());
         assertEquals(correlationId, deserialized.correlationId());
         assertEquals("", deserialized.clientId());
+        assertEquals(0, deserialized.headerVersion());
 
         Struct serialized = deserialized.toStruct();
         ByteBuffer serializedBuffer = toBuffer(serialized);
@@ -54,23 +55,22 @@ public void testSerdeControlledShutdownV0() {
     }
 
     @Test
-    public void testRequestHeader() {
+    public void testRequestHeaderV1() {
         RequestHeader header = new RequestHeader(ApiKeys.FIND_COORDINATOR, (short) 1, "", 10);
+        assertEquals(1, header.headerVersion());
         ByteBuffer buffer = toBuffer(header.toStruct());
+        assertEquals(10, buffer.remaining());
         RequestHeader deserialized = RequestHeader.parse(buffer);
         assertEquals(header, deserialized);
     }
 
     @Test
-    public void testRequestHeaderWithNullClientId() {
-        RequestHeader header = new RequestHeader(ApiKeys.FIND_COORDINATOR, (short) 1, null, 10);
-        Struct headerStruct = header.toStruct();
-        ByteBuffer buffer = toBuffer(headerStruct);
+    public void testRequestHeaderV2() {
+        RequestHeader header = new RequestHeader(ApiKeys.CREATE_DELEGATION_TOKEN, (short) 2, "", 10);
+        assertEquals(2, header.headerVersion());
+        ByteBuffer buffer = toBuffer(header.toStruct());
+        assertEquals(11, buffer.remaining());
         RequestHeader deserialized = RequestHeader.parse(buffer);
-        assertEquals(header.apiKey(), deserialized.apiKey());
-        assertEquals(header.apiVersion(), deserialized.apiVersion());
-        assertEquals(header.correlationId(), deserialized.correlationId());
-        assertEquals("", deserialized.clientId()); // null defaults to ""
+        assertEquals(header, deserialized);
     }
-
 }
diff --git a/clients/src/test/java/org/apache/kafka/common/requests/RequestResponseTest.java b/clients/src/test/java/org/apache/kafka/common/requests/RequestResponseTest.java
index 5d60086a84370..0b2012e160585 100644
--- a/clients/src/test/java/org/apache/kafka/common/requests/RequestResponseTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/requests/RequestResponseTest.java
@@ -18,6 +18,7 @@
 
 import org.apache.kafka.common.Node;
 import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.ElectionType;
 import org.apache.kafka.common.acl.AccessControlEntry;
 import org.apache.kafka.common.acl.AccessControlEntryFilter;
 import org.apache.kafka.common.acl.AclBinding;
@@ -32,29 +33,87 @@
 import org.apache.kafka.common.errors.SecurityDisabledException;
 import org.apache.kafka.common.errors.UnknownServerException;
 import org.apache.kafka.common.errors.UnsupportedVersionException;
-import org.apache.kafka.common.message.CreateTopicsRequestData;
+import org.apache.kafka.common.message.ApiVersionsRequestData;
+import org.apache.kafka.common.message.ApiVersionsResponseData;
+import org.apache.kafka.common.message.ApiVersionsResponseData.ApiVersionsResponseKey;
+import org.apache.kafka.common.message.ApiVersionsResponseData.ApiVersionsResponseKeyCollection;
+import org.apache.kafka.common.message.ControlledShutdownRequestData;
+import org.apache.kafka.common.message.ControlledShutdownResponseData.RemainingPartition;
+import org.apache.kafka.common.message.ControlledShutdownResponseData.RemainingPartitionCollection;
+import org.apache.kafka.common.message.ControlledShutdownResponseData;
+import org.apache.kafka.common.message.CreateDelegationTokenRequestData;
+import org.apache.kafka.common.message.CreateDelegationTokenRequestData.CreatableRenewers;
+import org.apache.kafka.common.message.CreateDelegationTokenResponseData;
 import org.apache.kafka.common.message.CreateTopicsRequestData.CreatableReplicaAssignment;
 import org.apache.kafka.common.message.CreateTopicsRequestData.CreatableTopic;
 import org.apache.kafka.common.message.CreateTopicsRequestData.CreateableTopicConfig;
-import org.apache.kafka.common.message.CreateTopicsResponseData;
+import org.apache.kafka.common.message.CreateTopicsRequestData;
 import org.apache.kafka.common.message.CreateTopicsResponseData.CreatableTopicResult;
+import org.apache.kafka.common.message.CreateTopicsResponseData;
+import org.apache.kafka.common.message.DeleteGroupsRequestData;
+import org.apache.kafka.common.message.DeleteGroupsResponseData;
+import org.apache.kafka.common.message.DeleteGroupsResponseData.DeletableGroupResult;
+import org.apache.kafka.common.message.DeleteGroupsResponseData.DeletableGroupResultCollection;
+import org.apache.kafka.common.message.DeleteTopicsRequestData;
+import org.apache.kafka.common.message.DeleteTopicsResponseData.DeletableTopicResult;
+import org.apache.kafka.common.message.DeleteTopicsResponseData;
 import org.apache.kafka.common.message.DescribeGroupsRequestData;
 import org.apache.kafka.common.message.DescribeGroupsResponseData;
-import org.apache.kafka.common.message.ElectPreferredLeadersRequestData;
-import org.apache.kafka.common.message.ElectPreferredLeadersRequestData.TopicPartitions;
-import org.apache.kafka.common.message.ElectPreferredLeadersResponseData;
-import org.apache.kafka.common.message.ElectPreferredLeadersResponseData.PartitionResult;
-import org.apache.kafka.common.message.ElectPreferredLeadersResponseData.ReplicaElectionResult;
-import org.apache.kafka.common.message.LeaveGroupRequestData;
+import org.apache.kafka.common.message.DescribeGroupsResponseData.DescribedGroup;
+import org.apache.kafka.common.message.ElectLeadersResponseData.PartitionResult;
+import org.apache.kafka.common.message.ElectLeadersResponseData.ReplicaElectionResult;
+import org.apache.kafka.common.message.ExpireDelegationTokenRequestData;
+import org.apache.kafka.common.message.ExpireDelegationTokenResponseData;
+import org.apache.kafka.common.message.FindCoordinatorRequestData;
+import org.apache.kafka.common.message.HeartbeatRequestData;
+import org.apache.kafka.common.message.HeartbeatResponseData;
+import org.apache.kafka.common.message.IncrementalAlterConfigsRequestData.AlterConfigsResource;
+import org.apache.kafka.common.message.IncrementalAlterConfigsRequestData.AlterableConfig;
+import org.apache.kafka.common.message.IncrementalAlterConfigsRequestData;
+import org.apache.kafka.common.message.IncrementalAlterConfigsResponseData.AlterConfigsResourceResponse;
+import org.apache.kafka.common.message.IncrementalAlterConfigsResponseData;
+import org.apache.kafka.common.message.AlterPartitionReassignmentsRequestData;
+import org.apache.kafka.common.message.AlterPartitionReassignmentsResponseData;
+import org.apache.kafka.common.message.ListPartitionReassignmentsRequestData;
+import org.apache.kafka.common.message.ListPartitionReassignmentsResponseData;
+import org.apache.kafka.common.message.InitProducerIdRequestData;
+import org.apache.kafka.common.message.InitProducerIdResponseData;
+import org.apache.kafka.common.message.JoinGroupRequestData;
+import org.apache.kafka.common.message.JoinGroupResponseData;
+import org.apache.kafka.common.message.LeaderAndIsrRequestData.LeaderAndIsrPartitionState;
+import org.apache.kafka.common.message.LeaderAndIsrResponseData;
+import org.apache.kafka.common.message.LeaveGroupRequestData.MemberIdentity;
 import org.apache.kafka.common.message.LeaveGroupResponseData;
+import org.apache.kafka.common.message.ListGroupsRequestData;
+import org.apache.kafka.common.message.ListGroupsResponseData;
+import org.apache.kafka.common.message.OffsetCommitRequestData;
+import org.apache.kafka.common.message.OffsetCommitResponseData;
+import org.apache.kafka.common.message.OffsetDeleteRequestData;
+import org.apache.kafka.common.message.OffsetDeleteRequestData.OffsetDeleteRequestPartition;
+import org.apache.kafka.common.message.OffsetDeleteRequestData.OffsetDeleteRequestTopic;
+import org.apache.kafka.common.message.OffsetDeleteRequestData.OffsetDeleteRequestTopicCollection;
+import org.apache.kafka.common.message.OffsetDeleteResponseData;
+import org.apache.kafka.common.message.OffsetDeleteResponseData.OffsetDeleteResponsePartition;
+import org.apache.kafka.common.message.OffsetDeleteResponseData.OffsetDeleteResponsePartitionCollection;
+import org.apache.kafka.common.message.OffsetDeleteResponseData.OffsetDeleteResponseTopic;
+import org.apache.kafka.common.message.OffsetDeleteResponseData.OffsetDeleteResponseTopicCollection;
+import org.apache.kafka.common.message.RenewDelegationTokenRequestData;
+import org.apache.kafka.common.message.RenewDelegationTokenResponseData;
 import org.apache.kafka.common.message.SaslAuthenticateRequestData;
 import org.apache.kafka.common.message.SaslAuthenticateResponseData;
 import org.apache.kafka.common.message.SaslHandshakeRequestData;
 import org.apache.kafka.common.message.SaslHandshakeResponseData;
+import org.apache.kafka.common.message.StopReplicaResponseData;
+import org.apache.kafka.common.message.TxnOffsetCommitRequestData;
+import org.apache.kafka.common.message.UpdateMetadataRequestData.UpdateMetadataBroker;
+import org.apache.kafka.common.message.UpdateMetadataRequestData.UpdateMetadataEndpoint;
+import org.apache.kafka.common.message.UpdateMetadataRequestData.UpdateMetadataPartitionState;
+import org.apache.kafka.common.message.UpdateMetadataResponseData;
 import org.apache.kafka.common.network.ListenerName;
 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.protocol.types.SchemaException;
 import org.apache.kafka.common.protocol.types.Struct;
 import org.apache.kafka.common.record.CompressionType;
 import org.apache.kafka.common.record.MemoryRecords;
@@ -63,8 +122,10 @@
 import org.apache.kafka.common.requests.CreateAclsRequest.AclCreation;
 import org.apache.kafka.common.requests.CreateAclsResponse.AclCreationResponse;
 import org.apache.kafka.common.requests.CreatePartitionsRequest.PartitionDetails;
+import org.apache.kafka.common.requests.CreateTopicsRequest.Builder;
 import org.apache.kafka.common.requests.DeleteAclsResponse.AclDeletionResult;
 import org.apache.kafka.common.requests.DeleteAclsResponse.AclFilterResponse;
+import org.apache.kafka.common.requests.FindCoordinatorRequest.CoordinatorType;
 import org.apache.kafka.common.resource.PatternType;
 import org.apache.kafka.common.resource.ResourcePattern;
 import org.apache.kafka.common.resource.ResourcePatternFilter;
@@ -95,11 +156,13 @@
 import java.util.Set;
 
 import static java.util.Arrays.asList;
-import static java.util.Collections.singletonList;
+import static org.apache.kafka.common.protocol.ApiKeys.FETCH;
 import static org.apache.kafka.common.requests.FetchMetadata.INVALID_SESSION_ID;
 import static org.apache.kafka.test.TestUtils.toBuffer;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
@@ -107,242 +170,266 @@ public class RequestResponseTest {
 
     @Test
     public void testSerialization() throws Exception {
-        checkRequest(createFindCoordinatorRequest(0));
-        checkRequest(createFindCoordinatorRequest(1));
-        checkErrorResponse(createFindCoordinatorRequest(0), new UnknownServerException());
-        checkErrorResponse(createFindCoordinatorRequest(1), new UnknownServerException());
-        checkResponse(createFindCoordinatorResponse(), 0);
-        checkResponse(createFindCoordinatorResponse(), 1);
-        checkRequest(createControlledShutdownRequest());
-        checkResponse(createControlledShutdownResponse(), 1);
-        checkErrorResponse(createControlledShutdownRequest(), new UnknownServerException());
-        checkErrorResponse(createControlledShutdownRequest(0), new UnknownServerException());
-        checkRequest(createFetchRequest(4));
-        checkResponse(createFetchResponse(), 4);
+        checkRequest(createFindCoordinatorRequest(0), true);
+        checkRequest(createFindCoordinatorRequest(1), true);
+        checkErrorResponse(createFindCoordinatorRequest(0), new UnknownServerException(), true);
+        checkErrorResponse(createFindCoordinatorRequest(1), new UnknownServerException(), true);
+        checkResponse(createFindCoordinatorResponse(), 0, true);
+        checkResponse(createFindCoordinatorResponse(), 1, true);
+        checkRequest(createControlledShutdownRequest(), true);
+        checkResponse(createControlledShutdownResponse(), 1, true);
+        checkErrorResponse(createControlledShutdownRequest(), new UnknownServerException(), true);
+        checkErrorResponse(createControlledShutdownRequest(0), new UnknownServerException(), true);
+        checkRequest(createFetchRequest(4), true);
+        checkResponse(createFetchResponse(), 4, true);
         List toForgetTopics = new ArrayList<>();
         toForgetTopics.add(new TopicPartition("foo", 0));
         toForgetTopics.add(new TopicPartition("foo", 2));
         toForgetTopics.add(new TopicPartition("bar", 0));
-        checkRequest(createFetchRequest(7, new FetchMetadata(123, 456), toForgetTopics));
-        checkResponse(createFetchResponse(123), 7);
-        checkResponse(createFetchResponse(Errors.FETCH_SESSION_ID_NOT_FOUND, 123), 7);
-        checkErrorResponse(createFetchRequest(4), new UnknownServerException());
-        checkRequest(createHeartBeatRequest());
-        checkErrorResponse(createHeartBeatRequest(), new UnknownServerException());
-        checkResponse(createHeartBeatResponse(), 0);
-        checkRequest(createJoinGroupRequest(1));
-        checkErrorResponse(createJoinGroupRequest(0), new UnknownServerException());
-        checkErrorResponse(createJoinGroupRequest(1), new UnknownServerException());
-        checkResponse(createJoinGroupResponse(), 0);
-        checkRequest(createLeaveGroupRequest());
-        checkErrorResponse(createLeaveGroupRequest(), new UnknownServerException());
-        checkResponse(createLeaveGroupResponse(), 0);
-        checkRequest(createListGroupsRequest());
-        checkErrorResponse(createListGroupsRequest(), new UnknownServerException());
-        checkResponse(createListGroupsResponse(), 0);
-        checkRequest(createDescribeGroupRequest());
-        checkErrorResponse(createDescribeGroupRequest(), new UnknownServerException());
-        checkResponse(createDescribeGroupResponse(), 0);
-        checkRequest(createDeleteGroupsRequest());
-        checkErrorResponse(createDeleteGroupsRequest(), new UnknownServerException());
-        checkResponse(createDeleteGroupsResponse(), 0);
-        checkRequest(createListOffsetRequest(1));
-        checkErrorResponse(createListOffsetRequest(1), new UnknownServerException());
-        checkResponse(createListOffsetResponse(1), 1);
-        checkRequest(createListOffsetRequest(2));
-        checkErrorResponse(createListOffsetRequest(2), new UnknownServerException());
-        checkResponse(createListOffsetResponse(2), 2);
-        checkRequest(MetadataRequest.Builder.allTopics().build((short) 2));
-        checkRequest(createMetadataRequest(1, singletonList("topic1")));
-        checkErrorResponse(createMetadataRequest(1, singletonList("topic1")), new UnknownServerException());
-        checkResponse(createMetadataResponse(), 2);
-        checkErrorResponse(createMetadataRequest(2, singletonList("topic1")), new UnknownServerException());
-        checkResponse(createMetadataResponse(), 3);
-        checkErrorResponse(createMetadataRequest(3, singletonList("topic1")), new UnknownServerException());
-        checkResponse(createMetadataResponse(), 4);
-        checkErrorResponse(createMetadataRequest(4, singletonList("topic1")), new UnknownServerException());
-        checkRequest(OffsetFetchRequest.forAllPartitions("group1"));
-        checkErrorResponse(OffsetFetchRequest.forAllPartitions("group1"), new NotCoordinatorException("Not Coordinator"));
-        checkRequest(createOffsetFetchRequest(0));
-        checkRequest(createOffsetFetchRequest(1));
-        checkRequest(createOffsetFetchRequest(2));
-        checkRequest(OffsetFetchRequest.forAllPartitions("group1"));
-        checkErrorResponse(createOffsetFetchRequest(0), new UnknownServerException());
-        checkErrorResponse(createOffsetFetchRequest(1), new UnknownServerException());
-        checkErrorResponse(createOffsetFetchRequest(2), new UnknownServerException());
-        checkResponse(createOffsetFetchResponse(), 0);
-        checkRequest(createProduceRequest(2));
-        checkErrorResponse(createProduceRequest(2), new UnknownServerException());
-        checkRequest(createProduceRequest(3));
-        checkErrorResponse(createProduceRequest(3), new UnknownServerException());
-        checkResponse(createProduceResponse(), 2);
-        checkRequest(createStopReplicaRequest(0, true));
-        checkRequest(createStopReplicaRequest(0, false));
-        checkErrorResponse(createStopReplicaRequest(0, true), new UnknownServerException());
-        checkRequest(createStopReplicaRequest(1, true));
-        checkRequest(createStopReplicaRequest(1, false));
-        checkErrorResponse(createStopReplicaRequest(1, true), new UnknownServerException());
-        checkResponse(createStopReplicaResponse(), 0);
-        checkRequest(createLeaderAndIsrRequest(0));
-        checkErrorResponse(createLeaderAndIsrRequest(0), new UnknownServerException());
-        checkRequest(createLeaderAndIsrRequest(1));
-        checkErrorResponse(createLeaderAndIsrRequest(1), new UnknownServerException());
-        checkResponse(createLeaderAndIsrResponse(), 0);
-        checkRequest(createSaslHandshakeRequest());
-        checkErrorResponse(createSaslHandshakeRequest(), new UnknownServerException());
-        checkResponse(createSaslHandshakeResponse(), 0);
-        checkRequest(createSaslAuthenticateRequest());
-        checkErrorResponse(createSaslAuthenticateRequest(), new UnknownServerException());
-        checkResponse(createSaslAuthenticateResponse(), 0);
-        checkResponse(createSaslAuthenticateResponse(), 1);
-        checkRequest(createApiVersionRequest());
-        checkErrorResponse(createApiVersionRequest(), new UnknownServerException());
-        checkResponse(createApiVersionResponse(), 0);
-        checkRequest(createCreateTopicRequest(0));
-        checkErrorResponse(createCreateTopicRequest(0), new UnknownServerException());
-        checkResponse(createCreateTopicResponse(), 0);
-        checkRequest(createCreateTopicRequest(1));
-        checkErrorResponse(createCreateTopicRequest(1), new UnknownServerException());
-        checkResponse(createCreateTopicResponse(), 1);
-        checkRequest(createDeleteTopicsRequest());
-        checkErrorResponse(createDeleteTopicsRequest(), new UnknownServerException());
-        checkResponse(createDeleteTopicsResponse(), 0);
-
-        checkRequest(createInitPidRequest());
-        checkErrorResponse(createInitPidRequest(), new UnknownServerException());
-        checkResponse(createInitPidResponse(), 0);
-
-        checkRequest(createAddPartitionsToTxnRequest());
-        checkResponse(createAddPartitionsToTxnResponse(), 0);
-        checkErrorResponse(createAddPartitionsToTxnRequest(), new UnknownServerException());
-        checkRequest(createAddOffsetsToTxnRequest());
-        checkResponse(createAddOffsetsToTxnResponse(), 0);
-        checkErrorResponse(createAddOffsetsToTxnRequest(), new UnknownServerException());
-        checkRequest(createEndTxnRequest());
-        checkResponse(createEndTxnResponse(), 0);
-        checkErrorResponse(createEndTxnRequest(), new UnknownServerException());
-        checkRequest(createWriteTxnMarkersRequest());
-        checkResponse(createWriteTxnMarkersResponse(), 0);
-        checkErrorResponse(createWriteTxnMarkersRequest(), new UnknownServerException());
-        checkRequest(createTxnOffsetCommitRequest());
-        checkResponse(createTxnOffsetCommitResponse(), 0);
-        checkErrorResponse(createTxnOffsetCommitRequest(), new UnknownServerException());
+        checkRequest(createFetchRequest(7, new FetchMetadata(123, 456), toForgetTopics), true);
+        checkResponse(createFetchResponse(123), 7, true);
+        checkResponse(createFetchResponse(Errors.FETCH_SESSION_ID_NOT_FOUND, 123), 7, true);
+        checkErrorResponse(createFetchRequest(4), new UnknownServerException(), true);
+        checkRequest(createHeartBeatRequest(), true);
+        checkErrorResponse(createHeartBeatRequest(), new UnknownServerException(), true);
+        checkResponse(createHeartBeatResponse(), 0, true);
+        checkRequest(createJoinGroupRequest(1), true);
+        checkErrorResponse(createJoinGroupRequest(0), new UnknownServerException(), true);
+        checkErrorResponse(createJoinGroupRequest(1), new UnknownServerException(), true);
+        checkResponse(createJoinGroupResponse(), 0, true);
+        checkRequest(createLeaveGroupRequest(), true);
+        checkErrorResponse(createLeaveGroupRequest(), new UnknownServerException(), true);
+        checkResponse(createLeaveGroupResponse(), 0, true);
+        checkRequest(createListGroupsRequest(), true);
+        checkErrorResponse(createListGroupsRequest(), new UnknownServerException(), true);
+        checkResponse(createListGroupsResponse(), 0, true);
+        checkRequest(createDescribeGroupRequest(), true);
+        checkErrorResponse(createDescribeGroupRequest(), new UnknownServerException(), true);
+        checkResponse(createDescribeGroupResponse(), 0, true);
+        checkRequest(createDeleteGroupsRequest(), true);
+        checkErrorResponse(createDeleteGroupsRequest(), new UnknownServerException(), true);
+        checkResponse(createDeleteGroupsResponse(), 0, true);
+        for (int i = 0; i < ApiKeys.LIST_OFFSETS.latestVersion(); i++) {
+            checkRequest(createListOffsetRequest(i), true);
+            checkErrorResponse(createListOffsetRequest(i), new UnknownServerException(), true);
+            checkResponse(createListOffsetResponse(i), i, true);
+        }
+        checkRequest(MetadataRequest.Builder.allTopics().build((short) 2), true);
+        checkRequest(createMetadataRequest(1, Collections.singletonList("topic1")), true);
+        checkErrorResponse(createMetadataRequest(1, Collections.singletonList("topic1")), new UnknownServerException(), true);
+        checkResponse(createMetadataResponse(), 2, true);
+        checkErrorResponse(createMetadataRequest(2, Collections.singletonList("topic1")), new UnknownServerException(), true);
+        checkResponse(createMetadataResponse(), 3, true);
+        checkErrorResponse(createMetadataRequest(3, Collections.singletonList("topic1")), new UnknownServerException(), true);
+        checkResponse(createMetadataResponse(), 4, true);
+        checkErrorResponse(createMetadataRequest(4, Collections.singletonList("topic1")), new UnknownServerException(), true);
+        checkRequest(OffsetFetchRequest.Builder.allTopicPartitions("group1").build(), true);
+        checkErrorResponse(OffsetFetchRequest.Builder.allTopicPartitions("group1").build(), new NotCoordinatorException("Not Coordinator"), true);
+        checkRequest(createOffsetFetchRequest(0), true);
+        checkRequest(createOffsetFetchRequest(1), true);
+        checkRequest(createOffsetFetchRequest(2), true);
+        checkRequest(OffsetFetchRequest.Builder.allTopicPartitions("group1").build(), true);
+        checkErrorResponse(createOffsetFetchRequest(0), new UnknownServerException(), true);
+        checkErrorResponse(createOffsetFetchRequest(1), new UnknownServerException(), true);
+        checkErrorResponse(createOffsetFetchRequest(2), new UnknownServerException(), true);
+        checkResponse(createOffsetFetchResponse(), 0, true);
+        checkRequest(createProduceRequest(2), true);
+        checkErrorResponse(createProduceRequest(2), new UnknownServerException(), true);
+        checkRequest(createProduceRequest(3), true);
+        checkErrorResponse(createProduceRequest(3), new UnknownServerException(), true);
+        checkResponse(createProduceResponse(), 2, true);
+        checkResponse(createProduceResponseWithErrorMessage(), 8, true);
+        checkRequest(createStopReplicaRequest(0, true), true);
+        checkRequest(createStopReplicaRequest(0, false), true);
+        checkErrorResponse(createStopReplicaRequest(0, true), new UnknownServerException(), true);
+        checkRequest(createStopReplicaRequest(1, true), true);
+        checkRequest(createStopReplicaRequest(1, false), true);
+        checkErrorResponse(createStopReplicaRequest(1, true), new UnknownServerException(), true);
+        checkResponse(createStopReplicaResponse(), 0, true);
+        checkRequest(createLeaderAndIsrRequest(0), true);
+        checkErrorResponse(createLeaderAndIsrRequest(0), new UnknownServerException(), false);
+        checkRequest(createLeaderAndIsrRequest(1), true);
+        checkErrorResponse(createLeaderAndIsrRequest(1), new UnknownServerException(), false);
+        checkRequest(createLeaderAndIsrRequest(2), true);
+        checkErrorResponse(createLeaderAndIsrRequest(2), new UnknownServerException(), false);
+        checkResponse(createLeaderAndIsrResponse(), 0, true);
+        checkRequest(createSaslHandshakeRequest(), true);
+        checkErrorResponse(createSaslHandshakeRequest(), new UnknownServerException(), true);
+        checkResponse(createSaslHandshakeResponse(), 0, true);
+        checkRequest(createSaslAuthenticateRequest(), true);
+        checkErrorResponse(createSaslAuthenticateRequest(), new UnknownServerException(), true);
+        checkResponse(createSaslAuthenticateResponse(), 0, true);
+        checkResponse(createSaslAuthenticateResponse(), 1, true);
+        checkRequest(createApiVersionRequest(), true);
+        checkErrorResponse(createApiVersionRequest(), new UnknownServerException(), true);
+        checkErrorResponse(createApiVersionRequest(), new UnsupportedVersionException("Not Supported"), true);
+        checkResponse(createApiVersionResponse(), 0, true);
+        checkResponse(createApiVersionResponse(), 1, true);
+        checkResponse(createApiVersionResponse(), 2, true);
+        checkResponse(createApiVersionResponse(), 3, true);
+        checkResponse(ApiVersionsResponse.DEFAULT_API_VERSIONS_RESPONSE, 0, true);
+        checkResponse(ApiVersionsResponse.DEFAULT_API_VERSIONS_RESPONSE, 1, true);
+        checkResponse(ApiVersionsResponse.DEFAULT_API_VERSIONS_RESPONSE, 2, true);
+        checkResponse(ApiVersionsResponse.DEFAULT_API_VERSIONS_RESPONSE, 3, true);
+
+        checkRequest(createCreateTopicRequest(0), true);
+        checkErrorResponse(createCreateTopicRequest(0), new UnknownServerException(), true);
+        checkResponse(createCreateTopicResponse(), 0, true);
+        checkRequest(createCreateTopicRequest(1), true);
+        checkErrorResponse(createCreateTopicRequest(1), new UnknownServerException(), true);
+        checkResponse(createCreateTopicResponse(), 1, true);
+        checkRequest(createDeleteTopicsRequest(), true);
+        checkErrorResponse(createDeleteTopicsRequest(), new UnknownServerException(), true);
+        checkResponse(createDeleteTopicsResponse(), 0, true);
+
+        checkRequest(createInitPidRequest(), true);
+        checkErrorResponse(createInitPidRequest(), new UnknownServerException(), true);
+        checkResponse(createInitPidResponse(), 0, true);
+
+        checkRequest(createAddPartitionsToTxnRequest(), true);
+        checkResponse(createAddPartitionsToTxnResponse(), 0, true);
+        checkErrorResponse(createAddPartitionsToTxnRequest(), new UnknownServerException(), true);
+        checkRequest(createAddOffsetsToTxnRequest(), true);
+        checkResponse(createAddOffsetsToTxnResponse(), 0, true);
+        checkErrorResponse(createAddOffsetsToTxnRequest(), new UnknownServerException(), true);
+        checkRequest(createEndTxnRequest(), true);
+        checkResponse(createEndTxnResponse(), 0, true);
+        checkErrorResponse(createEndTxnRequest(), new UnknownServerException(), true);
+        checkRequest(createWriteTxnMarkersRequest(), true);
+        checkResponse(createWriteTxnMarkersResponse(), 0, true);
+        checkErrorResponse(createWriteTxnMarkersRequest(), new UnknownServerException(), true);
+        checkRequest(createTxnOffsetCommitRequest(), true);
+        checkResponse(createTxnOffsetCommitResponse(), 0, true);
+        checkErrorResponse(createTxnOffsetCommitRequest(), new UnknownServerException(), true);
 
         checkOlderFetchVersions();
-        checkResponse(createMetadataResponse(), 0);
-        checkResponse(createMetadataResponse(), 1);
-        checkErrorResponse(createMetadataRequest(1, singletonList("topic1")), new UnknownServerException());
-        checkRequest(createOffsetCommitRequest(0));
-        checkErrorResponse(createOffsetCommitRequest(0), new UnknownServerException());
-        checkRequest(createOffsetCommitRequest(1));
-        checkErrorResponse(createOffsetCommitRequest(1), new UnknownServerException());
-        checkRequest(createOffsetCommitRequest(2));
-        checkErrorResponse(createOffsetCommitRequest(2), new UnknownServerException());
-        checkRequest(createOffsetCommitRequest(3));
-        checkErrorResponse(createOffsetCommitRequest(3), new UnknownServerException());
-        checkRequest(createOffsetCommitRequest(4));
-        checkErrorResponse(createOffsetCommitRequest(4), new UnknownServerException());
-        checkResponse(createOffsetCommitResponse(), 4);
-        checkRequest(createOffsetCommitRequest(5));
-        checkErrorResponse(createOffsetCommitRequest(5), new UnknownServerException());
-        checkResponse(createOffsetCommitResponse(), 5);
-        checkRequest(createJoinGroupRequest(0));
-        checkRequest(createUpdateMetadataRequest(0, null));
-        checkErrorResponse(createUpdateMetadataRequest(0, null), new UnknownServerException());
-        checkRequest(createUpdateMetadataRequest(1, null));
-        checkRequest(createUpdateMetadataRequest(1, "rack1"));
-        checkErrorResponse(createUpdateMetadataRequest(1, null), new UnknownServerException());
-        checkRequest(createUpdateMetadataRequest(2, "rack1"));
-        checkRequest(createUpdateMetadataRequest(2, null));
-        checkErrorResponse(createUpdateMetadataRequest(2, "rack1"), new UnknownServerException());
-        checkRequest(createUpdateMetadataRequest(3, "rack1"));
-        checkRequest(createUpdateMetadataRequest(3, null));
-        checkErrorResponse(createUpdateMetadataRequest(3, "rack1"), new UnknownServerException());
-        checkRequest(createUpdateMetadataRequest(4, "rack1"));
-        checkRequest(createUpdateMetadataRequest(4, null));
-        checkErrorResponse(createUpdateMetadataRequest(4, "rack1"), new UnknownServerException());
-        checkRequest(createUpdateMetadataRequest(5, "rack1"));
-        checkRequest(createUpdateMetadataRequest(5, null));
-        checkErrorResponse(createUpdateMetadataRequest(5, "rack1"), new UnknownServerException());
-        checkResponse(createUpdateMetadataResponse(), 0);
-        checkRequest(createListOffsetRequest(0));
-        checkErrorResponse(createListOffsetRequest(0), new UnknownServerException());
-        checkResponse(createListOffsetResponse(0), 0);
-        checkRequest(createLeaderEpochRequest());
-        checkResponse(createLeaderEpochResponse(), 0);
-        checkErrorResponse(createLeaderEpochRequest(), new UnknownServerException());
-        checkRequest(createAddPartitionsToTxnRequest());
-        checkErrorResponse(createAddPartitionsToTxnRequest(), new UnknownServerException());
-        checkResponse(createAddPartitionsToTxnResponse(), 0);
-        checkRequest(createAddOffsetsToTxnRequest());
-        checkErrorResponse(createAddOffsetsToTxnRequest(), new UnknownServerException());
-        checkResponse(createAddOffsetsToTxnResponse(), 0);
-        checkRequest(createEndTxnRequest());
-        checkErrorResponse(createEndTxnRequest(), new UnknownServerException());
-        checkResponse(createEndTxnResponse(), 0);
-        checkRequest(createWriteTxnMarkersRequest());
-        checkErrorResponse(createWriteTxnMarkersRequest(), new UnknownServerException());
-        checkResponse(createWriteTxnMarkersResponse(), 0);
-        checkRequest(createTxnOffsetCommitRequest());
-        checkErrorResponse(createTxnOffsetCommitRequest(), new UnknownServerException());
-        checkResponse(createTxnOffsetCommitResponse(), 0);
-        checkRequest(createListAclsRequest());
-        checkErrorResponse(createListAclsRequest(), new SecurityDisabledException("Security is not enabled."));
-        checkResponse(createDescribeAclsResponse(), ApiKeys.DESCRIBE_ACLS.latestVersion());
-        checkRequest(createCreateAclsRequest());
-        checkErrorResponse(createCreateAclsRequest(), new SecurityDisabledException("Security is not enabled."));
-        checkResponse(createCreateAclsResponse(), ApiKeys.CREATE_ACLS.latestVersion());
-        checkRequest(createDeleteAclsRequest());
-        checkErrorResponse(createDeleteAclsRequest(), new SecurityDisabledException("Security is not enabled."));
-        checkResponse(createDeleteAclsResponse(), ApiKeys.DELETE_ACLS.latestVersion());
-        checkRequest(createAlterConfigsRequest());
-        checkErrorResponse(createAlterConfigsRequest(), new UnknownServerException());
-        checkResponse(createAlterConfigsResponse(), 0);
-        checkRequest(createDescribeConfigsRequest(0));
-        checkRequest(createDescribeConfigsRequestWithConfigEntries(0));
-        checkErrorResponse(createDescribeConfigsRequest(0), new UnknownServerException());
-        checkResponse(createDescribeConfigsResponse(), 0);
-        checkRequest(createDescribeConfigsRequest(1));
-        checkRequest(createDescribeConfigsRequestWithConfigEntries(1));
-        checkErrorResponse(createDescribeConfigsRequest(1), new UnknownServerException());
-        checkResponse(createDescribeConfigsResponse(), 1);
+        checkResponse(createMetadataResponse(), 0, true);
+        checkResponse(createMetadataResponse(), 1, true);
+        checkErrorResponse(createMetadataRequest(1, Collections.singletonList("topic1")), new UnknownServerException(), true);
+        checkRequest(createOffsetCommitRequest(0), true);
+        checkErrorResponse(createOffsetCommitRequest(0), new UnknownServerException(), true);
+        checkRequest(createOffsetCommitRequest(1), true);
+        checkErrorResponse(createOffsetCommitRequest(1), new UnknownServerException(), true);
+        checkRequest(createOffsetCommitRequest(2), true);
+        checkErrorResponse(createOffsetCommitRequest(2), new UnknownServerException(), true);
+        checkRequest(createOffsetCommitRequest(3), true);
+        checkErrorResponse(createOffsetCommitRequest(3), new UnknownServerException(), true);
+        checkRequest(createOffsetCommitRequest(4), true);
+        checkErrorResponse(createOffsetCommitRequest(4), new UnknownServerException(), true);
+        checkResponse(createOffsetCommitResponse(), 4, true);
+        checkRequest(createOffsetCommitRequest(5), true);
+        checkErrorResponse(createOffsetCommitRequest(5), new UnknownServerException(), true);
+        checkResponse(createOffsetCommitResponse(), 5, true);
+        checkRequest(createJoinGroupRequest(0), true);
+        checkRequest(createUpdateMetadataRequest(0, null), false);
+        checkErrorResponse(createUpdateMetadataRequest(0, null), new UnknownServerException(), true);
+        checkRequest(createUpdateMetadataRequest(1, null), false);
+        checkRequest(createUpdateMetadataRequest(1, "rack1"), false);
+        checkErrorResponse(createUpdateMetadataRequest(1, null), new UnknownServerException(), true);
+        checkRequest(createUpdateMetadataRequest(2, "rack1"), false);
+        checkRequest(createUpdateMetadataRequest(2, null), false);
+        checkErrorResponse(createUpdateMetadataRequest(2, "rack1"), new UnknownServerException(), true);
+        checkRequest(createUpdateMetadataRequest(3, "rack1"), false);
+        checkRequest(createUpdateMetadataRequest(3, null), false);
+        checkErrorResponse(createUpdateMetadataRequest(3, "rack1"), new UnknownServerException(), true);
+        checkRequest(createUpdateMetadataRequest(4, "rack1"), false);
+        checkRequest(createUpdateMetadataRequest(4, null), false);
+        checkErrorResponse(createUpdateMetadataRequest(4, "rack1"), new UnknownServerException(), true);
+        checkRequest(createUpdateMetadataRequest(5, "rack1"), false);
+        checkRequest(createUpdateMetadataRequest(5, null), false);
+        checkErrorResponse(createUpdateMetadataRequest(5, "rack1"), new UnknownServerException(), true);
+        checkResponse(createUpdateMetadataResponse(), 0, true);
+        checkRequest(createListOffsetRequest(0), true);
+        checkErrorResponse(createListOffsetRequest(0), new UnknownServerException(), true);
+        checkResponse(createListOffsetResponse(0), 0, true);
+        checkRequest(createLeaderEpochRequestForReplica(0, 1), true);
+        checkRequest(createLeaderEpochRequestForConsumer(), true);
+        checkResponse(createLeaderEpochResponse(), 0, true);
+        checkErrorResponse(createLeaderEpochRequestForConsumer(), new UnknownServerException(), true);
+        checkRequest(createAddPartitionsToTxnRequest(), true);
+        checkErrorResponse(createAddPartitionsToTxnRequest(), new UnknownServerException(), true);
+        checkResponse(createAddPartitionsToTxnResponse(), 0, true);
+        checkRequest(createAddOffsetsToTxnRequest(), true);
+        checkErrorResponse(createAddOffsetsToTxnRequest(), new UnknownServerException(), true);
+        checkResponse(createAddOffsetsToTxnResponse(), 0, true);
+        checkRequest(createEndTxnRequest(), true);
+        checkErrorResponse(createEndTxnRequest(), new UnknownServerException(), true);
+        checkResponse(createEndTxnResponse(), 0, true);
+        checkRequest(createWriteTxnMarkersRequest(), true);
+        checkErrorResponse(createWriteTxnMarkersRequest(), new UnknownServerException(), true);
+        checkResponse(createWriteTxnMarkersResponse(), 0, true);
+        checkRequest(createTxnOffsetCommitRequest(), true);
+        checkErrorResponse(createTxnOffsetCommitRequest(), new UnknownServerException(), true);
+        checkResponse(createTxnOffsetCommitResponse(), 0, true);
+        checkRequest(createListAclsRequest(), true);
+        checkErrorResponse(createListAclsRequest(), new SecurityDisabledException("Security is not enabled."), true);
+        checkResponse(createDescribeAclsResponse(), ApiKeys.DESCRIBE_ACLS.latestVersion(), true);
+        checkRequest(createCreateAclsRequest(), true);
+        checkErrorResponse(createCreateAclsRequest(), new SecurityDisabledException("Security is not enabled."), true);
+        checkResponse(createCreateAclsResponse(), ApiKeys.CREATE_ACLS.latestVersion(), true);
+        checkRequest(createDeleteAclsRequest(), true);
+        checkErrorResponse(createDeleteAclsRequest(), new SecurityDisabledException("Security is not enabled."), true);
+        checkResponse(createDeleteAclsResponse(), ApiKeys.DELETE_ACLS.latestVersion(), true);
+        checkRequest(createAlterConfigsRequest(), false);
+        checkErrorResponse(createAlterConfigsRequest(), new UnknownServerException(), true);
+        checkResponse(createAlterConfigsResponse(), 0, false);
+        checkRequest(createDescribeConfigsRequest(0), true);
+        checkRequest(createDescribeConfigsRequestWithConfigEntries(0), false);
+        checkErrorResponse(createDescribeConfigsRequest(0), new UnknownServerException(), true);
+        checkResponse(createDescribeConfigsResponse(), 0, false);
+        checkRequest(createDescribeConfigsRequest(1), true);
+        checkRequest(createDescribeConfigsRequestWithConfigEntries(1), false);
+        checkErrorResponse(createDescribeConfigsRequest(1), new UnknownServerException(), true);
+        checkResponse(createDescribeConfigsResponse(), 1, false);
         checkDescribeConfigsResponseVersions();
-        checkRequest(createCreatePartitionsRequest());
-        checkRequest(createCreatePartitionsRequestWithAssignments());
-        checkErrorResponse(createCreatePartitionsRequest(), new InvalidTopicException());
-        checkResponse(createCreatePartitionsResponse(), 0);
-        checkRequest(createCreateTokenRequest());
-        checkErrorResponse(createCreateTokenRequest(), new UnknownServerException());
-        checkResponse(createCreateTokenResponse(), 0);
-        checkRequest(createDescribeTokenRequest());
-        checkErrorResponse(createDescribeTokenRequest(), new UnknownServerException());
-        checkResponse(createDescribeTokenResponse(), 0);
-        checkRequest(createExpireTokenRequest());
-        checkErrorResponse(createExpireTokenRequest(), new UnknownServerException());
-        checkResponse(createExpireTokenResponse(), 0);
-        checkRequest(createRenewTokenRequest());
-        checkErrorResponse(createRenewTokenRequest(), new UnknownServerException());
-        checkResponse(createRenewTokenResponse(), 0);
-        checkRequest(createElectPreferredLeadersRequest());
-        checkRequest(createElectPreferredLeadersRequestNullPartitions());
-        checkErrorResponse(createElectPreferredLeadersRequest(), new UnknownServerException());
-        checkResponse(createElectPreferredLeadersResponse(), 0);
+        checkRequest(createCreatePartitionsRequest(), true);
+        checkRequest(createCreatePartitionsRequestWithAssignments(), false);
+        checkErrorResponse(createCreatePartitionsRequest(), new InvalidTopicException(), true);
+        checkResponse(createCreatePartitionsResponse(), 0, true);
+        checkRequest(createCreateTokenRequest(), true);
+        checkErrorResponse(createCreateTokenRequest(), new UnknownServerException(), true);
+        checkResponse(createCreateTokenResponse(), 0, true);
+        checkRequest(createDescribeTokenRequest(), true);
+        checkErrorResponse(createDescribeTokenRequest(), new UnknownServerException(), true);
+        checkResponse(createDescribeTokenResponse(), 0, true);
+        checkRequest(createExpireTokenRequest(), true);
+        checkErrorResponse(createExpireTokenRequest(), new UnknownServerException(), true);
+        checkResponse(createExpireTokenResponse(), 0, true);
+        checkRequest(createRenewTokenRequest(), true);
+        checkErrorResponse(createRenewTokenRequest(), new UnknownServerException(), true);
+        checkResponse(createRenewTokenResponse(), 0, true);
+        checkRequest(createElectLeadersRequest(), true);
+        checkRequest(createElectLeadersRequestNullPartitions(), true);
+        checkErrorResponse(createElectLeadersRequest(), new UnknownServerException(), true);
+        checkResponse(createElectLeadersResponse(), 1, true);
+        checkRequest(createIncrementalAlterConfigsRequest(), true);
+        checkErrorResponse(createIncrementalAlterConfigsRequest(), new UnknownServerException(), true);
+        checkResponse(createIncrementalAlterConfigsResponse(), 0, true);
+        checkRequest(createAlterPartitionReassignmentsRequest(), true);
+        checkErrorResponse(createAlterPartitionReassignmentsRequest(), new UnknownServerException(), true);
+        checkResponse(createAlterPartitionReassignmentsResponse(), 0, true);
+        checkRequest(createListPartitionReassignmentsRequest(), true);
+        checkErrorResponse(createListPartitionReassignmentsRequest(), new UnknownServerException(), true);
+        checkResponse(createListPartitionReassignmentsResponse(), 0, true);
+        checkRequest(createOffsetDeleteRequest(), true);
+        checkErrorResponse(createOffsetDeleteRequest(), new UnknownServerException(), true);
+        checkResponse(createOffsetDeleteResponse(), 0, true);
     }
 
     @Test
     public void testResponseHeader() {
-        ResponseHeader header = createResponseHeader();
+        ResponseHeader header = createResponseHeader((short) 1);
         ByteBuffer buffer = toBuffer(header.toStruct());
-        ResponseHeader deserialized = ResponseHeader.parse(buffer);
+        ResponseHeader deserialized = ResponseHeader.parse(buffer, header.headerVersion());
         assertEquals(header.correlationId(), deserialized.correlationId());
     }
 
     private void checkOlderFetchVersions() throws Exception {
-        int latestVersion = ApiKeys.FETCH.latestVersion();
+        int latestVersion = FETCH.latestVersion();
         for (int i = 0; i < latestVersion; ++i) {
-            checkErrorResponse(createFetchRequest(i), new UnknownServerException());
-            checkRequest(createFetchRequest(i));
-            checkResponse(createFetchResponse(), i);
+            checkErrorResponse(createFetchRequest(i), new UnknownServerException(), true);
+            checkRequest(createFetchRequest(i), true);
+            checkResponse(createFetchResponse(), i, true);
         }
     }
 
@@ -376,34 +463,42 @@ private void checkDescribeConfigsResponseVersions() throws Exception {
         verifyDescribeConfigsResponse(response, deserialized1, 1);
     }
 
-    private void checkErrorResponse(AbstractRequest req, Throwable e) throws Exception {
-        checkResponse(req.getErrorResponse(e), req.version());
+    private void checkErrorResponse(AbstractRequest req, Throwable e, boolean checkEqualityAndHashCode) {
+        checkResponse(req.getErrorResponse(e), req.version(), checkEqualityAndHashCode);
     }
 
-    private void checkRequest(AbstractRequest req) throws Exception {
+    private void checkRequest(AbstractRequest req, boolean checkEqualityAndHashCode) {
         // Check that we can serialize, deserialize and serialize again
-        // We don't check for equality or hashCode because it is likely to fail for any request containing a HashMap
-        checkRequest(req, false);
-    }
-
-    private void checkRequest(AbstractRequest req, boolean checkEqualityAndHashCode) throws Exception {
-        // Check that we can serialize, deserialize and serialize again
-        // Check for equality and hashCode only if indicated
-        Struct struct = req.toStruct();
-        AbstractRequest deserialized = (AbstractRequest) deserialize(req, struct, req.version());
-        Struct struct2 = deserialized.toStruct();
-        if (checkEqualityAndHashCode) {
-            assertEquals(struct, struct2);
-            assertEquals(struct.hashCode(), struct2.hashCode());
+        // Check for equality and hashCode of the Struct only if indicated (it is likely to fail if any of the fields
+        // in the request is a HashMap with multiple elements since ordering of the elements may vary)
+        try {
+            Struct struct = req.toStruct();
+            AbstractRequest deserialized = AbstractRequest.parseRequest(req.api, req.version(), struct);
+            Struct struct2 = deserialized.toStruct();
+            if (checkEqualityAndHashCode) {
+                assertEquals(struct, struct2);
+                assertEquals(struct.hashCode(), struct2.hashCode());
+            }
+        } catch (Exception e) {
+            throw new RuntimeException("Failed to deserialize request " + req + " with type " + req.getClass(), e);
         }
     }
 
-    private void checkResponse(AbstractResponse response, int version) throws Exception {
+    private void checkResponse(AbstractResponse response, int version, boolean checkEqualityAndHashCode) {
         // Check that we can serialize, deserialize and serialize again
-        // We don't check for equality or hashCode because it is likely to fail for any response containing a HashMap
-        Struct struct = response.toStruct((short) version);
-        AbstractResponse deserialized = (AbstractResponse) deserialize(response, struct, (short) version);
-        Struct struct2 = deserialized.toStruct((short) version);
+        // Check for equality and hashCode of the Struct only if indicated (it is likely to fail if any of the fields
+        // in the response is a HashMap with multiple elements since ordering of the elements may vary)
+        try {
+            Struct struct = response.toStruct((short) version);
+            AbstractResponse deserialized = (AbstractResponse) deserialize(response, struct, (short) version);
+            Struct struct2 = deserialized.toStruct((short) version);
+            if (checkEqualityAndHashCode) {
+                assertEquals(struct, struct2);
+                assertEquals(struct.hashCode(), struct2.hashCode());
+            }
+        } catch (Exception e) {
+            throw new RuntimeException("Failed to deserialize response " + response + " with type " + response.getClass(), e);
+        }
     }
 
     private AbstractRequestResponse deserialize(AbstractRequestResponse req, Struct struct, short version) throws NoSuchMethodException, IllegalAccessException, InvocationTargetException {
@@ -414,7 +509,10 @@ private AbstractRequestResponse deserialize(AbstractRequestResponse req, Struct
 
     @Test(expected = UnsupportedVersionException.class)
     public void cannotUseFindCoordinatorV0ToFindTransactionCoordinator() {
-        FindCoordinatorRequest.Builder builder = new FindCoordinatorRequest.Builder(FindCoordinatorRequest.CoordinatorType.TRANSACTION, "foobar");
+        FindCoordinatorRequest.Builder builder = new FindCoordinatorRequest.Builder(
+                new FindCoordinatorRequestData()
+                    .setKeyType(CoordinatorType.TRANSACTION.id)
+                    .setKey("foobar"));
         builder.build((short) 0);
     }
 
@@ -465,57 +563,6 @@ public void produceRequestGetErrorResponseTest() {
         assertEquals(RecordBatch.NO_TIMESTAMP, partitionResponse.logAppendTime);
     }
 
-    @Test
-    public void produceResponseV5Test() {
-        Map responseData = new HashMap<>();
-        TopicPartition tp0 = new TopicPartition("test", 0);
-        responseData.put(tp0, new ProduceResponse.PartitionResponse(Errors.NONE,
-                10000, RecordBatch.NO_TIMESTAMP, 100));
-
-        ProduceResponse v5Response = new ProduceResponse(responseData, 10);
-        short version = 5;
-
-        ByteBuffer buffer = v5Response.serialize(version, new ResponseHeader(0));
-        buffer.rewind();
-
-        ResponseHeader.parse(buffer); // throw away.
-
-        Struct deserializedStruct = ApiKeys.PRODUCE.parseResponse(version, buffer);
-
-        ProduceResponse v5FromBytes = (ProduceResponse) AbstractResponse.parseResponse(ApiKeys.PRODUCE,
-                deserializedStruct, version);
-
-        assertEquals(1, v5FromBytes.responses().size());
-        assertTrue(v5FromBytes.responses().containsKey(tp0));
-        ProduceResponse.PartitionResponse partitionResponse = v5FromBytes.responses().get(tp0);
-        assertEquals(100, partitionResponse.logStartOffset);
-        assertEquals(10000, partitionResponse.baseOffset);
-        assertEquals(10, v5FromBytes.throttleTimeMs());
-        assertEquals(responseData, v5Response.responses());
-    }
-
-    @Test
-    public void produceResponseVersionTest() {
-        Map responseData = new HashMap<>();
-        responseData.put(new TopicPartition("test", 0), new ProduceResponse.PartitionResponse(Errors.NONE,
-                10000, RecordBatch.NO_TIMESTAMP, 100));
-        ProduceResponse v0Response = new ProduceResponse(responseData);
-        ProduceResponse v1Response = new ProduceResponse(responseData, 10);
-        ProduceResponse v2Response = new ProduceResponse(responseData, 10);
-        assertEquals("Throttle time must be zero", 0, v0Response.throttleTimeMs());
-        assertEquals("Throttle time must be 10", 10, v1Response.throttleTimeMs());
-        assertEquals("Throttle time must be 10", 10, v2Response.throttleTimeMs());
-        assertEquals("Should use schema version 0", ApiKeys.PRODUCE.responseSchema((short) 0),
-                v0Response.toStruct((short) 0).schema());
-        assertEquals("Should use schema version 1", ApiKeys.PRODUCE.responseSchema((short) 1),
-                v1Response.toStruct((short) 1).schema());
-        assertEquals("Should use schema version 2", ApiKeys.PRODUCE.responseSchema((short) 2),
-                v2Response.toStruct((short) 2).schema());
-        assertEquals("Response data does not match", responseData, v0Response.responses());
-        assertEquals("Response data does not match", responseData, v1Response.responses());
-        assertEquals("Response data does not match", responseData, v2Response.responses());
-    }
-
     @Test
     public void fetchResponseVersionTest() {
         LinkedHashMap> responseData = new LinkedHashMap<>();
@@ -523,15 +570,15 @@ public void fetchResponseVersionTest() {
         MemoryRecords records = MemoryRecords.readableRecords(ByteBuffer.allocate(10));
         responseData.put(new TopicPartition("test", 0), new FetchResponse.PartitionData<>(
                 Errors.NONE, 1000000, FetchResponse.INVALID_LAST_STABLE_OFFSET,
-                0L, null, records));
+                0L, Optional.empty(), null, records));
 
         FetchResponse v0Response = new FetchResponse<>(Errors.NONE, responseData, 0, INVALID_SESSION_ID);
         FetchResponse v1Response = new FetchResponse<>(Errors.NONE, responseData, 10, INVALID_SESSION_ID);
         assertEquals("Throttle time must be zero", 0, v0Response.throttleTimeMs());
         assertEquals("Throttle time must be 10", 10, v1Response.throttleTimeMs());
-        assertEquals("Should use schema version 0", ApiKeys.FETCH.responseSchema((short) 0),
+        assertEquals("Should use schema version 0", FETCH.responseSchema((short) 0),
                 v0Response.toStruct((short) 0).schema());
-        assertEquals("Should use schema version 1", ApiKeys.FETCH.responseSchema((short) 1),
+        assertEquals("Should use schema version 1", FETCH.responseSchema((short) 1),
                 v1Response.toStruct((short) 1).schema());
         assertEquals("Response data does not match", responseData, v0Response.responseData());
         assertEquals("Response data does not match", responseData, v1Response.responseData());
@@ -547,11 +594,11 @@ public void testFetchResponseV4() {
                 new FetchResponse.AbortedTransaction(15, 50)
         );
         responseData.put(new TopicPartition("bar", 0), new FetchResponse.PartitionData<>(Errors.NONE, 100000,
-                FetchResponse.INVALID_LAST_STABLE_OFFSET, FetchResponse.INVALID_LOG_START_OFFSET, abortedTransactions, records));
+                FetchResponse.INVALID_LAST_STABLE_OFFSET, FetchResponse.INVALID_LOG_START_OFFSET, Optional.empty(), abortedTransactions, records));
         responseData.put(new TopicPartition("bar", 1), new FetchResponse.PartitionData<>(Errors.NONE, 900000,
-                5, FetchResponse.INVALID_LOG_START_OFFSET, null, records));
+                5, FetchResponse.INVALID_LOG_START_OFFSET, Optional.empty(), null, records));
         responseData.put(new TopicPartition("foo", 0), new FetchResponse.PartitionData<>(Errors.NONE, 70000,
-                6, FetchResponse.INVALID_LOG_START_OFFSET, Collections.emptyList(), records));
+                6, FetchResponse.INVALID_LOG_START_OFFSET, Optional.empty(), Collections.emptyList(), records));
 
         FetchResponse response = new FetchResponse<>(Errors.NONE, responseData, 10, INVALID_SESSION_ID);
         FetchResponse deserialized = FetchResponse.parse(toBuffer(response.toStruct((short) 4)), (short) 4);
@@ -560,10 +607,10 @@ public void testFetchResponseV4() {
 
     @Test
     public void verifyFetchResponseFullWrites() throws Exception {
-        verifyFetchResponseFullWrite(ApiKeys.FETCH.latestVersion(), createFetchResponse(123));
-        verifyFetchResponseFullWrite(ApiKeys.FETCH.latestVersion(),
+        verifyFetchResponseFullWrite(FETCH.latestVersion(), createFetchResponse(123));
+        verifyFetchResponseFullWrite(FETCH.latestVersion(),
             createFetchResponse(Errors.FETCH_SESSION_ID_NOT_FOUND, 123));
-        for (short version = 0; version <= ApiKeys.FETCH.latestVersion(); version++) {
+        for (short version = 0; version <= FETCH.latestVersion(); version++) {
             verifyFetchResponseFullWrite(version, createFetchResponse());
         }
     }
@@ -571,7 +618,8 @@ public void verifyFetchResponseFullWrites() throws Exception {
     private void verifyFetchResponseFullWrite(short apiVersion, FetchResponse fetchResponse) throws Exception {
         int correlationId = 15;
 
-        Send send = fetchResponse.toSend("1", new ResponseHeader(correlationId), apiVersion);
+        short responseHeaderVersion = FETCH.responseHeaderVersion(apiVersion);
+        Send send = fetchResponse.toSend("1", new ResponseHeader(correlationId, responseHeaderVersion), apiVersion);
         ByteBufferChannel channel = new ByteBufferChannel(send.size());
         send.writeTo(channel);
         channel.close();
@@ -583,11 +631,11 @@ private void verifyFetchResponseFullWrite(short apiVersion, FetchResponse fetchR
         assertTrue(size > 0);
 
         // read the header
-        ResponseHeader responseHeader = ResponseHeader.parse(channel.buffer());
+        ResponseHeader responseHeader = ResponseHeader.parse(channel.buffer(), responseHeaderVersion);
         assertEquals(correlationId, responseHeader.correlationId());
 
         // read the body
-        Struct responseBody = ApiKeys.FETCH.responseSchema(apiVersion).read(buf);
+        Struct responseBody = FETCH.responseSchema(apiVersion).read(buf);
         assertEquals(fetchResponse.toStruct(apiVersion), responseBody);
 
         assertEquals(size, responseHeader.sizeOf() + responseBody.sizeOf());
@@ -601,7 +649,7 @@ public void testControlledShutdownResponse() {
         ByteBuffer buffer = toBuffer(struct);
         ControlledShutdownResponse deserialized = ControlledShutdownResponse.parse(buffer, version);
         assertEquals(response.error(), deserialized.error());
-        assertEquals(response.partitionsRemaining(), deserialized.partitionsRemaining());
+        assertEquals(response.data().remainingPartitions(), deserialized.data().remainingPartitions());
     }
 
     @Test(expected = UnsupportedVersionException.class)
@@ -609,6 +657,28 @@ public void testCreateTopicRequestV0FailsIfValidateOnly() {
         createCreateTopicRequest(0, true);
     }
 
+    @Test
+    public void testCreateTopicRequestV3FailsIfNoPartitionsOrReplicas() {
+        final UnsupportedVersionException exception = assertThrows(
+            UnsupportedVersionException.class, () -> {
+                CreateTopicsRequestData data = new CreateTopicsRequestData()
+                    .setTimeoutMs(123)
+                    .setValidateOnly(false);
+                data.topics().add(new CreatableTopic().
+                    setName("foo").
+                    setNumPartitions(CreateTopicsRequest.NO_NUM_PARTITIONS).
+                    setReplicationFactor((short) 1));
+                data.topics().add(new CreatableTopic().
+                    setName("bar").
+                    setNumPartitions(1).
+                    setReplicationFactor(CreateTopicsRequest.NO_REPLICATION_FACTOR));
+
+                new Builder(data).build((short) 3);
+            });
+        assertTrue(exception.getMessage().contains("supported in CreateTopicRequest version 4+"));
+        assertTrue(exception.getMessage().contains("[foo, bar]"));
+    }
+
     @Test
     public void testFetchRequestMaxBytesOldVersions() throws Exception {
         final short version = 1;
@@ -648,30 +718,122 @@ public void testJoinGroupRequestVersion0RebalanceTimeout() {
         final short version = 0;
         JoinGroupRequest jgr = createJoinGroupRequest(version);
         JoinGroupRequest jgr2 = new JoinGroupRequest(jgr.toStruct(), version);
-        assertEquals(jgr2.rebalanceTimeout(), jgr.rebalanceTimeout());
+        assertEquals(jgr2.data().rebalanceTimeoutMs(), jgr.data().rebalanceTimeoutMs());
     }
 
     @Test
     public void testOffsetFetchRequestBuilderToString() {
         String allTopicPartitionsString = OffsetFetchRequest.Builder.allTopicPartitions("someGroup").toString();
-        assertTrue(allTopicPartitionsString.contains(""));
+
+        assertTrue(allTopicPartitionsString.contains("groupId='someGroup', topics=null"));
         String string = new OffsetFetchRequest.Builder("group1",
-                singletonList(new TopicPartition("test11", 1))).toString();
+            Collections.singletonList(new TopicPartition("test11", 1))).toString();
         assertTrue(string.contains("test11"));
         assertTrue(string.contains("group1"));
     }
 
-    private ResponseHeader createResponseHeader() {
-        return new ResponseHeader(10);
+    @Test
+    public void testApiVersionsRequestBeforeV3Validation() {
+        for (short version = 0; version < 3; version++) {
+            ApiVersionsRequest request = new ApiVersionsRequest(new ApiVersionsRequestData(), version);
+            assertTrue(request.isValid());
+        }
+    }
+
+    @Test
+    public void testValidApiVersionsRequest() {
+        ApiVersionsRequest request;
+
+        request = new ApiVersionsRequest.Builder().build();
+        assertTrue(request.isValid());
+
+        request = new ApiVersionsRequest(new ApiVersionsRequestData()
+            .setClientSoftwareName("apache-kafka.java")
+            .setClientSoftwareVersion("0.0.0-SNAPSHOT"),
+            ApiKeys.API_VERSIONS.latestVersion()
+        );
+        assertTrue(request.isValid());
+    }
+
+    @Test
+    public void testInvalidApiVersionsRequest() {
+        testInvalidCase("java@apache_kafka", "0.0.0-SNAPSHOT");
+        testInvalidCase("apache-kafka-java", "0.0.0@java");
+        testInvalidCase("-apache-kafka-java", "0.0.0");
+        testInvalidCase("apache-kafka-java.", "0.0.0");
+    }
+
+    private void testInvalidCase(String name, String version) {
+        ApiVersionsRequest request = new ApiVersionsRequest(new ApiVersionsRequestData()
+            .setClientSoftwareName(name)
+            .setClientSoftwareVersion(version),
+            ApiKeys.API_VERSIONS.latestVersion()
+        );
+        assertFalse(request.isValid());
+    }
+
+    @Test
+    public void testApiVersionResponseWithUnsupportedError() {
+        ApiVersionsRequest request = new ApiVersionsRequest.Builder().build();
+        ApiVersionsResponse response = request.getErrorResponse(0, Errors.UNSUPPORTED_VERSION.exception());
+
+        assertEquals(Errors.UNSUPPORTED_VERSION.code(), response.data.errorCode());
+
+        ApiVersionsResponseKey apiVersion = response.data.apiKeys().find(ApiKeys.API_VERSIONS.id);
+        assertNotNull(apiVersion);
+        assertEquals(ApiKeys.API_VERSIONS.id, apiVersion.apiKey());
+        assertEquals(ApiKeys.API_VERSIONS.oldestVersion(), apiVersion.minVersion());
+        assertEquals(ApiKeys.API_VERSIONS.latestVersion(), apiVersion.maxVersion());
+    }
+
+    @Test
+    public void testApiVersionResponseWithNotUnsupportedError() {
+        ApiVersionsRequest request = new ApiVersionsRequest.Builder().build();
+        ApiVersionsResponse response = request.getErrorResponse(0, Errors.INVALID_REQUEST.exception());
+
+        assertEquals(response.data.errorCode(), Errors.INVALID_REQUEST.code());
+        assertTrue(response.data.apiKeys().isEmpty());
+    }
+
+    @Test
+    public void testApiVersionResponseStructParsingFallback() {
+        Struct struct = ApiVersionsResponse.DEFAULT_API_VERSIONS_RESPONSE.toStruct((short) 0);
+        ApiVersionsResponse response = ApiVersionsResponse.
+            fromStruct(struct, ApiKeys.API_VERSIONS.latestVersion());
+
+        assertEquals(Errors.NONE.code(), response.data.errorCode());
+    }
+
+    @Test(expected = SchemaException.class)
+    public void testApiVersionResponseStructParsingFallbackException() {
+        short version = 0;
+        ApiVersionsResponse.fromStruct(new Struct(ApiKeys.API_VERSIONS.requestSchema(version)), version);
+    }
+
+    @Test
+    public void testApiVersionResponseStructParsing() {
+        Struct struct = ApiVersionsResponse.DEFAULT_API_VERSIONS_RESPONSE.toStruct(ApiKeys.API_VERSIONS.latestVersion());
+        ApiVersionsResponse response = ApiVersionsResponse.
+            fromStruct(struct, ApiKeys.API_VERSIONS.latestVersion());
+
+        assertEquals(Errors.NONE.code(), response.data.errorCode());
+    }
+
+    private ResponseHeader createResponseHeader(short headerVersion) {
+        return new ResponseHeader(10, headerVersion);
     }
 
     private FindCoordinatorRequest createFindCoordinatorRequest(int version) {
-        return new FindCoordinatorRequest.Builder(FindCoordinatorRequest.CoordinatorType.GROUP, "test-group")
+        return new FindCoordinatorRequest.Builder(
+                new FindCoordinatorRequestData()
+                    .setKeyType(CoordinatorType.GROUP.id())
+                    .setKey("test-group"))
                 .build((short) version);
     }
 
     private FindCoordinatorResponse createFindCoordinatorResponse() {
-        return new FindCoordinatorResponse(Errors.NONE, new Node(10, "host1", 2014));
+        Node node = new Node(10, "host1", 2014);
+        return FindCoordinatorResponse.prepareResponse(Errors.NONE, node);
     }
 
     private FetchRequest createFetchRequest(int version, FetchMetadata metadata, List toForget) {
@@ -711,11 +873,11 @@ private FetchResponse createFetchResponse(int sessionId) {
         LinkedHashMap> responseData = new LinkedHashMap<>();
         MemoryRecords records = MemoryRecords.withRecords(CompressionType.NONE, new SimpleRecord("blah".getBytes()));
         responseData.put(new TopicPartition("test", 0), new FetchResponse.PartitionData<>(Errors.NONE,
-            1000000, FetchResponse.INVALID_LAST_STABLE_OFFSET, 0L, null, records));
+            1000000, FetchResponse.INVALID_LAST_STABLE_OFFSET, 0L, Optional.empty(), null, records));
         List abortedTransactions = Collections.singletonList(
             new FetchResponse.AbortedTransaction(234L, 999L));
         responseData.put(new TopicPartition("test", 1), new FetchResponse.PartitionData<>(Errors.NONE,
-            1000000, FetchResponse.INVALID_LAST_STABLE_OFFSET, 0L, abortedTransactions, MemoryRecords.EMPTY));
+            1000000, FetchResponse.INVALID_LAST_STABLE_OFFSET, 0L, Optional.empty(), abortedTransactions, MemoryRecords.EMPTY));
         return new FetchResponse<>(Errors.NONE, responseData, 25, sessionId);
     }
 
@@ -723,51 +885,93 @@ private FetchResponse createFetchResponse() {
         LinkedHashMap> responseData = new LinkedHashMap<>();
         MemoryRecords records = MemoryRecords.withRecords(CompressionType.NONE, new SimpleRecord("blah".getBytes()));
         responseData.put(new TopicPartition("test", 0), new FetchResponse.PartitionData<>(Errors.NONE,
-                1000000, FetchResponse.INVALID_LAST_STABLE_OFFSET, 0L, null, records));
+                1000000, FetchResponse.INVALID_LAST_STABLE_OFFSET, 0L, Optional.empty(), null, records));
 
         List abortedTransactions = Collections.singletonList(
                 new FetchResponse.AbortedTransaction(234L, 999L));
         responseData.put(new TopicPartition("test", 1), new FetchResponse.PartitionData<>(Errors.NONE,
-                1000000, FetchResponse.INVALID_LAST_STABLE_OFFSET, 0L, abortedTransactions, MemoryRecords.EMPTY));
+                1000000, FetchResponse.INVALID_LAST_STABLE_OFFSET, 0L, Optional.empty(), abortedTransactions, MemoryRecords.EMPTY));
 
         return new FetchResponse<>(Errors.NONE, responseData, 25, INVALID_SESSION_ID);
     }
 
     private HeartbeatRequest createHeartBeatRequest() {
-        return new HeartbeatRequest.Builder("group1", 1, "consumer1").build();
+        return new HeartbeatRequest.Builder(new HeartbeatRequestData()
+                .setGroupId("group1")
+                .setGenerationId(1)
+                .setMemberId("consumer1")).build();
     }
 
     private HeartbeatResponse createHeartBeatResponse() {
-        return new HeartbeatResponse(Errors.NONE);
+        return new HeartbeatResponse(new HeartbeatResponseData().setErrorCode(Errors.NONE.code()));
     }
 
     private JoinGroupRequest createJoinGroupRequest(int version) {
-        ByteBuffer metadata = ByteBuffer.wrap(new byte[] {});
-        List protocols = new ArrayList<>();
-        protocols.add(new JoinGroupRequest.ProtocolMetadata("consumer-range", metadata));
-        if (version == 0) {
-            return new JoinGroupRequest.Builder("group1", 30000, "consumer1", "consumer", protocols).
-                    build((short) version);
+        JoinGroupRequestData.JoinGroupRequestProtocolCollection protocols = new JoinGroupRequestData.JoinGroupRequestProtocolCollection(
+                Collections.singleton(
+                new JoinGroupRequestData.JoinGroupRequestProtocol()
+                        .setName("consumer-range")
+                        .setMetadata(new byte[0])).iterator()
+        );
+        if (version <= 4) {
+            return new JoinGroupRequest.Builder(
+                    new JoinGroupRequestData()
+                            .setGroupId("group1")
+                            .setSessionTimeoutMs(30000)
+                            .setMemberId("consumer1")
+                            .setGroupInstanceId(null)
+                            .setProtocolType("consumer")
+                            .setProtocols(protocols)
+                            .setRebalanceTimeoutMs(60000)) // v1 and above contains rebalance timeout
+                    .build((short) version);
         } else {
-            return new JoinGroupRequest.Builder("group1", 10000, "consumer1", "consumer", protocols).
-                    setRebalanceTimeout(60000).build();
+            return new JoinGroupRequest.Builder(
+                    new JoinGroupRequestData()
+                            .setGroupId("group1")
+                            .setSessionTimeoutMs(30000)
+                            .setMemberId("consumer1")
+                            .setGroupInstanceId("groupInstanceId") // v5 and above could set group instance id
+                            .setProtocolType("consumer")
+                            .setProtocols(protocols)
+                            .setRebalanceTimeoutMs(60000)) // v1 and above contains rebalance timeout
+                    .build((short) version);
         }
     }
 
     private JoinGroupResponse createJoinGroupResponse() {
-        Map members = new HashMap<>();
-        members.put("consumer1", ByteBuffer.wrap(new byte[]{}));
-        members.put("consumer2", ByteBuffer.wrap(new byte[]{}));
-        return new JoinGroupResponse(Errors.NONE, 1, "range", "consumer1", "leader", members);
+        List members = Arrays.asList(
+                new JoinGroupResponseData.JoinGroupResponseMember()
+                        .setMemberId("consumer1")
+                        .setMetadata(new byte[0]),
+                new JoinGroupResponseData.JoinGroupResponseMember()
+                        .setMemberId("consumer2")
+                        .setMetadata(new byte[0])
+        );
+
+        return new JoinGroupResponse(
+                new JoinGroupResponseData()
+                        .setErrorCode(Errors.NONE.code())
+                        .setGenerationId(1)
+                        .setProtocolName("range")
+                        .setLeader("leader")
+                        .setMemberId("consumer1")
+                        .setMembers(members)
+        );
     }
 
     private ListGroupsRequest createListGroupsRequest() {
-        return new ListGroupsRequest.Builder().build();
+        return new ListGroupsRequest.Builder(new ListGroupsRequestData()).build();
     }
 
     private ListGroupsResponse createListGroupsResponse() {
-        List groups = Collections.singletonList(new ListGroupsResponse.Group("test-group", "consumer"));
-        return new ListGroupsResponse(Errors.NONE, groups);
+        return new ListGroupsResponse(
+                new ListGroupsResponseData()
+                        .setErrorCode(Errors.NONE.code())
+                        .setGroups(Collections.singletonList(
+                                new ListGroupsResponseData.ListedGroup()
+                                        .setGroupId("test-group")
+                                        .setProtocolType("consumer")
+                )));
     }
 
     private DescribeGroupsRequest createDescribeGroupRequest() {
@@ -780,16 +984,24 @@ private DescribeGroupsResponse createDescribeGroupResponse() {
         String clientId = "consumer-1";
         String clientHost = "localhost";
         DescribeGroupsResponseData describeGroupsResponseData = new DescribeGroupsResponseData();
-        DescribeGroupsResponseData.DescribedGroupMember member = DescribeGroupsResponse.groupMember("memberId",
+        DescribeGroupsResponseData.DescribedGroupMember member = DescribeGroupsResponse.groupMember("memberId", null,
                 clientId, clientHost, new byte[0], new byte[0]);
-        DescribeGroupsResponseData.DescribedGroup metadata = DescribeGroupsResponse.groupMetadata("test-group", Errors.NONE,
-                "STABLE", "consumer", "roundrobin", asList(member), Collections.emptySet());
+        DescribedGroup metadata = DescribeGroupsResponse.groupMetadata("test-group",
+                                                                       Errors.NONE,
+                                                                       "STABLE",
+                                                                       "consumer",
+                                                                       "roundrobin",
+                                                                       Collections.singletonList(member),
+                                                                       Collections.emptySet());
         describeGroupsResponseData.groups().add(metadata);
         return new DescribeGroupsResponse(describeGroupsResponseData);
     }
 
     private LeaveGroupRequest createLeaveGroupRequest() {
-        return new LeaveGroupRequest.Builder(new LeaveGroupRequestData().setGroupId("group1").setMemberId("consumer1")).build();
+        return new LeaveGroupRequest.Builder(
+            "group1", Collections.singletonList(new MemberIdentity()
+                                                    .setMemberId("consumer1"))
+            ).build();
     }
 
     private LeaveGroupResponse createLeaveGroupResponse() {
@@ -797,13 +1009,21 @@ private LeaveGroupResponse createLeaveGroupResponse() {
     }
 
     private DeleteGroupsRequest createDeleteGroupsRequest() {
-        return new DeleteGroupsRequest.Builder(Collections.singleton("test-group")).build();
+        return new DeleteGroupsRequest.Builder(
+            new DeleteGroupsRequestData()
+                .setGroupsNames(Collections.singletonList("test-group"))
+        ).build();
     }
 
     private DeleteGroupsResponse createDeleteGroupsResponse() {
-        Map result = new HashMap<>();
-        result.put("test-group", Errors.NONE);
-        return new DeleteGroupsResponse(result);
+        DeletableGroupResultCollection result = new DeletableGroupResultCollection();
+        result.add(new DeletableGroupResult()
+                       .setGroupId("test-group")
+                       .setErrorCode(Errors.NONE.code()));
+        return new DeleteGroupsResponse(
+            new DeleteGroupsResponseData()
+                .setResults(result)
+        );
     }
 
     @SuppressWarnings("deprecation")
@@ -824,7 +1044,7 @@ private ListOffsetRequest createListOffsetRequest(int version) {
                     .forConsumer(true, IsolationLevel.READ_UNCOMMITTED)
                     .setTargetTimes(offsetData)
                     .build((short) version);
-        } else if (version == 2) {
+        } else if (version >= 2 && version <= 5) {
             Map offsetData = Collections.singletonMap(
                     new TopicPartition("test", 0),
                     new ListOffsetRequest.PartitionData(1000000L, Optional.of(5)));
@@ -844,7 +1064,7 @@ private ListOffsetResponse createListOffsetResponse(int version) {
             responseData.put(new TopicPartition("test", 0),
                     new ListOffsetResponse.PartitionData(Errors.NONE, asList(100L)));
             return new ListOffsetResponse(responseData);
-        } else if (version == 1 || version == 2) {
+        } else if (version >= 1 && version <= 5) {
             Map responseData = new HashMap<>();
             responseData.put(new TopicPartition("test", 0),
                     new ListOffsetResponse.PartitionData(Errors.NONE, 10000L, 100L, Optional.of(27)));
@@ -874,30 +1094,50 @@ private MetadataResponse createMetadataResponse() {
             asList(new MetadataResponse.PartitionMetadata(Errors.LEADER_NOT_AVAILABLE, 0, null,
                 Optional.empty(), replicas, isr, offlineReplicas))));
 
-        return new MetadataResponse(asList(node), null, MetadataResponse.NO_CONTROLLER_ID, allTopicMetadata);
+        return MetadataResponse.prepareResponse(asList(node), null, MetadataResponse.NO_CONTROLLER_ID, allTopicMetadata);
     }
 
-    @SuppressWarnings("deprecation")
     private OffsetCommitRequest createOffsetCommitRequest(int version) {
-        Map commitData = new HashMap<>();
-        commitData.put(new TopicPartition("test", 0), new OffsetCommitRequest.PartitionData(100,
-                RecordBatch.NO_PARTITION_LEADER_EPOCH, ""));
-        commitData.put(new TopicPartition("test", 1), new OffsetCommitRequest.PartitionData(200,
-                RecordBatch.NO_PARTITION_LEADER_EPOCH, null));
-        return new OffsetCommitRequest.Builder("group1", commitData)
-                .setGenerationId(100)
+        return new OffsetCommitRequest.Builder(new OffsetCommitRequestData()
+                .setGroupId("group1")
                 .setMemberId("consumer1")
-                .build((short) version);
+                .setGroupInstanceId(null)
+                .setGenerationId(100)
+                .setTopics(Collections.singletonList(
+                        new OffsetCommitRequestData.OffsetCommitRequestTopic()
+                                .setName("test")
+                                .setPartitions(Arrays.asList(
+                                        new OffsetCommitRequestData.OffsetCommitRequestPartition()
+                                                .setPartitionIndex(0)
+                                                .setCommittedOffset(100)
+                                                .setCommittedLeaderEpoch(RecordBatch.NO_PARTITION_LEADER_EPOCH)
+                                                .setCommittedMetadata(""),
+                                        new OffsetCommitRequestData.OffsetCommitRequestPartition()
+                                                .setPartitionIndex(1)
+                                                .setCommittedOffset(200)
+                                                .setCommittedLeaderEpoch(RecordBatch.NO_PARTITION_LEADER_EPOCH)
+                                                .setCommittedMetadata(null)
+                                ))
+                ))
+        ).build((short) version);
     }
 
     private OffsetCommitResponse createOffsetCommitResponse() {
-        Map responseData = new HashMap<>();
-        responseData.put(new TopicPartition("test", 0), Errors.NONE);
-        return new OffsetCommitResponse(responseData);
+        return new OffsetCommitResponse(new OffsetCommitResponseData()
+                .setTopics(Collections.singletonList(
+                        new OffsetCommitResponseData.OffsetCommitResponseTopic()
+                                .setName("test")
+                                .setPartitions(Collections.singletonList(
+                                        new OffsetCommitResponseData.OffsetCommitResponsePartition()
+                                                .setPartitionIndex(0)
+                                                .setErrorCode(Errors.NONE.code())
+                                ))
+                ))
+        );
     }
 
     private OffsetFetchRequest createOffsetFetchRequest(int version) {
-        return new OffsetFetchRequest.Builder("group1", singletonList(new TopicPartition("test11", 1)))
+        return new OffsetFetchRequest.Builder("group1", Collections.singletonList(new TopicPartition("test11", 1)))
                 .build((short) version);
     }
 
@@ -928,95 +1168,198 @@ private ProduceResponse createProduceResponse() {
         return new ProduceResponse(responseData, 0);
     }
 
+    private ProduceResponse createProduceResponseWithErrorMessage() {
+        Map responseData = new HashMap<>();
+        responseData.put(new TopicPartition("test", 0), new ProduceResponse.PartitionResponse(Errors.NONE,
+                10000, RecordBatch.NO_TIMESTAMP, 100, Collections.singletonList(new ProduceResponse.RecordError(0, "error message")),
+                "global error message"));
+        return new ProduceResponse(responseData, 0);
+    }
+
     private StopReplicaRequest createStopReplicaRequest(int version, boolean deletePartitions) {
         Set partitions = Utils.mkSet(new TopicPartition("test", 0));
-        return new StopReplicaRequest.Builder((short) version, 0, 1, 0, deletePartitions, partitions).build();
+        return new StopReplicaRequest.Builder((short) version, 0, 1, 0, 0, deletePartitions, partitions).build();
     }
 
     private StopReplicaResponse createStopReplicaResponse() {
-        Map responses = new HashMap<>();
-        responses.put(new TopicPartition("test", 0), Errors.NONE);
-        return new StopReplicaResponse(Errors.NONE, responses);
+        List partitions = new ArrayList<>();
+        partitions.add(new StopReplicaResponseData.StopReplicaPartitionError()
+            .setTopicName("test")
+            .setPartitionIndex(0)
+            .setErrorCode(Errors.NONE.code()));
+        return new StopReplicaResponse(new StopReplicaResponseData()
+            .setErrorCode(Errors.NONE.code())
+            .setPartitionErrors(partitions));
     }
 
     private ControlledShutdownRequest createControlledShutdownRequest() {
-        return new ControlledShutdownRequest.Builder(10, 0, ApiKeys.CONTROLLED_SHUTDOWN.latestVersion()).build();
+        ControlledShutdownRequestData data = new ControlledShutdownRequestData()
+                .setBrokerId(10)
+                .setBrokerEpoch(0L);
+        return new ControlledShutdownRequest.Builder(
+                data,
+                ApiKeys.CONTROLLED_SHUTDOWN.latestVersion()).build();
     }
 
     private ControlledShutdownRequest createControlledShutdownRequest(int version) {
-        return new ControlledShutdownRequest.Builder(10, 0, ApiKeys.CONTROLLED_SHUTDOWN.latestVersion()).build((short) version);
+        ControlledShutdownRequestData data = new ControlledShutdownRequestData()
+                .setBrokerId(10)
+                .setBrokerEpoch(0L);
+        return new ControlledShutdownRequest.Builder(
+                data,
+                ApiKeys.CONTROLLED_SHUTDOWN.latestVersion()).build((short) version);
     }
 
     private ControlledShutdownResponse createControlledShutdownResponse() {
-        Set topicPartitions = Utils.mkSet(
-                new TopicPartition("test2", 5),
-                new TopicPartition("test1", 10)
-        );
-        return new ControlledShutdownResponse(Errors.NONE, topicPartitions);
+        RemainingPartition p1 = new RemainingPartition()
+                .setTopicName("test2")
+                .setPartitionIndex(5);
+        RemainingPartition p2 = new RemainingPartition()
+                .setTopicName("test1")
+                .setPartitionIndex(10);
+        RemainingPartitionCollection pSet = new RemainingPartitionCollection();
+        pSet.add(p1);
+        pSet.add(p2);
+        ControlledShutdownResponseData data = new ControlledShutdownResponseData()
+                .setErrorCode(Errors.NONE.code())
+                .setRemainingPartitions(pSet);
+        return new ControlledShutdownResponse(data);
     }
 
     private LeaderAndIsrRequest createLeaderAndIsrRequest(int version) {
-        Map partitionStates = new HashMap<>();
+        List partitionStates = new ArrayList<>();
         List isr = asList(1, 2);
         List replicas = asList(1, 2, 3, 4);
-        partitionStates.put(new TopicPartition("topic5", 105),
-                new LeaderAndIsrRequest.PartitionState(0, 2, 1, new ArrayList<>(isr), 2, replicas, false));
-        partitionStates.put(new TopicPartition("topic5", 1),
-                new LeaderAndIsrRequest.PartitionState(1, 1, 1, new ArrayList<>(isr), 2, replicas, false));
-        partitionStates.put(new TopicPartition("topic20", 1),
-                new LeaderAndIsrRequest.PartitionState(1, 0, 1, new ArrayList<>(isr), 2, replicas, false));
+        partitionStates.add(new LeaderAndIsrPartitionState()
+            .setTopicName("topic5")
+            .setPartitionIndex(105)
+            .setControllerEpoch(0)
+            .setLeader(2)
+            .setLeaderEpoch(1)
+            .setIsr(isr)
+            .setZkVersion(2)
+            .setReplicas(replicas)
+            .setIsNew(false));
+        partitionStates.add(new LeaderAndIsrPartitionState()
+            .setTopicName("topic5")
+            .setPartitionIndex(1)
+            .setControllerEpoch(1)
+            .setLeader(1)
+            .setLeaderEpoch(1)
+            .setIsr(isr)
+            .setZkVersion(2)
+            .setReplicas(replicas)
+            .setIsNew(false));
+        partitionStates.add(new LeaderAndIsrPartitionState()
+            .setTopicName("topic20")
+            .setPartitionIndex(1)
+            .setControllerEpoch(1)
+            .setLeader(0)
+            .setLeaderEpoch(1)
+            .setIsr(isr)
+            .setZkVersion(2)
+            .setReplicas(replicas)
+            .setIsNew(false));
 
         Set leaders = Utils.mkSet(
                 new Node(0, "test0", 1223),
                 new Node(1, "test1", 1223)
         );
-        return new LeaderAndIsrRequest.Builder((short) version, 1, 10, 0, partitionStates, leaders).build();
+        return new LeaderAndIsrRequest.Builder((short) version, 1, 10, 0, 0, partitionStates, leaders).build();
     }
 
     private LeaderAndIsrResponse createLeaderAndIsrResponse() {
-        Map responses = new HashMap<>();
-        responses.put(new TopicPartition("test", 0), Errors.NONE);
-        return new LeaderAndIsrResponse(Errors.NONE, responses);
+        List partitions = new ArrayList<>();
+        partitions.add(new LeaderAndIsrResponseData.LeaderAndIsrPartitionError()
+            .setTopicName("test")
+            .setPartitionIndex(0)
+            .setErrorCode(Errors.NONE.code()));
+        return new LeaderAndIsrResponse(new LeaderAndIsrResponseData()
+            .setErrorCode(Errors.NONE.code())
+            .setPartitionErrors(partitions));
     }
 
     private UpdateMetadataRequest createUpdateMetadataRequest(int version, String rack) {
-        Map partitionStates = new HashMap<>();
+        List partitionStates = new ArrayList<>();
         List isr = asList(1, 2);
         List replicas = asList(1, 2, 3, 4);
         List offlineReplicas = asList();
-        partitionStates.put(new TopicPartition("topic5", 105),
-            new UpdateMetadataRequest.PartitionState(0, 2, 1, isr, 2, replicas, offlineReplicas));
-        partitionStates.put(new TopicPartition("topic5", 1),
-            new UpdateMetadataRequest.PartitionState(1, 1, 1, isr, 2, replicas, offlineReplicas));
-        partitionStates.put(new TopicPartition("topic20", 1),
-            new UpdateMetadataRequest.PartitionState(1, 0, 1, isr, 2, replicas, offlineReplicas));
+        partitionStates.add(new UpdateMetadataPartitionState()
+            .setTopicName("topic5")
+            .setPartitionIndex(105)
+            .setControllerEpoch(0)
+            .setLeader(2)
+            .setLeaderEpoch(1)
+            .setIsr(isr)
+            .setZkVersion(2)
+            .setReplicas(replicas)
+            .setOfflineReplicas(offlineReplicas));
+        partitionStates.add(new UpdateMetadataPartitionState()
+                .setTopicName("topic5")
+                .setPartitionIndex(1)
+                .setControllerEpoch(1)
+                .setLeader(1)
+                .setLeaderEpoch(1)
+                .setIsr(isr)
+                .setZkVersion(2)
+                .setReplicas(replicas)
+                .setOfflineReplicas(offlineReplicas));
+        partitionStates.add(new UpdateMetadataPartitionState()
+                .setTopicName("topic20")
+                .setPartitionIndex(1)
+                .setControllerEpoch(1)
+                .setLeader(0)
+                .setLeaderEpoch(1)
+                .setIsr(isr)
+                .setZkVersion(2)
+                .setReplicas(replicas)
+                .setOfflineReplicas(offlineReplicas));
 
         SecurityProtocol plaintext = SecurityProtocol.PLAINTEXT;
-        List endPoints1 = new ArrayList<>();
-        endPoints1.add(new UpdateMetadataRequest.EndPoint("host1", 1223, plaintext,
-                ListenerName.forSecurityProtocol(plaintext)));
-
-        List endPoints2 = new ArrayList<>();
-        endPoints2.add(new UpdateMetadataRequest.EndPoint("host1", 1244, plaintext,
-                ListenerName.forSecurityProtocol(plaintext)));
+        List endpoints1 = new ArrayList<>();
+        endpoints1.add(new UpdateMetadataEndpoint()
+            .setHost("host1")
+            .setPort(1223)
+            .setSecurityProtocol(plaintext.id)
+            .setListener(ListenerName.forSecurityProtocol(plaintext).value()));
+
+        List endpoints2 = new ArrayList<>();
+        endpoints2.add(new UpdateMetadataEndpoint()
+            .setHost("host1")
+            .setPort(1244)
+            .setSecurityProtocol(plaintext.id)
+            .setListener(ListenerName.forSecurityProtocol(plaintext).value()));
         if (version > 0) {
             SecurityProtocol ssl = SecurityProtocol.SSL;
-            endPoints2.add(new UpdateMetadataRequest.EndPoint("host2", 1234, ssl,
-                    ListenerName.forSecurityProtocol(ssl)));
-            endPoints2.add(new UpdateMetadataRequest.EndPoint("host2", 1334, ssl,
-                    new ListenerName("CLIENT")));
+            endpoints2.add(new UpdateMetadataEndpoint()
+                .setHost("host2")
+                .setPort(1234)
+                .setSecurityProtocol(ssl.id)
+                .setListener(ListenerName.forSecurityProtocol(ssl).value()));
+            endpoints2.add(new UpdateMetadataEndpoint()
+                .setHost("host2")
+                .setPort(1334)
+                .setSecurityProtocol(ssl.id));
+            if (version >= 3)
+                endpoints2.get(1).setListener("CLIENT");
         }
 
-        Set liveBrokers = Utils.mkSet(
-                new UpdateMetadataRequest.Broker(0, endPoints1, rack),
-                new UpdateMetadataRequest.Broker(1, endPoints2, rack)
+        List liveBrokers = Arrays.asList(
+            new UpdateMetadataBroker()
+                .setId(0)
+                .setEndpoints(endpoints1)
+                .setRack(rack),
+            new UpdateMetadataBroker()
+                .setId(1)
+                .setEndpoints(endpoints2)
+                .setRack(rack)
         );
-        return new UpdateMetadataRequest.Builder((short) version, 1, 10, 0, partitionStates,
-                liveBrokers).build();
+        return new UpdateMetadataRequest.Builder((short) version, 1, 10, 0, 0, partitionStates,
+            liveBrokers).build();
     }
 
     private UpdateMetadataResponse createUpdateMetadataResponse() {
-        return new UpdateMetadataResponse(Errors.NONE);
+        return new UpdateMetadataResponse(new UpdateMetadataResponseData().setErrorCode(Errors.NONE.code()));
     }
 
     private SaslHandshakeRequest createSaslHandshakeRequest() {
@@ -1027,7 +1370,7 @@ private SaslHandshakeRequest createSaslHandshakeRequest() {
     private SaslHandshakeResponse createSaslHandshakeResponse() {
         return new SaslHandshakeResponse(
                 new SaslHandshakeResponseData()
-                .setErrorCode(Errors.NONE.code()).setMechanisms(singletonList("GSSAPI")));
+                .setErrorCode(Errors.NONE.code()).setMechanisms(Collections.singletonList("GSSAPI")));
     }
 
     private SaslAuthenticateRequest createSaslAuthenticateRequest() {
@@ -1048,8 +1391,16 @@ private ApiVersionsRequest createApiVersionRequest() {
     }
 
     private ApiVersionsResponse createApiVersionResponse() {
-        List apiVersions = asList(new ApiVersionsResponse.ApiVersion((short) 0, (short) 0, (short) 2));
-        return new ApiVersionsResponse(Errors.NONE, apiVersions);
+        ApiVersionsResponseKeyCollection apiVersions = new ApiVersionsResponseKeyCollection();
+        apiVersions.add(new ApiVersionsResponseKey()
+            .setApiKey((short) 0)
+            .setMinVersion((short) 0)
+            .setMaxVersion((short) 2));
+
+        return new ApiVersionsResponse(new ApiVersionsResponseData()
+            .setErrorCode(Errors.NONE.code())
+            .setThrottleTimeMs(0)
+            .setApiKeys(apiVersions));
     }
 
     private CreateTopicsRequest createCreateTopicRequest(int version) {
@@ -1092,36 +1443,58 @@ private CreateTopicsResponse createCreateTopicResponse() {
     }
 
     private DeleteTopicsRequest createDeleteTopicsRequest() {
-        return new DeleteTopicsRequest.Builder(Utils.mkSet("my_t1", "my_t2"), 10000).build();
+        return new DeleteTopicsRequest.Builder(
+                new DeleteTopicsRequestData()
+                .setTopicNames(Arrays.asList("my_t1", "my_t2"))
+                .setTimeoutMs(1000)).build();
     }
 
     private DeleteTopicsResponse createDeleteTopicsResponse() {
-        Map errors = new HashMap<>();
-        errors.put("t1", Errors.INVALID_TOPIC_EXCEPTION);
-        errors.put("t2", Errors.TOPIC_AUTHORIZATION_FAILED);
-        return new DeleteTopicsResponse(errors);
+        DeleteTopicsResponseData data = new DeleteTopicsResponseData();
+        data.responses().add(new DeletableTopicResult()
+                .setName("t1")
+                .setErrorCode(Errors.INVALID_TOPIC_EXCEPTION.code()));
+        data.responses().add(new DeletableTopicResult()
+                .setName("t2")
+                .setErrorCode(Errors.TOPIC_AUTHORIZATION_FAILED.code()));
+        return new DeleteTopicsResponse(data);
     }
 
     private InitProducerIdRequest createInitPidRequest() {
-        return new InitProducerIdRequest.Builder(null, 100).build();
+        InitProducerIdRequestData requestData = new InitProducerIdRequestData()
+                .setTransactionalId(null)
+                .setTransactionTimeoutMs(100);
+        return new InitProducerIdRequest.Builder(requestData).build();
     }
 
     private InitProducerIdResponse createInitPidResponse() {
-        return new InitProducerIdResponse(0, Errors.NONE, 3332, (short) 3);
+        InitProducerIdResponseData responseData = new InitProducerIdResponseData()
+                .setErrorCode(Errors.NONE.code())
+                .setProducerEpoch((short) 3)
+                .setProducerId(3332)
+                .setThrottleTimeMs(0);
+        return new InitProducerIdResponse(responseData);
     }
 
-
-    private OffsetsForLeaderEpochRequest createLeaderEpochRequest() {
+    private Map createOffsetForLeaderEpochPartitionData() {
         Map epochs = new HashMap<>();
-
         epochs.put(new TopicPartition("topic1", 0),
                 new OffsetsForLeaderEpochRequest.PartitionData(Optional.of(0), 1));
         epochs.put(new TopicPartition("topic1", 1),
                 new OffsetsForLeaderEpochRequest.PartitionData(Optional.of(0), 1));
         epochs.put(new TopicPartition("topic2", 2),
                 new OffsetsForLeaderEpochRequest.PartitionData(Optional.empty(), 3));
+        return epochs;
+    }
+
+    private OffsetsForLeaderEpochRequest createLeaderEpochRequestForConsumer() {
+        Map epochs = createOffsetForLeaderEpochPartitionData();
+        return OffsetsForLeaderEpochRequest.Builder.forConsumer(epochs).build();
+    }
 
-        return new OffsetsForLeaderEpochRequest.Builder(ApiKeys.OFFSET_FOR_LEADER_EPOCH.latestVersion(), epochs).build();
+    private OffsetsForLeaderEpochRequest createLeaderEpochRequestForReplica(int version, int replicaId) {
+        Map epochs = createOffsetForLeaderEpochPartitionData();
+        return OffsetsForLeaderEpochRequest.Builder.forFollower((short) version, epochs, replicaId).build();
     }
 
     private OffsetsForLeaderEpochResponse createLeaderEpochResponse() {
@@ -1179,7 +1552,14 @@ private TxnOffsetCommitRequest createTxnOffsetCommitRequest() {
                     new TxnOffsetCommitRequest.CommittedOffset(100, null, Optional.empty()));
         offsets.put(new TopicPartition("topic", 74),
                 new TxnOffsetCommitRequest.CommittedOffset(100, "blah", Optional.of(27)));
-        return new TxnOffsetCommitRequest.Builder("transactionalId", "groupId", 21L, (short) 42, offsets).build();
+        return new TxnOffsetCommitRequest.Builder(
+            new TxnOffsetCommitRequestData()
+                .setTransactionalId("transactionalId")
+                .setGroupId("groupId")
+                .setProducerId(21L)
+                .setProducerEpoch((short) 42)
+                .setTopics(TxnOffsetCommitRequest.getTopics(offsets))
+            ).build();
     }
 
     private TxnOffsetCommitResponse createTxnOffsetCommitResponse() {
@@ -1252,7 +1632,7 @@ private DescribeConfigsRequest createDescribeConfigsRequestWithConfigEntries(int
         Map> resources = new HashMap<>();
         resources.put(new ConfigResource(ConfigResource.Type.BROKER, "0"), asList("foo", "bar"));
         resources.put(new ConfigResource(ConfigResource.Type.TOPIC, "topic"), null);
-        resources.put(new ConfigResource(ConfigResource.Type.TOPIC, "topic a"), Collections.emptyList());
+        resources.put(new ConfigResource(ConfigResource.Type.TOPIC, "topic a"), Collections.emptyList());
         return new DescribeConfigsRequest.Builder(resources).build((short) version);
     }
 
@@ -1314,31 +1694,60 @@ private CreatePartitionsResponse createCreatePartitionsResponse() {
     }
 
     private CreateDelegationTokenRequest createCreateTokenRequest() {
-        List renewers = new ArrayList<>();
-        renewers.add(SecurityUtils.parseKafkaPrincipal("User:user1"));
-        renewers.add(SecurityUtils.parseKafkaPrincipal("User:user2"));
-        return new CreateDelegationTokenRequest.Builder(renewers, System.currentTimeMillis()).build();
+        List renewers = new ArrayList<>();
+        renewers.add(new CreatableRenewers()
+                .setPrincipalType("User")
+                .setPrincipalName("user1"));
+        renewers.add(new CreatableRenewers()
+                .setPrincipalType("User")
+                .setPrincipalName("user2"));
+        return new CreateDelegationTokenRequest.Builder(new CreateDelegationTokenRequestData()
+                .setRenewers(renewers)
+                .setMaxLifetimeMs(System.currentTimeMillis())).build();
     }
 
     private CreateDelegationTokenResponse createCreateTokenResponse() {
-        return new CreateDelegationTokenResponse(20, Errors.NONE, SecurityUtils.parseKafkaPrincipal("User:user1"), System.currentTimeMillis(),
-            System.currentTimeMillis(), System.currentTimeMillis(), "token1", ByteBuffer.wrap("test".getBytes()));
+        CreateDelegationTokenResponseData data = new CreateDelegationTokenResponseData()
+                .setThrottleTimeMs(20)
+                .setErrorCode(Errors.NONE.code())
+                .setPrincipalType("User")
+                .setPrincipalName("user1")
+                .setIssueTimestampMs(System.currentTimeMillis())
+                .setExpiryTimestampMs(System.currentTimeMillis())
+                .setMaxTimestampMs(System.currentTimeMillis())
+                .setTokenId("token1")
+                .setHmac("test".getBytes());
+        return new CreateDelegationTokenResponse(data);
     }
 
     private RenewDelegationTokenRequest createRenewTokenRequest() {
-        return new RenewDelegationTokenRequest.Builder("test".getBytes(), System.currentTimeMillis()).build();
+        RenewDelegationTokenRequestData data = new RenewDelegationTokenRequestData()
+                .setHmac("test".getBytes())
+                .setRenewPeriodMs(System.currentTimeMillis());
+        return new RenewDelegationTokenRequest.Builder(data).build();
     }
 
     private RenewDelegationTokenResponse createRenewTokenResponse() {
-        return new RenewDelegationTokenResponse(20, Errors.NONE, System.currentTimeMillis());
+        RenewDelegationTokenResponseData data = new RenewDelegationTokenResponseData()
+                .setThrottleTimeMs(20)
+                .setErrorCode(Errors.NONE.code())
+                .setExpiryTimestampMs(System.currentTimeMillis());
+        return new RenewDelegationTokenResponse(data);
     }
 
     private ExpireDelegationTokenRequest createExpireTokenRequest() {
-        return new ExpireDelegationTokenRequest.Builder("test".getBytes(), System.currentTimeMillis()).build();
+        ExpireDelegationTokenRequestData data = new ExpireDelegationTokenRequestData()
+                .setHmac("test".getBytes())
+                .setExpiryTimePeriodMs(System.currentTimeMillis());
+        return new ExpireDelegationTokenRequest.Builder(data).build();
     }
 
     private ExpireDelegationTokenResponse createExpireTokenResponse() {
-        return new ExpireDelegationTokenResponse(20, Errors.NONE, System.currentTimeMillis());
+        ExpireDelegationTokenResponseData data = new ExpireDelegationTokenResponseData()
+                .setThrottleTimeMs(20)
+                .setErrorCode(Errors.NONE.code())
+                .setExpiryTimestampMs(System.currentTimeMillis());
+        return new ExpireDelegationTokenResponse(data);
     }
 
     private DescribeDelegationTokenRequest createDescribeTokenRequest() {
@@ -1367,32 +1776,158 @@ private DescribeDelegationTokenResponse createDescribeTokenResponse() {
         return new DescribeDelegationTokenResponse(20, Errors.NONE, tokenList);
     }
 
-    private ElectPreferredLeadersRequest createElectPreferredLeadersRequestNullPartitions() {
-        return new ElectPreferredLeadersRequest.Builder(
-                new ElectPreferredLeadersRequestData()
-                        .setTimeoutMs(100)
-                        .setTopicPartitions(null))
-                .build((short) 0);
+    private ElectLeadersRequest createElectLeadersRequestNullPartitions() {
+        return new ElectLeadersRequest.Builder(ElectionType.PREFERRED, null, 100).build((short) 1);
     }
 
-    private ElectPreferredLeadersRequest createElectPreferredLeadersRequest() {
-        ElectPreferredLeadersRequestData data = new ElectPreferredLeadersRequestData()
-                .setTimeoutMs(100);
-        data.topicPartitions().add(new TopicPartitions().setTopic("data").setPartitionId(asList(1, 2)));
-        return new ElectPreferredLeadersRequest.Builder(data).build((short) 0);
+    private ElectLeadersRequest createElectLeadersRequest() {
+        List partitions = asList(new TopicPartition("data", 1), new TopicPartition("data", 2));
+
+        return new ElectLeadersRequest.Builder(ElectionType.PREFERRED, partitions, 100).build((short) 1);
     }
 
-    private ElectPreferredLeadersResponse createElectPreferredLeadersResponse() {
-        ElectPreferredLeadersResponseData data = new ElectPreferredLeadersResponseData().setThrottleTimeMs(200);
-        ReplicaElectionResult resultsByTopic = new ReplicaElectionResult().setTopic("myTopic");
-        resultsByTopic.partitionResult().add(new PartitionResult().setPartitionId(0)
-                .setErrorCode(Errors.NONE.code())
-                .setErrorMessage(Errors.NONE.message()));
-        resultsByTopic.partitionResult().add(new PartitionResult().setPartitionId(1)
-                .setErrorCode(Errors.UNKNOWN_TOPIC_OR_PARTITION.code())
-                .setErrorMessage(Errors.UNKNOWN_TOPIC_OR_PARTITION.message()));
-        data.replicaElectionResults().add(resultsByTopic);
-        return new ElectPreferredLeadersResponse(data);
+    private ElectLeadersResponse createElectLeadersResponse() {
+        String topic = "myTopic";
+        List electionResults = new ArrayList<>();
+        ReplicaElectionResult electionResult = new ReplicaElectionResult();
+        electionResult.setTopic(topic);
+        // Add partition 1 result
+        PartitionResult partitionResult = new PartitionResult();
+        partitionResult.setPartitionId(0);
+        partitionResult.setErrorCode(ApiError.NONE.error().code());
+        partitionResult.setErrorMessage(ApiError.NONE.message());
+        electionResult.partitionResult().add(partitionResult);
+
+        // Add partition 2 result
+        partitionResult = new PartitionResult();
+        partitionResult.setPartitionId(1);
+        partitionResult.setErrorCode(Errors.UNKNOWN_TOPIC_OR_PARTITION.code());
+        partitionResult.setErrorMessage(Errors.UNKNOWN_TOPIC_OR_PARTITION.message());
+        electionResult.partitionResult().add(partitionResult);
+
+        return new ElectLeadersResponse(200, Errors.NONE.code(), electionResults);
+    }
+
+    private IncrementalAlterConfigsRequest createIncrementalAlterConfigsRequest() {
+        IncrementalAlterConfigsRequestData data = new IncrementalAlterConfigsRequestData();
+        AlterableConfig alterableConfig = new AlterableConfig()
+                .setName("retention.ms")
+                .setConfigOperation((byte) 0)
+                .setValue("100");
+        IncrementalAlterConfigsRequestData.AlterableConfigCollection alterableConfigs = new IncrementalAlterConfigsRequestData.AlterableConfigCollection();
+        alterableConfigs.add(alterableConfig);
+
+        data.resources().add(new AlterConfigsResource()
+                .setResourceName("testtopic")
+                .setResourceType(ResourceType.TOPIC.code())
+                .setConfigs(alterableConfigs));
+        return new IncrementalAlterConfigsRequest.Builder(data).build((short) 0);
+    }
+
+    private IncrementalAlterConfigsResponse createIncrementalAlterConfigsResponse() {
+        IncrementalAlterConfigsResponseData data = new IncrementalAlterConfigsResponseData();
+
+        data.responses().add(new AlterConfigsResourceResponse()
+                .setResourceName("testtopic")
+                .setResourceType(ResourceType.TOPIC.code())
+                .setErrorCode(Errors.INVALID_REQUEST.code())
+                .setErrorMessage("Duplicate Keys"));
+        return new IncrementalAlterConfigsResponse(data);
+    }
+
+    private AlterPartitionReassignmentsRequest createAlterPartitionReassignmentsRequest() {
+        AlterPartitionReassignmentsRequestData data = new AlterPartitionReassignmentsRequestData();
+        data.topics().add(
+                new AlterPartitionReassignmentsRequestData.ReassignableTopic().setName("topic").setPartitions(
+                        Collections.singletonList(
+                                new AlterPartitionReassignmentsRequestData.ReassignablePartition().setPartitionIndex(0).setReplicas(null)
+                        )
+                )
+        );
+        return new AlterPartitionReassignmentsRequest.Builder(data).build((short) 0);
+    }
+
+    private AlterPartitionReassignmentsResponse createAlterPartitionReassignmentsResponse() {
+        AlterPartitionReassignmentsResponseData data = new AlterPartitionReassignmentsResponseData();
+        data.responses().add(
+                new AlterPartitionReassignmentsResponseData.ReassignableTopicResponse()
+                        .setName("topic")
+                        .setPartitions(Collections.singletonList(
+                                new AlterPartitionReassignmentsResponseData.ReassignablePartitionResponse()
+                                        .setPartitionIndex(0)
+                                        .setErrorCode(Errors.NO_REASSIGNMENT_IN_PROGRESS.code())
+                                        .setErrorMessage("No reassignment is in progress for topic topic partition 0")
+                                )
+                        )
+        );
+        return new AlterPartitionReassignmentsResponse(data);
+    }
+
+    private ListPartitionReassignmentsRequest createListPartitionReassignmentsRequest() {
+        ListPartitionReassignmentsRequestData data = new ListPartitionReassignmentsRequestData();
+        data.setTopics(
+            Collections.singletonList(
+                new ListPartitionReassignmentsRequestData.ListPartitionReassignmentsTopics()
+                    .setName("topic")
+                    .setPartitionIndexes(Collections.singletonList(1))
+            )
+        );
+        return new ListPartitionReassignmentsRequest.Builder(data).build((short) 0);
+    }
+
+    private ListPartitionReassignmentsResponse createListPartitionReassignmentsResponse() {
+        ListPartitionReassignmentsResponseData data = new ListPartitionReassignmentsResponseData();
+        data.setTopics(Collections.singletonList(
+            new ListPartitionReassignmentsResponseData.OngoingTopicReassignment()
+                        .setName("topic")
+                        .setPartitions(Collections.singletonList(
+                                new ListPartitionReassignmentsResponseData.OngoingPartitionReassignment()
+                                        .setPartitionIndex(0)
+                                        .setReplicas(Arrays.asList(1, 2))
+                                        .setAddingReplicas(Collections.singletonList(2))
+                                        .setRemovingReplicas(Collections.singletonList(1))
+                                )
+                        )
+        ));
+        return new ListPartitionReassignmentsResponse(data);
+    }
+
+    private OffsetDeleteRequest createOffsetDeleteRequest() {
+        OffsetDeleteRequestTopicCollection topics = new OffsetDeleteRequestTopicCollection();
+        topics.add(new OffsetDeleteRequestTopic()
+            .setName("topic1")
+            .setPartitions(Collections.singletonList(
+                new OffsetDeleteRequestPartition()
+                    .setPartitionIndex(0)
+                )
+            )
+        );
+
+        OffsetDeleteRequestData data = new OffsetDeleteRequestData();
+        data.setGroupId("group1");
+        data.setTopics(topics);
+
+        return new OffsetDeleteRequest.Builder(data).build((short) 0);
+    }
+
+    private OffsetDeleteResponse createOffsetDeleteResponse() {
+        OffsetDeleteResponsePartitionCollection partitions = new OffsetDeleteResponsePartitionCollection();
+        partitions.add(new OffsetDeleteResponsePartition()
+            .setPartitionIndex(0)
+            .setErrorCode(Errors.NONE.code())
+        );
+
+        OffsetDeleteResponseTopicCollection topics = new OffsetDeleteResponseTopicCollection();
+        topics.add(new OffsetDeleteResponseTopic()
+            .setName("topic1")
+            .setPartitions(partitions)
+        );
+
+        OffsetDeleteResponseData data = new OffsetDeleteResponseData();
+        data.setErrorCode(Errors.NONE.code());
+        data.setTopics(topics);
+
+        return new OffsetDeleteResponse(data);
     }
 
 }
diff --git a/clients/src/test/java/org/apache/kafka/common/requests/StopReplicaRequestTest.java b/clients/src/test/java/org/apache/kafka/common/requests/StopReplicaRequestTest.java
new file mode 100644
index 0000000000000..ae1859ce1a73a
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/StopReplicaRequestTest.java
@@ -0,0 +1,65 @@
+/*
+ * 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.requests;
+
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.errors.ClusterAuthorizationException;
+import org.apache.kafka.common.errors.UnsupportedVersionException;
+import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.common.protocol.MessageTestUtil;
+import org.apache.kafka.test.TestUtils;
+import org.junit.Test;
+
+import java.util.Collections;
+import java.util.Set;
+
+import static org.apache.kafka.common.protocol.ApiKeys.STOP_REPLICA;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
+
+public class StopReplicaRequestTest {
+
+    @Test
+    public void testUnsupportedVersion() {
+        StopReplicaRequest.Builder builder = new StopReplicaRequest.Builder(
+                (short) (STOP_REPLICA.latestVersion() + 1),
+                0, 0, 0L, 0L, false, Collections.emptyList());
+        assertThrows(UnsupportedVersionException.class, builder::build);
+    }
+
+    @Test
+    public void testGetErrorResponse() {
+        for (short version = STOP_REPLICA.oldestVersion(); version < STOP_REPLICA.latestVersion(); version++) {
+            StopReplicaRequest.Builder builder = new StopReplicaRequest.Builder(version,
+                    0, 0, 0L, 0L, false, Collections.emptyList());
+            StopReplicaRequest request = builder.build();
+            StopReplicaResponse response = request.getErrorResponse(0,
+                    new ClusterAuthorizationException("Not authorized"));
+            assertEquals(Errors.CLUSTER_AUTHORIZATION_FAILED, response.error());
+        }
+    }
+
+    @Test
+    public void testStopReplicaRequestNormalization() {
+        Set tps = TestUtils.generateRandomTopicPartitions(10, 10);
+        StopReplicaRequest.Builder builder = new StopReplicaRequest.Builder((short) 5, 0, 0, 0, 0, false, tps);
+        assertTrue(MessageTestUtil.messageSize(builder.build((short) 1).data(), (short) 1) <
+            MessageTestUtil.messageSize(builder.build((short) 0).data(), (short) 0));
+    }
+
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/requests/StopReplicaResponseTest.java b/clients/src/test/java/org/apache/kafka/common/requests/StopReplicaResponseTest.java
index b6d4bdafb9b44..9f5c006785f86 100644
--- a/clients/src/test/java/org/apache/kafka/common/requests/StopReplicaResponseTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/requests/StopReplicaResponseTest.java
@@ -17,13 +17,16 @@
 package org.apache.kafka.common.requests;
 
 import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.message.StopReplicaResponseData;
+import org.apache.kafka.common.message.StopReplicaResponseData.StopReplicaPartitionError;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
 import org.apache.kafka.common.utils.Utils;
 import org.junit.Test;
 
+import java.util.ArrayList;
 import java.util.Collections;
-import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 
 import static org.junit.Assert.assertEquals;
@@ -33,7 +36,7 @@ public class StopReplicaResponseTest {
 
     @Test
     public void testErrorCountsFromGetErrorResponse() {
-        StopReplicaRequest request = new StopReplicaRequest.Builder(ApiKeys.STOP_REPLICA.latestVersion(), 15, 20, 0, false,
+        StopReplicaRequest request = new StopReplicaRequest.Builder(ApiKeys.STOP_REPLICA.latestVersion(), 15, 20, 0, 0, false,
                 Utils.mkSet(new TopicPartition("foo", 0), new TopicPartition("foo", 1))).build();
         StopReplicaResponse response = request.getErrorResponse(0, Errors.CLUSTER_AUTHORIZATION_FAILED.exception());
         assertEquals(Collections.singletonMap(Errors.CLUSTER_AUTHORIZATION_FAILED, 2), response.errorCounts());
@@ -41,19 +44,25 @@ public void testErrorCountsFromGetErrorResponse() {
 
     @Test
     public void testErrorCountsWithTopLevelError() {
-        Map errors = new HashMap<>();
-        errors.put(new TopicPartition("foo", 0), Errors.NONE);
-        errors.put(new TopicPartition("foo", 1), Errors.NOT_LEADER_FOR_PARTITION);
-        StopReplicaResponse response = new StopReplicaResponse(Errors.UNKNOWN_SERVER_ERROR, errors);
+        List errors = new ArrayList<>();
+        errors.add(new StopReplicaPartitionError().setTopicName("foo").setPartitionIndex(0));
+        errors.add(new StopReplicaPartitionError().setTopicName("foo").setPartitionIndex(1)
+            .setErrorCode(Errors.NOT_LEADER_FOR_PARTITION.code()));
+        StopReplicaResponse response = new StopReplicaResponse(new StopReplicaResponseData()
+            .setErrorCode(Errors.UNKNOWN_SERVER_ERROR.code())
+            .setPartitionErrors(errors));
         assertEquals(Collections.singletonMap(Errors.UNKNOWN_SERVER_ERROR, 2), response.errorCounts());
     }
 
     @Test
     public void testErrorCountsNoTopLevelError() {
-        Map errors = new HashMap<>();
-        errors.put(new TopicPartition("foo", 0), Errors.NONE);
-        errors.put(new TopicPartition("foo", 1), Errors.CLUSTER_AUTHORIZATION_FAILED);
-        StopReplicaResponse response = new StopReplicaResponse(Errors.NONE, errors);
+        List errors = new ArrayList<>();
+        errors.add(new StopReplicaPartitionError().setTopicName("foo").setPartitionIndex(0));
+        errors.add(new StopReplicaPartitionError().setTopicName("foo").setPartitionIndex(1)
+            .setErrorCode(Errors.CLUSTER_AUTHORIZATION_FAILED.code()));
+        StopReplicaResponse response = new StopReplicaResponse(new StopReplicaResponseData()
+            .setErrorCode(Errors.NONE.code())
+            .setPartitionErrors(errors));
         Map errorCounts = response.errorCounts();
         assertEquals(2, errorCounts.size());
         assertEquals(1, errorCounts.get(Errors.NONE).intValue());
@@ -62,14 +71,15 @@ public void testErrorCountsNoTopLevelError() {
 
     @Test
     public void testToString() {
-        Map errors = new HashMap<>();
-        errors.put(new TopicPartition("foo", 0), Errors.NONE);
-        errors.put(new TopicPartition("foo", 1), Errors.CLUSTER_AUTHORIZATION_FAILED);
-        StopReplicaResponse response = new StopReplicaResponse(Errors.NONE, errors);
+        List errors = new ArrayList<>();
+        errors.add(new StopReplicaPartitionError().setTopicName("foo").setPartitionIndex(0));
+        errors.add(new StopReplicaPartitionError().setTopicName("foo").setPartitionIndex(1)
+            .setErrorCode(Errors.CLUSTER_AUTHORIZATION_FAILED.code()));
+        StopReplicaResponse response = new StopReplicaResponse(new StopReplicaResponseData().setPartitionErrors(errors));
         String responseStr = response.toString();
         assertTrue(responseStr.contains(StopReplicaResponse.class.getSimpleName()));
         assertTrue(responseStr.contains(errors.toString()));
-        assertTrue(responseStr.contains(Errors.NONE.name()));
+        assertTrue(responseStr.contains("errorCode=" + Errors.NONE.code()));
     }
 
 }
diff --git a/clients/src/test/java/org/apache/kafka/common/requests/SyncGroupRequestTest.java b/clients/src/test/java/org/apache/kafka/common/requests/SyncGroupRequestTest.java
new file mode 100644
index 0000000000000..e4b40a45395a2
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/SyncGroupRequestTest.java
@@ -0,0 +1,34 @@
+/*
+ * 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.requests;
+
+import org.apache.kafka.common.errors.UnsupportedVersionException;
+import org.apache.kafka.common.message.SyncGroupRequestData;
+import org.junit.Test;
+
+public class SyncGroupRequestTest {
+
+    @Test(expected = UnsupportedVersionException.class)
+    public void testRequestVersionCompatibilityFailBuild() {
+        new SyncGroupRequest.Builder(
+                new SyncGroupRequestData()
+                        .setGroupId("groupId")
+                        .setMemberId("consumerId")
+                        .setGroupInstanceId("groupInstanceId")
+        ).build((short) 2);
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/requests/TxnOffsetCommitRequestTest.java b/clients/src/test/java/org/apache/kafka/common/requests/TxnOffsetCommitRequestTest.java
new file mode 100644
index 0000000000000..ff1d1ab872b6c
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/TxnOffsetCommitRequestTest.java
@@ -0,0 +1,109 @@
+/*
+ * 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.requests;
+
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.message.TxnOffsetCommitRequestData;
+import org.apache.kafka.common.message.TxnOffsetCommitRequestData.TxnOffsetCommitRequestPartition;
+import org.apache.kafka.common.message.TxnOffsetCommitRequestData.TxnOffsetCommitRequestTopic;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.common.requests.TxnOffsetCommitRequest.CommittedOffset;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+
+import static org.junit.Assert.assertEquals;
+
+public class TxnOffsetCommitRequestTest extends OffsetCommitRequestTest {
+
+    private static String transactionalId = "transactionalId";
+    private static int producerId = 10;
+    private static short producerEpoch = 1;
+
+    private static TxnOffsetCommitRequestData data;
+
+    @Before
+    @Override
+    public void setUp() {
+        super.setUp();
+        data = new TxnOffsetCommitRequestData()
+                   .setGroupId(groupId)
+                   .setTransactionalId(transactionalId)
+                   .setProducerId(producerId)
+                   .setProducerEpoch(producerEpoch)
+                   .setTopics(Arrays.asList(
+                       new TxnOffsetCommitRequestTopic()
+                           .setName(topicOne)
+                           .setPartitions(Collections.singletonList(
+                               new TxnOffsetCommitRequestPartition()
+                                   .setPartitionIndex(partitionOne)
+                                   .setCommittedOffset(offset)
+                                   .setCommittedLeaderEpoch(leaderEpoch)
+                                   .setCommittedMetadata(metadata)
+                           )),
+                       new TxnOffsetCommitRequestTopic()
+                           .setName(topicTwo)
+                           .setPartitions(Collections.singletonList(
+                               new TxnOffsetCommitRequestPartition()
+                                   .setPartitionIndex(partitionTwo)
+                                   .setCommittedOffset(offset)
+                                   .setCommittedLeaderEpoch(leaderEpoch)
+                                   .setCommittedMetadata(metadata)
+                           ))
+                   ));
+    }
+
+    @Test
+    @Override
+    public void testConstructor() {
+        Map expectedOffsets = new HashMap<>();
+        expectedOffsets.put(new TopicPartition(topicOne, partitionOne),
+                            new CommittedOffset(
+                                offset,
+                                metadata,
+                                Optional.of((int) leaderEpoch)));
+        expectedOffsets.put(new TopicPartition(topicTwo, partitionTwo),
+                            new CommittedOffset(
+                                offset,
+                                metadata,
+                                Optional.of((int) leaderEpoch)));
+
+        TxnOffsetCommitRequest.Builder builder = new TxnOffsetCommitRequest.Builder(data);
+        Map errorsMap = new HashMap<>();
+        errorsMap.put(new TopicPartition(topicOne, partitionOne), Errors.NOT_COORDINATOR);
+        errorsMap.put(new TopicPartition(topicTwo, partitionTwo), Errors.NOT_COORDINATOR);
+
+        for (short version = 0; version <= ApiKeys.TXN_OFFSET_COMMIT.latestVersion(); version++) {
+            TxnOffsetCommitRequest request = builder.build(version);
+            assertEquals(expectedOffsets, request.offsets());
+            assertEquals(data.topics(), TxnOffsetCommitRequest.getTopics(request.offsets()));
+
+            TxnOffsetCommitResponse response =
+                request.getErrorResponse(throttleTimeMs, Errors.NOT_COORDINATOR.exception());
+
+            assertEquals(errorsMap, response.errors());
+            assertEquals(Collections.singletonMap(Errors.NOT_COORDINATOR, 2), response.errorCounts());
+            assertEquals(throttleTimeMs, response.throttleTimeMs());
+        }
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/requests/TxnOffsetCommitResponseTest.java b/clients/src/test/java/org/apache/kafka/common/requests/TxnOffsetCommitResponseTest.java
new file mode 100644
index 0000000000000..be4f5df359a05
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/TxnOffsetCommitResponseTest.java
@@ -0,0 +1,67 @@
+/*
+ * 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.requests;
+
+import org.apache.kafka.common.message.TxnOffsetCommitResponseData;
+import org.apache.kafka.common.message.TxnOffsetCommitResponseData.TxnOffsetCommitResponsePartition;
+import org.apache.kafka.common.message.TxnOffsetCommitResponseData.TxnOffsetCommitResponseTopic;
+import org.apache.kafka.common.protocol.ApiKeys;
+
+import org.junit.Test;
+
+import java.util.Arrays;
+import java.util.Collections;
+
+import static org.junit.Assert.assertEquals;
+
+public class TxnOffsetCommitResponseTest extends OffsetCommitResponseTest {
+
+    @Test
+    @Override
+    public void testConstructorWithErrorResponse() {
+        TxnOffsetCommitResponse response = new TxnOffsetCommitResponse(throttleTimeMs, errorsMap);
+
+        assertEquals(errorsMap, response.errors());
+        assertEquals(expectedErrorCounts, response.errorCounts());
+        assertEquals(throttleTimeMs, response.throttleTimeMs());
+    }
+
+    @Test
+    @Override
+    public void testConstructorWithStruct() {
+        TxnOffsetCommitResponseData data = new TxnOffsetCommitResponseData()
+            .setThrottleTimeMs(throttleTimeMs)
+            .setTopics(Arrays.asList(
+                new TxnOffsetCommitResponseTopic().setPartitions(
+                    Collections.singletonList(new TxnOffsetCommitResponsePartition()
+                                                  .setPartitionIndex(partitionOne)
+                                                  .setErrorCode(errorOne.code()))),
+                new TxnOffsetCommitResponseTopic().setPartitions(
+                    Collections.singletonList(new TxnOffsetCommitResponsePartition()
+                                                  .setPartitionIndex(partitionTwo)
+                                                  .setErrorCode(errorTwo.code()))
+                    )
+            ));
+
+        for (short version = 0; version <= ApiKeys.TXN_OFFSET_COMMIT.latestVersion(); version++) {
+            TxnOffsetCommitResponse response = new TxnOffsetCommitResponse(data.toStruct(version), version);
+            assertEquals(expectedErrorCounts, response.errorCounts());
+            assertEquals(throttleTimeMs, response.throttleTimeMs());
+            assertEquals(version >= 1, response.shouldClientThrottle(version));
+        }
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/requests/UpdateMetadataRequestTest.java b/clients/src/test/java/org/apache/kafka/common/requests/UpdateMetadataRequestTest.java
new file mode 100644
index 0000000000000..4758936bfd97c
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/UpdateMetadataRequestTest.java
@@ -0,0 +1,213 @@
+/*
+ * 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.requests;
+
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.errors.ClusterAuthorizationException;
+import org.apache.kafka.common.errors.UnsupportedVersionException;
+import org.apache.kafka.common.message.UpdateMetadataRequestData;
+import org.apache.kafka.common.message.UpdateMetadataRequestData.UpdateMetadataBroker;
+import org.apache.kafka.common.message.UpdateMetadataRequestData.UpdateMetadataEndpoint;
+import org.apache.kafka.common.message.UpdateMetadataRequestData.UpdateMetadataPartitionState;
+import org.apache.kafka.common.network.ListenerName;
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
+import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.common.protocol.MessageTestUtil;
+import org.apache.kafka.common.security.auth.SecurityProtocol;
+import org.apache.kafka.test.TestUtils;
+import org.junit.Test;
+
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+import java.util.stream.Collectors;
+import java.util.stream.StreamSupport;
+
+import static java.util.Arrays.asList;
+import static java.util.Collections.emptyList;
+import static org.apache.kafka.common.protocol.ApiKeys.UPDATE_METADATA;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
+
+public class UpdateMetadataRequestTest {
+
+    @Test
+    public void testUnsupportedVersion() {
+        UpdateMetadataRequest.Builder builder = new UpdateMetadataRequest.Builder(
+                (short) (UPDATE_METADATA.latestVersion() + 1), 0, 0, 0, 0,
+                Collections.emptyList(), Collections.emptyList());
+        assertThrows(UnsupportedVersionException.class, builder::build);
+    }
+
+    @Test
+    public void testGetErrorResponse() {
+        for (short version = UPDATE_METADATA.oldestVersion(); version < UPDATE_METADATA.latestVersion(); version++) {
+            UpdateMetadataRequest.Builder builder = new UpdateMetadataRequest.Builder(
+                    version, 0, 0, 0, 0, Collections.emptyList(), Collections.emptyList());
+            UpdateMetadataRequest request = builder.build();
+            UpdateMetadataResponse response = request.getErrorResponse(0,
+                    new ClusterAuthorizationException("Not authorized"));
+            assertEquals(Errors.CLUSTER_AUTHORIZATION_FAILED, response.error());
+        }
+    }
+
+    /**
+     * Verifies the logic we have in UpdateMetadataRequest to present a unified interface across the various versions
+     * works correctly. For example, `UpdateMetadataPartitionState.topicName` is not serialiazed/deserialized in
+     * recent versions, but we set it manually so that we can always present the ungrouped partition states
+     * independently of the version.
+     */
+    @Test
+    public void testVersionLogic() {
+        for (short version = UPDATE_METADATA.oldestVersion(); version <= UPDATE_METADATA.latestVersion(); version++) {
+            List partitionStates = asList(
+                new UpdateMetadataPartitionState()
+                    .setTopicName("topic0")
+                    .setPartitionIndex(0)
+                    .setControllerEpoch(2)
+                    .setLeader(0)
+                    .setLeaderEpoch(10)
+                    .setIsr(asList(0, 1))
+                    .setZkVersion(10)
+                    .setReplicas(asList(0, 1, 2))
+                    .setOfflineReplicas(asList(2)),
+                new UpdateMetadataPartitionState()
+                    .setTopicName("topic0")
+                    .setPartitionIndex(1)
+                    .setControllerEpoch(2)
+                    .setLeader(1)
+                    .setLeaderEpoch(11)
+                    .setIsr(asList(1, 2, 3))
+                    .setZkVersion(11)
+                    .setReplicas(asList(1, 2, 3))
+                    .setOfflineReplicas(emptyList()),
+                new UpdateMetadataPartitionState()
+                    .setTopicName("topic1")
+                    .setPartitionIndex(0)
+                    .setControllerEpoch(2)
+                    .setLeader(2)
+                    .setLeaderEpoch(11)
+                    .setIsr(asList(2, 3))
+                    .setZkVersion(11)
+                    .setReplicas(asList(2, 3, 4))
+                    .setOfflineReplicas(emptyList())
+            );
+
+            List broker0Endpoints = new ArrayList<>();
+            broker0Endpoints.add(
+                new UpdateMetadataEndpoint()
+                    .setHost("host0")
+                    .setPort(9090)
+                    .setSecurityProtocol(SecurityProtocol.PLAINTEXT.id));
+
+            // Non plaintext endpoints are only supported from version 1
+            if (version >= 1) {
+                broker0Endpoints.add(new UpdateMetadataEndpoint()
+                    .setHost("host0")
+                    .setPort(9091)
+                    .setSecurityProtocol(SecurityProtocol.SSL.id));
+            }
+
+            // Custom listeners are only supported from version 3
+            if (version >= 3) {
+                broker0Endpoints.get(0).setListener("listener0");
+                broker0Endpoints.get(1).setListener("listener1");
+            }
+
+            List liveBrokers = asList(
+                new UpdateMetadataBroker()
+                    .setId(0)
+                    .setRack("rack0")
+                    .setEndpoints(broker0Endpoints),
+                new UpdateMetadataBroker()
+                    .setId(1)
+                    .setEndpoints(asList(
+                        new UpdateMetadataEndpoint()
+                            .setHost("host1")
+                            .setPort(9090)
+                            .setSecurityProtocol(SecurityProtocol.PLAINTEXT.id)
+                            .setListener("PLAINTEXT")
+                    ))
+            );
+
+            UpdateMetadataRequest request = new UpdateMetadataRequest.Builder(version, 1, 2, 3, 3,
+                partitionStates, liveBrokers).build();
+
+            assertEquals(new HashSet<>(partitionStates), iterableToSet(request.partitionStates()));
+            assertEquals(liveBrokers, request.liveBrokers());
+            assertEquals(1, request.controllerId());
+            assertEquals(2, request.controllerEpoch());
+            assertEquals(3, request.brokerEpoch());
+
+            ByteBuffer byteBuffer = MessageTestUtil.messageToByteBuffer(request.data(), request.version());
+            UpdateMetadataRequest deserializedRequest = new UpdateMetadataRequest(new UpdateMetadataRequestData(
+                    new ByteBufferAccessor(byteBuffer), version), version);
+
+            // Unset fields that are not supported in this version as the deserialized request won't have them
+
+            // Rack is only supported from version 2
+            if (version < 2) {
+                for (UpdateMetadataBroker liveBroker : liveBrokers)
+                    liveBroker.setRack("");
+            }
+
+            // Non plaintext listener name is only supported from version 3
+            if (version < 3) {
+                for (UpdateMetadataBroker liveBroker : liveBrokers) {
+                    for (UpdateMetadataEndpoint endpoint : liveBroker.endpoints()) {
+                        SecurityProtocol securityProtocol = SecurityProtocol.forId(endpoint.securityProtocol());
+                        endpoint.setListener(ListenerName.forSecurityProtocol(securityProtocol).value());
+                    }
+                }
+            }
+
+            // Offline replicas are only supported from version 4
+            if (version < 4)
+                partitionStates.get(0).setOfflineReplicas(emptyList());
+
+            assertEquals(new HashSet<>(partitionStates), iterableToSet(deserializedRequest.partitionStates()));
+            assertEquals(liveBrokers, deserializedRequest.liveBrokers());
+            assertEquals(1, request.controllerId());
+            assertEquals(2, request.controllerEpoch());
+            assertEquals(3, request.brokerEpoch());
+        }
+    }
+
+    @Test
+    public void testTopicPartitionGroupingSizeReduction() {
+        Set tps = TestUtils.generateRandomTopicPartitions(10, 10);
+        List partitionStates = new ArrayList<>();
+        for (TopicPartition tp : tps) {
+            partitionStates.add(new UpdateMetadataPartitionState()
+                .setTopicName(tp.topic())
+                .setPartitionIndex(tp.partition()));
+        }
+        UpdateMetadataRequest.Builder builder = new UpdateMetadataRequest.Builder((short) 5, 0, 0, 0, 0,
+                partitionStates, Collections.emptyList());
+
+        assertTrue(MessageTestUtil.messageSize(builder.build((short) 5).data(), (short) 5) <
+            MessageTestUtil.messageSize(builder.build((short) 4).data(), (short) 4));
+    }
+
+    private  Set iterableToSet(Iterable iterable) {
+        return StreamSupport.stream(iterable.spliterator(), false).collect(Collectors.toSet());
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/security/auth/DefaultKafkaPrincipalBuilderTest.java b/clients/src/test/java/org/apache/kafka/common/security/auth/DefaultKafkaPrincipalBuilderTest.java
index dd5087a84b387..dc1d6225e66d9 100644
--- a/clients/src/test/java/org/apache/kafka/common/security/auth/DefaultKafkaPrincipalBuilderTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/security/auth/DefaultKafkaPrincipalBuilderTest.java
@@ -30,8 +30,6 @@
 import javax.security.sasl.SaslServer;
 import java.net.InetAddress;
 import java.security.Principal;
-import java.util.Arrays;
-import java.util.List;
 
 import static org.junit.Assert.assertEquals;
 import static org.mockito.ArgumentMatchers.any;
@@ -143,7 +141,7 @@ public void testPrincipalWithSslPrincipalMapper() throws Exception {
                                         .thenReturn(new X500Principal("CN=duke, OU=JavaSoft, O=Sun Microsystems"))
                                         .thenReturn(new X500Principal("OU=JavaSoft, O=Sun Microsystems, C=US"));
 
-        List rules = Arrays.asList(
+        String rules = String.join(", ",
             "RULE:^CN=(.*),OU=ServiceUsers.*$/$1/L",
             "RULE:^CN=(.*),OU=(.*),O=(.*),L=(.*),ST=(.*),C=(.*)$/$1@$2/L",
             "RULE:^.*[Cc][Nn]=([a-zA-Z0-9.]*).*$/$1/U",
diff --git a/clients/src/test/java/org/apache/kafka/common/security/authenticator/ClientAuthenticationFailureTest.java b/clients/src/test/java/org/apache/kafka/common/security/authenticator/ClientAuthenticationFailureTest.java
index 61606ab8e898f..0a93c6afebc9f 100644
--- a/clients/src/test/java/org/apache/kafka/common/security/authenticator/ClientAuthenticationFailureTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/security/authenticator/ClientAuthenticationFailureTest.java
@@ -17,13 +17,15 @@
 package org.apache.kafka.common.security.authenticator;
 
 import org.apache.kafka.clients.CommonClientConfigs;
-import org.apache.kafka.clients.admin.AdminClient;
-import org.apache.kafka.clients.admin.DescribeTopicsResult;
+import org.apache.kafka.clients.admin.Admin;
+import org.apache.kafka.clients.admin.TopicDescription;
 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.clients.producer.RecordMetadata;
+import org.apache.kafka.common.KafkaFuture;
 import org.apache.kafka.common.config.SaslConfigs;
 import org.apache.kafka.common.config.internals.BrokerSecurityConfigs;
 import org.apache.kafka.common.errors.SaslAuthenticationException;
@@ -35,6 +37,7 @@
 import org.apache.kafka.common.serialization.StringDeserializer;
 import org.apache.kafka.common.serialization.StringSerializer;
 import org.apache.kafka.common.utils.MockTime;
+import org.apache.kafka.test.TestUtils;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
@@ -44,9 +47,9 @@
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.concurrent.Future;
 
-import static org.junit.Assert.assertTrue;
-import static org.junit.Assert.fail;
+import static org.junit.Assert.assertThrows;
 
 public class ClientAuthenticationFailureTest {
     private static MockTime time = new MockTime(50);
@@ -88,13 +91,10 @@ public void testConsumerWithInvalidCredentials() {
         StringDeserializer deserializer = new StringDeserializer();
 
         try (KafkaConsumer consumer = new KafkaConsumer<>(props, deserializer, deserializer)) {
-            consumer.subscribe(Arrays.asList(topic));
-            consumer.poll(Duration.ofSeconds(10));
-            fail("Expected an authentication error!");
-        } catch (SaslAuthenticationException e) {
-            // OK
-        } catch (Exception e) {
-            throw new AssertionError("Expected only an authentication error, but another error occurred.", e);
+            assertThrows(SaslAuthenticationException.class, () -> {
+                consumer.subscribe(Collections.singleton(topic));
+                consumer.poll(Duration.ofSeconds(10));
+            });
         }
     }
 
@@ -106,11 +106,8 @@ public void testProducerWithInvalidCredentials() {
 
         try (KafkaProducer producer = new KafkaProducer<>(props, serializer, serializer)) {
             ProducerRecord record = new ProducerRecord<>(topic, "message");
-            producer.send(record).get();
-            fail("Expected an authentication error!");
-        } catch (Exception e) {
-            assertTrue("Expected SaslAuthenticationException, got " + e.getCause().getClass(),
-                    e.getCause() instanceof SaslAuthenticationException);
+            Future future = producer.send(record);
+            TestUtils.assertFutureThrows(future, SaslAuthenticationException.class);
         }
     }
 
@@ -118,13 +115,9 @@ public void testProducerWithInvalidCredentials() {
     public void testAdminClientWithInvalidCredentials() {
         Map props = new HashMap<>(saslClientConfigs);
         props.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:" + server.port());
-        try (AdminClient client = AdminClient.create(props)) {
-            DescribeTopicsResult result = client.describeTopics(Collections.singleton("test"));
-            result.all().get();
-            fail("Expected an authentication error!");
-        } catch (Exception e) {
-            assertTrue("Expected SaslAuthenticationException, got " + e.getCause().getClass(),
-                    e.getCause() instanceof SaslAuthenticationException);
+        try (Admin client = Admin.create(props)) {
+            KafkaFuture> future = client.describeTopics(Collections.singleton("test")).all();
+            TestUtils.assertFutureThrows(future, SaslAuthenticationException.class);
         }
     }
 
@@ -137,10 +130,7 @@ public void testTransactionalProducerWithInvalidCredentials() {
         StringSerializer serializer = new StringSerializer();
 
         try (KafkaProducer producer = new KafkaProducer<>(props, serializer, serializer)) {
-            producer.initTransactions();
-            fail("Expected an authentication error!");
-        } catch (SaslAuthenticationException e) {
-            // expected exception
+            assertThrows(SaslAuthenticationException.class, producer::initTransactions);
         }
     }
 
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 773fffbb287f5..0a128c9449476 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
@@ -31,6 +31,7 @@
 import org.apache.kafka.common.security.TestSecurityConfig;
 import org.apache.kafka.common.security.auth.SecurityProtocol;
 import org.apache.kafka.common.utils.MockTime;
+import org.apache.kafka.test.TestSslUtils;
 import org.apache.kafka.test.TestUtils;
 import org.junit.After;
 import org.junit.Assert;
@@ -81,8 +82,8 @@ public static Collection data() {
     @Before
     public void setup() throws Exception {
         LoginManager.closeAll();
-        serverCertStores = new CertStores(true, "localhost");
-        clientCertStores = new CertStores(false, "localhost");
+        serverCertStores = new CertStores(true, "localhost", TestSslUtils.SSLProvider.DEFAULT);
+        clientCertStores = new CertStores(false, "localhost", TestSslUtils.SSLProvider.DEFAULT);
         saslServerConfigs = serverCertStores.getTrustingConfig(clientCertStores);
         saslClientConfigs = clientCertStores.getTrustingConfig(serverCertStores);
         credentialCache = new CredentialCache();
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 fc4eb031b35d6..9bc574fe15e96 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
@@ -25,15 +25,19 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Base64;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
-import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Random;
 import java.util.Base64.Encoder;
+import java.util.Set;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Function;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
 
 import javax.security.auth.Subject;
 import javax.security.auth.callback.Callback;
@@ -45,14 +49,21 @@
 import javax.security.auth.login.AppConfigurationEntry;
 import javax.security.auth.login.LoginContext;
 import javax.security.auth.login.LoginException;
+import javax.security.sasl.SaslClient;
 import javax.security.sasl.SaslException;
 
 import org.apache.kafka.clients.NetworkClient;
 import org.apache.kafka.common.KafkaException;
+import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.config.SaslConfigs;
 import org.apache.kafka.common.config.internals.BrokerSecurityConfigs;
 import org.apache.kafka.common.config.types.Password;
 import org.apache.kafka.common.errors.SaslAuthenticationException;
+import org.apache.kafka.common.message.ApiVersionsRequestData;
+import org.apache.kafka.common.message.ApiVersionsResponseData;
+import org.apache.kafka.common.message.ApiVersionsResponseData.ApiVersionsResponseKey;
+import org.apache.kafka.common.message.ApiVersionsResponseData.ApiVersionsResponseKeyCollection;
+import org.apache.kafka.common.message.RequestHeaderData;
 import org.apache.kafka.common.message.SaslAuthenticateRequestData;
 import org.apache.kafka.common.message.SaslHandshakeRequestData;
 import org.apache.kafka.common.network.CertStores;
@@ -70,13 +81,14 @@
 import org.apache.kafka.common.network.TransportLayer;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.common.protocol.types.SchemaException;
+import org.apache.kafka.common.requests.ListOffsetResponse;
 import org.apache.kafka.common.security.auth.Login;
 import org.apache.kafka.common.security.auth.SecurityProtocol;
 import org.apache.kafka.common.requests.AbstractRequest;
 import org.apache.kafka.common.requests.AbstractResponse;
 import org.apache.kafka.common.requests.ApiVersionsRequest;
 import org.apache.kafka.common.requests.ApiVersionsResponse;
-import org.apache.kafka.common.requests.ApiVersionsResponse.ApiVersion;
 import org.apache.kafka.common.requests.MetadataRequest;
 import org.apache.kafka.common.requests.RequestHeader;
 import org.apache.kafka.common.requests.ResponseHeader;
@@ -109,19 +121,31 @@
 
 import org.apache.kafka.common.utils.Time;
 import org.apache.kafka.common.utils.Utils;
+import org.apache.kafka.test.TestSslUtils.SSLProvider;
 import org.junit.After;
+import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
 
+import static org.apache.kafka.common.protocol.ApiKeys.LIST_OFFSETS;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
 /**
  * Tests for the Sasl authenticator. These use a test harness that runs a simple socket server that echos back responses.
  */
+@RunWith(Parameterized.class)
 public class SaslAuthenticatorTest {
+    private final SSLProvider provider;
+
+    public SaslAuthenticatorTest(SSLProvider provider) {
+        this.provider = provider;
+    }
 
     private static final long CONNECTIONS_MAX_REAUTH_MS_VALUE = 100L;
     private static final int BUFFER_SIZE = 4 * 1024;
@@ -141,8 +165,8 @@ public class SaslAuthenticatorTest {
     public void setup() throws Exception {
         LoginManager.closeAll();
         time = Time.SYSTEM;
-        serverCertStores = new CertStores(true, "localhost");
-        clientCertStores = new CertStores(false, "localhost");
+        serverCertStores = new CertStores(true, "localhost", provider);
+        clientCertStores = new CertStores(false, "localhost", provider);
         saslServerConfigs = serverCertStores.getTrustingConfig(clientCertStores);
         saslClientConfigs = clientCertStores.getTrustingConfig(serverCertStores);
         credentialCache = new CredentialCache();
@@ -385,7 +409,7 @@ public void testMultipleServerMechanisms() throws Exception {
             selector.connect(node3, new InetSocketAddress("127.0.0.1", server.port()), BUFFER_SIZE, BUFFER_SIZE);
             NetworkTestUtils.checkClientConnection(selector, node3, 100, 10);
             server.verifyAuthenticationMetrics(3, 0);
-            
+
             /*
              * Now re-authenticate the connections. First we have to sleep long enough so
              * that the next write will cause re-authentication, which we expect to succeed.
@@ -398,7 +422,7 @@ public void testMultipleServerMechanisms() throws Exception {
 
             NetworkTestUtils.checkClientConnection(selector3, node3, 100, 10);
             server.verifyReauthenticationMetrics(2, 0);
-            
+
         } finally {
             if (selector2 != null)
                 selector2.close();
@@ -680,13 +704,25 @@ public void testApiVersionsRequestWithUnsupportedVersion() throws Exception {
         // Send ApiVersionsRequest with unsupported version and validate error response.
         String node = "1";
         createClientConnection(SecurityProtocol.PLAINTEXT, node);
-        RequestHeader header = new RequestHeader(ApiKeys.API_VERSIONS, Short.MAX_VALUE, "someclient", 1);
+
+        RequestHeader header = new RequestHeader(new RequestHeaderData().
+                setRequestApiKey(ApiKeys.API_VERSIONS.id).
+                setRequestApiVersion(Short.MAX_VALUE).
+                setClientId("someclient").
+                setCorrelationId(1),
+                (short) 2);
         ApiVersionsRequest request = new ApiVersionsRequest.Builder().build();
         selector.send(request.toSend(node, header));
         ByteBuffer responseBuffer = waitForResponse();
-        ResponseHeader.parse(responseBuffer);
+        ResponseHeader.parse(responseBuffer, ApiKeys.API_VERSIONS.responseHeaderVersion((short) 0));
         ApiVersionsResponse response = ApiVersionsResponse.parse(responseBuffer, (short) 0);
-        assertEquals(Errors.UNSUPPORTED_VERSION, response.error());
+        assertEquals(Errors.UNSUPPORTED_VERSION.code(), response.data.errorCode());
+
+        ApiVersionsResponseKey apiVersion = response.data.apiKeys().find(ApiKeys.API_VERSIONS.id);
+        assertNotNull(apiVersion);
+        assertEquals(ApiKeys.API_VERSIONS.id, apiVersion.apiKey());
+        assertEquals(ApiKeys.API_VERSIONS.oldestVersion(), apiVersion.minVersion());
+        assertEquals(ApiKeys.API_VERSIONS.latestVersion(), apiVersion.maxVersion());
 
         // Send ApiVersionsRequest with a supported version. This should succeed.
         sendVersionRequestReceiveResponse(node);
@@ -696,6 +732,68 @@ public void testApiVersionsRequestWithUnsupportedVersion() throws Exception {
         authenticateUsingSaslPlainAndCheckConnection(node, handshakeVersion > 0);
     }
 
+    /**
+     * Tests that invalid ApiVersionRequest is handled by the server correctly and
+     * returns an INVALID_REQUEST error.
+     */
+    @Test
+    public void testInvalidApiVersionsRequest() throws Exception {
+        short handshakeVersion = ApiKeys.SASL_HANDSHAKE.latestVersion();
+        SecurityProtocol securityProtocol = SecurityProtocol.SASL_PLAINTEXT;
+        configureMechanisms("PLAIN", Arrays.asList("PLAIN"));
+        server = createEchoServer(securityProtocol);
+
+        // Send ApiVersionsRequest with invalid version and validate error response.
+        String node = "1";
+        short version = ApiKeys.API_VERSIONS.latestVersion();
+        createClientConnection(SecurityProtocol.PLAINTEXT, node);
+        RequestHeader header = new RequestHeader(ApiKeys.API_VERSIONS, version, "someclient", 1);
+        ApiVersionsRequest request = new ApiVersionsRequest(new ApiVersionsRequestData().
+                setClientSoftwareName("  ").
+                setClientSoftwareVersion("   "), version);
+        selector.send(request.toSend(node, header));
+        ByteBuffer responseBuffer = waitForResponse();
+        ResponseHeader.parse(responseBuffer, ApiKeys.API_VERSIONS.responseHeaderVersion(version));
+        ApiVersionsResponse response =
+            ApiVersionsResponse.parse(responseBuffer, version);
+        assertEquals(Errors.INVALID_REQUEST.code(), response.data.errorCode());
+
+        // Send ApiVersionsRequest with a supported version. This should succeed.
+        sendVersionRequestReceiveResponse(node);
+
+        // Test that client can authenticate successfully
+        sendHandshakeRequestReceiveResponse(node, handshakeVersion);
+        authenticateUsingSaslPlainAndCheckConnection(node, handshakeVersion > 0);
+    }
+
+    /**
+     * Tests that valid ApiVersionRequest is handled by the server correctly and
+     * returns an NONE error.
+     */
+    @Test
+    public void testValidApiVersionsRequest() throws Exception {
+        short handshakeVersion = ApiKeys.SASL_HANDSHAKE.latestVersion();
+        SecurityProtocol securityProtocol = SecurityProtocol.SASL_PLAINTEXT;
+        configureMechanisms("PLAIN", Arrays.asList("PLAIN"));
+        server = createEchoServer(securityProtocol);
+
+        // Send ApiVersionsRequest with valid version and validate error response.
+        String node = "1";
+        short version = ApiKeys.API_VERSIONS.latestVersion();
+        createClientConnection(SecurityProtocol.PLAINTEXT, node);
+        RequestHeader header = new RequestHeader(ApiKeys.API_VERSIONS, version, "someclient", 1);
+        ApiVersionsRequest request = new ApiVersionsRequest.Builder().build(version);
+        selector.send(request.toSend(node, header));
+        ByteBuffer responseBuffer = waitForResponse();
+        ResponseHeader.parse(responseBuffer, ApiKeys.API_VERSIONS.responseHeaderVersion(version));
+        ApiVersionsResponse response = ApiVersionsResponse.parse(responseBuffer, version);
+        assertEquals(Errors.NONE.code(), response.data.errorCode());
+
+        // Test that client can authenticate successfully
+        sendHandshakeRequestReceiveResponse(node, handshakeVersion);
+        authenticateUsingSaslPlainAndCheckConnection(node, handshakeVersion > 0);
+    }
+
     /**
      * Tests that unsupported version of SASL handshake request returns error
      * response and fails authentication. This test is similar to
@@ -714,7 +812,7 @@ public void testSaslHandshakeRequestWithUnsupportedVersion() throws Exception {
         createClientConnection(SecurityProtocol.PLAINTEXT, node1);
         SaslHandshakeRequest request = buildSaslHandshakeRequest("PLAIN", ApiKeys.SASL_HANDSHAKE.latestVersion());
         RequestHeader header = new RequestHeader(ApiKeys.SASL_HANDSHAKE, Short.MAX_VALUE, "someclient", 2);
-        
+
         selector.send(request.toSend(node1, header));
         // This test uses a non-SASL PLAINTEXT client in order to do manual handshake.
         // So the channel is in READY state.
@@ -1395,7 +1493,7 @@ public void testValidSaslOauthBearerMechanism() throws Exception {
         server = createEchoServer(securityProtocol);
         createAndCheckClientConnection(securityProtocol, node);
     }
-    
+
     /**
      * Re-authentication must fail if principal changes
      */
@@ -1428,6 +1526,47 @@ public void testCannotReauthenticateWithDifferentPrincipal() throws Exception {
             server.verifyReauthenticationMetrics(0, 1);
         }
     }
+
+    @Test
+    public void testCorrelationId() {
+        SaslClientAuthenticator authenticator = new SaslClientAuthenticator(
+              Collections.emptyMap(),
+              null,
+              "node",
+              null,
+              null,
+              null,
+              "plain",
+              false,
+              null,
+              null
+        ) {
+            @Override
+            SaslClient createSaslClient() {
+                return null;
+            }
+        };
+        int count = (SaslClientAuthenticator.MAX_RESERVED_CORRELATION_ID - SaslClientAuthenticator.MIN_RESERVED_CORRELATION_ID) * 2;
+        Set ids = IntStream.range(0, count)
+            .mapToObj(i -> authenticator.nextCorrelationId())
+            .collect(Collectors.toSet());
+        assertEquals(SaslClientAuthenticator.MAX_RESERVED_CORRELATION_ID - SaslClientAuthenticator.MIN_RESERVED_CORRELATION_ID + 1, ids.size());
+        ids.forEach(id -> {
+            assertTrue(id >= SaslClientAuthenticator.MIN_RESERVED_CORRELATION_ID);
+            assertTrue(SaslClientAuthenticator.isReserved(id));
+        });
+    }
+
+    @Test
+    public void testConvertListOffsetResponseToSaslHandshakeResponse() {
+        ListOffsetResponse response = new ListOffsetResponse(0, Collections.singletonMap(new TopicPartition("topic", 0),
+            new ListOffsetResponse.PartitionData(Errors.NONE, 0, 0, Optional.empty())));
+        ByteBuffer buffer = response.serialize(ApiKeys.LIST_OFFSETS, LIST_OFFSETS.latestVersion(), 0);
+        final RequestHeader header0 = new RequestHeader(LIST_OFFSETS, LIST_OFFSETS.latestVersion(), "id", SaslClientAuthenticator.MIN_RESERVED_CORRELATION_ID);
+        Assert.assertThrows(SchemaException.class, () -> NetworkClient.parseResponse(buffer.duplicate(), header0));
+        final RequestHeader header1 = new RequestHeader(LIST_OFFSETS, LIST_OFFSETS.latestVersion(), "id", 1);
+        Assert.assertThrows(IllegalStateException.class, () -> NetworkClient.parseResponse(buffer.duplicate(), header1));
+    }
     
     /**
      * Re-authentication must fail if mechanism changes
@@ -1513,7 +1652,7 @@ public void testCannotReauthenticateAgainFasterThanOneSecond() throws Exception
                     e.getMessage().matches(
                             ".*\\<\\[" + expectedResponseTextRegex + "]>.*\\<\\[" + receivedResponseTextRegex + "]>"));
             server.verifyReauthenticationMetrics(1, 0); // unchanged
-        } finally { 
+        } finally {
             selector.close();
             selector = null;
         }
@@ -1555,7 +1694,7 @@ public void testRepeatedValidSaslPlainOverSsl() throws Exception {
         }
         server.verifyReauthenticationMetrics(desiredNumReauthentications, 0);
     }
-    
+
     /**
      * Tests OAUTHBEARER client channels without tokens for the server.
      */
@@ -1674,15 +1813,25 @@ protected SaslServerAuthenticator buildServerAuthenticator(Map config
 
                     @Override
                     protected ApiVersionsResponse apiVersionsResponse() {
-                        List apiVersions = new ArrayList<>(ApiVersionsResponse.defaultApiVersionsResponse().apiVersions());
-                        for (Iterator it = apiVersions.iterator(); it.hasNext(); ) {
-                            ApiVersion apiVersion = it.next();
-                            if (apiVersion.apiKey == ApiKeys.SASL_AUTHENTICATE.id) {
-                                it.remove();
-                                break;
+                        ApiVersionsResponse defaultApiVersionResponse = ApiVersionsResponse.DEFAULT_API_VERSIONS_RESPONSE;
+                        ApiVersionsResponseKeyCollection apiVersions = new ApiVersionsResponseKeyCollection();
+                        for (ApiVersionsResponseKey apiVersion : defaultApiVersionResponse.data.apiKeys()) {
+                            if (apiVersion.apiKey() != ApiKeys.SASL_AUTHENTICATE.id) {
+                                // ApiVersionsResponseKey can NOT be reused in second ApiVersionsResponseKeyCollection
+                                // due to the internal pointers it contains.
+                                apiVersions.add(new ApiVersionsResponseKey()
+                                    .setApiKey(apiVersion.apiKey())
+                                    .setMinVersion(apiVersion.minVersion())
+                                    .setMaxVersion(apiVersion.maxVersion())
+                                );
                             }
+
                         }
-                        return new ApiVersionsResponse(0, Errors.NONE, apiVersions);
+                        ApiVersionsResponseData data = new ApiVersionsResponseData()
+                            .setErrorCode(Errors.NONE.code())
+                            .setThrottleTimeMs(0)
+                            .setApiKeys(apiVersions);
+                        return new ApiVersionsResponse(data);
                     }
 
                     @Override
@@ -1784,10 +1933,10 @@ private void testUnauthenticatedApiVersionsRequest(SecurityProtocol securityProt
 
         // Send ApiVersionsRequest and check response
         ApiVersionsResponse versionsResponse = sendVersionRequestReceiveResponse(node);
-        assertEquals(ApiKeys.SASL_HANDSHAKE.oldestVersion(), versionsResponse.apiVersion(ApiKeys.SASL_HANDSHAKE.id).minVersion);
-        assertEquals(ApiKeys.SASL_HANDSHAKE.latestVersion(), versionsResponse.apiVersion(ApiKeys.SASL_HANDSHAKE.id).maxVersion);
-        assertEquals(ApiKeys.SASL_AUTHENTICATE.oldestVersion(), versionsResponse.apiVersion(ApiKeys.SASL_AUTHENTICATE.id).minVersion);
-        assertEquals(ApiKeys.SASL_AUTHENTICATE.latestVersion(), versionsResponse.apiVersion(ApiKeys.SASL_AUTHENTICATE.id).maxVersion);
+        assertEquals(ApiKeys.SASL_HANDSHAKE.oldestVersion(), versionsResponse.apiVersion(ApiKeys.SASL_HANDSHAKE.id).minVersion());
+        assertEquals(ApiKeys.SASL_HANDSHAKE.latestVersion(), versionsResponse.apiVersion(ApiKeys.SASL_HANDSHAKE.id).maxVersion());
+        assertEquals(ApiKeys.SASL_AUTHENTICATE.oldestVersion(), versionsResponse.apiVersion(ApiKeys.SASL_AUTHENTICATE.id).minVersion());
+        assertEquals(ApiKeys.SASL_AUTHENTICATE.latestVersion(), versionsResponse.apiVersion(ApiKeys.SASL_AUTHENTICATE.id).maxVersion());
 
         // Send SaslHandshakeRequest and check response
         SaslHandshakeResponse handshakeResponse = sendHandshakeRequestReceiveResponse(node, saslHandshakeVersion);
@@ -1863,7 +2012,7 @@ private void createClientConnection(SecurityProtocol securityProtocol, String no
         InetSocketAddress addr = new InetSocketAddress("localhost", server.port());
         selector.connect(node, addr, BUFFER_SIZE, BUFFER_SIZE);
     }
-    
+
     private void checkClientConnection(String node) throws Exception {
         NetworkTestUtils.checkClientConnection(selector, node, 100, 10);
     }
@@ -1945,7 +2094,7 @@ private SaslHandshakeResponse sendHandshakeRequestReceiveResponse(String node, s
     private ApiVersionsResponse sendVersionRequestReceiveResponse(String node) throws Exception {
         ApiVersionsRequest handshakeRequest = createApiVersionsRequestV0();
         ApiVersionsResponse response =  (ApiVersionsResponse) sendKafkaRequestReceiveResponse(node, ApiKeys.API_VERSIONS, handshakeRequest);
-        assertEquals(Errors.NONE, response.error());
+        assertEquals(Errors.NONE.code(), response.data.errorCode());
         return response;
     }
 
@@ -2248,4 +2397,12 @@ protected SaslHandshakeRequest createSaslHandshakeRequest(short version) {
                 };
         }
     }
+
+    @Parameterized.Parameters(name = "SSLProvider={0}")
+    public static Collection data() {
+        Collection p = new ArrayList<>();
+        p.add(new Object[]{SSLProvider.DEFAULT});
+        p.add(new Object[]{SSLProvider.OPENSSL});
+        return p;
+    }
 }
diff --git a/clients/src/test/java/org/apache/kafka/common/security/kerberos/KerberosNameTest.java b/clients/src/test/java/org/apache/kafka/common/security/kerberos/KerberosNameTest.java
index 687f9370a6c9c..9dd44a14fb670 100644
--- a/clients/src/test/java/org/apache/kafka/common/security/kerberos/KerberosNameTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/security/kerberos/KerberosNameTest.java
@@ -86,6 +86,35 @@ public void testToLowerCase() throws Exception {
         assertEquals("user", shortNamer.shortName(name));
     }
 
+    @Test
+    public void testToUpperCase() throws Exception {
+        List rules = Arrays.asList(
+            "RULE:[1:$1]/U",
+            "RULE:[2:$1](Test.*)s/ABC///U",
+            "RULE:[2:$1](ABC.*)s/ABC/XYZ/g/U",
+            "RULE:[2:$1](App\\..*)s/App\\.(.*)/$1/g/U",
+            "RULE:[2:$1]/U",
+            "DEFAULT"
+        );
+
+        KerberosShortNamer shortNamer = KerberosShortNamer.fromUnparsedRules("REALM.COM", rules);
+
+        KerberosName name = KerberosName.parse("User@REALM.COM");
+        assertEquals("USER", shortNamer.shortName(name));
+
+        name = KerberosName.parse("TestABC/host@FOO.COM");
+        assertEquals("TEST", shortNamer.shortName(name));
+
+        name = KerberosName.parse("ABC_User_ABC/host@FOO.COM");
+        assertEquals("XYZ_USER_XYZ", shortNamer.shortName(name));
+
+        name = KerberosName.parse("App.SERVICE-name/example.com@REALM.COM");
+        assertEquals("SERVICE-NAME", shortNamer.shortName(name));
+
+        name = KerberosName.parse("User/root@REALM.COM");
+        assertEquals("USER", shortNamer.shortName(name));
+    }
+
     @Test
     public void testInvalidRules() {
         testInvalidRule(Arrays.asList("default"));
@@ -94,6 +123,9 @@ public void testInvalidRules() {
         testInvalidRule(Arrays.asList("DEFAULT/g"));
 
         testInvalidRule(Arrays.asList("rule:[1:$1]"));
+        testInvalidRule(Arrays.asList("rule:[1:$1]/L/U"));
+        testInvalidRule(Arrays.asList("rule:[1:$1]/U/L"));
+        testInvalidRule(Arrays.asList("rule:[1:$1]/LU"));
         testInvalidRule(Arrays.asList("RULE:[1:$1/L"));
         testInvalidRule(Arrays.asList("RULE:[1:$1]/l"));
         testInvalidRule(Arrays.asList("RULE:[2:$1](ABC.*)s/ABC/XYZ/L/g"));
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 bfe34c98382be..b4166a968ff8d 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
@@ -19,19 +19,25 @@
 import java.io.File;
 import java.nio.file.Files;
 import java.security.KeyStore;
+import java.util.Arrays;
 import java.util.Map;
 
 import javax.net.ssl.SSLContext;
 import javax.net.ssl.SSLEngine;
-import javax.net.ssl.SSLHandshakeException;
 
+import org.apache.kafka.common.KafkaException;
+import org.apache.kafka.common.config.ConfigException;
+import org.apache.kafka.common.config.SecurityConfig;
 import org.apache.kafka.common.config.SslConfigs;
 import org.apache.kafka.common.config.types.Password;
+import org.apache.kafka.common.security.ssl.mock.TestKeyManagerFactory;
+import org.apache.kafka.common.security.ssl.mock.TestProviderCreator;
+import org.apache.kafka.common.security.ssl.mock.TestTrustManagerFactory;
+import org.apache.kafka.common.utils.Utils;
 import org.apache.kafka.test.TestSslUtils;
 import org.apache.kafka.common.network.Mode;
 import org.junit.Test;
 
-import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotEquals;
@@ -40,26 +46,61 @@
 import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
+import java.security.Security;
 
-/**
- * A set of tests for the selector over ssl. These use a test harness that runs a simple socket server that echos back responses.
- */
 public class SslFactoryTest {
-
     @Test
     public void testSslFactoryConfiguration() throws Exception {
         File trustStoreFile = File.createTempFile("truststore", ".jks");
-        Map serverSslConfig = TestSslUtils.createSslConfig(false, true, Mode.SERVER, trustStoreFile, "server");
+        Map serverSslConfig =
+                TestSslUtils.createSslConfig(false, true, Mode.SERVER, trustStoreFile, "server");
         SslFactory sslFactory = new SslFactory(Mode.SERVER);
         sslFactory.configure(serverSslConfig);
         //host and port are hints
         SSLEngine engine = sslFactory.createSslEngine("localhost", 0);
         assertNotNull(engine);
-        String[] expectedProtocols = {"TLSv1.2"};
-        assertArrayEquals(expectedProtocols, engine.getEnabledProtocols());
+        assertEquals(Utils.mkSet("TLSv1.2"), Utils.mkSet(engine.getEnabledProtocols()));
         assertEquals(false, engine.getUseClientMode());
     }
 
+    @Test
+    public void testSslFactoryWithCustomKeyManagerConfiguration() {
+        TestProviderCreator testProviderCreator = new TestProviderCreator();
+        Map serverSslConfig = TestSslUtils.createSslConfig(
+                TestKeyManagerFactory.ALGORITHM,
+                TestTrustManagerFactory.ALGORITHM
+        );
+        serverSslConfig.put(SecurityConfig.SECURITY_PROVIDERS_CONFIG, testProviderCreator.getClass().getName());
+        SslFactory sslFactory = new SslFactory(Mode.SERVER);
+        sslFactory.configure(serverSslConfig);
+        assertNotNull("SslEngineBuilder not created", sslFactory.sslEngineBuilder());
+        Security.removeProvider(testProviderCreator.getProvider().getName());
+    }
+
+    @Test(expected = KafkaException.class)
+    public void testSslFactoryWithoutProviderClassConfiguration() {
+        // An exception is thrown as the algorithm is not registered through a provider
+        Map serverSslConfig = TestSslUtils.createSslConfig(
+                TestKeyManagerFactory.ALGORITHM,
+                TestTrustManagerFactory.ALGORITHM
+        );
+        SslFactory sslFactory = new SslFactory(Mode.SERVER);
+        sslFactory.configure(serverSslConfig);
+    }
+
+    @Test(expected = KafkaException.class)
+    public void testSslFactoryWithIncorrectProviderClassConfiguration() {
+        // An exception is thrown as the algorithm is not registered through a provider
+        Map serverSslConfig = TestSslUtils.createSslConfig(
+                TestKeyManagerFactory.ALGORITHM,
+                TestTrustManagerFactory.ALGORITHM
+        );
+        serverSslConfig.put(SecurityConfig.SECURITY_PROVIDERS_CONFIG,
+                "com.fake.ProviderClass1,com.fake.ProviderClass2");
+        SslFactory sslFactory = new SslFactory(Mode.SERVER);
+        sslFactory.configure(serverSslConfig);
+    }
+
     @Test
     public void testSslFactoryWithoutPasswordConfiguration() throws Exception {
         File trustStoreFile = File.createTempFile("truststore", ".jks");
@@ -77,7 +118,8 @@ public void testSslFactoryWithoutPasswordConfiguration() throws Exception {
     @Test
     public void testClientMode() throws Exception {
         File trustStoreFile = File.createTempFile("truststore", ".jks");
-        Map clientSslConfig = TestSslUtils.createSslConfig(false, true, Mode.CLIENT, trustStoreFile, "client");
+        Map clientSslConfig =
+                TestSslUtils.createSslConfig(false, true, Mode.CLIENT, trustStoreFile, "client");
         SslFactory sslFactory = new SslFactory(Mode.CLIENT);
         sslFactory.configure(clientSslConfig);
         //host and port are hints
@@ -88,92 +130,172 @@ public void testClientMode() throws Exception {
     @Test
     public void testReconfiguration() throws Exception {
         File trustStoreFile = File.createTempFile("truststore", ".jks");
-        Map sslConfig = TestSslUtils.createSslConfig(false, true, Mode.SERVER, trustStoreFile, "server");
+        Map sslConfig = TestSslUtils.
+                createSslConfig(false, true, Mode.SERVER, trustStoreFile, "server");
         SslFactory sslFactory = new SslFactory(Mode.SERVER);
         sslFactory.configure(sslConfig);
-        SSLContext sslContext = sslFactory.sslContext();
-        assertNotNull("SSL context not created", sslContext);
-        assertSame("SSL context recreated unnecessarily", sslContext, sslFactory.sslContext());
-        assertFalse(sslFactory.createSslEngine("localhost", 0).getUseClientMode());
+        SslEngineBuilder sslEngineBuilder = sslFactory.sslEngineBuilder();
+        assertNotNull("SslEngineBuilder not created", sslEngineBuilder);
 
-        // Verify that context is not recreated on reconfigure() if config and file are not changed
+        // Verify that SslEngineBuilder is not recreated on reconfigure() if config and
+        // file are not changed
         sslFactory.reconfigure(sslConfig);
-        assertSame("SSL context recreated unnecessarily", sslContext, sslFactory.sslContext());
+        assertSame("SslEngineBuilder recreated unnecessarily",
+                sslEngineBuilder, sslFactory.sslEngineBuilder());
 
-        // Verify that context is recreated on reconfigure() if config is changed
+        // Verify that the SslEngineBuilder is recreated on reconfigure() if config is changed
         trustStoreFile = File.createTempFile("truststore", ".jks");
         sslConfig = TestSslUtils.createSslConfig(false, true, Mode.SERVER, trustStoreFile, "server");
         sslFactory.reconfigure(sslConfig);
-        assertNotSame("SSL context not recreated", sslContext, sslFactory.sslContext());
-        sslContext = sslFactory.sslContext();
+        assertNotSame("SslEngineBuilder not recreated",
+                sslEngineBuilder, sslFactory.sslEngineBuilder());
+        sslEngineBuilder = sslFactory.sslEngineBuilder();
 
-        // Verify that context is recreated on reconfigure() if config is not changed, but truststore file was modified
+        // Verify that builder is recreated on reconfigure() if config is not changed, but truststore file was modified
         trustStoreFile.setLastModified(System.currentTimeMillis() + 10000);
         sslFactory.reconfigure(sslConfig);
-        assertNotSame("SSL context not recreated", sslContext, sslFactory.sslContext());
-        sslContext = sslFactory.sslContext();
+        assertNotSame("SslEngineBuilder not recreated",
+                sslEngineBuilder, sslFactory.sslEngineBuilder());
+        sslEngineBuilder = sslFactory.sslEngineBuilder();
 
-        // Verify that context is recreated on reconfigure() if config is not changed, but keystore file was modified
+        // Verify that builder is recreated on reconfigure() if config is not changed, but keystore file was modified
         File keyStoreFile = new File((String) sslConfig.get(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG));
         keyStoreFile.setLastModified(System.currentTimeMillis() + 10000);
         sslFactory.reconfigure(sslConfig);
-        assertNotSame("SSL context not recreated", sslContext, sslFactory.sslContext());
-        sslContext = sslFactory.sslContext();
+        assertNotSame("SslEngineBuilder not recreated",
+                sslEngineBuilder, sslFactory.sslEngineBuilder());
+        sslEngineBuilder = sslFactory.sslEngineBuilder();
+
+        // Verify that builder is recreated after validation on reconfigure() if config is not changed, but keystore file was modified
+        keyStoreFile.setLastModified(System.currentTimeMillis() + 15000);
+        sslFactory.validateReconfiguration(sslConfig);
+        sslFactory.reconfigure(sslConfig);
+        assertNotSame("SslEngineBuilder not recreated",
+                sslEngineBuilder, sslFactory.sslEngineBuilder());
+        sslEngineBuilder = sslFactory.sslEngineBuilder();
 
-        // Verify that the context is not recreated if modification time cannot be determined
+        // Verify that the builder is not recreated if modification time cannot be determined
         keyStoreFile.setLastModified(System.currentTimeMillis() + 20000);
         Files.delete(keyStoreFile.toPath());
         sslFactory.reconfigure(sslConfig);
-        assertSame("SSL context recreated unnecessarily", sslContext, sslFactory.sslContext());
+        assertSame("SslEngineBuilder recreated unnecessarily",
+                sslEngineBuilder, sslFactory.sslEngineBuilder());
     }
 
     @Test
-    public void testKeyStoreTrustStoreValidation() throws Exception {
+    public void testReconfigurationWithoutTruststore() throws Exception {
         File trustStoreFile = File.createTempFile("truststore", ".jks");
-        Map serverSslConfig = TestSslUtils.createSslConfig(false, true,
-                Mode.SERVER, trustStoreFile, "server");
+        Map sslConfig = TestSslUtils.
+            createSslConfig(false, true, Mode.SERVER, trustStoreFile, "server");
+        sslConfig.remove(SslConfigs.SSL_TRUSTSTORE_LOCATION_CONFIG);
+        sslConfig.remove(SslConfigs.SSL_TRUSTSTORE_PASSWORD_CONFIG);
+        sslConfig.remove(SslConfigs.SSL_TRUSTSTORE_TYPE_CONFIG);
         SslFactory sslFactory = new SslFactory(Mode.SERVER);
-        sslFactory.configure(serverSslConfig);
-        SSLContext sslContext = sslFactory.createSSLContext(sslKeyStore(serverSslConfig), null);
+        sslFactory.configure(sslConfig);
+        SSLContext sslContext = sslFactory.sslEngineBuilder().sslContext();
+        assertNotNull("SSL context not created", sslContext);
+        assertSame("SSL context recreated unnecessarily", sslContext,
+                sslFactory.sslEngineBuilder().sslContext());
+        assertFalse(sslFactory.createSslEngine("localhost", 0).getUseClientMode());
+
+        Map sslConfig2 = TestSslUtils.
+            createSslConfig(false, true, Mode.SERVER, trustStoreFile, "server");
+        try {
+            sslFactory.validateReconfiguration(sslConfig2);
+            fail("Truststore configured dynamically for listener without previous truststore");
+        } catch (ConfigException e) {
+            // Expected exception
+        }
+    }
+
+    @Test
+    public void testReconfigurationWithoutKeystore() throws Exception {
+        File trustStoreFile = File.createTempFile("truststore", ".jks");
+        Map sslConfig = TestSslUtils.
+                createSslConfig(false, true, Mode.SERVER, trustStoreFile, "server");
+        sslConfig.remove(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG);
+        sslConfig.remove(SslConfigs.SSL_KEYSTORE_PASSWORD_CONFIG);
+        sslConfig.remove(SslConfigs.SSL_KEYSTORE_TYPE_CONFIG);
+        SslFactory sslFactory = new SslFactory(Mode.SERVER);
+        sslFactory.configure(sslConfig);
+        SSLContext sslContext = sslFactory.sslEngineBuilder().sslContext();
         assertNotNull("SSL context not created", sslContext);
+        assertSame("SSL context recreated unnecessarily", sslContext,
+                sslFactory.sslEngineBuilder().sslContext());
+        assertFalse(sslFactory.createSslEngine("localhost", 0).getUseClientMode());
 
-        SSLContext sslContext2 = sslFactory.createSSLContext(null, sslTrustStore(serverSslConfig));
-        assertNotNull("SSL context not created", sslContext2);
+        File newTrustStoreFile = File.createTempFile("truststore", ".jks");
+        sslConfig = TestSslUtils.
+                createSslConfig(false, true, Mode.SERVER, newTrustStoreFile, "server");
+        sslConfig.remove(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG);
+        sslConfig.remove(SslConfigs.SSL_KEYSTORE_PASSWORD_CONFIG);
+        sslConfig.remove(SslConfigs.SSL_KEYSTORE_TYPE_CONFIG);
+        sslFactory.reconfigure(sslConfig);
+        assertNotSame("SSL context not recreated", sslContext,
+                sslFactory.sslEngineBuilder().sslContext());
 
-        SSLContext sslContext3 = sslFactory.createSSLContext(sslKeyStore(serverSslConfig), sslTrustStore(serverSslConfig));
-        assertNotNull("SSL context not created", sslContext3);
+        sslConfig = TestSslUtils.
+                createSslConfig(false, true, Mode.SERVER, newTrustStoreFile, "server");
+        try {
+            sslFactory.validateReconfiguration(sslConfig);
+            fail("Keystore configured dynamically for listener without previous keystore");
+        } catch (ConfigException e) {
+            // Expected exception
+        }
     }
 
     @Test
-    public void testUntrustedKeyStoreValidation() throws Exception {
+    public void testKeyStoreTrustStoreValidation() throws Exception {
         File trustStoreFile = File.createTempFile("truststore", ".jks");
         Map serverSslConfig = TestSslUtils.createSslConfig(false, true,
                 Mode.SERVER, trustStoreFile, "server");
-        Map untrustedConfig = TestSslUtils.createSslConfig(false, true,
-                Mode.SERVER, File.createTempFile("truststore", ".jks"), "server");
-        SslFactory sslFactory = new SslFactory(Mode.SERVER, null, true);
+        SslFactory sslFactory = new SslFactory(Mode.SERVER);
         sslFactory.configure(serverSslConfig);
-        try {
-            sslFactory.createSSLContext(sslKeyStore(untrustedConfig), null);
-            fail("Validation did not fail with untrusted keystore");
-        } catch (SSLHandshakeException e) {
-            // Expected exception
+        assertNotNull("SslEngineBuilder not created", sslFactory.sslEngineBuilder());
+    }
+
+    @Test
+    public void testUntrustedKeyStoreValidationFails() throws Exception {
+        File trustStoreFile1 = File.createTempFile("truststore1", ".jks");
+        File trustStoreFile2 = File.createTempFile("truststore2", ".jks");
+        Map sslConfig1 = TestSslUtils.createSslConfig(false, true,
+                Mode.SERVER, trustStoreFile1, "server");
+        Map sslConfig2 = TestSslUtils.createSslConfig(false, true,
+                Mode.SERVER, trustStoreFile2, "server");
+        SslFactory sslFactory = new SslFactory(Mode.SERVER, null, true);
+        for (String key : Arrays.asList(SslConfigs.SSL_TRUSTSTORE_LOCATION_CONFIG,
+                SslConfigs.SSL_TRUSTSTORE_PASSWORD_CONFIG,
+                SslConfigs.SSL_TRUSTSTORE_TYPE_CONFIG,
+                SslConfigs.SSL_TRUSTMANAGER_ALGORITHM_CONFIG)) {
+            sslConfig1.put(key, sslConfig2.get(key));
         }
         try {
-            sslFactory.createSSLContext(null, sslTrustStore(untrustedConfig));
+            sslFactory.configure(sslConfig1);
             fail("Validation did not fail with untrusted truststore");
-        } catch (SSLHandshakeException e) {
+        } catch (ConfigException e) {
             // Expected exception
         }
+    }
 
+    @Test
+    public void testKeystoreVerifiableUsingTruststore() throws Exception {
+        File trustStoreFile1 = File.createTempFile("truststore1", ".jks");
+        Map sslConfig1 = TestSslUtils.createSslConfig(false, true,
+                Mode.SERVER, trustStoreFile1, "server");
+        SslFactory sslFactory = new SslFactory(Mode.SERVER, null, true);
+        sslFactory.configure(sslConfig1);
+
+        File trustStoreFile2 = File.createTempFile("truststore2", ".jks");
+        Map sslConfig2 = TestSslUtils.createSslConfig(false, true,
+                Mode.SERVER, trustStoreFile2, "server");
         // Verify that `createSSLContext` fails even if certificate from new keystore is trusted by
         // the new truststore, if certificate is not trusted by the existing truststore on the `SslFactory`.
         // This is to prevent both keystores and truststores to be modified simultaneously on an inter-broker
         // listener to stores that may not work with other brokers where the update hasn't yet been performed.
         try {
-            sslFactory.createSSLContext(sslKeyStore(untrustedConfig), sslTrustStore(untrustedConfig));
-            fail("Validation did not fail with untrusted truststore");
-        } catch (SSLHandshakeException e) {
+            sslFactory.validateReconfiguration(sslConfig2);
+            fail("ValidateReconfiguration did not fail as expected");
+        } catch (ConfigException e) {
             // Expected exception
         }
     }
@@ -197,8 +319,8 @@ public void testCertificateEntriesValidation() throws Exception {
         assertNotEquals(SslFactory.CertificateEntries.create(ks1), SslFactory.CertificateEntries.create(ks3));
     }
 
-    private SslFactory.SecurityStore sslKeyStore(Map sslConfig) {
-        return new SslFactory.SecurityStore(
+    private SslEngineBuilder.SecurityStore sslKeyStore(Map sslConfig) {
+        return new SslEngineBuilder.SecurityStore(
                 (String) sslConfig.get(SslConfigs.SSL_KEYSTORE_TYPE_CONFIG),
                 (String) sslConfig.get(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG),
                 (Password) sslConfig.get(SslConfigs.SSL_KEYSTORE_PASSWORD_CONFIG),
@@ -206,13 +328,12 @@ private SslFactory.SecurityStore sslKeyStore(Map sslConfig) {
         );
     }
 
-    private SslFactory.SecurityStore sslTrustStore(Map sslConfig) {
-        return new SslFactory.SecurityStore(
+    private SslEngineBuilder.SecurityStore sslTrustStore(Map sslConfig) {
+        return new SslEngineBuilder.SecurityStore(
                 (String) sslConfig.get(SslConfigs.SSL_TRUSTSTORE_TYPE_CONFIG),
                 (String) sslConfig.get(SslConfigs.SSL_TRUSTSTORE_LOCATION_CONFIG),
                 (Password) sslConfig.get(SslConfigs.SSL_TRUSTSTORE_PASSWORD_CONFIG),
                 null
         );
     }
-
 }
diff --git a/clients/src/test/java/org/apache/kafka/common/security/ssl/SslPrincipalMapperTest.java b/clients/src/test/java/org/apache/kafka/common/security/ssl/SslPrincipalMapperTest.java
index 56ef977a1120a..7b7b67b1dbc3a 100644
--- a/clients/src/test/java/org/apache/kafka/common/security/ssl/SslPrincipalMapperTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/security/ssl/SslPrincipalMapperTest.java
@@ -18,9 +18,6 @@
 
 import org.junit.Test;
 
-import java.util.Arrays;
-import java.util.List;
-
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.fail;
 
@@ -28,59 +25,37 @@ public class SslPrincipalMapperTest {
 
     @Test
     public void testValidRules() {
-        testValidRule(Arrays.asList("DEFAULT"));
-        testValidRule(Arrays.asList("RULE:^CN=(.*?),OU=ServiceUsers.*$/$1/"));
-        testValidRule(Arrays.asList("RULE:^CN=(.*?),OU=ServiceUsers.*$/$1/L", "DEFAULT"));
-        testValidRule(Arrays.asList("RULE:^CN=(.*?),OU=(.*?),O=(.*?),L=(.*?),ST=(.*?),C=(.*?)$/$1@$2/"));
-        testValidRule(Arrays.asList("RULE:^.*[Cc][Nn]=([a-zA-Z0-9.]*).*$/$1/L"));
-        testValidRule(Arrays.asList("RULE:^cn=(.?),ou=(.?),dc=(.?),dc=(.?)$/$1@$2/U"));
+        testValidRule("DEFAULT");
+        testValidRule("RULE:^CN=(.*?),OU=ServiceUsers.*$/$1/");
+        testValidRule("RULE:^CN=(.*?),OU=ServiceUsers.*$/$1/L, DEFAULT");
+        testValidRule("RULE:^CN=(.*?),OU=(.*?),O=(.*?),L=(.*?),ST=(.*?),C=(.*?)$/$1@$2/");
+        testValidRule("RULE:^.*[Cc][Nn]=([a-zA-Z0-9.]*).*$/$1/L");
+        testValidRule("RULE:^cn=(.?),ou=(.?),dc=(.?),dc=(.?)$/$1@$2/U");
+
+        testValidRule("RULE:^CN=([^,ADEFLTU,]+)(,.*|$)/$1/");
+        testValidRule("RULE:^CN=([^,DEFAULT,]+)(,.*|$)/$1/");
     }
 
-    @Test
-    public void testValidSplitRules() {
-        testValidRule(Arrays.asList("DEFAULT"));
-        testValidRule(Arrays.asList("RULE:^CN=(.*?)", "OU=ServiceUsers.*$/$1/"));
-        testValidRule(Arrays.asList("RULE:^CN=(.*?)", "OU=ServiceUsers.*$/$1/L", "DEFAULT"));
-        testValidRule(Arrays.asList("RULE:^CN=(.*?)", "OU=(.*?),O=(.*?),L=(.*?)", "ST=(.*?)", "C=(.*?)$/$1@$2/"));
-        testValidRule(Arrays.asList("RULE:^.*[Cc][Nn]=([a-zA-Z0-9.]*).*$/$1/L"));
-        testValidRule(Arrays.asList("RULE:^cn=(.?)", "ou=(.?)", "dc=(.?)", "dc=(.?)$/$1@$2/U"));
-    }
-
-    private void testValidRule(List rules) {
+    private void testValidRule(String rules) {
         SslPrincipalMapper.fromRules(rules);
     }
 
     @Test
     public void testInvalidRules() {
-        testInvalidRule(Arrays.asList("default"));
-        testInvalidRule(Arrays.asList("DEFAUL"));
-        testInvalidRule(Arrays.asList("DEFAULT/L"));
-        testInvalidRule(Arrays.asList("DEFAULT/U"));
-
-        testInvalidRule(Arrays.asList("RULE:CN=(.*?),OU=ServiceUsers.*/$1"));
-        testInvalidRule(Arrays.asList("rule:^CN=(.*?),OU=ServiceUsers.*$/$1/"));
-        testInvalidRule(Arrays.asList("RULE:^CN=(.*?),OU=ServiceUsers.*$/$1/L/U"));
-        testInvalidRule(Arrays.asList("RULE:^CN=(.*?),OU=ServiceUsers.*$/L"));
-        testInvalidRule(Arrays.asList("RULE:^CN=(.*?),OU=ServiceUsers.*$/U"));
-        testInvalidRule(Arrays.asList("RULE:^CN=(.*?),OU=ServiceUsers.*$/LU"));
+        testInvalidRule("default");
+        testInvalidRule("DEFAUL");
+        testInvalidRule("DEFAULT/L");
+        testInvalidRule("DEFAULT/U");
+
+        testInvalidRule("RULE:CN=(.*?),OU=ServiceUsers.*/$1");
+        testInvalidRule("rule:^CN=(.*?),OU=ServiceUsers.*$/$1/");
+        testInvalidRule("RULE:^CN=(.*?),OU=ServiceUsers.*$/$1/L/U");
+        testInvalidRule("RULE:^CN=(.*?),OU=ServiceUsers.*$/L");
+        testInvalidRule("RULE:^CN=(.*?),OU=ServiceUsers.*$/U");
+        testInvalidRule("RULE:^CN=(.*?),OU=ServiceUsers.*$/LU");
     }
 
-    @Test
-    public void testInvalidSplitRules() {
-        testInvalidRule(Arrays.asList("default"));
-        testInvalidRule(Arrays.asList("DEFAUL"));
-        testInvalidRule(Arrays.asList("DEFAULT/L"));
-        testInvalidRule(Arrays.asList("DEFAULT/U"));
-
-        testInvalidRule(Arrays.asList("RULE:CN=(.*?)", "OU=ServiceUsers.*/$1"));
-        testInvalidRule(Arrays.asList("rule:^CN=(.*?)", "OU=ServiceUsers.*$/$1/"));
-        testInvalidRule(Arrays.asList("RULE:^CN=(.*?)", "OU=ServiceUsers.*$/$1/L/U"));
-        testInvalidRule(Arrays.asList("RULE:^CN=(.*?)", "OU=ServiceUsers.*$/L"));
-        testInvalidRule(Arrays.asList("RULE:^CN=(.*?)", "OU=ServiceUsers.*$/U"));
-        testInvalidRule(Arrays.asList("RULE:^CN=(.*?)", "OU=ServiceUsers.*$/LU"));
-    }
-
-    private void testInvalidRule(List rules) {
+    private void testInvalidRule(String rules) {
         try {
             System.out.println(SslPrincipalMapper.fromRules(rules));
             fail("should have thrown IllegalArgumentException");
@@ -90,7 +65,7 @@ private void testInvalidRule(List rules) {
 
     @Test
     public void testSslPrincipalMapper() throws Exception {
-        List rules = Arrays.asList(
+        String rules = String.join(", ",
             "RULE:^CN=(.*?),OU=ServiceUsers.*$/$1/L",
             "RULE:^CN=(.*?),OU=(.*?),O=(.*?),L=(.*?),ST=(.*?),C=(.*?)$/$1@$2/L",
             "RULE:^cn=(.*?),ou=(.*?),dc=(.*?),dc=(.*?)$/$1@$2/U",
@@ -107,4 +82,47 @@ public void testSslPrincipalMapper() throws Exception {
         assertEquals("OU=JavaSoft,O=Sun Microsystems,C=US", mapper.getName("OU=JavaSoft,O=Sun Microsystems,C=US"));
     }
 
+    private void testRulesSplitting(String expected, String rules) {
+        SslPrincipalMapper mapper = SslPrincipalMapper.fromRules(rules);
+        assertEquals(String.format("SslPrincipalMapper(rules = %s)", expected), mapper.toString());
+    }
+
+    @Test
+    public void testRulesSplitting() {
+        // seeing is believing
+        testRulesSplitting("[]", "");
+        testRulesSplitting("[DEFAULT]", "DEFAULT");
+        testRulesSplitting("[RULE:/]", "RULE://");
+        testRulesSplitting("[RULE:/.*]", "RULE:/.*/");
+        testRulesSplitting("[RULE:/.*/L]", "RULE:/.*/L");
+        testRulesSplitting("[RULE:/, DEFAULT]", "RULE://,DEFAULT");
+        testRulesSplitting("[RULE:/, DEFAULT]", "  RULE:// ,  DEFAULT  ");
+        testRulesSplitting("[RULE:   /     , DEFAULT]", "  RULE:   /     / ,  DEFAULT  ");
+        testRulesSplitting("[RULE:  /     /U, DEFAULT]", "  RULE:  /     /U   ,DEFAULT  ");
+        testRulesSplitting("[RULE:([A-Z]*)/$1/U, RULE:([a-z]+)/$1, DEFAULT]", "  RULE:([A-Z]*)/$1/U   ,RULE:([a-z]+)/$1/,   DEFAULT  ");
+
+        // empty rules are ignored
+        testRulesSplitting("[]", ",   , , ,      , , ,   ");
+        testRulesSplitting("[RULE:/, DEFAULT]", ",,RULE://,,,DEFAULT,,");
+        testRulesSplitting("[RULE: /   , DEFAULT]", ",  , RULE: /   /    ,,,   DEFAULT, ,   ");
+        testRulesSplitting("[RULE:   /  /U, DEFAULT]", "     ,  , RULE:   /  /U    ,,  ,DEFAULT, ,");
+
+        // escape sequences
+        testRulesSplitting("[RULE:\\/\\\\\\(\\)\\n\\t/\\/\\/]", "RULE:\\/\\\\\\(\\)\\n\\t/\\/\\//");
+        testRulesSplitting("[RULE:\\**\\/+/*/L, RULE:\\/*\\**/**]", "RULE:\\**\\/+/*/L,RULE:\\/*\\**/**/");
+
+        // rules rule
+        testRulesSplitting(
+            "[RULE:,RULE:,/,RULE:,\\//U, RULE:,/RULE:,, RULE:,RULE:,/L,RULE:,/L, RULE:, DEFAULT, /DEFAULT, DEFAULT]",
+            "RULE:,RULE:,/,RULE:,\\//U,RULE:,/RULE:,/,RULE:,RULE:,/L,RULE:,/L,RULE:, DEFAULT, /DEFAULT/,DEFAULT"
+        );
+    }
+
+    @Test
+    public void testCommaWithWhitespace() throws Exception {
+        String rules = "RULE:^CN=((\\\\, *|\\w)+)(,.*|$)/$1/,DEFAULT";
+
+        SslPrincipalMapper mapper = SslPrincipalMapper.fromRules(rules);
+        assertEquals("Tkac\\, Adam", mapper.getName("CN=Tkac\\, Adam,OU=ITZ,DC=geodis,DC=cz"));
+    }
 }
diff --git a/clients/src/test/java/org/apache/kafka/common/security/ssl/mock/TestKeyManagerFactory.java b/clients/src/test/java/org/apache/kafka/common/security/ssl/mock/TestKeyManagerFactory.java
new file mode 100644
index 0000000000000..dc686c246b52a
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/security/ssl/mock/TestKeyManagerFactory.java
@@ -0,0 +1,113 @@
+/*
+ * 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.security.ssl.mock;
+
+import javax.net.ssl.KeyManager;
+import javax.net.ssl.KeyManagerFactorySpi;
+import javax.net.ssl.ManagerFactoryParameters;
+import javax.net.ssl.X509ExtendedKeyManager;
+import java.io.File;
+import java.io.IOException;
+import java.net.Socket;
+import java.security.GeneralSecurityException;
+import java.security.KeyPair;
+import java.security.KeyStore;
+import java.security.Principal;
+import java.security.PrivateKey;
+import java.security.cert.X509Certificate;
+import java.util.HashMap;
+import java.util.Map;
+
+import org.apache.kafka.common.config.types.Password;
+import org.apache.kafka.test.TestSslUtils;
+import org.apache.kafka.test.TestSslUtils.CertificateBuilder;
+
+public class TestKeyManagerFactory extends KeyManagerFactorySpi {
+    public static final String ALGORITHM = "TestAlgorithm";
+
+    @Override
+    protected void engineInit(KeyStore keyStore, char[] chars) {
+
+    }
+
+    @Override
+    protected void engineInit(ManagerFactoryParameters managerFactoryParameters) {
+
+    }
+
+    @Override
+    protected KeyManager[] engineGetKeyManagers() {
+        return new KeyManager[] {new TestKeyManager()};
+    }
+
+    public static class TestKeyManager extends X509ExtendedKeyManager {
+
+        public static String mockTrustStoreFile;
+        public static final String ALIAS = "TestAlias";
+        private static final String CN = "localhost";
+        private static final String SIGNATURE_ALGORITHM = "RSA";
+        private KeyPair keyPair;
+        private X509Certificate certificate;
+
+        protected TestKeyManager() {
+            try {
+                this.keyPair = TestSslUtils.generateKeyPair(SIGNATURE_ALGORITHM);
+                CertificateBuilder certBuilder = new CertificateBuilder();
+                this.certificate = certBuilder.generate("CN=" + CN + ", O=A server", this.keyPair);
+                Map certificates = new HashMap<>();
+                certificates.put(ALIAS, certificate);
+                File trustStoreFile = File.createTempFile("testTrustStore", ".jks");
+                mockTrustStoreFile = trustStoreFile.getPath();
+                TestSslUtils.createTrustStore(mockTrustStoreFile, new Password(TestSslUtils.TRUST_STORE_PASSWORD), certificates);
+            } catch (IOException | GeneralSecurityException e) {
+                throw new RuntimeException(e);
+            }
+        }
+
+        @Override
+        public String[] getClientAliases(String s, Principal[] principals) {
+            return new String[] {ALIAS};
+        }
+
+        @Override
+        public String chooseClientAlias(String[] strings, Principal[] principals, Socket socket) {
+            return ALIAS;
+        }
+
+        @Override
+        public String[] getServerAliases(String s, Principal[] principals) {
+            return new String[] {ALIAS};
+        }
+
+        @Override
+        public String chooseServerAlias(String s, Principal[] principals, Socket socket) {
+            return ALIAS;
+        }
+
+        @Override
+        public X509Certificate[] getCertificateChain(String s) {
+            return new X509Certificate[] {this.certificate};
+        }
+
+        @Override
+        public PrivateKey getPrivateKey(String s) {
+            return this.keyPair.getPrivate();
+        }
+    }
+
+}
+
diff --git a/clients/src/test/java/org/apache/kafka/common/security/ssl/mock/TestPlainSaslServerProvider.java b/clients/src/test/java/org/apache/kafka/common/security/ssl/mock/TestPlainSaslServerProvider.java
new file mode 100644
index 0000000000000..5e6e82ea07e28
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/security/ssl/mock/TestPlainSaslServerProvider.java
@@ -0,0 +1,31 @@
+/*
+ * 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.security.ssl.mock;
+
+import java.security.Provider;
+
+public class TestPlainSaslServerProvider extends Provider {
+
+    public TestPlainSaslServerProvider() {
+        this("TestPlainSaslServerProvider", 0.1, "test plain sasl server provider");
+    }
+
+    protected TestPlainSaslServerProvider(String name, double version, String info) {
+        super(name, version, info);
+    }
+
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/security/ssl/mock/TestPlainSaslServerProviderCreator.java b/clients/src/test/java/org/apache/kafka/common/security/ssl/mock/TestPlainSaslServerProviderCreator.java
new file mode 100644
index 0000000000000..0ab927eb392ce
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/security/ssl/mock/TestPlainSaslServerProviderCreator.java
@@ -0,0 +1,34 @@
+/*
+ * 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.security.ssl.mock;
+
+import org.apache.kafka.common.security.auth.SecurityProviderCreator;
+
+import java.security.Provider;
+
+public class TestPlainSaslServerProviderCreator implements SecurityProviderCreator {
+
+    private TestPlainSaslServerProvider provider;
+
+    @Override
+    public Provider getProvider() {
+        if (provider == null) {
+            provider = new TestPlainSaslServerProvider();
+        }
+        return provider;
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/security/ssl/mock/TestProvider.java b/clients/src/test/java/org/apache/kafka/common/security/ssl/mock/TestProvider.java
new file mode 100644
index 0000000000000..fb44d3c994fed
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/security/ssl/mock/TestProvider.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.common.security.ssl.mock;
+
+import java.security.Provider;
+
+public class TestProvider extends Provider {
+
+    private static final String KEY_MANAGER_FACTORY = String.format("KeyManagerFactory.%s", TestKeyManagerFactory.ALGORITHM);
+    private static final String TRUST_MANAGER_FACTORY = String.format("TrustManagerFactory.%s", TestTrustManagerFactory.ALGORITHM);
+
+    public TestProvider() {
+        this("TestProvider", 0.1, "provider for test cases");
+    }
+
+    protected TestProvider(String name, double version, String info) {
+        super(name, version, info);
+        super.put(KEY_MANAGER_FACTORY, TestKeyManagerFactory.class.getName());
+        super.put(TRUST_MANAGER_FACTORY, TestTrustManagerFactory.class.getName());
+    }
+
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/security/ssl/mock/TestProviderCreator.java b/clients/src/test/java/org/apache/kafka/common/security/ssl/mock/TestProviderCreator.java
new file mode 100644
index 0000000000000..57c455a3e3c69
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/security/ssl/mock/TestProviderCreator.java
@@ -0,0 +1,34 @@
+/*
+ * 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.security.ssl.mock;
+
+import org.apache.kafka.common.security.auth.SecurityProviderCreator;
+
+import java.security.Provider;
+
+public class TestProviderCreator implements SecurityProviderCreator {
+
+    private TestProvider provider;
+
+    @Override
+    public Provider getProvider() {
+        if (provider == null) {
+            provider = new TestProvider();
+        }
+        return provider;
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/security/ssl/mock/TestScramSaslServerProvider.java b/clients/src/test/java/org/apache/kafka/common/security/ssl/mock/TestScramSaslServerProvider.java
new file mode 100644
index 0000000000000..c5e831028f7ab
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/security/ssl/mock/TestScramSaslServerProvider.java
@@ -0,0 +1,31 @@
+/*
+ * 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.security.ssl.mock;
+
+import java.security.Provider;
+
+public class TestScramSaslServerProvider extends Provider {
+
+    public TestScramSaslServerProvider() {
+        this("TestScramSaslServerProvider", 0.1, "test scram sasl server provider");
+    }
+
+    protected TestScramSaslServerProvider(String name, double version, String info) {
+        super(name, version, info);
+    }
+
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/security/ssl/mock/TestScramSaslServerProviderCreator.java b/clients/src/test/java/org/apache/kafka/common/security/ssl/mock/TestScramSaslServerProviderCreator.java
new file mode 100644
index 0000000000000..72eb8800d6089
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/security/ssl/mock/TestScramSaslServerProviderCreator.java
@@ -0,0 +1,34 @@
+/*
+ * 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.security.ssl.mock;
+
+import org.apache.kafka.common.security.auth.SecurityProviderCreator;
+
+import java.security.Provider;
+
+public class TestScramSaslServerProviderCreator implements SecurityProviderCreator {
+
+    private TestScramSaslServerProvider provider;
+
+    @Override
+    public Provider getProvider() {
+        if (provider == null) {
+            provider = new TestScramSaslServerProvider();
+        }
+        return provider;
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/security/ssl/mock/TestTrustManagerFactory.java b/clients/src/test/java/org/apache/kafka/common/security/ssl/mock/TestTrustManagerFactory.java
new file mode 100644
index 0000000000000..4115a5f8cbdc5
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/security/ssl/mock/TestTrustManagerFactory.java
@@ -0,0 +1,88 @@
+/*
+ * 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.security.ssl.mock;
+
+import javax.net.ssl.ManagerFactoryParameters;
+import javax.net.ssl.SSLEngine;
+import javax.net.ssl.TrustManager;
+import javax.net.ssl.TrustManagerFactorySpi;
+import javax.net.ssl.X509ExtendedTrustManager;
+import java.net.Socket;
+import java.security.KeyStore;
+import java.security.cert.CertificateException;
+import java.security.cert.X509Certificate;
+
+public class TestTrustManagerFactory extends TrustManagerFactorySpi {
+    public static final String ALGORITHM = "TestAlgorithm";
+
+    @Override
+    protected void engineInit(KeyStore keyStore) {
+
+    }
+
+    @Override
+    protected void engineInit(ManagerFactoryParameters managerFactoryParameters) {
+
+    }
+
+    @Override
+    protected TrustManager[] engineGetTrustManagers() {
+        return new TrustManager[] {new TestTrustManager()};
+    }
+
+    public static class TestTrustManager extends X509ExtendedTrustManager {
+
+        public static final String ALIAS = "TestAlias";
+
+        @Override
+        public void checkClientTrusted(X509Certificate[] x509Certificates, String s) throws CertificateException {
+
+        }
+
+        @Override
+        public void checkServerTrusted(X509Certificate[] x509Certificates, String s) throws CertificateException {
+
+        }
+
+        @Override
+        public X509Certificate[] getAcceptedIssuers() {
+            return new X509Certificate[0];
+        }
+
+        @Override
+        public void checkClientTrusted(X509Certificate[] x509Certificates, String s, Socket socket) throws CertificateException {
+
+        }
+
+        @Override
+        public void checkServerTrusted(X509Certificate[] x509Certificates, String s, Socket socket) throws CertificateException {
+
+        }
+
+        @Override
+        public void checkClientTrusted(X509Certificate[] x509Certificates, String s, SSLEngine sslEngine) throws CertificateException {
+
+        }
+
+        @Override
+        public void checkServerTrusted(X509Certificate[] x509Certificates, String s, SSLEngine sslEngine) throws CertificateException {
+
+        }
+    }
+
+}
+
diff --git a/clients/src/test/java/org/apache/kafka/common/utils/AppInfoParserTest.java b/clients/src/test/java/org/apache/kafka/common/utils/AppInfoParserTest.java
new file mode 100644
index 0000000000000..34dba81b1147a
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/utils/AppInfoParserTest.java
@@ -0,0 +1,88 @@
+/*
+ * 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.utils;
+
+import org.apache.kafka.common.metrics.Metrics;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+import javax.management.JMException;
+import javax.management.MBeanServer;
+import javax.management.MalformedObjectNameException;
+import javax.management.ObjectName;
+
+import java.lang.management.ManagementFactory;
+
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+public class AppInfoParserTest {
+    private static final String EXPECTED_COMMIT_VERSION = AppInfoParser.DEFAULT_VALUE;
+    private static final String EXPECTED_VERSION = AppInfoParser.DEFAULT_VALUE;
+    private static final Long EXPECTED_START_MS = 1552313875722L;
+    private static final String METRICS_PREFIX = "app-info-test";
+    private static final String METRICS_ID = "test";
+
+    private Metrics metrics;
+    private MBeanServer mBeanServer;
+
+    @Before
+    public void setUp() {
+        metrics = new Metrics(new MockTime(1));
+        mBeanServer = ManagementFactory.getPlatformMBeanServer();
+    }
+
+    @After
+    public void tearDown() {
+        metrics.close();
+    }
+
+    @Test
+    public void testRegisterAppInfoRegistersMetrics() throws JMException {
+        registerAppInfo();
+    }
+
+    @Test
+    public void testUnregisterAppInfoUnregistersMetrics() throws JMException {
+        registerAppInfo();
+        AppInfoParser.unregisterAppInfo(METRICS_PREFIX, METRICS_ID, metrics);
+
+        assertFalse(mBeanServer.isRegistered(expectedAppObjectName()));
+        assertNull(metrics.metric(metrics.metricName("commit-id", "app-info")));
+        assertNull(metrics.metric(metrics.metricName("version", "app-info")));
+        assertNull(metrics.metric(metrics.metricName("start-time-ms", "app-info")));
+    }
+
+    private void registerAppInfo() throws JMException {
+        assertEquals(EXPECTED_COMMIT_VERSION, AppInfoParser.getCommitId());
+        assertEquals(EXPECTED_VERSION, AppInfoParser.getVersion());
+
+        AppInfoParser.registerAppInfo(METRICS_PREFIX, METRICS_ID, metrics, EXPECTED_START_MS);
+
+        assertTrue(mBeanServer.isRegistered(expectedAppObjectName()));
+        assertEquals(EXPECTED_COMMIT_VERSION, metrics.metric(metrics.metricName("commit-id", "app-info")).metricValue());
+        assertEquals(EXPECTED_VERSION, metrics.metric(metrics.metricName("version", "app-info")).metricValue());
+        assertEquals(EXPECTED_START_MS, metrics.metric(metrics.metricName("start-time-ms", "app-info")).metricValue());
+    }
+
+    private ObjectName expectedAppObjectName() throws MalformedObjectNameException {
+        return new ObjectName(METRICS_PREFIX + ":type=app-info,id=" + METRICS_ID);
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/utils/ByteUtilsTest.java b/clients/src/test/java/org/apache/kafka/common/utils/ByteUtilsTest.java
index ce23a3378a857..f15e26d222b05 100644
--- a/clients/src/test/java/org/apache/kafka/common/utils/ByteUtilsTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/utils/ByteUtilsTest.java
@@ -33,11 +33,17 @@ public class ByteUtilsTest {
     private final byte x01 = 0x01;
     private final byte x02 = 0x02;
     private final byte x0F = 0x0f;
+    private final byte x07 = 0x07;
+    private final byte x08 = 0x08;
+    private final byte x3F = 0x3f;
+    private final byte x40 = 0x40;
     private final byte x7E = 0x7E;
     private final byte x7F = 0x7F;
     private final byte xFF = (byte) 0xff;
     private final byte x80 = (byte) 0x80;
     private final byte x81 = (byte) 0x81;
+    private final byte xBF = (byte) 0xbf;
+    private final byte xC0 = (byte) 0xc0;
     private final byte xFE = (byte) 0xfe;
 
     @Test
@@ -111,6 +117,24 @@ public void testWriteUnsignedIntLEToOutputStream() throws IOException {
         assertArrayEquals(new byte[] {(byte) 0xf1, (byte) 0xf2, (byte) 0xf3, (byte) 0xf4}, os2.toByteArray());
     }
 
+    @Test
+    public void testUnsignedVarintSerde() throws Exception {
+        assertUnsignedVarintSerde(0, new byte[] {x00});
+        assertUnsignedVarintSerde(-1, new byte[] {xFF, xFF, xFF, xFF, x0F});
+        assertUnsignedVarintSerde(1, new byte[] {x01});
+        assertUnsignedVarintSerde(63, new byte[] {x3F});
+        assertUnsignedVarintSerde(-64, new byte[] {xC0, xFF, xFF, xFF, x0F});
+        assertUnsignedVarintSerde(64, new byte[] {x40});
+        assertUnsignedVarintSerde(8191, new byte[] {xFF, x3F});
+        assertUnsignedVarintSerde(-8192, new byte[] {x80, xC0, xFF, xFF, x0F});
+        assertUnsignedVarintSerde(8192, new byte[] {x80, x40});
+        assertUnsignedVarintSerde(-8193, new byte[] {xFF, xBF, xFF, xFF, x0F});
+        assertUnsignedVarintSerde(1048575, new byte[] {xFF, xFF, x3F});
+        assertUnsignedVarintSerde(1048576, new byte[] {x80, x80, x40});
+        assertUnsignedVarintSerde(Integer.MAX_VALUE, new byte[] {xFF, xFF, xFF, xFF, x07});
+        assertUnsignedVarintSerde(Integer.MIN_VALUE, new byte[] {x80, x80, x80, x80, x08});
+    }
+
     @Test
     public void testVarintSerde() throws Exception {
         assertVarintSerde(0, new byte[] {x00});
@@ -197,6 +221,22 @@ public void testInvalidVarlong() {
         ByteUtils.readVarlong(buf);
     }
 
+    private void assertUnsignedVarintSerde(int value, byte[] expectedEncoding) throws IOException {
+        ByteBuffer buf = ByteBuffer.allocate(32);
+        ByteUtils.writeUnsignedVarint(value, buf);
+        buf.flip();
+        assertArrayEquals(expectedEncoding, Utils.toArray(buf));
+        assertEquals(value, ByteUtils.readUnsignedVarint(buf.duplicate()));
+
+        buf.rewind();
+        DataOutputStream out = new DataOutputStream(new ByteBufferOutputStream(buf));
+        ByteUtils.writeUnsignedVarint(value, out);
+        buf.flip();
+        assertArrayEquals(expectedEncoding, Utils.toArray(buf));
+        DataInputStream in = new DataInputStream(new ByteBufferInputStream(buf));
+        assertEquals(value, ByteUtils.readUnsignedVarint(in));
+    }
+
     private void assertVarintSerde(int value, byte[] expectedEncoding) throws IOException {
         ByteBuffer buf = ByteBuffer.allocate(32);
         ByteUtils.writeVarint(value, buf);
diff --git a/clients/src/test/java/org/apache/kafka/common/utils/BytesTest.java b/clients/src/test/java/org/apache/kafka/common/utils/BytesTest.java
new file mode 100644
index 0000000000000..bf7ec712ddca2
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/utils/BytesTest.java
@@ -0,0 +1,84 @@
+/*
+ * 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.utils;
+
+import org.junit.Test;
+
+import java.util.Comparator;
+import java.util.NavigableMap;
+import java.util.TreeMap;
+
+import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertEquals;
+
+public class BytesTest {
+
+    @Test
+    public void testIncrement() {
+        byte[] input = new byte[]{(byte) 0xAB, (byte) 0xCD, (byte) 0xFF};
+        byte[] expected = new byte[]{(byte) 0xAB, (byte) 0xCE, (byte) 0x00};
+        Bytes output = Bytes.increment(Bytes.wrap(input));
+        assertArrayEquals(output.get(), expected);
+    }
+
+    @Test
+    public void testIncrementUpperBoundary() {
+        byte[] input = new byte[]{(byte) 0xFF, (byte) 0xFF, (byte) 0xFF};
+        assertThrows(IndexOutOfBoundsException.class, () -> Bytes.increment(Bytes.wrap(input)));
+    }
+
+    @Test
+    public void testIncrementWithSubmap() {
+        final NavigableMap map = new TreeMap<>();
+        Bytes key1 = Bytes.wrap(new byte[]{(byte) 0xAA});
+        byte[] val = new byte[]{(byte) 0x00};
+        map.put(key1, val);
+
+        Bytes key2 = Bytes.wrap(new byte[]{(byte) 0xAA, (byte) 0xAA});
+        map.put(key2, val);
+
+        Bytes key3 = Bytes.wrap(new byte[]{(byte) 0xAA, (byte) 0x00, (byte) 0xFF, (byte) 0xFF, (byte) 0xFF});
+        map.put(key3, val);
+
+        Bytes key4 = Bytes.wrap(new byte[]{(byte) 0xAB, (byte) 0x00});
+        map.put(key4, val);
+
+        Bytes key5 = Bytes.wrap(new byte[]{(byte) 0x00, (byte) 0x00, (byte) 0x00, (byte) 0x01});
+        map.put(key5, val);
+
+        Bytes prefix = key1;
+        Bytes prefixEnd = Bytes.increment(prefix);
+
+        Comparator comparator = map.comparator();
+        final int result = comparator == null ? prefix.compareTo(prefixEnd) : comparator.compare(prefix, prefixEnd);
+        NavigableMap subMapResults;
+        if (result > 0) {
+            //Prefix increment would cause a wrap-around. Get the submap from toKey to the end of the map
+            subMapResults = map.tailMap(prefix, true);
+        } else {
+            subMapResults = map.subMap(prefix, true, prefixEnd, false);
+        }
+
+        NavigableMap subMapExpected = new TreeMap<>();
+        subMapExpected.put(key1, val);
+        subMapExpected.put(key2, val);
+        subMapExpected.put(key3, val);
+
+        assertEquals(subMapExpected.keySet(), subMapResults.keySet());
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/utils/CircularIteratorTest.java b/clients/src/test/java/org/apache/kafka/common/utils/CircularIteratorTest.java
new file mode 100644
index 0000000000000..a37c85e4beeae
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/utils/CircularIteratorTest.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.common.utils;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
+
+import java.util.Arrays;
+import java.util.Collections;
+
+import org.junit.Test;
+
+public class CircularIteratorTest {
+
+    @Test
+    public void testNullCollection() {
+        assertThrows(NullPointerException.class, () -> new CircularIterator<>(null));
+    }
+
+    @Test
+    public void testEmptyCollection() {
+        assertThrows(IllegalArgumentException.class, () -> new CircularIterator<>(Collections.emptyList()));
+    }
+
+    @Test()
+    public void testCycleCollection() {
+        final CircularIterator it = new CircularIterator<>(Arrays.asList("A", "B", null, "C"));
+
+        assertEquals("A", it.peek());
+        assertTrue(it.hasNext());
+        assertEquals("A", it.next());
+        assertEquals("B", it.peek());
+        assertTrue(it.hasNext());
+        assertEquals("B", it.next());
+        assertEquals(null, it.peek());
+        assertTrue(it.hasNext());
+        assertEquals(null, it.next());
+        assertEquals("C", it.peek());
+        assertTrue(it.hasNext());
+        assertEquals("C", it.next());
+        assertEquals("A", it.peek());
+        assertTrue(it.hasNext());
+        assertEquals("A", it.next());
+        assertEquals("B", it.peek());
+
+        // Check that peek does not have any side-effects
+        assertEquals("B", it.peek());
+    }
+
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/utils/FixedOrderMapTest.java b/clients/src/test/java/org/apache/kafka/common/utils/FixedOrderMapTest.java
new file mode 100644
index 0000000000000..7d3f3f7525677
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/utils/FixedOrderMapTest.java
@@ -0,0 +1,86 @@
+/*
+ * 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.utils;
+
+import org.hamcrest.CoreMatchers;
+import org.junit.Test;
+
+import java.util.Iterator;
+import java.util.Map;
+
+import static org.apache.kafka.common.utils.Utils.mkEntry;
+import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.junit.Assert.fail;
+
+public class FixedOrderMapTest {
+    @Test
+    public void shouldMaintainOrderWhenAdding() {
+        final FixedOrderMap map = new FixedOrderMap<>();
+        map.put("a", 0);
+        map.put("b", 1);
+        map.put("c", 2);
+        map.put("b", 3);
+        final Iterator> iterator = map.entrySet().iterator();
+        assertThat(iterator.next(), is(mkEntry("a", 0)));
+        assertThat(iterator.next(), is(mkEntry("b", 3)));
+        assertThat(iterator.next(), is(mkEntry("c", 2)));
+        assertThat(iterator.hasNext(), is(false));
+    }
+
+    @SuppressWarnings("deprecation")
+    @Test
+    public void shouldForbidRemove() {
+        final FixedOrderMap map = new FixedOrderMap<>();
+        map.put("a", 0);
+        try {
+            map.remove("a");
+            fail("expected exception");
+        } catch (final RuntimeException e) {
+            assertThat(e, CoreMatchers.instanceOf(UnsupportedOperationException.class));
+        }
+        assertThat(map.get("a"), is(0));
+    }
+
+    @SuppressWarnings("deprecation")
+    @Test
+    public void shouldForbidConditionalRemove() {
+        final FixedOrderMap map = new FixedOrderMap<>();
+        map.put("a", 0);
+        try {
+            map.remove("a", 0);
+            fail("expected exception");
+        } catch (final RuntimeException e) {
+            assertThat(e, CoreMatchers.instanceOf(UnsupportedOperationException.class));
+        }
+        assertThat(map.get("a"), is(0));
+    }
+
+    @SuppressWarnings("deprecation")
+    @Test
+    public void shouldForbidConditionalClear() {
+        final FixedOrderMap map = new FixedOrderMap<>();
+        map.put("a", 0);
+        try {
+            map.clear();
+            fail("expected exception");
+        } catch (final RuntimeException e) {
+            assertThat(e, CoreMatchers.instanceOf(UnsupportedOperationException.class));
+        }
+        assertThat(map.get("a"), is(0));
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/utils/FlattenedIteratorTest.java b/clients/src/test/java/org/apache/kafka/common/utils/FlattenedIteratorTest.java
new file mode 100644
index 0000000000000..95206cae473a8
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/utils/FlattenedIteratorTest.java
@@ -0,0 +1,115 @@
+/*
+ * 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.utils;
+
+import org.junit.Test;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.stream.Collectors;
+
+import static java.util.Arrays.asList;
+import static java.util.Collections.emptyList;
+import static org.junit.Assert.assertEquals;
+
+public class FlattenedIteratorTest {
+
+    @Test
+    public void testNestedLists() {
+        List> list = asList(
+            asList("foo", "a", "bc"),
+            asList("ddddd"),
+            asList("", "bar2", "baz45"));
+
+        Iterable flattenedIterable = () -> new FlattenedIterator<>(list.iterator(), l -> l.iterator());
+        List flattened = new ArrayList<>();
+        flattenedIterable.forEach(flattened::add);
+
+        assertEquals(list.stream().flatMap(l -> l.stream()).collect(Collectors.toList()), flattened);
+
+        // Ensure we can iterate multiple times
+        List flattened2 = new ArrayList<>();
+        flattenedIterable.forEach(flattened2::add);
+
+        assertEquals(flattened, flattened2);
+    }
+
+    @Test
+    public void testEmptyList() {
+        List> list = emptyList();
+
+        Iterable flattenedIterable = () -> new FlattenedIterator<>(list.iterator(), l -> l.iterator());
+        List flattened = new ArrayList<>();
+        flattenedIterable.forEach(flattened::add);
+
+        assertEquals(emptyList(), flattened);
+    }
+
+    @Test
+    public void testNestedSingleEmptyList() {
+        List> list = asList(emptyList());
+
+        Iterable flattenedIterable = () -> new FlattenedIterator<>(list.iterator(), l -> l.iterator());
+        List flattened = new ArrayList<>();
+        flattenedIterable.forEach(flattened::add);
+
+        assertEquals(emptyList(), flattened);
+    }
+
+    @Test
+    public void testEmptyListFollowedByNonEmpty() {
+        List> list = asList(
+            emptyList(),
+            asList("boo", "b", "de"));
+
+        Iterable flattenedIterable = () -> new FlattenedIterator<>(list.iterator(), l -> l.iterator());
+        List flattened = new ArrayList<>();
+        flattenedIterable.forEach(flattened::add);
+
+        assertEquals(list.stream().flatMap(l -> l.stream()).collect(Collectors.toList()), flattened);
+    }
+
+    @Test
+    public void testEmptyListInBetweenNonEmpty() {
+        List> list = asList(
+            asList("aadwdwdw"),
+            emptyList(),
+            asList("ee", "aa", "dd"));
+
+        Iterable flattenedIterable = () -> new FlattenedIterator<>(list.iterator(), l -> l.iterator());
+        List flattened = new ArrayList<>();
+        flattenedIterable.forEach(flattened::add);
+
+        assertEquals(list.stream().flatMap(l -> l.stream()).collect(Collectors.toList()), flattened);
+    }
+
+    @Test
+    public void testEmptyListAtTheEnd() {
+        List> list = asList(
+            asList("ee", "dd"),
+            asList("e"),
+            emptyList());
+
+        Iterable flattenedIterable = () -> new FlattenedIterator<>(list.iterator(), l -> l.iterator());
+        List flattened = new ArrayList<>();
+        flattenedIterable.forEach(flattened::add);
+
+        assertEquals(list.stream().flatMap(l -> l.stream()).collect(Collectors.toList()), flattened);
+    }
+
+}
+
diff --git a/clients/src/test/java/org/apache/kafka/common/utils/ImplicitLinkedHashCollectionTest.java b/clients/src/test/java/org/apache/kafka/common/utils/ImplicitLinkedHashCollectionTest.java
new file mode 100644
index 0000000000000..389c24e456b96
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/utils/ImplicitLinkedHashCollectionTest.java
@@ -0,0 +1,546 @@
+/*
+ * 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.utils;
+
+import org.junit.Assert;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.Timeout;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Iterator;
+import java.util.LinkedHashSet;
+import java.util.List;
+import java.util.ListIterator;
+import java.util.Random;
+import java.util.Set;
+
+import static org.junit.Assert.assertNotEquals;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.fail;
+
+/**
+ * A unit test for ImplicitLinkedHashCollection.
+ */
+public class ImplicitLinkedHashCollectionTest {
+    @Rule
+    final public Timeout globalTimeout = Timeout.millis(120000);
+
+    final static class TestElement implements ImplicitLinkedHashCollection.Element {
+        private int prev = ImplicitLinkedHashCollection.INVALID_INDEX;
+        private int next = ImplicitLinkedHashCollection.INVALID_INDEX;
+        private final int val;
+
+        TestElement(int val) {
+            this.val = val;
+        }
+
+        @Override
+        public int prev() {
+            return prev;
+        }
+
+        @Override
+        public void setPrev(int prev) {
+            this.prev = prev;
+        }
+
+        @Override
+        public int next() {
+            return next;
+        }
+
+        @Override
+        public void setNext(int next) {
+            this.next = next;
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) return true;
+            if ((o == null) || (o.getClass() != TestElement.class)) return false;
+            TestElement that = (TestElement) o;
+            return val == that.val;
+        }
+
+        @Override
+        public String toString() {
+            return "TestElement(" + val + ")";
+        }
+
+        @Override
+        public int hashCode() {
+            return val;
+        }
+    }
+
+    @Test
+    public void testNullForbidden() {
+        ImplicitLinkedHashMultiCollection multiColl = new ImplicitLinkedHashMultiCollection<>();
+        assertFalse(multiColl.add(null));
+    }
+
+    @Test
+    public void testInsertDelete() {
+        ImplicitLinkedHashCollection coll = new ImplicitLinkedHashCollection<>(100);
+        assertTrue(coll.add(new TestElement(1)));
+        TestElement second = new TestElement(2);
+        assertTrue(coll.add(second));
+        assertTrue(coll.add(new TestElement(3)));
+        assertFalse(coll.add(new TestElement(3)));
+        assertEquals(3, coll.size());
+        assertTrue(coll.contains(new TestElement(1)));
+        assertFalse(coll.contains(new TestElement(4)));
+        TestElement secondAgain = coll.find(new TestElement(2));
+        assertTrue(second == secondAgain);
+        assertTrue(coll.remove(new TestElement(1)));
+        assertFalse(coll.remove(new TestElement(1)));
+        assertEquals(2, coll.size());
+        coll.clear();
+        assertEquals(0, coll.size());
+    }
+
+    static void expectTraversal(Iterator iterator, Integer... sequence) {
+        int i = 0;
+        while (iterator.hasNext()) {
+            TestElement element = iterator.next();
+            Assert.assertTrue("Iterator yieled " + (i + 1) + " elements, but only " +
+                    sequence.length + " were expected.", i < sequence.length);
+            Assert.assertEquals("Iterator value number " + (i + 1) + " was incorrect.",
+                    sequence[i].intValue(), element.val);
+            i = i + 1;
+        }
+        Assert.assertTrue("Iterator yieled " + (i + 1) + " elements, but " +
+                sequence.length + " were expected.", i == sequence.length);
+    }
+
+    static void expectTraversal(Iterator iter, Iterator expectedIter) {
+        int i = 0;
+        while (iter.hasNext()) {
+            TestElement element = iter.next();
+            Assert.assertTrue("Iterator yieled " + (i + 1) + " elements, but only " +
+                    i + " were expected.", expectedIter.hasNext());
+            Integer expected = expectedIter.next();
+            Assert.assertEquals("Iterator value number " + (i + 1) + " was incorrect.",
+                    expected.intValue(), element.val);
+            i = i + 1;
+        }
+        Assert.assertFalse("Iterator yieled " + i + " elements, but at least " +
+                (i + 1) + " were expected.", expectedIter.hasNext());
+    }
+
+    @Test
+    public void testTraversal() {
+        ImplicitLinkedHashCollection coll = new ImplicitLinkedHashCollection<>();
+        expectTraversal(coll.iterator());
+        assertTrue(coll.add(new TestElement(2)));
+        expectTraversal(coll.iterator(), 2);
+        assertTrue(coll.add(new TestElement(1)));
+        expectTraversal(coll.iterator(), 2, 1);
+        assertTrue(coll.add(new TestElement(100)));
+        expectTraversal(coll.iterator(), 2, 1, 100);
+        assertTrue(coll.remove(new TestElement(1)));
+        expectTraversal(coll.iterator(), 2, 100);
+        assertTrue(coll.add(new TestElement(1)));
+        expectTraversal(coll.iterator(), 2, 100, 1);
+        Iterator iter = coll.iterator();
+        iter.next();
+        iter.next();
+        iter.remove();
+        iter.next();
+        assertFalse(iter.hasNext());
+        expectTraversal(coll.iterator(), 2, 1);
+        List list = new ArrayList<>();
+        list.add(new TestElement(1));
+        list.add(new TestElement(2));
+        assertTrue(coll.removeAll(list));
+        assertFalse(coll.removeAll(list));
+        expectTraversal(coll.iterator());
+        assertEquals(0, coll.size());
+        assertTrue(coll.isEmpty());
+    }
+
+    @Test
+    public void testSetViewGet() {
+        ImplicitLinkedHashCollection coll = new ImplicitLinkedHashCollection<>();
+        coll.add(new TestElement(1));
+        coll.add(new TestElement(2));
+        coll.add(new TestElement(3));
+
+        Set set = coll.valuesSet();
+        assertTrue(set.contains(new TestElement(1)));
+        assertTrue(set.contains(new TestElement(2)));
+        assertTrue(set.contains(new TestElement(3)));
+        assertEquals(3, set.size());
+    }
+
+    @Test
+    public void testSetViewModification() {
+        ImplicitLinkedHashCollection coll = new ImplicitLinkedHashCollection<>();
+        coll.add(new TestElement(1));
+        coll.add(new TestElement(2));
+        coll.add(new TestElement(3));
+
+        // Removal from set is reflected in collection
+        Set set = coll.valuesSet();
+        set.remove(new TestElement(1));
+        assertFalse(coll.contains(new TestElement(1)));
+        assertEquals(2, coll.size());
+
+        // Addition to set is reflected in collection
+        set.add(new TestElement(4));
+        assertTrue(coll.contains(new TestElement(4)));
+        assertEquals(3, coll.size());
+
+        // Removal from collection is reflected in set
+        coll.remove(new TestElement(2));
+        assertFalse(set.contains(new TestElement(2)));
+        assertEquals(2, set.size());
+
+        // Addition to collection is reflected in set
+        coll.add(new TestElement(5));
+        assertTrue(set.contains(new TestElement(5)));
+        assertEquals(3, set.size());
+
+        // Ordering in the collection is maintained
+        int val = 3;
+        for (TestElement e : coll) {
+            assertEquals(val, e.val);
+            ++val;
+        }
+    }
+
+    @Test
+    public void testListViewGet() {
+        ImplicitLinkedHashCollection coll = new ImplicitLinkedHashCollection<>();
+        coll.add(new TestElement(1));
+        coll.add(new TestElement(2));
+        coll.add(new TestElement(3));
+
+        List list = coll.valuesList();
+        assertEquals(1, list.get(0).val);
+        assertEquals(2, list.get(1).val);
+        assertEquals(3, list.get(2).val);
+        assertEquals(3, list.size());
+    }
+
+    @Test
+    public void testListViewModification() {
+        ImplicitLinkedHashCollection coll = new ImplicitLinkedHashCollection<>();
+        coll.add(new TestElement(1));
+        coll.add(new TestElement(2));
+        coll.add(new TestElement(3));
+
+        // Removal from list is reflected in collection
+        List list = coll.valuesList();
+        list.remove(1);
+        assertTrue(coll.contains(new TestElement(1)));
+        assertFalse(coll.contains(new TestElement(2)));
+        assertTrue(coll.contains(new TestElement(3)));
+        assertEquals(2, coll.size());
+
+        // Removal from collection is reflected in list
+        coll.remove(new TestElement(1));
+        assertEquals(3, list.get(0).val);
+        assertEquals(1, list.size());
+
+        // Addition to collection is reflected in list
+        coll.add(new TestElement(4));
+        assertEquals(3, list.get(0).val);
+        assertEquals(4, list.get(1).val);
+        assertEquals(2, list.size());
+    }
+
+    @Test
+    public void testEmptyListIterator() {
+        ImplicitLinkedHashCollection coll = new ImplicitLinkedHashCollection<>();
+        ListIterator iter = coll.valuesList().listIterator();
+        assertFalse(iter.hasNext());
+        assertFalse(iter.hasPrevious());
+        assertEquals(0, iter.nextIndex());
+        assertEquals(-1, iter.previousIndex());
+    }
+
+    @Test
+    public void testListIteratorCreation() {
+        ImplicitLinkedHashCollection coll = new ImplicitLinkedHashCollection<>();
+        coll.add(new TestElement(1));
+        coll.add(new TestElement(2));
+        coll.add(new TestElement(3));
+
+        // Iterator created at the start of the list should have a next but no prev
+        ListIterator iter = coll.valuesList().listIterator();
+        assertTrue(iter.hasNext());
+        assertFalse(iter.hasPrevious());
+        assertEquals(0, iter.nextIndex());
+        assertEquals(-1, iter.previousIndex());
+
+        // Iterator created in the middle of the list should have both a next and a prev
+        iter = coll.valuesList().listIterator(2);
+        assertTrue(iter.hasNext());
+        assertTrue(iter.hasPrevious());
+        assertEquals(2, iter.nextIndex());
+        assertEquals(1, iter.previousIndex());
+
+        // Iterator created at the end of the list should have a prev but no next
+        iter = coll.valuesList().listIterator(3);
+        assertFalse(iter.hasNext());
+        assertTrue(iter.hasPrevious());
+        assertEquals(3, iter.nextIndex());
+        assertEquals(2, iter.previousIndex());
+    }
+
+    @Test
+    public void testListIteratorTraversal() {
+        ImplicitLinkedHashCollection coll = new ImplicitLinkedHashCollection<>();
+        coll.add(new TestElement(1));
+        coll.add(new TestElement(2));
+        coll.add(new TestElement(3));
+        ListIterator iter = coll.valuesList().listIterator();
+
+        // Step the iterator forward to the end of the list
+        assertTrue(iter.hasNext());
+        assertFalse(iter.hasPrevious());
+        assertEquals(0, iter.nextIndex());
+        assertEquals(-1, iter.previousIndex());
+
+        assertEquals(1, iter.next().val);
+        assertTrue(iter.hasNext());
+        assertTrue(iter.hasPrevious());
+        assertEquals(1, iter.nextIndex());
+        assertEquals(0, iter.previousIndex());
+
+        assertEquals(2, iter.next().val);
+        assertTrue(iter.hasNext());
+        assertTrue(iter.hasPrevious());
+        assertEquals(2, iter.nextIndex());
+        assertEquals(1, iter.previousIndex());
+
+        assertEquals(3, iter.next().val);
+        assertFalse(iter.hasNext());
+        assertTrue(iter.hasPrevious());
+        assertEquals(3, iter.nextIndex());
+        assertEquals(2, iter.previousIndex());
+
+        // Step back to the middle of the list
+        assertEquals(3, iter.previous().val);
+        assertTrue(iter.hasNext());
+        assertTrue(iter.hasPrevious());
+        assertEquals(2, iter.nextIndex());
+        assertEquals(1, iter.previousIndex());
+
+        assertEquals(2, iter.previous().val);
+        assertTrue(iter.hasNext());
+        assertTrue(iter.hasPrevious());
+        assertEquals(1, iter.nextIndex());
+        assertEquals(0, iter.previousIndex());
+
+        // Step forward one and then back one, return value should remain the same
+        assertEquals(2, iter.next().val);
+        assertTrue(iter.hasNext());
+        assertTrue(iter.hasPrevious());
+        assertEquals(2, iter.nextIndex());
+        assertEquals(1, iter.previousIndex());
+
+        assertEquals(2, iter.previous().val);
+        assertTrue(iter.hasNext());
+        assertTrue(iter.hasPrevious());
+        assertEquals(1, iter.nextIndex());
+        assertEquals(0, iter.previousIndex());
+
+        // Step back to the front of the list
+        assertEquals(1, iter.previous().val);
+        assertTrue(iter.hasNext());
+        assertFalse(iter.hasPrevious());
+        assertEquals(0, iter.nextIndex());
+        assertEquals(-1, iter.previousIndex());
+    }
+
+    @Test
+    public void testListIteratorRemove() {
+        ImplicitLinkedHashCollection coll = new ImplicitLinkedHashCollection<>();
+        coll.add(new TestElement(1));
+        coll.add(new TestElement(2));
+        coll.add(new TestElement(3));
+        coll.add(new TestElement(4));
+        coll.add(new TestElement(5));
+
+        ListIterator iter = coll.valuesList().listIterator();
+        try {
+            iter.remove();
+            fail("Calling remove() without calling next() or previous() should raise an exception");
+        } catch (IllegalStateException e) {
+            // expected
+        }
+
+        // Remove after next()
+        iter.next();
+        iter.next();
+        iter.next();
+        iter.remove();
+        assertTrue(iter.hasNext());
+        assertTrue(iter.hasPrevious());
+        assertEquals(2, iter.nextIndex());
+        assertEquals(1, iter.previousIndex());
+
+        try {
+            iter.remove();
+            fail("Calling remove() twice without calling next() or previous() in between should raise an exception");
+        } catch (IllegalStateException e) {
+            // expected
+        }
+
+        // Remove after previous()
+        assertEquals(2, iter.previous().val);
+        iter.remove();
+        assertTrue(iter.hasNext());
+        assertTrue(iter.hasPrevious());
+        assertEquals(1, iter.nextIndex());
+        assertEquals(0, iter.previousIndex());
+
+        // Remove the first element of the list
+        assertEquals(1, iter.previous().val);
+        iter.remove();
+        assertTrue(iter.hasNext());
+        assertFalse(iter.hasPrevious());
+        assertEquals(0, iter.nextIndex());
+        assertEquals(-1, iter.previousIndex());
+
+        // Remove the last element of the list
+        assertEquals(4, iter.next().val);
+        assertEquals(5, iter.next().val);
+        iter.remove();
+        assertFalse(iter.hasNext());
+        assertTrue(iter.hasPrevious());
+        assertEquals(1, iter.nextIndex());
+        assertEquals(0, iter.previousIndex());
+
+        // Remove the final remaining element of the list
+        assertEquals(4, iter.previous().val);
+        iter.remove();
+        assertFalse(iter.hasNext());
+        assertFalse(iter.hasPrevious());
+        assertEquals(0, iter.nextIndex());
+        assertEquals(-1, iter.previousIndex());
+
+    }
+
+    @Test
+    public void testCollisions() {
+        ImplicitLinkedHashCollection coll = new ImplicitLinkedHashCollection<>(5);
+        assertEquals(11, coll.numSlots());
+        assertTrue(coll.add(new TestElement(11)));
+        assertTrue(coll.add(new TestElement(0)));
+        assertTrue(coll.add(new TestElement(22)));
+        assertTrue(coll.add(new TestElement(33)));
+        assertEquals(11, coll.numSlots());
+        expectTraversal(coll.iterator(), 11, 0, 22, 33);
+        assertTrue(coll.remove(new TestElement(22)));
+        expectTraversal(coll.iterator(), 11, 0, 33);
+        assertEquals(3, coll.size());
+        assertFalse(coll.isEmpty());
+    }
+
+    @Test
+    public void testEnlargement() {
+        ImplicitLinkedHashCollection coll = new ImplicitLinkedHashCollection<>(5);
+        assertEquals(11, coll.numSlots());
+        for (int i = 0; i < 6; i++) {
+            assertTrue(coll.add(new TestElement(i)));
+        }
+        assertEquals(23, coll.numSlots());
+        assertEquals(6, coll.size());
+        expectTraversal(coll.iterator(), 0, 1, 2, 3, 4, 5);
+        for (int i = 0; i < 6; i++) {
+            assertTrue("Failed to find element " + i, coll.contains(new TestElement(i)));
+        }
+        coll.remove(new TestElement(3));
+        assertEquals(23, coll.numSlots());
+        assertEquals(5, coll.size());
+        expectTraversal(coll.iterator(), 0, 1, 2, 4, 5);
+    }
+
+    @Test
+    public void testManyInsertsAndDeletes() {
+        Random random = new Random(123);
+        LinkedHashSet existing = new LinkedHashSet<>();
+        ImplicitLinkedHashCollection coll = new ImplicitLinkedHashCollection<>();
+        for (int i = 0; i < 100; i++) {
+            addRandomElement(random, existing, coll);
+            addRandomElement(random, existing, coll);
+            addRandomElement(random, existing, coll);
+            removeRandomElement(random, existing, coll);
+            expectTraversal(coll.iterator(), existing.iterator());
+        }
+    }
+
+    @Test
+    public void testEquals() {
+        ImplicitLinkedHashCollection coll1 = new ImplicitLinkedHashCollection<>();
+        coll1.add(new TestElement(1));
+        coll1.add(new TestElement(2));
+        coll1.add(new TestElement(3));
+
+        ImplicitLinkedHashCollection coll2 = new ImplicitLinkedHashCollection<>();
+        coll2.add(new TestElement(1));
+        coll2.add(new TestElement(2));
+        coll2.add(new TestElement(3));
+
+        ImplicitLinkedHashCollection coll3 = new ImplicitLinkedHashCollection<>();
+        coll3.add(new TestElement(1));
+        coll3.add(new TestElement(3));
+        coll3.add(new TestElement(2));
+
+        assertEquals(coll1, coll2);
+        assertNotEquals(coll1, coll3);
+        assertNotEquals(coll2, coll3);
+    }
+
+    @Test
+    public void testFindContainsRemoveOnEmptyCollection() {
+        ImplicitLinkedHashCollection coll = new ImplicitLinkedHashCollection<>();
+        assertNull(coll.find(new TestElement(2)));
+        assertFalse(coll.contains(new TestElement(2)));
+        assertFalse(coll.remove(new TestElement(2)));
+    }
+
+    private void addRandomElement(Random random, LinkedHashSet existing,
+                                  ImplicitLinkedHashCollection set) {
+        int next;
+        do {
+            next = random.nextInt();
+        } while (existing.contains(next));
+        existing.add(next);
+        set.add(new TestElement(next));
+    }
+
+    @SuppressWarnings("unlikely-arg-type")
+    private void removeRandomElement(Random random, Collection existing,
+                                     ImplicitLinkedHashCollection coll) {
+        int removeIdx = random.nextInt(existing.size());
+        Iterator iter = existing.iterator();
+        Integer element = null;
+        for (int i = 0; i <= removeIdx; i++) {
+            element = iter.next();
+        }
+        existing.remove(new TestElement(element));
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/utils/ImplicitLinkedHashMultiSetTest.java b/clients/src/test/java/org/apache/kafka/common/utils/ImplicitLinkedHashMultiCollectionTest.java
similarity index 83%
rename from clients/src/test/java/org/apache/kafka/common/utils/ImplicitLinkedHashMultiSetTest.java
rename to clients/src/test/java/org/apache/kafka/common/utils/ImplicitLinkedHashMultiCollectionTest.java
index 950deb8bfb657..ad87b55c493c4 100644
--- a/clients/src/test/java/org/apache/kafka/common/utils/ImplicitLinkedHashMultiSetTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/utils/ImplicitLinkedHashMultiCollectionTest.java
@@ -16,7 +16,7 @@
  */
 package org.apache.kafka.common.utils;
 
-import org.apache.kafka.common.utils.ImplicitLinkedHashSetTest.TestElement;
+import org.apache.kafka.common.utils.ImplicitLinkedHashCollectionTest.TestElement;
 import org.junit.Assert;
 import org.junit.Rule;
 import org.junit.Test;
@@ -28,25 +28,35 @@
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
 /**
- * A unit test for ImplicitLinkedHashMultiSet.
+ * A unit test for ImplicitLinkedHashMultiCollection.
  */
-public class ImplicitLinkedHashMultiSetTest {
+public class ImplicitLinkedHashMultiCollectionTest {
     @Rule
     final public Timeout globalTimeout = Timeout.millis(120000);
 
     @Test
     public void testNullForbidden() {
-        ImplicitLinkedHashMultiSet multiSet = new ImplicitLinkedHashMultiSet<>();
+        ImplicitLinkedHashMultiCollection multiSet = new ImplicitLinkedHashMultiCollection<>();
         assertFalse(multiSet.add(null));
     }
 
+    @Test
+    public void testFindFindAllContainsRemoveOnEmptyCollection() {
+        ImplicitLinkedHashMultiCollection coll = new ImplicitLinkedHashMultiCollection<>();
+        assertNull(coll.find(new TestElement(2)));
+        assertFalse(coll.contains(new TestElement(2)));
+        assertFalse(coll.remove(new TestElement(2)));
+        assertTrue(coll.findAll(new TestElement(2)).isEmpty());
+    }
+
     @Test
     public void testInsertDelete() {
-        ImplicitLinkedHashMultiSet multiSet = new ImplicitLinkedHashMultiSet<>(100);
+        ImplicitLinkedHashMultiCollection multiSet = new ImplicitLinkedHashMultiCollection<>(100);
         TestElement e1 = new TestElement(1);
         TestElement e2 = new TestElement(1);
         TestElement e3 = new TestElement(2);
@@ -64,7 +74,7 @@ public void testInsertDelete() {
 
     @Test
     public void testTraversal() {
-        ImplicitLinkedHashMultiSet multiSet = new ImplicitLinkedHashMultiSet<>();
+        ImplicitLinkedHashMultiCollection multiSet = new ImplicitLinkedHashMultiCollection<>();
         expectExactTraversal(multiSet.iterator());
         TestElement e1 = new TestElement(1);
         TestElement e2 = new TestElement(1);
@@ -96,7 +106,7 @@ static void expectExactTraversal(Iterator iterator, TestElement...
 
     @Test
     public void testEnlargement() {
-        ImplicitLinkedHashMultiSet multiSet = new ImplicitLinkedHashMultiSet<>(5);
+        ImplicitLinkedHashMultiCollection multiSet = new ImplicitLinkedHashMultiCollection<>(5);
         assertEquals(11, multiSet.numSlots());
         TestElement[] testElements = {
             new TestElement(100),
@@ -126,7 +136,7 @@ public void testEnlargement() {
     public void testManyInsertsAndDeletes() {
         Random random = new Random(123);
         LinkedList existing = new LinkedList<>();
-        ImplicitLinkedHashMultiSet multiSet = new ImplicitLinkedHashMultiSet<>();
+        ImplicitLinkedHashMultiCollection multiSet = new ImplicitLinkedHashMultiCollection<>();
         for (int i = 0; i < 100; i++) {
             for (int j = 0; j < 4; j++) {
                 TestElement testElement = new TestElement(random.nextInt());
diff --git a/clients/src/test/java/org/apache/kafka/common/utils/ImplicitLinkedHashSetTest.java b/clients/src/test/java/org/apache/kafka/common/utils/ImplicitLinkedHashSetTest.java
deleted file mode 100644
index 156eba23c0e5d..0000000000000
--- a/clients/src/test/java/org/apache/kafka/common/utils/ImplicitLinkedHashSetTest.java
+++ /dev/null
@@ -1,245 +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.common.utils;
-
-import org.junit.Assert;
-import org.junit.Rule;
-import org.junit.Test;
-import org.junit.rules.Timeout;
-
-import java.util.ArrayList;
-import java.util.Collection;
-import java.util.Iterator;
-import java.util.LinkedHashSet;
-import java.util.List;
-import java.util.Random;
-
-import static org.junit.Assert.assertTrue;
-import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertEquals;
-
-/**
- * A unit test for ImplicitLinkedHashSet.
- */
-public class ImplicitLinkedHashSetTest {
-    @Rule
-    final public Timeout globalTimeout = Timeout.millis(120000);
-
-    final static class TestElement implements ImplicitLinkedHashSet.Element {
-        private int prev = ImplicitLinkedHashSet.INVALID_INDEX;
-        private int next = ImplicitLinkedHashSet.INVALID_INDEX;
-        private final int val;
-
-        TestElement(int val) {
-            this.val = val;
-        }
-
-        @Override
-        public int prev() {
-            return prev;
-        }
-
-        @Override
-        public void setPrev(int prev) {
-            this.prev = prev;
-        }
-
-        @Override
-        public int next() {
-            return next;
-        }
-
-        @Override
-        public void setNext(int next) {
-            this.next = next;
-        }
-
-        @Override
-        public boolean equals(Object o) {
-            if (this == o) return true;
-            if ((o == null) || (o.getClass() != TestElement.class)) return false;
-            TestElement that = (TestElement) o;
-            return val == that.val;
-        }
-
-        @Override
-        public String toString() {
-            return "TestElement(" + val + ")";
-        }
-
-        @Override
-        public int hashCode() {
-            return val;
-        }
-    }
-
-    @Test
-    public void testNullForbidden() {
-        ImplicitLinkedHashMultiSet multiSet = new ImplicitLinkedHashMultiSet<>();
-        assertFalse(multiSet.add(null));
-    }
-
-    @Test
-    public void testInsertDelete() {
-        ImplicitLinkedHashSet set = new ImplicitLinkedHashSet<>(100);
-        assertTrue(set.add(new TestElement(1)));
-        TestElement second = new TestElement(2);
-        assertTrue(set.add(second));
-        assertTrue(set.add(new TestElement(3)));
-        assertFalse(set.add(new TestElement(3)));
-        assertEquals(3, set.size());
-        assertTrue(set.contains(new TestElement(1)));
-        assertFalse(set.contains(new TestElement(4)));
-        TestElement secondAgain = set.find(new TestElement(2));
-        assertTrue(second == secondAgain);
-        assertTrue(set.remove(new TestElement(1)));
-        assertFalse(set.remove(new TestElement(1)));
-        assertEquals(2, set.size());
-        set.clear();
-        assertEquals(0, set.size());
-    }
-
-    static void expectTraversal(Iterator iterator, Integer... sequence) {
-        int i = 0;
-        while (iterator.hasNext()) {
-            TestElement element = iterator.next();
-            Assert.assertTrue("Iterator yieled " + (i + 1) + " elements, but only " +
-                sequence.length + " were expected.", i < sequence.length);
-            Assert.assertEquals("Iterator value number " + (i + 1) + " was incorrect.",
-                sequence[i].intValue(), element.val);
-            i = i + 1;
-        }
-        Assert.assertTrue("Iterator yieled " + (i + 1) + " elements, but " +
-            sequence.length + " were expected.", i == sequence.length);
-    }
-
-    static void expectTraversal(Iterator iter, Iterator expectedIter) {
-        int i = 0;
-        while (iter.hasNext()) {
-            TestElement element = iter.next();
-            Assert.assertTrue("Iterator yieled " + (i + 1) + " elements, but only " +
-                i + " were expected.", expectedIter.hasNext());
-            Integer expected = expectedIter.next();
-            Assert.assertEquals("Iterator value number " + (i + 1) + " was incorrect.",
-                expected.intValue(), element.val);
-            i = i + 1;
-        }
-        Assert.assertFalse("Iterator yieled " + i + " elements, but at least " +
-            (i + 1) + " were expected.", expectedIter.hasNext());
-    }
-
-    @Test
-    public void testTraversal() {
-        ImplicitLinkedHashSet set = new ImplicitLinkedHashSet<>();
-        expectTraversal(set.iterator());
-        assertTrue(set.add(new TestElement(2)));
-        expectTraversal(set.iterator(), 2);
-        assertTrue(set.add(new TestElement(1)));
-        expectTraversal(set.iterator(), 2, 1);
-        assertTrue(set.add(new TestElement(100)));
-        expectTraversal(set.iterator(), 2, 1, 100);
-        assertTrue(set.remove(new TestElement(1)));
-        expectTraversal(set.iterator(), 2, 100);
-        assertTrue(set.add(new TestElement(1)));
-        expectTraversal(set.iterator(), 2, 100, 1);
-        Iterator iter = set.iterator();
-        iter.next();
-        iter.next();
-        iter.remove();
-        iter.next();
-        assertFalse(iter.hasNext());
-        expectTraversal(set.iterator(), 2, 1);
-        List list = new ArrayList<>();
-        list.add(new TestElement(1));
-        list.add(new TestElement(2));
-        assertTrue(set.removeAll(list));
-        assertFalse(set.removeAll(list));
-        expectTraversal(set.iterator());
-        assertEquals(0, set.size());
-        assertTrue(set.isEmpty());
-    }
-
-    @Test
-    public void testCollisions() {
-        ImplicitLinkedHashSet set = new ImplicitLinkedHashSet<>(5);
-        assertEquals(11, set.numSlots());
-        assertTrue(set.add(new TestElement(11)));
-        assertTrue(set.add(new TestElement(0)));
-        assertTrue(set.add(new TestElement(22)));
-        assertTrue(set.add(new TestElement(33)));
-        assertEquals(11, set.numSlots());
-        expectTraversal(set.iterator(), 11, 0, 22, 33);
-        assertTrue(set.remove(new TestElement(22)));
-        expectTraversal(set.iterator(), 11, 0, 33);
-        assertEquals(3, set.size());
-        assertFalse(set.isEmpty());
-    }
-
-    @Test
-    public void testEnlargement() {
-        ImplicitLinkedHashSet set = new ImplicitLinkedHashSet<>(5);
-        assertEquals(11, set.numSlots());
-        for (int i = 0; i < 6; i++) {
-            assertTrue(set.add(new TestElement(i)));
-        }
-        assertEquals(23, set.numSlots());
-        assertEquals(6, set.size());
-        expectTraversal(set.iterator(), 0, 1, 2, 3, 4, 5);
-        for (int i = 0; i < 6; i++) {
-            assertTrue("Failed to find element " + i, set.contains(new TestElement(i)));
-        }
-        set.remove(new TestElement(3));
-        assertEquals(23, set.numSlots());
-        assertEquals(5, set.size());
-        expectTraversal(set.iterator(), 0, 1, 2, 4, 5);
-    }
-
-    @Test
-    public void testManyInsertsAndDeletes() {
-        Random random = new Random(123);
-        LinkedHashSet existing = new LinkedHashSet<>();
-        ImplicitLinkedHashSet set = new ImplicitLinkedHashSet<>();
-        for (int i = 0; i < 100; i++) {
-            addRandomElement(random, existing, set);
-            addRandomElement(random, existing, set);
-            addRandomElement(random, existing, set);
-            removeRandomElement(random, existing, set);
-            expectTraversal(set.iterator(), existing.iterator());
-        }
-    }
-
-    private void addRandomElement(Random random, LinkedHashSet existing,
-                                  ImplicitLinkedHashSet set) {
-        int next;
-        do {
-            next = random.nextInt();
-        } while (existing.contains(next));
-        existing.add(next);
-        set.add(new TestElement(next));
-    }
-
-    private void removeRandomElement(Random random, Collection existing,
-                             ImplicitLinkedHashSet set) {
-        int removeIdx = random.nextInt(existing.size());
-        Iterator iter = existing.iterator();
-        Integer element = null;
-        for (int i = 0; i <= removeIdx; i++) {
-            element = iter.next();
-        }
-        existing.remove(new TestElement(element));
-    }
-}
diff --git a/clients/src/test/java/org/apache/kafka/common/utils/MappedIteratorTest.java b/clients/src/test/java/org/apache/kafka/common/utils/MappedIteratorTest.java
new file mode 100644
index 0000000000000..729bb25f70f8d
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/utils/MappedIteratorTest.java
@@ -0,0 +1,61 @@
+/*
+ * 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.utils;
+
+import org.junit.Test;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+
+import static java.util.Arrays.asList;
+import static java.util.Collections.emptyList;
+import static org.junit.Assert.assertEquals;
+
+public class MappedIteratorTest {
+
+    @Test
+    public void testStringToInteger() {
+        List list = asList("foo", "", "bar2", "baz45");
+        Function mapper = s -> s.length();
+
+        Iterable mappedIterable = () -> new MappedIterator<>(list.iterator(), mapper);
+        List mapped = new ArrayList<>();
+        mappedIterable.forEach(mapped::add);
+
+        assertEquals(list.stream().map(mapper).collect(Collectors.toList()), mapped);
+
+        // Ensure that we can iterate a second time
+        List mapped2 = new ArrayList<>();
+        mappedIterable.forEach(mapped2::add);
+        assertEquals(mapped, mapped2);
+    }
+
+    @Test
+    public void testEmptyList() {
+        List list = emptyList();
+        Function mapper = s -> s.length();
+
+        Iterable mappedIterable = () -> new MappedIterator<>(list.iterator(), mapper);
+        List mapped = new ArrayList<>();
+        mappedIterable.forEach(mapped::add);
+
+        assertEquals(emptyList(), mapped);
+    }
+
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/utils/SecurityUtilsTest.java b/clients/src/test/java/org/apache/kafka/common/utils/SecurityUtilsTest.java
index 273c13abbfe08..d182a00c21c2c 100644
--- a/clients/src/test/java/org/apache/kafka/common/utils/SecurityUtilsTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/utils/SecurityUtilsTest.java
@@ -16,13 +16,49 @@
  */
 package org.apache.kafka.common.utils;
 
+import org.apache.kafka.common.config.SecurityConfig;
+import org.apache.kafka.common.security.auth.SecurityProviderCreator;
 import org.apache.kafka.common.security.auth.KafkaPrincipal;
+import org.apache.kafka.common.security.ssl.mock.TestPlainSaslServerProviderCreator;
+import org.apache.kafka.common.security.ssl.mock.TestScramSaslServerProviderCreator;
+import org.hamcrest.MatcherAssert;
+import org.junit.After;
+import org.junit.Before;
 import org.junit.Test;
 
+import java.security.Provider;
+import java.security.Security;
+import java.util.HashMap;
+import java.util.Map;
+
 import static org.junit.Assert.assertEquals;
 
 public class SecurityUtilsTest {
 
+    private SecurityProviderCreator testScramSaslServerProviderCreator = new TestScramSaslServerProviderCreator();
+    private SecurityProviderCreator testPlainSaslServerProviderCreator = new TestPlainSaslServerProviderCreator();
+
+    private Provider testScramSaslServerProvider = testScramSaslServerProviderCreator.getProvider();
+    private Provider testPlainSaslServerProvider = testPlainSaslServerProviderCreator.getProvider();
+
+    private void clearTestProviders() {
+        Security.removeProvider(testScramSaslServerProvider.getName());
+        Security.removeProvider(testPlainSaslServerProvider.getName());
+    }
+
+    @Before
+    // Remove the providers if already added
+    public void setUp() {
+        clearTestProviders();
+    }
+
+    // Remove the providers after running test cases
+    @After
+    public void tearDown() {
+        clearTestProviders();
+    }
+
+
     @Test
     public void testPrincipalNameCanContainSeparator() {
         String name = "name:with:separator:in:it";
@@ -40,4 +76,31 @@ public void testParseKafkaPrincipalWithNonUserPrincipalType() {
         assertEquals(name, principal.getName());
     }
 
+    private int getProviderIndexFromName(String providerName, Provider[] providers) {
+        for (int index = 0; index < providers.length; index++) {
+            if (providers[index].getName().equals(providerName)) {
+                return index;
+            }
+        }
+        return -1;
+    }
+
+    // Tests if the custom providers configured are being added to the JVM correctly. These providers are
+    // expected to be added at the start of the list of available providers and with the relative ordering maintained
+    @Test
+    public void testAddCustomSecurityProvider() {
+        String customProviderClasses = testScramSaslServerProviderCreator.getClass().getName() + "," +
+                testPlainSaslServerProviderCreator.getClass().getName();
+        Map configs = new HashMap<>();
+        configs.put(SecurityConfig.SECURITY_PROVIDERS_CONFIG, customProviderClasses);
+        SecurityUtils.addConfiguredSecurityProviders(configs);
+
+        Provider[] providers = Security.getProviders();
+        int testScramSaslServerProviderIndex = getProviderIndexFromName(testScramSaslServerProvider.getName(), providers);
+        int testPlainSaslServerProviderIndex = getProviderIndexFromName(testPlainSaslServerProvider.getName(), providers);
+
+        // validations
+        MatcherAssert.assertThat(testScramSaslServerProvider.getName() + " testProvider not found at expected index", testScramSaslServerProviderIndex == 0);
+        MatcherAssert.assertThat(testPlainSaslServerProvider.getName() + " testProvider not found at expected index", testPlainSaslServerProviderIndex == 1);
+    }
 }
diff --git a/clients/src/test/java/org/apache/kafka/common/utils/UtilsTest.java b/clients/src/test/java/org/apache/kafka/common/utils/UtilsTest.java
index 172a992b1bda9..90cc7accc0676 100755
--- a/clients/src/test/java/org/apache/kafka/common/utils/UtilsTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/utils/UtilsTest.java
@@ -32,6 +32,7 @@
 import java.nio.file.StandardOpenOption;
 import java.util.Collections;
 import java.util.HashSet;
+import java.util.Map;
 import java.util.Random;
 import java.util.Set;
 import java.util.stream.Collectors;
@@ -43,6 +44,7 @@
 import static org.apache.kafka.common.utils.Utils.getHost;
 import static org.apache.kafka.common.utils.Utils.getPort;
 import static org.apache.kafka.common.utils.Utils.mkSet;
+import static org.apache.kafka.common.utils.Utils.murmur2;
 import static org.apache.kafka.common.utils.Utils.validHostPattern;
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
@@ -58,6 +60,21 @@
 
 public class UtilsTest {
 
+    @Test
+    public void testMurmur2() {
+        Map cases = new java.util.HashMap<>();
+        cases.put("21".getBytes(), -973932308);
+        cases.put("foobar".getBytes(), -790332482);
+        cases.put("a-little-bit-long-string".getBytes(), -985981536);
+        cases.put("a-little-bit-longer-string".getBytes(), -1486304829);
+        cases.put("lkjh234lh9fiuh90y23oiuhsafujhadof229phr9h19h89h8".getBytes(), -58897971);
+        cases.put(new byte[]{'a', 'b', 'c'}, 479470107);
+
+        for (Map.Entry c : cases.entrySet()) {
+            assertEquals((int) c.getValue(), murmur2((byte[]) c.getKey()));
+        }
+    }
+
     @Test
     public void testGetHost() {
         assertEquals("127.0.0.1", getHost("127.0.0.1:8000"));
diff --git a/clients/src/test/java/org/apache/kafka/test/MetricsBench.java b/clients/src/test/java/org/apache/kafka/test/MetricsBench.java
index 379db2f526d1d..93cbf6d2a8910 100644
--- a/clients/src/test/java/org/apache/kafka/test/MetricsBench.java
+++ b/clients/src/test/java/org/apache/kafka/test/MetricsBench.java
@@ -21,11 +21,11 @@
 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.Count;
 import org.apache.kafka.common.metrics.stats.Max;
 import org.apache.kafka.common.metrics.stats.Percentile;
 import org.apache.kafka.common.metrics.stats.Percentiles;
 import org.apache.kafka.common.metrics.stats.Percentiles.BucketSizing;
+import org.apache.kafka.common.metrics.stats.WindowedCount;
 
 public class MetricsBench {
 
@@ -37,7 +37,7 @@ public static void main(String[] args) {
             Sensor child = metrics.sensor("child", parent);
             for (Sensor sensor : Arrays.asList(parent, child)) {
                 sensor.add(metrics.metricName(sensor.name() + ".avg", "grp1"), new Avg());
-                sensor.add(metrics.metricName(sensor.name() + ".count", "grp1"), new Count());
+                sensor.add(metrics.metricName(sensor.name() + ".count", "grp1"), new WindowedCount());
                 sensor.add(metrics.metricName(sensor.name() + ".max", "grp1"), new Max());
                 sensor.add(new Percentiles(1024,
                         0.0,
diff --git a/clients/src/test/java/org/apache/kafka/test/MockSelector.java b/clients/src/test/java/org/apache/kafka/test/MockSelector.java
index 200d51115177c..90a83b0adc7cb 100644
--- a/clients/src/test/java/org/apache/kafka/test/MockSelector.java
+++ b/clients/src/test/java/org/apache/kafka/test/MockSelector.java
@@ -16,9 +16,9 @@
  */
 package org.apache.kafka.test;
 
+import org.apache.kafka.common.errors.AuthenticationException;
 import org.apache.kafka.common.network.ChannelState;
 import org.apache.kafka.common.network.NetworkReceive;
-import org.apache.kafka.common.network.NetworkSend;
 import org.apache.kafka.common.network.Selectable;
 import org.apache.kafka.common.network.Send;
 import org.apache.kafka.common.requests.ByteBufferChannel;
@@ -41,6 +41,7 @@ public class MockSelector implements Selectable {
     private final Time time;
     private final List initiatedSends = new ArrayList<>();
     private final List completedSends = new ArrayList<>();
+    private final List completedSendBuffers = new ArrayList<>();
     private final List completedReceives = new ArrayList<>();
     private final Map disconnected = new HashMap<>();
     private final List connected = new ArrayList<>();
@@ -87,18 +88,21 @@ public void serverDisconnect(String id) {
         close(id);
     }
 
+    public void serverAuthenticationFailed(String id) {
+        ChannelState authFailed = new ChannelState(ChannelState.State.AUTHENTICATION_FAILED,
+                new AuthenticationException("Authentication failed"), null);
+        this.disconnected.put(id, authFailed);
+        close(id);
+    }
+
     private void removeSendsForNode(String id, Collection sends) {
-        Iterator iter = sends.iterator();
-        while (iter.hasNext()) {
-            Send send = iter.next();
-            if (id.equals(send.destination()))
-                iter.remove();
-        }
+        sends.removeIf(send -> id.equals(send.destination()));
     }
 
     public void clear() {
         this.completedSends.clear();
         this.completedReceives.clear();
+        this.completedSendBuffers.clear();
         this.disconnected.clear();
         this.connected.clear();
     }
@@ -129,6 +133,7 @@ private void completeSend(Send send) throws IOException {
                 send.writeTo(discardChannel);
             }
             completedSends.add(send);
+            completedSendBuffers.add(discardChannel);
         }
     }
 
@@ -150,8 +155,8 @@ public List completedSends() {
         return completedSends;
     }
 
-    public void completeSend(NetworkSend send) {
-        this.completedSends.add(send);
+    public List completedSendBuffers() {
+        return completedSendBuffers;
     }
 
     @Override
diff --git a/clients/src/test/java/org/apache/kafka/test/TestCondition.java b/clients/src/test/java/org/apache/kafka/test/TestCondition.java
index e7b78cf71e6a7..b153c75938b34 100644
--- a/clients/src/test/java/org/apache/kafka/test/TestCondition.java
+++ b/clients/src/test/java/org/apache/kafka/test/TestCondition.java
@@ -23,5 +23,5 @@
 @FunctionalInterface
 public interface TestCondition {
 
-    boolean conditionMet();
+    boolean conditionMet() throws InterruptedException;
 }
diff --git a/clients/src/test/java/org/apache/kafka/test/TestSslUtils.java b/clients/src/test/java/org/apache/kafka/test/TestSslUtils.java
index b2de0e6a2b51c..bcdcb584e92bd 100644
--- a/clients/src/test/java/org/apache/kafka/test/TestSslUtils.java
+++ b/clients/src/test/java/org/apache/kafka/test/TestSslUtils.java
@@ -44,6 +44,8 @@
 import java.security.cert.X509Certificate;
 
 import org.apache.kafka.common.config.types.Password;
+import org.apache.kafka.common.security.ssl.BoringSslContextProvider;
+import org.apache.kafka.common.security.ssl.SimpleSslContextProvider;
 import org.bouncycastle.asn1.DEROctetString;
 import org.bouncycastle.asn1.x500.X500Name;
 import org.bouncycastle.asn1.x509.AlgorithmIdentifier;
@@ -70,6 +72,11 @@
 
 public class TestSslUtils {
 
+    public static final String TRUST_STORE_PASSWORD = "TrustStorePassword";
+    public enum SSLProvider {
+        DEFAULT, OPENSSL
+    }
+
     /**
      * Create a self-signed X.509 Certificate.
      * From http://bfo.com/blog/2011/03/08/odds_and_ends_creating_a_new_x_509_certificate.html.
@@ -151,7 +158,7 @@ public static  void createTrustStore(
     }
 
     private static Map createSslConfig(Mode mode, File keyStoreFile, Password password, Password keyPassword,
-                                                       File trustStoreFile, Password trustStorePassword) {
+                                                       File trustStoreFile, Password trustStorePassword, SSLProvider provider) {
         Map sslConfigs = new HashMap<>();
         sslConfigs.put(SslConfigs.SSL_PROTOCOL_CONFIG, "TLSv1.2"); // protocol to create SSLContext
 
@@ -168,6 +175,26 @@ private static Map createSslConfig(Mode mode, File keyStoreFile,
         sslConfigs.put(SslConfigs.SSL_TRUSTSTORE_TYPE_CONFIG, "JKS");
         sslConfigs.put(SslConfigs.SSL_TRUSTMANAGER_ALGORITHM_CONFIG, TrustManagerFactory.getDefaultAlgorithm());
 
+        List enabledProtocols  = new ArrayList<>();
+        enabledProtocols.add("TLSv1.2");
+        sslConfigs.put(SslConfigs.SSL_ENABLED_PROTOCOLS_CONFIG, enabledProtocols);
+        if (provider.equals(SSLProvider.OPENSSL)) {
+            sslConfigs.put(SslConfigs.SSL_CONTEXT_PROVIDER_CLASS_CONFIG, BoringSslContextProvider.class.getName());
+        } else {
+            sslConfigs.put(SslConfigs.SSL_CONTEXT_PROVIDER_CLASS_CONFIG, SimpleSslContextProvider.class.getName());
+        }
+
+        return sslConfigs;
+    }
+
+    public static Map createSslConfig(String keyManagerAlgorithm, String trustManagerAlgorithm) {
+        Map sslConfigs = new HashMap<>();
+        sslConfigs.put(SslConfigs.SSL_PROTOCOL_CONFIG, "TLSv1.2"); // protocol to create SSLContext
+        sslConfigs.put(SslConfigs.SSL_CONTEXT_PROVIDER_CLASS_CONFIG, SimpleSslContextProvider.class.getName());
+
+        sslConfigs.put(SslConfigs.SSL_KEYMANAGER_ALGORITHM_CONFIG, keyManagerAlgorithm);
+        sslConfigs.put(SslConfigs.SSL_TRUSTMANAGER_ALGORITHM_CONFIG, trustManagerAlgorithm);
+
         List enabledProtocols  = new ArrayList<>();
         enabledProtocols.add("TLSv1.2");
         sslConfigs.put(SslConfigs.SSL_ENABLED_PROTOCOLS_CONFIG, enabledProtocols);
@@ -183,17 +210,17 @@ public static  Map createSslConfig(boolean useClientCert, boolea
     public static  Map createSslConfig(boolean useClientCert, boolean trustStore,
             Mode mode, File trustStoreFile, String certAlias, String cn)
         throws IOException, GeneralSecurityException {
-        return createSslConfig(useClientCert, trustStore, mode, trustStoreFile, certAlias, cn, new CertificateBuilder());
+        return createSslConfig(useClientCert, trustStore, mode, trustStoreFile, certAlias, cn, new CertificateBuilder(), SSLProvider.DEFAULT);
     }
 
     public static  Map createSslConfig(boolean useClientCert, boolean trustStore,
-            Mode mode, File trustStoreFile, String certAlias, String cn, CertificateBuilder certBuilder)
+            Mode mode, File trustStoreFile, String certAlias, String cn, CertificateBuilder certBuilder, SSLProvider provider)
             throws IOException, GeneralSecurityException {
         Map certs = new HashMap<>();
         File keyStoreFile = null;
         Password password = mode == Mode.SERVER ? new Password("ServerPassword") : new Password("ClientPassword");
 
-        Password trustStorePassword = new Password("TrustStorePassword");
+        Password trustStorePassword = new Password(TRUST_STORE_PASSWORD);
 
         if (mode == Mode.CLIENT && useClientCert) {
             keyStoreFile = File.createTempFile("clientKS", ".jks");
@@ -216,7 +243,7 @@ public static  Map createSslConfig(boolean useClientCert, boolea
             trustStoreFile.deleteOnExit();
         }
 
-        return createSslConfig(mode, keyStoreFile, password, password, trustStoreFile, trustStorePassword);
+        return createSslConfig(mode, keyStoreFile, password, password, trustStoreFile, trustStorePassword, provider);
     }
 
     public static class CertificateBuilder {
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 3f9a1b76711f3..57db41335cfd6 100644
--- a/clients/src/test/java/org/apache/kafka/test/TestUtils.java
+++ b/clients/src/test/java/org/apache/kafka/test/TestUtils.java
@@ -21,6 +21,7 @@
 import org.apache.kafka.common.Cluster;
 import org.apache.kafka.common.Node;
 import org.apache.kafka.common.PartitionInfo;
+import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.internals.Topic;
 import org.apache.kafka.common.network.NetworkReceive;
 import org.apache.kafka.common.protocol.ApiKeys;
@@ -34,11 +35,13 @@
 
 import java.io.File;
 import java.io.IOException;
+import java.lang.reflect.Field;
 import java.nio.ByteBuffer;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.ArrayList;
 import java.util.Base64;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
@@ -52,13 +55,17 @@
 import java.util.UUID;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Future;
+import java.util.function.Consumer;
+import java.util.function.Function;
 import java.util.function.Supplier;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
 import static java.util.Arrays.asList;
+import static org.hamcrest.MatcherAssert.assertThat;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
@@ -77,6 +84,7 @@ public class TestUtils {
     /* A consistent random number generator to make tests repeatable */
     public static final Random SEEDED_RANDOM = new Random(192348092834L);
     public static final Random RANDOM = new Random();
+    public static final long DEFAULT_POLL_INTERVAL_MS = 100;
     public static final long DEFAULT_MAX_WAIT_MS = 15000;
 
     public static Cluster singletonCluster() {
@@ -102,7 +110,7 @@ public static Cluster clusterWith(final int nodes, final Map to
             for (int i = 0; i < partitions; i++)
                 parts.add(new PartitionInfo(topic, i, ns[i % ns.length], ns, ns));
         }
-        return new Cluster("kafka-cluster", asList(ns), parts, Collections.emptySet(), Topic.INTERNAL_TOPICS);
+        return new Cluster("kafka-cluster", asList(ns), parts, Collections.emptySet(), Collections.emptySet());
     }
 
     public static MetadataResponse metadataUpdateWith(final int numNodes,
@@ -113,20 +121,29 @@ public static MetadataResponse metadataUpdateWith(final int numNodes,
     public static MetadataResponse metadataUpdateWith(final String clusterId,
                                                       final int numNodes,
                                                       final Map topicPartitionCounts) {
-        return metadataUpdateWith(clusterId, numNodes, Collections.emptyMap(), topicPartitionCounts, MetadataResponse.PartitionMetadata::new);
+        return metadataUpdateWith(clusterId, numNodes, Collections.emptyMap(), topicPartitionCounts, tp -> null, MetadataResponse.PartitionMetadata::new);
     }
 
     public static MetadataResponse metadataUpdateWith(final String clusterId,
                                                       final int numNodes,
                                                       final Map topicErrors,
                                                       final Map topicPartitionCounts) {
-        return metadataUpdateWith(clusterId, numNodes, topicErrors, topicPartitionCounts, MetadataResponse.PartitionMetadata::new);
+        return metadataUpdateWith(clusterId, numNodes, topicErrors, topicPartitionCounts, tp -> null, MetadataResponse.PartitionMetadata::new);
     }
 
     public static MetadataResponse metadataUpdateWith(final String clusterId,
                                                       final int numNodes,
                                                       final Map topicErrors,
                                                       final Map topicPartitionCounts,
+                                                      final Function epochSupplier) {
+        return metadataUpdateWith(clusterId, numNodes, topicErrors, topicPartitionCounts, epochSupplier, MetadataResponse.PartitionMetadata::new);
+    }
+
+    public static MetadataResponse metadataUpdateWith(final String clusterId,
+                                                      final int numNodes,
+                                                      final Map topicErrors,
+                                                      final Map topicPartitionCounts,
+                                                      final Function epochSupplier,
                                                       final PartitionMetadataSupplier partitionSupplier) {
         final List nodes = new ArrayList<>(numNodes);
         for (int i = 0; i < numNodes; i++)
@@ -139,10 +156,11 @@ public static MetadataResponse metadataUpdateWith(final String clusterId,
 
             List partitionMetadata = new ArrayList<>(numPartitions);
             for (int i = 0; i < numPartitions; i++) {
+                TopicPartition tp = new TopicPartition(topic, i);
                 Node leader = nodes.get(i % nodes.size());
                 List replicas = Collections.singletonList(leader);
                 partitionMetadata.add(partitionSupplier.supply(
-                        Errors.NONE, i, leader, Optional.empty(), replicas, replicas, replicas));
+                        Errors.NONE, i, leader, Optional.ofNullable(epochSupplier.apply(tp)), replicas, replicas, replicas));
             }
 
             topicMetadata.add(new MetadataResponse.TopicMetadata(Errors.NONE, topic,
@@ -155,7 +173,7 @@ public static MetadataResponse metadataUpdateWith(final String clusterId,
                     Topic.isInternal(topic), Collections.emptyList()));
         }
 
-        return new MetadataResponse(nodes, clusterId, 0, topicMetadata);
+        return MetadataResponse.prepareResponse(nodes, clusterId, 0, topicMetadata);
     }
 
     @FunctionalInterface
@@ -337,7 +355,7 @@ public static void waitForCondition(final TestCondition testCondition, final Sup
     public static void waitForCondition(final TestCondition testCondition, final long maxWaitMs, String conditionDetails) throws InterruptedException {
         waitForCondition(testCondition, maxWaitMs, () -> conditionDetails);
     }
-    
+
     /**
      * Wait for condition to be met for at most {@code maxWaitMs} and throw assertion failure otherwise.
      * This should be used instead of {@code Thread.sleep} whenever possible as it allows a longer timeout to be used
@@ -345,20 +363,69 @@ public static void waitForCondition(final TestCondition testCondition, final lon
      * avoid transient failures due to slow or overloaded machines.
      */
     public static void waitForCondition(final TestCondition testCondition, final long maxWaitMs, Supplier conditionDetailsSupplier) throws InterruptedException {
-        final long startTime = System.currentTimeMillis();
+        String conditionDetailsSupplied = conditionDetailsSupplier != null ? conditionDetailsSupplier.get() : null;
+        String conditionDetails = conditionDetailsSupplied != null ? conditionDetailsSupplied : "";
+        retryOnExceptionWithTimeout(maxWaitMs, () -> {
+            assertThat("Condition not met within timeout " + maxWaitMs + ". " + conditionDetails,
+                testCondition.conditionMet());
+        });
+    }
 
-        boolean testConditionMet;
-        while (!(testConditionMet = testCondition.conditionMet()) && ((System.currentTimeMillis() - startTime) < maxWaitMs)) {
-            Thread.sleep(Math.min(maxWaitMs, 100L));
-        }
+    /**
+     * Wait for the given runnable to complete successfully, i.e. throw now {@link Exception}s or
+     * {@link AssertionError}s, or for the given timeout to expire. If the timeout expires then the
+     * last exception or assertion failure will be thrown thus providing context for the failure.
+     *
+     * @param timeoutMs the total time in milliseconds to wait for {@code runnable} to complete successfully.
+     * @param runnable the code to attempt to execute successfully.
+     * @throws InterruptedException if the current thread is interrupted while waiting for {@code runnable} to complete successfully.
+     */
+    public static void retryOnExceptionWithTimeout(final long timeoutMs,
+                                                   final ValuelessCallable runnable) throws InterruptedException {
+        retryOnExceptionWithTimeout(DEFAULT_POLL_INTERVAL_MS, timeoutMs, runnable);
+    }
 
-        // don't re-evaluate testCondition.conditionMet() because this might slow down some tests significantly (this
-        // could be avoided by making the implementations more robust, but we have a large number of such implementations
-        // and it's easier to simply avoid the issue altogether)
-        if (!testConditionMet) {
-            String conditionDetailsSupplied = conditionDetailsSupplier != null ? conditionDetailsSupplier.get() : null;
-            String conditionDetails = conditionDetailsSupplied != null ? conditionDetailsSupplied : "";
-            throw new AssertionError("Condition not met within timeout " + maxWaitMs + ". " + conditionDetails);
+    /**
+     * Wait for the given runnable to complete successfully, i.e. throw now {@link Exception}s or
+     * {@link AssertionError}s, or for the default timeout to expire. If the timeout expires then the
+     * last exception or assertion failure will be thrown thus providing context for the failure.
+     *
+     * @param runnable the code to attempt to execute successfully.
+     * @throws InterruptedException if the current thread is interrupted while waiting for {@code runnable} to complete successfully.
+     */
+    public static void retryOnExceptionWithTimeout(final ValuelessCallable runnable) throws InterruptedException {
+        retryOnExceptionWithTimeout(DEFAULT_POLL_INTERVAL_MS, DEFAULT_MAX_WAIT_MS, runnable);
+    }
+
+    /**
+     * Wait for the given runnable to complete successfully, i.e. throw now {@link Exception}s or
+     * {@link AssertionError}s, or for the given timeout to expire. If the timeout expires then the
+     * last exception or assertion failure will be thrown thus providing context for the failure.
+     *
+     * @param pollIntervalMs the interval in milliseconds to wait between invoking {@code runnable}.
+     * @param timeoutMs the total time in milliseconds to wait for {@code runnable} to complete successfully.
+     * @param runnable the code to attempt to execute successfully.
+     * @throws InterruptedException if the current thread is interrupted while waiting for {@code runnable} to complete successfully.
+     */
+    public static void retryOnExceptionWithTimeout(final long pollIntervalMs,
+                                                   final long timeoutMs,
+                                                   final ValuelessCallable runnable) throws InterruptedException {
+        final long expectedEnd = System.currentTimeMillis() + timeoutMs;
+
+        while (true) {
+            try {
+                runnable.call();
+                return;
+            } catch (final AssertionError t) {
+                if (expectedEnd <= System.currentTimeMillis()) {
+                    throw t;
+                }
+            } catch (final Exception e) {
+                if (expectedEnd <= System.currentTimeMillis()) {
+                    throw new AssertionError(e);
+                }
+            }
+            Thread.sleep(Math.min(pollIntervalMs, timeoutMs));
         }
     }
 
@@ -420,6 +487,10 @@ public static  List toList(Iterable iterable) {
         return list;
     }
 
+    public static  Set toSet(Collection collection) {
+        return new HashSet<>(collection);
+    }
+
     public static ByteBuffer toBuffer(Struct struct) {
         ByteBuffer buffer = ByteBuffer.allocate(struct.sizeOf());
         struct.writeTo(buffer);
@@ -427,6 +498,33 @@ public static ByteBuffer toBuffer(Struct struct) {
         return buffer;
     }
 
+    public static Set generateRandomTopicPartitions(int numTopic, int numPartitionPerTopic) {
+        Set tps = new HashSet<>();
+        for (int i = 0; i < numTopic; i++) {
+            String topic = randomString(32);
+            for (int j = 0; j < numPartitionPerTopic; j++) {
+                tps.add(new TopicPartition(topic, j));
+            }
+        }
+        return tps;
+    }
+
+    /**
+     * Assert that a future raises an expected exception cause type. Return the exception cause
+     * if the assertion succeeds; otherwise raise AssertionError.
+     *
+     * @param future The future to await
+     * @param exceptionCauseClass Class of the expected exception cause
+     * @param  Exception cause type parameter
+     * @return The caught exception cause
+     */
+    public static  T assertFutureThrows(Future future, Class exceptionCauseClass) {
+        ExecutionException exception = assertThrows(ExecutionException.class, future::get);
+        assertTrue("Unexpected exception cause " + exception.getCause(),
+                exceptionCauseClass.isInstance(exception.getCause()));
+        return exceptionCauseClass.cast(exception.getCause());
+    }
+
     public static void assertFutureError(Future future, Class exceptionClass)
         throws InterruptedException {
         try {
@@ -443,4 +541,23 @@ public static void assertFutureError(Future future, Class void assertOptional(Optional optional, Consumer assertion) {
+        if (optional.isPresent()) {
+            assertion.accept(optional.get());
+        } else {
+            fail("Missing value from Optional");
+        }
+    }
+
+    @SuppressWarnings("unchecked")
+    public static  T fieldValue(Object o, Class clazz, String fieldName)  {
+        try {
+            Field field = clazz.getDeclaredField(fieldName);
+            field.setAccessible(true);
+            return (T) field.get(o);
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+    }
 }
diff --git a/clients/src/test/java/org/apache/kafka/test/ValuelessCallable.java b/clients/src/test/java/org/apache/kafka/test/ValuelessCallable.java
new file mode 100644
index 0000000000000..62863b3f65ffe
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/test/ValuelessCallable.java
@@ -0,0 +1,25 @@
+/*
+ * 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.test;
+
+/**
+ * Like a {@link Runnable} that allows exceptions to be thrown or a {@link java.util.concurrent.Callable}
+ * that does not return a value.
+ */
+public interface ValuelessCallable {
+    void call() throws Exception;
+}
diff --git a/clients/src/test/resources/common/message/SimpleExampleMessage.json b/clients/src/test/resources/common/message/SimpleExampleMessage.json
new file mode 100644
index 0000000000000..b20c1c090cbc1
--- /dev/null
+++ b/clients/src/test/resources/common/message/SimpleExampleMessage.json
@@ -0,0 +1,25 @@
+// 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": "SimpleExampleMessage",
+  "type": "header",
+  "validVersions": "0-1",
+  "flexibleVersions": "1+",
+  "fields": [
+    { "name": "processId", "versions": "1+", "type": "uuid" },
+    { "name": "myTaggedIntArray", "type": "[]int32",
+        "versions": "1+", "taggedVersions": "1+", "tag": 0 }
+  ]
+}
diff --git a/config/connect-log4j.properties b/config/connect-log4j.properties
index 808addba4bbe7..f695e37eb348b 100644
--- a/config/connect-log4j.properties
+++ b/config/connect-log4j.properties
@@ -13,13 +13,31 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-log4j.rootLogger=INFO, stdout
-
+log4j.rootLogger=INFO, stdout, connectAppender
 
+# Send the logs to the console.
+#
 log4j.appender.stdout=org.apache.log4j.ConsoleAppender
 log4j.appender.stdout.layout=org.apache.log4j.PatternLayout
-log4j.appender.stdout.layout.ConversionPattern=[%d] %p %m (%c:%L)%n
+
+# Send the logs to a file, rolling the file at midnight local time. For example, the `File` option specifies the
+# location of the log files (e.g. ${kafka.logs.dir}/connect.log), and at midnight local time the file is closed
+# and copied in the same directory but with a filename that ends in the `DatePattern` option.
+#
+log4j.appender.connectAppender=org.apache.log4j.DailyRollingFileAppender
+log4j.appender.connectAppender.DatePattern='.'yyyy-MM-dd-HH
+log4j.appender.connectAppender.File=${kafka.logs.dir}/connect.log
+log4j.appender.connectAppender.layout=org.apache.log4j.PatternLayout
+
+# The `%X{connector.context}` parameter in the layout includes connector-specific and task-specific information
+# in the log message, where appropriate. This makes it easier to identify those log messages that apply to a
+# specific connector. Simply add this parameter to the log layout configuration below to include the contextual information.
+#
+connect.log.pattern=[%d] %p %m (%c:%L)%n
+#connect.log.pattern=[%d] %p %X{connector.context}%m (%c:%L)%n
+
+log4j.appender.stdout.layout.ConversionPattern=${connect.log.pattern}
+log4j.appender.connectAppender.layout.ConversionPattern=${connect.log.pattern}
 
 log4j.logger.org.apache.zookeeper=ERROR
-log4j.logger.org.I0Itec.zkclient=ERROR
 log4j.logger.org.reflections=ERROR
diff --git a/config/connect-mirror-maker.properties b/config/connect-mirror-maker.properties
new file mode 100644
index 0000000000000..16c1b791d3d5f
--- /dev/null
+++ b/config/connect-mirror-maker.properties
@@ -0,0 +1,37 @@
+# Licensed to the Apache Software Foundation (ASF) under A 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.
+# see org.apache.kafka.clients.consumer.ConsumerConfig for more details
+
+# Sample MirrorMaker 2.0 top-level configuration file
+# Run with ./bin/connect-mirror-maker.sh connect-mirror-maker.properties 
+
+# specify any number of cluster aliases
+clusters = A, B, C
+
+# connection information for each cluster
+A.bootstrap.servers = A_host1:9092, A_host2:9092, A_host3:9092
+B.bootstrap.servers = B_host1:9092, B_host2:9092, B_host3:9092
+C.bootstrap.servers = C_host1:9092, C_host2:9092, C_host3:9092
+
+# enable and configure individual replication flows
+A->B.enabled = true
+A->B.topics = foo-.*
+B->C.enabled = true
+B->C.topics = bar-.*
+
+# customize as needed
+# replication.policy.separator = _
+# sync.topic.acls.enabled = false
+# emit.heartbeats.interval.seconds = 5
diff --git a/config/log4j.properties b/config/log4j.properties
index 3ff3f9e4f6c11..5de1c1903c036 100644
--- a/config/log4j.properties
+++ b/config/log4j.properties
@@ -57,8 +57,7 @@ log4j.appender.authorizerAppender.File=${kafka.logs.dir}/kafka-authorizer.log
 log4j.appender.authorizerAppender.layout=org.apache.log4j.PatternLayout
 log4j.appender.authorizerAppender.layout.ConversionPattern=[%d] %p %m (%c)%n
 
-# Change the two lines below to adjust ZK client logging
-log4j.logger.org.I0Itec.zkclient.ZkClient=INFO
+# Change the line below to adjust ZK client logging
 log4j.logger.org.apache.zookeeper=INFO
 
 # Change the two lines below to adjust the general broker logging level (output to server.log and stdout)
diff --git a/config/server.properties b/config/server.properties
index 898d8ebc28c5c..46208b1523f63 100644
--- a/config/server.properties
+++ b/config/server.properties
@@ -70,7 +70,7 @@ num.recovery.threads.per.data.dir=1
 
 ############################# Internal Topic Settings  #############################
 # The replication factor for the group metadata internal topics "__consumer_offsets" and "__transaction_state"
-# For anything other than development testing, a value greater than 1 is recommended for to ensure availability such as 3.
+# For anything other than development testing, a value greater than 1 is recommended to ensure availability such as 3.
 offsets.topic.replication.factor=1
 transaction.state.log.replication.factor=1
 transaction.state.log.min.isr=1
@@ -133,4 +133,4 @@ zookeeper.connection.timeout.ms=6000
 # The default value for this is 3 seconds.
 # We override this to 0 here as it makes for a better out-of-the-box experience for development and testing.
 # However, in production environments the default value of 3 seconds is more suitable as this will help to avoid unnecessary, and potentially expensive, rebalances during application startup.
-group.initial.rebalance.delay.ms=0
\ No newline at end of file
+group.initial.rebalance.delay.ms=0
diff --git a/config/zookeeper.properties b/config/zookeeper.properties
index 74cbf90428f81..90f4332ec31cf 100644
--- a/config/zookeeper.properties
+++ b/config/zookeeper.properties
@@ -18,3 +18,7 @@ dataDir=/tmp/zookeeper
 clientPort=2181
 # disable the per-ip limit on the number of connections since this is a non-production config
 maxClientCnxns=0
+# Disable the adminserver by default to avoid port conflicts.
+# Set the port to something non-conflicting if choosing to enable this
+admin.enableServer=false
+# admin.serverPort=8080
diff --git a/connect/api/src/main/java/org/apache/kafka/connect/connector/policy/ConnectorClientConfigOverridePolicy.java b/connect/api/src/main/java/org/apache/kafka/connect/connector/policy/ConnectorClientConfigOverridePolicy.java
new file mode 100644
index 0000000000000..94e5fd6f5b0cf
--- /dev/null
+++ b/connect/api/src/main/java/org/apache/kafka/connect/connector/policy/ConnectorClientConfigOverridePolicy.java
@@ -0,0 +1,47 @@
+/*
+ * 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.connector.policy;
+
+import org.apache.kafka.common.Configurable;
+import org.apache.kafka.common.config.ConfigValue;
+
+import java.util.List;
+
+/**
+ * 
      An interface for enforcing a policy on overriding of client configs via the connector configs.
+ *
+ * 
      Common use cases are ability to provide principal per connector, sasl.jaas.config
+ * and/or enforcing that the producer/consumer configurations for optimizations are within acceptable ranges.
+ */
+public interface ConnectorClientConfigOverridePolicy extends Configurable, AutoCloseable {
+
+
+    /**
+     * Worker will invoke this while constructing the producer for the SourceConnectors,  DLQ for SinkConnectors and the consumer for the
+     * SinkConnectors to validate if all of the overridden client configurations are allowed per the
+     * policy implementation. This would also be invoked during the validate of connector configs via the Rest API.
+     *
+     * If there are any policy violations, the connector will not be started.
+     *
+     * @param connectorClientConfigRequest an instance of {@code ConnectorClientConfigRequest} that provides the configs to overridden and
+     *                                     its context; never {@code null}
+     * @return list of {@link ConfigValue} instances that describe each client configuration in the request and includes an 
+               {@link ConfigValue#errorMessages error} if the configuration is not allowed by the policy; never null
+     */
+    List validate(ConnectorClientConfigRequest connectorClientConfigRequest);
+}
diff --git a/connect/api/src/main/java/org/apache/kafka/connect/connector/policy/ConnectorClientConfigRequest.java b/connect/api/src/main/java/org/apache/kafka/connect/connector/policy/ConnectorClientConfigRequest.java
new file mode 100644
index 0000000000000..11f756bcb41e3
--- /dev/null
+++ b/connect/api/src/main/java/org/apache/kafka/connect/connector/policy/ConnectorClientConfigRequest.java
@@ -0,0 +1,101 @@
+/*
+ * 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.connector.policy;
+
+import org.apache.kafka.connect.connector.Connector;
+import org.apache.kafka.connect.health.ConnectorType;
+
+import java.util.Map;
+
+public class ConnectorClientConfigRequest {
+
+    private Map clientProps;
+    private ClientType  clientType;
+    private String connectorName;
+    private ConnectorType connectorType;
+    private Class connectorClass;
+
+    public ConnectorClientConfigRequest(
+        String connectorName,
+        ConnectorType connectorType,
+        Class connectorClass,
+        Map clientProps,
+        ClientType clientType) {
+        this.clientProps = clientProps;
+        this.clientType = clientType;
+        this.connectorName = connectorName;
+        this.connectorType = connectorType;
+        this.connectorClass = connectorClass;
+    }
+
+    /**
+     * Provides Config with prefix {@code producer.override.} for {@link ConnectorType#SOURCE}.
+     * Provides Config with prefix {@code consumer.override.} for {@link ConnectorType#SINK}.
+     * Provides Config with prefix {@code producer.override.} for {@link ConnectorType#SINK} for DLQ.
+     * Provides Config with prefix {@code admin.override.} for {@link ConnectorType#SINK} for DLQ.
+     *
+     * @return The client properties specified in the Connector Config with prefix {@code producer.override.} ,
+     * {@code consumer.override.} and {@code admin.override.}. The configs don't include the prefixes.
+     */
+    public Map clientProps() {
+        return clientProps;
+    }
+
+    /**
+     * {@link ClientType#PRODUCER} for {@link ConnectorType#SOURCE}
+     * {@link ClientType#CONSUMER} for {@link ConnectorType#SINK}
+     * {@link ClientType#PRODUCER} for DLQ in {@link ConnectorType#SINK}
+     * {@link ClientType#ADMIN} for DLQ  Topic Creation in {@link ConnectorType#SINK}
+     *
+     * @return enumeration specifying the client type that is being overriden by the worker; never null.
+     */
+    public ClientType clientType() {
+        return clientType;
+    }
+
+    /**
+     * Name of the connector specified in the connector config.
+     *
+     * @return name of the connector; never null.
+     */
+    public String connectorName() {
+        return connectorName;
+    }
+
+    /**
+     * Type of the Connector.
+     *
+     * @return enumeration specifying the type of the connector {@link ConnectorType#SINK} or {@link ConnectorType#SOURCE}.
+     */
+    public ConnectorType connectorType() {
+        return connectorType;
+    }
+
+    /**
+     * The class of the Connector.
+     *
+     * @return the class of the Connector being created; never null
+     */
+    public Class connectorClass() {
+        return connectorClass;
+    }
+
+    public enum ClientType {
+        PRODUCER, CONSUMER, ADMIN;
+    }
+}
\ No newline at end of file
diff --git a/connect/api/src/main/java/org/apache/kafka/connect/data/SchemaBuilder.java b/connect/api/src/main/java/org/apache/kafka/connect/data/SchemaBuilder.java
index fdcf05a7ac055..722f5fc89ed5a 100644
--- a/connect/api/src/main/java/org/apache/kafka/connect/data/SchemaBuilder.java
+++ b/connect/api/src/main/java/org/apache/kafka/connect/data/SchemaBuilder.java
@@ -382,6 +382,26 @@ public static SchemaBuilder map(Schema keySchema, Schema valueSchema) {
         return builder;
     }
 
+    static SchemaBuilder arrayOfNull() {
+        return new SchemaBuilder(Type.ARRAY);
+    }
+
+    static SchemaBuilder mapOfNull() {
+        return new SchemaBuilder(Type.MAP);
+    }
+
+    static SchemaBuilder mapWithNullKeys(Schema valueSchema) {
+        SchemaBuilder result = new SchemaBuilder(Type.MAP);
+        result.valueSchema = valueSchema;
+        return result;
+    }
+
+    static SchemaBuilder mapWithNullValues(Schema keySchema) {
+        SchemaBuilder result = new SchemaBuilder(Type.MAP);
+        result.keySchema = keySchema;
+        return result;
+    }
+
     @Override
     public Schema keySchema() {
         return keySchema;
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 c1bebdfe14e35..d99fbcabf86df 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
@@ -30,13 +30,17 @@
 import java.text.SimpleDateFormat;
 import java.text.StringCharacterIterator;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Base64;
 import java.util.Calendar;
+import java.util.Collections;
+import java.util.HashSet;
 import java.util.Iterator;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.NoSuchElementException;
+import java.util.Set;
 import java.util.TimeZone;
 import java.util.regex.Pattern;
 
@@ -70,9 +74,18 @@ public class Values {
     private static final String FALSE_LITERAL = Boolean.FALSE.toString();
     private static final long MILLIS_PER_DAY = 24 * 60 * 60 * 1000;
     private static final String NULL_VALUE = "null";
-    private static final String ISO_8601_DATE_FORMAT_PATTERN = "yyyy-MM-dd";
-    private static final String ISO_8601_TIME_FORMAT_PATTERN = "HH:mm:ss.SSS'Z'";
-    private static final String ISO_8601_TIMESTAMP_FORMAT_PATTERN = ISO_8601_DATE_FORMAT_PATTERN + "'T'" + ISO_8601_TIME_FORMAT_PATTERN;
+    static final String ISO_8601_DATE_FORMAT_PATTERN = "yyyy-MM-dd";
+    static final String ISO_8601_TIME_FORMAT_PATTERN = "HH:mm:ss.SSS'Z'";
+    static final String ISO_8601_TIMESTAMP_FORMAT_PATTERN = ISO_8601_DATE_FORMAT_PATTERN + "'T'" + ISO_8601_TIME_FORMAT_PATTERN;
+    private static final Set TEMPORAL_LOGICAL_TYPE_NAMES =
+            Collections.unmodifiableSet(
+                    new HashSet<>(
+                            Arrays.asList(Time.LOGICAL_NAME,
+                            Timestamp.LOGICAL_NAME,
+                            Date.LOGICAL_NAME
+                            )
+                    )
+            );
 
     private static final String QUOTE_DELIMITER = "\"";
     private static final String COMMA_DELIMITER = ",";
@@ -467,6 +480,9 @@ protected static Object convertTo(Schema toSchema, Schema fromSchema, Object val
                                 int days = (int) (millis / MILLIS_PER_DAY); // truncates
                                 return Date.toLogical(toSchema, days);
                             }
+                        } else {
+                            // There is no fromSchema, so no conversion is needed
+                            return value;
                         }
                     }
                     long numeric = asLong(value, fromSchema, null);
@@ -492,6 +508,9 @@ protected static Object convertTo(Schema toSchema, Schema fromSchema, Object val
                                 calendar.set(Calendar.DAY_OF_MONTH, 1);
                                 return Time.toLogical(toSchema, (int) calendar.getTimeInMillis());
                             }
+                        } else {
+                            // There is no fromSchema, so no conversion is needed
+                            return value;
                         }
                     }
                     long numeric = asLong(value, fromSchema, null);
@@ -523,6 +542,9 @@ protected static Object convertTo(Schema toSchema, Schema fromSchema, Object val
                             if (Timestamp.LOGICAL_NAME.equals(fromSchemaName)) {
                                 return value;
                             }
+                        } else {
+                            // There is no fromSchema, so no conversion is needed
+                            return value;
                         }
                     }
                     long numeric = asLong(value, fromSchema, null);
@@ -723,14 +745,30 @@ public static DateFormat dateFormatFor(java.util.Date value) {
         return new SimpleDateFormat(ISO_8601_TIMESTAMP_FORMAT_PATTERN);
     }
 
+    protected static boolean canParseSingleTokenLiteral(Parser parser, boolean embedded, String tokenLiteral) {
+        int startPosition = parser.mark();
+        // If the next token is what we expect, then either...
+        if (parser.canConsume(tokenLiteral)) {
+            //   ...we're reading an embedded value, in which case the next token will be handled appropriately
+            //      by the caller if it's something like an end delimiter for a map or array, or a comma to
+            //      separate multiple embedded values...
+            //   ...or it's being parsed as part of a top-level string, in which case, any other tokens should
+            //      cause use to stop parsing this single-token literal as such and instead just treat it like
+            //      a string. For example, the top-level string "true}" will be tokenized as the tokens "true" and
+            //      "}", but should ultimately be parsed as just the string "true}" instead of the boolean true.
+            if (embedded || !parser.hasNext()) {
+                return true;
+            }
+        }
+        parser.rewindTo(startPosition);
+        return false;
+    }
+
     protected static SchemaAndValue parse(Parser parser, boolean embedded) throws NoSuchElementException {
         if (!parser.hasNext()) {
             return null;
         }
         if (embedded) {
-            if (parser.canConsume(NULL_VALUE)) {
-                return null;
-            }
             if (parser.canConsume(QUOTE_DELIMITER)) {
                 StringBuilder sb = new StringBuilder();
                 while (parser.hasNext()) {
@@ -739,37 +777,61 @@ protected static SchemaAndValue parse(Parser parser, boolean embedded) throws No
                     }
                     sb.append(parser.next());
                 }
-                return new SchemaAndValue(Schema.STRING_SCHEMA, sb.toString());
+                String content = sb.toString();
+                // We can parse string literals as temporal logical types, but all others
+                // are treated as strings
+                SchemaAndValue parsed = parseString(content);
+                if (parsed != null && TEMPORAL_LOGICAL_TYPE_NAMES.contains(parsed.schema().name())) {
+                    return parsed;
+                }
+                return new SchemaAndValue(Schema.STRING_SCHEMA, content);
             }
         }
-        if (parser.canConsume(TRUE_LITERAL)) {
+
+        if (canParseSingleTokenLiteral(parser, embedded, NULL_VALUE)) {
+            return null;
+        }
+        if (canParseSingleTokenLiteral(parser, embedded, TRUE_LITERAL)) {
             return TRUE_SCHEMA_AND_VALUE;
         }
-        if (parser.canConsume(FALSE_LITERAL)) {
+        if (canParseSingleTokenLiteral(parser, embedded, FALSE_LITERAL)) {
             return FALSE_SCHEMA_AND_VALUE;
         }
+
         int startPosition = parser.mark();
+
         try {
             if (parser.canConsume(ARRAY_BEGIN_DELIMITER)) {
                 List result = new ArrayList<>();
                 Schema elementSchema = null;
                 while (parser.hasNext()) {
                     if (parser.canConsume(ARRAY_END_DELIMITER)) {
-                        Schema listSchema = null;
+                        Schema listSchema;
                         if (elementSchema != null) {
                             listSchema = SchemaBuilder.array(elementSchema).schema();
+                            result = alignListEntriesWithSchema(listSchema, result);
+                        } else {
+                            // Every value is null
+                            listSchema = SchemaBuilder.arrayOfNull().build();
                         }
-                        result = alignListEntriesWithSchema(listSchema, result);
                         return new SchemaAndValue(listSchema, result);
                     }
+
                     if (parser.canConsume(COMMA_DELIMITER)) {
                         throw new DataException("Unable to parse an empty array element: " + parser.original());
                     }
                     SchemaAndValue element = parse(parser, true);
                     elementSchema = commonSchemaFor(elementSchema, element);
-                    result.add(element.value());
-                    parser.canConsume(COMMA_DELIMITER);
+                    result.add(element != null ? element.value() : null);
+
+                    int currentPosition = parser.mark();
+                    if (parser.canConsume(ARRAY_END_DELIMITER)) {
+                        parser.rewindTo(currentPosition);
+                    } else if (!parser.canConsume(COMMA_DELIMITER)) {
+                        throw new DataException("Array elements missing '" + COMMA_DELIMITER + "' delimiter");
+                    }
                 }
+
                 // Missing either a comma or an end delimiter
                 if (COMMA_DELIMITER.equals(parser.previous())) {
                     throw new DataException("Array is missing element after ',': " + parser.original());
@@ -783,26 +845,36 @@ protected static SchemaAndValue parse(Parser parser, boolean embedded) throws No
                 Schema valueSchema = null;
                 while (parser.hasNext()) {
                     if (parser.canConsume(MAP_END_DELIMITER)) {
-                        Schema mapSchema = null;
+                        Schema mapSchema;
                         if (keySchema != null && valueSchema != null) {
-                            mapSchema = SchemaBuilder.map(keySchema, valueSchema).schema();
+                            mapSchema = SchemaBuilder.map(keySchema, valueSchema).build();
+                            result = alignMapKeysAndValuesWithSchema(mapSchema, result);
+                        } else if (keySchema != null) {
+                            mapSchema = SchemaBuilder.mapWithNullValues(keySchema);
+                            result = alignMapKeysWithSchema(mapSchema, result);
+                        } else {
+                            mapSchema = SchemaBuilder.mapOfNull().build();
                         }
-                        result = alignMapKeysAndValuesWithSchema(mapSchema, result);
                         return new SchemaAndValue(mapSchema, result);
                     }
+
                     if (parser.canConsume(COMMA_DELIMITER)) {
-                        throw new DataException("Unable to parse a map entry has no key or value: " + parser.original());
+                        throw new DataException("Unable to parse a map entry with no key or value: " + parser.original());
                     }
                     SchemaAndValue key = parse(parser, true);
                     if (key == null || key.value() == null) {
                         throw new DataException("Map entry may not have a null key: " + parser.original());
                     }
+
                     if (!parser.canConsume(ENTRY_DELIMITER)) {
-                        throw new DataException("Map entry is missing '=': " + parser.original());
+                        throw new DataException("Map entry is missing '" + ENTRY_DELIMITER
+                                                + "' at " + parser.position()
+                                                + " in " + parser.original());
                     }
                     SchemaAndValue value = parse(parser, true);
                     Object entryValue = value != null ? value.value() : null;
                     result.put(key.value(), entryValue);
+
                     parser.canConsume(COMMA_DELIMITER);
                     keySchema = commonSchemaFor(keySchema, key);
                     valueSchema = commonSchemaFor(valueSchema, value);
@@ -811,16 +883,42 @@ protected static SchemaAndValue parse(Parser parser, boolean embedded) throws No
                 if (COMMA_DELIMITER.equals(parser.previous())) {
                     throw new DataException("Map is missing element after ',': " + parser.original());
                 }
-                throw new DataException("Map is missing terminating ']': " + parser.original());
+                throw new DataException("Map is missing terminating '}': " + parser.original());
             }
         } catch (DataException e) {
-            LOG.debug("Unable to parse the value as a map; reverting to string", e);
+            LOG.trace("Unable to parse the value as a map or an array; reverting to string", e);
             parser.rewindTo(startPosition);
         }
-        String token = parser.next().trim();
-        assert !token.isEmpty(); // original can be empty string but is handled right away; no way for token to be empty here
+
+        String token = parser.next();
+        if (token.trim().isEmpty()) {
+            return new SchemaAndValue(Schema.STRING_SCHEMA, token);
+        }
+        token = token.trim();
+
         char firstChar = token.charAt(0);
         boolean firstCharIsDigit = Character.isDigit(firstChar);
+
+        // Temporal types are more restrictive, so try them first
+        if (firstCharIsDigit) {
+            // The time and timestamp literals may be split into 5 tokens since an unescaped colon
+            // is a delimiter. Check these first since the first of these tokens is a simple numeric
+            int position = parser.mark();
+            String remainder = parser.next(4);
+            if (remainder != null) {
+                String timeOrTimestampStr = token + remainder;
+                SchemaAndValue temporal = parseAsTemporal(timeOrTimestampStr);
+                if (temporal != null) {
+                    return temporal;
+                }
+            }
+            // No match was found using the 5 tokens, so rewind and see if the current token has a date, time, or timestamp
+            parser.rewindTo(position);
+            SchemaAndValue temporal = parseAsTemporal(token);
+            if (temporal != null) {
+                return temporal;
+            }
+        }
         if (firstCharIsDigit || firstChar == '+' || firstChar == '-') {
             try {
                 // Try to parse as a number ...
@@ -855,34 +953,42 @@ protected static SchemaAndValue parse(Parser parser, boolean embedded) throws No
                 // can't parse as a number
             }
         }
-        if (firstCharIsDigit) {
-            // Check for a date, time, or timestamp ...
-            int tokenLength = token.length();
-            if (tokenLength == ISO_8601_DATE_LENGTH) {
-                try {
-                    return new SchemaAndValue(Date.SCHEMA, new SimpleDateFormat(ISO_8601_DATE_FORMAT_PATTERN).parse(token));
-                } catch (ParseException e) {
-                    // not a valid date
-                }
-            } else if (tokenLength == ISO_8601_TIME_LENGTH) {
-                try {
-                    return new SchemaAndValue(Time.SCHEMA, new SimpleDateFormat(ISO_8601_TIME_FORMAT_PATTERN).parse(token));
-                } catch (ParseException e) {
-                    // not a valid date
-                }
-            } else if (tokenLength == ISO_8601_TIMESTAMP_LENGTH) {
-                try {
-                    return new SchemaAndValue(Timestamp.SCHEMA, new SimpleDateFormat(ISO_8601_TIMESTAMP_FORMAT_PATTERN).parse(token));
-                } catch (ParseException e) {
-                    // not a valid date
-                }
-            }
+        if (embedded) {
+            throw new DataException("Failed to parse embedded value");
         }
-        // At this point, the only thing this can be is a string. Embedded strings were processed above,
-        // so this is not embedded and we can use the original string...
+        // At this point, the only thing this non-embedded value can be is a string.
         return new SchemaAndValue(Schema.STRING_SCHEMA, parser.original());
     }
 
+    private static SchemaAndValue parseAsTemporal(String token) {
+        if (token == null) {
+            return null;
+        }
+        // If the colons were escaped, we'll see the escape chars and need to remove them
+        token = token.replace("\\:", ":");
+        int tokenLength = token.length();
+        if (tokenLength == ISO_8601_TIME_LENGTH) {
+            try {
+                return new SchemaAndValue(Time.SCHEMA, new SimpleDateFormat(ISO_8601_TIME_FORMAT_PATTERN).parse(token));
+            } catch (ParseException e) {
+              // not a valid date
+            }
+        } else if (tokenLength == ISO_8601_TIMESTAMP_LENGTH) {
+            try {
+                return new SchemaAndValue(Timestamp.SCHEMA, new SimpleDateFormat(ISO_8601_TIMESTAMP_FORMAT_PATTERN).parse(token));
+            } catch (ParseException e) {
+              // not a valid date
+            }
+        } else if (tokenLength == ISO_8601_DATE_LENGTH) {
+            try {
+                return new SchemaAndValue(Date.SCHEMA, new SimpleDateFormat(ISO_8601_DATE_FORMAT_PATTERN).parse(token));
+            } catch (ParseException e) {
+                // not a valid date
+            }
+        }
+        return null;
+    }
+
     protected static Schema commonSchemaFor(Schema previous, SchemaAndValue latest) {
         if (latest == null) {
             return previous;
@@ -953,9 +1059,6 @@ protected static Schema commonSchemaFor(Schema previous, SchemaAndValue latest)
     }
 
     protected static List alignListEntriesWithSchema(Schema schema, List input) {
-        if (schema == null) {
-            return input;
-        }
         Schema valueSchema = schema.valueSchema();
         List result = new ArrayList<>();
         for (Object value : input) {
@@ -966,9 +1069,6 @@ protected static List alignListEntriesWithSchema(Schema schema, List alignMapKeysAndValuesWithSchema(Schema mapSchema, Map input) {
-        if (mapSchema == null) {
-            return input;
-        }
         Schema keySchema = mapSchema.keySchema();
         Schema valueSchema = mapSchema.valueSchema();
         Map result = new LinkedHashMap<>();
@@ -980,6 +1080,16 @@ protected static Map alignMapKeysAndValuesWithSchema(Schema mapS
         return result;
     }
 
+    protected static Map alignMapKeysWithSchema(Schema mapSchema, Map input) {
+        Schema keySchema = mapSchema.keySchema();
+        Map result = new LinkedHashMap<>();
+        for (Map.Entry entry : input.entrySet()) {
+            Object newKey = convertTo(keySchema, null, entry.getKey());
+            result.put(newKey, entry.getValue());
+        }
+        return result;
+    }
+
     protected static class SchemaDetector {
         private Type knownType = null;
         private boolean optional = false;
@@ -1035,6 +1145,7 @@ public int mark() {
         public void rewindTo(int position) {
             iter.setIndex(position);
             nextToken = null;
+            previousToken = null;
         }
 
         public String original() {
@@ -1059,6 +1170,19 @@ public String next() {
             return previousToken;
         }
 
+        public String next(int n) {
+            int current = mark();
+            int start = mark();
+            for (int i = 0; i != n; ++i) {
+                if (!hasNext()) {
+                    rewindTo(start);
+                    return null;
+                }
+                next();
+            }
+            return original.substring(current, position());
+        }
+
         private String consumeNextToken() throws NoSuchElementException {
             boolean escaped = false;
             int start = iter.getIndex();
@@ -1123,12 +1247,13 @@ protected boolean isNext(String expected, boolean ignoreLeadingAndTrailingWhites
                 nextToken = consumeNextToken();
             }
             if (ignoreLeadingAndTrailingWhitespace) {
-                nextToken = nextToken.trim();
-                while (nextToken.isEmpty() && canConsumeNextToken()) {
-                    nextToken = consumeNextToken().trim();
+                while (nextToken.trim().isEmpty() && canConsumeNextToken()) {
+                    nextToken = consumeNextToken();
                 }
             }
-            return nextToken.equals(expected);
+            return ignoreLeadingAndTrailingWhitespace
+                ? nextToken.trim().equals(expected)
+                : nextToken.equals(expected);
         }
     }
 }
diff --git a/connect/api/src/main/java/org/apache/kafka/connect/header/ConnectHeaders.java b/connect/api/src/main/java/org/apache/kafka/connect/header/ConnectHeaders.java
index 185ba65bf0fb6..0b5c484b3533e 100644
--- a/connect/api/src/main/java/org/apache/kafka/connect/header/ConnectHeaders.java
+++ b/connect/api/src/main/java/org/apache/kafka/connect/header/ConnectHeaders.java
@@ -274,7 +274,7 @@ public Iterator
       iterator() {
     @Override
     public Headers remove(String key) {
         checkKey(key);
-        if (!headers.isEmpty()) {
+        if (!isEmpty()) {
             Iterator
       iterator = iterator();
             while (iterator.hasNext()) {
                 if (iterator.next().key().equals(key)) {
@@ -287,7 +287,7 @@ public Headers remove(String key) {
 
     @Override
     public Headers retainLatest() {
-        if (!headers.isEmpty()) {
+        if (!isEmpty()) {
             Set keys = new HashSet<>();
             ListIterator
       iter = headers.listIterator(headers.size());
             while (iter.hasPrevious()) {
@@ -304,7 +304,7 @@ public Headers retainLatest() {
     @Override
     public Headers retainLatest(String key) {
         checkKey(key);
-        if (!headers.isEmpty()) {
+        if (!isEmpty()) {
             boolean found = false;
             ListIterator
       iter = headers.listIterator(headers.size());
             while (iter.hasPrevious()) {
@@ -322,7 +322,7 @@ public Headers retainLatest(String key) {
     @Override
     public Headers apply(String key, HeaderTransform transform) {
         checkKey(key);
-        if (!headers.isEmpty()) {
+        if (!isEmpty()) {
             ListIterator
       iter = headers.listIterator();
             while (iter.hasNext()) {
                 Header orig = iter.next();
@@ -341,7 +341,7 @@ public Headers apply(String key, HeaderTransform transform) {
 
     @Override
     public Headers apply(HeaderTransform transform) {
-        if (!headers.isEmpty()) {
+        if (!isEmpty()) {
             ListIterator
       iter = headers.listIterator();
             while (iter.hasNext()) {
                 Header orig = iter.next();
diff --git a/connect/api/src/main/java/org/apache/kafka/connect/header/Headers.java b/connect/api/src/main/java/org/apache/kafka/connect/header/Headers.java
index 761abe3d3af98..d7bd77952cf24 100644
--- a/connect/api/src/main/java/org/apache/kafka/connect/header/Headers.java
+++ b/connect/api/src/main/java/org/apache/kafka/connect/header/Headers.java
@@ -254,7 +254,7 @@ public interface Headers extends Iterable
       {
     Headers retainLatest(String key);
 
     /**
-     * Removes all but the last {@Header} object with each key.
+     * Removes all but the last {@link Header} object with each key.
      *
      * @return this object to facilitate chaining multiple methods; never null
      */
diff --git a/connect/api/src/main/java/org/apache/kafka/connect/health/AbstractState.java b/connect/api/src/main/java/org/apache/kafka/connect/health/AbstractState.java
index a18c46334cf29..f707b3c3026ed 100644
--- a/connect/api/src/main/java/org/apache/kafka/connect/health/AbstractState.java
+++ b/connect/api/src/main/java/org/apache/kafka/connect/health/AbstractState.java
@@ -17,6 +17,8 @@
 
 package org.apache.kafka.connect.health;
 
+import java.util.Objects;
+
 /**
  * Provides the current status along with identifier for Connect worker and tasks.
  */
@@ -34,10 +36,10 @@ public abstract class AbstractState {
      * @param traceMessage  any error trace message associated with the connector or the task; may be null or empty
      */
     public AbstractState(String state, String workerId, String traceMessage) {
-        if (state != null && !state.trim().isEmpty()) {
+        if (state == null || state.trim().isEmpty()) {
             throw new IllegalArgumentException("State must not be null or empty");
         }
-        if (workerId != null && !workerId.trim().isEmpty()) {
+        if (workerId == null || workerId.trim().isEmpty()) {
             throw new IllegalArgumentException("Worker ID must not be null or empty");
         }
         this.state = state;
@@ -71,4 +73,21 @@ public String workerId() {
     public String traceMessage() {
         return traceMessage;
     }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o)
+            return true;
+        if (o == null || getClass() != o.getClass())
+            return false;
+        AbstractState that = (AbstractState) o;
+        return state.equals(that.state)
+            && Objects.equals(traceMessage, that.traceMessage)
+            && workerId.equals(that.workerId);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(state, traceMessage, workerId);
+    }
 }
diff --git a/connect/api/src/main/java/org/apache/kafka/connect/health/ConnectClusterDetails.java b/connect/api/src/main/java/org/apache/kafka/connect/health/ConnectClusterDetails.java
new file mode 100644
index 0000000000000..edde6ff657a7e
--- /dev/null
+++ b/connect/api/src/main/java/org/apache/kafka/connect/health/ConnectClusterDetails.java
@@ -0,0 +1,32 @@
+/*
+ * 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.health;
+
+/**
+ * Provides immutable Connect cluster information, such as the ID of the backing Kafka cluster. The
+ * Connect framework provides the implementation for this interface.
+ */
+public interface ConnectClusterDetails {
+
+    /**
+     * Get the cluster ID of the Kafka cluster backing this Connect cluster.
+     * 
+     * @return the cluster ID of the Kafka cluster backing this Connect cluster
+     **/
+    String kafkaClusterId();
+}
\ No newline at end of file
diff --git a/connect/api/src/main/java/org/apache/kafka/connect/health/ConnectClusterState.java b/connect/api/src/main/java/org/apache/kafka/connect/health/ConnectClusterState.java
index d4292efd0caf3..753ee1a7613cd 100644
--- a/connect/api/src/main/java/org/apache/kafka/connect/health/ConnectClusterState.java
+++ b/connect/api/src/main/java/org/apache/kafka/connect/health/ConnectClusterState.java
@@ -18,10 +18,13 @@
 package org.apache.kafka.connect.health;
 
 import java.util.Collection;
+import java.util.Map;
 
 /**
- * Provides the ability to lookup connector metadata and its health. This is made available to the {@link org.apache.kafka.connect.rest.ConnectRestExtension}
- * implementations. The Connect framework provides the implementation for this interface.
+ * Provides the ability to lookup connector metadata, including status and configurations, as well
+ * as immutable cluster information such as Kafka cluster ID. This is made available to
+ * {@link org.apache.kafka.connect.rest.ConnectRestExtension} implementations. The Connect framework
+ * provides the implementation for this interface.
  */
 public interface ConnectClusterState {
 
@@ -43,4 +46,27 @@ public interface ConnectClusterState {
      * @throws org.apache.kafka.connect.errors.NotFoundException if the requested connector can't be found
      */
     ConnectorHealth connectorHealth(String connName);
+
+    /**
+     * Lookup the current configuration of a connector. This provides the current snapshot of configuration by querying the underlying
+     * herder. A connector returned by previous invocation of {@link #connectors()} may no longer be available and could result in {@link
+     * org.apache.kafka.connect.errors.NotFoundException}.
+     *
+     * @param connName name of the connector
+     * @return the configuration of the connector for the connector name
+     * @throws org.apache.kafka.connect.errors.NotFoundException if the requested connector can't be found
+     * @throws java.lang.UnsupportedOperationException if the default implementation has not been overridden
+     */
+    default Map connectorConfig(String connName) {
+        throw new UnsupportedOperationException();
+    }
+
+    /**
+     * Get details about the setup of the Connect cluster.
+     * @return a {@link ConnectClusterDetails} object containing information about the cluster
+     * @throws java.lang.UnsupportedOperationException if the default implementation has not been overridden
+     **/
+    default ConnectClusterDetails clusterDetails() {
+        throw new UnsupportedOperationException();
+    }
 }
diff --git a/connect/api/src/main/java/org/apache/kafka/connect/health/ConnectorHealth.java b/connect/api/src/main/java/org/apache/kafka/connect/health/ConnectorHealth.java
index 3a9efd15372ed..12fa6b76aff1e 100644
--- a/connect/api/src/main/java/org/apache/kafka/connect/health/ConnectorHealth.java
+++ b/connect/api/src/main/java/org/apache/kafka/connect/health/ConnectorHealth.java
@@ -35,7 +35,7 @@ public ConnectorHealth(String name,
                            ConnectorState connectorState,
                            Map tasks,
                            ConnectorType type) {
-        if (name != null && !name.trim().isEmpty()) {
+        if (name == null || name.trim().isEmpty()) {
             throw new IllegalArgumentException("Connector name is required");
         }
         Objects.requireNonNull(connectorState, "connectorState can't be null");
@@ -83,4 +83,31 @@ public ConnectorType type() {
         return type;
     }
 
+    @Override
+    public boolean equals(Object o) {
+        if (this == o)
+            return true;
+        if (o == null || getClass() != o.getClass())
+            return false;
+        ConnectorHealth that = (ConnectorHealth) o;
+        return name.equals(that.name)
+            && connectorState.equals(that.connectorState)
+            && tasks.equals(that.tasks)
+            && type == that.type;
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(name, connectorState, tasks, type);
+    }
+
+    @Override
+    public String toString() {
+        return "ConnectorHealth{"
+            + "name='" + name + '\''
+            + ", connectorState=" + connectorState
+            + ", tasks=" + tasks
+            + ", type=" + type
+            + '}';
+    }
 }
diff --git a/connect/api/src/main/java/org/apache/kafka/connect/health/ConnectorState.java b/connect/api/src/main/java/org/apache/kafka/connect/health/ConnectorState.java
index d5571bc4ff5d0..63044265bb999 100644
--- a/connect/api/src/main/java/org/apache/kafka/connect/health/ConnectorState.java
+++ b/connect/api/src/main/java/org/apache/kafka/connect/health/ConnectorState.java
@@ -32,4 +32,13 @@ public class ConnectorState extends AbstractState {
     public ConnectorState(String state, String workerId, String traceMessage) {
         super(state, workerId, traceMessage);
     }
+
+    @Override
+    public String toString() {
+        return "ConnectorState{"
+            + "state='" + state() + '\''
+            + ", traceMessage='" + traceMessage() + '\''
+            + ", workerId='" + workerId() + '\''
+            + '}';
+    }
 }
diff --git a/connect/api/src/main/java/org/apache/kafka/connect/health/TaskState.java b/connect/api/src/main/java/org/apache/kafka/connect/health/TaskState.java
index 1c1be159970d9..ae78a5f3af990 100644
--- a/connect/api/src/main/java/org/apache/kafka/connect/health/TaskState.java
+++ b/connect/api/src/main/java/org/apache/kafka/connect/health/TaskState.java
@@ -50,20 +50,28 @@ public int taskId() {
 
     @Override
     public boolean equals(Object o) {
-        if (this == o) {
+        if (this == o)
             return true;
-        }
-        if (o == null || getClass() != o.getClass()) {
+        if (o == null || getClass() != o.getClass())
+            return false;
+        if (!super.equals(o))
             return false;
-        }
-
         TaskState taskState = (TaskState) o;
-
         return taskId == taskState.taskId;
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(taskId);
+        return Objects.hash(super.hashCode(), taskId);
+    }
+
+    @Override
+    public String toString() {
+        return "TaskState{"
+            + "taskId='" + taskId + '\''
+            + "state='" + state() + '\''
+            + ", traceMessage='" + traceMessage() + '\''
+            + ", workerId='" + workerId() + '\''
+            + '}';
     }
 }
diff --git a/connect/api/src/main/java/org/apache/kafka/connect/sink/SinkRecord.java b/connect/api/src/main/java/org/apache/kafka/connect/sink/SinkRecord.java
index 642b7167f83cb..12c7ee119abbb 100644
--- a/connect/api/src/main/java/org/apache/kafka/connect/sink/SinkRecord.java
+++ b/connect/api/src/main/java/org/apache/kafka/connect/sink/SinkRecord.java
@@ -88,7 +88,7 @@ public boolean equals(Object o) {
     @Override
     public int hashCode() {
         int result = super.hashCode();
-        result = 31 * result + (int) (kafkaOffset ^ (kafkaOffset >>> 32));
+        result = 31 * result + Long.hashCode(kafkaOffset);
         result = 31 * result + timestampType.hashCode();
         return result;
     }
diff --git a/connect/api/src/main/java/org/apache/kafka/connect/source/SourceRecord.java b/connect/api/src/main/java/org/apache/kafka/connect/source/SourceRecord.java
index 315f0f3a08847..2c390eed4e495 100644
--- a/connect/api/src/main/java/org/apache/kafka/connect/source/SourceRecord.java
+++ b/connect/api/src/main/java/org/apache/kafka/connect/source/SourceRecord.java
@@ -21,6 +21,7 @@
 import org.apache.kafka.connect.header.Header;
 
 import java.util.Map;
+import java.util.Objects;
 
 /**
  * 
      
@@ -113,12 +114,8 @@ public boolean equals(Object o) {
 
         SourceRecord that = (SourceRecord) o;
 
-        if (sourcePartition != null ? !sourcePartition.equals(that.sourcePartition) : that.sourcePartition != null)
-            return false;
-        if (sourceOffset != null ? !sourceOffset.equals(that.sourceOffset) : that.sourceOffset != null)
-            return false;
-
-        return true;
+        return Objects.equals(sourcePartition, that.sourcePartition) &&
+                Objects.equals(sourceOffset, that.sourceOffset);
     }
 
     @Override
diff --git a/connect/api/src/main/java/org/apache/kafka/connect/source/SourceTask.java b/connect/api/src/main/java/org/apache/kafka/connect/source/SourceTask.java
index 8767a62045e4e..4dea6ce80bdbc 100644
--- a/connect/api/src/main/java/org/apache/kafka/connect/source/SourceTask.java
+++ b/connect/api/src/main/java/org/apache/kafka/connect/source/SourceTask.java
@@ -17,6 +17,7 @@
 package org.apache.kafka.connect.source;
 
 import org.apache.kafka.connect.connector.Task;
+import org.apache.kafka.clients.producer.RecordMetadata;
 
 import java.util.List;
 import java.util.Map;
@@ -88,7 +89,13 @@ public void commit() throws InterruptedException {
 
     /**
      * 
      
-     * Commit an individual {@link SourceRecord} when the callback from the producer client is received, or if a record is filtered by a transformation.
+     * Commit an individual {@link SourceRecord} when the callback from the producer client is received. This method is
+     * also called when a record is filtered by a transformation, and thus will never be ACK'd by a broker.
+     * 
      
+     * 
      
+     * This is an alias for {@link commitRecord(SourceRecord, RecordMetadata)} for backwards compatibility. The default
+     * implementation of {@link commitRecord(SourceRecord, RecordMetadata)} just calls this method. It is not necessary
+     * to override both methods.
      * 
      
      * 
      
      * SourceTasks are not required to implement this functionality; Kafka Connect will record offsets
@@ -96,10 +103,37 @@ public void commit() throws InterruptedException {
      * in their own system.
      * 
      
      *
-     * @param record {@link SourceRecord} that was successfully sent via the producer.
+     * @param record {@link SourceRecord} that was successfully sent via the producer or filtered by a transformation
      * @throws InterruptedException
+     * @see commitRecord(SourceRecord, RecordMetadata)
      */
     public void commitRecord(SourceRecord record) throws InterruptedException {
         // This space intentionally left blank.
     }
+
+    /**
+     * 
      
+     * Commit an individual {@link SourceRecord} when the callback from the producer client is received. This method is
+     * also called when a record is filtered by a transformation, and thus will never be ACK'd by a broker. In this case
+     * {@code metadata} will be null.
+     * 
      
+     * 
      
+     * SourceTasks are not required to implement this functionality; Kafka Connect will record offsets
+     * automatically. This hook is provided for systems that also need to store offsets internally
+     * in their own system.
+     * 
      
+     * 
      
+     * The default implementation just calls @{link commitRecord(SourceRecord)}, which is a nop by default. It is
+     * not necessary to implement both methods.
+     * 
      
+     *
+     * @param record {@link SourceRecord} that was successfully sent via the producer or filtered by a transformation
+     * @param metadata {@link RecordMetadata} record metadata returned from the broker, or null if the record was filtered
+     * @throws InterruptedException
+     */
+    public void commitRecord(SourceRecord record, RecordMetadata metadata)
+            throws InterruptedException {
+        // by default, just call other method for backwards compatability
+        commitRecord(record);
+    }
 }
diff --git a/connect/api/src/main/java/org/apache/kafka/connect/storage/Converter.java b/connect/api/src/main/java/org/apache/kafka/connect/storage/Converter.java
index 507489345b9e9..29300c74eda49 100644
--- a/connect/api/src/main/java/org/apache/kafka/connect/storage/Converter.java
+++ b/connect/api/src/main/java/org/apache/kafka/connect/storage/Converter.java
@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.connect.storage;
 
+import org.apache.kafka.common.header.Headers;
 import org.apache.kafka.connect.data.Schema;
 import org.apache.kafka.connect.data.SchemaAndValue;
 
@@ -44,6 +45,23 @@ public interface Converter {
      */
     byte[] fromConnectData(String topic, Schema schema, Object value);
 
+    /**
+     * Convert a Kafka Connect data object to a native object for serialization,
+     * potentially using the supplied topic and headers in the record as necessary.
+     *
+     * 
      Connect uses this method directly, and for backward compatibility reasons this method
+     * by default will call the {@link #fromConnectData(String, Schema, Object)} method.
+     * Override this method to make use of the supplied headers.
      
+     * @param topic the topic associated with the data
+     * @param headers the headers associated with the data
+     * @param schema the schema for the value
+     * @param value the value to convert
+     * @return the serialized value
+     */
+    default byte[] fromConnectData(String topic, Headers headers, Schema schema, Object value) {
+        return fromConnectData(topic, schema, value);
+    }
+
     /**
      * Convert a native object to a Kafka Connect data object.
      * @param topic the topic associated with the data
@@ -51,4 +69,20 @@ public interface Converter {
      * @return an object containing the {@link Schema} and the converted value
      */
     SchemaAndValue toConnectData(String topic, byte[] value);
+
+    /**
+     * Convert a native object to a Kafka Connect data object,
+     * potentially using the supplied topic and headers in the record as necessary.
+     *
+     * 
      Connect uses this method directly, and for backward compatibility reasons this method
+     * by default will call the {@link #toConnectData(String, byte[])} method.
+     * Override this method to make use of the supplied headers.
      
+     * @param topic the topic associated with the data
+     * @param headers the headers associated with the data
+     * @param value the value to convert
+     * @return an object containing the {@link Schema} and the converted value
+     */
+    default SchemaAndValue toConnectData(String topic, Headers headers, byte[] value) {
+        return toConnectData(topic, value);
+    }
 }
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 162222db155c8..c437e46c25956 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
@@ -21,7 +21,11 @@
 import org.apache.kafka.connect.errors.DataException;
 import org.junit.Test;
 
+import java.text.SimpleDateFormat;
 import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
@@ -35,6 +39,8 @@
 
 public class ValuesTest {
 
+    private static final String WHITESPACE = "\n \t \t\n";
+
     private static final long MILLIS_PER_DAY = 24 * 60 * 60 * 1000;
 
     private static final Map STRING_MAP = new LinkedHashMap<>();
@@ -67,6 +73,149 @@ public class ValuesTest {
         INT_LIST.add(-987654321);
     }
 
+    @Test
+    public void shouldNotParseUnquotedEmbeddedMapKeysAsStrings() {
+        SchemaAndValue schemaAndValue = Values.parseString("{foo: 3}");
+        assertEquals(Type.STRING, schemaAndValue.schema().type());
+        assertEquals("{foo: 3}", schemaAndValue.value());
+    }
+
+    @Test
+    public void shouldNotParseUnquotedEmbeddedMapValuesAsStrings() {
+        SchemaAndValue schemaAndValue = Values.parseString("{3: foo}");
+        assertEquals(Type.STRING, schemaAndValue.schema().type());
+        assertEquals("{3: foo}", schemaAndValue.value());
+    }
+
+    @Test
+    public void shouldNotParseUnquotedArrayElementsAsStrings() {
+        SchemaAndValue schemaAndValue = Values.parseString("[foo]");
+        assertEquals(Type.STRING, schemaAndValue.schema().type());
+        assertEquals("[foo]", schemaAndValue.value());
+    }
+
+    @Test
+    public void shouldNotParseStringsBeginningWithNullAsStrings() {
+        SchemaAndValue schemaAndValue = Values.parseString("null=");
+        assertEquals(Type.STRING, schemaAndValue.schema().type());
+        assertEquals("null=", schemaAndValue.value());
+    }
+
+    @Test
+    public void shouldParseStringsBeginningWithTrueAsStrings() {
+        SchemaAndValue schemaAndValue = Values.parseString("true}");
+        assertEquals(Type.STRING, schemaAndValue.schema().type());
+        assertEquals("true}", schemaAndValue.value());
+    }
+
+    @Test
+    public void shouldParseStringsBeginningWithFalseAsStrings() {
+        SchemaAndValue schemaAndValue = Values.parseString("false]");
+        assertEquals(Type.STRING, schemaAndValue.schema().type());
+        assertEquals("false]", schemaAndValue.value());
+    }
+
+    @Test
+    public void shouldParseTrueAsBooleanIfSurroundedByWhitespace() {
+        SchemaAndValue schemaAndValue = Values.parseString(WHITESPACE + "true" + WHITESPACE);
+        assertEquals(Type.BOOLEAN, schemaAndValue.schema().type());
+        assertEquals(true, schemaAndValue.value());
+    }
+
+    @Test
+    public void shouldParseFalseAsBooleanIfSurroundedByWhitespace() {
+        SchemaAndValue schemaAndValue = Values.parseString(WHITESPACE + "false" + WHITESPACE);
+        assertEquals(Type.BOOLEAN, schemaAndValue.schema().type());
+        assertEquals(false, schemaAndValue.value());
+    }
+
+    @Test
+    public void shouldParseNullAsNullIfSurroundedByWhitespace() {
+        SchemaAndValue schemaAndValue = Values.parseString(WHITESPACE + "null" + WHITESPACE);
+        assertNull(schemaAndValue);
+    }
+
+    @Test
+    public void shouldParseBooleanLiteralsEmbeddedInArray() {
+        SchemaAndValue schemaAndValue = Values.parseString("[true, false]");
+        assertEquals(Type.ARRAY, schemaAndValue.schema().type());
+        assertEquals(Type.BOOLEAN, schemaAndValue.schema().valueSchema().type());
+        assertEquals(Arrays.asList(true, false), schemaAndValue.value());
+    }
+
+    @Test
+    public void shouldParseBooleanLiteralsEmbeddedInMap() {
+        SchemaAndValue schemaAndValue = Values.parseString("{true: false, false: true}");
+        assertEquals(Type.MAP, schemaAndValue.schema().type());
+        assertEquals(Type.BOOLEAN, schemaAndValue.schema().keySchema().type());
+        assertEquals(Type.BOOLEAN, schemaAndValue.schema().valueSchema().type());
+        Map expectedValue = new HashMap<>();
+        expectedValue.put(true, false);
+        expectedValue.put(false, true);
+        assertEquals(expectedValue, schemaAndValue.value());
+    }
+
+    @Test
+    public void shouldNotParseAsMapWithoutCommas() {
+        SchemaAndValue schemaAndValue = Values.parseString("{6:9 4:20}");
+        assertEquals(Type.STRING, schemaAndValue.schema().type());
+        assertEquals("{6:9 4:20}", schemaAndValue.value());
+    }
+
+    @Test
+    public void shouldNotParseAsArrayWithoutCommas() {
+        SchemaAndValue schemaAndValue = Values.parseString("[0 1 2]");
+        assertEquals(Type.STRING, schemaAndValue.schema().type());
+        assertEquals("[0 1 2]", schemaAndValue.value());
+    }
+
+    @Test
+    public void shouldParseEmptyMap() {
+        SchemaAndValue schemaAndValue = Values.parseString("{}");
+        assertEquals(Type.MAP, schemaAndValue.schema().type());
+        assertEquals(Collections.emptyMap(), schemaAndValue.value());
+    }
+
+    @Test
+    public void shouldParseEmptyArray() {
+        SchemaAndValue schemaAndValue = Values.parseString("[]");
+        assertEquals(Type.ARRAY, schemaAndValue.schema().type());
+        assertEquals(Collections.emptyList(), schemaAndValue.value());
+    }
+
+    @Test
+    public void shouldNotParseAsMapWithNullKeys() {
+        SchemaAndValue schemaAndValue = Values.parseString("{null: 3}");
+        assertEquals(Type.STRING, schemaAndValue.schema().type());
+        assertEquals("{null: 3}", schemaAndValue.value());
+    }
+
+    @Test
+    public void shouldParseNull() {
+        SchemaAndValue schemaAndValue = Values.parseString("null");
+        assertNull(schemaAndValue);
+    }
+
+    @Test
+    public void shouldConvertStringOfNull() {
+        assertRoundTrip(Schema.STRING_SCHEMA, "null");
+    }
+
+    @Test
+    public void shouldParseNullMapValues() {
+        SchemaAndValue schemaAndValue = Values.parseString("{3: null}");
+        assertEquals(Type.MAP, schemaAndValue.schema().type());
+        assertEquals(Type.INT8, schemaAndValue.schema().keySchema().type());
+        assertEquals(Collections.singletonMap((byte) 3, null), schemaAndValue.value());
+    }
+
+    @Test
+    public void shouldParseNullArrayElements() {
+        SchemaAndValue schemaAndValue = Values.parseString("[null]");
+        assertEquals(Type.ARRAY, schemaAndValue.schema().type());
+        assertEquals(Collections.singletonList(null), schemaAndValue.value());
+    }
+
     @Test
     public void shouldEscapeStringsWithEmbeddedQuotesAndBackslashes() {
         String original = "three\"blind\\\"mice";
@@ -214,7 +363,8 @@ public void shouldConvertStringOfListWithMixedElementTypesIntoListWithDifferentE
     public void shouldParseStringListWithMultipleElementTypesAndReturnListWithNoSchema() {
         String str = "[1, 2, 3, \"four\"]";
         SchemaAndValue result = Values.parseString(str);
-        assertNull(result.schema());
+        assertEquals(Type.ARRAY, result.schema().type());
+        assertNull(result.schema().valueSchema());
         List list = (List) result.value();
         assertEquals(4, list.size());
         assertEquals(1, ((Number) list.get(0)).intValue());
@@ -234,6 +384,146 @@ public void shouldParseStringListWithExtraDelimitersAndReturnString() {
         assertEquals(str, result.value());
     }
 
+    @Test
+    public void shouldParseTimestampStringAsTimestamp() throws Exception {
+        String str = "2019-08-23T14:34:54.346Z";
+        SchemaAndValue result = Values.parseString(str);
+        assertEquals(Type.INT64, result.schema().type());
+        assertEquals(Timestamp.LOGICAL_NAME, result.schema().name());
+        java.util.Date expected = new SimpleDateFormat(Values.ISO_8601_TIMESTAMP_FORMAT_PATTERN).parse(str);
+        assertEquals(expected, result.value());
+    }
+
+    @Test
+    public void shouldParseDateStringAsDate() throws Exception {
+        String str = "2019-08-23";
+        SchemaAndValue result = Values.parseString(str);
+        assertEquals(Type.INT32, result.schema().type());
+        assertEquals(Date.LOGICAL_NAME, result.schema().name());
+        java.util.Date expected = new SimpleDateFormat(Values.ISO_8601_DATE_FORMAT_PATTERN).parse(str);
+        assertEquals(expected, result.value());
+    }
+
+    @Test
+    public void shouldParseTimeStringAsDate() throws Exception {
+        String str = "14:34:54.346Z";
+        SchemaAndValue result = Values.parseString(str);
+        assertEquals(Type.INT32, result.schema().type());
+        assertEquals(Time.LOGICAL_NAME, result.schema().name());
+        java.util.Date expected = new SimpleDateFormat(Values.ISO_8601_TIME_FORMAT_PATTERN).parse(str);
+        assertEquals(expected, result.value());
+    }
+
+    @Test
+    public void shouldParseTimestampStringWithEscapedColonsAsTimestamp() throws Exception {
+        String str = "2019-08-23T14\\:34\\:54.346Z";
+        SchemaAndValue result = Values.parseString(str);
+        assertEquals(Type.INT64, result.schema().type());
+        assertEquals(Timestamp.LOGICAL_NAME, result.schema().name());
+        String expectedStr = "2019-08-23T14:34:54.346Z";
+        java.util.Date expected = new SimpleDateFormat(Values.ISO_8601_TIMESTAMP_FORMAT_PATTERN).parse(expectedStr);
+        assertEquals(expected, result.value());
+    }
+
+    @Test
+    public void shouldParseTimeStringWithEscapedColonsAsDate() throws Exception {
+        String str = "14\\:34\\:54.346Z";
+        SchemaAndValue result = Values.parseString(str);
+        assertEquals(Type.INT32, result.schema().type());
+        assertEquals(Time.LOGICAL_NAME, result.schema().name());
+        String expectedStr = "14:34:54.346Z";
+        java.util.Date expected = new SimpleDateFormat(Values.ISO_8601_TIME_FORMAT_PATTERN).parse(expectedStr);
+        assertEquals(expected, result.value());
+    }
+
+    @Test
+    public void shouldParseDateStringAsDateInArray() throws Exception {
+        String dateStr = "2019-08-23";
+        String arrayStr = "[" + dateStr + "]";
+        SchemaAndValue result = Values.parseString(arrayStr);
+        assertEquals(Type.ARRAY, result.schema().type());
+        Schema elementSchema = result.schema().valueSchema();
+        assertEquals(Type.INT32, elementSchema.type());
+        assertEquals(Date.LOGICAL_NAME, elementSchema.name());
+        java.util.Date expected = new SimpleDateFormat(Values.ISO_8601_DATE_FORMAT_PATTERN).parse(dateStr);
+        assertEquals(Collections.singletonList(expected), result.value());
+    }
+
+    @Test
+    public void shouldParseTimeStringAsTimeInArray() throws Exception {
+        String timeStr = "14:34:54.346Z";
+        String arrayStr = "[" + timeStr + "]";
+        SchemaAndValue result = Values.parseString(arrayStr);
+        assertEquals(Type.ARRAY, result.schema().type());
+        Schema elementSchema = result.schema().valueSchema();
+        assertEquals(Type.INT32, elementSchema.type());
+        assertEquals(Time.LOGICAL_NAME, elementSchema.name());
+        java.util.Date expected = new SimpleDateFormat(Values.ISO_8601_TIME_FORMAT_PATTERN).parse(timeStr);
+        assertEquals(Collections.singletonList(expected), result.value());
+    }
+
+    @Test
+    public void shouldParseTimestampStringAsTimestampInArray() throws Exception {
+        String tsStr = "2019-08-23T14:34:54.346Z";
+        String arrayStr = "[" + tsStr + "]";
+        SchemaAndValue result = Values.parseString(arrayStr);
+        assertEquals(Type.ARRAY, result.schema().type());
+        Schema elementSchema = result.schema().valueSchema();
+        assertEquals(Type.INT64, elementSchema.type());
+        assertEquals(Timestamp.LOGICAL_NAME, elementSchema.name());
+        java.util.Date expected = new SimpleDateFormat(Values.ISO_8601_TIMESTAMP_FORMAT_PATTERN).parse(tsStr);
+        assertEquals(Collections.singletonList(expected), result.value());
+    }
+
+    @Test
+    public void shouldParseMultipleTimestampStringAsTimestampInArray() throws Exception {
+        String tsStr1 = "2019-08-23T14:34:54.346Z";
+        String tsStr2 = "2019-01-23T15:12:34.567Z";
+        String tsStr3 = "2019-04-23T19:12:34.567Z";
+        String arrayStr = "[" + tsStr1 + "," + tsStr2 + ",   " + tsStr3 + "]";
+        SchemaAndValue result = Values.parseString(arrayStr);
+        assertEquals(Type.ARRAY, result.schema().type());
+        Schema elementSchema = result.schema().valueSchema();
+        assertEquals(Type.INT64, elementSchema.type());
+        assertEquals(Timestamp.LOGICAL_NAME, elementSchema.name());
+        java.util.Date expected1 = new SimpleDateFormat(Values.ISO_8601_TIMESTAMP_FORMAT_PATTERN).parse(tsStr1);
+        java.util.Date expected2 = new SimpleDateFormat(Values.ISO_8601_TIMESTAMP_FORMAT_PATTERN).parse(tsStr2);
+        java.util.Date expected3 = new SimpleDateFormat(Values.ISO_8601_TIMESTAMP_FORMAT_PATTERN).parse(tsStr3);
+        assertEquals(Arrays.asList(expected1, expected2, expected3), result.value());
+    }
+
+    @Test
+    public void shouldParseQuotedTimeStringAsTimeInMap() throws Exception {
+        String keyStr = "k1";
+        String timeStr = "14:34:54.346Z";
+        String mapStr = "{\"" + keyStr + "\":\"" + timeStr + "\"}";
+        SchemaAndValue result = Values.parseString(mapStr);
+        assertEquals(Type.MAP, result.schema().type());
+        Schema keySchema = result.schema().keySchema();
+        Schema valueSchema = result.schema().valueSchema();
+        assertEquals(Type.STRING, keySchema.type());
+        assertEquals(Type.INT32, valueSchema.type());
+        assertEquals(Time.LOGICAL_NAME, valueSchema.name());
+        java.util.Date expected = new SimpleDateFormat(Values.ISO_8601_TIME_FORMAT_PATTERN).parse(timeStr);
+        assertEquals(Collections.singletonMap(keyStr, expected), result.value());
+    }
+
+    @Test
+    public void shouldParseTimeStringAsTimeInMap() throws Exception {
+        String keyStr = "k1";
+        String timeStr = "14:34:54.346Z";
+        String mapStr = "{\"" + keyStr + "\":" + timeStr + "}";
+        SchemaAndValue result = Values.parseString(mapStr);
+        assertEquals(Type.MAP, result.schema().type());
+        Schema keySchema = result.schema().keySchema();
+        Schema valueSchema = result.schema().valueSchema();
+        assertEquals(Type.STRING, keySchema.type());
+        assertEquals(Type.INT32, valueSchema.type());
+        assertEquals(Time.LOGICAL_NAME, valueSchema.name());
+        java.util.Date expected = new SimpleDateFormat(Values.ISO_8601_TIME_FORMAT_PATTERN).parse(timeStr);
+        assertEquals(Collections.singletonMap(keyStr, expected), result.value());
+    }
+
     /**
      * This is technically invalid JSON, and we don't want to simply ignore the blank elements.
      */
@@ -256,7 +546,7 @@ public void shouldFailToConvertToListFromStringWithNonCommonElementTypeAndBlankE
      */
     @Test(expected = DataException.class)
     public void shouldFailToParseStringOfMapWithIntValuesWithBlankEntry() {
-        Values.convertToList(Schema.STRING_SCHEMA, " { \"foo\" :  1234567890 ,, \"bar\" : 0,  \"baz\" : -987654321 }  ");
+        Values.convertToMap(Schema.STRING_SCHEMA, " { \"foo\" :  1234567890 ,, \"bar\" : 0,  \"baz\" : -987654321 }  ");
     }
 
     /**
@@ -264,7 +554,7 @@ public void shouldFailToParseStringOfMapWithIntValuesWithBlankEntry() {
      */
     @Test(expected = DataException.class)
     public void shouldFailToParseStringOfMalformedMap() {
-        Values.convertToList(Schema.STRING_SCHEMA, " { \"foo\" :  1234567890 , \"a\", \"bar\" : 0,  \"baz\" : -987654321 }  ");
+        Values.convertToMap(Schema.STRING_SCHEMA, " { \"foo\" :  1234567890 , \"a\", \"bar\" : 0,  \"baz\" : -987654321 }  ");
     }
 
     /**
@@ -272,7 +562,7 @@ public void shouldFailToParseStringOfMalformedMap() {
      */
     @Test(expected = DataException.class)
     public void shouldFailToParseStringOfMapWithIntValuesWithOnlyBlankEntries() {
-        Values.convertToList(Schema.STRING_SCHEMA, " { ,,  , , }  ");
+        Values.convertToMap(Schema.STRING_SCHEMA, " { ,,  , , }  ");
     }
 
     /**
@@ -280,15 +570,21 @@ public void shouldFailToParseStringOfMapWithIntValuesWithOnlyBlankEntries() {
      */
     @Test(expected = DataException.class)
     public void shouldFailToParseStringOfMapWithIntValuesWithBlankEntries() {
-        Values.convertToList(Schema.STRING_SCHEMA, " { \"foo\" :  \"1234567890\" ,, \"bar\" : \"0\",  \"baz\" : \"boz\" }  ");
+        Values.convertToMap(Schema.STRING_SCHEMA, " { \"foo\" :  \"1234567890\" ,, \"bar\" : \"0\",  \"baz\" : \"boz\" }  ");
     }
 
-    /**
-     * Schema for Map requires a schema for key and value, but we have no key or value and Connect has no "any" type
-     */
-    @Test(expected = DataException.class)
-    public void shouldFailToParseStringOfEmptyMap() {
-        Values.convertToList(Schema.STRING_SCHEMA, " { }  ");
+    @Test
+    public void shouldConsumeMultipleTokens() {
+        String value = "a:b:c:d:e:f:g:h";
+        Parser parser = new Parser(value);
+        String firstFive = parser.next(5);
+        assertEquals("a:b:c", firstFive);
+        assertEquals(":", parser.next());
+        assertEquals("d", parser.next());
+        assertEquals(":", parser.next());
+        String lastEight = parser.next(8); // only 7 remain
+        assertNull(lastEight);
+        assertEquals("e", parser.next());
     }
 
     @Test
diff --git a/connect/api/src/test/java/org/apache/kafka/connect/header/ConnectHeadersTest.java b/connect/api/src/test/java/org/apache/kafka/connect/header/ConnectHeadersTest.java
index 343bc5d43f59a..72418ba47ffa7 100644
--- a/connect/api/src/test/java/org/apache/kafka/connect/header/ConnectHeadersTest.java
+++ b/connect/api/src/test/java/org/apache/kafka/connect/header/ConnectHeadersTest.java
@@ -118,6 +118,14 @@ public void shouldHaveToString() {
         assertNotNull(headers.toString());
     }
 
+    @Test
+    public void shouldRetainLatestWhenEmpty() {
+        headers.retainLatest(other);
+        headers.retainLatest(key);
+        headers.retainLatest();
+        assertTrue(headers.isEmpty());
+    }
+
     @Test
     public void shouldAddMultipleHeadersWithSameKeyAndRetainLatest() {
         populate(headers);
@@ -179,6 +187,12 @@ public void shouldNotAddHeadersWithObjectValuesAndMismatchedSchema() {
         attemptAndFailToAddHeader("k2", Schema.OPTIONAL_STRING_SCHEMA, 0L);
     }
 
+    @Test
+    public void shouldRemoveAllHeadersWithSameKeyWhenEmpty() {
+        headers.remove(key);
+        assertNoHeaderWithKey(key);
+    }
+
     @Test
     public void shouldRemoveAllHeadersWithSameKey() {
         populate(headers);
@@ -211,6 +225,13 @@ public void shouldRemoveAllHeaders() {
         assertTrue(headers.isEmpty());
     }
 
+    @Test
+    public void shouldTransformHeadersWhenEmpty() {
+        headers.apply(appendToKey("-suffix"));
+        headers.apply(key, appendToKey("-suffix"));
+        assertTrue(headers.isEmpty());
+    }
+
     @Test
     public void shouldTransformHeaders() {
         populate(headers);
@@ -544,4 +565,4 @@ protected void assertHeader(Header header, String key, Schema schema, Object val
         assertSame(schema, header.schema());
         assertSame(value, header.value());
     }
-}
\ No newline at end of file
+}
diff --git a/connect/api/src/test/java/org/apache/kafka/connect/storage/SimpleHeaderConverterTest.java b/connect/api/src/test/java/org/apache/kafka/connect/storage/SimpleHeaderConverterTest.java
index fdfdd32f139f3..58a17f57ba192 100644
--- a/connect/api/src/test/java/org/apache/kafka/connect/storage/SimpleHeaderConverterTest.java
+++ b/connect/api/src/test/java/org/apache/kafka/connect/storage/SimpleHeaderConverterTest.java
@@ -164,11 +164,15 @@ public void shouldConvertListWithIntegerValues() {
     }
 
     @Test
-    public void shouldConvertMapWithStringKeysAndMixedValuesToMapWithoutSchema() {
+    public void shouldConvertMapWithStringKeysAndMixedValuesToMap() {
         Map map = new LinkedHashMap<>();
         map.put("foo", "bar");
         map.put("baz", (short) 3456);
-        assertRoundTrip(null, map);
+        SchemaAndValue result = roundTrip(null, map);
+        assertEquals(Schema.Type.MAP, result.schema().type());
+        assertEquals(Schema.Type.STRING, result.schema().keySchema().type());
+        assertNull(result.schema().valueSchema());
+        assertEquals(map, result.value());
     }
 
     @Test
@@ -176,17 +180,29 @@ public void shouldConvertListWithMixedValuesToListWithoutSchema() {
         List list = new ArrayList<>();
         list.add("foo");
         list.add((short) 13344);
-        assertRoundTrip(null, list);
+        SchemaAndValue result = roundTrip(null, list);
+        assertEquals(Schema.Type.ARRAY, result.schema().type());
+        assertNull(result.schema().valueSchema());
+        assertEquals(list, result.value());
     }
 
     @Test
-    public void shouldConvertEmptyMapToMapWithoutSchema() {
-        assertRoundTrip(null, new LinkedHashMap<>());
+    public void shouldConvertEmptyMapToMap() {
+        Map map = new LinkedHashMap<>();
+        SchemaAndValue result = roundTrip(null, map);
+        assertEquals(Schema.Type.MAP, result.schema().type());
+        assertNull(result.schema().keySchema());
+        assertNull(result.schema().valueSchema());
+        assertEquals(map, result.value());
     }
 
     @Test
-    public void shouldConvertEmptyListToListWithoutSchema() {
-        assertRoundTrip(null, new ArrayList<>());
+    public void shouldConvertEmptyListToList() {
+        List list = new ArrayList<>();
+        SchemaAndValue result = roundTrip(null, list);
+        assertEquals(Schema.Type.ARRAY, result.schema().type());
+        assertNull(result.schema().valueSchema());
+        assertEquals(list, result.value());
     }
 
     protected SchemaAndValue roundTrip(Schema schema, Object input) {
diff --git a/connect/basic-auth-extension/src/main/java/org/apache/kafka/connect/rest/basic/auth/extension/JaasBasicAuthFilter.java b/connect/basic-auth-extension/src/main/java/org/apache/kafka/connect/rest/basic/auth/extension/JaasBasicAuthFilter.java
index 6167434b98031..d5b15c6c65a5f 100644
--- a/connect/basic-auth-extension/src/main/java/org/apache/kafka/connect/rest/basic/auth/extension/JaasBasicAuthFilter.java
+++ b/connect/basic-auth-extension/src/main/java/org/apache/kafka/connect/rest/basic/auth/extension/JaasBasicAuthFilter.java
@@ -17,6 +17,8 @@
 
 package org.apache.kafka.connect.rest.basic.auth.extension;
 
+import java.util.regex.Pattern;
+import javax.ws.rs.HttpMethod;
 import org.apache.kafka.common.config.ConfigException;
 
 import java.io.IOException;
@@ -35,18 +37,18 @@
 import javax.ws.rs.core.Response;
 
 public class JaasBasicAuthFilter implements ContainerRequestFilter {
-
     private static final String CONNECT_LOGIN_MODULE = "KafkaConnect";
     static final String AUTHORIZATION = "Authorization";
-
+    private static final Pattern TASK_REQUEST_PATTERN = Pattern.compile("/?connectors/([^/]+)/tasks/?");
     @Override
     public void filter(ContainerRequestContext requestContext) throws IOException {
-
         try {
-            LoginContext loginContext =
-                new LoginContext(CONNECT_LOGIN_MODULE, new BasicAuthCallBackHandler(
-                    requestContext.getHeaderString(AUTHORIZATION)));
-            loginContext.login();
+            if (!(requestContext.getMethod().equals(HttpMethod.POST) && TASK_REQUEST_PATTERN.matcher(requestContext.getUriInfo().getPath()).matches())) {
+                LoginContext loginContext =
+                    new LoginContext(CONNECT_LOGIN_MODULE, new BasicAuthCallBackHandler(
+                        requestContext.getHeaderString(AUTHORIZATION)));
+                loginContext.login();
+            }
         } catch (LoginException | ConfigException e) {
             requestContext.abortWith(
                 Response.status(Response.Status.UNAUTHORIZED)
diff --git a/connect/basic-auth-extension/src/test/java/org/apache/kafka/connect/rest/basic/auth/extension/JaasBasicAuthFilterTest.java b/connect/basic-auth-extension/src/test/java/org/apache/kafka/connect/rest/basic/auth/extension/JaasBasicAuthFilterTest.java
index 835bef0f37c14..fe5f8b9de5e4e 100644
--- a/connect/basic-auth-extension/src/test/java/org/apache/kafka/connect/rest/basic/auth/extension/JaasBasicAuthFilterTest.java
+++ b/connect/basic-auth-extension/src/test/java/org/apache/kafka/connect/rest/basic/auth/extension/JaasBasicAuthFilterTest.java
@@ -17,12 +17,15 @@
 
 package org.apache.kafka.connect.rest.basic.auth.extension;
 
+import javax.ws.rs.HttpMethod;
+import javax.ws.rs.core.UriInfo;
 import org.apache.kafka.common.security.JaasUtils;
 import org.easymock.EasyMock;
 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.MockStrict;
 import org.powermock.core.classloader.annotations.PowerMockIgnore;
 import org.powermock.modules.junit4.PowerMockRunner;
@@ -52,6 +55,9 @@ public class JaasBasicAuthFilterTest {
     private String previousJaasConfig;
     private Configuration previousConfiguration;
 
+    @MockStrict
+    private UriInfo uriInfo;
+
     @Before
     public void setup() {
         EasyMock.reset(requestContext);
@@ -137,7 +143,34 @@ public void testNoFileOption() throws IOException {
         jaasBasicAuthFilter.filter(requestContext);
     }
 
+    @Test
+    public void testPostWithoutAppropriateCredential() throws IOException {
+        EasyMock.expect(requestContext.getMethod()).andReturn(HttpMethod.POST);
+        EasyMock.expect(requestContext.getUriInfo()).andReturn(uriInfo);
+        EasyMock.expect(uriInfo.getPath()).andReturn("connectors/connName/tasks");
+
+        PowerMock.replayAll();
+        jaasBasicAuthFilter.filter(requestContext);
+        EasyMock.verify(requestContext);
+    }
+
+    @Test
+    public void testPostNotChangingConnectorTask() throws IOException {
+        EasyMock.expect(requestContext.getMethod()).andReturn(HttpMethod.POST);
+        EasyMock.expect(requestContext.getUriInfo()).andReturn(uriInfo);
+        EasyMock.expect(uriInfo.getPath()).andReturn("local:randomport/connectors/connName");
+        String authHeader = "Basic" + Base64.getEncoder().encodeToString(("user" + ":" + "password").getBytes());
+        EasyMock.expect(requestContext.getHeaderString(JaasBasicAuthFilter.AUTHORIZATION))
+            .andReturn(authHeader);
+        requestContext.abortWith(EasyMock.anyObject(Response.class));
+        EasyMock.expectLastCall();
+        PowerMock.replayAll();
+        jaasBasicAuthFilter.filter(requestContext);
+        EasyMock.verify(requestContext);
+    }
+
     private void setMock(String authorization, String username, String password, boolean exceptionCase) {
+        EasyMock.expect(requestContext.getMethod()).andReturn(HttpMethod.GET);
         String authHeader = authorization + " " + Base64.getEncoder().encodeToString((username + ":" + password).getBytes());
         EasyMock.expect(requestContext.getHeaderString(JaasBasicAuthFilter.AUTHORIZATION))
             .andReturn(authHeader);
@@ -152,7 +185,6 @@ private void setupJaasConfig(String loginModule, String credentialFilePath, bool
         File jaasConfigFile = File.createTempFile("ks-jaas-", ".conf");
         jaasConfigFile.deleteOnExit();
         previousJaasConfig = System.setProperty(JaasUtils.JAVA_LOGIN_CONFIG_PARAM, jaasConfigFile.getPath());
-
         List lines;
         lines = new ArrayList<>();
         lines.add(loginModule + " { org.apache.kafka.connect.rest.basic.auth.extension.PropertyFileLoginModule required ");
diff --git a/connect/json/src/main/java/org/apache/kafka/connect/json/DecimalFormat.java b/connect/json/src/main/java/org/apache/kafka/connect/json/DecimalFormat.java
new file mode 100644
index 0000000000000..b4a7fc55d4c61
--- /dev/null
+++ b/connect/json/src/main/java/org/apache/kafka/connect/json/DecimalFormat.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.connect.json;
+
+/**
+ * Represents the valid {@link org.apache.kafka.connect.data.Decimal} serialization formats
+ * in a {@link JsonConverter}.
+ */
+public enum DecimalFormat {
+
+    /**
+     * Serializes the JSON Decimal as a base-64 string. For example, serializing the value
+     * `10.2345` with the BASE64 setting will result in `"D3J5"`.
+     */
+    BASE64,
+
+    /**
+     * Serializes the JSON Decimal as a JSON number. For example, serializing the value
+     * `10.2345` with the NUMERIC setting will result in `10.2345`.
+     */
+    NUMERIC
+}
diff --git a/connect/json/src/main/java/org/apache/kafka/connect/json/JsonConverter.java b/connect/json/src/main/java/org/apache/kafka/connect/json/JsonConverter.java
index 546fcf0daaa68..7af202b7b9ff1 100644
--- a/connect/json/src/main/java/org/apache/kafka/connect/json/JsonConverter.java
+++ b/connect/json/src/main/java/org/apache/kafka/connect/json/JsonConverter.java
@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.connect.json;
 
+import com.fasterxml.jackson.databind.DeserializationFeature;
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.node.ArrayNode;
 import com.fasterxml.jackson.databind.node.JsonNodeFactory;
@@ -51,6 +52,8 @@
 import java.util.Iterator;
 import java.util.Map;
 
+import static org.apache.kafka.common.utils.Utils.mkSet;
+
 /**
  * Implementation of Converter that uses JSON to store schemas and objects. By default this converter will serialize Connect keys, values,
  * and headers with schemas, although this can be disabled with {@link JsonConverterConfig#SCHEMAS_ENABLE_CONFIG schemas.enable}
@@ -185,94 +188,116 @@ public Object convert(Schema schema, JsonNode value) {
         });
     }
 
-    // Convert values in Kafka Connect form into their logical types. These logical converters are discovered by logical type
+    // Convert values in Kafka Connect form into/from their logical types. These logical converters are discovered by logical type
     // names specified in the field
-    private static final HashMap TO_CONNECT_LOGICAL_CONVERTERS = new HashMap<>();
+    private static final HashMap LOGICAL_CONVERTERS = new HashMap<>();
+
+    private static final JsonNodeFactory JSON_NODE_FACTORY = JsonNodeFactory.withExactBigDecimals(true);
+
     static {
-        TO_CONNECT_LOGICAL_CONVERTERS.put(Decimal.LOGICAL_NAME, new LogicalTypeConverter() {
+        LOGICAL_CONVERTERS.put(Decimal.LOGICAL_NAME, new LogicalTypeConverter() {
             @Override
-            public Object convert(Schema schema, Object value) {
-                if (!(value instanceof byte[]))
-                    throw new DataException("Invalid type for Decimal, underlying representation should be bytes but was " + value.getClass());
-                return Decimal.toLogical(schema, (byte[]) value);
-            }
-        });
+            public JsonNode toJson(final Schema schema, final Object value, final JsonConverterConfig config) {
+                if (!(value instanceof BigDecimal))
+                    throw new DataException("Invalid type for Decimal, expected BigDecimal but was " + value.getClass());
 
-        TO_CONNECT_LOGICAL_CONVERTERS.put(Date.LOGICAL_NAME, new LogicalTypeConverter() {
-            @Override
-            public Object convert(Schema schema, Object value) {
-                if (!(value instanceof Integer))
-                    throw new DataException("Invalid type for Date, underlying representation should be int32 but was " + value.getClass());
-                return Date.toLogical(schema, (int) value);
+                final BigDecimal decimal = (BigDecimal) value;
+                switch (config.decimalFormat()) {
+                    case NUMERIC:
+                        return JSON_NODE_FACTORY.numberNode(decimal);
+                    case BASE64:
+                        return JSON_NODE_FACTORY.binaryNode(Decimal.fromLogical(schema, decimal));
+                    default:
+                        throw new DataException("Unexpected " + JsonConverterConfig.DECIMAL_FORMAT_CONFIG + ": " + config.decimalFormat());
+                }
             }
-        });
 
-        TO_CONNECT_LOGICAL_CONVERTERS.put(Time.LOGICAL_NAME, new LogicalTypeConverter() {
             @Override
-            public Object convert(Schema schema, Object value) {
-                if (!(value instanceof Integer))
-                    throw new DataException("Invalid type for Time, underlying representation should be int32 but was " + value.getClass());
-                return Time.toLogical(schema, (int) value);
+            public Object toConnect(final Schema schema, final JsonNode value) {
+                if (value.isNumber()) return value.decimalValue();
+                if (value.isBinary() || value.isTextual()) {
+                    try {
+                        return Decimal.toLogical(schema, value.binaryValue());
+                    } catch (Exception e) {
+                        throw new DataException("Invalid bytes for Decimal field", e);
+                    }
+                }
+
+                throw new DataException("Invalid type for Decimal, underlying representation should be numeric or bytes but was " + value.getNodeType());
             }
         });
 
-        TO_CONNECT_LOGICAL_CONVERTERS.put(Timestamp.LOGICAL_NAME, new LogicalTypeConverter() {
+        LOGICAL_CONVERTERS.put(Date.LOGICAL_NAME, new LogicalTypeConverter() {
             @Override
-            public Object convert(Schema schema, Object value) {
-                if (!(value instanceof Long))
-                    throw new DataException("Invalid type for Timestamp, underlying representation should be int64 but was " + value.getClass());
-                return Timestamp.toLogical(schema, (long) value);
+            public JsonNode toJson(final Schema schema, final Object value, final JsonConverterConfig config) {
+                if (!(value instanceof java.util.Date))
+                    throw new DataException("Invalid type for Date, expected Date but was " + value.getClass());
+                return JSON_NODE_FACTORY.numberNode(Date.fromLogical(schema, (java.util.Date) value));
             }
-        });
-    }
 
-    private static final HashMap TO_JSON_LOGICAL_CONVERTERS = new HashMap<>();
-    static {
-        TO_JSON_LOGICAL_CONVERTERS.put(Decimal.LOGICAL_NAME, new LogicalTypeConverter() {
             @Override
-            public Object convert(Schema schema, Object value) {
-                if (!(value instanceof BigDecimal))
-                    throw new DataException("Invalid type for Decimal, expected BigDecimal but was " + value.getClass());
-                return Decimal.fromLogical(schema, (BigDecimal) value);
+            public Object toConnect(final Schema schema, final JsonNode value) {
+                if (!(value.isInt()))
+                    throw new DataException("Invalid type for Date, underlying representation should be integer but was " + value.getNodeType());
+                return Date.toLogical(schema, value.intValue());
             }
         });
 
-        TO_JSON_LOGICAL_CONVERTERS.put(Date.LOGICAL_NAME, new LogicalTypeConverter() {
+        LOGICAL_CONVERTERS.put(Time.LOGICAL_NAME, new LogicalTypeConverter() {
             @Override
-            public Object convert(Schema schema, Object value) {
+            public JsonNode toJson(final Schema schema, final Object value, final JsonConverterConfig config) {
                 if (!(value instanceof java.util.Date))
-                    throw new DataException("Invalid type for Date, expected Date but was " + value.getClass());
-                return Date.fromLogical(schema, (java.util.Date) value);
+                    throw new DataException("Invalid type for Time, expected Date but was " + value.getClass());
+                return JSON_NODE_FACTORY.numberNode(Time.fromLogical(schema, (java.util.Date) value));
             }
-        });
 
-        TO_JSON_LOGICAL_CONVERTERS.put(Time.LOGICAL_NAME, new LogicalTypeConverter() {
             @Override
-            public Object convert(Schema schema, Object value) {
-                if (!(value instanceof java.util.Date))
-                    throw new DataException("Invalid type for Time, expected Date but was " + value.getClass());
-                return Time.fromLogical(schema, (java.util.Date) value);
+            public Object toConnect(final Schema schema, final JsonNode value) {
+                if (!(value.isInt()))
+                    throw new DataException("Invalid type for Time, underlying representation should be integer but was " + value.getNodeType());
+                return Time.toLogical(schema, value.intValue());
             }
         });
 
-        TO_JSON_LOGICAL_CONVERTERS.put(Timestamp.LOGICAL_NAME, new LogicalTypeConverter() {
+        LOGICAL_CONVERTERS.put(Timestamp.LOGICAL_NAME, new LogicalTypeConverter() {
             @Override
-            public Object convert(Schema schema, Object value) {
+            public JsonNode toJson(final Schema schema, final Object value, final JsonConverterConfig config) {
                 if (!(value instanceof java.util.Date))
                     throw new DataException("Invalid type for Timestamp, expected Date but was " + value.getClass());
-                return Timestamp.fromLogical(schema, (java.util.Date) value);
+                return JSON_NODE_FACTORY.numberNode(Timestamp.fromLogical(schema, (java.util.Date) value));
+            }
+
+            @Override
+            public Object toConnect(final Schema schema, final JsonNode value) {
+                if (!(value.isIntegralNumber()))
+                    throw new DataException("Invalid type for Timestamp, underlying representation should be integral but was " + value.getNodeType());
+                return Timestamp.toLogical(schema, value.longValue());
             }
         });
     }
 
-
-    private boolean enableSchemas = JsonConverterConfig.SCHEMAS_ENABLE_DEFAULT;
-    private int cacheSize = JsonConverterConfig.SCHEMAS_CACHE_SIZE_DEFAULT;
+    private JsonConverterConfig config;
     private Cache fromConnectSchemaCache;
     private Cache toConnectSchemaCache;
 
-    private final JsonSerializer serializer = new JsonSerializer();
-    private final JsonDeserializer deserializer = new JsonDeserializer();
+    private final JsonSerializer serializer;
+    private final JsonDeserializer deserializer;
+
+    public JsonConverter() {
+        serializer = new JsonSerializer(
+            mkSet(),
+            JSON_NODE_FACTORY
+        );
+
+        deserializer = new JsonDeserializer(
+            mkSet(
+                // this ensures that the JsonDeserializer maintains full precision on
+                // floating point numbers that cannot fit into float64
+                DeserializationFeature.USE_BIG_DECIMAL_FOR_FLOATS
+            ),
+            JSON_NODE_FACTORY
+        );
+    }
 
     @Override
     public ConfigDef config() {
@@ -281,16 +306,13 @@ public ConfigDef config() {
 
     @Override
     public void configure(Map configs) {
-        JsonConverterConfig config = new JsonConverterConfig(configs);
-        enableSchemas = config.schemasEnabled();
-        cacheSize = config.schemaCacheSize();
+        config = new JsonConverterConfig(configs);
 
-        boolean isKey = config.type() == ConverterType.KEY;
-        serializer.configure(configs, isKey);
-        deserializer.configure(configs, isKey);
+        serializer.configure(configs, config.type() == ConverterType.KEY);
+        deserializer.configure(configs, config.type() == ConverterType.KEY);
 
-        fromConnectSchemaCache = new SynchronizedCache<>(new LRUCache(cacheSize));
-        toConnectSchemaCache = new SynchronizedCache<>(new LRUCache(cacheSize));
+        fromConnectSchemaCache = new SynchronizedCache<>(new LRUCache<>(config.schemaCacheSize()));
+        toConnectSchemaCache = new SynchronizedCache<>(new LRUCache<>(config.schemaCacheSize()));
     }
 
     @Override
@@ -321,7 +343,7 @@ public byte[] fromConnectData(String topic, Schema schema, Object value) {
             return null;
         }
 
-        JsonNode jsonValue = enableSchemas ? convertToJsonWithEnvelope(schema, value) : convertToJsonWithoutEnvelope(schema, value);
+        JsonNode jsonValue = config.schemasEnabled() ? convertToJsonWithEnvelope(schema, value) : convertToJsonWithoutEnvelope(schema, value);
         try {
             return serializer.serialize(topic, jsonValue);
         } catch (SerializationException e) {
@@ -344,14 +366,14 @@ public SchemaAndValue toConnectData(String topic, byte[] value) {
             throw new DataException("Converting byte[] to Kafka Connect data failed due to serialization error: ", e);
         }
 
-        if (enableSchemas && (!jsonValue.isObject() || jsonValue.size() != 2 || !jsonValue.has(JsonSchema.ENVELOPE_SCHEMA_FIELD_NAME) || !jsonValue.has(JsonSchema.ENVELOPE_PAYLOAD_FIELD_NAME)))
+        if (config.schemasEnabled() && (!jsonValue.isObject() || jsonValue.size() != 2 || !jsonValue.has(JsonSchema.ENVELOPE_SCHEMA_FIELD_NAME) || !jsonValue.has(JsonSchema.ENVELOPE_PAYLOAD_FIELD_NAME)))
             throw new DataException("JsonConverter with schemas.enable requires \"schema\" and \"payload\" fields and may not contain additional fields." +
                     " If you are trying to deserialize plain JSON data, set schemas.enable=false in your converter configuration.");
 
         // The deserialized data should either be an envelope object containing the schema and the payload or the schema
         // was stripped during serialization and we need to fill in an all-encompassing schema.
-        if (!enableSchemas) {
-            ObjectNode envelope = JsonNodeFactory.instance.objectNode();
+        if (!config.schemasEnabled()) {
+            ObjectNode envelope = JSON_NODE_FACTORY.objectNode();
             envelope.set(JsonSchema.ENVELOPE_SCHEMA_FIELD_NAME, null);
             envelope.set(JsonSchema.ENVELOPE_PAYLOAD_FIELD_NAME, jsonValue);
             jsonValue = envelope;
@@ -402,17 +424,17 @@ public ObjectNode asJsonSchema(Schema schema) {
                 jsonSchema = JsonSchema.STRING_SCHEMA.deepCopy();
                 break;
             case ARRAY:
-                jsonSchema = JsonNodeFactory.instance.objectNode().put(JsonSchema.SCHEMA_TYPE_FIELD_NAME, JsonSchema.ARRAY_TYPE_NAME);
+                jsonSchema = JSON_NODE_FACTORY.objectNode().put(JsonSchema.SCHEMA_TYPE_FIELD_NAME, JsonSchema.ARRAY_TYPE_NAME);
                 jsonSchema.set(JsonSchema.ARRAY_ITEMS_FIELD_NAME, asJsonSchema(schema.valueSchema()));
                 break;
             case MAP:
-                jsonSchema = JsonNodeFactory.instance.objectNode().put(JsonSchema.SCHEMA_TYPE_FIELD_NAME, JsonSchema.MAP_TYPE_NAME);
+                jsonSchema = JSON_NODE_FACTORY.objectNode().put(JsonSchema.SCHEMA_TYPE_FIELD_NAME, JsonSchema.MAP_TYPE_NAME);
                 jsonSchema.set(JsonSchema.MAP_KEY_FIELD_NAME, asJsonSchema(schema.keySchema()));
                 jsonSchema.set(JsonSchema.MAP_VALUE_FIELD_NAME, asJsonSchema(schema.valueSchema()));
                 break;
             case STRUCT:
-                jsonSchema = JsonNodeFactory.instance.objectNode().put(JsonSchema.SCHEMA_TYPE_FIELD_NAME, JsonSchema.STRUCT_TYPE_NAME);
-                ArrayNode fields = JsonNodeFactory.instance.arrayNode();
+                jsonSchema = JSON_NODE_FACTORY.objectNode().put(JsonSchema.SCHEMA_TYPE_FIELD_NAME, JsonSchema.STRUCT_TYPE_NAME);
+                ArrayNode fields = JSON_NODE_FACTORY.arrayNode();
                 for (Field field : schema.fields()) {
                     ObjectNode fieldJsonSchema = asJsonSchema(field.schema()).deepCopy();
                     fieldJsonSchema.put(JsonSchema.STRUCT_FIELD_NAME_FIELD_NAME, field.name());
@@ -432,7 +454,7 @@ public ObjectNode asJsonSchema(Schema schema) {
         if (schema.doc() != null)
             jsonSchema.put(JsonSchema.SCHEMA_DOC_FIELD_NAME, schema.doc());
         if (schema.parameters() != null) {
-            ObjectNode jsonSchemaParams = JsonNodeFactory.instance.objectNode();
+            ObjectNode jsonSchemaParams = JSON_NODE_FACTORY.objectNode();
             for (Map.Entry prop : schema.parameters().entrySet())
                 jsonSchemaParams.put(prop.getKey(), prop.getValue());
             jsonSchema.set(JsonSchema.SCHEMA_PARAMETERS_FIELD_NAME, jsonSchemaParams);
@@ -578,22 +600,21 @@ private JsonNode convertToJsonWithoutEnvelope(Schema schema, Object value) {
      * Convert this object, in the org.apache.kafka.connect.data format, into a JSON object, returning both the schema
      * and the converted object.
      */
-    private static JsonNode convertToJson(Schema schema, Object logicalValue) {
-        if (logicalValue == null) {
+    private JsonNode convertToJson(Schema schema, Object value) {
+        if (value == null) {
             if (schema == null) // Any schema is valid and we don't have a default, so treat this as an optional schema
                 return null;
             if (schema.defaultValue() != null)
                 return convertToJson(schema, schema.defaultValue());
             if (schema.isOptional())
-                return JsonNodeFactory.instance.nullNode();
+                return JSON_NODE_FACTORY.nullNode();
             throw new DataException("Conversion error: null value for field that is required and has no default value");
         }
 
-        Object value = logicalValue;
         if (schema != null && schema.name() != null) {
-            LogicalTypeConverter logicalConverter = TO_JSON_LOGICAL_CONVERTERS.get(schema.name());
+            LogicalTypeConverter logicalConverter = LOGICAL_CONVERTERS.get(schema.name());
             if (logicalConverter != null)
-                value = logicalConverter.convert(schema, logicalValue);
+                return logicalConverter.toJson(schema, value, config);
         }
 
         try {
@@ -607,32 +628,32 @@ private static JsonNode convertToJson(Schema schema, Object logicalValue) {
             }
             switch (schemaType) {
                 case INT8:
-                    return JsonNodeFactory.instance.numberNode((Byte) value);
+                    return JSON_NODE_FACTORY.numberNode((Byte) value);
                 case INT16:
-                    return JsonNodeFactory.instance.numberNode((Short) value);
+                    return JSON_NODE_FACTORY.numberNode((Short) value);
                 case INT32:
-                    return JsonNodeFactory.instance.numberNode((Integer) value);
+                    return JSON_NODE_FACTORY.numberNode((Integer) value);
                 case INT64:
-                    return JsonNodeFactory.instance.numberNode((Long) value);
+                    return JSON_NODE_FACTORY.numberNode((Long) value);
                 case FLOAT32:
-                    return JsonNodeFactory.instance.numberNode((Float) value);
+                    return JSON_NODE_FACTORY.numberNode((Float) value);
                 case FLOAT64:
-                    return JsonNodeFactory.instance.numberNode((Double) value);
+                    return JSON_NODE_FACTORY.numberNode((Double) value);
                 case BOOLEAN:
-                    return JsonNodeFactory.instance.booleanNode((Boolean) value);
+                    return JSON_NODE_FACTORY.booleanNode((Boolean) value);
                 case STRING:
                     CharSequence charSeq = (CharSequence) value;
-                    return JsonNodeFactory.instance.textNode(charSeq.toString());
+                    return JSON_NODE_FACTORY.textNode(charSeq.toString());
                 case BYTES:
                     if (value instanceof byte[])
-                        return JsonNodeFactory.instance.binaryNode((byte[]) value);
+                        return JSON_NODE_FACTORY.binaryNode((byte[]) value);
                     else if (value instanceof ByteBuffer)
-                        return JsonNodeFactory.instance.binaryNode(((ByteBuffer) value).array());
+                        return JSON_NODE_FACTORY.binaryNode(((ByteBuffer) value).array());
                     else
                         throw new DataException("Invalid type for bytes type: " + value.getClass());
                 case ARRAY: {
                     Collection collection = (Collection) value;
-                    ArrayNode list = JsonNodeFactory.instance.arrayNode();
+                    ArrayNode list = JSON_NODE_FACTORY.arrayNode();
                     for (Object elem : collection) {
                         Schema valueSchema = schema == null ? null : schema.valueSchema();
                         JsonNode fieldValue = convertToJson(valueSchema, elem);
@@ -658,9 +679,9 @@ else if (value instanceof ByteBuffer)
                     ObjectNode obj = null;
                     ArrayNode list = null;
                     if (objectMode)
-                        obj = JsonNodeFactory.instance.objectNode();
+                        obj = JSON_NODE_FACTORY.objectNode();
                     else
-                        list = JsonNodeFactory.instance.arrayNode();
+                        list = JSON_NODE_FACTORY.arrayNode();
                     for (Map.Entry entry : map.entrySet()) {
                         Schema keySchema = schema == null ? null : schema.keySchema();
                         Schema valueSchema = schema == null ? null : schema.valueSchema();
@@ -670,7 +691,7 @@ else if (value instanceof ByteBuffer)
                         if (objectMode)
                             obj.set(mapKey.asText(), mapValue);
                         else
-                            list.add(JsonNodeFactory.instance.arrayNode().add(mapKey).add(mapValue));
+                            list.add(JSON_NODE_FACTORY.arrayNode().add(mapKey).add(mapValue));
                     }
                     return objectMode ? obj : list;
                 }
@@ -678,7 +699,7 @@ else if (value instanceof ByteBuffer)
                     Struct struct = (Struct) value;
                     if (!struct.schema().equals(schema))
                         throw new DataException("Mismatching schema.");
-                    ObjectNode obj = JsonNodeFactory.instance.objectNode();
+                    ObjectNode obj = JSON_NODE_FACTORY.objectNode();
                     for (Field field : schema.fields()) {
                         obj.set(field.name(), convertToJson(field.schema(), struct.get(field)));
                     }
@@ -698,7 +719,7 @@ private static Object convertToConnect(Schema schema, JsonNode jsonValue) {
         final Schema.Type schemaType;
         if (schema != null) {
             schemaType = schema.type();
-            if (jsonValue.isNull()) {
+            if (jsonValue == null || jsonValue.isNull()) {
                 if (schema.defaultValue() != null)
                     return schema.defaultValue(); // any logical type conversions should already have been applied
                 if (schema.isOptional())
@@ -742,13 +763,13 @@ private static Object convertToConnect(Schema schema, JsonNode jsonValue) {
         if (typeConverter == null)
             throw new DataException("Unknown schema type: " + String.valueOf(schemaType));
 
-        Object converted = typeConverter.convert(schema, jsonValue);
         if (schema != null && schema.name() != null) {
-            LogicalTypeConverter logicalConverter = TO_CONNECT_LOGICAL_CONVERTERS.get(schema.name());
+            LogicalTypeConverter logicalConverter = LOGICAL_CONVERTERS.get(schema.name());
             if (logicalConverter != null)
-                converted = logicalConverter.convert(schema, converted);
+                return logicalConverter.toConnect(schema, jsonValue);
         }
-        return converted;
+
+        return typeConverter.convert(schema, jsonValue);
     }
 
     private interface JsonToConnectTypeConverter {
@@ -756,6 +777,7 @@ private interface JsonToConnectTypeConverter {
     }
 
     private interface LogicalTypeConverter {
-        Object convert(Schema schema, Object value);
+        JsonNode toJson(Schema schema, Object value, JsonConverterConfig config);
+        Object toConnect(Schema schema, JsonNode value);
     }
 }
diff --git a/connect/json/src/main/java/org/apache/kafka/connect/json/JsonConverterConfig.java b/connect/json/src/main/java/org/apache/kafka/connect/json/JsonConverterConfig.java
index 7f1dda2c4c1b1..efb497991cbc9 100644
--- a/connect/json/src/main/java/org/apache/kafka/connect/json/JsonConverterConfig.java
+++ b/connect/json/src/main/java/org/apache/kafka/connect/json/JsonConverterConfig.java
@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.connect.json;
 
+import java.util.Locale;
 import org.apache.kafka.common.config.ConfigDef;
 import org.apache.kafka.common.config.ConfigDef.Importance;
 import org.apache.kafka.common.config.ConfigDef.Type;
@@ -39,6 +40,12 @@ public class JsonConverterConfig extends ConverterConfig {
     private static final String SCHEMAS_CACHE_SIZE_DOC = "The maximum number of schemas that can be cached in this converter instance.";
     private static final String SCHEMAS_CACHE_SIZE_DISPLAY = "Schema Cache Size";
 
+    public static final String DECIMAL_FORMAT_CONFIG = "decimal.format";
+    public static final String DECIMAL_FORMAT_DEFAULT = DecimalFormat.BASE64.name();
+    private static final String DECIMAL_FORMAT_DOC = "Controls which format this converter will serialize decimals in."
+        + " This value is case insensitive and can be either 'BASE64' (default) or 'NUMERIC'";
+    private static final String DECIMAL_FORMAT_DISPLAY = "Decimal Format";
+
     private final static ConfigDef CONFIG;
 
     static {
@@ -49,14 +56,32 @@ public class JsonConverterConfig extends ConverterConfig {
                       orderInGroup++, Width.MEDIUM, SCHEMAS_ENABLE_DISPLAY);
         CONFIG.define(SCHEMAS_CACHE_SIZE_CONFIG, Type.INT, SCHEMAS_CACHE_SIZE_DEFAULT, Importance.HIGH, SCHEMAS_CACHE_SIZE_DOC, group,
                       orderInGroup++, Width.MEDIUM, SCHEMAS_CACHE_SIZE_DISPLAY);
+
+        group = "Serialization";
+        orderInGroup = 0;
+        CONFIG.define(
+            DECIMAL_FORMAT_CONFIG, Type.STRING, DECIMAL_FORMAT_DEFAULT,
+            ConfigDef.CaseInsensitiveValidString.in(
+                DecimalFormat.BASE64.name(),
+                DecimalFormat.NUMERIC.name()),
+            Importance.LOW, DECIMAL_FORMAT_DOC, group, orderInGroup++,
+            Width.MEDIUM, DECIMAL_FORMAT_DISPLAY);
     }
 
     public static ConfigDef configDef() {
         return CONFIG;
     }
 
+    // cached config values
+    private final boolean schemasEnabled;
+    private final int schemaCacheSize;
+    private final DecimalFormat decimalFormat;
+
     public JsonConverterConfig(Map props) {
         super(CONFIG, props);
+        this.schemasEnabled = getBoolean(SCHEMAS_ENABLE_CONFIG);
+        this.schemaCacheSize = getInt(SCHEMAS_CACHE_SIZE_CONFIG);
+        this.decimalFormat = DecimalFormat.valueOf(getString(DECIMAL_FORMAT_CONFIG).toUpperCase(Locale.ROOT));
     }
 
     /**
@@ -65,7 +90,7 @@ public JsonConverterConfig(Map props) {
      * @return true if enabled, or false otherwise
      */
     public boolean schemasEnabled() {
-        return getBoolean(SCHEMAS_ENABLE_CONFIG);
+        return schemasEnabled;
     }
 
     /**
@@ -74,6 +99,16 @@ public boolean schemasEnabled() {
      * @return the cache size
      */
     public int schemaCacheSize() {
-        return getInt(SCHEMAS_CACHE_SIZE_CONFIG);
+        return schemaCacheSize;
     }
+
+    /**
+     * Get the serialization format for decimal types.
+     *
+     * @return the decimal serialization format
+     */
+    public DecimalFormat decimalFormat() {
+        return decimalFormat;
+    }
+
 }
diff --git a/connect/json/src/main/java/org/apache/kafka/connect/json/JsonDeserializer.java b/connect/json/src/main/java/org/apache/kafka/connect/json/JsonDeserializer.java
index b006e22721519..2e6e821b2d3bc 100644
--- a/connect/json/src/main/java/org/apache/kafka/connect/json/JsonDeserializer.java
+++ b/connect/json/src/main/java/org/apache/kafka/connect/json/JsonDeserializer.java
@@ -16,8 +16,12 @@
  */
 package org.apache.kafka.connect.json;
 
+import com.fasterxml.jackson.databind.DeserializationFeature;
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.node.JsonNodeFactory;
+import java.util.Collections;
+import java.util.Set;
 import org.apache.kafka.common.errors.SerializationException;
 import org.apache.kafka.common.serialization.Deserializer;
 
@@ -26,14 +30,29 @@
  * structured data without having associated Java classes. This deserializer also supports Connect schemas.
  */
 public class JsonDeserializer implements Deserializer {
-    private ObjectMapper objectMapper = new ObjectMapper();
+    private final ObjectMapper objectMapper = new ObjectMapper();
 
     /**
      * Default constructor needed by Kafka
      */
     public JsonDeserializer() {
+        this(Collections.emptySet(), JsonNodeFactory.withExactBigDecimals(true));
     }
 
+    /**
+     * A constructor that additionally specifies some {@link DeserializationFeature}
+     * for the deserializer
+     *
+     * @param deserializationFeatures the specified deserialization features
+     * @param jsonNodeFactory the json node factory to use.
+     */
+    JsonDeserializer(
+        final Set deserializationFeatures,
+        final JsonNodeFactory jsonNodeFactory
+    ) {
+        deserializationFeatures.forEach(objectMapper::enable);
+        objectMapper.setNodeFactory(jsonNodeFactory);
+    }
 
     @Override
     public JsonNode deserialize(String topic, byte[] bytes) {
diff --git a/connect/json/src/main/java/org/apache/kafka/connect/json/JsonSerializer.java b/connect/json/src/main/java/org/apache/kafka/connect/json/JsonSerializer.java
index 94ec0a83e8ce6..0f2b62bd0a40e 100644
--- a/connect/json/src/main/java/org/apache/kafka/connect/json/JsonSerializer.java
+++ b/connect/json/src/main/java/org/apache/kafka/connect/json/JsonSerializer.java
@@ -18,9 +18,14 @@
 
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.SerializationFeature;
+import com.fasterxml.jackson.databind.node.JsonNodeFactory;
 import org.apache.kafka.common.errors.SerializationException;
 import org.apache.kafka.common.serialization.Serializer;
 
+import java.util.Collections;
+import java.util.Set;
+
 /**
  * Serialize Jackson JsonNode tree model objects to UTF-8 JSON. Using the tree model allows handling arbitrarily
  * structured data without corresponding Java classes. This serializer also supports Connect schemas.
@@ -32,7 +37,22 @@ public class JsonSerializer implements Serializer {
      * Default constructor needed by Kafka
      */
     public JsonSerializer() {
+        this(Collections.emptySet(), JsonNodeFactory.withExactBigDecimals(true));
+    }
 
+    /**
+     * A constructor that additionally specifies some {@link SerializationFeature}
+     * for the serializer
+     *
+     * @param serializationFeatures the specified serialization features
+     * @param jsonNodeFactory the json node factory to use.
+     */
+    JsonSerializer(
+        final Set serializationFeatures,
+        final JsonNodeFactory jsonNodeFactory
+    ) {
+        serializationFeatures.forEach(objectMapper::enable);
+        objectMapper.setNodeFactory(jsonNodeFactory);
     }
 
     @Override
diff --git a/connect/json/src/test/java/org/apache/kafka/connect/json/JsonConverterConfigTest.java b/connect/json/src/test/java/org/apache/kafka/connect/json/JsonConverterConfigTest.java
new file mode 100644
index 0000000000000..590707bffab4b
--- /dev/null
+++ b/connect/json/src/test/java/org/apache/kafka/connect/json/JsonConverterConfigTest.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.connect.json;
+
+import static org.junit.Assert.assertEquals;
+
+import java.util.HashMap;
+import java.util.Map;
+import org.apache.kafka.connect.storage.ConverterConfig;
+import org.apache.kafka.connect.storage.ConverterType;
+import org.junit.Test;
+
+public class JsonConverterConfigTest {
+
+    @Test
+    public void shouldBeCaseInsensitiveForDecimalFormatConfig() {
+        final Map configValues = new HashMap<>();
+        configValues.put(ConverterConfig.TYPE_CONFIG, ConverterType.KEY.getName());
+        configValues.put(JsonConverterConfig.DECIMAL_FORMAT_CONFIG, "NuMeRiC");
+
+        final JsonConverterConfig config = new JsonConverterConfig(configValues);
+        assertEquals(config.decimalFormat(), DecimalFormat.NUMERIC);
+    }
+
+}
\ No newline at end of file
diff --git a/connect/json/src/test/java/org/apache/kafka/connect/json/JsonConverterTest.java b/connect/json/src/test/java/org/apache/kafka/connect/json/JsonConverterTest.java
index 003d7e610872c..2e189e2d584ae 100644
--- a/connect/json/src/test/java/org/apache/kafka/connect/json/JsonConverterTest.java
+++ b/connect/json/src/test/java/org/apache/kafka/connect/json/JsonConverterTest.java
@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.connect.json;
 
+import com.fasterxml.jackson.databind.DeserializationFeature;
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.fasterxml.jackson.databind.node.ArrayNode;
@@ -58,14 +59,18 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
 public class JsonConverterTest {
     private static final String TOPIC = "topic";
 
-    ObjectMapper objectMapper = new ObjectMapper();
-    JsonConverter converter = new JsonConverter();
+    private final ObjectMapper objectMapper = new ObjectMapper()
+        .enable(DeserializationFeature.USE_BIG_DECIMAL_FOR_FLOATS)
+        .setNodeFactory(JsonNodeFactory.withExactBigDecimals(true));
+
+    private final JsonConverter converter = new JsonConverter();
 
     @Before
     public void setUp() {
@@ -172,6 +177,15 @@ public void structToConnect() {
         assertEquals(new SchemaAndValue(expectedSchema, expected), converted);
     }
 
+    @Test
+    public void structWithOptionalFieldToConnect() {
+        byte[] structJson = "{ \"schema\": { \"type\": \"struct\", \"fields\": [{ \"field\":\"optional\", \"type\": \"string\", \"optional\": true }, {  \"field\": \"required\", \"type\": \"string\" }] }, \"payload\": { \"required\": \"required\" } }".getBytes();
+        Schema expectedSchema = SchemaBuilder.struct().field("optional", Schema.OPTIONAL_STRING_SCHEMA).field("required", Schema.STRING_SCHEMA).build();
+        Struct expected = new Struct(expectedSchema).put("required", "required");
+        SchemaAndValue converted = converter.toConnectData(TOPIC, structJson);
+        assertEquals(new SchemaAndValue(expectedSchema, expected), converted);
+    }
+
     @Test
     public void nullToConnect() {
         // When schemas are enabled, trying to decode a tombstone should be an empty envelope
@@ -252,6 +266,37 @@ public void decimalToConnectOptionalWithDefaultValue() {
         assertEquals(reference, schemaAndValue.value());
     }
 
+    @Test
+    public void numericDecimalToConnect() {
+        BigDecimal reference = new BigDecimal(new BigInteger("156"), 2);
+        Schema schema = Decimal.schema(2);
+        String msg = "{ \"schema\": { \"type\": \"bytes\", \"name\": \"org.apache.kafka.connect.data.Decimal\", \"version\": 1, \"parameters\": { \"scale\": \"2\" } }, \"payload\": 1.56 }";
+        SchemaAndValue schemaAndValue = converter.toConnectData(TOPIC, msg.getBytes());
+        assertEquals(schema, schemaAndValue.schema());
+        assertEquals(reference, schemaAndValue.value());
+    }
+
+    @Test
+    public void numericDecimalWithTrailingZerosToConnect() {
+        BigDecimal reference = new BigDecimal(new BigInteger("15600"), 4);
+        Schema schema = Decimal.schema(4);
+        String msg = "{ \"schema\": { \"type\": \"bytes\", \"name\": \"org.apache.kafka.connect.data.Decimal\", \"version\": 1, \"parameters\": { \"scale\": \"4\" } }, \"payload\": 1.5600 }";
+        SchemaAndValue schemaAndValue = converter.toConnectData(TOPIC, msg.getBytes());
+        assertEquals(schema, schemaAndValue.schema());
+        assertEquals(reference, schemaAndValue.value());
+    }
+
+    @Test
+    public void highPrecisionNumericDecimalToConnect() {
+        // this number is too big to be kept in a float64!
+        BigDecimal reference = new BigDecimal("1.23456789123456789");
+        Schema schema = Decimal.schema(17);
+        String msg = "{ \"schema\": { \"type\": \"bytes\", \"name\": \"org.apache.kafka.connect.data.Decimal\", \"version\": 1, \"parameters\": { \"scale\": \"17\" } }, \"payload\": 1.23456789123456789 }";
+        SchemaAndValue schemaAndValue = converter.toConnectData(TOPIC, msg.getBytes());
+        assertEquals(schema, schemaAndValue.schema());
+        assertEquals(reference, schemaAndValue.value());
+    }
+
     @Test
     public void dateToConnect() {
         Schema schema = Date.SCHEMA;
@@ -587,9 +632,40 @@ public void decimalToJson() throws IOException {
         validateEnvelope(converted);
         assertEquals(parse("{ \"type\": \"bytes\", \"optional\": false, \"name\": \"org.apache.kafka.connect.data.Decimal\", \"version\": 1, \"parameters\": { \"scale\": \"2\" } }"),
                 converted.get(JsonSchema.ENVELOPE_SCHEMA_FIELD_NAME));
+        assertTrue("expected node to be base64 text", converted.get(JsonSchema.ENVELOPE_PAYLOAD_FIELD_NAME).isTextual());
         assertArrayEquals(new byte[]{0, -100}, converted.get(JsonSchema.ENVELOPE_PAYLOAD_FIELD_NAME).binaryValue());
     }
 
+    @Test
+    public void decimalToNumericJson() {
+        converter.configure(Collections.singletonMap(JsonConverterConfig.DECIMAL_FORMAT_CONFIG, DecimalFormat.NUMERIC.name()), false);
+        JsonNode converted = parse(converter.fromConnectData(TOPIC, Decimal.schema(2), new BigDecimal(new BigInteger("156"), 2)));
+        validateEnvelope(converted);
+        assertEquals(parse("{ \"type\": \"bytes\", \"optional\": false, \"name\": \"org.apache.kafka.connect.data.Decimal\", \"version\": 1, \"parameters\": { \"scale\": \"2\" } }"),
+            converted.get(JsonSchema.ENVELOPE_SCHEMA_FIELD_NAME));
+        assertTrue("expected node to be numeric", converted.get(JsonSchema.ENVELOPE_PAYLOAD_FIELD_NAME).isNumber());
+        assertEquals(new BigDecimal("1.56"), converted.get(JsonSchema.ENVELOPE_PAYLOAD_FIELD_NAME).decimalValue());
+    }
+
+    @Test
+    public void decimalWithTrailingZerosToNumericJson() {
+        converter.configure(Collections.singletonMap(JsonConverterConfig.DECIMAL_FORMAT_CONFIG, DecimalFormat.NUMERIC.name()), false);
+        JsonNode converted = parse(converter.fromConnectData(TOPIC, Decimal.schema(4), new BigDecimal(new BigInteger("15600"), 4)));
+        validateEnvelope(converted);
+        assertEquals(parse("{ \"type\": \"bytes\", \"optional\": false, \"name\": \"org.apache.kafka.connect.data.Decimal\", \"version\": 1, \"parameters\": { \"scale\": \"4\" } }"),
+            converted.get(JsonSchema.ENVELOPE_SCHEMA_FIELD_NAME));
+        assertTrue("expected node to be numeric", converted.get(JsonSchema.ENVELOPE_PAYLOAD_FIELD_NAME).isNumber());
+        assertEquals(new BigDecimal("1.5600"), converted.get(JsonSchema.ENVELOPE_PAYLOAD_FIELD_NAME).decimalValue());
+    }
+
+    @Test
+    public void decimalToJsonWithoutSchema() {
+        assertThrows(
+            "expected data exception when serializing BigDecimal without schema",
+            DataException.class,
+            () -> converter.fromConnectData(TOPIC, null, new BigDecimal(new BigInteger("156"), 2)));
+    }
+
     @Test
     public void dateToJson() {
         GregorianCalendar calendar = new GregorianCalendar(1970, Calendar.JANUARY, 1, 0, 0, 0);
diff --git a/connect/mirror-client/src/main/java/org/apache/kafka/connect/mirror/Checkpoint.java b/connect/mirror-client/src/main/java/org/apache/kafka/connect/mirror/Checkpoint.java
new file mode 100644
index 0000000000000..74db7461ef0b6
--- /dev/null
+++ b/connect/mirror-client/src/main/java/org/apache/kafka/connect/mirror/Checkpoint.java
@@ -0,0 +1,184 @@
+/*
+ * 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.mirror;
+
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.protocol.types.Field;
+import org.apache.kafka.common.protocol.types.Schema;
+import org.apache.kafka.common.protocol.types.Struct;
+import org.apache.kafka.common.protocol.types.Type;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.OffsetAndMetadata;
+
+import java.util.Map;
+import java.util.HashMap;
+import java.nio.ByteBuffer;
+
+/** Checkpoint records emitted from MirrorCheckpointConnector. Encodes remote consumer group state. */
+public class Checkpoint {
+    public static final String TOPIC_KEY = "topic";
+    public static final String PARTITION_KEY = "partition";
+    public static final String CONSUMER_GROUP_ID_KEY = "group";
+    public static final String UPSTREAM_OFFSET_KEY = "upstreamOffset";
+    public static final String DOWNSTREAM_OFFSET_KEY = "offset";
+    public static final String METADATA_KEY = "metadata";
+    public static final String VERSION_KEY = "version";
+    public static final short VERSION = 0;
+
+    public static final Schema VALUE_SCHEMA_V0 = new Schema(
+            new Field(UPSTREAM_OFFSET_KEY, Type.INT64),
+            new Field(DOWNSTREAM_OFFSET_KEY, Type.INT64),
+            new Field(METADATA_KEY, Type.STRING));
+
+    public static final Schema KEY_SCHEMA = new Schema(
+            new Field(CONSUMER_GROUP_ID_KEY, Type.STRING),
+            new Field(TOPIC_KEY, Type.STRING),
+            new Field(PARTITION_KEY, Type.INT32));
+
+    public static final Schema HEADER_SCHEMA = new Schema(
+            new Field(VERSION_KEY, Type.INT16));
+
+    private String consumerGroupId;
+    private TopicPartition topicPartition;
+    private long upstreamOffset;
+    private long downstreamOffset;
+    private String metadata;
+
+    public Checkpoint(String consumerGroupId, TopicPartition topicPartition, long upstreamOffset,
+            long downstreamOffset, String metadata) {
+        this.consumerGroupId = consumerGroupId;
+        this.topicPartition = topicPartition;
+        this.upstreamOffset = upstreamOffset;
+        this.downstreamOffset = downstreamOffset;
+        this.metadata = metadata;
+    }
+
+    public String consumerGroupId() {
+        return consumerGroupId;
+    }
+
+    public TopicPartition topicPartition() {
+        return topicPartition;
+    }
+
+    public long upstreamOffset() {
+        return upstreamOffset;
+    }
+
+    public long downstreamOffset() {
+        return downstreamOffset;
+    }
+
+    public String metadata() {
+        return metadata;
+    }
+
+    public OffsetAndMetadata offsetAndMetadata() {
+        return new OffsetAndMetadata(downstreamOffset, metadata);
+    }
+
+    @Override
+    public String toString() {
+        return String.format("Checkpoint{consumerGroupId=%s, topicPartition=%s, "
+            + "upstreamOffset=%d, downstreamOffset=%d, metatadata=%s}",
+            consumerGroupId, topicPartition, upstreamOffset, downstreamOffset, metadata);
+    }
+
+    ByteBuffer serializeValue(short version) {
+        Struct header = headerStruct(version);
+        Schema valueSchema = valueSchema(version);
+        Struct valueStruct = valueStruct(valueSchema);
+        ByteBuffer buffer = ByteBuffer.allocate(HEADER_SCHEMA.sizeOf(header) + valueSchema.sizeOf(valueStruct));
+        HEADER_SCHEMA.write(buffer, header);
+        valueSchema.write(buffer, valueStruct);
+        buffer.flip();
+        return buffer;
+    }
+
+    ByteBuffer serializeKey() {
+        Struct struct = keyStruct();
+        ByteBuffer buffer = ByteBuffer.allocate(KEY_SCHEMA.sizeOf(struct));
+        KEY_SCHEMA.write(buffer, struct);
+        buffer.flip();
+        return buffer;
+    }
+
+    public static Checkpoint deserializeRecord(ConsumerRecord record) {
+        ByteBuffer value = ByteBuffer.wrap(record.value());
+        Struct header = HEADER_SCHEMA.read(value);
+        short version = header.getShort(VERSION_KEY);
+        Schema valueSchema = valueSchema(version);
+        Struct valueStruct = valueSchema.read(value);
+        long upstreamOffset = valueStruct.getLong(UPSTREAM_OFFSET_KEY);
+        long downstreamOffset = valueStruct.getLong(DOWNSTREAM_OFFSET_KEY);
+        String metadata = valueStruct.getString(METADATA_KEY);
+        Struct keyStruct = KEY_SCHEMA.read(ByteBuffer.wrap(record.key()));
+        String group = keyStruct.getString(CONSUMER_GROUP_ID_KEY);
+        String topic = keyStruct.getString(TOPIC_KEY);
+        int partition = keyStruct.getInt(PARTITION_KEY);
+        return new Checkpoint(group, new TopicPartition(topic, partition), upstreamOffset,
+            downstreamOffset, metadata);
+    }
+
+    private static Schema valueSchema(short version) {
+        assert version == 0;
+        return VALUE_SCHEMA_V0;
+    }
+
+    private Struct valueStruct(Schema schema) {
+        Struct struct = new Struct(schema);
+        struct.set(UPSTREAM_OFFSET_KEY, upstreamOffset);
+        struct.set(DOWNSTREAM_OFFSET_KEY, downstreamOffset);
+        struct.set(METADATA_KEY, metadata);
+        return struct;
+    }
+
+    private Struct keyStruct() {
+        Struct struct = new Struct(KEY_SCHEMA);
+        struct.set(CONSUMER_GROUP_ID_KEY, consumerGroupId);
+        struct.set(TOPIC_KEY, topicPartition.topic());
+        struct.set(PARTITION_KEY, topicPartition.partition());
+        return struct;
+    }
+
+    private Struct headerStruct(short version) {
+        Struct struct = new Struct(HEADER_SCHEMA);
+        struct.set(VERSION_KEY, version);
+        return struct;
+    }
+
+    Map connectPartition() {
+        Map partition = new HashMap<>();
+        partition.put(CONSUMER_GROUP_ID_KEY, consumerGroupId);
+        partition.put(TOPIC_KEY, topicPartition.topic());
+        partition.put(PARTITION_KEY, topicPartition.partition());
+        return partition;
+    }
+
+    static String unwrapGroup(Map connectPartition) {
+        return connectPartition.get(CONSUMER_GROUP_ID_KEY).toString();
+    }
+
+    byte[] recordKey() {
+        return serializeKey().array();
+    }
+
+    byte[] recordValue() {
+        return serializeValue(VERSION).array();
+    }
+};
+
diff --git a/connect/mirror-client/src/main/java/org/apache/kafka/connect/mirror/DefaultReplicationPolicy.java b/connect/mirror-client/src/main/java/org/apache/kafka/connect/mirror/DefaultReplicationPolicy.java
new file mode 100644
index 0000000000000..30d75348a8a4e
--- /dev/null
+++ b/connect/mirror-client/src/main/java/org/apache/kafka/connect/mirror/DefaultReplicationPolicy.java
@@ -0,0 +1,73 @@
+/*
+ * 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.mirror;
+
+import org.apache.kafka.common.Configurable;
+
+import java.util.Map;
+import java.util.regex.Pattern;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** Defines remote topics like "us-west.topic1". The separator is customizable and defaults to a period. */
+public class DefaultReplicationPolicy implements ReplicationPolicy, Configurable {
+    
+    private static final Logger log = LoggerFactory.getLogger(DefaultReplicationPolicy.class);
+
+    // In order to work with various metrics stores, we allow custom separators.
+    public static final String SEPARATOR_CONFIG = MirrorClientConfig.REPLICATION_POLICY_SEPARATOR;
+    public static final String SEPARATOR_DEFAULT = ".";
+
+    private String separator = SEPARATOR_DEFAULT;
+    private Pattern separatorPattern = Pattern.compile(Pattern.quote(SEPARATOR_DEFAULT));
+
+    @Override
+    public void configure(Map props) {
+        if (props.containsKey(SEPARATOR_CONFIG)) {
+            separator = (String) props.get(SEPARATOR_CONFIG);
+            log.info("Using custom remote topic separator: '{}'", separator);
+            separatorPattern = Pattern.compile(Pattern.quote(separator));
+        }
+    }
+
+    @Override
+    public String formatRemoteTopic(String sourceClusterAlias, String topic) {
+        return sourceClusterAlias + separator + topic;
+    }
+
+    @Override
+    public String topicSource(String topic) {
+        String[] parts = separatorPattern.split(topic);
+        if (parts.length < 2) {
+            // this is not a remote topic
+            return null;
+        } else {
+            return parts[0];
+        }
+    }
+
+    @Override
+    public String upstreamTopic(String topic) {
+        String source = topicSource(topic);
+        if (source == null) {
+            return null;
+        } else {
+            return topic.substring(source.length() + separator.length());
+        }
+    }
+}
diff --git a/connect/mirror-client/src/main/java/org/apache/kafka/connect/mirror/Heartbeat.java b/connect/mirror-client/src/main/java/org/apache/kafka/connect/mirror/Heartbeat.java
new file mode 100644
index 0000000000000..a34ce9efb326e
--- /dev/null
+++ b/connect/mirror-client/src/main/java/org/apache/kafka/connect/mirror/Heartbeat.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.connect.mirror;
+
+import org.apache.kafka.common.protocol.types.Field;
+import org.apache.kafka.common.protocol.types.Schema;
+import org.apache.kafka.common.protocol.types.Struct;
+import org.apache.kafka.common.protocol.types.Type;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+
+import java.util.Map;
+import java.util.HashMap;
+import java.nio.ByteBuffer;
+
+/** Heartbeat message sent from MirrorHeartbeatTask to target cluster. Heartbeats are always replicated. */
+public class Heartbeat {
+    public static final String SOURCE_CLUSTER_ALIAS_KEY = "sourceClusterAlias";
+    public static final String TARGET_CLUSTER_ALIAS_KEY = "targetClusterAlias";
+    public static final String TIMESTAMP_KEY = "timestamp";
+    public static final String VERSION_KEY = "version";
+    public static final short VERSION = 0;
+
+    public static final Schema VALUE_SCHEMA_V0 = new Schema(
+            new Field(TIMESTAMP_KEY, Type.INT64));
+
+    public static final Schema KEY_SCHEMA = new Schema(
+            new Field(SOURCE_CLUSTER_ALIAS_KEY, Type.STRING),
+            new Field(TARGET_CLUSTER_ALIAS_KEY, Type.STRING));
+
+    public static final Schema HEADER_SCHEMA = new Schema(
+            new Field(VERSION_KEY, Type.INT16));
+
+    private String sourceClusterAlias;
+    private String targetClusterAlias;
+    private long timestamp;
+
+    public Heartbeat(String sourceClusterAlias, String targetClusterAlias, long timestamp) {
+        this.sourceClusterAlias = sourceClusterAlias;
+        this.targetClusterAlias = targetClusterAlias;
+        this.timestamp = timestamp;
+    }
+
+    public String sourceClusterAlias() {
+        return sourceClusterAlias;
+    }
+
+    public String targetClusterAlias() {
+        return targetClusterAlias;
+    }
+
+    public long timestamp() {
+        return timestamp;
+    }
+
+    @Override
+    public String toString() {
+        return String.format("Heartbeat{sourceClusterAlias=%s, targetClusterAlias=%s, timestamp=%d}",
+            sourceClusterAlias, targetClusterAlias, timestamp);
+    }
+
+    ByteBuffer serializeValue(short version) {
+        Schema valueSchema = valueSchema(version);
+        Struct header = headerStruct(version);
+        Struct value = valueStruct(valueSchema);
+        ByteBuffer buffer = ByteBuffer.allocate(HEADER_SCHEMA.sizeOf(header) + valueSchema.sizeOf(value));
+        HEADER_SCHEMA.write(buffer, header);
+        valueSchema.write(buffer, value);
+        buffer.flip();
+        return buffer;
+    }
+
+    ByteBuffer serializeKey() {
+        Struct struct = keyStruct();
+        ByteBuffer buffer = ByteBuffer.allocate(KEY_SCHEMA.sizeOf(struct));
+        KEY_SCHEMA.write(buffer, struct);
+        buffer.flip();
+        return buffer;
+    }
+
+    public static Heartbeat deserializeRecord(ConsumerRecord record) {
+        ByteBuffer value = ByteBuffer.wrap(record.value());
+        Struct headerStruct = HEADER_SCHEMA.read(value);
+        short version = headerStruct.getShort(VERSION_KEY);
+        Struct valueStruct = valueSchema(version).read(value);
+        long timestamp = valueStruct.getLong(TIMESTAMP_KEY);
+        Struct keyStruct = KEY_SCHEMA.read(ByteBuffer.wrap(record.key()));
+        String sourceClusterAlias = keyStruct.getString(SOURCE_CLUSTER_ALIAS_KEY);
+        String targetClusterAlias = keyStruct.getString(TARGET_CLUSTER_ALIAS_KEY); 
+        return new Heartbeat(sourceClusterAlias, targetClusterAlias, timestamp);    
+    } 
+
+    private Struct headerStruct(short version) {
+        Struct struct = new Struct(HEADER_SCHEMA);
+        struct.set(VERSION_KEY, version);
+        return struct;
+    }
+
+    private Struct valueStruct(Schema schema) {
+        Struct struct = new Struct(schema);
+        struct.set(TIMESTAMP_KEY, timestamp);
+        return struct;
+    }
+
+    private Struct keyStruct() {
+        Struct struct = new Struct(KEY_SCHEMA);
+        struct.set(SOURCE_CLUSTER_ALIAS_KEY, sourceClusterAlias);
+        struct.set(TARGET_CLUSTER_ALIAS_KEY, targetClusterAlias);
+        return struct;
+    }
+
+    Map connectPartition() {
+        Map partition = new HashMap<>();
+        partition.put(SOURCE_CLUSTER_ALIAS_KEY, sourceClusterAlias);
+        partition.put(TARGET_CLUSTER_ALIAS_KEY, targetClusterAlias);
+        return partition;
+    }
+
+    byte[] recordKey() {
+        return serializeKey().array();
+    }
+
+    byte[] recordValue() {
+        return serializeValue(VERSION).array();
+    }
+
+    private static Schema valueSchema(short version) {
+        assert version == 0;
+        return VALUE_SCHEMA_V0;
+    }
+};
+
diff --git a/connect/mirror-client/src/main/java/org/apache/kafka/connect/mirror/MirrorClient.java b/connect/mirror-client/src/main/java/org/apache/kafka/connect/mirror/MirrorClient.java
new file mode 100644
index 0000000000000..17d18ecb58a27
--- /dev/null
+++ b/connect/mirror-client/src/main/java/org/apache/kafka/connect/mirror/MirrorClient.java
@@ -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 org.apache.kafka.connect.mirror;
+
+import org.apache.kafka.clients.admin.AdminClient;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.KafkaConsumer;
+import org.apache.kafka.common.serialization.ByteArrayDeserializer;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.KafkaException;
+import org.apache.kafka.clients.consumer.OffsetAndMetadata;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.protocol.types.SchemaException;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+ 
+import java.time.Duration;
+import java.util.Set;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.HashMap;
+import java.util.Collections;
+import java.util.Collection;
+import java.util.stream.Collectors;
+import java.util.concurrent.ExecutionException;
+
+/** Interprets MM2's internal topics (checkpoints, heartbeats) on a given cluster.
+ *  
       
+ *  Given a top-level "mm2.properties" configuration file, MirrorClients can be constructed
+ *  for individual clusters as follows:
+ *  
       
+ *  
      + *    MirrorMakerConfig mmConfig = new MirrorMakerConfig(props);
+ *    MirrorClientConfig mmClientConfig = mmConfig.clientConfig("some-cluster");
+ *    MirrorClient mmClient = new Mirrorclient(mmClientConfig);
+ *  
      
+ */
+public class MirrorClient implements AutoCloseable {
+    private static final Logger log = LoggerFactory.getLogger(MirrorClient.class);
+
+    private AdminClient adminClient;
+    private ReplicationPolicy replicationPolicy;
+    private Map consumerConfig;
+
+    public MirrorClient(Map props) {
+        this(new MirrorClientConfig(props));
+    }
+
+    public MirrorClient(MirrorClientConfig config) {
+        adminClient = AdminClient.create(config.adminConfig());
+        consumerConfig = config.consumerConfig();
+        replicationPolicy = config.replicationPolicy();
+    }
+
+    // for testing
+    MirrorClient(AdminClient adminClient, ReplicationPolicy replicationPolicy,
+            Map consumerConfig) {
+        this.adminClient = adminClient;
+        this.replicationPolicy = replicationPolicy;
+        this.consumerConfig = consumerConfig;
+    }
+
+    /** Close internal clients. */
+    public void close() {
+        adminClient.close();
+    }
+
+    /** Get the ReplicationPolicy instance used to interpret remote topics. This instance is constructed based on
+     *  relevant configuration properties, including {@code replication.policy.class}. */
+    public ReplicationPolicy replicationPolicy() {
+        return replicationPolicy;
+    }
+
+    /** Compute shortest number of hops from an upstream source cluster.
+     *  For example, given replication flow A->B->C, there are two hops from A to C.
+     *  Returns -1 if upstream cluster is unreachable.
+     */
+    public int replicationHops(String upstreamClusterAlias) throws InterruptedException {
+        return heartbeatTopics().stream()
+            .map(x -> countHopsForTopic(x, upstreamClusterAlias))
+            .filter(x -> x != -1)
+            .mapToInt(x -> x)
+            .min()
+            .orElse(-1);
+    }
+
+    /** Find all heartbeat topics on this cluster. Heartbeat topics are replicated from other clusters. */
+    public Set heartbeatTopics() throws InterruptedException {
+        return listTopics().stream()
+            .filter(this::isHeartbeatTopic)
+            .collect(Collectors.toSet());
+    }
+
+    /** Find all checkpoint topics on this cluster. */
+    public Set checkpointTopics() throws InterruptedException {
+        return listTopics().stream()
+            .filter(this::isCheckpointTopic)
+            .collect(Collectors.toSet());
+    }
+
+    /** Find upstream clusters, which may be multiple hops away, based on incoming heartbeats. */
+    public Set upstreamClusters() throws InterruptedException {
+        return listTopics().stream()
+            .filter(this::isHeartbeatTopic)
+            .flatMap(x -> allSources(x).stream())
+            .distinct()
+            .collect(Collectors.toSet());
+    }
+
+    /** Find all remote topics on this cluster. This does not include internal topics (heartbeats, checkpoints). */
+    public Set remoteTopics() throws InterruptedException {
+        return listTopics().stream()
+            .filter(this::isRemoteTopic)
+            .collect(Collectors.toSet());
+    }
+
+    /** Find all remote topics that have been replicated directly from the given source cluster. */
+    public Set remoteTopics(String source) throws InterruptedException {
+        return listTopics().stream()
+            .filter(this::isRemoteTopic)
+            .filter(x -> source.equals(replicationPolicy.topicSource(x)))
+            .distinct()
+            .collect(Collectors.toSet());
+    }
+
+    /** Translate a remote consumer group's offsets into corresponding local offsets. Topics are automatically
+     *  renamed according to the ReplicationPolicy.
+     *  @param consumerGroupId group ID of remote consumer group
+     *  @param remoteClusterAlias alias of remote cluster
+     *  @param timeout timeout
+     */
+    public Map remoteConsumerOffsets(String consumerGroupId,
+            String remoteClusterAlias, Duration timeout) {
+        long deadline = System.currentTimeMillis() + timeout.toMillis();
+        Map offsets = new HashMap<>();
+        KafkaConsumer consumer = new KafkaConsumer<>(consumerConfig,
+            new ByteArrayDeserializer(), new ByteArrayDeserializer());
+        try {
+            // checkpoint topics are not "remote topics", as they are not replicated. So we don't need
+            // to use ReplicationPolicy to create the checkpoint topic here.
+            String checkpointTopic = remoteClusterAlias + MirrorClientConfig.CHECKPOINTS_TOPIC_SUFFIX;
+            List checkpointAssignment =
+                Collections.singletonList(new TopicPartition(checkpointTopic, 0));
+            consumer.assign(checkpointAssignment);
+            consumer.seekToBeginning(checkpointAssignment);
+            while (System.currentTimeMillis() < deadline && !endOfStream(consumer, checkpointAssignment)) {
+                ConsumerRecords records = consumer.poll(timeout);
+                for (ConsumerRecord record : records) {
+                    try {
+                        Checkpoint checkpoint = Checkpoint.deserializeRecord(record);
+                        if (checkpoint.consumerGroupId().equals(consumerGroupId)) {
+                            offsets.put(checkpoint.topicPartition(), checkpoint.offsetAndMetadata());
+                        }
+                    } catch (SchemaException e) {
+                        log.info("Could not deserialize record. Skipping.", e);
+                    }
+                }
+            }
+            log.info("Consumed {} checkpoint records for {} from {}.", offsets.size(),
+                consumerGroupId, checkpointTopic);
+        } finally {
+            consumer.close();
+        }
+        return offsets;
+    }
+
+    Set listTopics() throws InterruptedException {
+        try {
+            return adminClient.listTopics().names().get();
+        } catch (ExecutionException e) {
+            throw new KafkaException(e.getCause());
+        }
+    }
+
+    int countHopsForTopic(String topic, String sourceClusterAlias) {
+        int hops = 0;
+        while (true) {
+            hops++;
+            String source = replicationPolicy.topicSource(topic);
+            if (source == null) {
+                return -1;
+            }
+            if (source.equals(sourceClusterAlias)) {
+                return hops;
+            }
+            topic = replicationPolicy.upstreamTopic(topic);
+        } 
+    }
+
+    boolean isHeartbeatTopic(String topic) {
+        // heartbeats are replicated, so we must use ReplicationPolicy here
+        return MirrorClientConfig.HEARTBEATS_TOPIC.equals(replicationPolicy.originalTopic(topic));
+    }
+
+    boolean isCheckpointTopic(String topic) {
+        // checkpoints are not replicated, so we don't need to use ReplicationPolicy here
+        return topic.endsWith(MirrorClientConfig.CHECKPOINTS_TOPIC_SUFFIX);
+    }
+
+    boolean isRemoteTopic(String topic) {
+        return !replicationPolicy.isInternalTopic(topic)
+            && replicationPolicy.topicSource(topic) != null;
+    }
+
+    Set allSources(String topic) {
+        Set sources = new HashSet<>();
+        String source = replicationPolicy.topicSource(topic);
+        while (source != null) {
+            sources.add(source);
+            topic = replicationPolicy.upstreamTopic(topic);
+            source = replicationPolicy.topicSource(topic);
+        }
+        return sources;
+    }
+
+    static private boolean endOfStream(Consumer consumer, Collection assignments) {
+        Map endOffsets = consumer.endOffsets(assignments);
+        for (TopicPartition topicPartition : assignments) {
+            if (consumer.position(topicPartition) < endOffsets.get(topicPartition)) {
+                return false;
+            }
+        }
+        return true;
+    }
+}
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
new file mode 100644
index 0000000000000..0c163d87fb6ec
--- /dev/null
+++ b/connect/mirror-client/src/main/java/org/apache/kafka/connect/mirror/MirrorClientConfig.java
@@ -0,0 +1,135 @@
+/*
+ * 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.mirror;
+
+import org.apache.kafka.common.config.AbstractConfig;
+import org.apache.kafka.common.config.ConfigDef;
+import org.apache.kafka.common.config.ConfigDef.Type;
+import org.apache.kafka.common.config.ConfigDef.Importance;
+import org.apache.kafka.clients.CommonClientConfigs;
+
+import java.util.Map;
+import java.util.HashMap;
+
+/** Configuration required for MirrorClient to talk to a given target cluster.
+ *  
      
+ *  Generally, these properties come from an mm2.properties configuration file
+ *  (@see MirrorMakerConfig.clientConfig):
+ *  
      
+ *  
      + *    MirrorMakerConfig mmConfig = new MirrorMakerConfig(props);
+ *    MirrorClientConfig mmClientConfig = mmConfig.clientConfig("some-cluster");
+ *  
      
+ *  
      
+ *  In addition to the properties defined here, sub-configs are supported for Admin, Consumer, and Producer clients.
+ *  For example:
+ *  
      
+ *  
      + *      bootstrap.servers = host1:9092
+ *      consumer.client.id = mm2-client
+ *      replication.policy.separator = __
+ *  
      
+ */
+public class MirrorClientConfig extends AbstractConfig {
+    public static final String REPLICATION_POLICY_CLASS = "replication.policy.class";
+    private static final String REPLICATION_POLICY_CLASS_DOC = "Class which defines the remote topic naming convention.";
+    public static final Class REPLICATION_POLICY_CLASS_DEFAULT = DefaultReplicationPolicy.class;
+    public static final String REPLICATION_POLICY_SEPARATOR = "replication.policy.separator";
+    private static final String REPLICATION_POLICY_SEPARATOR_DOC = "Separator used in remote topic naming convention.";
+    public static final String REPLICATION_POLICY_SEPARATOR_DEFAULT =
+        DefaultReplicationPolicy.SEPARATOR_DEFAULT;
+    
+    public static final String ADMIN_CLIENT_PREFIX = "admin.";
+    public static final String CONSUMER_CLIENT_PREFIX = "consumer.";
+    public static final String PRODUCER_CLIENT_PREFIX = "producer.";
+
+    static final String CHECKPOINTS_TOPIC_SUFFIX = ".checkpoints.internal"; // internal so not replicated
+    static final String HEARTBEATS_TOPIC = "heartbeats";
+ 
+    MirrorClientConfig(Map props) {
+        super(CONFIG_DEF, props, true);
+    }
+
+    public ReplicationPolicy replicationPolicy() {
+        return getConfiguredInstance(REPLICATION_POLICY_CLASS, ReplicationPolicy.class);
+    }
+
+    /** Sub-config for Admin clients. */
+    public Map adminConfig() {
+        return clientConfig(ADMIN_CLIENT_PREFIX);
+    }
+
+    /** Sub-config for Consumer clients. */
+    public Map consumerConfig() {
+        return clientConfig(CONSUMER_CLIENT_PREFIX);
+    }
+
+    /** Sub-config for Producer clients. */
+    public Map producerConfig() {
+        return clientConfig(PRODUCER_CLIENT_PREFIX);
+    }
+    
+    private Map clientConfig(String prefix) {
+        Map props = new HashMap<>();
+        props.putAll(valuesWithPrefixOverride(prefix));
+        props.keySet().retainAll(CLIENT_CONFIG_DEF.names());
+        props.entrySet().removeIf(x -> x.getValue() == null);
+        return props;
+    }
+
+    // Properties passed to internal Kafka clients
+    static final ConfigDef CLIENT_CONFIG_DEF = new ConfigDef()
+        .define(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG,
+            Type.LIST,
+            null,
+            Importance.HIGH,
+            CommonClientConfigs.BOOTSTRAP_SERVERS_DOC)
+        // security support
+        .define(CommonClientConfigs.SECURITY_PROTOCOL_CONFIG,
+            Type.STRING,
+            CommonClientConfigs.DEFAULT_SECURITY_PROTOCOL,
+            Importance.MEDIUM,
+            CommonClientConfigs.SECURITY_PROTOCOL_DOC)
+        .withClientSslSupport()
+        .withClientSaslSupport();
+ 
+    static final ConfigDef CONFIG_DEF = new ConfigDef()
+        .define(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG,
+            Type.STRING,
+            null,
+            Importance.HIGH,
+            CommonClientConfigs.BOOTSTRAP_SERVERS_DOC) 
+        .define(
+            REPLICATION_POLICY_CLASS,
+            ConfigDef.Type.CLASS,
+            REPLICATION_POLICY_CLASS_DEFAULT,
+            ConfigDef.Importance.LOW,
+            REPLICATION_POLICY_CLASS_DOC)
+        .define(
+            REPLICATION_POLICY_SEPARATOR,
+            ConfigDef.Type.STRING,
+            REPLICATION_POLICY_SEPARATOR_DEFAULT,
+            ConfigDef.Importance.LOW,
+            REPLICATION_POLICY_SEPARATOR_DOC)
+        .define(CommonClientConfigs.SECURITY_PROTOCOL_CONFIG,
+                Type.STRING,
+                CommonClientConfigs.DEFAULT_SECURITY_PROTOCOL,
+                Importance.MEDIUM,
+                CommonClientConfigs.SECURITY_PROTOCOL_DOC)
+        .withClientSslSupport()
+        .withClientSaslSupport();
+}
diff --git a/connect/mirror-client/src/main/java/org/apache/kafka/connect/mirror/RemoteClusterUtils.java b/connect/mirror-client/src/main/java/org/apache/kafka/connect/mirror/RemoteClusterUtils.java
new file mode 100644
index 0000000000000..f93431985455f
--- /dev/null
+++ b/connect/mirror-client/src/main/java/org/apache/kafka/connect/mirror/RemoteClusterUtils.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 org.apache.kafka.connect.mirror;
+
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.clients.consumer.OffsetAndMetadata;
+
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.TimeoutException;
+import java.time.Duration;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+ 
+/** Convenience methods for multi-cluster environments. Wraps MirrorClient (@see MirrorClient).
+ *  
      
+ *  Properties passed to these methods are used to construct internal Admin and Consumer clients.
+ *  Sub-configs like "admin.xyz" are also supported. For example:
+ *  
      
+ *  
      + *      bootstrap.servers = host1:9092
+ *      consumer.client.id = mm2-client
+ *  
      
+ *  
      
+ *  @see MirrorClientConfig for additional properties used by the internal MirrorClient.
+ *  
      
+ */
+public final class RemoteClusterUtils {
+    private static final Logger log = LoggerFactory.getLogger(RemoteClusterUtils.class);
+
+    // utility class
+    private RemoteClusterUtils() {}
+
+    /** Find shortest number of hops from an upstream cluster.
+     *  Returns -1 if the cluster is unreachable */ 
+    public static int replicationHops(Map properties, String upstreamClusterAlias)
+            throws InterruptedException, TimeoutException {
+        try (MirrorClient client = new MirrorClient(properties)) {
+            return client.replicationHops(upstreamClusterAlias);
+        }
+    }
+
+    /** Find all heartbeat topics */
+    public static Set heartbeatTopics(Map properties)
+            throws InterruptedException, TimeoutException {
+        try (MirrorClient client = new MirrorClient(properties)) {
+            return client.heartbeatTopics();
+        }
+    }
+
+    /** Find all checkpoint topics */
+    public static Set checkpointTopics(Map properties)
+            throws InterruptedException, TimeoutException {
+        try (MirrorClient client = new MirrorClient(properties)) {
+            return client.checkpointTopics();
+        }
+    }
+
+    /** Find all upstream clusters */
+    public static Set upstreamClusters(Map properties)
+            throws InterruptedException, TimeoutException {
+        try (MirrorClient client = new MirrorClient(properties)) {
+            return client.upstreamClusters();
+        }
+    }
+
+    /** Translate a remote consumer group's offsets into corresponding local offsets. Topics are automatically
+     *  renamed according to the ReplicationPolicy.
+     *  @param properties @see MirrorClientConfig
+     *  @param consumerGroupId group ID of remote consumer group
+     *  @param remoteClusterAlias alias of remote cluster
+     *  @param timeout timeout
+     */
+    public static Map translateOffsets(Map properties,
+            String remoteClusterAlias, String consumerGroupId, Duration timeout)
+            throws InterruptedException, TimeoutException {
+        try (MirrorClient client = new MirrorClient(properties)) {
+            return client.remoteConsumerOffsets(consumerGroupId, remoteClusterAlias, timeout);
+        }
+    }
+}
diff --git a/connect/mirror-client/src/main/java/org/apache/kafka/connect/mirror/ReplicationPolicy.java b/connect/mirror-client/src/main/java/org/apache/kafka/connect/mirror/ReplicationPolicy.java
new file mode 100644
index 0000000000000..11f73f50ceac5
--- /dev/null
+++ b/connect/mirror-client/src/main/java/org/apache/kafka/connect/mirror/ReplicationPolicy.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.connect.mirror;
+
+import org.apache.kafka.common.annotation.InterfaceStability;
+
+/** Defines which topics are "remote topics". e.g. "us-west.topic1". */
+@InterfaceStability.Evolving
+public interface ReplicationPolicy {
+
+    /** How to rename remote topics; generally should be like us-west.topic1. */
+    String formatRemoteTopic(String sourceClusterAlias, String topic);
+
+    /** Source cluster alias of given remote topic, e.g. "us-west" for "us-west.topic1".
+     *  Returns null if not a remote topic.
+     */
+    String topicSource(String topic);
+
+    /** Name of topic on the source cluster, e.g. "topic1" for "us-west.topic1".
+     *
+     *  Topics may be replicated multiple hops, so the immediately upstream topic
+     *  may itself be a remote topic.
+     *
+     *  Returns null if not a remote topic.
+     */
+    String upstreamTopic(String topic); 
+
+    /** The name of the original source-topic, which may have been replicated multiple hops.
+     *  Returns the topic if it is not a remote topic.
+     */
+    default String originalTopic(String topic) {
+        String upstream = upstreamTopic(topic);
+        if (upstream == null) {
+            return topic;
+        } else {
+            return originalTopic(upstream);
+        }
+    }
+
+    /** Internal topics are never replicated. */
+    default boolean isInternalTopic(String topic) {
+        return topic.endsWith(".internal") || topic.endsWith("-internal") || topic.startsWith("__")
+            || topic.startsWith(".");
+    }
+}
diff --git a/connect/mirror-client/src/main/java/org/apache/kafka/connect/mirror/SourceAndTarget.java b/connect/mirror-client/src/main/java/org/apache/kafka/connect/mirror/SourceAndTarget.java
new file mode 100644
index 0000000000000..f853dc40c16ac
--- /dev/null
+++ b/connect/mirror-client/src/main/java/org/apache/kafka/connect/mirror/SourceAndTarget.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.connect.mirror;
+
+/** Directional pair of clustes, where source is replicated to target. */
+public class SourceAndTarget {
+    private String source;
+    private String target;
+
+    public SourceAndTarget(String source, String target) {
+        this.source = source;
+        this.target = target;
+    }
+
+    public String source() {
+        return source;
+    }
+
+    public String target() {
+        return target;
+    }
+
+    @Override
+    public String toString() {
+        return source + "->" + target;
+    }
+
+    @Override
+    public int hashCode() {
+        return toString().hashCode();
+    }
+
+    @Override
+    public boolean equals(Object other) {
+        return other != null && toString().equals(other.toString());
+    }
+}
+
diff --git a/connect/mirror-client/src/test/java/org/apache/kafka/connect/mirror/MirrorClientTest.java b/connect/mirror-client/src/test/java/org/apache/kafka/connect/mirror/MirrorClientTest.java
new file mode 100644
index 0000000000000..c2536d5db8508
--- /dev/null
+++ b/connect/mirror-client/src/test/java/org/apache/kafka/connect/mirror/MirrorClientTest.java
@@ -0,0 +1,163 @@
+/*
+ * 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.mirror;
+
+import org.apache.kafka.common.Configurable;
+
+import java.util.Collections;
+import java.util.List;
+import java.util.Set;
+import java.util.HashSet;
+import java.util.Arrays;
+import java.util.concurrent.TimeoutException;
+
+import org.junit.Test;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.assertFalse;
+
+public class MirrorClientTest {
+
+    private static class FakeMirrorClient extends MirrorClient {
+
+        List topics;
+
+        FakeMirrorClient(List topics) {
+            super(null, new DefaultReplicationPolicy(), null);
+            this.topics = topics;
+        }
+
+        FakeMirrorClient() {
+            this(Collections.emptyList());
+        } 
+
+        @Override
+        protected Set listTopics() {
+            return new HashSet<>(topics);
+        }
+    }
+
+    @Test
+    public void testIsHeartbeatTopic() throws InterruptedException, TimeoutException {
+        MirrorClient client = new FakeMirrorClient();
+        assertTrue(client.isHeartbeatTopic("heartbeats"));
+        assertTrue(client.isHeartbeatTopic("source1.heartbeats"));
+        assertTrue(client.isHeartbeatTopic("source2.source1.heartbeats"));
+        assertFalse(client.isHeartbeatTopic("heartbeats!"));
+        assertFalse(client.isHeartbeatTopic("!heartbeats"));
+        assertFalse(client.isHeartbeatTopic("source1heartbeats"));
+        assertFalse(client.isHeartbeatTopic("source1-heartbeats"));
+    }
+
+    @Test
+    public void testIsCheckpointTopic() throws InterruptedException, TimeoutException {
+        MirrorClient client = new FakeMirrorClient();
+        assertTrue(client.isCheckpointTopic("source1.checkpoints.internal"));
+        assertFalse(client.isCheckpointTopic("checkpoints.internal"));
+        assertFalse(client.isCheckpointTopic("checkpoints-internal"));
+        assertFalse(client.isCheckpointTopic("checkpoints.internal!"));
+        assertFalse(client.isCheckpointTopic("!checkpoints.internal"));
+        assertFalse(client.isCheckpointTopic("source1checkpointsinternal"));
+    }
+
+    @Test
+    public void countHopsForTopicTest() throws InterruptedException, TimeoutException {
+        MirrorClient client = new FakeMirrorClient();
+        assertEquals(-1, client.countHopsForTopic("topic", "source"));
+        assertEquals(-1, client.countHopsForTopic("source", "source"));
+        assertEquals(-1, client.countHopsForTopic("sourcetopic", "source"));
+        assertEquals(-1, client.countHopsForTopic("source1.topic", "source2"));
+        assertEquals(1, client.countHopsForTopic("source1.topic", "source1"));
+        assertEquals(1, client.countHopsForTopic("source2.source1.topic", "source2"));
+        assertEquals(2, client.countHopsForTopic("source2.source1.topic", "source1"));
+        assertEquals(3, client.countHopsForTopic("source3.source2.source1.topic", "source1"));
+        assertEquals(-1, client.countHopsForTopic("source3.source2.source1.topic", "source4"));
+    }
+
+    @Test
+    public void heartbeatTopicsTest() throws InterruptedException, TimeoutException {
+        MirrorClient client = new FakeMirrorClient(Arrays.asList("topic1", "topic2", "heartbeats",
+            "source1.heartbeats", "source2.source1.heartbeats", "source3.heartbeats"));
+        Set heartbeatTopics = client.heartbeatTopics();
+        assertEquals(heartbeatTopics, new HashSet<>(Arrays.asList("heartbeats", "source1.heartbeats",
+            "source2.source1.heartbeats", "source3.heartbeats")));
+    }
+
+    @Test
+    public void checkpointsTopicsTest() throws InterruptedException, TimeoutException {
+        MirrorClient client = new FakeMirrorClient(Arrays.asList("topic1", "topic2", "checkpoints.internal",
+            "source1.checkpoints.internal", "source2.source1.checkpoints.internal", "source3.checkpoints.internal"));
+        Set checkpointTopics = client.checkpointTopics();
+        assertEquals(new HashSet<>(Arrays.asList("source1.checkpoints.internal",
+            "source2.source1.checkpoints.internal", "source3.checkpoints.internal")), checkpointTopics);
+    }
+
+    @Test
+    public void replicationHopsTest() throws InterruptedException, TimeoutException {
+        MirrorClient client = new FakeMirrorClient(Arrays.asList("topic1", "topic2", "heartbeats",
+            "source1.heartbeats", "source1.source2.heartbeats", "source3.heartbeats"));
+        assertEquals(1, client.replicationHops("source1"));
+        assertEquals(2, client.replicationHops("source2")); 
+        assertEquals(1, client.replicationHops("source3"));
+        assertEquals(-1, client.replicationHops("source4"));
+    }
+
+    @Test
+    public void upstreamClustersTest() throws InterruptedException {
+        MirrorClient client = new FakeMirrorClient(Arrays.asList("topic1", "topic2", "heartbeats",
+            "source1.heartbeats", "source1.source2.heartbeats", "source3.source4.source5.heartbeats"));
+        Set sources = client.upstreamClusters();
+        assertTrue(sources.contains("source1"));
+        assertTrue(sources.contains("source2"));
+        assertTrue(sources.contains("source3"));
+        assertTrue(sources.contains("source4"));
+        assertTrue(sources.contains("source5"));
+        assertFalse(sources.contains("sourceX"));
+        assertFalse(sources.contains(""));
+        assertFalse(sources.contains(null));
+    }
+
+    @Test
+    public void remoteTopicsTest() throws InterruptedException {
+        MirrorClient client = new FakeMirrorClient(Arrays.asList("topic1", "topic2", "topic3",
+            "source1.topic4", "source1.source2.topic5", "source3.source4.source5.topic6"));
+        Set remoteTopics = client.remoteTopics();
+        assertFalse(remoteTopics.contains("topic1"));
+        assertFalse(remoteTopics.contains("topic2"));
+        assertFalse(remoteTopics.contains("topic3"));
+        assertTrue(remoteTopics.contains("source1.topic4"));
+        assertTrue(remoteTopics.contains("source1.source2.topic5"));
+        assertTrue(remoteTopics.contains("source3.source4.source5.topic6"));
+    }
+
+    @Test
+    public void remoteTopicsSeparatorTest() throws InterruptedException {
+        MirrorClient client = new FakeMirrorClient(Arrays.asList("topic1", "topic2", "topic3",
+            "source1__topic4", "source1__source2__topic5", "source3__source4__source5__topic6"));
+        ((Configurable) client.replicationPolicy()).configure(
+            Collections.singletonMap("replication.policy.separator", "__"));
+        Set remoteTopics = client.remoteTopics();
+        assertFalse(remoteTopics.contains("topic1"));
+        assertFalse(remoteTopics.contains("topic2"));
+        assertFalse(remoteTopics.contains("topic3"));
+        assertTrue(remoteTopics.contains("source1__topic4"));
+        assertTrue(remoteTopics.contains("source1__source2__topic5"));
+        assertTrue(remoteTopics.contains("source3__source4__source5__topic6"));
+    }
+
+}
diff --git a/connect/mirror/README.md b/connect/mirror/README.md
new file mode 100644
index 0000000000000..68e3536d94c3d
--- /dev/null
+++ b/connect/mirror/README.md
@@ -0,0 +1,222 @@
+
+# MirrorMaker 2.0
+
+MM2 leverages the Connect framework to replicate topics between Kafka
+clusters. MM2 includes several new features, including:
+
+ - both topics and consumer groups are replicated
+ - topic configuration and ACLs are replicated
+ - cross-cluster offsets are synchronized
+ - partitioning is preserved
+
+## Replication flows
+
+MM2 replicates topics and consumer groups from upstream source clusters
+to downstream target clusters. These directional flows are notated
+`A->B`.
+
+It's possible to create complex replication topologies based on these
+`source->target` flows, including:
+
+ - *fan-out*, e.g. `K->A, K->B, K->C`
+ - *aggregation*, e.g. `A->K, B->K, C->K`
+ - *active/active*, e.g. `A->B, B->A`
+
+Each replication flow can be configured independently, e.g. to replicate
+specific topics or groups:
+
+    A->B.topics = topic-1, topic-2
+    A->B.groups = group-1, group-2
+
+By default, all topics and consumer groups are replicated (except
+blacklisted ones), across all enabled replication flows. Each
+replication flow must be explicitly enabled to begin replication:
+
+    A->B.enabled = true
+    B->A.enabled = true
+
+## Starting an MM2 process
+
+You can run any number of MM2 processes as needed. Any MM2 processes
+which are configured to replicate the same Kafka clusters will find each
+other, share configuration, load balance, etc.
+
+To start an MM2 process, first specify Kafka cluster information in a
+configuration file as follows:
+
+    # mm2.properties
+    clusters = us-west, us-east
+    us-west.bootstrap.servers = host1:9092
+    us-east.bootstrap.servers = host2:9092
+
+You can list any number of clusters this way.
+
+Optionally, you can override default MirrorMaker properties:
+
+    topics = .*
+    groups = group1, group2
+    emit.checkpoints.interval.seconds = 10
+
+These will apply to all replication flows. You can also override default
+properties for specific clusters or replication flows:
+
+    # configure a specific cluster
+    us-west.offset.storage.topic = mm2-offsets
+
+    # configure a specific source->target replication flow
+    us-west->us-east.emit.heartbeats = false
+
+Next, enable individual replication flows as follows:
+
+    us-west->us-east.enabled = true     # disabled by default
+
+Finally, launch one or more MirrorMaker processes with the `connect-mirror-maker.sh`
+script:
+
+    $ ./bin/connect-mirror-maker.sh mm2.properties
+
+## Multicluster environments
+
+MM2 supports replication between multiple Kafka clusters, whether in the
+same data center or across multiple data centers. A single MM2 cluster
+can span multiple data centers, but it is recommended to keep MM2's producers
+as close as possible to their target clusters. To do so, specify a subset
+of clusters for each MM2 node as follows:
+
+    # in west DC:
+    $ ./bin/connect-mirror-maker.sh mm2.properties --clusters west-1 west-2
+
+This signals to the node that the given clusters are nearby, and prevents the
+node from sending records or configuration to clusters in other data centers.
+
+### Example
+
+Say there are three data centers (west, east, north) with two Kafka
+clusters in each data center (west-1, west-2 etc). We can configure MM2
+for active/active replication within each data center, as well as cross data
+center replication (XDCR) as follows:
+
+    # mm2.properties
+    clusters: west-1, west-2, east-1, east-2, north-1, north-2
+
+    west-1.bootstrap.servers = ...
+    ---%<---
+
+    # active/active in west
+    west-1->west-2.enabled = true
+    west-2->west-1.enabled = true
+
+    # active/active in east
+    east-1->east-2.enabled = true
+    east-2->east-1.enabled = true
+
+    # active/active in north
+    north-1->north-2.enabled = true
+    north-2->north-1.enabled = true
+
+    # XDCR via west-1, east-1, north-1
+    west-1->east-1.enabled = true
+    west-1->north-1.enabled = true
+    east-1->west-1.enabled = true
+    east-1->north-1.enabled = true
+    north-1->west-1.enabled = true
+    north-1->east-1.enabled = true
+
+Then, launch MM2 in each data center as follows:
+
+    # in west:
+    $ ./bin/connect-mirror-maker.sh mm2.properties --clusters west-1 west-2
+
+    # in east:
+    $ ./bin/connect-mirror-maker.sh mm2.properties --clusters east-1 east-2
+
+    # in north:
+    $ ./bin/connect-mirror-maker.sh mm2.properties --clusters north-1 north-2
+    
+With this configuration, records produced to any cluster will be replicated
+within the data center, as well as across to other data centers. By providing
+the `--clusters` parameter, we ensure that each node only produces records to
+nearby clusters.
+
+N.B. that the `--clusters` parameter is not technically required here. MM2 will work fine without it; however, throughput may suffer from "producer lag" between
+data centers, and you may incur unnecessary data transfer costs.
+
+## Shared configuration
+
+MM2 processes share configuration via their target Kafka clusters. 
+For example, the following two processes would be racy:
+
+    # process1:
+    A->B.enabled = true
+    A->B.topics = foo
+
+    # process2:
+    A->B.enabled = true
+    A->B.topics = bar
+
+In this case, the two processes will share configuration via cluster `B`.
+Depending on which processes is elected "leader", the result will be
+that either `foo` or `bar` is replicated -- but not both. For this reason,
+it is important to keep configuration consistent across flows to the same
+target cluster. In most cases, your entire organization should use a single
+MM2 configuration file.
+
+## Remote topics
+
+MM2 employs a naming convention to ensure that records from different
+clusters are not written to the same partition. By default, replicated
+topics are renamed based on "source cluster aliases":
+
+    topic-1 --> source.topic-1
+
+This can be customized by overriding the `replication.policy.separator`
+property (default is a period). If you need more control over how
+remote topics are defined, you can implement a custom `ReplicationPolicy`
+and override `replication.policy.class` (default is
+`DefaultReplicationPolicy`).
+
+## Monitoring an MM2 process
+
+MM2 is built on the Connect framework and inherits all of Connect's metrics, e.g.
+`source-record-poll-rate`. In addition, MM2 produces its own metrics under the
+`kafka.connect.mirror` metric group. Metrics are tagged with the following properties:
+
+    - *target*: alias of target cluster
+    - *source*: alias of source cluster
+    - *topic*:  remote topic on target cluster 
+    - *partition*: partition being replicated
+
+Metrics are tracked for each *remote* topic. The source cluster can be inferred
+from the topic name. For example, replicating `topic1` from `A->B` will yield metrics
+like:
+
+    - `target=B`
+    - `topic=A.topic1`
+    - `partition=1`
+
+The following metrics are emitted:
+
+    # MBean: kafka.connect.mirror:type=MirrorSourceConnector,target=([-.w]+),topic=([-.w]+),partition=([0-9]+)
+
+    record-count            # number of records replicated source -> target
+    record-age-ms           # age of records when they are replicated
+    record-age-ms-min
+    record-age-ms-max
+    record-age-ms-avg
+    replication-latecny-ms  # time it takes records to propagate source->target
+    replication-latency-ms-min
+    replication-latency-ms-max
+    replication-latency-ms-avg
+    byte-rate               # average number of bytes/sec in replicated records
+
+
+    # MBean: kafka.connect.mirror:type=MirrorCheckpointConnector,source=([-.w]+),target=([-.w]+)
+
+    checkpoint-latency-ms   # time it takes to replicate consumer offsets
+    checkpoint-latency-ms-min
+    checkpoint-latency-ms-max
+    checkpoint-latency-ms-avg
+
+These metrics do not discern between created-at and log-append timestamps.
+
+
diff --git a/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/ConfigPropertyFilter.java b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/ConfigPropertyFilter.java
new file mode 100644
index 0000000000000..ec6b3b910710b
--- /dev/null
+++ b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/ConfigPropertyFilter.java
@@ -0,0 +1,37 @@
+/*
+ * 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.mirror;
+
+import org.apache.kafka.common.Configurable;
+import org.apache.kafka.common.annotation.InterfaceStability;
+import java.util.Map;
+
+/** Defines which topic configuration properties should be replicated. */
+@InterfaceStability.Evolving
+public interface ConfigPropertyFilter extends Configurable, AutoCloseable {
+
+    boolean shouldReplicateConfigProperty(String prop);
+
+    default void close() {
+        //nop
+    }
+
+    default void configure(Map props) {
+        //nop
+    }
+}
diff --git a/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/DefaultConfigPropertyFilter.java b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/DefaultConfigPropertyFilter.java
new file mode 100644
index 0000000000000..f51db1cbd7dbf
--- /dev/null
+++ b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/DefaultConfigPropertyFilter.java
@@ -0,0 +1,77 @@
+/*
+ * 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.mirror;
+
+import org.apache.kafka.common.config.AbstractConfig;
+import org.apache.kafka.common.config.ConfigDef;
+import org.apache.kafka.common.config.ConfigDef.Type;
+import org.apache.kafka.common.config.ConfigDef.Importance;
+
+import java.util.Map;
+import java.util.regex.Pattern;
+
+/** Uses a blacklist of property names or regexes. */
+public class DefaultConfigPropertyFilter implements ConfigPropertyFilter {
+    
+    public static final String CONFIG_PROPERTIES_BLACKLIST_CONFIG = "config.properties.blacklist";
+    private static final String CONFIG_PROPERTIES_BLACKLIST_DOC = "List of topic configuration properties and/or regexes "
+        + "that should not be replicated.";
+    public static final String CONFIG_PROPERTIES_BLACKLIST_DEFAULT = "follower\\.replication\\.throttled\\.replicas, "
+        + "leader\\.replication\\.throttled\\.replicas, "
+        + "message\\.timestamp\\.difference\\.max\\.ms, "
+        + "message\\.timestamp\\.type, "
+        + "unclean\\.leader\\.election\\.enable, "
+        + "min\\.insync\\.replicas";
+    private Pattern blacklistPattern = MirrorUtils.compilePatternList(CONFIG_PROPERTIES_BLACKLIST_DEFAULT);
+
+    @Override
+    public void configure(Map props) {
+        ConfigPropertyFilterConfig config = new ConfigPropertyFilterConfig(props);
+        blacklistPattern = config.blacklistPattern();
+    }
+
+    @Override
+    public void close() {
+    }
+
+    private boolean blacklisted(String prop) {
+        return blacklistPattern != null && blacklistPattern.matcher(prop).matches();
+    }
+
+    @Override
+    public boolean shouldReplicateConfigProperty(String prop) {
+        return !blacklisted(prop);
+    }
+
+    static class ConfigPropertyFilterConfig extends AbstractConfig {
+
+        static final ConfigDef DEF = new ConfigDef()
+            .define(CONFIG_PROPERTIES_BLACKLIST_CONFIG,
+                Type.LIST,
+                CONFIG_PROPERTIES_BLACKLIST_DEFAULT,
+                Importance.HIGH,
+                CONFIG_PROPERTIES_BLACKLIST_DOC);
+
+        ConfigPropertyFilterConfig(Map props) {
+            super(DEF, props, false);
+        }
+
+        Pattern blacklistPattern() {
+            return MirrorUtils.compilePatternList(getList(CONFIG_PROPERTIES_BLACKLIST_CONFIG));
+        }
+    }
+}
diff --git a/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/DefaultGroupFilter.java b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/DefaultGroupFilter.java
new file mode 100644
index 0000000000000..acf5115236dd7
--- /dev/null
+++ b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/DefaultGroupFilter.java
@@ -0,0 +1,91 @@
+/*
+ * 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.mirror;
+
+import org.apache.kafka.common.config.AbstractConfig;
+import org.apache.kafka.common.config.ConfigDef;
+import org.apache.kafka.common.config.ConfigDef.Type;
+import org.apache.kafka.common.config.ConfigDef.Importance;
+
+import java.util.Map;
+import java.util.regex.Pattern;
+
+/** Uses a whitelist and blacklist. */
+public class DefaultGroupFilter implements GroupFilter {
+
+    public static final String GROUPS_WHITELIST_CONFIG = "groups";
+    private static final String GROUPS_WHITELIST_DOC = "List of consumer group names and/or regexes to replicate.";
+    public static final String GROUPS_WHITELIST_DEFAULT = ".*";
+
+    public static final String GROUPS_BLACKLIST_CONFIG = "groups.blacklist";
+    private static final String GROUPS_BLACKLIST_DOC = "List of consumer group names and/or regexes that should not be replicated.";
+    public static final String GROUPS_BLACKLIST_DEFAULT = "console-consumer-.*, connect-.*, __.*";
+
+    private Pattern whitelistPattern;
+    private Pattern blacklistPattern;
+
+    @Override
+    public void configure(Map props) {
+        GroupFilterConfig config = new GroupFilterConfig(props);
+        whitelistPattern = config.whitelistPattern();
+        blacklistPattern = config.blacklistPattern();
+    }
+
+    @Override
+    public void close() {
+    }
+
+    private boolean whitelisted(String group) {
+        return whitelistPattern != null && whitelistPattern.matcher(group).matches();
+    }
+
+    private boolean blacklisted(String group) {
+        return blacklistPattern != null && blacklistPattern.matcher(group).matches();
+    }
+
+    @Override
+    public boolean shouldReplicateGroup(String group) {
+        return whitelisted(group) && !blacklisted(group);
+    }
+
+    static class GroupFilterConfig extends AbstractConfig {
+
+        static final ConfigDef DEF = new ConfigDef()
+            .define(GROUPS_WHITELIST_CONFIG,
+                Type.LIST,
+                GROUPS_WHITELIST_DEFAULT,
+                Importance.HIGH,
+                GROUPS_WHITELIST_DOC) 
+            .define(GROUPS_BLACKLIST_CONFIG,
+                Type.LIST,
+                GROUPS_BLACKLIST_DEFAULT,
+                Importance.HIGH,
+                GROUPS_BLACKLIST_DOC);
+
+        GroupFilterConfig(Map props) {
+            super(DEF, props, false);
+        }
+
+        Pattern whitelistPattern() {
+            return MirrorUtils.compilePatternList(getList(GROUPS_WHITELIST_CONFIG));
+        }
+
+        Pattern blacklistPattern() {
+            return MirrorUtils.compilePatternList(getList(GROUPS_BLACKLIST_CONFIG));
+        }
+    }
+}
diff --git a/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/DefaultTopicFilter.java b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/DefaultTopicFilter.java
new file mode 100644
index 0000000000000..308bdbfd8a8c4
--- /dev/null
+++ b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/DefaultTopicFilter.java
@@ -0,0 +1,91 @@
+/*
+ * 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.mirror;
+
+import org.apache.kafka.common.config.AbstractConfig;
+import org.apache.kafka.common.config.ConfigDef;
+import org.apache.kafka.common.config.ConfigDef.Type;
+import org.apache.kafka.common.config.ConfigDef.Importance;
+
+import java.util.Map;
+import java.util.regex.Pattern;
+
+/** Uses a whitelist and blacklist. */
+public class DefaultTopicFilter implements TopicFilter {
+    
+    public static final String TOPICS_WHITELIST_CONFIG = "topics";
+    private static final String TOPICS_WHITELIST_DOC = "List of topics and/or regexes to replicate.";
+    public static final String TOPICS_WHITELIST_DEFAULT = ".*";
+
+    public static final String TOPICS_BLACKLIST_CONFIG = "topics.blacklist";
+    private static final String TOPICS_BLACKLIST_DOC = "List of topics and/or regexes that should not be replicated.";
+    public static final String TOPICS_BLACKLIST_DEFAULT = ".*[\\-\\.]internal, .*\\.replica, __.*";
+
+    private Pattern whitelistPattern;
+    private Pattern blacklistPattern;
+
+    @Override
+    public void configure(Map props) {
+        TopicFilterConfig config = new TopicFilterConfig(props);
+        whitelistPattern = config.whitelistPattern();
+        blacklistPattern = config.blacklistPattern();
+    }
+
+    @Override
+    public void close() {
+    }
+
+    private boolean whitelisted(String topic) {
+        return whitelistPattern != null && whitelistPattern.matcher(topic).matches();
+    }
+
+    private boolean blacklisted(String topic) {
+        return blacklistPattern != null && blacklistPattern.matcher(topic).matches();
+    }
+
+    @Override
+    public boolean shouldReplicateTopic(String topic) {
+        return whitelisted(topic) && !blacklisted(topic);
+    }
+
+    static class TopicFilterConfig extends AbstractConfig {
+
+        static final ConfigDef DEF = new ConfigDef()
+            .define(TOPICS_WHITELIST_CONFIG,
+                Type.LIST,
+                TOPICS_WHITELIST_DEFAULT,
+                Importance.HIGH,
+                TOPICS_WHITELIST_DOC) 
+            .define(TOPICS_BLACKLIST_CONFIG,
+                Type.LIST,
+                TOPICS_BLACKLIST_DEFAULT,
+                Importance.HIGH,
+                TOPICS_BLACKLIST_DOC);
+
+        TopicFilterConfig(Map props) {
+            super(DEF, props, false);
+        }
+
+        Pattern whitelistPattern() {
+            return MirrorUtils.compilePatternList(getList(TOPICS_WHITELIST_CONFIG));
+        }
+
+        Pattern blacklistPattern() {
+            return MirrorUtils.compilePatternList(getList(TOPICS_BLACKLIST_CONFIG));
+        }
+    }
+}
diff --git a/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/GroupFilter.java b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/GroupFilter.java
new file mode 100644
index 0000000000000..0202dd5d2b358
--- /dev/null
+++ b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/GroupFilter.java
@@ -0,0 +1,37 @@
+/*
+ * 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.mirror;
+
+import org.apache.kafka.common.Configurable;
+import org.apache.kafka.common.annotation.InterfaceStability;
+import java.util.Map;
+
+/** Defines which consumer groups should be replicated. */
+@InterfaceStability.Evolving
+public interface GroupFilter extends Configurable, AutoCloseable {
+
+    boolean shouldReplicateGroup(String group);
+
+    default void close() {
+        //nop
+    }
+
+    default void configure(Map props) {
+        //nop
+    }
+}
diff --git a/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorCheckpointConnector.java b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorCheckpointConnector.java
new file mode 100644
index 0000000000000..a358584dff009
--- /dev/null
+++ b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorCheckpointConnector.java
@@ -0,0 +1,156 @@
+/*
+ * 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.mirror;
+
+import org.apache.kafka.clients.admin.AdminClient;
+import org.apache.kafka.clients.admin.ConsumerGroupListing;
+import org.apache.kafka.common.config.ConfigDef;
+import org.apache.kafka.common.utils.Utils;
+import org.apache.kafka.connect.connector.Task;
+import org.apache.kafka.connect.source.SourceConnector;
+import org.apache.kafka.connect.util.ConnectorUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.ExecutionException;
+import java.util.stream.Collectors;
+
+/** Replicate consumer group state between clusters. Emits checkpoint records.
+ *
+ *  @see MirrorConnectorConfig for supported config properties.
+ */
+public class MirrorCheckpointConnector extends SourceConnector {
+
+    private static final Logger log = LoggerFactory.getLogger(MirrorCheckpointConnector.class);
+
+    private Scheduler scheduler;
+    private MirrorConnectorConfig config;
+    private GroupFilter groupFilter;
+    private AdminClient sourceAdminClient;
+    private SourceAndTarget sourceAndTarget;
+    private String connectorName;
+    private List knownConsumerGroups = Collections.emptyList();
+
+    @Override
+    public void start(Map props) {
+        config = new MirrorConnectorConfig(props);
+        if (!config.enabled()) {
+            return;
+        }
+        connectorName = config.connectorName();
+        sourceAndTarget = new SourceAndTarget(config.sourceClusterAlias(), config.targetClusterAlias());
+        groupFilter = config.groupFilter();
+        sourceAdminClient = AdminClient.create(config.sourceAdminConfig());
+        scheduler = new Scheduler(MirrorCheckpointConnector.class, config.adminTimeout());
+        scheduler.execute(this::createInternalTopics, "creating internal topics");
+        scheduler.execute(this::loadInitialConsumerGroups, "loading initial consumer groups");
+        scheduler.scheduleRepeatingDelayed(this::refreshConsumerGroups, config.refreshGroupsInterval(),
+                "refreshing consumer groups");
+        log.info("Started {} with {} consumer groups.", connectorName, knownConsumerGroups.size());
+        log.debug("Started {} with consumer groups: {}", connectorName, knownConsumerGroups);
+    }
+
+    @Override
+    public void stop() {
+        if (!config.enabled()) {
+            return;
+        }
+        Utils.closeQuietly(scheduler, "scheduler");
+        Utils.closeQuietly(groupFilter, "group filter");
+        Utils.closeQuietly(sourceAdminClient, "source admin client");
+    }
+
+    @Override
+    public Class taskClass() {
+        return MirrorCheckpointTask.class;
+    }
+
+    // divide consumer groups among tasks
+    @Override
+    public List> taskConfigs(int maxTasks) {
+        if (!config.enabled() || knownConsumerGroups.isEmpty()) {
+            return Collections.emptyList();
+        }
+        int numTasks = Math.min(maxTasks, knownConsumerGroups.size());
+        return ConnectorUtils.groupPartitions(knownConsumerGroups, numTasks).stream()
+                .map(config::taskConfigForConsumerGroups)
+                .collect(Collectors.toList());
+    }
+
+    @Override
+    public ConfigDef config() {
+        return MirrorConnectorConfig.CONNECTOR_CONFIG_DEF;
+    }
+
+    @Override
+    public String version() {
+        return "1";
+    }
+
+    private void refreshConsumerGroups()
+            throws InterruptedException, ExecutionException {
+        List consumerGroups = findConsumerGroups();
+        Set newConsumerGroups = new HashSet<>();
+        newConsumerGroups.addAll(consumerGroups);
+        newConsumerGroups.removeAll(knownConsumerGroups);
+        Set deadConsumerGroups = new HashSet<>();
+        deadConsumerGroups.addAll(knownConsumerGroups);
+        deadConsumerGroups.removeAll(consumerGroups);
+        if (!newConsumerGroups.isEmpty() || !deadConsumerGroups.isEmpty()) {
+            log.info("Found {} consumer groups for {}. {} are new. {} were removed. Previously had {}.",
+                    consumerGroups.size(), sourceAndTarget, newConsumerGroups.size(), deadConsumerGroups.size(),
+                    knownConsumerGroups.size());
+            log.debug("Found new consumer groups: {}", newConsumerGroups);
+            knownConsumerGroups = consumerGroups;
+            context.requestTaskReconfiguration();
+        }
+    }
+
+    private void loadInitialConsumerGroups()
+            throws InterruptedException, ExecutionException {
+        knownConsumerGroups = findConsumerGroups();
+    }
+
+    private List findConsumerGroups()
+            throws InterruptedException, ExecutionException {
+        return listConsumerGroups().stream()
+                .filter(x -> !x.isSimpleConsumerGroup())
+                .map(x -> x.groupId())
+                .filter(this::shouldReplicate)
+                .collect(Collectors.toList());
+    }
+
+    private Collection listConsumerGroups()
+            throws InterruptedException, ExecutionException {
+        return sourceAdminClient.listConsumerGroups().valid().get();
+    }
+
+    private void createInternalTopics() {
+        MirrorUtils.createSinglePartitionCompactedTopic(config.checkpointsTopic(),
+            config.checkpointsTopicReplicationFactor(), config.targetAdminConfig());
+    } 
+
+    boolean shouldReplicate(String group) {
+        return groupFilter.shouldReplicateGroup(group);
+    }
+}
diff --git a/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorCheckpointTask.java b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorCheckpointTask.java
new file mode 100644
index 0000000000000..47a05693c3313
--- /dev/null
+++ b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorCheckpointTask.java
@@ -0,0 +1,193 @@
+/*
+ * 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.mirror;
+
+import org.apache.kafka.connect.source.SourceTask;
+import org.apache.kafka.connect.source.SourceRecord;
+import org.apache.kafka.connect.data.Schema;
+import org.apache.kafka.clients.admin.AdminClient;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.utils.Utils;
+import org.apache.kafka.clients.consumer.OffsetAndMetadata;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.Map;
+import java.util.List;
+import java.util.ArrayList;
+import java.util.Set;
+import java.util.Collections;
+import java.util.stream.Collectors;
+import java.util.concurrent.ExecutionException;
+import java.time.Duration;
+
+/** Emits checkpoints for upstream consumer groups. */
+public class MirrorCheckpointTask extends SourceTask {
+
+    private static final Logger log = LoggerFactory.getLogger(MirrorCheckpointTask.class);
+
+    private AdminClient sourceAdminClient;
+    private String sourceClusterAlias;
+    private String targetClusterAlias;
+    private String checkpointsTopic;
+    private Duration interval;
+    private Duration pollTimeout;
+    private Duration adminTimeout;
+    private TopicFilter topicFilter;
+    private Set consumerGroups;
+    private ReplicationPolicy replicationPolicy;
+    private OffsetSyncStore offsetSyncStore;
+    private boolean stopping;
+    private MirrorMetrics metrics;
+
+    public MirrorCheckpointTask() {}
+
+    // for testing
+    MirrorCheckpointTask(String sourceClusterAlias, String targetClusterAlias,
+            ReplicationPolicy replicationPolicy, OffsetSyncStore offsetSyncStore) {
+        this.sourceClusterAlias = sourceClusterAlias;
+        this.targetClusterAlias = targetClusterAlias;
+        this.replicationPolicy = replicationPolicy;
+        this.offsetSyncStore = offsetSyncStore;
+    }
+
+    @Override
+    public void start(Map props) {
+        MirrorTaskConfig config = new MirrorTaskConfig(props);
+        stopping = false;
+        sourceClusterAlias = config.sourceClusterAlias();
+        targetClusterAlias = config.targetClusterAlias();
+        consumerGroups = config.taskConsumerGroups();
+        checkpointsTopic = config.checkpointsTopic();
+        topicFilter = config.topicFilter();
+        replicationPolicy = config.replicationPolicy();
+        interval = config.emitCheckpointsInterval();
+        pollTimeout = config.consumerPollTimeout();
+        adminTimeout = config.adminTimeout();
+        offsetSyncStore = new OffsetSyncStore(config);
+        sourceAdminClient = AdminClient.create(config.sourceAdminConfig());
+        metrics = config.metrics();
+    }
+
+    @Override
+    public void commit() throws InterruptedException {
+        // nop
+    }
+
+    @Override
+    public void stop() {
+        long start = System.currentTimeMillis();
+        stopping = true;
+        Utils.closeQuietly(offsetSyncStore, "offset sync store");
+        Utils.closeQuietly(sourceAdminClient, "source admin client");
+        Utils.closeQuietly(metrics, "metrics");
+        log.info("Stopping {} took {} ms.", Thread.currentThread().getName(), System.currentTimeMillis() - start);
+    }
+
+    @Override
+    public String version() {
+        return "1";
+    }
+
+    @Override
+    public List poll() throws InterruptedException {
+        try { 
+            long deadline = System.currentTimeMillis() + interval.toMillis();
+            while (!stopping && System.currentTimeMillis() < deadline) {
+                offsetSyncStore.update(pollTimeout);
+            }
+            List records = new ArrayList<>();
+            for (String group : consumerGroups) {
+                records.addAll(checkpointsForGroup(group));
+            }
+            if (records.isEmpty()) {
+                // WorkerSourceTask expects non-zero batches or null
+                return null;
+            } else {
+                return records;
+            }
+        } catch (Throwable e) {
+            log.warn("Failure polling consumer state for checkpoints.", e);
+            return null;
+        }
+    }
+
+    private List checkpointsForGroup(String group) throws InterruptedException {
+        try {
+            long timestamp = System.currentTimeMillis();
+            return listConsumerGroupOffsets(group).entrySet().stream()
+                .filter(x -> shouldCheckpointTopic(x.getKey().topic()))
+                .map(x -> checkpoint(group, x.getKey(), x.getValue()))
+                .filter(x -> x.downstreamOffset() > 0)  // ignore offsets we cannot translate accurately
+                .map(x -> checkpointRecord(x, timestamp))
+                .collect(Collectors.toList());
+        } catch (ExecutionException e) {
+            log.error("Error querying offsets for consumer group {} on cluster {}.",  group, sourceClusterAlias, e);
+            return Collections.emptyList();
+        }
+    }
+
+    private Map listConsumerGroupOffsets(String group)
+            throws InterruptedException, ExecutionException {
+        if (stopping) {
+            // short circuit if stopping
+            return Collections.emptyMap();
+        }
+        return sourceAdminClient.listConsumerGroupOffsets(group).partitionsToOffsetAndMetadata().get();
+    }
+
+    Checkpoint checkpoint(String group, TopicPartition topicPartition,
+            OffsetAndMetadata offsetAndMetadata) {
+        long upstreamOffset = offsetAndMetadata.offset();
+        long downstreamOffset = offsetSyncStore.translateDownstream(topicPartition, upstreamOffset);
+        return new Checkpoint(group, renameTopicPartition(topicPartition),
+            upstreamOffset, downstreamOffset, offsetAndMetadata.metadata());
+    }
+
+    SourceRecord checkpointRecord(Checkpoint checkpoint, long timestamp) {
+        return new SourceRecord(
+            checkpoint.connectPartition(), MirrorUtils.wrapOffset(0),
+            checkpointsTopic, 0,
+            Schema.BYTES_SCHEMA, checkpoint.recordKey(),
+            Schema.BYTES_SCHEMA, checkpoint.recordValue(),
+            timestamp);
+    }
+
+    TopicPartition renameTopicPartition(TopicPartition upstreamTopicPartition) {
+        if (targetClusterAlias.equals(replicationPolicy.topicSource(upstreamTopicPartition.topic()))) {
+            // this topic came from the target cluster, so we rename like us-west.topic1 -> topic1
+            return new TopicPartition(replicationPolicy.originalTopic(upstreamTopicPartition.topic()),
+                upstreamTopicPartition.partition());
+        } else {
+            // rename like topic1 -> us-west.topic1
+            return new TopicPartition(replicationPolicy.formatRemoteTopic(sourceClusterAlias,
+                upstreamTopicPartition.topic()), upstreamTopicPartition.partition());
+        }
+    }
+
+    boolean shouldCheckpointTopic(String topic) {
+        return topicFilter.shouldReplicateTopic(topic);
+    }
+
+    @Override
+    public void commitRecord(SourceRecord record) {
+        metrics.checkpointLatency(MirrorUtils.unwrapPartition(record.sourcePartition()),
+            Checkpoint.unwrapGroup(record.sourcePartition()),
+            System.currentTimeMillis() - record.timestamp());
+    }
+}
diff --git a/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorConnectorConfig.java b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorConnectorConfig.java
new file mode 100644
index 0000000000000..d922eade5f372
--- /dev/null
+++ b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorConnectorConfig.java
@@ -0,0 +1,601 @@
+/*
+ * 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.mirror;
+
+import org.apache.kafka.common.config.AbstractConfig;
+import org.apache.kafka.common.config.ConfigDef;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.metrics.MetricsReporter;
+import org.apache.kafka.common.metrics.JmxReporter;
+import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.connect.runtime.ConnectorConfig;
+
+import java.util.Map;
+import java.util.HashMap;
+import java.util.List;
+import java.util.stream.Collectors;
+import java.time.Duration;
+
+/** Shared config properties used by MirrorSourceConnector, MirrorCheckpointConnector, and MirrorHeartbeatConnector.
+ *  
      
+ *  Generally, these properties are filled-in automatically by MirrorMaker based on a top-level mm2.properties file.
+ *  However, when running MM2 connectors as plugins on a Connect-as-a-Service cluster, these properties must be configured manually,
+ *  e.g. via the Connect REST API.
+ *  
      
+ *  
      
+ *  An example configuration when running on Connect (not via MirrorMaker driver):
+ *  
      
+ *  
      + *      {
+ *        "name": "MirrorSourceConnector",
+ *        "connector.class": "org.apache.kafka.connect.mirror.MirrorSourceConnector",
+ *        "replication.factor": "1",
+ *        "source.cluster.alias": "backup",
+ *        "target.cluster.alias": "primary",
+ *        "source.cluster.bootstrap.servers": "vip1:9092",
+ *        "target.cluster.bootstrap.servers": "vip2:9092",
+ *        "topics": ".*test-topic-.*",
+ *        "groups": "consumer-group-.*",
+ *        "emit.checkpoints.interval.seconds": "1",
+ *        "emit.heartbeats.interval.seconds": "1",
+ *        "sync.topic.acls.enabled": "false"
+ *      }
+ *  
      
+ */
+public class MirrorConnectorConfig extends AbstractConfig {
+
+    protected static final String ENABLED_SUFFIX = ".enabled";
+    protected static final String INTERVAL_SECONDS_SUFFIX = ".interval.seconds";
+
+    protected static final String REFRESH_TOPICS = "refresh.topics";
+    protected static final String REFRESH_GROUPS = "refresh.groups";
+    protected static final String SYNC_TOPIC_CONFIGS = "sync.topic.configs";
+    protected static final String SYNC_TOPIC_ACLS = "sync.topic.acls";
+    protected static final String EMIT_HEARTBEATS = "emit.heartbeats";
+    protected static final String EMIT_CHECKPOINTS = "emit.checkpoints";
+
+    public static final String ENABLED = "enabled";
+    private static final String ENABLED_DOC = "Whether to replicate source->target.";
+    public static final String SOURCE_CLUSTER_ALIAS = "source.cluster.alias";
+    private static final String SOURCE_CLUSTER_ALIAS_DOC = "Alias of source cluster";
+    public static final String TARGET_CLUSTER_ALIAS = "target.cluster.alias";
+    public static final String TARGET_CLUSTER_ALIAS_DEFAULT = "target";
+    private static final String TARGET_CLUSTER_ALIAS_DOC = "Alias of target cluster. Used in metrics reporting.";
+    public static final String REPLICATION_POLICY_CLASS = MirrorClientConfig.REPLICATION_POLICY_CLASS;
+    public static final Class REPLICATION_POLICY_CLASS_DEFAULT = MirrorClientConfig.REPLICATION_POLICY_CLASS_DEFAULT;
+    private static final String REPLICATION_POLICY_CLASS_DOC = "Class which defines the remote topic naming convention.";
+    public static final String REPLICATION_POLICY_SEPARATOR = MirrorClientConfig.REPLICATION_POLICY_SEPARATOR;
+    private static final String REPLICATION_POLICY_SEPARATOR_DOC = "Separator used in remote topic naming convention.";
+    public static final String REPLICATION_POLICY_SEPARATOR_DEFAULT =
+            MirrorClientConfig.REPLICATION_POLICY_SEPARATOR_DEFAULT;
+    public static final String REPLICATION_FACTOR = "replication.factor";
+    private static final String REPLICATION_FACTOR_DOC = "Replication factor for newly created remote topics.";
+    public static final int REPLICATION_FACTOR_DEFAULT = 2;
+    public static final String TOPICS = DefaultTopicFilter.TOPICS_WHITELIST_CONFIG;
+    public static final String TOPICS_DEFAULT = DefaultTopicFilter.TOPICS_WHITELIST_DEFAULT;
+    private static final String TOPICS_DOC = "Topics to replicate. Supports comma-separated topic names and regexes.";
+    public static final String TOPICS_BLACKLIST = DefaultTopicFilter.TOPICS_BLACKLIST_CONFIG;
+    public static final String TOPICS_BLACKLIST_DEFAULT = DefaultTopicFilter.TOPICS_BLACKLIST_DEFAULT;
+    private static final String TOPICS_BLACKLIST_DOC = "Blacklisted topics. Supports comma-separated topic names and regexes."
+            + " Blacklists take precedence over whitelists.";
+    public static final String GROUPS = DefaultGroupFilter.GROUPS_WHITELIST_CONFIG;
+    public static final String GROUPS_DEFAULT = DefaultGroupFilter.GROUPS_WHITELIST_DEFAULT;
+    private static final String GROUPS_DOC = "Consumer groups to replicate. Supports comma-separated group IDs and regexes.";
+    public static final String GROUPS_BLACKLIST = DefaultGroupFilter.GROUPS_BLACKLIST_CONFIG;
+    public static final String GROUPS_BLACKLIST_DEFAULT = DefaultGroupFilter.GROUPS_BLACKLIST_DEFAULT;
+    private static final String GROUPS_BLACKLIST_DOC = "Blacklisted groups. Supports comma-separated group IDs and regexes."
+            + " Blacklists take precedence over whitelists.";
+    public static final String CONFIG_PROPERTIES_BLACKLIST = DefaultConfigPropertyFilter.CONFIG_PROPERTIES_BLACKLIST_CONFIG;
+    public static final String CONFIG_PROPERTIES_BLACKLIST_DEFAULT = DefaultConfigPropertyFilter.CONFIG_PROPERTIES_BLACKLIST_DEFAULT;
+    private static final String CONFIG_PROPERTIES_BLACKLIST_DOC = "Topic config properties that should not be replicated. Supports "
+            + "comma-separated property names and regexes.";
+
+    public static final String HEARTBEATS_TOPIC_REPLICATION_FACTOR = "heartbeats.topic.replication.factor";
+    public static final String HEARTBEATS_TOPIC_REPLICATION_FACTOR_DOC = "Replication factor for heartbeats topic.";
+    public static final short HEARTBEATS_TOPIC_REPLICATION_FACTOR_DEFAULT = 3;
+
+    public static final String CHECKPOINTS_TOPIC_REPLICATION_FACTOR = "checkpoints.topic.replication.factor";
+    public static final String CHECKPOINTS_TOPIC_REPLICATION_FACTOR_DOC = "Replication factor for checkpoints topic.";
+    public static final short CHECKPOINTS_TOPIC_REPLICATION_FACTOR_DEFAULT = 3;
+
+    public static final String OFFSET_SYNCS_TOPIC_REPLICATION_FACTOR = "offset-syncs.topic.replication.factor";
+    public static final String OFFSET_SYNCS_TOPIC_REPLICATION_FACTOR_DOC = "Replication factor for offset-syncs topic.";
+    public static final short OFFSET_SYNCS_TOPIC_REPLICATION_FACTOR_DEFAULT = 3;
+
+    protected static final String TASK_TOPIC_PARTITIONS = "task.assigned.partitions";
+    protected static final String TASK_CONSUMER_GROUPS = "task.assigned.groups";
+
+    public static final String CONSUMER_POLL_TIMEOUT_MILLIS = "consumer.poll.timeout.ms";
+    private static final String CONSUMER_POLL_TIMEOUT_MILLIS_DOC = "Timeout when polling source cluster.";
+    public static final long CONSUMER_POLL_TIMEOUT_MILLIS_DEFAULT = 1000L;
+
+    public static final String ADMIN_TASK_TIMEOUT_MILLIS = "admin.timeout.ms";
+    private static final String ADMIN_TASK_TIMEOUT_MILLIS_DOC = "Timeout for administrative tasks, e.g. detecting new topics.";
+    public static final long ADMIN_TASK_TIMEOUT_MILLIS_DEFAULT = 60000L;
+
+    public static final String REFRESH_TOPICS_ENABLED = REFRESH_TOPICS + ENABLED_SUFFIX;
+    private static final String REFRESH_TOPICS_ENABLED_DOC = "Whether to periodically check for new topics and partitions.";
+    public static final boolean REFRESH_TOPICS_ENABLED_DEFAULT = true;
+    public static final String REFRESH_TOPICS_INTERVAL_SECONDS = REFRESH_TOPICS + INTERVAL_SECONDS_SUFFIX;
+    private static final String REFRESH_TOPICS_INTERVAL_SECONDS_DOC = "Frequency of topic refresh.";
+    public static final long REFRESH_TOPICS_INTERVAL_SECONDS_DEFAULT = 10 * 60;
+
+    public static final String REFRESH_GROUPS_ENABLED = REFRESH_GROUPS + ENABLED_SUFFIX;
+    private static final String REFRESH_GROUPS_ENABLED_DOC = "Whether to periodically check for new consumer groups.";
+    public static final boolean REFRESH_GROUPS_ENABLED_DEFAULT = true;
+    public static final String REFRESH_GROUPS_INTERVAL_SECONDS = REFRESH_GROUPS + INTERVAL_SECONDS_SUFFIX;
+    private static final String REFRESH_GROUPS_INTERVAL_SECONDS_DOC = "Frequency of group refresh.";
+    public static final long REFRESH_GROUPS_INTERVAL_SECONDS_DEFAULT = 10 * 60;
+
+    public static final String SYNC_TOPIC_CONFIGS_ENABLED = SYNC_TOPIC_CONFIGS + ENABLED_SUFFIX;
+    private static final String SYNC_TOPIC_CONFIGS_ENABLED_DOC = "Whether to periodically configure remote topics to match their corresponding upstream topics.";
+    public static final boolean SYNC_TOPIC_CONFIGS_ENABLED_DEFAULT = true;
+    public static final String SYNC_TOPIC_CONFIGS_INTERVAL_SECONDS = SYNC_TOPIC_CONFIGS + INTERVAL_SECONDS_SUFFIX;
+    private static final String SYNC_TOPIC_CONFIGS_INTERVAL_SECONDS_DOC = "Frequency of topic config sync.";
+    public static final long SYNC_TOPIC_CONFIGS_INTERVAL_SECONDS_DEFAULT = 10 * 60;
+
+    public static final String SYNC_TOPIC_ACLS_ENABLED = SYNC_TOPIC_ACLS + ENABLED_SUFFIX;
+    private static final String SYNC_TOPIC_ACLS_ENABLED_DOC = "Whether to periodically configure remote topic ACLs to match their corresponding upstream topics.";
+    public static final boolean SYNC_TOPIC_ACLS_ENABLED_DEFAULT = true;
+    public static final String SYNC_TOPIC_ACLS_INTERVAL_SECONDS = SYNC_TOPIC_ACLS + INTERVAL_SECONDS_SUFFIX;
+    private static final String SYNC_TOPIC_ACLS_INTERVAL_SECONDS_DOC = "Frequency of topic ACL sync.";
+    public static final long SYNC_TOPIC_ACLS_INTERVAL_SECONDS_DEFAULT = 10 * 60;
+
+    public static final String EMIT_HEARTBEATS_ENABLED = EMIT_HEARTBEATS + ENABLED_SUFFIX;
+    private static final String EMIT_HEARTBEATS_ENABLED_DOC = "Whether to emit heartbeats to target cluster.";
+    public static final boolean EMIT_HEARTBEATS_ENABLED_DEFAULT = true;
+    public static final String EMIT_HEARTBEATS_INTERVAL_SECONDS = EMIT_HEARTBEATS + INTERVAL_SECONDS_SUFFIX;
+    private static final String EMIT_HEARTBEATS_INTERVAL_SECONDS_DOC = "Frequency of heartbeats.";
+    public static final long EMIT_HEARTBEATS_INTERVAL_SECONDS_DEFAULT = 1;
+
+    public static final String EMIT_CHECKPOINTS_ENABLED = EMIT_CHECKPOINTS + ENABLED_SUFFIX;
+    private static final String EMIT_CHECKPOINTS_ENABLED_DOC = "Whether to replicate consumer offsets to target cluster.";
+    public static final boolean EMIT_CHECKPOINTS_ENABLED_DEFAULT = true;
+    public static final String EMIT_CHECKPOINTS_INTERVAL_SECONDS = EMIT_CHECKPOINTS + INTERVAL_SECONDS_SUFFIX;
+    private static final String EMIT_CHECKPOINTS_INTERVAL_SECONDS_DOC = "Frequency of checkpoints.";
+    public static final long EMIT_CHECKPOINTS_INTERVAL_SECONDS_DEFAULT = 60;
+
+    public static final String TOPIC_FILTER_CLASS = "topic.filter.class";
+    private static final String TOPIC_FILTER_CLASS_DOC = "TopicFilter to use. Selects topics to replicate.";
+    public static final Class TOPIC_FILTER_CLASS_DEFAULT = DefaultTopicFilter.class;
+    public static final String GROUP_FILTER_CLASS = "group.filter.class";
+    private static final String GROUP_FILTER_CLASS_DOC = "GroupFilter to use. Selects consumer groups to replicate.";
+    public static final Class GROUP_FILTER_CLASS_DEFAULT = DefaultGroupFilter.class;
+    public static final String CONFIG_PROPERTY_FILTER_CLASS = "config.property.filter.class";
+    private static final String CONFIG_PROPERTY_FILTER_CLASS_DOC = "ConfigPropertyFilter to use. Selects topic config "
+            + " properties to replicate.";
+    public static final Class CONFIG_PROPERTY_FILTER_CLASS_DEFAULT = DefaultConfigPropertyFilter.class;
+
+    public static final String OFFSET_LAG_MAX = "offset.lag.max";
+    private static final String OFFSET_LAG_MAX_DOC = "How out-of-sync a remote partition can be before it is resynced.";
+    public static final long OFFSET_LAG_MAX_DEFAULT = 100L;
+
+    protected static final String SOURCE_CLUSTER_PREFIX = MirrorMakerConfig.SOURCE_CLUSTER_PREFIX;
+    protected static final String TARGET_CLUSTER_PREFIX = MirrorMakerConfig.TARGET_CLUSTER_PREFIX;
+    protected static final String PRODUCER_CLIENT_PREFIX = "producer.";
+    protected static final String CONSUMER_CLIENT_PREFIX = "consumer.";
+    protected static final String ADMIN_CLIENT_PREFIX = "admin.";
+    protected static final String SOURCE_ADMIN_CLIENT_PREFIX = "source.admin.";
+    protected static final String TARGET_ADMIN_CLIENT_PREFIX = "target.admin.";
+
+    public MirrorConnectorConfig(Map props) {
+        this(CONNECTOR_CONFIG_DEF, props);
+    }
+
+    protected MirrorConnectorConfig(ConfigDef configDef, Map props) {
+        super(configDef, props, true);
+    }
+
+    String connectorName() {
+        return getString(ConnectorConfig.NAME_CONFIG);
+    }
+
+    boolean enabled() {
+        return getBoolean(ENABLED);
+    }
+
+    Duration consumerPollTimeout() {
+        return Duration.ofMillis(getLong(CONSUMER_POLL_TIMEOUT_MILLIS));
+    }
+
+    Duration adminTimeout() {
+        return Duration.ofMillis(getLong(ADMIN_TASK_TIMEOUT_MILLIS));
+    }
+
+    Map sourceProducerConfig() {
+        Map props = new HashMap<>();
+        props.putAll(originalsWithPrefix(SOURCE_CLUSTER_PREFIX));
+        props.keySet().retainAll(MirrorClientConfig.CLIENT_CONFIG_DEF.names());
+        props.putAll(originalsWithPrefix(PRODUCER_CLIENT_PREFIX));
+        return props;
+    }
+
+    Map sourceConsumerConfig() {
+        Map props = new HashMap<>();
+        props.putAll(originalsWithPrefix(SOURCE_CLUSTER_PREFIX));
+        props.keySet().retainAll(MirrorClientConfig.CLIENT_CONFIG_DEF.names());
+        props.putAll(originalsWithPrefix(CONSUMER_CLIENT_PREFIX));
+        props.put("enable.auto.commit", "false");
+        props.put("auto.offset.reset", "earliest");
+        return props;
+    }
+
+    Map taskConfigForTopicPartitions(List topicPartitions) {
+        Map props = originalsStrings();
+        String topicPartitionsString = topicPartitions.stream()
+                .map(MirrorUtils::encodeTopicPartition)
+                .collect(Collectors.joining(","));
+        props.put(TASK_TOPIC_PARTITIONS, topicPartitionsString);
+        return props;
+    }
+
+    Map taskConfigForConsumerGroups(List groups) {
+        Map props = originalsStrings();
+        props.put(TASK_CONSUMER_GROUPS, String.join(",", groups));
+        return props;
+    }
+
+    Map targetAdminConfig() {
+        Map props = new HashMap<>();
+        props.putAll(originalsWithPrefix(TARGET_CLUSTER_PREFIX));
+        props.keySet().retainAll(MirrorClientConfig.CLIENT_CONFIG_DEF.names());
+        props.putAll(originalsWithPrefix(ADMIN_CLIENT_PREFIX));
+        props.putAll(originalsWithPrefix(TARGET_ADMIN_CLIENT_PREFIX));
+        return props;
+    }
+
+    Map sourceAdminConfig() {
+        Map props = new HashMap<>();
+        props.putAll(originalsWithPrefix(SOURCE_CLUSTER_PREFIX));
+        props.keySet().retainAll(MirrorClientConfig.CLIENT_CONFIG_DEF.names());
+        props.putAll(originalsWithPrefix(ADMIN_CLIENT_PREFIX));
+        props.putAll(originalsWithPrefix(SOURCE_ADMIN_CLIENT_PREFIX));
+        return props;
+    }
+
+    List metricsReporters() {
+        List reporters = getConfiguredInstances(
+                CommonClientConfigs.METRIC_REPORTER_CLASSES_CONFIG, MetricsReporter.class);
+        reporters.add(new JmxReporter("kafka.connect.mirror"));
+        return reporters;
+    }
+
+    String sourceClusterAlias() {
+        return getString(SOURCE_CLUSTER_ALIAS);
+    }
+
+    String targetClusterAlias() {
+        return getString(TARGET_CLUSTER_ALIAS);
+    }
+
+    String offsetSyncsTopic() {
+        // ".internal" suffix ensures this doesn't get replicated
+        return "mm2-offset-syncs." + targetClusterAlias() + ".internal";
+    }
+
+    String heartbeatsTopic() {
+        return MirrorClientConfig.HEARTBEATS_TOPIC;
+    }
+
+    // e.g. source1.heartbeats
+    String targetHeartbeatsTopic() {
+        return replicationPolicy().formatRemoteTopic(sourceClusterAlias(), heartbeatsTopic());
+    }
+
+    String checkpointsTopic() {
+        // Checkpoint topics are not "remote topics", as they are not replicated, so we don't
+        // need to use ReplicationPolicy here.
+        return sourceClusterAlias() + MirrorClientConfig.CHECKPOINTS_TOPIC_SUFFIX;
+    }
+
+    long maxOffsetLag() {
+        return getLong(OFFSET_LAG_MAX);
+    }
+
+    Duration emitHeartbeatsInterval() {
+        if (getBoolean(EMIT_HEARTBEATS_ENABLED)) {
+            return Duration.ofSeconds(getLong(EMIT_HEARTBEATS_INTERVAL_SECONDS));
+        } else {
+            // negative interval to disable
+            return Duration.ofMillis(-1);
+        }
+    }
+
+    Duration emitCheckpointsInterval() {
+        if (getBoolean(EMIT_CHECKPOINTS_ENABLED)) {
+            return Duration.ofSeconds(getLong(EMIT_CHECKPOINTS_INTERVAL_SECONDS));
+        } else {
+            // negative interval to disable
+            return Duration.ofMillis(-1);
+        }
+    }
+
+    Duration refreshTopicsInterval() {
+        if (getBoolean(REFRESH_TOPICS_ENABLED)) {
+            return Duration.ofSeconds(getLong(REFRESH_TOPICS_INTERVAL_SECONDS));
+        } else {
+            // negative interval to disable
+            return Duration.ofMillis(-1);
+        }
+    }
+
+    Duration refreshGroupsInterval() {
+        if (getBoolean(REFRESH_GROUPS_ENABLED)) {
+            return Duration.ofSeconds(getLong(REFRESH_GROUPS_INTERVAL_SECONDS));
+        } else {
+            // negative interval to disable
+            return Duration.ofMillis(-1);
+        }
+    }
+
+    Duration syncTopicConfigsInterval() {
+        if (getBoolean(SYNC_TOPIC_CONFIGS_ENABLED)) {
+            return Duration.ofSeconds(getLong(SYNC_TOPIC_CONFIGS_INTERVAL_SECONDS));
+        } else {
+            // negative interval to disable
+            return Duration.ofMillis(-1);
+        }
+    }
+
+    Duration syncTopicAclsInterval() {
+        if (getBoolean(SYNC_TOPIC_ACLS_ENABLED)) {
+            return Duration.ofSeconds(getLong(SYNC_TOPIC_ACLS_INTERVAL_SECONDS));
+        } else {
+            // negative interval to disable
+            return Duration.ofMillis(-1);
+        }
+    }
+
+    ReplicationPolicy replicationPolicy() {
+        return getConfiguredInstance(REPLICATION_POLICY_CLASS, ReplicationPolicy.class);
+    }
+
+    int replicationFactor() {
+        return getInt(REPLICATION_FACTOR);
+    }
+
+    short heartbeatsTopicReplicationFactor() {
+        return getShort(HEARTBEATS_TOPIC_REPLICATION_FACTOR);
+    }
+
+    short checkpointsTopicReplicationFactor() {
+        return getShort(CHECKPOINTS_TOPIC_REPLICATION_FACTOR);
+    }
+
+    short offsetSyncsTopicReplicationFactor() {
+        return getShort(OFFSET_SYNCS_TOPIC_REPLICATION_FACTOR);
+    }
+
+    TopicFilter topicFilter() {
+        return getConfiguredInstance(TOPIC_FILTER_CLASS, TopicFilter.class);
+    }
+
+    GroupFilter groupFilter() {
+        return getConfiguredInstance(GROUP_FILTER_CLASS, GroupFilter.class);
+    }
+
+    ConfigPropertyFilter configPropertyFilter() {
+        return getConfiguredInstance(CONFIG_PROPERTY_FILTER_CLASS, ConfigPropertyFilter.class);
+    }
+
+    protected static final ConfigDef CONNECTOR_CONFIG_DEF = ConnectorConfig.configDef()
+            .define(
+                    ENABLED,
+                    ConfigDef.Type.BOOLEAN,
+                    true,
+                    ConfigDef.Importance.LOW,
+                    ENABLED_DOC)
+            .define(
+                    TOPICS,
+                    ConfigDef.Type.LIST,
+                    TOPICS_DEFAULT,
+                    ConfigDef.Importance.HIGH,
+                    TOPICS_DOC) 
+            .define(
+                    TOPICS_BLACKLIST,
+                    ConfigDef.Type.LIST,
+                    TOPICS_BLACKLIST_DEFAULT,
+                    ConfigDef.Importance.HIGH,
+                    TOPICS_BLACKLIST_DOC)
+            .define(
+                    GROUPS,
+                    ConfigDef.Type.LIST,
+                    GROUPS_DEFAULT,
+                    ConfigDef.Importance.HIGH,
+                    GROUPS_DOC) 
+            .define(
+                    GROUPS_BLACKLIST,
+                    ConfigDef.Type.LIST,
+                    GROUPS_BLACKLIST_DEFAULT,
+                    ConfigDef.Importance.HIGH,
+                    GROUPS_BLACKLIST_DOC)
+            .define(
+                    CONFIG_PROPERTIES_BLACKLIST,
+                    ConfigDef.Type.LIST,
+                    CONFIG_PROPERTIES_BLACKLIST_DEFAULT,
+                    ConfigDef.Importance.HIGH,
+                    CONFIG_PROPERTIES_BLACKLIST_DOC)
+            .define(
+                    TOPIC_FILTER_CLASS,
+                    ConfigDef.Type.CLASS,
+                    TOPIC_FILTER_CLASS_DEFAULT,
+                    ConfigDef.Importance.LOW,
+                    TOPIC_FILTER_CLASS_DOC)
+            .define(
+                    GROUP_FILTER_CLASS,
+                    ConfigDef.Type.CLASS,
+                    GROUP_FILTER_CLASS_DEFAULT,
+                    ConfigDef.Importance.LOW,
+                    GROUP_FILTER_CLASS_DOC)
+            .define(
+                    CONFIG_PROPERTY_FILTER_CLASS,
+                    ConfigDef.Type.CLASS,
+                    CONFIG_PROPERTY_FILTER_CLASS_DEFAULT,
+                    ConfigDef.Importance.LOW,
+                    CONFIG_PROPERTY_FILTER_CLASS_DOC)
+            .define(
+                    SOURCE_CLUSTER_ALIAS,
+                    ConfigDef.Type.STRING,
+                    ConfigDef.Importance.HIGH,
+                    SOURCE_CLUSTER_ALIAS_DOC)
+            .define(
+                    TARGET_CLUSTER_ALIAS,
+                    ConfigDef.Type.STRING,
+                    TARGET_CLUSTER_ALIAS_DEFAULT,
+                    ConfigDef.Importance.HIGH,
+                    TARGET_CLUSTER_ALIAS_DOC)
+            .define(
+                    CONSUMER_POLL_TIMEOUT_MILLIS,
+                    ConfigDef.Type.LONG,
+                    CONSUMER_POLL_TIMEOUT_MILLIS_DEFAULT,
+                    ConfigDef.Importance.LOW,
+                    CONSUMER_POLL_TIMEOUT_MILLIS_DOC)
+            .define(
+                    ADMIN_TASK_TIMEOUT_MILLIS,
+                    ConfigDef.Type.LONG,
+                    ADMIN_TASK_TIMEOUT_MILLIS_DEFAULT,
+                    ConfigDef.Importance.LOW,
+                    ADMIN_TASK_TIMEOUT_MILLIS_DOC)
+            .define(
+                    REFRESH_TOPICS_ENABLED,
+                    ConfigDef.Type.BOOLEAN,
+                    REFRESH_TOPICS_ENABLED_DEFAULT,
+                    ConfigDef.Importance.LOW,
+                    REFRESH_TOPICS_ENABLED_DOC)
+            .define(
+                    REFRESH_TOPICS_INTERVAL_SECONDS,
+                    ConfigDef.Type.LONG,
+                    REFRESH_TOPICS_INTERVAL_SECONDS_DEFAULT,
+                    ConfigDef.Importance.LOW,
+                    REFRESH_TOPICS_INTERVAL_SECONDS_DOC)
+            .define(
+                    REFRESH_GROUPS_ENABLED,
+                    ConfigDef.Type.BOOLEAN,
+                    REFRESH_GROUPS_ENABLED_DEFAULT,
+                    ConfigDef.Importance.LOW,
+                    REFRESH_GROUPS_ENABLED_DOC)
+            .define(
+                    REFRESH_GROUPS_INTERVAL_SECONDS,
+                    ConfigDef.Type.LONG,
+                    REFRESH_GROUPS_INTERVAL_SECONDS_DEFAULT,
+                    ConfigDef.Importance.LOW,
+                    REFRESH_GROUPS_INTERVAL_SECONDS_DOC)
+            .define(
+                    SYNC_TOPIC_CONFIGS_ENABLED,
+                    ConfigDef.Type.BOOLEAN,
+                    SYNC_TOPIC_CONFIGS_ENABLED_DEFAULT,
+                    ConfigDef.Importance.LOW,
+                    SYNC_TOPIC_CONFIGS_ENABLED_DOC)
+            .define(
+                    SYNC_TOPIC_CONFIGS_INTERVAL_SECONDS,
+                    ConfigDef.Type.LONG,
+                    SYNC_TOPIC_CONFIGS_INTERVAL_SECONDS_DEFAULT,
+                    ConfigDef.Importance.LOW,
+                    SYNC_TOPIC_CONFIGS_INTERVAL_SECONDS_DOC)
+            .define(
+                    SYNC_TOPIC_ACLS_ENABLED,
+                    ConfigDef.Type.BOOLEAN,
+                    SYNC_TOPIC_ACLS_ENABLED_DEFAULT,
+                    ConfigDef.Importance.LOW,
+                    SYNC_TOPIC_ACLS_ENABLED_DOC)
+            .define(
+                    SYNC_TOPIC_ACLS_INTERVAL_SECONDS,
+                    ConfigDef.Type.LONG,
+                    SYNC_TOPIC_ACLS_INTERVAL_SECONDS_DEFAULT,
+                    ConfigDef.Importance.LOW,
+                    SYNC_TOPIC_ACLS_INTERVAL_SECONDS_DOC)
+            .define(
+                    EMIT_HEARTBEATS_ENABLED,
+                    ConfigDef.Type.BOOLEAN,
+                    EMIT_HEARTBEATS_ENABLED_DEFAULT,
+                    ConfigDef.Importance.LOW,
+                    EMIT_HEARTBEATS_ENABLED_DOC)
+            .define(
+                    EMIT_HEARTBEATS_INTERVAL_SECONDS,
+                    ConfigDef.Type.LONG,
+                    EMIT_HEARTBEATS_INTERVAL_SECONDS_DEFAULT,
+                    ConfigDef.Importance.LOW,
+                    EMIT_HEARTBEATS_INTERVAL_SECONDS_DOC)
+            .define(
+                    EMIT_CHECKPOINTS_ENABLED,
+                    ConfigDef.Type.BOOLEAN,
+                    EMIT_CHECKPOINTS_ENABLED_DEFAULT,
+                    ConfigDef.Importance.LOW,
+                    EMIT_CHECKPOINTS_ENABLED_DOC)
+            .define(
+                    EMIT_CHECKPOINTS_INTERVAL_SECONDS,
+                    ConfigDef.Type.LONG,
+                    EMIT_CHECKPOINTS_INTERVAL_SECONDS_DEFAULT,
+                    ConfigDef.Importance.LOW,
+                    EMIT_CHECKPOINTS_INTERVAL_SECONDS_DOC)
+            .define(
+                    REPLICATION_POLICY_CLASS,
+                    ConfigDef.Type.CLASS,
+                    REPLICATION_POLICY_CLASS_DEFAULT,
+                    ConfigDef.Importance.LOW,
+                    REPLICATION_POLICY_CLASS_DOC)
+            .define(
+                    REPLICATION_POLICY_SEPARATOR,
+                    ConfigDef.Type.STRING,
+                    REPLICATION_POLICY_SEPARATOR_DEFAULT,
+                    ConfigDef.Importance.LOW,
+                    REPLICATION_POLICY_SEPARATOR_DOC)
+            .define(
+                    REPLICATION_FACTOR,
+                    ConfigDef.Type.INT,
+                    REPLICATION_FACTOR_DEFAULT,
+                    ConfigDef.Importance.LOW,
+                    REPLICATION_FACTOR_DOC)
+            .define(
+                    HEARTBEATS_TOPIC_REPLICATION_FACTOR,
+                    ConfigDef.Type.SHORT,
+                    HEARTBEATS_TOPIC_REPLICATION_FACTOR_DEFAULT,
+                    ConfigDef.Importance.LOW,
+                    HEARTBEATS_TOPIC_REPLICATION_FACTOR_DOC)
+            .define(
+                    CHECKPOINTS_TOPIC_REPLICATION_FACTOR,
+                    ConfigDef.Type.SHORT,
+                    CHECKPOINTS_TOPIC_REPLICATION_FACTOR_DEFAULT,
+                    ConfigDef.Importance.LOW,
+                    CHECKPOINTS_TOPIC_REPLICATION_FACTOR_DOC)
+            .define(
+                    OFFSET_SYNCS_TOPIC_REPLICATION_FACTOR,
+                    ConfigDef.Type.SHORT,
+                    OFFSET_SYNCS_TOPIC_REPLICATION_FACTOR_DEFAULT,
+                    ConfigDef.Importance.LOW,
+                    OFFSET_SYNCS_TOPIC_REPLICATION_FACTOR_DOC)
+            .define(
+                    OFFSET_LAG_MAX,
+                    ConfigDef.Type.LONG,
+                    OFFSET_LAG_MAX_DEFAULT,
+                    ConfigDef.Importance.LOW,
+                    OFFSET_LAG_MAX_DOC)
+            .define(
+                    CommonClientConfigs.METRIC_REPORTER_CLASSES_CONFIG,
+                    ConfigDef.Type.LIST,
+                    null,
+                    ConfigDef.Importance.LOW,
+                    CommonClientConfigs.METRIC_REPORTER_CLASSES_DOC)
+            .define(
+                    CommonClientConfigs.SECURITY_PROTOCOL_CONFIG,
+                    ConfigDef.Type.STRING,
+                    CommonClientConfigs.DEFAULT_SECURITY_PROTOCOL,
+                    ConfigDef.Importance.MEDIUM,
+                    CommonClientConfigs.SECURITY_PROTOCOL_DOC)
+            .withClientSslSupport()
+            .withClientSaslSupport();
+}
diff --git a/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorHeartbeatConnector.java b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorHeartbeatConnector.java
new file mode 100644
index 0000000000000..3942c8454ab72
--- /dev/null
+++ b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorHeartbeatConnector.java
@@ -0,0 +1,71 @@
+/*
+ * 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.mirror;
+
+import org.apache.kafka.connect.connector.Task;
+import org.apache.kafka.connect.source.SourceConnector;
+import org.apache.kafka.common.config.ConfigDef;
+import org.apache.kafka.common.utils.Utils;
+
+import java.util.Map;
+import java.util.List;
+import java.util.Collections;
+
+/** Emits heartbeats to Kafka.
+ */
+public class MirrorHeartbeatConnector extends SourceConnector {
+    private MirrorConnectorConfig config;
+    private Scheduler scheduler;
+
+    @Override
+    public void start(Map props) {
+        config = new MirrorConnectorConfig(props);
+        scheduler = new Scheduler(MirrorHeartbeatConnector.class, config.adminTimeout());
+        scheduler.execute(this::createInternalTopics, "creating internal topics");
+    }
+
+    @Override
+    public void stop() {
+        Utils.closeQuietly(scheduler, "scheduler");
+    }
+
+    @Override
+    public Class taskClass() {
+        return MirrorHeartbeatTask.class;
+    }
+
+    @Override
+    public List> taskConfigs(int maxTasks) {
+        // just need a single task
+        return Collections.singletonList(config.originalsStrings());
+    }
+
+    @Override
+    public ConfigDef config() {
+        return MirrorConnectorConfig.CONNECTOR_CONFIG_DEF;
+    }
+
+    @Override
+    public String version() {
+        return "1";
+    }
+
+    private void createInternalTopics() {
+        MirrorUtils.createSinglePartitionCompactedTopic(config.heartbeatsTopic(),
+            config.heartbeatsTopicReplicationFactor(), config.targetAdminConfig());
+    }
+}
diff --git a/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorHeartbeatTask.java b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorHeartbeatTask.java
new file mode 100644
index 0000000000000..6bfe441512450
--- /dev/null
+++ b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorHeartbeatTask.java
@@ -0,0 +1,84 @@
+/*
+ * 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.mirror;
+
+import org.apache.kafka.connect.source.SourceTask;
+import org.apache.kafka.connect.source.SourceRecord;
+import org.apache.kafka.connect.data.Schema;
+
+import java.util.Map;
+import java.util.List;
+import java.util.Collections;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.time.Duration;
+
+/** Emits heartbeats. */
+public class MirrorHeartbeatTask extends SourceTask {
+    private String sourceClusterAlias;
+    private String targetClusterAlias;
+    private String heartbeatsTopic;
+    private Duration interval;
+    private CountDownLatch stopped;
+
+    @Override
+    public void start(Map props) {
+        stopped = new CountDownLatch(1);
+        MirrorTaskConfig config = new MirrorTaskConfig(props);
+        sourceClusterAlias = config.sourceClusterAlias();
+        targetClusterAlias = config.targetClusterAlias();
+        heartbeatsTopic = config.heartbeatsTopic();
+        interval = config.emitHeartbeatsInterval();
+    }
+
+    @Override
+    public void commit() throws InterruptedException {
+        // nop
+    }
+
+    @Override
+    public void stop() {
+        stopped.countDown();
+    }
+
+    @Override
+    public String version() {
+        return "1";
+    }
+
+    @Override
+    public List poll() throws InterruptedException {
+        // pause to throttle, unless we've stopped
+        if (stopped.await(interval.toMillis(), TimeUnit.MILLISECONDS)) {
+            // SourceWorkerTask expects non-zero batches or null
+            return null;
+        }
+        long timestamp = System.currentTimeMillis();
+        Heartbeat heartbeat = new Heartbeat(sourceClusterAlias, targetClusterAlias, timestamp);
+        SourceRecord record = new SourceRecord(
+            heartbeat.connectPartition(), MirrorUtils.wrapOffset(0),
+            heartbeatsTopic, 0,
+            Schema.BYTES_SCHEMA, heartbeat.recordKey(),
+            Schema.BYTES_SCHEMA, heartbeat.recordValue(),
+            timestamp);
+        return Collections.singletonList(record);
+    }
+
+    @Override
+    public void commitRecord(SourceRecord record) {
+    }
+}
diff --git a/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorMaker.java b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorMaker.java
new file mode 100644
index 0000000000000..5ba689062072f
--- /dev/null
+++ b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorMaker.java
@@ -0,0 +1,309 @@
+/*
+ * 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.mirror;
+
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.common.utils.Utils;
+import org.apache.kafka.common.utils.Exit;
+import org.apache.kafka.connect.runtime.Herder;
+import org.apache.kafka.connect.runtime.isolation.Plugins;
+import org.apache.kafka.connect.runtime.Worker;
+import org.apache.kafka.connect.runtime.WorkerConfigTransformer;
+import org.apache.kafka.connect.runtime.distributed.DistributedConfig;
+import org.apache.kafka.connect.runtime.distributed.DistributedHerder;
+import org.apache.kafka.connect.runtime.distributed.NotLeaderException;
+import org.apache.kafka.connect.storage.KafkaOffsetBackingStore;
+import org.apache.kafka.connect.storage.StatusBackingStore;
+import org.apache.kafka.connect.storage.KafkaStatusBackingStore;
+import org.apache.kafka.connect.storage.ConfigBackingStore;
+import org.apache.kafka.connect.storage.KafkaConfigBackingStore;
+import org.apache.kafka.connect.storage.Converter;
+import org.apache.kafka.connect.util.ConnectUtils;
+import org.apache.kafka.connect.connector.policy.AllConnectorClientConfigOverridePolicy;
+import org.apache.kafka.connect.connector.policy.ConnectorClientConfigOverridePolicy;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import net.sourceforge.argparse4j.impl.Arguments;
+import net.sourceforge.argparse4j.inf.Namespace;
+import net.sourceforge.argparse4j.inf.ArgumentParser;
+import net.sourceforge.argparse4j.inf.ArgumentParserException;
+import net.sourceforge.argparse4j.ArgumentParsers;
+
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.Map;
+import java.util.HashMap;
+import java.util.Set;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Arrays;
+import java.util.Properties;
+import java.util.stream.Collectors;
+import java.io.File;
+
+/**
+ *  Entry point for "MirrorMaker 2.0".
+ *  
      
+ *  MirrorMaker runs a set of Connectors between multiple clusters, in order to replicate data, configuration,
+ *  ACL rules, and consumer group state.
+ *  
      
+ *  
      
+ *  Configuration is via a top-level "mm2.properties" file, which supports per-cluster and per-replication
+ *  sub-configs. Each source->target replication must be explicitly enabled. For example:
+ *  
      
+ *  
      + *    clusters = primary, backup
+ *    primary.bootstrap.servers = vip1:9092
+ *    backup.bootstrap.servers = vip2:9092
+ *    primary->backup.enabled = true
+ *    backup->primary.enabled = true
+ *  
      
+ *  
      
+ *  Run as follows:
+ *  
      
+ *  
      + *    ./bin/connect-mirror-maker.sh mm2.properties
+ *  
      
+ *  
      
+ *  Additional information and example configurations are provided in ./connect/mirror/README.md
+ *  
      
+ */
+public class MirrorMaker {
+    private static final Logger log = LoggerFactory.getLogger(MirrorMaker.class);
+
+    private static final long SHUTDOWN_TIMEOUT_SECONDS = 60L;
+    private static final ConnectorClientConfigOverridePolicy CLIENT_CONFIG_OVERRIDE_POLICY =
+            new AllConnectorClientConfigOverridePolicy();
+
+    private static final List CONNECTOR_CLASSES = Arrays.asList(
+        MirrorSourceConnector.class,
+        MirrorHeartbeatConnector.class,
+        MirrorCheckpointConnector.class);
+ 
+    private final Map herders = new HashMap<>();
+    private CountDownLatch startLatch;
+    private CountDownLatch stopLatch;
+    private final AtomicBoolean shutdown = new AtomicBoolean(false);
+    private final ShutdownHook shutdownHook;
+    private final String advertisedBaseUrl;
+    private final Time time;
+    private final MirrorMakerConfig config;
+    private final Set clusters;
+    private final Set herderPairs;
+
+    /**
+     * @param config    MM2 configuration from mm2.properties file
+     * @param clusters  target clusters for this node. These must match cluster
+     *                  aliases as defined in the config. If null or empty list,
+     *                  uses all clusters in the config.
+     * @param time      time source
+     */
+    public MirrorMaker(MirrorMakerConfig config, List clusters, Time time) {
+        log.debug("Kafka MirrorMaker instance created");
+        this.time = time;
+        this.advertisedBaseUrl = "NOTUSED";
+        this.config = config;
+        if (clusters != null && !clusters.isEmpty()) {
+            this.clusters = new HashSet<>(clusters);
+        } else {
+            // default to all clusters
+            this.clusters = config.clusters();
+        }
+        log.info("Targeting clusters {}", this.clusters);
+        this.herderPairs = config.clusterPairs().stream()
+            .filter(x -> this.clusters.contains(x.target()))
+            .collect(Collectors.toSet());
+        if (herderPairs.isEmpty()) {
+            throw new IllegalArgumentException("No source->target replication flows.");
+        }
+        this.herderPairs.forEach(x -> addHerder(x));
+        shutdownHook = new ShutdownHook();
+    }
+
+    /**
+     * @param config    MM2 configuration from mm2.properties file
+     * @param clusters  target clusters for this node. These must match cluster
+     *                  aliases as defined in the config. If null or empty list,
+     *                  uses all clusters in the config.
+     * @param time      time source
+     */
+    public MirrorMaker(Map config, List clusters, Time time) {
+        this(new MirrorMakerConfig(config), clusters, time);
+    }
+
+    public MirrorMaker(Map props, List clusters) {
+        this(props, clusters, Time.SYSTEM);
+    }
+
+    public MirrorMaker(Map props) {
+        this(props, null);
+    }
+
+
+    public void start() {
+        log.info("Kafka MirrorMaker starting with {} herders.", herders.size());
+        if (startLatch != null) {
+            throw new IllegalStateException("MirrorMaker instance already started");
+        }
+        startLatch = new CountDownLatch(herders.size());
+        stopLatch = new CountDownLatch(herders.size());
+        Runtime.getRuntime().addShutdownHook(shutdownHook);
+        for (Herder herder : herders.values()) {
+            try {
+                herder.start();
+            } finally {
+                startLatch.countDown();
+            }
+        }
+        log.info("Configuring connectors...");
+        herderPairs.forEach(x -> configureConnectors(x));
+        log.info("Kafka MirrorMaker started");
+    }
+
+    public void stop() {
+        boolean wasShuttingDown = shutdown.getAndSet(true);
+        if (!wasShuttingDown) {
+            log.info("Kafka MirrorMaker stopping");
+            for (Herder herder : herders.values()) {
+                try {
+                    herder.stop();
+                } finally {
+                    stopLatch.countDown();
+                }
+            }
+            log.info("Kafka MirrorMaker stopped.");
+        }
+    }
+
+    public void awaitStop() {
+        try {
+            stopLatch.await();
+        } catch (InterruptedException e) {
+            log.error("Interrupted waiting for MirrorMaker to shutdown");
+        }
+    }
+
+    private void configureConnector(SourceAndTarget sourceAndTarget, Class connectorClass) {
+        checkHerder(sourceAndTarget);
+        Map connectorProps = config.connectorBaseConfig(sourceAndTarget, connectorClass);
+        herders.get(sourceAndTarget)
+                .putConnectorConfig(connectorClass.getSimpleName(), connectorProps, true, (e, x) -> {
+                    if (e instanceof NotLeaderException) {
+                        // No way to determine if the connector is a leader or not beforehand.
+                        log.info("Connector {} is a follower. Using existing configuration.", sourceAndTarget);
+                    } else {
+                        log.info("Connector {} configured.", sourceAndTarget, e);
+                    }
+                });
+    }
+
+    private void checkHerder(SourceAndTarget sourceAndTarget) {
+        if (!herders.containsKey(sourceAndTarget)) {
+            throw new IllegalArgumentException("No herder for " + sourceAndTarget.toString());
+        }
+    }
+
+    private void configureConnectors(SourceAndTarget sourceAndTarget) {
+        CONNECTOR_CLASSES.forEach(x -> configureConnector(sourceAndTarget, x));
+    }
+
+    private void addHerder(SourceAndTarget sourceAndTarget) {
+        log.info("creating herder for " + sourceAndTarget.toString());
+        Map workerProps = config.workerConfig(sourceAndTarget);
+        String advertisedUrl = advertisedBaseUrl + "/" + sourceAndTarget.source();
+        String workerId = sourceAndTarget.toString();
+        Plugins plugins = new Plugins(workerProps);
+        plugins.compareAndSwapWithDelegatingLoader();
+        DistributedConfig distributedConfig = new DistributedConfig(workerProps);
+        String kafkaClusterId = ConnectUtils.lookupKafkaClusterId(distributedConfig);
+        KafkaOffsetBackingStore offsetBackingStore = new KafkaOffsetBackingStore();
+        offsetBackingStore.configure(distributedConfig);
+        Worker worker = new Worker(workerId, time, plugins, distributedConfig, offsetBackingStore, CLIENT_CONFIG_OVERRIDE_POLICY);
+        WorkerConfigTransformer configTransformer = worker.configTransformer();
+        Converter internalValueConverter = worker.getInternalValueConverter();
+        StatusBackingStore statusBackingStore = new KafkaStatusBackingStore(time, internalValueConverter);
+        statusBackingStore.configure(distributedConfig);
+        ConfigBackingStore configBackingStore = new KafkaConfigBackingStore(
+                internalValueConverter,
+                distributedConfig,
+                configTransformer);
+        Herder herder = new DistributedHerder(distributedConfig, time, worker,
+                kafkaClusterId, statusBackingStore, configBackingStore,
+                advertisedUrl, CLIENT_CONFIG_OVERRIDE_POLICY);
+        herders.put(sourceAndTarget, herder);
+    }
+
+    private class ShutdownHook extends Thread {
+        @Override
+        public void run() {
+            try {
+                if (!startLatch.await(SHUTDOWN_TIMEOUT_SECONDS, TimeUnit.SECONDS)) {
+                    log.error("Timed out in shutdown hook waiting for MirrorMaker startup to finish. Unable to shutdown cleanly.");
+                }
+            } catch (InterruptedException e) {
+                log.error("Interrupted in shutdown hook while waiting for MirrorMaker startup to finish. Unable to shutdown cleanly.");
+            } finally {
+                MirrorMaker.this.stop();
+            }
+        }
+    }
+
+    public static void main(String[] args) {
+        ArgumentParser parser = ArgumentParsers.newArgumentParser("connect-mirror-maker");
+        parser.description("MirrorMaker 2.0 driver");
+        parser.addArgument("config").type(Arguments.fileType().verifyCanRead())
+            .metavar("mm2.properties").required(true)
+            .help("MM2 configuration file.");
+        parser.addArgument("--clusters").nargs("+").metavar("CLUSTER").required(false)
+            .help("Target cluster to use for this node.");
+        Namespace ns;
+        try {
+            ns = parser.parseArgs(args);
+        } catch (ArgumentParserException e) {
+            parser.handleError(e);
+            Exit.exit(-1);
+            return;
+        }
+        File configFile = (File) ns.get("config");
+        List clusters = ns.getList("clusters");
+        try {
+            log.info("Kafka MirrorMaker initializing ...");
+
+            Properties props = Utils.loadProps(configFile.getPath());
+            Map config = Utils.propsToStringMap(props);
+            MirrorMaker mirrorMaker = new MirrorMaker(config, clusters, Time.SYSTEM);
+            
+            try {
+                mirrorMaker.start();
+            } catch (Exception e) {
+                log.error("Failed to start MirrorMaker", e);
+                mirrorMaker.stop();
+                Exit.exit(3);
+            }
+
+            mirrorMaker.awaitStop();
+
+        } catch (Throwable t) {
+            log.error("Stopping due to error", t);
+            Exit.exit(2);
+        }
+    }
+
+}
diff --git a/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorMakerConfig.java b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorMakerConfig.java
new file mode 100644
index 0000000000000..df5d38f7c53b2
--- /dev/null
+++ b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorMakerConfig.java
@@ -0,0 +1,255 @@
+/*
+ * 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.mirror;
+
+import org.apache.kafka.common.utils.Utils;
+import org.apache.kafka.common.config.AbstractConfig;
+import org.apache.kafka.common.config.ConfigDef;
+import org.apache.kafka.common.config.ConfigDef.Type;
+import org.apache.kafka.common.config.ConfigDef.Importance;
+import org.apache.kafka.common.config.provider.ConfigProvider;
+import org.apache.kafka.common.config.ConfigTransformer;
+import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.connect.runtime.WorkerConfig;
+import org.apache.kafka.connect.runtime.distributed.DistributedConfig;
+import org.apache.kafka.connect.runtime.isolation.Plugins;
+
+import java.util.Map;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Set;
+import java.util.HashSet;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.stream.Collectors;
+
+/** Top-level config describing replication flows between multiple Kafka clusters.
+ *
+ *  Supports cluster-level properties of the form cluster.x.y.z, and replication-level
+ *  properties of the form source->target.x.y.z.
+ *  e.g.
+ *
+ *      clusters = A, B, C
+ *      A.bootstrap.servers = aaa:9092
+ *      A.security.protocol = SSL
+ *      --->%---
+ *      A->B.enabled = true
+ *      A->B.producer.client.id = "A-B-producer"
+ *      --->%---
+ *
+ */
+public class MirrorMakerConfig extends AbstractConfig {
+
+    public static final String CLUSTERS_CONFIG = "clusters";
+    private static final String CLUSTERS_DOC = "List of cluster aliases.";
+    public static final String CONFIG_PROVIDERS_CONFIG = WorkerConfig.CONFIG_PROVIDERS_CONFIG;
+    private static final String CONFIG_PROVIDERS_DOC = "Names of ConfigProviders to use.";
+
+    private static final String NAME = "name";
+    private static final String CONNECTOR_CLASS = "connector.class";
+    private static final String SOURCE_CLUSTER_ALIAS = "source.cluster.alias";
+    private static final String TARGET_CLUSTER_ALIAS = "target.cluster.alias";
+    private static final String GROUP_ID_CONFIG = "group.id";
+    private static final String KEY_CONVERTER_CLASS_CONFIG = "key.converter";
+    private static final String VALUE_CONVERTER_CLASS_CONFIG = "value.converter";
+    private static final String HEADER_CONVERTER_CLASS_CONFIG = "header.converter";
+    private static final String BYTE_ARRAY_CONVERTER_CLASS =
+        "org.apache.kafka.connect.converters.ByteArrayConverter";
+    private static final String REPLICATION_FACTOR = "replication.factor";
+
+    static final String SOURCE_CLUSTER_PREFIX = "source.cluster.";
+    static final String TARGET_CLUSTER_PREFIX = "target.cluster.";
+
+    private final Plugins plugins;
+   
+    public MirrorMakerConfig(Map props) {
+        super(CONFIG_DEF, props, true);
+        plugins = new Plugins(originalsStrings());
+    }
+
+    public Set clusters() {
+        return new HashSet<>(getList(CLUSTERS_CONFIG));
+    }
+
+    public List clusterPairs() {
+        List pairs = new ArrayList<>();
+        Set clusters = clusters();
+        for (String source : clusters) {
+            for (String target : clusters) {
+                SourceAndTarget sourceAndTarget = new SourceAndTarget(source, target);
+                if (!source.equals(target)) {
+                    pairs.add(sourceAndTarget);
+                }
+            }
+        }
+        return pairs;
+    }
+
+    /** Construct a MirrorClientConfig from properties of the form cluster.x.y.z.
+      * Use to connect to a cluster based on the MirrorMaker top-level config file.
+      */
+    public MirrorClientConfig clientConfig(String cluster) {
+        Map props = new HashMap<>();
+        props.putAll(originalsStrings());
+        props.putAll(clusterProps(cluster));
+        return new MirrorClientConfig(transform(props));
+    }
+
+    // loads properties of the form cluster.x.y.z
+    Map clusterProps(String cluster) {
+        Map props = new HashMap<>();
+        Map strings = originalsStrings();
+
+        props.putAll(stringsWithPrefixStripped(cluster + "."));
+
+        for (String k : MirrorClientConfig.CLIENT_CONFIG_DEF.names()) {
+            String v = props.get(k);
+            if (v != null) {
+                props.putIfAbsent("producer." + k, v);
+                props.putIfAbsent("consumer." + k, v);
+                props.putIfAbsent("admin." + k, v);
+            }
+        }
+
+        for (String k : MirrorClientConfig.CLIENT_CONFIG_DEF.names()) {
+            String v = strings.get(k);
+            if (v != null) {
+                props.putIfAbsent("producer." + k, v);
+                props.putIfAbsent("consumer." + k, v);
+                props.putIfAbsent("admin." + k, v);
+                props.putIfAbsent(k, v);
+            }
+        }
+ 
+        return props;
+    }
+
+    // loads worker configs based on properties of the form x.y.z and cluster.x.y.z 
+    Map workerConfig(SourceAndTarget sourceAndTarget) {
+        Map props = new HashMap<>();
+        props.putAll(clusterProps(sourceAndTarget.target()));
+      
+        // Accept common top-level configs that are otherwise ignored by MM2.
+        // N.B. all other worker properties should be configured for specific herders,
+        // e.g. primary->backup.client.id
+        props.putAll(stringsWithPrefix("offset.storage"));
+        props.putAll(stringsWithPrefix("config.storage"));
+        props.putAll(stringsWithPrefix("status.storage"));
+        props.putAll(stringsWithPrefix("key.converter")); 
+        props.putAll(stringsWithPrefix("value.converter")); 
+        props.putAll(stringsWithPrefix("header.converter"));
+        props.putAll(stringsWithPrefix("task"));
+        props.putAll(stringsWithPrefix("worker"));
+ 
+        // transform any expression like ${provider:path:key}, since the worker doesn't do so
+        props = transform(props);
+        props.putAll(stringsWithPrefix(CONFIG_PROVIDERS_CONFIG));
+
+        // fill in reasonable defaults
+        props.putIfAbsent(GROUP_ID_CONFIG, sourceAndTarget.source() + "-mm2");
+        props.putIfAbsent(DistributedConfig.OFFSET_STORAGE_TOPIC_CONFIG, "mm2-offsets."
+                + sourceAndTarget.source() + ".internal");
+        props.putIfAbsent(DistributedConfig.STATUS_STORAGE_TOPIC_CONFIG, "mm2-status."
+                + sourceAndTarget.source() + ".internal");
+        props.putIfAbsent(DistributedConfig.CONFIG_TOPIC_CONFIG, "mm2-configs."
+                + sourceAndTarget.source() + ".internal");
+        props.putIfAbsent(KEY_CONVERTER_CLASS_CONFIG, BYTE_ARRAY_CONVERTER_CLASS); 
+        props.putIfAbsent(VALUE_CONVERTER_CLASS_CONFIG, BYTE_ARRAY_CONVERTER_CLASS); 
+        props.putIfAbsent(HEADER_CONVERTER_CLASS_CONFIG, BYTE_ARRAY_CONVERTER_CLASS);
+
+        return props;
+    }
+
+    // loads properties of the form cluster.x.y.z and source->target.x.y.z
+    Map connectorBaseConfig(SourceAndTarget sourceAndTarget, Class connectorClass) {
+        Map props = new HashMap<>();
+
+        props.putAll(originalsStrings());
+        props.keySet().retainAll(MirrorConnectorConfig.CONNECTOR_CONFIG_DEF.names());
+        
+        props.putAll(stringsWithPrefix(CONFIG_PROVIDERS_CONFIG));
+        
+        props.putAll(withPrefix(SOURCE_CLUSTER_PREFIX, clusterProps(sourceAndTarget.source())));
+        props.putAll(withPrefix(TARGET_CLUSTER_PREFIX, clusterProps(sourceAndTarget.target())));
+
+        props.putIfAbsent(NAME, connectorClass.getSimpleName());
+        props.putIfAbsent(CONNECTOR_CLASS, connectorClass.getName());
+        props.putIfAbsent(SOURCE_CLUSTER_ALIAS, sourceAndTarget.source());
+        props.putIfAbsent(TARGET_CLUSTER_ALIAS, sourceAndTarget.target());
+
+        // override with connector-level properties
+        props.putAll(stringsWithPrefixStripped(sourceAndTarget.source() + "->"
+            + sourceAndTarget.target() + "."));
+
+        // disabled by default
+        props.putIfAbsent(MirrorConnectorConfig.ENABLED, "false");
+
+        // don't transform -- the worker will handle transformation of Connector and Task configs
+        return props;
+    }
+
+    List configProviders() {
+        return getList(CONFIG_PROVIDERS_CONFIG);
+    } 
+
+    Map transform(Map props) {
+        // transform worker config according to config.providers
+        List providerNames = configProviders();
+        Map providers = new HashMap<>();
+        for (String name : providerNames) {
+            ConfigProvider configProvider = plugins.newConfigProvider(
+                    this,
+                    CONFIG_PROVIDERS_CONFIG + "." + name,
+                    Plugins.ClassLoaderUsage.PLUGINS
+            );
+            providers.put(name, configProvider);
+        }
+        ConfigTransformer transformer = new ConfigTransformer(providers);
+        Map transformed = transformer.transform(props).data();
+        providers.values().forEach(x -> Utils.closeQuietly(x, "config provider"));
+        return transformed;
+    }
+ 
+    protected static final ConfigDef CONFIG_DEF = new ConfigDef()
+            .define(CLUSTERS_CONFIG, Type.LIST, Importance.HIGH, CLUSTERS_DOC)
+            .define(CONFIG_PROVIDERS_CONFIG, Type.LIST, Collections.emptyList(), Importance.LOW, CONFIG_PROVIDERS_DOC)
+            // security support
+            .define(CommonClientConfigs.SECURITY_PROTOCOL_CONFIG,
+                Type.STRING,
+                CommonClientConfigs.DEFAULT_SECURITY_PROTOCOL,
+                Importance.MEDIUM,
+                CommonClientConfigs.SECURITY_PROTOCOL_DOC)
+            .withClientSslSupport()
+            .withClientSaslSupport();
+
+    private Map stringsWithPrefixStripped(String prefix) {
+        return originalsStrings().entrySet().stream()
+            .filter(x -> x.getKey().startsWith(prefix))
+            .collect(Collectors.toMap(x -> x.getKey().substring(prefix.length()), x -> x.getValue()));
+    }
+
+    private Map stringsWithPrefix(String prefix) {
+        Map strings = originalsStrings();
+        strings.keySet().removeIf(x -> !x.startsWith(prefix));
+        return strings;
+    } 
+
+    static Map withPrefix(String prefix, Map props) {
+        return props.entrySet().stream()
+            .collect(Collectors.toMap(x -> prefix + x.getKey(), x -> x.getValue()));
+    }
+}
diff --git a/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorMetrics.java b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorMetrics.java
new file mode 100644
index 0000000000000..ea9d2f732bd1d
--- /dev/null
+++ b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorMetrics.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.connect.mirror;
+
+import org.apache.kafka.common.MetricNameTemplate;
+import org.apache.kafka.common.metrics.Metrics;
+import org.apache.kafka.common.metrics.MetricsReporter;
+import org.apache.kafka.common.metrics.Sensor;
+import org.apache.kafka.common.metrics.stats.Value;
+import org.apache.kafka.common.metrics.stats.Min;
+import org.apache.kafka.common.metrics.stats.Max;
+import org.apache.kafka.common.metrics.stats.Avg;
+import org.apache.kafka.common.metrics.stats.Meter;
+import org.apache.kafka.common.TopicPartition;
+
+import java.util.Arrays;
+import java.util.Set;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.HashMap;
+import java.util.LinkedHashMap;
+import java.util.stream.Collectors;
+
+/** Metrics for replicated topic-partitions */
+class MirrorMetrics implements AutoCloseable {
+
+    private static final String SOURCE_CONNECTOR_GROUP = MirrorSourceConnector.class.getSimpleName();
+    private static final String CHECKPOINT_CONNECTOR_GROUP = MirrorCheckpointConnector.class.getSimpleName();
+
+    private static final Set PARTITION_TAGS = new HashSet<>(Arrays.asList("target", "topic", "partition"));
+    private static final Set GROUP_TAGS = new HashSet<>(Arrays.asList("source", "target", "group", "topic", "partition"));
+    
+    private static final MetricNameTemplate RECORD_COUNT = new MetricNameTemplate(
+            "record-count", SOURCE_CONNECTOR_GROUP,
+            "Number of source records replicated to the target cluster.", PARTITION_TAGS);
+    private static final MetricNameTemplate RECORD_RATE = new MetricNameTemplate(
+            "record-rate", SOURCE_CONNECTOR_GROUP,
+            "Average number of source records replicated to the target cluster per second.", PARTITION_TAGS);
+    private static final MetricNameTemplate RECORD_AGE = new MetricNameTemplate(
+            "record-age-ms", SOURCE_CONNECTOR_GROUP,
+            "The age of incoming source records when replicated to the target cluster.", PARTITION_TAGS);
+    private static final MetricNameTemplate RECORD_AGE_MAX = new MetricNameTemplate(
+            "record-age-ms-max", SOURCE_CONNECTOR_GROUP,
+            "The max age of incoming source records when replicated to the target cluster.", PARTITION_TAGS);
+    private static final MetricNameTemplate RECORD_AGE_MIN = new MetricNameTemplate(
+            "record-age-ms-min", SOURCE_CONNECTOR_GROUP,
+            "The min age of incoming source records when replicated to the target cluster.", PARTITION_TAGS);
+    private static final MetricNameTemplate RECORD_AGE_AVG = new MetricNameTemplate(
+            "record-age-ms-avg", SOURCE_CONNECTOR_GROUP,
+            "The average age of incoming source records when replicated to the target cluster.", PARTITION_TAGS);
+    private static final MetricNameTemplate BYTE_COUNT = new MetricNameTemplate(
+            "byte-count", SOURCE_CONNECTOR_GROUP,
+            "Number of bytes replicated to the target cluster.", PARTITION_TAGS);
+    private static final MetricNameTemplate BYTE_RATE = new MetricNameTemplate(
+            "byte-rate", SOURCE_CONNECTOR_GROUP,
+            "Average number of bytes replicated per second.", PARTITION_TAGS);
+    private static final MetricNameTemplate REPLICATION_LATENCY = new MetricNameTemplate(
+            "replication-latency-ms", SOURCE_CONNECTOR_GROUP,
+            "Time it takes records to replicate from source to target cluster.", PARTITION_TAGS);
+    private static final MetricNameTemplate REPLICATION_LATENCY_MAX = new MetricNameTemplate(
+            "replication-latency-ms-max", SOURCE_CONNECTOR_GROUP,
+            "Max time it takes records to replicate from source to target cluster.", PARTITION_TAGS);
+    private static final MetricNameTemplate REPLICATION_LATENCY_MIN = new MetricNameTemplate(
+            "replication-latency-ms-min", SOURCE_CONNECTOR_GROUP,
+            "Min time it takes records to replicate from source to target cluster.", PARTITION_TAGS);
+    private static final MetricNameTemplate REPLICATION_LATENCY_AVG = new MetricNameTemplate(
+            "replication-latency-ms-avg", SOURCE_CONNECTOR_GROUP,
+            "Average time it takes records to replicate from source to target cluster.", PARTITION_TAGS);
+
+    private static final MetricNameTemplate CHECKPOINT_LATENCY = new MetricNameTemplate(
+            "checkpoint-latency-ms", CHECKPOINT_CONNECTOR_GROUP,
+            "Time it takes consumer group offsets to replicate from source to target cluster.", GROUP_TAGS);
+    private static final MetricNameTemplate CHECKPOINT_LATENCY_MAX = new MetricNameTemplate(
+            "checkpoint-latency-ms-max", CHECKPOINT_CONNECTOR_GROUP,
+            "Max time it takes consumer group offsets to replicate from source to target cluster.", GROUP_TAGS);
+    private static final MetricNameTemplate CHECKPOINT_LATENCY_MIN = new MetricNameTemplate(
+            "checkpoint-latency-ms-min", CHECKPOINT_CONNECTOR_GROUP,
+            "Min time it takes consumer group offsets to replicate from source to target cluster.", GROUP_TAGS);
+    private static final MetricNameTemplate CHECKPOINT_LATENCY_AVG = new MetricNameTemplate(
+            "checkpoint-latency-ms-avg", CHECKPOINT_CONNECTOR_GROUP,
+            "Average time it takes consumer group offsets to replicate from source to target cluster.", GROUP_TAGS);
+
+
+    private final Metrics metrics; 
+    private final Map partitionMetrics; 
+    private final Map groupMetrics = new HashMap<>();
+    private final String source;
+    private final String target;
+    private final Set groups;
+
+    MirrorMetrics(MirrorTaskConfig taskConfig) {
+        this.target = taskConfig.targetClusterAlias();
+        this.source = taskConfig.sourceClusterAlias();
+        this.groups = taskConfig.taskConsumerGroups();
+        this.metrics = new Metrics();
+
+        // for side-effect
+        metrics.sensor("record-count");
+        metrics.sensor("byte-rate");
+        metrics.sensor("record-age");
+        metrics.sensor("replication-latency");
+
+        ReplicationPolicy replicationPolicy = taskConfig.replicationPolicy();
+        partitionMetrics = taskConfig.taskTopicPartitions().stream()
+            .map(x -> new TopicPartition(replicationPolicy.formatRemoteTopic(source, x.topic()), x.partition()))
+            .collect(Collectors.toMap(x -> x, x -> new PartitionMetrics(x)));
+
+    }
+
+    @Override
+    public void close() {
+        metrics.close();
+    }
+
+    void countRecord(TopicPartition topicPartition) {
+        partitionMetrics.get(topicPartition).recordSensor.record();
+    }
+
+    void recordAge(TopicPartition topicPartition, long ageMillis) {
+        partitionMetrics.get(topicPartition).recordAgeSensor.record((double) ageMillis);
+    }
+
+    void replicationLatency(TopicPartition topicPartition, long millis) {
+        partitionMetrics.get(topicPartition).replicationLatencySensor.record((double) millis);
+    }
+
+    void recordBytes(TopicPartition topicPartition, long bytes) {
+        partitionMetrics.get(topicPartition).byteSensor.record((double) bytes);
+    }
+
+    void checkpointLatency(TopicPartition topicPartition, String group, long millis) {
+        group(topicPartition, group).checkpointLatencySensor.record((double) millis);
+    }
+
+    GroupMetrics group(TopicPartition topicPartition, String group) {
+        return groupMetrics.computeIfAbsent(String.join("-", topicPartition.toString(), group),
+            x -> new GroupMetrics(topicPartition, group));
+    }
+
+    void addReporter(MetricsReporter reporter) {
+        metrics.addReporter(reporter);
+    }
+
+    private class PartitionMetrics {
+        private final Sensor recordSensor;
+        private final Sensor byteSensor;
+        private final Sensor recordAgeSensor;
+        private final Sensor replicationLatencySensor;
+
+        PartitionMetrics(TopicPartition topicPartition) {
+            String prefix = topicPartition.topic() + "-" + topicPartition.partition() + "-";
+
+            Map tags = new LinkedHashMap<>();
+            tags.put("target", target); 
+            tags.put("topic", topicPartition.topic());
+            tags.put("partition", Integer.toString(topicPartition.partition()));
+
+            recordSensor = metrics.sensor(prefix + "records-sent");
+            recordSensor.add(new Meter(metrics.metricInstance(RECORD_RATE, tags), metrics.metricInstance(RECORD_COUNT, tags)));
+
+            byteSensor = metrics.sensor(prefix + "bytes-sent");
+            byteSensor.add(new Meter(metrics.metricInstance(BYTE_RATE, tags), metrics.metricInstance(BYTE_COUNT, tags)));
+
+            recordAgeSensor = metrics.sensor(prefix + "record-age");
+            recordAgeSensor.add(metrics.metricInstance(RECORD_AGE, tags), new Value());
+            recordAgeSensor.add(metrics.metricInstance(RECORD_AGE_MAX, tags), new Max());
+            recordAgeSensor.add(metrics.metricInstance(RECORD_AGE_MIN, tags), new Min());
+            recordAgeSensor.add(metrics.metricInstance(RECORD_AGE_AVG, tags), new Avg());
+
+            replicationLatencySensor = metrics.sensor(prefix + "replication-latency");
+            replicationLatencySensor.add(metrics.metricInstance(REPLICATION_LATENCY, tags), new Value());
+            replicationLatencySensor.add(metrics.metricInstance(REPLICATION_LATENCY_MAX, tags), new Max());
+            replicationLatencySensor.add(metrics.metricInstance(REPLICATION_LATENCY_MIN, tags), new Min());
+            replicationLatencySensor.add(metrics.metricInstance(REPLICATION_LATENCY_AVG, tags), new Avg());
+        }
+    }
+
+    private class GroupMetrics {
+        private final Sensor checkpointLatencySensor;
+
+        GroupMetrics(TopicPartition topicPartition, String group) {
+            Map tags = new LinkedHashMap<>();
+            tags.put("source", source); 
+            tags.put("target", target); 
+            tags.put("group", group);
+            tags.put("topic", topicPartition.topic());
+            tags.put("partition", Integer.toString(topicPartition.partition()));
+ 
+            checkpointLatencySensor = metrics.sensor("checkpoint-latency");
+            checkpointLatencySensor.add(metrics.metricInstance(CHECKPOINT_LATENCY, tags), new Value());
+            checkpointLatencySensor.add(metrics.metricInstance(CHECKPOINT_LATENCY_MAX, tags), new Max());
+            checkpointLatencySensor.add(metrics.metricInstance(CHECKPOINT_LATENCY_MIN, tags), new Min());
+            checkpointLatencySensor.add(metrics.metricInstance(CHECKPOINT_LATENCY_AVG, tags), new Avg());
+        }
+    }
+}
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
new file mode 100644
index 0000000000000..081bedc78a0f4
--- /dev/null
+++ b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorSourceConnector.java
@@ -0,0 +1,390 @@
+/*
+ * 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.mirror;
+
+import org.apache.kafka.connect.connector.Task;
+import org.apache.kafka.connect.source.SourceConnector;
+import org.apache.kafka.connect.util.ConnectorUtils;
+import org.apache.kafka.common.config.ConfigDef;
+import org.apache.kafka.common.config.ConfigResource;
+import org.apache.kafka.common.acl.AclBinding;
+import org.apache.kafka.common.acl.AclBindingFilter;
+import org.apache.kafka.common.acl.AccessControlEntry;
+import org.apache.kafka.common.acl.AccessControlEntryFilter;
+import org.apache.kafka.common.acl.AclPermissionType;
+import org.apache.kafka.common.acl.AclOperation;
+import org.apache.kafka.common.resource.ResourceType;
+import org.apache.kafka.common.resource.ResourcePattern;
+import org.apache.kafka.common.resource.ResourcePatternFilter;
+import org.apache.kafka.common.resource.PatternType;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.errors.InvalidPartitionsException;
+import org.apache.kafka.common.utils.Utils;
+import org.apache.kafka.clients.admin.AdminClient;
+import org.apache.kafka.clients.admin.TopicDescription;
+import org.apache.kafka.clients.admin.Config;
+import org.apache.kafka.clients.admin.ConfigEntry;
+import org.apache.kafka.clients.admin.NewPartitions;
+import org.apache.kafka.clients.admin.NewTopic;
+import org.apache.kafka.clients.admin.CreateTopicsOptions;
+
+import java.util.Map;
+import java.util.List;
+import java.util.Set;
+import java.util.HashSet;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+import java.util.concurrent.ExecutionException;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** Replicate data, configuration, and ACLs between clusters.
+ *
+ *  @see MirrorConnectorConfig for supported config properties.
+ */
+public class MirrorSourceConnector extends SourceConnector {
+
+    private static final Logger log = LoggerFactory.getLogger(MirrorSourceConnector.class);
+    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 Scheduler scheduler;
+    private MirrorConnectorConfig config;
+    private SourceAndTarget sourceAndTarget;
+    private String connectorName;
+    private TopicFilter topicFilter;
+    private ConfigPropertyFilter configPropertyFilter;
+    private List knownTopicPartitions = Collections.emptyList();
+    private Set knownTargetTopics = Collections.emptySet();
+    private ReplicationPolicy replicationPolicy;
+    private int replicationFactor;
+    private AdminClient sourceAdminClient;
+    private AdminClient targetAdminClient;
+
+    public MirrorSourceConnector() {
+        // nop
+    }
+
+    // visible for testing
+    MirrorSourceConnector(SourceAndTarget sourceAndTarget, ReplicationPolicy replicationPolicy,
+            TopicFilter topicFilter, ConfigPropertyFilter configPropertyFilter) {
+        this.sourceAndTarget = sourceAndTarget;
+        this.replicationPolicy = replicationPolicy;
+        this.topicFilter = topicFilter;
+        this.configPropertyFilter = configPropertyFilter;
+    } 
+
+    @Override
+    public void start(Map props) {
+        long start = System.currentTimeMillis();
+        config = new MirrorConnectorConfig(props);
+        if (!config.enabled()) {
+            return;
+        }
+        connectorName = config.connectorName();
+        sourceAndTarget = new SourceAndTarget(config.sourceClusterAlias(), config.targetClusterAlias());
+        topicFilter = config.topicFilter();
+        configPropertyFilter = config.configPropertyFilter();
+        replicationPolicy = config.replicationPolicy();
+        replicationFactor = config.replicationFactor();
+        sourceAdminClient = AdminClient.create(config.sourceAdminConfig());
+        targetAdminClient = AdminClient.create(config.targetAdminConfig());
+        scheduler = new Scheduler(MirrorSourceConnector.class, config.adminTimeout());
+        scheduler.execute(this::createOffsetSyncsTopic, "creating upstream offset-syncs topic");
+        scheduler.execute(this::loadTopicPartitions, "loading initial set of topic-partitions");
+        scheduler.execute(this::createTopicPartitions, "creating downstream topic-partitions");
+        scheduler.execute(this::refreshKnownTargetTopics, "refreshing known target topics");
+        scheduler.scheduleRepeating(this::syncTopicAcls, config.syncTopicAclsInterval(), "syncing topic ACLs");
+        scheduler.scheduleRepeating(this::syncTopicConfigs, config.syncTopicConfigsInterval(),
+            "syncing topic configs");
+        scheduler.scheduleRepeatingDelayed(this::refreshTopicPartitions, config.refreshTopicsInterval(),
+            "refreshing topics");
+        log.info("Started {} with {} topic-partitions.", connectorName, knownTopicPartitions.size());
+        log.info("Starting {} took {} ms.", connectorName, System.currentTimeMillis() - start);
+    }
+
+    @Override
+    public void stop() {
+        long start = System.currentTimeMillis();
+        if (!config.enabled()) {
+            return;
+        }
+        Utils.closeQuietly(scheduler, "scheduler");
+        Utils.closeQuietly(topicFilter, "topic filter");
+        Utils.closeQuietly(configPropertyFilter, "config property filter");
+        Utils.closeQuietly(sourceAdminClient, "source admin client");
+        Utils.closeQuietly(targetAdminClient, "target admin client");
+        log.info("Stopping {} took {} ms.", connectorName, System.currentTimeMillis() - start);
+    }
+
+    @Override
+    public Class taskClass() {
+        return MirrorSourceTask.class;
+    }
+
+    // divide topic-partitions among tasks
+    @Override
+    public List> taskConfigs(int maxTasks) {
+        if (!config.enabled() || knownTopicPartitions.isEmpty()) {
+            return Collections.emptyList();
+        }
+        int numTasks = Math.min(maxTasks, knownTopicPartitions.size());
+        return ConnectorUtils.groupPartitions(knownTopicPartitions, numTasks).stream()
+            .map(config::taskConfigForTopicPartitions)
+            .collect(Collectors.toList());
+    }
+
+    @Override
+    public ConfigDef config() {
+        return MirrorConnectorConfig.CONNECTOR_CONFIG_DEF;
+    }
+
+    @Override
+    public String version() {
+        return "1";
+    }
+
+    private List findTopicPartitions()
+            throws InterruptedException, ExecutionException {
+        Set topics = listTopics(sourceAdminClient).stream()
+            .filter(this::shouldReplicateTopic)
+            .collect(Collectors.toSet());
+        return describeTopics(topics).stream()
+            .flatMap(MirrorSourceConnector::expandTopicDescription)
+            .collect(Collectors.toList());
+    }
+
+    private void refreshTopicPartitions()
+            throws InterruptedException, ExecutionException {
+        List topicPartitions = findTopicPartitions();
+        Set newTopicPartitions = new HashSet<>();
+        newTopicPartitions.addAll(topicPartitions);
+        newTopicPartitions.removeAll(knownTopicPartitions);
+        Set deadTopicPartitions = new HashSet<>();
+        deadTopicPartitions.addAll(knownTopicPartitions);
+        deadTopicPartitions.removeAll(topicPartitions);
+        if (!newTopicPartitions.isEmpty() || !deadTopicPartitions.isEmpty()) {
+            log.info("Found {} topic-partitions on {}. {} are new. {} were removed. Previously had {}.",
+                    topicPartitions.size(), sourceAndTarget.source(), newTopicPartitions.size(), 
+                    deadTopicPartitions.size(), knownTopicPartitions.size());
+            log.trace("Found new topic-partitions: {}", newTopicPartitions);
+            knownTopicPartitions = topicPartitions;
+            knownTargetTopics = findExistingTargetTopics(); 
+            createTopicPartitions();
+            context.requestTaskReconfiguration();
+        }
+    }
+
+    private void loadTopicPartitions()
+            throws InterruptedException, ExecutionException {
+        knownTopicPartitions = findTopicPartitions();
+        knownTargetTopics = findExistingTargetTopics(); 
+    }
+
+    private void refreshKnownTargetTopics()
+            throws InterruptedException, ExecutionException {
+        knownTargetTopics = findExistingTargetTopics();
+    }
+
+    private Set findExistingTargetTopics()
+            throws InterruptedException, ExecutionException {
+        return listTopics(targetAdminClient).stream()
+            .filter(x -> sourceAndTarget.source().equals(replicationPolicy.topicSource(x)))
+            .collect(Collectors.toSet());
+    }
+
+    private Set topicsBeingReplicated() {
+        return knownTopicPartitions.stream()
+            .map(x -> x.topic())
+            .distinct()
+            .filter(x -> knownTargetTopics.contains(formatRemoteTopic(x)))
+            .collect(Collectors.toSet());
+    }
+
+    private void syncTopicAcls()
+            throws InterruptedException, ExecutionException {
+        List bindings = listTopicAclBindings().stream()
+            .filter(x -> x.pattern().resourceType() == ResourceType.TOPIC)
+            .filter(x -> x.pattern().patternType() == PatternType.LITERAL)
+            .filter(this::shouldReplicateAcl)
+            .filter(x -> shouldReplicateTopic(x.pattern().name()))
+            .map(this::targetAclBinding)
+            .collect(Collectors.toList());
+        updateTopicAcls(bindings);
+    }
+
+    private void syncTopicConfigs()
+            throws InterruptedException, ExecutionException {
+        Map sourceConfigs = describeTopicConfigs(topicsBeingReplicated());
+        Map targetConfigs = sourceConfigs.entrySet().stream()
+            .collect(Collectors.toMap(x -> formatRemoteTopic(x.getKey()), x -> targetConfig(x.getValue())));
+        updateTopicConfigs(targetConfigs);
+    }
+
+    private void createOffsetSyncsTopic() {
+        MirrorUtils.createSinglePartitionCompactedTopic(config.offsetSyncsTopic(), config.offsetSyncsTopicReplicationFactor(), config.sourceAdminConfig());
+    }
+
+    private void createTopicPartitions()
+            throws InterruptedException, ExecutionException {
+        Map partitionCounts = knownTopicPartitions.stream()
+            .collect(Collectors.groupingBy(x -> x.topic(), Collectors.counting())).entrySet().stream()
+            .collect(Collectors.toMap(x -> formatRemoteTopic(x.getKey()), x -> x.getValue()));
+        List newTopics = partitionCounts.entrySet().stream()
+            .filter(x -> !knownTargetTopics.contains(x.getKey()))
+            .map(x -> new NewTopic(x.getKey(), x.getValue().intValue(), (short) replicationFactor))
+            .collect(Collectors.toList());
+        Map newPartitions = partitionCounts.entrySet().stream()
+            .filter(x -> knownTargetTopics.contains(x.getKey()))
+            .collect(Collectors.toMap(x -> x.getKey(), x -> NewPartitions.increaseTo(x.getValue().intValue())));
+        targetAdminClient.createTopics(newTopics, new CreateTopicsOptions()).values().forEach((k, v) -> v.whenComplete((x, e) -> {
+            if (e != null) {
+                log.warn("Could not create topic {}.", k, e);
+            } else {
+                log.info("Created remote topic {} with {} partitions.", k, partitionCounts.get(k));
+            }
+        }));
+        targetAdminClient.createPartitions(newPartitions).values().forEach((k, v) -> v.whenComplete((x, e) -> {
+            if (e instanceof InvalidPartitionsException) {
+                // swallow, this is normal
+            } else if (e != null) {
+                log.warn("Could not create topic-partitions for {}.", k, e);
+            } else {
+                log.info("Increased size of {} to {} partitions.", k, partitionCounts.get(k));
+            }
+        }));
+    }
+
+    private Set listTopics(AdminClient adminClient)
+            throws InterruptedException, ExecutionException {
+        return adminClient.listTopics().names().get();
+    }
+
+    private Collection listTopicAclBindings()
+            throws InterruptedException, ExecutionException {
+        return sourceAdminClient.describeAcls(ANY_TOPIC_ACL).values().get();
+    }
+
+    private Collection describeTopics(Collection topics)
+            throws InterruptedException, ExecutionException {
+        return sourceAdminClient.describeTopics(topics).all().get().values();
+    }
+
+    @SuppressWarnings("deprecation")
+    // use deprecated alterConfigs API for broker compatibility back to 0.11.0
+    private void updateTopicConfigs(Map topicConfigs)
+            throws InterruptedException, ExecutionException {
+        Map configs = topicConfigs.entrySet().stream()
+            .collect(Collectors.toMap(x ->
+                new ConfigResource(ConfigResource.Type.TOPIC, x.getKey()), x -> x.getValue()));
+        log.trace("Syncing configs for {} topics.", configs.size());
+        targetAdminClient.alterConfigs(configs).values().forEach((k, v) -> v.whenComplete((x, e) -> {
+            if (e != null) {
+                log.warn("Could not alter configuration of topic {}.", k.name(), e);
+            }
+        }));
+    }
+
+    private void updateTopicAcls(List bindings)
+            throws InterruptedException, ExecutionException {
+        log.trace("Syncing {} topic ACL bindings.", bindings.size());
+        targetAdminClient.createAcls(bindings).values().forEach((k, v) -> v.whenComplete((x, e) -> {
+            if (e != null) {
+                log.warn("Could not sync ACL of topic {}.", k.pattern().name(), e);
+            }
+        }));
+    }
+
+    private static Stream expandTopicDescription(TopicDescription description) {
+        String topic = description.name();
+        return description.partitions().stream()
+            .map(x -> new TopicPartition(topic, x.partition()));
+    }
+
+    private Map describeTopicConfigs(Set topics)
+            throws InterruptedException, ExecutionException {
+        Set resources = topics.stream()
+            .map(x -> new ConfigResource(ConfigResource.Type.TOPIC, x))
+            .collect(Collectors.toSet());
+        return sourceAdminClient.describeConfigs(resources).all().get().entrySet().stream()
+            .collect(Collectors.toMap(x -> x.getKey().name(), x -> x.getValue()));
+    }
+
+    Config targetConfig(Config sourceConfig) {
+        List entries = sourceConfig.entries().stream()
+            .filter(x -> !x.isDefault() && !x.isReadOnly() && !x.isSensitive())
+            .filter(x -> x.source() != ConfigEntry.ConfigSource.STATIC_BROKER_CONFIG)
+            .filter(x -> shouldReplicateTopicConfigurationProperty(x.name()))
+            .collect(Collectors.toList());
+        return new Config(entries);
+    }
+
+    private static AccessControlEntry downgradeAllowAllACL(AccessControlEntry entry) {
+        return new AccessControlEntry(entry.principal(), entry.host(), AclOperation.READ, entry.permissionType());
+    }
+
+    AclBinding targetAclBinding(AclBinding sourceAclBinding) {
+        String targetTopic = formatRemoteTopic(sourceAclBinding.pattern().name());
+        final AccessControlEntry entry;
+        if (sourceAclBinding.entry().permissionType() == AclPermissionType.ALLOW
+                && sourceAclBinding.entry().operation() == AclOperation.ALL) {
+            entry = downgradeAllowAllACL(sourceAclBinding.entry());
+        } else {
+            entry = sourceAclBinding.entry();
+        }
+        return new AclBinding(new ResourcePattern(ResourceType.TOPIC, targetTopic, PatternType.LITERAL), entry);
+    }
+
+    boolean shouldReplicateTopic(String topic) {
+        return (topicFilter.shouldReplicateTopic(topic) || isHeartbeatTopic(topic))
+            && !replicationPolicy.isInternalTopic(topic) && !isCycle(topic);
+    }
+
+    boolean shouldReplicateAcl(AclBinding aclBinding) {
+        return !(aclBinding.entry().permissionType() == AclPermissionType.ALLOW
+            && aclBinding.entry().operation() == AclOperation.WRITE);
+    }
+
+    boolean shouldReplicateTopicConfigurationProperty(String property) {
+        return configPropertyFilter.shouldReplicateConfigProperty(property);
+    }
+
+    // Recurse upstream to detect cycles, i.e. whether this topic is already on the target cluster
+    boolean isCycle(String topic) {
+        String source = replicationPolicy.topicSource(topic);
+        if (source == null) {
+            return false;
+        } else if (source.equals(sourceAndTarget.target())) {
+            return true;
+        } else {
+            return isCycle(replicationPolicy.upstreamTopic(topic));
+        }
+    }
+
+    // e.g. heartbeats, us-west.heartbeats
+    boolean isHeartbeatTopic(String topic) {
+        return MirrorClientConfig.HEARTBEATS_TOPIC.equals(replicationPolicy.originalTopic(topic));
+    }
+
+    String formatRemoteTopic(String topic) {
+        return replicationPolicy.formatRemoteTopic(sourceAndTarget.source(), topic);
+    }
+}
diff --git a/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorSourceTask.java b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorSourceTask.java
new file mode 100644
index 0000000000000..0b864768c9f7d
--- /dev/null
+++ b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorSourceTask.java
@@ -0,0 +1,293 @@
+/*
+ * 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.mirror;
+
+import org.apache.kafka.connect.data.Schema;
+import org.apache.kafka.connect.source.SourceTask;
+import org.apache.kafka.connect.source.SourceRecord;
+import org.apache.kafka.connect.header.Headers;
+import org.apache.kafka.connect.header.ConnectHeaders;
+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.RecordMetadata;
+import org.apache.kafka.clients.producer.ProducerRecord;
+import org.apache.kafka.clients.producer.KafkaProducer;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.header.Header;
+import org.apache.kafka.common.errors.WakeupException;
+import org.apache.kafka.common.KafkaException;
+import org.apache.kafka.common.utils.Utils;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.Map;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Set;
+import java.util.ArrayList;
+import java.util.stream.Collectors;
+import java.util.concurrent.Semaphore;
+import java.time.Duration;
+
+/** Replicates a set of topic-partitions. */
+public class MirrorSourceTask extends SourceTask {
+
+    private static final Logger log = LoggerFactory.getLogger(MirrorSourceTask.class);
+
+    private static final int MAX_OUTSTANDING_OFFSET_SYNCS = 10;
+
+    private KafkaConsumer consumer;
+    private KafkaProducer offsetProducer;
+    private String sourceClusterAlias;
+    private String offsetSyncsTopic;
+    private Duration pollTimeout;
+    private long maxOffsetLag;
+    private Map partitionStates;
+    private ReplicationPolicy replicationPolicy;
+    private MirrorMetrics metrics;
+    private boolean stopping = false;
+    private Semaphore outstandingOffsetSyncs;
+    private Semaphore consumerAccess;
+
+    public MirrorSourceTask() {}
+
+    // for testing
+    MirrorSourceTask(String sourceClusterAlias, ReplicationPolicy replicationPolicy, long maxOffsetLag) {
+        this.sourceClusterAlias = sourceClusterAlias;
+        this.replicationPolicy = replicationPolicy;
+        this.maxOffsetLag = maxOffsetLag;
+    }
+
+    @Override
+    public void start(Map props) {
+        MirrorTaskConfig config = new MirrorTaskConfig(props);
+        outstandingOffsetSyncs = new Semaphore(MAX_OUTSTANDING_OFFSET_SYNCS);
+        consumerAccess = new Semaphore(1);  // let one thread at a time access the consumer
+        sourceClusterAlias = config.sourceClusterAlias();
+        metrics = config.metrics();
+        pollTimeout = config.consumerPollTimeout();
+        maxOffsetLag = config.maxOffsetLag();
+        replicationPolicy = config.replicationPolicy();
+        partitionStates = new HashMap<>();
+        offsetSyncsTopic = config.offsetSyncsTopic();
+        consumer = MirrorUtils.newConsumer(config.sourceConsumerConfig());
+        offsetProducer = MirrorUtils.newProducer(config.sourceProducerConfig());
+        Set taskTopicPartitions = config.taskTopicPartitions();
+        Map topicPartitionOffsets = loadOffsets(taskTopicPartitions);
+        consumer.assign(topicPartitionOffsets.keySet());
+        log.info("Starting with {} previously uncommitted partitions.", topicPartitionOffsets.entrySet().stream()
+            .filter(x -> x.getValue() == 0L).count());
+        log.trace("Seeking offsets: {}", topicPartitionOffsets);
+        topicPartitionOffsets.forEach(consumer::seek);
+        log.info("{} replicating {} topic-partitions {}->{}: {}.", Thread.currentThread().getName(),
+            taskTopicPartitions.size(), sourceClusterAlias, config.targetClusterAlias(), taskTopicPartitions);
+    }
+
+    @Override
+    public void commit() {
+        // nop
+    }
+
+    @Override
+    public void stop() {
+        long start = System.currentTimeMillis();
+        stopping = true;
+        consumer.wakeup();
+        try {
+            consumerAccess.acquire();
+        } catch (InterruptedException e) {
+            log.warn("Interrupted waiting for access to consumer. Will try closing anyway."); 
+        }
+        Utils.closeQuietly(consumer, "source consumer");
+        Utils.closeQuietly(offsetProducer, "offset producer");
+        Utils.closeQuietly(metrics, "metrics");
+        log.info("Stopping {} took {} ms.", Thread.currentThread().getName(), System.currentTimeMillis() - start);
+    }
+   
+    @Override
+    public String version() {
+        return "1";
+    }
+
+    @Override
+    public List poll() {
+        if (!consumerAccess.tryAcquire()) {
+            return null;
+        }
+        if (stopping) {
+            return null;
+        }
+        try {
+            ConsumerRecords records = consumer.poll(pollTimeout);
+            List sourceRecords = new ArrayList<>(records.count());
+            for (ConsumerRecord record : records) {
+                SourceRecord converted = convertRecord(record);
+                sourceRecords.add(converted);
+                TopicPartition topicPartition = new TopicPartition(converted.topic(), converted.kafkaPartition());
+                metrics.recordAge(topicPartition, System.currentTimeMillis() - record.timestamp());
+                metrics.recordBytes(topicPartition, byteSize(record.value()));
+            }
+            if (sourceRecords.isEmpty()) {
+                // WorkerSourceTasks expects non-zero batch size
+                return null;
+            } else {
+                log.trace("Polled {} records from {}.", sourceRecords.size(), records.partitions());
+                return sourceRecords;
+            }
+        } catch (WakeupException e) {
+            return null;
+        } catch (KafkaException e) {
+            log.warn("Failure during poll.", e);
+            return null;
+        } catch (Throwable e)  {
+            log.error("Failure during poll.", e);
+            // allow Connect to deal with the exception
+            throw e;
+        } finally {
+            consumerAccess.release();
+        }
+    }
+ 
+    @Override
+    public void commitRecord(SourceRecord record, RecordMetadata metadata) {
+        try {
+            if (stopping) {
+                return;
+            }
+            if (!metadata.hasOffset()) {
+                log.error("RecordMetadata has no offset -- can't sync offsets for {}.", record.topic());
+                return;
+            }
+            TopicPartition topicPartition = new TopicPartition(record.topic(), record.kafkaPartition());
+            long latency = System.currentTimeMillis() - record.timestamp();
+            metrics.countRecord(topicPartition);
+            metrics.replicationLatency(topicPartition, latency);
+            TopicPartition sourceTopicPartition = MirrorUtils.unwrapPartition(record.sourcePartition());
+            long upstreamOffset = MirrorUtils.unwrapOffset(record.sourceOffset());
+            long downstreamOffset = metadata.offset();
+            maybeSyncOffsets(sourceTopicPartition, upstreamOffset, downstreamOffset);
+        } catch (Throwable e) {
+            log.warn("Failure committing record.", e);
+        }
+    }
+
+    // updates partition state and sends OffsetSync if necessary
+    private void maybeSyncOffsets(TopicPartition topicPartition, long upstreamOffset,
+            long downstreamOffset) {
+        PartitionState partitionState =
+            partitionStates.computeIfAbsent(topicPartition, x -> new PartitionState(maxOffsetLag));
+        if (partitionState.update(upstreamOffset, downstreamOffset)) {
+            sendOffsetSync(topicPartition, upstreamOffset, downstreamOffset);
+        }
+    }
+
+    // sends OffsetSync record upstream to internal offsets topic
+    private void sendOffsetSync(TopicPartition topicPartition, long upstreamOffset,
+            long downstreamOffset) {
+        if (!outstandingOffsetSyncs.tryAcquire()) {
+            // Too many outstanding offset syncs.
+            return;
+        }
+        OffsetSync offsetSync = new OffsetSync(topicPartition, upstreamOffset, downstreamOffset);
+        ProducerRecord record = new ProducerRecord<>(offsetSyncsTopic, 0,
+                offsetSync.recordKey(), offsetSync.recordValue());
+        offsetProducer.send(record, (x, e) -> {
+            if (e != null) {
+                log.error("Failure sending offset sync.", e);
+            } else {
+                log.trace("Sync'd offsets for {}: {}=={}", topicPartition,
+                    upstreamOffset, downstreamOffset);
+            }
+            outstandingOffsetSyncs.release();
+        });
+    }
+ 
+    private Map loadOffsets(Set topicPartitions) {
+        return topicPartitions.stream().collect(Collectors.toMap(x -> x, x -> loadOffset(x)));
+    }
+
+    private Long loadOffset(TopicPartition topicPartition) {
+        Map wrappedPartition = MirrorUtils.wrapPartition(topicPartition, sourceClusterAlias);
+        Map wrappedOffset = context.offsetStorageReader().offset(wrappedPartition);
+        return MirrorUtils.unwrapOffset(wrappedOffset) + 1;
+    }
+
+    // visible for testing 
+    SourceRecord convertRecord(ConsumerRecord record) {
+        String targetTopic = formatRemoteTopic(record.topic());
+        Headers headers = convertHeaders(record);
+        return new SourceRecord(
+                MirrorUtils.wrapPartition(new TopicPartition(record.topic(), record.partition()), sourceClusterAlias),
+                MirrorUtils.wrapOffset(record.offset()),
+                targetTopic, record.partition(),
+                Schema.OPTIONAL_BYTES_SCHEMA, record.key(),
+                Schema.BYTES_SCHEMA, record.value(),
+                record.timestamp(), headers);
+    }
+
+    private Headers convertHeaders(ConsumerRecord record) {
+        ConnectHeaders headers = new ConnectHeaders();
+        for (Header header : record.headers()) {
+            headers.addBytes(header.key(), header.value());
+        }
+        return headers;
+    }
+
+    private String formatRemoteTopic(String topic) {
+        return replicationPolicy.formatRemoteTopic(sourceClusterAlias, topic);
+    }
+
+    private static int byteSize(byte[] bytes) {
+        if (bytes == null) {
+            return 0;
+        } else {
+            return bytes.length;
+        }
+    }
+
+    static class PartitionState {
+        long previousUpstreamOffset = -1L;
+        long previousDownstreamOffset = -1L;
+        long lastSyncUpstreamOffset = -1L;
+        long lastSyncDownstreamOffset = -1L;
+        long maxOffsetLag;
+
+        PartitionState(long maxOffsetLag) {
+            this.maxOffsetLag = maxOffsetLag;
+        }
+
+        // true if we should emit an offset sync
+        boolean update(long upstreamOffset, long downstreamOffset) {
+            boolean shouldSyncOffsets = false;
+            long upstreamStep = upstreamOffset - lastSyncUpstreamOffset;
+            long downstreamTargetOffset = lastSyncDownstreamOffset + upstreamStep;
+            if (lastSyncDownstreamOffset == -1L
+                    || downstreamOffset - downstreamTargetOffset >= maxOffsetLag
+                    || upstreamOffset - previousUpstreamOffset != 1L
+                    || downstreamOffset < previousDownstreamOffset) {
+                lastSyncUpstreamOffset = upstreamOffset;
+                lastSyncDownstreamOffset = downstreamOffset;
+                shouldSyncOffsets = true;
+            }
+            previousUpstreamOffset = upstreamOffset;
+            previousDownstreamOffset = downstreamOffset;
+            return shouldSyncOffsets;
+        }
+    }
+}
diff --git a/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorTaskConfig.java b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorTaskConfig.java
new file mode 100644
index 0000000000000..73024f5914f4e
--- /dev/null
+++ b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorTaskConfig.java
@@ -0,0 +1,75 @@
+/*
+ * 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.mirror;
+
+import org.apache.kafka.common.config.ConfigDef;
+import org.apache.kafka.common.TopicPartition;
+
+import java.util.Map;
+import java.util.Set;
+import java.util.List;
+import java.util.HashSet;
+import java.util.Collections;
+import java.util.stream.Collectors;
+
+public class MirrorTaskConfig extends MirrorConnectorConfig {
+
+    private static final String TASK_TOPIC_PARTITIONS_DOC = "Topic-partitions assigned to this task to replicate.";
+    private static final String TASK_CONSUMER_GROUPS_DOC = "Consumer groups assigned to this task to replicate.";
+
+    public MirrorTaskConfig(Map props) {
+        super(TASK_CONFIG_DEF, props);
+    }
+
+    Set taskTopicPartitions() {
+        List fields = getList(TASK_TOPIC_PARTITIONS);
+        if (fields == null || fields.isEmpty()) {
+            return Collections.emptySet();
+        }
+        return fields.stream()
+            .map(MirrorUtils::decodeTopicPartition)
+            .collect(Collectors.toSet());
+    }
+
+    Set taskConsumerGroups() {
+        List fields = getList(TASK_CONSUMER_GROUPS);
+        if (fields == null || fields.isEmpty()) {
+            return Collections.emptySet();
+        }
+        return new HashSet<>(fields);
+    } 
+
+    MirrorMetrics metrics() {
+        MirrorMetrics metrics = new MirrorMetrics(this);
+        metricsReporters().forEach(metrics::addReporter);
+        return metrics;
+    }
+ 
+    protected static final ConfigDef TASK_CONFIG_DEF = new ConfigDef(CONNECTOR_CONFIG_DEF)
+        .define(
+            TASK_TOPIC_PARTITIONS,
+            ConfigDef.Type.LIST,
+            null,
+            ConfigDef.Importance.LOW,
+            TASK_TOPIC_PARTITIONS_DOC)
+        .define(
+            TASK_CONSUMER_GROUPS,
+            ConfigDef.Type.LIST,
+            null,
+            ConfigDef.Importance.LOW,
+            TASK_CONSUMER_GROUPS_DOC);
+}
diff --git a/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorUtils.java b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorUtils.java
new file mode 100644
index 0000000000000..f15dda854c7a0
--- /dev/null
+++ b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorUtils.java
@@ -0,0 +1,116 @@
+/*
+ * 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.mirror;
+
+import org.apache.kafka.clients.admin.NewTopic;
+import org.apache.kafka.clients.producer.KafkaProducer;
+import org.apache.kafka.clients.consumer.KafkaConsumer;
+import org.apache.kafka.common.serialization.ByteArrayDeserializer;
+import org.apache.kafka.common.serialization.ByteArraySerializer;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.connect.util.TopicAdmin;
+
+import java.util.Arrays;
+import java.util.Map;
+import java.util.List;
+import java.util.HashMap;
+import java.util.Collections;
+import java.util.regex.Pattern;
+
+/** Internal utility methods. */
+final class MirrorUtils {
+
+    // utility class
+    private MirrorUtils() {}
+
+    static KafkaProducer newProducer(Map props) {
+        return new KafkaProducer<>(props, new ByteArraySerializer(), new ByteArraySerializer());
+    }
+
+    static KafkaConsumer newConsumer(Map props) {
+        return new KafkaConsumer<>(props, new ByteArrayDeserializer(), new ByteArrayDeserializer());
+    }
+
+    static String encodeTopicPartition(TopicPartition topicPartition) {
+        return topicPartition.toString();
+    }
+
+    static Map wrapPartition(TopicPartition topicPartition, String sourceClusterAlias) {
+        Map wrapped = new HashMap<>();
+        wrapped.put("topic", topicPartition.topic());
+        wrapped.put("partition", topicPartition.partition());
+        wrapped.put("cluster", sourceClusterAlias);
+        return wrapped;
+    }
+
+    static Map wrapOffset(long offset) {
+        return Collections.singletonMap("offset", offset);
+    }
+
+    static TopicPartition unwrapPartition(Map wrapped) {
+        String topic = (String) wrapped.get("topic");
+        int partition = (Integer) wrapped.get("partition");
+        return new TopicPartition(topic, partition);
+    }
+
+    static Long unwrapOffset(Map wrapped) {
+        if (wrapped == null || wrapped.get("offset") == null) {
+            return -1L;
+        }
+        return (Long) wrapped.get("offset");
+    }
+
+    static TopicPartition decodeTopicPartition(String topicPartitionString) {
+        int sep = topicPartitionString.lastIndexOf('-');
+        String topic = topicPartitionString.substring(0, sep);
+        String partitionString = topicPartitionString.substring(sep + 1);
+        int partition = Integer.parseInt(partitionString);
+        return new TopicPartition(topic, partition);
+    }
+
+    // returns null if given empty list
+    static Pattern compilePatternList(List fields) {
+        if (fields.isEmpty()) {
+            // The empty pattern matches _everything_, but a blank
+            // config property should match _nothing_.
+            return null;
+        } else {
+            String joined = String.join("|", fields);
+            return Pattern.compile(joined);
+        }
+    }
+
+    static Pattern compilePatternList(String fields) {
+        return compilePatternList(Arrays.asList(fields.split("\\W*,\\W*")));
+    }
+
+    static void createCompactedTopic(String topicName, short partitions, short replicationFactor, Map adminProps) {
+        NewTopic topicDescription = TopicAdmin.defineTopic(topicName).
+                compacted().
+                partitions(partitions).
+                replicationFactor(replicationFactor).
+                build();
+
+        try (TopicAdmin admin = new TopicAdmin(adminProps)) {
+            admin.createTopics(topicDescription);
+        }
+    }
+
+    static void createSinglePartitionCompactedTopic(String topicName, short replicationFactor, Map adminProps) {
+        createCompactedTopic(topicName, (short) 1, replicationFactor, adminProps);
+    }
+}
diff --git a/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/OffsetSync.java b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/OffsetSync.java
new file mode 100644
index 0000000000000..abdc64cc3c754
--- /dev/null
+++ b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/OffsetSync.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.connect.mirror;
+
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.protocol.types.Field;
+import org.apache.kafka.common.protocol.types.Schema;
+import org.apache.kafka.common.protocol.types.Struct;
+import org.apache.kafka.common.protocol.types.Type;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+
+import java.nio.ByteBuffer;
+
+public class OffsetSync {
+    public static final String TOPIC_KEY = "topic";
+    public static final String PARTITION_KEY = "partition";
+    public static final String UPSTREAM_OFFSET_KEY = "upstreamOffset";
+    public static final String DOWNSTREAM_OFFSET_KEY = "offset";
+
+    public static final Schema VALUE_SCHEMA = new Schema(
+            new Field(UPSTREAM_OFFSET_KEY, Type.INT64),
+            new Field(DOWNSTREAM_OFFSET_KEY, Type.INT64));
+
+    public static final Schema KEY_SCHEMA = new Schema(
+            new Field(TOPIC_KEY, Type.STRING),
+            new Field(PARTITION_KEY, Type.INT32));
+
+    private TopicPartition topicPartition;
+    private long upstreamOffset;
+    private long downstreamOffset;
+
+    public OffsetSync(TopicPartition topicPartition, long upstreamOffset, long downstreamOffset) {
+        this.topicPartition = topicPartition;
+        this.upstreamOffset = upstreamOffset;
+        this.downstreamOffset = downstreamOffset;
+    }
+
+    public TopicPartition topicPartition() {
+        return topicPartition;
+    }
+
+    public long upstreamOffset() {
+        return upstreamOffset;
+    }
+
+    public long downstreamOffset() {
+        return downstreamOffset;
+    }
+
+    @Override
+    public String toString() {
+        return String.format("OffsetSync{topicPartition=%s, upstreamOffset=%d, downstreamOffset=%d}",
+            topicPartition, upstreamOffset, downstreamOffset);
+    }
+
+    ByteBuffer serializeValue() {
+        Struct struct = valueStruct();
+        ByteBuffer buffer = ByteBuffer.allocate(VALUE_SCHEMA.sizeOf(struct));
+        VALUE_SCHEMA.write(buffer, struct);
+        buffer.flip();
+        return buffer;
+    }
+
+    ByteBuffer serializeKey() {
+        Struct struct = keyStruct();
+        ByteBuffer buffer = ByteBuffer.allocate(KEY_SCHEMA.sizeOf(struct));
+        KEY_SCHEMA.write(buffer, struct);
+        buffer.flip();
+        return buffer;
+    }
+
+    static OffsetSync deserializeRecord(ConsumerRecord record) {
+        Struct keyStruct = KEY_SCHEMA.read(ByteBuffer.wrap(record.key()));
+        String topic = keyStruct.getString(TOPIC_KEY);
+        int partition = keyStruct.getInt(PARTITION_KEY);
+        
+        Struct valueStruct = VALUE_SCHEMA.read(ByteBuffer.wrap(record.value()));
+        long upstreamOffset = valueStruct.getLong(UPSTREAM_OFFSET_KEY);
+        long downstreamOffset = valueStruct.getLong(DOWNSTREAM_OFFSET_KEY);
+        
+        return new OffsetSync(new TopicPartition(topic, partition), upstreamOffset, downstreamOffset);
+    } 
+
+    private Struct valueStruct() {
+        Struct struct = new Struct(VALUE_SCHEMA);
+        struct.set(UPSTREAM_OFFSET_KEY, upstreamOffset);
+        struct.set(DOWNSTREAM_OFFSET_KEY, downstreamOffset);
+        return struct;
+    }
+
+    private Struct keyStruct() {
+        Struct struct = new Struct(KEY_SCHEMA);
+        struct.set(TOPIC_KEY, topicPartition.topic());
+        struct.set(PARTITION_KEY, topicPartition.partition());
+        return struct;
+    }
+
+    byte[] recordKey() {
+        return serializeKey().array();
+    }
+
+    byte[] recordValue() {
+        return serializeValue().array();
+    }
+};
+
diff --git a/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/OffsetSyncStore.java b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/OffsetSyncStore.java
new file mode 100644
index 0000000000000..fff1abd1cf080
--- /dev/null
+++ b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/OffsetSyncStore.java
@@ -0,0 +1,84 @@
+/*
+ * 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.mirror;
+
+import org.apache.kafka.clients.consumer.KafkaConsumer;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.errors.WakeupException;
+import org.apache.kafka.common.serialization.ByteArrayDeserializer;
+import org.apache.kafka.common.utils.Utils;
+
+import java.util.Map;
+import java.util.HashMap;
+import java.util.Collections;
+import java.time.Duration;
+
+/** Used internally by MirrorMaker. Stores offset syncs and performs offset translation. */
+class OffsetSyncStore implements AutoCloseable {
+    private KafkaConsumer consumer;
+    private Map offsetSyncs = new HashMap<>();
+    private TopicPartition offsetSyncTopicPartition;
+
+    OffsetSyncStore(MirrorConnectorConfig config) {
+        consumer = new KafkaConsumer<>(config.sourceConsumerConfig(),
+            new ByteArrayDeserializer(), new ByteArrayDeserializer());
+        offsetSyncTopicPartition = new TopicPartition(config.offsetSyncsTopic(), 0);
+        consumer.assign(Collections.singleton(offsetSyncTopicPartition));
+    }
+
+    // for testing
+    OffsetSyncStore(KafkaConsumer consumer, TopicPartition offsetSyncTopicPartition) {
+        this.consumer = consumer;
+        this.offsetSyncTopicPartition = offsetSyncTopicPartition;
+    }
+
+    long translateDownstream(TopicPartition sourceTopicPartition, long upstreamOffset) {
+        OffsetSync offsetSync = latestOffsetSync(sourceTopicPartition);
+        if (offsetSync.upstreamOffset() > upstreamOffset) {
+            // Offset is too far in the past to translate accurately
+            return -1;
+        }
+        long upstreamStep = upstreamOffset - offsetSync.upstreamOffset();
+        return offsetSync.downstreamOffset() + upstreamStep;
+    }
+
+    // poll and handle records
+    synchronized void update(Duration pollTimeout) {
+        try {
+            consumer.poll(pollTimeout).forEach(this::handleRecord);
+        } catch (WakeupException e) {
+            // swallow
+        }
+    }
+
+    public synchronized void close() {
+        consumer.wakeup();
+        Utils.closeQuietly(consumer, "offset sync store consumer");
+    }
+
+    protected void handleRecord(ConsumerRecord record) {
+        OffsetSync offsetSync = OffsetSync.deserializeRecord(record);
+        TopicPartition sourceTopicPartition = offsetSync.topicPartition();
+        offsetSyncs.put(sourceTopicPartition, offsetSync);
+    }
+
+    private OffsetSync latestOffsetSync(TopicPartition topicPartition) {
+        return offsetSyncs.computeIfAbsent(topicPartition, x -> new OffsetSync(topicPartition,
+            -1, -1));
+    }
+}
diff --git a/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/Scheduler.java b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/Scheduler.java
new file mode 100644
index 0000000000000..cac9a80c0fd89
--- /dev/null
+++ b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/Scheduler.java
@@ -0,0 +1,115 @@
+/*
+ * 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.mirror;
+
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeoutException;
+import java.time.Duration;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+class Scheduler implements AutoCloseable {
+    private static Logger log = LoggerFactory.getLogger(Scheduler.class);
+
+    private final String name;
+    private final ScheduledExecutorService executor = Executors.newSingleThreadScheduledExecutor();
+    private final Duration timeout;
+    private boolean closed = false;
+
+    Scheduler(String name, Duration timeout) {
+        this.name = name;
+        this.timeout = timeout;
+    }
+
+    Scheduler(Class clazz, Duration timeout) {
+        this("Scheduler for " + clazz.getSimpleName(), timeout);
+    }
+
+    void scheduleRepeating(Task task, Duration interval, String description) {
+        if (interval.toMillis() < 0L) {
+            return;
+        }
+        executor.scheduleAtFixedRate(() -> executeThread(task, description), 0, interval.toMillis(), TimeUnit.MILLISECONDS);
+    }
+ 
+    void scheduleRepeatingDelayed(Task task, Duration interval, String description) {
+        if (interval.toMillis() < 0L) {
+            return;
+        }
+        executor.scheduleAtFixedRate(() -> executeThread(task, description), interval.toMillis(),
+            interval.toMillis(), TimeUnit.MILLISECONDS);
+    }
+
+    void execute(Task task, String description) {
+        try {
+            executor.submit(() -> executeThread(task, description)).get(timeout.toMillis(), TimeUnit.MILLISECONDS);
+        } catch (InterruptedException e) {
+            log.warn("{} was interrupted running task: {}", name, description);
+        } catch (TimeoutException e) {
+            log.error("{} timed out running task: {}", name, description);
+        } catch (Throwable e) {
+            log.error("{} caught exception in task: {}", name, description, e);
+        }
+    } 
+
+    public void close() {
+        closed = true;
+        executor.shutdown();
+        try {
+            boolean terminated = executor.awaitTermination(timeout.toMillis(), TimeUnit.MILLISECONDS);
+            if (!terminated) {
+                log.error("{} timed out during shutdown of internal scheduler.", name);
+            }
+        } catch (InterruptedException e) {
+            log.warn("{} was interrupted during shutdown of internal scheduler.", name);
+        }
+    }
+
+    interface Task {
+        void run() throws InterruptedException, ExecutionException;
+    }
+
+    private void run(Task task, String description) {
+        try {
+            long start = System.currentTimeMillis();
+            task.run();
+            long elapsed = System.currentTimeMillis() - start;
+            log.info("{} took {} ms", description, elapsed);
+            if (elapsed > timeout.toMillis()) {
+                log.warn("{} took too long ({} ms) running task: {}", name, elapsed, description);
+            }
+        } catch (InterruptedException e) {
+            log.warn("{} was interrupted running task: {}", name, description);
+        } catch (Throwable e) {
+            log.error("{} caught exception in scheduled task: {}", name, description, e);
+        }
+    }
+
+    private void executeThread(Task task, String description) {
+        Thread.currentThread().setName(description);
+        if (closed) {
+            log.info("{} skipping task due to shutdown: {}", name, description);
+            return;
+        }
+        run(task, description);
+    }
+}
+
diff --git a/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/TopicFilter.java b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/TopicFilter.java
new file mode 100644
index 0000000000000..f13453f116850
--- /dev/null
+++ b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/TopicFilter.java
@@ -0,0 +1,37 @@
+/*
+ * 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.mirror;
+
+import org.apache.kafka.common.Configurable;
+import org.apache.kafka.common.annotation.InterfaceStability;
+import java.util.Map;
+
+/** Defines which topics should be replicated. */
+@InterfaceStability.Evolving
+public interface TopicFilter extends Configurable, AutoCloseable {
+
+    boolean shouldReplicateTopic(String topic);
+
+    default void close() {
+        //nop
+    }
+
+    default void configure(Map props) {
+        //nop
+    }
+}
diff --git a/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/CheckpointTest.java b/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/CheckpointTest.java
new file mode 100644
index 0000000000000..1a3f21011817a
--- /dev/null
+++ b/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/CheckpointTest.java
@@ -0,0 +1,40 @@
+/*
+ * 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.mirror;
+
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.common.TopicPartition;
+
+import org.junit.Test;
+
+import static org.junit.Assert.assertEquals;
+
+public class CheckpointTest {
+
+    @Test
+    public void testSerde() {
+        Checkpoint checkpoint = new Checkpoint("group-1", new TopicPartition("topic-2", 3), 4, 5, "metadata-6");
+        byte[] key = checkpoint.recordKey();
+        byte[] value = checkpoint.recordValue();
+        ConsumerRecord record = new ConsumerRecord<>("any-topic", 7, 8, key, value);
+        Checkpoint deserialized = Checkpoint.deserializeRecord(record);
+        assertEquals(checkpoint.consumerGroupId(), deserialized.consumerGroupId());
+        assertEquals(checkpoint.topicPartition(), deserialized.topicPartition());
+        assertEquals(checkpoint.upstreamOffset(), deserialized.upstreamOffset());
+        assertEquals(checkpoint.downstreamOffset(), deserialized.downstreamOffset());
+    }
+}
diff --git a/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/HeartbeatTest.java b/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/HeartbeatTest.java
new file mode 100644
index 0000000000000..adc657862a9e4
--- /dev/null
+++ b/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/HeartbeatTest.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.connect.mirror;
+
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+
+import org.junit.Test;
+
+import static org.junit.Assert.assertEquals;
+
+public class HeartbeatTest {
+
+    @Test
+    public void testSerde() {
+        Heartbeat heartbeat = new Heartbeat("source-1", "target-2", 1234567890L);
+        byte[] key = heartbeat.recordKey();
+        byte[] value = heartbeat.recordValue();
+        ConsumerRecord record = new ConsumerRecord<>("any-topic", 6, 7, key, value);
+        Heartbeat deserialized = Heartbeat.deserializeRecord(record);
+        assertEquals(heartbeat.sourceClusterAlias(), deserialized.sourceClusterAlias());
+        assertEquals(heartbeat.targetClusterAlias(), deserialized.targetClusterAlias());
+        assertEquals(heartbeat.timestamp(), deserialized.timestamp());
+    }
+}
diff --git a/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/MirrorCheckpointTaskTest.java b/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/MirrorCheckpointTaskTest.java
new file mode 100644
index 0000000000000..a66574444d65f
--- /dev/null
+++ b/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/MirrorCheckpointTaskTest.java
@@ -0,0 +1,67 @@
+/*
+ * 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.mirror;
+
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.clients.consumer.OffsetAndMetadata;
+import org.apache.kafka.connect.source.SourceRecord;
+
+import org.junit.Test;
+
+import static org.junit.Assert.assertEquals;
+
+public class MirrorCheckpointTaskTest {
+
+    @Test
+    public void testDownstreamTopicRenaming() {
+        MirrorCheckpointTask mirrorCheckpointTask = new MirrorCheckpointTask("source1", "target2",
+            new DefaultReplicationPolicy(), null);
+        assertEquals(new TopicPartition("source1.topic3", 4),
+            mirrorCheckpointTask.renameTopicPartition(new TopicPartition("topic3", 4)));
+        assertEquals(new TopicPartition("topic3", 5),
+            mirrorCheckpointTask.renameTopicPartition(new TopicPartition("target2.topic3", 5)));
+        assertEquals(new TopicPartition("source1.source6.topic7", 8),
+            mirrorCheckpointTask.renameTopicPartition(new TopicPartition("source6.topic7", 8)));
+    }
+
+    @Test
+    public void testCheckpoint() {
+        OffsetSyncStoreTest.FakeOffsetSyncStore offsetSyncStore = new OffsetSyncStoreTest.FakeOffsetSyncStore();
+        MirrorCheckpointTask mirrorCheckpointTask = new MirrorCheckpointTask("source1", "target2",
+            new DefaultReplicationPolicy(), offsetSyncStore);
+        offsetSyncStore.sync(new TopicPartition("topic1", 2), 3L, 4L);
+        offsetSyncStore.sync(new TopicPartition("target2.topic5", 6), 7L, 8L);
+        Checkpoint checkpoint1 = mirrorCheckpointTask.checkpoint("group9", new TopicPartition("topic1", 2),
+            new OffsetAndMetadata(10, null));
+        SourceRecord sourceRecord1 = mirrorCheckpointTask.checkpointRecord(checkpoint1, 123L);
+        assertEquals(new TopicPartition("source1.topic1", 2), checkpoint1.topicPartition());
+        assertEquals("group9", checkpoint1.consumerGroupId());
+        assertEquals("group9", Checkpoint.unwrapGroup(sourceRecord1.sourcePartition()));
+        assertEquals(10, checkpoint1.upstreamOffset());
+        assertEquals(11, checkpoint1.downstreamOffset());
+        assertEquals(123L, sourceRecord1.timestamp().longValue());
+        Checkpoint checkpoint2 = mirrorCheckpointTask.checkpoint("group11", new TopicPartition("target2.topic5", 6),
+            new OffsetAndMetadata(12, null));
+        SourceRecord sourceRecord2 = mirrorCheckpointTask.checkpointRecord(checkpoint2, 234L);
+        assertEquals(new TopicPartition("topic5", 6), checkpoint2.topicPartition());
+        assertEquals("group11", checkpoint2.consumerGroupId());
+        assertEquals("group11", Checkpoint.unwrapGroup(sourceRecord2.sourcePartition()));
+        assertEquals(12, checkpoint2.upstreamOffset());
+        assertEquals(13, checkpoint2.downstreamOffset());
+        assertEquals(234L, sourceRecord2.timestamp().longValue());
+    }
+}
diff --git a/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/MirrorConnectorConfigTest.java b/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/MirrorConnectorConfigTest.java
new file mode 100644
index 0000000000000..8e99779e8f7c0
--- /dev/null
+++ b/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/MirrorConnectorConfigTest.java
@@ -0,0 +1,134 @@
+/*
+ * 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.mirror;
+
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.config.ConfigDef;
+import org.junit.Test;
+
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.List;
+import java.util.Map;
+import java.util.HashMap;
+import java.util.HashSet;
+
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertEquals;
+
+public class MirrorConnectorConfigTest {
+
+    private Map makeProps(String... keyValues) {
+        Map props = new HashMap<>();
+        props.put("name", "ConnectorName");
+        props.put("connector.class", "ConnectorClass");
+        props.put("source.cluster.alias", "source1");
+        props.put("target.cluster.alias", "target2");
+        for (int i = 0; i < keyValues.length; i += 2) {
+            props.put(keyValues[i], keyValues[i + 1]);
+        }
+        return props;
+    }
+
+    @Test
+    public void testTaskConfigTopicPartitions() {
+        List topicPartitions = Arrays.asList(new TopicPartition("topic-1", 2),
+            new TopicPartition("topic-3", 4), new TopicPartition("topic-5", 6));
+        MirrorConnectorConfig config = new MirrorConnectorConfig(makeProps());
+        Map props = config.taskConfigForTopicPartitions(topicPartitions);
+        MirrorTaskConfig taskConfig = new MirrorTaskConfig(props);
+        assertEquals(taskConfig.taskTopicPartitions(), new HashSet<>(topicPartitions));
+    }
+
+    @Test
+    public void testTaskConfigConsumerGroups() {
+        List groups = Arrays.asList("consumer-1", "consumer-2", "consumer-3");
+        MirrorConnectorConfig config = new MirrorConnectorConfig(makeProps());
+        Map props = config.taskConfigForConsumerGroups(groups);
+        MirrorTaskConfig taskConfig = new MirrorTaskConfig(props);
+        assertEquals(taskConfig.taskConsumerGroups(), new HashSet<>(groups));
+    }
+
+    @Test
+    public void testTopicMatching() {
+        MirrorConnectorConfig config = new MirrorConnectorConfig(makeProps("topics", "topic1"));
+        assertTrue(config.topicFilter().shouldReplicateTopic("topic1"));
+        assertFalse(config.topicFilter().shouldReplicateTopic("topic2"));
+    }
+
+    @Test
+    public void testGroupMatching() {
+        MirrorConnectorConfig config = new MirrorConnectorConfig(makeProps("groups", "group1"));
+        assertTrue(config.groupFilter().shouldReplicateGroup("group1"));
+        assertFalse(config.groupFilter().shouldReplicateGroup("group2"));
+    }
+
+    @Test
+    public void testConfigPropertyMatching() {
+        MirrorConnectorConfig config = new MirrorConnectorConfig(
+            makeProps("config.properties.blacklist", "prop2"));
+        assertTrue(config.configPropertyFilter().shouldReplicateConfigProperty("prop1"));
+        assertFalse(config.configPropertyFilter().shouldReplicateConfigProperty("prop2"));
+    }
+
+    @Test
+    public void testNoTopics() {
+        MirrorConnectorConfig config = new MirrorConnectorConfig(makeProps("topics", ""));
+        assertFalse(config.topicFilter().shouldReplicateTopic("topic1"));
+        assertFalse(config.topicFilter().shouldReplicateTopic("topic2"));
+        assertFalse(config.topicFilter().shouldReplicateTopic(""));
+    }
+
+    @Test
+    public void testAllTopics() {
+        MirrorConnectorConfig config = new MirrorConnectorConfig(makeProps("topics", ".*"));
+        assertTrue(config.topicFilter().shouldReplicateTopic("topic1"));
+        assertTrue(config.topicFilter().shouldReplicateTopic("topic2"));
+    }
+
+    @Test
+    public void testListOfTopics() {
+        MirrorConnectorConfig config = new MirrorConnectorConfig(makeProps("topics", "topic1, topic2"));
+        assertTrue(config.topicFilter().shouldReplicateTopic("topic1"));
+        assertTrue(config.topicFilter().shouldReplicateTopic("topic2"));
+        assertFalse(config.topicFilter().shouldReplicateTopic("topic3"));
+    }
+
+    @Test
+    public void testNonMutationOfConfigDef() {
+        Collection taskSpecificProperties = Arrays.asList(
+            MirrorConnectorConfig.TASK_TOPIC_PARTITIONS,
+            MirrorConnectorConfig.TASK_CONSUMER_GROUPS
+        );
+
+        // Sanity check to make sure that these properties are actually defined for the task config,
+        // and that the task config class has been loaded and statically initialized by the JVM
+        ConfigDef taskConfigDef = MirrorTaskConfig.TASK_CONFIG_DEF;
+        taskSpecificProperties.forEach(taskSpecificProperty -> assertTrue(
+            taskSpecificProperty + " should be defined for task ConfigDef",
+            taskConfigDef.names().contains(taskSpecificProperty)
+        ));
+
+        // Ensure that the task config class hasn't accidentally modified the connector config
+        ConfigDef connectorConfigDef = MirrorConnectorConfig.CONNECTOR_CONFIG_DEF;
+        taskSpecificProperties.forEach(taskSpecificProperty -> assertFalse(
+            taskSpecificProperty + " should not be defined for connector ConfigDef",
+            connectorConfigDef.names().contains(taskSpecificProperty)
+        ));
+    }
+}
diff --git a/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/MirrorConnectorsIntegrationTest.java b/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/MirrorConnectorsIntegrationTest.java
new file mode 100644
index 0000000000000..30924ef84bff4
--- /dev/null
+++ b/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/MirrorConnectorsIntegrationTest.java
@@ -0,0 +1,326 @@
+/*
+ * 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.mirror;
+
+import org.apache.kafka.clients.admin.Admin;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.OffsetAndMetadata;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.connect.util.clusters.EmbeddedConnectCluster;
+import org.apache.kafka.connect.util.clusters.EmbeddedKafkaCluster;
+import org.apache.kafka.test.IntegrationTest;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.experimental.categories.Category;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.time.Duration;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Properties;
+import java.util.Set;
+
+import static org.apache.kafka.test.TestUtils.waitForCondition;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+/**
+ * Tests MM2 replication and failover/failback logic.
+ *
+ * MM2 is configured with active/active replication between two Kafka clusters. Tests validate that
+ * records sent to either cluster arrive at the other cluster. Then, a consumer group is migrated from
+ * one cluster to the other and back. Tests validate that consumer offsets are translated and replicated
+ * between clusters during this failover and failback.
+ */
+@Category(IntegrationTest.class)
+public class MirrorConnectorsIntegrationTest {
+
+    private static final Logger log = LoggerFactory.getLogger(MirrorConnectorsIntegrationTest.class);
+
+    private static final int NUM_RECORDS_PRODUCED = 100;  // to save trees
+    private static final int NUM_PARTITIONS = 10;
+    private static final int RECORD_TRANSFER_DURATION_MS = 10_000;
+    private static final int CHECKPOINT_DURATION_MS = 20_000;
+
+    private MirrorMakerConfig mm2Config; 
+    private EmbeddedConnectCluster primary;
+    private EmbeddedConnectCluster backup;
+
+    @Before
+    public void setup() throws InterruptedException {
+        Properties brokerProps = new Properties();
+        brokerProps.put("auto.create.topics.enable", "false");
+
+        Map mm2Props = new HashMap<>();
+        mm2Props.put("clusters", "primary, backup");
+        mm2Props.put("max.tasks", "10");
+        mm2Props.put("topics", "test-topic-.*, primary.test-topic-.*, backup.test-topic-.*");
+        mm2Props.put("groups", "consumer-group-.*");
+        mm2Props.put("primary->backup.enabled", "true");
+        mm2Props.put("backup->primary.enabled", "true");
+        mm2Props.put("sync.topic.acls.enabled", "false");
+        mm2Props.put("emit.checkpoints.interval.seconds", "1");
+        mm2Props.put("emit.heartbeats.interval.seconds", "1");
+        mm2Props.put("refresh.topics.interval.seconds", "1");
+        mm2Props.put("refresh.groups.interval.seconds", "1");
+        mm2Props.put("checkpoints.topic.replication.factor", "1");
+        mm2Props.put("heartbeats.topic.replication.factor", "1");
+        mm2Props.put("offset-syncs.topic.replication.factor", "1");
+        mm2Props.put("config.storage.replication.factor", "1");
+        mm2Props.put("offset.storage.replication.factor", "1");
+        mm2Props.put("status.storage.replication.factor", "1");
+        mm2Props.put("replication.factor", "1");
+        
+        mm2Config = new MirrorMakerConfig(mm2Props); 
+        Map primaryWorkerProps = mm2Config.workerConfig(new SourceAndTarget("backup", "primary"));
+        Map backupWorkerProps = mm2Config.workerConfig(new SourceAndTarget("primary", "backup"));
+
+        primary = new EmbeddedConnectCluster.Builder()
+                .name("primary-connect-cluster")
+                .numWorkers(3)
+                .numBrokers(1)
+                .brokerProps(brokerProps)
+                .workerProps(primaryWorkerProps)
+                .build();
+
+        backup = new EmbeddedConnectCluster.Builder()
+                .name("backup-connect-cluster")
+                .numWorkers(3)
+                .numBrokers(1)
+                .brokerProps(brokerProps)
+                .workerProps(backupWorkerProps)
+                .build();
+
+        primary.start();
+        primary.assertions().assertAtLeastNumWorkersAreUp(3,
+                "Workers of primary-connect-cluster did not start in time.");
+        backup.start();
+        primary.assertions().assertAtLeastNumWorkersAreUp(3,
+                "Workers of backup-connect-cluster did not start in time.");
+
+        // create these topics before starting the connectors so we don't need to wait for discovery
+        primary.kafka().createTopic("test-topic-1", NUM_PARTITIONS);
+        primary.kafka().createTopic("backup.test-topic-1", 1);
+        primary.kafka().createTopic("heartbeats", 1);
+        backup.kafka().createTopic("test-topic-1", NUM_PARTITIONS);
+        backup.kafka().createTopic("primary.test-topic-1", 1);
+        backup.kafka().createTopic("heartbeats", 1);
+
+        for (int i = 0; i < NUM_RECORDS_PRODUCED; i++) {
+            primary.kafka().produce("test-topic-1", i % NUM_PARTITIONS, "key", "message-1-" + i);
+            backup.kafka().produce("test-topic-1", i % NUM_PARTITIONS, "key", "message-2-" + i);
+        }
+
+        // create consumers before starting the connectors so we don't need to wait for discovery
+        Consumer consumer1 = primary.kafka().createConsumerAndSubscribeTo(Collections.singletonMap(
+            "group.id", "consumer-group-1"), "test-topic-1", "backup.test-topic-1");
+        consumer1.poll(Duration.ofMillis(500));
+        consumer1.commitSync();
+        consumer1.close();
+
+        Consumer consumer2 = backup.kafka().createConsumerAndSubscribeTo(Collections.singletonMap(
+            "group.id", "consumer-group-1"), "test-topic-1", "primary.test-topic-1");
+        consumer2.poll(Duration.ofMillis(500));
+        consumer2.commitSync();
+        consumer2.close();
+
+        log.info("primary REST service: {}", primary.endpointForResource("connectors"));
+        log.info("backup REST service: {}", backup.endpointForResource("connectors"));
+ 
+        log.info("primary brokers: {}", primary.kafka().bootstrapServers());
+        log.info("backup brokers: {}", backup.kafka().bootstrapServers());
+        
+        // now that the brokers are running, we can finish setting up the Connectors
+        mm2Props.put("primary.bootstrap.servers", primary.kafka().bootstrapServers());
+        mm2Props.put("backup.bootstrap.servers", backup.kafka().bootstrapServers());
+        mm2Config = new MirrorMakerConfig(mm2Props);
+
+        // we wait for the connector and tasks to come up for each connector, so that when we do the
+        // actual testing, we are certain that the tasks are up and running; this will prevent
+        // flaky tests where the connector and tasks didn't start up in time for the tests to be
+        // run
+        Set connectorNames = new HashSet<>(Arrays.asList("MirrorSourceConnector",
+                "MirrorCheckpointConnector", "MirrorHeartbeatConnector"));
+
+        backup.configureConnector("MirrorSourceConnector", mm2Config.connectorBaseConfig(new SourceAndTarget("primary", "backup"),
+            MirrorSourceConnector.class));
+
+        backup.configureConnector("MirrorCheckpointConnector", mm2Config.connectorBaseConfig(new SourceAndTarget("primary", "backup"),
+            MirrorCheckpointConnector.class));
+
+        backup.configureConnector("MirrorHeartbeatConnector", mm2Config.connectorBaseConfig(new SourceAndTarget("primary", "backup"),
+            MirrorHeartbeatConnector.class));
+
+        waitUntilMirrorMakerIsRunning(backup, connectorNames);
+
+        primary.configureConnector("MirrorSourceConnector", mm2Config.connectorBaseConfig(new SourceAndTarget("backup", "primary"),
+            MirrorSourceConnector.class));
+
+        primary.configureConnector("MirrorCheckpointConnector", mm2Config.connectorBaseConfig(new SourceAndTarget("backup", "primary"),
+            MirrorCheckpointConnector.class));
+
+        primary.configureConnector("MirrorHeartbeatConnector", mm2Config.connectorBaseConfig(new SourceAndTarget("backup", "primary"),
+            MirrorHeartbeatConnector.class));
+
+        waitUntilMirrorMakerIsRunning(primary, connectorNames);
+    }
+
+
+    private void waitUntilMirrorMakerIsRunning(EmbeddedConnectCluster connectCluster,
+        Set connNames) throws InterruptedException {
+        for (String connector : connNames) {
+            connectCluster.assertions().assertConnectorAndAtLeastNumTasksAreRunning(connector, 1,
+                    "Connector " + connector + " tasks did not start in time on cluster: " + connectCluster);
+        }
+    }
+
+    @After
+    public void close() {
+        for (String x : primary.connectors()) {
+            primary.deleteConnector(x);
+        }
+        for (String x : backup.connectors()) {
+            backup.deleteConnector(x);
+        }
+        deleteAllTopics(primary.kafka());
+        deleteAllTopics(backup.kafka());
+        primary.stop();
+        backup.stop();
+    }
+
+    @Test
+    public void testReplication() throws InterruptedException {
+        MirrorClient primaryClient = new MirrorClient(mm2Config.clientConfig("primary"));
+        MirrorClient backupClient = new MirrorClient(mm2Config.clientConfig("backup"));
+
+        assertEquals("Records were not produced to primary cluster.", NUM_RECORDS_PRODUCED,
+            primary.kafka().consume(NUM_RECORDS_PRODUCED, RECORD_TRANSFER_DURATION_MS, "test-topic-1").count());
+        assertEquals("Records were not replicated to backup cluster.", NUM_RECORDS_PRODUCED,
+            backup.kafka().consume(NUM_RECORDS_PRODUCED, RECORD_TRANSFER_DURATION_MS, "primary.test-topic-1").count());
+        assertEquals("Records were not produced to backup cluster.", NUM_RECORDS_PRODUCED,
+            backup.kafka().consume(NUM_RECORDS_PRODUCED, RECORD_TRANSFER_DURATION_MS, "test-topic-1").count());
+        assertEquals("Records were not replicated to primary cluster.", NUM_RECORDS_PRODUCED,
+            primary.kafka().consume(NUM_RECORDS_PRODUCED, RECORD_TRANSFER_DURATION_MS, "backup.test-topic-1").count());
+        assertEquals("Primary cluster doesn't have all records from both clusters.", NUM_RECORDS_PRODUCED * 2,
+            primary.kafka().consume(NUM_RECORDS_PRODUCED * 2, RECORD_TRANSFER_DURATION_MS, "backup.test-topic-1", "test-topic-1").count());
+        assertEquals("Backup cluster doesn't have all records from both clusters.", NUM_RECORDS_PRODUCED * 2,
+            backup.kafka().consume(NUM_RECORDS_PRODUCED * 2, RECORD_TRANSFER_DURATION_MS, "primary.test-topic-1", "test-topic-1").count());
+        assertTrue("Heartbeats were not emitted to primary cluster.", primary.kafka().consume(1,
+            RECORD_TRANSFER_DURATION_MS, "heartbeats").count() > 0);
+        assertTrue("Heartbeats were not emitted to backup cluster.", backup.kafka().consume(1,
+            RECORD_TRANSFER_DURATION_MS, "heartbeats").count() > 0);
+        assertTrue("Heartbeats were not replicated downstream to backup cluster.", backup.kafka().consume(1,
+            RECORD_TRANSFER_DURATION_MS, "primary.heartbeats").count() > 0);
+        assertTrue("Heartbeats were not replicated downstream to primary cluster.", primary.kafka().consume(1,
+            RECORD_TRANSFER_DURATION_MS, "backup.heartbeats").count() > 0);
+        assertTrue("Did not find upstream primary cluster.", backupClient.upstreamClusters().contains("primary"));
+        assertEquals("Did not calculate replication hops correctly.", 1, backupClient.replicationHops("primary"));
+        assertTrue("Did not find upstream backup cluster.", primaryClient.upstreamClusters().contains("backup"));
+        assertEquals("Did not calculate replication hops correctly.", 1, primaryClient.replicationHops("backup"));
+        assertTrue("Checkpoints were not emitted downstream to backup cluster.", backup.kafka().consume(1,
+            CHECKPOINT_DURATION_MS, "primary.checkpoints.internal").count() > 0);
+
+        Map backupOffsets = backupClient.remoteConsumerOffsets("consumer-group-1", "primary",
+            Duration.ofMillis(CHECKPOINT_DURATION_MS));
+
+        assertTrue("Offsets not translated downstream to backup cluster. Found: " + backupOffsets, backupOffsets.containsKey(
+            new TopicPartition("primary.test-topic-1", 0)));
+
+        // Failover consumer group to backup cluster.
+        Consumer consumer1 = backup.kafka().createConsumer(Collections.singletonMap("group.id", "consumer-group-1"));
+        consumer1.assign(backupOffsets.keySet());
+        backupOffsets.forEach(consumer1::seek);
+        consumer1.poll(Duration.ofMillis(500));
+        consumer1.commitSync();
+
+        assertTrue("Consumer failedover to zero offset.", consumer1.position(new TopicPartition("primary.test-topic-1", 0)) > 0);
+        assertTrue("Consumer failedover beyond expected offset.", consumer1.position(
+            new TopicPartition("primary.test-topic-1", 0)) <= NUM_RECORDS_PRODUCED);
+        assertTrue("Checkpoints were not emitted upstream to primary cluster.", primary.kafka().consume(1,
+            CHECKPOINT_DURATION_MS, "backup.checkpoints.internal").count() > 0);
+
+        consumer1.close();
+
+        waitForCondition(() -> {
+            try {
+                return primaryClient.remoteConsumerOffsets("consumer-group-1", "backup",
+                    Duration.ofMillis(CHECKPOINT_DURATION_MS)).containsKey(new TopicPartition("backup.test-topic-1", 0));
+            } catch (Throwable e) {
+                return false;
+            }
+        }, CHECKPOINT_DURATION_MS, "Offsets not translated downstream to primary cluster.");
+
+        waitForCondition(() -> {
+            try {
+                return primaryClient.remoteConsumerOffsets("consumer-group-1", "backup",
+                    Duration.ofMillis(CHECKPOINT_DURATION_MS)).containsKey(new TopicPartition("test-topic-1", 0));
+            } catch (Throwable e) {
+                return false;
+            }
+        }, CHECKPOINT_DURATION_MS, "Offsets not translated upstream to primary cluster.");
+
+        Map primaryOffsets = primaryClient.remoteConsumerOffsets("consumer-group-1", "backup",
+                Duration.ofMillis(CHECKPOINT_DURATION_MS));
+ 
+        // Failback consumer group to primary cluster
+        Consumer consumer2 = primary.kafka().createConsumer(Collections.singletonMap("group.id", "consumer-group-1"));
+        consumer2.assign(primaryOffsets.keySet());
+        primaryOffsets.forEach(consumer2::seek);
+        consumer2.poll(Duration.ofMillis(500));
+
+        assertTrue("Consumer failedback to zero upstream offset.", consumer2.position(new TopicPartition("test-topic-1", 0)) > 0);
+        assertTrue("Consumer failedback to zero downstream offset.", consumer2.position(new TopicPartition("backup.test-topic-1", 0)) > 0);
+        assertTrue("Consumer failedback beyond expected upstream offset.", consumer2.position(
+            new TopicPartition("test-topic-1", 0)) <= NUM_RECORDS_PRODUCED);
+        assertTrue("Consumer failedback beyond expected downstream offset.", consumer2.position(
+            new TopicPartition("backup.test-topic-1", 0)) <= NUM_RECORDS_PRODUCED);
+        
+        consumer2.close();
+      
+        // create more matching topics
+        primary.kafka().createTopic("test-topic-2", NUM_PARTITIONS);
+        backup.kafka().createTopic("test-topic-3", NUM_PARTITIONS);
+
+        for (int i = 0; i < NUM_RECORDS_PRODUCED; i++) {
+            primary.kafka().produce("test-topic-2", 0, "key", "message-2-" + i);
+            backup.kafka().produce("test-topic-3", 0, "key", "message-3-" + i);
+        }
+
+        assertEquals("Records were not produced to primary cluster.", NUM_RECORDS_PRODUCED,
+            primary.kafka().consume(NUM_RECORDS_PRODUCED, RECORD_TRANSFER_DURATION_MS, "test-topic-2").count());
+        assertEquals("Records were not produced to backup cluster.", NUM_RECORDS_PRODUCED,
+            backup.kafka().consume(NUM_RECORDS_PRODUCED, RECORD_TRANSFER_DURATION_MS, "test-topic-3").count());
+ 
+        assertEquals("New topic was not replicated to primary cluster.", NUM_RECORDS_PRODUCED,
+            primary.kafka().consume(NUM_RECORDS_PRODUCED, 2 * RECORD_TRANSFER_DURATION_MS, "backup.test-topic-3").count());
+        assertEquals("New topic was not replicated to backup cluster.", NUM_RECORDS_PRODUCED,
+            backup.kafka().consume(NUM_RECORDS_PRODUCED, 2 * RECORD_TRANSFER_DURATION_MS, "primary.test-topic-2").count());
+    }
+
+    private void deleteAllTopics(EmbeddedKafkaCluster cluster) {
+        Admin client = cluster.createAdminClient();
+        try {
+            client.deleteTopics(client.listTopics().names().get());
+        } catch (Throwable e) {
+        }
+    }
+}
diff --git a/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/MirrorMakerConfigTest.java b/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/MirrorMakerConfigTest.java
new file mode 100644
index 0000000000000..b618e37c3b3ba
--- /dev/null
+++ b/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/MirrorMakerConfigTest.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.connect.mirror;
+
+import org.apache.kafka.common.config.types.Password;
+import org.apache.kafka.common.config.provider.ConfigProvider;
+import org.apache.kafka.common.config.ConfigData;
+import org.apache.kafka.common.metrics.FakeMetricsReporter;
+
+import org.junit.Test;
+
+import java.util.Map;
+import java.util.Set;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Arrays;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+
+public class MirrorMakerConfigTest {
+
+    private Map makeProps(String... keyValues) {
+        Map props = new HashMap<>();
+        for (int i = 0; i < keyValues.length; i += 2) {
+            props.put(keyValues[i], keyValues[i + 1]);
+        }
+        return props;
+    }
+
+    @Test
+    public void testClusterConfigProperties() {
+        MirrorMakerConfig mirrorConfig = new MirrorMakerConfig(makeProps(
+            "clusters", "a, b",
+            "a.bootstrap.servers", "servers-one",
+            "b.bootstrap.servers", "servers-two",
+            "security.protocol", "SASL",
+            "replication.factor", "4"));
+        Map connectorProps = mirrorConfig.connectorBaseConfig(new SourceAndTarget("a", "b"),
+            MirrorSourceConnector.class);
+        assertEquals("source.cluster.bootstrap.servers is set",
+            "servers-one", connectorProps.get("source.cluster.bootstrap.servers"));
+        assertEquals("target.cluster.bootstrap.servers is set",
+            "servers-two", connectorProps.get("target.cluster.bootstrap.servers"));
+        assertEquals("top-level security.protocol is passed through to connector config",
+            "SASL", connectorProps.get("security.protocol"));
+    }
+
+    @Test
+    public void testReplicationConfigProperties() {
+        MirrorMakerConfig mirrorConfig = new MirrorMakerConfig(makeProps(
+            "clusters", "a, b",
+            "a->b.tasks.max", "123"));
+        Map connectorProps = mirrorConfig.connectorBaseConfig(new SourceAndTarget("a", "b"),
+            MirrorSourceConnector.class);
+        assertEquals("connector props should include tasks.max",
+            "123", connectorProps.get("tasks.max"));
+    }
+
+    @Test
+    public void testClientConfigProperties() {
+        MirrorMakerConfig mirrorConfig = new MirrorMakerConfig(makeProps(
+            "clusters", "a, b",
+            "config.providers", "fake",
+            "config.providers.fake.class", FakeConfigProvider.class.getName(),
+            "replication.policy.separator", "__",
+            "ssl.truststore.password", "secret1",
+            "ssl.key.password", "${fake:secret:password}",  // resolves to "secret2"
+            "security.protocol", "SSL", 
+            "a.security.protocol", "PLAINTEXT", 
+            "a.producer.security.protocol", "SASL", 
+            "a.bootstrap.servers", "one:9092, two:9092",
+            "metrics.reporter", FakeMetricsReporter.class.getName(),
+            "a.metrics.reporter", FakeMetricsReporter.class.getName(),
+            "b->a.metrics.reporter", FakeMetricsReporter.class.getName(),
+            "a.xxx", "yyy",
+            "xxx", "zzz"));
+        MirrorClientConfig aClientConfig = mirrorConfig.clientConfig("a");
+        MirrorClientConfig bClientConfig = mirrorConfig.clientConfig("b");
+        assertEquals("replication.policy.separator is picked up in MirrorClientConfig",
+            "__", aClientConfig.getString("replication.policy.separator"));
+        assertEquals("replication.policy.separator is honored",
+            "b__topic1", aClientConfig.replicationPolicy().formatRemoteTopic("b", "topic1"));
+        assertEquals("client configs include boostrap.servers",
+            "one:9092, two:9092", aClientConfig.adminConfig().get("bootstrap.servers"));
+        assertEquals("client configs include security.protocol",
+            "PLAINTEXT", aClientConfig.adminConfig().get("security.protocol"));
+        assertEquals("producer configs include security.protocol",
+            "SASL", aClientConfig.producerConfig().get("security.protocol"));
+        assertFalse("unknown properties aren't included in client configs",
+            aClientConfig.adminConfig().containsKey("xxx"));
+        assertFalse("top-leve metrics reporters aren't included in client configs",
+            aClientConfig.adminConfig().containsKey("metric.reporters"));
+        assertEquals("security properties are picked up in MirrorClientConfig",
+            "secret1", aClientConfig.getPassword("ssl.truststore.password").value());
+        assertEquals("client configs include top-level security properties",
+            "secret1", ((Password) aClientConfig.adminConfig().get("ssl.truststore.password")).value());
+        assertEquals("security properties are translated from external sources",
+            "secret2", aClientConfig.getPassword("ssl.key.password").value());
+        assertEquals("client configs are translated from external sources",
+            "secret2", ((Password) aClientConfig.adminConfig().get("ssl.key.password")).value());
+        assertFalse("client configs should not include metrics reporter",
+            aClientConfig.producerConfig().containsKey("metrics.reporter"));
+        assertFalse("client configs should not include metrics reporter",
+            bClientConfig.adminConfig().containsKey("metrics.reporter"));
+    }
+
+    @Test
+    public void testIncludesConnectorConfigProperties() {
+        MirrorMakerConfig mirrorConfig = new MirrorMakerConfig(makeProps(
+            "clusters", "a, b",
+            "tasks.max", "100",
+            "topics", "topic-1",
+            "groups", "group-2",
+            "replication.policy.separator", "__",
+            "config.properties.blacklist", "property-3",
+            "metric.reporters", "FakeMetricsReporter",
+            "topic.filter.class", DefaultTopicFilter.class.getName(),
+            "xxx", "yyy"));
+        SourceAndTarget sourceAndTarget = new SourceAndTarget("source", "target");
+        Map connectorProps = mirrorConfig.connectorBaseConfig(sourceAndTarget,
+            MirrorSourceConnector.class);
+        MirrorConnectorConfig connectorConfig = new MirrorConnectorConfig(connectorProps);
+        assertEquals("Connector properties like tasks.max should be passed through to underlying Connectors.",
+            100, (int) connectorConfig.getInt("tasks.max"));
+        assertEquals("Topics whitelist should be passed through to underlying Connectors.",
+            Arrays.asList("topic-1"), connectorConfig.getList("topics"));
+        assertEquals("Groups whitelist should be passed through to underlying Connectors.",
+            Arrays.asList("group-2"), connectorConfig.getList("groups"));
+        assertEquals("Config properties blacklist should be passed through to underlying Connectors.",
+            Arrays.asList("property-3"), connectorConfig.getList("config.properties.blacklist"));
+        assertEquals("Metrics reporters should be passed through to underlying Connectors.",
+            Arrays.asList("FakeMetricsReporter"), connectorConfig.getList("metric.reporters"));
+        assertEquals("Filters should be passed through to underlying Connectors.",
+            "DefaultTopicFilter", connectorConfig.getClass("topic.filter.class").getSimpleName());
+        assertEquals("replication policy separator should be passed through to underlying Connectors.",
+            "__", connectorConfig.getString("replication.policy.separator"));
+        assertFalse("Unknown properties should not be passed through to Connectors.",
+            connectorConfig.originals().containsKey("xxx"));
+    }
+
+    @Test
+    public void testIncludesTopicFilterProperties() {
+        MirrorMakerConfig mirrorConfig = new MirrorMakerConfig(makeProps(
+            "clusters", "a, b",
+            "source->target.topics", "topic1, topic2",
+            "source->target.topics.blacklist", "topic3"));
+        SourceAndTarget sourceAndTarget = new SourceAndTarget("source", "target");
+        Map connectorProps = mirrorConfig.connectorBaseConfig(sourceAndTarget,
+            MirrorSourceConnector.class);
+        DefaultTopicFilter.TopicFilterConfig filterConfig = 
+            new DefaultTopicFilter.TopicFilterConfig(connectorProps);
+        assertEquals("source->target.topics should be passed through to TopicFilters.",
+            Arrays.asList("topic1", "topic2"), filterConfig.getList("topics"));
+        assertEquals("source->target.topics.blacklist should be passed through to TopicFilters.",
+            Arrays.asList("topic3"), filterConfig.getList("topics.blacklist"));
+    }
+
+    @Test
+    public void testWorkerConfigs() {
+        MirrorMakerConfig mirrorConfig = new MirrorMakerConfig(makeProps(
+            "clusters", "a, b",
+            "config.providers", "fake",
+            "config.providers.fake.class", FakeConfigProvider.class.getName(),
+            "replication.policy.separator", "__",
+            "offset.storage.replication.factor", "123",
+            "b.status.storage.replication.factor", "456",
+            "b.producer.client.id", "client-one",
+            "b.security.protocol", "PLAINTEXT",
+            "b.producer.security.protocol", "SASL",
+            "ssl.truststore.password", "secret1",
+            "ssl.key.password", "${fake:secret:password}",  // resolves to "secret2"
+            "b.xxx", "yyy"));
+        SourceAndTarget a = new SourceAndTarget("b", "a");
+        SourceAndTarget b = new SourceAndTarget("a", "b");
+        Map aProps = mirrorConfig.workerConfig(a);
+        assertEquals("123", aProps.get("offset.storage.replication.factor"));
+        Map bProps = mirrorConfig.workerConfig(b);
+        assertEquals("456", bProps.get("status.storage.replication.factor"));
+        assertEquals("producer props should be passed through to worker producer config: " + bProps,
+            "client-one", bProps.get("producer.client.id"));
+        assertEquals("replication-level security props should be passed through to worker producer config",
+            "SASL", bProps.get("producer.security.protocol"));
+        assertEquals("replication-level security props should be passed through to worker producer config",
+            "SASL", bProps.get("producer.security.protocol"));
+        assertEquals("replication-level security props should be passed through to worker consumer config",
+            "PLAINTEXT", bProps.get("consumer.security.protocol"));
+        assertEquals("security properties should be passed through to worker config: " + bProps,
+            "secret1", bProps.get("ssl.truststore.password"));
+        assertEquals("security properties should be passed through to worker producer config: " + bProps,
+            "secret1", bProps.get("producer.ssl.truststore.password"));
+        assertEquals("security properties should be transformed in worker config",
+            "secret2", bProps.get("ssl.key.password"));
+        assertEquals("security properties should be transformed in worker producer config",
+            "secret2", bProps.get("producer.ssl.key.password"));
+    }
+
+    public static class FakeConfigProvider implements ConfigProvider {
+
+        Map secrets = Collections.singletonMap("password", "secret2");
+
+        @Override
+        public void configure(Map props) {
+        }
+
+        @Override
+        public void close() {
+        }
+
+        @Override
+        public ConfigData get(String path) {
+            return new ConfigData(secrets);
+        }
+
+        @Override
+        public ConfigData get(String path, Set keys) {
+            return get(path);
+        }
+    }
+}
diff --git a/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/MirrorSourceConnectorTest.java b/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/MirrorSourceConnectorTest.java
new file mode 100644
index 0000000000000..b1ccef8158df7
--- /dev/null
+++ b/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/MirrorSourceConnectorTest.java
@@ -0,0 +1,115 @@
+/*
+ * 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.mirror;
+
+import org.apache.kafka.common.acl.AccessControlEntry;
+import org.apache.kafka.common.acl.AclBinding;
+import org.apache.kafka.common.acl.AclOperation;
+import org.apache.kafka.common.acl.AclPermissionType;
+import org.apache.kafka.common.resource.PatternType;
+import org.apache.kafka.common.resource.ResourcePattern;
+import org.apache.kafka.common.resource.ResourceType;
+import org.apache.kafka.clients.admin.Config;
+import org.apache.kafka.clients.admin.ConfigEntry;
+
+import org.junit.Test;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.assertFalse;
+
+import java.util.ArrayList;
+
+public class MirrorSourceConnectorTest {
+
+    @Test
+    public void testReplicatesHeartbeatsByDefault() {
+        MirrorSourceConnector connector = new MirrorSourceConnector(new SourceAndTarget("source", "target"), 
+            new DefaultReplicationPolicy(), new DefaultTopicFilter(), new DefaultConfigPropertyFilter());
+        assertTrue("should replicate heartbeats", connector.shouldReplicateTopic("heartbeats"));
+        assertTrue("should replicate upstream heartbeats", connector.shouldReplicateTopic("us-west.heartbeats"));
+    }
+
+    @Test
+    public void testReplicatesHeartbeatsDespiteFilter() {
+        MirrorSourceConnector connector = new MirrorSourceConnector(new SourceAndTarget("source", "target"),
+            new DefaultReplicationPolicy(), x -> false, new DefaultConfigPropertyFilter());
+        assertTrue("should replicate heartbeats", connector.shouldReplicateTopic("heartbeats"));
+        assertTrue("should replicate upstream heartbeats", connector.shouldReplicateTopic("us-west.heartbeats"));
+    }
+
+    @Test
+    public void testNoCycles() {
+        MirrorSourceConnector connector = new MirrorSourceConnector(new SourceAndTarget("source", "target"),
+            new DefaultReplicationPolicy(), x -> true, x -> true);
+        assertFalse("should not allow cycles", connector.shouldReplicateTopic("target.topic1"));
+        assertFalse("should not allow cycles", connector.shouldReplicateTopic("target.source.topic1"));
+        assertFalse("should not allow cycles", connector.shouldReplicateTopic("source.target.topic1"));
+        assertTrue("should allow anything else", connector.shouldReplicateTopic("topic1"));
+        assertTrue("should allow anything else", connector.shouldReplicateTopic("source.topic1"));
+    }
+
+    @Test
+    public void testAclFiltering() {
+        MirrorSourceConnector connector = new MirrorSourceConnector(new SourceAndTarget("source", "target"),
+            new DefaultReplicationPolicy(), x -> true, x -> true);
+        assertFalse("should not replicate ALLOW WRITE", connector.shouldReplicateAcl(
+            new AclBinding(new ResourcePattern(ResourceType.TOPIC, "test_topic", PatternType.LITERAL),
+            new AccessControlEntry("kafka", "", AclOperation.WRITE, AclPermissionType.ALLOW))));
+        assertTrue("should replicate ALLOW ALL", connector.shouldReplicateAcl(
+            new AclBinding(new ResourcePattern(ResourceType.TOPIC, "test_topic", PatternType.LITERAL),
+            new AccessControlEntry("kafka", "", AclOperation.ALL, AclPermissionType.ALLOW))));
+    }
+
+    @Test
+    public void testAclTransformation() {
+        MirrorSourceConnector connector = new MirrorSourceConnector(new SourceAndTarget("source", "target"),
+            new DefaultReplicationPolicy(), x -> true, x -> true);
+        AclBinding allowAllAclBinding = new AclBinding(
+            new ResourcePattern(ResourceType.TOPIC, "test_topic", PatternType.LITERAL),
+            new AccessControlEntry("kafka", "", AclOperation.ALL, AclPermissionType.ALLOW));
+        AclBinding processedAllowAllAclBinding = connector.targetAclBinding(allowAllAclBinding);
+        String expectedRemoteTopicName = "source" + DefaultReplicationPolicy.SEPARATOR_DEFAULT
+            + allowAllAclBinding.pattern().name();
+        assertTrue("should change topic name",
+            processedAllowAllAclBinding.pattern().name().equals(expectedRemoteTopicName));
+        assertTrue("should change ALL to READ", processedAllowAllAclBinding.entry().operation() == AclOperation.READ);
+        assertTrue("should not change ALLOW",
+            processedAllowAllAclBinding.entry().permissionType() == AclPermissionType.ALLOW);
+
+        AclBinding denyAllAclBinding = new AclBinding(
+            new ResourcePattern(ResourceType.TOPIC, "test_topic", PatternType.LITERAL),
+            new AccessControlEntry("kafka", "", AclOperation.ALL, AclPermissionType.DENY));
+        AclBinding processedDenyAllAclBinding = connector.targetAclBinding(denyAllAclBinding);
+        assertTrue("should not change ALL", processedDenyAllAclBinding.entry().operation() == AclOperation.ALL);
+        assertTrue("should not change DENY",
+            processedDenyAllAclBinding.entry().permissionType() == AclPermissionType.DENY);
+    }
+    
+    @Test
+    public void testConfigPropertyFiltering() {
+        MirrorSourceConnector connector = new MirrorSourceConnector(new SourceAndTarget("source", "target"),
+            new DefaultReplicationPolicy(), x -> true, new DefaultConfigPropertyFilter());
+        ArrayList entries = new ArrayList<>();
+        entries.add(new ConfigEntry("name-1", "value-1"));
+        entries.add(new ConfigEntry("min.insync.replicas", "2"));
+        Config config = new Config(entries);
+        Config targetConfig = connector.targetConfig(config);
+        assertTrue("should replicate properties", targetConfig.entries().stream()
+            .anyMatch(x -> x.name().equals("name-1")));
+        assertFalse("should not replicate blacklisted properties", targetConfig.entries().stream()
+            .anyMatch(x -> x.name().equals("min.insync.replicas")));
+    }
+}
diff --git a/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/MirrorSourceTaskTest.java b/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/MirrorSourceTaskTest.java
new file mode 100644
index 0000000000000..003511760fa26
--- /dev/null
+++ b/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/MirrorSourceTaskTest.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.connect.mirror;
+
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.common.record.TimestampType;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.header.Headers;
+import org.apache.kafka.common.header.internals.RecordHeaders;
+import org.apache.kafka.connect.source.SourceRecord;
+
+import org.junit.Test;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.assertFalse;
+
+public class MirrorSourceTaskTest {
+
+    @Test
+    public void testSerde() {
+        byte[] key = new byte[]{'a', 'b', 'c', 'd', 'e'};
+        byte[] value = new byte[]{'f', 'g', 'h', 'i', 'j', 'k'};
+        Headers headers = new RecordHeaders();
+        headers.add("header1", new byte[]{'l', 'm', 'n', 'o'});
+        headers.add("header2", new byte[]{'p', 'q', 'r', 's', 't'});
+        ConsumerRecord consumerRecord = new ConsumerRecord<>("topic1", 2, 3L, 4L,
+            TimestampType.CREATE_TIME, 0L, 5, 6, key, value, headers);
+        MirrorSourceTask mirrorSourceTask = new MirrorSourceTask("cluster7",
+            new DefaultReplicationPolicy(), 50);
+        SourceRecord sourceRecord = mirrorSourceTask.convertRecord(consumerRecord);
+        assertEquals("cluster7.topic1", sourceRecord.topic());
+        assertEquals(2, sourceRecord.kafkaPartition().intValue());
+        assertEquals(new TopicPartition("topic1", 2), MirrorUtils.unwrapPartition(sourceRecord.sourcePartition()));
+        assertEquals(3L, MirrorUtils.unwrapOffset(sourceRecord.sourceOffset()).longValue());
+        assertEquals(4L, sourceRecord.timestamp().longValue());
+        assertEquals(key, sourceRecord.key());
+        assertEquals(value, sourceRecord.value());
+        assertEquals(headers.lastHeader("header1").value(), sourceRecord.headers().lastWithName("header1").value());
+        assertEquals(headers.lastHeader("header2").value(), sourceRecord.headers().lastWithName("header2").value());
+    }
+
+    @Test
+    public void testOffsetSync() {
+        MirrorSourceTask.PartitionState partitionState = new MirrorSourceTask.PartitionState(50);
+
+        assertTrue("always emit offset sync on first update",
+            partitionState.update(0, 100));
+        assertTrue("upstream offset skipped -> resync",
+            partitionState.update(2, 102));
+        assertFalse("no sync",
+            partitionState.update(3, 152));
+        assertFalse("no sync",
+            partitionState.update(4, 153));
+        assertFalse("no sync",
+            partitionState.update(5, 154));
+        assertTrue("one past target offset",
+            partitionState.update(6, 205));
+        assertTrue("upstream reset",
+            partitionState.update(2, 206));
+        assertFalse("no sync",
+            partitionState.update(3, 207));
+        assertTrue("downstream reset",
+                partitionState.update(4, 3));
+        assertFalse("no sync",
+            partitionState.update(5, 4));
+    }
+
+    @Test
+    public void testZeroOffsetSync() {
+        MirrorSourceTask.PartitionState partitionState = new MirrorSourceTask.PartitionState(0);
+
+        // if max offset lag is zero, should always emit offset syncs
+        assertTrue(partitionState.update(0, 100));
+        assertTrue(partitionState.update(2, 102));
+        assertTrue(partitionState.update(3, 153));
+        assertTrue(partitionState.update(4, 154));
+        assertTrue(partitionState.update(5, 155));
+        assertTrue(partitionState.update(6, 207));
+        assertTrue(partitionState.update(2, 208));
+        assertTrue(partitionState.update(3, 209));
+        assertTrue(partitionState.update(4, 3));
+        assertTrue(partitionState.update(5, 4));
+    }
+}
diff --git a/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/OffsetSyncStoreTest.java b/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/OffsetSyncStoreTest.java
new file mode 100644
index 0000000000000..19954cd8fcaf4
--- /dev/null
+++ b/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/OffsetSyncStoreTest.java
@@ -0,0 +1,67 @@
+/*
+ * 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.mirror;
+
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.common.TopicPartition;
+
+import org.junit.Test;
+
+import static org.junit.Assert.assertEquals;
+
+public class OffsetSyncStoreTest {
+
+    static TopicPartition tp = new TopicPartition("topic1", 2);
+
+    static class FakeOffsetSyncStore extends OffsetSyncStore {
+
+        FakeOffsetSyncStore() {
+            super(null, null);
+        }
+
+        void sync(TopicPartition topicPartition, long upstreamOffset, long downstreamOffset) {
+            OffsetSync offsetSync = new OffsetSync(topicPartition, upstreamOffset, downstreamOffset);
+            byte[] key = offsetSync.recordKey();
+            byte[] value = offsetSync.recordValue();
+            ConsumerRecord record = new ConsumerRecord<>("test.offsets.internal", 0, 3, key, value);
+            handleRecord(record);
+        }
+    }
+
+    @Test
+    public void testOffsetTranslation() {
+        FakeOffsetSyncStore store = new FakeOffsetSyncStore();
+
+        store.sync(tp, 100, 200);
+        assertEquals(store.translateDownstream(tp, 150), 250);
+
+        // Translate exact offsets
+        store.sync(tp, 150, 251);
+        assertEquals(store.translateDownstream(tp, 150), 251);
+
+        // Use old offset (5) prior to any sync -> can't translate
+        assertEquals(-1, store.translateDownstream(tp, 5));
+
+        // Downstream offsets reset
+        store.sync(tp, 200, 10);
+        assertEquals(store.translateDownstream(tp, 200), 10);
+
+        // Upstream offsets reset
+        store.sync(tp, 20, 20);
+        assertEquals(store.translateDownstream(tp, 20), 20);
+    }
+}
diff --git a/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/OffsetSyncTest.java b/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/OffsetSyncTest.java
new file mode 100644
index 0000000000000..5dc472964b345
--- /dev/null
+++ b/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/OffsetSyncTest.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.connect.mirror;
+
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.common.TopicPartition;
+
+import org.junit.Test;
+
+import static org.junit.Assert.assertEquals;
+
+public class OffsetSyncTest {
+
+    @Test
+    public void testSerde() {
+        OffsetSync offsetSync = new OffsetSync(new TopicPartition("topic-1", 2), 3, 4);
+        byte[] key = offsetSync.recordKey();
+        byte[] value = offsetSync.recordValue();
+        ConsumerRecord record = new ConsumerRecord<>("any-topic", 6, 7, key, value);
+        OffsetSync deserialized = OffsetSync.deserializeRecord(record);
+        assertEquals(offsetSync.topicPartition(), deserialized.topicPartition());
+        assertEquals(offsetSync.upstreamOffset(), deserialized.upstreamOffset());
+        assertEquals(offsetSync.downstreamOffset(), deserialized.downstreamOffset());
+    }
+}
diff --git a/connect/mirror/src/test/resources/log4j.properties b/connect/mirror/src/test/resources/log4j.properties
new file mode 100644
index 0000000000000..a2ac021dfab98
--- /dev/null
+++ b/connect/mirror/src/test/resources/log4j.properties
@@ -0,0 +1,34 @@
+##
+# 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.
+##
+log4j.rootLogger=ERROR, stdout
+
+log4j.appender.stdout=org.apache.log4j.ConsoleAppender
+log4j.appender.stdout.layout=org.apache.log4j.PatternLayout
+#
+# The `%X{connector.context}` parameter in the layout includes connector-specific and task-specific information
+# in the log message, where appropriate. This makes it easier to identify those log messages that apply to a
+# specific connector. Simply add this parameter to the log layout configuration below to include the contextual information.
+#
+log4j.appender.stdout.layout.ConversionPattern=[%d] %p %X{connector.context}%m (%c:%L)%n
+#
+# The following line includes no MDC context parameters:
+#log4j.appender.stdout.layout.ConversionPattern=[%d] %p %m (%c:%L)%n (%t)
+
+log4j.logger.org.reflections=OFF
+log4j.logger.kafka=OFF
+log4j.logger.state.change.logger=OFF
+log4j.logger.org.apache.kafka.connect.mirror=INFO
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/cli/ConnectDistributed.java b/connect/runtime/src/main/java/org/apache/kafka/connect/cli/ConnectDistributed.java
index a6c6d98facacc..22c1ad82d6138 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/cli/ConnectDistributed.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/cli/ConnectDistributed.java
@@ -19,9 +19,10 @@
 import org.apache.kafka.common.utils.Exit;
 import org.apache.kafka.common.utils.Time;
 import org.apache.kafka.common.utils.Utils;
+import org.apache.kafka.connect.connector.policy.ConnectorClientConfigOverridePolicy;
 import org.apache.kafka.connect.runtime.Connect;
-import org.apache.kafka.connect.runtime.HerderProvider;
 import org.apache.kafka.connect.runtime.Worker;
+import org.apache.kafka.connect.runtime.WorkerConfig;
 import org.apache.kafka.connect.runtime.WorkerConfigTransformer;
 import org.apache.kafka.connect.runtime.WorkerInfo;
 import org.apache.kafka.connect.runtime.distributed.DistributedConfig;
@@ -95,8 +96,7 @@ public Connect startConnect(Map workerProps) {
         log.debug("Kafka cluster ID: {}", kafkaClusterId);
 
         RestServer rest = new RestServer(config);
-        HerderProvider provider = new HerderProvider();
-        rest.start(provider, plugins);
+        rest.initializeServer();
 
         URI advertisedUrl = rest.advertisedUrl();
         String workerId = advertisedUrl.getHost() + ":" + advertisedUrl.getPort();
@@ -104,7 +104,11 @@ public Connect startConnect(Map workerProps) {
         KafkaOffsetBackingStore offsetBackingStore = new KafkaOffsetBackingStore();
         offsetBackingStore.configure(config);
 
-        Worker worker = new Worker(workerId, time, plugins, config, offsetBackingStore);
+        ConnectorClientConfigOverridePolicy connectorClientConfigOverridePolicy = plugins.newPlugin(
+                config.getString(WorkerConfig.CONNECTOR_CLIENT_POLICY_CLASS_CONFIG),
+                config, ConnectorClientConfigOverridePolicy.class);
+
+        Worker worker = new Worker(workerId, time, plugins, config, offsetBackingStore, connectorClientConfigOverridePolicy);
         WorkerConfigTransformer configTransformer = worker.configTransformer();
 
         Converter internalValueConverter = worker.getInternalValueConverter();
@@ -118,14 +122,12 @@ public Connect startConnect(Map workerProps) {
 
         DistributedHerder herder = new DistributedHerder(config, time, worker,
                 kafkaClusterId, statusBackingStore, configBackingStore,
-                advertisedUrl.toString());
+                advertisedUrl.toString(), connectorClientConfigOverridePolicy);
 
         final Connect connect = new Connect(herder, rest);
         log.info("Kafka Connect distributed worker initialization took {}ms", time.hiResClockMs() - initStart);
         try {
             connect.start();
-            // herder has initialized now, and ready to be used by the RestServer.
-            provider.setHerder(herder);
         } catch (Exception e) {
             log.error("Failed to start Connect", e);
             connect.stop();
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/cli/ConnectStandalone.java b/connect/runtime/src/main/java/org/apache/kafka/connect/cli/ConnectStandalone.java
index dd1cf0f7a307b..cf7b93bd7c838 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/cli/ConnectStandalone.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/cli/ConnectStandalone.java
@@ -19,11 +19,12 @@
 import org.apache.kafka.common.utils.Exit;
 import org.apache.kafka.common.utils.Time;
 import org.apache.kafka.common.utils.Utils;
+import org.apache.kafka.connect.connector.policy.ConnectorClientConfigOverridePolicy;
 import org.apache.kafka.connect.runtime.Connect;
 import org.apache.kafka.connect.runtime.ConnectorConfig;
 import org.apache.kafka.connect.runtime.Herder;
-import org.apache.kafka.connect.runtime.HerderProvider;
 import org.apache.kafka.connect.runtime.Worker;
+import org.apache.kafka.connect.runtime.WorkerConfig;
 import org.apache.kafka.connect.runtime.WorkerInfo;
 import org.apache.kafka.connect.runtime.isolation.Plugins;
 import org.apache.kafka.connect.runtime.rest.RestServer;
@@ -83,22 +84,23 @@ public static void main(String[] args) {
             log.debug("Kafka cluster ID: {}", kafkaClusterId);
 
             RestServer rest = new RestServer(config);
-            HerderProvider provider = new HerderProvider();
-            rest.start(provider, plugins);
+            rest.initializeServer();
 
             URI advertisedUrl = rest.advertisedUrl();
             String workerId = advertisedUrl.getHost() + ":" + advertisedUrl.getPort();
 
-            Worker worker = new Worker(workerId, time, plugins, config, new FileOffsetBackingStore());
+            ConnectorClientConfigOverridePolicy connectorClientConfigOverridePolicy = plugins.newPlugin(
+                config.getString(WorkerConfig.CONNECTOR_CLIENT_POLICY_CLASS_CONFIG),
+                config, ConnectorClientConfigOverridePolicy.class);
+            Worker worker = new Worker(workerId, time, plugins, config, new FileOffsetBackingStore(),
+                                       connectorClientConfigOverridePolicy);
 
-            Herder herder = new StandaloneHerder(worker, kafkaClusterId);
+            Herder herder = new StandaloneHerder(worker, kafkaClusterId, connectorClientConfigOverridePolicy);
             final Connect connect = new Connect(herder, rest);
             log.info("Kafka Connect standalone worker initialization took {}ms", time.hiResClockMs() - initStart);
 
             try {
                 connect.start();
-                // herder has initialized now, and ready to be used by the RestServer.
-                provider.setHerder(herder);
                 for (final String connectorPropsFile : Arrays.copyOfRange(args, 1, args.length)) {
                     Map connectorProps = Utils.propsToStringMap(Utils.loadProps(connectorPropsFile));
                     FutureCallback> cb = new FutureCallback<>(new Callback>() {
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/connector/policy/AbstractConnectorClientConfigOverridePolicy.java b/connect/runtime/src/main/java/org/apache/kafka/connect/connector/policy/AbstractConnectorClientConfigOverridePolicy.java
new file mode 100644
index 0000000000000..3c310db966c04
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/connector/policy/AbstractConnectorClientConfigOverridePolicy.java
@@ -0,0 +1,57 @@
+/*
+ * 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.connector.policy;
+
+import org.apache.kafka.common.config.ConfigValue;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+public abstract class AbstractConnectorClientConfigOverridePolicy implements ConnectorClientConfigOverridePolicy {
+
+    @Override
+    public void close() throws Exception {
+
+    }
+
+    @Override
+    public final List validate(ConnectorClientConfigRequest connectorClientConfigRequest) {
+        Map inputConfig = connectorClientConfigRequest.clientProps();
+        return inputConfig.entrySet().stream().map(this::configValue).collect(Collectors.toList());
+    }
+
+    protected ConfigValue configValue(Map.Entry configEntry) {
+        ConfigValue configValue =
+            new ConfigValue(configEntry.getKey(), configEntry.getValue(), new ArrayList<>(), new ArrayList());
+        validate(configValue);
+        return configValue;
+    }
+
+    protected void validate(ConfigValue configValue) {
+        if (!isAllowed(configValue)) {
+            configValue.addErrorMessage("The '" + policyName() + "' policy does not allow '" + configValue.name()
+                                        + "' to be overridden in the connector configuration.");
+        }
+    }
+
+    protected abstract String policyName();
+
+    protected abstract boolean isAllowed(ConfigValue configValue);
+}
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/connector/policy/AllConnectorClientConfigOverridePolicy.java b/connect/runtime/src/main/java/org/apache/kafka/connect/connector/policy/AllConnectorClientConfigOverridePolicy.java
new file mode 100644
index 0000000000000..e8088570ffd8f
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/connector/policy/AllConnectorClientConfigOverridePolicy.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.connect.connector.policy;
+
+import org.apache.kafka.common.config.ConfigValue;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.Map;
+
+/**
+ * Allows all client configurations to be overridden via the connector configs by setting {@code connector.client.config.override.policy} to {@code All}
+ */
+public class AllConnectorClientConfigOverridePolicy extends AbstractConnectorClientConfigOverridePolicy {
+    private static final Logger log = LoggerFactory.getLogger(AllConnectorClientConfigOverridePolicy.class);
+
+    @Override
+    protected String policyName() {
+        return "All";
+    }
+
+    @Override
+    protected boolean isAllowed(ConfigValue configValue) {
+        return true;
+    }
+
+    @Override
+    public void configure(Map configs) {
+        log.info("Setting up All Policy for ConnectorClientConfigOverride. This will allow all client configurations to be overridden");
+    }
+}
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/connector/policy/NoneConnectorClientConfigOverridePolicy.java b/connect/runtime/src/main/java/org/apache/kafka/connect/connector/policy/NoneConnectorClientConfigOverridePolicy.java
new file mode 100644
index 0000000000000..9b414c4323ee0
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/connector/policy/NoneConnectorClientConfigOverridePolicy.java
@@ -0,0 +1,47 @@
+/*
+ * 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.connector.policy;
+
+import org.apache.kafka.common.config.ConfigValue;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.Map;
+
+/**
+ * Disallow any client configuration to be overridden via the connector configs by setting {@code connector.client.config.override.policy} to {@code None}.
+ * This is the default behavior.
+ */
+public class NoneConnectorClientConfigOverridePolicy extends AbstractConnectorClientConfigOverridePolicy {
+    private static final Logger log = LoggerFactory.getLogger(NoneConnectorClientConfigOverridePolicy.class);
+
+    @Override
+    protected String policyName() {
+        return "None";
+    }
+
+    @Override
+    protected boolean isAllowed(ConfigValue configValue) {
+        return false;
+    }
+
+    @Override
+    public void configure(Map configs) {
+        log.info("Setting up None Policy for ConnectorClientConfigOverride. This will disallow any client configuration to be overridden");
+    }
+}
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/connector/policy/PrincipalConnectorClientConfigOverridePolicy.java b/connect/runtime/src/main/java/org/apache/kafka/connect/connector/policy/PrincipalConnectorClientConfigOverridePolicy.java
new file mode 100644
index 0000000000000..492c5a9599d54
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/connector/policy/PrincipalConnectorClientConfigOverridePolicy.java
@@ -0,0 +1,57 @@
+/*
+ * 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.connector.policy;
+
+import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.common.config.ConfigValue;
+import org.apache.kafka.common.config.SaslConfigs;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.Map;
+import java.util.Set;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+/**
+ * Allows all {@code sasl} configurations to be overridden via the connector configs by setting {@code connector.client.config.override.policy} to
+ * {@code Principal}. This allows to set a principal per connector.
+ */
+public class PrincipalConnectorClientConfigOverridePolicy extends AbstractConnectorClientConfigOverridePolicy {
+    private static final Logger log = LoggerFactory.getLogger(PrincipalConnectorClientConfigOverridePolicy.class);
+
+    private static final Set ALLOWED_CONFIG =
+        Stream.of(SaslConfigs.SASL_JAAS_CONFIG, SaslConfigs.SASL_MECHANISM, CommonClientConfigs.SECURITY_PROTOCOL_CONFIG).
+            collect(Collectors.toSet());
+
+    @Override
+    protected String policyName() {
+        return "Principal";
+    }
+
+    @Override
+    protected boolean isAllowed(ConfigValue configValue) {
+        return ALLOWED_CONFIG.contains(configValue.name());
+    }
+
+    @Override
+    public void configure(Map configs) {
+        log.info("Setting up Principal policy for ConnectorClientConfigOverride. This will allow `sasl` client configuration to be "
+                 + "overridden.");
+    }
+}
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/AbstractHerder.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/AbstractHerder.java
index 82fdeccc96ba4..c1881503ff6fd 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/AbstractHerder.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/AbstractHerder.java
@@ -16,6 +16,8 @@
  */
 package org.apache.kafka.connect.runtime;
 
+import org.apache.kafka.clients.producer.ProducerConfig;
+import org.apache.kafka.common.config.AbstractConfig;
 import org.apache.kafka.common.config.Config;
 import org.apache.kafka.common.config.ConfigDef;
 import org.apache.kafka.common.config.ConfigDef.ConfigKey;
@@ -23,6 +25,8 @@
 import org.apache.kafka.common.config.ConfigTransformer;
 import org.apache.kafka.common.config.ConfigValue;
 import org.apache.kafka.connect.connector.Connector;
+import org.apache.kafka.connect.connector.policy.ConnectorClientConfigOverridePolicy;
+import org.apache.kafka.connect.connector.policy.ConnectorClientConfigRequest;
 import org.apache.kafka.connect.errors.NotFoundException;
 import org.apache.kafka.connect.runtime.distributed.ClusterConfigState;
 import org.apache.kafka.connect.runtime.isolation.Plugins;
@@ -87,6 +91,7 @@ public abstract class AbstractHerder implements Herder, TaskStatus.Listener, Con
     private final String kafkaClusterId;
     protected final StatusBackingStore statusBackingStore;
     protected final ConfigBackingStore configBackingStore;
+    private final ConnectorClientConfigOverridePolicy connectorClientConfigOverridePolicy;
 
     private Map tempConnectors = new ConcurrentHashMap<>();
 
@@ -94,13 +99,15 @@ public AbstractHerder(Worker worker,
                           String workerId,
                           String kafkaClusterId,
                           StatusBackingStore statusBackingStore,
-                          ConfigBackingStore configBackingStore) {
+                          ConfigBackingStore configBackingStore,
+                          ConnectorClientConfigOverridePolicy connectorClientConfigOverridePolicy) {
         this.worker = worker;
         this.worker.herder = this;
         this.workerId = workerId;
         this.kafkaClusterId = kafkaClusterId;
         this.statusBackingStore = statusBackingStore;
         this.configBackingStore = configBackingStore;
+        this.connectorClientConfigOverridePolicy = connectorClientConfigOverridePolicy;
     }
 
     @Override
@@ -180,10 +187,15 @@ public void onPause(ConnectorTaskId id) {
     @Override
     public void onDeletion(String connector) {
         for (TaskStatus status : statusBackingStore.getAll(connector))
-            statusBackingStore.put(new TaskStatus(status.id(), TaskStatus.State.DESTROYED, workerId, generation()));
+            onDeletion(status.id());
         statusBackingStore.put(new ConnectorStatus(connector, ConnectorStatus.State.DESTROYED, workerId, generation()));
     }
 
+    @Override
+    public void onDeletion(ConnectorTaskId id) {
+        statusBackingStore.put(new TaskStatus(id, TaskStatus.State.DESTROYED, workerId, generation()));
+    }
+
     @Override
     public void pauseConnector(String connector) {
         if (!configBackingStore.contains(connector))
@@ -208,6 +220,27 @@ public Plugins plugins() {
      */
     protected abstract Map config(String connName);
 
+    @Override
+    public Collection connectors() {
+        return configBackingStore.snapshot().connectors();
+    }
+
+    @Override
+    public ConnectorInfo connectorInfo(String connector) {
+        final ClusterConfigState configState = configBackingStore.snapshot();
+
+        if (!configState.contains(connector))
+            return null;
+        Map config = configState.rawConnectorConfig(connector);
+
+        return new ConnectorInfo(
+            connector,
+            config,
+            configState.tasks(connector),
+            connectorTypeForClass(config.get(ConnectorConfig.CONNECTOR_CLASS_CONFIG))
+        );
+    }
+
     @Override
     public ConnectorStateInfo connectorStatus(String connName) {
         ConnectorStatus connector = statusBackingStore.get(connName);
@@ -259,14 +292,17 @@ public ConfigInfos validateConnectorConfig(Map connectorProps) {
             throw new BadRequestException("Connector config " + connectorProps + " contains no connector type");
 
         Connector connector = getConnector(connType);
+        org.apache.kafka.connect.health.ConnectorType connectorType;
         ClassLoader savedLoader = plugins().compareAndSwapLoaders(connector);
         try {
             ConfigDef baseConfigDef;
             if (connector instanceof SourceConnector) {
                 baseConfigDef = SourceConnectorConfig.configDef();
+                connectorType = org.apache.kafka.connect.health.ConnectorType.SOURCE;
             } else {
                 baseConfigDef = SinkConnectorConfig.configDef();
                 SinkConnectorConfig.validate(connectorProps);
+                connectorType = org.apache.kafka.connect.health.ConnectorType.SINK;
             }
             ConfigDef enrichedConfigDef = ConnectorConfig.enrich(plugins(), baseConfigDef, connectorProps, false);
             Map validatedConnectorConfig = validateBasicConnectorConfig(
@@ -279,33 +315,135 @@ public ConfigInfos validateConnectorConfig(Map connectorProps) {
             Set allGroups = new LinkedHashSet<>(enrichedConfigDef.groups());
 
             // do custom connector-specific validation
-            Config config = connector.validate(connectorProps);
-            if (null == config) {
+            ConfigDef configDef = connector.config();
+            if (null == configDef) {
                 throw new BadRequestException(
-                    String.format(
-                        "%s.validate() must return a Config that is not null.",
-                        connector.getClass().getName()
-                    )
+                        String.format(
+                                "%s.config() must return a ConfigDef that is not null.",
+                                connector.getClass().getName()
+                        )
                 );
             }
-            ConfigDef configDef = connector.config();
-            if (null == configDef) {
+            Config config = connector.validate(connectorProps);
+            if (null == config) {
                 throw new BadRequestException(
-                    String.format(
-                        "%s.config() must return a ConfigDef that is not null.",
-                        connector.getClass().getName()
-                    )
+                        String.format(
+                                "%s.validate() must return a Config that is not null.",
+                                connector.getClass().getName()
+                        )
                 );
             }
             configKeys.putAll(configDef.configKeys());
             allGroups.addAll(configDef.groups());
             configValues.addAll(config.configValues());
-            return generateResult(connType, configKeys, configValues, new ArrayList<>(allGroups));
+            ConfigInfos configInfos =  generateResult(connType, configKeys, configValues, new ArrayList<>(allGroups));
+
+            AbstractConfig connectorConfig = new AbstractConfig(new ConfigDef(), connectorProps);
+            String connName = connectorProps.get(ConnectorConfig.NAME_CONFIG);
+            ConfigInfos producerConfigInfos = null;
+            ConfigInfos consumerConfigInfos = null;
+            ConfigInfos adminConfigInfos = null;
+            if (connectorType.equals(org.apache.kafka.connect.health.ConnectorType.SOURCE)) {
+                producerConfigInfos = validateClientOverrides(connName,
+                                                              ConnectorConfig.CONNECTOR_CLIENT_PRODUCER_OVERRIDES_PREFIX,
+                                                              connectorConfig,
+                                                              ProducerConfig.configDef(),
+                                                              connector.getClass(),
+                                                              connectorType,
+                                                              ConnectorClientConfigRequest.ClientType.PRODUCER,
+                                                              connectorClientConfigOverridePolicy);
+                return mergeConfigInfos(connType, configInfos, producerConfigInfos);
+            } else {
+                consumerConfigInfos = validateClientOverrides(connName,
+                                                              ConnectorConfig.CONNECTOR_CLIENT_CONSUMER_OVERRIDES_PREFIX,
+                                                              connectorConfig,
+                                                              ProducerConfig.configDef(),
+                                                              connector.getClass(),
+                                                              connectorType,
+                                                              ConnectorClientConfigRequest.ClientType.CONSUMER,
+                                                              connectorClientConfigOverridePolicy);
+                // check if topic for dead letter queue exists
+                String topic = connectorProps.get(SinkConnectorConfig.DLQ_TOPIC_NAME_CONFIG);
+                if (topic != null && !topic.isEmpty()) {
+                    adminConfigInfos = validateClientOverrides(connName,
+                                                               ConnectorConfig.CONNECTOR_CLIENT_ADMIN_OVERRIDES_PREFIX,
+                                                               connectorConfig,
+                                                               ProducerConfig.configDef(),
+                                                               connector.getClass(),
+                                                               connectorType,
+                                                               ConnectorClientConfigRequest.ClientType.ADMIN,
+                                                               connectorClientConfigOverridePolicy);
+                }
+
+            }
+            return mergeConfigInfos(connType, configInfos, producerConfigInfos, consumerConfigInfos, adminConfigInfos);
         } finally {
             Plugins.compareAndSwapLoaders(savedLoader);
         }
     }
 
+    private static ConfigInfos mergeConfigInfos(String connType, ConfigInfos... configInfosList) {
+        int errorCount = 0;
+        List configInfoList = new LinkedList<>();
+        Set groups = new LinkedHashSet<>();
+        for (ConfigInfos configInfos : configInfosList) {
+            if (configInfos != null) {
+                errorCount += configInfos.errorCount();
+                configInfoList.addAll(configInfos.values());
+                groups.addAll(configInfos.groups());
+            }
+        }
+        return new ConfigInfos(connType, errorCount, new ArrayList<>(groups), configInfoList);
+    }
+
+    private static ConfigInfos validateClientOverrides(String connName,
+                                                      String prefix,
+                                                      AbstractConfig connectorConfig,
+                                                      ConfigDef configDef,
+                                                      Class connectorClass,
+                                                      org.apache.kafka.connect.health.ConnectorType connectorType,
+                                                      ConnectorClientConfigRequest.ClientType clientType,
+                                                      ConnectorClientConfigOverridePolicy connectorClientConfigOverridePolicy) {
+        int errorCount = 0;
+        List configInfoList = new LinkedList<>();
+        Map configKeys = configDef.configKeys();
+        Set groups = new LinkedHashSet<>();
+        Map clientConfigs = new HashMap<>();
+        for (Map.Entry rawClientConfig : connectorConfig.originalsWithPrefix(prefix).entrySet()) {
+            String configName = rawClientConfig.getKey();
+            Object rawConfigValue = rawClientConfig.getValue();
+            ConfigKey configKey = configDef.configKeys().get(configName);
+            Object parsedConfigValue = configKey != null
+                ? ConfigDef.parseType(configName, rawConfigValue, configKey.type)
+                : rawConfigValue;
+            clientConfigs.put(configName, parsedConfigValue);
+        }
+        ConnectorClientConfigRequest connectorClientConfigRequest = new ConnectorClientConfigRequest(
+            connName, connectorType, connectorClass, clientConfigs, clientType);
+        List configValues = connectorClientConfigOverridePolicy.validate(connectorClientConfigRequest);
+        if (configValues != null) {
+            for (ConfigValue validatedConfigValue : configValues) {
+                ConfigKey configKey = configKeys.get(validatedConfigValue.name());
+                ConfigKeyInfo configKeyInfo = null;
+                if (configKey != null) {
+                    if (configKey.group != null) {
+                        groups.add(configKey.group);
+                    }
+                    configKeyInfo = convertConfigKey(configKey, prefix);
+                }
+
+                ConfigValue configValue = new ConfigValue(prefix + validatedConfigValue.name(), validatedConfigValue.value(),
+                                                          validatedConfigValue.recommendedValues(), validatedConfigValue.errorMessages());
+                if (configValue.errorMessages().size() > 0) {
+                    errorCount++;
+                }
+                ConfigValueInfo configValueInfo = convertConfigValue(configValue, configKey != null ? configKey.type : null);
+                configInfoList.add(new ConfigInfo(configKeyInfo, configValueInfo));
+            }
+        }
+        return new ConfigInfos(connectorClass.toString(), errorCount, new ArrayList<>(groups), configInfoList);
+    }
+
     // public for testing
     public static ConfigInfos generateResult(String connType, Map configKeys, List configValues, List groups) {
         int errorCount = 0;
@@ -337,7 +475,11 @@ public static ConfigInfos generateResult(String connType, Map
     }
 
     private static ConfigKeyInfo convertConfigKey(ConfigKey configKey) {
-        String name = configKey.name;
+        return convertConfigKey(configKey, "");
+    }
+
+    private static ConfigKeyInfo convertConfigKey(ConfigKey configKey, String prefix) {
+        String name = prefix + configKey.name;
         Type type = configKey.type;
         String typeName = configKey.type.name();
 
@@ -419,7 +561,7 @@ protected final boolean maybeAddConfigErrors(
             callback.onCompletion(
                 new BadRequestException(
                     messages.append(
-                        "\nYou can also find the above list of errors at the endpoint `/{connectorType}/config/validate`"
+                        "\nYou can also find the above list of errors at the endpoint `/connector-plugins/{connectorType}/config/validate`"
                     ).toString()
                 ), null
             );
@@ -461,12 +603,13 @@ public static List> reverseTransform(String connName,
         return result;
     }
 
-    private static Set keysWithVariableValues(Map rawConfig, Pattern pattern) {
+    // Visible for testing
+    static Set keysWithVariableValues(Map rawConfig, Pattern pattern) {
         Set keys = new HashSet<>();
         for (Map.Entry config : rawConfig.entrySet()) {
             if (config.getValue() != null) {
                 Matcher matcher = pattern.matcher(config.getValue());
-                if (matcher.matches()) {
+                if (matcher.find()) {
                     keys.add(config.getKey());
                 }
             }
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/Connect.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/Connect.java
index 965046cccf003..773869a5ba21a 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/Connect.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/Connect.java
@@ -51,6 +51,7 @@ public void start() {
             Runtime.getRuntime().addShutdownHook(shutdownHook);
 
             herder.start();
+            rest.initializeResources(herder);
 
             log.info("Kafka Connect started");
         } finally {
@@ -87,6 +88,10 @@ public URI restUrl() {
         return rest.serverUrl();
     }
 
+    public URI adminUrl() {
+        return rest.adminUrl();
+    }
+
     private class ShutdownHook extends Thread {
         @Override
         public void run() {
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/ConnectMetrics.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/ConnectMetrics.java
index cf82f867baf7d..d09b06d545fc4 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/ConnectMetrics.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/ConnectMetrics.java
@@ -81,7 +81,7 @@ public ConnectMetrics(String workerId, Time time, int numSamples, long sampleWin
         reporters.add(new JmxReporter(JMX_PREFIX));
         this.metrics = new Metrics(metricConfig, reporters, time);
         LOG.debug("Registering Connect metrics with JMX for worker '{}'", workerId);
-        AppInfoParser.registerAppInfo(JMX_PREFIX, workerId, metrics);
+        AppInfoParser.registerAppInfo(JMX_PREFIX, workerId, metrics, time.milliseconds());
     }
 
     /**
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/ConnectMetricsRegistry.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/ConnectMetricsRegistry.java
index 04699ea9ecbfe..6996dacb390a7 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/ConnectMetricsRegistry.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/ConnectMetricsRegistry.java
@@ -20,8 +20,10 @@
 
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.LinkedHashSet;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 
 public class ConnectMetricsRegistry {
@@ -41,6 +43,12 @@ public class ConnectMetricsRegistry {
     public final MetricNameTemplate connectorType;
     public final MetricNameTemplate connectorClass;
     public final MetricNameTemplate connectorVersion;
+    public final MetricNameTemplate connectorTotalTaskCount;
+    public final MetricNameTemplate connectorRunningTaskCount;
+    public final MetricNameTemplate connectorPausedTaskCount;
+    public final MetricNameTemplate connectorFailedTaskCount;
+    public final MetricNameTemplate connectorUnassignedTaskCount;
+    public final MetricNameTemplate connectorDestroyedTaskCount;
     public final MetricNameTemplate taskStatus;
     public final MetricNameTemplate taskRunningRatio;
     public final MetricNameTemplate taskPauseRatio;
@@ -87,6 +95,7 @@ public class ConnectMetricsRegistry {
     public final MetricNameTemplate taskStartupSuccessPercentage;
     public final MetricNameTemplate taskStartupFailureTotal;
     public final MetricNameTemplate taskStartupFailurePercentage;
+    public final MetricNameTemplate connectProtocol;
     public final MetricNameTemplate leaderName;
     public final MetricNameTemplate epoch;
     public final MetricNameTemplate rebalanceCompletedTotal;
@@ -103,6 +112,8 @@ public class ConnectMetricsRegistry {
     public final MetricNameTemplate dlqProduceFailures;
     public final MetricNameTemplate lastErrorTimestamp;
 
+    public Map connectorStatusMetrics;
+
     public ConnectMetricsRegistry() {
         this(new LinkedHashSet());
     }
@@ -288,9 +299,35 @@ public ConnectMetricsRegistry(Set tags) {
         taskStartupFailurePercentage = createTemplate("task-startup-failure-percentage", WORKER_GROUP_NAME,
                                                       "The average percentage of this worker's tasks starts that failed.", workerTags);
 
+        Set workerConnectorTags = new LinkedHashSet<>(tags);
+        workerConnectorTags.add(CONNECTOR_TAG_NAME);
+        connectorTotalTaskCount = createTemplate("connector-total-task-count", WORKER_GROUP_NAME,
+            "The number of tasks of the connector on the worker.", workerConnectorTags);
+        connectorRunningTaskCount = createTemplate("connector-running-task-count", WORKER_GROUP_NAME,
+            "The number of running tasks of the connector on the worker.", workerConnectorTags);
+        connectorPausedTaskCount = createTemplate("connector-paused-task-count", WORKER_GROUP_NAME,
+            "The number of paused tasks of the connector on the worker.", workerConnectorTags);
+        connectorFailedTaskCount = createTemplate("connector-failed-task-count", WORKER_GROUP_NAME,
+            "The number of failed tasks of the connector on the worker.", workerConnectorTags);
+        connectorUnassignedTaskCount = createTemplate("connector-unassigned-task-count",
+            WORKER_GROUP_NAME,
+            "The number of unassigned tasks of the connector on the worker.", workerConnectorTags);
+        connectorDestroyedTaskCount = createTemplate("connector-destroyed-task-count",
+            WORKER_GROUP_NAME,
+            "The number of destroyed tasks of the connector on the worker.", workerConnectorTags);
+
+        connectorStatusMetrics = new HashMap<>();
+        connectorStatusMetrics.put(connectorRunningTaskCount, TaskStatus.State.RUNNING);
+        connectorStatusMetrics.put(connectorPausedTaskCount, TaskStatus.State.PAUSED);
+        connectorStatusMetrics.put(connectorFailedTaskCount, TaskStatus.State.FAILED);
+        connectorStatusMetrics.put(connectorUnassignedTaskCount, TaskStatus.State.UNASSIGNED);
+        connectorStatusMetrics.put(connectorDestroyedTaskCount, TaskStatus.State.DESTROYED);
+        connectorStatusMetrics = Collections.unmodifiableMap(connectorStatusMetrics);
+
         /***** Worker rebalance level *****/
         Set rebalanceTags = new LinkedHashSet<>(tags);
 
+        connectProtocol = createTemplate("connect-protocol", WORKER_REBALANCE_GROUP_NAME, "The Connect protocol used by this cluster", rebalanceTags);
         leaderName = createTemplate("leader-name", WORKER_REBALANCE_GROUP_NAME, "The name of the group leader.", rebalanceTags);
         epoch = createTemplate("epoch", WORKER_REBALANCE_GROUP_NAME, "The epoch or generation number of this worker.", rebalanceTags);
         rebalanceCompletedTotal = createTemplate("completed-rebalances-total", WORKER_REBALANCE_GROUP_NAME,
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/ConnectorConfig.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/ConnectorConfig.java
index 8889aadbae18a..b7340ae1e8e03 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/ConnectorConfig.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/ConnectorConfig.java
@@ -29,6 +29,7 @@
 import org.apache.kafka.connect.runtime.isolation.Plugins;
 import org.apache.kafka.connect.transforms.Transformation;
 
+import java.lang.reflect.Modifier;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
@@ -37,6 +38,8 @@
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 import static org.apache.kafka.common.config.ConfigDef.NonEmptyStringWithoutControlChars.nonEmptyStringWithoutControlChars;
 import static org.apache.kafka.common.config.ConfigDef.Range.atLeast;
@@ -140,6 +143,11 @@ public class ConnectorConfig extends AbstractConfig {
             "a failure. This is 'false' by default, which will prevent record keys, values, and headers from being written to log files, " +
             "although some information such as topic and partition number will still be logged.";
 
+
+    public static final String CONNECTOR_CLIENT_PRODUCER_OVERRIDES_PREFIX = "producer.override.";
+    public static final String CONNECTOR_CLIENT_CONSUMER_OVERRIDES_PREFIX = "consumer.override.";
+    public static final String CONNECTOR_CLIENT_ADMIN_OVERRIDES_PREFIX = "admin.override.";
+
     private final EnrichedConnectorConfig enrichedConfig;
     private static class EnrichedConnectorConfig extends AbstractConfig {
         EnrichedConnectorConfig(ConfigDef configDef, Map props) {
@@ -324,11 +332,25 @@ static ConfigDef getConfigDefFromTransformation(String key, Class transformat
         if (transformationCls == null || !Transformation.class.isAssignableFrom(transformationCls)) {
             throw new ConfigException(key, String.valueOf(transformationCls), "Not a Transformation");
         }
+        if (Modifier.isAbstract(transformationCls.getModifiers())) {
+            String childClassNames = Stream.of(transformationCls.getClasses())
+                .filter(transformationCls::isAssignableFrom)
+                .filter(c -> !Modifier.isAbstract(c.getModifiers()))
+                .filter(c -> Modifier.isPublic(c.getModifiers()))
+                .map(Class::getName)
+                .collect(Collectors.joining(", "));
+            String message = childClassNames.trim().isEmpty() ?
+                "Transformation is abstract and cannot be created." :
+                "Transformation is abstract and cannot be created. Did you mean " + childClassNames + "?";
+            throw new ConfigException(key, String.valueOf(transformationCls), message);
+        }
         Transformation transformation;
         try {
             transformation = transformationCls.asSubclass(Transformation.class).newInstance();
         } catch (Exception e) {
-            throw new ConfigException(key, String.valueOf(transformationCls), "Error getting config definition from Transformation: " + e.getMessage());
+            ConfigException exception = new ConfigException(key, String.valueOf(transformationCls), "Error getting config definition from Transformation: " + e.getMessage());
+            exception.initCause(e);
+            throw exception;
         }
         ConfigDef configDef = transformation.config();
         if (null == configDef) {
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/Herder.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/Herder.java
index c572e20b52f2b..7915682790317 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/Herder.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/Herder.java
@@ -17,6 +17,7 @@
 package org.apache.kafka.connect.runtime;
 
 import org.apache.kafka.connect.runtime.isolation.Plugins;
+import org.apache.kafka.connect.runtime.rest.InternalRequestSignature;
 import org.apache.kafka.connect.runtime.rest.entities.ConfigInfos;
 import org.apache.kafka.connect.runtime.rest.entities.ConnectorInfo;
 import org.apache.kafka.connect.runtime.rest.entities.ConnectorStateInfo;
@@ -120,8 +121,22 @@ public interface Herder {
      * @param connName connector to update
      * @param configs list of configurations
      * @param callback callback to invoke upon completion
+     * @param requestSignature the signature of the request made for this task (re-)configuration;
+     *                         may be null if no signature was provided
      */
-    void putTaskConfigs(String connName, List> configs, Callback callback);
+    void putTaskConfigs(String connName, List> configs, Callback callback, InternalRequestSignature requestSignature);
+
+    /**
+     * Get a list of connectors currently running in this cluster.
+     * @returns A list of connector names
+     */
+    Collection connectors();
+
+    /**
+     * Get the definition and status of a connector.
+     * @param connName name of the connector
+     */
+    ConnectorInfo connectorInfo(String connName);
 
     /**
      * Lookup the current status of a connector.
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/HerderProvider.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/HerderProvider.java
deleted file mode 100644
index 42c0925a704a2..0000000000000
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/HerderProvider.java
+++ /dev/null
@@ -1,68 +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.connect.errors.ConnectException;
-
-import java.util.concurrent.CountDownLatch;
-import java.util.concurrent.TimeUnit;
-
-/**
- * A supplier for {@link Herder}s.
- */
-public class HerderProvider {
-
-    private final CountDownLatch initialized = new CountDownLatch(1);
-    volatile Herder herder = null;
-
-    public HerderProvider() {
-    }
-
-    /**
-     * Create a herder provider with a herder.
-     * @param herder the herder that will be supplied to threads waiting on this provider
-     */
-    public HerderProvider(Herder herder) {
-        this.herder = herder;
-        initialized.countDown();
-    }
-
-    /**
-     * @return the contained herder.
-     * @throws ConnectException if a herder was not available within a duration of calling this method
-     */
-    public Herder get() {
-        try {
-            // wait for herder to be initialized
-            if (!initialized.await(1, TimeUnit.MINUTES)) {
-                throw new ConnectException("Timed out waiting for herder to be initialized.");
-            }
-        } catch (InterruptedException e) {
-            throw new ConnectException("Interrupted while waiting for herder to be initialized.", e);
-        }
-        return herder;
-    }
-
-    /**
-     * @param herder set a herder, and signal to all threads waiting on get().
-     */
-    public void setHerder(Herder herder) {
-        this.herder = herder;
-        initialized.countDown();
-    }
-
-}
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/SessionKey.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/SessionKey.java
new file mode 100644
index 0000000000000..ab5476e4b9000
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/SessionKey.java
@@ -0,0 +1,73 @@
+/*
+ * 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 javax.crypto.SecretKey;
+import java.util.Objects;
+
+/**
+ * A session key, which can be used to validate internal REST requests between workers.
+ */
+public class SessionKey {
+
+    private final SecretKey key;
+    private final long creationTimestamp;
+
+    /**
+     * Create a new session key with the given key value and creation timestamp
+     * @param key the actual cryptographic key to use for request validation; may not be null
+     * @param creationTimestamp the time at which the key was generated
+     */
+    public SessionKey(SecretKey key, long creationTimestamp) {
+        this.key = Objects.requireNonNull(key, "Key may not be null");
+        this.creationTimestamp = creationTimestamp;
+    }
+
+    /**
+     * Get the cryptographic key to use for request validation.
+     *
+     * @return the cryptographic key; may not be null
+     */
+    public SecretKey key() {
+        return key;
+    }
+
+    /**
+     * Get the time at which the key was generated.
+     *
+     * @return the time at which the key was generated
+     */
+    public long creationTimestamp() {
+        return creationTimestamp;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o)
+            return true;
+        if (o == null || getClass() != o.getClass())
+            return false;
+        SessionKey that = (SessionKey) o;
+        return creationTimestamp == that.creationTimestamp
+            && key.equals(that.key);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(key, creationTimestamp);
+    }
+}
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/SinkConnectorConfig.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/SinkConnectorConfig.java
index d8912a1350213..0ed1b2b9608c8 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/SinkConnectorConfig.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/SinkConnectorConfig.java
@@ -18,12 +18,17 @@
 
 import org.apache.kafka.common.config.ConfigDef;
 import org.apache.kafka.common.config.ConfigDef.Importance;
+import org.apache.kafka.common.config.ConfigDef.Type;
 import org.apache.kafka.common.config.ConfigException;
 import org.apache.kafka.connect.runtime.isolation.Plugins;
 import org.apache.kafka.connect.sink.SinkTask;
 import org.apache.kafka.connect.transforms.util.RegexValidator;
 
+import java.util.Collections;
+import java.util.List;
 import java.util.Map;
+import java.util.regex.Pattern;
+import java.util.stream.Collectors;
 
 /**
  * Configuration needed for all sink connectors
@@ -86,6 +91,7 @@ public SinkConnectorConfig(Plugins plugins, Map props) {
     public static void validate(Map props) {
         final boolean hasTopicsConfig = hasTopicsConfig(props);
         final boolean hasTopicsRegexConfig = hasTopicsRegexConfig(props);
+        final boolean hasDlqTopicConfig = hasDlqTopicConfig(props);
 
         if (hasTopicsConfig && hasTopicsRegexConfig) {
             throw new ConfigException(SinkTask.TOPICS_CONFIG + " and " + SinkTask.TOPICS_REGEX_CONFIG +
@@ -96,6 +102,25 @@ public static void validate(Map props) {
             throw new ConfigException("Must configure one of " +
                 SinkTask.TOPICS_CONFIG + " or " + SinkTask.TOPICS_REGEX_CONFIG);
         }
+
+        if (hasDlqTopicConfig) {
+            String dlqTopic = props.get(DLQ_TOPIC_NAME_CONFIG).trim();
+            if (hasTopicsConfig) {
+                List topics = parseTopicsList(props);
+                if (topics.contains(dlqTopic)) {
+                    throw new ConfigException(String.format("The DLQ topic '%s' may not be included in the list of "
+                            + "topics ('%s=%s') consumed by the connector", dlqTopic, SinkTask.TOPICS_REGEX_CONFIG, topics));
+                }
+            }
+            if (hasTopicsRegexConfig) {
+                String topicsRegexStr = props.get(SinkTask.TOPICS_REGEX_CONFIG);
+                Pattern pattern = Pattern.compile(topicsRegexStr);
+                if (pattern.matcher(dlqTopic).matches()) {
+                    throw new ConfigException(String.format("The DLQ topic '%s' may not be included in the regex matching the "
+                            + "topics ('%s=%s') consumed by the connector", dlqTopic, SinkTask.TOPICS_REGEX_CONFIG, topicsRegexStr));
+                }
+            }
+        }
     }
 
     public static boolean hasTopicsConfig(Map props) {
@@ -108,6 +133,24 @@ public static boolean hasTopicsRegexConfig(Map props) {
         return topicsRegexStr != null && !topicsRegexStr.trim().isEmpty();
     }
 
+    public static boolean hasDlqTopicConfig(Map props) {
+        String dqlTopicStr = props.get(DLQ_TOPIC_NAME_CONFIG);
+        return dqlTopicStr != null && !dqlTopicStr.trim().isEmpty();
+    }
+
+    @SuppressWarnings("unchecked")
+    public static List parseTopicsList(Map props) {
+        List topics = (List) ConfigDef.parseType(TOPICS_CONFIG, props.get(TOPICS_CONFIG), Type.LIST);
+        if (topics == null) {
+            return Collections.emptyList();
+        }
+        return topics
+                .stream()
+                .filter(topic -> !topic.isEmpty())
+                .distinct()
+                .collect(Collectors.toList());
+    }
+
     public String dlqTopicName() {
         return getString(DLQ_TOPIC_NAME_CONFIG);
     }
@@ -121,6 +164,6 @@ public boolean isDlqContextHeadersEnabled() {
     }
 
     public static void main(String[] args) {
-        System.out.println(config.toHtmlTable());
+        System.out.println(config.toHtml());
     }
 }
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/SourceConnectorConfig.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/SourceConnectorConfig.java
index ad891b6945033..a0ee46893acaa 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/SourceConnectorConfig.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/SourceConnectorConfig.java
@@ -34,6 +34,6 @@ public SourceConnectorConfig(Plugins plugins, Map props) {
     }
 
     public static void main(String[] args) {
-        System.out.println(config.toHtmlTable());
+        System.out.println(config.toHtml());
     }
 }
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/SourceTaskOffsetCommitter.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/SourceTaskOffsetCommitter.java
index c50280932a4cc..8e8d3fa056cfd 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/SourceTaskOffsetCommitter.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/SourceTaskOffsetCommitter.java
@@ -18,6 +18,7 @@
 
 import org.apache.kafka.connect.errors.ConnectException;
 import org.apache.kafka.connect.util.ConnectorTaskId;
+import org.apache.kafka.connect.util.LoggingContext;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -79,7 +80,9 @@ public void schedule(final ConnectorTaskId id, final WorkerSourceTask workerTask
         ScheduledFuture commitFuture = commitExecutorService.scheduleWithFixedDelay(new Runnable() {
             @Override
             public void run() {
-                commit(workerTask);
+                try (LoggingContext loggingContext = LoggingContext.forOffsets(id)) {
+                    commit(workerTask);
+                }
             }
         }, commitIntervalMs, commitIntervalMs, TimeUnit.MILLISECONDS);
         committers.put(id, commitFuture);
@@ -90,7 +93,7 @@ public void remove(ConnectorTaskId id) {
         if (task == null)
             return;
 
-        try {
+        try (LoggingContext loggingContext = LoggingContext.forTask(id)) {
             task.cancel(false);
             if (!task.isDone())
                 task.get();
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/TaskStatus.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/TaskStatus.java
index 5ee90415b9660..62bb1c717112b 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/TaskStatus.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/TaskStatus.java
@@ -61,5 +61,12 @@ public interface Listener {
          */
         void onShutdown(ConnectorTaskId id);
 
+        /**
+         * Invoked after the task has been deleted. Can be called if the
+         * connector tasks have been reduced, or if the connector itself has
+         * been deleted.
+         * @param id The id of the task
+         */
+        void onDeletion(ConnectorTaskId id);
     }
 }
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 673bd4e0c3d03..3cf4d01ac1cd6 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
@@ -16,22 +16,27 @@
  */
 package org.apache.kafka.connect.runtime;
 
+import org.apache.kafka.clients.admin.AdminClientConfig;
 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.common.MetricName;
+import org.apache.kafka.common.MetricNameTemplate;
+import org.apache.kafka.common.config.ConfigValue;
 import org.apache.kafka.common.config.provider.ConfigProvider;
 import org.apache.kafka.common.metrics.Sensor;
+import org.apache.kafka.common.metrics.stats.CumulativeSum;
 import org.apache.kafka.common.metrics.stats.Frequencies;
-import org.apache.kafka.common.metrics.stats.Total;
 import org.apache.kafka.common.utils.Time;
 import org.apache.kafka.common.utils.Utils;
 import org.apache.kafka.connect.connector.Connector;
 import org.apache.kafka.connect.connector.ConnectorContext;
 import org.apache.kafka.connect.connector.Task;
+import org.apache.kafka.connect.connector.policy.ConnectorClientConfigOverridePolicy;
+import org.apache.kafka.connect.connector.policy.ConnectorClientConfigRequest;
 import org.apache.kafka.connect.errors.ConnectException;
-import org.apache.kafka.connect.runtime.ConnectMetrics.LiteralSupplier;
+import org.apache.kafka.connect.health.ConnectorType;
 import org.apache.kafka.connect.runtime.ConnectMetrics.MetricGroup;
 import org.apache.kafka.connect.runtime.distributed.ClusterConfigState;
 import org.apache.kafka.connect.runtime.errors.DeadLetterQueueReporter;
@@ -45,13 +50,14 @@
 import org.apache.kafka.connect.sink.SinkTask;
 import org.apache.kafka.connect.source.SourceRecord;
 import org.apache.kafka.connect.source.SourceTask;
+import org.apache.kafka.connect.storage.CloseableOffsetStorageReader;
 import org.apache.kafka.connect.storage.Converter;
 import org.apache.kafka.connect.storage.HeaderConverter;
 import org.apache.kafka.connect.storage.OffsetBackingStore;
-import org.apache.kafka.connect.storage.OffsetStorageReader;
 import org.apache.kafka.connect.storage.OffsetStorageReaderImpl;
 import org.apache.kafka.connect.storage.OffsetStorageWriter;
 import org.apache.kafka.connect.util.ConnectorTaskId;
+import org.apache.kafka.connect.util.LoggingContext;
 import org.apache.kafka.connect.util.SinkUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -67,6 +73,7 @@
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
+import java.util.stream.Collectors;
 
 
 /**
@@ -88,6 +95,7 @@ public class Worker {
     private final Plugins plugins;
     private final ConnectMetrics metrics;
     private final WorkerMetricsGroup workerMetricsGroup;
+    private ConnectorStatusMetricsGroup connectorStatusMetricsGroup;
     private final WorkerConfig config;
     private final Converter internalKeyConverter;
     private final Converter internalValueConverter;
@@ -96,22 +104,36 @@ public class Worker {
     private final ConcurrentMap connectors = new ConcurrentHashMap<>();
     private final ConcurrentMap tasks = new ConcurrentHashMap<>();
     private SourceTaskOffsetCommitter sourceTaskOffsetCommitter;
-    private WorkerConfigTransformer workerConfigTransformer;
+    private final WorkerConfigTransformer workerConfigTransformer;
+    private final ConnectorClientConfigOverridePolicy connectorClientConfigOverridePolicy;
 
-    @SuppressWarnings("deprecation")
     public Worker(
+        String workerId,
+        Time time,
+        Plugins plugins,
+        WorkerConfig config,
+        OffsetBackingStore offsetBackingStore,
+        ConnectorClientConfigOverridePolicy connectorClientConfigOverridePolicy) {
+        this(workerId, time, plugins, config, offsetBackingStore, Executors.newCachedThreadPool(), connectorClientConfigOverridePolicy);
+    }
+
+    @SuppressWarnings("deprecation")
+    Worker(
             String workerId,
             Time time,
             Plugins plugins,
             WorkerConfig config,
-            OffsetBackingStore offsetBackingStore
+            OffsetBackingStore offsetBackingStore,
+            ExecutorService executorService,
+            ConnectorClientConfigOverridePolicy connectorClientConfigOverridePolicy
     ) {
         this.metrics = new ConnectMetrics(workerId, config, time);
-        this.executor = Executors.newCachedThreadPool();
+        this.executor = executorService;
         this.workerId = workerId;
         this.time = time;
         this.plugins = plugins;
         this.config = config;
+        this.connectorClientConfigOverridePolicy = connectorClientConfigOverridePolicy;
         this.workerMetricsGroup = new WorkerMetricsGroup(metrics);
 
         // Internal converters are required properties, thus getClass won't return null.
@@ -164,6 +186,8 @@ public void start() {
         offsetBackingStore.start();
         sourceTaskOffsetCommitter = new SourceTaskOffsetCommitter(config);
 
+        connectorStatusMetricsGroup = new ConnectorStatusMetricsGroup(metrics, tasks, herder);
+
         log.info("Worker started");
     }
 
@@ -195,6 +219,9 @@ public void stop() {
         log.info("Worker stopped");
 
         workerMetricsGroup.close();
+        connectorStatusMetricsGroup.close();
+
+        workerConfigTransformer.close();
     }
 
     /**
@@ -214,38 +241,40 @@ public boolean startConnector(
             ConnectorStatus.Listener statusListener,
             TargetState initialState
     ) {
-        if (connectors.containsKey(connName))
-            throw new ConnectException("Connector with name " + connName + " already exists");
-
-        final WorkerConnector workerConnector;
-        ClassLoader savedLoader = plugins.currentThreadLoader();
-        try {
-            final ConnectorConfig connConfig = new ConnectorConfig(plugins, connProps);
-            final String connClass = connConfig.getString(ConnectorConfig.CONNECTOR_CLASS_CONFIG);
-            log.info("Creating connector {} of type {}", connName, connClass);
-            final Connector connector = plugins.newConnector(connClass);
-            workerConnector = new WorkerConnector(connName, connector, ctx, metrics,  statusListener);
-            log.info("Instantiated connector {} with version {} of type {}", connName, connector.version(), connector.getClass());
-            savedLoader = plugins.compareAndSwapLoaders(connector);
-            workerConnector.initialize(connConfig);
-            workerConnector.transitionTo(initialState);
-            Plugins.compareAndSwapLoaders(savedLoader);
-        } catch (Throwable t) {
-            log.error("Failed to start connector {}", connName, t);
-            // Can't be put in a finally block because it needs to be swapped before the call on
-            // statusListener
-            Plugins.compareAndSwapLoaders(savedLoader);
-            workerMetricsGroup.recordConnectorStartupFailure();
-            statusListener.onFailure(connName, t);
-            return false;
-        }
+        try (LoggingContext loggingContext = LoggingContext.forConnector(connName)) {
+            if (connectors.containsKey(connName))
+                throw new ConnectException("Connector with name " + connName + " already exists");
+
+            final WorkerConnector workerConnector;
+            ClassLoader savedLoader = plugins.currentThreadLoader();
+            try {
+                final ConnectorConfig connConfig = new ConnectorConfig(plugins, connProps);
+                final String connClass = connConfig.getString(ConnectorConfig.CONNECTOR_CLASS_CONFIG);
+                log.info("Creating connector {} of type {}", connName, connClass);
+                final Connector connector = plugins.newConnector(connClass);
+                workerConnector = new WorkerConnector(connName, connector, ctx, metrics, statusListener);
+                log.info("Instantiated connector {} with version {} of type {}", connName, connector.version(), connector.getClass());
+                savedLoader = plugins.compareAndSwapLoaders(connector);
+                workerConnector.initialize(connConfig);
+                workerConnector.transitionTo(initialState);
+                Plugins.compareAndSwapLoaders(savedLoader);
+            } catch (Throwable t) {
+                log.error("Failed to start connector {}", connName, t);
+                // Can't be put in a finally block because it needs to be swapped before the call on
+                // statusListener
+                Plugins.compareAndSwapLoaders(savedLoader);
+                workerMetricsGroup.recordConnectorStartupFailure();
+                statusListener.onFailure(connName, t);
+                return false;
+            }
 
-        WorkerConnector existing = connectors.putIfAbsent(connName, workerConnector);
-        if (existing != null)
-            throw new ConnectException("Connector with name " + connName + " already exists");
+            WorkerConnector existing = connectors.putIfAbsent(connName, workerConnector);
+            if (existing != null)
+                throw new ConnectException("Connector with name " + connName + " already exists");
 
-        log.info("Finished creating connector {}", connName);
-        workerMetricsGroup.recordConnectorStartupSuccess();
+            log.info("Finished creating connector {}", connName);
+            workerMetricsGroup.recordConnectorStartupSuccess();
+        }
         return true;
     }
 
@@ -277,35 +306,37 @@ public boolean isSinkConnector(String connName) {
      * @return a list of updated tasks properties.
      */
     public List> connectorTaskConfigs(String connName, ConnectorConfig connConfig) {
-        log.trace("Reconfiguring connector tasks for {}", connName);
-
-        WorkerConnector workerConnector = connectors.get(connName);
-        if (workerConnector == null)
-            throw new ConnectException("Connector " + connName + " not found in this worker.");
-
-        int maxTasks = connConfig.getInt(ConnectorConfig.TASKS_MAX_CONFIG);
-        Map connOriginals = connConfig.originalsStrings();
-
-        Connector connector = workerConnector.connector();
         List> result = new ArrayList<>();
-        ClassLoader savedLoader = plugins.currentThreadLoader();
-        try {
-            savedLoader = plugins.compareAndSwapLoaders(connector);
-            String taskClassName = connector.taskClass().getName();
-            for (Map taskProps : connector.taskConfigs(maxTasks)) {
-                // Ensure we don't modify the connector's copy of the config
-                Map taskConfig = new HashMap<>(taskProps);
-                taskConfig.put(TaskConfig.TASK_CLASS_CONFIG, taskClassName);
-                if (connOriginals.containsKey(SinkTask.TOPICS_CONFIG)) {
-                    taskConfig.put(SinkTask.TOPICS_CONFIG, connOriginals.get(SinkTask.TOPICS_CONFIG));
-                }
-                if (connOriginals.containsKey(SinkTask.TOPICS_REGEX_CONFIG)) {
-                    taskConfig.put(SinkTask.TOPICS_REGEX_CONFIG, connOriginals.get(SinkTask.TOPICS_REGEX_CONFIG));
+        try (LoggingContext loggingContext = LoggingContext.forConnector(connName)) {
+            log.trace("Reconfiguring connector tasks for {}", connName);
+
+            WorkerConnector workerConnector = connectors.get(connName);
+            if (workerConnector == null)
+                throw new ConnectException("Connector " + connName + " not found in this worker.");
+
+            int maxTasks = connConfig.getInt(ConnectorConfig.TASKS_MAX_CONFIG);
+            Map connOriginals = connConfig.originalsStrings();
+
+            Connector connector = workerConnector.connector();
+            ClassLoader savedLoader = plugins.currentThreadLoader();
+            try {
+                savedLoader = plugins.compareAndSwapLoaders(connector);
+                String taskClassName = connector.taskClass().getName();
+                for (Map taskProps : connector.taskConfigs(maxTasks)) {
+                    // Ensure we don't modify the connector's copy of the config
+                    Map taskConfig = new HashMap<>(taskProps);
+                    taskConfig.put(TaskConfig.TASK_CLASS_CONFIG, taskClassName);
+                    if (connOriginals.containsKey(SinkTask.TOPICS_CONFIG)) {
+                        taskConfig.put(SinkTask.TOPICS_CONFIG, connOriginals.get(SinkTask.TOPICS_CONFIG));
+                    }
+                    if (connOriginals.containsKey(SinkTask.TOPICS_REGEX_CONFIG)) {
+                        taskConfig.put(SinkTask.TOPICS_REGEX_CONFIG, connOriginals.get(SinkTask.TOPICS_REGEX_CONFIG));
+                    }
+                    result.add(taskConfig);
                 }
-                result.add(taskConfig);
+            } finally {
+                Plugins.compareAndSwapLoaders(savedLoader);
             }
-        } finally {
-            Plugins.compareAndSwapLoaders(savedLoader);
         }
 
         return result;
@@ -325,23 +356,25 @@ private void stopConnectors() {
      * @return true if the connector belonged to this worker and was successfully stopped.
      */
     public boolean stopConnector(String connName) {
-        log.info("Stopping connector {}", connName);
+        try (LoggingContext loggingContext = LoggingContext.forConnector(connName)) {
+            log.info("Stopping connector {}", connName);
 
-        WorkerConnector workerConnector = connectors.remove(connName);
-        if (workerConnector == null) {
-            log.warn("Ignoring stop request for unowned connector {}", connName);
-            return false;
-        }
+            WorkerConnector workerConnector = connectors.remove(connName);
+            if (workerConnector == null) {
+                log.warn("Ignoring stop request for unowned connector {}", connName);
+                return false;
+            }
 
-        ClassLoader savedLoader = plugins.currentThreadLoader();
-        try {
-            savedLoader = plugins.compareAndSwapLoaders(workerConnector.connector());
-            workerConnector.shutdown();
-        } finally {
-            Plugins.compareAndSwapLoaders(savedLoader);
-        }
+            ClassLoader savedLoader = plugins.currentThreadLoader();
+            try {
+                savedLoader = plugins.compareAndSwapLoaders(workerConnector.connector());
+                workerConnector.shutdown();
+            } finally {
+                Plugins.compareAndSwapLoaders(savedLoader);
+            }
 
-        log.info("Stopped connector {}", connName);
+            log.info("Stopped connector {}", connName);
+        }
         return true;
     }
 
@@ -383,84 +416,80 @@ public boolean startTask(
             TaskStatus.Listener statusListener,
             TargetState initialState
     ) {
-        log.info("Creating task {}", id);
-
-        if (tasks.containsKey(id))
-            throw new ConnectException("Task already exists in this worker: " + id);
-
         final WorkerTask workerTask;
-        ClassLoader savedLoader = plugins.currentThreadLoader();
-        try {
-            final ConnectorConfig connConfig = new ConnectorConfig(plugins, connProps);
-            String connType = connConfig.getString(ConnectorConfig.CONNECTOR_CLASS_CONFIG);
-            ClassLoader connectorLoader = plugins.delegatingLoader().connectorLoader(connType);
-            savedLoader = Plugins.compareAndSwapLoaders(connectorLoader);
-            final TaskConfig taskConfig = new TaskConfig(taskProps);
-            final Class taskClass = taskConfig.getClass(TaskConfig.TASK_CLASS_CONFIG).asSubclass(Task.class);
-            final Task task = plugins.newTask(taskClass);
-            log.info("Instantiated task {} with version {} of type {}", id, task.version(), taskClass.getName());
-
-            // By maintaining connector's specific class loader for this thread here, we first
-            // search for converters within the connector dependencies.
-            // If any of these aren't found, that means the connector didn't configure specific converters,
-            // so we should instantiate based upon the worker configuration
-            Converter keyConverter = plugins.newConverter(
-                    connConfig,
-                    WorkerConfig.KEY_CONVERTER_CLASS_CONFIG,
-                    ClassLoaderUsage.CURRENT_CLASSLOADER
-            );
-            Converter valueConverter = plugins.newConverter(
-                    connConfig,
-                    WorkerConfig.VALUE_CONVERTER_CLASS_CONFIG,
-                    ClassLoaderUsage.CURRENT_CLASSLOADER
-            );
-            HeaderConverter headerConverter = plugins.newHeaderConverter(
-                    connConfig,
-                    WorkerConfig.HEADER_CONVERTER_CLASS_CONFIG,
-                    ClassLoaderUsage.CURRENT_CLASSLOADER
-            );
-            if (keyConverter == null) {
-                keyConverter = plugins.newConverter(config, WorkerConfig.KEY_CONVERTER_CLASS_CONFIG, ClassLoaderUsage.PLUGINS);
-                log.info("Set up the key converter {} for task {} using the worker config", keyConverter.getClass(), id);
-            } else {
-                log.info("Set up the key converter {} for task {} using the connector config", keyConverter.getClass(), id);
-            }
-            if (valueConverter == null) {
-                valueConverter = plugins.newConverter(config, WorkerConfig.VALUE_CONVERTER_CLASS_CONFIG, ClassLoaderUsage.PLUGINS);
-                log.info("Set up the value converter {} for task {} using the worker config", valueConverter.getClass(), id);
-            } else {
-                log.info("Set up the value converter {} for task {} using the connector config", valueConverter.getClass(), id);
-            }
-            if (headerConverter == null) {
-                headerConverter = plugins.newHeaderConverter(config, WorkerConfig.HEADER_CONVERTER_CLASS_CONFIG, ClassLoaderUsage.PLUGINS);
-                log.info("Set up the header converter {} for task {} using the worker config", headerConverter.getClass(), id);
-            } else {
-                log.info("Set up the header converter {} for task {} using the connector config", headerConverter.getClass(), id);
-            }
+        try (LoggingContext loggingContext = LoggingContext.forTask(id)) {
+            log.info("Creating task {}", id);
+
+            if (tasks.containsKey(id))
+                throw new ConnectException("Task already exists in this worker: " + id);
+
+            connectorStatusMetricsGroup.recordTaskAdded(id);
+            ClassLoader savedLoader = plugins.currentThreadLoader();
+            try {
+                String connType = connProps.get(ConnectorConfig.CONNECTOR_CLASS_CONFIG);
+                ClassLoader connectorLoader = plugins.delegatingLoader().connectorLoader(connType);
+                savedLoader = Plugins.compareAndSwapLoaders(connectorLoader);
+                final ConnectorConfig connConfig = new ConnectorConfig(plugins, connProps);
+                final TaskConfig taskConfig = new TaskConfig(taskProps);
+                final Class taskClass = taskConfig.getClass(TaskConfig.TASK_CLASS_CONFIG).asSubclass(Task.class);
+                final Task task = plugins.newTask(taskClass);
+                log.info("Instantiated task {} with version {} of type {}", id, task.version(), taskClass.getName());
+
+                // By maintaining connector's specific class loader for this thread here, we first
+                // search for converters within the connector dependencies.
+                // If any of these aren't found, that means the connector didn't configure specific converters,
+                // so we should instantiate based upon the worker configuration
+                Converter keyConverter = plugins.newConverter(connConfig, WorkerConfig.KEY_CONVERTER_CLASS_CONFIG, ClassLoaderUsage
+                                                                                                                           .CURRENT_CLASSLOADER);
+                Converter valueConverter = plugins.newConverter(connConfig, WorkerConfig.VALUE_CONVERTER_CLASS_CONFIG, ClassLoaderUsage.CURRENT_CLASSLOADER);
+                HeaderConverter headerConverter = plugins.newHeaderConverter(connConfig, WorkerConfig.HEADER_CONVERTER_CLASS_CONFIG,
+                                                                             ClassLoaderUsage.CURRENT_CLASSLOADER);
+                if (keyConverter == null) {
+                    keyConverter = plugins.newConverter(config, WorkerConfig.KEY_CONVERTER_CLASS_CONFIG, ClassLoaderUsage.PLUGINS);
+                    log.info("Set up the key converter {} for task {} using the worker config", keyConverter.getClass(), id);
+                } else {
+                    log.info("Set up the key converter {} for task {} using the connector config", keyConverter.getClass(), id);
+                }
+                if (valueConverter == null) {
+                    valueConverter = plugins.newConverter(config, WorkerConfig.VALUE_CONVERTER_CLASS_CONFIG, ClassLoaderUsage.PLUGINS);
+                    log.info("Set up the value converter {} for task {} using the worker config", valueConverter.getClass(), id);
+                } else {
+                    log.info("Set up the value converter {} for task {} using the connector config", valueConverter.getClass(), id);
+                }
+                if (headerConverter == null) {
+                    headerConverter = plugins.newHeaderConverter(config, WorkerConfig.HEADER_CONVERTER_CLASS_CONFIG, ClassLoaderUsage
+                                                                                                                             .PLUGINS);
+                    log.info("Set up the header converter {} for task {} using the worker config", headerConverter.getClass(), id);
+                } else {
+                    log.info("Set up the header converter {} for task {} using the connector config", headerConverter.getClass(), id);
+                }
 
-            workerTask = buildWorkerTask(configState, connConfig, id, task, statusListener, initialState, keyConverter, valueConverter, headerConverter, connectorLoader);
-            workerTask.initialize(taskConfig);
-            Plugins.compareAndSwapLoaders(savedLoader);
-        } catch (Throwable t) {
-            log.error("Failed to start task {}", id, t);
-            // Can't be put in a finally block because it needs to be swapped before the call on
-            // statusListener
-            Plugins.compareAndSwapLoaders(savedLoader);
-            workerMetricsGroup.recordTaskFailure();
-            statusListener.onFailure(id, t);
-            return false;
-        }
+                workerTask = buildWorkerTask(configState, connConfig, id, task, statusListener, initialState, keyConverter, valueConverter,
+                                             headerConverter, connectorLoader);
+                workerTask.initialize(taskConfig);
+                Plugins.compareAndSwapLoaders(savedLoader);
+            } catch (Throwable t) {
+                log.error("Failed to start task {}", id, t);
+                // Can't be put in a finally block because it needs to be swapped before the call on
+                // statusListener
+                Plugins.compareAndSwapLoaders(savedLoader);
+                connectorStatusMetricsGroup.recordTaskRemoved(id);
+                workerMetricsGroup.recordTaskFailure();
+                statusListener.onFailure(id, t);
+                return false;
+            }
 
-        WorkerTask existing = tasks.putIfAbsent(id, workerTask);
-        if (existing != null)
-            throw new ConnectException("Task already exists in this worker: " + id);
+            WorkerTask existing = tasks.putIfAbsent(id, workerTask);
+            if (existing != null)
+                throw new ConnectException("Task already exists in this worker: " + id);
 
-        executor.submit(workerTask);
-        if (workerTask instanceof WorkerSourceTask) {
-            sourceTaskOffsetCommitter.schedule(id, (WorkerSourceTask) workerTask);
+            executor.submit(workerTask);
+            if (workerTask instanceof WorkerSourceTask) {
+                sourceTaskOffsetCommitter.schedule(id, (WorkerSourceTask) workerTask);
+            }
+            workerMetricsGroup.recordTaskSuccess();
+            return true;
         }
-        workerMetricsGroup.recordTaskSuccess();
-        return true;
     }
 
     private WorkerTask buildWorkerTask(ClusterConfigState configState,
@@ -474,7 +503,8 @@ private WorkerTask buildWorkerTask(ClusterConfigState configState,
                                        HeaderConverter headerConverter,
                                        ClassLoader loader) {
         ErrorHandlingMetrics errorHandlingMetrics = errorHandlingMetrics(id);
-
+        final Class connectorClass = plugins.connectorClass(
+            connConfig.getString(ConnectorConfig.CONNECTOR_CLASS_CONFIG));
         RetryWithToleranceOperator retryWithToleranceOperator = new RetryWithToleranceOperator(connConfig.errorRetryTimeout(),
                 connConfig.errorMaxDelayInMillis(), connConfig.errorToleranceType(), Time.SYSTEM);
         retryWithToleranceOperator.metrics(errorHandlingMetrics);
@@ -484,11 +514,12 @@ private WorkerTask buildWorkerTask(ClusterConfigState configState,
             retryWithToleranceOperator.reporters(sourceTaskReporters(id, connConfig, errorHandlingMetrics));
             TransformationChain transformationChain = new TransformationChain<>(connConfig.transformations(), retryWithToleranceOperator);
             log.info("Initializing: {}", transformationChain);
-            OffsetStorageReader offsetReader = new OffsetStorageReaderImpl(offsetBackingStore, id.connector(),
+            CloseableOffsetStorageReader offsetReader = new OffsetStorageReaderImpl(offsetBackingStore, id.connector(),
                     internalKeyConverter, internalValueConverter);
             OffsetStorageWriter offsetWriter = new OffsetStorageWriter(offsetBackingStore, id.connector(),
                     internalKeyConverter, internalValueConverter);
-            Map producerProps = producerConfigs(config);
+            Map producerProps = producerConfigs(id, "connector-producer-" + id, config, connConfig, connectorClass,
+                                                                connectorClientConfigOverridePolicy);
             KafkaProducer producer = new KafkaProducer<>(producerProps);
 
             // Note we pass the configState as it performs dynamic transformations under the covers
@@ -499,44 +530,62 @@ private WorkerTask buildWorkerTask(ClusterConfigState configState,
             TransformationChain transformationChain = new TransformationChain<>(connConfig.transformations(), retryWithToleranceOperator);
             log.info("Initializing: {}", transformationChain);
             SinkConnectorConfig sinkConfig = new SinkConnectorConfig(plugins, connConfig.originalsStrings());
-            retryWithToleranceOperator.reporters(sinkTaskReporters(id, sinkConfig, errorHandlingMetrics));
+            retryWithToleranceOperator.reporters(sinkTaskReporters(id, sinkConfig, errorHandlingMetrics, connectorClass));
 
-            Map consumerProps = consumerConfigs(id, config);
+            Map consumerProps = consumerConfigs(id, config, connConfig, connectorClass, connectorClientConfigOverridePolicy);
             KafkaConsumer consumer = new KafkaConsumer<>(consumerProps);
 
             return new WorkerSinkTask(id, (SinkTask) task, statusListener, initialState, config, configState, metrics, keyConverter,
                                       valueConverter, headerConverter, transformationChain, consumer, loader, time,
                                       retryWithToleranceOperator);
         } else {
-            log.error("Tasks must be a subclass of either SourceTask or SinkTask", task);
+            log.error("Tasks must be a subclass of either SourceTask or SinkTask and current is {}", task);
             throw new ConnectException("Tasks must be a subclass of either SourceTask or SinkTask");
         }
     }
 
-    static Map producerConfigs(WorkerConfig config) {
+    static Map producerConfigs(ConnectorTaskId id,
+                                               String defaultClientId,
+                                               WorkerConfig config,
+                                               ConnectorConfig connConfig,
+                                               Class  connectorClass,
+                                               ConnectorClientConfigOverridePolicy connectorClientConfigOverridePolicy) {
         Map producerProps = new HashMap<>();
         producerProps.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, Utils.join(config.getList(WorkerConfig.BOOTSTRAP_SERVERS_CONFIG), ","));
         producerProps.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.ByteArraySerializer");
         producerProps.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.ByteArraySerializer");
-        // These settings are designed to ensure there is no data loss. They *may* be overridden via configs passed to the
-        // worker, but this may compromise the delivery guarantees of Kafka Connect.
+        // These settings will execute infinite retries on retriable exceptions. They *may* be overridden via configs passed to the worker,
+        // but this may compromise the delivery guarantees of Kafka Connect.
         producerProps.put(ProducerConfig.REQUEST_TIMEOUT_MS_CONFIG, Integer.toString(Integer.MAX_VALUE));
         producerProps.put(ProducerConfig.MAX_BLOCK_MS_CONFIG, Long.toString(Long.MAX_VALUE));
         producerProps.put(ProducerConfig.ACKS_CONFIG, "all");
         producerProps.put(ProducerConfig.MAX_IN_FLIGHT_REQUESTS_PER_CONNECTION, "1");
         producerProps.put(ProducerConfig.DELIVERY_TIMEOUT_MS_CONFIG, Integer.toString(Integer.MAX_VALUE));
+        producerProps.put(ProducerConfig.CLIENT_ID_CONFIG, defaultClientId);
         // User-specified overrides
         producerProps.putAll(config.originalsWithPrefix("producer."));
+
+        // Connector-specified overrides
+        Map producerOverrides =
+            connectorClientConfigOverrides(id, connConfig, connectorClass, ConnectorConfig.CONNECTOR_CLIENT_PRODUCER_OVERRIDES_PREFIX,
+                                           ConnectorType.SOURCE, ConnectorClientConfigRequest.ClientType.PRODUCER,
+                                           connectorClientConfigOverridePolicy);
+        producerProps.putAll(producerOverrides);
+
         return producerProps;
     }
 
-
-    static Map consumerConfigs(ConnectorTaskId id, WorkerConfig config) {
+    static Map consumerConfigs(ConnectorTaskId id,
+                                               WorkerConfig config,
+                                               ConnectorConfig connConfig,
+                                               Class connectorClass,
+                                               ConnectorClientConfigOverridePolicy connectorClientConfigOverridePolicy) {
         // Include any unknown worker configs so consumer configs can be set globally on the worker
         // and through to the task
         Map consumerProps = new HashMap<>();
 
         consumerProps.put(ConsumerConfig.GROUP_ID_CONFIG, SinkUtils.consumerGroupId(id.connector()));
+        consumerProps.put(ConsumerConfig.CLIENT_ID_CONFIG, "connector-consumer-" + id);
         consumerProps.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG,
                   Utils.join(config.getList(WorkerConfig.BOOTSTRAP_SERVERS_CONFIG), ","));
         consumerProps.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "false");
@@ -545,15 +594,81 @@ static Map consumerConfigs(ConnectorTaskId id, WorkerConfig conf
         consumerProps.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.ByteArrayDeserializer");
 
         consumerProps.putAll(config.originalsWithPrefix("consumer."));
+        // Connector-specified overrides
+        Map consumerOverrides =
+            connectorClientConfigOverrides(id, connConfig, connectorClass, ConnectorConfig.CONNECTOR_CLIENT_CONSUMER_OVERRIDES_PREFIX,
+                                           ConnectorType.SINK, ConnectorClientConfigRequest.ClientType.CONSUMER,
+                                           connectorClientConfigOverridePolicy);
+        consumerProps.putAll(consumerOverrides);
+
         return consumerProps;
     }
 
+    static Map adminConfigs(ConnectorTaskId id,
+                                            WorkerConfig config,
+                                            ConnectorConfig connConfig,
+                                            Class connectorClass,
+                                            ConnectorClientConfigOverridePolicy connectorClientConfigOverridePolicy) {
+        Map adminProps = new HashMap<>();
+        // Use the top-level worker configs to retain backwards compatibility with older releases which
+        // did not require a prefix for connector admin client configs in the worker configuration file
+        // Ignore configs that begin with "admin." since those will be added next (with the prefix stripped)
+        // and those that begin with "producer." and "consumer.", since we know they aren't intended for
+        // the admin client
+        Map nonPrefixedWorkerConfigs = config.originals().entrySet().stream()
+            .filter(e -> !e.getKey().startsWith("admin.")
+                && !e.getKey().startsWith("producer.")
+                && !e.getKey().startsWith("consumer."))
+            .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
+        adminProps.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG,
+            Utils.join(config.getList(WorkerConfig.BOOTSTRAP_SERVERS_CONFIG), ","));
+        adminProps.putAll(nonPrefixedWorkerConfigs);
+
+        // Admin client-specific overrides in the worker config
+        adminProps.putAll(config.originalsWithPrefix("admin."));
+
+        // Connector-specified overrides
+        Map adminOverrides =
+            connectorClientConfigOverrides(id, connConfig, connectorClass, ConnectorConfig.CONNECTOR_CLIENT_ADMIN_OVERRIDES_PREFIX,
+                                           ConnectorType.SINK, ConnectorClientConfigRequest.ClientType.ADMIN,
+                                           connectorClientConfigOverridePolicy);
+        adminProps.putAll(adminOverrides);
+
+        return adminProps;
+    }
+
+    private static Map connectorClientConfigOverrides(ConnectorTaskId id,
+                                                                      ConnectorConfig connConfig,
+                                                                      Class connectorClass,
+                                                                      String clientConfigPrefix,
+                                                                      ConnectorType connectorType,
+                                                                      ConnectorClientConfigRequest.ClientType clientType,
+                                                                      ConnectorClientConfigOverridePolicy connectorClientConfigOverridePolicy) {
+        Map clientOverrides = connConfig.originalsWithPrefix(clientConfigPrefix);
+        ConnectorClientConfigRequest connectorClientConfigRequest = new ConnectorClientConfigRequest(
+            id.connector(),
+            connectorType,
+            connectorClass,
+            clientOverrides,
+            clientType
+        );
+        List configValues = connectorClientConfigOverridePolicy.validate(connectorClientConfigRequest);
+        List errorConfigs = configValues.stream().
+            filter(configValue -> configValue.errorMessages().size() > 0).collect(Collectors.toList());
+        // These should be caught when the herder validates the connector configuration, but just in case
+        if (errorConfigs.size() > 0) {
+            throw new ConnectException("Client Config Overrides not allowed " + errorConfigs);
+        }
+        return clientOverrides;
+    }
+
     ErrorHandlingMetrics errorHandlingMetrics(ConnectorTaskId id) {
         return new ErrorHandlingMetrics(id, metrics);
     }
 
     private List sinkTaskReporters(ConnectorTaskId id, SinkConnectorConfig connConfig,
-                                                  ErrorHandlingMetrics errorHandlingMetrics) {
+                                                  ErrorHandlingMetrics errorHandlingMetrics,
+                                                  Class connectorClass) {
         ArrayList reporters = new ArrayList<>();
         LogReporter logReporter = new LogReporter(id, connConfig, errorHandlingMetrics);
         reporters.add(logReporter);
@@ -561,8 +676,10 @@ private List sinkTaskReporters(ConnectorTaskId id, SinkConnectorC
         // check if topic for dead letter queue exists
         String topic = connConfig.dlqTopicName();
         if (topic != null && !topic.isEmpty()) {
-            Map producerProps = producerConfigs(config);
-            DeadLetterQueueReporter reporter = DeadLetterQueueReporter.createAndSetup(config, id, connConfig, producerProps, errorHandlingMetrics);
+            Map producerProps = producerConfigs(id, "connector-dlq-producer-" + id, config, connConfig, connectorClass,
+                                                                connectorClientConfigOverridePolicy);
+            Map adminProps = adminConfigs(id, config, connConfig, connectorClass, connectorClientConfigOverridePolicy);
+            DeadLetterQueueReporter reporter = DeadLetterQueueReporter.createAndSetup(adminProps, id, connConfig, producerProps, errorHandlingMetrics);
             reporters.add(reporter);
         }
 
@@ -579,22 +696,24 @@ private List sourceTaskReporters(ConnectorTaskId id, ConnectorCon
     }
 
     private void stopTask(ConnectorTaskId taskId) {
-        WorkerTask task = tasks.get(taskId);
-        if (task == null) {
-            log.warn("Ignoring stop request for unowned task {}", taskId);
-            return;
-        }
+        try (LoggingContext loggingContext = LoggingContext.forTask(taskId)) {
+            WorkerTask task = tasks.get(taskId);
+            if (task == null) {
+                log.warn("Ignoring stop request for unowned task {}", taskId);
+                return;
+            }
 
-        log.info("Stopping task {}", task.id());
-        if (task instanceof WorkerSourceTask)
-            sourceTaskOffsetCommitter.remove(task.id());
+            log.info("Stopping task {}", task.id());
+            if (task instanceof WorkerSourceTask)
+                sourceTaskOffsetCommitter.remove(task.id());
 
-        ClassLoader savedLoader = plugins.currentThreadLoader();
-        try {
-            savedLoader = Plugins.compareAndSwapLoaders(task.loader());
-            task.stop();
-        } finally {
-            Plugins.compareAndSwapLoaders(savedLoader);
+            ClassLoader savedLoader = plugins.currentThreadLoader();
+            try {
+                savedLoader = Plugins.compareAndSwapLoaders(task.loader());
+                task.stop();
+            } finally {
+                Plugins.compareAndSwapLoaders(savedLoader);
+            }
         }
     }
 
@@ -607,15 +726,20 @@ private void stopTasks(Collection ids) {
     }
 
     private void awaitStopTask(ConnectorTaskId taskId, long timeout) {
-        WorkerTask task = tasks.remove(taskId);
-        if (task == null) {
-            log.warn("Ignoring await stop request for non-present task {}", taskId);
-            return;
-        }
+        try (LoggingContext loggingContext = LoggingContext.forTask(taskId)) {
+            WorkerTask task = tasks.remove(taskId);
+            if (task == null) {
+                log.warn("Ignoring await stop request for non-present task {}", taskId);
+                return;
+            }
 
-        if (!task.awaitStop(timeout)) {
-            log.error("Graceful stop of task {} failed.", task.id());
-            task.cancel();
+            connectorStatusMetricsGroup.recordTaskRemoved(taskId);
+            if (!task.awaitStop(timeout)) {
+                log.error("Graceful stop of task {} failed.", task.id());
+                task.cancel();
+            } else {
+                log.debug("Graceful stop of task {} succeeded.", task.id());
+            }
         }
     }
 
@@ -723,10 +847,84 @@ private void transitionTo(Object connectorOrTask, TargetState state, ClassLoader
         }
     }
 
+    ConnectorStatusMetricsGroup connectorStatusMetricsGroup() {
+        return connectorStatusMetricsGroup;
+    }
+
     WorkerMetricsGroup workerMetricsGroup() {
         return workerMetricsGroup;
     }
 
+    static class ConnectorStatusMetricsGroup {
+        private final ConnectMetrics connectMetrics;
+        private final ConnectMetricsRegistry registry;
+        private final ConcurrentMap connectorStatusMetrics = new ConcurrentHashMap<>();
+        private final Herder herder;
+        private final ConcurrentMap tasks;
+
+
+        protected ConnectorStatusMetricsGroup(
+            ConnectMetrics connectMetrics, ConcurrentMap tasks, Herder herder) {
+            this.connectMetrics = connectMetrics;
+            this.registry = connectMetrics.registry();
+            this.tasks = tasks;
+            this.herder = herder;
+        }
+
+        protected ConnectMetrics.LiteralSupplier taskCounter(String connName) {
+            return now -> tasks.keySet()
+                .stream()
+                .filter(taskId -> taskId.connector().equals(connName))
+                .count();
+        }
+
+        protected ConnectMetrics.LiteralSupplier taskStatusCounter(String connName, TaskStatus.State state) {
+            return now -> tasks.values()
+                .stream()
+                .filter(task ->
+                    task.id().connector().equals(connName) &&
+                    herder.taskStatus(task.id()).state().equalsIgnoreCase(state.toString()))
+                .count();
+        }
+
+        protected synchronized void recordTaskAdded(ConnectorTaskId connectorTaskId) {
+            if (connectorStatusMetrics.containsKey(connectorTaskId.connector())) {
+                return;
+            }
+
+            String connName = connectorTaskId.connector();
+
+            MetricGroup metricGroup = connectMetrics.group(registry.workerGroupName(),
+                registry.connectorTagName(), connName);
+
+            metricGroup.addValueMetric(registry.connectorTotalTaskCount, taskCounter(connName));
+            for (Map.Entry statusMetric : registry.connectorStatusMetrics
+                .entrySet()) {
+                metricGroup.addValueMetric(statusMetric.getKey(), taskStatusCounter(connName,
+                    statusMetric.getValue()));
+            }
+            connectorStatusMetrics.put(connectorTaskId.connector(), metricGroup);
+        }
+
+        protected synchronized void recordTaskRemoved(ConnectorTaskId connectorTaskId) {
+            // Unregister connector task count metric if we remove the last task of the connector
+            if (tasks.keySet().stream().noneMatch(id -> id.connector().equals(connectorTaskId.connector()))) {
+                connectorStatusMetrics.get(connectorTaskId.connector()).close();
+                connectorStatusMetrics.remove(connectorTaskId.connector());
+            }
+        }
+
+        protected synchronized void close() {
+            for (MetricGroup metricGroup: connectorStatusMetrics.values()) {
+                metricGroup.close();
+            }
+        }
+
+        protected MetricGroup metricGroup(String connectorId) {
+            return connectorStatusMetrics.get(connectorId);
+        }
+    }
+
     class WorkerMetricsGroup {
         private final MetricGroup metricGroup;
         private final Sensor connectorStartupAttempts;
@@ -742,18 +940,8 @@ public WorkerMetricsGroup(ConnectMetrics connectMetrics) {
             ConnectMetricsRegistry registry = connectMetrics.registry();
             metricGroup = connectMetrics.group(registry.workerGroupName());
 
-            metricGroup.addValueMetric(registry.connectorCount, new LiteralSupplier() {
-                @Override
-                public Double metricValue(long now) {
-                    return (double) connectors.size();
-                }
-            });
-            metricGroup.addValueMetric(registry.taskCount, new LiteralSupplier() {
-                @Override
-                public Double metricValue(long now) {
-                    return (double) tasks.size();
-                }
-            });
+            metricGroup.addValueMetric(registry.connectorCount, now -> (double) connectors.size());
+            metricGroup.addValueMetric(registry.taskCount, now -> (double) tasks.size());
 
             MetricName connectorFailurePct = metricGroup.metricName(registry.connectorStartupFailurePercentage);
             MetricName connectorSuccessPct = metricGroup.metricName(registry.connectorStartupSuccessPercentage);
@@ -762,13 +950,13 @@ public Double metricValue(long now) {
             connectorStartupResults.add(connectorStartupResultFrequencies);
 
             connectorStartupAttempts = metricGroup.sensor("connector-startup-attempts");
-            connectorStartupAttempts.add(metricGroup.metricName(registry.connectorStartupAttemptsTotal), new Total());
+            connectorStartupAttempts.add(metricGroup.metricName(registry.connectorStartupAttemptsTotal), new CumulativeSum());
 
             connectorStartupSuccesses = metricGroup.sensor("connector-startup-successes");
-            connectorStartupSuccesses.add(metricGroup.metricName(registry.connectorStartupSuccessTotal), new Total());
+            connectorStartupSuccesses.add(metricGroup.metricName(registry.connectorStartupSuccessTotal), new CumulativeSum());
 
             connectorStartupFailures = metricGroup.sensor("connector-startup-failures");
-            connectorStartupFailures.add(metricGroup.metricName(registry.connectorStartupFailureTotal), new Total());
+            connectorStartupFailures.add(metricGroup.metricName(registry.connectorStartupFailureTotal), new CumulativeSum());
 
             MetricName taskFailurePct = metricGroup.metricName(registry.taskStartupFailurePercentage);
             MetricName taskSuccessPct = metricGroup.metricName(registry.taskStartupSuccessPercentage);
@@ -777,13 +965,13 @@ public Double metricValue(long now) {
             taskStartupResults.add(taskStartupResultFrequencies);
 
             taskStartupAttempts = metricGroup.sensor("task-startup-attempts");
-            taskStartupAttempts.add(metricGroup.metricName(registry.taskStartupAttemptsTotal), new Total());
+            taskStartupAttempts.add(metricGroup.metricName(registry.taskStartupAttemptsTotal), new CumulativeSum());
 
             taskStartupSuccesses = metricGroup.sensor("task-startup-successes");
-            taskStartupSuccesses.add(metricGroup.metricName(registry.taskStartupSuccessTotal), new Total());
+            taskStartupSuccesses.add(metricGroup.metricName(registry.taskStartupSuccessTotal), new CumulativeSum());
 
             taskStartupFailures = metricGroup.sensor("task-startup-failures");
-            taskStartupFailures.add(metricGroup.metricName(registry.taskStartupFailureTotal), new Total());
+            taskStartupFailures.add(metricGroup.metricName(registry.taskStartupFailureTotal), new CumulativeSum());
         }
 
         void close() {
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 ff4f43f558fb7..837cfe5c19c1f 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
@@ -22,6 +22,7 @@
 import org.apache.kafka.common.config.ConfigDef;
 import org.apache.kafka.common.config.ConfigDef.Importance;
 import org.apache.kafka.common.config.ConfigDef.Type;
+import org.apache.kafka.common.config.ConfigException;
 import org.apache.kafka.common.config.internals.BrokerSecurityConfigs;
 import org.apache.kafka.common.metrics.Sensor;
 import org.apache.kafka.connect.json.JsonConverter;
@@ -36,6 +37,7 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.regex.Pattern;
 
 import static org.apache.kafka.common.config.ConfigDef.Range.atLeast;
@@ -186,6 +188,14 @@ public class WorkerConfig extends AbstractConfig {
         + "The default value of the Access-Control-Allow-Methods header allows cross origin requests for GET, POST and HEAD.";
     protected static final String ACCESS_CONTROL_ALLOW_METHODS_DEFAULT = "";
 
+    public static final String ADMIN_LISTENERS_CONFIG = "admin.listeners";
+    protected static final String ADMIN_LISTENERS_DOC = "List of comma-separated URIs the Admin REST API will listen on." +
+            " The supported protocols are HTTP and HTTPS." +
+            " An empty or blank string will disable this feature." +
+            " The default behavior is to use the regular listener (specified by the 'listeners' property).";
+    protected static final List ADMIN_LISTENERS_DEFAULT = null;
+    public static final String ADMIN_LISTENERS_HTTPS_CONFIGS_PREFIX = "admin.listeners.https.";
+
     public static final String PLUGIN_PATH_CONFIG = "plugin.path";
     protected static final String PLUGIN_PATH_DOC = "List of paths separated by commas (,) that "
             + "contain plugins (connectors, converters, transformations). The list should consist"
@@ -196,7 +206,10 @@ public class WorkerConfig extends AbstractConfig {
             + "plugins and their dependencies\n"
             + "Note: symlinks will be followed to discover dependencies or plugins.\n"
             + "Examples: plugin.path=/usr/local/share/java,/usr/local/share/kafka/plugins,"
-            + "/opt/connectors";
+            + "/opt/connectors\n" 
+            + "Do not use config provider variables in this property, since the raw path is used "
+            + "by the worker's scanner before config providers are initialized and used to "
+            + "replace variables.";
 
     public static final String CONFIG_PROVIDERS_CONFIG = "config.providers";
     protected static final String CONFIG_PROVIDERS_DOC =
@@ -212,6 +225,14 @@ public class WorkerConfig extends AbstractConfig {
             + "ConnectRestExtension allows you to inject into Connect's REST API user defined resources like filters. "
             + "Typically used to add custom capability like logging, security, etc. ";
 
+    public static final String CONNECTOR_CLIENT_POLICY_CLASS_CONFIG = "connector.client.config.override.policy";
+    public static final String CONNECTOR_CLIENT_POLICY_CLASS_DOC =
+        "Class name or alias of implementation of ConnectorClientConfigOverridePolicy. Defines what client configurations can be "
+        + "overriden by the connector. The default implementation is `None`. The other possible policies in the framework include `All` "
+        + "and `Principal`. ";
+    public static final String CONNECTOR_CLIENT_POLICY_CLASS_DEFAULT = "None";
+
+
     public static final String METRICS_SAMPLE_WINDOW_MS_CONFIG = CommonClientConfigs.METRICS_SAMPLE_WINDOW_MS_CONFIG;
     public static final String METRICS_NUM_SAMPLES_CONFIG = CommonClientConfigs.METRICS_NUM_SAMPLES_CONFIG;
     public static final String METRICS_RECORDING_LEVEL_CONFIG = CommonClientConfigs.METRICS_RECORDING_LEVEL_CONFIG;
@@ -289,7 +310,13 @@ protected static ConfigDef baseConfigDef() {
                         Collections.emptyList(),
                         Importance.LOW, CONFIG_PROVIDERS_DOC)
                 .define(REST_EXTENSION_CLASSES_CONFIG, Type.LIST, "",
-                        Importance.LOW, REST_EXTENSION_CLASSES_DOC);
+                        Importance.LOW, REST_EXTENSION_CLASSES_DOC)
+                .define(ADMIN_LISTENERS_CONFIG, Type.LIST, null,
+                        new AdminListenersValidator(), Importance.LOW, ADMIN_LISTENERS_DOC)
+                .define(CONNECTOR_CLIENT_POLICY_CLASS_CONFIG, Type.STRING, CONNECTOR_CLIENT_POLICY_CLASS_DEFAULT,
+                        Importance.MEDIUM, CONNECTOR_CLIENT_POLICY_CLASS_DOC)
+                // security support
+                .withClientSslSupport();
     }
 
     private void logInternalConverterDeprecationWarnings(Map props) {
@@ -344,6 +371,27 @@ private void logDeprecatedProperty(String propName, String propValue, String def
         }
     }
 
+    private void logPluginPathConfigProviderWarning(Map rawOriginals) {
+        String rawPluginPath = rawOriginals.get(PLUGIN_PATH_CONFIG);
+        // Can't use AbstractConfig::originalsStrings here since some values may be null, which
+        // causes that method to fail
+        String transformedPluginPath = Objects.toString(originals().get(PLUGIN_PATH_CONFIG));
+        if (!Objects.equals(rawPluginPath, transformedPluginPath)) {
+            log.warn(
+                "Variables cannot be used in the 'plugin.path' property, since the property is "
+                + "used by plugin scanning before the config providers that replace the " 
+                + "variables are initialized. The raw value '{}' was used for plugin scanning, as " 
+                + "opposed to the transformed value '{}', and this may cause unexpected results.",
+                rawPluginPath,
+                transformedPluginPath
+            );
+        }
+    }
+
+    public Integer getRebalanceTimeout() {
+        return null;
+    }
+
     @Override
     protected Map postProcessParsedConfig(final Map parsedValues) {
         return CommonClientConfigs.postProcessReconnectBackoffConfigs(this, parsedValues);
@@ -359,6 +407,34 @@ public static List pluginLocations(Map props) {
     public WorkerConfig(ConfigDef definition, Map props) {
         super(definition, props);
         logInternalConverterDeprecationWarnings(props);
+        logPluginPathConfigProviderWarning(props);
+    }
+
+    private static class AdminListenersValidator implements ConfigDef.Validator {
+        @Override
+        public void ensureValid(String name, Object value) {
+            if (value == null) {
+                return;
+            }
+
+            if (!(value instanceof List)) {
+                throw new ConfigException("Invalid value type (list expected).");
+            }
+
+            List items = (List) value;
+            if (items.isEmpty()) {
+                return;
+            }
+
+            for (Object item: items) {
+                if (!(item instanceof String)) {
+                    throw new ConfigException("Invalid type for admin listener (expected String).");
+                }
+                if (((String) item).trim().isEmpty()) {
+                    throw new ConfigException("Empty listener found when parsing list.");
+                }
+            }
+        }
     }
 
 }
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerConfigTransformer.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerConfigTransformer.java
index 3373d5ce32886..318626bd5cade 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerConfigTransformer.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerConfigTransformer.java
@@ -20,7 +20,9 @@
 import org.apache.kafka.common.config.provider.ConfigProvider;
 import org.apache.kafka.common.config.ConfigTransformer;
 import org.apache.kafka.common.config.ConfigTransformerResult;
+import org.apache.kafka.common.utils.Utils;
 import org.apache.kafka.connect.runtime.Herder.ConfigReloadAction;
+import org.apache.kafka.connect.util.Callback;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -33,15 +35,17 @@
  * A wrapper class to perform configuration transformations and schedule reloads for any
  * retrieved TTL values.
  */
-public class WorkerConfigTransformer {
+public class WorkerConfigTransformer implements AutoCloseable {
     private static final Logger log = LoggerFactory.getLogger(WorkerConfigTransformer.class);
 
     private final Worker worker;
     private final ConfigTransformer configTransformer;
     private final ConcurrentMap> requests = new ConcurrentHashMap<>();
+    private final Map configProviders;
 
     public WorkerConfigTransformer(Worker worker, Map configProviders) {
         this.worker = worker;
+        this.configProviders = configProviders;
         this.configTransformer = new ConfigTransformer(configProviders);
     }
 
@@ -86,7 +90,20 @@ private void scheduleReload(String connectorName, String path, long ttl) {
             }
         }
         log.info("Scheduling a restart of connector {} in {} ms", connectorName, ttl);
-        HerderRequest request = worker.herder().restartConnector(ttl, connectorName, null);
+        Callback cb = new Callback() {
+            @Override
+            public void onCompletion(Throwable error, Void result) {
+                if (error != null) {
+                    log.error("Unexpected error during connector restart: ", error);
+                }
+            }
+        };
+        HerderRequest request = worker.herder().restartConnector(ttl, connectorName, cb);
         connectorRequests.put(path, request);
     }
+
+    @Override
+    public void close() {
+        configProviders.values().forEach(x -> Utils.closeQuietly(x, "config provider"));
+    }
 }
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerConnector.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerConnector.java
index e5b990fc7abf7..7923943d121f7 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerConnector.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerConnector.java
@@ -18,7 +18,6 @@
 
 import org.apache.kafka.connect.connector.Connector;
 import org.apache.kafka.connect.connector.ConnectorContext;
-import org.apache.kafka.connect.runtime.ConnectMetrics.LiteralSupplier;
 import org.apache.kafka.connect.runtime.ConnectMetrics.MetricGroup;
 import org.apache.kafka.connect.sink.SinkConnector;
 import org.apache.kafka.connect.source.SourceConnector;
@@ -76,7 +75,7 @@ public WorkerConnector(String connName,
     public void initialize(ConnectorConfig connectorConfig) {
         try {
             this.config = connectorConfig.originalsStrings();
-            log.debug("{} Initializing connector {} with config {}", this, connName, config);
+            log.debug("{} Initializing connector {}", this, connName);
             if (isSinkConnector()) {
                 SinkConnectorConfig.validate(config);
             }
@@ -257,12 +256,7 @@ public ConnectorMetricsGroup(ConnectMetrics connectMetrics, AbstractStatus.State
             metricGroup.addImmutableValueMetric(registry.connectorType, connectorType());
             metricGroup.addImmutableValueMetric(registry.connectorClass, connector.getClass().getName());
             metricGroup.addImmutableValueMetric(registry.connectorVersion, connector.version());
-            metricGroup.addValueMetric(registry.connectorStatus, new LiteralSupplier() {
-                @Override
-                public String metricValue(long now) {
-                    return state.toString().toLowerCase(Locale.getDefault());
-                }
-            });
+            metricGroup.addValueMetric(registry.connectorStatus, now -> state.toString().toLowerCase(Locale.getDefault()));
         }
 
         public void close() {
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 a112bfa41aefa..0fa28e6383cff 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
@@ -27,11 +27,13 @@
 import org.apache.kafka.common.errors.WakeupException;
 import org.apache.kafka.common.metrics.Sensor;
 import org.apache.kafka.common.metrics.stats.Avg;
+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.Total;
 import org.apache.kafka.common.metrics.stats.Value;
 import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.common.utils.Utils;
+import org.apache.kafka.common.utils.Utils.UncheckedCloseable;
 import org.apache.kafka.connect.data.SchemaAndValue;
 import org.apache.kafka.connect.errors.ConnectException;
 import org.apache.kafka.connect.errors.RetriableException;
@@ -52,7 +54,6 @@
 
 import java.time.Duration;
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.List;
@@ -171,6 +172,7 @@ protected void close() {
         } catch (Throwable t) {
             log.warn("Could not close transformation chain", t);
         }
+        Utils.closeQuietly(retryWithToleranceOperator, "retry operator");
     }
 
     @Override
@@ -187,13 +189,11 @@ public void transitionTo(TargetState state) {
     @Override
     public void execute() {
         initializeAndStart();
-        try {
+        // Make sure any uncommitted data has been committed and the task has
+        // a chance to clean up its state
+        try (UncheckedCloseable suppressible = this::closePartitions) {
             while (!isStopping())
                 iteration();
-        } finally {
-            // Make sure any uncommitted data has been committed and the task has
-            // a chance to clean up its state
-            closePartitions();
         }
     }
 
@@ -206,7 +206,7 @@ protected void iteration() {
             // Maybe commit
             if (!committing && (context.isCommitRequested() || now >= nextCommit)) {
                 commitOffsets(now, false);
-                nextCommit += offsetCommitIntervalMs;
+                nextCommit = now + offsetCommitIntervalMs;
                 context.clearCommitRequest();
             }
 
@@ -286,9 +286,9 @@ protected void initializeAndStart() {
         SinkConnectorConfig.validate(taskConfig);
 
         if (SinkConnectorConfig.hasTopicsConfig(taskConfig)) {
-            String[] topics = taskConfig.get(SinkTask.TOPICS_CONFIG).split(",");
-            consumer.subscribe(Arrays.asList(topics), new HandleRebalance());
-            log.debug("{} Initializing and starting task for topics {}", this, topics);
+            List topics = SinkConnectorConfig.parseTopicsList(taskConfig);
+            consumer.subscribe(topics, new HandleRebalance());
+            log.debug("{} Initializing and starting task for topics {}", this, Utils.join(topics, ", "));
         } else {
             String topicsRegexStr = taskConfig.get(SinkTask.TOPICS_REGEX_CONFIG);
             Pattern pattern = Pattern.compile(topicsRegexStr);
@@ -481,10 +481,10 @@ private void convertMessages(ConsumerRecords msgs) {
     }
 
     private SinkRecord convertAndTransformRecord(final ConsumerRecord msg) {
-        SchemaAndValue keyAndSchema = retryWithToleranceOperator.execute(() -> keyConverter.toConnectData(msg.topic(), msg.key()),
+        SchemaAndValue keyAndSchema = retryWithToleranceOperator.execute(() -> keyConverter.toConnectData(msg.topic(), msg.headers(), msg.key()),
                 Stage.KEY_CONVERTER, keyConverter.getClass());
 
-        SchemaAndValue valueAndSchema = retryWithToleranceOperator.execute(() -> valueConverter.toConnectData(msg.topic(), msg.value()),
+        SchemaAndValue valueAndSchema = retryWithToleranceOperator.execute(() -> valueConverter.toConnectData(msg.topic(), msg.headers(), msg.value()),
                 Stage.VALUE_CONVERTER, valueConverter.getClass());
 
         Headers headers = retryWithToleranceOperator.execute(() -> convertHeadersFor(msg), Stage.HEADER_CONVERTER, headerConverter.getClass());
@@ -612,6 +612,11 @@ SinkTaskMetricsGroup sinkTaskMetricsGroup() {
         return sinkTaskMetricsGroup;
     }
 
+    // Visible for testing
+    long getNextCommit() {
+        return nextCommit;
+    }
+
     private class HandleRebalance implements ConsumerRebalanceListener {
         @Override
         public void onPartitionsAssigned(Collection partitions) {
@@ -700,11 +705,11 @@ public SinkTaskMetricsGroup(ConnectorTaskId id, ConnectMetrics connectMetrics) {
 
             sinkRecordRead = metricGroup.sensor("sink-record-read");
             sinkRecordRead.add(metricGroup.metricName(registry.sinkRecordReadRate), new Rate());
-            sinkRecordRead.add(metricGroup.metricName(registry.sinkRecordReadTotal), new Total());
+            sinkRecordRead.add(metricGroup.metricName(registry.sinkRecordReadTotal), new CumulativeSum());
 
             sinkRecordSend = metricGroup.sensor("sink-record-send");
             sinkRecordSend.add(metricGroup.metricName(registry.sinkRecordSendRate), new Rate());
-            sinkRecordSend.add(metricGroup.metricName(registry.sinkRecordSendTotal), new Total());
+            sinkRecordSend.add(metricGroup.metricName(registry.sinkRecordSendTotal), new CumulativeSum());
 
             sinkRecordActiveCount = metricGroup.sensor("sink-record-active-count");
             sinkRecordActiveCount.add(metricGroup.metricName(registry.sinkRecordActiveCount), new Value());
@@ -719,11 +724,11 @@ public SinkTaskMetricsGroup(ConnectorTaskId id, ConnectMetrics connectMetrics) {
 
             offsetCompletion = metricGroup.sensor("offset-commit-completion");
             offsetCompletion.add(metricGroup.metricName(registry.sinkRecordOffsetCommitCompletionRate), new Rate());
-            offsetCompletion.add(metricGroup.metricName(registry.sinkRecordOffsetCommitCompletionTotal), new Total());
+            offsetCompletion.add(metricGroup.metricName(registry.sinkRecordOffsetCommitCompletionTotal), new CumulativeSum());
 
             offsetCompletionSkip = metricGroup.sensor("offset-commit-completion-skip");
             offsetCompletionSkip.add(metricGroup.metricName(registry.sinkRecordOffsetCommitSkipRate), new Rate());
-            offsetCompletionSkip.add(metricGroup.metricName(registry.sinkRecordOffsetCommitSkipTotal), new Total());
+            offsetCompletionSkip.add(metricGroup.metricName(registry.sinkRecordOffsetCommitSkipTotal), new CumulativeSum());
 
             putBatchTime = metricGroup.sensor("put-batch-time");
             putBatchTime.add(metricGroup.metricName(registry.sinkRecordPutBatchTimeMax), new Max());
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerSourceTask.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerSourceTask.java
index 71e026c717686..c3739b5efd2e6 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerSourceTask.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerSourceTask.java
@@ -21,16 +21,17 @@
 import org.apache.kafka.clients.producer.ProducerRecord;
 import org.apache.kafka.clients.producer.RecordMetadata;
 import org.apache.kafka.common.KafkaException;
-import org.apache.kafka.common.errors.RetriableException;
 import org.apache.kafka.common.header.internals.RecordHeaders;
 import org.apache.kafka.common.metrics.Sensor;
 import org.apache.kafka.common.metrics.stats.Avg;
+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.Total;
 import org.apache.kafka.common.metrics.stats.Value;
 import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.common.utils.Utils;
 import org.apache.kafka.connect.errors.ConnectException;
+import org.apache.kafka.connect.errors.RetriableException;
 import org.apache.kafka.connect.header.Header;
 import org.apache.kafka.connect.header.Headers;
 import org.apache.kafka.connect.runtime.ConnectMetrics.MetricGroup;
@@ -39,9 +40,9 @@
 import org.apache.kafka.connect.runtime.errors.Stage;
 import org.apache.kafka.connect.source.SourceRecord;
 import org.apache.kafka.connect.source.SourceTask;
+import org.apache.kafka.connect.storage.CloseableOffsetStorageReader;
 import org.apache.kafka.connect.storage.Converter;
 import org.apache.kafka.connect.storage.HeaderConverter;
-import org.apache.kafka.connect.storage.OffsetStorageReader;
 import org.apache.kafka.connect.storage.OffsetStorageWriter;
 import org.apache.kafka.connect.util.ConnectUtils;
 import org.apache.kafka.connect.util.ConnectorTaskId;
@@ -57,6 +58,7 @@
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicReference;
 
 /**
  * WorkerTask that uses a SourceTask to ingest data into Kafka.
@@ -74,10 +76,11 @@ class WorkerSourceTask extends WorkerTask {
     private final HeaderConverter headerConverter;
     private final TransformationChain transformationChain;
     private KafkaProducer producer;
-    private final OffsetStorageReader offsetReader;
+    private final CloseableOffsetStorageReader offsetReader;
     private final OffsetStorageWriter offsetWriter;
     private final Time time;
     private final SourceTaskMetricsGroup sourceTaskMetricsGroup;
+    private final AtomicReference producerSendException;
 
     private List toSend;
     private boolean lastSendFailed; // Whether the last send failed *synchronously*, i.e. never made it into the producer's RecordAccumulator
@@ -103,7 +106,7 @@ public WorkerSourceTask(ConnectorTaskId id,
                             HeaderConverter headerConverter,
                             TransformationChain transformationChain,
                             KafkaProducer producer,
-                            OffsetStorageReader offsetReader,
+                            CloseableOffsetStorageReader offsetReader,
                             OffsetStorageWriter offsetWriter,
                             WorkerConfig workerConfig,
                             ClusterConfigState configState,
@@ -133,6 +136,7 @@ public WorkerSourceTask(ConnectorTaskId id,
         this.flushing = false;
         this.stopRequestedLatch = new CountDownLatch(1);
         this.sourceTaskMetricsGroup = new SourceTaskMetricsGroup(id, connectMetrics);
+        this.producerSendException = new AtomicReference<>();
     }
 
     @Override
@@ -162,6 +166,7 @@ protected void close() {
         } catch (Throwable t) {
             log.warn("Could not close transformation chain", t);
         }
+        Utils.closeQuietly(retryWithToleranceOperator, "retry operator");
     }
 
     @Override
@@ -169,6 +174,12 @@ protected void releaseResources() {
         sourceTaskMetricsGroup.close();
     }
 
+    @Override
+    public void cancel() {
+        super.cancel();
+        offsetReader.close();
+    }
+
     @Override
     public void stop() {
         super.stop();
@@ -215,6 +226,8 @@ public void execute() {
                     continue;
                 }
 
+                maybeThrowProducerSendException();
+
                 if (toSend == null) {
                     log.trace("{} Nothing to send to Kafka. Polling source for additional records", this);
                     long start = time.milliseconds();
@@ -225,7 +238,7 @@ public void execute() {
                 }
                 if (toSend == null)
                     continue;
-                log.debug("{} About to send " + toSend.size() + " records to Kafka", this);
+                log.trace("{} About to send {} records to Kafka", this, toSend.size());
                 if (!sendRecords())
                     stopRequestedLatch.await(SEND_FAILED_BACKOFF_MS, TimeUnit.MILLISECONDS);
             }
@@ -240,10 +253,19 @@ public void execute() {
         }
     }
 
+    private void maybeThrowProducerSendException() {
+        if (producerSendException.get() != null) {
+            throw new ConnectException(
+                "Unrecoverable exception from producer send callback",
+                producerSendException.get()
+            );
+        }
+    }
+
     protected List poll() throws InterruptedException {
         try {
             return task.poll();
-        } catch (RetriableException e) {
+        } catch (RetriableException | org.apache.kafka.common.errors.RetriableException e) {
             log.warn("{} failed to poll records from SourceTask. Will retry operation.", this, e);
             // Do nothing. Let the framework poll whenever it's ready.
             return null;
@@ -264,10 +286,10 @@ private ProducerRecord convertTransformedRecord(SourceRecord rec
 
         RecordHeaders headers = retryWithToleranceOperator.execute(() -> convertHeaderFor(record), Stage.HEADER_CONVERTER, headerConverter.getClass());
 
-        byte[] key = retryWithToleranceOperator.execute(() -> keyConverter.fromConnectData(record.topic(), record.keySchema(), record.key()),
+        byte[] key = retryWithToleranceOperator.execute(() -> keyConverter.fromConnectData(record.topic(), headers, record.keySchema(), record.key()),
                 Stage.KEY_CONVERTER, keyConverter.getClass());
 
-        byte[] value = retryWithToleranceOperator.execute(() -> valueConverter.fromConnectData(record.topic(), record.valueSchema(), record.value()),
+        byte[] value = retryWithToleranceOperator.execute(() -> valueConverter.fromConnectData(record.topic(), headers, record.valueSchema(), record.value()),
                 Stage.VALUE_CONVERTER, valueConverter.getClass());
 
         if (retryWithToleranceOperator.failed()) {
@@ -286,15 +308,17 @@ private ProducerRecord convertTransformedRecord(SourceRecord rec
     private boolean sendRecords() {
         int processed = 0;
         recordBatch(toSend.size());
-        final SourceRecordWriteCounter counter = new SourceRecordWriteCounter(toSend.size(), sourceTaskMetricsGroup);
+        final SourceRecordWriteCounter counter =
+                toSend.size() > 0 ? new SourceRecordWriteCounter(toSend.size(), sourceTaskMetricsGroup) : null;
         for (final SourceRecord preTransformRecord : toSend) {
+            maybeThrowProducerSendException();
 
             retryWithToleranceOperator.sourceRecord(preTransformRecord);
             final SourceRecord record = transformationChain.apply(preTransformRecord);
             final ProducerRecord producerRecord = convertTransformedRecord(record);
             if (producerRecord == null || retryWithToleranceOperator.failed()) {
                 counter.skipRecord();
-                commitTaskRecord(preTransformRecord);
+                commitTaskRecord(preTransformRecord, null);
                 continue;
             }
 
@@ -322,26 +346,22 @@ private boolean sendRecords() {
                             @Override
                             public void onCompletion(RecordMetadata recordMetadata, Exception e) {
                                 if (e != null) {
-                                    // Given the default settings for zero data loss, this should basically never happen --
-                                    // between "infinite" retries, indefinite blocking on full buffers, and "infinite" request
-                                    // timeouts, callbacks with exceptions should never be invoked in practice. If the
-                                    // user overrode these settings, the best we can do is notify them of the failure via
-                                    // logging.
-                                    log.error("{} failed to send record to {}: {}", WorkerSourceTask.this, topic, e);
+                                    log.error("{} failed to send record to {}:", WorkerSourceTask.this, topic, e);
                                     log.debug("{} Failed record: {}", WorkerSourceTask.this, preTransformRecord);
+                                    producerSendException.compareAndSet(null, e);
                                 } else {
+                                    recordSent(producerRecord);
+                                    counter.completeRecord();
                                     log.trace("{} Wrote record successfully: topic {} partition {} offset {}",
                                             WorkerSourceTask.this,
                                             recordMetadata.topic(), recordMetadata.partition(),
                                             recordMetadata.offset());
-                                    commitTaskRecord(preTransformRecord);
+                                    commitTaskRecord(preTransformRecord, recordMetadata);
                                 }
-                                recordSent(producerRecord);
-                                counter.completeRecord();
                             }
                         });
                 lastSendFailed = false;
-            } catch (RetriableException e) {
+            } catch (org.apache.kafka.common.errors.RetriableException e) {
                 log.warn("{} Failed to send {}, backing off before retrying:", this, producerRecord, e);
                 toSend = toSend.subList(processed, toSend.size());
                 lastSendFailed = true;
@@ -370,9 +390,9 @@ private RecordHeaders convertHeaderFor(SourceRecord record) {
         return result;
     }
 
-    private void commitTaskRecord(SourceRecord record) {
+    private void commitTaskRecord(SourceRecord record, RecordMetadata metadata) {
         try {
-            task.commitRecord(record);
+            task.commitRecord(record, metadata);
         } catch (Throwable t) {
             log.error("{} Exception thrown while calling task.commitRecord()", this, t);
         }
@@ -591,11 +611,11 @@ public SourceTaskMetricsGroup(ConnectorTaskId id, ConnectMetrics connectMetrics)
 
             sourceRecordPoll = metricGroup.sensor("source-record-poll");
             sourceRecordPoll.add(metricGroup.metricName(registry.sourceRecordPollRate), new Rate());
-            sourceRecordPoll.add(metricGroup.metricName(registry.sourceRecordPollTotal), new Total());
+            sourceRecordPoll.add(metricGroup.metricName(registry.sourceRecordPollTotal), new CumulativeSum());
 
             sourceRecordWrite = metricGroup.sensor("source-record-write");
             sourceRecordWrite.add(metricGroup.metricName(registry.sourceRecordWriteRate), new Rate());
-            sourceRecordWrite.add(metricGroup.metricName(registry.sourceRecordWriteTotal), new Total());
+            sourceRecordWrite.add(metricGroup.metricName(registry.sourceRecordWriteTotal), new CumulativeSum());
 
             pollTime = metricGroup.sensor("poll-batch-time");
             pollTime.add(metricGroup.metricName(registry.sourceRecordPollBatchTimeMax), new Max());
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerTask.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerTask.java
index 9cecb3d258231..05d18ff0cff36 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerTask.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerTask.java
@@ -31,6 +31,7 @@
 import org.apache.kafka.connect.runtime.errors.RetryWithToleranceOperator;
 import org.apache.kafka.connect.runtime.isolation.Plugins;
 import org.apache.kafka.connect.util.ConnectorTaskId;
+import org.apache.kafka.connect.util.LoggingContext;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -215,27 +216,32 @@ protected synchronized void onResume() {
 
     @Override
     public void run() {
-        ClassLoader savedLoader = Plugins.compareAndSwapLoaders(loader);
-        String savedName = Thread.currentThread().getName();
-        try {
-            Thread.currentThread().setName(THREAD_NAME_PREFIX + id);
-            doRun();
-            onShutdown();
-        } catch (Throwable t) {
-            onFailure(t);
+        // Clear all MDC parameters, in case this thread is being reused
+        LoggingContext.clear();
 
-            if (t instanceof Error)
-                throw (Error) t;
-        } finally {
+        try (LoggingContext loggingContext = LoggingContext.forTask(id())) {
+            ClassLoader savedLoader = Plugins.compareAndSwapLoaders(loader);
+            String savedName = Thread.currentThread().getName();
             try {
-                Thread.currentThread().setName(savedName);
-                Plugins.compareAndSwapLoaders(savedLoader);
-                shutdownLatch.countDown();
+                Thread.currentThread().setName(THREAD_NAME_PREFIX + id);
+                doRun();
+                onShutdown();
+            } catch (Throwable t) {
+                onFailure(t);
+
+                if (t instanceof Error)
+                    throw (Error) t;
             } finally {
                 try {
-                    releaseResources();
+                    Thread.currentThread().setName(savedName);
+                    Plugins.compareAndSwapLoaders(savedLoader);
+                    shutdownLatch.countDown();
                 } finally {
-                    taskMetricsGroup.close();
+                    try {
+                        releaseResources();
+                    } finally {
+                        taskMetricsGroup.close();
+                    }
                 }
             }
         }
@@ -409,6 +415,12 @@ public void onShutdown(ConnectorTaskId id) {
             delegateListener.onShutdown(id);
         }
 
+        @Override
+        public void onDeletion(ConnectorTaskId id) {
+            taskStateTimer.changeState(State.DESTROYED, time.milliseconds());
+            delegateListener.onDeletion(id);
+        }
+
         public void recordState(TargetState state) {
             switch (state) {
                 case STARTED:
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/ClusterConfigState.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/ClusterConfigState.java
index fc6a50d2fc078..a85a8e69e9f99 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/ClusterConfigState.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/ClusterConfigState.java
@@ -17,6 +17,7 @@
 package org.apache.kafka.connect.runtime.distributed;
 
 import org.apache.kafka.common.config.provider.ConfigProvider;
+import org.apache.kafka.connect.runtime.SessionKey;
 import org.apache.kafka.connect.runtime.WorkerConfigTransformer;
 import org.apache.kafka.connect.runtime.TargetState;
 import org.apache.kafka.connect.util.ConnectorTaskId;
@@ -37,6 +38,7 @@ public class ClusterConfigState {
     public static final long NO_OFFSET = -1;
     public static final ClusterConfigState EMPTY = new ClusterConfigState(
             NO_OFFSET,
+            null,
             Collections.emptyMap(),
             Collections.>emptyMap(),
             Collections.emptyMap(),
@@ -44,6 +46,7 @@ public class ClusterConfigState {
             Collections.emptySet());
 
     private final long offset;
+    private final SessionKey sessionKey;
     private final Map connectorTaskCounts;
     private final Map> connectorConfigs;
     private final Map connectorTargetStates;
@@ -52,12 +55,14 @@ public class ClusterConfigState {
     private final WorkerConfigTransformer configTransformer;
 
     public ClusterConfigState(long offset,
+                              SessionKey sessionKey,
                               Map connectorTaskCounts,
                               Map> connectorConfigs,
                               Map connectorTargetStates,
                               Map> taskConfigs,
                               Set inconsistentConnectors) {
         this(offset,
+                sessionKey,
                 connectorTaskCounts,
                 connectorConfigs,
                 connectorTargetStates,
@@ -67,6 +72,7 @@ public ClusterConfigState(long offset,
     }
 
     public ClusterConfigState(long offset,
+                              SessionKey sessionKey,
                               Map connectorTaskCounts,
                               Map> connectorConfigs,
                               Map connectorTargetStates,
@@ -74,6 +80,7 @@ public ClusterConfigState(long offset,
                               Set inconsistentConnectors,
                               WorkerConfigTransformer configTransformer) {
         this.offset = offset;
+        this.sessionKey = sessionKey;
         this.connectorTaskCounts = connectorTaskCounts;
         this.connectorConfigs = connectorConfigs;
         this.connectorTargetStates = connectorTargetStates;
@@ -91,6 +98,14 @@ public long offset() {
         return offset;
     }
 
+    /**
+     * Get the latest session key from the config state
+     * @return the {@link SessionKey session key}; may be null if no key has been read yet
+     */
+    public SessionKey sessionKey() {
+        return sessionKey;
+    }
+
     /**
      * Check whether this snapshot contains configuration for a connector.
      * @param connector name of the connector
@@ -229,6 +244,7 @@ public Set inconsistentConnectors() {
     public String toString() {
         return "ClusterConfigState{" +
                 "offset=" + offset +
+                ", sessionKey=" + (sessionKey != null ? "[hidden]" : "null") +
                 ", connectorTaskCounts=" + connectorTaskCounts +
                 ", connectorConfigs=" + connectorConfigs +
                 ", taskConfigs=" + taskConfigs +
@@ -242,6 +258,7 @@ public boolean equals(Object o) {
         if (o == null || getClass() != o.getClass()) return false;
         ClusterConfigState that = (ClusterConfigState) o;
         return offset == that.offset &&
+                Objects.equals(sessionKey, that.sessionKey) &&
                 Objects.equals(connectorTaskCounts, that.connectorTaskCounts) &&
                 Objects.equals(connectorConfigs, that.connectorConfigs) &&
                 Objects.equals(connectorTargetStates, that.connectorTargetStates) &&
@@ -254,6 +271,7 @@ public boolean equals(Object o) {
     public int hashCode() {
         return Objects.hash(
                 offset,
+                sessionKey,
                 connectorTaskCounts,
                 connectorConfigs,
                 connectorTargetStates,
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/ConnectAssignor.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/ConnectAssignor.java
new file mode 100644
index 0000000000000..752e62e680a5e
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/ConnectAssignor.java
@@ -0,0 +1,43 @@
+/*
+ * 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.distributed;
+
+import org.apache.kafka.common.message.JoinGroupResponseData;
+
+import java.nio.ByteBuffer;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * An assignor that computes a distribution of connectors and tasks among the workers of the group
+ * that performs rebalancing.
+ */
+public interface ConnectAssignor {
+    /**
+     * Based on the member metadata and the information stored in the worker coordinator this
+     * method computes an assignment of connectors and tasks among the members of the worker group.
+     *
+     * @param leaderId the leader of the group
+     * @param protocol the protocol type; for Connect assignors this is normally "connect"
+     * @param allMemberMetadata the metadata of all the active workers of the group
+     * @param coordinator the worker coordinator that runs this assignor
+     * @return the assignment of connectors and tasks to workers
+     */
+    Map performAssignment(String leaderId, String protocol,
+                                              List allMemberMetadata,
+                                              WorkerCoordinator coordinator);
+}
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/ConnectProtocol.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/ConnectProtocol.java
index dbb33bcbeb26c..15fc6059a627b 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/ConnectProtocol.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/ConnectProtocol.java
@@ -26,10 +26,17 @@
 
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
 import java.util.HashSet;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
+
+import static org.apache.kafka.common.message.JoinGroupRequestData.JoinGroupRequestProtocol;
+import static org.apache.kafka.common.message.JoinGroupRequestData.JoinGroupRequestProtocolCollection;
+import static org.apache.kafka.connect.runtime.distributed.ConnectProtocolCompatibility.EAGER;
 
 /**
  * This class implements the protocol for Kafka Connect workers in a group. It includes the format of worker state used when
@@ -50,19 +57,53 @@ public class ConnectProtocol {
     public static final short CONNECT_PROTOCOL_V0 = 0;
     public static final Schema CONNECT_PROTOCOL_HEADER_SCHEMA = new Schema(
             new Field(VERSION_KEY_NAME, Type.INT16));
+
+    /**
+     * Connect Protocol Header V0:
+     * 
      +     *   Version            => Int16
+     * 
      
+     */
     private static final Struct CONNECT_PROTOCOL_HEADER_V0 = new Struct(CONNECT_PROTOCOL_HEADER_SCHEMA)
             .set(VERSION_KEY_NAME, CONNECT_PROTOCOL_V0);
 
+    /**
+     * Config State V0:
+     * 
      +     *   Url                => [String]
+     *   ConfigOffset       => Int64
+     * 
      
+     */
     public static final Schema CONFIG_STATE_V0 = new Schema(
             new Field(URL_KEY_NAME, Type.STRING),
             new Field(CONFIG_OFFSET_KEY_NAME, Type.INT64));
 
-    // Assignments for each worker are a set of connectors and tasks. These are categorized by connector ID. A sentinel
-    // task ID (CONNECTOR_TASK) is used to indicate the connector itself (i.e. that the assignment includes
-    // responsibility for running the Connector instance in addition to any tasks it generates).
+    /**
+     * Connector Assignment V0:
+     * 
      +     *   Connector          => [String]
+     *   Tasks              => [Int32]
+     * 
      
+     *
+     * 
      Assignments for each worker are a set of connectors and tasks. These are categorized by
+     * connector ID. A sentinel task ID (CONNECTOR_TASK) is used to indicate the connector itself
+     * (i.e. that the assignment includes responsibility for running the Connector instance in
+     * addition to any tasks it generates).
      
+     */
     public static final Schema CONNECTOR_ASSIGNMENT_V0 = new Schema(
             new Field(CONNECTOR_KEY_NAME, Type.STRING),
             new Field(TASKS_KEY_NAME, new ArrayOf(Type.INT32)));
+
+    /**
+     * Assignment V0:
+     * 
      +     *   Error              => Int16
+     *   Leader             => [String]
+     *   LeaderUrl          => [String]
+     *   ConfigOffset       => Int64
+     *   Assignment         => [Connector Assignment]
+     * 
      
+     */
     public static final Schema ASSIGNMENT_V0 = new Schema(
             new Field(ERROR_KEY_NAME, Type.INT16),
             new Field(LEADER_KEY_NAME, Type.STRING),
@@ -70,6 +111,18 @@ public class ConnectProtocol {
             new Field(CONFIG_OFFSET_KEY_NAME, Type.INT64),
             new Field(ASSIGNMENT_KEY_NAME, new ArrayOf(CONNECTOR_ASSIGNMENT_V0)));
 
+    /**
+     * The fields are serialized in sequence as follows:
+     * Subscription V0:
+     * 
      +     *   Version            => Int16
+     *   Url                => [String]
+     *   ConfigOffset       => Int64
+     * 
      
+     *
+     * @param workerState the current state of the worker metadata
+     * @return the serialized state of the worker metadata
+     */
     public static ByteBuffer serializeMetadata(WorkerState workerState) {
         Struct struct = new Struct(CONFIG_STATE_V0);
         struct.set(URL_KEY_NAME, workerState.url());
@@ -81,6 +134,29 @@ public static ByteBuffer serializeMetadata(WorkerState workerState) {
         return buffer;
     }
 
+    /**
+     * Returns the collection of Connect protocols that are supported by this version along
+     * with their serialized metadata. The protocols are ordered by preference.
+     *
+     * @param workerState the current state of the worker metadata
+     * @return the collection of Connect protocol metadata
+     */
+    public static JoinGroupRequestProtocolCollection metadataRequest(WorkerState workerState) {
+        return new JoinGroupRequestProtocolCollection(Collections.singleton(
+                new JoinGroupRequestProtocol()
+                        .setName(EAGER.protocol())
+                        .setMetadata(ConnectProtocol.serializeMetadata(workerState).array()))
+                .iterator());
+    }
+
+    /**
+     * Given a byte buffer that contains protocol metadata return the deserialized form of the
+     * metadata.
+     *
+     * @param buffer A buffer containing the protocols metadata
+     * @return the deserialized metadata
+     * @throws SchemaException on incompatible Connect protocol version
+     */
     public static WorkerState deserializeMetadata(ByteBuffer buffer) {
         Struct header = CONNECT_PROTOCOL_HEADER_SCHEMA.read(buffer);
         Short version = header.getShort(VERSION_KEY_NAME);
@@ -91,6 +167,18 @@ public static WorkerState deserializeMetadata(ByteBuffer buffer) {
         return new WorkerState(url, configOffset);
     }
 
+    /**
+     * The fields are serialized in sequence as follows:
+     * Complete Assignment V0:
+     * 
      +     *   Version            => Int16
+     *   Error              => Int16
+     *   Leader             => [String]
+     *   LeaderUrl          => [String]
+     *   ConfigOffset       => Int64
+     *   Assignment         => [Connector Assignment]
+     * 
      
+     */
     public static ByteBuffer serializeAssignment(Assignment assignment) {
         Struct struct = new Struct(ASSIGNMENT_V0);
         struct.set(ERROR_KEY_NAME, assignment.error());
@@ -98,10 +186,10 @@ public static ByteBuffer serializeAssignment(Assignment assignment) {
         struct.set(LEADER_URL_KEY_NAME, assignment.leaderUrl());
         struct.set(CONFIG_OFFSET_KEY_NAME, assignment.offset());
         List taskAssignments = new ArrayList<>();
-        for (Map.Entry> connectorEntry : assignment.asMap().entrySet()) {
+        for (Map.Entry> connectorEntry : assignment.asMap().entrySet()) {
             Struct taskAssignment = new Struct(CONNECTOR_ASSIGNMENT_V0);
             taskAssignment.set(CONNECTOR_KEY_NAME, connectorEntry.getKey());
-            List tasks = connectorEntry.getValue();
+            Collection tasks = connectorEntry.getValue();
             taskAssignment.set(TASKS_KEY_NAME, tasks.toArray());
             taskAssignments.add(taskAssignment);
         }
@@ -114,6 +202,14 @@ public static ByteBuffer serializeAssignment(Assignment assignment) {
         return buffer;
     }
 
+    /**
+     * Given a byte buffer that contains an assignment as defined by this protocol, return the
+     * deserialized form of the assignment.
+     *
+     * @param buffer the buffer containing a serialized assignment
+     * @return the deserialized assignment
+     * @throws SchemaException on incompatible Connect protocol version
+     */
     public static Assignment deserializeAssignment(ByteBuffer buffer) {
         Struct header = CONNECT_PROTOCOL_HEADER_SCHEMA.read(buffer);
         Short version = header.getShort(VERSION_KEY_NAME);
@@ -139,6 +235,9 @@ public static Assignment deserializeAssignment(ByteBuffer buffer) {
         return new Assignment(error, leader, leaderUrl, offset, connectorIds, taskIds);
     }
 
+    /**
+     * A class that captures the deserialized form of a worker's metadata.
+     */
     public static class WorkerState {
         private final String url;
         private final long offset;
@@ -152,6 +251,11 @@ public String url() {
             return url;
         }
 
+        /**
+         * The most up-to-date (maximum) configuration offset according known to this worker.
+         *
+         * @return the configuration offset
+         */
         public long offset() {
             return offset;
         }
@@ -165,6 +269,9 @@ public String toString() {
         }
     }
 
+    /**
+     * The basic assignment of connectors and tasks introduced with V0 version of the Connect protocol.
+     */
     public static class Assignment {
         public static final short NO_ERROR = 0;
         // Configuration offsets mismatched in a way that the leader could not resolve. Workers should read to the end
@@ -175,49 +282,92 @@ public static class Assignment {
         private final String leader;
         private final String leaderUrl;
         private final long offset;
-        private final List connectorIds;
-        private final List taskIds;
+        private final Collection connectorIds;
+        private final Collection taskIds;
 
         /**
          * Create an assignment indicating responsibility for the given connector instances and task Ids.
-         * @param connectorIds list of connectors that the worker should instantiate and run
-         * @param taskIds list of task IDs that the worker should instantiate and run
+         *
+         * @param error error code for this assignment; {@code ConnectProtocol.Assignment.NO_ERROR}
+         *              indicates no error during assignment
+         * @param leader Connect group's leader Id; may be null only on the empty assignment
+         * @param leaderUrl Connect group's leader URL; may be null only on the empty assignment
+         * @param configOffset the most up-to-date configuration offset according to this assignment
+         * @param connectorIds list of connectors that the worker should instantiate and run; may not be null
+         * @param taskIds list of task IDs that the worker should instantiate and run; may not be null
          */
         public Assignment(short error, String leader, String leaderUrl, long configOffset,
-                          List connectorIds, List taskIds) {
+                          Collection connectorIds, Collection taskIds) {
             this.error = error;
             this.leader = leader;
             this.leaderUrl = leaderUrl;
             this.offset = configOffset;
-            this.taskIds = taskIds;
-            this.connectorIds = connectorIds;
+            this.connectorIds = Objects.requireNonNull(connectorIds,
+                    "Assigned connector IDs may be empty but not null");
+            this.taskIds = Objects.requireNonNull(taskIds,
+                    "Assigned task IDs may be empty but not null");
         }
 
+        /**
+         * Return the error code of this assignment; 0 signals successful assignment ({@code ConnectProtocol.Assignment.NO_ERROR}).
+         *
+         * @return the error code of the assignment
+         */
         public short error() {
             return error;
         }
 
+        /**
+         * Return the ID of the leader Connect worker in this assignment.
+         *
+         * @return the ID of the leader
+         */
         public String leader() {
             return leader;
         }
 
+        /**
+         * Return the URL to which the leader accepts requests from other members of the group.
+         *
+         * @return the leader URL
+         */
         public String leaderUrl() {
             return leaderUrl;
         }
 
+        /**
+         * Check if this assignment failed.
+         *
+         * @return true if this assignment failed; false otherwise
+         */
         public boolean failed() {
             return error != NO_ERROR;
         }
 
+        /**
+         * Return the most up-to-date offset in the configuration topic according to this assignment
+         *
+         * @return the configuration topic
+         */
         public long offset() {
             return offset;
         }
 
-        public List connectors() {
+        /**
+         * The connectors included in this assignment.
+         *
+         * @return the connectors
+         */
+        public Collection connectors() {
             return connectorIds;
         }
 
-        public List tasks() {
+        /**
+         * The tasks included in this assignment.
+         *
+         * @return the tasks
+         */
+        public Collection tasks() {
             return taskIds;
         }
 
@@ -233,11 +383,11 @@ public String toString() {
                     '}';
         }
 
-        private Map> asMap() {
+        protected Map> asMap() {
             // Using LinkedHashMap preserves the ordering, which is helpful for tests and debugging
-            Map> taskMap = new LinkedHashMap<>();
+            Map> taskMap = new LinkedHashMap<>();
             for (String connectorId : new HashSet<>(connectorIds)) {
-                List connectorTasks = taskMap.get(connectorId);
+                Collection connectorTasks = taskMap.get(connectorId);
                 if (connectorTasks == null) {
                     connectorTasks = new ArrayList<>();
                     taskMap.put(connectorId, connectorTasks);
@@ -246,7 +396,7 @@ private Map> asMap() {
             }
             for (ConnectorTaskId taskId : taskIds) {
                 String connectorId = taskId.connector();
-                List connectorTasks = taskMap.get(connectorId);
+                Collection connectorTasks = taskMap.get(connectorId);
                 if (connectorTasks == null) {
                     connectorTasks = new ArrayList<>();
                     taskMap.put(connectorId, connectorTasks);
@@ -264,5 +414,4 @@ private static void checkVersionCompatibility(short version) {
 
         // otherwise, assume versions can be parsed as V0
     }
-
 }
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/ConnectProtocolCompatibility.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/ConnectProtocolCompatibility.java
new file mode 100644
index 0000000000000..d618fe255345d
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/ConnectProtocolCompatibility.java
@@ -0,0 +1,148 @@
+/*
+ * 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.distributed;
+
+import java.util.Arrays;
+import java.util.Locale;
+
+import static org.apache.kafka.connect.runtime.distributed.ConnectProtocol.CONNECT_PROTOCOL_V0;
+import static org.apache.kafka.connect.runtime.distributed.IncrementalCooperativeConnectProtocol.CONNECT_PROTOCOL_V1;
+import static org.apache.kafka.connect.runtime.distributed.IncrementalCooperativeConnectProtocol.CONNECT_PROTOCOL_V2;
+
+/**
+ * An enumeration of the modes available to the worker to signal which Connect protocols are
+ * enabled at any time.
+ *
+ * {@code EAGER} signifies that this worker only supports prompt release of assigned connectors
+ * and tasks in every rebalance. Corresponds to Connect protocol V0.
+ *
+ * {@code COMPATIBLE} signifies that this worker supports both eager and incremental cooperative
+ * Connect protocols and will use the version that is elected by the Kafka broker coordinator
+ * during rebalance.
+ *
+ * {@code SESSIONED} signifies that this worker supports all of the above protocols in addition to
+ * a protocol that uses incremental cooperative rebalancing for worker assignment and uses session
+ * keys distributed via the config topic to verify internal REST requests
+ */
+public enum ConnectProtocolCompatibility {
+    EAGER {
+        @Override
+        public String protocol() {
+            return "default";
+        }
+
+        @Override
+        public short protocolVersion() {
+            return CONNECT_PROTOCOL_V0;
+        }
+    },
+
+    COMPATIBLE {
+        @Override
+        public String protocol() {
+            return "compatible";
+        }
+
+        @Override
+        public short protocolVersion() {
+            return CONNECT_PROTOCOL_V1;
+        }
+    },
+
+    SESSIONED {
+        @Override
+        public String protocol() {
+            return "sessioned";
+        }
+
+        @Override
+        public short protocolVersion() {
+            return CONNECT_PROTOCOL_V2;
+        }
+    };
+
+    /**
+     * Return the enum that corresponds to the name that is given as an argument;
+     * if no mapping is found {@code IllegalArgumentException} is thrown.
+     *
+     * @param name the name of the protocol compatibility mode
+     * @return the enum that corresponds to the protocol compatibility mode
+     */
+    public static ConnectProtocolCompatibility compatibility(String name) {
+        return Arrays.stream(ConnectProtocolCompatibility.values())
+                .filter(mode -> mode.name().equalsIgnoreCase(name))
+                .findFirst()
+                .orElseThrow(() -> new IllegalArgumentException(
+                        "Unknown Connect protocol compatibility mode: " + name));
+    }
+
+    /**
+     * Return the enum that corresponds to the Connect protocol version that is given as an argument;
+     * if no mapping is found {@code IllegalArgumentException} is thrown.
+     *
+     * @param protocolVersion the version of the protocol; for example,
+     * {@link ConnectProtocol#CONNECT_PROTOCOL_V0 CONNECT_PROTOCOL_V0}. May not be null
+     * @return the enum that corresponds to the protocol compatibility mode
+     */
+    public static ConnectProtocolCompatibility fromProtocolVersion(short protocolVersion) {
+        switch (protocolVersion) {
+            case CONNECT_PROTOCOL_V0:
+                return EAGER;
+            case CONNECT_PROTOCOL_V1:
+                return COMPATIBLE;
+            case CONNECT_PROTOCOL_V2:
+                return SESSIONED;
+            default:
+                throw new IllegalArgumentException("Unknown Connect protocol version: " + protocolVersion);
+        }
+    }
+
+    @Override
+    public String toString() {
+        return name().toLowerCase(Locale.ROOT);
+    }
+
+    /**
+     * Return the version of the protocol for this mode.
+     *
+     * @return the protocol version
+     */
+    public abstract short protocolVersion();
+
+    /**
+     * Return the name of the protocol that this mode will use in {@code ProtocolMetadata}.
+     *
+     * @return the protocol name
+     */
+    public abstract String protocol();
+
+    /**
+     * Return the enum that corresponds to the protocol name that is given as an argument;
+     * if no mapping is found {@code IllegalArgumentException} is thrown.
+     *
+     * @param protocolName the name of the connect protocol
+     * @return the enum that corresponds to the protocol compatibility mode that supports the
+     * given protocol
+     */
+    public static ConnectProtocolCompatibility fromProtocol(String protocolName) {
+        return Arrays.stream(ConnectProtocolCompatibility.values())
+                .filter(mode -> mode.protocol().equalsIgnoreCase(protocolName))
+                .findFirst()
+                .orElseThrow(() -> new IllegalArgumentException(
+                        "Not found Connect protocol compatibility mode for protocol: " + protocolName));
+    }
+}
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/DistributedConfig.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/DistributedConfig.java
index dc9017beeda75..00f6663724b9d 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/DistributedConfig.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/DistributedConfig.java
@@ -18,15 +18,24 @@
 
 import org.apache.kafka.clients.CommonClientConfigs;
 import org.apache.kafka.common.config.ConfigDef;
+import org.apache.kafka.common.config.ConfigException;
+import org.apache.kafka.common.utils.Utils;
 import org.apache.kafka.connect.runtime.WorkerConfig;
 
+import javax.crypto.KeyGenerator;
+import javax.crypto.Mac;
+import java.security.InvalidParameterException;
+import java.security.NoSuchAlgorithmException;
+import java.util.Collections;
+import java.util.List;
 import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
 
 import static org.apache.kafka.common.config.ConfigDef.Range.atLeast;
+import static org.apache.kafka.common.config.ConfigDef.Range.between;
 
 public class DistributedConfig extends WorkerConfig {
-    private static final ConfigDef CONFIG;
-
     /*
      * NOTE: DO NOT CHANGE EITHER CONFIG STRINGS OR THEIR JAVA VARIABLE NAMES AS
      * THESE ARE PART OF THE PUBLIC API AND CHANGE WILL BREAK USER CODE.
@@ -35,13 +44,13 @@ public class DistributedConfig extends WorkerConfig {
     /**
      * group.id
      */
-    public static final String GROUP_ID_CONFIG = "group.id";
+    public static final String GROUP_ID_CONFIG = CommonClientConfigs.GROUP_ID_CONFIG;
     private static final String GROUP_ID_DOC = "A unique string that identifies the Connect cluster group this worker belongs to.";
 
     /**
      * session.timeout.ms
      */
-    public static final String SESSION_TIMEOUT_MS_CONFIG = "session.timeout.ms";
+    public static final String SESSION_TIMEOUT_MS_CONFIG = CommonClientConfigs.SESSION_TIMEOUT_MS_CONFIG;
     private static final String SESSION_TIMEOUT_MS_DOC = "The timeout used to detect worker failures. " +
             "The worker sends periodic heartbeats to indicate its liveness to the broker. If no heartbeats are " +
             "received by the broker before the expiration of this session timeout, then the broker will remove the " +
@@ -52,7 +61,7 @@ public class DistributedConfig extends WorkerConfig {
     /**
      * heartbeat.interval.ms
      */
-    public static final String HEARTBEAT_INTERVAL_MS_CONFIG = "heartbeat.interval.ms";
+    public static final String HEARTBEAT_INTERVAL_MS_CONFIG = CommonClientConfigs.HEARTBEAT_INTERVAL_MS_CONFIG;
     private static final String HEARTBEAT_INTERVAL_MS_DOC = "The expected time between heartbeats to the group " +
             "coordinator when using Kafka's group management facilities. Heartbeats are used to ensure that the " +
             "worker's session stays active and to facilitate rebalancing when new members join or leave the group. " +
@@ -62,11 +71,8 @@ public class DistributedConfig extends WorkerConfig {
     /**
      * rebalance.timeout.ms
      */
-    public static final String REBALANCE_TIMEOUT_MS_CONFIG = "rebalance.timeout.ms";
-    private static final String REBALANCE_TIMEOUT_MS_DOC = "The maximum allowed time for each worker to join the group " +
-            "once a rebalance has begun. This is basically a limit on the amount of time needed for all tasks to " +
-            "flush any pending data and commit offsets. If the timeout is exceeded, then the worker will be removed " +
-            "from the group, which will cause offset commit failures.";
+    public static final String REBALANCE_TIMEOUT_MS_CONFIG = CommonClientConfigs.REBALANCE_TIMEOUT_MS_CONFIG;
+    private static final String REBALANCE_TIMEOUT_MS_DOC = CommonClientConfigs.REBALANCE_TIMEOUT_MS_DOC;
 
     /**
      * worker.sync.timeout.ms
@@ -132,147 +138,303 @@ public class DistributedConfig extends WorkerConfig {
     public static final String STATUS_STORAGE_REPLICATION_FACTOR_CONFIG = "status.storage.replication.factor";
     private static final String STATUS_STORAGE_REPLICATION_FACTOR_CONFIG_DOC = "Replication factor used when creating the status storage topic";
 
-    static {
-        CONFIG = baseConfigDef()
-                .define(GROUP_ID_CONFIG,
-                        ConfigDef.Type.STRING,
-                        ConfigDef.Importance.HIGH,
-                        GROUP_ID_DOC)
-                .define(SESSION_TIMEOUT_MS_CONFIG,
-                        ConfigDef.Type.INT,
-                        10000,
-                        ConfigDef.Importance.HIGH,
-                        SESSION_TIMEOUT_MS_DOC)
-                .define(REBALANCE_TIMEOUT_MS_CONFIG,
-                        ConfigDef.Type.INT,
-                        60000,
-                        ConfigDef.Importance.HIGH,
-                        REBALANCE_TIMEOUT_MS_DOC)
-                .define(HEARTBEAT_INTERVAL_MS_CONFIG,
-                        ConfigDef.Type.INT,
-                        3000,
-                        ConfigDef.Importance.HIGH,
-                        HEARTBEAT_INTERVAL_MS_DOC)
-                .define(CommonClientConfigs.METADATA_MAX_AGE_CONFIG,
-                        ConfigDef.Type.LONG,
-                        5 * 60 * 1000,
-                        atLeast(0),
-                        ConfigDef.Importance.LOW,
-                        CommonClientConfigs.METADATA_MAX_AGE_DOC)
-                .define(CommonClientConfigs.CLIENT_ID_CONFIG,
-                        ConfigDef.Type.STRING,
-                        "",
-                        ConfigDef.Importance.LOW,
-                        CommonClientConfigs.CLIENT_ID_DOC)
-                .define(CommonClientConfigs.SEND_BUFFER_CONFIG,
-                        ConfigDef.Type.INT,
-                        128 * 1024,
-                        atLeast(0),
-                        ConfigDef.Importance.MEDIUM,
-                        CommonClientConfigs.SEND_BUFFER_DOC)
-                .define(CommonClientConfigs.RECEIVE_BUFFER_CONFIG,
-                        ConfigDef.Type.INT,
-                        32 * 1024,
-                        atLeast(0),
-                        ConfigDef.Importance.MEDIUM,
-                        CommonClientConfigs.RECEIVE_BUFFER_DOC)
-                .define(CommonClientConfigs.RECONNECT_BACKOFF_MS_CONFIG,
-                        ConfigDef.Type.LONG,
-                        50L,
-                        atLeast(0L),
-                        ConfigDef.Importance.LOW,
-                        CommonClientConfigs.RECONNECT_BACKOFF_MS_DOC)
-                .define(CommonClientConfigs.RECONNECT_BACKOFF_MAX_MS_CONFIG,
-                        ConfigDef.Type.LONG,
-                        1000L,
-                        atLeast(0L),
-                        ConfigDef.Importance.LOW,
-                        CommonClientConfigs.RECONNECT_BACKOFF_MAX_MS_DOC)
-                .define(CommonClientConfigs.RETRY_BACKOFF_MS_CONFIG,
-                        ConfigDef.Type.LONG,
-                        100L,
-                        atLeast(0L),
-                        ConfigDef.Importance.LOW,
-                        CommonClientConfigs.RETRY_BACKOFF_MS_DOC)
-                .define(CommonClientConfigs.REQUEST_TIMEOUT_MS_CONFIG,
-                        ConfigDef.Type.INT,
-                        40 * 1000,
-                        atLeast(0),
-                        ConfigDef.Importance.MEDIUM,
-                        CommonClientConfigs.REQUEST_TIMEOUT_MS_DOC)
-                        /* default is set to be a bit lower than the server default (10 min), to avoid both client and server closing connection at same time */
-                .define(CommonClientConfigs.CONNECTIONS_MAX_IDLE_MS_CONFIG,
-                        ConfigDef.Type.LONG,
-                        9 * 60 * 1000,
-                        ConfigDef.Importance.MEDIUM,
-                        CommonClientConfigs.CONNECTIONS_MAX_IDLE_MS_DOC)
-                // security support
-                .define(CommonClientConfigs.SECURITY_PROTOCOL_CONFIG,
-                        ConfigDef.Type.STRING,
-                        CommonClientConfigs.DEFAULT_SECURITY_PROTOCOL,
-                        ConfigDef.Importance.MEDIUM,
-                        CommonClientConfigs.SECURITY_PROTOCOL_DOC)
-                .withClientSslSupport()
-                .withClientSaslSupport()
-                .define(WORKER_SYNC_TIMEOUT_MS_CONFIG,
-                        ConfigDef.Type.INT,
-                        3000,
-                        ConfigDef.Importance.MEDIUM,
-                        WORKER_SYNC_TIMEOUT_MS_DOC)
-                .define(WORKER_UNSYNC_BACKOFF_MS_CONFIG,
-                        ConfigDef.Type.INT,
-                        WORKER_UNSYNC_BACKOFF_MS_DEFAULT,
-                        ConfigDef.Importance.MEDIUM,
-                        WORKER_UNSYNC_BACKOFF_MS_DOC)
-                .define(OFFSET_STORAGE_TOPIC_CONFIG,
-                        ConfigDef.Type.STRING,
-                        ConfigDef.Importance.HIGH,
-                        OFFSET_STORAGE_TOPIC_CONFIG_DOC)
-                .define(OFFSET_STORAGE_PARTITIONS_CONFIG,
-                        ConfigDef.Type.INT,
-                        25,
-                        atLeast(1),
-                        ConfigDef.Importance.LOW,
-                        OFFSET_STORAGE_PARTITIONS_CONFIG_DOC)
-                .define(OFFSET_STORAGE_REPLICATION_FACTOR_CONFIG,
-                        ConfigDef.Type.SHORT,
-                        (short) 3,
-                        atLeast(1),
-                        ConfigDef.Importance.LOW,
-                        OFFSET_STORAGE_REPLICATION_FACTOR_CONFIG_DOC)
-                .define(CONFIG_TOPIC_CONFIG,
-                        ConfigDef.Type.STRING,
-                        ConfigDef.Importance.HIGH,
-                        CONFIG_TOPIC_CONFIG_DOC)
-                .define(CONFIG_STORAGE_REPLICATION_FACTOR_CONFIG,
-                        ConfigDef.Type.SHORT,
-                        (short) 3,
-                        atLeast(1),
-                        ConfigDef.Importance.LOW,
-                        CONFIG_STORAGE_REPLICATION_FACTOR_CONFIG_DOC)
-                .define(STATUS_STORAGE_TOPIC_CONFIG,
-                        ConfigDef.Type.STRING,
-                        ConfigDef.Importance.HIGH,
-                        STATUS_STORAGE_TOPIC_CONFIG_DOC)
-                .define(STATUS_STORAGE_PARTITIONS_CONFIG,
-                        ConfigDef.Type.INT,
-                        5,
-                        atLeast(1),
-                        ConfigDef.Importance.LOW,
-                        STATUS_STORAGE_PARTITIONS_CONFIG_DOC)
-                .define(STATUS_STORAGE_REPLICATION_FACTOR_CONFIG,
-                        ConfigDef.Type.SHORT,
-                        (short) 3,
-                        atLeast(1),
-                        ConfigDef.Importance.LOW,
-                        STATUS_STORAGE_REPLICATION_FACTOR_CONFIG_DOC);
+    /**
+     * connect.protocol
+     */
+    public static final String CONNECT_PROTOCOL_CONFIG = "connect.protocol";
+    public static final String CONNECT_PROTOCOL_DOC = "Compatibility mode for Kafka Connect Protocol";
+    public static final String CONNECT_PROTOCOL_DEFAULT = ConnectProtocolCompatibility.SESSIONED.toString();
+
+    /**
+     * scheduled.rebalance.max.delay.ms
+     */
+    public static final String SCHEDULED_REBALANCE_MAX_DELAY_MS_CONFIG = "scheduled.rebalance.max.delay.ms";
+    public static final String SCHEDULED_REBALANCE_MAX_DELAY_MS_DOC = "The maximum delay that is "
+            + "scheduled in order to wait for the return of one or more departed workers before "
+            + "rebalancing and reassigning their connectors and tasks to the group. During this "
+            + "period the connectors and tasks of the departed workers remain unassigned";
+    public static final int SCHEDULED_REBALANCE_MAX_DELAY_MS_DEFAULT = Math.toIntExact(TimeUnit.SECONDS.toMillis(300));
+
+    public static final String INTER_WORKER_KEY_GENERATION_ALGORITHM_CONFIG = "inter.worker.key.generation.algorithm";
+    public static final String INTER_WORKER_KEY_GENERATION_ALGORITHM_DOC = "The algorithm to use for generating internal request keys";
+    public static final String INTER_WORKER_KEY_GENERATION_ALGORITHM_DEFAULT = "HmacSHA256";
+
+    public static final String INTER_WORKER_KEY_SIZE_CONFIG = "inter.worker.key.size";
+    public static final String INTER_WORKER_KEY_SIZE_DOC = "The size of the key to use for signing internal requests, in bits. "
+        + "If null, the default key size for the key generation algorithm will be used.";
+    public static final Long INTER_WORKER_KEY_SIZE_DEFAULT = null;
+
+    public static final String INTER_WORKER_KEY_TTL_MS_CONFIG = "inter.worker.key.ttl.ms";
+    public static final String INTER_WORKER_KEY_TTL_MS_MS_DOC = "The TTL of generated session keys used for "
+        + "internal request validation (in milliseconds)";
+    public static final int INTER_WORKER_KEY_TTL_MS_MS_DEFAULT = Math.toIntExact(TimeUnit.HOURS.toMillis(1));
+
+    public static final String INTER_WORKER_SIGNATURE_ALGORITHM_CONFIG = "inter.worker.signature.algorithm";
+    public static final String INTER_WORKER_SIGNATURE_ALGORITHM_DOC = "The algorithm used to sign internal requests";
+    public static final String INTER_WORKER_SIGNATURE_ALGORITHM_DEFAULT = "HmacSHA256";
+
+    public static final String INTER_WORKER_VERIFICATION_ALGORITHMS_CONFIG = "inter.worker.verification.algorithms";
+    public static final String INTER_WORKER_VERIFICATION_ALGORITHMS_DOC = "A list of permitted algorithms for verifying internal requests";
+    public static final List INTER_WORKER_VERIFICATION_ALGORITHMS_DEFAULT = Collections.singletonList(INTER_WORKER_SIGNATURE_ALGORITHM_DEFAULT);
+
+    @SuppressWarnings("unchecked")
+    private static final ConfigDef CONFIG = baseConfigDef()
+            .define(GROUP_ID_CONFIG,
+                    ConfigDef.Type.STRING,
+                    ConfigDef.Importance.HIGH,
+                    GROUP_ID_DOC)
+            .define(SESSION_TIMEOUT_MS_CONFIG,
+                    ConfigDef.Type.INT,
+                    Math.toIntExact(TimeUnit.SECONDS.toMillis(10)),
+                    ConfigDef.Importance.HIGH,
+                    SESSION_TIMEOUT_MS_DOC)
+            .define(REBALANCE_TIMEOUT_MS_CONFIG,
+                    ConfigDef.Type.INT,
+                    Math.toIntExact(TimeUnit.MINUTES.toMillis(1)),
+                    ConfigDef.Importance.HIGH,
+                    REBALANCE_TIMEOUT_MS_DOC)
+            .define(HEARTBEAT_INTERVAL_MS_CONFIG,
+                    ConfigDef.Type.INT,
+                    Math.toIntExact(TimeUnit.SECONDS.toMillis(3)),
+                    ConfigDef.Importance.HIGH,
+                    HEARTBEAT_INTERVAL_MS_DOC)
+            .define(CommonClientConfigs.METADATA_MAX_AGE_CONFIG,
+                    ConfigDef.Type.LONG,
+                    TimeUnit.MINUTES.toMillis(5),
+                    atLeast(0),
+                    ConfigDef.Importance.LOW,
+                    CommonClientConfigs.METADATA_MAX_AGE_DOC)
+            .define(CommonClientConfigs.CLIENT_ID_CONFIG,
+                    ConfigDef.Type.STRING,
+                    "",
+                    ConfigDef.Importance.LOW,
+                    CommonClientConfigs.CLIENT_ID_DOC)
+            .define(CommonClientConfigs.SEND_BUFFER_CONFIG,
+                    ConfigDef.Type.INT,
+                    128 * 1024,
+                    atLeast(0),
+                    ConfigDef.Importance.MEDIUM,
+                    CommonClientConfigs.SEND_BUFFER_DOC)
+            .define(CommonClientConfigs.RECEIVE_BUFFER_CONFIG,
+                    ConfigDef.Type.INT,
+                    32 * 1024,
+                    atLeast(0),
+                    ConfigDef.Importance.MEDIUM,
+                    CommonClientConfigs.RECEIVE_BUFFER_DOC)
+            .define(CommonClientConfigs.RECONNECT_BACKOFF_MS_CONFIG,
+                    ConfigDef.Type.LONG,
+                    50L,
+                    atLeast(0L),
+                    ConfigDef.Importance.LOW,
+                    CommonClientConfigs.RECONNECT_BACKOFF_MS_DOC)
+            .define(CommonClientConfigs.RECONNECT_BACKOFF_MAX_MS_CONFIG,
+                    ConfigDef.Type.LONG,
+                    TimeUnit.SECONDS.toMillis(1),
+                    atLeast(0L),
+                    ConfigDef.Importance.LOW,
+                    CommonClientConfigs.RECONNECT_BACKOFF_MAX_MS_DOC)
+            .define(CommonClientConfigs.RETRY_BACKOFF_MS_CONFIG,
+                    ConfigDef.Type.LONG,
+                    100L,
+                    atLeast(0L),
+                    ConfigDef.Importance.LOW,
+                    CommonClientConfigs.RETRY_BACKOFF_MS_DOC)
+            .define(CommonClientConfigs.REQUEST_TIMEOUT_MS_CONFIG,
+                    ConfigDef.Type.INT,
+                    Math.toIntExact(TimeUnit.SECONDS.toMillis(40)),
+                    atLeast(0),
+                    ConfigDef.Importance.MEDIUM,
+                    CommonClientConfigs.REQUEST_TIMEOUT_MS_DOC)
+                    /* default is set to be a bit lower than the server default (10 min), to avoid both client and server closing connection at same time */
+            .define(CommonClientConfigs.CONNECTIONS_MAX_IDLE_MS_CONFIG,
+                    ConfigDef.Type.LONG,
+                    TimeUnit.MINUTES.toMillis(9),
+                    ConfigDef.Importance.MEDIUM,
+                    CommonClientConfigs.CONNECTIONS_MAX_IDLE_MS_DOC)
+            // security support
+            .define(CommonClientConfigs.SECURITY_PROTOCOL_CONFIG,
+                    ConfigDef.Type.STRING,
+                    CommonClientConfigs.DEFAULT_SECURITY_PROTOCOL,
+                    ConfigDef.Importance.MEDIUM,
+                    CommonClientConfigs.SECURITY_PROTOCOL_DOC)
+            .withClientSaslSupport()
+            .define(WORKER_SYNC_TIMEOUT_MS_CONFIG,
+                    ConfigDef.Type.INT,
+                    3000,
+                    ConfigDef.Importance.MEDIUM,
+                    WORKER_SYNC_TIMEOUT_MS_DOC)
+            .define(WORKER_UNSYNC_BACKOFF_MS_CONFIG,
+                    ConfigDef.Type.INT,
+                    WORKER_UNSYNC_BACKOFF_MS_DEFAULT,
+                    ConfigDef.Importance.MEDIUM,
+                    WORKER_UNSYNC_BACKOFF_MS_DOC)
+            .define(OFFSET_STORAGE_TOPIC_CONFIG,
+                    ConfigDef.Type.STRING,
+                    ConfigDef.Importance.HIGH,
+                    OFFSET_STORAGE_TOPIC_CONFIG_DOC)
+            .define(OFFSET_STORAGE_PARTITIONS_CONFIG,
+                    ConfigDef.Type.INT,
+                    25,
+                    atLeast(1),
+                    ConfigDef.Importance.LOW,
+                    OFFSET_STORAGE_PARTITIONS_CONFIG_DOC)
+            .define(OFFSET_STORAGE_REPLICATION_FACTOR_CONFIG,
+                    ConfigDef.Type.SHORT,
+                    (short) 3,
+                    atLeast(1),
+                    ConfigDef.Importance.LOW,
+                    OFFSET_STORAGE_REPLICATION_FACTOR_CONFIG_DOC)
+            .define(CONFIG_TOPIC_CONFIG,
+                    ConfigDef.Type.STRING,
+                    ConfigDef.Importance.HIGH,
+                    CONFIG_TOPIC_CONFIG_DOC)
+            .define(CONFIG_STORAGE_REPLICATION_FACTOR_CONFIG,
+                    ConfigDef.Type.SHORT,
+                    (short) 3,
+                    atLeast(1),
+                    ConfigDef.Importance.LOW,
+                    CONFIG_STORAGE_REPLICATION_FACTOR_CONFIG_DOC)
+            .define(STATUS_STORAGE_TOPIC_CONFIG,
+                    ConfigDef.Type.STRING,
+                    ConfigDef.Importance.HIGH,
+                    STATUS_STORAGE_TOPIC_CONFIG_DOC)
+            .define(STATUS_STORAGE_PARTITIONS_CONFIG,
+                    ConfigDef.Type.INT,
+                    5,
+                    atLeast(1),
+                    ConfigDef.Importance.LOW,
+                    STATUS_STORAGE_PARTITIONS_CONFIG_DOC)
+            .define(STATUS_STORAGE_REPLICATION_FACTOR_CONFIG,
+                    ConfigDef.Type.SHORT,
+                    (short) 3,
+                    atLeast(1),
+                    ConfigDef.Importance.LOW,
+                    STATUS_STORAGE_REPLICATION_FACTOR_CONFIG_DOC)
+            .define(CONNECT_PROTOCOL_CONFIG,
+                    ConfigDef.Type.STRING,
+                    CONNECT_PROTOCOL_DEFAULT,
+                    ConfigDef.LambdaValidator.with(
+                        (name, value) -> {
+                            try {
+                                ConnectProtocolCompatibility.compatibility((String) value);
+                            } catch (Throwable t) {
+                                throw new ConfigException(name, value, "Invalid Connect protocol "
+                                        + "compatibility");
+                            }
+                        },
+                        () -> "[" + Utils.join(ConnectProtocolCompatibility.values(), ", ") + "]"),
+                    ConfigDef.Importance.LOW,
+                    CONNECT_PROTOCOL_DOC)
+            .define(SCHEDULED_REBALANCE_MAX_DELAY_MS_CONFIG,
+                    ConfigDef.Type.INT,
+                    SCHEDULED_REBALANCE_MAX_DELAY_MS_DEFAULT,
+                    between(0, Integer.MAX_VALUE),
+                    ConfigDef.Importance.LOW,
+                    SCHEDULED_REBALANCE_MAX_DELAY_MS_DOC)
+            .define(INTER_WORKER_KEY_TTL_MS_CONFIG,
+                    ConfigDef.Type.INT,
+                    INTER_WORKER_KEY_TTL_MS_MS_DEFAULT,
+                    between(0, Integer.MAX_VALUE),
+                    ConfigDef.Importance.LOW,
+                    INTER_WORKER_KEY_TTL_MS_MS_DOC)
+            .define(INTER_WORKER_KEY_GENERATION_ALGORITHM_CONFIG,
+                    ConfigDef.Type.STRING,
+                    INTER_WORKER_KEY_GENERATION_ALGORITHM_DEFAULT,
+                    ConfigDef.LambdaValidator.with(
+                        (name, value) -> validateKeyAlgorithm(name, (String) value),
+                        () -> "Any KeyGenerator algorithm supported by the worker JVM"
+                    ),
+                    ConfigDef.Importance.LOW,
+                    INTER_WORKER_KEY_GENERATION_ALGORITHM_DOC)
+            .define(INTER_WORKER_KEY_SIZE_CONFIG,
+                    ConfigDef.Type.INT,
+                    INTER_WORKER_KEY_SIZE_DEFAULT,
+                    ConfigDef.Importance.LOW,
+                    INTER_WORKER_KEY_SIZE_DOC)
+            .define(INTER_WORKER_SIGNATURE_ALGORITHM_CONFIG,
+                    ConfigDef.Type.STRING,
+                    INTER_WORKER_SIGNATURE_ALGORITHM_DEFAULT,
+                    ConfigDef.LambdaValidator.with(
+                        (name, value) -> validateSignatureAlgorithm(name, (String) value),
+                        () -> "Any MAC algorithm supported by the worker JVM"),
+                    ConfigDef.Importance.LOW,
+                    INTER_WORKER_SIGNATURE_ALGORITHM_DOC)
+            .define(INTER_WORKER_VERIFICATION_ALGORITHMS_CONFIG,
+                    ConfigDef.Type.LIST,
+                    INTER_WORKER_VERIFICATION_ALGORITHMS_DEFAULT,
+                    ConfigDef.LambdaValidator.with(
+                        (name, value) -> validateSignatureAlgorithms(name, (List) value),
+                        () -> "A list of one or more MAC algorithms, each supported by the worker JVM"
+                    ),
+                    ConfigDef.Importance.LOW,
+                    INTER_WORKER_VERIFICATION_ALGORITHMS_DOC);
+
+    @Override
+    public Integer getRebalanceTimeout() {
+        return getInt(DistributedConfig.REBALANCE_TIMEOUT_MS_CONFIG);
     }
 
     public DistributedConfig(Map props) {
         super(CONFIG, props);
+        getInternalRequestKeyGenerator(); // Check here for a valid key size + key algorithm to fail fast if either are invalid
+        validateKeyAlgorithmAndVerificationAlgorithms();
     }
 
     public static void main(String[] args) {
-        System.out.println(CONFIG.toHtmlTable());
+        System.out.println(CONFIG.toHtml());
+    }
+
+    public KeyGenerator getInternalRequestKeyGenerator() {
+        try {
+            KeyGenerator result = KeyGenerator.getInstance(getString(INTER_WORKER_KEY_GENERATION_ALGORITHM_CONFIG));
+            Optional.ofNullable(getInt(INTER_WORKER_KEY_SIZE_CONFIG)).ifPresent(result::init);
+            return result;
+        } catch (NoSuchAlgorithmException | InvalidParameterException e) {
+            throw new ConfigException(String.format(
+                "Unable to create key generator with algorithm %s and key size %d: %s",
+                getString(INTER_WORKER_KEY_GENERATION_ALGORITHM_CONFIG),
+                getInt(INTER_WORKER_KEY_SIZE_CONFIG),
+                e.getMessage()
+            ));
+        }
+    }
+
+    private void validateKeyAlgorithmAndVerificationAlgorithms() {
+        String keyAlgorithm = getString(INTER_WORKER_KEY_GENERATION_ALGORITHM_CONFIG);
+        List verificationAlgorithms = getList(INTER_WORKER_VERIFICATION_ALGORITHMS_CONFIG);
+        if (!verificationAlgorithms.contains(keyAlgorithm)) {
+            throw new ConfigException(
+                INTER_WORKER_KEY_GENERATION_ALGORITHM_CONFIG,
+                keyAlgorithm,
+                String.format("Key generation algorithm must be present in %s list", INTER_WORKER_VERIFICATION_ALGORITHMS_CONFIG)
+            );
+        }
+    }
+
+    private static void validateSignatureAlgorithms(String configName, List algorithms) {
+        if (algorithms.isEmpty()) {
+            throw new ConfigException(
+                configName,
+                algorithms,
+                "At least one signature verification algorithm must be provided"
+            );
+        }
+        algorithms.forEach(algorithm -> validateSignatureAlgorithm(configName, algorithm));
+    }
+
+    private static void validateSignatureAlgorithm(String configName, String algorithm) {
+        try {
+            Mac.getInstance(algorithm);
+        } catch (NoSuchAlgorithmException e) {
+            throw new ConfigException(configName, algorithm, e.getMessage());
+        }
+    }
+
+    private static void validateKeyAlgorithm(String configName, String algorithm) {
+        try {
+            KeyGenerator.getInstance(algorithm);
+        } catch (NoSuchAlgorithmException e) {
+            throw new ConfigException(configName, algorithm, e.getMessage());
+        }
     }
 }
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/DistributedHerder.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/DistributedHerder.java
index 711b6c9dfc80a..91139c2241e87 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/DistributedHerder.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/DistributedHerder.java
@@ -16,18 +16,20 @@
  */
 package org.apache.kafka.connect.runtime.distributed;
 
+import org.apache.kafka.clients.CommonClientConfigs;
 import org.apache.kafka.common.config.ConfigDef;
 import org.apache.kafka.common.config.ConfigValue;
 import org.apache.kafka.common.errors.WakeupException;
 import org.apache.kafka.common.metrics.Sensor;
 import org.apache.kafka.common.metrics.stats.Avg;
+import org.apache.kafka.common.metrics.stats.CumulativeSum;
 import org.apache.kafka.common.metrics.stats.Max;
-import org.apache.kafka.common.metrics.stats.Total;
 import org.apache.kafka.common.utils.Exit;
+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.connect.connector.Connector;
 import org.apache.kafka.connect.connector.ConnectorContext;
+import org.apache.kafka.connect.connector.policy.ConnectorClientConfigOverridePolicy;
 import org.apache.kafka.connect.errors.AlreadyExistsException;
 import org.apache.kafka.connect.errors.ConnectException;
 import org.apache.kafka.connect.errors.NotFoundException;
@@ -39,14 +41,19 @@
 import org.apache.kafka.connect.runtime.ConnectorConfig;
 import org.apache.kafka.connect.runtime.HerderConnectorContext;
 import org.apache.kafka.connect.runtime.HerderRequest;
+import org.apache.kafka.connect.runtime.SessionKey;
 import org.apache.kafka.connect.runtime.SinkConnectorConfig;
 import org.apache.kafka.connect.runtime.SourceConnectorConfig;
 import org.apache.kafka.connect.runtime.TargetState;
+import org.apache.kafka.connect.runtime.TaskStatus;
 import org.apache.kafka.connect.runtime.Worker;
+import org.apache.kafka.connect.runtime.rest.InternalRequestSignature;
 import org.apache.kafka.connect.runtime.rest.RestClient;
 import org.apache.kafka.connect.runtime.rest.RestServer;
 import org.apache.kafka.connect.runtime.rest.entities.ConnectorInfo;
 import org.apache.kafka.connect.runtime.rest.entities.TaskInfo;
+import org.apache.kafka.connect.runtime.rest.errors.BadRequestException;
+import org.apache.kafka.connect.runtime.rest.errors.ConnectRestException;
 import org.apache.kafka.connect.sink.SinkConnector;
 import org.apache.kafka.connect.storage.ConfigBackingStore;
 import org.apache.kafka.connect.storage.StatusBackingStore;
@@ -54,8 +61,10 @@
 import org.apache.kafka.connect.util.ConnectorTaskId;
 import org.apache.kafka.connect.util.SinkUtils;
 import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 
+import javax.crypto.KeyGenerator;
+import javax.crypto.SecretKey;
+import javax.ws.rs.core.Response;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
@@ -76,7 +85,14 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.stream.Collectors;
+
+import static org.apache.kafka.connect.runtime.distributed.ConnectProtocol.CONNECT_PROTOCOL_V0;
+import static org.apache.kafka.connect.runtime.distributed.ConnectProtocolCompatibility.EAGER;
+import static org.apache.kafka.connect.runtime.distributed.IncrementalCooperativeConnectProtocol.CONNECT_PROTOCOL_V2;
 
 /**
  * 
      
@@ -108,12 +124,14 @@
  * 
      
  */
 public class DistributedHerder extends AbstractHerder implements Runnable {
-    private static final Logger log = LoggerFactory.getLogger(DistributedHerder.class);
+    private static final AtomicInteger CONNECT_CLIENT_ID_SEQUENCE = new AtomicInteger(1);
+    private final Logger log;
 
     private static final long FORWARD_REQUEST_SHUTDOWN_TIMEOUT_MS = TimeUnit.SECONDS.toMillis(10);
     private static final long START_AND_STOP_SHUTDOWN_TIMEOUT_MS = TimeUnit.SECONDS.toMillis(1);
     private static final long RECONFIGURE_CONNECTOR_TASKS_BACKOFF_MS = 250;
     private static final int START_STOP_THREAD_POOL_SIZE = 8;
+    private static final short BACKOFF_RETRIES = 5;
 
     private final AtomicLong requestSeqNum = new AtomicLong();
 
@@ -124,6 +142,10 @@ public class DistributedHerder extends AbstractHerder implements Runnable {
     private final int workerSyncTimeoutMs;
     private final long workerTasksShutdownTimeoutMs;
     private final int workerUnsyncBackoffMs;
+    private final int keyRotationIntervalMs;
+    private final String requestSignatureAlgorithm;
+    private final List keySignatureVerificationAlgorithms;
+    private final KeyGenerator keyGenerator;
 
     private final ExecutorService herderExecutor;
     private final ExecutorService forwardRequestExecutor;
@@ -134,7 +156,9 @@ public class DistributedHerder extends AbstractHerder implements Runnable {
     // Track enough information about the current membership state to be able to determine which requests via the API
     // and the from other nodes are safe to process
     private boolean rebalanceResolved;
-    private ConnectProtocol.Assignment assignment;
+    private ExtendedAssignment runningAssignment = ExtendedAssignment.empty();
+    private Set tasksToRestart = new HashSet<>();
+    private ExtendedAssignment assignment;
     private boolean canReadConfigs;
     private ClusterConfigState configState;
 
@@ -144,11 +168,17 @@ public class DistributedHerder extends AbstractHerder implements Runnable {
     // Config updates can be collected and applied together when possible. Also, we need to take care to rebalance when
     // needed (e.g. task reconfiguration, which requires everyone to coordinate offset commits).
     private Set connectorConfigUpdates = new HashSet<>();
+    private Set taskConfigUpdates = new HashSet<>();
     // Similarly collect target state changes (when observed by the config storage listener) for handling in the
     // herder's main thread.
     private Set connectorTargetStateChanges = new HashSet<>();
     private boolean needsReconfigRebalance;
     private volatile int generation;
+    private volatile long scheduledRebalance;
+    private SecretKey sessionKey;
+    private volatile long keyExpiration;
+    private short currentProtocolVersion;
+    private short backoffRetries;
 
     private final DistributedConfig config;
 
@@ -158,8 +188,10 @@ public DistributedHerder(DistributedConfig config,
                              String kafkaClusterId,
                              StatusBackingStore statusBackingStore,
                              ConfigBackingStore configBackingStore,
-                             String restUrl) {
-        this(config, worker, worker.workerId(), kafkaClusterId, statusBackingStore, configBackingStore, null, restUrl, worker.metrics(), time);
+                             String restUrl,
+                             ConnectorClientConfigOverridePolicy connectorClientConfigOverridePolicy) {
+        this(config, worker, worker.workerId(), kafkaClusterId, statusBackingStore, configBackingStore, null, restUrl, worker.metrics(),
+             time, connectorClientConfigOverridePolicy);
         configBackingStore.setUpdateListener(new ConfigUpdateListener());
     }
 
@@ -173,8 +205,9 @@ public DistributedHerder(DistributedConfig config,
                       WorkerGroupMember member,
                       String restUrl,
                       ConnectMetrics metrics,
-                      Time time) {
-        super(worker, workerId, kafkaClusterId, statusBackingStore, configBackingStore);
+                      Time time,
+                      ConnectorClientConfigOverridePolicy connectorClientConfigOverridePolicy) {
+        super(worker, workerId, kafkaClusterId, statusBackingStore, configBackingStore, connectorClientConfigOverridePolicy);
 
         this.time = time;
         this.herderMetrics = new HerderMetrics(metrics);
@@ -182,12 +215,26 @@ public DistributedHerder(DistributedConfig config,
         this.workerSyncTimeoutMs = config.getInt(DistributedConfig.WORKER_SYNC_TIMEOUT_MS_CONFIG);
         this.workerTasksShutdownTimeoutMs = config.getLong(DistributedConfig.TASK_SHUTDOWN_GRACEFUL_TIMEOUT_MS_CONFIG);
         this.workerUnsyncBackoffMs = config.getInt(DistributedConfig.WORKER_UNSYNC_BACKOFF_MS_CONFIG);
-        this.member = member != null ? member : new WorkerGroupMember(config, restUrl, this.configBackingStore, new RebalanceListener(), time);
+        this.requestSignatureAlgorithm = config.getString(DistributedConfig.INTER_WORKER_SIGNATURE_ALGORITHM_CONFIG);
+        this.keyRotationIntervalMs = config.getInt(DistributedConfig.INTER_WORKER_KEY_TTL_MS_CONFIG);
+        this.keySignatureVerificationAlgorithms = config.getList(DistributedConfig.INTER_WORKER_VERIFICATION_ALGORITHMS_CONFIG);
+        this.keyGenerator = config.getInternalRequestKeyGenerator();
+
+        String clientIdConfig = config.getString(CommonClientConfigs.CLIENT_ID_CONFIG);
+        String clientId = clientIdConfig.length() <= 0 ? "connect-" + CONNECT_CLIENT_ID_SEQUENCE.getAndIncrement() : clientIdConfig;
+        LogContext logContext = new LogContext("[Worker clientId=" + clientId + ", groupId=" + this.workerGroupId + "] ");
+        log = logContext.logger(DistributedHerder.class);
+
+        this.member = member != null
+                      ? member
+                      : new WorkerGroupMember(config, restUrl, this.configBackingStore,
+                              new RebalanceListener(time), time, clientId, logContext);
+
         this.herderExecutor = new ThreadPoolExecutor(1, 1, 0L, TimeUnit.MILLISECONDS, new LinkedBlockingDeque(1),
                 new ThreadFactory() {
                     @Override
                     public Thread newThread(Runnable herder) {
-                        return new Thread(herder, "DistributedHerder");
+                        return new Thread(herder, "DistributedHerder-" + clientId);
                     }
                 });
         this.forwardRequestExecutor = Executors.newSingleThreadExecutor();
@@ -199,6 +246,26 @@ public Thread newThread(Runnable herder) {
         rebalanceResolved = true; // If we still need to follow up after a rebalance occurred, starting up tasks
         needsReconfigRebalance = false;
         canReadConfigs = true; // We didn't try yet, but Configs are readable until proven otherwise
+        scheduledRebalance = Long.MAX_VALUE;
+        keyExpiration = Long.MAX_VALUE;
+        sessionKey = null;
+        backoffRetries = BACKOFF_RETRIES;
+
+        currentProtocolVersion = ConnectProtocolCompatibility.compatibility(
+            config.getString(DistributedConfig.CONNECT_PROTOCOL_CONFIG)
+        ).protocolVersion();
+        if (!internalRequestValidationEnabled(currentProtocolVersion)) {
+            log.warn(
+                "Internal request verification will be disabled for this cluster as this worker's {} configuration has been set to '{}'. "
+                    + "If this is not intentional, either remove the '{}' configuration from the worker config file or change its value "
+                    + "to '{}'. If this configuration is left as-is, the cluster will be insecure; for more information, see KIP-507: "
+                    + "https://cwiki.apache.org/confluence/display/KAFKA/KIP-507%3A+Securing+Internal+Connect+REST+Endpoints",
+                DistributedConfig.CONNECT_PROTOCOL_CONFIG,
+                config.getString(DistributedConfig.CONNECT_PROTOCOL_CONFIG),
+                DistributedConfig.CONNECT_PROTOCOL_CONFIG,
+                ConnectProtocolCompatibility.SESSIONED.name()
+            );
+        }
     }
 
     @Override
@@ -252,8 +319,17 @@ public void tick() {
             return;
         }
 
+        long now = time.milliseconds();
+
+        if (checkForKeyRotation(now)) {
+            keyExpiration = Long.MAX_VALUE;
+            configBackingStore.putSessionKey(new SessionKey(
+                keyGenerator.generateKey(),
+                now
+            ));
+        }
+
         // Process any external requests
-        final long now = time.milliseconds();
         long nextRequestTimeoutMs = Long.MAX_VALUE;
         while (true) {
             final DistributedHerderRequest next = peekWithoutException();
@@ -274,50 +350,60 @@ public void tick() {
             }
         }
 
+        if (scheduledRebalance < Long.MAX_VALUE) {
+            nextRequestTimeoutMs = Math.min(nextRequestTimeoutMs, Math.max(scheduledRebalance - now, 0));
+            rebalanceResolved = false;
+            log.debug("Scheduled rebalance at: {} (now: {} nextRequestTimeoutMs: {}) ",
+                    scheduledRebalance, now, nextRequestTimeoutMs);
+        }
+        if (internalRequestValidationEnabled() && keyExpiration < Long.MAX_VALUE) {
+            nextRequestTimeoutMs = Math.min(nextRequestTimeoutMs, Math.max(keyExpiration - now, 0));
+            log.debug("Scheduled next key rotation at: {} (now: {} nextRequestTimeoutMs: {}) ",
+                    keyExpiration, now, nextRequestTimeoutMs);
+        }
+
         // Process any configuration updates
-        Set connectorConfigUpdatesCopy = null;
-        Set connectorTargetStateChangesCopy = null;
-        synchronized (this) {
-            if (needsReconfigRebalance || !connectorConfigUpdates.isEmpty() || !connectorTargetStateChanges.isEmpty()) {
-                // Connector reconfigs only need local updates since there is no coordination between workers required.
-                // However, if connectors were added or removed, work needs to be rebalanced since we have more work
-                // items to distribute among workers.
-                configState = configBackingStore.snapshot();
-
-                if (needsReconfigRebalance) {
-                    // Task reconfigs require a rebalance. Request the rebalance, clean out state, and then restart
-                    // this loop, which will then ensure the rebalance occurs without any other requests being
-                    // processed until it completes.
-                    member.requestRejoin();
-                    // Any connector config updates or target state changes will be addressed during the rebalance too
-                    connectorConfigUpdates.clear();
-                    connectorTargetStateChanges.clear();
-                    needsReconfigRebalance = false;
-                    return;
-                } else {
-                    if (!connectorConfigUpdates.isEmpty()) {
-                        // We can't start/stop while locked since starting connectors can cause task updates that will
-                        // require writing configs, which in turn make callbacks into this class from another thread that
-                        // require acquiring a lock. This leads to deadlock. Instead, just copy the info we need and process
-                        // the updates after unlocking.
-                        connectorConfigUpdatesCopy = connectorConfigUpdates;
-                        connectorConfigUpdates = new HashSet<>();
-                    }
+        AtomicReference> connectorConfigUpdatesCopy = new AtomicReference<>();
+        AtomicReference> connectorTargetStateChangesCopy = new AtomicReference<>();
+        AtomicReference> taskConfigUpdatesCopy = new AtomicReference<>();
+
+        boolean shouldReturn;
+        if (member.currentProtocolVersion() == CONNECT_PROTOCOL_V0) {
+            shouldReturn = updateConfigsWithEager(connectorConfigUpdatesCopy,
+                    connectorTargetStateChangesCopy);
+            // With eager protocol we should return immediately if needsReconfigRebalance has
+            // been set to retain the old workflow
+            if (shouldReturn) {
+                return;
+            }
 
-                    if (!connectorTargetStateChanges.isEmpty()) {
-                        // Similarly for target state changes which can cause connectors to be restarted
-                        connectorTargetStateChangesCopy = connectorTargetStateChanges;
-                        connectorTargetStateChanges = new HashSet<>();
-                    }
-                }
+            if (connectorConfigUpdatesCopy.get() != null) {
+                processConnectorConfigUpdates(connectorConfigUpdatesCopy.get());
             }
-        }
 
-        if (connectorConfigUpdatesCopy != null)
-            processConnectorConfigUpdates(connectorConfigUpdatesCopy);
+            if (connectorTargetStateChangesCopy.get() != null) {
+                processTargetStateChanges(connectorTargetStateChangesCopy.get());
+            }
+        } else {
+            shouldReturn = updateConfigsWithIncrementalCooperative(connectorConfigUpdatesCopy,
+                    connectorTargetStateChangesCopy, taskConfigUpdatesCopy);
 
-        if (connectorTargetStateChangesCopy != null)
-            processTargetStateChanges(connectorTargetStateChangesCopy);
+            if (connectorConfigUpdatesCopy.get() != null) {
+                processConnectorConfigUpdates(connectorConfigUpdatesCopy.get());
+            }
+
+            if (connectorTargetStateChangesCopy.get() != null) {
+                processTargetStateChanges(connectorTargetStateChangesCopy.get());
+            }
+
+            if (taskConfigUpdatesCopy.get() != null) {
+                processTaskConfigUpdatesWithIncrementalCooperative(taskConfigUpdatesCopy.get());
+            }
+
+            if (shouldReturn) {
+                return;
+            }
+        }
 
         // Let the group take any actions it needs to
         try {
@@ -329,6 +415,120 @@ public void tick() {
         }
     }
 
+    private synchronized boolean checkForKeyRotation(long now) {
+        if (internalRequestValidationEnabled()) {
+            if (isLeader()) {
+                if (sessionKey == null) {
+                    log.debug("Internal request signing is enabled but no session key has been distributed yet. "
+                        + "Distributing new key now.");
+                    return true;
+                } else if (keyExpiration <= now) {
+                    log.debug("Existing key has expired. Distributing new key now.");
+                    return true;
+                } else if (!sessionKey.getAlgorithm().equals(keyGenerator.getAlgorithm())
+                        || sessionKey.getEncoded().length != keyGenerator.generateKey().getEncoded().length) {
+                    log.debug("Previously-distributed key uses different algorithm/key size "
+                        + "than required by current worker configuration. Distributing new key now.");
+                    return true;
+                }
+            } else if (sessionKey == null && configState.sessionKey() != null) {
+                // This happens on startup for follower workers; the snapshot contains the session key,
+                // but no callback in the config update listener has been fired for it yet.
+                sessionKey = configState.sessionKey().key();
+            }
+        }
+        return false;
+    }
+
+    private synchronized boolean updateConfigsWithEager(AtomicReference> connectorConfigUpdatesCopy,
+                                                        AtomicReference> connectorTargetStateChangesCopy) {
+        // This branch is here to avoid creating a snapshot if not needed
+        if (needsReconfigRebalance
+                || !connectorConfigUpdates.isEmpty()
+                || !connectorTargetStateChanges.isEmpty()) {
+            // Connector reconfigs only need local updates since there is no coordination between workers required.
+            // However, if connectors were added or removed, work needs to be rebalanced since we have more work
+            // items to distribute among workers.
+            configState = configBackingStore.snapshot();
+
+            if (needsReconfigRebalance) {
+                // Task reconfigs require a rebalance. Request the rebalance, clean out state, and then restart
+                // this loop, which will then ensure the rebalance occurs without any other requests being
+                // processed until it completes.
+                log.debug("Requesting rebalance due to reconfiguration of tasks (needsReconfigRebalance: {})",
+                        needsReconfigRebalance);
+                member.requestRejoin();
+                needsReconfigRebalance = false;
+                // Any connector config updates or target state changes will be addressed during the rebalance too
+                connectorConfigUpdates.clear();
+                connectorTargetStateChanges.clear();
+                return true;
+            } else {
+                if (!connectorConfigUpdates.isEmpty()) {
+                    // We can't start/stop while locked since starting connectors can cause task updates that will
+                    // require writing configs, which in turn make callbacks into this class from another thread that
+                    // require acquiring a lock. This leads to deadlock. Instead, just copy the info we need and process
+                    // the updates after unlocking.
+                    connectorConfigUpdatesCopy.set(connectorConfigUpdates);
+                    connectorConfigUpdates = new HashSet<>();
+                }
+
+                if (!connectorTargetStateChanges.isEmpty()) {
+                    // Similarly for target state changes which can cause connectors to be restarted
+                    connectorTargetStateChangesCopy.set(connectorTargetStateChanges);
+                    connectorTargetStateChanges = new HashSet<>();
+                }
+            }
+        }
+        return false;
+    }
+
+    private synchronized boolean updateConfigsWithIncrementalCooperative(AtomicReference> connectorConfigUpdatesCopy,
+                                                                         AtomicReference> connectorTargetStateChangesCopy,
+                                                                         AtomicReference> taskConfigUpdatesCopy) {
+        boolean retValue = false;
+        // This branch is here to avoid creating a snapshot if not needed
+        if (needsReconfigRebalance
+                || !connectorConfigUpdates.isEmpty()
+                || !connectorTargetStateChanges.isEmpty()
+                || !taskConfigUpdates.isEmpty()) {
+            // Connector reconfigs only need local updates since there is no coordination between workers required.
+            // However, if connectors were added or removed, work needs to be rebalanced since we have more work
+            // items to distribute among workers.
+            configState = configBackingStore.snapshot();
+
+            if (needsReconfigRebalance) {
+                log.debug("Requesting rebalance due to reconfiguration of tasks (needsReconfigRebalance: {})",
+                        needsReconfigRebalance);
+                member.requestRejoin();
+                needsReconfigRebalance = false;
+                retValue = true;
+            }
+
+            if (!connectorConfigUpdates.isEmpty()) {
+                // We can't start/stop while locked since starting connectors can cause task updates that will
+                // require writing configs, which in turn make callbacks into this class from another thread that
+                // require acquiring a lock. This leads to deadlock. Instead, just copy the info we need and process
+                // the updates after unlocking.
+                connectorConfigUpdatesCopy.set(connectorConfigUpdates);
+                connectorConfigUpdates = new HashSet<>();
+            }
+
+            if (!connectorTargetStateChanges.isEmpty()) {
+                // Similarly for target state changes which can cause connectors to be restarted
+                connectorTargetStateChangesCopy.set(connectorTargetStateChanges);
+                connectorTargetStateChanges = new HashSet<>();
+            }
+
+            if (!taskConfigUpdates.isEmpty()) {
+                // Similarly for task config updates
+                taskConfigUpdatesCopy.set(taskConfigUpdates);
+                taskConfigUpdates = new HashSet<>();
+            }
+        }
+        return retValue;
+    }
+
     private void processConnectorConfigUpdates(Set connectorConfigUpdates) {
         // If we only have connector config updates, we can just bounce the updated connectors that are
         // currently assigned to this worker.
@@ -361,10 +561,25 @@ private void processTargetStateChanges(Set connectorTargetStateChanges)
             // additionally, if the worker is running the connector itself, then we need to
             // request reconfiguration to ensure that config changes while paused take effect
             if (targetState == TargetState.STARTED)
-                reconfigureConnectorTasksWithRetry(connector);
+                reconfigureConnectorTasksWithRetry(time.milliseconds(), connector);
         }
     }
 
+    private void processTaskConfigUpdatesWithIncrementalCooperative(Set taskConfigUpdates) {
+        Set localTasks = assignment == null
+                                          ? Collections.emptySet()
+                                          : new HashSet<>(assignment.tasks());
+        Set connectorsWhoseTasksToStop = taskConfigUpdates.stream()
+                .map(ConnectorTaskId::connector).collect(Collectors.toSet());
+
+        List tasksToStop = localTasks.stream()
+                .filter(taskId -> connectorsWhoseTasksToStop.contains(taskId.connector()))
+                .collect(Collectors.toList());
+        log.info("Handling task config update by restarting tasks {}", tasksToStop);
+        worker.stopAndAwaitTasks(tasksToStop);
+        tasksToRestart.addAll(tasksToStop);
+    }
+
     // public for testing
     public void halt() {
         synchronized (this) {
@@ -451,10 +666,7 @@ public Void call() throws Exception {
                         if (!configState.contains(connName)) {
                             callback.onCompletion(new NotFoundException("Connector " + connName + " not found"), null);
                         } else {
-                            Map config = configState.rawConnectorConfig(connName);
-                            callback.onCompletion(null, new ConnectorInfo(connName, config,
-                                configState.tasks(connName),
-                                connectorTypeForClass(config.get(ConnectorConfig.CONNECTOR_CLASS_CONFIG))));
+                            callback.onCompletion(null, connectorInfo(connName));
                         }
                         return null;
                     }
@@ -574,7 +786,7 @@ public void requestTaskReconfiguration(final String connName) {
                 new Callable() {
                     @Override
                     public Void call() throws Exception {
-                        reconfigureConnectorTasksWithRetry(connName);
+                        reconfigureConnectorTasksWithRetry(time.milliseconds(), connName);
                         return null;
                     }
                 },
@@ -619,8 +831,36 @@ public Void call() throws Exception {
     }
 
     @Override
-    public void putTaskConfigs(final String connName, final List> configs, final Callback callback) {
+    public void putTaskConfigs(final String connName, final List> configs, final Callback callback, InternalRequestSignature requestSignature) {
         log.trace("Submitting put task configuration request {}", connName);
+        if (internalRequestValidationEnabled()) {
+            ConnectRestException requestValidationError = null;
+            if (requestSignature == null) {
+                requestValidationError = new BadRequestException("Internal request missing required signature");
+            } else if (!keySignatureVerificationAlgorithms.contains(requestSignature.keyAlgorithm())) {
+                requestValidationError = new BadRequestException(String.format(
+                    "This worker does not support the '%s' key signing algorithm used by other workers. " 
+                        + "This worker is currently configured to use: %s. " 
+                        + "Check that all workers' configuration files permit the same set of signature algorithms, " 
+                        + "and correct any misconfigured worker and restart it.",
+                    requestSignature.keyAlgorithm(),
+                    keySignatureVerificationAlgorithms
+                ));
+            } else {
+                synchronized (this) {
+                    if (!requestSignature.isValid(sessionKey)) {
+                        requestValidationError = new ConnectRestException(
+                            Response.Status.FORBIDDEN,
+                            "Internal request contained invalid signature."
+                        );
+                    }
+                }
+            }
+            if (requestValidationError != null) {
+                callback.onCompletion(requestValidationError, null);
+                return;
+            }
+        }
 
         addRequest(
                 new Callable() {
@@ -738,13 +978,15 @@ private String leaderUrl() {
 
     /**
      * Handle post-assignment operations, either trying to resolve issues that kept assignment from completing, getting
-     * this node into sync and its work started. Since
+     * this node into sync and its work started.
      *
      * @return false if we couldn't finish
      */
     private boolean handleRebalanceCompleted() {
-        if (this.rebalanceResolved)
+        if (rebalanceResolved) {
+            log.trace("Returning early because rebalance is marked as resolved (rebalanceResolved: true)");
             return true;
+        }
 
         // We need to handle a variety of cases after a rebalance:
         // 1. Assignment failed
@@ -777,6 +1019,15 @@ private boolean handleRebalanceCompleted() {
             }
         }
 
+        long now = time.milliseconds();
+        if (scheduledRebalance <= now) {
+            log.debug("Requesting rebalance because scheduled rebalance timeout has been reached "
+                    + "(now: {} scheduledRebalance: {}", scheduledRebalance, now);
+
+            needsRejoin = true;
+            scheduledRebalance = Long.MAX_VALUE;
+        }
+
         if (needsReadToEnd) {
             // Force exiting this method to avoid creating any connectors/tasks and require immediate rejoining if
             // we timed out. This should only happen if we failed to read configuration for long enough,
@@ -808,6 +1059,13 @@ private boolean handleRebalanceCompleted() {
         // guarantees we'll attempt to rejoin before executing this method again.
         herderMetrics.rebalanceSucceeded(time.milliseconds());
         rebalanceResolved = true;
+
+        if (!assignment.revokedConnectors().isEmpty() || !assignment.revokedTasks().isEmpty()) {
+            assignment.revokedConnectors().clear();
+            assignment.revokedTasks().clear();
+            member.requestRejoin();
+            return false;
+        }
         return true;
     }
 
@@ -822,19 +1080,41 @@ private boolean readConfigToEnd(long timeoutMs) {
             configBackingStore.refresh(timeoutMs, TimeUnit.MILLISECONDS);
             configState = configBackingStore.snapshot();
             log.info("Finished reading to end of log and updated config snapshot, new config log offset: {}", configState.offset());
+            backoffRetries = BACKOFF_RETRIES;
             return true;
         } catch (TimeoutException e) {
             // in case reading the log takes too long, leave the group to ensure a quick rebalance (although by default we should be out of the group already)
             // and back off to avoid a tight loop of rejoin-attempt-to-catch-up-leave
             log.warn("Didn't reach end of config log quickly enough", e);
-            member.maybeLeaveGroup();
+            member.maybeLeaveGroup("taking too long to read the log");
             backoff(workerUnsyncBackoffMs);
             return false;
         }
     }
 
     private void backoff(long ms) {
-        Utils.sleep(ms);
+        if (ConnectProtocolCompatibility.fromProtocolVersion(currentProtocolVersion) == EAGER) {
+            time.sleep(ms);
+            return;
+        }
+
+        if (backoffRetries > 0) {
+            int rebalanceDelayFraction =
+                    config.getInt(DistributedConfig.SCHEDULED_REBALANCE_MAX_DELAY_MS_CONFIG) / 10 / backoffRetries;
+            time.sleep(rebalanceDelayFraction);
+            --backoffRetries;
+            return;
+        }
+
+        ExtendedAssignment runningAssignmentSnapshot;
+        synchronized (this) {
+            runningAssignmentSnapshot = ExtendedAssignment.duplicate(runningAssignment);
+        }
+        log.info("Revoking current running assignment {} because after {} retries the worker "
+                + "has not caught up with the latest Connect cluster updates",
+                runningAssignmentSnapshot, BACKOFF_RETRIES);
+        member.revokeAssignment(runningAssignmentSnapshot);
+        backoffRetries = BACKOFF_RETRIES;
     }
 
     private void startAndStop(Collection> callables) {
@@ -847,19 +1127,53 @@ private void startAndStop(Collection> callables) {
 
     private void startWork() {
         // Start assigned connectors and tasks
-        log.info("Starting connectors and tasks using config offset {}", assignment.offset());
         List> callables = new ArrayList<>();
-        for (String connectorName : assignment.connectors()) {
-            callables.add(getConnectorStartingCallable(connectorName));
-        }
 
-        for (ConnectorTaskId taskId : assignment.tasks()) {
-            callables.add(getTaskStartingCallable(taskId));
+        // The sets in runningAssignment may change when onRevoked is called voluntarily by this
+        // herder (e.g. when a broker coordinator failure is detected). Otherwise the
+        // runningAssignment is always replaced by the assignment here.
+        synchronized (this) {
+            log.info("Starting connectors and tasks using config offset {}", assignment.offset());
+            log.debug("Received assignment: {}", assignment);
+            log.debug("Currently running assignment: {}", runningAssignment);
+
+            for (String connectorName : assignmentDifference(assignment.connectors(), runningAssignment.connectors())) {
+                callables.add(getConnectorStartingCallable(connectorName));
+            }
+
+            // These tasks have been stopped by this worker due to task reconfiguration. In order to
+            // restart them, they are removed just before the overall task startup from the set of
+            // currently running tasks. Therefore, they'll be restarted only if they are included in
+            // the assignment that was just received after rebalancing.
+            log.debug("Tasks to restart from currently running assignment: {}", tasksToRestart);
+            runningAssignment.tasks().removeAll(tasksToRestart);
+            tasksToRestart.clear();
+            for (ConnectorTaskId taskId : assignmentDifference(assignment.tasks(), runningAssignment.tasks())) {
+                callables.add(getTaskStartingCallable(taskId));
+            }
         }
+
         startAndStop(callables);
+
+        synchronized (this) {
+            runningAssignment = member.currentProtocolVersion() == CONNECT_PROTOCOL_V0
+                                ? ExtendedAssignment.empty()
+                                : assignment;
+        }
+
         log.info("Finished starting connectors and tasks");
     }
 
+    // arguments should assignment collections (connectors or tasks) and should not be null
+    private static  Collection assignmentDifference(Collection update, Collection running) {
+        if (running.isEmpty()) {
+            return update;
+        }
+        HashSet diff = new HashSet<>(update);
+        diff.removeAll(running);
+        return diff;
+    }
+
     private boolean startTask(ConnectorTaskId taskId) {
         log.info("Starting task {}", taskId);
         return worker.startTask(
@@ -911,7 +1225,7 @@ private boolean startConnector(String connectorName) {
         // task configs if they are actually different from the existing ones to avoid unnecessary updates when this is
         // just restoring an existing connector.
         if (started && initialState == TargetState.STARTED)
-            reconfigureConnectorTasksWithRetry(connectorName);
+            reconfigureConnectorTasksWithRetry(time.milliseconds(), connectorName);
 
         return started;
     }
@@ -946,7 +1260,7 @@ public Void call() throws Exception {
         };
     }
 
-    private void reconfigureConnectorTasksWithRetry(final String connName) {
+    private void reconfigureConnectorTasksWithRetry(long initialRequestTime, final String connName) {
         reconfigureConnector(connName, new Callback() {
             @Override
             public void onCompletion(Throwable error, Void result) {
@@ -955,19 +1269,25 @@ public void onCompletion(Throwable error, Void result) {
                 // never makes progress. The retry has to run through a DistributedHerderRequest since this callback could be happening
                 // from the HTTP request forwarding thread.
                 if (error != null) {
-                    log.error("Failed to reconfigure connector's tasks, retrying after backoff:", error);
+                    if (isPossibleExpiredKeyException(initialRequestTime, error)) {
+                        log.debug("Failed to reconfigure connector's tasks, possibly due to expired session key. Retrying after backoff");
+                    } else {
+                        log.error("Failed to reconfigure connector's tasks, retrying after backoff:", error);
+                    }
                     addRequest(RECONFIGURE_CONNECTOR_TASKS_BACKOFF_MS,
                             new Callable() {
                                 @Override
                                 public Void call() throws Exception {
-                                    reconfigureConnectorTasksWithRetry(connName);
+                                    reconfigureConnectorTasksWithRetry(initialRequestTime, connName);
                                     return null;
                                 }
                             }, new Callback() {
                                 @Override
                                 public void onCompletion(Throwable error, Void result) {
-                                    log.error("Unexpected error during connector task reconfiguration: ", error);
-                                    log.error("Task reconfiguration for {} failed unexpectedly, this connector will not be properly reconfigured unless manually triggered.", connName);
+                                    if (error != null) {
+                                        log.error("Unexpected error during connector task reconfiguration: ", error);
+                                        log.error("Task reconfiguration for {} failed unexpectedly, this connector will not be properly reconfigured unless manually triggered.", connName);
+                                    }
                                 }
                             }
                     );
@@ -976,6 +1296,15 @@ public void onCompletion(Throwable error, Void result) {
         });
     }
 
+    boolean isPossibleExpiredKeyException(long initialRequestTime, Throwable error) {
+        if (error instanceof ConnectRestException) {
+            ConnectRestException connectError = (ConnectRestException) error;
+            return connectError.statusCode() == Response.Status.FORBIDDEN.getStatusCode()
+                && initialRequestTime + TimeUnit.MINUTES.toMillis(1) >= time.milliseconds();
+        }
+        return false;
+    }
+
     // Updates configurations for a connector by requesting them from the connector, filling in parameters provided
     // by the system, then checks whether any configs have actually changed before submitting the new configs to storage
     private void reconfigureConnector(final String connName, final Callback cb) {
@@ -1031,7 +1360,8 @@ public void run() {
                                     return;
                                 }
                                 String reconfigUrl = RestServer.urlJoin(leaderUrl, "/connectors/" + connName + "/tasks");
-                                RestClient.httpRequest(reconfigUrl, "POST", rawTaskProps, null, config);
+                                log.trace("Forwarding task configurations for connector {} to leader", connName);
+                                RestClient.httpRequest(reconfigUrl, "POST", null, rawTaskProps, null, config, sessionKey, requestSignatureAlgorithm);
                                 cb.onCompletion(null, null);
                             } catch (ConnectException e) {
                                 log.error("Request to leader to reconfigure connector tasks failed", e);
@@ -1068,6 +1398,14 @@ DistributedHerderRequest addRequest(long delayMs, Callable action, Callbac
         return req;
     }
 
+    private boolean internalRequestValidationEnabled() {
+        return internalRequestValidationEnabled(member.currentProtocolVersion());
+    }
+
+    private static boolean internalRequestValidationEnabled(short protocolVersion) {
+        return protocolVersion >= CONNECT_PROTOCOL_V2;
+    }
+
     private DistributedHerderRequest peekWithoutException() {
         try {
             return requests.isEmpty() ? null : requests.first();
@@ -1111,12 +1449,17 @@ public void onConnectorConfigUpdate(String connector) {
         public void onTaskConfigUpdate(Collection tasks) {
             log.info("Tasks {} configs updated", tasks);
 
-            // Stage the update and wake up the work thread. No need to record the set of tasks here because task reconfigs
-            // always need a rebalance to ensure offsets get committed.
+            // Stage the update and wake up the work thread.
+            // The set of tasks is recorder for incremental cooperative rebalancing, in which
+            // tasks don't get restarted unless they are balanced between workers.
+            // With eager rebalancing there's no need to record the set of tasks because task reconfigs
+            // always need a rebalance to ensure offsets get committed. In eager rebalancing the
+            // recorded set of tasks remains unused.
             // TODO: As an optimization, some task config updates could avoid a rebalance. In particular, single-task
             // connectors clearly don't need any coordination.
             synchronized (DistributedHerder.this) {
                 needsReconfigRebalance = true;
+                taskConfigUpdates.addAll(tasks);
             }
             member.wakeup();
         }
@@ -1130,6 +1473,21 @@ public void onConnectorTargetStateChange(String connector) {
             }
             member.wakeup();
         }
+
+        @Override
+        public void onSessionKeyUpdate(SessionKey sessionKey) {
+            log.info("Session key updated");
+
+            synchronized (DistributedHerder.this) {
+                DistributedHerder.this.sessionKey = sessionKey.key();
+                // Track the expiration of the key if and only if this worker is the leader
+                // Followers will receive rotated keys from the follower and won't be responsible for
+                // tracking expiration and distributing new keys themselves
+                if (isLeader() && keyRotationIntervalMs > 0) {
+                    DistributedHerder.this.keyExpiration = sessionKey.creationTimestamp() + keyRotationIntervalMs;
+                }
+            }
+        }
     }
 
     class DistributedHerderRequest implements HerderRequest, Comparable {
@@ -1200,32 +1558,87 @@ private void updateDeletedConnectorStatus() {
         }
     }
 
+    private void updateDeletedTaskStatus() {
+        ClusterConfigState snapshot = configBackingStore.snapshot();
+        for (String connector : statusBackingStore.connectors()) {
+            Set remainingTasks = new HashSet<>(snapshot.tasks(connector));
+            
+            statusBackingStore.getAll(connector).stream()
+                .map(TaskStatus::id)
+                .filter(task -> !remainingTasks.contains(task))
+                .forEach(this::onDeletion);
+        }
+    }
+
     protected HerderMetrics herderMetrics() {
         return herderMetrics;
     }
 
     // Rebalances are triggered internally from the group member, so these are always executed in the work thread.
     public class RebalanceListener implements WorkerRebalanceListener {
+        private final Time time;
+        RebalanceListener(Time time) {
+            this.time = time;
+        }
+
         @Override
-        public void onAssigned(ConnectProtocol.Assignment assignment, int generation) {
+        public void onAssigned(ExtendedAssignment assignment, int generation) {
             // This callback just logs the info and saves it. The actual response is handled in the main loop, which
             // ensures the group member's logic for rebalancing can complete, potentially long-running steps to
             // catch up (or backoff if we fail) not executed in a callback, and so we'll be able to invoke other
             // group membership actions (e.g., we may need to explicitly leave the group if we cannot handle the
             // assigned tasks).
-            log.info("Joined group and got assignment: {}", assignment);
+            short priorProtocolVersion = currentProtocolVersion;
+            DistributedHerder.this.currentProtocolVersion = member.currentProtocolVersion();
+            log.info(
+                "Joined group at generation {} with protocol version {} and got assignment: {} with rebalance delay: {}",
+                generation,
+                DistributedHerder.this.currentProtocolVersion,
+                assignment,
+                assignment.delay()
+            );
             synchronized (DistributedHerder.this) {
                 DistributedHerder.this.assignment = assignment;
                 DistributedHerder.this.generation = generation;
+                int delay = assignment.delay();
+                DistributedHerder.this.scheduledRebalance = delay > 0
+                    ? time.milliseconds() + delay
+                    : Long.MAX_VALUE;
+
+                boolean requestValidationWasEnabled = internalRequestValidationEnabled(priorProtocolVersion);
+                boolean requestValidationNowEnabled = internalRequestValidationEnabled(currentProtocolVersion);
+                if (requestValidationNowEnabled != requestValidationWasEnabled) {
+                    // Internal request verification has been switched on or off; let the user know
+                    if (requestValidationNowEnabled) {
+                        log.info("Internal request validation has been re-enabled");
+                    } else {
+                        log.warn(
+                            "The protocol used by this Connect cluster has been downgraded from '{}' to '{}' and internal request "
+                                + "validation is now disabled. This is most likely caused by a new worker joining the cluster with an "
+                                + "older protocol specified for the {} configuration; if this is not intentional, either remove the {} "
+                                + "configuration from that worker's config file, or change its value to '{}'. If this configuration is "
+                                + "left as-is, the cluster will be insecure; for more information, see KIP-507: "
+                                + "https://cwiki.apache.org/confluence/display/KAFKA/KIP-507%3A+Securing+Internal+Connect+REST+Endpoints",
+                            ConnectProtocolCompatibility.fromProtocolVersion(priorProtocolVersion),
+                            ConnectProtocolCompatibility.fromProtocolVersion(DistributedHerder.this.currentProtocolVersion),
+                            DistributedConfig.CONNECT_PROTOCOL_CONFIG,
+                            DistributedConfig.CONNECT_PROTOCOL_CONFIG,
+                            ConnectProtocolCompatibility.SESSIONED.name()
+                        );
+                    }
+                }
+
                 rebalanceResolved = false;
                 herderMetrics.rebalanceStarted(time.milliseconds());
             }
 
-            // Delete the statuses of all connectors removed prior to the start of this rebalance. This has to
-            // be done after the rebalance completes to avoid race conditions as the previous generation attempts
-            // to change the state to UNASSIGNED after tasks have been stopped.
-            if (isLeader())
+            // Delete the statuses of all connectors and tasks removed prior to the start of this rebalance. This
+            // has to be done after the rebalance completes to avoid race conditions as the previous generation
+            // attempts to change the state to UNASSIGNED after tasks have been stopped.
+            if (isLeader()) {
                 updateDeletedConnectorStatus();
+                updateDeletedTaskStatus();
+            }
 
             // We *must* interrupt any poll() call since this could occur when the poll starts, and we might then
             // sleep in the poll() for a long time. Forcing a wakeup ensures we'll get to process this event in the
@@ -1235,16 +1648,10 @@ public void onAssigned(ConnectProtocol.Assignment assignment, int generation) {
 
         @Override
         public void onRevoked(String leader, Collection connectors, Collection tasks) {
-            log.info("Rebalance started");
-
             // Note that since we don't reset the assignment, we don't revoke leadership here. During a rebalance,
             // it is still important to have a leader that can write configs, offsets, etc.
 
             if (rebalanceResolved) {
-                // TODO: Technically we don't have to stop connectors at all until we know they've really been removed from
-                // this worker. Instead, we can let them continue to run but buffer any update requests (which should be
-                // rare anyway). This would avoid a steady stream of start/stop, which probably also includes lots of
-                // unnecessary repeated connections to the source/sink system.
                 List> callables = new ArrayList<>();
                 for (final String connectorName : connectors) {
                     callables.add(getConnectorStoppingCallable(connectorName));
@@ -1260,12 +1667,20 @@ public void onRevoked(String leader, Collection connectors, Collection() {
+                @Override
+                public String metricValue(long now) {
+                    return ConnectProtocolCompatibility.fromProtocolVersion(member.currentProtocolVersion()).name();
+                }
+            });
             metricGroup.addValueMetric(registry.leaderName, new LiteralSupplier() {
                 @Override
                 public String metricValue(long now) {
@@ -1304,7 +1725,7 @@ public Double metricValue(long now) {
             });
 
             rebalanceCompletedCounts = metricGroup.sensor("completed-rebalance-count");
-            rebalanceCompletedCounts.add(metricGroup.metricName(registry.rebalanceCompletedTotal), new Total());
+            rebalanceCompletedCounts.add(metricGroup.metricName(registry.rebalanceCompletedTotal), new CumulativeSum());
 
             rebalanceTime = metricGroup.sensor("rebalance-time");
             rebalanceTime.add(metricGroup.metricName(registry.rebalanceTimeMax), new Max());
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/EagerAssignor.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/EagerAssignor.java
new file mode 100644
index 0000000000000..b6dbd09147805
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/EagerAssignor.java
@@ -0,0 +1,182 @@
+/*
+ * 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.distributed;
+
+import org.apache.kafka.common.utils.CircularIterator;
+import org.apache.kafka.common.utils.LogContext;
+import org.apache.kafka.connect.util.ConnectorTaskId;
+import org.slf4j.Logger;
+
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import static org.apache.kafka.common.message.JoinGroupResponseData.JoinGroupResponseMember;
+import static org.apache.kafka.connect.runtime.distributed.ConnectProtocol.Assignment;
+import static org.apache.kafka.connect.runtime.distributed.WorkerCoordinator.LeaderState;
+
+
+/**
+ * An assignor that computes a unweighted round-robin distribution of connectors and tasks. The
+ * connectors are assigned to the workers first, followed by the tasks. This is to avoid
+ * load imbalance when several 1-task connectors are running, given that a connector is usually
+ * more lightweight than a task.
+ *
+ * Note that this class is NOT thread-safe.
+ */
+public class EagerAssignor implements ConnectAssignor {
+    private final Logger log;
+
+    public EagerAssignor(LogContext logContext) {
+        this.log = logContext.logger(EagerAssignor.class);
+    }
+
+    @Override
+    public Map performAssignment(String leaderId, String protocol,
+                                                     List allMemberMetadata,
+                                                     WorkerCoordinator coordinator) {
+        log.debug("Performing task assignment");
+        Map memberConfigs = new HashMap<>();
+        for (JoinGroupResponseMember member : allMemberMetadata)
+            memberConfigs.put(member.memberId(), IncrementalCooperativeConnectProtocol.deserializeMetadata(ByteBuffer.wrap(member.metadata())));
+
+        long maxOffset = findMaxMemberConfigOffset(memberConfigs, coordinator);
+        Long leaderOffset = ensureLeaderConfig(maxOffset, coordinator);
+        if (leaderOffset == null)
+            return fillAssignmentsAndSerialize(memberConfigs.keySet(), Assignment.CONFIG_MISMATCH,
+                    leaderId, memberConfigs.get(leaderId).url(), maxOffset,
+                    new HashMap<>(), new HashMap<>());
+        return performTaskAssignment(leaderId, leaderOffset, memberConfigs, coordinator);
+    }
+
+    private Long ensureLeaderConfig(long maxOffset, WorkerCoordinator coordinator) {
+        // If this leader is behind some other members, we can't do assignment
+        if (coordinator.configSnapshot().offset() < maxOffset) {
+            // We might be able to take a new snapshot to catch up immediately and avoid another round of syncing here.
+            // Alternatively, if this node has already passed the maximum reported by any other member of the group, it
+            // is also safe to use this newer state.
+            ClusterConfigState updatedSnapshot = coordinator.configFreshSnapshot();
+            if (updatedSnapshot.offset() < maxOffset) {
+                log.info("Was selected to perform assignments, but do not have latest config found in sync request. " +
+                        "Returning an empty configuration to trigger re-sync.");
+                return null;
+            } else {
+                coordinator.configSnapshot(updatedSnapshot);
+                return updatedSnapshot.offset();
+            }
+        }
+        return maxOffset;
+    }
+
+    private Map performTaskAssignment(String leaderId, long maxOffset,
+                                                          Map memberConfigs,
+                                                          WorkerCoordinator coordinator) {
+        Map> connectorAssignments = new HashMap<>();
+        Map> taskAssignments = new HashMap<>();
+
+        // Perform round-robin task assignment. Assign all connectors and then all tasks because assigning both the
+        // connector and its tasks can lead to very uneven distribution of work in some common cases (e.g. for connectors
+        // that generate only 1 task each; in a cluster of 2 or an even # of nodes, only even nodes will be assigned
+        // connectors and only odd nodes will be assigned tasks, but tasks are, on average, actually more resource
+        // intensive than connectors).
+        List connectorsSorted = sorted(coordinator.configSnapshot().connectors());
+        CircularIterator memberIt = new CircularIterator<>(sorted(memberConfigs.keySet()));
+        for (String connectorId : connectorsSorted) {
+            String connectorAssignedTo = memberIt.next();
+            log.trace("Assigning connector {} to {}", connectorId, connectorAssignedTo);
+            Collection memberConnectors = connectorAssignments.get(connectorAssignedTo);
+            if (memberConnectors == null) {
+                memberConnectors = new ArrayList<>();
+                connectorAssignments.put(connectorAssignedTo, memberConnectors);
+            }
+            memberConnectors.add(connectorId);
+        }
+        for (String connectorId : connectorsSorted) {
+            for (ConnectorTaskId taskId : sorted(coordinator.configSnapshot().tasks(connectorId))) {
+                String taskAssignedTo = memberIt.next();
+                log.trace("Assigning task {} to {}", taskId, taskAssignedTo);
+                Collection memberTasks = taskAssignments.get(taskAssignedTo);
+                if (memberTasks == null) {
+                    memberTasks = new ArrayList<>();
+                    taskAssignments.put(taskAssignedTo, memberTasks);
+                }
+                memberTasks.add(taskId);
+            }
+        }
+
+        coordinator.leaderState(new LeaderState(memberConfigs, connectorAssignments, taskAssignments));
+
+        return fillAssignmentsAndSerialize(memberConfigs.keySet(), Assignment.NO_ERROR,
+                leaderId, memberConfigs.get(leaderId).url(), maxOffset, connectorAssignments, taskAssignments);
+    }
+
+    private Map fillAssignmentsAndSerialize(Collection members,
+                                                                short error,
+                                                                String leaderId,
+                                                                String leaderUrl,
+                                                                long maxOffset,
+                                                                Map> connectorAssignments,
+                                                                Map> taskAssignments) {
+
+        Map groupAssignment = new HashMap<>();
+        for (String member : members) {
+            Collection connectors = connectorAssignments.get(member);
+            if (connectors == null) {
+                connectors = Collections.emptyList();
+            }
+            Collection tasks = taskAssignments.get(member);
+            if (tasks == null) {
+                tasks = Collections.emptyList();
+            }
+            Assignment assignment = new Assignment(error, leaderId, leaderUrl, maxOffset, connectors, tasks);
+            log.debug("Assignment: {} -> {}", member, assignment);
+            groupAssignment.put(member, ConnectProtocol.serializeAssignment(assignment));
+        }
+        log.debug("Finished assignment");
+        return groupAssignment;
+    }
+
+    private long findMaxMemberConfigOffset(Map memberConfigs,
+                                           WorkerCoordinator coordinator) {
+        // The new config offset is the maximum seen by any member. We always perform assignment using this offset,
+        // even if some members have fallen behind. The config offset used to generate the assignment is included in
+        // the response so members that have fallen behind will not use the assignment until they have caught up.
+        Long maxOffset = null;
+        for (Map.Entry stateEntry : memberConfigs.entrySet()) {
+            long memberRootOffset = stateEntry.getValue().offset();
+            if (maxOffset == null)
+                maxOffset = memberRootOffset;
+            else
+                maxOffset = Math.max(maxOffset, memberRootOffset);
+        }
+
+        log.debug("Max config offset root: {}, local snapshot config offsets root: {}",
+                  maxOffset, coordinator.configSnapshot().offset());
+        return maxOffset;
+    }
+
+    private static > List sorted(Collection members) {
+        List res = new ArrayList<>(members);
+        Collections.sort(res);
+        return res;
+    }
+
+}
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/ExtendedAssignment.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/ExtendedAssignment.java
new file mode 100644
index 0000000000000..e544407ca0912
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/ExtendedAssignment.java
@@ -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 org.apache.kafka.connect.runtime.distributed;
+
+import org.apache.kafka.common.protocol.types.Struct;
+import org.apache.kafka.connect.util.ConnectorTaskId;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.LinkedHashMap;
+import java.util.LinkedHashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.stream.Collectors;
+
+import static org.apache.kafka.connect.runtime.distributed.ConnectProtocol.ASSIGNMENT_KEY_NAME;
+import static org.apache.kafka.connect.runtime.distributed.ConnectProtocol.CONFIG_OFFSET_KEY_NAME;
+import static org.apache.kafka.connect.runtime.distributed.ConnectProtocol.CONNECTOR_KEY_NAME;
+import static org.apache.kafka.connect.runtime.distributed.ConnectProtocol.CONNECTOR_TASK;
+import static org.apache.kafka.connect.runtime.distributed.ConnectProtocol.ERROR_KEY_NAME;
+import static org.apache.kafka.connect.runtime.distributed.ConnectProtocol.LEADER_KEY_NAME;
+import static org.apache.kafka.connect.runtime.distributed.ConnectProtocol.LEADER_URL_KEY_NAME;
+import static org.apache.kafka.connect.runtime.distributed.ConnectProtocol.TASKS_KEY_NAME;
+import static org.apache.kafka.connect.runtime.distributed.IncrementalCooperativeConnectProtocol.ASSIGNMENT_V1;
+import static org.apache.kafka.connect.runtime.distributed.IncrementalCooperativeConnectProtocol.CONNECTOR_ASSIGNMENT_V1;
+import static org.apache.kafka.connect.runtime.distributed.IncrementalCooperativeConnectProtocol.CONNECT_PROTOCOL_V1;
+import static org.apache.kafka.connect.runtime.distributed.IncrementalCooperativeConnectProtocol.REVOKED_KEY_NAME;
+import static org.apache.kafka.connect.runtime.distributed.IncrementalCooperativeConnectProtocol.SCHEDULED_DELAY_KEY_NAME;
+
+/**
+ * The extended assignment of connectors and tasks that includes revoked connectors and tasks
+ * as well as a scheduled rebalancing delay.
+ */
+public class ExtendedAssignment extends ConnectProtocol.Assignment {
+    private final short version;
+    private final Collection revokedConnectorIds;
+    private final Collection revokedTaskIds;
+    private final int delay;
+
+    private static final ExtendedAssignment EMPTY = new ExtendedAssignment(
+            CONNECT_PROTOCOL_V1, ConnectProtocol.Assignment.NO_ERROR, null, null, -1,
+            Collections.emptyList(), Collections.emptyList(), Collections.emptyList(), Collections.emptyList(), 0);
+
+    /**
+     * Create an assignment indicating responsibility for the given connector instances and task Ids.
+     *
+     * @param version Connect protocol version
+     * @param error error code for this assignment; {@code ConnectProtocol.Assignment.NO_ERROR}
+     *              indicates no error during assignment
+     * @param leader Connect group's leader Id; may be null only on the empty assignment
+     * @param leaderUrl Connect group's leader URL; may be null only on the empty assignment
+     * @param configOffset the offset in the config topic that this assignment is corresponding to
+     * @param connectorIds list of connectors that the worker should instantiate and run; may not be null
+     * @param taskIds list of task IDs that the worker should instantiate and run; may not be null
+     * @param revokedConnectorIds list of connectors that the worker should stop running; may not be null
+     * @param revokedTaskIds list of task IDs that the worker should stop running; may not be null
+     * @param delay the scheduled delay after which the worker should rejoin the group
+     */
+    public ExtendedAssignment(short version, short error, String leader, String leaderUrl, long configOffset,
+                             Collection connectorIds, Collection taskIds,
+                             Collection revokedConnectorIds, Collection revokedTaskIds,
+                             int delay) {
+        super(error, leader, leaderUrl, configOffset, connectorIds, taskIds);
+        this.version = version;
+        this.revokedConnectorIds = Objects.requireNonNull(revokedConnectorIds,
+                "Revoked connector IDs may be empty but not null");
+        this.revokedTaskIds = Objects.requireNonNull(revokedTaskIds,
+                "Revoked task IDs may be empty but not null");
+        this.delay = delay;
+    }
+
+    public static ExtendedAssignment duplicate(ExtendedAssignment assignment) {
+        return new ExtendedAssignment(
+                assignment.version(),
+                assignment.error(),
+                assignment.leader(),
+                assignment.leaderUrl(),
+                assignment.offset(),
+                new LinkedHashSet<>(assignment.connectors()),
+                new LinkedHashSet<>(assignment.tasks()),
+                new LinkedHashSet<>(assignment.revokedConnectors()),
+                new LinkedHashSet<>(assignment.revokedTasks()),
+                assignment.delay());
+    }
+
+    /**
+     * Return the version of the connect protocol that this assignment belongs to.
+     *
+     * @return the connect protocol version of this assignment
+     */
+    public short version() {
+        return version;
+    }
+
+    /**
+     * Return the IDs of the connectors that are revoked by this assignment.
+     *
+     * @return the revoked connector IDs; empty if there are no revoked connectors
+     */
+    public Collection revokedConnectors() {
+        return revokedConnectorIds;
+    }
+
+    /**
+     * Return the IDs of the tasks that are revoked by this assignment.
+     *
+     * @return the revoked task IDs; empty if there are no revoked tasks
+     */
+    public Collection revokedTasks() {
+        return revokedTaskIds;
+    }
+
+    /**
+     * Return the delay for the rebalance that is scheduled by this assignment.
+     *
+     * @return the scheduled delay
+     */
+    public int delay() {
+        return delay;
+    }
+
+    /**
+     * Return an empty assignment.
+     *
+     * @return an empty assignment
+     */
+    public static ExtendedAssignment empty() {
+        return EMPTY;
+    }
+
+    @Override
+    public String toString() {
+        return "Assignment{" +
+                "error=" + error() +
+                ", leader='" + leader() + '\'' +
+                ", leaderUrl='" + leaderUrl() + '\'' +
+                ", offset=" + offset() +
+                ", connectorIds=" + connectors() +
+                ", taskIds=" + tasks() +
+                ", revokedConnectorIds=" + revokedConnectorIds +
+                ", revokedTaskIds=" + revokedTaskIds +
+                ", delay=" + delay +
+                '}';
+    }
+
+    private Map> revokedAsMap() {
+        if (revokedConnectorIds == null && revokedTaskIds == null) {
+            return null;
+        }
+        // Using LinkedHashMap preserves the ordering, which is helpful for tests and debugging
+        Map> taskMap = new LinkedHashMap<>();
+        Optional.ofNullable(revokedConnectorIds)
+                .orElseGet(Collections::emptyList)
+                .stream()
+                .distinct()
+                .forEachOrdered(connectorId -> {
+                    Collection connectorTasks =
+                            taskMap.computeIfAbsent(connectorId, v -> new ArrayList<>());
+                    connectorTasks.add(CONNECTOR_TASK);
+                });
+
+        Optional.ofNullable(revokedTaskIds)
+                .orElseGet(Collections::emptyList)
+                .forEach(taskId -> {
+                    String connectorId = taskId.connector();
+                    Collection connectorTasks =
+                            taskMap.computeIfAbsent(connectorId, v -> new ArrayList<>());
+                    connectorTasks.add(taskId.task());
+                });
+        return taskMap;
+    }
+
+    /**
+     * Return the {@code Struct} that corresponds to this assignment.
+     *
+     * @return the assignment struct
+     */
+    public Struct toStruct() {
+        Collection assigned = taskAssignments(asMap());
+        Collection revoked = taskAssignments(revokedAsMap());
+        return new Struct(ASSIGNMENT_V1)
+                .set(ERROR_KEY_NAME, error())
+                .set(LEADER_KEY_NAME, leader())
+                .set(LEADER_URL_KEY_NAME, leaderUrl())
+                .set(CONFIG_OFFSET_KEY_NAME, offset())
+                .set(ASSIGNMENT_KEY_NAME, assigned != null ? assigned.toArray() : null)
+                .set(REVOKED_KEY_NAME, revoked != null ? revoked.toArray() : null)
+                .set(SCHEDULED_DELAY_KEY_NAME, delay);
+    }
+
+    /**
+     * Given a {@code Struct} that encodes an assignment return the assignment object.
+     *
+     * @param struct a struct representing an assignment
+     * @return the assignment
+     */
+    public static ExtendedAssignment fromStruct(short version, Struct struct) {
+        return struct == null
+               ? null
+               : new ExtendedAssignment(
+                       version,
+                       struct.getShort(ERROR_KEY_NAME),
+                       struct.getString(LEADER_KEY_NAME),
+                       struct.getString(LEADER_URL_KEY_NAME),
+                       struct.getLong(CONFIG_OFFSET_KEY_NAME),
+                       extractConnectors(struct, ASSIGNMENT_KEY_NAME),
+                       extractTasks(struct, ASSIGNMENT_KEY_NAME),
+                       extractConnectors(struct, REVOKED_KEY_NAME),
+                       extractTasks(struct, REVOKED_KEY_NAME),
+                       struct.getInt(SCHEDULED_DELAY_KEY_NAME));
+    }
+
+    private static Collection taskAssignments(Map> assignments) {
+        return assignments == null
+               ? null
+               : assignments.entrySet().stream()
+                       .map(connectorEntry -> {
+                           Struct taskAssignment = new Struct(CONNECTOR_ASSIGNMENT_V1);
+                           taskAssignment.set(CONNECTOR_KEY_NAME, connectorEntry.getKey());
+                           taskAssignment.set(TASKS_KEY_NAME, connectorEntry.getValue().toArray());
+                           return taskAssignment;
+                       }).collect(Collectors.toList());
+    }
+
+    private static Collection extractConnectors(Struct struct, String key) {
+        assert REVOKED_KEY_NAME.equals(key) || ASSIGNMENT_KEY_NAME.equals(key);
+
+        Object[] connectors = struct.getArray(key);
+        if (connectors == null) {
+            return Collections.emptyList();
+        }
+        List connectorIds = new ArrayList<>();
+        for (Object structObj : connectors) {
+            Struct assignment = (Struct) structObj;
+            String connector = assignment.getString(CONNECTOR_KEY_NAME);
+            for (Object taskIdObj : assignment.getArray(TASKS_KEY_NAME)) {
+                Integer taskId = (Integer) taskIdObj;
+                if (taskId == CONNECTOR_TASK) {
+                    connectorIds.add(connector);
+                }
+            }
+        }
+        return connectorIds;
+    }
+
+    private static Collection extractTasks(Struct struct, String key) {
+        assert REVOKED_KEY_NAME.equals(key) || ASSIGNMENT_KEY_NAME.equals(key);
+
+        Object[] tasks = struct.getArray(key);
+        if (tasks == null) {
+            return Collections.emptyList();
+        }
+        List tasksIds = new ArrayList<>();
+        for (Object structObj : tasks) {
+            Struct assignment = (Struct) structObj;
+            String connector = assignment.getString(CONNECTOR_KEY_NAME);
+            for (Object taskIdObj : assignment.getArray(TASKS_KEY_NAME)) {
+                Integer taskId = (Integer) taskIdObj;
+                if (taskId != CONNECTOR_TASK) {
+                    tasksIds.add(new ConnectorTaskId(connector, taskId));
+                }
+            }
+        }
+        return tasksIds;
+    }
+
+}
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/ExtendedWorkerState.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/ExtendedWorkerState.java
new file mode 100644
index 0000000000000..663979bdc48ca
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/ExtendedWorkerState.java
@@ -0,0 +1,48 @@
+/*
+ * 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.distributed;
+
+/**
+ * A class that captures the deserialized form of a worker's metadata.
+ */
+public class ExtendedWorkerState extends ConnectProtocol.WorkerState {
+    private final ExtendedAssignment assignment;
+
+    public ExtendedWorkerState(String url, long offset, ExtendedAssignment assignment) {
+        super(url, offset);
+        this.assignment = assignment != null ? assignment : ExtendedAssignment.empty();
+    }
+
+    /**
+     * This method returns which was the assignment of connectors and tasks on a worker at the
+     * moment that its state was captured by this class.
+     *
+     * @return the assignment of connectors and tasks
+     */
+    public ExtendedAssignment assignment() {
+        return assignment;
+    }
+
+    @Override
+    public String toString() {
+        return "WorkerState{" +
+                "url='" + url() + '\'' +
+                ", offset=" + offset() +
+                ", " + assignment +
+                '}';
+    }
+}
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/IncrementalCooperativeAssignor.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/IncrementalCooperativeAssignor.java
new file mode 100644
index 0000000000000..744a099730097
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/IncrementalCooperativeAssignor.java
@@ -0,0 +1,749 @@
+/*
+ * 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.distributed;
+
+import org.apache.kafka.common.utils.LogContext;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.connect.runtime.distributed.WorkerCoordinator.ConnectorsAndTasks;
+import org.apache.kafka.connect.runtime.distributed.WorkerCoordinator.WorkerLoad;
+import org.apache.kafka.connect.util.ConnectorTaskId;
+import org.slf4j.Logger;
+
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.LinkedHashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import static org.apache.kafka.common.message.JoinGroupResponseData.JoinGroupResponseMember;
+import static org.apache.kafka.connect.runtime.distributed.ConnectProtocol.Assignment;
+import static org.apache.kafka.connect.runtime.distributed.IncrementalCooperativeConnectProtocol.CONNECT_PROTOCOL_V1;
+import static org.apache.kafka.connect.runtime.distributed.IncrementalCooperativeConnectProtocol.CONNECT_PROTOCOL_V2;
+import static org.apache.kafka.connect.runtime.distributed.WorkerCoordinator.LeaderState;
+
+/**
+ * An assignor that computes a distribution of connectors and tasks according to the incremental
+ * cooperative strategy for rebalancing. {@see
+ * https://cwiki.apache.org/confluence/display/KAFKA/KIP-415%3A+Incremental+Cooperative
+ * +Rebalancing+in+Kafka+Connect} for a description of the assignment policy.
+ *
+ * Note that this class is NOT thread-safe.
+ */
+public class IncrementalCooperativeAssignor implements ConnectAssignor {
+    private final Logger log;
+    private final Time time;
+    private final int maxDelay;
+    private ConnectorsAndTasks previousAssignment;
+    private ConnectorsAndTasks previousRevocation;
+    private boolean canRevoke;
+    // visible for testing
+    protected final Set candidateWorkersForReassignment;
+    protected long scheduledRebalance;
+    protected int delay;
+    protected int previousGenerationId;
+    protected Set previousMembers;
+
+    public IncrementalCooperativeAssignor(LogContext logContext, Time time, int maxDelay) {
+        this.log = logContext.logger(IncrementalCooperativeAssignor.class);
+        this.time = time;
+        this.maxDelay = maxDelay;
+        this.previousAssignment = ConnectorsAndTasks.EMPTY;
+        this.previousRevocation = new ConnectorsAndTasks.Builder().build();
+        this.canRevoke = true;
+        this.scheduledRebalance = 0;
+        this.candidateWorkersForReassignment = new LinkedHashSet<>();
+        this.delay = 0;
+        this.previousGenerationId = -1;
+        this.previousMembers = Collections.emptySet();
+    }
+
+    @Override
+    public Map performAssignment(String leaderId, String protocol,
+                                                     List allMemberMetadata,
+                                                     WorkerCoordinator coordinator) {
+        log.debug("Performing task assignment");
+
+        Map memberConfigs = new HashMap<>();
+        for (JoinGroupResponseMember member : allMemberMetadata) {
+            memberConfigs.put(
+                    member.memberId(),
+                    IncrementalCooperativeConnectProtocol.deserializeMetadata(ByteBuffer.wrap(member.metadata())));
+        }
+        log.debug("Member configs: {}", memberConfigs);
+
+        // The new config offset is the maximum seen by any member. We always perform assignment using this offset,
+        // even if some members have fallen behind. The config offset used to generate the assignment is included in
+        // the response so members that have fallen behind will not use the assignment until they have caught up.
+        long maxOffset = memberConfigs.values().stream().map(ExtendedWorkerState::offset).max(Long::compare).get();
+        log.debug("Max config offset root: {}, local snapshot config offsets root: {}",
+                  maxOffset, coordinator.configSnapshot().offset());
+
+        short protocolVersion = memberConfigs.values().stream()
+            .allMatch(state -> state.assignment().version() == CONNECT_PROTOCOL_V2)
+                ? CONNECT_PROTOCOL_V2
+                : CONNECT_PROTOCOL_V1;
+
+        Long leaderOffset = ensureLeaderConfig(maxOffset, coordinator);
+        if (leaderOffset == null) {
+            Map assignments = fillAssignments(
+                    memberConfigs.keySet(), Assignment.CONFIG_MISMATCH,
+                    leaderId, memberConfigs.get(leaderId).url(), maxOffset, Collections.emptyMap(),
+                    Collections.emptyMap(), Collections.emptyMap(), 0, protocolVersion);
+            return serializeAssignments(assignments);
+        }
+        return performTaskAssignment(leaderId, leaderOffset, memberConfigs, coordinator, protocolVersion);
+    }
+
+    private Long ensureLeaderConfig(long maxOffset, WorkerCoordinator coordinator) {
+        // If this leader is behind some other members, we can't do assignment
+        if (coordinator.configSnapshot().offset() < maxOffset) {
+            // We might be able to take a new snapshot to catch up immediately and avoid another round of syncing here.
+            // Alternatively, if this node has already passed the maximum reported by any other member of the group, it
+            // is also safe to use this newer state.
+            ClusterConfigState updatedSnapshot = coordinator.configFreshSnapshot();
+            if (updatedSnapshot.offset() < maxOffset) {
+                log.info("Was selected to perform assignments, but do not have latest config found in sync request. "
+                         + "Returning an empty configuration to trigger re-sync.");
+                return null;
+            } else {
+                coordinator.configSnapshot(updatedSnapshot);
+                return updatedSnapshot.offset();
+            }
+        }
+        return maxOffset;
+    }
+
+    /**
+     * Performs task assignment based on the incremental cooperative connect protocol.
+     * Read more on the design and implementation in:
+     * {@see https://cwiki.apache.org/confluence/display/KAFKA/KIP-415%3A+Incremental+Cooperative+Rebalancing+in+Kafka+Connect}
+     *
+     * @param leaderId the ID of the group leader
+     * @param maxOffset the latest known offset of the configuration topic
+     * @param memberConfigs the metadata of all the members of the group as gather in the current
+     * round of rebalancing
+     * @param coordinator the worker coordinator instance that provide the configuration snapshot
+     * and get assigned the leader state during this assignment
+     * @param protocolVersion the Connect subprotocol version
+     * @return the serialized assignment of tasks to the whole group, including assigned or
+     * revoked tasks
+     */
+    protected Map performTaskAssignment(String leaderId, long maxOffset,
+                                                            Map memberConfigs,
+                                                            WorkerCoordinator coordinator, short protocolVersion) {
+        log.debug("Performing task assignment during generation: {} with memberId: {}",
+                coordinator.generationId(), coordinator.memberId());
+
+        // Base set: The previous assignment of connectors-and-tasks is a standalone snapshot that
+        // can be used to calculate derived sets
+        log.debug("Previous assignments: {}", previousAssignment);
+        int lastCompletedGenerationId = coordinator.lastCompletedGenerationId();
+        if (previousGenerationId != lastCompletedGenerationId) {
+            log.debug("Clearing the view of previous assignments due to generation mismatch between "
+                    + "previous generation ID {} and last completed generation ID {}. This can "
+                    + "happen if the leader fails to sync the assignment within a rebalancing round. "
+                    + "The following view of previous assignments might be outdated and will be "
+                    + "ignored by the leader in the current computation of new assignments. "
+                    + "Possibly outdated previous assignments: {}",
+                    previousGenerationId, lastCompletedGenerationId, previousAssignment);
+            this.previousAssignment = ConnectorsAndTasks.EMPTY;
+        }
+
+        ClusterConfigState snapshot = coordinator.configSnapshot();
+        Set configuredConnectors = new TreeSet<>(snapshot.connectors());
+        Set configuredTasks = configuredConnectors.stream()
+                .flatMap(c -> snapshot.tasks(c).stream())
+                .collect(Collectors.toSet());
+
+        // Base set: The set of configured connectors-and-tasks is a standalone snapshot that can
+        // be used to calculate derived sets
+        ConnectorsAndTasks configured = new ConnectorsAndTasks.Builder()
+                .with(configuredConnectors, configuredTasks).build();
+        log.debug("Configured assignments: {}", configured);
+
+        // Base set: The set of active connectors-and-tasks is a standalone snapshot that can be
+        // used to calculate derived sets
+        ConnectorsAndTasks activeAssignments = assignment(memberConfigs);
+        log.debug("Active assignments: {}", activeAssignments);
+
+        // This means that a previous revocation did not take effect. In this case, reset
+        // appropriately and be ready to re-apply revocation of tasks
+        if (!previousRevocation.isEmpty()) {
+            if (previousRevocation.connectors().stream().anyMatch(c -> activeAssignments.connectors().contains(c))
+                    || previousRevocation.tasks().stream().anyMatch(t -> activeAssignments.tasks().contains(t))) {
+                previousAssignment = activeAssignments;
+                canRevoke = true;
+            }
+            previousRevocation.connectors().clear();
+            previousRevocation.tasks().clear();
+        }
+
+        // Derived set: The set of deleted connectors-and-tasks is a derived set from the set
+        // difference of previous - configured
+        ConnectorsAndTasks deleted = diff(previousAssignment, configured);
+        log.debug("Deleted assignments: {}", deleted);
+
+        // Derived set: The set of remaining active connectors-and-tasks is a derived set from the
+        // set difference of active - deleted
+        ConnectorsAndTasks remainingActive = diff(activeAssignments, deleted);
+        log.debug("Remaining (excluding deleted) active assignments: {}", remainingActive);
+
+        // Derived set: The set of lost or unaccounted connectors-and-tasks is a derived set from
+        // the set difference of previous - active - deleted
+        ConnectorsAndTasks lostAssignments = diff(previousAssignment, activeAssignments, deleted);
+        log.debug("Lost assignments: {}", lostAssignments);
+
+        // Derived set: The set of new connectors-and-tasks is a derived set from the set
+        // difference of configured - previous - active
+        ConnectorsAndTasks newSubmissions = diff(configured, previousAssignment, activeAssignments);
+        log.debug("New assignments: {}", newSubmissions);
+
+        // A collection of the complete assignment
+        List completeWorkerAssignment = workerAssignment(memberConfigs, ConnectorsAndTasks.EMPTY);
+        log.debug("Complete (ignoring deletions) worker assignments: {}", completeWorkerAssignment);
+
+        // Per worker connector assignments without removing deleted connectors yet
+        Map> connectorAssignments =
+                completeWorkerAssignment.stream().collect(Collectors.toMap(WorkerLoad::worker, WorkerLoad::connectors));
+        log.debug("Complete (ignoring deletions) connector assignments: {}", connectorAssignments);
+
+        // Per worker task assignments without removing deleted connectors yet
+        Map> taskAssignments =
+                completeWorkerAssignment.stream().collect(Collectors.toMap(WorkerLoad::worker, WorkerLoad::tasks));
+        log.debug("Complete (ignoring deletions) task assignments: {}", taskAssignments);
+
+        // A collection of the current assignment excluding the connectors-and-tasks to be deleted
+        List currentWorkerAssignment = workerAssignment(memberConfigs, deleted);
+
+        Map toRevoke = computeDeleted(deleted, connectorAssignments, taskAssignments);
+        log.debug("Connector and task to delete assignments: {}", toRevoke);
+
+        // Revoking redundant connectors/tasks if the the workers have duplicate assignments
+        toRevoke.putAll(computeDuplicatedAssignments(memberConfigs, connectorAssignments, taskAssignments));
+        log.debug("Connector and task to revoke assignments (include duplicated assignments): {}", toRevoke);
+
+        // Recompute the complete assignment excluding the deleted connectors-and-tasks
+        completeWorkerAssignment = workerAssignment(memberConfigs, deleted);
+        connectorAssignments =
+                completeWorkerAssignment.stream().collect(Collectors.toMap(WorkerLoad::worker, WorkerLoad::connectors));
+        taskAssignments =
+                completeWorkerAssignment.stream().collect(Collectors.toMap(WorkerLoad::worker, WorkerLoad::tasks));
+
+        handleLostAssignments(lostAssignments, newSubmissions, completeWorkerAssignment, memberConfigs);
+
+        // Do not revoke resources for re-assignment while a delayed rebalance is active
+        // Also we do not revoke in two consecutive rebalances by the same leader
+        canRevoke = delay == 0 && canRevoke;
+
+        // Compute the connectors-and-tasks to be revoked for load balancing without taking into
+        // account the deleted ones.
+        log.debug("Can leader revoke tasks in this assignment? {} (delay: {})", canRevoke, delay);
+        if (canRevoke) {
+            Map toExplicitlyRevoke =
+                    performTaskRevocation(activeAssignments, currentWorkerAssignment);
+
+            log.debug("Connector and task to revoke assignments: {}", toRevoke);
+
+            toExplicitlyRevoke.forEach(
+                (worker, assignment) -> {
+                    ConnectorsAndTasks existing = toRevoke.computeIfAbsent(
+                        worker,
+                        v -> new ConnectorsAndTasks.Builder().build());
+                    existing.connectors().addAll(assignment.connectors());
+                    existing.tasks().addAll(assignment.tasks());
+                }
+            );
+            canRevoke = toExplicitlyRevoke.size() == 0;
+        } else {
+            canRevoke = delay == 0;
+        }
+
+        assignConnectors(completeWorkerAssignment, newSubmissions.connectors());
+        assignTasks(completeWorkerAssignment, newSubmissions.tasks());
+        log.debug("Current complete assignments: {}", currentWorkerAssignment);
+        log.debug("New complete assignments: {}", completeWorkerAssignment);
+
+        Map> currentConnectorAssignments =
+                currentWorkerAssignment.stream().collect(Collectors.toMap(WorkerLoad::worker, WorkerLoad::connectors));
+        Map> currentTaskAssignments =
+                currentWorkerAssignment.stream().collect(Collectors.toMap(WorkerLoad::worker, WorkerLoad::tasks));
+        Map> incrementalConnectorAssignments =
+                diff(connectorAssignments, currentConnectorAssignments);
+        Map> incrementalTaskAssignments =
+                diff(taskAssignments, currentTaskAssignments);
+
+        log.debug("Incremental connector assignments: {}", incrementalConnectorAssignments);
+        log.debug("Incremental task assignments: {}", incrementalTaskAssignments);
+
+        coordinator.leaderState(new LeaderState(memberConfigs, connectorAssignments, taskAssignments));
+
+        Map assignments =
+                fillAssignments(memberConfigs.keySet(), Assignment.NO_ERROR, leaderId,
+                                memberConfigs.get(leaderId).url(), maxOffset, incrementalConnectorAssignments,
+                                incrementalTaskAssignments, toRevoke, delay, protocolVersion);
+        previousAssignment = computePreviousAssignment(toRevoke, connectorAssignments, taskAssignments, lostAssignments);
+        previousGenerationId = coordinator.generationId();
+        previousMembers = memberConfigs.keySet();
+        log.debug("Actual assignments: {}", assignments);
+        return serializeAssignments(assignments);
+    }
+
+    private Map computeDeleted(ConnectorsAndTasks deleted,
+                                                           Map> connectorAssignments,
+                                                           Map> taskAssignments) {
+        // Connector to worker reverse lookup map
+        Map connectorOwners = WorkerCoordinator.invertAssignment(connectorAssignments);
+        // Task to worker reverse lookup map
+        Map taskOwners = WorkerCoordinator.invertAssignment(taskAssignments);
+
+        Map toRevoke = new HashMap<>();
+        // Add the connectors that have been deleted to the revoked set
+        deleted.connectors().forEach(c ->
+                toRevoke.computeIfAbsent(
+                    connectorOwners.get(c),
+                    v -> new ConnectorsAndTasks.Builder().build()
+                ).connectors().add(c));
+        // Add the tasks that have been deleted to the revoked set
+        deleted.tasks().forEach(t ->
+                toRevoke.computeIfAbsent(
+                    taskOwners.get(t),
+                    v -> new ConnectorsAndTasks.Builder().build()
+                ).tasks().add(t));
+        log.debug("Connectors and tasks to delete assignments: {}", toRevoke);
+        return toRevoke;
+    }
+
+    private ConnectorsAndTasks computePreviousAssignment(Map toRevoke,
+                                                         Map> connectorAssignments,
+                                                         Map> taskAssignments,
+                                                         ConnectorsAndTasks lostAssignments) {
+        ConnectorsAndTasks previousAssignment = new ConnectorsAndTasks.Builder().with(
+                connectorAssignments.values().stream().flatMap(Collection::stream).collect(Collectors.toSet()),
+                taskAssignments.values() .stream() .flatMap(Collection::stream).collect(Collectors.toSet()))
+                .build();
+
+        for (ConnectorsAndTasks revoked : toRevoke.values()) {
+            previousAssignment.connectors().removeAll(revoked.connectors());
+            previousAssignment.tasks().removeAll(revoked.tasks());
+            previousRevocation.connectors().addAll(revoked.connectors());
+            previousRevocation.tasks().addAll(revoked.tasks());
+        }
+
+        // Depends on the previous assignment's collections being sets at the moment.
+        // TODO: make it independent
+        previousAssignment.connectors().addAll(lostAssignments.connectors());
+        previousAssignment.tasks().addAll(lostAssignments.tasks());
+
+        return previousAssignment;
+    }
+
+    private ConnectorsAndTasks duplicatedAssignments(Map memberConfigs) {
+        Set connectors = memberConfigs.entrySet().stream()
+                .flatMap(memberConfig -> memberConfig.getValue().assignment().connectors().stream())
+                .collect(Collectors.groupingBy(Function.identity(), Collectors.counting()))
+                .entrySet().stream()
+                .filter(entry -> entry.getValue() > 1L)
+                .map(entry -> entry.getKey())
+                .collect(Collectors.toSet());
+
+        Set tasks = memberConfigs.values().stream()
+                .flatMap(state -> state.assignment().tasks().stream())
+                .collect(Collectors.groupingBy(Function.identity(), Collectors.counting()))
+                .entrySet().stream()
+                .filter(entry -> entry.getValue() > 1L)
+                .map(entry -> entry.getKey())
+                .collect(Collectors.toSet());
+        return new ConnectorsAndTasks.Builder().with(connectors, tasks).build();
+    }
+
+    private Map computeDuplicatedAssignments(Map memberConfigs,
+                                             Map> connectorAssignments,
+                                             Map> taskAssignment) {
+        ConnectorsAndTasks duplicatedAssignments = duplicatedAssignments(memberConfigs);
+        log.debug("Duplicated assignments: {}", duplicatedAssignments);
+
+        Map toRevoke = new HashMap<>();
+        if (!duplicatedAssignments.connectors().isEmpty()) {
+            connectorAssignments.entrySet().stream()
+                    .forEach(entry -> {
+                        Set duplicatedConnectors = new HashSet<>(duplicatedAssignments.connectors());
+                        duplicatedConnectors.retainAll(entry.getValue());
+                        if (!duplicatedConnectors.isEmpty()) {
+                            toRevoke.computeIfAbsent(
+                                entry.getKey(),
+                                v -> new ConnectorsAndTasks.Builder().build()
+                            ).connectors().addAll(duplicatedConnectors);
+                        }
+                    });
+        }
+        if (!duplicatedAssignments.tasks().isEmpty()) {
+            taskAssignment.entrySet().stream()
+                    .forEach(entry -> {
+                        Set duplicatedTasks = new HashSet<>(duplicatedAssignments.tasks());
+                        duplicatedTasks.retainAll(entry.getValue());
+                        if (!duplicatedTasks.isEmpty()) {
+                            toRevoke.computeIfAbsent(
+                                entry.getKey(),
+                                v -> new ConnectorsAndTasks.Builder().build()
+                            ).tasks().addAll(duplicatedTasks);
+                        }
+                    });
+        }
+        return toRevoke;
+    }
+
+    // visible for testing
+    protected void handleLostAssignments(ConnectorsAndTasks lostAssignments,
+                                         ConnectorsAndTasks newSubmissions,
+                                         List completeWorkerAssignment,
+                                         Map memberConfigs) {
+        if (lostAssignments.isEmpty()) {
+            resetDelay();
+            return;
+        }
+
+        final long now = time.milliseconds();
+        log.debug("Found the following connectors and tasks missing from previous assignments: "
+                + lostAssignments);
+
+        if (scheduledRebalance <= 0 && memberConfigs.keySet().containsAll(previousMembers)) {
+            log.debug("No worker seems to have departed the group during the rebalance. The "
+                    + "missing assignments that the leader is detecting are probably due to some "
+                    + "workers failing to receive the new assignments in the previous rebalance. "
+                    + "Will reassign missing tasks as new tasks");
+            newSubmissions.connectors().addAll(lostAssignments.connectors());
+            newSubmissions.tasks().addAll(lostAssignments.tasks());
+            return;
+        }
+
+        if (scheduledRebalance > 0 && now >= scheduledRebalance) {
+            // delayed rebalance expired and it's time to assign resources
+            log.debug("Delayed rebalance expired. Reassigning lost tasks");
+            Optional candidateWorkerLoad = Optional.empty();
+            if (!candidateWorkersForReassignment.isEmpty()) {
+                candidateWorkerLoad = pickCandidateWorkerForReassignment(completeWorkerAssignment);
+            }
+
+            if (candidateWorkerLoad.isPresent()) {
+                WorkerLoad workerLoad = candidateWorkerLoad.get();
+                log.debug("A candidate worker has been found to assign lost tasks: {}", workerLoad.worker());
+                lostAssignments.connectors().forEach(workerLoad::assign);
+                lostAssignments.tasks().forEach(workerLoad::assign);
+            } else {
+                log.debug("No single candidate worker was found to assign lost tasks. Treating lost tasks as new tasks");
+                newSubmissions.connectors().addAll(lostAssignments.connectors());
+                newSubmissions.tasks().addAll(lostAssignments.tasks());
+            }
+            resetDelay();
+        } else {
+            candidateWorkersForReassignment
+                    .addAll(candidateWorkersForReassignment(completeWorkerAssignment));
+            if (now < scheduledRebalance) {
+                // a delayed rebalance is in progress, but it's not yet time to reassign
+                // unaccounted resources
+                delay = calculateDelay(now);
+                log.debug("Delayed rebalance in progress. Task reassignment is postponed. New computed rebalance delay: {}", delay);
+            } else {
+                // This means scheduledRebalance == 0
+                // We could also also extract the current minimum delay from the group, to make
+                // independent of consecutive leader failures, but this optimization is skipped
+                // at the moment
+                delay = maxDelay;
+                log.debug("Resetting rebalance delay to the max: {}. scheduledRebalance: {} now: {} diff scheduledRebalance - now: {}",
+                        delay, scheduledRebalance, now, scheduledRebalance - now);
+            }
+            scheduledRebalance = now + delay;
+        }
+    }
+
+    private void resetDelay() {
+        candidateWorkersForReassignment.clear();
+        scheduledRebalance = 0;
+        if (delay != 0) {
+            log.debug("Resetting delay from previous value: {} to 0", delay);
+        }
+        delay = 0;
+    }
+
+    private Set candidateWorkersForReassignment(List completeWorkerAssignment) {
+        return completeWorkerAssignment.stream()
+                .filter(WorkerLoad::isEmpty)
+                .map(WorkerLoad::worker)
+                .collect(Collectors.toSet());
+    }
+
+    private Optional pickCandidateWorkerForReassignment(List completeWorkerAssignment) {
+        Map activeWorkers = completeWorkerAssignment.stream()
+                .collect(Collectors.toMap(WorkerLoad::worker, Function.identity()));
+        return candidateWorkersForReassignment.stream()
+                .map(activeWorkers::get)
+                .filter(Objects::nonNull)
+                .findFirst();
+    }
+
+    /**
+     * Task revocation is based on an rough estimation of the lower average number of tasks before
+     * and after new workers join the group. If no new workers join, no revocation takes place.
+     * Based on this estimation, tasks are revoked until the new floor average is reached for
+     * each existing worker. The revoked tasks, once assigned to the new workers will maintain
+     * a balanced load among the group.
+     *
+     * @param activeAssignments
+     * @param completeWorkerAssignment
+     * @return
+     */
+    private Map performTaskRevocation(ConnectorsAndTasks activeAssignments,
+                                                                  Collection completeWorkerAssignment) {
+        int totalActiveConnectorsNum = activeAssignments.connectors().size();
+        int totalActiveTasksNum = activeAssignments.tasks().size();
+        Collection existingWorkers = completeWorkerAssignment.stream()
+                .filter(wl -> wl.size() > 0)
+                .collect(Collectors.toList());
+        int existingWorkersNum = existingWorkers.size();
+        int totalWorkersNum = completeWorkerAssignment.size();
+        int newWorkersNum = totalWorkersNum - existingWorkersNum;
+
+        if (log.isDebugEnabled()) {
+            completeWorkerAssignment.forEach(wl -> log.debug(
+                    "Per worker current load size; worker: {} connectors: {} tasks: {}",
+                    wl.worker(), wl.connectorsSize(), wl.tasksSize()));
+        }
+
+        Map revoking = new HashMap<>();
+        // If there are no new workers, or no existing workers to revoke tasks from return early
+        // after logging the status
+        if (!(newWorkersNum > 0 && existingWorkersNum > 0)) {
+            log.debug("No task revocation required; workers with existing load: {} workers with "
+                    + "no load {} total workers {}",
+                    existingWorkersNum, newWorkersNum, totalWorkersNum);
+            // This is intentionally empty but mutable, because the map is used to include deleted
+            // connectors and tasks as well
+            return revoking;
+        }
+
+        log.debug("Task revocation is required; workers with existing load: {} workers with "
+                + "no load {} total workers {}",
+                existingWorkersNum, newWorkersNum, totalWorkersNum);
+
+        // We have at least one worker assignment (the leader itself) so totalWorkersNum can't be 0
+        log.debug("Previous rounded down (floor) average number of connectors per worker {}", totalActiveConnectorsNum / existingWorkersNum);
+        int floorConnectors = totalActiveConnectorsNum / totalWorkersNum;
+        log.debug("New rounded down (floor) average number of connectors per worker {}", floorConnectors);
+
+        log.debug("Previous rounded down (floor) average number of tasks per worker {}", totalActiveTasksNum / existingWorkersNum);
+        int floorTasks = totalActiveTasksNum / totalWorkersNum;
+        log.debug("New rounded down (floor) average number of tasks per worker {}", floorTasks);
+
+        int numToRevoke = floorConnectors;
+        for (WorkerLoad existing : existingWorkers) {
+            Iterator connectors = existing.connectors().iterator();
+            for (int i = existing.connectorsSize(); i > floorConnectors && numToRevoke > 0; --i, --numToRevoke) {
+                ConnectorsAndTasks resources = revoking.computeIfAbsent(
+                    existing.worker(),
+                    w -> new ConnectorsAndTasks.Builder().build());
+                resources.connectors().add(connectors.next());
+            }
+            if (numToRevoke == 0) {
+                break;
+            }
+        }
+
+        numToRevoke = floorTasks;
+        for (WorkerLoad existing : existingWorkers) {
+            Iterator tasks = existing.tasks().iterator();
+            for (int i = existing.tasksSize(); i > floorTasks && numToRevoke > 0; --i, --numToRevoke) {
+                ConnectorsAndTasks resources = revoking.computeIfAbsent(
+                    existing.worker(),
+                    w -> new ConnectorsAndTasks.Builder().build());
+                resources.tasks().add(tasks.next());
+            }
+            if (numToRevoke == 0) {
+                break;
+            }
+        }
+
+        return revoking;
+    }
+
+    private Map fillAssignments(Collection members, short error,
+                                                            String leaderId, String leaderUrl, long maxOffset,
+                                                            Map> connectorAssignments,
+                                                            Map> taskAssignments,
+                                                            Map revoked,
+                                                            int delay, short protocolVersion) {
+        Map groupAssignment = new HashMap<>();
+        for (String member : members) {
+            Collection connectorsToStart = connectorAssignments.getOrDefault(member, Collections.emptyList());
+            Collection tasksToStart = taskAssignments.getOrDefault(member, Collections.emptyList());
+            Collection connectorsToStop = revoked.getOrDefault(member, ConnectorsAndTasks.EMPTY).connectors();
+            Collection tasksToStop = revoked.getOrDefault(member, ConnectorsAndTasks.EMPTY).tasks();
+            ExtendedAssignment assignment =
+                    new ExtendedAssignment(protocolVersion, error, leaderId, leaderUrl, maxOffset,
+                            connectorsToStart, tasksToStart, connectorsToStop, tasksToStop, delay);
+            log.debug("Filling assignment: {} -> {}", member, assignment);
+            groupAssignment.put(member, assignment);
+        }
+        log.debug("Finished assignment");
+        return groupAssignment;
+    }
+
+    /**
+     * From a map of workers to assignment object generate the equivalent map of workers to byte
+     * buffers of serialized assignments.
+     *
+     * @param assignments the map of worker assignments
+     * @return the serialized map of assignments to workers
+     */
+    protected Map serializeAssignments(Map assignments) {
+        return assignments.entrySet()
+                .stream()
+                .collect(Collectors.toMap(
+                    Map.Entry::getKey,
+                    e -> IncrementalCooperativeConnectProtocol.serializeAssignment(e.getValue())));
+    }
+
+    private static ConnectorsAndTasks diff(ConnectorsAndTasks base,
+                                           ConnectorsAndTasks... toSubtract) {
+        Collection connectors = new TreeSet<>(base.connectors());
+        Collection tasks = new TreeSet<>(base.tasks());
+        for (ConnectorsAndTasks sub : toSubtract) {
+            connectors.removeAll(sub.connectors());
+            tasks.removeAll(sub.tasks());
+        }
+        return new ConnectorsAndTasks.Builder().with(connectors, tasks).build();
+    }
+
+    private static  Map> diff(Map> base,
+                                                       Map> toSubtract) {
+        Map> incremental = new HashMap<>();
+        for (Map.Entry> entry : base.entrySet()) {
+            List values = new ArrayList<>(entry.getValue());
+            values.removeAll(toSubtract.get(entry.getKey()));
+            incremental.put(entry.getKey(), values);
+        }
+        return incremental;
+    }
+
+    private ConnectorsAndTasks assignment(Map memberConfigs) {
+        log.debug("Received assignments: {}", memberConfigs);
+        Set connectors = memberConfigs.values()
+                .stream()
+                .flatMap(state -> state.assignment().connectors().stream())
+                .collect(Collectors.toSet());
+        Set tasks = memberConfigs.values()
+                .stream()
+                .flatMap(state -> state.assignment().tasks().stream())
+                .collect(Collectors.toSet());
+        return new ConnectorsAndTasks.Builder().with(connectors, tasks).build();
+    }
+
+    private int calculateDelay(long now) {
+        long diff = scheduledRebalance - now;
+        return diff > 0 ? (int) Math.min(diff, maxDelay) : 0;
+    }
+
+    /**
+     * Perform a round-robin assignment of connectors to workers with existing worker load. This
+     * assignment tries to balance the load between workers, by assigning connectors to workers
+     * that have equal load, starting with the least loaded workers.
+     *
+     * @param workerAssignment the current worker assignment; assigned connectors are added to this list
+     * @param connectors the connectors to be assigned
+     */
+    protected void assignConnectors(List workerAssignment, Collection connectors) {
+        workerAssignment.sort(WorkerLoad.connectorComparator());
+        WorkerLoad first = workerAssignment.get(0);
+
+        Iterator load = connectors.iterator();
+        while (load.hasNext()) {
+            int firstLoad = first.connectorsSize();
+            int upTo = IntStream.range(0, workerAssignment.size())
+                    .filter(i -> workerAssignment.get(i).connectorsSize() > firstLoad)
+                    .findFirst()
+                    .orElse(workerAssignment.size());
+            for (WorkerLoad worker : workerAssignment.subList(0, upTo)) {
+                String connector = load.next();
+                log.debug("Assigning connector {} to {}", connector, worker.worker());
+                worker.assign(connector);
+                if (!load.hasNext()) {
+                    break;
+                }
+            }
+        }
+    }
+
+    /**
+     * Perform a round-robin assignment of tasks to workers with existing worker load. This
+     * assignment tries to balance the load between workers, by assigning tasks to workers that
+     * have equal load, starting with the least loaded workers.
+     *
+     * @param workerAssignment the current worker assignment; assigned tasks are added to this list
+     * @param tasks the tasks to be assigned
+     */
+    protected void assignTasks(List workerAssignment, Collection tasks) {
+        workerAssignment.sort(WorkerLoad.taskComparator());
+        WorkerLoad first = workerAssignment.get(0);
+
+        Iterator load = tasks.iterator();
+        while (load.hasNext()) {
+            int firstLoad = first.tasksSize();
+            int upTo = IntStream.range(0, workerAssignment.size())
+                    .filter(i -> workerAssignment.get(i).tasksSize() > firstLoad)
+                    .findFirst()
+                    .orElse(workerAssignment.size());
+            for (WorkerLoad worker : workerAssignment.subList(0, upTo)) {
+                ConnectorTaskId task = load.next();
+                log.debug("Assigning task {} to {}", task, worker.worker());
+                worker.assign(task);
+                if (!load.hasNext()) {
+                    break;
+                }
+            }
+        }
+    }
+
+    private static List workerAssignment(Map memberConfigs,
+                                                     ConnectorsAndTasks toExclude) {
+        ConnectorsAndTasks ignore = new ConnectorsAndTasks.Builder()
+                .with(new HashSet<>(toExclude.connectors()), new HashSet<>(toExclude.tasks()))
+                .build();
+
+        return memberConfigs.entrySet().stream()
+                .map(e -> new WorkerLoad.Builder(e.getKey()).with(
+                        e.getValue().assignment().connectors().stream()
+                                .filter(v -> !ignore.connectors().contains(v))
+                                .collect(Collectors.toList()),
+                        e.getValue().assignment().tasks().stream()
+                                .filter(v -> !ignore.tasks().contains(v))
+                                .collect(Collectors.toList())
+                        ).build()
+                ).collect(Collectors.toList());
+    }
+
+}
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/IncrementalCooperativeConnectProtocol.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/IncrementalCooperativeConnectProtocol.java
new file mode 100644
index 0000000000000..6bcf9be65eb64
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/IncrementalCooperativeConnectProtocol.java
@@ -0,0 +1,274 @@
+/*
+ * 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.distributed;
+
+import org.apache.kafka.common.protocol.types.ArrayOf;
+import org.apache.kafka.common.protocol.types.Field;
+import org.apache.kafka.common.protocol.types.Schema;
+import org.apache.kafka.common.protocol.types.SchemaException;
+import org.apache.kafka.common.protocol.types.Struct;
+import org.apache.kafka.common.protocol.types.Type;
+
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.List;
+
+import static org.apache.kafka.common.message.JoinGroupRequestData.JoinGroupRequestProtocol;
+import static org.apache.kafka.common.message.JoinGroupRequestData.JoinGroupRequestProtocolCollection;
+import static org.apache.kafka.common.protocol.types.Type.NULLABLE_BYTES;
+import static org.apache.kafka.connect.runtime.distributed.ConnectProtocol.ASSIGNMENT_KEY_NAME;
+import static org.apache.kafka.connect.runtime.distributed.ConnectProtocol.CONFIG_OFFSET_KEY_NAME;
+import static org.apache.kafka.connect.runtime.distributed.ConnectProtocol.CONFIG_STATE_V0;
+import static org.apache.kafka.connect.runtime.distributed.ConnectProtocol.CONNECTOR_ASSIGNMENT_V0;
+import static org.apache.kafka.connect.runtime.distributed.ConnectProtocol.CONNECT_PROTOCOL_HEADER_SCHEMA;
+import static org.apache.kafka.connect.runtime.distributed.ConnectProtocol.CONNECT_PROTOCOL_V0;
+import static org.apache.kafka.connect.runtime.distributed.ConnectProtocol.ERROR_KEY_NAME;
+import static org.apache.kafka.connect.runtime.distributed.ConnectProtocol.LEADER_KEY_NAME;
+import static org.apache.kafka.connect.runtime.distributed.ConnectProtocol.LEADER_URL_KEY_NAME;
+import static org.apache.kafka.connect.runtime.distributed.ConnectProtocol.URL_KEY_NAME;
+import static org.apache.kafka.connect.runtime.distributed.ConnectProtocol.VERSION_KEY_NAME;
+import static org.apache.kafka.connect.runtime.distributed.ConnectProtocolCompatibility.COMPATIBLE;
+import static org.apache.kafka.connect.runtime.distributed.ConnectProtocolCompatibility.EAGER;
+import static org.apache.kafka.connect.runtime.distributed.ConnectProtocolCompatibility.SESSIONED;
+
+
+/**
+ * This class implements a group protocol for Kafka Connect workers that support incremental and
+ * cooperative rebalancing of connectors and tasks. It includes the format of worker state used when
+ * joining the group and distributing assignments, and the format of assignments of connectors
+ * and tasks to workers.
+ */
+public class IncrementalCooperativeConnectProtocol {
+    public static final String ALLOCATION_KEY_NAME = "allocation";
+    public static final String REVOKED_KEY_NAME = "revoked";
+    public static final String SCHEDULED_DELAY_KEY_NAME = "delay";
+    public static final short CONNECT_PROTOCOL_V1 = 1;
+    public static final short CONNECT_PROTOCOL_V2 = 2;
+    public static final boolean TOLERATE_MISSING_FIELDS_WITH_DEFAULTS = true;
+
+    /**
+     * Connect Protocol Header V1:
+     * 
      +     *   Version            => Int16
+     * 
      
+     */
+    private static final Struct CONNECT_PROTOCOL_HEADER_V1 = new Struct(CONNECT_PROTOCOL_HEADER_SCHEMA)
+            .set(VERSION_KEY_NAME, CONNECT_PROTOCOL_V1);
+
+    /**
+     * Connect Protocol Header V2:
+     * 
      +     *   Version            => Int16
+     * 
      
+     * The V2 protocol is schematically identical to V1, but is used to signify that internal request
+     * verification and distribution of session keys is enabled (for more information, see KIP-507:
+     * https://cwiki.apache.org/confluence/display/KAFKA/KIP-507%3A+Securing+Internal+Connect+REST+Endpoints)
+     */
+    private static final Struct CONNECT_PROTOCOL_HEADER_V2 = new Struct(CONNECT_PROTOCOL_HEADER_SCHEMA)
+        .set(VERSION_KEY_NAME, CONNECT_PROTOCOL_V2);
+
+
+    /**
+     * Config State V1:
+     * 
      +     *   Url                => [String]
+     *   ConfigOffset       => Int64
+     * 
      
+     */
+    public static final Schema CONFIG_STATE_V1 = CONFIG_STATE_V0;
+
+    /**
+     * Allocation V1
+     * 
      +     *   Current Assignment => [Byte]
+     * 
      
+     */
+    public static final Schema ALLOCATION_V1 = new Schema(
+            TOLERATE_MISSING_FIELDS_WITH_DEFAULTS,
+            new Field(ALLOCATION_KEY_NAME, NULLABLE_BYTES, null, true, null));
+
+    /**
+     *
+     * Connector Assignment V1:
+     * 
      +     *   Connector          => [String]
+     *   Tasks              => [Int32]
+     * 
      
+     *
+     * 
      Assignments for each worker are a set of connectors and tasks. These are categorized by
+     * connector ID. A sentinel task ID (CONNECTOR_TASK) is used to indicate the connector itself
+     * (i.e. that the assignment includes responsibility for running the Connector instance in
+     * addition to any tasks it generates).
      
+     */
+    public static final Schema CONNECTOR_ASSIGNMENT_V1 = CONNECTOR_ASSIGNMENT_V0;
+
+    /**
+     * Raw (non versioned) assignment V1:
+     * 
      +     *   Error              => Int16
+     *   Leader             => [String]
+     *   LeaderUrl          => [String]
+     *   ConfigOffset       => Int64
+     *   Assignment         => [Connector Assignment]
+     *   Revoked            => [Connector Assignment]
+     *   ScheduledDelay     => Int32
+     * 
      
+     */
+    public static final Schema ASSIGNMENT_V1 = new Schema(
+            TOLERATE_MISSING_FIELDS_WITH_DEFAULTS,
+            new Field(ERROR_KEY_NAME, Type.INT16),
+            new Field(LEADER_KEY_NAME, Type.STRING),
+            new Field(LEADER_URL_KEY_NAME, Type.STRING),
+            new Field(CONFIG_OFFSET_KEY_NAME, Type.INT64),
+            new Field(ASSIGNMENT_KEY_NAME, ArrayOf.nullable(CONNECTOR_ASSIGNMENT_V1), null, true, null),
+            new Field(REVOKED_KEY_NAME, ArrayOf.nullable(CONNECTOR_ASSIGNMENT_V1), null, true, null),
+            new Field(SCHEDULED_DELAY_KEY_NAME, Type.INT32, null, 0));
+
+    /**
+     * The fields are serialized in sequence as follows:
+     * Subscription V1:
+     * 
      +     *   Version            => Int16
+     *   Url                => [String]
+     *   ConfigOffset       => Int64
+     *   Current Assignment => [Byte]
+     * 
      
+     */
+    public static ByteBuffer serializeMetadata(ExtendedWorkerState workerState, boolean sessioned) {
+        Struct configState = new Struct(CONFIG_STATE_V1)
+                .set(URL_KEY_NAME, workerState.url())
+                .set(CONFIG_OFFSET_KEY_NAME, workerState.offset());
+        // Not a big issue if we embed the protocol version with the assignment in the metadata
+        Struct allocation = new Struct(ALLOCATION_V1)
+                .set(ALLOCATION_KEY_NAME, serializeAssignment(workerState.assignment()));
+        Struct connectProtocolHeader = sessioned ? CONNECT_PROTOCOL_HEADER_V2 : CONNECT_PROTOCOL_HEADER_V1;
+        ByteBuffer buffer = ByteBuffer.allocate(connectProtocolHeader.sizeOf()
+                                                + CONFIG_STATE_V1.sizeOf(configState)
+                                                + ALLOCATION_V1.sizeOf(allocation));
+        connectProtocolHeader.writeTo(buffer);
+        CONFIG_STATE_V1.write(buffer, configState);
+        ALLOCATION_V1.write(buffer, allocation);
+        buffer.flip();
+        return buffer;
+    }
+
+    /**
+     * Returns the collection of Connect protocols that are supported by this version along
+     * with their serialized metadata. The protocols are ordered by preference.
+     *
+     * @param workerState the current state of the worker metadata
+     * @param sessioned whether the {@link ConnectProtocolCompatibility#SESSIONED} protocol should
+     *                  be included in the collection of supported protocols
+     * @return the collection of Connect protocol metadata
+     */
+    public static JoinGroupRequestProtocolCollection metadataRequest(ExtendedWorkerState workerState, boolean sessioned) {
+        // Order matters in terms of protocol preference
+        List joinGroupRequestProtocols = new ArrayList<>();
+        if (sessioned) {
+            joinGroupRequestProtocols.add(new JoinGroupRequestProtocol()
+                .setName(SESSIONED.protocol())
+                .setMetadata(IncrementalCooperativeConnectProtocol.serializeMetadata(workerState, true).array())
+            );
+        }
+        joinGroupRequestProtocols.add(new JoinGroupRequestProtocol()
+                        .setName(COMPATIBLE.protocol())
+                        .setMetadata(IncrementalCooperativeConnectProtocol.serializeMetadata(workerState, false).array())
+        );
+        joinGroupRequestProtocols.add(new JoinGroupRequestProtocol()
+                        .setName(EAGER.protocol())
+                        .setMetadata(ConnectProtocol.serializeMetadata(workerState).array())
+        );
+        return new JoinGroupRequestProtocolCollection(joinGroupRequestProtocols.iterator());
+    }
+
+    /**
+     * Given a byte buffer that contains protocol metadata return the deserialized form of the
+     * metadata.
+     *
+     * @param buffer A buffer containing the protocols metadata
+     * @return the deserialized metadata
+     * @throws SchemaException on incompatible Connect protocol version
+     */
+    public static ExtendedWorkerState deserializeMetadata(ByteBuffer buffer) {
+        Struct header = CONNECT_PROTOCOL_HEADER_SCHEMA.read(buffer);
+        Short version = header.getShort(VERSION_KEY_NAME);
+        checkVersionCompatibility(version);
+        Struct configState = CONFIG_STATE_V1.read(buffer);
+        long configOffset = configState.getLong(CONFIG_OFFSET_KEY_NAME);
+        String url = configState.getString(URL_KEY_NAME);
+        Struct allocation = ALLOCATION_V1.read(buffer);
+        // Protocol version is embedded with the assignment in the metadata
+        ExtendedAssignment assignment = deserializeAssignment(allocation.getBytes(ALLOCATION_KEY_NAME));
+        return new ExtendedWorkerState(url, configOffset, assignment);
+    }
+
+    /**
+     * The fields are serialized in sequence as follows:
+     * Complete Assignment V1:
+     * 
      +     *   Version            => Int16
+     *   Error              => Int16
+     *   Leader             => [String]
+     *   LeaderUrl          => [String]
+     *   ConfigOffset       => Int64
+     *   Assignment         => [Connector Assignment]
+     *   Revoked            => [Connector Assignment]
+     *   ScheduledDelay     => Int32
+     * 
      
+     */
+    public static ByteBuffer serializeAssignment(ExtendedAssignment assignment) {
+        // comparison depends on reference equality for now
+        if (assignment == null || ExtendedAssignment.empty().equals(assignment)) {
+            return null;
+        }
+        Struct struct = assignment.toStruct();
+        ByteBuffer buffer = ByteBuffer.allocate(CONNECT_PROTOCOL_HEADER_V1.sizeOf()
+                                                + ASSIGNMENT_V1.sizeOf(struct));
+        CONNECT_PROTOCOL_HEADER_V1.writeTo(buffer);
+        ASSIGNMENT_V1.write(buffer, struct);
+        buffer.flip();
+        return buffer;
+    }
+
+    /**
+     * Given a byte buffer that contains an assignment as defined by this protocol, return the
+     * deserialized form of the assignment.
+     *
+     * @param buffer the buffer containing a serialized assignment
+     * @return the deserialized assignment
+     * @throws SchemaException on incompatible Connect protocol version
+     */
+    public static ExtendedAssignment deserializeAssignment(ByteBuffer buffer) {
+        if (buffer == null) {
+            return null;
+        }
+        Struct header = CONNECT_PROTOCOL_HEADER_SCHEMA.read(buffer);
+        Short version = header.getShort(VERSION_KEY_NAME);
+        checkVersionCompatibility(version);
+        Struct struct = ASSIGNMENT_V1.read(buffer);
+        return ExtendedAssignment.fromStruct(version, struct);
+    }
+
+    private static void checkVersionCompatibility(short version) {
+        // check for invalid versions
+        if (version < CONNECT_PROTOCOL_V0)
+            throw new SchemaException("Unsupported subscription version: " + version);
+
+        // otherwise, assume versions can be parsed
+    }
+
+}
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/WorkerCoordinator.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/WorkerCoordinator.java
index 103a323eba063..800c1d361e3ad 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/WorkerCoordinator.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/WorkerCoordinator.java
@@ -18,12 +18,11 @@
 
 import org.apache.kafka.clients.consumer.internals.AbstractCoordinator;
 import org.apache.kafka.clients.consumer.internals.ConsumerNetworkClient;
+import org.apache.kafka.clients.GroupRebalanceConfig;
 import org.apache.kafka.common.metrics.Measurable;
 import org.apache.kafka.common.metrics.MetricConfig;
 import org.apache.kafka.common.metrics.Metrics;
 import org.apache.kafka.common.requests.JoinGroupRequest;
-import org.apache.kafka.common.requests.JoinGroupRequest.ProtocolMetadata;
-import org.apache.kafka.common.utils.CircularIterator;
 import org.apache.kafka.common.utils.LogContext;
 import org.apache.kafka.common.utils.Time;
 import org.apache.kafka.common.utils.Timer;
@@ -36,55 +35,60 @@
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.Comparator;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
+
+import static org.apache.kafka.common.message.JoinGroupRequestData.JoinGroupRequestProtocolCollection;
+import static org.apache.kafka.common.message.JoinGroupResponseData.JoinGroupResponseMember;
+import static org.apache.kafka.connect.runtime.distributed.ConnectProtocolCompatibility.EAGER;
 
 /**
  * This class manages the coordination process with the Kafka group coordinator on the broker for managing assignments
  * to workers.
  */
-public final class WorkerCoordinator extends AbstractCoordinator implements Closeable {
+public class WorkerCoordinator extends AbstractCoordinator implements Closeable {
     // Currently doesn't support multiple task assignment strategies, so we just fill in a default value
     public static final String DEFAULT_SUBPROTOCOL = "default";
 
     private final Logger log;
     private final String restUrl;
     private final ConfigBackingStore configStorage;
-    private ConnectProtocol.Assignment assignmentSnapshot;
+    private volatile ExtendedAssignment assignmentSnapshot;
     private ClusterConfigState configSnapshot;
     private final WorkerRebalanceListener listener;
+    private final ConnectProtocolCompatibility protocolCompatibility;
     private LeaderState leaderState;
 
     private boolean rejoinRequested;
+    private volatile ConnectProtocolCompatibility currentConnectProtocol;
+    private volatile int lastCompletedGenerationId;
+    private final ConnectAssignor eagerAssignor;
+    private final ConnectAssignor incrementalAssignor;
+    private final int coordinatorDiscoveryTimeoutMs;
 
     /**
      * Initialize the coordination manager.
      */
-    public WorkerCoordinator(LogContext logContext,
+    public WorkerCoordinator(GroupRebalanceConfig config,
+                             LogContext logContext,
                              ConsumerNetworkClient client,
-                             String groupId,
-                             int rebalanceTimeoutMs,
-                             int sessionTimeoutMs,
-                             int heartbeatIntervalMs,
                              Metrics metrics,
                              String metricGrpPrefix,
                              Time time,
-                             long retryBackoffMs,
                              String restUrl,
                              ConfigBackingStore configStorage,
-                             WorkerRebalanceListener listener) {
-        super(logContext,
+                             WorkerRebalanceListener listener,
+                             ConnectProtocolCompatibility protocolCompatibility,
+                             int maxDelay) {
+        super(config,
+              logContext,
               client,
-              groupId,
-              rebalanceTimeoutMs,
-              sessionTimeoutMs,
-              heartbeatIntervalMs,
               metrics,
               metricGrpPrefix,
-              time,
-              retryBackoffMs,
-              true);
+              time);
         this.log = logContext.logger(WorkerCoordinator.class);
         this.restUrl = restUrl;
         this.configStorage = configStorage;
@@ -92,6 +96,12 @@ public WorkerCoordinator(LogContext logContext,
         new WorkerCoordinatorMetrics(metrics, metricGrpPrefix);
         this.listener = listener;
         this.rejoinRequested = false;
+        this.protocolCompatibility = protocolCompatibility;
+        this.incrementalAssignor = new IncrementalCooperativeAssignor(logContext, time, maxDelay);
+        this.eagerAssignor = new EagerAssignor(logContext);
+        this.currentConnectProtocol = protocolCompatibility;
+        this.coordinatorDiscoveryTimeoutMs = config.heartbeatIntervalMs;
+        this.lastCompletedGenerationId = Generation.NO_GENERATION.generationId;
     }
 
     @Override
@@ -118,7 +128,20 @@ public void poll(long timeout) {
 
         do {
             if (coordinatorUnknown()) {
-                ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
+                log.debug("Broker coordinator is marked unknown. Attempting discovery with a timeout of {}ms",
+                        coordinatorDiscoveryTimeoutMs);
+                if (ensureCoordinatorReady(time.timer(coordinatorDiscoveryTimeoutMs))) {
+                    log.debug("Broker coordinator is ready");
+                } else {
+                    log.debug("Can not connect to broker coordinator");
+                    final ExtendedAssignment localAssignmentSnapshot = assignmentSnapshot;
+                    if (localAssignmentSnapshot != null && !localAssignmentSnapshot.failed()) {
+                        log.info("Broker coordinator was unreachable for {}ms. Revoking previous assignment {} to " +
+                                "avoid running tasks while not being a member the group", coordinatorDiscoveryTimeoutMs, localAssignmentSnapshot);
+                        listener.onRevoked(localAssignmentSnapshot.leader(), localAssignmentSnapshot.connectors(), localAssignmentSnapshot.tasks());
+                        assignmentSnapshot = null;
+                    }
+                }
                 now = time.milliseconds();
             }
 
@@ -144,177 +167,181 @@ public void poll(long timeout) {
     }
 
     @Override
-    public List metadata() {
+    public JoinGroupRequestProtocolCollection metadata() {
         configSnapshot = configStorage.snapshot();
-        ConnectProtocol.WorkerState workerState = new ConnectProtocol.WorkerState(restUrl, configSnapshot.offset());
-        ByteBuffer metadata = ConnectProtocol.serializeMetadata(workerState);
-        return Collections.singletonList(new ProtocolMetadata(DEFAULT_SUBPROTOCOL, metadata));
+        final ExtendedAssignment localAssignmentSnapshot = assignmentSnapshot;
+        ExtendedWorkerState workerState = new ExtendedWorkerState(restUrl, configSnapshot.offset(), localAssignmentSnapshot);
+        switch (protocolCompatibility) {
+            case EAGER:
+                return ConnectProtocol.metadataRequest(workerState);
+            case COMPATIBLE:
+                return IncrementalCooperativeConnectProtocol.metadataRequest(workerState, false);
+            case SESSIONED:
+                return IncrementalCooperativeConnectProtocol.metadataRequest(workerState, true);
+            default:
+                throw new IllegalStateException("Unknown Connect protocol compatibility mode " + protocolCompatibility);
+        }
     }
 
     @Override
     protected void onJoinComplete(int generation, String memberId, String protocol, ByteBuffer memberAssignment) {
-        assignmentSnapshot = ConnectProtocol.deserializeAssignment(memberAssignment);
+        ExtendedAssignment newAssignment = IncrementalCooperativeConnectProtocol.deserializeAssignment(memberAssignment);
+        log.debug("Deserialized new assignment: {}", newAssignment);
+        currentConnectProtocol = ConnectProtocolCompatibility.fromProtocol(protocol);
         // At this point we always consider ourselves to be a member of the cluster, even if there was an assignment
         // error (the leader couldn't make the assignment) or we are behind the config and cannot yet work on our assigned
         // tasks. It's the responsibility of the code driving this process to decide how to react (e.g. trying to get
         // up to date, try to rejoin again, leaving the group and backing off, etc.).
         rejoinRequested = false;
-        listener.onAssigned(assignmentSnapshot, generation);
-    }
-
-    @Override
-    protected Map performAssignment(String leaderId, String protocol, Map allMemberMetadata) {
-        log.debug("Performing task assignment");
-
-        Map memberConfigs = new HashMap<>();
-        for (Map.Entry entry : allMemberMetadata.entrySet())
-            memberConfigs.put(entry.getKey(), ConnectProtocol.deserializeMetadata(entry.getValue()));
-
-        long maxOffset = findMaxMemberConfigOffset(memberConfigs);
-        Long leaderOffset = ensureLeaderConfig(maxOffset);
-        if (leaderOffset == null)
-            return fillAssignmentsAndSerialize(memberConfigs.keySet(), ConnectProtocol.Assignment.CONFIG_MISMATCH,
-                    leaderId, memberConfigs.get(leaderId).url(), maxOffset,
-                    new HashMap>(), new HashMap>());
-        return performTaskAssignment(leaderId, leaderOffset, memberConfigs);
-    }
-
-    private long findMaxMemberConfigOffset(Map memberConfigs) {
-        // The new config offset is the maximum seen by any member. We always perform assignment using this offset,
-        // even if some members have fallen behind. The config offset used to generate the assignment is included in
-        // the response so members that have fallen behind will not use the assignment until they have caught up.
-        Long maxOffset = null;
-        for (Map.Entry stateEntry : memberConfigs.entrySet()) {
-            long memberRootOffset = stateEntry.getValue().offset();
-            if (maxOffset == null)
-                maxOffset = memberRootOffset;
-            else
-                maxOffset = Math.max(maxOffset, memberRootOffset);
-        }
-
-        log.debug("Max config offset root: {}, local snapshot config offsets root: {}",
-                maxOffset, configSnapshot.offset());
-        return maxOffset;
-    }
-
-    private Long ensureLeaderConfig(long maxOffset) {
-        // If this leader is behind some other members, we can't do assignment
-        if (configSnapshot.offset() < maxOffset) {
-            // We might be able to take a new snapshot to catch up immediately and avoid another round of syncing here.
-            // Alternatively, if this node has already passed the maximum reported by any other member of the group, it
-            // is also safe to use this newer state.
-            ClusterConfigState updatedSnapshot = configStorage.snapshot();
-            if (updatedSnapshot.offset() < maxOffset) {
-                log.info("Was selected to perform assignments, but do not have latest config found in sync request. " +
-                        "Returning an empty configuration to trigger re-sync.");
-                return null;
-            } else {
-                configSnapshot = updatedSnapshot;
-                return configSnapshot.offset();
+        if (currentConnectProtocol != EAGER) {
+            if (!newAssignment.revokedConnectors().isEmpty() || !newAssignment.revokedTasks().isEmpty()) {
+                listener.onRevoked(newAssignment.leader(), newAssignment.revokedConnectors(), newAssignment.revokedTasks());
             }
-        }
-
-        return maxOffset;
-    }
 
-    private Map performTaskAssignment(String leaderId, long maxOffset, Map memberConfigs) {
-        Map> connectorAssignments = new HashMap<>();
-        Map> taskAssignments = new HashMap<>();
-
-        // Perform round-robin task assignment. Assign all connectors and then all tasks because assigning both the
-        // connector and its tasks can lead to very uneven distribution of work in some common cases (e.g. for connectors
-        // that generate only 1 task each; in a cluster of 2 or an even # of nodes, only even nodes will be assigned
-        // connectors and only odd nodes will be assigned tasks, but tasks are, on average, actually more resource
-        // intensive than connectors).
-        List connectorsSorted = sorted(configSnapshot.connectors());
-        CircularIterator memberIt = new CircularIterator<>(sorted(memberConfigs.keySet()));
-        for (String connectorId : connectorsSorted) {
-            String connectorAssignedTo = memberIt.next();
-            log.trace("Assigning connector {} to {}", connectorId, connectorAssignedTo);
-            List memberConnectors = connectorAssignments.get(connectorAssignedTo);
-            if (memberConnectors == null) {
-                memberConnectors = new ArrayList<>();
-                connectorAssignments.put(connectorAssignedTo, memberConnectors);
+            final ExtendedAssignment localAssignmentSnapshot = assignmentSnapshot;
+            if (localAssignmentSnapshot != null) {
+                localAssignmentSnapshot.connectors().removeAll(newAssignment.revokedConnectors());
+                localAssignmentSnapshot.tasks().removeAll(newAssignment.revokedTasks());
+                log.debug("After revocations snapshot of assignment: {}", localAssignmentSnapshot);
+                newAssignment.connectors().addAll(localAssignmentSnapshot.connectors());
+                newAssignment.tasks().addAll(localAssignmentSnapshot.tasks());
             }
-            memberConnectors.add(connectorId);
+            log.debug("Augmented new assignment: {}", newAssignment);
         }
-        for (String connectorId : connectorsSorted) {
-            for (ConnectorTaskId taskId : sorted(configSnapshot.tasks(connectorId))) {
-                String taskAssignedTo = memberIt.next();
-                log.trace("Assigning task {} to {}", taskId, taskAssignedTo);
-                List memberTasks = taskAssignments.get(taskAssignedTo);
-                if (memberTasks == null) {
-                    memberTasks = new ArrayList<>();
-                    taskAssignments.put(taskAssignedTo, memberTasks);
-                }
-                memberTasks.add(taskId);
-            }
-        }
-
-        this.leaderState = new LeaderState(memberConfigs, connectorAssignments, taskAssignments);
-
-        return fillAssignmentsAndSerialize(memberConfigs.keySet(), ConnectProtocol.Assignment.NO_ERROR,
-                leaderId, memberConfigs.get(leaderId).url(), maxOffset, connectorAssignments, taskAssignments);
+        assignmentSnapshot = newAssignment;
+        lastCompletedGenerationId = generation;
+        listener.onAssigned(newAssignment, generation);
     }
 
-    private Map fillAssignmentsAndSerialize(Collection members,
-                                                                short error,
-                                                                String leaderId,
-                                                                String leaderUrl,
-                                                                long maxOffset,
-                                                                Map> connectorAssignments,
-                                                                Map> taskAssignments) {
-
-        Map groupAssignment = new HashMap<>();
-        for (String member : members) {
-            List connectors = connectorAssignments.get(member);
-            if (connectors == null)
-                connectors = Collections.emptyList();
-            List tasks = taskAssignments.get(member);
-            if (tasks == null)
-                tasks = Collections.emptyList();
-            ConnectProtocol.Assignment assignment = new ConnectProtocol.Assignment(error, leaderId, leaderUrl, maxOffset, connectors, tasks);
-            log.debug("Assignment: {} -> {}", member, assignment);
-            groupAssignment.put(member, ConnectProtocol.serializeAssignment(assignment));
-        }
-        log.debug("Finished assignment");
-        return groupAssignment;
+    @Override
+    protected Map performAssignment(String leaderId, String protocol, List allMemberMetadata) {
+        return ConnectProtocolCompatibility.fromProtocol(protocol) == EAGER
+               ? eagerAssignor.performAssignment(leaderId, protocol, allMemberMetadata, this)
+               : incrementalAssignor.performAssignment(leaderId, protocol, allMemberMetadata, this);
     }
 
     @Override
     protected void onJoinPrepare(int generation, String memberId) {
-        this.leaderState = null;
-        log.debug("Revoking previous assignment {}", assignmentSnapshot);
-        if (assignmentSnapshot != null && !assignmentSnapshot.failed())
-            listener.onRevoked(assignmentSnapshot.leader(), assignmentSnapshot.connectors(), assignmentSnapshot.tasks());
+        log.info("Rebalance started");
+        leaderState(null);
+        final ExtendedAssignment localAssignmentSnapshot = assignmentSnapshot;
+        if (currentConnectProtocol == EAGER) {
+            log.debug("Revoking previous assignment {}", localAssignmentSnapshot);
+            if (localAssignmentSnapshot != null && !localAssignmentSnapshot.failed())
+                listener.onRevoked(localAssignmentSnapshot.leader(), localAssignmentSnapshot.connectors(), localAssignmentSnapshot.tasks());
+        } else {
+            log.debug("Cooperative rebalance triggered. Keeping assignment {} until it's "
+                      + "explicitly revoked.", localAssignmentSnapshot);
+        }
     }
 
     @Override
     protected boolean rejoinNeededOrPending() {
-        return super.rejoinNeededOrPending() || (assignmentSnapshot == null || assignmentSnapshot.failed()) || rejoinRequested;
+        final ExtendedAssignment localAssignmentSnapshot = assignmentSnapshot;
+        return super.rejoinNeededOrPending() || (localAssignmentSnapshot == null || localAssignmentSnapshot.failed()) || rejoinRequested;
     }
 
+    @Override
     public String memberId() {
-        Generation generation = generation();
+        Generation generation = generationIfStable();
         if (generation != null)
             return generation.memberId;
         return JoinGroupRequest.UNKNOWN_MEMBER_ID;
     }
 
+    /**
+     * Return the current generation. The generation refers to this worker's knowledge with
+     * respect to which  generation is the latest one and, therefore, this information is local.
+     *
+     * @return the generation ID or -1 if no generation is defined
+     */
+    public int generationId() {
+        return super.generation().generationId;
+    }
+
+    /**
+     * Return id that corresponds to the group generation that was active when the last join was successful
+     *
+     * @return the generation ID of the last group that was joined successfully by this member or -1 if no generation
+     * was stable at that point
+     */
+    public int lastCompletedGenerationId() {
+        return lastCompletedGenerationId;
+    }
+
+    public void revokeAssignment(ExtendedAssignment assignment) {
+        listener.onRevoked(assignment.leader(), assignment.connectors(), assignment.tasks());
+    }
+
     private boolean isLeader() {
-        return assignmentSnapshot != null && memberId().equals(assignmentSnapshot.leader());
+        final ExtendedAssignment localAssignmentSnapshot = assignmentSnapshot;
+        return localAssignmentSnapshot != null && memberId().equals(localAssignmentSnapshot.leader());
     }
 
     public String ownerUrl(String connector) {
         if (rejoinNeededOrPending() || !isLeader())
             return null;
-        return leaderState.ownerUrl(connector);
+        return leaderState().ownerUrl(connector);
     }
 
     public String ownerUrl(ConnectorTaskId task) {
         if (rejoinNeededOrPending() || !isLeader())
             return null;
-        return leaderState.ownerUrl(task);
+        return leaderState().ownerUrl(task);
+    }
+
+    /**
+     * Get an up-to-date snapshot of the cluster configuration.
+     *
+     * @return the state of the cluster configuration; the result is not locally cached
+     */
+    public ClusterConfigState configFreshSnapshot() {
+        return configStorage.snapshot();
+    }
+
+    /**
+     * Get a snapshot of the cluster configuration.
+     *
+     * @return the state of the cluster configuration
+     */
+    public ClusterConfigState configSnapshot() {
+        return configSnapshot;
+    }
+
+    /**
+     * Set the state of the cluster configuration to this worker coordinator.
+     *
+     * @param update the updated state of the cluster configuration
+     */
+    public void configSnapshot(ClusterConfigState update) {
+        configSnapshot = update;
+    }
+
+    /**
+     * Get the leader state stored in this worker coordinator.
+     *
+     * @return the leader state
+     */
+    private LeaderState leaderState() {
+        return leaderState;
+    }
+
+    /**
+     * Store the leader state to this worker coordinator.
+     *
+     * @param update the updated leader state
+     */
+    public void leaderState(LeaderState update) {
+        leaderState = update;
+    }
+
+    /**
+     * Get the version of the connect protocol that is currently active in the group of workers.
+     *
+     * @return the current connect protocol version
+     */
+    public short currentProtocolVersion() {
+        return currentConnectProtocol.protocolVersion();
     }
 
     private class WorkerCoordinatorMetrics {
@@ -326,14 +353,22 @@ public WorkerCoordinatorMetrics(Metrics metrics, String metricGrpPrefix) {
             Measurable numConnectors = new Measurable() {
                 @Override
                 public double measure(MetricConfig config, long now) {
-                    return assignmentSnapshot.connectors().size();
+                    final ExtendedAssignment localAssignmentSnapshot = assignmentSnapshot;
+                    if (localAssignmentSnapshot == null) {
+                        return 0.0;
+                    }
+                    return localAssignmentSnapshot.connectors().size();
                 }
             };
 
             Measurable numTasks = new Measurable() {
                 @Override
                 public double measure(MetricConfig config, long now) {
-                    return assignmentSnapshot.tasks().size();
+                    final ExtendedAssignment localAssignmentSnapshot = assignmentSnapshot;
+                    if (localAssignmentSnapshot == null) {
+                        return 0.0;
+                    }
+                    return localAssignmentSnapshot.tasks().size();
                 }
             };
 
@@ -346,15 +381,9 @@ public double measure(MetricConfig config, long now) {
         }
     }
 
-    private static > List sorted(Collection members) {
-        List res = new ArrayList<>(members);
-        Collections.sort(res);
-        return res;
-    }
-
-    private static  Map invertAssignment(Map> assignment) {
+    public static  Map invertAssignment(Map> assignment) {
         Map inverted = new HashMap<>();
-        for (Map.Entry> assignmentEntry : assignment.entrySet()) {
+        for (Map.Entry> assignmentEntry : assignment.entrySet()) {
             K key = assignmentEntry.getKey();
             for (V value : assignmentEntry.getValue())
                 inverted.put(value, key);
@@ -362,14 +391,14 @@ private static  Map invertAssignment(Map> assignment) {
         return inverted;
     }
 
-    private static class LeaderState {
-        private final Map allMembers;
+    public static class LeaderState {
+        private final Map allMembers;
         private final Map connectorOwners;
         private final Map taskOwners;
 
-        public LeaderState(Map allMembers,
-                           Map> connectorAssignment,
-                           Map> taskAssignment) {
+        public LeaderState(Map allMembers,
+                           Map> connectorAssignment,
+                           Map> taskAssignment) {
             this.allMembers = allMembers;
             this.connectorOwners = invertAssignment(connectorAssignment);
             this.taskOwners = invertAssignment(taskAssignment);
@@ -391,4 +420,194 @@ private String ownerUrl(String connector) {
 
     }
 
+    public static class ConnectorsAndTasks {
+        public static final ConnectorsAndTasks EMPTY =
+                new ConnectorsAndTasks(Collections.emptyList(), Collections.emptyList());
+
+        private final Collection connectors;
+        private final Collection tasks;
+
+        private ConnectorsAndTasks(Collection connectors, Collection tasks) {
+            this.connectors = connectors;
+            this.tasks = tasks;
+        }
+
+        public static class Builder {
+            private Collection withConnectors;
+            private Collection withTasks;
+
+            public Builder() {
+            }
+
+            public ConnectorsAndTasks.Builder withCopies(Collection connectors,
+                                                         Collection tasks) {
+                withConnectors = new ArrayList<>(connectors);
+                withTasks = new ArrayList<>(tasks);
+                return this;
+            }
+
+            public ConnectorsAndTasks.Builder with(Collection connectors,
+                                                   Collection tasks) {
+                withConnectors = new ArrayList<>(connectors);
+                withTasks = new ArrayList<>(tasks);
+                return this;
+            }
+
+            public ConnectorsAndTasks build() {
+                return new ConnectorsAndTasks(
+                        withConnectors != null ? withConnectors : new ArrayList<>(),
+                        withTasks != null ? withTasks : new ArrayList<>());
+            }
+        }
+
+        public Collection connectors() {
+            return connectors;
+        }
+
+        public Collection tasks() {
+            return tasks;
+        }
+
+        public int size() {
+            return connectors.size() + tasks.size();
+        }
+
+        public boolean isEmpty() {
+            return connectors.isEmpty() && tasks.isEmpty();
+        }
+
+        @Override
+        public String toString() {
+            return "{ connectorIds=" + connectors + ", taskIds=" + tasks + '}';
+        }
+    }
+
+    public static class WorkerLoad {
+        private final String worker;
+        private final Collection connectors;
+        private final Collection tasks;
+
+        private WorkerLoad(
+                String worker,
+                Collection connectors,
+                Collection tasks
+        ) {
+            this.worker = worker;
+            this.connectors = connectors;
+            this.tasks = tasks;
+        }
+
+        public static class Builder {
+            private String withWorker;
+            private Collection withConnectors;
+            private Collection withTasks;
+
+            public Builder(String worker) {
+                this.withWorker = Objects.requireNonNull(worker, "worker cannot be null");
+            }
+
+            public WorkerLoad.Builder withCopies(Collection connectors,
+                                                 Collection tasks) {
+                withConnectors = new ArrayList<>(
+                        Objects.requireNonNull(connectors, "connectors may be empty but not null"));
+                withTasks = new ArrayList<>(
+                        Objects.requireNonNull(tasks, "tasks may be empty but not null"));
+                return this;
+            }
+
+            public WorkerLoad.Builder with(Collection connectors,
+                                           Collection tasks) {
+                withConnectors = Objects.requireNonNull(connectors,
+                        "connectors may be empty but not null");
+                withTasks = Objects.requireNonNull(tasks, "tasks may be empty but not null");
+                return this;
+            }
+
+            public WorkerLoad build() {
+                return new WorkerLoad(
+                        withWorker,
+                        withConnectors != null ? withConnectors : new ArrayList<>(),
+                        withTasks != null ? withTasks : new ArrayList<>());
+            }
+        }
+
+        public String worker() {
+            return worker;
+        }
+
+        public Collection connectors() {
+            return connectors;
+        }
+
+        public Collection tasks() {
+            return tasks;
+        }
+
+        public int connectorsSize() {
+            return connectors.size();
+        }
+
+        public int tasksSize() {
+            return tasks.size();
+        }
+
+        public void assign(String connector) {
+            connectors.add(connector);
+        }
+
+        public void assign(ConnectorTaskId task) {
+            tasks.add(task);
+        }
+
+        public int size() {
+            return connectors.size() + tasks.size();
+        }
+
+        public boolean isEmpty() {
+            return connectors.isEmpty() && tasks.isEmpty();
+        }
+
+        public static Comparator connectorComparator() {
+            return (left, right) -> {
+                int res = left.connectors.size() - right.connectors.size();
+                return res != 0 ? res : left.worker == null
+                                        ? right.worker == null ? 0 : -1
+                                        : left.worker.compareTo(right.worker);
+            };
+        }
+
+        public static Comparator taskComparator() {
+            return (left, right) -> {
+                int res = left.tasks.size() - right.tasks.size();
+                return res != 0 ? res : left.worker == null
+                                        ? right.worker == null ? 0 : -1
+                                        : left.worker.compareTo(right.worker);
+            };
+        }
+
+        @Override
+        public String toString() {
+            return "{ worker=" + worker + ", connectorIds=" + connectors + ", taskIds=" + tasks + '}';
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) {
+                return true;
+            }
+            if (!(o instanceof WorkerLoad)) {
+                return false;
+            }
+            WorkerLoad that = (WorkerLoad) o;
+            return worker.equals(that.worker) &&
+                    connectors.equals(that.connectors) &&
+                    tasks.equals(that.tasks);
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(worker, connectors, tasks);
+        }
+    }
+
 }
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/WorkerGroupMember.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/WorkerGroupMember.java
index 6591be085efb1..739575e38904c 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/WorkerGroupMember.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/WorkerGroupMember.java
@@ -23,6 +23,7 @@
 import org.apache.kafka.clients.Metadata;
 import org.apache.kafka.clients.NetworkClient;
 import org.apache.kafka.clients.consumer.internals.ConsumerNetworkClient;
+import org.apache.kafka.clients.GroupRebalanceConfig;
 import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.internals.ClusterResourceListeners;
 import org.apache.kafka.common.metrics.JmxReporter;
@@ -34,6 +35,7 @@
 import org.apache.kafka.common.utils.AppInfoParser;
 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.connect.storage.ConfigBackingStore;
 import org.apache.kafka.connect.util.ConnectorTaskId;
 import org.slf4j.Logger;
@@ -44,7 +46,6 @@
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.TimeUnit;
-import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 
 /**
@@ -54,8 +55,6 @@
  * higher level operations in response to group membership events being handled by the herder.
  */
 public class WorkerGroupMember {
-
-    private static final AtomicInteger CONNECT_CLIENT_ID_SEQUENCE = new AtomicInteger(1);
     private static final String JMX_PREFIX = "kafka.connect";
 
     private final Logger log;
@@ -73,15 +72,12 @@ public WorkerGroupMember(DistributedConfig config,
                              String restUrl,
                              ConfigBackingStore configStorage,
                              WorkerRebalanceListener listener,
-                             Time time) {
+                             Time time,
+                             String clientId,
+                             LogContext logContext) {
         try {
             this.time = time;
-
-            String clientIdConfig = config.getString(CommonClientConfigs.CLIENT_ID_CONFIG);
-            clientId = clientIdConfig.length() <= 0 ? "connect-" + CONNECT_CLIENT_ID_SEQUENCE.getAndIncrement() : clientIdConfig;
-            String groupId = config.getString(DistributedConfig.GROUP_ID_CONFIG);
-
-            LogContext logContext = new LogContext("[Worker clientId=" + clientId + ", groupId=" + groupId + "] ");
+            this.clientId = clientId;
             this.log = logContext.logger(WorkerGroupMember.class);
 
             Map metricsTags = new LinkedHashMap<>();
@@ -100,7 +96,7 @@ public WorkerGroupMember(DistributedConfig config,
             List addresses = ClientUtils.parseAndValidateAddresses(
                     config.getList(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG),
                     config.getString(CommonClientConfigs.CLIENT_DNS_LOOKUP_CONFIG));
-            this.metadata.bootstrap(addresses, time.milliseconds());
+            this.metadata.bootstrap(addresses);
             String metricGrpPrefix = "connect";
             ChannelBuilder channelBuilder = ClientUtils.createChannelBuilder(config, time);
             NetworkClient netClient = new NetworkClient(
@@ -127,21 +123,19 @@ public WorkerGroupMember(DistributedConfig config,
                     config.getInt(CommonClientConfigs.REQUEST_TIMEOUT_MS_CONFIG),
                     Integer.MAX_VALUE);
             this.coordinator = new WorkerCoordinator(
+                    new GroupRebalanceConfig(config, GroupRebalanceConfig.ProtocolType.CONNECT),
                     logContext,
                     this.client,
-                    groupId,
-                    config.getInt(DistributedConfig.REBALANCE_TIMEOUT_MS_CONFIG),
-                    config.getInt(DistributedConfig.SESSION_TIMEOUT_MS_CONFIG),
-                    config.getInt(DistributedConfig.HEARTBEAT_INTERVAL_MS_CONFIG),
                     metrics,
                     metricGrpPrefix,
                     this.time,
-                    retryBackoffMs,
                     restUrl,
                     configStorage,
-                    listener);
+                    listener,
+                    ConnectProtocolCompatibility.compatibility(config.getString(DistributedConfig.CONNECT_PROTOCOL_CONFIG)),
+                    config.getInt(DistributedConfig.SCHEDULED_REBALANCE_MAX_DELAY_MS_CONFIG));
 
-            AppInfoParser.registerAppInfo(JMX_PREFIX, clientId, metrics);
+            AppInfoParser.registerAppInfo(JMX_PREFIX, clientId, metrics, time.milliseconds());
             log.debug("Connect group member created");
         } catch (Throwable t) {
             // call close methods if internal objects are already constructed
@@ -157,6 +151,10 @@ public void stop() {
         stop(false);
     }
 
+    /**
+     * Ensure that the connection to the broker coordinator is up and that the worker is an
+     * active member of the group.
+     */
     public void ensureActive() {
         coordinator.poll(0);
     }
@@ -189,8 +187,8 @@ public void requestRejoin() {
         coordinator.requestRejoin();
     }
 
-    public void maybeLeaveGroup() {
-        coordinator.maybeLeaveGroup();
+    public void maybeLeaveGroup(String leaveReason) {
+        coordinator.maybeLeaveGroup(leaveReason);
     }
 
     public String ownerUrl(String connector) {
@@ -201,13 +199,26 @@ public String ownerUrl(ConnectorTaskId task) {
         return coordinator.ownerUrl(task);
     }
 
+    /**
+     * Get the version of the connect protocol that is currently active in the group of workers.
+     *
+     * @return the current connect protocol version
+     */
+    public short currentProtocolVersion() {
+        return coordinator.currentProtocolVersion();
+    }
+
+    public void revokeAssignment(ExtendedAssignment assignment) {
+        coordinator.revokeAssignment(assignment);
+    }
+
     private void stop(boolean swallowException) {
         log.trace("Stopping the Connect group member.");
         AtomicReference firstException = new AtomicReference<>();
         this.stopped = true;
-        ClientUtils.closeQuietly(coordinator, "coordinator", firstException);
-        ClientUtils.closeQuietly(metrics, "consumer metrics", firstException);
-        ClientUtils.closeQuietly(client, "consumer network client", firstException);
+        Utils.closeQuietly(coordinator, "coordinator", firstException);
+        Utils.closeQuietly(metrics, "consumer metrics", firstException);
+        Utils.closeQuietly(client, "consumer network client", firstException);
         AppInfoParser.unregisterAppInfo(JMX_PREFIX, clientId, metrics);
         if (firstException.get() != null && !swallowException)
             throw new KafkaException("Failed to stop the Connect group member", firstException.get());
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/WorkerRebalanceListener.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/WorkerRebalanceListener.java
index 6ff5ce48353a4..93d03272813b6 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/WorkerRebalanceListener.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/WorkerRebalanceListener.java
@@ -25,13 +25,15 @@
  */
 public interface WorkerRebalanceListener {
     /**
-     * Invoked when a new assignment is created by joining the Connect worker group. This is invoked for both successful
-     * and unsuccessful assignments.
+     * Invoked when a new assignment is created by joining the Connect worker group. This is
+     * invoked for both successful and unsuccessful assignments.
      */
-    void onAssigned(ConnectProtocol.Assignment assignment, int generation);
+    void onAssigned(ExtendedAssignment assignment, int generation);
 
     /**
-     * Invoked when a rebalance operation starts, revoking ownership for the set of connectors and tasks.
+     * Invoked when a rebalance operation starts, revoking ownership for the set of connectors
+     * and tasks. Depending on the Connect protocol version, the collection of revoked connectors
+     * or tasks might refer to all or some of the connectors and tasks running on the worker.
      */
     void onRevoked(String leader, Collection connectors, Collection tasks);
 }
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/errors/DeadLetterQueueReporter.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/errors/DeadLetterQueueReporter.java
index 231226997833d..20ed2f2d577f2 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/errors/DeadLetterQueueReporter.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/errors/DeadLetterQueueReporter.java
@@ -16,7 +16,7 @@
  */
 package org.apache.kafka.connect.runtime.errors;
 
-import org.apache.kafka.clients.admin.AdminClient;
+import org.apache.kafka.clients.admin.Admin;
 import org.apache.kafka.clients.admin.NewTopic;
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.clients.producer.KafkaProducer;
@@ -26,7 +26,6 @@
 import org.apache.kafka.common.record.RecordBatch;
 import org.apache.kafka.connect.errors.ConnectException;
 import org.apache.kafka.connect.runtime.SinkConnectorConfig;
-import org.apache.kafka.connect.runtime.WorkerConfig;
 import org.apache.kafka.connect.util.ConnectorTaskId;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -71,13 +70,13 @@ public class DeadLetterQueueReporter implements ErrorReporter {
 
     private KafkaProducer kafkaProducer;
 
-    public static DeadLetterQueueReporter createAndSetup(WorkerConfig workerConfig,
+    public static DeadLetterQueueReporter createAndSetup(Map adminProps,
                                                          ConnectorTaskId id,
                                                          SinkConnectorConfig sinkConfig, Map producerProps,
                                                          ErrorHandlingMetrics errorHandlingMetrics) {
         String topic = sinkConfig.dlqTopicName();
 
-        try (AdminClient admin = AdminClient.create(workerConfig.originals())) {
+        try (Admin admin = Admin.create(adminProps)) {
             if (!admin.listTopics().names().get().contains(topic)) {
                 log.error("Topic {} doesn't exist. Will attempt to create topic.", topic);
                 NewTopic schemaTopicRequest = new NewTopic(topic, DLQ_NUM_DESIRED_PARTITIONS, sinkConfig.dlqTopicReplicationFactor());
@@ -205,4 +204,9 @@ private byte[] toBytes(String value) {
             return null;
         }
     }
+
+    @Override
+    public void close() {
+        kafkaProducer.close();
+    }
 }
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/errors/ErrorHandlingMetrics.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/errors/ErrorHandlingMetrics.java
index c589012d02cc2..0deecd129a06e 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/errors/ErrorHandlingMetrics.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/errors/ErrorHandlingMetrics.java
@@ -17,7 +17,7 @@
 package org.apache.kafka.connect.runtime.errors;
 
 import org.apache.kafka.common.metrics.Sensor;
-import org.apache.kafka.common.metrics.stats.Total;
+import org.apache.kafka.common.metrics.stats.CumulativeSum;
 import org.apache.kafka.common.utils.SystemTime;
 import org.apache.kafka.common.utils.Time;
 import org.apache.kafka.connect.runtime.ConnectMetrics;
@@ -62,25 +62,25 @@ public ErrorHandlingMetrics(ConnectorTaskId id, ConnectMetrics connectMetrics) {
         metricGroup.close();
 
         recordProcessingFailures = metricGroup.sensor("total-record-failures");
-        recordProcessingFailures.add(metricGroup.metricName(registry.recordProcessingFailures), new Total());
+        recordProcessingFailures.add(metricGroup.metricName(registry.recordProcessingFailures), new CumulativeSum());
 
         recordProcessingErrors = metricGroup.sensor("total-record-errors");
-        recordProcessingErrors.add(metricGroup.metricName(registry.recordProcessingErrors), new Total());
+        recordProcessingErrors.add(metricGroup.metricName(registry.recordProcessingErrors), new CumulativeSum());
 
         recordsSkipped = metricGroup.sensor("total-records-skipped");
-        recordsSkipped.add(metricGroup.metricName(registry.recordsSkipped), new Total());
+        recordsSkipped.add(metricGroup.metricName(registry.recordsSkipped), new CumulativeSum());
 
         retries = metricGroup.sensor("total-retries");
-        retries.add(metricGroup.metricName(registry.retries), new Total());
+        retries.add(metricGroup.metricName(registry.retries), new CumulativeSum());
 
         errorsLogged = metricGroup.sensor("total-errors-logged");
-        errorsLogged.add(metricGroup.metricName(registry.errorsLogged), new Total());
+        errorsLogged.add(metricGroup.metricName(registry.errorsLogged), new CumulativeSum());
 
         dlqProduceRequests = metricGroup.sensor("deadletterqueue-produce-requests");
-        dlqProduceRequests.add(metricGroup.metricName(registry.dlqProduceRequests), new Total());
+        dlqProduceRequests.add(metricGroup.metricName(registry.dlqProduceRequests), new CumulativeSum());
 
         dlqProduceFailures = metricGroup.sensor("deadletterqueue-produce-failures");
-        dlqProduceFailures.add(metricGroup.metricName(registry.dlqProduceFailures), new Total());
+        dlqProduceFailures.add(metricGroup.metricName(registry.dlqProduceFailures), new CumulativeSum());
 
         metricGroup.addValueMetric(registry.lastErrorTimestamp, now -> lastErrorTime);
     }
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/errors/ErrorReporter.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/errors/ErrorReporter.java
index 58336163fbf4a..5eaa42716b6b3 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/errors/ErrorReporter.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/errors/ErrorReporter.java
@@ -19,7 +19,7 @@
 /**
  * Report an error using the information contained in the {@link ProcessingContext}.
  */
-public interface ErrorReporter {
+public interface ErrorReporter extends AutoCloseable {
 
     /**
      * Report an error.
@@ -28,4 +28,6 @@ public interface ErrorReporter {
      */
     void report(ProcessingContext context);
 
+    @Override
+    default void close() { }
 }
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/errors/ProcessingContext.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/errors/ProcessingContext.java
index f826d74fe70cb..e7fb03185f85e 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/errors/ProcessingContext.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/errors/ProcessingContext.java
@@ -18,6 +18,7 @@
 
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.common.record.TimestampType;
+import org.apache.kafka.connect.errors.ConnectException;
 import org.apache.kafka.connect.source.SourceRecord;
 
 import java.util.Collection;
@@ -28,7 +29,7 @@
  * Contains all the metadata related to the currently evaluating operation. Only one instance of this class is meant
  * to exist per task in a JVM.
  */
-class ProcessingContext {
+class ProcessingContext implements AutoCloseable {
 
     private Collection reporters = Collections.emptyList();
 
@@ -216,4 +217,19 @@ public void reporters(Collection reporters) {
         this.reporters = reporters;
     }
 
+    @Override
+    public void close() {
+        ConnectException e = null;
+        for (ErrorReporter reporter : reporters) {
+            try {
+                reporter.close();
+            } catch (Throwable t) {
+                e = e != null ? e : new ConnectException("Failed to close all reporters");
+                e.addSuppressed(t);
+            }
+        }
+        if (e != null) {
+            throw e;
+        }
+    }
 }
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/errors/RetryWithToleranceOperator.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/errors/RetryWithToleranceOperator.java
index 2513514475582..4e627ef2e70c9 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/errors/RetryWithToleranceOperator.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/errors/RetryWithToleranceOperator.java
@@ -51,7 +51,7 @@
  * then it is wrapped into a ConnectException and rethrown to the caller.
  * 
      
  */
-public class RetryWithToleranceOperator {
+public class RetryWithToleranceOperator implements AutoCloseable {
 
     private static final Logger log = LoggerFactory.getLogger(RetryWithToleranceOperator.class);
 
@@ -270,4 +270,9 @@ public void consumerRecord(ConsumerRecord consumedMessage) {
     public boolean failed() {
         return this.context.failed();
     }
+
+    @Override
+    public void close() {
+        this.context.close();
+    }
 }
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/health/ConnectClusterDetailsImpl.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/health/ConnectClusterDetailsImpl.java
new file mode 100644
index 0000000000000..09f09bd7d383c
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/health/ConnectClusterDetailsImpl.java
@@ -0,0 +1,34 @@
+/*
+ * 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.health;
+
+import org.apache.kafka.connect.health.ConnectClusterDetails;
+
+public class ConnectClusterDetailsImpl implements ConnectClusterDetails {
+
+    private final String kafkaClusterId;
+
+    public ConnectClusterDetailsImpl(String kafkaClusterId) {
+        this.kafkaClusterId = kafkaClusterId;
+    }
+
+    @Override
+    public String kafkaClusterId() {
+        return kafkaClusterId;
+    }
+}
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/health/ConnectClusterStateImpl.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/health/ConnectClusterStateImpl.java
index ea93a72d5006e..6b7285df50b97 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/health/ConnectClusterStateImpl.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/health/ConnectClusterStateImpl.java
@@ -17,45 +17,55 @@
 
 package org.apache.kafka.connect.runtime.health;
 
+import org.apache.kafka.connect.errors.ConnectException;
+import org.apache.kafka.connect.health.ConnectClusterDetails;
 import org.apache.kafka.connect.health.ConnectClusterState;
 import org.apache.kafka.connect.health.ConnectorHealth;
 import org.apache.kafka.connect.health.ConnectorState;
 import org.apache.kafka.connect.health.ConnectorType;
 import org.apache.kafka.connect.health.TaskState;
-import org.apache.kafka.connect.runtime.HerderProvider;
+import org.apache.kafka.connect.runtime.Herder;
 import org.apache.kafka.connect.runtime.rest.entities.ConnectorStateInfo;
-import org.apache.kafka.connect.util.Callback;
+import org.apache.kafka.connect.util.FutureCallback;
 
-import java.util.ArrayList;
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 
 public class ConnectClusterStateImpl implements ConnectClusterState {
+    
+    private final long herderRequestTimeoutMs;
+    private final ConnectClusterDetails clusterDetails;
+    private final Herder herder;
 
-    private HerderProvider herderProvider;
-
-    public ConnectClusterStateImpl(HerderProvider herderProvider) {
-        this.herderProvider = herderProvider;
+    public ConnectClusterStateImpl(
+        long connectorsTimeoutMs,
+        ConnectClusterDetails clusterDetails,
+        Herder herder
+    ) {
+        this.herderRequestTimeoutMs = connectorsTimeoutMs;
+        this.clusterDetails = clusterDetails;
+        this.herder = herder;
     }
 
     @Override
     public Collection connectors() {
-        final Collection connectors = new ArrayList<>();
-        herderProvider.get().connectors(new Callback>() {
-            @Override
-            public void onCompletion(Throwable error, Collection result) {
-                connectors.addAll(result);
-            }
-        });
-        return connectors;
+        FutureCallback> connectorsCallback = new FutureCallback<>();
+        herder.connectors(connectorsCallback);
+        try {
+            return connectorsCallback.get(herderRequestTimeoutMs, TimeUnit.MILLISECONDS);
+        } catch (InterruptedException | ExecutionException | TimeoutException e) {
+            throw new ConnectException("Failed to retrieve list of connectors", e);
+        }
     }
 
     @Override
     public ConnectorHealth connectorHealth(String connName) {
-
-        ConnectorStateInfo state = herderProvider.get().connectorStatus(connName);
+        ConnectorStateInfo state = herder.connectorStatus(connName);
         ConnectorState connectorState = new ConnectorState(
             state.connector().state(),
             state.connector().workerId(),
@@ -71,6 +81,25 @@ public ConnectorHealth connectorHealth(String connName) {
         return connectorHealth;
     }
 
+    @Override
+    public Map connectorConfig(String connName) {
+        FutureCallback> connectorConfigCallback = new FutureCallback<>();
+        herder.connectorConfig(connName, connectorConfigCallback);
+        try {
+            return new HashMap<>(connectorConfigCallback.get(herderRequestTimeoutMs, TimeUnit.MILLISECONDS));
+        } catch (InterruptedException | ExecutionException | TimeoutException e) {
+            throw new ConnectException(
+                String.format("Failed to retrieve configuration for connector '%s'", connName),
+                e
+            );
+        }
+    }
+
+    @Override
+    public ConnectClusterDetails clusterDetails() {
+        return clusterDetails;
+    }
+
     private Map taskStates(List states) {
 
         Map taskStates = new HashMap<>();
@@ -78,7 +107,7 @@ private Map taskStates(List st
         for (ConnectorStateInfo.TaskState state : states) {
             taskStates.put(
                 state.id(),
-                new TaskState(state.id(), state.workerId(), state.state(), state.trace())
+                new TaskState(state.id(), state.state(), state.workerId(), state.trace())
             );
         }
         return taskStates;
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/isolation/DelegatingClassLoader.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/isolation/DelegatingClassLoader.java
index 460df39db3dea..8d964b4c94a6d 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/isolation/DelegatingClassLoader.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/isolation/DelegatingClassLoader.java
@@ -19,6 +19,7 @@
 import org.apache.kafka.common.config.provider.ConfigProvider;
 import org.apache.kafka.connect.components.Versioned;
 import org.apache.kafka.connect.connector.Connector;
+import org.apache.kafka.connect.connector.policy.ConnectorClientConfigOverridePolicy;
 import org.apache.kafka.connect.rest.ConnectRestExtension;
 import org.apache.kafka.connect.storage.Converter;
 import org.apache.kafka.connect.storage.HeaderConverter;
@@ -72,6 +73,7 @@ public class DelegatingClassLoader extends URLClassLoader {
     private final SortedSet> transformations;
     private final SortedSet> configProviders;
     private final SortedSet> restExtensions;
+    private final SortedSet> connectorClientConfigPolicies;
     private final List pluginPaths;
 
     private static final String MANIFEST_PREFIX = "META-INF/services/";
@@ -91,6 +93,7 @@ public DelegatingClassLoader(List pluginPaths, ClassLoader parent) {
         this.transformations = new TreeSet<>();
         this.configProviders = new TreeSet<>();
         this.restExtensions = new TreeSet<>();
+        this.connectorClientConfigPolicies = new TreeSet<>();
     }
 
     public DelegatingClassLoader(List pluginPaths) {
@@ -125,6 +128,29 @@ public Set> restExtensions() {
         return restExtensions;
     }
 
+    public Set> connectorClientConfigPolicies() {
+        return connectorClientConfigPolicies;
+    }
+
+    /**
+     * Retrieve the PluginClassLoader associated with a plugin class
+     * @param name The fully qualified class name of the plugin
+     * @return the PluginClassLoader that should be used to load this, or null if the plugin is not isolated.
+     */
+    public PluginClassLoader pluginClassLoader(String name) {
+        if (!PluginUtils.shouldLoadInIsolation(name)) {
+            return null;
+        }
+        SortedMap, ClassLoader> inner = pluginLoaders.get(name);
+        if (inner == null) {
+            return null;
+        }
+        ClassLoader pluginLoader = inner.get(inner.lastKey());
+        return pluginLoader instanceof PluginClassLoader
+               ? (PluginClassLoader) pluginLoader
+               : null;
+    }
+
     public ClassLoader connectorLoader(Connector connector) {
         return connectorLoader(connector.getClass().getName());
     }
@@ -134,8 +160,8 @@ public ClassLoader connectorLoader(String connectorClassOrAlias) {
         String fullName = aliases.containsKey(connectorClassOrAlias)
                           ? aliases.get(connectorClassOrAlias)
                           : connectorClassOrAlias;
-        SortedMap, ClassLoader> inner = pluginLoaders.get(fullName);
-        if (inner == null) {
+        PluginClassLoader classLoader = pluginClassLoader(fullName);
+        if (classLoader == null) {
             log.error(
                     "Plugin class loader for connector: '{}' was not found. Returning: {}",
                     connectorClassOrAlias,
@@ -143,7 +169,7 @@ public ClassLoader connectorLoader(String connectorClassOrAlias) {
             );
             return this;
         }
-        return inner.get(inner.lastKey());
+        return classLoader;
     }
 
     private static PluginClassLoader newPluginClassLoader(
@@ -249,6 +275,8 @@ private void scanUrlsAndAddPlugins(
             configProviders.addAll(plugins.configProviders());
             addPlugins(plugins.restExtensions(), loader);
             restExtensions.addAll(plugins.restExtensions());
+            addPlugins(plugins.connectorClientConfigPolicies(), loader);
+            connectorClientConfigPolicies.addAll(plugins.connectorClientConfigPolicies());
         }
 
         loadJdbcDrivers(loader);
@@ -304,7 +332,8 @@ private PluginScanResult scanPluginPath(
                 getPluginDesc(reflections, HeaderConverter.class, loader),
                 getPluginDesc(reflections, Transformation.class, loader),
                 getServiceLoaderPluginDesc(ConfigProvider.class, loader),
-                getServiceLoaderPluginDesc(ConnectRestExtension.class, loader)
+                getServiceLoaderPluginDesc(ConnectRestExtension.class, loader),
+                getServiceLoaderPluginDesc(ConnectorClientConfigOverridePolicy.class, loader)
         );
     }
 
@@ -328,10 +357,16 @@ private  Collection> getPluginDesc(
 
     @SuppressWarnings("unchecked")
     private  Collection> getServiceLoaderPluginDesc(Class klass, ClassLoader loader) {
-        ServiceLoader serviceLoader = ServiceLoader.load(klass, loader);
+        ClassLoader savedLoader = Plugins.compareAndSwapLoaders(loader);
         Collection> result = new ArrayList<>();
-        for (T pluginImpl : serviceLoader) {
-            result.add(new PluginDesc<>((Class) pluginImpl.getClass(), versionFor(pluginImpl), loader));
+        try {
+            ServiceLoader serviceLoader = ServiceLoader.load(klass, loader);
+            for (T pluginImpl : serviceLoader) {
+                result.add(new PluginDesc<>((Class) pluginImpl.getClass(),
+                    versionFor(pluginImpl), loader));
+            }
+        } finally {
+            Plugins.compareAndSwapLoaders(savedLoader);
         }
         return result;
     }
@@ -347,19 +382,11 @@ private static  String versionFor(Class pluginKlass) throws Ille
 
     @Override
     protected Class loadClass(String name, boolean resolve) throws ClassNotFoundException {
-        if (!PluginUtils.shouldLoadInIsolation(name)) {
-            // There are no paths in this classloader, will attempt to load with the parent.
-            return super.loadClass(name, resolve);
-        }
-
         String fullName = aliases.containsKey(name) ? aliases.get(name) : name;
-        SortedMap, ClassLoader> inner = pluginLoaders.get(fullName);
-        if (inner != null) {
-            ClassLoader pluginLoader = inner.get(inner.lastKey());
+        PluginClassLoader pluginLoader = pluginClassLoader(fullName);
+        if (pluginLoader != null) {
             log.trace("Retrieving loaded class '{}' from '{}'", fullName, pluginLoader);
-            return pluginLoader instanceof PluginClassLoader
-                   ? ((PluginClassLoader) pluginLoader).loadClass(fullName, resolve)
-                   : super.loadClass(fullName, resolve);
+            return pluginLoader.loadClass(fullName, resolve);
         }
 
         return super.loadClass(fullName, resolve);
@@ -371,6 +398,7 @@ private void addAllAliases() {
         addAliases(headerConverters);
         addAliases(transformations);
         addAliases(restExtensions);
+        addAliases(connectorClientConfigPolicies);
     }
 
     private  void addAliases(Collection> plugins) {
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/isolation/PluginScanResult.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/isolation/PluginScanResult.java
index ef077b3e7af19..e64a96c6f00a0 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/isolation/PluginScanResult.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/isolation/PluginScanResult.java
@@ -18,12 +18,15 @@
 
 import org.apache.kafka.common.config.provider.ConfigProvider;
 import org.apache.kafka.connect.connector.Connector;
+import org.apache.kafka.connect.connector.policy.ConnectorClientConfigOverridePolicy;
 import org.apache.kafka.connect.rest.ConnectRestExtension;
 import org.apache.kafka.connect.storage.Converter;
 import org.apache.kafka.connect.storage.HeaderConverter;
 import org.apache.kafka.connect.transforms.Transformation;
 
+import java.util.Arrays;
 import java.util.Collection;
+import java.util.List;
 
 public class PluginScanResult {
     private final Collection> connectors;
@@ -32,6 +35,9 @@ public class PluginScanResult {
     private final Collection> transformations;
     private final Collection> configProviders;
     private final Collection> restExtensions;
+    private final Collection> connectorClientConfigPolicies;
+
+    private final List allPlugins;
 
     public PluginScanResult(
             Collection> connectors,
@@ -39,7 +45,8 @@ public PluginScanResult(
             Collection> headerConverters,
             Collection> transformations,
             Collection> configProviders,
-            Collection> restExtensions
+            Collection> restExtensions,
+            Collection> connectorClientConfigPolicies
     ) {
         this.connectors = connectors;
         this.converters = converters;
@@ -47,6 +54,10 @@ public PluginScanResult(
         this.transformations = transformations;
         this.configProviders = configProviders;
         this.restExtensions = restExtensions;
+        this.connectorClientConfigPolicies = connectorClientConfigPolicies;
+        this.allPlugins =
+            Arrays.asList(connectors, converters, headerConverters, transformations, configProviders,
+                          connectorClientConfigPolicies);
     }
 
     public Collection> connectors() {
@@ -73,12 +84,15 @@ public Collection> restExtensions() {
         return restExtensions;
     }
 
+    public Collection> connectorClientConfigPolicies() {
+        return connectorClientConfigPolicies;
+    }
+
     public boolean isEmpty() {
-        return connectors().isEmpty()
-               && converters().isEmpty()
-               && headerConverters().isEmpty()
-               && transformations().isEmpty()
-               && configProviders().isEmpty()
-               && restExtensions().isEmpty();
+        boolean isEmpty = true;
+        for (Collection plugins : allPlugins) {
+            isEmpty = isEmpty && plugins.isEmpty();
+        }
+        return isEmpty;
     }
 }
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/isolation/PluginType.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/isolation/PluginType.java
index 2833b4c4ba0bf..8b42f59ef9368 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/isolation/PluginType.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/isolation/PluginType.java
@@ -18,6 +18,7 @@
 
 import org.apache.kafka.common.config.provider.ConfigProvider;
 import org.apache.kafka.connect.connector.Connector;
+import org.apache.kafka.connect.connector.policy.ConnectorClientConfigOverridePolicy;
 import org.apache.kafka.connect.rest.ConnectRestExtension;
 import org.apache.kafka.connect.sink.SinkConnector;
 import org.apache.kafka.connect.source.SourceConnector;
@@ -34,6 +35,7 @@ public enum PluginType {
     TRANSFORMATION(Transformation.class),
     CONFIGPROVIDER(ConfigProvider.class),
     REST_EXTENSION(ConnectRestExtension.class),
+    CONNECTOR_CLIENT_CONFIG_OVERRIDE_POLICY(ConnectorClientConfigOverridePolicy.class),
     UNKNOWN(Object.class);
 
     private Class klass;
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/isolation/PluginUtils.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/isolation/PluginUtils.java
index 8d2a3cedc1764..acb26fb9ebcc3 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/isolation/PluginUtils.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/isolation/PluginUtils.java
@@ -123,14 +123,21 @@ public class PluginUtils {
             + "|org\\.slf4j"
             + ")\\..*$");
 
+    // If the base interface or class that will be used to identify Connect plugins resides within
+    // the same java package as the plugins that need to be loaded in isolation (and thus are
+    // added to the WHITELIST), then this base interface or class needs to be excluded in the
+    // regular expression pattern
     private static final Pattern WHITELIST = Pattern.compile("^org\\.apache\\.kafka\\.(?:connect\\.(?:"
             + "transforms\\.(?!Transformation$).*"
             + "|json\\..*"
             + "|file\\..*"
+            + "|mirror\\..*"
+            + "|mirror-client\\..*"
             + "|converters\\..*"
             + "|storage\\.StringConverter"
             + "|storage\\.SimpleHeaderConverter"
             + "|rest\\.basic\\.auth\\.extension\\.BasicAuthSecurityRestExtension"
+            + "|connector\\.policy\\.(?!ConnectorClientConfig(?:OverridePolicy|Request(?:\\$ClientType)?)$).*"
             + ")"
             + "|common\\.config\\.provider\\.(?!ConfigProvider$).*"
             + ")$");
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/isolation/Plugins.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/isolation/Plugins.java
index e7cb16db1bc83..d068a03ceecde 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/isolation/Plugins.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/isolation/Plugins.java
@@ -17,7 +17,6 @@
 package org.apache.kafka.connect.runtime.isolation;
 
 import org.apache.kafka.common.Configurable;
-import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.config.AbstractConfig;
 import org.apache.kafka.common.config.provider.ConfigProvider;
 import org.apache.kafka.common.utils.Utils;
@@ -72,13 +71,37 @@ private static  String pluginNames(Collection> plugins) {
     }
 
     protected static  T newPlugin(Class klass) {
+        // KAFKA-8340: The thread classloader is used during static initialization and must be
+        // set to the plugin's classloader during instantiation
+        ClassLoader savedLoader = compareAndSwapLoaders(klass.getClassLoader());
         try {
             return Utils.newInstance(klass);
         } catch (Throwable t) {
             throw new ConnectException("Instantiation error", t);
+        } finally {
+            compareAndSwapLoaders(savedLoader);
         }
     }
 
+    @SuppressWarnings("unchecked")
+    protected  Class pluginClassFromConfig(
+            AbstractConfig config,
+            String propertyName,
+            Class pluginClass,
+            Collection> plugins
+    ) {
+        Class klass = config.getClass(propertyName);
+        if (pluginClass.isAssignableFrom(klass)) {
+            return (Class) klass;
+        }
+        throw new ConnectException(
+            "Failed to find any class that implements " + pluginClass.getSimpleName()
+                + " for the config "
+                + propertyName + ", available classes are: "
+                + pluginNames(plugins)
+        );
+    }
+
     @SuppressWarnings("unchecked")
     protected static  Class pluginClass(
             DelegatingClassLoader loader,
@@ -149,6 +172,11 @@ public Set> configProviders() {
     }
 
     public Connector newConnector(String connectorClassOrAlias) {
+        Class klass = connectorClass(connectorClassOrAlias);
+        return newPlugin(klass);
+    }
+
+    public Class connectorClass(String connectorClassOrAlias) {
         Class klass;
         try {
             klass = pluginClass(
@@ -188,7 +216,7 @@ public Connector newConnector(String connectorClassOrAlias) {
             PluginDesc entry = matches.get(0);
             klass = entry.pluginClass();
         }
-        return newPlugin(klass);
+        return klass;
     }
 
     public Task newTask(Class taskClass) {
@@ -210,18 +238,17 @@ public Converter newConverter(AbstractConfig config, String classPropertyName, C
             // it does not represent an internal converter (which has a default available)
             return null;
         }
-        Converter plugin = null;
+        Class klass = null;
         switch (classLoaderUsage) {
             case CURRENT_CLASSLOADER:
                 // Attempt to load first with the current classloader, and plugins as a fallback.
                 // Note: we can't use config.getConfiguredInstance because Converter doesn't implement Configurable, and even if it did
                 // we have to remove the property prefixes before calling config(...) and we still always want to call Converter.config.
-                plugin = getInstance(config, classPropertyName, Converter.class);
+                klass = pluginClassFromConfig(config, classPropertyName, Converter.class, delegatingLoader.converters());
                 break;
             case PLUGINS:
                 // Attempt to load with the plugin class loader, which uses the current classloader as a fallback
                 String converterClassOrAlias = config.getClass(classPropertyName).getName();
-                Class klass;
                 try {
                     klass = pluginClass(delegatingLoader, converterClassOrAlias, Converter.class);
                 } catch (ClassNotFoundException e) {
@@ -231,11 +258,10 @@ public Converter newConverter(AbstractConfig config, String classPropertyName, C
                             + pluginNames(delegatingLoader.converters())
                     );
                 }
-                plugin = newPlugin(klass);
                 break;
         }
-        if (plugin == null) {
-            throw new ConnectException("Unable to instantiate the Converter specified in '" + classPropertyName + "'");
+        if (klass == null) {
+            throw new ConnectException("Unable to initialize the Converter specified in '" + classPropertyName + "'");
         }
 
         // Determine whether this is a key or value converter based upon the supplied property name ...
@@ -246,13 +272,13 @@ public Converter newConverter(AbstractConfig config, String classPropertyName, C
         // Configure the Converter using only the old configuration mechanism ...
         String configPrefix = classPropertyName + ".";
         Map converterConfig = config.originalsWithPrefix(configPrefix);
-        log.debug("Configuring the {} converter with configuration:{}{}",
-                  isKeyConverter ? "key" : "value", System.lineSeparator(), converterConfig);
+        log.debug("Configuring the {} converter with configuration keys:{}{}",
+                  isKeyConverter ? "key" : "value", System.lineSeparator(), converterConfig.keySet());
 
         // Have to override schemas.enable from true to false for internal JSON converters
         // Don't have to warn the user about anything since all deprecation warnings take place in the
         // WorkerConfig class
-        if (plugin instanceof JsonConverter && isInternalConverter(classPropertyName)) {
+        if (JsonConverter.class.isAssignableFrom(klass) && isInternalConverter(classPropertyName)) {
             // If they haven't explicitly specified values for internal.key.converter.schemas.enable
             // or internal.value.converter.schemas.enable, we can safely default them to false
             if (!converterConfig.containsKey(JsonConverterConfig.SCHEMAS_ENABLE_CONFIG)) {
@@ -260,7 +286,14 @@ public Converter newConverter(AbstractConfig config, String classPropertyName, C
             }
         }
 
-        plugin.configure(converterConfig, isKeyConverter);
+        Converter plugin;
+        ClassLoader savedLoader = compareAndSwapLoaders(klass.getClassLoader());
+        try {
+            plugin = newPlugin(klass);
+            plugin.configure(converterConfig, isKeyConverter);
+        } finally {
+            compareAndSwapLoaders(savedLoader);
+        }
         return plugin;
     }
 
@@ -275,7 +308,7 @@ public Converter newConverter(AbstractConfig config, String classPropertyName, C
      * @throws ConnectException if the {@link HeaderConverter} implementation class could not be found
      */
     public HeaderConverter newHeaderConverter(AbstractConfig config, String classPropertyName, ClassLoaderUsage classLoaderUsage) {
-        HeaderConverter plugin = null;
+        Class klass = null;
         switch (classLoaderUsage) {
             case CURRENT_CLASSLOADER:
                 if (!config.originals().containsKey(classPropertyName)) {
@@ -285,13 +318,12 @@ public HeaderConverter newHeaderConverter(AbstractConfig config, String classPro
                 // Attempt to load first with the current classloader, and plugins as a fallback.
                 // Note: we can't use config.getConfiguredInstance because we have to remove the property prefixes
                 // before calling config(...)
-                plugin = getInstance(config, classPropertyName, HeaderConverter.class);
+                klass = pluginClassFromConfig(config, classPropertyName, HeaderConverter.class, delegatingLoader.headerConverters());
                 break;
             case PLUGINS:
                 // Attempt to load with the plugin class loader, which uses the current classloader as a fallback.
                 // Note that there will always be at least a default header converter for the worker
                 String converterClassOrAlias = config.getClass(classPropertyName).getName();
-                Class klass;
                 try {
                     klass = pluginClass(
                             delegatingLoader,
@@ -306,17 +338,24 @@ public HeaderConverter newHeaderConverter(AbstractConfig config, String classPro
                                     + pluginNames(delegatingLoader.headerConverters())
                     );
                 }
-                plugin = newPlugin(klass);
         }
-        if (plugin == null) {
-            throw new ConnectException("Unable to instantiate the Converter specified in '" + classPropertyName + "'");
+        if (klass == null) {
+            throw new ConnectException("Unable to initialize the HeaderConverter specified in '" + classPropertyName + "'");
         }
 
         String configPrefix = classPropertyName + ".";
         Map converterConfig = config.originalsWithPrefix(configPrefix);
         converterConfig.put(ConverterConfig.TYPE_CONFIG, ConverterType.HEADER.getName());
-        log.debug("Configuring the header converter with configuration:{}{}", System.lineSeparator(), converterConfig);
-        plugin.configure(converterConfig);
+        log.debug("Configuring the header converter with configuration keys:{}{}", System.lineSeparator(), converterConfig.keySet());
+
+        HeaderConverter plugin;
+        ClassLoader savedLoader = compareAndSwapLoaders(klass.getClassLoader());
+        try {
+            plugin = newPlugin(klass);
+            plugin.configure(converterConfig);
+        } finally {
+            compareAndSwapLoaders(savedLoader);
+        }
         return plugin;
     }
 
@@ -327,16 +366,15 @@ public ConfigProvider newConfigProvider(AbstractConfig config, String providerPr
             // This configuration does not define the config provider via the specified property name
             return null;
         }
-        ConfigProvider plugin = null;
+        Class klass = null;
         switch (classLoaderUsage) {
             case CURRENT_CLASSLOADER:
                 // Attempt to load first with the current classloader, and plugins as a fallback.
-                plugin = getInstance(config, classPropertyName, ConfigProvider.class);
+                klass = pluginClassFromConfig(config, classPropertyName, ConfigProvider.class, delegatingLoader.configProviders());
                 break;
             case PLUGINS:
                 // Attempt to load with the plugin class loader, which uses the current classloader as a fallback
                 String configProviderClassOrAlias = originalConfig.get(classPropertyName);
-                Class klass;
                 try {
                     klass = pluginClass(delegatingLoader, configProviderClassOrAlias, ConfigProvider.class);
                 } catch (ClassNotFoundException e) {
@@ -346,17 +384,24 @@ public ConfigProvider newConfigProvider(AbstractConfig config, String providerPr
                                     + pluginNames(delegatingLoader.configProviders())
                     );
                 }
-                plugin = newPlugin(klass);
                 break;
         }
-        if (plugin == null) {
-            throw new ConnectException("Unable to instantiate the ConfigProvider specified in '" + classPropertyName + "'");
+        if (klass == null) {
+            throw new ConnectException("Unable to initialize the ConfigProvider specified in '" + classPropertyName + "'");
         }
 
         // Configure the ConfigProvider
         String configPrefix = providerPrefix + ".param.";
         Map configProviderConfig = config.originalsWithPrefix(configPrefix);
-        plugin.configure(configProviderConfig);
+
+        ConfigProvider plugin;
+        ClassLoader savedLoader = compareAndSwapLoaders(klass.getClassLoader());
+        try {
+            plugin = newPlugin(klass);
+            plugin.configure(configProviderConfig);
+        } finally {
+            compareAndSwapLoaders(savedLoader);
+        }
         return plugin;
     }
 
@@ -390,42 +435,25 @@ public  T newPlugin(String klassName, AbstractConfig config, Class pluginK
                                        + "name matches %s", pluginKlass, klassName);
             throw new ConnectException(msg);
         }
-        plugin = newPlugin(klass);
-        if (plugin == null) {
-            throw new ConnectException("Unable to instantiate '" + klassName + "'");
-        }
-        if (plugin instanceof Versioned) {
-            Versioned versionedPlugin = (Versioned) plugin;
-            if (versionedPlugin.version() == null || versionedPlugin.version().trim().isEmpty()) {
-                throw new ConnectException("Version not defined for '" + klassName + "'");
+        ClassLoader savedLoader = compareAndSwapLoaders(klass.getClassLoader());
+        try {
+            plugin = newPlugin(klass);
+            if (plugin instanceof Versioned) {
+                Versioned versionedPlugin = (Versioned) plugin;
+                if (versionedPlugin.version() == null || versionedPlugin.version().trim()
+                    .isEmpty()) {
+                    throw new ConnectException("Version not defined for '" + klassName + "'");
+                }
             }
-        }
-        if (plugin instanceof Configurable) {
-            ((Configurable) plugin).configure(config.originals());
+            if (plugin instanceof Configurable) {
+                ((Configurable) plugin).configure(config.originals());
+            }
+        } finally {
+            compareAndSwapLoaders(savedLoader);
         }
         return plugin;
     }
 
-    /**
-     * Get an instance of the give class specified by the given configuration key.
-     *
-     * @param key The configuration key for the class
-     * @param t The interface the class should implement
-     * @return A instance of the class
-     */
-    private  T getInstance(AbstractConfig config, String key, Class t) {
-        Class c = config.getClass(key);
-        if (c == null) {
-            return null;
-        }
-        // Instantiate the class, but we don't know if the class extends the supplied type
-        Object o = Utils.newInstance(c);
-        if (!t.isInstance(o)) {
-            throw new KafkaException(c.getName() + " is not an instance of " + t.getName());
-        }
-        return t.cast(o);
-    }
-
     public > Transformation newTranformations(
             String transformationClassOrAlias
     ) {
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/InternalRequestSignature.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/InternalRequestSignature.java
new file mode 100644
index 0000000000000..d59425b13f6e2
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/InternalRequestSignature.java
@@ -0,0 +1,148 @@
+/*
+ * 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.rest;
+
+import org.apache.kafka.connect.errors.ConnectException;
+import org.apache.kafka.connect.runtime.rest.errors.BadRequestException;
+import org.eclipse.jetty.client.api.Request;
+
+import javax.crypto.Mac;
+import javax.crypto.SecretKey;
+import javax.ws.rs.core.HttpHeaders;
+import java.security.InvalidKeyException;
+import java.security.NoSuchAlgorithmException;
+import java.util.Arrays;
+import java.util.Base64;
+import java.util.Objects;
+
+public class InternalRequestSignature {
+
+    public static final String SIGNATURE_HEADER = "X-Connect-Authorization";
+    public static final String SIGNATURE_ALGORITHM_HEADER = "X-Connect-Request-Signature-Algorithm";
+
+    private final byte[] requestBody;
+    private final Mac mac;
+    private final byte[] requestSignature;
+
+    /**
+     * Add a signature to a request.
+     * @param key the key to sign the request with; may not be null
+     * @param requestBody the body of the request; may not be null
+     * @param signatureAlgorithm the algorithm to use to sign the request; may not be null
+     * @param request the request to add the signature to; may not be null
+     */
+    public static void addToRequest(SecretKey key, byte[] requestBody, String signatureAlgorithm, Request request) {
+        Mac mac;
+        try {
+            mac = mac(signatureAlgorithm);
+        }  catch (NoSuchAlgorithmException e) {
+            throw new ConnectException(e);
+        }
+        byte[] requestSignature = sign(mac, key, requestBody);
+        request.header(InternalRequestSignature.SIGNATURE_HEADER, Base64.getEncoder().encodeToString(requestSignature))
+               .header(InternalRequestSignature.SIGNATURE_ALGORITHM_HEADER, signatureAlgorithm);
+    }
+
+    /**
+     * Extract a signature from a request.
+     * @param requestBody the body of the request; may not be null
+     * @param headers the headers for the request; may be null
+     * @return the signature extracted from the request, or null if one or more request signature
+     * headers was not present
+     */
+    public static InternalRequestSignature fromHeaders(byte[] requestBody, HttpHeaders headers) {
+        if (headers == null) {
+            return null;
+        }
+
+        String signatureAlgorithm = headers.getHeaderString(SIGNATURE_ALGORITHM_HEADER);
+        String encodedSignature = headers.getHeaderString(SIGNATURE_HEADER);
+        if (signatureAlgorithm == null || encodedSignature == null) {
+            return null;
+        }
+
+        Mac mac;
+        try {
+            mac = mac(signatureAlgorithm);
+        } catch (NoSuchAlgorithmException e) {
+            throw new BadRequestException(e.getMessage());
+        }
+
+        byte[] decodedSignature;
+        try {
+            decodedSignature = Base64.getDecoder().decode(encodedSignature);
+        } catch (IllegalArgumentException e) {
+            throw new BadRequestException(e.getMessage());
+        }
+
+        return new InternalRequestSignature(
+            requestBody,
+            mac,
+            decodedSignature
+        );
+    }
+
+    // Public for testing
+    public InternalRequestSignature(byte[] requestBody, Mac mac, byte[] requestSignature) {
+        this.requestBody = requestBody;
+        this.mac = mac;
+        this.requestSignature = requestSignature;
+    }
+
+    public String keyAlgorithm() {
+        return mac.getAlgorithm();
+    }
+
+    public boolean isValid(SecretKey key) {
+        return Arrays.equals(sign(mac, key, requestBody), requestSignature);
+    }
+
+    private static Mac mac(String signatureAlgorithm) throws NoSuchAlgorithmException {
+        return Mac.getInstance(signatureAlgorithm);
+    }
+
+    private static byte[] sign(Mac mac, SecretKey key, byte[] requestBody) {
+        try {
+            mac.init(key);
+        } catch (InvalidKeyException e) {
+            throw new ConnectException(e);
+        }
+        return mac.doFinal(requestBody);
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o)
+            return true;
+        if (o == null || getClass() != o.getClass())
+            return false;
+        InternalRequestSignature that = (InternalRequestSignature) o;
+        return Arrays.equals(requestBody, that.requestBody)
+            && mac.getAlgorithm().equals(that.mac.getAlgorithm())
+            && mac.getMacLength() == that.mac.getMacLength()
+            && mac.getProvider().equals(that.mac.getProvider())
+            && Arrays.equals(requestSignature, that.requestSignature);
+    }
+
+    @Override
+    public int hashCode() {
+        int result = Objects.hash(mac);
+        result = 31 * result + Arrays.hashCode(requestBody);
+        result = 31 * result + Arrays.hashCode(requestSignature);
+        return result;
+    }
+}
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 15e8418a30c90..b38125971656c 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
@@ -19,6 +19,10 @@
 
 import com.fasterxml.jackson.core.type.TypeReference;
 import com.fasterxml.jackson.databind.ObjectMapper;
+
+import javax.crypto.SecretKey;
+import javax.ws.rs.core.HttpHeaders;
+
 import org.apache.kafka.connect.runtime.WorkerConfig;
 import org.apache.kafka.connect.runtime.rest.entities.ErrorMessage;
 import org.apache.kafka.connect.runtime.rest.errors.ConnectRestException;
@@ -50,17 +54,39 @@ public class RestClient {
      *
      * @param url             HTTP connection will be established with this url.
      * @param method          HTTP method ("GET", "POST", "PUT", etc.)
+     * @param headers         HTTP headers from REST endpoint
      * @param requestBodyData Object to serialize as JSON and send in the request body.
      * @param responseFormat  Expected format of the response to the HTTP request.
      * @param              The type of the deserialized response to the HTTP request.
      * @return The deserialized response to the HTTP request, or null if no data is expected.
      */
-    public static  HttpResponse httpRequest(String url, String method, Object requestBodyData,
+    public static  HttpResponse httpRequest(String url, String method, HttpHeaders headers, Object requestBodyData,
                                                   TypeReference responseFormat, WorkerConfig config) {
+        return httpRequest(url, method, headers, requestBodyData, responseFormat, config, null, null);
+    }
+
+    /**
+     * Sends HTTP request to remote REST server
+     *
+     * @param url                       HTTP connection will be established with this url.
+     * @param method                    HTTP method ("GET", "POST", "PUT", etc.)
+     * @param headers                   HTTP headers from REST endpoint
+     * @param requestBodyData           Object to serialize as JSON and send in the request body.
+     * @param responseFormat            Expected format of the response to the HTTP request.
+     * @param                        The type of the deserialized response to the HTTP request.
+     * @param sessionKey                The key to sign the request with (intended for internal requests only);
+     *                                  may be null if the request doesn't need to be signed
+     * @param requestSignatureAlgorithm The algorithm to sign the request with (intended for internal requests only);
+     *                                  may be null if the request doesn't need to be signed
+     * @return The deserialized response to the HTTP request, or null if no data is expected.
+     */
+    public static  HttpResponse httpRequest(String url, String method, HttpHeaders headers, Object requestBodyData,
+                                                  TypeReference responseFormat, WorkerConfig config,
+                                                  SecretKey sessionKey, String requestSignatureAlgorithm) {
         HttpClient client;
 
         if (url.startsWith("https://")) {
-            client = new HttpClient(SSLUtils.createSslContextFactory(config, true));
+            client = new HttpClient(SSLUtils.createClientSideSslContextFactory(config));
         } else {
             client = new HttpClient();
         }
@@ -82,8 +108,18 @@ public static  HttpResponse httpRequest(String url, String method, Object
             req.method(method);
             req.accept("application/json");
             req.agent("kafka-connect");
+            addHeadersToRequest(headers, req);
+
             if (serializedBody != null) {
                 req.content(new StringContentProvider(serializedBody, StandardCharsets.UTF_8), "application/json");
+                if (sessionKey != null && requestSignatureAlgorithm != null) {
+                    InternalRequestSignature.addToRequest(
+                        sessionKey,
+                        serializedBody.getBytes(StandardCharsets.UTF_8),
+                        requestSignatureAlgorithm,
+                        req
+                    );
+                }
             }
 
             ContentResponse res = req.send();
@@ -107,12 +143,26 @@ public static  HttpResponse httpRequest(String url, String method, Object
             log.error("IO error forwarding REST request: ", e);
             throw new ConnectRestException(Response.Status.INTERNAL_SERVER_ERROR, "IO Error trying to forward REST request: " + e.getMessage(), e);
         } finally {
-            if (client != null)
-                try {
-                    client.stop();
-                } catch (Exception e) {
-                    log.error("Failed to stop HTTP client", e);
-                }
+            try {
+                client.stop();
+            } catch (Exception e) {
+                log.error("Failed to stop HTTP client", e);
+            }
+        }
+    }
+
+
+    /**
+     * Extract headers from REST call and add to client request
+     * @param headers         Headers from REST endpoint
+     * @param req             The client request to modify
+     */
+    private static void addHeadersToRequest(HttpHeaders headers, Request req) {
+        if (headers != null) {
+            String credentialAuthorization = headers.getHeaderString(HttpHeaders.AUTHORIZATION);
+            if (credentialAuthorization != null) {
+                req.header(HttpHeaders.AUTHORIZATION, credentialAuthorization);
+            }
         }
     }
 
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 5cc31cd5cca3c..02b4677521fb1 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
@@ -19,24 +19,27 @@
 import com.fasterxml.jackson.jaxrs.json.JacksonJsonProvider;
 import org.apache.kafka.common.config.ConfigException;
 import org.apache.kafka.connect.errors.ConnectException;
+import org.apache.kafka.connect.health.ConnectClusterDetails;
 import org.apache.kafka.connect.rest.ConnectRestExtension;
 import org.apache.kafka.connect.rest.ConnectRestExtensionContext;
-import org.apache.kafka.connect.runtime.HerderProvider;
+import org.apache.kafka.connect.runtime.Herder;
 import org.apache.kafka.connect.runtime.WorkerConfig;
+import org.apache.kafka.connect.runtime.health.ConnectClusterDetailsImpl;
 import org.apache.kafka.connect.runtime.health.ConnectClusterStateImpl;
-import org.apache.kafka.connect.runtime.isolation.Plugins;
 import org.apache.kafka.connect.runtime.rest.errors.ConnectExceptionMapper;
 import org.apache.kafka.connect.runtime.rest.resources.ConnectorPluginsResource;
 import org.apache.kafka.connect.runtime.rest.resources.ConnectorsResource;
+import org.apache.kafka.connect.runtime.rest.resources.LoggingResource;
 import org.apache.kafka.connect.runtime.rest.resources.RootResource;
 import org.apache.kafka.connect.runtime.rest.util.SSLUtils;
 import org.eclipse.jetty.server.Connector;
+import org.eclipse.jetty.server.CustomRequestLog;
 import org.eclipse.jetty.server.Handler;
 import org.eclipse.jetty.server.Server;
 import org.eclipse.jetty.server.ServerConnector;
-import org.eclipse.jetty.server.Slf4jRequestLog;
+import org.eclipse.jetty.server.Slf4jRequestLogWriter;
+import org.eclipse.jetty.server.handler.ContextHandlerCollection;
 import org.eclipse.jetty.server.handler.DefaultHandler;
-import org.eclipse.jetty.server.handler.HandlerCollection;
 import org.eclipse.jetty.server.handler.RequestLogHandler;
 import org.eclipse.jetty.server.handler.StatisticsHandler;
 import org.eclipse.jetty.servlet.FilterHolder;
@@ -62,12 +65,17 @@
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
+import static org.apache.kafka.connect.runtime.WorkerConfig.ADMIN_LISTENERS_HTTPS_CONFIGS_PREFIX;
+
 /**
  * Embedded server for the REST API that provides the control plane for Kafka Connect workers.
  */
 public class RestServer {
     private static final Logger log = LoggerFactory.getLogger(RestServer.class);
 
+    // Used to distinguish between Admin connectors and regular REST API connectors when binding admin handlers
+    private static final String ADMIN_SERVER_CONNECTOR_NAME = "Admin";
+
     private static final Pattern LISTENER_PATTERN = Pattern.compile("^(.*)://\\[?([0-9a-zA-Z\\-%._:]*)\\]?:(-?[0-9]+)");
     private static final long GRACEFUL_SHUTDOWN_TIMEOUT_MS = 60 * 1000;
 
@@ -75,6 +83,7 @@ public class RestServer {
     private static final String PROTOCOL_HTTPS = "https";
 
     private final WorkerConfig config;
+    private ContextHandlerCollection handlers;
     private Server jettyServer;
 
     private List connectRestExtensions = Collections.emptyList();
@@ -86,10 +95,12 @@ public RestServer(WorkerConfig config) {
         this.config = config;
 
         List listeners = parseListeners();
+        List adminListeners = config.getList(WorkerConfig.ADMIN_LISTENERS_CONFIG);
 
         jettyServer = new Server();
+        handlers = new ContextHandlerCollection();
 
-        createConnectors(listeners);
+        createConnectors(listeners, adminListeners);
     }
 
     @SuppressWarnings("deprecation")
@@ -110,24 +121,39 @@ List parseListeners() {
     /**
      * Adds Jetty connector for each configured listener
      */
-    public void createConnectors(List listeners) {
+    public void createConnectors(List listeners, List adminListeners) {
         List connectors = new ArrayList<>();
 
         for (String listener : listeners) {
             if (!listener.isEmpty()) {
                 Connector connector = createConnector(listener);
                 connectors.add(connector);
-                log.info("Added connector for " + listener);
+                log.info("Added connector for {}", listener);
             }
         }
 
         jettyServer.setConnectors(connectors.toArray(new Connector[connectors.size()]));
+
+        if (adminListeners != null && !adminListeners.isEmpty()) {
+            for (String adminListener : adminListeners) {
+                Connector conn = createConnector(adminListener, true);
+                jettyServer.addConnector(conn);
+                log.info("Added admin connector for {}", adminListener);
+            }
+        }
     }
 
     /**
-     * Creates Jetty connector according to configuration
+     * Creates regular (non-admin) Jetty connector according to configuration
      */
     public Connector createConnector(String listener) {
+        return createConnector(listener, false);
+    }
+
+    /**
+     * Creates Jetty connector according to configuration
+     */
+    public Connector createConnector(String listener, boolean isAdmin) {
         Matcher listenerMatcher = LISTENER_PATTERN.matcher(listener);
 
         if (!listenerMatcher.matches())
@@ -144,12 +170,25 @@ public Connector createConnector(String listener) {
         ServerConnector connector;
 
         if (PROTOCOL_HTTPS.equals(protocol)) {
-            SslContextFactory ssl = SSLUtils.createSslContextFactory(config);
+            SslContextFactory ssl;
+            if (isAdmin) {
+                ssl = SSLUtils.createServerSideSslContextFactory(config, ADMIN_LISTENERS_HTTPS_CONFIGS_PREFIX);
+            } else {
+                ssl = SSLUtils.createServerSideSslContextFactory(config);
+            }
             connector = new ServerConnector(jettyServer, ssl);
-            connector.setName(String.format("%s_%s%d", PROTOCOL_HTTPS, hostname, port));
+            if (!isAdmin) {
+                connector.setName(String.format("%s_%s%d", PROTOCOL_HTTPS, hostname, port));
+            }
         } else {
             connector = new ServerConnector(jettyServer);
-            connector.setName(String.format("%s_%s%d", PROTOCOL_HTTP, hostname, port));
+            if (!isAdmin) {
+                connector.setName(String.format("%s_%s%d", PROTOCOL_HTTP, hostname, port));
+            }
+        }
+
+        if (isAdmin) {
+            connector.setName(ADMIN_SERVER_CONNECTOR_NAME);
         }
 
         if (!hostname.isEmpty())
@@ -160,27 +199,79 @@ public Connector createConnector(String listener) {
         return connector;
     }
 
-    public void start(HerderProvider herderProvider, Plugins plugins) {
-        log.info("Starting REST server");
+    public void initializeServer() {
+        log.info("Initializing REST server");
+
+        /* Needed for graceful shutdown as per `setStopTimeout` documentation */
+        StatisticsHandler statsHandler = new StatisticsHandler();
+        statsHandler.setHandler(handlers);
+        jettyServer.setHandler(statsHandler);
+        jettyServer.setStopTimeout(GRACEFUL_SHUTDOWN_TIMEOUT_MS);
+        jettyServer.setStopAtShutdown(true);
+
+        try {
+            jettyServer.start();
+        } catch (Exception e) {
+            throw new ConnectException("Unable to initialize REST server", e);
+        }
+
+        log.info("REST server listening at " + jettyServer.getURI() + ", advertising URL " + advertisedUrl());
+        log.info("REST admin endpoints at " + adminUrl());
+    }
+
+    public void initializeResources(Herder herder) {
+        log.info("Initializing REST resources");
 
         ResourceConfig resourceConfig = new ResourceConfig();
         resourceConfig.register(new JacksonJsonProvider());
 
-        resourceConfig.register(new RootResource(herderProvider));
-        resourceConfig.register(new ConnectorsResource(herderProvider, config));
-        resourceConfig.register(new ConnectorPluginsResource(herderProvider));
+        resourceConfig.register(new RootResource(herder));
+        resourceConfig.register(new ConnectorsResource(herder, config));
+        resourceConfig.register(new ConnectorPluginsResource(herder));
 
         resourceConfig.register(ConnectExceptionMapper.class);
         resourceConfig.property(ServerProperties.WADL_FEATURE_DISABLE, true);
 
-        registerRestExtensions(herderProvider, plugins, resourceConfig);
+        registerRestExtensions(herder, resourceConfig);
+
+        List adminListeners = config.getList(WorkerConfig.ADMIN_LISTENERS_CONFIG);
+        ResourceConfig adminResourceConfig;
+        if (adminListeners == null) {
+            log.info("Adding admin resources to main listener");
+            adminResourceConfig = resourceConfig;
+            adminResourceConfig.register(new LoggingResource());
+        } else if (adminListeners.size() > 0) {
+            // TODO: we need to check if these listeners are same as 'listeners'
+            // TODO: the following code assumes that they are different
+            log.info("Adding admin resources to admin listener");
+            adminResourceConfig = new ResourceConfig();
+            adminResourceConfig.register(new JacksonJsonProvider());
+            adminResourceConfig.register(new LoggingResource());
+            adminResourceConfig.register(ConnectExceptionMapper.class);
+        } else {
+            log.info("Skipping adding admin resources");
+            // set up adminResource but add no handlers to it
+            adminResourceConfig = resourceConfig;
+        }
 
         ServletContainer servletContainer = new ServletContainer(resourceConfig);
         ServletHolder servletHolder = new ServletHolder(servletContainer);
+        List contextHandlers = new ArrayList<>();
 
         ServletContextHandler context = new ServletContextHandler(ServletContextHandler.SESSIONS);
         context.setContextPath("/");
         context.addServlet(servletHolder, "/*");
+        contextHandlers.add(context);
+
+        ServletContextHandler adminContext = null;
+        if (adminResourceConfig != resourceConfig) {
+            adminContext = new ServletContextHandler(ServletContextHandler.SESSIONS);
+            ServletHolder adminServletHolder = new ServletHolder(new ServletContainer(adminResourceConfig));
+            adminContext.setContextPath("/");
+            adminContext.addServlet(adminServletHolder, "/*");
+            adminContext.setVirtualHosts(new String[]{"@" + ADMIN_SERVER_CONNECTOR_NAME});
+            contextHandlers.add(adminContext);
+        }
 
         String allowedOrigins = config.getString(WorkerConfig.ACCESS_CONTROL_ALLOW_ORIGIN_CONFIG);
         if (allowedOrigins != null && !allowedOrigins.trim().isEmpty()) {
@@ -195,28 +286,31 @@ public void start(HerderProvider herderProvider, Plugins plugins) {
         }
 
         RequestLogHandler requestLogHandler = new RequestLogHandler();
-        Slf4jRequestLog requestLog = new Slf4jRequestLog();
-        requestLog.setLoggerName(RestServer.class.getCanonicalName());
-        requestLog.setLogLatency(true);
+        Slf4jRequestLogWriter slf4jRequestLogWriter = new Slf4jRequestLogWriter();
+        slf4jRequestLogWriter.setLoggerName(RestServer.class.getCanonicalName());
+        CustomRequestLog requestLog = new CustomRequestLog(slf4jRequestLogWriter, CustomRequestLog.EXTENDED_NCSA_FORMAT + " %msT");
         requestLogHandler.setRequestLog(requestLog);
 
-        HandlerCollection handlers = new HandlerCollection();
-        handlers.setHandlers(new Handler[]{context, new DefaultHandler(), requestLogHandler});
-
-        /* Needed for graceful shutdown as per `setStopTimeout` documentation */
-        StatisticsHandler statsHandler = new StatisticsHandler();
-        statsHandler.setHandler(handlers);
-        jettyServer.setHandler(statsHandler);
-        jettyServer.setStopTimeout(GRACEFUL_SHUTDOWN_TIMEOUT_MS);
-        jettyServer.setStopAtShutdown(true);
+        contextHandlers.add(new DefaultHandler());
+        contextHandlers.add(requestLogHandler);
 
+        handlers.setHandlers(contextHandlers.toArray(new Handler[]{}));
         try {
-            jettyServer.start();
+            context.start();
         } catch (Exception e) {
-            throw new ConnectException("Unable to start REST server", e);
+            throw new ConnectException("Unable to initialize REST resources", e);
         }
 
-        log.info("REST server listening at " + jettyServer.getURI() + ", advertising URL " + advertisedUrl());
+        if (adminResourceConfig != resourceConfig) {
+            try {
+                log.debug("Starting admin context");
+                adminContext.start();
+            } catch (Exception e) {
+                throw new ConnectException("Unable to initialize Admin REST resources", e);
+            }
+        }
+
+        log.info("REST resources initialized; server is started and ready to handle requests");
     }
 
     public URI serverUrl() {
@@ -237,9 +331,8 @@ public void stop() {
             jettyServer.stop();
             jettyServer.join();
         } catch (Exception e) {
-            throw new ConnectException("Unable to stop REST server", e);
-        } finally {
             jettyServer.destroy();
+            throw new ConnectException("Unable to stop REST server", e);
         }
 
         log.info("REST server stopped");
@@ -247,7 +340,8 @@ public void stop() {
 
     /**
      * Get the URL to advertise to other workers and clients. This uses the default connector from the embedded Jetty
-     * server, unless overrides for advertised hostname and/or port are provided via configs.
+     * server, unless overrides for advertised hostname and/or port are provided via configs. {@link #initializeServer()}
+     * must be invoked successfully before calling this method.
      */
     public URI advertisedUrl() {
         UriBuilder builder = UriBuilder.fromUri(jettyServer.getURI());
@@ -273,6 +367,34 @@ else if (serverConnector != null && serverConnector.getPort() > 0)
         return builder.build();
     }
 
+    /**
+     * @return the admin url for this worker. can be null if admin endpoints are disabled.
+     */
+    public URI adminUrl() {
+        ServerConnector adminConnector = null;
+        for (Connector connector : jettyServer.getConnectors()) {
+            if (ADMIN_SERVER_CONNECTOR_NAME.equals(connector.getName()))
+                adminConnector = (ServerConnector) connector;
+        }
+
+        if (adminConnector == null) {
+            List adminListeners = config.getList(WorkerConfig.ADMIN_LISTENERS_CONFIG);
+            if (adminListeners == null) {
+                return advertisedUrl();
+            } else if (adminListeners.isEmpty()) {
+                return null;
+            } else {
+                log.error("No admin connector found for listeners {}", adminListeners);
+                return null;
+            }
+        }
+
+        UriBuilder builder = UriBuilder.fromUri(jettyServer.getURI());
+        builder.port(adminConnector.getLocalPort());
+
+        return builder.build();
+    }
+
     String determineAdvertisedProtocol() {
         String advertisedSecurityProtocol = config.getString(WorkerConfig.REST_ADVERTISED_LISTENER_CONFIG);
         if (advertisedSecurityProtocol == null) {
@@ -294,24 +416,48 @@ else if (listeners.contains(String.format("%s://", PROTOCOL_HTTPS)))
         }
     }
 
+    /**
+     * Locate a Jetty connector for the standard (non-admin) REST API that uses the given protocol.
+     * @param protocol the protocol for the connector (e.g., "http" or "https").
+     * @return a {@link ServerConnector} for the server that uses the requested protocol, or
+     * {@code null} if none exist.
+     */
     ServerConnector findConnector(String protocol) {
         for (Connector connector : jettyServer.getConnectors()) {
-            if (connector.getName().startsWith(protocol))
+            String connectorName = connector.getName();
+            // We set the names for these connectors when instantiating them, beginning with the
+            // protocol for the connector and then an underscore ("_"). We rely on that format here
+            // when trying to locate a connector with the requested protocol; if the naming format
+            // for the connectors we create is ever changed, we'll need to adjust the logic here
+            // accordingly.
+            if (connectorName.startsWith(protocol + "_") && !ADMIN_SERVER_CONNECTOR_NAME.equals(connectorName))
                 return (ServerConnector) connector;
         }
 
         return null;
     }
 
-    void registerRestExtensions(HerderProvider provider, Plugins plugins, ResourceConfig resourceConfig) {
-        connectRestExtensions = plugins.newPlugins(
+    void registerRestExtensions(Herder herder, ResourceConfig resourceConfig) {
+        connectRestExtensions = herder.plugins().newPlugins(
             config.getList(WorkerConfig.REST_EXTENSION_CLASSES_CONFIG),
             config, ConnectRestExtension.class);
 
+        long herderRequestTimeoutMs = ConnectorsResource.REQUEST_TIMEOUT_MS;
+
+        Integer rebalanceTimeoutMs = config.getRebalanceTimeout();
+
+        if (rebalanceTimeoutMs != null) {
+            herderRequestTimeoutMs = Math.min(herderRequestTimeoutMs, rebalanceTimeoutMs.longValue());
+        }
+
+        ConnectClusterDetails connectClusterDetails = new ConnectClusterDetailsImpl(
+            herder.kafkaClusterId()
+        );
+
         ConnectRestExtensionContext connectRestExtensionContext =
             new ConnectRestExtensionContextImpl(
                 new ConnectRestConfigurable(resourceConfig),
-                new ConnectClusterStateImpl(provider)
+                new ConnectClusterStateImpl(herderRequestTimeoutMs, connectClusterDetails, herder)
             );
         for (ConnectRestExtension connectRestExtension : connectRestExtensions) {
             connectRestExtension.register(connectRestExtensionContext);
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/entities/ServerInfo.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/entities/ServerInfo.java
index a12751c7ae91e..e5c55533a31ab 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/entities/ServerInfo.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/entities/ServerInfo.java
@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.connect.runtime.rest.entities;
 
+import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import org.apache.kafka.common.utils.AppInfoParser;
 
@@ -24,12 +25,19 @@ public class ServerInfo {
     private final String commit;
     private final String kafkaClusterId;
 
-    public ServerInfo(String kafkaClusterId) {
-        this.version = AppInfoParser.getVersion();
-        this.commit = AppInfoParser.getCommitId();
+    @JsonCreator
+    private ServerInfo(@JsonProperty("version") String version,
+                       @JsonProperty("commit") String commit,
+                       @JsonProperty("kafka_cluster_id") String kafkaClusterId) {
+        this.version = version;
+        this.commit = commit;
         this.kafkaClusterId = kafkaClusterId;
     }
 
+    public ServerInfo(String kafkaClusterId) {
+        this(AppInfoParser.getVersion(), AppInfoParser.getCommitId(), kafkaClusterId);
+    }
+
     @JsonProperty
     public String version() {
         return version;
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/resources/ConnectorPluginsResource.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/resources/ConnectorPluginsResource.java
index 87f25b29cb51a..24eb93b8c0d5f 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/resources/ConnectorPluginsResource.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/resources/ConnectorPluginsResource.java
@@ -18,7 +18,7 @@
 
 import org.apache.kafka.connect.connector.Connector;
 import org.apache.kafka.connect.runtime.ConnectorConfig;
-import org.apache.kafka.connect.runtime.HerderProvider;
+import org.apache.kafka.connect.runtime.Herder;
 import org.apache.kafka.connect.runtime.isolation.PluginDesc;
 import org.apache.kafka.connect.runtime.rest.entities.ConfigInfos;
 import org.apache.kafka.connect.runtime.rest.entities.ConnectorPluginInfo;
@@ -49,7 +49,7 @@
 public class ConnectorPluginsResource {
 
     private static final String ALIAS_SUFFIX = "Connector";
-    private final HerderProvider herderProvider;
+    private final Herder herder;
     private final List connectorPlugins;
 
     private static final List> CONNECTOR_EXCLUDES = Arrays.asList(
@@ -58,8 +58,8 @@ public class ConnectorPluginsResource {
             SchemaSourceConnector.class
     );
 
-    public ConnectorPluginsResource(HerderProvider herderProvider) {
-        this.herderProvider = herderProvider;
+    public ConnectorPluginsResource(Herder herder) {
+        this.herder = herder;
         this.connectorPlugins = new ArrayList<>();
     }
 
@@ -78,7 +78,7 @@ public ConfigInfos validateConfigs(
             );
         }
 
-        return herderProvider.get().validateConnectorConfig(connectorConfig);
+        return herder.validateConnectorConfig(connectorConfig);
     }
 
     @GET
@@ -90,7 +90,7 @@ public List listConnectorPlugins() {
     // TODO: improve once plugins are allowed to be added/removed during runtime.
     private synchronized List getConnectorPlugins() {
         if (connectorPlugins.isEmpty()) {
-            for (PluginDesc plugin : herderProvider.get().plugins().connectors()) {
+            for (PluginDesc plugin : herder.plugins().connectors()) {
                 if (!CONNECTOR_EXCLUDES.contains(plugin.pluginClass())) {
                     connectorPlugins.add(new ConnectorPluginInfo(plugin));
                 }
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/resources/ConnectorsResource.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/resources/ConnectorsResource.java
index 29a8c39028ef7..8728e1cd1748e 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/resources/ConnectorsResource.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/resources/ConnectorsResource.java
@@ -17,12 +17,17 @@
 package org.apache.kafka.connect.runtime.rest.resources;
 
 import com.fasterxml.jackson.core.type.TypeReference;
+
+import javax.ws.rs.core.HttpHeaders;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import org.apache.kafka.connect.errors.NotFoundException;
 import org.apache.kafka.connect.runtime.ConnectorConfig;
 import org.apache.kafka.connect.runtime.Herder;
-import org.apache.kafka.connect.runtime.HerderProvider;
 import org.apache.kafka.connect.runtime.WorkerConfig;
 import org.apache.kafka.connect.runtime.distributed.RebalanceNeededException;
 import org.apache.kafka.connect.runtime.distributed.RequestTargetException;
+import org.apache.kafka.connect.runtime.rest.InternalRequestSignature;
 import org.apache.kafka.connect.runtime.rest.RestClient;
 import org.apache.kafka.connect.runtime.rest.entities.ConnectorInfo;
 import org.apache.kafka.connect.runtime.rest.entities.ConnectorStateInfo;
@@ -45,11 +50,14 @@
 import javax.ws.rs.PathParam;
 import javax.ws.rs.Produces;
 import javax.ws.rs.QueryParam;
+import javax.ws.rs.core.Context;
 import javax.ws.rs.core.MediaType;
 import javax.ws.rs.core.Response;
 import javax.ws.rs.core.UriBuilder;
+import javax.ws.rs.core.UriInfo;
+
 import java.net.URI;
-import java.util.Collection;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ExecutionException;
@@ -61,39 +69,66 @@
 @Consumes(MediaType.APPLICATION_JSON)
 public class ConnectorsResource {
     private static final Logger log = LoggerFactory.getLogger(ConnectorsResource.class);
+    private static final TypeReference>> TASK_CONFIGS_TYPE =
+        new TypeReference>>() { };
 
     // TODO: This should not be so long. However, due to potentially long rebalances that may have to wait a full
     // session timeout to complete, during which we cannot serve some requests. Ideally we could reduce this, but
     // we need to consider all possible scenarios this could fail. It might be ok to fail with a timeout in rare cases,
     // but currently a worker simply leaving the group can take this long as well.
-    private static final long REQUEST_TIMEOUT_MS = 90 * 1000;
+    public static final long REQUEST_TIMEOUT_MS = 90 * 1000;
 
-    private final HerderProvider herderProvider;
+    private final Herder herder;
     private final WorkerConfig config;
     @javax.ws.rs.core.Context
     private ServletContext context;
 
-    public ConnectorsResource(HerderProvider herder, WorkerConfig config) {
-        this.herderProvider = herder;
+    public ConnectorsResource(Herder herder, WorkerConfig config) {
+        this.herder = herder;
         this.config = config;
     }
 
-    private Herder herder() {
-        return herderProvider.get();
-    }
-
     @GET
     @Path("/")
-    public Collection listConnectors(final @QueryParam("forward") Boolean forward) throws Throwable {
-        FutureCallback> cb = new FutureCallback<>();
-        herder().connectors(cb);
-        return completeOrForwardRequest(cb, "/connectors", "GET", null, new TypeReference>() {
-        }, forward);
+    public Response listConnectors(
+        final @Context UriInfo uriInfo,
+        final @Context HttpHeaders headers
+    ) throws Throwable {
+        if (uriInfo.getQueryParameters().containsKey("expand")) {
+            Map> out = new HashMap<>();
+            for (String connector : herder.connectors()) {
+                try {
+                    Map connectorExpansions = new HashMap<>();
+                    for (String expansion : uriInfo.getQueryParameters().get("expand")) {
+                        switch (expansion) {
+                            case "status":
+                                connectorExpansions.put("status", herder.connectorStatus(connector));
+                                break;
+                            case "info":
+                                connectorExpansions.put("info", herder.connectorInfo(connector));
+                                break;
+                            default:
+                                log.info("Ignoring unknown expanion type {}", expansion);
+                        }
+                    }
+                    out.put(connector, connectorExpansions);
+                } catch (NotFoundException e) {
+                    // this likely means that a connector has been removed while we look its info up
+                    // we can just not include this connector in the return entity
+                    log.debug("Unable to get connector info for {} on this worker", connector);
+                }
+
+            }
+            return Response.ok(out).build();
+        } else {
+            return Response.ok(herder.connectors()).build();
+        }
     }
 
     @POST
     @Path("/")
     public Response createConnector(final @QueryParam("forward") Boolean forward,
+                                    final @Context HttpHeaders headers,
                                     final CreateConnectorRequest createRequest) throws Throwable {
         // Trim leading and trailing whitespaces from the connector name, replace null with empty string
         // if no name element present to keep validation within validator (NonEmptyStringWithoutControlChars
@@ -104,8 +139,8 @@ public Response createConnector(final @QueryParam("forward") Boolean forward,
         checkAndPutConnectorConfigName(name, configs);
 
         FutureCallback> cb = new FutureCallback<>();
-        herder().putConnectorConfig(name, configs, false, cb);
-        Herder.Created info = completeOrForwardRequest(cb, "/connectors", "POST", createRequest,
+        herder.putConnectorConfig(name, configs, false, cb);
+        Herder.Created info = completeOrForwardRequest(cb, "/connectors", "POST", headers, createRequest,
                 new TypeReference() { }, new CreatedConnectorInfoTranslator(), forward);
 
         URI location = UriBuilder.fromUri("/connectors").path(name).build();
@@ -115,38 +150,41 @@ public Response createConnector(final @QueryParam("forward") Boolean forward,
     @GET
     @Path("/{connector}")
     public ConnectorInfo getConnector(final @PathParam("connector") String connector,
+                                      final @Context HttpHeaders headers,
                                       final @QueryParam("forward") Boolean forward) throws Throwable {
         FutureCallback cb = new FutureCallback<>();
-        herder().connectorInfo(connector, cb);
-        return completeOrForwardRequest(cb, "/connectors/" + connector, "GET", null, forward);
+        herder.connectorInfo(connector, cb);
+        return completeOrForwardRequest(cb, "/connectors/" + connector, "GET", headers, null, forward);
     }
 
     @GET
     @Path("/{connector}/config")
     public Map getConnectorConfig(final @PathParam("connector") String connector,
+                                                  final @Context HttpHeaders headers,
                                                   final @QueryParam("forward") Boolean forward) throws Throwable {
         FutureCallback> cb = new FutureCallback<>();
-        herder().connectorConfig(connector, cb);
-        return completeOrForwardRequest(cb, "/connectors/" + connector + "/config", "GET", null, forward);
+        herder.connectorConfig(connector, cb);
+        return completeOrForwardRequest(cb, "/connectors/" + connector + "/config", "GET", headers, null, forward);
     }
 
     @GET
     @Path("/{connector}/status")
     public ConnectorStateInfo getConnectorStatus(final @PathParam("connector") String connector) throws Throwable {
-        return herder().connectorStatus(connector);
+        return herder.connectorStatus(connector);
     }
 
     @PUT
     @Path("/{connector}/config")
     public Response putConnectorConfig(final @PathParam("connector") String connector,
+                                       final @Context HttpHeaders headers,
                                        final @QueryParam("forward") Boolean forward,
                                        final Map connectorConfig) throws Throwable {
         FutureCallback> cb = new FutureCallback<>();
         checkAndPutConnectorConfigName(connector, connectorConfig);
 
-        herder().putConnectorConfig(connector, connectorConfig, true, cb);
+        herder.putConnectorConfig(connector, connectorConfig, true, cb);
         Herder.Created createdInfo = completeOrForwardRequest(cb, "/connectors/" + connector + "/config",
-                "PUT", connectorConfig, new TypeReference() { }, new CreatedConnectorInfoTranslator(), forward);
+                "PUT", headers, connectorConfig, new TypeReference() { }, new CreatedConnectorInfoTranslator(), forward);
         Response.ResponseBuilder response;
         if (createdInfo.created()) {
             URI location = UriBuilder.fromUri("/connectors").path(connector).build();
@@ -160,71 +198,78 @@ public Response putConnectorConfig(final @PathParam("connector") String connecto
     @POST
     @Path("/{connector}/restart")
     public void restartConnector(final @PathParam("connector") String connector,
+                                 final @Context HttpHeaders headers,
                                  final @QueryParam("forward") Boolean forward) throws Throwable {
         FutureCallback cb = new FutureCallback<>();
-        herder().restartConnector(connector, cb);
-        completeOrForwardRequest(cb, "/connectors/" + connector + "/restart", "POST", null, forward);
+        herder.restartConnector(connector, cb);
+        completeOrForwardRequest(cb, "/connectors/" + connector + "/restart", "POST", headers, null, forward);
     }
 
     @PUT
     @Path("/{connector}/pause")
-    public Response pauseConnector(@PathParam("connector") String connector) {
-        herder().pauseConnector(connector);
+    public Response pauseConnector(@PathParam("connector") String connector, final @Context HttpHeaders headers) {
+        herder.pauseConnector(connector);
         return Response.accepted().build();
     }
 
     @PUT
     @Path("/{connector}/resume")
     public Response resumeConnector(@PathParam("connector") String connector) {
-        herder().resumeConnector(connector);
+        herder.resumeConnector(connector);
         return Response.accepted().build();
     }
 
     @GET
     @Path("/{connector}/tasks")
     public List getTaskConfigs(final @PathParam("connector") String connector,
+                                         final @Context HttpHeaders headers,
                                          final @QueryParam("forward") Boolean forward) throws Throwable {
         FutureCallback> cb = new FutureCallback<>();
-        herder().taskConfigs(connector, cb);
-        return completeOrForwardRequest(cb, "/connectors/" + connector + "/tasks", "GET", null, new TypeReference>() {
+        herder.taskConfigs(connector, cb);
+        return completeOrForwardRequest(cb, "/connectors/" + connector + "/tasks", "GET", headers, null, new TypeReference>() {
         }, forward);
     }
 
     @POST
     @Path("/{connector}/tasks")
     public void putTaskConfigs(final @PathParam("connector") String connector,
+                               final @Context HttpHeaders headers,
                                final @QueryParam("forward") Boolean forward,
-                               final List> taskConfigs) throws Throwable {
+                               final byte[] requestBody) throws Throwable {
+        List> taskConfigs = new ObjectMapper().readValue(requestBody, TASK_CONFIGS_TYPE);
         FutureCallback cb = new FutureCallback<>();
-        herder().putTaskConfigs(connector, taskConfigs, cb);
-        completeOrForwardRequest(cb, "/connectors/" + connector + "/tasks", "POST", taskConfigs, forward);
+        herder.putTaskConfigs(connector, taskConfigs, cb, InternalRequestSignature.fromHeaders(requestBody, headers));
+        completeOrForwardRequest(cb, "/connectors/" + connector + "/tasks", "POST", headers, taskConfigs, forward);
     }
 
     @GET
     @Path("/{connector}/tasks/{task}/status")
     public ConnectorStateInfo.TaskState getTaskStatus(final @PathParam("connector") String connector,
+                                                      final @Context HttpHeaders headers,
                                                       final @PathParam("task") Integer task) throws Throwable {
-        return herder().taskStatus(new ConnectorTaskId(connector, task));
+        return herder.taskStatus(new ConnectorTaskId(connector, task));
     }
 
     @POST
     @Path("/{connector}/tasks/{task}/restart")
     public void restartTask(final @PathParam("connector") String connector,
                             final @PathParam("task") Integer task,
+                            final @Context HttpHeaders headers,
                             final @QueryParam("forward") Boolean forward) throws Throwable {
         FutureCallback cb = new FutureCallback<>();
         ConnectorTaskId taskId = new ConnectorTaskId(connector, task);
-        herder().restartTask(taskId, cb);
-        completeOrForwardRequest(cb, "/connectors/" + connector + "/tasks/" + task + "/restart", "POST", null, forward);
+        herder.restartTask(taskId, cb);
+        completeOrForwardRequest(cb, "/connectors/" + connector + "/tasks/" + task + "/restart", "POST", headers, null, forward);
     }
 
     @DELETE
     @Path("/{connector}")
     public void destroyConnector(final @PathParam("connector") String connector,
+                                 final @Context HttpHeaders headers,
                                  final @QueryParam("forward") Boolean forward) throws Throwable {
         FutureCallback> cb = new FutureCallback<>();
-        herder().deleteConnectorConfig(connector, cb);
-        completeOrForwardRequest(cb, "/connectors/" + connector, "DELETE", null, forward);
+        herder.deleteConnectorConfig(connector, cb);
+        completeOrForwardRequest(cb, "/connectors/" + connector, "DELETE", headers, null, forward);
     }
 
     // Check whether the connector name from the url matches the one (if there is one) provided in the connectorconfig
@@ -244,6 +289,7 @@ private void checkAndPutConnectorConfigName(String connectorName, Map T completeOrForwardRequest(FutureCallback cb,
                                               String path,
                                               String method,
+                                              HttpHeaders headers,
                                               Object body,
                                               TypeReference resultType,
                                               Translator translator,
@@ -260,13 +306,20 @@ private  T completeOrForwardRequest(FutureCallback cb,
                     // this gives two total hops to resolve the request before giving up.
                     boolean recursiveForward = forward == null;
                     RequestTargetException targetException = (RequestTargetException) cause;
-                    String forwardUrl = UriBuilder.fromUri(targetException.forwardUrl())
+                    String forwardedUrl = targetException.forwardUrl();
+                    if (forwardedUrl == null) {
+                        // the target didn't know of the leader at this moment.
+                        throw new ConnectRestException(Response.Status.CONFLICT.getStatusCode(),
+                                "Cannot complete request momentarily due to no known leader URL, "
+                                + "likely because a rebalance was underway.");
+                    }
+                    String forwardUrl = UriBuilder.fromUri(forwardedUrl)
                             .path(path)
                             .queryParam("forward", recursiveForward)
                             .build()
                             .toString();
                     log.debug("Forwarding request {} {} {}", forwardUrl, method, body);
-                    return translator.translate(RestClient.httpRequest(forwardUrl, method, body, resultType, config));
+                    return translator.translate(RestClient.httpRequest(forwardUrl, method, headers, body, resultType, config));
                 } else {
                     // we should find the right target for the query within two hops, so if
                     // we don't, it probably means that a rebalance has taken place.
@@ -288,14 +341,14 @@ private  T completeOrForwardRequest(FutureCallback cb,
         }
     }
 
-    private  T completeOrForwardRequest(FutureCallback cb, String path, String method, Object body,
+    private  T completeOrForwardRequest(FutureCallback cb, String path, String method, HttpHeaders headers, Object body,
                                            TypeReference resultType, Boolean forward) throws Throwable {
-        return completeOrForwardRequest(cb, path, method, body, resultType, new IdentityTranslator(), forward);
+        return completeOrForwardRequest(cb, path, method, headers, body, resultType, new IdentityTranslator(), forward);
     }
 
-    private  T completeOrForwardRequest(FutureCallback cb, String path, String method,
+    private  T completeOrForwardRequest(FutureCallback cb, String path, String method, HttpHeaders headers,
                                            Object body, Boolean forward) throws Throwable {
-        return completeOrForwardRequest(cb, path, method, body, null, new IdentityTranslator(), forward);
+        return completeOrForwardRequest(cb, path, method, headers, body, null, new IdentityTranslator(), forward);
     }
 
     private interface Translator {
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/resources/LoggingResource.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/resources/LoggingResource.java
new file mode 100644
index 0000000000000..05796600a936b
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/resources/LoggingResource.java
@@ -0,0 +1,206 @@
+/*
+ * 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.rest.resources;
+
+import org.apache.kafka.connect.errors.NotFoundException;
+import org.apache.kafka.connect.runtime.rest.errors.BadRequestException;
+import org.apache.log4j.Level;
+import org.apache.log4j.LogManager;
+import org.apache.log4j.Logger;
+
+import javax.ws.rs.Consumes;
+import javax.ws.rs.GET;
+import javax.ws.rs.PUT;
+import javax.ws.rs.Path;
+import javax.ws.rs.PathParam;
+import javax.ws.rs.Produces;
+import javax.ws.rs.core.MediaType;
+import javax.ws.rs.core.Response;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.Enumeration;
+import java.util.List;
+import java.util.Locale;
+import java.util.Map;
+import java.util.Objects;
+import java.util.TreeMap;
+
+/**
+ * A set of endpoints to adjust the log levels of runtime loggers.
+ */
+@Path("/admin/loggers")
+@Produces(MediaType.APPLICATION_JSON)
+@Consumes(MediaType.APPLICATION_JSON)
+public class LoggingResource {
+
+    /**
+     * Log4j uses "root" (case insensitive) as name of the root logger.
+     */
+    private static final String ROOT_LOGGER_NAME = "root";
+
+    /**
+     * List the current loggers that have their levels explicitly set and their log levels.
+     *
+     * @return a list of current loggers and their levels.
+     */
+    @GET
+    @Path("/")
+    public Response listLoggers() {
+        Map> loggers = new TreeMap<>();
+        Enumeration enumeration = currentLoggers();
+        Collections.list(enumeration)
+                .stream()
+                .filter(logger -> logger.getLevel() != null)
+                .forEach(logger -> loggers.put(logger.getName(), levelToMap(logger)));
+
+        Logger root = rootLogger();
+        if (root.getLevel() != null) {
+            loggers.put(ROOT_LOGGER_NAME, levelToMap(root));
+        }
+
+        return Response.ok(loggers).build();
+    }
+
+    /**
+     * Get the log level of a named logger.
+     *
+     * @param namedLogger name of a logger
+     * @return level of the logger, effective level if the level was not explicitly set.
+     */
+    @GET
+    @Path("/{logger}")
+    public Response getLogger(final @PathParam("logger") String namedLogger) {
+        Objects.requireNonNull(namedLogger, "require non-null name");
+
+        Logger logger = null;
+        if (ROOT_LOGGER_NAME.equalsIgnoreCase(namedLogger)) {
+            logger = rootLogger();
+        } else {
+            Enumeration en = currentLoggers();
+            // search within existing loggers for the given name.
+            // using LogManger.getLogger() will create a logger if it doesn't exist
+            // (potential leak since these don't get cleaned up).
+            while (en.hasMoreElements()) {
+                Logger l = en.nextElement();
+                if (namedLogger.equals(l.getName())) {
+                    logger = l;
+                    break;
+                }
+            }
+        }
+        if (logger == null) {
+            throw new NotFoundException("Logger " + namedLogger + " not found.");
+        } else {
+            return Response.ok(effectiveLevelToMap(logger)).build();
+        }
+    }
+
+
+    /**
+     * Adjust level of a named logger. if name corresponds to an ancestor, then the log level is applied to all child loggers.
+     *
+     * @param namedLogger name of the logger
+     * @param levelMap a map that is expected to contain one key 'level', and a value that is one of the log4j levels:
+     *                 DEBUG, ERROR, FATAL, INFO, TRACE, WARN
+     * @return names of loggers whose levels were modified
+     */
+    @PUT
+    @Path("/{logger}")
+    public Response setLevel(final @PathParam("logger") String namedLogger,
+                             final Map levelMap) {
+        String desiredLevelStr = levelMap.get("level");
+        if (desiredLevelStr == null) {
+            throw new BadRequestException("Desired 'level' parameter was not specified in request.");
+        }
+
+        Level level = Level.toLevel(desiredLevelStr.toUpperCase(Locale.ROOT), null);
+        if (level == null) {
+            throw new NotFoundException("invalid log level '" + desiredLevelStr + "'.");
+        }
+
+        List childLoggers;
+        if (ROOT_LOGGER_NAME.equalsIgnoreCase(namedLogger)) {
+            childLoggers = Collections.list(currentLoggers());
+            childLoggers.add(rootLogger());
+        } else {
+            childLoggers = new ArrayList<>();
+            Logger ancestorLogger = lookupLogger(namedLogger);
+            Enumeration en = currentLoggers();
+            boolean present = false;
+            while (en.hasMoreElements()) {
+                Logger current = (Logger) en.nextElement();
+                if (current.getName().startsWith(namedLogger)) {
+                    childLoggers.add(current);
+                }
+                if (namedLogger.equals(current.getName())) {
+                    present = true;
+                }
+            }
+            if (!present) {
+                childLoggers.add(ancestorLogger);
+            }
+        }
+
+        List modifiedLoggerNames = new ArrayList<>();
+        for (Logger logger: childLoggers) {
+            logger.setLevel(level);
+            modifiedLoggerNames.add(logger.getName());
+        }
+        Collections.sort(modifiedLoggerNames);
+
+        return Response.ok(modifiedLoggerNames).build();
+    }
+
+    protected Logger lookupLogger(String namedLogger) {
+        return LogManager.getLogger(namedLogger);
+    }
+
+    @SuppressWarnings("unchecked")
+    protected Enumeration currentLoggers() {
+        return LogManager.getCurrentLoggers();
+    }
+
+    protected Logger rootLogger() {
+        return LogManager.getRootLogger();
+    }
+
+    /**
+     *
+     * Map representation of a logger's effective log level.
+     *
+     * @param logger a non-null log4j logger
+     * @return a singleton map whose key is level and the value is the string representation of the logger's effective log level.
+     */
+    private static Map effectiveLevelToMap(Logger logger) {
+        Level level = logger.getLevel();
+        if (level == null) {
+            level = logger.getEffectiveLevel();
+        }
+        return Collections.singletonMap("level", String.valueOf(level));
+    }
+
+    /**
+     *
+     * Map representation of a logger's log level.
+     *
+     * @param logger a non-null log4j logger
+     * @return a singleton map whose key is level and the value is the string representation of the logger's log level.
+     */
+    private static Map levelToMap(Logger logger) {
+        return Collections.singletonMap("level", String.valueOf(logger.getLevel()));
+    }
+}
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/resources/RootResource.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/resources/RootResource.java
index 56516cd410924..9666bf15954f9 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/resources/RootResource.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/resources/RootResource.java
@@ -16,7 +16,7 @@
  */
 package org.apache.kafka.connect.runtime.rest.resources;
 
-import org.apache.kafka.connect.runtime.HerderProvider;
+import org.apache.kafka.connect.runtime.Herder;
 import org.apache.kafka.connect.runtime.rest.entities.ServerInfo;
 
 import javax.ws.rs.GET;
@@ -28,15 +28,15 @@
 @Produces(MediaType.APPLICATION_JSON)
 public class RootResource {
 
-    private final HerderProvider herder;
+    private final Herder herder;
 
-    public RootResource(HerderProvider herder) {
+    public RootResource(Herder herder) {
         this.herder = herder;
     }
 
     @GET
     @Path("/")
     public ServerInfo serverInfo() {
-        return new ServerInfo(herder.get().kafkaClusterId());
+        return new ServerInfo(herder.kafkaClusterId());
     }
 }
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 a50a0b0439c91..6b391d96adaa7 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
@@ -34,28 +34,42 @@ public class SSLUtils {
 
     private static final Pattern COMMA_WITH_WHITESPACE = Pattern.compile("\\s*,\\s*");
 
+
     /**
-     * Configures SSL/TLS for HTTPS Jetty Server / Client
+     * Configures SSL/TLS for HTTPS Jetty Server using configs with the given prefix
      */
-    public static SslContextFactory createSslContextFactory(WorkerConfig config) {
-        return createSslContextFactory(config, false);
+    public static SslContextFactory createServerSideSslContextFactory(WorkerConfig config, String prefix) {
+        Map sslConfigValues = config.valuesWithPrefixAllOrNothing(prefix);
+
+        final SslContextFactory.Server ssl = new SslContextFactory.Server();
+
+        configureSslContextFactoryKeyStore(ssl, sslConfigValues);
+        configureSslContextFactoryTrustStore(ssl, sslConfigValues);
+        configureSslContextFactoryAlgorithms(ssl, sslConfigValues);
+        configureSslContextFactoryAuthentication(ssl, sslConfigValues);
+
+        return ssl;
     }
 
     /**
-     * Configures SSL/TLS for HTTPS Jetty Server / Client
+     * Configures SSL/TLS for HTTPS Jetty Server
      */
-    public static SslContextFactory createSslContextFactory(WorkerConfig config, boolean client) {
+    public static SslContextFactory createServerSideSslContextFactory(WorkerConfig config) {
+        return createServerSideSslContextFactory(config, "listeners.https.");
+    }
+
+    /**
+     * Configures SSL/TLS for HTTPS Jetty Client
+     */
+    public static SslContextFactory createClientSideSslContextFactory(WorkerConfig config) {
         Map sslConfigValues = config.valuesWithPrefixAllOrNothing("listeners.https.");
 
-        SslContextFactory ssl = new SslContextFactory();
+        final SslContextFactory.Client ssl = new SslContextFactory.Client();
 
         configureSslContextFactoryKeyStore(ssl, sslConfigValues);
         configureSslContextFactoryTrustStore(ssl, sslConfigValues);
         configureSslContextFactoryAlgorithms(ssl, sslConfigValues);
-        configureSslContextFactoryAuthentication(ssl, sslConfigValues);
-
-        if (client)
-            configureSslContextFactoryEndpointIdentification(ssl, sslConfigValues);
+        configureSslContextFactoryEndpointIdentification(ssl, sslConfigValues);
 
         return ssl;
     }
@@ -140,7 +154,7 @@ protected static void configureSslContextFactoryEndpointIdentification(SslContex
     /**
      * Configures Authentication related settings in SslContextFactory
      */
-    protected static void configureSslContextFactoryAuthentication(SslContextFactory ssl, Map sslConfigValues) {
+    protected static void configureSslContextFactoryAuthentication(SslContextFactory.Server ssl, Map sslConfigValues) {
         String sslClientAuth = (String) getOrDefault(sslConfigValues, BrokerSecurityConfigs.SSL_CLIENT_AUTH_CONFIG, "none");
         switch (sslClientAuth) {
             case "requested":
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/standalone/StandaloneHerder.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/standalone/StandaloneHerder.java
index 172c9b293a9c4..6c5398f8e4573 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/standalone/StandaloneHerder.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/standalone/StandaloneHerder.java
@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.connect.runtime.standalone;
 
+import org.apache.kafka.connect.connector.policy.ConnectorClientConfigOverridePolicy;
 import org.apache.kafka.connect.errors.AlreadyExistsException;
 import org.apache.kafka.connect.errors.ConnectException;
 import org.apache.kafka.connect.errors.NotFoundException;
@@ -23,11 +24,13 @@
 import org.apache.kafka.connect.runtime.ConnectorConfig;
 import org.apache.kafka.connect.runtime.HerderConnectorContext;
 import org.apache.kafka.connect.runtime.HerderRequest;
+import org.apache.kafka.connect.runtime.SessionKey;
 import org.apache.kafka.connect.runtime.SinkConnectorConfig;
 import org.apache.kafka.connect.runtime.SourceConnectorConfig;
 import org.apache.kafka.connect.runtime.TargetState;
 import org.apache.kafka.connect.runtime.Worker;
 import org.apache.kafka.connect.runtime.distributed.ClusterConfigState;
+import org.apache.kafka.connect.runtime.rest.InternalRequestSignature;
 import org.apache.kafka.connect.runtime.rest.entities.ConnectorInfo;
 import org.apache.kafka.connect.runtime.rest.entities.TaskInfo;
 import org.apache.kafka.connect.storage.ConfigBackingStore;
@@ -62,12 +65,14 @@ public class StandaloneHerder extends AbstractHerder {
 
     private ClusterConfigState configState;
 
-    public StandaloneHerder(Worker worker, String kafkaClusterId) {
+    public StandaloneHerder(Worker worker, String kafkaClusterId,
+                            ConnectorClientConfigOverridePolicy connectorClientConfigOverridePolicy) {
         this(worker,
                 worker.workerId(),
                 kafkaClusterId,
                 new MemoryStatusBackingStore(),
-                new MemoryConfigBackingStore(worker.configTransformer()));
+                new MemoryConfigBackingStore(worker.configTransformer()),
+             connectorClientConfigOverridePolicy);
     }
 
     // visible for testing
@@ -75,8 +80,9 @@ public StandaloneHerder(Worker worker, String kafkaClusterId) {
                      String workerId,
                      String kafkaClusterId,
                      StatusBackingStore statusBackingStore,
-                     MemoryConfigBackingStore configBackingStore) {
-        super(worker, workerId, kafkaClusterId, statusBackingStore, configBackingStore);
+                     MemoryConfigBackingStore configBackingStore,
+                     ConnectorClientConfigOverridePolicy connectorClientConfigOverridePolicy) {
+        super(worker, workerId, kafkaClusterId, statusBackingStore, configBackingStore, connectorClientConfigOverridePolicy);
         this.configState = ClusterConfigState.EMPTY;
         this.requestExecutorService = Executors.newSingleThreadScheduledExecutor();
         configBackingStore.setUpdateListener(new ConfigUpdateListener());
@@ -103,7 +109,7 @@ public synchronized void stop() {
         // There's no coordination/hand-off to do here since this is all standalone. Instead, we
         // should just clean up the stuff we normally would, i.e. cleanly checkpoint and shutdown all
         // the tasks.
-        for (String connName : configState.connectors()) {
+        for (String connName : connectors()) {
             removeConnectorTasks(connName);
             worker.stopConnector(connName);
         }
@@ -118,12 +124,12 @@ public int generation() {
 
     @Override
     public synchronized void connectors(Callback> callback) {
-        callback.onCompletion(null, configState.connectors());
+        callback.onCompletion(null, connectors());
     }
-
+    
     @Override
     public synchronized void connectorInfo(String connName, Callback callback) {
-        ConnectorInfo connectorInfo = createConnectorInfo(connName);
+        ConnectorInfo connectorInfo = connectorInfo(connName);
         if (connectorInfo == null) {
             callback.onCompletion(new NotFoundException("Connector " + connName + " not found"), null);
             return;
@@ -237,7 +243,7 @@ public synchronized void taskConfigs(String connName, Callback> c
     }
 
     @Override
-    public void putTaskConfigs(String connName, List> configs, Callback callback) {
+    public void putTaskConfigs(String connName, List> configs, Callback callback, InternalRequestSignature requestSignature) {
         throw new UnsupportedOperationException("Kafka Connect in standalone mode does not support externally setting task configurations.");
     }
 
@@ -313,6 +319,7 @@ private void removeConnectorTasks(String connName) {
         if (!tasks.isEmpty()) {
             worker.stopAndAwaitTasks(tasks);
             configBackingStore.removeTaskConfigs(connName);
+            tasks.forEach(this::onDeletion);
         }
     }
 
@@ -375,6 +382,11 @@ public void onConnectorTargetStateChange(String connector) {
                     updateConnectorTasks(connector);
             }
         }
+
+        @Override
+        public void onSessionKeyUpdate(SessionKey sessionKey) {
+            // no-op
+        }
     }
 
     static class StandaloneHerderRequest implements HerderRequest {
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/storage/CloseableOffsetStorageReader.java b/connect/runtime/src/main/java/org/apache/kafka/connect/storage/CloseableOffsetStorageReader.java
new file mode 100644
index 0000000000000..b90273936bfb0
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/storage/CloseableOffsetStorageReader.java
@@ -0,0 +1,33 @@
+/*
+ * 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 java.io.Closeable;
+import java.util.Collection;
+import java.util.Map;
+import java.util.concurrent.Future;
+
+public interface CloseableOffsetStorageReader extends Closeable, OffsetStorageReader {
+
+    /**
+     * {@link Future#cancel(boolean) Cancel} all outstanding offset read requests, and throw an
+     * exception in all current and future calls to {@link #offsets(Collection)} and
+     * {@link #offset(Map)}. This is useful for unblocking task threads which need to shut down but
+     * are blocked on offset reads.
+     */
+    void close();
+}
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/storage/ConfigBackingStore.java b/connect/runtime/src/main/java/org/apache/kafka/connect/storage/ConfigBackingStore.java
index b8fd64380cdb7..f8a6f70fd2a7f 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/storage/ConfigBackingStore.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/storage/ConfigBackingStore.java
@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.connect.storage;
 
+import org.apache.kafka.connect.runtime.SessionKey;
 import org.apache.kafka.connect.runtime.TargetState;
 import org.apache.kafka.connect.runtime.distributed.ClusterConfigState;
 import org.apache.kafka.connect.util.ConnectorTaskId;
@@ -88,6 +89,8 @@ public interface ConfigBackingStore {
      */
     void putTargetState(String connector, TargetState state);
 
+    void putSessionKey(SessionKey sessionKey);
+
     /**
      * Set an update listener to get notifications when there are config/target state
      * changes.
@@ -119,6 +122,12 @@ interface UpdateListener {
          * @param connector name of the connector
          */
         void onConnectorTargetStateChange(String connector);
+
+        /**
+         * Invoked when the leader has distributed a new session key
+         * @param sessionKey the {@link SessionKey session key}
+         */
+        void onSessionKeyUpdate(SessionKey sessionKey);
     }
 
 }
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 8cf47fe04799f..4c6429ce9304f 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
@@ -33,6 +33,7 @@
 import org.apache.kafka.connect.data.Struct;
 import org.apache.kafka.connect.errors.ConnectException;
 import org.apache.kafka.connect.errors.DataException;
+import org.apache.kafka.connect.runtime.SessionKey;
 import org.apache.kafka.connect.runtime.TargetState;
 import org.apache.kafka.connect.runtime.WorkerConfig;
 import org.apache.kafka.connect.runtime.WorkerConfigTransformer;
@@ -45,8 +46,10 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import javax.crypto.spec.SecretKeySpec;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Base64;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
@@ -176,6 +179,8 @@ public static String COMMIT_TASKS_KEY(String connectorName) {
         return COMMIT_TASKS_PREFIX + connectorName;
     }
 
+    public static final String SESSION_KEY_KEY = "session-key";
+
     // Note that while using real serialization for values as we have here, but ad hoc string serialization for keys,
     // isn't ideal, we use this approach because it avoids any potential problems with schema evolution or
     // converter/serializer changes causing keys to change. We need to absolutely ensure that the keys remain precisely
@@ -190,6 +195,13 @@ public static String COMMIT_TASKS_KEY(String connectorName) {
     public static final Schema TARGET_STATE_V0 = SchemaBuilder.struct()
             .field("state", Schema.STRING_SCHEMA)
             .build();
+    // The key is logically a byte array, but we can't use the JSON converter to (de-)serialize that without a schema.
+    // So instead, we base 64-encode it before serializing and decode it after deserializing.
+    public static final Schema SESSION_KEY_V0 = SchemaBuilder.struct()
+            .field("key", Schema.STRING_SCHEMA)
+            .field("algorithm", Schema.STRING_SCHEMA)
+            .field("creation-timestamp", Schema.INT64_SCHEMA)
+            .build();
 
     private static final long READ_TO_END_TIMEOUT_MS = 30000;
 
@@ -216,6 +228,8 @@ public static String COMMIT_TASKS_KEY(String connectorName) {
     // The most recently read offset. This does not take into account deferred task updates/commits, so we may have
     // outstanding data to be applied.
     private volatile long offset;
+    // The most recently read session key, to use for validating internal REST requests.
+    private volatile SessionKey sessionKey;
 
     // Connector -> Map[ConnectorTaskId -> Configs]
     private final Map>> deferredTaskUpdates = new HashMap<>();
@@ -249,6 +263,16 @@ public void start() {
         // Before startup, callbacks are *not* invoked. You can grab a snapshot after starting -- just take care that
         // updates can continue to occur in the background
         configLog.start();
+
+        int partitionCount = configLog.partitionCount();
+        if (partitionCount > 1) {
+            String msg = String.format("Topic '%s' supplied via the '%s' property is required "
+                    + "to have a single partition in order to guarantee consistency of "
+                    + "connector configurations, but found %d partitions.",
+                    topic, DistributedConfig.CONFIG_TOPIC_CONFIG, partitionCount);
+            throw new ConfigException(msg);
+        }
+
         started = true;
         log.info("Started KafkaConfigBackingStore");
     }
@@ -266,10 +290,11 @@ public void stop() {
     @Override
     public ClusterConfigState snapshot() {
         synchronized (lock) {
-            // Doing a shallow copy of the data is safe here because the complex nested data that is copied should all be
-            // immutable configs
+            // Only a shallow copy is performed here; in order to avoid accidentally corrupting the worker's view
+            // of the config topic, any nested structures should be copied before making modifications
             return new ClusterConfigState(
                     offset,
+                    sessionKey,
                     new HashMap<>(connectorTaskCounts),
                     new HashMap<>(connectorConfigs),
                     new HashMap<>(connectorTargetStates),
@@ -409,6 +434,23 @@ public void putTargetState(String connector, TargetState state) {
         configLog.send(TARGET_STATE_KEY(connector), serializedTargetState);
     }
 
+    @Override
+    public void putSessionKey(SessionKey sessionKey) {
+        log.debug("Distributing new session key");
+        Struct sessionKeyStruct = new Struct(SESSION_KEY_V0);
+        sessionKeyStruct.put("key", Base64.getEncoder().encodeToString(sessionKey.key().getEncoded()));
+        sessionKeyStruct.put("algorithm", sessionKey.key().getAlgorithm());
+        sessionKeyStruct.put("creation-timestamp", sessionKey.creationTimestamp());
+        byte[] serializedSessionKey = converter.fromConnectData(topic, SESSION_KEY_V0, sessionKeyStruct);
+        try {
+            configLog.send(SESSION_KEY_KEY, serializedSessionKey);
+            configLog.readToEnd().get(READ_TO_END_TIMEOUT_MS, TimeUnit.MILLISECONDS);
+        } catch (InterruptedException | ExecutionException | TimeoutException e) {
+            log.error("Failed to write session key to Kafka: ", e);
+            throw new ConnectException("Error writing session key to Kafka", e);
+        }
+    }
+
     // package private for testing
     KafkaBasedLog setupAndCreateKafkaBasedLog(String topic, final WorkerConfig config) {
         Map originals = config.originals();
@@ -516,6 +558,8 @@ public void onCompletion(Throwable error, ConsumerRecord record)
                         // Connector deletion will be written as a null value
                         log.info("Successfully processed removal of connector '{}'", connectorName);
                         connectorConfigs.remove(connectorName);
+                        connectorTaskCounts.remove(connectorName);
+                        taskConfigs.keySet().removeIf(taskId -> taskId.connector().equals(connectorName));
                         removed = true;
                     } else {
                         // Connector configs can be applied and callbacks invoked immediately
@@ -620,7 +664,7 @@ public void onCompletion(Throwable error, ConsumerRecord record)
                     } else {
                         if (deferred != null) {
                             taskConfigs.putAll(deferred);
-                            updatedTasks.addAll(taskConfigs.keySet());
+                            updatedTasks.addAll(deferred.keySet());
                         }
                         inconsistent.remove(connectorName);
                     }
@@ -635,6 +679,45 @@ public void onCompletion(Throwable error, ConsumerRecord record)
 
                 if (started)
                     updateListener.onTaskConfigUpdate(updatedTasks);
+            } else if (record.key().equals(SESSION_KEY_KEY)) {
+                synchronized (lock) {
+                    if (value.value() == null) {
+                        log.error("Ignoring session key because it is unexpectedly null");
+                        return;
+                    }
+                    if (!(value.value() instanceof Map)) {
+                        log.error("Ignoring session key because the value is not a Map but is {}", value.value().getClass());
+                        return;
+                    }
+
+                    Map valueAsMap = (Map) value.value();
+
+                    Object sessionKey = valueAsMap.get("key");
+                    if (!(sessionKey instanceof String)) {
+                        log.error("Invalid data for session key 'key' field should be a String but is {}", sessionKey.getClass());
+                        return;
+                    }
+                    byte[] key = Base64.getDecoder().decode((String) sessionKey);
+
+                    Object keyAlgorithm = valueAsMap.get("algorithm");
+                    if (!(keyAlgorithm instanceof String)) {
+                        log.error("Invalid data for session key 'algorithm' field should be a String but it is {}", keyAlgorithm.getClass());
+                        return;
+                    }
+
+                    Object creationTimestamp = valueAsMap.get("creation-timestamp");
+                    if (!(creationTimestamp instanceof Long)) {
+                        log.error("Invalid data for session key 'creation-timestamp' field should be a long but it is {}", creationTimestamp.getClass());
+                        return;
+                    }
+                    KafkaConfigBackingStore.this.sessionKey = new SessionKey(
+                        new SecretKeySpec(key, (String) keyAlgorithm),
+                        (long) creationTimestamp
+                    );
+
+                    if (started)
+                        updateListener.onSessionKeyUpdate(KafkaConfigBackingStore.this.sessionKey);
+                }
             } else {
                 log.error("Discarding config update record with invalid key: {}", record.key());
             }
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 195c498edb76f..22bcb7a5ae2e6 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
@@ -118,9 +118,8 @@ public void stop() {
     }
 
     @Override
-    public Future> get(final Collection keys,
-                                                   final Callback> callback) {
-        ConvertingFutureCallback> future = new ConvertingFutureCallback>(callback) {
+    public Future> get(final Collection keys) {
+        ConvertingFutureCallback> future = new ConvertingFutureCallback>() {
             @Override
             public Map convert(Void result) {
                 Map values = new HashMap<>();
@@ -230,6 +229,4 @@ public synchronized Void get(long timeout, TimeUnit unit) throws InterruptedExce
             return null;
         }
     }
-
-
 }
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/storage/MemoryConfigBackingStore.java b/connect/runtime/src/main/java/org/apache/kafka/connect/storage/MemoryConfigBackingStore.java
index 7e7d62ba3c982..90039b9c01618 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/storage/MemoryConfigBackingStore.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/storage/MemoryConfigBackingStore.java
@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.connect.storage;
 
+import org.apache.kafka.connect.runtime.SessionKey;
 import org.apache.kafka.connect.runtime.TargetState;
 import org.apache.kafka.connect.runtime.WorkerConfigTransformer;
 import org.apache.kafka.connect.runtime.distributed.ClusterConfigState;
@@ -68,6 +69,7 @@ public synchronized ClusterConfigState snapshot() {
 
         return new ClusterConfigState(
                 ClusterConfigState.NO_OFFSET,
+                null,
                 connectorTaskCounts,
                 connectorConfigs,
                 connectorTargetStates,
@@ -143,6 +145,11 @@ public synchronized void putTargetState(String connector, TargetState state) {
             updateListener.onConnectorTargetStateChange(connector);
     }
 
+    @Override
+    public void putSessionKey(SessionKey sessionKey) {
+        // no-op
+    }
+
     @Override
     public synchronized void setUpdateListener(UpdateListener listener) {
         this.updateListener = listener;
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/storage/MemoryOffsetBackingStore.java b/connect/runtime/src/main/java/org/apache/kafka/connect/storage/MemoryOffsetBackingStore.java
index ab8130b6ef776..72439e7d687b3 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/storage/MemoryOffsetBackingStore.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/storage/MemoryOffsetBackingStore.java
@@ -75,9 +75,7 @@ public void stop() {
     }
 
     @Override
-    public Future> get(
-            final Collection keys,
-            final Callback> callback) {
+    public Future> get(final Collection keys) {
         return executor.submit(new Callable>() {
             @Override
             public Map call() throws Exception {
@@ -85,8 +83,6 @@ public Map call() throws Exception {
                 for (ByteBuffer key : keys) {
                     result.put(key, data.get(key));
                 }
-                if (callback != null)
-                    callback.onCompletion(null, result);
                 return result;
             }
         });
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/storage/OffsetBackingStore.java b/connect/runtime/src/main/java/org/apache/kafka/connect/storage/OffsetBackingStore.java
index 9998164ddf5bf..1e4375b7d8eff 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/storage/OffsetBackingStore.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/storage/OffsetBackingStore.java
@@ -53,12 +53,9 @@ public interface OffsetBackingStore {
     /**
      * Get the values for the specified keys
      * @param keys list of keys to look up
-     * @param callback callback to invoke on completion
      * @return future for the resulting map from key to value
      */
-    Future> get(
-            Collection keys,
-            Callback> callback);
+    Future> get(Collection keys);
 
     /**
      * Set the specified keys and values.
@@ -66,8 +63,7 @@ Future> get(
      * @param callback callback to invoke on completion
      * @return void future for the operation
      */
-    Future set(Map values,
-                            Callback callback);
+    Future set(Map values, Callback callback);
 
     /**
      * Configure class with the given key-value pairs
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/storage/OffsetStorageReaderImpl.java b/connect/runtime/src/main/java/org/apache/kafka/connect/storage/OffsetStorageReaderImpl.java
index 9f926dc5040aa..a1eea43103a39 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/storage/OffsetStorageReaderImpl.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/storage/OffsetStorageReaderImpl.java
@@ -26,20 +26,27 @@
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.CancellationException;
+import java.util.concurrent.Future;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 /**
  * Implementation of OffsetStorageReader. Unlike OffsetStorageWriter which is implemented
  * directly, the interface is only separate from this implementation because it needs to be
  * included in the public API package.
  */
-public class OffsetStorageReaderImpl implements OffsetStorageReader {
+public class OffsetStorageReaderImpl implements CloseableOffsetStorageReader {
     private static final Logger log = LoggerFactory.getLogger(OffsetStorageReaderImpl.class);
 
     private final OffsetBackingStore backingStore;
     private final String namespace;
     private final Converter keyConverter;
     private final Converter valueConverter;
+    private final AtomicBoolean closed;
+    private final Set>> offsetReadFutures;
 
     public OffsetStorageReaderImpl(OffsetBackingStore backingStore, String namespace,
                                    Converter keyConverter, Converter valueConverter) {
@@ -47,6 +54,8 @@ public OffsetStorageReaderImpl(OffsetBackingStore backingStore, String namespace
         this.namespace = namespace;
         this.keyConverter = keyConverter;
         this.valueConverter = valueConverter;
+        this.closed = new AtomicBoolean(false);
+        this.offsetReadFutures = new HashSet<>();
     }
 
     @Override
@@ -76,7 +85,30 @@ public  Map, Map> offsets(Collection serialized value from backing store
         Map raw;
         try {
-            raw = backingStore.get(serializedToOriginal.keySet(), null).get();
+            Future> offsetReadFuture;
+            synchronized (offsetReadFutures) {
+                if (closed.get()) {
+                    throw new ConnectException(
+                        "Offset reader is closed. This is likely because the task has already been "
+                            + "scheduled to stop but has taken longer than the graceful shutdown "
+                            + "period to do so.");
+                }
+                offsetReadFuture = backingStore.get(serializedToOriginal.keySet());
+                offsetReadFutures.add(offsetReadFuture);
+            }
+
+            try {
+                raw = offsetReadFuture.get();
+            } catch (CancellationException e) {
+                throw new ConnectException(
+                    "Offset reader closed while attempting to read offsets. This is likely because "
+                        + "the task was been scheduled to stop but has taken longer than the "
+                        + "graceful shutdown period to do so.");
+            } finally {
+                synchronized (offsetReadFutures) {
+                    offsetReadFutures.remove(offsetReadFuture);
+                }
+            }
         } catch (Exception e) {
             log.error("Failed to fetch offsets from namespace {}: ", namespace, e);
             throw new ConnectException("Failed to fetch offsets.", e);
@@ -108,4 +140,19 @@ public  Map, Map> offsets(Collection> offsetReadFuture : offsetReadFutures) {
+                    try {
+                        offsetReadFuture.cancel(true);
+                    } catch (Throwable t) {
+                        log.error("Failed to cancel offset read future", t);
+                    }
+                }
+                offsetReadFutures.clear();
+            }
+        }
+    }
 }
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 b76e7d49b5d5c..77d7728e2434b 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
@@ -74,7 +74,7 @@ private static void printTransformationHtml(PrintStream out, DocInfo docInfo) {
 
         out.println("
      ");
 
-        out.println(docInfo.configDef.toHtmlTable());
+        out.println(docInfo.configDef.toHtml());
 
         out.println("");
     }
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/util/ConnectUtils.java b/connect/runtime/src/main/java/org/apache/kafka/connect/util/ConnectUtils.java
index 9f30236fdee47..e1e4a874035ab 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/util/ConnectUtils.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/util/ConnectUtils.java
@@ -16,9 +16,9 @@
  */
 package org.apache.kafka.connect.util;
 
-import org.apache.kafka.clients.admin.AdminClient;
+import org.apache.kafka.clients.admin.Admin;
 import org.apache.kafka.common.KafkaFuture;
-import org.apache.kafka.common.record.InvalidRecordException;
+import org.apache.kafka.common.InvalidRecordException;
 import org.apache.kafka.common.record.RecordBatch;
 import org.apache.kafka.connect.errors.ConnectException;
 import org.apache.kafka.connect.runtime.WorkerConfig;
@@ -41,12 +41,12 @@ else if (timestamp == RecordBatch.NO_TIMESTAMP)
 
     public static String lookupKafkaClusterId(WorkerConfig config) {
         log.info("Creating Kafka admin client");
-        try (AdminClient adminClient = AdminClient.create(config.originals())) {
+        try (Admin adminClient = Admin.create(config.originals())) {
             return lookupKafkaClusterId(adminClient);
         }
     }
 
-    static String lookupKafkaClusterId(AdminClient adminClient) {
+    static String lookupKafkaClusterId(Admin adminClient) {
         log.debug("Looking up Kafka cluster ID");
         try {
             KafkaFuture clusterIdFuture = adminClient.describeCluster().clusterId();
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/util/ConnectorTaskId.java b/connect/runtime/src/main/java/org/apache/kafka/connect/util/ConnectorTaskId.java
index 03a51f2b25aa1..1b69bd0179550 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/util/ConnectorTaskId.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/util/ConnectorTaskId.java
@@ -20,6 +20,7 @@
 import com.fasterxml.jackson.annotation.JsonProperty;
 
 import java.io.Serializable;
+import java.util.Objects;
 
 /**
  * Unique ID for a single task. It includes a unique connector ID and a task ID that is unique within
@@ -56,10 +57,8 @@ public boolean equals(Object o) {
 
         if (task != that.task)
             return false;
-        if (connector != null ? !connector.equals(that.connector) : that.connector != null)
-            return false;
 
-        return true;
+        return Objects.equals(connector, that.connector);
     }
 
     @Override
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/util/ConvertingFutureCallback.java b/connect/runtime/src/main/java/org/apache/kafka/connect/util/ConvertingFutureCallback.java
index d5abed9385cc0..e15c38ea4c4ae 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/util/ConvertingFutureCallback.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/util/ConvertingFutureCallback.java
@@ -16,6 +16,9 @@
  */
 package org.apache.kafka.connect.util;
 
+import org.apache.kafka.connect.errors.ConnectException;
+
+import java.util.concurrent.CancellationException;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Future;
@@ -24,10 +27,15 @@
 
 public abstract class ConvertingFutureCallback implements Callback, Future {
 
-    private Callback underlying;
-    private CountDownLatch finishedLatch;
-    private T result = null;
-    private Throwable exception = null;
+    private final Callback underlying;
+    private final CountDownLatch finishedLatch;
+    private volatile T result = null;
+    private volatile Throwable exception = null;
+    private volatile boolean cancelled = false;
+
+    public ConvertingFutureCallback() {
+        this(null);
+    }
 
     public ConvertingFutureCallback(Callback underlying) {
         this.underlying = underlying;
@@ -38,21 +46,46 @@ public ConvertingFutureCallback(Callback underlying) {
 
     @Override
     public void onCompletion(Throwable error, U result) {
-        this.exception = error;
-        this.result = convert(result);
-        if (underlying != null)
-            underlying.onCompletion(error, this.result);
-        finishedLatch.countDown();
+        synchronized (this) {
+            if (isDone()) {
+                return;
+            }
+            
+            if (error != null) {
+                this.exception = error;
+            } else {
+                this.result = convert(result);
+            }
+
+            if (underlying != null)
+                underlying.onCompletion(error, this.result);
+            finishedLatch.countDown();
+        }
     }
 
     @Override
-    public boolean cancel(boolean b) {
+    public boolean cancel(boolean mayInterruptIfRunning) {
+        synchronized (this) {
+            if (isDone()) {
+                return false;
+            }
+            if (mayInterruptIfRunning) {
+                this.cancelled = true;
+                finishedLatch.countDown();
+                return true;
+            }
+        }
+        try {
+            finishedLatch.await();
+        } catch (InterruptedException e) {
+            throw new ConnectException("Interrupted while waiting for task to complete", e);
+        }
         return false;
     }
 
     @Override
     public boolean isCancelled() {
-        return false;
+        return cancelled;
     }
 
     @Override
@@ -75,6 +108,9 @@ public T get(long l, TimeUnit timeUnit)
     }
 
     private T result() throws ExecutionException {
+        if (cancelled) {
+            throw new CancellationException();
+        }
         if (exception != null) {
             throw new ExecutionException(exception);
         }
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 9d77d21767a89..69d2588fdea7a 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
@@ -28,6 +28,7 @@
 import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.PartitionInfo;
 import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.errors.TimeoutException;
 import org.apache.kafka.common.errors.WakeupException;
 import org.apache.kafka.common.utils.Time;
 import org.apache.kafka.common.utils.Utils;
@@ -73,6 +74,7 @@ public class KafkaBasedLog {
 
     private Time time;
     private final String topic;
+    private int partitionCount;
     private final Map producerConfigs;
     private final Map consumerConfigs;
     private final Callback> consumedCallback;
@@ -144,6 +146,7 @@ public void start() {
 
         for (PartitionInfo partition : partitionInfos)
             partitions.add(new TopicPartition(partition.topic(), partition.partition()));
+        partitionCount = partitions.size();
         consumer.assign(partitions);
 
         // Always consume from the beginning of all partitions. Necessary to ensure that we don't use committed offsets
@@ -237,6 +240,9 @@ public void send(K key, V value, org.apache.kafka.clients.producer.Callback call
         producer.send(new ProducerRecord<>(topic, key, value), callback);
     }
 
+    public int partitionCount() {
+        return partitionCount;
+    }
 
     private Producer createProducer() {
         // Always require producer acks to all to ensure durable writes
@@ -280,9 +286,15 @@ private void readToLogEnd() {
             Iterator> it = endOffsets.entrySet().iterator();
             while (it.hasNext()) {
                 Map.Entry entry = it.next();
-                if (consumer.position(entry.getKey()) >= entry.getValue())
+                TopicPartition topicPartition = entry.getKey();
+                long endOffset = entry.getValue();
+                long lastConsumedOffset = consumer.position(topicPartition);
+                if (lastConsumedOffset >= endOffset) {
+                    log.trace("Read to end offset {} for {}", endOffset, topicPartition);
                     it.remove();
-                else {
+                } else {
+                    log.trace("Behind end offset {} for {}; last-read offset is {}",
+                            endOffset, topicPartition, lastConsumedOffset);
                     poll(Integer.MAX_VALUE);
                     break;
                 }
@@ -312,6 +324,10 @@ public void run() {
                         try {
                             readToLogEnd();
                             log.trace("Finished read to end log for topic {}", topic);
+                        } catch (TimeoutException e) {
+                            log.warn("Timeout while reading log to end for topic '{}'. Retrying automatically. " +
+                                "This may occur when brokers are unavailable or unreachable. Reason: {}", topic, e.getMessage());
+                            continue;
                         } catch (WakeupException e) {
                             // Either received another get() call and need to retry reading to end of log or stop() was
                             // called. Both are handled by restarting this loop.
@@ -340,4 +356,4 @@ public void run() {
             }
         }
     }
-}
\ No newline at end of file
+}
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/util/LoggingContext.java b/connect/runtime/src/main/java/org/apache/kafka/connect/util/LoggingContext.java
new file mode 100644
index 0000000000000..8df5f9c39d048
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/util/LoggingContext.java
@@ -0,0 +1,227 @@
+/*
+ * 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.util;
+
+import org.slf4j.MDC;
+
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Map;
+import java.util.Objects;
+
+/**
+ * A utility for defining Mapped Diagnostic Context (MDC) for SLF4J logs.
+ *
+ * 
      {@link LoggingContext} instances should be created in a try-with-resources block to ensure
+ * that the logging context is properly closed. The only exception is the logging context created
+ * upon thread creation that is to be used for the entire lifetime of the thread.
+ *
+ * 
      Any logger created on the thread will inherit the MDC context, so this mechanism is ideal for
+ * providing additional information in the log messages without requiring connector
+ * implementations to use a specific Connect API or SLF4J API. {@link LoggingContext#close()}
+ * will also properly restore the Connect MDC parameters to their state just prior to when the
+ * LoggingContext was created. Use {@link #clear()} to remove all MDC parameters from the
+ * current thread context.
+ *
+ * 
      Compare this approach to {@link org.apache.kafka.common.utils.LogContext}, which must be
+ * used to create a new {@link org.slf4j.Logger} instance pre-configured with the desired prefix.
+ * Currently LogContext does not allow the prefix to be changed, and it requires that all
+ * components use the LogContext to create their Logger instance.
+ */
+public final class LoggingContext implements AutoCloseable {
+
+    /**
+     * The name of the Mapped Diagnostic Context (MDC) key that defines the context for a connector.
+     */
+    public static final String CONNECTOR_CONTEXT = "connector.context";
+
+    public static final Collection ALL_CONTEXTS = Collections.singleton(CONNECTOR_CONTEXT);
+
+    /**
+     * The Scope values used by Connect when specifying the context.
+     */
+    public enum Scope {
+        /**
+         * The scope value for the worker as it starts a connector.
+         */
+        WORKER("worker"),
+
+        /**
+         * The scope value for Task implementations.
+         */
+        TASK("task"),
+
+        /**
+         * The scope value for committing offsets.
+         */
+        OFFSETS("offsets"),
+
+        /**
+         * The scope value for validating connector configurations.
+         */
+        VALIDATE("validate");
+
+        private final String text;
+        Scope(String value) {
+            this.text = value;
+        }
+
+        @Override
+        public String toString() {
+            return text;
+        }
+    }
+
+    /**
+     * Clear all MDC parameters.
+     */
+    public static void clear() {
+        MDC.clear();
+    }
+
+    /**
+     * Modify the current {@link MDC} logging context to set the {@link #CONNECTOR_CONTEXT connector context} to include the
+     * supplied name and the {@link Scope#WORKER} scope.
+     *
+     * @param connectorName the connector name; may not be null
+     */
+    public static LoggingContext forConnector(String connectorName) {
+        Objects.requireNonNull(connectorName);
+        LoggingContext context = new LoggingContext();
+        MDC.put(CONNECTOR_CONTEXT, prefixFor(connectorName, Scope.WORKER, null));
+        return context;
+    }
+
+    /**
+     * Modify the current {@link MDC} logging context to set the {@link #CONNECTOR_CONTEXT connector context} to include the
+     * supplied connector name and the {@link Scope#VALIDATE} scope.
+     *
+     * @param connectorName the connector name
+     */
+    public static LoggingContext forValidation(String connectorName) {
+        LoggingContext context = new LoggingContext();
+        MDC.put(CONNECTOR_CONTEXT, prefixFor(connectorName, Scope.VALIDATE, null));
+        return context;
+    }
+
+    /**
+     * Modify the current {@link MDC} logging context to set the {@link #CONNECTOR_CONTEXT connector context} to include the
+     * connector name and task number using the supplied {@link ConnectorTaskId}, and to set the scope to {@link Scope#TASK}.
+     *
+     * @param id the connector task ID; may not be null
+     */
+    public static LoggingContext forTask(ConnectorTaskId id) {
+        Objects.requireNonNull(id);
+        LoggingContext context = new LoggingContext();
+        MDC.put(CONNECTOR_CONTEXT, prefixFor(id.connector(), Scope.TASK, id.task()));
+        return context;
+    }
+
+    /**
+     * Modify the current {@link MDC} logging context to set the {@link #CONNECTOR_CONTEXT connector context} to include the
+     * connector name and task number using the supplied {@link ConnectorTaskId}, and to set the scope to {@link Scope#OFFSETS}.
+     *
+     * @param id the connector task ID; may not be null
+     */
+    public static LoggingContext forOffsets(ConnectorTaskId id) {
+        Objects.requireNonNull(id);
+        LoggingContext context = new LoggingContext();
+        MDC.put(CONNECTOR_CONTEXT, prefixFor(id.connector(), Scope.OFFSETS, id.task()));
+        return context;
+    }
+
+    /**
+     * Return the prefix that uses the specified connector name, task number, and scope. The
+     * format is as follows:
+     *
+     * 
      +     *     [<connectorName>|<scope>]<sp>
+     * 
      
+     *
+     * where "<connectorName>" is the name of the connector,
+     * "<sp>" indicates a trailing space, and
+     * "<scope>" is one of the following:
+     *
+     * 
        
+     *   
      • "task-n" for the operation of the numbered task, including calling the
+     *      task methods and the producer/consumer; here "n" is the 0-based task number
+     *   
      • "task-n|offset" for the committing of source offsets for the numbered
+     *       task; here "n" is the * zero-based task number
+     *   
      • "worker" for the creation and usage of connector instances
+     * 
      
+     *
+     * 
      The following are examples of the connector context for a connector named "my-connector":
+     *
+     * 
        
+     *   
      • `[my-connector|worker]` - used on log messages where the Connect worker is
+     *     validating the configuration for or starting/stopping the "local-file-source" connector
+     *     via the SourceConnector / SinkConnector implementation methods.
+     *   
      • `[my-connector|task-0]` - used on log messages where the Connect worker is executing
+     *     task 0 of the "local-file-source" connector, including calling any of the SourceTask /
+     *     SinkTask implementation methods, processing the messages for/from the task, and
+     *     calling the task's * producer/consumer.
+     *   
      • `[my-connector|task-0|offsets]` - used on log messages where the Connect worker is
+     *       committing * source offsets for task 0 of the "local-file-source" connector.
+     * 
      
+     *
+     * @param connectorName the name of the connector; may not be null
+     * @param scope the scope; may not be null
+     * @param taskNumber the 0-based task number; may be null if there is no associated task
+     * @return the prefix; never null
+     */
+    protected static String prefixFor(String connectorName, Scope scope, Integer taskNumber) {
+        StringBuilder sb = new StringBuilder();
+        sb.append("[");
+        sb.append(connectorName);
+        if (taskNumber != null) {
+            // There is a task number, so this is a task
+            sb.append("|");
+            sb.append(Scope.TASK.toString());
+            sb.append("-");
+            sb.append(taskNumber.toString());
+        }
+        // Append non-task scopes (e.g., worker and offset)
+        if (scope != Scope.TASK) {
+            sb.append("|");
+            sb.append(scope.toString());
+        }
+        sb.append("] ");
+        return sb.toString();
+    }
+
+    private final Map previous;
+
+    private LoggingContext() {
+        previous = MDC.getCopyOfContextMap(); // may be null!
+    }
+
+    /**
+     * Close this logging context, restoring the Connect {@link MDC} parameters back to the state
+     * just before this context was created. This does not affect other MDC parameters set by
+     * connectors or tasks.
+     */
+    @Override
+    public void close() {
+        for (String param : ALL_CONTEXTS) {
+            if (previous != null && previous.containsKey(param)) {
+                MDC.put(param, previous.get(param));
+            } else {
+                MDC.remove(param);
+            }
+        }
+    }
+}
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/util/TopicAdmin.java b/connect/runtime/src/main/java/org/apache/kafka/connect/util/TopicAdmin.java
index ad21561baf259..e644e805ea613 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/util/TopicAdmin.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/util/TopicAdmin.java
@@ -16,12 +16,13 @@
  */
 package org.apache.kafka.connect.util;
 
-import org.apache.kafka.clients.admin.AdminClient;
+import org.apache.kafka.clients.admin.Admin;
 import org.apache.kafka.clients.admin.AdminClientConfig;
 import org.apache.kafka.clients.admin.CreateTopicsOptions;
 import org.apache.kafka.clients.admin.NewTopic;
 import org.apache.kafka.common.KafkaFuture;
 import org.apache.kafka.common.errors.ClusterAuthorizationException;
+import org.apache.kafka.common.errors.TopicAuthorizationException;
 import org.apache.kafka.common.errors.TimeoutException;
 import org.apache.kafka.common.errors.TopicExistsException;
 import org.apache.kafka.common.errors.UnsupportedVersionException;
@@ -38,7 +39,7 @@
 import java.util.concurrent.ExecutionException;
 
 /**
- * Utility to simplify creating and managing topics via the {@link org.apache.kafka.clients.admin.AdminClient}.
+ * Utility to simplify creating and managing topics via the {@link Admin}.
  */
 public class TopicAdmin implements AutoCloseable {
 
@@ -157,19 +158,19 @@ public static NewTopicBuilder defineTopic(String topicName) {
 
     private static final Logger log = LoggerFactory.getLogger(TopicAdmin.class);
     private final Map adminConfig;
-    private final AdminClient admin;
+    private final Admin admin;
 
     /**
      * Create a new topic admin component with the given configuration.
      *
-     * @param adminConfig the configuration for the {@link AdminClient}
+     * @param adminConfig the configuration for the {@link Admin}
      */
     public TopicAdmin(Map adminConfig) {
-        this(adminConfig, AdminClient.create(adminConfig));
+        this(adminConfig, Admin.create(adminConfig));
     }
 
     // visible for testing
-    TopicAdmin(Map adminConfig, AdminClient adminClient) {
+    TopicAdmin(Map adminConfig, Admin adminClient) {
         this.admin = adminClient;
         this.adminConfig = adminConfig != null ? adminConfig : Collections.emptyMap();
     }
@@ -235,17 +236,23 @@ public Set createTopics(NewTopic... topics) {
                     continue;
                 }
                 if (cause instanceof UnsupportedVersionException) {
-                    log.debug("Unable to create topic(s) '{}' since the brokers at {} do not support the CreateTopics API.",
+                    log.debug("Unable to create topic(s) '{}' since the brokers at {} do not support the CreateTopics API." +
                             " Falling back to assume topic(s) exist or will be auto-created by the broker.",
                             topicNameList, bootstrapServers);
                     return Collections.emptySet();
                 }
                 if (cause instanceof ClusterAuthorizationException) {
-                    log.debug("Not authorized to create topic(s) '{}'." +
+                    log.debug("Not authorized to create topic(s) '{}' upon the brokers {}." +
                             " Falling back to assume topic(s) exist or will be auto-created by the broker.",
                             topicNameList, bootstrapServers);
                     return Collections.emptySet();
                 }
+                if (cause instanceof TopicAuthorizationException) {
+                    log.debug("Not authorized to create topic(s) '{}' upon the brokers {}." +
+                                    " Falling back to assume topic(s) exist or will be auto-created by the broker.",
+                            topicNameList, bootstrapServers);
+                    return Collections.emptySet();
+                }
                 if (cause instanceof TimeoutException) {
                     // Timed out waiting for the operation to complete
                     throw new ConnectException("Timed out while checking for or creating topic(s) '" + topicNameList + "'." +
diff --git a/connect/runtime/src/main/resources/META-INF/services/org.apache.kafka.connect.connector.policy.ConnectorClientConfigOverridePolicy b/connect/runtime/src/main/resources/META-INF/services/org.apache.kafka.connect.connector.policy.ConnectorClientConfigOverridePolicy
new file mode 100644
index 0000000000000..8b76ce452b659
--- /dev/null
+++ b/connect/runtime/src/main/resources/META-INF/services/org.apache.kafka.connect.connector.policy.ConnectorClientConfigOverridePolicy
@@ -0,0 +1,18 @@
+ # 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.
+
+org.apache.kafka.connect.connector.policy.AllConnectorClientConfigOverridePolicy
+org.apache.kafka.connect.connector.policy.PrincipalConnectorClientConfigOverridePolicy
+org.apache.kafka.connect.connector.policy.NoneConnectorClientConfigOverridePolicy
\ No newline at end of file
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/connector/policy/BaseConnectorClientConfigOverridePolicyTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/connector/policy/BaseConnectorClientConfigOverridePolicyTest.java
new file mode 100644
index 0000000000000..28fee73a93966
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/connector/policy/BaseConnectorClientConfigOverridePolicyTest.java
@@ -0,0 +1,59 @@
+/*
+ * 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.connector.policy;
+
+import org.apache.kafka.common.config.ConfigValue;
+import org.apache.kafka.connect.health.ConnectorType;
+import org.apache.kafka.connect.runtime.WorkerTest;
+import org.junit.Assert;
+
+import java.util.List;
+import java.util.Map;
+
+public abstract class BaseConnectorClientConfigOverridePolicyTest {
+
+    protected abstract ConnectorClientConfigOverridePolicy  policyToTest();
+
+    protected void testValidOverride(Map clientConfig) {
+        List configValues = configValues(clientConfig);
+        assertNoError(configValues);
+    }
+
+    protected void testInvalidOverride(Map clientConfig) {
+        List configValues = configValues(clientConfig);
+        assertError(configValues);
+    }
+
+    private List configValues(Map clientConfig) {
+        ConnectorClientConfigRequest connectorClientConfigRequest = new ConnectorClientConfigRequest(
+            "test",
+            ConnectorType.SOURCE,
+            WorkerTest.WorkerTestConnector.class,
+            clientConfig,
+            ConnectorClientConfigRequest.ClientType.PRODUCER);
+        return policyToTest().validate(connectorClientConfigRequest);
+    }
+
+    protected void assertNoError(List configValues) {
+        Assert.assertTrue(configValues.stream().allMatch(configValue -> configValue.errorMessages().size() == 0));
+    }
+
+    protected void assertError(List configValues) {
+        Assert.assertTrue(configValues.stream().anyMatch(configValue -> configValue.errorMessages().size() > 0));
+    }
+}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/connector/policy/NoneConnectorClientConfigOverridePolicyTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/connector/policy/NoneConnectorClientConfigOverridePolicyTest.java
new file mode 100644
index 0000000000000..2c7b0789d8022
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/connector/policy/NoneConnectorClientConfigOverridePolicyTest.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.connect.connector.policy;
+
+import org.apache.kafka.clients.producer.ProducerConfig;
+import org.apache.kafka.common.config.SaslConfigs;
+import org.junit.Test;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
+public class NoneConnectorClientConfigOverridePolicyTest extends BaseConnectorClientConfigOverridePolicyTest {
+
+    ConnectorClientConfigOverridePolicy noneConnectorClientConfigOverridePolicy = new NoneConnectorClientConfigOverridePolicy();
+
+    @Test
+    public void testNoOverrides() {
+        testValidOverride(Collections.emptyMap());
+    }
+
+    @Test
+    public void testWithOverrides() {
+        Map clientConfig = new HashMap<>();
+        clientConfig.put(SaslConfigs.SASL_JAAS_CONFIG, "test");
+        clientConfig.put(ProducerConfig.ACKS_CONFIG, "none");
+        testInvalidOverride(clientConfig);
+    }
+
+    @Override
+    protected ConnectorClientConfigOverridePolicy policyToTest() {
+        return noneConnectorClientConfigOverridePolicy;
+    }
+}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/connector/policy/PrincipalConnectorClientConfigOverridePolicyTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/connector/policy/PrincipalConnectorClientConfigOverridePolicyTest.java
new file mode 100644
index 0000000000000..0e79c8a6ca990
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/connector/policy/PrincipalConnectorClientConfigOverridePolicyTest.java
@@ -0,0 +1,50 @@
+/*
+ * 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.connector.policy;
+
+import org.apache.kafka.clients.producer.ProducerConfig;
+import org.apache.kafka.common.config.SaslConfigs;
+import org.junit.Test;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
+public class PrincipalConnectorClientConfigOverridePolicyTest extends BaseConnectorClientConfigOverridePolicyTest {
+
+    ConnectorClientConfigOverridePolicy principalConnectorClientConfigOverridePolicy = new PrincipalConnectorClientConfigOverridePolicy();
+
+    @Test
+    public void testPrincipalOnly() {
+        Map clientConfig = Collections.singletonMap(SaslConfigs.SASL_JAAS_CONFIG, "test");
+        testValidOverride(clientConfig);
+    }
+
+    @Test
+    public void testPrincipalPlusOtherConfigs() {
+        Map clientConfig = new HashMap<>();
+        clientConfig.put(SaslConfigs.SASL_JAAS_CONFIG, "test");
+        clientConfig.put(ProducerConfig.ACKS_CONFIG, "none");
+        testInvalidOverride(clientConfig);
+    }
+
+    @Override
+    protected ConnectorClientConfigOverridePolicy policyToTest() {
+        return principalConnectorClientConfigOverridePolicy;
+    }
+}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ConnectIntegrationTestUtils.java b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ConnectIntegrationTestUtils.java
new file mode 100644
index 0000000000000..058dbe206fe0d
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ConnectIntegrationTestUtils.java
@@ -0,0 +1,43 @@
+/*
+ * 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.integration;
+
+import org.junit.rules.TestRule;
+import org.junit.rules.TestWatcher;
+import org.junit.runner.Description;
+import org.slf4j.Logger;
+
+/**
+ * A utility class for Connect's integration tests
+ */
+public class ConnectIntegrationTestUtils {
+    public static TestRule newTestWatcher(Logger log) {
+        return new TestWatcher() {
+            @Override
+            protected void starting(Description description) {
+                super.starting(description);
+                log.info("Starting test {}", description.getMethodName());
+            }
+
+            @Override
+            protected void finished(Description description) {
+                super.finished(description);
+                log.info("Finished test {}", description.getMethodName());
+            }
+        };
+    }
+}
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
new file mode 100644
index 0000000000000..99c8c4c9a7ed5
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ConnectWorkerIntegrationTest.java
@@ -0,0 +1,295 @@
+/*
+ * 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.integration;
+
+import org.apache.kafka.connect.runtime.distributed.DistributedConfig;
+import org.apache.kafka.connect.storage.StringConverter;
+import org.apache.kafka.connect.util.clusters.EmbeddedConnectCluster;
+import org.apache.kafka.connect.util.clusters.WorkerHandle;
+import org.apache.kafka.test.IntegrationTest;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.experimental.categories.Category;
+import org.junit.rules.TestRule;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Properties;
+import java.util.Set;
+import java.util.concurrent.TimeUnit;
+
+import static org.apache.kafka.clients.CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG;
+import static org.apache.kafka.connect.runtime.ConnectorConfig.CONNECTOR_CLASS_CONFIG;
+import static org.apache.kafka.connect.runtime.ConnectorConfig.CONNECTOR_CLIENT_PRODUCER_OVERRIDES_PREFIX;
+import static org.apache.kafka.connect.runtime.ConnectorConfig.KEY_CONVERTER_CLASS_CONFIG;
+import static org.apache.kafka.connect.runtime.ConnectorConfig.TASKS_MAX_CONFIG;
+import static org.apache.kafka.connect.runtime.ConnectorConfig.VALUE_CONVERTER_CLASS_CONFIG;
+import static org.apache.kafka.connect.runtime.WorkerConfig.CONNECTOR_CLIENT_POLICY_CLASS_CONFIG;
+import static org.apache.kafka.connect.runtime.WorkerConfig.OFFSET_COMMIT_INTERVAL_MS_CONFIG;
+import static org.apache.kafka.connect.util.clusters.EmbeddedConnectClusterAssertions.CONNECTOR_SETUP_DURATION_MS;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+/**
+ * Test simple operations on the workers of a Connect cluster.
+ */
+@Category(IntegrationTest.class)
+public class ConnectWorkerIntegrationTest {
+    private static final Logger log = LoggerFactory.getLogger(ConnectWorkerIntegrationTest.class);
+
+    private static final int NUM_TOPIC_PARTITIONS = 3;
+    private static final long OFFSET_COMMIT_INTERVAL_MS = TimeUnit.SECONDS.toMillis(30);
+    private static final int NUM_WORKERS = 3;
+    private static final String CONNECTOR_NAME = "simple-source";
+
+    private EmbeddedConnectCluster.Builder connectBuilder;
+    private EmbeddedConnectCluster connect;
+    Map workerProps = new HashMap<>();
+    Properties brokerProps = new Properties();
+
+    @Rule
+    public TestRule watcher = ConnectIntegrationTestUtils.newTestWatcher(log);
+
+    @Before
+    public void setup() {
+        // setup Connect worker properties
+        workerProps.put(OFFSET_COMMIT_INTERVAL_MS_CONFIG, String.valueOf(OFFSET_COMMIT_INTERVAL_MS));
+        workerProps.put(CONNECTOR_CLIENT_POLICY_CLASS_CONFIG, "All");
+
+        // setup Kafka broker properties
+        brokerProps.put("auto.create.topics.enable", String.valueOf(false));
+
+        // build a Connect cluster backed by Kafka and Zk
+        connectBuilder = new EmbeddedConnectCluster.Builder()
+                .name("connect-cluster")
+                .numWorkers(NUM_WORKERS)
+                .workerProps(workerProps)
+                .brokerProps(brokerProps)
+                .maskExitProcedures(true); // true is the default, setting here as example
+    }
+
+    @After
+    public void close() {
+        // stop all Connect, Kafka and Zk threads.
+        connect.stop();
+    }
+
+    /**
+     * Simple test case to add and then remove a worker from the embedded Connect cluster while
+     * running a simple source connector.
+     */
+    @Test
+    public void testAddAndRemoveWorker() throws Exception {
+        connect = connectBuilder.build();
+        // start the clusters
+        connect.start();
+
+        int numTasks = 4;
+        // create test topic
+        connect.kafka().createTopic("test-topic", NUM_TOPIC_PARTITIONS);
+
+        // set up props for the source connector
+        Map props = new HashMap<>();
+        props.put(CONNECTOR_CLASS_CONFIG, MonitorableSourceConnector.class.getSimpleName());
+        props.put(TASKS_MAX_CONFIG, String.valueOf(numTasks));
+        props.put("throughput", String.valueOf(1));
+        props.put("messages.per.poll", String.valueOf(10));
+        props.put(KEY_CONVERTER_CLASS_CONFIG, StringConverter.class.getName());
+        props.put(VALUE_CONVERTER_CLASS_CONFIG, StringConverter.class.getName());
+
+        connect.assertions().assertAtLeastNumWorkersAreUp(NUM_WORKERS,
+                "Initial group of workers did not start in time.");
+
+        // start a source connector
+        connect.configureConnector(CONNECTOR_NAME, props);
+
+        connect.assertions().assertConnectorAndAtLeastNumTasksAreRunning(CONNECTOR_NAME, numTasks,
+                "Connector tasks did not start in time.");
+
+        WorkerHandle extraWorker = connect.addWorker();
+
+        connect.assertions().assertAtLeastNumWorkersAreUp(NUM_WORKERS + 1,
+                "Expanded group of workers did not start in time.");
+
+        connect.assertions().assertConnectorAndAtLeastNumTasksAreRunning(CONNECTOR_NAME, numTasks,
+                "Connector tasks are not all in running state.");
+
+        Set workers = connect.activeWorkers();
+        assertTrue(workers.contains(extraWorker));
+
+        connect.removeWorker(extraWorker);
+
+        connect.assertions().assertExactlyNumWorkersAreUp(NUM_WORKERS,
+                "Group of workers did not shrink in time.");
+
+        workers = connect.activeWorkers();
+        assertFalse(workers.contains(extraWorker));
+    }
+
+    /**
+     * Verify that a failed task can be restarted successfully.
+     */
+    @Test
+    public void testRestartFailedTask() throws Exception {
+        connect = connectBuilder.build();
+        // start the clusters
+        connect.start();
+
+        int numTasks = 1;
+
+        // Properties for the source connector. The task should fail at startup due to the bad broker address.
+        Map connectorProps = new HashMap<>();
+        connectorProps.put(CONNECTOR_CLASS_CONFIG, MonitorableSourceConnector.class.getName());
+        connectorProps.put(TASKS_MAX_CONFIG, Objects.toString(numTasks));
+        connectorProps.put(CONNECTOR_CLIENT_PRODUCER_OVERRIDES_PREFIX + BOOTSTRAP_SERVERS_CONFIG, "nobrokerrunningatthisaddress");
+
+        connect.assertions().assertExactlyNumWorkersAreUp(NUM_WORKERS,
+                "Initial group of workers did not start in time.");
+
+        // Try to start the connector and its single task.
+        connect.configureConnector(CONNECTOR_NAME, connectorProps);
+
+        connect.assertions().assertConnectorIsRunningAndTasksHaveFailed(CONNECTOR_NAME, numTasks,
+                "Connector tasks did not fail in time");
+
+        // Reconfigure the connector without the bad broker address.
+        connectorProps.remove(CONNECTOR_CLIENT_PRODUCER_OVERRIDES_PREFIX + BOOTSTRAP_SERVERS_CONFIG);
+        connect.configureConnector(CONNECTOR_NAME, connectorProps);
+
+        // Restart the failed task
+        String taskRestartEndpoint = connect.endpointForResource(
+            String.format("connectors/%s/tasks/0/restart", CONNECTOR_NAME));
+        connect.requestPost(taskRestartEndpoint, "", Collections.emptyMap());
+
+        // Ensure the task started successfully this time
+        connect.assertions().assertConnectorAndAtLeastNumTasksAreRunning(CONNECTOR_NAME, numTasks,
+            "Connector tasks are not all in running state.");
+    }
+
+    /**
+     * Verify that a set of tasks restarts correctly after a broker goes offline and back online
+     */
+    @Test
+    public void testBrokerCoordinator() throws Exception {
+        ConnectorHandle connectorHandle = RuntimeHandles.get().connectorHandle(CONNECTOR_NAME);
+        workerProps.put(DistributedConfig.SCHEDULED_REBALANCE_MAX_DELAY_MS_CONFIG, String.valueOf(5000));
+        connect = connectBuilder.workerProps(workerProps).build();
+        // start the clusters
+        connect.start();
+        int numTasks = 4;
+        // create test topic
+        connect.kafka().createTopic("test-topic", NUM_TOPIC_PARTITIONS);
+
+        // set up props for the source connector
+        Map props = new HashMap<>();
+        props.put(CONNECTOR_CLASS_CONFIG, MonitorableSourceConnector.class.getSimpleName());
+        props.put(TASKS_MAX_CONFIG, String.valueOf(numTasks));
+        props.put("topic", "test-topic");
+        props.put("throughput", String.valueOf(1));
+        props.put("messages.per.poll", String.valueOf(10));
+        props.put(KEY_CONVERTER_CLASS_CONFIG, StringConverter.class.getName());
+        props.put(VALUE_CONVERTER_CLASS_CONFIG, StringConverter.class.getName());
+
+        connect.assertions().assertAtLeastNumWorkersAreUp(NUM_WORKERS,
+                "Initial group of workers did not start in time.");
+
+        // start a source connector
+        connect.configureConnector(CONNECTOR_NAME, props);
+
+        connect.assertions().assertConnectorAndAtLeastNumTasksAreRunning(CONNECTOR_NAME, numTasks,
+                "Connector tasks did not start in time.");
+
+        // expect that the connector will be stopped once the coordinator is detected to be down
+        StartAndStopLatch stopLatch = connectorHandle.expectedStops(1, false);
+
+        connect.kafka().stopOnlyKafka();
+
+        connect.assertions().assertExactlyNumWorkersAreUp(NUM_WORKERS,
+                "Group of workers did not remain the same after broker shutdown");
+
+        // Allow for the workers to discover that the coordinator is unavailable, wait is
+        // heartbeat timeout * 2 + 4sec
+        Thread.sleep(TimeUnit.SECONDS.toMillis(10));
+
+        // Wait for the connector to be stopped
+        assertTrue("Failed to stop connector and tasks after coordinator failure within "
+                        + CONNECTOR_SETUP_DURATION_MS + "ms",
+                stopLatch.await(CONNECTOR_SETUP_DURATION_MS, TimeUnit.MILLISECONDS));
+
+        StartAndStopLatch startLatch = connectorHandle.expectedStarts(1, false);
+        connect.kafka().startOnlyKafkaOnSamePorts();
+
+        // Allow for the kafka brokers to come back online
+        Thread.sleep(TimeUnit.SECONDS.toMillis(10));
+
+        connect.assertions().assertExactlyNumWorkersAreUp(NUM_WORKERS,
+                "Group of workers did not remain the same within the designated time.");
+
+        // Allow for the workers to rebalance and reach a steady state
+        Thread.sleep(TimeUnit.SECONDS.toMillis(10));
+
+        connect.assertions().assertConnectorAndAtLeastNumTasksAreRunning(CONNECTOR_NAME, numTasks,
+                "Connector tasks did not start in time.");
+
+        // Expect that the connector has started again
+        assertTrue("Failed to stop connector and tasks after coordinator failure within "
+                        + CONNECTOR_SETUP_DURATION_MS + "ms",
+                startLatch.await(CONNECTOR_SETUP_DURATION_MS, TimeUnit.MILLISECONDS));
+    }
+
+    /**
+     * Verify that the number of tasks listed in the REST API is updated correctly after changes to
+     * the "tasks.max" connector configuration.
+     */
+    @Test
+    public void testTaskStatuses() throws Exception {
+        connect = connectBuilder.build();
+        // start the clusters
+        connect.start();
+
+        connect.assertions().assertAtLeastNumWorkersAreUp(NUM_WORKERS,
+            "Initial group of workers did not start in time.");
+
+        // base connector props
+        Map connectorProps = new HashMap<>();
+        connectorProps.put(CONNECTOR_CLASS_CONFIG, MonitorableSourceConnector.class.getSimpleName());
+
+        // start the connector with only one task
+        final int initialNumTasks = 1;
+        connectorProps.put(TASKS_MAX_CONFIG, String.valueOf(initialNumTasks));
+        connect.configureConnector(CONNECTOR_NAME, connectorProps);
+        connect.assertions().assertConnectorAndExactlyNumTasksAreRunning(CONNECTOR_NAME, initialNumTasks, "Connector tasks did not start in time");
+
+        // then reconfigure it to use more tasks
+        final int increasedNumTasks = 5;
+        connectorProps.put(TASKS_MAX_CONFIG, String.valueOf(increasedNumTasks));
+        connect.configureConnector(CONNECTOR_NAME, connectorProps);
+        connect.assertions().assertConnectorAndExactlyNumTasksAreRunning(CONNECTOR_NAME, increasedNumTasks, "Connector task statuses did not update in time.");
+
+        // then reconfigure it to use fewer tasks
+        final int decreasedNumTasks = 3;
+        connectorProps.put(TASKS_MAX_CONFIG, String.valueOf(decreasedNumTasks));
+        connect.configureConnector(CONNECTOR_NAME, connectorProps);
+        connect.assertions().assertConnectorAndExactlyNumTasksAreRunning(CONNECTOR_NAME, decreasedNumTasks, "Connector task statuses did not update in time.");
+    }
+}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ConnectorCientPolicyIntegrationTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ConnectorCientPolicyIntegrationTest.java
new file mode 100644
index 0000000000000..00f541beee6a7
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ConnectorCientPolicyIntegrationTest.java
@@ -0,0 +1,149 @@
+/*
+ * 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.integration;
+
+import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.common.config.SaslConfigs;
+import org.apache.kafka.connect.runtime.ConnectorConfig;
+import org.apache.kafka.connect.runtime.WorkerConfig;
+import org.apache.kafka.connect.runtime.rest.errors.ConnectRestException;
+import org.apache.kafka.connect.storage.StringConverter;
+import org.apache.kafka.connect.util.clusters.EmbeddedConnectCluster;
+import org.apache.kafka.test.IntegrationTest;
+import org.junit.After;
+import org.junit.Test;
+import org.junit.experimental.categories.Category;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Properties;
+
+import static org.apache.kafka.connect.runtime.ConnectorConfig.CONNECTOR_CLASS_CONFIG;
+import static org.apache.kafka.connect.runtime.ConnectorConfig.KEY_CONVERTER_CLASS_CONFIG;
+import static org.apache.kafka.connect.runtime.ConnectorConfig.TASKS_MAX_CONFIG;
+import static org.apache.kafka.connect.runtime.ConnectorConfig.VALUE_CONVERTER_CLASS_CONFIG;
+import static org.apache.kafka.connect.runtime.SinkConnectorConfig.TOPICS_CONFIG;
+import static org.apache.kafka.connect.runtime.WorkerConfig.OFFSET_COMMIT_INTERVAL_MS_CONFIG;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.fail;
+
+@Category(IntegrationTest.class)
+public class ConnectorCientPolicyIntegrationTest {
+
+    private static final int NUM_TASKS = 1;
+    private static final int NUM_WORKERS = 1;
+    private static final String CONNECTOR_NAME = "simple-conn";
+
+    @After
+    public void close() {
+    }
+
+    @Test
+    public void testCreateWithOverridesForNonePolicy() throws Exception {
+        Map props = basicConnectorConfig();
+        props.put(ConnectorConfig.CONNECTOR_CLIENT_CONSUMER_OVERRIDES_PREFIX + SaslConfigs.SASL_JAAS_CONFIG, "sasl");
+        assertFailCreateConnector("None", props);
+    }
+
+    @Test
+    public void testCreateWithNotAllowedOverridesForPrincipalPolicy() throws Exception {
+        Map props = basicConnectorConfig();
+        props.put(ConnectorConfig.CONNECTOR_CLIENT_CONSUMER_OVERRIDES_PREFIX + SaslConfigs.SASL_JAAS_CONFIG, "sasl");
+        props.put(ConnectorConfig.CONNECTOR_CLIENT_CONSUMER_OVERRIDES_PREFIX + ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "latest");
+        assertFailCreateConnector("Principal", props);
+    }
+
+    @Test
+    public void testCreateWithAllowedOverridesForPrincipalPolicy() throws Exception {
+        Map props = basicConnectorConfig();
+        props.put(ConnectorConfig.CONNECTOR_CLIENT_CONSUMER_OVERRIDES_PREFIX + CommonClientConfigs.SECURITY_PROTOCOL_CONFIG, "PLAINTEXT");
+        assertPassCreateConnector("Principal", props);
+    }
+
+    @Test
+    public void testCreateWithAllowedOverridesForAllPolicy() throws Exception {
+        // setup up props for the sink connector
+        Map props = basicConnectorConfig();
+        props.put(ConnectorConfig.CONNECTOR_CLIENT_CONSUMER_OVERRIDES_PREFIX + CommonClientConfigs.CLIENT_ID_CONFIG, "test");
+        assertPassCreateConnector("All", props);
+    }
+
+    private EmbeddedConnectCluster connectClusterWithPolicy(String policy) throws InterruptedException {
+        // setup Connect worker properties
+        Map workerProps = new HashMap<>();
+        workerProps.put(OFFSET_COMMIT_INTERVAL_MS_CONFIG, String.valueOf(5_000));
+        workerProps.put(WorkerConfig.CONNECTOR_CLIENT_POLICY_CLASS_CONFIG, policy);
+
+        // setup Kafka broker properties
+        Properties exampleBrokerProps = new Properties();
+        exampleBrokerProps.put("auto.create.topics.enable", "false");
+
+        // build a Connect cluster backed by Kafka and Zk
+        EmbeddedConnectCluster connect = new EmbeddedConnectCluster.Builder()
+            .name("connect-cluster")
+            .numWorkers(NUM_WORKERS)
+            .numBrokers(1)
+            .workerProps(workerProps)
+            .brokerProps(exampleBrokerProps)
+            .build();
+
+        // start the clusters
+        connect.start();
+        connect.assertions().assertAtLeastNumWorkersAreUp(NUM_WORKERS,
+                "Initial group of workers did not start in time.");
+
+        return connect;
+    }
+
+    private void assertFailCreateConnector(String policy, Map props) throws InterruptedException {
+        EmbeddedConnectCluster connect = connectClusterWithPolicy(policy);
+        try {
+            connect.configureConnector(CONNECTOR_NAME, props);
+            fail("Shouldn't be able to create connector");
+        } catch (ConnectRestException e) {
+            assertEquals(e.statusCode(), 400);
+        } finally {
+            connect.stop();
+        }
+    }
+
+    private void assertPassCreateConnector(String policy, Map props) throws InterruptedException {
+        EmbeddedConnectCluster connect = connectClusterWithPolicy(policy);
+        try {
+            connect.configureConnector(CONNECTOR_NAME, props);
+            connect.assertions().assertConnectorAndAtLeastNumTasksAreRunning(CONNECTOR_NAME, NUM_TASKS,
+                    "Connector tasks did not start in time.");
+        } catch (ConnectRestException e) {
+            fail("Should be able to create connector");
+        } finally {
+            connect.stop();
+        }
+    }
+
+
+    public Map basicConnectorConfig() {
+        Map props = new HashMap<>();
+        props.put(CONNECTOR_CLASS_CONFIG, MonitorableSinkConnector.class.getSimpleName());
+        props.put(TASKS_MAX_CONFIG, String.valueOf(NUM_TASKS));
+        props.put(TOPICS_CONFIG, "test-topic");
+        props.put(KEY_CONVERTER_CLASS_CONFIG, StringConverter.class.getName());
+        props.put(VALUE_CONVERTER_CLASS_CONFIG, StringConverter.class.getName());
+        return props;
+    }
+
+}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ConnectorHandle.java b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ConnectorHandle.java
index e59691b843d09..a4a461288301b 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ConnectorHandle.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ConnectorHandle.java
@@ -21,10 +21,13 @@
 import org.slf4j.LoggerFactory;
 
 import java.util.Collection;
+import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
 
 /**
  * A handle to a connector executing in a Connect cluster.
@@ -35,9 +38,12 @@ public class ConnectorHandle {
 
     private final String connectorName;
     private final Map taskHandles = new ConcurrentHashMap<>();
+    private final StartAndStopCounter startAndStopCounter = new StartAndStopCounter();
 
     private CountDownLatch recordsRemainingLatch;
+    private CountDownLatch recordsToCommitLatch;
     private int expectedRecords = -1;
+    private int expectedCommits = -1;
 
     public ConnectorHandle(String connectorName) {
         this.connectorName = connectorName;
@@ -54,6 +60,20 @@ public TaskHandle taskHandle(String taskId) {
         return taskHandles.computeIfAbsent(taskId, k -> new TaskHandle(this, taskId));
     }
 
+    /**
+     * Get the connector's name corresponding to this handle.
+     *
+     * @return the connector's name
+     */
+    public String name() {
+        return connectorName;
+    }
+
+    /**
+     * Get the list of tasks handles monitored by this connector handle.
+     *
+     * @return the task handle list
+     */
     public Collection tasks() {
         return taskHandles.values();
     }
@@ -69,13 +89,23 @@ public void deleteTask(String taskId) {
     }
 
     /**
-     * Set the number of expected records for this task.
+     * Set the number of expected records for this connector.
+     *
+     * @param expected number of records
+     */
+    public void expectedRecords(int expected) {
+        expectedRecords = expected;
+        recordsRemainingLatch = new CountDownLatch(expected);
+    }
+
+    /**
+     * Set the number of expected commits performed by this connector.
      *
-     * @param expectedRecords number of records
+     * @param expected number of commits
      */
-    public void expectedRecords(int expectedRecords) {
-        this.expectedRecords = expectedRecords;
-        this.recordsRemainingLatch = new CountDownLatch(expectedRecords);
+    public void expectedCommits(int expected) {
+        expectedCommits = expected;
+        recordsToCommitLatch = new CountDownLatch(expected);
     }
 
     /**
@@ -88,25 +118,208 @@ public void record() {
     }
 
     /**
-     * Wait for this task to receive the expected number of records.
+     * Record arrival of a batch of messages at the connector.
+     *
+     * @param batchSize the number of messages
+     */
+    public void record(int batchSize) {
+        if (recordsRemainingLatch != null) {
+            IntStream.range(0, batchSize).forEach(i -> recordsRemainingLatch.countDown());
+        }
+    }
+
+    /**
+     * Record a message commit from the connector.
+     */
+    public void commit() {
+        if (recordsToCommitLatch != null) {
+            recordsToCommitLatch.countDown();
+        }
+    }
+
+    /**
+     * Record commit on a batch of messages from the connector.
+     *
+     * @param batchSize the number of messages
+     */
+    public void commit(int batchSize) {
+        if (recordsToCommitLatch != null) {
+            IntStream.range(0, batchSize).forEach(i -> recordsToCommitLatch.countDown());
+        }
+    }
+
+    /**
+     * Wait for this connector to meet the expected number of records as defined by {@code
+     * expectedRecords}.
      *
-     * @param consumeMaxDurationMs max duration to wait for records
+     * @param timeout max duration to wait for records
      * @throws InterruptedException if another threads interrupts this one while waiting for records
      */
-    public void awaitRecords(int consumeMaxDurationMs) throws InterruptedException {
+    public void awaitRecords(long timeout) throws InterruptedException {
         if (recordsRemainingLatch == null || expectedRecords < 0) {
-            throw new IllegalStateException("expectedRecords() was not set for this task?");
+            throw new IllegalStateException("expectedRecords() was not set for this connector?");
         }
-        if (!recordsRemainingLatch.await(consumeMaxDurationMs, TimeUnit.MILLISECONDS)) {
-            String msg = String.format("Insufficient records seen by connector %s in %d millis. Records expected=%d, actual=%d",
+        if (!recordsRemainingLatch.await(timeout, TimeUnit.MILLISECONDS)) {
+            String msg = String.format(
+                    "Insufficient records seen by connector %s in %d millis. Records expected=%d, actual=%d",
                     connectorName,
-                    consumeMaxDurationMs,
+                    timeout,
                     expectedRecords,
                     expectedRecords - recordsRemainingLatch.getCount());
             throw new DataException(msg);
         }
     }
 
+     /**
+     * Wait for this connector to meet the expected number of commits as defined by {@code
+     * expectedCommits}.
+     *
+     * @param  timeout duration to wait for commits
+     * @throws InterruptedException if another threads interrupts this one while waiting for commits
+     */
+    public void awaitCommits(long timeout) throws InterruptedException {
+        if (recordsToCommitLatch == null || expectedCommits < 0) {
+            throw new IllegalStateException("expectedCommits() was not set for this connector?");
+        }
+        if (!recordsToCommitLatch.await(timeout, TimeUnit.MILLISECONDS)) {
+            String msg = String.format(
+                    "Insufficient records committed by connector %s in %d millis. Records expected=%d, actual=%d",
+                    connectorName,
+                    timeout,
+                    expectedCommits,
+                    expectedCommits - recordsToCommitLatch.getCount());
+            throw new DataException(msg);
+        }
+    }
+
+    /**
+     * Record that this connector has been started. This should be called by the connector under
+     * test.
+     *
+     * @see #expectedStarts(int)
+     */
+    public void recordConnectorStart() {
+        startAndStopCounter.recordStart();
+    }
+
+    /**
+     * Record that this connector has been stopped. This should be called by the connector under
+     * test.
+     *
+     * @see #expectedStarts(int)
+     */
+    public void recordConnectorStop() {
+        startAndStopCounter.recordStop();
+    }
+
+    /**
+     * Obtain a {@link StartAndStopLatch} that can be used to wait until the connector using this handle
+     * and all tasks using {@link TaskHandle} have completed the expected number of
+     * starts, starting the counts at the time this method is called.
+     *
+     * 
      A test can call this method, specifying the number of times the connector and tasks
+     * will each be stopped and started from that point (typically {@code expectedStarts(1)}).
+     * The test should then change the connector or otherwise cause the connector to restart one or
+     * more times, and then can call {@link StartAndStopLatch#await(long, TimeUnit)} to wait up to a
+     * specified duration for the connector and all tasks to be started at least the specified
+     * number of times.
+     *
+     * 
      This method does not track the number of times the connector and tasks are stopped, and
+     * only tracks the number of times the connector and tasks are started.
+     *
+     * @param expectedStarts the minimum number of starts that are expected once this method is
+     *                       called
+     * @return the latch that can be used to wait for the starts to complete; never null
+     */
+    public StartAndStopLatch expectedStarts(int expectedStarts) {
+        return expectedStarts(expectedStarts, true);
+    }
+
+    /**
+     * Obtain a {@link StartAndStopLatch} that can be used to wait until the connector using this handle
+     * and optionally all tasks using {@link TaskHandle} have completed the expected number of
+     * starts, starting the counts at the time this method is called.
+     *
+     * 
      A test can call this method, specifying the number of times the connector and tasks
+     * will each be stopped and started from that point (typically {@code expectedStarts(1)}).
+     * The test should then change the connector or otherwise cause the connector to restart one or
+     * more times, and then can call {@link StartAndStopLatch#await(long, TimeUnit)} to wait up to a
+     * specified duration for the connector and all tasks to be started at least the specified
+     * number of times.
+     *
+     * 
      This method does not track the number of times the connector and tasks are stopped, and
+     * only tracks the number of times the connector and tasks are started.
+     *
+     * @param expectedStarts the minimum number of starts that are expected once this method is
+     *                       called
+     * @param includeTasks  true if the latch should also wait for the tasks to be stopped the
+     *                      specified minimum number of times
+     * @return the latch that can be used to wait for the starts to complete; never null
+     */
+    public StartAndStopLatch expectedStarts(int expectedStarts, boolean includeTasks) {
+        List taskLatches = null;
+        if (includeTasks) {
+            taskLatches = taskHandles.values().stream()
+                                     .map(task -> task.expectedStarts(expectedStarts))
+                                     .collect(Collectors.toList());
+        }
+        return startAndStopCounter.expectedStarts(expectedStarts, taskLatches);
+    }
+
+    /**
+     * Obtain a {@link StartAndStopLatch} that can be used to wait until the connector using this handle
+     * and optionally all tasks using {@link TaskHandle} have completed the minimum number of
+     * stops, starting the counts at the time this method is called.
+     *
+     * 
      A test can call this method, specifying the number of times the connector and tasks
+     * will each be stopped from that point (typically {@code expectedStops(1)}).
+     * The test should then change the connector or otherwise cause the connector to stop (or
+     * restart) one or more times, and then can call
+     * {@link StartAndStopLatch#await(long, TimeUnit)} to wait up to a specified duration for the
+     * connector and all tasks to be started at least the specified number of times.
+     *
+     * 
      This method does not track the number of times the connector and tasks are stopped, and
+     * only tracks the number of times the connector and tasks are started.
+     *
+     * @param expectedStops the minimum number of starts that are expected once this method is
+     *                      called
+     * @return the latch that can be used to wait for the starts to complete; never null
+     */
+    public StartAndStopLatch expectedStops(int expectedStops) {
+        return expectedStops(expectedStops, true);
+    }
+
+    /**
+     * Obtain a {@link StartAndStopLatch} that can be used to wait until the connector using this handle
+     * and optionally all tasks using {@link TaskHandle} have completed the minimum number of
+     * stops, starting the counts at the time this method is called.
+     *
+     * 
      A test can call this method, specifying the number of times the connector and tasks
+     * will each be stopped from that point (typically {@code expectedStops(1)}).
+     * The test should then change the connector or otherwise cause the connector to stop (or
+     * restart) one or more times, and then can call
+     * {@link StartAndStopLatch#await(long, TimeUnit)} to wait up to a specified duration for the
+     * connector and all tasks to be started at least the specified number of times.
+     *
+     * 
      This method does not track the number of times the connector and tasks are stopped, and
+     * only tracks the number of times the connector and tasks are started.
+     *
+     * @param expectedStops the minimum number of starts that are expected once this method is
+     *                      called
+     * @param includeTasks  true if the latch should also wait for the tasks to be stopped the
+     *                      specified minimum number of times
+     * @return the latch that can be used to wait for the starts to complete; never null
+     */
+    public StartAndStopLatch expectedStops(int expectedStops, boolean includeTasks) {
+        List taskLatches = null;
+        if (includeTasks) {
+            taskLatches = taskHandles.values().stream()
+                                     .map(task -> task.expectedStops(expectedStops))
+                                     .collect(Collectors.toList());
+        }
+        return startAndStopCounter.expectedStops(expectedStops, taskLatches);
+    }
+
     @Override
     public String toString() {
         return "ConnectorHandle{" +
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ErrorHandlingIntegrationTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ErrorHandlingIntegrationTest.java
index 5f7cfc93082ea..8963b8cb7ca3a 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ErrorHandlingIntegrationTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ErrorHandlingIntegrationTest.java
@@ -34,9 +34,9 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.io.IOException;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.concurrent.TimeUnit;
 
 import static org.apache.kafka.connect.runtime.ConnectorConfig.CONNECTOR_CLASS_CONFIG;
 import static org.apache.kafka.connect.runtime.ConnectorConfig.ERRORS_LOG_ENABLE_CONFIG;
@@ -68,6 +68,7 @@ public class ErrorHandlingIntegrationTest {
 
     private static final Logger log = LoggerFactory.getLogger(ErrorHandlingIntegrationTest.class);
 
+    private static final int NUM_WORKERS = 1;
     private static final String DLQ_TOPIC = "my-connector-errors";
     private static final String CONNECTOR_NAME = "error-conn";
     private static final String TASK_ID = "error-conn-0";
@@ -75,19 +76,21 @@ public class ErrorHandlingIntegrationTest {
     private static final int EXPECTED_CORRECT_RECORDS = 19;
     private static final int EXPECTED_INCORRECT_RECORDS = 1;
     private static final int NUM_TASKS = 1;
-    private static final int CONNECTOR_SETUP_DURATION_MS = 5000;
-    private static final int CONSUME_MAX_DURATION_MS = 5000;
+    private static final long CONNECTOR_SETUP_DURATION_MS = TimeUnit.SECONDS.toMillis(60);
+    private static final long CONSUME_MAX_DURATION_MS = TimeUnit.SECONDS.toMillis(30);
 
     private EmbeddedConnectCluster connect;
     private ConnectorHandle connectorHandle;
 
     @Before
-    public void setup() throws IOException {
+    public void setup() throws InterruptedException {
         // setup Connect cluster with defaults
         connect = new EmbeddedConnectCluster.Builder().build();
 
         // start Connect cluster
         connect.start();
+        connect.assertions().assertAtLeastNumWorkersAreUp(NUM_WORKERS,
+                "Initial group of workers did not start in time.");
 
         // get connector handles before starting test.
         connectorHandle = RuntimeHandles.get().connectorHandle(CONNECTOR_NAME);
@@ -106,7 +109,7 @@ public void testSkipRetryAndDLQWithHeaders() throws Exception {
 
         // setup connector config
         Map props = new HashMap<>();
-        props.put(CONNECTOR_CLASS_CONFIG, "MonitorableSink");
+        props.put(CONNECTOR_CLASS_CONFIG, MonitorableSinkConnector.class.getSimpleName());
         props.put(TASKS_MAX_CONFIG, String.valueOf(NUM_TASKS));
         props.put(TOPICS_CONFIG, "test-topic");
         props.put(KEY_CONVERTER_CLASS_CONFIG, StringConverter.class.getName());
@@ -133,6 +136,8 @@ public void testSkipRetryAndDLQWithHeaders() throws Exception {
         connectorHandle.taskHandle(TASK_ID).expectedRecords(EXPECTED_CORRECT_RECORDS);
 
         connect.configureConnector(CONNECTOR_NAME, props);
+        connect.assertions().assertConnectorAndAtLeastNumTasksAreRunning(CONNECTOR_NAME, NUM_TASKS,
+                "Connector tasks did not start in time.");
 
         waitForCondition(this::checkForPartitionAssignment,
                 CONNECTOR_SETUP_DURATION_MS,
@@ -171,6 +176,9 @@ public void testSkipRetryAndDLQWithHeaders() throws Exception {
         }
 
         connect.deleteConnector(CONNECTOR_NAME);
+        connect.assertions().assertConnectorAndTasksAreStopped(CONNECTOR_NAME,
+                "Connector tasks did not stop in time.");
+
     }
 
     /**
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ExampleConnectIntegrationTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ExampleConnectIntegrationTest.java
index 0648e9ff59ac3..5d36486f018f1 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ExampleConnectIntegrationTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ExampleConnectIntegrationTest.java
@@ -22,15 +22,17 @@
 import org.apache.kafka.test.IntegrationTest;
 import org.junit.After;
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.experimental.categories.Category;
+import org.junit.rules.TestRule;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.io.IOException;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Properties;
+import java.util.concurrent.TimeUnit;
 
 import static org.apache.kafka.connect.runtime.ConnectorConfig.CONNECTOR_CLASS_CONFIG;
 import static org.apache.kafka.connect.runtime.ConnectorConfig.KEY_CONVERTER_CLASS_CONFIG;
@@ -40,6 +42,7 @@
 import static org.apache.kafka.connect.runtime.WorkerConfig.OFFSET_COMMIT_INTERVAL_MS_CONFIG;
 import static org.apache.kafka.test.TestUtils.waitForCondition;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
 
 /**
  * An example integration test that demonstrates how to setup an integration test for Connect.
@@ -54,28 +57,32 @@ public class ExampleConnectIntegrationTest {
 
     private static final int NUM_RECORDS_PRODUCED = 2000;
     private static final int NUM_TOPIC_PARTITIONS = 3;
-    private static final int CONSUME_MAX_DURATION_MS = 5000;
-    private static final int CONNECTOR_SETUP_DURATION_MS = 15000;
+    private static final long RECORD_TRANSFER_DURATION_MS = TimeUnit.SECONDS.toMillis(30);
+    private static final long CONNECTOR_SETUP_DURATION_MS = TimeUnit.SECONDS.toMillis(60);
     private static final int NUM_TASKS = 3;
+    private static final int NUM_WORKERS = 3;
     private static final String CONNECTOR_NAME = "simple-conn";
 
     private EmbeddedConnectCluster connect;
     private ConnectorHandle connectorHandle;
 
+    @Rule
+    public TestRule watcher = ConnectIntegrationTestUtils.newTestWatcher(log);
+
     @Before
-    public void setup() throws IOException {
+    public void setup() {
         // setup Connect worker properties
         Map exampleWorkerProps = new HashMap<>();
-        exampleWorkerProps.put(OFFSET_COMMIT_INTERVAL_MS_CONFIG, "30000");
+        exampleWorkerProps.put(OFFSET_COMMIT_INTERVAL_MS_CONFIG, String.valueOf(5_000));
 
         // setup Kafka broker properties
         Properties exampleBrokerProps = new Properties();
         exampleBrokerProps.put("auto.create.topics.enable", "false");
 
-        // build a Connect cluster backed by Kakfa and Zk
+        // build a Connect cluster backed by Kafka and Zk
         connect = new EmbeddedConnectCluster.Builder()
-                .name("example-cluster")
-                .numWorkers(3)
+                .name("connect-cluster")
+                .numWorkers(NUM_WORKERS)
                 .numBrokers(1)
                 .workerProps(exampleWorkerProps)
                 .brokerProps(exampleBrokerProps)
@@ -93,7 +100,7 @@ public void close() {
         // delete connector handle
         RuntimeHandles.get().deleteConnector(CONNECTOR_NAME);
 
-        // stop all Connect, Kakfa and Zk threads.
+        // stop all Connect, Kafka and Zk threads.
         connect.stop();
     }
 
@@ -102,13 +109,13 @@ public void close() {
      * records, and start up a sink connector which will consume these records.
      */
     @Test
-    public void testProduceConsumeConnector() throws Exception {
+    public void testSinkConnector() throws Exception {
         // create test topic
         connect.kafka().createTopic("test-topic", NUM_TOPIC_PARTITIONS);
 
         // setup up props for the sink connector
         Map props = new HashMap<>();
-        props.put(CONNECTOR_CLASS_CONFIG, "MonitorableSink");
+        props.put(CONNECTOR_CLASS_CONFIG, MonitorableSinkConnector.class.getSimpleName());
         props.put(TASKS_MAX_CONFIG, String.valueOf(NUM_TASKS));
         props.put(TOPICS_CONFIG, "test-topic");
         props.put(KEY_CONVERTER_CLASS_CONFIG, StringConverter.class.getName());
@@ -117,6 +124,9 @@ public void testProduceConsumeConnector() throws Exception {
         // expect all records to be consumed by the connector
         connectorHandle.expectedRecords(NUM_RECORDS_PRODUCED);
 
+        // expect all records to be consumed by the connector
+        connectorHandle.expectedCommits(NUM_RECORDS_PRODUCED);
+
         // start a sink connector
         connect.configureConnector(CONNECTOR_NAME, props);
 
@@ -131,10 +141,55 @@ public void testProduceConsumeConnector() throws Exception {
 
         // consume all records from the source topic or fail, to ensure that they were correctly produced.
         assertEquals("Unexpected number of records consumed", NUM_RECORDS_PRODUCED,
-                connect.kafka().consume(NUM_RECORDS_PRODUCED, CONSUME_MAX_DURATION_MS, "test-topic").count());
+                connect.kafka().consume(NUM_RECORDS_PRODUCED, RECORD_TRANSFER_DURATION_MS, "test-topic").count());
 
         // wait for the connector tasks to consume all records.
-        connectorHandle.awaitRecords(CONSUME_MAX_DURATION_MS);
+        connectorHandle.awaitRecords(RECORD_TRANSFER_DURATION_MS);
+
+        // wait for the connector tasks to commit all records.
+        connectorHandle.awaitCommits(RECORD_TRANSFER_DURATION_MS);
+
+        // delete connector
+        connect.deleteConnector(CONNECTOR_NAME);
+    }
+
+    /**
+     * Simple test case to configure and execute an embedded Connect cluster. The test will produce and consume
+     * records, and start up a sink connector which will consume these records.
+     */
+    @Test
+    public void testSourceConnector() throws Exception {
+        // create test topic
+        connect.kafka().createTopic("test-topic", NUM_TOPIC_PARTITIONS);
+
+        // setup up props for the sink connector
+        Map props = new HashMap<>();
+        props.put(CONNECTOR_CLASS_CONFIG, MonitorableSourceConnector.class.getSimpleName());
+        props.put(TASKS_MAX_CONFIG, String.valueOf(NUM_TASKS));
+        props.put("topic", "test-topic");
+        props.put("throughput", String.valueOf(500));
+        props.put(KEY_CONVERTER_CLASS_CONFIG, StringConverter.class.getName());
+        props.put(VALUE_CONVERTER_CLASS_CONFIG, StringConverter.class.getName());
+
+        // expect all records to be produced by the connector
+        connectorHandle.expectedRecords(NUM_RECORDS_PRODUCED);
+
+        // expect all records to be produced by the connector
+        connectorHandle.expectedCommits(NUM_RECORDS_PRODUCED);
+
+        // start a source connector
+        connect.configureConnector(CONNECTOR_NAME, props);
+
+        // wait for the connector tasks to produce enough records
+        connectorHandle.awaitRecords(RECORD_TRANSFER_DURATION_MS);
+
+        // wait for the connector tasks to commit enough records
+        connectorHandle.awaitCommits(RECORD_TRANSFER_DURATION_MS);
+
+        // consume all records from the source topic or fail, to ensure that they were correctly produced
+        int recordNum = connect.kafka().consume(NUM_RECORDS_PRODUCED, RECORD_TRANSFER_DURATION_MS, "test-topic").count();
+        assertTrue("Not enough records produced by source connector. Expected at least: " + NUM_RECORDS_PRODUCED + " + but got " + recordNum,
+                recordNum >= NUM_RECORDS_PRODUCED);
 
         // delete connector
         connect.deleteConnector(CONNECTOR_NAME);
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/MonitorableSinkConnector.java b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/MonitorableSinkConnector.java
index 23a8d99e84edc..ed6f6b96f98c2 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/MonitorableSinkConnector.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/MonitorableSinkConnector.java
@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.connect.integration;
 
+import org.apache.kafka.clients.consumer.OffsetAndMetadata;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.config.ConfigDef;
 import org.apache.kafka.connect.connector.Task;
@@ -28,11 +29,15 @@
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 /**
- * A connector to be used in integration tests. This class provides methods to find task instances
+ * A sink connector that is used in Apache Kafka integration tests to verify the behavior of the
+ * Connect framework, but that can be used in other integration tests as a simple connector that
+ * consumes and counts records. This class provides methods to find task instances
  * which are initiated by the embedded connector, and wait for them to consume a desired number of
  * messages.
  */
@@ -41,11 +46,16 @@ public class MonitorableSinkConnector extends TestSinkConnector {
     private static final Logger log = LoggerFactory.getLogger(MonitorableSinkConnector.class);
 
     private String connectorName;
+    private Map commonConfigs;
+    private ConnectorHandle connectorHandle;
 
     @Override
     public void start(Map props) {
+        connectorHandle = RuntimeHandles.get().connectorHandle(props.get("name"));
         connectorName = props.get("name");
+        commonConfigs = props;
         log.info("Starting connector {}", props.get("name"));
+        connectorHandle.recordConnectorStart();
     }
 
     @Override
@@ -57,7 +67,7 @@ public Class taskClass() {
     public List> taskConfigs(int maxTasks) {
         List> configs = new ArrayList<>();
         for (int i = 0; i < maxTasks; i++) {
-            Map config = new HashMap<>();
+            Map config = new HashMap<>(commonConfigs);
             config.put("connector.name", connectorName);
             config.put("task.id", connectorName + "-" + i);
             configs.add(config);
@@ -67,6 +77,7 @@ public List> taskConfigs(int maxTasks) {
 
     @Override
     public void stop() {
+        connectorHandle.recordConnectorStop();
     }
 
     @Override
@@ -79,6 +90,15 @@ public static class MonitorableSinkTask extends SinkTask {
         private String connectorName;
         private String taskId;
         private TaskHandle taskHandle;
+        private Set assignments;
+        private Map committedOffsets;
+        private Map> cachedTopicPartitions;
+
+        public MonitorableSinkTask() {
+            this.assignments = new HashSet<>();
+            this.committedOffsets = new HashMap<>();
+            this.cachedTopicPartitions = new HashMap<>();
+        }
 
         @Override
         public String version() {
@@ -91,12 +111,13 @@ public void start(Map props) {
             connectorName = props.get("connector.name");
             taskHandle = RuntimeHandles.get().connectorHandle(connectorName).taskHandle(taskId);
             log.debug("Starting task {}", taskId);
+            taskHandle.recordTaskStart();
         }
 
         @Override
         public void open(Collection partitions) {
             log.debug("Opening {} partitions", partitions.size());
-            super.open(partitions);
+            assignments.addAll(partitions);
             taskHandle.partitionsAssigned(partitions.size());
         }
 
@@ -104,12 +125,35 @@ public void open(Collection partitions) {
         public void put(Collection records) {
             for (SinkRecord rec : records) {
                 taskHandle.record();
+                TopicPartition tp = cachedTopicPartitions
+                        .computeIfAbsent(rec.topic(), v -> new HashMap<>())
+                        .computeIfAbsent(rec.kafkaPartition(), v -> new TopicPartition(rec.topic(), rec.kafkaPartition()));
+                committedOffsets.put(tp, committedOffsets.getOrDefault(tp, 0L) + 1);
                 log.trace("Task {} obtained record (key='{}' value='{}')", taskId, rec.key(), rec.value());
             }
         }
 
+        @Override
+        public Map preCommit(Map offsets) {
+            for (TopicPartition tp : assignments) {
+                Long recordsSinceLastCommit = committedOffsets.get(tp);
+                if (recordsSinceLastCommit == null) {
+                    log.warn("preCommit was called with topic-partition {} that is not included "
+                            + "in the assignments of this task {}", tp, assignments);
+                } else {
+                    taskHandle.commit(recordsSinceLastCommit.intValue());
+                    log.error("Forwarding to framework request to commit additional {} for {}",
+                            recordsSinceLastCommit, tp);
+                    taskHandle.commit((int) (long) recordsSinceLastCommit);
+                    committedOffsets.put(tp, 0L);
+                }
+            }
+            return offsets;
+        }
+
         @Override
         public void stop() {
+            taskHandle.recordTaskStop();
         }
     }
 }
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/MonitorableSourceConnector.java b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/MonitorableSourceConnector.java
new file mode 100644
index 0000000000000..fbd763a928035
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/MonitorableSourceConnector.java
@@ -0,0 +1,166 @@
+/*
+ * 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.integration;
+
+import org.apache.kafka.common.config.ConfigDef;
+import org.apache.kafka.clients.producer.RecordMetadata;
+import org.apache.kafka.connect.connector.Task;
+import org.apache.kafka.connect.data.Schema;
+import org.apache.kafka.connect.runtime.TestSourceConnector;
+import org.apache.kafka.connect.source.SourceRecord;
+import org.apache.kafka.connect.source.SourceTask;
+import org.apache.kafka.tools.ThroughputThrottler;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.stream.Collectors;
+import java.util.stream.LongStream;
+
+/**
+ * A source connector that is used in Apache Kafka integration tests to verify the behavior of
+ * the Connect framework, but that can be used in other integration tests as a simple connector
+ * that generates records of a fixed structure. The rate of record production can be adjusted
+ * through the configs 'throughput' and 'messages.per.poll'
+ */
+public class MonitorableSourceConnector extends TestSourceConnector {
+    private static final Logger log = LoggerFactory.getLogger(MonitorableSourceConnector.class);
+
+    public static final String TOPIC_CONFIG = "topic";
+    private String connectorName;
+    private ConnectorHandle connectorHandle;
+    private Map commonConfigs;
+
+    @Override
+    public void start(Map props) {
+        connectorHandle = RuntimeHandles.get().connectorHandle(props.get("name"));
+        connectorName = connectorHandle.name();
+        commonConfigs = props;
+        log.info("Started {} connector {}", this.getClass().getSimpleName(), connectorName);
+        connectorHandle.recordConnectorStart();
+    }
+
+    @Override
+    public Class taskClass() {
+        return MonitorableSourceTask.class;
+    }
+
+    @Override
+    public List> taskConfigs(int maxTasks) {
+        List> configs = new ArrayList<>();
+        for (int i = 0; i < maxTasks; i++) {
+            Map config = new HashMap<>(commonConfigs);
+            config.put("connector.name", connectorName);
+            config.put("task.id", connectorName + "-" + i);
+            configs.add(config);
+        }
+        return configs;
+    }
+
+    @Override
+    public void stop() {
+        log.info("Stopped {} connector {}", this.getClass().getSimpleName(), connectorName);
+        connectorHandle.recordConnectorStop();
+    }
+
+    @Override
+    public ConfigDef config() {
+        log.info("Configured {} connector {}", this.getClass().getSimpleName(), connectorName);
+        return new ConfigDef();
+    }
+
+    public static class MonitorableSourceTask extends SourceTask {
+        private String connectorName;
+        private String taskId;
+        private String topicName;
+        private TaskHandle taskHandle;
+        private volatile boolean stopped;
+        private long startingSeqno;
+        private long seqno;
+        private long throughput;
+        private int batchSize;
+        private ThroughputThrottler throttler;
+
+        @Override
+        public String version() {
+            return "unknown";
+        }
+
+        @Override
+        public void start(Map props) {
+            taskId = props.get("task.id");
+            connectorName = props.get("connector.name");
+            topicName = props.getOrDefault(TOPIC_CONFIG, "sequential-topic");
+            throughput = Long.valueOf(props.getOrDefault("throughput", "-1"));
+            batchSize = Integer.valueOf(props.getOrDefault("messages.per.poll", "1"));
+            taskHandle = RuntimeHandles.get().connectorHandle(connectorName).taskHandle(taskId);
+            Map offset = Optional.ofNullable(
+                    context.offsetStorageReader().offset(Collections.singletonMap("task.id", taskId)))
+                    .orElse(Collections.emptyMap());
+            startingSeqno = Optional.ofNullable((Long) offset.get("saved")).orElse(0L);
+            log.info("Started {} task {} with properties {}", this.getClass().getSimpleName(), taskId, props);
+            throttler = new ThroughputThrottler(throughput, System.currentTimeMillis());
+            taskHandle.recordTaskStart();
+        }
+
+        @Override
+        public List poll() {
+            if (!stopped) {
+                if (throttler.shouldThrottle(seqno - startingSeqno, System.currentTimeMillis())) {
+                    throttler.throttle();
+                }
+                taskHandle.record(batchSize);
+                return LongStream.range(0, batchSize)
+                        .mapToObj(i -> new SourceRecord(
+                                Collections.singletonMap("task.id", taskId),
+                                Collections.singletonMap("saved", ++seqno),
+                                topicName,
+                                null,
+                                Schema.STRING_SCHEMA,
+                                "key-" + taskId + "-" + seqno,
+                                Schema.STRING_SCHEMA,
+                                "value-" + taskId + "-" + seqno))
+                        .collect(Collectors.toList());
+            }
+            return null;
+        }
+
+        @Override
+        public void commit() {
+            log.info("Task {} committing offsets", taskId);
+            //TODO: save progress outside the offset topic, potentially in the task handle
+        }
+
+        @Override
+        public void commitRecord(SourceRecord record, RecordMetadata metadata) {
+            log.trace("Committing record: {}", record);
+            taskHandle.commit();
+        }
+
+        @Override
+        public void stop() {
+            log.info("Stopped {} task {}", this.getClass().getSimpleName(), taskId);
+            stopped = true;
+            taskHandle.recordTaskStop();
+        }
+    }
+}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/RebalanceSourceConnectorsIntegrationTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/RebalanceSourceConnectorsIntegrationTest.java
new file mode 100644
index 0000000000000..be8ce6136a351
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/RebalanceSourceConnectorsIntegrationTest.java
@@ -0,0 +1,333 @@
+/*
+ * 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.integration;
+
+import org.apache.kafka.connect.runtime.rest.entities.ConnectorStateInfo;
+import org.apache.kafka.connect.storage.StringConverter;
+import org.apache.kafka.connect.util.clusters.EmbeddedConnectCluster;
+import org.apache.kafka.test.IntegrationTest;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.experimental.categories.Category;
+import org.junit.rules.TestRule;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Properties;
+import java.util.concurrent.TimeUnit;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import static org.apache.kafka.connect.integration.MonitorableSourceConnector.TOPIC_CONFIG;
+import static org.apache.kafka.connect.runtime.ConnectorConfig.CONNECTOR_CLASS_CONFIG;
+import static org.apache.kafka.connect.runtime.ConnectorConfig.KEY_CONVERTER_CLASS_CONFIG;
+import static org.apache.kafka.connect.runtime.ConnectorConfig.TASKS_MAX_CONFIG;
+import static org.apache.kafka.connect.runtime.ConnectorConfig.VALUE_CONVERTER_CLASS_CONFIG;
+import static org.apache.kafka.connect.runtime.WorkerConfig.OFFSET_COMMIT_INTERVAL_MS_CONFIG;
+import static org.apache.kafka.connect.runtime.distributed.ConnectProtocolCompatibility.COMPATIBLE;
+import static org.apache.kafka.connect.runtime.distributed.DistributedConfig.CONNECT_PROTOCOL_CONFIG;
+import static org.apache.kafka.connect.runtime.distributed.DistributedConfig.SCHEDULED_REBALANCE_MAX_DELAY_MS_CONFIG;
+import static org.apache.kafka.test.TestUtils.waitForCondition;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotEquals;
+import static org.junit.Assert.assertTrue;
+
+/**
+ * Integration tests for incremental cooperative rebalancing between Connect workers
+ */
+@Category(IntegrationTest.class)
+public class RebalanceSourceConnectorsIntegrationTest {
+
+    private static final Logger log = LoggerFactory.getLogger(RebalanceSourceConnectorsIntegrationTest.class);
+
+    private static final int NUM_TOPIC_PARTITIONS = 3;
+    private static final long CONNECTOR_SETUP_DURATION_MS = TimeUnit.SECONDS.toMillis(30);
+    private static final long WORKER_SETUP_DURATION_MS = TimeUnit.SECONDS.toMillis(60);
+    private static final int NUM_WORKERS = 3;
+    private static final int NUM_TASKS = 4;
+    private static final String CONNECTOR_NAME = "seq-source1";
+    private static final String TOPIC_NAME = "sequential-topic";
+
+    private EmbeddedConnectCluster connect;
+
+    @Rule
+    public TestRule watcher = ConnectIntegrationTestUtils.newTestWatcher(log);
+
+    @Before
+    public void setup() {
+        // setup Connect worker properties
+        Map workerProps = new HashMap<>();
+        workerProps.put(CONNECT_PROTOCOL_CONFIG, COMPATIBLE.toString());
+        workerProps.put(OFFSET_COMMIT_INTERVAL_MS_CONFIG, String.valueOf(TimeUnit.SECONDS.toMillis(30)));
+        workerProps.put(SCHEDULED_REBALANCE_MAX_DELAY_MS_CONFIG, String.valueOf(TimeUnit.SECONDS.toMillis(30)));
+
+        // setup Kafka broker properties
+        Properties brokerProps = new Properties();
+        brokerProps.put("auto.create.topics.enable", "false");
+
+        // build a Connect cluster backed by Kafka and Zk
+        connect = new EmbeddedConnectCluster.Builder()
+                .name("connect-cluster")
+                .numWorkers(NUM_WORKERS)
+                .numBrokers(1)
+                .workerProps(workerProps)
+                .brokerProps(brokerProps)
+                .build();
+
+        // start the clusters
+        connect.start();
+    }
+
+    @After
+    public void close() {
+        // stop all Connect, Kafka and Zk threads.
+        connect.stop();
+    }
+
+    @Test
+    public void testStartTwoConnectors() throws Exception {
+        // create test topic
+        connect.kafka().createTopic(TOPIC_NAME, NUM_TOPIC_PARTITIONS);
+
+        // setup up props for the source connector
+        Map props = new HashMap<>();
+        props.put(CONNECTOR_CLASS_CONFIG, MonitorableSourceConnector.class.getSimpleName());
+        props.put(TASKS_MAX_CONFIG, String.valueOf(NUM_TASKS));
+        props.put("throughput", String.valueOf(1));
+        props.put("messages.per.poll", String.valueOf(10));
+        props.put(TOPIC_CONFIG, TOPIC_NAME);
+        props.put(KEY_CONVERTER_CLASS_CONFIG, StringConverter.class.getName());
+        props.put(VALUE_CONVERTER_CLASS_CONFIG, StringConverter.class.getName());
+
+        connect.assertions().assertAtLeastNumWorkersAreUp(NUM_WORKERS,
+                "Connect workers did not start in time.");
+
+        // start a source connector
+        connect.configureConnector(CONNECTOR_NAME, props);
+
+        connect.assertions().assertConnectorAndAtLeastNumTasksAreRunning(CONNECTOR_NAME, NUM_TASKS,
+                "Connector tasks did not start in time.");
+
+        // start a source connector
+        connect.configureConnector("another-source", props);
+
+        connect.assertions().assertConnectorAndAtLeastNumTasksAreRunning(CONNECTOR_NAME, NUM_TASKS,
+                "Connector tasks did not start in time.");
+
+        connect.assertions().assertConnectorAndAtLeastNumTasksAreRunning("another-source", 4,
+                "Connector tasks did not start in time.");
+    }
+
+    @Test
+    public void testReconfigConnector() throws Exception {
+        ConnectorHandle connectorHandle = RuntimeHandles.get().connectorHandle(CONNECTOR_NAME);
+
+        // create test topic
+        String anotherTopic = "another-topic";
+        connect.kafka().createTopic(TOPIC_NAME, NUM_TOPIC_PARTITIONS);
+        connect.kafka().createTopic(anotherTopic, NUM_TOPIC_PARTITIONS);
+
+        // setup up props for the source connector
+        Map props = new HashMap<>();
+        props.put(CONNECTOR_CLASS_CONFIG, MonitorableSourceConnector.class.getSimpleName());
+        props.put(TASKS_MAX_CONFIG, String.valueOf(NUM_TASKS));
+        props.put("throughput", String.valueOf(1));
+        props.put("messages.per.poll", String.valueOf(10));
+        props.put(TOPIC_CONFIG, TOPIC_NAME);
+        props.put(KEY_CONVERTER_CLASS_CONFIG, StringConverter.class.getName());
+        props.put(VALUE_CONVERTER_CLASS_CONFIG, StringConverter.class.getName());
+
+        connect.assertions().assertAtLeastNumWorkersAreUp(NUM_WORKERS,
+                "Connect workers did not start in time.");
+
+        // start a source connector
+        connect.configureConnector(CONNECTOR_NAME, props);
+
+        connect.assertions().assertConnectorAndAtLeastNumTasksAreRunning(CONNECTOR_NAME, NUM_TASKS,
+                "Connector tasks did not start in time.");
+
+        int numRecordsProduced = 100;
+        long recordTransferDurationMs = TimeUnit.SECONDS.toMillis(30);
+
+        // consume all records from the source topic or fail, to ensure that they were correctly produced
+        int recordNum = connect.kafka().consume(numRecordsProduced, recordTransferDurationMs, TOPIC_NAME).count();
+        assertTrue("Not enough records produced by source connector. Expected at least: " + numRecordsProduced + " + but got " + recordNum,
+                recordNum >= numRecordsProduced);
+
+        // expect that we're going to restart the connector and its tasks
+        StartAndStopLatch restartLatch = connectorHandle.expectedStarts(1);
+
+        // Reconfigure the source connector by changing the Kafka topic used as output
+        props.put(TOPIC_CONFIG, anotherTopic);
+        connect.configureConnector(CONNECTOR_NAME, props);
+
+        // Wait for the connector *and tasks* to be restarted
+        assertTrue("Failed to alter connector configuration and see connector and tasks restart "
+                   + "within " + CONNECTOR_SETUP_DURATION_MS + "ms",
+                restartLatch.await(CONNECTOR_SETUP_DURATION_MS, TimeUnit.MILLISECONDS));
+
+        // And wait for the Connect to show the connectors and tasks are running
+        connect.assertions().assertConnectorAndAtLeastNumTasksAreRunning(CONNECTOR_NAME, NUM_TASKS,
+                "Connector tasks did not start in time.");
+
+        // consume all records from the source topic or fail, to ensure that they were correctly produced
+        recordNum = connect.kafka().consume(numRecordsProduced, recordTransferDurationMs, anotherTopic).count();
+        assertTrue("Not enough records produced by source connector. Expected at least: " + numRecordsProduced + " + but got " + recordNum,
+                recordNum >= numRecordsProduced);
+    }
+
+    @Test
+    public void testDeleteConnector() throws Exception {
+        // create test topic
+        connect.kafka().createTopic(TOPIC_NAME, NUM_TOPIC_PARTITIONS);
+
+        // setup up props for the source connector
+        Map props = new HashMap<>();
+        props.put(CONNECTOR_CLASS_CONFIG, MonitorableSourceConnector.class.getSimpleName());
+        props.put(TASKS_MAX_CONFIG, String.valueOf(NUM_TASKS));
+        props.put("throughput", String.valueOf(1));
+        props.put("messages.per.poll", String.valueOf(10));
+        props.put(TOPIC_CONFIG, TOPIC_NAME);
+        props.put(KEY_CONVERTER_CLASS_CONFIG, StringConverter.class.getName());
+        props.put(VALUE_CONVERTER_CLASS_CONFIG, StringConverter.class.getName());
+
+        connect.assertions().assertAtLeastNumWorkersAreUp(NUM_WORKERS,
+                "Connect workers did not start in time.");
+
+        // start several source connectors
+        IntStream.range(0, 4).forEachOrdered(i -> connect.configureConnector(CONNECTOR_NAME + i, props));
+
+        connect.assertions().assertConnectorAndAtLeastNumTasksAreRunning(CONNECTOR_NAME + 3, NUM_TASKS,
+                "Connector tasks did not start in time.");
+
+        // delete connector
+        connect.deleteConnector(CONNECTOR_NAME + 3);
+
+        connect.assertions().assertConnectorAndTasksAreStopped(CONNECTOR_NAME + 3,
+                "Connector tasks did not stop in time.");
+
+        waitForCondition(this::assertConnectorAndTasksAreUnique,
+                WORKER_SETUP_DURATION_MS, "Connect and tasks are imbalanced between the workers.");
+    }
+
+    @Test
+    public void testAddingWorker() throws Exception {
+        // create test topic
+        connect.kafka().createTopic(TOPIC_NAME, NUM_TOPIC_PARTITIONS);
+
+        // setup up props for the source connector
+        Map props = new HashMap<>();
+        props.put(CONNECTOR_CLASS_CONFIG, MonitorableSourceConnector.class.getSimpleName());
+        props.put(TASKS_MAX_CONFIG, String.valueOf(NUM_TASKS));
+        props.put("throughput", String.valueOf(1));
+        props.put("messages.per.poll", String.valueOf(10));
+        props.put(TOPIC_CONFIG, TOPIC_NAME);
+        props.put(KEY_CONVERTER_CLASS_CONFIG, StringConverter.class.getName());
+        props.put(VALUE_CONVERTER_CLASS_CONFIG, StringConverter.class.getName());
+
+        connect.assertions().assertAtLeastNumWorkersAreUp(NUM_WORKERS,
+                "Connect workers did not start in time.");
+
+        // start a source connector
+        IntStream.range(0, 4).forEachOrdered(i -> connect.configureConnector(CONNECTOR_NAME + i, props));
+
+        connect.assertions().assertConnectorAndAtLeastNumTasksAreRunning(CONNECTOR_NAME + 3, NUM_TASKS,
+                "Connector tasks did not start in time.");
+
+        connect.addWorker();
+
+        connect.assertions().assertAtLeastNumWorkersAreUp(NUM_WORKERS + 1,
+                "Connect workers did not start in time.");
+
+        connect.assertions().assertConnectorAndAtLeastNumTasksAreRunning(CONNECTOR_NAME + 3, NUM_TASKS,
+                "Connector tasks did not start in time.");
+
+        waitForCondition(this::assertConnectorAndTasksAreUnique,
+                WORKER_SETUP_DURATION_MS, "Connect and tasks are imbalanced between the workers.");
+    }
+
+    @Test
+    public void testRemovingWorker() throws Exception {
+        // create test topic
+        connect.kafka().createTopic(TOPIC_NAME, NUM_TOPIC_PARTITIONS);
+
+        // setup up props for the source connector
+        Map props = new HashMap<>();
+        props.put(CONNECTOR_CLASS_CONFIG, MonitorableSourceConnector.class.getSimpleName());
+        props.put(TASKS_MAX_CONFIG, String.valueOf(NUM_TASKS));
+        props.put("throughput", String.valueOf(1));
+        props.put("messages.per.poll", String.valueOf(10));
+        props.put(TOPIC_CONFIG, TOPIC_NAME);
+        props.put(KEY_CONVERTER_CLASS_CONFIG, StringConverter.class.getName());
+        props.put(VALUE_CONVERTER_CLASS_CONFIG, StringConverter.class.getName());
+
+        connect.assertions().assertExactlyNumWorkersAreUp(NUM_WORKERS,
+                "Connect workers did not start in time.");
+
+        // start a source connector
+        IntStream.range(0, 4).forEachOrdered(i -> connect.configureConnector(CONNECTOR_NAME + i, props));
+
+        connect.assertions().assertConnectorAndAtLeastNumTasksAreRunning(CONNECTOR_NAME + 3, NUM_TASKS,
+                "Connector tasks did not start in time.");
+
+        connect.removeWorker();
+
+        connect.assertions().assertExactlyNumWorkersAreUp(NUM_WORKERS - 1,
+                "Connect workers did not start in time.");
+
+        waitForCondition(this::assertConnectorAndTasksAreUnique,
+                WORKER_SETUP_DURATION_MS, "Connect and tasks are imbalanced between the workers.");
+    }
+
+    private boolean assertConnectorAndTasksAreUnique() {
+        try {
+            Map> connectors = new HashMap<>();
+            Map> tasks = new HashMap<>();
+            for (String connector : connect.connectors()) {
+                ConnectorStateInfo info = connect.connectorStatus(connector);
+                connectors.computeIfAbsent(info.connector().workerId(), k -> new ArrayList<>())
+                        .add(connector);
+                info.tasks().forEach(
+                    t -> tasks.computeIfAbsent(t.workerId(), k -> new ArrayList<>())
+                           .add(connector + "-" + t.id()));
+            }
+
+            int maxConnectors = connectors.values().stream().mapToInt(Collection::size).max().orElse(0);
+            int maxTasks = tasks.values().stream().mapToInt(Collection::size).max().orElse(0);
+
+            assertNotEquals("Found no connectors running!", maxConnectors, 0);
+            assertNotEquals("Found no tasks running!", maxTasks, 0);
+            assertEquals("Connector assignments are not unique: " + connectors,
+                    connectors.values().size(),
+                    connectors.values().stream().distinct().collect(Collectors.toList()).size());
+            assertEquals("Task assignments are not unique: " + tasks,
+                    tasks.values().size(),
+                    tasks.values().stream().distinct().collect(Collectors.toList()).size());
+            return true;
+        } catch (Exception e) {
+            log.error("Could not check connector state info.", e);
+            return false;
+        }
+    }
+
+}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/RestExtensionIntegrationTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/RestExtensionIntegrationTest.java
new file mode 100644
index 0000000000000..6ec86bd9342b1
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/RestExtensionIntegrationTest.java
@@ -0,0 +1,217 @@
+/*
+ * 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.integration;
+
+import org.apache.kafka.connect.errors.ConnectException;
+import org.apache.kafka.connect.health.ConnectClusterState;
+import org.apache.kafka.connect.health.ConnectorHealth;
+import org.apache.kafka.connect.health.ConnectorState;
+import org.apache.kafka.connect.health.ConnectorType;
+import org.apache.kafka.connect.health.TaskState;
+import org.apache.kafka.connect.rest.ConnectRestExtension;
+import org.apache.kafka.connect.rest.ConnectRestExtensionContext;
+import org.apache.kafka.connect.util.clusters.EmbeddedConnectCluster;
+import org.apache.kafka.connect.util.clusters.WorkerHandle;
+import org.apache.kafka.test.IntegrationTest;
+import org.junit.After;
+import org.junit.Test;
+import org.junit.experimental.categories.Category;
+
+import javax.ws.rs.GET;
+import javax.ws.rs.Path;
+import javax.ws.rs.core.Response;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.TimeUnit;
+
+import static javax.ws.rs.core.Response.Status.BAD_REQUEST;
+import static org.apache.kafka.connect.runtime.ConnectorConfig.CONNECTOR_CLASS_CONFIG;
+import static org.apache.kafka.connect.runtime.ConnectorConfig.NAME_CONFIG;
+import static org.apache.kafka.connect.runtime.ConnectorConfig.TASKS_MAX_CONFIG;
+import static org.apache.kafka.connect.runtime.SinkConnectorConfig.TOPICS_CONFIG;
+import static org.apache.kafka.connect.runtime.WorkerConfig.REST_EXTENSION_CLASSES_CONFIG;
+import static org.apache.kafka.test.TestUtils.waitForCondition;
+import static org.junit.Assert.assertEquals;
+
+/**
+ * A simple integration test to ensure that REST extensions are registered correctly.
+ */
+@Category(IntegrationTest.class)
+public class RestExtensionIntegrationTest {
+
+    private static final long REST_EXTENSION_REGISTRATION_TIMEOUT_MS = TimeUnit.MINUTES.toMillis(1);
+    private static final long CONNECTOR_HEALTH_AND_CONFIG_TIMEOUT_MS = TimeUnit.MINUTES.toMillis(1);
+    private static final int NUM_WORKERS = 1;
+
+    private EmbeddedConnectCluster connect;
+
+    @Test
+    public void testRestExtensionApi() throws InterruptedException {
+        // setup Connect worker properties
+        Map workerProps = new HashMap<>();
+        workerProps.put(REST_EXTENSION_CLASSES_CONFIG, IntegrationTestRestExtension.class.getName());
+
+        // build a Connect cluster backed by Kafka and Zk
+        connect = new EmbeddedConnectCluster.Builder()
+            .name("connect-cluster")
+            .numWorkers(NUM_WORKERS)
+            .numBrokers(1)
+            .workerProps(workerProps)
+            .build();
+
+        // start the clusters
+        connect.start();
+
+        connect.assertions().assertAtLeastNumWorkersAreUp(NUM_WORKERS,
+                "Initial group of workers did not start in time.");
+
+        WorkerHandle worker = connect.workers().stream()
+            .findFirst()
+            .orElseThrow(() -> new AssertionError("At least one worker handle should be available"));
+
+        waitForCondition(
+            this::extensionIsRegistered,
+            REST_EXTENSION_REGISTRATION_TIMEOUT_MS,
+            "REST extension was never registered"
+        );
+
+        ConnectorHandle connectorHandle = RuntimeHandles.get().connectorHandle("test-conn");
+        try {
+            // setup up props for the connector
+            Map connectorProps = new HashMap<>();
+            connectorProps.put(CONNECTOR_CLASS_CONFIG, MonitorableSinkConnector.class.getSimpleName());
+            connectorProps.put(TASKS_MAX_CONFIG, String.valueOf(1));
+            connectorProps.put(TOPICS_CONFIG, "test-topic");
+
+            // start a connector
+            connectorHandle.taskHandle(connectorHandle.name() + "-0");
+            StartAndStopLatch connectorStartLatch = connectorHandle.expectedStarts(1);
+            connect.configureConnector(connectorHandle.name(), connectorProps);
+            connect.assertions().assertConnectorAndAtLeastNumTasksAreRunning(connectorHandle.name(), 1,
+                    "Connector tasks did not start in time.");
+            connectorStartLatch.await(CONNECTOR_HEALTH_AND_CONFIG_TIMEOUT_MS, TimeUnit.MILLISECONDS);
+
+            String workerId = String.format("%s:%d", worker.url().getHost(), worker.url().getPort());
+            ConnectorHealth expectedHealth = new ConnectorHealth(
+                connectorHandle.name(),
+                new ConnectorState(
+                    "RUNNING",
+                    workerId,
+                    null
+                ),
+                Collections.singletonMap(
+                    0,
+                    new TaskState(0, "RUNNING", workerId, null)
+                ),
+                ConnectorType.SINK
+            );
+
+            connectorProps.put(NAME_CONFIG, connectorHandle.name());
+
+            // Test the REST extension API; specifically, that the connector's health and configuration
+            // are available to the REST extension we registered and that they contain expected values
+            waitForCondition(
+                () -> verifyConnectorHealthAndConfig(connectorHandle.name(), expectedHealth, connectorProps),
+                CONNECTOR_HEALTH_AND_CONFIG_TIMEOUT_MS,
+                "Connector health and/or config was never accessible by the REST extension"
+            );
+        } finally {
+            RuntimeHandles.get().deleteConnector(connectorHandle.name());
+        }
+    }
+
+    @After
+    public void close() {
+        // stop all Connect, Kafka and Zk threads.
+        connect.stop();
+        IntegrationTestRestExtension.instance = null;
+    }
+
+    private boolean extensionIsRegistered() {
+        try {
+            String extensionUrl = connect.endpointForResource("integration-test-rest-extension/registered");
+            Response response = connect.requestGet(extensionUrl);
+            return response.getStatus() < BAD_REQUEST.getStatusCode();
+        } catch (ConnectException e) {
+            return false;
+        }
+    }
+
+    private boolean verifyConnectorHealthAndConfig(
+        String connectorName,
+        ConnectorHealth expectedHealth,
+        Map expectedConfig
+    ) {
+        ConnectClusterState clusterState =
+            IntegrationTestRestExtension.instance.restPluginContext.clusterState();
+        
+        ConnectorHealth actualHealth = clusterState.connectorHealth(connectorName);
+        if (actualHealth.tasksState().isEmpty()) {
+            // Happens if the task has been started but its status has not yet been picked up from
+            // the status topic by the worker.
+            return false;
+        }
+        Map actualConfig = clusterState.connectorConfig(connectorName);
+
+        assertEquals(expectedConfig, actualConfig);
+        assertEquals(expectedHealth, actualHealth);
+
+        return true;
+    }
+
+    public static class IntegrationTestRestExtension implements ConnectRestExtension {
+        private static IntegrationTestRestExtension instance;
+
+        public ConnectRestExtensionContext restPluginContext;
+
+        @Override
+        public void register(ConnectRestExtensionContext restPluginContext) {
+            instance = this;
+            this.restPluginContext = restPluginContext;
+            // Immediately request a list of connectors to confirm that the context and its fields
+            // has been fully initialized and there is no risk of deadlock
+            restPluginContext.clusterState().connectors();
+            // Install a new REST resource that can be used to confirm that the extension has been
+            // successfully registered
+            restPluginContext.configurable().register(new IntegrationTestRestExtensionResource());
+        }
+    
+        @Override
+        public void close() {
+        }
+    
+        @Override
+        public void configure(Map configs) {
+        }
+    
+        @Override
+        public String version() {
+            return "test";
+        }
+
+        @Path("integration-test-rest-extension")
+        public static class IntegrationTestRestExtensionResource {
+
+            @GET
+            @Path("/registered")
+            public boolean isRegistered() {
+                return true;
+            }
+        }
+    }
+}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/SessionedProtocolIntegrationTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/SessionedProtocolIntegrationTest.java
new file mode 100644
index 0000000000000..8956a86e7c73c
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/SessionedProtocolIntegrationTest.java
@@ -0,0 +1,167 @@
+/*
+ * 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.integration;
+
+import org.apache.kafka.connect.runtime.distributed.ConnectProtocolCompatibility;
+import org.apache.kafka.connect.storage.StringConverter;
+import org.apache.kafka.connect.util.clusters.EmbeddedConnectCluster;
+import org.apache.kafka.test.IntegrationTest;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Ignore;
+import org.junit.Test;
+import org.junit.experimental.categories.Category;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.TimeUnit;
+
+import static javax.ws.rs.core.Response.Status.BAD_REQUEST;
+import static javax.ws.rs.core.Response.Status.FORBIDDEN;
+import static org.apache.kafka.connect.runtime.ConnectorConfig.CONNECTOR_CLASS_CONFIG;
+import static org.apache.kafka.connect.runtime.ConnectorConfig.KEY_CONVERTER_CLASS_CONFIG;
+import static org.apache.kafka.connect.runtime.ConnectorConfig.TASKS_MAX_CONFIG;
+import static org.apache.kafka.connect.runtime.ConnectorConfig.VALUE_CONVERTER_CLASS_CONFIG;
+import static org.apache.kafka.connect.runtime.SinkConnectorConfig.TOPICS_CONFIG;
+import static org.apache.kafka.connect.runtime.distributed.DistributedConfig.CONNECT_PROTOCOL_CONFIG;
+import static org.apache.kafka.connect.runtime.rest.InternalRequestSignature.SIGNATURE_ALGORITHM_HEADER;
+import static org.apache.kafka.connect.runtime.rest.InternalRequestSignature.SIGNATURE_HEADER;
+import static org.junit.Assert.assertEquals;
+
+/**
+ * A simple integration test to ensure that internal request validation becomes enabled with the
+ * "sessioned" protocol.
+ */
+@Category(IntegrationTest.class)
+public class SessionedProtocolIntegrationTest {
+
+    private static final Logger log = LoggerFactory.getLogger(SessionedProtocolIntegrationTest.class);
+
+    private static final String CONNECTOR_NAME = "connector";
+    private static final long CONNECTOR_SETUP_DURATION_MS = 60000;
+
+    private EmbeddedConnectCluster connect;
+    private ConnectorHandle connectorHandle;
+
+    @Before
+    public void setup() {
+        // setup Connect worker properties
+        Map workerProps = new HashMap<>();
+        workerProps.put(CONNECT_PROTOCOL_CONFIG, ConnectProtocolCompatibility.SESSIONED.protocol());
+
+        // build a Connect cluster backed by Kafka and Zk
+        connect = new EmbeddedConnectCluster.Builder()
+            .name("connect-cluster")
+            .numWorkers(2)
+            .numBrokers(1)
+            .workerProps(workerProps)
+            .build();
+
+        // start the clusters
+        connect.start();
+
+        // get a handle to the connector
+        connectorHandle = RuntimeHandles.get().connectorHandle(CONNECTOR_NAME);
+    }
+
+    @After
+    public void close() {
+        // stop all Connect, Kafka and Zk threads.
+        connect.stop();
+    }
+
+    @Test
+    @Ignore
+    // TODO: This test runs fine locally but fails on Jenkins. Ignoring for now, should revisit when
+    //       possible.
+    public void ensureInternalEndpointIsSecured() throws Throwable {
+        final String connectorTasksEndpoint = connect.endpointForResource(String.format(
+            "connectors/%s/tasks",
+            CONNECTOR_NAME
+        ));
+        final Map emptyHeaders = new HashMap<>();
+        final Map invalidSignatureHeaders = new HashMap<>();
+        invalidSignatureHeaders.put(SIGNATURE_HEADER, "S2Fma2Flc3F1ZQ==");
+        invalidSignatureHeaders.put(SIGNATURE_ALGORITHM_HEADER, "HmacSHA256");
+
+        // We haven't created the connector yet, but this should still return a 400 instead of a 404
+        // if the endpoint is secured
+        log.info(
+            "Making a POST request to the {} endpoint with no connector started and no signature header; " 
+                + "expecting 400 error response",
+            connectorTasksEndpoint
+        );
+        assertEquals(
+            BAD_REQUEST.getStatusCode(),
+            connect.requestPost(connectorTasksEndpoint, "[]", emptyHeaders).getStatus()
+        );
+
+        // Try again, but with an invalid signature
+        log.info(
+            "Making a POST request to the {} endpoint with no connector started and an invalid signature header; "
+                + "expecting 403 error response",
+            connectorTasksEndpoint
+        );
+        assertEquals(
+            FORBIDDEN.getStatusCode(),
+            connect.requestPost(connectorTasksEndpoint, "[]", invalidSignatureHeaders).getStatus()
+        );
+
+        // Create the connector now
+        // setup up props for the sink connector
+        Map connectorProps = new HashMap<>();
+        connectorProps.put(CONNECTOR_CLASS_CONFIG, MonitorableSinkConnector.class.getSimpleName());
+        connectorProps.put(TASKS_MAX_CONFIG, String.valueOf(1));
+        connectorProps.put(TOPICS_CONFIG, "test-topic");
+        connectorProps.put(KEY_CONVERTER_CLASS_CONFIG, StringConverter.class.getName());
+        connectorProps.put(VALUE_CONVERTER_CLASS_CONFIG, StringConverter.class.getName());
+
+        // start a sink connector
+        log.info("Starting the {} connector", CONNECTOR_NAME);
+        StartAndStopLatch startLatch = connectorHandle.expectedStarts(1);
+        connect.configureConnector(CONNECTOR_NAME, connectorProps);
+        startLatch.await(CONNECTOR_SETUP_DURATION_MS, TimeUnit.MILLISECONDS);
+
+
+        // Verify the exact same behavior, after starting the connector
+
+        // We haven't created the connector yet, but this should still return a 400 instead of a 404
+        // if the endpoint is secured
+        log.info(
+            "Making a POST request to the {} endpoint with the connector started and no signature header; "
+                + "expecting 400 error response",
+            connectorTasksEndpoint
+        );
+        assertEquals(
+            BAD_REQUEST.getStatusCode(),
+            connect.requestPost(connectorTasksEndpoint, "[]", emptyHeaders).getStatus()
+        );
+
+        // Try again, but with an invalid signature
+        log.info(
+            "Making a POST request to the {} endpoint with the connector started and an invalid signature header; "
+                + "expecting 403 error response",
+            connectorTasksEndpoint
+        );
+        assertEquals(
+            FORBIDDEN.getStatusCode(),
+            connect.requestPost(connectorTasksEndpoint, "[]", invalidSignatureHeaders).getStatus()
+        );
+    }
+}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/StartAndStopCounter.java b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/StartAndStopCounter.java
new file mode 100644
index 0000000000000..9ddfa1f685c97
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/StartAndStopCounter.java
@@ -0,0 +1,179 @@
+/*
+ * 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.integration;
+
+import java.util.List;
+import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import org.apache.kafka.common.utils.Time;
+
+public class StartAndStopCounter {
+
+    private final AtomicInteger startCounter = new AtomicInteger(0);
+    private final AtomicInteger stopCounter = new AtomicInteger(0);
+    private final List restartLatches = new CopyOnWriteArrayList<>();
+    private final Time clock;
+
+    public StartAndStopCounter() {
+        this(Time.SYSTEM);
+    }
+
+    public StartAndStopCounter(Time clock) {
+        this.clock = clock != null ? clock : Time.SYSTEM;
+    }
+
+    /**
+     * Record a start.
+     */
+    public void recordStart() {
+        startCounter.incrementAndGet();
+        restartLatches.forEach(StartAndStopLatch::recordStart);
+    }
+
+    /**
+     * Record a stop.
+     */
+    public void recordStop() {
+        stopCounter.incrementAndGet();
+        restartLatches.forEach(StartAndStopLatch::recordStop);
+    }
+
+    /**
+     * Get the number of starts.
+     *
+     * @return the number of starts
+     */
+    public int starts() {
+        return startCounter.get();
+    }
+
+    /**
+     * Get the number of stops.
+     *
+     * @return the number of stops
+     */
+    public int stops() {
+        return stopCounter.get();
+    }
+
+    /**
+     * Obtain a {@link StartAndStopLatch} that can be used to wait until the expected number of restarts
+     * has been completed.
+     *
+     * @param expectedStarts   the expected number of starts; may be 0
+     * @param expectedStops    the expected number of stops; may be 0
+     * @return the latch; never null
+     */
+    public StartAndStopLatch expectedRestarts(int expectedStarts, int expectedStops) {
+        return expectedRestarts(expectedStarts, expectedStops, null);
+    }
+
+    /**
+     * Obtain a {@link StartAndStopLatch} that can be used to wait until the expected number of restarts
+     * has been completed.
+     *
+     * @param expectedStarts   the expected number of starts; may be 0
+     * @param expectedStops    the expected number of stops; may be 0
+     * @param dependents       any dependent latches that must also complete in order for the
+     *                         resulting latch to complete
+     * @return the latch; never null
+     */
+    public StartAndStopLatch expectedRestarts(int expectedStarts, int expectedStops, List dependents) {
+        StartAndStopLatch latch = new StartAndStopLatch(expectedStarts, expectedStops, this::remove, dependents, clock);
+        restartLatches.add(latch);
+        return latch;
+    }
+
+    /**
+     * Obtain a {@link StartAndStopLatch} that can be used to wait until the expected number of restarts
+     * has been completed.
+     *
+     * @param expectedRestarts the expected number of restarts
+     * @return the latch; never null
+     */
+    public StartAndStopLatch expectedRestarts(int expectedRestarts) {
+        return expectedRestarts(expectedRestarts, expectedRestarts);
+    }
+
+    /**
+     * Obtain a {@link StartAndStopLatch} that can be used to wait until the expected number of restarts
+     * has been completed.
+     *
+     * @param expectedRestarts the expected number of restarts
+     * @param dependents       any dependent latches that must also complete in order for the
+     *                         resulting latch to complete
+     * @return the latch; never null
+     */
+    public StartAndStopLatch expectedRestarts(int expectedRestarts, List dependents) {
+        return expectedRestarts(expectedRestarts, expectedRestarts, dependents);
+    }
+
+    /**
+     * Obtain a {@link StartAndStopLatch} that can be used to wait until the expected number of starts
+     * has been completed.
+     *
+     * @param expectedStarts the expected number of starts
+     * @return the latch; never null
+     */
+    public StartAndStopLatch expectedStarts(int expectedStarts) {
+        return expectedRestarts(expectedStarts, 0);
+    }
+
+    /**
+     * Obtain a {@link StartAndStopLatch} that can be used to wait until the expected number of starts
+     * has been completed.
+     *
+     * @param expectedStarts the expected number of starts
+     * @param dependents     any dependent latches that must also complete in order for the
+     *                       resulting latch to complete
+     * @return the latch; never null
+     */
+    public StartAndStopLatch expectedStarts(int expectedStarts, List dependents) {
+        return expectedRestarts(expectedStarts, 0, dependents);
+    }
+
+
+    /**
+     * Obtain a {@link StartAndStopLatch} that can be used to wait until the expected number of
+     * stops has been completed.
+     *
+     * @param expectedStops the expected number of stops
+     * @return the latch; never null
+     */
+    public StartAndStopLatch expectedStops(int expectedStops) {
+        return expectedRestarts(0, expectedStops);
+    }
+
+    /**
+     * Obtain a {@link StartAndStopLatch} that can be used to wait until the expected number of
+     * stops has been completed.
+     *
+     * @param expectedStops the expected number of stops
+     * @param dependents    any dependent latches that must also complete in order for the
+     *                      resulting latch to complete
+     * @return the latch; never null
+     */
+    public StartAndStopLatch expectedStops(int expectedStops, List dependents) {
+        return expectedRestarts(0, expectedStops, dependents);
+    }
+
+    protected void remove(StartAndStopLatch restartLatch) {
+        restartLatches.remove(restartLatch);
+    }
+}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/StartAndStopCounterTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/StartAndStopCounterTest.java
new file mode 100644
index 0000000000000..7820a6d94d8e3
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/StartAndStopCounterTest.java
@@ -0,0 +1,116 @@
+/*
+ * 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.integration;
+
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+
+import org.apache.kafka.common.utils.MockTime;
+import org.apache.kafka.common.utils.Time;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+public class StartAndStopCounterTest {
+
+    private StartAndStopCounter counter;
+    private Time clock;
+    private ExecutorService waiters;
+    private StartAndStopLatch latch;
+
+    @Before
+    public void setup() {
+        clock = new MockTime();
+        counter = new StartAndStopCounter(clock);
+    }
+
+    @After
+    public void teardown() {
+        if (waiters != null) {
+            try {
+                waiters.shutdownNow();
+            } finally {
+                waiters = null;
+            }
+        }
+    }
+
+    @Test
+    public void shouldRecordStarts() {
+        assertEquals(0, counter.starts());
+        counter.recordStart();
+        assertEquals(1, counter.starts());
+        counter.recordStart();
+        assertEquals(2, counter.starts());
+        assertEquals(2, counter.starts());
+    }
+
+    @Test
+    public void shouldRecordStops() {
+        assertEquals(0, counter.stops());
+        counter.recordStop();
+        assertEquals(1, counter.stops());
+        counter.recordStop();
+        assertEquals(2, counter.stops());
+        assertEquals(2, counter.stops());
+    }
+
+    @Test
+    public void shouldExpectRestarts() throws Exception {
+        waiters = Executors.newSingleThreadExecutor();
+
+        latch = counter.expectedRestarts(1);
+        Future future = asyncAwait(100, TimeUnit.MILLISECONDS);
+
+        clock.sleep(1000);
+        counter.recordStop();
+        counter.recordStart();
+        assertTrue(future.get(200, TimeUnit.MILLISECONDS));
+        assertTrue(future.isDone());
+    }
+    @Test
+    public void shouldFailToWaitForRestartThatNeverHappens() throws Exception {
+        waiters = Executors.newSingleThreadExecutor();
+
+        latch = counter.expectedRestarts(1);
+        Future future = asyncAwait(100, TimeUnit.MILLISECONDS);
+
+        clock.sleep(1000);
+        // Record a stop but NOT a start
+        counter.recordStop();
+        assertFalse(future.get(200, TimeUnit.MILLISECONDS));
+        assertTrue(future.isDone());
+    }
+
+    private Future asyncAwait(long duration, TimeUnit unit) {
+        return waiters.submit(() -> {
+            try {
+                return latch.await(duration, unit);
+            } catch (InterruptedException e) {
+                Thread.interrupted();
+                return false;
+            }
+        });
+    }
+}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/StartAndStopLatch.java b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/StartAndStopLatch.java
new file mode 100644
index 0000000000000..b77007c5f2e13
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/StartAndStopLatch.java
@@ -0,0 +1,118 @@
+/*
+ * 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.integration;
+
+import java.util.List;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.function.Consumer;
+
+import org.apache.kafka.common.utils.Time;
+
+/**
+ * A latch that can be used to count down the number of times a connector and/or tasks have
+ * been started and stopped.
+ */
+public class StartAndStopLatch {
+    private final CountDownLatch startLatch;
+    private final CountDownLatch stopLatch;
+    private final List dependents;
+    private final Consumer uponCompletion;
+    private final Time clock;
+
+    StartAndStopLatch(int expectedStarts, int expectedStops, Consumer uponCompletion,
+                 List dependents, Time clock) {
+        this.startLatch = new CountDownLatch(expectedStarts < 0 ? 0 : expectedStarts);
+        this.stopLatch = new CountDownLatch(expectedStops < 0 ? 0 : expectedStops);
+        this.dependents = dependents;
+        this.uponCompletion = uponCompletion;
+        this.clock = clock;
+    }
+
+    protected void recordStart() {
+        startLatch.countDown();
+    }
+
+    protected void recordStop() {
+        stopLatch.countDown();
+    }
+
+    /**
+     * Causes the current thread to wait until the latch has counted down the starts and
+     * stops to zero, unless the thread is {@linkplain Thread#interrupt interrupted},
+     * or the specified waiting time elapses.
+     *
+     * 
      If the current counts are zero then this method returns immediately
+     * with the value {@code true}.
+     *
+     * 
      If the current count is greater than zero then the current
+     * thread becomes disabled for thread scheduling purposes and lies
+     * dormant until one of three things happen:
+     * 
        
+     * 
      • The counts reach zero due to invocations of the {@link #recordStart()} and
+     * {@link #recordStop()} methods; or
+     * 
      • Some other thread {@linkplain Thread#interrupt interrupts}
+     * the current thread; or
+     * 
      • The specified waiting time elapses.
+     * 
      
+     *
+     * 
      If the count reaches zero then the method returns with the
+     * value {@code true}.
+     *
+     * 
      If the current thread:
+     * 
        
+     * 
      • has its interrupted status set on entry to this method; or
+     * 
      • is {@linkplain Thread#interrupt interrupted} while waiting,
+     * 
      
+     * then {@link InterruptedException} is thrown and the current thread's
+     * interrupted status is cleared.
+     *
+     * 
      If the specified waiting time elapses then the value {@code false}
+     * is returned.  If the time is less than or equal to zero, the method
+     * will not wait at all.
+     *
+     * @param timeout the maximum time to wait
+     * @param unit    the time unit of the {@code timeout} argument
+     * @return {@code true} if the counts reached zero and {@code false}
+     *         if the waiting time elapsed before the counts reached zero
+     * @throws InterruptedException if the current thread is interrupted
+     *         while waiting
+     */
+    public boolean await(long timeout, TimeUnit unit) throws InterruptedException {
+        final long start = clock.milliseconds();
+        final long end = start + unit.toMillis(timeout);
+        if (!startLatch.await(end - start, TimeUnit.MILLISECONDS)) {
+            return false;
+        }
+        if (!stopLatch.await(end - clock.milliseconds(), TimeUnit.MILLISECONDS)) {
+            return false;
+        }
+
+        if (dependents != null) {
+            for (StartAndStopLatch dependent : dependents) {
+                if (!dependent.await(end - clock.milliseconds(), TimeUnit.MILLISECONDS)) {
+                    return false;
+                }
+            }
+        }
+        if (uponCompletion != null) {
+            uponCompletion.accept(this);
+        }
+        return true;
+    }
+}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/StartAndStopLatchTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/StartAndStopLatchTest.java
new file mode 100644
index 0000000000000..d2732ea4bb28d
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/StartAndStopLatchTest.java
@@ -0,0 +1,137 @@
+/*
+ * 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.integration;
+
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+import org.apache.kafka.common.utils.MockTime;
+import org.apache.kafka.common.utils.Time;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+public class StartAndStopLatchTest {
+
+    private Time clock;
+    private StartAndStopLatch latch;
+    private List dependents;
+    private AtomicBoolean completed = new AtomicBoolean();
+    private ExecutorService waiters;
+    private Future future;
+
+    @Before
+    public void setup() {
+        clock = new MockTime();
+        waiters = Executors.newSingleThreadExecutor();
+    }
+
+    @After
+    public void teardown() {
+        if (waiters != null) {
+            waiters.shutdownNow();
+        }
+    }
+
+    @Test
+    public void shouldReturnFalseWhenAwaitingForStartToNeverComplete() throws Throwable {
+        latch = new StartAndStopLatch(1, 1, this::complete, dependents, clock);
+        future = asyncAwait(100);
+        clock.sleep(10);
+        assertFalse(future.get(200, TimeUnit.MILLISECONDS));
+        assertTrue(future.isDone());
+    }
+
+    @Test
+    public void shouldReturnFalseWhenAwaitingForStopToNeverComplete() throws Throwable {
+        latch = new StartAndStopLatch(1, 1, this::complete, dependents, clock);
+        future = asyncAwait(100);
+        latch.recordStart();
+        clock.sleep(10);
+        assertFalse(future.get(200, TimeUnit.MILLISECONDS));
+        assertTrue(future.isDone());
+    }
+
+    @Test
+    public void shouldReturnTrueWhenAwaitingForStartAndStopToComplete() throws Throwable {
+        latch = new StartAndStopLatch(1, 1, this::complete, dependents, clock);
+        future = asyncAwait(100);
+        latch.recordStart();
+        latch.recordStop();
+        clock.sleep(10);
+        assertTrue(future.get(200, TimeUnit.MILLISECONDS));
+        assertTrue(future.isDone());
+    }
+
+    @Test
+    public void shouldReturnFalseWhenAwaitingForDependentLatchToComplete() throws Throwable {
+        StartAndStopLatch depLatch = new StartAndStopLatch(1, 1, this::complete, null, clock);
+        dependents = Collections.singletonList(depLatch);
+        latch = new StartAndStopLatch(1, 1, this::complete, dependents, clock);
+
+        future = asyncAwait(100);
+        latch.recordStart();
+        latch.recordStop();
+        clock.sleep(10);
+        assertFalse(future.get(200, TimeUnit.MILLISECONDS));
+        assertTrue(future.isDone());
+    }
+
+    @Test
+    public void shouldReturnTrueWhenAwaitingForStartAndStopAndDependentLatch() throws Throwable {
+        StartAndStopLatch depLatch = new StartAndStopLatch(1, 1, this::complete, null, clock);
+        dependents = Collections.singletonList(depLatch);
+        latch = new StartAndStopLatch(1, 1, this::complete, dependents, clock);
+
+        future = asyncAwait(100);
+        latch.recordStart();
+        latch.recordStop();
+        depLatch.recordStart();
+        depLatch.recordStop();
+        clock.sleep(10);
+        assertTrue(future.get(200, TimeUnit.MILLISECONDS));
+        assertTrue(future.isDone());
+    }
+
+    private Future asyncAwait(long duration) {
+        return asyncAwait(duration, TimeUnit.MILLISECONDS);
+    }
+
+    private Future asyncAwait(long duration, TimeUnit unit) {
+        return waiters.submit(() -> {
+            try {
+                return latch.await(duration, unit);
+            } catch (InterruptedException e) {
+                Thread.interrupted();
+                return false;
+            }
+        });
+    }
+
+    private void complete(StartAndStopLatch latch) {
+        completed.set(true);
+    }
+}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/TaskHandle.java b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/TaskHandle.java
index de3d9240d1be7..1159cb8fcc714 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/TaskHandle.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/TaskHandle.java
@@ -23,6 +23,7 @@
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.IntStream;
 
 /**
  * A handle to an executing task in a worker. Use this class to record progress, for example: number of records seen
@@ -35,9 +36,12 @@ public class TaskHandle {
     private final String taskId;
     private final ConnectorHandle connectorHandle;
     private final AtomicInteger partitionsAssigned = new AtomicInteger(0);
+    private final StartAndStopCounter startAndStopCounter = new StartAndStopCounter();
 
     private CountDownLatch recordsRemainingLatch;
+    private CountDownLatch recordsToCommitLatch;
     private int expectedRecords = -1;
+    private int expectedCommits = -1;
 
     public TaskHandle(ConnectorHandle connectorHandle, String taskId) {
         log.info("Created task {} for connector {}", taskId, connectorHandle);
@@ -46,7 +50,7 @@ public TaskHandle(ConnectorHandle connectorHandle, String taskId) {
     }
 
     /**
-     * Record a message arrival at the task.
+     * Record a message arrival at the task and the connector overall.
      */
     public void record() {
         if (recordsRemainingLatch != null) {
@@ -55,14 +59,58 @@ public void record() {
         connectorHandle.record();
     }
 
+    /**
+     * Record arrival of a batch of messages at the task and the connector overall.
+     *
+     * @param batchSize the number of messages
+     */
+    public void record(int batchSize) {
+        if (recordsRemainingLatch != null) {
+            IntStream.range(0, batchSize).forEach(i -> recordsRemainingLatch.countDown());
+        }
+        connectorHandle.record(batchSize);
+    }
+
+    /**
+     * Record a message commit from the task and the connector overall.
+     */
+    public void commit() {
+        if (recordsToCommitLatch != null) {
+            recordsToCommitLatch.countDown();
+        }
+        connectorHandle.commit();
+    }
+
+    /**
+     * Record commit on a batch of messages from the task and the connector overall.
+     *
+     * @param batchSize the number of messages
+     */
+    public void commit(int batchSize) {
+        if (recordsToCommitLatch != null) {
+            IntStream.range(0, batchSize).forEach(i -> recordsToCommitLatch.countDown());
+        }
+        connectorHandle.commit(batchSize);
+    }
+
     /**
      * Set the number of expected records for this task.
      *
-     * @param expectedRecords number of records
+     * @param expected number of records
+     */
+    public void expectedRecords(int expected) {
+        expectedRecords = expected;
+        recordsRemainingLatch = new CountDownLatch(expected);
+    }
+
+    /**
+     * Set the number of expected record commits performed by this task.
+     *
+     * @param expected number of commits
      */
-    public void expectedRecords(int expectedRecords) {
-        this.expectedRecords = expectedRecords;
-        this.recordsRemainingLatch = new CountDownLatch(expectedRecords);
+    public void expectedCommits(int expected) {
+        expectedRecords = expected;
+        recordsToCommitLatch = new CountDownLatch(expected);
     }
 
     /**
@@ -82,24 +130,111 @@ public int partitionsAssigned() {
     }
 
     /**
-     * Wait for this task to receive the expected number of records.
+     * Wait up to the specified number of milliseconds for this task to meet the expected number of
+     * records as defined by {@code expectedRecords}.
      *
-     * @param consumeMaxDurationMs max duration to wait for records
+     * @param timeoutMillis number of milliseconds to wait for records
      * @throws InterruptedException if another threads interrupts this one while waiting for records
      */
-    public void awaitRecords(int consumeMaxDurationMs) throws InterruptedException {
+    public void awaitRecords(long timeoutMillis) throws InterruptedException {
+        awaitRecords(timeoutMillis, TimeUnit.MILLISECONDS);
+    }
+
+    /**
+     * Wait up to the specified timeout for this task to meet the expected number of records as
+     * defined by {@code expectedRecords}.
+     *
+     * @param timeout duration to wait for records
+     * @param unit    the unit of duration; may not be null
+     * @throws InterruptedException if another threads interrupts this one while waiting for records
+     */
+    public void awaitRecords(long timeout, TimeUnit unit) throws InterruptedException {
         if (recordsRemainingLatch == null) {
             throw new IllegalStateException("Illegal state encountered. expectedRecords() was not set for this task?");
         }
-        if (!recordsRemainingLatch.await(consumeMaxDurationMs, TimeUnit.MILLISECONDS)) {
-            String msg = String.format("Insufficient records seen by task %s in %d millis. Records expected=%d, actual=%d",
+        if (!recordsRemainingLatch.await(timeout, unit)) {
+            String msg = String.format(
+                    "Insufficient records seen by task %s in %d millis. Records expected=%d, actual=%d",
                     taskId,
-                    consumeMaxDurationMs,
+                    unit.toMillis(timeout),
                     expectedRecords,
                     expectedRecords - recordsRemainingLatch.getCount());
             throw new DataException(msg);
         }
-        log.debug("Task {} saw {} records, expected {} records", taskId, expectedRecords - recordsRemainingLatch.getCount(), expectedRecords);
+        log.debug("Task {} saw {} records, expected {} records",
+                  taskId, expectedRecords - recordsRemainingLatch.getCount(), expectedRecords);
+    }
+
+    /**
+     * Wait up to the specified timeout in milliseconds for this task to meet the expected number
+     * of commits as defined by {@code expectedCommits}.
+     *
+     * @param timeoutMillis number of milliseconds to wait for commits
+     * @throws InterruptedException if another threads interrupts this one while waiting for commits
+     */
+    public void awaitCommits(long timeoutMillis) throws InterruptedException {
+        awaitCommits(timeoutMillis, TimeUnit.MILLISECONDS);
+    }
+
+    /**
+     * Wait up to the specified timeout for this task to meet the expected number of commits as
+     * defined by {@code expectedCommits}.
+     *
+     * @param timeout duration to wait for commits
+     * @param unit    the unit of duration; may not be null
+     * @throws InterruptedException if another threads interrupts this one while waiting for commits
+     */
+    public void awaitCommits(long timeout, TimeUnit unit) throws InterruptedException {
+        if (recordsToCommitLatch == null) {
+            throw new IllegalStateException("Illegal state encountered. expectedRecords() was not set for this task?");
+        }
+        if (!recordsToCommitLatch.await(timeout, unit)) {
+            String msg = String.format(
+                    "Insufficient records seen by task %s in %d millis. Records expected=%d, actual=%d",
+                    taskId,
+                    unit.toMillis(timeout),
+                    expectedCommits,
+                    expectedCommits - recordsToCommitLatch.getCount());
+            throw new DataException(msg);
+        }
+        log.debug("Task {} saw {} records, expected {} records",
+                  taskId, expectedCommits - recordsToCommitLatch.getCount(), expectedCommits);
+    }
+
+    /**
+     * Record that this task has been stopped. This should be called by the task.
+     */
+    public void recordTaskStart() {
+        startAndStopCounter.recordStart();
+    }
+
+    /**
+     * Record that this task has been stopped. This should be called by the task.
+     */
+    public void recordTaskStop() {
+        startAndStopCounter.recordStop();
+    }
+
+    /**
+     * Obtain a {@link StartAndStopLatch} that can be used to wait until this task has completed the
+     * expected number of starts.
+     *
+     * @param expectedStarts    the expected number of starts
+     * @return the latch; never null
+     */
+    public StartAndStopLatch expectedStarts(int expectedStarts) {
+        return startAndStopCounter.expectedStarts(expectedStarts);
+    }
+
+    /**
+     * Obtain a {@link StartAndStopLatch} that can be used to wait until this task has completed the
+     * expected number of starts.
+     *
+     * @param expectedStops    the expected number of stops
+     * @return the latch; never null
+     */
+    public StartAndStopLatch expectedStops(int expectedStops) {
+        return startAndStopCounter.expectedStops(expectedStops);
     }
 
     @Override
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 fc557c641b776..fe57d57972c39 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
@@ -16,13 +16,23 @@
  */
 package org.apache.kafka.connect.runtime;
 
+import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.clients.producer.ProducerConfig;
 import org.apache.kafka.common.config.ConfigDef;
 import org.apache.kafka.common.config.ConfigException;
+import org.apache.kafka.common.config.ConfigTransformer;
+import org.apache.kafka.common.config.SaslConfigs;
+import org.apache.kafka.common.security.oauthbearer.internals.unsecured.OAuthBearerUnsecuredLoginCallbackHandler;
 import org.apache.kafka.connect.connector.ConnectRecord;
 import org.apache.kafka.connect.connector.Connector;
+import org.apache.kafka.connect.connector.policy.AllConnectorClientConfigOverridePolicy;
+import org.apache.kafka.connect.connector.policy.ConnectorClientConfigOverridePolicy;
+import org.apache.kafka.connect.connector.policy.NoneConnectorClientConfigOverridePolicy;
+import org.apache.kafka.connect.connector.policy.PrincipalConnectorClientConfigOverridePolicy;
 import org.apache.kafka.connect.runtime.distributed.ClusterConfigState;
 import org.apache.kafka.connect.runtime.isolation.PluginDesc;
 import org.apache.kafka.connect.runtime.isolation.Plugins;
+import org.apache.kafka.connect.runtime.rest.entities.ConfigInfo;
 import org.apache.kafka.connect.runtime.rest.entities.ConfigInfos;
 import org.apache.kafka.connect.runtime.rest.entities.ConnectorStateInfo;
 import org.apache.kafka.connect.runtime.rest.entities.ConnectorType;
@@ -51,7 +61,9 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.stream.Collectors;
 
+import static org.apache.kafka.connect.runtime.AbstractHerder.keysWithVariableValues;
 import static org.powermock.api.easymock.PowerMock.verifyAll;
 import static org.powermock.api.easymock.PowerMock.replayAll;
 import static org.easymock.EasyMock.strictMock;
@@ -105,10 +117,10 @@ public class AbstractHerderTest {
         TASK_CONFIGS_MAP.put(TASK1, TASK_CONFIG);
         TASK_CONFIGS_MAP.put(TASK2, TASK_CONFIG);
     }
-    private static final ClusterConfigState SNAPSHOT = new ClusterConfigState(1, Collections.singletonMap(CONN1, 3),
+    private static final ClusterConfigState SNAPSHOT = new ClusterConfigState(1, null, Collections.singletonMap(CONN1, 3),
             Collections.singletonMap(CONN1, CONN1_CONFIG), Collections.singletonMap(CONN1, TargetState.STARTED),
             TASK_CONFIGS_MAP, Collections.emptySet());
-    private static final ClusterConfigState SNAPSHOT_NO_TASKS = new ClusterConfigState(1, Collections.singletonMap(CONN1, 3),
+    private static final ClusterConfigState SNAPSHOT_NO_TASKS = new ClusterConfigState(1, null, Collections.singletonMap(CONN1, 3),
             Collections.singletonMap(CONN1, CONN1_CONFIG), Collections.singletonMap(CONN1, TargetState.STARTED),
             Collections.emptyMap(), Collections.emptySet());
 
@@ -116,6 +128,7 @@ public class AbstractHerderTest {
     private final String kafkaClusterId = "I4ZmrWqfT2e-upky_4fdPA";
     private final int generation = 5;
     private final String connector = "connector";
+    private final ConnectorClientConfigOverridePolicy noneConnectorClientConfigOverridePolicy = new NoneConnectorClientConfigOverridePolicy();
 
     @MockStrict private Worker worker;
     @MockStrict private WorkerConfigTransformer transformer;
@@ -124,13 +137,66 @@ public class AbstractHerderTest {
     @MockStrict private ConfigBackingStore configStore;
     @MockStrict private StatusBackingStore statusStore;
 
+    @Test
+    public void testConnectors() {
+        AbstractHerder herder = partialMockBuilder(AbstractHerder.class)
+            .withConstructor(
+                Worker.class,
+                String.class,
+                String.class,
+                StatusBackingStore.class,
+                ConfigBackingStore.class,
+                ConnectorClientConfigOverridePolicy.class
+            )
+            .withArgs(worker, workerId, kafkaClusterId, statusStore, configStore, noneConnectorClientConfigOverridePolicy)
+            .addMockedMethod("generation")
+            .createMock();
+
+        EasyMock.expect(herder.generation()).andStubReturn(generation);
+        EasyMock.expect(herder.config(connector)).andReturn(null);
+        EasyMock.expect(configStore.snapshot()).andReturn(SNAPSHOT);
+        replayAll();
+        assertEquals(Collections.singleton(CONN1), new HashSet<>(herder.connectors()));
+        PowerMock.verifyAll();
+    }
+
+    @Test
+    public void testConnectorStatus() {
+        ConnectorTaskId taskId = new ConnectorTaskId(connector, 0);
+        AbstractHerder herder = partialMockBuilder(AbstractHerder.class)
+            .withConstructor(
+                Worker.class,
+                String.class,
+                String.class,
+                StatusBackingStore.class,
+                ConfigBackingStore.class,
+                ConnectorClientConfigOverridePolicy.class
+            )
+            .withArgs(worker, workerId, kafkaClusterId, statusStore, configStore, noneConnectorClientConfigOverridePolicy)
+            .addMockedMethod("generation")
+            .createMock();
+
+        EasyMock.expect(herder.generation()).andStubReturn(generation);
+        EasyMock.expect(herder.config(connector)).andReturn(null);
+        EasyMock.expect(statusStore.get(connector))
+            .andReturn(new ConnectorStatus(connector, AbstractStatus.State.RUNNING, workerId, generation));
+        EasyMock.expect(statusStore.getAll(connector))
+            .andReturn(Collections.singletonList(
+                new TaskStatus(taskId, AbstractStatus.State.UNASSIGNED, workerId, generation)));
+
+        replayAll();
+        ConnectorStateInfo csi = herder.connectorStatus(connector);
+        PowerMock.verifyAll();
+    }
+
     @Test
     public void connectorStatus() {
         ConnectorTaskId taskId = new ConnectorTaskId(connector, 0);
 
         AbstractHerder herder = partialMockBuilder(AbstractHerder.class)
-                .withConstructor(Worker.class, String.class, String.class, StatusBackingStore.class, ConfigBackingStore.class)
-                .withArgs(worker, workerId, kafkaClusterId, statusStore, configStore)
+                .withConstructor(Worker.class, String.class, String.class, StatusBackingStore.class, ConfigBackingStore.class,
+                                 ConnectorClientConfigOverridePolicy.class)
+                .withArgs(worker, workerId, kafkaClusterId, statusStore, configStore, noneConnectorClientConfigOverridePolicy)
                 .addMockedMethod("generation")
                 .createMock();
 
@@ -169,8 +235,9 @@ public void taskStatus() {
         String workerId = "workerId";
 
         AbstractHerder herder = partialMockBuilder(AbstractHerder.class)
-                .withConstructor(Worker.class, String.class, String.class, StatusBackingStore.class, ConfigBackingStore.class)
-                .withArgs(worker, workerId, kafkaClusterId, statusStore, configStore)
+                .withConstructor(Worker.class, String.class, String.class, StatusBackingStore.class, ConfigBackingStore.class,
+                                 ConnectorClientConfigOverridePolicy.class)
+                .withArgs(worker, workerId, kafkaClusterId, statusStore, configStore, noneConnectorClientConfigOverridePolicy)
                 .addMockedMethod("generation")
                 .createMock();
 
@@ -203,7 +270,7 @@ public TaskStatus answer() throws Throwable {
 
     @Test(expected = BadRequestException.class)
     public void testConfigValidationEmptyConfig() {
-        AbstractHerder herder = createConfigValidationHerder(TestSourceConnector.class);
+        AbstractHerder herder = createConfigValidationHerder(TestSourceConnector.class, noneConnectorClientConfigOverridePolicy);
         replayAll();
 
         herder.validateConnectorConfig(new HashMap());
@@ -213,7 +280,7 @@ public void testConfigValidationEmptyConfig() {
 
     @Test()
     public void testConfigValidationMissingName() {
-        AbstractHerder herder = createConfigValidationHerder(TestSourceConnector.class);
+        AbstractHerder herder = createConfigValidationHerder(TestSourceConnector.class, noneConnectorClientConfigOverridePolicy);
         replayAll();
 
         Map config = Collections.singletonMap(ConnectorConfig.CONNECTOR_CLASS_CONFIG, TestSourceConnector.class.getName());
@@ -238,7 +305,7 @@ public void testConfigValidationMissingName() {
 
     @Test(expected = ConfigException.class)
     public void testConfigValidationInvalidTopics() {
-        AbstractHerder herder = createConfigValidationHerder(TestSinkConnector.class);
+        AbstractHerder herder = createConfigValidationHerder(TestSinkConnector.class, noneConnectorClientConfigOverridePolicy);
         replayAll();
 
         Map config = new HashMap<>();
@@ -251,9 +318,39 @@ public void testConfigValidationInvalidTopics() {
         verifyAll();
     }
 
+    @Test(expected = ConfigException.class)
+    public void testConfigValidationTopicsWithDlq() {
+        AbstractHerder herder = createConfigValidationHerder(TestSinkConnector.class, noneConnectorClientConfigOverridePolicy);
+        replayAll();
+
+        Map config = new HashMap<>();
+        config.put(ConnectorConfig.CONNECTOR_CLASS_CONFIG, TestSinkConnector.class.getName());
+        config.put(SinkConnectorConfig.TOPICS_CONFIG, "topic1");
+        config.put(SinkConnectorConfig.DLQ_TOPIC_NAME_CONFIG, "topic1");
+
+        herder.validateConnectorConfig(config);
+
+        verifyAll();
+    }
+
+    @Test(expected = ConfigException.class)
+    public void testConfigValidationTopicsRegexWithDlq() {
+        AbstractHerder herder = createConfigValidationHerder(TestSinkConnector.class, noneConnectorClientConfigOverridePolicy);
+        replayAll();
+
+        Map config = new HashMap<>();
+        config.put(ConnectorConfig.CONNECTOR_CLASS_CONFIG, TestSinkConnector.class.getName());
+        config.put(SinkConnectorConfig.TOPICS_REGEX_CONFIG, "topic.*");
+        config.put(SinkConnectorConfig.DLQ_TOPIC_NAME_CONFIG, "topic1");
+
+        herder.validateConnectorConfig(config);
+
+        verifyAll();
+    }
+
     @Test()
     public void testConfigValidationTransformsExtendResults() {
-        AbstractHerder herder = createConfigValidationHerder(TestSourceConnector.class);
+        AbstractHerder herder = createConfigValidationHerder(TestSourceConnector.class, noneConnectorClientConfigOverridePolicy);
 
         // 2 transform aliases defined -> 2 plugin lookups
         Set> transformations = new HashSet<>();
@@ -299,6 +396,93 @@ public void testConfigValidationTransformsExtendResults() {
         verifyAll();
     }
 
+    @Test()
+    public void testConfigValidationPrincipalOnlyOverride() {
+        AbstractHerder herder = createConfigValidationHerder(TestSourceConnector.class, new PrincipalConnectorClientConfigOverridePolicy());
+        replayAll();
+
+        Map config = new HashMap<>();
+        config.put(ConnectorConfig.CONNECTOR_CLASS_CONFIG, TestSourceConnector.class.getName());
+        config.put(ConnectorConfig.NAME_CONFIG, "connector-name");
+        config.put("required", "value"); // connector required config
+        String ackConfigKey = producerOverrideKey(ProducerConfig.ACKS_CONFIG);
+        String saslConfigKey = producerOverrideKey(SaslConfigs.SASL_JAAS_CONFIG);
+        config.put(ackConfigKey, "none");
+        config.put(saslConfigKey, "jaas_config");
+
+        ConfigInfos result = herder.validateConnectorConfig(config);
+        assertEquals(herder.connectorTypeForClass(config.get(ConnectorConfig.CONNECTOR_CLASS_CONFIG)), ConnectorType.SOURCE);
+
+        // We expect there to be errors due to now allowed override policy for ACKS.... Note that these assertions depend heavily on
+        // the config fields for SourceConnectorConfig, but we expect these to change rarely.
+        assertEquals(TestSourceConnector.class.getName(), result.name());
+        // Each transform also gets its own group
+        List expectedGroups = Arrays.asList(
+            ConnectorConfig.COMMON_GROUP,
+            ConnectorConfig.TRANSFORMS_GROUP,
+            ConnectorConfig.ERROR_GROUP
+        );
+        assertEquals(expectedGroups, result.groups());
+        assertEquals(1, result.errorCount());
+        // Base connector config has 13 fields, connector's configs add 2, and 2 producer overrides
+        assertEquals(17, result.values().size());
+        assertTrue(result.values().stream().anyMatch(
+            configInfo -> ackConfigKey.equals(configInfo.configValue().name()) && !configInfo.configValue().errors().isEmpty()));
+        assertTrue(result.values().stream().anyMatch(
+            configInfo -> saslConfigKey.equals(configInfo.configValue().name()) && configInfo.configValue().errors().isEmpty()));
+
+        verifyAll();
+    }
+
+    @Test
+    public void testConfigValidationAllOverride() {
+        AbstractHerder herder = createConfigValidationHerder(TestSourceConnector.class, new AllConnectorClientConfigOverridePolicy());
+        replayAll();
+
+        Map config = new HashMap<>();
+        config.put(ConnectorConfig.CONNECTOR_CLASS_CONFIG, TestSourceConnector.class.getName());
+        config.put(ConnectorConfig.NAME_CONFIG, "connector-name");
+        config.put("required", "value"); // connector required config
+        // Try to test a variety of configuration types: string, int, long, boolean, list, class
+        String protocolConfigKey = producerOverrideKey(CommonClientConfigs.SECURITY_PROTOCOL_CONFIG);
+        config.put(protocolConfigKey, "SASL_PLAINTEXT");
+        String maxRequestSizeConfigKey = producerOverrideKey(ProducerConfig.MAX_REQUEST_SIZE_CONFIG);
+        config.put(maxRequestSizeConfigKey, "420");
+        String maxBlockConfigKey = producerOverrideKey(ProducerConfig.MAX_BLOCK_MS_CONFIG);
+        config.put(maxBlockConfigKey, "28980");
+        String idempotenceConfigKey = producerOverrideKey(ProducerConfig.ENABLE_IDEMPOTENCE_CONFIG);
+        config.put(idempotenceConfigKey, "true");
+        String bootstrapServersConfigKey = producerOverrideKey(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG);
+        config.put(bootstrapServersConfigKey, "SASL_PLAINTEXT://localhost:12345,SASL_PLAINTEXT://localhost:23456");
+        String loginCallbackHandlerConfigKey = producerOverrideKey(SaslConfigs.SASL_LOGIN_CALLBACK_HANDLER_CLASS);
+        config.put(loginCallbackHandlerConfigKey, OAuthBearerUnsecuredLoginCallbackHandler.class.getName());
+
+        final Set overriddenClientConfigs = new HashSet<>();
+        overriddenClientConfigs.add(protocolConfigKey);
+        overriddenClientConfigs.add(maxRequestSizeConfigKey);
+        overriddenClientConfigs.add(maxBlockConfigKey);
+        overriddenClientConfigs.add(idempotenceConfigKey);
+        overriddenClientConfigs.add(bootstrapServersConfigKey);
+        overriddenClientConfigs.add(loginCallbackHandlerConfigKey);
+
+        ConfigInfos result = herder.validateConnectorConfig(config);
+        assertEquals(herder.connectorTypeForClass(config.get(ConnectorConfig.CONNECTOR_CLASS_CONFIG)), ConnectorType.SOURCE);
+
+        Map validatedOverriddenClientConfigs = new HashMap<>();
+        for (ConfigInfo configInfo : result.values()) {
+            String configName = configInfo.configKey().name();
+            if (overriddenClientConfigs.contains(configName)) {
+                validatedOverriddenClientConfigs.put(configName, configInfo.configValue().value());
+            }
+        }
+        Map rawOverriddenClientConfigs = config.entrySet().stream()
+            .filter(e -> overriddenClientConfigs.contains(e.getKey()))
+            .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
+
+        assertEquals(rawOverriddenClientConfigs, validatedOverriddenClientConfigs);
+        verifyAll();
+    }
+
     @Test
     public void testReverseTransformConfigs() {
         // Construct a task config with constant values for TEST_KEY and TEST_KEY2
@@ -322,15 +506,43 @@ public void testReverseTransformConfigs() {
         assertFalse(reverseTransformed.get(0).containsKey(TEST_KEY3));
     }
 
-    private AbstractHerder createConfigValidationHerder(Class connectorClass) {
+    @Test
+    public void testConfigProviderRegex() {
+        testConfigProviderRegex("\"${::}\"");
+        testConfigProviderRegex("${::}");
+        testConfigProviderRegex("\"${:/a:somevar}\"");
+        testConfigProviderRegex("\"${file::somevar}\"");
+        testConfigProviderRegex("${file:/a/b/c:}");
+        testConfigProviderRegex("${file:/tmp/somefile.txt:somevar}");
+        testConfigProviderRegex("\"${file:/tmp/somefile.txt:somevar}\"");
+        testConfigProviderRegex("plain.PlainLoginModule required username=\"${file:/tmp/somefile.txt:somevar}\"");
+        testConfigProviderRegex("plain.PlainLoginModule required username=${file:/tmp/somefile.txt:somevar}");
+        testConfigProviderRegex("plain.PlainLoginModule required username=${file:/tmp/somefile.txt:somevar} not null");
+        testConfigProviderRegex("plain.PlainLoginModule required username=${file:/tmp/somefile.txt:somevar} password=${file:/tmp/somefile.txt:othervar}");
+        testConfigProviderRegex("plain.PlainLoginModule required username", false);
+    }
+
+    private void testConfigProviderRegex(String rawConnConfig) {
+        testConfigProviderRegex(rawConnConfig, true);
+    }
+    
+    private void testConfigProviderRegex(String rawConnConfig, boolean expected) {
+        Set keys = keysWithVariableValues(Collections.singletonMap("key", rawConnConfig), ConfigTransformer.DEFAULT_PATTERN);
+        boolean actual = keys != null && !keys.isEmpty() && keys.contains("key");
+        assertEquals(String.format("%s should have matched regex", rawConnConfig), expected, actual);
+    }
+
+    private AbstractHerder createConfigValidationHerder(Class connectorClass,
+                                                        ConnectorClientConfigOverridePolicy connectorClientConfigOverridePolicy) {
 
 
         ConfigBackingStore configStore = strictMock(ConfigBackingStore.class);
         StatusBackingStore statusStore = strictMock(StatusBackingStore.class);
 
         AbstractHerder herder = partialMockBuilder(AbstractHerder.class)
-                .withConstructor(Worker.class, String.class, String.class, StatusBackingStore.class, ConfigBackingStore.class)
-                .withArgs(worker, workerId, kafkaClusterId, statusStore, configStore)
+                .withConstructor(Worker.class, String.class, String.class, StatusBackingStore.class, ConfigBackingStore.class,
+                                 ConnectorClientConfigOverridePolicy.class)
+                .withArgs(worker, workerId, kafkaClusterId, statusStore, configStore, connectorClientConfigOverridePolicy)
                 .addMockedMethod("generation")
                 .createMock();
         EasyMock.expect(herder.generation()).andStubReturn(generation);
@@ -380,4 +592,8 @@ private abstract class BogusSourceConnector extends SourceConnector {
 
     private abstract class BogusSourceTask extends SourceTask {
     }
+
+    private static String producerOverrideKey(String config) {
+        return ConnectorConfig.CONNECTOR_CLIENT_PRODUCER_OVERRIDES_PREFIX + config;
+    }
 }
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/ConnectorConfigTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/ConnectorConfigTest.java
index fe1bf26465468..f674a8e22f244 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/ConnectorConfigTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/ConnectorConfigTest.java
@@ -177,4 +177,77 @@ public void multipleTransforms() {
         assertEquals(84, ((SimpleTransformation) transformations.get(1)).magicNumber);
     }
 
+    @Test
+    public void abstractTransform() {
+        Map props = new HashMap<>();
+        props.put("name", "test");
+        props.put("connector.class", TestConnector.class.getName());
+        props.put("transforms", "a");
+        props.put("transforms.a.type", AbstractTransformation.class.getName());
+        try {
+            new ConnectorConfig(MOCK_PLUGINS, props);
+        } catch (ConfigException ex) {
+            assertTrue(
+                ex.getMessage().contains("Transformation is abstract and cannot be created.")
+            );
+        }
+    }
+    @Test
+    public void abstractKeyValueTransform() {
+        Map props = new HashMap<>();
+        props.put("name", "test");
+        props.put("connector.class", TestConnector.class.getName());
+        props.put("transforms", "a");
+        props.put("transforms.a.type", AbstractKeyValueTransformation.class.getName());
+        try {
+            new ConnectorConfig(MOCK_PLUGINS, props);
+        } catch (ConfigException ex) {
+            assertTrue(
+                ex.getMessage().contains("Transformation is abstract and cannot be created.")
+            );
+            assertTrue(
+                ex.getMessage().contains(AbstractKeyValueTransformation.Key.class.getName())
+            );
+            assertTrue(
+                ex.getMessage().contains(AbstractKeyValueTransformation.Value.class.getName())
+            );
+        }
+    }
+
+    public static abstract class AbstractTransformation> implements Transformation  {
+
+    }
+
+    public static abstract class AbstractKeyValueTransformation> implements Transformation  {
+        @Override
+        public R apply(R record) {
+            return null;
+        }
+
+        @Override
+        public ConfigDef config() {
+            return new ConfigDef();
+        }
+
+        @Override
+        public void close() {
+
+        }
+
+        @Override
+        public void configure(Map configs) {
+
+        }
+
+
+        public static class Key extends AbstractKeyValueTransformation {
+
+
+        }
+        public static class Value extends AbstractKeyValueTransformation {
+
+        }
+    }
+
+
 }
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/ErrorHandlingTaskTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/ErrorHandlingTaskTest.java
index 5d223f4713ec0..bb42fc6433882 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/ErrorHandlingTaskTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/ErrorHandlingTaskTest.java
@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.connect.runtime;
 
+import java.util.Arrays;
 import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.clients.consumer.ConsumerRecords;
@@ -33,6 +34,7 @@
 import org.apache.kafka.connect.json.JsonConverter;
 import org.apache.kafka.connect.runtime.distributed.ClusterConfigState;
 import org.apache.kafka.connect.runtime.errors.ErrorHandlingMetrics;
+import org.apache.kafka.connect.runtime.errors.ErrorReporter;
 import org.apache.kafka.connect.runtime.errors.LogReporter;
 import org.apache.kafka.connect.runtime.errors.RetryWithToleranceOperator;
 import org.apache.kafka.connect.runtime.errors.ToleranceType;
@@ -46,7 +48,7 @@
 import org.apache.kafka.connect.source.SourceTask;
 import org.apache.kafka.connect.storage.Converter;
 import org.apache.kafka.connect.storage.HeaderConverter;
-import org.apache.kafka.connect.storage.OffsetStorageReader;
+import org.apache.kafka.connect.storage.OffsetStorageReaderImpl;
 import org.apache.kafka.connect.storage.OffsetStorageWriter;
 import org.apache.kafka.connect.transforms.Transformation;
 import org.apache.kafka.connect.transforms.util.SimpleConfig;
@@ -127,7 +129,7 @@ public class ErrorHandlingTaskTest {
     private KafkaProducer producer;
 
     @Mock
-    OffsetStorageReader offsetReader;
+    OffsetStorageReaderImpl offsetReader;
     @Mock
     OffsetStorageWriter offsetWriter;
 
@@ -162,6 +164,93 @@ public void tearDown() {
         }
     }
 
+    @Test
+    public void testSinkTasksCloseErrorReporters() throws Exception {
+        ErrorReporter reporter = EasyMock.mock(ErrorReporter.class);
+
+        RetryWithToleranceOperator retryWithToleranceOperator = operator();
+        retryWithToleranceOperator.metrics(errorHandlingMetrics);
+        retryWithToleranceOperator.reporters(singletonList(reporter));
+
+        createSinkTask(initialState, retryWithToleranceOperator);
+
+        expectInitializeTask();
+        reporter.close();
+        EasyMock.expectLastCall();
+        sinkTask.stop();
+        EasyMock.expectLastCall();
+
+        consumer.close();
+        EasyMock.expectLastCall();
+
+        PowerMock.replayAll();
+
+        workerSinkTask.initialize(TASK_CONFIG);
+        workerSinkTask.initializeAndStart();
+        workerSinkTask.close();
+
+        PowerMock.verifyAll();
+    }
+
+    @Test
+    public void testSourceTasksCloseErrorReporters() {
+        ErrorReporter reporter = EasyMock.mock(ErrorReporter.class);
+
+        RetryWithToleranceOperator retryWithToleranceOperator = operator();
+        retryWithToleranceOperator.metrics(errorHandlingMetrics);
+        retryWithToleranceOperator.reporters(singletonList(reporter));
+
+        createSourceTask(initialState, retryWithToleranceOperator);
+
+        sourceTask.stop();
+        PowerMock.expectLastCall();
+
+        producer.close(EasyMock.anyObject());
+        PowerMock.expectLastCall();
+
+        reporter.close();
+        EasyMock.expectLastCall();
+
+        PowerMock.replayAll();
+
+        workerSourceTask.initialize(TASK_CONFIG);
+        workerSourceTask.close();
+
+        PowerMock.verifyAll();
+    }
+
+    @Test
+    public void testCloseErrorReportersExceptionPropagation() {
+        ErrorReporter reporterA = EasyMock.mock(ErrorReporter.class);
+        ErrorReporter reporterB = EasyMock.mock(ErrorReporter.class);
+
+        RetryWithToleranceOperator retryWithToleranceOperator = operator();
+        retryWithToleranceOperator.metrics(errorHandlingMetrics);
+        retryWithToleranceOperator.reporters(Arrays.asList(reporterA, reporterB));
+
+        createSourceTask(initialState, retryWithToleranceOperator);
+
+        sourceTask.stop();
+        PowerMock.expectLastCall();
+
+        producer.close(EasyMock.anyObject());
+        PowerMock.expectLastCall();
+
+        // Even though the reporters throw exceptions, they should both still be closed.
+        reporterA.close();
+        EasyMock.expectLastCall().andThrow(new RuntimeException());
+
+        reporterB.close();
+        EasyMock.expectLastCall().andThrow(new RuntimeException());
+
+        PowerMock.replayAll();
+
+        workerSourceTask.initialize(TASK_CONFIG);
+        workerSourceTask.close();
+
+        PowerMock.verifyAll();
+    }
+
     @Test
     public void testErrorHandlingInSinkTasks() throws Exception {
         Map reportProps = new HashMap<>();
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/SourceTaskOffsetCommitterTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/SourceTaskOffsetCommitterTest.java
index baf0d8e0ebeef..fb7ed82af7399 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/SourceTaskOffsetCommitterTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/SourceTaskOffsetCommitterTest.java
@@ -147,6 +147,8 @@ public void testRemove() throws Exception {
         EasyMock.expect(taskFuture.cancel(eq(false))).andReturn(false);
         EasyMock.expect(taskFuture.isDone()).andReturn(false);
         EasyMock.expect(taskFuture.get()).andReturn(null);
+        EasyMock.expect(taskId.connector()).andReturn("MyConnector");
+        EasyMock.expect(taskId.task()).andReturn(1);
         PowerMock.replayAll();
 
         committers.put(taskId, taskFuture);
@@ -160,6 +162,8 @@ public void testRemove() throws Exception {
         EasyMock.expect(taskFuture.cancel(eq(false))).andReturn(false);
         EasyMock.expect(taskFuture.isDone()).andReturn(false);
         EasyMock.expect(taskFuture.get()).andThrow(new CancellationException());
+        EasyMock.expect(taskId.connector()).andReturn("MyConnector");
+        EasyMock.expect(taskId.task()).andReturn(1);
         mockLog.trace(EasyMock.anyString(), EasyMock.anyObject());
         PowerMock.expectLastCall();
         PowerMock.replayAll();
@@ -175,6 +179,8 @@ public void testRemove() throws Exception {
         EasyMock.expect(taskFuture.cancel(eq(false))).andReturn(false);
         EasyMock.expect(taskFuture.isDone()).andReturn(false);
         EasyMock.expect(taskFuture.get()).andThrow(new InterruptedException());
+        EasyMock.expect(taskId.connector()).andReturn("MyConnector");
+        EasyMock.expect(taskId.task()).andReturn(1);
         PowerMock.replayAll();
 
         try {
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/TestConverterWithHeaders.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/TestConverterWithHeaders.java
new file mode 100644
index 0000000000000..91e0999d2904e
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/TestConverterWithHeaders.java
@@ -0,0 +1,80 @@
+/*
+ * 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 java.io.UnsupportedEncodingException;
+import java.util.Map;
+import org.apache.kafka.common.header.Header;
+import org.apache.kafka.common.header.Headers;
+import org.apache.kafka.connect.data.Schema;
+import org.apache.kafka.connect.data.SchemaAndValue;
+import org.apache.kafka.connect.errors.DataException;
+import org.apache.kafka.connect.storage.Converter;
+
+/**
+ * This is a simple Converter implementation that uses "encoding" header to encode/decode strings via provided charset name
+ */
+public class TestConverterWithHeaders implements Converter {
+    private static final String HEADER_ENCODING = "encoding";
+
+    @Override
+    public void configure(Map configs, boolean isKey) {
+
+    }
+
+    @Override
+    public SchemaAndValue toConnectData(String topic, Headers headers, byte[] value) {
+        String encoding = extractEncoding(headers);
+
+        try {
+            return new SchemaAndValue(Schema.STRING_SCHEMA, new String(value, encoding));
+        } catch (UnsupportedEncodingException e) {
+            throw new DataException("Unsupported encoding: " + encoding, e);
+        }
+    }
+
+    @Override
+    public byte[] fromConnectData(String topic, Headers headers, Schema schema, Object value) {
+        String encoding = extractEncoding(headers);
+
+        try {
+            return ((String) value).getBytes(encoding);
+        } catch (UnsupportedEncodingException e) {
+            throw new DataException("Unsupported encoding: " + encoding, e);
+        }
+    }
+
+    private String extractEncoding(Headers headers) {
+        Header header = headers.lastHeader(HEADER_ENCODING);
+        if (header == null) {
+            throw new DataException("Header '" + HEADER_ENCODING + "' is required!");
+        }
+
+        return new String(header.value());
+    }
+
+
+    @Override
+    public SchemaAndValue toConnectData(String topic, byte[] value) {
+        throw new DataException("Headers are required for this converter!");
+    }
+
+    @Override
+    public byte[] fromConnectData(String topic, Schema schema, Object value) {
+        throw new DataException("Headers are required for this converter!");
+    }
+}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerConfigTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerConfigTest.java
new file mode 100644
index 0000000000000..33416b99e04c5
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerConfigTest.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.connect.runtime;
+
+import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.common.config.ConfigException;
+import org.junit.Test;
+
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.Map;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
+
+public class WorkerConfigTest {
+
+    @Test
+    public void testAdminListenersConfigAllowedValues() {
+        Map props = baseProps();
+
+        // no value set for "admin.listeners"
+        WorkerConfig config = new WorkerConfig(WorkerConfig.baseConfigDef(), props);
+        assertNull("Default value should be null.", config.getList(WorkerConfig.ADMIN_LISTENERS_CONFIG));
+
+        props.put(WorkerConfig.ADMIN_LISTENERS_CONFIG, "");
+        config = new WorkerConfig(WorkerConfig.baseConfigDef(), props);
+        assertTrue(config.getList(WorkerConfig.ADMIN_LISTENERS_CONFIG).isEmpty());
+
+        props.put(WorkerConfig.ADMIN_LISTENERS_CONFIG, "http://a.b:9999, https://a.b:7812");
+        config = new WorkerConfig(WorkerConfig.baseConfigDef(), props);
+        assertEquals(config.getList(WorkerConfig.ADMIN_LISTENERS_CONFIG), Arrays.asList("http://a.b:9999", "https://a.b:7812"));
+
+        new WorkerConfig(WorkerConfig.baseConfigDef(), props);
+    }
+
+    @Test(expected = ConfigException.class)
+    public void testAdminListenersNotAllowingEmptyStrings() {
+        Map props = baseProps();
+        props.put(WorkerConfig.ADMIN_LISTENERS_CONFIG, "http://a.b:9999,");
+        new WorkerConfig(WorkerConfig.baseConfigDef(), props);
+    }
+
+    @Test(expected = ConfigException.class)
+    public void testAdminListenersNotAllowingBlankStrings() {
+        Map props = baseProps();
+        props.put(WorkerConfig.ADMIN_LISTENERS_CONFIG, "http://a.b:9999, ,https://a.b:9999");
+        new WorkerConfig(WorkerConfig.baseConfigDef(), props);
+    }
+
+    private Map baseProps() {
+        Map props = new HashMap<>();
+        props.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092");
+        props.put(WorkerConfig.KEY_CONVERTER_CLASS_CONFIG, "org.apache.kafka.connect.json.JsonConverter");
+        props.put(WorkerConfig.VALUE_CONVERTER_CLASS_CONFIG, "org.apache.kafka.connect.json.JsonConverter");
+        return props;
+    }
+
+}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerConfigTransformerTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerConfigTransformerTest.java
index e30acb15ed4bd..6f4bda66904d7 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerConfigTransformerTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerConfigTransformerTest.java
@@ -20,6 +20,7 @@
 import org.apache.kafka.common.config.ConfigData;
 import org.apache.kafka.common.config.provider.ConfigProvider;
 import org.easymock.EasyMock;
+import static org.easymock.EasyMock.eq;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -34,6 +35,7 @@
 
 import static org.apache.kafka.connect.runtime.ConnectorConfig.CONFIG_RELOAD_ACTION_CONFIG;
 import static org.apache.kafka.connect.runtime.ConnectorConfig.CONFIG_RELOAD_ACTION_NONE;
+import static org.easymock.EasyMock.notNull;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNull;
 import static org.powermock.api.easymock.PowerMock.replayAll;
@@ -84,8 +86,7 @@ public void testReplaceVariableWithTTL() {
     @Test
     public void testReplaceVariableWithTTLAndScheduleRestart() {
         EasyMock.expect(worker.herder()).andReturn(herder);
-        EasyMock.expect(herder.restartConnector(1L, MY_CONNECTOR, null)).andReturn(requestId);
-
+        EasyMock.expect(herder.restartConnector(eq(1L), eq(MY_CONNECTOR), notNull())).andReturn(requestId);
         replayAll();
 
         Map result = configTransformer.transform(MY_CONNECTOR, Collections.singletonMap(MY_KEY, "${test:testPath:testKeyWithTTL}"));
@@ -95,13 +96,13 @@ public void testReplaceVariableWithTTLAndScheduleRestart() {
     @Test
     public void testReplaceVariableWithTTLFirstCancelThenScheduleRestart() {
         EasyMock.expect(worker.herder()).andReturn(herder);
-        EasyMock.expect(herder.restartConnector(1L, MY_CONNECTOR, null)).andReturn(requestId);
+        EasyMock.expect(herder.restartConnector(eq(1L), eq(MY_CONNECTOR), notNull())).andReturn(requestId);
 
         EasyMock.expect(worker.herder()).andReturn(herder);
         EasyMock.expectLastCall();
         requestId.cancel();
         EasyMock.expectLastCall();
-        EasyMock.expect(herder.restartConnector(10L, MY_CONNECTOR, null)).andReturn(requestId);
+        EasyMock.expect(herder.restartConnector(eq(10L), eq(MY_CONNECTOR), notNull())).andReturn(requestId);
 
         replayAll();
 
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 3e047ff9ed5ce..c6d25a4b82224 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
@@ -16,6 +16,8 @@
  */
 package org.apache.kafka.connect.runtime;
 
+import java.util.Arrays;
+import java.util.Iterator;
 import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.clients.consumer.ConsumerRecords;
@@ -25,11 +27,15 @@
 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.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.distributed.ClusterConfigState;
@@ -42,6 +48,7 @@
 import org.apache.kafka.connect.sink.SinkTask;
 import org.apache.kafka.connect.storage.Converter;
 import org.apache.kafka.connect.storage.HeaderConverter;
+import org.apache.kafka.connect.storage.StringConverter;
 import org.apache.kafka.connect.util.ConnectorTaskId;
 import org.easymock.Capture;
 import org.easymock.CaptureType;
@@ -77,6 +84,7 @@
 
 import static java.util.Arrays.asList;
 import static java.util.Collections.singleton;
+import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotEquals;
@@ -162,6 +170,10 @@ public void setUp() {
     }
 
     private void createTask(TargetState initialState) {
+        createTask(initialState, keyConverter, valueConverter, headerConverter);
+    }
+
+    private void createTask(TargetState initialState, Converter keyConverter, Converter valueConverter, HeaderConverter headerConverter) {
         workerTask = new WorkerSinkTask(
             taskId, sinkTask, statusListener, initialState, workerConfig, ClusterConfigState.EMPTY, metrics,
             keyConverter, valueConverter, headerConverter,
@@ -576,6 +588,10 @@ public Void answer() throws Throwable {
         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.>getInternalState(workerTask, "lastCommittedOffsets"));
@@ -585,6 +601,14 @@ public Void answer() throws Throwable {
         assertFalse(sinkTaskContext.getValue().isCommitRequested()); // should have been cleared
         assertEquals(offsets, Whitebox.>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);
@@ -821,6 +845,88 @@ public void run() {
         PowerMock.verifyAll();
     }
 
+    @Test
+    public void testSinkTasksHandleCloseErrors() throws Exception {
+        createTask(initialState);
+        expectInitializeTask();
+
+        // 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);
+        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();
+
+        // 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);
+        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.
@@ -1252,6 +1358,85 @@ public void testMetricsGroup() {
         assertEquals(30, metrics.currentMetricValueAsDouble(group1.metricGroup(), "put-batch-max-time-ms"), 0.001d);
     }
 
+    @Test
+    public void testHeaders() throws Exception {
+        Headers headers = new RecordHeaders();
+        headers.add("header_key", "header_value".getBytes());
+
+        createTask(initialState);
+
+        expectInitializeTask();
+        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();
+        TestConverterWithHeaders testConverter = new TestConverterWithHeaders();
+
+        createTask(initialState, stringConverter, testConverter, stringConverter);
+
+        expectInitializeTask();
+        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,
+                0L, 0, 0, keyA.getBytes(), valueA.getBytes(encodingA), headersA),
+            new ConsumerRecord<>(TOPIC, PARTITION, FIRST_OFFSET + recordsReturnedTp1 + 2, RecordBatch.NO_TIMESTAMP, TimestampType.NO_TIMESTAMP_TYPE,
+                0L, 0, 0, keyB.getBytes(), valueB.getBytes(encodingB), headersB)
+        ));
+
+        expectTransformation(2, null);
+
+        Capture> 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 iterator = records.getValue().iterator();
+
+        SinkRecord recordA = iterator.next();
+        assertEquals(keyA, recordA.key());
+        assertEquals(valueA, (String) recordA.value());
+
+        SinkRecord recordB = iterator.next();
+        assertEquals(keyB, recordB.key());
+        assertEquals(valueB, (String) recordB.value());
+
+        PowerMock.verifyAll();
+    }
+
     private void expectInitializeTask() throws Exception {
         consumer.subscribe(EasyMock.eq(asList(TOPIC)), EasyMock.capture(rebalanceListener));
         PowerMock.expectLastCall();
@@ -1334,17 +1519,25 @@ private void expectConsumerWakeup() {
     }
 
     private void expectConsumerPoll(final int numMessages) {
-        expectConsumerPoll(numMessages, RecordBatch.NO_TIMESTAMP, TimestampType.NO_TIMESTAMP_TYPE);
+        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(
                 new IAnswer>() {
                     @Override
                     public ConsumerRecords answer() throws Throwable {
                         List> records = new ArrayList<>();
                         for (int i = 0; i < numMessages; i++)
-                            records.add(new ConsumerRecord<>(TOPIC, PARTITION, FIRST_OFFSET + recordsReturnedTp1 + i, timestamp, timestampType, 0L, 0, 0, RAW_KEY, RAW_VALUE));
+                            records.add(new ConsumerRecord<>(TOPIC, PARTITION, FIRST_OFFSET + recordsReturnedTp1 + i, timestamp, timestampType, 0L, 0, 0, RAW_KEY, RAW_VALUE, headers));
                         recordsReturnedTp1 += numMessages;
                         return new ConsumerRecords<>(
                                 numMessages > 0 ?
@@ -1355,14 +1548,40 @@ public ConsumerRecords answer() throws Throwable {
                 });
     }
 
+    private void expectConsumerPoll(List> records) {
+        EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andAnswer(
+            new IAnswer>() {
+                @Override
+                public ConsumerRecords answer() throws Throwable {
+                    return new ConsumerRecords<>(
+                        records.isEmpty() ?
+                            Collections.>>emptyMap() :
+                            Collections.singletonMap(new TopicPartition(TOPIC, PARTITION), records)
+                    );
+                }
+            });
+    }
+
     private void expectConversionAndTransformation(final int numMessages) {
         expectConversionAndTransformation(numMessages, null);
     }
 
     private void expectConversionAndTransformation(final int numMessages, final String topicPrefix) {
-        EasyMock.expect(keyConverter.toConnectData(TOPIC, RAW_KEY)).andReturn(new SchemaAndValue(KEY_SCHEMA, KEY)).times(numMessages);
-        EasyMock.expect(valueConverter.toConnectData(TOPIC, RAW_VALUE)).andReturn(new SchemaAndValue(VALUE_SCHEMA, VALUE)).times(numMessages);
+        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 recordCapture = EasyMock.newCapture();
         EasyMock.expect(transformationChain.apply(EasyMock.capture(recordCapture)))
                 .andAnswer(new IAnswer() {
@@ -1377,7 +1596,8 @@ public SinkRecord answer() {
                                        origRecord.key(),
                                        origRecord.valueSchema(),
                                        origRecord.value(),
-                                       origRecord.timestamp()
+                                       origRecord.timestamp(),
+                                       origRecord.headers()
                                )
                                : origRecord;
                     }
@@ -1460,6 +1680,10 @@ private void assertMetrics(int minimumPollCountExpected) {
         double sendTotal = metrics.currentMetricValueAsDouble(sinkTaskGroup, "sink-record-send-total");
     }
 
+    private RecordHeaders emptyHeaders() {
+        return new RecordHeaders();
+    }
+
     private abstract static class TestSinkTask extends SinkTask  {
     }
 }
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskThreadedTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskThreadedTest.java
index 6e2b01ce7bc19..ab20dce2d26bd 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskThreadedTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskThreadedTest.java
@@ -23,6 +23,7 @@
 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.header.internals.RecordHeaders;
 import org.apache.kafka.common.record.TimestampType;
 import org.apache.kafka.common.utils.Time;
 import org.apache.kafka.connect.data.Schema;
@@ -572,8 +573,8 @@ public ConsumerRecords answer() throws Throwable {
                         return records;
                     }
                 });
-        EasyMock.expect(keyConverter.toConnectData(TOPIC, RAW_KEY)).andReturn(new SchemaAndValue(KEY_SCHEMA, KEY)).anyTimes();
-        EasyMock.expect(valueConverter.toConnectData(TOPIC, RAW_VALUE)).andReturn(new SchemaAndValue(VALUE_SCHEMA, VALUE)).anyTimes();
+        EasyMock.expect(keyConverter.toConnectData(TOPIC, emptyHeaders(), RAW_KEY)).andReturn(new SchemaAndValue(KEY_SCHEMA, KEY)).anyTimes();
+        EasyMock.expect(valueConverter.toConnectData(TOPIC, emptyHeaders(), RAW_VALUE)).andReturn(new SchemaAndValue(VALUE_SCHEMA, VALUE)).anyTimes();
 
         final Capture recordCapture = EasyMock.newCapture();
         EasyMock.expect(transformationChain.apply(EasyMock.capture(recordCapture))).andAnswer(
@@ -606,8 +607,8 @@ public ConsumerRecords answer() throws Throwable {
                         return records;
                     }
                 });
-        EasyMock.expect(keyConverter.toConnectData(TOPIC, RAW_KEY)).andReturn(new SchemaAndValue(KEY_SCHEMA, KEY));
-        EasyMock.expect(valueConverter.toConnectData(TOPIC, RAW_VALUE)).andReturn(new SchemaAndValue(VALUE_SCHEMA, VALUE));
+        EasyMock.expect(keyConverter.toConnectData(TOPIC, emptyHeaders(), RAW_KEY)).andReturn(new SchemaAndValue(KEY_SCHEMA, KEY));
+        EasyMock.expect(valueConverter.toConnectData(TOPIC, emptyHeaders(), RAW_VALUE)).andReturn(new SchemaAndValue(VALUE_SCHEMA, VALUE));
         sinkTask.put(EasyMock.anyObject(Collection.class));
         return EasyMock.expectLastCall();
     }
@@ -651,8 +652,8 @@ public ConsumerRecords answer() throws Throwable {
         consumer.seek(TOPIC_PARTITION, startOffset);
         EasyMock.expectLastCall();
 
-        EasyMock.expect(keyConverter.toConnectData(TOPIC, RAW_KEY)).andReturn(new SchemaAndValue(KEY_SCHEMA, KEY));
-        EasyMock.expect(valueConverter.toConnectData(TOPIC, RAW_VALUE)).andReturn(new SchemaAndValue(VALUE_SCHEMA, VALUE));
+        EasyMock.expect(keyConverter.toConnectData(TOPIC, emptyHeaders(), RAW_KEY)).andReturn(new SchemaAndValue(KEY_SCHEMA, KEY));
+        EasyMock.expect(valueConverter.toConnectData(TOPIC, emptyHeaders(), RAW_VALUE)).andReturn(new SchemaAndValue(VALUE_SCHEMA, VALUE));
         sinkTask.put(EasyMock.anyObject(Collection.class));
         return EasyMock.expectLastCall();
     }
@@ -694,6 +695,10 @@ public Object answer() throws Throwable {
         return capturedCallback;
     }
 
+    private RecordHeaders emptyHeaders() {
+        return new RecordHeaders();
+    }
+
     private static abstract class TestSinkTask extends SinkTask {
     }
 }
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSourceTaskTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSourceTaskTest.java
index 24a13c2be267d..ea3716d207111 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSourceTaskTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSourceTaskTest.java
@@ -16,14 +16,22 @@
  */
 package org.apache.kafka.connect.runtime;
 
+import java.nio.ByteBuffer;
 import org.apache.kafka.clients.producer.KafkaProducer;
 import org.apache.kafka.clients.producer.ProducerRecord;
 import org.apache.kafka.clients.producer.RecordMetadata;
 import org.apache.kafka.common.MetricName;
 import org.apache.kafka.common.TopicPartition;
-import org.apache.kafka.common.record.InvalidRecordException;
+import org.apache.kafka.common.InvalidRecordException;
+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.errors.TopicAuthorizationException;
 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.header.ConnectHeaders;
+import org.apache.kafka.connect.errors.ConnectException;
 import org.apache.kafka.connect.runtime.ConnectMetrics.MetricGroup;
 import org.apache.kafka.connect.runtime.WorkerSourceTask.SourceTaskMetricsGroup;
 import org.apache.kafka.connect.runtime.distributed.ClusterConfigState;
@@ -33,10 +41,11 @@
 import org.apache.kafka.connect.source.SourceRecord;
 import org.apache.kafka.connect.source.SourceTask;
 import org.apache.kafka.connect.source.SourceTaskContext;
+import org.apache.kafka.connect.storage.CloseableOffsetStorageReader;
 import org.apache.kafka.connect.storage.Converter;
 import org.apache.kafka.connect.storage.HeaderConverter;
-import org.apache.kafka.connect.storage.OffsetStorageReader;
 import org.apache.kafka.connect.storage.OffsetStorageWriter;
+import org.apache.kafka.connect.storage.StringConverter;
 import org.apache.kafka.connect.util.Callback;
 import org.apache.kafka.connect.util.ConnectorTaskId;
 import org.apache.kafka.connect.util.ThreadedTest;
@@ -75,6 +84,7 @@
 import static org.junit.Assert.assertTrue;
 
 @PowerMockIgnore({"javax.management.*",
+                  "org.apache.log4j.*",
                   "org.apache.kafka.connect.runtime.isolation.*"})
 @RunWith(PowerMockRunner.class)
 public class WorkerSourceTaskTest extends ThreadedTest {
@@ -104,7 +114,7 @@ public class WorkerSourceTaskTest extends ThreadedTest {
     @Mock private HeaderConverter headerConverter;
     @Mock private TransformationChain transformationChain;
     @Mock private KafkaProducer producer;
-    @Mock private OffsetStorageReader offsetReader;
+    @Mock private CloseableOffsetStorageReader offsetReader;
     @Mock private OffsetStorageWriter offsetWriter;
     @Mock private ClusterConfigState clusterConfigState;
     private WorkerSourceTask workerTask;
@@ -150,6 +160,10 @@ private void createWorkerTask() {
     }
 
     private void createWorkerTask(TargetState initialState) {
+        createWorkerTask(initialState, keyConverter, valueConverter, headerConverter);
+    }
+
+    private void createWorkerTask(TargetState initialState, Converter keyConverter, Converter valueConverter, HeaderConverter headerConverter) {
         workerTask = new WorkerSourceTask(taskId, sourceTask, statusListener, initialState, keyConverter, valueConverter, headerConverter,
                 transformationChain, producer, offsetReader, offsetWriter, config, clusterConfigState, metrics, plugins.delegatingLoader(), Time.SYSTEM,
                 RetryWithToleranceOperatorTest.NOOP_OPERATOR);
@@ -337,6 +351,51 @@ public List answer() throws Throwable {
         PowerMock.verifyAll();
     }
 
+    @Test
+    public void testPollReturnsNoRecords() throws Exception {
+        // Test that the task handles an empty list of records
+        createWorkerTask();
+
+        sourceTask.initialize(EasyMock.anyObject(SourceTaskContext.class));
+        EasyMock.expectLastCall();
+        sourceTask.start(TASK_PROPS);
+        EasyMock.expectLastCall();
+        statusListener.onStartup(taskId);
+        EasyMock.expectLastCall();
+
+        // We'll wait for some data, then trigger a flush
+        final CountDownLatch pollLatch = expectEmptyPolls(1, new AtomicInteger());
+        expectOffsetFlush(true);
+
+        sourceTask.stop();
+        EasyMock.expectLastCall();
+        expectOffsetFlush(true);
+
+        statusListener.onShutdown(taskId);
+        EasyMock.expectLastCall();
+
+        producer.close(EasyMock.anyObject(Duration.class));
+        EasyMock.expectLastCall();
+
+        transformationChain.close();
+        EasyMock.expectLastCall();
+
+        PowerMock.replayAll();
+
+        workerTask.initialize(TASK_CONFIG);
+        Future taskFuture = executor.submit(workerTask);
+
+        assertTrue(awaitLatch(pollLatch));
+        assertTrue(workerTask.commitOffsets());
+        workerTask.stop();
+        assertTrue(workerTask.awaitStop(1000));
+
+        taskFuture.get();
+        assertPollMetrics(0);
+
+        PowerMock.verifyAll();
+    }
+
     @Test
     public void testCommit() throws Exception {
         // Test that the task commits properly when prompted
@@ -542,6 +601,21 @@ public void testSendRecordsRetries() throws Exception {
         PowerMock.verifyAll();
     }
 
+    @Test(expected = ConnectException.class)
+    public void testSendRecordsProducerCallbackFail() throws Exception {
+        createWorkerTask();
+
+        SourceRecord record1 = new SourceRecord(PARTITION, OFFSET, "topic", 1, KEY_SCHEMA, KEY, RECORD_SCHEMA, RECORD);
+        SourceRecord record2 = new SourceRecord(PARTITION, OFFSET, "topic", 2, KEY_SCHEMA, KEY, RECORD_SCHEMA, RECORD);
+
+        expectSendRecordProducerCallbackFail();
+
+        PowerMock.replayAll();
+
+        Whitebox.setInternalState(workerTask, "toSend", Arrays.asList(record1, record2));
+        Whitebox.invokeMethod(workerTask, "sendRecords");
+    }
+
     @Test
     public void testSendRecordsTaskCommitRecordFail() throws Exception {
         createWorkerTask();
@@ -619,6 +693,20 @@ public Object answer() throws Throwable {
         PowerMock.verifyAll();
     }
 
+    @Test
+    public void testCancel() {
+        createWorkerTask();
+
+        offsetReader.close();
+        PowerMock.expectLastCall();
+
+        PowerMock.replayAll();
+
+        workerTask.cancel();
+
+        PowerMock.verifyAll();
+    }
+
     @Test
     public void testMetricsGroup() {
         SourceTaskMetricsGroup group = new SourceTaskMetricsGroup(taskId, metrics);
@@ -665,6 +753,98 @@ public void testMetricsGroup() {
         assertEquals(1800.0, metrics.currentMetricValueAsDouble(group1.metricGroup(), "source-record-active-count"), 0.001d);
     }
 
+    @Test
+    public void testHeaders() throws Exception {
+        Headers headers = new RecordHeaders();
+        headers.add("header_key", "header_value".getBytes());
+
+        org.apache.kafka.connect.header.Headers connectHeaders = new ConnectHeaders();
+        connectHeaders.add("header_key", new SchemaAndValue(Schema.STRING_SCHEMA, "header_value"));
+
+        createWorkerTask();
+
+        List records = new ArrayList<>();
+        records.add(new SourceRecord(PARTITION, OFFSET, "topic", null, KEY_SCHEMA, KEY, RECORD_SCHEMA, RECORD, null, connectHeaders));
+
+        Capture> sent = expectSendRecord(true, false, true, true, true, headers);
+
+        PowerMock.replayAll();
+
+        Whitebox.setInternalState(workerTask, "toSend", records);
+        Whitebox.invokeMethod(workerTask, "sendRecords");
+        assertEquals(SERIALIZED_KEY, sent.getValue().key());
+        assertEquals(SERIALIZED_RECORD, sent.getValue().value());
+        assertEquals(headers, sent.getValue().headers());
+
+        PowerMock.verifyAll();
+    }
+
+    @Test
+    public void testHeadersWithCustomConverter() throws Exception {
+        StringConverter stringConverter = new StringConverter();
+        TestConverterWithHeaders testConverter = new TestConverterWithHeaders();
+
+        createWorkerTask(TargetState.STARTED, stringConverter, testConverter, stringConverter);
+
+        List records = new ArrayList<>();
+
+        String stringA = "Árvíztűrő tükörfúrógép";
+        org.apache.kafka.connect.header.Headers headersA = new ConnectHeaders();
+        String encodingA = "latin2";
+        headersA.addString("encoding", encodingA);
+
+        records.add(new SourceRecord(PARTITION, OFFSET, "topic", null, Schema.STRING_SCHEMA, "a", Schema.STRING_SCHEMA, stringA, null, headersA));
+
+        String stringB = "Тестовое сообщение";
+        org.apache.kafka.connect.header.Headers headersB = new ConnectHeaders();
+        String encodingB = "koi8_r";
+        headersB.addString("encoding", encodingB);
+
+        records.add(new SourceRecord(PARTITION, OFFSET, "topic", null, Schema.STRING_SCHEMA, "b", Schema.STRING_SCHEMA, stringB, null, headersB));
+
+        Capture> sentRecordA = expectSendRecord(false, false, true, true, false, null);
+        Capture> sentRecordB = expectSendRecord(false, false, true, true, false, null);
+
+        PowerMock.replayAll();
+
+        Whitebox.setInternalState(workerTask, "toSend", records);
+        Whitebox.invokeMethod(workerTask, "sendRecords");
+
+        assertEquals(ByteBuffer.wrap("a".getBytes()), ByteBuffer.wrap(sentRecordA.getValue().key()));
+        assertEquals(
+            ByteBuffer.wrap(stringA.getBytes(encodingA)),
+            ByteBuffer.wrap(sentRecordA.getValue().value())
+        );
+        assertEquals(encodingA, new String(sentRecordA.getValue().headers().lastHeader("encoding").value()));
+
+        assertEquals(ByteBuffer.wrap("b".getBytes()), ByteBuffer.wrap(sentRecordB.getValue().key()));
+        assertEquals(
+            ByteBuffer.wrap(stringB.getBytes(encodingB)),
+            ByteBuffer.wrap(sentRecordB.getValue().value())
+        );
+        assertEquals(encodingB, new String(sentRecordB.getValue().headers().lastHeader("encoding").value()));
+
+        PowerMock.verifyAll();
+    }
+
+    private CountDownLatch expectEmptyPolls(int minimum, final AtomicInteger count) throws InterruptedException {
+        final CountDownLatch latch = new CountDownLatch(minimum);
+        // Note that we stub these to allow any number of calls because the thread will continue to
+        // run. The count passed in + latch returned just makes sure we get *at least* that number of
+        // calls
+        EasyMock.expect(sourceTask.poll())
+                .andStubAnswer(new IAnswer>() {
+                    @Override
+                    public List answer() throws Throwable {
+                        count.incrementAndGet();
+                        latch.countDown();
+                        Thread.sleep(10);
+                        return Collections.emptyList();
+                    }
+                });
+        return latch;
+    }
+
     private CountDownLatch expectPolls(int minimum, final AtomicInteger count) throws InterruptedException {
         final CountDownLatch latch = new CountDownLatch(minimum);
         // Note that we stub these to allow any number of calls because the thread will continue to
@@ -691,7 +871,7 @@ private CountDownLatch expectPolls(int count) throws InterruptedException {
 
     @SuppressWarnings("unchecked")
     private void expectSendRecordSyncFailure(Throwable error) throws InterruptedException {
-        expectConvertKeyValue(false);
+        expectConvertHeadersAndKeyValue(false);
         expectApplyTransformationChain(false);
 
         offsetWriter.offset(PARTITION, OFFSET);
@@ -711,17 +891,35 @@ private Capture> expectSendRecordOnce(boolean isR
         return expectSendRecordTaskCommitRecordSucceed(false, isRetry);
     }
 
+    private Capture> expectSendRecordProducerCallbackFail() throws InterruptedException {
+        return expectSendRecord(false, false, false, false, true, emptyHeaders());
+    }
+
     private Capture> expectSendRecordTaskCommitRecordSucceed(boolean anyTimes, boolean isRetry) throws InterruptedException {
-        return expectSendRecord(anyTimes, isRetry, true);
+        return expectSendRecord(anyTimes, isRetry, true, true, true, emptyHeaders());
     }
 
     private Capture> expectSendRecordTaskCommitRecordFail(boolean anyTimes, boolean isRetry) throws InterruptedException {
-        return expectSendRecord(anyTimes, isRetry, false);
+        return expectSendRecord(anyTimes, isRetry, true, false, true, emptyHeaders());
     }
 
-    @SuppressWarnings("unchecked")
     private Capture> expectSendRecord(boolean anyTimes, boolean isRetry, boolean succeed) throws InterruptedException {
-        expectConvertKeyValue(anyTimes);
+        return expectSendRecord(anyTimes, isRetry, succeed, true, true, emptyHeaders());
+    }
+
+    @SuppressWarnings("unchecked")
+    private Capture> expectSendRecord(
+        boolean anyTimes,
+        boolean isRetry,
+        boolean sendSuccess,
+        boolean commitSuccess,
+        boolean isMockedConverters,
+        Headers headers
+    ) throws InterruptedException {
+        if (isMockedConverters) {
+            expectConvertHeadersAndKeyValue(anyTimes, headers);
+        }
+
         expectApplyTransformationChain(anyTimes);
 
         Capture> sent = EasyMock.newCapture();
@@ -737,15 +935,19 @@ private Capture> expectSendRecord(boolean anyTime
 
         // 2. Converted data passed to the producer, which will need callbacks invoked for flush to work
         IExpectationSetters> expect = EasyMock.expect(
-                producer.send(EasyMock.capture(sent),
-                        EasyMock.capture(producerCallbacks)));
+            producer.send(EasyMock.capture(sent),
+                EasyMock.capture(producerCallbacks)));
         IAnswer> expectResponse = new IAnswer>() {
             @Override
             public Future answer() throws Throwable {
                 synchronized (producerCallbacks) {
                     for (org.apache.kafka.clients.producer.Callback cb : producerCallbacks.getValues()) {
-                        cb.onCompletion(new RecordMetadata(new TopicPartition("foo", 0), 0, 0,
-                                                           0L, 0L, 0, 0), null);
+                        if (sendSuccess) {
+                            cb.onCompletion(new RecordMetadata(new TopicPartition("foo", 0), 0, 0,
+                                0L, 0L, 0, 0), null);
+                        } else {
+                            cb.onCompletion(null, new TopicAuthorizationException("foo"));
+                        }
                     }
                     producerCallbacks.reset();
                 }
@@ -757,19 +959,32 @@ public Future answer() throws Throwable {
         else
             expect.andAnswer(expectResponse);
 
-        // 3. As a result of a successful producer send callback, we'll notify the source task of the record commit
-        expectTaskCommitRecord(anyTimes, succeed);
+        if (sendSuccess) {
+            // 3. As a result of a successful producer send callback, we'll notify the source task of the record commit
+            expectTaskCommitRecordWithOffset(anyTimes, commitSuccess);
+        }
 
         return sent;
     }
 
-    private void expectConvertKeyValue(boolean anyTimes) {
-        IExpectationSetters convertKeyExpect = EasyMock.expect(keyConverter.fromConnectData(TOPIC, KEY_SCHEMA, KEY));
+    private void expectConvertHeadersAndKeyValue(boolean anyTimes) {
+        expectConvertHeadersAndKeyValue(anyTimes, emptyHeaders());
+    }
+
+    private void expectConvertHeadersAndKeyValue(boolean anyTimes, Headers headers) {
+        for (Header header : headers) {
+            IExpectationSetters convertHeaderExpect = EasyMock.expect(headerConverter.fromConnectHeader(TOPIC, header.key(), Schema.STRING_SCHEMA, new String(header.value())));
+            if (anyTimes)
+                convertHeaderExpect.andStubReturn(header.value());
+            else
+                convertHeaderExpect.andReturn(header.value());
+        }
+        IExpectationSetters convertKeyExpect = EasyMock.expect(keyConverter.fromConnectData(TOPIC, headers, KEY_SCHEMA, KEY));
         if (anyTimes)
             convertKeyExpect.andStubReturn(SERIALIZED_KEY);
         else
             convertKeyExpect.andReturn(SERIALIZED_KEY);
-        IExpectationSetters convertValueExpect = EasyMock.expect(valueConverter.fromConnectData(TOPIC, RECORD_SCHEMA, RECORD));
+        IExpectationSetters convertValueExpect = EasyMock.expect(valueConverter.fromConnectData(TOPIC, headers, RECORD_SCHEMA, RECORD));
         if (anyTimes)
             convertValueExpect.andStubReturn(SERIALIZED_RECORD);
         else
@@ -795,8 +1010,8 @@ public SourceRecord answer() {
             });
     }
 
-    private void expectTaskCommitRecord(boolean anyTimes, boolean succeed) throws InterruptedException {
-        sourceTask.commitRecord(EasyMock.anyObject(SourceRecord.class));
+    private void expectTaskCommitRecordWithOffset(boolean anyTimes, boolean succeed) throws InterruptedException {
+        sourceTask.commitRecord(EasyMock.anyObject(SourceRecord.class), EasyMock.anyObject(RecordMetadata.class));
         IExpectationSetters expect = EasyMock.expectLastCall();
         if (!succeed) {
             expect = expect.andThrow(new RuntimeException("Error committing record in source task"));
@@ -871,6 +1086,10 @@ private void assertPollMetrics(int minimumPollCountExpected) {
         }
     }
 
+    private RecordHeaders emptyHeaders() {
+        return new RecordHeaders();
+    }
+
     private abstract static class TestSourceTask extends SourceTask {
     }
 
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 8f15c87c5016c..e7ffd607f9431 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
@@ -24,11 +24,15 @@
 import org.apache.kafka.common.config.AbstractConfig;
 import org.apache.kafka.common.config.ConfigDef;
 import org.apache.kafka.common.config.ConfigException;
+import org.apache.kafka.common.config.provider.MockFileConfigProvider;
 import org.apache.kafka.common.utils.MockTime;
 import org.apache.kafka.common.utils.Time;
 import org.apache.kafka.connect.connector.Connector;
 import org.apache.kafka.connect.connector.ConnectorContext;
 import org.apache.kafka.connect.connector.Task;
+import org.apache.kafka.connect.connector.policy.AllConnectorClientConfigOverridePolicy;
+import org.apache.kafka.connect.connector.policy.ConnectorClientConfigOverridePolicy;
+import org.apache.kafka.connect.connector.policy.NoneConnectorClientConfigOverridePolicy;
 import org.apache.kafka.connect.data.Schema;
 import org.apache.kafka.connect.data.SchemaAndValue;
 import org.apache.kafka.connect.errors.ConnectException;
@@ -42,6 +46,7 @@
 import org.apache.kafka.connect.runtime.isolation.PluginClassLoader;
 import org.apache.kafka.connect.runtime.isolation.Plugins;
 import org.apache.kafka.connect.runtime.isolation.Plugins.ClassLoaderUsage;
+import org.apache.kafka.connect.runtime.rest.entities.ConnectorStateInfo;
 import org.apache.kafka.connect.runtime.standalone.StandaloneConfig;
 import org.apache.kafka.connect.sink.SinkTask;
 import org.apache.kafka.connect.source.SourceRecord;
@@ -60,6 +65,7 @@
 import org.junit.runner.RunWith;
 import org.powermock.api.easymock.PowerMock;
 import org.powermock.api.easymock.annotation.Mock;
+import org.powermock.api.easymock.annotation.MockNice;
 import org.powermock.api.easymock.annotation.MockStrict;
 import org.powermock.core.classloader.annotations.PowerMockIgnore;
 import org.powermock.core.classloader.annotations.PrepareForTest;
@@ -71,6 +77,10 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.UUID;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.ExecutorService;
 
 import static org.apache.kafka.connect.runtime.errors.RetryWithToleranceOperatorTest.NOOP_OPERATOR;
 import static org.easymock.EasyMock.anyObject;
@@ -78,6 +88,7 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
 import static org.junit.Assert.fail;
 
 @RunWith(PowerMockRunner.class)
@@ -88,6 +99,8 @@ public class WorkerTest extends ThreadedTest {
     private static final String CONNECTOR_ID = "test-connector";
     private static final ConnectorTaskId TASK_ID = new ConnectorTaskId("job", 0);
     private static final String WORKER_ID = "localhost:8083";
+    private final ConnectorClientConfigOverridePolicy noneConnectorClientConfigOverridePolicy = new NoneConnectorClientConfigOverridePolicy();
+    private final ConnectorClientConfigOverridePolicy allConnectorClientConfigOverridePolicy = new AllConnectorClientConfigOverridePolicy();
 
     private Map workerProps = new HashMap<>();
     private WorkerConfig config;
@@ -109,6 +122,7 @@ public class WorkerTest extends ThreadedTest {
     @MockStrict
     private ConnectorStatus.Listener connectorStatusListener;
 
+    @Mock private Herder herder;
     @Mock private Connector connector;
     @Mock private ConnectorContext ctx;
     @Mock private TestSourceTask task;
@@ -118,6 +132,9 @@ public class WorkerTest extends ThreadedTest {
     @Mock private Converter taskKeyConverter;
     @Mock private Converter taskValueConverter;
     @Mock private HeaderConverter taskHeaderConverter;
+    @Mock private ExecutorService executorService;
+    @MockNice private ConnectorConfig connectorConfig;
+    private String mockFileProviderTestId;
 
     @Before
     public void setup() {
@@ -130,6 +147,10 @@ public void setup() {
         workerProps.put("internal.value.converter.schemas.enable", "false");
         workerProps.put("offset.storage.file.filename", "/tmp/connect.offsets");
         workerProps.put(CommonClientConfigs.METRIC_REPORTER_CLASSES_CONFIG, MockMetricsReporter.class.getName());
+        workerProps.put("config.providers", "file");
+        workerProps.put("config.providers.file.class", MockFileConfigProvider.class.getName());
+        mockFileProviderTestId = UUID.randomUUID().toString();
+        workerProps.put("config.providers.file.param.testId", mockFileProviderTestId);
         config = new StandaloneConfig(workerProps);
 
         defaultProducerConfigs.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092");
@@ -173,9 +194,11 @@ public void testStartAndStopConnector() {
 
         EasyMock.expect(connector.version()).andReturn("1.0");
 
+        expectFileConfigProvider();
         EasyMock.expect(plugins.compareAndSwapLoaders(connector))
                 .andReturn(delegatingLoader)
                 .times(2);
+
         connector.initialize(anyObject(ConnectorContext.class));
         EasyMock.expectLastCall();
         connector.start(props);
@@ -198,7 +221,7 @@ public void testStartAndStopConnector() {
 
         PowerMock.replayAll();
 
-        worker = new Worker(WORKER_ID, new MockTime(), plugins, config, offsetBackingStore);
+        worker = new Worker(WORKER_ID, new MockTime(), plugins, config, offsetBackingStore, noneConnectorClientConfigOverridePolicy);
         worker.start();
 
         assertEquals(Collections.emptySet(), worker.connectorNames());
@@ -221,12 +244,24 @@ public void testStartAndStopConnector() {
         assertStatistics(worker, 0, 0);
 
         PowerMock.verifyAll();
+        MockFileConfigProvider.assertClosed(mockFileProviderTestId);
+    }
+
+    private void expectFileConfigProvider() {
+        EasyMock.expect(plugins.newConfigProvider(EasyMock.anyObject(),
+                    EasyMock.eq("config.providers.file"), EasyMock.anyObject()))
+                .andAnswer(() -> {
+                    MockFileConfigProvider mockFileConfigProvider = new MockFileConfigProvider();
+                    mockFileConfigProvider.configure(Collections.singletonMap("testId", mockFileProviderTestId));
+                    return mockFileConfigProvider;
+                });
     }
 
     @Test
     public void testStartConnectorFailure() {
         expectConverters();
         expectStartStorage();
+        expectFileConfigProvider();
 
         Map props = new HashMap<>();
         props.put(SinkConnectorConfig.TOPICS_CONFIG, "foo,bar");
@@ -249,7 +284,7 @@ public void testStartConnectorFailure() {
 
         PowerMock.replayAll();
 
-        worker = new Worker(WORKER_ID, new MockTime(), plugins, config, offsetBackingStore);
+        worker = new Worker(WORKER_ID, new MockTime(), plugins, config, offsetBackingStore, noneConnectorClientConfigOverridePolicy);
         worker.start();
 
         assertStatistics(worker, 0, 0);
@@ -271,6 +306,7 @@ public void testStartConnectorFailure() {
     public void testAddConnectorByAlias() {
         expectConverters();
         expectStartStorage();
+        expectFileConfigProvider();
 
         EasyMock.expect(plugins.currentThreadLoader()).andReturn(delegatingLoader).times(2);
         EasyMock.expect(plugins.newConnector("WorkerTestConnector")).andReturn(connector);
@@ -309,7 +345,7 @@ public void testAddConnectorByAlias() {
 
         PowerMock.replayAll();
 
-        worker = new Worker(WORKER_ID, new MockTime(), plugins, config, offsetBackingStore);
+        worker = new Worker(WORKER_ID, new MockTime(), plugins, config, offsetBackingStore, noneConnectorClientConfigOverridePolicy);
         worker.start();
 
         assertStatistics(worker, 0, 0);
@@ -335,6 +371,7 @@ public void testAddConnectorByAlias() {
     public void testAddConnectorByShortAlias() {
         expectConverters();
         expectStartStorage();
+        expectFileConfigProvider();
 
         EasyMock.expect(plugins.currentThreadLoader()).andReturn(delegatingLoader).times(2);
         EasyMock.expect(plugins.newConnector("WorkerTest")).andReturn(connector);
@@ -373,7 +410,7 @@ public void testAddConnectorByShortAlias() {
 
         PowerMock.replayAll();
 
-        worker = new Worker(WORKER_ID, new MockTime(), plugins, config, offsetBackingStore);
+        worker = new Worker(WORKER_ID, new MockTime(), plugins, config, offsetBackingStore, noneConnectorClientConfigOverridePolicy);
         worker.start();
 
         assertStatistics(worker, 0, 0);
@@ -396,10 +433,11 @@ public void testAddConnectorByShortAlias() {
     public void testStopInvalidConnector() {
         expectConverters();
         expectStartStorage();
+        expectFileConfigProvider();
 
         PowerMock.replayAll();
 
-        worker = new Worker(WORKER_ID, new MockTime(), plugins, config, offsetBackingStore);
+        worker = new Worker(WORKER_ID, new MockTime(), plugins, config, offsetBackingStore, noneConnectorClientConfigOverridePolicy);
         worker.start();
 
         worker.stopConnector(CONNECTOR_ID);
@@ -411,6 +449,7 @@ public void testStopInvalidConnector() {
     public void testReconfigureConnectorTasks() {
         expectConverters();
         expectStartStorage();
+        expectFileConfigProvider();
 
         EasyMock.expect(plugins.currentThreadLoader()).andReturn(delegatingLoader).times(3);
         EasyMock.expect(plugins.newConnector(WorkerTestConnector.class.getName()))
@@ -456,7 +495,7 @@ public void testReconfigureConnectorTasks() {
 
         PowerMock.replayAll();
 
-        worker = new Worker(WORKER_ID, new MockTime(), plugins, config, offsetBackingStore);
+        worker = new Worker(WORKER_ID, new MockTime(), plugins, config, offsetBackingStore, noneConnectorClientConfigOverridePolicy);
         worker.start();
 
         assertStatistics(worker, 0, 0);
@@ -499,6 +538,7 @@ public void testReconfigureConnectorTasks() {
     public void testAddRemoveTask() throws Exception {
         expectConverters();
         expectStartStorage();
+        expectFileConfigProvider();
 
         EasyMock.expect(workerTask.id()).andStubReturn(TASK_ID);
 
@@ -543,13 +583,11 @@ public void testAddRemoveTask() throws Exception {
         expectTaskValueConverters(ClassLoaderUsage.CURRENT_CLASSLOADER, taskValueConverter);
         expectTaskHeaderConverter(ClassLoaderUsage.CURRENT_CLASSLOADER, taskHeaderConverter);
 
-        workerTask.run();
-        EasyMock.expectLastCall();
+        EasyMock.expect(executorService.submit(workerTask)).andReturn(null);
 
         EasyMock.expect(plugins.delegatingLoader()).andReturn(delegatingLoader);
         EasyMock.expect(delegatingLoader.connectorLoader(WorkerTestConnector.class.getName()))
                 .andReturn(pluginLoader);
-
         EasyMock.expect(Plugins.compareAndSwapLoaders(pluginLoader)).andReturn(delegatingLoader)
                 .times(2);
 
@@ -557,7 +595,8 @@ public void testAddRemoveTask() throws Exception {
 
         EasyMock.expect(Plugins.compareAndSwapLoaders(delegatingLoader)).andReturn(pluginLoader)
                 .times(2);
-
+        plugins.connectorClass(WorkerTestConnector.class.getName());
+        EasyMock.expectLastCall().andReturn(WorkerTestConnector.class);
         // Remove
         workerTask.stop();
         EasyMock.expectLastCall();
@@ -568,7 +607,8 @@ public void testAddRemoveTask() throws Exception {
 
         PowerMock.replayAll();
 
-        worker = new Worker(WORKER_ID, new MockTime(), plugins, config, offsetBackingStore);
+        worker = new Worker(WORKER_ID, new MockTime(), plugins, config, offsetBackingStore, executorService,
+                            noneConnectorClientConfigOverridePolicy);
         worker.start();
         assertStatistics(worker, 0, 0);
         assertStartupStatistics(worker, 0, 0, 0, 0);
@@ -589,10 +629,176 @@ public void testAddRemoveTask() throws Exception {
         PowerMock.verifyAll();
     }
 
+    @Test
+    public void testTaskStatusMetricsStatuses() throws Exception {
+        expectConverters();
+        expectStartStorage();
+        expectFileConfigProvider();
+
+        EasyMock.expect(workerTask.id()).andStubReturn(TASK_ID);
+
+        EasyMock.expect(plugins.currentThreadLoader()).andReturn(delegatingLoader).times(2);
+        PowerMock.expectNew(WorkerSourceTask.class, EasyMock.eq(TASK_ID),
+            EasyMock.eq(task),
+            anyObject(TaskStatus.Listener.class),
+            EasyMock.eq(TargetState.STARTED),
+            anyObject(JsonConverter.class),
+            anyObject(JsonConverter.class),
+            anyObject(JsonConverter.class),
+            EasyMock.eq(new TransformationChain<>(Collections.emptyList(), NOOP_OPERATOR)),
+            anyObject(KafkaProducer.class),
+            anyObject(OffsetStorageReader.class),
+            anyObject(OffsetStorageWriter.class),
+            EasyMock.eq(config),
+            anyObject(ClusterConfigState.class),
+            anyObject(ConnectMetrics.class),
+            anyObject(ClassLoader.class),
+            anyObject(Time.class),
+            anyObject(RetryWithToleranceOperator.class)).andReturn(workerTask);
+        Map origProps = new HashMap<>();
+        origProps.put(TaskConfig.TASK_CLASS_CONFIG, TestSourceTask.class.getName());
+
+        TaskConfig taskConfig = new TaskConfig(origProps);
+        // We should expect this call, but the pluginLoader being swapped in is only mocked.
+        // EasyMock.expect(pluginLoader.loadClass(TestSourceTask.class.getName()))
+        //        .andReturn((Class) TestSourceTask.class);
+        EasyMock.expect(plugins.newTask(TestSourceTask.class)).andReturn(task);
+        EasyMock.expect(task.version()).andReturn("1.0");
+
+        workerTask.initialize(taskConfig);
+        EasyMock.expectLastCall();
+
+        // Expect that the worker will create converters and will find them using the current classloader ...
+        assertNotNull(taskKeyConverter);
+        assertNotNull(taskValueConverter);
+        assertNotNull(taskHeaderConverter);
+        expectTaskKeyConverters(ClassLoaderUsage.CURRENT_CLASSLOADER, taskKeyConverter);
+        expectTaskValueConverters(ClassLoaderUsage.CURRENT_CLASSLOADER, taskValueConverter);
+        expectTaskHeaderConverter(ClassLoaderUsage.CURRENT_CLASSLOADER, taskHeaderConverter);
+
+        EasyMock.expect(executorService.submit(workerTask)).andReturn(null);
+
+        EasyMock.expect(plugins.delegatingLoader()).andReturn(delegatingLoader);
+        EasyMock.expect(delegatingLoader.connectorLoader(WorkerTestConnector.class.getName()))
+            .andReturn(pluginLoader);
+        EasyMock.expect(Plugins.compareAndSwapLoaders(pluginLoader)).andReturn(delegatingLoader)
+            .times(2);
+
+        EasyMock.expect(Plugins.compareAndSwapLoaders(delegatingLoader)).andReturn(pluginLoader)
+            .times(2);
+        plugins.connectorClass(WorkerTestConnector.class.getName());
+        EasyMock.expectLastCall().andReturn(WorkerTestConnector.class);
+
+        EasyMock.expect(workerTask.awaitStop(EasyMock.anyLong())).andStubReturn(true);
+        EasyMock.expectLastCall();
+
+        // Each time we check the task metrics, the worker will call the herder
+        herder.taskStatus(TASK_ID);
+        EasyMock.expectLastCall()
+            .andReturn(new ConnectorStateInfo.TaskState(0, "RUNNING", "worker", "msg"));
+
+        herder.taskStatus(TASK_ID);
+        EasyMock.expectLastCall()
+            .andReturn(new ConnectorStateInfo.TaskState(0, "PAUSED", "worker", "msg"));
+
+        herder.taskStatus(TASK_ID);
+        EasyMock.expectLastCall()
+            .andReturn(new ConnectorStateInfo.TaskState(0, "FAILED", "worker", "msg"));
+
+        herder.taskStatus(TASK_ID);
+        EasyMock.expectLastCall()
+            .andReturn(new ConnectorStateInfo.TaskState(0, "DESTROYED", "worker", "msg"));
+
+        herder.taskStatus(TASK_ID);
+        EasyMock.expectLastCall()
+            .andReturn(new ConnectorStateInfo.TaskState(0, "UNASSIGNED", "worker", "msg"));
+
+        // Called when we stop the worker
+        EasyMock.expect(workerTask.loader()).andReturn(pluginLoader);
+        workerTask.stop();
+        EasyMock.expectLastCall();
+
+        PowerMock.replayAll();
+
+        worker = new Worker(WORKER_ID,
+            new MockTime(),
+            plugins,
+            config,
+            offsetBackingStore,
+            executorService,
+            noneConnectorClientConfigOverridePolicy);
+
+        worker.herder = herder;
+
+        worker.start();
+        assertStatistics(worker, 0, 0);
+        assertStartupStatistics(worker, 0, 0, 0, 0);
+        assertEquals(Collections.emptySet(), worker.taskIds());
+        worker.startTask(
+            TASK_ID,
+            ClusterConfigState.EMPTY,
+            anyConnectorConfigMap(),
+            origProps,
+            taskStatusListener,
+            TargetState.STARTED);
+
+        assertStatusMetrics(1L, "connector-running-task-count");
+        assertStatusMetrics(1L, "connector-paused-task-count");
+        assertStatusMetrics(1L, "connector-failed-task-count");
+        assertStatusMetrics(1L, "connector-destroyed-task-count");
+        assertStatusMetrics(1L, "connector-unassigned-task-count");
+
+        worker.stopAndAwaitTask(TASK_ID);
+        assertStatusMetrics(0L, "connector-running-task-count");
+        assertStatusMetrics(0L, "connector-paused-task-count");
+        assertStatusMetrics(0L, "connector-failed-task-count");
+        assertStatusMetrics(0L, "connector-destroyed-task-count");
+        assertStatusMetrics(0L, "connector-unassigned-task-count");
+
+        PowerMock.verifyAll();
+    }
+
+    @Test
+    public void testConnectorStatusMetricsGroup_taskStatusCounter() {
+        ConcurrentMap tasks = new ConcurrentHashMap<>();
+        tasks.put(new ConnectorTaskId("c1", 0), workerTask);
+        tasks.put(new ConnectorTaskId("c1", 1), workerTask);
+        tasks.put(new ConnectorTaskId("c2", 0), workerTask);
+
+        expectConverters();
+        expectStartStorage();
+        expectFileConfigProvider();
+
+        EasyMock.expect(Plugins.compareAndSwapLoaders(pluginLoader)).andReturn(delegatingLoader);
+        EasyMock.expect(Plugins.compareAndSwapLoaders(pluginLoader)).andReturn(delegatingLoader);
+
+        EasyMock.expect(Plugins.compareAndSwapLoaders(delegatingLoader)).andReturn(pluginLoader);
+
+        taskStatusListener.onFailure(EasyMock.eq(TASK_ID), EasyMock.anyObject());
+        EasyMock.expectLastCall();
+
+        PowerMock.replayAll();
+
+        worker = new Worker(WORKER_ID,
+            new MockTime(),
+            plugins,
+            config,
+            offsetBackingStore,
+            noneConnectorClientConfigOverridePolicy);
+
+        Worker.ConnectorStatusMetricsGroup metricGroup = new Worker.ConnectorStatusMetricsGroup(
+            worker.metrics(), tasks, herder
+        );
+        assertEquals(2L, (long) metricGroup.taskCounter("c1").metricValue(0L));
+        assertEquals(1L, (long) metricGroup.taskCounter("c2").metricValue(0L));
+        assertEquals(0L, (long) metricGroup.taskCounter("fakeConnector").metricValue(0L));
+    }
+
     @Test
     public void testStartTaskFailure() {
         expectConverters();
         expectStartStorage();
+        expectFileConfigProvider();
 
         Map origProps = new HashMap<>();
         origProps.put(TaskConfig.TASK_CLASS_CONFIG, "missing.From.This.Workers.Classpath");
@@ -619,7 +825,7 @@ public void testStartTaskFailure() {
 
         PowerMock.replayAll();
 
-        worker = new Worker(WORKER_ID, new MockTime(), plugins, config, offsetBackingStore);
+        worker = new Worker(WORKER_ID, new MockTime(), plugins, config, offsetBackingStore, noneConnectorClientConfigOverridePolicy);
         worker.start();
         assertStatistics(worker, 0, 0);
         assertStartupStatistics(worker, 0, 0, 0, 0);
@@ -638,6 +844,7 @@ public void testStartTaskFailure() {
     public void testCleanupTasksOnStop() throws Exception {
         expectConverters();
         expectStartStorage();
+        expectFileConfigProvider();
 
         EasyMock.expect(workerTask.id()).andStubReturn(TASK_ID);
 
@@ -685,8 +892,7 @@ public void testCleanupTasksOnStop() throws Exception {
         expectTaskHeaderConverter(ClassLoaderUsage.CURRENT_CLASSLOADER, null);
         expectTaskHeaderConverter(ClassLoaderUsage.PLUGINS, taskHeaderConverter);
 
-        workerTask.run();
-        EasyMock.expectLastCall();
+        EasyMock.expect(executorService.submit(workerTask)).andReturn(null);
 
         EasyMock.expect(plugins.delegatingLoader()).andReturn(delegatingLoader);
         EasyMock.expect(delegatingLoader.connectorLoader(WorkerTestConnector.class.getName()))
@@ -699,7 +905,8 @@ public void testCleanupTasksOnStop() throws Exception {
 
         EasyMock.expect(Plugins.compareAndSwapLoaders(delegatingLoader)).andReturn(pluginLoader)
                 .times(2);
-
+        plugins.connectorClass(WorkerTestConnector.class.getName());
+        EasyMock.expectLastCall().andReturn(WorkerTestConnector.class);
         // Remove on Worker.stop()
         workerTask.stop();
         EasyMock.expectLastCall();
@@ -712,7 +919,8 @@ public void testCleanupTasksOnStop() throws Exception {
 
         PowerMock.replayAll();
 
-        worker = new Worker(WORKER_ID, new MockTime(), plugins, config, offsetBackingStore);
+        worker = new Worker(WORKER_ID, new MockTime(), plugins, config, offsetBackingStore, executorService,
+                            noneConnectorClientConfigOverridePolicy);
         worker.start();
         assertStatistics(worker, 0, 0);
         worker.startTask(TASK_ID, ClusterConfigState.EMPTY, anyConnectorConfigMap(), origProps, taskStatusListener, TargetState.STARTED);
@@ -727,6 +935,7 @@ public void testCleanupTasksOnStop() throws Exception {
     public void testConverterOverrides() throws Exception {
         expectConverters();
         expectStartStorage();
+        expectFileConfigProvider();
 
         EasyMock.expect(workerTask.id()).andStubReturn(TASK_ID);
 
@@ -778,8 +987,7 @@ public void testConverterOverrides() throws Exception {
         expectTaskHeaderConverter(ClassLoaderUsage.CURRENT_CLASSLOADER, null);
         expectTaskHeaderConverter(ClassLoaderUsage.PLUGINS, taskHeaderConverter);
 
-        workerTask.run();
-        EasyMock.expectLastCall();
+        EasyMock.expect(executorService.submit(workerTask)).andReturn(null);
 
         EasyMock.expect(plugins.delegatingLoader()).andReturn(delegatingLoader);
         EasyMock.expect(delegatingLoader.connectorLoader(WorkerTestConnector.class.getName()))
@@ -792,6 +1000,8 @@ public void testConverterOverrides() throws Exception {
 
         EasyMock.expect(Plugins.compareAndSwapLoaders(delegatingLoader)).andReturn(pluginLoader)
                 .times(2);
+        plugins.connectorClass(WorkerTestConnector.class.getName());
+        EasyMock.expectLastCall().andReturn(WorkerTestConnector.class);
 
         // Remove
         workerTask.stop();
@@ -803,7 +1013,8 @@ public void testConverterOverrides() throws Exception {
 
         PowerMock.replayAll();
 
-        worker = new Worker(WORKER_ID, new MockTime(), plugins, config, offsetBackingStore);
+        worker = new Worker(WORKER_ID, new MockTime(), plugins, config, offsetBackingStore, executorService,
+                            noneConnectorClientConfigOverridePolicy);
         worker.start();
         assertStatistics(worker, 0, 0);
         assertEquals(Collections.emptySet(), worker.taskIds());
@@ -829,7 +1040,13 @@ public void testConverterOverrides() throws Exception {
 
     @Test
     public void testProducerConfigsWithoutOverrides() {
-        assertEquals(defaultProducerConfigs, Worker.producerConfigs(config));
+        EasyMock.expect(connectorConfig.originalsWithPrefix(ConnectorConfig.CONNECTOR_CLIENT_PRODUCER_OVERRIDES_PREFIX)).andReturn(
+            new HashMap());
+        PowerMock.replayAll();
+        Map expectedConfigs = new HashMap<>(defaultProducerConfigs);
+        expectedConfigs.put("client.id", "connector-producer-job-0");
+        assertEquals(expectedConfigs,
+                     Worker.producerConfigs(TASK_ID, "connector-producer-" + TASK_ID, config, connectorConfig, null, noneConnectorClientConfigOverridePolicy));
     }
 
     @Test
@@ -837,19 +1054,52 @@ public void testProducerConfigsWithOverrides() {
         Map props = new HashMap<>(workerProps);
         props.put("producer.acks", "-1");
         props.put("producer.linger.ms", "1000");
+        props.put("producer.client.id", "producer-test-id");
         WorkerConfig configWithOverrides = new StandaloneConfig(props);
 
         Map expectedConfigs = new HashMap<>(defaultProducerConfigs);
         expectedConfigs.put("acks", "-1");
         expectedConfigs.put("linger.ms", "1000");
-        assertEquals(expectedConfigs, Worker.producerConfigs(configWithOverrides));
+        expectedConfigs.put("client.id", "producer-test-id");
+        EasyMock.expect(connectorConfig.originalsWithPrefix(ConnectorConfig.CONNECTOR_CLIENT_PRODUCER_OVERRIDES_PREFIX)).andReturn(
+            new HashMap());
+        PowerMock.replayAll();
+        assertEquals(expectedConfigs,
+                     Worker.producerConfigs(TASK_ID, "connector-producer-" + TASK_ID, configWithOverrides, connectorConfig, null, allConnectorClientConfigOverridePolicy));
+    }
+
+    @Test
+    public void testProducerConfigsWithClientOverrides() {
+        Map props = new HashMap<>(workerProps);
+        props.put("producer.acks", "-1");
+        props.put("producer.linger.ms", "1000");
+        props.put("producer.client.id", "producer-test-id");
+        WorkerConfig configWithOverrides = new StandaloneConfig(props);
+
+        Map expectedConfigs = new HashMap<>(defaultProducerConfigs);
+        expectedConfigs.put("acks", "-1");
+        expectedConfigs.put("linger.ms", "5000");
+        expectedConfigs.put("batch.size", "1000");
+        expectedConfigs.put("client.id", "producer-test-id");
+        Map connConfig = new HashMap();
+        connConfig.put("linger.ms", "5000");
+        connConfig.put("batch.size", "1000");
+        EasyMock.expect(connectorConfig.originalsWithPrefix(ConnectorConfig.CONNECTOR_CLIENT_PRODUCER_OVERRIDES_PREFIX))
+            .andReturn(connConfig);
+        PowerMock.replayAll();
+        assertEquals(expectedConfigs,
+                     Worker.producerConfigs(TASK_ID, "connector-producer-" + TASK_ID, configWithOverrides, connectorConfig, null, allConnectorClientConfigOverridePolicy));
     }
 
     @Test
     public void testConsumerConfigsWithoutOverrides() {
         Map expectedConfigs = new HashMap<>(defaultConsumerConfigs);
         expectedConfigs.put("group.id", "connect-test");
-        assertEquals(expectedConfigs, Worker.consumerConfigs(new ConnectorTaskId("test", 1), config));
+        expectedConfigs.put("client.id", "connector-consumer-test-1");
+        EasyMock.expect(connectorConfig.originalsWithPrefix(ConnectorConfig.CONNECTOR_CLIENT_CONSUMER_OVERRIDES_PREFIX)).andReturn(new HashMap<>());
+        PowerMock.replayAll();
+        assertEquals(expectedConfigs, Worker.consumerConfigs(new ConnectorTaskId("test", 1), config, connectorConfig,
+                                                             null, noneConnectorClientConfigOverridePolicy));
     }
 
     @Test
@@ -857,16 +1107,114 @@ public void testConsumerConfigsWithOverrides() {
         Map props = new HashMap<>(workerProps);
         props.put("consumer.auto.offset.reset", "latest");
         props.put("consumer.max.poll.records", "1000");
+        props.put("consumer.client.id", "consumer-test-id");
         WorkerConfig configWithOverrides = new StandaloneConfig(props);
 
         Map expectedConfigs = new HashMap<>(defaultConsumerConfigs);
         expectedConfigs.put("group.id", "connect-test");
         expectedConfigs.put("auto.offset.reset", "latest");
         expectedConfigs.put("max.poll.records", "1000");
-        assertEquals(expectedConfigs, Worker.consumerConfigs(new ConnectorTaskId("test", 1), configWithOverrides));
+        expectedConfigs.put("client.id", "consumer-test-id");
+        EasyMock.expect(connectorConfig.originalsWithPrefix(ConnectorConfig.CONNECTOR_CLIENT_CONSUMER_OVERRIDES_PREFIX)).andReturn(new HashMap<>());
+        PowerMock.replayAll();
+        assertEquals(expectedConfigs, Worker.consumerConfigs(new ConnectorTaskId("test", 1), configWithOverrides, connectorConfig,
+                                                             null, noneConnectorClientConfigOverridePolicy));
+
+    }
+
+    @Test
+    public void testConsumerConfigsWithClientOverrides() {
+        Map props = new HashMap<>(workerProps);
+        props.put("consumer.auto.offset.reset", "latest");
+        props.put("consumer.max.poll.records", "5000");
+        WorkerConfig configWithOverrides = new StandaloneConfig(props);
+
+        Map expectedConfigs = new HashMap<>(defaultConsumerConfigs);
+        expectedConfigs.put("group.id", "connect-test");
+        expectedConfigs.put("auto.offset.reset", "latest");
+        expectedConfigs.put("max.poll.records", "5000");
+        expectedConfigs.put("max.poll.interval.ms", "1000");
+        expectedConfigs.put("client.id", "connector-consumer-test-1");
+        Map connConfig = new HashMap();
+        connConfig.put("max.poll.records", "5000");
+        connConfig.put("max.poll.interval.ms", "1000");
+        EasyMock.expect(connectorConfig.originalsWithPrefix(ConnectorConfig.CONNECTOR_CLIENT_CONSUMER_OVERRIDES_PREFIX))
+            .andReturn(connConfig);
+        PowerMock.replayAll();
+        assertEquals(expectedConfigs, Worker.consumerConfigs(new ConnectorTaskId("test", 1), configWithOverrides, connectorConfig,
+                                                             null, allConnectorClientConfigOverridePolicy));
+    }
+
+    @Test(expected = ConnectException.class)
+    public void testConsumerConfigsClientOverridesWithNonePolicy() {
+        Map props = new HashMap<>(workerProps);
+        props.put("consumer.auto.offset.reset", "latest");
+        props.put("consumer.max.poll.records", "5000");
+        WorkerConfig configWithOverrides = new StandaloneConfig(props);
+
+        Map connConfig = new HashMap();
+        connConfig.put("max.poll.records", "5000");
+        connConfig.put("max.poll.interval.ms", "1000");
+        EasyMock.expect(connectorConfig.originalsWithPrefix(ConnectorConfig.CONNECTOR_CLIENT_CONSUMER_OVERRIDES_PREFIX))
+            .andReturn(connConfig);
+        PowerMock.replayAll();
+        Worker.consumerConfigs(new ConnectorTaskId("test", 1), configWithOverrides, connectorConfig,
+                               null, noneConnectorClientConfigOverridePolicy);
+    }
+
+    @Test
+    public void testAdminConfigsClientOverridesWithAllPolicy() {
+        Map props = new HashMap<>(workerProps);
+        props.put("admin.client.id", "testid");
+        props.put("admin.metadata.max.age.ms", "5000");
+        props.put("producer.bootstrap.servers", "cbeauho.com");
+        props.put("consumer.bootstrap.servers", "localhost:4761");
+        WorkerConfig configWithOverrides = new StandaloneConfig(props);
+
+        Map connConfig = new HashMap();
+        connConfig.put("metadata.max.age.ms", "10000");
+
+        Map expectedConfigs = new HashMap<>(workerProps);
+
+        expectedConfigs.put("bootstrap.servers", "localhost:9092");
+        expectedConfigs.put("client.id", "testid");
+        expectedConfigs.put("metadata.max.age.ms", "10000");
+
+        EasyMock.expect(connectorConfig.originalsWithPrefix(ConnectorConfig.CONNECTOR_CLIENT_ADMIN_OVERRIDES_PREFIX))
+            .andReturn(connConfig);
+        PowerMock.replayAll();
+        assertEquals(expectedConfigs, Worker.adminConfigs(new ConnectorTaskId("test", 1), configWithOverrides, connectorConfig,
+                                                             null, allConnectorClientConfigOverridePolicy));
+    }
+
+    @Test(expected = ConnectException.class)
+    public void testAdminConfigsClientOverridesWithNonePolicy() {
+        Map props = new HashMap<>(workerProps);
+        props.put("admin.client.id", "testid");
+        props.put("admin.metadata.max.age.ms", "5000");
+        WorkerConfig configWithOverrides = new StandaloneConfig(props);
+
+        Map connConfig = new HashMap();
+        connConfig.put("metadata.max.age.ms", "10000");
+
+        EasyMock.expect(connectorConfig.originalsWithPrefix(ConnectorConfig.CONNECTOR_CLIENT_ADMIN_OVERRIDES_PREFIX))
+            .andReturn(connConfig);
+        PowerMock.replayAll();
+        Worker.adminConfigs(new ConnectorTaskId("test", 1), configWithOverrides, connectorConfig,
+                                                          null, noneConnectorClientConfigOverridePolicy);
+    }
+
+    private void assertStatusMetrics(long expected, String metricName) {
+        MetricGroup statusMetrics = worker.connectorStatusMetricsGroup().metricGroup(TASK_ID.connector());
+        if (expected == 0L) {
+            assertNull(statusMetrics);
+            return;
+        }
+        assertEquals(expected, (long) MockConnectMetrics.currentMetricValue(worker.metrics(), statusMetrics, metricName));
     }
 
     private void assertStatistics(Worker worker, int connectors, int tasks) {
+        assertStatusMetrics((long) tasks, "connector-total-task-count");
         MetricGroup workerMetrics = worker.workerMetricsGroup().metricGroup();
         assertEquals(connectors, MockConnectMetrics.currentMetricValueAsDouble(worker.metrics(), workerMetrics, "connector-count"), 0.0001d);
         assertEquals(tasks, MockConnectMetrics.currentMetricValueAsDouble(worker.metrics(), workerMetrics, "task-count"), 0.0001d);
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerTestUtils.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerTestUtils.java
new file mode 100644
index 0000000000000..ed77018f2883b
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerTestUtils.java
@@ -0,0 +1,193 @@
+/*
+ * 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.connect.runtime.distributed.ClusterConfigState;
+import org.apache.kafka.connect.runtime.distributed.ExtendedAssignment;
+import org.apache.kafka.connect.runtime.distributed.ExtendedWorkerState;
+import org.apache.kafka.connect.util.ConnectorTaskId;
+
+import java.util.AbstractMap.SimpleEntry;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import static org.apache.kafka.connect.runtime.distributed.WorkerCoordinator.WorkerLoad;
+import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+
+public class WorkerTestUtils {
+
+    public static WorkerLoad emptyWorkerLoad(String worker) {
+        return new WorkerLoad.Builder(worker).build();
+    }
+
+    public WorkerLoad workerLoad(String worker, int connectorStart, int connectorNum,
+                                  int taskStart, int taskNum) {
+        return new WorkerLoad.Builder(worker).with(
+                newConnectors(connectorStart, connectorStart + connectorNum),
+                newTasks(taskStart, taskStart + taskNum)).build();
+    }
+
+    public static List newConnectors(int start, int end) {
+        return IntStream.range(start, end)
+                .mapToObj(i -> "connector" + i)
+                .collect(Collectors.toList());
+    }
+
+    public static List newTasks(int start, int end) {
+        return IntStream.range(start, end)
+                .mapToObj(i -> new ConnectorTaskId("task", i))
+                .collect(Collectors.toList());
+    }
+
+    public static ClusterConfigState clusterConfigState(long offset,
+                                                        int connectorNum,
+                                                        int taskNum) {
+        return new ClusterConfigState(
+                offset,
+                null,
+                connectorTaskCounts(1, connectorNum, taskNum),
+                connectorConfigs(1, connectorNum),
+                connectorTargetStates(1, connectorNum, TargetState.STARTED),
+                taskConfigs(0, connectorNum, connectorNum * taskNum),
+                Collections.emptySet());
+    }
+
+    public static Map memberConfigs(String givenLeader,
+                                                                 long givenOffset,
+                                                                 Map givenAssignments) {
+        return givenAssignments.entrySet().stream()
+                .collect(Collectors.toMap(
+                    Map.Entry::getKey,
+                    e -> new ExtendedWorkerState(expectedLeaderUrl(givenLeader), givenOffset, e.getValue())));
+    }
+
+    public static Map memberConfigs(String givenLeader,
+                                                                 long givenOffset,
+                                                                 int start,
+                                                                 int connectorNum) {
+        return IntStream.range(start, connectorNum + 1)
+                .mapToObj(i -> new SimpleEntry<>("worker" + i, new ExtendedWorkerState(expectedLeaderUrl(givenLeader), givenOffset, null)))
+                .collect(Collectors.toMap(SimpleEntry::getKey, SimpleEntry::getValue));
+    }
+
+    public static Map connectorTaskCounts(int start,
+                                                           int connectorNum,
+                                                           int taskCounts) {
+        return IntStream.range(start, connectorNum + 1)
+                .mapToObj(i -> new SimpleEntry<>("connector" + i, taskCounts))
+                .collect(Collectors.toMap(SimpleEntry::getKey, SimpleEntry::getValue));
+    }
+
+    public static Map> connectorConfigs(int start, int connectorNum) {
+        return IntStream.range(start, connectorNum + 1)
+                .mapToObj(i -> new SimpleEntry<>("connector" + i, new HashMap()))
+                .collect(Collectors.toMap(SimpleEntry::getKey, SimpleEntry::getValue));
+    }
+
+    public static Map connectorTargetStates(int start,
+                                                                 int connectorNum,
+                                                                 TargetState state) {
+        return IntStream.range(start, connectorNum + 1)
+                .mapToObj(i -> new SimpleEntry<>("connector" + i, state))
+                .collect(Collectors.toMap(SimpleEntry::getKey, SimpleEntry::getValue));
+    }
+
+    public static Map> taskConfigs(int start,
+                                                                        int connectorNum,
+                                                                        int taskNum) {
+        return IntStream.range(start, taskNum + 1)
+                .mapToObj(i -> new SimpleEntry<>(
+                        new ConnectorTaskId("connector" + i / connectorNum + 1, i),
+                        new HashMap())
+                ).collect(Collectors.toMap(SimpleEntry::getKey, SimpleEntry::getValue));
+    }
+
+    public static String expectedLeaderUrl(String givenLeader) {
+        return "http://" + givenLeader + ":8083";
+    }
+
+    public static void assertAssignment(String expectedLeader,
+                                        long expectedOffset,
+                                        List expectedAssignedConnectors,
+                                        int expectedAssignedTaskNum,
+                                        List expectedRevokedConnectors,
+                                        int expectedRevokedTaskNum,
+                                        ExtendedAssignment assignment) {
+        assertAssignment(false, expectedLeader, expectedOffset,
+                expectedAssignedConnectors, expectedAssignedTaskNum,
+                expectedRevokedConnectors, expectedRevokedTaskNum,
+                0,
+                assignment);
+    }
+
+    public static void assertAssignment(String expectedLeader,
+                                        long expectedOffset,
+                                        List expectedAssignedConnectors,
+                                        int expectedAssignedTaskNum,
+                                        List expectedRevokedConnectors,
+                                        int expectedRevokedTaskNum,
+                                        int expectedDelay,
+                                        ExtendedAssignment assignment) {
+        assertAssignment(false, expectedLeader, expectedOffset,
+                expectedAssignedConnectors, expectedAssignedTaskNum,
+                expectedRevokedConnectors, expectedRevokedTaskNum,
+                expectedDelay,
+                assignment);
+    }
+
+    public static void assertAssignment(boolean expectFailed,
+                                        String expectedLeader,
+                                        long expectedOffset,
+                                        List expectedAssignedConnectors,
+                                        int expectedAssignedTaskNum,
+                                        List expectedRevokedConnectors,
+                                        int expectedRevokedTaskNum,
+                                        int expectedDelay,
+                                        ExtendedAssignment assignment) {
+        assertNotNull("Assignment can't be null", assignment);
+
+        assertEquals("Wrong status in " + assignment, expectFailed, assignment.failed());
+
+        assertEquals("Wrong leader in " + assignment, expectedLeader, assignment.leader());
+
+        assertEquals("Wrong leaderUrl in " + assignment, expectedLeaderUrl(expectedLeader),
+                assignment.leaderUrl());
+
+        assertEquals("Wrong offset in " + assignment, expectedOffset, assignment.offset());
+
+        assertThat("Wrong set of assigned connectors in " + assignment,
+                assignment.connectors(), is(expectedAssignedConnectors));
+
+        assertEquals("Wrong number of assigned tasks in " + assignment,
+                expectedAssignedTaskNum, assignment.tasks().size());
+
+        assertThat("Wrong set of revoked connectors in " + assignment,
+                assignment.revokedConnectors(), is(expectedRevokedConnectors));
+
+        assertEquals("Wrong number of revoked tasks in " + assignment,
+                expectedRevokedTaskNum, assignment.revokedTasks().size());
+
+        assertEquals("Wrong rebalance delay in " + assignment, expectedDelay, assignment.delay());
+    }
+}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/distributed/ConnectProtocolCompatibilityTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/distributed/ConnectProtocolCompatibilityTest.java
new file mode 100644
index 0000000000000..a3144c04d3462
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/distributed/ConnectProtocolCompatibilityTest.java
@@ -0,0 +1,259 @@
+/*
+ * 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.distributed;
+
+import org.apache.kafka.connect.runtime.TargetState;
+import org.apache.kafka.connect.storage.KafkaConfigBackingStore;
+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.mockito.Mock;
+import org.mockito.junit.MockitoJUnit;
+import org.mockito.junit.MockitoRule;
+
+import java.nio.ByteBuffer;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+
+import static org.apache.kafka.connect.runtime.distributed.IncrementalCooperativeConnectProtocol.CONNECT_PROTOCOL_V1;
+import static org.junit.Assert.assertEquals;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
+import static org.mockito.Mockito.when;
+
+public class ConnectProtocolCompatibilityTest {
+    private static final String LEADER_URL = "leaderUrl:8083";
+
+    private String connectorId1 = "connector1";
+    private String connectorId2 = "connector2";
+    private String connectorId3 = "connector3";
+    private ConnectorTaskId taskId1x0 = new ConnectorTaskId(connectorId1, 0);
+    private ConnectorTaskId taskId1x1 = new ConnectorTaskId(connectorId1, 1);
+    private ConnectorTaskId taskId2x0 = new ConnectorTaskId(connectorId2, 0);
+    private ConnectorTaskId taskId3x0 = new ConnectorTaskId(connectorId3, 0);
+
+    @Rule
+    public MockitoRule rule = MockitoJUnit.rule();
+
+    @Mock
+    private KafkaConfigBackingStore configStorage;
+    private ClusterConfigState configState;
+
+    @Before
+    public void setup() {
+        configStorage = mock(KafkaConfigBackingStore.class);
+        configState = new ClusterConfigState(
+                1L,
+                null,
+                Collections.singletonMap(connectorId1, 1),
+                Collections.singletonMap(connectorId1, new HashMap<>()),
+                Collections.singletonMap(connectorId1, TargetState.STARTED),
+                Collections.singletonMap(taskId1x0, new HashMap<>()),
+                Collections.emptySet());
+    }
+
+    @After
+    public void teardown() {
+        verifyNoMoreInteractions(configStorage);
+    }
+
+    @Test
+    public void testEagerToEagerMetadata() {
+        when(configStorage.snapshot()).thenReturn(configState);
+        ExtendedWorkerState workerState = new ExtendedWorkerState(LEADER_URL, configStorage.snapshot().offset(), null);
+        ByteBuffer metadata = ConnectProtocol.serializeMetadata(workerState);
+        ConnectProtocol.WorkerState state = ConnectProtocol.deserializeMetadata(metadata);
+        assertEquals(LEADER_URL, state.url());
+        assertEquals(1, state.offset());
+        verify(configStorage).snapshot();
+    }
+
+    @Test
+    public void testCoopToCoopMetadata() {
+        when(configStorage.snapshot()).thenReturn(configState);
+        ExtendedWorkerState workerState = new ExtendedWorkerState(LEADER_URL, configStorage.snapshot().offset(), null);
+        ByteBuffer metadata = IncrementalCooperativeConnectProtocol.serializeMetadata(workerState, false);
+        ExtendedWorkerState state = IncrementalCooperativeConnectProtocol.deserializeMetadata(metadata);
+        assertEquals(LEADER_URL, state.url());
+        assertEquals(1, state.offset());
+        verify(configStorage).snapshot();
+    }
+
+    @Test
+    public void testSessionedToCoopMetadata() {
+        when(configStorage.snapshot()).thenReturn(configState);
+        ExtendedWorkerState workerState = new ExtendedWorkerState(LEADER_URL, configStorage.snapshot().offset(), null);
+        ByteBuffer metadata = IncrementalCooperativeConnectProtocol.serializeMetadata(workerState, true);
+        ExtendedWorkerState state = IncrementalCooperativeConnectProtocol.deserializeMetadata(metadata);
+        assertEquals(LEADER_URL, state.url());
+        assertEquals(1, state.offset());
+        verify(configStorage).snapshot();
+    }
+
+    @Test
+    public void testSessionedToEagerMetadata() {
+        when(configStorage.snapshot()).thenReturn(configState);
+        ExtendedWorkerState workerState = new ExtendedWorkerState(LEADER_URL, configStorage.snapshot().offset(), null);
+        ByteBuffer metadata = IncrementalCooperativeConnectProtocol.serializeMetadata(workerState, true);
+        ConnectProtocol.WorkerState state = ConnectProtocol.deserializeMetadata(metadata);
+        assertEquals(LEADER_URL, state.url());
+        assertEquals(1, state.offset());
+        verify(configStorage).snapshot();
+    }
+
+    @Test
+    public void testCoopToEagerMetadata() {
+        when(configStorage.snapshot()).thenReturn(configState);
+        ExtendedWorkerState workerState = new ExtendedWorkerState(LEADER_URL, configStorage.snapshot().offset(), null);
+        ByteBuffer metadata = IncrementalCooperativeConnectProtocol.serializeMetadata(workerState, false);
+        ConnectProtocol.WorkerState state = ConnectProtocol.deserializeMetadata(metadata);
+        assertEquals(LEADER_URL, state.url());
+        assertEquals(1, state.offset());
+        verify(configStorage).snapshot();
+    }
+
+    @Test
+    public void testEagerToCoopMetadata() {
+        when(configStorage.snapshot()).thenReturn(configState);
+        ConnectProtocol.WorkerState workerState = new ConnectProtocol.WorkerState(LEADER_URL, configStorage.snapshot().offset());
+        ByteBuffer metadata = ConnectProtocol.serializeMetadata(workerState);
+        ConnectProtocol.WorkerState state = IncrementalCooperativeConnectProtocol.deserializeMetadata(metadata);
+        assertEquals(LEADER_URL, state.url());
+        assertEquals(1, state.offset());
+        verify(configStorage).snapshot();
+    }
+
+    @Test
+    public void testEagerToEagerAssignment() {
+        ConnectProtocol.Assignment assignment = new ConnectProtocol.Assignment(
+                ConnectProtocol.Assignment.NO_ERROR, "leader", LEADER_URL, 1L,
+                Arrays.asList(connectorId1, connectorId3), Arrays.asList(taskId2x0));
+
+        ByteBuffer leaderBuf = ConnectProtocol.serializeAssignment(assignment);
+        ConnectProtocol.Assignment leaderAssignment = ConnectProtocol.deserializeAssignment(leaderBuf);
+        assertEquals(false, leaderAssignment.failed());
+        assertEquals("leader", leaderAssignment.leader());
+        assertEquals(1, leaderAssignment.offset());
+        assertEquals(Arrays.asList(connectorId1, connectorId3), leaderAssignment.connectors());
+        assertEquals(Collections.singletonList(taskId2x0), leaderAssignment.tasks());
+
+        ConnectProtocol.Assignment assignment2 = new ConnectProtocol.Assignment(
+                ConnectProtocol.Assignment.NO_ERROR, "member", LEADER_URL, 1L,
+                Arrays.asList(connectorId2), Arrays.asList(taskId1x0, taskId3x0));
+
+        ByteBuffer memberBuf = ConnectProtocol.serializeAssignment(assignment2);
+        ConnectProtocol.Assignment memberAssignment = ConnectProtocol.deserializeAssignment(memberBuf);
+        assertEquals(false, memberAssignment.failed());
+        assertEquals("member", memberAssignment.leader());
+        assertEquals(1, memberAssignment.offset());
+        assertEquals(Collections.singletonList(connectorId2), memberAssignment.connectors());
+        assertEquals(Arrays.asList(taskId1x0, taskId3x0), memberAssignment.tasks());
+    }
+
+    @Test
+    public void testCoopToCoopAssignment() {
+        ExtendedAssignment assignment = new ExtendedAssignment(
+                CONNECT_PROTOCOL_V1, ConnectProtocol.Assignment.NO_ERROR, "leader", LEADER_URL, 1L,
+                Arrays.asList(connectorId1, connectorId3), Arrays.asList(taskId2x0),
+                Collections.emptyList(), Collections.emptyList(), 0);
+
+        ByteBuffer leaderBuf = IncrementalCooperativeConnectProtocol.serializeAssignment(assignment);
+        ConnectProtocol.Assignment leaderAssignment = ConnectProtocol.deserializeAssignment(leaderBuf);
+        assertEquals(false, leaderAssignment.failed());
+        assertEquals("leader", leaderAssignment.leader());
+        assertEquals(1, leaderAssignment.offset());
+        assertEquals(Arrays.asList(connectorId1, connectorId3), leaderAssignment.connectors());
+        assertEquals(Collections.singletonList(taskId2x0), leaderAssignment.tasks());
+
+        ExtendedAssignment assignment2 = new ExtendedAssignment(
+                CONNECT_PROTOCOL_V1, ConnectProtocol.Assignment.NO_ERROR, "member", LEADER_URL, 1L,
+                Arrays.asList(connectorId2), Arrays.asList(taskId1x0, taskId3x0),
+                Collections.emptyList(), Collections.emptyList(), 0);
+
+        ByteBuffer memberBuf = ConnectProtocol.serializeAssignment(assignment2);
+        ConnectProtocol.Assignment memberAssignment =
+                IncrementalCooperativeConnectProtocol.deserializeAssignment(memberBuf);
+        assertEquals(false, memberAssignment.failed());
+        assertEquals("member", memberAssignment.leader());
+        assertEquals(1, memberAssignment.offset());
+        assertEquals(Collections.singletonList(connectorId2), memberAssignment.connectors());
+        assertEquals(Arrays.asList(taskId1x0, taskId3x0), memberAssignment.tasks());
+    }
+
+    @Test
+    public void testEagerToCoopAssignment() {
+        ConnectProtocol.Assignment assignment = new ConnectProtocol.Assignment(
+                ConnectProtocol.Assignment.NO_ERROR, "leader", LEADER_URL, 1L,
+                Arrays.asList(connectorId1, connectorId3), Arrays.asList(taskId2x0));
+
+        ByteBuffer leaderBuf = ConnectProtocol.serializeAssignment(assignment);
+        ConnectProtocol.Assignment leaderAssignment =
+                IncrementalCooperativeConnectProtocol.deserializeAssignment(leaderBuf);
+        assertEquals(false, leaderAssignment.failed());
+        assertEquals("leader", leaderAssignment.leader());
+        assertEquals(1, leaderAssignment.offset());
+        assertEquals(Arrays.asList(connectorId1, connectorId3), leaderAssignment.connectors());
+        assertEquals(Collections.singletonList(taskId2x0), leaderAssignment.tasks());
+
+        ConnectProtocol.Assignment assignment2 = new ConnectProtocol.Assignment(
+                ConnectProtocol.Assignment.NO_ERROR, "member", LEADER_URL, 1L,
+                Arrays.asList(connectorId2), Arrays.asList(taskId1x0, taskId3x0));
+
+        ByteBuffer memberBuf = ConnectProtocol.serializeAssignment(assignment2);
+        ConnectProtocol.Assignment memberAssignment =
+                IncrementalCooperativeConnectProtocol.deserializeAssignment(memberBuf);
+        assertEquals(false, memberAssignment.failed());
+        assertEquals("member", memberAssignment.leader());
+        assertEquals(1, memberAssignment.offset());
+        assertEquals(Collections.singletonList(connectorId2), memberAssignment.connectors());
+        assertEquals(Arrays.asList(taskId1x0, taskId3x0), memberAssignment.tasks());
+    }
+
+    @Test
+    public void testCoopToEagerAssignment() {
+        ExtendedAssignment assignment = new ExtendedAssignment(
+                CONNECT_PROTOCOL_V1, ConnectProtocol.Assignment.NO_ERROR, "leader", LEADER_URL, 1L,
+                Arrays.asList(connectorId1, connectorId3), Arrays.asList(taskId2x0),
+                Collections.emptyList(), Collections.emptyList(), 0);
+
+        ByteBuffer leaderBuf = IncrementalCooperativeConnectProtocol.serializeAssignment(assignment);
+        ConnectProtocol.Assignment leaderAssignment = ConnectProtocol.deserializeAssignment(leaderBuf);
+        assertEquals(false, leaderAssignment.failed());
+        assertEquals("leader", leaderAssignment.leader());
+        assertEquals(1, leaderAssignment.offset());
+        assertEquals(Arrays.asList(connectorId1, connectorId3), leaderAssignment.connectors());
+        assertEquals(Collections.singletonList(taskId2x0), leaderAssignment.tasks());
+
+        ExtendedAssignment assignment2 = new ExtendedAssignment(
+                CONNECT_PROTOCOL_V1, ConnectProtocol.Assignment.NO_ERROR, "member", LEADER_URL, 1L,
+                Arrays.asList(connectorId2), Arrays.asList(taskId1x0, taskId3x0),
+                Collections.emptyList(), Collections.emptyList(), 0);
+
+        ByteBuffer memberBuf = IncrementalCooperativeConnectProtocol.serializeAssignment(assignment2);
+        ConnectProtocol.Assignment memberAssignment = ConnectProtocol.deserializeAssignment(memberBuf);
+        assertEquals(false, memberAssignment.failed());
+        assertEquals("member", memberAssignment.leader());
+        assertEquals(1, memberAssignment.offset());
+        assertEquals(Collections.singletonList(connectorId2), memberAssignment.connectors());
+        assertEquals(Arrays.asList(taskId1x0, taskId3x0), memberAssignment.tasks());
+    }
+
+}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/distributed/DistributedConfigTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/distributed/DistributedConfigTest.java
new file mode 100644
index 0000000000000..630465b912f7b
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/distributed/DistributedConfigTest.java
@@ -0,0 +1,108 @@
+/*
+ * 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.distributed;
+
+import org.apache.kafka.common.config.ConfigException;
+import org.junit.Test;
+
+import javax.crypto.KeyGenerator;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertThrows;
+
+public class DistributedConfigTest {
+
+    public Map configs() {
+        Map result = new HashMap<>();
+        result.put(DistributedConfig.GROUP_ID_CONFIG, "connect-cluster");
+        result.put(DistributedConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092");
+        result.put(DistributedConfig.CONFIG_TOPIC_CONFIG, "connect-configs");
+        result.put(DistributedConfig.OFFSET_STORAGE_TOPIC_CONFIG, "connect-offsets");
+        result.put(DistributedConfig.STATUS_STORAGE_TOPIC_CONFIG, "connect-status");
+        result.put(DistributedConfig.KEY_CONVERTER_CLASS_CONFIG, "org.apache.kafka.connect.json.JsonConverter");
+        result.put(DistributedConfig.VALUE_CONVERTER_CLASS_CONFIG, "org.apache.kafka.connect.json.JsonConverter");
+        return result;
+    }
+
+    @Test
+    public void shouldCreateKeyGeneratorWithDefaultSettings() {
+        DistributedConfig config = new DistributedConfig(configs());
+        assertNotNull(config.getInternalRequestKeyGenerator());
+    }
+
+    @Test
+    public void shouldCreateKeyGeneratorWithSpecificSettings() {
+        final String algorithm = "HmacSHA1";
+        Map configs = configs();
+        configs.put(DistributedConfig.INTER_WORKER_KEY_GENERATION_ALGORITHM_CONFIG, algorithm);
+        configs.put(DistributedConfig.INTER_WORKER_KEY_SIZE_CONFIG, "512");
+        configs.put(DistributedConfig.INTER_WORKER_VERIFICATION_ALGORITHMS_CONFIG, algorithm);
+        DistributedConfig config = new DistributedConfig(configs);
+        KeyGenerator keyGenerator = config.getInternalRequestKeyGenerator();
+        assertNotNull(keyGenerator);
+        assertEquals(algorithm, keyGenerator.getAlgorithm());
+        assertEquals(512 / 8, keyGenerator.generateKey().getEncoded().length);
+    }
+
+    @Test(expected = ConfigException.class)
+    public void shouldFailWithEmptyListOfVerificationAlgorithms() {
+        Map configs = configs();
+        configs.put(DistributedConfig.INTER_WORKER_VERIFICATION_ALGORITHMS_CONFIG, "");
+        new DistributedConfig(configs);
+    }
+
+    @Test(expected = ConfigException.class)
+    public void shouldFailIfKeyAlgorithmNotInVerificationAlgorithmsList() {
+        Map configs = configs();
+        configs.put(DistributedConfig.INTER_WORKER_KEY_GENERATION_ALGORITHM_CONFIG, "HmacSHA1");
+        configs.put(DistributedConfig.INTER_WORKER_VERIFICATION_ALGORITHMS_CONFIG, "HmacSHA256");
+        new DistributedConfig(configs);
+    }
+
+    @Test(expected = ConfigException.class)
+    public void shouldFailWithInvalidKeyAlgorithm() {
+        Map configs = configs();
+        configs.put(DistributedConfig.INTER_WORKER_KEY_GENERATION_ALGORITHM_CONFIG, "not-actually-a-key-algorithm");
+        new DistributedConfig(configs);
+    }
+
+    @Test(expected = ConfigException.class)
+    public void shouldFailWithInvalidKeySize() {
+        Map configs = configs();
+        configs.put(DistributedConfig.INTER_WORKER_KEY_SIZE_CONFIG, "0");
+        new DistributedConfig(configs);
+    }
+
+    @Test
+    public void shouldValidateAllVerificationAlgorithms() {
+        List algorithms =
+            new ArrayList<>(Arrays.asList("HmacSHA1", "HmacSHA256", "HmacMD5", "bad-algorithm"));
+        Map configs = configs();
+        for (int i = 0; i < algorithms.size(); i++) {
+            configs.put(DistributedConfig.INTER_WORKER_VERIFICATION_ALGORITHMS_CONFIG, String.join(",", algorithms));
+            assertThrows(ConfigException.class, () -> new DistributedConfig(configs));
+            algorithms.add(algorithms.remove(0));
+        }
+    }
+}
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 b8ec0f1e57fe2..81208d398b415 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
@@ -23,6 +23,8 @@
 import org.apache.kafka.common.utils.MockTime;
 import org.apache.kafka.connect.connector.Connector;
 import org.apache.kafka.connect.connector.ConnectorContext;
+import org.apache.kafka.connect.connector.policy.ConnectorClientConfigOverridePolicy;
+import org.apache.kafka.connect.connector.policy.NoneConnectorClientConfigOverridePolicy;
 import org.apache.kafka.connect.errors.AlreadyExistsException;
 import org.apache.kafka.connect.errors.NotFoundException;
 import org.apache.kafka.connect.runtime.ConnectMetrics.MetricGroup;
@@ -39,10 +41,12 @@
 import org.apache.kafka.connect.runtime.isolation.DelegatingClassLoader;
 import org.apache.kafka.connect.runtime.isolation.PluginClassLoader;
 import org.apache.kafka.connect.runtime.isolation.Plugins;
+import org.apache.kafka.connect.runtime.rest.InternalRequestSignature;
 import org.apache.kafka.connect.runtime.rest.entities.ConnectorInfo;
 import org.apache.kafka.connect.runtime.rest.entities.ConnectorType;
 import org.apache.kafka.connect.runtime.rest.entities.TaskInfo;
 import org.apache.kafka.connect.runtime.rest.errors.BadRequestException;
+import org.apache.kafka.connect.runtime.rest.errors.ConnectRestException;
 import org.apache.kafka.connect.sink.SinkConnector;
 import org.apache.kafka.connect.source.SourceConnector;
 import org.apache.kafka.connect.source.SourceTask;
@@ -76,14 +80,21 @@
 import java.util.concurrent.TimeoutException;
 
 import static java.util.Collections.singletonList;
+import static javax.ws.rs.core.Response.Status.FORBIDDEN;
+import static org.apache.kafka.connect.runtime.distributed.ConnectProtocol.CONNECT_PROTOCOL_V0;
+import static org.apache.kafka.connect.runtime.distributed.IncrementalCooperativeConnectProtocol.CONNECT_PROTOCOL_V1;
+import static org.apache.kafka.connect.runtime.distributed.IncrementalCooperativeConnectProtocol.CONNECT_PROTOCOL_V2;
+import static org.easymock.EasyMock.capture;
+import static org.easymock.EasyMock.newCapture;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
 @SuppressWarnings("deprecation")
 @RunWith(PowerMockRunner.class)
 @PrepareForTest({DistributedHerder.class, Plugins.class})
-@PowerMockIgnore("javax.management.*")
+@PowerMockIgnore({"javax.management.*", "javax.crypto.*"})
 public class DistributedHerderTest {
     private static final Map HERDER_CONFIG = new HashMap<>();
     static {
@@ -140,13 +151,13 @@ public class DistributedHerderTest {
         TASK_CONFIGS_MAP.put(TASK1, TASK_CONFIG);
         TASK_CONFIGS_MAP.put(TASK2, TASK_CONFIG);
     }
-    private static final ClusterConfigState SNAPSHOT = new ClusterConfigState(1, Collections.singletonMap(CONN1, 3),
+    private static final ClusterConfigState SNAPSHOT = new ClusterConfigState(1, null, Collections.singletonMap(CONN1, 3),
             Collections.singletonMap(CONN1, CONN1_CONFIG), Collections.singletonMap(CONN1, TargetState.STARTED),
             TASK_CONFIGS_MAP, Collections.emptySet());
-    private static final ClusterConfigState SNAPSHOT_PAUSED_CONN1 = new ClusterConfigState(1, Collections.singletonMap(CONN1, 3),
+    private static final ClusterConfigState SNAPSHOT_PAUSED_CONN1 = new ClusterConfigState(1, null, Collections.singletonMap(CONN1, 3),
             Collections.singletonMap(CONN1, CONN1_CONFIG), Collections.singletonMap(CONN1, TargetState.PAUSED),
             TASK_CONFIGS_MAP, Collections.emptySet());
-    private static final ClusterConfigState SNAPSHOT_UPDATED_CONN1_CONFIG = new ClusterConfigState(1, Collections.singletonMap(CONN1, 3),
+    private static final ClusterConfigState SNAPSHOT_UPDATED_CONN1_CONFIG = new ClusterConfigState(1, null, Collections.singletonMap(CONN1, 3),
             Collections.singletonMap(CONN1, CONN1_CONFIG_UPDATED), Collections.singletonMap(CONN1, TargetState.STARTED),
             TASK_CONFIGS_MAP, Collections.emptySet());
 
@@ -174,6 +185,9 @@ public class DistributedHerderTest {
 
     private SinkConnectorConfig conn1SinkConfig;
     private SinkConnectorConfig conn1SinkConfigUpdated;
+    private short connectProtocolVersion;
+    private final ConnectorClientConfigOverridePolicy
+        noneConnectorClientConfigOverridePolicy = new NoneConnectorClientConfigOverridePolicy();
 
     @Before
     public void setUp() throws Exception {
@@ -182,13 +196,16 @@ public void setUp() throws Exception {
         worker = PowerMock.createMock(Worker.class);
         EasyMock.expect(worker.isSinkConnector(CONN1)).andStubReturn(Boolean.TRUE);
 
+        // Default to the old protocol unless specified otherwise
+        connectProtocolVersion = CONNECT_PROTOCOL_V0;
+
         herder = PowerMock.createPartialMock(DistributedHerder.class,
-                new String[]{"backoff", "connectorTypeForClass", "updateDeletedConnectorStatus"},
+                new String[]{"connectorTypeForClass", "updateDeletedConnectorStatus", "updateDeletedTaskStatus"},
                 new DistributedConfig(HERDER_CONFIG), worker, WORKER_ID, KAFKA_CLUSTER_ID,
-                statusBackingStore, configBackingStore, member, MEMBER_URL, metrics, time);
+                statusBackingStore, configBackingStore, member, MEMBER_URL, metrics, time, noneConnectorClientConfigOverridePolicy);
 
         configUpdateListener = herder.new ConfigUpdateListener();
-        rebalanceListener = herder.new RebalanceListener();
+        rebalanceListener = herder.new RebalanceListener(time);
         plugins = PowerMock.createMock(Plugins.class);
         conn1SinkConfig = new SinkConnectorConfig(plugins, CONN1_CONFIG);
         conn1SinkConfigUpdated = new SinkConnectorConfig(plugins, CONN1_CONFIG_UPDATED);
@@ -197,6 +214,7 @@ public void setUp() throws Exception {
         delegatingLoader = PowerMock.createMock(DelegatingClassLoader.class);
         PowerMock.mockStatic(Plugins.class);
         PowerMock.expectPrivate(herder, "updateDeletedConnectorStatus").andVoid().anyTimes();
+        PowerMock.expectPrivate(herder, "updateDeletedTaskStatus").andVoid().anyTimes();
     }
 
     @After
@@ -208,6 +226,7 @@ public void tearDown() {
     public void testJoinAssignment() throws Exception {
         // Join group and get assignment
         EasyMock.expect(member.memberId()).andStubReturn("member");
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V0);
         EasyMock.expect(worker.getPlugins()).andReturn(plugins);
         expectRebalance(1, Arrays.asList(CONN1), Arrays.asList(TASK1));
         expectPostRebalanceCatchup(SNAPSHOT);
@@ -236,6 +255,7 @@ public void testJoinAssignment() throws Exception {
     public void testRebalance() throws Exception {
         // Join group and get assignment
         EasyMock.expect(member.memberId()).andStubReturn("member");
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V0);
         EasyMock.expect(worker.getPlugins()).andReturn(plugins);
         expectRebalance(1, Arrays.asList(CONN1), Arrays.asList(TASK1));
         expectPostRebalanceCatchup(SNAPSHOT);
@@ -279,10 +299,157 @@ public void testRebalance() throws Exception {
         PowerMock.verifyAll();
     }
 
+    @Test
+    public void testIncrementalCooperativeRebalanceForNewMember() throws Exception {
+        connectProtocolVersion = CONNECT_PROTOCOL_V1;
+        // Join group. First rebalance contains revocations from other members. For the new
+        // member the assignment should be empty
+        EasyMock.expect(member.memberId()).andStubReturn("member");
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V1);
+        expectRebalance(1, Collections.emptyList(), Collections.emptyList());
+        expectPostRebalanceCatchup(SNAPSHOT);
+
+        member.poll(EasyMock.anyInt());
+        PowerMock.expectLastCall();
+
+        // The new member got its assignment
+        expectRebalance(Collections.emptyList(), Collections.emptyList(),
+                ConnectProtocol.Assignment.NO_ERROR,
+                1, Arrays.asList(CONN1), Arrays.asList(TASK1), 0);
+
+        EasyMock.expect(worker.getPlugins()).andReturn(plugins);
+        // and the new assignment started
+        worker.startConnector(EasyMock.eq(CONN1), EasyMock.anyObject(), EasyMock.anyObject(),
+                EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED));
+        PowerMock.expectLastCall().andReturn(true);
+        EasyMock.expect(worker.isRunning(CONN1)).andReturn(true);
+        EasyMock.expect(worker.connectorTaskConfigs(CONN1, conn1SinkConfig)).andReturn(TASK_CONFIGS);
+
+        worker.startTask(EasyMock.eq(TASK1), EasyMock.anyObject(), EasyMock.anyObject(), EasyMock.anyObject(),
+                EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED));
+        PowerMock.expectLastCall().andReturn(true);
+        member.poll(EasyMock.anyInt());
+        PowerMock.expectLastCall();
+
+        PowerMock.replayAll();
+
+        time.sleep(1000L);
+        assertStatistics(0, 0, 0, Double.POSITIVE_INFINITY);
+        herder.tick();
+
+        time.sleep(2000L);
+        assertStatistics(3, 1, 100, 2000);
+        herder.tick();
+
+        time.sleep(3000L);
+        assertStatistics(3, 2, 100, 3000);
+
+        PowerMock.verifyAll();
+    }
+
+    @Test
+    public void testIncrementalCooperativeRebalanceForExistingMember() throws Exception {
+        connectProtocolVersion = CONNECT_PROTOCOL_V1;
+        // Join group. First rebalance contains revocations because a new member joined.
+        EasyMock.expect(member.memberId()).andStubReturn("member");
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V1);
+        expectRebalance(Arrays.asList(CONN1), Arrays.asList(TASK1),
+                ConnectProtocol.Assignment.NO_ERROR, 1,
+                Collections.emptyList(), Collections.emptyList(), 0);
+        member.requestRejoin();
+        PowerMock.expectLastCall();
+
+        expectPostRebalanceCatchup(SNAPSHOT);
+
+        // In the second rebalance the new member gets its assignment and this member has no
+        // assignments or revocations
+        expectRebalance(1, Collections.emptyList(), Collections.emptyList());
+
+        member.poll(EasyMock.anyInt());
+        PowerMock.expectLastCall();
+
+        PowerMock.replayAll();
+
+        time.sleep(1000L);
+        assertStatistics(0, 0, 0, Double.POSITIVE_INFINITY);
+        herder.tick();
+
+        time.sleep(2000L);
+        assertStatistics(3, 1, 100, 2000);
+        herder.tick();
+
+        time.sleep(3000L);
+        assertStatistics(3, 2, 100, 3000);
+
+        PowerMock.verifyAll();
+    }
+
+    @Test
+    public void testIncrementalCooperativeRebalanceWithDelay() throws Exception {
+        connectProtocolVersion = CONNECT_PROTOCOL_V1;
+        // Join group. First rebalance contains some assignments but also a delay, because a
+        // member was detected missing
+        int rebalanceDelay = 10_000;
+
+        EasyMock.expect(member.memberId()).andStubReturn("member");
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V1);
+        expectRebalance(Collections.emptyList(), Collections.emptyList(),
+                ConnectProtocol.Assignment.NO_ERROR, 1,
+                Collections.emptyList(), Arrays.asList(TASK2),
+                rebalanceDelay);
+        expectPostRebalanceCatchup(SNAPSHOT);
+
+        worker.startTask(EasyMock.eq(TASK2), EasyMock.anyObject(), EasyMock.anyObject(), EasyMock.anyObject(),
+                EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED));
+        PowerMock.expectLastCall().andReturn(true);
+        member.poll(EasyMock.anyInt());
+        PowerMock.expectLastCall().andAnswer(() -> {
+            time.sleep(9900L);
+            return null;
+        });
+
+        // Request to re-join because the scheduled rebalance delay has been reached
+        member.requestRejoin();
+        PowerMock.expectLastCall();
+
+        // The member got its assignment and revocation
+        expectRebalance(Collections.emptyList(), Collections.emptyList(),
+                ConnectProtocol.Assignment.NO_ERROR,
+                1, Arrays.asList(CONN1), Arrays.asList(TASK1), 0);
+
+        EasyMock.expect(worker.getPlugins()).andReturn(plugins);
+        // and the new assignment started
+        worker.startConnector(EasyMock.eq(CONN1), EasyMock.anyObject(), EasyMock.anyObject(),
+                EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED));
+        PowerMock.expectLastCall().andReturn(true);
+        EasyMock.expect(worker.isRunning(CONN1)).andReturn(true);
+        EasyMock.expect(worker.connectorTaskConfigs(CONN1, conn1SinkConfig)).andReturn(TASK_CONFIGS);
+
+        worker.startTask(EasyMock.eq(TASK1), EasyMock.anyObject(), EasyMock.anyObject(), EasyMock.anyObject(),
+                EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED));
+        PowerMock.expectLastCall().andReturn(true);
+        member.poll(EasyMock.anyInt());
+        PowerMock.expectLastCall();
+
+        PowerMock.replayAll();
+
+        time.sleep(1000L);
+        assertStatistics(0, 0, 0, Double.POSITIVE_INFINITY);
+        herder.tick();
+
+        herder.tick();
+
+        time.sleep(2000L);
+        assertStatistics(3, 2, 100, 2000);
+
+        PowerMock.verifyAll();
+    }
+
     @Test
     public void testRebalanceFailedConnector() throws Exception {
         // Join group and get assignment
         EasyMock.expect(member.memberId()).andStubReturn("member");
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V0);
         EasyMock.expect(worker.getPlugins()).andReturn(plugins);
         expectRebalance(1, Arrays.asList(CONN1), Arrays.asList(TASK1));
         expectPostRebalanceCatchup(SNAPSHOT);
@@ -351,6 +518,7 @@ public void testHaltCleansUpWorker() {
     @Test
     public void testCreateConnector() throws Exception {
         EasyMock.expect(member.memberId()).andStubReturn("leader");
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V0);
         expectRebalance(1, Collections.emptyList(), Collections.emptyList());
         expectPostRebalanceCatchup(SNAPSHOT);
 
@@ -360,7 +528,7 @@ public void testCreateConnector() throws Exception {
         // config validation
         Connector connectorMock = PowerMock.createMock(SourceConnector.class);
         EasyMock.expect(worker.configTransformer()).andReturn(transformer).times(2);
-        final Capture> configCapture = EasyMock.newCapture();
+        final Capture> configCapture = newCapture();
         EasyMock.expect(transformer.transform(EasyMock.capture(configCapture))).andAnswer(configCapture::getValue);
         EasyMock.expect(worker.getPlugins()).andReturn(plugins).times(3);
         EasyMock.expect(plugins.compareAndSwapLoaders(connectorMock)).andReturn(delegatingLoader);
@@ -394,6 +562,7 @@ public void testCreateConnector() throws Exception {
     @Test
     public void testCreateConnectorFailedBasicValidation() throws Exception {
         EasyMock.expect(member.memberId()).andStubReturn("leader");
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V0);
         expectRebalance(1, Collections.emptyList(), Collections.emptyList());
         expectPostRebalanceCatchup(SNAPSHOT);
 
@@ -406,7 +575,7 @@ public void testCreateConnectorFailedBasicValidation() throws Exception {
         // config validation
         Connector connectorMock = PowerMock.createMock(SourceConnector.class);
         EasyMock.expect(worker.configTransformer()).andReturn(transformer).times(2);
-        final Capture> configCapture = EasyMock.newCapture();
+        final Capture> configCapture = newCapture();
         EasyMock.expect(transformer.transform(EasyMock.capture(configCapture))).andAnswer(configCapture::getValue);
         EasyMock.expect(worker.getPlugins()).andReturn(plugins).times(3);
         EasyMock.expect(plugins.compareAndSwapLoaders(connectorMock)).andReturn(delegatingLoader);
@@ -420,7 +589,7 @@ public void testCreateConnectorFailedBasicValidation() throws Exception {
 
         // CONN2 creation should fail
 
-        Capture error = EasyMock.newCapture();
+        Capture error = newCapture();
         putConnectorCallback.onCompletion(EasyMock.capture(error), EasyMock.>isNull());
         PowerMock.expectLastCall();
 
@@ -445,6 +614,7 @@ public void testCreateConnectorFailedBasicValidation() throws Exception {
     @Test
     public void testCreateConnectorFailedCustomValidation() throws Exception {
         EasyMock.expect(member.memberId()).andStubReturn("leader");
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V0);
         expectRebalance(1, Collections.emptyList(), Collections.emptyList());
         expectPostRebalanceCatchup(SNAPSHOT);
 
@@ -454,7 +624,7 @@ public void testCreateConnectorFailedCustomValidation() throws Exception {
         // config validation
         Connector connectorMock = PowerMock.createMock(SourceConnector.class);
         EasyMock.expect(worker.configTransformer()).andReturn(transformer).times(2);
-        final Capture> configCapture = EasyMock.newCapture();
+        final Capture> configCapture = newCapture();
         EasyMock.expect(transformer.transform(EasyMock.capture(configCapture))).andAnswer(configCapture::getValue);
         EasyMock.expect(worker.getPlugins()).andReturn(plugins).times(3);
         EasyMock.expect(plugins.compareAndSwapLoaders(connectorMock)).andReturn(delegatingLoader);
@@ -471,7 +641,7 @@ public void testCreateConnectorFailedCustomValidation() throws Exception {
 
         // CONN2 creation should fail
 
-        Capture error = EasyMock.newCapture();
+        Capture error = newCapture();
         putConnectorCallback.onCompletion(EasyMock.capture(error), EasyMock.>isNull());
         PowerMock.expectLastCall();
 
@@ -497,6 +667,7 @@ public void testCreateConnectorFailedCustomValidation() throws Exception {
     @Test
     public void testConnectorNameConflictsWithWorkerGroupId() throws Exception {
         EasyMock.expect(member.memberId()).andStubReturn("leader");
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V0);
         expectRebalance(1, Collections.emptyList(), Collections.emptyList());
         expectPostRebalanceCatchup(SNAPSHOT);
 
@@ -509,7 +680,7 @@ public void testConnectorNameConflictsWithWorkerGroupId() throws Exception {
         // config validation
         Connector connectorMock = PowerMock.createMock(SinkConnector.class);
         EasyMock.expect(worker.configTransformer()).andReturn(transformer).times(2);
-        final Capture> configCapture = EasyMock.newCapture();
+        final Capture> configCapture = newCapture();
         EasyMock.expect(transformer.transform(EasyMock.capture(configCapture))).andAnswer(configCapture::getValue);
         EasyMock.expect(worker.getPlugins()).andReturn(plugins).times(3);
         EasyMock.expect(plugins.compareAndSwapLoaders(connectorMock)).andReturn(delegatingLoader);
@@ -521,7 +692,7 @@ public void testConnectorNameConflictsWithWorkerGroupId() throws Exception {
         // CONN2 creation should fail because the worker group id (connect-test-group) conflicts with
         // the consumer group id we would use for this sink
 
-        Capture error = EasyMock.newCapture();
+        Capture error = newCapture();
         putConnectorCallback.onCompletion(EasyMock.capture(error), EasyMock.isNull(Herder.Created.class));
         PowerMock.expectLastCall();
 
@@ -546,8 +717,9 @@ public void testConnectorNameConflictsWithWorkerGroupId() throws Exception {
     @Test
     public void testCreateConnectorAlreadyExists() throws Exception {
         EasyMock.expect(member.memberId()).andStubReturn("leader");
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V0);
         EasyMock.expect(worker.configTransformer()).andReturn(transformer).times(2);
-        final Capture> configCapture = EasyMock.newCapture();
+        final Capture> configCapture = newCapture();
         EasyMock.expect(transformer.transform(EasyMock.capture(configCapture))).andAnswer(configCapture::getValue);
         EasyMock.expect(worker.getPlugins()).andReturn(plugins);
         EasyMock.expect(plugins.newConnector(EasyMock.anyString())).andReturn(null);
@@ -577,6 +749,7 @@ public void testCreateConnectorAlreadyExists() throws Exception {
     @Test
     public void testDestroyConnector() throws Exception {
         EasyMock.expect(member.memberId()).andStubReturn("leader");
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V0);
         // Start with one connector
         EasyMock.expect(worker.getPlugins()).andReturn(plugins);
         expectRebalance(1, Arrays.asList(CONN1), Collections.emptyList());
@@ -615,6 +788,7 @@ public void testRestartConnector() throws Exception {
 
         // get the initial assignment
         EasyMock.expect(member.memberId()).andStubReturn("leader");
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V0);
         EasyMock.expect(worker.getPlugins()).andReturn(plugins);
         expectRebalance(1, singletonList(CONN1), Collections.emptyList());
         expectPostRebalanceCatchup(SNAPSHOT);
@@ -656,6 +830,7 @@ public void testRestartConnector() throws Exception {
     public void testRestartUnknownConnector() throws Exception {
         // get the initial assignment
         EasyMock.expect(member.memberId()).andStubReturn("leader");
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V0);
         expectRebalance(1, Collections.emptyList(), Collections.emptyList());
         expectPostRebalanceCatchup(SNAPSHOT);
         member.poll(EasyMock.anyInt());
@@ -690,6 +865,7 @@ public void testRestartUnknownConnector() throws Exception {
     public void testRestartConnectorRedirectToLeader() throws Exception {
         // get the initial assignment
         EasyMock.expect(member.memberId()).andStubReturn("member");
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V0);
         expectRebalance(1, Collections.emptyList(), Collections.emptyList());
         expectPostRebalanceCatchup(SNAPSHOT);
         member.poll(EasyMock.anyInt());
@@ -724,6 +900,7 @@ public void testRestartConnectorRedirectToLeader() throws Exception {
     public void testRestartConnectorRedirectToOwner() throws Exception {
         // get the initial assignment
         EasyMock.expect(member.memberId()).andStubReturn("leader");
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V0);
         expectRebalance(1, Collections.emptyList(), Collections.emptyList());
         expectPostRebalanceCatchup(SNAPSHOT);
         member.poll(EasyMock.anyInt());
@@ -771,6 +948,7 @@ public void testRestartTask() throws Exception {
 
         // get the initial assignment
         EasyMock.expect(member.memberId()).andStubReturn("leader");
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V0);
         expectRebalance(1, Collections.emptyList(), singletonList(TASK0));
         expectPostRebalanceCatchup(SNAPSHOT);
         member.poll(EasyMock.anyInt());
@@ -808,6 +986,7 @@ public void testRestartTask() throws Exception {
     public void testRestartUnknownTask() throws Exception {
         // get the initial assignment
         EasyMock.expect(member.memberId()).andStubReturn("member");
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V0);
         expectRebalance(1, Collections.emptyList(), Collections.emptyList());
         expectPostRebalanceCatchup(SNAPSHOT);
         member.poll(EasyMock.anyInt());
@@ -854,6 +1033,7 @@ public void testRequestProcessingOrder() {
     public void testRestartTaskRedirectToLeader() throws Exception {
         // get the initial assignment
         EasyMock.expect(member.memberId()).andStubReturn("member");
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V0);
         expectRebalance(1, Collections.emptyList(), Collections.emptyList());
         expectPostRebalanceCatchup(SNAPSHOT);
         member.poll(EasyMock.anyInt());
@@ -888,6 +1068,7 @@ public void testRestartTaskRedirectToLeader() throws Exception {
     public void testRestartTaskRedirectToOwner() throws Exception {
         // get the initial assignment
         EasyMock.expect(member.memberId()).andStubReturn("leader");
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V0);
         expectRebalance(1, Collections.emptyList(), Collections.emptyList());
         expectPostRebalanceCatchup(SNAPSHOT);
         member.poll(EasyMock.anyInt());
@@ -926,6 +1107,7 @@ public void testRestartTaskRedirectToOwner() throws Exception {
     public void testConnectorConfigAdded() {
         // If a connector was added, we need to rebalance
         EasyMock.expect(member.memberId()).andStubReturn("member");
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V0);
 
         // join, no configs so no need to catch up on config topic
         expectRebalance(-1, Collections.emptyList(), Collections.emptyList());
@@ -943,6 +1125,7 @@ public void testConnectorConfigAdded() {
         // Performs rebalance and gets new assignment
         expectRebalance(Collections.emptyList(), Collections.emptyList(),
                 ConnectProtocol.Assignment.NO_ERROR, 1, Arrays.asList(CONN1), Collections.emptyList());
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V0);
         worker.startConnector(EasyMock.eq(CONN1), EasyMock.>anyObject(), EasyMock.anyObject(),
                 EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED));
         PowerMock.expectLastCall().andReturn(true);
@@ -967,11 +1150,13 @@ public void testConnectorConfigUpdate() throws Exception {
         // Connector config can be applied without any rebalance
 
         EasyMock.expect(member.memberId()).andStubReturn("member");
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V0);
         EasyMock.expect(worker.connectorNames()).andStubReturn(Collections.singleton(CONN1));
 
         // join
         expectRebalance(1, Arrays.asList(CONN1), Collections.emptyList());
         expectPostRebalanceCatchup(SNAPSHOT);
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V0);
         worker.startConnector(EasyMock.eq(CONN1), EasyMock.>anyObject(), EasyMock.anyObject(),
                 EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED));
         PowerMock.expectLastCall().andReturn(true);
@@ -988,6 +1173,7 @@ public void testConnectorConfigUpdate() throws Exception {
         EasyMock.expect(configBackingStore.snapshot()).andReturn(SNAPSHOT); // for this test, it doesn't matter if we use the same config snapshot
         worker.stopConnector(CONN1);
         PowerMock.expectLastCall().andReturn(true);
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V0);
         worker.startConnector(EasyMock.eq(CONN1), EasyMock.>anyObject(), EasyMock.anyObject(),
                 EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED));
         PowerMock.expectLastCall().andReturn(true);
@@ -1011,11 +1197,13 @@ public void testConnectorPaused() throws Exception {
         // ensure that target state changes are propagated to the worker
 
         EasyMock.expect(member.memberId()).andStubReturn("member");
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V0);
         EasyMock.expect(worker.connectorNames()).andStubReturn(Collections.singleton(CONN1));
 
         // join
         expectRebalance(1, Arrays.asList(CONN1), Collections.emptyList());
         expectPostRebalanceCatchup(SNAPSHOT);
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V0);
         worker.startConnector(EasyMock.eq(CONN1), EasyMock.>anyObject(), EasyMock.anyObject(),
                 EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED));
         PowerMock.expectLastCall().andReturn(true);
@@ -1051,11 +1239,13 @@ public void testConnectorPaused() throws Exception {
     @Test
     public void testConnectorResumed() throws Exception {
         EasyMock.expect(member.memberId()).andStubReturn("member");
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V0);
         EasyMock.expect(worker.connectorNames()).andStubReturn(Collections.singleton(CONN1));
 
         // start with the connector paused
         expectRebalance(1, Arrays.asList(CONN1), Collections.emptyList());
         expectPostRebalanceCatchup(SNAPSHOT_PAUSED_CONN1);
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V0);
         worker.startConnector(EasyMock.eq(CONN1), EasyMock.>anyObject(), EasyMock.anyObject(),
                 EasyMock.eq(herder), EasyMock.eq(TargetState.PAUSED));
         PowerMock.expectLastCall().andReturn(true);
@@ -1094,6 +1284,7 @@ public void testConnectorResumed() throws Exception {
     @Test
     public void testUnknownConnectorPaused() throws Exception {
         EasyMock.expect(member.memberId()).andStubReturn("member");
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V0);
         EasyMock.expect(worker.connectorNames()).andStubReturn(Collections.singleton(CONN1));
 
         // join
@@ -1131,6 +1322,7 @@ public void testConnectorPausedRunningTaskOnly() throws Exception {
         // changes to the worker so that tasks will transition correctly
 
         EasyMock.expect(member.memberId()).andStubReturn("member");
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V0);
         EasyMock.expect(worker.connectorNames()).andStubReturn(Collections.emptySet());
 
         // join
@@ -1171,6 +1363,7 @@ public void testConnectorResumedRunningTaskOnly() throws Exception {
         // changes to the worker so that tasks will transition correctly
 
         EasyMock.expect(member.memberId()).andStubReturn("member");
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V0);
         EasyMock.expect(worker.connectorNames()).andStubReturn(Collections.emptySet());
 
         // join
@@ -1211,6 +1404,7 @@ public void testConnectorResumedRunningTaskOnly() throws Exception {
     public void testTaskConfigAdded() {
         // Task config always requires rebalance
         EasyMock.expect(member.memberId()).andStubReturn("member");
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V0);
 
         // join
         expectRebalance(-1, Collections.emptyList(), Collections.emptyList());
@@ -1249,21 +1443,22 @@ public void testTaskConfigAdded() {
     public void testJoinLeaderCatchUpFails() throws Exception {
         // Join group and as leader fail to do assignment
         EasyMock.expect(member.memberId()).andStubReturn("leader");
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V0);
         expectRebalance(Collections.emptyList(), Collections.emptyList(),
                 ConnectProtocol.Assignment.CONFIG_MISMATCH, 1, Collections.emptyList(),
                 Collections.emptyList());
         // Reading to end of log times out
         configBackingStore.refresh(EasyMock.anyLong(), EasyMock.anyObject(TimeUnit.class));
         EasyMock.expectLastCall().andThrow(new TimeoutException());
-        member.maybeLeaveGroup();
+        member.maybeLeaveGroup(EasyMock.eq("taking too long to read the log"));
         EasyMock.expectLastCall();
-        PowerMock.expectPrivate(herder, "backoff", DistributedConfig.WORKER_UNSYNC_BACKOFF_MS_DEFAULT);
         member.requestRejoin();
 
         // After backoff, restart the process and this time succeed
         expectRebalance(1, Arrays.asList(CONN1), Arrays.asList(TASK1));
         expectPostRebalanceCatchup(SNAPSHOT);
 
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V0);
         worker.startConnector(EasyMock.eq(CONN1), EasyMock.>anyObject(), EasyMock.anyObject(),
                 EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED));
         PowerMock.expectLastCall().andReturn(true);
@@ -1278,14 +1473,198 @@ public void testJoinLeaderCatchUpFails() throws Exception {
 
         PowerMock.replayAll();
 
+        long before = time.milliseconds();
+        int workerUnsyncBackoffMs = DistributedConfig.WORKER_UNSYNC_BACKOFF_MS_DEFAULT;
+        int coordinatorDiscoveryTimeoutMs = 100;
         herder.tick();
+        assertEquals(before + coordinatorDiscoveryTimeoutMs + workerUnsyncBackoffMs, time.milliseconds());
+
         time.sleep(1000L);
         assertStatistics("leaderUrl", true, 3, 0, Double.NEGATIVE_INFINITY, Double.POSITIVE_INFINITY);
 
+        before = time.milliseconds();
         herder.tick();
+        assertEquals(before + coordinatorDiscoveryTimeoutMs, time.milliseconds());
         time.sleep(2000L);
         assertStatistics("leaderUrl", false, 3, 1, 100, 2000L);
 
+        PowerMock.verifyAll();
+    }
+
+    @Test
+    public void testJoinLeaderCatchUpRetriesForIncrementalCooperative() throws Exception {
+        connectProtocolVersion = CONNECT_PROTOCOL_V1;
+
+        // Join group and as leader fail to do assignment
+        EasyMock.expect(member.memberId()).andStubReturn("leader");
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V1);
+        expectRebalance(1, Arrays.asList(CONN1), Arrays.asList(TASK1));
+        expectPostRebalanceCatchup(SNAPSHOT);
+
+        member.poll(EasyMock.anyInt());
+        PowerMock.expectLastCall();
+
+        // The leader got its assignment
+        expectRebalance(Collections.emptyList(), Collections.emptyList(),
+                ConnectProtocol.Assignment.NO_ERROR,
+                1, Arrays.asList(CONN1), Arrays.asList(TASK1), 0);
+
+        EasyMock.expect(worker.getPlugins()).andReturn(plugins);
+        // and the new assignment started
+        worker.startConnector(EasyMock.eq(CONN1), EasyMock.anyObject(), EasyMock.anyObject(),
+                EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED));
+        PowerMock.expectLastCall().andReturn(true);
+        EasyMock.expect(worker.isRunning(CONN1)).andReturn(true);
+        EasyMock.expect(worker.connectorTaskConfigs(CONN1, conn1SinkConfig)).andReturn(TASK_CONFIGS);
+
+        worker.startTask(EasyMock.eq(TASK1), EasyMock.anyObject(), EasyMock.anyObject(), EasyMock.anyObject(),
+                EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED));
+        PowerMock.expectLastCall().andReturn(true);
+        member.poll(EasyMock.anyInt());
+        PowerMock.expectLastCall();
+
+        // Another rebalance is triggered but this time it fails to read to the max offset and
+        // triggers a re-sync
+        expectRebalance(Collections.emptyList(), Collections.emptyList(),
+                ConnectProtocol.Assignment.CONFIG_MISMATCH, 1, Collections.emptyList(),
+                Collections.emptyList());
+
+        // The leader will retry a few times to read to the end of the config log
+        int retries = 2;
+        member.requestRejoin();
+        for (int i = retries; i >= 0; --i) {
+            // Reading to end of log times out
+            configBackingStore.refresh(EasyMock.anyLong(), EasyMock.anyObject(TimeUnit.class));
+            EasyMock.expectLastCall().andThrow(new TimeoutException());
+            member.maybeLeaveGroup(EasyMock.eq("taking too long to read the log"));
+            EasyMock.expectLastCall();
+        }
+
+        // After a few retries succeed to read the log to the end
+        expectRebalance(Collections.emptyList(), Collections.emptyList(),
+                ConnectProtocol.Assignment.NO_ERROR,
+                1, Arrays.asList(CONN1), Arrays.asList(TASK1), 0);
+        expectPostRebalanceCatchup(SNAPSHOT);
+        member.poll(EasyMock.anyInt());
+        PowerMock.expectLastCall();
+
+        PowerMock.replayAll();
+
+        assertStatistics(0, 0, 0, Double.POSITIVE_INFINITY);
+        herder.tick();
+
+        time.sleep(2000L);
+        assertStatistics(3, 1, 100, 2000);
+        herder.tick();
+
+        long before;
+        int coordinatorDiscoveryTimeoutMs = 100;
+        int maxRetries = 5;
+        for (int i = maxRetries; i >= maxRetries - retries; --i) {
+            before = time.milliseconds();
+            int workerUnsyncBackoffMs =
+                    DistributedConfig.SCHEDULED_REBALANCE_MAX_DELAY_MS_DEFAULT / 10 / i;
+            herder.tick();
+            assertEquals(before + coordinatorDiscoveryTimeoutMs + workerUnsyncBackoffMs, time.milliseconds());
+            coordinatorDiscoveryTimeoutMs = 0;
+        }
+
+        before = time.milliseconds();
+        coordinatorDiscoveryTimeoutMs = 100;
+        herder.tick();
+        assertEquals(before + coordinatorDiscoveryTimeoutMs, time.milliseconds());
+
+        PowerMock.verifyAll();
+    }
+
+    @Test
+    public void testJoinLeaderCatchUpFailsForIncrementalCooperative() throws Exception {
+        connectProtocolVersion = CONNECT_PROTOCOL_V1;
+
+        // Join group and as leader fail to do assignment
+        EasyMock.expect(member.memberId()).andStubReturn("leader");
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V1);
+        expectRebalance(1, Arrays.asList(CONN1), Arrays.asList(TASK1));
+        expectPostRebalanceCatchup(SNAPSHOT);
+
+        member.poll(EasyMock.anyInt());
+        PowerMock.expectLastCall();
+
+        // The leader got its assignment
+        expectRebalance(Collections.emptyList(), Collections.emptyList(),
+                ConnectProtocol.Assignment.NO_ERROR,
+                1, Arrays.asList(CONN1), Arrays.asList(TASK1), 0);
+
+        EasyMock.expect(worker.getPlugins()).andReturn(plugins);
+        // and the new assignment started
+        worker.startConnector(EasyMock.eq(CONN1), EasyMock.anyObject(), EasyMock.anyObject(),
+                EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED));
+        PowerMock.expectLastCall().andReturn(true);
+        EasyMock.expect(worker.isRunning(CONN1)).andReturn(true);
+        EasyMock.expect(worker.connectorTaskConfigs(CONN1, conn1SinkConfig)).andReturn(TASK_CONFIGS);
+
+        worker.startTask(EasyMock.eq(TASK1), EasyMock.anyObject(), EasyMock.anyObject(), EasyMock.anyObject(),
+                EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED));
+        PowerMock.expectLastCall().andReturn(true);
+        member.poll(EasyMock.anyInt());
+        PowerMock.expectLastCall();
+
+        // Another rebalance is triggered but this time it fails to read to the max offset and
+        // triggers a re-sync
+        expectRebalance(Collections.emptyList(), Collections.emptyList(),
+                ConnectProtocol.Assignment.CONFIG_MISMATCH, 1, Collections.emptyList(),
+                Collections.emptyList());
+
+        // The leader will exhaust the retries while trying to read to the end of the config log
+        int maxRetries = 5;
+        member.requestRejoin();
+        for (int i = maxRetries; i >= 0; --i) {
+            // Reading to end of log times out
+            configBackingStore.refresh(EasyMock.anyLong(), EasyMock.anyObject(TimeUnit.class));
+            EasyMock.expectLastCall().andThrow(new TimeoutException());
+            member.maybeLeaveGroup(EasyMock.eq("taking too long to read the log"));
+            EasyMock.expectLastCall();
+        }
+
+        Capture assignmentCapture = newCapture();
+        member.revokeAssignment(capture(assignmentCapture));
+        PowerMock.expectLastCall();
+
+        // After a complete backoff and a revocation of running tasks rejoin and this time succeed
+        // The worker gets back the assignment that had given up
+        expectRebalance(Collections.emptyList(), Collections.emptyList(),
+                ConnectProtocol.Assignment.NO_ERROR,
+                1, Arrays.asList(CONN1), Arrays.asList(TASK1), 0);
+        expectPostRebalanceCatchup(SNAPSHOT);
+        member.poll(EasyMock.anyInt());
+        PowerMock.expectLastCall();
+
+        PowerMock.replayAll();
+
+        assertStatistics(0, 0, 0, Double.POSITIVE_INFINITY);
+        herder.tick();
+
+        time.sleep(2000L);
+        assertStatistics(3, 1, 100, 2000);
+        herder.tick();
+
+        long before;
+        int coordinatorDiscoveryTimeoutMs = 100;
+        for (int i = maxRetries; i > 0; --i) {
+            before = time.milliseconds();
+            int workerUnsyncBackoffMs =
+                    DistributedConfig.SCHEDULED_REBALANCE_MAX_DELAY_MS_DEFAULT / 10 / i;
+            herder.tick();
+            assertEquals(before + coordinatorDiscoveryTimeoutMs + workerUnsyncBackoffMs, time.milliseconds());
+            coordinatorDiscoveryTimeoutMs = 0;
+        }
+
+        before = time.milliseconds();
+        herder.tick();
+        assertEquals(before, time.milliseconds());
+        assertEquals(Collections.singleton(CONN1), assignmentCapture.getValue().connectors());
+        assertEquals(Collections.singleton(TASK1), assignmentCapture.getValue().tasks());
+        herder.tick();
 
         PowerMock.verifyAll();
     }
@@ -1293,14 +1672,16 @@ public void testJoinLeaderCatchUpFails() throws Exception {
     @Test
     public void testAccessors() throws Exception {
         EasyMock.expect(member.memberId()).andStubReturn("leader");
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V0);
         EasyMock.expect(worker.getPlugins()).andReturn(plugins).anyTimes();
         expectRebalance(1, Collections.emptyList(), Collections.emptyList());
+        EasyMock.expect(configBackingStore.snapshot()).andReturn(SNAPSHOT).times(2);
 
         WorkerConfigTransformer configTransformer = EasyMock.mock(WorkerConfigTransformer.class);
         EasyMock.expect(configTransformer.transform(EasyMock.eq(CONN1), EasyMock.anyObject()))
             .andThrow(new AssertionError("Config transformation should not occur when requesting connector or task info"));
         EasyMock.replay(configTransformer);
-        ClusterConfigState snapshotWithTransform = new ClusterConfigState(1, Collections.singletonMap(CONN1, 3),
+        ClusterConfigState snapshotWithTransform = new ClusterConfigState(1, null, Collections.singletonMap(CONN1, 3),
             Collections.singletonMap(CONN1, CONN1_CONFIG), Collections.singletonMap(CONN1, TargetState.STARTED),
             TASK_CONFIGS_MAP, Collections.emptySet(), configTransformer);
 
@@ -1314,6 +1695,7 @@ public void testAccessors() throws Exception {
         PowerMock.expectLastCall();
 
 
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V0);
         PowerMock.replayAll();
 
         FutureCallback> listConnectorsCb = new FutureCallback<>();
@@ -1349,6 +1731,7 @@ public void testPutConnectorConfig() throws Exception {
         EasyMock.expect(member.memberId()).andStubReturn("leader");
         expectRebalance(1, Arrays.asList(CONN1), Collections.emptyList());
         expectPostRebalanceCatchup(SNAPSHOT);
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V0);
         worker.startConnector(EasyMock.eq(CONN1), EasyMock.>anyObject(), EasyMock.anyObject(),
                 EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED));
         PowerMock.expectLastCall().andReturn(true);
@@ -1368,7 +1751,7 @@ public void testPutConnectorConfig() throws Exception {
         // config validation
         Connector connectorMock = PowerMock.createMock(SourceConnector.class);
         EasyMock.expect(worker.configTransformer()).andReturn(transformer).times(2);
-        final Capture> configCapture = EasyMock.newCapture();
+        final Capture> configCapture = newCapture();
         EasyMock.expect(transformer.transform(EasyMock.capture(configCapture))).andAnswer(configCapture::getValue);
         EasyMock.expect(worker.getPlugins()).andReturn(plugins).anyTimes();
         EasyMock.expect(plugins.compareAndSwapLoaders(connectorMock)).andReturn(delegatingLoader);
@@ -1376,6 +1759,7 @@ public void testPutConnectorConfig() throws Exception {
         EasyMock.expect(connectorMock.config()).andReturn(new ConfigDef());
         EasyMock.expect(connectorMock.validate(CONN1_CONFIG_UPDATED)).andReturn(new Config(Collections.emptyList()));
         EasyMock.expect(Plugins.compareAndSwapLoaders(delegatingLoader)).andReturn(pluginLoader);
+        EasyMock.expect(configBackingStore.snapshot()).andReturn(SNAPSHOT);
 
         configBackingStore.putConnectorConfig(CONN1, CONN1_CONFIG_UPDATED);
         PowerMock.expectLastCall().andAnswer(new IAnswer() {
@@ -1388,9 +1772,10 @@ public Object answer() throws Throwable {
         });
         // As a result of reconfig, should need to update snapshot. With only connector updates, we'll just restart
         // connector without rebalance
-        EasyMock.expect(configBackingStore.snapshot()).andReturn(SNAPSHOT_UPDATED_CONN1_CONFIG);
+        EasyMock.expect(configBackingStore.snapshot()).andReturn(SNAPSHOT_UPDATED_CONN1_CONFIG).times(2);
         worker.stopConnector(CONN1);
         PowerMock.expectLastCall().andReturn(true);
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V0);
         worker.startConnector(EasyMock.eq(CONN1), EasyMock.>anyObject(), EasyMock.anyObject(),
                 EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED));
         PowerMock.expectLastCall().andReturn(true);
@@ -1433,16 +1818,143 @@ public Object answer() throws Throwable {
         PowerMock.verifyAll();
     }
 
+
+    @Test
+    public void testPutTaskConfigsSignatureNotRequiredV0() {
+        Callback taskConfigCb = EasyMock.mock(Callback.class);
+
+        member.wakeup();
+        EasyMock.expectLastCall().once();
+        EasyMock.expect(member.currentProtocolVersion()).andReturn(CONNECT_PROTOCOL_V0).anyTimes();
+        PowerMock.replayAll(taskConfigCb);
+
+        herder.putTaskConfigs(CONN1, TASK_CONFIGS, taskConfigCb, null);
+
+        PowerMock.verifyAll();
+    }
+    @Test
+    public void testPutTaskConfigsSignatureNotRequiredV1() {
+        Callback taskConfigCb = EasyMock.mock(Callback.class);
+
+        member.wakeup();
+        EasyMock.expectLastCall().once();
+        EasyMock.expect(member.currentProtocolVersion()).andReturn(CONNECT_PROTOCOL_V1).anyTimes();
+        PowerMock.replayAll(taskConfigCb);
+
+        herder.putTaskConfigs(CONN1, TASK_CONFIGS, taskConfigCb, null);
+
+        PowerMock.verifyAll();
+    }
+
+    @Test
+    public void testPutTaskConfigsMissingRequiredSignature() {
+        Callback taskConfigCb = EasyMock.mock(Callback.class);
+        Capture errorCapture = Capture.newInstance();
+        taskConfigCb.onCompletion(EasyMock.capture(errorCapture), EasyMock.eq(null));
+        EasyMock.expectLastCall().once();
+
+        EasyMock.expect(member.currentProtocolVersion()).andReturn(CONNECT_PROTOCOL_V2).anyTimes();
+        PowerMock.replayAll(taskConfigCb);
+
+        herder.putTaskConfigs(CONN1, TASK_CONFIGS, taskConfigCb, null);
+
+        PowerMock.verifyAll();
+        assertTrue(errorCapture.getValue() instanceof BadRequestException);
+    }
+
+    @Test
+    public void testPutTaskConfigsDisallowedSignatureAlgorithm() {
+        Callback taskConfigCb = EasyMock.mock(Callback.class);
+        Capture errorCapture = Capture.newInstance();
+        taskConfigCb.onCompletion(EasyMock.capture(errorCapture), EasyMock.eq(null));
+        EasyMock.expectLastCall().once();
+
+        EasyMock.expect(member.currentProtocolVersion()).andReturn(CONNECT_PROTOCOL_V2).anyTimes();
+
+        InternalRequestSignature signature = EasyMock.mock(InternalRequestSignature.class);
+        EasyMock.expect(signature.keyAlgorithm()).andReturn("HmacSHA489").anyTimes();
+
+        PowerMock.replayAll(taskConfigCb, signature);
+
+        herder.putTaskConfigs(CONN1, TASK_CONFIGS, taskConfigCb, signature);
+
+        PowerMock.verifyAll();
+        assertTrue(errorCapture.getValue() instanceof BadRequestException);
+    }
+
+    @Test
+    public void testPutTaskConfigsInvalidSignature() {
+        Callback taskConfigCb = EasyMock.mock(Callback.class);
+        Capture errorCapture = Capture.newInstance();
+        taskConfigCb.onCompletion(EasyMock.capture(errorCapture), EasyMock.eq(null));
+        EasyMock.expectLastCall().once();
+
+        EasyMock.expect(member.currentProtocolVersion()).andReturn(CONNECT_PROTOCOL_V2).anyTimes();
+
+        InternalRequestSignature signature = EasyMock.mock(InternalRequestSignature.class);
+        EasyMock.expect(signature.keyAlgorithm()).andReturn("HmacSHA256").anyTimes();
+        EasyMock.expect(signature.isValid(EasyMock.anyObject())).andReturn(false).anyTimes();
+
+        PowerMock.replayAll(taskConfigCb, signature);
+
+        herder.putTaskConfigs(CONN1, TASK_CONFIGS, taskConfigCb, signature);
+
+        PowerMock.verifyAll();
+        assertTrue(errorCapture.getValue() instanceof ConnectRestException);
+        assertEquals(FORBIDDEN.getStatusCode(), ((ConnectRestException) errorCapture.getValue()).statusCode());
+    }
+
+    @Test
+    public void testPutTaskConfigsValidRequiredSignature() {
+        Callback taskConfigCb = EasyMock.mock(Callback.class);
+
+        member.wakeup();
+        EasyMock.expectLastCall().once();
+        EasyMock.expect(member.currentProtocolVersion()).andReturn(CONNECT_PROTOCOL_V2).anyTimes();
+
+        InternalRequestSignature signature = EasyMock.mock(InternalRequestSignature.class);
+        EasyMock.expect(signature.keyAlgorithm()).andReturn("HmacSHA256").anyTimes();
+        EasyMock.expect(signature.isValid(EasyMock.anyObject())).andReturn(true).anyTimes();
+
+        PowerMock.replayAll(taskConfigCb, signature);
+
+        herder.putTaskConfigs(CONN1, TASK_CONFIGS, taskConfigCb, signature);
+
+        PowerMock.verifyAll();
+    }
+
+    @Test
+    public void testKeyExceptionDetection() {
+        assertFalse(herder.isPossibleExpiredKeyException(
+            time.milliseconds(),
+            new RuntimeException()
+        ));
+        assertFalse(herder.isPossibleExpiredKeyException(
+            time.milliseconds(),
+            new BadRequestException("")
+        ));
+        assertFalse(herder.isPossibleExpiredKeyException(
+            time.milliseconds() - TimeUnit.MINUTES.toMillis(2),
+            new ConnectRestException(FORBIDDEN.getStatusCode(), "")
+        ));
+        assertTrue(herder.isPossibleExpiredKeyException(
+            time.milliseconds(),
+            new ConnectRestException(FORBIDDEN.getStatusCode(), "")
+        ));
+    }
+
     @Test
     public void testInconsistentConfigs() {
         // FIXME: if we have inconsistent configs, we need to request forced reconfig + write of the connector's task configs
         // This requires inter-worker communication, so needs the REST API
     }
 
+
     private void expectRebalance(final long offset,
                                  final List assignedConnectors,
                                  final List assignedTasks) {
-        expectRebalance(null, null, ConnectProtocol.Assignment.NO_ERROR, offset, assignedConnectors, assignedTasks);
+        expectRebalance(Collections.emptyList(), Collections.emptyList(),
+                ConnectProtocol.Assignment.NO_ERROR, offset, assignedConnectors, assignedTasks, 0);
     }
 
     // Handles common initial part of rebalance callback. Does not handle instantiation of connectors and tasks.
@@ -1452,33 +1964,56 @@ private void expectRebalance(final Collection revokedConnectors,
                                  final long offset,
                                  final List assignedConnectors,
                                  final List assignedTasks) {
+        expectRebalance(revokedConnectors, revokedTasks, error, offset, assignedConnectors, assignedTasks, 0);
+    }
+
+    // Handles common initial part of rebalance callback. Does not handle instantiation of connectors and tasks.
+    private void expectRebalance(final Collection revokedConnectors,
+                                 final List revokedTasks,
+                                 final short error,
+                                 final long offset,
+                                 final List assignedConnectors,
+                                 final List assignedTasks,
+                                 int delay) {
         member.ensureActive();
         PowerMock.expectLastCall().andAnswer(new IAnswer() {
             @Override
             public Object answer() throws Throwable {
-                if (revokedConnectors != null)
+                ExtendedAssignment assignment;
+                if (!revokedConnectors.isEmpty() || !revokedTasks.isEmpty()) {
                     rebalanceListener.onRevoked("leader", revokedConnectors, revokedTasks);
-                ConnectProtocol.Assignment assignment = new ConnectProtocol.Assignment(
-                        error, "leader", "leaderUrl", offset, assignedConnectors, assignedTasks);
+                }
+
+                if (connectProtocolVersion == CONNECT_PROTOCOL_V0) {
+                    assignment = new ExtendedAssignment(
+                            connectProtocolVersion, error, "leader", "leaderUrl", offset,
+                            assignedConnectors, assignedTasks,
+                            Collections.emptyList(), Collections.emptyList(), 0);
+                } else {
+                    assignment = new ExtendedAssignment(
+                            connectProtocolVersion, error, "leader", "leaderUrl", offset,
+                            assignedConnectors, assignedTasks,
+                            new ArrayList<>(revokedConnectors), new ArrayList<>(revokedTasks), delay);
+                }
                 rebalanceListener.onAssigned(assignment, 3);
                 time.sleep(100L);
                 return null;
             }
         });
 
-        if (revokedConnectors != null) {
+        if (!revokedConnectors.isEmpty()) {
             for (String connector : revokedConnectors) {
                 worker.stopConnector(connector);
                 PowerMock.expectLastCall().andReturn(true);
             }
         }
 
-        if (revokedTasks != null && !revokedTasks.isEmpty()) {
+        if (!revokedTasks.isEmpty()) {
             worker.stopAndAwaitTask(EasyMock.anyObject(ConnectorTaskId.class));
             PowerMock.expectLastCall();
         }
 
-        if (revokedConnectors != null) {
+        if (!revokedConnectors.isEmpty()) {
             statusBackingStore.flush();
             PowerMock.expectLastCall();
         }
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/distributed/IncrementalCooperativeAssignorTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/distributed/IncrementalCooperativeAssignorTest.java
new file mode 100644
index 0000000000000..684da46bdadb1
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/distributed/IncrementalCooperativeAssignorTest.java
@@ -0,0 +1,1463 @@
+/*
+ * 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.distributed;
+
+import org.apache.kafka.clients.consumer.internals.RequestFuture;
+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.connect.runtime.TargetState;
+import org.apache.kafka.connect.runtime.distributed.WorkerCoordinator.ConnectorsAndTasks;
+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.mockito.ArgumentCaptor;
+import org.mockito.Captor;
+import org.mockito.Mock;
+import org.mockito.Mockito;
+import org.mockito.junit.MockitoJUnit;
+import org.mockito.junit.MockitoRule;
+
+import java.util.AbstractMap.SimpleEntry;
+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.Set;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import static org.apache.kafka.connect.runtime.distributed.IncrementalCooperativeConnectProtocol.CONNECT_PROTOCOL_V1;
+import static org.apache.kafka.connect.runtime.distributed.IncrementalCooperativeConnectProtocol.CONNECT_PROTOCOL_V2;
+import static org.apache.kafka.connect.runtime.distributed.WorkerCoordinator.WorkerLoad;
+import static org.hamcrest.CoreMatchers.hasItems;
+import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+import static org.junit.runners.Parameterized.Parameter;
+import static org.junit.runners.Parameterized.Parameters;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
+import static org.mockito.Mockito.when;
+
+public class IncrementalCooperativeAssignorTest {
+    @Rule
+    public MockitoRule rule = MockitoJUnit.rule();
+
+    @Mock
+    private WorkerCoordinator coordinator;
+
+    @Captor
+    ArgumentCaptor> assignmentsCapture;
+
+    @Parameters
+    public static Iterable mode() {
+        return Arrays.asList(new Object[][] {{CONNECT_PROTOCOL_V1, CONNECT_PROTOCOL_V2}});
+    }
+
+    @Parameter
+    public short protocolVersion;
+
+    private ClusterConfigState configState;
+    private Map memberConfigs;
+    private Map expectedMemberConfigs;
+    private long offset;
+    private String leader;
+    private String leaderUrl;
+    private Time time;
+    private int rebalanceDelay;
+    private IncrementalCooperativeAssignor assignor;
+    private int rebalanceNum;
+    Map assignments;
+    Map returnedAssignments;
+
+    @Before
+    public void setup() {
+        leader = "worker1";
+        leaderUrl = expectedLeaderUrl(leader);
+        offset = 10;
+        configState = clusterConfigState(offset, 2, 4);
+        memberConfigs = memberConfigs(leader, offset, 1, 1);
+        time = Time.SYSTEM;
+        rebalanceDelay = DistributedConfig.SCHEDULED_REBALANCE_MAX_DELAY_MS_DEFAULT;
+        assignments = new HashMap<>();
+        initAssignor();
+    }
+
+    @After
+    public void teardown() {
+        verifyNoMoreInteractions(coordinator);
+    }
+
+    public void initAssignor() {
+        assignor = Mockito.spy(new IncrementalCooperativeAssignor(
+                new LogContext(),
+                time,
+                rebalanceDelay));
+        assignor.previousGenerationId = 1000;
+    }
+
+    @Test
+    public void testTaskAssignmentWhenWorkerJoins() {
+        when(coordinator.configSnapshot()).thenReturn(configState);
+        doReturn(Collections.EMPTY_MAP).when(assignor).serializeAssignments(assignmentsCapture.capture());
+
+        // First assignment with 1 worker and 2 connectors configured but not yet assigned
+        expectGeneration();
+        assignor.performTaskAssignment(leader, offset, memberConfigs, coordinator, protocolVersion);
+        ++rebalanceNum;
+        returnedAssignments = assignmentsCapture.getValue();
+        assertDelay(0, returnedAssignments);
+        expectedMemberConfigs = memberConfigs(leader, offset, returnedAssignments);
+        assertNoReassignments(memberConfigs, expectedMemberConfigs);
+        assertAssignment(2, 8, 0, 0, "worker1");
+
+        // Second assignment with a second worker joining and all connectors running on previous worker
+        applyAssignments(returnedAssignments);
+        memberConfigs = memberConfigs(leader, offset, assignments);
+        memberConfigs.put("worker2", new ExtendedWorkerState(leaderUrl, offset, null));
+        expectGeneration();
+        assignor.performTaskAssignment(leader, offset, memberConfigs, coordinator, protocolVersion);
+        ++rebalanceNum;
+        returnedAssignments = assignmentsCapture.getValue();
+        assertDelay(0, returnedAssignments);
+        expectedMemberConfigs = memberConfigs(leader, offset, returnedAssignments);
+        assertNoReassignments(memberConfigs, expectedMemberConfigs);
+        assertAssignment(0, 0, 1, 4, "worker1", "worker2");
+
+        // Third assignment after revocations
+        applyAssignments(returnedAssignments);
+        memberConfigs = memberConfigs(leader, offset, assignments);
+        expectGeneration();
+        assignor.performTaskAssignment(leader, offset, memberConfigs, coordinator, protocolVersion);
+        ++rebalanceNum;
+        returnedAssignments = assignmentsCapture.getValue();
+        assertDelay(0, returnedAssignments);
+        expectedMemberConfigs = memberConfigs(leader, offset, returnedAssignments);
+        assertNoReassignments(memberConfigs, expectedMemberConfigs);
+        assertAssignment(1, 4, 0, 0, "worker1", "worker2");
+
+        // A fourth rebalance should not change assignments
+        applyAssignments(returnedAssignments);
+        memberConfigs = memberConfigs(leader, offset, assignments);
+        expectGeneration();
+        assignor.performTaskAssignment(leader, offset, memberConfigs, coordinator, protocolVersion);
+        ++rebalanceNum;
+        returnedAssignments = assignmentsCapture.getValue();
+        assertDelay(0, returnedAssignments);
+        expectedMemberConfigs = memberConfigs(leader, offset, returnedAssignments);
+        assertNoReassignments(memberConfigs, expectedMemberConfigs);
+        assertAssignment(0, 0, 0, 0, "worker1", "worker2");
+
+        verify(coordinator, times(rebalanceNum)).configSnapshot();
+        verify(coordinator, times(rebalanceNum)).leaderState(any());
+        verify(coordinator, times(2 * rebalanceNum)).generationId();
+        verify(coordinator, times(rebalanceNum)).memberId();
+        verify(coordinator, times(rebalanceNum)).lastCompletedGenerationId();
+    }
+
+    @Test
+    public void testTaskAssignmentWhenWorkerLeavesPermanently() {
+        // Customize assignor for this test case
+        time = new MockTime();
+        initAssignor();
+
+        when(coordinator.configSnapshot()).thenReturn(configState);
+        doReturn(Collections.EMPTY_MAP).when(assignor).serializeAssignments(assignmentsCapture.capture());
+
+        // First assignment with 2 workers and 2 connectors configured but not yet assigned
+        memberConfigs.put("worker2", new ExtendedWorkerState(leaderUrl, offset, null));
+        expectGeneration();
+        assignor.performTaskAssignment(leader, offset, memberConfigs, coordinator, protocolVersion);
+        ++rebalanceNum;
+        returnedAssignments = assignmentsCapture.getValue();
+        assertDelay(0, returnedAssignments);
+        expectedMemberConfigs = memberConfigs(leader, offset, returnedAssignments);
+        assertNoReassignments(memberConfigs, expectedMemberConfigs);
+        assertAssignment(2, 8, 0, 0, "worker1", "worker2");
+
+        // Second assignment with only one worker remaining in the group. The worker that left the
+        // group was a follower. No re-assignments take place immediately and the count
+        // down for the rebalance delay starts
+        applyAssignments(returnedAssignments);
+        assignments.remove("worker2");
+        memberConfigs = memberConfigs(leader, offset, assignments);
+        expectGeneration();
+        assignor.performTaskAssignment(leader, offset, memberConfigs, coordinator, protocolVersion);
+        ++rebalanceNum;
+        returnedAssignments = assignmentsCapture.getValue();
+        assertDelay(rebalanceDelay, returnedAssignments);
+        expectedMemberConfigs = memberConfigs(leader, offset, returnedAssignments);
+        assertNoReassignments(memberConfigs, expectedMemberConfigs);
+        assertAssignment(0, 0, 0, 0, "worker1");
+
+        time.sleep(rebalanceDelay / 2);
+
+        // Third (incidental) assignment with still only one worker in the group. Max delay has not
+        // been reached yet
+        applyAssignments(returnedAssignments);
+        memberConfigs = memberConfigs(leader, offset, assignments);
+        expectGeneration();
+        assignor.performTaskAssignment(leader, offset, memberConfigs, coordinator, protocolVersion);
+        ++rebalanceNum;
+        returnedAssignments = assignmentsCapture.getValue();
+        assertDelay(rebalanceDelay / 2, returnedAssignments);
+        expectedMemberConfigs = memberConfigs(leader, offset, returnedAssignments);
+        assertNoReassignments(memberConfigs, expectedMemberConfigs);
+        assertAssignment(0, 0, 0, 0, "worker1");
+
+        time.sleep(rebalanceDelay / 2 + 1);
+
+        // Fourth assignment after delay expired
+        applyAssignments(returnedAssignments);
+        memberConfigs = memberConfigs(leader, offset, assignments);
+        expectGeneration();
+        assignor.performTaskAssignment(leader, offset, memberConfigs, coordinator, protocolVersion);
+        ++rebalanceNum;
+        returnedAssignments = assignmentsCapture.getValue();
+        assertDelay(0, returnedAssignments);
+        expectedMemberConfigs = memberConfigs(leader, offset, returnedAssignments);
+        assertNoReassignments(memberConfigs, expectedMemberConfigs);
+        assertAssignment(1, 4, 0, 0, "worker1");
+
+        verify(coordinator, times(rebalanceNum)).configSnapshot();
+        verify(coordinator, times(rebalanceNum)).leaderState(any());
+        verify(coordinator, times(2 * rebalanceNum)).generationId();
+        verify(coordinator, times(rebalanceNum)).memberId();
+        verify(coordinator, times(rebalanceNum)).lastCompletedGenerationId();
+    }
+
+    @Test
+    public void testTaskAssignmentWhenWorkerBounces() {
+        // Customize assignor for this test case
+        time = new MockTime();
+        initAssignor();
+
+        when(coordinator.configSnapshot()).thenReturn(configState);
+        doReturn(Collections.EMPTY_MAP).when(assignor).serializeAssignments(assignmentsCapture.capture());
+
+        // First assignment with 2 workers and 2 connectors configured but not yet assigned
+        memberConfigs.put("worker2", new ExtendedWorkerState(leaderUrl, offset, null));
+        expectGeneration();
+        assignor.performTaskAssignment(leader, offset, memberConfigs, coordinator, protocolVersion);
+        ++rebalanceNum;
+        returnedAssignments = assignmentsCapture.getValue();
+        assertDelay(0, returnedAssignments);
+        expectedMemberConfigs = memberConfigs(leader, offset, returnedAssignments);
+        assertNoReassignments(memberConfigs, expectedMemberConfigs);
+        assertAssignment(2, 8, 0, 0, "worker1", "worker2");
+
+        // Second assignment with only one worker remaining in the group. The worker that left the
+        // group was a follower. No re-assignments take place immediately and the count
+        // down for the rebalance delay starts
+        applyAssignments(returnedAssignments);
+        assignments.remove("worker2");
+        memberConfigs = memberConfigs(leader, offset, assignments);
+        expectGeneration();
+        assignor.performTaskAssignment(leader, offset, memberConfigs, coordinator, protocolVersion);
+        ++rebalanceNum;
+        returnedAssignments = assignmentsCapture.getValue();
+        assertDelay(rebalanceDelay, returnedAssignments);
+        expectedMemberConfigs = memberConfigs(leader, offset, returnedAssignments);
+        assertNoReassignments(memberConfigs, expectedMemberConfigs);
+        assertAssignment(0, 0, 0, 0, "worker1");
+
+        time.sleep(rebalanceDelay / 2);
+
+        // Third (incidental) assignment with still only one worker in the group. Max delay has not
+        // been reached yet
+        applyAssignments(returnedAssignments);
+        memberConfigs = memberConfigs(leader, offset, assignments);
+        expectGeneration();
+        assignor.performTaskAssignment(leader, offset, memberConfigs, coordinator, protocolVersion);
+        ++rebalanceNum;
+        returnedAssignments = assignmentsCapture.getValue();
+        assertDelay(rebalanceDelay / 2, returnedAssignments);
+        expectedMemberConfigs = memberConfigs(leader, offset, returnedAssignments);
+        assertNoReassignments(memberConfigs, expectedMemberConfigs);
+        assertAssignment(0, 0, 0, 0, "worker1");
+
+        time.sleep(rebalanceDelay / 4);
+
+        // Fourth assignment with the second worker returning before the delay expires
+        // Since the delay is still active, lost assignments are not reassigned yet
+        applyAssignments(returnedAssignments);
+        memberConfigs = memberConfigs(leader, offset, assignments);
+        memberConfigs.put("worker2", new ExtendedWorkerState(leaderUrl, offset, null));
+        expectGeneration();
+        assignor.performTaskAssignment(leader, offset, memberConfigs, coordinator, protocolVersion);
+        ++rebalanceNum;
+        returnedAssignments = assignmentsCapture.getValue();
+        assertDelay(rebalanceDelay / 4, returnedAssignments);
+        expectedMemberConfigs = memberConfigs(leader, offset, returnedAssignments);
+        assertNoReassignments(memberConfigs, expectedMemberConfigs);
+        assertAssignment(0, 0, 0, 0, "worker1", "worker2");
+
+        time.sleep(rebalanceDelay / 4);
+
+        // Fifth assignment with the same two workers. The delay has expired, so the lost
+        // assignments ought to be assigned to the worker that has appeared as returned.
+        applyAssignments(returnedAssignments);
+        memberConfigs = memberConfigs(leader, offset, assignments);
+        expectGeneration();
+        assignor.performTaskAssignment(leader, offset, memberConfigs, coordinator, protocolVersion);
+        ++rebalanceNum;
+        returnedAssignments = assignmentsCapture.getValue();
+        assertDelay(0, returnedAssignments);
+        expectedMemberConfigs = memberConfigs(leader, offset, returnedAssignments);
+        assertNoReassignments(memberConfigs, expectedMemberConfigs);
+        assertAssignment(1, 4, 0, 0, "worker1", "worker2");
+
+        verify(coordinator, times(rebalanceNum)).configSnapshot();
+        verify(coordinator, times(rebalanceNum)).leaderState(any());
+        verify(coordinator, times(2 * rebalanceNum)).generationId();
+        verify(coordinator, times(rebalanceNum)).memberId();
+        verify(coordinator, times(rebalanceNum)).lastCompletedGenerationId();
+    }
+
+    @Test
+    public void testTaskAssignmentWhenLeaderLeavesPermanently() {
+        // Customize assignor for this test case
+        time = new MockTime();
+        initAssignor();
+
+        when(coordinator.configSnapshot()).thenReturn(configState);
+        doReturn(Collections.EMPTY_MAP).when(assignor).serializeAssignments(assignmentsCapture.capture());
+
+        // First assignment with 3 workers and 2 connectors configured but not yet assigned
+        memberConfigs.put("worker2", new ExtendedWorkerState(leaderUrl, offset, null));
+        memberConfigs.put("worker3", new ExtendedWorkerState(leaderUrl, offset, null));
+        expectGeneration();
+        assignor.performTaskAssignment(leader, offset, memberConfigs, coordinator, protocolVersion);
+        ++rebalanceNum;
+        returnedAssignments = assignmentsCapture.getValue();
+        assertDelay(0, returnedAssignments);
+        expectedMemberConfigs = memberConfigs(leader, offset, returnedAssignments);
+        assertNoReassignments(memberConfigs, expectedMemberConfigs);
+        assertAssignment(2, 8, 0, 0, "worker1", "worker2", "worker3");
+
+        // Second assignment with two workers remaining in the group. The worker that left the
+        // group was the leader. The new leader has no previous assignments and is not tracking a
+        // delay upon a leader's exit
+        applyAssignments(returnedAssignments);
+        assignments.remove("worker1");
+        leader = "worker2";
+        leaderUrl = expectedLeaderUrl(leader);
+        memberConfigs = memberConfigs(leader, offset, assignments);
+        // The fact that the leader bounces means that the assignor starts from a clean slate
+        initAssignor();
+
+        // Capture needs to be reset to point to the new assignor
+        doReturn(Collections.EMPTY_MAP).when(assignor).serializeAssignments(assignmentsCapture.capture());
+        expectGeneration();
+        assignor.performTaskAssignment(leader, offset, memberConfigs, coordinator, protocolVersion);
+        ++rebalanceNum;
+        returnedAssignments = assignmentsCapture.getValue();
+        assertDelay(0, returnedAssignments);
+        expectedMemberConfigs = memberConfigs(leader, offset, returnedAssignments);
+        assertNoReassignments(memberConfigs, expectedMemberConfigs);
+        assertAssignment(1, 3, 0, 0, "worker2", "worker3");
+
+        // Third (incidental) assignment with still only one worker in the group.
+        applyAssignments(returnedAssignments);
+        memberConfigs = memberConfigs(leader, offset, assignments);
+        expectGeneration();
+        assignor.performTaskAssignment(leader, offset, memberConfigs, coordinator, protocolVersion);
+        ++rebalanceNum;
+        returnedAssignments = assignmentsCapture.getValue();
+        expectedMemberConfigs = memberConfigs(leader, offset, returnedAssignments);
+        assertNoReassignments(memberConfigs, expectedMemberConfigs);
+        assertAssignment(0, 0, 0, 0, "worker2", "worker3");
+
+        verify(coordinator, times(rebalanceNum)).configSnapshot();
+        verify(coordinator, times(rebalanceNum)).leaderState(any());
+        verify(coordinator, times(2 * rebalanceNum)).generationId();
+        verify(coordinator, times(rebalanceNum)).memberId();
+        verify(coordinator, times(rebalanceNum)).lastCompletedGenerationId();
+    }
+
+    @Test
+    public void testTaskAssignmentWhenLeaderBounces() {
+        // Customize assignor for this test case
+        time = new MockTime();
+        initAssignor();
+
+        when(coordinator.configSnapshot()).thenReturn(configState);
+        doReturn(Collections.EMPTY_MAP).when(assignor).serializeAssignments(assignmentsCapture.capture());
+
+        // First assignment with 3 workers and 2 connectors configured but not yet assigned
+        memberConfigs.put("worker2", new ExtendedWorkerState(leaderUrl, offset, null));
+        memberConfigs.put("worker3", new ExtendedWorkerState(leaderUrl, offset, null));
+        expectGeneration();
+        assignor.performTaskAssignment(leader, offset, memberConfigs, coordinator, protocolVersion);
+        ++rebalanceNum;
+        returnedAssignments = assignmentsCapture.getValue();
+        assertDelay(0, returnedAssignments);
+        expectedMemberConfigs = memberConfigs(leader, offset, returnedAssignments);
+        assertNoReassignments(memberConfigs, expectedMemberConfigs);
+        assertAssignment(2, 8, 0, 0, "worker1", "worker2", "worker3");
+
+        // Second assignment with two workers remaining in the group. The worker that left the
+        // group was the leader. The new leader has no previous assignments and is not tracking a
+        // delay upon a leader's exit
+        applyAssignments(returnedAssignments);
+        assignments.remove("worker1");
+        leader = "worker2";
+        leaderUrl = expectedLeaderUrl(leader);
+        memberConfigs = memberConfigs(leader, offset, assignments);
+        // The fact that the leader bounces means that the assignor starts from a clean slate
+        initAssignor();
+
+        // Capture needs to be reset to point to the new assignor
+        doReturn(Collections.EMPTY_MAP).when(assignor).serializeAssignments(assignmentsCapture.capture());
+        expectGeneration();
+        assignor.performTaskAssignment(leader, offset, memberConfigs, coordinator, protocolVersion);
+        ++rebalanceNum;
+        returnedAssignments = assignmentsCapture.getValue();
+        assertDelay(0, returnedAssignments);
+        expectedMemberConfigs = memberConfigs(leader, offset, returnedAssignments);
+        assertNoReassignments(memberConfigs, expectedMemberConfigs);
+        assertAssignment(1, 3, 0, 0, "worker2", "worker3");
+
+        // Third assignment with the previous leader returning as a follower. In this case, the
+        // arrival of the previous leader is treated as an arrival of a new worker. Reassignment
+        // happens immediately, first with a revocation
+        applyAssignments(returnedAssignments);
+        memberConfigs = memberConfigs(leader, offset, assignments);
+        memberConfigs.put("worker1", new ExtendedWorkerState(leaderUrl, offset, null));
+        expectGeneration();
+        assignor.performTaskAssignment(leader, offset, memberConfigs, coordinator, protocolVersion);
+        ++rebalanceNum;
+        returnedAssignments = assignmentsCapture.getValue();
+        expectedMemberConfigs = memberConfigs(leader, offset, returnedAssignments);
+        assertNoReassignments(memberConfigs, expectedMemberConfigs);
+        assertAssignment(0, 0, 0, 2, "worker1", "worker2", "worker3");
+
+        // Fourth assignment after revocations
+        applyAssignments(returnedAssignments);
+        memberConfigs = memberConfigs(leader, offset, assignments);
+        expectGeneration();
+        assignor.performTaskAssignment(leader, offset, memberConfigs, coordinator, protocolVersion);
+        ++rebalanceNum;
+        returnedAssignments = assignmentsCapture.getValue();
+        assertDelay(0, returnedAssignments);
+        expectedMemberConfigs = memberConfigs(leader, offset, returnedAssignments);
+        assertNoReassignments(memberConfigs, expectedMemberConfigs);
+        assertAssignment(0, 2, 0, 0, "worker1", "worker2", "worker3");
+
+        verify(coordinator, times(rebalanceNum)).configSnapshot();
+        verify(coordinator, times(rebalanceNum)).leaderState(any());
+        verify(coordinator, times(2 * rebalanceNum)).generationId();
+        verify(coordinator, times(rebalanceNum)).memberId();
+        verify(coordinator, times(rebalanceNum)).lastCompletedGenerationId();
+    }
+
+    @Test
+    public void testTaskAssignmentWhenFirstAssignmentAttemptFails() {
+        // Customize assignor for this test case
+        time = new MockTime();
+        initAssignor();
+
+        when(coordinator.configSnapshot()).thenReturn(configState);
+        doThrow(new RuntimeException("Unable to send computed assignment with SyncGroupRequest"))
+                .when(assignor).serializeAssignments(assignmentsCapture.capture());
+
+        // First assignment with 2 workers and 2 connectors configured but not yet assigned
+        memberConfigs.put("worker2", new ExtendedWorkerState(leaderUrl, offset, null));
+        try {
+            expectGeneration();
+            assignor.performTaskAssignment(leader, offset, memberConfigs, coordinator, protocolVersion);
+        } catch (RuntimeException e) {
+            RequestFuture.failure(e);
+        }
+        ++rebalanceNum;
+        returnedAssignments = assignmentsCapture.getValue();
+        expectedMemberConfigs = memberConfigs(leader, offset, returnedAssignments);
+        // This was the assignment that should have been sent, but didn't make it all the way
+        assertDelay(0, returnedAssignments);
+        assertNoReassignments(memberConfigs, expectedMemberConfigs);
+        assertAssignment(2, 8, 0, 0, "worker1", "worker2");
+
+        // Second assignment happens with members returning the same assignments (memberConfigs)
+        // as the first time. The assignor detects that the number of members did not change and
+        // avoids the rebalance delay, treating the lost assignments as new assignments.
+        doReturn(Collections.EMPTY_MAP).when(assignor).serializeAssignments(assignmentsCapture.capture());
+        expectGeneration();
+        assignor.performTaskAssignment(leader, offset, memberConfigs, coordinator, protocolVersion);
+        ++rebalanceNum;
+        returnedAssignments = assignmentsCapture.getValue();
+        assertDelay(0, returnedAssignments);
+        expectedMemberConfigs = memberConfigs(leader, offset, returnedAssignments);
+        assertNoReassignments(memberConfigs, expectedMemberConfigs);
+        assertAssignment(2, 8, 0, 0, "worker1", "worker2");
+
+        verify(coordinator, times(rebalanceNum)).configSnapshot();
+        verify(coordinator, times(rebalanceNum)).leaderState(any());
+        verify(coordinator, times(2 * rebalanceNum)).generationId();
+        verify(coordinator, times(rebalanceNum)).memberId();
+        verify(coordinator, times(rebalanceNum)).lastCompletedGenerationId();
+    }
+
+    @Test
+    public void testTaskAssignmentWhenSubsequentAssignmentAttemptFails() {
+        // Customize assignor for this test case
+        time = new MockTime();
+        initAssignor();
+
+        when(coordinator.configSnapshot()).thenReturn(configState);
+        doReturn(Collections.EMPTY_MAP).when(assignor).serializeAssignments(assignmentsCapture.capture());
+
+        // First assignment with 2 workers and 2 connectors configured but not yet assigned
+        memberConfigs.put("worker2", new ExtendedWorkerState(leaderUrl, offset, null));
+        expectGeneration();
+        assignor.performTaskAssignment(leader, offset, memberConfigs, coordinator, protocolVersion);
+        ++rebalanceNum;
+        returnedAssignments = assignmentsCapture.getValue();
+        assertDelay(0, returnedAssignments);
+        expectedMemberConfigs = memberConfigs(leader, offset, returnedAssignments);
+        assertNoReassignments(memberConfigs, expectedMemberConfigs);
+        assertAssignment(2, 8, 0, 0, "worker1", "worker2");
+
+        when(coordinator.configSnapshot()).thenReturn(configState);
+        doThrow(new RuntimeException("Unable to send computed assignment with SyncGroupRequest"))
+                .when(assignor).serializeAssignments(assignmentsCapture.capture());
+
+        // Second assignment triggered by a third worker joining. The computed assignment should
+        // revoke tasks from the existing group. But the assignment won't be correctly delivered.
+        applyAssignments(returnedAssignments);
+        memberConfigs = memberConfigs(leader, offset, assignments);
+        memberConfigs.put("worker3", new ExtendedWorkerState(leaderUrl, offset, null));
+        try {
+            expectGeneration();
+            assignor.performTaskAssignment(leader, offset, memberConfigs, coordinator, protocolVersion);
+        } catch (RuntimeException e) {
+            RequestFuture.failure(e);
+        }
+        ++rebalanceNum;
+        returnedAssignments = assignmentsCapture.getValue();
+        expectedMemberConfigs = memberConfigs(leader, offset, returnedAssignments);
+        // This was the assignment that should have been sent, but didn't make it all the way
+        assertDelay(0, returnedAssignments);
+        assertNoReassignments(memberConfigs, expectedMemberConfigs);
+        assertAssignment(0, 0, 0, 2, "worker1", "worker2", "worker3");
+
+        // Third assignment happens with members returning the same assignments (memberConfigs)
+        // as the first time.
+        doReturn(Collections.EMPTY_MAP).when(assignor).serializeAssignments(assignmentsCapture.capture());
+        expectGeneration();
+        assignor.performTaskAssignment(leader, offset, memberConfigs, coordinator, protocolVersion);
+        ++rebalanceNum;
+        returnedAssignments = assignmentsCapture.getValue();
+        expectedMemberConfigs = memberConfigs(leader, offset, returnedAssignments);
+        assertDelay(0, returnedAssignments);
+        assertNoReassignments(memberConfigs, expectedMemberConfigs);
+        assertAssignment(0, 0, 0, 2, "worker1", "worker2", "worker3");
+
+        verify(coordinator, times(rebalanceNum)).configSnapshot();
+        verify(coordinator, times(rebalanceNum)).leaderState(any());
+        verify(coordinator, times(2 * rebalanceNum)).generationId();
+        verify(coordinator, times(rebalanceNum)).memberId();
+        verify(coordinator, times(rebalanceNum)).lastCompletedGenerationId();
+    }
+
+    @Test
+    public void testTaskAssignmentWhenSubsequentAssignmentAttemptFailsOutsideTheAssignor() {
+        // Customize assignor for this test case
+        time = new MockTime();
+        initAssignor();
+
+        expectGeneration();
+        when(coordinator.configSnapshot()).thenReturn(configState);
+        doReturn(Collections.EMPTY_MAP).when(assignor).serializeAssignments(assignmentsCapture.capture());
+
+        // First assignment with 2 workers and 2 connectors configured but not yet assigned
+        memberConfigs.put("worker2", new ExtendedWorkerState(leaderUrl, offset, null));
+        expectGeneration();
+        assignor.performTaskAssignment(leader, offset, memberConfigs, coordinator, protocolVersion);
+        ++rebalanceNum;
+        returnedAssignments = assignmentsCapture.getValue();
+        assertDelay(0, returnedAssignments);
+        expectedMemberConfigs = memberConfigs(leader, offset, returnedAssignments);
+        assertNoReassignments(memberConfigs, expectedMemberConfigs);
+        assertAssignment(2, 8, 0, 0, "worker1", "worker2");
+
+        // Second assignment triggered by a third worker joining. The computed assignment should
+        // revoke tasks from the existing group. But the assignment won't be correctly delivered
+        // and sync group with fail on the leader worker.
+        applyAssignments(returnedAssignments);
+        memberConfigs = memberConfigs(leader, offset, assignments);
+        memberConfigs.put("worker3", new ExtendedWorkerState(leaderUrl, offset, null));
+        when(coordinator.generationId())
+                .thenReturn(assignor.previousGenerationId + 1)
+                .thenReturn(assignor.previousGenerationId + 1);
+        when(coordinator.lastCompletedGenerationId()).thenReturn(assignor.previousGenerationId - 1);
+        assignor.performTaskAssignment(leader, offset, memberConfigs, coordinator, protocolVersion);
+        ++rebalanceNum;
+        returnedAssignments = assignmentsCapture.getValue();
+        expectedMemberConfigs = memberConfigs(leader, offset, returnedAssignments);
+        // This was the assignment that should have been sent, but didn't make it all the way
+        assertDelay(0, returnedAssignments);
+        assertNoReassignments(memberConfigs, expectedMemberConfigs);
+        assertAssignment(0, 0, 0, 2, "worker1", "worker2", "worker3");
+
+        // Third assignment happens with members returning the same assignments (memberConfigs)
+        // as the first time.
+        when(coordinator.lastCompletedGenerationId()).thenReturn(assignor.previousGenerationId - 1);
+        doReturn(Collections.EMPTY_MAP).when(assignor).serializeAssignments(assignmentsCapture.capture());
+        expectGeneration();
+        assignor.performTaskAssignment(leader, offset, memberConfigs, coordinator, protocolVersion);
+        ++rebalanceNum;
+        returnedAssignments = assignmentsCapture.getValue();
+        expectedMemberConfigs = memberConfigs(leader, offset, returnedAssignments);
+        assertDelay(0, returnedAssignments);
+        assertNoReassignments(memberConfigs, expectedMemberConfigs);
+        assertAssignment(0, 0, 0, 2, "worker1", "worker2", "worker3");
+
+        verify(coordinator, times(rebalanceNum)).configSnapshot();
+        verify(coordinator, times(rebalanceNum)).leaderState(any());
+        verify(coordinator, times(2 * rebalanceNum)).generationId();
+        verify(coordinator, times(rebalanceNum)).memberId();
+        verify(coordinator, times(rebalanceNum)).lastCompletedGenerationId();
+    }
+
+    @Test
+    public void testTaskAssignmentWhenConnectorsAreDeleted() {
+        configState = clusterConfigState(offset, 3, 4);
+        when(coordinator.configSnapshot()).thenReturn(configState);
+        doReturn(Collections.EMPTY_MAP).when(assignor).serializeAssignments(assignmentsCapture.capture());
+
+        // First assignment with 1 worker and 2 connectors configured but not yet assigned
+        memberConfigs.put("worker2", new ExtendedWorkerState(leaderUrl, offset, null));
+        expectGeneration();
+        assignor.performTaskAssignment(leader, offset, memberConfigs, coordinator, protocolVersion);
+        ++rebalanceNum;
+        returnedAssignments = assignmentsCapture.getValue();
+        assertDelay(0, returnedAssignments);
+        expectedMemberConfigs = memberConfigs(leader, offset, returnedAssignments);
+        assertNoReassignments(memberConfigs, expectedMemberConfigs);
+        assertAssignment(3, 12, 0, 0, "worker1", "worker2");
+
+        // Second assignment with an updated config state that reflects removal of a connector
+        configState = clusterConfigState(offset + 1, 2, 4);
+        when(coordinator.configSnapshot()).thenReturn(configState);
+        applyAssignments(returnedAssignments);
+        memberConfigs = memberConfigs(leader, offset, assignments);
+        expectGeneration();
+        assignor.performTaskAssignment(leader, offset, memberConfigs, coordinator, protocolVersion);
+        ++rebalanceNum;
+        returnedAssignments = assignmentsCapture.getValue();
+        assertDelay(0, returnedAssignments);
+        expectedMemberConfigs = memberConfigs(leader, offset, returnedAssignments);
+        assertNoReassignments(memberConfigs, expectedMemberConfigs);
+        assertAssignment(0, 0, 1, 4, "worker1", "worker2");
+
+        verify(coordinator, times(rebalanceNum)).configSnapshot();
+        verify(coordinator, times(rebalanceNum)).leaderState(any());
+        verify(coordinator, times(2 * rebalanceNum)).generationId();
+        verify(coordinator, times(rebalanceNum)).memberId();
+        verify(coordinator, times(rebalanceNum)).lastCompletedGenerationId();
+    }
+
+    @Test
+    public void testAssignConnectorsWhenBalanced() {
+        int num = 2;
+        List existingAssignment = IntStream.range(0, 3)
+                .mapToObj(i -> workerLoad("worker" + i, i * num, num, i * num, num))
+                .collect(Collectors.toList());
+
+        List expectedAssignment = existingAssignment.stream()
+                .map(wl -> new WorkerLoad.Builder(wl.worker()).withCopies(wl.connectors(), wl.tasks()).build())
+                .collect(Collectors.toList());
+        expectedAssignment.get(0).connectors().addAll(Arrays.asList("connector6", "connector9"));
+        expectedAssignment.get(1).connectors().addAll(Arrays.asList("connector7", "connector10"));
+        expectedAssignment.get(2).connectors().addAll(Arrays.asList("connector8"));
+
+        List newConnectors = newConnectors(6, 11);
+        assignor.assignConnectors(existingAssignment, newConnectors);
+        assertEquals(expectedAssignment, existingAssignment);
+    }
+
+    @Test
+    public void testAssignTasksWhenBalanced() {
+        int num = 2;
+        List existingAssignment = IntStream.range(0, 3)
+                .mapToObj(i -> workerLoad("worker" + i, i * num, num, i * num, num))
+                .collect(Collectors.toList());
+
+        List expectedAssignment = existingAssignment.stream()
+                .map(wl -> new WorkerLoad.Builder(wl.worker()).withCopies(wl.connectors(), wl.tasks()).build())
+                .collect(Collectors.toList());
+
+        expectedAssignment.get(0).connectors().addAll(Arrays.asList("connector6", "connector9"));
+        expectedAssignment.get(1).connectors().addAll(Arrays.asList("connector7", "connector10"));
+        expectedAssignment.get(2).connectors().addAll(Arrays.asList("connector8"));
+
+        expectedAssignment.get(0).tasks().addAll(Arrays.asList(new ConnectorTaskId("task", 6), new ConnectorTaskId("task", 9)));
+        expectedAssignment.get(1).tasks().addAll(Arrays.asList(new ConnectorTaskId("task", 7), new ConnectorTaskId("task", 10)));
+        expectedAssignment.get(2).tasks().addAll(Arrays.asList(new ConnectorTaskId("task", 8)));
+
+        List newConnectors = newConnectors(6, 11);
+        assignor.assignConnectors(existingAssignment, newConnectors);
+        List newTasks = newTasks(6, 11);
+        assignor.assignTasks(existingAssignment, newTasks);
+        assertEquals(expectedAssignment, existingAssignment);
+    }
+
+    @Test
+    public void testAssignConnectorsWhenImbalanced() {
+        List existingAssignment = new ArrayList<>();
+        existingAssignment.add(workerLoad("worker0", 0, 2, 0, 2));
+        existingAssignment.add(workerLoad("worker1", 2, 3, 2, 3));
+        existingAssignment.add(workerLoad("worker2", 5, 4, 5, 4));
+        existingAssignment.add(emptyWorkerLoad("worker3"));
+
+        List newConnectors = newConnectors(9, 24);
+        List newTasks = newTasks(9, 24);
+        assignor.assignConnectors(existingAssignment, newConnectors);
+        assignor.assignTasks(existingAssignment, newTasks);
+        for (WorkerLoad worker : existingAssignment) {
+            assertEquals(6, worker.connectorsSize());
+            assertEquals(6, worker.tasksSize());
+        }
+    }
+
+    @Test
+    public void testLostAssignmentHandlingWhenWorkerBounces() {
+        // Customize assignor for this test case
+        time = new MockTime();
+        initAssignor();
+
+        assertTrue(assignor.candidateWorkersForReassignment.isEmpty());
+        assertEquals(0, assignor.scheduledRebalance);
+        assertEquals(0, assignor.delay);
+
+        Map configuredAssignment = new HashMap<>();
+        configuredAssignment.put("worker0", workerLoad("worker0", 0, 2, 0, 4));
+        configuredAssignment.put("worker1", workerLoad("worker1", 2, 2, 4, 4));
+        configuredAssignment.put("worker2", workerLoad("worker2", 4, 2, 8, 4));
+        memberConfigs = memberConfigs(leader, offset, 0, 2);
+
+        ConnectorsAndTasks newSubmissions = new ConnectorsAndTasks.Builder().build();
+
+        // No lost assignments
+        assignor.handleLostAssignments(new ConnectorsAndTasks.Builder().build(),
+                newSubmissions,
+                new ArrayList<>(configuredAssignment.values()),
+                memberConfigs);
+
+        assertThat("Wrong set of workers for reassignments",
+                Collections.emptySet(),
+                is(assignor.candidateWorkersForReassignment));
+        assertEquals(0, assignor.scheduledRebalance);
+        assertEquals(0, assignor.delay);
+
+        assignor.previousMembers = new HashSet<>(memberConfigs.keySet());
+        String flakyWorker = "worker1";
+        WorkerLoad lostLoad = workerLoad(flakyWorker, 2, 2, 4, 4);
+        memberConfigs.remove(flakyWorker);
+
+        ConnectorsAndTasks lostAssignments = new ConnectorsAndTasks.Builder()
+                .withCopies(lostLoad.connectors(), lostLoad.tasks()).build();
+
+        // Lost assignments detected - No candidate worker has appeared yet (worker with no assignments)
+        assignor.handleLostAssignments(lostAssignments, newSubmissions,
+                new ArrayList<>(configuredAssignment.values()), memberConfigs);
+
+        assertThat("Wrong set of workers for reassignments",
+                Collections.emptySet(),
+                is(assignor.candidateWorkersForReassignment));
+        assertEquals(time.milliseconds() + rebalanceDelay, assignor.scheduledRebalance);
+        assertEquals(rebalanceDelay, assignor.delay);
+
+        assignor.previousMembers = new HashSet<>(memberConfigs.keySet());
+        time.sleep(rebalanceDelay / 2);
+        rebalanceDelay /= 2;
+
+        // A new worker (probably returning worker) has joined
+        configuredAssignment.put(flakyWorker, new WorkerLoad.Builder(flakyWorker).build());
+        memberConfigs.put(flakyWorker, new ExtendedWorkerState(leaderUrl, offset, null));
+        assignor.handleLostAssignments(lostAssignments, newSubmissions,
+                new ArrayList<>(configuredAssignment.values()), memberConfigs);
+
+        assertThat("Wrong set of workers for reassignments",
+                Collections.singleton(flakyWorker),
+                is(assignor.candidateWorkersForReassignment));
+        assertEquals(time.milliseconds() + rebalanceDelay, assignor.scheduledRebalance);
+        assertEquals(rebalanceDelay, assignor.delay);
+
+        assignor.previousMembers = new HashSet<>(memberConfigs.keySet());
+        time.sleep(rebalanceDelay);
+
+        // The new worker has still no assignments
+        assignor.handleLostAssignments(lostAssignments, newSubmissions,
+                new ArrayList<>(configuredAssignment.values()), memberConfigs);
+
+        assertTrue("Wrong assignment of lost connectors",
+                configuredAssignment.getOrDefault(flakyWorker, new WorkerLoad.Builder(flakyWorker).build())
+                        .connectors()
+                        .containsAll(lostAssignments.connectors()));
+        assertTrue("Wrong assignment of lost tasks",
+                configuredAssignment.getOrDefault(flakyWorker, new WorkerLoad.Builder(flakyWorker).build())
+                        .tasks()
+                        .containsAll(lostAssignments.tasks()));
+        assertThat("Wrong set of workers for reassignments",
+                Collections.emptySet(),
+                is(assignor.candidateWorkersForReassignment));
+        assertEquals(0, assignor.scheduledRebalance);
+        assertEquals(0, assignor.delay);
+    }
+
+    @Test
+    public void testLostAssignmentHandlingWhenWorkerLeavesPermanently() {
+        // Customize assignor for this test case
+        time = new MockTime();
+        initAssignor();
+
+        assertTrue(assignor.candidateWorkersForReassignment.isEmpty());
+        assertEquals(0, assignor.scheduledRebalance);
+        assertEquals(0, assignor.delay);
+
+        Map configuredAssignment = new HashMap<>();
+        configuredAssignment.put("worker0", workerLoad("worker0", 0, 2, 0, 4));
+        configuredAssignment.put("worker1", workerLoad("worker1", 2, 2, 4, 4));
+        configuredAssignment.put("worker2", workerLoad("worker2", 4, 2, 8, 4));
+        memberConfigs = memberConfigs(leader, offset, 0, 2);
+
+        ConnectorsAndTasks newSubmissions = new ConnectorsAndTasks.Builder().build();
+
+        // No lost assignments
+        assignor.handleLostAssignments(new ConnectorsAndTasks.Builder().build(),
+                newSubmissions,
+                new ArrayList<>(configuredAssignment.values()),
+                memberConfigs);
+
+        assertThat("Wrong set of workers for reassignments",
+                Collections.emptySet(),
+                is(assignor.candidateWorkersForReassignment));
+        assertEquals(0, assignor.scheduledRebalance);
+        assertEquals(0, assignor.delay);
+
+        assignor.previousMembers = new HashSet<>(memberConfigs.keySet());
+        String removedWorker = "worker1";
+        WorkerLoad lostLoad = workerLoad(removedWorker, 2, 2, 4, 4);
+        memberConfigs.remove(removedWorker);
+
+        ConnectorsAndTasks lostAssignments = new ConnectorsAndTasks.Builder()
+                .withCopies(lostLoad.connectors(), lostLoad.tasks()).build();
+
+        // Lost assignments detected - No candidate worker has appeared yet (worker with no assignments)
+        assignor.handleLostAssignments(lostAssignments, newSubmissions,
+                new ArrayList<>(configuredAssignment.values()), memberConfigs);
+
+        assertThat("Wrong set of workers for reassignments",
+                Collections.emptySet(),
+                is(assignor.candidateWorkersForReassignment));
+        assertEquals(time.milliseconds() + rebalanceDelay, assignor.scheduledRebalance);
+        assertEquals(rebalanceDelay, assignor.delay);
+
+        assignor.previousMembers = new HashSet<>(memberConfigs.keySet());
+        time.sleep(rebalanceDelay / 2);
+        rebalanceDelay /= 2;
+
+        // No new worker has joined
+        assignor.handleLostAssignments(lostAssignments, newSubmissions,
+                new ArrayList<>(configuredAssignment.values()), memberConfigs);
+
+        assertThat("Wrong set of workers for reassignments",
+                Collections.emptySet(),
+                is(assignor.candidateWorkersForReassignment));
+        assertEquals(time.milliseconds() + rebalanceDelay, assignor.scheduledRebalance);
+        assertEquals(rebalanceDelay, assignor.delay);
+
+        time.sleep(rebalanceDelay);
+
+        assignor.handleLostAssignments(lostAssignments, newSubmissions,
+                new ArrayList<>(configuredAssignment.values()), memberConfigs);
+
+        assertTrue("Wrong assignment of lost connectors",
+                newSubmissions.connectors().containsAll(lostAssignments.connectors()));
+        assertTrue("Wrong assignment of lost tasks",
+                newSubmissions.tasks().containsAll(lostAssignments.tasks()));
+        assertThat("Wrong set of workers for reassignments",
+                Collections.emptySet(),
+                is(assignor.candidateWorkersForReassignment));
+        assertEquals(0, assignor.scheduledRebalance);
+        assertEquals(0, assignor.delay);
+    }
+
+    @Test
+    public void testLostAssignmentHandlingWithMoreThanOneCandidates() {
+        // Customize assignor for this test case
+        time = new MockTime();
+        initAssignor();
+
+        assertTrue(assignor.candidateWorkersForReassignment.isEmpty());
+        assertEquals(0, assignor.scheduledRebalance);
+        assertEquals(0, assignor.delay);
+
+        Map configuredAssignment = new HashMap<>();
+        configuredAssignment.put("worker0", workerLoad("worker0", 0, 2, 0, 4));
+        configuredAssignment.put("worker1", workerLoad("worker1", 2, 2, 4, 4));
+        configuredAssignment.put("worker2", workerLoad("worker2", 4, 2, 8, 4));
+        memberConfigs = memberConfigs(leader, offset, 0, 2);
+
+        ConnectorsAndTasks newSubmissions = new ConnectorsAndTasks.Builder().build();
+
+        // No lost assignments
+        assignor.handleLostAssignments(new ConnectorsAndTasks.Builder().build(),
+                newSubmissions,
+                new ArrayList<>(configuredAssignment.values()),
+                memberConfigs);
+
+        assertThat("Wrong set of workers for reassignments",
+                Collections.emptySet(),
+                is(assignor.candidateWorkersForReassignment));
+        assertEquals(0, assignor.scheduledRebalance);
+        assertEquals(0, assignor.delay);
+
+        assignor.previousMembers = new HashSet<>(memberConfigs.keySet());
+        String flakyWorker = "worker1";
+        WorkerLoad lostLoad = workerLoad(flakyWorker, 2, 2, 4, 4);
+        memberConfigs.remove(flakyWorker);
+        String newWorker = "worker3";
+
+        ConnectorsAndTasks lostAssignments = new ConnectorsAndTasks.Builder()
+                .withCopies(lostLoad.connectors(), lostLoad.tasks()).build();
+
+        // Lost assignments detected - A new worker also has joined that is not the returning worker
+        configuredAssignment.put(newWorker, new WorkerLoad.Builder(newWorker).build());
+        memberConfigs.put(newWorker, new ExtendedWorkerState(leaderUrl, offset, null));
+        assignor.handleLostAssignments(lostAssignments, newSubmissions,
+                new ArrayList<>(configuredAssignment.values()), memberConfigs);
+
+        assertThat("Wrong set of workers for reassignments",
+                Collections.singleton(newWorker),
+                is(assignor.candidateWorkersForReassignment));
+        assertEquals(time.milliseconds() + rebalanceDelay, assignor.scheduledRebalance);
+        assertEquals(rebalanceDelay, assignor.delay);
+
+        assignor.previousMembers = new HashSet<>(memberConfigs.keySet());
+        time.sleep(rebalanceDelay / 2);
+        rebalanceDelay /= 2;
+
+        // Now two new workers have joined
+        configuredAssignment.put(flakyWorker, new WorkerLoad.Builder(flakyWorker).build());
+        memberConfigs.put(flakyWorker, new ExtendedWorkerState(leaderUrl, offset, null));
+        assignor.handleLostAssignments(lostAssignments, newSubmissions,
+                new ArrayList<>(configuredAssignment.values()), memberConfigs);
+
+        Set expectedWorkers = new HashSet<>();
+        expectedWorkers.addAll(Arrays.asList(newWorker, flakyWorker));
+        assertThat("Wrong set of workers for reassignments",
+                expectedWorkers,
+                is(assignor.candidateWorkersForReassignment));
+        assertEquals(time.milliseconds() + rebalanceDelay, assignor.scheduledRebalance);
+        assertEquals(rebalanceDelay, assignor.delay);
+
+        assignor.previousMembers = new HashSet<>(memberConfigs.keySet());
+        time.sleep(rebalanceDelay);
+
+        // The new workers have new assignments, other than the lost ones
+        configuredAssignment.put(flakyWorker, workerLoad(flakyWorker, 6, 2, 8, 4));
+        configuredAssignment.put(newWorker, workerLoad(newWorker, 8, 2, 12, 4));
+        // we don't reflect these new assignments in memberConfigs currently because they are not
+        // used in handleLostAssignments method
+        assignor.handleLostAssignments(lostAssignments, newSubmissions,
+                new ArrayList<>(configuredAssignment.values()), memberConfigs);
+
+        // newWorker joined first, so should be picked up first as a candidate for reassignment
+        assertTrue("Wrong assignment of lost connectors",
+                configuredAssignment.getOrDefault(newWorker, new WorkerLoad.Builder(flakyWorker).build())
+                        .connectors()
+                        .containsAll(lostAssignments.connectors()));
+        assertTrue("Wrong assignment of lost tasks",
+                configuredAssignment.getOrDefault(newWorker, new WorkerLoad.Builder(flakyWorker).build())
+                        .tasks()
+                        .containsAll(lostAssignments.tasks()));
+        assertThat("Wrong set of workers for reassignments",
+                Collections.emptySet(),
+                is(assignor.candidateWorkersForReassignment));
+        assertEquals(0, assignor.scheduledRebalance);
+        assertEquals(0, assignor.delay);
+    }
+
+    @Test
+    public void testLostAssignmentHandlingWhenWorkerBouncesBackButFinallyLeaves() {
+        // Customize assignor for this test case
+        time = new MockTime();
+        initAssignor();
+
+        assertTrue(assignor.candidateWorkersForReassignment.isEmpty());
+        assertEquals(0, assignor.scheduledRebalance);
+        assertEquals(0, assignor.delay);
+
+        Map configuredAssignment = new HashMap<>();
+        configuredAssignment.put("worker0", workerLoad("worker0", 0, 2, 0, 4));
+        configuredAssignment.put("worker1", workerLoad("worker1", 2, 2, 4, 4));
+        configuredAssignment.put("worker2", workerLoad("worker2", 4, 2, 8, 4));
+        memberConfigs = memberConfigs(leader, offset, 0, 2);
+
+        ConnectorsAndTasks newSubmissions = new ConnectorsAndTasks.Builder().build();
+
+        // No lost assignments
+        assignor.handleLostAssignments(new ConnectorsAndTasks.Builder().build(),
+                newSubmissions,
+                new ArrayList<>(configuredAssignment.values()),
+                memberConfigs);
+
+        assertThat("Wrong set of workers for reassignments",
+                Collections.emptySet(),
+                is(assignor.candidateWorkersForReassignment));
+        assertEquals(0, assignor.scheduledRebalance);
+        assertEquals(0, assignor.delay);
+
+        assignor.previousMembers = new HashSet<>(memberConfigs.keySet());
+        String veryFlakyWorker = "worker1";
+        WorkerLoad lostLoad = workerLoad(veryFlakyWorker, 2, 2, 4, 4);
+        memberConfigs.remove(veryFlakyWorker);
+
+        ConnectorsAndTasks lostAssignments = new ConnectorsAndTasks.Builder()
+                .withCopies(lostLoad.connectors(), lostLoad.tasks()).build();
+
+        // Lost assignments detected - No candidate worker has appeared yet (worker with no assignments)
+        assignor.handleLostAssignments(lostAssignments, newSubmissions,
+                new ArrayList<>(configuredAssignment.values()), memberConfigs);
+
+        assertThat("Wrong set of workers for reassignments",
+                Collections.emptySet(),
+                is(assignor.candidateWorkersForReassignment));
+        assertEquals(time.milliseconds() + rebalanceDelay, assignor.scheduledRebalance);
+        assertEquals(rebalanceDelay, assignor.delay);
+
+        assignor.previousMembers = new HashSet<>(memberConfigs.keySet());
+        time.sleep(rebalanceDelay / 2);
+        rebalanceDelay /= 2;
+
+        // A new worker (probably returning worker) has joined
+        configuredAssignment.put(veryFlakyWorker, new WorkerLoad.Builder(veryFlakyWorker).build());
+        memberConfigs.put(veryFlakyWorker, new ExtendedWorkerState(leaderUrl, offset, null));
+        assignor.handleLostAssignments(lostAssignments, newSubmissions,
+                new ArrayList<>(configuredAssignment.values()), memberConfigs);
+
+        assertThat("Wrong set of workers for reassignments",
+                Collections.singleton(veryFlakyWorker),
+                is(assignor.candidateWorkersForReassignment));
+        assertEquals(time.milliseconds() + rebalanceDelay, assignor.scheduledRebalance);
+        assertEquals(rebalanceDelay, assignor.delay);
+
+        assignor.previousMembers = new HashSet<>(memberConfigs.keySet());
+        time.sleep(rebalanceDelay);
+
+        // The returning worker leaves permanently after joining briefly during the delay
+        configuredAssignment.remove(veryFlakyWorker);
+        memberConfigs.remove(veryFlakyWorker);
+        assignor.handleLostAssignments(lostAssignments, newSubmissions,
+                new ArrayList<>(configuredAssignment.values()), memberConfigs);
+
+        assertTrue("Wrong assignment of lost connectors",
+                newSubmissions.connectors().containsAll(lostAssignments.connectors()));
+        assertTrue("Wrong assignment of lost tasks",
+                newSubmissions.tasks().containsAll(lostAssignments.tasks()));
+        assertThat("Wrong set of workers for reassignments",
+                Collections.emptySet(),
+                is(assignor.candidateWorkersForReassignment));
+        assertEquals(0, assignor.scheduledRebalance);
+        assertEquals(0, assignor.delay);
+    }
+
+    @Test
+    public void testTaskAssignmentWhenTasksDuplicatedInWorkerAssignment() {
+        when(coordinator.configSnapshot()).thenReturn(configState);
+        doReturn(Collections.EMPTY_MAP).when(assignor).serializeAssignments(assignmentsCapture.capture());
+
+        // First assignment with 1 worker and 2 connectors configured but not yet assigned
+        assignor.performTaskAssignment(leader, offset, memberConfigs, coordinator, protocolVersion);
+        ++rebalanceNum;
+        returnedAssignments = assignmentsCapture.getValue();
+        assertDelay(0, returnedAssignments);
+        expectedMemberConfigs = memberConfigs(leader, offset, returnedAssignments);
+        assertNoReassignments(memberConfigs, expectedMemberConfigs);
+        assertAssignment(2, 8, 0, 0, "worker1");
+
+        // Second assignment with a second worker with duplicate assignment joining and all connectors running on previous worker
+        applyAssignments(returnedAssignments);
+        memberConfigs = memberConfigs(leader, offset, assignments);
+        ExtendedAssignment duplicatedWorkerAssignment = newExpandableAssignment();
+        duplicatedWorkerAssignment.connectors().addAll(newConnectors(1, 2));
+        duplicatedWorkerAssignment.tasks().addAll(newTasks("connector1", 0, 4));
+        memberConfigs.put("worker2", new ExtendedWorkerState(leaderUrl, offset, duplicatedWorkerAssignment));
+        assignor.performTaskAssignment(leader, offset, memberConfigs, coordinator, protocolVersion);
+        ++rebalanceNum;
+        returnedAssignments = assignmentsCapture.getValue();
+        assertDelay(0, returnedAssignments);
+        expectedMemberConfigs = memberConfigs(leader, offset, returnedAssignments);
+        assertNoReassignments(memberConfigs, expectedMemberConfigs);
+        assertAssignment(0, 0, 2, 8, "worker1", "worker2");
+
+        // Third assignment after revocations
+        applyAssignments(returnedAssignments);
+        memberConfigs = memberConfigs(leader, offset, assignments);
+        assignor.performTaskAssignment(leader, offset, memberConfigs, coordinator, protocolVersion);
+        ++rebalanceNum;
+        returnedAssignments = assignmentsCapture.getValue();
+        assertDelay(0, returnedAssignments);
+        expectedMemberConfigs = memberConfigs(leader, offset, returnedAssignments);
+        assertNoReassignments(memberConfigs, expectedMemberConfigs);
+        assertAssignment(1, 4, 0, 2, "worker1", "worker2");
+
+        // fourth rebalance after revocations
+        applyAssignments(returnedAssignments);
+        memberConfigs = memberConfigs(leader, offset, assignments);
+        assignor.performTaskAssignment(leader, offset, memberConfigs, coordinator, protocolVersion);
+        ++rebalanceNum;
+        returnedAssignments = assignmentsCapture.getValue();
+        assertDelay(0, returnedAssignments);
+        expectedMemberConfigs = memberConfigs(leader, offset, returnedAssignments);
+        assertNoReassignments(memberConfigs, expectedMemberConfigs);
+        assertAssignment(0, 2, 0, 0, "worker1", "worker2");
+
+        // Fifth rebalance should not change assignments
+        applyAssignments(returnedAssignments);
+        memberConfigs = memberConfigs(leader, offset, assignments);
+        assignor.performTaskAssignment(leader, offset, memberConfigs, coordinator, protocolVersion);
+        ++rebalanceNum;
+        returnedAssignments = assignmentsCapture.getValue();
+        assertDelay(0, returnedAssignments);
+        expectedMemberConfigs = memberConfigs(leader, offset, returnedAssignments);
+        assertNoReassignments(memberConfigs, expectedMemberConfigs);
+        assertAssignment(0, 0, 0, 0, "worker1", "worker2");
+
+        verify(coordinator, times(rebalanceNum)).configSnapshot();
+        verify(coordinator, times(rebalanceNum)).leaderState(any());
+        verify(coordinator, times(2 * rebalanceNum)).generationId();
+        verify(coordinator, times(rebalanceNum)).memberId();
+        verify(coordinator, times(rebalanceNum)).lastCompletedGenerationId();
+    }
+
+    @Test
+    public void testDuplicatedAssignmentHandleWhenTheDuplicatedAssignmentsDeleted() {
+        when(coordinator.configSnapshot()).thenReturn(configState);
+        doReturn(Collections.EMPTY_MAP).when(assignor).serializeAssignments(assignmentsCapture.capture());
+
+        // First assignment with 1 worker and 2 connectors configured but not yet assigned
+        assignor.performTaskAssignment(leader, offset, memberConfigs, coordinator, protocolVersion);
+        ++rebalanceNum;
+        returnedAssignments = assignmentsCapture.getValue();
+        assertDelay(0, returnedAssignments);
+        expectedMemberConfigs = memberConfigs(leader, offset, returnedAssignments);
+        assertNoReassignments(memberConfigs, expectedMemberConfigs);
+        assertAssignment(2, 8, 0, 0, "worker1");
+
+        //delete connector1
+        configState = clusterConfigState(offset, 2, 1, 4);
+        when(coordinator.configSnapshot()).thenReturn(configState);
+
+        // Second assignment with a second worker with duplicate assignment joining and the duplicated assignment is deleted at the same time
+        applyAssignments(returnedAssignments);
+        memberConfigs = memberConfigs(leader, offset, assignments);
+        ExtendedAssignment duplicatedWorkerAssignment = newExpandableAssignment();
+        duplicatedWorkerAssignment.connectors().addAll(newConnectors(1, 2));
+        duplicatedWorkerAssignment.tasks().addAll(newTasks("connector1", 0, 4));
+        memberConfigs.put("worker2", new ExtendedWorkerState(leaderUrl, offset, duplicatedWorkerAssignment));
+        assignor.performTaskAssignment(leader, offset, memberConfigs, coordinator, protocolVersion);
+        ++rebalanceNum;
+        returnedAssignments = assignmentsCapture.getValue();
+        assertDelay(0, returnedAssignments);
+        expectedMemberConfigs = memberConfigs(leader, offset, returnedAssignments);
+        assertNoReassignments(memberConfigs, expectedMemberConfigs);
+        assertAssignment(0, 0, 2, 8, "worker1", "worker2");
+
+        // Third assignment after revocations
+        applyAssignments(returnedAssignments);
+        memberConfigs = memberConfigs(leader, offset, assignments);
+        assignor.performTaskAssignment(leader, offset, memberConfigs, coordinator, protocolVersion);
+        ++rebalanceNum;
+        returnedAssignments = assignmentsCapture.getValue();
+        assertDelay(0, returnedAssignments);
+        expectedMemberConfigs = memberConfigs(leader, offset, returnedAssignments);
+        assertNoReassignments(memberConfigs, expectedMemberConfigs);
+        assertAssignment(0, 0, 0, 2, "worker1", "worker2");
+
+        // fourth rebalance after revocations
+        applyAssignments(returnedAssignments);
+        memberConfigs = memberConfigs(leader, offset, assignments);
+        assignor.performTaskAssignment(leader, offset, memberConfigs, coordinator, protocolVersion);
+        ++rebalanceNum;
+        returnedAssignments = assignmentsCapture.getValue();
+        assertDelay(0, returnedAssignments);
+        expectedMemberConfigs = memberConfigs(leader, offset, returnedAssignments);
+        assertNoReassignments(memberConfigs, expectedMemberConfigs);
+        assertAssignment(0, 2, 0, 0, "worker1", "worker2");
+
+        // Fifth rebalance should not change assignments
+        applyAssignments(returnedAssignments);
+        memberConfigs = memberConfigs(leader, offset, assignments);
+        assignor.performTaskAssignment(leader, offset, memberConfigs, coordinator, protocolVersion);
+        ++rebalanceNum;
+        returnedAssignments = assignmentsCapture.getValue();
+        assertDelay(0, returnedAssignments);
+        expectedMemberConfigs = memberConfigs(leader, offset, returnedAssignments);
+        assertNoReassignments(memberConfigs, expectedMemberConfigs);
+        assertAssignment(0, 0, 0, 0, "worker1", "worker2");
+
+        verify(coordinator, times(rebalanceNum)).configSnapshot();
+        verify(coordinator, times(rebalanceNum)).leaderState(any());
+        verify(coordinator, times(2 * rebalanceNum)).generationId();
+        verify(coordinator, times(rebalanceNum)).memberId();
+        verify(coordinator, times(rebalanceNum)).lastCompletedGenerationId();
+    }
+
+    private WorkerLoad emptyWorkerLoad(String worker) {
+        return new WorkerLoad.Builder(worker).build();
+    }
+
+    private WorkerLoad workerLoad(String worker, int connectorStart, int connectorNum,
+                                  int taskStart, int taskNum) {
+        return new WorkerLoad.Builder(worker).with(
+                newConnectors(connectorStart, connectorStart + connectorNum),
+                newTasks(taskStart, taskStart + taskNum)).build();
+    }
+
+    private static List newConnectors(int start, int end) {
+        return IntStream.range(start, end)
+                .mapToObj(i -> "connector" + i)
+                .collect(Collectors.toList());
+    }
+
+    private static List newTasks(int start, int end) {
+        return newTasks("task", start, end);
+    }
+
+    private static List newTasks(String connectorName, int start, int end) {
+        return IntStream.range(start, end)
+                .mapToObj(i -> new ConnectorTaskId(connectorName, i))
+                .collect(Collectors.toList());
+    }
+
+    private static ClusterConfigState clusterConfigState(long offset,
+                                                         int connectorNum,
+                                                         int taskNum) {
+        return clusterConfigState(offset, 1, connectorNum, taskNum);
+    }
+
+    private static ClusterConfigState clusterConfigState(long offset,
+                                                         int connectorStart,
+                                                         int connectorNum,
+                                                         int taskNum) {
+        int connectorNumEnd = connectorStart + connectorNum - 1;
+        return new ClusterConfigState(
+                offset,
+                null,
+                connectorTaskCounts(connectorStart, connectorNumEnd, taskNum),
+                connectorConfigs(connectorStart, connectorNumEnd),
+                connectorTargetStates(connectorStart, connectorNumEnd, TargetState.STARTED),
+                taskConfigs(0, connectorNum, connectorNum * taskNum),
+                Collections.emptySet());
+    }
+
+    private static Map memberConfigs(String givenLeader,
+                                                                  long givenOffset,
+                                                                  Map givenAssignments) {
+        return givenAssignments.entrySet().stream()
+                .collect(Collectors.toMap(
+                    Map.Entry::getKey,
+                    e -> new ExtendedWorkerState(expectedLeaderUrl(givenLeader), givenOffset, e.getValue())));
+    }
+
+    private static Map memberConfigs(String givenLeader,
+                                                                  long givenOffset,
+                                                                  int start,
+                                                                  int connectorNum) {
+        return IntStream.range(start, connectorNum + 1)
+                .mapToObj(i -> new SimpleEntry<>("worker" + i, new ExtendedWorkerState(expectedLeaderUrl(givenLeader), givenOffset, null)))
+                .collect(Collectors.toMap(SimpleEntry::getKey, SimpleEntry::getValue));
+    }
+
+    private static Map connectorTaskCounts(int start,
+                                                            int connectorNum,
+                                                            int taskCounts) {
+        return IntStream.range(start, connectorNum + 1)
+                .mapToObj(i -> new SimpleEntry<>("connector" + i, taskCounts))
+                .collect(Collectors.toMap(SimpleEntry::getKey, SimpleEntry::getValue));
+    }
+
+    private static Map> connectorConfigs(int start, int connectorNum) {
+        return IntStream.range(start, connectorNum + 1)
+                .mapToObj(i -> new SimpleEntry<>("connector" + i, new HashMap()))
+                .collect(Collectors.toMap(SimpleEntry::getKey, SimpleEntry::getValue));
+    }
+
+    private static Map connectorTargetStates(int start,
+                                                                  int connectorNum,
+                                                                  TargetState state) {
+        return IntStream.range(start, connectorNum + 1)
+                .mapToObj(i -> new SimpleEntry<>("connector" + i, state))
+                .collect(Collectors.toMap(SimpleEntry::getKey, SimpleEntry::getValue));
+    }
+
+    private static Map> taskConfigs(int start,
+                                                                         int connectorNum,
+                                                                         int taskNum) {
+        return IntStream.range(start, taskNum + 1)
+                .mapToObj(i -> new SimpleEntry<>(
+                        new ConnectorTaskId("connector" + i / connectorNum + 1, i),
+                        new HashMap())
+                ).collect(Collectors.toMap(SimpleEntry::getKey, SimpleEntry::getValue));
+    }
+
+    private void applyAssignments(Map newAssignments) {
+        newAssignments.forEach((k, v) -> {
+            assignments.computeIfAbsent(k, noop -> newExpandableAssignment())
+                    .connectors()
+                    .removeAll(v.revokedConnectors());
+            assignments.computeIfAbsent(k, noop -> newExpandableAssignment())
+                    .connectors()
+                    .addAll(v.connectors());
+            assignments.computeIfAbsent(k, noop -> newExpandableAssignment())
+                    .tasks()
+                    .removeAll(v.revokedTasks());
+            assignments.computeIfAbsent(k, noop -> newExpandableAssignment())
+                    .tasks()
+                    .addAll(v.tasks());
+        });
+    }
+
+    private ExtendedAssignment newExpandableAssignment() {
+        return new ExtendedAssignment(
+                protocolVersion,
+                ConnectProtocol.Assignment.NO_ERROR,
+                leader,
+                leaderUrl,
+                offset,
+                new ArrayList<>(),
+                new ArrayList<>(),
+                new ArrayList<>(),
+                new ArrayList<>(),
+                0);
+    }
+
+    private static String expectedLeaderUrl(String givenLeader) {
+        return "http://" + givenLeader + ":8083";
+    }
+
+    private void assertAssignment(int connectorNum, int taskNum,
+                                  int revokedConnectorNum, int revokedTaskNum,
+                                  String... workers) {
+        assertAssignment(leader, connectorNum, taskNum, revokedConnectorNum, revokedTaskNum, workers);
+    }
+
+    private void assertAssignment(String expectedLeader, int connectorNum, int taskNum,
+                                  int revokedConnectorNum, int revokedTaskNum,
+                                  String... workers) {
+        assertThat("Wrong number of workers",
+                expectedMemberConfigs.keySet().size(),
+                is(workers.length));
+        assertThat("Wrong set of workers",
+                new ArrayList<>(expectedMemberConfigs.keySet()), hasItems(workers));
+        assertThat("Wrong number of assigned connectors",
+                expectedMemberConfigs.values().stream().map(v -> v.assignment().connectors().size()).reduce(0, Integer::sum),
+                is(connectorNum));
+        assertThat("Wrong number of assigned tasks",
+                expectedMemberConfigs.values().stream().map(v -> v.assignment().tasks().size()).reduce(0, Integer::sum),
+                is(taskNum));
+        assertThat("Wrong number of revoked connectors",
+                expectedMemberConfigs.values().stream().map(v -> v.assignment().revokedConnectors().size()).reduce(0, Integer::sum),
+                is(revokedConnectorNum));
+        assertThat("Wrong number of revoked tasks",
+                expectedMemberConfigs.values().stream().map(v -> v.assignment().revokedTasks().size()).reduce(0, Integer::sum),
+                is(revokedTaskNum));
+        assertThat("Wrong leader in assignments",
+                expectedMemberConfigs.values().stream().map(v -> v.assignment().leader()).distinct().collect(Collectors.joining(", ")),
+                is(expectedLeader));
+        assertThat("Wrong leaderUrl in assignments",
+                expectedMemberConfigs.values().stream().map(v -> v.assignment().leaderUrl()).distinct().collect(Collectors.joining(", ")),
+                is(expectedLeaderUrl(expectedLeader)));
+    }
+
+    private void assertDelay(int expectedDelay, Map newAssignments) {
+        newAssignments.values().stream()
+                .forEach(a -> assertEquals(
+                        "Wrong rebalance delay in " + a, expectedDelay, a.delay()));
+    }
+
+    private void assertNoReassignments(Map existingAssignments,
+                                       Map newAssignments) {
+        assertNoDuplicateInAssignment(existingAssignments);
+        assertNoDuplicateInAssignment(newAssignments);
+
+        List existingConnectors = existingAssignments.values().stream()
+                .flatMap(a -> a.assignment().connectors().stream())
+                .collect(Collectors.toList());
+        List newConnectors = newAssignments.values().stream()
+                .flatMap(a -> a.assignment().connectors().stream())
+                .collect(Collectors.toList());
+
+        List existingTasks = existingAssignments.values().stream()
+                .flatMap(a -> a.assignment().tasks().stream())
+                .collect(Collectors.toList());
+
+        List newTasks = newAssignments.values().stream()
+                .flatMap(a -> a.assignment().tasks().stream())
+                .collect(Collectors.toList());
+
+        existingConnectors.retainAll(newConnectors);
+        assertThat("Found connectors in new assignment that already exist in current assignment",
+                Collections.emptyList(),
+                is(existingConnectors));
+        existingTasks.retainAll(newTasks);
+        assertThat("Found tasks in new assignment that already exist in current assignment",
+                Collections.emptyList(),
+                is(existingConnectors));
+    }
+
+    private void assertNoDuplicateInAssignment(Map existingAssignment) {
+        List existingConnectors = existingAssignment.values().stream()
+                .flatMap(a -> a.assignment().connectors().stream())
+                .collect(Collectors.toList());
+        Set existingUniqueConnectors = new HashSet<>(existingConnectors);
+        existingConnectors.removeAll(existingUniqueConnectors);
+        assertThat("Connectors should be unique in assignments but duplicates where found",
+                Collections.emptyList(),
+                is(existingConnectors));
+
+        List existingTasks = existingAssignment.values().stream()
+                .flatMap(a -> a.assignment().tasks().stream())
+                .collect(Collectors.toList());
+        Set existingUniqueTasks = new HashSet<>(existingTasks);
+        existingTasks.removeAll(existingUniqueTasks);
+        assertThat("Tasks should be unique in assignments but duplicates where found",
+                Collections.emptyList(),
+                is(existingTasks));
+    }
+
+    private void expectGeneration() {
+        when(coordinator.generationId())
+                .thenReturn(assignor.previousGenerationId + 1)
+                .thenReturn(assignor.previousGenerationId + 1);
+        when(coordinator.lastCompletedGenerationId()).thenReturn(assignor.previousGenerationId);
+    }
+}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/distributed/WorkerCoordinatorIncrementalTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/distributed/WorkerCoordinatorIncrementalTest.java
new file mode 100644
index 0000000000000..04c764cb89f1e
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/distributed/WorkerCoordinatorIncrementalTest.java
@@ -0,0 +1,583 @@
+/*
+ * 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.distributed;
+
+import org.apache.kafka.clients.GroupRebalanceConfig;
+import org.apache.kafka.clients.Metadata;
+import org.apache.kafka.clients.MockClient;
+import org.apache.kafka.clients.consumer.internals.ConsumerNetworkClient;
+import org.apache.kafka.common.Node;
+import org.apache.kafka.common.internals.ClusterResourceListeners;
+import org.apache.kafka.common.metrics.Metrics;
+import org.apache.kafka.common.utils.LogContext;
+import org.apache.kafka.common.utils.MockTime;
+import org.apache.kafka.connect.storage.KafkaConfigBackingStore;
+import org.apache.kafka.connect.util.ConnectorTaskId;
+import org.apache.kafka.test.TestUtils;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+import org.mockito.Mock;
+import org.mockito.junit.MockitoJUnit;
+import org.mockito.junit.MockitoRule;
+
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+
+import static org.apache.kafka.common.message.JoinGroupRequestData.JoinGroupRequestProtocol;
+import static org.apache.kafka.common.message.JoinGroupRequestData.JoinGroupRequestProtocolCollection;
+import static org.apache.kafka.common.message.JoinGroupResponseData.JoinGroupResponseMember;
+import static org.apache.kafka.connect.runtime.WorkerTestUtils.assertAssignment;
+import static org.apache.kafka.connect.runtime.WorkerTestUtils.clusterConfigState;
+import static org.apache.kafka.connect.runtime.distributed.ConnectProtocol.WorkerState;
+import static org.apache.kafka.connect.runtime.distributed.ConnectProtocolCompatibility.COMPATIBLE;
+import static org.apache.kafka.connect.runtime.distributed.ConnectProtocolCompatibility.EAGER;
+import static org.apache.kafka.connect.runtime.distributed.ConnectProtocolCompatibility.SESSIONED;
+import static org.apache.kafka.connect.runtime.distributed.IncrementalCooperativeConnectProtocol.CONNECT_PROTOCOL_V1;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotEquals;
+import static org.junit.Assert.assertTrue;
+import static org.junit.runners.Parameterized.Parameter;
+import static org.junit.runners.Parameterized.Parameters;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
+import static org.mockito.Mockito.when;
+
+@RunWith(value = Parameterized.class)
+public class WorkerCoordinatorIncrementalTest {
+    @Rule
+    public MockitoRule rule = MockitoJUnit.rule();
+
+    private String connectorId1 = "connector1";
+    private String connectorId2 = "connector2";
+    private String connectorId3 = "connector3";
+    private ConnectorTaskId taskId1x0 = new ConnectorTaskId(connectorId1, 0);
+    private ConnectorTaskId taskId1x1 = new ConnectorTaskId(connectorId1, 1);
+    private ConnectorTaskId taskId2x0 = new ConnectorTaskId(connectorId2, 0);
+    private ConnectorTaskId taskId3x0 = new ConnectorTaskId(connectorId3, 0);
+
+    private String groupId = "test-group";
+    private int sessionTimeoutMs = 10;
+    private int rebalanceTimeoutMs = 60;
+    private int heartbeatIntervalMs = 2;
+    private long retryBackoffMs = 100;
+    private int requestTimeoutMs = 1000;
+    private MockTime time;
+    private MockClient client;
+    private Node node;
+    private Metadata metadata;
+    private Metrics metrics;
+    private ConsumerNetworkClient consumerClient;
+    private MockRebalanceListener rebalanceListener;
+    @Mock
+    private KafkaConfigBackingStore configStorage;
+    private GroupRebalanceConfig rebalanceConfig;
+    private WorkerCoordinator coordinator;
+    private int rebalanceDelay = DistributedConfig.SCHEDULED_REBALANCE_MAX_DELAY_MS_DEFAULT;
+
+    private String leaderId;
+    private String memberId;
+    private String anotherMemberId;
+    private String leaderUrl;
+    private String memberUrl;
+    private String anotherMemberUrl;
+    private int generationId;
+    private long offset;
+
+    private int configStorageCalls;
+
+    private ClusterConfigState configState1;
+    private ClusterConfigState configState2;
+    private ClusterConfigState configStateSingleTaskConnectors;
+
+    // Arguments are:
+    // - Protocol type
+    // - Expected metadata size
+    @Parameters
+    public static Iterable mode() {
+        return Arrays.asList(new Object[][]{{COMPATIBLE, 2}, {SESSIONED, 3}});
+    }
+
+    @Parameter
+    public ConnectProtocolCompatibility compatibility;
+
+    @Parameter(1)
+    public int expectedMetadataSize;
+
+    @Before
+    public void setup() {
+        LogContext loggerFactory = new LogContext();
+
+        this.time = new MockTime();
+        this.metadata = new Metadata(0, Long.MAX_VALUE, loggerFactory, new ClusterResourceListeners());
+        this.client = new MockClient(time, metadata);
+        this.client.updateMetadata(TestUtils.metadataUpdateWith(1, Collections.singletonMap("topic", 1)));
+        this.node = metadata.fetch().nodes().get(0);
+        this.consumerClient = new ConsumerNetworkClient(loggerFactory, client, metadata, time,
+                retryBackoffMs, requestTimeoutMs, heartbeatIntervalMs);
+        this.metrics = new Metrics(time);
+        this.rebalanceListener = new MockRebalanceListener();
+
+        this.leaderId = "worker1";
+        this.memberId = "worker2";
+        this.anotherMemberId = "worker3";
+        this.leaderUrl = expectedUrl(leaderId);
+        this.memberUrl = expectedUrl(memberId);
+        this.anotherMemberUrl = expectedUrl(anotherMemberId);
+        this.generationId = 3;
+        this.offset = 10L;
+
+        this.configStorageCalls = 0;
+
+        this.rebalanceConfig = new GroupRebalanceConfig(sessionTimeoutMs,
+                                                        rebalanceTimeoutMs,
+                                                        heartbeatIntervalMs,
+                                                        groupId,
+                                                        Optional.empty(),
+                                                        retryBackoffMs,
+                                                        true);
+        this.coordinator = new WorkerCoordinator(rebalanceConfig,
+                                                 loggerFactory,
+                                                 consumerClient,
+                                                 metrics,
+                                                 "worker" + groupId,
+                                                 time,
+                                                 expectedUrl(leaderId),
+                                                 configStorage,
+                                                 rebalanceListener,
+                                                 compatibility,
+                                                 rebalanceDelay);
+
+        configState1 = clusterConfigState(offset, 2, 4);
+    }
+
+    @After
+    public void teardown() {
+        this.metrics.close();
+        verifyNoMoreInteractions(configStorage);
+    }
+
+    private static String expectedUrl(String member) {
+        return "http://" + member + ":8083";
+    }
+
+    // We only test functionality unique to WorkerCoordinator. Other functionality is already
+    // well tested via the tests that cover AbstractCoordinator & ConsumerCoordinator.
+
+    @Test
+    public void testMetadata() {
+        when(configStorage.snapshot()).thenReturn(configState1);
+
+        JoinGroupRequestProtocolCollection serialized = coordinator.metadata();
+        assertEquals(expectedMetadataSize, serialized.size());
+
+        Iterator protocolIterator = serialized.iterator();
+        assertTrue(protocolIterator.hasNext());
+        JoinGroupRequestProtocol defaultMetadata = protocolIterator.next();
+        assertEquals(compatibility.protocol(), defaultMetadata.name());
+        WorkerState state = IncrementalCooperativeConnectProtocol
+                .deserializeMetadata(ByteBuffer.wrap(defaultMetadata.metadata()));
+        assertEquals(offset, state.offset());
+
+        verify(configStorage, times(1)).snapshot();
+    }
+
+    @Test
+    public void testMetadataWithExistingAssignment() {
+        when(configStorage.snapshot()).thenReturn(configState1);
+
+        ExtendedAssignment assignment = new ExtendedAssignment(
+                CONNECT_PROTOCOL_V1, ExtendedAssignment.NO_ERROR, leaderId, leaderUrl, configState1.offset(),
+                Collections.singletonList(connectorId1), Arrays.asList(taskId1x0, taskId2x0),
+                Collections.emptyList(), Collections.emptyList(), 0);
+        ByteBuffer buf = IncrementalCooperativeConnectProtocol.serializeAssignment(assignment);
+        // Using onJoinComplete to register the protocol selection decided by the the broker
+        // coordinator as well as an existing previous assignment that the call to metadata will
+        // include with v1 but not with v0
+        coordinator.onJoinComplete(generationId, memberId, compatibility.protocol(), buf);
+
+        JoinGroupRequestProtocolCollection serialized = coordinator.metadata();
+        assertEquals(expectedMetadataSize, serialized.size());
+
+        Iterator protocolIterator = serialized.iterator();
+        assertTrue(protocolIterator.hasNext());
+        JoinGroupRequestProtocol selectedMetadata = protocolIterator.next();
+        assertEquals(compatibility.protocol(), selectedMetadata.name());
+        ExtendedWorkerState state = IncrementalCooperativeConnectProtocol
+                .deserializeMetadata(ByteBuffer.wrap(selectedMetadata.metadata()));
+        assertEquals(offset, state.offset());
+        assertNotEquals(ExtendedAssignment.empty(), state.assignment());
+        assertEquals(Collections.singletonList(connectorId1), state.assignment().connectors());
+        assertEquals(Arrays.asList(taskId1x0, taskId2x0), state.assignment().tasks());
+
+        verify(configStorage, times(1)).snapshot();
+    }
+
+    @Test
+    public void testMetadataWithExistingAssignmentButOlderProtocolSelection() {
+        when(configStorage.snapshot()).thenReturn(configState1);
+
+        ExtendedAssignment assignment = new ExtendedAssignment(
+                CONNECT_PROTOCOL_V1, ExtendedAssignment.NO_ERROR, leaderId, leaderUrl, configState1.offset(),
+                Collections.singletonList(connectorId1), Arrays.asList(taskId1x0, taskId2x0),
+                Collections.emptyList(), Collections.emptyList(), 0);
+        ByteBuffer buf = IncrementalCooperativeConnectProtocol.serializeAssignment(assignment);
+        // Using onJoinComplete to register the protocol selection decided by the the broker
+        // coordinator as well as an existing previous assignment that the call to metadata will
+        // include with v1 but not with v0
+        coordinator.onJoinComplete(generationId, memberId, EAGER.protocol(), buf);
+
+        JoinGroupRequestProtocolCollection serialized = coordinator.metadata();
+        assertEquals(expectedMetadataSize, serialized.size());
+
+        Iterator protocolIterator = serialized.iterator();
+        assertTrue(protocolIterator.hasNext());
+        JoinGroupRequestProtocol selectedMetadata = protocolIterator.next();
+        assertEquals(compatibility.protocol(), selectedMetadata.name());
+        ExtendedWorkerState state = IncrementalCooperativeConnectProtocol
+                .deserializeMetadata(ByteBuffer.wrap(selectedMetadata.metadata()));
+        assertEquals(offset, state.offset());
+        assertNotEquals(ExtendedAssignment.empty(), state.assignment());
+
+        verify(configStorage, times(1)).snapshot();
+    }
+
+    @Test
+    public void testTaskAssignmentWhenWorkerJoins() {
+        when(configStorage.snapshot()).thenReturn(configState1);
+
+        coordinator.metadata();
+        ++configStorageCalls;
+
+        List responseMembers = new ArrayList<>();
+        addJoinGroupResponseMember(responseMembers, leaderId, offset, null);
+        addJoinGroupResponseMember(responseMembers, memberId, offset, null);
+
+        Map result = coordinator.performAssignment(leaderId, compatibility.protocol(), responseMembers);
+
+        ExtendedAssignment leaderAssignment = deserializeAssignment(result, leaderId);
+        assertAssignment(leaderId, offset,
+                Collections.singletonList(connectorId1), 4,
+                Collections.emptyList(), 0,
+                leaderAssignment);
+
+        ExtendedAssignment memberAssignment = deserializeAssignment(result, memberId);
+        assertAssignment(leaderId, offset,
+                Collections.singletonList(connectorId2), 4,
+                Collections.emptyList(), 0,
+                memberAssignment);
+
+        coordinator.metadata();
+        ++configStorageCalls;
+
+        responseMembers = new ArrayList<>();
+        addJoinGroupResponseMember(responseMembers, leaderId, offset, leaderAssignment);
+        addJoinGroupResponseMember(responseMembers, memberId, offset, memberAssignment);
+        addJoinGroupResponseMember(responseMembers, anotherMemberId, offset, null);
+
+        result = coordinator.performAssignment(leaderId, compatibility.protocol(), responseMembers);
+
+        leaderAssignment = deserializeAssignment(result, leaderId);
+        assertAssignment(leaderId, offset,
+                Collections.emptyList(), 0,
+                Collections.emptyList(), 2,
+                leaderAssignment);
+
+        memberAssignment = deserializeAssignment(result, memberId);
+        assertAssignment(leaderId, offset,
+                Collections.emptyList(), 0,
+                Collections.emptyList(), 0,
+                memberAssignment);
+
+        ExtendedAssignment anotherMemberAssignment = deserializeAssignment(result, anotherMemberId);
+        assertAssignment(leaderId, offset,
+                Collections.emptyList(), 0,
+                Collections.emptyList(), 0,
+                anotherMemberAssignment);
+
+        verify(configStorage, times(configStorageCalls)).snapshot();
+    }
+
+    @Test
+    public void testTaskAssignmentWhenWorkerLeavesPermanently() {
+        when(configStorage.snapshot()).thenReturn(configState1);
+
+        // First assignment distributes configured connectors and tasks
+        coordinator.metadata();
+        ++configStorageCalls;
+
+        List responseMembers = new ArrayList<>();
+        addJoinGroupResponseMember(responseMembers, leaderId, offset, null);
+        addJoinGroupResponseMember(responseMembers, memberId, offset, null);
+        addJoinGroupResponseMember(responseMembers, anotherMemberId, offset, null);
+
+        Map result = coordinator.performAssignment(leaderId, compatibility.protocol(), responseMembers);
+
+        ExtendedAssignment leaderAssignment = deserializeAssignment(result, leaderId);
+        assertAssignment(leaderId, offset,
+                Collections.singletonList(connectorId1), 3,
+                Collections.emptyList(), 0,
+                leaderAssignment);
+
+        ExtendedAssignment memberAssignment = deserializeAssignment(result, memberId);
+        assertAssignment(leaderId, offset,
+                Collections.singletonList(connectorId2), 3,
+                Collections.emptyList(), 0,
+                memberAssignment);
+
+        ExtendedAssignment anotherMemberAssignment = deserializeAssignment(result, anotherMemberId);
+        assertAssignment(leaderId, offset,
+                Collections.emptyList(), 2,
+                Collections.emptyList(), 0,
+                anotherMemberAssignment);
+
+        // Second rebalance detects a worker is missing
+        coordinator.metadata();
+        ++configStorageCalls;
+
+        // Mark everyone as in sync with configState1
+        responseMembers = new ArrayList<>();
+        addJoinGroupResponseMember(responseMembers, leaderId, offset, leaderAssignment);
+        addJoinGroupResponseMember(responseMembers, memberId, offset, memberAssignment);
+
+        result = coordinator.performAssignment(leaderId, compatibility.protocol(), responseMembers);
+
+        leaderAssignment = deserializeAssignment(result, leaderId);
+        assertAssignment(leaderId, offset,
+                Collections.emptyList(), 0,
+                Collections.emptyList(), 0,
+                rebalanceDelay,
+                leaderAssignment);
+
+        memberAssignment = deserializeAssignment(result, memberId);
+        assertAssignment(leaderId, offset,
+                Collections.emptyList(), 0,
+                Collections.emptyList(), 0,
+                rebalanceDelay,
+                memberAssignment);
+
+        rebalanceDelay /= 2;
+        time.sleep(rebalanceDelay);
+
+        // A third rebalance before the delay expires won't change the assignments
+        result = coordinator.performAssignment(leaderId, compatibility.protocol(), responseMembers);
+
+        leaderAssignment = deserializeAssignment(result, leaderId);
+        assertAssignment(leaderId, offset,
+                Collections.emptyList(), 0,
+                Collections.emptyList(), 0,
+                rebalanceDelay,
+                leaderAssignment);
+
+        memberAssignment = deserializeAssignment(result, memberId);
+        assertAssignment(leaderId, offset,
+                Collections.emptyList(), 0,
+                Collections.emptyList(), 0,
+                rebalanceDelay,
+                memberAssignment);
+
+        time.sleep(rebalanceDelay + 1);
+
+        // A rebalance after the delay expires re-assigns the lost tasks
+        result = coordinator.performAssignment(leaderId, compatibility.protocol(), responseMembers);
+
+        leaderAssignment = deserializeAssignment(result, leaderId);
+        assertAssignment(leaderId, offset,
+                Collections.emptyList(), 1,
+                Collections.emptyList(), 0,
+                leaderAssignment);
+
+        memberAssignment = deserializeAssignment(result, memberId);
+        assertAssignment(leaderId, offset,
+                Collections.emptyList(), 1,
+                Collections.emptyList(), 0,
+                memberAssignment);
+
+        verify(configStorage, times(configStorageCalls)).snapshot();
+    }
+
+    @Test
+    public void testTaskAssignmentWhenWorkerBounces() {
+        when(configStorage.snapshot()).thenReturn(configState1);
+
+        // First assignment distributes configured connectors and tasks
+        coordinator.metadata();
+        ++configStorageCalls;
+
+        List responseMembers = new ArrayList<>();
+        addJoinGroupResponseMember(responseMembers, leaderId, offset, null);
+        addJoinGroupResponseMember(responseMembers, memberId, offset, null);
+        addJoinGroupResponseMember(responseMembers, anotherMemberId, offset, null);
+
+        Map result = coordinator.performAssignment(leaderId, compatibility.protocol(), responseMembers);
+
+        ExtendedAssignment leaderAssignment = deserializeAssignment(result, leaderId);
+        assertAssignment(leaderId, offset,
+                Collections.singletonList(connectorId1), 3,
+                Collections.emptyList(), 0,
+                leaderAssignment);
+
+        ExtendedAssignment memberAssignment = deserializeAssignment(result, memberId);
+        assertAssignment(leaderId, offset,
+                Collections.singletonList(connectorId2), 3,
+                Collections.emptyList(), 0,
+                memberAssignment);
+
+        ExtendedAssignment anotherMemberAssignment = deserializeAssignment(result, anotherMemberId);
+        assertAssignment(leaderId, offset,
+                Collections.emptyList(), 2,
+                Collections.emptyList(), 0,
+                anotherMemberAssignment);
+
+        // Second rebalance detects a worker is missing
+        coordinator.metadata();
+        ++configStorageCalls;
+
+        responseMembers = new ArrayList<>();
+        addJoinGroupResponseMember(responseMembers, leaderId, offset, leaderAssignment);
+        addJoinGroupResponseMember(responseMembers, memberId, offset, memberAssignment);
+
+        result = coordinator.performAssignment(leaderId, compatibility.protocol(), responseMembers);
+
+        leaderAssignment = deserializeAssignment(result, leaderId);
+        assertAssignment(leaderId, offset,
+                Collections.emptyList(), 0,
+                Collections.emptyList(), 0,
+                rebalanceDelay,
+                leaderAssignment);
+
+        memberAssignment = deserializeAssignment(result, memberId);
+        assertAssignment(leaderId, offset,
+                Collections.emptyList(), 0,
+                Collections.emptyList(), 0,
+                rebalanceDelay,
+                memberAssignment);
+
+        rebalanceDelay /= 2;
+        time.sleep(rebalanceDelay);
+
+        // A third rebalance before the delay expires won't change the assignments even if the
+        // member returns in the meantime
+        addJoinGroupResponseMember(responseMembers, anotherMemberId, offset, null);
+        result = coordinator.performAssignment(leaderId, compatibility.protocol(), responseMembers);
+
+        leaderAssignment = deserializeAssignment(result, leaderId);
+        assertAssignment(leaderId, offset,
+                Collections.emptyList(), 0,
+                Collections.emptyList(), 0,
+                rebalanceDelay,
+                leaderAssignment);
+
+        memberAssignment = deserializeAssignment(result, memberId);
+        assertAssignment(leaderId, offset,
+                Collections.emptyList(), 0,
+                Collections.emptyList(), 0,
+                rebalanceDelay,
+                memberAssignment);
+
+        anotherMemberAssignment = deserializeAssignment(result, anotherMemberId);
+        assertAssignment(leaderId, offset,
+                Collections.emptyList(), 0,
+                Collections.emptyList(), 0,
+                rebalanceDelay,
+                anotherMemberAssignment);
+
+        time.sleep(rebalanceDelay + 1);
+
+        result = coordinator.performAssignment(leaderId, compatibility.protocol(), responseMembers);
+
+        // A rebalance after the delay expires re-assigns the lost tasks the the returning member
+        leaderAssignment = deserializeAssignment(result, leaderId);
+        assertAssignment(leaderId, offset,
+                Collections.emptyList(), 0,
+                Collections.emptyList(), 0,
+                leaderAssignment);
+
+        memberAssignment = deserializeAssignment(result, memberId);
+        assertAssignment(leaderId, offset,
+                Collections.emptyList(), 0,
+                Collections.emptyList(), 0,
+                memberAssignment);
+
+        anotherMemberAssignment = deserializeAssignment(result, anotherMemberId);
+        assertAssignment(leaderId, offset,
+                Collections.emptyList(), 2,
+                Collections.emptyList(), 0,
+                anotherMemberAssignment);
+
+        verify(configStorage, times(configStorageCalls)).snapshot();
+    }
+
+    private static class MockRebalanceListener implements WorkerRebalanceListener {
+        public ExtendedAssignment assignment = null;
+
+        public String revokedLeader;
+        public Collection revokedConnectors = Collections.emptyList();
+        public Collection revokedTasks = Collections.emptyList();
+
+        public int revokedCount = 0;
+        public int assignedCount = 0;
+
+        @Override
+        public void onAssigned(ExtendedAssignment assignment, int generation) {
+            this.assignment = assignment;
+            assignedCount++;
+        }
+
+        @Override
+        public void onRevoked(String leader, Collection connectors, Collection tasks) {
+            if (connectors.isEmpty() && tasks.isEmpty()) {
+                return;
+            }
+            this.revokedLeader = leader;
+            this.revokedConnectors = connectors;
+            this.revokedTasks = tasks;
+            revokedCount++;
+        }
+    }
+
+    private static ExtendedAssignment deserializeAssignment(Map assignment,
+                                                           String member) {
+        return IncrementalCooperativeConnectProtocol.deserializeAssignment(assignment.get(member));
+    }
+
+    private void addJoinGroupResponseMember(List responseMembers,
+                                                   String member,
+                                                   long offset,
+                                                   ExtendedAssignment assignment) {
+        responseMembers.add(new JoinGroupResponseMember()
+                .setMemberId(member)
+                .setMetadata(
+                    IncrementalCooperativeConnectProtocol.serializeMetadata(
+                        new ExtendedWorkerState(expectedUrl(member), offset, assignment),
+                        compatibility != COMPATIBLE
+                    ).array()
+                )
+        );
+    }
+}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/distributed/WorkerCoordinatorTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/distributed/WorkerCoordinatorTest.java
index 68d236f14aeb7..05e83fb80231f 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/distributed/WorkerCoordinatorTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/distributed/WorkerCoordinatorTest.java
@@ -16,21 +16,25 @@
  */
 package org.apache.kafka.connect.runtime.distributed;
 
+import org.apache.kafka.clients.GroupRebalanceConfig;
 import org.apache.kafka.clients.Metadata;
 import org.apache.kafka.clients.MockClient;
 import org.apache.kafka.clients.consumer.internals.ConsumerNetworkClient;
 import org.apache.kafka.common.Node;
 import org.apache.kafka.common.internals.ClusterResourceListeners;
+import org.apache.kafka.common.message.JoinGroupRequestData;
+import org.apache.kafka.common.message.JoinGroupResponseData;
+import org.apache.kafka.common.message.SyncGroupResponseData;
 import org.apache.kafka.common.metrics.Metrics;
 import org.apache.kafka.common.protocol.Errors;
 import org.apache.kafka.common.requests.AbstractRequest;
 import org.apache.kafka.common.requests.FindCoordinatorResponse;
-import org.apache.kafka.common.requests.JoinGroupRequest.ProtocolMetadata;
 import org.apache.kafka.common.requests.JoinGroupResponse;
 import org.apache.kafka.common.requests.SyncGroupRequest;
 import org.apache.kafka.common.requests.SyncGroupResponse;
 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.connect.runtime.TargetState;
 import org.apache.kafka.connect.storage.KafkaConfigBackingStore;
 import org.apache.kafka.connect.util.ConnectorTaskId;
@@ -40,20 +44,28 @@
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
 import org.powermock.api.easymock.PowerMock;
-import org.powermock.reflect.Whitebox;
 
 import java.nio.ByteBuffer;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import static org.junit.runners.Parameterized.Parameter;
+import static org.junit.runners.Parameterized.Parameters;
 
+@RunWith(value = Parameterized.class)
 public class WorkerCoordinatorTest {
 
     private static final String LEADER_URL = "leaderUrl:8083";
@@ -80,12 +92,29 @@ public class WorkerCoordinatorTest {
     private ConsumerNetworkClient consumerClient;
     private MockRebalanceListener rebalanceListener;
     @Mock private KafkaConfigBackingStore configStorage;
+    private GroupRebalanceConfig rebalanceConfig;
     private WorkerCoordinator coordinator;
 
     private ClusterConfigState configState1;
     private ClusterConfigState configState2;
     private ClusterConfigState configStateSingleTaskConnectors;
 
+    // Arguments are:
+    // - Protocol type
+    // - Expected metadata size
+    @Parameters
+    public static Iterable mode() {
+        return Arrays.asList(new Object[][]{
+                {ConnectProtocolCompatibility.EAGER, 1},
+                {ConnectProtocolCompatibility.COMPATIBLE, 2}});
+    }
+
+    @Parameter
+    public ConnectProtocolCompatibility compatibility;
+
+    @Parameter(1)
+    public int expectedMetadataSize;
+
     @Before
     public void setup() {
         LogContext logContext = new LogContext();
@@ -99,24 +128,28 @@ public void setup() {
         this.metrics = new Metrics(time);
         this.rebalanceListener = new MockRebalanceListener();
         this.configStorage = PowerMock.createMock(KafkaConfigBackingStore.class);
-
-        this.coordinator = new WorkerCoordinator(
-                logContext,
-                consumerClient,
-                groupId,
-                rebalanceTimeoutMs,
-                sessionTimeoutMs,
-                heartbeatIntervalMs,
-                metrics,
-                "consumer" + groupId,
-                time,
-                retryBackoffMs,
-                LEADER_URL,
-                configStorage,
-                rebalanceListener);
+        this.rebalanceConfig = new GroupRebalanceConfig(sessionTimeoutMs,
+                                                        rebalanceTimeoutMs,
+                                                        heartbeatIntervalMs,
+                                                        groupId,
+                                                        Optional.empty(),
+                                                        retryBackoffMs,
+                                                        true);
+        this.coordinator = new WorkerCoordinator(rebalanceConfig,
+                                                 logContext,
+                                                 consumerClient,
+                                                 metrics,
+                                                 "consumer" + groupId,
+                                                 time,
+                                                 LEADER_URL,
+                                                 configStorage,
+                                                 rebalanceListener,
+                                                 compatibility,
+                                                 0);
 
         configState1 = new ClusterConfigState(
                 1L,
+                null,
                 Collections.singletonMap(connectorId1, 1),
                 Collections.singletonMap(connectorId1, (Map) new HashMap()),
                 Collections.singletonMap(connectorId1, TargetState.STARTED),
@@ -139,6 +172,7 @@ public void setup() {
         configState2TaskConfigs.put(taskId2x0, new HashMap());
         configState2 = new ClusterConfigState(
                 2L,
+                null,
                 configState2ConnectorTaskCounts,
                 configState2ConnectorConfigs,
                 configState2TargetStates,
@@ -164,6 +198,7 @@ public void setup() {
         configStateSingleTaskConnectorsTaskConfigs.put(taskId3x0, new HashMap());
         configStateSingleTaskConnectors = new ClusterConfigState(
                 2L,
+                null,
                 configStateSingleTaskConnectorsConnectorTaskCounts,
                 configStateSingleTaskConnectorsConnectorConfigs,
                 configStateSingleTaskConnectorsTargetStates,
@@ -186,12 +221,15 @@ public void testMetadata() {
 
         PowerMock.replayAll();
 
-        List serialized = coordinator.metadata();
-        assertEquals(1, serialized.size());
+        JoinGroupRequestData.JoinGroupRequestProtocolCollection serialized = coordinator.metadata();
+        assertEquals(expectedMetadataSize, serialized.size());
 
-        ProtocolMetadata defaultMetadata = serialized.get(0);
-        assertEquals(WorkerCoordinator.DEFAULT_SUBPROTOCOL, defaultMetadata.name());
-        ConnectProtocol.WorkerState state = ConnectProtocol.deserializeMetadata(defaultMetadata.metadata());
+        Iterator protocolIterator = serialized.iterator();
+        assertTrue(protocolIterator.hasNext());
+        JoinGroupRequestData.JoinGroupRequestProtocol defaultMetadata = protocolIterator.next();
+        assertEquals(compatibility.protocol(), defaultMetadata.name());
+        ConnectProtocol.WorkerState state = ConnectProtocol.deserializeMetadata(
+                ByteBuffer.wrap(defaultMetadata.metadata()));
         assertEquals(1, state.offset());
 
         PowerMock.verifyAll();
@@ -205,7 +243,7 @@ public void testNormalJoinGroupLeader() {
 
         final String consumerId = "leader";
 
-        client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+        client.prepareResponse(FindCoordinatorResponse.prepareResponse(Errors.NONE, node));
         coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
         // normal join group
@@ -217,9 +255,9 @@ public void testNormalJoinGroupLeader() {
             @Override
             public boolean matches(AbstractRequest body) {
                 SyncGroupRequest sync = (SyncGroupRequest) body;
-                return sync.memberId().equals(consumerId) &&
-                        sync.generationId() == 1 &&
-                        sync.groupAssignment().containsKey(consumerId);
+                return sync.data.memberId().equals(consumerId) &&
+                        sync.data.generationId() == 1 &&
+                        sync.groupAssignments().containsKey(consumerId);
             }
         }, syncGroupResponse(ConnectProtocol.Assignment.NO_ERROR, "leader", 1L, Collections.singletonList(connectorId1),
                 Collections.emptyList(), Errors.NONE));
@@ -245,7 +283,7 @@ public void testNormalJoinGroupFollower() {
 
         final String memberId = "member";
 
-        client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+        client.prepareResponse(FindCoordinatorResponse.prepareResponse(Errors.NONE, node));
         coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
         // normal join group
@@ -254,9 +292,9 @@ public void testNormalJoinGroupFollower() {
             @Override
             public boolean matches(AbstractRequest body) {
                 SyncGroupRequest sync = (SyncGroupRequest) body;
-                return sync.memberId().equals(memberId) &&
-                        sync.generationId() == 1 &&
-                        sync.groupAssignment().isEmpty();
+                return sync.data.memberId().equals(memberId) &&
+                        sync.data.generationId() == 1 &&
+                        sync.data.assignments().isEmpty();
             }
         }, syncGroupResponse(ConnectProtocol.Assignment.NO_ERROR, "leader", 1L, Collections.emptyList(),
                 Collections.singletonList(taskId1x0), Errors.NONE));
@@ -286,7 +324,7 @@ public void testJoinLeaderCannotAssign() {
 
         final String memberId = "member";
 
-        client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+        client.prepareResponse(FindCoordinatorResponse.prepareResponse(Errors.NONE, node));
         coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
         // config mismatch results in assignment error
@@ -295,9 +333,9 @@ public void testJoinLeaderCannotAssign() {
             @Override
             public boolean matches(AbstractRequest body) {
                 SyncGroupRequest sync = (SyncGroupRequest) body;
-                return sync.memberId().equals(memberId) &&
-                        sync.generationId() == 1 &&
-                        sync.groupAssignment().isEmpty();
+                return sync.data.memberId().equals(memberId) &&
+                        sync.data.generationId() == 1 &&
+                        sync.data.assignments().isEmpty();
             }
         };
         client.prepareResponse(matcher, syncGroupResponse(ConnectProtocol.Assignment.CONFIG_MISMATCH, "leader", 10L,
@@ -317,7 +355,7 @@ public void testRejoinGroup() {
 
         PowerMock.replayAll();
 
-        client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+        client.prepareResponse(FindCoordinatorResponse.prepareResponse(Errors.NONE, node));
         coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
         // join the group once
@@ -364,11 +402,17 @@ public void testLeaderPerformAssignment1() throws Exception {
         // Prime the current configuration state
         coordinator.metadata();
 
-        Map configs = new HashMap<>();
         // Mark everyone as in sync with configState1
-        configs.put("leader", ConnectProtocol.serializeMetadata(new ConnectProtocol.WorkerState(LEADER_URL, 1L)));
-        configs.put("member", ConnectProtocol.serializeMetadata(new ConnectProtocol.WorkerState(MEMBER_URL, 1L)));
-        Map result = Whitebox.invokeMethod(coordinator, "performAssignment", "leader", WorkerCoordinator.DEFAULT_SUBPROTOCOL, configs);
+        List responseMembers = new ArrayList<>();
+        responseMembers.add(new JoinGroupResponseData.JoinGroupResponseMember()
+                .setMemberId("leader")
+                .setMetadata(ConnectProtocol.serializeMetadata(new ConnectProtocol.WorkerState(LEADER_URL, 1L)).array())
+        );
+        responseMembers.add(new JoinGroupResponseData.JoinGroupResponseMember()
+                .setMemberId("member")
+                .setMetadata(ConnectProtocol.serializeMetadata(new ConnectProtocol.WorkerState(MEMBER_URL, 1L)).array())
+        );
+        Map result = coordinator.performAssignment("leader", WorkerCoordinator.DEFAULT_SUBPROTOCOL, responseMembers);
 
         // configState1 has 1 connector, 1 task
         ConnectProtocol.Assignment leaderAssignment = ConnectProtocol.deserializeAssignment(result.get("leader"));
@@ -400,11 +444,18 @@ public void testLeaderPerformAssignment2() throws Exception {
         // Prime the current configuration state
         coordinator.metadata();
 
-        Map configs = new HashMap<>();
         // Mark everyone as in sync with configState1
-        configs.put("leader", ConnectProtocol.serializeMetadata(new ConnectProtocol.WorkerState(LEADER_URL, 1L)));
-        configs.put("member", ConnectProtocol.serializeMetadata(new ConnectProtocol.WorkerState(MEMBER_URL, 1L)));
-        Map result = Whitebox.invokeMethod(coordinator, "performAssignment", "leader", WorkerCoordinator.DEFAULT_SUBPROTOCOL, configs);
+        List responseMembers = new ArrayList<>();
+        responseMembers.add(new JoinGroupResponseData.JoinGroupResponseMember()
+                .setMemberId("leader")
+                .setMetadata(ConnectProtocol.serializeMetadata(new ConnectProtocol.WorkerState(LEADER_URL, 1L)).array())
+        );
+        responseMembers.add(new JoinGroupResponseData.JoinGroupResponseMember()
+                .setMemberId("member")
+                .setMetadata(ConnectProtocol.serializeMetadata(new ConnectProtocol.WorkerState(MEMBER_URL, 1L)).array())
+        );
+
+        Map result = coordinator.performAssignment("leader", WorkerCoordinator.DEFAULT_SUBPROTOCOL, responseMembers);
 
         // configState2 has 2 connector, 3 tasks and should trigger round robin assignment
         ConnectProtocol.Assignment leaderAssignment = ConnectProtocol.deserializeAssignment(result.get("leader"));
@@ -436,11 +487,18 @@ public void testLeaderPerformAssignmentSingleTaskConnectors() throws Exception {
         // Prime the current configuration state
         coordinator.metadata();
 
-        Map configs = new HashMap<>();
         // Mark everyone as in sync with configState1
-        configs.put("leader", ConnectProtocol.serializeMetadata(new ConnectProtocol.WorkerState(LEADER_URL, 1L)));
-        configs.put("member", ConnectProtocol.serializeMetadata(new ConnectProtocol.WorkerState(MEMBER_URL, 1L)));
-        Map result = Whitebox.invokeMethod(coordinator, "performAssignment", "leader", WorkerCoordinator.DEFAULT_SUBPROTOCOL, configs);
+        List responseMembers = new ArrayList<>();
+        responseMembers.add(new JoinGroupResponseData.JoinGroupResponseMember()
+                .setMemberId("leader")
+                .setMetadata(ConnectProtocol.serializeMetadata(new ConnectProtocol.WorkerState(LEADER_URL, 1L)).array())
+        );
+        responseMembers.add(new JoinGroupResponseData.JoinGroupResponseMember()
+                .setMemberId("member")
+                .setMetadata(ConnectProtocol.serializeMetadata(new ConnectProtocol.WorkerState(MEMBER_URL, 1L)).array())
+        );
+
+        Map result = coordinator.performAssignment("leader", WorkerCoordinator.DEFAULT_SUBPROTOCOL, responseMembers);
 
         // Round robin assignment when there are the same number of connectors and tasks should result in each being
         // evenly distributed across the workers, i.e. round robin assignment of connectors first, then followed by tasks
@@ -461,38 +519,53 @@ public void testLeaderPerformAssignmentSingleTaskConnectors() throws Exception {
         PowerMock.verifyAll();
     }
 
-
-    private FindCoordinatorResponse groupCoordinatorResponse(Node node, Errors error) {
-        return new FindCoordinatorResponse(error, node);
-    }
-
     private JoinGroupResponse joinGroupLeaderResponse(int generationId, String memberId,
                                            Map configOffsets, Errors error) {
-        Map metadata = new HashMap<>();
+        List metadata = new ArrayList<>();
         for (Map.Entry configStateEntry : configOffsets.entrySet()) {
             // We need a member URL, but it doesn't matter for the purposes of this test. Just set it to the member ID
             String memberUrl = configStateEntry.getKey();
             long configOffset = configStateEntry.getValue();
             ByteBuffer buf = ConnectProtocol.serializeMetadata(new ConnectProtocol.WorkerState(memberUrl, configOffset));
-            metadata.put(configStateEntry.getKey(), buf);
+            metadata.add(new JoinGroupResponseData.JoinGroupResponseMember()
+                    .setMemberId(configStateEntry.getKey())
+                    .setMetadata(buf.array())
+            );
         }
-        return new JoinGroupResponse(error, generationId, WorkerCoordinator.DEFAULT_SUBPROTOCOL, memberId, memberId, metadata);
+        return new JoinGroupResponse(
+                new JoinGroupResponseData().setErrorCode(error.code())
+                .setGenerationId(generationId)
+                .setProtocolName(WorkerCoordinator.DEFAULT_SUBPROTOCOL)
+                .setLeader(memberId)
+                .setMemberId(memberId)
+                .setMembers(metadata)
+        );
     }
 
     private JoinGroupResponse joinGroupFollowerResponse(int generationId, String memberId, String leaderId, Errors error) {
-        return new JoinGroupResponse(error, generationId, WorkerCoordinator.DEFAULT_SUBPROTOCOL, memberId, leaderId,
-                Collections.emptyMap());
+        return new JoinGroupResponse(
+                new JoinGroupResponseData().setErrorCode(error.code())
+                        .setGenerationId(generationId)
+                        .setProtocolName(WorkerCoordinator.DEFAULT_SUBPROTOCOL)
+                        .setLeader(leaderId)
+                        .setMemberId(memberId)
+                        .setMembers(Collections.emptyList())
+        );
     }
 
     private SyncGroupResponse syncGroupResponse(short assignmentError, String leader, long configOffset, List connectorIds,
                                      List taskIds, Errors error) {
         ConnectProtocol.Assignment assignment = new ConnectProtocol.Assignment(assignmentError, leader, LEADER_URL, configOffset, connectorIds, taskIds);
         ByteBuffer buf = ConnectProtocol.serializeAssignment(assignment);
-        return new SyncGroupResponse(error, buf);
+        return new SyncGroupResponse(
+                new SyncGroupResponseData()
+                        .setErrorCode(error.code())
+                        .setAssignment(Utils.toArray(buf))
+        );
     }
 
     private static class MockRebalanceListener implements WorkerRebalanceListener {
-        public ConnectProtocol.Assignment assignment = null;
+        public ExtendedAssignment assignment = null;
 
         public String revokedLeader;
         public Collection revokedConnectors;
@@ -502,13 +575,16 @@ private static class MockRebalanceListener implements WorkerRebalanceListener {
         public int assignedCount = 0;
 
         @Override
-        public void onAssigned(ConnectProtocol.Assignment assignment, int generation) {
+        public void onAssigned(ExtendedAssignment assignment, int generation) {
             this.assignment = assignment;
             assignedCount++;
         }
 
         @Override
         public void onRevoked(String leader, Collection connectors, Collection tasks) {
+            if (connectors.isEmpty() && tasks.isEmpty()) {
+                return;
+            }
             this.revokedLeader = leader;
             this.revokedConnectors = connectors;
             this.revokedTasks = tasks;
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 00a922f76ad97..f01cd4995c052 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
@@ -146,6 +146,20 @@ public void testReportDLQTwice() {
         PowerMock.verifyAll();
     }
 
+    @Test
+    public void testCloseDLQ() {
+        DeadLetterQueueReporter deadLetterQueueReporter = new DeadLetterQueueReporter(
+            producer, config(singletonMap(SinkConnectorConfig.DLQ_TOPIC_NAME_CONFIG, DLQ_TOPIC)), TASK_ID, errorHandlingMetrics);
+
+        producer.close();
+        EasyMock.expectLastCall();
+        replay(producer);
+
+        deadLetterQueueReporter.close();
+
+        PowerMock.verifyAll();
+    }
+
     @Test
     public void testLogOnDisabledLogReporter() {
         LogReporter logReporter = new LogReporter(TASK_ID, config(emptyMap()), errorHandlingMetrics);
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/health/ConnectClusterStateImplTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/health/ConnectClusterStateImplTest.java
new file mode 100644
index 0000000000000..d8a7e49ee8104
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/health/ConnectClusterStateImplTest.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.connect.runtime.health;
+
+import org.apache.kafka.connect.errors.ConnectException;
+import org.apache.kafka.connect.runtime.Herder;
+import org.apache.kafka.connect.util.Callback;
+import org.easymock.Capture;
+import org.easymock.EasyMock;
+import org.easymock.IAnswer;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.powermock.api.easymock.annotation.Mock;
+import org.powermock.modules.junit4.PowerMockRunner;
+
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Map;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotSame;
+import static org.junit.Assert.assertThrows;
+
+@RunWith(PowerMockRunner.class)
+public class ConnectClusterStateImplTest {
+    protected static final String KAFKA_CLUSTER_ID = "franzwashere";
+
+    @Mock
+    protected Herder herder;
+    protected ConnectClusterStateImpl connectClusterState;
+    protected long herderRequestTimeoutMs = TimeUnit.SECONDS.toMillis(10);
+    protected Collection expectedConnectors;
+    
+    @Before
+    public void setUp() {
+        expectedConnectors = Arrays.asList("sink1", "source1", "source2");
+        connectClusterState = new ConnectClusterStateImpl(
+            herderRequestTimeoutMs,
+            new ConnectClusterDetailsImpl(KAFKA_CLUSTER_ID),
+            herder
+        );
+    }
+    
+    @Test
+    public void connectors() {
+        Capture>> callback = EasyMock.newCapture();
+        herder.connectors(EasyMock.capture(callback));
+        EasyMock.expectLastCall().andAnswer(new IAnswer() {
+            @Override
+            public Void answer() {
+                callback.getValue().onCompletion(null, expectedConnectors);
+                return null;
+            }
+        });
+        EasyMock.replay(herder);
+        assertEquals(expectedConnectors, connectClusterState.connectors());
+    }
+
+    @Test
+    public void connectorConfig() {
+        final String connName = "sink6";
+        final Map expectedConfig = Collections.singletonMap("key", "value");
+        Capture>> callback = EasyMock.newCapture();
+        herder.connectorConfig(EasyMock.eq(connName), EasyMock.capture(callback));
+        EasyMock.expectLastCall().andAnswer(new IAnswer() {
+            @Override
+            public Void answer() {
+                callback.getValue().onCompletion(null, expectedConfig);
+                return null;
+            }
+        });
+        EasyMock.replay(herder);
+        Map actualConfig = connectClusterState.connectorConfig(connName);
+        assertEquals(expectedConfig, actualConfig);
+        assertNotSame(
+            "Config should be copied in order to avoid mutation by REST extensions",
+            expectedConfig,
+            actualConfig
+        );
+    }
+
+    @Test
+    public void kafkaClusterId() {
+        assertEquals(KAFKA_CLUSTER_ID, connectClusterState.clusterDetails().kafkaClusterId());
+    }
+
+    @Test
+    public void connectorsFailure() {
+        Capture>> callback = EasyMock.newCapture();
+        herder.connectors(EasyMock.capture(callback));
+        EasyMock.expectLastCall().andAnswer(new IAnswer() {
+            @Override
+            public Void answer() {
+                Throwable timeout = new TimeoutException();
+                callback.getValue().onCompletion(timeout, null);
+                return null;
+            }
+        });
+        EasyMock.replay(herder);
+        assertThrows(ConnectException.class, connectClusterState::connectors);
+    }
+}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/isolation/DelegatingClassLoaderTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/isolation/DelegatingClassLoaderTest.java
index 5c06eaa7ba11f..3e346bb824623 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/isolation/DelegatingClassLoaderTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/isolation/DelegatingClassLoaderTest.java
@@ -17,8 +17,10 @@
 
 package org.apache.kafka.connect.runtime.isolation;
 
+import java.util.Collections;
 import org.junit.Test;
 
+import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 
@@ -38,4 +40,25 @@ public void testOtherResources() {
             DelegatingClassLoader.serviceLoaderManifestForPlugin("META-INF/services/org.apache.kafka.connect.transforms.Transformation"));
         assertFalse(DelegatingClassLoader.serviceLoaderManifestForPlugin("resource/version.properties"));
     }
+
+    @Test(expected = ClassNotFoundException.class)
+    public void testLoadingUnloadedPluginClass() throws ClassNotFoundException {
+        TestPlugins.assertAvailable();
+        DelegatingClassLoader classLoader = new DelegatingClassLoader(Collections.emptyList());
+        classLoader.initLoaders();
+        for (String pluginClassName : TestPlugins.pluginClasses()) {
+            classLoader.loadClass(pluginClassName);
+        }
+    }
+
+    @Test
+    public void testLoadingPluginClass() throws ClassNotFoundException {
+        TestPlugins.assertAvailable();
+        DelegatingClassLoader classLoader = new DelegatingClassLoader(TestPlugins.pluginPath());
+        classLoader.initLoaders();
+        for (String pluginClassName : TestPlugins.pluginClasses()) {
+            assertNotNull(classLoader.loadClass(pluginClassName));
+            assertNotNull(classLoader.pluginClassLoader(pluginClassName));
+        }
+    }
 }
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/isolation/PluginUtilsTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/isolation/PluginUtilsTest.java
index f9a2d8f4afac5..7ac08cc13bf34 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/isolation/PluginUtilsTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/isolation/PluginUtilsTest.java
@@ -70,7 +70,7 @@ public void testThirdPartyClasses() {
     }
 
     @Test
-    public void testConnectFrameworkClasses() {
+    public void testKafkaDependencyClasses() {
         assertFalse(PluginUtils.shouldLoadInIsolation("org.apache.kafka.common."));
         assertFalse(PluginUtils.shouldLoadInIsolation(
                 "org.apache.kafka.common.config.AbstractConfig")
@@ -81,30 +81,6 @@ public void testConnectFrameworkClasses() {
         assertFalse(PluginUtils.shouldLoadInIsolation(
                 "org.apache.kafka.common.serialization.Deserializer")
         );
-        assertFalse(PluginUtils.shouldLoadInIsolation("org.apache.kafka.connect."));
-        assertFalse(PluginUtils.shouldLoadInIsolation(
-                "org.apache.kafka.connect.connector.Connector")
-        );
-        assertFalse(PluginUtils.shouldLoadInIsolation(
-                "org.apache.kafka.connect.source.SourceConnector")
-        );
-        assertFalse(PluginUtils.shouldLoadInIsolation(
-                "org.apache.kafka.connect.sink.SinkConnector")
-        );
-        assertFalse(PluginUtils.shouldLoadInIsolation("org.apache.kafka.connect.connector.Task"));
-        assertFalse(PluginUtils.shouldLoadInIsolation(
-                "org.apache.kafka.connect.source.SourceTask")
-        );
-        assertFalse(PluginUtils.shouldLoadInIsolation("org.apache.kafka.connect.sink.SinkTask"));
-        assertFalse(PluginUtils.shouldLoadInIsolation(
-                "org.apache.kafka.connect.transforms.Transformation")
-        );
-        assertFalse(PluginUtils.shouldLoadInIsolation(
-                "org.apache.kafka.connect.storage.Converter")
-        );
-        assertFalse(PluginUtils.shouldLoadInIsolation(
-                "org.apache.kafka.connect.storage.OffsetBackingStore")
-        );
         assertFalse(PluginUtils.shouldLoadInIsolation(
                 "org.apache.kafka.clients.producer.ProducerConfig")
         );
@@ -114,62 +90,266 @@ public void testConnectFrameworkClasses() {
         assertFalse(PluginUtils.shouldLoadInIsolation(
                 "org.apache.kafka.clients.admin.KafkaAdminClient")
         );
-        assertFalse(PluginUtils.shouldLoadInIsolation(
-                "org.apache.kafka.connect.rest.ConnectRestExtension")
-        );
     }
 
     @Test
-    public void testAllowedConnectFrameworkClasses() {
-        assertTrue(PluginUtils.shouldLoadInIsolation("org.apache.kafka.connect.transforms."));
-        assertTrue(PluginUtils.shouldLoadInIsolation(
-                "org.apache.kafka.connect.transforms.ExtractField")
-        );
-        assertTrue(PluginUtils.shouldLoadInIsolation(
-                "org.apache.kafka.connect.transforms.ExtractField$Key")
-        );
-        assertTrue(PluginUtils.shouldLoadInIsolation("org.apache.kafka.connect.json."));
-        assertTrue(PluginUtils.shouldLoadInIsolation(
-                "org.apache.kafka.connect.json.JsonConverter")
-        );
-        assertTrue(PluginUtils.shouldLoadInIsolation(
-                "org.apache.kafka.connect.json.JsonConverter$21")
-        );
-        assertTrue(PluginUtils.shouldLoadInIsolation("org.apache.kafka.connect.file."));
-        assertTrue(PluginUtils.shouldLoadInIsolation(
-                "org.apache.kafka.connect.file.FileStreamSourceTask")
-        );
-        assertTrue(PluginUtils.shouldLoadInIsolation(
-                "org.apache.kafka.connect.file.FileStreamSinkConnector")
-        );
-        assertTrue(PluginUtils.shouldLoadInIsolation("org.apache.kafka.connect.converters."));
-        assertTrue(PluginUtils.shouldLoadInIsolation(
-                "org.apache.kafka.connect.converters.ByteArrayConverter")
-        );
-        assertTrue(PluginUtils.shouldLoadInIsolation(
-                "org.apache.kafka.connect.converters.DoubleConverter")
-        );
-        assertTrue(PluginUtils.shouldLoadInIsolation(
-                "org.apache.kafka.connect.converters.FloatConverter")
-        );
-        assertTrue(PluginUtils.shouldLoadInIsolation(
-                "org.apache.kafka.connect.converters.IntegerConverter")
-        );
-        assertTrue(PluginUtils.shouldLoadInIsolation(
-                "org.apache.kafka.connect.converters.LongConverter")
-        );
-        assertTrue(PluginUtils.shouldLoadInIsolation(
-                "org.apache.kafka.connect.converters.ShortConverter")
-        );
+    public void testConnectApiClasses() {
+        List apiClasses = Arrays.asList(
+            // Enumerate all packages and classes
+            "org.apache.kafka.connect.",
+            "org.apache.kafka.connect.components.",
+            "org.apache.kafka.connect.components.Versioned",
+            //"org.apache.kafka.connect.connector.policy.", isolated by default
+            "org.apache.kafka.connect.connector.policy.ConnectorClientConfigOverridePolicy",
+            "org.apache.kafka.connect.connector.policy.ConnectorClientConfigRequest",
+            "org.apache.kafka.connect.connector.policy.ConnectorClientConfigRequest$ClientType",
+            "org.apache.kafka.connect.connector.",
+            "org.apache.kafka.connect.connector.Connector",
+            "org.apache.kafka.connect.connector.ConnectorContext",
+            "org.apache.kafka.connect.connector.ConnectRecord",
+            "org.apache.kafka.connect.connector.Task",
+            "org.apache.kafka.connect.data.",
+            "org.apache.kafka.connect.data.ConnectSchema",
+            "org.apache.kafka.connect.data.Date",
+            "org.apache.kafka.connect.data.Decimal",
+            "org.apache.kafka.connect.data.Field",
+            "org.apache.kafka.connect.data.Schema",
+            "org.apache.kafka.connect.data.SchemaAndValue",
+            "org.apache.kafka.connect.data.SchemaBuilder",
+            "org.apache.kafka.connect.data.SchemaProjector",
+            "org.apache.kafka.connect.data.Struct",
+            "org.apache.kafka.connect.data.Time",
+            "org.apache.kafka.connect.data.Timestamp",
+            "org.apache.kafka.connect.data.Values",
+            "org.apache.kafka.connect.errors.",
+            "org.apache.kafka.connect.errors.AlreadyExistsException",
+            "org.apache.kafka.connect.errors.ConnectException",
+            "org.apache.kafka.connect.errors.DataException",
+            "org.apache.kafka.connect.errors.IllegalWorkerStateException",
+            "org.apache.kafka.connect.errors.NotFoundException",
+            "org.apache.kafka.connect.errors.RetriableException",
+            "org.apache.kafka.connect.errors.SchemaBuilderException",
+            "org.apache.kafka.connect.errors.SchemaProjectorException",
+            "org.apache.kafka.connect.header.",
+            "org.apache.kafka.connect.header.ConnectHeader",
+            "org.apache.kafka.connect.header.ConnectHeaders",
+            "org.apache.kafka.connect.header.Header",
+            "org.apache.kafka.connect.header.Headers",
+            "org.apache.kafka.connect.health.",
+            "org.apache.kafka.connect.health.AbstractState",
+            "org.apache.kafka.connect.health.ConnectClusterDetails",
+            "org.apache.kafka.connect.health.ConnectClusterState",
+            "org.apache.kafka.connect.health.ConnectorHealth",
+            "org.apache.kafka.connect.health.ConnectorState",
+            "org.apache.kafka.connect.health.ConnectorType",
+            "org.apache.kafka.connect.health.TaskState",
+            "org.apache.kafka.connect.rest.",
+            "org.apache.kafka.connect.rest.ConnectRestExtension",
+            "org.apache.kafka.connect.rest.ConnectRestExtensionContext",
+            "org.apache.kafka.connect.sink.",
+            "org.apache.kafka.connect.sink.SinkConnector",
+            "org.apache.kafka.connect.sink.SinkRecord",
+            "org.apache.kafka.connect.sink.SinkTask",
+            "org.apache.kafka.connect.sink.SinkTaskContext",
+            "org.apache.kafka.connect.sink.ErrantRecordReporter",
+            "org.apache.kafka.connect.source.",
+            "org.apache.kafka.connect.source.SourceConnector",
+            "org.apache.kafka.connect.source.SourceRecord",
+            "org.apache.kafka.connect.source.SourceTask",
+            "org.apache.kafka.connect.source.SourceTaskContext",
+            "org.apache.kafka.connect.storage.",
+            "org.apache.kafka.connect.storage.Converter",
+            "org.apache.kafka.connect.storage.ConverterConfig",
+            "org.apache.kafka.connect.storage.ConverterType",
+            "org.apache.kafka.connect.storage.HeaderConverter",
+            "org.apache.kafka.connect.storage.OffsetStorageReader",
+            //"org.apache.kafka.connect.storage.SimpleHeaderConverter", explicitly isolated
+            //"org.apache.kafka.connect.storage.StringConverter", explicitly isolated
+            "org.apache.kafka.connect.storage.StringConverterConfig",
+            //"org.apache.kafka.connect.transforms.", isolated by default
+            "org.apache.kafka.connect.transforms.Transformation",
+            "org.apache.kafka.connect.util.",
+            "org.apache.kafka.connect.util.ConnectorUtils"
+        );
+        // Classes in the API should never be loaded in isolation.
+        for (String clazz : apiClasses) {
+            assertFalse(
+                clazz + " from 'api' is loaded in isolation but should not be",
+                PluginUtils.shouldLoadInIsolation(clazz)
+            );
+        }
+    }
+
+    @Test
+    public void testConnectRuntimeClasses() {
+        // Only list packages, because there are too many classes.
+        List runtimeClasses = Arrays.asList(
+            "org.apache.kafka.connect.cli.",
+            //"org.apache.kafka.connect.connector.policy.", isolated by default
+            //"org.apache.kafka.connect.converters.", isolated by default
+            "org.apache.kafka.connect.runtime.",
+            "org.apache.kafka.connect.runtime.distributed",
+            "org.apache.kafka.connect.runtime.errors",
+            "org.apache.kafka.connect.runtime.health",
+            "org.apache.kafka.connect.runtime.isolation",
+            "org.apache.kafka.connect.runtime.rest.",
+            "org.apache.kafka.connect.runtime.rest.entities.",
+            "org.apache.kafka.connect.runtime.rest.errors.",
+            "org.apache.kafka.connect.runtime.rest.resources.",
+            "org.apache.kafka.connect.runtime.rest.util.",
+            "org.apache.kafka.connect.runtime.standalone.",
+            "org.apache.kafka.connect.runtime.rest.",
+            "org.apache.kafka.connect.storage.",
+            "org.apache.kafka.connect.tools.",
+            "org.apache.kafka.connect.util."
+        );
+        for (String clazz : runtimeClasses) {
+            assertFalse(
+                clazz + " from 'runtime' is loaded in isolation but should not be",
+                PluginUtils.shouldLoadInIsolation(clazz)
+            );
+        }
+    }
+
+    @Test
+    public void testAllowedRuntimeClasses() {
+        List jsonConverterClasses = Arrays.asList(
+            "org.apache.kafka.connect.connector.policy.",
+            "org.apache.kafka.connect.connector.policy.AbstractConnectorClientConfigOverridePolicy",
+            "org.apache.kafka.connect.connector.policy.AllConnectorClientConfigOverridePolicy",
+            "org.apache.kafka.connect.connector.policy.NoneConnectorClientConfigOverridePolicy",
+            "org.apache.kafka.connect.connector.policy.PrincipalConnectorClientConfigOverridePolicy",
+            "org.apache.kafka.connect.converters.",
+            "org.apache.kafka.connect.converters.ByteArrayConverter",
+            "org.apache.kafka.connect.converters.DoubleConverter",
+            "org.apache.kafka.connect.converters.FloatConverter",
+            "org.apache.kafka.connect.converters.IntegerConverter",
+            "org.apache.kafka.connect.converters.LongConverter",
+            "org.apache.kafka.connect.converters.NumberConverter",
+            "org.apache.kafka.connect.converters.NumberConverterConfig",
+            "org.apache.kafka.connect.converters.ShortConverter",
+            //"org.apache.kafka.connect.storage.", not isolated by default
+            "org.apache.kafka.connect.storage.StringConverter",
+            "org.apache.kafka.connect.storage.SimpleHeaderConverter"
+        );
+        for (String clazz : jsonConverterClasses) {
+            assertTrue(
+                clazz + " from 'runtime' is not loaded in isolation but should be",
+                PluginUtils.shouldLoadInIsolation(clazz)
+            );
+        }
+    }
+
+    @Test
+    public void testTransformsClasses() {
+        List transformsClasses = Arrays.asList(
+            "org.apache.kafka.connect.transforms.",
+            "org.apache.kafka.connect.transforms.util.",
+            "org.apache.kafka.connect.transforms.util.NonEmptyListValidator",
+            "org.apache.kafka.connect.transforms.util.RegexValidator",
+            "org.apache.kafka.connect.transforms.util.Requirements",
+            "org.apache.kafka.connect.transforms.util.SchemaUtil",
+            "org.apache.kafka.connect.transforms.util.SimpleConfig",
+            "org.apache.kafka.connect.transforms.Cast",
+            "org.apache.kafka.connect.transforms.Cast$Key",
+            "org.apache.kafka.connect.transforms.Cast$Value",
+            "org.apache.kafka.connect.transforms.ExtractField",
+            "org.apache.kafka.connect.transforms.ExtractField$Key",
+            "org.apache.kafka.connect.transforms.ExtractField$Value",
+            "org.apache.kafka.connect.transforms.Flatten",
+            "org.apache.kafka.connect.transforms.Flatten$Key",
+            "org.apache.kafka.connect.transforms.Flatten$Value",
+            "org.apache.kafka.connect.transforms.HoistField",
+            "org.apache.kafka.connect.transforms.HoistField$Key",
+            "org.apache.kafka.connect.transforms.HoistField$Key",
+            "org.apache.kafka.connect.transforms.InsertField",
+            "org.apache.kafka.connect.transforms.InsertField$Key",
+            "org.apache.kafka.connect.transforms.InsertField$Value",
+            "org.apache.kafka.connect.transforms.MaskField",
+            "org.apache.kafka.connect.transforms.MaskField$Key",
+            "org.apache.kafka.connect.transforms.MaskField$Value",
+            "org.apache.kafka.connect.transforms.RegexRouter",
+            "org.apache.kafka.connect.transforms.ReplaceField",
+            "org.apache.kafka.connect.transforms.ReplaceField$Key",
+            "org.apache.kafka.connect.transforms.ReplaceField$Value",
+            "org.apache.kafka.connect.transforms.SetSchemaMetadata",
+            "org.apache.kafka.connect.transforms.SetSchemaMetadata$Key",
+            "org.apache.kafka.connect.transforms.SetSchemaMetadata$Value",
+            "org.apache.kafka.connect.transforms.TimestampConverter",
+            "org.apache.kafka.connect.transforms.TimestampConverter$Key",
+            "org.apache.kafka.connect.transforms.TimestampConverter$Value",
+            "org.apache.kafka.connect.transforms.TimestampRouter",
+            "org.apache.kafka.connect.transforms.TimestampRouter$Key",
+            "org.apache.kafka.connect.transforms.TimestampRouter$Value",
+            "org.apache.kafka.connect.transforms.ValueToKey"
+        );
+        for (String clazz : transformsClasses) {
+            assertTrue(
+                clazz + " from 'transforms' is not loaded in isolation but should be",
+                PluginUtils.shouldLoadInIsolation(clazz)
+            );
+        }
+    }
+
+    @Test
+    public void testAllowedJsonConverterClasses() {
+        List jsonConverterClasses = Arrays.asList(
+            "org.apache.kafka.connect.json.",
+            "org.apache.kafka.connect.json.DecimalFormat",
+            "org.apache.kafka.connect.json.JsonConverter",
+            "org.apache.kafka.connect.json.JsonConverterConfig",
+            "org.apache.kafka.connect.json.JsonDeserializer",
+            "org.apache.kafka.connect.json.JsonSchema",
+            "org.apache.kafka.connect.json.JsonSerializer"
+        );
+        for (String clazz : jsonConverterClasses) {
+            assertTrue(
+                clazz + " from 'json' is not loaded in isolation but should be",
+                PluginUtils.shouldLoadInIsolation(clazz)
+            );
+        }
+    }
+
+    @Test
+    public void testAllowedFileConnectors() {
+        List jsonConverterClasses = Arrays.asList(
+            "org.apache.kafka.connect.file.",
+            "org.apache.kafka.connect.file.FileStreamSinkConnector",
+            "org.apache.kafka.connect.file.FileStreamSinkTask",
+            "org.apache.kafka.connect.file.FileStreamSourceConnector",
+            "org.apache.kafka.connect.file.FileStreamSourceTask"
+        );
+        for (String clazz : jsonConverterClasses) {
+            assertTrue(
+                clazz + " from 'file' is not loaded in isolation but should be",
+                PluginUtils.shouldLoadInIsolation(clazz)
+            );
+        }
+    }
+
+    @Test
+    public void testAllowedBasicAuthExtensionClasses() {
+        List basicAuthExtensionClasses = Arrays.asList(
+            "org.apache.kafka.connect.rest.basic.auth.extension.BasicAuthSecurityRestExtension"
+            //"org.apache.kafka.connect.rest.basic.auth.extension.JaasBasicAuthFilter", TODO fix?
+            //"org.apache.kafka.connect.rest.basic.auth.extension.PropertyFileLoginModule" TODO fix?
+        );
+        for (String clazz : basicAuthExtensionClasses) {
+            assertTrue(
+                clazz + " from 'basic-auth-extension' is not loaded in isolation but should be",
+                PluginUtils.shouldLoadInIsolation(clazz)
+            );
+        }
+    }
+
+    @Test
+    public void testMirrorClasses() {
         assertTrue(PluginUtils.shouldLoadInIsolation(
-                "org.apache.kafka.connect.storage.StringConverter")
+                "org.apache.kafka.connect.mirror.MirrorSourceTask")
         );
         assertTrue(PluginUtils.shouldLoadInIsolation(
-                "org.apache.kafka.connect.storage.SimpleHeaderConverter")
+                "org.apache.kafka.connect.mirror.MirrorSourceConnector")
         );
-        assertTrue(PluginUtils.shouldLoadInIsolation(
-            "org.apache.kafka.connect.rest.basic.auth.extension.BasicAuthSecurityRestExtension"
-        ));
     }
 
     @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 1100910df5f43..4d304c710d838 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
@@ -17,11 +17,16 @@
 
 package org.apache.kafka.connect.runtime.isolation;
 
+import java.util.Collections;
+import java.util.Map.Entry;
 import org.apache.kafka.common.Configurable;
 import org.apache.kafka.common.config.AbstractConfig;
 import org.apache.kafka.common.config.ConfigDef;
+import org.apache.kafka.common.config.ConfigException;
+import org.apache.kafka.common.config.provider.ConfigProvider;
 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.json.JsonConverter;
 import org.apache.kafka.connect.json.JsonConverterConfig;
 import org.apache.kafka.connect.rest.ConnectRestExtension;
@@ -34,7 +39,6 @@
 import org.apache.kafka.connect.storage.HeaderConverter;
 import org.apache.kafka.connect.storage.SimpleHeaderConverter;
 import org.junit.Before;
-import org.junit.BeforeClass;
 import org.junit.Test;
 
 import java.io.IOException;
@@ -45,31 +49,26 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertTrue;
 
 public class PluginsTest {
 
-    private static Map pluginProps;
-    private static Plugins plugins;
+    private Plugins plugins;
     private Map props;
     private AbstractConfig config;
     private TestConverter converter;
     private TestHeaderConverter headerConverter;
     private TestInternalConverter internalConverter;
 
-    @BeforeClass
-    public static void beforeAll() {
-        pluginProps = new HashMap<>();
-
-        // Set up the plugins to have no additional plugin directories.
-        // This won't allow us to test classpath isolation, but it will allow us to test some of the utility methods.
-        pluginProps.put(WorkerConfig.PLUGIN_PATH_CONFIG, "");
-        plugins = new Plugins(pluginProps);
-    }
-
     @SuppressWarnings("deprecation")
     @Before
     public void setup() {
+        Map pluginProps = new HashMap<>();
+
+        // Set up the plugins with some test plugins to test isolation
+        pluginProps.put(WorkerConfig.PLUGIN_PATH_CONFIG, String.join(",", TestPlugins.pluginPath()));
+        plugins = new Plugins(pluginProps);
         props = new HashMap<>(pluginProps);
         props.put(WorkerConfig.KEY_CONVERTER_CLASS_CONFIG, TestConverter.class.getName());
         props.put(WorkerConfig.VALUE_CONVERTER_CLASS_CONFIG, TestConverter.class.getName());
@@ -185,6 +184,186 @@ public void shouldInstantiateAndConfigureDefaultHeaderConverter() {
         assertTrue(headerConverter instanceof SimpleHeaderConverter);
     }
 
+    @Test(expected = ConnectException.class)
+    public void shouldThrowIfPluginThrows() {
+        TestPlugins.assertAvailable();
+
+        plugins.newPlugin(
+            TestPlugins.ALWAYS_THROW_EXCEPTION,
+            new AbstractConfig(new ConfigDef(), Collections.emptyMap()),
+            Converter.class
+        );
+    }
+
+    @Test
+    public void shouldShareStaticValuesBetweenSamePlugin() {
+        // Plugins are not isolated from other instances of their own class.
+        TestPlugins.assertAvailable();
+        Converter firstPlugin = plugins.newPlugin(
+            TestPlugins.ALIASED_STATIC_FIELD,
+            new AbstractConfig(new ConfigDef(), Collections.emptyMap()),
+            Converter.class
+        );
+
+        assertInstanceOf(SamplingTestPlugin.class, firstPlugin, "Cannot collect samples");
+
+        Converter secondPlugin = plugins.newPlugin(
+            TestPlugins.ALIASED_STATIC_FIELD,
+            new AbstractConfig(new ConfigDef(), Collections.emptyMap()),
+            Converter.class
+        );
+
+        assertInstanceOf(SamplingTestPlugin.class, secondPlugin, "Cannot collect samples");
+        assertSame(
+            ((SamplingTestPlugin) firstPlugin).otherSamples(),
+            ((SamplingTestPlugin) secondPlugin).otherSamples()
+        );
+    }
+
+    @Test
+    public void newPluginShouldServiceLoadWithPluginClassLoader() {
+        TestPlugins.assertAvailable();
+        Converter plugin = plugins.newPlugin(
+            TestPlugins.SERVICE_LOADER,
+            new AbstractConfig(new ConfigDef(), Collections.emptyMap()),
+            Converter.class
+        );
+
+        assertInstanceOf(SamplingTestPlugin.class, plugin, "Cannot collect samples");
+        Map samples = ((SamplingTestPlugin) plugin).flatten();
+        // Assert that the service loaded subclass is found in both environments
+        assertTrue(samples.containsKey("ServiceLoadedSubclass.static"));
+        assertTrue(samples.containsKey("ServiceLoadedSubclass.dynamic"));
+        assertPluginClassLoaderAlwaysActive(samples);
+    }
+
+    @Test
+    public void newPluginShouldInstantiateWithPluginClassLoader() {
+        TestPlugins.assertAvailable();
+        Converter plugin = plugins.newPlugin(
+            TestPlugins.ALIASED_STATIC_FIELD,
+            new AbstractConfig(new ConfigDef(), Collections.emptyMap()),
+            Converter.class
+        );
+
+        assertInstanceOf(SamplingTestPlugin.class, plugin, "Cannot collect samples");
+        Map samples = ((SamplingTestPlugin) plugin).flatten();
+        assertPluginClassLoaderAlwaysActive(samples);
+    }
+
+    @Test(expected = ConfigException.class)
+    public void shouldFailToFindConverterInCurrentClassloader() {
+        TestPlugins.assertAvailable();
+        props.put(WorkerConfig.KEY_CONVERTER_CLASS_CONFIG, TestPlugins.SAMPLING_CONVERTER);
+        createConfig();
+    }
+
+    @Test
+    public void newConverterShouldConfigureWithPluginClassLoader() {
+        TestPlugins.assertAvailable();
+        props.put(WorkerConfig.KEY_CONVERTER_CLASS_CONFIG, TestPlugins.SAMPLING_CONVERTER);
+        ClassLoader classLoader = plugins.delegatingLoader().pluginClassLoader(TestPlugins.SAMPLING_CONVERTER);
+        ClassLoader savedLoader = Plugins.compareAndSwapLoaders(classLoader);
+        createConfig();
+        Plugins.compareAndSwapLoaders(savedLoader);
+
+        Converter plugin = plugins.newConverter(
+            config,
+            WorkerConfig.KEY_CONVERTER_CLASS_CONFIG,
+            ClassLoaderUsage.PLUGINS
+        );
+
+        assertInstanceOf(SamplingTestPlugin.class, plugin, "Cannot collect samples");
+        Map samples = ((SamplingTestPlugin) plugin).flatten();
+        assertTrue(samples.containsKey("configure"));
+        assertPluginClassLoaderAlwaysActive(samples);
+    }
+
+    @Test
+    public void newConfigProviderShouldConfigureWithPluginClassLoader() {
+        TestPlugins.assertAvailable();
+        String providerPrefix = "some.provider";
+        props.put(providerPrefix + ".class", TestPlugins.SAMPLING_CONFIG_PROVIDER);
+
+        PluginClassLoader classLoader = plugins.delegatingLoader().pluginClassLoader(TestPlugins.SAMPLING_CONFIG_PROVIDER);
+        assertNotNull(classLoader);
+        ClassLoader savedLoader = Plugins.compareAndSwapLoaders(classLoader);
+        createConfig();
+        Plugins.compareAndSwapLoaders(savedLoader);
+
+        ConfigProvider plugin = plugins.newConfigProvider(
+            config,
+            providerPrefix,
+            ClassLoaderUsage.PLUGINS
+        );
+
+        assertInstanceOf(SamplingTestPlugin.class, plugin, "Cannot collect samples");
+        Map samples = ((SamplingTestPlugin) plugin).flatten();
+        assertTrue(samples.containsKey("configure"));
+        assertPluginClassLoaderAlwaysActive(samples);
+    }
+
+    @Test
+    public void newHeaderConverterShouldConfigureWithPluginClassLoader() {
+        TestPlugins.assertAvailable();
+        props.put(WorkerConfig.HEADER_CONVERTER_CLASS_CONFIG, TestPlugins.SAMPLING_HEADER_CONVERTER);
+        ClassLoader classLoader = plugins.delegatingLoader().pluginClassLoader(TestPlugins.SAMPLING_HEADER_CONVERTER);
+        ClassLoader savedLoader = Plugins.compareAndSwapLoaders(classLoader);
+        createConfig();
+        Plugins.compareAndSwapLoaders(savedLoader);
+
+        HeaderConverter plugin = plugins.newHeaderConverter(
+            config,
+            WorkerConfig.HEADER_CONVERTER_CLASS_CONFIG,
+            ClassLoaderUsage.PLUGINS
+        );
+
+        assertInstanceOf(SamplingTestPlugin.class, plugin, "Cannot collect samples");
+        Map samples = ((SamplingTestPlugin) plugin).flatten();
+        assertTrue(samples.containsKey("configure")); // HeaderConverter::configure was called
+        assertPluginClassLoaderAlwaysActive(samples);
+    }
+
+    @Test
+    public void newPluginsShouldConfigureWithPluginClassLoader() {
+        TestPlugins.assertAvailable();
+        List configurables = plugins.newPlugins(
+            Collections.singletonList(TestPlugins.SAMPLING_CONFIGURABLE),
+            config,
+            Configurable.class
+        );
+        assertEquals(1, configurables.size());
+        Configurable plugin = configurables.get(0);
+
+        assertInstanceOf(SamplingTestPlugin.class, plugin, "Cannot collect samples");
+        Map samples = ((SamplingTestPlugin) plugin).flatten();
+        assertTrue(samples.containsKey("configure")); // Configurable::configure was called
+        assertPluginClassLoaderAlwaysActive(samples);
+    }
+
+    public static void assertPluginClassLoaderAlwaysActive(Map samples) {
+        for (Entry e : samples.entrySet()) {
+            String sampleName = "\"" + e.getKey() + "\" (" + e.getValue() + ")";
+            assertInstanceOf(
+                PluginClassLoader.class,
+                e.getValue().staticClassloader(),
+                sampleName + " has incorrect static classloader"
+            );
+            assertInstanceOf(
+                PluginClassLoader.class,
+                e.getValue().classloader(),
+                sampleName + " has incorrect dynamic classloader"
+            );
+        }
+    }
+
+    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/isolation/SamplingTestPlugin.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/isolation/SamplingTestPlugin.java
new file mode 100644
index 0000000000000..bcf8881898e8e
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/isolation/SamplingTestPlugin.java
@@ -0,0 +1,122 @@
+/*
+ * 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.isolation;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Map.Entry;
+
+/**
+ * Base class for plugins so we can sample information about their initialization
+ */
+public abstract class SamplingTestPlugin {
+
+    /**
+     * @return the ClassLoader used to statically initialize this plugin class
+     */
+    public abstract ClassLoader staticClassloader();
+
+    /**
+     * @return the ClassLoader used to initialize this plugin instance
+     */
+    public abstract ClassLoader classloader();
+
+    /**
+     * @return a group of other SamplingTestPlugin instances known by this plugin
+     * This should only return direct children, and not reference this instance directly
+     */
+    public Map otherSamples() {
+        return Collections.emptyMap();
+    }
+
+    /**
+     * @return a flattened list of child samples including this entry keyed as "this"
+     */
+    public Map flatten() {
+        Map out = new HashMap<>();
+        Map otherSamples = otherSamples();
+        if (otherSamples != null) {
+            for (Entry child : otherSamples.entrySet()) {
+                for (Entry flattened : child.getValue().flatten().entrySet()) {
+                    String key = child.getKey();
+                    if (flattened.getKey().length() > 0) {
+                        key += "." + flattened.getKey();
+                    }
+                    out.put(key, flattened.getValue());
+                }
+            }
+        }
+        out.put("", this);
+        return out;
+    }
+
+    /**
+     * Log the parent method call as a child sample.
+     * Stores only the last invocation of each method if there are multiple invocations.
+     * @param samples The collection of samples to which this method call should be added
+     */
+    public void logMethodCall(Map samples) {
+        StackTraceElement[] stackTraces = Thread.currentThread().getStackTrace();
+        if (stackTraces.length < 2) {
+            return;
+        }
+        // 0 is inside getStackTrace
+        // 1 is this method
+        // 2 is our caller method
+        StackTraceElement caller = stackTraces[2];
+
+        samples.put(caller.getMethodName(), new MethodCallSample(
+            caller,
+            Thread.currentThread().getContextClassLoader(),
+            getClass().getClassLoader()
+        ));
+    }
+
+    public static class MethodCallSample extends SamplingTestPlugin {
+
+        private final StackTraceElement caller;
+        private final ClassLoader staticClassLoader;
+        private final ClassLoader dynamicClassLoader;
+
+        public MethodCallSample(
+            StackTraceElement caller,
+            ClassLoader staticClassLoader,
+            ClassLoader dynamicClassLoader
+        ) {
+            this.caller = caller;
+            this.staticClassLoader = staticClassLoader;
+            this.dynamicClassLoader = dynamicClassLoader;
+        }
+
+        @Override
+        public ClassLoader staticClassloader() {
+            return staticClassLoader;
+        }
+
+        @Override
+        public ClassLoader classloader() {
+            return dynamicClassLoader;
+        }
+
+        @Override
+        public String toString() {
+            return caller.toString();
+        }
+    }
+}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/isolation/TestPlugins.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/isolation/TestPlugins.java
new file mode 100644
index 0000000000000..9561ffb0f5b14
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/isolation/TestPlugins.java
@@ -0,0 +1,267 @@
+/*
+ * 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.isolation;
+
+import java.io.BufferedInputStream;
+import java.io.File;
+import java.io.FileInputStream;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.StringWriter;
+import java.net.URL;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.jar.Attributes;
+import java.util.jar.JarEntry;
+import java.util.jar.JarOutputStream;
+import java.util.jar.Manifest;
+import java.util.stream.Collectors;
+import javax.tools.JavaCompiler;
+import javax.tools.StandardJavaFileManager;
+import javax.tools.ToolProvider;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Utility class for constructing test plugins for Connect.
+ *
+ * 
      Plugins are built from their source under resources/test-plugins/ and placed into temporary
+ * jar files that are deleted when the process exits.
+ *
+ * 
      To add a plugin, create the source files in the resource tree, and edit this class to build
+ * that plugin during initialization. For example, the plugin class {@literal package.Class} should
+ * be placed in {@literal resources/test-plugins/something/package/Class.java} and loaded using
+ * {@code createPluginJar("something")}. The class name, contents, and plugin directory can take
+ * any value you need for testing.
+ *
+ * 
      To use this class in your tests, make sure to first call
+ * {@link TestPlugins#assertAvailable()} to verify that the plugins initialized correctly.
+ * Otherwise, exceptions during the plugin build are not propagated, and may invalidate your test.
+ * You can access the list of plugin jars for assembling a {@literal plugin.path}, and reference
+ * the names of the different plugins directly via the exposed constants.
+ */
+public class TestPlugins {
+
+    /**
+     * Class name of a plugin which will always throw an exception during loading
+     */
+    public static final String ALWAYS_THROW_EXCEPTION = "test.plugins.AlwaysThrowException";
+    /**
+     * Class name of a plugin which samples information about its initialization.
+     */
+    public static final String ALIASED_STATIC_FIELD = "test.plugins.AliasedStaticField";
+    /**
+     * Class name of a {@link org.apache.kafka.connect.storage.Converter}
+     * which samples information about its method calls.
+     */
+    public static final String SAMPLING_CONVERTER = "test.plugins.SamplingConverter";
+    /**
+     * Class name of a {@link org.apache.kafka.common.Configurable}
+     * which samples information about its method calls.
+     */
+    public static final String SAMPLING_CONFIGURABLE = "test.plugins.SamplingConfigurable";
+    /**
+     * Class name of a {@link org.apache.kafka.connect.storage.HeaderConverter}
+     * which samples information about its method calls.
+     */
+    public static final String SAMPLING_HEADER_CONVERTER = "test.plugins.SamplingHeaderConverter";
+    /**
+     * Class name of a {@link org.apache.kafka.common.config.provider.ConfigProvider}
+     * which samples information about its method calls.
+     */
+    public static final String SAMPLING_CONFIG_PROVIDER = "test.plugins.SamplingConfigProvider";
+    /**
+     * Class name of a plugin which uses a {@link java.util.ServiceLoader}
+     * to load internal classes, and samples information about their initialization.
+     */
+    public static final String SERVICE_LOADER = "test.plugins.ServiceLoaderPlugin";
+
+    private static final Logger log = LoggerFactory.getLogger(TestPlugins.class);
+    private static final Map PLUGIN_JARS;
+    private static final Throwable INITIALIZATION_EXCEPTION;
+
+    static {
+        Throwable err = null;
+        HashMap pluginJars = new HashMap<>();
+        try {
+            pluginJars.put(ALWAYS_THROW_EXCEPTION, createPluginJar("always-throw-exception"));
+            pluginJars.put(ALIASED_STATIC_FIELD, createPluginJar("aliased-static-field"));
+            pluginJars.put(SAMPLING_CONVERTER, createPluginJar("sampling-converter"));
+            pluginJars.put(SAMPLING_CONFIGURABLE, createPluginJar("sampling-configurable"));
+            pluginJars.put(SAMPLING_HEADER_CONVERTER, createPluginJar("sampling-header-converter"));
+            pluginJars.put(SAMPLING_CONFIG_PROVIDER, createPluginJar("sampling-config-provider"));
+            pluginJars.put(SERVICE_LOADER, createPluginJar("service-loader"));
+        } catch (Throwable e) {
+            log.error("Could not set up plugin test jars", e);
+            err = e;
+        }
+        PLUGIN_JARS = Collections.unmodifiableMap(pluginJars);
+        INITIALIZATION_EXCEPTION = err;
+    }
+
+    /**
+     * Ensure that the test plugin JARs were assembled without error before continuing.
+     * @throws AssertionError if any plugin failed to load, or no plugins were loaded.
+     */
+    public static void assertAvailable() throws AssertionError {
+        if (INITIALIZATION_EXCEPTION != null) {
+            throw new AssertionError("TestPlugins did not initialize completely",
+                INITIALIZATION_EXCEPTION);
+        }
+        if (PLUGIN_JARS.isEmpty()) {
+            throw new AssertionError("No test plugins loaded");
+        }
+    }
+
+    /**
+     * A list of jar files containing test plugins
+     * @return A list of plugin jar filenames
+     */
+    public static List pluginPath() {
+        return PLUGIN_JARS.values()
+            .stream()
+            .map(File::getPath)
+            .collect(Collectors.toList());
+    }
+
+    /**
+     * Get all of the classes that were successfully built by this class
+     * @return A list of plugin class names
+     */
+    public static List pluginClasses() {
+        return new ArrayList<>(PLUGIN_JARS.keySet());
+    }
+
+    private static File createPluginJar(String resourceDir) throws IOException {
+        Path inputDir = resourceDirectoryPath("test-plugins/" + resourceDir);
+        Path binDir = Files.createTempDirectory(resourceDir + ".bin.");
+        compileJavaSources(inputDir, binDir);
+        File jarFile = Files.createTempFile(resourceDir + ".", ".jar").toFile();
+        try (JarOutputStream jar = openJarFile(jarFile)) {
+            writeJar(jar, inputDir);
+            writeJar(jar, binDir);
+        }
+        removeDirectory(binDir);
+        jarFile.deleteOnExit();
+        return jarFile;
+    }
+
+    private static Path resourceDirectoryPath(String resourceDir) throws IOException {
+        URL resource = Thread.currentThread()
+            .getContextClassLoader()
+            .getResource(resourceDir);
+        if (resource == null) {
+            throw new IOException("Could not find test plugin resource: " + resourceDir);
+        }
+        File file = new File(resource.getFile());
+        if (!file.isDirectory()) {
+            throw new IOException("Resource is not a directory: " + resourceDir);
+        }
+        if (!file.canRead()) {
+            throw new IOException("Resource directory is not readable: " + resourceDir);
+        }
+        return file.toPath();
+    }
+
+    private static JarOutputStream openJarFile(File jarFile) throws IOException {
+        Manifest manifest = new Manifest();
+        manifest.getMainAttributes().put(Attributes.Name.MANIFEST_VERSION, "1.0");
+        return new JarOutputStream(new FileOutputStream(jarFile), manifest);
+    }
+
+    private static void removeDirectory(Path binDir) throws IOException {
+        List classFiles = Files.walk(binDir)
+            .sorted(Comparator.reverseOrder())
+            .map(Path::toFile)
+            .collect(Collectors.toList());
+        for (File classFile : classFiles) {
+            if (!classFile.delete()) {
+                throw new IOException("Could not delete: " + classFile);
+            }
+        }
+    }
+
+    /**
+     * Compile a directory of .java source files into .class files
+     * .class files are placed into the same directory as their sources.
+     *
+     * 
      Dependencies between source files in this directory are resolved against one another
+     * and the classes present in the test environment.
+     * See https://stackoverflow.com/questions/1563909/ for more information.
+     * Additional dependencies in your plugins should be added as test scope to :connect:runtime.
+     * @param sourceDir Directory containing java source files
+     * @throws IOException if the files cannot be compiled
+     */
+    private static void compileJavaSources(Path sourceDir, Path binDir) throws IOException {
+        JavaCompiler compiler = ToolProvider.getSystemJavaCompiler();
+        List sourceFiles = Files.walk(sourceDir)
+            .filter(Files::isRegularFile)
+            .filter(path -> path.toFile().getName().endsWith(".java"))
+            .map(Path::toFile)
+            .collect(Collectors.toList());
+        StringWriter writer = new StringWriter();
+        List options = Arrays.asList(
+            "-d", binDir.toString() // Write class output to a different directory.
+        );
+
+        try (StandardJavaFileManager fileManager = compiler.getStandardFileManager(null, null, null)) {
+            boolean success = compiler.getTask(
+                writer,
+                fileManager,
+                null,
+                options,
+                null,
+                fileManager.getJavaFileObjectsFromFiles(sourceFiles)
+            ).call();
+            if (!success) {
+                throw new RuntimeException("Failed to compile test plugin:\n" + writer);
+            }
+        }
+    }
+
+    private static void writeJar(JarOutputStream jar, Path inputDir) throws IOException {
+        List paths = Files.walk(inputDir)
+            .filter(Files::isRegularFile)
+            .filter(path -> !path.toFile().getName().endsWith(".java"))
+            .collect(Collectors.toList());
+        for (Path path : paths) {
+            try (InputStream in = new BufferedInputStream(new FileInputStream(path.toFile()))) {
+                jar.putNextEntry(new JarEntry(
+                    inputDir.relativize(path)
+                        .toFile()
+                        .getPath()
+                        .replace(File.separator, "/")
+                ));
+                byte[] buffer = new byte[1024];
+                for (int count; (count = in.read(buffer)) != -1; ) {
+                    jar.write(buffer, 0, count);
+                }
+                jar.closeEntry();
+            }
+        }
+    }
+
+}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/rest/InternalRequestSignatureTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/rest/InternalRequestSignatureTest.java
new file mode 100644
index 0000000000000..89012ee474ec6
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/rest/InternalRequestSignatureTest.java
@@ -0,0 +1,151 @@
+/*
+ * 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.rest;
+
+import org.apache.kafka.connect.errors.ConnectException;
+import org.apache.kafka.connect.runtime.rest.errors.BadRequestException;
+import org.eclipse.jetty.client.api.Request;
+import org.junit.Test;
+import org.mockito.ArgumentCaptor;
+
+import javax.crypto.Mac;
+import javax.crypto.SecretKey;
+import javax.crypto.spec.SecretKeySpec;
+import javax.ws.rs.core.HttpHeaders;
+
+import java.util.Base64;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+public class InternalRequestSignatureTest {
+
+    private static final byte[] REQUEST_BODY =
+            "[{\"config\":\"value\"},{\"config\":\"other_value\"}]".getBytes();
+    private static final String SIGNATURE_ALGORITHM = "HmacSHA256";
+    private static final SecretKey KEY = new SecretKeySpec(
+        new byte[] {
+            109, 116, -111, 49, -94, 25, -103, 44, -99, -118, 53, -69, 87, -124, 5, 48,
+            89, -105, -2, 58, -92, 87, 67, 49, -125, -79, -39, -126, -51, -53, -85, 57
+        }, "HmacSHA256"
+    );
+    private static final byte[] SIGNATURE = new byte[] {
+        42, -3, 127, 57, 43, 49, -51, -43, 72, -62, -10, 120, 123, 125, 26, -65,
+        36, 72, 86, -71, -32, 13, -8, 115, 85, 73, -65, -112, 6, 68, 41, -50
+    };
+    private static final String ENCODED_SIGNATURE = Base64.getEncoder().encodeToString(SIGNATURE);
+
+    @Test
+    public void fromHeadersShouldReturnNullOnNullHeaders() {
+        assertNull(InternalRequestSignature.fromHeaders(REQUEST_BODY, null));
+    }
+
+    @Test
+    public void fromHeadersShouldReturnNullIfSignatureHeaderMissing() {
+        assertNull(InternalRequestSignature.fromHeaders(REQUEST_BODY, internalRequestHeaders(null, SIGNATURE_ALGORITHM)));
+    }
+
+    @Test
+    public void fromHeadersShouldReturnNullIfSignatureAlgorithmHeaderMissing() {
+        assertNull(InternalRequestSignature.fromHeaders(REQUEST_BODY, internalRequestHeaders(ENCODED_SIGNATURE, null)));
+    }
+
+    @Test(expected = BadRequestException.class)
+    public void fromHeadersShouldThrowExceptionOnInvalidSignatureAlgorithm() {
+        InternalRequestSignature.fromHeaders(REQUEST_BODY, internalRequestHeaders(ENCODED_SIGNATURE, "doesn'texist"));
+    }
+
+    @Test(expected = BadRequestException.class)
+    public void fromHeadersShouldThrowExceptionOnInvalidBase64Signature() {
+        InternalRequestSignature.fromHeaders(REQUEST_BODY, internalRequestHeaders("not valid base 64", SIGNATURE_ALGORITHM));
+    }
+
+    @Test
+    public void fromHeadersShouldReturnNonNullResultOnValidSignatureAndSignatureAlgorithm() {
+        InternalRequestSignature signature =
+                InternalRequestSignature.fromHeaders(REQUEST_BODY, internalRequestHeaders(ENCODED_SIGNATURE, SIGNATURE_ALGORITHM));
+        assertNotNull(signature);
+        assertNotNull(signature.keyAlgorithm());
+    }
+
+    @Test(expected = ConnectException.class)
+    public void addToRequestShouldThrowExceptionOnInvalidSignatureAlgorithm() {
+        Request request = mock(Request.class);
+        InternalRequestSignature.addToRequest(KEY, REQUEST_BODY, "doesn'texist", request);
+    }
+
+    @Test
+    public void addToRequestShouldAddHeadersOnValidSignatureAlgorithm() {
+        Request request = mock(Request.class);
+        ArgumentCaptor signatureCapture = ArgumentCaptor.forClass(String.class);
+        ArgumentCaptor signatureAlgorithmCapture = ArgumentCaptor.forClass(String.class);
+        when(request.header(
+                eq(InternalRequestSignature.SIGNATURE_HEADER),
+                signatureCapture.capture()
+            )).thenReturn(request);
+        when(request.header(
+                eq(InternalRequestSignature.SIGNATURE_ALGORITHM_HEADER),
+                signatureAlgorithmCapture.capture()
+            )).thenReturn(request);
+
+        InternalRequestSignature.addToRequest(KEY, REQUEST_BODY, SIGNATURE_ALGORITHM, request);
+
+        assertEquals(
+            "Request should have valid base 64-encoded signature added as header",
+            ENCODED_SIGNATURE,
+            signatureCapture.getValue()
+        );
+        assertEquals(
+            "Request should have provided signature algorithm added as header",
+            SIGNATURE_ALGORITHM,
+            signatureAlgorithmCapture.getValue()
+        );
+    }
+
+    @Test
+    public void testSignatureValidation() throws Exception {
+        Mac mac = Mac.getInstance(SIGNATURE_ALGORITHM);
+
+        InternalRequestSignature signature = new InternalRequestSignature(REQUEST_BODY, mac, SIGNATURE);
+        assertTrue(signature.isValid(KEY));
+
+        signature = InternalRequestSignature.fromHeaders(REQUEST_BODY, internalRequestHeaders(ENCODED_SIGNATURE, SIGNATURE_ALGORITHM));
+        assertTrue(signature.isValid(KEY));
+
+        signature = new InternalRequestSignature("[{\"different_config\":\"different_value\"}]".getBytes(), mac, SIGNATURE);
+        assertFalse(signature.isValid(KEY));
+
+        signature = new InternalRequestSignature(REQUEST_BODY, mac, "bad signature".getBytes());
+        assertFalse(signature.isValid(KEY));
+    }
+
+    private static HttpHeaders internalRequestHeaders(String signature, String signatureAlgorithm) {
+        HttpHeaders result = mock(HttpHeaders.class);
+        when(result.getHeaderString(eq(InternalRequestSignature.SIGNATURE_HEADER)))
+            .thenReturn(signature);
+        when(result.getHeaderString(eq(InternalRequestSignature.SIGNATURE_ALGORITHM_HEADER)))
+            .thenReturn(signatureAlgorithm);
+        return result;
+    }
+}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/rest/RestServerTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/rest/RestServerTest.java
index 7b985546bad5e..575c4dae2a2cd 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/rest/RestServerTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/rest/RestServerTest.java
@@ -16,22 +16,25 @@
  */
 package org.apache.kafka.connect.runtime.rest;
 
+import com.fasterxml.jackson.core.type.TypeReference;
+import com.fasterxml.jackson.databind.ObjectMapper;
 import org.apache.http.HttpHost;
 import org.apache.http.HttpRequest;
 import org.apache.http.client.methods.CloseableHttpResponse;
 import org.apache.http.client.methods.HttpGet;
 import org.apache.http.client.methods.HttpOptions;
+import org.apache.http.client.methods.HttpPut;
+import org.apache.http.entity.StringEntity;
+import org.apache.http.impl.client.BasicResponseHandler;
 import org.apache.http.impl.client.CloseableHttpClient;
 import org.apache.http.impl.client.HttpClients;
 import org.apache.kafka.clients.CommonClientConfigs;
 import org.apache.kafka.connect.rest.ConnectRestExtension;
 import org.apache.kafka.connect.runtime.Herder;
-import org.apache.kafka.connect.runtime.HerderProvider;
 import org.apache.kafka.connect.runtime.WorkerConfig;
 import org.apache.kafka.connect.runtime.distributed.DistributedConfig;
 import org.apache.kafka.connect.runtime.isolation.Plugins;
-import org.apache.kafka.connect.util.Callback;
-import org.easymock.Capture;
+import org.apache.kafka.connect.runtime.standalone.StandaloneConfig;
 import org.easymock.EasyMock;
 import org.junit.After;
 import org.junit.Assert;
@@ -41,19 +44,26 @@
 import org.powermock.api.easymock.annotation.MockStrict;
 import org.powermock.core.classloader.annotations.PowerMockIgnore;
 import org.powermock.modules.junit4.PowerMockRunner;
+import org.slf4j.LoggerFactory;
 
 import javax.ws.rs.core.MediaType;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
 import java.util.Arrays;
-import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
 
+import static org.apache.kafka.connect.runtime.WorkerConfig.ADMIN_LISTENERS_CONFIG;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
+
 @RunWith(PowerMockRunner.class)
-@PowerMockIgnore({"javax.net.ssl.*", "javax.security.*"})
+@PowerMockIgnore({"javax.net.ssl.*", "javax.security.*", "javax.crypto.*"})
 public class RestServerTest {
 
     @MockStrict
@@ -62,6 +72,8 @@ public class RestServerTest {
     private Plugins plugins;
     private RestServer server;
 
+    protected static final String KAFKA_CLUSTER_ID = "Xbafgnagvar";
+
     @After
     public void tearDown() {
         server.stop();
@@ -79,6 +91,7 @@ private Map baseWorkerProps() {
         workerProps.put(WorkerConfig.INTERNAL_KEY_CONVERTER_CLASS_CONFIG, "org.apache.kafka.connect.json.JsonConverter");
         workerProps.put(WorkerConfig.INTERNAL_VALUE_CONVERTER_CLASS_CONFIG, "org.apache.kafka.connect.json.JsonConverter");
         workerProps.put(DistributedConfig.OFFSET_STORAGE_TOPIC_CONFIG, "connect-offsets");
+        workerProps.put(WorkerConfig.LISTENERS_CONFIG, "HTTP://localhost:0");
 
         return workerProps;
     }
@@ -106,6 +119,7 @@ public void testParseListeners() {
 
         // Build listener from hostname and port
         configMap = new HashMap<>(baseWorkerProps());
+        configMap.remove(WorkerConfig.LISTENERS_CONFIG);
         configMap.put(WorkerConfig.REST_HOST_NAME_CONFIG, "my-hostname");
         configMap.put(WorkerConfig.REST_PORT_CONFIG, "8080");
         config = new DistributedConfig(configMap);
@@ -116,7 +130,7 @@ public void testParseListeners() {
     @SuppressWarnings("deprecation")
     @Test
     public void testAdvertisedUri() {
-        // Advertised URI from listeenrs without protocol
+        // Advertised URI from listeners without protocol
         Map configMap = new HashMap<>(baseWorkerProps());
         configMap.put(WorkerConfig.LISTENERS_CONFIG, "http://localhost:8080,https://localhost:8443");
         DistributedConfig config = new DistributedConfig(configMap);
@@ -154,11 +168,20 @@ public void testAdvertisedUri() {
 
         // listener from hostname and port
         configMap = new HashMap<>(baseWorkerProps());
+        configMap.remove(WorkerConfig.LISTENERS_CONFIG);
         configMap.put(WorkerConfig.REST_HOST_NAME_CONFIG, "my-hostname");
         configMap.put(WorkerConfig.REST_PORT_CONFIG, "8080");
         config = new DistributedConfig(configMap);
         server = new RestServer(config);
         Assert.assertEquals("http://my-hostname:8080/", server.advertisedUrl().toString());
+
+        // correct listener is chosen when https listener is configured before http listener and advertised listener is http
+        configMap = new HashMap<>(baseWorkerProps());
+        configMap.put(WorkerConfig.LISTENERS_CONFIG, "https://encrypted-localhost:42069,http://plaintext-localhost:4761");
+        configMap.put(WorkerConfig.REST_ADVERTISED_LISTENER_CONFIG, "http");
+        config = new DistributedConfig(configMap);
+        server = new RestServer(config);
+        Assert.assertEquals("http://plaintext-localhost:4761/", server.advertisedUrl().toString());
     }
 
     @Test
@@ -166,6 +189,7 @@ public void testOptionsDoesNotIncludeWadlOutput() throws IOException {
         Map configMap = new HashMap<>(baseWorkerProps());
         DistributedConfig workerConfig = new DistributedConfig(configMap);
 
+        EasyMock.expect(herder.kafkaClusterId()).andReturn(KAFKA_CLUSTER_ID);
         EasyMock.expect(herder.plugins()).andStubReturn(plugins);
         EasyMock.expect(plugins.newPlugins(Collections.emptyList(),
             workerConfig,
@@ -174,7 +198,8 @@ public void testOptionsDoesNotIncludeWadlOutput() throws IOException {
         PowerMock.replayAll();
 
         server = new RestServer(workerConfig);
-        server.start(new HerderProvider(herder), herder.plugins());
+        server.initializeServer();
+        server.initializeResources(herder);
 
         HttpOptions request = new HttpOptions("/connectors");
         request.addHeader("Content-Type", MediaType.WILDCARD);
@@ -201,23 +226,20 @@ public void checkCORSRequest(String corsDomain, String origin, String expectedHe
         workerProps.put(WorkerConfig.ACCESS_CONTROL_ALLOW_METHODS_CONFIG, method);
         WorkerConfig workerConfig = new DistributedConfig(workerProps);
 
+        EasyMock.expect(herder.kafkaClusterId()).andReturn(KAFKA_CLUSTER_ID);
         EasyMock.expect(herder.plugins()).andStubReturn(plugins);
         EasyMock.expect(plugins.newPlugins(Collections.emptyList(),
                                            workerConfig,
                                            ConnectRestExtension.class))
             .andStubReturn(Collections.emptyList());
 
-        final Capture>> connectorsCallback = EasyMock.newCapture();
-        herder.connectors(EasyMock.capture(connectorsCallback));
-        PowerMock.expectLastCall().andAnswer(() -> {
-            connectorsCallback.getValue().onCompletion(null, Arrays.asList("a", "b"));
-            return null;
-        });
+        EasyMock.expect(herder.connectors()).andReturn(Arrays.asList("a", "b"));
 
         PowerMock.replayAll();
 
         server = new RestServer(workerConfig);
-        server.start(new HerderProvider(herder), herder.plugins());
+        server.initializeServer();
+        server.initializeResources(herder);
         HttpRequest request = new HttpGet("/connectors");
         request.addHeader("Referer", origin + "/page");
         request.addHeader("Origin", origin);
@@ -251,4 +273,158 @@ public void checkCORSRequest(String corsDomain, String origin, String expectedHe
         }
         PowerMock.verifyAll();
     }
+
+    @Test
+    public void testStandaloneConfig() throws IOException  {
+        Map workerProps = baseWorkerProps();
+        workerProps.put("offset.storage.file.filename", "/tmp");
+        WorkerConfig workerConfig = new StandaloneConfig(workerProps);
+
+        EasyMock.expect(herder.kafkaClusterId()).andReturn(KAFKA_CLUSTER_ID);
+        EasyMock.expect(herder.plugins()).andStubReturn(plugins);
+        EasyMock.expect(plugins.newPlugins(Collections.emptyList(),
+            workerConfig,
+            ConnectRestExtension.class)).andStubReturn(Collections.emptyList());
+
+        EasyMock.expect(herder.connectors()).andReturn(Arrays.asList("a", "b"));
+
+        PowerMock.replayAll();
+
+        server = new RestServer(workerConfig);
+        server.initializeServer();
+        server.initializeResources(herder);
+        HttpRequest request = new HttpGet("/connectors");
+        CloseableHttpClient httpClient = HttpClients.createMinimal();
+        HttpHost httpHost = new HttpHost(server.advertisedUrl().getHost(), server.advertisedUrl().getPort());
+        CloseableHttpResponse response = httpClient.execute(httpHost, request);
+
+        Assert.assertEquals(200, response.getStatusLine().getStatusCode());
+    }
+
+    @Test
+    public void testLoggersEndpointWithDefaults() throws IOException {
+        Map configMap = new HashMap<>(baseWorkerProps());
+        DistributedConfig workerConfig = new DistributedConfig(configMap);
+
+        EasyMock.expect(herder.kafkaClusterId()).andReturn(KAFKA_CLUSTER_ID);
+        EasyMock.expect(herder.plugins()).andStubReturn(plugins);
+        EasyMock.expect(plugins.newPlugins(Collections.emptyList(),
+                workerConfig,
+                ConnectRestExtension.class))
+                .andStubReturn(Collections.emptyList());
+        PowerMock.replayAll();
+
+        // create some loggers in the process
+        LoggerFactory.getLogger("a.b.c.s.W");
+
+        server = new RestServer(workerConfig);
+        server.initializeServer();
+        server.initializeResources(herder);
+
+        ObjectMapper mapper = new ObjectMapper();
+
+        String host = server.advertisedUrl().getHost();
+        int port = server.advertisedUrl().getPort();
+
+        executePut(host, port, "/admin/loggers/a.b.c.s.W", "{\"level\": \"INFO\"}");
+
+        String responseStr = executeGet(host, port, "/admin/loggers");
+        Map> loggers = mapper.readValue(responseStr, new TypeReference>>() {
+        });
+        assertNotNull("expected non null response for /admin/loggers" + prettyPrint(loggers), loggers);
+        assertTrue("expect at least 1 logger. instead found " + prettyPrint(loggers), loggers.size() >= 1);
+        assertEquals("expected to find logger a.b.c.s.W set to INFO level", loggers.get("a.b.c.s.W").get("level"), "INFO");
+    }
+
+    @Test
+    public void testIndependentAdminEndpoint() throws IOException {
+        Map configMap = new HashMap<>(baseWorkerProps());
+        configMap.put(ADMIN_LISTENERS_CONFIG, "http://localhost:0");
+
+        DistributedConfig workerConfig = new DistributedConfig(configMap);
+
+        EasyMock.expect(herder.kafkaClusterId()).andReturn(KAFKA_CLUSTER_ID);
+        EasyMock.expect(herder.plugins()).andStubReturn(plugins);
+        EasyMock.expect(plugins.newPlugins(Collections.emptyList(),
+                workerConfig,
+                ConnectRestExtension.class))
+                .andStubReturn(Collections.emptyList());
+        PowerMock.replayAll();
+
+        // create some loggers in the process
+        LoggerFactory.getLogger("a.b.c.s.W");
+        LoggerFactory.getLogger("a.b.c.p.X");
+        LoggerFactory.getLogger("a.b.c.p.Y");
+        LoggerFactory.getLogger("a.b.c.p.Z");
+
+        server = new RestServer(workerConfig);
+        server.initializeServer();
+        server.initializeResources(herder);
+
+        assertNotEquals(server.advertisedUrl(), server.adminUrl());
+
+        executeGet(server.adminUrl().getHost(), server.adminUrl().getPort(), "/admin/loggers");
+
+        HttpRequest request = new HttpGet("/admin/loggers");
+        CloseableHttpClient httpClient = HttpClients.createMinimal();
+        HttpHost httpHost = new HttpHost(server.advertisedUrl().getHost(), server.advertisedUrl().getPort());
+        CloseableHttpResponse response = httpClient.execute(httpHost, request);
+        Assert.assertEquals(404, response.getStatusLine().getStatusCode());
+    }
+
+    @Test
+    public void testDisableAdminEndpoint() throws IOException {
+        Map configMap = new HashMap<>(baseWorkerProps());
+        configMap.put(ADMIN_LISTENERS_CONFIG, "");
+
+        DistributedConfig workerConfig = new DistributedConfig(configMap);
+
+        EasyMock.expect(herder.kafkaClusterId()).andReturn(KAFKA_CLUSTER_ID);
+        EasyMock.expect(herder.plugins()).andStubReturn(plugins);
+        EasyMock.expect(plugins.newPlugins(Collections.emptyList(),
+                workerConfig,
+                ConnectRestExtension.class))
+                .andStubReturn(Collections.emptyList());
+        PowerMock.replayAll();
+
+        server = new RestServer(workerConfig);
+        server.initializeServer();
+        server.initializeResources(herder);
+
+        assertNull(server.adminUrl());
+
+        HttpRequest request = new HttpGet("/admin/loggers");
+        CloseableHttpClient httpClient = HttpClients.createMinimal();
+        HttpHost httpHost = new HttpHost(server.advertisedUrl().getHost(), server.advertisedUrl().getPort());
+        CloseableHttpResponse response = httpClient.execute(httpHost, request);
+        Assert.assertEquals(404, response.getStatusLine().getStatusCode());
+    }
+
+    private String executeGet(String host, int port, String endpoint) throws IOException {
+        HttpRequest request = new HttpGet(endpoint);
+        CloseableHttpClient httpClient = HttpClients.createMinimal();
+        HttpHost httpHost = new HttpHost(host, port);
+        CloseableHttpResponse response = httpClient.execute(httpHost, request);
+
+        Assert.assertEquals(200, response.getStatusLine().getStatusCode());
+        return new BasicResponseHandler().handleResponse(response);
+    }
+
+    private String executePut(String host, int port, String endpoint, String jsonBody) throws IOException {
+        HttpPut request = new HttpPut(endpoint);
+        StringEntity entity = new StringEntity(jsonBody, "UTF-8");
+        entity.setContentType("application/json");
+        request.setEntity(entity);
+        CloseableHttpClient httpClient = HttpClients.createMinimal();
+        HttpHost httpHost = new HttpHost(host, port);
+        CloseableHttpResponse response = httpClient.execute(httpHost, request);
+
+        Assert.assertEquals(200, response.getStatusLine().getStatusCode());
+        return new BasicResponseHandler().handleResponse(response);
+    }
+
+    private static String prettyPrint(Map map) throws IOException {
+        ObjectMapper mapper = new ObjectMapper();
+        return mapper.writerWithDefaultPrettyPrinter().writeValueAsString(map);
+    }
 }
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/rest/resources/ConnectorPluginsResourceTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/rest/resources/ConnectorPluginsResourceTest.java
index 684064d30df4d..2882355eccbec 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/rest/resources/ConnectorPluginsResourceTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/rest/resources/ConnectorPluginsResourceTest.java
@@ -18,6 +18,7 @@
 
 import com.fasterxml.jackson.core.type.TypeReference;
 import com.fasterxml.jackson.databind.ObjectMapper;
+import javax.ws.rs.core.HttpHeaders;
 import org.apache.kafka.common.config.Config;
 import org.apache.kafka.common.config.ConfigDef;
 import org.apache.kafka.common.config.ConfigDef.Importance;
@@ -30,7 +31,6 @@
 import org.apache.kafka.connect.runtime.AbstractHerder;
 import org.apache.kafka.connect.runtime.ConnectorConfig;
 import org.apache.kafka.connect.runtime.Herder;
-import org.apache.kafka.connect.runtime.HerderProvider;
 import org.apache.kafka.connect.runtime.TestSinkConnector;
 import org.apache.kafka.connect.runtime.TestSourceConnector;
 import org.apache.kafka.connect.runtime.WorkerConfig;
@@ -183,11 +183,11 @@ public class ConnectorPluginsResourceTest {
     @Before
     public void setUp() throws Exception {
         PowerMock.mockStatic(RestClient.class,
-                RestClient.class.getMethod("httpRequest", String.class, String.class, Object.class, TypeReference.class, WorkerConfig.class));
+                RestClient.class.getMethod("httpRequest", String.class, String.class, HttpHeaders.class, Object.class, TypeReference.class, WorkerConfig.class));
 
         plugins = PowerMock.createMock(Plugins.class);
         herder = PowerMock.createMock(AbstractHerder.class);
-        connectorPluginsResource = new ConnectorPluginsResource(new HerderProvider(herder));
+        connectorPluginsResource = new ConnectorPluginsResource(herder);
     }
 
     private void expectPlugins() {
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/rest/resources/ConnectorsResourceTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/rest/resources/ConnectorsResourceTest.java
index 5a520744bcd26..f604c80ceec69 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/rest/resources/ConnectorsResourceTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/rest/resources/ConnectorsResourceTest.java
@@ -18,20 +18,25 @@
 
 import com.fasterxml.jackson.core.type.TypeReference;
 
+import javax.crypto.Mac;
+import javax.ws.rs.core.HttpHeaders;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
 import org.apache.kafka.connect.errors.AlreadyExistsException;
-import org.apache.kafka.connect.errors.ConnectException;
 import org.apache.kafka.connect.errors.NotFoundException;
 import org.apache.kafka.connect.runtime.ConnectorConfig;
 import org.apache.kafka.connect.runtime.Herder;
-import org.apache.kafka.connect.runtime.HerderProvider;
 import org.apache.kafka.connect.runtime.WorkerConfig;
 import org.apache.kafka.connect.runtime.distributed.NotAssignedException;
 import org.apache.kafka.connect.runtime.distributed.NotLeaderException;
+import org.apache.kafka.connect.runtime.rest.InternalRequestSignature;
 import org.apache.kafka.connect.runtime.rest.RestClient;
 import org.apache.kafka.connect.runtime.rest.entities.ConnectorInfo;
+import org.apache.kafka.connect.runtime.rest.entities.ConnectorStateInfo;
 import org.apache.kafka.connect.runtime.rest.entities.ConnectorType;
 import org.apache.kafka.connect.runtime.rest.entities.CreateConnectorRequest;
 import org.apache.kafka.connect.runtime.rest.entities.TaskInfo;
+import org.apache.kafka.connect.runtime.rest.errors.ConnectRestException;
 import org.apache.kafka.connect.util.Callback;
 import org.apache.kafka.connect.util.ConnectorTaskId;
 import org.easymock.Capture;
@@ -48,9 +53,15 @@
 import org.powermock.modules.junit4.PowerMockRunner;
 
 import javax.ws.rs.BadRequestException;
+import javax.ws.rs.core.MultivaluedHashMap;
+import javax.ws.rs.core.MultivaluedMap;
+import javax.ws.rs.core.UriInfo;
+
+import java.io.IOException;
 import java.net.URI;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Base64;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
@@ -59,10 +70,12 @@
 import java.util.Map;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
 
 @RunWith(PowerMockRunner.class)
 @PrepareForTest(RestClient.class)
-@PowerMockIgnore("javax.management.*")
+@PowerMockIgnore({"javax.management.*", "javax.crypto.*"})
 @SuppressWarnings("unchecked")
 public class ConnectorsResourceTest {
     // Note trailing / and that we do *not* use LEADER_URL to construct our reference values. This checks that we handle
@@ -76,6 +89,7 @@ public class ConnectorsResourceTest {
     private static final String CONNECTOR_NAME_PADDING_WHITESPACES = "   " + CONNECTOR_NAME + "  \n  ";
     private static final Boolean FORWARD = true;
     private static final Map CONNECTOR_CONFIG_SPECIAL_CHARS = new HashMap<>();
+    private static final HttpHeaders NULL_HEADERS = null;
     static {
         CONNECTOR_CONFIG_SPECIAL_CHARS.put("name", CONNECTOR_NAME_SPECIAL_CHARS);
         CONNECTOR_CONFIG_SPECIAL_CHARS.put("sample_config", "test_config");
@@ -122,12 +136,18 @@ public class ConnectorsResourceTest {
     @Mock
     private Herder herder;
     private ConnectorsResource connectorsResource;
+    private UriInfo forward;
 
     @Before
     public void setUp() throws NoSuchMethodException {
         PowerMock.mockStatic(RestClient.class,
-                RestClient.class.getMethod("httpRequest", String.class, String.class, Object.class, TypeReference.class, WorkerConfig.class));
-        connectorsResource = new ConnectorsResource(new HerderProvider(herder), null);
+                RestClient.class.getMethod("httpRequest", String.class, String.class, HttpHeaders.class, Object.class, TypeReference.class, WorkerConfig.class));
+        connectorsResource = new ConnectorsResource(herder, null);
+        forward = EasyMock.mock(UriInfo.class);
+        MultivaluedMap queryParams = new MultivaluedHashMap<>();
+        queryParams.putSingle("forward", "true");
+        EasyMock.expect(forward.getQueryParameters()).andReturn(queryParams).anyTimes();
+        EasyMock.replay(forward);
     }
 
     private static final Map getConnectorConfig(Map mapToClone) {
@@ -138,12 +158,11 @@ private static final Map getConnectorConfig(Map
     @Test
     public void testListConnectors() throws Throwable {
         final Capture>> cb = Capture.newInstance();
-        herder.connectors(EasyMock.capture(cb));
-        expectAndCallbackResult(cb, Arrays.asList(CONNECTOR2_NAME, CONNECTOR_NAME));
+        EasyMock.expect(herder.connectors()).andReturn(Arrays.asList(CONNECTOR2_NAME, CONNECTOR_NAME));
 
         PowerMock.replayAll();
 
-        Collection connectors = connectorsResource.listConnectors(FORWARD);
+        Collection connectors = (Collection) connectorsResource.listConnectors(forward, NULL_HEADERS).getEntity();
         // Ordering isn't guaranteed, compare sets
         assertEquals(new HashSet<>(Arrays.asList(CONNECTOR_NAME, CONNECTOR2_NAME)), new HashSet<>(connectors));
 
@@ -151,36 +170,107 @@ public void testListConnectors() throws Throwable {
     }
 
     @Test
-    public void testListConnectorsNotLeader() throws Throwable {
-        final Capture>> cb = Capture.newInstance();
-        herder.connectors(EasyMock.capture(cb));
-        expectAndCallbackNotLeaderException(cb);
-        // Should forward request
-        EasyMock.expect(RestClient.httpRequest(EasyMock.eq("http://leader:8083/connectors?forward=false"), EasyMock.eq("GET"),
-                EasyMock.isNull(), EasyMock.anyObject(TypeReference.class), EasyMock.anyObject(WorkerConfig.class)))
-                .andReturn(new RestClient.HttpResponse<>(200, new HashMap(), Arrays.asList(CONNECTOR2_NAME, CONNECTOR_NAME)));
+    public void testExpandConnectorsStatus() throws Throwable {
+        EasyMock.expect(herder.connectors()).andReturn(Arrays.asList(CONNECTOR2_NAME, CONNECTOR_NAME));
+        ConnectorStateInfo connector = EasyMock.mock(ConnectorStateInfo.class);
+        ConnectorStateInfo connector2 = EasyMock.mock(ConnectorStateInfo.class);
+        EasyMock.expect(herder.connectorStatus(CONNECTOR2_NAME)).andReturn(connector2);
+        EasyMock.expect(herder.connectorStatus(CONNECTOR_NAME)).andReturn(connector);
+
+        forward = EasyMock.mock(UriInfo.class);
+        MultivaluedMap queryParams = new MultivaluedHashMap<>();
+        queryParams.putSingle("expand", "status");
+        EasyMock.expect(forward.getQueryParameters()).andReturn(queryParams).anyTimes();
+        EasyMock.replay(forward);
 
         PowerMock.replayAll();
 
-        Collection connectors = connectorsResource.listConnectors(FORWARD);
+        Map> expanded = (Map>) connectorsResource.listConnectors(forward, NULL_HEADERS).getEntity();
         // Ordering isn't guaranteed, compare sets
-        assertEquals(new HashSet<>(Arrays.asList(CONNECTOR_NAME, CONNECTOR2_NAME)), new HashSet<>(connectors));
+        assertEquals(new HashSet<>(Arrays.asList(CONNECTOR_NAME, CONNECTOR2_NAME)), expanded.keySet());
+        assertEquals(connector2, expanded.get(CONNECTOR2_NAME).get("status"));
+        assertEquals(connector, expanded.get(CONNECTOR_NAME).get("status"));
+        PowerMock.verifyAll();
+    }
+
+    @Test
+    public void testExpandConnectorsInfo() throws Throwable {
+        EasyMock.expect(herder.connectors()).andReturn(Arrays.asList(CONNECTOR2_NAME, CONNECTOR_NAME));
+        ConnectorInfo connector = EasyMock.mock(ConnectorInfo.class);
+        ConnectorInfo connector2 = EasyMock.mock(ConnectorInfo.class);
+        EasyMock.expect(herder.connectorInfo(CONNECTOR2_NAME)).andReturn(connector2);
+        EasyMock.expect(herder.connectorInfo(CONNECTOR_NAME)).andReturn(connector);
+
+        forward = EasyMock.mock(UriInfo.class);
+        MultivaluedMap queryParams = new MultivaluedHashMap<>();
+        queryParams.putSingle("expand", "info");
+        EasyMock.expect(forward.getQueryParameters()).andReturn(queryParams).anyTimes();
+        EasyMock.replay(forward);
 
+        PowerMock.replayAll();
+
+        Map> expanded = (Map>) connectorsResource.listConnectors(forward, NULL_HEADERS).getEntity();
+        // Ordering isn't guaranteed, compare sets
+        assertEquals(new HashSet<>(Arrays.asList(CONNECTOR_NAME, CONNECTOR2_NAME)), expanded.keySet());
+        assertEquals(connector2, expanded.get(CONNECTOR2_NAME).get("info"));
+        assertEquals(connector, expanded.get(CONNECTOR_NAME).get("info"));
         PowerMock.verifyAll();
     }
 
-    @Test(expected = ConnectException.class)
-    public void testListConnectorsNotSynced() throws Throwable {
-        final Capture>> cb = Capture.newInstance();
-        herder.connectors(EasyMock.capture(cb));
-        expectAndCallbackException(cb, new ConnectException("not synced"));
+    @Test
+    public void testFullExpandConnectors() throws Throwable {
+        EasyMock.expect(herder.connectors()).andReturn(Arrays.asList(CONNECTOR2_NAME, CONNECTOR_NAME));
+        ConnectorInfo connectorInfo = EasyMock.mock(ConnectorInfo.class);
+        ConnectorInfo connectorInfo2 = EasyMock.mock(ConnectorInfo.class);
+        EasyMock.expect(herder.connectorInfo(CONNECTOR2_NAME)).andReturn(connectorInfo2);
+        EasyMock.expect(herder.connectorInfo(CONNECTOR_NAME)).andReturn(connectorInfo);
+        ConnectorStateInfo connector = EasyMock.mock(ConnectorStateInfo.class);
+        ConnectorStateInfo connector2 = EasyMock.mock(ConnectorStateInfo.class);
+        EasyMock.expect(herder.connectorStatus(CONNECTOR2_NAME)).andReturn(connector2);
+        EasyMock.expect(herder.connectorStatus(CONNECTOR_NAME)).andReturn(connector);
+
+        forward = EasyMock.mock(UriInfo.class);
+        MultivaluedMap queryParams = new MultivaluedHashMap<>();
+        queryParams.put("expand", Arrays.asList("info", "status"));
+        EasyMock.expect(forward.getQueryParameters()).andReturn(queryParams).anyTimes();
+        EasyMock.replay(forward);
 
         PowerMock.replayAll();
 
-        // throws
-        connectorsResource.listConnectors(FORWARD);
+        Map> expanded = (Map>) connectorsResource.listConnectors(forward, NULL_HEADERS).getEntity();
+        // Ordering isn't guaranteed, compare sets
+        assertEquals(new HashSet<>(Arrays.asList(CONNECTOR_NAME, CONNECTOR2_NAME)), expanded.keySet());
+        assertEquals(connectorInfo2, expanded.get(CONNECTOR2_NAME).get("info"));
+        assertEquals(connectorInfo, expanded.get(CONNECTOR_NAME).get("info"));
+        assertEquals(connector2, expanded.get(CONNECTOR2_NAME).get("status"));
+        assertEquals(connector, expanded.get(CONNECTOR_NAME).get("status"));
+        PowerMock.verifyAll();
     }
 
+    @Test
+    public void testExpandConnectorsWithConnectorNotFound() throws Throwable {
+        EasyMock.expect(herder.connectors()).andReturn(Arrays.asList(CONNECTOR2_NAME, CONNECTOR_NAME));
+        ConnectorStateInfo connector = EasyMock.mock(ConnectorStateInfo.class);
+        ConnectorStateInfo connector2 = EasyMock.mock(ConnectorStateInfo.class);
+        EasyMock.expect(herder.connectorStatus(CONNECTOR2_NAME)).andReturn(connector2);
+        EasyMock.expect(herder.connectorStatus(CONNECTOR_NAME)).andThrow(EasyMock.mock(NotFoundException.class));
+
+        forward = EasyMock.mock(UriInfo.class);
+        MultivaluedMap queryParams = new MultivaluedHashMap<>();
+        queryParams.putSingle("expand", "status");
+        EasyMock.expect(forward.getQueryParameters()).andReturn(queryParams).anyTimes();
+        EasyMock.replay(forward);
+
+        PowerMock.replayAll();
+
+        Map> expanded = (Map>) connectorsResource.listConnectors(forward, NULL_HEADERS).getEntity();
+        // Ordering isn't guaranteed, compare sets
+        assertEquals(Collections.singleton(CONNECTOR2_NAME), expanded.keySet());
+        assertEquals(connector2, expanded.get(CONNECTOR2_NAME).get("status"));
+        PowerMock.verifyAll();
+    }
+
+
     @Test
     public void testCreateConnector() throws Throwable {
         CreateConnectorRequest body = new CreateConnectorRequest(CONNECTOR_NAME, Collections.singletonMap(ConnectorConfig.NAME_CONFIG, CONNECTOR_NAME));
@@ -192,7 +282,7 @@ public void testCreateConnector() throws Throwable {
 
         PowerMock.replayAll();
 
-        connectorsResource.createConnector(FORWARD, body);
+        connectorsResource.createConnector(FORWARD, NULL_HEADERS, body);
 
         PowerMock.verifyAll();
     }
@@ -205,19 +295,57 @@ public void testCreateConnectorNotLeader() throws Throwable {
         herder.putConnectorConfig(EasyMock.eq(CONNECTOR_NAME), EasyMock.eq(body.config()), EasyMock.eq(false), EasyMock.capture(cb));
         expectAndCallbackNotLeaderException(cb);
         // Should forward request
-        EasyMock.expect(RestClient.httpRequest(EasyMock.eq("http://leader:8083/connectors?forward=false"), EasyMock.eq("POST"), EasyMock.eq(body), EasyMock.anyObject(), EasyMock.anyObject(WorkerConfig.class)))
+        EasyMock.expect(RestClient.httpRequest(EasyMock.eq("http://leader:8083/connectors?forward=false"), EasyMock.eq("POST"), EasyMock.isNull(), EasyMock.eq(body), EasyMock.anyObject(), EasyMock.anyObject(WorkerConfig.class)))
                 .andReturn(new RestClient.HttpResponse<>(201, new HashMap(), new ConnectorInfo(CONNECTOR_NAME, CONNECTOR_CONFIG, CONNECTOR_TASK_NAMES,
                     ConnectorType.SOURCE)));
 
         PowerMock.replayAll();
 
-        connectorsResource.createConnector(FORWARD, body);
+        connectorsResource.createConnector(FORWARD, NULL_HEADERS, body);
 
         PowerMock.verifyAll();
 
 
     }
 
+    @Test
+    public void testCreateConnectorWithHeaderAuthorization() throws Throwable {
+        CreateConnectorRequest body = new CreateConnectorRequest(CONNECTOR_NAME, Collections.singletonMap(ConnectorConfig.NAME_CONFIG, CONNECTOR_NAME));
+        final Capture>> cb = Capture.newInstance();
+        HttpHeaders httpHeaders = EasyMock.mock(HttpHeaders.class);
+        EasyMock.expect(httpHeaders.getHeaderString("Authorization")).andReturn("Basic YWxhZGRpbjpvcGVuc2VzYW1l").times(1);
+        EasyMock.replay(httpHeaders);
+        herder.putConnectorConfig(EasyMock.eq(CONNECTOR_NAME), EasyMock.eq(body.config()), EasyMock.eq(false), EasyMock.capture(cb));
+        expectAndCallbackResult(cb, new Herder.Created<>(true, new ConnectorInfo(CONNECTOR_NAME, CONNECTOR_CONFIG,
+            CONNECTOR_TASK_NAMES, ConnectorType.SOURCE)));
+
+        PowerMock.replayAll();
+
+        connectorsResource.createConnector(FORWARD, httpHeaders, body);
+
+        PowerMock.verifyAll();
+    }
+
+
+
+    @Test
+    public void testCreateConnectorWithoutHeaderAuthorization() throws Throwable {
+        CreateConnectorRequest body = new CreateConnectorRequest(CONNECTOR_NAME, Collections.singletonMap(ConnectorConfig.NAME_CONFIG, CONNECTOR_NAME));
+        final Capture>> cb = Capture.newInstance();
+        HttpHeaders httpHeaders = EasyMock.mock(HttpHeaders.class);
+        EasyMock.expect(httpHeaders.getHeaderString("Authorization")).andReturn(null).times(1);
+        EasyMock.replay(httpHeaders);
+        herder.putConnectorConfig(EasyMock.eq(CONNECTOR_NAME), EasyMock.eq(body.config()), EasyMock.eq(false), EasyMock.capture(cb));
+        expectAndCallbackResult(cb, new Herder.Created<>(true, new ConnectorInfo(CONNECTOR_NAME, CONNECTOR_CONFIG,
+            CONNECTOR_TASK_NAMES, ConnectorType.SOURCE)));
+
+        PowerMock.replayAll();
+
+        connectorsResource.createConnector(FORWARD, httpHeaders, body);
+
+        PowerMock.verifyAll();
+    }
+
     @Test(expected = AlreadyExistsException.class)
     public void testCreateConnectorExists() throws Throwable {
         CreateConnectorRequest body = new CreateConnectorRequest(CONNECTOR_NAME, Collections.singletonMap(ConnectorConfig.NAME_CONFIG, CONNECTOR_NAME));
@@ -228,7 +356,7 @@ public void testCreateConnectorExists() throws Throwable {
 
         PowerMock.replayAll();
 
-        connectorsResource.createConnector(FORWARD, body);
+        connectorsResource.createConnector(FORWARD, NULL_HEADERS, body);
 
         PowerMock.verifyAll();
     }
@@ -247,7 +375,7 @@ public void testCreateConnectorNameTrimWhitespaces() throws Throwable {
 
         PowerMock.replayAll();
 
-        connectorsResource.createConnector(FORWARD, bodyIn);
+        connectorsResource.createConnector(FORWARD, NULL_HEADERS, bodyIn);
 
         PowerMock.verifyAll();
     }
@@ -266,7 +394,7 @@ public void testCreateConnectorNameAllWhitespaces() throws Throwable {
 
         PowerMock.replayAll();
 
-        connectorsResource.createConnector(FORWARD, bodyIn);
+        connectorsResource.createConnector(FORWARD, NULL_HEADERS, bodyIn);
 
         PowerMock.verifyAll();
     }
@@ -285,7 +413,7 @@ public void testCreateConnectorNoName() throws Throwable {
 
         PowerMock.replayAll();
 
-        connectorsResource.createConnector(FORWARD, bodyIn);
+        connectorsResource.createConnector(FORWARD, NULL_HEADERS, bodyIn);
 
         PowerMock.verifyAll();
     }
@@ -298,7 +426,7 @@ public void testDeleteConnector() throws Throwable {
 
         PowerMock.replayAll();
 
-        connectorsResource.destroyConnector(CONNECTOR_NAME, FORWARD);
+        connectorsResource.destroyConnector(CONNECTOR_NAME, NULL_HEADERS, FORWARD);
 
         PowerMock.verifyAll();
     }
@@ -309,12 +437,12 @@ public void testDeleteConnectorNotLeader() throws Throwable {
         herder.deleteConnectorConfig(EasyMock.eq(CONNECTOR_NAME), EasyMock.capture(cb));
         expectAndCallbackNotLeaderException(cb);
         // Should forward request
-        EasyMock.expect(RestClient.httpRequest("http://leader:8083/connectors/" + CONNECTOR_NAME + "?forward=false", "DELETE", null, null, null))
+        EasyMock.expect(RestClient.httpRequest("http://leader:8083/connectors/" + CONNECTOR_NAME + "?forward=false", "DELETE", NULL_HEADERS, null, null, null))
                 .andReturn(new RestClient.HttpResponse<>(204, new HashMap(), null));
 
         PowerMock.replayAll();
 
-        connectorsResource.destroyConnector(CONNECTOR_NAME, FORWARD);
+        connectorsResource.destroyConnector(CONNECTOR_NAME, NULL_HEADERS, FORWARD);
 
         PowerMock.verifyAll();
     }
@@ -328,7 +456,7 @@ public void testDeleteConnectorNotFound() throws Throwable {
 
         PowerMock.replayAll();
 
-        connectorsResource.destroyConnector(CONNECTOR_NAME, FORWARD);
+        connectorsResource.destroyConnector(CONNECTOR_NAME, NULL_HEADERS, FORWARD);
 
         PowerMock.verifyAll();
     }
@@ -342,7 +470,7 @@ public void testGetConnector() throws Throwable {
 
         PowerMock.replayAll();
 
-        ConnectorInfo connInfo = connectorsResource.getConnector(CONNECTOR_NAME, FORWARD);
+        ConnectorInfo connInfo = connectorsResource.getConnector(CONNECTOR_NAME, NULL_HEADERS, FORWARD);
         assertEquals(new ConnectorInfo(CONNECTOR_NAME, CONNECTOR_CONFIG, CONNECTOR_TASK_NAMES, ConnectorType.SOURCE),
             connInfo);
 
@@ -357,7 +485,7 @@ public void testGetConnectorConfig() throws Throwable {
 
         PowerMock.replayAll();
 
-        Map connConfig = connectorsResource.getConnectorConfig(CONNECTOR_NAME, FORWARD);
+        Map connConfig = connectorsResource.getConnectorConfig(CONNECTOR_NAME, NULL_HEADERS, FORWARD);
         assertEquals(CONNECTOR_CONFIG, connConfig);
 
         PowerMock.verifyAll();
@@ -371,7 +499,7 @@ public void testGetConnectorConfigConnectorNotFound() throws Throwable {
 
         PowerMock.replayAll();
 
-        connectorsResource.getConnectorConfig(CONNECTOR_NAME, FORWARD);
+        connectorsResource.getConnectorConfig(CONNECTOR_NAME, NULL_HEADERS, FORWARD);
 
         PowerMock.verifyAll();
     }
@@ -385,7 +513,7 @@ public void testPutConnectorConfig() throws Throwable {
 
         PowerMock.replayAll();
 
-        connectorsResource.putConnectorConfig(CONNECTOR_NAME, FORWARD, CONNECTOR_CONFIG);
+        connectorsResource.putConnectorConfig(CONNECTOR_NAME, NULL_HEADERS, FORWARD, CONNECTOR_CONFIG);
 
         PowerMock.verifyAll();
     }
@@ -401,7 +529,7 @@ public void testCreateConnectorWithSpecialCharsInName() throws Throwable {
 
         PowerMock.replayAll();
 
-        String rspLocation = connectorsResource.createConnector(FORWARD, body).getLocation().toString();
+        String rspLocation = connectorsResource.createConnector(FORWARD, NULL_HEADERS, body).getLocation().toString();
         String decoded = new URI(rspLocation).getPath();
         Assert.assertEquals("/connectors/" + CONNECTOR_NAME_SPECIAL_CHARS, decoded);
 
@@ -419,7 +547,7 @@ public void testCreateConnectorWithControlSequenceInName() throws Throwable {
 
         PowerMock.replayAll();
 
-        String rspLocation = connectorsResource.createConnector(FORWARD, body).getLocation().toString();
+        String rspLocation = connectorsResource.createConnector(FORWARD, NULL_HEADERS, body).getLocation().toString();
         String decoded = new URI(rspLocation).getPath();
         Assert.assertEquals("/connectors/" + CONNECTOR_NAME_CONTROL_SEQUENCES1, decoded);
 
@@ -436,7 +564,7 @@ public void testPutConnectorConfigWithSpecialCharsInName() throws Throwable {
 
         PowerMock.replayAll();
 
-        String rspLocation = connectorsResource.putConnectorConfig(CONNECTOR_NAME_SPECIAL_CHARS, FORWARD, CONNECTOR_CONFIG_SPECIAL_CHARS).getLocation().toString();
+        String rspLocation = connectorsResource.putConnectorConfig(CONNECTOR_NAME_SPECIAL_CHARS, NULL_HEADERS, FORWARD, CONNECTOR_CONFIG_SPECIAL_CHARS).getLocation().toString();
         String decoded = new URI(rspLocation).getPath();
         Assert.assertEquals("/connectors/" + CONNECTOR_NAME_SPECIAL_CHARS, decoded);
 
@@ -453,7 +581,7 @@ public void testPutConnectorConfigWithControlSequenceInName() throws Throwable {
 
         PowerMock.replayAll();
 
-        String rspLocation = connectorsResource.putConnectorConfig(CONNECTOR_NAME_CONTROL_SEQUENCES1, FORWARD, CONNECTOR_CONFIG_CONTROL_SEQUENCES).getLocation().toString();
+        String rspLocation = connectorsResource.putConnectorConfig(CONNECTOR_NAME_CONTROL_SEQUENCES1, NULL_HEADERS, FORWARD, CONNECTOR_CONFIG_CONTROL_SEQUENCES).getLocation().toString();
         String decoded = new URI(rspLocation).getPath();
         Assert.assertEquals("/connectors/" + CONNECTOR_NAME_CONTROL_SEQUENCES1, decoded);
 
@@ -464,7 +592,7 @@ public void testPutConnectorConfigWithControlSequenceInName() throws Throwable {
     public void testPutConnectorConfigNameMismatch() throws Throwable {
         Map connConfig = new HashMap<>(CONNECTOR_CONFIG);
         connConfig.put(ConnectorConfig.NAME_CONFIG, "mismatched-name");
-        connectorsResource.putConnectorConfig(CONNECTOR_NAME, FORWARD, connConfig);
+        connectorsResource.putConnectorConfig(CONNECTOR_NAME, NULL_HEADERS, FORWARD, connConfig);
     }
 
     @Test(expected = BadRequestException.class)
@@ -472,7 +600,7 @@ public void testCreateConnectorConfigNameMismatch() throws Throwable {
         Map connConfig = new HashMap<>();
         connConfig.put(ConnectorConfig.NAME_CONFIG, "mismatched-name");
         CreateConnectorRequest request = new CreateConnectorRequest(CONNECTOR_NAME, connConfig);
-        connectorsResource.createConnector(FORWARD, request);
+        connectorsResource.createConnector(FORWARD, NULL_HEADERS, request);
     }
 
     @Test
@@ -483,7 +611,7 @@ public void testGetConnectorTaskConfigs() throws Throwable {
 
         PowerMock.replayAll();
 
-        List taskInfos = connectorsResource.getTaskConfigs(CONNECTOR_NAME, FORWARD);
+        List taskInfos = connectorsResource.getTaskConfigs(CONNECTOR_NAME, NULL_HEADERS, FORWARD);
         assertEquals(TASK_INFOS, taskInfos);
 
         PowerMock.verifyAll();
@@ -497,33 +625,82 @@ public void testGetConnectorTaskConfigsConnectorNotFound() throws Throwable {
 
         PowerMock.replayAll();
 
-        connectorsResource.getTaskConfigs(CONNECTOR_NAME, FORWARD);
+        connectorsResource.getTaskConfigs(CONNECTOR_NAME, NULL_HEADERS, FORWARD);
 
         PowerMock.verifyAll();
     }
 
     @Test
-    public void testPutConnectorTaskConfigs() throws Throwable {
+    public void testPutConnectorTaskConfigsNoInternalRequestSignature() throws Throwable {
         final Capture> cb = Capture.newInstance();
-        herder.putTaskConfigs(EasyMock.eq(CONNECTOR_NAME), EasyMock.eq(TASK_CONFIGS), EasyMock.capture(cb));
+        herder.putTaskConfigs(
+            EasyMock.eq(CONNECTOR_NAME),
+            EasyMock.eq(TASK_CONFIGS),
+            EasyMock.capture(cb),
+            EasyMock.anyObject(InternalRequestSignature.class)
+        );
         expectAndCallbackResult(cb, null);
 
         PowerMock.replayAll();
 
-        connectorsResource.putTaskConfigs(CONNECTOR_NAME, FORWARD, TASK_CONFIGS);
+        connectorsResource.putTaskConfigs(CONNECTOR_NAME, NULL_HEADERS, FORWARD, serializeAsBytes(TASK_CONFIGS));
+
+        PowerMock.verifyAll();
+    }
+
+    @Test
+    public void testPutConnectorTaskConfigsWithInternalRequestSignature() throws Throwable {
+        final String signatureAlgorithm = "HmacSHA256";
+        final String encodedSignature = "Kv1/OSsxzdVIwvZ4e30avyRIVrngDfhzVUm/kAZEKc4=";
+
+        final Capture> cb = Capture.newInstance();
+        final Capture signatureCapture = Capture.newInstance();
+        herder.putTaskConfigs(
+            EasyMock.eq(CONNECTOR_NAME),
+            EasyMock.eq(TASK_CONFIGS),
+            EasyMock.capture(cb),
+            EasyMock.capture(signatureCapture)
+        );
+        expectAndCallbackResult(cb, null);
+
+        HttpHeaders headers = EasyMock.mock(HttpHeaders.class);
+        EasyMock.expect(headers.getHeaderString(InternalRequestSignature.SIGNATURE_ALGORITHM_HEADER))
+            .andReturn(signatureAlgorithm)
+            .once();
+        EasyMock.expect(headers.getHeaderString(InternalRequestSignature.SIGNATURE_HEADER))
+            .andReturn(encodedSignature)
+            .once();
+
+        PowerMock.replayAll(headers);
+
+        connectorsResource.putTaskConfigs(CONNECTOR_NAME, headers, FORWARD, serializeAsBytes(TASK_CONFIGS));
 
         PowerMock.verifyAll();
+        InternalRequestSignature expectedSignature = new InternalRequestSignature(
+            serializeAsBytes(TASK_CONFIGS),
+            Mac.getInstance(signatureAlgorithm),
+            Base64.getDecoder().decode(encodedSignature)
+        );
+        assertEquals(
+            expectedSignature,
+            signatureCapture.getValue()
+        );
     }
 
     @Test(expected = NotFoundException.class)
     public void testPutConnectorTaskConfigsConnectorNotFound() throws Throwable {
         final Capture> cb = Capture.newInstance();
-        herder.putTaskConfigs(EasyMock.eq(CONNECTOR_NAME), EasyMock.eq(TASK_CONFIGS), EasyMock.capture(cb));
+        herder.putTaskConfigs(
+            EasyMock.eq(CONNECTOR_NAME),
+            EasyMock.eq(TASK_CONFIGS),
+            EasyMock.capture(cb),
+            EasyMock.anyObject(InternalRequestSignature.class)
+        );
         expectAndCallbackException(cb, new NotFoundException("not found"));
 
         PowerMock.replayAll();
 
-        connectorsResource.putTaskConfigs(CONNECTOR_NAME, FORWARD, TASK_CONFIGS);
+        connectorsResource.putTaskConfigs(CONNECTOR_NAME, NULL_HEADERS, FORWARD, serializeAsBytes(TASK_CONFIGS));
 
         PowerMock.verifyAll();
     }
@@ -536,7 +713,7 @@ public void testRestartConnectorNotFound() throws Throwable {
 
         PowerMock.replayAll();
 
-        connectorsResource.restartConnector(CONNECTOR_NAME, FORWARD);
+        connectorsResource.restartConnector(CONNECTOR_NAME, NULL_HEADERS, FORWARD);
 
         PowerMock.verifyAll();
     }
@@ -548,12 +725,12 @@ public void testRestartConnectorLeaderRedirect() throws Throwable {
         expectAndCallbackNotLeaderException(cb);
 
         EasyMock.expect(RestClient.httpRequest(EasyMock.eq("http://leader:8083/connectors/" + CONNECTOR_NAME + "/restart?forward=true"),
-                EasyMock.eq("POST"), EasyMock.isNull(), EasyMock.anyObject(), EasyMock.anyObject(WorkerConfig.class)))
+                EasyMock.eq("POST"), EasyMock.isNull(), EasyMock.isNull(), EasyMock.anyObject(), EasyMock.anyObject(WorkerConfig.class)))
                 .andReturn(new RestClient.HttpResponse<>(202, new HashMap(), null));
 
         PowerMock.replayAll();
 
-        connectorsResource.restartConnector(CONNECTOR_NAME, null);
+        connectorsResource.restartConnector(CONNECTOR_NAME, NULL_HEADERS, null);
 
         PowerMock.verifyAll();
     }
@@ -566,12 +743,12 @@ public void testRestartConnectorOwnerRedirect() throws Throwable {
         expectAndCallbackException(cb, new NotAssignedException("not owner test", ownerUrl));
 
         EasyMock.expect(RestClient.httpRequest(EasyMock.eq("http://owner:8083/connectors/" + CONNECTOR_NAME + "/restart?forward=false"),
-                EasyMock.eq("POST"), EasyMock.isNull(), EasyMock.anyObject(), EasyMock.anyObject(WorkerConfig.class)))
+                EasyMock.eq("POST"), EasyMock.isNull(), EasyMock.isNull(), EasyMock.anyObject(), EasyMock.anyObject(WorkerConfig.class)))
                 .andReturn(new RestClient.HttpResponse<>(202, new HashMap(), null));
 
         PowerMock.replayAll();
 
-        connectorsResource.restartConnector(CONNECTOR_NAME, true);
+        connectorsResource.restartConnector(CONNECTOR_NAME, NULL_HEADERS, true);
 
         PowerMock.verifyAll();
     }
@@ -585,7 +762,7 @@ public void testRestartTaskNotFound() throws Throwable {
 
         PowerMock.replayAll();
 
-        connectorsResource.restartTask(CONNECTOR_NAME, 0, FORWARD);
+        connectorsResource.restartTask(CONNECTOR_NAME, 0, NULL_HEADERS, FORWARD);
 
         PowerMock.verifyAll();
     }
@@ -599,12 +776,12 @@ public void testRestartTaskLeaderRedirect() throws Throwable {
         expectAndCallbackNotLeaderException(cb);
 
         EasyMock.expect(RestClient.httpRequest(EasyMock.eq("http://leader:8083/connectors/" + CONNECTOR_NAME + "/tasks/0/restart?forward=true"),
-                EasyMock.eq("POST"), EasyMock.isNull(), EasyMock.anyObject(), EasyMock.anyObject(WorkerConfig.class)))
+                EasyMock.eq("POST"), EasyMock.isNull(), EasyMock.isNull(), EasyMock.anyObject(), EasyMock.anyObject(WorkerConfig.class)))
                 .andReturn(new RestClient.HttpResponse<>(202, new HashMap(), null));
 
         PowerMock.replayAll();
 
-        connectorsResource.restartTask(CONNECTOR_NAME, 0, null);
+        connectorsResource.restartTask(CONNECTOR_NAME, 0, NULL_HEADERS, null);
 
         PowerMock.verifyAll();
     }
@@ -619,16 +796,36 @@ public void testRestartTaskOwnerRedirect() throws Throwable {
         expectAndCallbackException(cb, new NotAssignedException("not owner test", ownerUrl));
 
         EasyMock.expect(RestClient.httpRequest(EasyMock.eq("http://owner:8083/connectors/" + CONNECTOR_NAME + "/tasks/0/restart?forward=false"),
-                EasyMock.eq("POST"), EasyMock.isNull(), EasyMock.anyObject(), EasyMock.anyObject(WorkerConfig.class)))
+                EasyMock.eq("POST"), EasyMock.isNull(), EasyMock.isNull(), EasyMock.anyObject(), EasyMock.anyObject(WorkerConfig.class)))
                 .andReturn(new RestClient.HttpResponse<>(202, new HashMap(), null));
 
         PowerMock.replayAll();
 
-        connectorsResource.restartTask(CONNECTOR_NAME, 0, true);
+        connectorsResource.restartTask(CONNECTOR_NAME, 0, NULL_HEADERS, true);
 
         PowerMock.verifyAll();
     }
 
+    @Test
+    public void testCompleteOrForwardWithErrorAndNoForwardUrl() throws Throwable {
+        final Capture>> cb = Capture.newInstance();
+        herder.deleteConnectorConfig(EasyMock.eq(CONNECTOR_NAME), EasyMock.capture(cb));
+        String leaderUrl = null;
+        expectAndCallbackException(cb, new NotLeaderException("not leader", leaderUrl));
+
+        PowerMock.replayAll();
+
+        ConnectRestException e = assertThrows(ConnectRestException.class, () -> {
+            connectorsResource.destroyConnector(CONNECTOR_NAME, NULL_HEADERS, FORWARD);
+        });
+        assertTrue(e.getMessage().contains("no known leader URL"));
+        PowerMock.verifyAll();
+    }
+
+    private  byte[] serializeAsBytes(final T value) throws IOException {
+        return new ObjectMapper().writeValueAsBytes(value);
+    }
+
     private   void expectAndCallbackResult(final Capture> cb, final T value) {
         PowerMock.expectLastCall().andAnswer(new IAnswer() {
             @Override
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/rest/resources/LoggingResourceTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/rest/resources/LoggingResourceTest.java
new file mode 100644
index 0000000000000..155eae92781af
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/rest/resources/LoggingResourceTest.java
@@ -0,0 +1,195 @@
+/*
+ * 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.rest.resources;
+
+import org.apache.kafka.connect.errors.NotFoundException;
+import org.apache.kafka.connect.runtime.rest.errors.BadRequestException;
+import org.apache.log4j.Hierarchy;
+import org.apache.log4j.Level;
+import org.apache.log4j.Logger;
+import org.junit.Test;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.Enumeration;
+import java.util.List;
+import java.util.Map;
+import java.util.Vector;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+@SuppressWarnings("unchecked")
+public class LoggingResourceTest {
+
+    @Test
+    public void getLoggersIgnoresNullLevelsTest() {
+        LoggingResource loggingResource = mock(LoggingResource.class);
+        Logger root = new Logger("root") {
+        };
+        Logger a = new Logger("a") {
+        };
+        a.setLevel(null);
+        Logger b = new Logger("b") {
+        };
+        b.setLevel(Level.INFO);
+        when(loggingResource.currentLoggers()).thenReturn(loggers(a, b));
+        when(loggingResource.rootLogger()).thenReturn(root);
+        when(loggingResource.listLoggers()).thenCallRealMethod();
+        Map> loggers = (Map>) loggingResource.listLoggers().getEntity();
+        assertEquals(1, loggers.size());
+        assertEquals("INFO", loggers.get("b").get("level"));
+    }
+
+    @Test
+    public void getLoggerFallsbackToEffectiveLogLevelTest() {
+        LoggingResource loggingResource = mock(LoggingResource.class);
+        Logger root = new Logger("root") {
+        };
+        root.setLevel(Level.ERROR);
+        Hierarchy hierarchy = new Hierarchy(root);
+        Logger a = hierarchy.getLogger("a");
+        a.setLevel(null);
+        Logger b = hierarchy.getLogger("b");
+        b.setLevel(Level.INFO);
+        when(loggingResource.currentLoggers()).thenReturn(loggers(a, b));
+        when(loggingResource.rootLogger()).thenReturn(root);
+        when(loggingResource.getLogger(any())).thenCallRealMethod();
+        Map level = (Map) loggingResource.getLogger("a").getEntity();
+        assertEquals(1, level.size());
+        assertEquals("ERROR", level.get("level"));
+    }
+
+    @Test(expected = NotFoundException.class)
+    public void getUnknownLoggerTest() {
+        LoggingResource loggingResource = mock(LoggingResource.class);
+        Logger root = new Logger("root") {
+        };
+        root.setLevel(Level.ERROR);
+        Hierarchy hierarchy = new Hierarchy(root);
+        Logger a = hierarchy.getLogger("a");
+        a.setLevel(null);
+        Logger b = hierarchy.getLogger("b");
+        b.setLevel(Level.INFO);
+        when(loggingResource.currentLoggers()).thenReturn(loggers(a, b));
+        when(loggingResource.rootLogger()).thenReturn(root);
+        when(loggingResource.getLogger(any())).thenCallRealMethod();
+        loggingResource.getLogger("c");
+    }
+
+    @Test
+    public void setLevelTest() {
+        LoggingResource loggingResource = mock(LoggingResource.class);
+        Logger root = new Logger("root") {
+        };
+        root.setLevel(Level.ERROR);
+        Hierarchy hierarchy = new Hierarchy(root);
+        Logger p = hierarchy.getLogger("a.b.c.p");
+        Logger x = hierarchy.getLogger("a.b.c.p.X");
+        Logger y = hierarchy.getLogger("a.b.c.p.Y");
+        Logger z = hierarchy.getLogger("a.b.c.p.Z");
+        Logger w = hierarchy.getLogger("a.b.c.s.W");
+        x.setLevel(Level.INFO);
+        y.setLevel(Level.INFO);
+        z.setLevel(Level.INFO);
+        w.setLevel(Level.INFO);
+        when(loggingResource.currentLoggers()).thenReturn(loggers(x, y, z, w));
+        when(loggingResource.lookupLogger("a.b.c.p")).thenReturn(p);
+        when(loggingResource.rootLogger()).thenReturn(root);
+        when(loggingResource.setLevel(any(), any())).thenCallRealMethod();
+        List modified = (List) loggingResource.setLevel("a.b.c.p", Collections.singletonMap("level", "DEBUG")).getEntity();
+        assertEquals(4, modified.size());
+        assertEquals(Arrays.asList("a.b.c.p", "a.b.c.p.X", "a.b.c.p.Y", "a.b.c.p.Z"), modified);
+        assertEquals(p.getLevel(), Level.DEBUG);
+        assertEquals(x.getLevel(), Level.DEBUG);
+        assertEquals(y.getLevel(), Level.DEBUG);
+        assertEquals(z.getLevel(), Level.DEBUG);
+    }
+
+    @Test
+    public void setRootLevelTest() {
+        LoggingResource loggingResource = mock(LoggingResource.class);
+        Logger root = new Logger("root") {
+        };
+        root.setLevel(Level.ERROR);
+        Hierarchy hierarchy = new Hierarchy(root);
+        Logger p = hierarchy.getLogger("a.b.c.p");
+        Logger x = hierarchy.getLogger("a.b.c.p.X");
+        Logger y = hierarchy.getLogger("a.b.c.p.Y");
+        Logger z = hierarchy.getLogger("a.b.c.p.Z");
+        Logger w = hierarchy.getLogger("a.b.c.s.W");
+        x.setLevel(Level.INFO);
+        y.setLevel(Level.INFO);
+        z.setLevel(Level.INFO);
+        w.setLevel(Level.INFO);
+        when(loggingResource.currentLoggers()).thenReturn(loggers(x, y, z, w));
+        when(loggingResource.lookupLogger("a.b.c.p")).thenReturn(p);
+        when(loggingResource.rootLogger()).thenReturn(root);
+        when(loggingResource.setLevel(any(), any())).thenCallRealMethod();
+        List modified = (List) loggingResource.setLevel("root", Collections.singletonMap("level", "DEBUG")).getEntity();
+        assertEquals(5, modified.size());
+        assertEquals(Arrays.asList("a.b.c.p.X", "a.b.c.p.Y", "a.b.c.p.Z", "a.b.c.s.W", "root"), modified);
+        assertNull(p.getLevel());
+        assertEquals(root.getLevel(), Level.DEBUG);
+        assertEquals(w.getLevel(), Level.DEBUG);
+        assertEquals(x.getLevel(), Level.DEBUG);
+        assertEquals(y.getLevel(), Level.DEBUG);
+        assertEquals(z.getLevel(), Level.DEBUG);
+    }
+
+    @Test(expected = BadRequestException.class)
+    public void setLevelWithEmptyArgTest() {
+        LoggingResource loggingResource = mock(LoggingResource.class);
+        Logger root = new Logger("root") {
+        };
+        root.setLevel(Level.ERROR);
+        Hierarchy hierarchy = new Hierarchy(root);
+        Logger a = hierarchy.getLogger("a");
+        a.setLevel(null);
+        Logger b = hierarchy.getLogger("b");
+        b.setLevel(Level.INFO);
+        when(loggingResource.currentLoggers()).thenReturn(loggers(a, b));
+        when(loggingResource.rootLogger()).thenReturn(root);
+        when(loggingResource.setLevel(any(), any())).thenCallRealMethod();
+        loggingResource.setLevel("@root", Collections.emptyMap());
+    }
+
+    @Test(expected = NotFoundException.class)
+    public void setLevelWithInvalidArgTest() {
+        LoggingResource loggingResource = mock(LoggingResource.class);
+        Logger root = new Logger("root") {
+        };
+        root.setLevel(Level.ERROR);
+        Hierarchy hierarchy = new Hierarchy(root);
+        Logger a = hierarchy.getLogger("a");
+        a.setLevel(null);
+        Logger b = hierarchy.getLogger("b");
+        b.setLevel(Level.INFO);
+        when(loggingResource.currentLoggers()).thenReturn(loggers(a, b));
+        when(loggingResource.rootLogger()).thenReturn(root);
+        when(loggingResource.setLevel(any(), any())).thenCallRealMethod();
+        loggingResource.setLevel("@root", Collections.singletonMap("level", "HIGH"));
+    }
+
+    private Enumeration loggers(Logger... loggers) {
+        return new Vector<>(Arrays.asList(loggers)).elements();
+    }
+
+}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/rest/resources/RootResourceTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/rest/resources/RootResourceTest.java
index be80e28f42da9..4e928a370372f 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/rest/resources/RootResourceTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/rest/resources/RootResourceTest.java
@@ -19,7 +19,6 @@
 import org.apache.kafka.clients.admin.MockAdminClient;
 import org.apache.kafka.common.utils.AppInfoParser;
 import org.apache.kafka.connect.runtime.Herder;
-import org.apache.kafka.connect.runtime.HerderProvider;
 import org.apache.kafka.connect.runtime.rest.entities.ServerInfo;
 import org.easymock.EasyMock;
 import org.easymock.EasyMockRunner;
@@ -40,7 +39,7 @@ public class RootResourceTest extends EasyMockSupport {
 
     @Before
     public void setUp() {
-        rootResource = new RootResource(new HerderProvider(herder));
+        rootResource = new RootResource(herder);
     }
 
     @Test
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 63595d653698c..8959a6c6e0b1f 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
@@ -59,7 +59,7 @@ public void testGetOrDefault() {
     }
 
     @Test
-    public void testCreateSslContextFactory() {
+    public void testCreateServerSideSslContextFactory() {
         Map configMap = new HashMap<>(DEFAULT_CONFIG);
         configMap.put("ssl.keystore.location", "/path/to/keystore");
         configMap.put("ssl.keystore.password", "123456");
@@ -79,7 +79,7 @@ public void testCreateSslContextFactory() {
         configMap.put("ssl.trustmanager.algorithm", "PKIX");
 
         DistributedConfig config = new DistributedConfig(configMap);
-        SslContextFactory ssl = SSLUtils.createSslContextFactory(config);
+        SslContextFactory ssl = SSLUtils.createServerSideSslContextFactory(config);
 
         Assert.assertEquals("file:///path/to/keystore", ssl.getKeyStorePath());
         Assert.assertEquals("file:///path/to/truststore", ssl.getTrustStorePath());
@@ -87,6 +87,7 @@ public void testCreateSslContextFactory() {
         Assert.assertArrayEquals(new String[] {"SSL_RSA_WITH_RC4_128_SHA", "SSL_RSA_WITH_RC4_128_MD5"}, ssl.getIncludeCipherSuites());
         Assert.assertEquals("SHA1PRNG", ssl.getSecureRandomAlgorithm());
         Assert.assertTrue(ssl.getNeedClientAuth());
+        Assert.assertFalse(ssl.getWantClientAuth());
         Assert.assertEquals("JKS", ssl.getKeyStoreType());
         Assert.assertEquals("JKS", ssl.getTrustStoreType());
         Assert.assertEquals("TLS", ssl.getProtocol());
@@ -96,7 +97,75 @@ public void testCreateSslContextFactory() {
     }
 
     @Test
-    public void testCreateSslContextFactoryDefaultValues() {
+    public void testCreateClientSideSslContextFactory() {
+        Map configMap = new HashMap<>(DEFAULT_CONFIG);
+        configMap.put("ssl.keystore.location", "/path/to/keystore");
+        configMap.put("ssl.keystore.password", "123456");
+        configMap.put("ssl.key.password", "123456");
+        configMap.put("ssl.truststore.location", "/path/to/truststore");
+        configMap.put("ssl.truststore.password", "123456");
+        configMap.put("ssl.provider", "SunJSSE");
+        configMap.put("ssl.cipher.suites", "SSL_RSA_WITH_RC4_128_SHA,SSL_RSA_WITH_RC4_128_MD5");
+        configMap.put("ssl.secure.random.implementation", "SHA1PRNG");
+        configMap.put("ssl.client.auth", "required");
+        configMap.put("ssl.endpoint.identification.algorithm", "HTTPS");
+        configMap.put("ssl.keystore.type", "JKS");
+        configMap.put("ssl.protocol", "TLS");
+        configMap.put("ssl.truststore.type", "JKS");
+        configMap.put("ssl.enabled.protocols", "TLSv1.2,TLSv1.1,TLSv1");
+        configMap.put("ssl.keymanager.algorithm", "SunX509");
+        configMap.put("ssl.trustmanager.algorithm", "PKIX");
+
+        DistributedConfig config = new DistributedConfig(configMap);
+        SslContextFactory 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());
+        Assert.assertArrayEquals(new String[] {"TLSv1.2", "TLSv1.1", "TLSv1"}, ssl.getIncludeProtocols());
+        Assert.assertEquals("SunX509", ssl.getKeyManagerFactoryAlgorithm());
+        Assert.assertEquals("PKIX", ssl.getTrustManagerFactoryAlgorithm());
+    }
+
+    @Test
+    public void testCreateServerSideSslContextFactoryDefaultValues() {
+        Map configMap = new HashMap<>(DEFAULT_CONFIG);
+        configMap.put(StandaloneConfig.OFFSET_STORAGE_FILE_FILENAME_CONFIG, "/tmp/offset/file");
+        configMap.put(WorkerConfig.KEY_CONVERTER_CLASS_CONFIG, "org.apache.kafka.connect.json.JsonConverter");
+        configMap.put(WorkerConfig.VALUE_CONVERTER_CLASS_CONFIG, "org.apache.kafka.connect.json.JsonConverter");
+        configMap.put(WorkerConfig.INTERNAL_KEY_CONVERTER_CLASS_CONFIG, "org.apache.kafka.connect.json.JsonConverter");
+        configMap.put(WorkerConfig.INTERNAL_VALUE_CONVERTER_CLASS_CONFIG, "org.apache.kafka.connect.json.JsonConverter");
+        configMap.put("ssl.keystore.location", "/path/to/keystore");
+        configMap.put("ssl.keystore.password", "123456");
+        configMap.put("ssl.key.password", "123456");
+        configMap.put("ssl.truststore.location", "/path/to/truststore");
+        configMap.put("ssl.truststore.password", "123456");
+        configMap.put("ssl.provider", "SunJSSE");
+        configMap.put("ssl.cipher.suites", "SSL_RSA_WITH_RC4_128_SHA,SSL_RSA_WITH_RC4_128_MD5");
+        configMap.put("ssl.secure.random.implementation", "SHA1PRNG");
+
+        DistributedConfig config = new DistributedConfig(configMap);
+        SslContextFactory ssl = SSLUtils.createServerSideSslContextFactory(config);
+
+        Assert.assertEquals(SslConfigs.DEFAULT_SSL_KEYSTORE_TYPE, ssl.getKeyStoreType());
+        Assert.assertEquals(SslConfigs.DEFAULT_SSL_TRUSTSTORE_TYPE, ssl.getTrustStoreType());
+        Assert.assertEquals(SslConfigs.DEFAULT_SSL_PROTOCOL, ssl.getProtocol());
+        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());
+    }
+
+    @Test
+    public void testCreateClientSideSslContextFactoryDefaultValues() {
         Map configMap = new HashMap<>(DEFAULT_CONFIG);
         configMap.put(StandaloneConfig.OFFSET_STORAGE_FILE_FILENAME_CONFIG, "/tmp/offset/file");
         configMap.put(WorkerConfig.KEY_CONVERTER_CLASS_CONFIG, "org.apache.kafka.connect.json.JsonConverter");
@@ -113,7 +182,7 @@ public void testCreateSslContextFactoryDefaultValues() {
         configMap.put("ssl.secure.random.implementation", "SHA1PRNG");
 
         DistributedConfig config = new DistributedConfig(configMap);
-        SslContextFactory ssl = SSLUtils.createSslContextFactory(config);
+        SslContextFactory ssl = SSLUtils.createClientSideSslContextFactory(config);
 
         Assert.assertEquals(SslConfigs.DEFAULT_SSL_KEYSTORE_TYPE, ssl.getKeyStoreType());
         Assert.assertEquals(SslConfigs.DEFAULT_SSL_TRUSTSTORE_TYPE, ssl.getTrustStoreType());
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/standalone/StandaloneConfigTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/standalone/StandaloneConfigTest.java
new file mode 100644
index 0000000000000..e2e886f7925f7
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/standalone/StandaloneConfigTest.java
@@ -0,0 +1,88 @@
+/*
+ * 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.standalone;
+
+import org.apache.kafka.common.config.ConfigDef;
+import org.apache.kafka.common.config.SslConfigs;
+import org.apache.kafka.common.config.types.Password;
+import org.apache.kafka.connect.runtime.WorkerConfig;
+import org.junit.Test;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+import static org.junit.Assert.assertEquals;
+
+public class StandaloneConfigTest {
+
+    private static final String HTTPS_LISTENER_PREFIX = "listeners.https.";
+
+    private Map sslProps() {
+        return new HashMap() {
+            {
+                put(SslConfigs.SSL_KEY_PASSWORD_CONFIG, new Password("ssl_key_password"));
+                put(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG, "ssl_keystore");
+                put(SslConfigs.SSL_KEYSTORE_PASSWORD_CONFIG, new Password("ssl_keystore_password"));
+                put(SslConfigs.SSL_TRUSTSTORE_LOCATION_CONFIG, "ssl_truststore");
+                put(SslConfigs.SSL_TRUSTSTORE_PASSWORD_CONFIG, new Password("ssl_truststore_password"));
+            }
+        };
+    }
+
+    private Map baseWorkerProps() {
+        return new HashMap() {
+            {
+                put(WorkerConfig.KEY_CONVERTER_CLASS_CONFIG, "org.apache.kafka.connect.json.JsonConverter");
+                put(WorkerConfig.VALUE_CONVERTER_CLASS_CONFIG, "org.apache.kafka.connect.json.JsonConverter");
+                put(StandaloneConfig.OFFSET_STORAGE_FILE_FILENAME_CONFIG, "/tmp/foo");
+            }
+        };
+    }
+
+    private static Map withStringValues(Map inputs, String prefix) {
+        return ConfigDef.convertToStringMapWithPasswordValues(inputs).entrySet().stream()
+            .collect(Collectors.toMap(
+                entry -> prefix + entry.getKey(),
+                Map.Entry::getValue
+            ));
+    }
+
+    @Test
+    public void testRestServerPrefixedSslConfigs() {
+        Map workerProps = baseWorkerProps();
+        Map expectedSslProps = sslProps();
+        workerProps.putAll(withStringValues(expectedSslProps, HTTPS_LISTENER_PREFIX));
+
+        StandaloneConfig config = new StandaloneConfig(workerProps);
+        assertEquals(expectedSslProps, config.valuesWithPrefixAllOrNothing(HTTPS_LISTENER_PREFIX));
+    }
+
+    @Test
+    public void testRestServerNonPrefixedSslConfigs() {
+        Map props = baseWorkerProps();
+        Map expectedSslProps = sslProps();
+        props.putAll(withStringValues(expectedSslProps, ""));
+
+        StandaloneConfig config = new StandaloneConfig(props);
+        Map actualProps = config.valuesWithPrefixAllOrNothing(HTTPS_LISTENER_PREFIX)
+            .entrySet().stream()
+            .filter(entry -> expectedSslProps.containsKey(entry.getKey()))
+            .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
+        assertEquals(expectedSslProps, actualProps);
+    }
+}
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 8aa1c706f7af1..feb0f995b7f45 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
@@ -22,6 +22,8 @@
 import org.apache.kafka.connect.connector.Connector;
 import org.apache.kafka.connect.connector.ConnectorContext;
 import org.apache.kafka.connect.connector.Task;
+import org.apache.kafka.connect.connector.policy.ConnectorClientConfigOverridePolicy;
+import org.apache.kafka.connect.connector.policy.NoneConnectorClientConfigOverridePolicy;
 import org.apache.kafka.connect.errors.AlreadyExistsException;
 import org.apache.kafka.connect.errors.ConnectException;
 import org.apache.kafka.connect.errors.NotFoundException;
@@ -111,11 +113,15 @@ private enum SourceSink {
     @Mock protected Callback> createCallback;
     @Mock protected StatusBackingStore statusBackingStore;
 
+    private final ConnectorClientConfigOverridePolicy
+        noneConnectorClientConfigOverridePolicy = new NoneConnectorClientConfigOverridePolicy();
+
+
     @Before
     public void setup() {
         worker = PowerMock.createMock(Worker.class);
         herder = PowerMock.createPartialMock(StandaloneHerder.class, new String[]{"connectorTypeForClass"},
-            worker, WORKER_ID, KAFKA_CLUSTER_ID, statusBackingStore, new MemoryConfigBackingStore(transformer));
+            worker, WORKER_ID, KAFKA_CLUSTER_ID, statusBackingStore, new MemoryConfigBackingStore(transformer), noneConnectorClientConfigOverridePolicy);
         plugins = PowerMock.createMock(Plugins.class);
         pluginLoader = PowerMock.createMock(PluginClassLoader.class);
         delegatingLoader = PowerMock.createMock(DelegatingClassLoader.class);
@@ -260,6 +266,7 @@ public void testDestroyConnector() throws Exception {
 
         EasyMock.expect(statusBackingStore.getAll(CONNECTOR_NAME)).andReturn(Collections.emptyList());
         statusBackingStore.put(new ConnectorStatus(CONNECTOR_NAME, AbstractStatus.State.DESTROYED, WORKER_ID, 0));
+        statusBackingStore.put(new TaskStatus(new ConnectorTaskId(CONNECTOR_NAME, 0), TaskStatus.State.DESTROYED, WORKER_ID, 0));
 
         expectDestroy();
 
@@ -355,6 +362,7 @@ public void testRestartTask() throws Exception {
 
         ClusterConfigState configState = new ClusterConfigState(
                 -1,
+                null,
                 Collections.singletonMap(CONNECTOR_NAME, 1),
                 Collections.singletonMap(CONNECTOR_NAME, connectorConfig),
                 Collections.singletonMap(CONNECTOR_NAME, TargetState.STARTED),
@@ -389,6 +397,7 @@ public void testRestartTaskFailureOnStart() throws Exception {
 
         ClusterConfigState configState = new ClusterConfigState(
                 -1,
+                null,
                 Collections.singletonMap(CONNECTOR_NAME, 1),
                 Collections.singletonMap(CONNECTOR_NAME, connectorConfig),
                 Collections.singletonMap(CONNECTOR_NAME, TargetState.STARTED),
@@ -426,6 +435,8 @@ public void testCreateAndStop() throws Exception {
         // herder.stop() should stop any running connectors and tasks even if destroyConnector was not invoked
         expectStop();
 
+        statusBackingStore.put(new TaskStatus(new ConnectorTaskId(CONNECTOR_NAME, 0), AbstractStatus.State.DESTROYED, WORKER_ID, 0));
+
         statusBackingStore.stop();
         EasyMock.expectLastCall();
         worker.stop();
@@ -564,7 +575,8 @@ public void testPutTaskConfigs() {
 
         herder.putTaskConfigs(CONNECTOR_NAME,
                 Arrays.asList(singletonMap("config", "value")),
-                cb);
+                cb,
+                null);
 
         PowerMock.verifyAll();
     }
@@ -608,7 +620,7 @@ public void testCorruptConfig() {
             capture.getValue().getMessage(),
             "Connector configuration is invalid and contains the following 1 error(s):\n" +
                 error + "\n" +
-                "You can also find the above list of errors at the endpoint `/{connectorType}/config/validate`"
+                "You can also find the above list of errors at the endpoint `/connector-plugins/{connectorType}/config/validate`"
         );
 
         PowerMock.verifyAll();
@@ -639,6 +651,7 @@ private void expectAdd(SourceSink sourceSink) {
 
         ClusterConfigState configState = new ClusterConfigState(
                 -1,
+                null,
                 Collections.singletonMap(CONNECTOR_NAME, 1),
                 Collections.singletonMap(CONNECTOR_NAME, connectorConfig(sourceSink)),
                 Collections.singletonMap(CONNECTOR_NAME, TargetState.STARTED),
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/storage/FileOffsetBackingStoreTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/storage/FileOffsetBackingStoreTest.java
index df955f8813d53..1fdb91a2dcb2d 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/storage/FileOffsetBackingStoreTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/storage/FileOffsetBackingStoreTest.java
@@ -71,12 +71,11 @@ public void teardown() {
     @Test
     public void testGetSet() throws Exception {
         Callback setCallback = expectSuccessfulSetCallback();
-        Callback> getCallback = expectSuccessfulGetCallback();
         PowerMock.replayAll();
 
         store.set(firstSet, setCallback).get();
 
-        Map values = store.get(Arrays.asList(buffer("key"), buffer("bad")), getCallback).get();
+        Map values = store.get(Arrays.asList(buffer("key"), buffer("bad"))).get();
         assertEquals(buffer("value"), values.get(buffer("key")));
         assertEquals(null, values.get(buffer("bad")));
 
@@ -86,7 +85,6 @@ public void testGetSet() throws Exception {
     @Test
     public void testSaveRestore() throws Exception {
         Callback setCallback = expectSuccessfulSetCallback();
-        Callback> getCallback = expectSuccessfulGetCallback();
         PowerMock.replayAll();
 
         store.set(firstSet, setCallback).get();
@@ -96,7 +94,7 @@ public void testSaveRestore() throws Exception {
         FileOffsetBackingStore restore = new FileOffsetBackingStore();
         restore.configure(config);
         restore.start();
-        Map values = restore.get(Arrays.asList(buffer("key")), getCallback).get();
+        Map values = restore.get(Arrays.asList(buffer("key"))).get();
         assertEquals(buffer("value"), values.get(buffer("key")));
 
         PowerMock.verifyAll();
@@ -113,12 +111,4 @@ private Callback expectSuccessfulSetCallback() {
         PowerMock.expectLastCall();
         return setCallback;
     }
-
-    @SuppressWarnings("unchecked")
-    private Callback> expectSuccessfulGetCallback() {
-        Callback> getCallback = PowerMock.createMock(Callback.class);
-        getCallback.onCompletion(EasyMock.isNull(Throwable.class), EasyMock.anyObject(Map.class));
-        PowerMock.expectLastCall();
-        return getCallback;
-    }
 }
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 68c447a1aebdd..fc1b814fec1ac 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
@@ -22,6 +22,7 @@
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.clients.producer.ProducerConfig;
 import org.apache.kafka.common.record.TimestampType;
+import org.apache.kafka.common.config.ConfigException;
 import org.apache.kafka.connect.data.Field;
 import org.apache.kafka.connect.data.Schema;
 import org.apache.kafka.connect.data.SchemaAndValue;
@@ -57,12 +58,14 @@
 import java.util.concurrent.TimeUnit;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 
 @RunWith(PowerMockRunner.class)
 @PrepareForTest(KafkaConfigBackingStore.class)
-@PowerMockIgnore("javax.management.*")
+@PowerMockIgnore({"javax.management.*", "javax.crypto.*"})
 @SuppressWarnings({"unchecked", "deprecation"})
 public class KafkaConfigBackingStoreTest {
     private static final String TOPIC = "connect-configs";
@@ -155,6 +158,7 @@ public void setUp() {
     public void testStartStop() throws Exception {
         expectConfigure();
         expectStart(Collections.emptyList(), Collections.emptyMap());
+        expectPartitionCount(1);
         expectStop();
 
         PowerMock.replayAll();
@@ -198,6 +202,7 @@ public void testPutConnectorConfig() throws Exception {
         configUpdateListener.onConnectorConfigRemove(CONNECTOR_IDS.get(1));
         EasyMock.expectLastCall();
 
+        expectPartitionCount(1);
         expectStop();
 
         PowerMock.replayAll();
@@ -266,6 +271,7 @@ public void testPutTaskConfigs() throws Exception {
         serializedConfigs.put(COMMIT_TASKS_CONFIG_KEYS.get(0), CONFIGS_SERIALIZED.get(2));
         expectReadToEnd(serializedConfigs);
 
+        expectPartitionCount(1);
         expectStop();
 
         PowerMock.replayAll();
@@ -302,6 +308,97 @@ public void testPutTaskConfigs() throws Exception {
         PowerMock.verifyAll();
     }
 
+    @Test
+    public void testPutTaskConfigsStartsOnlyReconfiguredTasks() throws Exception {
+        expectConfigure();
+        expectStart(Collections.emptyList(), Collections.emptyMap());
+
+        // Task configs should read to end, write to the log, read to end, write root, then read to end again
+        expectReadToEnd(new LinkedHashMap<>());
+        expectConvertWriteRead(
+                TASK_CONFIG_KEYS.get(0), KafkaConfigBackingStore.TASK_CONFIGURATION_V0, CONFIGS_SERIALIZED.get(0),
+                "properties", SAMPLE_CONFIGS.get(0));
+        expectConvertWriteRead(
+                TASK_CONFIG_KEYS.get(1), KafkaConfigBackingStore.TASK_CONFIGURATION_V0, CONFIGS_SERIALIZED.get(1),
+                "properties", SAMPLE_CONFIGS.get(1));
+        expectReadToEnd(new LinkedHashMap());
+        expectConvertWriteRead(
+                COMMIT_TASKS_CONFIG_KEYS.get(0), KafkaConfigBackingStore.CONNECTOR_TASKS_COMMIT_V0, CONFIGS_SERIALIZED.get(2),
+                "tasks", 2); // Starts with 0 tasks, after update has 2
+        // As soon as root is rewritten, we should see a callback notifying us that we reconfigured some tasks
+        configUpdateListener.onTaskConfigUpdate(Arrays.asList(TASK_IDS.get(0), TASK_IDS.get(1)));
+        EasyMock.expectLastCall();
+
+        // Records to be read by consumer as it reads to the end of the log
+        LinkedHashMap serializedConfigs = new LinkedHashMap<>();
+        serializedConfigs.put(TASK_CONFIG_KEYS.get(0), CONFIGS_SERIALIZED.get(0));
+        serializedConfigs.put(TASK_CONFIG_KEYS.get(1), CONFIGS_SERIALIZED.get(1));
+        serializedConfigs.put(COMMIT_TASKS_CONFIG_KEYS.get(0), CONFIGS_SERIALIZED.get(2));
+        expectReadToEnd(serializedConfigs);
+
+        // Task configs should read to end, write to the log, read to end, write root, then read to end again
+        expectReadToEnd(new LinkedHashMap<>());
+        expectConvertWriteRead(
+                TASK_CONFIG_KEYS.get(2), KafkaConfigBackingStore.TASK_CONFIGURATION_V0, CONFIGS_SERIALIZED.get(3),
+                "properties", SAMPLE_CONFIGS.get(2));
+        expectReadToEnd(new LinkedHashMap<>());
+        expectConvertWriteRead(
+                COMMIT_TASKS_CONFIG_KEYS.get(1), KafkaConfigBackingStore.CONNECTOR_TASKS_COMMIT_V0, CONFIGS_SERIALIZED.get(4),
+                "tasks", 1); // Starts with 2 tasks, after update has 3
+
+        // As soon as root is rewritten, we should see a callback notifying us that we reconfigured some tasks
+        configUpdateListener.onTaskConfigUpdate(Arrays.asList(TASK_IDS.get(2)));
+        EasyMock.expectLastCall();
+
+        // Records to be read by consumer as it reads to the end of the log
+        serializedConfigs = new LinkedHashMap<>();
+        serializedConfigs.put(TASK_CONFIG_KEYS.get(2), CONFIGS_SERIALIZED.get(3));
+        serializedConfigs.put(COMMIT_TASKS_CONFIG_KEYS.get(1), CONFIGS_SERIALIZED.get(4));
+        expectReadToEnd(serializedConfigs);
+
+        expectPartitionCount(1);
+        expectStop();
+
+        PowerMock.replayAll();
+
+        configStorage.setupAndCreateKafkaBasedLog(TOPIC, DEFAULT_DISTRIBUTED_CONFIG);
+        configStorage.start();
+
+        // Bootstrap as if we had already added the connector, but no tasks had been added yet
+        whiteboxAddConnector(CONNECTOR_IDS.get(0), SAMPLE_CONFIGS.get(0), Collections.emptyList());
+        whiteboxAddConnector(CONNECTOR_IDS.get(1), SAMPLE_CONFIGS.get(1), Collections.emptyList());
+
+        // Null before writing
+        ClusterConfigState configState = configStorage.snapshot();
+        assertEquals(-1, configState.offset());
+        assertNull(configState.taskConfig(TASK_IDS.get(0)));
+        assertNull(configState.taskConfig(TASK_IDS.get(1)));
+
+        // Writing task configs should block until all the writes have been performed and the root record update
+        // has completed
+        List> taskConfigs = Arrays.asList(SAMPLE_CONFIGS.get(0), SAMPLE_CONFIGS.get(1));
+        configStorage.putTaskConfigs("connector1", taskConfigs);
+        taskConfigs = Collections.singletonList(SAMPLE_CONFIGS.get(2));
+        configStorage.putTaskConfigs("connector2", taskConfigs);
+
+        // Validate root config by listing all connectors and tasks
+        configState = configStorage.snapshot();
+        assertEquals(5, configState.offset());
+        String connectorName1 = CONNECTOR_IDS.get(0);
+        String connectorName2 = CONNECTOR_IDS.get(1);
+        assertEquals(Arrays.asList(connectorName1, connectorName2), new ArrayList<>(configState.connectors()));
+        assertEquals(Arrays.asList(TASK_IDS.get(0), TASK_IDS.get(1)), configState.tasks(connectorName1));
+        assertEquals(Collections.singletonList(TASK_IDS.get(2)), configState.tasks(connectorName2));
+        assertEquals(SAMPLE_CONFIGS.get(0), configState.taskConfig(TASK_IDS.get(0)));
+        assertEquals(SAMPLE_CONFIGS.get(1), configState.taskConfig(TASK_IDS.get(1)));
+        assertEquals(SAMPLE_CONFIGS.get(2), configState.taskConfig(TASK_IDS.get(2)));
+        assertEquals(Collections.EMPTY_SET, configState.inconsistentConnectors());
+
+        configStorage.stop();
+
+        PowerMock.verifyAll();
+    }
+
     @Test
     public void testPutTaskConfigsZeroTasks() throws Exception {
         expectConfigure();
@@ -321,6 +418,7 @@ public void testPutTaskConfigsZeroTasks() throws Exception {
         serializedConfigs.put(COMMIT_TASKS_CONFIG_KEYS.get(0), CONFIGS_SERIALIZED.get(0));
         expectReadToEnd(serializedConfigs);
 
+        expectPartitionCount(1);
         expectStop();
 
         PowerMock.replayAll();
@@ -374,6 +472,7 @@ public void testRestoreTargetState() throws Exception {
 
         // Shouldn't see any callbacks since this is during startup
 
+        expectPartitionCount(1);
         expectStop();
 
         PowerMock.replayAll();
@@ -416,6 +515,7 @@ public void testBackgroundUpdateTargetState() throws Exception {
         configUpdateListener.onConnectorTargetStateChange(CONNECTOR_IDS.get(0));
         EasyMock.expectLastCall();
 
+        expectPartitionCount(1);
         expectStop();
 
         PowerMock.replayAll();
@@ -447,7 +547,7 @@ public void testBackgroundConnectorDeletion() throws Exception {
         LinkedHashMap deserialized = new LinkedHashMap<>();
         deserialized.put(CONFIGS_SERIALIZED.get(0), CONNECTOR_CONFIG_STRUCTS.get(0));
         deserialized.put(CONFIGS_SERIALIZED.get(1), TASK_CONFIG_STRUCTS.get(0));
-        deserialized.put(CONFIGS_SERIALIZED.get(2), TASK_CONFIG_STRUCTS.get(0));
+        deserialized.put(CONFIGS_SERIALIZED.get(2), TASK_CONFIG_STRUCTS.get(1));
         deserialized.put(CONFIGS_SERIALIZED.get(3), TASKS_COMMIT_STRUCT_TWO_TASK_CONNECTOR);
         logOffset = 5;
 
@@ -466,6 +566,7 @@ public void testBackgroundConnectorDeletion() throws Exception {
         configUpdateListener.onConnectorConfigRemove(CONNECTOR_IDS.get(0));
         EasyMock.expectLastCall();
 
+        expectPartitionCount(1);
         expectStop();
 
         PowerMock.replayAll();
@@ -476,8 +577,17 @@ public void testBackgroundConnectorDeletion() throws Exception {
         // Should see a single connector with initial state paused
         ClusterConfigState configState = configStorage.snapshot();
         assertEquals(TargetState.STARTED, configState.targetState(CONNECTOR_IDS.get(0)));
+        assertEquals(SAMPLE_CONFIGS.get(0), configState.connectorConfig(CONNECTOR_IDS.get(0)));
+        assertEquals(SAMPLE_CONFIGS.subList(0, 2), configState.allTaskConfigs(CONNECTOR_IDS.get(0)));
+        assertEquals(2, configState.taskCount(CONNECTOR_IDS.get(0)));
 
         configStorage.refresh(0, TimeUnit.SECONDS);
+        configState = configStorage.snapshot();
+        // Connector should now be removed from the snapshot
+        assertFalse(configState.contains(CONNECTOR_IDS.get(0)));
+        // Task configs for the deleted connector should also be removed from the snapshot
+        assertEquals(Collections.emptyList(), configState.allTaskConfigs(CONNECTOR_IDS.get(0)));
+        assertEquals(0, configState.taskCount(CONNECTOR_IDS.get(0)));
 
         configStorage.stop();
 
@@ -502,6 +612,7 @@ public void testRestoreTargetStateUnexpectedDeletion() throws Exception {
         logOffset = 5;
 
         expectStart(existingRecords, deserialized);
+        expectPartitionCount(1);
 
         // Shouldn't see any callbacks since this is during startup
 
@@ -549,6 +660,7 @@ public void testRestore() throws Exception {
         deserialized.put(CONFIGS_SERIALIZED.get(6), TASK_CONFIG_STRUCTS.get(1));
         logOffset = 7;
         expectStart(existingRecords, deserialized);
+        expectPartitionCount(1);
 
         // Shouldn't see any callbacks since this is during startup
 
@@ -603,6 +715,7 @@ public void testRestoreConnectorDeletion() throws Exception {
 
         logOffset = 6;
         expectStart(existingRecords, deserialized);
+        expectPartitionCount(1);
 
         // Shouldn't see any callbacks since this is during startup
 
@@ -650,6 +763,7 @@ public void testRestoreZeroTasks() throws Exception {
         deserialized.put(CONFIGS_SERIALIZED.get(7), TASKS_COMMIT_STRUCT_ZERO_TASK_CONNECTOR);
         logOffset = 8;
         expectStart(existingRecords, deserialized);
+        expectPartitionCount(1);
 
         // Shouldn't see any callbacks since this is during startup
 
@@ -697,6 +811,7 @@ public void testPutTaskConfigsDoesNotResolveAllInconsistencies() throws Exceptio
         deserialized.put(CONFIGS_SERIALIZED.get(5), TASK_CONFIG_STRUCTS.get(1));
         logOffset = 6;
         expectStart(existingRecords, deserialized);
+        expectPartitionCount(1);
 
         // Successful attempt to write new task config
         expectReadToEnd(new LinkedHashMap());
@@ -751,6 +866,22 @@ public void testPutTaskConfigsDoesNotResolveAllInconsistencies() throws Exceptio
         PowerMock.verifyAll();
     }
 
+    @Test
+    public void testExceptionOnStartWhenConfigTopicHasMultiplePartitions() throws Exception {
+        expectConfigure();
+        expectStart(Collections.emptyList(), Collections.emptyMap());
+
+        expectPartitionCount(2);
+
+        PowerMock.replayAll();
+
+        configStorage.setupAndCreateKafkaBasedLog(TOPIC, DEFAULT_DISTRIBUTED_CONFIG);
+        ConfigException e = assertThrows(ConfigException.class, () -> configStorage.start());
+        assertTrue(e.getMessage().contains("required to have a single partition"));
+
+        PowerMock.verifyAll();
+    }
+
     private void expectConfigure() throws Exception {
         PowerMock.expectPrivate(configStorage, "createKafkaBasedLog",
                 EasyMock.capture(capturedTopic), EasyMock.capture(capturedProducerProps),
@@ -759,6 +890,11 @@ private void expectConfigure() throws Exception {
                 .andReturn(storeLog);
     }
 
+    private void expectPartitionCount(int partitionCount) {
+        EasyMock.expect(storeLog.partitionCount())
+                .andReturn(partitionCount);
+    }
+
     // If non-empty, deserializations should be a LinkedHashMap
     private void expectStart(final List> preexistingRecords,
                              final Map deserializations) {
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 ff9f2c9d1ba2f..c04363b974605 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
@@ -59,7 +59,7 @@
 
 @RunWith(PowerMockRunner.class)
 @PrepareForTest(KafkaOffsetBackingStore.class)
-@PowerMockIgnore("javax.management.*")
+@PowerMockIgnore({"javax.management.*", "javax.crypto.*"})
 @SuppressWarnings({"unchecked", "deprecation"})
 public class KafkaOffsetBackingStoreTest {
     private static final String TOPIC = "connect-offsets";
@@ -217,17 +217,10 @@ public Object answer() throws Throwable {
         store.start();
 
         // Getting from empty store should return nulls
-        final AtomicBoolean getInvokedAndPassed = new AtomicBoolean(false);
-        store.get(Arrays.asList(TP0_KEY, TP1_KEY), new Callback>() {
-            @Override
-            public void onCompletion(Throwable error, Map result) {
-                // Since we didn't read them yet, these will be null
-                assertEquals(null, result.get(TP0_KEY));
-                assertEquals(null, result.get(TP1_KEY));
-                getInvokedAndPassed.set(true);
-            }
-        }).get(10000, TimeUnit.MILLISECONDS);
-        assertTrue(getInvokedAndPassed.get());
+        Map offsets = store.get(Arrays.asList(TP0_KEY, TP1_KEY)).get(10000, TimeUnit.MILLISECONDS);
+        // Since we didn't read them yet, these will be null
+        assertNull(offsets.get(TP0_KEY));
+        assertNull(offsets.get(TP1_KEY));
 
         // Set some offsets
         Map toSet = new HashMap<>();
@@ -250,28 +243,14 @@ public void onCompletion(Throwable error, Void result) {
         assertTrue(invoked.get());
 
         // Getting data should read to end of our published data and return it
-        final AtomicBoolean secondGetInvokedAndPassed = new AtomicBoolean(false);
-        store.get(Arrays.asList(TP0_KEY, TP1_KEY), new Callback>() {
-            @Override
-            public void onCompletion(Throwable error, Map result) {
-                assertEquals(TP0_VALUE, result.get(TP0_KEY));
-                assertEquals(TP1_VALUE, result.get(TP1_KEY));
-                secondGetInvokedAndPassed.set(true);
-            }
-        }).get(10000, TimeUnit.MILLISECONDS);
-        assertTrue(secondGetInvokedAndPassed.get());
+        offsets = store.get(Arrays.asList(TP0_KEY, TP1_KEY)).get(10000, TimeUnit.MILLISECONDS);
+        assertEquals(TP0_VALUE, offsets.get(TP0_KEY));
+        assertEquals(TP1_VALUE, offsets.get(TP1_KEY));
 
         // Getting data should read to end of our published data and return it
-        final AtomicBoolean thirdGetInvokedAndPassed = new AtomicBoolean(false);
-        store.get(Arrays.asList(TP0_KEY, TP1_KEY), new Callback>() {
-            @Override
-            public void onCompletion(Throwable error, Map result) {
-                assertEquals(TP0_VALUE_NEW, result.get(TP0_KEY));
-                assertEquals(TP1_VALUE_NEW, result.get(TP1_KEY));
-                thirdGetInvokedAndPassed.set(true);
-            }
-        }).get(10000, TimeUnit.MILLISECONDS);
-        assertTrue(thirdGetInvokedAndPassed.get());
+        offsets = store.get(Arrays.asList(TP0_KEY, TP1_KEY)).get(10000, TimeUnit.MILLISECONDS);
+        assertEquals(TP0_VALUE_NEW, offsets.get(TP0_KEY));
+        assertEquals(TP1_VALUE_NEW, offsets.get(TP1_KEY));
 
         store.stop();
 
@@ -329,16 +308,9 @@ public void onCompletion(Throwable error, Void result) {
         assertTrue(invoked.get());
 
         // Getting data should read to end of our published data and return it
-        final AtomicBoolean secondGetInvokedAndPassed = new AtomicBoolean(false);
-        store.get(Arrays.asList(null, TP1_KEY), new Callback>() {
-            @Override
-            public void onCompletion(Throwable error, Map result) {
-                assertEquals(TP0_VALUE, result.get(null));
-                assertNull(result.get(TP1_KEY));
-                secondGetInvokedAndPassed.set(true);
-            }
-        }).get(10000, TimeUnit.MILLISECONDS);
-        assertTrue(secondGetInvokedAndPassed.get());
+        Map offsets = store.get(Arrays.asList(null, TP1_KEY)).get(10000, TimeUnit.MILLISECONDS);
+        assertEquals(TP0_VALUE, offsets.get(null));
+        assertNull(offsets.get(TP1_KEY));
 
         store.stop();
 
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/util/ConvertingFutureCallbackTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/util/ConvertingFutureCallbackTest.java
new file mode 100644
index 0000000000000..9535003c863c3
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/util/ConvertingFutureCallbackTest.java
@@ -0,0 +1,242 @@
+/*
+ * 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.util;
+
+import org.junit.Before;
+import org.junit.Test;
+
+import java.util.concurrent.CancellationException;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
+
+public class ConvertingFutureCallbackTest {
+
+    private ExecutorService executor;
+
+    @Before
+    public void setup() {
+        executor = Executors.newSingleThreadExecutor();
+    }
+  
+    @Test
+    public void shouldConvertBeforeGetOnSuccessfulCompletion() throws Exception {
+        final Object expectedConversion = new Object();
+        TestConvertingFutureCallback testCallback = new TestConvertingFutureCallback();
+        testCallback.onCompletion(null, expectedConversion);
+        assertEquals(1, testCallback.numberOfConversions());
+        assertEquals(expectedConversion, testCallback.get());
+    }
+
+    @Test
+    public void shouldConvertOnlyOnceBeforeGetOnSuccessfulCompletion() throws Exception {
+        final Object expectedConversion = new Object();
+        TestConvertingFutureCallback testCallback = new TestConvertingFutureCallback();
+        testCallback.onCompletion(null, expectedConversion);
+        testCallback.onCompletion(null, 69);
+        testCallback.cancel(true);
+        testCallback.onCompletion(new RuntimeException(), null);
+        assertEquals(1, testCallback.numberOfConversions());
+        assertEquals(expectedConversion, testCallback.get());
+    }
+
+    @Test
+    public void shouldNotConvertBeforeGetOnFailedCompletion() throws Exception {
+        final Throwable expectedError = new Throwable();
+        TestConvertingFutureCallback testCallback = new TestConvertingFutureCallback();
+        testCallback.onCompletion(expectedError, null);
+        assertEquals(0, testCallback.numberOfConversions());
+        try {
+            testCallback.get();
+            fail("Expected ExecutionException");
+        } catch (ExecutionException e) {
+            assertEquals(expectedError, e.getCause());
+        }
+    }
+
+    @Test
+    public void shouldRecordOnlyFirstErrorBeforeGetOnFailedCompletion() throws Exception {
+        final Throwable expectedError = new Throwable();
+        TestConvertingFutureCallback testCallback = new TestConvertingFutureCallback();
+        testCallback.onCompletion(expectedError, null);
+        testCallback.onCompletion(new RuntimeException(), null);
+        testCallback.cancel(true);
+        testCallback.onCompletion(null, "420");
+        assertEquals(0, testCallback.numberOfConversions());
+        try {
+            testCallback.get();
+            fail("Expected ExecutionException");
+        } catch (ExecutionException e) {
+            assertEquals(expectedError, e.getCause());
+        }
+    }
+  
+    @Test(expected = CancellationException.class)
+    public void shouldCancelBeforeGetIfMayCancelWhileRunning() throws Exception {
+        TestConvertingFutureCallback testCallback = new TestConvertingFutureCallback();
+        assertTrue(testCallback.cancel(true));
+        testCallback.get();
+    }
+
+    @Test
+    public void shouldBlockUntilSuccessfulCompletion() throws Exception {
+        AtomicReference testThreadException = new AtomicReference<>();
+        TestConvertingFutureCallback testCallback = new TestConvertingFutureCallback();
+        final Object expectedConversion = new Object();
+        executor.submit(() -> {
+            try {
+                testCallback.waitForGet();
+                testCallback.onCompletion(null, expectedConversion);
+            } catch (Exception e) {
+                testThreadException.compareAndSet(null, e);
+            }
+        });
+        assertFalse(testCallback.isDone());
+        assertEquals(expectedConversion, testCallback.get());
+        assertEquals(1, testCallback.numberOfConversions());
+        assertTrue(testCallback.isDone());
+        if (testThreadException.get() != null) {
+            throw testThreadException.get();
+        }
+    }
+
+    @Test
+    public void shouldBlockUntilFailedCompletion() throws Exception {
+        AtomicReference testThreadException = new AtomicReference<>();
+        TestConvertingFutureCallback testCallback = new TestConvertingFutureCallback();
+        final Throwable expectedError = new Throwable();
+        executor.submit(() -> {
+            try {
+                testCallback.waitForGet();
+                testCallback.onCompletion(expectedError, null);
+            } catch (Exception e) {
+                testThreadException.compareAndSet(null, e);
+            }
+        });
+        assertFalse(testCallback.isDone());
+        try {
+            testCallback.get();
+            fail("Expected ExecutionException");
+        } catch (ExecutionException e) {
+            assertEquals(expectedError, e.getCause());
+        }
+        assertEquals(0, testCallback.numberOfConversions());
+        assertTrue(testCallback.isDone());
+        if (testThreadException.get() != null) {
+            throw testThreadException.get();
+        }
+    }
+
+    @Test(expected = CancellationException.class)
+    public void shouldBlockUntilCancellation() throws Exception {
+        AtomicReference testThreadException = new AtomicReference<>();
+        TestConvertingFutureCallback testCallback = new TestConvertingFutureCallback();
+        executor.submit(() -> {
+            try {
+                testCallback.waitForGet();
+                testCallback.cancel(true);
+            } catch (Exception e) {
+                testThreadException.compareAndSet(null, e);
+            }
+        });
+        assertFalse(testCallback.isDone());
+        testCallback.get();
+        if (testThreadException.get() != null) {
+            throw testThreadException.get();
+        }
+    }
+
+    @Test
+    public void shouldNotCancelIfMayNotCancelWhileRunning() throws Exception {
+        AtomicReference testThreadException = new AtomicReference<>();
+        TestConvertingFutureCallback testCallback = new TestConvertingFutureCallback();
+        final Object expectedConversion = new Object();
+        executor.submit(() -> {
+            try {
+                testCallback.waitForCancel();
+                testCallback.onCompletion(null, expectedConversion);
+            } catch (Exception e) {
+                testThreadException.compareAndSet(null, e);
+            }
+        });
+        assertFalse(testCallback.isCancelled());
+        assertFalse(testCallback.isDone());
+        testCallback.cancel(false);
+        assertFalse(testCallback.isCancelled());
+        assertTrue(testCallback.isDone());
+        assertEquals(expectedConversion, testCallback.get());
+        assertEquals(1, testCallback.numberOfConversions());
+        if (testThreadException.get() != null) {
+            throw testThreadException.get();
+        }
+    }
+  
+    protected static class TestConvertingFutureCallback extends ConvertingFutureCallback {
+        private AtomicInteger numberOfConversions = new AtomicInteger();
+        private CountDownLatch getInvoked = new CountDownLatch(1);
+        private CountDownLatch cancelInvoked = new CountDownLatch(1);
+    
+        public int numberOfConversions() {
+            return numberOfConversions.get();
+        }
+
+        public void waitForGet() throws InterruptedException {
+            getInvoked.await();
+        }
+
+        public void waitForCancel() throws InterruptedException {
+            cancelInvoked.await();
+        }
+    
+        @Override
+        public Object convert(Object result) {
+            numberOfConversions.incrementAndGet();
+            return result;
+        }
+
+        @Override
+        public Object get() throws InterruptedException, ExecutionException {
+            getInvoked.countDown();
+            return super.get();
+        }
+
+        @Override
+        public Object get(
+            long duration,
+            TimeUnit unit
+        ) throws InterruptedException, ExecutionException, TimeoutException {
+            getInvoked.countDown();
+            return super.get(duration, unit);
+        }
+
+        @Override
+        public boolean cancel(boolean mayInterruptIfRunning) {
+            cancelInvoked.countDown();
+            return super.cancel(mayInterruptIfRunning);
+        }
+    }
+}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/util/KafkaBasedLogTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/util/KafkaBasedLogTest.java
index 1af6e343521a6..080f9434d13f6 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/util/KafkaBasedLogTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/util/KafkaBasedLogTest.java
@@ -30,6 +30,7 @@
 import org.apache.kafka.common.PartitionInfo;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.errors.LeaderNotAvailableException;
+import org.apache.kafka.common.errors.TimeoutException;
 import org.apache.kafka.common.errors.WakeupException;
 import org.apache.kafka.common.protocol.Errors;
 import org.apache.kafka.common.record.TimestampType;
@@ -370,7 +371,7 @@ public void run() {
     }
 
     @Test
-    public void testConsumerError() throws Exception {
+    public void testPollConsumerError() throws Exception {
         expectStart();
         expectStop();
 
@@ -388,7 +389,7 @@ public void run() {
                 consumer.schedulePollTask(new Runnable() {
                     @Override
                     public void run() {
-                        consumer.setException(Errors.COORDINATOR_NOT_AVAILABLE.exception());
+                        consumer.setPollException(Errors.COORDINATOR_NOT_AVAILABLE.exception());
                     }
                 });
 
@@ -423,6 +424,77 @@ public void run() {
         PowerMock.verifyAll();
     }
 
+    @Test
+    public void testGetOffsetsConsumerErrorOnReadToEnd() throws Exception {
+        expectStart();
+
+        // Producer flushes when read to log end is called
+        producer.flush();
+        PowerMock.expectLastCall();
+
+        expectStop();
+
+        PowerMock.replayAll();
+        final CountDownLatch finishedLatch = new CountDownLatch(1);
+        Map endOffsets = new HashMap<>();
+        endOffsets.put(TP0, 0L);
+        endOffsets.put(TP1, 0L);
+        consumer.updateEndOffsets(endOffsets);
+        store.start();
+        final AtomicBoolean getInvoked = new AtomicBoolean(false);
+        final FutureCallback readEndFutureCallback = new FutureCallback<>(new Callback() {
+            @Override
+            public void onCompletion(Throwable error, Void result) {
+                getInvoked.set(true);
+            }
+        });
+        consumer.schedulePollTask(new Runnable() {
+            @Override
+            public void run() {
+                // Once we're synchronized in a poll, start the read to end and schedule the exact set of poll events
+                // that should follow. This readToEnd call will immediately wakeup this consumer.poll() call without
+                // returning any data.
+                Map newEndOffsets = new HashMap<>();
+                newEndOffsets.put(TP0, 1L);
+                newEndOffsets.put(TP1, 1L);
+                consumer.updateEndOffsets(newEndOffsets);
+                // Set exception to occur when getting offsets to read log to end.  It'll be caught in the work thread,
+                // which will retry and eventually get the correct offsets and read log to end.
+                consumer.setOffsetsException(new TimeoutException("Failed to get offsets by times"));
+                store.readToEnd(readEndFutureCallback);
+
+                // Should keep polling until it reaches current log end offset for all partitions
+                consumer.scheduleNopPollTask();
+                consumer.scheduleNopPollTask();
+                consumer.schedulePollTask(new Runnable() {
+                    @Override
+                    public void run() {
+                        consumer.addRecord(new ConsumerRecord<>(TOPIC, 0, 0, 0L, TimestampType.CREATE_TIME, 0L, 0, 0, TP0_KEY, TP0_VALUE));
+                        consumer.addRecord(new ConsumerRecord<>(TOPIC, 1, 0, 0L, TimestampType.CREATE_TIME, 0L, 0, 0, TP0_KEY, TP0_VALUE_NEW));
+                    }
+                });
+
+                consumer.schedulePollTask(new Runnable() {
+                    @Override
+                    public void run() {
+                        finishedLatch.countDown();
+                    }
+                });
+            }
+        });
+        readEndFutureCallback.get(10000, TimeUnit.MILLISECONDS);
+        assertTrue(getInvoked.get());
+        assertTrue(finishedLatch.await(10000, TimeUnit.MILLISECONDS));
+        assertEquals(CONSUMER_ASSIGNMENT, consumer.assignment());
+        assertEquals(1L, consumer.position(TP0));
+
+        store.stop();
+
+        assertFalse(Whitebox.getInternalState(store, "thread").isAlive());
+        assertTrue(consumer.closed());
+        PowerMock.verifyAll();
+    }
+
     @Test
     public void testProducerError() throws Exception {
         expectStart();
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/util/LoggingContextTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/util/LoggingContextTest.java
new file mode 100644
index 0000000000000..347e508ecf0e6
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/util/LoggingContextTest.java
@@ -0,0 +1,207 @@
+/*
+ * 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.util;
+
+import org.apache.kafka.connect.util.LoggingContext.Scope;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.slf4j.MDC;
+
+import java.util.HashMap;
+import java.util.Map;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
+
+public class LoggingContextTest {
+
+    private static final Logger log = LoggerFactory.getLogger(LoggingContextTest.class);
+
+    private static final String CONNECTOR_NAME = "MyConnector";
+    private static final ConnectorTaskId TASK_ID1 = new ConnectorTaskId(CONNECTOR_NAME, 1);
+    private static final String EXTRA_KEY1 = "extra.key.1";
+    private static final String EXTRA_VALUE1 = "value1";
+    private static final String EXTRA_KEY2 = "extra.key.2";
+    private static final String EXTRA_VALUE2 = "value2";
+    private static final String EXTRA_KEY3 = "extra.key.3";
+    private static final String EXTRA_VALUE3 = "value3";
+
+    private Map mdc;
+
+    @Before
+    public void setup() {
+        mdc = new HashMap<>();
+        Map existing = MDC.getCopyOfContextMap();
+        if (existing != null) {
+            mdc.putAll(existing);
+        }
+        MDC.put(EXTRA_KEY1, EXTRA_VALUE1);
+        MDC.put(EXTRA_KEY2, EXTRA_VALUE2);
+    }
+
+    @After
+    public void tearDown() {
+        MDC.clear();
+        MDC.setContextMap(mdc);
+    }
+
+    @Test(expected = NullPointerException.class)
+    public void shouldNotAllowNullConnectorNameForConnectorContext() {
+        LoggingContext.forConnector(null);
+    }
+
+    @Test(expected = NullPointerException.class)
+    public void shouldNotAllowNullTaskIdForTaskContext() {
+        LoggingContext.forTask(null);
+    }
+
+    @Test(expected = NullPointerException.class)
+    public void shouldNotAllowNullTaskIdForOffsetContext() {
+        LoggingContext.forOffsets(null);
+    }
+
+    @Test
+    public void shouldCreateAndCloseLoggingContextEvenWithNullContextMap() {
+        MDC.clear();
+        assertMdc(null, null, null);
+        try (LoggingContext loggingContext = LoggingContext.forConnector(CONNECTOR_NAME)) {
+            assertMdc(CONNECTOR_NAME, null, Scope.WORKER);
+            log.info("Starting Connector");
+        }
+        assertMdc(null, null, null);
+    }
+
+    @Test
+    public void shouldCreateConnectorLoggingContext() {
+        assertMdcExtrasUntouched();
+        assertMdc(null, null, null);
+
+        try (LoggingContext loggingContext = LoggingContext.forConnector(CONNECTOR_NAME)) {
+            assertMdc(CONNECTOR_NAME, null, Scope.WORKER);
+            log.info("Starting Connector");
+        }
+
+        assertMdcExtrasUntouched();
+        assertMdc(null, null, null);
+    }
+
+    @Test
+    public void shouldCreateTaskLoggingContext() {
+        assertMdcExtrasUntouched();
+        try (LoggingContext loggingContext = LoggingContext.forTask(TASK_ID1)) {
+            assertMdc(TASK_ID1.connector(), TASK_ID1.task(), Scope.TASK);
+            log.info("Running task");
+        }
+
+        assertMdcExtrasUntouched();
+        assertMdc(null, null, null);
+    }
+
+    @Test
+    public void shouldCreateOffsetsLoggingContext() {
+        assertMdcExtrasUntouched();
+        try (LoggingContext loggingContext = LoggingContext.forOffsets(TASK_ID1)) {
+            assertMdc(TASK_ID1.connector(), TASK_ID1.task(), Scope.OFFSETS);
+            log.info("Running task");
+        }
+
+        assertMdcExtrasUntouched();
+        assertMdc(null, null, null);
+    }
+
+    @Test
+    public void shouldAllowNestedLoggingContexts() {
+        assertMdcExtrasUntouched();
+        assertMdc(null, null, null);
+        try (LoggingContext loggingContext1 = LoggingContext.forConnector(CONNECTOR_NAME)) {
+            assertMdc(CONNECTOR_NAME, null, Scope.WORKER);
+            log.info("Starting Connector");
+            // Set the extra MDC parameter, as if the connector were
+            MDC.put(EXTRA_KEY3, EXTRA_VALUE3);
+            assertConnectorMdcSet();
+
+            try (LoggingContext loggingContext2 = LoggingContext.forTask(TASK_ID1)) {
+                assertMdc(TASK_ID1.connector(), TASK_ID1.task(), Scope.TASK);
+                log.info("Starting task");
+                // The extra connector-specific MDC parameter should still be set
+                assertConnectorMdcSet();
+
+                try (LoggingContext loggingContext3 = LoggingContext.forOffsets(TASK_ID1)) {
+                    assertMdc(TASK_ID1.connector(), TASK_ID1.task(), Scope.OFFSETS);
+                    assertConnectorMdcSet();
+                    log.info("Offsets for task");
+                }
+
+                assertMdc(TASK_ID1.connector(), TASK_ID1.task(), Scope.TASK);
+                log.info("Stopping task");
+                // The extra connector-specific MDC parameter should still be set
+                assertConnectorMdcSet();
+            }
+
+            assertMdc(CONNECTOR_NAME, null, Scope.WORKER);
+            log.info("Stopping Connector");
+            // The extra connector-specific MDC parameter should still be set
+            assertConnectorMdcSet();
+        }
+        assertMdcExtrasUntouched();
+        assertMdc(null, null, null);
+
+        // The extra connector-specific MDC parameter should still be set
+        assertConnectorMdcSet();
+
+        LoggingContext.clear();
+        assertConnectorMdcUnset();
+    }
+
+    protected void assertMdc(String connectorName, Integer taskId, Scope scope) {
+        String context = MDC.get(LoggingContext.CONNECTOR_CONTEXT);
+        if (context != null) {
+            assertEquals(
+                "Context should begin with connector name when the connector name is non-null",
+                connectorName != null,
+                context.startsWith("[" + connectorName)
+            );
+            if (scope != null) {
+                assertTrue("Context should contain the scope", context.contains(scope.toString()));
+            }
+            if (taskId != null) {
+                assertTrue("Context should contain the taskId", context.contains(taskId.toString()));
+            }
+        } else {
+            assertNull("No logging context found, expected null connector name", connectorName);
+            assertNull("No logging context found, expected null task ID", taskId);
+            assertNull("No logging context found, expected null scope", scope);
+        }
+    }
+
+    protected void assertMdcExtrasUntouched() {
+        assertEquals(EXTRA_VALUE1, MDC.get(EXTRA_KEY1));
+        assertEquals(EXTRA_VALUE2, MDC.get(EXTRA_KEY2));
+    }
+
+    protected void assertConnectorMdcSet() {
+        assertEquals(EXTRA_VALUE3, MDC.get(EXTRA_KEY3));
+    }
+
+    protected void assertConnectorMdcUnset() {
+        assertEquals(null, MDC.get(EXTRA_KEY3));
+    }
+}
\ No newline at end of file
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 64ddfebf0e18b..cfcab32c7754c 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
@@ -72,6 +72,18 @@ public void returnNullWithClusterAuthorizationFailure() {
         }
     }
 
+    @Test
+    public void returnNullWithTopicAuthorizationFailure() {
+        final NewTopic newTopic = TopicAdmin.defineTopic("myTopic").partitions(1).compacted().build();
+        Cluster cluster = createCluster(1);
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(new MockTime(), cluster)) {
+            env.kafkaClient().prepareResponse(createTopicResponseWithTopicAuthorizationException(newTopic));
+            TopicAdmin admin = new TopicAdmin(null, env.adminClient());
+            boolean created = admin.createTopic(newTopic);
+            assertFalse(created);
+        }
+    }
+
     @Test
     public void shouldNotCreateTopicWhenItAlreadyExists() {
         NewTopic newTopic = TopicAdmin.defineTopic("myTopic").partitions(1).compacted().build();
@@ -136,6 +148,10 @@ private CreateTopicsResponse createTopicResponseWithClusterAuthorizationExceptio
         return createTopicResponse(new ApiError(Errors.CLUSTER_AUTHORIZATION_FAILED, "Not authorized to create topic(s)"), topics);
     }
 
+    private CreateTopicsResponse createTopicResponseWithTopicAuthorizationException(NewTopic... topics) {
+        return createTopicResponse(new ApiError(Errors.TOPIC_AUTHORIZATION_FAILED, "Not authorized to create topic(s)"), topics);
+    }
+
     private CreateTopicsResponse createTopicResponse(ApiError error, NewTopic... topics) {
         if (error == null) error = new ApiError(Errors.NONE, "");
         CreateTopicsResponseData response = new CreateTopicsResponseData();
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/util/clusters/EmbeddedConnectCluster.java b/connect/runtime/src/test/java/org/apache/kafka/connect/util/clusters/EmbeddedConnectCluster.java
index b660a1dfa5b94..961b1d840c63d 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/util/clusters/EmbeddedConnectCluster.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/util/clusters/EmbeddedConnectCluster.java
@@ -17,25 +17,32 @@
 package org.apache.kafka.connect.util.clusters;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
-import org.apache.kafka.connect.cli.ConnectDistributed;
+import org.apache.kafka.common.utils.Exit;
 import org.apache.kafka.connect.errors.ConnectException;
-import org.apache.kafka.connect.runtime.Connect;
 import org.apache.kafka.connect.runtime.rest.entities.ConnectorStateInfo;
+import org.apache.kafka.connect.runtime.rest.entities.ServerInfo;
 import org.apache.kafka.connect.runtime.rest.errors.ConnectRestException;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import javax.servlet.http.HttpServletResponse;
 import javax.ws.rs.core.Response;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStreamWriter;
 import java.net.HttpURLConnection;
 import java.net.URL;
+import java.util.Collection;
+import java.util.Collections;
 import java.util.HashMap;
+import java.util.Iterator;
+import java.util.LinkedHashSet;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Properties;
+import java.util.Set;
 import java.util.UUID;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Collectors;
 
 import static org.apache.kafka.clients.consumer.ConsumerConfig.GROUP_ID_CONFIG;
 import static org.apache.kafka.connect.runtime.ConnectorConfig.KEY_CONVERTER_CLASS_CONFIG;
@@ -52,7 +59,8 @@
 
 /**
  * Start an embedded connect worker. Internally, this class will spin up a Kafka and Zk cluster, setup any tmp
- * directories and clean up them on them.
+ * directories and clean up them on them. Methods on the same {@code EmbeddedConnectCluster} are
+ * not guaranteed to be thread-safe.
  */
 public class EmbeddedConnectCluster {
 
@@ -63,24 +71,66 @@ public class EmbeddedConnectCluster {
     private static final Properties DEFAULT_BROKER_CONFIG = new Properties();
     private static final String REST_HOST_NAME = "localhost";
 
-    private final Connect[] connectCluster;
+    private static final String DEFAULT_WORKER_NAME_PREFIX = "connect-worker-";
+
+    private final Set connectCluster;
     private final EmbeddedKafkaCluster kafkaCluster;
     private final Map workerProps;
     private final String connectClusterName;
     private final int numBrokers;
-
-    private EmbeddedConnectCluster(String name, Map workerProps, int numWorkers, int numBrokers, Properties brokerProps) {
+    private final int numInitialWorkers;
+    private final boolean maskExitProcedures;
+    private final String workerNamePrefix;
+    private final AtomicInteger nextWorkerId = new AtomicInteger(0);
+    private final EmbeddedConnectClusterAssertions assertions;
+
+    private EmbeddedConnectCluster(String name, Map workerProps, int numWorkers,
+                                   int numBrokers, Properties brokerProps,
+                                   boolean maskExitProcedures) {
         this.workerProps = workerProps;
         this.connectClusterName = name;
         this.numBrokers = numBrokers;
         this.kafkaCluster = new EmbeddedKafkaCluster(numBrokers, brokerProps);
-        this.connectCluster = new Connect[numWorkers];
+        this.connectCluster = new LinkedHashSet<>();
+        this.numInitialWorkers = numWorkers;
+        this.maskExitProcedures = maskExitProcedures;
+        // leaving non-configurable for now
+        this.workerNamePrefix = DEFAULT_WORKER_NAME_PREFIX;
+        this.assertions = new EmbeddedConnectClusterAssertions(this);
     }
 
+    /**
+     * A more graceful way to handle abnormal exit of services in integration tests.
+     */
+    public Exit.Procedure exitProcedure = (code, message) -> {
+        if (code != 0) {
+            String exitMessage = "Abrupt service exit with code " + code + " and message " + message;
+            log.warn(exitMessage);
+            throw new UngracefulShutdownException(exitMessage);
+        }
+        Exit.exit(0, message);
+    };
+
+    /**
+     * A more graceful way to handle abnormal halt of services in integration tests.
+     */
+    public Exit.Procedure haltProcedure = (code, message) -> {
+        if (code != 0) {
+            String haltMessage = "Abrupt service halt with code " + code + " and message " + message;
+            log.warn(haltMessage);
+            throw new UngracefulShutdownException(haltMessage);
+        }
+        Exit.halt(0, message);
+    };
+
     /**
      * Start the connect cluster and the embedded Kafka and Zookeeper cluster.
      */
-    public void start() throws IOException {
+    public void start() {
+        if (maskExitProcedures) {
+            Exit.setExitProcedure(exitProcedure);
+            Exit.setHaltProcedure(haltProcedure);
+        }
         kafkaCluster.before();
         startConnect();
     }
@@ -88,28 +138,79 @@ public void start() throws IOException {
     /**
      * Stop the connect cluster and the embedded Kafka and Zookeeper cluster.
      * Clean up any temp directories created locally.
+     *
+     * @throws RuntimeException if Kafka brokers fail to stop
      */
     public void stop() {
-        for (Connect worker : this.connectCluster) {
-            try {
-                worker.stop();
-            } catch (Exception e) {
-                log.error("Could not stop connect", e);
-                throw new RuntimeException("Could not stop worker", e);
-            }
-        }
-
+        connectCluster.forEach(this::stopWorker);
         try {
             kafkaCluster.after();
+        } catch (UngracefulShutdownException e) {
+            log.warn("Kafka did not shutdown gracefully");
         } catch (Exception e) {
             log.error("Could not stop kafka", e);
             throw new RuntimeException("Could not stop brokers", e);
+        } finally {
+            Exit.resetExitProcedure();
+            Exit.resetHaltProcedure();
+        }
+    }
+
+    /**
+     * Provision and start an additional worker to the Connect cluster.
+     *
+     * @return the worker handle of the worker that was provisioned
+     */
+    public WorkerHandle addWorker() {
+        WorkerHandle worker = WorkerHandle.start(workerNamePrefix + nextWorkerId.getAndIncrement(), workerProps);
+        connectCluster.add(worker);
+        log.info("Started worker {}", worker);
+        return worker;
+    }
+
+    /**
+     * Decommission one of the workers from this Connect cluster. Which worker is removed is
+     * implementation dependent and selection is not guaranteed to be consistent. Use this method
+     * when you don't care which worker stops.
+     *
+     * @see #removeWorker(WorkerHandle)
+     */
+    public void removeWorker() {
+        WorkerHandle toRemove = null;
+        for (Iterator it = connectCluster.iterator(); it.hasNext(); toRemove = it.next()) {
+        }
+        removeWorker(toRemove);
+    }
+
+    /**
+     * Decommission a specific worker from this Connect cluster.
+     *
+     * @param worker the handle of the worker to remove from the cluster
+     * @throws IllegalStateException if the Connect cluster has no workers
+     */
+    public void removeWorker(WorkerHandle worker) {
+        if (connectCluster.isEmpty()) {
+            throw new IllegalStateException("Cannot remove worker. Cluster is empty");
+        }
+        stopWorker(worker);
+        connectCluster.remove(worker);
+    }
+
+    private void stopWorker(WorkerHandle worker) {
+        try {
+            log.info("Stopping worker {}", worker);
+            worker.stop();
+        } catch (UngracefulShutdownException e) {
+            log.warn("Worker {} did not shutdown gracefully", worker);
+        } catch (Exception e) {
+            log.error("Could not stop connect", e);
+            throw new RuntimeException("Could not stop worker", e);
         }
     }
 
     @SuppressWarnings("deprecation")
     public void startConnect() {
-        log.info("Starting Connect cluster with {} workers. clusterName {}", connectCluster.length, connectClusterName);
+        log.info("Starting Connect cluster '{}' with {} workers", connectClusterName, numInitialWorkers);
 
         workerProps.put(BOOTSTRAP_SERVERS_CONFIG, kafka().bootstrapServers());
         workerProps.put(REST_HOST_NAME_CONFIG, REST_HOST_NAME);
@@ -126,71 +227,165 @@ public void startConnect() {
         putIfAbsent(workerProps, KEY_CONVERTER_CLASS_CONFIG, "org.apache.kafka.connect.storage.StringConverter");
         putIfAbsent(workerProps, VALUE_CONVERTER_CLASS_CONFIG, "org.apache.kafka.connect.storage.StringConverter");
 
-        for (int i = 0; i < connectCluster.length; i++) {
-            connectCluster[i] = new ConnectDistributed().startConnect(workerProps);
+        for (int i = 0; i < numInitialWorkers; i++) {
+            addWorker();
         }
     }
 
+    /**
+     * Get the workers that are up and running.
+     *
+     * @return the list of handles of the online workers
+     */
+    public Set activeWorkers() {
+        ObjectMapper mapper = new ObjectMapper();
+        return connectCluster.stream()
+                .filter(w -> {
+                    try {
+                        mapper.readerFor(ServerInfo.class)
+                                .readValue(responseToString(requestGet(w.url().toString())));
+                        return true;
+                    } catch (ConnectException | IOException e) {
+                        // Worker failed to respond. Consider it's offline
+                        return false;
+                    }
+                })
+                .collect(Collectors.toSet());
+    }
+
+    /**
+     * Get the provisioned workers.
+     *
+     * @return the list of handles of the provisioned workers
+     */
+    public Set workers() {
+        return new LinkedHashSet<>(connectCluster);
+    }
+
     /**
      * Configure a connector. If the connector does not already exist, a new one will be created and
      * the given configuration will be applied to it.
      *
      * @param connName   the name of the connector
      * @param connConfig the intended configuration
-     * @throws IOException          if call to the REST api fails.
-     * @throws ConnectRestException if REST api returns error status
+     * @throws ConnectRestException if the REST api returns error status
+     * @throws ConnectException if the configuration fails to be serialized or if the request could not be sent
      */
-    public void configureConnector(String connName, Map connConfig) throws IOException {
+    public String configureConnector(String connName, Map connConfig) {
         String url = endpointForResource(String.format("connectors/%s/config", connName));
         ObjectMapper mapper = new ObjectMapper();
-        int status;
+        String content;
         try {
-            String content = mapper.writeValueAsString(connConfig);
-            status = executePut(url, content);
+            content = mapper.writeValueAsString(connConfig);
         } catch (IOException e) {
-            log.error("Could not execute PUT request to " + url, e);
-            throw e;
+            throw new ConnectException("Could not serialize connector configuration and execute PUT request");
         }
-        if (status >= HttpServletResponse.SC_BAD_REQUEST) {
-            throw new ConnectRestException(status, "Could not execute PUT request");
+        Response response = requestPut(url, content);
+        if (response.getStatus() < Response.Status.BAD_REQUEST.getStatusCode()) {
+            return responseToString(response);
         }
+        throw new ConnectRestException(response.getStatus(),
+                "Could not execute PUT request. Error response: " + responseToString(response));
     }
 
     /**
      * Delete an existing connector.
      *
      * @param connName name of the connector to be deleted
-     * @throws IOException if call to the REST api fails.
+     * @throws ConnectRestException if the REST api returns error status
+     * @throws ConnectException for any other error.
      */
-    public void deleteConnector(String connName) throws IOException {
+    public void deleteConnector(String connName) {
         String url = endpointForResource(String.format("connectors/%s", connName));
-        int status = executeDelete(url);
-        if (status >= HttpServletResponse.SC_BAD_REQUEST) {
-            throw new ConnectRestException(status, "Could not execute DELETE request.");
+        Response response = requestDelete(url);
+        if (response.getStatus() >= Response.Status.BAD_REQUEST.getStatusCode()) {
+            throw new ConnectRestException(response.getStatus(),
+                    "Could not execute DELETE request. Error response: " + responseToString(response));
+        }
+    }
+
+    /**
+     * Get the connector names of the connectors currently running on this cluster.
+     *
+     * @return the list of connector names
+     * @throws ConnectRestException if the HTTP request to the REST API failed with a valid status code.
+     * @throws ConnectException for any other error.
+     */
+    public Collection connectors() {
+        ObjectMapper mapper = new ObjectMapper();
+        String url = endpointForResource("connectors");
+        Response response = requestGet(url);
+        if (response.getStatus() < Response.Status.BAD_REQUEST.getStatusCode()) {
+            try {
+                return mapper.readerFor(Collection.class).readValue(responseToString(response));
+            } catch (IOException e) {
+                log.error("Could not parse connector list from response: {}",
+                        responseToString(response), e
+                );
+                throw new ConnectException("Could not not parse connector list", e);
+            }
         }
+        throw new ConnectRestException(response.getStatus(),
+                "Could not read connector list. Error response: " + responseToString(response));
     }
 
     /**
      * Get the status for a connector running in this cluster.
      *
      * @param connectorName name of the connector
-     * @return an instance of {@link ConnectorStateInfo} populated with state informaton of the connector and it's tasks.
+     * @return an instance of {@link ConnectorStateInfo} populated with state information of the connector and its tasks.
      * @throws ConnectRestException if the HTTP request to the REST API failed with a valid status code.
      * @throws ConnectException for any other error.
      */
     public ConnectorStateInfo connectorStatus(String connectorName) {
         ObjectMapper mapper = new ObjectMapper();
         String url = endpointForResource(String.format("connectors/%s/status", connectorName));
+        Response response = requestGet(url);
         try {
-            return mapper.readerFor(ConnectorStateInfo.class).readValue(executeGet(url));
+            if (response.getStatus() < Response.Status.BAD_REQUEST.getStatusCode()) {
+                return mapper.readerFor(ConnectorStateInfo.class)
+                        .readValue(responseToString(response));
+            }
         } catch (IOException e) {
-            log.error("Could not read connector state", e);
-            throw new ConnectException("Could not read connector state", e);
+            log.error("Could not read connector state from response: {}",
+                    responseToString(response), e);
+            throw new ConnectException("Could not not parse connector state", e);
         }
+        throw new ConnectRestException(response.getStatus(),
+                "Could not read connector state. Error response: " + responseToString(response));
+    }
+
+    /**
+     * Get the full URL of the admin endpoint that corresponds to the given REST resource
+     *
+     * @param resource the resource under the worker's admin endpoint
+     * @return the admin endpoint URL
+     * @throws ConnectException if no admin REST endpoint is available
+     */
+    public String adminEndpoint(String resource) {
+        String url = connectCluster.stream()
+                .map(WorkerHandle::adminUrl)
+                .filter(Objects::nonNull)
+                .findFirst()
+                .orElseThrow(() -> new ConnectException("Admin endpoint is disabled."))
+                .toString();
+        return url + resource;
     }
 
-    private String endpointForResource(String resource) {
-        String url = String.valueOf(connectCluster[0].restUrl());
+    /**
+     * Get the full URL of the endpoint that corresponds to the given REST resource
+     *
+     * @param resource the resource under the worker's admin endpoint
+     * @return the admin endpoint URL
+     * @throws ConnectException if no REST endpoint is available
+     */
+    public String endpointForResource(String resource) {
+        String url = connectCluster.stream()
+                .map(WorkerHandle::url)
+                .filter(Objects::nonNull)
+                .findFirst()
+                .orElseThrow(() -> new ConnectException("Connect workers have not been provisioned"))
+                .toString();
         return url + resource;
     }
 
@@ -200,68 +395,170 @@ private static void putIfAbsent(Map props, String propertyKey, S
         }
     }
 
+    /**
+     * Return the handle to the Kafka cluster this Connect cluster connects to.
+     *
+     * @return the Kafka cluster handle
+     */
     public EmbeddedKafkaCluster kafka() {
         return kafkaCluster;
     }
 
-    public int executePut(String url, String body) throws IOException {
-        log.debug("Executing PUT request to URL={}. Payload={}", url, body);
-        HttpURLConnection httpCon = (HttpURLConnection) new URL(url).openConnection();
-        httpCon.setDoOutput(true);
-        httpCon.setRequestProperty("Content-Type", "application/json");
-        httpCon.setRequestMethod("PUT");
-        try (OutputStreamWriter out = new OutputStreamWriter(httpCon.getOutputStream())) {
-            out.write(body);
-        }
-        try (InputStream is = httpCon.getInputStream()) {
-            int c;
-            StringBuilder response = new StringBuilder();
-            while ((c = is.read()) != -1) {
-                response.append((char) c);
-            }
-            log.info("Put response for URL={} is {}", url, response);
-        }
-        return httpCon.getResponseCode();
+    /**
+     * Execute a GET request on the given URL.
+     *
+     * @param url the HTTP endpoint
+     * @return the response to the GET request
+     * @throws ConnectException if execution of the GET request fails
+     * @deprecated Use {@link #requestGet(String)} instead.
+     */
+    @Deprecated
+    public String executeGet(String url) {
+        return responseToString(requestGet(url));
     }
 
     /**
      * Execute a GET request on the given URL.
      *
      * @param url the HTTP endpoint
-     * @return response body encoded as a String
-     * @throws ConnectRestException if the HTTP request fails with a valid status code
-     * @throws IOException for any other I/O error.
+     * @return the response to the GET request
+     * @throws ConnectException if execution of the GET request fails
      */
-    public String executeGet(String url) throws IOException {
-        log.debug("Executing GET request to URL={}.", url);
-        HttpURLConnection httpCon = (HttpURLConnection) new URL(url).openConnection();
-        httpCon.setDoOutput(true);
-        httpCon.setRequestMethod("GET");
-        try (InputStream is = httpCon.getInputStream()) {
-            int c;
-            StringBuilder response = new StringBuilder();
-            while ((c = is.read()) != -1) {
-                response.append((char) c);
+    public Response requestGet(String url) {
+        return requestHttpMethod(url, null, Collections.emptyMap(), "GET");
+    }
+
+    /**
+     * Execute a PUT request on the given URL.
+     *
+     * @param url the HTTP endpoint
+     * @param body the payload of the PUT request
+     * @return the response to the PUT request
+     * @throws ConnectException if execution of the PUT request fails
+     * @deprecated Use {@link #requestPut(String, String)} instead.
+     */
+    @Deprecated
+    public int executePut(String url, String body) {
+        return requestPut(url, body).getStatus();
+    }
+
+    /**
+     * Execute a PUT request on the given URL.
+     *
+     * @param url the HTTP endpoint
+     * @param body the payload of the PUT request
+     * @return the response to the PUT request
+     * @throws ConnectException if execution of the PUT request fails
+     */
+    public Response requestPut(String url, String body) {
+        return requestHttpMethod(url, body, Collections.emptyMap(), "PUT");
+    }
+
+    /**
+     * Execute a POST request on the given URL.
+     *
+     * @param url the HTTP endpoint
+     * @param body the payload of the POST request
+     * @param headers a map that stores the POST request headers
+     * @return the response to the POST request
+     * @throws ConnectException if execution of the POST request fails
+     * @deprecated Use {@link #requestPost(String, String, java.util.Map)} instead.
+     */
+    @Deprecated
+    public int executePost(String url, String body, Map headers) {
+        return requestPost(url, body, headers).getStatus();
+    }
+
+    /**
+     * Execute a POST request on the given URL.
+     *
+     * @param url the HTTP endpoint
+     * @param body the payload of the POST request
+     * @param headers a map that stores the POST request headers
+     * @return the response to the POST request
+     * @throws ConnectException if execution of the POST request fails
+     */
+    public Response requestPost(String url, String body, Map headers) {
+        return requestHttpMethod(url, body, headers, "POST");
+    }
+
+    /**
+     * Execute a DELETE request on the given URL.
+     *
+     * @param url the HTTP endpoint
+     * @return the response to the DELETE request
+     * @throws ConnectException if execution of the DELETE request fails
+     * @deprecated Use {@link #requestDelete(String)} instead.
+     */
+    @Deprecated
+    public int executeDelete(String url) {
+        return requestDelete(url).getStatus();
+    }
+
+    /**
+     * Execute a DELETE request on the given URL.
+     *
+     * @param url the HTTP endpoint
+     * @return the response to the DELETE request
+     * @throws ConnectException if execution of the DELETE request fails
+     */
+    public Response requestDelete(String url) {
+        return requestHttpMethod(url, null, Collections.emptyMap(), "DELETE");
+    }
+
+    /**
+     * A general method that executes an HTTP request on a given URL.
+     *
+     * @param url the HTTP endpoint
+     * @param body the payload of the request; null if there isn't one
+     * @param headers a map that stores the request headers; empty if there are no headers
+     * @param httpMethod the name of the HTTP method to execute
+     * @return the response to the HTTP request
+     * @throws ConnectException if execution of the HTTP method fails
+     */
+    protected Response requestHttpMethod(String url, String body, Map headers,
+                                      String httpMethod) {
+        log.debug("Executing {} request to URL={}." + (body != null ? " Payload={}" : ""),
+                httpMethod, url, body);
+        try {
+            HttpURLConnection httpCon = (HttpURLConnection) new URL(url).openConnection();
+            httpCon.setDoOutput(true);
+            httpCon.setRequestMethod(httpMethod);
+            if (body != null) {
+                httpCon.setRequestProperty("Content-Type", "application/json");
+                headers.forEach(httpCon::setRequestProperty);
+                try (OutputStreamWriter out = new OutputStreamWriter(httpCon.getOutputStream())) {
+                    out.write(body);
+                }
             }
-            log.debug("Get response for URL={} is {}", url, response);
-            return response.toString();
-        } catch (IOException e) {
-            Response.Status status = Response.Status.fromStatusCode(httpCon.getResponseCode());
-            if (status != null) {
-                throw new ConnectRestException(status, "Invalid endpoint: " + url, e);
+            try (InputStream is = httpCon.getResponseCode() < HttpURLConnection.HTTP_BAD_REQUEST
+                                  ? httpCon.getInputStream()
+                                  : httpCon.getErrorStream()
+            ) {
+                String responseEntity = responseToString(is);
+                log.info("{} response for URL={} is {}",
+                        httpMethod, url, responseEntity.isEmpty() ? "empty" : responseEntity);
+                return Response.status(Response.Status.fromStatusCode(httpCon.getResponseCode()))
+                        .entity(responseEntity)
+                        .build();
             }
-            // invalid response code, re-throw the IOException.
-            throw e;
+        } catch (IOException e) {
+            log.error("Could not execute " + httpMethod + " request to " + url, e);
+            throw new ConnectException(e);
         }
     }
 
-    public int executeDelete(String url) throws IOException {
-        log.debug("Executing DELETE request to URL={}", url);
-        HttpURLConnection httpCon = (HttpURLConnection) new URL(url).openConnection();
-        httpCon.setDoOutput(true);
-        httpCon.setRequestMethod("DELETE");
-        httpCon.connect();
-        return httpCon.getResponseCode();
+    private String responseToString(Response response) {
+        return response == null ? "empty" : (String) response.getEntity();
+    }
+
+    private String responseToString(InputStream stream) throws IOException {
+        int c;
+        StringBuilder response = new StringBuilder();
+        while ((c = stream.read()) != -1) {
+            response.append((char) c);
+        }
+        return response.toString();
     }
 
     public static class Builder {
@@ -270,6 +567,7 @@ public static class Builder {
         private int numWorkers = DEFAULT_NUM_WORKERS;
         private int numBrokers = DEFAULT_NUM_BROKERS;
         private Properties brokerProps = DEFAULT_BROKER_CONFIG;
+        private boolean maskExitProcedures = true;
 
         public Builder name(String name) {
             this.name = name;
@@ -296,9 +594,36 @@ public Builder brokerProps(Properties brokerProps) {
             return this;
         }
 
+        /**
+         * In the event of ungraceful shutdown, embedded clusters call exit or halt with non-zero
+         * exit statuses. Exiting with a non-zero status forces a test to fail and is hard to
+         * handle. Because graceful exit is usually not required during a test and because
+         * depending on such an exit increases flakiness, this setting allows masking
+         * exit and halt procedures by using a runtime exception instead. Customization of the
+         * exit and halt procedures is possible through {@code exitProcedure} and {@code
+         * haltProcedure} respectively.
+         *
+         * @param mask if false, exit and halt procedures remain unchanged; true is the default.
+         * @return the builder for this cluster
+         */
+        public Builder maskExitProcedures(boolean mask) {
+            this.maskExitProcedures = mask;
+            return this;
+        }
+
         public EmbeddedConnectCluster build() {
-            return new EmbeddedConnectCluster(name, workerProps, numWorkers, numBrokers, brokerProps);
+            return new EmbeddedConnectCluster(name, workerProps, numWorkers, numBrokers,
+                    brokerProps, maskExitProcedures);
         }
     }
 
+    /**
+     * Return the available assertions for this Connect cluster
+     *
+     * @return the assertions object
+     */
+    public EmbeddedConnectClusterAssertions assertions() {
+        return assertions;
+    }
+
 }
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/util/clusters/EmbeddedConnectClusterAssertions.java b/connect/runtime/src/test/java/org/apache/kafka/connect/util/clusters/EmbeddedConnectClusterAssertions.java
new file mode 100644
index 0000000000000..04bfd66081d96
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/util/clusters/EmbeddedConnectClusterAssertions.java
@@ -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 org.apache.kafka.connect.util.clusters;
+
+import org.apache.kafka.connect.runtime.AbstractStatus;
+import org.apache.kafka.connect.runtime.rest.entities.ConnectorStateInfo;
+import org.apache.kafka.connect.runtime.rest.errors.ConnectRestException;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.ws.rs.core.Response;
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+import java.util.function.BiFunction;
+
+import static org.apache.kafka.test.TestUtils.waitForCondition;
+
+/**
+ * A set of common assertions that can be applied to a Connect cluster during integration testing
+ */
+public class EmbeddedConnectClusterAssertions {
+
+    private static final Logger log = LoggerFactory.getLogger(EmbeddedConnectClusterAssertions.class);
+    public static final long WORKER_SETUP_DURATION_MS = TimeUnit.SECONDS.toMillis(60);
+    public static final long CONNECTOR_SETUP_DURATION_MS = TimeUnit.SECONDS.toMillis(30);
+
+    private final EmbeddedConnectCluster connect;
+
+    EmbeddedConnectClusterAssertions(EmbeddedConnectCluster connect) {
+        this.connect = connect;
+    }
+
+    /**
+     * Assert that at least the requested number of workers are up and running.
+     *
+     * @param numWorkers the number of online workers
+     */
+    public void assertAtLeastNumWorkersAreUp(int numWorkers, String detailMessage) throws InterruptedException {
+        try {
+            waitForCondition(
+                () -> checkWorkersUp(numWorkers, (actual, expected) -> actual >= expected).orElse(false),
+                WORKER_SETUP_DURATION_MS,
+                "Didn't meet the minimum requested number of online workers: " + numWorkers);
+        } catch (AssertionError e) {
+            throw new AssertionError(detailMessage, e);
+        }
+    }
+
+    /**
+     * Assert that at least the requested number of workers are up and running.
+     *
+     * @param numWorkers the number of online workers
+     */
+    public void assertExactlyNumWorkersAreUp(int numWorkers, String detailMessage) throws InterruptedException {
+        try {
+            waitForCondition(
+                () -> checkWorkersUp(numWorkers, (actual, expected) -> actual == expected).orElse(false),
+                WORKER_SETUP_DURATION_MS,
+                "Didn't meet the exact requested number of online workers: " + numWorkers);
+        } catch (AssertionError e) {
+            throw new AssertionError(detailMessage, e);
+        }
+    }
+
+    /**
+     * Confirm that the requested number of workers are up and running.
+     *
+     * @param numWorkers the number of online workers
+     * @return true if at least {@code numWorkers} are up; false otherwise
+     */
+    protected Optional checkWorkersUp(int numWorkers, BiFunction comp) {
+        try {
+            int numUp = connect.activeWorkers().size();
+            return Optional.of(comp.apply(numUp, numWorkers));
+        } catch (Exception e) {
+            log.error("Could not check active workers.", e);
+            return Optional.empty();
+        }
+    }
+
+    /**
+     * Assert that a connector is running with at least the given number of tasks all in running state
+     *
+     * @param connectorName the connector name
+     * @param numTasks the number of tasks
+     * @param detailMessage
+     * @throws InterruptedException
+     */
+    public void assertConnectorAndAtLeastNumTasksAreRunning(String connectorName, int numTasks, String detailMessage)
+            throws InterruptedException {
+        try {
+            waitForCondition(
+                () -> checkConnectorState(
+                    connectorName,
+                    AbstractStatus.State.RUNNING,
+                    numTasks,
+                    AbstractStatus.State.RUNNING,
+                    (actual, expected) -> actual >= expected
+                ).orElse(false),
+                CONNECTOR_SETUP_DURATION_MS,
+                "The connector or at least " + numTasks + " of tasks are not running.");
+        } catch (AssertionError e) {
+            throw new AssertionError(detailMessage, e);
+        }
+    }
+
+    /**
+     * Assert that a connector is running with at least the given number of tasks all in running state
+     *
+     * @param connectorName the connector name
+     * @param numTasks the number of tasks
+     * @param detailMessage the assertion message
+     * @throws InterruptedException
+     */
+    public void assertConnectorAndExactlyNumTasksAreRunning(String connectorName, int numTasks, String detailMessage)
+            throws InterruptedException {
+        try {
+            waitForCondition(
+                () -> checkConnectorState(
+                    connectorName,
+                    AbstractStatus.State.RUNNING,
+                    numTasks,
+                    AbstractStatus.State.RUNNING,
+                    (actual, expected) -> actual == expected
+                ).orElse(false),
+                CONNECTOR_SETUP_DURATION_MS,
+                "The connector or exactly " + numTasks + " tasks are not running.");
+        } catch (AssertionError e) {
+            throw new AssertionError(detailMessage, e);
+        }
+    }
+
+    /**
+     * Assert that a connector is running, that it has a specific number of tasks, and that all of
+     * its tasks are in the FAILED state.
+     *
+     * @param connectorName the connector name
+     * @param numTasks the number of tasks
+     * @param detailMessage the assertion message
+     * @throws InterruptedException
+     */
+    public void assertConnectorIsRunningAndTasksHaveFailed(String connectorName, int numTasks, String detailMessage)
+            throws InterruptedException {
+        try {
+            waitForCondition(
+                () -> checkConnectorState(
+                    connectorName,
+                    AbstractStatus.State.RUNNING,
+                    numTasks,
+                    AbstractStatus.State.FAILED,
+                    (actual, expected) -> actual >= expected
+                ).orElse(false),
+                CONNECTOR_SETUP_DURATION_MS,
+                "Either the connector is not running or not all the " + numTasks + " tasks have failed.");
+        } catch (AssertionError e) {
+            throw new AssertionError(detailMessage, e);
+        }
+    }
+
+    /**
+     * Assert that a connector and its tasks are not running.
+     *
+     * @param connectorName the connector name
+     * @param detailMessage the assertion message
+     * @throws InterruptedException
+     */
+    public void assertConnectorAndTasksAreStopped(String connectorName, String detailMessage)
+            throws InterruptedException {
+        try {
+            waitForCondition(
+                () -> checkConnectorAndTasksAreStopped(connectorName),
+                CONNECTOR_SETUP_DURATION_MS,
+                "At least the connector or one of its tasks is still running");
+        } catch (AssertionError e) {
+            throw new AssertionError(detailMessage, e);
+        }
+    }
+
+    /**
+     * Check whether the connector or any of its tasks are still in RUNNING state
+     *
+     * @param connectorName the connector
+     * @return true if the connector and all the tasks are not in RUNNING state; false otherwise
+     */
+    protected boolean checkConnectorAndTasksAreStopped(String connectorName) {
+        ConnectorStateInfo info;
+        try {
+            info = connect.connectorStatus(connectorName);
+        } catch (ConnectRestException e) {
+            return e.statusCode() == Response.Status.NOT_FOUND.getStatusCode();
+        } catch (Exception e) {
+            log.error("Could not check connector state info.", e);
+            return false;
+        }
+        if (info == null) {
+            return true;
+        }
+        return !info.connector().state().equals(AbstractStatus.State.RUNNING.toString())
+                && info.tasks().stream().noneMatch(s -> s.state().equals(AbstractStatus.State.RUNNING.toString()));
+    }
+
+    /**
+     * Check whether the given connector state matches the current state of the connector and
+     * whether it has at least the given number of tasks, with all the tasks matching the given
+     * task state.
+     * @param connectorName the connector
+     * @param connectorState
+     * @param numTasks the expected number of tasks
+     * @param tasksState
+     * @return true if the connector and tasks are in RUNNING state; false otherwise
+     */
+    protected Optional checkConnectorState(
+            String connectorName,
+            AbstractStatus.State connectorState,
+            int numTasks,
+            AbstractStatus.State tasksState,
+            BiFunction comp
+    ) {
+        try {
+            ConnectorStateInfo info = connect.connectorStatus(connectorName);
+            boolean result = info != null
+                    && comp.apply(info.tasks().size(), numTasks)
+                    && info.connector().state().equals(connectorState.toString())
+                    && info.tasks().stream().allMatch(s -> s.state().equals(tasksState.toString()));
+            return Optional.of(result);
+        } catch (Exception e) {
+            log.error("Could not check connector state info.", e);
+            return Optional.empty();
+        }
+    }
+
+}
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 44643950b4ba6..b365e651ee473 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
@@ -23,7 +23,7 @@
 import kafka.utils.TestUtils;
 import kafka.zk.EmbeddedZookeeper;
 import org.apache.kafka.clients.CommonClientConfigs;
-import org.apache.kafka.clients.admin.AdminClient;
+import org.apache.kafka.clients.admin.Admin;
 import org.apache.kafka.clients.admin.AdminClientConfig;
 import org.apache.kafka.clients.admin.NewTopic;
 import org.apache.kafka.clients.consumer.ConsumerRecord;
@@ -81,6 +81,8 @@ public class EmbeddedKafkaCluster extends ExternalResource {
     private final KafkaServer[] brokers;
     private final Properties brokerConfig;
     private final Time time = new MockTime();
+    private final int[] currentBrokerPorts;
+    private final String[] currentBrokerLogDirs;
 
     private EmbeddedZookeeper zookeeper = null;
     private ListenerName listenerName = new ListenerName("PLAINTEXT");
@@ -89,11 +91,13 @@ public class EmbeddedKafkaCluster extends ExternalResource {
     public EmbeddedKafkaCluster(final int numBrokers,
                                 final Properties brokerConfig) {
         brokers = new KafkaServer[numBrokers];
+        currentBrokerPorts = new int[numBrokers];
+        currentBrokerLogDirs = new String[numBrokers];
         this.brokerConfig = brokerConfig;
     }
 
     @Override
-    protected void before() throws IOException {
+    protected void before() {
         start();
     }
 
@@ -102,11 +106,26 @@ protected void after() {
         stop();
     }
 
-    private void start() throws IOException {
+    /**
+     * Starts the Kafka cluster alone using the ports that were assigned during initialization of
+     * the harness.
+     *
+     * @throws ConnectException if a directory to store the data cannot be created
+     */
+    public void startOnlyKafkaOnSamePorts() {
+        start(currentBrokerPorts, currentBrokerLogDirs);
+    }
+
+    private void start() {
+        // pick a random port
         zookeeper = new EmbeddedZookeeper();
+        Arrays.fill(currentBrokerPorts, 0);
+        Arrays.fill(currentBrokerLogDirs, null);
+        start(currentBrokerPorts, currentBrokerLogDirs);
+    }
 
+    private void start(int[] brokerPorts, String[] logDirs) {
         brokerConfig.put(KafkaConfig$.MODULE$.ZkConnectProp(), zKConnectString());
-        brokerConfig.put(KafkaConfig$.MODULE$.PortProp(), 0); // pick a random port
 
         putIfAbsent(brokerConfig, KafkaConfig$.MODULE$.HostNameProp(), "localhost");
         putIfAbsent(brokerConfig, KafkaConfig$.MODULE$.DeleteTopicEnableProp(), true);
@@ -121,8 +140,11 @@ private void start() throws IOException {
 
         for (int i = 0; i < brokers.length; i++) {
             brokerConfig.put(KafkaConfig$.MODULE$.BrokerIdProp(), i);
-            brokerConfig.put(KafkaConfig$.MODULE$.LogDirProp(), createLogDir());
+            currentBrokerLogDirs[i] = logDirs[i] == null ? createLogDir() : currentBrokerLogDirs[i];
+            brokerConfig.put(KafkaConfig$.MODULE$.LogDirProp(), currentBrokerLogDirs[i]);
+            brokerConfig.put(KafkaConfig$.MODULE$.PortProp(), brokerPorts[i]);
             brokers[i] = TestUtils.createServer(new KafkaConfig(brokerConfig, true), time);
+            currentBrokerPorts[i] = brokers[i].boundPort(listenerName);
         }
 
         Map producerProps = new HashMap<>();
@@ -132,8 +154,15 @@ private void start() throws IOException {
         producer = new KafkaProducer<>(producerProps);
     }
 
+    public void stopOnlyKafka() {
+        stop(false, false);
+    }
+
     private void stop() {
+        stop(true, true);
+    }
 
+    private void stop(boolean deleteLogDirs, boolean stopZK) {
         try {
             producer.close();
         } catch (Exception e) {
@@ -151,19 +180,24 @@ private void stop() {
             }
         }
 
-        for (KafkaServer broker : brokers) {
-            try {
-                log.info("Cleaning up kafka log dirs at {}", broker.config().logDirs());
-                CoreUtils.delete(broker.config().logDirs());
-            } catch (Throwable t) {
-                String msg = String.format("Could not clean up log dirs for broker at %s", address(broker));
-                log.error(msg, t);
-                throw new RuntimeException(msg, t);
+        if (deleteLogDirs) {
+            for (KafkaServer broker : brokers) {
+                try {
+                    log.info("Cleaning up kafka log dirs at {}", broker.config().logDirs());
+                    CoreUtils.delete(broker.config().logDirs());
+                } catch (Throwable t) {
+                    String msg = String.format("Could not clean up log dirs for broker at %s",
+                            address(broker));
+                    log.error(msg, t);
+                    throw new RuntimeException(msg, t);
+                }
             }
         }
 
         try {
-            zookeeper.shutdown();
+            if (stopZK) {
+                zookeeper.shutdown();
+            }
         } catch (Throwable t) {
             String msg = String.format("Could not shutdown zookeeper at %s", zKConnectString());
             log.error(msg, t);
@@ -177,10 +211,15 @@ private static void putIfAbsent(final Properties props, final String propertyKey
         }
     }
 
-    private String createLogDir() throws IOException {
+    private String createLogDir() {
         TemporaryFolder tmpFolder = new TemporaryFolder();
-        tmpFolder.create();
-        return tmpFolder.newFolder().getAbsolutePath();
+        try {
+            tmpFolder.create();
+            return tmpFolder.newFolder().getAbsolutePath();
+        } catch (IOException e) {
+            log.error("Unable to create temporary log directory", e);
+            throw new ConnectException("Unable to create temporary log directory", e);
+        }
     }
 
     public String bootstrapServers() {
@@ -234,7 +273,7 @@ public void createTopic(String topic, int partitions, int replication, Map workerProperties) {
+        return new WorkerHandle(name, new ConnectDistributed().startConnect(workerProperties));
+    }
+
+    /**
+     * Stop this worker.
+     */
+    public void stop() {
+        worker.stop();
+    }
+
+    /**
+     * Get the workers's name corresponding to this handle.
+     *
+     * @return the worker's name
+     */
+    public String name() {
+        return workerName;
+    }
+
+    /**
+     * Get the workers's url that accepts requests to its REST endpoint.
+     *
+     * @return the worker's url
+     */
+    public URI url() {
+        return worker.restUrl();
+    }
+
+    /**
+     * Get the workers's url that accepts requests to its Admin REST endpoint.
+     *
+     * @return the worker's admin url
+     */
+    public URI adminUrl() {
+        return worker.adminUrl();
+    }
+
+    @Override
+    public String toString() {
+        return "WorkerHandle{" +
+                "workerName='" + workerName + '\'' +
+                "workerURL='" + worker.restUrl() + '\'' +
+                '}';
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+        if (!(o instanceof WorkerHandle)) {
+            return false;
+        }
+        WorkerHandle that = (WorkerHandle) o;
+        return Objects.equals(workerName, that.workerName) &&
+                Objects.equals(worker, that.worker);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(workerName, worker);
+    }
+}
diff --git a/connect/runtime/src/test/resources/log4j.properties b/connect/runtime/src/test/resources/log4j.properties
index 1feedb89721bc..176692deb7b2b 100644
--- a/connect/runtime/src/test/resources/log4j.properties
+++ b/connect/runtime/src/test/resources/log4j.properties
@@ -14,11 +14,22 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 ##
-log4j.rootLogger=OFF, stdout
+log4j.rootLogger=INFO, stdout
 
 log4j.appender.stdout=org.apache.log4j.ConsoleAppender
 log4j.appender.stdout.layout=org.apache.log4j.PatternLayout
-log4j.appender.stdout.layout.ConversionPattern=[%d] (%t) %p %m (%c:%L)%n
+#
+# The `%X{connector.context}` parameter in the layout includes connector-specific and task-specific information
+# in the log message, where appropriate. This makes it easier to identify those log messages that apply to a
+# specific connector. Simply add this parameter to the log layout configuration below to include the contextual information.
+#
+log4j.appender.stdout.layout.ConversionPattern=[%d] %p %X{connector.context}%m (%c:%L)%n
+#
+# The following line includes no MDC context parameters:
+#log4j.appender.stdout.layout.ConversionPattern=[%d] %p %m (%c:%L)%n (%t)
 
 log4j.logger.org.reflections=ERROR
-log4j.logger.org.apache.kafka=ERROR
+log4j.logger.kafka=WARN
+log4j.logger.org.apache.kafka.connect=DEBUG
+log4j.logger.org.apache.kafka.connect.runtime.distributed=DEBUG
+log4j.logger.org.apache.kafka.connect.integration=DEBUG
diff --git a/connect/runtime/src/test/resources/test-plugins/aliased-static-field/test/plugins/AliasedStaticField.java b/connect/runtime/src/test/resources/test-plugins/aliased-static-field/test/plugins/AliasedStaticField.java
new file mode 100644
index 0000000000000..d865f4e91ba86
--- /dev/null
+++ b/connect/runtime/src/test/resources/test-plugins/aliased-static-field/test/plugins/AliasedStaticField.java
@@ -0,0 +1,75 @@
+/*
+ * 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 test.plugins;
+
+import java.util.Map;
+import java.util.HashMap;
+import org.apache.kafka.connect.data.Schema;
+import org.apache.kafka.connect.data.SchemaAndValue;
+import org.apache.kafka.connect.storage.Converter;
+import org.apache.kafka.connect.runtime.isolation.SamplingTestPlugin;
+
+/**
+ * Samples data about its initialization environment for later analysis
+ * Samples are shared between instances of the same class in a static variable
+ */
+public class AliasedStaticField extends SamplingTestPlugin implements Converter {
+
+  private static final Map SAMPLES;
+  private static final ClassLoader STATIC_CLASS_LOADER;
+  private final ClassLoader classloader;
+
+  static {
+    SAMPLES = new HashMap<>();
+    STATIC_CLASS_LOADER = Thread.currentThread().getContextClassLoader();
+  }
+
+  {
+    classloader = Thread.currentThread().getContextClassLoader();
+  }
+
+  @Override
+  public void configure(final Map configs, final boolean isKey) {
+
+  }
+
+  @Override
+  public byte[] fromConnectData(final String topic, final Schema schema, final Object value) {
+    return new byte[0];
+  }
+
+  @Override
+  public SchemaAndValue toConnectData(final String topic, final byte[] value) {
+    return null;
+  }
+
+  @Override
+  public ClassLoader staticClassloader() {
+    return STATIC_CLASS_LOADER;
+  }
+
+  @Override
+  public ClassLoader classloader() {
+    return classloader;
+  }
+
+  @Override
+  public Map otherSamples() {
+      return SAMPLES;
+  }
+}
diff --git a/connect/runtime/src/test/resources/test-plugins/always-throw-exception/test/plugins/AlwaysThrowException.java b/connect/runtime/src/test/resources/test-plugins/always-throw-exception/test/plugins/AlwaysThrowException.java
new file mode 100644
index 0000000000000..858f3ed5eacbe
--- /dev/null
+++ b/connect/runtime/src/test/resources/test-plugins/always-throw-exception/test/plugins/AlwaysThrowException.java
@@ -0,0 +1,53 @@
+/*
+ * 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 test.plugins;
+
+import java.util.Map;
+import org.apache.kafka.connect.data.Schema;
+import org.apache.kafka.connect.data.SchemaAndValue;
+import org.apache.kafka.connect.runtime.isolation.SamplingTestPlugin;
+import org.apache.kafka.connect.storage.Converter;
+
+/**
+ * Unconditionally throw an exception during static initialization.
+ */
+public class AlwaysThrowException implements Converter {
+
+    static {
+        setup();
+    }
+
+    public static void setup() {
+        throw new RuntimeException("I always throw an exception");
+    }
+
+    @Override
+    public void configure(final Map configs, final boolean isKey) {
+
+    }
+
+    @Override
+    public byte[] fromConnectData(final String topic, final Schema schema, final Object value) {
+        return new byte[0];
+    }
+
+    @Override
+    public SchemaAndValue toConnectData(final String topic, final byte[] value) {
+        return null;
+    }
+}
diff --git a/connect/runtime/src/test/resources/test-plugins/sampling-config-provider/META-INF/services/org.apache.kafka.common.config.provider.ConfigProvider b/connect/runtime/src/test/resources/test-plugins/sampling-config-provider/META-INF/services/org.apache.kafka.common.config.provider.ConfigProvider
new file mode 100644
index 0000000000000..62d8df254bbc3
--- /dev/null
+++ b/connect/runtime/src/test/resources/test-plugins/sampling-config-provider/META-INF/services/org.apache.kafka.common.config.provider.ConfigProvider
@@ -0,0 +1,16 @@
+ # 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.
+
+test.plugins.SamplingConfigProvider
diff --git a/connect/runtime/src/test/resources/test-plugins/sampling-config-provider/test/plugins/SamplingConfigProvider.java b/connect/runtime/src/test/resources/test-plugins/sampling-config-provider/test/plugins/SamplingConfigProvider.java
new file mode 100644
index 0000000000000..df8285eba9a1a
--- /dev/null
+++ b/connect/runtime/src/test/resources/test-plugins/sampling-config-provider/test/plugins/SamplingConfigProvider.java
@@ -0,0 +1,101 @@
+/*
+ * 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 test.plugins;
+
+import java.util.Set;
+import java.util.Map;
+import java.util.HashMap;
+import org.apache.kafka.common.config.provider.ConfigProvider;
+import org.apache.kafka.common.config.ConfigData;
+import org.apache.kafka.common.config.ConfigChangeCallback;
+import org.apache.kafka.connect.data.Schema;
+import org.apache.kafka.connect.data.SchemaAndValue;
+import org.apache.kafka.connect.storage.Converter;
+import org.apache.kafka.connect.runtime.isolation.SamplingTestPlugin;
+import org.apache.kafka.connect.storage.HeaderConverter;
+
+/**
+ * Samples data about its initialization environment for later analysis
+ */
+public class SamplingConfigProvider extends SamplingTestPlugin implements ConfigProvider {
+
+  private static final ClassLoader STATIC_CLASS_LOADER;
+  private final ClassLoader classloader;
+  private Map samples;
+
+  static {
+    STATIC_CLASS_LOADER = Thread.currentThread().getContextClassLoader();
+  }
+
+  {
+    samples = new HashMap<>();
+    classloader = Thread.currentThread().getContextClassLoader();
+  }
+
+  @Override
+  public ConfigData get(String path) {
+    logMethodCall(samples);
+    return null;
+  }
+
+  @Override
+  public ConfigData get(String path, Set keys) {
+    logMethodCall(samples);
+    return null;
+  }
+
+  @Override
+  public void subscribe(String path, Set keys, ConfigChangeCallback callback) {
+    logMethodCall(samples);
+  }
+
+  @Override
+  public void unsubscribe(String path, Set keys, ConfigChangeCallback callback) {
+    logMethodCall(samples);
+  }
+
+  @Override
+  public void unsubscribeAll() {
+    logMethodCall(samples);
+  }
+
+  @Override
+  public void configure(final Map configs) {
+    logMethodCall(samples);
+  }
+
+  @Override
+  public void close() {
+    logMethodCall(samples);
+  }
+
+  @Override
+  public ClassLoader staticClassloader() {
+    return STATIC_CLASS_LOADER;
+  }
+
+  @Override
+  public ClassLoader classloader() {
+    return classloader;
+  }
+
+  @Override
+  public Map otherSamples() {
+    return samples;
+  }
+}
diff --git a/connect/runtime/src/test/resources/test-plugins/sampling-configurable/test/plugins/SamplingConfigurable.java b/connect/runtime/src/test/resources/test-plugins/sampling-configurable/test/plugins/SamplingConfigurable.java
new file mode 100644
index 0000000000000..a917f2f2ca1a8
--- /dev/null
+++ b/connect/runtime/src/test/resources/test-plugins/sampling-configurable/test/plugins/SamplingConfigurable.java
@@ -0,0 +1,79 @@
+/*
+ * 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 test.plugins;
+
+import java.util.Map;
+import java.util.HashMap;
+import org.apache.kafka.common.Configurable;
+import org.apache.kafka.connect.data.Schema;
+import org.apache.kafka.connect.data.SchemaAndValue;
+import org.apache.kafka.connect.storage.Converter;
+import org.apache.kafka.connect.runtime.isolation.SamplingTestPlugin;
+
+/**
+ * Samples data about its initialization environment for later analysis
+ */
+public class SamplingConfigurable extends SamplingTestPlugin implements Converter, Configurable {
+
+  private static final ClassLoader STATIC_CLASS_LOADER;
+  private final ClassLoader classloader;
+  private Map samples;
+
+  static {
+    STATIC_CLASS_LOADER = Thread.currentThread().getContextClassLoader();
+  }
+
+  {
+    samples = new HashMap<>();
+    classloader = Thread.currentThread().getContextClassLoader();
+  }
+
+  @Override
+  public void configure(final Map configs) {
+    logMethodCall(samples);
+  }
+
+  @Override
+  public void configure(final Map configs, final boolean isKey) {
+  }
+
+  @Override
+  public byte[] fromConnectData(final String topic, final Schema schema, final Object value) {
+    return new byte[0];
+  }
+
+  @Override
+  public SchemaAndValue toConnectData(final String topic, final byte[] value) {
+    return null;
+  }
+
+  @Override
+  public ClassLoader staticClassloader() {
+    return STATIC_CLASS_LOADER;
+  }
+
+  @Override
+  public ClassLoader classloader() {
+    return classloader;
+  }
+
+  @Override
+  public Map otherSamples() {
+    return samples;
+  }
+}
diff --git a/connect/runtime/src/test/resources/test-plugins/sampling-converter/test/plugins/SamplingConverter.java b/connect/runtime/src/test/resources/test-plugins/sampling-converter/test/plugins/SamplingConverter.java
new file mode 100644
index 0000000000000..39109a1d4e573
--- /dev/null
+++ b/connect/runtime/src/test/resources/test-plugins/sampling-converter/test/plugins/SamplingConverter.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 test.plugins;
+
+import java.util.Map;
+import java.util.HashMap;
+import org.apache.kafka.connect.data.Schema;
+import org.apache.kafka.connect.data.SchemaAndValue;
+import org.apache.kafka.connect.storage.Converter;
+import org.apache.kafka.connect.runtime.isolation.SamplingTestPlugin;
+
+/**
+ * Samples data about its initialization environment for later analysis
+ */
+public class SamplingConverter extends SamplingTestPlugin implements Converter {
+
+  private static final ClassLoader STATIC_CLASS_LOADER;
+  private final ClassLoader classloader;
+  private Map samples;
+
+  static {
+    STATIC_CLASS_LOADER = Thread.currentThread().getContextClassLoader();
+  }
+
+  {
+    samples = new HashMap<>();
+    classloader = Thread.currentThread().getContextClassLoader();
+  }
+
+  @Override
+  public void configure(final Map configs, final boolean isKey) {
+    logMethodCall(samples);
+  }
+
+  @Override
+  public byte[] fromConnectData(final String topic, final Schema schema, final Object value) {
+    logMethodCall(samples);
+    return new byte[0];
+  }
+
+  @Override
+  public SchemaAndValue toConnectData(final String topic, final byte[] value) {
+    logMethodCall(samples);
+    return null;
+  }
+
+  @Override
+  public ClassLoader staticClassloader() {
+    return STATIC_CLASS_LOADER;
+  }
+
+  @Override
+  public ClassLoader classloader() {
+    return classloader;
+  }
+
+  @Override
+  public Map otherSamples() {
+    return samples;
+  }
+}
diff --git a/connect/runtime/src/test/resources/test-plugins/sampling-header-converter/test/plugins/SamplingHeaderConverter.java b/connect/runtime/src/test/resources/test-plugins/sampling-header-converter/test/plugins/SamplingHeaderConverter.java
new file mode 100644
index 0000000000000..11a1e28e7278c
--- /dev/null
+++ b/connect/runtime/src/test/resources/test-plugins/sampling-header-converter/test/plugins/SamplingHeaderConverter.java
@@ -0,0 +1,89 @@
+/*
+ * 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 test.plugins;
+
+import java.util.Map;
+import java.util.HashMap;
+import org.apache.kafka.common.config.ConfigDef;
+import org.apache.kafka.connect.data.Schema;
+import org.apache.kafka.connect.data.SchemaAndValue;
+import org.apache.kafka.connect.storage.Converter;
+import org.apache.kafka.connect.runtime.isolation.SamplingTestPlugin;
+import org.apache.kafka.connect.storage.HeaderConverter;
+
+/**
+ * Samples data about its initialization environment for later analysis
+ */
+public class SamplingHeaderConverter extends SamplingTestPlugin implements HeaderConverter {
+
+  private static final ClassLoader STATIC_CLASS_LOADER;
+  private final ClassLoader classloader;
+  private Map samples;
+
+  static {
+    STATIC_CLASS_LOADER = Thread.currentThread().getContextClassLoader();
+  }
+
+  {
+    samples = new HashMap<>();
+    classloader = Thread.currentThread().getContextClassLoader();
+  }
+
+  @Override
+  public SchemaAndValue toConnectHeader(String topic, String headerKey, byte[] value) {
+    logMethodCall(samples);
+    return null;
+  }
+
+  @Override
+  public byte[] fromConnectHeader(String topic, String headerKey, Schema schema, Object value) {
+    logMethodCall(samples);
+    return new byte[0];
+  }
+
+  @Override
+  public ConfigDef config() {
+    logMethodCall(samples);
+    return null;
+  }
+
+  @Override
+  public void configure(final Map configs) {
+    logMethodCall(samples);
+  }
+
+  @Override
+  public void close() {
+    logMethodCall(samples);
+  }
+
+  @Override
+  public ClassLoader staticClassloader() {
+    return STATIC_CLASS_LOADER;
+  }
+
+  @Override
+  public ClassLoader classloader() {
+    return classloader;
+  }
+
+  @Override
+  public Map otherSamples() {
+    return samples;
+  }
+}
diff --git a/connect/runtime/src/test/resources/test-plugins/service-loader/META-INF/services/test.plugins.ServiceLoadedClass b/connect/runtime/src/test/resources/test-plugins/service-loader/META-INF/services/test.plugins.ServiceLoadedClass
new file mode 100644
index 0000000000000..b8db8656487d2
--- /dev/null
+++ b/connect/runtime/src/test/resources/test-plugins/service-loader/META-INF/services/test.plugins.ServiceLoadedClass
@@ -0,0 +1,16 @@
+ # 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.
+
+test.plugins.ServiceLoadedSubclass
\ No newline at end of file
diff --git a/core/src/test/scala/unit/kafka/controller/ControllerTestUtils.scala b/connect/runtime/src/test/resources/test-plugins/service-loader/test/plugins/ServiceLoadedClass.java
similarity index 56%
rename from core/src/test/scala/unit/kafka/controller/ControllerTestUtils.scala
rename to connect/runtime/src/test/resources/test-plugins/service-loader/test/plugins/ServiceLoadedClass.java
index 84b956df79b9b..98677ed43d65d 100644
--- a/core/src/test/scala/unit/kafka/controller/ControllerTestUtils.scala
+++ b/connect/runtime/src/test/resources/test-plugins/service-loader/test/plugins/ServiceLoadedClass.java
@@ -14,22 +14,35 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package kafka.controller
-
-import org.easymock.{EasyMock, IAnswer}
-
-object ControllerTestUtils {
-
-  /** Since ControllerEvent is sealed, return a subclass of ControllerEvent created with EasyMock */
-  def createMockControllerEvent(controllerState: ControllerState, process: () => Unit): ControllerEvent = {
-    val mockEvent: ControllerEvent = EasyMock.createNiceMock(classOf[ControllerEvent])
-    EasyMock.expect(mockEvent.state).andReturn(controllerState)
-    EasyMock.expect(mockEvent.process()).andAnswer(new IAnswer[Unit]() {
-      def answer(): Unit = {
-        process()
-      }
-    })
-    EasyMock.replay(mockEvent)
-    mockEvent
+
+package test.plugins;
+
+import org.apache.kafka.connect.runtime.isolation.SamplingTestPlugin;
+
+/**
+ * Superclass for service loaded classes
+ */
+public class ServiceLoadedClass extends SamplingTestPlugin {
+
+  private static final ClassLoader STATIC_CLASS_LOADER;
+  private final ClassLoader classloader;
+
+  static {
+    STATIC_CLASS_LOADER = Thread.currentThread().getContextClassLoader();
   }
+
+  {
+    classloader = Thread.currentThread().getContextClassLoader();
+  }
+
+  @Override
+  public ClassLoader staticClassloader() {
+    return STATIC_CLASS_LOADER;
+  }
+
+  @Override
+  public ClassLoader classloader() {
+    return classloader;
+  }
+
 }
diff --git a/connect/runtime/src/test/resources/test-plugins/service-loader/test/plugins/ServiceLoadedSubclass.java b/connect/runtime/src/test/resources/test-plugins/service-loader/test/plugins/ServiceLoadedSubclass.java
new file mode 100644
index 0000000000000..cfc6b6f9cfc94
--- /dev/null
+++ b/connect/runtime/src/test/resources/test-plugins/service-loader/test/plugins/ServiceLoadedSubclass.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 test.plugins;
+
+/**
+ * Instance of a service loaded class
+ */
+public class ServiceLoadedSubclass extends ServiceLoadedClass {
+
+  private static final ClassLoader STATIC_CLASS_LOADER;
+  private final ClassLoader classloader;
+
+  static {
+    STATIC_CLASS_LOADER = Thread.currentThread().getContextClassLoader();
+  }
+
+  {
+    classloader = Thread.currentThread().getContextClassLoader();
+  }
+
+  @Override
+  public ClassLoader staticClassloader() {
+    return STATIC_CLASS_LOADER;
+  }
+
+  @Override
+  public ClassLoader classloader() {
+    return classloader;
+  }
+
+}
diff --git a/connect/runtime/src/test/resources/test-plugins/service-loader/test/plugins/ServiceLoaderPlugin.java b/connect/runtime/src/test/resources/test-plugins/service-loader/test/plugins/ServiceLoaderPlugin.java
new file mode 100644
index 0000000000000..e6371baf56aad
--- /dev/null
+++ b/connect/runtime/src/test/resources/test-plugins/service-loader/test/plugins/ServiceLoaderPlugin.java
@@ -0,0 +1,85 @@
+/*
+ * 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 test.plugins;
+
+import java.util.Map;
+import java.util.HashMap;
+import java.util.ServiceLoader;
+import java.util.Iterator;
+import org.apache.kafka.connect.data.Schema;
+import org.apache.kafka.connect.data.SchemaAndValue;
+import org.apache.kafka.connect.storage.Converter;
+import org.apache.kafka.connect.runtime.isolation.SamplingTestPlugin;
+
+/**
+ * Samples data about its initialization environment for later analysis
+ */
+public class ServiceLoaderPlugin extends SamplingTestPlugin implements Converter {
+
+  private static final ClassLoader STATIC_CLASS_LOADER;
+  private static final Map SAMPLES;
+  private final ClassLoader classloader;
+
+  static {
+    STATIC_CLASS_LOADER = Thread.currentThread().getContextClassLoader();
+    SAMPLES = new HashMap<>();
+    Iterator it = ServiceLoader.load(ServiceLoadedClass.class).iterator();
+    while (it.hasNext()) {
+      ServiceLoadedClass loaded = it.next();
+      SAMPLES.put(loaded.getClass().getSimpleName() + ".static", loaded);
+    }
+  }
+
+  {
+    classloader = Thread.currentThread().getContextClassLoader();
+    Iterator it = ServiceLoader.load(ServiceLoadedClass.class).iterator();
+    while (it.hasNext()) {
+      ServiceLoadedClass loaded = it.next();
+      SAMPLES.put(loaded.getClass().getSimpleName() + ".dynamic", loaded);
+    }
+  }
+
+  @Override
+  public void configure(final Map configs, final boolean isKey) {
+  }
+
+  @Override
+  public byte[] fromConnectData(final String topic, final Schema schema, final Object value) {
+    return new byte[0];
+  }
+
+  @Override
+  public SchemaAndValue toConnectData(final String topic, final byte[] value) {
+    return null;
+  }
+
+  @Override
+  public ClassLoader staticClassloader() {
+    return STATIC_CLASS_LOADER;
+  }
+
+  @Override
+  public ClassLoader classloader() {
+    return classloader;
+  }
+
+  @Override
+  public Map otherSamples() {
+    return SAMPLES;
+  }
+}
diff --git a/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/Cast.java b/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/Cast.java
index 3dc6dc7535717..475b98e88b3a7 100644
--- a/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/Cast.java
+++ b/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/Cast.java
@@ -24,10 +24,14 @@
 import org.apache.kafka.common.config.ConfigException;
 import org.apache.kafka.connect.connector.ConnectRecord;
 import org.apache.kafka.connect.data.ConnectSchema;
+import org.apache.kafka.connect.data.Date;
 import org.apache.kafka.connect.data.Field;
 import org.apache.kafka.connect.data.Schema;
+import org.apache.kafka.connect.data.Schema.Type;
 import org.apache.kafka.connect.data.SchemaBuilder;
 import org.apache.kafka.connect.data.Struct;
+import org.apache.kafka.connect.data.Time;
+import org.apache.kafka.connect.data.Timestamp;
 import org.apache.kafka.connect.data.Values;
 import org.apache.kafka.connect.errors.DataException;
 import org.apache.kafka.connect.transforms.util.SchemaUtil;
@@ -222,7 +226,18 @@ private SchemaBuilder convertFieldType(Schema.Type type) {
             default:
                 throw new DataException("Unexpected type in Cast transformation: " + type);
         }
+    }
 
+    private static Object encodeLogicalType(Schema schema, Object value) {
+        switch (schema.name()) {
+            case Date.LOGICAL_NAME:
+                return Date.fromLogical(schema, (java.util.Date) value);
+            case Time.LOGICAL_NAME:
+                return Time.fromLogical(schema, (java.util.Date) value);
+            case Timestamp.LOGICAL_NAME:
+                return Timestamp.fromLogical(schema, (java.util.Date) value);
+        }
+        return value;
     }
 
     private static Object castValueToType(Schema schema, Object value, Schema.Type targetType) {
@@ -238,6 +253,11 @@ private static Object castValueToType(Schema schema, Object value, Schema.Type t
             // Ensure the type we are trying to cast from is supported
             validCastType(inferredType, FieldType.INPUT);
 
+            // Perform logical type encoding to their internal representation.
+            if (schema != null && schema.name() != null && targetType != Type.STRING) {
+                value = encodeLogicalType(schema, value);
+            }
+
             switch (targetType) {
                 case INT8:
                     return castToInt8(value);
diff --git a/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/ExtractField.java b/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/ExtractField.java
index eb8c357f3ff16..bd3cbd9288c26 100644
--- a/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/ExtractField.java
+++ b/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/ExtractField.java
@@ -18,6 +18,7 @@
 
 import org.apache.kafka.common.config.ConfigDef;
 import org.apache.kafka.connect.connector.ConnectRecord;
+import org.apache.kafka.connect.data.Field;
 import org.apache.kafka.connect.data.Schema;
 import org.apache.kafka.connect.data.Struct;
 import org.apache.kafka.connect.transforms.util.SimpleConfig;
@@ -58,7 +59,13 @@ public R apply(R record) {
             return newRecord(record, null, value == null ? null : value.get(fieldName));
         } else {
             final Struct value = requireStructOrNull(operatingValue(record), PURPOSE);
-            return newRecord(record, schema.field(fieldName).schema(), value == null ? null : value.get(fieldName));
+            Field field = schema.field(fieldName);
+
+            if (field == null) {
+                throw new IllegalArgumentException("Unknown field: " + fieldName);
+            }
+
+            return newRecord(record, field.schema(), value == null ? null : value.get(fieldName));
         }
     }
 
diff --git a/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/Flatten.java b/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/Flatten.java
index c5e4000975324..dbd8a6830390a 100644
--- a/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/Flatten.java
+++ b/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/Flatten.java
@@ -35,7 +35,7 @@
 import java.util.Map;
 
 import static org.apache.kafka.connect.transforms.util.Requirements.requireMap;
-import static org.apache.kafka.connect.transforms.util.Requirements.requireStruct;
+import static org.apache.kafka.connect.transforms.util.Requirements.requireStructOrNull;
 
 public abstract class Flatten> implements Transformation {
 
@@ -69,7 +69,9 @@ public void configure(Map props) {
 
     @Override
     public R apply(R record) {
-        if (operatingSchema(record) == null) {
+        if (operatingValue(record) == null) {
+            return record;
+        } else if (operatingSchema(record) == null) {
             return applySchemaless(record);
         } else {
             return applyWithSchema(record);
@@ -136,20 +138,24 @@ private void applySchemaless(Map originalRecord, String fieldNam
     }
 
     private R applyWithSchema(R record) {
-        final Struct value = requireStruct(operatingValue(record), PURPOSE);
+        final Struct value = requireStructOrNull(operatingValue(record), PURPOSE);
 
-        Schema updatedSchema = schemaUpdateCache.get(value.schema());
+        Schema schema = operatingSchema(record);
+        Schema updatedSchema = schemaUpdateCache.get(schema);
         if (updatedSchema == null) {
-            final SchemaBuilder builder = SchemaUtil.copySchemaBasics(value.schema(), SchemaBuilder.struct());
-            Struct defaultValue = (Struct) value.schema().defaultValue();
-            buildUpdatedSchema(value.schema(), "", builder, value.schema().isOptional(), defaultValue);
+            final SchemaBuilder builder = SchemaUtil.copySchemaBasics(schema, SchemaBuilder.struct());
+            Struct defaultValue = (Struct) schema.defaultValue();
+            buildUpdatedSchema(schema, "", builder, schema.isOptional(), defaultValue);
             updatedSchema = builder.build();
-            schemaUpdateCache.put(value.schema(), updatedSchema);
+            schemaUpdateCache.put(schema, updatedSchema);
+        }
+        if (value == null) {
+            return newRecord(record, updatedSchema, null);
+        } else {
+            final Struct updatedValue = new Struct(updatedSchema);
+            buildWithSchema(value, "", updatedValue);
+            return newRecord(record, updatedSchema, updatedValue);
         }
-
-        final Struct updatedValue = new Struct(updatedSchema);
-        buildWithSchema(value, "", updatedValue);
-        return newRecord(record, updatedSchema, updatedValue);
     }
 
     /**
@@ -190,7 +196,7 @@ private void buildUpdatedSchema(Schema schema, String fieldNamePrefix, SchemaBui
                     break;
                 default:
                     throw new DataException("Flatten transformation does not support " + field.schema().type()
-                            + " for record without schemas (for field " + fieldName + ").");
+                            + " for record with schemas (for field " + fieldName + ").");
             }
         }
     }
@@ -216,6 +222,9 @@ private Schema convertFieldSchema(Schema orig, boolean optional, Object defaultF
     }
 
     private void buildWithSchema(Struct record, String fieldNamePrefix, Struct newRecord) {
+        if (record == null) {
+            return;
+        }
         for (Field field : record.schema().fields()) {
             final String fieldName = fieldName(fieldNamePrefix, field.name());
             switch (field.schema().type()) {
@@ -235,7 +244,7 @@ private void buildWithSchema(Struct record, String fieldNamePrefix, Struct newRe
                     break;
                 default:
                     throw new DataException("Flatten transformation does not support " + field.schema().type()
-                            + " for record without schemas (for field " + fieldName + ").");
+                            + " for record with schemas (for field " + fieldName + ").");
             }
         }
     }
diff --git a/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/InsertField.java b/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/InsertField.java
index 5e472a907c3f4..8f0c5aa1f63fb 100644
--- a/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/InsertField.java
+++ b/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/InsertField.java
@@ -127,7 +127,9 @@ public void configure(Map props) {
 
     @Override
     public R apply(R record) {
-        if (operatingSchema(record) == null) {
+        if (operatingValue(record) == null) {
+            return record;
+        } else if (operatingSchema(record) == null) {
             return applySchemaless(record);
         } else {
             return applyWithSchema(record);
diff --git a/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/ReplaceField.java b/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/ReplaceField.java
index f071bdad3156a..3d9abc25bc731 100644
--- a/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/ReplaceField.java
+++ b/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/ReplaceField.java
@@ -123,7 +123,9 @@ String reverseRenamed(String fieldName) {
 
     @Override
     public R apply(R record) {
-        if (operatingSchema(record) == null) {
+        if (operatingValue(record) == null) {
+            return record;
+        } else if (operatingSchema(record) == null) {
             return applySchemaless(record);
         } else {
             return applyWithSchema(record);
diff --git a/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/TimestampConverter.java b/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/TimestampConverter.java
index 85574415f4559..f32253ec51b6d 100644
--- a/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/TimestampConverter.java
+++ b/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/TimestampConverter.java
@@ -47,7 +47,7 @@
 import java.util.TimeZone;
 
 import static org.apache.kafka.connect.transforms.util.Requirements.requireMap;
-import static org.apache.kafka.connect.transforms.util.Requirements.requireStruct;
+import static org.apache.kafka.connect.transforms.util.Requirements.requireStructOrNull;
 
 public abstract class TimestampConverter> implements Transformation {
 
@@ -85,6 +85,10 @@ public abstract class TimestampConverter> implements
 
     private static final TimeZone UTC = TimeZone.getTimeZone("UTC");
 
+    public static final Schema OPTIONAL_DATE_SCHEMA = org.apache.kafka.connect.data.Date.builder().optional().schema();
+    public static final Schema OPTIONAL_TIMESTAMP_SCHEMA = Timestamp.builder().optional().schema();
+    public static final Schema OPTIONAL_TIME_SCHEMA = Time.builder().optional().schema();
+
     private interface TimestampTranslator {
         /**
          * Convert from the type-specific format to the universal java.util.Date format
@@ -94,7 +98,7 @@ private interface TimestampTranslator {
         /**
          * Get the schema for this format.
          */
-        Schema typeSchema();
+        Schema typeSchema(boolean isOptional);
 
         /**
          * Convert from the universal java.util.Date format to the type-specific format
@@ -118,8 +122,8 @@ public Date toRaw(Config config, Object orig) {
             }
 
             @Override
-            public Schema typeSchema() {
-                return Schema.STRING_SCHEMA;
+            public Schema typeSchema(boolean isOptional) {
+                return isOptional ? Schema.OPTIONAL_STRING_SCHEMA : Schema.STRING_SCHEMA;
             }
 
             @Override
@@ -139,8 +143,8 @@ public Date toRaw(Config config, Object orig) {
             }
 
             @Override
-            public Schema typeSchema() {
-                return Schema.INT64_SCHEMA;
+            public Schema typeSchema(boolean isOptional) {
+                return isOptional ? Schema.OPTIONAL_INT64_SCHEMA : Schema.INT64_SCHEMA;
             }
 
             @Override
@@ -159,8 +163,8 @@ public Date toRaw(Config config, Object orig) {
             }
 
             @Override
-            public Schema typeSchema() {
-                return org.apache.kafka.connect.data.Date.SCHEMA;
+            public Schema typeSchema(boolean isOptional) {
+                return isOptional ? OPTIONAL_DATE_SCHEMA : org.apache.kafka.connect.data.Date.SCHEMA;
             }
 
             @Override
@@ -185,8 +189,8 @@ public Date toRaw(Config config, Object orig) {
             }
 
             @Override
-            public Schema typeSchema() {
-                return Time.SCHEMA;
+            public Schema typeSchema(boolean isOptional) {
+                return isOptional ? OPTIONAL_TIME_SCHEMA : Time.SCHEMA;
             }
 
             @Override
@@ -212,8 +216,8 @@ public Date toRaw(Config config, Object orig) {
             }
 
             @Override
-            public Schema typeSchema() {
-                return Timestamp.SCHEMA;
+            public Schema typeSchema(boolean isOptional) {
+                return isOptional ? OPTIONAL_TIMESTAMP_SCHEMA : Timestamp.SCHEMA;
             }
 
             @Override
@@ -330,16 +334,16 @@ private R applyWithSchema(R record) {
         if (config.field.isEmpty()) {
             Object value = operatingValue(record);
             // New schema is determined by the requested target timestamp type
-            Schema updatedSchema = TRANSLATORS.get(config.type).typeSchema();
+            Schema updatedSchema = TRANSLATORS.get(config.type).typeSchema(schema.isOptional());
             return newRecord(record, updatedSchema, convertTimestamp(value, timestampTypeFromSchema(schema)));
         } else {
-            final Struct value = requireStruct(operatingValue(record), PURPOSE);
-            Schema updatedSchema = schemaUpdateCache.get(value.schema());
+            final Struct value = requireStructOrNull(operatingValue(record), PURPOSE);
+            Schema updatedSchema = schemaUpdateCache.get(schema);
             if (updatedSchema == null) {
                 SchemaBuilder builder = SchemaUtil.copySchemaBasics(schema, SchemaBuilder.struct());
                 for (Field field : schema.fields()) {
                     if (field.name().equals(config.field)) {
-                        builder.field(field.name(), TRANSLATORS.get(config.type).typeSchema());
+                        builder.field(field.name(), TRANSLATORS.get(config.type).typeSchema(field.schema().isOptional()));
                     } else {
                         builder.field(field.name(), field.schema());
                     }
@@ -361,6 +365,9 @@ private R applyWithSchema(R record) {
     }
 
     private Struct applyValueWithSchema(Struct value, Schema updatedSchema) {
+        if (value == null) {
+            return null;
+        }
         Struct updatedValue = new Struct(updatedSchema);
         for (Field field : value.schema().fields()) {
             final Object updatedFieldValue;
@@ -375,11 +382,11 @@ private Struct applyValueWithSchema(Struct value, Schema updatedSchema) {
     }
 
     private R applySchemaless(R record) {
-        if (config.field.isEmpty()) {
-            Object value = operatingValue(record);
-            return newRecord(record, null, convertTimestamp(value));
+        Object rawValue = operatingValue(record);
+        if (rawValue == null || config.field.isEmpty()) {
+            return newRecord(record, null, convertTimestamp(rawValue));
         } else {
-            final Map value = requireMap(operatingValue(record), PURPOSE);
+            final Map value = requireMap(rawValue, PURPOSE);
             final HashMap updatedValue = new HashMap<>(value);
             updatedValue.put(config.field, convertTimestamp(value.get(config.field)));
             return newRecord(record, null, updatedValue);
@@ -424,11 +431,14 @@ private String inferTimestampType(Object timestamp) {
 
     /**
      * Convert the given timestamp to the target timestamp format.
-     * @param timestamp the input timestamp
+     * @param timestamp the input timestamp, may be null
      * @param timestampFormat the format of the timestamp, or null if the format should be inferred
      * @return the converted timestamp
      */
     private Object convertTimestamp(Object timestamp, String timestampFormat) {
+        if (timestamp == null) {
+            return null;
+        }
         if (timestampFormat == null) {
             timestampFormat = inferTimestampType(timestamp);
         }
diff --git a/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/ValueToKey.java b/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/ValueToKey.java
index a9d960103a0b8..b2226d38fca87 100644
--- a/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/ValueToKey.java
+++ b/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/ValueToKey.java
@@ -21,9 +21,11 @@
 import org.apache.kafka.common.cache.SynchronizedCache;
 import org.apache.kafka.common.config.ConfigDef;
 import org.apache.kafka.connect.connector.ConnectRecord;
+import org.apache.kafka.connect.data.Field;
 import org.apache.kafka.connect.data.Schema;
 import org.apache.kafka.connect.data.SchemaBuilder;
 import org.apache.kafka.connect.data.Struct;
+import org.apache.kafka.connect.errors.DataException;
 import org.apache.kafka.connect.transforms.util.NonEmptyListValidator;
 import org.apache.kafka.connect.transforms.util.SimpleConfig;
 
@@ -82,8 +84,11 @@ private R applyWithSchema(R record) {
         if (keySchema == null) {
             final SchemaBuilder keySchemaBuilder = SchemaBuilder.struct();
             for (String field : fields) {
-                final Schema fieldSchema = value.schema().field(field).schema();
-                keySchemaBuilder.field(field, fieldSchema);
+                final Field fieldFromValue = value.schema().field(field);
+                if (fieldFromValue == null) {
+                    throw new DataException("Field does not exist: " + field);
+                }
+                keySchemaBuilder.field(field, fieldFromValue.schema());
             }
             keySchema = keySchemaBuilder.build();
             valueToKeySchemaCache.put(value.schema(), keySchema);
diff --git a/connect/transforms/src/test/java/org/apache/kafka/connect/transforms/CastTest.java b/connect/transforms/src/test/java/org/apache/kafka/connect/transforms/CastTest.java
index c568afb9ff452..a28aa28c6d72e 100644
--- a/connect/transforms/src/test/java/org/apache/kafka/connect/transforms/CastTest.java
+++ b/connect/transforms/src/test/java/org/apache/kafka/connect/transforms/CastTest.java
@@ -17,11 +17,16 @@
 
 package org.apache.kafka.connect.transforms;
 
+import java.util.Arrays;
+import java.util.List;
+import java.util.concurrent.TimeUnit;
 import org.apache.kafka.common.config.ConfigException;
 import org.apache.kafka.connect.data.Decimal;
 import org.apache.kafka.connect.data.Schema;
+import org.apache.kafka.connect.data.Schema.Type;
 import org.apache.kafka.connect.data.SchemaBuilder;
 import org.apache.kafka.connect.data.Struct;
+import org.apache.kafka.connect.data.Time;
 import org.apache.kafka.connect.data.Timestamp;
 import org.apache.kafka.connect.data.Values;
 import org.apache.kafka.connect.errors.DataException;
@@ -42,7 +47,8 @@
 public class CastTest {
     private final Cast xformKey = new Cast.Key<>();
     private final Cast xformValue = new Cast.Value<>();
-    private static final long MILLIS_PER_DAY = 24 * 60 * 60 * 1000;
+    private static final long MILLIS_PER_HOUR = TimeUnit.HOURS.toMillis(1);
+    private static final long MILLIS_PER_DAY = TimeUnit.DAYS.toMillis(1);
 
     @After
     public void teardown() {
@@ -320,6 +326,97 @@ public void castWholeRecordValueSchemalessUnsupportedType() {
         xformValue.apply(new SourceRecord(null, null, "topic", 0, null, Collections.singletonList("foo")));
     }
 
+    @Test
+    public void castLogicalToPrimitive() {
+        List specParts = Arrays.asList(
+            "date_to_int32:int32",  // Cast to underlying representation
+            "timestamp_to_int64:int64",  // Cast to underlying representation
+            "time_to_int64:int64",  // Cast to wider datatype than underlying representation
+            "decimal_to_int32:int32",  // Cast to narrower datatype with data loss
+            "timestamp_to_float64:float64",  // loss of precision casting to double
+            "null_timestamp_to_int32:int32"
+        );
+
+        Date day = new Date(MILLIS_PER_DAY);
+        xformValue.configure(Collections.singletonMap(Cast.SPEC_CONFIG,
+            String.join(",", specParts)));
+
+        SchemaBuilder builder = SchemaBuilder.struct();
+        builder.field("date_to_int32", org.apache.kafka.connect.data.Date.SCHEMA);
+        builder.field("timestamp_to_int64", Timestamp.SCHEMA);
+        builder.field("time_to_int64", Time.SCHEMA);
+        builder.field("decimal_to_int32", Decimal.schema(new BigDecimal((long) Integer.MAX_VALUE + 1).scale()));
+        builder.field("timestamp_to_float64", Timestamp.SCHEMA);
+        builder.field("null_timestamp_to_int32", Timestamp.builder().optional().build());
+
+        Schema supportedTypesSchema = builder.build();
+
+        Struct recordValue = new Struct(supportedTypesSchema);
+        recordValue.put("date_to_int32", day);
+        recordValue.put("timestamp_to_int64", new Date(0));
+        recordValue.put("time_to_int64", new Date(1));
+        recordValue.put("decimal_to_int32", new BigDecimal((long) Integer.MAX_VALUE + 1));
+        recordValue.put("timestamp_to_float64", new Date(Long.MAX_VALUE));
+        recordValue.put("null_timestamp_to_int32", null);
+
+        SourceRecord transformed = xformValue.apply(
+            new SourceRecord(null, null, "topic", 0,
+                supportedTypesSchema, recordValue));
+
+        assertEquals(1, ((Struct) transformed.value()).get("date_to_int32"));
+        assertEquals(0L, ((Struct) transformed.value()).get("timestamp_to_int64"));
+        assertEquals(1L, ((Struct) transformed.value()).get("time_to_int64"));
+        assertEquals(Integer.MIN_VALUE, ((Struct) transformed.value()).get("decimal_to_int32"));
+        assertEquals(9.223372036854776E18, ((Struct) transformed.value()).get("timestamp_to_float64"));
+        assertNull(((Struct) transformed.value()).get("null_timestamp_to_int32"));
+
+        Schema transformedSchema = ((Struct) transformed.value()).schema();
+        assertEquals(Type.INT32, transformedSchema.field("date_to_int32").schema().type());
+        assertEquals(Type.INT64, transformedSchema.field("timestamp_to_int64").schema().type());
+        assertEquals(Type.INT64, transformedSchema.field("time_to_int64").schema().type());
+        assertEquals(Type.INT32, transformedSchema.field("decimal_to_int32").schema().type());
+        assertEquals(Type.FLOAT64, transformedSchema.field("timestamp_to_float64").schema().type());
+        assertEquals(Type.INT32, transformedSchema.field("null_timestamp_to_int32").schema().type());
+    }
+
+    @Test
+    public void castLogicalToString() {
+        Date date = new Date(MILLIS_PER_DAY);
+        Date time = new Date(MILLIS_PER_HOUR);
+        Date timestamp = new Date();
+
+        xformValue.configure(Collections.singletonMap(Cast.SPEC_CONFIG,
+            "date:string,decimal:string,time:string,timestamp:string"));
+
+        SchemaBuilder builder = SchemaBuilder.struct();
+        builder.field("date", org.apache.kafka.connect.data.Date.SCHEMA);
+        builder.field("decimal", Decimal.schema(new BigDecimal(1982).scale()));
+        builder.field("time", Time.SCHEMA);
+        builder.field("timestamp", Timestamp.SCHEMA);
+
+        Schema supportedTypesSchema = builder.build();
+
+        Struct recordValue = new Struct(supportedTypesSchema);
+        recordValue.put("date", date);
+        recordValue.put("decimal", new BigDecimal(1982));
+        recordValue.put("time", time);
+        recordValue.put("timestamp", timestamp);
+
+        SourceRecord transformed = xformValue.apply(
+            new SourceRecord(null, null, "topic", 0,
+                supportedTypesSchema, recordValue));
+
+        assertEquals(Values.dateFormatFor(date).format(date), ((Struct) transformed.value()).get("date"));
+        assertEquals("1982", ((Struct) transformed.value()).get("decimal"));
+        assertEquals(Values.dateFormatFor(time).format(time), ((Struct) transformed.value()).get("time"));
+        assertEquals(Values.dateFormatFor(timestamp).format(timestamp), ((Struct) transformed.value()).get("timestamp"));
+
+        Schema transformedSchema = ((Struct) transformed.value()).schema();
+        assertEquals(Type.STRING, transformedSchema.field("date").schema().type());
+        assertEquals(Type.STRING, transformedSchema.field("decimal").schema().type());
+        assertEquals(Type.STRING, transformedSchema.field("time").schema().type());
+        assertEquals(Type.STRING, transformedSchema.field("timestamp").schema().type());
+    }
 
     @Test
     public void castFieldsWithSchema() {
@@ -338,7 +435,7 @@ public void castFieldsWithSchema() {
         builder.field("boolean", Schema.BOOLEAN_SCHEMA);
         builder.field("string", Schema.STRING_SCHEMA);
         builder.field("bigdecimal", Decimal.schema(new BigDecimal(42).scale()));
-        builder.field("date", Timestamp.SCHEMA);
+        builder.field("date", org.apache.kafka.connect.data.Date.SCHEMA);
         builder.field("optional", Schema.OPTIONAL_FLOAT32_SCHEMA);
         builder.field("timestamp", Timestamp.SCHEMA);
         Schema supportedTypesSchema = builder.build();
diff --git a/connect/transforms/src/test/java/org/apache/kafka/connect/transforms/ExtractFieldTest.java b/connect/transforms/src/test/java/org/apache/kafka/connect/transforms/ExtractFieldTest.java
index acb0beb0049ec..a78c6d5804b20 100644
--- a/connect/transforms/src/test/java/org/apache/kafka/connect/transforms/ExtractFieldTest.java
+++ b/connect/transforms/src/test/java/org/apache/kafka/connect/transforms/ExtractFieldTest.java
@@ -28,6 +28,7 @@
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNull;
+import static org.junit.Assert.fail;
 
 public class ExtractFieldTest {
     private final ExtractField xform = new ExtractField.Key<>();
@@ -86,4 +87,30 @@ public void testNullWithSchema() {
         assertNull(transformedRecord.key());
     }
 
+    @Test
+    public void nonExistentFieldSchemalessShouldReturnNull() {
+        xform.configure(Collections.singletonMap("field", "nonexistent"));
+
+        final SinkRecord record = new SinkRecord("test", 0, null, Collections.singletonMap("magic", 42), null, null, 0);
+        final SinkRecord transformedRecord = xform.apply(record);
+
+        assertNull(transformedRecord.keySchema());
+        assertNull(transformedRecord.key());
+    }
+
+    @Test
+    public void nonExistentFieldWithSchemaShouldFail() {
+        xform.configure(Collections.singletonMap("field", "nonexistent"));
+
+        final Schema keySchema = SchemaBuilder.struct().field("magic", Schema.INT32_SCHEMA).build();
+        final Struct key = new Struct(keySchema).put("magic", 42);
+        final SinkRecord record = new SinkRecord("test", 0, keySchema, key, null, null, 0);
+
+        try {
+            xform.apply(record);
+            fail("Expected exception wasn't raised");
+        } catch (IllegalArgumentException iae) {
+            assertEquals("Unknown field: nonexistent", iae.getMessage());
+        }
+    }
 }
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 b0549fbfc9d4d..d044338f5e25d 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
@@ -182,6 +182,46 @@ public void testOptionalFieldStruct() {
         assertNull(transformedStruct.get("B.opt_int32"));
     }
 
+    @Test
+    public void testOptionalStruct() {
+        xformValue.configure(Collections.emptyMap());
+
+        SchemaBuilder builder = SchemaBuilder.struct().optional();
+        builder.field("opt_int32", Schema.OPTIONAL_INT32_SCHEMA);
+        Schema schema = builder.build();
+
+        SourceRecord transformed = xformValue.apply(new SourceRecord(null, null,
+            "topic", 0,
+            schema, null));
+
+        assertEquals(Schema.Type.STRUCT, transformed.valueSchema().type());
+        assertNull(transformed.value());
+    }
+
+    @Test
+    public void testOptionalNestedStruct() {
+        xformValue.configure(Collections.emptyMap());
+
+        SchemaBuilder builder = SchemaBuilder.struct().optional();
+        builder.field("opt_int32", Schema.OPTIONAL_INT32_SCHEMA);
+        Schema supportedTypesSchema = builder.build();
+
+        builder = SchemaBuilder.struct();
+        builder.field("B", supportedTypesSchema);
+        Schema oneLevelNestedSchema = builder.build();
+
+        Struct oneLevelNestedStruct = new Struct(oneLevelNestedSchema);
+        oneLevelNestedStruct.put("B", null);
+
+        SourceRecord transformed = xformValue.apply(new SourceRecord(null, null,
+            "topic", 0,
+            oneLevelNestedSchema, oneLevelNestedStruct));
+
+        assertEquals(Schema.Type.STRUCT, transformed.valueSchema().type());
+        Struct transformedStruct = (Struct) transformed.value();
+        assertNull(transformedStruct.get("B.opt_int32"));
+    }
+
     @Test
     public void testOptionalFieldMap() {
         xformValue.configure(Collections.emptyMap());
@@ -257,4 +297,29 @@ public void testOptionalAndDefaultValuesNested() {
         Schema transformedOptFieldSchema = SchemaBuilder.string().optional().defaultValue("child_default").build();
         assertEquals(transformedOptFieldSchema, transformedSchema.field("opt_field").schema());
     }
+
+    @Test
+    public void tombstoneEventWithoutSchemaShouldPassThrough() {
+        xformValue.configure(Collections.emptyMap());
+
+        final SourceRecord record = new SourceRecord(null, null, "test", 0,
+                null, null);
+        final SourceRecord transformedRecord = xformValue.apply(record);
+
+        assertEquals(null, transformedRecord.value());
+        assertEquals(null, transformedRecord.valueSchema());
+    }
+
+    @Test
+    public void tombstoneEventWithSchemaShouldPassThrough() {
+        xformValue.configure(Collections.emptyMap());
+
+        final Schema simpleStructSchema = SchemaBuilder.struct().name("name").version(1).doc("doc").field("magic", Schema.OPTIONAL_INT64_SCHEMA).build();
+        final SourceRecord record = new SourceRecord(null, null, "test", 0,
+                simpleStructSchema, null);
+        final SourceRecord transformedRecord = xformValue.apply(record);
+
+        assertEquals(null, transformedRecord.value());
+        assertEquals(simpleStructSchema, transformedRecord.valueSchema());
+    }
 }
diff --git a/connect/transforms/src/test/java/org/apache/kafka/connect/transforms/InsertFieldTest.java b/connect/transforms/src/test/java/org/apache/kafka/connect/transforms/InsertFieldTest.java
index a0a09752a4ae9..dc6611c21c096 100644
--- a/connect/transforms/src/test/java/org/apache/kafka/connect/transforms/InsertFieldTest.java
+++ b/connect/transforms/src/test/java/org/apache/kafka/connect/transforms/InsertFieldTest.java
@@ -33,17 +33,18 @@
 import static org.junit.Assert.assertSame;
 
 public class InsertFieldTest {
-    private InsertField xform = new InsertField.Value<>();
+    private InsertField xformKey = new InsertField.Key<>();
+    private InsertField xformValue = new InsertField.Value<>();
 
     @After
     public void teardown() {
-        xform.close();
+        xformValue.close();
     }
 
     @Test(expected = DataException.class)
     public void topLevelStructRequired() {
-        xform.configure(Collections.singletonMap("topic.field", "topic_field"));
-        xform.apply(new SourceRecord(null, null, "", 0, Schema.INT32_SCHEMA, 42));
+        xformValue.configure(Collections.singletonMap("topic.field", "topic_field"));
+        xformValue.apply(new SourceRecord(null, null, "", 0, Schema.INT32_SCHEMA, 42));
     }
 
     @Test
@@ -55,13 +56,13 @@ public void copySchemaAndInsertConfiguredFields() {
         props.put("static.field", "instance_id");
         props.put("static.value", "my-instance-id");
 
-        xform.configure(props);
+        xformValue.configure(props);
 
         final Schema simpleStructSchema = SchemaBuilder.struct().name("name").version(1).doc("doc").field("magic", Schema.OPTIONAL_INT64_SCHEMA).build();
         final Struct simpleStruct = new Struct(simpleStructSchema).put("magic", 42L);
 
         final SourceRecord record = new SourceRecord(null, null, "test", 0, simpleStructSchema, simpleStruct);
-        final SourceRecord transformedRecord = xform.apply(record);
+        final SourceRecord transformedRecord = xformValue.apply(record);
 
         assertEquals(simpleStructSchema.name(), transformedRecord.valueSchema().name());
         assertEquals(simpleStructSchema.version(), transformedRecord.valueSchema().version());
@@ -83,7 +84,7 @@ public void copySchemaAndInsertConfiguredFields() {
         assertEquals("my-instance-id", ((Struct) transformedRecord.value()).getString("instance_id"));
 
         // Exercise caching
-        final SourceRecord transformedRecord2 = xform.apply(
+        final SourceRecord transformedRecord2 = xformValue.apply(
                 new SourceRecord(null, null, "test", 1, simpleStructSchema, new Struct(simpleStructSchema)));
         assertSame(transformedRecord.valueSchema(), transformedRecord2.valueSchema());
     }
@@ -97,18 +98,103 @@ public void schemalessInsertConfiguredFields() {
         props.put("static.field", "instance_id");
         props.put("static.value", "my-instance-id");
 
-        xform.configure(props);
+        xformValue.configure(props);
 
         final SourceRecord record = new SourceRecord(null, null, "test", 0,
                 null, Collections.singletonMap("magic", 42L));
 
-        final SourceRecord transformedRecord = xform.apply(record);
+        final SourceRecord transformedRecord = xformValue.apply(record);
 
-        assertEquals(42L, ((Map) transformedRecord.value()).get("magic"));
-        assertEquals("test", ((Map) transformedRecord.value()).get("topic_field"));
-        assertEquals(0, ((Map) transformedRecord.value()).get("partition_field"));
-        assertEquals(null, ((Map) transformedRecord.value()).get("timestamp_field"));
-        assertEquals("my-instance-id", ((Map) transformedRecord.value()).get("instance_id"));
+        assertEquals(42L, ((Map) transformedRecord.value()).get("magic"));
+        assertEquals("test", ((Map) transformedRecord.value()).get("topic_field"));
+        assertEquals(0, ((Map) transformedRecord.value()).get("partition_field"));
+        assertEquals(null, ((Map) transformedRecord.value()).get("timestamp_field"));
+        assertEquals("my-instance-id", ((Map) transformedRecord.value()).get("instance_id"));
     }
 
+
+    @Test
+    public void insertConfiguredFieldsIntoTombstoneEventWithoutSchemaLeavesValueUnchanged() {
+        final Map props = new HashMap<>();
+        props.put("topic.field", "topic_field!");
+        props.put("partition.field", "partition_field");
+        props.put("timestamp.field", "timestamp_field?");
+        props.put("static.field", "instance_id");
+        props.put("static.value", "my-instance-id");
+
+        xformValue.configure(props);
+
+        final SourceRecord record = new SourceRecord(null, null, "test", 0,
+                null, null);
+
+        final SourceRecord transformedRecord = xformValue.apply(record);
+
+        assertEquals(null, transformedRecord.value());
+        assertEquals(null, transformedRecord.valueSchema());
+    }
+
+    @Test
+    public void insertConfiguredFieldsIntoTombstoneEventWithSchemaLeavesValueUnchanged() {
+        final Map props = new HashMap<>();
+        props.put("topic.field", "topic_field!");
+        props.put("partition.field", "partition_field");
+        props.put("timestamp.field", "timestamp_field?");
+        props.put("static.field", "instance_id");
+        props.put("static.value", "my-instance-id");
+
+        xformValue.configure(props);
+
+        final Schema simpleStructSchema = SchemaBuilder.struct().name("name").version(1).doc("doc").field("magic", Schema.OPTIONAL_INT64_SCHEMA).build();
+
+        final SourceRecord record = new SourceRecord(null, null, "test", 0,
+                simpleStructSchema, null);
+
+        final SourceRecord transformedRecord = xformValue.apply(record);
+
+        assertEquals(null, transformedRecord.value());
+        assertEquals(simpleStructSchema, transformedRecord.valueSchema());
+    }
+
+    @Test
+    public void insertKeyFieldsIntoTombstoneEvent() {
+        final Map props = new HashMap<>();
+        props.put("topic.field", "topic_field!");
+        props.put("partition.field", "partition_field");
+        props.put("timestamp.field", "timestamp_field?");
+        props.put("static.field", "instance_id");
+        props.put("static.value", "my-instance-id");
+
+        xformKey.configure(props);
+
+        final SourceRecord record = new SourceRecord(null, null, "test", 0,
+            null, Collections.singletonMap("magic", 42L), null, null);
+
+        final SourceRecord transformedRecord = xformKey.apply(record);
+
+        assertEquals(42L, ((Map) transformedRecord.key()).get("magic"));
+        assertEquals("test", ((Map) transformedRecord.key()).get("topic_field"));
+        assertEquals(0, ((Map) transformedRecord.key()).get("partition_field"));
+        assertEquals(null, ((Map) transformedRecord.key()).get("timestamp_field"));
+        assertEquals("my-instance-id", ((Map) transformedRecord.key()).get("instance_id"));
+        assertEquals(null, transformedRecord.value());
+    }
+
+    @Test
+    public void insertIntoNullKeyLeavesRecordUnchanged() {
+        final Map props = new HashMap<>();
+        props.put("topic.field", "topic_field!");
+        props.put("partition.field", "partition_field");
+        props.put("timestamp.field", "timestamp_field?");
+        props.put("static.field", "instance_id");
+        props.put("static.value", "my-instance-id");
+
+        xformKey.configure(props);
+
+        final SourceRecord record = new SourceRecord(null, null, "test", 0,
+            null, null, null, Collections.singletonMap("magic", 42L));
+
+        final SourceRecord transformedRecord = xformKey.apply(record);
+
+        assertSame(record, transformedRecord);
+    }
 }
diff --git a/connect/transforms/src/test/java/org/apache/kafka/connect/transforms/ReplaceFieldTest.java b/connect/transforms/src/test/java/org/apache/kafka/connect/transforms/ReplaceFieldTest.java
index 6a1a13ada9dde..7ab00ed8da8c3 100644
--- a/connect/transforms/src/test/java/org/apache/kafka/connect/transforms/ReplaceFieldTest.java
+++ b/connect/transforms/src/test/java/org/apache/kafka/connect/transforms/ReplaceFieldTest.java
@@ -27,6 +27,7 @@
 import java.util.Map;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
 
 public class ReplaceFieldTest {
     private ReplaceField xform = new ReplaceField.Value<>();
@@ -36,6 +37,43 @@ public void teardown() {
         xform.close();
     }
 
+    @Test
+    public void tombstoneSchemaless() {
+        final Map props = new HashMap<>();
+        props.put("whitelist", "abc,foo");
+        props.put("renames", "abc:xyz,foo:bar");
+
+        xform.configure(props);
+
+        final SinkRecord record = new SinkRecord("test", 0, null, null, null, null, 0);
+        final SinkRecord transformedRecord = xform.apply(record);
+
+        assertNull(transformedRecord.value());
+        assertNull(transformedRecord.valueSchema());
+    }
+
+    @Test
+    public void tombstoneWithSchema() {
+        final Map props = new HashMap<>();
+        props.put("whitelist", "abc,foo");
+        props.put("renames", "abc:xyz,foo:bar");
+
+        xform.configure(props);
+
+        final Schema schema = SchemaBuilder.struct()
+            .field("dont", Schema.STRING_SCHEMA)
+            .field("abc", Schema.INT32_SCHEMA)
+            .field("foo", Schema.BOOLEAN_SCHEMA)
+            .field("etc", Schema.STRING_SCHEMA)
+            .build();
+
+        final SinkRecord record = new SinkRecord("test", 0, null, null, schema, null, 0);
+        final SinkRecord transformedRecord = xform.apply(record);
+
+        assertNull(transformedRecord.value());
+        assertEquals(schema, transformedRecord.valueSchema());
+    }
+
     @Test
     public void schemaless() {
         final Map props = new HashMap<>();
diff --git a/connect/transforms/src/test/java/org/apache/kafka/connect/transforms/TimestampConverterTest.java b/connect/transforms/src/test/java/org/apache/kafka/connect/transforms/TimestampConverterTest.java
index 475066f74e281..3a1920ed528ce 100644
--- a/connect/transforms/src/test/java/org/apache/kafka/connect/transforms/TimestampConverterTest.java
+++ b/connect/transforms/src/test/java/org/apache/kafka/connect/transforms/TimestampConverterTest.java
@@ -35,6 +35,7 @@
 import java.util.Map;
 import java.util.TimeZone;
 
+import static org.apache.kafka.connect.transforms.util.Requirements.requireStruct;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNull;
 
@@ -105,13 +106,12 @@ public void testConfigInvalidFormat() {
         xformValue.configure(config);
     }
 
-
     // Conversions without schemas (most flexible Timestamp -> other types)
 
     @Test
     public void testSchemalessIdentity() {
         xformValue.configure(Collections.singletonMap(TimestampConverter.TARGET_TYPE_CONFIG, "Timestamp"));
-        SourceRecord transformed = xformValue.apply(new SourceRecord(null, null, "topic", 0, null, DATE_PLUS_TIME.getTime()));
+        SourceRecord transformed = xformValue.apply(createRecordSchemaless(DATE_PLUS_TIME.getTime()));
 
         assertNull(transformed.valueSchema());
         assertEquals(DATE_PLUS_TIME.getTime(), transformed.value());
@@ -120,7 +120,7 @@ public void testSchemalessIdentity() {
     @Test
     public void testSchemalessTimestampToDate() {
         xformValue.configure(Collections.singletonMap(TimestampConverter.TARGET_TYPE_CONFIG, "Date"));
-        SourceRecord transformed = xformValue.apply(new SourceRecord(null, null, "topic", 0, null, DATE_PLUS_TIME.getTime()));
+        SourceRecord transformed = xformValue.apply(createRecordSchemaless(DATE_PLUS_TIME.getTime()));
 
         assertNull(transformed.valueSchema());
         assertEquals(DATE.getTime(), transformed.value());
@@ -129,7 +129,7 @@ public void testSchemalessTimestampToDate() {
     @Test
     public void testSchemalessTimestampToTime() {
         xformValue.configure(Collections.singletonMap(TimestampConverter.TARGET_TYPE_CONFIG, "Time"));
-        SourceRecord transformed = xformValue.apply(new SourceRecord(null, null, "topic", 0, null, DATE_PLUS_TIME.getTime()));
+        SourceRecord transformed = xformValue.apply(createRecordSchemaless(DATE_PLUS_TIME.getTime()));
 
         assertNull(transformed.valueSchema());
         assertEquals(TIME.getTime(), transformed.value());
@@ -138,7 +138,7 @@ public void testSchemalessTimestampToTime() {
     @Test
     public void testSchemalessTimestampToUnix() {
         xformValue.configure(Collections.singletonMap(TimestampConverter.TARGET_TYPE_CONFIG, "unix"));
-        SourceRecord transformed = xformValue.apply(new SourceRecord(null, null, "topic", 0, null, DATE_PLUS_TIME.getTime()));
+        SourceRecord transformed = xformValue.apply(createRecordSchemaless(DATE_PLUS_TIME.getTime()));
 
         assertNull(transformed.valueSchema());
         assertEquals(DATE_PLUS_TIME_UNIX, transformed.value());
@@ -150,7 +150,7 @@ public void testSchemalessTimestampToString() {
         config.put(TimestampConverter.TARGET_TYPE_CONFIG, "string");
         config.put(TimestampConverter.FORMAT_CONFIG, STRING_DATE_FMT);
         xformValue.configure(config);
-        SourceRecord transformed = xformValue.apply(new SourceRecord(null, null, "topic", 0, null, DATE_PLUS_TIME.getTime()));
+        SourceRecord transformed = xformValue.apply(createRecordSchemaless(DATE_PLUS_TIME.getTime()));
 
         assertNull(transformed.valueSchema());
         assertEquals(DATE_PLUS_TIME_STRING, transformed.value());
@@ -162,7 +162,7 @@ public void testSchemalessTimestampToString() {
     @Test
     public void testSchemalessDateToTimestamp() {
         xformValue.configure(Collections.singletonMap(TimestampConverter.TARGET_TYPE_CONFIG, "Timestamp"));
-        SourceRecord transformed = xformValue.apply(new SourceRecord(null, null, "topic", 0, null, DATE.getTime()));
+        SourceRecord transformed = xformValue.apply(createRecordSchemaless(DATE.getTime()));
 
         assertNull(transformed.valueSchema());
         // No change expected since the source type is coarser-grained
@@ -172,7 +172,7 @@ public void testSchemalessDateToTimestamp() {
     @Test
     public void testSchemalessTimeToTimestamp() {
         xformValue.configure(Collections.singletonMap(TimestampConverter.TARGET_TYPE_CONFIG, "Timestamp"));
-        SourceRecord transformed = xformValue.apply(new SourceRecord(null, null, "topic", 0, null, TIME.getTime()));
+        SourceRecord transformed = xformValue.apply(createRecordSchemaless(TIME.getTime()));
 
         assertNull(transformed.valueSchema());
         // No change expected since the source type is coarser-grained
@@ -182,7 +182,7 @@ public void testSchemalessTimeToTimestamp() {
     @Test
     public void testSchemalessUnixToTimestamp() {
         xformValue.configure(Collections.singletonMap(TimestampConverter.TARGET_TYPE_CONFIG, "Timestamp"));
-        SourceRecord transformed = xformValue.apply(new SourceRecord(null, null, "topic", 0, null, DATE_PLUS_TIME_UNIX));
+        SourceRecord transformed = xformValue.apply(createRecordSchemaless(DATE_PLUS_TIME_UNIX));
 
         assertNull(transformed.valueSchema());
         assertEquals(DATE_PLUS_TIME.getTime(), transformed.value());
@@ -194,7 +194,7 @@ public void testSchemalessStringToTimestamp() {
         config.put(TimestampConverter.TARGET_TYPE_CONFIG, "Timestamp");
         config.put(TimestampConverter.FORMAT_CONFIG, STRING_DATE_FMT);
         xformValue.configure(config);
-        SourceRecord transformed = xformValue.apply(new SourceRecord(null, null, "topic", 0, null, DATE_PLUS_TIME_STRING));
+        SourceRecord transformed = xformValue.apply(createRecordSchemaless(DATE_PLUS_TIME_STRING));
 
         assertNull(transformed.valueSchema());
         assertEquals(DATE_PLUS_TIME.getTime(), transformed.value());
@@ -206,7 +206,7 @@ public void testSchemalessStringToTimestamp() {
     @Test
     public void testWithSchemaIdentity() {
         xformValue.configure(Collections.singletonMap(TimestampConverter.TARGET_TYPE_CONFIG, "Timestamp"));
-        SourceRecord transformed = xformValue.apply(new SourceRecord(null, null, "topic", 0, Timestamp.SCHEMA, DATE_PLUS_TIME.getTime()));
+        SourceRecord transformed = xformValue.apply(createRecordWithSchema(Timestamp.SCHEMA, DATE_PLUS_TIME.getTime()));
 
         assertEquals(Timestamp.SCHEMA, transformed.valueSchema());
         assertEquals(DATE_PLUS_TIME.getTime(), transformed.value());
@@ -215,7 +215,7 @@ public void testWithSchemaIdentity() {
     @Test
     public void testWithSchemaTimestampToDate() {
         xformValue.configure(Collections.singletonMap(TimestampConverter.TARGET_TYPE_CONFIG, "Date"));
-        SourceRecord transformed = xformValue.apply(new SourceRecord(null, null, "topic", 0, Timestamp.SCHEMA, DATE_PLUS_TIME.getTime()));
+        SourceRecord transformed = xformValue.apply(createRecordWithSchema(Timestamp.SCHEMA, DATE_PLUS_TIME.getTime()));
 
         assertEquals(Date.SCHEMA, transformed.valueSchema());
         assertEquals(DATE.getTime(), transformed.value());
@@ -224,7 +224,7 @@ public void testWithSchemaTimestampToDate() {
     @Test
     public void testWithSchemaTimestampToTime() {
         xformValue.configure(Collections.singletonMap(TimestampConverter.TARGET_TYPE_CONFIG, "Time"));
-        SourceRecord transformed = xformValue.apply(new SourceRecord(null, null, "topic", 0, Timestamp.SCHEMA, DATE_PLUS_TIME.getTime()));
+        SourceRecord transformed = xformValue.apply(createRecordWithSchema(Timestamp.SCHEMA, DATE_PLUS_TIME.getTime()));
 
         assertEquals(Time.SCHEMA, transformed.valueSchema());
         assertEquals(TIME.getTime(), transformed.value());
@@ -233,7 +233,7 @@ public void testWithSchemaTimestampToTime() {
     @Test
     public void testWithSchemaTimestampToUnix() {
         xformValue.configure(Collections.singletonMap(TimestampConverter.TARGET_TYPE_CONFIG, "unix"));
-        SourceRecord transformed = xformValue.apply(new SourceRecord(null, null, "topic", 0, Timestamp.SCHEMA, DATE_PLUS_TIME.getTime()));
+        SourceRecord transformed = xformValue.apply(createRecordWithSchema(Timestamp.SCHEMA, DATE_PLUS_TIME.getTime()));
 
         assertEquals(Schema.INT64_SCHEMA, transformed.valueSchema());
         assertEquals(DATE_PLUS_TIME_UNIX, transformed.value());
@@ -245,19 +245,70 @@ public void testWithSchemaTimestampToString() {
         config.put(TimestampConverter.TARGET_TYPE_CONFIG, "string");
         config.put(TimestampConverter.FORMAT_CONFIG, STRING_DATE_FMT);
         xformValue.configure(config);
-        SourceRecord transformed = xformValue.apply(new SourceRecord(null, null, "topic", 0, Timestamp.SCHEMA, DATE_PLUS_TIME.getTime()));
+        SourceRecord transformed = xformValue.apply(createRecordWithSchema(Timestamp.SCHEMA, DATE_PLUS_TIME.getTime()));
 
         assertEquals(Schema.STRING_SCHEMA, transformed.valueSchema());
         assertEquals(DATE_PLUS_TIME_STRING, transformed.value());
     }
 
+    // Null-value conversions schemaless
+
+    @Test
+    public void testSchemalessNullValueToString() {
+        testSchemalessNullValueConversion("string");
+        testSchemalessNullFieldConversion("string");
+    }
+    @Test
+    public void testSchemalessNullValueToDate() {
+        testSchemalessNullValueConversion("Date");
+        testSchemalessNullFieldConversion("Date");
+    }
+    @Test
+    public void testSchemalessNullValueToTimestamp() {
+        testSchemalessNullValueConversion("Timestamp");
+        testSchemalessNullFieldConversion("Timestamp");
+    }
+    @Test
+    public void testSchemalessNullValueToUnix() {
+        testSchemalessNullValueConversion("unix");
+        testSchemalessNullFieldConversion("unix");
+    }
+
+    @Test
+    public void testSchemalessNullValueToTime() {
+        testSchemalessNullValueConversion("Time");
+        testSchemalessNullFieldConversion("Time");
+    }
+
+    private void testSchemalessNullValueConversion(String targetType) {
+        Map config = new HashMap<>();
+        config.put(TimestampConverter.TARGET_TYPE_CONFIG, targetType);
+        config.put(TimestampConverter.FORMAT_CONFIG, STRING_DATE_FMT);
+        xformValue.configure(config);
+        SourceRecord transformed = xformValue.apply(createRecordSchemaless(null));
+
+        assertNull(transformed.valueSchema());
+        assertNull(transformed.value());
+    }
+
+    private void testSchemalessNullFieldConversion(String targetType) {
+        Map config = new HashMap<>();
+        config.put(TimestampConverter.TARGET_TYPE_CONFIG, targetType);
+        config.put(TimestampConverter.FORMAT_CONFIG, STRING_DATE_FMT);
+        config.put(TimestampConverter.FIELD_CONFIG, "ts");
+        xformValue.configure(config);
+        SourceRecord transformed = xformValue.apply(createRecordSchemaless(null));
+
+        assertNull(transformed.valueSchema());
+        assertNull(transformed.value());
+    }
 
     // Conversions with schemas (core types -> most flexible Timestamp format)
 
     @Test
     public void testWithSchemaDateToTimestamp() {
         xformValue.configure(Collections.singletonMap(TimestampConverter.TARGET_TYPE_CONFIG, "Timestamp"));
-        SourceRecord transformed = xformValue.apply(new SourceRecord(null, null, "topic", 0, Date.SCHEMA, DATE.getTime()));
+        SourceRecord transformed = xformValue.apply(createRecordWithSchema(Date.SCHEMA, DATE.getTime()));
 
         assertEquals(Timestamp.SCHEMA, transformed.valueSchema());
         // No change expected since the source type is coarser-grained
@@ -267,7 +318,7 @@ public void testWithSchemaDateToTimestamp() {
     @Test
     public void testWithSchemaTimeToTimestamp() {
         xformValue.configure(Collections.singletonMap(TimestampConverter.TARGET_TYPE_CONFIG, "Timestamp"));
-        SourceRecord transformed = xformValue.apply(new SourceRecord(null, null, "topic", 0, Time.SCHEMA, TIME.getTime()));
+        SourceRecord transformed = xformValue.apply(createRecordWithSchema(Time.SCHEMA, TIME.getTime()));
 
         assertEquals(Timestamp.SCHEMA, transformed.valueSchema());
         // No change expected since the source type is coarser-grained
@@ -277,7 +328,7 @@ public void testWithSchemaTimeToTimestamp() {
     @Test
     public void testWithSchemaUnixToTimestamp() {
         xformValue.configure(Collections.singletonMap(TimestampConverter.TARGET_TYPE_CONFIG, "Timestamp"));
-        SourceRecord transformed = xformValue.apply(new SourceRecord(null, null, "topic", 0, Schema.INT64_SCHEMA, DATE_PLUS_TIME_UNIX));
+        SourceRecord transformed = xformValue.apply(createRecordWithSchema(Schema.INT64_SCHEMA, DATE_PLUS_TIME_UNIX));
 
         assertEquals(Timestamp.SCHEMA, transformed.valueSchema());
         assertEquals(DATE_PLUS_TIME.getTime(), transformed.value());
@@ -289,12 +340,145 @@ public void testWithSchemaStringToTimestamp() {
         config.put(TimestampConverter.TARGET_TYPE_CONFIG, "Timestamp");
         config.put(TimestampConverter.FORMAT_CONFIG, STRING_DATE_FMT);
         xformValue.configure(config);
-        SourceRecord transformed = xformValue.apply(new SourceRecord(null, null, "topic", 0, Schema.STRING_SCHEMA, DATE_PLUS_TIME_STRING));
+        SourceRecord transformed = xformValue.apply(createRecordWithSchema(Schema.STRING_SCHEMA, DATE_PLUS_TIME_STRING));
 
         assertEquals(Timestamp.SCHEMA, transformed.valueSchema());
         assertEquals(DATE_PLUS_TIME.getTime(), transformed.value());
     }
 
+    // Null-value conversions with schema
+
+    @Test
+    public void testWithSchemaNullValueToTimestamp() {
+        testWithSchemaNullValueConversion("Timestamp", Schema.OPTIONAL_INT64_SCHEMA, TimestampConverter.OPTIONAL_TIMESTAMP_SCHEMA);
+        testWithSchemaNullValueConversion("Timestamp", TimestampConverter.OPTIONAL_TIME_SCHEMA, TimestampConverter.OPTIONAL_TIMESTAMP_SCHEMA);
+        testWithSchemaNullValueConversion("Timestamp", TimestampConverter.OPTIONAL_DATE_SCHEMA, TimestampConverter.OPTIONAL_TIMESTAMP_SCHEMA);
+        testWithSchemaNullValueConversion("Timestamp", Schema.OPTIONAL_STRING_SCHEMA, TimestampConverter.OPTIONAL_TIMESTAMP_SCHEMA);
+        testWithSchemaNullValueConversion("Timestamp", TimestampConverter.OPTIONAL_TIMESTAMP_SCHEMA, TimestampConverter.OPTIONAL_TIMESTAMP_SCHEMA);
+    }
+
+    @Test
+    public void testWithSchemaNullFieldToTimestamp() {
+        testWithSchemaNullFieldConversion("Timestamp", Schema.OPTIONAL_INT64_SCHEMA, TimestampConverter.OPTIONAL_TIMESTAMP_SCHEMA);
+        testWithSchemaNullFieldConversion("Timestamp", TimestampConverter.OPTIONAL_TIME_SCHEMA, TimestampConverter.OPTIONAL_TIMESTAMP_SCHEMA);
+        testWithSchemaNullFieldConversion("Timestamp", TimestampConverter.OPTIONAL_DATE_SCHEMA, TimestampConverter.OPTIONAL_TIMESTAMP_SCHEMA);
+        testWithSchemaNullFieldConversion("Timestamp", Schema.OPTIONAL_STRING_SCHEMA, TimestampConverter.OPTIONAL_TIMESTAMP_SCHEMA);
+        testWithSchemaNullFieldConversion("Timestamp", TimestampConverter.OPTIONAL_TIMESTAMP_SCHEMA, TimestampConverter.OPTIONAL_TIMESTAMP_SCHEMA);
+    }
+
+    @Test
+    public void testWithSchemaNullValueToUnix() {
+        testWithSchemaNullValueConversion("unix", Schema.OPTIONAL_INT64_SCHEMA, Schema.OPTIONAL_INT64_SCHEMA);
+        testWithSchemaNullValueConversion("unix", TimestampConverter.OPTIONAL_TIME_SCHEMA, Schema.OPTIONAL_INT64_SCHEMA);
+        testWithSchemaNullValueConversion("unix", TimestampConverter.OPTIONAL_DATE_SCHEMA, Schema.OPTIONAL_INT64_SCHEMA);
+        testWithSchemaNullValueConversion("unix", Schema.OPTIONAL_STRING_SCHEMA, Schema.OPTIONAL_INT64_SCHEMA);
+        testWithSchemaNullValueConversion("unix", TimestampConverter.OPTIONAL_TIMESTAMP_SCHEMA, Schema.OPTIONAL_INT64_SCHEMA);
+    }
+
+    @Test
+    public void testWithSchemaNullFieldToUnix() {
+        testWithSchemaNullFieldConversion("unix", Schema.OPTIONAL_INT64_SCHEMA, Schema.OPTIONAL_INT64_SCHEMA);
+        testWithSchemaNullFieldConversion("unix", TimestampConverter.OPTIONAL_TIME_SCHEMA, Schema.OPTIONAL_INT64_SCHEMA);
+        testWithSchemaNullFieldConversion("unix", TimestampConverter.OPTIONAL_DATE_SCHEMA, Schema.OPTIONAL_INT64_SCHEMA);
+        testWithSchemaNullFieldConversion("unix", Schema.OPTIONAL_STRING_SCHEMA, Schema.OPTIONAL_INT64_SCHEMA);
+        testWithSchemaNullFieldConversion("unix", TimestampConverter.OPTIONAL_TIMESTAMP_SCHEMA, Schema.OPTIONAL_INT64_SCHEMA);
+    }
+
+    @Test
+    public void testWithSchemaNullValueToTime() {
+        testWithSchemaNullValueConversion("Time", Schema.OPTIONAL_INT64_SCHEMA, TimestampConverter.OPTIONAL_TIME_SCHEMA);
+        testWithSchemaNullValueConversion("Time", TimestampConverter.OPTIONAL_TIME_SCHEMA, TimestampConverter.OPTIONAL_TIME_SCHEMA);
+        testWithSchemaNullValueConversion("Time", TimestampConverter.OPTIONAL_DATE_SCHEMA, TimestampConverter.OPTIONAL_TIME_SCHEMA);
+        testWithSchemaNullValueConversion("Time", Schema.OPTIONAL_STRING_SCHEMA, TimestampConverter.OPTIONAL_TIME_SCHEMA);
+        testWithSchemaNullValueConversion("Time", TimestampConverter.OPTIONAL_TIMESTAMP_SCHEMA, TimestampConverter.OPTIONAL_TIME_SCHEMA);
+    }
+
+    @Test
+    public void testWithSchemaNullFieldToTime() {
+        testWithSchemaNullFieldConversion("Time", Schema.OPTIONAL_INT64_SCHEMA, TimestampConverter.OPTIONAL_TIME_SCHEMA);
+        testWithSchemaNullFieldConversion("Time", TimestampConverter.OPTIONAL_TIME_SCHEMA, TimestampConverter.OPTIONAL_TIME_SCHEMA);
+        testWithSchemaNullFieldConversion("Time", TimestampConverter.OPTIONAL_DATE_SCHEMA, TimestampConverter.OPTIONAL_TIME_SCHEMA);
+        testWithSchemaNullFieldConversion("Time", Schema.OPTIONAL_STRING_SCHEMA, TimestampConverter.OPTIONAL_TIME_SCHEMA);
+        testWithSchemaNullFieldConversion("Time", TimestampConverter.OPTIONAL_TIMESTAMP_SCHEMA, TimestampConverter.OPTIONAL_TIME_SCHEMA);
+    }
+
+    @Test
+    public void testWithSchemaNullValueToDate() {
+        testWithSchemaNullValueConversion("Date", Schema.OPTIONAL_INT64_SCHEMA, TimestampConverter.OPTIONAL_DATE_SCHEMA);
+        testWithSchemaNullValueConversion("Date", TimestampConverter.OPTIONAL_TIME_SCHEMA, TimestampConverter.OPTIONAL_DATE_SCHEMA);
+        testWithSchemaNullValueConversion("Date", TimestampConverter.OPTIONAL_DATE_SCHEMA, TimestampConverter.OPTIONAL_DATE_SCHEMA);
+        testWithSchemaNullValueConversion("Date", Schema.OPTIONAL_STRING_SCHEMA, TimestampConverter.OPTIONAL_DATE_SCHEMA);
+        testWithSchemaNullValueConversion("Date", TimestampConverter.OPTIONAL_TIMESTAMP_SCHEMA, TimestampConverter.OPTIONAL_DATE_SCHEMA);
+    }
+
+    @Test
+    public void testWithSchemaNullFieldToDate() {
+        testWithSchemaNullFieldConversion("Date", Schema.OPTIONAL_INT64_SCHEMA, TimestampConverter.OPTIONAL_DATE_SCHEMA);
+        testWithSchemaNullFieldConversion("Date", TimestampConverter.OPTIONAL_TIME_SCHEMA, TimestampConverter.OPTIONAL_DATE_SCHEMA);
+        testWithSchemaNullFieldConversion("Date", TimestampConverter.OPTIONAL_DATE_SCHEMA, TimestampConverter.OPTIONAL_DATE_SCHEMA);
+        testWithSchemaNullFieldConversion("Date", Schema.OPTIONAL_STRING_SCHEMA, TimestampConverter.OPTIONAL_DATE_SCHEMA);
+        testWithSchemaNullFieldConversion("Date", TimestampConverter.OPTIONAL_TIMESTAMP_SCHEMA, TimestampConverter.OPTIONAL_DATE_SCHEMA);
+    }
+
+    @Test
+    public void testWithSchemaNullValueToString() {
+        testWithSchemaNullValueConversion("string", Schema.OPTIONAL_INT64_SCHEMA, Schema.OPTIONAL_STRING_SCHEMA);
+        testWithSchemaNullValueConversion("string", TimestampConverter.OPTIONAL_TIME_SCHEMA, Schema.OPTIONAL_STRING_SCHEMA);
+        testWithSchemaNullValueConversion("string", TimestampConverter.OPTIONAL_DATE_SCHEMA, Schema.OPTIONAL_STRING_SCHEMA);
+        testWithSchemaNullValueConversion("string", Schema.OPTIONAL_STRING_SCHEMA, Schema.OPTIONAL_STRING_SCHEMA);
+        testWithSchemaNullValueConversion("string", TimestampConverter.OPTIONAL_TIMESTAMP_SCHEMA, Schema.OPTIONAL_STRING_SCHEMA);
+    }
+
+    @Test
+    public void testWithSchemaNullFieldToString() {
+        testWithSchemaNullFieldConversion("string", Schema.OPTIONAL_INT64_SCHEMA, Schema.OPTIONAL_STRING_SCHEMA);
+        testWithSchemaNullFieldConversion("string", TimestampConverter.OPTIONAL_TIME_SCHEMA, Schema.OPTIONAL_STRING_SCHEMA);
+        testWithSchemaNullFieldConversion("string", TimestampConverter.OPTIONAL_DATE_SCHEMA, Schema.OPTIONAL_STRING_SCHEMA);
+        testWithSchemaNullFieldConversion("string", Schema.OPTIONAL_STRING_SCHEMA, Schema.OPTIONAL_STRING_SCHEMA);
+        testWithSchemaNullFieldConversion("string", TimestampConverter.OPTIONAL_TIMESTAMP_SCHEMA, Schema.OPTIONAL_STRING_SCHEMA);
+    }
+
+    private void testWithSchemaNullValueConversion(String targetType, Schema originalSchema, Schema expectedSchema) {
+        Map config = new HashMap<>();
+        config.put(TimestampConverter.TARGET_TYPE_CONFIG, targetType);
+        config.put(TimestampConverter.FORMAT_CONFIG, STRING_DATE_FMT);
+        xformValue.configure(config);
+        SourceRecord transformed = xformValue.apply(createRecordWithSchema(originalSchema, null));
+
+        assertEquals(expectedSchema, transformed.valueSchema());
+        assertNull(transformed.value());
+    }
+
+    private void testWithSchemaNullFieldConversion(String targetType, Schema originalSchema, Schema expectedSchema) {
+        Map config = new HashMap<>();
+        config.put(TimestampConverter.TARGET_TYPE_CONFIG, targetType);
+        config.put(TimestampConverter.FORMAT_CONFIG, STRING_DATE_FMT);
+        config.put(TimestampConverter.FIELD_CONFIG, "ts");
+        xformValue.configure(config);
+        SchemaBuilder structSchema = SchemaBuilder.struct()
+                .field("ts", originalSchema)
+                .field("other", Schema.STRING_SCHEMA);
+
+        SchemaBuilder expectedStructSchema = SchemaBuilder.struct()
+                .field("ts", expectedSchema)
+                .field("other", Schema.STRING_SCHEMA);
+
+        Struct original = new Struct(structSchema);
+        original.put("ts", null);
+        original.put("other", "test");
+
+        // Struct field is null
+        SourceRecord transformed = xformValue.apply(createRecordWithSchema(structSchema.build(), original));
+
+        assertEquals(expectedStructSchema.build(), transformed.valueSchema());
+        assertNull(requireStruct(transformed.value(), "").get("ts"));
+
+        // entire Struct is null
+        transformed = xformValue.apply(createRecordWithSchema(structSchema.optional().build(), null));
+
+        assertEquals(expectedStructSchema.optional().build(), transformed.valueSchema());
+        assertNull(transformed.value());
+    }
 
     // Convert field instead of entire key/value
 
@@ -306,7 +490,7 @@ public void testSchemalessFieldConversion() {
         xformValue.configure(config);
 
         Object value = Collections.singletonMap("ts", DATE_PLUS_TIME.getTime());
-        SourceRecord transformed = xformValue.apply(new SourceRecord(null, null, "topic", 0, null, value));
+        SourceRecord transformed = xformValue.apply(createRecordSchemaless(value));
 
         assertNull(transformed.valueSchema());
         assertEquals(Collections.singletonMap("ts", DATE.getTime()), transformed.value());
@@ -328,7 +512,7 @@ public void testWithSchemaFieldConversion() {
         original.put("ts", DATE_PLUS_TIME_UNIX);
         original.put("other", "test");
 
-        SourceRecord transformed = xformValue.apply(new SourceRecord(null, null, "topic", 0, structWithTimestampFieldSchema, original));
+        SourceRecord transformed = xformValue.apply(createRecordWithSchema(structWithTimestampFieldSchema, original));
 
         Schema expectedSchema = SchemaBuilder.struct()
                 .field("ts", Timestamp.SCHEMA)
@@ -351,4 +535,11 @@ public void testKey() {
         assertEquals(DATE_PLUS_TIME.getTime(), transformed.key());
     }
 
+    private SourceRecord createRecordWithSchema(Schema schema, Object value) {
+        return new SourceRecord(null, null, "topic", 0, schema, value);
+    }
+
+    private SourceRecord createRecordSchemaless(Object value) {
+        return createRecordWithSchema(null, value);
+    }
 }
diff --git a/connect/transforms/src/test/java/org/apache/kafka/connect/transforms/ValueToKeyTest.java b/connect/transforms/src/test/java/org/apache/kafka/connect/transforms/ValueToKeyTest.java
index e2dfa1773b96f..5854658aa29b3 100644
--- a/connect/transforms/src/test/java/org/apache/kafka/connect/transforms/ValueToKeyTest.java
+++ b/connect/transforms/src/test/java/org/apache/kafka/connect/transforms/ValueToKeyTest.java
@@ -19,6 +19,7 @@
 import org.apache.kafka.connect.data.Schema;
 import org.apache.kafka.connect.data.SchemaBuilder;
 import org.apache.kafka.connect.data.Struct;
+import org.apache.kafka.connect.errors.DataException;
 import org.apache.kafka.connect.sink.SinkRecord;
 import org.junit.After;
 import org.junit.Test;
@@ -28,6 +29,7 @@
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertThrows;
 
 public class ValueToKeyTest {
     private final ValueToKey xform = new ValueToKey<>();
@@ -88,4 +90,20 @@ public void withSchema() {
         assertEquals(expectedKey, transformedRecord.key());
     }
 
+    @Test
+    public void nonExistingField() {
+        xform.configure(Collections.singletonMap("fields", "not_exist"));
+
+        final Schema valueSchema = SchemaBuilder.struct()
+            .field("a", Schema.INT32_SCHEMA)
+            .build();
+
+        final Struct value = new Struct(valueSchema);
+        value.put("a", 1);
+
+        final SinkRecord record = new SinkRecord("", 0, null, null, valueSchema, value, 0);
+
+        DataException actual = assertThrows(DataException.class, () -> xform.apply(record));
+        assertEquals("Field does not exist: not_exist", actual.getMessage());
+    }
 }
diff --git a/core/src/main/scala/kafka/Kafka.scala b/core/src/main/scala/kafka/Kafka.scala
index c47aa529c172e..5950b7145181b 100755
--- a/core/src/main/scala/kafka/Kafka.scala
+++ b/core/src/main/scala/kafka/Kafka.scala
@@ -34,11 +34,20 @@ object Kafka extends Logging {
     val overrideOpt = optionParser.accepts("override", "Optional property that should override values set in server.properties file")
       .withRequiredArg()
       .ofType(classOf[String])
+    // This is just to make the parameter show up in the help output, we are not actually using this due the
+    // fact that this class ignores the first parameter which is interpreted as positional and mandatory
+    // but would not be mandatory if --version is specified
+    // This is a bit of an ugly crutch till we get a chance to rework the entire command line parsing
+    optionParser.accepts("version", "Print version information and exit.")
 
-    if (args.length == 0) {
+    if (args.length == 0 || args.contains("--help")) {
       CommandLineUtils.printUsageAndDie(optionParser, "USAGE: java [options] %s server.properties [--override property=value]*".format(classOf[KafkaServer].getSimpleName()))
     }
 
+    if (args.contains("--version")) {
+      CommandLineUtils.printVersionAndDie()
+    }
+
     val props = Utils.loadProps(args(0))
 
     if (args.length > 1) {
diff --git a/core/src/main/scala/kafka/admin/AclCommand.scala b/core/src/main/scala/kafka/admin/AclCommand.scala
index f945b254b1d22..0ed0703a95698 100644
--- a/core/src/main/scala/kafka/admin/AclCommand.scala
+++ b/core/src/main/scala/kafka/admin/AclCommand.scala
@@ -22,16 +22,21 @@ import java.util.Properties
 import joptsimple._
 import joptsimple.util.EnumConverter
 import kafka.security.auth._
+import kafka.security.authorizer.AuthorizerUtils
 import kafka.server.KafkaConfig
 import kafka.utils._
-import org.apache.kafka.clients.admin.{AdminClientConfig, AdminClient => JAdminClient}
+import org.apache.kafka.clients.admin.{Admin, AdminClientConfig, AdminClient => JAdminClient}
 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.common.resource.{PatternType, ResourcePattern, ResourcePatternFilter, Resource => JResource, ResourceType => JResourceType}
 import org.apache.kafka.common.security.JaasUtils
 import org.apache.kafka.common.security.auth.KafkaPrincipal
-import org.apache.kafka.common.utils.{SecurityUtils, Utils}
+import org.apache.kafka.common.utils.{Utils, SecurityUtils => JSecurityUtils}
+import org.apache.kafka.server.authorizer.{Authorizer => JAuthorizer}
 
 import scala.collection.JavaConverters._
+import scala.compat.java8.OptionConverters._
 import scala.collection.mutable
 import scala.io.StdIn
 
@@ -41,7 +46,7 @@ object AclCommand extends Logging {
 
   private val Newline = scala.util.Properties.lineSeparator
 
-  def main(args: Array[String]) {
+  def main(args: Array[String]): Unit = {
 
     val opts = new AclCommandOptions(args)
 
@@ -53,7 +58,18 @@ object AclCommand extends Logging {
       if (opts.options.has(opts.bootstrapServerOpt)) {
         new AdminClientService(opts)
       } else {
-        new AuthorizerService(opts)
+        val authorizerClass = if (opts.options.has(opts.authorizerOpt)) {
+          val className = opts.options.valueOf(opts.authorizerOpt)
+          Class.forName(className, true, Utils.getContextOrKafkaClassLoader)
+        } else
+          classOf[SimpleAclAuthorizer]
+
+        if (classOf[JAuthorizer].isAssignableFrom(authorizerClass))
+          new JAuthorizerService(authorizerClass.asSubclass(classOf[JAuthorizer]), opts)
+        else if (classOf[Authorizer].isAssignableFrom(authorizerClass))
+          new AuthorizerService(authorizerClass.asSubclass(classOf[Authorizer]), opts)
+        else
+          throw new IllegalArgumentException(s"Authorizer $authorizerClass does not implement ${classOf[Authorizer]} or ${classOf[JAuthorizer]}.")
       }
     }
 
@@ -80,7 +96,7 @@ object AclCommand extends Logging {
 
   class AdminClientService(val opts: AclCommandOptions) extends AclCommandService with Logging {
 
-    private def withAdminClient(opts: AclCommandOptions)(f: JAdminClient => Unit) {
+    private def withAdminClient(opts: AclCommandOptions)(f: Admin => Unit): Unit = {
       val props = if (opts.options.has(opts.commandConfigOpt))
         Utils.loadProps(opts.options.valueOf(opts.commandConfigOpt))
       else
@@ -99,9 +115,8 @@ object AclCommand extends Logging {
       val resourceToAcl = getResourceToAcls(opts)
       withAdminClient(opts) { adminClient =>
         for ((resource, acls) <- resourceToAcl) {
-          val resourcePattern = resource.toPattern
-          println(s"Adding ACLs for resource `$resourcePattern`: $Newline ${acls.map("\t" + _).mkString(Newline)} $Newline")
-          val aclBindings = acls.map(acl => new AclBinding(resourcePattern, getAccessControlEntry(acl))).asJavaCollection
+          println(s"Adding ACLs for resource `$resource`: $Newline ${acls.map("\t" + _).mkString(Newline)} $Newline")
+          val aclBindings = acls.map(acl => new AclBinding(resource, acl)).asJavaCollection
           adminClient.createAcls(aclBindings).all().get()
         }
 
@@ -149,24 +164,16 @@ object AclCommand extends Logging {
       }
     }
 
-    private def getAccessControlEntry(acl: Acl): AccessControlEntry = {
-      new AccessControlEntry(acl.principal.toString, acl.host, acl.operation.toJava, acl.permissionType.toJava)
-    }
-
-    private def removeAcls(adminClient: JAdminClient, acls: Set[Acl], filter: ResourcePatternFilter): Unit = {
+    private def removeAcls(adminClient: Admin, acls: Set[AccessControlEntry], filter: ResourcePatternFilter): Unit = {
       if (acls.isEmpty)
         adminClient.deleteAcls(List(new AclBindingFilter(filter, AccessControlEntryFilter.ANY)).asJava).all().get()
       else {
-        val aclBindingFilters = acls.map(acl => new AclBindingFilter(filter, getAccessControlEntryFilter(acl))).toList.asJava
+        val aclBindingFilters = acls.map(acl => new AclBindingFilter(filter, acl.toFilter)).toList.asJava
         adminClient.deleteAcls(aclBindingFilters).all().get()
       }
     }
 
-    private def getAccessControlEntryFilter(acl: Acl): AccessControlEntryFilter = {
-      new AccessControlEntryFilter(acl.principal.toString, acl.host, acl.operation.toJava, acl.permissionType.toJava)
-    }
-
-    private def getAcls(adminClient: JAdminClient, filters: Set[ResourcePatternFilter]): Map[ResourcePattern, Set[AccessControlEntry]] = {
+    private def getAcls(adminClient: Admin, filters: Set[ResourcePatternFilter]): Map[ResourcePattern, Set[AccessControlEntry]] = {
       val aclBindings =
         if (filters.isEmpty) adminClient.describeAcls(AclBindingFilter.ANY).values().get().asScala.toList
         else {
@@ -183,9 +190,10 @@ object AclCommand extends Logging {
     }
   }
 
-  class AuthorizerService(val opts: AclCommandOptions) extends AclCommandService with Logging {
+  @deprecated("Use JAuthorizerService", "Since 2.4")
+  class AuthorizerService(val authorizerClass: Class[_ <: Authorizer], val opts: AclCommandOptions) extends AclCommandService with Logging {
 
-    private def withAuthorizer()(f: Authorizer => Unit) {
+    private def withAuthorizer()(f: Authorizer => Unit): Unit = {
       val defaultProps = Map(KafkaConfig.ZkEnableSecureAclsProp -> JaasUtils.isZkSecurityEnabled)
       val authorizerProperties =
         if (opts.options.has(opts.authorizerPropertiesOpt)) {
@@ -195,12 +203,7 @@ object AclCommand extends Logging {
           defaultProps
         }
 
-      val authorizerClass = if (opts.options.has(opts.authorizerOpt))
-        opts.options.valueOf(opts.authorizerOpt)
-      else
-        classOf[SimpleAclAuthorizer].getName
-
-      val authZ = CoreUtils.createObject[Authorizer](authorizerClass)
+      val authZ = Utils.newInstance(authorizerClass)
       try {
         authZ.configure(authorizerProperties.asJava)
         f(authZ)
@@ -213,7 +216,7 @@ object AclCommand extends Logging {
       withAuthorizer() { authorizer =>
         for ((resource, acls) <- resourceToAcl) {
           println(s"Adding ACLs for resource `$resource`: $Newline ${acls.map("\t" + _).mkString(Newline)} $Newline")
-          authorizer.addAcls(acls, resource)
+          authorizer.addAcls(acls.map(AuthorizerUtils.convertToAcl), AuthorizerUtils.convertToResource(resource))
         }
 
         listAcls()
@@ -258,12 +261,12 @@ object AclCommand extends Logging {
       }
     }
 
-    private def removeAcls(authorizer: Authorizer, acls: Set[Acl], filter: ResourcePatternFilter) {
+    private def removeAcls(authorizer: Authorizer, acls: Set[AccessControlEntry], filter: ResourcePatternFilter): Unit = {
       getAcls(authorizer, filter)
         .keys
         .foreach(resource =>
           if (acls.isEmpty) authorizer.removeAcls(resource)
-          else authorizer.removeAcls(acls, resource)
+          else authorizer.removeAcls(acls.map(AuthorizerUtils.convertToAcl), resource)
         )
     }
 
@@ -280,20 +283,136 @@ object AclCommand extends Logging {
     private def getAcls(authorizer: Authorizer, filter: ResourcePatternFilter,
                         listPrincipal: Option[KafkaPrincipal] = None): Map[Resource, Set[Acl]] =
       if (listPrincipal.isEmpty)
-        authorizer.getAcls().filter { case (resource, acl) => filter.matches(resource.toPattern) }
+        authorizer.getAcls().filter { case (resource, _) => filter.matches(resource.toPattern) }
       else
-        authorizer.getAcls(listPrincipal.get).filter { case (resource, acl) => filter.matches(resource.toPattern) }
+        authorizer.getAcls(listPrincipal.get).filter { case (resource, _) => filter.matches(resource.toPattern) }
+
+  }
+
+  class JAuthorizerService(val authorizerClass: Class[_ <: JAuthorizer], val opts: AclCommandOptions) extends AclCommandService with Logging {
+
+    private def withAuthorizer()(f: JAuthorizer => Unit) {
+      val defaultProps = Map(KafkaConfig.ZkEnableSecureAclsProp -> JaasUtils.isZkSecurityEnabled)
+      val authorizerProperties =
+        if (opts.options.has(opts.authorizerPropertiesOpt)) {
+          val authorizerProperties = opts.options.valuesOf(opts.authorizerPropertiesOpt).asScala
+          defaultProps ++ CommandLineUtils.parseKeyValueArgs(authorizerProperties, acceptMissingValue = false).asScala
+        } else {
+          defaultProps
+        }
+
+      val authZ = Utils.newInstance(authorizerClass)
+      try {
+        authZ.configure(authorizerProperties.asJava)
+        f(authZ)
+      }
+      finally CoreUtils.swallow(authZ.close(), this)
+    }
+
+    def addAcls(): Unit = {
+      val resourceToAcl = getResourceToAcls(opts)
+      withAuthorizer() { authorizer =>
+        for ((resource, acls) <- resourceToAcl) {
+          println(s"Adding ACLs for resource `$resource`: $Newline ${acls.map("\t" + _).mkString(Newline)} $Newline")
+          val aclBindings = acls.map(acl => new AclBinding(resource, acl))
+          authorizer.createAcls(null,aclBindings.toList.asJava).asScala.map(_.toCompletableFuture.get).foreach { result =>
+            result.exception.asScala.foreach { exception =>
+              println(s"Error while adding ACLs: ${exception.getMessage}")
+              println(Utils.stackTrace(exception))
+            }
+          }
+        }
+
+        listAcls()
+      }
+    }
+
+    def removeAcls(): Unit = {
+      withAuthorizer() { authorizer =>
+        val filterToAcl = getResourceFilterToAcls(opts)
+
+        for ((filter, acls) <- filterToAcl) {
+          if (acls.isEmpty) {
+            if (confirmAction(opts, s"Are you sure you want to delete all ACLs for resource filter `$filter`? (y/n)"))
+              removeAcls(authorizer, acls, filter)
+          } else {
+            if (confirmAction(opts, s"Are you sure you want to remove ACLs: $Newline ${acls.map("\t" + _).mkString(Newline)} $Newline from resource filter `$filter`? (y/n)"))
+              removeAcls(authorizer, acls, filter)
+          }
+        }
+
+        listAcls()
+      }
+    }
+
+    def listAcls(): Unit = {
+      withAuthorizer() { authorizer =>
+        val filters = getResourceFilter(opts, dieIfNoResourceFound = false)
+        val listPrincipals = getPrincipals(opts, opts.listPrincipalsOpt)
+        val resourceToAcls = getAcls(authorizer, filters)
 
+        if (listPrincipals.isEmpty) {
+          for ((resource, acls) <- resourceToAcls)
+            println(s"Current ACLs for resource `$resource`: $Newline ${acls.map("\t" + _).mkString(Newline)} $Newline")
+        } else {
+          listPrincipals.foreach(principal => {
+            println(s"ACLs for principal `$principal`")
+            val filteredResourceToAcls =  resourceToAcls.mapValues(acls =>
+              acls.filter(acl => principal.toString.equals(acl.principal))).filter(entry => entry._2.nonEmpty)
+
+            for ((resource, acls) <- filteredResourceToAcls)
+              println(s"Current ACLs for resource `$resource`: $Newline ${acls.map("\t" + _).mkString(Newline)} $Newline")
+          })
+        }
+      }
+    }
+
+    private def removeAcls(authorizer: JAuthorizer, acls: Set[AccessControlEntry], filter: ResourcePatternFilter): Unit = {
+      val result = if (acls.isEmpty)
+        authorizer.deleteAcls(null, List(new AclBindingFilter(filter, AccessControlEntryFilter.ANY)).asJava)
+      else {
+        val aclBindingFilters = acls.map(acl => new AclBindingFilter(filter, acl.toFilter)).toList.asJava
+        authorizer.deleteAcls(null, aclBindingFilters)
+      }
+      result.asScala.map(_.toCompletableFuture.get).foreach { result =>
+        result.exception.asScala.foreach { exception =>
+          println(s"Error while removing ACLs: ${exception.getMessage}")
+          println(Utils.stackTrace(exception))
+        }
+        result.aclBindingDeleteResults.asScala.foreach { deleteResult =>
+          deleteResult.exception.asScala.foreach { exception =>
+            println(s"Error while removing ACLs: ${exception.getMessage}")
+            println(Utils.stackTrace(exception))
+          }
+        }
+      }
+    }
+
+    private def getAcls(authorizer: JAuthorizer, filters: Set[ResourcePatternFilter]): Map[ResourcePattern, Set[AccessControlEntry]] = {
+      val aclBindings =
+        if (filters.isEmpty) authorizer.acls(AclBindingFilter.ANY).asScala
+        else {
+          val results = for (filter <- filters) yield {
+            authorizer.acls(new AclBindingFilter(filter, AccessControlEntryFilter.ANY)).asScala
+          }
+          results.reduceLeft(_ ++ _)
+        }
+
+      val resourceToAcls = mutable.Map[ResourcePattern, Set[AccessControlEntry]]().withDefaultValue(Set())
+
+      aclBindings.foreach(aclBinding => resourceToAcls(aclBinding.pattern()) = resourceToAcls(aclBinding.pattern()) + aclBinding.entry())
+      resourceToAcls.toMap
+    }
   }
 
-  private def getResourceToAcls(opts: AclCommandOptions): Map[Resource, Set[Acl]] = {
+  private def getResourceToAcls(opts: AclCommandOptions): Map[ResourcePattern, Set[AccessControlEntry]] = {
     val patternType: PatternType = opts.options.valueOf(opts.resourcePatternType)
     if (!patternType.isSpecific)
       CommandLineUtils.printUsageAndDie(opts.parser, s"A '--resource-pattern-type' value of '$patternType' is not valid when adding acls.")
 
     val resourceToAcl = getResourceFilterToAcls(opts).map {
       case (filter, acls) =>
-        Resource(ResourceType.fromJava(filter.resourceType()), filter.name(), filter.patternType()) -> acls
+        new ResourcePattern(filter.resourceType(), filter.name(), filter.patternType()) -> acls
     }
 
     if (resourceToAcl.values.exists(_.isEmpty))
@@ -302,8 +421,8 @@ object AclCommand extends Logging {
     resourceToAcl
   }
 
-  private def getResourceFilterToAcls(opts: AclCommandOptions): Map[ResourcePatternFilter, Set[Acl]] = {
-    var resourceToAcls = Map.empty[ResourcePatternFilter, Set[Acl]]
+  private def getResourceFilterToAcls(opts: AclCommandOptions): Map[ResourcePatternFilter, Set[AccessControlEntry]] = {
+    var resourceToAcls = Map.empty[ResourcePatternFilter, Set[AccessControlEntry]]
 
     //if none of the --producer or --consumer options are specified , just construct ACLs from CLI options.
     if (!opts.options.has(opts.producerOpt) && !opts.options.has(opts.consumerOpt)) {
@@ -315,33 +434,33 @@ object AclCommand extends Logging {
       resourceToAcls ++= getProducerResourceFilterToAcls(opts)
 
     if (opts.options.has(opts.consumerOpt))
-      resourceToAcls ++= getConsumerResourceFilterToAcls(opts).map { case (k, v) => k -> (v ++ resourceToAcls.getOrElse(k, Set.empty[Acl])) }
+      resourceToAcls ++= getConsumerResourceFilterToAcls(opts).map { case (k, v) => k -> (v ++ resourceToAcls.getOrElse(k, Set.empty[AccessControlEntry])) }
 
     validateOperation(opts, resourceToAcls)
 
     resourceToAcls
   }
 
-  private def getProducerResourceFilterToAcls(opts: AclCommandOptions): Map[ResourcePatternFilter, Set[Acl]] = {
+  private def getProducerResourceFilterToAcls(opts: AclCommandOptions): Map[ResourcePatternFilter, Set[AccessControlEntry]] = {
     val filters = getResourceFilter(opts)
 
     val topics: Set[ResourcePatternFilter] = filters.filter(_.resourceType == JResourceType.TOPIC)
     val transactionalIds: Set[ResourcePatternFilter] = filters.filter(_.resourceType == JResourceType.TRANSACTIONAL_ID)
     val enableIdempotence = opts.options.has(opts.idempotentOpt)
 
-    val topicAcls = getAcl(opts, Set(Write, Describe, Create))
-    val transactionalIdAcls = getAcl(opts, Set(Write, Describe))
+    val topicAcls = getAcl(opts, Set(WRITE, DESCRIBE, CREATE))
+    val transactionalIdAcls = getAcl(opts, Set(WRITE, DESCRIBE))
 
     //Write, Describe, Create permission on topics, Write, Describe on transactionalIds
     topics.map(_ -> topicAcls).toMap ++
       transactionalIds.map(_ -> transactionalIdAcls).toMap ++
         (if (enableIdempotence)
-          Map(ClusterResourceFilter -> getAcl(opts, Set(IdempotentWrite)))
+          Map(ClusterResourceFilter -> getAcl(opts, Set(IDEMPOTENT_WRITE)))
         else
           Map.empty)
   }
 
-  private def getConsumerResourceFilterToAcls(opts: AclCommandOptions): Map[ResourcePatternFilter, Set[Acl]] = {
+  private def getConsumerResourceFilterToAcls(opts: AclCommandOptions): Map[ResourcePatternFilter, Set[AccessControlEntry]] = {
     val filters = getResourceFilter(opts)
 
     val topics: Set[ResourcePatternFilter] = filters.filter(_.resourceType == JResourceType.TOPIC)
@@ -349,19 +468,19 @@ object AclCommand extends Logging {
 
     //Read, Describe on topic, Read on consumerGroup
 
-    val acls = getAcl(opts, Set(Read, Describe))
+    val acls = getAcl(opts, Set(READ, DESCRIBE))
 
-    topics.map(_ -> acls).toMap[ResourcePatternFilter, Set[Acl]] ++
-      groups.map(_ -> getAcl(opts, Set(Read))).toMap[ResourcePatternFilter, Set[Acl]]
+    topics.map(_ -> acls).toMap[ResourcePatternFilter, Set[AccessControlEntry]] ++
+      groups.map(_ -> getAcl(opts, Set(READ))).toMap[ResourcePatternFilter, Set[AccessControlEntry]]
   }
 
-  private def getCliResourceFilterToAcls(opts: AclCommandOptions): Map[ResourcePatternFilter, Set[Acl]] = {
+  private def getCliResourceFilterToAcls(opts: AclCommandOptions): Map[ResourcePatternFilter, Set[AccessControlEntry]] = {
     val acls = getAcl(opts)
     val filters = getResourceFilter(opts)
     filters.map(_ -> acls).toMap
   }
 
-  private def getAcl(opts: AclCommandOptions, operations: Set[Operation]): Set[Acl] = {
+  private def getAcl(opts: AclCommandOptions, operations: Set[AclOperation]): Set[AccessControlEntry] = {
     val allowedPrincipals = getPrincipals(opts, opts.allowPrincipalsOpt)
 
     val deniedPrincipals = getPrincipals(opts, opts.denyPrincipalsOpt)
@@ -370,28 +489,28 @@ object AclCommand extends Logging {
 
     val deniedHosts = getHosts(opts, opts.denyHostsOpt, opts.denyPrincipalsOpt)
 
-    val acls = new collection.mutable.HashSet[Acl]
+    val acls = new collection.mutable.HashSet[AccessControlEntry]
     if (allowedHosts.nonEmpty && allowedPrincipals.nonEmpty)
-      acls ++= getAcls(allowedPrincipals, Allow, operations, allowedHosts)
+      acls ++= getAcls(allowedPrincipals, ALLOW, operations, allowedHosts)
 
     if (deniedHosts.nonEmpty && deniedPrincipals.nonEmpty)
-      acls ++= getAcls(deniedPrincipals, Deny, operations, deniedHosts)
+      acls ++= getAcls(deniedPrincipals, DENY, operations, deniedHosts)
 
     acls.toSet
   }
 
-  private def getAcl(opts: AclCommandOptions): Set[Acl] = {
-    val operations = opts.options.valuesOf(opts.operationsOpt).asScala.map(operation => Operation.fromString(operation.trim)).toSet
+  private def getAcl(opts: AclCommandOptions): Set[AccessControlEntry] = {
+    val operations = opts.options.valuesOf(opts.operationsOpt).asScala.map(operation => Operation.fromString(operation.trim)).map(_.toJava).toSet
     getAcl(opts, operations)
   }
 
-  def getAcls(principals: Set[KafkaPrincipal], permissionType: PermissionType, operations: Set[Operation],
-              hosts: Set[String]): Set[Acl] = {
+  def getAcls(principals: Set[KafkaPrincipal], permissionType: AclPermissionType, operations: Set[AclOperation],
+              hosts: Set[String]): Set[AccessControlEntry] = {
     for {
       principal <- principals
       operation <- operations
       host <- hosts
-    } yield new Acl(principal, permissionType, host, operation)
+    } yield new AccessControlEntry(principal.toString, host, operation, permissionType)
   }
 
   private def getHosts(opts: AclCommandOptions, hostOptionSpec: ArgumentAcceptingOptionSpec[String],
@@ -406,7 +525,7 @@ object AclCommand extends Logging {
 
   private def getPrincipals(opts: AclCommandOptions, principalOptionSpec: ArgumentAcceptingOptionSpec[String]): Set[KafkaPrincipal] = {
     if (opts.options.has(principalOptionSpec))
-      opts.options.valuesOf(principalOptionSpec).asScala.map(s => SecurityUtils.parseKafkaPrincipal(s.trim)).toSet
+      opts.options.valuesOf(principalOptionSpec).asScala.map(s => JSecurityUtils.parseKafkaPrincipal(s.trim)).toSet
     else
       Set.empty[KafkaPrincipal]
   }
@@ -444,10 +563,10 @@ object AclCommand extends Logging {
     StdIn.readLine().equalsIgnoreCase("y")
   }
 
-  private def validateOperation(opts: AclCommandOptions, resourceToAcls: Map[ResourcePatternFilter, Set[Acl]]): Unit = {
+  private def validateOperation(opts: AclCommandOptions, resourceToAcls: Map[ResourcePatternFilter, Set[AccessControlEntry]]): Unit = {
     for ((resource, acls) <- resourceToAcls) {
       val validOps = ResourceType.fromJava(resource.resourceType).supportedOperations + All
-      if ((acls.map(_.operation) -- validOps).nonEmpty)
+      if ((acls.map(_.operation) -- validOps.map(_.toJava)).nonEmpty)
         CommandLineUtils.printUsageAndDie(opts.parser, s"ResourceType ${resource.resourceType} only supports operations ${validOps.mkString(",")}")
     }
   }
@@ -573,7 +692,7 @@ object AclCommand extends Logging {
 
     options = parser.parse(args: _*)
 
-    def checkArgs() {
+    def checkArgs(): Unit = {
       if (options.has(bootstrapServerOpt) && options.has(authorizerOpt))
         CommandLineUtils.printUsageAndDie(parser, "Only one of --bootstrap-server or --authorizer must be specified")
 
diff --git a/core/src/main/scala/kafka/admin/AdminClient.scala b/core/src/main/scala/kafka/admin/AdminClient.scala
deleted file mode 100644
index 17716d3762f67..0000000000000
--- a/core/src/main/scala/kafka/admin/AdminClient.scala
+++ /dev/null
@@ -1,488 +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.IOException
-import java.nio.ByteBuffer
-import java.util.{Collections, Properties}
-import java.util.concurrent.atomic.AtomicInteger
-import java.util.concurrent.{ConcurrentLinkedQueue, Future, TimeUnit}
-
-import kafka.common.KafkaException
-import kafka.coordinator.group.GroupOverview
-import kafka.utils.Logging
-import org.apache.kafka.clients._
-import org.apache.kafka.clients.consumer.internals.{ConsumerNetworkClient, ConsumerProtocol, RequestFuture}
-import org.apache.kafka.common.config.ConfigDef.ValidString._
-import org.apache.kafka.common.config.ConfigDef.{Importance, Type}
-import org.apache.kafka.common.config.{AbstractConfig, ConfigDef}
-import org.apache.kafka.common.errors.{AuthenticationException, TimeoutException}
-import org.apache.kafka.common.internals.ClusterResourceListeners
-import org.apache.kafka.common.message.{DescribeGroupsRequestData, DescribeGroupsResponseData}
-import org.apache.kafka.common.metrics.Metrics
-import org.apache.kafka.common.network.Selector
-import org.apache.kafka.common.protocol.{ApiKeys, Errors}
-import org.apache.kafka.common.requests._
-import org.apache.kafka.common.requests.ApiVersionsResponse.ApiVersion
-import org.apache.kafka.common.requests.OffsetFetchResponse
-import org.apache.kafka.common.utils.LogContext
-import org.apache.kafka.common.utils.{KafkaThread, Time}
-import org.apache.kafka.common.{Node, TopicPartition}
-
-import scala.collection.JavaConverters._
-import scala.util.{Failure, Success, Try}
-
-/**
-  * A Scala administrative client for Kafka which supports managing and inspecting topics, brokers,
-  * and configurations. This client is deprecated, and will be replaced by org.apache.kafka.clients.admin.AdminClient.
-  */
-@deprecated("This class is deprecated in favour of org.apache.kafka.clients.admin.AdminClient and it will be removed in " +
-  "a future release.", since = "0.11.0")
-class AdminClient(val time: Time,
-                  val requestTimeoutMs: Int,
-                  val retryBackoffMs: Long,
-                  val client: ConsumerNetworkClient,
-                  val bootstrapBrokers: List[Node]) extends Logging {
-
-  @volatile var running: Boolean = true
-  val pendingFutures = new ConcurrentLinkedQueue[RequestFuture[ClientResponse]]()
-
-  val networkThread = new KafkaThread("admin-client-network-thread", new Runnable {
-    override def run() {
-      try {
-        while (running)
-          client.poll(time.timer(Long.MaxValue))
-      } catch {
-        case t : Throwable =>
-          error("admin-client-network-thread exited", t)
-      } finally {
-        pendingFutures.asScala.foreach { future =>
-          try {
-            future.raise(Errors.UNKNOWN_SERVER_ERROR)
-          } catch {
-            case _: IllegalStateException => // It is OK if the future has been completed
-          }
-        }
-        pendingFutures.clear()
-      }
-    }
-  }, true)
-
-  networkThread.start()
-
-  private def send(target: Node,
-                   api: ApiKeys,
-                   request: AbstractRequest.Builder[_ <: AbstractRequest]): AbstractResponse = {
-    val future: RequestFuture[ClientResponse] = client.send(target, request)
-    pendingFutures.add(future)
-    future.awaitDone(Long.MaxValue, TimeUnit.MILLISECONDS)
-    pendingFutures.remove(future)
-    if (future.succeeded())
-      future.value().responseBody()
-    else
-      throw future.exception()
-  }
-
-  private def sendAnyNode(api: ApiKeys, request: AbstractRequest.Builder[_ <: AbstractRequest]): AbstractResponse = {
-    bootstrapBrokers.foreach { broker =>
-      try {
-        return send(broker, api, request)
-      } catch {
-        case e: AuthenticationException =>
-          throw e
-        case e: Exception =>
-          debug(s"Request $api failed against node $broker", e)
-      }
-    }
-    throw new RuntimeException(s"Request $api failed on brokers $bootstrapBrokers")
-  }
-
-  def findCoordinator(groupId: String, timeoutMs: Long = 0): Node = {
-    val requestBuilder = new FindCoordinatorRequest.Builder(FindCoordinatorRequest.CoordinatorType.GROUP, groupId)
-
-    def sendRequest: Try[FindCoordinatorResponse] =
-      Try(sendAnyNode(ApiKeys.FIND_COORDINATOR, requestBuilder).asInstanceOf[FindCoordinatorResponse])
-
-    val startTime = time.milliseconds
-    var response = sendRequest
-
-    while ((response.isFailure || response.get.error == Errors.COORDINATOR_NOT_AVAILABLE) &&
-      (time.milliseconds - startTime < timeoutMs)) {
-
-      Thread.sleep(retryBackoffMs)
-      response = sendRequest
-    }
-
-    def timeoutException(cause: Throwable) =
-      throw new TimeoutException("The consumer group command timed out while waiting for group to initialize: ", cause)
-
-    response match {
-      case Failure(exception) => throw timeoutException(exception)
-      case Success(response) =>
-        if (response.error == Errors.COORDINATOR_NOT_AVAILABLE)
-          throw timeoutException(response.error.exception)
-        response.error.maybeThrow()
-        response.node
-    }
-  }
-
-  def listGroups(node: Node): List[GroupOverview] = {
-    val response = send(node, ApiKeys.LIST_GROUPS, new ListGroupsRequest.Builder()).asInstanceOf[ListGroupsResponse]
-    response.error.maybeThrow()
-    response.groups.asScala.map(group => GroupOverview(group.groupId, group.protocolType)).toList
-  }
-
-  def getApiVersions(node: Node): List[ApiVersion] = {
-    val response = send(node, ApiKeys.API_VERSIONS, new ApiVersionsRequest.Builder()).asInstanceOf[ApiVersionsResponse]
-    response.error.maybeThrow()
-    response.apiVersions.asScala.toList
-  }
-
-  /**
-   * Wait until there is a non-empty list of brokers in the cluster.
-   */
-  def awaitBrokers() {
-    var nodes = List[Node]()
-    do {
-      nodes = findAllBrokers()
-      if (nodes.isEmpty)
-        Thread.sleep(50)
-    } while (nodes.isEmpty)
-  }
-
-  def findAllBrokers(): List[Node] = {
-    val request = MetadataRequest.Builder.allTopics()
-    val response = sendAnyNode(ApiKeys.METADATA, request).asInstanceOf[MetadataResponse]
-    val errors = response.errors
-    if (!errors.isEmpty)
-      debug(s"Metadata request contained errors: $errors")
-    response.cluster.nodes.asScala.toList
-  }
-
-  def listAllGroups(): Map[Node, List[GroupOverview]] = {
-    findAllBrokers().map { broker =>
-      broker -> {
-        try {
-          listGroups(broker)
-        } catch {
-          case e: Exception =>
-            debug(s"Failed to find groups from broker $broker", e)
-            List[GroupOverview]()
-        }
-      }
-    }.toMap
-  }
-
-  def listAllConsumerGroups(): Map[Node, List[GroupOverview]] = {
-    listAllGroups().mapValues { groups =>
-      groups.filter(isConsumerGroup)
-    }
-  }
-
-  def listAllGroupsFlattened(): List[GroupOverview] = {
-    listAllGroups().values.flatten.toList
-  }
-
-  def listAllConsumerGroupsFlattened(): List[GroupOverview] = {
-    listAllGroupsFlattened().filter(isConsumerGroup)
-  }
-
-  private def isConsumerGroup(group: GroupOverview): Boolean = {
-    // Consumer groups which are using group management use the "consumer" protocol type.
-    // Consumer groups which are only using offset storage will have an empty protocol type.
-    group.protocolType.isEmpty || group.protocolType == ConsumerProtocol.PROTOCOL_TYPE
-  }
-
-  def listGroupOffsets(groupId: String): Map[TopicPartition, Long] = {
-    val coordinator = findCoordinator(groupId)
-    val responseBody = send(coordinator, ApiKeys.OFFSET_FETCH, OffsetFetchRequest.Builder.allTopicPartitions(groupId))
-    val response = responseBody.asInstanceOf[OffsetFetchResponse]
-    if (response.hasError)
-      throw response.error.exception
-    response.maybeThrowFirstPartitionError()
-    response.responseData.asScala.map { case (tp, partitionData) => (tp, partitionData.offset) }.toMap
-  }
-
-  def listAllBrokerVersionInfo(): Map[Node, Try[NodeApiVersions]] =
-    findAllBrokers().map { broker =>
-      broker -> Try[NodeApiVersions](new NodeApiVersions(getApiVersions(broker).asJava))
-    }.toMap
-
-  /**
-   * Case class used to represent a consumer of a consumer group
-   */
-  case class ConsumerSummary(consumerId: String,
-                             clientId: String,
-                             host: String,
-                             assignment: List[TopicPartition])
-
-  /**
-   * Case class used to represent group metadata (including the group coordinator) for the DescribeGroup API
-   */
-  case class ConsumerGroupSummary(state: String,
-                                  assignmentStrategy: String,
-                                  consumers: Option[List[ConsumerSummary]],
-                                  coordinator: Node)
-
-  def describeConsumerGroupHandler(coordinator: Node, groupId: String): DescribeGroupsResponseData.DescribedGroup = {
-    val responseBody = send(coordinator, ApiKeys.DESCRIBE_GROUPS,
-      new DescribeGroupsRequest.Builder(new DescribeGroupsRequestData().setGroups(Collections.singletonList(groupId))))
-    val response = responseBody.asInstanceOf[DescribeGroupsResponse]
-    val metadata = response.data().groups().asScala.find(group => groupId.equals(group.groupId()))
-      .getOrElse(throw new KafkaException(s"Response from broker contained no metadata for group $groupId"))
-    metadata
-  }
-
-  def describeConsumerGroup(groupId: String, timeoutMs: Long = 0): ConsumerGroupSummary = {
-
-    def isValidConsumerGroupResponse(metadata: DescribeGroupsResponseData.DescribedGroup): Boolean =
-      metadata.errorCode() == Errors.NONE.code() && (metadata.groupState() == "Dead" ||
-        metadata.groupState() == "Empty" || metadata.protocolType == ConsumerProtocol.PROTOCOL_TYPE)
-
-    val startTime = time.milliseconds
-    val coordinator = findCoordinator(groupId, timeoutMs)
-    var metadata = describeConsumerGroupHandler(coordinator, groupId)
-
-    while (!isValidConsumerGroupResponse(metadata) && time.milliseconds - startTime < timeoutMs) {
-      debug(s"The consumer group response for group '$groupId' is invalid. Retrying the request as the group is initializing ...")
-      Thread.sleep(retryBackoffMs)
-      metadata = describeConsumerGroupHandler(coordinator, groupId)
-    }
-
-    if (!isValidConsumerGroupResponse(metadata))
-      throw new TimeoutException("The consumer group command timed out while waiting for group to initialize")
-
-    val consumers = metadata.members.asScala.map { consumer =>
-      ConsumerSummary(consumer.memberId, consumer.clientId, consumer.clientHost, metadata.groupState() match {
-        case "Stable" =>
-          val assignment = ConsumerProtocol.deserializeAssignment(ByteBuffer.wrap(consumer.memberAssignment))
-          assignment.partitions.asScala.toList
-        case _ =>
-          List()
-      })
-    }.toList
-
-    ConsumerGroupSummary(metadata.groupState(), metadata.protocolData(), Some(consumers), coordinator)
-  }
-
-  def deleteConsumerGroups(groups: List[String]): Map[String, Errors] = {
-
-    def coordinatorLookup(group: String): Either[Node, Errors] = {
-      try {
-        Left(findCoordinator(group))
-      } catch {
-        case e: Throwable =>
-          if (e.isInstanceOf[TimeoutException])
-            Right(Errors.COORDINATOR_NOT_AVAILABLE)
-          else
-            Right(Errors.forException(e))
-      }
-    }
-
-    var errors: Map[String, Errors] = Map()
-    var groupsPerCoordinator: Map[Node, List[String]] = Map()
-
-    groups.foreach { group =>
-      coordinatorLookup(group) match {
-        case Right(error) =>
-          errors += group -> error
-        case Left(coordinator) =>
-          groupsPerCoordinator.get(coordinator) match {
-            case Some(gList) =>
-              val gListNew = group :: gList
-              groupsPerCoordinator += coordinator -> gListNew
-            case None =>
-              groupsPerCoordinator += coordinator -> List(group)
-          }
-      }
-    }
-
-    groupsPerCoordinator.foreach { case (coordinator, groups) =>
-      val responseBody = send(coordinator, ApiKeys.DELETE_GROUPS, new DeleteGroupsRequest.Builder(groups.toSet.asJava))
-      val response = responseBody.asInstanceOf[DeleteGroupsResponse]
-      groups.foreach {
-        case group if response.hasError(group) => errors += group -> response.errors.get(group)
-        case group => errors += group -> Errors.NONE
-      }
-    }
-
-    errors
-  }
-
-  def close() {
-    running = false
-    try {
-      client.close()
-    } catch {
-      case e: IOException =>
-        error("Exception closing nioSelector:", e)
-    }
-  }
-
-}
-
-/*
- * CompositeFuture assumes that the future object in the futures list does not raise error
- */
-class CompositeFuture[T](time: Time,
-                         defaultResults: Map[TopicPartition, T],
-                         futures: List[RequestFuture[Map[TopicPartition, T]]]) extends Future[Map[TopicPartition, T]] {
-
-  override def isCancelled = false
-
-  override def cancel(interrupt: Boolean) = false
-
-  override def get(): Map[TopicPartition, T] = {
-    get(Long.MaxValue, TimeUnit.MILLISECONDS)
-  }
-
-  override def get(timeout: Long, unit: TimeUnit): Map[TopicPartition, T] = {
-    val start: Long = time.milliseconds()
-    val timeoutMs = unit.toMillis(timeout)
-    var remaining: Long = timeoutMs
-
-    val observedResults = futures.flatMap { future =>
-      val elapsed = time.milliseconds() - start
-      remaining = if (timeoutMs - elapsed > 0) timeoutMs - elapsed else 0L
-
-      if (future.awaitDone(remaining, TimeUnit.MILLISECONDS)) future.value()
-      else Map.empty[TopicPartition, T]
-    }.toMap
-
-    defaultResults ++ observedResults
-  }
-
-  override def isDone: Boolean = {
-    futures.forall(_.isDone)
-  }
-}
-
-@deprecated("This class is deprecated in favour of org.apache.kafka.clients.admin.AdminClient and it will be removed in " +
-  "a future release.", since = "0.11.0")
-object AdminClient {
-  val DefaultConnectionMaxIdleMs = 9 * 60 * 1000
-  val DefaultRequestTimeoutMs = 5000
-  val DefaultMaxInFlightRequestsPerConnection = 100
-  val DefaultReconnectBackoffMs = 50
-  val DefaultReconnectBackoffMax = 50
-  val DefaultSendBufferBytes = 128 * 1024
-  val DefaultReceiveBufferBytes = 32 * 1024
-  val DefaultRetryBackoffMs = 100
-
-  val AdminClientIdSequence = new AtomicInteger(1)
-  val AdminConfigDef = {
-    val config = new ConfigDef()
-      .define(
-        CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG,
-        Type.LIST,
-        Importance.HIGH,
-        CommonClientConfigs.BOOTSTRAP_SERVERS_DOC)
-      .define(CommonClientConfigs.CLIENT_DNS_LOOKUP_CONFIG,
-        Type.STRING,
-        ClientDnsLookup.DEFAULT.toString,
-        in(ClientDnsLookup.DEFAULT.toString,
-           ClientDnsLookup.USE_ALL_DNS_IPS.toString,
-           ClientDnsLookup.RESOLVE_CANONICAL_BOOTSTRAP_SERVERS_ONLY.toString),
-        Importance.MEDIUM,
-        CommonClientConfigs.CLIENT_DNS_LOOKUP_DOC)
-      .define(
-        CommonClientConfigs.SECURITY_PROTOCOL_CONFIG,
-        ConfigDef.Type.STRING,
-        CommonClientConfigs.DEFAULT_SECURITY_PROTOCOL,
-        ConfigDef.Importance.MEDIUM,
-        CommonClientConfigs.SECURITY_PROTOCOL_DOC)
-      .define(
-        CommonClientConfigs.REQUEST_TIMEOUT_MS_CONFIG,
-        ConfigDef.Type.INT,
-        DefaultRequestTimeoutMs,
-        ConfigDef.Importance.MEDIUM,
-        CommonClientConfigs.REQUEST_TIMEOUT_MS_DOC)
-      .define(
-        CommonClientConfigs.RETRY_BACKOFF_MS_CONFIG,
-        ConfigDef.Type.LONG,
-        DefaultRetryBackoffMs,
-        ConfigDef.Importance.MEDIUM,
-        CommonClientConfigs.RETRY_BACKOFF_MS_DOC)
-      .withClientSslSupport()
-      .withClientSaslSupport()
-    config
-  }
-
-  class AdminConfig(originals: Map[_,_]) extends AbstractConfig(AdminConfigDef, originals.asJava, false)
-
-  def createSimplePlaintext(brokerUrl: String): AdminClient = {
-    val config = Map(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG -> brokerUrl)
-    create(new AdminConfig(config))
-  }
-
-  def create(props: Properties): AdminClient = create(props.asScala.toMap)
-
-  def create(props: Map[String, _]): AdminClient = create(new AdminConfig(props))
-
-  def create(config: AdminConfig): AdminClient = {
-    val clientId = "admin-" + AdminClientIdSequence.getAndIncrement()
-    val logContext = new LogContext(s"[LegacyAdminClient clientId=$clientId] ")
-    val time = Time.SYSTEM
-    val metrics = new Metrics(time)
-    val metadata = new Metadata(100L, 60 * 60 * 1000L, logContext,
-      new ClusterResourceListeners)
-    val channelBuilder = ClientUtils.createChannelBuilder(config, time)
-    val requestTimeoutMs = config.getInt(CommonClientConfigs.REQUEST_TIMEOUT_MS_CONFIG)
-    val retryBackoffMs = config.getLong(CommonClientConfigs.RETRY_BACKOFF_MS_CONFIG)
-
-    val brokerUrls = config.getList(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG)
-    val clientDnsLookup = config.getString(CommonClientConfigs.CLIENT_DNS_LOOKUP_CONFIG)
-    val brokerAddresses = ClientUtils.parseAndValidateAddresses(brokerUrls, clientDnsLookup)
-    metadata.bootstrap(brokerAddresses, time.milliseconds())
-
-    val selector = new Selector(
-      DefaultConnectionMaxIdleMs,
-      metrics,
-      time,
-      "admin",
-      channelBuilder,
-      logContext)
-
-    val networkClient = new NetworkClient(
-      selector,
-      metadata,
-      clientId,
-      DefaultMaxInFlightRequestsPerConnection,
-      DefaultReconnectBackoffMs,
-      DefaultReconnectBackoffMax,
-      DefaultSendBufferBytes,
-      DefaultReceiveBufferBytes,
-      requestTimeoutMs,
-      ClientDnsLookup.DEFAULT,
-      time,
-      true,
-      new ApiVersions,
-      logContext)
-
-    val highLevelClient = new ConsumerNetworkClient(
-      logContext,
-      networkClient,
-      metadata,
-      time,
-      retryBackoffMs,
-      requestTimeoutMs,
-      Integer.MAX_VALUE)
-
-    new AdminClient(
-      time,
-      requestTimeoutMs,
-      retryBackoffMs,
-      highLevelClient,
-      metadata.fetch.nodes.asScala.toList)
-  }
-}
diff --git a/core/src/main/scala/kafka/admin/AdminUtils.scala b/core/src/main/scala/kafka/admin/AdminUtils.scala
index cd479691294e1..ce4052dbe990f 100644
--- a/core/src/main/scala/kafka/admin/AdminUtils.scala
+++ b/core/src/main/scala/kafka/admin/AdminUtils.scala
@@ -17,110 +17,84 @@
 
 package kafka.admin
 
-import kafka.log.LogConfig
-import kafka.server.{ConfigEntityName, ConfigType, DynamicConfig}
-import kafka.utils._
-import kafka.utils.ZkUtils._
 import java.util.Random
-import java.util.Properties
 
-import kafka.common.TopicAlreadyMarkedForDeletionException
-import org.apache.kafka.common.errors._
+import kafka.utils.Logging
+import org.apache.kafka.common.errors.{InvalidPartitionsException, InvalidReplicationFactorException}
 
-import collection.{Map, Set, mutable, _}
-import scala.collection.JavaConverters._
-import org.I0Itec.zkclient.exception.ZkNodeExistsException
-import org.apache.kafka.common.internals.Topic
+import collection.{Map, mutable, _}
 
-@deprecated("This class is deprecated and will be replaced by kafka.zk.AdminZkClient.", "1.1.0")
-trait AdminUtilities {
-  def changeTopicConfig(zkUtils: ZkUtils, topic: String, configs: Properties)
-  def changeClientIdConfig(zkUtils: ZkUtils, clientId: String, configs: Properties)
-  def changeUserOrUserClientIdConfig(zkUtils: ZkUtils, sanitizedEntityName: String, configs: Properties)
-  def changeBrokerConfig(zkUtils: ZkUtils, brokerIds: Seq[Int], configs: Properties)
-
-  def changeConfigs(zkUtils: ZkUtils, entityType: String, entityName: String, configs: Properties): Unit = {
-
-    def parseBroker(broker: String): Int = {
-      try broker.toInt
-      catch {
-        case _: NumberFormatException =>
-          throw new IllegalArgumentException(s"Error parsing broker $broker. The broker's Entity Name must be a single integer value")
-      }
-    }
-
-    entityType match {
-      case ConfigType.Topic => changeTopicConfig(zkUtils, entityName, configs)
-      case ConfigType.Client => changeClientIdConfig(zkUtils, entityName, configs)
-      case ConfigType.User => changeUserOrUserClientIdConfig(zkUtils, entityName, configs)
-      case ConfigType.Broker => changeBrokerConfig(zkUtils, Seq(parseBroker(entityName)), configs)
-      case _ => throw new IllegalArgumentException(s"$entityType is not a known entityType. Should be one of ${ConfigType.Topic}, ${ConfigType.Client}, ${ConfigType.Broker}")
-    }
-  }
-
-  def fetchEntityConfig(zkUtils: ZkUtils, entityType: String, entityName: String): Properties
-}
-
-object AdminUtils extends Logging with AdminUtilities {
+object AdminUtils extends Logging {
   val rand = new Random
   val AdminClientId = "__admin_client"
-  val EntityConfigChangeZnodePrefix = "config_change_"
 
   /**
    * There are 3 goals of replica assignment:
    *
-   * 1. Spread the replicas evenly among brokers.
-   * 2. For partitions assigned to a particular broker, their other replicas are spread over the other brokers.
-   * 3. If all brokers have rack information, assign the replicas for each partition to different racks if possible
+   * 
        
+   * 
      1.  Spread the replicas evenly among brokers.
        2. 
+   * 
      2.  For partitions assigned to a particular broker, their other replicas are spread over the other brokers.
        3. 
+   * 
      3.  If all brokers have rack information, assign the replicas for each partition to different racks if possible
        4. 
+   * 
      
    *
    * To achieve this goal for replica assignment without considering racks, we:
-   * 1. Assign the first replica of each partition by round-robin, starting from a random position in the broker list.
-   * 2. Assign the remaining replicas of each partition with an increasing shift.
+   * 
        
+   * 
      1.  Assign the first replica of each partition by round-robin, starting from a random position in the broker list.
        2. 
+   * 
      2.  Assign the remaining replicas of each partition with an increasing shift.
        3. 
+   * 
      
    *
    * Here is an example of assigning
-   * broker-0  broker-1  broker-2  broker-3  broker-4
-   * p0        p1        p2        p3        p4       (1st replica)
-   * p5        p6        p7        p8        p9       (1st replica)
-   * p4        p0        p1        p2        p3       (2nd replica)
-   * p8        p9        p5        p6        p7       (2nd replica)
-   * p3        p4        p0        p1        p2       (3nd replica)
-   * p7        p8        p9        p5        p6       (3nd replica)
-   *
+   * 
      
+   * 
+   * 
+   * 
+   * 
+   * 
+   * 
+   * 
+   * 
      broker-0broker-1broker-2broker-3broker-4 
      p0      p1      p2      p3      p4      (1st replica)
      p5      p6      p7      p8      p9      (1st replica)
      p4      p0      p1      p2      p3      (2nd replica)
      p8      p9      p5      p6      p7      (2nd replica)
      p3      p4      p0      p1      p2      (3nd replica)
      p7      p8      p9      p5      p6      (3nd replica)
      
+   *
+   * 
      
    * To create rack aware assignment, this API will first create a rack alternated broker list. For example,
-   * from this brokerID -> rack mapping:
-   *
+   * from this brokerID -> rack mapping:
      
    * 0 -> "rack1", 1 -> "rack3", 2 -> "rack3", 3 -> "rack2", 4 -> "rack2", 5 -> "rack1"
-   *
+   * 

      
+   * 
      
    * The rack alternated list will be:
-   *
+   * 
      
    * 0, 3, 1, 5, 4, 2
-   *
+   * 

      
+   * 
      
    * Then an easy round-robin assignment can be applied. Assume 6 partitions with replication factor of 3, the assignment
    * will be:
-   *
-   * 0 -> 0,3,1
-   * 1 -> 3,1,5
-   * 2 -> 1,5,4
-   * 3 -> 5,4,2
-   * 4 -> 4,2,0
-   * 5 -> 2,0,3
-   *
+   * 
      
+   * 0 -> 0,3,1 
      
+   * 1 -> 3,1,5 
      
+   * 2 -> 1,5,4 
      
+   * 3 -> 5,4,2 
      
+   * 4 -> 4,2,0 
      
+   * 5 -> 2,0,3 
      
+   * 
      
+   * 
      
    * Once it has completed the first round-robin, if there are more partitions to assign, the algorithm will start
    * shifting the followers. This is to ensure we will not always get the same set of sequences.
    * In this case, if there is another partition to assign (partition #6), the assignment will be:
-   *
+   * 
      
    * 6 -> 0,4,2 (instead of repeating 0,3,1 as partition 0)
-   *
+   * 

      
+   * 
      
    * The rack aware assignment always chooses the 1st replica of the partition using round robin on the rack alternated
    * broker list. For rest of the replicas, it will be biased towards brokers on racks that do not have
    * any replica assignment, until every rack has a replica. Then the assignment will go back to round-robin on
    * the broker list.
-   *
+   * 
      
+   * 
      
+   * 
      
    * As the result, if the number of replicas is equal to or greater than the number of racks, it will ensure that
    * each rack will get at least one replica. Otherwise, each rack will get at most one replica. In a perfect
    * situation where the number of replicas is the same as the number of racks and each rack has the same number of
    * brokers, it guarantees that the replica distribution is even across brokers and racks.
-   *
+   * 
      
    * @return a Map from partition id to replica ids
    * @throws AdminOperationException If rack information is supplied but it is incomplete, or if it is not possible to
    *                                 assign each replica to a unique rack.
@@ -256,391 +230,6 @@ object AdminUtils extends Logging with AdminUtilities {
       .groupBy { case (rack, _) => rack }
       .map { case (rack, rackAndIdList) => (rack, rackAndIdList.map { case (_, id) => id }.sorted) }
   }
- /**
-  * Add partitions to existing topic with optional replica assignment
-  *
-  * @param zkUtils Zookeeper utilities
-  * @param topic Topic for adding partitions to
-  * @param existingAssignment A map from partition id to its assigned replicas
-  * @param allBrokers All brokers in the cluster
-  * @param numPartitions Number of partitions to be set
-  * @param replicaAssignment Manual replica assignment, or none
-  * @param validateOnly If true, validate the parameters without actually adding the partitions
-  * @return the updated replica assignment
-  */
-  @deprecated("This method is deprecated and will be replaced by kafka.zk.AdminZkClient.", "1.1.0")
-  def addPartitions(zkUtils: ZkUtils,
-                    topic: String,
-                    existingAssignment: Map[Int, Seq[Int]],
-                    allBrokers: Seq[BrokerMetadata],
-                    numPartitions: Int = 1,
-                    replicaAssignment: Option[Map[Int, Seq[Int]]] = None,
-                    validateOnly: Boolean = false): Map[Int, Seq[Int]] = {
-    val existingAssignmentPartition0 = existingAssignment.getOrElse(0,
-      throw new AdminOperationException(
-        s"Unexpected existing replica assignment for topic '$topic', partition id 0 is missing. " +
-          s"Assignment: $existingAssignment"))
-
-    val partitionsToAdd = numPartitions - existingAssignment.size
-    if (partitionsToAdd <= 0)
-      throw new InvalidPartitionsException(
-        s"The number of partitions for a topic can only be increased. " +
-          s"Topic $topic currently has ${existingAssignment.size} partitions, " +
-          s"$numPartitions would not be an increase.")
-
-    replicaAssignment.foreach { proposedReplicaAssignment =>
-      validateReplicaAssignment(proposedReplicaAssignment, existingAssignmentPartition0,
-        allBrokers.map(_.id).toSet)
-    }
-
-    val proposedAssignmentForNewPartitions = replicaAssignment.getOrElse {
-      val startIndex = math.max(0, allBrokers.indexWhere(_.id >= existingAssignmentPartition0.head))
-      AdminUtils.assignReplicasToBrokers(allBrokers, partitionsToAdd, existingAssignmentPartition0.size,
-        startIndex, existingAssignment.size)
-    }
-    val proposedAssignment = existingAssignment ++ proposedAssignmentForNewPartitions
-    if (!validateOnly) {
-      info(s"Creating $partitionsToAdd partitions for '$topic' with the following replica assignment: " +
-        s"$proposedAssignmentForNewPartitions.")
-      // add the combined new list
-      AdminUtils.createOrUpdateTopicPartitionAssignmentPathInZK(zkUtils, topic, proposedAssignment, update = true)
-    }
-    proposedAssignment
-
-  }
-
-  /**
-    * Parse a replica assignment string of the form:
-    * {{{
-    * broker_id_for_part1_replica1:broker_id_for_part1_replica2,
-    * broker_id_for_part2_replica1:broker_id_for_part2_replica2,
-    * ...
-    * }}}
-    */
-  def parseReplicaAssignment(replicaAssignmentsString: String, startPartitionId: Int): Map[Int, Seq[Int]] = {
-    val assignmentStrings = replicaAssignmentsString.split(",")
-    val assignmentMap = mutable.Map[Int, Seq[Int]]()
-    var partitionId = startPartitionId
-    for (assignmentString <- assignmentStrings) {
-      val brokerIds = assignmentString.split(":").map(_.trim.toInt).toSeq
-      assignmentMap.put(partitionId, brokerIds)
-      partitionId = partitionId + 1
-    }
-    assignmentMap
-  }
-
-  private def validateReplicaAssignment(replicaAssignment: Map[Int, Seq[Int]],
-                                        existingAssignmentPartition0: Seq[Int],
-                                        availableBrokerIds: Set[Int]): Unit = {
-
-    replicaAssignment.foreach { case (partitionId, replicas) =>
-      if (replicas.isEmpty)
-        throw new InvalidReplicaAssignmentException(
-          s"Cannot have replication factor of 0 for partition id $partitionId.")
-      if (replicas.size != replicas.toSet.size)
-        throw new InvalidReplicaAssignmentException(
-          s"Duplicate brokers not allowed in replica assignment: " +
-            s"${replicas.mkString(", ")} for partition id $partitionId.")
-      if (!replicas.toSet.subsetOf(availableBrokerIds))
-        throw new BrokerNotAvailableException(
-          s"Some brokers specified for partition id $partitionId are not available. " +
-            s"Specified brokers: ${replicas.mkString(", ")}, " +
-            s"available brokers: ${availableBrokerIds.mkString(", ")}.")
-      partitionId -> replicas.size
-    }
-    val badRepFactors = replicaAssignment.collect {
-      case (partition, replicas) if replicas.size != existingAssignmentPartition0.size => partition -> replicas.size
-    }
-    if (badRepFactors.nonEmpty) {
-      val sortedBadRepFactors = badRepFactors.toSeq.sortBy { case (partitionId, _) => partitionId }
-      val partitions = sortedBadRepFactors.map { case (partitionId, _) => partitionId }
-      val repFactors = sortedBadRepFactors.map { case (_, rf) => rf }
-      throw new InvalidReplicaAssignmentException(s"Inconsistent replication factor between partitions, " +
-        s"partition 0 has ${existingAssignmentPartition0.size} while partitions [${partitions.mkString(", ")}] have " +
-        s"replication factors [${repFactors.mkString(", ")}], respectively.")
-    }
-  }
-
-  @deprecated("This method is deprecated and will be replaced by kafka.zk.AdminZkClient.", "1.1.0")
-  def deleteTopic(zkUtils: ZkUtils, topic: String) {
-    if (topicExists(zkUtils, topic)) {
-      try {
-        zkUtils.createPersistentPath(getDeleteTopicPath(topic))
-      } catch {
-        case _: ZkNodeExistsException => throw new TopicAlreadyMarkedForDeletionException(
-          "topic %s is already marked for deletion".format(topic))
-        case e2: Throwable => throw new AdminOperationException(e2)
-      }
-    } else {
-      throw new UnknownTopicOrPartitionException(s"Topic `$topic` to delete does not exist")
-    }
-  }
-
-  @deprecated("This method is deprecated and will be replaced by kafka.zk.AdminZkClient.", "1.1.0")
-  def topicExists(zkUtils: ZkUtils, topic: String): Boolean =
-    zkUtils.pathExists(getTopicPath(topic))
-
-  @deprecated("This method is deprecated and will be replaced by kafka.zk.AdminZkClient.", "1.1.0")
-  def getBrokerMetadatas(zkUtils: ZkUtils, rackAwareMode: RackAwareMode = RackAwareMode.Enforced,
-                         brokerList: Option[Seq[Int]] = None): Seq[BrokerMetadata] = {
-    val allBrokers = zkUtils.getAllBrokersInCluster()
-    val brokers = brokerList.map(brokerIds => allBrokers.filter(b => brokerIds.contains(b.id))).getOrElse(allBrokers)
-    val brokersWithRack = brokers.filter(_.rack.nonEmpty)
-    if (rackAwareMode == RackAwareMode.Enforced && brokersWithRack.nonEmpty && brokersWithRack.size < brokers.size) {
-      throw new AdminOperationException("Not all brokers have rack information. Add --disable-rack-aware in command line" +
-        " to make replica assignment without rack information.")
-    }
-    val brokerMetadatas = rackAwareMode match {
-      case RackAwareMode.Disabled => brokers.map(broker => BrokerMetadata(broker.id, None))
-      case RackAwareMode.Safe if brokersWithRack.size < brokers.size =>
-        brokers.map(broker => BrokerMetadata(broker.id, None))
-      case _ => brokers.map(broker => BrokerMetadata(broker.id, broker.rack))
-    }
-    brokerMetadatas.sortBy(_.id)
-  }
-
-  @deprecated("This method is deprecated and will be replaced by kafka.zk.AdminZkClient.", "1.1.0")
-  def createTopic(zkUtils: ZkUtils,
-                  topic: String,
-                  partitions: Int,
-                  replicationFactor: Int,
-                  topicConfig: Properties = new Properties,
-                  rackAwareMode: RackAwareMode = RackAwareMode.Enforced) {
-    val brokerMetadatas = getBrokerMetadatas(zkUtils, rackAwareMode)
-    val replicaAssignment = AdminUtils.assignReplicasToBrokers(brokerMetadatas, partitions, replicationFactor)
-    AdminUtils.createOrUpdateTopicPartitionAssignmentPathInZK(zkUtils, topic, replicaAssignment, topicConfig)
-  }
-
-  @deprecated("This method is deprecated and will be replaced by kafka.zk.AdminZkClient.", "1.1.0")
-  def validateCreateOrUpdateTopic(zkUtils: ZkUtils,
-                                  topic: String,
-                                  partitionReplicaAssignment: Map[Int, Seq[Int]],
-                                  config: Properties,
-                                  update: Boolean): Unit = {
-    // validate arguments
-    Topic.validate(topic)
-
-    if (!update) {
-      if (topicExists(zkUtils, topic))
-        throw new TopicExistsException(s"Topic '$topic' already exists.")
-      else if (Topic.hasCollisionChars(topic)) {
-        val allTopics = zkUtils.getAllTopics()
-        // check again in case the topic was created in the meantime, otherwise the
-        // topic could potentially collide with itself
-        if (allTopics.contains(topic))
-          throw new TopicExistsException(s"Topic '$topic' already exists.")
-        val collidingTopics = allTopics.filter(Topic.hasCollision(topic, _))
-        if (collidingTopics.nonEmpty) {
-          throw new InvalidTopicException(s"Topic '$topic' collides with existing topics: ${collidingTopics.mkString(", ")}")
-        }
-      }
-    }
-
-    if (partitionReplicaAssignment.values.map(_.size).toSet.size != 1)
-      throw new InvalidReplicaAssignmentException("All partitions should have the same number of replicas")
-
-    partitionReplicaAssignment.values.foreach(reps =>
-      if (reps.size != reps.toSet.size)
-        throw new InvalidReplicaAssignmentException("Duplicate replica assignment found: " + partitionReplicaAssignment)
-    )
-
-
-    // Configs only matter if a topic is being created. Changing configs via AlterTopic is not supported
-    if (!update)
-      LogConfig.validate(config)
-  }
-
-  @deprecated("This method is deprecated and will be replaced by kafka.zk.AdminZkClient.", "1.1.0")
-  def createOrUpdateTopicPartitionAssignmentPathInZK(zkUtils: ZkUtils,
-                                                     topic: String,
-                                                     partitionReplicaAssignment: Map[Int, Seq[Int]],
-                                                     config: Properties = new Properties,
-                                                     update: Boolean = false) {
-    validateCreateOrUpdateTopic(zkUtils, topic, partitionReplicaAssignment, config, update)
-
-    // Configs only matter if a topic is being created. Changing configs via AlterTopic is not supported
-    if (!update) {
-      // write out the config if there is any, this isn't transactional with the partition assignments
-      writeEntityConfig(zkUtils, getEntityConfigPath(ConfigType.Topic, topic), config)
-    }
-
-    // create the partition assignment
-    writeTopicPartitionAssignment(zkUtils, topic, partitionReplicaAssignment, update)
-  }
-
-  @deprecated("This method is deprecated and will be replaced by kafka.zk.AdminZkClient.", "1.1.0")
-  private def writeTopicPartitionAssignment(zkUtils: ZkUtils, topic: String, replicaAssignment: Map[Int, Seq[Int]], update: Boolean) {
-    try {
-      val zkPath = getTopicPath(topic)
-      val jsonPartitionData = zkUtils.replicaAssignmentZkData(replicaAssignment.map(e => e._1.toString -> e._2))
-
-      if (!update) {
-        info(s"Topic creation $jsonPartitionData")
-        zkUtils.createPersistentPath(zkPath, jsonPartitionData)
-      } else {
-        info(s"Topic update $jsonPartitionData")
-        zkUtils.updatePersistentPath(zkPath, jsonPartitionData)
-      }
-      debug("Updated path %s with %s for replica assignment".format(zkPath, jsonPartitionData))
-    } catch {
-      case _: ZkNodeExistsException => throw new TopicExistsException(s"Topic '$topic' already exists.")
-      case e2: Throwable => throw new AdminOperationException(e2.toString)
-    }
-  }
-
-  /**
-   * Update the config for a client and create a change notification so the change will propagate to other brokers.
-   * If clientId is , default clientId config is updated. ClientId configs are used only if 
-   * and  configs are not specified.
-   *
-   * @param zkUtils Zookeeper utilities used to write the config to ZK
-   * @param sanitizedClientId: The sanitized clientId for which configs are being changed
-   * @param configs: The final set of configs that will be applied to the topic. If any new configs need to be added or
-   *                 existing configs need to be deleted, it should be done prior to invoking this API
-   *
-   */
-   @deprecated("This method is deprecated and will be replaced by kafka.zk.AdminZkClient.", "1.1.0")
-  def changeClientIdConfig(zkUtils: ZkUtils, sanitizedClientId: String, configs: Properties) {
-    DynamicConfig.Client.validate(configs)
-    changeEntityConfig(zkUtils, ConfigType.Client, sanitizedClientId, configs)
-  }
-
-  /**
-   * Update the config for a  or  and create a change notification so the change will propagate to other brokers.
-   * User and/or clientId components of the path may be , indicating that the configuration is the default
-   * value to be applied if a more specific override is not configured.
-   *
-   * @param zkUtils Zookeeper utilities used to write the config to ZK
-   * @param sanitizedEntityName:  or /clients/
-   * @param configs: The final set of configs that will be applied to the topic. If any new configs need to be added or
-   *                 existing configs need to be deleted, it should be done prior to invoking this API
-   *
-   */
-   @deprecated("This method is deprecated and will be replaced by kafka.zk.AdminZkClient.", "1.1.0")
-  def changeUserOrUserClientIdConfig(zkUtils: ZkUtils, sanitizedEntityName: String, configs: Properties) {
-    if (sanitizedEntityName == ConfigEntityName.Default || sanitizedEntityName.contains("/clients"))
-      DynamicConfig.Client.validate(configs)
-    else
-      DynamicConfig.User.validate(configs)
-    changeEntityConfig(zkUtils, ConfigType.User, sanitizedEntityName, configs)
-  }
-
-  @deprecated("This method is deprecated and will be replaced by kafka.zk.AdminZkClient.", "1.1.0")
-  def validateTopicConfig(zkUtils: ZkUtils, topic: String, configs: Properties): Unit = {
-    Topic.validate(topic)
-    if (!topicExists(zkUtils, topic))
-      throw new AdminOperationException("Topic \"%s\" does not exist.".format(topic))
-    // remove the topic overrides
-    LogConfig.validate(configs)
-  }
-
-  /**
-   * Update the config for an existing topic and create a change notification so the change will propagate to other brokers
-   *
-   * @param zkUtils Zookeeper utilities used to write the config to ZK
-   * @param topic: The topic for which configs are being changed
-   * @param configs: The final set of configs that will be applied to the topic. If any new configs need to be added or
-   *                 existing configs need to be deleted, it should be done prior to invoking this API
-   *
-   */
-   @deprecated("This method is deprecated and will be replaced by kafka.zk.AdminZkClient.", "1.1.0")
-  def changeTopicConfig(zkUtils: ZkUtils, topic: String, configs: Properties) {
-    validateTopicConfig(zkUtils, topic, configs)
-    changeEntityConfig(zkUtils, ConfigType.Topic, topic, configs)
-  }
-
-  /**
-    * Override the broker config on some set of brokers. These overrides will be persisted between sessions, and will
-    * override any defaults entered in the broker's config files
-    *
-    * @param zkUtils: Zookeeper utilities used to write the config to ZK
-    * @param brokers: The list of brokers to apply config changes to
-    * @param configs: The config to change, as properties
-    */
-   @deprecated("This method is deprecated and will be replaced by kafka.zk.AdminZkClient.", "1.1.0")
-  def changeBrokerConfig(zkUtils: ZkUtils, brokers: Seq[Int], configs: Properties): Unit = {
-    DynamicConfig.Broker.validate(configs)
-    brokers.foreach { broker =>
-      changeEntityConfig(zkUtils, ConfigType.Broker, broker.toString, configs)
-    }
-  }
-
-  @deprecated("This method is deprecated and will be replaced by kafka.zk.AdminZkClient.", "1.1.0")
-  private def changeEntityConfig(zkUtils: ZkUtils, rootEntityType: String, fullSanitizedEntityName: String, configs: Properties) {
-    val sanitizedEntityPath = rootEntityType + '/' + fullSanitizedEntityName
-    val entityConfigPath = getEntityConfigPath(rootEntityType, fullSanitizedEntityName)
-    // write the new config--may not exist if there were previously no overrides
-    writeEntityConfig(zkUtils, entityConfigPath, configs)
-
-    // create the change notification
-    val seqNode = ZkUtils.ConfigChangesPath + "/" + EntityConfigChangeZnodePrefix
-    val content = Json.legacyEncodeAsString(getConfigChangeZnodeData(sanitizedEntityPath))
-    zkUtils.createSequentialPersistentPath(seqNode, content)
-  }
-
-  def getConfigChangeZnodeData(sanitizedEntityPath: String): Map[String, Any] = {
-    Map("version" -> 2, "entity_path" -> sanitizedEntityPath)
-  }
-
-  /**
-   * Write out the entity config to zk, if there is any
-   */
-  @deprecated("This method is deprecated and will be replaced by kafka.zk.AdminZkClient.", "1.1.0")
-  private def writeEntityConfig(zkUtils: ZkUtils, entityPath: String, config: Properties) {
-    val map = Map("version" -> 1, "config" -> config.asScala)
-    zkUtils.updatePersistentPath(entityPath, Json.legacyEncodeAsString(map))
-  }
-
-  /**
-   * Read the entity (topic, broker, client, user or ) config (if any) from zk
-   * sanitizedEntityName is , , ,  or /clients/.
-   */
-  @deprecated("This method is deprecated and will be replaced by kafka.zk.AdminZkClient.", "1.1.0")
-  def fetchEntityConfig(zkUtils: ZkUtils, rootEntityType: String, sanitizedEntityName: String): Properties = {
-    val entityConfigPath = getEntityConfigPath(rootEntityType, sanitizedEntityName)
-    // readDataMaybeNull returns Some(null) if the path exists, but there is no data
-    val str = zkUtils.readDataMaybeNull(entityConfigPath)._1.orNull
-    val props = new Properties()
-    if (str != null) {
-      Json.parseFull(str).foreach { jsValue =>
-        val jsObject = jsValue.asJsonObjectOption.getOrElse {
-          throw new IllegalArgumentException(s"Unexpected value in config: $str, entity_config_path: $entityConfigPath")
-        }
-        require(jsObject("version").to[Int] == 1)
-        val config = jsObject.get("config").flatMap(_.asJsonObjectOption).getOrElse {
-          throw new IllegalArgumentException(s"Invalid $entityConfigPath config: $str")
-        }
-        config.iterator.foreach { case (k, v) => props.setProperty(k, v.to[String]) }
-      }
-    }
-    props
-  }
-
-  @deprecated("This method is deprecated and will be replaced by kafka.zk.AdminZkClient.", "1.1.0")
-  def fetchAllTopicConfigs(zkUtils: ZkUtils): Map[String, Properties] =
-    zkUtils.getAllTopics().map(topic => (topic, fetchEntityConfig(zkUtils, ConfigType.Topic, topic))).toMap
-
-  @deprecated("This method is deprecated and will be replaced by kafka.zk.AdminZkClient.", "1.1.0")
-  def fetchAllEntityConfigs(zkUtils: ZkUtils, entityType: String): Map[String, Properties] =
-    zkUtils.getAllEntitiesWithConfig(entityType).map(entity => (entity, fetchEntityConfig(zkUtils, entityType, entity))).toMap
-
-  @deprecated("This method is deprecated and will be replaced by kafka.zk.AdminZkClient.", "1.1.0")
-  def fetchAllChildEntityConfigs(zkUtils: ZkUtils, rootEntityType: String, childEntityType: String): Map[String, Properties] = {
-    def entityPaths(zkUtils: ZkUtils, rootPath: Option[String]): Seq[String] = {
-      val root = rootPath match {
-        case Some(path) => rootEntityType + '/' + path
-        case None => rootEntityType
-      }
-      val entityNames = zkUtils.getAllEntitiesWithConfig(root)
-      rootPath match {
-        case Some(path) => entityNames.map(entityName => path + '/' + entityName)
-        case None => entityNames
-      }
-    }
-    entityPaths(zkUtils, None)
-      .flatMap(entity => entityPaths(zkUtils, Some(entity + '/' + childEntityType)))
-      .map(entityPath => (entityPath, fetchEntityConfig(zkUtils, rootEntityType, entityPath))).toMap
-  }
 
   private def replicaIndex(firstReplicaIndex: Int, secondReplicaShift: Int, replicaIndex: Int, nBrokers: Int): Int = {
     val shift = 1 + (secondReplicaShift + replicaIndex) % (nBrokers - 1)
diff --git a/core/src/main/scala/kafka/admin/BrokerApiVersionsCommand.scala b/core/src/main/scala/kafka/admin/BrokerApiVersionsCommand.scala
index 88acf7ffcea95..92cdb9ee02396 100644
--- a/core/src/main/scala/kafka/admin/BrokerApiVersionsCommand.scala
+++ b/core/src/main/scala/kafka/admin/BrokerApiVersionsCommand.scala
@@ -18,13 +18,32 @@
 package kafka.admin
 
 import java.io.PrintStream
+import java.io.IOException
 import java.util.Properties
+import java.util.concurrent.atomic.AtomicInteger
+import java.util.concurrent.{ConcurrentLinkedQueue, TimeUnit}
 
 import kafka.utils.{CommandDefaultOptions, CommandLineUtils}
-import org.apache.kafka.clients.CommonClientConfigs
+import kafka.utils.Logging
 import org.apache.kafka.common.utils.Utils
+import org.apache.kafka.clients.{ApiVersions, ClientDnsLookup, ClientResponse, ClientUtils, CommonClientConfigs, Metadata, NetworkClient, NodeApiVersions}
+import org.apache.kafka.clients.consumer.internals.{ConsumerNetworkClient, RequestFuture}
+import org.apache.kafka.common.config.ConfigDef.ValidString._
+import org.apache.kafka.common.config.ConfigDef.{Importance, Type}
+import org.apache.kafka.common.config.{AbstractConfig, ConfigDef}
+import org.apache.kafka.common.errors.AuthenticationException
+import org.apache.kafka.common.internals.ClusterResourceListeners
+import org.apache.kafka.common.metrics.Metrics
+import org.apache.kafka.common.network.Selector
+import org.apache.kafka.common.protocol.{ApiKeys, Errors}
+import org.apache.kafka.common.utils.LogContext
+import org.apache.kafka.common.utils.{KafkaThread, Time}
+import org.apache.kafka.common.Node
+import org.apache.kafka.common.message.ApiVersionsResponseData.ApiVersionsResponseKeyCollection
+import org.apache.kafka.common.requests.{AbstractRequest, AbstractResponse, ApiVersionsRequest, ApiVersionsResponse, MetadataRequest, MetadataResponse}
 
-import scala.util.{Failure, Success}
+import scala.collection.JavaConverters._
+import scala.util.{Failure, Success, Try}
 
 /**
  * A command for retrieving broker version information.
@@ -73,10 +92,234 @@ object BrokerApiVersionsCommand {
     options = parser.parse(args : _*)
     checkArgs()
 
-    def checkArgs() {
+    def checkArgs(): Unit = {
       CommandLineUtils.printHelpAndExitIfNeeded(this, "This tool helps to retrieve broker version information.")
       // check required args
       CommandLineUtils.checkRequiredArgs(parser, options, bootstrapServerOpt)
     }
   }
+
+  // org.apache.kafka.clients.admin.AdminClient doesn't currently expose a way to retrieve the supported api versions.
+  // We inline the bits we need from kafka.admin.AdminClient so that we can delete it.
+  private class AdminClient(val time: Time,
+                            val requestTimeoutMs: Int,
+                            val retryBackoffMs: Long,
+                            val client: ConsumerNetworkClient,
+                            val bootstrapBrokers: List[Node]) extends Logging {
+
+    @volatile var running: Boolean = true
+    val pendingFutures = new ConcurrentLinkedQueue[RequestFuture[ClientResponse]]()
+
+    val networkThread = new KafkaThread("admin-client-network-thread", new Runnable {
+      override def run(): Unit = {
+        try {
+          while (running)
+            client.poll(time.timer(Long.MaxValue))
+        } catch {
+          case t : Throwable =>
+            error("admin-client-network-thread exited", t)
+        } finally {
+          pendingFutures.asScala.foreach { future =>
+            try {
+              future.raise(Errors.UNKNOWN_SERVER_ERROR)
+            } catch {
+              case _: IllegalStateException => // It is OK if the future has been completed
+            }
+          }
+          pendingFutures.clear()
+        }
+      }
+    }, true)
+
+    networkThread.start()
+
+    private def send(target: Node,
+                     api: ApiKeys,
+                     request: AbstractRequest.Builder[_ <: AbstractRequest]): AbstractResponse = {
+      val future: RequestFuture[ClientResponse] = client.send(target, request)
+      pendingFutures.add(future)
+      future.awaitDone(Long.MaxValue, TimeUnit.MILLISECONDS)
+      pendingFutures.remove(future)
+      if (future.succeeded())
+        future.value().responseBody()
+      else
+        throw future.exception()
+    }
+
+    private def sendAnyNode(api: ApiKeys, request: AbstractRequest.Builder[_ <: AbstractRequest]): AbstractResponse = {
+      bootstrapBrokers.foreach { broker =>
+        try {
+          return send(broker, api, request)
+        } catch {
+          case e: AuthenticationException =>
+            throw e
+          case e: Exception =>
+            debug(s"Request $api failed against node $broker", e)
+        }
+      }
+      throw new RuntimeException(s"Request $api failed on brokers $bootstrapBrokers")
+    }
+
+    private def getApiVersions(node: Node): ApiVersionsResponseKeyCollection = {
+      val response = send(node, ApiKeys.API_VERSIONS, new ApiVersionsRequest.Builder()).asInstanceOf[ApiVersionsResponse]
+      Errors.forCode(response.data.errorCode()).maybeThrow()
+      response.data.apiKeys()
+    }
+
+    /**
+     * Wait until there is a non-empty list of brokers in the cluster.
+     */
+    def awaitBrokers(): Unit = {
+      var nodes = List[Node]()
+      do {
+        nodes = findAllBrokers()
+        if (nodes.isEmpty)
+          Thread.sleep(50)
+      } while (nodes.isEmpty)
+    }
+
+    private def findAllBrokers(): List[Node] = {
+      val request = MetadataRequest.Builder.allTopics()
+      val response = sendAnyNode(ApiKeys.METADATA, request).asInstanceOf[MetadataResponse]
+      val errors = response.errors
+      if (!errors.isEmpty)
+        debug(s"Metadata request contained errors: $errors")
+      response.cluster.nodes.asScala.toList
+    }
+
+    def listAllBrokerVersionInfo(): Map[Node, Try[NodeApiVersions]] =
+      findAllBrokers().map { broker =>
+        broker -> Try[NodeApiVersions](new NodeApiVersions(getApiVersions(broker)))
+      }.toMap
+
+    def close(): Unit = {
+      running = false
+      try {
+        client.close()
+      } catch {
+        case e: IOException =>
+          error("Exception closing nioSelector:", e)
+      }
+    }
+
+  }
+
+  private object AdminClient {
+    val DefaultConnectionMaxIdleMs = 9 * 60 * 1000
+    val DefaultRequestTimeoutMs = 5000
+    val DefaultMaxInFlightRequestsPerConnection = 100
+    val DefaultReconnectBackoffMs = 50
+    val DefaultReconnectBackoffMax = 50
+    val DefaultSendBufferBytes = 128 * 1024
+    val DefaultReceiveBufferBytes = 32 * 1024
+    val DefaultRetryBackoffMs = 100
+
+    val AdminClientIdSequence = new AtomicInteger(1)
+    val AdminConfigDef = {
+      val config = new ConfigDef()
+        .define(
+          CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG,
+          Type.LIST,
+          Importance.HIGH,
+          CommonClientConfigs.BOOTSTRAP_SERVERS_DOC)
+        .define(CommonClientConfigs.CLIENT_DNS_LOOKUP_CONFIG,
+          Type.STRING,
+          ClientDnsLookup.DEFAULT.toString,
+          in(ClientDnsLookup.DEFAULT.toString,
+            ClientDnsLookup.USE_ALL_DNS_IPS.toString,
+            ClientDnsLookup.RESOLVE_CANONICAL_BOOTSTRAP_SERVERS_ONLY.toString),
+          Importance.MEDIUM,
+          CommonClientConfigs.CLIENT_DNS_LOOKUP_DOC)
+        .define(
+          CommonClientConfigs.SECURITY_PROTOCOL_CONFIG,
+          ConfigDef.Type.STRING,
+          CommonClientConfigs.DEFAULT_SECURITY_PROTOCOL,
+          ConfigDef.Importance.MEDIUM,
+          CommonClientConfigs.SECURITY_PROTOCOL_DOC)
+        .define(
+          CommonClientConfigs.REQUEST_TIMEOUT_MS_CONFIG,
+          ConfigDef.Type.INT,
+          DefaultRequestTimeoutMs,
+          ConfigDef.Importance.MEDIUM,
+          CommonClientConfigs.REQUEST_TIMEOUT_MS_DOC)
+        .define(
+          CommonClientConfigs.RETRY_BACKOFF_MS_CONFIG,
+          ConfigDef.Type.LONG,
+          DefaultRetryBackoffMs,
+          ConfigDef.Importance.MEDIUM,
+          CommonClientConfigs.RETRY_BACKOFF_MS_DOC)
+        .withClientSslSupport()
+        .withClientSaslSupport()
+      config
+    }
+
+    class AdminConfig(originals: Map[_,_]) extends AbstractConfig(AdminConfigDef, originals.asJava, false)
+
+    def createSimplePlaintext(brokerUrl: String): AdminClient = {
+      val config = Map(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG -> brokerUrl)
+      create(new AdminConfig(config))
+    }
+
+    def create(props: Properties): AdminClient = create(props.asScala.toMap)
+
+    def create(props: Map[String, _]): AdminClient = create(new AdminConfig(props))
+
+    def create(config: AdminConfig): AdminClient = {
+      val clientId = "admin-" + AdminClientIdSequence.getAndIncrement()
+      val logContext = new LogContext(s"[LegacyAdminClient clientId=$clientId] ")
+      val time = Time.SYSTEM
+      val metrics = new Metrics(time)
+      val metadata = new Metadata(100L, 60 * 60 * 1000L, logContext,
+        new ClusterResourceListeners)
+      val channelBuilder = ClientUtils.createChannelBuilder(config, time)
+      val requestTimeoutMs = config.getInt(CommonClientConfigs.REQUEST_TIMEOUT_MS_CONFIG)
+      val retryBackoffMs = config.getLong(CommonClientConfigs.RETRY_BACKOFF_MS_CONFIG)
+
+      val brokerUrls = config.getList(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG)
+      val clientDnsLookup = config.getString(CommonClientConfigs.CLIENT_DNS_LOOKUP_CONFIG)
+      val brokerAddresses = ClientUtils.parseAndValidateAddresses(brokerUrls, clientDnsLookup)
+      metadata.bootstrap(brokerAddresses)
+
+      val selector = new Selector(
+        DefaultConnectionMaxIdleMs,
+        metrics,
+        time,
+        "admin",
+        channelBuilder,
+        logContext)
+
+      val networkClient = new NetworkClient(
+        selector,
+        metadata,
+        clientId,
+        DefaultMaxInFlightRequestsPerConnection,
+        DefaultReconnectBackoffMs,
+        DefaultReconnectBackoffMax,
+        DefaultSendBufferBytes,
+        DefaultReceiveBufferBytes,
+        requestTimeoutMs,
+        ClientDnsLookup.DEFAULT,
+        time,
+        true,
+        new ApiVersions,
+        logContext)
+
+      val highLevelClient = new ConsumerNetworkClient(
+        logContext,
+        networkClient,
+        metadata,
+        time,
+        retryBackoffMs,
+        requestTimeoutMs,
+        Integer.MAX_VALUE)
+
+      new AdminClient(
+        time,
+        requestTimeoutMs,
+        retryBackoffMs,
+        highLevelClient,
+        metadata.fetch.nodes.asScala.toList)
+    }
+  }
+
 }
diff --git a/core/src/main/scala/kafka/admin/ConfigCommand.scala b/core/src/main/scala/kafka/admin/ConfigCommand.scala
index f1aa9ff5bde0f..9203334e4615e 100644
--- a/core/src/main/scala/kafka/admin/ConfigCommand.scala
+++ b/core/src/main/scala/kafka/admin/ConfigCommand.scala
@@ -28,7 +28,7 @@ import kafka.utils.{CommandDefaultOptions, CommandLineUtils, Exit, PasswordEncod
 import kafka.utils.Implicits._
 import kafka.zk.{AdminZkClient, KafkaZkClient}
 import org.apache.kafka.clients.CommonClientConfigs
-import org.apache.kafka.clients.admin.{AlterConfigsOptions, ConfigEntry, DescribeConfigsOptions, AdminClient => JAdminClient, Config => JConfig}
+import org.apache.kafka.clients.admin.{Admin, AlterConfigOp, AlterConfigsOptions, ConfigEntry, DescribeConfigsOptions, AdminClient => JAdminClient, Config => JConfig}
 import org.apache.kafka.common.config.ConfigResource
 import org.apache.kafka.common.config.types.Password
 import org.apache.kafka.common.errors.InvalidConfigurationException
@@ -41,7 +41,6 @@ import scala.collection._
 
 
 /**
- * This script can be used to change configs for topics/clients/brokers dynamically
  * This script can be used to change configs for topics/clients/users/brokers dynamically
  * An entity described or altered by the command may be one of:
  * 
        
@@ -50,12 +49,15 @@ import scala.collection._
  *     
      •  user: --entity-type users --entity-name 
  *     
      •  : --entity-type users --entity-name  --entity-type clients --entity-name 
  *     
      •  broker: --entity-type brokers --entity-name 
+ *     
      •  broker-logger: --entity-type broker-loggers --entity-name 
  * 
      
  * --entity-default may be used instead of --entity-name when describing or altering default configuration for users and clients.
  *
  */
 object ConfigCommand extends Config {
 
+  val BrokerLoggerConfigType = "broker-loggers"
+  val BrokerSupportedConfigTypes = Seq(ConfigType.Broker, BrokerLoggerConfigType)
   val DefaultScramIterations = 4096
   // Dynamic broker configs can only be updated using the new AdminClient once brokers have started
   // so that configs may be fully validated. Prior to starting brokers, updates may be performed using
@@ -88,6 +90,7 @@ object ConfigCommand extends Config {
         Exit.exit(1)
 
       case t: Throwable =>
+        logger.debug(s"Error while executing config command with args '${args.mkString(" ")}'", t)
         System.err.println(s"Error while executing config command with args '${args.mkString(" ")}'")
         t.printStackTrace(System.err)
         Exit.exit(1)
@@ -108,7 +111,7 @@ object ConfigCommand extends Config {
     }
   }
 
-  private[admin] def alterConfig(zkClient: KafkaZkClient, opts: ConfigCommandOptions, adminZkClient: AdminZkClient) {
+  private[admin] def alterConfig(zkClient: KafkaZkClient, opts: ConfigCommandOptions, adminZkClient: AdminZkClient): Unit = {
     val configsToBeAdded = parseConfigsToBeAdded(opts)
     val configsToBeDeleted = parseConfigsToBeDeleted(opts)
     val entity = parseEntity(opts)
@@ -125,7 +128,7 @@ object ConfigCommand extends Config {
       if (dynamicBrokerConfigs.nonEmpty) {
         val perBrokerConfig = entityName != ConfigEntityName.Default
         val errorMessage = s"--bootstrap-server option must be specified to update broker configs $dynamicBrokerConfigs."
-        val info = "Broker configuraton updates using ZooKeeper are supported for bootstrapping before brokers" +
+        val info = "Broker configuration updates using ZooKeeper are supported for bootstrapping before brokers" +
           " are started to enable encrypted password configs to be stored in ZooKeeper."
         if (perBrokerConfig) {
           adminZkClient.parseBroker(entityName).foreach { brokerId =>
@@ -154,7 +157,7 @@ object ConfigCommand extends Config {
     println(s"Completed Updating config for entity: $entity.")
   }
 
-  private def preProcessScramCredentials(configsToBeAdded: Properties) {
+  private def preProcessScramCredentials(configsToBeAdded: Properties): Unit = {
     def scramCredential(mechanism: ScramMechanism, credentialStr: String): String = {
       val pattern = "(?:iterations=([0-9]*),)?password=(.*)".r
       val (iterations, password) = credentialStr match {
@@ -191,9 +194,9 @@ object ConfigCommand extends Config {
    * Password configs are encrypted using the secret `KafkaConfig.PasswordEncoderSecretProp`.
    * The secret is removed from `configsToBeAdded` and will not be persisted in ZooKeeper.
    */
-  private def preProcessBrokerConfigs(configsToBeAdded: Properties, perBrokerConfig: Boolean) {
+  private def preProcessBrokerConfigs(configsToBeAdded: Properties, perBrokerConfig: Boolean): Unit = {
     val passwordEncoderConfigs = new Properties
-    passwordEncoderConfigs ++= configsToBeAdded.asScala.filterKeys(_.startsWith("password.encoder."))
+    passwordEncoderConfigs ++= configsToBeAdded.asScala.filter { case (key, _) => key.startsWith("password.encoder.") }
     if (!passwordEncoderConfigs.isEmpty) {
       info(s"Password encoder configs ${passwordEncoderConfigs.keySet} will be used for encrypting" +
         " passwords, but will not be stored in ZooKeeper.")
@@ -218,7 +221,7 @@ object ConfigCommand extends Config {
     }
   }
 
-  private def describeConfig(zkClient: KafkaZkClient, opts: ConfigCommandOptions, adminZkClient: AdminZkClient) {
+  private def describeConfig(zkClient: KafkaZkClient, opts: ConfigCommandOptions, adminZkClient: AdminZkClient): Unit = {
     val configEntity = parseEntity(opts)
     val describeAllUsers = configEntity.root.entityType == ConfigType.User && !configEntity.root.sanitizedName.isDefined && !configEntity.child.isDefined
     val entities = configEntity.getAllEntities(zkClient)
@@ -273,49 +276,61 @@ object ConfigCommand extends Config {
     val adminClient = JAdminClient.create(props)
     val entityName = if (opts.options.has(opts.entityName))
       opts.options.valueOf(opts.entityName)
-    else if (opts.options.has(opts.entityDefault))
+    else // default entity
       ""
-    else
-      throw new IllegalArgumentException("At least one of --entity-name or --entity-default must be specified with --bootstrap-server")
 
     val entityTypes = opts.options.valuesOf(opts.entityType).asScala
     if (entityTypes.size != 1)
-      throw new IllegalArgumentException("Exactly one --entity-type must be specified with --bootstrap-server")
-    if (entityTypes.head != ConfigType.Broker)
-      throw new IllegalArgumentException(s"--zookeeper option must be specified for entity-type $entityTypes")
+      throw new IllegalArgumentException(s"Exactly one --entity-type (out of ${BrokerSupportedConfigTypes.mkString(",")}) must be specified with --bootstrap-server")
 
     try {
       if (opts.options.has(opts.alterOpt))
-        alterBrokerConfig(adminClient, opts, entityName)
+        alterBrokerConfig(adminClient, opts, entityTypes.head, entityName)
       else if (opts.options.has(opts.describeOpt))
-        describeBrokerConfig(adminClient, opts, entityName)
+        describeBrokerConfig(adminClient, opts, entityTypes.head, entityName)
     } finally {
       adminClient.close()
     }
   }
 
-  private[admin] def alterBrokerConfig(adminClient: JAdminClient, opts: ConfigCommandOptions, entityName: String) {
+  private[admin] def alterBrokerConfig(adminClient: Admin, opts: ConfigCommandOptions,
+                                       entityType: String, entityName: String): Unit = {
     val configsToBeAdded = parseConfigsToBeAdded(opts).asScala.map { case (k, v) => (k, new ConfigEntry(k, v)) }
     val configsToBeDeleted = parseConfigsToBeDeleted(opts)
 
-    // compile the final set of configs
-    val configResource = new ConfigResource(ConfigResource.Type.BROKER, entityName)
-    val oldConfig = brokerConfig(adminClient, entityName, includeSynonyms = false)
+    if (entityType == ConfigType.Broker) {
+      val configResource = new ConfigResource(ConfigResource.Type.BROKER, entityName)
+      val oldConfig = brokerConfig(adminClient, entityName, includeSynonyms = false)
         .map { entry => (entry.name, entry) }.toMap
 
-    // fail the command if any of the configs to be deleted does not exist
-    val invalidConfigs = configsToBeDeleted.filterNot(oldConfig.contains)
-    if (invalidConfigs.nonEmpty)
-      throw new InvalidConfigurationException(s"Invalid config(s): ${invalidConfigs.mkString(",")}")
-
-    val newEntries = oldConfig ++ configsToBeAdded -- configsToBeDeleted
-    val sensitiveEntries = newEntries.filter(_._2.value == null)
-    if (sensitiveEntries.nonEmpty)
-      throw new InvalidConfigurationException(s"All sensitive broker config entries must be specified for --alter, missing entries: ${sensitiveEntries.keySet}")
-    val newConfig = new JConfig(newEntries.asJava.values)
-
-    val alterOptions = new AlterConfigsOptions().timeoutMs(30000).validateOnly(false)
-    adminClient.alterConfigs(Map(configResource -> newConfig).asJava, alterOptions).all().get(60, TimeUnit.SECONDS)
+      // fail the command if any of the configs to be deleted does not exist
+      val invalidConfigs = configsToBeDeleted.filterNot(oldConfig.contains)
+      if (invalidConfigs.nonEmpty)
+        throw new InvalidConfigurationException(s"Invalid config(s): ${invalidConfigs.mkString(",")}")
+
+      val newEntries = oldConfig ++ configsToBeAdded -- configsToBeDeleted
+      val sensitiveEntries = newEntries.filter(_._2.value == null)
+      if (sensitiveEntries.nonEmpty)
+        throw new InvalidConfigurationException(s"All sensitive broker config entries must be specified for --alter, missing entries: ${sensitiveEntries.keySet}")
+      val newConfig = new JConfig(newEntries.asJava.values)
+
+      val alterOptions = new AlterConfigsOptions().timeoutMs(30000).validateOnly(false)
+      adminClient.alterConfigs(Map(configResource -> newConfig).asJava, alterOptions).all().get(60, TimeUnit.SECONDS)
+    } else if (entityType == BrokerLoggerConfigType) {
+      val configResource = new ConfigResource(ConfigResource.Type.BROKER_LOGGER, entityName)
+      val validLoggers = brokerLoggerConfigs(adminClient, entityName).map(_.name)
+      // fail the command if any of the configured broker loggers do not exist
+      val invalidBrokerLoggers = configsToBeDeleted.filterNot(validLoggers.contains) ++ configsToBeAdded.keys.filterNot(validLoggers.contains)
+      if (invalidBrokerLoggers.nonEmpty)
+        throw new InvalidConfigurationException(s"Invalid broker logger(s): ${invalidBrokerLoggers.mkString(",")}")
+
+      val alterOptions = new AlterConfigsOptions().timeoutMs(30000).validateOnly(false)
+      val alterLogLevelEntries = (configsToBeAdded.values.map(new AlterConfigOp(_, AlterConfigOp.OpType.SET))
+        ++ configsToBeDeleted.map { k => new AlterConfigOp(new ConfigEntry(k, ""), AlterConfigOp.OpType.DELETE) }
+      ).asJavaCollection
+
+      adminClient.incrementalAlterConfigs(Map(configResource -> alterLogLevelEntries).asJava, alterOptions).all().get(60, TimeUnit.SECONDS)
+    }
 
     if (entityName.nonEmpty)
       println(s"Completed updating config for broker: $entityName.")
@@ -323,8 +338,13 @@ object ConfigCommand extends Config {
       println(s"Completed updating default config for brokers in the cluster,")
   }
 
-  private def describeBrokerConfig(adminClient: JAdminClient, opts: ConfigCommandOptions, entityName: String) {
-    val configs = brokerConfig(adminClient, entityName, includeSynonyms = true)
+  private def describeBrokerConfig(adminClient: Admin, opts: ConfigCommandOptions,
+                                   entityType: String, entityName: String): Unit = {
+    val configs = if (entityType == ConfigType.Broker)
+      brokerConfig(adminClient, entityName, includeSynonyms = true)
+    else // broker logger
+      brokerLoggerConfigs(adminClient, entityName)
+
     if (entityName.nonEmpty)
       println(s"Configs for broker $entityName are:")
     else
@@ -335,7 +355,7 @@ object ConfigCommand extends Config {
     }
   }
 
-  private def brokerConfig(adminClient: JAdminClient, entityName: String, includeSynonyms: Boolean): Seq[ConfigEntry] = {
+  private def brokerConfig(adminClient: Admin, entityName: String, includeSynonyms: Boolean): Seq[ConfigEntry] = {
     val configResource = new ConfigResource(ConfigResource.Type.BROKER, entityName)
     val configSource = if (!entityName.isEmpty)
       ConfigEntry.ConfigSource.DYNAMIC_BROKER_CONFIG
@@ -348,6 +368,15 @@ object ConfigCommand extends Config {
       .toSeq
   }
 
+  /**
+    * Returns all the valid broker logger configurations
+    */
+  private def brokerLoggerConfigs(adminClient: Admin, entityName: String): Seq[ConfigEntry] = {
+    val configResource = new ConfigResource(ConfigResource.Type.BROKER_LOGGER, entityName)
+    val configs = adminClient.describeConfigs(Collections.singleton(configResource)).all.get(30, TimeUnit.SECONDS)
+    configs.get(configResource).entries.asScala.toSeq
+  }
+
   case class Entity(entityType: String, sanitizedName: Option[String]) {
     val entityPath = sanitizedName match {
       case Some(n) => entityType + "/" + n
@@ -444,7 +473,7 @@ object ConfigCommand extends Config {
     if (opts.options.has(opts.alterOpt) && names.size != types.size)
       throw new IllegalArgumentException("--entity-name or --entity-default must be specified with each --entity-type for --alter")
 
-    val reverse = types.size == 2 && types(0) == ConfigType.Client
+    val reverse = types.size == 2 && types.head == ConfigType.Client
     val entityTypes = if (reverse) types.reverse else types
     val sortedNames = (if (reverse && names.length == 2) names.reverse else names).iterator
 
@@ -482,7 +511,7 @@ object ConfigCommand extends Config {
       .ofType(classOf[String])
     val alterOpt = parser.accepts("alter", "Alter the configuration for the entity.")
     val describeOpt = parser.accepts("describe", "List configs for the given entity.")
-    val entityType = parser.accepts("entity-type", "Type of entity (topics/clients/users/brokers)")
+    val entityType = parser.accepts("entity-type", "Type of entity (topics/clients/users/brokers/broker-loggers)")
             .withRequiredArg
             .ofType(classOf[String])
     val entityName = parser.accepts("entity-name", "Name of entity (topic name/client id/user principal name/broker id)")
@@ -493,9 +522,9 @@ object ConfigCommand extends Config {
     val nl = System.getProperty("line.separator")
     val addConfig = parser.accepts("add-config", "Key Value pairs of configs to add. Square brackets can be used to group values which contain commas: 'k1=v1,k2=[v1,v2,v2],k3=v3'. The following is a list of valid configurations: " +
             "For entity-type '" + ConfigType.Topic + "': " + LogConfig.configNames.map("\t" + _).mkString(nl, nl, nl) +
-            "For entity-type '" + ConfigType.Broker + "': " + DynamicConfig.Broker.names.asScala.map("\t" + _).mkString(nl, nl, nl) +
-            "For entity-type '" + ConfigType.User + "': " + DynamicConfig.User.names.asScala.map("\t" + _).mkString(nl, nl, nl) +
-            "For entity-type '" + ConfigType.Client + "': " + DynamicConfig.Client.names.asScala.map("\t" + _).mkString(nl, nl, nl) +
+            "For entity-type '" + ConfigType.Broker + "': " + DynamicConfig.Broker.names.asScala.toSeq.sorted.map("\t" + _).mkString(nl, nl, nl) +
+            "For entity-type '" + ConfigType.User + "': " + DynamicConfig.User.names.asScala.toSeq.sorted.map("\t" + _).mkString(nl, nl, nl) +
+            "For entity-type '" + ConfigType.Client + "': " + DynamicConfig.Client.names.asScala.toSeq.sorted.map("\t" + _).mkString(nl, nl, nl) +
             s"Entity types '${ConfigType.User}' and '${ConfigType.Client}' may be specified together to update config for clients of a specific user.")
             .withRequiredArg
             .ofType(classOf[String])
@@ -508,41 +537,63 @@ object ConfigCommand extends Config {
 
     val allOpts: Set[OptionSpec[_]] = Set(alterOpt, describeOpt, entityType, entityName, addConfig, deleteConfig, helpOpt)
 
-    def checkArgs() {
+    def checkArgs(): Unit = {
       // should have exactly one action
       val actions = Seq(alterOpt, describeOpt).count(options.has _)
       if(actions != 1)
         CommandLineUtils.printUsageAndDie(parser, "Command must include exactly one action: --describe, --alter")
-
       // check required args
       CommandLineUtils.checkInvalidArgs(parser, options, alterOpt, Set(describeOpt))
       CommandLineUtils.checkInvalidArgs(parser, options, describeOpt, Set(alterOpt, addConfig, deleteConfig))
+
       val entityTypeVals = options.valuesOf(entityType).asScala
+      val (allowedEntityTypes, connectOptString) = if (options.has(bootstrapServerOpt))
+        (BrokerSupportedConfigTypes, "--bootstrap-server")
+      else
+        (ConfigType.all, "--zookeeper")
+
+      entityTypeVals.foreach(entityTypeVal =>
+        if (!allowedEntityTypes.contains(entityTypeVal))
+          throw new IllegalArgumentException(s"Invalid entity-type $entityTypeVal, --entity-type must be one of ${allowedEntityTypes.mkString(",")} with the $connectOptString argument")
+      )
+      if (entityTypeVals.isEmpty)
+        throw new IllegalArgumentException("At least one --entity-type must be specified")
+      else if (entityTypeVals.size > 1 && !entityTypeVals.toSet.equals(Set(ConfigType.User, ConfigType.Client)))
+        throw new IllegalArgumentException(s"Only '${ConfigType.User}' and '${ConfigType.Client}' entity types may be specified together")
 
-      if (options.has(bootstrapServerOpt) == options.has(zkConnectOpt))
+      if (!options.has(bootstrapServerOpt) && !options.has(zkConnectOpt))
+        throw new IllegalArgumentException("One of the required --bootstrap-server or --zookeeper arguments must be specified")
+      else if (options.has(bootstrapServerOpt) && options.has(zkConnectOpt))
         throw new IllegalArgumentException("Only one of --bootstrap-server or --zookeeper must be specified")
+      else if (options.has(bootstrapServerOpt) && !options.has(entityName) && !options.has(entityDefault))
+        throw new IllegalArgumentException(s"At least one of --entity-name or --entity-default must be specified with --bootstrap-server")
+
+      if (options.has(entityName) && (entityTypeVals.contains(ConfigType.Broker) || entityTypeVals.contains(BrokerLoggerConfigType))) {
+        val brokerId = options.valueOf(entityName)
+        try brokerId.toInt catch {
+          case _: NumberFormatException =>
+            throw new IllegalArgumentException(s"The entity name for ${entityTypeVals.head} must be a valid integer broker id , but it is: $brokerId")
+        }
+      }
+
       if (entityTypeVals.contains(ConfigType.Client) || entityTypeVals.contains(ConfigType.Topic) || entityTypeVals.contains(ConfigType.User))
         CommandLineUtils.checkRequiredArgs(parser, options, zkConnectOpt, entityType)
-      if(options.has(alterOpt)) {
+
+      if (options.has(describeOpt) && entityTypeVals.contains(BrokerLoggerConfigType) && !options.has(entityName))
+        throw new IllegalArgumentException(s"--entity-name must be specified with --describe of ${entityTypeVals.mkString(",")}")
+
+      if (options.has(alterOpt)) {
         if (entityTypeVals.contains(ConfigType.User) || entityTypeVals.contains(ConfigType.Client) || entityTypeVals.contains(ConfigType.Broker)) {
           if (!options.has(entityName) && !options.has(entityDefault))
             throw new IllegalArgumentException("--entity-name or --entity-default must be specified with --alter of users, clients or brokers")
         } else if (!options.has(entityName))
-            throw new IllegalArgumentException(s"--entity-name must be specified with --alter of ${entityTypeVals}")
+          throw new IllegalArgumentException(s"--entity-name must be specified with --alter of ${entityTypeVals.mkString(",")}")
 
         val isAddConfigPresent: Boolean = options.has(addConfig)
         val isDeleteConfigPresent: Boolean = options.has(deleteConfig)
         if(! isAddConfigPresent && ! isDeleteConfigPresent)
           throw new IllegalArgumentException("At least one of --add-config or --delete-config must be specified with --alter")
       }
-      entityTypeVals.foreach(entityTypeVal =>
-        if (!ConfigType.all.contains(entityTypeVal))
-          throw new IllegalArgumentException(s"Invalid entity-type ${entityTypeVal}, --entity-type must be one of ${ConfigType.all}")
-      )
-      if (entityTypeVals.isEmpty)
-        throw new IllegalArgumentException("At least one --entity-type must be specified")
-      else if (entityTypeVals.size > 1 && !entityTypeVals.toSet.equals(Set(ConfigType.User, ConfigType.Client)))
-        throw new IllegalArgumentException(s"Only '${ConfigType.User}' and '${ConfigType.Client}' entity types may be specified together")
     }
   }
 }
diff --git a/core/src/main/scala/kafka/admin/ConsumerGroupCommand.scala b/core/src/main/scala/kafka/admin/ConsumerGroupCommand.scala
index b40a33b980d87..eeb93d65c6616 100755
--- a/core/src/main/scala/kafka/admin/ConsumerGroupCommand.scala
+++ b/core/src/main/scala/kafka/admin/ConsumerGroupCommand.scala
@@ -19,10 +19,12 @@ package kafka.admin
 
 import java.text.{ParseException, SimpleDateFormat}
 import java.time.{Duration, Instant}
-import java.util
 import java.util.Properties
+import java.util.concurrent.ExecutionException
 
-import joptsimple.OptionSpec
+import com.fasterxml.jackson.dataformat.csv.CsvMapper
+import com.fasterxml.jackson.module.scala.DefaultScalaModule
+import com.fasterxml.jackson.module.scala.experimental.ScalaObjectMapper
 import kafka.utils._
 import org.apache.kafka.clients.admin._
 import org.apache.kafka.clients.consumer.{ConsumerConfig, KafkaConsumer, OffsetAndMetadata}
@@ -33,20 +35,25 @@ import org.apache.kafka.common.{KafkaException, Node, TopicPartition}
 
 import scala.collection.JavaConverters._
 import scala.collection.mutable.ListBuffer
-import scala.collection.{Seq, Set}
+import scala.collection.{Map, Seq, Set, immutable, mutable}
 import scala.util.{Failure, Success, Try}
+import joptsimple.OptionSpec
+import org.apache.kafka.common.protocol.Errors
+
+import scala.collection.immutable.TreeMap
+import scala.reflect.ClassTag
 
 object ConsumerGroupCommand extends Logging {
 
-  def main(args: Array[String]) {
+  def main(args: Array[String]): Unit = {
     val opts = new ConsumerGroupCommandOptions(args)
 
     CommandLineUtils.printHelpAndExitIfNeeded(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).count(opts.options.has)
+    val actions = Seq(opts.listOpt, opts.describeOpt, opts.deleteOpt, opts.resetOffsetsOpt, opts.deleteOffsetsOpt).count(opts.options.has)
     if (actions != 1)
-      CommandLineUtils.printUsageAndDie(opts.parser, "Command must include exactly one action: --list, --describe, --delete, --reset-offsets")
+      CommandLineUtils.printUsageAndDie(opts.parser, "Command must include exactly one action: --list, --describe, --delete, --reset-offsets, --delete-offsets")
 
     opts.checkArgs()
 
@@ -56,17 +63,20 @@ object ConsumerGroupCommand extends Logging {
       if (opts.options.has(opts.listOpt))
         consumerGroupService.listGroups().foreach(println(_))
       else if (opts.options.has(opts.describeOpt))
-        consumerGroupService.describeGroup()
+        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.exportOffsetsToReset(offsetsToReset)
+          val exported = consumerGroupService.exportOffsetsToCsv(offsetsToReset)
           println(exported)
         } else
           printOffsetsToReset(offsetsToReset)
       }
+      else if (opts.options.has(opts.deleteOffsetsOpt)) {
+        consumerGroupService.deleteOffsets()
+      }
     } catch {
       case e: Throwable =>
         printError(s"Executing consumer group command failed due to ${e.getMessage}", Some(e))
@@ -78,7 +88,7 @@ object ConsumerGroupCommand extends Logging {
   val MISSING_COLUMN_VALUE = "-"
 
   def printError(msg: String, e: Option[Throwable] = None): Unit = {
-    println(s"Error: $msg")
+    println(s"\nError: $msg")
     e.foreach(_.printStackTrace())
   }
 
@@ -95,15 +105,18 @@ object ConsumerGroupCommand extends Logging {
     date.getTime
   }
 
-  def printOffsetsToReset(groupAssignmentsToReset: Map[TopicPartition, OffsetAndMetadata]): Unit = {
-    println("\n%-30s %-10s %-15s".format("TOPIC", "PARTITION", "NEW-OFFSET"))
-
-    groupAssignmentsToReset.foreach {
-      case (consumerAssignment, offsetAndMetadata) =>
-        println("%-30s %-10s %-15s".format(
-          consumerAssignment.topic,
-          consumerAssignment.partition,
-          offsetAndMetadata.offset))
+  def printOffsetsToReset(groupAssignmentsToReset: Map[String, Map[TopicPartition, OffsetAndMetadata]]): Unit = {
+    if (groupAssignmentsToReset.nonEmpty)
+      println("\n%-30s %-30s %-10s %-15s".format("GROUP", "TOPIC", "PARTITION", "NEW-OFFSET"))
+    for {
+      (groupId, assignment) <- groupAssignmentsToReset
+      (consumerAssignment, offsetAndMetadata) <- assignment
+    } {
+      println("%-30s %-30s %-10s %-15s".format(
+        groupId,
+        consumerAssignment.topic,
+        consumerAssignment.partition,
+        offsetAndMetadata.offset))
     }
   }
 
@@ -117,19 +130,62 @@ object ConsumerGroupCommand extends Logging {
 
   private[admin] case class GroupState(group: String, coordinator: Node, assignmentStrategy: String, state: String, numMembers: Int)
 
-  class ConsumerGroupService(val opts: ConsumerGroupCommandOptions) {
+  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("group", "topic", "partition", "offset")
+  }
+  private[admin] object CsvRecordNoGroup {
+    val fields = Array("topic", "partition", "offset")
+  }
+  // Example: CsvUtils().readerFor[CsvRecordWithoutGroup]
+  private[admin] case class CsvUtils() {
+    val mapper = new CsvMapper with ScalaObjectMapper
+    mapper.registerModule(DefaultScalaModule)
+    def readerFor[T <: CsvRecord: ClassTag] = {
+      val schema = getSchema[T]
+      val clazz = implicitly[ClassTag[T]].runtimeClass
+      mapper.readerFor(clazz).`with`(schema)
+    }
+    def writerFor[T <: CsvRecord: ClassTag] = {
+      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 = clazz match {
+        case _ if classOf[CsvRecordWithGroup] == clazz => CsvRecordWithGroup.fields
+        case _ if classOf[CsvRecordNoGroup]   == clazz => CsvRecordNoGroup.fields
+      }
+      val schema = mapper.schemaFor(clazz).sortedBy(fields: _*)
+      schema
+    }
+  }
 
-    private val adminClient = createAdminClient()
+  class ConsumerGroupService(val opts: ConsumerGroupCommandOptions,
+                             private[admin] val configOverrides: Map[String, String] = Map.empty) {
 
-    // `consumer` is only needed for `describe`, so we instantiate it lazily
-    private var consumer: KafkaConsumer[String, String] = _
+    private val adminClient = createAdminClient(configOverrides)
 
-    def listGroups(): List[String] = {
-      val result = adminClient.listConsumerGroups(
-        withTimeoutMs(new ListConsumerGroupsOptions))
+    // `consumers` are only needed for `describe`, so we instantiate them lazily
+    private lazy val consumers: mutable.Map[String, KafkaConsumer[String, String]] = mutable.Map.empty
 
-        val listings = result.all.get.asScala
-        listings.map(_.groupId).toList
+    // 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(): List[String] = {
+      val result = adminClient.listConsumerGroups(withTimeoutMs(new ListConsumerGroupsOptions))
+      val listings = result.all.get.asScala
+      listings.map(_.groupId).toList
     }
 
     private def shouldPrintMemberState(group: String, state: Option[String], numRows: Option[Int]): Boolean = {
@@ -143,9 +199,9 @@ object ConsumerGroupCommand extends Logging {
           case Some("Dead") =>
             printError(s"Consumer group '$group' does not exist.")
           case Some("Empty") =>
-            Console.err.println(s"Consumer group '$group' has no active members.")
+            Console.err.println(s"\nConsumer group '$group' has no active members.")
           case Some("PreparingRebalance") | Some("CompletingRebalance") =>
-            Console.err.println(s"Warning: Consumer group '$group' is rebalancing.")
+            Console.err.println(s"\nWarning: Consumer group '$group' is rebalancing.")
           case Some("Stable") =>
           case other =>
             // the control should never reach here
@@ -157,113 +213,127 @@ object ConsumerGroupCommand extends Logging {
 
     private def size(colOpt: Option[Seq[Object]]): Option[Int] = colOpt.map(_.size)
 
-    private def printOffsets(group: String, state: Option[String], assignments: Option[Seq[PartitionAssignmentState]]): Unit = {
-      if (shouldPrintMemberState(group, state, size(assignments))) {
-        // find proper columns width
-        var (maxTopicLen, maxConsumerIdLen, maxHostLen) = (15, 15, 15)
-        assignments match {
-          case None => // do nothing
-          case Some(consumerAssignments) =>
-            consumerAssignments.foreach { consumerAssignment =>
-              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)
-            }
-        }
+    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)
+              }
+          }
 
-        println(s"\n%${-maxTopicLen}s %-10s %-15s %-15s %-15s %${-maxConsumerIdLen}s %${-maxHostLen}s %s"
-          .format("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(s"%-${maxTopicLen}s %-10s %-15s %-15s %-15s %-${maxConsumerIdLen}s %-${maxHostLen}s %s".format(
-                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)))
-            }
+          println(s"\n%${-maxGroupLen}s %${-maxTopicLen}s %-10s %-15s %-15s %-15s %${-maxConsumerIdLen}s %${-maxHostLen}s %s"
+            .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(s"%${-maxGroupLen}s %${-maxTopicLen}s %-10s %-15s %-15s %-15s %${-maxConsumerIdLen}s %${-maxHostLen}s %s".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(group: String, state: Option[String], assignments: Option[Seq[MemberAssignmentState]], verbose: Boolean): Unit = {
-      if (shouldPrintMemberState(group, state, size(assignments))) {
-        // find proper columns width
-        var (maxConsumerIdLen, maxHostLen, maxClientIdLen) = (15, 15, 15)
-        assignments match {
-          case None => // do nothing
-          case Some(memberAssignments) =>
-            memberAssignments.foreach { memberAssignment =>
-              maxConsumerIdLen = Math.max(maxConsumerIdLen, memberAssignment.consumerId.length)
-              maxHostLen = Math.max(maxHostLen, memberAssignment.host.length)
-              maxClientIdLen = Math.max(maxClientIdLen, memberAssignment.clientId.length)
-            }
-        }
+    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, maxHostLen, maxClientIdLen) = (15, 15, 15, 15)
+          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)
+                maxHostLen = Math.max(maxHostLen, memberAssignment.host.length)
+                maxClientIdLen = Math.max(maxClientIdLen, memberAssignment.clientId.length)
+              }
+          }
 
-        print(s"\n%${-maxConsumerIdLen}s %${-maxHostLen}s %${-maxClientIdLen}s %-15s "
-          .format("CONSUMER-ID", "HOST", "CLIENT-ID", "#PARTITIONS"))
-        if (verbose)
-          print("%s".format("ASSIGNMENT"))
-        println()
-
-        assignments match {
-          case None => // do nothing
-          case Some(memberAssignments) =>
-            memberAssignments.foreach { memberAssignment =>
-              print(s"%${-maxConsumerIdLen}s %${-maxHostLen}s %${-maxClientIdLen}s %-15s ".format(
-                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"\n%${-maxGroupLen}s %${-maxConsumerIdLen}s %${-maxHostLen}s %${-maxClientIdLen}s %-15s "
+            .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 =>
+                print(s"%${-maxGroupLen}s %${-maxConsumerIdLen}s %${-maxHostLen}s %${-maxClientIdLen}s %-15s ".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))
                 }
-                print("%s".format(partitions))
+                println()
               }
-              println()
-            }
+          }
         }
       }
     }
 
-    private def printState(group: String, state: GroupState): Unit = {
-      if (shouldPrintMemberState(group, Some(state.state), Some(1))) {
-        val coordinator = s"${state.coordinator.host}:${state.coordinator.port} (${state.coordinator.idString})"
-        val coordinatorColLen = Math.max(25, coordinator.length)
-        print(s"\n%${-coordinatorColLen}s %-25s %-20s %s".format("COORDINATOR (ID)", "ASSIGNMENT-STRATEGY", "STATE", "#MEMBERS"))
-        print(s"\n%${-coordinatorColLen}s %-25s %-20s %s".format(coordinator, state.assignmentStrategy, state.state, state.numMembers))
-        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)
+          print(s"\n%${-coordinatorColLen}s %-25s %-20s %-15s %s".format("GROUP", "COORDINATOR (ID)", "ASSIGNMENT-STRATEGY", "STATE", "#MEMBERS"))
+          print(s"\n%${-coordinatorColLen}s %-25s %-20s %-15s %s".format(state.group, coordinator, state.assignmentStrategy, state.state, state.numMembers))
+          println()
+        }
       }
     }
 
-    def describeGroup(): Unit = {
-      val group = opts.options.valuesOf(opts.groupOpt).asScala.head
+    def describeGroups(): Unit = {
+      val groupIds =
+        if (opts.options.has(opts.allGroupsOpt)) listGroups()
+        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 = collectGroupOffsets()
-        printOffsets(group, offsets._1, offsets._2)
+        val offsets = collectGroupsOffsets(groupIds)
+        printOffsets(offsets)
       } else if (membersOptPresent) {
-        val members = collectGroupMembers(opts.options.has(opts.verboseOpt))
-        printMembers(group, members._1, members._2, opts.options.has(opts.verboseOpt))
-      } else
-        printState(group, collectGroupState())
+        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] = {
+                                          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)
@@ -290,109 +360,229 @@ object ConsumerGroupCommand extends Logging {
           getLag(offset, logEndOffsetOpt), consumerIdOpt, hostOpt, clientIdOpt, logEndOffsetOpt)
       }
 
-      getLogEndOffsets(topicPartitions).map {
-        case (topicPartition, LogOffsetResult.LogOffset(offset)) => getDescribePartitionResult(topicPartition, Some(offset))
-        case (topicPartition, _) => getDescribePartitionResult(topicPartition, None)
+      getLogEndOffsets(group, 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[TopicPartition, OffsetAndMetadata] = {
-      val groupId = opts.options.valueOf(opts.groupOpt)
+    def resetOffsets(): Map[String, Map[TopicPartition, OffsetAndMetadata]] = {
+      val groupIds =
+        if (opts.options.has(opts.allGroupsOpt)) listGroups()
+        else opts.options.valuesOf(opts.groupOpt).asScala
+
       val consumerGroups = adminClient.describeConsumerGroups(
-        util.Arrays.asList(groupId),
+        groupIds.asJava,
         withTimeoutMs(new DescribeConsumerGroupsOptions)
       ).describedGroups()
 
-      val group = consumerGroups.get(groupId).get
-      group.state.toString match {
-        case "Empty" | "Dead" =>
-          val partitionsToReset = getPartitionsToReset(groupId)
-          val preparedOffsets = prepareOffsetsToReset(groupId, partitionsToReset)
+      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)
+                  getConsumer(groupId).commitSync(preparedOffsets.asJava)
+                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]) = {
+      var partitionLevelResult: Map[TopicPartition, Throwable] = mutable.HashMap()
+
+      val (topicWithPartitions, topicWithoutPartitions) = topics.partition(_.contains(":"))
 
-          // 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)
-            getConsumer.commitSync(preparedOffsets.asJava)
-          preparedOffsets
-        case currentState =>
-          printError(s"Assignments can only be reset if the group '$groupId' is inactive, but the current state is $currentState.")
-          Map.empty
+      val knownPartitions = topicWithPartitions.flatMap { topicArg =>
+        val split = topicArg.split(":")
+        split(1).split(",").map { partition =>
+          new TopicPartition(split(0), partition.toInt)
+        }
       }
-    }
 
-    /**
-      * Returns the state of the specified consumer group and partition assignment states
-      */
-    def collectGroupOffsets(): (Option[String], Option[Seq[PartitionAssignmentState]]) = {
-      val groupId = opts.options.valueOf(opts.groupOpt)
-      val consumerGroup = adminClient.describeConsumerGroups(
-        List(groupId).asJava,
-        withTimeoutMs(new DescribeConsumerGroupsOptions())
-      ).describedGroups.get(groupId).get
-
-      val state = consumerGroup.state
-      val committedOffsets = getCommittedOffsets(groupId).asScala.toMap
-      var assignedTopicPartitions = ListBuffer[TopicPartition]()
-      val rowsWithConsumer = consumerGroup.members.asScala.filter(!_.assignment.topicPartitions.isEmpty).toSeq
-        .sortWith(_.assignment.topicPartitions.size > _.assignment.topicPartitions.size).flatMap { consumerSummary =>
-        val topicPartitions = consumerSummary.assignment.topicPartitions.asScala
-        assignedTopicPartitions = assignedTopicPartitions ++ topicPartitions
-        val partitionOffsets = consumerSummary.assignment.topicPartitions.asScala
-          .map { topicPartition =>
-            topicPartition -> committedOffsets.get(topicPartition).map(_.offset)
-          }.toMap
+      // Get the partitions of topics that the user did not explicitly specify the partitions
+      val describeTopicsResult = adminClient.describeTopics(
+        topicWithoutPartitions.asJava,
+        withTimeoutMs(new DescribeTopicsOptions))
 
-        collectConsumerAssignment(groupId, Option(consumerGroup.coordinator), topicPartitions.toList,
-          partitionOffsets, Some(s"${consumerSummary.consumerId}"), Some(s"${consumerSummary.host}"),
-          Some(s"${consumerSummary.clientId}"))
+      val unknownPartitions = describeTopicsResult.values().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 rowsWithoutConsumer = committedOffsets.filterKeys(!assignedTopicPartitions.contains(_)).flatMap {
-        case (topicPartition, offset) =>
-          collectConsumerAssignment(
-            groupId,
-            Option(consumerGroup.coordinator),
-            Seq(topicPartition),
-            Map(topicPartition -> Some(offset.offset)),
-                                  Some(MISSING_COLUMN_VALUE),
-                                  Some(MISSING_COLUMN_VALUE),
-                                  Some(MISSING_COLUMN_VALUE))
+      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)
       }
 
-      (Some(state.toString), Some(rowsWithConsumer ++ rowsWithoutConsumer))
+      partitions.foreach { partition =>
+        Try(deleteResult.partitionResult(partition).get()) match {
+          case Success(_) => partitionLevelResult += partition -> null
+          case Failure(e) => partitionLevelResult += partition -> e
+        }
+      }
+
+      (topLevelException, partitionLevelResult)
     }
 
-    private[admin] def collectGroupMembers(verbose: Boolean): (Option[String], Option[Seq[MemberAssignmentState]]) = {
+    def deleteOffsets(): Unit = {
       val groupId = opts.options.valueOf(opts.groupOpt)
-      val consumerGroups = adminClient.describeConsumerGroups(
-        List(groupId).asJava,
+      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")
+      }
+
+      println("\n%-30s %-15s %-15s".format("TOPIC", "PARTITION", "STATUS"))
+      partitionLevelResult.toList.sortBy(t => t._1.topic + t._1.partition.toString).foreach { case (tp, error) =>
+        println("%-30s %-15s %-15s".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()
+      ).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
+      */
+    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)
+        var assignedTopicPartitions = ListBuffer[TopicPartition]()
+        val rowsWithConsumer = consumerGroup.members.asScala.filter(!_.assignment.topicPartitions.isEmpty).toSeq
+          .sortWith(_.assignment.topicPartitions.size > _.assignment.topicPartitions.size).flatMap { consumerSummary =>
+          val topicPartitions = consumerSummary.assignment.topicPartitions.asScala
+          assignedTopicPartitions = assignedTopicPartitions ++ topicPartitions
+          val partitionOffsets = consumerSummary.assignment.topicPartitions.asScala
+            .map { topicPartition =>
+              topicPartition -> committedOffsets.get(topicPartition).map(_.offset)
+            }.toMap
+          collectConsumerAssignment(groupId, Option(consumerGroup.coordinator), topicPartitions.toList,
+            partitionOffsets, Some(s"${consumerSummary.consumerId}"), Some(s"${consumerSummary.host}"),
+            Some(s"${consumerSummary.clientId}"))
+        }
+        val rowsWithoutConsumer = committedOffsets.filterKeys(!assignedTopicPartitions.contains(_)).flatMap {
+          case (topicPartition, offset) =>
+            collectConsumerAssignment(
+              groupId,
+              Option(consumerGroup.coordinator),
+              Seq(topicPartition),
+              Map(topicPartition -> Some(offset.offset)),
+              Some(MISSING_COLUMN_VALUE),
+              Some(MISSING_COLUMN_VALUE),
+              Some(MISSING_COLUMN_VALUE)).toSeq
+        }
+        groupId -> (Some(state.toString), Some(rowsWithConsumer ++ rowsWithoutConsumer))
+      }).toMap
 
-      val group = consumerGroups.get(groupId).get
-      val state = group.state
+      groupOffsets
+    }
 
-      (Some(state.toString),
-        Option(group.members().asScala.map {
-          consumer => MemberAssignmentState(groupId, consumer.consumerId, consumer.host, consumer.clientId, consumer.assignment.topicPartitions.size(),
-            if (verbose) consumer.assignment.topicPartitions.asScala.toList else List())
-        }.toList))
+    private[admin] def collectGroupMembers(groupId: String, verbose: Boolean): (Option[String], Option[Seq[MemberAssignmentState]]) = {
+      collectGroupsMembers(Seq(groupId), verbose)(groupId)
     }
 
-    private[admin] def collectGroupState(): GroupState = {
-      val groupId = opts.options.valueOf(opts.groupOpt)
-      val consumerGroups = adminClient.describeConsumerGroups(
-        util.Arrays.asList(groupId),
-        withTimeoutMs(new DescribeConsumerGroupsOptions)
-      ).describedGroups()
+    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.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)
+    }
 
-      val group = consumerGroups.get(groupId).get
-      GroupState(groupId, group.coordinator, group.partitionAssignor(),
-        group.state.toString, group.members().size)
+    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 offsets = getConsumer.endOffsets(topicPartitions.asJava)
+    private def getLogEndOffsets(groupId: String, topicPartitions: Seq[TopicPartition]): Map[TopicPartition, LogOffsetResult] = {
+      val offsets = getConsumer(groupId).endOffsets(topicPartitions.asJava)
       topicPartitions.map { topicPartition =>
         Option(offsets.get(topicPartition)) match {
           case Some(logEndOffset) => topicPartition -> LogOffsetResult.LogOffset(logEndOffset)
@@ -401,8 +591,8 @@ object ConsumerGroupCommand extends Logging {
       }.toMap
     }
 
-    private def getLogStartOffsets(topicPartitions: Seq[TopicPartition]): Map[TopicPartition, LogOffsetResult] = {
-      val offsets = getConsumer.beginningOffsets(topicPartitions.asJava)
+    private def getLogStartOffsets(groupId: String, topicPartitions: Seq[TopicPartition]): Map[TopicPartition, LogOffsetResult] = {
+      val offsets = getConsumer(groupId).beginningOffsets(topicPartitions.asJava)
       topicPartitions.map { topicPartition =>
         Option(offsets.get(topicPartition)) match {
           case Some(logStartOffset) => topicPartition -> LogOffsetResult.LogOffset(logStartOffset)
@@ -411,8 +601,8 @@ object ConsumerGroupCommand extends Logging {
       }.toMap
     }
 
-    private def getLogTimestampOffsets(topicPartitions: Seq[TopicPartition], timestamp: java.lang.Long): Map[TopicPartition, LogOffsetResult] = {
-      val consumer = getConsumer
+    private def getLogTimestampOffsets(groupId: String, topicPartitions: Seq[TopicPartition], timestamp: java.lang.Long): Map[TopicPartition, LogOffsetResult] = {
+      val consumer = getConsumer(groupId)
       consumer.assign(topicPartitions.asJava)
 
       val (successfulOffsetsForTimes, unsuccessfulOffsetsForTimes) =
@@ -422,32 +612,35 @@ object ConsumerGroupCommand extends Logging {
         case (topicPartition, offsetAndTimestamp) => topicPartition -> LogOffsetResult.LogOffset(offsetAndTimestamp.offset)
       }.toMap
 
-      successfulLogTimestampOffsets ++ getLogEndOffsets(unsuccessfulOffsetsForTimes.keySet.toSeq)
+      successfulLogTimestampOffsets ++ getLogEndOffsets(groupId, unsuccessfulOffsetsForTimes.keySet.toSeq)
     }
 
-    def close() {
+    def close(): Unit = {
       adminClient.close()
-      if (consumer != null) consumer.close()
+      consumers.values.foreach(consumer =>
+        Option(consumer).foreach(_.close())
+      )
     }
 
-    private def createAdminClient(): admin.AdminClient = {
+    private 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.foreach { case (k, v) => props.put(k, v)}
       admin.AdminClient.create(props)
     }
 
-    private def getConsumer = {
-      if (consumer == null)
-        consumer = createConsumer()
-      consumer
+    private def getConsumer(groupId: String) = {
+      if (consumers.get(groupId).isEmpty)
+        consumers.update(groupId, createConsumer(groupId))
+      consumers(groupId)
     }
 
-    private def createConsumer(): KafkaConsumer[String, String] = {
+    private def createConsumer(groupId: String): KafkaConsumer[String, String] = {
       val properties = new Properties()
       val deserializer = (new StringDeserializer).getClass.getName
       val brokerUrl = opts.options.valueOf(opts.bootstrapServerOpt)
       properties.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, brokerUrl)
-      properties.put(ConsumerConfig.GROUP_ID_CONFIG, opts.options.valueOf(opts.groupOpt))
+      properties.put(ConsumerConfig.GROUP_ID_CONFIG, groupId)
       properties.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "false")
       properties.put(ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG, "30000")
       properties.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, deserializer)
@@ -467,22 +660,21 @@ object ConsumerGroupCommand extends Logging {
       options.timeoutMs(t)
     }
 
-    private def parseTopicPartitionsToReset(topicArgs: Seq[String]): Seq[TopicPartition] = topicArgs.flatMap {
+    private def parseTopicPartitionsToReset(groupId: String, topicArgs: Seq[String]): Seq[TopicPartition] = topicArgs.flatMap {
       case topicArg if topicArg.contains(":") =>
         val topicPartitions = topicArg.split(":")
         val topic = topicPartitions(0)
         topicPartitions(1).split(",").map(partition => new TopicPartition(topic, partition.toInt))
-      case topic => getConsumer.partitionsFor(topic).asScala
+      case topic => getConsumer(groupId).partitionsFor(topic).asScala
         .map(partitionInfo => new TopicPartition(topic, partitionInfo.partition))
     }
 
     private def getPartitionsToReset(groupId: String): Seq[TopicPartition] = {
       if (opts.options.has(opts.allTopicsOpt)) {
-        val allTopicPartitions = getCommittedOffsets(groupId).keySet().asScala.toSeq
-        allTopicPartitions
+        getCommittedOffsets(groupId).keys.toSeq
       } else if (opts.options.has(opts.topicOpt)) {
         val topics = opts.options.valuesOf(opts.topicOpt).asScala
-        parseTopicPartitionsToReset(topics)
+        parseTopicPartitionsToReset(groupId, topics)
       } else {
         if (opts.options.has(opts.resetFromFileOpt))
           Nil
@@ -491,31 +683,53 @@ object ConsumerGroupCommand extends Logging {
       }
     }
 
-    private def getCommittedOffsets(groupId: String) = {
+    private def getCommittedOffsets(groupId: String): Map[TopicPartition, OffsetAndMetadata] = {
       adminClient.listConsumerGroupOffsets(
         groupId,
         withTimeoutMs(new ListConsumerGroupOffsetsOptions)
-      ).partitionsToOffsetAndMetadata.get
+      ).partitionsToOffsetAndMetadata.get.asScala
     }
 
-    private def parseResetPlan(resetPlanCsv: String): Map[TopicPartition, OffsetAndMetadata] = {
-      resetPlanCsv.split("\n")
-        .map { line =>
-          val Array(topic, partition, offset) = line.split(",").map(_.trim)
-          val topicPartition = new TopicPartition(topic, partition.toInt)
-          val offsetAndMetadata = new OffsetAndMetadata(offset.toLong)
-          (topicPartition, offsetAndMetadata)
-        }.toMap
+    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] = {
+    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 {
+        checkOffsetsRange(groupId, partitionsToReset.map((_, offset)).toMap).map {
           case (topicPartition, newOffset) => (topicPartition, new OffsetAndMetadata(newOffset))
         }
       } else if (opts.options.has(opts.resetToEarliestOpt)) {
-        val logStartOffsets = getLogStartOffsets(partitionsToReset)
+        val logStartOffsets = getLogStartOffsets(groupId, partitionsToReset)
         partitionsToReset.map { topicPartition =>
           logStartOffsets.get(topicPartition) match {
             case Some(LogOffsetResult.LogOffset(offset)) => (topicPartition, new OffsetAndMetadata(offset))
@@ -523,7 +737,7 @@ object ConsumerGroupCommand extends Logging {
           }
         }.toMap
       } else if (opts.options.has(opts.resetToLatestOpt)) {
-        val logEndOffsets = getLogEndOffsets(partitionsToReset)
+        val logEndOffsets = getLogEndOffsets(groupId, partitionsToReset)
         partitionsToReset.map { topicPartition =>
           logEndOffsets.get(topicPartition) match {
             case Some(LogOffsetResult.LogOffset(offset)) => (topicPartition, new OffsetAndMetadata(offset))
@@ -534,16 +748,16 @@ object ConsumerGroupCommand extends Logging {
         val currentCommittedOffsets = getCommittedOffsets(groupId)
         val requestedOffsets = partitionsToReset.map { topicPartition =>
           val shiftBy = opts.options.valueOf(opts.resetShiftByOpt)
-          val currentOffset = currentCommittedOffsets.asScala.getOrElse(topicPartition,
+          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 {
+        checkOffsetsRange(groupId, requestedOffsets).map {
           case (topicPartition, newOffset) => (topicPartition, new OffsetAndMetadata(newOffset))
         }
       } else if (opts.options.has(opts.resetToDatetimeOpt)) {
         val timestamp = convertTimestamp(opts.options.valueOf(opts.resetToDatetimeOpt))
-        val logTimestampOffsets = getLogTimestampOffsets(partitionsToReset, timestamp)
+        val logTimestampOffsets = getLogTimestampOffsets(groupId, partitionsToReset, timestamp)
         partitionsToReset.map { topicPartition =>
           val logTimestampOffset = logTimestampOffsets.get(topicPartition)
           logTimestampOffset match {
@@ -555,8 +769,9 @@ object ConsumerGroupCommand extends Logging {
         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)
+        val logTimestampOffsets = getLogTimestampOffsets(groupId, partitionsToReset, timestamp)
         partitionsToReset.map { topicPartition =>
           val logTimestampOffset = logTimestampOffsets.get(topicPartition)
           logTimestampOffset match {
@@ -564,16 +779,20 @@ object ConsumerGroupCommand extends Logging {
             case _ => CommandLineUtils.printUsageAndDie(opts.parser, s"Error getting offset by timestamp of topic partition: $topicPartition")
           }
         }.toMap
-      } else if (opts.options.has(opts.resetFromFileOpt)) {
-        val resetPlanPath = opts.options.valueOf(opts.resetFromFileOpt)
-        val resetPlanCsv = Utils.readFileAsString(resetPlanPath)
-        val resetPlan = parseResetPlan(resetPlanCsv)
-        val requestedOffsets = resetPlan.keySet.map { topicPartition =>
-          (topicPartition, resetPlan(topicPartition).offset)
-        }.toMap
-        checkOffsetsRange(requestedOffsets).map {
-          case (topicPartition, newOffset) => (topicPartition, new OffsetAndMetadata(newOffset))
-        }
+      } else if (resetPlanFromFile.isDefined) {
+        resetPlanFromFile.map(resetPlan => resetPlan.get(groupId).map { resetPlanForGroup =>
+          val requestedOffsets = resetPlanForGroup.keySet.map { topicPartition =>
+            topicPartition -> resetPlanForGroup(topicPartition).offset
+          }.toMap
+          checkOffsetsRange(groupId, 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) =
@@ -581,12 +800,12 @@ object ConsumerGroupCommand extends Logging {
 
         val preparedOffsetsForPartitionsWithCommittedOffset = partitionsToResetWithCommittedOffset.map { topicPartition =>
           (topicPartition, new OffsetAndMetadata(currentCommittedOffsets.get(topicPartition) match {
-            case offset if offset != null => offset.offset
-            case _ => throw new IllegalStateException(s"Expected a valid current offset for topic partition: $topicPartition")
+            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 {
+        val preparedOffsetsForPartitionsWithoutCommittedOffset = getLogEndOffsets(groupId, partitionsToResetWithoutCommittedOffset).map {
           case (topicPartition, LogOffsetResult.LogOffset(offset)) => (topicPartition, new OffsetAndMetadata(offset))
           case (topicPartition, _) => CommandLineUtils.printUsageAndDie(opts.parser, s"Error getting ending offset of topic partition: $topicPartition")
         }
@@ -597,9 +816,9 @@ object ConsumerGroupCommand extends Logging {
       }
     }
 
-    private def checkOffsetsRange(requestedOffsets: Map[TopicPartition, Long]) = {
-      val logStartOffsets = getLogStartOffsets(requestedOffsets.keySet.toSeq)
-      val logEndOffsets = getLogEndOffsets(requestedOffsets.keySet.toSeq)
+    private def checkOffsetsRange(groupId: String, requestedOffsets: Map[TopicPartition, Long]) = {
+      val logStartOffsets = getLogStartOffsets(groupId, requestedOffsets.keySet.toSeq)
+      val logEndOffsets = getLogEndOffsets(groupId, requestedOffsets.keySet.toSeq)
       requestedOffsets.map { case (topicPartition, offset) => (topicPartition,
         logEndOffsets.get(topicPartition) match {
           case Some(LogOffsetResult.LogOffset(endOffset)) if offset > endOffset =>
@@ -620,19 +839,33 @@ object ConsumerGroupCommand extends Logging {
       }
     }
 
-    def exportOffsetsToReset(assignmentsToReset: Map[TopicPartition, OffsetAndMetadata]): String = {
-      val rows = assignmentsToReset.map { case (k,v) => s"${k.topic},${k.partition},${v.offset}" }(collection.breakOut): List[String]
-      rows.foldRight("")(_ + "\n" + _)
+    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 groupsToDelete = opts.options.valuesOf(opts.groupOpt).asScala.toList
-      val deletedGroups = adminClient.deleteConsumerGroups(
-        groupsToDelete.asJava,
+      val groupIds =
+        if (opts.options.has(opts.allGroupsOpt)) listGroups()
+        else opts.options.valuesOf(opts.groupOpt).asScala
+
+      val groupsToDelete = adminClient.deleteConsumerGroups(
+        groupIds.asJava,
         withTimeoutMs(new DeleteConsumerGroupsOptions)
       ).deletedGroups().asScala
 
-      val result = deletedGroups.mapValues { f =>
+      val result = groupsToDelete.mapValues { f =>
         Try(f.get) match {
           case _: Success[_] => null
           case Failure(e) => e
@@ -676,6 +909,7 @@ object ConsumerGroupCommand extends Logging {
     val AllTopicsDoc = "Consider all topics assigned to a group in the `reset-offsets` process."
     val ListDoc = "List all consumer groups."
     val DescribeDoc = "Describe consumer group and list offset lag (number of messages not yet processed) related to given group."
+    val AllGroupsDoc = "Apply to all consumer groups."
     val nl = System.getProperty("line.separator")
     val DeleteDoc = "Pass in groups to delete topic partition offsets and ownership information " +
       "over the entire consumer group. For instance --group g1 --group g2"
@@ -709,6 +943,7 @@ object ConsumerGroupCommand extends Logging {
       "Example: --bootstrap-server localhost:9092 --describe --group group1 --offsets"
     val StateDoc = "Describe the group state. This option may be used with '--describe' and '--bootstrap-server' options only." + nl +
       "Example: --bootstrap-server localhost:9092 --describe --group group1 --state"
+    val DeleteOffsetsDoc = "Delete offsets of consumer group. Supports one consumer group at the time, and multiple topics."
 
     val bootstrapServerOpt = parser.accepts("bootstrap-server", BootstrapServerDoc)
                                    .withRequiredArg
@@ -725,6 +960,7 @@ object ConsumerGroupCommand extends Logging {
     val allTopicsOpt = parser.accepts("all-topics", AllTopicsDoc)
     val listOpt = parser.accepts("list", ListDoc)
     val describeOpt = parser.accepts("describe", DescribeDoc)
+    val allGroupsOpt = parser.accepts("all-groups", AllGroupsDoc)
     val deleteOpt = parser.accepts("delete", DeleteDoc)
     val timeoutMsOpt = parser.accepts("timeout", TimeoutMsDoc)
                              .withRequiredArg
@@ -736,6 +972,7 @@ object ConsumerGroupCommand extends Logging {
                                   .describedAs("command config property file")
                                   .ofType(classOf[String])
     val resetOffsetsOpt = parser.accepts("reset-offsets", ResetOffsetsDoc)
+    val deleteOffsetsOpt = parser.accepts("delete-offsets", DeleteOffsetsDoc)
     val dryRunOpt = parser.accepts("dry-run", DryRunDoc)
     val executeOpt = parser.accepts("execute", ExecuteDoc)
     val exportOpt = parser.accepts("export", ExportDoc)
@@ -775,28 +1012,39 @@ object ConsumerGroupCommand extends Logging {
 
     options = parser.parse(args : _*)
 
-    val describeOptPresent = options.has(describeOpt)
-
-    val allConsumerGroupLevelOpts: Set[OptionSpec[_]] = Set(listOpt, describeOpt, deleteOpt, resetOffsetsOpt)
+    val allGroupSelectionScopeOpts: Set[OptionSpec[_]] = Set(groupOpt, allGroupsOpt)
+    val allConsumerGroupLevelOpts: Set[OptionSpec[_]]  = Set(listOpt, describeOpt, deleteOpt, resetOffsetsOpt)
     val allResetOffsetScenarioOpts: Set[OptionSpec[_]] = Set(resetToOffsetOpt, resetShiftByOpt,
       resetToDatetimeOpt, resetByDurationOpt, resetToEarliestOpt, resetToLatestOpt, resetToCurrentOpt, resetFromFileOpt)
+    val allDeleteOffsetsOpts: Set[OptionSpec[_]] = Set(groupOpt, topicOpt)
 
-    def checkArgs() {
-      // check required args
-      if (options.has(timeoutMsOpt) && !describeOptPresent)
-        debug(s"Option $timeoutMsOpt is applicable only when $describeOpt is used.")
+    def checkArgs(): Unit = {
 
       CommandLineUtils.checkRequiredArgs(parser, options, bootstrapServerOpt)
 
-      if (options.has(deleteOpt) && options.has(topicOpt))
+      if (options.has(describeOpt)) {
+        if (!options.has(groupOpt) && !options.has(allGroupsOpt))
+          CommandLineUtils.printUsageAndDie(parser,
+            s"Option $describeOpt takes one of these options: ${allGroupSelectionScopeOpts.mkString(", ")}")
+      } 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.printUsageAndDie(parser,
+            s"Option $deleteOpt takes one of these options: ${allGroupSelectionScopeOpts.mkString(", ")}")
+        if (options.has(topicOpt))
           CommandLineUtils.printUsageAndDie(parser, s"The consumer does not support topic-specific offset " +
             "deletion from a consumer group.")
+      }
 
-      if (describeOptPresent)
-        CommandLineUtils.checkRequiredArgs(parser, options, groupOpt)
-
-      if (options.has(deleteOpt) && !options.has(groupOpt))
-        CommandLineUtils.printUsageAndDie(parser, s"Option $deleteOpt takes $groupOpt")
+      if (options.has(deleteOffsetsOpt)) {
+        if (!options.has(groupOpt) || !options.has(topicOpt))
+          CommandLineUtils.printUsageAndDie(parser,
+            s"Option $deleteOffsetsOpt takes the following options: ${allDeleteOffsetsOpts.mkString(", ")}")
+      }
 
       if (options.has(resetOffsetsOpt)) {
         if (options.has(dryRunOpt) && options.has(executeOpt))
@@ -809,18 +1057,20 @@ object ConsumerGroupCommand extends Logging {
             "if you are scripting this command and want to keep the current default behavior without prompting.")
         }
 
-        CommandLineUtils.checkRequiredArgs(parser, options, groupOpt)
-        CommandLineUtils.checkInvalidArgs(parser, options, resetToOffsetOpt, allResetOffsetScenarioOpts - resetToOffsetOpt)
+        if (!options.has(groupOpt) && !options.has(allGroupsOpt))
+          CommandLineUtils.printUsageAndDie(parser,
+            s"Option $resetOffsetsOpt takes one of these options: ${allGroupSelectionScopeOpts.mkString(", ")}")
+        CommandLineUtils.checkInvalidArgs(parser, options, resetToOffsetOpt,   allResetOffsetScenarioOpts - resetToOffsetOpt)
         CommandLineUtils.checkInvalidArgs(parser, options, resetToDatetimeOpt, allResetOffsetScenarioOpts - resetToDatetimeOpt)
         CommandLineUtils.checkInvalidArgs(parser, options, resetByDurationOpt, allResetOffsetScenarioOpts - resetByDurationOpt)
         CommandLineUtils.checkInvalidArgs(parser, options, resetToEarliestOpt, allResetOffsetScenarioOpts - resetToEarliestOpt)
-        CommandLineUtils.checkInvalidArgs(parser, options, resetToLatestOpt, allResetOffsetScenarioOpts - resetToLatestOpt)
-        CommandLineUtils.checkInvalidArgs(parser, options, resetToCurrentOpt, allResetOffsetScenarioOpts - resetToCurrentOpt)
-        CommandLineUtils.checkInvalidArgs(parser, options, resetShiftByOpt, allResetOffsetScenarioOpts - resetShiftByOpt)
-        CommandLineUtils.checkInvalidArgs(parser, options, resetFromFileOpt, allResetOffsetScenarioOpts - resetFromFileOpt)
+        CommandLineUtils.checkInvalidArgs(parser, options, resetToLatestOpt,   allResetOffsetScenarioOpts - resetToLatestOpt)
+        CommandLineUtils.checkInvalidArgs(parser, options, resetToCurrentOpt,  allResetOffsetScenarioOpts - resetToCurrentOpt)
+        CommandLineUtils.checkInvalidArgs(parser, options, resetShiftByOpt,    allResetOffsetScenarioOpts - resetShiftByOpt)
+        CommandLineUtils.checkInvalidArgs(parser, options, resetFromFileOpt,   allResetOffsetScenarioOpts - resetFromFileOpt)
       }
 
-      // check invalid args
+      CommandLineUtils.checkInvalidArgs(parser, options, groupOpt, allGroupSelectionScopeOpts - groupOpt)
       CommandLineUtils.checkInvalidArgs(parser, options, groupOpt, allConsumerGroupLevelOpts - describeOpt - deleteOpt - resetOffsetsOpt)
       CommandLineUtils.checkInvalidArgs(parser, options, topicOpt, allConsumerGroupLevelOpts - deleteOpt - resetOffsetsOpt)
     }
diff --git a/core/src/main/scala/kafka/admin/DelegationTokenCommand.scala b/core/src/main/scala/kafka/admin/DelegationTokenCommand.scala
index 48371805523de..316698baea397 100644
--- a/core/src/main/scala/kafka/admin/DelegationTokenCommand.scala
+++ b/core/src/main/scala/kafka/admin/DelegationTokenCommand.scala
@@ -24,7 +24,7 @@ import java.util.Base64
 import joptsimple.ArgumentAcceptingOptionSpec
 import kafka.utils.{CommandDefaultOptions, CommandLineUtils, Exit, Logging}
 import org.apache.kafka.clients.CommonClientConfigs
-import org.apache.kafka.clients.admin.{CreateDelegationTokenOptions, DescribeDelegationTokenOptions, ExpireDelegationTokenOptions, RenewDelegationTokenOptions, AdminClient => JAdminClient}
+import org.apache.kafka.clients.admin.{Admin, CreateDelegationTokenOptions, DescribeDelegationTokenOptions, ExpireDelegationTokenOptions, RenewDelegationTokenOptions, AdminClient => JAdminClient}
 import org.apache.kafka.common.security.auth.KafkaPrincipal
 import org.apache.kafka.common.security.token.delegation.DelegationToken
 import org.apache.kafka.common.utils.{SecurityUtils, Utils}
@@ -72,7 +72,7 @@ object DelegationTokenCommand extends Logging {
     }
   }
 
-  def createToken(adminClient: JAdminClient, opts: DelegationTokenCommandOptions): DelegationToken = {
+  def createToken(adminClient: Admin, opts: DelegationTokenCommandOptions): DelegationToken = {
     val renewerPrincipals = getPrincipals(opts, opts.renewPrincipalsOpt).getOrElse(new util.LinkedList[KafkaPrincipal]())
     val maxLifeTimeMs = opts.options.valueOf(opts.maxLifeTimeOpt).longValue
 
@@ -108,7 +108,7 @@ object DelegationTokenCommand extends Logging {
       None
   }
 
-  def renewToken(adminClient: JAdminClient, opts: DelegationTokenCommandOptions): Long = {
+  def renewToken(adminClient: Admin, opts: DelegationTokenCommandOptions): Long = {
     val hmac = opts.options.valueOf(opts.hmacOpt)
     val renewTimePeriodMs = opts.options.valueOf(opts.renewTimePeriodOpt).longValue()
     println("Calling renew token operation with hmac :" + hmac +" , renew-time-period :"+ renewTimePeriodMs)
@@ -119,7 +119,7 @@ object DelegationTokenCommand extends Logging {
     expiryTimeStamp
   }
 
-  def expireToken(adminClient: JAdminClient, opts: DelegationTokenCommandOptions): Long = {
+  def expireToken(adminClient: Admin, opts: DelegationTokenCommandOptions): Long = {
     val hmac = opts.options.valueOf(opts.hmacOpt)
     val expiryTimePeriodMs = opts.options.valueOf(opts.expiryTimePeriodOpt).longValue()
     println("Calling expire token operation with hmac :" + hmac +" , expire-time-period : "+ expiryTimePeriodMs)
@@ -130,7 +130,7 @@ object DelegationTokenCommand extends Logging {
     expiryTimeStamp
   }
 
-  def describeToken(adminClient: JAdminClient, opts: DelegationTokenCommandOptions): List[DelegationToken] = {
+  def describeToken(adminClient: Admin, opts: DelegationTokenCommandOptions): List[DelegationToken] = {
     val ownerPrincipals = getPrincipals(opts, opts.ownerPrincipalsOpt)
     if (ownerPrincipals.isEmpty)
       println("Calling describe token operation for current user.")
@@ -143,7 +143,7 @@ object DelegationTokenCommand extends Logging {
     tokens
   }
 
-  private def createAdminClient(opts: DelegationTokenCommandOptions): JAdminClient = {
+  private def createAdminClient(opts: DelegationTokenCommandOptions): Admin = {
     val props = Utils.loadProps(opts.options.valueOf(opts.commandConfigOpt))
     props.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, opts.options.valueOf(opts.bootstrapServerOpt))
     JAdminClient.create(props)
@@ -196,7 +196,7 @@ object DelegationTokenCommand extends Logging {
 
     options = parser.parse(args : _*)
 
-    def checkArgs() {
+    def checkArgs(): Unit = {
       // check required args
       CommandLineUtils.checkRequiredArgs(parser, options, bootstrapServerOpt, commandConfigOpt)
 
diff --git a/core/src/main/scala/kafka/admin/DeleteRecordsCommand.scala b/core/src/main/scala/kafka/admin/DeleteRecordsCommand.scala
index 10977c2ec839b..3a342b7967a26 100644
--- a/core/src/main/scala/kafka/admin/DeleteRecordsCommand.scala
+++ b/core/src/main/scala/kafka/admin/DeleteRecordsCommand.scala
@@ -29,6 +29,7 @@ import org.apache.kafka.common.TopicPartition
 import org.apache.kafka.common.utils.Utils
 
 import scala.collection.JavaConverters._
+import scala.collection.Seq
 
 /**
  * A command for delete records of the given partitions down to the specified offset.
@@ -99,7 +100,7 @@ object DeleteRecordsCommand {
     adminClient.close()
   }
 
-  private def createAdminClient(opts: DeleteRecordsCommandOptions): admin.AdminClient = {
+  private def createAdminClient(opts: DeleteRecordsCommandOptions): admin.Admin = {
     val props = if (opts.options.has(opts.commandConfigOpt))
       Utils.loadProps(opts.options.valueOf(opts.commandConfigOpt))
     else
diff --git a/core/src/main/scala/kafka/admin/LeaderElectionCommand.scala b/core/src/main/scala/kafka/admin/LeaderElectionCommand.scala
new file mode 100644
index 0000000000000..ff349f6f3612e
--- /dev/null
+++ b/core/src/main/scala/kafka/admin/LeaderElectionCommand.scala
@@ -0,0 +1,287 @@
+/**
+ * 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 java.util.concurrent.ExecutionException
+import joptsimple.util.EnumConverter
+import kafka.common.AdminCommandFailedException
+import kafka.utils.CommandDefaultOptions
+import kafka.utils.CommandLineUtils
+import kafka.utils.CoreUtils
+import kafka.utils.Json
+import kafka.utils.Logging
+import org.apache.kafka.clients.admin.{Admin, AdminClientConfig, AdminClient => JAdminClient}
+import org.apache.kafka.common.ElectionType
+import org.apache.kafka.common.TopicPartition
+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.Utils
+import scala.collection.JavaConverters._
+import scala.collection.mutable
+import scala.concurrent.duration._
+
+object LeaderElectionCommand extends Logging {
+  def main(args: Array[String]): Unit = {
+    run(args, 30.second)
+  }
+
+  def run(args: Array[String], timeout: Duration): Unit = {
+    val commandOptions = new LeaderElectionCommandOptions(args)
+    CommandLineUtils.printHelpAndExitIfNeeded(
+      commandOptions,
+      "This tool attempts to elect a new leader for a set of topic partitions. The type of elections supported are preferred replicas and unclean replicas."
+    )
+
+    validate(commandOptions)
+
+    val electionType = commandOptions.options.valueOf(commandOptions.electionType)
+
+    val jsonFileTopicPartitions = Option(commandOptions.options.valueOf(commandOptions.pathToJsonFile)).map { path  =>
+      parseReplicaElectionData(Utils.readFileAsString(path))
+    }
+
+    val singleTopicPartition = (
+      Option(commandOptions.options.valueOf(commandOptions.topic)),
+      Option(commandOptions.options.valueOf(commandOptions.partition))
+    ) match {
+      case (Some(topic), Some(partition)) => Some(Set(new TopicPartition(topic, partition)))
+      case _ => None
+    }
+
+    /* Note: No need to look at --all-topic-partitions as we want this to be None if it is use.
+     * The validate function should be checking that this option is required if the --topic and --path-to-json-file
+     * are not specified.
+     */
+    val topicPartitions = jsonFileTopicPartitions.orElse(singleTopicPartition)
+
+    val adminClient = {
+      val props = Option(commandOptions.options.valueOf(commandOptions.adminClientConfig)).map { config =>
+        Utils.loadProps(config)
+      }.getOrElse(new Properties())
+
+      props.setProperty(
+        AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG,
+        commandOptions.options.valueOf(commandOptions.bootstrapServer)
+      )
+      props.setProperty(AdminClientConfig.REQUEST_TIMEOUT_MS_CONFIG, timeout.toMillis.toString)
+
+      JAdminClient.create(props)
+    }
+
+    try {
+      electLeaders(adminClient, electionType, topicPartitions)
+    } finally {
+      adminClient.close()
+    }
+  }
+
+  private[this] def parseReplicaElectionData(jsonString: String): Set[TopicPartition] = {
+    Json.parseFull(jsonString) match {
+      case Some(js) =>
+        js.asJsonObject.get("partitions") match {
+          case Some(partitionsList) =>
+            val partitionsRaw = partitionsList.asJsonArray.iterator.map(_.asJsonObject)
+            val partitions = partitionsRaw.map { p =>
+              val topic = p("topic").to[String]
+              val partition = p("partition").to[Int]
+              new TopicPartition(topic, partition)
+            }.toBuffer
+            val duplicatePartitions = CoreUtils.duplicates(partitions)
+            if (duplicatePartitions.nonEmpty) {
+              throw new AdminOperationException(
+                s"Replica election data contains duplicate partitions: ${duplicatePartitions.mkString(",")}"
+              )
+            }
+            partitions.toSet
+          case None => throw new AdminOperationException("Replica election data is missing \"partitions\" field")
+        }
+      case None => throw new AdminOperationException("Replica election data is empty")
+    }
+  }
+
+  private[this] def electLeaders(
+    client: Admin,
+    electionType: ElectionType,
+    topicPartitions: Option[Set[TopicPartition]]
+  ): Unit = {
+    val electionResults = try {
+      val partitions = topicPartitions.map(_.asJava).orNull
+      debug(s"Calling AdminClient.electLeaders($electionType, $partitions)")
+      client.electLeaders(electionType, partitions).partitions.get.asScala
+    } catch {
+      case e: ExecutionException =>
+        e.getCause match {
+          case cause: TimeoutException =>
+            val message = "Timeout waiting for election results"
+            println(message)
+            throw new AdminCommandFailedException(message, cause)
+          case cause: ClusterAuthorizationException =>
+            val message = "Not authorized to perform leader election"
+            println(message)
+            throw new AdminCommandFailedException(message, cause)
+          case _ =>
+            throw e
+        }
+      case e: Throwable =>
+        println("Error while making request")
+        throw e
+    }
+
+    val succeeded = mutable.Set.empty[TopicPartition]
+    val noop = mutable.Set.empty[TopicPartition]
+    val failed = mutable.Map.empty[TopicPartition, Throwable]
+
+    electionResults.foreach[Unit] { case (topicPartition, error) =>
+      if (error.isPresent) {
+        error.get match {
+          case _: ElectionNotNeededException => noop += topicPartition
+          case _ => failed += topicPartition -> error.get
+        }
+      } else {
+        succeeded += topicPartition
+      }
+    }
+
+    if (succeeded.nonEmpty) {
+      val partitions = succeeded.mkString(", ")
+      println(s"Successfully completed leader election ($electionType) for partitions $partitions")
+    }
+
+    if (noop.nonEmpty) {
+      val partitions = succeeded.mkString(", ")
+      println(s"Valid replica already elected for partitions $partitions")
+    }
+
+    if (failed.nonEmpty) {
+      val rootException = new AdminCommandFailedException(s"${failed.size} replica(s) could not be elected")
+      failed.foreach { case (topicPartition, exception) =>
+        println(s"Error completing leader election ($electionType) for partition: $topicPartition: $exception")
+        rootException.addSuppressed(exception)
+      }
+      throw rootException
+    }
+  }
+
+  private[this] def validate(commandOptions: LeaderElectionCommandOptions): Unit = {
+    // required options: --bootstrap-server and --election-type
+    var missingOptions = List.empty[String]
+    if (!commandOptions.options.has(commandOptions.bootstrapServer)) {
+      missingOptions = commandOptions.bootstrapServer.options().get(0) :: missingOptions
+    }
+
+    if (!commandOptions.options.has(commandOptions.electionType)) {
+      missingOptions = commandOptions.electionType.options().get(0) :: missingOptions
+    }
+
+    if (missingOptions.nonEmpty) {
+      throw new AdminCommandFailedException(s"Missing required option(s): ${missingOptions.mkString(", ")}")
+    }
+
+    // One and only one is required: --topic, --all-topic-partitions or --path-to-json-file
+    val mutuallyExclusiveOptions = Seq(
+      commandOptions.topic,
+      commandOptions.allTopicPartitions,
+      commandOptions.pathToJsonFile
+    )
+
+    mutuallyExclusiveOptions.count(commandOptions.options.has) match {
+      case 1 => // This is the only correct configuration, don't throw an exception
+      case _ =>
+        throw new AdminCommandFailedException(
+          "One and only one of the following options is required: " +
+          s"${mutuallyExclusiveOptions.map(_.options.get(0)).mkString(", ")}"
+        )
+    }
+
+    // --partition if and only if --topic is used
+    (
+      commandOptions.options.has(commandOptions.topic),
+      commandOptions.options.has(commandOptions.partition)
+    ) match {
+      case (true, false) =>
+        throw new AdminCommandFailedException(
+          s"Missing required option(s): ${commandOptions.partition.options.get(0)}"
+        )
+      case (false, true) =>
+        throw new AdminCommandFailedException(
+          s"Option ${commandOptions.partition.options.get(0)} is only allowed if " +
+          s"${commandOptions.topic.options.get(0)} is used"
+        )
+      case _ => // Ignore; we have a valid configuration
+    }
+  }
+}
+
+private final class LeaderElectionCommandOptions(args: Array[String]) extends CommandDefaultOptions(args) {
+  val bootstrapServer = parser
+    .accepts(
+      "bootstrap-server",
+      "A hostname and port for the broker to connect to, in the form host:port. Multiple comma separated URLs can be given. REQUIRED.")
+    .withRequiredArg
+    .describedAs("host:port")
+    .ofType(classOf[String])
+  val adminClientConfig = parser
+    .accepts(
+      "admin.config",
+      "Configuration properties files to pass to the admin client")
+    .withRequiredArg
+    .describedAs("config file")
+    .ofType(classOf[String])
+
+  val pathToJsonFile = parser
+    .accepts(
+      "path-to-json-file",
+      "The JSON file with the list  of partition for which leader elections should be performed. This is an example format. \n{\"partitions\":\n\t[{\"topic\": \"foo\", \"partition\": 1},\n\t {\"topic\": \"foobar\", \"partition\": 2}]\n}\nNot allowed if --all-topic-partitions or --topic flags are specified.")
+    .withRequiredArg
+    .describedAs("Path to JSON file")
+    .ofType(classOf[String])
+
+  val topic = parser
+    .accepts(
+      "topic",
+      "Name of topic for which to perform an election. Not allowed if --path-to-json-file or --all-topic-partitions is specified.")
+    .withRequiredArg
+    .describedAs("topic name")
+    .ofType(classOf[String])
+
+  val partition = parser
+    .accepts(
+      "partition",
+      "Partition id for which to perform an election. REQUIRED if --topic is specified.")
+    .withRequiredArg
+    .describedAs("partition id")
+    .ofType(classOf[Integer])
+
+  val allTopicPartitions = parser
+    .accepts(
+      "all-topic-partitions",
+      "Perform election on all of the eligible topic partitions based on the type of election (see the --election-type flag). Not allowed if --topic or --path-to-json-file is specified.")
+
+  val electionType = parser
+    .accepts(
+      "election-type",
+      "Type of election to attempt. Possible values are \"preferred\" for preferred leader election or \"unclean\" for unclean leader election. If preferred election is selection, the election is only performed if the current leader is not the preferred leader for the topic partition. If unclean election is selected, the election is only performed if there are no leader for the topic partition. REQUIRED.")
+    .withRequiredArg
+    .describedAs("election type")
+    .withValuesConvertedBy(ElectionTypeConverter)
+
+  options = parser.parse(args: _*)
+}
+
+final object ElectionTypeConverter extends EnumConverter[ElectionType](classOf[ElectionType]) { }
diff --git a/core/src/main/scala/kafka/admin/LogDirsCommand.scala b/core/src/main/scala/kafka/admin/LogDirsCommand.scala
index 3beb29bf349be..40d39e1152b6b 100644
--- a/core/src/main/scala/kafka/admin/LogDirsCommand.scala
+++ b/core/src/main/scala/kafka/admin/LogDirsCommand.scala
@@ -21,7 +21,7 @@ import java.io.PrintStream
 import java.util.Properties
 
 import kafka.utils.{CommandDefaultOptions, CommandLineUtils, Json}
-import org.apache.kafka.clients.admin.{AdminClientConfig, DescribeLogDirsResult, AdminClient => JAdminClient}
+import org.apache.kafka.clients.admin.{Admin, AdminClientConfig, DescribeLogDirsResult, AdminClient => JAdminClient}
 import org.apache.kafka.common.requests.DescribeLogDirsResponse.LogDirInfo
 import org.apache.kafka.common.utils.Utils
 
@@ -48,7 +48,7 @@ object LogDirsCommand {
 
         out.println("Querying brokers for log directories information")
         val describeLogDirsResult: DescribeLogDirsResult = adminClient.describeLogDirs(brokerList.map(Integer.valueOf).toSeq.asJava)
-        val logDirInfosByBroker = describeLogDirsResult.all.get().asScala.mapValues(_.asScala)
+        val logDirInfosByBroker = describeLogDirsResult.all.get().asScala.mapValues(_.asScala).toMap
 
         out.println(s"Received log directory information from brokers ${brokerList.mkString(",")}")
         out.println(formatAsJson(logDirInfosByBroker, topicList.toSet))
@@ -82,7 +82,7 @@ object LogDirsCommand {
         ).asJava)
     }
 
-    private def createAdminClient(opts: LogDirsCommandOptions): JAdminClient = {
+    private def createAdminClient(opts: LogDirsCommandOptions): Admin = {
         val props = if (opts.options.has(opts.commandConfigOpt))
             Utils.loadProps(opts.options.valueOf(opts.commandConfigOpt))
         else
diff --git a/core/src/main/scala/kafka/admin/PreferredReplicaLeaderElectionCommand.scala b/core/src/main/scala/kafka/admin/PreferredReplicaLeaderElectionCommand.scala
index 8740ed45b89c5..04765a8b20600 100755
--- a/core/src/main/scala/kafka/admin/PreferredReplicaLeaderElectionCommand.scala
+++ b/core/src/main/scala/kafka/admin/PreferredReplicaLeaderElectionCommand.scala
@@ -16,32 +16,34 @@
  */
 package kafka.admin
 
+import collection.JavaConverters._
+import collection._
 import java.util.Properties
 import java.util.concurrent.ExecutionException
-
 import joptsimple.OptionSpecBuilder
 import kafka.common.AdminCommandFailedException
 import kafka.utils._
 import kafka.zk.KafkaZkClient
 import org.apache.kafka.clients.admin.AdminClientConfig
+import org.apache.kafka.common.ElectionType
+import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.errors.ClusterAuthorizationException
+import org.apache.kafka.common.errors.ElectionNotNeededException
 import org.apache.kafka.common.errors.TimeoutException
-
-import collection.JavaConverters._
-import org.apache.kafka.common.utils.{Time, Utils}
 import org.apache.kafka.common.security.JaasUtils
-import org.apache.kafka.common.{KafkaFuture, TopicPartition}
+import org.apache.kafka.common.utils.Time
+import org.apache.kafka.common.utils.Utils
 import org.apache.zookeeper.KeeperException.NodeExistsException
 
-import collection._
-
 object PreferredReplicaLeaderElectionCommand extends Logging {
 
   def main(args: Array[String]): Unit = {
-
     val timeout = 30000
     run(args, timeout)
   }
+
   def run(args: Array[String], timeout: Int = 30000): Unit = {
+    println("This tool is deprecated. Please use kafka-leader-election tool. Tracking issue: KAFKA-8405")
     val commandOpts = new PreferredReplicaLeaderElectionCommandOptions(args)
     CommandLineUtils.printHelpAndExitIfNeeded(commandOpts, "This tool helps to causes leadership for each partition to be transferred back to the 'preferred replica'," +
       " it can be used to balance leadership among the servers.")
@@ -103,7 +105,7 @@ object PreferredReplicaLeaderElectionCommand extends Logging {
   }
 
   def writePreferredReplicaElectionData(zkClient: KafkaZkClient,
-                                        partitionsUndergoingPreferredReplicaElection: Set[TopicPartition]) {
+                                        partitionsUndergoingPreferredReplicaElection: Set[TopicPartition]): Unit = {
     try {
       zkClient.createPreferredReplicaElection(partitionsUndergoingPreferredReplicaElection.toSet)
       println("Created preferred replica election path with %s".format(partitionsUndergoingPreferredReplicaElection.mkString(",")))
@@ -148,7 +150,6 @@ object PreferredReplicaLeaderElectionCommand extends Logging {
       .describedAs("config file")
       .ofType(classOf[String])
 
-    parser.accepts("")
     options = parser.parse(args: _*)
   }
 
@@ -157,8 +158,8 @@ object PreferredReplicaLeaderElectionCommand extends Logging {
     /** Elect the preferred leader for the given {@code partitionsForElection}.
       * If the given {@code partitionsForElection} are None then elect the preferred leader for all partitions.
       */
-    def electPreferredLeaders(partitionsForElection: Option[Set[TopicPartition]]) : Unit
-    def close() : Unit
+    def electPreferredLeaders(partitionsForElection: Option[Set[TopicPartition]]): Unit
+    def close(): Unit
   }
 
   class ZkCommand(zkConnect: String, isSecure: Boolean, timeout: Int)
@@ -168,7 +169,7 @@ object PreferredReplicaLeaderElectionCommand extends Logging {
     val time = Time.SYSTEM
     zkClient = KafkaZkClient(zkConnect, isSecure, timeout, timeout, Int.MaxValue, time)
 
-    override def electPreferredLeaders(partitionsFromUser: Option[Set[TopicPartition]]) {
+    override def electPreferredLeaders(partitionsFromUser: Option[Set[TopicPartition]]): Unit = {
       try {
         val topics =
           partitionsFromUser match {
@@ -210,71 +211,68 @@ object PreferredReplicaLeaderElectionCommand extends Logging {
 
     val adminClient = org.apache.kafka.clients.admin.AdminClient.create(adminClientProps)
 
-    /**
-      * Wait until the given future has completed, then return whether it completed exceptionally.
-      * Because KafkaFuture.isCompletedExceptionally doesn't wait for a result
-      */
-    private def completedExceptionally[T](future: KafkaFuture[T]): Boolean = {
-      try {
-        future.get()
-        false
-      } catch {
-        case (_: Throwable) =>
-          true
-      }
-    }
-
     override def electPreferredLeaders(partitionsFromUser: Option[Set[TopicPartition]]): Unit = {
       val partitions = partitionsFromUser match {
         case Some(partitionsFromUser) => partitionsFromUser.asJava
         case None => null
       }
-      debug(s"Calling AdminClient.electPreferredLeaders($partitions)")
-      val result = adminClient.electPreferredLeaders(partitions)
-      // wait for all results
-
-      val attemptedPartitions = partitionsFromUser match {
-        case Some(partitionsFromUser) => partitions.asScala
-        case None => try {
-          result.partitions().get.asScala
-        } catch {
-          case e: ExecutionException =>
-            val cause = e.getCause
-            if (cause.isInstanceOf[TimeoutException]) {
-              // We timed out, or don't even know the attempted partitions
-              println("Timeout waiting for election results")
-            }
-            throw new AdminCommandFailedException(null, cause)
-          case e: Throwable =>
-            // We don't even know the attempted partitions
-            println("Error while making request")
-            e.printStackTrace()
-            return
-        }
-      }
+      debug(s"Calling AdminClient.electLeaders(ElectionType.PREFERRED, $partitions)")
 
-      val (exceptional, ok) = attemptedPartitions.map(tp => tp -> result.partitionResult(tp)).
-        partition { case (_, partitionResult) => completedExceptionally(partitionResult) }
+      val electionResults = try {
+        adminClient.electLeaders(ElectionType.PREFERRED, partitions).partitions.get.asScala
+      } catch {
+        case e: ExecutionException =>
+          val cause = e.getCause
+          if (cause.isInstanceOf[TimeoutException]) {
+            println("Timeout waiting for election results")
+            throw new AdminCommandFailedException("Timeout waiting for election results", cause)
+          } else if (cause.isInstanceOf[ClusterAuthorizationException]) {
+            println(s"Not authorized to perform leader election")
+            throw new AdminCommandFailedException("Not authorized to perform leader election", cause)
+          }
 
-      if (!ok.isEmpty) {
-        println(s"Successfully completed preferred replica election for partitions ${ok.map{ case (tp, future) => tp }.mkString(", ")}")
+          throw e
+        case e: Throwable =>
+          // We don't even know the attempted partitions
+          println("Error while making request")
+          e.printStackTrace()
+          return
       }
-      if (!exceptional.isEmpty) {
-        val adminException = new AdminCommandFailedException(
-          s"${exceptional.size} preferred replica(s) could not be elected")
-        for ((partition, void) <- exceptional) {
-          val exception = try {
-            void.get()
-            new AdminCommandFailedException("Exceptional future with no exception")
-          } catch {
-            case e: ExecutionException => e.getCause
+
+      val succeeded = mutable.Set.empty[TopicPartition]
+      val noop = mutable.Set.empty[TopicPartition]
+      val failed = mutable.Map.empty[TopicPartition, Throwable]
+
+      electionResults.foreach[Unit] { case (topicPartition, error) =>
+        if (error.isPresent) {
+          if (error.get.isInstanceOf[ElectionNotNeededException]) {
+            noop += topicPartition
+          } else {
+            failed += topicPartition -> error.get
           }
-          println(s"Error completing preferred replica election for partition $partition: $exception")
-          adminException.addSuppressed(exception)
+        } else {
+          succeeded += topicPartition
         }
-        throw adminException
       }
 
+      if (!succeeded.isEmpty) {
+        val partitions = succeeded.mkString(", ")
+        println(s"Successfully completed preferred leader election for partitions $partitions")
+      }
+
+      if (!noop.isEmpty) {
+        val partitions = succeeded.mkString(", ")
+        println(s"Preferred replica already elected for partitions $partitions")
+      }
+
+      if (!failed.isEmpty) {
+        val rootException = new AdminCommandFailedException(s"${failed.size} preferred replica(s) could not be elected")
+        failed.foreach { case (topicPartition, exception) =>
+          println(s"Error completing preferred leader election for partition: $topicPartition: $exception")
+          rootException.addSuppressed(exception)
+        }
+        throw rootException
+      }
     }
 
     override def close(): Unit = {
@@ -285,7 +283,7 @@ object PreferredReplicaLeaderElectionCommand extends Logging {
 }
 
 class PreferredReplicaLeaderElectionCommand(zkClient: KafkaZkClient, partitionsFromUser: scala.collection.Set[TopicPartition]) {
-  def moveLeaderToPreferredReplica() = {
+  def moveLeaderToPreferredReplica(): Unit = {
     try {
       val topics = partitionsFromUser.map(_.topic).toSet
       val partitionsFromZk = zkClient.getPartitionsForTopics(topics).flatMap { case (topic, partitions) =>
diff --git a/core/src/main/scala/kafka/admin/ReassignPartitionsCommand.scala b/core/src/main/scala/kafka/admin/ReassignPartitionsCommand.scala
index c108f07ee67b2..934191ac39e2a 100755
--- a/core/src/main/scala/kafka/admin/ReassignPartitionsCommand.scala
+++ b/core/src/main/scala/kafka/admin/ReassignPartitionsCommand.scala
@@ -27,7 +27,7 @@ import kafka.utils._
 import kafka.utils.json.JsonValue
 import kafka.zk.{AdminZkClient, KafkaZkClient}
 import org.apache.kafka.clients.admin.DescribeReplicaLogDirsResult.ReplicaLogDirInfo
-import org.apache.kafka.clients.admin.{AdminClientConfig, AlterReplicaLogDirsOptions, AdminClient => JAdminClient}
+import org.apache.kafka.clients.admin.{Admin, AdminClientConfig, AlterReplicaLogDirsOptions, AdminClient => JAdminClient}
 import org.apache.kafka.common.errors.ReplicaNotAvailableException
 import org.apache.kafka.common.security.JaasUtils
 import org.apache.kafka.common.utils.{Time, Utils}
@@ -70,7 +70,7 @@ object ReassignPartitionsCommand extends Logging {
     } finally zkClient.close()
   }
 
-  private def createAdminClient(opts: ReassignPartitionsCommandOptions): Option[JAdminClient] = {
+  private def createAdminClient(opts: ReassignPartitionsCommandOptions): Option[Admin] = {
     if (opts.options.has(opts.bootstrapServerOpt)) {
       val props = if (opts.options.has(opts.commandConfigOpt))
         Utils.loadProps(opts.options.valueOf(opts.commandConfigOpt))
@@ -84,13 +84,13 @@ object ReassignPartitionsCommand extends Logging {
     }
   }
 
-  def verifyAssignment(zkClient: KafkaZkClient, adminClientOpt: Option[JAdminClient], opts: ReassignPartitionsCommandOptions) {
+  def verifyAssignment(zkClient: KafkaZkClient, adminClientOpt: Option[Admin], opts: ReassignPartitionsCommandOptions): Unit = {
     val jsonFile = opts.options.valueOf(opts.reassignmentJsonFileOpt)
     val jsonString = Utils.readFileAsString(jsonFile)
     verifyAssignment(zkClient, adminClientOpt, jsonString)
   }
 
-  def verifyAssignment(zkClient: KafkaZkClient, adminClientOpt: Option[JAdminClient], jsonString: String): Unit = {
+  def verifyAssignment(zkClient: KafkaZkClient, adminClientOpt: Option[Admin], jsonString: String): Unit = {
     println("Status of partition reassignment: ")
     val adminZkClient = new AdminZkClient(zkClient)
     val (partitionsToBeReassigned, replicaAssignment) = parsePartitionReassignmentData(jsonString)
@@ -160,7 +160,7 @@ object ReassignPartitionsCommand extends Logging {
     }
   }
 
-  def generateAssignment(zkClient: KafkaZkClient, opts: ReassignPartitionsCommandOptions) {
+  def generateAssignment(zkClient: KafkaZkClient, opts: ReassignPartitionsCommandOptions): Unit = {
     val topicsToMoveJsonFile = opts.options.valueOf(opts.topicsToMoveJsonFileOpt)
     val brokerListToReassign = opts.options.valueOf(opts.brokerListOpt).split(',').map(_.toInt)
     val duplicateReassignments = CoreUtils.duplicates(brokerListToReassign)
@@ -196,7 +196,7 @@ object ReassignPartitionsCommand extends Logging {
     (partitionsToBeReassigned, currentAssignment)
   }
 
-  def executeAssignment(zkClient: KafkaZkClient, adminClientOpt: Option[JAdminClient], opts: ReassignPartitionsCommandOptions) {
+  def executeAssignment(zkClient: KafkaZkClient, adminClientOpt: Option[Admin], opts: ReassignPartitionsCommandOptions): Unit = {
     val reassignmentJsonFile =  opts.options.valueOf(opts.reassignmentJsonFileOpt)
     val reassignmentJsonString = Utils.readFileAsString(reassignmentJsonFile)
     val interBrokerThrottle = opts.options.valueOf(opts.interBrokerThrottleOpt)
@@ -205,7 +205,7 @@ object ReassignPartitionsCommand extends Logging {
     executeAssignment(zkClient, adminClientOpt, reassignmentJsonString, Throttle(interBrokerThrottle, replicaAlterLogDirsThrottle), timeoutMs)
   }
 
-  def executeAssignment(zkClient: KafkaZkClient, adminClientOpt: Option[JAdminClient], reassignmentJsonString: String, throttle: Throttle, timeoutMs: Long = 10000L) {
+  def executeAssignment(zkClient: KafkaZkClient, adminClientOpt: Option[Admin], reassignmentJsonString: String, throttle: Throttle, timeoutMs: Long = 10000L): Unit = {
     val (partitionAssignment, replicaAssignment) = parseAndValidate(zkClient, reassignmentJsonString)
     val adminZkClient = new AdminZkClient(zkClient)
     val reassignPartitionsCommand = new ReassignPartitionsCommand(zkClient, adminClientOpt, partitionAssignment.toMap, replicaAssignment, adminZkClient)
@@ -371,7 +371,7 @@ object ReassignPartitionsCommand extends Logging {
     }.toMap
   }
 
-  private def checkIfReplicaReassignmentSucceeded(adminClientOpt: Option[JAdminClient], replicaAssignment: Map[TopicPartitionReplica, String])
+  private def checkIfReplicaReassignmentSucceeded(adminClientOpt: Option[Admin], replicaAssignment: Map[TopicPartitionReplica, String])
   :Map[TopicPartitionReplica, ReassignmentStatus] = {
 
     val replicaLogDirInfos = {
@@ -502,7 +502,7 @@ object ReassignPartitionsCommand extends Logging {
 }
 
 class ReassignPartitionsCommand(zkClient: KafkaZkClient,
-                                adminClientOpt: Option[JAdminClient],
+                                adminClientOpt: Option[Admin],
                                 proposedPartitionAssignment: Map[TopicPartition, Seq[Int]],
                                 proposedReplicaAssignment: Map[TopicPartitionReplica, String] = Map.empty,
                                 adminZkClient: AdminZkClient)
@@ -531,7 +531,7 @@ class ReassignPartitionsCommand(zkClient: KafkaZkClient,
     * Limit the throttle on currently moving replicas. Note that this command can use used to alter the throttle, but
     * it may not alter all limits originally set, if some of the brokers have completed their rebalance.
     */
-  def maybeLimit(throttle: Throttle) {
+  def maybeLimit(throttle: Throttle): Unit = {
     if (throttle.interBrokerLimit >= 0 || throttle.replicaAlterLogDirsLimit >= 0) {
       val existingBrokers = existingAssignment().values.flatten.toSeq
       val proposedBrokers = proposedPartitionAssignment.values.flatten.toSeq ++ proposedReplicaAssignment.keys.toSeq.map(_.brokerId())
@@ -596,7 +596,7 @@ class ReassignPartitionsCommand(zkClient: KafkaZkClient,
     }.mkString(",")
 
   private def alterReplicaLogDirsIgnoreReplicaNotAvailable(replicaAssignment: Map[TopicPartitionReplica, String],
-                                                           adminClient: JAdminClient,
+                                                           adminClient: Admin,
                                                            timeoutMs: Long): Set[TopicPartitionReplica] = {
     val alterReplicaLogDirsResult = adminClient.alterReplicaLogDirs(replicaAssignment.asJava, new AlterReplicaLogDirsOptions().timeoutMs(timeoutMs.toInt))
     val replicasAssignedToFutureDir = alterReplicaLogDirsResult.values().asScala.flatMap { case (replica, future) => {
@@ -640,7 +640,8 @@ class ReassignPartitionsCommand(zkClient: KafkaZkClient,
         val replicasAssignedToFutureDir = mutable.Set.empty[TopicPartitionReplica]
         while (remainingTimeMs > 0 && replicasAssignedToFutureDir.size < proposedReplicaAssignment.size) {
           replicasAssignedToFutureDir ++= alterReplicaLogDirsIgnoreReplicaNotAvailable(
-            proposedReplicaAssignment.filterKeys(replica => !replicasAssignedToFutureDir.contains(replica)), adminClientOpt.get, remainingTimeMs)
+            proposedReplicaAssignment.filter { case (replica, _) => !replicasAssignedToFutureDir.contains(replica) },
+            adminClientOpt.get, remainingTimeMs)
           Thread.sleep(100)
           remainingTimeMs = startTimeMs + timeoutMs - System.currentTimeMillis()
         }
@@ -648,7 +649,7 @@ class ReassignPartitionsCommand(zkClient: KafkaZkClient,
       }
     } catch {
       case _: NodeExistsException =>
-        val partitionsBeingReassigned = zkClient.getPartitionReassignment
+        val partitionsBeingReassigned = zkClient.getPartitionReassignment()
         throw new AdminCommandFailedException("Partition reassignment currently in " +
           "progress for %s. Aborting operation".format(partitionsBeingReassigned))
     }
diff --git a/core/src/main/scala/kafka/admin/TopicCommand.scala b/core/src/main/scala/kafka/admin/TopicCommand.scala
index a4fa20f324b66..0ceb9aaa03605 100755
--- a/core/src/main/scala/kafka/admin/TopicCommand.scala
+++ b/core/src/main/scala/kafka/admin/TopicCommand.scala
@@ -28,10 +28,10 @@ import kafka.utils.Implicits._
 import kafka.utils._
 import kafka.zk.{AdminZkClient, KafkaZkClient}
 import org.apache.kafka.clients.CommonClientConfigs
-import org.apache.kafka.clients.admin.{ListTopicsOptions, NewPartitions, NewTopic, AdminClient => JAdminClient}
-import org.apache.kafka.common.TopicPartition
-import org.apache.kafka.common.config.{ConfigResource, TopicConfig}
+import org.apache.kafka.clients.admin.{Admin, ConfigEntry, ListTopicsOptions, NewPartitions, NewTopic, AdminClient => JAdminClient, Config => JConfig}
+import org.apache.kafka.common.{Node, TopicPartition, TopicPartitionInfo}
 import org.apache.kafka.common.config.ConfigResource.Type
+import org.apache.kafka.common.config.{ConfigResource, TopicConfig}
 import org.apache.kafka.common.errors.{InvalidTopicException, TopicExistsException}
 import org.apache.kafka.common.internals.Topic
 import org.apache.kafka.common.security.JaasUtils
@@ -40,6 +40,7 @@ import org.apache.zookeeper.KeeperException.NodeExistsException
 
 import scala.collection.JavaConverters._
 import scala.collection._
+import scala.compat.java8.OptionConverters._
 import scala.io.StdIn
 
 object TopicCommand extends Logging {
@@ -77,12 +78,10 @@ object TopicCommand extends Logging {
     }
   }
 
-
-
   class CommandTopicPartition(opts: TopicCommandOptions) {
     val name: String = opts.topic.get
     val partitions: Option[Integer] = opts.partitions
-    val replicationFactor: Integer = opts.replicationFactor.getOrElse(-1)
+    val replicationFactor: Option[Integer] = opts.replicationFactor
     val replicaAssignment: Option[Map[Int, List[Int]]] = opts.replicaAssignment
     val configsToAdd: Properties = parseTopicConfigsToBeAdded(opts)
     val configsToDelete: Seq[String] = parseTopicConfigsToBeDeleted(opts)
@@ -93,43 +92,96 @@ object TopicCommand extends Logging {
     def ifTopicDoesntExist(): Boolean = opts.ifNotExists
   }
 
-  case class PartitionDescription(
-                                   topic: String,
-                                   partition: Int,
-                                   leader: Option[Int],
-                                   assignedReplicas: Seq[Int],
-                                   isr: Seq[Int],
-                                   minIsrCount: Int,
-                                   markedForDeletion: Boolean,
-                                   describeConfigs: Boolean)
+  case class TopicDescription(topic: String,
+                              numPartitions: Int,
+                              replicationFactor: Int,
+                              config: JConfig,
+                              markedForDeletion: Boolean) {
+
+    def printDescription(): Unit = {
+      val configsAsString = config.entries.asScala.filter(!_.isDefault).map { ce => s"${ce.name}=${ce.value}" }.mkString(",")
+      print(s"Topic: $topic")
+      print(s"\tPartitionCount: $numPartitions")
+      print(s"\tReplicationFactor: $replicationFactor")
+      print(s"\tConfigs: $configsAsString")
+      print(if (markedForDeletion) "\tMarkedForDeletion: true" else "")
+      println()
+    }
+  }
 
-  class DescribeOptions(opts: TopicCommandOptions, liveBrokers: Set[Int]) {
-    val describeConfigs: Boolean = !opts.reportUnavailablePartitions && !opts.reportUnderReplicatedPartitions && !opts.reportUnderMinIsrPartitions
-    val describePartitions: Boolean = !opts.reportOverriddenConfigs
-    private def hasUnderReplicatedPartitions(partitionDescription: PartitionDescription) = {
-      partitionDescription.isr.size < partitionDescription.assignedReplicas.size
+  case class PartitionDescription(topic: String,
+                                  info: TopicPartitionInfo,
+                                  config: Option[JConfig],
+                                  markedForDeletion: Boolean) {
+
+    private def minIsrCount: Option[Int] = {
+      config.map(_.get(TopicConfig.MIN_IN_SYNC_REPLICAS_CONFIG).value.toInt)
+    }
+
+    def hasUnderReplicatedPartitions: Boolean = {
+      info.isr.size < info.replicas.size
     }
-    private def shouldPrintUnderReplicatedPartitions(partitionDescription: PartitionDescription) = {
-      opts.reportUnderReplicatedPartitions && hasUnderReplicatedPartitions(partitionDescription)
+
+    private def hasLeader: Boolean = {
+      info.leader != null
     }
-    private def hasUnavailablePartitions(partitionDescription: PartitionDescription) = {
-      partitionDescription.leader.isEmpty || !liveBrokers.contains(partitionDescription.leader.get)
+
+    def hasUnderMinIsrPartitions: Boolean = {
+      !hasLeader || minIsrCount.exists(info.isr.size < _)
     }
-    private def shouldPrintUnavailablePartitions(partitionDescription: PartitionDescription) = {
-      opts.reportUnavailablePartitions && hasUnavailablePartitions(partitionDescription)
+
+    def isAtMinIsrPartitions: Boolean =  {
+      minIsrCount.contains(info.isr.size)
     }
-    private def hasUnderMinIsrPartitions(partitionDescription: PartitionDescription) = {
-      partitionDescription.isr.size < partitionDescription.minIsrCount
+
+    def hasUnavailablePartitions(liveBrokers: Set[Int]): Boolean = {
+      !hasLeader || !liveBrokers.contains(info.leader.id)
     }
-    private def shouldPrintUnderMinIsrPartitions(partitionDescription: PartitionDescription) = {
-      opts.reportUnderMinIsrPartitions && hasUnderMinIsrPartitions(partitionDescription)
+
+    def printDescription(): Unit = {
+      print("\tTopic: " + topic)
+      print("\tPartition: " + info.partition)
+      print("\tLeader: " + (if (hasLeader) info.leader.id else "none"))
+      print("\tReplicas: " + info.replicas.asScala.map(_.id).mkString(","))
+      print("\tIsr: " + info.isr.asScala.map(_.id).mkString(","))
+      print(if (markedForDeletion) "\tMarkedForDeletion: true" else "")
+      println()
     }
 
-    def shouldPrintTopicPartition(partitionDesc: PartitionDescription): Boolean = {
+  }
+
+  class DescribeOptions(opts: TopicCommandOptions, liveBrokers: Set[Int]) {
+    val describeConfigs: Boolean =
+      !opts.reportUnavailablePartitions &&
+      !opts.reportUnderReplicatedPartitions &&
+      !opts.reportUnderMinIsrPartitions &&
+      !opts.reportAtMinIsrPartitions
+    val describePartitions: Boolean = !opts.reportOverriddenConfigs
+
+    private def shouldPrintUnderReplicatedPartitions(partitionDescription: PartitionDescription): Boolean = {
+      opts.reportUnderReplicatedPartitions && partitionDescription.hasUnderReplicatedPartitions
+    }
+    private def shouldPrintUnavailablePartitions(partitionDescription: PartitionDescription): Boolean = {
+      opts.reportUnavailablePartitions && partitionDescription.hasUnavailablePartitions(liveBrokers)
+    }
+    private def shouldPrintUnderMinIsrPartitions(partitionDescription: PartitionDescription): Boolean = {
+      opts.reportUnderMinIsrPartitions && partitionDescription.hasUnderMinIsrPartitions
+    }
+    private def shouldPrintAtMinIsrPartitions(partitionDescription: PartitionDescription): Boolean = {
+      opts.reportAtMinIsrPartitions && partitionDescription.isAtMinIsrPartitions
+    }
+
+    private def shouldPrintTopicPartition(partitionDesc: PartitionDescription): Boolean = {
       describeConfigs ||
         shouldPrintUnderReplicatedPartitions(partitionDesc) ||
         shouldPrintUnavailablePartitions(partitionDesc) ||
-        shouldPrintUnderMinIsrPartitions(partitionDesc)
+        shouldPrintUnderMinIsrPartitions(partitionDesc) ||
+        shouldPrintAtMinIsrPartitions(partitionDesc)
+    }
+
+    def maybePrintPartitionDescription(desc: PartitionDescription): Unit = {
+      if (shouldPrintTopicPartition(desc))
+        desc.printDescription()
     }
   }
 
@@ -141,16 +193,16 @@ object TopicCommand extends Logging {
           "collide. To avoid issues it is best to use either, but not both.")
       createTopic(topic)
     }
-    def createTopic(topic: CommandTopicPartition)
-    def listTopics(opts: TopicCommandOptions)
-    def alterTopic(opts: TopicCommandOptions)
-    def describeTopic(opts: TopicCommandOptions)
-    def deleteTopic(opts: TopicCommandOptions)
+    def createTopic(topic: CommandTopicPartition): Unit
+    def listTopics(opts: TopicCommandOptions): Unit
+    def alterTopic(opts: TopicCommandOptions): Unit
+    def describeTopic(opts: TopicCommandOptions): Unit
+    def deleteTopic(opts: TopicCommandOptions): Unit
     def getTopics(topicWhitelist: Option[String], excludeInternalTopics: Boolean = false): Seq[String]
   }
 
   object AdminClientTopicService {
-    def createAdminClient(commandConfig: Properties, bootstrapServer: Option[String]): JAdminClient = {
+    def createAdminClient(commandConfig: Properties, bootstrapServer: Option[String]): Admin = {
       bootstrapServer match {
         case Some(serverList) => commandConfig.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, serverList)
         case None =>
@@ -162,17 +214,24 @@ object TopicCommand extends Logging {
       new AdminClientTopicService(createAdminClient(commandConfig, bootstrapServer))
   }
 
-  case class AdminClientTopicService private (adminClient: JAdminClient) extends TopicService {
+  case class AdminClientTopicService private (adminClient: Admin) extends TopicService {
 
     override def createTopic(topic: CommandTopicPartition): Unit = {
-      if (topic.replicationFactor > Short.MaxValue)
-        throw new IllegalArgumentException(s"The replication factor's maximum value must be smaller or equal to ${Short.MaxValue}")
+      if (topic.replicationFactor.exists(rf => rf > Short.MaxValue || rf < 1))
+        throw new IllegalArgumentException(s"The replication factor must be between 1 and ${Short.MaxValue} inclusive")
+      if (topic.partitions.exists(partitions => partitions < 1))
+        throw new IllegalArgumentException(s"The partitions must be greater than 0")
 
       if (!adminClient.listTopics().names().get().contains(topic.name)) {
         val newTopic = if (topic.hasReplicaAssignment)
           new NewTopic(topic.name, asJavaReplicaReassignment(topic.replicaAssignment.get))
-        else
-          new NewTopic(topic.name, topic.partitions.get, topic.replicationFactor.shortValue())
+        else {
+          new NewTopic(
+            topic.name,
+            topic.partitions.asJava,
+            topic.replicationFactor.map(_.toShort).map(Short.box).asJava)
+        }
+
         val configsMap = topic.configsToAdd.stringPropertyNames()
           .asScala
           .map(name => name -> topic.configsToAdd.getProperty(name))
@@ -193,7 +252,7 @@ object TopicCommand extends Logging {
     override def alterTopic(opts: TopicCommandOptions): Unit = {
       val topic = new CommandTopicPartition(opts)
       val topics = getTopics(opts.topic, opts.excludeInternalTopics)
-      ensureTopicExists(topics)
+      ensureTopicExists(topics, opts.topic)
       val topicsInfo = adminClient.describeTopics(topics.asJavaCollection).values()
       adminClient.createPartitions(topics.map {topicName =>
         if (topic.hasReplicaAssignment) {
@@ -216,35 +275,24 @@ object TopicCommand extends Logging {
       val describeOptions = new DescribeOptions(opts, liveBrokers.toSet)
 
       for (td <- topicDescriptions) {
-        val sortedPartitions = td.partitions().asScala.sortBy(_.partition())
+        val topicName = td.name
+        val config = allConfigs.get(new ConfigResource(Type.TOPIC, topicName)).get()
+        val sortedPartitions = td.partitions.asScala.sortBy(_.partition)
+
         if (describeOptions.describeConfigs) {
-          val config = allConfigs.get(new ConfigResource(Type.TOPIC, td.name())).get()
           val hasNonDefault = config.entries().asScala.exists(!_.isDefault)
           if (!opts.reportOverriddenConfigs || hasNonDefault) {
             val numPartitions = td.partitions().size
-            val replicationFactor = td.partitions().iterator().next().replicas().size
-            val configsAsString = config.entries().asScala.filter(!_.isDefault).map { ce => s"${ce.name()}=${ce.value()}" }.mkString(",")
-            println(s"Topic:${td.name()}\tPartitionCount:$numPartitions\tReplicationFactor:$replicationFactor\tConfigs:$configsAsString")
+            val replicationFactor = td.partitions.iterator.next().replicas.size
+            val topicDesc = TopicDescription(topicName, numPartitions, replicationFactor, config, markedForDeletion = false)
+            topicDesc.printDescription()
           }
         }
+
         if (describeOptions.describePartitions) {
-          val computedMinIsrCount = if (opts.reportUnderMinIsrPartitions)
-            allConfigs.get(new ConfigResource(ConfigResource.Type.TOPIC, td.name())).get().get(TopicConfig.MIN_IN_SYNC_REPLICAS_CONFIG).value().toInt else 0
           for (partition <- sortedPartitions) {
-            val partitionDesc = PartitionDescription(
-              topic = td.name(),
-              partition.partition(),
-              leader = Option(partition.leader()).map(_.id()),
-              assignedReplicas = partition.replicas().asScala.map(_.id()),
-              isr = partition.isr().asScala.map(_.id()),
-              minIsrCount = computedMinIsrCount,
-              markedForDeletion = false,
-              describeOptions.describeConfigs
-            )
-
-            if (describeOptions.shouldPrintTopicPartition(partitionDesc)) {
-              printPartition(partitionDesc)
-            }
+            val partitionDesc = PartitionDescription(topicName, partition, Some(config), markedForDeletion = false)
+            describeOptions.maybePrintPartitionDescription(partitionDesc)
           }
         }
       }
@@ -252,7 +300,7 @@ object TopicCommand extends Logging {
 
     override def deleteTopic(opts: TopicCommandOptions): Unit = {
       val topics = getTopics(opts.topic, opts.excludeInternalTopics)
-      ensureTopicExists(topics)
+      ensureTopicExists(topics, opts.topic)
       adminClient.deleteTopics(topics.asJavaCollection).all().get()
     }
 
@@ -282,7 +330,7 @@ object TopicCommand extends Logging {
         if (topic.hasReplicaAssignment)
           adminZkClient.createTopicWithAssignment(topic.name, topic.configsToAdd, topic.replicaAssignment.get)
         else
-          adminZkClient.createTopic(topic.name, topic.partitions.get, topic.replicationFactor, topic.configsToAdd, topic.rackAwareMode)
+          adminZkClient.createTopic(topic.name, topic.partitions.get, topic.replicationFactor.get, topic.configsToAdd, topic.rackAwareMode)
         println(s"Created topic ${topic.name}.")
       } catch  {
         case e: TopicExistsException => if (!topic.ifTopicDoesntExist()) throw e
@@ -302,7 +350,7 @@ object TopicCommand extends Logging {
     override def alterTopic(opts: TopicCommandOptions): Unit = {
       val topics = getTopics(opts.topic, opts.excludeInternalTopics)
       val tp = new CommandTopicPartition(opts)
-      ensureTopicExists(topics, opts.ifExists)
+      ensureTopicExists(topics, opts.topic, !opts.ifExists)
       val adminZkClient = new AdminZkClient(zkClient)
       topics.foreach { topic =>
         val configs = adminZkClient.fetchEntityConfig(ConfigType.Topic, topic)
@@ -318,13 +366,13 @@ object TopicCommand extends Logging {
         }
 
         if(tp.hasPartitions) {
-          if (topic == Topic.GROUP_METADATA_TOPIC_NAME) {
-            throw new IllegalArgumentException("The number of partitions for the offsets topic cannot be changed.")
+          if (Topic.isInternal(topic)) {
+            throw new IllegalArgumentException(s"The number of partitions for the internal topic $topic cannot be changed.")
           }
           println("WARNING: If partitions are increased for a topic that has a key, the partition " +
             "logic or ordering of the messages will be affected")
-          val existingAssignment = zkClient.getReplicaAssignmentForTopics(immutable.Set(topic)).map {
-            case (topicPartition, replicas) => topicPartition.partition -> replicas
+          val existingAssignment = zkClient.getFullReplicaAssignmentForTopics(immutable.Set(topic)).map {
+            case (topicPartition, assignment) => topicPartition.partition -> assignment
           }
           if (existingAssignment.isEmpty)
             throw new InvalidTopicException(s"The topic $topic does not exist")
@@ -339,10 +387,10 @@ object TopicCommand extends Logging {
 
     override def describeTopic(opts: TopicCommandOptions): Unit = {
       val topics = getTopics(opts.topic, opts.excludeInternalTopics)
-      val topicOptWithExits = opts.topic.isDefined && opts.ifExists
-      ensureTopicExists(topics, topicOptWithExits)
-      val liveBrokers = zkClient.getAllBrokersInCluster.map(_.id).toSet
-      val describeOptions = new DescribeOptions(opts, liveBrokers)
+      ensureTopicExists(topics, opts.topic, !opts.ifExists)
+      val liveBrokers = zkClient.getAllBrokersInCluster.map(broker => broker.id -> broker).toMap
+      val liveBrokerIds = liveBrokers.keySet
+      val describeOptions = new DescribeOptions(opts, liveBrokerIds)
       val adminZkClient = new AdminZkClient(zkClient)
 
       for (topic <- topics) {
@@ -353,29 +401,34 @@ object TopicCommand extends Logging {
               val configs = adminZkClient.fetchEntityConfig(ConfigType.Topic, topic).asScala
               if (!opts.reportOverriddenConfigs || configs.nonEmpty) {
                 val numPartitions = topicPartitionAssignment.size
-                val replicationFactor = topicPartitionAssignment.head._2.size
-                val configsAsString = configs.map { case (k, v) => s"$k=$v" }.mkString(",")
-                val markedForDeletionString = if (markedForDeletion) "\tMarkedForDeletion:true" else ""
-                println(s"Topic:$topic\tPartitionCount:$numPartitions\tReplicationFactor:$replicationFactor\tConfigs:$configsAsString$markedForDeletionString")
+                val replicationFactor = topicPartitionAssignment.head._2.replicas.size
+                val config = new JConfig(configs.map{ case (k, v) => new ConfigEntry(k, v) }.asJavaCollection)
+                val topicDesc = TopicDescription(topic, numPartitions, replicationFactor, config, markedForDeletion)
+                topicDesc.printDescription()
               }
             }
             if (describeOptions.describePartitions) {
-              for ((partitionId, assignedReplicas) <- topicPartitionAssignment.toSeq.sortBy(_._1)) {
-                val leaderIsrEpoch = zkClient.getTopicPartitionState(new TopicPartition(topic, partitionId))
-                val partitionDesc = PartitionDescription(
-                  topic,
-                  partitionId,
-                  leader = if (leaderIsrEpoch.isEmpty) None else Option(leaderIsrEpoch.get.leaderAndIsr.leader),
-                  assignedReplicas,
-                  isr = if (leaderIsrEpoch.isEmpty) Seq.empty[Int] else leaderIsrEpoch.get.leaderAndIsr.isr,
-                  minIsrCount = 0,
-                  markedForDeletion,
-                  describeOptions.describeConfigs
-                )
-
-                if (describeOptions.shouldPrintTopicPartition(partitionDesc)) {
-                  printPartition(partitionDesc)
+              for ((partitionId, replicaAssignment) <- topicPartitionAssignment.toSeq.sortBy(_._1)) {
+                val assignedReplicas = replicaAssignment.replicas
+                val tp = new TopicPartition(topic, partitionId)
+                val (leaderOpt, isr) =  zkClient.getTopicPartitionState(tp).map(_.leaderAndIsr) match {
+                  case Some(leaderAndIsr) => (leaderAndIsr.leaderOpt, leaderAndIsr.isr)
+                  case None => (None, Seq.empty[Int])
                 }
+
+                def asNode(brokerId: Int): Node = {
+                  liveBrokers.get(brokerId) match {
+                    case Some(broker) => broker.node(broker.endPoints.head.listenerName)
+                    case None => new Node(brokerId, "", -1)
+                  }
+                }
+
+                val info = new TopicPartitionInfo(partitionId, leaderOpt.map(asNode).orNull,
+                  assignedReplicas.map(asNode).toList.asJava,
+                  isr.map(asNode).toList.asJava)
+
+                val partitionDesc = PartitionDescription(topic, info, config = None, markedForDeletion)
+                describeOptions.maybePrintPartitionDescription(partitionDesc)
               }
             }
           case None =>
@@ -386,7 +439,7 @@ object TopicCommand extends Logging {
 
     override def deleteTopic(opts: TopicCommandOptions): Unit = {
       val topics = getTopics(opts.topic, opts.excludeInternalTopics)
-      ensureTopicExists(topics, opts.ifExists)
+      ensureTopicExists(topics, opts.topic, !opts.ifExists)
       topics.foreach { topic =>
         try {
           if (Topic.isInternal(topic)) {
@@ -401,14 +454,14 @@ object TopicCommand extends Logging {
             println(s"Topic $topic is already marked for deletion.")
           case e: AdminOperationException =>
             throw e
-          case _: Throwable =>
-            throw new AdminOperationException(s"Error while deleting topic $topic")
+          case e: Throwable =>
+            throw new AdminOperationException(s"Error while deleting topic $topic", e)
         }
       }
     }
 
     override def getTopics(topicWhitelist: Option[String], excludeInternalTopics: Boolean = false): Seq[String] = {
-      val allTopics = zkClient.getAllTopicsInCluster.sorted
+      val allTopics = zkClient.getAllTopicsInCluster.toSeq.sorted
       doGetTopics(allTopics, topicWhitelist, excludeInternalTopics)
     }
 
@@ -418,29 +471,20 @@ object TopicCommand extends Logging {
   /**
     * ensures topic existence and throws exception if topic doesn't exist
     *
-    * @param opts
-    * @param topics
-    * @param topicOptWithExists
+    * @param foundTopics Topics that were found to match the requested topic name.
+    * @param requestedTopic Name of the topic that was requested.
+    * @param requireTopicExists Indicates if the topic needs to exist for the operation to be successful.
+    *                           If set to true, the command will throw an exception if the topic with the
+    *                           requested name does not exist.
     */
-  private def ensureTopicExists(topics: Seq[String], topicOptWithExists: Boolean = false) = {
-    if (topics.isEmpty && !topicOptWithExists) {
+  private def ensureTopicExists(foundTopics: Seq[String], requestedTopic: Option[String], requireTopicExists: Boolean = true) = {
+    // If no topic name was mentioned, do not need to throw exception.
+    if (requestedTopic.isDefined && requireTopicExists && foundTopics.isEmpty) {
       // If given topic doesn't exist then throw exception
-      throw new IllegalArgumentException(s"Topics in [${topics.mkString(",")}] does not exist")
+      throw new IllegalArgumentException(s"Topic '${requestedTopic.get}' does not exist as expected")
     }
   }
 
-  private def printPartition(tp: PartitionDescription): Unit = {
-    val markedForDeletionString =
-    if (tp.markedForDeletion && !tp.describeConfigs) "\tMarkedForDeletion: true" else ""
-    print("\tTopic: " + tp.topic)
-    print("\tPartition: " + tp.partition)
-    print("\tLeader: " + (if(tp.leader.isDefined) tp.leader.get else "none"))
-    print("\tReplicas: " + tp.assignedReplicas.mkString(","))
-    print("\tIsr: " + tp.isr.mkString(","))
-    print(markedForDeletionString)
-    println()
-  }
-
   private def doGetTopics(allTopics: Seq[String], topicWhitelist: Option[String], excludeInternalTopics: Boolean): Seq[String] = {
     if (topicWhitelist.isDefined) {
       val topicsFilter = Whitelist(topicWhitelist.get)
@@ -531,11 +575,11 @@ object TopicCommand extends Logging {
                            .describedAs("name")
                            .ofType(classOf[String])
     private val partitionsOpt = parser.accepts("partitions", "The number of partitions for the topic being created or " +
-      "altered (WARNING: If partitions are increased for a topic that has a key, the partition logic or ordering of the messages will be affected")
+      "altered (WARNING: If partitions are increased for a topic that has a key, the partition logic or ordering of the messages will be affected). If not supplied for create, defaults to the cluster default.")
                            .withRequiredArg
                            .describedAs("# of partitions")
                            .ofType(classOf[java.lang.Integer])
-    private val replicationFactorOpt = parser.accepts("replication-factor", "The replication factor for each partition in the topic being created.")
+    private val replicationFactorOpt = parser.accepts("replication-factor", "The replication factor for each partition in the topic being created. If not supplied, defaults to the cluster default.")
                            .withRequiredArg
                            .describedAs("replication factor")
                            .ofType(classOf[java.lang.Integer])
@@ -545,29 +589,34 @@ object TopicCommand extends Logging {
                                         "broker_id_for_part2_replica1 : broker_id_for_part2_replica2 , ...")
                            .ofType(classOf[String])
     private val reportUnderReplicatedPartitionsOpt = parser.accepts("under-replicated-partitions",
-                                                            "if set when describing topics, only show under replicated partitions")
+      "if set when describing topics, only show under replicated partitions")
     private val reportUnavailablePartitionsOpt = parser.accepts("unavailable-partitions",
-                                                            "if set when describing topics, only show partitions whose leader is not available")
+      "if set when describing topics, only show partitions whose leader is not available")
     private val reportUnderMinIsrPartitionsOpt = parser.accepts("under-min-isr-partitions",
-                                                            "if set when describing topics, only show partitions whose isr count is less than the configured minimum. Not supported with the --zookeeper option.")
+      "if set when describing topics, only show partitions whose isr count is less than the configured minimum. Not supported with the --zookeeper option.")
+    private val reportAtMinIsrPartitionsOpt = parser.accepts("at-min-isr-partitions",
+      "if set when describing topics, only show partitions whose isr count is equal to the configured minimum. Not supported with the --zookeeper option.")
     private val topicsWithOverridesOpt = parser.accepts("topics-with-overrides",
-                                                "if set when describing topics, only show topics that have overridden configs")
+      "if set when describing topics, only show topics that have overridden configs")
     private val ifExistsOpt = parser.accepts("if-exists",
-                                     "if set when altering or deleting or describing topics, the action will only execute if the topic exists. Not supported with the --bootstrap-server option.")
+      "if set when altering or deleting or describing topics, the action will only execute if the topic exists. Not supported with the --bootstrap-server option.")
     private val ifNotExistsOpt = parser.accepts("if-not-exists",
-                                        "if set when creating topics, the action will only execute if the topic does not already exist. Not supported with the --bootstrap-server option.")
+      "if set when creating topics, the action will only execute if the topic does not already exist. Not supported with the --bootstrap-server option.")
 
     private val disableRackAware = parser.accepts("disable-rack-aware", "Disable rack aware replica assignment")
 
     // This is not currently used, but we keep it for compatibility
     parser.accepts("force", "Suppress console prompts")
 
-    private val excludeInternalTopicOpt = parser.accepts("exclude-internal", "exclude internal topics when running list or describe command. The internal topics will be listed by default")
+    private val excludeInternalTopicOpt = parser.accepts("exclude-internal",
+      "exclude internal topics when running list or describe command. The internal topics will be listed by default")
 
     options = parser.parse(args : _*)
 
     private val allTopicLevelOpts: Set[OptionSpec[_]] = Set(alterOpt, createOpt, describeOpt, listOpt, deleteOpt)
 
+    private val allReplicationReportOpts: Set[OptionSpec[_]] = Set(reportUnderReplicatedPartitionsOpt, reportUnderMinIsrPartitionsOpt, reportAtMinIsrPartitionsOpt, reportUnavailablePartitionsOpt)
+
     def has(builder: OptionSpec[_]): Boolean = options.has(builder)
     def valueAsOption[A](option: OptionSpec[A], defaultValue: Option[A] = None): Option[A] = if (has(option)) Some(options.valueOf(option)) else defaultValue
     def valuesAsOption[A](option: OptionSpec[A], defaultValue: Option[util.List[A]] = None): Option[util.List[A]] = if (has(option)) Some(options.valuesOf(option)) else defaultValue
@@ -593,6 +642,7 @@ object TopicCommand extends Logging {
     def reportUnderReplicatedPartitions: Boolean = has(reportUnderReplicatedPartitionsOpt)
     def reportUnavailablePartitions: Boolean = has(reportUnavailablePartitionsOpt)
     def reportUnderMinIsrPartitions: Boolean = has(reportUnderMinIsrPartitionsOpt)
+    def reportAtMinIsrPartitions: Boolean = has(reportAtMinIsrPartitionsOpt)
     def reportOverriddenConfigs: Boolean = has(topicsWithOverridesOpt)
     def ifExists: Boolean = has(ifExistsOpt)
     def ifNotExists: Boolean = has(ifNotExistsOpt)
@@ -600,7 +650,7 @@ object TopicCommand extends Logging {
     def topicConfig: Option[util.List[String]] = valuesAsOption(configOpt)
     def configsToDelete: Option[util.List[String]] = valuesAsOption(deleteConfigOpt)
 
-    def checkArgs() {
+    def checkArgs(): Unit = {
       if (args.length == 0)
         CommandLineUtils.printUsageAndDie(parser, "Create, delete, describe, or change a topic.")
 
@@ -621,14 +671,15 @@ object TopicCommand extends Logging {
         CommandLineUtils.checkRequiredArgs(parser, options, topicOpt)
       if (!has(listOpt) && !has(describeOpt))
         CommandLineUtils.checkRequiredArgs(parser, options, topicOpt)
-      if (has(createOpt) && !has(replicaAssignmentOpt))
+      if (has(createOpt) && !has(replicaAssignmentOpt) && has(zkConnectOpt))
         CommandLineUtils.checkRequiredArgs(parser, options, partitionsOpt, replicationFactorOpt)
-      if (has(bootstrapServerOpt) && has(alterOpt))
+      if (has(bootstrapServerOpt) && has(alterOpt)) {
+        CommandLineUtils.checkInvalidArgsSet(parser, options, Set(bootstrapServerOpt, configOpt), Set(alterOpt))
         CommandLineUtils.checkRequiredArgs(parser, options, partitionsOpt)
+      }
 
       // check invalid args
       CommandLineUtils.checkInvalidArgs(parser, options, configOpt, allTopicLevelOpts -- Set(alterOpt, createOpt))
-      CommandLineUtils.checkInvalidArgsSet(parser, options, Set(bootstrapServerOpt, configOpt), Set(alterOpt))
       CommandLineUtils.checkInvalidArgs(parser, options, deleteConfigOpt, allTopicLevelOpts -- Set(alterOpt) ++ Set(bootstrapServerOpt))
       CommandLineUtils.checkInvalidArgs(parser, options, partitionsOpt, allTopicLevelOpts -- Set(alterOpt, createOpt))
       CommandLineUtils.checkInvalidArgs(parser, options, replicationFactorOpt, allTopicLevelOpts -- Set(createOpt))
@@ -636,13 +687,15 @@ object TopicCommand extends Logging {
       if(options.has(createOpt))
           CommandLineUtils.checkInvalidArgs(parser, options, replicaAssignmentOpt, Set(partitionsOpt, replicationFactorOpt))
       CommandLineUtils.checkInvalidArgs(parser, options, reportUnderReplicatedPartitionsOpt,
-        allTopicLevelOpts -- Set(describeOpt) + reportUnderMinIsrPartitionsOpt + reportUnavailablePartitionsOpt + topicsWithOverridesOpt)
+        allTopicLevelOpts -- Set(describeOpt) ++ allReplicationReportOpts - reportUnderReplicatedPartitionsOpt + topicsWithOverridesOpt)
       CommandLineUtils.checkInvalidArgs(parser, options, reportUnderMinIsrPartitionsOpt,
-        allTopicLevelOpts -- Set(describeOpt) + reportUnderReplicatedPartitionsOpt + reportUnavailablePartitionsOpt + topicsWithOverridesOpt + zkConnectOpt)
+        allTopicLevelOpts -- Set(describeOpt) ++ allReplicationReportOpts - reportUnderMinIsrPartitionsOpt + topicsWithOverridesOpt + zkConnectOpt)
+      CommandLineUtils.checkInvalidArgs(parser, options, reportAtMinIsrPartitionsOpt,
+        allTopicLevelOpts -- Set(describeOpt) ++ allReplicationReportOpts - reportAtMinIsrPartitionsOpt + topicsWithOverridesOpt + zkConnectOpt)
       CommandLineUtils.checkInvalidArgs(parser, options, reportUnavailablePartitionsOpt,
-        allTopicLevelOpts -- Set(describeOpt) + reportUnderReplicatedPartitionsOpt + reportUnderReplicatedPartitionsOpt + topicsWithOverridesOpt)
+        allTopicLevelOpts -- Set(describeOpt) ++ allReplicationReportOpts - reportUnavailablePartitionsOpt + topicsWithOverridesOpt)
       CommandLineUtils.checkInvalidArgs(parser, options, topicsWithOverridesOpt,
-        allTopicLevelOpts -- Set(describeOpt) + reportUnderReplicatedPartitionsOpt + reportUnderMinIsrPartitionsOpt + reportUnavailablePartitionsOpt)
+        allTopicLevelOpts -- Set(describeOpt) ++ allReplicationReportOpts)
       CommandLineUtils.checkInvalidArgs(parser, options, ifExistsOpt, allTopicLevelOpts -- Set(alterOpt, deleteOpt, describeOpt) ++ Set(bootstrapServerOpt))
       CommandLineUtils.checkInvalidArgs(parser, options, ifNotExistsOpt, allTopicLevelOpts -- Set(createOpt) ++ Set(bootstrapServerOpt))
       CommandLineUtils.checkInvalidArgs(parser, options, excludeInternalTopicOpt, allTopicLevelOpts -- Set(listOpt, describeOpt))
diff --git a/core/src/main/scala/kafka/admin/ZkSecurityMigrator.scala b/core/src/main/scala/kafka/admin/ZkSecurityMigrator.scala
index 5cab801bf3fd8..0790fde01c89e 100644
--- a/core/src/main/scala/kafka/admin/ZkSecurityMigrator.scala
+++ b/core/src/main/scala/kafka/admin/ZkSecurityMigrator.scala
@@ -19,7 +19,6 @@ package kafka.admin
 
 import kafka.utils.{CommandDefaultOptions, CommandLineUtils, Logging}
 import kafka.zk.{KafkaZkClient, ZkData, ZkSecurityMigratorUtils}
-import org.I0Itec.zkclient.exception.ZkException
 import org.apache.kafka.common.security.JaasUtils
 import org.apache.kafka.common.utils.Time
 import org.apache.zookeeper.AsyncCallback.{ChildrenCallback, StatCallback}
@@ -62,7 +61,7 @@ object ZkSecurityMigrator extends Logging {
                       + "znodes as part of the process of setting up ZooKeeper "
                       + "authentication.")
 
-  def run(args: Array[String]) {
+  def run(args: Array[String]): Unit = {
     val jaasFile = System.getProperty(JaasUtils.JAVA_LOGIN_CONFIG_PARAM)
     val opts = new ZkSecurityMigratorOptions(args)
 
@@ -100,7 +99,7 @@ object ZkSecurityMigrator extends Logging {
     migrator.run()
   }
 
-  def main(args: Array[String]) {
+  def main(args: Array[String]): Unit = {
     try {
       run(args)
     } catch {
@@ -160,7 +159,7 @@ class ZkSecurityMigrator(zkClient: KafkaZkClient) extends Logging {
     def processResult(rc: Int,
                       path: String,
                       ctx: Object,
-                      children: java.util.List[String]) {
+                      children: java.util.List[String]): Unit = {
       val zkHandle = zkSecurityMigratorUtils.currentZooKeeper
       val promise = ctx.asInstanceOf[Promise[String]]
       Code.get(rc) match {
@@ -182,10 +181,10 @@ class ZkSecurityMigrator(zkClient: KafkaZkClient) extends Logging {
           // Starting a new session isn't really a problem, but it'd complicate
           // the logic of the tool, so we quit and let the user re-run it.
           System.out.println("ZooKeeper session expired while changing ACLs")
-          promise failure ZkException.create(KeeperException.create(Code.get(rc)))
+          promise failure KeeperException.create(Code.get(rc))
         case _ =>
           System.out.println("Unexpected return code: %d".format(rc))
-          promise failure ZkException.create(KeeperException.create(Code.get(rc)))
+          promise failure KeeperException.create(Code.get(rc))
       }
     }
   }
@@ -194,7 +193,7 @@ class ZkSecurityMigrator(zkClient: KafkaZkClient) extends Logging {
     def processResult(rc: Int,
                       path: String,
                       ctx: Object,
-                      stat: Stat) {
+                      stat: Stat): Unit = {
       val zkHandle = zkSecurityMigratorUtils.currentZooKeeper
       val promise = ctx.asInstanceOf[Promise[String]]
 
@@ -211,10 +210,10 @@ class ZkSecurityMigrator(zkClient: KafkaZkClient) extends Logging {
           // Starting a new session isn't really a problem, but it'd complicate
           // the logic of the tool, so we quit and let the user re-run it.
           System.out.println("ZooKeeper session expired while changing ACLs")
-          promise failure ZkException.create(KeeperException.create(Code.get(rc)))
+          promise failure KeeperException.create(Code.get(rc))
         case _ =>
           System.out.println("Unexpected return code: %d".format(rc))
-          promise failure ZkException.create(KeeperException.create(Code.get(rc)))
+          promise failure KeeperException.create(Code.get(rc))
       }
     }
   }
diff --git a/core/src/main/scala/kafka/api/ApiUtils.scala b/core/src/main/scala/kafka/api/ApiUtils.scala
index 4a0c8b0c636b7..1644358f7bb3a 100644
--- a/core/src/main/scala/kafka/api/ApiUtils.scala
+++ b/core/src/main/scala/kafka/api/ApiUtils.scala
@@ -45,7 +45,7 @@ object ApiUtils {
    * @param buffer The buffer to write to
    * @param string The string to write
    */
-  def writeShortString(buffer: ByteBuffer, string: String) {
+  def writeShortString(buffer: ByteBuffer, string: String): Unit = {
     if(string == null) {
       buffer.putShort(-1)
     } else {
diff --git a/core/src/main/scala/kafka/api/ApiVersion.scala b/core/src/main/scala/kafka/api/ApiVersion.scala
index 8f78a960cf91d..fb69cba087d89 100644
--- a/core/src/main/scala/kafka/api/ApiVersion.scala
+++ b/core/src/main/scala/kafka/api/ApiVersion.scala
@@ -86,7 +86,17 @@ object ApiVersion {
     // LeaderAdnIsrRequest V2, UpdateMetadataRequest V5, StopReplicaRequest V1
     KAFKA_2_2_IV0,
     // New error code for ListOffsets when a new leader is lagging behind former HW (KIP-207)
-    KAFKA_2_2_IV1
+    KAFKA_2_2_IV1,
+    // Introduced static membership.
+    KAFKA_2_3_IV0,
+    // Add rack_id to FetchRequest, preferred_read_replica to FetchResponse, and replica_id to OffsetsForLeaderRequest
+    KAFKA_2_3_IV1,
+    // Add cacheable broker epoch to UpdateMetadataRequest
+    KAFKA_2_3_IV2,
+    // Add adding_replicas and removing_replicas fields to LeaderAndIsrRequest
+    KAFKA_2_4_IV0,
+    // Flexible version support in inter-broker APIs
+    KAFKA_2_4_IV1
   )
 
   // Map keys are the union of the short and full versions
@@ -298,6 +308,42 @@ case object KAFKA_2_2_IV1 extends DefaultApiVersion {
   val id: Int = 21
 }
 
+case object KAFKA_2_3_IV0 extends DefaultApiVersion {
+  val shortVersion: String = "2.3"
+  val subVersion = "IV0"
+  val recordVersion = RecordVersion.V2
+  val id: Int = 22
+}
+
+case object KAFKA_2_3_IV1 extends DefaultApiVersion {
+  val shortVersion: String = "2.3"
+  val subVersion = "IV1"
+  val recordVersion = RecordVersion.V2
+  val id: Int = 23
+}
+
+case object KAFKA_2_3_IV2 extends DefaultApiVersion {
+  val shortVersion: String = "2.3"
+  val subVersion = "IV2"
+  val recordVersion = RecordVersion.V2
+  val id: Int = 24
+}
+
+case object KAFKA_2_4_IV0 extends DefaultApiVersion {
+  val shortVersion: String = "2.4"
+  val subVersion = "IV0"
+  val recordVersion = RecordVersion.V2
+  val id: Int = 25
+}
+
+case object KAFKA_2_4_IV1 extends DefaultApiVersion {
+  val shortVersion: String = "2.4"
+  val subVersion = "IV1"
+  val recordVersion = RecordVersion.V2
+  val id: Int = 26
+}
+
+
 object ApiVersionValidator extends Validator {
 
   override def ensureValid(name: String, value: Any): Unit = {
diff --git a/core/src/main/scala/kafka/api/LeaderAndIsr.scala b/core/src/main/scala/kafka/api/LeaderAndIsr.scala
index cb59575ea7a4b..ebfa710d3de00 100644
--- a/core/src/main/scala/kafka/api/LeaderAndIsr.scala
+++ b/core/src/main/scala/kafka/api/LeaderAndIsr.scala
@@ -40,6 +40,10 @@ case class LeaderAndIsr(leader: Int,
 
   def newEpochAndZkVersion = newLeaderAndIsr(leader, isr)
 
+  def leaderOpt: Option[Int] = {
+    if (leader == LeaderAndIsr.NoLeader) None else Some(leader)
+  }
+
   override def toString: String = {
     s"LeaderAndIsr(leader=$leader, leaderEpoch=$leaderEpoch, isr=$isr, zkVersion=$zkVersion)"
   }
diff --git a/core/src/main/scala/kafka/api/package.scala b/core/src/main/scala/kafka/api/package.scala
new file mode 100644
index 0000000000000..98d004542292f
--- /dev/null
+++ b/core/src/main/scala/kafka/api/package.scala
@@ -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 kafka
+
+import org.apache.kafka.common.ElectionType
+import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.requests.ElectLeadersRequest
+import scala.collection.JavaConverters._
+
+package object api {
+  implicit final class ElectLeadersRequestOps(val self: ElectLeadersRequest) extends AnyVal {
+    def topicPartitions: Set[TopicPartition] = {
+      if (self.data.topicPartitions == null) {
+        Set.empty
+      } else {
+        self.data.topicPartitions.asScala.iterator.flatMap { topicPartition =>
+          topicPartition.partitionId.asScala.map { partitionId =>
+            new TopicPartition(topicPartition.topic, partitionId)
+          }
+        }.toSet
+      }
+    }
+
+    def electionType: ElectionType = {
+      if (self.version == 0) {
+        ElectionType.PREFERRED
+      } else {
+        ElectionType.valueOf(self.data.electionType)
+      }
+    }
+  }
+}
diff --git a/core/src/main/scala/kafka/cluster/Broker.scala b/core/src/main/scala/kafka/cluster/Broker.scala
index 010698227daae..e8e621f03555d 100755
--- a/core/src/main/scala/kafka/cluster/Broker.scala
+++ b/core/src/main/scala/kafka/cluster/Broker.scala
@@ -17,10 +17,24 @@
 
 package kafka.cluster
 
+import java.util
+
 import kafka.common.BrokerEndPointNotAvailableException
-import org.apache.kafka.common.Node
+import kafka.server.KafkaConfig
+import org.apache.kafka.common.{ClusterResource, Endpoint, Node}
 import org.apache.kafka.common.network.ListenerName
 import org.apache.kafka.common.security.auth.SecurityProtocol
+import org.apache.kafka.server.authorizer.AuthorizerServerInfo
+
+import scala.collection.Seq
+import scala.collection.JavaConverters._
+
+object Broker {
+  private[cluster] case class ServerInfo(clusterResource: ClusterResource,
+                                         brokerId: Int,
+                                         endpoints: util.List[Endpoint],
+                                         interBrokerEndpoint: Endpoint) extends AuthorizerServerInfo
+}
 
 /**
  * A Kafka broker.
@@ -57,9 +71,19 @@ case class Broker(id: Int, endPoints: Seq[EndPoint], rack: Option[String]) {
     endPointsMap.get(listenerName).map(endpoint => new Node(id, endpoint.host, endpoint.port, rack.orNull))
 
   def brokerEndPoint(listenerName: ListenerName): BrokerEndPoint = {
-    val endpoint = endPointsMap.getOrElse(listenerName,
-      throw new BrokerEndPointNotAvailableException(s"End point with listener name ${listenerName.value} not found for broker $id"))
+    val endpoint = endPoint(listenerName)
     new BrokerEndPoint(id, endpoint.host, endpoint.port)
   }
 
+  def endPoint(listenerName: ListenerName): EndPoint = {
+    endPointsMap.getOrElse(listenerName,
+      throw new BrokerEndPointNotAvailableException(s"End point with listener name ${listenerName.value} not found for broker $id"))
+  }
+
+  def toServerInfo(clusterId: String, config: KafkaConfig): AuthorizerServerInfo = {
+    val clusterResource: ClusterResource = new ClusterResource(clusterId)
+    val interBrokerEndpoint: Endpoint = endPoint(config.interBrokerListenerName).toJava
+    val brokerEndpoints: util.List[Endpoint] = endPoints.toList.map(_.toJava).asJava
+    Broker.ServerInfo(clusterResource, id, brokerEndpoints, interBrokerEndpoint)
+  }
 }
diff --git a/core/src/main/scala/kafka/cluster/EndPoint.scala b/core/src/main/scala/kafka/cluster/EndPoint.scala
index 2bca5c8797673..2f8229a38865c 100644
--- a/core/src/main/scala/kafka/cluster/EndPoint.scala
+++ b/core/src/main/scala/kafka/cluster/EndPoint.scala
@@ -17,7 +17,7 @@
 
 package kafka.cluster
 
-import org.apache.kafka.common.KafkaException
+import org.apache.kafka.common.{Endpoint => JEndpoint, KafkaException}
 import org.apache.kafka.common.network.ListenerName
 import org.apache.kafka.common.security.auth.SecurityProtocol
 import org.apache.kafka.common.utils.Utils
@@ -71,4 +71,8 @@ case class EndPoint(host: String, port: Int, listenerName: ListenerName, securit
         Utils.formatAddress(host, port)
     listenerName.value + "://" + hostport
   }
+
+  def toJava: JEndpoint = {
+    new JEndpoint(listenerName.value, securityProtocol, host, port)
+  }
 }
diff --git a/core/src/main/scala/kafka/cluster/Partition.scala b/core/src/main/scala/kafka/cluster/Partition.scala
index 698fb486fdcb5..2d28aa12b69bc 100755
--- a/core/src/main/scala/kafka/cluster/Partition.scala
+++ b/core/src/main/scala/kafka/cluster/Partition.scala
@@ -16,21 +16,23 @@
  */
 package kafka.cluster
 
-import java.util.Optional
+import com.yammer.metrics.core.Gauge
 import java.util.concurrent.locks.ReentrantReadWriteLock
+import java.util.{Optional, Properties}
 
-import com.yammer.metrics.core.Gauge
-import kafka.api.{ApiVersion, LeaderAndIsr, Request}
+import kafka.api.{ApiVersion, LeaderAndIsr}
 import kafka.common.UnexpectedAppendOffsetException
 import kafka.controller.KafkaController
 import kafka.log._
 import kafka.metrics.KafkaMetricsGroup
 import kafka.server._
+import kafka.server.checkpoints.OffsetCheckpoints
 import kafka.utils.CoreUtils.{inReadLock, inWriteLock}
 import kafka.utils._
 import kafka.zk.{AdminZkClient, KafkaZkClient}
 import org.apache.kafka.common.TopicPartition
 import org.apache.kafka.common.errors._
+import org.apache.kafka.common.message.LeaderAndIsrRequestData.LeaderAndIsrPartitionState
 import org.apache.kafka.common.protocol.Errors
 import org.apache.kafka.common.protocol.Errors._
 import org.apache.kafka.common.record.FileRecords.TimestampAndOffset
@@ -40,21 +42,116 @@ import org.apache.kafka.common.requests._
 import org.apache.kafka.common.utils.Time
 
 import scala.collection.JavaConverters._
-import scala.collection.Map
+import scala.collection.{Map, Seq}
+
+trait PartitionStateStore {
+  def fetchTopicConfig(): Properties
+  def shrinkIsr(controllerEpoch: Int, leaderAndIsr: LeaderAndIsr): Option[Int]
+  def expandIsr(controllerEpoch: Int, leaderAndIsr: LeaderAndIsr): Option[Int]
+}
+
+class ZkPartitionStateStore(topicPartition: TopicPartition,
+                            zkClient: KafkaZkClient,
+                            replicaManager: ReplicaManager) extends PartitionStateStore {
+
+  override def fetchTopicConfig(): Properties = {
+    val adminZkClient = new AdminZkClient(zkClient)
+    adminZkClient.fetchEntityConfig(ConfigType.Topic, topicPartition.topic)
+  }
+
+  override def shrinkIsr(controllerEpoch: Int, leaderAndIsr: LeaderAndIsr): Option[Int] = {
+    val newVersionOpt = updateIsr(controllerEpoch, leaderAndIsr)
+    if (newVersionOpt.isDefined)
+      replicaManager.isrShrinkRate.mark()
+    newVersionOpt
+  }
+
+  override def expandIsr(controllerEpoch: Int, leaderAndIsr: LeaderAndIsr): Option[Int] = {
+    val newVersionOpt = updateIsr(controllerEpoch, leaderAndIsr)
+    if (newVersionOpt.isDefined)
+      replicaManager.isrExpandRate.mark()
+    newVersionOpt
+  }
+
+  private def updateIsr(controllerEpoch: Int, leaderAndIsr: LeaderAndIsr): Option[Int] = {
+    val (updateSucceeded, newVersion) = ReplicationUtils.updateLeaderAndIsr(zkClient, topicPartition,
+      leaderAndIsr, controllerEpoch)
+
+    if (updateSucceeded) {
+      replicaManager.recordIsrChange(topicPartition)
+      Some(newVersion)
+    } else {
+      replicaManager.failedIsrUpdatesRate.mark()
+      None
+    }
+  }
+}
+
+class DelayedOperations(topicPartition: TopicPartition,
+                        produce: DelayedOperationPurgatory[DelayedProduce],
+                        fetch: DelayedOperationPurgatory[DelayedFetch],
+                        deleteRecords: DelayedOperationPurgatory[DelayedDeleteRecords]) {
 
-object Partition {
+  def checkAndCompleteAll(): Unit = {
+    val requestKey = TopicPartitionOperationKey(topicPartition)
+    fetch.checkAndComplete(requestKey)
+    produce.checkAndComplete(requestKey)
+    deleteRecords.checkAndComplete(requestKey)
+  }
+
+  def checkAndCompleteFetch(): Unit = {
+    fetch.checkAndComplete(TopicPartitionOperationKey(topicPartition))
+  }
+
+  def checkAndCompleteProduce(): Unit = {
+    produce.checkAndComplete(TopicPartitionOperationKey(topicPartition))
+  }
+
+  def checkAndCompleteDeleteRecords(): Unit = {
+    deleteRecords.checkAndComplete(TopicPartitionOperationKey(topicPartition))
+  }
+
+  def numDelayedDelete: Int = deleteRecords.numDelayed
+
+  def numDelayedFetch: Int = fetch.numDelayed
+
+  def numDelayedProduce: Int = produce.numDelayed
+}
+
+object Partition extends KafkaMetricsGroup {
   def apply(topicPartition: TopicPartition,
             time: Time,
             replicaManager: ReplicaManager): Partition = {
+    val zkIsrBackingStore = new ZkPartitionStateStore(
+      topicPartition,
+      replicaManager.zkClient,
+      replicaManager)
+
+    val delayedOperations = new DelayedOperations(
+      topicPartition,
+      replicaManager.delayedProducePurgatory,
+      replicaManager.delayedFetchPurgatory,
+      replicaManager.delayedDeleteRecordsPurgatory)
+
     new Partition(topicPartition,
-      isOffline = false,
       replicaLagTimeMaxMs = replicaManager.config.replicaLagTimeMaxMs,
       interBrokerProtocolVersion = replicaManager.config.interBrokerProtocolVersion,
       localBrokerId = replicaManager.config.brokerId,
       time = time,
-      replicaManager = replicaManager,
-      logManager = replicaManager.logManager,
-      zkClient = replicaManager.zkClient)
+      stateStore = zkIsrBackingStore,
+      delayedOperations = delayedOperations,
+      metadataCache = replicaManager.metadataCache,
+      logManager = replicaManager.logManager)
+  }
+
+  def removeMetrics(topicPartition: TopicPartition): Unit = {
+    val tags = Map("topic" -> topicPartition.topic, "partition" -> topicPartition.partition.toString)
+    removeMetric("UnderReplicated", tags)
+    removeMetric("UnderMinIsr", tags)
+    removeMetric("InSyncReplicasCount", tags)
+    removeMetric("ReplicasCount", tags)
+    removeMetric("LastStableOffsetLag", tags)
+    removeMetric("AtMinIsr", tags)
   }
 }
 
@@ -62,20 +159,19 @@ object Partition {
  * Data structure that represents a topic partition. The leader maintains the AR, ISR, CUR, RAR
  */
 class Partition(val topicPartition: TopicPartition,
-                val isOffline: Boolean,
-                private val replicaLagTimeMaxMs: Long,
-                private val interBrokerProtocolVersion: ApiVersion,
-                private val localBrokerId: Int,
-                private val time: Time,
-                private val replicaManager: ReplicaManager,
-                private val logManager: LogManager,
-                private val zkClient: KafkaZkClient) extends Logging with KafkaMetricsGroup {
+                replicaLagTimeMaxMs: Long,
+                interBrokerProtocolVersion: ApiVersion,
+                localBrokerId: Int,
+                time: Time,
+                stateStore: PartitionStateStore,
+                delayedOperations: DelayedOperations,
+                metadataCache: MetadataCache,
+                logManager: LogManager) extends Logging with KafkaMetricsGroup {
 
   def topic: String = topicPartition.topic
   def partitionId: Int = topicPartition.partition
 
-  // allReplicasMap includes both assigned replicas and the future replica if there is ongoing replica movement
-  private val allReplicasMap = new Pool[Int, Replica]
+  private val remoteReplicasMap = new Pool[Int, Replica]
   // The read lock is only required when multiple reads are executed and needs to be in a consistent manner
   private val leaderIsrUpdateLock = new ReentrantReadWriteLock
   private var zkVersion: Int = LeaderAndIsr.initialZKVersion
@@ -84,7 +180,17 @@ class Partition(val topicPartition: TopicPartition,
   // defined when this broker is leader for partition
   @volatile private var leaderEpochStartOffsetOpt: Option[Long] = None
   @volatile var leaderReplicaIdOpt: Option[Int] = None
-  @volatile var inSyncReplicas: Set[Replica] = Set.empty[Replica]
+  @volatile var inSyncReplicaIds = Set.empty[Int]
+  // An ordered sequence of all the valid broker ids that were assigned to this topic partition
+  @volatile var allReplicaIds = Seq.empty[Int]
+
+  // Logs belonging to this partition. Majority of time it will be only one log, but if log directory
+  // is getting changed (as a result of ReplicaAlterLogDirs command), we may have two logs until copy
+  // completes and a switch to new location is performed.
+  // log and futureLog variables defined below are used to capture this
+  @volatile var log: Option[Log] = None
+  // If ReplicaAlterLogDir command is in progress, this is future location of the log
+  @volatile var futureLog: Option[Log] = None
 
   /* Epoch of the controller that last changed the leader. This needs to be initialized correctly upon broker startup.
    * One way of doing that is through the controller's start replica state change command. When a new broker starts up
@@ -94,69 +200,77 @@ class Partition(val topicPartition: TopicPartition,
   private var controllerEpoch: Int = KafkaController.InitialControllerEpoch
   this.logIdent = s"[Partition $topicPartition broker=$localBrokerId] "
 
-  private def isReplicaLocal(replicaId: Int) : Boolean = replicaId == localBrokerId || replicaId == Request.FutureLocalReplicaId
-
   private val tags = Map("topic" -> topic, "partition" -> partitionId.toString)
 
-  // Do not create metrics if this partition is ReplicaManager.OfflinePartition
-  if (!isOffline) {
-    newGauge("UnderReplicated",
-      new Gauge[Int] {
-        def value = {
-          if (isUnderReplicated) 1 else 0
-        }
-      },
-      tags
-    )
-
-    newGauge("InSyncReplicasCount",
-      new Gauge[Int] {
-        def value = {
-          if (isLeaderReplicaLocal) inSyncReplicas.size else 0
-        }
-      },
-      tags
-    )
-
-    newGauge("UnderMinIsr",
-      new Gauge[Int] {
-        def value = {
-          if (isUnderMinIsr) 1 else 0
-        }
-      },
-      tags
-    )
-
-    newGauge("ReplicasCount",
-      new Gauge[Int] {
-        def value = {
-          if (isLeaderReplicaLocal) assignedReplicas.size else 0
-        }
-      },
-      tags
-    )
-
-    newGauge("LastStableOffsetLag",
-      new Gauge[Long] {
-        def value = {
-          leaderReplicaIfLocal.map { replica =>
-            replica.highWatermark.messageOffset - replica.lastStableOffset.messageOffset
-          }.getOrElse(0)
-        }
-      },
-      tags
-    )
-  }
-
-  private def isLeaderReplicaLocal: Boolean = leaderReplicaIfLocal.isDefined
+  newGauge("UnderReplicated",
+    new Gauge[Int] {
+      def value: Int = {
+        if (isUnderReplicated) 1 else 0
+      }
+    },
+    tags
+  )
+
+  newGauge("InSyncReplicasCount",
+    new Gauge[Int] {
+      def value: Int = {
+        if (isLeader) inSyncReplicaIds.size else 0
+      }
+    },
+    tags
+  )
+
+  newGauge("UnderMinIsr",
+    new Gauge[Int] {
+      def value: Int = {
+        if (isUnderMinIsr) 1 else 0
+      }
+    },
+    tags
+  )
+
+  newGauge("AtMinIsr",
+    new Gauge[Int] {
+      def value: Int = {
+        if (isAtMinIsr) 1 else 0
+      }
+    },
+    tags
+  )
+
+  newGauge("ReplicasCount",
+    new Gauge[Int] {
+      def value: Int = {
+        if (isLeader) allReplicaIds.size else 0
+      }
+    },
+    tags
+  )
+
+  newGauge("LastStableOffsetLag",
+    new Gauge[Long] {
+      def value: Long = {
+        log.map(_.lastStableOffsetLag).getOrElse(0)
+      }
+    },
+    tags
+  )
 
   def isUnderReplicated: Boolean =
-    isLeaderReplicaLocal && inSyncReplicas.size < assignedReplicas.size
+    isLeader && inSyncReplicaIds.size < allReplicaIds.size
 
   def isUnderMinIsr: Boolean = {
-    leaderReplicaIfLocal match {
-      case Some(leaderReplica) =>
-        inSyncReplicas.size < leaderReplica.log.get.config.minInSyncReplicas
+    leaderLogIfLocal.exists { inSyncReplicaIds.size < _.config.minInSyncReplicas }
+  }
+
+  def isAtMinIsr: Boolean = {
+    leaderLogIfLocal.exists { inSyncReplicaIds.size == _.config.minInSyncReplicas }
+  }
+
+  def isOneAboveMinIsr: Boolean = {
+    leaderLogIfLocal match {
+      case Some(leaderLog) =>
+        inSyncReplicaIds.size == leaderLog.config.minInSyncReplicas + 1
       case None =>
         false
     }
@@ -167,50 +281,80 @@ class Partition(val topicPartition: TopicPartition,
     * does not exist. This method assumes that the current replica has already been created.
     *
     * @param logDir log directory
+    * @param highWatermarkCheckpoints Checkpoint to load initial high watermark from
     * @return true iff the future replica is created
     */
-  def maybeCreateFutureReplica(logDir: String): Boolean = {
+  def maybeCreateFutureReplica(logDir: String, highWatermarkCheckpoints: OffsetCheckpoints): Boolean = {
     // The writeLock is needed to make sure that while the caller checks the log directory of the
     // current replica and the existence of the future replica, no other thread can update the log directory of the
     // current replica or remove the future replica.
     inWriteLock(leaderIsrUpdateLock) {
-      val currentReplica = localReplicaOrException
-      if (currentReplica.log.get.dir.getParent == logDir)
+      val currentLogDir = localLogOrException.parentDir
+      if (currentLogDir == logDir) {
+        info(s"Current log directory $currentLogDir is same as requested log dir $logDir. " +
+          s"Skipping future replica creation.")
         false
-      else {
-        futureLocalReplica match {
-          case Some(replica) =>
-            val futureReplicaLogDir = replica.log.get.dir.getParent
-            if (futureReplicaLogDir != logDir)
-              throw new IllegalStateException(s"The future log dir $futureReplicaLogDir of $topicPartition is " +
+      } else {
+        futureLog match {
+          case Some(partitionFutureLog) =>
+            val futureLogDir = partitionFutureLog.parentDir
+            if (futureLogDir != logDir)
+              throw new IllegalStateException(s"The future log dir $futureLogDir of $topicPartition is " +
                 s"different from the requested log dir $logDir")
             false
           case None =>
-            getOrCreateReplica(Request.FutureLocalReplicaId, isNew = false)
+            createLogIfNotExists(isNew = false, isFutureReplica = true, highWatermarkCheckpoints)
             true
         }
       }
     }
   }
 
-  def getOrCreateReplica(replicaId: Int, isNew: Boolean = false): Replica = {
-    allReplicasMap.getAndMaybePut(replicaId, {
-      if (isReplicaLocal(replicaId)) {
-        val adminZkClient = new AdminZkClient(zkClient)
-        val props = adminZkClient.fetchEntityConfig(ConfigType.Topic, topic)
-        val config = LogConfig.fromProps(logManager.currentDefaultConfig.originals, props)
-        val log = logManager.getOrCreateLog(topicPartition, config, isNew, replicaId == Request.FutureLocalReplicaId)
-        val checkpoint = replicaManager.highWatermarkCheckpoints(log.dir.getParent)
-        val offsetMap = checkpoint.read()
-        if (!offsetMap.contains(topicPartition))
-          info(s"No checkpointed highwatermark is found for partition $topicPartition")
-        val offset = math.min(offsetMap.getOrElse(topicPartition, 0L), log.logEndOffset)
-        new Replica(replicaId, topicPartition, time, offset, Some(log))
-      } else new Replica(replicaId, topicPartition, time)
-    })
+  def createLogIfNotExists(isNew: Boolean, isFutureReplica: Boolean, offsetCheckpoints: OffsetCheckpoints): Unit = {
+    isFutureReplica match {
+      case true if futureLog.isEmpty =>
+        val log = createLog(isNew, isFutureReplica, offsetCheckpoints)
+        this.futureLog = Option(log)
+      case false if log.isEmpty =>
+        val log = createLog(isNew, isFutureReplica, offsetCheckpoints)
+        this.log = Option(log)
+      case _ => trace(s"${if (isFutureReplica) "Future Log" else "Log"} already exists.")
+    }
   }
 
-  def getReplica(replicaId: Int): Option[Replica] = Option(allReplicasMap.get(replicaId))
+  // Visible for testing
+  private[cluster] def createLog(isNew: Boolean, isFutureReplica: Boolean, offsetCheckpoints: OffsetCheckpoints): Log = {
+    def fetchLogConfig: LogConfig = {
+      val props = stateStore.fetchTopicConfig()
+      LogConfig.fromProps(logManager.currentDefaultConfig.originals, props)
+    }
+
+    def updateHighWatermark(log: Log) = {
+      val checkpointHighWatermark = offsetCheckpoints.fetch(log.parentDir, topicPartition).getOrElse {
+        info(s"No checkpointed highwatermark is found for partition $topicPartition")
+        0L
+      }
+      val initialHighWatermark = log.updateHighWatermark(checkpointHighWatermark)
+      info(s"Log loaded for partition $topicPartition with initial high watermark $initialHighWatermark")
+    }
+
+    logManager.initializingLog(topicPartition)
+    var maybeLog: Option[Log] = None
+    try {
+      val log = logManager.getOrCreateLog(topicPartition, () => fetchLogConfig, isNew, isFutureReplica)
+      maybeLog = Some(log)
+      updateHighWatermark(log)
+      log
+    } finally {
+      logManager.finishedInitializingLog(topicPartition, maybeLog, () => fetchLogConfig)
+    }
+  }
+
+  def getReplica(replicaId: Int): Option[Replica] = Option(remoteReplicasMap.get(replicaId))
+
+  private def getReplicaOrException(replicaId: Int): Replica = getReplica(replicaId).getOrElse{
+    throw new ReplicaNotAvailableException(s"Replica with id $replicaId is not available on broker $localBrokerId")
+  }
 
   private def checkCurrentLeaderEpoch(remoteLeaderEpochOpt: Optional[Integer]): Errors = {
     if (!remoteLeaderEpochOpt.isPresent) {
@@ -227,22 +371,21 @@ class Partition(val topicPartition: TopicPartition,
     }
   }
 
-  private def getLocalReplica(replicaId: Int,
-                              currentLeaderEpoch: Optional[Integer],
-                              requireLeader: Boolean): Either[Replica, Errors] = {
+  private def getLocalLog(currentLeaderEpoch: Optional[Integer],
+                          requireLeader: Boolean): Either[Log, Errors] = {
     checkCurrentLeaderEpoch(currentLeaderEpoch) match {
       case Errors.NONE =>
-        if (requireLeader && !leaderReplicaIdOpt.contains(localBrokerId)) {
+        if (requireLeader && !isLeader) {
           Right(Errors.NOT_LEADER_FOR_PARTITION)
         } else {
-          val replica = allReplicasMap.get(replicaId)
-          if (replica == null) {
-            if (requireLeader)
-              Right(Errors.NOT_LEADER_FOR_PARTITION)
-            else
-              Right(Errors.REPLICA_NOT_AVAILABLE)
-          } else {
-            Left(replica)
+          log match {
+            case Some(partitionLog) =>
+              Left(partitionLog)
+            case _ =>
+              if (requireLeader)
+                Right(Errors.NOT_LEADER_FOR_PARTITION)
+              else
+                Right(Errors.REPLICA_NOT_AVAILABLE)
           }
         }
       case error =>
@@ -250,67 +393,57 @@ class Partition(val topicPartition: TopicPartition,
     }
   }
 
-  def localReplica: Option[Replica] = getReplica(localBrokerId)
-
-  def localReplicaOrException: Replica = localReplica.getOrElse {
-    throw new ReplicaNotAvailableException(s"Replica for partition $topicPartition is not available " +
+  def localLogOrException: Log = log.getOrElse {
+    throw new ReplicaNotAvailableException(s"Log for partition $topicPartition is not available " +
       s"on broker $localBrokerId")
   }
 
-  def futureLocalReplica: Option[Replica] = getReplica(Request.FutureLocalReplicaId)
-
-  def futureLocalReplicaOrException: Replica = futureLocalReplica.getOrElse {
-    throw new ReplicaNotAvailableException(s"Future replica for partition $topicPartition is not available " +
+  def futureLocalLogOrException: Log = futureLog.getOrElse {
+    throw new ReplicaNotAvailableException(s"Future log for partition $topicPartition is not available " +
       s"on broker $localBrokerId")
   }
 
-  def leaderReplicaIfLocal: Option[Replica] = {
-    if (leaderReplicaIdOpt.contains(localBrokerId))
-      localReplica
-    else
-      None
+  def leaderLogIfLocal: Option[Log] = {
+    log.filter(_ => isLeader)
   }
 
-  private def localReplicaWithEpochOrException(currentLeaderEpoch: Optional[Integer],
-                                               requireLeader: Boolean): Replica = {
-    getLocalReplica(localBrokerId, currentLeaderEpoch, requireLeader) match {
-      case Left(replica) => replica
+  /**
+   * Returns true if this node is currently leader for the Partition.
+   */
+  def isLeader: Boolean = leaderReplicaIdOpt.contains(localBrokerId)
+
+  private def localLogWithEpochOrException(currentLeaderEpoch: Optional[Integer],
+                                           requireLeader: Boolean): Log = {
+    getLocalLog(currentLeaderEpoch, requireLeader) match {
+      case Left(localLog) => localLog
       case Right(error) =>
-        throw error.exception(s"Failed to find ${if (requireLeader) "leader " else ""} for " +
+        throw error.exception(s"Failed to find ${if (requireLeader) "leader " else ""} log for " +
           s"partition $topicPartition with leader epoch $currentLeaderEpoch. The current leader " +
           s"is $leaderReplicaIdOpt and the current epoch $leaderEpoch")
     }
   }
 
-  def addReplicaIfNotExists(replica: Replica): Replica =
-    allReplicasMap.putIfNotExists(replica.brokerId, replica)
-
-  def assignedReplicas: Set[Replica] =
-    allReplicasMap.values.filter(replica => Request.isValidBrokerId(replica.brokerId)).toSet
-
-  def allReplicas: Set[Replica] =
-    allReplicasMap.values.toSet
-
-  private def removeReplica(replicaId: Int) {
-    allReplicasMap.remove(replicaId)
+  // Visible for testing -- Used by unit tests to set log for this partition
+  def setLog(log: Log, isFutureLog: Boolean): Unit = {
+    if (isFutureLog)
+      futureLog = Some(log)
+    else
+      this.log = Some(log)
   }
 
+  // remoteReplicas will be called in the hot path, and must be inexpensive
+  def remoteReplicas: Iterable[Replica] =
+    remoteReplicasMap.values
+
   def futureReplicaDirChanged(newDestinationDir: String): Boolean = {
     inReadLock(leaderIsrUpdateLock) {
-      futureLocalReplica match {
-        case Some(replica) =>
-          if (replica.log.get.dir.getParent != newDestinationDir)
-            true
-          else
-            false
-        case None => false
-      }
+      futureLog.exists(_.parentDir != newDestinationDir)
     }
   }
 
-  def removeFutureLocalReplica(deleteFromLogDir: Boolean = true) {
+  def removeFutureLocalReplica(deleteFromLogDir: Boolean = true): Unit = {
     inWriteLock(leaderIsrUpdateLock) {
-      allReplicasMap.remove(Request.FutureLocalReplicaId)
+      futureLog = None
       if (deleteFromLogDir)
         logManager.asyncDelete(topicPartition, isFuture = true)
     }
@@ -320,39 +453,42 @@ class Partition(val topicPartition: TopicPartition,
   // Only ReplicaAlterDirThread will call this method and ReplicaAlterDirThread should remove the partition
   // from its partitionStates if this method returns true
   def maybeReplaceCurrentWithFutureReplica(): Boolean = {
-    val replica = localReplicaOrException
-    val futureReplicaLEO = futureLocalReplica.map(_.logEndOffset)
-    if (futureReplicaLEO.contains(replica.logEndOffset)) {
+    val localReplicaLEO = localLogOrException.logEndOffset
+    val futureReplicaLEO = futureLog.map(_.logEndOffset)
+    if (futureReplicaLEO.contains(localReplicaLEO)) {
       // The write lock is needed to make sure that while ReplicaAlterDirThread checks the LEO of the
       // current replica, no other thread can update LEO of the current replica via log truncation or log append operation.
       inWriteLock(leaderIsrUpdateLock) {
-        futureLocalReplica match {
-          case Some(futureReplica) =>
-            if (replica.logEndOffset == futureReplica.logEndOffset) {
+        futureLog match {
+          case Some(futurePartitionLog) =>
+            if (log.exists(_.logEndOffset == futurePartitionLog.logEndOffset)) {
               logManager.replaceCurrentWithFutureLog(topicPartition)
-              replica.log = futureReplica.log
-              futureReplica.log = None
-              allReplicasMap.remove(Request.FutureLocalReplicaId)
+              log = futureLog
+              removeFutureLocalReplica(false)
               true
             } else false
           case None =>
             // Future replica is removed by a non-ReplicaAlterLogDirsThread before this method is called
             // In this case the partition should have been removed from state of the ReplicaAlterLogDirsThread
-            // Return false so that ReplicaAlterLogDirsThread does not have to remove this partition from the state again to avoid race condition
+            // Return false so that ReplicaAlterLogDirsThread does not have to remove this partition from the
+            // state again to avoid race condition
             false
         }
       }
     } else false
   }
 
-  def delete() {
+  def delete(): Unit = {
     // need to hold the lock to prevent appendMessagesToLeader() from hitting I/O exceptions due to log being deleted
     inWriteLock(leaderIsrUpdateLock) {
-      allReplicasMap.clear()
-      inSyncReplicas = Set.empty[Replica]
+      remoteReplicasMap.clear()
+      allReplicaIds = Seq.empty
+      log = None
+      futureLog = None
+      inSyncReplicaIds = Set.empty
       leaderReplicaIdOpt = None
       leaderEpochStartOffsetOpt = None
-      removePartitionMetrics()
+      Partition.removeMetrics(topicPartition)
       logManager.asyncDelete(topicPartition)
       if (logManager.getLog(topicPartition, isFuture = true).isDefined)
         logManager.asyncDelete(topicPartition, isFuture = true)
@@ -366,57 +502,60 @@ class Partition(val topicPartition: TopicPartition,
    * from the time when this broker was the leader last time) and setting the new leader and ISR.
    * If the leader replica id does not change, return false to indicate the replica manager.
    */
-  def makeLeader(controllerId: Int, partitionStateInfo: LeaderAndIsrRequest.PartitionState, correlationId: Int): Boolean = {
+  def makeLeader(partitionState: LeaderAndIsrPartitionState,
+                 highWatermarkCheckpoints: OffsetCheckpoints): Boolean = {
     val (leaderHWIncremented, isNewLeader) = inWriteLock(leaderIsrUpdateLock) {
-      val newAssignedReplicas = partitionStateInfo.basePartitionState.replicas.asScala.map(_.toInt)
       // record the epoch of the controller that made the leadership decision. This is useful while updating the isr
       // to maintain the decision maker controller's epoch in the zookeeper path
-      controllerEpoch = partitionStateInfo.basePartitionState.controllerEpoch
-      // add replicas that are new
-      val newInSyncReplicas = partitionStateInfo.basePartitionState.isr.asScala.map(r => getOrCreateReplica(r, partitionStateInfo.isNew)).toSet
-      // remove assigned replicas that have been removed by the controller
-      (assignedReplicas.map(_.brokerId) -- newAssignedReplicas).foreach(removeReplica)
-      inSyncReplicas = newInSyncReplicas
-      newAssignedReplicas.foreach(id => getOrCreateReplica(id, partitionStateInfo.isNew))
-
-      val leaderReplica = localReplicaOrException
-      val leaderEpochStartOffset = leaderReplica.logEndOffset
-      info(s"$topicPartition starts at Leader Epoch ${partitionStateInfo.basePartitionState.leaderEpoch} from " +
-        s"offset $leaderEpochStartOffset. Previous Leader Epoch was: $leaderEpoch")
+      controllerEpoch = partitionState.controllerEpoch
+
+      updateAssignmentAndIsr(
+        assignment = partitionState.replicas.asScala.iterator.map(_.toInt).toSeq,
+        isr = partitionState.isr.asScala.iterator.map(_.toInt).toSet
+      )
+      createLogIfNotExists(partitionState.isNew, isFutureReplica = false, highWatermarkCheckpoints)
+
+      val leaderLog = localLogOrException
+      val leaderEpochStartOffset = leaderLog.logEndOffset
+      info(s"$topicPartition starts at leader epoch ${partitionState.leaderEpoch} from " +
+        s"offset $leaderEpochStartOffset with high watermark ${leaderLog.highWatermark}. " +
+        s"Previous leader epoch was $leaderEpoch.")
 
       //We cache the leader epoch here, persisting it only if it's local (hence having a log dir)
-      leaderEpoch = partitionStateInfo.basePartitionState.leaderEpoch
+      leaderEpoch = partitionState.leaderEpoch
       leaderEpochStartOffsetOpt = Some(leaderEpochStartOffset)
-      zkVersion = partitionStateInfo.basePartitionState.zkVersion
+      zkVersion = partitionState.zkVersion
 
       // In the case of successive leader elections in a short time period, a follower may have
       // entries in its log from a later epoch than any entry in the new leader's log. In order
       // to ensure that these followers can truncate to the right offset, we must cache the new
       // leader epoch and the start offset since it should be larger than any epoch that a follower
       // would try to query.
-      leaderReplica.log.foreach { log =>
-        log.maybeAssignEpochStartOffset(leaderEpoch, leaderEpochStartOffset)
-      }
+      leaderLog.maybeAssignEpochStartOffset(leaderEpoch, leaderEpochStartOffset)
 
-      val isNewLeader = !leaderReplicaIdOpt.contains(localBrokerId)
-      val curLeaderLogEndOffset = leaderReplica.logEndOffset
+      val isNewLeader = !isLeader
       val curTimeMs = time.milliseconds
       // initialize lastCaughtUpTime of replicas as well as their lastFetchTimeMs and lastFetchLeaderLogEndOffset.
-      (assignedReplicas - leaderReplica).foreach { replica =>
-        val lastCaughtUpTimeMs = if (inSyncReplicas.contains(replica)) curTimeMs else 0L
-        replica.resetLastCaughtUpTime(curLeaderLogEndOffset, curTimeMs, lastCaughtUpTimeMs)
+      remoteReplicas.foreach { replica =>
+        val lastCaughtUpTimeMs = if (inSyncReplicaIds.contains(replica.brokerId)) curTimeMs else 0L
+        replica.resetLastCaughtUpTime(leaderEpochStartOffset, curTimeMs, lastCaughtUpTimeMs)
       }
 
       if (isNewLeader) {
-        // construct the high watermark metadata for the new leader replica
-        leaderReplica.convertHWToLocalOffsetMetadata()
         // mark local replica as the leader after converting hw
         leaderReplicaIdOpt = Some(localBrokerId)
         // reset log end offset for remote replicas
-        assignedReplicas.filter(_.brokerId != localBrokerId).foreach(_.updateLogReadResult(LogReadResult.UnknownLogReadResult))
+        remoteReplicas.foreach { replica =>
+          replica.updateFetchState(
+            followerFetchOffsetMetadata = LogOffsetMetadata.UnknownOffsetMetadata,
+            followerStartOffset = Log.UnknownOffset,
+            followerFetchTimeMs = 0L,
+            leaderEndOffset = Log.UnknownOffset,
+            lastSentHighwatermark = 0L)
+        }
       }
       // we may need to increment high watermark since ISR could be down to 1
-      (maybeIncrementLeaderHW(leaderReplica), isNewLeader)
+      (maybeIncrementLeaderHW(leaderLog), isNewLeader)
     }
     // some delayed operations may be unblocked after HW changed
     if (leaderHWIncremented)
@@ -430,22 +569,24 @@ class Partition(val topicPartition: TopicPartition,
    *  greater (that is, no updates have been missed), return false to indicate to the
     * replica manager that state is already correct and the become-follower steps can be skipped
    */
-  def makeFollower(controllerId: Int, partitionStateInfo: LeaderAndIsrRequest.PartitionState, correlationId: Int): Boolean = {
+  def makeFollower(partitionState: LeaderAndIsrPartitionState,
+                   highWatermarkCheckpoints: OffsetCheckpoints): Boolean = {
     inWriteLock(leaderIsrUpdateLock) {
-      val newAssignedReplicas = partitionStateInfo.basePartitionState.replicas.asScala.map(_.toInt)
-      val newLeaderBrokerId = partitionStateInfo.basePartitionState.leader
+      val newLeaderBrokerId = partitionState.leader
       val oldLeaderEpoch = leaderEpoch
       // record the epoch of the controller that made the leadership decision. This is useful while updating the isr
       // to maintain the decision maker controller's epoch in the zookeeper path
-      controllerEpoch = partitionStateInfo.basePartitionState.controllerEpoch
-      // add replicas that are new
-      newAssignedReplicas.foreach(r => getOrCreateReplica(r, partitionStateInfo.isNew))
-      // remove assigned replicas that have been removed by the controller
-      (assignedReplicas.map(_.brokerId) -- newAssignedReplicas).foreach(removeReplica)
-      inSyncReplicas = Set.empty[Replica]
-      leaderEpoch = partitionStateInfo.basePartitionState.leaderEpoch
+      controllerEpoch = partitionState.controllerEpoch
+
+      updateAssignmentAndIsr(
+        assignment = partitionState.replicas.asScala.iterator.map(_.toInt).toSeq,
+        isr = Set.empty[Int]
+      )
+      createLogIfNotExists(partitionState.isNew, isFutureReplica = false, highWatermarkCheckpoints)
+
+      leaderEpoch = partitionState.leaderEpoch
       leaderEpochStartOffsetOpt = None
-      zkVersion = partitionStateInfo.basePartitionState.zkVersion
+      zkVersion = partitionState.zkVersion
 
       if (leaderReplicaIdOpt.contains(newLeaderBrokerId) && leaderEpoch == oldLeaderEpoch) {
         false
@@ -460,28 +601,79 @@ class Partition(val topicPartition: TopicPartition,
    * Update the follower's state in the leader based on the last fetch request. See
    * [[kafka.cluster.Replica#updateLogReadResult]] for details.
    *
-   * @return true if the leader's log start offset or high watermark have been updated
+   * @return true if the follower's fetch state was updated, false if the followerId is not recognized
    */
-  def updateReplicaLogReadResult(replica: Replica, logReadResult: LogReadResult): Boolean = {
-    val replicaId = replica.brokerId
-    // No need to calculate low watermark if there is no delayed DeleteRecordsRequest
-    val oldLeaderLW = if (replicaManager.delayedDeleteRecordsPurgatory.delayed > 0) lowWatermarkIfLeader else -1L
-    replica.updateLogReadResult(logReadResult)
-    val newLeaderLW = if (replicaManager.delayedDeleteRecordsPurgatory.delayed > 0) lowWatermarkIfLeader else -1L
-    // check if the LW of the partition has incremented
-    // since the replica's logStartOffset may have incremented
-    val leaderLWIncremented = newLeaderLW > oldLeaderLW
-    // check if we need to expand ISR to include this replica
-    // if it is not in the ISR yet
-    val leaderHWIncremented = maybeExpandIsr(replicaId, logReadResult)
-
-    val result = leaderLWIncremented || leaderHWIncremented
-    // some delayed operations may be unblocked after HW or LW changed
-    if (result)
-      tryCompleteDelayedRequests()
+  def updateFollowerFetchState(followerId: Int,
+                               followerFetchOffsetMetadata: LogOffsetMetadata,
+                               followerStartOffset: Long,
+                               followerFetchTimeMs: Long,
+                               leaderEndOffset: Long,
+                               lastSentHighwatermark: Long): Boolean = {
+    getReplica(followerId) match {
+      case Some(followerReplica) =>
+        // No need to calculate low watermark if there is no delayed DeleteRecordsRequest
+        val oldLeaderLW = if (delayedOperations.numDelayedDelete > 0) lowWatermarkIfLeader else -1L
+        val prevFollowerEndOffset = followerReplica.logEndOffset
+        followerReplica.updateFetchState(
+          followerFetchOffsetMetadata,
+          followerStartOffset,
+          followerFetchTimeMs,
+          leaderEndOffset,
+          lastSentHighwatermark)
+
+        val newLeaderLW = if (delayedOperations.numDelayedDelete > 0) lowWatermarkIfLeader else -1L
+        // check if the LW of the partition has incremented
+        // since the replica's logStartOffset may have incremented
+        val leaderLWIncremented = newLeaderLW > oldLeaderLW
+
+        // check if we need to expand ISR to include this replica
+        // if it is not in the ISR yet
+        if (!inSyncReplicaIds(followerId))
+          maybeExpandIsr(followerReplica, followerFetchTimeMs)
+
+        // check if the HW of the partition can now be incremented
+        // since the replica may already be in the ISR and its LEO has just incremented
+        val leaderHWIncremented = if (prevFollowerEndOffset != followerReplica.logEndOffset) {
+          leaderLogIfLocal.exists(leaderLog => maybeIncrementLeaderHW(leaderLog, followerFetchTimeMs))
+        } else {
+          false
+        }
+
+        // some delayed operations may be unblocked after HW or LW changed
+        if (leaderLWIncremented || leaderHWIncremented)
+          tryCompleteDelayedRequests()
 
-    debug(s"Recorded replica $replicaId log end offset (LEO) position ${logReadResult.info.fetchOffsetMetadata.messageOffset}.")
-    result
+        debug(s"Recorded replica $followerId log end offset (LEO) position " +
+          s"${followerFetchOffsetMetadata.messageOffset} and log start offset $followerStartOffset.")
+        true
+
+      case None =>
+        false
+    }
+  }
+
+  /**
+   * Stores the topic partition assignment and ISR.
+   * It creates a new Replica object for any new remote broker. The isr parameter is
+   * expected to be a subset of the assignment parameter.
+   *
+   * Note: public visibility for tests.
+   *
+   * @param assignment An ordered sequence of all the broker ids that were assigned to this
+   *                   topic partition
+   * @param isr The set of broker ids that are known to be insync with the leader
+   */
+  def updateAssignmentAndIsr(assignment: Seq[Int], isr: Set[Int]): Unit = {
+    val replicaSet = assignment.toSet
+    val removedReplicas = remoteReplicasMap.keys -- replicaSet
+
+    assignment
+      .filter(_ != localBrokerId)
+      .foreach(id => remoteReplicasMap.getAndMaybePut(id, new Replica(id, topicPartition)))
+    removedReplicas.foreach(remoteReplicasMap.remove)
+    allReplicaIds = assignment
+
+    inSyncReplicaIds = isr
   }
 
   /**
@@ -497,36 +689,29 @@ class Partition(val topicPartition: TopicPartition,
    * whether a replica is in-sync, we only check HW.
    *
    * This function can be triggered when a replica's LEO has incremented.
-   *
-   * @return true if the high watermark has been updated
    */
-  def maybeExpandIsr(replicaId: Int, logReadResult: LogReadResult): Boolean = {
+  private def maybeExpandIsr(followerReplica: Replica, followerFetchTimeMs: Long): Unit = {
     inWriteLock(leaderIsrUpdateLock) {
       // check if this replica needs to be added to the ISR
-      leaderReplicaIfLocal match {
-        case Some(leaderReplica) =>
-          val replica = getReplica(replicaId).get
-          val leaderHW = leaderReplica.highWatermark
-          val fetchOffset = logReadResult.info.fetchOffsetMetadata.messageOffset
-          if (!inSyncReplicas.contains(replica) &&
-             assignedReplicas.map(_.brokerId).contains(replicaId) &&
-             replica.logEndOffsetMetadata.offsetDiff(leaderHW) >= 0 &&
-             leaderEpochStartOffsetOpt.exists(fetchOffset >= _)) {
-            val newInSyncReplicas = inSyncReplicas + replica
-            info(s"Expanding ISR from ${inSyncReplicas.map(_.brokerId).mkString(",")} " +
-              s"to ${newInSyncReplicas.map(_.brokerId).mkString(",")}")
-            // update ISR in ZK and cache
-            updateIsr(newInSyncReplicas)
-            replicaManager.isrExpandRate.mark()
-          }
-          // check if the HW of the partition can now be incremented
-          // since the replica may already be in the ISR and its LEO has just incremented
-          maybeIncrementLeaderHW(leaderReplica, logReadResult.fetchTimeMs)
-        case None => false // nothing to do if no longer leader
+      leaderLogIfLocal.foreach { leaderLog =>
+        val leaderHighwatermark = leaderLog.highWatermark
+        if (!inSyncReplicaIds.contains(followerReplica.brokerId) && isFollowerInSync(followerReplica, leaderHighwatermark)) {
+          val newInSyncReplicaIds = inSyncReplicaIds + followerReplica.brokerId
+          info(s"Expanding ISR from ${inSyncReplicaIds.mkString(",")} " +
+            s"to ${newInSyncReplicaIds.mkString(",")}")
+
+          // update ISR in ZK and cache
+          expandIsr(newInSyncReplicaIds)
+        }
       }
     }
   }
 
+  private def isFollowerInSync(followerReplica: Replica, highWatermark: Long): Boolean = {
+    val followerEndOffset = followerReplica.logEndOffset
+    followerEndOffset >= highWatermark && leaderEpochStartOffsetOpt.exists(followerEndOffset >= _)
+  }
+
   /*
    * Returns a tuple where the first element is a boolean indicating whether enough replicas reached `requiredOffset`
    * and the second element is an error (which would be `Errors.NONE` for no error).
@@ -536,27 +721,33 @@ class Partition(val topicPartition: TopicPartition,
    * produce request.
    */
   def checkEnoughReplicasReachOffset(requiredOffset: Long): (Boolean, Errors) = {
-    leaderReplicaIfLocal match {
-      case Some(leaderReplica) =>
+    leaderLogIfLocal match {
+      case Some(leaderLog) =>
         // keep the current immutable replica list reference
-        val curInSyncReplicas = inSyncReplicas
+        val curInSyncReplicaIds = inSyncReplicaIds
 
         if (isTraceEnabled) {
-          def logEndOffsetString(r: Replica) = s"broker ${r.brokerId}: ${r.logEndOffset}"
-          val (ackedReplicas, awaitingReplicas) = curInSyncReplicas.partition { replica =>
-            replica.logEndOffset >= requiredOffset
+          def logEndOffsetString: ((Int, Long)) => String = {
+            case (brokerId, logEndOffset) => s"broker $brokerId: $logEndOffset"
           }
-          trace(s"Progress awaiting ISR acks for offset $requiredOffset: acked: ${ackedReplicas.map(logEndOffsetString)}, " +
+
+          val curInSyncReplicaObjects = (curInSyncReplicaIds - localBrokerId).map(getReplicaOrException)
+          val replicaInfo = curInSyncReplicaObjects.map(replica => (replica.brokerId, replica.logEndOffset))
+          val localLogInfo = (localBrokerId, localLogOrException.logEndOffset)
+          val (ackedReplicas, awaitingReplicas) = (replicaInfo + localLogInfo).partition { _._2 >= requiredOffset}
+
+          trace(s"Progress awaiting ISR acks for offset $requiredOffset: " +
+            s"acked: ${ackedReplicas.map(logEndOffsetString)}, " +
             s"awaiting ${awaitingReplicas.map(logEndOffsetString)}")
         }
 
-        val minIsr = leaderReplica.log.get.config.minInSyncReplicas
-        if (leaderReplica.highWatermark.messageOffset >= requiredOffset) {
+        val minIsr = leaderLog.config.minInSyncReplicas
+        if (leaderLog.highWatermark >= requiredOffset) {
           /*
            * The topic may be configured not to accept messages if there are not enough replicas in ISR
            * in this scenario the request was already appended locally and then added to the purgatory before the ISR was shrunk
            */
-          if (minIsr <= curInSyncReplicas.size)
+          if (minIsr <= curInSyncReplicaIds.size)
             (true, Errors.NONE)
           else
             (true, Errors.NOT_ENOUGH_REPLICAS_AFTER_APPEND)
@@ -585,25 +776,36 @@ class Partition(val topicPartition: TopicPartition,
    * Note There is no need to acquire the leaderIsrUpdate lock here
    * since all callers of this private API acquire that lock
    */
-  private def maybeIncrementLeaderHW(leaderReplica: Replica, curTime: Long = time.milliseconds): Boolean = {
-    val allLogEndOffsets = assignedReplicas.filter { replica =>
-      curTime - replica.lastCaughtUpTimeMs <= replicaLagTimeMaxMs || inSyncReplicas.contains(replica)
-    }.map(_.logEndOffsetMetadata)
-    val newHighWatermark = allLogEndOffsets.min(new LogOffsetMetadata.OffsetOrdering)
-    val oldHighWatermark = leaderReplica.highWatermark
-
-    // Ensure that the high watermark increases monotonically. We also update the high watermark when the new
-    // offset metadata is on a newer segment, which occurs whenever the log is rolled to a new segment.
-    if (oldHighWatermark.messageOffset < newHighWatermark.messageOffset ||
-      (oldHighWatermark.messageOffset == newHighWatermark.messageOffset && oldHighWatermark.onOlderSegment(newHighWatermark))) {
-      leaderReplica.highWatermark = newHighWatermark
-      debug(s"High watermark updated to $newHighWatermark")
-      true
-    } else {
-      def logEndOffsetString(r: Replica) = s"replica ${r.brokerId}: ${r.logEndOffsetMetadata}"
-      debug(s"Skipping update high watermark since new hw $newHighWatermark is not larger than old hw $oldHighWatermark. " +
-        s"All current LEOs are ${assignedReplicas.map(logEndOffsetString)}")
-      false
+  private def maybeIncrementLeaderHW(leaderLog: Log, curTime: Long = time.milliseconds): Boolean = {
+    inReadLock(leaderIsrUpdateLock) {
+      // maybeIncrementLeaderHW is in the hot path, the following code is written to
+      // avoid unnecessary collection generation
+      var newHighWatermark = leaderLog.logEndOffsetMetadata
+      remoteReplicasMap.values.foreach { replica =>
+        if (replica.logEndOffsetMetadata.messageOffset < newHighWatermark.messageOffset &&
+          (curTime - replica.lastCaughtUpTimeMs <= replicaLagTimeMaxMs || inSyncReplicaIds.contains(replica.brokerId))) {
+          newHighWatermark = replica.logEndOffsetMetadata
+        }
+      }
+
+      leaderLog.maybeIncrementHighWatermark(newHighWatermark) match {
+        case Some(oldHighWatermark) =>
+          debug(s"High watermark updated from $oldHighWatermark to $newHighWatermark")
+          true
+
+        case None =>
+          def logEndOffsetString: ((Int, LogOffsetMetadata)) => String = {
+            case (brokerId, logEndOffsetMetadata) => s"replica $brokerId: $logEndOffsetMetadata"
+          }
+
+          if (isTraceEnabled) {
+            val replicaInfo = remoteReplicas.map(replica => (replica.brokerId, replica.logEndOffsetMetadata)).toSet
+            val localLogInfo = (localBrokerId, localLogOrException.logEndOffsetMetadata)
+            trace(s"Skipping update high watermark since new hw $newHighWatermark is not larger than old value. " +
+              s"All current LEOs are ${(replicaInfo + localLogInfo).map(logEndOffsetString)}")
+          }
+          false
+      }
     }
   }
 
@@ -613,49 +815,55 @@ class Partition(val topicPartition: TopicPartition,
    * Low watermark will increase when the leader broker receives either FetchRequest or DeleteRecordsRequest.
    */
   def lowWatermarkIfLeader: Long = {
-    if (!isLeaderReplicaLocal)
+    if (!isLeader)
       throw new NotLeaderForPartitionException(s"Leader not local for partition $topicPartition on broker $localBrokerId")
-    val logStartOffsets = allReplicas.collect {
-      case replica if replicaManager.metadataCache.isBrokerAlive(replica.brokerId) || replica.brokerId == Request.FutureLocalReplicaId => replica.logStartOffset
+
+    // lowWatermarkIfLeader may be called many times when a DeleteRecordsRequest is outstanding,
+    // care has been taken to avoid generating unnecessary collections in this code
+    var lowWaterMark = localLogOrException.logStartOffset
+    remoteReplicas.foreach { replica =>
+      if (metadataCache.getAliveBroker(replica.brokerId).nonEmpty && replica.logStartOffset < lowWaterMark) {
+        lowWaterMark = replica.logStartOffset
+      }
+    }
+
+    futureLog match {
+      case Some(partitionFutureLog) =>
+        Math.min(lowWaterMark, partitionFutureLog.logStartOffset)
+      case None =>
+        lowWaterMark
     }
-    CoreUtils.min(logStartOffsets, 0L)
   }
 
   /**
    * Try to complete any pending requests. This should be called without holding the leaderIsrUpdateLock.
    */
-  private def tryCompleteDelayedRequests() {
-    val requestKey = new TopicPartitionOperationKey(topicPartition)
-    replicaManager.tryCompleteDelayedFetch(requestKey)
-    replicaManager.tryCompleteDelayedProduce(requestKey)
-    replicaManager.tryCompleteDelayedDeleteRecords(requestKey)
-  }
+  private def tryCompleteDelayedRequests(): Unit = delayedOperations.checkAndCompleteAll()
 
-  def maybeShrinkIsr(replicaMaxLagTimeMs: Long) {
+  def maybeShrinkIsr(replicaMaxLagTimeMs: Long): Unit = {
     val leaderHWIncremented = inWriteLock(leaderIsrUpdateLock) {
-      leaderReplicaIfLocal match {
-        case Some(leaderReplica) =>
-          val outOfSyncReplicas = getOutOfSyncReplicas(leaderReplica, replicaMaxLagTimeMs)
-          if (outOfSyncReplicas.nonEmpty) {
-            val newInSyncReplicas = inSyncReplicas -- outOfSyncReplicas
-            assert(newInSyncReplicas.nonEmpty)
+      leaderLogIfLocal match {
+        case Some(leaderLog) =>
+          val outOfSyncReplicaIds = getOutOfSyncReplicas(replicaMaxLagTimeMs)
+          if (outOfSyncReplicaIds.nonEmpty) {
+            val newInSyncReplicaIds = inSyncReplicaIds -- outOfSyncReplicaIds
+            assert(newInSyncReplicaIds.nonEmpty)
             info("Shrinking ISR from %s to %s. Leader: (highWatermark: %d, endOffset: %d). Out of sync replicas: %s."
-              .format(inSyncReplicas.map(_.brokerId).mkString(","),
-                newInSyncReplicas.map(_.brokerId).mkString(","),
-                leaderReplica.highWatermark.messageOffset,
-                leaderReplica.logEndOffset,
-                outOfSyncReplicas.map { replica =>
-                  s"(brokerId: ${replica.brokerId}, endOffset: ${replica.logEndOffset})"
+              .format(inSyncReplicaIds.mkString(","),
+                newInSyncReplicaIds.mkString(","),
+                leaderLog.highWatermark,
+                leaderLog.logEndOffset,
+                outOfSyncReplicaIds.map { replicaId =>
+                  s"(brokerId: $replicaId, endOffset: ${getReplicaOrException(replicaId).logEndOffset})"
                 }.mkString(" ")
               )
             )
 
             // update ISR in zk and in cache
-            updateIsr(newInSyncReplicas)
-            replicaManager.isrShrinkRate.mark()
+            shrinkIsr(newInSyncReplicaIds)
 
             // we may need to increment high watermark since ISR could be down to 1
-            maybeIncrementLeaderHW(leaderReplica)
+            maybeIncrementLeaderHW(leaderLog)
           } else {
             false
           }
@@ -669,7 +877,16 @@ class Partition(val topicPartition: TopicPartition,
       tryCompleteDelayedRequests()
   }
 
-  def getOutOfSyncReplicas(leaderReplica: Replica, maxLagMs: Long): Set[Replica] = {
+  private def isFollowerOutOfSync(replicaId: Int,
+                                  leaderEndOffset: Long,
+                                  currentTimeMs: Long,
+                                  maxLagMs: Long): Boolean = {
+    val followerReplica = getReplicaOrException(replicaId)
+    followerReplica.logEndOffset != leaderEndOffset &&
+      (currentTimeMs - followerReplica.lastCaughtUpTimeMs) > maxLagMs
+  }
+
+  def getOutOfSyncReplicas(maxLagMs: Long): Set[Int] = {
     /**
      * If the follower already has the same leo as the leader, it will not be considered as out-of-sync,
      * otherwise there are two cases that will be handled here -
@@ -682,14 +899,10 @@ class Partition(val topicPartition: TopicPartition,
      * is violated, that replica is considered to be out of sync
      *
      **/
-    val candidateReplicas = inSyncReplicas - leaderReplica
-
-    val laggingReplicas = candidateReplicas.filter(r =>
-      r.logEndOffset != leaderReplica.logEndOffset && (time.milliseconds - r.lastCaughtUpTimeMs) > maxLagMs)
-    if (laggingReplicas.nonEmpty)
-      debug("Lagging replicas are %s".format(laggingReplicas.map(_.brokerId).mkString(",")))
-
-    laggingReplicas
+    val candidateReplicaIds = inSyncReplicaIds - localBrokerId
+    val currentTimeMs = time.milliseconds()
+    val leaderEndOffset = localLogOrException.logEndOffset
+    candidateReplicaIds.filter(replicaId => isFollowerOutOfSync(replicaId, leaderEndOffset, currentTimeMs, maxLagMs))
   }
 
   private def doAppendRecordsToFollowerOrFutureReplica(records: MemoryRecords, isFuture: Boolean): Option[LogAppendInfo] = {
@@ -699,13 +912,11 @@ class Partition(val topicPartition: TopicPartition,
       if (isFuture) {
         // Note the replica may be undefined if it is removed by a non-ReplicaAlterLogDirsThread before
         // this method is called
-        futureLocalReplica.map { replica =>
-          replica.log.get.appendAsFollower(records)
-        }
+        futureLog.map { _.appendAsFollower(records) }
       } else {
         // The read lock is needed to prevent the follower replica from being updated while ReplicaAlterDirThread
         // is executing maybeDeleteAndSwapFutureReplica() to replace follower replica with the future replica.
-        Some(localReplicaOrException.log.get.appendAsFollower(records))
+        Some(localLogOrException.appendAsFollower(records))
       }
     }
   }
@@ -715,9 +926,9 @@ class Partition(val topicPartition: TopicPartition,
       doAppendRecordsToFollowerOrFutureReplica(records, isFuture)
     } catch {
       case e: UnexpectedAppendOffsetException =>
-        val replica = if (isFuture) futureLocalReplicaOrException else localReplicaOrException
-        val logEndOffset = replica.logEndOffset
-        if (logEndOffset == replica.logStartOffset &&
+        val log = if (isFuture) futureLocalLogOrException else localLogOrException
+        val logEndOffset = log.logEndOffset
+        if (logEndOffset == log.logStartOffset &&
             e.firstOffset < logEndOffset && e.lastOffset >= logEndOffset) {
           // This may happen if the log start offset on the leader (or current replica) falls in
           // the middle of the batch due to delete records request and the follower tries to
@@ -727,7 +938,7 @@ class Partition(val topicPartition: TopicPartition,
           // (base offset of the batch), which will move recoveryPoint backwards, so we will need
           // to checkpoint the new recovery point before we append
           val replicaName = if (isFuture) "future replica" else "follower"
-          info(s"Unexpected offset in append to $topicPartition. First offset ${e.firstOffset} is less than log start offset ${replica.logStartOffset}." +
+          info(s"Unexpected offset in append to $topicPartition. First offset ${e.firstOffset} is less than log start offset ${log.logStartOffset}." +
                s" Since this is the first record to be appended to the $replicaName's log, will start the log from offset ${e.firstOffset}.")
           truncateFullyAndStartAt(e.firstOffset, isFuture)
           doAppendRecordsToFollowerOrFutureReplica(records, isFuture)
@@ -736,25 +947,24 @@ class Partition(val topicPartition: TopicPartition,
     }
   }
 
-  def appendRecordsToLeader(records: MemoryRecords, isFromClient: Boolean, requiredAcks: Int = 0): LogAppendInfo = {
+  def appendRecordsToLeader(records: MemoryRecords, origin: AppendOrigin, requiredAcks: Int): LogAppendInfo = {
     val (info, leaderHWIncremented) = inReadLock(leaderIsrUpdateLock) {
-      leaderReplicaIfLocal match {
-        case Some(leaderReplica) =>
-          val log = leaderReplica.log.get
-          val minIsr = log.config.minInSyncReplicas
-          val inSyncSize = inSyncReplicas.size
+      leaderLogIfLocal match {
+        case Some(leaderLog) =>
+          val minIsr = leaderLog.config.minInSyncReplicas
+          val inSyncSize = inSyncReplicaIds.size
 
           // Avoid writing to leader if there are not enough insync replicas to make it safe
           if (inSyncSize < minIsr && requiredAcks == -1) {
-            throw new NotEnoughReplicasException(s"The size of the current ISR ${inSyncReplicas.map(_.brokerId)} " +
+            throw new NotEnoughReplicasException(s"The size of the current ISR $inSyncReplicaIds " +
               s"is insufficient to satisfy the min.isr requirement of $minIsr for partition $topicPartition")
           }
 
-          val info = log.appendAsLeader(records, leaderEpoch = this.leaderEpoch, isFromClient,
+          val info = leaderLog.appendAsLeader(records, leaderEpoch = this.leaderEpoch, origin,
             interBrokerProtocolVersion)
 
           // we may need to increment high watermark since ISR could be down to 1
-          (info, maybeIncrementLeaderHW(leaderReplica))
+          (info, maybeIncrementLeaderHW(leaderLog))
 
         case None =>
           throw new NotLeaderForPartitionException("Leader not local for partition %s on broker %d"
@@ -767,7 +977,7 @@ class Partition(val topicPartition: TopicPartition,
       tryCompleteDelayedRequests()
     else {
       // probably unblock some follower fetch requests since log end offset has been updated
-      replicaManager.tryCompleteDelayedFetch(new TopicPartitionOperationKey(topicPartition))
+      delayedOperations.checkAndCompleteFetch()
     }
 
     info
@@ -780,34 +990,15 @@ class Partition(val topicPartition: TopicPartition,
                   fetchOnlyFromLeader: Boolean,
                   minOneMessage: Boolean): LogReadInfo = inReadLock(leaderIsrUpdateLock) {
     // decide whether to only fetch from leader
-    val localReplica = localReplicaWithEpochOrException(currentLeaderEpoch, fetchOnlyFromLeader)
-
-    /* Read the LogOffsetMetadata prior to performing the read from the log.
-     * We use the LogOffsetMetadata to determine if a particular replica is in-sync or not.
-     * Using the log end offset after performing the read can lead to a race condition
-     * where data gets appended to the log immediately after the replica has consumed from it
-     * This can cause a replica to always be out of sync.
-     */
-    val initialHighWatermark = localReplica.highWatermark.messageOffset
-    val initialLogStartOffset = localReplica.logStartOffset
-    val initialLogEndOffset = localReplica.logEndOffset
-    val initialLastStableOffset = localReplica.lastStableOffset.messageOffset
-
-    val maxOffsetOpt = fetchIsolation match {
-      case FetchLogEnd => None
-      case FetchHighWatermark => Some(initialHighWatermark)
-      case FetchTxnCommitted => Some(initialLastStableOffset)
-    }
+    val localLog = localLogWithEpochOrException(currentLeaderEpoch, fetchOnlyFromLeader)
 
-    val fetchedData = localReplica.log match {
-      case Some(log) =>
-        log.read(fetchOffset, maxBytes, maxOffsetOpt, minOneMessage,
-          includeAbortedTxns = fetchIsolation == FetchTxnCommitted)
-
-      case None =>
-        error(s"Leader does not have a local log")
-        FetchDataInfo(LogOffsetMetadata.UnknownOffsetMetadata, MemoryRecords.EMPTY)
-    }
+    // Note we use the log end offset prior to the read. This ensures that any appends following
+    // the fetch do not prevent a follower from coming into sync.
+    val initialHighWatermark = localLog.highWatermark
+    val initialLogStartOffset = localLog.logStartOffset
+    val initialLogEndOffset = localLog.logEndOffset
+    val initialLastStableOffset = localLog.lastStableOffset
+    val fetchedData = localLog.read(fetchOffset, maxBytes, fetchIsolation, minOneMessage)
 
     LogReadInfo(
       fetchedData = fetchedData,
@@ -822,12 +1013,12 @@ class Partition(val topicPartition: TopicPartition,
                               currentLeaderEpoch: Optional[Integer],
                               fetchOnlyFromLeader: Boolean): Option[TimestampAndOffset] = inReadLock(leaderIsrUpdateLock) {
     // decide whether to only fetch from leader
-    val localReplica = localReplicaWithEpochOrException(currentLeaderEpoch, fetchOnlyFromLeader)
+    val localLog = localLogWithEpochOrException(currentLeaderEpoch, fetchOnlyFromLeader)
 
     val lastFetchableOffset = isolationLevel match {
-      case Some(IsolationLevel.READ_COMMITTED) => localReplica.lastStableOffset.messageOffset
-      case Some(IsolationLevel.READ_UNCOMMITTED) => localReplica.highWatermark.messageOffset
-      case None => localReplica.logEndOffset
+      case Some(IsolationLevel.READ_COMMITTED) => localLog.lastStableOffset
+      case Some(IsolationLevel.READ_UNCOMMITTED) => localLog.highWatermark
+      case None => localLog.logEndOffset
     }
 
     val epochLogString = if(currentLeaderEpoch.isPresent) {
@@ -839,10 +1030,10 @@ class Partition(val topicPartition: TopicPartition,
     // Only consider throwing an error if we get a client request (isolationLevel is defined) and the start offset
     // is lagging behind the high watermark
     val maybeOffsetsError: Option[ApiException] = leaderEpochStartOffsetOpt
-      .filter(epochStart => isolationLevel.isDefined && epochStart > localReplica.highWatermark.messageOffset)
+      .filter(epochStart => isolationLevel.isDefined && epochStart > localLog.highWatermark)
       .map(epochStart => Errors.OFFSET_NOT_AVAILABLE.exception(s"Failed to fetch offsets for " +
         s"partition $topicPartition with leader $epochLogString as this partition's " +
-        s"high watermark (${localReplica.highWatermark.messageOffset}) is lagging behind the " +
+        s"high watermark (${localLog.highWatermark}) is lagging behind the " +
         s"start offset from the beginning of this epoch ($epochStart)."))
 
     def getOffsetByTimestamp: Option[TimestampAndOffset] = {
@@ -866,37 +1057,21 @@ class Partition(val topicPartition: TopicPartition,
   def fetchOffsetSnapshot(currentLeaderEpoch: Optional[Integer],
                           fetchOnlyFromLeader: Boolean): LogOffsetSnapshot = inReadLock(leaderIsrUpdateLock) {
     // decide whether to only fetch from leader
-    val localReplica = localReplicaWithEpochOrException(currentLeaderEpoch, fetchOnlyFromLeader)
-    localReplica.offsetSnapshot
-  }
-
-  def fetchOffsetSnapshotOrError(currentLeaderEpoch: Optional[Integer],
-                                 fetchOnlyFromLeader: Boolean): Either[LogOffsetSnapshot, Errors] = {
-    inReadLock(leaderIsrUpdateLock) {
-      getLocalReplica(localBrokerId, currentLeaderEpoch, fetchOnlyFromLeader)
-        .left.map(_.offsetSnapshot)
-    }
+    val localLog = localLogWithEpochOrException(currentLeaderEpoch, fetchOnlyFromLeader)
+    localLog.fetchOffsetSnapshot
   }
 
   def legacyFetchOffsetsForTimestamp(timestamp: Long,
                                      maxNumOffsets: Int,
                                      isFromConsumer: Boolean,
                                      fetchOnlyFromLeader: Boolean): Seq[Long] = inReadLock(leaderIsrUpdateLock) {
-    val localReplica = localReplicaWithEpochOrException(Optional.empty(), fetchOnlyFromLeader)
-    val allOffsets = logManager.getLog(topicPartition) match {
-      case Some(log) =>
-        log.legacyFetchOffsetsBefore(timestamp, maxNumOffsets)
-      case None =>
-        if (timestamp == ListOffsetRequest.LATEST_TIMESTAMP || timestamp == ListOffsetRequest.EARLIEST_TIMESTAMP)
-          Seq(0L)
-        else
-          Nil
-    }
+    val localLog = localLogWithEpochOrException(Optional.empty(), fetchOnlyFromLeader)
+    val allOffsets = localLog.legacyFetchOffsetsBefore(timestamp, maxNumOffsets)
 
     if (!isFromConsumer) {
       allOffsets
     } else {
-      val hw = localReplica.highWatermark.messageOffset
+      val hw = localLog.highWatermark
       if (allOffsets.exists(_ > hw))
         hw +: allOffsets.dropWhile(_ > hw)
       else
@@ -906,7 +1081,7 @@ class Partition(val topicPartition: TopicPartition,
 
   def logStartOffset: Long = {
     inReadLock(leaderIsrUpdateLock) {
-      leaderReplicaIfLocal.map(_.log.get.logStartOffset).getOrElse(-1)
+      leaderLogIfLocal.map(_.logStartOffset).getOrElse(-1)
     }
   }
 
@@ -917,20 +1092,17 @@ class Partition(val topicPartition: TopicPartition,
    * Return low watermark of the partition.
    */
   def deleteRecordsOnLeader(offset: Long): LogDeleteRecordsResult = inReadLock(leaderIsrUpdateLock) {
-    leaderReplicaIfLocal match {
-      case Some(leaderReplica) =>
-        if (!leaderReplica.log.get.config.delete)
-          throw new PolicyViolationException(s"Records of partition $topicPartition can not be deleted due to the configured policy")
-
+    leaderLogIfLocal match {
+      case Some(leaderLog) =>
         val convertedOffset = if (offset == DeleteRecordsRequest.HIGH_WATERMARK)
-          leaderReplica.highWatermark.messageOffset
+          leaderLog.highWatermark
         else
           offset
 
         if (convertedOffset < 0)
           throw new OffsetOutOfRangeException(s"The offset $convertedOffset for partition $topicPartition is not valid")
 
-        leaderReplica.maybeIncrementLogStartOffset(convertedOffset)
+        leaderLog.maybeIncrementLogStartOffset(convertedOffset)
         LogDeleteRecordsResult(
           requestedOffset = convertedOffset,
           lowWatermark = lowWatermarkIfLeader)
@@ -945,7 +1117,7 @@ class Partition(val topicPartition: TopicPartition,
     * @param offset offset to be used for truncation
     * @param isFuture True iff the truncation should be performed on the future log of this partition
     */
-  def truncateTo(offset: Long, isFuture: Boolean) {
+  def truncateTo(offset: Long, isFuture: Boolean): Unit = {
     // The read lock is needed to prevent the follower replica from being truncated while ReplicaAlterDirThread
     // is executing maybeDeleteAndSwapFutureReplica() to replace follower replica with the future replica.
     inReadLock(leaderIsrUpdateLock) {
@@ -959,7 +1131,7 @@ class Partition(val topicPartition: TopicPartition,
     * @param newOffset The new offset to start the log with
     * @param isFuture True iff the truncation should be performed on the future log of this partition
     */
-  def truncateFullyAndStartAt(newOffset: Long, isFuture: Boolean) {
+  def truncateFullyAndStartAt(newOffset: Long, isFuture: Boolean): Unit = {
     // The read lock is needed to prevent the follower replica from being truncated while ReplicaAlterDirThread
     // is executing maybeDeleteAndSwapFutureReplica() to replace follower replica with the future replica.
     inReadLock(leaderIsrUpdateLock) {
@@ -984,10 +1156,10 @@ class Partition(val topicPartition: TopicPartition,
                                leaderEpoch: Int,
                                fetchOnlyFromLeader: Boolean): EpochEndOffset = {
     inReadLock(leaderIsrUpdateLock) {
-      val localReplicaOrError = getLocalReplica(localBrokerId, currentLeaderEpoch, fetchOnlyFromLeader)
-      localReplicaOrError match {
-        case Left(replica) =>
-          replica.endOffsetForEpoch(leaderEpoch) match {
+      val localLogOrError = getLocalLog(currentLeaderEpoch, fetchOnlyFromLeader)
+      localLogOrError match {
+        case Left(localLog) =>
+          localLog.endOffsetForEpoch(leaderEpoch) match {
             case Some(epochAndOffset) => new EpochEndOffset(NONE, epochAndOffset.leaderEpoch, epochAndOffset.offset)
             case None => new EpochEndOffset(NONE, UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET)
           }
@@ -997,48 +1169,45 @@ class Partition(val topicPartition: TopicPartition,
     }
   }
 
-  private def updateIsr(newIsr: Set[Replica]) {
-    val newLeaderAndIsr = new LeaderAndIsr(localBrokerId, leaderEpoch, newIsr.map(_.brokerId).toList, zkVersion)
-    val (updateSucceeded, newVersion) = ReplicationUtils.updateLeaderAndIsr(zkClient, topicPartition, newLeaderAndIsr,
-      controllerEpoch)
+  private def expandIsr(newIsr: Set[Int]): Unit = {
+    val newLeaderAndIsr = new LeaderAndIsr(localBrokerId, leaderEpoch, newIsr.toList, zkVersion)
+    val zkVersionOpt = stateStore.expandIsr(controllerEpoch, newLeaderAndIsr)
+    maybeUpdateIsrAndVersion(newIsr, zkVersionOpt)
+  }
 
-    if (updateSucceeded) {
-      replicaManager.recordIsrChange(topicPartition)
-      inSyncReplicas = newIsr
-      zkVersion = newVersion
-      trace("ISR updated to [%s] and zkVersion updated to [%d]".format(newIsr.mkString(","), zkVersion))
-    } else {
-      replicaManager.failedIsrUpdatesRate.mark()
-      info("Cached zkVersion [%d] not equal to that in zookeeper, skip updating ISR".format(zkVersion))
-    }
+  private def shrinkIsr(newIsr: Set[Int]): Unit = {
+    val newLeaderAndIsr = new LeaderAndIsr(localBrokerId, leaderEpoch, newIsr.toList, zkVersion)
+    val zkVersionOpt = stateStore.shrinkIsr(controllerEpoch, newLeaderAndIsr)
+    maybeUpdateIsrAndVersion(newIsr, zkVersionOpt)
   }
 
-  /**
-   * remove deleted log metrics
-   */
-  def removePartitionMetrics() {
-    removeMetric("UnderReplicated", tags)
-    removeMetric("UnderMinIsr", tags)
-    removeMetric("InSyncReplicasCount", tags)
-    removeMetric("ReplicasCount", tags)
-    removeMetric("LastStableOffsetLag", tags)
+  private def maybeUpdateIsrAndVersion(isr: Set[Int], zkVersionOpt: Option[Int]): Unit = {
+    zkVersionOpt match {
+      case Some(newVersion) =>
+        inSyncReplicaIds = isr
+        zkVersion = newVersion
+        info("ISR updated to [%s] and zkVersion updated to [%d]".format(isr.mkString(","), zkVersion))
+
+      case None =>
+        info(s"Cached zkVersion $zkVersion not equal to that in zookeeper, skip updating ISR")
+    }
   }
 
   override def equals(that: Any): Boolean = that match {
-    case other: Partition => partitionId == other.partitionId && topic == other.topic && isOffline == other.isOffline
+    case other: Partition => partitionId == other.partitionId && topic == other.topic
     case _ => false
   }
 
   override def hashCode: Int =
-    31 + topic.hashCode + 17 * partitionId + (if (isOffline) 1 else 0)
+    31 + topic.hashCode + 17 * partitionId
 
-  override def toString(): String = {
+  override def toString: String = {
     val partitionString = new StringBuilder
     partitionString.append("Topic: " + topic)
     partitionString.append("; Partition: " + partitionId)
     partitionString.append("; Leader: " + leaderReplicaIdOpt)
-    partitionString.append("; AllReplicas: " + allReplicasMap.keys.mkString(","))
-    partitionString.append("; InSyncReplicas: " + inSyncReplicas.map(_.brokerId).mkString(","))
+    partitionString.append("; AllReplicaIds: " + allReplicaIds.mkString(","))
+    partitionString.append("; InSyncReplicaIds: " + inSyncReplicaIds.mkString(","))
     partitionString.toString
   }
 }
diff --git a/core/src/main/scala/kafka/cluster/Replica.scala b/core/src/main/scala/kafka/cluster/Replica.scala
index 32055d7a5e30d..f9de7ba784c0b 100644
--- a/core/src/main/scala/kafka/cluster/Replica.scala
+++ b/core/src/main/scala/kafka/cluster/Replica.scala
@@ -17,26 +17,18 @@
 
 package kafka.cluster
 
-import kafka.log.{Log, LogOffsetSnapshot}
+import kafka.log.{Log}
 import kafka.utils.Logging
-import kafka.server.{LogOffsetMetadata, LogReadResult, OffsetAndEpoch}
-import org.apache.kafka.common.{KafkaException, TopicPartition}
-import org.apache.kafka.common.errors.OffsetOutOfRangeException
-import org.apache.kafka.common.utils.Time
-
-class Replica(val brokerId: Int,
-              val topicPartition: TopicPartition,
-              time: Time = Time.SYSTEM,
-              initialHighWatermarkValue: Long = 0L,
-              @volatile var log: Option[Log] = None) extends Logging {
-  // the high watermark offset value, in non-leader replicas only its message offsets are kept
-  @volatile private[this] var highWatermarkMetadata = new LogOffsetMetadata(initialHighWatermarkValue)
+import kafka.server.{LogOffsetMetadata}
+import org.apache.kafka.common.{TopicPartition}
+
+class Replica(val brokerId: Int, val topicPartition: TopicPartition) extends Logging {
   // the log end offset value, kept in all replicas;
   // for local replica it is the log's end offset, for remote replicas its value is only updated by follower fetch
   @volatile private[this] var _logEndOffsetMetadata = LogOffsetMetadata.UnknownOffsetMetadata
   // the log start offset value, kept in all replicas;
   // for local replica it is the log's start offset, for remote replicas its value is only updated by follower fetch
-  @volatile private[this] var _logStartOffset = Log.UnknownLogStartOffset
+  @volatile private[this] var _logStartOffset = Log.UnknownOffset
 
   // The log end offset value at the time the leader received the last FetchRequest from this follower
   // This is used to determine the lastCaughtUpTimeMs of the follower
@@ -50,12 +42,19 @@ class Replica(val brokerId: Int,
   // the LEO of leader at time t. This is used to determine the lag of this follower and ISR of this partition.
   @volatile private[this] var _lastCaughtUpTimeMs = 0L
 
-  def isLocal: Boolean = log.isDefined
+  // highWatermark is the leader's high watermark after the most recent FetchRequest from this follower. This is
+  // used to determine the maximum HW this follower knows about. See KIP-392
+  @volatile private[this] var _lastSentHighWatermark = 0L
+
+  def logStartOffset: Long = _logStartOffset
+
+  def logEndOffsetMetadata: LogOffsetMetadata = _logEndOffsetMetadata
+
+  def logEndOffset: Long = logEndOffsetMetadata.messageOffset
 
   def lastCaughtUpTimeMs: Long = _lastCaughtUpTimeMs
 
-  info(s"Replica loaded for partition $topicPartition with initial high watermark $initialHighWatermarkValue")
-  log.foreach(_.onHighWatermarkIncremented(initialHighWatermarkValue))
+  def lastSentHighWatermark: Long = _lastSentHighWatermark
 
   /*
    * If the FetchRequest reads up to the log end offset of the leader when the current fetch request is received,
@@ -69,165 +68,64 @@ class Replica(val brokerId: Int,
    * fetch request is always smaller than the leader's LEO, which can happen if small produce requests are received at
    * high frequency.
    */
-  def updateLogReadResult(logReadResult: LogReadResult) {
-    if (logReadResult.info.fetchOffsetMetadata.messageOffset >= logReadResult.leaderLogEndOffset)
-      _lastCaughtUpTimeMs = math.max(_lastCaughtUpTimeMs, logReadResult.fetchTimeMs)
-    else if (logReadResult.info.fetchOffsetMetadata.messageOffset >= lastFetchLeaderLogEndOffset)
+  def updateFetchState(followerFetchOffsetMetadata: LogOffsetMetadata,
+                       followerStartOffset: Long,
+                       followerFetchTimeMs: Long,
+                       leaderEndOffset: Long,
+                       lastSentHighwatermark: Long): Unit = {
+    if (followerFetchOffsetMetadata.messageOffset >= leaderEndOffset)
+      _lastCaughtUpTimeMs = math.max(_lastCaughtUpTimeMs, followerFetchTimeMs)
+    else if (followerFetchOffsetMetadata.messageOffset >= lastFetchLeaderLogEndOffset)
       _lastCaughtUpTimeMs = math.max(_lastCaughtUpTimeMs, lastFetchTimeMs)
 
-    logStartOffset = logReadResult.followerLogStartOffset
-    logEndOffsetMetadata = logReadResult.info.fetchOffsetMetadata
-    lastFetchLeaderLogEndOffset = logReadResult.leaderLogEndOffset
-    lastFetchTimeMs = logReadResult.fetchTimeMs
-  }
-
-  def resetLastCaughtUpTime(curLeaderLogEndOffset: Long, curTimeMs: Long, lastCaughtUpTimeMs: Long) {
-    lastFetchLeaderLogEndOffset = curLeaderLogEndOffset
-    lastFetchTimeMs = curTimeMs
-    _lastCaughtUpTimeMs = lastCaughtUpTimeMs
-  }
-
-  private def logEndOffsetMetadata_=(newLogEndOffset: LogOffsetMetadata) {
-    if (isLocal) {
-      throw new KafkaException(s"Should not set log end offset on partition $topicPartition's local replica $brokerId")
-    } else {
-      _logEndOffsetMetadata = newLogEndOffset
-      trace(s"Setting log end offset for replica $brokerId for partition $topicPartition to [${_logEndOffsetMetadata}]")
-    }
-  }
-
-  def latestEpoch: Option[Int] = {
-    if (isLocal) {
-      log.get.latestEpoch
-    } else {
-      throw new KafkaException(s"Cannot get latest epoch of non-local replica of $topicPartition")
-    }
-  }
-
-  def endOffsetForEpoch(leaderEpoch: Int): Option[OffsetAndEpoch] = {
-    if (isLocal) {
-      log.get.endOffsetForEpoch(leaderEpoch)
-    } else {
-      throw new KafkaException(s"Cannot lookup end offset for epoch of non-local replica of $topicPartition")
-    }
-  }
-
-  def logEndOffsetMetadata: LogOffsetMetadata =
-    if (isLocal)
-      log.get.logEndOffsetMetadata
-    else
-      _logEndOffsetMetadata
-
-  def logEndOffset: Long =
-    logEndOffsetMetadata.messageOffset
-
-  /**
-   * Increment the log start offset if the new offset is greater than the previous log start offset. The replica
-   * must be local and the new log start offset must be lower than the current high watermark.
-   */
-  def maybeIncrementLogStartOffset(newLogStartOffset: Long) {
-    if (isLocal) {
-      if (newLogStartOffset > highWatermark.messageOffset)
-        throw new OffsetOutOfRangeException(s"Cannot increment the log start offset to $newLogStartOffset of partition $topicPartition " +
-          s"since it is larger than the high watermark ${highWatermark.messageOffset}")
-      log.get.maybeIncrementLogStartOffset(newLogStartOffset)
-    } else {
-      throw new KafkaException(s"Should not try to delete records on partition $topicPartition's non-local replica $brokerId")
-    }
-  }
-
-  private def logStartOffset_=(newLogStartOffset: Long) {
-    if (isLocal) {
-      throw new KafkaException(s"Should not set log start offset on partition $topicPartition's local replica $brokerId " +
-                               s"without attempting to delete records of the log")
-    } else {
-      _logStartOffset = newLogStartOffset
-      trace(s"Setting log start offset for remote replica $brokerId for partition $topicPartition to [$newLogStartOffset]")
-    }
-  }
-
-  def logStartOffset: Long =
-    if (isLocal)
-      log.get.logStartOffset
-    else
-      _logStartOffset
-
-  def highWatermark_=(newHighWatermark: LogOffsetMetadata) {
-    if (isLocal) {
-      if (newHighWatermark.messageOffset < 0)
-        throw new IllegalArgumentException("High watermark offset should be non-negative")
-
-      highWatermarkMetadata = newHighWatermark
-      log.foreach(_.onHighWatermarkIncremented(newHighWatermark.messageOffset))
-      trace(s"Setting high watermark for replica $brokerId partition $topicPartition to [$newHighWatermark]")
-    } else {
-      throw new KafkaException(s"Should not set high watermark on partition $topicPartition's non-local replica $brokerId")
-    }
+    _logStartOffset = followerStartOffset
+    _logEndOffsetMetadata = followerFetchOffsetMetadata
+    lastFetchLeaderLogEndOffset = leaderEndOffset
+    lastFetchTimeMs = followerFetchTimeMs
+    updateLastSentHighWatermark(lastSentHighwatermark)
+    trace(s"Updated state of replica to $this")
   }
 
-  def highWatermark: LogOffsetMetadata = highWatermarkMetadata
-
   /**
-   * The last stable offset (LSO) is defined as the first offset such that all lower offsets have been "decided."
-   * Non-transactional messages are considered decided immediately, but transactional messages are only decided when
-   * the corresponding COMMIT or ABORT marker is written. This implies that the last stable offset will be equal
-   * to the high watermark if there are no transactional messages in the log. Note also that the LSO cannot advance
-   * beyond the high watermark.
-   */
-  def lastStableOffset: LogOffsetMetadata = {
-    log.map { log =>
-      log.firstUnstableOffset match {
-        case Some(offsetMetadata) if offsetMetadata.messageOffset < highWatermark.messageOffset => offsetMetadata
-        case _ => highWatermark
-      }
-    }.getOrElse(throw new KafkaException(s"Cannot fetch last stable offset on partition $topicPartition's " +
-      s"non-local replica $brokerId"))
-  }
-
-  /*
-   * Convert hw to local offset metadata by reading the log at the hw offset.
-   * If the hw offset is out of range, return the first offset of the first log segment as the offset metadata.
-   */
-  def convertHWToLocalOffsetMetadata() {
-    if (isLocal) {
-      highWatermarkMetadata = log.get.convertToOffsetMetadata(highWatermarkMetadata.messageOffset).getOrElse {
-        log.get.convertToOffsetMetadata(logStartOffset).getOrElse {
-          val firstSegmentOffset = log.get.logSegments.head.baseOffset
-          new LogOffsetMetadata(firstSegmentOffset, firstSegmentOffset, 0)
-        }
-      }
-    } else {
-      throw new KafkaException(s"Should not construct complete high watermark on partition $topicPartition's non-local replica $brokerId")
-    }
-  }
-
-  def offsetSnapshot: LogOffsetSnapshot = {
-    LogOffsetSnapshot(
-      logStartOffset = logStartOffset,
-      logEndOffset = logEndOffsetMetadata,
-      highWatermark =  highWatermark,
-      lastStableOffset = lastStableOffset)
+    * Update the high watermark of this remote replica. This is used to track what we think is the last known HW to
+    * a remote follower. Since this is recorded when we send a response, there is no way to guarantee that the follower
+    * actually receives this HW. So we consider this to be an upper bound on what the follower knows.
+    *
+    * When handling fetches, the last sent high watermark for a replica is checked to see if we should return immediately
+    * in order to propagate the HW more expeditiously. See KIP-392
+    */
+  private def updateLastSentHighWatermark(highWatermark: Long): Unit = {
+    _lastSentHighWatermark = highWatermark
+    trace(s"Updated HW of replica to $highWatermark")
   }
 
-  override def equals(that: Any): Boolean = that match {
-    case other: Replica => brokerId == other.brokerId && topicPartition == other.topicPartition
-    case _ => false
+  def resetLastCaughtUpTime(curLeaderLogEndOffset: Long, curTimeMs: Long, lastCaughtUpTimeMs: Long): Unit = {
+    lastFetchLeaderLogEndOffset = curLeaderLogEndOffset
+    lastFetchTimeMs = curTimeMs
+    _lastCaughtUpTimeMs = lastCaughtUpTimeMs
+    trace(s"Reset state of replica to $this")
   }
 
-  override def hashCode: Int = 31 + topicPartition.hashCode + 17 * brokerId
-
   override def toString: String = {
     val replicaString = new StringBuilder
     replicaString.append("Replica(replicaId=" + brokerId)
     replicaString.append(s", topic=${topicPartition.topic}")
     replicaString.append(s", partition=${topicPartition.partition}")
-    replicaString.append(s", isLocal=$isLocal")
     replicaString.append(s", lastCaughtUpTimeMs=$lastCaughtUpTimeMs")
-    if (isLocal) {
-      replicaString.append(s", highWatermark=$highWatermark")
-      replicaString.append(s", lastStableOffset=$lastStableOffset")
-    }
+    replicaString.append(s", logStartOffset=$logStartOffset")
+    replicaString.append(s", logEndOffset=$logEndOffset")
+    replicaString.append(s", logEndOffsetMetadata=$logEndOffsetMetadata")
+    replicaString.append(s", lastFetchLeaderLogEndOffset=$lastFetchLeaderLogEndOffset")
+    replicaString.append(s", lastFetchTimeMs=$lastFetchTimeMs")
+    replicaString.append(s", lastSentHighWatermark=$lastSentHighWatermark")
     replicaString.append(")")
     replicaString.toString
   }
+
+  override def equals(that: Any): Boolean = that match {
+    case other: Replica => brokerId == other.brokerId && topicPartition == other.topicPartition
+    case _ => false
+  }
+
+  override def hashCode: Int = 31 + topicPartition.hashCode + 17 * brokerId
 }
diff --git a/core/src/main/scala/kafka/common/Config.scala b/core/src/main/scala/kafka/common/Config.scala
index 4110ba7b07315..f56cca8bd0528 100644
--- a/core/src/main/scala/kafka/common/Config.scala
+++ b/core/src/main/scala/kafka/common/Config.scala
@@ -23,7 +23,7 @@ import org.apache.kafka.common.errors.InvalidConfigurationException
 
 trait Config extends Logging {
 
-  def validateChars(prop: String, value: String) {
+  def validateChars(prop: String, value: String): Unit = {
     val legalChars = "[a-zA-Z0-9\\._\\-]"
     val rgx = new Regex(legalChars + "*")
 
diff --git a/core/src/main/scala/kafka/common/TopicAndPartition.scala b/core/src/main/scala/kafka/common/InconsistentBrokerMetadataException.scala
similarity index 71%
rename from core/src/main/scala/kafka/common/TopicAndPartition.scala
rename to core/src/main/scala/kafka/common/InconsistentBrokerMetadataException.scala
index 6c276952f5aa7..2b11512e44ce2 100644
--- a/core/src/main/scala/kafka/common/TopicAndPartition.scala
+++ b/core/src/main/scala/kafka/common/InconsistentBrokerMetadataException.scala
@@ -1,7 +1,3 @@
-package kafka.common
-
-import org.apache.kafka.common.TopicPartition
-
 /**
  * Licensed to the Apache Software Foundation (ASF) under one or more
  * contributor license agreements.  See the NOTICE file distributed with
@@ -19,12 +15,13 @@ import org.apache.kafka.common.TopicPartition
  * limitations under the License.
  */
 
+package kafka.common
+
 /**
- * Convenience case class since (topic, partition) pairs are ubiquitous.
+ * Indicates the BrokerMetadata stored in logDirs is not consistent across logDirs.
  */
-case class TopicAndPartition(topic: String, partition: Int) {
-
-  def this(topicPartition: TopicPartition) = this(topicPartition.topic, topicPartition.partition)
-
-  override def toString: String = s"$topic-$partition"
+class InconsistentBrokerMetadataException(message: String, cause: Throwable) extends RuntimeException(message, cause) {
+  def this(message: String) = this(message, null)
+  def this(cause: Throwable) = this(null, cause)
+  def this() = this(null, null)
 }
diff --git a/core/src/main/scala/kafka/tools/ZooKeeperMainWrapper.scala b/core/src/main/scala/kafka/common/InconsistentClusterIdException.scala
similarity index 56%
rename from core/src/main/scala/kafka/tools/ZooKeeperMainWrapper.scala
rename to core/src/main/scala/kafka/common/InconsistentClusterIdException.scala
index 14d2ceb8a647e..6868dd8780de9 100644
--- a/core/src/main/scala/kafka/tools/ZooKeeperMainWrapper.scala
+++ b/core/src/main/scala/kafka/common/InconsistentClusterIdException.scala
@@ -1,4 +1,4 @@
-/*
+/**
  * 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.
@@ -6,7 +6,7 @@
  * (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
+ *    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,
@@ -15,26 +15,13 @@
  * limitations under the License.
  */
 
-package kafka.tools
-
-import kafka.utils.Exit
-import org.apache.zookeeper.ZooKeeperMain
-
-class ZooKeeperMainWrapper(args: Array[String]) extends ZooKeeperMain(args) {
-  def runCmd(): Unit = {
-    processCmd(this.cl)
-    Exit.exit(0)
-  }
-}
+package kafka.common
 
 /**
- * ZooKeeper 3.4.6 broke being able to pass commands on command line.
- * See ZOOKEEPER-1897.  This class is a hack to restore this facility.
+ * Indicates the clusterId stored in logDirs is not consistent with the clusterIs stored in ZK.
  */
-object ZooKeeperMainWrapper {
-
-  def main(args: Array[String]): Unit = {
-    val main: ZooKeeperMainWrapper = new ZooKeeperMainWrapper(args)
-    main.runCmd()
-  }
+class InconsistentClusterIdException(message: String, cause: Throwable) extends RuntimeException(message, cause) {
+  def this(message: String) = this(message, null)
+  def this(cause: Throwable) = this(null, cause)
+  def this() = this(null, null)
 }
diff --git a/core/src/main/scala/kafka/common/InterBrokerSendThread.scala b/core/src/main/scala/kafka/common/InterBrokerSendThread.scala
index aedaac79d2774..1551704203a91 100644
--- a/core/src/main/scala/kafka/common/InterBrokerSendThread.scala
+++ b/core/src/main/scala/kafka/common/InterBrokerSendThread.scala
@@ -51,7 +51,7 @@ abstract class InterBrokerSendThread(name: String,
     awaitShutdown()
   }
 
-  override def doWork() {
+  override def doWork(): Unit = {
     var now = time.milliseconds()
 
     generateRequests().foreach { request =>
diff --git a/core/src/main/scala/kafka/common/KafkaException.scala b/core/src/main/scala/kafka/common/KafkaException.scala
index 61b3ba316e0e7..9c34dd9bd78b4 100644
--- a/core/src/main/scala/kafka/common/KafkaException.scala
+++ b/core/src/main/scala/kafka/common/KafkaException.scala
@@ -19,9 +19,8 @@ package kafka.common
 /**
  * Usage of this class is discouraged. Use org.apache.kafka.common.KafkaException instead.
  *
- * This class will be removed once ZkUtils and the kafka.security.auth classes are removed.
- * The former is internal, but widely used, so we are leaving it in the codebase for now.
-*/
+ * This class will be removed once kafka.security.auth classes are removed.
+ */
 class KafkaException(message: String, t: Throwable) extends RuntimeException(message, t) {
   def this(message: String) = this(message, null)
   def this(t: Throwable) = this("", t)
diff --git a/core/src/main/scala/kafka/common/MessageFormatter.scala b/core/src/main/scala/kafka/common/MessageFormatter.scala
index ef3c7238b7b11..9f6c3fa94adc5 100644
--- a/core/src/main/scala/kafka/common/MessageFormatter.scala
+++ b/core/src/main/scala/kafka/common/MessageFormatter.scala
@@ -30,10 +30,10 @@ import org.apache.kafka.clients.consumer.ConsumerRecord
   */
 trait MessageFormatter {
 
-  def init(props: Properties) {}
+  def init(props: Properties): Unit = {}
 
   def writeTo(consumerRecord: ConsumerRecord[Array[Byte], Array[Byte]], output: PrintStream): Unit
 
-  def close() {}
+  def close(): Unit = {}
 
 }
diff --git a/core/src/main/scala/kafka/common/MessageReader.scala b/core/src/main/scala/kafka/common/MessageReader.scala
index 56b55ce891f9e..de456e16ae532 100644
--- a/core/src/main/scala/kafka/common/MessageReader.scala
+++ b/core/src/main/scala/kafka/common/MessageReader.scala
@@ -30,10 +30,10 @@ import org.apache.kafka.clients.producer.ProducerRecord
   */
 trait MessageReader {
 
-  def init(inputStream: InputStream, props: Properties) {}
+  def init(inputStream: InputStream, props: Properties): Unit = {}
 
   def readMessage(): ProducerRecord[Array[Byte], Array[Byte]]
 
-  def close() {}
+  def close(): Unit = {}
 
 }
diff --git a/core/src/main/scala/kafka/common/AppInfo.scala b/core/src/main/scala/kafka/common/RecordValidationException.scala
similarity index 52%
rename from core/src/main/scala/kafka/common/AppInfo.scala
rename to core/src/main/scala/kafka/common/RecordValidationException.scala
index f77bdf55a69b0..2acdf84c00c06 100644
--- a/core/src/main/scala/kafka/common/AppInfo.scala
+++ b/core/src/main/scala/kafka/common/RecordValidationException.scala
@@ -6,7 +6,7 @@
  * (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
+ *    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,
@@ -17,38 +17,9 @@
 
 package kafka.common
 
-import com.yammer.metrics.core.Gauge
-import kafka.metrics.KafkaMetricsGroup
-import org.apache.kafka.common.utils.AppInfoParser
+import org.apache.kafka.common.errors.ApiException
+import org.apache.kafka.common.requests.ProduceResponse.RecordError
 
-object AppInfo extends KafkaMetricsGroup {
-  private var isRegistered = false
-  private val lock = new Object()
-
-  def registerInfo(): Unit = {
-    lock.synchronized {
-      if (isRegistered) {
-        return
-      }
-    }
-
-    newGauge("Version",
-      new Gauge[String] {
-        def value = {
-          AppInfoParser.getVersion()
-        }
-      })
-
-    newGauge("CommitID",
-      new Gauge[String] {
-        def value = {
-          AppInfoParser.getCommitId()
-        }
-      })
-
-    lock.synchronized {
-      isRegistered = true
-    }
-
-  }
+class RecordValidationException(val invalidException: ApiException,
+                                val recordErrors: List[RecordError]) extends RuntimeException {
 }
diff --git a/core/src/main/scala/kafka/common/ZkNodeChangeNotificationListener.scala b/core/src/main/scala/kafka/common/ZkNodeChangeNotificationListener.scala
index 65c350632b2ab..42341cf72e47b 100644
--- a/core/src/main/scala/kafka/common/ZkNodeChangeNotificationListener.scala
+++ b/core/src/main/scala/kafka/common/ZkNodeChangeNotificationListener.scala
@@ -25,13 +25,14 @@ import kafka.zk.{KafkaZkClient, StateChangeHandlers}
 import kafka.zookeeper.{StateChangeHandler, ZNodeChildChangeHandler}
 import org.apache.kafka.common.utils.Time
 
+import scala.collection.Seq
 import scala.util.{Failure, Try}
 
 /**
  * Handle the notificationMessage.
  */
 trait NotificationHandler {
-  def processNotification(notificationMessage: Array[Byte])
+  def processNotification(notificationMessage: Array[Byte]): Unit
 }
 
 /**
@@ -59,7 +60,7 @@ class ZkNodeChangeNotificationListener(private val zkClient: KafkaZkClient,
   private val thread = new ChangeEventProcessThread(s"$seqNodeRoot-event-process-thread")
   private val isClosed = new AtomicBoolean(false)
 
-  def init() {
+  def init(): Unit = {
     zkClient.registerStateChangeHandler(ZkStateChangeHandler)
     zkClient.registerZNodeChildChangeHandler(ChangeNotificationHandler)
     addChangeNotification()
@@ -77,7 +78,7 @@ class ZkNodeChangeNotificationListener(private val zkClient: KafkaZkClient,
   /**
    * Process notifications
    */
-  private def processNotifications() {
+  private def processNotifications(): Unit = {
     try {
       val notifications = zkClient.getChildren(seqNodeRoot).sorted
       if (notifications.nonEmpty) {
@@ -125,7 +126,7 @@ class ZkNodeChangeNotificationListener(private val zkClient: KafkaZkClient,
    * @param now
    * @param notifications
    */
-  private def purgeObsoleteNotifications(now: Long, notifications: Seq[String]) {
+  private def purgeObsoleteNotifications(now: Long, notifications: Seq[String]): Unit = {
     for (notification <- notifications.sorted) {
       val notificationNode = seqNodeRoot + "/" + notification
       val (data, stat) = zkClient.getDataAndStat(notificationNode)
diff --git a/core/src/main/scala/kafka/controller/ControllerChannelManager.scala b/core/src/main/scala/kafka/controller/ControllerChannelManager.scala
index 3776b6903041a..a1dad5bcc48af 100755
--- a/core/src/main/scala/kafka/controller/ControllerChannelManager.scala
+++ b/core/src/main/scala/kafka/controller/ControllerChannelManager.scala
@@ -26,32 +26,38 @@ import kafka.metrics.KafkaMetricsGroup
 import kafka.server.KafkaConfig
 import kafka.utils._
 import org.apache.kafka.clients._
+import org.apache.kafka.common.message.LeaderAndIsrRequestData.LeaderAndIsrPartitionState
 import org.apache.kafka.common.metrics.Metrics
 import org.apache.kafka.common.network._
-import org.apache.kafka.common.protocol.ApiKeys
-import org.apache.kafka.common.requests.UpdateMetadataRequest.EndPoint
+import org.apache.kafka.common.protocol.{ApiKeys, Errors}
+import org.apache.kafka.common.message.UpdateMetadataRequestData.{UpdateMetadataBroker, UpdateMetadataEndpoint, UpdateMetadataPartitionState}
 import org.apache.kafka.common.requests._
 import org.apache.kafka.common.security.JaasContext
 import org.apache.kafka.common.security.auth.SecurityProtocol
 import org.apache.kafka.common.utils.{LogContext, Time}
-import org.apache.kafka.common.{KafkaException, Node, TopicPartition}
+import org.apache.kafka.common.{KafkaException, Node, Reconfigurable, TopicPartition}
 
 import scala.collection.JavaConverters._
-import scala.collection.mutable.HashMap
-import scala.collection.{Set, mutable}
+import scala.collection.mutable.{HashMap, ListBuffer}
+import scala.collection.{Seq, Map, Set, mutable}
 
 object ControllerChannelManager {
   val QueueSizeMetricName = "QueueSize"
   val RequestRateAndQueueTimeMetricName = "RequestRateAndQueueTimeMs"
 }
 
-class ControllerChannelManager(controllerContext: ControllerContext, config: KafkaConfig, time: Time, metrics: Metrics,
-                               stateChangeLogger: StateChangeLogger, threadNamePrefix: Option[String] = None) extends Logging with KafkaMetricsGroup {
+class ControllerChannelManager(controllerContext: ControllerContext,
+                               config: KafkaConfig,
+                               time: Time,
+                               metrics: Metrics,
+                               stateChangeLogger: StateChangeLogger,
+                               threadNamePrefix: Option[String] = None) extends Logging with KafkaMetricsGroup {
   import ControllerChannelManager._
+
   protected val brokerStateInfo = new HashMap[Int, ControllerBrokerStateInfo]
   private val brokerLock = new Object
   this.logIdent = "[Channel manager on controller " + config.brokerId + "]: "
-
+  val brokerResponseSensors: mutable.Map[ApiKeys, BrokerResponseTimeStats] = mutable.HashMap.empty
   newGauge(
     "TotalQueueSize",
     new Gauge[Int] {
@@ -61,34 +67,49 @@ class ControllerChannelManager(controllerContext: ControllerContext, config: Kaf
     }
   )
 
-  controllerContext.liveBrokers.foreach(addNewBroker)
-
   def startup() = {
+    controllerContext.liveOrShuttingDownBrokers.foreach(addNewBroker)
+
     brokerLock synchronized {
       brokerStateInfo.foreach(brokerState => startRequestSendThread(brokerState._1))
     }
+    initBrokerResponseSensors()
   }
 
   def shutdown() = {
     brokerLock synchronized {
-      brokerStateInfo.values.foreach(removeExistingBroker)
+      brokerStateInfo.values.toList.foreach(removeExistingBroker)
     }
+    removeBrokerResponseSensors()
   }
 
-  def sendRequest(brokerId: Int, apiKey: ApiKeys, request: AbstractControlRequest.Builder[_ <: AbstractControlRequest],
-                  callback: AbstractResponse => Unit = null) {
+  def initBrokerResponseSensors(): Unit = {
+    Array(ApiKeys.STOP_REPLICA, ApiKeys.LEADER_AND_ISR, ApiKeys.UPDATE_METADATA).foreach { k: ApiKeys =>
+      brokerResponseSensors.put(k, new BrokerResponseTimeStats(k))
+    }
+  }
+
+  def removeBrokerResponseSensors(): Unit = {
+    brokerResponseSensors.keySet.foreach { k: ApiKeys =>
+      brokerResponseSensors(k).removeMetrics()
+      brokerResponseSensors.remove(k)
+    }
+  }
+
+  def sendRequest(brokerId: Int, request: AbstractControlRequest.Builder[_ <: AbstractControlRequest],
+                  callback: AbstractResponse => Unit = null): Unit = {
     brokerLock synchronized {
       val stateInfoOpt = brokerStateInfo.get(brokerId)
       stateInfoOpt match {
         case Some(stateInfo) =>
-          stateInfo.messageQueue.put(QueueItem(apiKey, request, callback, time.milliseconds()))
+          stateInfo.messageQueue.put(QueueItem(request.apiKey, request, callback, time.milliseconds()))
         case None =>
           warn(s"Not sending request $request to broker $brokerId, since it is offline.")
       }
     }
   }
 
-  def addBroker(broker: Broker) {
+  def addBroker(broker: Broker): Unit = {
     // be careful here. Maybe the startup() API has already started the request send thread
     brokerLock synchronized {
       if (!brokerStateInfo.contains(broker.id)) {
@@ -98,20 +119,20 @@ class ControllerChannelManager(controllerContext: ControllerContext, config: Kaf
     }
   }
 
-  def removeBroker(brokerId: Int) {
+  def removeBroker(brokerId: Int): Unit = {
     brokerLock synchronized {
       removeExistingBroker(brokerStateInfo(brokerId))
     }
   }
 
-  private def addNewBroker(broker: Broker) {
+  private def addNewBroker(broker: Broker): Unit = {
     val messageQueue = new LinkedBlockingQueue[QueueItem]
     debug(s"Controller ${config.brokerId} trying to connect to broker ${broker.id}")
     val controllerToBrokerListenerName = config.controlPlaneListenerName.getOrElse(config.interBrokerListenerName)
     val controllerToBrokerSecurityProtocol = config.controlPlaneSecurityProtocol.getOrElse(config.interBrokerSecurityProtocol)
     val brokerNode = broker.node(controllerToBrokerListenerName)
     val logContext = new LogContext(s"[Controller id=${config.brokerId}, targetBrokerId=${brokerNode.idString}] ")
-    val networkClient = {
+    val (networkClient, reconfigurableChannelBuilder) = {
       val channelBuilder = ChannelBuilders.clientChannelBuilder(
         controllerToBrokerSecurityProtocol,
         JaasContext.Type.SERVER,
@@ -121,6 +142,12 @@ class ControllerChannelManager(controllerContext: ControllerContext, config: Kaf
         time,
         config.saslInterBrokerHandshakeRequestEnable
       )
+      val reconfigurableChannelBuilder = channelBuilder match {
+        case reconfigurable: Reconfigurable =>
+          config.addReconfigurable(reconfigurable)
+          Some(reconfigurable)
+        case _ => None
+      }
       val selector = new Selector(
         NetworkReceive.UNLIMITED,
         Selector.NO_IDLE_TIMEOUT_MS,
@@ -132,7 +159,7 @@ class ControllerChannelManager(controllerContext: ControllerContext, config: Kaf
         channelBuilder,
         logContext
       )
-      new NetworkClient(
+      val networkClient = new NetworkClient(
         selector,
         new ManualMetadataUpdater(Seq(brokerNode).asJava),
         config.brokerId.toString,
@@ -148,6 +175,7 @@ class ControllerChannelManager(controllerContext: ControllerContext, config: Kaf
         new ApiVersions,
         logContext
       )
+      (networkClient, reconfigurableChannelBuilder)
     }
     val threadName = threadNamePrefix match {
       case None => s"Controller-${config.brokerId}-to-broker-${broker.id}-send-thread"
@@ -159,7 +187,7 @@ class ControllerChannelManager(controllerContext: ControllerContext, config: Kaf
     )
 
     val requestThread = new RequestSendThread(config.brokerId, controllerContext, messageQueue, networkClient,
-      brokerNode, config, time, requestRateAndQueueTimeMetrics, stateChangeLogger, threadName)
+      brokerNode, config, time, requestRateAndQueueTimeMetrics, stateChangeLogger, threadName, this)
     requestThread.setDaemon(false)
 
     val queueSizeGauge = newGauge(
@@ -171,17 +199,18 @@ class ControllerChannelManager(controllerContext: ControllerContext, config: Kaf
     )
 
     brokerStateInfo.put(broker.id, ControllerBrokerStateInfo(networkClient, brokerNode, messageQueue,
-      requestThread, queueSizeGauge, requestRateAndQueueTimeMetrics))
+      requestThread, queueSizeGauge, requestRateAndQueueTimeMetrics, reconfigurableChannelBuilder))
   }
 
   private def brokerMetricTags(brokerId: Int) = Map("broker-id" -> brokerId.toString)
 
-  private def removeExistingBroker(brokerState: ControllerBrokerStateInfo) {
+  private def removeExistingBroker(brokerState: ControllerBrokerStateInfo): Unit = {
     try {
       // Shutdown the RequestSendThread before closing the NetworkClient to avoid the concurrent use of the
       // non-threadsafe classes as described in KAFKA-4959.
       // The call to shutdownLatch.await() in ShutdownableThread.shutdown() serves as a synchronization barrier that
       // hands off the NetworkClient from the RequestSendThread to the ZkEventThread.
+      brokerState.reconfigurableChannelBuilder.foreach(config.removeReconfigurable)
       brokerState.requestSendThread.shutdown()
       brokerState.networkClient.close()
       brokerState.messageQueue.clear()
@@ -193,7 +222,7 @@ class ControllerChannelManager(controllerContext: ControllerContext, config: Kaf
     }
   }
 
-  protected def startRequestSendThread(brokerId: Int) {
+  protected def startRequestSendThread(brokerId: Int): Unit = {
     val requestThread = brokerStateInfo(brokerId).requestSendThread
     if (requestThread.getState == Thread.State.NEW)
       requestThread.start()
@@ -212,7 +241,8 @@ class RequestSendThread(val controllerId: Int,
                         val time: Time,
                         val requestRateAndQueueTimeMetrics: Timer,
                         val stateChangeLogger: StateChangeLogger,
-                        name: String)
+                        name: String,
+                        val controllerChannelManager: ControllerChannelManager)
   extends ShutdownableThread(name = name) {
 
   logIdent = s"[RequestSendThread controllerId=$controllerId] "
@@ -224,7 +254,9 @@ class RequestSendThread(val controllerId: Int,
     def backoff(): Unit = pause(100, TimeUnit.MILLISECONDS)
 
     val QueueItem(apiKey, requestBuilder, callback, enqueueTimeMs) = queue.take()
-    requestRateAndQueueTimeMetrics.update(time.milliseconds() - enqueueTimeMs, TimeUnit.MILLISECONDS)
+    var queueTimeMs = time.milliseconds() - enqueueTimeMs
+    var remoteTimeMs: Long = 0
+    requestRateAndQueueTimeMetrics.update(queueTimeMs, TimeUnit.MILLISECONDS)
 
     var clientResponse: ClientResponse = null
     try {
@@ -242,6 +274,7 @@ class RequestSendThread(val controllerId: Int,
               time.milliseconds(), true)
             clientResponse = NetworkClientUtils.sendAndReceive(networkClient, clientRequest, time)
             isSendSuccessful = true
+            remoteTimeMs = time.milliseconds() - enqueueTimeMs - queueTimeMs
           }
         } catch {
           case e: Throwable => // if the send was not successful, reconnect to broker and resend the message
@@ -267,6 +300,7 @@ class RequestSendThread(val controllerId: Int,
         if (callback != null) {
           callback(response)
         }
+        controllerChannelManager.brokerResponseSensors(api).update(queueTimeMs, remoteTimeMs)
       }
     } catch {
       case e: Throwable =>
@@ -303,15 +337,43 @@ class RequestSendThread(val controllerId: Int,
   }
 }
 
-class ControllerBrokerRequestBatch(controller: KafkaController, stateChangeLogger: StateChangeLogger) extends  Logging {
-  val controllerContext = controller.controllerContext
-  val controllerId: Int = controller.config.brokerId
-  val leaderAndIsrRequestMap = mutable.Map.empty[Int, mutable.Map[TopicPartition, LeaderAndIsrRequest.PartitionState]]
-  val stopReplicaRequestMap = mutable.Map.empty[Int, Seq[StopReplicaRequestInfo]]
+class ControllerBrokerRequestBatch(config: KafkaConfig,
+                                   controllerChannelManager: ControllerChannelManager,
+                                   controllerEventManager: ControllerEventManager,
+                                   controllerContext: ControllerContext,
+                                   stateChangeLogger: StateChangeLogger)
+  extends AbstractControllerBrokerRequestBatch(config, controllerContext, stateChangeLogger) {
+
+  def sendEvent(event: ControllerEvent): Unit = {
+    controllerEventManager.put(event)
+  }
+
+  def sendRequest(brokerId: Int,
+                  request: AbstractControlRequest.Builder[_ <: AbstractControlRequest],
+                  callback: AbstractResponse => Unit = null): Unit = {
+    controllerChannelManager.sendRequest(brokerId, request, callback)
+  }
+
+}
+
+case class StopReplicaRequestInfo(replica: PartitionAndReplica, deletePartition: Boolean)
+
+abstract class AbstractControllerBrokerRequestBatch(config: KafkaConfig,
+                                                    controllerContext: ControllerContext,
+                                                    stateChangeLogger: StateChangeLogger) extends  Logging {
+  val controllerId: Int = config.brokerId
+  val leaderAndIsrRequestMap = mutable.Map.empty[Int, mutable.Map[TopicPartition, LeaderAndIsrPartitionState]]
+  val stopReplicaRequestMap = mutable.Map.empty[Int, ListBuffer[StopReplicaRequestInfo]]
   val updateMetadataRequestBrokerSet = mutable.Set.empty[Int]
-  val updateMetadataRequestPartitionInfoMap = mutable.Map.empty[TopicPartition, UpdateMetadataRequest.PartitionState]
+  val updateMetadataRequestPartitionInfoMap = mutable.Map.empty[TopicPartition, UpdateMetadataPartitionState]
+
+  def sendEvent(event: ControllerEvent): Unit
+
+  def sendRequest(brokerId: Int,
+                  request: AbstractControlRequest.Builder[_ <: AbstractControlRequest],
+                  callback: AbstractResponse => Unit = null): Unit
 
-  def newBatch() {
+  def newBatch(): Unit = {
     // raise error if the previous batch is not empty
     if (leaderAndIsrRequestMap.nonEmpty)
       throw new IllegalStateException("Controller to broker state change requests batch is not empty while creating " +
@@ -325,66 +387,72 @@ class ControllerBrokerRequestBatch(controller: KafkaController, stateChangeLogge
         s"$updateMetadataRequestPartitionInfoMap might be lost ")
   }
 
-  def clear() {
+  def clear(): Unit = {
     leaderAndIsrRequestMap.clear()
     stopReplicaRequestMap.clear()
     updateMetadataRequestBrokerSet.clear()
     updateMetadataRequestPartitionInfoMap.clear()
   }
 
-  def addLeaderAndIsrRequestForBrokers(brokerIds: Seq[Int], topicPartition: TopicPartition,
+  def addLeaderAndIsrRequestForBrokers(brokerIds: Seq[Int],
+                                       topicPartition: TopicPartition,
                                        leaderIsrAndControllerEpoch: LeaderIsrAndControllerEpoch,
-                                       replicas: Seq[Int], isNew: Boolean) {
+                                       replicaAssignment: ReplicaAssignment,
+                                       isNew: Boolean): Unit = {
 
     brokerIds.filter(_ >= 0).foreach { brokerId =>
       val result = leaderAndIsrRequestMap.getOrElseUpdate(brokerId, mutable.Map.empty)
       val alreadyNew = result.get(topicPartition).exists(_.isNew)
-      result.put(topicPartition, new LeaderAndIsrRequest.PartitionState(leaderIsrAndControllerEpoch.controllerEpoch,
-        leaderIsrAndControllerEpoch.leaderAndIsr.leader,
-        leaderIsrAndControllerEpoch.leaderAndIsr.leaderEpoch,
-        leaderIsrAndControllerEpoch.leaderAndIsr.isr.map(Integer.valueOf).asJava,
-        leaderIsrAndControllerEpoch.leaderAndIsr.zkVersion,
-        replicas.map(Integer.valueOf).asJava,
-        isNew || alreadyNew))
+      val leaderAndIsr = leaderIsrAndControllerEpoch.leaderAndIsr
+      result.put(topicPartition, new LeaderAndIsrPartitionState()
+        .setTopicName(topicPartition.topic)
+        .setPartitionIndex(topicPartition.partition)
+        .setControllerEpoch(leaderIsrAndControllerEpoch.controllerEpoch)
+        .setLeader(leaderAndIsr.leader)
+        .setLeaderEpoch(leaderAndIsr.leaderEpoch)
+        .setIsr(leaderAndIsr.isr.map(Integer.valueOf).asJava)
+        .setZkVersion(leaderAndIsr.zkVersion)
+        .setReplicas(replicaAssignment.replicas.map(Integer.valueOf).asJava)
+        .setAddingReplicas(replicaAssignment.addingReplicas.map(Integer.valueOf).asJava)
+        .setRemovingReplicas(replicaAssignment.removingReplicas.map(Integer.valueOf).asJava)
+        .setIsNew(isNew || alreadyNew))
     }
 
     addUpdateMetadataRequestForBrokers(controllerContext.liveOrShuttingDownBrokerIds.toSeq, Set(topicPartition))
   }
 
-  def addStopReplicaRequestForBrokers(brokerIds: Seq[Int], topicPartition: TopicPartition, deletePartition: Boolean,
-                                      callback: (AbstractResponse, Int) => Unit) {
+  def addStopReplicaRequestForBrokers(brokerIds: Seq[Int],
+                                      topicPartition: TopicPartition,
+                                      deletePartition: Boolean): Unit = {
     brokerIds.filter(_ >= 0).foreach { brokerId =>
-      stopReplicaRequestMap.getOrElseUpdate(brokerId, Seq.empty[StopReplicaRequestInfo])
-      val v = stopReplicaRequestMap(brokerId)
-      stopReplicaRequestMap(brokerId) = v :+ StopReplicaRequestInfo(PartitionAndReplica(topicPartition, brokerId),
-        deletePartition, (r: AbstractResponse) => callback(r, brokerId))
+      val stopReplicaInfos = stopReplicaRequestMap.getOrElseUpdate(brokerId, ListBuffer.empty[StopReplicaRequestInfo])
+      stopReplicaInfos.append(StopReplicaRequestInfo(PartitionAndReplica(topicPartition, brokerId), deletePartition))
     }
   }
 
   /** Send UpdateMetadataRequest to the given brokers for the given partitions and partitions that are being deleted */
   def addUpdateMetadataRequestForBrokers(brokerIds: Seq[Int],
-                                         partitions: collection.Set[TopicPartition]) {
+                                         partitions: collection.Set[TopicPartition]): Unit = {
 
-    def updateMetadataRequestPartitionInfo(partition: TopicPartition, beingDeleted: Boolean) {
-      val leaderIsrAndControllerEpochOpt = controllerContext.partitionLeadershipInfo.get(partition)
-      leaderIsrAndControllerEpochOpt match {
-        case Some(l @ LeaderIsrAndControllerEpoch(leaderAndIsr, controllerEpoch)) =>
+    def updateMetadataRequestPartitionInfo(partition: TopicPartition, beingDeleted: Boolean): Unit = {
+      controllerContext.partitionLeadershipInfo.get(partition) match {
+        case Some(LeaderIsrAndControllerEpoch(leaderAndIsr, controllerEpoch)) =>
           val replicas = controllerContext.partitionReplicaAssignment(partition)
           val offlineReplicas = replicas.filter(!controllerContext.isReplicaOnline(_, partition))
-          val leaderIsrAndControllerEpoch = if (beingDeleted) {
-            val leaderDuringDelete = LeaderAndIsr.duringDelete(leaderAndIsr.isr)
-            LeaderIsrAndControllerEpoch(leaderDuringDelete, controllerEpoch)
-          } else {
-            l
-          }
-
-          val partitionStateInfo = new UpdateMetadataRequest.PartitionState(leaderIsrAndControllerEpoch.controllerEpoch,
-            leaderIsrAndControllerEpoch.leaderAndIsr.leader,
-            leaderIsrAndControllerEpoch.leaderAndIsr.leaderEpoch,
-            leaderIsrAndControllerEpoch.leaderAndIsr.isr.map(Integer.valueOf).asJava,
-            leaderIsrAndControllerEpoch.leaderAndIsr.zkVersion,
-            replicas.map(Integer.valueOf).asJava,
-            offlineReplicas.map(Integer.valueOf).asJava)
+          val updatedLeaderAndIsr =
+            if (beingDeleted) LeaderAndIsr.duringDelete(leaderAndIsr.isr)
+            else leaderAndIsr
+
+          val partitionStateInfo = new UpdateMetadataPartitionState()
+            .setTopicName(partition.topic)
+            .setPartitionIndex(partition.partition)
+            .setControllerEpoch(controllerEpoch)
+            .setLeader(updatedLeaderAndIsr.leader)
+            .setLeaderEpoch(updatedLeaderAndIsr.leaderEpoch)
+            .setIsr(updatedLeaderAndIsr.isr.map(Integer.valueOf).asJava)
+            .setZkVersion(updatedLeaderAndIsr.zkVersion)
+            .setReplicas(replicas.map(Integer.valueOf).asJava)
+            .setOfflineReplicas(offlineReplicas.map(Integer.valueOf).asJava)
           updateMetadataRequestPartitionInfoMap.put(partition, partitionStateInfo)
 
         case None =>
@@ -394,109 +462,160 @@ class ControllerBrokerRequestBatch(controller: KafkaController, stateChangeLogge
 
     updateMetadataRequestBrokerSet ++= brokerIds.filter(_ >= 0)
     partitions.foreach(partition => updateMetadataRequestPartitionInfo(partition,
-      beingDeleted = controller.topicDeletionManager.topicsToBeDeleted.contains(partition.topic)))
+      beingDeleted = controllerContext.topicsToBeDeleted.contains(partition.topic)))
   }
 
-  def sendRequestsToBrokers(controllerEpoch: Int) {
-    try {
-      val stateChangeLog = stateChangeLogger.withControllerEpoch(controllerEpoch)
-
-      val leaderAndIsrRequestVersion: Short =
-        if (controller.config.interBrokerProtocolVersion >= KAFKA_2_2_IV0) 2
-        else if (controller.config.interBrokerProtocolVersion >= KAFKA_1_0_IV0) 1
-        else 0
-
-      leaderAndIsrRequestMap.filterKeys(controllerContext.liveOrShuttingDownBrokerIds.contains).foreach {
-        case (broker, leaderAndIsrPartitionStates) =>
-          leaderAndIsrPartitionStates.foreach {
-            case (topicPartition, state) =>
-              val typeOfRequest =
-                if (broker == state.basePartitionState.leader) "become-leader"
-                else "become-follower"
-              stateChangeLog.trace(s"Sending $typeOfRequest LeaderAndIsr request $state to broker $broker for partition $topicPartition")
-          }
-          val leaderIds = leaderAndIsrPartitionStates.map(_._2.basePartitionState.leader).toSet
-          val leaders = controllerContext.liveOrShuttingDownBrokers.filter(b => leaderIds.contains(b.id)).map {
-            _.node(controller.config.interBrokerListenerName)
+  private def sendLeaderAndIsrRequest(controllerEpoch: Int, stateChangeLog: StateChangeLogger): Unit = {
+    val leaderAndIsrRequestVersion: Short =
+      if (config.interBrokerProtocolVersion >= KAFKA_2_4_IV1) 5
+      else if (config.interBrokerProtocolVersion >= KAFKA_2_4_IV0) 4
+      else if (config.interBrokerProtocolVersion >= KAFKA_2_3_IV2) 3
+      else if (config.interBrokerProtocolVersion >= KAFKA_2_2_IV0) 2
+      else if (config.interBrokerProtocolVersion >= KAFKA_1_0_IV0) 1
+      else 0
+
+    val maxBrokerEpoch = controllerContext.maxBrokerEpoch
+    leaderAndIsrRequestMap.filterKeys(controllerContext.liveOrShuttingDownBrokerIds.contains).foreach {
+      case (broker, leaderAndIsrPartitionStates) =>
+        if (stateChangeLog.isTraceEnabled) {
+          leaderAndIsrPartitionStates.foreach { case (topicPartition, state) =>
+            val typeOfRequest =
+              if (broker == state.leader) "become-leader"
+              else "become-follower"
+            stateChangeLog.trace(s"Sending $typeOfRequest LeaderAndIsr request $state to broker $broker for partition $topicPartition")
           }
-          val brokerEpoch = controllerContext.liveBrokerIdAndEpochs(broker)
-          val leaderAndIsrRequestBuilder = new LeaderAndIsrRequest.Builder(leaderAndIsrRequestVersion, controllerId, controllerEpoch,
-            brokerEpoch, leaderAndIsrPartitionStates.asJava, leaders.asJava)
-          controller.sendRequest(broker, ApiKeys.LEADER_AND_ISR, leaderAndIsrRequestBuilder,
-            (r: AbstractResponse) => controller.eventManager.put(controller.LeaderAndIsrResponseReceived(r, broker)))
+        }
+        val leaderIds = leaderAndIsrPartitionStates.map(_._2.leader).toSet
+        val leaders = controllerContext.liveOrShuttingDownBrokers.filter(b => leaderIds.contains(b.id)).map {
+          _.node(config.interBrokerListenerName)
+        }
+        val brokerEpoch = controllerContext.liveBrokerIdAndEpochs(broker)
 
-      }
-      leaderAndIsrRequestMap.clear()
+        val leaderAndIsrRequestBuilder = new LeaderAndIsrRequest.Builder(leaderAndIsrRequestVersion, controllerId, controllerEpoch,
+          brokerEpoch, maxBrokerEpoch, leaderAndIsrPartitionStates.values.toBuffer.asJava, leaders.asJava)
+        sendRequest(broker, leaderAndIsrRequestBuilder, (r: AbstractResponse) => sendEvent(LeaderAndIsrResponseReceived(r, broker)))
 
-      updateMetadataRequestPartitionInfoMap.foreach { case (tp, partitionState) =>
-        stateChangeLog.trace(s"Sending UpdateMetadata request $partitionState to brokers $updateMetadataRequestBrokerSet " +
-          s"for partition $tp")
-      }
+    }
+    leaderAndIsrRequestMap.clear()
+  }
 
-      val partitionStates = Map.empty ++ updateMetadataRequestPartitionInfoMap
-      val updateMetadataRequestVersion: Short =
-        if (controller.config.interBrokerProtocolVersion >= KAFKA_2_2_IV0) 5
-        else if (controller.config.interBrokerProtocolVersion >= KAFKA_1_0_IV0) 4
-        else if (controller.config.interBrokerProtocolVersion >= KAFKA_0_10_2_IV0) 3
-        else if (controller.config.interBrokerProtocolVersion >= KAFKA_0_10_0_IV1) 2
-        else if (controller.config.interBrokerProtocolVersion >= KAFKA_0_9_0) 1
-        else 0
-
-      val liveBrokers = if (updateMetadataRequestVersion == 0) {
-        // Version 0 of UpdateMetadataRequest only supports PLAINTEXT.
-        controllerContext.liveOrShuttingDownBrokers.map { broker =>
-          val securityProtocol = SecurityProtocol.PLAINTEXT
-          val listenerName = ListenerName.forSecurityProtocol(securityProtocol)
-          val node = broker.node(listenerName)
-          val endPoints = Seq(new EndPoint(node.host, node.port, securityProtocol, listenerName))
-          new UpdateMetadataRequest.Broker(broker.id, endPoints.asJava, broker.rack.orNull)
-        }
+  private def sendUpdateMetadataRequests(controllerEpoch: Int, stateChangeLog: StateChangeLogger): Unit = {
+    updateMetadataRequestPartitionInfoMap.foreach { case (tp, partitionState) =>
+      stateChangeLog.trace(s"Sending UpdateMetadata request $partitionState to brokers $updateMetadataRequestBrokerSet " +
+        s"for partition $tp")
+    }
+
+    val partitionStates = updateMetadataRequestPartitionInfoMap.values.toBuffer
+    val updateMetadataRequestVersion: Short =
+      if (config.interBrokerProtocolVersion >= KAFKA_2_4_IV1) 7
+      else if (config.interBrokerProtocolVersion >= KAFKA_2_3_IV2) 6
+      else if (config.interBrokerProtocolVersion >= KAFKA_2_2_IV0) 5
+      else if (config.interBrokerProtocolVersion >= KAFKA_1_0_IV0) 4
+      else if (config.interBrokerProtocolVersion >= KAFKA_0_10_2_IV0) 3
+      else if (config.interBrokerProtocolVersion >= KAFKA_0_10_0_IV1) 2
+      else if (config.interBrokerProtocolVersion >= KAFKA_0_9_0) 1
+      else 0
+
+    val liveBrokers = controllerContext.liveOrShuttingDownBrokers.iterator.map { broker =>
+      val endpoints = if (updateMetadataRequestVersion == 0) {
+        // Version 0 of UpdateMetadataRequest only supports PLAINTEXT
+        val securityProtocol = SecurityProtocol.PLAINTEXT
+        val listenerName = ListenerName.forSecurityProtocol(securityProtocol)
+        val node = broker.node(listenerName)
+        Seq(new UpdateMetadataEndpoint()
+          .setHost(node.host)
+          .setPort(node.port)
+          .setSecurityProtocol(securityProtocol.id)
+          .setListener(listenerName.value))
       } else {
-        controllerContext.liveOrShuttingDownBrokers.map { broker =>
-          val endPoints = broker.endPoints.map { endPoint =>
-            new UpdateMetadataRequest.EndPoint(endPoint.host, endPoint.port, endPoint.securityProtocol, endPoint.listenerName)
-          }
-          new UpdateMetadataRequest.Broker(broker.id, endPoints.asJava, broker.rack.orNull)
+        broker.endPoints.map { endpoint =>
+          new UpdateMetadataEndpoint()
+            .setHost(endpoint.host)
+            .setPort(endpoint.port)
+            .setSecurityProtocol(endpoint.securityProtocol.id)
+            .setListener(endpoint.listenerName.value)
         }
       }
+      new UpdateMetadataBroker()
+        .setId(broker.id)
+        .setEndpoints(endpoints.asJava)
+        .setRack(broker.rack.orNull)
+    }.toBuffer
+
+    if (updateMetadataRequestVersion >= 6) {
+      // We should only create one copy UpdateMetadataRequest that should apply to all brokers.
+      // The goal is to reduce memory footprint on the controller.
+      val maxBrokerEpoch = controllerContext.maxBrokerEpoch
+      val updateMetadataRequest = new UpdateMetadataRequest.Builder(updateMetadataRequestVersion, controllerId, controllerEpoch,
+        AbstractControlRequest.UNKNOWN_BROKER_EPOCH, maxBrokerEpoch, partitionStates.asJava, liveBrokers.asJava)
 
+      updateMetadataRequestBrokerSet.intersect(controllerContext.liveOrShuttingDownBrokerIds).foreach { broker =>
+        sendRequest(broker, updateMetadataRequest)
+      }
+    } else {
       updateMetadataRequestBrokerSet.intersect(controllerContext.liveOrShuttingDownBrokerIds).foreach { broker =>
         val brokerEpoch = controllerContext.liveBrokerIdAndEpochs(broker)
         val updateMetadataRequest = new UpdateMetadataRequest.Builder(updateMetadataRequestVersion, controllerId, controllerEpoch,
-          brokerEpoch, partitionStates.asJava, liveBrokers.asJava)
-        controller.sendRequest(broker, ApiKeys.UPDATE_METADATA, updateMetadataRequest, null)
+          brokerEpoch, AbstractControlRequest.UNKNOWN_BROKER_EPOCH, partitionStates.asJava, liveBrokers.asJava)
+        sendRequest(broker, updateMetadataRequest)
       }
-      updateMetadataRequestBrokerSet.clear()
-      updateMetadataRequestPartitionInfoMap.clear()
+    }
 
-      val stopReplicaRequestVersion: Short =
-        if (controller.config.interBrokerProtocolVersion >= KAFKA_2_2_IV0) 1
-        else 0
 
-      stopReplicaRequestMap.filterKeys(controllerContext.liveOrShuttingDownBrokerIds.contains).foreach { case (broker, replicaInfoList) =>
-        val stopReplicaWithDelete = replicaInfoList.filter(_.deletePartition).map(_.replica).toSet
-        val stopReplicaWithoutDelete = replicaInfoList.filterNot(_.deletePartition).map(_.replica).toSet
-        debug(s"The stop replica request (delete = true) sent to broker $broker is ${stopReplicaWithDelete.mkString(",")}")
-        debug(s"The stop replica request (delete = false) sent to broker $broker is ${stopReplicaWithoutDelete.mkString(",")}")
+    updateMetadataRequestBrokerSet.clear()
+    updateMetadataRequestPartitionInfoMap.clear()
+  }
 
-        val (replicasToGroup, replicasToNotGroup) = replicaInfoList.partition(r => !r.deletePartition && r.callback == null)
-        val brokerEpoch = controllerContext.liveBrokerIdAndEpochs(broker)
+  private def sendStopReplicaRequests(controllerEpoch: Int): Unit = {
+    val stopReplicaRequestVersion: Short =
+      if (config.interBrokerProtocolVersion >= KAFKA_2_4_IV1) 3
+      else if (config.interBrokerProtocolVersion >= KAFKA_2_3_IV2) 2
+      else if (config.interBrokerProtocolVersion >= KAFKA_2_2_IV0) 1
+      else 0
+
+    def stopReplicaPartitionDeleteResponseCallback(brokerId: Int)(response: AbstractResponse): Unit = {
+      val stopReplicaResponse = response.asInstanceOf[StopReplicaResponse]
+      val partitionErrorsForDeletingTopics = stopReplicaResponse.partitionErrors.asScala.iterator.filter { pe =>
+        controllerContext.isTopicDeletionInProgress(pe.topicName)
+      }.map(pe => new TopicPartition(pe.topicName, pe.partitionIndex) -> Errors.forCode(pe.errorCode)).toMap
+
+      if (partitionErrorsForDeletingTopics.nonEmpty)
+        sendEvent(TopicDeletionStopReplicaResponseReceived(brokerId, stopReplicaResponse.error, partitionErrorsForDeletingTopics))
+    }
 
-        // Send one StopReplicaRequest for all partitions that require neither delete nor callback. This potentially
-        // changes the order in which the requests are sent for the same partitions, but that's OK.
-        val stopReplicaRequest = new StopReplicaRequest.Builder(stopReplicaRequestVersion, controllerId, controllerEpoch,
-          brokerEpoch, false,
-          replicasToGroup.map(_.replica.topicPartition).toSet.asJava)
-        controller.sendRequest(broker, ApiKeys.STOP_REPLICA, stopReplicaRequest)
-
-        replicasToNotGroup.foreach { r =>
-          val stopReplicaRequest = new StopReplicaRequest.Builder(stopReplicaRequestVersion,
-            controllerId, controllerEpoch, brokerEpoch, r.deletePartition,
-            Set(r.replica.topicPartition).asJava)
-          controller.sendRequest(broker, ApiKeys.STOP_REPLICA, stopReplicaRequest, r.callback)
-        }
+    def createStopReplicaRequest(brokerEpoch: Long, maxBrokerEpoch: Long, requests: Seq[StopReplicaRequestInfo], deletePartitions: Boolean): StopReplicaRequest.Builder = {
+      val partitions = requests.map(_.replica.topicPartition).asJava
+      new StopReplicaRequest.Builder(stopReplicaRequestVersion, controllerId, controllerEpoch,
+        brokerEpoch, maxBrokerEpoch, deletePartitions, partitions)
+    }
+
+    val maxBrokerEpoch = controllerContext.maxBrokerEpoch
+    stopReplicaRequestMap.filterKeys(controllerContext.liveOrShuttingDownBrokerIds.contains).foreach { case (brokerId, replicaInfoList) =>
+      val (stopReplicaWithDelete, stopReplicaWithoutDelete) = replicaInfoList.partition(r => r.deletePartition)
+      val brokerEpoch = controllerContext.liveBrokerIdAndEpochs(brokerId)
+
+      if (stopReplicaWithDelete.nonEmpty) {
+        debug(s"The stop replica request (delete = true) sent to broker $brokerId is ${stopReplicaWithDelete.mkString(",")}")
+        val stopReplicaRequest = createStopReplicaRequest(brokerEpoch, maxBrokerEpoch, stopReplicaWithDelete, deletePartitions = true)
+        val callback = stopReplicaPartitionDeleteResponseCallback(brokerId) _
+        sendRequest(brokerId, stopReplicaRequest, callback)
+      }
+
+      if (stopReplicaWithoutDelete.nonEmpty) {
+        debug(s"The stop replica request (delete = false) sent to broker $brokerId is ${stopReplicaWithoutDelete.mkString(",")}")
+        val stopReplicaRequest = createStopReplicaRequest(brokerEpoch, maxBrokerEpoch, stopReplicaWithoutDelete, deletePartitions = false)
+        sendRequest(brokerId, stopReplicaRequest)
       }
-      stopReplicaRequestMap.clear()
+    }
+    stopReplicaRequestMap.clear()
+  }
+
+  def sendRequestsToBrokers(controllerEpoch: Int): Unit = {
+    try {
+      val stateChangeLog = stateChangeLogger.withControllerEpoch(controllerEpoch)
+      sendLeaderAndIsrRequest(controllerEpoch, stateChangeLog)
+      sendUpdateMetadataRequests(controllerEpoch, stateChangeLog)
+      sendStopReplicaRequests(controllerEpoch)
     } catch {
       case e: Throwable =>
         if (leaderAndIsrRequestMap.nonEmpty) {
@@ -521,8 +640,25 @@ case class ControllerBrokerStateInfo(networkClient: NetworkClient,
                                      messageQueue: BlockingQueue[QueueItem],
                                      requestSendThread: RequestSendThread,
                                      queueSizeGauge: Gauge[Int],
-                                     requestRateAndTimeMetrics: Timer)
+                                     requestRateAndTimeMetrics: Timer,
+                                     reconfigurableChannelBuilder: Option[Reconfigurable])
+
+
+class BrokerResponseTimeStats(val key: ApiKeys) extends KafkaMetricsGroup {
+  // Records time for request waits on local send thread queue
+  val brokerRequestQueueTime = newHistogram("brokerRequestQueueTimeMs", true, responseTimeTags)
+  // Records time for controller to send request and receive response
+  val brokerRequestRemoteTime = newHistogram("brokerRequestRemoteTimeMs", true, responseTimeTags)
 
-case class StopReplicaRequestInfo(replica: PartitionAndReplica, deletePartition: Boolean, callback: AbstractResponse => Unit)
+  def responseTimeTags = Map("request" -> key.toString)
 
-class Callbacks(val stopReplicaResponseCallback: (AbstractResponse, Int) => Unit = (_, _ ) => ())
+  def update(queueTime: Long, remoteTime: Long): Unit = {
+    brokerRequestQueueTime.update(queueTime)
+    brokerRequestRemoteTime.update(remoteTime)
+  }
+
+  def removeMetrics(): Unit = {
+    removeMetric("brokerRequestQueueTimeMs", responseTimeTags)
+    removeMetric("brokerRequestRemoteTimeMs", responseTimeTags)
+  }
+}
diff --git a/core/src/main/scala/kafka/controller/ControllerContext.scala b/core/src/main/scala/kafka/controller/ControllerContext.scala
index c3bcc5282d12d..3e1a5a9769e0c 100644
--- a/core/src/main/scala/kafka/controller/ControllerContext.scala
+++ b/core/src/main/scala/kafka/controller/ControllerContext.scala
@@ -20,91 +20,197 @@ package kafka.controller
 import kafka.cluster.Broker
 import org.apache.kafka.common.TopicPartition
 
-import scala.collection.{Seq, Set, mutable}
+import scala.collection.{Map, Seq, Set, mutable}
 
-class ControllerContext {
-  val stats = new ControllerStats
+object ReplicaAssignment {
+  def fromOldAndNewReplicas(oldReplicas: Seq[Int], newReplicas: Seq[Int]): ReplicaAssignment = {
+    val fullReplicaSet = (newReplicas ++ oldReplicas).distinct
+    ReplicaAssignment(
+      fullReplicaSet,
+      fullReplicaSet.diff(oldReplicas),
+      fullReplicaSet.diff(newReplicas)
+    )
+  }
+
+  def apply(replicas: Seq[Int]): ReplicaAssignment = {
+    apply(replicas, Seq.empty, Seq.empty)
+  }
+}
+
+case class ReplicaAssignment(replicas: Seq[Int],
+                             addingReplicas: Seq[Int],
+                             removingReplicas: Seq[Int]) {
+
+  lazy val originReplicas: Seq[Int] = replicas.diff(addingReplicas)
+  lazy val targetReplicas: Seq[Int] = replicas.diff(removingReplicas)
+
+  def isBeingReassigned: Boolean = {
+    addingReplicas.nonEmpty || removingReplicas.nonEmpty
+  }
+
+  def reassignTo(newReplicas: Seq[Int]): ReplicaAssignment = {
+    ReplicaAssignment.fromOldAndNewReplicas(originReplicas, newReplicas)
+  }
 
-  var controllerChannelManager: ControllerChannelManager = null
+  override def toString: String = s"ReplicaAssignment(" +
+    s"replicas=${replicas.mkString(",")}, " +
+    s"addingReplicas=${addingReplicas.mkString(",")}, " +
+    s"removingReplicas=${removingReplicas.mkString(",")})"
+}
 
+class ControllerContext {
+  val stats = new ControllerStats
+  var offlinePartitionCount = 0
   var shuttingDownBrokerIds: mutable.Set[Int] = mutable.Set.empty
+  private var liveBrokers: Set[Broker] = Set.empty
+  private var liveBrokerEpochs: Map[Int, Long] = Map.empty
   var epoch: Int = KafkaController.InitialControllerEpoch
   var epochZkVersion: Int = KafkaController.InitialControllerEpochZkVersion
+
   var allTopics: Set[String] = Set.empty
-  private val partitionReplicaAssignmentUnderlying: mutable.Map[String, mutable.Map[Int, Seq[Int]]] = mutable.Map.empty
-  val partitionLeadershipInfo: mutable.Map[TopicPartition, LeaderIsrAndControllerEpoch] = mutable.Map.empty
-  val partitionsBeingReassigned: mutable.Map[TopicPartition, ReassignedPartitionsContext] = mutable.Map.empty
+  val partitionAssignments = mutable.Map.empty[String, mutable.Map[Int, ReplicaAssignment]]
+  val partitionLeadershipInfo = mutable.Map.empty[TopicPartition, LeaderIsrAndControllerEpoch]
+  val partitionsBeingReassigned = mutable.Set.empty[TopicPartition]
+  val partitionStates = mutable.Map.empty[TopicPartition, PartitionState]
+  val replicaStates = mutable.Map.empty[PartitionAndReplica, ReplicaState]
   val replicasOnOfflineDirs: mutable.Map[Int, Set[TopicPartition]] = mutable.Map.empty
 
-  private var liveBrokersUnderlying: Set[Broker] = Set.empty
-  private var liveBrokerIdAndEpochsUnderlying: Map[Int, Long] = Map.empty
+  val topicsToBeDeleted = mutable.Set.empty[String]
 
-  def partitionReplicaAssignment(topicPartition: TopicPartition): Seq[Int] = {
-    partitionReplicaAssignmentUnderlying.getOrElse(topicPartition.topic, mutable.Map.empty)
-      .getOrElse(topicPartition.partition, Seq.empty)
-  }
+  /** The following topicsWithDeletionStarted variable is used to properly update the offlinePartitionCount metric.
+   * When a topic is going through deletion, we don't want to keep track of its partition state
+   * changes in the offlinePartitionCount metric. This goal means if some partitions of a topic are already
+   * in OfflinePartition state when deletion starts, we need to change the corresponding partition
+   * states to NonExistentPartition first before starting the deletion.
+   *
+   * However we can NOT change partition states to NonExistentPartition at the time of enqueuing topics
+   * for deletion. The reason is that when a topic is enqueued for deletion, it may be ineligible for
+   * deletion due to ongoing partition reassignments. Hence there might be a delay between enqueuing
+   * a topic for deletion and the actual start of deletion. In this delayed interval, partitions may still
+   * transition to or out of the OfflinePartition state.
+   *
+   * Hence we decide to change partition states to NonExistentPartition only when the actual deletion have started.
+   * For topics whose deletion have actually started, we keep track of them in the following topicsWithDeletionStarted
+   * variable. And once a topic is in the topicsWithDeletionStarted set, we are sure there will no longer
+   * be partition reassignments to any of its partitions, and only then it's safe to move its partitions to
+   * NonExistentPartition state. Once a topic is in the topicsWithDeletionStarted set, we will stop monitoring
+   * its partition state changes in the offlinePartitionCount metric
+   */
+  val topicsWithDeletionStarted = mutable.Set.empty[String]
+  val topicsIneligibleForDeletion = mutable.Set.empty[String]
+
+  @volatile var livePreferredControllerIds: Set[Int] = Set.empty
 
   private def clearTopicsState(): Unit = {
     allTopics = Set.empty
-    partitionReplicaAssignmentUnderlying.clear()
+    partitionAssignments.clear()
     partitionLeadershipInfo.clear()
     partitionsBeingReassigned.clear()
     replicasOnOfflineDirs.clear()
+    partitionStates.clear()
+    offlinePartitionCount = 0
+    replicaStates.clear()
+  }
+
+  def partitionReplicaAssignment(topicPartition: TopicPartition): Seq[Int] = {
+    partitionAssignments.getOrElse(topicPartition.topic, mutable.Map.empty)
+      .get(topicPartition.partition) match {
+        case Some(partitionAssignment) => partitionAssignment.replicas
+        case None => Seq.empty
+      }
+  }
+
+  def partitionFullReplicaAssignment(topicPartition: TopicPartition): ReplicaAssignment = {
+    partitionAssignments.getOrElse(topicPartition.topic, mutable.Map.empty).get(topicPartition.partition) match {
+      case Some(partitionAssignment) => partitionAssignment
+      case None => ReplicaAssignment(Seq(), Seq(), Seq())
+    }
   }
 
   def updatePartitionReplicaAssignment(topicPartition: TopicPartition, newReplicas: Seq[Int]): Unit = {
-    partitionReplicaAssignmentUnderlying.getOrElseUpdate(topicPartition.topic, mutable.Map.empty)
-      .put(topicPartition.partition, newReplicas)
+    val assignments = partitionAssignments.getOrElseUpdate(topicPartition.topic, mutable.Map.empty)
+    val newAssignment = assignments.get(topicPartition.partition) match {
+      case Some(partitionAssignment) =>
+        ReplicaAssignment(
+          newReplicas,
+          partitionAssignment.addingReplicas,
+          partitionAssignment.removingReplicas
+        )
+      case None =>
+        ReplicaAssignment(
+          newReplicas,
+          Seq.empty,
+          Seq.empty
+        )
+    }
+    updatePartitionFullReplicaAssignment(topicPartition, newAssignment)
+  }
+
+  def updatePartitionFullReplicaAssignment(topicPartition: TopicPartition, newAssignment: ReplicaAssignment): Unit = {
+    val assignments = partitionAssignments.getOrElseUpdate(topicPartition.topic, mutable.Map.empty)
+    assignments.put(topicPartition.partition, newAssignment)
   }
 
   def partitionReplicaAssignmentForTopic(topic : String): Map[TopicPartition, Seq[Int]] = {
-    partitionReplicaAssignmentUnderlying.getOrElse(topic, Map.empty).map {
-      case (partition, replicas) => (new TopicPartition(topic, partition), replicas)
+    partitionAssignments.getOrElse(topic, Map.empty).map {
+      case (partition, assignment) => (new TopicPartition(topic, partition), assignment.replicas)
+    }.toMap
+  }
+
+  def partitionFullReplicaAssignmentForTopic(topic : String): Map[TopicPartition, ReplicaAssignment] = {
+    partitionAssignments.getOrElse(topic, Map.empty).map {
+      case (partition, assignment) => (new TopicPartition(topic, partition), assignment)
     }.toMap
   }
 
   def allPartitions: Set[TopicPartition] = {
-    partitionReplicaAssignmentUnderlying.flatMap {
+    partitionAssignments.flatMap {
       case (topic, topicReplicaAssignment) => topicReplicaAssignment.map {
         case (partition, _) => new TopicPartition(topic, partition)
       }
     }.toSet
   }
 
-  def setLiveBrokerAndEpochs(brokerAndEpochs: Map[Broker, Long]) {
-    liveBrokersUnderlying = brokerAndEpochs.keySet
-    liveBrokerIdAndEpochsUnderlying =
+  def setLiveBrokerAndEpochs(brokerAndEpochs: Map[Broker, Long]): Unit = {
+    liveBrokers = brokerAndEpochs.keySet
+    liveBrokerEpochs =
       brokerAndEpochs map { case (broker, brokerEpoch) => (broker.id, brokerEpoch)}
   }
 
   def addLiveBrokersAndEpochs(brokerAndEpochs: Map[Broker, Long]): Unit = {
-    liveBrokersUnderlying = liveBrokersUnderlying ++ brokerAndEpochs.keySet
-    liveBrokerIdAndEpochsUnderlying = liveBrokerIdAndEpochsUnderlying ++
+    liveBrokers = liveBrokers ++ brokerAndEpochs.keySet
+    liveBrokerEpochs = liveBrokerEpochs ++
       (brokerAndEpochs map { case (broker, brokerEpoch) => (broker.id, brokerEpoch)})
   }
 
-  def removeLiveBrokersAndEpochs(brokerIds : Set[Int]): Unit = {
-    liveBrokersUnderlying = liveBrokersUnderlying.filter(broker => !brokerIds.contains(broker.id))
-    liveBrokerIdAndEpochsUnderlying = liveBrokerIdAndEpochsUnderlying.filterKeys(id => !brokerIds.contains(id))
+  def removeLiveBrokers(brokerIds: Set[Int]): Unit = {
+    liveBrokers = liveBrokers.filter(broker => !brokerIds.contains(broker.id))
+    liveBrokerEpochs = liveBrokerEpochs.filter { case (id, _) => !brokerIds.contains(id) }
   }
 
-  def updateBrokerMetadata(oldMetadata: Option[Broker], newMetadata: Option[Broker]): Unit = {
-    liveBrokersUnderlying = liveBrokersUnderlying -- oldMetadata ++ newMetadata
+  def updateBrokerMetadata(oldMetadata: Broker, newMetadata: Broker): Unit = {
+    liveBrokers -= oldMetadata
+    liveBrokers += newMetadata
   }
 
-  // getter
-  def liveBrokers = liveBrokersUnderlying.filter(broker => !shuttingDownBrokerIds.contains(broker.id))
-  def liveBrokerIds = liveBrokerIdAndEpochsUnderlying.keySet -- shuttingDownBrokerIds
+  def setLivePreferredControllerIds(preferredControllerIds: Set[Int]): Unit = {
+    livePreferredControllerIds = preferredControllerIds
+  }
 
-  def liveOrShuttingDownBrokerIds = liveBrokerIdAndEpochsUnderlying.keySet
-  def liveOrShuttingDownBrokers = liveBrokersUnderlying
+  // getter
+  def liveBrokerIds: Set[Int] = liveBrokerEpochs.keySet -- shuttingDownBrokerIds
+  def liveOrShuttingDownBrokerIds: Set[Int] = liveBrokerEpochs.keySet
+  def liveOrShuttingDownBrokers: Set[Broker] = liveBrokers
+  def liveBrokerIdAndEpochs: Map[Int, Long] = liveBrokerEpochs
+  def maxBrokerEpoch: Long = liveBrokerEpochs.values.max
+  def liveOrShuttingDownBroker(brokerId: Int): Option[Broker] = liveOrShuttingDownBrokers.find(_.id == brokerId)
 
-  def liveBrokerIdAndEpochs = liveBrokerIdAndEpochsUnderlying
+  def getLivePreferredControllerIds : Set[Int] = livePreferredControllerIds
 
   def partitionsOnBroker(brokerId: Int): Set[TopicPartition] = {
-    partitionReplicaAssignmentUnderlying.flatMap {
+    partitionAssignments.flatMap {
       case (topic, topicReplicaAssignment) => topicReplicaAssignment.filter {
-        case (_, replicas) => replicas.contains(brokerId)
+        case (_, partitionAssignment) => partitionAssignment.replicas.contains(brokerId)
       }.map {
         case (partition, _) => new TopicPartition(topic, partition)
       }
@@ -121,9 +227,9 @@ class ControllerContext {
 
   def replicasOnBrokers(brokerIds: Set[Int]): Set[PartitionAndReplica] = {
     brokerIds.flatMap { brokerId =>
-      partitionReplicaAssignmentUnderlying.flatMap {
+      partitionAssignments.flatMap {
         case (topic, topicReplicaAssignment) => topicReplicaAssignment.collect {
-          case (partition, replicas)  if replicas.contains(brokerId) =>
+          case (partition, partitionAssignment) if partitionAssignment.replicas.contains(brokerId) =>
             PartitionAndReplica(new TopicPartition(topic, partition), brokerId)
         }
       }
@@ -131,13 +237,13 @@ class ControllerContext {
   }
 
   def replicasForTopic(topic: String): Set[PartitionAndReplica] = {
-    partitionReplicaAssignmentUnderlying.getOrElse(topic, mutable.Map.empty).flatMap {
-      case (partition, replicas) => replicas.map(r => PartitionAndReplica(new TopicPartition(topic, partition), r))
+    partitionAssignments.getOrElse(topic, mutable.Map.empty).flatMap {
+      case (partition, assignment) => assignment.replicas.map(r => PartitionAndReplica(new TopicPartition(topic, partition), r))
     }.toSet
   }
 
   def partitionsForTopic(topic: String): collection.Set[TopicPartition] = {
-    partitionReplicaAssignmentUnderlying.getOrElse(topic, mutable.Map.empty).map {
+    partitionAssignments.getOrElse(topic, mutable.Map.empty).map {
       case (partition, _) => new TopicPartition(topic, partition)
     }.toSet
   }
@@ -148,6 +254,28 @@ class ControllerContext {
     }
   }
 
+  /**
+    * Get all online and offline replicas.
+    *
+    * @return a tuple consisting of first the online replicas and followed by the offline replicas
+    */
+  def onlineAndOfflineReplicas: (Set[PartitionAndReplica], Set[PartitionAndReplica]) = {
+    val onlineReplicas = mutable.Set.empty[PartitionAndReplica]
+    val offlineReplicas = mutable.Set.empty[PartitionAndReplica]
+    for ((topic, partitionAssignments) <- partitionAssignments;
+         (partitionId, assignment) <- partitionAssignments) {
+      val partition = new TopicPartition(topic, partitionId)
+      for (replica <- assignment.replicas) {
+        val partitionAndReplica = PartitionAndReplica(partition, replica)
+        if (isReplicaOnline(replica, partition))
+          onlineReplicas.add(partitionAndReplica)
+        else
+          offlineReplicas.add(partitionAndReplica)
+      }
+    }
+    (onlineReplicas, offlineReplicas)
+  }
+
   def replicasForPartition(partitions: collection.Set[TopicPartition]): collection.Set[PartitionAndReplica] = {
     partitions.flatMap { p =>
       val replicas = partitionReplicaAssignment(p)
@@ -156,10 +284,9 @@ class ControllerContext {
   }
 
   def resetContext(): Unit = {
-    if (controllerChannelManager != null) {
-      controllerChannelManager.shutdown()
-      controllerChannelManager = null
-    }
+    topicsToBeDeleted.clear()
+    topicsWithDeletionStarted.clear()
+    topicsIneligibleForDeletion.clear()
     shuttingDownBrokerIds.clear()
     epoch = 0
     epochZkVersion = 0
@@ -169,10 +296,130 @@ class ControllerContext {
 
   def removeTopic(topic: String): Unit = {
     allTopics -= topic
-    partitionReplicaAssignmentUnderlying.remove(topic)
+    partitionAssignments.remove(topic)
+    partitionStates.foreach {
+      case (topicPartition, _) if topicPartition.topic == topic => partitionStates.remove(topicPartition)
+      case _ =>
+    }
     partitionLeadershipInfo.foreach {
       case (topicPartition, _) if topicPartition.topic == topic => partitionLeadershipInfo.remove(topicPartition)
       case _ =>
     }
   }
+
+  def queueTopicDeletion(topics: Set[String]): Unit = {
+    topicsToBeDeleted ++= topics
+  }
+
+  def beginTopicDeletion(topics: Set[String]): Unit = {
+    topicsWithDeletionStarted ++= topics
+  }
+
+  def isTopicDeletionInProgress(topic: String): Boolean = {
+    topicsWithDeletionStarted.contains(topic)
+  }
+
+  def isTopicQueuedUpForDeletion(topic: String): Boolean = {
+    topicsToBeDeleted.contains(topic)
+  }
+
+  def isTopicEligibleForDeletion(topic: String): Boolean = {
+    topicsToBeDeleted.contains(topic) && !topicsIneligibleForDeletion.contains(topic)
+  }
+
+  def topicsQueuedForDeletion: Set[String] = {
+    topicsToBeDeleted
+  }
+
+  def replicasInState(topic: String, state: ReplicaState): Set[PartitionAndReplica] = {
+    replicasForTopic(topic).filter(replica => replicaStates(replica) == state).toSet
+  }
+
+  def areAllReplicasInState(topic: String, state: ReplicaState): Boolean = {
+    replicasForTopic(topic).forall(replica => replicaStates(replica) == state)
+  }
+
+  def isAnyReplicaInState(topic: String, state: ReplicaState): Boolean = {
+    replicasForTopic(topic).exists(replica => replicaStates(replica) == state)
+  }
+
+  def checkValidReplicaStateChange(replicas: Seq[PartitionAndReplica], targetState: ReplicaState): (Seq[PartitionAndReplica], Seq[PartitionAndReplica]) = {
+    replicas.partition(replica => isValidReplicaStateTransition(replica, targetState))
+  }
+
+  def checkValidPartitionStateChange(partitions: Seq[TopicPartition], targetState: PartitionState): (Seq[TopicPartition], Seq[TopicPartition]) = {
+    partitions.partition(p => isValidPartitionStateTransition(p, targetState))
+  }
+
+  def putReplicaState(replica: PartitionAndReplica, state: ReplicaState): Unit = {
+    replicaStates.put(replica, state)
+  }
+
+  def removeReplicaState(replica: PartitionAndReplica): Unit = {
+    replicaStates.remove(replica)
+  }
+
+  def putReplicaStateIfNotExists(replica: PartitionAndReplica, state: ReplicaState): Unit = {
+    replicaStates.getOrElseUpdate(replica, state)
+  }
+
+  def putPartitionState(partition: TopicPartition, targetState: PartitionState): Unit = {
+    val currentState = partitionStates.put(partition, targetState).getOrElse(NonExistentPartition)
+    updatePartitionStateMetrics(partition, currentState, targetState)
+  }
+
+  def excludeDeletingTopicFromOfflinePartitionCount(topic: String): Unit = {
+    if (isTopicQueuedUpForDeletion(topic)) {
+      offlinePartitionCount = offlinePartitionCount -
+        partitionsForTopic(topic).count(partition => partitionState(partition) == OfflinePartition)
+    }
+  }
+
+  private def updatePartitionStateMetrics(partition: TopicPartition,
+                                          currentState: PartitionState,
+                                          targetState: PartitionState): Unit = {
+    if (!isTopicQueuedUpForDeletion(partition.topic)) {
+      if (currentState != OfflinePartition && targetState == OfflinePartition) {
+        offlinePartitionCount = offlinePartitionCount + 1
+      } else if (currentState == OfflinePartition && targetState != OfflinePartition) {
+        offlinePartitionCount = offlinePartitionCount - 1
+      }
+    }
+  }
+
+  def putPartitionStateIfNotExists(partition: TopicPartition, state: PartitionState): Unit = {
+    if (partitionStates.getOrElseUpdate(partition, state) == state)
+      updatePartitionStateMetrics(partition, NonExistentPartition, state)
+  }
+
+  def replicaState(replica: PartitionAndReplica): ReplicaState = {
+    replicaStates(replica)
+  }
+
+  def partitionState(partition: TopicPartition): PartitionState = {
+    partitionStates(partition)
+  }
+
+  def partitionsInState(state: PartitionState): Set[TopicPartition] = {
+    partitionStates.filter { case (_, s) => s == state }.keySet.toSet
+  }
+
+  def partitionsInStates(states: Set[PartitionState]): Set[TopicPartition] = {
+    partitionStates.filter { case (_, s) => states.contains(s) }.keySet.toSet
+  }
+
+  def partitionsInState(topic: String, state: PartitionState): Set[TopicPartition] = {
+    partitionsForTopic(topic).filter { partition => state == partitionState(partition) }.toSet
+  }
+
+  def partitionsInStates(topic: String, states: Set[PartitionState]): Set[TopicPartition] = {
+    partitionsForTopic(topic).filter { partition => states.contains(partitionState(partition)) }.toSet
+  }
+
+  private def isValidReplicaStateTransition(replica: PartitionAndReplica, targetState: ReplicaState): Boolean =
+    targetState.validPreviousStates.contains(replicaStates(replica))
+
+  private def isValidPartitionStateTransition(partition: TopicPartition, targetState: PartitionState): Boolean =
+    targetState.validPreviousStates.contains(partitionStates(partition))
+
 }
diff --git a/core/src/main/scala/kafka/controller/ControllerEventManager.scala b/core/src/main/scala/kafka/controller/ControllerEventManager.scala
index a456ce32895a9..1ce594d2d26b6 100644
--- a/core/src/main/scala/kafka/controller/ControllerEventManager.scala
+++ b/core/src/main/scala/kafka/controller/ControllerEventManager.scala
@@ -17,14 +17,14 @@
 
 package kafka.controller
 
-import java.util.concurrent.LinkedBlockingQueue
+import java.util.concurrent.atomic.AtomicBoolean
+import java.util.concurrent.{CountDownLatch, LinkedBlockingQueue}
 import java.util.concurrent.locks.ReentrantLock
 
 import com.yammer.metrics.core.Gauge
 import kafka.metrics.{KafkaMetricsGroup, KafkaTimer}
 import kafka.utils.CoreUtils.inLock
 import kafka.utils.ShutdownableThread
-import org.apache.kafka.common.errors.ControllerMovedException
 import org.apache.kafka.common.utils.Time
 
 import scala.collection._
@@ -32,22 +32,58 @@ import scala.collection.JavaConverters._
 
 object ControllerEventManager {
   val ControllerEventThreadName = "controller-event-thread"
+  val EventQueueTimeMetricName = "EventQueueTimeMs"
+  val EventQueueSizeMetricName = "EventQueueSize"
 }
-class ControllerEventManager(controllerId: Int, rateAndTimeMetrics: Map[ControllerState, KafkaTimer],
-                             eventProcessedListener: ControllerEvent => Unit,
-                             controllerMovedListener: () => Unit) extends KafkaMetricsGroup {
+
+trait ControllerEventProcessor {
+  def process(event: ControllerEvent): Unit
+  def preempt(event: ControllerEvent): Unit
+}
+
+class QueuedEvent(val event: ControllerEvent,
+                  val enqueueTimeMs: Long) {
+  val processingStarted = new CountDownLatch(1)
+  val spent = new AtomicBoolean(false)
+
+  def process(processor: ControllerEventProcessor): Unit = {
+    if (spent.getAndSet(true))
+      return
+    processingStarted.countDown()
+    processor.process(event)
+  }
+
+  def preempt(processor: ControllerEventProcessor): Unit = {
+    if (spent.getAndSet(true))
+      return
+    processor.preempt(event)
+  }
+
+  def awaitProcessing(): Unit = {
+    processingStarted.await()
+  }
+
+  override def toString: String = {
+    s"QueuedEvent(event=$event, enqueueTimeMs=$enqueueTimeMs)"
+  }
+}
+
+class ControllerEventManager(controllerId: Int,
+                             processor: ControllerEventProcessor,
+                             time: Time,
+                             rateAndTimeMetrics: Map[ControllerState, KafkaTimer]) extends KafkaMetricsGroup {
+  import ControllerEventManager._
 
   @volatile private var _state: ControllerState = ControllerState.Idle
   private val putLock = new ReentrantLock()
-  private val queue = new LinkedBlockingQueue[ControllerEvent]
+  private val queue = new LinkedBlockingQueue[QueuedEvent]
   // Visible for test
-  private[controller] val thread = new ControllerEventThread(ControllerEventManager.ControllerEventThreadName)
-  private val time = Time.SYSTEM
+  private[controller] val thread = new ControllerEventThread(ControllerEventThreadName)
 
-  private val eventQueueTimeHist = newHistogram("EventQueueTimeMs")
+  private val eventQueueTimeHist = newHistogram(EventQueueTimeMetricName)
 
   newGauge(
-    "EventQueueSize",
+    EventQueueSizeMetricName,
     new Gauge[Int] {
       def value: Int = {
         queue.size()
@@ -55,55 +91,53 @@ class ControllerEventManager(controllerId: Int, rateAndTimeMetrics: Map[Controll
     }
   )
 
-
   def state: ControllerState = _state
 
   def start(): Unit = thread.start()
 
   def close(): Unit = {
-    thread.initiateShutdown()
-    clearAndPut(KafkaController.ShutdownEventThread)
-    thread.awaitShutdown()
+    try {
+      thread.initiateShutdown()
+      clearAndPut(ShutdownEventThread)
+      thread.awaitShutdown()
+    } finally {
+      removeMetric(EventQueueTimeMetricName)
+      removeMetric(EventQueueSizeMetricName)
+    }
   }
 
-  def put(event: ControllerEvent): Unit = inLock(putLock) {
-    queue.put(event)
+  def put(event: ControllerEvent): QueuedEvent = inLock(putLock) {
+    val queuedEvent = new QueuedEvent(event, time.milliseconds())
+    queue.put(queuedEvent)
+    queuedEvent
   }
 
-  def clearAndPut(event: ControllerEvent): Unit = inLock(putLock) {
-    queue.asScala.foreach(evt =>
-      if (evt.isInstanceOf[PreemptableControllerEvent])
-        evt.asInstanceOf[PreemptableControllerEvent].preempt()
-    )
+  def clearAndPut(event: ControllerEvent): QueuedEvent = inLock(putLock) {
+    queue.asScala.foreach(_.preempt(processor))
     queue.clear()
     put(event)
   }
 
+  def isEmpty: Boolean = queue.isEmpty
+
   class ControllerEventThread(name: String) extends ShutdownableThread(name = name, isInterruptible = false) {
     logIdent = s"[ControllerEventThread controllerId=$controllerId] "
 
     override def doWork(): Unit = {
-      queue.take() match {
-        case KafkaController.ShutdownEventThread => // The shutting down of the thread has been initiated at this point. Ignore this event.
+      val dequeued = queue.take()
+      dequeued.event match {
+        case ShutdownEventThread => // The shutting down of the thread has been initiated at this point. Ignore this event.
         case controllerEvent =>
           _state = controllerEvent.state
 
-          eventQueueTimeHist.update(time.milliseconds() - controllerEvent.enqueueTimeMs)
+          eventQueueTimeHist.update(time.milliseconds() - dequeued.enqueueTimeMs)
 
           try {
             rateAndTimeMetrics(state).time {
-              controllerEvent.process()
+              dequeued.process(processor)
             }
           } catch {
-            case e: ControllerMovedException =>
-              info(s"Controller moved to another broker when processing $controllerEvent.", e)
-              controllerMovedListener()
-            case e: Throwable => error(s"Error processing event $controllerEvent", e)
-          }
-
-          try eventProcessedListener(controllerEvent)
-          catch {
-            case e: Throwable => error(s"Error while invoking listener for processed event $controllerEvent", e)
+            case e: Throwable => error(s"Uncaught error processing event $controllerEvent", e)
           }
 
           _state = ControllerState.Idle
diff --git a/core/src/main/scala/kafka/controller/ControllerState.scala b/core/src/main/scala/kafka/controller/ControllerState.scala
index aa41c7f54586a..f734acbc6e307 100644
--- a/core/src/main/scala/kafka/controller/ControllerState.scala
+++ b/core/src/main/scala/kafka/controller/ControllerState.scala
@@ -58,8 +58,10 @@ object ControllerState {
     def value = 4
   }
 
-  case object PartitionReassignment extends ControllerState {
+  case object AlterPartitionReassignment extends ControllerState {
     def value = 5
+
+    override def rateAndTimeMetricName: Option[String] = Some("PartitionReassignmentRateAndTimeMs")
   }
 
   case object AutoLeaderBalance extends ControllerState {
@@ -98,7 +100,20 @@ object ControllerState {
     def value = 14
   }
 
+  case object ListPartitionReassignment extends ControllerState {
+    def value = 15
+  }
+
+  case object TopicDeletionFlagChange extends ControllerState {
+    def value = 16
+  }
+
+  case object PreferredControllerChange extends ControllerState {
+    def value = 17
+  }
+
   val values: Seq[ControllerState] = Seq(Idle, ControllerChange, BrokerChange, TopicChange, TopicDeletion,
-    PartitionReassignment, AutoLeaderBalance, ManualLeaderBalance, ControlledShutdown, IsrChange, LeaderAndIsrResponseReceived,
-    LogDirChange, ControllerShutdown, UncleanLeaderElectionEnable, TopicUncleanLeaderElectionEnable)
+    AlterPartitionReassignment, AutoLeaderBalance, ManualLeaderBalance, ControlledShutdown, IsrChange, LeaderAndIsrResponseReceived,
+    LogDirChange, ControllerShutdown, UncleanLeaderElectionEnable, TopicUncleanLeaderElectionEnable, ListPartitionReassignment,
+    TopicDeletionFlagChange, PreferredControllerChange)
 }
diff --git a/core/src/main/scala/kafka/controller/Election.scala b/core/src/main/scala/kafka/controller/Election.scala
new file mode 100644
index 0000000000000..5a59143205bfb
--- /dev/null
+++ b/core/src/main/scala/kafka/controller/Election.scala
@@ -0,0 +1,163 @@
+/*
+ * 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.controller
+
+import kafka.api.LeaderAndIsr
+import kafka.utils.Logging
+import org.apache.kafka.common.TopicPartition
+
+import scala.collection.Seq
+
+case class ElectionResult(topicPartition: TopicPartition, leaderAndIsr: Option[LeaderAndIsr], liveReplicas: Seq[Int])
+
+object Election extends Logging {
+
+  private def leaderForOffline(partition: TopicPartition,
+                               leaderAndIsrOpt: Option[LeaderAndIsr],
+                               uncleanLeaderElectionEnabled: Boolean,
+                               controllerContext: ControllerContext): ElectionResult = {
+
+    val assignment = controllerContext.partitionReplicaAssignment(partition)
+    val liveReplicas = assignment.filter(replica => controllerContext.isReplicaOnline(replica, partition))
+    leaderAndIsrOpt match {
+      case Some(leaderAndIsr) =>
+        val isr = leaderAndIsr.isr
+        val leaderOpt = PartitionLeaderElectionAlgorithms.offlinePartitionLeaderElection(
+          assignment, isr, liveReplicas.toSet, uncleanLeaderElectionEnabled)
+        val newLeaderAndIsrOpt = leaderOpt.map { case (leader, uncleanElection) =>
+          if (uncleanElection) {
+            controllerContext.stats.uncleanLeaderElectionRate.mark()
+            warn(s"Unclean leader election. Partition $partition has been assigned leader $leader from deposed " +
+              s"leader ${leaderAndIsr.leader}.")
+          }
+          val newIsr = if (isr.contains(leader)) isr.filter(replica => controllerContext.isReplicaOnline(replica, partition))
+          else List(leader)
+          leaderAndIsr.newLeaderAndIsr(leader, newIsr)
+        }
+        ElectionResult(partition, newLeaderAndIsrOpt, liveReplicas)
+
+      case None =>
+        ElectionResult(partition, None, liveReplicas)
+    }
+  }
+
+  /**
+   * Elect leaders for new or offline partitions.
+   *
+   * @param controllerContext Context with the current state of the cluster
+   * @param partitionsWithUncleanLeaderElectionState A sequence of tuples representing the partitions
+   *                                                 that need election, their leader/ISR state, and whether
+   *                                                 or not unclean leader election is enabled
+   *
+   * @return The election results
+   */
+  def leaderForOffline(
+    controllerContext: ControllerContext,
+    partitionsWithUncleanLeaderElectionState: Seq[(TopicPartition, Option[LeaderAndIsr], Boolean)]
+  ): Seq[ElectionResult] = {
+    partitionsWithUncleanLeaderElectionState.map {
+      case (partition, leaderAndIsrOpt, uncleanLeaderElectionEnabled) =>
+        leaderForOffline(partition, leaderAndIsrOpt, uncleanLeaderElectionEnabled, controllerContext)
+    }
+  }
+
+  private def leaderForReassign(partition: TopicPartition,
+                                leaderAndIsr: LeaderAndIsr,
+                                controllerContext: ControllerContext): ElectionResult = {
+    val targetReplicas = controllerContext.partitionFullReplicaAssignment(partition).targetReplicas
+    val liveReplicas = targetReplicas.filter(replica => controllerContext.isReplicaOnline(replica, partition))
+    val isr = leaderAndIsr.isr
+    val leaderOpt = PartitionLeaderElectionAlgorithms.reassignPartitionLeaderElection(targetReplicas, isr, liveReplicas.toSet)
+    val newLeaderAndIsrOpt = leaderOpt.map(leader => leaderAndIsr.newLeader(leader))
+    ElectionResult(partition, newLeaderAndIsrOpt, targetReplicas)
+  }
+
+  /**
+   * Elect leaders for partitions that are undergoing reassignment.
+   *
+   * @param controllerContext Context with the current state of the cluster
+   * @param leaderAndIsrs A sequence of tuples representing the partitions that need election
+   *                                     and their respective leader/ISR states
+   *
+   * @return The election results
+   */
+  def leaderForReassign(controllerContext: ControllerContext,
+                        leaderAndIsrs: Seq[(TopicPartition, LeaderAndIsr)]): Seq[ElectionResult] = {
+    leaderAndIsrs.map { case (partition, leaderAndIsr) =>
+      leaderForReassign(partition, leaderAndIsr, controllerContext)
+    }
+  }
+
+  private def leaderForPreferredReplica(partition: TopicPartition,
+                                        leaderAndIsr: LeaderAndIsr,
+                                        controllerContext: ControllerContext): ElectionResult = {
+    val assignment = controllerContext.partitionReplicaAssignment(partition)
+    val liveReplicas = assignment.filter(replica => controllerContext.isReplicaOnline(replica, partition))
+    val isr = leaderAndIsr.isr
+    val leaderOpt = PartitionLeaderElectionAlgorithms.preferredReplicaPartitionLeaderElection(assignment, isr, liveReplicas.toSet)
+    val newLeaderAndIsrOpt = leaderOpt.map(leader => leaderAndIsr.newLeader(leader))
+    ElectionResult(partition, newLeaderAndIsrOpt, assignment)
+  }
+
+  /**
+   * Elect preferred leaders.
+   *
+   * @param controllerContext Context with the current state of the cluster
+   * @param leaderAndIsrs A sequence of tuples representing the partitions that need election
+   *                                     and their respective leader/ISR states
+   *
+   * @return The election results
+   */
+  def leaderForPreferredReplica(controllerContext: ControllerContext,
+                                leaderAndIsrs: Seq[(TopicPartition, LeaderAndIsr)]): Seq[ElectionResult] = {
+    leaderAndIsrs.map { case (partition, leaderAndIsr) =>
+      leaderForPreferredReplica(partition, leaderAndIsr, controllerContext)
+    }
+  }
+
+  private def leaderForControlledShutdown(partition: TopicPartition,
+                                          leaderAndIsr: LeaderAndIsr,
+                                          shuttingDownBrokerIds: Set[Int],
+                                          controllerContext: ControllerContext): ElectionResult = {
+    val assignment = controllerContext.partitionReplicaAssignment(partition)
+    val liveOrShuttingDownReplicas = assignment.filter(replica =>
+      controllerContext.isReplicaOnline(replica, partition, includeShuttingDownBrokers = true))
+    val isr = leaderAndIsr.isr
+    val leaderOpt = PartitionLeaderElectionAlgorithms.controlledShutdownPartitionLeaderElection(assignment, isr,
+      liveOrShuttingDownReplicas.toSet, shuttingDownBrokerIds)
+    val newIsr = isr.filter(replica => !shuttingDownBrokerIds.contains(replica))
+    val newLeaderAndIsrOpt = leaderOpt.map(leader => leaderAndIsr.newLeaderAndIsr(leader, newIsr))
+    ElectionResult(partition, newLeaderAndIsrOpt, liveOrShuttingDownReplicas)
+  }
+
+  /**
+   * Elect leaders for partitions whose current leaders are shutting down.
+   *
+   * @param controllerContext Context with the current state of the cluster
+   * @param leaderAndIsrs A sequence of tuples representing the partitions that need election
+   *                                     and their respective leader/ISR states
+   *
+   * @return The election results
+   */
+  def leaderForControlledShutdown(controllerContext: ControllerContext,
+                                  leaderAndIsrs: Seq[(TopicPartition, LeaderAndIsr)]): Seq[ElectionResult] = {
+    val shuttingDownBrokerIds = controllerContext.shuttingDownBrokerIds.toSet
+    leaderAndIsrs.map { case (partition, leaderAndIsr) =>
+      leaderForControlledShutdown(partition, leaderAndIsr, shuttingDownBrokerIds, controllerContext)
+    }
+  }
+}
diff --git a/core/src/main/scala/kafka/controller/KafkaController.scala b/core/src/main/scala/kafka/controller/KafkaController.scala
index ea23beb419234..e8b744227b968 100644
--- a/core/src/main/scala/kafka/controller/KafkaController.scala
+++ b/core/src/main/scala/kafka/controller/KafkaController.scala
@@ -16,55 +16,92 @@
  */
 package kafka.controller
 
-import java.util.concurrent.atomic.AtomicBoolean
-import java.util.concurrent.{CountDownLatch, TimeUnit}
+import java.util.concurrent.TimeUnit
 
 import com.yammer.metrics.core.Gauge
-import kafka.admin.AdminOperationException
+import kafka.admin.{AdminOperationException, AdminUtils}
 import kafka.api._
 import kafka.common._
+import kafka.controller.KafkaController.{AlterReassignmentsCallback, ElectLeadersCallback, ListReassignmentsCallback}
 import kafka.metrics.{KafkaMetricsGroup, KafkaTimer}
 import kafka.server._
 import kafka.utils._
 import kafka.zk.KafkaZkClient.UpdateLeaderAndIsrResult
 import kafka.zk._
 import kafka.zookeeper.{StateChangeHandler, ZNodeChangeHandler, ZNodeChildChangeHandler}
-import org.apache.kafka.common.{KafkaException, TopicPartition}
+import org.apache.kafka.common.ElectionType
+import org.apache.kafka.common.KafkaException
+import org.apache.kafka.common.TopicPartition
 import org.apache.kafka.common.errors.{BrokerNotAvailableException, ControllerMovedException, StaleBrokerEpochException}
+import org.apache.kafka.common.errors.PolicyViolationException
 import org.apache.kafka.common.metrics.Metrics
-import org.apache.kafka.common.protocol.{ApiKeys, Errors}
-import org.apache.kafka.common.requests.{AbstractControlRequest, AbstractResponse, ApiError, LeaderAndIsrResponse, StopReplicaResponse}
+import org.apache.kafka.common.protocol.Errors
+import org.apache.kafka.common.requests.{AbstractControlRequest, AbstractResponse, ApiError, LeaderAndIsrResponse}
 import org.apache.kafka.common.utils.Time
 import org.apache.zookeeper.KeeperException
 import org.apache.zookeeper.KeeperException.Code
+import org.apache.kafka.server.policy.CreateTopicPolicy
 
-import scala.collection._
+import scala.collection.JavaConverters._
+import scala.collection.{Map, Seq, Set, immutable, mutable}
+import scala.collection.mutable.ArrayBuffer
 import scala.util.{Failure, Try}
 
+sealed trait ElectionTrigger
+final case object AutoTriggered extends ElectionTrigger
+final case object ZkTriggered extends ElectionTrigger
+final case object AdminClientTriggered extends ElectionTrigger
+
 object KafkaController extends Logging {
   val InitialControllerEpoch = 0
   val InitialControllerEpochZkVersion = 0
 
-  /**
-   * ControllerEventThread will shutdown once it sees this event
-   */
-  private[controller] case object ShutdownEventThread extends ControllerEvent {
-    def state = ControllerState.ControllerShutdown
-    override def process(): Unit = ()
-  }
+  type ElectLeadersCallback = Map[TopicPartition, Either[ApiError, Int]] => Unit
+  type ListReassignmentsCallback = Either[Map[TopicPartition, ReplicaAssignment], ApiError] => Unit
+  type AlterReassignmentsCallback = Either[Map[TopicPartition, ApiError], ApiError] => Unit
 
-  // Used only by test
-  private[controller] case class AwaitOnLatch(latch: CountDownLatch) extends ControllerEvent {
-    override def state: ControllerState = ControllerState.ControllerChange
-    override def process(): Unit = latch.await()
+  def satisfiesLiCreateTopicPolicy(createTopicPolicy : Option[CreateTopicPolicy], zkClient : KafkaZkClient,
+    topic : String, partitionsAssignment : collection.Map[Int, ReplicaAssignment]): Boolean = {
+    try {
+      createTopicPolicy match {
+        case Some(policy) =>
+          if (policy.isInstanceOf[LiCreateTopicPolicy]) {
+            import scala.collection.JavaConverters._
+            val jPartitionAssignment = partitionsAssignment.map { case(partition, replicaAssignment) =>
+              (new Integer(partition), seqAsJavaListConverter(replicaAssignment.replicas.map{e => new Integer(e)}).asJava)
+            }
+            // Use min size of all replica lists as a stand in for replicationFactor. Generally replicas sizes should be
+            // the same, but minBy gets us the worst case.
+            val replicationFactor = partitionsAssignment.minBy(_._2.replicas.size)._1.toShort
+            policy.validate(new CreateTopicPolicy.RequestMetadata(topic, partitionsAssignment.size, replicationFactor,
+              jPartitionAssignment.asJava, new java.util.HashMap[String, String]()))
+          }
+          true
+        case None =>
+          true
+      }
+    } catch {
+      case e : PolicyViolationException => {
+        if (zkClient.getTopicPartitions(topic).isEmpty) {
+          info(e.getMessage)
+          false
+        } else true
+      }
+    }
   }
-
 }
 
-class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Time, metrics: Metrics,
-                      initialBrokerInfo: BrokerInfo, initialBrokerEpoch: Long, tokenManager: DelegationTokenManager,
-                      threadNamePrefix: Option[String] = None) extends Logging with KafkaMetricsGroup {
-
+class KafkaController(val config: KafkaConfig,
+                      zkClient: KafkaZkClient,
+                      time: Time,
+                      metrics: Metrics,
+                      initialBrokerInfo: BrokerInfo,
+                      initialBrokerEpoch: Long,
+                      tokenManager: DelegationTokenManager,
+                      threadNamePrefix: Option[String] = None)
+  extends ControllerEventProcessor with Logging with KafkaMetricsGroup {
+
+  val adminZkClient = new AdminZkClient(zkClient)
   this.logIdent = s"[Controller id=${config.brokerId}] "
 
   @volatile private var brokerInfo = initialBrokerInfo
@@ -72,37 +109,50 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
 
   private val stateChangeLogger = new StateChangeLogger(config.brokerId, inControllerContext = true, None)
   val controllerContext = new ControllerContext
+  var controllerChannelManager = new ControllerChannelManager(controllerContext, config, time, metrics,
+    stateChangeLogger, threadNamePrefix)
 
   // have a separate scheduler for the controller to be able to start and stop independently of the kafka server
   // visible for testing
   private[controller] val kafkaScheduler = new KafkaScheduler(1)
 
   // visible for testing
-  private[controller] val eventManager = new ControllerEventManager(config.brokerId,
-    controllerContext.stats.rateAndTimeMetrics, _ => updateMetrics(), () => maybeResign())
-
-  val topicDeletionManager = new TopicDeletionManager(this, eventManager, zkClient)
-  private val brokerRequestBatch = new ControllerBrokerRequestBatch(this, stateChangeLogger)
-  val replicaStateMachine = new ReplicaStateMachine(config, stateChangeLogger, controllerContext, topicDeletionManager, zkClient, mutable.Map.empty, new ControllerBrokerRequestBatch(this, stateChangeLogger))
-  val partitionStateMachine = new PartitionStateMachine(config, stateChangeLogger, controllerContext, zkClient, mutable.Map.empty, new ControllerBrokerRequestBatch(this, stateChangeLogger))
-  partitionStateMachine.setTopicDeletionManager(topicDeletionManager)
-
-  private val controllerChangeHandler = new ControllerChangeHandler(this, eventManager)
-  private val brokerChangeHandler = new BrokerChangeHandler(this, eventManager)
+  private[controller] val eventManager = new ControllerEventManager(config.brokerId, this, time,
+    controllerContext.stats.rateAndTimeMetrics)
+
+  private val brokerRequestBatch = new ControllerBrokerRequestBatch(config, controllerChannelManager,
+    eventManager, controllerContext, stateChangeLogger)
+  val replicaStateMachine: ReplicaStateMachine = new ZkReplicaStateMachine(config, stateChangeLogger, controllerContext, zkClient,
+    new ControllerBrokerRequestBatch(config, controllerChannelManager, eventManager, controllerContext, stateChangeLogger))
+  val partitionStateMachine: PartitionStateMachine = new ZkPartitionStateMachine(config, stateChangeLogger, controllerContext, zkClient,
+    new ControllerBrokerRequestBatch(config, controllerChannelManager, eventManager, controllerContext, stateChangeLogger))
+  val topicDeletionManager = new TopicDeletionManager(config, controllerContext, replicaStateMachine,
+    partitionStateMachine, new ControllerDeletionClient(this, zkClient))
+
+  private val controllerChangeHandler = new ControllerChangeHandler(eventManager)
+  private val brokerChangeHandler = new BrokerChangeHandler(eventManager)
+  private val preferredControllerChangeHandler = new PreferredControllerChangeHandler(eventManager)
   private val brokerModificationsHandlers: mutable.Map[Int, BrokerModificationsHandler] = mutable.Map.empty
-  private val topicChangeHandler = new TopicChangeHandler(this, eventManager)
-  private val topicDeletionHandler = new TopicDeletionHandler(this, eventManager)
+  private val topicChangeHandler = new TopicChangeHandler(eventManager)
+  private val topicDeletionHandler = new TopicDeletionHandler(eventManager)
   private val partitionModificationsHandlers: mutable.Map[String, PartitionModificationsHandler] = mutable.Map.empty
-  private val partitionReassignmentHandler = new PartitionReassignmentHandler(this, eventManager)
-  private val preferredReplicaElectionHandler = new PreferredReplicaElectionHandler(this, eventManager)
-  private val isrChangeNotificationHandler = new IsrChangeNotificationHandler(this, eventManager)
-  private val logDirEventNotificationHandler = new LogDirEventNotificationHandler(this, eventManager)
+  private val partitionReassignmentHandler = new PartitionReassignmentHandler(eventManager)
+  private val preferredReplicaElectionHandler = new PreferredReplicaElectionHandler(eventManager)
+  private val isrChangeNotificationHandler = new IsrChangeNotificationHandler(eventManager)
+  private val logDirEventNotificationHandler = new LogDirEventNotificationHandler(eventManager)
+  private val topicDeletionFlagHandler = new TopicDeletionFlagHandler(this, eventManager)
 
   @volatile private var activeControllerId = -1
   @volatile private var offlinePartitionCount = 0
-  @volatile private var preferredReplicaImbalanceCount = 0
   @volatile private var globalTopicCount = 0
   @volatile private var globalPartitionCount = 0
+  @volatile private var topicsToDeleteCount = 0
+  @volatile private var replicasToDeleteCount = 0
+  @volatile private var ineligibleTopicsToDeleteCount = 0
+  @volatile private var ineligibleReplicasToDeleteCount = 0
+
+  private val createTopicPolicy =
+    Option(config.getConfiguredInstance(KafkaConfig.CreateTopicPolicyClassNameProp, classOf[CreateTopicPolicy]))
 
   /* single-thread scheduler to clean expired tokens */
   private val tokenCleanScheduler = new KafkaScheduler(threads = 1, threadNamePrefix = "delegation-token-cleaner")
@@ -115,16 +165,23 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
   )
 
   newGauge(
-    "OfflinePartitionsCount",
+    "ActivePreferredControllerCount",
     new Gauge[Int] {
-      def value: Int = offlinePartitionCount
+      def value = if (isActive && config.preferredController) 1 else 0
     }
   )
 
   newGauge(
-    "PreferredReplicaImbalanceCount",
+    "StandbyPreferredControllerCount",
+    new Gauge[Int] {
+      def value = if (!isActive && config.preferredController) 1 else 0
+    }
+  )
+
+  newGauge(
+    "OfflinePartitionsCount",
     new Gauge[Int] {
-      def value: Int = preferredReplicaImbalanceCount
+      def value: Int = offlinePartitionCount
     }
   )
 
@@ -149,6 +206,41 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
     }
   )
 
+  newGauge(
+    "TopicsToDeleteCount",
+    new Gauge[Int] {
+      def value: Int = topicsToDeleteCount
+    }
+  )
+
+  newGauge(
+    "ReplicasToDeleteCount",
+    new Gauge[Int] {
+      def value: Int = replicasToDeleteCount
+    }
+  )
+
+  newGauge(
+    "TopicsIneligibleToDeleteCount",
+    new Gauge[Int] {
+      def value: Int = ineligibleTopicsToDeleteCount
+    }
+  )
+
+  newGauge(
+    "ReplicasIneligibleToDeleteCount",
+    new Gauge[Int] {
+      def value: Int = ineligibleReplicasToDeleteCount
+    }
+  )
+
+    newGauge(
+    "MaintenanceBrokerCount",
+    new Gauge[Int] {
+      def value: Int = if (isActive) config.getMaintenanceBrokerList.size else 0
+    }
+  )
+
   /**
    * Returns true if this broker is the current controller.
    */
@@ -158,6 +250,11 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
 
   def epoch: Int = controllerContext.epoch
 
+  /**
+    * @return brokers that don't accept new replicas, including maintenance brokers and preferred controller brokers
+    */
+  def partitionUnassignableBrokerIds: Seq[Int] = config.getMaintenanceBrokerList ++ controllerContext.getLivePreferredControllerIds
+
   /**
    * Invoked when the controller module of a Kafka server is started up. This does not assume that the current broker
    * is the controller. It merely registers the session expiration listener and starts the controller leader
@@ -170,12 +267,11 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
         eventManager.put(RegisterBrokerAndReelect)
       }
       override def beforeInitializingSession(): Unit = {
-        val expireEvent = new Expire
-        eventManager.clearAndPut(expireEvent)
+        val queuedEvent = eventManager.clearAndPut(Expire)
 
         // Block initialization of the new session until the expiration event is being handled,
         // which ensures that all pending events have been processed before creating the new session
-        expireEvent.waitUntilProcessingStarted()
+        queuedEvent.awaitProcessing()
       }
     })
     eventManager.put(Startup)
@@ -221,6 +317,10 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
     }
   }
 
+  private[kafka] def enablePreferredControllerFallback(): Unit = {
+    eventManager.put(PreferredControllerChange)
+  }
+
   private def state: ControllerState = eventManager.state
 
   /**
@@ -234,14 +334,14 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
    * If it encounters any unexpected exception/error while becoming controller, it resigns as the current controller.
    * This ensures another controller election will be triggered and there will always be an actively serving controller
    */
-  private def onControllerFailover() {
+  private def onControllerFailover(): Unit = {
     info("Registering handlers")
 
     // before reading source of truth from zookeeper, register the listeners to get broker/topic callbacks
     val childChangeHandlers = Seq(brokerChangeHandler, topicChangeHandler, topicDeletionHandler, logDirEventNotificationHandler,
       isrChangeNotificationHandler)
     childChangeHandlers.foreach(zkClient.registerZNodeChildChangeHandler)
-    val nodeChangeHandlers = Seq(preferredReplicaElectionHandler, partitionReassignmentHandler)
+    val nodeChangeHandlers = Seq(preferredReplicaElectionHandler, partitionReassignmentHandler, topicDeletionFlagHandler)
     nodeChangeHandlers.foreach(zkClient.registerZNodeChangeHandlerAndCheckExistence)
 
     info("Deleting log dir event notifications")
@@ -266,10 +366,11 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
     partitionStateMachine.startup()
 
     info(s"Ready to serve as the new controller with epoch $epoch")
-    maybeTriggerPartitionReassignment(controllerContext.partitionsBeingReassigned.keySet)
+
+    initializePartitionReassignments()
     topicDeletionManager.tryTopicDeletion()
     val pendingPreferredReplicaElections = fetchPendingPreferredReplicaElections()
-    onPreferredReplicaElection(pendingPreferredReplicaElections, ZkTriggered)
+    onReplicaElection(pendingPreferredReplicaElections, ElectionType.PREFERRED, ZkTriggered)
     info("Starting the controller scheduler")
     kafkaScheduler.startup()
     if (config.autoLeaderRebalanceEnable) {
@@ -295,7 +396,7 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
    * This callback is invoked by the zookeeper leader elector when the current broker resigns as the controller. This is
    * required to clean up internal controller data structures
    */
-  private def onControllerResignation() {
+  private def onControllerResignation(): Unit = {
     debug("Resigning")
     // de-register listeners
     zkClient.unregisterZNodeChildChangeHandler(isrChangeNotificationHandler.path)
@@ -304,15 +405,15 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
     zkClient.unregisterZNodeChildChangeHandler(logDirEventNotificationHandler.path)
     unregisterBrokerModificationsHandler(brokerModificationsHandlers.keySet)
 
-    // reset topic deletion manager
-    topicDeletionManager.reset()
-
     // shutdown leader rebalance scheduler
     kafkaScheduler.shutdown()
     offlinePartitionCount = 0
-    preferredReplicaImbalanceCount = 0
     globalTopicCount = 0
     globalPartitionCount = 0
+    topicsToDeleteCount = 0
+    replicasToDeleteCount = 0
+    ineligibleTopicsToDeleteCount = 0
+    ineligibleReplicasToDeleteCount = 0
 
     // stop token expiry check scheduler
     if (tokenCleanScheduler.isStarted)
@@ -325,10 +426,12 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
     zkClient.unregisterZNodeChildChangeHandler(topicChangeHandler.path)
     unregisterPartitionModificationsHandlers(partitionModificationsHandlers.keys.toSeq)
     zkClient.unregisterZNodeChildChangeHandler(topicDeletionHandler.path)
+    zkClient.unregisterZNodeChangeHandler(topicDeletionFlagHandler.path)
     // shutdown replica state machine
     replicaStateMachine.shutdown()
     zkClient.unregisterZNodeChildChangeHandler(brokerChangeHandler.path)
 
+    controllerChannelManager.shutdown()
     controllerContext.resetContext()
 
     info("Resigned")
@@ -340,7 +443,7 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
    * to all these brokers to query the state of their replicas. Replicas with an offline log directory respond with
    * KAFKA_STORAGE_ERROR, which will be handled by the LeaderAndIsrResponseReceived event.
    */
-  private def onBrokerLogDirFailure(brokerIds: Seq[Int]) {
+  private def onBrokerLogDirFailure(brokerIds: Seq[Int]): Unit = {
     // send LeaderAndIsrRequest for all replicas on those brokers to see if they are still online.
     info(s"Handling log directory failure for brokers ${brokerIds.mkString(",")}")
     val replicasOnBrokers = controllerContext.replicasOnBrokers(brokerIds.toSet)
@@ -361,7 +464,7 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
    * 2. Even if we do refresh the cache, there is no guarantee that by the time the leader and ISR request reaches
    *    every broker that it is still valid.  Brokers check the leader epoch to determine validity of the request.
    */
-  private def onBrokerStartup(newBrokers: Seq[Int]) {
+  private def onBrokerStartup(newBrokers: Seq[Int]): Unit = {
     info(s"New broker startup callback for ${newBrokers.mkString(",")}")
     newBrokers.foreach(controllerContext.replicasOnOfflineDirs.remove)
     val newBrokersSet = newBrokers.toSet
@@ -381,26 +484,33 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
     // to see if these brokers can become leaders for some/all of those
     partitionStateMachine.triggerOnlinePartitionStateChange()
     // check if reassignment of some partitions need to be restarted
-    val partitionsWithReplicasOnNewBrokers = controllerContext.partitionsBeingReassigned.filter {
-      case (_, reassignmentContext) => reassignmentContext.newReplicas.exists(newBrokersSet.contains)
+    maybeResumeReassignments { (_, assignment) =>
+      assignment.targetReplicas.exists(newBrokersSet.contains)
     }
-    partitionsWithReplicasOnNewBrokers.foreach { case (tp, context) => onPartitionReassignment(tp, context) }
     // check if topic deletion needs to be resumed. If at least one replica that belongs to the topic being deleted exists
     // on the newly restarted brokers, there is a chance that topic deletion can resume
     val replicasForTopicsToBeDeleted = allReplicasOnNewBrokers.filter(p => topicDeletionManager.isTopicQueuedUpForDeletion(p.topic))
     if (replicasForTopicsToBeDeleted.nonEmpty) {
       info(s"Some replicas ${replicasForTopicsToBeDeleted.mkString(",")} for topics scheduled for deletion " +
-        s"${topicDeletionManager.topicsToBeDeleted.mkString(",")} are on the newly restarted brokers " +
+        s"${controllerContext.topicsToBeDeleted.mkString(",")} are on the newly restarted brokers " +
         s"${newBrokers.mkString(",")}. Signaling restart of topic deletion for these topics")
       topicDeletionManager.resumeDeletionForTopics(replicasForTopicsToBeDeleted.map(_.topic))
     }
     registerBrokerModificationsHandler(newBrokers)
   }
 
+  private def maybeResumeReassignments(shouldResume: (TopicPartition, ReplicaAssignment) => Boolean): Unit = {
+    controllerContext.partitionsBeingReassigned.foreach { tp =>
+      val currentAssignment = controllerContext.partitionFullReplicaAssignment(tp)
+      if (shouldResume(tp, currentAssignment))
+        onPartitionReassignment(tp, currentAssignment)
+    }
+  }
+
   private def registerBrokerModificationsHandler(brokerIds: Iterable[Int]): Unit = {
     debug(s"Register BrokerModifications handler for $brokerIds")
     brokerIds.foreach { brokerId =>
-      val brokerModificationsHandler = new BrokerModificationsHandler(this, eventManager, brokerId)
+      val brokerModificationsHandler = new BrokerModificationsHandler(eventManager, brokerId)
       zkClient.registerZNodeChangeHandlerAndCheckExistence(brokerModificationsHandler)
       brokerModificationsHandlers.put(brokerId, brokerModificationsHandler)
     }
@@ -417,7 +527,7 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
    * This callback is invoked by the replica state machine's broker change listener with the list of failed brokers
    * as input. It will call onReplicaBecomeOffline(...) with the list of replicas on those failed brokers as input.
    */
-  private def onBrokerFailure(deadBrokers: Seq[Int]) {
+  private def onBrokerFailure(deadBrokers: Seq[Int]): Unit = {
     info(s"Broker failure callback for ${deadBrokers.mkString(",")}")
     deadBrokers.foreach(controllerContext.replicasOnOfflineDirs.remove)
     val deadBrokersThatWereShuttingDown =
@@ -430,7 +540,7 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
     unregisterBrokerModificationsHandler(deadBrokers)
   }
 
-  private def onBrokerUpdate(updatedBrokerId: Int) {
+  private def onBrokerUpdate(updatedBrokerId: Int): Unit = {
     info(s"Broker info update callback for $updatedBrokerId")
     sendUpdateMetadataRequest(controllerContext.liveOrShuttingDownBrokerIds.toSeq, Set.empty)
   }
@@ -482,90 +592,117 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
    * 1. Move the newly created partitions to the NewPartition state
    * 2. Move the newly created partitions from NewPartition->OnlinePartition state
    */
-  private def onNewPartitionCreation(newPartitions: Set[TopicPartition]) {
+  private def onNewPartitionCreation(newPartitions: Set[TopicPartition]): Unit = {
     info(s"New partition creation callback for ${newPartitions.mkString(",")}")
     partitionStateMachine.handleStateChanges(newPartitions.toSeq, NewPartition)
     replicaStateMachine.handleStateChanges(controllerContext.replicasForPartition(newPartitions).toSeq, NewReplica)
-    partitionStateMachine.handleStateChanges(newPartitions.toSeq, OnlinePartition, Option(OfflinePartitionLeaderElectionStrategy))
+    partitionStateMachine.handleStateChanges(
+      newPartitions.toSeq,
+      OnlinePartition,
+      Some(OfflinePartitionLeaderElectionStrategy(false))
+    )
     replicaStateMachine.handleStateChanges(controllerContext.replicasForPartition(newPartitions).toSeq, OnlineReplica)
   }
 
   /**
-   * This callback is invoked by the reassigned partitions listener. When an admin command initiates a partition
-   * reassignment, it creates the /admin/reassign_partitions path that triggers the zookeeper listener.
+   * This callback is invoked:
+   * 1. By the AlterPartitionReassignments API
+   * 2. By the reassigned partitions listener which is triggered when the /admin/reassign/partitions znode is created
+   * 3. When an ongoing reassignment finishes - this is detected by a change in the partition's ISR znode
+   * 4. Whenever a new broker comes up which is part of an ongoing reassignment
+   * 5. On controller startup/failover
+   *
    * Reassigning replicas for a partition goes through a few steps listed in the code.
-   * RAR = Reassigned replicas
-   * OAR = Original list of replicas for partition
-   * AR = current assigned replicas
+   * RS = current assigned replica set
+   * ORS = Original replica set for partition
+   * TRS = Reassigned (target) replica set
+   * AR = The replicas we are adding as part of this reassignment
+   * RR = The replicas we are removing as part of this reassignment
+   *
+   * A reassignment may have up to three phases, each with its own steps:
+
+   * Phase U (Assignment update): Regardless of the trigger, the first step is in the reassignment process
+   * is to update the existing assignment state. We always update the state in Zookeeper before
+   * we update memory so that it can be resumed upon controller fail-over.
    *
-   * 1. Update AR in ZK with OAR + RAR.
-   * 2. Send LeaderAndIsr request to every replica in OAR + RAR (with AR as OAR + RAR). We do this by forcing an update
-   *    of the leader epoch in zookeeper.
-   * 3. Start new replicas RAR - OAR by moving replicas in RAR - OAR to NewReplica state.
-   * 4. Wait until all replicas in RAR are in sync with the leader.
-   * 5  Move all replicas in RAR to OnlineReplica state.
-   * 6. Set AR to RAR in memory.
-   * 7. If the leader is not in RAR, elect a new leader from RAR. If new leader needs to be elected from RAR, a LeaderAndIsr
-   *    will be sent. If not, then leader epoch will be incremented in zookeeper and a LeaderAndIsr request will be sent.
-   *    In any case, the LeaderAndIsr request will have AR = RAR. This will prevent the leader from adding any replica in
-   *    RAR - OAR back in the isr.
-   * 8. Move all replicas in OAR - RAR to OfflineReplica state. As part of OfflineReplica state change, we shrink the
-   *    isr to remove OAR - RAR in zookeeper and send a LeaderAndIsr ONLY to the Leader to notify it of the shrunk isr.
-   *    After that, we send a StopReplica (delete = false) to the replicas in OAR - RAR.
-   * 9. Move all replicas in OAR - RAR to NonExistentReplica state. This will send a StopReplica (delete = true) to
-   *    the replicas in OAR - RAR to physically delete the replicas on disk.
-   * 10. Update AR in ZK with RAR.
-   * 11. Update the /admin/reassign_partitions path in ZK to remove this partition.
-   * 12. After electing leader, the replicas and isr information changes. So resend the update metadata request to every broker.
+   *   U1. Update ZK with RS = ORS + TRS, AR = TRS - ORS, RR = ORS - TRS.
+   *   U2. Update memory with RS = ORS + TRS, AR = TRS - ORS and RR = ORS - TRS
+   *   U3. If we are cancelling or replacing an existing reassignment, send StopReplica to all members
+   *       of AR in the original reassignment if they are not in TRS from the new assignment
    *
-   * For example, if OAR = {1, 2, 3} and RAR = {4,5,6}, the values in the assigned replica (AR) and leader/isr path in ZK
-   * may go through the following transition.
-   * AR                 leader/isr
-   * {1,2,3}            1/{1,2,3}           (initial state)
-   * {1,2,3,4,5,6}      1/{1,2,3}           (step 2)
-   * {1,2,3,4,5,6}      1/{1,2,3,4,5,6}     (step 4)
-   * {1,2,3,4,5,6}      4/{1,2,3,4,5,6}     (step 7)
-   * {1,2,3,4,5,6}      4/{4,5,6}           (step 8)
-   * {4,5,6}            4/{4,5,6}           (step 10)
+   * To complete the reassignment, we need to bring the new replicas into sync, so depending on the state
+   * of the ISR, we will execute one of the following steps.
    *
-   * Note that we have to update AR in ZK with RAR last since it's the only place where we store OAR persistently.
+   * Phase A (when TRS != ISR): The reassignment is not yet complete
+   *
+   *   A1. Bump the leader epoch for the partition and send LeaderAndIsr updates to RS.
+   *   A2. Start new replicas AR by moving replicas in AR to NewReplica state.
+   *
+   * Phase B (when TRS = ISR): The reassignment is complete
+   *
+   *   B1. Move all replicas in AR to OnlineReplica state.
+   *   B2. Set RS = TRS, AR = [], RR = [] in memory.
+   *   B3. Send a LeaderAndIsr request with RS = TRS. This will prevent the leader from adding any replica in TRS - ORS back in the isr.
+   *       If the current leader is not in TRS or isn't alive, we move the leader to a new replica in TRS.
+   *       We may send the LeaderAndIsr to more than the TRS replicas due to the
+   *       way the partition state machine works (it reads replicas from ZK)
+   *   B4. Move all replicas in RR to OfflineReplica state. As part of OfflineReplica state change, we shrink the
+   *       isr to remove RR in ZooKeeper and send a LeaderAndIsr ONLY to the Leader to notify it of the shrunk isr.
+   *       After that, we send a StopReplica (delete = false) to the replicas in RR.
+   *   B5. Move all replicas in RR to NonExistentReplica state. This will send a StopReplica (delete = true) to
+   *       the replicas in RR to physically delete the replicas on disk.
+   *   B6. Update ZK with RS=TRS, AR=[], RR=[].
+   *   B7. Remove the ISR reassign listener and maybe update the /admin/reassign_partitions path in ZK to remove this partition from it if present.
+   *   B8. After electing leader, the replicas and isr information changes. So resend the update metadata request to every broker.
+   *
+   * In general, there are two goals we want to aim for:
+   * 1. Every replica present in the replica set of a LeaderAndIsrRequest gets the request sent to it
+   * 2. Replicas that are removed from a partition's assignment get StopReplica sent to them
+   *
+   * For example, if ORS = {1,2,3} and TRS = {4,5,6}, the values in the topic and leader/isr paths in ZK
+   * may go through the following transitions.
+   * RS                AR          RR          leader     isr
+   * {1,2,3}           {}          {}          1          {1,2,3}           (initial state)
+   * {4,5,6,1,2,3}     {4,5,6}     {1,2,3}     1          {1,2,3}           (step A2)
+   * {4,5,6,1,2,3}     {4,5,6}     {1,2,3}     1          {1,2,3,4,5,6}     (phase B)
+   * {4,5,6,1,2,3}     {4,5,6}     {1,2,3}     4          {1,2,3,4,5,6}     (step B3)
+   * {4,5,6,1,2,3}     {4,5,6}     {1,2,3}     4          {4,5,6}           (step B4)
+   * {4,5,6}           {}          {}          4          {4,5,6}           (step B6)
+   *
+   * Note that we have to update RS in ZK with TRS last since it's the only place where we store ORS persistently.
    * This way, if the controller crashes before that step, we can still recover.
    */
-  private def onPartitionReassignment(topicPartition: TopicPartition, reassignedPartitionContext: ReassignedPartitionsContext) {
-    val reassignedReplicas = reassignedPartitionContext.newReplicas
-    if (!areReplicasInIsr(topicPartition, reassignedReplicas)) {
-      info(s"New replicas ${reassignedReplicas.mkString(",")} for partition $topicPartition being reassigned not yet " +
-        "caught up with the leader")
-      val newReplicasNotInOldReplicaList = reassignedReplicas.toSet -- controllerContext.partitionReplicaAssignment(topicPartition).toSet
-      val newAndOldReplicas = (reassignedPartitionContext.newReplicas ++ controllerContext.partitionReplicaAssignment(topicPartition)).toSet
-      //1. Update AR in ZK with OAR + RAR.
-      updateAssignedReplicasForPartition(topicPartition, newAndOldReplicas.toSeq)
-      //2. Send LeaderAndIsr request to every replica in OAR + RAR (with AR as OAR + RAR).
-      updateLeaderEpochAndSendRequest(topicPartition, controllerContext.partitionReplicaAssignment(topicPartition),
-        newAndOldReplicas.toSeq)
-      //3. replicas in RAR - OAR -> NewReplica
-      startNewReplicasForReassignedPartition(topicPartition, reassignedPartitionContext, newReplicasNotInOldReplicaList)
-      info(s"Waiting for new replicas ${reassignedReplicas.mkString(",")} for partition ${topicPartition} being " +
-        "reassigned to catch up with the leader")
+  private def onPartitionReassignment(topicPartition: TopicPartition, reassignment: ReplicaAssignment): Unit = {
+    // While a reassignment is in progress, deletion is not allowed
+    topicDeletionManager.markTopicIneligibleForDeletion(Set(topicPartition.topic), reason = "topic reassignment in progress")
+
+    updateCurrentReassignment(topicPartition, reassignment)
+
+    val addingReplicas = reassignment.addingReplicas
+    val removingReplicas = reassignment.removingReplicas
+
+    if (!isReassignmentComplete(topicPartition, reassignment)) {
+      // A1. Send LeaderAndIsr request to every replica in ORS + TRS (with the new RS, AR and RR).
+      updateLeaderEpochAndSendRequest(topicPartition, reassignment)
+      // A2. replicas in AR -> NewReplica
+      startNewReplicasForReassignedPartition(topicPartition, addingReplicas)
     } else {
-      //4. Wait until all replicas in RAR are in sync with the leader.
-      val oldReplicas = controllerContext.partitionReplicaAssignment(topicPartition).toSet -- reassignedReplicas.toSet
-      //5. replicas in RAR -> OnlineReplica
-      reassignedReplicas.foreach { replica =>
-        replicaStateMachine.handleStateChanges(Seq(new PartitionAndReplica(topicPartition, replica)), OnlineReplica)
-      }
-      //6. Set AR to RAR in memory.
-      //7. Send LeaderAndIsr request with a potential new leader (if current leader not in RAR) and
-      //   a new AR (using RAR) and same isr to every broker in RAR
-      moveReassignedPartitionLeaderIfRequired(topicPartition, reassignedPartitionContext)
-      //8. replicas in OAR - RAR -> Offline (force those replicas out of isr)
-      //9. replicas in OAR - RAR -> NonExistentReplica (force those replicas to be deleted)
-      stopOldReplicasOfReassignedPartition(topicPartition, reassignedPartitionContext, oldReplicas)
-      //10. Update AR in ZK with RAR.
-      updateAssignedReplicasForPartition(topicPartition, reassignedReplicas)
-      //11. Update the /admin/reassign_partitions path in ZK to remove this partition.
-      removePartitionsFromReassignedPartitions(Set(topicPartition))
-      //12. After electing leader, the replicas and isr information changes, so resend the update metadata request to every broker
+      // B1. replicas in AR -> OnlineReplica
+      replicaStateMachine.handleStateChanges(addingReplicas.map(PartitionAndReplica(topicPartition, _)), OnlineReplica)
+      // B2. Set RS = TRS, AR = [], RR = [] in memory.
+      val completedReassignment = ReplicaAssignment(reassignment.targetReplicas)
+      controllerContext.updatePartitionFullReplicaAssignment(topicPartition, completedReassignment)
+      // B3. Send LeaderAndIsr request with a potential new leader (if current leader not in TRS) and
+      //   a new RS (using TRS) and same isr to every broker in ORS + TRS or TRS
+      moveReassignedPartitionLeaderIfRequired(topicPartition, completedReassignment)
+      // B4. replicas in RR -> Offline (force those replicas out of isr)
+      // B5. replicas in RR -> NonExistentReplica (force those replicas to be deleted)
+      stopRemovedReplicasOfReassignedPartition(topicPartition, removingReplicas)
+      // B6. Update ZK with RS = TRS, AR = [], RR = [].
+      updateReplicaAssignmentForPartition(topicPartition, completedReassignment)
+      // B7. Remove the ISR reassign listener and maybe update the /admin/reassign_partitions path in ZK to remove this partition from it.
+      removePartitionFromReassigningPartitions(topicPartition, completedReassignment)
+      // B8. After electing a leader in B3, the replicas and isr information changes, so resend the update metadata request to every broker
       sendUpdateMetadataRequest(controllerContext.liveOrShuttingDownBrokerIds.toSeq, Set(topicPartition))
       // signal delete topic thread if reassignment for some partitions belonging to topics being deleted just completed
       topicDeletionManager.resumeDeletionForTopics(Set(topicPartition.topic))
@@ -573,116 +710,163 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
   }
 
   /**
-   * Trigger partition reassignment for the provided partitions if the assigned replicas are not the same as the
-   * reassigned replicas (as defined in `ControllerContext.partitionsBeingReassigned`) and if the topic has not been
-   * deleted.
+   * Update the current assignment state in Zookeeper and in memory. If a reassignment is already in
+   * progress, then the new reassignment will supplant it and some replicas will be shutdown.
+   *
+   * Note that due to the way we compute the original replica set, we cannot guarantee that a
+   * cancellation will restore the original replica order. Target replicas are always listed
+   * first in the replica set in the desired order, which means we have no way to get to the
+   * original order if the reassignment overlaps with the current assignment. For example,
+   * with an initial assignment of [1, 2, 3] and a reassignment of [3, 4, 2], then the replicas
+   * will be encoded as [3, 4, 2, 1] while the reassignment is in progress. If the reassignment
+   * is cancelled, there is no way to restore the original order.
+   *
+   * @param topicPartition The reassigning partition
+   * @param reassignment The new reassignment
+   *
+   * @return The updated assignment state
+   */
+  private def updateCurrentReassignment(topicPartition: TopicPartition, reassignment: ReplicaAssignment): Unit = {
+    val currentAssignment = controllerContext.partitionFullReplicaAssignment(topicPartition)
+
+    if (currentAssignment != reassignment) {
+      debug(s"Updating assignment of partition $topicPartition from $currentAssignment to $reassignment")
+
+      // U1. Update assignment state in zookeeper
+      updateReplicaAssignmentForPartition(topicPartition, reassignment)
+      // U2. Update assignment state in memory
+      controllerContext.updatePartitionFullReplicaAssignment(topicPartition, reassignment)
+
+      // If there is a reassignment already in progress, then some of the currently adding replicas
+      // may be eligible for immediate removal, in which case we need to stop the replicas.
+      val unneededReplicas = currentAssignment.replicas.diff(reassignment.replicas)
+      if (unneededReplicas.nonEmpty)
+        stopRemovedReplicasOfReassignedPartition(topicPartition, unneededReplicas)
+    }
+
+    val reassignIsrChangeHandler = new PartitionReassignmentIsrChangeHandler(eventManager, topicPartition)
+    zkClient.registerZNodeChangeHandler(reassignIsrChangeHandler)
+
+    controllerContext.partitionsBeingReassigned.add(topicPartition)
+  }
+
+  /**
+   * Trigger a partition reassignment provided that the topic exists and is not being deleted.
+   *
+   * This is called when a reassignment is initially received either through Zookeeper or through the
+   * AlterPartitionReassignments API
    *
-   * `partitionsBeingReassigned` must be populated with all partitions being reassigned before this method is invoked
-   * as explained in the method documentation of `removePartitionFromReassignedPartitions` (which is invoked by this
-   * method).
+   * The `partitionsBeingReassigned` field in the controller context will be updated by this
+   * call after the reassignment completes validation and is successfully stored in the topic
+   * assignment zNode.
    *
-   * @throws IllegalStateException if a partition is not in `partitionsBeingReassigned`
+   * @param reassignments The reassignments to begin processing
+   * @return A map of any errors in the reassignment. If the error is NONE for a given partition,
+   *         then the reassignment was submitted successfully.
    */
-  private def maybeTriggerPartitionReassignment(topicPartitions: Set[TopicPartition]) {
-    val partitionsToBeRemovedFromReassignment = scala.collection.mutable.Set.empty[TopicPartition]
-    topicPartitions.foreach { tp =>
-      if (topicDeletionManager.isTopicQueuedUpForDeletion(tp.topic)) {
+  private def maybeTriggerPartitionReassignment(reassignments: Map[TopicPartition, ReplicaAssignment]): Map[TopicPartition, ApiError] = {
+    reassignments.map { case (tp, reassignment) =>
+      val topic = tp.topic
+
+      val apiError = if (topicDeletionManager.isTopicQueuedUpForDeletion(topic)) {
         info(s"Skipping reassignment of $tp since the topic is currently being deleted")
-        partitionsToBeRemovedFromReassignment.add(tp)
+        new ApiError(Errors.UNKNOWN_TOPIC_OR_PARTITION, "The partition does not exist.")
       } else {
-        val reassignedPartitionContext = controllerContext.partitionsBeingReassigned.get(tp).getOrElse {
-          throw new IllegalStateException(s"Initiating reassign replicas for partition $tp not present in " +
-            s"partitionsBeingReassigned: ${controllerContext.partitionsBeingReassigned.mkString(", ")}")
-        }
-        val newReplicas = reassignedPartitionContext.newReplicas
-        val topic = tp.topic
         val assignedReplicas = controllerContext.partitionReplicaAssignment(tp)
         if (assignedReplicas.nonEmpty) {
-          if (assignedReplicas == newReplicas) {
-            info(s"Partition $tp to be reassigned is already assigned to replicas " +
-              s"${newReplicas.mkString(",")}. Ignoring request for partition reassignment.")
-            partitionsToBeRemovedFromReassignment.add(tp)
-          } else {
-            try {
-              info(s"Handling reassignment of partition $tp to new replicas ${newReplicas.mkString(",")}")
-              // first register ISR change listener
-              reassignedPartitionContext.registerReassignIsrChangeHandler(zkClient)
-              // mark topic ineligible for deletion for the partitions being reassigned
-              topicDeletionManager.markTopicIneligibleForDeletion(Set(topic))
-              onPartitionReassignment(tp, reassignedPartitionContext)
-            } catch {
-              case e: ControllerMovedException =>
-                error(s"Error completing reassignment of partition $tp because controller has moved to another broker", e)
-                throw e
-              case e: Throwable =>
-                error(s"Error completing reassignment of partition $tp", e)
-                // remove the partition from the admin path to unblock the admin client
-                partitionsToBeRemovedFromReassignment.add(tp)
-            }
+          try {
+            onPartitionReassignment(tp, reassignment)
+            ApiError.NONE
+          } catch {
+            case e: ControllerMovedException =>
+              info(s"Failed completing reassignment of partition $tp because controller has moved to another broker")
+              throw e
+            case e: Throwable =>
+              error(s"Error completing reassignment of partition $tp", e)
+              new ApiError(Errors.UNKNOWN_SERVER_ERROR)
           }
         } else {
-            error(s"Ignoring request to reassign partition $tp that doesn't exist.")
-            partitionsToBeRemovedFromReassignment.add(tp)
+            new ApiError(Errors.UNKNOWN_TOPIC_OR_PARTITION, "The partition does not exist.")
         }
       }
+
+      tp -> apiError
     }
-    removePartitionsFromReassignedPartitions(partitionsToBeRemovedFromReassignment)
   }
 
-  sealed trait ElectionType
-  object AutoTriggered extends ElectionType
-  object ZkTriggered extends ElectionType
-  object AdminClientTriggered extends ElectionType
-
   /**
-    * Attempt to elect the preferred replica as leader for each of the given partitions.
-    * @param partitions The partitions to have their preferred leader elected
-    * @param electionType The election type
-    * @return A map of failed elections where keys are partitions which had an error and the corresponding value is
-    *         the exception that was thrown.
+    * Attempt to elect a replica as leader for each of the given partitions.
+    * @param partitions The partitions to have a new leader elected
+    * @param electionType The type of election to perform
+    * @param electionTrigger The reason for tigger this election
+    * @return A map of failed and successful elections. The keys are the topic partitions and the corresponding values are
+    *         either the exception that was thrown or new leader & ISR.
     */
-  private def onPreferredReplicaElection(partitions: Set[TopicPartition],
-                                         electionType: ElectionType): Map[TopicPartition, Throwable] = {
-    info(s"Starting preferred replica leader election for partitions ${partitions.mkString(",")}")
+  private[this] def onReplicaElection(
+    partitions: Set[TopicPartition],
+    electionType: ElectionType,
+    electionTrigger: ElectionTrigger
+  ): Map[TopicPartition, Either[Throwable, LeaderAndIsr]] = {
+    info(s"Starting replica leader election ($electionType) for partitions ${partitions.mkString(",")} triggered by $electionTrigger")
     try {
-      val results = partitionStateMachine.handleStateChanges(partitions.toSeq, OnlinePartition,
-        Option(PreferredReplicaPartitionLeaderElectionStrategy))
-      if (electionType != AdminClientTriggered) {
-        results.foreach { case (tp, throwable) =>
-          if (throwable.isInstanceOf[ControllerMovedException]) {
-            error(s"Error completing preferred replica leader election for partition $tp because controller has moved to another broker.", throwable)
-            throw throwable
-          } else {
-            error(s"Error completing preferred replica leader election for partition $tp", throwable)
-          }
+      val strategy = electionType match {
+        case ElectionType.PREFERRED => PreferredReplicaPartitionLeaderElectionStrategy
+        case ElectionType.UNCLEAN =>
+          /* Let's be conservative and only trigger unclean election if the election type is unclean and it was
+           * triggered by the admin client
+           */
+          OfflinePartitionLeaderElectionStrategy(allowUnclean = electionTrigger == AdminClientTriggered)
+      }
+
+      val results = partitionStateMachine.handleStateChanges(
+        partitions.toSeq,
+        OnlinePartition,
+        Some(strategy)
+      )
+      if (electionTrigger != AdminClientTriggered) {
+        results.foreach {
+          case (tp, Left(throwable)) =>
+            if (throwable.isInstanceOf[ControllerMovedException]) {
+              info(s"Error completing replica leader election ($electionType) for partition $tp because controller has moved to another broker.", throwable)
+              throw throwable
+            } else {
+              error(s"Error completing replica leader election ($electionType) for partition $tp", throwable)
+            }
+          case (_, Right(_)) => // Ignored; No need to log or throw exception for the success cases
         }
       }
-      return results;
+
+      results
     } finally {
-      if (electionType != AdminClientTriggered)
-        removePartitionsFromPreferredReplicaElection(partitions, electionType == AutoTriggered)
+      if (electionTrigger != AdminClientTriggered) {
+        removePartitionsFromPreferredReplicaElection(partitions, electionTrigger == AutoTriggered)
+      }
     }
   }
 
-  private def initializeControllerContext() {
+  private def initializeControllerContext(): Unit = {
     // update controller cache with delete topic information
     val curBrokerAndEpochs = zkClient.getAllBrokerAndEpochsInCluster
     controllerContext.setLiveBrokerAndEpochs(curBrokerAndEpochs)
     info(s"Initialized broker epochs cache: ${controllerContext.liveBrokerIdAndEpochs}")
-    controllerContext.allTopics = zkClient.getAllTopicsInCluster.toSet
+
+    controllerContext.allTopics = zkClient.getAllTopicsInCluster
+    rearrangePartitionReplicaAssignmentForNewTopics(controllerContext.allTopics.toSet)
     registerPartitionModificationsHandlers(controllerContext.allTopics.toSeq)
-    zkClient.getReplicaAssignmentForTopics(controllerContext.allTopics.toSet).foreach {
-      case (topicPartition, assignedReplicas) => controllerContext.updatePartitionReplicaAssignment(topicPartition, assignedReplicas)
+    getReplicaAssignmentPolicyCompliant(controllerContext.allTopics.toSet).foreach {
+      case (topicPartition, replicaAssignment) =>
+        controllerContext.updatePartitionFullReplicaAssignment(topicPartition, replicaAssignment)
+        if (replicaAssignment.isBeingReassigned)
+          controllerContext.partitionsBeingReassigned.add(topicPartition)
     }
     controllerContext.partitionLeadershipInfo.clear()
     controllerContext.shuttingDownBrokerIds = mutable.Set.empty[Int]
     // register broker modifications handlers
-    registerBrokerModificationsHandler(controllerContext.liveBrokers.map(_.id))
+    registerBrokerModificationsHandler(controllerContext.liveOrShuttingDownBrokerIds)
     // update the leader and isr cache for all existing partitions from Zookeeper
     updateLeaderAndIsrCache()
     // start the channel manager
-    startChannelManager()
-    initializePartitionReassignment()
+    controllerChannelManager.startup()
     info(s"Currently active brokers in the cluster: ${controllerContext.liveBrokerIds}")
     info(s"Currently shutting brokers in the cluster: ${controllerContext.shuttingDownBrokerIds}")
     info(s"Current list of topics in the cluster: ${controllerContext.allTopics}")
@@ -708,14 +892,16 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
     pendingPreferredReplicaElections
   }
 
-  private def initializePartitionReassignment() {
-    // read the partitions being reassigned from zookeeper path /admin/reassign_partitions
-    val partitionsBeingReassigned = zkClient.getPartitionReassignment
-    info(s"Partitions being reassigned: $partitionsBeingReassigned")
-
-    controllerContext.partitionsBeingReassigned ++= partitionsBeingReassigned.iterator.map { case (tp, newReplicas) =>
-      val reassignIsrChangeHandler = new PartitionReassignmentIsrChangeHandler(this, eventManager, tp)
-      tp -> new ReassignedPartitionsContext(newReplicas, reassignIsrChangeHandler)
+  /**
+   * Initialize pending reassignments. This includes reassignments sent through /admin/reassign_partitions,
+   * which will supplant any API reassignments already in progress.
+   */
+  private def initializePartitionReassignments(): Unit = {
+    // New reassignments may have been submitted through Zookeeper while the controller was failing over
+    val zkPartitionsResumed = processZkPartitionReassignment()
+    // We may also have some API-based reassignments that need to be restarted
+    maybeResumeReassignments { (tp, _) =>
+      !zkPartitionsResumed.contains(tp)
     }
   }
 
@@ -725,122 +911,147 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
       val replicasForTopic = controllerContext.replicasForTopic(topic)
       replicasForTopic.exists(r => !controllerContext.isReplicaOnline(r.replica, r.topicPartition))
     }}
-    val topicsForWhichPartitionReassignmentIsInProgress = controllerContext.partitionsBeingReassigned.keySet.map(_.topic)
+    val topicsForWhichPartitionReassignmentIsInProgress = controllerContext.partitionsBeingReassigned.map(_.topic)
     val topicsIneligibleForDeletion = topicsWithOfflineReplicas | topicsForWhichPartitionReassignmentIsInProgress
     info(s"List of topics to be deleted: ${topicsToBeDeleted.mkString(",")}")
     info(s"List of topics ineligible for deletion: ${topicsIneligibleForDeletion.mkString(",")}")
     (topicsToBeDeleted, topicsIneligibleForDeletion)
   }
 
-  private def startChannelManager() {
-    controllerContext.controllerChannelManager = new ControllerChannelManager(controllerContext, config, time, metrics,
-      stateChangeLogger, threadNamePrefix)
-    controllerContext.controllerChannelManager.startup()
-  }
-
-  private def updateLeaderAndIsrCache(partitions: Seq[TopicPartition] = controllerContext.allPartitions.toSeq) {
+  private def updateLeaderAndIsrCache(partitions: Seq[TopicPartition] = controllerContext.allPartitions.toSeq): Unit = {
     val leaderIsrAndControllerEpochs = zkClient.getTopicPartitionStates(partitions)
     leaderIsrAndControllerEpochs.foreach { case (partition, leaderIsrAndControllerEpoch) =>
       controllerContext.partitionLeadershipInfo.put(partition, leaderIsrAndControllerEpoch)
     }
   }
 
-  private def areReplicasInIsr(partition: TopicPartition, replicas: Seq[Int]): Boolean = {
-    zkClient.getTopicPartitionStates(Seq(partition)).get(partition).exists { leaderIsrAndControllerEpoch =>
-      replicas.forall(leaderIsrAndControllerEpoch.leaderAndIsr.isr.contains)
+  private def isReassignmentComplete(partition: TopicPartition, assignment: ReplicaAssignment): Boolean = {
+    if (!assignment.isBeingReassigned) {
+      true
+    } else {
+      zkClient.getTopicPartitionStates(Seq(partition)).get(partition).exists { leaderIsrAndControllerEpoch =>
+        val isr = leaderIsrAndControllerEpoch.leaderAndIsr.isr.toSet
+        val targetReplicas = assignment.targetReplicas.toSet
+        targetReplicas.subsetOf(isr)
+      }
+    }
+  }
+
+  // Rearrange partition and replica assignment for new topics that get assigned to
+  // maintenance brokers that do not take new partitions
+  private def rearrangePartitionReplicaAssignmentForNewTopics(topics: Set[String]) {
+    try {
+      val noNewPartitionBrokers = partitionUnassignableBrokerIds
+      if (noNewPartitionBrokers.nonEmpty) {
+        val newTopics = zkClient.getPartitionNodeNonExistsTopics(topics.toSet)
+        val newTopicsToBeArranged = zkClient.getPartitionAssignmentForTopics(newTopics).filter {
+          case (_, partitionMap) =>
+            partitionMap.exists {
+              case (_, assignedReplicas) =>
+                assignedReplicas.replicas.intersect(noNewPartitionBrokers).nonEmpty
+            }
+        }
+        newTopicsToBeArranged.foreach {
+          case (topic, partitionMap) =>
+            val numPartitions = partitionMap.size
+            val numReplica = partitionMap.head._2.replicas.size
+            val brokers = controllerContext.liveOrShuttingDownBrokers.map { b => kafka.admin.BrokerMetadata(b.id, b.rack) }.toSeq
+
+            val replicaAssignment = adminZkClient.assignReplicasToAvailableBrokers(brokers, noNewPartitionBrokers.toSet, numPartitions, numReplica)
+            adminZkClient.writeTopicPartitionAssignment(topic, replicaAssignment.mapValues(ReplicaAssignment(_)).toMap, true)
+            info(s"Rearrange partition and replica assignment for topic [$topic]")
+        }
+      }
+    } catch {
+      case e =>
+        error("Error during rearranging partition and replica assignment for new topics for maintenance brokers :" + e.getMessage)
     }
   }
 
   private def moveReassignedPartitionLeaderIfRequired(topicPartition: TopicPartition,
-                                                      reassignedPartitionContext: ReassignedPartitionsContext) {
-    val reassignedReplicas = reassignedPartitionContext.newReplicas
+                                                      newAssignment: ReplicaAssignment): Unit = {
+    val reassignedReplicas = newAssignment.replicas
     val currentLeader = controllerContext.partitionLeadershipInfo(topicPartition).leaderAndIsr.leader
-    // change the assigned replica list to just the reassigned replicas in the cache so it gets sent out on the LeaderAndIsr
-    // request to the current or new leader. This will prevent it from adding the old replicas to the ISR
-    val oldAndNewReplicas = controllerContext.partitionReplicaAssignment(topicPartition)
-    controllerContext.updatePartitionReplicaAssignment(topicPartition, reassignedReplicas)
-    if (!reassignedPartitionContext.newReplicas.contains(currentLeader)) {
+
+    if (!reassignedReplicas.contains(currentLeader)) {
       info(s"Leader $currentLeader for partition $topicPartition being reassigned, " +
         s"is not in the new list of replicas ${reassignedReplicas.mkString(",")}. Re-electing leader")
       // move the leader to one of the alive and caught up new replicas
-      partitionStateMachine.handleStateChanges(Seq(topicPartition), OnlinePartition, Option(ReassignPartitionLeaderElectionStrategy))
+      partitionStateMachine.handleStateChanges(Seq(topicPartition), OnlinePartition, Some(ReassignPartitionLeaderElectionStrategy))
+    } else if (controllerContext.isReplicaOnline(currentLeader, topicPartition)) {
+      info(s"Leader $currentLeader for partition $topicPartition being reassigned, " +
+        s"is already in the new list of replicas ${reassignedReplicas.mkString(",")} and is alive")
+      // shrink replication factor and update the leader epoch in zookeeper to use on the next LeaderAndIsrRequest
+      updateLeaderEpochAndSendRequest(topicPartition, newAssignment)
     } else {
-      // check if the leader is alive or not
-      if (controllerContext.isReplicaOnline(currentLeader, topicPartition)) {
-        info(s"Leader $currentLeader for partition $topicPartition being reassigned, " +
-          s"is already in the new list of replicas ${reassignedReplicas.mkString(",")} and is alive")
-        // shrink replication factor and update the leader epoch in zookeeper to use on the next LeaderAndIsrRequest
-        updateLeaderEpochAndSendRequest(topicPartition, oldAndNewReplicas, reassignedReplicas)
-      } else {
-        info(s"Leader $currentLeader for partition $topicPartition being reassigned, " +
-          s"is already in the new list of replicas ${reassignedReplicas.mkString(",")} but is dead")
-        partitionStateMachine.handleStateChanges(Seq(topicPartition), OnlinePartition, Option(ReassignPartitionLeaderElectionStrategy))
-      }
+      info(s"Leader $currentLeader for partition $topicPartition being reassigned, " +
+        s"is already in the new list of replicas ${reassignedReplicas.mkString(",")} but is dead")
+      partitionStateMachine.handleStateChanges(Seq(topicPartition), OnlinePartition, Some(ReassignPartitionLeaderElectionStrategy))
     }
   }
 
-  private def stopOldReplicasOfReassignedPartition(topicPartition: TopicPartition,
-                                                   reassignedPartitionContext: ReassignedPartitionsContext,
-                                                   oldReplicas: Set[Int]) {
+  private def stopRemovedReplicasOfReassignedPartition(topicPartition: TopicPartition,
+                                                       removedReplicas: Seq[Int]): Unit = {
     // first move the replica to offline state (the controller removes it from the ISR)
-    val replicasToBeDeleted = oldReplicas.map(PartitionAndReplica(topicPartition, _))
-    replicaStateMachine.handleStateChanges(replicasToBeDeleted.toSeq, OfflineReplica)
+    val replicasToBeDeleted = removedReplicas.map(PartitionAndReplica(topicPartition, _))
+    replicaStateMachine.handleStateChanges(replicasToBeDeleted, OfflineReplica)
     // send stop replica command to the old replicas
-    replicaStateMachine.handleStateChanges(replicasToBeDeleted.toSeq, ReplicaDeletionStarted)
+    replicaStateMachine.handleStateChanges(replicasToBeDeleted, ReplicaDeletionStarted)
     // TODO: Eventually partition reassignment could use a callback that does retries if deletion failed
-    replicaStateMachine.handleStateChanges(replicasToBeDeleted.toSeq, ReplicaDeletionSuccessful)
-    replicaStateMachine.handleStateChanges(replicasToBeDeleted.toSeq, NonExistentReplica)
+    replicaStateMachine.handleStateChanges(replicasToBeDeleted, ReplicaDeletionSuccessful)
+    replicaStateMachine.handleStateChanges(replicasToBeDeleted, NonExistentReplica)
   }
 
-  private def updateAssignedReplicasForPartition(partition: TopicPartition,
-                                                 replicas: Seq[Int]) {
-    controllerContext.updatePartitionReplicaAssignment(partition, replicas)
-    val setDataResponse = zkClient.setTopicAssignmentRaw(partition.topic, controllerContext.partitionReplicaAssignmentForTopic(partition.topic), controllerContext.epochZkVersion)
+  private def updateReplicaAssignmentForPartition(topicPartition: TopicPartition, assignment: ReplicaAssignment): Unit = {
+    var topicAssignment = controllerContext.partitionFullReplicaAssignmentForTopic(topicPartition.topic)
+    topicAssignment += topicPartition -> assignment
+
+    val setDataResponse = zkClient.setTopicAssignmentRaw(topicPartition.topic, topicAssignment, controllerContext.epochZkVersion)
     setDataResponse.resultCode match {
       case Code.OK =>
-        info(s"Updated assigned replicas for partition $partition being reassigned to ${replicas.mkString(",")}")
-        // update the assigned replica list after a successful zookeeper write
-        controllerContext.updatePartitionReplicaAssignment(partition, replicas)
-      case Code.NONODE => throw new IllegalStateException(s"Topic ${partition.topic} doesn't exist")
+        info(s"Successfully updated assignment of partition $topicPartition to $assignment")
+      case Code.NONODE =>
+        throw new IllegalStateException(s"Failed to update assignment for $topicPartition since the topic " +
+          "has no current assignment")
       case _ => throw new KafkaException(setDataResponse.resultException.get)
     }
   }
 
-  private def startNewReplicasForReassignedPartition(topicPartition: TopicPartition,
-                                                     reassignedPartitionContext: ReassignedPartitionsContext,
-                                                     newReplicas: Set[Int]) {
+  private def startNewReplicasForReassignedPartition(topicPartition: TopicPartition, newReplicas: Seq[Int]): Unit = {
     // send the start replica request to the brokers in the reassigned replicas list that are not in the assigned
     // replicas list
     newReplicas.foreach { replica =>
-      replicaStateMachine.handleStateChanges(Seq(new PartitionAndReplica(topicPartition, replica)), NewReplica)
+      replicaStateMachine.handleStateChanges(Seq(PartitionAndReplica(topicPartition, replica)), NewReplica)
     }
   }
 
-  private def updateLeaderEpochAndSendRequest(partition: TopicPartition, replicasToReceiveRequest: Seq[Int], newAssignedReplicas: Seq[Int]) {
+  private def updateLeaderEpochAndSendRequest(topicPartition: TopicPartition,
+                                              assignment: ReplicaAssignment): Unit = {
     val stateChangeLog = stateChangeLogger.withControllerEpoch(controllerContext.epoch)
-    updateLeaderEpoch(partition) match {
+    updateLeaderEpoch(topicPartition) match {
       case Some(updatedLeaderIsrAndControllerEpoch) =>
         try {
           brokerRequestBatch.newBatch()
-          brokerRequestBatch.addLeaderAndIsrRequestForBrokers(replicasToReceiveRequest, partition,
-            updatedLeaderIsrAndControllerEpoch, newAssignedReplicas, isNew = false)
+          brokerRequestBatch.addLeaderAndIsrRequestForBrokers(assignment.replicas, topicPartition,
+            updatedLeaderIsrAndControllerEpoch, assignment, isNew = false)
           brokerRequestBatch.sendRequestsToBrokers(controllerContext.epoch)
         } catch {
           case e: IllegalStateException =>
             handleIllegalState(e)
         }
-        stateChangeLog.trace(s"Sent LeaderAndIsr request $updatedLeaderIsrAndControllerEpoch with new assigned replica " +
-          s"list ${newAssignedReplicas.mkString(",")} to leader ${updatedLeaderIsrAndControllerEpoch.leaderAndIsr.leader} " +
-          s"for partition being reassigned $partition")
+        stateChangeLog.trace(s"Sent LeaderAndIsr request $updatedLeaderIsrAndControllerEpoch with " +
+          s"new replica assignment $assignment to leader ${updatedLeaderIsrAndControllerEpoch.leaderAndIsr.leader} " +
+          s"for partition being reassigned $topicPartition")
+
       case None => // fail the reassignment
-        stateChangeLog.error("Failed to send LeaderAndIsr request with new assigned replica list " +
-          s"${newAssignedReplicas.mkString( ",")} to leader for partition being reassigned $partition")
+        stateChangeLog.error(s"Failed to send LeaderAndIsr request with new replica assignment " +
+          s"$assignment to leader for partition being reassigned $topicPartition")
     }
   }
 
   private def registerPartitionModificationsHandlers(topics: Seq[String]) = {
     topics.foreach { topic =>
-      val partitionModificationsHandler = new PartitionModificationsHandler(this, eventManager, topic)
+      val partitionModificationsHandler = new PartitionModificationsHandler(eventManager, topic)
       partitionModificationsHandlers.put(topic, partitionModificationsHandler)
     }
     partitionModificationsHandlers.values.foreach(zkClient.registerZNodeChangeHandler)
@@ -852,46 +1063,57 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
     }
   }
 
-  private def unregisterPartitionReassignmentIsrChangeHandlers() {
-    controllerContext.partitionsBeingReassigned.values.foreach(_.unregisterReassignIsrChangeHandler(zkClient))
+  private def unregisterPartitionReassignmentIsrChangeHandlers(): Unit = {
+    controllerContext.partitionsBeingReassigned.foreach { tp =>
+      val path = TopicPartitionStateZNode.path(tp)
+      zkClient.unregisterZNodeChangeHandler(path)
+    }
+  }
+
+  private def removePartitionFromReassigningPartitions(topicPartition: TopicPartition,
+                                                       assignment: ReplicaAssignment): Unit = {
+    if (controllerContext.partitionsBeingReassigned.contains(topicPartition)) {
+      val path = TopicPartitionStateZNode.path(topicPartition)
+      zkClient.unregisterZNodeChangeHandler(path)
+      maybeRemoveFromZkReassignment((tp, replicas) => tp == topicPartition && replicas == assignment.replicas)
+      controllerContext.partitionsBeingReassigned.remove(topicPartition)
+    } else {
+      throw new IllegalStateException("Cannot remove a reassigning partition because it is not present in memory")
+    }
   }
 
   /**
-   * Remove partition from partitions being reassigned in ZooKeeper and ControllerContext. If the partition reassignment
-   * is complete (i.e. there is no other partition with a reassignment in progress), the reassign_partitions znode
-   * is deleted.
+   * Remove partitions from an active zk-based reassignment (if one exists).
    *
-   * `ControllerContext.partitionsBeingReassigned` must be populated with all partitions being reassigned before this
-   * method is invoked to avoid premature deletion of the `reassign_partitions` znode.
+   * @param shouldRemoveReassignment Predicate indicating which partition reassignments should be removed
    */
-  private def removePartitionsFromReassignedPartitions(partitionsToBeRemoved: Set[TopicPartition]) {
-    partitionsToBeRemoved.map(controllerContext.partitionsBeingReassigned).foreach { reassignContext =>
-      reassignContext.unregisterReassignIsrChangeHandler(zkClient)
-    }
-
-    val updatedPartitionsBeingReassigned = controllerContext.partitionsBeingReassigned -- partitionsToBeRemoved
+  private[controller] def maybeRemoveFromZkReassignment(shouldRemoveReassignment: (TopicPartition, Seq[Int]) => Boolean): Unit = {
+    if (!zkClient.reassignPartitionsInProgress())
+      return
 
-    info(s"Removing partitions $partitionsToBeRemoved from the list of reassigned partitions in zookeeper")
+    val reassigningPartitions = zkClient.getPartitionReassignment()
+    val (removingPartitions, updatedPartitionsBeingReassigned) = reassigningPartitions.partition { case (tp, replicas) =>
+      shouldRemoveReassignment(tp, replicas)
+    }
+    info(s"Removing partitions $removingPartitions from the list of reassigned partitions in zookeeper")
 
     // write the new list to zookeeper
     if (updatedPartitionsBeingReassigned.isEmpty) {
       info(s"No more partitions need to be reassigned. Deleting zk path ${ReassignPartitionsZNode.path}")
       zkClient.deletePartitionReassignment(controllerContext.epochZkVersion)
       // Ensure we detect future reassignments
-      eventManager.put(PartitionReassignment)
+      eventManager.put(ZkPartitionReassignment)
     } else {
-      val reassignment = updatedPartitionsBeingReassigned.mapValues(_.newReplicas)
-      try zkClient.setOrCreatePartitionReassignment(reassignment, controllerContext.epochZkVersion)
-      catch {
+      try {
+        zkClient.setOrCreatePartitionReassignment(updatedPartitionsBeingReassigned, controllerContext.epochZkVersion)
+      } catch {
         case e: KeeperException => throw new AdminOperationException(e)
       }
     }
-
-    controllerContext.partitionsBeingReassigned --= partitionsToBeRemoved
   }
 
   private def removePartitionsFromPreferredReplicaElection(partitionsToBeRemoved: Set[TopicPartition],
-                                                           isTriggeredByAutoRebalance : Boolean) {
+                                                           isTriggeredByAutoRebalance : Boolean): Unit = {
     for (partition <- partitionsToBeRemoved) {
       // check the status
       val currentLeader = controllerContext.partitionLeadershipInfo(partition).leaderAndIsr.leader
@@ -906,22 +1128,17 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
     if (!isTriggeredByAutoRebalance) {
       zkClient.deletePreferredReplicaElection(controllerContext.epochZkVersion)
       // Ensure we detect future preferred replica leader elections
-      eventManager.put(PreferredReplicaLeaderElection(None))
+      eventManager.put(ReplicaLeaderElection(None, ElectionType.PREFERRED, ZkTriggered))
     }
   }
 
-  private[controller] def sendRequest(brokerId: Int, apiKey: ApiKeys, request: AbstractControlRequest.Builder[_ <: AbstractControlRequest],
-                                      callback: AbstractResponse => Unit = null) = {
-    controllerContext.controllerChannelManager.sendRequest(brokerId, apiKey, request, callback)
-  }
-
   /**
    * Send the leader information for selected partitions to selected brokers so that they can correctly respond to
    * metadata requests
    *
    * @param brokers The brokers that the update metadata request should be sent to
    */
-  private[controller] def sendUpdateMetadataRequest(brokers: Seq[Int], partitions: Set[TopicPartition]) {
+  private[controller] def sendUpdateMetadataRequest(brokers: Seq[Int], partitions: Set[TopicPartition]): Unit = {
     try {
       brokerRequestBatch.newBatch()
       brokerRequestBatch.addUpdateMetadataRequestForBrokers(brokers, partitions)
@@ -956,16 +1173,19 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
           // assigned replica list
           val newLeaderAndIsr = leaderAndIsr.newEpochAndZkVersion
           // update the new leadership decision in zookeeper or retry
-          val UpdateLeaderAndIsrResult(successfulUpdates, _, failedUpdates) =
+          val UpdateLeaderAndIsrResult(finishedUpdates, _) =
             zkClient.updateLeaderAndIsr(immutable.Map(partition -> newLeaderAndIsr), epoch, controllerContext.epochZkVersion)
-          if (successfulUpdates.contains(partition)) {
-            val finalLeaderAndIsr = successfulUpdates(partition)
-            finalLeaderIsrAndControllerEpoch = Some(LeaderIsrAndControllerEpoch(finalLeaderAndIsr, epoch))
-            info(s"Updated leader epoch for partition $partition to ${finalLeaderAndIsr.leaderEpoch}")
-            true
-          } else if (failedUpdates.contains(partition)) {
-            throw failedUpdates(partition)
-          } else false
+
+          finishedUpdates.get(partition) match {
+            case Some(Right(leaderAndIsr)) =>
+              val leaderIsrAndControllerEpoch = LeaderIsrAndControllerEpoch(leaderAndIsr, epoch)
+              controllerContext.partitionLeadershipInfo.put(partition, leaderIsrAndControllerEpoch)
+              finalLeaderIsrAndControllerEpoch = Some(leaderIsrAndControllerEpoch)
+              info(s"Updated leader epoch for partition $partition to ${leaderAndIsr.leaderEpoch}")
+              true
+            case Some(Left(e)) => throw e
+            case None => false
+          }
         case None =>
           throw new IllegalStateException(s"Cannot update leader epoch for partition $partition as " +
             "leaderAndIsr path is empty. This could mean we somehow tried to reassign a partition that doesn't exist")
@@ -1005,185 +1225,159 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
           controllerContext.partitionsBeingReassigned.isEmpty &&
           !topicDeletionManager.isTopicQueuedUpForDeletion(tp.topic) &&
           controllerContext.allTopics.contains(tp.topic))
-        onPreferredReplicaElection(candidatePartitions.toSet, AutoTriggered)
+        onReplicaElection(candidatePartitions.toSet, ElectionType.PREFERRED, AutoTriggered)
       }
     }
   }
 
-  case object AutoPreferredReplicaLeaderElection extends ControllerEvent {
-
-    def state = ControllerState.AutoLeaderBalance
-
-    override def process(): Unit = {
-      if (!isActive) return
-      try {
-        checkAndTriggerAutoLeaderRebalance()
-      } finally {
-        scheduleAutoLeaderRebalanceTask(delay = config.leaderImbalanceCheckIntervalSeconds, unit = TimeUnit.SECONDS)
-      }
+  private def processAutoPreferredReplicaLeaderElection(): Unit = {
+    if (!isActive) return
+    try {
+      info("Processing automatic preferred replica leader election")
+      checkAndTriggerAutoLeaderRebalance()
+    } finally {
+      scheduleAutoLeaderRebalanceTask(delay = config.leaderImbalanceCheckIntervalSeconds, unit = TimeUnit.SECONDS)
     }
   }
 
-  case object UncleanLeaderElectionEnable extends ControllerEvent {
-
-    def state = ControllerState.UncleanLeaderElectionEnable
-
-    override def process(): Unit = {
-      if (!isActive) return
-      partitionStateMachine.triggerOnlinePartitionStateChange()
-    }
+  private def processUncleanLeaderElectionEnable(): Unit = {
+    if (!isActive) return
+    info("Unclean leader election has been enabled by default")
+    partitionStateMachine.triggerOnlinePartitionStateChange()
   }
 
-  case class TopicUncleanLeaderElectionEnable(topic: String) extends ControllerEvent {
-
-    def state = ControllerState.TopicUncleanLeaderElectionEnable
-
-    override def process(): Unit = {
-      if (!isActive) return
-      partitionStateMachine.triggerOnlinePartitionStateChange(topic)
-    }
+  private def processTopicUncleanLeaderElectionEnable(topic: String): Unit = {
+    if (!isActive) return
+    info(s"Unclean leader election has been enabled for topic $topic")
+    partitionStateMachine.triggerOnlinePartitionStateChange(topic)
   }
 
-  case class ControlledShutdown(id: Int, brokerEpoch: Long, controlledShutdownCallback: Try[Set[TopicPartition]] => Unit) extends PreemptableControllerEvent {
-
-    def state = ControllerState.ControlledShutdown
-
-    override def handlePreempt(): Unit = {
-      controlledShutdownCallback(Failure(new ControllerMovedException("Controller moved to another broker")))
-    }
-
-    override def handleProcess(): Unit = {
-      val controlledShutdownResult = Try { doControlledShutdown(id) }
-      controlledShutdownCallback(controlledShutdownResult)
-    }
+  private def preemptControlledShutdown(id: Int, brokerEpoch: Long, controlledShutdownCallback: Try[Set[TopicPartition]] => Unit): Unit = {
+    controlledShutdownCallback(Failure(new ControllerMovedException("Controller moved to another broker")))
+  }
 
-    private def doControlledShutdown(id: Int): Set[TopicPartition] = {
-      if (!isActive) {
-        throw new ControllerMovedException("Controller moved to another broker. Aborting controlled shutdown")
-      }
+  private def processControlledShutdown(id: Int, brokerEpoch: Long, controlledShutdownCallback: Try[Set[TopicPartition]] => Unit): Unit = {
+    val controlledShutdownResult = Try { doControlledShutdown(id, brokerEpoch) }
+    controlledShutdownCallback(controlledShutdownResult)
+  }
 
-      // broker epoch in the request is unknown if the controller hasn't been upgraded to use KIP-380
-      // so we will keep the previous behavior and don't reject the request
-      if (brokerEpoch != AbstractControlRequest.UNKNOWN_BROKER_EPOCH) {
-        val cachedBrokerEpoch = controllerContext.liveBrokerIdAndEpochs(id)
-        if (brokerEpoch < cachedBrokerEpoch) {
-          val stateBrokerEpochErrorMessage = "Received controlled shutdown request from an old broker epoch " +
-            s"$brokerEpoch for broker $id. Current broker epoch is $cachedBrokerEpoch."
-          info(stateBrokerEpochErrorMessage)
-          throw new StaleBrokerEpochException(stateBrokerEpochErrorMessage)
-        }
+  private def doControlledShutdown(id: Int, brokerEpoch: Long): Set[TopicPartition] = {
+    if (!isActive) {
+      throw new ControllerMovedException("Controller moved to another broker. Aborting controlled shutdown")
+    }
+
+    // broker epoch in the request is unknown if the controller hasn't been upgraded to use KIP-380
+    // so we will keep the previous behavior and don't reject the request
+    if (brokerEpoch != AbstractControlRequest.UNKNOWN_BROKER_EPOCH) {
+      val cachedBrokerEpoch = controllerContext.liveBrokerIdAndEpochs(id)
+      if (brokerEpoch < cachedBrokerEpoch) {
+        val stateBrokerEpochErrorMessage = "Received controlled shutdown request from an old broker epoch " +
+          s"$brokerEpoch for broker $id. Current broker epoch is $cachedBrokerEpoch."
+        info(stateBrokerEpochErrorMessage)
+        throw new StaleBrokerEpochException(stateBrokerEpochErrorMessage)
       }
+    }
 
-      info(s"Shutting down broker $id")
+    info(s"Shutting down broker $id")
 
-      if (!controllerContext.liveOrShuttingDownBrokerIds.contains(id))
-        throw new BrokerNotAvailableException(s"Broker id $id does not exist.")
+    if (!controllerContext.liveOrShuttingDownBrokerIds.contains(id))
+      throw new BrokerNotAvailableException(s"Broker id $id does not exist.")
 
-      controllerContext.shuttingDownBrokerIds.add(id)
-      debug(s"All shutting down brokers: ${controllerContext.shuttingDownBrokerIds.mkString(",")}")
-      debug(s"Live brokers: ${controllerContext.liveBrokerIds.mkString(",")}")
+    controllerContext.shuttingDownBrokerIds.add(id)
+    debug(s"All shutting down brokers: ${controllerContext.shuttingDownBrokerIds.mkString(",")}")
+    debug(s"Live brokers: ${controllerContext.liveBrokerIds.mkString(",")}")
 
-      val partitionsToActOn = controllerContext.partitionsOnBroker(id).filter { partition =>
-        controllerContext.partitionReplicaAssignment(partition).size > 1 &&
-          controllerContext.partitionLeadershipInfo.contains(partition) &&
-          !topicDeletionManager.isTopicQueuedUpForDeletion(partition.topic)
-      }
-      val (partitionsLedByBroker, partitionsFollowedByBroker) = partitionsToActOn.partition { partition =>
-        controllerContext.partitionLeadershipInfo(partition).leaderAndIsr.leader == id
-      }
-      partitionStateMachine.handleStateChanges(partitionsLedByBroker.toSeq, OnlinePartition, Option(ControlledShutdownPartitionLeaderElectionStrategy))
-      try {
-        brokerRequestBatch.newBatch()
-        partitionsFollowedByBroker.foreach { partition =>
-          brokerRequestBatch.addStopReplicaRequestForBrokers(Seq(id), partition, deletePartition = false,
-            (_, _) => ())
-        }
-        brokerRequestBatch.sendRequestsToBrokers(epoch)
-      } catch {
-        case e: IllegalStateException =>
-          handleIllegalState(e)
-      }
-      // If the broker is a follower, updates the isr in ZK and notifies the current leader
-      replicaStateMachine.handleStateChanges(partitionsFollowedByBroker.map(partition =>
-        PartitionAndReplica(partition, id)).toSeq, OfflineReplica)
-      def replicatedPartitionsBrokerLeads() = {
-        trace(s"All leaders = ${controllerContext.partitionLeadershipInfo.mkString(",")}")
-        controllerContext.partitionLeadershipInfo.filter {
-          case (topicPartition, leaderIsrAndControllerEpoch) =>
-            !topicDeletionManager.isTopicQueuedUpForDeletion(topicPartition.topic) &&
-              leaderIsrAndControllerEpoch.leaderAndIsr.leader == id &&
-              controllerContext.partitionReplicaAssignment(topicPartition).size > 1
-        }.keys
+    val partitionsToActOn = controllerContext.partitionsOnBroker(id).filter { partition =>
+      controllerContext.partitionReplicaAssignment(partition).size > 1 &&
+        controllerContext.partitionLeadershipInfo.contains(partition) &&
+        !topicDeletionManager.isTopicQueuedUpForDeletion(partition.topic)
+    }
+    val (partitionsLedByBroker, partitionsFollowedByBroker) = partitionsToActOn.partition { partition =>
+      controllerContext.partitionLeadershipInfo(partition).leaderAndIsr.leader == id
+    }
+    partitionStateMachine.handleStateChanges(partitionsLedByBroker.toSeq, OnlinePartition, Some(ControlledShutdownPartitionLeaderElectionStrategy))
+    try {
+      brokerRequestBatch.newBatch()
+      partitionsFollowedByBroker.foreach { partition =>
+        brokerRequestBatch.addStopReplicaRequestForBrokers(Seq(id), partition, deletePartition = false)
       }
-      replicatedPartitionsBrokerLeads().toSet
+      brokerRequestBatch.sendRequestsToBrokers(epoch)
+    } catch {
+      case e: IllegalStateException =>
+        handleIllegalState(e)
     }
+    // If the broker is a follower, updates the isr in ZK and notifies the current leader
+    replicaStateMachine.handleStateChanges(partitionsFollowedByBroker.map(partition =>
+      PartitionAndReplica(partition, id)).toSeq, OfflineReplica)
+    def replicatedPartitionsBrokerLeads() = {
+      trace(s"All leaders = ${controllerContext.partitionLeadershipInfo.mkString(",")}")
+      controllerContext.partitionLeadershipInfo.filter {
+        case (topicPartition, leaderIsrAndControllerEpoch) =>
+          !topicDeletionManager.isTopicQueuedUpForDeletion(topicPartition.topic) &&
+            leaderIsrAndControllerEpoch.leaderAndIsr.leader == id &&
+            controllerContext.partitionReplicaAssignment(topicPartition).size > 1
+      }.keys
+    }
+    replicatedPartitionsBrokerLeads().toSet
   }
 
-  case class LeaderAndIsrResponseReceived(LeaderAndIsrResponseObj: AbstractResponse, brokerId: Int) extends ControllerEvent {
-
-    def state = ControllerState.LeaderAndIsrResponseReceived
+  private def processLeaderAndIsrResponseReceived(leaderAndIsrResponseObj: AbstractResponse, brokerId: Int): Unit = {
+    if (!isActive) return
+    val leaderAndIsrResponse = leaderAndIsrResponseObj.asInstanceOf[LeaderAndIsrResponse]
 
-    override def process(): Unit = {
-      import JavaConverters._
-      if (!isActive) return
-      val leaderAndIsrResponse = LeaderAndIsrResponseObj.asInstanceOf[LeaderAndIsrResponse]
+    if (leaderAndIsrResponse.error != Errors.NONE) {
+      stateChangeLogger.error(s"Received error in LeaderAndIsr response $leaderAndIsrResponse from broker $brokerId")
+      return
+    }
 
-      if (leaderAndIsrResponse.error != Errors.NONE) {
-        stateChangeLogger.error(s"Received error in LeaderAndIsr response $leaderAndIsrResponse from broker $brokerId")
-        return
-      }
+    val offlineReplicas = new ArrayBuffer[TopicPartition]()
+    val onlineReplicas = new ArrayBuffer[TopicPartition]()
 
-      val offlineReplicas = leaderAndIsrResponse.responses.asScala.collect {
-        case (tp, error) if error == Errors.KAFKA_STORAGE_ERROR => tp
-      }
-      val onlineReplicas = leaderAndIsrResponse.responses.asScala.collect {
-        case (tp, error) if error == Errors.NONE => tp
-      }
-      val previousOfflineReplicas = controllerContext.replicasOnOfflineDirs.getOrElse(brokerId, Set.empty[TopicPartition])
-      val currentOfflineReplicas = previousOfflineReplicas -- onlineReplicas ++ offlineReplicas
-      controllerContext.replicasOnOfflineDirs.put(brokerId, currentOfflineReplicas)
-      val newOfflineReplicas = currentOfflineReplicas -- previousOfflineReplicas
-
-      if (newOfflineReplicas.nonEmpty) {
-        stateChangeLogger.info(s"Mark replicas ${newOfflineReplicas.mkString(",")} on broker $brokerId as offline")
-        onReplicasBecomeOffline(newOfflineReplicas.map(PartitionAndReplica(_, brokerId)))
-      }
+    leaderAndIsrResponse.partitions.asScala.foreach { partition =>
+      val tp = new TopicPartition(partition.topicName, partition.partitionIndex)
+      if (partition.errorCode == Errors.KAFKA_STORAGE_ERROR.code)
+        offlineReplicas += tp
+      else if (partition.errorCode == Errors.NONE.code)
+        onlineReplicas += tp
     }
-  }
-
-  case class TopicDeletionStopReplicaResponseReceived(stopReplicaResponseObj: AbstractResponse, replicaId: Int) extends ControllerEvent {
 
-    def state = ControllerState.TopicDeletion
+    val previousOfflineReplicas = controllerContext.replicasOnOfflineDirs.getOrElse(brokerId, Set.empty[TopicPartition])
+    val currentOfflineReplicas = previousOfflineReplicas -- onlineReplicas ++ offlineReplicas
+    controllerContext.replicasOnOfflineDirs.put(brokerId, currentOfflineReplicas)
+    val newOfflineReplicas = currentOfflineReplicas -- previousOfflineReplicas
 
-    override def process(): Unit = {
-      import JavaConverters._
-      if (!isActive) return
-      val stopReplicaResponse = stopReplicaResponseObj.asInstanceOf[StopReplicaResponse]
-      debug(s"Delete topic callback invoked for $stopReplicaResponse")
-      val responseMap = stopReplicaResponse.responses.asScala
-      val partitionsInError =
-        if (stopReplicaResponse.error != Errors.NONE) responseMap.keySet
-        else responseMap.filter { case (_, error) => error != Errors.NONE }.keySet
-      val replicasInError = partitionsInError.map(PartitionAndReplica(_, replicaId))
-      // move all the failed replicas to ReplicaDeletionIneligible
-      topicDeletionManager.failReplicaDeletion(replicasInError)
-      if (replicasInError.size != responseMap.size) {
-        // some replicas could have been successfully deleted
-        val deletedReplicas = responseMap.keySet -- partitionsInError
-        topicDeletionManager.completeReplicaDeletion(deletedReplicas.map(PartitionAndReplica(_, replicaId)))
-      }
+    if (newOfflineReplicas.nonEmpty) {
+      stateChangeLogger.info(s"Mark replicas ${newOfflineReplicas.mkString(",")} on broker $brokerId as offline")
+      onReplicasBecomeOffline(newOfflineReplicas.map(PartitionAndReplica(_, brokerId)))
     }
   }
 
-  case object Startup extends ControllerEvent {
+  private def processTopicDeletionStopReplicaResponseReceived(replicaId: Int,
+                                                              requestError: Errors,
+                                                              partitionErrors: Map[TopicPartition, Errors]): Unit = {
+    if (!isActive) return
+    debug(s"Delete topic callback invoked on StopReplica response received from broker $replicaId: " +
+      s"request error = $requestError, partition errors = $partitionErrors")
 
-    def state = ControllerState.ControllerChange
+    val partitionsInError = if (requestError != Errors.NONE)
+      partitionErrors.keySet
+    else
+      partitionErrors.filter { case (_, error) => error != Errors.NONE }.keySet
 
-    override def process(): Unit = {
-      zkClient.registerZNodeChangeHandlerAndCheckExistence(controllerChangeHandler)
-      elect()
+    val replicasInError = partitionsInError.map(PartitionAndReplica(_, replicaId))
+    // move all the failed replicas to ReplicaDeletionIneligible
+    topicDeletionManager.failReplicaDeletion(replicasInError)
+    if (replicasInError.size != partitionErrors.size) {
+      // some replicas could have been successfully deleted
+      val deletedReplicas = partitionErrors.keySet -- partitionsInError
+      topicDeletionManager.completeReplicaDeletion(deletedReplicas.map(PartitionAndReplica(_, replicaId)))
     }
+  }
 
+  private def processStartup(): Unit = {
+    zkClient.registerZNodeChangeHandlerAndCheckExistence(controllerChangeHandler)
+    zkClient.registerZNodeChildChangeHandler(preferredControllerChangeHandler)
+    elect()
   }
 
   private def updateMetrics(): Unit = {
@@ -1191,25 +1385,30 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
       if (!isActive) {
         0
       } else {
-        partitionStateMachine.offlinePartitionCount
-      }
-
-    preferredReplicaImbalanceCount =
-      if (!isActive) {
-        0
-      } else {
-        controllerContext.allPartitions.count { topicPartition =>
-          val replicas = controllerContext.partitionReplicaAssignment(topicPartition)
-          val preferredReplica = replicas.head
-          val leadershipInfo = controllerContext.partitionLeadershipInfo.get(topicPartition)
-          leadershipInfo.map(_.leaderAndIsr.leader != preferredReplica).getOrElse(false) &&
-            !topicDeletionManager.isTopicQueuedUpForDeletion(topicPartition.topic)
-        }
+        controllerContext.offlinePartitionCount
       }
 
     globalTopicCount = if (!isActive) 0 else controllerContext.allTopics.size
 
     globalPartitionCount = if (!isActive) 0 else controllerContext.partitionLeadershipInfo.size
+
+    topicsToDeleteCount = if (!isActive) 0 else controllerContext.topicsToBeDeleted.size
+
+    replicasToDeleteCount = if (!isActive) 0 else controllerContext.topicsToBeDeleted.map { topic =>
+      // For each enqueued topic, count the number of replicas that are not yet deleted
+      controllerContext.replicasForTopic(topic).count { replica =>
+        controllerContext.replicaState(replica) != ReplicaDeletionSuccessful
+      }
+    }.sum
+
+    ineligibleTopicsToDeleteCount = if (!isActive) 0 else controllerContext.topicsIneligibleForDeletion.size
+
+    ineligibleReplicasToDeleteCount = if (!isActive) 0 else controllerContext.topicsToBeDeleted.map { topic =>
+      // For each enqueued topic, count the number of replicas that are ineligible
+      controllerContext.replicasForTopic(topic).count { replica =>
+        controllerContext.replicaState(replica) == ReplicaDeletionIneligible
+      }
+    }.sum
   }
 
   // visible for testing
@@ -1248,6 +1447,10 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
   }
 
   private def elect(): Unit = {
+
+    // refresh preferred controller nodes
+    controllerContext.setLivePreferredControllerIds(zkClient.getPreferredControllerList.toSet)
+
     activeControllerId = zkClient.getControllerId.getOrElse(-1)
     /*
      * We can get here during the initial startup and the handleDeleted ZK callback. Because of the potential race condition,
@@ -1259,6 +1462,16 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
       return
     }
 
+    if (!config.preferredController) {
+      /*
+       * A broker that is not preferred controller only elects itself when there is no preferred controllers in cluster
+       * and allowPreferredControllerFallback is set to true.
+       */
+      if (!config.allowPreferredControllerFallback || controllerContext.getLivePreferredControllerIds.nonEmpty) {
+        return
+      }
+    }
+
     try {
       val (epoch, epochZkVersion) = zkClient.registerControllerAndIncrementControllerEpoch(config.brokerId)
       controllerContext.epoch = epoch
@@ -1285,473 +1498,722 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
     }
   }
 
-  case object BrokerChange extends ControllerEvent {
-    override def state: ControllerState = ControllerState.BrokerChange
-
-    override def process(): Unit = {
-      if (!isActive) return
-      val curBrokerAndEpochs = zkClient.getAllBrokerAndEpochsInCluster
-      val curBrokerIdAndEpochs = curBrokerAndEpochs map { case (broker, epoch) => (broker.id, epoch) }
-      val curBrokerIds = curBrokerIdAndEpochs.keySet
-      val liveOrShuttingDownBrokerIds = controllerContext.liveOrShuttingDownBrokerIds
-      val newBrokerIds = curBrokerIds -- liveOrShuttingDownBrokerIds
-      val deadBrokerIds = liveOrShuttingDownBrokerIds -- curBrokerIds
-      val bouncedBrokerIds = (curBrokerIds & liveOrShuttingDownBrokerIds)
-        .filter(brokerId => curBrokerIdAndEpochs(brokerId) > controllerContext.liveBrokerIdAndEpochs(brokerId))
-      val newBrokerAndEpochs = curBrokerAndEpochs.filterKeys(broker => newBrokerIds.contains(broker.id))
-      val bouncedBrokerAndEpochs = curBrokerAndEpochs.filterKeys(broker => bouncedBrokerIds.contains(broker.id))
-      val newBrokerIdsSorted = newBrokerIds.toSeq.sorted
-      val deadBrokerIdsSorted = deadBrokerIds.toSeq.sorted
-      val liveBrokerIdsSorted = curBrokerIds.toSeq.sorted
-      val bouncedBrokerIdsSorted = bouncedBrokerIds.toSeq.sorted
-      info(s"Newly added brokers: ${newBrokerIdsSorted.mkString(",")}, " +
-        s"deleted brokers: ${deadBrokerIdsSorted.mkString(",")}, " +
-        s"bounced brokers: ${bouncedBrokerIdsSorted.mkString(",")}, " +
-        s"all live brokers: ${liveBrokerIdsSorted.mkString(",")}")
-
-      newBrokerAndEpochs.keySet.foreach(controllerContext.controllerChannelManager.addBroker)
-      bouncedBrokerIds.foreach(controllerContext.controllerChannelManager.removeBroker)
-      bouncedBrokerAndEpochs.keySet.foreach(controllerContext.controllerChannelManager.addBroker)
-      deadBrokerIds.foreach(controllerContext.controllerChannelManager.removeBroker)
-      if (newBrokerIds.nonEmpty) {
-        controllerContext.addLiveBrokersAndEpochs(newBrokerAndEpochs)
-        onBrokerStartup(newBrokerIdsSorted)
-      }
-      if (bouncedBrokerIds.nonEmpty) {
-        controllerContext.removeLiveBrokersAndEpochs(bouncedBrokerIds)
-        onBrokerFailure(bouncedBrokerIdsSorted)
-        controllerContext.addLiveBrokersAndEpochs(bouncedBrokerAndEpochs)
-        onBrokerStartup(bouncedBrokerIdsSorted)
-      }
-      if (deadBrokerIds.nonEmpty) {
-        controllerContext.removeLiveBrokersAndEpochs(deadBrokerIds)
-        onBrokerFailure(deadBrokerIdsSorted)
-      }
+  private def processBrokerChange(): Unit = {
+    if (!isActive) return
+    val curBrokerAndEpochs = zkClient.getAllBrokerAndEpochsInCluster
+    val curBrokerIdAndEpochs = curBrokerAndEpochs map { case (broker, epoch) => (broker.id, epoch) }
+    val curBrokerIds = curBrokerIdAndEpochs.keySet
+    val liveOrShuttingDownBrokerIds = controllerContext.liveOrShuttingDownBrokerIds
+    val newBrokerIds = curBrokerIds -- liveOrShuttingDownBrokerIds
+    val deadBrokerIds = liveOrShuttingDownBrokerIds -- curBrokerIds
+    val bouncedBrokerIds = (curBrokerIds & liveOrShuttingDownBrokerIds)
+      .filter(brokerId => curBrokerIdAndEpochs(brokerId) > controllerContext.liveBrokerIdAndEpochs(brokerId))
+    val newBrokerAndEpochs = curBrokerAndEpochs.filter { case (broker, _) => newBrokerIds.contains(broker.id) }
+    val bouncedBrokerAndEpochs = curBrokerAndEpochs.filter { case (broker, _) => bouncedBrokerIds.contains(broker.id) }
+    val newBrokerIdsSorted = newBrokerIds.toSeq.sorted
+    val deadBrokerIdsSorted = deadBrokerIds.toSeq.sorted
+    val liveBrokerIdsSorted = curBrokerIds.toSeq.sorted
+    val bouncedBrokerIdsSorted = bouncedBrokerIds.toSeq.sorted
+    info(s"Newly added brokers: ${newBrokerIdsSorted.mkString(",")}, " +
+      s"deleted brokers: ${deadBrokerIdsSorted.mkString(",")}, " +
+      s"bounced brokers: ${bouncedBrokerIdsSorted.mkString(",")}, " +
+      s"all live brokers: ${liveBrokerIdsSorted.mkString(",")}")
+
+    newBrokerAndEpochs.keySet.foreach(controllerChannelManager.addBroker)
+    bouncedBrokerIds.foreach(controllerChannelManager.removeBroker)
+    bouncedBrokerAndEpochs.keySet.foreach(controllerChannelManager.addBroker)
+    deadBrokerIds.foreach(controllerChannelManager.removeBroker)
+    if (newBrokerIds.nonEmpty) {
+      controllerContext.addLiveBrokersAndEpochs(newBrokerAndEpochs)
+      onBrokerStartup(newBrokerIdsSorted)
+    }
+    if (bouncedBrokerIds.nonEmpty) {
+      controllerContext.removeLiveBrokers(bouncedBrokerIds)
+      onBrokerFailure(bouncedBrokerIdsSorted)
+      controllerContext.addLiveBrokersAndEpochs(bouncedBrokerAndEpochs)
+      onBrokerStartup(bouncedBrokerIdsSorted)
+    }
+    if (deadBrokerIds.nonEmpty) {
+      controllerContext.removeLiveBrokers(deadBrokerIds)
+      onBrokerFailure(deadBrokerIdsSorted)
+    }
+
+    if (newBrokerIds.nonEmpty || deadBrokerIds.nonEmpty || bouncedBrokerIds.nonEmpty) {
+      info(s"Updated broker epochs cache: ${controllerContext.liveBrokerIdAndEpochs}")
+    }
+  }
 
-      if (newBrokerIds.nonEmpty || deadBrokerIds.nonEmpty || bouncedBrokerIds.nonEmpty) {
-        info(s"Updated broker epochs cache: ${controllerContext.liveBrokerIdAndEpochs}")
+  /**
+    * PreferredControllerChange event should be handled by all brokers to perform:
+    * 1) update preferred controllers
+    * 2) non-preferred active controller resign
+    * 3) elect controller in the absence of preferred controllers. When the last preferred controller goes offline,
+    * ControllerChange event might be fired before the PreferredControllerChange event. When ControllerChange event
+    * is processed, a broker may still see the preferred controller under preferred controller znode and not to elect
+    * itself as controller. Therefore, when PreferredControllerChange event is processed,
+    * a controller election needs to be attempted.
+    */
+  private def processPreferredControllerChange(): Unit = {
+    // refresh the preferred controller list and reset zookeeper watch by reading the list
+    controllerContext.setLivePreferredControllerIds(zkClient.getPreferredControllerList.toSet)
+    if (isActive) {
+      if (!config.preferredController && controllerContext.getLivePreferredControllerIds.nonEmpty) {
+        // resign if a non-preferred active controller detects any live preferred controller
+        info(s"Resign when detecting preferred controllers.")
+        triggerControllerMove()
+      }
+    } else {
+      if (controllerContext.getLivePreferredControllerIds.isEmpty) {
+        // trigger election in the absence of preferred controllers
+        processReelect()
       }
     }
   }
 
-  case class BrokerModifications(brokerId: Int) extends ControllerEvent {
-    override def state: ControllerState = ControllerState.BrokerChange
-
-    override def process(): Unit = {
-      if (!isActive) return
-      val newMetadata = zkClient.getBroker(brokerId)
-      val oldMetadata = controllerContext.liveBrokers.find(_.id == brokerId)
-      if (newMetadata.nonEmpty && oldMetadata.nonEmpty && newMetadata.map(_.endPoints) != oldMetadata.map(_.endPoints)) {
-        info(s"Updated broker: ${newMetadata.get}")
-
+  private def processBrokerModification(brokerId: Int): Unit = {
+    if (!isActive) return
+    val newMetadataOpt = zkClient.getBroker(brokerId)
+    val oldMetadataOpt = controllerContext.liveOrShuttingDownBroker(brokerId)
+    if (newMetadataOpt.nonEmpty && oldMetadataOpt.nonEmpty) {
+      val oldMetadata = oldMetadataOpt.get
+      val newMetadata = newMetadataOpt.get
+      if (newMetadata.endPoints != oldMetadata.endPoints) {
+        info(s"Updated broker metadata: $oldMetadata -> $newMetadata")
         controllerContext.updateBrokerMetadata(oldMetadata, newMetadata)
         onBrokerUpdate(brokerId)
       }
     }
   }
 
-  case object TopicChange extends ControllerEvent {
-    override def state: ControllerState = ControllerState.TopicChange
+  private def processTopicChange(): Unit = {
+    if (!isActive) return
+    val topics = zkClient.getAllTopicsInCluster
+    val newTopics = topics -- controllerContext.allTopics
+    val deletedTopics = controllerContext.allTopics -- topics
+    controllerContext.allTopics = topics
+    rearrangePartitionReplicaAssignmentForNewTopics(newTopics)
 
-    override def process(): Unit = {
-      if (!isActive) return
-      val topics = zkClient.getAllTopicsInCluster.toSet
-      val newTopics = topics -- controllerContext.allTopics
-      val deletedTopics = controllerContext.allTopics -- topics
-      controllerContext.allTopics = topics
-
-      registerPartitionModificationsHandlers(newTopics.toSeq)
-      val addedPartitionReplicaAssignment = zkClient.getReplicaAssignmentForTopics(newTopics)
-      deletedTopics.foreach(controllerContext.removeTopic)
-      addedPartitionReplicaAssignment.foreach {
-        case (topicAndPartition, newReplicas) => controllerContext.updatePartitionReplicaAssignment(topicAndPartition, newReplicas)
-      }
-      info(s"New topics: [$newTopics], deleted topics: [$deletedTopics], new partition replica assignment " +
-        s"[$addedPartitionReplicaAssignment]")
-      if (addedPartitionReplicaAssignment.nonEmpty)
-        onNewPartitionCreation(addedPartitionReplicaAssignment.keySet)
+    registerPartitionModificationsHandlers(newTopics.toSeq)
+    val addedPartitionReplicaAssignment = getReplicaAssignmentPolicyCompliant(newTopics)
+    deletedTopics.foreach(controllerContext.removeTopic)
+    addedPartitionReplicaAssignment.foreach {
+      case (topicAndPartition, newReplicaAssignment) => controllerContext.updatePartitionFullReplicaAssignment(topicAndPartition, newReplicaAssignment)
     }
+    info(s"New topics: [$newTopics], deleted topics: [$deletedTopics], new partition replica assignment " +
+      s"[$addedPartitionReplicaAssignment]")
+    if (addedPartitionReplicaAssignment.nonEmpty)
+      onNewPartitionCreation(addedPartitionReplicaAssignment.keySet)
   }
 
-  case object LogDirEventNotification extends ControllerEvent {
-    override def state: ControllerState = ControllerState.LogDirChange
-
-    override def process(): Unit = {
-      if (!isActive) return
-      val sequenceNumbers = zkClient.getAllLogDirEventNotifications
-      try {
-        val brokerIds = zkClient.getBrokerIdsFromLogDirEvents(sequenceNumbers)
-        onBrokerLogDirFailure(brokerIds)
-      } finally {
-        // delete processed children
-        zkClient.deleteLogDirEventNotifications(sequenceNumbers, controllerContext.epochZkVersion)
-      }
+  private def processLogDirEventNotification(): Unit = {
+    if (!isActive) return
+    val sequenceNumbers = zkClient.getAllLogDirEventNotifications
+    try {
+      val brokerIds = zkClient.getBrokerIdsFromLogDirEvents(sequenceNumbers)
+      onBrokerLogDirFailure(brokerIds)
+    } finally {
+      // delete processed children
+      zkClient.deleteLogDirEventNotifications(sequenceNumbers, controllerContext.epochZkVersion)
     }
   }
 
-  case class PartitionModifications(topic: String) extends ControllerEvent {
-    override def state: ControllerState = ControllerState.TopicChange
-
-    def restorePartitionReplicaAssignment(topic: String, newPartitionReplicaAssignment : immutable.Map[TopicPartition, Seq[Int]]): Unit = {
+  private def processPartitionModifications(topic: String): Unit = {
+    def restorePartitionReplicaAssignment(topic: String, newPartitionReplicaAssignment: Map[TopicPartition, Seq[Int]]): Unit = {
       info("Restoring the partition replica assignment for topic %s".format(topic))
 
       val existingPartitions = zkClient.getChildren(TopicPartitionsZNode.path(topic))
-      val existingPartitionReplicaAssignment = newPartitionReplicaAssignment.filter(p =>
-        existingPartitions.contains(p._1.partition.toString))
+      val existingPartitionReplicaAssignment = newPartitionReplicaAssignment
+        .filter(p => existingPartitions.contains(p._1.partition.toString))
+        .map { case (tp, _) =>
+          tp -> controllerContext.partitionFullReplicaAssignment(tp)
+      }.toMap
+
+      zkClient.setTopicAssignment(topic,
+        existingPartitionReplicaAssignment,
+        controllerContext.epochZkVersion)
+    }
 
-      zkClient.setTopicAssignment(topic, existingPartitionReplicaAssignment, controllerContext.epochZkVersion)
+    if (!isActive) return
+    val partitionReplicaAssignment = zkClient.getReplicaAssignmentForTopics(immutable.Set(topic))
+    val partitionsToBeAdded = partitionReplicaAssignment.filter { case (topicPartition, _) =>
+      controllerContext.partitionReplicaAssignment(topicPartition).isEmpty
     }
+    if (topicDeletionManager.isTopicQueuedUpForDeletion(topic))
+      if (partitionsToBeAdded.nonEmpty) {
+        warn("Skipping adding partitions %s for topic %s since it is currently being deleted"
+          .format(partitionsToBeAdded.map(_._1.partition).mkString(","), topic))
 
-    override def process(): Unit = {
-      if (!isActive) return
-      val partitionReplicaAssignment = zkClient.getReplicaAssignmentForTopics(immutable.Set(topic))
-      val partitionsToBeAdded = partitionReplicaAssignment.filter { case (topicPartition, _) =>
-        controllerContext.partitionReplicaAssignment(topicPartition).isEmpty
+        restorePartitionReplicaAssignment(topic, partitionReplicaAssignment)
+      } else {
+        // This can happen if existing partition replica assignment are restored to prevent increasing partition count during topic deletion
+        info("Ignoring partition change during topic deletion as no new partitions are added")
+      }
+    else {
+      if (partitionsToBeAdded.nonEmpty) {
+        info(s"New partitions to be added $partitionsToBeAdded")
+        partitionsToBeAdded.foreach { case (topicPartition, assignedReplicas) =>
+          controllerContext.updatePartitionReplicaAssignment(topicPartition, assignedReplicas)
+        }
+        onNewPartitionCreation(partitionsToBeAdded.keySet)
       }
-      if (topicDeletionManager.isTopicQueuedUpForDeletion(topic))
-        if (partitionsToBeAdded.nonEmpty) {
-          warn("Skipping adding partitions %s for topic %s since it is currently being deleted"
-            .format(partitionsToBeAdded.map(_._1.partition).mkString(","), topic))
+    }
+  }
 
-          restorePartitionReplicaAssignment(topic, partitionReplicaAssignment)
-        } else {
-          // This can happen if existing partition replica assignment are restored to prevent increasing partition count during topic deletion
-          info("Ignoring partition change during topic deletion as no new partitions are added")
+  private def processTopicDeletion(): Unit = {
+    if (!isActive) return
+    var topicsToBeDeleted = zkClient.getTopicDeletions.toSet
+    debug(s"Delete topics listener fired for topics ${topicsToBeDeleted.mkString(",")} to be deleted")
+    val nonExistentTopics = topicsToBeDeleted -- controllerContext.allTopics
+    if (nonExistentTopics.nonEmpty) {
+      warn(s"Ignoring request to delete non-existing topics ${nonExistentTopics.mkString(",")}")
+      zkClient.deleteTopicDeletions(nonExistentTopics.toSeq, controllerContext.epochZkVersion)
+    }
+    topicsToBeDeleted --= nonExistentTopics
+    if (topicDeletionManager.isDeleteTopicEnabled) {
+      if (topicsToBeDeleted.nonEmpty) {
+        info(s"Starting topic deletion for topics ${topicsToBeDeleted.mkString(",")}")
+        // mark topic ineligible for deletion if other state changes are in progress
+        topicsToBeDeleted.foreach { topic =>
+          val partitionReassignmentInProgress =
+            controllerContext.partitionsBeingReassigned.map(_.topic).contains(topic)
+          if (partitionReassignmentInProgress)
+            topicDeletionManager.markTopicIneligibleForDeletion(Set(topic),
+              reason = "topic reassignment in progress")
         }
+        // add topic to deletion list
+        topicDeletionManager.enqueueTopicsForDeletion(topicsToBeDeleted)
+      }
+    } else {
+      // If delete topic is disabled remove entries under zookeeper path : /admin/delete_topics
+      info(s"Removing $topicsToBeDeleted since delete topic is disabled")
+      zkClient.deleteTopicDeletions(topicsToBeDeleted.toSeq, controllerContext.epochZkVersion)
+    }
+  }
+
+  private def processTopicDeletionFlagChange(reset: Boolean = false): Unit = {
+    info("Process TopicDeletionFlagChange event")
+    if (!isActive) return
+    if (reset)
+      topicDeletionManager.resetDeleteTopicEnabled()
+    else {
+      val topicDeletionFlag = zkClient.getTopicDeletionFlag
+      if (!topicDeletionFlag.equalsIgnoreCase("true") && !topicDeletionFlag.equalsIgnoreCase("false")) {
+        info(s"Overwrite ${DeleteTopicFlagZNode.path} to ${topicDeletionManager.isDeleteTopicEnabled}")
+        zkClient.setTopicDeletionFlag(topicDeletionManager.isDeleteTopicEnabled.toString)
+      }
       else {
-        if (partitionsToBeAdded.nonEmpty) {
-          info(s"New partitions to be added $partitionsToBeAdded")
-          partitionsToBeAdded.foreach { case (topicPartition, assignedReplicas) =>
-            controllerContext.updatePartitionReplicaAssignment(topicPartition, assignedReplicas)
-          }
-          onNewPartitionCreation(partitionsToBeAdded.keySet)
-        }
+        info(s"Set isDeleteTopicEnabled flag to $topicDeletionFlag")
+        topicDeletionManager.isDeleteTopicEnabled = topicDeletionFlag.toBoolean
       }
     }
   }
 
-  case object TopicDeletion extends ControllerEvent {
-    override def state: ControllerState = ControllerState.TopicDeletion
+  private def processZkPartitionReassignment(): Set[TopicPartition] = {
+    // We need to register the watcher if the path doesn't exist in order to detect future
+    // reassignments and we get the `path exists` check for free
+    if (isActive && zkClient.registerZNodeChangeHandlerAndCheckExistence(partitionReassignmentHandler)) {
+      val reassignmentResults = mutable.Map.empty[TopicPartition, ApiError]
+      val partitionsToReassign = mutable.Map.empty[TopicPartition, ReplicaAssignment]
 
-    override def process(): Unit = {
-      if (!isActive) return
-      var topicsToBeDeleted = zkClient.getTopicDeletions.toSet
-      debug(s"Delete topics listener fired for topics ${topicsToBeDeleted.mkString(",")} to be deleted")
-      val nonExistentTopics = topicsToBeDeleted -- controllerContext.allTopics
-      if (nonExistentTopics.nonEmpty) {
-        warn(s"Ignoring request to delete non-existing topics ${nonExistentTopics.mkString(",")}")
-        zkClient.deleteTopicDeletions(nonExistentTopics.toSeq, controllerContext.epochZkVersion)
-      }
-      topicsToBeDeleted --= nonExistentTopics
-      if (config.deleteTopicEnable) {
-        if (topicsToBeDeleted.nonEmpty) {
-          info(s"Starting topic deletion for topics ${topicsToBeDeleted.mkString(",")}")
-          // mark topic ineligible for deletion if other state changes are in progress
-          topicsToBeDeleted.foreach { topic =>
-            val partitionReassignmentInProgress =
-              controllerContext.partitionsBeingReassigned.keySet.map(_.topic).contains(topic)
-            if (partitionReassignmentInProgress)
-              topicDeletionManager.markTopicIneligibleForDeletion(Set(topic))
-          }
-          // add topic to deletion list
-          topicDeletionManager.enqueueTopicsForDeletion(topicsToBeDeleted)
+      zkClient.getPartitionReassignment().foreach { case (tp, targetReplicas) =>
+        maybeBuildReassignment(tp, Some(targetReplicas)) match {
+          case Some(context) => partitionsToReassign.put(tp, context)
+          case None => reassignmentResults.put(tp, new ApiError(Errors.NO_REASSIGNMENT_IN_PROGRESS))
         }
-      } else {
-        // If delete topic is disabled remove entries under zookeeper path : /admin/delete_topics
-        info(s"Removing $topicsToBeDeleted since delete topic is disabled")
-        zkClient.deleteTopicDeletions(topicsToBeDeleted.toSeq, controllerContext.epochZkVersion)
       }
+
+      reassignmentResults ++= maybeTriggerPartitionReassignment(partitionsToReassign)
+      val (partitionsReassigned, partitionsFailed) = reassignmentResults.partition(_._2.error == Errors.NONE)
+      if (partitionsFailed.nonEmpty) {
+        warn(s"Failed reassignment through zk with the following errors: $partitionsFailed")
+        maybeRemoveFromZkReassignment((tp, _) => partitionsFailed.contains(tp))
+      }
+      partitionsReassigned.keySet
+    } else {
+      Set.empty
     }
   }
 
-  case object PartitionReassignment extends ControllerEvent {
-    override def state: ControllerState = ControllerState.PartitionReassignment
+  /**
+   * Process a partition reassignment from the AlterPartitionReassignment API. If there is an
+   * existing reassignment through zookeeper for any of the requested partitions, they will be
+   * cancelled prior to beginning the new reassignment. Any zk-based reassignment for partitions
+   * which are NOT included in this call will not be affected.
+   *
+   * @param reassignments Map of reassignments passed through the AlterReassignments API. A null value
+   *                      means that we should cancel an in-progress reassignment.
+   * @param callback Callback to send AlterReassignments response
+   */
+  private def processApiPartitionReassignment(reassignments: Map[TopicPartition, Option[Seq[Int]]],
+                                              callback: AlterReassignmentsCallback): Unit = {
+    if (!isActive) {
+      callback(Right(new ApiError(Errors.NOT_CONTROLLER)))
+    } else {
+      val reassignmentResults = mutable.Map.empty[TopicPartition, ApiError]
+      val partitionsToReassign = mutable.Map.empty[TopicPartition, ReplicaAssignment]
+
+      reassignments.foreach { case (tp, targetReplicas) =>
+        if (replicasAreValid(tp, targetReplicas)) {
+          maybeBuildReassignment(tp, targetReplicas) match {
+            case Some(context) => partitionsToReassign.put(tp, context)
+            case None => reassignmentResults.put(tp, new ApiError(Errors.NO_REASSIGNMENT_IN_PROGRESS))
+          }
+        } else {
+          reassignmentResults.put(tp, new ApiError(Errors.INVALID_REPLICA_ASSIGNMENT))
+        }
+      }
 
-    override def process(): Unit = {
-      if (!isActive) return
+      // The latest reassignment (whether by API or through zk) always takes precedence,
+      // so remove from active zk reassignment (if one exists)
+      maybeRemoveFromZkReassignment((tp, _) => partitionsToReassign.contains(tp))
 
-      // We need to register the watcher if the path doesn't exist in order to detect future reassignments and we get
-      // the `path exists` check for free
-      if (zkClient.registerZNodeChangeHandlerAndCheckExistence(partitionReassignmentHandler)) {
-        val partitionReassignment = zkClient.getPartitionReassignment
+      reassignmentResults ++= maybeTriggerPartitionReassignment(partitionsToReassign)
+      callback(Left(reassignmentResults))
+    }
+  }
 
-        // Populate `partitionsBeingReassigned` with all partitions being reassigned before invoking
-        // `maybeTriggerPartitionReassignment` (see method documentation for the reason)
-        partitionReassignment.foreach { case (tp, newReplicas) =>
-          val reassignIsrChangeHandler = new PartitionReassignmentIsrChangeHandler(KafkaController.this, eventManager,
-            tp)
-          controllerContext.partitionsBeingReassigned.put(tp, ReassignedPartitionsContext(newReplicas, reassignIsrChangeHandler))
+  private def replicasAreValid(topicPartition: TopicPartition, replicasOpt: Option[Seq[Int]]): Boolean = {
+    replicasOpt match {
+      case Some(replicas) =>
+        val replicaSet = replicas.toSet
+        if (replicas.isEmpty || replicas.size != replicaSet.size)
+          false
+        else if (replicas.exists(_ < 0))
+          false
+        else {
+          // Ensure that any new replicas are among the live brokers
+          val currentAssignment = controllerContext.partitionFullReplicaAssignment(topicPartition)
+          val newAssignment = currentAssignment.reassignTo(replicas)
+          newAssignment.addingReplicas.toSet.subsetOf(controllerContext.liveBrokerIds)
         }
 
-        maybeTriggerPartitionReassignment(partitionReassignment.keySet)
+      case None => true
+    }
+  }
+
+  private def maybeBuildReassignment(topicPartition: TopicPartition,
+                                     targetReplicasOpt: Option[Seq[Int]]): Option[ReplicaAssignment] = {
+    val replicaAssignment = controllerContext.partitionFullReplicaAssignment(topicPartition)
+    if (replicaAssignment.isBeingReassigned) {
+      val targetReplicas = targetReplicasOpt.getOrElse(replicaAssignment.originReplicas)
+      Some(replicaAssignment.reassignTo(targetReplicas))
+    } else {
+      targetReplicasOpt.map { targetReplicas =>
+        replicaAssignment.reassignTo(targetReplicas)
       }
     }
   }
 
-  case class PartitionReassignmentIsrChange(partition: TopicPartition) extends ControllerEvent {
-    override def state: ControllerState = ControllerState.PartitionReassignment
-
-    override def process(): Unit = {
-      if (!isActive) return
-      // check if this partition is still being reassigned or not
-      controllerContext.partitionsBeingReassigned.get(partition).foreach { reassignedPartitionContext =>
-        val reassignedReplicas = reassignedPartitionContext.newReplicas.toSet
-        zkClient.getTopicPartitionStates(Seq(partition)).get(partition) match {
-          case Some(leaderIsrAndControllerEpoch) => // check if new replicas have joined ISR
-            val leaderAndIsr = leaderIsrAndControllerEpoch.leaderAndIsr
-            val caughtUpReplicas = reassignedReplicas & leaderAndIsr.isr.toSet
-            if (caughtUpReplicas == reassignedReplicas) {
-              // resume the partition reassignment process
-              info(s"${caughtUpReplicas.size}/${reassignedReplicas.size} replicas have caught up with the leader for " +
-                s"partition $partition being reassigned. Resuming partition reassignment")
-              onPartitionReassignment(partition, reassignedPartitionContext)
-            }
-            else {
-              info(s"${caughtUpReplicas.size}/${reassignedReplicas.size} replicas have caught up with the leader for " +
-                s"partition $partition being reassigned. Replica(s) " +
-                s"${(reassignedReplicas -- leaderAndIsr.isr.toSet).mkString(",")} still need to catch up")
-            }
-          case None => error(s"Error handling reassignment of partition $partition to replicas " +
-                         s"${reassignedReplicas.mkString(",")} as it was never created")
-        }
+  private def processPartitionReassignmentIsrChange(topicPartition: TopicPartition): Unit = {
+    if (!isActive) return
+
+    if (controllerContext.partitionsBeingReassigned.contains(topicPartition)) {
+      val reassignment = controllerContext.partitionFullReplicaAssignment(topicPartition)
+      if (isReassignmentComplete(topicPartition, reassignment)) {
+        // resume the partition reassignment process
+        info(s"Target replicas ${reassignment.targetReplicas} have all caught up with the leader for " +
+          s"reassigning partition $topicPartition")
+        onPartitionReassignment(topicPartition, reassignment)
       }
     }
   }
 
-  case object IsrChangeNotification extends ControllerEvent {
-    override def state: ControllerState = ControllerState.IsrChange
+  private def processListPartitionReassignments(partitionsOpt: Option[Set[TopicPartition]], callback: ListReassignmentsCallback): Unit = {
+    if (!isActive) {
+      callback(Right(new ApiError(Errors.NOT_CONTROLLER)))
+    } else {
+      val results: mutable.Map[TopicPartition, ReplicaAssignment] = mutable.Map.empty
+      val partitionsToList = partitionsOpt match {
+        case Some(partitions) => partitions
+        case None => controllerContext.partitionsBeingReassigned
+      }
 
-    override def process(): Unit = {
-      if (!isActive) return
-      val sequenceNumbers = zkClient.getAllIsrChangeNotifications
-      try {
-        val partitions = zkClient.getPartitionsFromIsrChangeNotifications(sequenceNumbers)
-        if (partitions.nonEmpty) {
-          updateLeaderAndIsrCache(partitions)
-          processUpdateNotifications(partitions)
+      partitionsToList.foreach { tp =>
+        val assignment = controllerContext.partitionFullReplicaAssignment(tp)
+        if (assignment.isBeingReassigned) {
+          results += tp -> assignment
         }
-      } finally {
-        // delete the notifications
-        zkClient.deleteIsrChangeNotifications(sequenceNumbers, controllerContext.epochZkVersion)
       }
+
+      callback(Left(results))
     }
+  }
 
-    private def processUpdateNotifications(partitions: Seq[TopicPartition]) {
+  private def processIsrChangeNotification(): Unit = {
+    def processUpdateNotifications(partitions: Seq[TopicPartition]): Unit = {
       val liveBrokers: Seq[Int] = controllerContext.liveOrShuttingDownBrokerIds.toSeq
       debug(s"Sending MetadataRequest to Brokers: $liveBrokers for TopicPartitions: $partitions")
       sendUpdateMetadataRequest(liveBrokers, partitions.toSet)
     }
+
+    if (!isActive) return
+    val sequenceNumbers = zkClient.getAllIsrChangeNotifications
+    try {
+      val partitions = zkClient.getPartitionsFromIsrChangeNotifications(sequenceNumbers)
+      if (partitions.nonEmpty) {
+        updateLeaderAndIsrCache(partitions)
+        processUpdateNotifications(partitions)
+      }
+    } finally {
+      // delete the notifications
+      zkClient.deleteIsrChangeNotifications(sequenceNumbers, controllerContext.epochZkVersion)
+    }
   }
 
-  type ElectPreferredLeadersCallback = (Map[TopicPartition, Int], Map[TopicPartition, ApiError])=>Unit
+  def electLeaders(
+    partitions: Set[TopicPartition],
+    electionType: ElectionType,
+    callback: ElectLeadersCallback
+  ): Unit = {
+    eventManager.put(ReplicaLeaderElection(Some(partitions), electionType, AdminClientTriggered, callback))
+  }
 
-  def electPreferredLeaders(partitions: Set[TopicPartition], callback: ElectPreferredLeadersCallback = { (_,_) => }): Unit =
-    eventManager.put(PreferredReplicaLeaderElection(Some(partitions), AdminClientTriggered, callback))
+  def listPartitionReassignments(partitions: Option[Set[TopicPartition]],
+                                 callback: ListReassignmentsCallback): Unit = {
+    eventManager.put(ListPartitionReassignments(partitions, callback))
+  }
 
-  case class PreferredReplicaLeaderElection(partitionsFromAdminClientOpt: Option[Set[TopicPartition]],
-                                            electionType: ElectionType = ZkTriggered,
-                                            callback: ElectPreferredLeadersCallback = (_,_) =>{}) extends PreemptableControllerEvent {
-    override def state: ControllerState = ControllerState.ManualLeaderBalance
+  def alterPartitionReassignments(partitions: Map[TopicPartition, Option[Seq[Int]]],
+                                  callback: AlterReassignmentsCallback): Unit = {
+    eventManager.put(ApiPartitionReassignment(partitions, callback))
+  }
 
-    override def handlePreempt(): Unit = {
-      callback(Map.empty, partitionsFromAdminClientOpt match {
-        case Some(partitions) => partitions.map(partition => partition -> new ApiError(Errors.NOT_CONTROLLER, null)).toMap
-        case None => Map.empty
-      })
-    }
+  private def preemptReplicaLeaderElection(
+    partitionsFromAdminClientOpt: Option[Set[TopicPartition]],
+    callback: ElectLeadersCallback
+  ): Unit = {
+    callback(
+      partitionsFromAdminClientOpt.fold(Map.empty[TopicPartition, Either[ApiError, Int]]) { partitions =>
+        partitions.iterator.map(partition => partition -> Left(new ApiError(Errors.NOT_CONTROLLER, null))).toMap
+      }
+    )
+  }
 
-    override def handleProcess(): Unit = {
-      if (!isActive) {
-        callback(Map.empty, partitionsFromAdminClientOpt match {
-          case Some(partitions) => partitions.map(partition => partition -> new ApiError(Errors.NOT_CONTROLLER, null)).toMap
-          case None => Map.empty
-        })
-      } else {
-        // We need to register the watcher if the path doesn't exist in order to detect future preferred replica
-        // leader elections and we get the `path exists` check for free
-        if (electionType == AdminClientTriggered || zkClient.registerZNodeChangeHandlerAndCheckExistence(preferredReplicaElectionHandler)) {
-          val partitions = partitionsFromAdminClientOpt match {
-            case Some(partitions) => partitions
-            case None => zkClient.getPreferredReplicaElection
-          }
+  private def processReplicaLeaderElection(
+    partitionsFromAdminClientOpt: Option[Set[TopicPartition]],
+    electionType: ElectionType,
+    electionTrigger: ElectionTrigger,
+    callback: ElectLeadersCallback
+  ): Unit = {
+    if (!isActive) {
+      callback(partitionsFromAdminClientOpt.fold(Map.empty[TopicPartition, Either[ApiError, Int]]) { partitions =>
+        partitions.iterator.map(partition => partition -> Left(new ApiError(Errors.NOT_CONTROLLER, null))).toMap
+      })
+    } else {
+      // We need to register the watcher if the path doesn't exist in order to detect future preferred replica
+      // leader elections and we get the `path exists` check for free
+      if (electionTrigger == AdminClientTriggered || zkClient.registerZNodeChangeHandlerAndCheckExistence(preferredReplicaElectionHandler)) {
+        val partitions = partitionsFromAdminClientOpt match {
+          case Some(partitions) => partitions
+          case None => zkClient.getPreferredReplicaElection
+        }
 
-          val (validPartitions, invalidPartitions) = partitions.partition(tp => controllerContext.allPartitions.contains(tp))
-          invalidPartitions.foreach { p =>
-            info(s"Skipping preferred replica leader election for partition ${p} since it doesn't exist.")
-          }
+        val (knownPartitions, unknownPartitions) = partitions.partition(tp => controllerContext.allPartitions.contains(tp))
+        unknownPartitions.foreach { p =>
+          info(s"Skipping replica leader election ($electionType) for partition $p by $electionTrigger since it doesn't exist.")
+        }
 
-          val (partitionsBeingDeleted, livePartitions) = validPartitions.partition(partition =>
+        val (partitionsBeingDeleted, livePartitions) = knownPartitions.partition(partition =>
             topicDeletionManager.isTopicQueuedUpForDeletion(partition.topic))
-          if (partitionsBeingDeleted.nonEmpty) {
-            warn(s"Skipping preferred replica election for partitions $partitionsBeingDeleted " +
-              s"since the respective topics are being deleted")
-          }
-          // partition those where preferred is already leader
-          val (electablePartitions, alreadyPreferred) = livePartitions.partition { partition =>
-            val assignedReplicas = controllerContext.partitionReplicaAssignment(partition)
-            val preferredReplica = assignedReplicas.head
-            val currentLeader = controllerContext.partitionLeadershipInfo(partition).leaderAndIsr.leader
-            currentLeader != preferredReplica
-          }
+        if (partitionsBeingDeleted.nonEmpty) {
+          warn(s"Skipping replica leader election ($electionType) for partitions $partitionsBeingDeleted " +
+            s"by $electionTrigger since the respective topics are being deleted")
+        }
 
-          val electionErrors = onPreferredReplicaElection(electablePartitions, electionType)
-          val successfulPartitions = electablePartitions -- electionErrors.keySet
-          val results = electionErrors.map { case (partition, ex) =>
-            val apiError = if (ex.isInstanceOf[StateChangeFailedException])
-              new ApiError(Errors.PREFERRED_LEADER_NOT_AVAILABLE, ex.getMessage)
-            else
-              ApiError.fromThrowable(ex)
-            partition -> apiError
-          } ++
-            alreadyPreferred.map(_ -> ApiError.NONE) ++
-            partitionsBeingDeleted.map(_ -> new ApiError(Errors.INVALID_TOPIC_EXCEPTION, "The topic is being deleted")) ++
-            invalidPartitions.map ( tp => tp -> new ApiError(Errors.UNKNOWN_TOPIC_OR_PARTITION, s"The partition does not exist.")
-            )
-          debug(s"PreferredReplicaLeaderElection waiting: $successfulPartitions, results: $results")
-          callback(successfulPartitions.map(
-              tp => tp->controllerContext.partitionReplicaAssignment(tp).head).toMap,
-            results)
+        // partition those that have a valid leader
+        val (electablePartitions, alreadyValidLeader) = livePartitions.partition { partition =>
+          electionType match {
+            case ElectionType.PREFERRED =>
+              val assignedReplicas = controllerContext.partitionReplicaAssignment(partition)
+              val preferredReplica = assignedReplicas.head
+              val currentLeader = controllerContext.partitionLeadershipInfo(partition).leaderAndIsr.leader
+              currentLeader != preferredReplica
+
+            case ElectionType.UNCLEAN =>
+              val currentLeader = controllerContext.partitionLeadershipInfo(partition).leaderAndIsr.leader
+              currentLeader == LeaderAndIsr.NoLeader || !controllerContext.liveBrokerIds.contains(currentLeader)
+          }
         }
+
+        val results = onReplicaElection(electablePartitions, electionType, electionTrigger).map {
+          case (k, Left(ex)) =>
+            if (ex.isInstanceOf[StateChangeFailedException]) {
+              val error = if (electionType == ElectionType.PREFERRED) {
+                Errors.PREFERRED_LEADER_NOT_AVAILABLE
+              } else {
+                Errors.ELIGIBLE_LEADERS_NOT_AVAILABLE
+              }
+              k -> Left(new ApiError(error, ex.getMessage))
+            } else {
+              k -> Left(ApiError.fromThrowable(ex))
+            }
+          case (k, Right(leaderAndIsr)) => k -> Right(leaderAndIsr.leader)
+        } ++
+        alreadyValidLeader.map(_ -> Left(new ApiError(Errors.ELECTION_NOT_NEEDED))) ++
+        partitionsBeingDeleted.map(
+          _ -> Left(new ApiError(Errors.INVALID_TOPIC_EXCEPTION, "The topic is being deleted"))
+        ) ++
+        unknownPartitions.map(
+          _ -> Left(new ApiError(Errors.UNKNOWN_TOPIC_OR_PARTITION, "The partition does not exist."))
+        )
+
+        debug(s"Waiting for any successful result for election type ($electionType) by $electionTrigger for partitions: $results")
+        callback(results)
       }
     }
   }
 
-  case object ControllerChange extends ControllerEvent {
-    override def state = ControllerState.ControllerChange
-
-    override def process(): Unit = {
-      maybeResign()
-    }
+  private def processControllerChange(): Unit = {
+    maybeResign()
   }
 
-  case object Reelect extends ControllerEvent {
-    override def state = ControllerState.ControllerChange
+  private def processReelect(): Unit = {
+    maybeResign()
+    elect()
+  }
 
-    override def process(): Unit = {
-      maybeResign()
-      elect()
+  private def processRegisterBrokerAndReelect(): Unit = {
+    if (config.preferredController) {
+      zkClient.registerPreferredControllerId(brokerInfo.broker.id)
     }
+    _brokerEpoch = zkClient.registerBroker(brokerInfo)
+    processReelect()
   }
 
-  case object RegisterBrokerAndReelect extends ControllerEvent {
-    override def state: ControllerState = ControllerState.ControllerChange
+  private def processExpire(): Unit = {
+    activeControllerId = -1
+    onControllerResignation()
+  }
+
+  // For any topics that fail to meet min RF requirement, generate new valid partition assignment and reset ZNode
+  private def fixTopicsFailingPolicy(topicsReplicaAssignment : Map[String, Map[Int, ReplicaAssignment]]) : Unit = {
+    if (topicsReplicaAssignment.isEmpty) return
 
-    override def process(): Unit = {
-      _brokerEpoch = zkClient.registerBroker(brokerInfo)
-      Reelect.process()
+    val replicationFactor = config.defaultReplicationFactor
+    val brokers = controllerContext.liveOrShuttingDownBrokers.map { sb => kafka.admin.BrokerMetadata(sb.id, sb.rack) }.toSeq
+    val noNewPartitionBrokers = partitionUnassignableBrokerIds.toSet
+
+    topicsReplicaAssignment.foreach{
+      case(topic, partitionAssignment) => {
+        val numPartitions = partitionAssignment.size
+        val assignment = adminZkClient.assignReplicasToAvailableBrokers(brokers, noNewPartitionBrokers, numPartitions, replicationFactor)
+          .map{ case(partition, replicas) => {
+            (new TopicPartition(topic, partition), new ReplicaAssignment(replicas, Seq.empty[Int], Seq.empty[Int]))
+          }}.toMap
+        zkClient.setTopicAssignment(topic, assignment, controllerContext.epochZkVersion)
+        info(s"Updated topic [$topic] with $assignment for replica assignment")
+      }
     }
   }
 
-  // We can't make this a case object due to the countDownLatch field
-  class Expire extends ControllerEvent {
-    private val processingStarted = new CountDownLatch(1)
-    override def state = ControllerState.ControllerChange
+  // Reset partition replica assignment for topics, if any, that fail replication factor check
+  private def getReplicaAssignmentPolicyCompliant(topics : immutable.Set[String]) : Map[TopicPartition, ReplicaAssignment] = {
+    val replicaAssignments = zkClient.getPartitionAssignmentForTopics(topics)
+    val (topicAssignmentSucceedingPolicy, topicAssignmentFailingPolicy) = replicaAssignments.partition(
+      assignment => KafkaController.satisfiesLiCreateTopicPolicy(createTopicPolicy, zkClient, assignment._1, assignment._2))
 
-    override def process(): Unit = {
-      processingStarted.countDown()
-      activeControllerId = -1
-      onControllerResignation()
+    var retTopicAssignment = topicAssignmentSucceedingPolicy
+    val topicsFailingPolicy = topicAssignmentFailingPolicy.keySet
+    if (!topicAssignmentFailingPolicy.isEmpty) {
+      // Since fixTopicsFailingPolicy() will trigger PartitionModification event, need to temporarily unregister
+      // event handler
+      unregisterPartitionModificationsHandlers(topicsFailingPolicy.toSeq)
+      fixTopicsFailingPolicy(topicAssignmentFailingPolicy)
+      registerPartitionModificationsHandlers(topicsFailingPolicy.toSeq)
+
+      // the new partition assignments should be valid now
+      retTopicAssignment ++= zkClient.getPartitionAssignmentForTopics(topicsFailingPolicy.toSet)
+    }
+    retTopicAssignment.flatMap{ case(topic, partitionAssignment) =>
+      partitionAssignment.map{
+        case(partition, replicas)=> (new TopicPartition(topic, partition), replicas)
+      }
     }
+  }
 
-    def waitUntilProcessingStarted(): Unit = {
-      processingStarted.await()
+  override def process(event: ControllerEvent): Unit = {
+    try {
+      event match {
+        case event: MockEvent =>
+          // Used only in test cases
+          event.process()
+        case ShutdownEventThread =>
+          error("Received a ShutdownEventThread event. This type of event is supposed to be handle by ControllerEventThread")
+        case AutoPreferredReplicaLeaderElection =>
+          processAutoPreferredReplicaLeaderElection()
+        case ReplicaLeaderElection(partitions, electionType, electionTrigger, callback) =>
+          processReplicaLeaderElection(partitions, electionType, electionTrigger, callback)
+        case UncleanLeaderElectionEnable =>
+          processUncleanLeaderElectionEnable()
+        case TopicUncleanLeaderElectionEnable(topic) =>
+          processTopicUncleanLeaderElectionEnable(topic)
+        case ControlledShutdown(id, brokerEpoch, callback) =>
+          processControlledShutdown(id, brokerEpoch, callback)
+        case LeaderAndIsrResponseReceived(response, brokerId) =>
+          processLeaderAndIsrResponseReceived(response, brokerId)
+        case TopicDeletionStopReplicaResponseReceived(replicaId, requestError, partitionErrors) =>
+          processTopicDeletionStopReplicaResponseReceived(replicaId, requestError, partitionErrors)
+        case BrokerChange =>
+          processBrokerChange()
+        case PreferredControllerChange =>
+          processPreferredControllerChange()
+        case BrokerModifications(brokerId) =>
+          processBrokerModification(brokerId)
+        case ControllerChange =>
+          processControllerChange()
+        case Reelect =>
+          processReelect()
+        case RegisterBrokerAndReelect =>
+          processRegisterBrokerAndReelect()
+        case Expire =>
+          processExpire()
+        case TopicChange =>
+          processTopicChange()
+        case LogDirEventNotification =>
+          processLogDirEventNotification()
+        case PartitionModifications(topic) =>
+          processPartitionModifications(topic)
+        case TopicDeletion =>
+          processTopicDeletion()
+        case ApiPartitionReassignment(reassignments, callback) =>
+          processApiPartitionReassignment(reassignments, callback)
+        case ZkPartitionReassignment =>
+          processZkPartitionReassignment()
+        case ListPartitionReassignments(partitions, callback) =>
+          processListPartitionReassignments(partitions, callback)
+        case TopicDeletionFlagChange(reset) =>
+          processTopicDeletionFlagChange(reset)
+        case PartitionReassignmentIsrChange(partition) =>
+          processPartitionReassignmentIsrChange(partition)
+        case IsrChangeNotification =>
+          processIsrChangeNotification()
+        case Startup =>
+          processStartup()
+      }
+    } catch {
+      case e: ControllerMovedException =>
+        info(s"Controller moved to another broker when processing $event.", e)
+        maybeResign()
+      case e: Throwable =>
+        error(s"Error processing event $event", e)
+    } finally {
+      updateMetrics()
     }
   }
 
+  override def preempt(event: ControllerEvent): Unit = {
+    event match {
+      case ReplicaLeaderElection(partitions, _, _, callback) =>
+        preemptReplicaLeaderElection(partitions, callback)
+      case ControlledShutdown(id, brokerEpoch, callback) =>
+        preemptControlledShutdown(id, brokerEpoch, callback)
+      case _ =>
+    }
+  }
 }
 
-class BrokerChangeHandler(controller: KafkaController, eventManager: ControllerEventManager) extends ZNodeChildChangeHandler {
+class BrokerChangeHandler(eventManager: ControllerEventManager) extends ZNodeChildChangeHandler {
   override val path: String = BrokerIdsZNode.path
 
   override def handleChildChange(): Unit = {
-    eventManager.put(controller.BrokerChange)
+    eventManager.put(BrokerChange)
   }
 }
 
-class BrokerModificationsHandler(controller: KafkaController, eventManager: ControllerEventManager, brokerId: Int) extends ZNodeChangeHandler {
+class PreferredControllerChangeHandler(eventManager: ControllerEventManager) extends ZNodeChildChangeHandler {
+  override val path: String = PreferredControllersZNode.path
+
+  override def handleChildChange(): Unit = {
+    eventManager.put(PreferredControllerChange)
+  }
+}
+
+class BrokerModificationsHandler(eventManager: ControllerEventManager, brokerId: Int) extends ZNodeChangeHandler {
   override val path: String = BrokerIdZNode.path(brokerId)
 
   override def handleDataChange(): Unit = {
-    eventManager.put(controller.BrokerModifications(brokerId))
+    eventManager.put(BrokerModifications(brokerId))
   }
 }
 
-class TopicChangeHandler(controller: KafkaController, eventManager: ControllerEventManager) extends ZNodeChildChangeHandler {
+class TopicChangeHandler(eventManager: ControllerEventManager) extends ZNodeChildChangeHandler {
   override val path: String = TopicsZNode.path
 
-  override def handleChildChange(): Unit = eventManager.put(controller.TopicChange)
+  override def handleChildChange(): Unit = eventManager.put(TopicChange)
 }
 
-class LogDirEventNotificationHandler(controller: KafkaController, eventManager: ControllerEventManager) extends ZNodeChildChangeHandler {
+class LogDirEventNotificationHandler(eventManager: ControllerEventManager) extends ZNodeChildChangeHandler {
   override val path: String = LogDirEventNotificationZNode.path
 
-  override def handleChildChange(): Unit = eventManager.put(controller.LogDirEventNotification)
+  override def handleChildChange(): Unit = eventManager.put(LogDirEventNotification)
 }
 
 object LogDirEventNotificationHandler {
   val Version: Long = 1L
 }
 
-class PartitionModificationsHandler(controller: KafkaController, eventManager: ControllerEventManager, topic: String) extends ZNodeChangeHandler {
+class PartitionModificationsHandler(eventManager: ControllerEventManager, topic: String) extends ZNodeChangeHandler {
   override val path: String = TopicZNode.path(topic)
 
-  override def handleDataChange(): Unit = eventManager.put(controller.PartitionModifications(topic))
+  override def handleDataChange(): Unit = eventManager.put(PartitionModifications(topic))
 }
 
-class TopicDeletionHandler(controller: KafkaController, eventManager: ControllerEventManager) extends ZNodeChildChangeHandler {
+class TopicDeletionHandler(eventManager: ControllerEventManager) extends ZNodeChildChangeHandler {
   override val path: String = DeleteTopicsZNode.path
 
-  override def handleChildChange(): Unit = eventManager.put(controller.TopicDeletion)
+  override def handleChildChange(): Unit = eventManager.put(TopicDeletion)
+}
+/**
+  * Listener for /topic_deletion_flag znode.
+  *   If the data of the znode is set to true/false, it will trigger the in memory isDeleteTopicEnabled to be set accordingly.
+  *   If the znode data cannot be converted to boolean, it will overwrite znode with the previous valid value.
+  *   If the znode path is deleted, it will reset the in memory isDeleteTopicEnabled to the config value.
+  */
+class TopicDeletionFlagHandler(controller: KafkaController, eventManager: ControllerEventManager) extends ZNodeChangeHandler {
+  override val path: String = DeleteTopicFlagZNode.path
+
+  override def handleDataChange(): Unit = eventManager.put(TopicDeletionFlagChange())
+
+  override def handleDeletion(): Unit = eventManager.put(TopicDeletionFlagChange(true))
 }
 
-class PartitionReassignmentHandler(controller: KafkaController, eventManager: ControllerEventManager) extends ZNodeChangeHandler {
+class PartitionReassignmentHandler(eventManager: ControllerEventManager) extends ZNodeChangeHandler {
   override val path: String = ReassignPartitionsZNode.path
 
   // Note that the event is also enqueued when the znode is deleted, but we do it explicitly instead of relying on
   // handleDeletion(). This approach is more robust as it doesn't depend on the watcher being re-registered after
   // it's consumed during data changes (we ensure re-registration when the znode is deleted).
-  override def handleCreation(): Unit = eventManager.put(controller.PartitionReassignment)
+  override def handleCreation(): Unit = eventManager.put(ZkPartitionReassignment)
 }
 
-class PartitionReassignmentIsrChangeHandler(controller: KafkaController, eventManager: ControllerEventManager, partition: TopicPartition) extends ZNodeChangeHandler {
+class PartitionReassignmentIsrChangeHandler(eventManager: ControllerEventManager, partition: TopicPartition) extends ZNodeChangeHandler {
   override val path: String = TopicPartitionStateZNode.path(partition)
 
-  override def handleDataChange(): Unit = eventManager.put(controller.PartitionReassignmentIsrChange(partition))
+  override def handleDataChange(): Unit = eventManager.put(PartitionReassignmentIsrChange(partition))
 }
 
-class IsrChangeNotificationHandler(controller: KafkaController, eventManager: ControllerEventManager) extends ZNodeChildChangeHandler {
+class IsrChangeNotificationHandler(eventManager: ControllerEventManager) extends ZNodeChildChangeHandler {
   override val path: String = IsrChangeNotificationZNode.path
 
-  override def handleChildChange(): Unit = eventManager.put(controller.IsrChangeNotification)
+  override def handleChildChange(): Unit = eventManager.put(IsrChangeNotification)
 }
 
 object IsrChangeNotificationHandler {
   val Version: Long = 1L
 }
 
-class PreferredReplicaElectionHandler(controller: KafkaController, eventManager: ControllerEventManager) extends ZNodeChangeHandler {
+class PreferredReplicaElectionHandler(eventManager: ControllerEventManager) extends ZNodeChangeHandler {
   override val path: String = PreferredReplicaElectionZNode.path
 
-  override def handleCreation(): Unit = eventManager.put(controller.PreferredReplicaLeaderElection(None))
+  override def handleCreation(): Unit = eventManager.put(ReplicaLeaderElection(None, ElectionType.PREFERRED, ZkTriggered))
 }
 
-class ControllerChangeHandler(controller: KafkaController, eventManager: ControllerEventManager) extends ZNodeChangeHandler {
+class ControllerChangeHandler(eventManager: ControllerEventManager) extends ZNodeChangeHandler {
   override val path: String = ControllerZNode.path
 
-  override def handleCreation(): Unit = eventManager.put(controller.ControllerChange)
-  override def handleDeletion(): Unit = eventManager.put(controller.Reelect)
-  override def handleDataChange(): Unit = eventManager.put(controller.ControllerChange)
-}
-
-case class ReassignedPartitionsContext(var newReplicas: Seq[Int] = Seq.empty,
-                                       val reassignIsrChangeHandler: PartitionReassignmentIsrChangeHandler) {
-
-  def registerReassignIsrChangeHandler(zkClient: KafkaZkClient): Unit =
-    zkClient.registerZNodeChangeHandler(reassignIsrChangeHandler)
-
-  def unregisterReassignIsrChangeHandler(zkClient: KafkaZkClient): Unit =
-    zkClient.unregisterZNodeChangeHandler(reassignIsrChangeHandler.path)
-
+  override def handleCreation(): Unit = eventManager.put(ControllerChange)
+  override def handleDeletion(): Unit = eventManager.put(Reelect)
+  override def handleDataChange(): Unit = eventManager.put(ControllerChange)
 }
 
 case class PartitionAndReplica(topicPartition: TopicPartition, replica: Int) {
@@ -1786,30 +2248,127 @@ private[controller] class ControllerStats extends KafkaMetricsGroup {
 }
 
 sealed trait ControllerEvent {
-  val enqueueTimeMs: Long = Time.SYSTEM.milliseconds()
-
   def state: ControllerState
-  def process(): Unit
 }
 
-/**
-  * A `ControllerEvent`, such as one with a client callback, which needs specific handling in the event of ZK session expiration.
-  */
-sealed trait PreemptableControllerEvent extends ControllerEvent {
+case object ControllerChange extends ControllerEvent {
+  override def state = ControllerState.ControllerChange
+}
 
-  val spent = new AtomicBoolean(false)
+case object Reelect extends ControllerEvent {
+  override def state = ControllerState.ControllerChange
+}
 
-  final def preempt(): Unit = {
-    if (!spent.getAndSet(true))
-      handlePreempt()
-  }
+case object RegisterBrokerAndReelect extends ControllerEvent {
+  override def state: ControllerState = ControllerState.ControllerChange
+}
 
-  final def process(): Unit = {
-    if (!spent.getAndSet(true))
-      handleProcess()
-  }
+case object Expire extends ControllerEvent {
+  override def state = ControllerState.ControllerChange
+}
+
+case object ShutdownEventThread extends ControllerEvent {
+  def state = ControllerState.ControllerShutdown
+}
 
-  def handlePreempt(): Unit
+case object AutoPreferredReplicaLeaderElection extends ControllerEvent {
+  def state = ControllerState.AutoLeaderBalance
+}
+
+case object UncleanLeaderElectionEnable extends ControllerEvent {
+  def state = ControllerState.UncleanLeaderElectionEnable
+}
+
+case class TopicUncleanLeaderElectionEnable(topic: String) extends ControllerEvent {
+  def state = ControllerState.TopicUncleanLeaderElectionEnable
+}
+
+case class ControlledShutdown(id: Int, brokerEpoch: Long, controlledShutdownCallback: Try[Set[TopicPartition]] => Unit) extends ControllerEvent {
+  def state = ControllerState.ControlledShutdown
+}
+
+case class LeaderAndIsrResponseReceived(LeaderAndIsrResponseObj: AbstractResponse, brokerId: Int) extends ControllerEvent {
+  def state = ControllerState.LeaderAndIsrResponseReceived
+}
+
+case class TopicDeletionStopReplicaResponseReceived(replicaId: Int,
+                                                    requestError: Errors,
+                                                    partitionErrors: Map[TopicPartition, Errors]) extends ControllerEvent {
+  def state = ControllerState.TopicDeletion
+}
+
+case object Startup extends ControllerEvent {
+  def state = ControllerState.ControllerChange
+}
+
+case object BrokerChange extends ControllerEvent {
+  override def state: ControllerState = ControllerState.BrokerChange
+}
+
+case object PreferredControllerChange extends ControllerEvent {
+  override def state: ControllerState = ControllerState.PreferredControllerChange
+}
+
+case class BrokerModifications(brokerId: Int) extends ControllerEvent {
+  override def state: ControllerState = ControllerState.BrokerChange
+}
+
+case object TopicChange extends ControllerEvent {
+  override def state: ControllerState = ControllerState.TopicChange
+}
 
-  def handleProcess(): Unit
+case object LogDirEventNotification extends ControllerEvent {
+  override def state: ControllerState = ControllerState.LogDirChange
+}
+
+case class PartitionModifications(topic: String) extends ControllerEvent {
+  override def state: ControllerState = ControllerState.TopicChange
+}
+
+case object TopicDeletion extends ControllerEvent {
+  override def state: ControllerState = ControllerState.TopicDeletion
+}
+
+case object ZkPartitionReassignment extends ControllerEvent {
+  override def state: ControllerState = ControllerState.AlterPartitionReassignment
+}
+
+case class ApiPartitionReassignment(reassignments: Map[TopicPartition, Option[Seq[Int]]],
+                                    callback: AlterReassignmentsCallback) extends ControllerEvent {
+  override def state: ControllerState = ControllerState.AlterPartitionReassignment
+}
+
+case class PartitionReassignmentIsrChange(partition: TopicPartition) extends ControllerEvent {
+  override def state: ControllerState = ControllerState.AlterPartitionReassignment
+}
+
+case object IsrChangeNotification extends ControllerEvent {
+  override def state: ControllerState = ControllerState.IsrChange
+}
+
+case class TopicDeletionFlagChange(reset: Boolean = false) extends ControllerEvent {
+  def state = ControllerState.TopicDeletionFlagChange
+}
+
+case class ReplicaLeaderElection(
+  partitionsFromAdminClientOpt: Option[Set[TopicPartition]],
+  electionType: ElectionType,
+  electionTrigger: ElectionTrigger,
+  callback: ElectLeadersCallback = _ => {}
+) extends ControllerEvent {
+  override def state: ControllerState = ControllerState.ManualLeaderBalance
+}
+
+/**
+  * @param partitionsOpt - an Optional set of partitions. If not present, all reassigning partitions are to be listed
+  */
+case class ListPartitionReassignments(partitionsOpt: Option[Set[TopicPartition]],
+                                      callback: ListReassignmentsCallback) extends ControllerEvent {
+  override def state: ControllerState = ControllerState.ListPartitionReassignment
+}
+
+
+// Used only in test cases
+abstract class MockEvent(val state: ControllerState) extends ControllerEvent {
+  def process(): Unit
 }
diff --git a/core/src/main/scala/kafka/controller/PartitionStateMachine.scala b/core/src/main/scala/kafka/controller/PartitionStateMachine.scala
index ad739797179c0..b1ecedfcd0684 100755
--- a/core/src/main/scala/kafka/controller/PartitionStateMachine.scala
+++ b/core/src/main/scala/kafka/controller/PartitionStateMachine.scala
@@ -18,73 +18,68 @@ package kafka.controller
 
 import kafka.api.LeaderAndIsr
 import kafka.common.StateChangeFailedException
+import kafka.controller.Election._
 import kafka.server.KafkaConfig
 import kafka.utils.Logging
-import kafka.zk.{KafkaZkClient, TopicPartitionStateZNode}
+import kafka.zk.KafkaZkClient
 import kafka.zk.KafkaZkClient.UpdateLeaderAndIsrResult
+import kafka.zk.TopicPartitionStateZNode
 import org.apache.kafka.common.TopicPartition
 import org.apache.kafka.common.errors.ControllerMovedException
 import org.apache.zookeeper.KeeperException
 import org.apache.zookeeper.KeeperException.Code
+import scala.collection.{Map, Seq, mutable}
 
-import scala.collection.mutable
-
-
-/**
- * This class represents the state machine for partitions. It defines the states that a partition can be in, and
- * transitions to move the partition to another legal state. The different states that a partition can be in are -
- * 1. NonExistentPartition: This state indicates that the partition was either never created or was created and then
- *                          deleted. Valid previous state, if one exists, is OfflinePartition
- * 2. NewPartition        : After creation, the partition is in the NewPartition state. In this state, the partition should have
- *                          replicas assigned to it, but no leader/isr yet. Valid previous states are NonExistentPartition
- * 3. OnlinePartition     : Once a leader is elected for a partition, it is in the OnlinePartition state.
- *                          Valid previous states are NewPartition/OfflinePartition
- * 4. OfflinePartition    : If, after successful leader election, the leader for partition dies, then the partition
- *                          moves to the OfflinePartition state. Valid previous states are NewPartition/OnlinePartition
- */
-class PartitionStateMachine(config: KafkaConfig,
-                            stateChangeLogger: StateChangeLogger,
-                            controllerContext: ControllerContext,
-                            zkClient: KafkaZkClient,
-                            partitionState: mutable.Map[TopicPartition, PartitionState],
-                            controllerBrokerRequestBatch: ControllerBrokerRequestBatch) extends Logging {
-  private val controllerId = config.brokerId
-
-  private var topicDeletionManager: TopicDeletionManager = _
-
-  this.logIdent = s"[PartitionStateMachine controllerId=$controllerId] "
-
-  var offlinePartitionCount = 0
-
+abstract class PartitionStateMachine(controllerContext: ControllerContext) extends Logging {
   /**
    * Invoked on successful controller election.
    */
-  def startup() {
+  def startup(): Unit = {
     info("Initializing partition state")
     initializePartitionState()
     info("Triggering online partition state changes")
     triggerOnlinePartitionStateChange()
-    info(s"Started partition state machine with initial state -> $partitionState")
+    debug(s"Started partition state machine with initial state -> ${controllerContext.partitionStates}")
   }
 
   /**
    * Invoked on controller shutdown.
    */
-  def shutdown() {
-    partitionState.clear()
-    offlinePartitionCount = 0
+  def shutdown(): Unit = {
     info("Stopped partition state machine")
   }
 
-  def setTopicDeletionManager(topicDeletionManager: TopicDeletionManager) {
-    this.topicDeletionManager = topicDeletionManager
+  /**
+   * This API invokes the OnlinePartition state change on all partitions in either the NewPartition or OfflinePartition
+   * state. This is called on a successful controller election and on broker changes
+   */
+  def triggerOnlinePartitionStateChange(): Unit = {
+    val partitions = controllerContext.partitionsInStates(Set(OfflinePartition, NewPartition))
+    triggerOnlineStateChangeForPartitions(partitions)
+  }
+
+  def triggerOnlinePartitionStateChange(topic: String): Unit = {
+    val partitions = controllerContext.partitionsInStates(topic, Set(OfflinePartition, NewPartition))
+    triggerOnlineStateChangeForPartitions(partitions)
+  }
+
+  private def triggerOnlineStateChangeForPartitions(partitions: collection.Set[TopicPartition]): Unit = {
+    // try to move all partitions in NewPartition or OfflinePartition state to OnlinePartition state except partitions
+    // that belong to topics to be deleted
+    val partitionsToTrigger = partitions.filter { partition =>
+      !controllerContext.isTopicQueuedUpForDeletion(partition.topic)
+    }.toSeq
+
+    handleStateChanges(partitionsToTrigger, OnlinePartition, Some(OfflinePartitionLeaderElectionStrategy(false)))
+    // TODO: If handleStateChanges catches an exception, it is not enough to bail out and log an error.
+    // It is important to trigger leader election for those partitions.
   }
 
   /**
    * Invoked on startup of the partition's state machine to set the initial state for all existing partitions in
    * zookeeper
    */
-  private def initializePartitionState() {
+  private def initializePartitionState(): Unit = {
     for (topicPartition <- controllerContext.allPartitions) {
       // check if leader and isr path exists for partition. If not, then it is in NEW state
       controllerContext.partitionLeadershipInfo.get(topicPartition) match {
@@ -92,86 +87,92 @@ class PartitionStateMachine(config: KafkaConfig,
           // else, check if the leader for partition is alive. If yes, it is in Online state, else it is in Offline state
           if (controllerContext.isReplicaOnline(currentLeaderIsrAndEpoch.leaderAndIsr.leader, topicPartition))
           // leader is alive
-            changeStateTo(topicPartition, NonExistentPartition, OnlinePartition)
+            controllerContext.putPartitionState(topicPartition, OnlinePartition)
           else
-            changeStateTo(topicPartition, NonExistentPartition, OfflinePartition)
+            controllerContext.putPartitionState(topicPartition, OfflinePartition)
         case None =>
-          changeStateTo(topicPartition, NonExistentPartition, NewPartition)
+          controllerContext.putPartitionState(topicPartition, NewPartition)
       }
     }
   }
 
-  /**
-   * This API invokes the OnlinePartition state change on all partitions in either the NewPartition or OfflinePartition
-   * state. This is called on a successful controller election and on broker changes
-   */
-  def triggerOnlinePartitionStateChange() {
-    triggerOnlinePartitionStateChange(partitionState.toMap)
+  def handleStateChanges(
+    partitions: Seq[TopicPartition],
+    targetState: PartitionState
+  ): Map[TopicPartition, Either[Throwable, LeaderAndIsr]] = {
+    handleStateChanges(partitions, targetState, None)
   }
 
-  def triggerOnlinePartitionStateChange(topic: String) {
-    triggerOnlinePartitionStateChange(partitionState.filterKeys(p => p.topic.equals(topic)).toMap)
-  }
+  def handleStateChanges(
+    partitions: Seq[TopicPartition],
+    targetState: PartitionState,
+    leaderElectionStrategy: Option[PartitionLeaderElectionStrategy]
+  ): Map[TopicPartition, Either[Throwable, LeaderAndIsr]]
 
-  def triggerOnlinePartitionStateChange(partitionState: Map[TopicPartition, PartitionState]) {
-    // try to move all partitions in NewPartition or OfflinePartition state to OnlinePartition state except partitions
-    // that belong to topics to be deleted
-    val partitionsToTrigger = partitionState.filter { case (partition, partitionState) =>
-      !topicDeletionManager.isTopicQueuedUpForDeletion(partition.topic) &&
-        (partitionState.equals(OfflinePartition) || partitionState.equals(NewPartition))
-    }.keys.toSeq
-    handleStateChanges(partitionsToTrigger, OnlinePartition, Option(OfflinePartitionLeaderElectionStrategy))
-    // TODO: If handleStateChanges catches an exception, it is not enough to bail out and log an error.
-    // It is important to trigger leader election for those partitions.
-  }
+}
+
+/**
+ * This class represents the state machine for partitions. It defines the states that a partition can be in, and
+ * transitions to move the partition to another legal state. The different states that a partition can be in are -
+ * 1. NonExistentPartition: This state indicates that the partition was either never created or was created and then
+ *                          deleted. Valid previous state, if one exists, is OfflinePartition
+ * 2. NewPartition        : After creation, the partition is in the NewPartition state. In this state, the partition should have
+ *                          replicas assigned to it, but no leader/isr yet. Valid previous states are NonExistentPartition
+ * 3. OnlinePartition     : Once a leader is elected for a partition, it is in the OnlinePartition state.
+ *                          Valid previous states are NewPartition/OfflinePartition
+ * 4. OfflinePartition    : If, after successful leader election, the leader for partition dies, then the partition
+ *                          moves to the OfflinePartition state. Valid previous states are NewPartition/OnlinePartition
+ */
+class ZkPartitionStateMachine(config: KafkaConfig,
+                              stateChangeLogger: StateChangeLogger,
+                              controllerContext: ControllerContext,
+                              zkClient: KafkaZkClient,
+                              controllerBrokerRequestBatch: ControllerBrokerRequestBatch)
+  extends PartitionStateMachine(controllerContext) {
+
+  private val controllerId = config.brokerId
+  this.logIdent = s"[PartitionStateMachine controllerId=$controllerId] "
 
   /**
-    * Try to change the state of the given partitions to the given targetState, using the given
-    * partitionLeaderElectionStrategyOpt if a leader election is required.
-    * @param partitions The partitions
-    * @param targetState The state
-    * @param partitionLeaderElectionStrategyOpt The leader election strategy if a leader election is required.
-    * @return partitions and corresponding throwable for those partitions which could not transition to the given state
-    */
-  def handleStateChanges(partitions: Seq[TopicPartition], targetState: PartitionState,
-                         partitionLeaderElectionStrategyOpt: Option[PartitionLeaderElectionStrategy] = None): Map[TopicPartition, Throwable] = {
+   * Try to change the state of the given partitions to the given targetState, using the given
+   * partitionLeaderElectionStrategyOpt if a leader election is required.
+   * @param partitions The partitions
+   * @param targetState The state
+   * @param partitionLeaderElectionStrategyOpt The leader election strategy if a leader election is required.
+   * @return A map of failed and successful elections when targetState is OnlinePartitions. The keys are the
+   *         topic partitions and the corresponding values are either the exception that was thrown or new
+   *         leader & ISR.
+   */
+  override def handleStateChanges(
+    partitions: Seq[TopicPartition],
+    targetState: PartitionState,
+    partitionLeaderElectionStrategyOpt: Option[PartitionLeaderElectionStrategy]
+  ): Map[TopicPartition, Either[Throwable, LeaderAndIsr]] = {
     if (partitions.nonEmpty) {
       try {
         controllerBrokerRequestBatch.newBatch()
-        val errors = doHandleStateChanges(partitions, targetState, partitionLeaderElectionStrategyOpt)
+        val result = doHandleStateChanges(
+          partitions,
+          targetState,
+          partitionLeaderElectionStrategyOpt
+        )
         controllerBrokerRequestBatch.sendRequestsToBrokers(controllerContext.epoch)
-        errors
+        result
       } catch {
         case e: ControllerMovedException =>
           error(s"Controller moved to another broker when moving some partitions to $targetState state", e)
           throw e
         case e: Throwable =>
           error(s"Error while moving some partitions to $targetState state", e)
-          partitions.map { _ -> e }.toMap
+          partitions.iterator.map(_ -> Left(e)).toMap
       }
     } else {
-      Map.empty[TopicPartition, Throwable]
+      Map.empty
     }
   }
 
-
-  def partitionsInState(state: PartitionState): Set[TopicPartition] = {
-    partitionState.filter { case (_, s) => s == state }.keySet.toSet
-  }
-
-  private def changeStateTo(partition: TopicPartition, currentState: PartitionState, targetState: PartitionState): Unit = {
-    partitionState.put(partition, targetState)
-    updateControllerMetrics(partition, currentState, targetState)
-  }
-
-  private def updateControllerMetrics(partition: TopicPartition, currentState: PartitionState, targetState: PartitionState) : Unit = {
-    if (!topicDeletionManager.isTopicWithDeletionStarted(partition.topic)) {
-      if (currentState != OfflinePartition && targetState == OfflinePartition) {
-        offlinePartitionCount = offlinePartitionCount + 1
-      } else if (currentState == OfflinePartition && targetState != OfflinePartition) {
-        offlinePartitionCount = offlinePartitionCount - 1
-      }
-    }
+  private def partitionState(partition: TopicPartition): PartitionState = {
+    controllerContext.partitionState(partition)
   }
 
   /**
@@ -195,19 +196,26 @@ class PartitionStateMachine(config: KafkaConfig,
    * --nothing other than marking the partition state as NonExistentPartition
    * @param partitions  The partitions for which the state transition is invoked
    * @param targetState The end state that the partition should be moved to
+   * @return A map of failed and successful elections when targetState is OnlinePartitions. The keys are the
+   *         topic partitions and the corresponding values are either the exception that was thrown or new
+   *         leader & ISR.
    */
-  private def doHandleStateChanges(partitions: Seq[TopicPartition], targetState: PartitionState,
-                           partitionLeaderElectionStrategyOpt: Option[PartitionLeaderElectionStrategy]): Map[TopicPartition, Throwable] = {
+  private def doHandleStateChanges(
+    partitions: Seq[TopicPartition],
+    targetState: PartitionState,
+    partitionLeaderElectionStrategyOpt: Option[PartitionLeaderElectionStrategy]
+  ): Map[TopicPartition, Either[Throwable, LeaderAndIsr]] = {
     val stateChangeLog = stateChangeLogger.withControllerEpoch(controllerContext.epoch)
-    partitions.foreach(partition => partitionState.getOrElseUpdate(partition, NonExistentPartition))
-    val (validPartitions, invalidPartitions) = partitions.partition(partition => isValidTransition(partition, targetState))
+    partitions.foreach(partition => controllerContext.putPartitionStateIfNotExists(partition, NonExistentPartition))
+    val (validPartitions, invalidPartitions) = controllerContext.checkValidPartitionStateChange(partitions, targetState)
     invalidPartitions.foreach(partition => logInvalidTransition(partition, targetState))
+
     targetState match {
       case NewPartition =>
         validPartitions.foreach { partition =>
           stateChangeLog.trace(s"Changed partition $partition state from ${partitionState(partition)} to $targetState with " +
             s"assigned replicas ${controllerContext.partitionReplicaAssignment(partition).mkString(",")}")
-          changeStateTo(partition, partitionState(partition), NewPartition)
+          controllerContext.putPartitionState(partition, NewPartition)
         }
         Map.empty
       case OnlinePartition =>
@@ -218,30 +226,40 @@ class PartitionStateMachine(config: KafkaConfig,
           successfulInitializations.foreach { partition =>
             stateChangeLog.trace(s"Changed partition $partition from ${partitionState(partition)} to $targetState with state " +
               s"${controllerContext.partitionLeadershipInfo(partition).leaderAndIsr}")
-            changeStateTo(partition, partitionState(partition), OnlinePartition)
+            controllerContext.putPartitionState(partition, OnlinePartition)
           }
         }
         if (partitionsToElectLeader.nonEmpty) {
-          val (successfulElections, failedElections) = electLeaderForPartitions(partitionsToElectLeader, partitionLeaderElectionStrategyOpt.get)
-          successfulElections.foreach { partition =>
-            stateChangeLog.trace(s"Changed partition $partition from ${partitionState(partition)} to $targetState with state " +
-              s"${controllerContext.partitionLeadershipInfo(partition).leaderAndIsr}")
-            changeStateTo(partition, partitionState(partition), OnlinePartition)
+          val electionResults = electLeaderForPartitions(
+            partitionsToElectLeader,
+            partitionLeaderElectionStrategyOpt.getOrElse(
+              throw new IllegalArgumentException("Election strategy is a required field when the target state is OnlinePartition")
+            )
+          )
+
+          electionResults.foreach {
+            case (partition, Right(leaderAndIsr)) =>
+              stateChangeLog.trace(
+                s"Changed partition $partition from ${partitionState(partition)} to $targetState with state $leaderAndIsr"
+              )
+              controllerContext.putPartitionState(partition, OnlinePartition)
+            case (_, Left(_)) => // Ignore; no need to update partition state on election error
           }
-          failedElections
+
+          electionResults
         } else {
           Map.empty
         }
       case OfflinePartition =>
         validPartitions.foreach { partition =>
           stateChangeLog.trace(s"Changed partition $partition state from ${partitionState(partition)} to $targetState")
-          changeStateTo(partition, partitionState(partition), OfflinePartition)
+          controllerContext.putPartitionState(partition, OfflinePartition)
         }
         Map.empty
       case NonExistentPartition =>
         validPartitions.foreach { partition =>
           stateChangeLog.trace(s"Changed partition $partition state from ${partitionState(partition)} to $targetState")
-          changeStateTo(partition, partitionState(partition), NonExistentPartition)
+          controllerContext.putPartitionState(partition, NonExistentPartition)
         }
         Map.empty
     }
@@ -290,7 +308,7 @@ class PartitionStateMachine(config: KafkaConfig,
       if (code == Code.OK) {
         controllerContext.partitionLeadershipInfo.put(partition, leaderIsrAndControllerEpoch)
         controllerBrokerRequestBatch.addLeaderAndIsrRequestForBrokers(leaderIsrAndControllerEpoch.leaderAndIsr.isr,
-          partition, leaderIsrAndControllerEpoch, controllerContext.partitionReplicaAssignment(partition), isNew = true)
+          partition, leaderIsrAndControllerEpoch, controllerContext.partitionFullReplicaAssignment(partition), isNew = true)
         successfulInitializations += partition
       } else {
         logFailedStateChange(partition, NewPartition, OnlinePartition, code)
@@ -303,24 +321,33 @@ class PartitionStateMachine(config: KafkaConfig,
    * Repeatedly attempt to elect leaders for multiple partitions until there are no more remaining partitions to retry.
    * @param partitions The partitions that we're trying to elect leaders for.
    * @param partitionLeaderElectionStrategy The election strategy to use.
-   * @return A pair with first element of which is the partitions that successfully had a leader elected
-    *        and the second element a map of failed partition to the corresponding thrown exception.
+   * @return A map of failed and successful elections. The keys are the topic partitions and the corresponding values are
+   *         either the exception that was thrown or new leader & ISR.
    */
-  private def electLeaderForPartitions(partitions: Seq[TopicPartition],
-                                       partitionLeaderElectionStrategy: PartitionLeaderElectionStrategy): (Seq[TopicPartition], Map[TopicPartition, Throwable]) = {
-    val successfulElections = mutable.Buffer.empty[TopicPartition]
+  private def electLeaderForPartitions(
+    partitions: Seq[TopicPartition],
+    partitionLeaderElectionStrategy: PartitionLeaderElectionStrategy
+  ): Map[TopicPartition, Either[Throwable, LeaderAndIsr]] = {
     var remaining = partitions
-    var failures = Map.empty[TopicPartition, Throwable]
+    val finishedElections = mutable.Map.empty[TopicPartition, Either[Throwable, LeaderAndIsr]]
+
     while (remaining.nonEmpty) {
-      val (success, updatesToRetry, failedElections) = doElectLeaderForPartitions(partitions, partitionLeaderElectionStrategy)
+      val (finished, updatesToRetry) = doElectLeaderForPartitions(remaining, partitionLeaderElectionStrategy)
       remaining = updatesToRetry
-      successfulElections ++= success
-      failedElections.foreach { case (partition, e) =>
-        logFailedStateChange(partition, partitionState(partition), OnlinePartition, e)
+
+      finished.foreach {
+        case (partition, Left(e)) =>
+          logFailedStateChange(partition, partitionState(partition), OnlinePartition, e)
+        case (_, Right(_)) => // Ignore; success so no need to log failed state change
       }
-      failures ++= failedElections
+
+      finishedElections ++= finished
+
+      if (remaining.nonEmpty)
+        logger.info(s"Retrying leader election with strategy $partitionLeaderElectionStrategy for partitions $remaining")
     }
-    (successfulElections, failures)
+
+    finishedElections.toMap
   }
 
   /**
@@ -329,154 +356,146 @@ class PartitionStateMachine(config: KafkaConfig,
    *
    * @param partitions The partitions that we're trying to elect leaders for.
    * @param partitionLeaderElectionStrategy The election strategy to use.
-   * @return A tuple of three values:
-   *         1. The partitions that successfully had a leader elected.
+   * @return A tuple of two values:
+   *         1. The partitions and the expected leader and isr that successfully had a leader elected. And exceptions
+   *         corresponding to failed elections that should not be retried.
    *         2. The partitions that we should retry due to a zookeeper BADVERSION conflict. Version conflicts can occur if
    *         the partition leader updated partition state while the controller attempted to update partition state.
-   *         3. Exceptions corresponding to failed elections that should not be retried.
    */
-  private def doElectLeaderForPartitions(partitions: Seq[TopicPartition], partitionLeaderElectionStrategy: PartitionLeaderElectionStrategy):
-  (Seq[TopicPartition], Seq[TopicPartition], Map[TopicPartition, Exception]) = {
+  private def doElectLeaderForPartitions(
+    partitions: Seq[TopicPartition],
+    partitionLeaderElectionStrategy: PartitionLeaderElectionStrategy
+  ): (Map[TopicPartition, Either[Exception, LeaderAndIsr]], Seq[TopicPartition]) = {
     val getDataResponses = try {
       zkClient.getTopicPartitionStatesRaw(partitions)
     } catch {
       case e: Exception =>
-        return (Seq.empty, Seq.empty, partitions.map(_ -> e).toMap)
+        return (partitions.iterator.map(_ -> Left(e)).toMap, Seq.empty)
     }
-    val failedElections = mutable.Map.empty[TopicPartition, Exception]
-    val leaderIsrAndControllerEpochPerPartition = mutable.Buffer.empty[(TopicPartition, LeaderIsrAndControllerEpoch)]
+    val failedElections = mutable.Map.empty[TopicPartition, Either[Exception, LeaderAndIsr]]
+    val validLeaderAndIsrs = mutable.Buffer.empty[(TopicPartition, LeaderAndIsr)]
+
     getDataResponses.foreach { getDataResponse =>
       val partition = getDataResponse.ctx.get.asInstanceOf[TopicPartition]
       val currState = partitionState(partition)
       if (getDataResponse.resultCode == Code.OK) {
-        val leaderIsrAndControllerEpochOpt = TopicPartitionStateZNode.decode(getDataResponse.data, getDataResponse.stat)
-        if (leaderIsrAndControllerEpochOpt.isEmpty) {
-          val exception = new StateChangeFailedException(s"LeaderAndIsr information doesn't exist for partition $partition in $currState state")
-          failedElections.put(partition, exception)
+        TopicPartitionStateZNode.decode(getDataResponse.data, getDataResponse.stat) match {
+          case Some(leaderIsrAndControllerEpoch) =>
+            if (leaderIsrAndControllerEpoch.controllerEpoch > controllerContext.epoch) {
+              val failMsg = s"Aborted leader election for partition $partition since the LeaderAndIsr path was " +
+                s"already written by another controller. This probably means that the current controller $controllerId went through " +
+                s"a soft failure and another controller was elected with epoch ${leaderIsrAndControllerEpoch.controllerEpoch}."
+              failedElections.put(partition, Left(new StateChangeFailedException(failMsg)))
+            } else {
+              validLeaderAndIsrs += partition -> leaderIsrAndControllerEpoch.leaderAndIsr
+            }
+
+          case None =>
+            val exception = new StateChangeFailedException(s"LeaderAndIsr information doesn't exist for partition $partition in $currState state")
+            failedElections.put(partition, Left(exception))
         }
-        leaderIsrAndControllerEpochPerPartition += partition -> leaderIsrAndControllerEpochOpt.get
+
       } else if (getDataResponse.resultCode == Code.NONODE) {
         val exception = new StateChangeFailedException(s"LeaderAndIsr information doesn't exist for partition $partition in $currState state")
-        failedElections.put(partition, exception)
+        failedElections.put(partition, Left(exception))
       } else {
-        failedElections.put(partition, getDataResponse.resultException.get)
+        failedElections.put(partition, Left(getDataResponse.resultException.get))
       }
     }
-    val (invalidPartitionsForElection, validPartitionsForElection) = leaderIsrAndControllerEpochPerPartition.partition { case (_, leaderIsrAndControllerEpoch) =>
-      leaderIsrAndControllerEpoch.controllerEpoch > controllerContext.epoch
-    }
-    invalidPartitionsForElection.foreach { case (partition, leaderIsrAndControllerEpoch) =>
-      val failMsg = s"aborted leader election for partition $partition since the LeaderAndIsr path was " +
-        s"already written by another controller. This probably means that the current controller $controllerId went through " +
-        s"a soft failure and another controller was elected with epoch ${leaderIsrAndControllerEpoch.controllerEpoch}."
-      failedElections.put(partition, new StateChangeFailedException(failMsg))
-    }
-    if (validPartitionsForElection.isEmpty) {
-      return (Seq.empty, Seq.empty, failedElections.toMap)
+
+    if (validLeaderAndIsrs.isEmpty) {
+      return (failedElections.toMap, Seq.empty)
     }
-    val shuttingDownBrokers  = controllerContext.shuttingDownBrokerIds.toSet
+
     val (partitionsWithoutLeaders, partitionsWithLeaders) = partitionLeaderElectionStrategy match {
-      case OfflinePartitionLeaderElectionStrategy =>
-        leaderForOffline(validPartitionsForElection).partition { case (_, newLeaderAndIsrOpt, _) => newLeaderAndIsrOpt.isEmpty }
+      case OfflinePartitionLeaderElectionStrategy(allowUnclean) =>
+        val partitionsWithUncleanLeaderElectionState = collectUncleanLeaderElectionState(
+          validLeaderAndIsrs,
+          allowUnclean
+        )
+        leaderForOffline(controllerContext, partitionsWithUncleanLeaderElectionState).partition(_.leaderAndIsr.isEmpty)
       case ReassignPartitionLeaderElectionStrategy =>
-        leaderForReassign(validPartitionsForElection).partition { case (_, newLeaderAndIsrOpt, _) => newLeaderAndIsrOpt.isEmpty }
+        leaderForReassign(controllerContext, validLeaderAndIsrs).partition(_.leaderAndIsr.isEmpty)
       case PreferredReplicaPartitionLeaderElectionStrategy =>
-        leaderForPreferredReplica(validPartitionsForElection).partition { case (_, newLeaderAndIsrOpt, _) => newLeaderAndIsrOpt.isEmpty }
+        leaderForPreferredReplica(controllerContext, validLeaderAndIsrs).partition(_.leaderAndIsr.isEmpty)
       case ControlledShutdownPartitionLeaderElectionStrategy =>
-        leaderForControlledShutdown(validPartitionsForElection, shuttingDownBrokers).partition { case (_, newLeaderAndIsrOpt, _) => newLeaderAndIsrOpt.isEmpty }
+        leaderForControlledShutdown(controllerContext, validLeaderAndIsrs).partition(_.leaderAndIsr.isEmpty)
     }
-    partitionsWithoutLeaders.foreach { case (partition, _, _) =>
+    partitionsWithoutLeaders.foreach { electionResult =>
+      val partition = electionResult.topicPartition
       val failMsg = s"Failed to elect leader for partition $partition under strategy $partitionLeaderElectionStrategy"
-      failedElections.put(partition, new StateChangeFailedException(failMsg))
+      failedElections.put(partition, Left(new StateChangeFailedException(failMsg)))
     }
-    val recipientsPerPartition = partitionsWithLeaders.map { case (partition, _, recipients) => partition -> recipients }.toMap
-    val adjustedLeaderAndIsrs = partitionsWithLeaders.map { case (partition, leaderAndIsrOpt, _) => partition -> leaderAndIsrOpt.get }.toMap
-    val UpdateLeaderAndIsrResult(successfulUpdates, updatesToRetry, failedUpdates) = zkClient.updateLeaderAndIsr(
+    val recipientsPerPartition = partitionsWithLeaders.map(result => result.topicPartition -> result.liveReplicas).toMap
+    val adjustedLeaderAndIsrs = partitionsWithLeaders.map(result => result.topicPartition -> result.leaderAndIsr.get).toMap
+    val UpdateLeaderAndIsrResult(finishedUpdates, updatesToRetry) = zkClient.updateLeaderAndIsr(
       adjustedLeaderAndIsrs, controllerContext.epoch, controllerContext.epochZkVersion)
-    successfulUpdates.foreach { case (partition, leaderAndIsr) =>
-      val replicas = controllerContext.partitionReplicaAssignment(partition)
-      val leaderIsrAndControllerEpoch = LeaderIsrAndControllerEpoch(leaderAndIsr, controllerContext.epoch)
-      controllerContext.partitionLeadershipInfo.put(partition, leaderIsrAndControllerEpoch)
-      controllerBrokerRequestBatch.addLeaderAndIsrRequestForBrokers(recipientsPerPartition(partition), partition,
-        leaderIsrAndControllerEpoch, replicas, isNew = false)
+    finishedUpdates.foreach { case (partition, result) =>
+      result.right.foreach { leaderAndIsr =>
+        val replicaAssignment = controllerContext.partitionFullReplicaAssignment(partition)
+        val leaderIsrAndControllerEpoch = LeaderIsrAndControllerEpoch(leaderAndIsr, controllerContext.epoch)
+        controllerContext.partitionLeadershipInfo.put(partition, leaderIsrAndControllerEpoch)
+        controllerBrokerRequestBatch.addLeaderAndIsrRequestForBrokers(recipientsPerPartition(partition), partition,
+          leaderIsrAndControllerEpoch, replicaAssignment, isNew = false)
+      }
     }
-    (successfulUpdates.keys.toSeq, updatesToRetry, failedElections.toMap ++ failedUpdates)
+
+    (finishedUpdates ++ failedElections, updatesToRetry)
   }
 
-  private def leaderForOffline(leaderIsrAndControllerEpochs: Seq[(TopicPartition, LeaderIsrAndControllerEpoch)]):
-  Seq[(TopicPartition, Option[LeaderAndIsr], Seq[Int])] = {
-    val (partitionsWithNoLiveInSyncReplicas, partitionsWithLiveInSyncReplicas) = leaderIsrAndControllerEpochs.partition { case (partition, leaderIsrAndControllerEpoch) =>
-      val liveInSyncReplicas = leaderIsrAndControllerEpoch.leaderAndIsr.isr.filter(replica => controllerContext.isReplicaOnline(replica, partition))
-      liveInSyncReplicas.isEmpty
+  /* For the provided set of topic partition and partition sync state it attempts to determine if unclean
+   * leader election should be performed. Unclean election should be performed if there are no live
+   * replica which are in sync and unclean leader election is allowed (allowUnclean parameter is true or
+   * the topic has been configured to allow unclean election).
+   *
+   * @param leaderIsrAndControllerEpochs set of partition to determine if unclean leader election should be
+   *                                     allowed
+   * @param allowUnclean whether to allow unclean election without having to read the topic configuration
+   * @return a sequence of three element tuple:
+   *         1. topic partition
+   *         2. leader, isr and controller epoc. Some means election should be performed
+   *         3. allow unclean
+   */
+  private def collectUncleanLeaderElectionState(
+    leaderAndIsrs: Seq[(TopicPartition, LeaderAndIsr)],
+    allowUnclean: Boolean
+  ): Seq[(TopicPartition, Option[LeaderAndIsr], Boolean)] = {
+    val (partitionsWithNoLiveInSyncReplicas, partitionsWithLiveInSyncReplicas) = leaderAndIsrs.partition {
+      case (partition, leaderAndIsr) =>
+        val liveInSyncReplicas = leaderAndIsr.isr.filter(controllerContext.isReplicaOnline(_, partition))
+        liveInSyncReplicas.isEmpty
     }
-    val (logConfigs, failed) = zkClient.getLogConfigs(partitionsWithNoLiveInSyncReplicas.map { case (partition, _) => partition.topic }, config.originals())
-    val partitionsWithUncleanLeaderElectionState = partitionsWithNoLiveInSyncReplicas.map { case (partition, leaderIsrAndControllerEpoch) =>
-      if (failed.contains(partition.topic)) {
-        logFailedStateChange(partition, partitionState(partition), OnlinePartition, failed(partition.topic))
-        (partition, None, false)
-      } else {
-        (partition, Option(leaderIsrAndControllerEpoch), logConfigs(partition.topic).uncleanLeaderElectionEnable.booleanValue())
+
+    val electionForPartitionWithoutLiveReplicas = if (allowUnclean) {
+      partitionsWithNoLiveInSyncReplicas.map { case (partition, leaderAndIsr) =>
+        (partition, Option(leaderAndIsr), true)
       }
-    } ++ partitionsWithLiveInSyncReplicas.map { case (partition, leaderIsrAndControllerEpoch) => (partition, Option(leaderIsrAndControllerEpoch), false) }
-    partitionsWithUncleanLeaderElectionState.map { case (partition, leaderIsrAndControllerEpochOpt, uncleanLeaderElectionEnabled) =>
-      val assignment = controllerContext.partitionReplicaAssignment(partition)
-      val liveReplicas = assignment.filter(replica => controllerContext.isReplicaOnline(replica, partition))
-      if (leaderIsrAndControllerEpochOpt.nonEmpty) {
-        val leaderIsrAndControllerEpoch = leaderIsrAndControllerEpochOpt.get
-        val isr = leaderIsrAndControllerEpoch.leaderAndIsr.isr
-        val leaderOpt = PartitionLeaderElectionAlgorithms.offlinePartitionLeaderElection(assignment, isr, liveReplicas.toSet, uncleanLeaderElectionEnabled, controllerContext)
-        val newLeaderAndIsrOpt = leaderOpt.map { leader =>
-          val newIsr = if (isr.contains(leader)) isr.filter(replica => controllerContext.isReplicaOnline(replica, partition))
-          else List(leader)
-          leaderIsrAndControllerEpoch.leaderAndIsr.newLeaderAndIsr(leader, newIsr)
+    } else {
+      val (logConfigs, failed) = zkClient.getLogConfigs(
+        partitionsWithNoLiveInSyncReplicas.iterator.map { case (partition, _) => partition.topic }.toSet,
+        config.originals()
+      )
+
+      partitionsWithNoLiveInSyncReplicas.map { case (partition, leaderAndIsr) =>
+        if (failed.contains(partition.topic)) {
+          logFailedStateChange(partition, partitionState(partition), OnlinePartition, failed(partition.topic))
+          (partition, None, false)
+        } else {
+          (
+            partition,
+            Option(leaderAndIsr),
+            logConfigs(partition.topic).uncleanLeaderElectionEnable.booleanValue()
+          )
         }
-        (partition, newLeaderAndIsrOpt, liveReplicas)
-      } else {
-        (partition, None, liveReplicas)
       }
     }
-  }
-
-  private def leaderForReassign(leaderIsrAndControllerEpochs: Seq[(TopicPartition, LeaderIsrAndControllerEpoch)]):
-  Seq[(TopicPartition, Option[LeaderAndIsr], Seq[Int])] = {
-    leaderIsrAndControllerEpochs.map { case (partition, leaderIsrAndControllerEpoch) =>
-      val reassignment = controllerContext.partitionsBeingReassigned(partition).newReplicas
-      val liveReplicas = reassignment.filter(replica => controllerContext.isReplicaOnline(replica, partition))
-      val isr = leaderIsrAndControllerEpoch.leaderAndIsr.isr
-      val leaderOpt = PartitionLeaderElectionAlgorithms.reassignPartitionLeaderElection(reassignment, isr, liveReplicas.toSet)
-      val newLeaderAndIsrOpt = leaderOpt.map(leader => leaderIsrAndControllerEpoch.leaderAndIsr.newLeader(leader))
-      (partition, newLeaderAndIsrOpt, reassignment)
-    }
-  }
-
-  private def leaderForPreferredReplica(leaderIsrAndControllerEpochs: Seq[(TopicPartition, LeaderIsrAndControllerEpoch)]):
-  Seq[(TopicPartition, Option[LeaderAndIsr], Seq[Int])] = {
-    leaderIsrAndControllerEpochs.map { case (partition, leaderIsrAndControllerEpoch) =>
-      val assignment = controllerContext.partitionReplicaAssignment(partition)
-      val liveReplicas = assignment.filter(replica => controllerContext.isReplicaOnline(replica, partition))
-      val isr = leaderIsrAndControllerEpoch.leaderAndIsr.isr
-      val leaderOpt = PartitionLeaderElectionAlgorithms.preferredReplicaPartitionLeaderElection(assignment, isr, liveReplicas.toSet)
-      val newLeaderAndIsrOpt = leaderOpt.map(leader => leaderIsrAndControllerEpoch.leaderAndIsr.newLeader(leader))
-      (partition, newLeaderAndIsrOpt, assignment)
-    }
-  }
 
-  private def leaderForControlledShutdown(leaderIsrAndControllerEpochs: Seq[(TopicPartition, LeaderIsrAndControllerEpoch)], shuttingDownBrokers: Set[Int]):
-  Seq[(TopicPartition, Option[LeaderAndIsr], Seq[Int])] = {
-    leaderIsrAndControllerEpochs.map { case (partition, leaderIsrAndControllerEpoch) =>
-      val assignment = controllerContext.partitionReplicaAssignment(partition)
-      val liveOrShuttingDownReplicas = assignment.filter(replica => controllerContext.isReplicaOnline(replica, partition, includeShuttingDownBrokers = true))
-      val isr = leaderIsrAndControllerEpoch.leaderAndIsr.isr
-      val leaderOpt = PartitionLeaderElectionAlgorithms.controlledShutdownPartitionLeaderElection(assignment, isr, liveOrShuttingDownReplicas.toSet, shuttingDownBrokers)
-      val newIsr = isr.filter(replica => !controllerContext.shuttingDownBrokerIds.contains(replica))
-      val newLeaderAndIsrOpt = leaderOpt.map(leader => leaderIsrAndControllerEpoch.leaderAndIsr.newLeaderAndIsr(leader, newIsr))
-      (partition, newLeaderAndIsrOpt, liveOrShuttingDownReplicas)
+    electionForPartitionWithoutLiveReplicas ++
+    partitionsWithLiveInSyncReplicas.map { case (partition, leaderAndIsr) =>
+      (partition, Option(leaderAndIsr), false)
     }
   }
 
-  private def isValidTransition(partition: TopicPartition, targetState: PartitionState) =
-    targetState.validPreviousStates.contains(partitionState(partition))
-
   private def logInvalidTransition(partition: TopicPartition, targetState: PartitionState): Unit = {
     val currState = partitionState(partition)
     val e = new IllegalStateException(s"Partition $partition should be in one of " +
@@ -497,16 +516,22 @@ class PartitionStateMachine(config: KafkaConfig,
 }
 
 object PartitionLeaderElectionAlgorithms {
-  def offlinePartitionLeaderElection(assignment: Seq[Int], isr: Seq[Int], liveReplicas: Set[Int], uncleanLeaderElectionEnabled: Boolean, controllerContext: ControllerContext): Option[Int] = {
-    assignment.find(id => liveReplicas.contains(id) && isr.contains(id)).orElse {
-      if (uncleanLeaderElectionEnabled) {
-        val leaderOpt = assignment.find(liveReplicas.contains)
-        if (!leaderOpt.isEmpty)
-          controllerContext.stats.uncleanLeaderElectionRate.mark()
-        leaderOpt
-      } else {
-        None
-      }
+
+  /**
+   * @return Optionally, a tuple (replica, flag) where flag is a boolean indicating if unclean leader election was
+   *         used to replace the replica.
+   */
+  def offlinePartitionLeaderElection(assignment: Seq[Int], isr: Seq[Int], liveReplicas: Set[Int], uncleanLeaderElectionEnabled: Boolean): Option[(Int, Boolean)] = {
+    assignment.find(id => liveReplicas.contains(id) && isr.contains(id)) match {
+      case Some(replicaId) => Some(replicaId, false)
+      case None =>  if (uncleanLeaderElectionEnabled) {
+                     assignment.find(liveReplicas.contains) match {
+                        case Some(uncleanReplicaId) => Some(uncleanReplicaId, true)
+                        case None => None
+                      }
+                  } else {
+                    None
+                  }
     }
   }
 
@@ -524,10 +549,10 @@ object PartitionLeaderElectionAlgorithms {
 }
 
 sealed trait PartitionLeaderElectionStrategy
-case object OfflinePartitionLeaderElectionStrategy extends PartitionLeaderElectionStrategy
-case object ReassignPartitionLeaderElectionStrategy extends PartitionLeaderElectionStrategy
-case object PreferredReplicaPartitionLeaderElectionStrategy extends PartitionLeaderElectionStrategy
-case object ControlledShutdownPartitionLeaderElectionStrategy extends PartitionLeaderElectionStrategy
+final case class OfflinePartitionLeaderElectionStrategy(allowUnclean: Boolean) extends PartitionLeaderElectionStrategy
+final case object ReassignPartitionLeaderElectionStrategy extends PartitionLeaderElectionStrategy
+final case object PreferredReplicaPartitionLeaderElectionStrategy extends PartitionLeaderElectionStrategy
+final case object ControlledShutdownPartitionLeaderElectionStrategy extends PartitionLeaderElectionStrategy
 
 sealed trait PartitionState {
   def state: Byte
diff --git a/core/src/main/scala/kafka/controller/ReplicaStateMachine.scala b/core/src/main/scala/kafka/controller/ReplicaStateMachine.scala
index 433ab5668379e..5fada1b326b99 100644
--- a/core/src/main/scala/kafka/controller/ReplicaStateMachine.scala
+++ b/core/src/main/scala/kafka/controller/ReplicaStateMachine.scala
@@ -20,59 +20,33 @@ import kafka.api.LeaderAndIsr
 import kafka.common.StateChangeFailedException
 import kafka.server.KafkaConfig
 import kafka.utils.Logging
-import kafka.zk.{KafkaZkClient, TopicPartitionStateZNode}
+import kafka.zk.KafkaZkClient
 import kafka.zk.KafkaZkClient.UpdateLeaderAndIsrResult
+import kafka.zk.TopicPartitionStateZNode
 import org.apache.kafka.common.TopicPartition
 import org.apache.kafka.common.errors.ControllerMovedException
 import org.apache.zookeeper.KeeperException.Code
+import scala.collection.{Seq, mutable}
 
-import scala.collection.mutable
-
-/**
- * This class represents the state machine for replicas. It defines the states that a replica can be in, and
- * transitions to move the replica to another legal state. The different states that a replica can be in are -
- * 1. NewReplica        : The controller can create new replicas during partition reassignment. In this state, a
- *                        replica can only get become follower state change request.  Valid previous
- *                        state is NonExistentReplica
- * 2. OnlineReplica     : Once a replica is started and part of the assigned replicas for its partition, it is in this
- *                        state. In this state, it can get either become leader or become follower state change requests.
- *                        Valid previous state are NewReplica, OnlineReplica or OfflineReplica
- * 3. OfflineReplica    : If a replica dies, it moves to this state. This happens when the broker hosting the replica
- *                        is down. Valid previous state are NewReplica, OnlineReplica
- * 4. ReplicaDeletionStarted: If replica deletion starts, it is moved to this state. Valid previous state is OfflineReplica
- * 5. ReplicaDeletionSuccessful: If replica responds with no error code in response to a delete replica request, it is
- *                        moved to this state. Valid previous state is ReplicaDeletionStarted
- * 6. ReplicaDeletionIneligible: If replica deletion fails, it is moved to this state. Valid previous state is ReplicaDeletionStarted
- * 7. NonExistentReplica: If a replica is deleted successfully, it is moved to this state. Valid previous state is
- *                        ReplicaDeletionSuccessful
- */
-class ReplicaStateMachine(config: KafkaConfig,
-                          stateChangeLogger: StateChangeLogger,
-                          controllerContext: ControllerContext,
-                          topicDeletionManager: TopicDeletionManager,
-                          zkClient: KafkaZkClient,
-                          replicaState: mutable.Map[PartitionAndReplica, ReplicaState],
-                          controllerBrokerRequestBatch: ControllerBrokerRequestBatch) extends Logging {
-  private val controllerId = config.brokerId
-
-  this.logIdent = s"[ReplicaStateMachine controllerId=$controllerId] "
-
+abstract class ReplicaStateMachine(controllerContext: ControllerContext) extends Logging {
   /**
    * Invoked on successful controller election.
    */
-  def startup() {
+  def startup(): Unit = {
     info("Initializing replica state")
     initializeReplicaState()
     info("Triggering online replica state changes")
-    handleStateChanges(controllerContext.allLiveReplicas().toSeq, OnlineReplica)
-    info(s"Started replica state machine with initial state -> $replicaState")
+    val (onlineReplicas, offlineReplicas) = controllerContext.onlineAndOfflineReplicas
+    handleStateChanges(onlineReplicas.toSeq, OnlineReplica)
+    info("Triggering offline replica state changes")
+    handleStateChanges(offlineReplicas.toSeq, OfflineReplica)
+    debug(s"Started replica state machine with initial state -> ${controllerContext.replicaStates}")
   }
 
   /**
    * Invoked on controller shutdown.
    */
-  def shutdown() {
-    replicaState.clear()
+  def shutdown(): Unit = {
     info("Stopped replica state machine")
   }
 
@@ -80,30 +54,61 @@ class ReplicaStateMachine(config: KafkaConfig,
    * Invoked on startup of the replica's state machine to set the initial state for replicas of all existing partitions
    * in zookeeper
    */
-  private def initializeReplicaState() {
+  private def initializeReplicaState(): Unit = {
     controllerContext.allPartitions.foreach { partition =>
       val replicas = controllerContext.partitionReplicaAssignment(partition)
       replicas.foreach { replicaId =>
         val partitionAndReplica = PartitionAndReplica(partition, replicaId)
-        if (controllerContext.isReplicaOnline(replicaId, partition))
-          replicaState.put(partitionAndReplica, OnlineReplica)
-        else
-        // mark replicas on dead brokers as failed for topic deletion, if they belong to a topic to be deleted.
-        // This is required during controller failover since during controller failover a broker can go down,
-        // so the replicas on that broker should be moved to ReplicaDeletionIneligible to be on the safer side.
-          replicaState.put(partitionAndReplica, ReplicaDeletionIneligible)
+        if (controllerContext.isReplicaOnline(replicaId, partition)) {
+          controllerContext.putReplicaState(partitionAndReplica, OnlineReplica)
+        } else {
+          // mark replicas on dead brokers as failed for topic deletion, if they belong to a topic to be deleted.
+          // This is required during controller failover since during controller failover a broker can go down,
+          // so the replicas on that broker should be moved to ReplicaDeletionIneligible to be on the safer side.
+          controllerContext.putReplicaState(partitionAndReplica, ReplicaDeletionIneligible)
+        }
       }
     }
   }
 
-  def handleStateChanges(replicas: Seq[PartitionAndReplica], targetState: ReplicaState,
-                         callbacks: Callbacks = new Callbacks()): Unit = {
+  def handleStateChanges(replicas: Seq[PartitionAndReplica], targetState: ReplicaState): Unit
+}
+
+/**
+ * This class represents the state machine for replicas. It defines the states that a replica can be in, and
+ * transitions to move the replica to another legal state. The different states that a replica can be in are -
+ * 1. NewReplica        : The controller can create new replicas during partition reassignment. In this state, a
+ *                        replica can only get become follower state change request.  Valid previous
+ *                        state is NonExistentReplica
+ * 2. OnlineReplica     : Once a replica is started and part of the assigned replicas for its partition, it is in this
+ *                        state. In this state, it can get either become leader or become follower state change requests.
+ *                        Valid previous state are NewReplica, OnlineReplica or OfflineReplica
+ * 3. OfflineReplica    : If a replica dies, it moves to this state. This happens when the broker hosting the replica
+ *                        is down. Valid previous state are NewReplica, OnlineReplica
+ * 4. ReplicaDeletionStarted: If replica deletion starts, it is moved to this state. Valid previous state is OfflineReplica
+ * 5. ReplicaDeletionSuccessful: If replica responds with no error code in response to a delete replica request, it is
+ *                        moved to this state. Valid previous state is ReplicaDeletionStarted
+ * 6. ReplicaDeletionIneligible: If replica deletion fails, it is moved to this state. Valid previous states are
+ *                        ReplicaDeletionStarted and OfflineReplica
+ * 7. NonExistentReplica: If a replica is deleted successfully, it is moved to this state. Valid previous state is
+ *                        ReplicaDeletionSuccessful
+ */
+class ZkReplicaStateMachine(config: KafkaConfig,
+                            stateChangeLogger: StateChangeLogger,
+                            controllerContext: ControllerContext,
+                            zkClient: KafkaZkClient,
+                            controllerBrokerRequestBatch: ControllerBrokerRequestBatch)
+  extends ReplicaStateMachine(controllerContext) with Logging {
+
+  private val controllerId = config.brokerId
+  this.logIdent = s"[ReplicaStateMachine controllerId=$controllerId] "
+
+  override def handleStateChanges(replicas: Seq[PartitionAndReplica], targetState: ReplicaState): Unit = {
     if (replicas.nonEmpty) {
       try {
         controllerBrokerRequestBatch.newBatch()
-        replicas.groupBy(_.replica).map { case (replicaId, replicas) =>
-          val partitions = replicas.map(_.topicPartition)
-          doHandleStateChanges(replicaId, partitions, targetState, callbacks)
+        replicas.groupBy(_.replica).foreach { case (replicaId, replicas) =>
+          doHandleStateChanges(replicaId, replicas, targetState)
         }
         controllerBrokerRequestBatch.sendRequestsToBrokers(controllerContext.epoch)
       } catch {
@@ -147,42 +152,45 @@ class ReplicaStateMachine(config: KafkaConfig,
    * -- remove the replica from the in memory partition replica assignment cache
    *
    * @param replicaId The replica for which the state transition is invoked
-   * @param partitions The partitions on this replica for which the state transition is invoked
+   * @param replicas The partitions on this replica for which the state transition is invoked
    * @param targetState The end state that the replica should be moved to
    */
-  private def doHandleStateChanges(replicaId: Int, partitions: Seq[TopicPartition], targetState: ReplicaState,
-                                   callbacks: Callbacks): Unit = {
-    val replicas = partitions.map(partition => PartitionAndReplica(partition, replicaId))
-    replicas.foreach(replica => replicaState.getOrElseUpdate(replica, NonExistentReplica))
-    val (validReplicas, invalidReplicas) = replicas.partition(replica => isValidTransition(replica, targetState))
+  private def doHandleStateChanges(replicaId: Int, replicas: Seq[PartitionAndReplica], targetState: ReplicaState): Unit = {
+    replicas.foreach(replica => controllerContext.putReplicaStateIfNotExists(replica, NonExistentReplica))
+    val (validReplicas, invalidReplicas) = controllerContext.checkValidReplicaStateChange(replicas, targetState)
     invalidReplicas.foreach(replica => logInvalidTransition(replica, targetState))
+
     targetState match {
       case NewReplica =>
         validReplicas.foreach { replica =>
           val partition = replica.topicPartition
+          val currentState = controllerContext.replicaState(replica)
+
           controllerContext.partitionLeadershipInfo.get(partition) match {
             case Some(leaderIsrAndControllerEpoch) =>
               if (leaderIsrAndControllerEpoch.leaderAndIsr.leader == replicaId) {
                 val exception = new StateChangeFailedException(s"Replica $replicaId for partition $partition cannot be moved to NewReplica state as it is being requested to become leader")
-                logFailedStateChange(replica, replicaState(replica), OfflineReplica, exception)
+                logFailedStateChange(replica, currentState, OfflineReplica, exception)
               } else {
                 controllerBrokerRequestBatch.addLeaderAndIsrRequestForBrokers(Seq(replicaId),
                   replica.topicPartition,
                   leaderIsrAndControllerEpoch,
-                  controllerContext.partitionReplicaAssignment(replica.topicPartition),
+                  controllerContext.partitionFullReplicaAssignment(replica.topicPartition),
                   isNew = true)
-                logSuccessfulTransition(replicaId, partition, replicaState(replica), NewReplica)
-                replicaState.put(replica, NewReplica)
+                logSuccessfulTransition(replicaId, partition, currentState, NewReplica)
+                controllerContext.putReplicaState(replica, NewReplica)
               }
             case None =>
-              logSuccessfulTransition(replicaId, partition, replicaState(replica), NewReplica)
-              replicaState.put(replica, NewReplica)
+              logSuccessfulTransition(replicaId, partition, currentState, NewReplica)
+              controllerContext.putReplicaState(replica, NewReplica)
           }
         }
       case OnlineReplica =>
         validReplicas.foreach { replica =>
           val partition = replica.topicPartition
-          replicaState(replica) match {
+          val currentState = controllerContext.replicaState(replica)
+
+          currentState match {
             case NewReplica =>
               val assignment = controllerContext.partitionReplicaAssignment(partition)
               if (!assignment.contains(replicaId)) {
@@ -194,64 +202,67 @@ class ReplicaStateMachine(config: KafkaConfig,
                   controllerBrokerRequestBatch.addLeaderAndIsrRequestForBrokers(Seq(replicaId),
                     replica.topicPartition,
                     leaderIsrAndControllerEpoch,
-                    controllerContext.partitionReplicaAssignment(partition), isNew = false)
+                    controllerContext.partitionFullReplicaAssignment(partition), isNew = false)
                 case None =>
               }
           }
-          logSuccessfulTransition(replicaId, partition, replicaState(replica), OnlineReplica)
-          replicaState.put(replica, OnlineReplica)
+          logSuccessfulTransition(replicaId, partition, currentState, OnlineReplica)
+          controllerContext.putReplicaState(replica, OnlineReplica)
         }
       case OfflineReplica =>
         validReplicas.foreach { replica =>
-          controllerBrokerRequestBatch.addStopReplicaRequestForBrokers(Seq(replicaId), replica.topicPartition,
-            deletePartition = false, (_, _) => ())
+          controllerBrokerRequestBatch.addStopReplicaRequestForBrokers(Seq(replicaId), replica.topicPartition, deletePartition = false)
         }
         val (replicasWithLeadershipInfo, replicasWithoutLeadershipInfo) = validReplicas.partition { replica =>
           controllerContext.partitionLeadershipInfo.contains(replica.topicPartition)
         }
         val updatedLeaderIsrAndControllerEpochs = removeReplicasFromIsr(replicaId, replicasWithLeadershipInfo.map(_.topicPartition))
         updatedLeaderIsrAndControllerEpochs.foreach { case (partition, leaderIsrAndControllerEpoch) =>
-          if (!topicDeletionManager.isTopicQueuedUpForDeletion(partition.topic)) {
+          if (!controllerContext.isTopicQueuedUpForDeletion(partition.topic)) {
             val recipients = controllerContext.partitionReplicaAssignment(partition).filterNot(_ == replicaId)
             controllerBrokerRequestBatch.addLeaderAndIsrRequestForBrokers(recipients,
               partition,
               leaderIsrAndControllerEpoch,
-              controllerContext.partitionReplicaAssignment(partition), isNew = false)
+              controllerContext.partitionFullReplicaAssignment(partition), isNew = false)
           }
           val replica = PartitionAndReplica(partition, replicaId)
-          logSuccessfulTransition(replicaId, partition, replicaState(replica), OfflineReplica)
-          replicaState.put(replica, OfflineReplica)
+          val currentState = controllerContext.replicaState(replica)
+          logSuccessfulTransition(replicaId, partition, currentState, OfflineReplica)
+          controllerContext.putReplicaState(replica, OfflineReplica)
         }
 
         replicasWithoutLeadershipInfo.foreach { replica =>
-          logSuccessfulTransition(replicaId, replica.topicPartition, replicaState(replica), OfflineReplica)
-          replicaState.put(replica, OfflineReplica)
+          val currentState = controllerContext.replicaState(replica)
+          logSuccessfulTransition(replicaId, replica.topicPartition, currentState, OfflineReplica)
+          controllerBrokerRequestBatch.addUpdateMetadataRequestForBrokers(controllerContext.liveOrShuttingDownBrokerIds.toSeq, Set(replica.topicPartition))
+          controllerContext.putReplicaState(replica, OfflineReplica)
         }
       case ReplicaDeletionStarted =>
         validReplicas.foreach { replica =>
-          logSuccessfulTransition(replicaId, replica.topicPartition, replicaState(replica), ReplicaDeletionStarted)
-          replicaState.put(replica, ReplicaDeletionStarted)
-          controllerBrokerRequestBatch.addStopReplicaRequestForBrokers(Seq(replicaId),
-            replica.topicPartition,
-            deletePartition = true,
-            callbacks.stopReplicaResponseCallback)
+          val currentState = controllerContext.replicaState(replica)
+          logSuccessfulTransition(replicaId, replica.topicPartition, currentState, ReplicaDeletionStarted)
+          controllerContext.putReplicaState(replica, ReplicaDeletionStarted)
+          controllerBrokerRequestBatch.addStopReplicaRequestForBrokers(Seq(replicaId), replica.topicPartition, deletePartition = true)
         }
       case ReplicaDeletionIneligible =>
         validReplicas.foreach { replica =>
-          logSuccessfulTransition(replicaId, replica.topicPartition, replicaState(replica), ReplicaDeletionIneligible)
-          replicaState.put(replica, ReplicaDeletionIneligible)
+          val currentState = controllerContext.replicaState(replica)
+          logSuccessfulTransition(replicaId, replica.topicPartition, currentState, ReplicaDeletionIneligible)
+          controllerContext.putReplicaState(replica, ReplicaDeletionIneligible)
         }
       case ReplicaDeletionSuccessful =>
         validReplicas.foreach { replica =>
-          logSuccessfulTransition(replicaId, replica.topicPartition, replicaState(replica), ReplicaDeletionSuccessful)
-          replicaState.put(replica, ReplicaDeletionSuccessful)
+          val currentState = controllerContext.replicaState(replica)
+          logSuccessfulTransition(replicaId, replica.topicPartition, currentState, ReplicaDeletionSuccessful)
+          controllerContext.putReplicaState(replica, ReplicaDeletionSuccessful)
         }
       case NonExistentReplica =>
         validReplicas.foreach { replica =>
+          val currentState = controllerContext.replicaState(replica)
           val currentAssignedReplicas = controllerContext.partitionReplicaAssignment(replica.topicPartition)
           controllerContext.updatePartitionReplicaAssignment(replica.topicPartition, currentAssignedReplicas.filterNot(_ == replica.replica))
-          logSuccessfulTransition(replicaId, replica.topicPartition, replicaState(replica), NonExistentReplica)
-          replicaState.remove(replica)
+          logSuccessfulTransition(replicaId, replica.topicPartition, currentState, NonExistentReplica)
+          controllerContext.removeReplicaState(replica)
         }
     }
   }
@@ -263,17 +274,23 @@ class ReplicaStateMachine(config: KafkaConfig,
    * @param partitions The partitions from which we're trying to remove the replica from isr
    * @return The updated LeaderIsrAndControllerEpochs of all partitions for which we successfully removed the replica from isr.
    */
-  private def removeReplicasFromIsr(replicaId: Int, partitions: Seq[TopicPartition]):
-  Map[TopicPartition, LeaderIsrAndControllerEpoch] = {
+  private def removeReplicasFromIsr(
+    replicaId: Int,
+    partitions: Seq[TopicPartition]
+  ): Map[TopicPartition, LeaderIsrAndControllerEpoch] = {
     var results = Map.empty[TopicPartition, LeaderIsrAndControllerEpoch]
     var remaining = partitions
     while (remaining.nonEmpty) {
-      val (successfulRemovals, removalsToRetry, failedRemovals) = doRemoveReplicasFromIsr(replicaId, remaining)
-      results ++= successfulRemovals
+      val (finishedRemoval, removalsToRetry) = doRemoveReplicasFromIsr(replicaId, remaining)
       remaining = removalsToRetry
-      failedRemovals.foreach { case (partition, e) =>
-        val replica = PartitionAndReplica(partition, replicaId)
-        logFailedStateChange(replica, replicaState(replica), OfflineReplica, e)
+
+      finishedRemoval.foreach {
+        case (partition, Left(e)) =>
+            val replica = PartitionAndReplica(partition, replicaId)
+            val currentState = controllerContext.replicaState(replica)
+            logFailedStateChange(replica, currentState, OfflineReplica, e)
+        case (partition, Right(leaderIsrAndEpoch)) =>
+          results += partition -> leaderIsrAndEpoch
       }
     }
     results
@@ -285,114 +302,126 @@ class ReplicaStateMachine(config: KafkaConfig,
    *
    * @param replicaId The replica being removed from isr of multiple partitions
    * @param partitions The partitions from which we're trying to remove the replica from isr
-   * @return A tuple of three values:
-   *         1. The updated LeaderIsrAndControllerEpochs of all partitions for which we successfully removed the replica from isr.
+   * @return A tuple of two elements:
+   *         1. The updated Right[LeaderIsrAndControllerEpochs] of all partitions for which we successfully
+   *         removed the replica from isr. Or Left[Exception] corresponding to failed removals that should
+   *         not be retried
    *         2. The partitions that we should retry due to a zookeeper BADVERSION conflict. Version conflicts can occur if
    *         the partition leader updated partition state while the controller attempted to update partition state.
-   *         3. Exceptions corresponding to failed removals that should not be retried.
    */
-  private def doRemoveReplicasFromIsr(replicaId: Int, partitions: Seq[TopicPartition]):
-  (Map[TopicPartition, LeaderIsrAndControllerEpoch],
-    Seq[TopicPartition],
-    Map[TopicPartition, Exception]) = {
-    val (leaderAndIsrs, partitionsWithNoLeaderAndIsrInZk, failedStateReads) = getTopicPartitionStatesFromZk(partitions)
-    val (leaderAndIsrsWithReplica, leaderAndIsrsWithoutReplica) = leaderAndIsrs.partition { case (_, leaderAndIsr) => leaderAndIsr.isr.contains(replicaId) }
-    val adjustedLeaderAndIsrs = leaderAndIsrsWithReplica.mapValues { leaderAndIsr =>
-      val newLeader = if (replicaId == leaderAndIsr.leader) LeaderAndIsr.NoLeader else leaderAndIsr.leader
-      val adjustedIsr = if (leaderAndIsr.isr.size == 1) leaderAndIsr.isr else leaderAndIsr.isr.filter(_ != replicaId)
-      leaderAndIsr.newLeaderAndIsr(newLeader, adjustedIsr)
+  private def doRemoveReplicasFromIsr(
+    replicaId: Int,
+    partitions: Seq[TopicPartition]
+  ): (Map[TopicPartition, Either[Exception, LeaderIsrAndControllerEpoch]], Seq[TopicPartition]) = {
+    val (leaderAndIsrs, partitionsWithNoLeaderAndIsrInZk) = getTopicPartitionStatesFromZk(partitions)
+    val (leaderAndIsrsWithReplica, leaderAndIsrsWithoutReplica) = leaderAndIsrs.partition { case (_, result) =>
+      result.right.map { leaderAndIsr =>
+        leaderAndIsr.isr.contains(replicaId)
+      }.right.getOrElse(false)
     }
-    val UpdateLeaderAndIsrResult(successfulUpdates, updatesToRetry, failedUpdates) = zkClient.updateLeaderAndIsr(
-      adjustedLeaderAndIsrs, controllerContext.epoch, controllerContext.epochZkVersion)
-    val exceptionsForPartitionsWithNoLeaderAndIsrInZk = partitionsWithNoLeaderAndIsrInZk.flatMap { partition =>
-      if (!topicDeletionManager.isTopicQueuedUpForDeletion(partition.topic)) {
-        val exception = new StateChangeFailedException(s"Failed to change state of replica $replicaId for partition $partition since the leader and isr path in zookeeper is empty")
-        Option(partition -> exception)
-      } else None
-    }.toMap
-    val leaderIsrAndControllerEpochs = (leaderAndIsrsWithoutReplica ++ successfulUpdates).map { case (partition, leaderAndIsr) =>
-      val leaderIsrAndControllerEpoch = LeaderIsrAndControllerEpoch(leaderAndIsr, controllerContext.epoch)
-      controllerContext.partitionLeadershipInfo.put(partition, leaderIsrAndControllerEpoch)
-      partition -> leaderIsrAndControllerEpoch
+
+    val adjustedLeaderAndIsrs: Map[TopicPartition, LeaderAndIsr] = leaderAndIsrsWithReplica.flatMap {
+      case (partition, result) =>
+        result.right.toOption.map { leaderAndIsr =>
+          val newLeader = if (replicaId == leaderAndIsr.leader) LeaderAndIsr.NoLeader else leaderAndIsr.leader
+          val adjustedIsr = if (leaderAndIsr.isr.size == 1) leaderAndIsr.isr else leaderAndIsr.isr.filter(_ != replicaId)
+          partition -> leaderAndIsr.newLeaderAndIsr(newLeader, adjustedIsr)
+        }
     }
-    (leaderIsrAndControllerEpochs, updatesToRetry, failedStateReads ++ exceptionsForPartitionsWithNoLeaderAndIsrInZk ++ failedUpdates)
+
+    val UpdateLeaderAndIsrResult(finishedPartitions, updatesToRetry) = zkClient.updateLeaderAndIsr(
+      adjustedLeaderAndIsrs,
+      controllerContext.epoch,
+      controllerContext.epochZkVersion
+    )
+
+    val exceptionsForPartitionsWithNoLeaderAndIsrInZk: Map[TopicPartition, Either[Exception, LeaderIsrAndControllerEpoch]] =
+      partitionsWithNoLeaderAndIsrInZk.iterator.flatMap { partition =>
+        if (!controllerContext.isTopicQueuedUpForDeletion(partition.topic)) {
+          val exception = new StateChangeFailedException(
+            s"Failed to change state of replica $replicaId for partition $partition since the leader and isr " +
+            "path in zookeeper is empty"
+          )
+          Option(partition -> Left(exception))
+        } else None
+      }.toMap
+
+    val leaderIsrAndControllerEpochs: Map[TopicPartition, Either[Exception, LeaderIsrAndControllerEpoch]] =
+      (leaderAndIsrsWithoutReplica ++ finishedPartitions).map { case (partition, result: Either[Exception, LeaderAndIsr]) =>
+        (
+          partition,
+          result.right.map { leaderAndIsr =>
+            val leaderIsrAndControllerEpoch = LeaderIsrAndControllerEpoch(leaderAndIsr, controllerContext.epoch)
+            controllerContext.partitionLeadershipInfo.put(partition, leaderIsrAndControllerEpoch)
+            leaderIsrAndControllerEpoch
+          }
+          )
+      }
+
+    (
+      leaderIsrAndControllerEpochs ++ exceptionsForPartitionsWithNoLeaderAndIsrInZk,
+      updatesToRetry
+    )
   }
 
   /**
    * Gets the partition state from zookeeper
    * @param partitions the partitions whose state we want from zookeeper
-   * @return A tuple of three values:
-   *         1. The LeaderAndIsrs of partitions whose state we successfully read from zookeeper
+   * @return A tuple of two values:
+   *         1. The Right(LeaderAndIsrs) of partitions whose state we successfully read from zookeeper.
+   *         The Left(Exception) to failed zookeeper lookups or states whose controller epoch exceeds our current epoch
    *         2. The partitions that had no leader and isr state in zookeeper. This happens if the controller
    *         didn't finish partition initialization.
-   *         3. Exceptions corresponding to failed zookeeper lookups or states whose controller epoch exceeds our current epoch.
    */
-  private def getTopicPartitionStatesFromZk(partitions: Seq[TopicPartition]):
-  (Map[TopicPartition, LeaderAndIsr],
-    Seq[TopicPartition],
-    Map[TopicPartition, Exception]) = {
-    val leaderAndIsrs = mutable.Map.empty[TopicPartition, LeaderAndIsr]
-    val partitionsWithNoLeaderAndIsrInZk = mutable.Buffer.empty[TopicPartition]
-    val failed = mutable.Map.empty[TopicPartition, Exception]
+  private def getTopicPartitionStatesFromZk(
+    partitions: Seq[TopicPartition]
+  ): (Map[TopicPartition, Either[Exception, LeaderAndIsr]], Seq[TopicPartition]) = {
     val getDataResponses = try {
       zkClient.getTopicPartitionStatesRaw(partitions)
     } catch {
       case e: Exception =>
-        partitions.foreach(partition => failed.put(partition, e))
-        return (leaderAndIsrs.toMap, partitionsWithNoLeaderAndIsrInZk, failed.toMap)
+        return (partitions.iterator.map(_ -> Left(e)).toMap, Seq.empty)
     }
-    getDataResponses.foreach { getDataResponse =>
+
+    val partitionsWithNoLeaderAndIsrInZk = mutable.Buffer.empty[TopicPartition]
+    val result = mutable.Map.empty[TopicPartition, Either[Exception, LeaderAndIsr]]
+
+    getDataResponses.foreach[Unit] { getDataResponse =>
       val partition = getDataResponse.ctx.get.asInstanceOf[TopicPartition]
       if (getDataResponse.resultCode == Code.OK) {
-        val leaderIsrAndControllerEpochOpt = TopicPartitionStateZNode.decode(getDataResponse.data, getDataResponse.stat)
-        if (leaderIsrAndControllerEpochOpt.isEmpty) {
-          partitionsWithNoLeaderAndIsrInZk += partition
-        } else {
-          val leaderIsrAndControllerEpoch = leaderIsrAndControllerEpochOpt.get
-          if (leaderIsrAndControllerEpoch.controllerEpoch > controllerContext.epoch) {
-            val exception = new StateChangeFailedException("Leader and isr path written by another controller. This probably" +
-              s"means the current controller with epoch ${controllerContext.epoch} went through a soft failure and another " +
-              s"controller was elected with epoch ${leaderIsrAndControllerEpoch.controllerEpoch}. Aborting state change by this controller")
-            failed.put(partition, exception)
-          } else {
-            leaderAndIsrs.put(partition, leaderIsrAndControllerEpoch.leaderAndIsr)
-          }
+        TopicPartitionStateZNode.decode(getDataResponse.data, getDataResponse.stat) match {
+          case None =>
+            partitionsWithNoLeaderAndIsrInZk += partition
+          case Some(leaderIsrAndControllerEpoch) =>
+            if (leaderIsrAndControllerEpoch.controllerEpoch > controllerContext.epoch) {
+              val exception = new StateChangeFailedException(
+                "Leader and isr path written by another controller. This probably " +
+                s"means the current controller with epoch ${controllerContext.epoch} went through a soft failure and " +
+                s"another controller was elected with epoch ${leaderIsrAndControllerEpoch.controllerEpoch}. Aborting " +
+                "state change by this controller"
+              )
+              result += (partition -> Left(exception))
+            } else {
+              result += (partition -> Right(leaderIsrAndControllerEpoch.leaderAndIsr))
+            }
         }
       } else if (getDataResponse.resultCode == Code.NONODE) {
         partitionsWithNoLeaderAndIsrInZk += partition
       } else {
-        failed.put(partition, getDataResponse.resultException.get)
+        result += (partition -> Left(getDataResponse.resultException.get))
       }
     }
-    (leaderAndIsrs.toMap, partitionsWithNoLeaderAndIsrInZk, failed.toMap)
-  }
-
-  def isAtLeastOneReplicaInDeletionStartedState(topic: String): Boolean = {
-    controllerContext.replicasForTopic(topic).exists(replica => replicaState(replica) == ReplicaDeletionStarted)
-  }
-
-  def replicasInState(topic: String, state: ReplicaState): Set[PartitionAndReplica] = {
-    replicaState.filter { case (replica, s) => replica.topic.equals(topic) && s == state }.keySet.toSet
-  }
 
-  def areAllReplicasForTopicDeleted(topic: String): Boolean = {
-    controllerContext.replicasForTopic(topic).forall(replica => replicaState(replica) == ReplicaDeletionSuccessful)
+    (result.toMap, partitionsWithNoLeaderAndIsrInZk)
   }
 
-  def isAnyReplicaInState(topic: String, state: ReplicaState): Boolean = {
-    replicaState.exists { case (replica, s) => replica.topic.equals(topic) && s == state}
-  }
-
-  private def isValidTransition(replica: PartitionAndReplica, targetState: ReplicaState) =
-    targetState.validPreviousStates.contains(replicaState(replica))
-
   private def logSuccessfulTransition(replicaId: Int, partition: TopicPartition, currState: ReplicaState, targetState: ReplicaState): Unit = {
     stateChangeLogger.withControllerEpoch(controllerContext.epoch)
       .trace(s"Changed state of replica $replicaId for partition $partition from $currState to $targetState")
   }
 
   private def logInvalidTransition(replica: PartitionAndReplica, targetState: ReplicaState): Unit = {
-    val currState = replicaState(replica)
+    val currState = controllerContext.replicaState(replica)
     val e = new IllegalStateException(s"Replica $replica should be in the ${targetState.validPreviousStates.mkString(",")} " +
       s"states before moving to $targetState state. Instead it is in $currState state")
     logFailedStateChange(replica, currState, targetState, e)
@@ -437,7 +466,7 @@ case object ReplicaDeletionSuccessful extends ReplicaState {
 
 case object ReplicaDeletionIneligible extends ReplicaState {
   val state: Byte = 6
-  val validPreviousStates: Set[ReplicaState] = Set(ReplicaDeletionStarted)
+  val validPreviousStates: Set[ReplicaState] = Set(OfflineReplica, ReplicaDeletionStarted)
 }
 
 case object NonExistentReplica extends ReplicaState {
diff --git a/core/src/main/scala/kafka/controller/TopicDeletionManager.scala b/core/src/main/scala/kafka/controller/TopicDeletionManager.scala
index 1ef79be794f89..d33af05023c50 100755
--- a/core/src/main/scala/kafka/controller/TopicDeletionManager.scala
+++ b/core/src/main/scala/kafka/controller/TopicDeletionManager.scala
@@ -16,11 +16,55 @@
  */
 package kafka.controller
 
+import kafka.server.KafkaConfig
 import kafka.utils.Logging
 import kafka.zk.KafkaZkClient
 import org.apache.kafka.common.TopicPartition
+import org.apache.zookeeper.KeeperException.{NoNodeException, NodeExistsException}
 
-import scala.collection.{Set, mutable}
+import scala.collection.Set
+
+trait DeletionClient {
+  def deleteTopic(topic: String, epochZkVersion: Int): Unit
+  def deleteTopicDeletions(topics: Seq[String], epochZkVersion: Int): Unit
+  def mutePartitionModifications(topic: String): Unit
+  def sendMetadataUpdate(partitions: Set[TopicPartition]): Unit
+  def createDeleteTopicFlagPath(): Unit
+  def getTopicDeletionFlag(): String
+  def removePartitionsFromReassignedPartitions(partitionsToBeRemoved: Set[TopicPartition]): Unit
+}
+
+class ControllerDeletionClient(controller: KafkaController, zkClient: KafkaZkClient) extends DeletionClient {
+  override def deleteTopic(topic: String, epochZkVersion: Int): Unit = {
+    zkClient.deleteTopicZNode(topic, epochZkVersion)
+    zkClient.deleteTopicConfigs(Seq(topic), epochZkVersion)
+    zkClient.deleteTopicDeletions(Seq(topic), epochZkVersion)
+  }
+
+  override def deleteTopicDeletions(topics: Seq[String], epochZkVersion: Int): Unit = {
+    zkClient.deleteTopicDeletions(topics, epochZkVersion)
+  }
+
+  override def mutePartitionModifications(topic: String): Unit = {
+    controller.unregisterPartitionModificationsHandlers(Seq(topic))
+  }
+
+  override def sendMetadataUpdate(partitions: Set[TopicPartition]): Unit = {
+    controller.sendUpdateMetadataRequest(controller.controllerContext.liveOrShuttingDownBrokerIds.toSeq, partitions)
+  }
+
+  override def createDeleteTopicFlagPath(): Unit = {
+    zkClient.createDeleteTopicFlagPath
+  }
+
+  override def getTopicDeletionFlag(): String = {
+    zkClient.getTopicDeletionFlag
+  }
+
+  override def removePartitionsFromReassignedPartitions(partitionsToBeRemoved: Set[TopicPartition]): Unit = {
+    controller.maybeRemoveFromZkReassignment((tp, _) => partitionsToBeRemoved.contains(tp))
+  }
+}
 
 /**
  * This manages the state machine for topic deletion.
@@ -55,44 +99,34 @@ import scala.collection.{Set, mutable}
  *    it marks the topic for deletion retry.
  * @param controller
  */
-class TopicDeletionManager(controller: KafkaController,
-                           eventManager: ControllerEventManager,
-                           zkClient: KafkaZkClient) extends Logging {
-  this.logIdent = s"[Topic Deletion Manager ${controller.config.brokerId}], "
-  val controllerContext = controller.controllerContext
-  val isDeleteTopicEnabled = controller.config.deleteTopicEnable
-  val topicsToBeDeleted = mutable.Set.empty[String]
-  /** The following topicsWithDeletionStarted variable is used to properly update the offlinePartitionCount metric.
-    * When a topic is going through deletion, we don't want to keep track of its partition state
-    * changes in the offlinePartitionCount metric, see the PartitionStateMachine#updateControllerMetrics
-    * for detailed logic. This goal means if some partitions of a topic are already
-    * in OfflinePartition state when deletion starts, we need to change the corresponding partition
-    * states to NonExistentPartition first before starting the deletion.
-    *
-    * However we can NOT change partition states to NonExistentPartition at the time of enqueuing topics
-    * for deletion. The reason is that when a topic is enqueued for deletion, it may be ineligible for
-    * deletion due to ongoing partition reassignments. Hence there might be a delay between enqueuing
-    * a topic for deletion and the actual start of deletion. In this delayed interval, partitions may still
-    * transition to or out of the OfflinePartition state.
-    *
-    * Hence we decide to change partition states to NonExistentPartition only when the actual deletion have started.
-    * For topics whose deletion have actually started, we keep track of them in the following topicsWithDeletionStarted
-    * variable. And once a topic is in the topicsWithDeletionStarted set, we are sure there will no longer
-    * be partition reassignments to any of its partitions, and only then it's safe to move its partitions to
-    * NonExistentPartition state. Once a topic is in the topicsWithDeletionStarted set, we will stop monitoring
-    * its partition state changes in the offlinePartitionCount metric
-    */
-  val topicsWithDeletionStarted = mutable.Set.empty[String]
-  val topicsIneligibleForDeletion = mutable.Set.empty[String]
+class TopicDeletionManager(config: KafkaConfig,
+                           controllerContext: ControllerContext,
+                           replicaStateMachine: ReplicaStateMachine,
+                           partitionStateMachine: PartitionStateMachine,
+                           client: DeletionClient) extends Logging {
+  this.logIdent = s"[Topic Deletion Manager ${config.brokerId}] "
+  var isDeleteTopicEnabled: Boolean = config.deleteTopicEnable
+
+  // Try to create the znode for delete topic flag
+  try {
+    client.createDeleteTopicFlagPath()
+  } catch {
+    case _: NodeExistsException =>
+  }
 
   def init(initialTopicsToBeDeleted: Set[String], initialTopicsIneligibleForDeletion: Set[String]): Unit = {
+    info(s"Initializing manager with initial deletions: $initialTopicsToBeDeleted, " +
+      s"initial ineligible deletions: $initialTopicsIneligibleForDeletion")
+
+    isDeleteTopicEnabled = getDeleteTopicEnabled()
+
     if (isDeleteTopicEnabled) {
-      topicsToBeDeleted ++= initialTopicsToBeDeleted
-      topicsIneligibleForDeletion ++= initialTopicsIneligibleForDeletion & topicsToBeDeleted
+      controllerContext.queueTopicDeletion(initialTopicsToBeDeleted)
+      controllerContext.topicsIneligibleForDeletion ++= initialTopicsIneligibleForDeletion & controllerContext.topicsToBeDeleted
     } else {
       // if delete topic is disabled clean the topic entries under /admin/delete_topics
       info(s"Removing $initialTopicsToBeDeleted since delete topic is disabled")
-      zkClient.deleteTopicDeletions(initialTopicsToBeDeleted.toSeq, controllerContext.epochZkVersion)
+      client.deleteTopicDeletions(initialTopicsToBeDeleted.toSeq, controllerContext.epochZkVersion)
     }
   }
 
@@ -102,26 +136,17 @@ class TopicDeletionManager(controller: KafkaController,
     }
   }
 
-  /**
-   * Invoked when the current controller resigns. At this time, all state for topic deletion should be cleared.
-   */
-  def reset() {
-    if (isDeleteTopicEnabled) {
-      topicsToBeDeleted.clear()
-      topicsWithDeletionStarted.clear()
-      topicsIneligibleForDeletion.clear()
-    }
-  }
-
   /**
    * Invoked by the child change listener on /admin/delete_topics to queue up the topics for deletion. The topic gets added
    * to the topicsToBeDeleted list and only gets removed from the list when the topic deletion has completed successfully
    * i.e. all replicas of all partitions of that topic are deleted successfully.
    * @param topics Topics that should be deleted
    */
-  def enqueueTopicsForDeletion(topics: Set[String]) {
+  def enqueueTopicsForDeletion(topics: Set[String]): Unit = {
     if (isDeleteTopicEnabled) {
-      topicsToBeDeleted ++= topics
+      val newTopicsToBeDeleted = topics -- controllerContext.topicsToBeDeleted
+      controllerContext.queueTopicDeletion(newTopicsToBeDeleted)
+      newTopicsToBeDeleted.foreach(controllerContext.excludeDeletingTopicFromOfflinePartitionCount)
       resumeDeletions()
     }
   }
@@ -132,11 +157,11 @@ class TopicDeletionManager(controller: KafkaController,
    * 2. Partition reassignment completes. Any partitions belonging to topics queued up for deletion finished reassignment
    * @param topics Topics for which deletion can be resumed
    */
-  def resumeDeletionForTopics(topics: Set[String] = Set.empty) {
+  def resumeDeletionForTopics(topics: Set[String] = Set.empty): Unit = {
     if (isDeleteTopicEnabled) {
-      val topicsToResumeDeletion = topics & topicsToBeDeleted
+      val topicsToResumeDeletion = topics & controllerContext.topicsToBeDeleted
       if (topicsToResumeDeletion.nonEmpty) {
-        topicsIneligibleForDeletion --= topicsToResumeDeletion
+        controllerContext.topicsIneligibleForDeletion --= topicsToResumeDeletion
         resumeDeletions()
       }
     }
@@ -149,14 +174,14 @@ class TopicDeletionManager(controller: KafkaController,
    * ineligible for deletion until further notice.
    * @param replicas Replicas for which deletion has failed
    */
-  def failReplicaDeletion(replicas: Set[PartitionAndReplica]) {
+  def failReplicaDeletion(replicas: Set[PartitionAndReplica]): Unit = {
     if (isDeleteTopicEnabled) {
       val replicasThatFailedToDelete = replicas.filter(r => isTopicQueuedUpForDeletion(r.topic))
       if (replicasThatFailedToDelete.nonEmpty) {
         val topics = replicasThatFailedToDelete.map(_.topic)
         debug(s"Deletion failed for replicas ${replicasThatFailedToDelete.mkString(",")}. Halting deletion for topics $topics")
-        controller.replicaStateMachine.handleStateChanges(replicasThatFailedToDelete.toSeq, ReplicaDeletionIneligible)
-        markTopicIneligibleForDeletion(topics)
+        replicaStateMachine.handleStateChanges(replicasThatFailedToDelete.toSeq, ReplicaDeletionIneligible)
+        markTopicIneligibleForDeletion(topics, reason = "replica deletion failure")
         resumeDeletions()
       }
     }
@@ -168,39 +193,32 @@ class TopicDeletionManager(controller: KafkaController,
    * 2. partition reassignment in progress for some partitions of the topic
    * @param topics Topics that should be marked ineligible for deletion. No op if the topic is was not previously queued up for deletion
    */
-  def markTopicIneligibleForDeletion(topics: Set[String]) {
+  def markTopicIneligibleForDeletion(topics: Set[String], reason: => String): Unit = {
     if (isDeleteTopicEnabled) {
-      val newTopicsToHaltDeletion = topicsToBeDeleted & topics
-      topicsIneligibleForDeletion ++= newTopicsToHaltDeletion
+      val newTopicsToHaltDeletion = controllerContext.topicsToBeDeleted & topics
+      controllerContext.topicsIneligibleForDeletion ++= newTopicsToHaltDeletion
       if (newTopicsToHaltDeletion.nonEmpty)
-        info(s"Halted deletion of topics ${newTopicsToHaltDeletion.mkString(",")}")
+        info(s"Halted deletion of topics ${newTopicsToHaltDeletion.mkString(",")} due to $reason")
     }
   }
 
   private def isTopicIneligibleForDeletion(topic: String): Boolean = {
     if (isDeleteTopicEnabled) {
-      topicsIneligibleForDeletion.contains(topic)
+      controllerContext.topicsIneligibleForDeletion.contains(topic)
     } else
       true
   }
 
   private def isTopicDeletionInProgress(topic: String): Boolean = {
     if (isDeleteTopicEnabled) {
-      controller.replicaStateMachine.isAtLeastOneReplicaInDeletionStartedState(topic)
-    } else
-      false
-  }
-
-  def isTopicWithDeletionStarted(topic: String) = {
-    if (isDeleteTopicEnabled) {
-      topicsWithDeletionStarted.contains(topic)
+      controllerContext.isAnyReplicaInState(topic, ReplicaDeletionStarted)
     } else
       false
   }
 
   def isTopicQueuedUpForDeletion(topic: String): Boolean = {
     if (isDeleteTopicEnabled) {
-      topicsToBeDeleted.contains(topic)
+      controllerContext.isTopicQueuedUpForDeletion(topic)
     } else
       false
   }
@@ -211,10 +229,10 @@ class TopicDeletionManager(controller: KafkaController,
    * the topic if all replicas of a topic have been successfully deleted
    * @param replicas Replicas that were successfully deleted by the broker
    */
-  def completeReplicaDeletion(replicas: Set[PartitionAndReplica]) {
+  def completeReplicaDeletion(replicas: Set[PartitionAndReplica]): Unit = {
     val successfullyDeletedReplicas = replicas.filter(r => isTopicQueuedUpForDeletion(r.topic))
     debug(s"Deletion successfully completed for replicas ${successfullyDeletedReplicas.mkString(",")}")
-    controller.replicaStateMachine.handleStateChanges(successfullyDeletedReplicas.toSeq, ReplicaDeletionSuccessful)
+    replicaStateMachine.handleStateChanges(successfullyDeletedReplicas.toSeq, ReplicaDeletionSuccessful)
     resumeDeletions()
   }
 
@@ -227,7 +245,9 @@ class TopicDeletionManager(controller: KafkaController,
    * @return Whether or not deletion can be retried for the topic
    */
   private def isTopicEligibleForDeletion(topic: String): Boolean = {
-    topicsToBeDeleted.contains(topic) && (!isTopicDeletionInProgress(topic) && !isTopicIneligibleForDeletion(topic))
+    controllerContext.isTopicQueuedUpForDeletion(topic) &&
+      !isTopicDeletionInProgress(topic) &&
+      !isTopicIneligibleForDeletion(topic)
   }
 
   /**
@@ -235,28 +255,36 @@ class TopicDeletionManager(controller: KafkaController,
    * To ensure a successful retry, reset states for respective replicas from ReplicaDeletionIneligible to OfflineReplica state
    *@param topic Topic for which deletion should be retried
    */
-  private def markTopicForDeletionRetry(topic: String) {
+  private def retryDeletionForIneligibleReplicas(topic: String): Unit = {
     // reset replica states from ReplicaDeletionIneligible to OfflineReplica
-    val failedReplicas = controller.replicaStateMachine.replicasInState(topic, ReplicaDeletionIneligible)
-    info(s"Retrying delete topic for topic $topic since replicas ${failedReplicas.mkString(",")} were not successfully deleted")
-    controller.replicaStateMachine.handleStateChanges(failedReplicas.toSeq, OfflineReplica)
+    val failedReplicas = controllerContext.replicasInState(topic, ReplicaDeletionIneligible)
+    info(s"Retrying deletion of topic $topic since replicas ${failedReplicas.mkString(",")} were not successfully deleted")
+    replicaStateMachine.handleStateChanges(failedReplicas.toSeq, OfflineReplica)
   }
 
-  private def completeDeleteTopic(topic: String) {
+  private def completeDeleteTopic(topic: String): Unit = {
+    // abort pending partition reassignments for deleted topic
+    abortPartitionReassignmentForTopic(topic)
     // deregister partition change listener on the deleted topic. This is to prevent the partition change listener
     // firing before the new topic listener when a deleted topic gets auto created
-    controller.unregisterPartitionModificationsHandlers(Seq(topic))
-    val replicasForDeletedTopic = controller.replicaStateMachine.replicasInState(topic, ReplicaDeletionSuccessful)
+    client.mutePartitionModifications(topic)
+    val replicasForDeletedTopic = controllerContext.replicasInState(topic, ReplicaDeletionSuccessful)
     // controller will remove this replica from the state machine as well as its partition assignment cache
-    controller.replicaStateMachine.handleStateChanges(replicasForDeletedTopic.toSeq, NonExistentReplica)
-    topicsToBeDeleted -= topic
-    topicsWithDeletionStarted -= topic
-    zkClient.deleteTopicZNode(topic, controllerContext.epochZkVersion)
-    zkClient.deleteTopicConfigs(Seq(topic), controllerContext.epochZkVersion)
-    zkClient.deleteTopicDeletions(Seq(topic), controllerContext.epochZkVersion)
+    replicaStateMachine.handleStateChanges(replicasForDeletedTopic.toSeq, NonExistentReplica)
+    controllerContext.topicsToBeDeleted -= topic
+    controllerContext.topicsWithDeletionStarted -= topic
+    client.deleteTopic(topic, controllerContext.epochZkVersion)
     controllerContext.removeTopic(topic)
   }
 
+  private def abortPartitionReassignmentForTopic(topic: String): Unit = {
+    val partitionsBeingReassignedForDeletedTopic =
+      controllerContext.partitionsBeingReassigned.filter(_.topic().equals(topic))
+    if (partitionsBeingReassignedForDeletedTopic.nonEmpty) {
+      client.removePartitionsFromReassignedPartitions(partitionsBeingReassignedForDeletedTopic)
+    }
+  }
+
   /**
    * Invoked with the list of topics to be deleted
    * It invokes onPartitionDeletion for all partitions of a topic.
@@ -264,21 +292,21 @@ class TopicDeletionManager(controller: KafkaController,
    * {@link LeaderAndIsr#LeaderDuringDelete}. This lets each broker know that this topic is being deleted and can be
    * removed from their caches.
    */
-  private def onTopicDeletion(topics: Set[String]) {
+  private def onTopicDeletion(topics: Set[String]): Unit = {
     info(s"Topic deletion callback for ${topics.mkString(",")}")
     // send update metadata so that brokers stop serving data for topics to be deleted
     val partitions = topics.flatMap(controllerContext.partitionsForTopic)
-    val unseenTopicsForDeletion = topics -- topicsWithDeletionStarted
+    val unseenTopicsForDeletion = topics -- controllerContext.topicsWithDeletionStarted
     if (unseenTopicsForDeletion.nonEmpty) {
       val unseenPartitionsForDeletion = unseenTopicsForDeletion.flatMap(controllerContext.partitionsForTopic)
-      controller.partitionStateMachine.handleStateChanges(unseenPartitionsForDeletion.toSeq, OfflinePartition)
-      controller.partitionStateMachine.handleStateChanges(unseenPartitionsForDeletion.toSeq, NonExistentPartition)
-      // adding of unseenTopicsForDeletion to topicsBeingDeleted must be done after the partition state changes
-      // to make sure the offlinePartitionCount metric is properly updated
-      topicsWithDeletionStarted ++= unseenTopicsForDeletion
+      partitionStateMachine.handleStateChanges(unseenPartitionsForDeletion.toSeq, OfflinePartition)
+      partitionStateMachine.handleStateChanges(unseenPartitionsForDeletion.toSeq, NonExistentPartition)
+      // adding of unseenTopicsForDeletion to topics with deletion started must be done after the partition
+      // state changes to make sure the offlinePartitionCount metric is properly updated
+      controllerContext.beginTopicDeletion(unseenTopicsForDeletion)
     }
 
-    controller.sendUpdateMetadataRequest(controllerContext.liveOrShuttingDownBrokerIds.toSeq, partitions)
+    client.sendMetadataUpdate(partitions)
     topics.foreach { topic =>
       onPartitionDeletion(controllerContext.partitionsForTopic(topic))
     }
@@ -298,25 +326,23 @@ class TopicDeletionManager(controller: KafkaController,
    * 1. Move all dead replicas directly to ReplicaDeletionIneligible state. Also mark the respective topics ineligible
    *    for deletion if some replicas are dead since it won't complete successfully anyway
    * 2. Move all alive replicas to ReplicaDeletionStarted state so they can be deleted successfully
-   *@param replicasForTopicsToBeDeleted
+   * @param replicasForTopicsToBeDeleted
    */
-  private def startReplicaDeletion(replicasForTopicsToBeDeleted: Set[PartitionAndReplica]) {
+  private def startReplicaDeletion(replicasForTopicsToBeDeleted: Set[PartitionAndReplica]): Unit = {
     replicasForTopicsToBeDeleted.groupBy(_.topic).keys.foreach { topic =>
       val aliveReplicasForTopic = controllerContext.allLiveReplicas().filter(p => p.topic == topic)
       val deadReplicasForTopic = replicasForTopicsToBeDeleted -- aliveReplicasForTopic
-      val successfullyDeletedReplicas = controller.replicaStateMachine.replicasInState(topic, ReplicaDeletionSuccessful)
+      val successfullyDeletedReplicas = controllerContext.replicasInState(topic, ReplicaDeletionSuccessful)
       val replicasForDeletionRetry = aliveReplicasForTopic -- successfullyDeletedReplicas
       // move dead replicas directly to failed state
-      controller.replicaStateMachine.handleStateChanges(deadReplicasForTopic.toSeq, ReplicaDeletionIneligible, new Callbacks())
+      replicaStateMachine.handleStateChanges(deadReplicasForTopic.toSeq, ReplicaDeletionIneligible)
       // send stop replica to all followers that are not in the OfflineReplica state so they stop sending fetch requests to the leader
-      controller.replicaStateMachine.handleStateChanges(replicasForDeletionRetry.toSeq, OfflineReplica, new Callbacks())
+      replicaStateMachine.handleStateChanges(replicasForDeletionRetry.toSeq, OfflineReplica)
       debug(s"Deletion started for replicas ${replicasForDeletionRetry.mkString(",")}")
-      controller.replicaStateMachine.handleStateChanges(replicasForDeletionRetry.toSeq, ReplicaDeletionStarted,
-        new Callbacks(stopReplicaResponseCallback = (stopReplicaResponseObj, replicaId) =>
-          eventManager.put(controller.TopicDeletionStopReplicaResponseReceived(stopReplicaResponseObj, replicaId))))
+      replicaStateMachine.handleStateChanges(replicasForDeletionRetry.toSeq, ReplicaDeletionStarted)
       if (deadReplicasForTopic.nonEmpty) {
         debug(s"Dead Replicas (${deadReplicasForTopic.mkString(",")}) found for topic $topic")
-        markTopicIneligibleForDeletion(Set(topic))
+        markTopicIneligibleForDeletion(Set(topic), reason = "offline replicas")
       }
     }
   }
@@ -332,49 +358,54 @@ class TopicDeletionManager(controller: KafkaController,
    * 3. Move all replicas to ReplicaDeletionStarted state. This will send StopReplicaRequest with deletePartition=true. And
    *    will delete all persistent data from all replicas of the respective partitions
    */
-  private def onPartitionDeletion(partitionsToBeDeleted: Set[TopicPartition]) {
+  private def onPartitionDeletion(partitionsToBeDeleted: Set[TopicPartition]): Unit = {
     info(s"Partition deletion callback for ${partitionsToBeDeleted.mkString(",")}")
     val replicasPerPartition = controllerContext.replicasForPartition(partitionsToBeDeleted)
     startReplicaDeletion(replicasPerPartition)
   }
 
   private def resumeDeletions(): Unit = {
-    val topicsQueuedForDeletion = Set.empty[String] ++ topicsToBeDeleted
-
+    val topicsQueuedForDeletion = Set.empty[String] ++ controllerContext.topicsToBeDeleted
     if (topicsQueuedForDeletion.nonEmpty)
       info(s"Handling deletion for topics ${topicsQueuedForDeletion.mkString(",")}")
 
     topicsQueuedForDeletion.foreach { topic =>
       // if all replicas are marked as deleted successfully, then topic deletion is done
-      if (controller.replicaStateMachine.areAllReplicasForTopicDeleted(topic)) {
+      if (controllerContext.areAllReplicasInState(topic, ReplicaDeletionSuccessful)) {
         // clear up all state for this topic from controller cache and zookeeper
         completeDeleteTopic(topic)
         info(s"Deletion of topic $topic successfully completed")
-      } else {
-        if (controller.replicaStateMachine.isAtLeastOneReplicaInDeletionStartedState(topic)) {
-          // ignore since topic deletion is in progress
-          val replicasInDeletionStartedState = controller.replicaStateMachine.replicasInState(topic, ReplicaDeletionStarted)
-          val replicaIds = replicasInDeletionStartedState.map(_.replica)
-          val partitions = replicasInDeletionStartedState.map(_.topicPartition)
-          info(s"Deletion for replicas ${replicaIds.mkString(",")} for partition ${partitions.mkString(",")} of topic $topic in progress")
-        } else {
-          // if you come here, then no replica is in TopicDeletionStarted and all replicas are not in
-          // TopicDeletionSuccessful. That means, that either given topic haven't initiated deletion
-          // or there is at least one failed replica (which means topic deletion should be retried).
-          if (controller.replicaStateMachine.isAnyReplicaInState(topic, ReplicaDeletionIneligible)) {
-            // mark topic for deletion retry
-            markTopicForDeletionRetry(topic)
-          }
+      } else if (!controllerContext.isAnyReplicaInState(topic, ReplicaDeletionStarted)) {
+        // if you come here, then no replica is in TopicDeletionStarted and all replicas are not in
+        // TopicDeletionSuccessful. That means, that either given topic haven't initiated deletion
+        // or there is at least one failed replica (which means topic deletion should be retried).
+        if (controllerContext.isAnyReplicaInState(topic, ReplicaDeletionIneligible)) {
+          retryDeletionForIneligibleReplicas(topic)
         }
       }
+
       // Try delete topic if it is eligible for deletion.
       if (isTopicEligibleForDeletion(topic)) {
         info(s"Deletion of topic $topic (re)started")
         // topic deletion will be kicked off
         onTopicDeletion(Set(topic))
-      } else if (isTopicIneligibleForDeletion(topic)) {
-        info(s"Not retrying deletion of topic $topic at this time since it is marked ineligible for deletion")
       }
     }
   }
+
+  private def getDeleteTopicEnabled(): Boolean = {
+    try {
+      val deleteTopicFlag = client.getTopicDeletionFlag
+      if (deleteTopicFlag == null || (!deleteTopicFlag.equalsIgnoreCase("true") && !deleteTopicFlag.equalsIgnoreCase("false")))
+        isDeleteTopicEnabled
+      else deleteTopicFlag.toBoolean
+    } catch {
+      case _: NoNodeException => config.deleteTopicEnable
+    }
+  }
+
+  def resetDeleteTopicEnabled(): Unit = {
+    info("Reset isDeleteTopicEnabled flag to %s".format(config.deleteTopicEnable))
+    isDeleteTopicEnabled = config.deleteTopicEnable
+  }
 }
diff --git a/core/src/main/scala/kafka/coordinator/group/DelayedHeartbeat.scala b/core/src/main/scala/kafka/coordinator/group/DelayedHeartbeat.scala
index 09c5eea58ea3e..3f402d9b5f747 100644
--- a/core/src/main/scala/kafka/coordinator/group/DelayedHeartbeat.scala
+++ b/core/src/main/scala/kafka/coordinator/group/DelayedHeartbeat.scala
@@ -27,11 +27,10 @@ private[group] class DelayedHeartbeat(coordinator: GroupCoordinator,
                                       group: GroupMetadata,
                                       memberId: String,
                                       isPending: Boolean,
-                                      deadline: Long,
                                       timeoutMs: Long)
   extends DelayedOperation(timeoutMs, Some(group.lock)) {
 
-  override def tryComplete(): Boolean = coordinator.tryCompleteHeartbeat(group, memberId, isPending, deadline, forceComplete _)
-  override def onExpiration() = coordinator.onExpireHeartbeat(group, memberId, isPending, deadline)
+  override def tryComplete(): Boolean = coordinator.tryCompleteHeartbeat(group, memberId, isPending, forceComplete _)
+  override def onExpiration() = coordinator.onExpireHeartbeat(group, memberId, isPending)
   override def onComplete() = coordinator.onCompleteHeartbeat()
 }
diff --git a/core/src/main/scala/kafka/coordinator/group/GroupCoordinator.scala b/core/src/main/scala/kafka/coordinator/group/GroupCoordinator.scala
index ead05ba49b9e4..c877fb5982172 100644
--- a/core/src/main/scala/kafka/coordinator/group/GroupCoordinator.scala
+++ b/core/src/main/scala/kafka/coordinator/group/GroupCoordinator.scala
@@ -27,6 +27,10 @@ import kafka.utils.Logging
 import kafka.zk.KafkaZkClient
 import org.apache.kafka.common.TopicPartition
 import org.apache.kafka.common.internals.Topic
+import org.apache.kafka.common.message.JoinGroupResponseData.JoinGroupResponseMember
+import org.apache.kafka.common.metrics.Metrics
+import org.apache.kafka.common.message.LeaveGroupRequestData.MemberIdentity
+import org.apache.kafka.common.metrics.stats.Meter
 import org.apache.kafka.common.protocol.{ApiKeys, Errors}
 import org.apache.kafka.common.record.RecordBatch.{NO_PRODUCER_EPOCH, NO_PRODUCER_ID}
 import org.apache.kafka.common.requests._
@@ -53,11 +57,33 @@ class GroupCoordinator(val brokerId: Int,
                        val groupManager: GroupMetadataManager,
                        val heartbeatPurgatory: DelayedOperationPurgatory[DelayedHeartbeat],
                        val joinPurgatory: DelayedOperationPurgatory[DelayedJoin],
-                       time: Time) extends Logging {
+                       time: Time,
+                       metrics: Metrics) extends Logging {
   import GroupCoordinator._
 
   type JoinCallback = JoinGroupResult => Unit
-  type SyncCallback = (Array[Byte], Errors) => Unit
+  type SyncCallback = SyncGroupResult => Unit
+
+  /* setup metrics */
+  val offsetDeletionSensor = metrics.sensor("OffsetDeletions")
+
+  offsetDeletionSensor.add(new Meter(
+    metrics.metricName("offset-deletion-rate",
+      "group-coordinator-metrics",
+      "The rate of administrative deleted offsets"),
+    metrics.metricName("offset-deletion-count",
+      "group-coordinator-metrics",
+      "The total number of administrative deleted offsets")))
+
+  val groupCompletedRebalanceSensor = metrics.sensor("CompletedRebalances")
+
+  groupCompletedRebalanceSensor.add(new Meter(
+    metrics.metricName("group-completed-rebalance-rate",
+      "group-coordinator-metrics",
+      "The rate of completed rebalance"),
+    metrics.metricName("group-completed-rebalance-count",
+      "group-coordinator-metrics",
+      "The total number of completed rebalance")))
 
   this.logIdent = "[GroupCoordinator " + brokerId + "]: "
 
@@ -68,6 +94,9 @@ class GroupCoordinator(val brokerId: Int,
     props.put(LogConfig.CleanupPolicyProp, LogConfig.Compact)
     props.put(LogConfig.SegmentBytesProp, offsetConfig.offsetsTopicSegmentBytes.toString)
     props.put(LogConfig.CompressionTypeProp, ProducerCompressionCodec.name)
+    props.put(LogConfig.MaxMessageBytesProp, offsetConfig.offsetsTopicMaxMessageBytes.toString)
+    props.put(LogConfig.MinInSyncReplicasProp, offsetConfig.offsetsTopicMinInSyncReplicas.toString)
+    props.put(LogConfig.MinCompactionLagMsProp, offsetConfig.offsetsTopicMinCompactionLagMs.toString)
     props
   }
 
@@ -79,7 +108,7 @@ class GroupCoordinator(val brokerId: Int,
   /**
    * Startup logic executed at the same time when the server starts up.
    */
-  def startup(enableMetadataExpiration: Boolean = true) {
+  def startup(enableMetadataExpiration: Boolean = true): Unit = {
     info("Starting up.")
     groupManager.startup(enableMetadataExpiration)
     isActive.set(true)
@@ -90,7 +119,7 @@ class GroupCoordinator(val brokerId: Int,
    * Shutdown logic executed at the same time when server shuts down.
    * Ordering of actions should be reversed from the startup process.
    */
-  def shutdown() {
+  def shutdown(): Unit = {
     info("Shutting down.")
     isActive.set(false)
     groupManager.shutdown()
@@ -101,6 +130,7 @@ class GroupCoordinator(val brokerId: Int,
 
   def handleJoinGroup(groupId: String,
                       memberId: String,
+                      groupInstanceId: Option[String],
                       requireKnownMemberId: Boolean,
                       clientId: String,
                       clientHost: String,
@@ -126,22 +156,22 @@ class GroupCoordinator(val brokerId: Int,
           // exist we should reject the request.
           if (isUnknownMember) {
             val group = groupManager.addGroup(new GroupMetadata(groupId, Empty, time))
-            doUnknownJoinGroup(group, requireKnownMemberId, clientId, clientHost, rebalanceTimeoutMs, sessionTimeoutMs, protocolType, protocols, responseCallback)
+            doUnknownJoinGroup(group, groupInstanceId, requireKnownMemberId, clientId, clientHost, rebalanceTimeoutMs, sessionTimeoutMs, protocolType, protocols, responseCallback)
           } else {
             responseCallback(joinError(memberId, Errors.UNKNOWN_MEMBER_ID))
           }
-
         case Some(group) =>
           group.inLock {
             if ((groupIsOverCapacity(group)
                   && group.has(memberId) && !group.get(memberId).isAwaitingJoin) // oversized group, need to shed members that haven't joined yet
                 || (isUnknownMember && group.size >= groupConfig.groupMaxSize)) {
               group.remove(memberId)
+              group.removeStaticMember(groupInstanceId)
               responseCallback(joinError(JoinGroupRequest.UNKNOWN_MEMBER_ID, Errors.GROUP_MAX_SIZE_REACHED))
             } else if (isUnknownMember) {
-              doUnknownJoinGroup(group, requireKnownMemberId, clientId, clientHost, rebalanceTimeoutMs, sessionTimeoutMs, protocolType, protocols, responseCallback)
+              doUnknownJoinGroup(group, groupInstanceId, requireKnownMemberId, clientId, clientHost, rebalanceTimeoutMs, sessionTimeoutMs, protocolType, protocols, responseCallback)
             } else {
-              doJoinGroup(group, memberId, clientId, clientHost, rebalanceTimeoutMs, sessionTimeoutMs, protocolType, protocols, responseCallback)
+              doJoinGroup(group, memberId, groupInstanceId, clientId, clientHost, rebalanceTimeoutMs, sessionTimeoutMs, protocolType, protocols, responseCallback)
             }
 
             // attempt to complete JoinGroup
@@ -149,11 +179,12 @@ class GroupCoordinator(val brokerId: Int,
               joinPurgatory.checkAndComplete(GroupKey(group.groupId))
             }
           }
+        }
       }
     }
-  }
 
   private def doUnknownJoinGroup(group: GroupMetadata,
+                                 groupInstanceId: Option[String],
                                  requireKnownMemberId: Boolean,
                                  clientId: String,
                                  clientHost: String,
@@ -167,22 +198,61 @@ class GroupCoordinator(val brokerId: Int,
         // if the group is marked as dead, it means some other thread has just removed the group
         // from the coordinator metadata; it is likely that the group has migrated to some other
         // coordinator OR the group is in a transient unstable phase. Let the member retry
-        // joining without the specified member id.
-        responseCallback(joinError(JoinGroupRequest.UNKNOWN_MEMBER_ID, Errors.UNKNOWN_MEMBER_ID))
+        // finding the correct coordinator and rejoin.
+        responseCallback(joinError(JoinGroupRequest.UNKNOWN_MEMBER_ID, Errors.COORDINATOR_NOT_AVAILABLE))
       } else if (!group.supportsProtocols(protocolType, MemberMetadata.plainProtocolSet(protocols))) {
         responseCallback(joinError(JoinGroupRequest.UNKNOWN_MEMBER_ID, Errors.INCONSISTENT_GROUP_PROTOCOL))
       } else {
-        val newMemberId = clientId + "-" + group.generateMemberIdSuffix
-
-        if (requireKnownMemberId) {
-          // If member id required, register the member in the pending member list
-          // and send back a response to call for another join group request with allocated member id.
+        val newMemberId = group.generateMemberId(clientId, groupInstanceId)
+
+        if (group.hasStaticMember(groupInstanceId)) {
+          val oldMemberId = group.getStaticMemberId(groupInstanceId)
+          info(s"Static member $groupInstanceId with unknown member id rejoins, assigning new member id $newMemberId, while " +
+            s"old member $oldMemberId will be removed.")
+
+          val currentLeader = group.leaderOrNull
+          val member = group.replaceGroupInstance(oldMemberId, newMemberId, groupInstanceId)
+          // Heartbeat of old member id will expire without effect since the group no longer contains that member id.
+          // New heartbeat shall be scheduled with new member id.
+          completeAndScheduleNextHeartbeatExpiration(group, member)
+
+          val knownStaticMember = group.get(newMemberId)
+          group.updateMember(knownStaticMember, protocols, responseCallback)
+
+          group.currentState match {
+            case Stable | CompletingRebalance =>
+              info(s"Static member joins during ${group.currentState} stage will not trigger rebalance.")
+              group.maybeInvokeJoinCallback(member, JoinGroupResult(
+                members = List.empty,
+                memberId = newMemberId,
+                generationId = group.generationId,
+                subProtocol = group.protocolOrNull,
+                // We want to avoid current leader performing trivial assignment while the group
+                // is in stable/awaiting sync stage, because the new assignment in leader's next sync call
+                // won't be broadcast by a stable/awaiting sync group. This could be guaranteed by
+                // always returning the old leader id so that the current leader won't assume itself
+                // as a leader based on the returned message, since the new member.id won't match
+                // returned leader id, therefore no assignment will be performed.
+                leaderId = currentLeader,
+                error = Errors.NONE))
+            case Empty | Dead =>
+              throw new IllegalStateException(s"Group ${group.groupId} was not supposed to be " +
+                s"in the state ${group.currentState} when the unknown static member $groupInstanceId rejoins.")
+            case PreparingRebalance =>
+          }
+        } else if (requireKnownMemberId) {
+            // If member id required (dynamic membership), register the member in the pending member list
+            // and send back a response to call for another join group request with allocated member id.
+          debug(s"Dynamic member with unknown member id rejoins group ${group.groupId} in " +
+              s"${group.currentState} state. Created a new member id $newMemberId and request the member to rejoin with this id.")
           group.addPendingMember(newMemberId)
           addPendingMemberExpiration(group, newMemberId, sessionTimeoutMs)
           responseCallback(joinError(newMemberId, Errors.MEMBER_ID_REQUIRED))
         } else {
-          addMemberAndRebalance(rebalanceTimeoutMs, sessionTimeoutMs, newMemberId, clientId, clientHost, protocolType,
-            protocols, group, responseCallback)
+          debug(s"Dynamic member with unknown member id rejoins group ${group.groupId} in " +
+            s"${group.currentState} state. Created a new member id $newMemberId for this member and add to the group.")
+          addMemberAndRebalance(rebalanceTimeoutMs, sessionTimeoutMs, newMemberId, groupInstanceId,
+            clientId, clientHost, protocolType, protocols, group, responseCallback)
         }
       }
     }
@@ -190,82 +260,95 @@ class GroupCoordinator(val brokerId: Int,
   
   private def doJoinGroup(group: GroupMetadata,
                           memberId: String,
+                          groupInstanceId: Option[String],
                           clientId: String,
                           clientHost: String,
                           rebalanceTimeoutMs: Int,
                           sessionTimeoutMs: Int,
                           protocolType: String,
                           protocols: List[(String, Array[Byte])],
-                          responseCallback: JoinCallback) {
+                          responseCallback: JoinCallback): Unit = {
     group.inLock {
       if (group.is(Dead)) {
         // if the group is marked as dead, it means some other thread has just removed the group
         // from the coordinator metadata; this is likely that the group has migrated to some other
         // coordinator OR the group is in a transient unstable phase. Let the member retry
-        // joining without the specified member id.
-        responseCallback(joinError(memberId, Errors.UNKNOWN_MEMBER_ID))
+        // finding the correct coordinator and rejoin.
+        responseCallback(joinError(memberId, Errors.COORDINATOR_NOT_AVAILABLE))
       } else if (!group.supportsProtocols(protocolType, MemberMetadata.plainProtocolSet(protocols))) {
         responseCallback(joinError(memberId, Errors.INCONSISTENT_GROUP_PROTOCOL))
       } else if (group.isPendingMember(memberId)) {
-        // A rejoining pending member will be accepted.
-        addMemberAndRebalance(rebalanceTimeoutMs, sessionTimeoutMs, memberId, clientId, clientHost, protocolType,
-          protocols, group, responseCallback)
-      } else if (!group.has(memberId)) {
-        // if the member trying to register with a un-recognized id, send the response to let
-        // it reset its member id and retry.
-        responseCallback(joinError(memberId, Errors.UNKNOWN_MEMBER_ID))
+        // A rejoining pending member will be accepted. Note that pending member will never be a static member.
+        if (groupInstanceId.isDefined) {
+          throw new IllegalStateException(s"the static member $groupInstanceId was not expected to be assigned " +
+            s"into pending member bucket with member id $memberId")
+        } else {
+          addMemberAndRebalance(rebalanceTimeoutMs, sessionTimeoutMs, memberId, groupInstanceId,
+            clientId, clientHost, protocolType, protocols, group, responseCallback)
+        }
       } else {
-        group.currentState match {
-          case PreparingRebalance =>
-            val member = group.get(memberId)
-            updateMemberAndRebalance(group, member, protocols, responseCallback)
+        val groupInstanceIdNotFound = groupInstanceId.isDefined && !group.hasStaticMember(groupInstanceId)
+        if (group.isStaticMemberFenced(memberId, groupInstanceId)) {
+          // given member id doesn't match with the groupInstanceId. Inform duplicate instance to shut down immediately.
+          responseCallback(joinError(memberId, Errors.FENCED_INSTANCE_ID))
+        } else if (!group.has(memberId) || groupInstanceIdNotFound) {
+            // If the dynamic member trying to register with an unrecognized id, or
+            // the static member joins with unknown group instance id, send the response to let
+            // it reset its member id and retry.
+          responseCallback(joinError(memberId, Errors.UNKNOWN_MEMBER_ID))
+        } else {
+          val member = group.get(memberId)
 
-          case CompletingRebalance =>
-            val member = group.get(memberId)
-            if (member.matches(protocols)) {
-              // member is joining with the same metadata (which could be because it failed to
-              // receive the initial JoinGroup response), so just return current group information
-              // for the current generation.
-              responseCallback(JoinGroupResult(
-                members = if (group.isLeader(memberId)) {
-                  group.currentMemberMetadata
-                } else {
-                  Map.empty
-                },
-                memberId = memberId,
-                generationId = group.generationId,
-                subProtocol = group.protocolOrNull,
-                leaderId = group.leaderOrNull,
-                error = Errors.NONE))
-            } else {
-              // member has changed metadata, so force a rebalance
+          group.currentState match {
+            case PreparingRebalance =>
               updateMemberAndRebalance(group, member, protocols, responseCallback)
-            }
 
-          case Stable =>
-            val member = group.get(memberId)
-            if (group.isLeader(memberId) || !member.matches(protocols)) {
-              // force a rebalance if a member has changed metadata or if the leader sends JoinGroup.
-              // The latter allows the leader to trigger rebalances for changes affecting assignment
-              // which do not affect the member metadata (such as topic metadata changes for the consumer)
-              updateMemberAndRebalance(group, member, protocols, responseCallback)
-            } else {
-              // for followers with no actual change to their metadata, just return group information
-              // for the current generation which will allow them to issue SyncGroup
-              responseCallback(JoinGroupResult(
-                members = Map.empty,
-                memberId = memberId,
-                generationId = group.generationId,
-                subProtocol = group.protocolOrNull,
-                leaderId = group.leaderOrNull,
-                error = Errors.NONE))
-            }
+            case CompletingRebalance =>
+              if (member.matches(protocols)) {
+                // member is joining with the same metadata (which could be because it failed to
+                // receive the initial JoinGroup response), so just return current group information
+                // for the current generation.
+                responseCallback(JoinGroupResult(
+                  members = if (group.isLeader(memberId)) {
+                    group.currentMemberMetadata
+                  } else {
+                    List.empty
+                  },
+                  memberId = memberId,
+                  generationId = group.generationId,
+                  subProtocol = group.protocolOrNull,
+                  leaderId = group.leaderOrNull,
+                  error = Errors.NONE))
+              } else {
+                // member has changed metadata, so force a rebalance
+                updateMemberAndRebalance(group, member, protocols, responseCallback)
+              }
 
-          case Empty | Dead =>
-            // Group reaches unexpected state. Let the joining member reset their generation and rejoin.
-            warn(s"Attempt to add rejoining member $memberId of group ${group.groupId} in " +
-              s"unexpected group state ${group.currentState}")
-            responseCallback(joinError(memberId, Errors.UNKNOWN_MEMBER_ID))
+            case Stable =>
+              val member = group.get(memberId)
+              if (group.isLeader(memberId) || !member.matches(protocols)) {
+                // force a rebalance if a member has changed metadata or if the leader sends JoinGroup.
+                // The latter allows the leader to trigger rebalances for changes affecting assignment
+                // which do not affect the member metadata (such as topic metadata changes for the consumer)
+                updateMemberAndRebalance(group, member, protocols, responseCallback)
+              } else {
+                // for followers with no actual change to their metadata, just return group information
+                // for the current generation which will allow them to issue SyncGroup
+                responseCallback(JoinGroupResult(
+                  members = List.empty,
+                  memberId = memberId,
+                  generationId = group.generationId,
+                  subProtocol = group.protocolOrNull,
+                  leaderId = group.leaderOrNull,
+                  error = Errors.NONE))
+              }
+
+            case Empty | Dead =>
+              // Group reaches unexpected state. Let the joining member reset their generation and rejoin.
+              warn(s"Attempt to add rejoining member $memberId of group ${group.groupId} in " +
+                s"unexpected group state ${group.currentState}")
+              responseCallback(joinError(memberId, Errors.UNKNOWN_MEMBER_ID))
+          }
         }
       }
     }
@@ -274,6 +357,7 @@ class GroupCoordinator(val brokerId: Int,
   def handleSyncGroup(groupId: String,
                       generation: Int,
                       memberId: String,
+                      groupInstanceId: Option[String],
                       groupAssignment: Map[String, Array[Byte]],
                       responseCallback: SyncCallback): Unit = {
     validateGroupStatus(groupId, ApiKeys.SYNC_GROUP) match {
@@ -282,14 +366,14 @@ class GroupCoordinator(val brokerId: Int,
         // group will need to start over at JoinGroup. By returning rebalance in progress, the consumer
         // will attempt to rejoin without needing to rediscover the coordinator. Note that we cannot
         // return COORDINATOR_LOAD_IN_PROGRESS since older clients do not expect the error.
-        responseCallback(Array.empty, Errors.REBALANCE_IN_PROGRESS)
+        responseCallback(SyncGroupResult(Array.empty, Errors.REBALANCE_IN_PROGRESS))
 
-      case Some(error) => responseCallback(Array.empty, error)
+      case Some(error) => responseCallback(SyncGroupResult(Array.empty, error))
 
       case None =>
         groupManager.getGroup(groupId) match {
-          case None => responseCallback(Array.empty, Errors.UNKNOWN_MEMBER_ID)
-          case Some(group) => doSyncGroup(group, generation, memberId, groupAssignment, responseCallback)
+          case None => responseCallback(SyncGroupResult(Array.empty, Errors.UNKNOWN_MEMBER_ID))
+          case Some(group) => doSyncGroup(group, generation, memberId, groupInstanceId, groupAssignment, responseCallback)
         }
     }
   }
@@ -297,20 +381,29 @@ class GroupCoordinator(val brokerId: Int,
   private def doSyncGroup(group: GroupMetadata,
                           generationId: Int,
                           memberId: String,
+                          groupInstanceId: Option[String],
                           groupAssignment: Map[String, Array[Byte]],
-                          responseCallback: SyncCallback) {
+                          responseCallback: SyncCallback): Unit = {
     group.inLock {
-      if (!group.has(memberId)) {
-        responseCallback(Array.empty, Errors.UNKNOWN_MEMBER_ID)
+      if (group.is(Dead)) {
+        // if the group is marked as dead, it means some other thread has just removed the group
+        // from the coordinator metadata; this is likely that the group has migrated to some other
+        // coordinator OR the group is in a transient unstable phase. Let the member retry
+        // finding the correct coordinator and rejoin.
+        responseCallback(SyncGroupResult(Array.empty, Errors.COORDINATOR_NOT_AVAILABLE))
+      } else if (group.isStaticMemberFenced(memberId, groupInstanceId)) {
+        responseCallback(SyncGroupResult(Array.empty, Errors.FENCED_INSTANCE_ID))
+      } else if (!group.has(memberId)) {
+        responseCallback(SyncGroupResult(Array.empty, Errors.UNKNOWN_MEMBER_ID))
       } else if (generationId != group.generationId) {
-        responseCallback(Array.empty, Errors.ILLEGAL_GENERATION)
+        responseCallback(SyncGroupResult(Array.empty, Errors.ILLEGAL_GENERATION))
       } else {
         group.currentState match {
-          case Empty | Dead =>
-            responseCallback(Array.empty, Errors.UNKNOWN_MEMBER_ID)
+          case Empty =>
+            responseCallback(SyncGroupResult(Array.empty, Errors.UNKNOWN_MEMBER_ID))
 
           case PreparingRebalance =>
-            responseCallback(Array.empty, Errors.REBALANCE_IN_PROGRESS)
+            responseCallback(SyncGroupResult(Array.empty, Errors.REBALANCE_IN_PROGRESS))
 
           case CompletingRebalance =>
             group.get(memberId).awaitingSyncCallback = responseCallback
@@ -339,52 +432,74 @@ class GroupCoordinator(val brokerId: Int,
                   }
                 }
               })
+              groupCompletedRebalanceSensor.record()
             }
 
           case Stable =>
             // if the group is stable, we just return the current assignment
             val memberMetadata = group.get(memberId)
-            responseCallback(memberMetadata.assignment, Errors.NONE)
+            responseCallback(SyncGroupResult(memberMetadata.assignment, Errors.NONE))
             completeAndScheduleNextHeartbeatExpiration(group, group.get(memberId))
+
+          case Dead =>
+            throw new IllegalStateException(s"Reached unexpected condition for Dead group ${group.groupId}")
         }
       }
     }
   }
 
-  def handleLeaveGroup(groupId: String, memberId: String, responseCallback: Errors => Unit): Unit = {
-    validateGroupStatus(groupId, ApiKeys.LEAVE_GROUP).foreach { error =>
-      responseCallback(error)
-      return
-    }
-
-    groupManager.getGroup(groupId) match {
+  def handleLeaveGroup(groupId: String,
+                       leavingMembers: List[MemberIdentity],
+                       responseCallback: LeaveGroupResult => Unit): Unit = {
+    validateGroupStatus(groupId, ApiKeys.LEAVE_GROUP) match {
+      case Some(error) =>
+        responseCallback(leaveError(error, List.empty))
       case None =>
-        // if the group is marked as dead, it means some other thread has just removed the group
-        // from the coordinator metadata; it is likely that the group has migrated to some other
-        // coordinator OR the group is in a transient unstable phase. Let the consumer to retry
-        // joining without specified consumer id,
-        responseCallback(Errors.UNKNOWN_MEMBER_ID)
-
-      case Some(group) =>
-        group.inLock {
-          if (group.is(Dead)) {
-            responseCallback(Errors.UNKNOWN_MEMBER_ID)
-          } else if (group.isPendingMember(memberId)) {
-            // if a pending member is leaving, it needs to be removed from the pending list, heartbeat cancelled
-            // and if necessary, prompt a JoinGroup completion.
-            info(s"Pending member $memberId is leaving group ${group.groupId}.")
-            removePendingMemberAndUpdateGroup(group, memberId)
-            heartbeatPurgatory.checkAndComplete(MemberKey(group.groupId, memberId))
-            responseCallback(Errors.NONE)
-          } else if (!group.has(memberId)) {
-            responseCallback(Errors.UNKNOWN_MEMBER_ID)
-          } else {
-            val member = group.get(memberId)
-            removeHeartbeatForLeavingMember(group, member)
-            info(s"Member ${member.memberId} in group ${group.groupId} has left, removing it from the group")
-            removeMemberAndUpdateGroup(group, member, s"removing member $memberId on LeaveGroup")
-            responseCallback(Errors.NONE)
-          }
+        groupManager.getGroup(groupId) match {
+          case None =>
+            responseCallback(leaveError(Errors.NONE, leavingMembers.map {leavingMember =>
+              memberLeaveError(leavingMember, Errors.UNKNOWN_MEMBER_ID)
+            }))
+          case Some(group) =>
+            group.inLock {
+              if (group.is(Dead)) {
+                responseCallback(leaveError(Errors.COORDINATOR_NOT_AVAILABLE, List.empty))
+              } else {
+                val memberErrors = leavingMembers.map { leavingMember =>
+                  val memberId = leavingMember.memberId
+                  val groupInstanceId = Option(leavingMember.groupInstanceId)
+                  if (memberId != JoinGroupRequest.UNKNOWN_MEMBER_ID
+                    && group.isStaticMemberFenced(memberId, groupInstanceId)) {
+                    memberLeaveError(leavingMember, Errors.FENCED_INSTANCE_ID)
+                  } else if (group.isPendingMember(memberId)) {
+                    if (groupInstanceId.isDefined) {
+                      throw new IllegalStateException(s"the static member $groupInstanceId was not expected to be leaving " +
+                        s"from pending member bucket with member id $memberId")
+                    } else {
+                      // if a pending member is leaving, it needs to be removed from the pending list, heartbeat cancelled
+                      // and if necessary, prompt a JoinGroup completion.
+                      info(s"Pending member $memberId is leaving group ${group.groupId}.")
+                      removePendingMemberAndUpdateGroup(group, memberId)
+                      heartbeatPurgatory.checkAndComplete(MemberKey(group.groupId, memberId))
+                      memberLeaveError(leavingMember, Errors.NONE)
+                    }
+                  } else if (!group.has(memberId) && !group.hasStaticMember(groupInstanceId)) {
+                    memberLeaveError(leavingMember, Errors.UNKNOWN_MEMBER_ID)
+                  } else {
+                    val member = if (group.hasStaticMember(groupInstanceId))
+                      group.get(group.getStaticMemberId(groupInstanceId))
+                    else
+                      group.get(memberId)
+                    removeHeartbeatForLeavingMember(group, member)
+                    info(s"Member[group.instance.id ${member.groupInstanceId}, member.id ${member.memberId}] " +
+                      s"in group ${group.groupId} has left, removing it from the group")
+                    removeMemberAndUpdateGroup(group, member, s"removing member $memberId on LeaveGroup")
+                    memberLeaveError(leavingMember, Errors.NONE)
+                  }
+                }
+                responseCallback(leaveError(Errors.NONE, memberErrors))
+              }
+            }
         }
     }
   }
@@ -412,7 +527,7 @@ class GroupCoordinator(val brokerId: Int,
                   case Empty =>
                     group.transitionTo(Dead)
                     groupsEligibleForDeletion :+= group
-                  case _ =>
+                  case Stable | PreparingRebalance | CompletingRebalance =>
                     groupErrors += groupId -> Errors.NON_EMPTY_GROUP
                 }
               }
@@ -430,10 +545,67 @@ class GroupCoordinator(val brokerId: Int,
     groupErrors
   }
 
+  def handleDeleteOffsets(groupId: String, partitions: Seq[TopicPartition]): (Errors, Map[TopicPartition, Errors]) = {
+    var groupError: Errors = Errors.NONE
+    var partitionErrors: Map[TopicPartition, Errors] = Map()
+    var partitionsEligibleForDeletion: Seq[TopicPartition] = Seq()
+
+    validateGroupStatus(groupId, ApiKeys.OFFSET_DELETE) match {
+      case Some(error) =>
+        groupError = error
+
+      case None =>
+        groupManager.getGroup(groupId) match {
+          case None =>
+            groupError = if (groupManager.groupNotExists(groupId))
+              Errors.GROUP_ID_NOT_FOUND else Errors.NOT_COORDINATOR
+
+          case Some(group) =>
+            group.inLock {
+              group.currentState match {
+                case Dead =>
+                  groupError = if (groupManager.groupNotExists(groupId))
+                    Errors.GROUP_ID_NOT_FOUND else Errors.NOT_COORDINATOR
+
+                case Empty =>
+                  partitionsEligibleForDeletion = partitions
+
+                case PreparingRebalance | CompletingRebalance | Stable if group.isConsumerGroup =>
+                  val (consumed, notConsumed) =
+                    partitions.partition(tp => group.isSubscribedToTopic(tp.topic()))
+
+                  partitionsEligibleForDeletion = notConsumed
+                  partitionErrors = consumed.map(_ -> Errors.GROUP_SUBSCRIBED_TO_TOPIC).toMap
+
+                case _ =>
+                  groupError = Errors.NON_EMPTY_GROUP
+              }
+            }
+
+            if (partitionsEligibleForDeletion.nonEmpty) {
+              val offsetsRemoved = groupManager.cleanupGroupMetadata(Seq(group), group => {
+                group.removeOffsets(partitionsEligibleForDeletion)
+              })
+
+              partitionErrors ++= partitionsEligibleForDeletion.map(_ -> Errors.NONE).toMap
+
+              offsetDeletionSensor.record(offsetsRemoved)
+
+              info(s"The following offsets of the group $groupId were deleted: ${partitionsEligibleForDeletion.mkString(", ")}. " +
+                s"A total of $offsetsRemoved offsets were removed.")
+            }
+        }
+    }
+
+    // If there is a group error, the partition errors is empty
+    groupError -> partitionErrors
+  }
+
   def handleHeartbeat(groupId: String,
                       memberId: String,
+                      groupInstanceId: Option[String],
                       generationId: Int,
-                      responseCallback: Errors => Unit) {
+                      responseCallback: Errors => Unit): Unit = {
     validateGroupStatus(groupId, ApiKeys.HEARTBEAT).foreach { error =>
       if (error == Errors.COORDINATOR_LOAD_IN_PROGRESS)
         // the group is still loading, so respond just blindly
@@ -448,44 +620,39 @@ class GroupCoordinator(val brokerId: Int,
         responseCallback(Errors.UNKNOWN_MEMBER_ID)
 
       case Some(group) => group.inLock {
-        group.currentState match {
-          case Dead =>
-            // if the group is marked as dead, it means some other thread has just removed the group
-            // from the coordinator metadata; it is likely that the group has migrated to some other
-            // coordinator OR the group is in a transient unstable phase. Let the member retry
-            // joining without the specified member id.
-            responseCallback(Errors.UNKNOWN_MEMBER_ID)
+        if (group.is(Dead)) {
+          // if the group is marked as dead, it means some other thread has just removed the group
+          // from the coordinator metadata; this is likely that the group has migrated to some other
+          // coordinator OR the group is in a transient unstable phase. Let the member retry
+          // finding the correct coordinator and rejoin.
+          responseCallback(Errors.COORDINATOR_NOT_AVAILABLE)
+        } else if (group.isStaticMemberFenced(memberId, groupInstanceId)) {
+          responseCallback(Errors.FENCED_INSTANCE_ID)
+        } else if (!group.has(memberId)) {
+          responseCallback(Errors.UNKNOWN_MEMBER_ID)
+        } else if (generationId != group.generationId) {
+          responseCallback(Errors.ILLEGAL_GENERATION)
+        } else {
+          group.currentState match {
+            case Empty =>
+              responseCallback(Errors.UNKNOWN_MEMBER_ID)
 
-          case Empty =>
-            responseCallback(Errors.UNKNOWN_MEMBER_ID)
+            case CompletingRebalance =>
+                responseCallback(Errors.REBALANCE_IN_PROGRESS)
 
-          case CompletingRebalance =>
-            if (!group.has(memberId))
-              responseCallback(Errors.UNKNOWN_MEMBER_ID)
-            else
-              responseCallback(Errors.REBALANCE_IN_PROGRESS)
+            case PreparingRebalance =>
+                val member = group.get(memberId)
+                completeAndScheduleNextHeartbeatExpiration(group, member)
+                responseCallback(Errors.REBALANCE_IN_PROGRESS)
 
-          case PreparingRebalance =>
-            if (!group.has(memberId)) {
-              responseCallback(Errors.UNKNOWN_MEMBER_ID)
-            } else if (generationId != group.generationId) {
-              responseCallback(Errors.ILLEGAL_GENERATION)
-            } else {
-              val member = group.get(memberId)
-              completeAndScheduleNextHeartbeatExpiration(group, member)
-              responseCallback(Errors.REBALANCE_IN_PROGRESS)
-            }
+            case Stable =>
+                val member = group.get(memberId)
+                completeAndScheduleNextHeartbeatExpiration(group, member)
+                responseCallback(Errors.NONE)
 
-          case Stable =>
-            if (!group.has(memberId)) {
-              responseCallback(Errors.UNKNOWN_MEMBER_ID)
-            } else if (generationId != group.generationId) {
-              responseCallback(Errors.ILLEGAL_GENERATION)
-            } else {
-              val member = group.get(memberId)
-              completeAndScheduleNextHeartbeatExpiration(group, member)
-              responseCallback(Errors.NONE)
-            }
+            case Dead =>
+              throw new IllegalStateException(s"Reached unexpected condition for Dead group $groupId")
+          }
         }
       }
     }
@@ -497,37 +664,38 @@ class GroupCoordinator(val brokerId: Int,
                              offsetMetadata: immutable.Map[TopicPartition, OffsetAndMetadata],
                              responseCallback: immutable.Map[TopicPartition, Errors] => Unit): Unit = {
     validateGroupStatus(groupId, ApiKeys.TXN_OFFSET_COMMIT) match {
-      case Some(error) => responseCallback(offsetMetadata.mapValues(_ => error))
+      case Some(error) => responseCallback(offsetMetadata.map { case (k, _) => k -> error })
       case None =>
         val group = groupManager.getGroup(groupId).getOrElse {
           groupManager.addGroup(new GroupMetadata(groupId, Empty, time))
         }
-        doCommitOffsets(group, NoMemberId, NoGeneration, producerId, producerEpoch, offsetMetadata, responseCallback)
+        doCommitOffsets(group, NoMemberId, None, NoGeneration, producerId, producerEpoch, offsetMetadata, responseCallback)
     }
   }
 
   def handleCommitOffsets(groupId: String,
                           memberId: String,
+                          groupInstanceId: Option[String],
                           generationId: Int,
                           offsetMetadata: immutable.Map[TopicPartition, OffsetAndMetadata],
-                          responseCallback: immutable.Map[TopicPartition, Errors] => Unit) {
+                          responseCallback: immutable.Map[TopicPartition, Errors] => Unit): Unit = {
     validateGroupStatus(groupId, ApiKeys.OFFSET_COMMIT) match {
-      case Some(error) => responseCallback(offsetMetadata.mapValues(_ => error))
+      case Some(error) => responseCallback(offsetMetadata.map { case (k, _) => k -> error })
       case None =>
         groupManager.getGroup(groupId) match {
           case None =>
             if (generationId < 0) {
               // the group is not relying on Kafka for group management, so allow the commit
               val group = groupManager.addGroup(new GroupMetadata(groupId, Empty, time))
-              doCommitOffsets(group, memberId, generationId, NO_PRODUCER_ID, NO_PRODUCER_EPOCH,
+              doCommitOffsets(group, memberId, groupInstanceId, generationId, NO_PRODUCER_ID, NO_PRODUCER_EPOCH,
                 offsetMetadata, responseCallback)
             } else {
               // or this is a request coming from an older generation. either way, reject the commit
-              responseCallback(offsetMetadata.mapValues(_ => Errors.ILLEGAL_GENERATION))
+              responseCallback(offsetMetadata.map { case (k, _) => k -> Errors.ILLEGAL_GENERATION })
             }
 
           case Some(group) =>
-            doCommitOffsets(group, memberId, generationId, NO_PRODUCER_ID, NO_PRODUCER_EPOCH,
+            doCommitOffsets(group, memberId, groupInstanceId, generationId, NO_PRODUCER_ID, NO_PRODUCER_EPOCH,
               offsetMetadata, responseCallback)
         }
     }
@@ -535,7 +703,7 @@ class GroupCoordinator(val brokerId: Int,
 
   def scheduleHandleTxnCompletion(producerId: Long,
                                   offsetsPartitions: Iterable[TopicPartition],
-                                  transactionResult: TransactionResult) {
+                                  transactionResult: TransactionResult): Unit = {
     require(offsetsPartitions.forall(_.topic == Topic.GROUP_METADATA_TOPIC_NAME))
     val isCommit = transactionResult == TransactionResult.COMMIT
     groupManager.scheduleHandleTxnCompletion(producerId, offsetsPartitions.map(_.partition).toSet, isCommit)
@@ -543,28 +711,48 @@ class GroupCoordinator(val brokerId: Int,
 
   private def doCommitOffsets(group: GroupMetadata,
                               memberId: String,
+                              groupInstanceId: Option[String],
                               generationId: Int,
                               producerId: Long,
                               producerEpoch: Short,
                               offsetMetadata: immutable.Map[TopicPartition, OffsetAndMetadata],
-                              responseCallback: immutable.Map[TopicPartition, Errors] => Unit) {
+                              responseCallback: immutable.Map[TopicPartition, Errors] => Unit): Unit = {
     group.inLock {
       if (group.is(Dead)) {
-        responseCallback(offsetMetadata.mapValues(_ => Errors.UNKNOWN_MEMBER_ID))
+        // if the group is marked as dead, it means some other thread has just removed the group
+        // from the coordinator metadata; it is likely that the group has migrated to some other
+        // coordinator OR the group is in a transient unstable phase. Let the member retry
+        // finding the correct coordinator and rejoin.
+        responseCallback(offsetMetadata.map { case (k, _) => k -> Errors.COORDINATOR_NOT_AVAILABLE })
+      } else if (group.isStaticMemberFenced(memberId, groupInstanceId)) {
+        responseCallback(offsetMetadata.map { case (k, _) => k -> Errors.FENCED_INSTANCE_ID })
       } else if ((generationId < 0 && group.is(Empty)) || (producerId != NO_PRODUCER_ID)) {
         // The group is only using Kafka to store offsets.
         // Also, for transactional offset commits we don't need to validate group membership and the generation.
         groupManager.storeOffsets(group, memberId, offsetMetadata, responseCallback, producerId, producerEpoch)
-      } else if (group.is(CompletingRebalance)) {
-        responseCallback(offsetMetadata.mapValues(_ => Errors.REBALANCE_IN_PROGRESS))
       } else if (!group.has(memberId)) {
-        responseCallback(offsetMetadata.mapValues(_ => Errors.UNKNOWN_MEMBER_ID))
+        responseCallback(offsetMetadata.map { case (k, _) => k -> Errors.UNKNOWN_MEMBER_ID })
       } else if (generationId != group.generationId) {
-        responseCallback(offsetMetadata.mapValues(_ => Errors.ILLEGAL_GENERATION))
+        responseCallback(offsetMetadata.map { case (k, _) => k -> Errors.ILLEGAL_GENERATION })
       } else {
-        val member = group.get(memberId)
-        completeAndScheduleNextHeartbeatExpiration(group, member)
-        groupManager.storeOffsets(group, memberId, offsetMetadata, responseCallback)
+        group.currentState match {
+          case Stable | PreparingRebalance =>
+            // During PreparingRebalance phase, we still allow a commit request since we rely
+            // on heartbeat response to eventually notify the rebalance in progress signal to the consumer
+            val member = group.get(memberId)
+            completeAndScheduleNextHeartbeatExpiration(group, member)
+            groupManager.storeOffsets(group, memberId, offsetMetadata, responseCallback)
+
+          case CompletingRebalance =>
+            // We should not receive a commit request if the group has not completed rebalance;
+            // but since the consumer's member.id and generation is valid, it means it has received
+            // the latest group generation information from the JoinResponse.
+            // So let's return a REBALANCE_IN_PROGRESS to let consumer handle it gracefully.
+            responseCallback(offsetMetadata.map { case (k, _) => k -> Errors.REBALANCE_IN_PROGRESS })
+
+          case _ =>
+            throw new RuntimeException(s"Logic error: unexpected group state ${group.currentState}")
+        }
       }
     }
   }
@@ -604,7 +792,7 @@ class GroupCoordinator(val brokerId: Int,
     }
   }
 
-  def handleDeletedPartitions(topicPartitions: Seq[TopicPartition]) {
+  def handleDeletedPartitions(topicPartitions: Seq[TopicPartition]): Unit = {
     val offsetsRemoved = groupManager.cleanupGroupMetadata(groupManager.currentGroups, group => {
       group.removeOffsets(topicPartitions)
     })
@@ -639,7 +827,7 @@ class GroupCoordinator(val brokerId: Int,
       None
   }
 
-  private def onGroupUnloaded(group: GroupMetadata) {
+  private def onGroupUnloaded(group: GroupMetadata): Unit = {
     group.inLock {
       info(s"Unloading group metadata for ${group.groupId} with generation ${group.generationId}")
       val previousState = group.currentState
@@ -656,17 +844,14 @@ class GroupCoordinator(val brokerId: Int,
 
         case Stable | CompletingRebalance =>
           for (member <- group.allMemberMetadata) {
-            if (member.awaitingSyncCallback != null) {
-              member.awaitingSyncCallback(Array.empty[Byte], Errors.NOT_COORDINATOR)
-              member.awaitingSyncCallback = null
-            }
+            group.maybeInvokeSyncCallback(member, SyncGroupResult(Array.empty, Errors.NOT_COORDINATOR))
             heartbeatPurgatory.checkAndComplete(MemberKey(member.groupId, member.memberId))
           }
       }
     }
   }
 
-  private def onGroupLoaded(group: GroupMetadata) {
+  private def onGroupLoaded(group: GroupMetadata): Unit = {
     group.inLock {
       info(s"Loading group metadata for ${group.groupId} with generation ${group.generationId}")
       assert(group.is(Stable) || group.is(Empty))
@@ -678,32 +863,39 @@ class GroupCoordinator(val brokerId: Int,
     }
   }
 
-  def handleGroupImmigration(offsetTopicPartitionId: Int) {
+  /**
+   * Load cached state from the given partition and begin handling requests for groups which map to it.
+   *
+   * @param offsetTopicPartitionId The partition we are now leading
+   */
+  def onElection(offsetTopicPartitionId: Int): Unit = {
     groupManager.scheduleLoadGroupAndOffsets(offsetTopicPartitionId, onGroupLoaded)
   }
 
-  def handleGroupEmigration(offsetTopicPartitionId: Int) {
+  /**
+   * Unload cached state for the given partition and stop handling requests for groups which map to it.
+   *
+   * @param offsetTopicPartitionId The partition we are no longer leading
+   */
+  def onResignation(offsetTopicPartitionId: Int): Unit = {
     groupManager.removeGroupsForPartition(offsetTopicPartitionId, onGroupUnloaded)
   }
 
-  private def setAndPropagateAssignment(group: GroupMetadata, assignment: Map[String, Array[Byte]]) {
+  private def setAndPropagateAssignment(group: GroupMetadata, assignment: Map[String, Array[Byte]]): Unit = {
     assert(group.is(CompletingRebalance))
     group.allMemberMetadata.foreach(member => member.assignment = assignment(member.memberId))
     propagateAssignment(group, Errors.NONE)
   }
 
-  private def resetAndPropagateAssignmentError(group: GroupMetadata, error: Errors) {
+  private def resetAndPropagateAssignmentError(group: GroupMetadata, error: Errors): Unit = {
     assert(group.is(CompletingRebalance))
-    group.allMemberMetadata.foreach(_.assignment = Array.empty[Byte])
+    group.allMemberMetadata.foreach(_.assignment = Array.empty)
     propagateAssignment(group, error)
   }
 
-  private def propagateAssignment(group: GroupMetadata, error: Errors) {
+  private def propagateAssignment(group: GroupMetadata, error: Errors): Unit = {
     for (member <- group.allMemberMetadata) {
-      if (member.awaitingSyncCallback != null) {
-        member.awaitingSyncCallback(member.assignment, error)
-        member.awaitingSyncCallback = null
-
+      if (group.maybeInvokeSyncCallback(member, SyncGroupResult(member.assignment, error))) {
         // reset the session timeout for members after propagating the member's assignment.
         // This is because if any member's session expired while we were still awaiting either
         // the leader sync group or the storage callback, its expiration will be ignored and no
@@ -713,46 +905,36 @@ class GroupCoordinator(val brokerId: Int,
     }
   }
 
-  private def joinError(memberId: String, error: Errors): JoinGroupResult = {
-    JoinGroupResult(
-      members = Map.empty,
-      memberId = memberId,
-      generationId = GroupCoordinator.NoGeneration,
-      subProtocol = GroupCoordinator.NoProtocol,
-      leaderId = GroupCoordinator.NoLeader,
-      error = error)
-  }
-
   /**
    * Complete existing DelayedHeartbeats for the given member and schedule the next one
    */
-  private def completeAndScheduleNextHeartbeatExpiration(group: GroupMetadata, member: MemberMetadata) {
+  private def completeAndScheduleNextHeartbeatExpiration(group: GroupMetadata, member: MemberMetadata): Unit = {
     completeAndScheduleNextExpiration(group, member, member.sessionTimeoutMs)
   }
 
-  private def completeAndScheduleNextExpiration(group: GroupMetadata, member: MemberMetadata, timeoutMs: Long) {
-    // complete current heartbeat expectation
-    member.latestHeartbeat = time.milliseconds()
+  private def completeAndScheduleNextExpiration(group: GroupMetadata, member: MemberMetadata, timeoutMs: Long): Unit = {
     val memberKey = MemberKey(member.groupId, member.memberId)
+
+    // complete current heartbeat expectation
+    member.heartbeatSatisfied = true
     heartbeatPurgatory.checkAndComplete(memberKey)
 
     // reschedule the next heartbeat expiration deadline
-    val deadline = member.latestHeartbeat + timeoutMs
-    val delayedHeartbeat = new DelayedHeartbeat(this, group, member.memberId, isPending = false, deadline, timeoutMs)
+    member.heartbeatSatisfied = false
+    val delayedHeartbeat = new DelayedHeartbeat(this, group, member.memberId, isPending = false, timeoutMs)
     heartbeatPurgatory.tryCompleteElseWatch(delayedHeartbeat, Seq(memberKey))
   }
 
   /**
     * Add pending member expiration to heartbeat purgatory
     */
-  private def addPendingMemberExpiration(group: GroupMetadata, pendingMemberId: String, timeoutMs: Long) {
+  private def addPendingMemberExpiration(group: GroupMetadata, pendingMemberId: String, timeoutMs: Long): Unit = {
     val pendingMemberKey = MemberKey(group.groupId, pendingMemberId)
-    val deadline = time.milliseconds() + timeoutMs
-    val delayedHeartbeat = new DelayedHeartbeat(this, group, pendingMemberId, isPending = true, deadline, timeoutMs)
+    val delayedHeartbeat = new DelayedHeartbeat(this, group, pendingMemberId, isPending = true, timeoutMs)
     heartbeatPurgatory.tryCompleteElseWatch(delayedHeartbeat, Seq(pendingMemberKey))
   }
 
-  private def removeHeartbeatForLeavingMember(group: GroupMetadata, member: MemberMetadata) {
+  private def removeHeartbeatForLeavingMember(group: GroupMetadata, member: MemberMetadata): Unit = {
     member.isLeaving = true
     val memberKey = MemberKey(member.groupId, member.memberId)
     heartbeatPurgatory.checkAndComplete(memberKey)
@@ -761,13 +943,15 @@ class GroupCoordinator(val brokerId: Int,
   private def addMemberAndRebalance(rebalanceTimeoutMs: Int,
                                     sessionTimeoutMs: Int,
                                     memberId: String,
+                                    groupInstanceId: Option[String],
                                     clientId: String,
                                     clientHost: String,
                                     protocolType: String,
                                     protocols: List[(String, Array[Byte])],
                                     group: GroupMetadata,
-                                    callback: JoinCallback) {
-    val member = new MemberMetadata(memberId, group.groupId, clientId, clientHost, rebalanceTimeoutMs,
+                                    callback: JoinCallback): Unit = {
+    val member = new MemberMetadata(memberId, group.groupId, groupInstanceId,
+      clientId, clientHost, rebalanceTimeoutMs,
       sessionTimeoutMs, protocolType, protocols)
 
     member.isNew = true
@@ -786,26 +970,29 @@ class GroupCoordinator(val brokerId: Int,
     // for new members. If the new member is still there, we expect it to retry.
     completeAndScheduleNextExpiration(group, member, NewMemberJoinTimeoutMs)
 
-    group.removePendingMember(memberId)
-    maybePrepareRebalance(group, s"Adding new member $memberId")
+    if (member.isStaticMember)
+      group.addStaticMember(groupInstanceId, memberId)
+    else
+      group.removePendingMember(memberId)
+    maybePrepareRebalance(group, s"Adding new member $memberId with group instanceid $groupInstanceId")
   }
 
   private def updateMemberAndRebalance(group: GroupMetadata,
                                        member: MemberMetadata,
                                        protocols: List[(String, Array[Byte])],
-                                       callback: JoinCallback) {
+                                       callback: JoinCallback): Unit = {
     group.updateMember(member, protocols, callback)
     maybePrepareRebalance(group, s"Updating metadata for member ${member.memberId}")
   }
 
-  private def maybePrepareRebalance(group: GroupMetadata, reason: String) {
+  private def maybePrepareRebalance(group: GroupMetadata, reason: String): Unit = {
     group.inLock {
       if (group.canRebalance)
         prepareRebalance(group, reason)
     }
   }
 
-  private def prepareRebalance(group: GroupMetadata, reason: String) {
+  private def prepareRebalance(group: GroupMetadata, reason: String): Unit = {
     // if any members are awaiting sync, cancel their request and have them rejoin
     if (group.is(CompletingRebalance))
       resetAndPropagateAssignmentError(group, Errors.REBALANCE_IN_PROGRESS)
@@ -829,13 +1016,14 @@ class GroupCoordinator(val brokerId: Int,
     joinPurgatory.tryCompleteElseWatch(delayedRebalance, Seq(groupKey))
   }
 
-  private def removeMemberAndUpdateGroup(group: GroupMetadata, member: MemberMetadata, reason: String) {
+  private def removeMemberAndUpdateGroup(group: GroupMetadata, member: MemberMetadata, reason: String): Unit = {
     // New members may timeout with a pending JoinGroup while the group is still rebalancing, so we have
     // to invoke the callback before removing the member. We return UNKNOWN_MEMBER_ID so that the consumer
     // will retry the JoinGroup request if is still active.
     group.maybeInvokeJoinCallback(member, joinError(NoMemberId, Errors.UNKNOWN_MEMBER_ID))
 
     group.remove(member.memberId)
+    group.removeStaticMember(member.groupInstanceId)
 
     group.currentState match {
       case Dead | Empty =>
@@ -844,7 +1032,7 @@ class GroupCoordinator(val brokerId: Int,
     }
   }
 
-  private def removePendingMemberAndUpdateGroup(group: GroupMetadata, memberId: String) {
+  private def removePendingMemberAndUpdateGroup(group: GroupMetadata, memberId: String): Unit = {
     group.removePendingMember(memberId)
 
     if (group.is(PreparingRebalance)) {
@@ -860,20 +1048,31 @@ class GroupCoordinator(val brokerId: Int,
     }
   }
 
-  def onExpireJoin() {
+  def onExpireJoin(): Unit = {
     // TODO: add metrics for restabilize timeouts
   }
 
-  def onCompleteJoin(group: GroupMetadata) {
+  def onCompleteJoin(group: GroupMetadata): Unit = {
     group.inLock {
-      // remove any members who haven't joined the group yet
-      group.notYetRejoinedMembers.foreach { failedMember =>
+      // remove dynamic members who haven't joined the group yet
+      group.notYetRejoinedMembers.filterNot(_.isStaticMember) foreach { failedMember =>
         removeHeartbeatForLeavingMember(group, failedMember)
         group.remove(failedMember.memberId)
+        group.removeStaticMember(failedMember.groupInstanceId)
         // TODO: cut the socket connection to the client
       }
 
-      if (!group.is(Dead)) {
+      if (group.is(Dead)) {
+        info(s"Group ${group.groupId} is dead, skipping rebalance stage")
+      } else if (!group.maybeElectNewJoinedLeader() && group.allMembers.nonEmpty) {
+        // If all members are not rejoining, we will postpone the completion
+        // of rebalance preparing stage, and send out another delayed operation
+        // until session timeout removes all the non-responsive members.
+        error(s"Group ${group.groupId} could not complete rebalance because no members rejoined")
+        joinPurgatory.tryCompleteElseWatch(
+          new DelayedJoin(this, group, group.rebalanceTimeoutMs),
+          Seq(GroupKey(group.groupId)))
+      } else {
         group.initNextGeneration()
         if (group.is(Empty)) {
           info(s"Group ${group.groupId} with generation ${group.generationId} is now empty " +
@@ -893,12 +1092,11 @@ class GroupCoordinator(val brokerId: Int,
 
           // trigger the awaiting join group response callback for all the members after rebalancing
           for (member <- group.allMemberMetadata) {
-            assert(member.awaitingJoinCallback != null)
             val joinResult = JoinGroupResult(
               members = if (group.isLeader(member.memberId)) {
                 group.currentMemberMetadata
               } else {
-                Map.empty
+                List.empty
               },
               memberId = member.memberId,
               generationId = group.generationId,
@@ -915,33 +1113,47 @@ class GroupCoordinator(val brokerId: Int,
     }
   }
 
-  def tryCompleteHeartbeat(group: GroupMetadata, memberId: String, isPending: Boolean, heartbeatDeadline: Long, forceComplete: () => Boolean) = {
+  def tryCompleteHeartbeat(group: GroupMetadata,
+                           memberId: String,
+                           isPending: Boolean,
+                           forceComplete: () => Boolean): Boolean = {
     group.inLock {
-      if (isPending) {
+      // The group has been unloaded and invalid, we should complete the heartbeat.
+      if (group.is(Dead)) {
+        forceComplete()
+      } else if (isPending) {
         // complete the heartbeat if the member has joined the group
         if (group.has(memberId)) {
           forceComplete()
         } else false
-      }
-      else {
-        val member = group.get(memberId)
-        if (member.shouldKeepAlive(heartbeatDeadline) || member.isLeaving) {
-          forceComplete()
-        } else false
-      }
+      } else if (shouldCompleteNonPendingHeartbeat(group, memberId)) {
+        forceComplete()
+      } else false
     }
   }
 
-  def onExpireHeartbeat(group: GroupMetadata, memberId: String, isPending: Boolean, heartbeatDeadline: Long) {
+  def shouldCompleteNonPendingHeartbeat(group: GroupMetadata, memberId: String): Boolean = {
+    if (group.has(memberId)) {
+      val member = group.get(memberId)
+      member.hasSatisfiedHeartbeat || member.isLeaving
+    } else {
+      info(s"Member id $memberId was not found in ${group.groupId} during heartbeat completion check")
+      true
+    }
+  }
+
+  def onExpireHeartbeat(group: GroupMetadata, memberId: String, isPending: Boolean): Unit = {
     group.inLock {
-      if (isPending) {
+      if (group.is(Dead)) {
+        info(s"Received notification of heartbeat expiration for member $memberId after group ${group.groupId} had already been unloaded or deleted.")
+      } else if (isPending) {
         info(s"Pending member $memberId in group ${group.groupId} has been removed after session timeout expiration.")
         removePendingMemberAndUpdateGroup(group, memberId)
       } else if (!group.has(memberId)) {
         debug(s"Member $memberId has already been removed from the group.")
       } else {
         val member = group.get(memberId)
-        if (!member.shouldKeepAlive(heartbeatDeadline)) {
+        if (!member.hasSatisfiedHeartbeat) {
           info(s"Member ${member.memberId} in group ${group.groupId} has failed, removing it from the group")
           removeMemberAndUpdateGroup(group, member, s"removing member ${member.memberId} on heartbeat expiration")
         }
@@ -949,7 +1161,7 @@ class GroupCoordinator(val brokerId: Int,
     }
   }
 
-  def onCompleteHeartbeat() {
+  def onCompleteHeartbeat(): Unit = {
     // TODO: add metrics for complete heartbeats
   }
 
@@ -980,10 +1192,11 @@ object GroupCoordinator {
   def apply(config: KafkaConfig,
             zkClient: KafkaZkClient,
             replicaManager: ReplicaManager,
-            time: Time): GroupCoordinator = {
+            time: Time,
+            metrics: Metrics): GroupCoordinator = {
     val heartbeatPurgatory = DelayedOperationPurgatory[DelayedHeartbeat]("Heartbeat", config.brokerId)
     val joinPurgatory = DelayedOperationPurgatory[DelayedJoin]("Rebalance", config.brokerId)
-    apply(config, zkClient, replicaManager, heartbeatPurgatory, joinPurgatory, time)
+    apply(config, zkClient, replicaManager, heartbeatPurgatory, joinPurgatory, time, metrics)
   }
 
   private[group] def offsetConfig(config: KafkaConfig) = OffsetConfig(
@@ -996,7 +1209,11 @@ object GroupCoordinator {
     offsetsTopicReplicationFactor = config.offsetsTopicReplicationFactor,
     offsetsTopicCompressionCodec = config.offsetsTopicCompressionCodec,
     offsetCommitTimeoutMs = config.offsetCommitTimeoutMs,
-    offsetCommitRequiredAcks = config.offsetCommitRequiredAcks
+    offsetCommitRequiredAcks = config.offsetCommitRequiredAcks,
+    offsetsTopicMaxMessageBytes = config.offsetsTopicMaxMessageBytes,
+    offsetsTopicMinInSyncReplicas = config.offsetsTopicMinInSyncReplicas,
+    offsetsTopicMinCompactionLagMs = config.offsetsTopicMinCompactionLagMs
+
   )
 
   def apply(config: KafkaConfig,
@@ -1004,7 +1221,8 @@ object GroupCoordinator {
             replicaManager: ReplicaManager,
             heartbeatPurgatory: DelayedOperationPurgatory[DelayedHeartbeat],
             joinPurgatory: DelayedOperationPurgatory[DelayedJoin],
-            time: Time): GroupCoordinator = {
+            time: Time,
+            metrics: Metrics): GroupCoordinator = {
     val offsetConfig = this.offsetConfig(config)
     val groupConfig = GroupConfig(groupMinSessionTimeoutMs = config.groupMinSessionTimeoutMs,
       groupMaxSessionTimeoutMs = config.groupMaxSessionTimeoutMs,
@@ -1012,10 +1230,34 @@ object GroupCoordinator {
       groupInitialRebalanceDelayMs = config.groupInitialRebalanceDelay)
 
     val groupMetadataManager = new GroupMetadataManager(config.brokerId, config.interBrokerProtocolVersion,
-      offsetConfig, replicaManager, zkClient, time)
-    new GroupCoordinator(config.brokerId, groupConfig, offsetConfig, groupMetadataManager, heartbeatPurgatory, joinPurgatory, time)
+      offsetConfig, replicaManager, zkClient, time, metrics)
+    new GroupCoordinator(config.brokerId, groupConfig, offsetConfig, groupMetadataManager, heartbeatPurgatory, joinPurgatory, time, metrics)
+  }
+
+  def joinError(memberId: String, error: Errors): JoinGroupResult = {
+    JoinGroupResult(
+      members = List.empty,
+      memberId = memberId,
+      generationId = GroupCoordinator.NoGeneration,
+      subProtocol = GroupCoordinator.NoProtocol,
+      leaderId = GroupCoordinator.NoLeader,
+      error = error)
+  }
+
+  private def memberLeaveError(memberIdentity: MemberIdentity,
+                               error: Errors): LeaveMemberResponse = {
+    LeaveMemberResponse(
+      memberId = memberIdentity.memberId,
+      groupInstanceId = Option(memberIdentity.groupInstanceId),
+      error = error)
   }
 
+  private def leaveError(topLevelError: Errors,
+                         memberResponses: List[LeaveMemberResponse]): LeaveGroupResult = {
+    LeaveGroupResult(
+      topLevelError = topLevelError,
+      memberResponses = memberResponses)
+  }
 }
 
 case class GroupConfig(groupMinSessionTimeoutMs: Int,
@@ -1023,9 +1265,19 @@ case class GroupConfig(groupMinSessionTimeoutMs: Int,
                        groupMaxSize: Int,
                        groupInitialRebalanceDelayMs: Int)
 
-case class JoinGroupResult(members: Map[String, Array[Byte]],
+case class JoinGroupResult(members: List[JoinGroupResponseMember],
                            memberId: String,
                            generationId: Int,
                            subProtocol: String,
                            leaderId: String,
                            error: Errors)
+
+case class SyncGroupResult(memberAssignment: Array[Byte],
+                           error: Errors)
+
+case class LeaveMemberResponse(memberId: String,
+                               groupInstanceId: Option[String],
+                               error: Errors)
+
+case class LeaveGroupResult(topLevelError: Errors,
+                            memberResponses : List[LeaveMemberResponse])
diff --git a/core/src/main/scala/kafka/coordinator/group/GroupMetadata.scala b/core/src/main/scala/kafka/coordinator/group/GroupMetadata.scala
index 5cfd42038b0d0..b8237686fecb5 100644
--- a/core/src/main/scala/kafka/coordinator/group/GroupMetadata.scala
+++ b/core/src/main/scala/kafka/coordinator/group/GroupMetadata.scala
@@ -16,15 +16,21 @@
  */
 package kafka.coordinator.group
 
+import java.nio.ByteBuffer
 import java.util.UUID
 import java.util.concurrent.locks.ReentrantLock
 
 import kafka.common.OffsetAndMetadata
 import kafka.utils.{CoreUtils, Logging, nonthreadsafe}
+import org.apache.kafka.clients.consumer.internals.ConsumerProtocol
 import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.message.JoinGroupResponseData.JoinGroupResponseMember
+import org.apache.kafka.common.protocol.Errors
+import org.apache.kafka.common.protocol.types.SchemaException
 import org.apache.kafka.common.utils.Time
 
 import scala.collection.{Seq, immutable, mutable}
+import scala.collection.JavaConverters._
 
 private[group] sealed trait GroupState
 
@@ -128,9 +134,17 @@ private object GroupMetadata {
     group.protocol = Option(protocol)
     group.leaderId = Option(leaderId)
     group.currentStateTimestamp = currentStateTimestamp
-    members.foreach(group.add(_, null))
+    members.foreach(member => {
+      group.add(member, null)
+      if (member.isStaticMember) {
+        group.addStaticMember(member.groupInstanceId, member.memberId)
+      }
+    })
+    group.subscribedTopics = group.computeSubscribedTopics()
     group
   }
+
+  private val MemberIdDelimiter = "-"
 }
 
 /**
@@ -154,7 +168,7 @@ case class GroupSummary(state: String,
   * being materialized.
   */
 case class CommitRecordMetadataAndOffset(appendedBatchOffset: Option[Long], offsetAndMetadata: OffsetAndMetadata) {
-  def olderThan(that: CommitRecordMetadataAndOffset) : Boolean = appendedBatchOffset.get < that.appendedBatchOffset.get
+  def olderThan(that: CommitRecordMetadataAndOffset): Boolean = appendedBatchOffset.get < that.appendedBatchOffset.get
 }
 
 /**
@@ -184,6 +198,8 @@ private[group] class GroupMetadata(val groupId: String, initialState: GroupState
   private var protocol: Option[String] = None
 
   private val members = new mutable.HashMap[String, MemberMetadata]
+  // Static membership mapping [key: group.instance.id, value: member.id]
+  private val staticMembers = new mutable.HashMap[String, String]
   private val pendingMembers = new mutable.HashSet[String]
   private var numMembersAwaitingJoin = 0
   private val supportedProtocols = new mutable.HashMap[String, Integer]().withDefaultValue(0)
@@ -193,6 +209,10 @@ private[group] class GroupMetadata(val groupId: String, initialState: GroupState
   private var receivedTransactionalOffsetCommits = false
   private var receivedConsumerOffsetCommits = false
 
+  // When protocolType == `consumer`, a set of subscribed topics is maintained. The set is
+  // computed when a new generation is created or when the group is restored from the log.
+  private var subscribedTopics: Option[Set[String]] = None
+
   var newMemberAdded: Boolean = false
 
   def inLock[T](fun: => T): T = CoreUtils.inLock(lock)(fun)
@@ -208,7 +228,9 @@ private[group] class GroupMetadata(val groupId: String, initialState: GroupState
   def protocolOrNull: String = protocol.orNull
   def currentStateTimestampOrDefault: Long = currentStateTimestamp.getOrElse(-1)
 
-  def add(member: MemberMetadata, callback: JoinCallback = null) {
+  def isConsumerGroup: Boolean = protocolType.contains(ConsumerProtocol.PROTOCOL_TYPE)
+
+  def add(member: MemberMetadata, callback: JoinCallback = null): Unit = {
     if (members.isEmpty)
       this.protocolType = Some(member.protocolType)
 
@@ -225,7 +247,7 @@ private[group] class GroupMetadata(val groupId: String, initialState: GroupState
       numMembersAwaitingJoin += 1
   }
 
-  def remove(memberId: String) {
+  def remove(memberId: String): Unit = {
     members.remove(memberId).foreach { member =>
       member.supportedProtocols.foreach{ case (protocol, _) => supportedProtocols(protocol) -= 1 }
       if (member.isAwaitingJoin)
@@ -236,12 +258,100 @@ private[group] class GroupMetadata(val groupId: String, initialState: GroupState
       leaderId = members.keys.headOption
   }
 
+  /**
+    * Check whether current leader is rejoined. If not, try to find another joined member to be
+    * new leader. Return false if
+    *   1. the group is currently empty (has no designated leader)
+    *   2. no member rejoined
+    */
+  def maybeElectNewJoinedLeader(): Boolean = {
+    leaderId.exists { currentLeaderId =>
+      val currentLeader = get(currentLeaderId)
+      if (!currentLeader.isAwaitingJoin) {
+        members.find(_._2.isAwaitingJoin) match {
+          case Some((anyJoinedMemberId, anyJoinedMember)) =>
+            leaderId = Option(anyJoinedMemberId)
+            info(s"Group leader [member.id: ${currentLeader.memberId}, " +
+              s"group.instance.id: ${currentLeader.groupInstanceId}] failed to join " +
+              s"before rebalance timeout, while new leader $anyJoinedMember was elected.")
+            true
+
+          case None =>
+            info(s"Group leader [member.id: ${currentLeader.memberId}, " +
+              s"group.instance.id: ${currentLeader.groupInstanceId}] failed to join " +
+              s"before rebalance timeout, and the group couldn't proceed to next generation" +
+              s"because no member joined.")
+            false
+        }
+      } else {
+        true
+      }
+    }
+  }
+
+  /**
+    * [For static members only]: Replace the old member id with the new one,
+    * keep everything else unchanged and return the updated member.
+    */
+  def replaceGroupInstance(oldMemberId: String,
+                           newMemberId: String,
+                           groupInstanceId: Option[String]): MemberMetadata = {
+    if(groupInstanceId.isEmpty) {
+      throw new IllegalArgumentException(s"unexpected null group.instance.id in replaceGroupInstance")
+    }
+    val oldMember = members.remove(oldMemberId)
+      .getOrElse(throw new IllegalArgumentException(s"Cannot replace non-existing member id $oldMemberId"))
+
+    // Fence potential duplicate member immediately if someone awaits join/sync callback.
+    maybeInvokeJoinCallback(oldMember, JoinGroupResult(
+      members = List.empty,
+      memberId = oldMemberId,
+      generationId = GroupCoordinator.NoGeneration,
+      subProtocol = GroupCoordinator.NoProtocol,
+      leaderId = GroupCoordinator.NoLeader,
+      error = Errors.FENCED_INSTANCE_ID))
+
+    maybeInvokeSyncCallback(oldMember, SyncGroupResult(
+      Array.empty, Errors.FENCED_INSTANCE_ID
+    ))
+
+    oldMember.memberId = newMemberId
+    members.put(newMemberId, oldMember)
+
+    if (isLeader(oldMemberId))
+      leaderId = Some(newMemberId)
+    addStaticMember(groupInstanceId, newMemberId)
+    oldMember
+  }
+
   def isPendingMember(memberId: String): Boolean = pendingMembers.contains(memberId) && !has(memberId)
 
   def addPendingMember(memberId: String) = pendingMembers.add(memberId)
 
   def removePendingMember(memberId: String) = pendingMembers.remove(memberId)
 
+  def hasStaticMember(groupInstanceId: Option[String]) = groupInstanceId.isDefined && staticMembers.contains(groupInstanceId.get)
+
+  def getStaticMemberId(groupInstanceId: Option[String]) = {
+    if(groupInstanceId.isEmpty) {
+      throw new IllegalArgumentException(s"unexpected null group.instance.id in getStaticMemberId")
+    }
+    staticMembers(groupInstanceId.get)
+  }
+
+  def addStaticMember(groupInstanceId: Option[String], newMemberId: String) = {
+    if(groupInstanceId.isEmpty) {
+      throw new IllegalArgumentException(s"unexpected null group.instance.id in addStaticMember")
+    }
+    staticMembers.put(groupInstanceId.get, newMemberId)
+  }
+
+  def removeStaticMember(groupInstanceId: Option[String]) = {
+    if (groupInstanceId.isDefined) {
+      staticMembers.remove(groupInstanceId.get)
+    }
+  }
+
   def currentState = state
 
   def notYetRejoinedMembers = members.values.filter(!_.isAwaitingJoin).toList
@@ -250,6 +360,8 @@ private[group] class GroupMetadata(val groupId: String, initialState: GroupState
 
   def allMembers = members.keySet
 
+  def allStaticMembers = staticMembers.keySet
+
   def numPending = pendingMembers.size
 
   def allMemberMetadata = members.values.toList
@@ -258,11 +370,35 @@ private[group] class GroupMetadata(val groupId: String, initialState: GroupState
     timeout.max(member.rebalanceTimeoutMs)
   }
 
-  def generateMemberIdSuffix = UUID.randomUUID().toString
+  def generateMemberId(clientId: String,
+                       groupInstanceId: Option[String]): String = {
+    groupInstanceId match {
+      case None =>
+        clientId + GroupMetadata.MemberIdDelimiter + UUID.randomUUID().toString
+      case Some(instanceId) =>
+        instanceId + GroupMetadata.MemberIdDelimiter + UUID.randomUUID().toString
+    }
+  }
+
+  /**
+    * Verify the member.id is up to date for static members. Return true if both conditions met:
+    *   1. given member is a known static member to group
+    *   2. group stored member.id doesn't match with given member.id
+    */
+  def isStaticMemberFenced(memberId: String,
+                           groupInstanceId: Option[String]): Boolean = {
+    if (hasStaticMember(groupInstanceId)
+      && getStaticMemberId(groupInstanceId) != memberId) {
+        error(s"given member.id $memberId is identified as a known static member ${groupInstanceId.get}," +
+          s"but not matching the expected member.id ${getStaticMemberId(groupInstanceId)}")
+        true
+    } else
+      false
+  }
 
   def canRebalance = GroupMetadata.validPreviousStates(PreparingRebalance).contains(state)
 
-  def transitionTo(groupState: GroupState) {
+  def transitionTo(groupState: GroupState): Unit = {
     assertValidTransition(groupState)
     state = groupState
     currentStateTimestamp = Some(time.milliseconds())
@@ -298,6 +434,54 @@ private[group] class GroupMetadata(val groupId: String, initialState: GroupState
       protocolType.contains(memberProtocolType) && memberProtocols.exists(supportedProtocols(_) == members.size)
   }
 
+  def getSubscribedTopics: Option[Set[String]] = subscribedTopics
+
+  /**
+   * Returns true if the consumer group is actively subscribed to the topic. When the consumer
+   * group does not know, because the information is not available yet or because the it has
+   * failed to parse the Consumer Protocol, it returns true to be safe.
+   */
+  def isSubscribedToTopic(topic: String): Boolean = subscribedTopics match {
+    case Some(topics) => topics.contains(topic)
+    case None => true
+  }
+
+  /**
+   * Collects the set of topics that the members are subscribed to when the Protocol Type is equal
+   * to 'consumer'. None is returned if
+   * - the protocol type is not equal to 'consumer';
+   * - the protocol is not defined yet; or
+   * - the protocol metadata does not comply with the schema.
+   */
+  private[group] def computeSubscribedTopics(): Option[Set[String]] = {
+    protocolType match {
+      case Some(ConsumerProtocol.PROTOCOL_TYPE) if members.nonEmpty && protocol.isDefined =>
+        try {
+          Some(
+            members.map { case (_, member) =>
+              // The consumer protocol is parsed with V0 which is the based prefix of all versions.
+              // This way the consumer group manager does not depend on any specific existing or
+              // future versions of the consumer protocol. VO must prefix all new versions.
+              val buffer = ByteBuffer.wrap(member.metadata(protocol.get))
+              ConsumerProtocol.deserializeVersion(buffer)
+              ConsumerProtocol.deserializeSubscriptionV0(buffer).topics.asScala.toSet
+            }.reduceLeft(_ ++ _)
+          )
+        } catch {
+          case e: SchemaException => {
+            warn(s"Failed to parse Consumer Protocol ${ConsumerProtocol.PROTOCOL_TYPE}:${protocol.get} " +
+              s"of group $groupId. Consumer group coordinator is not aware of the subscribed topics.", e)
+            None
+          }
+        }
+
+      case Some(ConsumerProtocol.PROTOCOL_TYPE) if members.isEmpty =>
+        Option(Set.empty)
+
+      case _ => None
+    }
+  }
+
   def updateMember(member: MemberMetadata,
                    protocols: List[(String, Array[Byte])],
                    callback: JoinCallback) = {
@@ -314,7 +498,7 @@ private[group] class GroupMetadata(val groupId: String, initialState: GroupState
   }
 
   def maybeInvokeJoinCallback(member: MemberMetadata,
-                              joinGroupResult: JoinGroupResult) : Unit = {
+                              joinGroupResult: JoinGroupResult): Unit = {
     if (member.isAwaitingJoin) {
       member.awaitingJoinCallback(joinGroupResult)
       member.awaitingJoinCallback = null
@@ -322,25 +506,44 @@ private[group] class GroupMetadata(val groupId: String, initialState: GroupState
     }
   }
 
+  /**
+    * @return true if a sync callback actually performs.
+    */
+  def maybeInvokeSyncCallback(member: MemberMetadata,
+                              syncGroupResult: SyncGroupResult): Boolean = {
+    if (member.isAwaitingSync) {
+      member.awaitingSyncCallback(syncGroupResult)
+      member.awaitingSyncCallback = null
+      true
+    } else {
+      false
+    }
+  }
+
   def initNextGeneration() = {
-    assert(notYetRejoinedMembers == List.empty[MemberMetadata])
     if (members.nonEmpty) {
       generationId += 1
       protocol = Some(selectProtocol)
+      subscribedTopics = computeSubscribedTopics()
       transitionTo(CompletingRebalance)
     } else {
       generationId += 1
       protocol = None
+      subscribedTopics = computeSubscribedTopics()
       transitionTo(Empty)
     }
     receivedConsumerOffsetCommits = false
     receivedTransactionalOffsetCommits = false
   }
 
-  def currentMemberMetadata: Map[String, Array[Byte]] = {
+  def currentMemberMetadata: List[JoinGroupResponseMember] = {
     if (is(Dead) || is(PreparingRebalance))
       throw new IllegalStateException("Cannot obtain member metadata for group in state %s".format(state))
-    members.map{ case (memberId, memberMetadata) => (memberId, memberMetadata.metadata(protocol.get))}.toMap
+    members.map{ case (memberId, memberMetadata) => new JoinGroupResponseMember()
+        .setMemberId(memberId)
+        .setGroupInstanceId(memberMetadata.groupInstanceId.orNull)
+        .setMetadata(memberMetadata.metadata(protocol.get))
+    }.toList
   }
 
   def summary: GroupSummary = {
@@ -362,12 +565,12 @@ private[group] class GroupMetadata(val groupId: String, initialState: GroupState
   }
 
   def initializeOffsets(offsets: collection.Map[TopicPartition, CommitRecordMetadataAndOffset],
-                        pendingTxnOffsets: Map[Long, mutable.Map[TopicPartition, CommitRecordMetadataAndOffset]]) {
+                        pendingTxnOffsets: Map[Long, mutable.Map[TopicPartition, CommitRecordMetadataAndOffset]]): Unit = {
     this.offsets ++= offsets
     this.pendingTransactionalOffsetCommits ++= pendingTxnOffsets
   }
 
-  def onOffsetCommitAppend(topicPartition: TopicPartition, offsetWithCommitRecordMetadata: CommitRecordMetadataAndOffset) {
+  def onOffsetCommitAppend(topicPartition: TopicPartition, offsetWithCommitRecordMetadata: CommitRecordMetadataAndOffset): Unit = {
     if (pendingOffsetCommits.contains(topicPartition)) {
       if (offsetWithCommitRecordMetadata.appendedBatchOffset.isEmpty)
         throw new IllegalStateException("Cannot complete offset commit write without providing the metadata of the record " +
@@ -392,12 +595,12 @@ private[group] class GroupMetadata(val groupId: String, initialState: GroupState
     }
   }
 
-  def prepareOffsetCommit(offsets: Map[TopicPartition, OffsetAndMetadata]) {
+  def prepareOffsetCommit(offsets: Map[TopicPartition, OffsetAndMetadata]): Unit = {
     receivedConsumerOffsetCommits = true
     pendingOffsetCommits ++= offsets
   }
 
-  def prepareTxnOffsetCommit(producerId: Long, offsets: Map[TopicPartition, OffsetAndMetadata]) {
+  def prepareTxnOffsetCommit(producerId: Long, offsets: Map[TopicPartition, OffsetAndMetadata]): Unit = {
     trace(s"TxnOffsetCommit for producer $producerId and group $groupId with offsets $offsets is pending")
     receivedTransactionalOffsetCommits = true
     val producerOffsets = pendingTransactionalOffsetCommits.getOrElseUpdate(producerId,
@@ -429,7 +632,7 @@ private[group] class GroupMetadata(val groupId: String, initialState: GroupState
   }
 
   def onTxnOffsetCommitAppend(producerId: Long, topicPartition: TopicPartition,
-                              commitRecordMetadataAndOffset: CommitRecordMetadataAndOffset) {
+                              commitRecordMetadataAndOffset: CommitRecordMetadataAndOffset): Unit = {
     pendingTransactionalOffsetCommits.get(producerId) match {
       case Some(pendingOffset) =>
         if (pendingOffset.contains(topicPartition)
@@ -486,11 +689,13 @@ private[group] class GroupMetadata(val groupId: String, initialState: GroupState
     }.toMap
   }
 
-  def removeExpiredOffsets(currentTimestamp: Long, offsetRetentionMs: Long) : Map[TopicPartition, OffsetAndMetadata] = {
+  def removeExpiredOffsets(currentTimestamp: Long, offsetRetentionMs: Long): Map[TopicPartition, OffsetAndMetadata] = {
 
-    def getExpiredOffsets(baseTimestamp: CommitRecordMetadataAndOffset => Long): Map[TopicPartition, OffsetAndMetadata] = {
+    def getExpiredOffsets(baseTimestamp: CommitRecordMetadataAndOffset => Long,
+                          subscribedTopics: Set[String] = Set.empty): Map[TopicPartition, OffsetAndMetadata] = {
       offsets.filter {
         case (topicPartition, commitRecordMetadataAndOffset) =>
+          !subscribedTopics.contains(topicPartition.topic()) &&
           !pendingOffsetCommits.contains(topicPartition) && {
             commitRecordMetadataAndOffset.offsetAndMetadata.expireTimestamp match {
               case None =>
@@ -514,8 +719,20 @@ private[group] class GroupMetadata(val groupId: String, initialState: GroupState
         //   expire all offsets with no pending offset commit;
         // - if there is no current state timestamp (old group metadata schema) and retention period has passed
         //   since the last commit timestamp, expire the offset
-        getExpiredOffsets(commitRecordMetadataAndOffset =>
-          currentStateTimestamp.getOrElse(commitRecordMetadataAndOffset.offsetAndMetadata.commitTimestamp))
+        getExpiredOffsets(
+          commitRecordMetadataAndOffset => currentStateTimestamp
+            .getOrElse(commitRecordMetadataAndOffset.offsetAndMetadata.commitTimestamp)
+        )
+
+      case Some(ConsumerProtocol.PROTOCOL_TYPE) if subscribedTopics.isDefined =>
+        // consumers exist in the group =>
+        // - if the group is aware of the subscribed topics and retention period had passed since the
+        //   the last commit timestamp, expire the offset. offset with pending offset commit are not
+        //   expired
+        getExpiredOffsets(
+          _.offsetAndMetadata.commitTimestamp,
+          subscribedTopics.get
+        )
 
       case None =>
         // protocolType is None => standalone (simple) consumer, that uses Kafka for offset storage only
@@ -546,7 +763,7 @@ private[group] class GroupMetadata(val groupId: String, initialState: GroupState
 
   def hasOffsets = offsets.nonEmpty || pendingOffsetCommits.nonEmpty || pendingTransactionalOffsetCommits.nonEmpty
 
-  private def assertValidTransition(targetState: GroupState) {
+  private def assertValidTransition(targetState: GroupState): Unit = {
     if (!GroupMetadata.validPreviousStates(targetState).contains(state))
       throw new IllegalStateException("Group %s should be in the %s states before moving to %s state. Instead it is in %s state"
         .format(groupId, GroupMetadata.validPreviousStates(targetState).mkString(","), targetState, state))
diff --git a/core/src/main/scala/kafka/coordinator/group/GroupMetadataManager.scala b/core/src/main/scala/kafka/coordinator/group/GroupMetadataManager.scala
index 7b244987d7634..ce042691fdb80 100644
--- a/core/src/main/scala/kafka/coordinator/group/GroupMetadataManager.scala
+++ b/core/src/main/scala/kafka/coordinator/group/GroupMetadataManager.scala
@@ -26,34 +26,39 @@ import java.util.concurrent.atomic.AtomicBoolean
 import java.util.concurrent.locks.ReentrantLock
 
 import com.yammer.metrics.core.Gauge
-import kafka.api.{ApiVersion, KAFKA_0_10_1_IV0, KAFKA_2_1_IV0, KAFKA_2_1_IV1}
+import kafka.api.{ApiVersion, KAFKA_0_10_1_IV0, KAFKA_2_1_IV0, KAFKA_2_1_IV1, KAFKA_2_3_IV0}
 import kafka.common.{MessageFormatter, OffsetAndMetadata}
+import kafka.log.AppendOrigin
 import kafka.metrics.KafkaMetricsGroup
-import kafka.server.ReplicaManager
+import kafka.server.{FetchLogEnd, ReplicaManager}
 import kafka.utils.CoreUtils.inLock
 import kafka.utils._
 import kafka.zk.KafkaZkClient
 import org.apache.kafka.clients.consumer.ConsumerRecord
-import org.apache.kafka.common.{KafkaException, TopicPartition}
+import org.apache.kafka.clients.consumer.internals.ConsumerProtocol
 import org.apache.kafka.common.internals.Topic
+import org.apache.kafka.common.metrics.Metrics
+import org.apache.kafka.common.metrics.stats.{Avg, Max, Meter}
 import org.apache.kafka.common.protocol.Errors
 import org.apache.kafka.common.protocol.types.Type._
 import org.apache.kafka.common.protocol.types._
 import org.apache.kafka.common.record._
-import org.apache.kafka.common.requests.{OffsetCommitRequest, OffsetFetchResponse}
 import org.apache.kafka.common.requests.ProduceResponse.PartitionResponse
+import org.apache.kafka.common.requests.{OffsetCommitRequest, OffsetFetchResponse}
 import org.apache.kafka.common.utils.{Time, Utils}
+import org.apache.kafka.common.{KafkaException, TopicPartition}
 
 import scala.collection.JavaConverters._
 import scala.collection._
-import scala.collection.mutable.ListBuffer
+import scala.collection.mutable.ArrayBuffer
 
 class GroupMetadataManager(brokerId: Int,
                            interBrokerProtocolVersion: ApiVersion,
                            config: OffsetConfig,
                            replicaManager: ReplicaManager,
                            zkClient: KafkaZkClient,
-                           time: Time) extends Logging with KafkaMetricsGroup {
+                           time: Time,
+                           metrics: Metrics) extends Logging with KafkaMetricsGroup {
 
   private val compressionType: CompressionType = CompressionType.forId(config.offsetsTopicCompressionCodec.codec)
 
@@ -82,6 +87,36 @@ class GroupMetadataManager(brokerId: Int,
    * We use this structure to quickly find the groups which need to be updated by the commit/abort marker. */
   private val openGroupsForProducer = mutable.HashMap[Long, mutable.Set[String]]()
 
+  /* setup metrics*/
+  val partitionLoadSensor = metrics.sensor("PartitionLoadTime")
+
+  partitionLoadSensor.add(metrics.metricName("partition-load-time-max",
+    "group-coordinator-metrics",
+    "The max time it took to load the partitions in the last 30sec"), new Max())
+  partitionLoadSensor.add(metrics.metricName("partition-load-time-avg",
+    "group-coordinator-metrics",
+    "The avg time it took to load the partitions in the last 30sec"), new Avg())
+
+  val offsetCommitsSensor = metrics.sensor("OffsetCommits")
+
+  offsetCommitsSensor.add(new Meter(
+    metrics.metricName("offset-commit-rate",
+      "group-coordinator-metrics",
+      "The rate of committed offsets"),
+    metrics.metricName("offset-commit-count",
+      "group-coordinator-metrics",
+      "The total number of committed offsets")))
+
+  val offsetExpiredSensor = metrics.sensor("OffsetExpired")
+
+  offsetExpiredSensor.add(new Meter(
+    metrics.metricName("offset-expiration-rate",
+      "group-coordinator-metrics",
+      "The rate of expired offsets"),
+    metrics.metricName("offset-expiration-count",
+      "group-coordinator-metrics",
+      "The total number of expired offsets")))
+
   this.logIdent = s"[GroupMetadataManager brokerId=$brokerId] "
 
   private def recreateGauge[T](name: String, gauge: Gauge[T]): Gauge[T] = {
@@ -136,7 +171,7 @@ class GroupMetadataManager(brokerId: Int,
       })
     })
 
-  def startup(enableMetadataExpiration: Boolean) {
+  def startup(enableMetadataExpiration: Boolean): Unit = {
     scheduler.startup()
     if (enableMetadataExpiration) {
       scheduler.schedule(name = "delete-expired-group-metadata",
@@ -217,7 +252,7 @@ class GroupMetadataManager(brokerId: Int,
         val generationId = group.generationId
 
         // set the callback function to insert the created group into cache after log append completed
-        def putCacheCallback(responseStatus: Map[TopicPartition, PartitionResponse]) {
+        def putCacheCallback(responseStatus: Map[TopicPartition, PartitionResponse]): Unit = {
           // the append response should only contain the topics partition
           if (responseStatus.size != 1 || !responseStatus.contains(groupMetadataPartition))
             throw new IllegalStateException("Append status %s should only have one partition %s"
@@ -281,7 +316,7 @@ class GroupMetadataManager(brokerId: Int,
       timeout = config.offsetCommitTimeoutMs.toLong,
       requiredAcks = config.offsetCommitRequiredAcks,
       internalTopicsAllowed = true,
-      isFromClient = false,
+      origin = AppendOrigin.Coordinator,
       entriesPerPartition = records,
       delayedProduceLock = Some(group.lock),
       responseCallback = callback)
@@ -312,7 +347,7 @@ class GroupMetadataManager(brokerId: Int,
     // construct the message set to append
     if (filteredOffsetMetadata.isEmpty) {
       // compute the final error codes for the commit response
-      val commitStatus = offsetMetadata.mapValues(_ => Errors.OFFSET_METADATA_TOO_LARGE)
+      val commitStatus = offsetMetadata.map { case (k, _) => k -> Errors.OFFSET_METADATA_TOO_LARGE }
       responseCallback(commitStatus)
       None
     } else {
@@ -340,12 +375,15 @@ class GroupMetadataManager(brokerId: Int,
           val entries = Map(offsetTopicPartition -> builder.build())
 
           // set the callback function to insert offsets into cache after log append completed
-          def putCacheCallback(responseStatus: Map[TopicPartition, PartitionResponse]) {
+          def putCacheCallback(responseStatus: Map[TopicPartition, PartitionResponse]): Unit = {
             // the append response should only contain the topics partition
             if (responseStatus.size != 1 || !responseStatus.contains(offsetTopicPartition))
               throw new IllegalStateException("Append status %s should only have one partition %s"
                 .format(responseStatus, offsetTopicPartition))
 
+            // record the number of offsets committed to the log
+            offsetCommitsSensor.record(records.size)
+
             // construct the commit response status and insert
             // the offset and metadata to cache if the append status has no error
             val status = responseStatus(offsetTopicPartition)
@@ -484,7 +522,7 @@ class GroupMetadataManager(brokerId: Int,
   /**
    * Asynchronously read the partition from the offsets topic and populate the cache
    */
-  def scheduleLoadGroupAndOffsets(offsetsPartition: Int, onGroupLoaded: GroupMetadata => Unit) {
+  def scheduleLoadGroupAndOffsets(offsetsPartition: Int, onGroupLoaded: GroupMetadata => Unit): Unit = {
     val topicPartition = new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, offsetsPartition)
     if (addLoadingPartition(offsetsPartition)) {
       info(s"Scheduling loading of offsets and group metadata from $topicPartition")
@@ -494,11 +532,14 @@ class GroupMetadataManager(brokerId: Int,
     }
   }
 
-  private[group] def loadGroupsAndOffsets(topicPartition: TopicPartition, onGroupLoaded: GroupMetadata => Unit) {
+  private[group] def loadGroupsAndOffsets(topicPartition: TopicPartition, onGroupLoaded: GroupMetadata => Unit): Unit = {
     try {
       val startMs = time.milliseconds()
       doLoadGroupsAndOffsets(topicPartition, onGroupLoaded)
-      info(s"Finished loading offsets and group metadata from $topicPartition in ${time.milliseconds() - startMs} milliseconds.")
+      val endMs = time.milliseconds()
+      val timeLapse = endMs - startMs
+      partitionLoadSensor.record(timeLapse, endMs, false)
+      info(s"Finished loading offsets and group metadata from $topicPartition in $timeLapse milliseconds.")
     } catch {
       case t: Throwable => error(s"Error loading offsets from $topicPartition", t)
     } finally {
@@ -509,7 +550,7 @@ class GroupMetadataManager(brokerId: Int,
     }
   }
 
-  private def doLoadGroupsAndOffsets(topicPartition: TopicPartition, onGroupLoaded: GroupMetadata => Unit) {
+  private def doLoadGroupsAndOffsets(topicPartition: TopicPartition, onGroupLoaded: GroupMetadata => Unit): Unit = {
     def logEndOffset: Long = replicaManager.getLogEndOffset(topicPartition).getOrElse(-1L)
 
     replicaManager.getLog(topicPartition) match {
@@ -517,20 +558,28 @@ class GroupMetadataManager(brokerId: Int,
         warn(s"Attempted to load offsets and group metadata from $topicPartition, but found no log")
 
       case Some(log) =>
-        var currOffset = log.logStartOffset
-
-        // buffer may not be needed if records are read from memory
-        var buffer = ByteBuffer.allocate(0)
-
-        // loop breaks if leader changes at any time during the load, since getHighWatermark is -1
         val loadedOffsets = mutable.Map[GroupTopicPartition, CommitRecordMetadataAndOffset]()
         val pendingOffsets = mutable.Map[Long, mutable.Map[GroupTopicPartition, CommitRecordMetadataAndOffset]]()
         val loadedGroups = mutable.Map[String, GroupMetadata]()
         val removedGroups = mutable.Set[String]()
 
-        while (currOffset < logEndOffset && !shuttingDown.get()) {
-          val fetchDataInfo = log.read(currOffset, config.loadBufferSize, maxOffset = None,
-            minOneMessage = true, includeAbortedTxns = false)
+        // buffer may not be needed if records are read from memory
+        var buffer = ByteBuffer.allocate(0)
+
+        // loop breaks if leader changes at any time during the load, since logEndOffset is -1
+        var currOffset = log.logStartOffset
+
+        // loop breaks if no records have been read, since the end of the log has been reached
+        var readAtLeastOneRecord = true
+
+        while (currOffset < logEndOffset && readAtLeastOneRecord && !shuttingDown.get()) {
+          val fetchDataInfo = log.read(currOffset,
+            maxLength = config.loadBufferSize,
+            isolation = FetchLogEnd,
+            minOneMessage = true)
+
+          readAtLeastOneRecord = fetchDataInfo.records.sizeInBytes > 0
+
           val memRecords = fetchDataInfo.records match {
             case records: MemoryRecords => records
             case fileRecords: FileRecords =>
@@ -688,12 +737,12 @@ class GroupMetadataManager(brokerId: Int,
    * @param offsetsPartition Groups belonging to this partition of the offsets topic will be deleted from the cache.
    */
   def removeGroupsForPartition(offsetsPartition: Int,
-                               onGroupUnloaded: GroupMetadata => Unit) {
+                               onGroupUnloaded: GroupMetadata => Unit): Unit = {
     val topicPartition = new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, offsetsPartition)
     info(s"Scheduling unloading of offsets and group metadata from $topicPartition")
     scheduler.schedule(topicPartition.toString, () => removeGroupsAndOffsets)
 
-    def removeGroupsAndOffsets() {
+    def removeGroupsAndOffsets(): Unit = {
       var numOffsetsRemoved = 0
       var numGroupsRemoved = 0
 
@@ -724,6 +773,7 @@ class GroupMetadataManager(brokerId: Int,
     val numOffsetsRemoved = cleanupGroupMetadata(groupMetadataCache.values, group => {
       group.removeExpiredOffsets(currentTimestamp, config.offsetsRetentionMs)
     })
+    offsetExpiredSensor.record(numOffsetsRemoved)
     info(s"Removed $numOffsetsRemoved expired offsets in ${time.milliseconds() - currentTimestamp} milliseconds.")
   }
 
@@ -757,7 +807,7 @@ class GroupMetadataManager(brokerId: Int,
         val timestamp = time.milliseconds()
 
           replicaManager.nonOfflinePartition(appendPartition).foreach { partition =>
-            val tombstones = ListBuffer.empty[SimpleRecord]
+            val tombstones = ArrayBuffer.empty[SimpleRecord]
             removedOffsets.foreach { case (topicPartition, offsetAndMetadata) =>
               trace(s"Removing expired/deleted offset and metadata for $groupId, $topicPartition: $offsetAndMetadata")
               val commitKey = GroupMetadataManager.offsetCommitKey(groupId, topicPartition)
@@ -780,8 +830,8 @@ class GroupMetadataManager(brokerId: Int,
               try {
                 // do not need to require acks since even if the tombstone is lost,
                 // it will be appended again in the next purge cycle
-                val records = MemoryRecords.withRecords(magicValue, 0L, compressionType, timestampType, tombstones: _*)
-                partition.appendRecordsToLeader(records, isFromClient = false, requiredAcks = 0)
+                val records = MemoryRecords.withRecords(magicValue, 0L, compressionType, timestampType, tombstones.toArray: _*)
+                partition.appendRecordsToLeader(records, origin = AppendOrigin.Coordinator, requiredAcks = 0)
 
                 offsetsRemoved += removedOffsets.size
                 trace(s"Successfully appended ${tombstones.size} tombstones to $appendPartition for expired/deleted " +
@@ -863,7 +913,7 @@ class GroupMetadataManager(brokerId: Int,
   }
 
 
-  def shutdown() {
+  def shutdown(): Unit = {
     shuttingDown.set(true)
     if (scheduler.isStarted)
       scheduler.shutdown()
@@ -893,7 +943,7 @@ class GroupMetadataManager(brokerId: Int,
    *
    * NOTE: this is for test only
    */
-  private[group] def addPartitionOwnership(partition: Int) {
+  private[group] def addPartitionOwnership(partition: Int): Unit = {
     inLock(partitionLock) {
       ownedPartitions.add(partition)
     }
@@ -977,6 +1027,7 @@ object GroupMetadataManager {
   private val GROUP_KEY_GROUP_FIELD = GROUP_METADATA_KEY_SCHEMA.get("group")
 
   private val MEMBER_ID_KEY = "member_id"
+  private val GROUP_INSTANCE_ID_KEY = "group_instance_id"
   private val CLIENT_ID_KEY = "client_id"
   private val CLIENT_HOST_KEY = "client_host"
   private val REBALANCE_TIMEOUT_KEY = "rebalance_timeout"
@@ -1003,6 +1054,16 @@ object GroupMetadataManager {
 
   private val MEMBER_METADATA_V2 = MEMBER_METADATA_V1
 
+  private val MEMBER_METADATA_V3 = new Schema(
+    new Field(MEMBER_ID_KEY, STRING),
+    new Field(GROUP_INSTANCE_ID_KEY, NULLABLE_STRING),
+    new Field(CLIENT_ID_KEY, STRING),
+    new Field(CLIENT_HOST_KEY, STRING),
+    new Field(REBALANCE_TIMEOUT_KEY, INT32),
+    new Field(SESSION_TIMEOUT_KEY, INT32),
+    new Field(SUBSCRIPTION_KEY, BYTES),
+    new Field(ASSIGNMENT_KEY, BYTES))
+
   private val PROTOCOL_TYPE_KEY = "protocol_type"
   private val GENERATION_KEY = "generation"
   private val PROTOCOL_KEY = "protocol"
@@ -1032,6 +1093,14 @@ object GroupMetadataManager {
     new Field(CURRENT_STATE_TIMESTAMP_KEY, INT64),
     new Field(MEMBERS_KEY, new ArrayOf(MEMBER_METADATA_V2)))
 
+  private val GROUP_METADATA_VALUE_SCHEMA_V3 = new Schema(
+    new Field(PROTOCOL_TYPE_KEY, STRING),
+    new Field(GENERATION_KEY, INT32),
+    new Field(PROTOCOL_KEY, NULLABLE_STRING),
+    new Field(LEADER_KEY, NULLABLE_STRING),
+    new Field(CURRENT_STATE_TIMESTAMP_KEY, INT64),
+    new Field(MEMBERS_KEY, new ArrayOf(MEMBER_METADATA_V3)))
+
   // map of versions to key schemas as data types
   private val MESSAGE_TYPE_SCHEMAS = Map(
     0 -> OFFSET_COMMIT_KEY_SCHEMA,
@@ -1049,7 +1118,8 @@ object GroupMetadataManager {
   private val GROUP_VALUE_SCHEMAS = Map(
     0 -> GROUP_METADATA_VALUE_SCHEMA_V0,
     1 -> GROUP_METADATA_VALUE_SCHEMA_V1,
-    2 -> GROUP_METADATA_VALUE_SCHEMA_V2)
+    2 -> GROUP_METADATA_VALUE_SCHEMA_V2,
+    3 -> GROUP_METADATA_VALUE_SCHEMA_V3)
 
   private val CURRENT_OFFSET_KEY_SCHEMA = schemaForKey(CURRENT_OFFSET_KEY_SCHEMA_VERSION)
   private val CURRENT_GROUP_KEY_SCHEMA = schemaForKey(CURRENT_GROUP_KEY_SCHEMA_VERSION)
@@ -1058,7 +1128,7 @@ object GroupMetadataManager {
     val schemaOpt = MESSAGE_TYPE_SCHEMAS.get(version)
     schemaOpt match {
       case Some(schema) => schema
-      case _ => throw new KafkaException("Unknown offset schema version " + version)
+      case _ => throw new KafkaException("Unknown message key schema version " + version)
     }
   }
 
@@ -1083,8 +1153,7 @@ object GroupMetadataManager {
    *
    * @return key for offset commit message
    */
-  private[group] def offsetCommitKey(group: String,
-                                     topicPartition: TopicPartition): Array[Byte] = {
+  def offsetCommitKey(group: String, topicPartition: TopicPartition): Array[Byte] = {
     val key = new Struct(CURRENT_OFFSET_KEY_SCHEMA)
     key.set(OFFSET_KEY_GROUP_FIELD, group)
     key.set(OFFSET_KEY_TOPIC_FIELD, topicPartition.topic)
@@ -1101,7 +1170,7 @@ object GroupMetadataManager {
    *
    * @return key bytes for group metadata message
    */
-  private[group] def groupMetadataKey(group: String): Array[Byte] = {
+  def groupMetadataKey(group: String): Array[Byte] = {
     val key = new Struct(CURRENT_GROUP_KEY_SCHEMA)
     key.set(GROUP_KEY_GROUP_FIELD, group)
 
@@ -1118,8 +1187,8 @@ object GroupMetadataManager {
    * @param apiVersion the api version
    * @return payload for offset commit message
    */
-  private[group] def offsetCommitValue(offsetAndMetadata: OffsetAndMetadata,
-                                       apiVersion: ApiVersion): Array[Byte] = {
+  def offsetCommitValue(offsetAndMetadata: OffsetAndMetadata,
+                        apiVersion: ApiVersion): Array[Byte] = {
     // generate commit value according to schema version
     val (version, value) = {
       if (apiVersion < KAFKA_2_1_IV0 || offsetAndMetadata.expireTimestamp.nonEmpty) {
@@ -1163,17 +1232,19 @@ object GroupMetadataManager {
    * @param apiVersion the api version
    * @return payload for offset commit message
    */
-  private[group] def groupMetadataValue(groupMetadata: GroupMetadata,
-                                        assignment: Map[String, Array[Byte]],
-                                        apiVersion: ApiVersion): Array[Byte] = {
+  def groupMetadataValue(groupMetadata: GroupMetadata,
+                         assignment: Map[String, Array[Byte]],
+                         apiVersion: ApiVersion): Array[Byte] = {
 
     val (version, value) = {
       if (apiVersion < KAFKA_0_10_1_IV0)
         (0.toShort, new Struct(GROUP_METADATA_VALUE_SCHEMA_V0))
       else if (apiVersion < KAFKA_2_1_IV0)
         (1.toShort, new Struct(GROUP_METADATA_VALUE_SCHEMA_V1))
-      else
+      else if (apiVersion < KAFKA_2_3_IV0)
         (2.toShort, new Struct(GROUP_METADATA_VALUE_SCHEMA_V2))
+      else
+        (3.toShort, new Struct(GROUP_METADATA_VALUE_SCHEMA_V3))
     }
 
     value.set(PROTOCOL_TYPE_KEY, groupMetadata.protocolType.getOrElse(""))
@@ -1194,6 +1265,9 @@ object GroupMetadataManager {
       if (version > 0)
         memberStruct.set(REBALANCE_TIMEOUT_KEY, memberMetadata.rebalanceTimeoutMs)
 
+      if (version >= 3)
+        memberStruct.set(GROUP_INSTANCE_ID_KEY, memberMetadata.groupInstanceId.orNull)
+
       // The group is non-empty, so the current protocol must be defined
       val protocol = groupMetadata.protocolOrNull
       if (protocol == null)
@@ -1312,7 +1386,7 @@ object GroupMetadataManager {
       val valueSchema = schemaForGroupValue(version)
       val value = valueSchema.read(buffer)
 
-      if (version >= 0 && version <= 2) {
+      if (version >= 0 && version <= 3) {
         val generationId = value.get(GENERATION_KEY).asInstanceOf[Int]
         val protocolType = value.get(PROTOCOL_TYPE_KEY).asInstanceOf[String]
         val protocol = value.get(PROTOCOL_KEY).asInstanceOf[String]
@@ -1333,13 +1407,18 @@ object GroupMetadataManager {
         val members = memberMetadataArray.map { memberMetadataObj =>
           val memberMetadata = memberMetadataObj.asInstanceOf[Struct]
           val memberId = memberMetadata.get(MEMBER_ID_KEY).asInstanceOf[String]
+          val groupInstanceId =
+            if (version >= 3)
+              Some(memberMetadata.get(GROUP_INSTANCE_ID_KEY).asInstanceOf[String])
+            else
+              None
           val clientId = memberMetadata.get(CLIENT_ID_KEY).asInstanceOf[String]
           val clientHost = memberMetadata.get(CLIENT_HOST_KEY).asInstanceOf[String]
           val sessionTimeout = memberMetadata.get(SESSION_TIMEOUT_KEY).asInstanceOf[Int]
           val rebalanceTimeout = if (version == 0) sessionTimeout else memberMetadata.get(REBALANCE_TIMEOUT_KEY).asInstanceOf[Int]
           val subscription = Utils.toArray(memberMetadata.get(SUBSCRIPTION_KEY).asInstanceOf[ByteBuffer])
 
-          val member = new MemberMetadata(memberId, groupId, clientId, clientHost, rebalanceTimeout, sessionTimeout,
+          val member = new MemberMetadata(memberId, groupId, groupInstanceId, clientId, clientHost, rebalanceTimeout, sessionTimeout,
             protocolType, List((protocol, subscription)))
           member.assignment = Utils.toArray(memberMetadata.get(ASSIGNMENT_KEY).asInstanceOf[ByteBuffer])
           member
@@ -1354,7 +1433,7 @@ object GroupMetadataManager {
   // Formatter for use with tools such as console consumer: Consumer should also set exclude.internal.topics to false.
   // (specify --formatter "kafka.coordinator.group.GroupMetadataManager\$OffsetsMessageFormatter" when consuming __consumer_offsets)
   class OffsetsMessageFormatter extends MessageFormatter {
-    def writeTo(consumerRecord: ConsumerRecord[Array[Byte], Array[Byte]], output: PrintStream) {
+    def writeTo(consumerRecord: ConsumerRecord[Array[Byte], Array[Byte]], output: PrintStream): Unit = {
       Option(consumerRecord.key).map(key => GroupMetadataManager.readMessageKey(ByteBuffer.wrap(key))).foreach {
         // Only print if the message is an offset record.
         // We ignore the timestamp of the message because GroupMetadataMessage has its own timestamp.
@@ -1375,7 +1454,7 @@ object GroupMetadataManager {
 
   // Formatter for use with tools to read group metadata history
   class GroupMetadataMessageFormatter extends MessageFormatter {
-    def writeTo(consumerRecord: ConsumerRecord[Array[Byte], Array[Byte]], output: PrintStream) {
+    def writeTo(consumerRecord: ConsumerRecord[Array[Byte], Array[Byte]], output: PrintStream): Unit = {
       Option(consumerRecord.key).map(key => GroupMetadataManager.readMessageKey(ByteBuffer.wrap(key))).foreach {
         // Only print if the message is a group metadata record.
         // We ignore the timestamp of the message because GroupMetadataMessage has its own timestamp.
@@ -1394,6 +1473,83 @@ object GroupMetadataManager {
     }
   }
 
+  /**
+   * Exposed for printing records using [[kafka.tools.DumpLogSegments]]
+   */
+  def formatRecordKeyAndValue(record: Record): (Option[String], Option[String]) = {
+    if (!record.hasKey) {
+      throw new KafkaException("Failed to decode message using offset topic decoder (message had a missing key)")
+    } else {
+      GroupMetadataManager.readMessageKey(record.key) match {
+        case offsetKey: OffsetKey => parseOffsets(offsetKey, record.value)
+        case groupMetadataKey: GroupMetadataKey => parseGroupMetadata(groupMetadataKey, record.value)
+        case _ => throw new KafkaException("Failed to decode message using offset topic decoder (message had an invalid key)")
+      }
+    }
+  }
+
+  private def parseOffsets(offsetKey: OffsetKey, payload: ByteBuffer): (Option[String], Option[String]) = {
+    val groupId = offsetKey.key.group
+    val topicPartition = offsetKey.key.topicPartition
+    val keyString = s"offset_commit::group=$groupId,partition=$topicPartition"
+
+    val offset = GroupMetadataManager.readOffsetMessageValue(payload)
+    val valueString = if (offset == null) {
+      ""
+    } else {
+      if (offset.metadata.isEmpty)
+        s"offset=${offset.offset}"
+      else
+        s"offset=${offset.offset},metadata=${offset.metadata}"
+    }
+
+    (Some(keyString), Some(valueString))
+  }
+
+  private def parseGroupMetadata(groupMetadataKey: GroupMetadataKey, payload: ByteBuffer): (Option[String], Option[String]) = {
+    val groupId = groupMetadataKey.key
+    val keyString = s"group_metadata::group=$groupId"
+
+    val group = GroupMetadataManager.readGroupMessageValue(groupId, payload, Time.SYSTEM)
+    val valueString = if (group == null)
+      ""
+    else {
+      val protocolType = group.protocolType.getOrElse("")
+
+      val assignment = group.allMemberMetadata.map { member =>
+        if (protocolType == ConsumerProtocol.PROTOCOL_TYPE) {
+          val partitionAssignment = ConsumerProtocol.deserializeAssignment(ByteBuffer.wrap(member.assignment))
+          val userData = Option(partitionAssignment.userData)
+            .map(Utils.toArray)
+            .map(hex)
+            .getOrElse("")
+
+          if (userData.isEmpty)
+            s"${member.memberId}=${partitionAssignment.partitions()}"
+          else
+            s"${member.memberId}=${partitionAssignment.partitions()}:$userData"
+        } else {
+          s"${member.memberId}=${hex(member.assignment)}"
+        }
+      }.mkString("{", ",", "}")
+
+      Json.encodeAsString(Map(
+        "protocolType" -> protocolType,
+        "protocol" -> group.protocolOrNull,
+        "generationId" -> group.generationId,
+        "assignment" -> assignment
+      ).asJava)
+    }
+    (Some(keyString), Some(valueString))
+  }
+
+  private def hex(bytes: Array[Byte]): String = {
+    if (bytes.isEmpty)
+      ""
+    else
+      "%X".format(BigInt(1, bytes))
+  }
+
 }
 
 case class GroupTopicPartition(group: String, topicPartition: TopicPartition) {
diff --git a/core/src/main/scala/kafka/coordinator/group/MemberMetadata.scala b/core/src/main/scala/kafka/coordinator/group/MemberMetadata.scala
index 1932f423bac92..e9ec23e519234 100644
--- a/core/src/main/scala/kafka/coordinator/group/MemberMetadata.scala
+++ b/core/src/main/scala/kafka/coordinator/group/MemberMetadata.scala
@@ -20,10 +20,9 @@ package kafka.coordinator.group
 import java.util
 
 import kafka.utils.nonthreadsafe
-import org.apache.kafka.common.protocol.Errors
-
 
 case class MemberSummary(memberId: String,
+                         groupInstanceId: Option[String],
                          clientId: String,
                          clientHost: String,
                          metadata: Array[Byte],
@@ -54,8 +53,9 @@ private object MemberMetadata {
  *                            and the group transitions to stable
  */
 @nonthreadsafe
-private[group] class MemberMetadata(val memberId: String,
+private[group] class MemberMetadata(var memberId: String,
                                     val groupId: String,
+                                    val groupInstanceId: Option[String],
                                     val clientId: String,
                                     val clientHost: String,
                                     val rebalanceTimeoutMs: Int,
@@ -65,12 +65,20 @@ private[group] class MemberMetadata(val memberId: String,
 
   var assignment: Array[Byte] = Array.empty[Byte]
   var awaitingJoinCallback: JoinGroupResult => Unit = null
-  var awaitingSyncCallback: (Array[Byte], Errors) => Unit = null
-  var latestHeartbeat: Long = -1
+  var awaitingSyncCallback: SyncGroupResult => Unit = null
   var isLeaving: Boolean = false
   var isNew: Boolean = false
+  val isStaticMember: Boolean = groupInstanceId.isDefined
+
+  // This variable is used to track heartbeat completion through the delayed
+  // heartbeat purgatory. When scheduling a new heartbeat expiration, we set
+  // this value to `false`. Upon receiving the heartbeat (or any other event
+  // indicating the liveness of the client), we set it to `true` so that the
+  // delayed heartbeat can be completed.
+  var heartbeatSatisfied: Boolean = false
 
   def isAwaitingJoin = awaitingJoinCallback != null
+  def isAwaitingSync = awaitingSyncCallback != null
 
   /**
    * Get metadata corresponding to the provided protocol.
@@ -83,11 +91,17 @@ private[group] class MemberMetadata(val memberId: String,
     }
   }
 
-  def shouldKeepAlive(deadlineMs: Long): Boolean = {
-    if (isAwaitingJoin)
-      !isNew || latestHeartbeat + GroupCoordinator.NewMemberJoinTimeoutMs > deadlineMs
-    else awaitingSyncCallback != null ||
-      latestHeartbeat + sessionTimeoutMs > deadlineMs
+  def hasSatisfiedHeartbeat: Boolean = {
+    if (isNew) {
+      // New members can be expired while awaiting join, so we have to check this first
+      heartbeatSatisfied
+    } else if (isAwaitingJoin || isAwaitingSync) {
+      // Members that are awaiting a rebalance automatically satisfy expected heartbeats
+      true
+    } else {
+      // Otherwise we require the next heartbeat
+      heartbeatSatisfied
+    }
   }
 
   /**
@@ -107,11 +121,11 @@ private[group] class MemberMetadata(val memberId: String,
   }
 
   def summary(protocol: String): MemberSummary = {
-    MemberSummary(memberId, clientId, clientHost, metadata(protocol), assignment)
+    MemberSummary(memberId, groupInstanceId, clientId, clientHost, metadata(protocol), assignment)
   }
 
   def summaryNoMetadata(): MemberSummary = {
-    MemberSummary(memberId, clientId, clientHost, Array.empty[Byte], Array.empty[Byte])
+    MemberSummary(memberId, groupInstanceId, clientId, clientHost, Array.empty[Byte], Array.empty[Byte])
   }
 
   /**
@@ -129,6 +143,7 @@ private[group] class MemberMetadata(val memberId: String,
   override def toString: String = {
     "MemberMetadata(" +
       s"memberId=$memberId, " +
+      s"groupInstanceId=$groupInstanceId, " +
       s"clientId=$clientId, " +
       s"clientHost=$clientHost, " +
       s"sessionTimeoutMs=$sessionTimeoutMs, " +
@@ -136,5 +151,4 @@ private[group] class MemberMetadata(val memberId: String,
       s"supportedProtocols=${supportedProtocols.map(_._1)}, " +
       ")"
   }
-
 }
diff --git a/core/src/main/scala/kafka/coordinator/group/OffsetConfig.scala b/core/src/main/scala/kafka/coordinator/group/OffsetConfig.scala
index 55ec590852cd0..f4a22d4da6ddf 100644
--- a/core/src/main/scala/kafka/coordinator/group/OffsetConfig.scala
+++ b/core/src/main/scala/kafka/coordinator/group/OffsetConfig.scala
@@ -36,6 +36,9 @@ import kafka.message.{CompressionCodec, NoCompressionCodec}
  *                              commit or this timeout is reached. (Similar to the producer request timeout.)
  * @param offsetCommitRequiredAcks The required acks before the commit can be accepted. In general, the default (-1)
  *                                 should not be overridden.
+ * @param offsetsTopicMaxMessageBytes The maximum record batch size for the offset commit topic
+ * @param offsetsTopicMinInSyncReplicas The minimum number of replicas that must acknowledged a write for the write to be considered successful
+ * @param offsetsTopicMinCompactionLagMs The minimum time a message will stay un-compacted in the log
  */
 case class OffsetConfig(maxMetadataSize: Int = OffsetConfig.DefaultMaxMetadataSize,
                         loadBufferSize: Int = OffsetConfig.DefaultLoadBufferSize,
@@ -46,7 +49,10 @@ case class OffsetConfig(maxMetadataSize: Int = OffsetConfig.DefaultMaxMetadataSi
                         offsetsTopicReplicationFactor: Short = OffsetConfig.DefaultOffsetsTopicReplicationFactor,
                         offsetsTopicCompressionCodec: CompressionCodec = OffsetConfig.DefaultOffsetsTopicCompressionCodec,
                         offsetCommitTimeoutMs: Int = OffsetConfig.DefaultOffsetCommitTimeoutMs,
-                        offsetCommitRequiredAcks: Short = OffsetConfig.DefaultOffsetCommitRequiredAcks)
+                        offsetCommitRequiredAcks: Short = OffsetConfig.DefaultOffsetCommitRequiredAcks,
+                        offsetsTopicMaxMessageBytes: Int = OffsetConfig.DefaultOffsetsTopicMaxMessageBytes,
+                        offsetsTopicMinInSyncReplicas: Int = OffsetConfig.DefaultOffsetsTopicMinInSyncReplicas,
+                        offsetsTopicMinCompactionLagMs: Long = OffsetConfig.DefaultOffsetsTopicMinCompactionLagMs)
 
 object OffsetConfig {
   val DefaultMaxMetadataSize = 4096
@@ -59,4 +65,7 @@ object OffsetConfig {
   val DefaultOffsetsTopicCompressionCodec = NoCompressionCodec
   val DefaultOffsetCommitTimeoutMs = 5000
   val DefaultOffsetCommitRequiredAcks = (-1).toShort
+  val DefaultOffsetsTopicMaxMessageBytes = 20 * 1024 * 1024
+  val DefaultOffsetsTopicMinInSyncReplicas = 1
+  val DefaultOffsetsTopicMinCompactionLagMs = 0L
 }
\ No newline at end of file
diff --git a/core/src/main/scala/kafka/coordinator/transaction/ProducerIdManager.scala b/core/src/main/scala/kafka/coordinator/transaction/ProducerIdManager.scala
index 5c22c8e2cb23a..b12dbf8124f91 100644
--- a/core/src/main/scala/kafka/coordinator/transaction/ProducerIdManager.scala
+++ b/core/src/main/scala/kafka/coordinator/transaction/ProducerIdManager.scala
@@ -150,7 +150,7 @@ class ProducerIdManager(val brokerId: Int, val zkClient: KafkaZkClient) extends
     }
   }
 
-  def shutdown() {
+  def shutdown(): Unit = {
     info(s"Shutdown complete: last producerId assigned $nextProducerId")
   }
 }
diff --git a/core/src/main/scala/kafka/coordinator/transaction/TransactionCoordinator.scala b/core/src/main/scala/kafka/coordinator/transaction/TransactionCoordinator.scala
index 9d4eed69fd371..c7601ef47f32c 100644
--- a/core/src/main/scala/kafka/coordinator/transaction/TransactionCoordinator.scala
+++ b/core/src/main/scala/kafka/coordinator/transaction/TransactionCoordinator.scala
@@ -55,7 +55,7 @@ object TransactionCoordinator {
     // we do not need to turn on reaper thread since no tasks will be expired and there are no completed tasks to be purged
     val txnMarkerPurgatory = DelayedOperationPurgatory[DelayedTxnMarker]("txn-marker-purgatory", config.brokerId,
       reaperEnabled = false, timerEnabled = false)
-    val txnStateManager = new TransactionStateManager(config.brokerId, zkClient, scheduler, replicaManager, txnConfig, time)
+    val txnStateManager = new TransactionStateManager(config.brokerId, zkClient, scheduler, replicaManager, txnConfig, time, metrics)
 
     val logContext = new LogContext(s"[TransactionCoordinator id=${config.brokerId}] ")
     val txnMarkerChannelManager = TransactionMarkerChannelManager(config, metrics, metadataCache, txnStateManager,
@@ -129,7 +129,7 @@ class TransactionCoordinator(brokerId: Int,
             state = Empty,
             topicPartitions = collection.mutable.Set.empty[TopicPartition],
             txnLastUpdateTimestamp = time.milliseconds())
-          txnManager.putTransactionStateIfNotExists(transactionalId, createdMetadata)
+          txnManager.putTransactionStateIfNotExists(createdMetadata)
 
         case Some(epochAndTxnMetadata) => Right(epochAndTxnMetadata)
       }
@@ -274,13 +274,37 @@ class TransactionCoordinator(brokerId: Int,
     }
   }
 
-  def handleTxnImmigration(txnTopicPartitionId: Int, coordinatorEpoch: Int) {
+  /**
+   * Load state from the given partition and begin handling requests for groups which map to this partition.
+   *
+   * @param txnTopicPartitionId The partition that we are now leading
+   * @param coordinatorEpoch The partition coordinator (or leader) epoch from the received LeaderAndIsr request
+   */
+  def onElection(txnTopicPartitionId: Int, coordinatorEpoch: Int): Unit = {
+    // The operations performed during immigration must be resilient to any previous errors we saw or partial state we
+    // left off during the unloading phase. Ensure we remove all associated state for this partition before we continue
+    // loading it.
+    txnMarkerChannelManager.removeMarkersForTxnTopicPartition(txnTopicPartitionId)
+
+    // Now load the partition.
     txnManager.loadTransactionsForTxnTopicPartition(txnTopicPartitionId, coordinatorEpoch, txnMarkerChannelManager.addTxnMarkersToSend)
   }
 
-  def handleTxnEmigration(txnTopicPartitionId: Int, coordinatorEpoch: Int) {
-      txnManager.removeTransactionsForTxnTopicPartition(txnTopicPartitionId, coordinatorEpoch)
-      txnMarkerChannelManager.removeMarkersForTxnTopicPartition(txnTopicPartitionId)
+  /**
+   * Clear coordinator caches for the given partition after giving up leadership.
+   *
+   * @param txnTopicPartitionId The partition that we are no longer leading
+   * @param coordinatorEpoch The partition coordinator (or leader) epoch, which may be absent if we
+   *                         are resigning after receiving a StopReplica request from the controller
+   */
+  def onResignation(txnTopicPartitionId: Int, coordinatorEpoch: Option[Int]): Unit = {
+    coordinatorEpoch match {
+      case Some(epoch) =>
+        txnManager.removeTransactionsForTxnTopicPartition(txnTopicPartitionId, epoch)
+      case None =>
+        txnManager.removeTransactionsForTxnTopicPartition(txnTopicPartitionId)
+    }
+    txnMarkerChannelManager.removeMarkersForTxnTopicPartition(txnTopicPartitionId)
   }
 
   private def logInvalidStateTransitionAndReturnError(transactionalId: String,
@@ -444,47 +468,52 @@ class TransactionCoordinator(brokerId: Int,
   def partitionFor(transactionalId: String): Int = txnManager.partitionFor(transactionalId)
 
   private def abortTimedOutTransactions(): Unit = {
+    def onComplete(txnIdAndPidEpoch: TransactionalIdAndProducerIdEpoch)(error: Errors): Unit = {
+      error match {
+        case Errors.NONE =>
+          info("Completed rollback of ongoing transaction for transactionalId " +
+            s"${txnIdAndPidEpoch.transactionalId} due to timeout")
+
+        case error@(Errors.INVALID_PRODUCER_ID_MAPPING |
+                    Errors.INVALID_PRODUCER_EPOCH |
+                    Errors.CONCURRENT_TRANSACTIONS) =>
+          debug(s"Rollback of ongoing transaction for transactionalId ${txnIdAndPidEpoch.transactionalId} " +
+            s"has been cancelled due to error $error")
+
+        case error =>
+          warn(s"Rollback of ongoing transaction for transactionalId ${txnIdAndPidEpoch.transactionalId} " +
+            s"failed due to error $error")
+      }
+    }
+
     txnManager.timedOutTransactions().foreach { txnIdAndPidEpoch =>
-      txnManager.getTransactionState(txnIdAndPidEpoch.transactionalId).right.flatMap {
+      txnManager.getTransactionState(txnIdAndPidEpoch.transactionalId).right.foreach {
         case None =>
-          error(s"Could not find transaction metadata when trying to timeout transaction with transactionalId " +
-            s"${txnIdAndPidEpoch.transactionalId}. ProducerId: ${txnIdAndPidEpoch.producerId}. ProducerEpoch: " +
-            s"${txnIdAndPidEpoch.producerEpoch}")
-          Left(Errors.INVALID_TXN_STATE)
+          error(s"Could not find transaction metadata when trying to timeout transaction for $txnIdAndPidEpoch")
 
         case Some(epochAndTxnMetadata) =>
           val txnMetadata = epochAndTxnMetadata.transactionMetadata
-          val transitMetadata = txnMetadata.inLock {
+          val transitMetadataOpt = txnMetadata.inLock {
             if (txnMetadata.producerId != txnIdAndPidEpoch.producerId) {
               error(s"Found incorrect producerId when expiring transactionalId: ${txnIdAndPidEpoch.transactionalId}. " +
                 s"Expected producerId: ${txnIdAndPidEpoch.producerId}. Found producerId: " +
                 s"${txnMetadata.producerId}")
-              Left(Errors.INVALID_PRODUCER_ID_MAPPING)
+              None
             } else if (txnMetadata.pendingTransitionInProgress) {
-              Left(Errors.CONCURRENT_TRANSACTIONS)
+              debug(s"Skipping abort of timed out transaction $txnIdAndPidEpoch since there is a " +
+                "pending state transition")
+              None
             } else {
-              Right(txnMetadata.prepareFenceProducerEpoch())
+              Some(txnMetadata.prepareFenceProducerEpoch())
             }
           }
-          transitMetadata match {
-            case Right(txnTransitMetadata) =>
-              handleEndTransaction(txnMetadata.transactionalId,
-                txnTransitMetadata.producerId,
-                txnTransitMetadata.producerEpoch,
-                TransactionResult.ABORT,
-                {
-                  case Errors.NONE =>
-                    info(s"Completed rollback ongoing transaction of transactionalId: ${txnIdAndPidEpoch.transactionalId} due to timeout")
-                  case e @ (Errors.INVALID_PRODUCER_ID_MAPPING |
-                            Errors.INVALID_PRODUCER_EPOCH |
-                            Errors.CONCURRENT_TRANSACTIONS) =>
-                    debug(s"Rolling back ongoing transaction of transactionalId: ${txnIdAndPidEpoch.transactionalId} has aborted due to ${e.exceptionName}")
-                  case e =>
-                    warn(s"Rolling back ongoing transaction of transactionalId: ${txnIdAndPidEpoch.transactionalId} failed due to ${e.exceptionName}")
-                })
-              Right(txnTransitMetadata)
-            case (error) =>
-              Left(error)
+
+          transitMetadataOpt.foreach { txnTransitMetadata =>
+            handleEndTransaction(txnMetadata.transactionalId,
+              txnTransitMetadata.producerId,
+              txnTransitMetadata.producerEpoch,
+              TransactionResult.ABORT,
+              onComplete(txnIdAndPidEpoch))
           }
       }
     }
@@ -493,11 +522,11 @@ class TransactionCoordinator(brokerId: Int,
   /**
    * Startup logic executed at the same time when the server starts up.
    */
-  def startup(enableTransactionalIdExpiration: Boolean = true) {
+  def startup(enableTransactionalIdExpiration: Boolean = true): Unit = {
     info("Starting up.")
     scheduler.startup()
     scheduler.schedule("transaction-abort",
-      () => abortTimedOutTransactions,
+      abortTimedOutTransactions,
       txnConfig.abortTimedOutTransactionsIntervalMs,
       txnConfig.abortTimedOutTransactionsIntervalMs
     )
@@ -513,7 +542,7 @@ class TransactionCoordinator(brokerId: Int,
    * Shutdown logic executed at the same time when server shuts down.
    * Ordering of actions should be reversed from the startup process.
    */
-  def shutdown() {
+  def shutdown(): Unit = {
     info("Shutting down.")
     isActive.set(false)
     scheduler.shutdown()
diff --git a/core/src/main/scala/kafka/coordinator/transaction/TransactionLog.scala b/core/src/main/scala/kafka/coordinator/transaction/TransactionLog.scala
index f74e8ad2dfa6a..2109095f08f35 100644
--- a/core/src/main/scala/kafka/coordinator/transaction/TransactionLog.scala
+++ b/core/src/main/scala/kafka/coordinator/transaction/TransactionLog.scala
@@ -16,16 +16,16 @@
  */
 package kafka.coordinator.transaction
 
-import kafka.common.MessageFormatter
-import org.apache.kafka.clients.consumer.ConsumerRecord
-import org.apache.kafka.common.{KafkaException, TopicPartition}
-import org.apache.kafka.common.protocol.types.Type._
-import org.apache.kafka.common.protocol.types._
 import java.io.PrintStream
 import java.nio.ByteBuffer
 import java.nio.charset.StandardCharsets
 
-import org.apache.kafka.common.record.CompressionType
+import kafka.common.MessageFormatter
+import org.apache.kafka.clients.consumer.ConsumerRecord
+import org.apache.kafka.common.protocol.types.Type._
+import org.apache.kafka.common.protocol.types._
+import org.apache.kafka.common.record.{CompressionType, Record}
+import org.apache.kafka.common.{KafkaException, TopicPartition}
 
 import scala.collection.mutable
 
@@ -130,7 +130,7 @@ object TransactionLog {
     *
     * @return key bytes
     */
-  private[coordinator] def keyToBytes(transactionalId: String): Array[Byte] = {
+  private[transaction] def keyToBytes(transactionalId: String): Array[Byte] = {
     import KeySchema._
     val key = new Struct(CURRENT)
     key.set(TXN_ID_FIELD, transactionalId)
@@ -146,7 +146,7 @@ object TransactionLog {
     *
     * @return value payload bytes
     */
-  private[coordinator] def valueToBytes(txnMetadata: TxnTransitMetadata): Array[Byte] = {
+  private[transaction] def valueToBytes(txnMetadata: TxnTransitMetadata): Array[Byte] = {
     import ValueSchema._
     val value = new Struct(Current)
     value.set(ProducerIdField, txnMetadata.producerId)
@@ -204,9 +204,9 @@ object TransactionLog {
     *
     * @return a transaction metadata object from the message
     */
-  def readTxnRecordValue(transactionalId: String, buffer: ByteBuffer): TransactionMetadata = {
+  def readTxnRecordValue(transactionalId: String, buffer: ByteBuffer): Option[TransactionMetadata] = {
     if (buffer == null) { // tombstone
-      null
+      None
     } else {
       import ValueSchema._
       val version = buffer.getShort
@@ -243,7 +243,7 @@ object TransactionLog {
           }
         }
 
-        transactionMetadata
+        Some(transactionMetadata)
       } else {
         throw new IllegalStateException(s"Unknown version $version from the transaction log message value")
       }
@@ -252,20 +252,43 @@ object TransactionLog {
 
   // Formatter for use with tools to read transaction log messages
   class TransactionLogMessageFormatter extends MessageFormatter {
-    def writeTo(consumerRecord: ConsumerRecord[Array[Byte], Array[Byte]], output: PrintStream) {
+    def writeTo(consumerRecord: ConsumerRecord[Array[Byte], Array[Byte]], output: PrintStream): Unit = {
       Option(consumerRecord.key).map(key => readTxnRecordKey(ByteBuffer.wrap(key))).foreach { txnKey =>
         val transactionalId = txnKey.transactionalId
         val value = consumerRecord.value
-        val producerIdMetadata =
-          if (value == null) "NULL"
-          else readTxnRecordValue(transactionalId, ByteBuffer.wrap(value))
+        val producerIdMetadata = if (value == null)
+          None
+        else
+          readTxnRecordValue(transactionalId, ByteBuffer.wrap(value))
         output.write(transactionalId.getBytes(StandardCharsets.UTF_8))
         output.write("::".getBytes(StandardCharsets.UTF_8))
-        output.write(producerIdMetadata.toString.getBytes(StandardCharsets.UTF_8))
+        output.write(producerIdMetadata.getOrElse("NULL").toString.getBytes(StandardCharsets.UTF_8))
         output.write("\n".getBytes(StandardCharsets.UTF_8))
       }
     }
   }
+
+  /**
+   * Exposed for printing records using [[kafka.tools.DumpLogSegments]]
+   */
+  def formatRecordKeyAndValue(record: Record): (Option[String], Option[String]) = {
+    val txnKey = TransactionLog.readTxnRecordKey(record.key)
+    val keyString = s"transaction_metadata::transactionalId=${txnKey.transactionalId}"
+
+    val valueString = TransactionLog.readTxnRecordValue(txnKey.transactionalId, record.value) match {
+      case None => ""
+
+      case Some(txnMetadata) => s"producerId:${txnMetadata.producerId}," +
+        s"producerEpoch:${txnMetadata.producerEpoch}," +
+        s"state=${txnMetadata.state}," +
+        s"partitions=${txnMetadata.topicPartitions.mkString("[", ",", "]")}," +
+        s"txnLastUpdateTimestamp=${txnMetadata.txnLastUpdateTimestamp}," +
+        s"txnTimeoutMs=${txnMetadata.txnTimeoutMs}"
+    }
+
+    (Some(keyString), Some(valueString))
+  }
+
 }
 
 case class TxnKey(version: Short, transactionalId: String) {
diff --git a/core/src/main/scala/kafka/coordinator/transaction/TransactionMarkerChannelManager.scala b/core/src/main/scala/kafka/coordinator/transaction/TransactionMarkerChannelManager.scala
index f49fa7b2a34d9..34cc6391193d3 100644
--- a/core/src/main/scala/kafka/coordinator/transaction/TransactionMarkerChannelManager.scala
+++ b/core/src/main/scala/kafka/coordinator/transaction/TransactionMarkerChannelManager.scala
@@ -22,7 +22,7 @@ import kafka.metrics.KafkaMetricsGroup
 import kafka.server.{DelayedOperationPurgatory, KafkaConfig, MetadataCache}
 import kafka.utils.{CoreUtils, Logging}
 import org.apache.kafka.clients._
-import org.apache.kafka.common.{Node, TopicPartition}
+import org.apache.kafka.common.{Node, Reconfigurable, TopicPartition}
 import org.apache.kafka.common.metrics.Metrics
 import org.apache.kafka.common.network._
 import org.apache.kafka.common.requests.{TransactionResult, WriteTxnMarkersRequest}
@@ -54,6 +54,10 @@ object TransactionMarkerChannelManager {
       time,
       config.saslInterBrokerHandshakeRequestEnable
     )
+    channelBuilder match {
+      case reconfigurable: Reconfigurable => config.addReconfigurable(reconfigurable)
+      case _ =>
+    }
     val selector = new Selector(
       NetworkReceive.UNLIMITED,
       config.connectionsMaxIdleMs,
@@ -169,7 +173,7 @@ class TransactionMarkerChannelManager(config: KafkaConfig,
   // visible for testing
   private[transaction] def queueForUnknownBroker = markersQueueForUnknownBroker
 
-  private[transaction] def addMarkersForBroker(broker: Node, txnTopicPartition: Int, txnIdAndMarker: TxnIdAndMarkerEntry) {
+  private[transaction] def addMarkersForBroker(broker: Node, txnTopicPartition: Int, txnIdAndMarker: TxnIdAndMarkerEntry): Unit = {
     val brokerId = broker.id
 
     // we do not synchronize on the update of the broker node with the enqueuing,
diff --git a/core/src/main/scala/kafka/coordinator/transaction/TransactionMarkerRequestCompletionHandler.scala b/core/src/main/scala/kafka/coordinator/transaction/TransactionMarkerRequestCompletionHandler.scala
index fefe767eb40b8..f655770ce6b0c 100644
--- a/core/src/main/scala/kafka/coordinator/transaction/TransactionMarkerRequestCompletionHandler.scala
+++ b/core/src/main/scala/kafka/coordinator/transaction/TransactionMarkerRequestCompletionHandler.scala
@@ -146,7 +146,8 @@ class TransactionMarkerRequestCompletionHandler(brokerId: Int,
                          Errors.NOT_LEADER_FOR_PARTITION |
                          Errors.NOT_ENOUGH_REPLICAS |
                          Errors.NOT_ENOUGH_REPLICAS_AFTER_APPEND |
-                         Errors.REQUEST_TIMED_OUT => // these are retriable errors
+                         Errors.REQUEST_TIMED_OUT |
+                         Errors.KAFKA_STORAGE_ERROR => // these are retriable errors
 
                       info(s"Sending $transactionalId's transaction marker for partition $topicPartition has failed with error ${error.exceptionName}, retrying " +
                         s"with current coordinator epoch ${epochAndMetadata.coordinatorEpoch}")
diff --git a/core/src/main/scala/kafka/coordinator/transaction/TransactionStateManager.scala b/core/src/main/scala/kafka/coordinator/transaction/TransactionStateManager.scala
index b45953f861eea..36bc9659427ea 100644
--- a/core/src/main/scala/kafka/coordinator/transaction/TransactionStateManager.scala
+++ b/core/src/main/scala/kafka/coordinator/transaction/TransactionStateManager.scala
@@ -22,23 +22,24 @@ import java.util.concurrent.TimeUnit
 import java.util.concurrent.atomic.AtomicBoolean
 import java.util.concurrent.locks.ReentrantReadWriteLock
 
-import kafka.log.LogConfig
+import kafka.log.{AppendOrigin, LogConfig}
 import kafka.message.UncompressedCodec
-import kafka.server.Defaults
-import kafka.server.ReplicaManager
+import kafka.server.{Defaults, FetchLogEnd, ReplicaManager}
 import kafka.utils.CoreUtils.{inReadLock, inWriteLock}
 import kafka.utils.{Logging, Pool, Scheduler}
 import kafka.zk.KafkaZkClient
-import org.apache.kafka.common.{KafkaException, TopicPartition}
 import org.apache.kafka.common.internals.Topic
+import org.apache.kafka.common.metrics.Metrics
+import org.apache.kafka.common.metrics.stats.{Avg, Max}
 import org.apache.kafka.common.protocol.Errors
 import org.apache.kafka.common.record.{FileRecords, MemoryRecords, SimpleRecord}
 import org.apache.kafka.common.requests.ProduceResponse.PartitionResponse
 import org.apache.kafka.common.requests.TransactionResult
 import org.apache.kafka.common.utils.{Time, Utils}
+import org.apache.kafka.common.{KafkaException, TopicPartition}
 
-import scala.collection.mutable
 import scala.collection.JavaConverters._
+import scala.collection.mutable
 
 
 object TransactionStateManager {
@@ -71,7 +72,8 @@ class TransactionStateManager(brokerId: Int,
                               scheduler: Scheduler,
                               replicaManager: ReplicaManager,
                               config: TransactionConfig,
-                              time: Time) extends Logging {
+                              time: Time,
+                              metrics: Metrics) extends Logging {
 
   this.logIdent = "[Transaction State Manager " + brokerId + "]: "
 
@@ -84,23 +86,28 @@ class TransactionStateManager(brokerId: Int,
   private val stateLock = new ReentrantReadWriteLock()
 
   /** partitions of transaction topic that are being loaded, state lock should be called BEFORE accessing this set */
-  private val loadingPartitions: mutable.Set[TransactionPartitionAndLeaderEpoch] = mutable.Set()
-
-  /** partitions of transaction topic that are being removed, state lock should be called BEFORE accessing this set */
-  private val leavingPartitions: mutable.Set[TransactionPartitionAndLeaderEpoch] = mutable.Set()
+  private[transaction] val loadingPartitions: mutable.Set[TransactionPartitionAndLeaderEpoch] = mutable.Set()
 
   /** transaction metadata cache indexed by assigned transaction topic partition ids */
-  private val transactionMetadataCache: mutable.Map[Int, TxnMetadataCacheEntry] = mutable.Map()
+  private[transaction] val transactionMetadataCache: mutable.Map[Int, TxnMetadataCacheEntry] = mutable.Map()
 
   /** number of partitions for the transaction log topic */
   private val transactionTopicPartitionCount = getTransactionTopicPartitionCount
 
+  /** setup metrics*/
+  private val partitionLoadSensor = metrics.sensor("PartitionLoadTime")
+
+  partitionLoadSensor.add(metrics.metricName("partition-load-time-max",
+    "transaction-coordinator-metrics",
+    "The max time it took to load the partitions in the last 30sec"), new Max())
+  partitionLoadSensor.add(metrics.metricName("partition-load-time-avg",
+    "transaction-coordinator-metrics",
+    "The avg time it took to load the partitions in the last 30sec"), new Avg())
+
   // visible for testing only
   private[transaction] def addLoadingPartition(partitionId: Int, coordinatorEpoch: Int): Unit = {
     val partitionAndLeaderEpoch = TransactionPartitionAndLeaderEpoch(partitionId, coordinatorEpoch)
-
     inWriteLock(stateLock) {
-      leavingPartitions.remove(partitionAndLeaderEpoch)
       loadingPartitions.add(partitionAndLeaderEpoch)
     }
   }
@@ -113,9 +120,7 @@ class TransactionStateManager(brokerId: Int,
   def timedOutTransactions(): Iterable[TransactionalIdAndProducerIdEpoch] = {
     val now = time.milliseconds()
     inReadLock(stateLock) {
-      transactionMetadataCache.filter { case (txnPartitionId, _) =>
-        !leavingPartitions.exists(_.txnPartitionId == txnPartitionId)
-      }.flatMap { case (_, entry) =>
+      transactionMetadataCache.flatMap { case (_, entry) =>
         entry.metadataPerTransactionalId.filter { case (_, txnMetadata) =>
           if (txnMetadata.pendingTransitionInProgress) {
             false
@@ -133,7 +138,7 @@ class TransactionStateManager(brokerId: Int,
     }
   }
 
-  def enableTransactionalIdExpiration() {
+  def enableTransactionalIdExpiration(): Unit = {
     scheduler.schedule("transactionalId-expiration", () => {
       val now = time.milliseconds()
       inReadLock(stateLock) {
@@ -199,7 +204,7 @@ class TransactionStateManager(brokerId: Int,
           config.requestTimeoutMs,
           TransactionLog.EnforcedRequiredAcks,
           internalTopicsAllowed = true,
-          isFromClient = false,
+          origin = AppendOrigin.Coordinator,
           recordsPerPartition,
           removeFromCacheCallback,
           Some(stateLock.readLock)
@@ -212,10 +217,10 @@ class TransactionStateManager(brokerId: Int,
   def getTransactionState(transactionalId: String): Either[Errors, Option[CoordinatorEpochAndTxnMetadata]] =
     getAndMaybeAddTransactionState(transactionalId, None)
 
-  def putTransactionStateIfNotExists(transactionalId: String,
-                                     txnMetadata: TransactionMetadata): Either[Errors, CoordinatorEpochAndTxnMetadata] =
-    getAndMaybeAddTransactionState(transactionalId, Some(txnMetadata))
+  def putTransactionStateIfNotExists(txnMetadata: TransactionMetadata): Either[Errors, CoordinatorEpochAndTxnMetadata] = {
+    getAndMaybeAddTransactionState(txnMetadata.transactionalId, Some(txnMetadata))
       .right.map(_.getOrElse(throw new IllegalStateException(s"Unexpected empty transaction metadata returned while putting $txnMetadata")))
+  }
 
   /**
    * Get the transaction metadata associated with the given transactional id, or an error if
@@ -230,8 +235,6 @@ class TransactionStateManager(brokerId: Int,
       val partitionId = partitionFor(transactionalId)
       if (loadingPartitions.exists(_.txnPartitionId == partitionId))
         Left(Errors.COORDINATOR_LOAD_IN_PROGRESS)
-      else if (leavingPartitions.exists(_.txnPartitionId == partitionId))
-        Left(Errors.NOT_COORDINATOR)
       else {
         transactionMetadataCache.get(partitionId) match {
           case Some(cacheEntry) =>
@@ -287,22 +290,29 @@ class TransactionStateManager(brokerId: Int,
 
     replicaManager.getLog(topicPartition) match {
       case None =>
-        warn(s"Attempted to load offsets and group metadata from $topicPartition, but found no log")
+        warn(s"Attempted to load transaction metadata from $topicPartition, but found no log")
 
       case Some(log) =>
         // buffer may not be needed if records are read from memory
         var buffer = ByteBuffer.allocate(0)
 
-        // loop breaks if leader changes at any time during the load, since getHighWatermark is -1
+        // loop breaks if leader changes at any time during the load, since logEndOffset is -1
         var currOffset = log.logStartOffset
 
+        // loop breaks if no records have been read, since the end of the log has been reached
+        var readAtLeastOneRecord = true
+
         try {
-          while (currOffset < logEndOffset
-            && !shuttingDown.get()
-            && inReadLock(stateLock) {loadingPartitions.exists { idAndEpoch: TransactionPartitionAndLeaderEpoch =>
+          while (currOffset < logEndOffset && readAtLeastOneRecord && !shuttingDown.get() && inReadLock(stateLock) {
+            loadingPartitions.exists { idAndEpoch: TransactionPartitionAndLeaderEpoch =>
               idAndEpoch.txnPartitionId == topicPartition.partition && idAndEpoch.coordinatorEpoch == coordinatorEpoch}}) {
-            val fetchDataInfo = log.read(currOffset, config.transactionLogLoadBufferSize, maxOffset = None,
-              minOneMessage = true, includeAbortedTxns = false)
+            val fetchDataInfo = log.read(currOffset,
+              maxLength = config.transactionLogLoadBufferSize,
+              isolation = FetchLogEnd,
+              minOneMessage = true)
+
+            readAtLeastOneRecord = fetchDataInfo.records.sizeInBytes > 0
+
             val memRecords = fetchDataInfo.records match {
               case records: MemoryRecords => records
               case fileRecords: FileRecords =>
@@ -312,7 +322,7 @@ class TransactionStateManager(brokerId: Int,
                 // minOneMessage = true in the above log.read means that the buffer may need to be grown to ensure progress can be made
                 if (buffer.capacity < bytesNeeded) {
                   if (config.transactionLogLoadBufferSize < bytesNeeded)
-                    warn(s"Loaded offsets and group metadata from $topicPartition with buffer larger ($bytesNeeded bytes) than " +
+                    warn(s"Loaded transaction metadata from $topicPartition with buffer larger ($bytesNeeded bytes) than " +
                       s"configured transaction.state.log.load.buffer.size (${config.transactionLogLoadBufferSize} bytes)")
 
                   buffer = ByteBuffer.allocate(bytesNeeded)
@@ -330,17 +340,19 @@ class TransactionStateManager(brokerId: Int,
                 val txnKey = TransactionLog.readTxnRecordKey(record.key)
                 // load transaction metadata along with transaction state
                 val transactionalId = txnKey.transactionalId
-                if (!record.hasValue) {
-                  loadedTransactions.remove(transactionalId)
-                } else {
-                  val txnMetadata = TransactionLog.readTxnRecordValue(transactionalId, record.value)
-                  loadedTransactions.put(transactionalId, txnMetadata)
+                TransactionLog.readTxnRecordValue(transactionalId, record.value) match {
+                  case None =>
+                    loadedTransactions.remove(transactionalId)
+                  case Some(txnMetadata) =>
+                    loadedTransactions.put(transactionalId, txnMetadata)
                 }
                 currOffset = batch.nextOffset
               }
             }
-
-            info(s"Finished loading ${loadedTransactions.size} transaction metadata from $topicPartition in ${time.milliseconds() - startMs} milliseconds")
+            val endMs = time.milliseconds()
+            val timeLapse = endMs - startMs
+            partitionLoadSensor.record(timeLapse, endMs, false)
+            info(s"Finished loading ${loadedTransactions.size} transaction metadata from $topicPartition in $timeLapse milliseconds")
           }
         } catch {
           case t: Throwable => error(s"Error loading transactions from transaction log $topicPartition", t)
@@ -352,41 +364,36 @@ class TransactionStateManager(brokerId: Int,
 
   /**
    * Add a transaction topic partition into the cache
-   *
-   * Make it package-private to be used only for unit tests.
    */
-  private[transaction] def addLoadedTransactionsToCache(txnTopicPartition: Int, coordinatorEpoch: Int, metadataPerTransactionalId: Pool[String, TransactionMetadata]): Unit = {
-    val txnMetadataCacheEntry = TxnMetadataCacheEntry(coordinatorEpoch, metadataPerTransactionalId)
-    val currentTxnMetadataCacheEntry = transactionMetadataCache.put(txnTopicPartition, txnMetadataCacheEntry)
-
-    if (currentTxnMetadataCacheEntry.isDefined) {
-      val coordinatorEpoch = currentTxnMetadataCacheEntry.get.coordinatorEpoch
-      val metadataPerTxnId = currentTxnMetadataCacheEntry.get.metadataPerTransactionalId
-      val errorMsg = s"The metadata cache for txn partition $txnTopicPartition has already exist with epoch $coordinatorEpoch " +
-        s"and ${metadataPerTxnId.size} entries while trying to add to it; " +
-        s"this should not happen"
-      fatal(errorMsg)
-      throw new IllegalStateException(errorMsg)
+  private[transaction] def addLoadedTransactionsToCache(txnTopicPartition: Int,
+                                                        coordinatorEpoch: Int,
+                                                        loadedTransactions: Pool[String, TransactionMetadata]): Unit = {
+    val txnMetadataCacheEntry = TxnMetadataCacheEntry(coordinatorEpoch, loadedTransactions)
+    val previousTxnMetadataCacheEntryOpt = transactionMetadataCache.put(txnTopicPartition, txnMetadataCacheEntry)
+
+    previousTxnMetadataCacheEntryOpt.foreach { previousTxnMetadataCacheEntry =>
+      warn(s"Unloaded transaction metadata $previousTxnMetadataCacheEntry from $txnTopicPartition as part of " +
+        s"loading metadata at epoch $coordinatorEpoch")
     }
   }
 
   /**
-   * When this broker becomes a leader for a transaction log partition, load this partition and
-   * populate the transaction metadata cache with the transactional ids.
+   * When this broker becomes a leader for a transaction log partition, load this partition and populate the transaction
+   * metadata cache with the transactional ids. This operation must be resilient to any partial state left off from
+   * the previous loading / unloading operation.
    */
-  def loadTransactionsForTxnTopicPartition(partitionId: Int, coordinatorEpoch: Int, sendTxnMarkers: SendTxnMarkersCallback) {
-    validateTransactionTopicPartitionCountIsStable()
-
+  def loadTransactionsForTxnTopicPartition(partitionId: Int, coordinatorEpoch: Int, sendTxnMarkers: SendTxnMarkersCallback): Unit = {
     val topicPartition = new TopicPartition(Topic.TRANSACTION_STATE_TOPIC_NAME, partitionId)
     val partitionAndLeaderEpoch = TransactionPartitionAndLeaderEpoch(partitionId, coordinatorEpoch)
 
     inWriteLock(stateLock) {
-      leavingPartitions.remove(partitionAndLeaderEpoch)
       loadingPartitions.add(partitionAndLeaderEpoch)
     }
 
-    def loadTransactions() {
-      info(s"Loading transaction metadata from $topicPartition")
+    def loadTransactions(): Unit = {
+      info(s"Loading transaction metadata from $topicPartition at epoch $coordinatorEpoch")
+      validateTransactionTopicPartitionCountIsStable()
+
       val loadedTransactions = loadTransactionMetadata(topicPartition, coordinatorEpoch)
 
       inWriteLock(stateLock) {
@@ -406,7 +413,7 @@ class TransactionStateManager(brokerId: Int,
                     transactionsPendingForCompletion +=
                       TransactionalIdCoordinatorEpochAndTransitMetadata(transactionalId, coordinatorEpoch, TransactionResult.COMMIT, txnMetadata, txnMetadata.prepareComplete(time.milliseconds()))
                   case _ =>
-                    // nothing need to be done
+                    // nothing needs to be done
                 }
               }
           }
@@ -421,44 +428,42 @@ class TransactionStateManager(brokerId: Int,
           }
         }
       }
+
+      info(s"Completed loading transaction metadata from $topicPartition for coordinator epoch $coordinatorEpoch")
     }
 
-    scheduler.schedule(s"load-txns-for-partition-$topicPartition", () => loadTransactions)
+    scheduler.schedule(s"load-txns-for-partition-$topicPartition", loadTransactions)
+  }
+
+  def removeTransactionsForTxnTopicPartition(partitionId: Int): Unit = {
+    val topicPartition = new TopicPartition(Topic.TRANSACTION_STATE_TOPIC_NAME, partitionId)
+    inWriteLock(stateLock) {
+      loadingPartitions.retain(_.txnPartitionId != partitionId)
+      transactionMetadataCache.remove(partitionId).foreach { txnMetadataCacheEntry =>
+        info(s"Unloaded transaction metadata $txnMetadataCacheEntry for $topicPartition following " +
+          s"local partition deletion")
+      }
+    }
   }
 
   /**
    * When this broker becomes a follower for a transaction log partition, clear out the cache for corresponding transactional ids
    * that belong to that partition.
    */
-  def removeTransactionsForTxnTopicPartition(partitionId: Int, coordinatorEpoch: Int) {
-    validateTransactionTopicPartitionCountIsStable()
-
+  def removeTransactionsForTxnTopicPartition(partitionId: Int, coordinatorEpoch: Int): Unit = {
     val topicPartition = new TopicPartition(Topic.TRANSACTION_STATE_TOPIC_NAME, partitionId)
     val partitionAndLeaderEpoch = TransactionPartitionAndLeaderEpoch(partitionId, coordinatorEpoch)
 
     inWriteLock(stateLock) {
       loadingPartitions.remove(partitionAndLeaderEpoch)
-      leavingPartitions.add(partitionAndLeaderEpoch)
-    }
+      transactionMetadataCache.remove(partitionId) match {
+        case Some(txnMetadataCacheEntry) =>
+          info(s"Unloaded transaction metadata $txnMetadataCacheEntry for $topicPartition on become-follower transition")
 
-    def removeTransactions() {
-      inWriteLock(stateLock) {
-        if (leavingPartitions.contains(partitionAndLeaderEpoch)) {
-          transactionMetadataCache.remove(partitionId) match {
-            case Some(txnMetadataCacheEntry) =>
-              info(s"Removed ${txnMetadataCacheEntry.metadataPerTransactionalId.size} cached transaction metadata for $topicPartition on follower transition")
-
-            case None =>
-              info(s"Trying to remove cached transaction metadata for $topicPartition on follower transition but there is no entries remaining; " +
-                s"it is likely that another process for removing the cached entries has just executed earlier before")
-          }
-
-          leavingPartitions.remove(partitionAndLeaderEpoch)
-        }
+        case None =>
+          info(s"No cached transaction metadata found for $topicPartition during become-follower transition")
       }
     }
-
-    scheduler.schedule(s"remove-txns-for-partition-$topicPartition", () => removeTransactions)
   }
 
   private def validateTransactionTopicPartitionCountIsStable(): Unit = {
@@ -622,7 +627,7 @@ class TransactionStateManager(brokerId: Int,
                 newMetadata.txnTimeoutMs.toLong,
                 TransactionLog.EnforcedRequiredAcks,
                 internalTopicsAllowed = true,
-                isFromClient = false,
+                origin = AppendOrigin.Coordinator,
                 recordsPerPartition,
                 updateCacheCallback,
                 delayedProduceLock = Some(stateLock.readLock))
@@ -633,7 +638,7 @@ class TransactionStateManager(brokerId: Int,
     }
   }
 
-  def shutdown() {
+  def shutdown(): Unit = {
     shuttingDown.set(true)
     loadingPartitions.clear()
     transactionMetadataCache.clear()
@@ -643,7 +648,12 @@ class TransactionStateManager(brokerId: Int,
 }
 
 
-private[transaction] case class TxnMetadataCacheEntry(coordinatorEpoch: Int, metadataPerTransactionalId: Pool[String, TransactionMetadata])
+private[transaction] case class TxnMetadataCacheEntry(coordinatorEpoch: Int,
+                                                      metadataPerTransactionalId: Pool[String, TransactionMetadata]) {
+  override def toString: String = {
+    s"TxnMetadataCacheEntry(coordinatorEpoch=$coordinatorEpoch, numTransactionalEntries=${metadataPerTransactionalId.size})"
+  }
+}
 
 private[transaction] case class CoordinatorEpochAndTxnMetadata(coordinatorEpoch: Int, transactionMetadata: TransactionMetadata)
 
@@ -658,7 +668,11 @@ private[transaction] case class TransactionConfig(transactionalIdExpirationMs: I
                                                   removeExpiredTransactionalIdsIntervalMs: Int = TransactionStateManager.DefaultRemoveExpiredTransactionalIdsIntervalMs,
                                                   requestTimeoutMs: Int = Defaults.RequestTimeoutMs)
 
-case class TransactionalIdAndProducerIdEpoch(transactionalId: String, producerId: Long, producerEpoch: Short)
+case class TransactionalIdAndProducerIdEpoch(transactionalId: String, producerId: Long, producerEpoch: Short) {
+  override def toString: String = {
+    s"(transactionalId=$transactionalId, producerId=$producerId, producerEpoch=$producerEpoch)"
+  }
+}
 
 case class TransactionPartitionAndLeaderEpoch(txnPartitionId: Int, coordinatorEpoch: Int)
 
diff --git a/core/src/main/scala/kafka/log/AbstractIndex.scala b/core/src/main/scala/kafka/log/AbstractIndex.scala
index c69e78344dc67..9fabb6322e128 100644
--- a/core/src/main/scala/kafka/log/AbstractIndex.scala
+++ b/core/src/main/scala/kafka/log/AbstractIndex.scala
@@ -34,12 +34,12 @@ import scala.math.ceil
 /**
  * The abstract index class which holds entry format agnostic methods.
  *
- * @param file The index file
+ * @param _file The index file
  * @param baseOffset the base offset of the segment that this index is corresponding to.
  * @param maxIndexSize The maximum index size in bytes.
  */
-abstract class AbstractIndex[K, V](@volatile var file: File, val baseOffset: Long,
-                                   val maxIndexSize: Int = -1, val writable: Boolean) extends Closeable {
+abstract class AbstractIndex(@volatile private var _file: File, val baseOffset: Long, val maxIndexSize: Int = -1,
+                             val writable: Boolean) extends Closeable {
   import AbstractIndex._
 
   // Length of the index file
@@ -144,23 +144,27 @@ abstract class AbstractIndex[K, V](@volatile var file: File, val baseOffset: Lon
    * The maximum number of entries this index can hold
    */
   @volatile
-  private[this] var _maxEntries = mmap.limit() / entrySize
+  private[this] var _maxEntries: Int = mmap.limit() / entrySize
 
   /** The number of entries in this index */
   @volatile
-  protected var _entries = mmap.position() / entrySize
+  protected var _entries: Int = mmap.position() / entrySize
 
   /**
    * True iff there are no more slots available in this index
    */
   def isFull: Boolean = _entries >= _maxEntries
 
+  def file: File = _file
+
   def maxEntries: Int = _maxEntries
 
   def entries: Int = _entries
 
   def length: Long = _length
 
+  def updateParentDir(parentDir: File): Unit = _file = new File(parentDir, file.getName)
+
   /**
    * Reset the size of the memory map and the underneath file. This is used in two kinds of cases: (1) in
    * trimToValidSize() which is called at closing the segment or new segment being rolled; (2) at
@@ -175,6 +179,8 @@ abstract class AbstractIndex[K, V](@volatile var file: File, val baseOffset: Lon
       val roundedNewSize = roundDownToExactMultiple(newSize, entrySize)
 
       if (_length == roundedNewSize) {
+        if (isDebugEnabled)
+          debug(s"Index ${file.getAbsolutePath} was not resized because it already has size $roundedNewSize")
         false
       } else {
         val raf = new RandomAccessFile(file, "rw")
@@ -189,6 +195,9 @@ abstract class AbstractIndex[K, V](@volatile var file: File, val baseOffset: Lon
           mmap = raf.getChannel().map(FileChannel.MapMode.READ_WRITE, 0, roundedNewSize)
           _maxEntries = mmap.limit() / entrySize
           mmap.position(position)
+          if (isDebugEnabled)
+            debug(s"Resized ${file.getAbsolutePath} to $roundedNewSize, position is ${mmap.position()} " +
+                  s"and limit is ${mmap.limit()}")
           true
         } finally {
           CoreUtils.swallow(raf.close(), AbstractIndex)
@@ -202,15 +211,15 @@ abstract class AbstractIndex[K, V](@volatile var file: File, val baseOffset: Lon
    *
    * @throws IOException if rename fails
    */
-  def renameTo(f: File) {
+  def renameTo(f: File): Unit = {
     try Utils.atomicMoveWithFallback(file.toPath, f.toPath)
-    finally file = f
+    finally _file = f
   }
 
   /**
    * Flush the data in the index to disk
    */
-  def flush() {
+  def flush(): Unit = {
     inLock(lock) {
       mmap.force()
     }
@@ -232,7 +241,7 @@ abstract class AbstractIndex[K, V](@volatile var file: File, val baseOffset: Lon
    * Trim this segment to fit just the valid entries, deleting all trailing unwritten bytes from
    * the file.
    */
-  def trimToValidSize() {
+  def trimToValidSize(): Unit = {
     inLock(lock) {
       resize(entrySize * _entries)
     }
@@ -241,10 +250,10 @@ abstract class AbstractIndex[K, V](@volatile var file: File, val baseOffset: Lon
   /**
    * The number of bytes actually used by this index
    */
-  def sizeInBytes = entrySize * _entries
+  def sizeInBytes: Int = entrySize * _entries
 
   /** Close the index */
-  def close() {
+  def close(): Unit = {
     trimToValidSize()
     closeHandler()
   }
@@ -315,7 +324,7 @@ abstract class AbstractIndex[K, V](@volatile var file: File, val baseOffset: Lon
   /**
    * Forcefully free the buffer's mmap.
    */
-  protected[log] def forceUnmap() {
+  protected[log] def forceUnmap(): Unit = {
     try MappedByteBuffers.unmap(file.getAbsolutePath, mmap)
     finally mmap = null // Accessing unmapped mmap crashes JVM by SEGV so we null it out to be safe
   }
@@ -424,7 +433,7 @@ abstract class AbstractIndex[K, V](@volatile var file: File, val baseOffset: Lon
 }
 
 object AbstractIndex extends Logging {
-  override val loggerName: String = classOf[AbstractIndex[_, _]].getName
+  override val loggerName: String = classOf[AbstractIndex].getName
 }
 
 object IndexSearchType extends Enumeration {
diff --git a/core/src/main/scala/kafka/log/LazyIndex.scala b/core/src/main/scala/kafka/log/LazyIndex.scala
new file mode 100644
index 0000000000000..a5a7c34a6e5b7
--- /dev/null
+++ b/core/src/main/scala/kafka/log/LazyIndex.scala
@@ -0,0 +1,166 @@
+/**
+ * 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.log
+
+import java.io.File
+import java.nio.file.{Files, NoSuchFileException}
+import java.util.concurrent.locks.ReentrantLock
+
+import LazyIndex._
+import kafka.utils.CoreUtils.inLock
+import kafka.utils.threadsafe
+import org.apache.kafka.common.utils.Utils
+
+/**
+  * A wrapper over an `AbstractIndex` instance that provides a mechanism to defer loading
+  * (i.e. memory mapping) the underlying index until it is accessed for the first time via the
+  * `get` method.
+  *
+  * In addition, this class exposes a number of methods (e.g. updateParentDir, renameTo, close,
+  * etc.) that provide the desired behavior without causing the index to be loaded. If the index
+  * had previously been loaded, the methods in this class simply delegate to the relevant method in
+  * the index.
+  *
+  * This is an important optimization with regards to broker start-up and shutdown time if it has a
+  * large number of segments.
+  *
+  * Methods of this class are thread safe. Make sure to check `AbstractIndex` subclasses
+  * documentation to establish their thread safety.
+  *
+  * @param loadIndex A function that takes a `File` pointing to an index and returns a loaded
+  *                  `AbstractIndex` instance.
+  */
+@threadsafe
+class LazyIndex[T <: AbstractIndex] private (@volatile private var indexWrapper: IndexWrapper, loadIndex: File => T) {
+
+  private val lock = new ReentrantLock()
+
+  def file: File = indexWrapper.file
+
+  def get: T = {
+    indexWrapper match {
+      case indexValue: IndexValue[T] => indexValue.index
+      case _: IndexFile =>
+        inLock(lock) {
+          indexWrapper match {
+            case indexValue: IndexValue[T] => indexValue.index
+            case indexFile: IndexFile =>
+              val indexValue = new IndexValue(loadIndex(indexFile.file))
+              indexWrapper = indexValue
+              indexValue.index
+          }
+        }
+    }
+  }
+
+  def updateParentDir(parentDir: File): Unit = {
+    inLock(lock) {
+      indexWrapper.updateParentDir(parentDir)
+    }
+  }
+
+  def renameTo(f: File): Unit = {
+    inLock(lock) {
+      indexWrapper.renameTo(f)
+    }
+  }
+
+  def deleteIfExists(): Boolean = {
+    inLock(lock) {
+      indexWrapper.deleteIfExists()
+    }
+  }
+
+  def close(): Unit = {
+    inLock(lock) {
+      indexWrapper.close()
+    }
+  }
+
+  def closeHandler(): Unit = {
+    inLock(lock) {
+      indexWrapper.closeHandler()
+    }
+  }
+
+}
+
+object LazyIndex {
+
+  def forOffset(file: File, baseOffset: Long, maxIndexSize: Int = -1, writable: Boolean = true): LazyIndex[OffsetIndex] =
+    new LazyIndex(new IndexFile(file), file => new OffsetIndex(file, baseOffset, maxIndexSize, writable))
+
+  def forTime(file: File, baseOffset: Long, maxIndexSize: Int = -1, writable: Boolean = true): LazyIndex[TimeIndex] =
+    new LazyIndex(new IndexFile(file), file => new TimeIndex(file, baseOffset, maxIndexSize, writable))
+
+  private sealed trait IndexWrapper {
+
+    def file: File
+
+    def updateParentDir(f: File): Unit
+
+    def renameTo(f: File): Unit
+
+    def deleteIfExists(): Boolean
+
+    def close(): Unit
+
+    def closeHandler(): Unit
+
+  }
+
+  private class IndexFile(@volatile private var _file: File) extends IndexWrapper {
+
+    def file: File = _file
+
+    def updateParentDir(parentDir: File): Unit = _file = new File(parentDir, file.getName)
+
+    def renameTo(f: File): Unit = {
+      try Utils.atomicMoveWithFallback(file.toPath, f.toPath)
+      catch {
+        case _: NoSuchFileException if !file.exists => ()
+      }
+      finally _file = f
+    }
+
+    def deleteIfExists(): Boolean = Files.deleteIfExists(file.toPath)
+
+    def close(): Unit = ()
+
+    def closeHandler(): Unit = ()
+
+  }
+
+  private class IndexValue[T <: AbstractIndex](val index: T) extends IndexWrapper {
+
+    def file: File = index.file
+
+    def updateParentDir(parentDir: File): Unit = index.updateParentDir(parentDir)
+
+    def renameTo(f: File): Unit = index.renameTo(f)
+
+    def deleteIfExists(): Boolean = index.deleteIfExists()
+
+    def close(): Unit = index.close()
+
+    def closeHandler(): Unit = index.closeHandler()
+
+  }
+
+}
+
diff --git a/core/src/main/scala/kafka/log/Log.scala b/core/src/main/scala/kafka/log/Log.scala
index 5ad3c3e58102a..f2466de358dec 100644
--- a/core/src/main/scala/kafka/log/Log.scala
+++ b/core/src/main/scala/kafka/log/Log.scala
@@ -34,15 +34,16 @@ import kafka.message.{BrokerCompressionCodec, CompressionCodec, NoCompressionCod
 import kafka.metrics.KafkaMetricsGroup
 import kafka.server.checkpoints.LeaderEpochCheckpointFile
 import kafka.server.epoch.LeaderEpochFileCache
-import kafka.server.{BrokerTopicStats, FetchDataInfo, LogDirFailureChannel, LogOffsetMetadata, OffsetAndEpoch}
+import kafka.server.{BrokerTopicStats, FetchDataInfo, FetchHighWatermark, FetchIsolation, FetchLogEnd, FetchTxnCommitted, LogDirFailureChannel, LogOffsetMetadata, OffsetAndEpoch}
 import kafka.utils._
 import org.apache.kafka.common.errors._
 import org.apache.kafka.common.record.FileRecords.TimestampAndOffset
 import org.apache.kafka.common.record._
 import org.apache.kafka.common.requests.FetchResponse.AbortedTransaction
+import org.apache.kafka.common.requests.ProduceResponse.RecordError
 import org.apache.kafka.common.requests.{EpochEndOffset, ListOffsetRequest}
 import org.apache.kafka.common.utils.{Time, Utils}
-import org.apache.kafka.common.{KafkaException, TopicPartition}
+import org.apache.kafka.common.{InvalidRecordException, KafkaException, TopicPartition}
 
 import scala.collection.JavaConverters._
 import scala.collection.mutable.{ArrayBuffer, ListBuffer}
@@ -54,7 +55,18 @@ object LogAppendInfo {
 
   def unknownLogAppendInfoWithLogStartOffset(logStartOffset: Long): LogAppendInfo =
     LogAppendInfo(None, -1, RecordBatch.NO_TIMESTAMP, -1L, RecordBatch.NO_TIMESTAMP, logStartOffset,
-      RecordConversionStats.EMPTY, NoCompressionCodec, NoCompressionCodec, -1, -1, offsetsMonotonic = false, -1L)
+      RecordConversionStats.EMPTY, NoCompressionCodec, NoCompressionCodec, -1, -1,
+      offsetsMonotonic = false, -1L)
+
+  /**
+   * In ProduceResponse V8+, we add two new fields record_errors and error_message (see KIP-467).
+   * For any record failures with InvalidTimestamp or InvalidRecordException, we construct a LogAppendInfo object like the one
+   * in unknownLogAppendInfoWithLogStartOffset, but with additiona fields recordErrors and errorMessage
+   */
+  def unknownLogAppendInfoWithAdditionalInfo(logStartOffset: Long, recordErrors: List[RecordError], errorMessage: String): LogAppendInfo =
+    LogAppendInfo(None, -1, RecordBatch.NO_TIMESTAMP, -1L, RecordBatch.NO_TIMESTAMP, logStartOffset,
+      RecordConversionStats.EMPTY, NoCompressionCodec, NoCompressionCodec, -1, -1,
+      offsetsMonotonic = false, -1L, recordErrors, errorMessage)
 }
 
 /**
@@ -87,7 +99,10 @@ case class LogAppendInfo(var firstOffset: Option[Long],
                          shallowCount: Int,
                          validBytes: Int,
                          offsetsMonotonic: Boolean,
-                         lastOffsetOfFirstBatch: Long) {
+                         lastOffsetOfFirstBatch: Long,
+                         recordErrors: List[RecordError] = List(),
+                         errorMessage: String = null,
+                         var recompressedBatchCount: Long = 0) {
   /**
    * Get the first offset if it exists, else get the last offset of the first batch
    * For magic versions 2 and newer, this method will return first offset. For magic versions
@@ -159,7 +174,7 @@ case class RollParams(maxSegmentMs: Long,
 
 object RollParams {
   def apply(config: LogConfig, appendInfo: LogAppendInfo, messagesSize: Int, now: Long): RollParams = {
-   new RollParams(config.segmentMs,
+   new RollParams(config.maxSegmentMs,
      config.segmentSize,
      appendInfo.maxTimestamp,
      appendInfo.lastOffset,
@@ -175,7 +190,7 @@ object RollParams {
  * New log segments are created according to a configurable policy that controls the size in bytes or time interval
  * for a given segment.
  *
- * @param dir The directory in which log segments are created.
+ * @param _dir The directory in which log segments are created.
  * @param config The log configuration settings
  * @param logStartOffset The earliest offset allowed to be exposed to kafka client.
  *                       The logStartOffset can be updated by :
@@ -196,7 +211,7 @@ object RollParams {
  * @param producerIdExpirationCheckIntervalMs How often to check for producer ids which need to be expired
  */
 @threadsafe
-class Log(@volatile var dir: File,
+class Log(@volatile private var _dir: File,
           @volatile var config: LogConfig,
           @volatile var logStartOffset: Long,
           @volatile var recoveryPoint: Long,
@@ -215,44 +230,17 @@ class Log(@volatile var dir: File,
 
   /* A lock that guards all modifications to the log */
   private val lock = new Object
+
   // The memory mapped buffer for index files of this log will be closed with either delete() or closeHandlers()
   // After memory mapped buffer is closed, no disk IO operation should be performed for this log
   @volatile private var isMemoryMappedBufferClosed = false
 
+  // Cache value of parent directory to avoid allocations in hot paths like ReplicaManager.checkpointHighWatermarks
+  @volatile private var _parentDir: String = dir.getParent
+
   /* last time it was flushed */
   private val lastFlushedTime = new AtomicLong(time.milliseconds)
 
-  def initFileSize: Int = {
-    if (config.preallocate)
-      config.segmentSize
-    else
-      0
-  }
-
-  def updateConfig(updatedKeys: Set[String], newConfig: LogConfig): Unit = {
-    if ((updatedKeys.contains(LogConfig.RetentionMsProp)
-      || updatedKeys.contains(LogConfig.MessageTimestampDifferenceMaxMsProp))
-      && topicPartition.partition == 0  // generate warnings only for one partition of each topic
-      && newConfig.retentionMs < newConfig.messageTimestampDifferenceMaxMs)
-      warn(s"${LogConfig.RetentionMsProp} for topic ${topicPartition.topic} is set to ${newConfig.retentionMs}. It is smaller than " +
-        s"${LogConfig.MessageTimestampDifferenceMaxMsProp}'s value ${newConfig.messageTimestampDifferenceMaxMs}. " +
-        s"This may result in frequent log rolling.")
-    val oldConfig = this.config
-    this.config = newConfig
-    if (updatedKeys.contains(LogConfig.MessageFormatVersionProp)) {
-      val oldRecordVersion = oldConfig.messageFormatVersion.recordVersion
-      val newRecordVersion = newConfig.messageFormatVersion.recordVersion
-      if (newRecordVersion.precedes(oldRecordVersion))
-        warn(s"Record format version has been downgraded from $oldRecordVersion to $newRecordVersion.")
-      initializeLeaderEpochCache()
-    }
-  }
-
-  private def checkIfMemoryMappedBufferClosed(): Unit = {
-    if (isMemoryMappedBufferClosed)
-      throw new KafkaStorageException(s"The memory mapped buffer for log of $topicPartition is already closed")
-  }
-
   @volatile private var nextOffsetMetadata: LogOffsetMetadata = _
 
   /* The earliest offset which is part of an incomplete transaction. This is used to compute the
@@ -266,14 +254,14 @@ class Log(@volatile var dir: File,
    * that this could result in disagreement between replicas depending on when they began replicating the log.
    * In the worst case, the LSO could be seen by a consumer to go backwards.
    */
-  @volatile var firstUnstableOffset: Option[LogOffsetMetadata] = None
+  @volatile private var firstUnstableOffsetMetadata: Option[LogOffsetMetadata] = None
 
   /* Keep track of the current high watermark in order to ensure that segments containing offsets at or above it are
    * not eligible for deletion. This means that the active segment is only eligible for deletion if the high watermark
    * equals the log end offset (which may never happen for a partition under consistent load). This is needed to
    * prevent the log start offset (which is exposed in fetch responses) from getting ahead of the high watermark.
    */
-  @volatile private var replicaHighWatermark: Option[Long] = None
+  @volatile private var highWatermarkMetadata: LogOffsetMetadata = LogOffsetMetadata(logStartOffset)
 
   /* the actual segments of the log */
   private val segments: ConcurrentNavigableMap[java.lang.Long, LogSegment] = new ConcurrentSkipListMap[java.lang.Long, LogSegment]
@@ -292,11 +280,11 @@ class Log(@volatile var dir: File,
     val nextOffset = loadSegments()
 
     /* Calculate the offset of the next message */
-    nextOffsetMetadata = new LogOffsetMetadata(nextOffset, activeSegment.baseOffset, activeSegment.size)
+    nextOffsetMetadata = LogOffsetMetadata(nextOffset, activeSegment.baseOffset, activeSegment.size)
 
     leaderEpochCache.foreach(_.truncateFromEnd(nextOffsetMetadata.messageOffset))
 
-    logStartOffset = math.max(logStartOffset, segments.firstEntry.getValue.baseOffset)
+    updateLogStartOffset(math.max(logStartOffset, segments.firstEntry.getValue.baseOffset))
 
     // The earliest leader epoch may not be flushed during a hard failure. Recover it here.
     leaderEpochCache.foreach(_.truncateFromStart(logStartOffset))
@@ -311,6 +299,181 @@ class Log(@volatile var dir: File,
       s"log end offset $logEndOffset in ${time.milliseconds() - startMs} ms")
   }
 
+  def dir: File = _dir
+
+  def parentDir: String = _parentDir
+
+  def parentDirFile: File = new File(_parentDir)
+
+  def initFileSize: Int = {
+    if (config.preallocate)
+      config.segmentSize
+    else
+      0
+  }
+
+  def updateConfig(newConfig: LogConfig): Unit = {
+    val oldConfig = this.config
+    this.config = newConfig
+    val oldRecordVersion = oldConfig.messageFormatVersion.recordVersion
+    val newRecordVersion = newConfig.messageFormatVersion.recordVersion
+    if (newRecordVersion.precedes(oldRecordVersion))
+      warn(s"Record format version has been downgraded from $oldRecordVersion to $newRecordVersion.")
+    if (newRecordVersion.value != oldRecordVersion.value)
+      initializeLeaderEpochCache()
+  }
+
+  private def checkIfMemoryMappedBufferClosed(): Unit = {
+    if (isMemoryMappedBufferClosed)
+      throw new KafkaStorageException(s"The memory mapped buffer for log of $topicPartition is already closed")
+  }
+
+  def highWatermark: Long = highWatermarkMetadata.messageOffset
+
+  /**
+   * Update the high watermark to a new offset. The new high watermark will be lower
+   * bounded by the log start offset and upper bounded by the log end offset.
+   *
+   * This is intended to be called when initializing the high watermark or when updating
+   * it on a follower after receiving a Fetch response from the leader.
+   *
+   * @param hw the suggested new value for the high watermark
+   * @return the updated high watermark offset
+   */
+  def updateHighWatermark(hw: Long): Long = {
+    val newHighWatermark = if (hw < logStartOffset)
+      logStartOffset
+    else if (hw > logEndOffset)
+      logEndOffset
+    else
+      hw
+    updateHighWatermarkMetadata(LogOffsetMetadata(newHighWatermark))
+    newHighWatermark
+  }
+
+  /**
+   * Update the high watermark to a new value if and only if it is larger than the old value. It is
+   * an error to update to a value which is larger than the log end offset.
+   *
+   * This method is intended to be used by the leader to update the high watermark after follower
+   * fetch offsets have been updated.
+   *
+   * @return the old high watermark, if updated by the new value
+   */
+  def maybeIncrementHighWatermark(newHighWatermark: LogOffsetMetadata): Option[LogOffsetMetadata] = {
+    if (newHighWatermark.messageOffset > logEndOffset)
+      throw new IllegalArgumentException(s"High watermark $newHighWatermark update exceeds current " +
+        s"log end offset $logEndOffsetMetadata")
+
+    lock.synchronized {
+      val oldHighWatermark = fetchHighWatermarkMetadata
+
+      // Ensure that the high watermark increases monotonically. We also update the high watermark when the new
+      // offset metadata is on a newer segment, which occurs whenever the log is rolled to a new segment.
+      if (oldHighWatermark.messageOffset < newHighWatermark.messageOffset ||
+        (oldHighWatermark.messageOffset == newHighWatermark.messageOffset && oldHighWatermark.onOlderSegment(newHighWatermark))) {
+        updateHighWatermarkMetadata(newHighWatermark)
+        Some(oldHighWatermark)
+      } else {
+        None
+      }
+    }
+  }
+
+  /**
+   * Get the offset and metadata for the current high watermark. If offset metadata is not
+   * known, this will do a lookup in the index and cache the result.
+   */
+  private def fetchHighWatermarkMetadata: LogOffsetMetadata = {
+    checkIfMemoryMappedBufferClosed()
+
+    val offsetMetadata = highWatermarkMetadata
+    if (offsetMetadata.messageOffsetOnly) {
+      lock.synchronized {
+        val fullOffset = convertToOffsetMetadataOrThrow(highWatermark)
+        updateHighWatermarkMetadata(fullOffset)
+        fullOffset
+      }
+    } else {
+      offsetMetadata
+    }
+  }
+
+  private def updateHighWatermarkMetadata(newHighWatermark: LogOffsetMetadata): Unit = {
+    if (newHighWatermark.messageOffset < 0)
+      throw new IllegalArgumentException("High watermark offset should be non-negative")
+
+    lock synchronized {
+      highWatermarkMetadata = newHighWatermark
+      producerStateManager.onHighWatermarkUpdated(newHighWatermark.messageOffset)
+      maybeIncrementFirstUnstableOffset()
+    }
+    trace(s"Setting high watermark $newHighWatermark")
+  }
+
+  /**
+   * Get the first unstable offset. Unlike the last stable offset, which is always defined,
+   * the first unstable offset only exists if there are transactions in progress.
+   *
+   * @return the first unstable offset, if it exists
+   */
+  private[log] def firstUnstableOffset: Option[Long] = firstUnstableOffsetMetadata.map(_.messageOffset)
+
+  private def fetchLastStableOffsetMetadata: LogOffsetMetadata = {
+    checkIfMemoryMappedBufferClosed()
+
+    // cache the current high watermark to avoid a concurrent update invalidating the range check
+    val highWatermarkMetadata = fetchHighWatermarkMetadata
+
+    firstUnstableOffsetMetadata match {
+      case Some(offsetMetadata) if offsetMetadata.messageOffset < highWatermarkMetadata.messageOffset =>
+        if (offsetMetadata.messageOffsetOnly) {
+          lock synchronized {
+            val fullOffset = convertToOffsetMetadataOrThrow(offsetMetadata.messageOffset)
+            if (firstUnstableOffsetMetadata.contains(offsetMetadata))
+              firstUnstableOffsetMetadata = Some(fullOffset)
+            fullOffset
+          }
+        } else {
+          offsetMetadata
+        }
+      case _ => highWatermarkMetadata
+    }
+  }
+
+  /**
+   * The last stable offset (LSO) is defined as the first offset such that all lower offsets have been "decided."
+   * Non-transactional messages are considered decided immediately, but transactional messages are only decided when
+   * the corresponding COMMIT or ABORT marker is written. This implies that the last stable offset will be equal
+   * to the high watermark if there are no transactional messages in the log. Note also that the LSO cannot advance
+   * beyond the high watermark.
+   */
+  def lastStableOffset: Long = {
+    firstUnstableOffsetMetadata match {
+      case Some(offsetMetadata) if offsetMetadata.messageOffset < highWatermark => offsetMetadata.messageOffset
+      case _ => highWatermark
+    }
+  }
+
+  def lastStableOffsetLag: Long = highWatermark - lastStableOffset
+
+  /**
+    * Fully materialize and return an offset snapshot including segment position info. This method will update
+    * the LogOffsetMetadata for the high watermark and last stable offset if they are message-only. Throws an
+    * offset out of range error if the segment info cannot be loaded.
+    */
+  def fetchOffsetSnapshot: LogOffsetSnapshot = {
+    val lastStable = fetchLastStableOffsetMetadata
+    val highWatermark = fetchHighWatermarkMetadata
+
+    LogOffsetSnapshot(
+      logStartOffset,
+      logEndOffsetMetadata,
+      highWatermark,
+      lastStable
+    )
+  }
+
   private val tags = {
     val maybeFutureTag = if (isFuture) Map("is-future" -> "true") else Map.empty[String, String]
     Map("topic" -> topicPartition.topic, "partition" -> topicPartition.partition.toString) ++ maybeFutureTag
@@ -340,7 +503,7 @@ class Log(@volatile var dir: File,
     },
     tags)
 
-  scheduler.schedule(name = "PeriodicProducerExpirationCheck", fun = () => {
+  val producerExpireCheck = scheduler.schedule(name = "PeriodicProducerExpirationCheck", fun = () => {
     lock synchronized {
       producerStateManager.removeExpiredProducers(time.milliseconds)
     }
@@ -570,17 +733,7 @@ class Log(@volatile var dir: File,
     // before the swap file is restored as the new segment file.
     completeSwapOperations(swapFiles)
 
-    if (logSegments.isEmpty) {
-      // no existing segments, create a new mutable segment beginning at offset 0
-      addSegment(LogSegment.open(dir = dir,
-        baseOffset = 0,
-        config,
-        time = time,
-        fileAlreadyExists = false,
-        initFileSize = this.initFileSize,
-        preallocate = config.preallocate))
-      0
-    } else if (!dir.getAbsolutePath.endsWith(Log.DeleteDirSuffix)) {
+    if (!dir.getAbsolutePath.endsWith(Log.DeleteDirSuffix)) {
       val nextOffset = retryOnOffsetOverflow {
         recoverLog()
       }
@@ -588,11 +741,44 @@ class Log(@volatile var dir: File,
       // reset the index size of the currently active log segment to allow more entries
       activeSegment.resizeIndexes(config.maxIndexSize)
       nextOffset
-    } else 0
+    } else {
+       if (logSegments.isEmpty) {
+          addSegment(LogSegment.open(dir = dir,
+            baseOffset = 0,
+            config,
+            time = time,
+            fileAlreadyExists = false,
+            initFileSize = this.initFileSize,
+            preallocate = false))
+       }
+      0
+    }
+  }
+
+  private def updateLogEndOffset(offset: Long): Unit = {
+    nextOffsetMetadata = LogOffsetMetadata(offset, activeSegment.baseOffset, activeSegment.size)
+
+    // Update the high watermark in case it has gotten ahead of the log end offset following a truncation
+    // or if a new segment has been rolled and the offset metadata needs to be updated.
+    if (highWatermark >= offset) {
+      updateHighWatermarkMetadata(nextOffsetMetadata)
+    }
+
+    if (this.recoveryPoint > offset) {
+      this.recoveryPoint = offset
+    }
   }
 
-  private def updateLogEndOffset(messageOffset: Long) {
-    nextOffsetMetadata = new LogOffsetMetadata(messageOffset, activeSegment.baseOffset, activeSegment.size)
+  private def updateLogStartOffset(offset: Long): Unit = {
+    logStartOffset = offset
+
+    if (highWatermark < offset) {
+      updateHighWatermark(offset)
+    }
+
+    if (this.recoveryPoint < offset) {
+      this.recoveryPoint = offset
+    }
   }
 
   /**
@@ -605,8 +791,10 @@ class Log(@volatile var dir: File,
     // if we have the clean shutdown marker, skip recovery
     if (!hasCleanShutdownFile) {
       // okay we need to actually recover this log
-      val unflushed = logSegments(this.recoveryPoint, Long.MaxValue).iterator
-      while (unflushed.hasNext) {
+      val unflushed = logSegments(this.recoveryPoint, Long.MaxValue).toIterator
+      var truncated = false
+
+      while (unflushed.hasNext && !truncated) {
         val segment = unflushed.next
         info(s"Recovering unflushed segment ${segment.baseOffset}")
         val truncatedBytes =
@@ -622,10 +810,32 @@ class Log(@volatile var dir: File,
         if (truncatedBytes > 0) {
           // we had an invalid message, delete all remaining log
           warn(s"Corruption found in segment ${segment.baseOffset}, truncating to offset ${segment.readNextOffset}")
-          unflushed.foreach(deleteSegment)
+          removeAndDeleteSegments(unflushed.toList, asyncDelete = true)
+          truncated = true
         }
       }
     }
+
+    if (logSegments.nonEmpty) {
+      val logEndOffset = activeSegment.readNextOffset
+      if (logEndOffset < logStartOffset) {
+        warn(s"Deleting all segments because logEndOffset ($logEndOffset) is smaller than logStartOffset ($logStartOffset). " +
+          "This could happen if segment files were deleted from the file system.")
+        removeAndDeleteSegments(logSegments, asyncDelete = true)
+      }
+    }
+
+    if (logSegments.isEmpty) {
+      // no existing segments, create a new mutable segment beginning at logStartOffset
+      addSegment(LogSegment.open(dir = dir,
+        baseOffset = logStartOffset,
+        config,
+        time = time,
+        fileAlreadyExists = false,
+        initFileSize = this.initFileSize,
+        preallocate = config.preallocate))
+    }
+
     recoveryPoint = activeSegment.readNextOffset
     recoveryPoint
   }
@@ -679,6 +889,8 @@ class Log(@volatile var dir: File,
       // and we can skip the loading. This is an optimization for users which are not yet using
       // idempotent/transactional features yet.
       if (lastOffset > producerStateManager.mapEndOffset && !isEmptyBeforeTruncation) {
+        val segmentOfLastOffset = floorLogSegment(lastOffset)
+
         logSegments(producerStateManager.mapEndOffset, lastOffset).foreach { segment =>
           val startOffset = Utils.max(segment.baseOffset, producerStateManager.mapEndOffset, logStartOffset)
           producerStateManager.updateMapEndOffset(startOffset)
@@ -686,7 +898,18 @@ class Log(@volatile var dir: File,
           if (offsetsToSnapshot.contains(Some(segment.baseOffset)))
             producerStateManager.takeSnapshot()
 
-          val fetchDataInfo = segment.read(startOffset, Some(lastOffset), Int.MaxValue)
+          val maxPosition = if (segmentOfLastOffset.contains(segment)) {
+            Option(segment.translateOffset(lastOffset))
+              .map(_.position)
+              .getOrElse(segment.size)
+          } else {
+            segment.size
+          }
+
+          val fetchDataInfo = segment.read(startOffset,
+            maxSize = Int.MaxValue,
+            maxPosition = maxPosition,
+            minOneMessage = false)
           if (fetchDataInfo != null)
             loadProducersFromLog(producerStateManager, fetchDataInfo.records)
         }
@@ -698,7 +921,7 @@ class Log(@volatile var dir: File,
 
   private def loadProducerState(lastOffset: Long, reloadFromCleanShutdown: Boolean): Unit = lock synchronized {
     rebuildProducerState(lastOffset, reloadFromCleanShutdown, producerStateManager)
-    updateFirstUnstableOffset()
+    maybeIncrementFirstUnstableOffset()
   }
 
   private def loadProducersFromLog(producerStateManager: ProducerStateManager, records: Records): Unit = {
@@ -706,7 +929,10 @@ class Log(@volatile var dir: File,
     val completedTxns = ListBuffer.empty[CompletedTxn]
     records.batches.asScala.foreach { batch =>
       if (batch.hasProducerId) {
-        val maybeCompletedTxn = updateProducers(batch, loadedProducers, isFromClient = false)
+        val maybeCompletedTxn = updateProducers(batch,
+          loadedProducers,
+          firstOffsetMetadata = None,
+          origin = AppendOrigin.Replication)
         maybeCompletedTxn.foreach(completedTxns += _)
       }
     }
@@ -720,6 +946,14 @@ class Log(@volatile var dir: File,
     }
   }
 
+  private[log] def lastRecordsOfActiveProducers: Map[Long, LastRecord] = lock synchronized {
+    producerStateManager.activeProducers.map { case (producerId, producerIdEntry) =>
+      val lastDataOffset = if (producerIdEntry.lastDataOffset >= 0 ) Some(producerIdEntry.lastDataOffset) else None
+      val lastRecord = LastRecord(lastDataOffset, producerIdEntry.producerEpoch)
+      producerId -> lastRecord
+    }
+  }
+
   /**
    * Check if we have the "clean shutdown" file
    */
@@ -735,10 +969,11 @@ class Log(@volatile var dir: File,
    * Close this log.
    * The memory mapped buffer for index files of this log will be left open until the log is deleted.
    */
-  def close() {
+  def close(): Unit = {
     debug("Closing log")
     lock synchronized {
       checkIfMemoryMappedBufferClosed()
+      producerExpireCheck.cancel(true)
       maybeHandleIOException(s"Error while renaming dir for $topicPartition in dir ${dir.getParent}") {
         // We take a snapshot at the last written offset to hopefully avoid the need to scan the log
         // after restarting and to ensure that we cannot inadvertently hit the upgrade optimization
@@ -754,14 +989,15 @@ class Log(@volatile var dir: File,
    *
    * @throws KafkaStorageException if rename fails
    */
-  def renameDir(name: String) {
+  def renameDir(name: String): Unit = {
     lock synchronized {
       maybeHandleIOException(s"Error while renaming dir for $topicPartition in log dir ${dir.getParent}") {
         val renamedDir = new File(dir.getParent, name)
         Utils.atomicMoveWithFallback(dir.toPath, renamedDir.toPath)
         if (renamedDir != dir) {
-          dir = renamedDir
-          logSegments.foreach(_.updateDir(renamedDir))
+          _dir = renamedDir
+          _parentDir = renamedDir.getParent
+          logSegments.foreach(_.updateParentDir(renamedDir))
           producerStateManager.logDir = dir
           // re-initialize leader epoch cache so that LeaderEpochCheckpointFile.checkpoint can correctly reference
           // the checkpoint file in renamed log directory
@@ -774,7 +1010,7 @@ class Log(@volatile var dir: File,
   /**
    * Close file handlers used by log but don't write to disk. This is called if the log directory is offline
    */
-  def closeHandlers() {
+  def closeHandlers(): Unit = {
     debug("Closing handlers")
     lock synchronized {
       logSegments.foreach(_.closeHandlers())
@@ -786,14 +1022,16 @@ class Log(@volatile var dir: File,
    * Append this message set to the active segment of the log, assigning offsets and Partition Leader Epochs
    *
    * @param records The records to append
-   * @param isFromClient Whether or not this append is from a producer
+   * @param origin Declares the origin of the append which affects required validations
    * @param interBrokerProtocolVersion Inter-broker message protocol version
    * @throws KafkaStorageException If the append fails due to an I/O error.
    * @return Information about the appended messages including the first and last offset.
    */
-  def appendAsLeader(records: MemoryRecords, leaderEpoch: Int, isFromClient: Boolean = true,
+  def appendAsLeader(records: MemoryRecords,
+                     leaderEpoch: Int,
+                     origin: AppendOrigin = AppendOrigin.Client,
                      interBrokerProtocolVersion: ApiVersion = ApiVersion.latestVersion): LogAppendInfo = {
-    append(records, isFromClient, interBrokerProtocolVersion, assignOffsets = true, leaderEpoch)
+    append(records, origin, interBrokerProtocolVersion, assignOffsets = true, leaderEpoch)
   }
 
   /**
@@ -804,9 +1042,85 @@ class Log(@volatile var dir: File,
    * @return Information about the appended messages including the first and last offset.
    */
   def appendAsFollower(records: MemoryRecords): LogAppendInfo = {
-    append(records, isFromClient = false, interBrokerProtocolVersion = ApiVersion.latestVersion, assignOffsets = false, leaderEpoch = -1)
+    append(records,
+      origin = AppendOrigin.Replication,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      assignOffsets = false,
+      leaderEpoch = -1)
   }
 
+  /**
+   *  Append records to the active segment of the log.
+   *
+   * This method accepts a MemoryRecords with a buffer that contains one or more valid MemoryRecords and calls
+   * appendSingleBatch for each individually. This allows multiple batches to be aggregated in a single
+   * record_set, but still be processed without forcing offset assignment.
+   *
+   * @param records The log records to append
+   * @param origin Declares the origin of the append which affects required validations
+   * @param interBrokerProtocolVersion Inter-broker message protocol version
+   * @param assignOffsets Should the log assign offsets to this message set or blindly apply what it is given
+   * @param leaderEpoch The partition's leader epoch which will be applied to messages when offsets are assigned on the leader
+   * @throws KafkaStorageException If the append fails due to an I/O error.
+   * @throws OffsetsOutOfOrderException If out of order offsets found in 'records'
+   * @throws UnexpectedAppendOffsetException If the first or last offset in append is less than next offset
+   * @return Information about the appended messages including the first and last offset.
+   */
+  private def append(records: MemoryRecords,
+                                origin: AppendOrigin,
+                                interBrokerProtocolVersion: ApiVersion,
+                                assignOffsets: Boolean,
+                                leaderEpoch: Int): LogAppendInfo = {
+    val batches = records.batches()
+
+    if (batches.asScala.isEmpty) {
+      return analyzeAndValidateRecords(records, origin)
+    }
+    val remainingBytes = records.buffer().duplicate()
+    batches.asScala.map(b => {
+      // Create the current record set using only the first batch
+      val batchSize = b.sizeInBytes()
+      val batchBuffer = remainingBytes.slice()
+      batchBuffer.limit(batchSize)
+      val batchRecords = MemoryRecords.readableRecords(batchBuffer)
+
+      // Advance the position in remaining bytes
+      remainingBytes.position(remainingBytes.position() + batchSize)
+
+      // Append the single record batch to the log
+      appendSingleBatch(batchRecords, origin, interBrokerProtocolVersion, assignOffsets, leaderEpoch)
+    }).reduceLeft((info1, info2) => {
+      // Don't assume that the max timestamp is always in the later batch
+      var maxTimestamp = info1.maxTimestamp
+      var offsetofMaxTimestamp = info1.offsetOfMaxTimestamp
+      if (info2.maxTimestamp > info1.maxTimestamp) {
+        maxTimestamp = info2.maxTimestamp
+        offsetofMaxTimestamp = info2.offsetOfMaxTimestamp
+      }
+      var combinedRecordConversionStats = info1.recordConversionStats
+      combinedRecordConversionStats.add(info2.recordConversionStats)
+
+      // Combine the LogAppendInfo to maintain an overall result to return
+      LogAppendInfo(
+        info1.firstOffset,
+        info2.lastOffset,
+        maxTimestamp,
+        offsetofMaxTimestamp,
+        info1.logAppendTime,
+        info1.logStartOffset,
+        combinedRecordConversionStats,
+        info1.sourceCodec,
+        info1.targetCodec,
+        info1.shallowCount + info2.shallowCount,
+        info1.validBytes + info2.validBytes,
+        info1.offsetsMonotonic && info2.offsetsMonotonic,
+        info1.lastOffsetOfFirstBatch,
+        info1.recordErrors ++ info2.recordErrors,
+        info1.errorMessage + info2.errorMessage,
+        info1.recompressedBatchCount + info2.recompressedBatchCount
+      )
+    })
+  }
   /**
    * Append this message set to the active segment of the log, rolling over to a fresh segment if necessary.
    *
@@ -814,7 +1128,7 @@ class Log(@volatile var dir: File,
    * however if the assignOffsets=false flag is passed we will only check that the existing offsets are valid.
    *
    * @param records The log records to append
-   * @param isFromClient Whether or not this append is from a producer
+   * @param origin Declares the origin of the append which affects required validations
    * @param interBrokerProtocolVersion Inter-broker message protocol version
    * @param assignOffsets Should the log assign offsets to this message set or blindly apply what it is given
    * @param leaderEpoch The partition's leader epoch which will be applied to messages when offsets are assigned on the leader
@@ -823,9 +1137,13 @@ class Log(@volatile var dir: File,
    * @throws UnexpectedAppendOffsetException If the first or last offset in append is less than next offset
    * @return Information about the appended messages including the first and last offset.
    */
-  private def append(records: MemoryRecords, isFromClient: Boolean, interBrokerProtocolVersion: ApiVersion, assignOffsets: Boolean, leaderEpoch: Int): LogAppendInfo = {
+  private def appendSingleBatch(records: MemoryRecords,
+                     origin: AppendOrigin,
+                     interBrokerProtocolVersion: ApiVersion,
+                     assignOffsets: Boolean,
+                     leaderEpoch: Int): LogAppendInfo = {
     maybeHandleIOException(s"Error while appending records to $topicPartition in dir ${dir.getParent}") {
-      val appendInfo = analyzeAndValidateRecords(records, isFromClient = isFromClient)
+      val appendInfo = analyzeAndValidateRecords(records, origin)
 
       // return if we have no valid messages or if this is a duplicate of the last appended entry
       if (appendInfo.shallowCount == 0)
@@ -844,6 +1162,7 @@ class Log(@volatile var dir: File,
           val now = time.milliseconds
           val validateAndOffsetAssignResult = try {
             LogValidator.validateMessagesAndAssignOffsets(validRecords,
+              topicPartition,
               offset,
               time,
               now,
@@ -854,8 +1173,10 @@ class Log(@volatile var dir: File,
               config.messageTimestampType,
               config.messageTimestampDifferenceMaxMs,
               leaderEpoch,
-              isFromClient,
-              interBrokerProtocolVersion)
+              origin,
+              interBrokerProtocolVersion,
+              brokerTopicStats,
+              config.producerBatchDecompressionEnable)
           } catch {
             case e: IOException =>
               throw new KafkaException(s"Error validating messages while appending to log $name", e)
@@ -865,6 +1186,9 @@ class Log(@volatile var dir: File,
           appendInfo.offsetOfMaxTimestamp = validateAndOffsetAssignResult.shallowOffsetOfMaxTimestamp
           appendInfo.lastOffset = offset.value - 1
           appendInfo.recordConversionStats = validateAndOffsetAssignResult.recordConversionStats
+          // update stats for compressed/decompressed batch count
+          if (validateAndOffsetAssignResult.recompressApplied)
+            appendInfo.recompressedBatchCount = 1
           if (config.messageTimestampType == TimestampType.LOG_APPEND_TIME)
             appendInfo.logAppendTime = now
 
@@ -909,8 +1233,17 @@ class Log(@volatile var dir: File,
 
         // update the epoch cache with the epoch stamped onto the message by the leader
         validRecords.batches.asScala.foreach { batch =>
-          if (batch.magic >= RecordBatch.MAGIC_VALUE_V2)
+          if (batch.magic >= RecordBatch.MAGIC_VALUE_V2) {
             maybeAssignEpochStartOffset(batch.partitionLeaderEpoch, batch.baseOffset)
+          } else {
+            // In partial upgrade scenarios, we may get a temporary regression to the message format. In
+            // order to ensure the safety of leader election, we clear the epoch cache so that we revert
+            // to truncation by high watermark after the next leader election.
+            leaderEpochCache.filter(_.nonEmpty).foreach { cache =>
+              warn(s"Clearing leader epoch cache after unexpected append with message format v${batch.magic}")
+              cache.clearAndFlush()
+            }
+          }
         }
 
         // check messages set size may be exceed config.segmentSize
@@ -919,9 +1252,19 @@ class Log(@volatile var dir: File,
             s"to partition $topicPartition, which exceeds the maximum configured segment size of ${config.segmentSize}.")
         }
 
+        // maybe roll the log if this segment is full
+        val segment = maybeRoll(validRecords.sizeInBytes, appendInfo)
+
+        val logOffsetMetadata = LogOffsetMetadata(
+          messageOffset = appendInfo.firstOrLastOffsetOfFirstBatch,
+          segmentBaseOffset = segment.baseOffset,
+          relativePositionInSegment = segment.size)
+
         // now that we have valid records, offsets assigned, and timestamps updated, we need to
         // validate the idempotent/transactional state of the producers and collect some metadata
-        val (updatedProducers, completedTxns, maybeDuplicate) = analyzeAndValidateProducerState(validRecords, isFromClient)
+        val (updatedProducers, completedTxns, maybeDuplicate) = analyzeAndValidateProducerState(
+          logOffsetMetadata, validRecords, origin)
+
         maybeDuplicate.foreach { duplicate =>
           appendInfo.firstOffset = Some(duplicate.firstOffset)
           appendInfo.lastOffset = duplicate.lastOffset
@@ -930,41 +1273,38 @@ class Log(@volatile var dir: File,
           return appendInfo
         }
 
-        // maybe roll the log if this segment is full
-        val segment = maybeRoll(validRecords.sizeInBytes, appendInfo)
-
-        val logOffsetMetadata = LogOffsetMetadata(
-          messageOffset = appendInfo.firstOrLastOffsetOfFirstBatch,
-          segmentBaseOffset = segment.baseOffset,
-          relativePositionInSegment = segment.size)
-
         segment.append(largestOffset = appendInfo.lastOffset,
           largestTimestamp = appendInfo.maxTimestamp,
           shallowOffsetOfMaxTimestamp = appendInfo.offsetOfMaxTimestamp,
           records = validRecords)
 
+        // Increment the log end offset. We do this immediately after the append because a
+        // write to the transaction index below may fail and we want to ensure that the offsets
+        // of future appends still grow monotonically. The resulting transaction index inconsistency
+        // will be cleaned up after the log directory is recovered. Note that the end offset of the
+        // ProducerStateManager will not be updated and the last stable offset will not advance
+        // if the append to the transaction index fails.
+        updateLogEndOffset(appendInfo.lastOffset + 1)
+
         // update the producer state
-        for ((_, producerAppendInfo) <- updatedProducers) {
-          producerAppendInfo.maybeCacheTxnFirstOffsetMetadata(logOffsetMetadata)
+        for (producerAppendInfo <- updatedProducers.values) {
           producerStateManager.update(producerAppendInfo)
         }
 
         // update the transaction index with the true last stable offset. The last offset visible
         // to consumers using READ_COMMITTED will be limited by this value and the high watermark.
         for (completedTxn <- completedTxns) {
-          val lastStableOffset = producerStateManager.completeTxn(completedTxn)
+          val lastStableOffset = producerStateManager.lastStableOffset(completedTxn)
           segment.updateTxnIndex(completedTxn, lastStableOffset)
+          producerStateManager.completeTxn(completedTxn)
         }
 
         // always update the last producer id map offset so that the snapshot reflects the current offset
         // even if there isn't any idempotent data being written
         producerStateManager.updateMapEndOffset(appendInfo.lastOffset + 1)
 
-        // increment the log end offset
-        updateLogEndOffset(appendInfo.lastOffset + 1)
-
         // update the first unstable offset (which is used to compute LSO)
-        updateFirstUnstableOffset()
+        maybeIncrementFirstUnstableOffset()
 
         trace(s"Appended message set with last offset: ${appendInfo.lastOffset}, " +
           s"first offset: ${appendInfo.firstOffset}, " +
@@ -997,35 +1337,30 @@ class Log(@volatile var dir: File,
     }
   }
 
-  def onHighWatermarkIncremented(highWatermark: Long): Unit = {
-    lock synchronized {
-      replicaHighWatermark = Some(highWatermark)
-      producerStateManager.onHighWatermarkUpdated(highWatermark)
-      updateFirstUnstableOffset()
-    }
-  }
-
-  private def updateFirstUnstableOffset(): Unit = lock synchronized {
+  private def maybeIncrementFirstUnstableOffset(): Unit = lock synchronized {
     checkIfMemoryMappedBufferClosed()
+
     val updatedFirstStableOffset = producerStateManager.firstUnstableOffset match {
       case Some(logOffsetMetadata) if logOffsetMetadata.messageOffsetOnly || logOffsetMetadata.messageOffset < logStartOffset =>
         val offset = math.max(logOffsetMetadata.messageOffset, logStartOffset)
-        val segment = segments.floorEntry(offset).getValue
-        val position  = segment.translateOffset(offset)
-        Some(LogOffsetMetadata(offset, segment.baseOffset, position.position))
+        Some(convertToOffsetMetadataOrThrow(offset))
       case other => other
     }
 
-    if (updatedFirstStableOffset != this.firstUnstableOffset) {
+    if (updatedFirstStableOffset != this.firstUnstableOffsetMetadata) {
       debug(s"First unstable offset updated to $updatedFirstStableOffset")
-      this.firstUnstableOffset = updatedFirstStableOffset
+      this.firstUnstableOffsetMetadata = updatedFirstStableOffset
     }
   }
 
   /**
    * Increment the log start offset if the provided offset is larger.
    */
-  def maybeIncrementLogStartOffset(newLogStartOffset: Long) {
+  def maybeIncrementLogStartOffset(newLogStartOffset: Long): Unit = {
+    if (newLogStartOffset > highWatermark)
+      throw new OffsetOutOfRangeException(s"Cannot increment the log start offset to $newLogStartOffset of partition $topicPartition " +
+        s"since it is larger than the high watermark $highWatermark")
+
     // We don't have to write the log start offset to log-start-offset-checkpoint immediately.
     // The deleteRecordsOffset may be lost only if all in-sync replicas of this broker are shutdown
     // in an unclean manner within log.flush.start.offset.checkpoint.interval.ms. The chance of this happening is low.
@@ -1034,32 +1369,47 @@ class Log(@volatile var dir: File,
         checkIfMemoryMappedBufferClosed()
         if (newLogStartOffset > logStartOffset) {
           info(s"Incrementing log start offset to $newLogStartOffset")
-          logStartOffset = newLogStartOffset
+          updateLogStartOffset(newLogStartOffset)
           leaderEpochCache.foreach(_.truncateFromStart(logStartOffset))
-          producerStateManager.truncateHead(logStartOffset)
-          updateFirstUnstableOffset()
+          producerStateManager.truncateHead(newLogStartOffset)
+          maybeIncrementFirstUnstableOffset()
         }
       }
     }
   }
 
-  private def analyzeAndValidateProducerState(records: MemoryRecords, isFromClient: Boolean):
+  private def analyzeAndValidateProducerState(appendOffsetMetadata: LogOffsetMetadata,
+                                              records: MemoryRecords,
+                                              origin: AppendOrigin):
   (mutable.Map[Long, ProducerAppendInfo], List[CompletedTxn], Option[BatchMetadata]) = {
     val updatedProducers = mutable.Map.empty[Long, ProducerAppendInfo]
     val completedTxns = ListBuffer.empty[CompletedTxn]
-    for (batch <- records.batches.asScala if batch.hasProducerId) {
-      val maybeLastEntry = producerStateManager.lastEntry(batch.producerId)
-
-      // if this is a client produce request, there will be up to 5 batches which could have been duplicated.
-      // If we find a duplicate, we return the metadata of the appended batch to the client.
-      if (isFromClient) {
-        maybeLastEntry.flatMap(_.findDuplicateBatch(batch)).foreach { duplicate =>
-          return (updatedProducers, completedTxns.toList, Some(duplicate))
+    var relativePositionInSegment = appendOffsetMetadata.relativePositionInSegment
+
+    for (batch <- records.batches.asScala) {
+      if (batch.hasProducerId) {
+        val maybeLastEntry = producerStateManager.lastEntry(batch.producerId)
+
+        // if this is a client produce request, there will be up to 5 batches which could have been duplicated.
+        // If we find a duplicate, we return the metadata of the appended batch to the client.
+        if (origin == AppendOrigin.Client) {
+          maybeLastEntry.flatMap(_.findDuplicateBatch(batch)).foreach { duplicate =>
+            return (updatedProducers, completedTxns.toList, Some(duplicate))
+          }
         }
+
+        // We cache offset metadata for the start of each transaction. This allows us to
+        // compute the last stable offset without relying on additional index lookups.
+        val firstOffsetMetadata = if (batch.isTransactional)
+          Some(LogOffsetMetadata(batch.baseOffset, appendOffsetMetadata.segmentBaseOffset, relativePositionInSegment))
+        else
+          None
+
+        val maybeCompletedTxn = updateProducers(batch, updatedProducers, firstOffsetMetadata, origin)
+        maybeCompletedTxn.foreach(completedTxns += _)
       }
 
-      val maybeCompletedTxn = updateProducers(batch, updatedProducers, isFromClient = isFromClient)
-      maybeCompletedTxn.foreach(completedTxns += _)
+      relativePositionInSegment += batch.sizeInBytes
     }
     (updatedProducers, completedTxns.toList, None)
   }
@@ -1082,7 +1432,7 @@ class Log(@volatile var dir: File,
    * 
    •  Whether any compression codec is used (if many are used, then the last one is given)
    * 
    */
-  private def analyzeAndValidateRecords(records: MemoryRecords, isFromClient: Boolean): LogAppendInfo = {
+  private def analyzeAndValidateRecords(records: MemoryRecords, origin: AppendOrigin): LogAppendInfo = {
     var shallowMessageCount = 0
     var validBytesCount = 0
     var firstOffset: Option[Long] = None
@@ -1095,11 +1445,6 @@ class Log(@volatile var dir: File,
     var lastOffsetOfFirstBatch = -1L
 
     for (batch <- records.batches.asScala) {
-      // we only validate V2 and higher to avoid potential compatibility issues with older clients
-      if (batch.magic >= RecordBatch.MAGIC_VALUE_V2 && isFromClient && batch.baseOffset != 0)
-        throw new InvalidRecordException(s"The baseOffset of the record batch in the append to $topicPartition should " +
-          s"be 0, but it is ${batch.baseOffset}")
-
       // update the first offset if on the first message. For magic versions older than 2, we use the last offset
       // to avoid the need to decompress the data (the last offset can be obtained directly from the wrapper message).
       // For magic version 2, we can get the first offset directly from the batch header.
@@ -1130,7 +1475,10 @@ class Log(@volatile var dir: File,
       }
 
       // check the validity of the message by checking CRC
-      batch.ensureValid()
+      if (!batch.isValid) {
+        brokerTopicStats.allTopicsStats.invalidMessageCrcRecordsPerSec.mark()
+        throw new CorruptRecordException(s"Record is corrupt (stored crc = ${batch.checksum()}) in topic partition $topicPartition.")
+      }
 
       if (batch.maxTimestamp > maxTimestamp) {
         maxTimestamp = batch.maxTimestamp
@@ -1153,10 +1501,11 @@ class Log(@volatile var dir: File,
 
   private def updateProducers(batch: RecordBatch,
                               producers: mutable.Map[Long, ProducerAppendInfo],
-                              isFromClient: Boolean): Option[CompletedTxn] = {
+                              firstOffsetMetadata: Option[LogOffsetMetadata],
+                              origin: AppendOrigin): Option[CompletedTxn] = {
     val producerId = batch.producerId
-    val appendInfo = producers.getOrElseUpdate(producerId, producerStateManager.prepareUpdate(producerId, isFromClient))
-    appendInfo.append(batch)
+    val appendInfo = producers.getOrElseUpdate(producerId, producerStateManager.prepareUpdate(producerId, origin))
+    appendInfo.append(batch, firstOffsetMetadata)
   }
 
   /**
@@ -1181,43 +1530,59 @@ class Log(@volatile var dir: File,
     }
   }
 
+  private def emptyFetchDataInfo(fetchOffsetMetadata: LogOffsetMetadata,
+                                 includeAbortedTxns: Boolean): FetchDataInfo = {
+    val abortedTransactions =
+      if (includeAbortedTxns) Some(List.empty[AbortedTransaction])
+      else None
+    FetchDataInfo(fetchOffsetMetadata,
+      MemoryRecords.EMPTY,
+      firstEntryIncomplete = false,
+      abortedTransactions = abortedTransactions)
+  }
+
   /**
    * Read messages from the log.
    *
    * @param startOffset The offset to begin reading at
    * @param maxLength The maximum number of bytes to read
-   * @param maxOffset The offset to read up to, exclusive. (i.e. this offset NOT included in the resulting message set)
+   * @param isolation The fetch isolation, which controls the maximum offset we are allowed to read
    * @param minOneMessage If this is true, the first message will be returned even if it exceeds `maxLength` (if one exists)
-   * @param includeAbortedTxns Whether or not to lookup aborted transactions for fetched data
    * @throws OffsetOutOfRangeException If startOffset is beyond the log end offset or before the log start offset
    * @return The fetch data information including fetch starting offset metadata and messages read.
    */
   def read(startOffset: Long,
            maxLength: Int,
-           maxOffset: Option[Long],
-           minOneMessage: Boolean,
-           includeAbortedTxns: Boolean): FetchDataInfo = {
+           isolation: FetchIsolation,
+           minOneMessage: Boolean): FetchDataInfo = {
     maybeHandleIOException(s"Exception while reading from $topicPartition in dir ${dir.getParent}") {
       trace(s"Reading $maxLength bytes from offset $startOffset of length $size bytes")
 
-      // Because we don't use lock for reading, the synchronization is a little bit tricky.
-      // We create the local variables to avoid race conditions with updates to the log.
-      val currentNextOffsetMetadata = nextOffsetMetadata
-      val next = currentNextOffsetMetadata.messageOffset
-      if (startOffset == next) {
-        val abortedTransactions =
-          if (includeAbortedTxns) Some(List.empty[AbortedTransaction])
-          else None
-        return FetchDataInfo(currentNextOffsetMetadata, MemoryRecords.EMPTY, firstEntryIncomplete = false,
-          abortedTransactions = abortedTransactions)
-      }
+      val includeAbortedTxns = isolation == FetchTxnCommitted
 
+      // Because we don't use the lock for reading, the synchronization is a little bit tricky.
+      // We create the local variables to avoid race conditions with updates to the log.
+      val endOffsetMetadata = nextOffsetMetadata
+      val endOffset = endOffsetMetadata.messageOffset
       var segmentEntry = segments.floorEntry(startOffset)
 
       // return error on attempt to read beyond the log end offset or read below log start offset
-      if (startOffset > next || segmentEntry == null || startOffset < logStartOffset)
+      if (startOffset > endOffset || segmentEntry == null || startOffset < logStartOffset)
         throw new OffsetOutOfRangeException(s"Received request for offset $startOffset for partition $topicPartition, " +
-          s"but we only have log segments in the range $logStartOffset to $next.")
+          s"but we only have log segments in the range $logStartOffset to $endOffset.")
+
+      val maxOffsetMetadata = isolation match {
+        case FetchLogEnd => endOffsetMetadata
+        case FetchHighWatermark => fetchHighWatermarkMetadata
+        case FetchTxnCommitted => fetchLastStableOffsetMetadata
+      }
+
+      if (startOffset == maxOffsetMetadata.messageOffset) {
+        return emptyFetchDataInfo(maxOffsetMetadata, includeAbortedTxns)
+      } else if (startOffset > maxOffsetMetadata.messageOffset) {
+        val startOffsetMetadata = convertToOffsetMetadataOrThrow(startOffset)
+        return emptyFetchDataInfo(startOffsetMetadata, includeAbortedTxns)
+      }
 
       // Do the read on the segment with a base offset less than the target offset
       // but if that segment doesn't contain any messages with an offset greater than that
@@ -1225,24 +1590,16 @@ class Log(@volatile var dir: File,
       while (segmentEntry != null) {
         val segment = segmentEntry.getValue
 
-        // If the fetch occurs on the active segment, there might be a race condition where two fetch requests occur after
-        // the message is appended but before the nextOffsetMetadata is updated. In that case the second fetch may
-        // cause OffsetOutOfRangeException. To solve that, we cap the reading up to exposed position instead of the log
-        // end of the active segment.
         val maxPosition = {
-          if (segmentEntry == segments.lastEntry) {
-            val exposedPos = nextOffsetMetadata.relativePositionInSegment.toLong
-            // Check the segment again in case a new segment has just rolled out.
-            if (segmentEntry != segments.lastEntry)
-            // New log segment has rolled out, we can read up to the file end.
-              segment.size
-            else
-              exposedPos
+          // Use the max offset position if it is on this segment; otherwise, the segment size is the limit.
+          if (maxOffsetMetadata.segmentBaseOffset == segment.baseOffset) {
+            maxOffsetMetadata.relativePositionInSegment
           } else {
             segment.size
           }
         }
-        val fetchInfo = segment.read(startOffset, maxOffset, maxLength, maxPosition, minOneMessage)
+
+        val fetchInfo = segment.read(startOffset, maxLength, maxPosition, minOneMessage)
         if (fetchInfo == null) {
           segmentEntry = segments.higherEntry(segmentEntry.getKey)
         } else {
@@ -1409,20 +1766,15 @@ class Log(@volatile var dir: File,
   }
 
   /**
-   * Given a message offset, find its corresponding offset metadata in the log.
-   * If the message offset is out of range, return None to the caller.
-   */
-  def convertToOffsetMetadata(offset: Long): Option[LogOffsetMetadata] = {
-    try {
-      val fetchDataInfo = read(offset,
-        maxLength = 1,
-        maxOffset = None,
-        minOneMessage = false,
-        includeAbortedTxns = false)
-      Some(fetchDataInfo.fetchOffsetMetadata)
-    } catch {
-      case _: OffsetOutOfRangeException => None
-    }
+    * Given a message offset, find its corresponding offset metadata in the log.
+    * If the message offset is out of range, throw an OffsetOutOfRangeException
+    */
+  private def convertToOffsetMetadataOrThrow(offset: Long): LogOffsetMetadata = {
+    val fetchDataInfo = read(offset,
+      maxLength = 1,
+      isolation = FetchLogEnd,
+      minOneMessage = false)
+    fetchDataInfo.fetchOffsetMetadata
   }
 
   /**
@@ -1452,7 +1804,7 @@ class Log(@volatile var dir: File,
         lock synchronized {
           checkIfMemoryMappedBufferClosed()
           // remove the segments for lookups
-          deletable.foreach(deleteSegment)
+          removeAndDeleteSegments(deletable, asyncDelete = true)
           maybeIncrementLogStartOffset(segments.firstEntry.getValue.baseOffset)
         }
       }
@@ -1473,10 +1825,9 @@ class Log(@volatile var dir: File,
    * @return the segments ready to be deleted
    */
   private def deletableSegments(predicate: (LogSegment, Option[LogSegment]) => Boolean): Iterable[LogSegment] = {
-    if (segments.isEmpty || replicaHighWatermark.isEmpty) {
+    if (segments.isEmpty) {
       Seq.empty
     } else {
-      val highWatermark = replicaHighWatermark.get
       val deletable = ArrayBuffer.empty[LogSegment]
       var segmentEntry = segments.firstEntry
       while (segmentEntry != null) {
@@ -1536,7 +1887,8 @@ class Log(@volatile var dir: File,
 
   private def deleteLogStartOffsetBreachedSegments(): Int = {
     def shouldDelete(segment: LogSegment, nextSegmentOpt: Option[LogSegment]) =
-      nextSegmentOpt.exists(_.baseOffset <= logStartOffset)
+      nextSegmentOpt.exists(_.baseOffset <= logStartOffset) ||
+        (nextSegmentOpt.isEmpty && logEndOffset == logStartOffset)
 
     deleteOldSegments(shouldDelete, reason = s"log start offset $logStartOffset breach")
   }
@@ -1629,7 +1981,7 @@ class Log(@volatile var dir: File,
                  s"=max(provided offset = $expectedNextOffset, LEO = $logEndOffset) while it already " +
                  s"exists and is active with size 0. Size of time index: ${activeSegment.timeIndex.entries}," +
                  s" size of offset index: ${activeSegment.offsetIndex.entries}.")
-            deleteSegment(activeSegment)
+            removeAndDeleteSegments(Seq(activeSegment), asyncDelete = true)
           } else {
             throw new KafkaException(s"Trying to roll a new log segment for topic partition $topicPartition with start offset $newOffset" +
                                      s" =max(provided offset = $expectedNextOffset, LEO = $logEndOffset) while it already exists. Existing " +
@@ -1668,9 +2020,11 @@ class Log(@volatile var dir: File,
           initFileSize = initFileSize,
           preallocate = config.preallocate)
         addSegment(segment)
+
         // We need to update the segment base offset and append position data of the metadata when log rolls.
         // The next offset should not change.
         updateLogEndOffset(nextOffsetMetadata.messageOffset)
+
         // schedule an asynchronous flush of the old segment
         scheduler.schedule("flush-log", () => flush(newOffset), delay = 0L)
 
@@ -1696,7 +2050,7 @@ class Log(@volatile var dir: File,
    *
    * @param offset The offset to flush up to (non-inclusive); the new recovery point
    */
-  def flush(offset: Long) : Unit = {
+  def flush(offset: Long): Unit = {
     maybeHandleIOException(s"Error while flushing log for $topicPartition in dir ${dir.getParent} with offset $offset") {
       if (offset <= this.recoveryPoint)
         return
@@ -1754,13 +2108,13 @@ class Log(@volatile var dir: File,
   /**
    * Completely delete this log directory and all contents from the file system with no delay
    */
-  private[log] def delete() {
+  private[log] def delete(): Unit = {
     maybeHandleIOException(s"Error while deleting log for $topicPartition in dir ${dir.getParent}") {
       lock synchronized {
         checkIfMemoryMappedBufferClosed()
         removeLogMetrics()
-        logSegments.foreach(_.deleteIfExists())
-        segments.clear()
+        producerExpireCheck.cancel(true)
+        removeAndDeleteSegments(logSegments, asyncDelete = false)
         leaderEpochCache.foreach(_.clear())
         Utils.delete(dir)
         // File handlers will be closed if this log is deleted
@@ -1811,11 +2165,10 @@ class Log(@volatile var dir: File,
             truncateFullyAndStartAt(targetOffset)
           } else {
             val deletable = logSegments.filter(segment => segment.baseOffset > targetOffset)
-            deletable.foreach(deleteSegment)
+            removeAndDeleteSegments(deletable, asyncDelete = true)
             activeSegment.truncateTo(targetOffset)
             updateLogEndOffset(targetOffset)
-            this.recoveryPoint = math.min(targetOffset, this.recoveryPoint)
-            this.logStartOffset = math.min(targetOffset, this.logStartOffset)
+            updateLogStartOffset(math.min(targetOffset, this.logStartOffset))
             leaderEpochCache.foreach(_.truncateFromEnd(targetOffset))
             loadProducerState(targetOffset, reloadFromCleanShutdown = false)
           }
@@ -1830,13 +2183,12 @@ class Log(@volatile var dir: File,
    *
    *  @param newOffset The new offset to start the log with
    */
-  private[log] def truncateFullyAndStartAt(newOffset: Long) {
+  private[log] def truncateFullyAndStartAt(newOffset: Long): Unit = {
     maybeHandleIOException(s"Error while truncating the entire log for $topicPartition in dir ${dir.getParent}") {
       debug(s"Truncate and start at offset $newOffset")
       lock synchronized {
         checkIfMemoryMappedBufferClosed()
-        val segmentsToDelete = logSegments.toList
-        segmentsToDelete.foreach(deleteSegment)
+        removeAndDeleteSegments(logSegments, asyncDelete = true)
         addSegment(LogSegment.open(dir,
           baseOffset = newOffset,
           config = config,
@@ -1849,10 +2201,8 @@ class Log(@volatile var dir: File,
 
         producerStateManager.truncate()
         producerStateManager.updateMapEndOffset(newOffset)
-        updateFirstUnstableOffset()
-
-        this.recoveryPoint = math.min(newOffset, this.recoveryPoint)
-        this.logStartOffset = newOffset
+        maybeIncrementFirstUnstableOffset()
+        updateLogStartOffset(newOffset)
       }
     }
   }
@@ -1874,44 +2224,86 @@ class Log(@volatile var dir: File,
 
   /**
    * Get all segments beginning with the segment that includes "from" and ending with the segment
-   * that includes up to "to-1" or the end of the log (if to > logEndOffset)
+   * that includes up to "to-1" or the end of the log (if to > logEndOffset).
    */
   def logSegments(from: Long, to: Long): Iterable[LogSegment] = {
+    if (from == to) {
+      // Handle non-segment-aligned empty sets
+      List.empty[LogSegment]
+    } else if (to < from) {
+      throw new IllegalArgumentException(s"Invalid log segment range: requested segments in $topicPartition " +
+        s"from offset $from which is greater than limit offset $to")
+    } else {
+      lock synchronized {
+        val view = Option(segments.floorKey(from)).map { floor =>
+          segments.subMap(floor, to)
+        }.getOrElse(segments.headMap(to))
+        view.values.asScala
+      }
+    }
+  }
+
+  def nonActiveLogSegmentsFrom(from: Long): Iterable[LogSegment] = {
     lock synchronized {
-      val view = Option(segments.floorKey(from)).map { floor =>
-        segments.subMap(floor, to)
-      }.getOrElse(segments.headMap(to))
-      view.values.asScala
+      if (from > activeSegment.baseOffset)
+        Seq.empty
+      else
+        logSegments(from, activeSegment.baseOffset)
     }
   }
 
-  override def toString = "Log(" + dir + ")"
+  /**
+   * Get the largest log segment with a base offset less than or equal to the given offset, if one exists.
+   * @return the optional log segment
+   */
+  private def floorLogSegment(offset: Long): Option[LogSegment] = {
+    Option(segments.floorEntry(offset)).map(_.getValue)
+  }
+
+  override def toString: String = {
+    val logString = new StringBuilder
+    logString.append(s"Log(dir=$dir")
+    logString.append(s", topic=${topicPartition.topic}")
+    logString.append(s", partition=${topicPartition.partition}")
+    logString.append(s", highWatermark=$highWatermark")
+    logString.append(s", lastStableOffset=$lastStableOffset")
+    logString.append(s", logStartOffset=$logStartOffset")
+    logString.append(s", logEndOffset=$logEndOffset")
+    logString.append(")")
+    logString.toString
+  }
 
   /**
-   * This method performs an asynchronous log segment delete by doing the following:
+   * This method deletes the given log segments by doing the following for each of them:
    * 
        
    *   
      1. It removes the segment from the segment map so that it will no longer be used for reads.
    *   
      2. It renames the index and log files by appending .deleted to the respective file name
-   *   
      3. It schedules an asynchronous delete operation to occur in the future
+   *   
      4. It can either schedule an asynchronous delete operation to occur in the future or perform the deletion synchronously
    * 
      
-   * This allows reads to happen concurrently without synchronization and without the possibility of physically
-   * deleting a file while it is being read from.
+   * Asynchronous deletion allows reads to happen concurrently without synchronization and without the possibility of
+   * physically deleting a file while it is being read.
    *
    * This method does not need to convert IOException to KafkaStorageException because it is either called before all logs are loaded
    * or the immediate caller will catch and handle IOException
    *
-   * @param segment The log segment to schedule for deletion
+   * @param segments The log segments to schedule for deletion
+   * @param asyncDelete Whether the segment files should be deleted asynchronously
    */
-  private def deleteSegment(segment: LogSegment) {
-    info(s"Scheduling log segment [baseOffset ${segment.baseOffset}, size ${segment.size}] for deletion.")
+  private def removeAndDeleteSegments(segments: Iterable[LogSegment], asyncDelete: Boolean): Unit = {
     lock synchronized {
-      segments.remove(segment.baseOffset)
-      asyncDeleteSegment(segment)
+      // As most callers hold an iterator into the `segments` collection and `removeAndDeleteSegment` mutates it by
+      // removing the deleted segment, we should force materialization of the iterator here, so that results of the
+      // iteration remain valid and deterministic.
+      val toDelete = segments.toList
+      toDelete.foreach { segment =>
+        this.segments.remove(segment.baseOffset)
+      }
+      deleteSegmentFiles(toDelete, asyncDelete)
     }
   }
 
   /**
-   * Perform an asynchronous delete on the given file.
+   * Perform physical deletion for the given file. Allows the file to be deleted asynchronously or synchronously.
    *
    * This method assumes that the file exists and the method is not thread-safe.
    *
@@ -1920,15 +2312,22 @@ class Log(@volatile var dir: File,
    *
    * @throws IOException if the file can't be renamed and still exists
    */
-  private def asyncDeleteSegment(segment: LogSegment) {
-    segment.changeFileSuffixes("", Log.DeletedFileSuffix)
-    def deleteSeg() {
-      info(s"Deleting segment ${segment.baseOffset}")
+  private def deleteSegmentFiles(segments: Iterable[LogSegment], asyncDelete: Boolean): Unit = {
+    segments.foreach(_.changeFileSuffixes("", Log.DeletedFileSuffix))
+
+    def deleteSegments(): Unit = {
+      info(s"Deleting segments $segments")
       maybeHandleIOException(s"Error while deleting segments for $topicPartition in dir ${dir.getParent}") {
-        segment.deleteIfExists()
+        segments.foreach(_.deleteIfExists())
       }
     }
-    scheduler.schedule("delete-file", deleteSeg _, delay = config.fileDeleteDelayMs)
+
+    if (asyncDelete) {
+      info(s"Scheduling segments for deletion $segments")
+      scheduler.schedule("delete-file", () => deleteSegments, delay = config.fileDeleteDelayMs)
+    } else {
+      deleteSegments()
+    }
   }
 
   /**
@@ -1963,7 +2362,7 @@ class Log(@volatile var dir: File,
    * @param oldSegments The old log segments to delete from the log
    * @param isRecoveredSwapFile true if the new segment was created from a swap file during recovery after a crash
    */
-  private[log] def replaceSegments(newSegments: Seq[LogSegment], oldSegments: Seq[LogSegment], isRecoveredSwapFile: Boolean = false) {
+  private[log] def replaceSegments(newSegments: Seq[LogSegment], oldSegments: Seq[LogSegment], isRecoveredSwapFile: Boolean = false): Unit = {
     lock synchronized {
       val sortedNewSegments = newSegments.sortBy(_.baseOffset)
       // Some old segments may have been removed from index and scheduled for async deletion after the caller reads segments
@@ -1983,14 +2382,28 @@ class Log(@volatile var dir: File,
         // remove the index entry
         if (seg.baseOffset != sortedNewSegments.head.baseOffset)
           segments.remove(seg.baseOffset)
-        // delete segment
-        asyncDeleteSegment(seg)
+        // delete segment files
+        deleteSegmentFiles(List(seg), asyncDelete = true)
       }
       // okay we are safe now, remove the swap suffix
       sortedNewSegments.foreach(_.changeFileSuffixes(Log.SwapFileSuffix, ""))
     }
   }
 
+  /**
+    * This function does not acquire Log.lock. The caller has to make sure log segments don't get deleted during
+    * this call, and also protects against calling this function on the same segment in parallel.
+    *
+    * Currently, it is used by LogCleaner threads on log compact non-active segments only with LogCleanerManager's lock
+    * to ensure no other logcleaner threads and retention thread can work on the same segment.
+    */
+  private[log] def getFirstBatchTimestampForSegments(segments: Iterable[LogSegment]): Iterable[Long] = {
+    segments.map {
+      segment =>
+        segment.getFirstBatchTimestamp()
+    }
+  }
+
   /**
    * remove deleted log metrics
    */
@@ -2137,10 +2550,10 @@ object Log {
   /** a directory that is used for future partition */
   val FutureDirSuffix = "-future"
 
-  private val DeleteDirPattern = Pattern.compile(s"^(\\S+)-(\\S+)\\.(\\S+)$DeleteDirSuffix")
-  private val FutureDirPattern = Pattern.compile(s"^(\\S+)-(\\S+)\\.(\\S+)$FutureDirSuffix")
+  private[log] val DeleteDirPattern = Pattern.compile(s"^(\\S+)-(\\S+)\\.(\\S+)$DeleteDirSuffix")
+  private[log] val FutureDirPattern = Pattern.compile(s"^(\\S+)-(\\S+)\\.(\\S+)$FutureDirSuffix")
 
-  val UnknownLogStartOffset = -1L
+  val UnknownOffset = -1L
 
   def apply(dir: File,
             config: LogConfig,
@@ -2184,11 +2597,16 @@ object Log {
     new File(dir, filenamePrefixFromOffset(offset) + LogFileSuffix + suffix)
 
   /**
-   * Return a directory name to rename the log directory to for async deletion. The name will be in the following
-   * format: topic-partition.uniqueId-delete where topic, partition and uniqueId are variables.
+   * Return a directory name to rename the log directory to for async deletion.
+   * The name will be in the following format: "topic-partitionId.uniqueId-delete".
+   * If the topic name is too long, it will be truncated to prevent the total name
+   * from exceeding 255 characters.
    */
   def logDeleteDirName(topicPartition: TopicPartition): String = {
-    logDirNameWithSuffix(topicPartition, DeleteDirSuffix)
+    val uniqueId = java.util.UUID.randomUUID.toString.replaceAll("-", "")
+    val suffix = s"-${topicPartition.partition()}.${uniqueId}${DeleteDirSuffix}"
+    val prefixLength = Math.min(topicPartition.topic().size, 255 - suffix.size)
+    s"${topicPartition.topic().substring(0, prefixLength)}${suffix}"
   }
 
   /**
diff --git a/core/src/main/scala/kafka/log/LogCleaner.scala b/core/src/main/scala/kafka/log/LogCleaner.scala
index 7fc70dabfd3bc..27be737bc86e7 100644
--- a/core/src/main/scala/kafka/log/LogCleaner.scala
+++ b/core/src/main/scala/kafka/log/LogCleaner.scala
@@ -36,7 +36,8 @@ import org.apache.kafka.common.record._
 import org.apache.kafka.common.utils.Time
 
 import scala.collection.JavaConverters._
-import scala.collection.{Iterable, Set, mutable}
+import scala.collection.mutable.ListBuffer
+import scala.collection.{Iterable, Seq, Set, mutable}
 import scala.util.control.ControlThrowable
 
 /**
@@ -87,12 +88,14 @@ import scala.util.control.ControlThrowable
  * @param initialConfig Initial configuration parameters for the cleaner. Actual config may be dynamically updated.
  * @param logDirs The directories where offset checkpoints reside
  * @param logs The pool of logs
+ * @param retentionCheckMs The frequency that the log cleaner threads checks whether any log is eligible for deletion
  * @param time A way to control the passage of time
  */
 class LogCleaner(initialConfig: CleanerConfig,
                  val logDirs: Seq[File],
                  val logs: Pool[TopicPartition, Log],
                  val logDirFailureChannel: LogDirFailureChannel,
+                 val retentionCheckMs: Long,
                  time: Time = Time.SYSTEM) extends Logging with KafkaMetricsGroup with BrokerReconfigurable
 {
 
@@ -100,7 +103,7 @@ class LogCleaner(initialConfig: CleanerConfig,
   @volatile private var config = initialConfig
 
   /* for managing the state of partitions being cleaned. package-private to allow access in tests */
-  private[log] val cleanerManager = new LogCleanerManager(logDirs, logs, logDirFailureChannel)
+  private[log] val cleanerManager = new LogCleanerManager(logDirs, logs, logDirFailureChannel, retentionCheckMs)
 
   /* a throttle used to limit the I/O of all the cleaner threads to a user-specified maximum rate */
   private val throttler = new Throttler(desiredRatePerSec = config.maxIoBytesPerSecond,
@@ -110,8 +113,7 @@ class LogCleaner(initialConfig: CleanerConfig,
                                         "bytes",
                                         time = time)
 
-  /* the threads */
-  private val cleaners = mutable.ArrayBuffer[CleanerThread]()
+  private[log] val cleaners = mutable.ArrayBuffer[CleanerThread]()
 
   /* a metric to track the maximum utilization of any thread's buffer in the last cleaning */
   newGauge("max-buffer-utilization-percent",
@@ -132,11 +134,30 @@ class LogCleaner(initialConfig: CleanerConfig,
            new Gauge[Int] {
              def value: Int = cleaners.map(_.lastStats).map(_.elapsedSecs).max.toInt
            })
+  // a metric to track delay between the time when a log is required to be compacted
+  // as determined by max compaction lag and the time of last cleaner run.
+  newGauge("max-compaction-delay-secs",
+          new Gauge[Int] {
+          def value: Int = Math.max(0, (cleaners.map(_.lastPreCleanStats).map(_.maxCompactionDelayMs).max / 1000).toInt)
+          })
+
+  newGauge("DeadThreadCount",
+    new Gauge[Int] {
+      def value: Int = deadThreadCount
+    })
+
+  private[log] def deadThreadCount: Int = cleaners.count(_.isThreadFailed)
+  /* a metric to track the number of cleaner threads alive */
+  newGauge("live-cleaner-thread-count",
+          new Gauge[Int] {
+            def value: Int = cleaners.count(_.asInstanceOf[Thread].isAlive)
+          })
+
 
   /**
    * Start the background cleaning
    */
-  def startup() {
+  def startup(): Unit = {
     info("Starting the log cleaner")
     (0 until config.numThreads).foreach { i =>
       val cleaner = new CleanerThread(i)
@@ -148,7 +169,7 @@ class LogCleaner(initialConfig: CleanerConfig,
   /**
    * Stop the background cleaning
    */
-  def shutdown() {
+  def shutdown(): Unit = {
     info("Shutting down the log cleaner.")
     cleaners.foreach(_.shutdown())
     cleaners.clear()
@@ -185,14 +206,14 @@ class LogCleaner(initialConfig: CleanerConfig,
    *  Abort the cleaning of a particular partition, if it's in progress. This call blocks until the cleaning of
    *  the partition is aborted.
    */
-  def abortCleaning(topicPartition: TopicPartition) {
+  def abortCleaning(topicPartition: TopicPartition): Unit = {
     cleanerManager.abortCleaning(topicPartition)
   }
 
   /**
    * Update checkpoint file, removing topics and partitions that no longer exist
    */
-  def updateCheckpoints(dataDir: File) {
+  def updateCheckpoints(dataDir: File): Unit = {
     cleanerManager.updateCheckpoints(dataDir, update=None)
   }
 
@@ -200,14 +221,14 @@ class LogCleaner(initialConfig: CleanerConfig,
     cleanerManager.alterCheckpointDir(topicPartition, sourceLogDir, destLogDir)
   }
 
-  def handleLogDirFailure(dir: String) {
+  def handleLogDirFailure(dir: String): Unit = {
     cleanerManager.handleLogDirFailure(dir)
   }
 
   /**
    * Truncate cleaner offset checkpoint for the given partition if its checkpointed offset is larger than the given offset
    */
-  def maybeTruncateCheckpoint(dataDir: File, topicPartition: TopicPartition, offset: Long) {
+  def maybeTruncateCheckpoint(dataDir: File, topicPartition: TopicPartition, offset: Long): Unit = {
     cleanerManager.maybeTruncateCheckpoint(dataDir, topicPartition, offset)
   }
 
@@ -215,14 +236,14 @@ class LogCleaner(initialConfig: CleanerConfig,
    *  Abort the cleaning of a particular partition if it's in progress, and pause any future cleaning of this partition.
    *  This call blocks until the cleaning of the partition is aborted and paused.
    */
-  def abortAndPauseCleaning(topicPartition: TopicPartition) {
+  def abortAndPauseCleaning(topicPartition: TopicPartition): Unit = {
     cleanerManager.abortAndPauseCleaning(topicPartition)
   }
 
   /**
     *  Resume the cleaning of paused partitions.
     */
-  def resumeCleaning(topicPartitions: Iterable[TopicPartition]) {
+  def resumeCleaning(topicPartitions: Iterable[TopicPartition]): Unit = {
     cleanerManager.resumeCleaning(topicPartitions)
   }
 
@@ -266,7 +287,7 @@ class LogCleaner(initialConfig: CleanerConfig,
    * The cleaner threads do the actual log cleaning. Each thread processes does its cleaning repeatedly by
    * choosing the dirtiest log, cleaning it, and then swapping in the cleaned segments.
    */
-  private class CleanerThread(threadId: Int)
+  private[log] class CleanerThread(threadId: Int)
     extends ShutdownableThread(name = s"kafka-log-cleaner-thread-$threadId", isInterruptible = false) {
 
     protected override def loggerName = classOf[LogCleaner].getName
@@ -285,8 +306,9 @@ class LogCleaner(initialConfig: CleanerConfig,
                               checkDone = checkDone)
 
     @volatile var lastStats: CleanerStats = new CleanerStats()
+    @volatile var lastPreCleanStats: PreCleanStats = new PreCleanStats()
 
-    private def checkDone(topicPartition: TopicPartition) {
+    private def checkDone(topicPartition: TopicPartition): Unit = {
       if (!isRunning)
         throw new ThreadShutdownException
       cleanerManager.checkCleaningAborted(topicPartition)
@@ -296,52 +318,60 @@ class LogCleaner(initialConfig: CleanerConfig,
      * The main loop for the cleaner thread
      * Clean a log if there is a dirty log available, otherwise sleep for a bit
      */
-    override def doWork() {
-      val cleaned = cleanFilthiestLog()
+    override def doWork(): Unit = {
+      val cleaned = tryCleanFilthiestLog()
       if (!cleaned)
         pause(config.backOffMs, TimeUnit.MILLISECONDS)
     }
 
     /**
-      * Cleans a log if there is a dirty log available
-      * @return whether a log was cleaned
-      */
-    private def cleanFilthiestLog(): Boolean = {
-      var currentLog: Option[Log] = None
-
+     * Cleans a log if there is a dirty log available
+     * @return whether a log was cleaned
+     */
+    private def tryCleanFilthiestLog(): Boolean = {
       try {
-        val cleaned = cleanerManager.grabFilthiestCompactedLog(time) match {
-          case None =>
-            false
-          case Some(cleanable) =>
-            // there's a log, clean it
-            currentLog = Some(cleanable.log)
+        cleanFilthiestLog()
+      } catch {
+        case e: LogCleaningException =>
+          warn(s"Unexpected exception thrown when cleaning log ${e.log}. Marking its partition (${e.log.topicPartition}) as uncleanable", e)
+          cleanerManager.markPartitionUncleanable(e.log.parentDir, e.log.topicPartition)
+
+          false
+      }
+    }
+
+    @throws(classOf[LogCleaningException])
+    private def cleanFilthiestLog(): Boolean = {
+      val preCleanStats = new PreCleanStats()
+      val cleaned = cleanerManager.grabFilthiestCompactedLog(time, preCleanStats) match {
+        case None =>
+          false
+        case Some(cleanable) =>
+          // there's a log, clean it
+          this.lastPreCleanStats = preCleanStats
+          try {
             cleanLog(cleanable)
             true
-        }
-        val deletable: Iterable[(TopicPartition, Log)] = cleanerManager.deletableLogs()
-        try {
-          deletable.foreach { case (topicPartition, log) =>
-            currentLog = Some(log)
+          } catch {
+            case e @ (_: ThreadShutdownException | _: ControlThrowable) => throw e
+            case e: Exception => throw new LogCleaningException(cleanable.log, e.getMessage, e)
+          }
+      }
+      val deletable: Iterable[(TopicPartition, Log)] = cleanerManager.deletableLogs()
+      try {
+        deletable.foreach { case (_, log) =>
+          try {
             log.deleteOldSegments()
+          } catch {
+            case e @ (_: ThreadShutdownException | _: ControlThrowable) => throw e
+            case e: Exception => throw new LogCleaningException(log, e.getMessage, e)
           }
-        } finally  {
-          cleanerManager.doneDeleting(deletable.map(_._1))
         }
-
-        cleaned
-      } catch {
-        case e @ (_: ThreadShutdownException | _: ControlThrowable) => throw e
-        case e: Exception =>
-          if (currentLog.isEmpty) {
-            throw new IllegalStateException("currentLog cannot be empty on an unexpected exception", e)
-          }
-          val erroneousLog = currentLog.get
-          warn(s"Unexpected exception thrown when cleaning log $erroneousLog. Marking its partition (${erroneousLog.topicPartition}) as uncleanable", e)
-          cleanerManager.markPartitionUncleanable(erroneousLog.dir.getParent, erroneousLog.topicPartition)
-
-          false
+      } finally  {
+        cleanerManager.doneDeleting(deletable.map(_._1))
       }
+
+      cleaned
     }
 
     private def cleanLog(cleanable: LogToClean): Unit = {
@@ -354,18 +384,18 @@ class LogCleaner(initialConfig: CleanerConfig,
         case _: LogCleaningAbortedException => // task can be aborted, let it go.
         case _: KafkaStorageException => // partition is already offline. let it go.
         case e: IOException =>
-          val logDirectory = cleanable.log.dir.getParent
+          val logDirectory = cleanable.log.parentDir
           val msg = s"Failed to clean up log for ${cleanable.topicPartition} in dir ${logDirectory} due to IOException"
           logDirFailureChannel.maybeAddOfflineLogDir(logDirectory, msg, e)
       } finally {
-        cleanerManager.doneCleaning(cleanable.topicPartition, cleanable.log.dir.getParentFile, endOffset)
+        cleanerManager.doneCleaning(cleanable.topicPartition, cleanable.log.parentDirFile, endOffset)
       }
     }
 
     /**
      * Log out statistics on a single run of the cleaner.
      */
-    def recordStats(id: Int, name: String, from: Long, to: Long, stats: CleanerStats) {
+    def recordStats(id: Int, name: String, from: Long, to: Long, stats: CleanerStats): Unit = {
       this.lastStats = stats
       def mb(bytes: Double) = bytes / (1024*1024)
       val message =
@@ -386,6 +416,9 @@ class LogCleaner(initialConfig: CleanerConfig,
         "\t%.1f%% size reduction (%.1f%% fewer messages)%n".format(100.0 * (1.0 - stats.bytesWritten.toDouble/stats.bytesRead),
                                                                    100.0 * (1.0 - stats.messagesWritten.toDouble/stats.messagesRead))
       info(message)
+      if (lastPreCleanStats.delayedPartitions > 0) {
+        info("\tCleanable partitions: %d, Delayed partitions: %d, max delay: %d".format(lastPreCleanStats.cleanablePartitions, lastPreCleanStats.delayedPartitions, lastPreCleanStats.maxCompactionDelayMs))
+      }
       if (stats.invalidMessagesRead > 0) {
         warn("\tFound %d invalid messages during compaction.".format(stats.invalidMessagesRead))
       }
@@ -423,17 +456,6 @@ object LogCleaner {
       fileSuffix = Log.CleanedFileSuffix, initFileSize = log.initFileSize, preallocate = log.config.preallocate)
   }
 
-  /**
-    * Given the first dirty offset and an uncleanable offset, calculates the total cleanable bytes for this log
-    * @return the biggest uncleanable offset and the total amount of cleanable bytes
-    */
-  def calculateCleanableBytes(log: Log, firstDirtyOffset: Long, uncleanableOffset: Long): (Long, Long) = {
-    val firstUncleanableSegment = log.logSegments(uncleanableOffset, log.activeSegment.baseOffset).headOption.getOrElse(log.activeSegment)
-    val firstUncleanableOffset = firstUncleanableSegment.baseOffset
-    val cleanableBytes = log.logSegments(firstDirtyOffset, math.max(firstDirtyOffset, firstUncleanableOffset)).map(_.size.toLong).sum
-
-    (firstUncleanableOffset, cleanableBytes)
-  }
 }
 
 /**
@@ -454,7 +476,7 @@ private[log] class Cleaner(val id: Int,
                            dupBufferLoadFactor: Double,
                            throttler: Throttler,
                            time: Time,
-                           checkDone: (TopicPartition) => Unit) extends Logging {
+                           checkDone: TopicPartition => Unit) extends Logging {
 
   protected override def loggerName = classOf[LogCleaner].getName
 
@@ -508,8 +530,12 @@ private[log] class Cleaner(val id: Int,
 
     // group the segments and clean the groups
     info("Cleaning log %s (cleaning prior to %s, discarding tombstones prior to %s)...".format(log.name, new Date(cleanableHorizonMs), new Date(deleteHorizonMs)))
-    for (group <- groupSegmentsBySize(log.logSegments(0, endOffset), log.config.segmentSize, log.config.maxIndexSize, cleanable.firstUncleanableOffset))
-      cleanSegments(log, group, offsetMap, deleteHorizonMs, stats)
+    val transactionMetadata = new CleanedTransactionMetadata
+
+    val groupedSegments = groupSegmentsBySize(log.logSegments(0, endOffset), log.config.segmentSize,
+      log.config.maxIndexSize, cleanable.firstUncleanableOffset)
+    for (group <- groupedSegments)
+      cleanSegments(log, group, offsetMap, deleteHorizonMs, stats, transactionMetadata)
 
     // record buffer utilization
     stats.bufferUtilization = offsetMap.utilization
@@ -527,19 +553,25 @@ private[log] class Cleaner(val id: Int,
    * @param map The offset map to use for cleaning segments
    * @param deleteHorizonMs The time to retain delete tombstones
    * @param stats Collector for cleaning statistics
+   * @param transactionMetadata State of ongoing transactions which is carried between the cleaning
+   *                            of the grouped segments
    */
   private[log] def cleanSegments(log: Log,
                                  segments: Seq[LogSegment],
                                  map: OffsetMap,
                                  deleteHorizonMs: Long,
-                                 stats: CleanerStats) {
+                                 stats: CleanerStats,
+                                 transactionMetadata: CleanedTransactionMetadata): Unit = {
     // create a new segment with a suffix appended to the name of the log and indexes
     val cleaned = LogCleaner.createNewCleanedSegment(log, segments.head.baseOffset)
+    transactionMetadata.cleanedIndex = Some(cleaned.txnIndex)
 
     try {
       // clean segments into the new destination segment
       val iter = segments.iterator
       var currentSegmentOpt: Option[LogSegment] = Some(iter.next())
+      val lastOffsetOfActiveProducers = log.lastRecordsOfActiveProducers
+
       while (currentSegmentOpt.isDefined) {
         val currentSegment = currentSegmentOpt.get
         val nextSegmentOpt = if (iter.hasNext) Some(iter.next()) else None
@@ -547,15 +579,16 @@ private[log] class Cleaner(val id: Int,
         val startOffset = currentSegment.baseOffset
         val upperBoundOffset = nextSegmentOpt.map(_.baseOffset).getOrElse(map.latestOffset + 1)
         val abortedTransactions = log.collectAbortedTransactions(startOffset, upperBoundOffset)
-        val transactionMetadata = CleanedTransactionMetadata(abortedTransactions, Some(cleaned.txnIndex))
+        transactionMetadata.addAbortedTransactions(abortedTransactions)
 
-        val retainDeletes = currentSegment.lastModified > deleteHorizonMs
-        info(s"Cleaning segment $startOffset in log ${log.name} (largest timestamp ${new Date(currentSegment.largestTimestamp)}) " +
-          s"into ${cleaned.baseOffset}, ${if(retainDeletes) "retaining" else "discarding"} deletes.")
+        val retainDeletesAndTxnMarkers = currentSegment.lastModified > deleteHorizonMs
+        info(s"Cleaning $currentSegment in log ${log.name} into ${cleaned.baseOffset} " +
+          s"with deletion horizon $deleteHorizonMs, " +
+          s"${if(retainDeletesAndTxnMarkers) "retaining" else "discarding"} deletes.")
 
         try {
-          cleanInto(log.topicPartition, currentSegment.log, cleaned, map, retainDeletes, log.config.maxMessageSize,
-            transactionMetadata, log.activeProducersWithLastSequence, stats)
+          cleanInto(log.topicPartition, currentSegment.log, cleaned, map, retainDeletesAndTxnMarkers, log.config.maxMessageSize,
+            transactionMetadata, lastOffsetOfActiveProducers, stats)
         } catch {
           case e: LogSegmentOffsetOverflowException =>
             // Split the current segment. It's also safest to abort the current cleaning process, so that we retry from
@@ -596,7 +629,7 @@ private[log] class Cleaner(val id: Int,
    * @param sourceRecords The dirty log segment
    * @param dest The cleaned log segment
    * @param map The key=>offset mapping
-   * @param retainDeletes Should delete tombstones be retained while cleaning this segment
+   * @param retainDeletesAndTxnMarkers Should tombstones and markers be retained while cleaning this segment
    * @param maxLogMessageSize The maximum message size of the corresponding topic
    * @param stats Collector for cleaning statistics
    */
@@ -604,22 +637,35 @@ private[log] class Cleaner(val id: Int,
                              sourceRecords: FileRecords,
                              dest: LogSegment,
                              map: OffsetMap,
-                             retainDeletes: Boolean,
+                             retainDeletesAndTxnMarkers: Boolean,
                              maxLogMessageSize: Int,
                              transactionMetadata: CleanedTransactionMetadata,
-                             activeProducers: Map[Long, Int],
-                             stats: CleanerStats) {
-    val logCleanerFilter = new RecordFilter {
+                             lastRecordsOfActiveProducers: Map[Long, LastRecord],
+                             stats: CleanerStats): Unit = {
+    val logCleanerFilter: RecordFilter = new RecordFilter {
       var discardBatchRecords: Boolean = _
 
       override def checkBatchRetention(batch: RecordBatch): BatchRetention = {
         // we piggy-back on the tombstone retention logic to delay deletion of transaction markers.
         // note that we will never delete a marker until all the records from that transaction are removed.
-        discardBatchRecords = shouldDiscardBatch(batch, transactionMetadata, retainTxnMarkers = retainDeletes)
+        discardBatchRecords = shouldDiscardBatch(batch, transactionMetadata, retainTxnMarkers = retainDeletesAndTxnMarkers)
+
+        def isBatchLastRecordOfProducer: Boolean = {
+          // We retain the batch in order to preserve the state of active producers. There are three cases:
+          // 1) The producer is no longer active, which means we can delete all records for that producer.
+          // 2) The producer is still active and has a last data offset. We retain the batch that contains
+          //    this offset since it also contains the last sequence number for this producer.
+          // 3) The last entry in the log is a transaction marker. We retain this marker since it has the
+          //    last producer epoch, which is needed to ensure fencing.
+          lastRecordsOfActiveProducers.get(batch.producerId).exists { lastRecord =>
+            lastRecord.lastDataOffset match {
+              case Some(offset) => batch.lastOffset == offset
+              case None => batch.isControlBatch && batch.producerEpoch == lastRecord.producerEpoch
+            }
+          }
+        }
 
-        // check if the batch contains the last sequence number for the producer. if so, we cannot
-        // remove the batch just yet or the producer may see an out of sequence error.
-        if (batch.hasProducerId && activeProducers.get(batch.producerId).contains(batch.lastSequence))
+        if (batch.hasProducerId && isBatchLastRecordOfProducer)
           BatchRetention.RETAIN_EMPTY
         else if (discardBatchRecords)
           BatchRetention.DELETE
@@ -632,7 +678,7 @@ private[log] class Cleaner(val id: Int,
           // The batch is only retained to preserve producer sequence information; the records can be removed
           false
         else
-          Cleaner.this.shouldRetainRecord(map, retainDeletes, batch, record, stats)
+          Cleaner.this.shouldRetainRecord(map, retainDeletesAndTxnMarkers, batch, record, stats)
       }
     }
 
@@ -731,13 +777,14 @@ private[log] class Cleaner(val id: Int,
     if (record.hasKey) {
       val key = record.key
       val foundOffset = map.get(key)
-      /* two cases in which we can get rid of a message:
-       *   1) if there exists a message with the same key but higher offset
-       *   2) if the message is a delete "tombstone" marker and enough time has passed
+      /* First,the message must have the latest offset for the key
+       * then there are two cases in which we can retain a message:
+       *   1) The message has value
+       *   2) The message doesn't has value but it can't be deleted now.
        */
-      val redundant = foundOffset >= 0 && record.offset < foundOffset
-      val obsoleteDelete = !retainDeletes && !record.hasValue
-      !redundant && !obsoleteDelete
+      val latestOffsetForKey = record.offset() >= foundOffset
+      val isRetainedValue = record.hasValue || retainDeletes
+      latestOffsetForKey && isRetainedValue
     } else {
       stats.invalidMessage()
       false
@@ -747,7 +794,7 @@ private[log] class Cleaner(val id: Int,
   /**
    * Double the I/O buffer capacity
    */
-  def growBuffers(maxLogMessageSize: Int) {
+  def growBuffers(maxLogMessageSize: Int): Unit = {
     val maxBufferSize = math.max(maxLogMessageSize, maxIoBufferSize)
     if(readBuffer.capacity >= maxBufferSize || writeBuffer.capacity >= maxBufferSize)
       throw new IllegalStateException("This log contains a message larger than maximum allowable size of %s.".format(maxBufferSize))
@@ -760,7 +807,7 @@ private[log] class Cleaner(val id: Int,
   /**
    * Restore the I/O buffer capacity to its original size
    */
-  def restoreBuffers() {
+  def restoreBuffers(): Unit = {
     if(this.readBuffer.capacity > this.ioBufferSize)
       this.readBuffer = ByteBuffer.allocate(this.ioBufferSize)
     if(this.writeBuffer.capacity > this.ioBufferSize)
@@ -837,21 +884,27 @@ private[log] class Cleaner(val id: Int,
                                   start: Long,
                                   end: Long,
                                   map: OffsetMap,
-                                  stats: CleanerStats) {
+                                  stats: CleanerStats): Unit = {
     map.clear()
     val dirty = log.logSegments(start, end).toBuffer
+    val nextSegmentStartOffsets = new ListBuffer[Long]
+    if (dirty.nonEmpty) {
+      for (nextSegment <- dirty.tail) nextSegmentStartOffsets.append(nextSegment.baseOffset)
+      nextSegmentStartOffsets.append(end)
+    }
     info("Building offset map for log %s for %d segments in offset range [%d, %d).".format(log.name, dirty.size, start, end))
 
+    val transactionMetadata = new CleanedTransactionMetadata
     val abortedTransactions = log.collectAbortedTransactions(start, end)
-    val transactionMetadata = CleanedTransactionMetadata(abortedTransactions)
+    transactionMetadata.addAbortedTransactions(abortedTransactions)
 
     // Add all the cleanable dirty segments. We must take at least map.slots * load_factor,
     // but we may be able to fit more (if there is lots of duplication in the dirty section of the log)
     var full = false
-    for (segment <- dirty if !full) {
+    for ( (segment, nextSegmentStartOffset) <- dirty.zip(nextSegmentStartOffsets) if !full) {
       checkDone(log.topicPartition)
 
-      full = buildOffsetMapForSegment(log.topicPartition, segment, map, start, log.config.maxMessageSize,
+      full = buildOffsetMapForSegment(log.topicPartition, segment, map, start, nextSegmentStartOffset, log.config.maxMessageSize,
         transactionMetadata, stats)
       if (full)
         debug("Offset map is full, %d segments fully mapped, segment with base offset %d is partially mapped".format(dirty.indexOf(segment), segment.baseOffset))
@@ -872,6 +925,7 @@ private[log] class Cleaner(val id: Int,
                                        segment: LogSegment,
                                        map: OffsetMap,
                                        startOffset: Long,
+                                       nextSegmentStartOffset: Long,
                                        maxLogMessageSize: Int,
                                        transactionMetadata: CleanedTransactionMetadata,
                                        stats: CleanerStats): Boolean = {
@@ -925,11 +979,34 @@ private[log] class Cleaner(val id: Int,
       if(position == startPosition)
         growBuffersOrFail(segment.log, position, maxLogMessageSize, records)
     }
+
+    // In the case of offsets gap, fast forward to latest expected offset in this segment.
+    map.updateLatestOffset(nextSegmentStartOffset - 1L)
+
     restoreBuffers()
     false
   }
 }
 
+/**
+  * A simple struct for collecting pre-clean stats
+  */
+private class PreCleanStats() {
+  var maxCompactionDelayMs = 0L
+  var delayedPartitions = 0
+  var cleanablePartitions = 0
+
+  def updateMaxCompactionDelay(delayMs: Long): Unit = {
+    maxCompactionDelayMs = Math.max(maxCompactionDelayMs, delayMs)
+    if (delayMs > 0) {
+      delayedPartitions += 1
+    }
+  }
+  def recordCleanablePartitions(numOfCleanables: Int): Unit = {
+    cleanablePartitions = numOfCleanables
+  }
+}
+
 /**
  * A simple struct for collecting stats about log cleaning
  */
@@ -946,33 +1023,33 @@ private class CleanerStats(time: Time = Time.SYSTEM) {
   var messagesWritten = 0L
   var bufferUtilization = 0.0d
 
-  def readMessages(messagesRead: Int, bytesRead: Int) {
+  def readMessages(messagesRead: Int, bytesRead: Int): Unit = {
     this.messagesRead += messagesRead
     this.bytesRead += bytesRead
   }
 
-  def invalidMessage() {
+  def invalidMessage(): Unit = {
     invalidMessagesRead += 1
   }
 
-  def recopyMessages(messagesWritten: Int, bytesWritten: Int) {
+  def recopyMessages(messagesWritten: Int, bytesWritten: Int): Unit = {
     this.messagesWritten += messagesWritten
     this.bytesWritten += bytesWritten
   }
 
-  def indexMessagesRead(size: Int) {
+  def indexMessagesRead(size: Int): Unit = {
     mapMessagesRead += size
   }
 
-  def indexBytesRead(size: Int) {
+  def indexBytesRead(size: Int): Unit = {
     mapBytesRead += size
   }
 
-  def indexDone() {
+  def indexDone(): Unit = {
     mapCompleteTime = time.milliseconds
   }
 
-  def allDone() {
+  def allDone(): Unit = {
     endTime = time.milliseconds
   }
 
@@ -983,39 +1060,41 @@ private class CleanerStats(time: Time = Time.SYSTEM) {
 }
 
 /**
- * Helper class for a log, its topic/partition, the first cleanable position, and the first uncleanable dirty position
- */
-private case class LogToClean(topicPartition: TopicPartition, log: Log, firstDirtyOffset: Long, uncleanableOffset: Long) extends Ordered[LogToClean] {
+  * Helper class for a log, its topic/partition, the first cleanable position, the first uncleanable dirty position,
+  * and whether it needs compaction immediately.
+  */
+private case class LogToClean(topicPartition: TopicPartition,
+                              log: Log,
+                              firstDirtyOffset: Long,
+                              uncleanableOffset: Long,
+                              needCompactionNow: Boolean = false) extends Ordered[LogToClean] {
   val cleanBytes = log.logSegments(-1, firstDirtyOffset).map(_.size.toLong).sum
-  val (firstUncleanableOffset, cleanableBytes) = LogCleaner.calculateCleanableBytes(log, firstDirtyOffset, uncleanableOffset)
+  val (firstUncleanableOffset, cleanableBytes) = LogCleanerManager.calculateCleanableBytes(log, firstDirtyOffset, uncleanableOffset)
   val totalBytes = cleanBytes + cleanableBytes
   val cleanableRatio = cleanableBytes / totalBytes.toDouble
   override def compare(that: LogToClean): Int = math.signum(this.cleanableRatio - that.cleanableRatio).toInt
 }
 
-private[log] object CleanedTransactionMetadata {
-  def apply(abortedTransactions: List[AbortedTxn],
-            transactionIndex: Option[TransactionIndex] = None): CleanedTransactionMetadata = {
-    val queue = mutable.PriorityQueue.empty[AbortedTxn](new Ordering[AbortedTxn] {
-      override def compare(x: AbortedTxn, y: AbortedTxn): Int = x.firstOffset compare y.firstOffset
-    }.reverse)
-    queue ++= abortedTransactions
-    new CleanedTransactionMetadata(queue, transactionIndex)
-  }
-
-  val Empty = CleanedTransactionMetadata(List.empty[AbortedTxn])
-}
-
 /**
- * This is a helper class to facilitate tracking transaction state while cleaning the log. It is initialized
- * with the aborted transactions from the transaction index and its state is updated as the cleaner iterates through
- * the log during a round of cleaning. This class is responsible for deciding when transaction markers can
- * be removed and is therefore also responsible for updating the cleaned transaction index accordingly.
+ * This is a helper class to facilitate tracking transaction state while cleaning the log. It maintains a set
+ * of the ongoing aborted and committed transactions as the cleaner is working its way through the log. This
+ * class is responsible for deciding when transaction markers can be removed and is therefore also responsible
+ * for updating the cleaned transaction index accordingly.
  */
-private[log] class CleanedTransactionMetadata(val abortedTransactions: mutable.PriorityQueue[AbortedTxn],
-                                              val transactionIndex: Option[TransactionIndex] = None) {
-  val ongoingCommittedTxns = mutable.Set.empty[Long]
-  val ongoingAbortedTxns = mutable.Map.empty[Long, AbortedTransactionMetadata]
+private[log] class CleanedTransactionMetadata {
+  private val ongoingCommittedTxns = mutable.Set.empty[Long]
+  private val ongoingAbortedTxns = mutable.Map.empty[Long, AbortedTransactionMetadata]
+  // Minheap of aborted transactions sorted by the transaction first offset
+  private val abortedTransactions = mutable.PriorityQueue.empty[AbortedTxn](new Ordering[AbortedTxn] {
+    override def compare(x: AbortedTxn, y: AbortedTxn): Int = x.firstOffset compare y.firstOffset
+  }.reverse)
+
+  // Output cleaned index to write retained aborted transactions
+  var cleanedIndex: Option[TransactionIndex] = None
+
+  def addAbortedTransactions(abortedTransactions: List[AbortedTxn]): Unit = {
+    this.abortedTransactions ++= abortedTransactions
+  }
 
   /**
    * Update the cleaned transaction state with a control batch that has just been traversed by the cleaner.
@@ -1032,9 +1111,11 @@ private[log] class CleanedTransactionMetadata(val abortedTransactions: mutable.P
       controlType match {
         case ControlRecordType.ABORT =>
           ongoingAbortedTxns.remove(producerId) match {
-            // Retain the marker until all batches from the transaction have been removed
+            // Retain the marker until all batches from the transaction have been removed.
+            // We may retain a record from an aborted transaction if it is the last entry
+            // written by a given producerId.
             case Some(abortedTxnMetadata) if abortedTxnMetadata.lastObservedBatchOffset.isDefined =>
-              transactionIndex.foreach(_.append(abortedTxnMetadata.abortedTxn))
+              cleanedIndex.foreach(_.append(abortedTxnMetadata.abortedTxn))
               false
             case _ => true
           }
@@ -1054,7 +1135,7 @@ private[log] class CleanedTransactionMetadata(val abortedTransactions: mutable.P
   private def consumeAbortedTxnsUpTo(offset: Long): Unit = {
     while (abortedTransactions.headOption.exists(_.firstOffset <= offset)) {
       val abortedTxn = abortedTransactions.dequeue()
-      ongoingAbortedTxns += abortedTxn.producerId -> new AbortedTransactionMetadata(abortedTxn)
+      ongoingAbortedTxns.getOrElseUpdate(abortedTxn.producerId, new AbortedTransactionMetadata(abortedTxn))
     }
   }
 
@@ -1082,4 +1163,6 @@ private[log] class CleanedTransactionMetadata(val abortedTransactions: mutable.P
 
 private class AbortedTransactionMetadata(val abortedTxn: AbortedTxn) {
   var lastObservedBatchOffset: Option[Long] = None
+
+  override def toString: String = s"(txn: $abortedTxn, lastOffset: $lastObservedBatchOffset)"
 }
diff --git a/core/src/main/scala/kafka/log/LogCleanerManager.scala b/core/src/main/scala/kafka/log/LogCleanerManager.scala
index e4559b8323cd5..99287cf403fc0 100755
--- a/core/src/main/scala/kafka/log/LogCleanerManager.scala
+++ b/core/src/main/scala/kafka/log/LogCleanerManager.scala
@@ -22,7 +22,7 @@ import java.util.concurrent.TimeUnit
 import java.util.concurrent.locks.ReentrantLock
 
 import com.yammer.metrics.core.Gauge
-import kafka.common.LogCleaningAbortedException
+import kafka.common.{KafkaException, LogCleaningAbortedException}
 import kafka.metrics.KafkaMetricsGroup
 import kafka.server.LogDirFailureChannel
 import kafka.server.checkpoints.OffsetCheckpointFile
@@ -32,13 +32,17 @@ import org.apache.kafka.common.TopicPartition
 import org.apache.kafka.common.utils.Time
 import org.apache.kafka.common.errors.KafkaStorageException
 
-import scala.collection.{Iterable, immutable, mutable}
+import scala.collection.{Iterable, Seq, mutable}
 
 private[log] sealed trait LogCleaningState
 private[log] case object LogCleaningInProgress extends LogCleaningState
 private[log] case object LogCleaningAborted extends LogCleaningState
 private[log] case class LogCleaningPaused(pausedCount: Int) extends LogCleaningState
 
+private[log] class LogCleaningException(val log: Log,
+                                        private val message: String,
+                                        private val cause: Throwable) extends KafkaException(message, cause)
+
 /**
   * This class manages the state of each partition being cleaned.
   * LogCleaningState defines the cleaning states that a TopicPartition can be in.
@@ -57,7 +61,10 @@ private[log] case class LogCleaningPaused(pausedCount: Int) extends LogCleaningS
   */
 private[log] class LogCleanerManager(val logDirs: Seq[File],
                                      val logs: Pool[TopicPartition, Log],
-                                     val logDirFailureChannel: LogDirFailureChannel) extends Logging with KafkaMetricsGroup {
+                                     val logDirFailureChannel: LogDirFailureChannel,
+                                     val retentionCheckMs: Long) extends Logging with KafkaMetricsGroup {
+  import LogCleanerManager._
+
 
   protected override def loggerName = classOf[LogCleaner].getName
 
@@ -81,6 +88,9 @@ private[log] class LogCleanerManager(val logDirs: Seq[File],
   /* for coordinating the pausing and the cleaning of a partition */
   private val pausedCleaningCond = lock.newCondition()
 
+  /* last time the cleaner threads check logs for retention*/
+  private var lastRetentionCheckTime: Long = Time.SYSTEM.milliseconds()
+
   /* gauges for tracking the number of partitions marked as uncleanable for each log directory */
   for (dir <- logDirs) {
     newGauge(
@@ -91,31 +101,32 @@ private[log] class LogCleanerManager(val logDirs: Seq[File],
   }
 
   /* gauges for tracking the number of uncleanable bytes from uncleanable partitions for each log directory */
-    for (dir <- logDirs) {
-      newGauge(
-        "uncleanable-bytes",
-        new Gauge[Long] {
-          def value = {
-            inLock(lock) {
-              uncleanablePartitions.get(dir.getAbsolutePath) match {
-                case Some(partitions) => {
-                  val lastClean = allCleanerCheckpoints
-                  val now = Time.SYSTEM.milliseconds
-                  partitions.map { tp =>
-                    val log = logs.get(tp)
-                    val (firstDirtyOffset, firstUncleanableDirtyOffset) = LogCleanerManager.cleanableOffsets(log, tp, lastClean, now)
-                    val (_, uncleanableBytes) = LogCleaner.calculateCleanableBytes(log, firstDirtyOffset, firstUncleanableDirtyOffset)
-                    uncleanableBytes
-                  }.sum
-                }
-                case _ => 0
+  for (dir <- logDirs) {
+    newGauge(
+      "uncleanable-bytes",
+      new Gauge[Long] {
+        def value = {
+          inLock(lock) {
+            uncleanablePartitions.get(dir.getAbsolutePath) match {
+              case Some(partitions) => {
+                val lastClean = allCleanerCheckpoints
+                val now = Time.SYSTEM.milliseconds
+                partitions.map { tp =>
+                  val log = logs.get(tp)
+                  val lastCleanOffset = lastClean.get(tp)
+                  val offsetsToClean = cleanableOffsets(log, lastCleanOffset, now)
+                  val (_, uncleanableBytes) = calculateCleanableBytes(log, offsetsToClean.firstDirtyOffset, offsetsToClean.firstUncleanableDirtyOffset)
+                  uncleanableBytes
+                }.sum
               }
+              case _ => 0
             }
           }
-        },
-        Map("logDirectory" -> dir.getAbsolutePath)
-      )
-    }
+        }
+      },
+      Map("logDirectory" -> dir.getAbsolutePath)
+    )
+  }
 
   /* a gauge for tracking the cleanable ratio of the dirtiest log */
   @volatile private var dirtiestLogCleanableRatio = 0.0
@@ -165,7 +176,7 @@ private[log] class LogCleanerManager(val logDirs: Seq[File],
     * each time from the full set of logs to allow logs to be dynamically added to the pool of logs
     * the log manager maintains.
     */
-  def grabFilthiestCompactedLog(time: Time): Option[LogToClean] = {
+  def grabFilthiestCompactedLog(time: Time, preCleanStats: PreCleanStats = new PreCleanStats()): Option[LogToClean] = {
     inLock(lock) {
       val now = time.milliseconds
       this.timeOfLastRun = now
@@ -178,17 +189,31 @@ private[log] class LogCleanerManager(val logDirs: Seq[File],
           inProgress.contains(topicPartition) || isUncleanablePartition(log, topicPartition)
       }.map {
         case (topicPartition, log) => // create a LogToClean instance for each
-          val (firstDirtyOffset, firstUncleanableDirtyOffset) = LogCleanerManager.cleanableOffsets(log, topicPartition,
-            lastClean, now)
-          LogToClean(topicPartition, log, firstDirtyOffset, firstUncleanableDirtyOffset)
+          try {
+            val lastCleanOffset = lastClean.get(topicPartition)
+            val offsetsToClean = cleanableOffsets(log, lastCleanOffset, now)
+            // update checkpoint for logs with invalid checkpointed offsets
+            if (offsetsToClean.forceUpdateCheckpoint)
+              updateCheckpoints(log.parentDirFile, Option(topicPartition, offsetsToClean.firstDirtyOffset))
+            val compactionDelayMs = maxCompactionDelay(log, offsetsToClean.firstDirtyOffset, now)
+            preCleanStats.updateMaxCompactionDelay(compactionDelayMs)
+
+            LogToClean(topicPartition, log, offsetsToClean.firstDirtyOffset, offsetsToClean.firstUncleanableDirtyOffset, compactionDelayMs > 0)
+          } catch {
+            case e: Throwable => throw new LogCleaningException(log,
+              s"Failed to calculate log cleaning stats for partition $topicPartition", e)
+          }
       }.filter(ltc => ltc.totalBytes > 0) // skip any empty logs
 
       this.dirtiestLogCleanableRatio = if (dirtyLogs.nonEmpty) dirtyLogs.max.cleanableRatio else 0
-      // and must meet the minimum threshold for dirty byte ratio
-      val cleanableLogs = dirtyLogs.filter(ltc => ltc.cleanableRatio > ltc.log.config.minCleanableRatio)
+      // and must meet the minimum threshold for dirty byte ratio or have some bytes required to be compacted
+      val cleanableLogs = dirtyLogs.filter { ltc =>
+        (ltc.needCompactionNow && ltc.cleanableBytes > 0) || ltc.cleanableRatio > ltc.log.config.minCleanableRatio
+      }
       if(cleanableLogs.isEmpty) {
         None
       } else {
+        preCleanStats.recordCleanablePartitions(cleanableLogs.size)
         val filthiest = cleanableLogs.max
         inProgress.put(filthiest.topicPartition, LogCleaningInProgress)
         Some(filthiest)
@@ -225,12 +250,17 @@ private[log] class LogCleanerManager(val logDirs: Seq[File],
     */
   def deletableLogs(): Iterable[(TopicPartition, Log)] = {
     inLock(lock) {
-      val toClean = logs.filter { case (topicPartition, log) =>
-        !inProgress.contains(topicPartition) && log.config.compact &&
-          !isUncleanablePartition(log, topicPartition)
+      val now = Time.SYSTEM.milliseconds()
+      if (now - lastRetentionCheckTime >= retentionCheckMs) {
+        lastRetentionCheckTime = now
+        val toClean = logs.filter { case (topicPartition, log) =>
+          !inProgress.contains(topicPartition) && log.config.compact && !isUncleanablePartition(log, topicPartition)
+        }
+        toClean.foreach { case (tp, _) => inProgress.put(tp, LogCleaningInProgress) }
+        toClean
       }
-      toClean.foreach { case (tp, _) => inProgress.put(tp, LogCleaningInProgress) }
-      toClean
+      else
+        Iterable.empty
     }
 
   }
@@ -240,7 +270,7 @@ private[log] class LogCleanerManager(val logDirs: Seq[File],
    *  the partition is aborted.
    *  This is implemented by first abortAndPausing and then resuming the cleaning of the partition.
    */
-  def abortCleaning(topicPartition: TopicPartition) {
+  def abortCleaning(topicPartition: TopicPartition): Unit = {
     inLock(lock) {
       abortAndPauseCleaning(topicPartition)
       resumeCleaning(Seq(topicPartition))
@@ -260,7 +290,7 @@ private[log] class LogCleanerManager(val logDirs: Seq[File],
    *  6. If the partition is already paused, a new call to this function
    *     will increase the paused count by one.
    */
-  def abortAndPauseCleaning(topicPartition: TopicPartition) {
+  def abortAndPauseCleaning(topicPartition: TopicPartition): Unit = {
     inLock(lock) {
       inProgress.get(topicPartition) match {
         case None =>
@@ -283,7 +313,7 @@ private[log] class LogCleanerManager(val logDirs: Seq[File],
     *  Resume the cleaning of paused partitions.
     *  Each call of this function will undo one pause.
     */
-  def resumeCleaning(topicPartitions: Iterable[TopicPartition]){
+  def resumeCleaning(topicPartitions: Iterable[TopicPartition]): Unit = {
     inLock(lock) {
       topicPartitions.foreach {
         topicPartition =>
@@ -326,7 +356,7 @@ private[log] class LogCleanerManager(val logDirs: Seq[File],
       case None => false
       case Some(state) =>
         state match {
-          case LogCleaningPaused(s) =>
+          case _: LogCleaningPaused =>
             true
           case _ =>
             false
@@ -337,19 +367,19 @@ private[log] class LogCleanerManager(val logDirs: Seq[File],
   /**
    *  Check if the cleaning for a partition is aborted. If so, throw an exception.
    */
-  def checkCleaningAborted(topicPartition: TopicPartition) {
+  def checkCleaningAborted(topicPartition: TopicPartition): Unit = {
     inLock(lock) {
       if (isCleaningInState(topicPartition, LogCleaningAborted))
         throw new LogCleaningAbortedException()
     }
   }
 
-  def updateCheckpoints(dataDir: File, update: Option[(TopicPartition,Long)]) {
+  def updateCheckpoints(dataDir: File, update: Option[(TopicPartition,Long)]): Unit = {
     inLock(lock) {
       val checkpoint = checkpoints(dataDir)
       if (checkpoint != null) {
         try {
-          val existing = checkpoint.read().filterKeys(logs.keys) ++ update
+          val existing = checkpoint.read().filter { case (k, _) => logs.keys.contains(k) } ++ update
           checkpoint.write(existing)
         } catch {
           case e: KafkaStorageException =>
@@ -366,7 +396,7 @@ private[log] class LogCleanerManager(val logDirs: Seq[File],
           case Some(offset) =>
             // Remove this partition from the checkpoint file in the source log directory
             updateCheckpoints(sourceLogDir, None)
-            // Add offset for this partition to the checkpoint file in the source log directory
+            // Add offset for this partition to the checkpoint file in the destination log directory
             updateCheckpoints(destLogDir, Option(topicPartition, offset))
           case None =>
         }
@@ -383,14 +413,14 @@ private[log] class LogCleanerManager(val logDirs: Seq[File],
     }
   }
 
-  def handleLogDirFailure(dir: String) {
+  def handleLogDirFailure(dir: String): Unit = {
     info(s"Stopping cleaning logs in dir $dir")
     inLock(lock) {
-      checkpoints = checkpoints.filterKeys(_.getAbsolutePath != dir)
+      checkpoints = checkpoints.filter { case (k, _) => k.getAbsolutePath != dir }
     }
   }
 
-  def maybeTruncateCheckpoint(dataDir: File, topicPartition: TopicPartition, offset: Long) {
+  def maybeTruncateCheckpoint(dataDir: File, topicPartition: TopicPartition, offset: Long): Unit = {
     inLock(lock) {
       if (logs.get(topicPartition).config.compact) {
         val checkpoint = checkpoints(dataDir)
@@ -406,7 +436,7 @@ private[log] class LogCleanerManager(val logDirs: Seq[File],
   /**
    * Save out the endOffset and remove the given log from the in-progress set, if not aborted.
    */
-  def doneCleaning(topicPartition: TopicPartition, dataDir: File, endOffset: Long) {
+  def doneCleaning(topicPartition: TopicPartition, dataDir: File, endOffset: Long): Unit = {
     inLock(lock) {
       inProgress.get(topicPartition) match {
         case Some(LogCleaningInProgress) =>
@@ -465,47 +495,86 @@ private[log] class LogCleanerManager(val logDirs: Seq[File],
 
   private def isUncleanablePartition(log: Log, topicPartition: TopicPartition): Boolean = {
     inLock(lock) {
-      uncleanablePartitions.get(log.dir.getParent).exists(partitions => partitions.contains(topicPartition))
+      uncleanablePartitions.get(log.parentDir).exists(partitions => partitions.contains(topicPartition))
     }
   }
 }
 
+/**
+ * Helper class for the range of cleanable dirty offsets of a log and whether to update the checkpoint associated with
+ * the log
+ *
+ * @param firstDirtyOffset the lower (inclusive) offset to begin cleaning from
+ * @param firstUncleanableDirtyOffset the upper(exclusive) offset to clean to
+ * @param forceUpdateCheckpoint whether to update the checkpoint associated with this log. if true, checkpoint should be
+ *                             reset to firstDirtyOffset
+ */
+private case class OffsetsToClean(firstDirtyOffset: Long,
+                                  firstUncleanableDirtyOffset: Long,
+                                  forceUpdateCheckpoint: Boolean = false) {
+}
+
 private[log] object LogCleanerManager extends Logging {
 
+
   def isCompactAndDelete(log: Log): Boolean = {
     log.config.compact && log.config.delete
   }
 
+  /**
+    * get max delay between the time when log is required to be compacted as determined
+    * by maxCompactionLagMs and the current time.
+    */
+  def maxCompactionDelay(log: Log, firstDirtyOffset: Long, now: Long) : Long = {
+    val dirtyNonActiveSegments = log.nonActiveLogSegmentsFrom(firstDirtyOffset)
+    val firstBatchTimestamps = log.getFirstBatchTimestampForSegments(dirtyNonActiveSegments).filter(_ > 0)
+
+    val earliestDirtySegmentTimestamp = {
+      if (firstBatchTimestamps.nonEmpty)
+        firstBatchTimestamps.min
+      else Long.MaxValue
+    }
+
+    val maxCompactionLagMs = math.max(log.config.maxCompactionLagMs, 0L)
+    val cleanUntilTime = now - maxCompactionLagMs
+
+    if (earliestDirtySegmentTimestamp < cleanUntilTime)
+      cleanUntilTime - earliestDirtySegmentTimestamp
+    else
+      0L
+  }
 
   /**
     * Returns the range of dirty offsets that can be cleaned.
     *
     * @param log the log
-    * @param lastClean the map of checkpointed offsets
+    * @param lastCleanOffset the last checkpointed offset
     * @param now the current time in milliseconds of the cleaning operation
-    * @return the lower (inclusive) and upper (exclusive) offsets
+    * @return OffsetsToClean containing offsets for cleanable portion of log and whether the log checkpoint needs updating
     */
-  def cleanableOffsets(log: Log, topicPartition: TopicPartition, lastClean: immutable.Map[TopicPartition, Long], now: Long): (Long, Long) = {
-
-    // the checkpointed offset, ie., the first offset of the next dirty segment
-    val lastCleanOffset: Option[Long] = lastClean.get(topicPartition)
-
+  def cleanableOffsets(log: Log, lastCleanOffset: Option[Long], now: Long): OffsetsToClean = {
     // If the log segments are abnormally truncated and hence the checkpointed offset is no longer valid;
     // reset to the log starting offset and log the error
-    val logStartOffset = log.logSegments.head.baseOffset
-    val firstDirtyOffset = {
-      val offset = lastCleanOffset.getOrElse(logStartOffset)
-      if (offset < logStartOffset) {
-        // don't bother with the warning if compact and delete are enabled.
-        if (!isCompactAndDelete(log))
-          warn(s"Resetting first dirty offset of ${log.name} to log start offset $logStartOffset since the checkpointed offset $offset is invalid.")
-        logStartOffset
+    val (firstDirtyOffset, forceUpdateCheckpoint) = {
+      val logStartOffset = log.logStartOffset
+      val checkpointDirtyOffset = lastCleanOffset.getOrElse(logStartOffset)
+
+      if (checkpointDirtyOffset < logStartOffset) {
+        debug(s"Resetting first dirty offset of ${log.name} to log start offset $logStartOffset " +
+            s"since the checkpointed offset $checkpointDirtyOffset is invalid.")
+        (logStartOffset, true)
+      } else if (checkpointDirtyOffset > log.logEndOffset) {
+        // The dirty offset has gotten ahead of the log end offset. This could happen if there was data
+        // corruption at the end of the log. We conservatively assume that the full log needs cleaning.
+        warn(s"The last checkpoint dirty offset for partition ${log.name} is $checkpointDirtyOffset, " +
+          s"which is larger than the log end offset ${log.logEndOffset}. Resetting to the log start offset $logStartOffset.")
+        (logStartOffset, true)
       } else {
-        offset
+        (checkpointDirtyOffset, false)
       }
     }
 
-    val compactionLagMs = math.max(log.config.compactionLagMs, 0L)
+    val minCompactionLagMs = math.max(log.config.compactionLagMs, 0L)
 
     // find first segment that cannot be cleaned
     // neither the active segment, nor segments with any messages closer to the head of the log than the minimum compaction lag time
@@ -513,25 +582,42 @@ private[log] object LogCleanerManager extends Logging {
     val firstUncleanableDirtyOffset: Long = Seq(
 
       // we do not clean beyond the first unstable offset
-      log.firstUnstableOffset.map(_.messageOffset),
+      log.firstUnstableOffset,
 
       // the active segment is always uncleanable
       Option(log.activeSegment.baseOffset),
 
       // the first segment whose largest message timestamp is within a minimum time lag from now
-      if (compactionLagMs > 0) {
+      if (minCompactionLagMs > 0) {
         // dirty log segments
-        val dirtyNonActiveSegments = log.logSegments(firstDirtyOffset, log.activeSegment.baseOffset)
+        val dirtyNonActiveSegments = log.nonActiveLogSegmentsFrom(firstDirtyOffset)
         dirtyNonActiveSegments.find { s =>
-          val isUncleanable = s.largestTimestamp > now - compactionLagMs
-          debug(s"Checking if log segment may be cleaned: log='${log.name}' segment.baseOffset=${s.baseOffset} segment.largestTimestamp=${s.largestTimestamp}; now - compactionLag=${now - compactionLagMs}; is uncleanable=$isUncleanable")
+          val isUncleanable = s.largestTimestamp > now - minCompactionLagMs
+          debug(s"Checking if log segment may be cleaned: log='${log.name}' segment.baseOffset=${s.baseOffset} " +
+            s"segment.largestTimestamp=${s.largestTimestamp}; now - compactionLag=${now - minCompactionLagMs}; " +
+            s"is uncleanable=$isUncleanable")
           isUncleanable
         }.map(_.baseOffset)
       } else None
     ).flatten.min
 
-    debug(s"Finding range of cleanable offsets for log=${log.name} topicPartition=$topicPartition. Last clean offset=$lastCleanOffset now=$now => firstDirtyOffset=$firstDirtyOffset firstUncleanableOffset=$firstUncleanableDirtyOffset activeSegment.baseOffset=${log.activeSegment.baseOffset}")
+    debug(s"Finding range of cleanable offsets for log=${log.name}. Last clean offset=$lastCleanOffset " +
+      s"now=$now => firstDirtyOffset=$firstDirtyOffset firstUncleanableOffset=$firstUncleanableDirtyOffset " +
+      s"activeSegment.baseOffset=${log.activeSegment.baseOffset}")
+
+    OffsetsToClean(firstDirtyOffset, math.max(firstDirtyOffset, firstUncleanableDirtyOffset), forceUpdateCheckpoint)
+  }
+
+  /**
+   * Given the first dirty offset and an uncleanable offset, calculates the total cleanable bytes for this log
+   * @return the biggest uncleanable offset and the total amount of cleanable bytes
+   */
+  def calculateCleanableBytes(log: Log, firstDirtyOffset: Long, uncleanableOffset: Long): (Long, Long) = {
+    val firstUncleanableSegment = log.nonActiveLogSegmentsFrom(uncleanableOffset).headOption.getOrElse(log.activeSegment)
+    val firstUncleanableOffset = firstUncleanableSegment.baseOffset
+    val cleanableBytes = log.logSegments(math.min(firstDirtyOffset, firstUncleanableOffset), firstUncleanableOffset).map(_.size.toLong).sum
 
-    (firstDirtyOffset, firstUncleanableDirtyOffset)
+    (firstUncleanableOffset, cleanableBytes)
   }
+
 }
diff --git a/core/src/main/scala/kafka/log/LogConfig.scala b/core/src/main/scala/kafka/log/LogConfig.scala
index ab58949fd5d46..ff4c83dfba21f 100755
--- a/core/src/main/scala/kafka/log/LogConfig.scala
+++ b/core/src/main/scala/kafka/log/LogConfig.scala
@@ -46,6 +46,7 @@ object Defaults {
   val FileDeleteDelayMs = kafka.server.Defaults.LogDeleteDelayMs
   val DeleteRetentionMs = kafka.server.Defaults.LogCleanerDeleteRetentionMs
   val MinCompactionLagMs = kafka.server.Defaults.LogCleanerMinCompactionLagMs
+  val MaxCompactionLagMs = kafka.server.Defaults.LogCleanerMaxCompactionLagMs
   val MinCleanableDirtyRatio = kafka.server.Defaults.LogCleanerMinCleanRatio
 
   @deprecated(message = "This is a misleading variable name as it actually refers to the 'delete' cleanup policy. Use " +
@@ -64,6 +65,7 @@ object Defaults {
   val FollowerReplicationThrottledReplicas = Collections.emptyList[String]()
   val MaxIdMapSnapshots = kafka.server.Defaults.MaxIdMapSnapshots
   val MessageDownConversionEnable = kafka.server.Defaults.MessageDownConversionEnable
+  val ProducerBatchDecompressionEnable = kafka.server.Defaults.ProducerBatchDecompressionEnable
 }
 
 case class LogConfig(props: java.util.Map[_, _], overriddenConfigs: Set[String] = Set.empty)
@@ -85,6 +87,7 @@ case class LogConfig(props: java.util.Map[_, _], overriddenConfigs: Set[String]
   val fileDeleteDelayMs = getLong(LogConfig.FileDeleteDelayMsProp)
   val deleteRetentionMs = getLong(LogConfig.DeleteRetentionMsProp)
   val compactionLagMs = getLong(LogConfig.MinCompactionLagMsProp)
+  val maxCompactionLagMs = getLong(LogConfig.MaxCompactionLagMsProp)
   val minCleanableRatio = getDouble(LogConfig.MinCleanableDirtyRatioProp)
   val compact = getList(LogConfig.CleanupPolicyProp).asScala.map(_.toLowerCase(Locale.ROOT)).contains(LogConfig.Compact)
   val delete = getList(LogConfig.CleanupPolicyProp).asScala.map(_.toLowerCase(Locale.ROOT)).contains(LogConfig.Delete)
@@ -98,15 +101,21 @@ case class LogConfig(props: java.util.Map[_, _], overriddenConfigs: Set[String]
   val LeaderReplicationThrottledReplicas = getList(LogConfig.LeaderReplicationThrottledReplicasProp)
   val FollowerReplicationThrottledReplicas = getList(LogConfig.FollowerReplicationThrottledReplicasProp)
   val messageDownConversionEnable = getBoolean(LogConfig.MessageDownConversionEnableProp)
+  val producerBatchDecompressionEnable = getBoolean(LogConfig.ProducerBatchDecompressionEnableProp)
 
   def randomSegmentJitter: Long =
     if (segmentJitterMs == 0) 0 else Utils.abs(scala.util.Random.nextInt()) % math.min(segmentJitterMs, segmentMs)
+
+  def maxSegmentMs: Long = {
+    if (compact && maxCompactionLagMs > 0) math.min(maxCompactionLagMs, segmentMs)
+    else segmentMs
+  }
 }
 
 object LogConfig {
 
-  def main(args: Array[String]) {
-    println(configDef.toHtmlTable)
+  def main(args: Array[String]): Unit = {
+    println(configDef.toHtml)
   }
 
   val SegmentBytesProp = TopicConfig.SEGMENT_BYTES_CONFIG
@@ -121,6 +130,7 @@ object LogConfig {
   val IndexIntervalBytesProp = TopicConfig.INDEX_INTERVAL_BYTES_CONFIG
   val DeleteRetentionMsProp = TopicConfig.DELETE_RETENTION_MS_CONFIG
   val MinCompactionLagMsProp = TopicConfig.MIN_COMPACTION_LAG_MS_CONFIG
+  val MaxCompactionLagMsProp = TopicConfig.MAX_COMPACTION_LAG_MS_CONFIG
   val FileDeleteDelayMsProp = TopicConfig.FILE_DELETE_DELAY_MS_CONFIG
   val MinCleanableDirtyRatioProp = TopicConfig.MIN_CLEANABLE_DIRTY_RATIO_CONFIG
   val CleanupPolicyProp = TopicConfig.CLEANUP_POLICY_CONFIG
@@ -152,6 +162,7 @@ object LogConfig {
   val FileDeleteDelayMsDoc = TopicConfig.FILE_DELETE_DELAY_MS_DOC
   val DeleteRetentionMsDoc = TopicConfig.DELETE_RETENTION_MS_DOC
   val MinCompactionLagMsDoc = TopicConfig.MIN_COMPACTION_LAG_MS_DOC
+  val MaxCompactionLagMsDoc = TopicConfig.MAX_COMPACTION_LAG_MS_DOC
   val MinCleanableRatioDoc = TopicConfig.MIN_CLEANABLE_DIRTY_RATIO_DOC
   val CompactDoc = TopicConfig.CLEANUP_POLICY_DOC
   val UncleanLeaderElectionEnableDoc = TopicConfig.UNCLEAN_LEADER_ELECTION_ENABLE_DOC
@@ -163,6 +174,8 @@ object LogConfig {
   val MessageTimestampDifferenceMaxMsDoc = TopicConfig.MESSAGE_TIMESTAMP_DIFFERENCE_MAX_MS_DOC
   val MessageDownConversionEnableDoc = TopicConfig.MESSAGE_DOWNCONVERSION_ENABLE_DOC
 
+  val ProducerBatchDecompressionEnableProp = "producer.batch.decompression.enable"
+
   val LeaderReplicationThrottledReplicasDoc = "A list of replicas for which log replication should be throttled on " +
     "the leader side. The list should describe a set of replicas in the form " +
     "[PartitionId]:[BrokerId],[PartitionId]:[BrokerId]:... or alternatively the wildcard '*' can be used to throttle " +
@@ -172,9 +185,24 @@ object LogConfig {
     "[PartitionId]:[BrokerId],[PartitionId]:[BrokerId]:... or alternatively the wildcard '*' can be used to throttle " +
     "all replicas for this topic."
 
-  private class LogConfigDef extends ConfigDef {
+  val ProducerBatchDecompressionEnableDoc = "This configuration controls whether a compressed batch sent by a producer " +
+    "will be decompressed to perform validation of individual records in the batch. The default policy is to decompress " +
+    "all batches.  When set to false, decompression will be skipped if the following three conditions are met : " +
+    "i) Batch sent by producer is greater than record format version V1  ii) Record format version of batch sent by producer " +
+    "is the same as the broker message format version. iii) The relative offsets of the records in the compressed batch are " +
+    "monotonically increasing by 1 starting from 0"
+
+  private[log] val ServerDefaultHeaderName = "Server Default Property"
+
+  // Package private for testing
+  private[log] class LogConfigDef(base: ConfigDef) extends ConfigDef(base) {
+    def this() = this(new ConfigDef)
 
     private final val serverDefaultConfigNames = mutable.Map[String, String]()
+    base match {
+      case b: LogConfigDef => serverDefaultConfigNames ++= b.serverDefaultConfigNames
+      case _ =>
+    }
 
     def define(name: String, defType: ConfigDef.Type, defaultValue: Any, validator: Validator,
                importance: ConfigDef.Importance, doc: String, serverDefaultConfigName: String): LogConfigDef = {
@@ -197,11 +225,12 @@ object LogConfig {
       this
     }
 
-    override def headers = List("Name", "Description", "Type", "Default", "Valid Values", "Server Default Property", "Importance").asJava
+    override def headers = List("Name", "Description", "Type", "Default", "Valid Values", ServerDefaultHeaderName,
+      "Importance").asJava
 
     override def getConfigValue(key: ConfigKey, headerName: String): String = {
       headerName match {
-        case "Server Default Property" => serverDefaultConfigNames.get(key.name).get
+        case ServerDefaultHeaderName => serverDefaultConfigNames.getOrElse(key.name, null)
         case _ => super.getConfigValue(key, headerName)
       }
     }
@@ -209,6 +238,9 @@ object LogConfig {
     def serverConfigName(configName: String): Option[String] = serverDefaultConfigNames.get(configName)
   }
 
+  // Package private for testing, return a copy since it's a mutable global variable
+  private[log] def configDefCopy: LogConfigDef = new LogConfigDef(configDef)
+
   private val configDef: LogConfigDef = {
     import org.apache.kafka.common.config.ConfigDef.Importance._
     import org.apache.kafka.common.config.ConfigDef.Range._
@@ -242,6 +274,8 @@ object LogConfig {
         DeleteRetentionMsDoc, KafkaConfig.LogCleanerDeleteRetentionMsProp)
       .define(MinCompactionLagMsProp, LONG, Defaults.MinCompactionLagMs, atLeast(0), MEDIUM, MinCompactionLagMsDoc,
         KafkaConfig.LogCleanerMinCompactionLagMsProp)
+      .define(MaxCompactionLagMsProp, LONG, Defaults.MaxCompactionLagMs, atLeast(1), MEDIUM, MaxCompactionLagMsDoc,
+        KafkaConfig.LogCleanerMaxCompactionLagMsProp)
       .define(FileDeleteDelayMsProp, LONG, Defaults.FileDeleteDelayMs, atLeast(0), MEDIUM, FileDeleteDelayMsDoc,
         KafkaConfig.LogDeleteDelayMsProp)
       .define(MinCleanableDirtyRatioProp, DOUBLE, Defaults.MinCleanableDirtyRatio, between(0, 1), MEDIUM,
@@ -268,6 +302,8 @@ object LogConfig {
         FollowerReplicationThrottledReplicasDoc, FollowerReplicationThrottledReplicasProp)
       .define(MessageDownConversionEnableProp, BOOLEAN, Defaults.MessageDownConversionEnable, LOW,
         MessageDownConversionEnableDoc, KafkaConfig.LogMessageDownConversionEnableProp)
+      .define(ProducerBatchDecompressionEnableProp, BOOLEAN, Defaults.ProducerBatchDecompressionEnable, LOW,
+        ProducerBatchDecompressionEnableDoc, KafkaConfig.ProducerBatchDecompressionEnableProp)
   }
 
   def apply(): LogConfig = LogConfig(new Properties())
@@ -290,19 +326,31 @@ object LogConfig {
   /**
    * Check that property names are valid
    */
-  def validateNames(props: Properties) {
+  def validateNames(props: Properties): Unit = {
     val names = configNames
     for(name <- props.asScala.keys)
       if (!names.contains(name))
         throw new InvalidConfigurationException(s"Unknown topic config name: $name")
   }
 
+  private[kafka] def configKeys: Map[String, ConfigKey] = configDef.configKeys.asScala
+
+  def validateValues(props: java.util.Map[_, _]): Unit = {
+    val minCompactionLag =  props.get(MinCompactionLagMsProp).asInstanceOf[Long]
+    val maxCompactionLag =  props.get(MaxCompactionLagMsProp).asInstanceOf[Long]
+    if (minCompactionLag > maxCompactionLag) {
+      throw new InvalidConfigurationException(s"conflict topic config setting $MinCompactionLagMsProp " +
+        s"($minCompactionLag) > $MaxCompactionLagMsProp ($maxCompactionLag)")
+    }
+  }
+
   /**
    * Check that the given properties contain only valid log config names and that all values can be parsed and are valid
    */
-  def validate(props: Properties) {
+  def validate(props: Properties): Unit = {
     validateNames(props)
-    configDef.parse(props)
+    val valueMaps = configDef.parse(props)
+    validateValues(valueMaps)
   }
 
   /**
@@ -322,6 +370,7 @@ object LogConfig {
     IndexIntervalBytesProp -> KafkaConfig.LogIndexIntervalBytesProp,
     DeleteRetentionMsProp -> KafkaConfig.LogCleanerDeleteRetentionMsProp,
     MinCompactionLagMsProp -> KafkaConfig.LogCleanerMinCompactionLagMsProp,
+    MaxCompactionLagMsProp -> KafkaConfig.LogCleanerMaxCompactionLagMsProp,
     FileDeleteDelayMsProp -> KafkaConfig.LogDeleteDelayMsProp,
     MinCleanableDirtyRatioProp -> KafkaConfig.LogCleanerMinCleanRatioProp,
     CleanupPolicyProp -> KafkaConfig.LogCleanupPolicyProp,
diff --git a/core/src/main/scala/kafka/log/LogManager.scala b/core/src/main/scala/kafka/log/LogManager.scala
index cae47f724e679..326feb99d911c 100755
--- a/core/src/main/scala/kafka/log/LogManager.scala
+++ b/core/src/main/scala/kafka/log/LogManager.scala
@@ -34,6 +34,7 @@ import org.apache.kafka.common.errors.{KafkaStorageException, LogDirNotFoundExce
 import scala.collection.JavaConverters._
 import scala.collection._
 import scala.collection.mutable.ArrayBuffer
+import scala.util.{Failure, Success, Try}
 
 /**
  * The entry point to the kafka log management subsystem. The log manager is responsible for log creation, retrieval, and cleaning.
@@ -81,6 +82,13 @@ class LogManager(logDirs: Seq[File],
   @volatile private var _currentDefaultConfig = initialDefaultConfig
   @volatile private var numRecoveryThreadsPerDataDir = recoveryThreadsPerDataDir
 
+  // This map contains all partitions whose logs are getting loaded and initialized. If log configuration
+  // of these partitions get updated at the same time, the corresponding entry in this map is set to "true",
+  // which triggers a config reload after initialization is finished (to get the latest config value).
+  // See KAFKA-8813 for more detail on the race condition
+  // Visible for testing
+  private[log] val partitionsInitializing = new ConcurrentHashMap[TopicPartition, Boolean]().asScala
+
   def reconfigureDefaultLogConfig(logConfig: LogConfig): Unit = {
     this._currentDefaultConfig = logConfig
   }
@@ -103,7 +111,7 @@ class LogManager(logDirs: Seq[File],
   private val preferredLogDirs = new ConcurrentHashMap[TopicPartition, String]()
 
   private def offlineLogDirs: Iterable[File] = {
-    val logDirsSet = mutable.Set[File](logDirs: _*)
+    val logDirsSet = mutable.Set[File]() ++= logDirs
     _liveLogDirs.asScala.foreach(logDirsSet -=)
     logDirsSet
   }
@@ -113,7 +121,7 @@ class LogManager(logDirs: Seq[File],
   // public, so we can access this from kafka.admin.DeleteTopicTest
   val cleaner: LogCleaner =
     if(cleanerConfig.enableCleaner)
-      new LogCleaner(cleanerConfig, liveLogDirs, currentLogs, logDirFailureChannel, time = time)
+      new LogCleaner(cleanerConfig, liveLogDirs, currentLogs, logDirFailureChannel, retentionCheckMs, time = time)
     else
       null
 
@@ -187,7 +195,7 @@ class LogManager(logDirs: Seq[File],
   }
 
   // dir should be an absolute path
-  def handleLogDirFailure(dir: String) {
+  def handleLogDirFailure(dir: String): Unit = {
     info(s"Stopping serving logs in dir $dir")
     logCreationOrDeletionLock synchronized {
       _liveLogDirs.remove(new File(dir))
@@ -202,7 +210,7 @@ class LogManager(logDirs: Seq[File],
         cleaner.handleLogDirFailure(dir)
 
       val offlineCurrentTopicPartitions = currentLogs.collect {
-        case (tp, log) if log.dir.getParent == dir => tp
+        case (tp, log) if log.parentDir == dir => tp
       }
       offlineCurrentTopicPartitions.foreach { topicPartition => {
         val removedLog = currentLogs.remove(topicPartition)
@@ -213,7 +221,7 @@ class LogManager(logDirs: Seq[File],
       }}
 
       val offlineFutureTopicPartitions = futureLogs.collect {
-        case (tp, log) if log.dir.getParent == dir => tp
+        case (tp, log) if log.parentDir == dir => tp
       }
       offlineFutureTopicPartitions.foreach { topicPartition => {
         val removedLog = futureLogs.remove(topicPartition)
@@ -285,7 +293,7 @@ class LogManager(logDirs: Seq[File],
       }
       if (previous != null) {
         if (log.isFuture)
-          throw new IllegalStateException("Duplicate log directories found: %s, %s!".format(log.dir.getAbsolutePath, previous.dir.getAbsolutePath))
+          throw new IllegalStateException(s"Duplicate log directories found: ${log.dir.getAbsolutePath}, ${previous.dir.getAbsolutePath}")
         else
           throw new IllegalStateException(s"Duplicate log directories for $topicPartition are found in both ${log.dir.getAbsolutePath} " +
             s"and ${previous.dir.getAbsolutePath}. It is likely because log directory failure happened while broker was " +
@@ -387,7 +395,7 @@ class LogManager(logDirs: Seq[File],
   /**
    *  Start the background threads to flush logs and do log cleanup
    */
-  def startup() {
+  def startup(): Unit = {
     /* Schedule the cleanup task to delete old logs */
     if (scheduler != null) {
       info("Starting log cleanup with a period of %d ms.".format(retentionCheckMs))
@@ -424,7 +432,7 @@ class LogManager(logDirs: Seq[File],
   /**
    * Close all the logs
    */
-  def shutdown() {
+  def shutdown(): Unit = {
     info("Shutting down.")
 
     removeMetric("OfflineLogDirectoryCount")
@@ -496,7 +504,7 @@ class LogManager(logDirs: Seq[File],
    * @param partitionOffsets Partition logs that need to be truncated
    * @param isFuture True iff the truncation should be performed on the future log of the specified partitions
    */
-  def truncateTo(partitionOffsets: Map[TopicPartition, Long], isFuture: Boolean) {
+  def truncateTo(partitionOffsets: Map[TopicPartition, Long], isFuture: Boolean): Unit = {
     val affectedLogs = ArrayBuffer.empty[Log]
     for ((topicPartition, truncateOffset) <- partitionOffsets) {
       val log = {
@@ -515,7 +523,7 @@ class LogManager(logDirs: Seq[File],
           if (log.truncateTo(truncateOffset))
             affectedLogs += log
           if (needToStopCleaner && !isFuture)
-            cleaner.maybeTruncateCheckpoint(log.dir.getParentFile, topicPartition, log.activeSegment.baseOffset)
+            cleaner.maybeTruncateCheckpoint(log.parentDirFile, topicPartition, log.activeSegment.baseOffset)
         } finally {
           if (needToStopCleaner && !isFuture) {
             cleaner.resumeCleaning(Seq(topicPartition))
@@ -525,8 +533,8 @@ class LogManager(logDirs: Seq[File],
       }
     }
 
-    for ((dir, logs) <- affectedLogs.groupBy(_.dir.getParentFile)) {
-      checkpointRecoveryOffsetsAndCleanSnapshot(dir, logs)
+    for ((dir, logs) <- affectedLogs.groupBy(_.parentDirFile)) {
+      checkpointRecoveryOffsetsAndCleanSnapshot(dir, ArrayBuffer.empty)
     }
   }
 
@@ -537,7 +545,7 @@ class LogManager(logDirs: Seq[File],
    * @param newOffset The new offset to start the log with
    * @param isFuture True iff the truncation should be performed on the future log of the specified partition
    */
-  def truncateFullyAndStartAt(topicPartition: TopicPartition, newOffset: Long, isFuture: Boolean) {
+  def truncateFullyAndStartAt(topicPartition: TopicPartition, newOffset: Long, isFuture: Boolean): Unit = {
     val log = {
       if (isFuture)
         futureLogs.get(topicPartition)
@@ -552,7 +560,7 @@ class LogManager(logDirs: Seq[File],
       try {
         log.truncateFullyAndStartAt(newOffset)
         if (cleaner != null && !isFuture) {
-          cleaner.maybeTruncateCheckpoint(log.dir.getParentFile, topicPartition, log.activeSegment.baseOffset)
+          cleaner.maybeTruncateCheckpoint(log.parentDirFile, topicPartition, log.activeSegment.baseOffset)
         }
       } finally {
         if (cleaner != null && !isFuture) {
@@ -560,7 +568,7 @@ class LogManager(logDirs: Seq[File],
           info(s"Compaction for partition $topicPartition is resumed")
         }
       }
-      checkpointRecoveryOffsetsAndCleanSnapshot(log.dir.getParentFile, Seq(log))
+      checkpointRecoveryOffsetsAndCleanSnapshot(log.parentDirFile, Seq(log))
     }
   }
 
@@ -568,7 +576,7 @@ class LogManager(logDirs: Seq[File],
    * Write out the current recovery point for all logs to a text file in the log directory
    * to avoid recovering the whole log on startup.
    */
-  def checkpointLogRecoveryOffsets() {
+  def checkpointLogRecoveryOffsets(): Unit = {
     logsByDir.foreach { case (dir, partitionToLogMap) =>
       liveLogDirs.find(_.getAbsolutePath.equals(dir)).foreach { f =>
         checkpointRecoveryOffsetsAndCleanSnapshot(f, partitionToLogMap.values.toSeq)
@@ -580,7 +588,7 @@ class LogManager(logDirs: Seq[File],
    * Write out the current log start offset for all logs to a text file in the log directory
    * to avoid exposing data that have been deleted by DeleteRecordsRequest
    */
-  def checkpointLogStartOffsets() {
+  def checkpointLogStartOffsets(): Unit = {
     liveLogDirs.foreach(checkpointLogStartOffsetsInDir)
   }
 
@@ -607,7 +615,7 @@ class LogManager(logDirs: Seq[File],
       partitionToLog <- logsByDir.get(dir.getAbsolutePath)
       checkpoint <- recoveryPointCheckpoints.get(dir)
     } {
-      checkpoint.write(partitionToLog.mapValues(_.recoveryPoint))
+      checkpoint.write(partitionToLog.map { case (tp, log) => tp -> log.recoveryPoint })
     }
   }
 
@@ -620,9 +628,9 @@ class LogManager(logDirs: Seq[File],
       checkpoint <- logStartOffsetCheckpoints.get(dir)
     } {
       try {
-        val logStartOffsets = partitionToLog.filter { case (_, log) =>
-          log.logStartOffset > log.logSegments.head.baseOffset
-        }.mapValues(_.logStartOffset)
+        val logStartOffsets = partitionToLog.collect {
+          case (k, log) if log.logStartOffset > log.logSegments.head.baseOffset => k -> log.logStartOffset
+        }
         checkpoint.write(logStartOffsets)
       } catch {
         case e: IOException =>
@@ -634,8 +642,8 @@ class LogManager(logDirs: Seq[File],
   // The logDir should be an absolute path
   def maybeUpdatePreferredLogDir(topicPartition: TopicPartition, logDir: String): Unit = {
     // Do not cache the preferred log directory if either the current log or the future log for this partition exists in the specified logDir
-    if (!getLog(topicPartition).exists(_.dir.getParent == logDir) &&
-        !getLog(topicPartition, isFuture = true).exists(_.dir.getParent == logDir))
+    if (!getLog(topicPartition).exists(_.parentDir == logDir) &&
+        !getLog(topicPartition, isFuture = true).exists(_.parentDir == logDir))
       preferredLogDirs.put(topicPartition, logDir)
   }
 
@@ -658,81 +666,140 @@ class LogManager(logDirs: Seq[File],
       Option(currentLogs.get(topicPartition))
   }
 
+  /**
+   * Method to indicate that logs are getting initialized for the partition passed in as argument.
+   * This method should always be followed by [[kafka.log.LogManager#finishedInitializingLog]] to indicate that log
+   * initialization is done.
+   */
+  def initializingLog(topicPartition: TopicPartition): Unit = {
+    partitionsInitializing(topicPartition) = false
+  }
+
+  /**
+   * Mark the partition configuration for all partitions that are getting initialized for topic
+   * as dirty. That will result in reloading of configuration once initialization is done.
+   */
+  def topicConfigUpdated(topic: String): Unit = {
+    partitionsInitializing.keys.filter(_.topic() == topic).foreach {
+      topicPartition => partitionsInitializing.replace(topicPartition, false, true)
+    }
+  }
+
+  /**
+   * Mark all in progress partitions having dirty configuration if broker configuration is updated.
+   */
+  def brokerConfigUpdated(): Unit = {
+    partitionsInitializing.keys.foreach {
+      topicPartition => partitionsInitializing.replace(topicPartition, false, true)
+    }
+  }
+
+  /**
+   * Method to indicate that the log initialization for the partition passed in as argument is
+   * finished. This method should follow a call to [[kafka.log.LogManager#initializingLog]].
+   *
+   * It will retrieve the topic configs a second time if they were updated while the
+   * relevant log was being loaded.
+   */
+  def finishedInitializingLog(topicPartition: TopicPartition,
+                              maybeLog: Option[Log],
+                              fetchLogConfig: () => LogConfig): Unit = {
+    val removedValue = partitionsInitializing.remove(topicPartition)
+    if (removedValue.contains(true))
+      maybeLog.foreach(_.updateConfig(fetchLogConfig()))
+  }
+
   /**
    * If the log already exists, just return a copy of the existing log
    * Otherwise if isNew=true or if there is no offline log directory, create a log for the given topic and the given partition
    * Otherwise throw KafkaStorageException
    *
    * @param topicPartition The partition whose log needs to be returned or created
-   * @param config The configuration of the log that should be applied for log creation
+   * @param loadConfig A function to retrieve the log config, this is only called if the log is created
    * @param isNew Whether the replica should have existed on the broker or not
    * @param isFuture True iff the future log of the specified partition should be returned or created
    * @throws KafkaStorageException if isNew=false, log is not found in the cache and there is offline log directory on the broker
    */
-  def getOrCreateLog(topicPartition: TopicPartition, config: LogConfig, isNew: Boolean = false, isFuture: Boolean = false): Log = {
+  def getOrCreateLog(topicPartition: TopicPartition, loadConfig: () => LogConfig, isNew: Boolean = false, isFuture: Boolean = false): Log = {
     logCreationOrDeletionLock synchronized {
       getLog(topicPartition, isFuture).getOrElse {
         // create the log if it has not already been created in another thread
         if (!isNew && offlineLogDirs.nonEmpty)
           throw new KafkaStorageException(s"Can not create log for $topicPartition because log directories ${offlineLogDirs.mkString(",")} are offline")
 
-        val logDir = {
+        val logDirs: List[File] = {
           val preferredLogDir = preferredLogDirs.get(topicPartition)
 
           if (isFuture) {
             if (preferredLogDir == null)
               throw new IllegalStateException(s"Can not create the future log for $topicPartition without having a preferred log directory")
-            else if (getLog(topicPartition).get.dir.getParent == preferredLogDir)
+            else if (getLog(topicPartition).get.parentDir == preferredLogDir)
               throw new IllegalStateException(s"Can not create the future log for $topicPartition in the current log directory of this partition")
           }
 
           if (preferredLogDir != null)
-            preferredLogDir
+            List(new File(preferredLogDir))
           else
-            nextLogDir().getAbsolutePath
+            nextLogDirs()
         }
-        if (!isLogDirOnline(logDir))
-          throw new KafkaStorageException(s"Can not create log for $topicPartition because log directory $logDir is offline")
-
-        try {
-          val dir = {
-            if (isFuture)
-              new File(logDir, Log.logFutureDirName(topicPartition))
-            else
-              new File(logDir, Log.logDirName(topicPartition))
-          }
-          Files.createDirectories(dir.toPath)
-
-          val log = Log(
-            dir = dir,
-            config = config,
-            logStartOffset = 0L,
-            recoveryPoint = 0L,
-            maxProducerIdExpirationMs = maxPidExpirationMs,
-            producerIdExpirationCheckIntervalMs = LogManager.ProducerIdExpirationCheckIntervalMs,
-            scheduler = scheduler,
-            time = time,
-            brokerTopicStats = brokerTopicStats,
-            logDirFailureChannel = logDirFailureChannel)
 
+        val logDirName = {
           if (isFuture)
-            futureLogs.put(topicPartition, log)
+            Log.logFutureDirName(topicPartition)
           else
-            currentLogs.put(topicPartition, log)
+            Log.logDirName(topicPartition)
+        }
 
-          info(s"Created log for partition $topicPartition in $logDir with properties " +
-            s"{${config.originals.asScala.mkString(", ")}}.")
-          // Remove the preferred log dir since it has already been satisfied
-          preferredLogDirs.remove(topicPartition)
+        val logDir = logDirs
+          .toStream // to prevent actually mapping the whole list, lazy map
+          .map(createLogDirectory(_, logDirName))
+          .find(_.isSuccess)
+          .getOrElse(Failure(new KafkaStorageException("No log directories available. Tried " + logDirs.map(_.getAbsolutePath).mkString(", "))))
+          .get // If Failure, will throw
+
+        val config = loadConfig()
+        val log = Log(
+          dir = logDir,
+          config = config,
+          logStartOffset = 0L,
+          recoveryPoint = 0L,
+          maxProducerIdExpirationMs = maxPidExpirationMs,
+          producerIdExpirationCheckIntervalMs = LogManager.ProducerIdExpirationCheckIntervalMs,
+          scheduler = scheduler,
+          time = time,
+          brokerTopicStats = brokerTopicStats,
+          logDirFailureChannel = logDirFailureChannel)
 
-          log
-        } catch {
-          case e: IOException =>
-            val msg = s"Error while creating log for $topicPartition in dir $logDir"
-            logDirFailureChannel.maybeAddOfflineLogDir(logDir, msg, e)
-            throw new KafkaStorageException(msg, e)
-        }
+        if (isFuture)
+          futureLogs.put(topicPartition, log)
+        else
+          currentLogs.put(topicPartition, log)
+
+        info(s"Created log for partition $topicPartition in $logDir with properties " + s"{${config.originals.asScala.mkString(", ")}}.")
+        // Remove the preferred log dir since it has already been satisfied
+        preferredLogDirs.remove(topicPartition)
+
+        log
+      }
+    }
+  }
+
+  private[log] def createLogDirectory(logDir: File, logDirName: String): Try[File] = {
+    val logDirPath = logDir.getAbsolutePath
+    if (isLogDirOnline(logDirPath)) {
+      val dir = new File(logDirPath, logDirName)
+      try {
+        Files.createDirectories(dir.toPath)
+        Success(dir)
+      } catch {
+        case e: IOException =>
+          val msg = s"Error while creating log for $logDirName in dir $logDirPath"
+          logDirFailureChannel.maybeAddOfflineLogDir(logDirPath, msg, e)
+          warn(msg, e)
+          Failure(new KafkaStorageException(msg, e))
       }
+    } else {
+      Failure(new KafkaStorageException(s"Can not create log $logDirName because log directory $logDirPath is offline"))
     }
   }
 
@@ -762,7 +829,7 @@ class LogManager(logDirs: Seq[File],
             info(s"Deleted log for partition ${removedLog.topicPartition} in ${removedLog.dir.getAbsolutePath}.")
           } catch {
             case e: KafkaStorageException =>
-              error(s"Exception while deleting $removedLog in dir ${removedLog.dir.getParent}.", e)
+              error(s"Exception while deleting $removedLog in dir ${removedLog.parentDir}.", e)
           }
         }
       }
@@ -803,12 +870,14 @@ class LogManager(logDirs: Seq[File],
         throw new KafkaStorageException(s"The future replica for $topicPartition is offline")
 
       destLog.renameDir(Log.logDirName(topicPartition))
+      destLog.updateHighWatermark(sourceLog.highWatermark)
+
       // Now that future replica has been successfully renamed to be the current replica
       // Update the cached map and log cleaner as appropriate.
       futureLogs.remove(topicPartition)
       currentLogs.put(topicPartition, destLog)
       if (cleaner != null) {
-        cleaner.alterCheckpointDir(topicPartition, sourceLog.dir.getParentFile, destLog.dir.getParentFile)
+        cleaner.alterCheckpointDir(topicPartition, sourceLog.parentDirFile, destLog.parentDirFile)
         cleaner.resumeCleaning(Seq(topicPartition))
         info(s"Compaction for partition $topicPartition is resumed")
       }
@@ -818,8 +887,8 @@ class LogManager(logDirs: Seq[File],
         // Now that replica in source log directory has been successfully renamed for deletion.
         // Close the log, update checkpoint files, and enqueue this log to be deleted.
         sourceLog.close()
-        checkpointRecoveryOffsetsAndCleanSnapshot(sourceLog.dir.getParentFile, ArrayBuffer.empty)
-        checkpointLogStartOffsetsInDir(sourceLog.dir.getParentFile)
+        checkpointRecoveryOffsetsAndCleanSnapshot(sourceLog.parentDirFile, ArrayBuffer.empty)
+        checkpointLogStartOffsetsInDir(sourceLog.parentDirFile)
         addLogToBeDeleted(sourceLog)
       } catch {
         case e: KafkaStorageException =>
@@ -853,11 +922,11 @@ class LogManager(logDirs: Seq[File],
       //We need to wait until there is no more cleaning task on the log to be deleted before actually deleting it.
       if (cleaner != null && !isFuture) {
         cleaner.abortCleaning(topicPartition)
-        cleaner.updateCheckpoints(removedLog.dir.getParentFile)
+        cleaner.updateCheckpoints(removedLog.parentDirFile)
       }
       removedLog.renameDir(Log.logDeleteDirName(topicPartition))
-      checkpointRecoveryOffsetsAndCleanSnapshot(removedLog.dir.getParentFile, ArrayBuffer.empty)
-      checkpointLogStartOffsetsInDir(removedLog.dir.getParentFile)
+      checkpointRecoveryOffsetsAndCleanSnapshot(removedLog.parentDirFile, ArrayBuffer.empty)
+      checkpointLogStartOffsetsInDir(removedLog.parentDirFile)
       addLogToBeDeleted(removedLog)
       info(s"Log for partition ${removedLog.topicPartition} is renamed to ${removedLog.dir.getAbsolutePath} and is scheduled for deletion")
     } else if (offlineLogDirs.nonEmpty) {
@@ -867,22 +936,23 @@ class LogManager(logDirs: Seq[File],
   }
 
   /**
-   * Choose the next directory in which to create a log. Currently this is done
-   * by calculating the number of partitions in each directory and then choosing the
-   * data directory with the fewest partitions.
+   * Provides the full ordered list of suggested directories for the next partition.
+   * Currently this is done by calculating the number of partitions in each directory and then sorting the
+   * data directories by fewest partitions.
    */
-  private def nextLogDir(): File = {
+  private def nextLogDirs(): List[File] = {
     if(_liveLogDirs.size == 1) {
-      _liveLogDirs.peek()
+      List(_liveLogDirs.peek())
     } else {
       // count the number of logs in each parent directory (including 0 for empty directories
-      val logCounts = allLogs.groupBy(_.dir.getParent).mapValues(_.size)
+      val logCounts = allLogs.groupBy(_.parentDir).mapValues(_.size)
       val zeros = _liveLogDirs.asScala.map(dir => (dir.getPath, 0)).toMap
       val dirCounts = (zeros ++ logCounts).toBuffer
 
       // choose the directory with the least logs in it
-      val leastLoaded = dirCounts.sortBy(_._2).head
-      new File(leastLoaded._1)
+      dirCounts.sortBy(_._2).map {
+        case (path: String, _: Int) => new File(path)
+      }.toList
     }
   }
 
@@ -890,7 +960,7 @@ class LogManager(logDirs: Seq[File],
    * Delete any eligible logs. Return the number of segments deleted.
    * Only consider logs that are not compacted.
    */
-  def cleanupLogs() {
+  def cleanupLogs(): Unit = {
     debug("Beginning log cleanup...")
     var total = 0
     val startMs = time.milliseconds
@@ -936,9 +1006,9 @@ class LogManager(logDirs: Seq[File],
   def allLogs: Iterable[Log] = currentLogs.values ++ futureLogs.values
 
   def logsByTopic(topic: String): Seq[Log] = {
-    (currentLogs.toList ++ futureLogs.toList).filter { case (topicPartition, _) =>
-      topicPartition.topic() == topic
-    }.map { case (_, log) => log }
+    (currentLogs.toList ++ futureLogs.toList).collect {
+      case (topicPartition, log) if topicPartition.topic == topic => log
+    }
   }
 
   /**
@@ -946,7 +1016,7 @@ class LogManager(logDirs: Seq[File],
    */
   private def logsByDir: Map[String, Map[TopicPartition, Log]] = {
     (this.currentLogs.toList ++ this.futureLogs.toList).toMap
-      .groupBy { case (_, log) => log.dir.getParent }
+      .groupBy { case (_, log) => log.parentDir }
   }
 
   // logDir should be an absolute path
@@ -994,10 +1064,15 @@ object LogManager {
             brokerTopicStats: BrokerTopicStats,
             logDirFailureChannel: LogDirFailureChannel): LogManager = {
     val defaultProps = KafkaServer.copyKafkaConfigToLog(config)
+
+    LogConfig.validateValues(defaultProps)
     val defaultLogConfig = LogConfig(defaultProps)
 
     // read the log configurations from zookeeper
-    val (topicConfigs, failed) = zkClient.getLogConfigs(zkClient.getAllTopicsInCluster, defaultProps)
+    val (topicConfigs, failed) = zkClient.getLogConfigs(
+      zkClient.getAllTopicsInCluster,
+      defaultProps
+    )
     if (!failed.isEmpty) throw failed.head._2
 
     val cleanerConfig = LogCleaner.cleanerConfig(config)
@@ -1012,7 +1087,7 @@ object LogManager {
       flushRecoveryOffsetCheckpointMs = config.logFlushOffsetCheckpointIntervalMs,
       flushStartOffsetCheckpointMs = config.logFlushStartOffsetCheckpointIntervalMs,
       retentionCheckMs = config.logCleanupIntervalMs,
-      maxPidExpirationMs = config.transactionIdExpirationMs,
+      maxPidExpirationMs = config.transactionalIdExpirationMs,
       scheduler = kafkaScheduler,
       brokerState = brokerState,
       brokerTopicStats = brokerTopicStats,
diff --git a/core/src/main/scala/kafka/log/LogSegment.scala b/core/src/main/scala/kafka/log/LogSegment.scala
index 9168ce01dd5cb..6d3beb7107617 100755
--- a/core/src/main/scala/kafka/log/LogSegment.scala
+++ b/core/src/main/scala/kafka/log/LogSegment.scala
@@ -26,6 +26,7 @@ import kafka.metrics.{KafkaMetricsGroup, KafkaTimer}
 import kafka.server.epoch.LeaderEpochFileCache
 import kafka.server.{FetchDataInfo, LogOffsetMetadata}
 import kafka.utils._
+import org.apache.kafka.common.InvalidRecordException
 import org.apache.kafka.common.errors.CorruptRecordException
 import org.apache.kafka.common.record.FileRecords.{LogOffsetPosition, TimestampAndOffset}
 import org.apache.kafka.common.record._
@@ -43,8 +44,8 @@ import scala.math._
  * A segment with a base offset of [base_offset] would be stored in two files, a [base_offset].index and a [base_offset].log file.
  *
  * @param log The file records containing log entries
- * @param offsetIndex The offset index
- * @param timeIndex The timestamp index
+ * @param lazyOffsetIndex The offset index
+ * @param lazyTimeIndex The timestamp index
  * @param txnIndex The transaction index
  * @param baseOffset A lower bound on the offsets in this segment
  * @param indexIntervalBytes The approximate number of bytes between entries in the index
@@ -53,8 +54,8 @@ import scala.math._
  */
 @nonthreadsafe
 class LogSegment private[log] (val log: FileRecords,
-                               val lazyOffsetIndex: LazyOffsetIndex,
-                               val lazyTimeIndex: LazyTimeIndex,
+                               val lazyOffsetIndex: LazyIndex[OffsetIndex],
+                               val lazyTimeIndex: LazyIndex[TimeIndex],
                                val txnIndex: TransactionIndex,
                                val baseOffset: Long,
                                val indexIntervalBytes: Int,
@@ -95,8 +96,9 @@ class LogSegment private[log] (val log: FileRecords,
   /* the number of bytes since we last added an entry in the offset index */
   private var bytesSinceLastIndexEntry = 0
 
-  /* The timestamp we used for time based log rolling */
-  private var rollingBasedTimestamp: Option[Long] = None
+  // The timestamp we used for time based log rolling and for ensuring max compaction delay
+  // volatile for LogCleaner to see the update
+  @volatile private var rollingBasedTimestamp: Option[Long] = None
 
   /* The maximum timestamp we see so far */
   @volatile private var _maxTimestampSoFar: Option[Long] = None
@@ -233,7 +235,7 @@ class LogSegment private[log] (val log: FileRecords,
   }
 
   @nonthreadsafe
-  def updateTxnIndex(completedTxn: CompletedTxn, lastStableOffset: Long) {
+  def updateTxnIndex(completedTxn: CompletedTxn, lastStableOffset: Long): Unit = {
     if (completedTxn.isAborted) {
       trace(s"Writing aborted transaction $completedTxn to transaction index, last stable offset is $lastStableOffset")
       txnIndex.append(new AbortedTxn(completedTxn, lastStableOffset))
@@ -243,12 +245,13 @@ class LogSegment private[log] (val log: FileRecords,
   private def updateProducerState(producerStateManager: ProducerStateManager, batch: RecordBatch): Unit = {
     if (batch.hasProducerId) {
       val producerId = batch.producerId
-      val appendInfo = producerStateManager.prepareUpdate(producerId, isFromClient = false)
-      val maybeCompletedTxn = appendInfo.append(batch)
+      val appendInfo = producerStateManager.prepareUpdate(producerId, origin = AppendOrigin.Replication)
+      val maybeCompletedTxn = appendInfo.append(batch, firstOffsetMetadataOpt = None)
       producerStateManager.update(appendInfo)
       maybeCompletedTxn.foreach { completedTxn =>
-        val lastStableOffset = producerStateManager.completeTxn(completedTxn)
+        val lastStableOffset = producerStateManager.lastStableOffset(completedTxn)
         updateTxnIndex(completedTxn, lastStableOffset)
+        producerStateManager.completeTxn(completedTxn)
       }
     }
     producerStateManager.updateMapEndOffset(batch.lastOffset + 1)
@@ -277,7 +280,6 @@ class LogSegment private[log] (val log: FileRecords,
    * no more than maxSize bytes and will end before maxOffset if a maxOffset is specified.
    *
    * @param startOffset A lower bound on the first offset to include in the message set we read
-   * @param maxOffset An optional maximum offset for the message set we read
    * @param maxSize The maximum number of bytes to include in the message set we read
    * @param maxPosition The maximum position in the log segment that should be exposed for read
    * @param minOneMessage If this is true, the first message will be returned even if it exceeds `maxSize` (if one exists)
@@ -286,12 +288,13 @@ class LogSegment private[log] (val log: FileRecords,
    *         or null if the startOffset is larger than the largest offset in this log
    */
   @threadsafe
-  def read(startOffset: Long, maxOffset: Option[Long], maxSize: Int, maxPosition: Long = size,
+  def read(startOffset: Long,
+           maxSize: Int,
+           maxPosition: Long = size,
            minOneMessage: Boolean = false): FetchDataInfo = {
     if (maxSize < 0)
       throw new IllegalArgumentException(s"Invalid max size $maxSize for log read from segment $log")
 
-    val logSize = log.sizeInBytes // this may change, need to save a consistent copy
     val startOffsetAndSize = translateOffset(startOffset)
 
     // if the start position is already off the end of the log, return null
@@ -299,7 +302,7 @@ class LogSegment private[log] (val log: FileRecords,
       return null
 
     val startPosition = startOffsetAndSize.position
-    val offsetMetadata = new LogOffsetMetadata(startOffset, this.baseOffset, startPosition)
+    val offsetMetadata = LogOffsetMetadata(startOffset, this.baseOffset, startPosition)
 
     val adjustedMaxSize =
       if (minOneMessage) math.max(maxSize, startOffsetAndSize.size)
@@ -310,25 +313,7 @@ class LogSegment private[log] (val log: FileRecords,
       return FetchDataInfo(offsetMetadata, MemoryRecords.EMPTY)
 
     // calculate the length of the message set to read based on whether or not they gave us a maxOffset
-    val fetchSize: Int = maxOffset match {
-      case None =>
-        // no max offset, just read until the max position
-        min((maxPosition - startPosition).toInt, adjustedMaxSize)
-      case Some(offset) =>
-        // there is a max offset, translate it to a file position and use that to calculate the max read size;
-        // when the leader of a partition changes, it's possible for the new leader's high watermark to be less than the
-        // true high watermark in the previous leader for a short window. In this window, if a consumer fetches on an
-        // offset between new leader's high watermark and the log end offset, we want to return an empty response.
-        if (offset < startOffset)
-          return FetchDataInfo(offsetMetadata, MemoryRecords.EMPTY, firstEntryIncomplete = false)
-        val mapping = translateOffset(offset, startPosition)
-        val endPosition =
-          if (mapping == null)
-            logSize // the max offset is off the end of the log, use the end of the file
-          else
-            mapping.position
-        min(min(maxPosition, endPosition) - startPosition, adjustedMaxSize).toInt
-    }
+    val fetchSize: Int = min((maxPosition - startPosition).toInt, adjustedMaxSize)
 
     FetchDataInfo(offsetMetadata, log.slice(startPosition, fetchSize),
       firstEntryIncomplete = adjustedMaxSize < startOffsetAndSize.size)
@@ -383,9 +368,9 @@ class LogSegment private[log] (val log: FileRecords,
         }
       }
     } catch {
-      case e: CorruptRecordException =>
-        warn("Found invalid messages in log segment %s at byte offset %d: %s."
-          .format(log.file.getAbsolutePath, validBytes, e.getMessage))
+      case e@ (_: CorruptRecordException | _: InvalidRecordException) =>
+        warn("Found invalid messages in log segment %s at byte offset %d: %s. %s"
+          .format(log.file.getAbsolutePath, validBytes, e.getMessage, e.getCause))
     }
     val truncated = log.sizeInBytes - validBytes
     if (truncated > 0)
@@ -399,7 +384,7 @@ class LogSegment private[log] (val log: FileRecords,
     truncated
   }
 
-  private def loadLargestTimestamp() {
+  private def loadLargestTimestamp(): Unit = {
     // Get the last time index entry. If the time index is empty, it will return (-1, baseOffset)
     val lastTimeIndexEntry = timeIndex.lastEntry
     maxTimestampSoFar = lastTimeIndexEntry.timestamp
@@ -425,7 +410,11 @@ class LogSegment private[log] (val log: FileRecords,
   def collectAbortedTxns(fetchOffset: Long, upperBoundOffset: Long): TxnIndexSearchResult =
     txnIndex.collectAbortedTxns(fetchOffset, upperBoundOffset)
 
-  override def toString = "LogSegment(baseOffset=" + baseOffset + ", size=" + size + ")"
+  override def toString: String = "LogSegment(baseOffset=" + baseOffset +
+    ", size=" + size +
+    ", lastModifiedTime=" + lastModified +
+    ", largestTime=" + largestTimestamp +
+    ")"
 
   /**
    * Truncate off all index and log entries with offsets >= the given offset.
@@ -465,7 +454,7 @@ class LogSegment private[log] (val log: FileRecords,
    */
   @threadsafe
   def readNextOffset: Long = {
-    val fetchData = read(offsetIndex.lastOffset, None, log.sizeInBytes)
+    val fetchData = read(offsetIndex.lastOffset, log.sizeInBytes)
     if (fetchData == null)
       baseOffset
     else
@@ -478,7 +467,7 @@ class LogSegment private[log] (val log: FileRecords,
    * Flush this log segment to disk
    */
   @threadsafe
-  def flush() {
+  def flush(): Unit = {
     LogFlushStats.logFlushTimer.time {
       log.flush()
       offsetIndex.flush()
@@ -491,21 +480,21 @@ class LogSegment private[log] (val log: FileRecords,
    * Update the directory reference for the log and indices in this segment. This would typically be called after a
    * directory is renamed.
    */
-  def updateDir(dir: File): Unit = {
-    log.setFile(new File(dir, log.file.getName))
-    lazyOffsetIndex.file = new File(dir, lazyOffsetIndex.file.getName)
-    lazyTimeIndex.file = new File(dir, lazyTimeIndex.file.getName)
-    txnIndex.file = new File(dir, txnIndex.file.getName)
+  def updateParentDir(dir: File): Unit = {
+    log.updateParentDir(dir)
+    lazyOffsetIndex.updateParentDir(dir)
+    lazyTimeIndex.updateParentDir(dir)
+    txnIndex.updateParentDir(dir)
   }
 
   /**
-   * Change the suffix for the index and log file for this log segment
+   * Change the suffix for the index and log files for this log segment
    * IOException from this method should be handled by the caller
    */
-  def changeFileSuffixes(oldSuffix: String, newSuffix: String) {
+  def changeFileSuffixes(oldSuffix: String, newSuffix: String): Unit = {
     log.renameTo(new File(CoreUtils.replaceSuffix(log.file.getPath, oldSuffix, newSuffix)))
-    offsetIndex.renameTo(new File(CoreUtils.replaceSuffix(lazyOffsetIndex.file.getPath, oldSuffix, newSuffix)))
-    timeIndex.renameTo(new File(CoreUtils.replaceSuffix(lazyTimeIndex.file.getPath, oldSuffix, newSuffix)))
+    lazyOffsetIndex.renameTo(new File(CoreUtils.replaceSuffix(lazyOffsetIndex.file.getPath, oldSuffix, newSuffix)))
+    lazyTimeIndex.renameTo(new File(CoreUtils.replaceSuffix(lazyTimeIndex.file.getPath, oldSuffix, newSuffix)))
     txnIndex.renameTo(new File(CoreUtils.replaceSuffix(txnIndex.file.getPath, oldSuffix, newSuffix)))
   }
 
@@ -514,13 +503,25 @@ class LogSegment private[log] (val log: FileRecords,
    *
    * The time index entry appended will be used to decide when to delete the segment.
    */
-  def onBecomeInactiveSegment() {
+  def onBecomeInactiveSegment(): Unit = {
     timeIndex.maybeAppend(maxTimestampSoFar, offsetOfMaxTimestampSoFar, skipFullCheck = true)
     offsetIndex.trimToValidSize()
     timeIndex.trimToValidSize()
     log.trim()
   }
 
+  /**
+    * If not previously loaded,
+    * load the timestamp of the first message into memory.
+    */
+  private def loadFirstBatchTimestamp(): Unit = {
+    if (rollingBasedTimestamp.isEmpty) {
+      val iter = log.batches.iterator()
+      if (iter.hasNext)
+        rollingBasedTimestamp = Some(iter.next().maxTimestamp)
+    }
+  }
+
   /**
    * The time this segment has waited to be rolled.
    * If the first message batch has a timestamp we use its timestamp to determine when to roll a segment. A segment
@@ -532,17 +533,24 @@ class LogSegment private[log] (val log: FileRecords,
    */
   def timeWaitedForRoll(now: Long, messageTimestamp: Long) : Long = {
     // Load the timestamp of the first message into memory
-    if (rollingBasedTimestamp.isEmpty) {
-      val iter = log.batches.iterator()
-      if (iter.hasNext)
-        rollingBasedTimestamp = Some(iter.next().maxTimestamp)
-    }
+    loadFirstBatchTimestamp()
     rollingBasedTimestamp match {
       case Some(t) if t >= 0 => messageTimestamp - t
       case _ => now - created
     }
   }
 
+  /**
+    * @return the first batch timestamp if the timestamp is available. Otherwise return Long.MaxValue
+    */
+  def getFirstBatchTimestamp() : Long = {
+    loadFirstBatchTimestamp()
+    rollingBasedTimestamp match {
+      case Some(t) if t >= 0 => t
+      case _ => Long.MaxValue
+    }
+  }
+
   /**
    * Search the message offset based on timestamp and offset.
    *
@@ -576,10 +584,12 @@ class LogSegment private[log] (val log: FileRecords,
   /**
    * Close this log segment
    */
-  def close() {
-    CoreUtils.swallow(timeIndex.maybeAppend(maxTimestampSoFar, offsetOfMaxTimestampSoFar, skipFullCheck = true), this)
-    CoreUtils.swallow(offsetIndex.close(), this)
-    CoreUtils.swallow(timeIndex.close(), this)
+  def close(): Unit = {
+    if (_maxTimestampSoFar.nonEmpty || _offsetOfMaxTimestampSoFar.nonEmpty)
+      CoreUtils.swallow(timeIndex.maybeAppend(maxTimestampSoFar, offsetOfMaxTimestampSoFar,
+        skipFullCheck = true), this)
+    CoreUtils.swallow(lazyOffsetIndex.close(), this)
+    CoreUtils.swallow(lazyTimeIndex.close(), this)
     CoreUtils.swallow(log.close(), this)
     CoreUtils.swallow(txnIndex.close(), this)
   }
@@ -587,9 +597,9 @@ class LogSegment private[log] (val log: FileRecords,
   /**
     * Close file handlers used by the log segment but don't write to disk. This is used when the disk may have failed
     */
-  def closeHandlers() {
-    CoreUtils.swallow(offsetIndex.closeHandler(), this)
-    CoreUtils.swallow(timeIndex.closeHandler(), this)
+  def closeHandlers(): Unit = {
+    CoreUtils.swallow(lazyOffsetIndex.closeHandler(), this)
+    CoreUtils.swallow(lazyTimeIndex.closeHandler(), this)
     CoreUtils.swallow(log.closeHandlers(), this)
     CoreUtils.swallow(txnIndex.close(), this)
   }
@@ -597,7 +607,7 @@ class LogSegment private[log] (val log: FileRecords,
   /**
    * Delete this log segment from the filesystem.
    */
-  def deleteIfExists() {
+  def deleteIfExists(): Unit = {
     def delete(delete: () => Boolean, fileType: String, file: File, logIfMissing: Boolean): Unit = {
       try {
         if (delete())
@@ -612,8 +622,8 @@ class LogSegment private[log] (val log: FileRecords,
 
     CoreUtils.tryAll(Seq(
       () => delete(log.deleteIfExists _, "log", log.file, logIfMissing = true),
-      () => delete(offsetIndex.deleteIfExists _, "offset index", lazyOffsetIndex.file, logIfMissing = true),
-      () => delete(timeIndex.deleteIfExists _, "time index", lazyTimeIndex.file, logIfMissing = true),
+      () => delete(lazyOffsetIndex.deleteIfExists _, "offset index", lazyOffsetIndex.file, logIfMissing = true),
+      () => delete(lazyTimeIndex.deleteIfExists _, "time index", lazyTimeIndex.file, logIfMissing = true),
       () => delete(txnIndex.deleteIfExists _, "transaction index", txnIndex.file, logIfMissing = false)
     ))
   }
@@ -647,8 +657,8 @@ object LogSegment {
     val maxIndexSize = config.maxIndexSize
     new LogSegment(
       FileRecords.open(Log.logFile(dir, baseOffset, fileSuffix), fileAlreadyExists, initFileSize, preallocate),
-      new LazyOffsetIndex(Log.offsetIndexFile(dir, baseOffset, fileSuffix), baseOffset = baseOffset, maxIndexSize = maxIndexSize),
-      new LazyTimeIndex(Log.timeIndexFile(dir, baseOffset, fileSuffix), baseOffset = baseOffset, maxIndexSize = maxIndexSize),
+      LazyIndex.forOffset(Log.offsetIndexFile(dir, baseOffset, fileSuffix), baseOffset = baseOffset, maxIndexSize = maxIndexSize),
+      LazyIndex.forTime(Log.timeIndexFile(dir, baseOffset, fileSuffix), baseOffset = baseOffset, maxIndexSize = maxIndexSize),
       new TransactionIndex(baseOffset, Log.transactionIndexFile(dir, baseOffset, fileSuffix)),
       baseOffset,
       indexIntervalBytes = config.indexInterval,
diff --git a/core/src/main/scala/kafka/log/LogValidator.scala b/core/src/main/scala/kafka/log/LogValidator.scala
index f7c7d04688560..a80b73c99fc9a 100644
--- a/core/src/main/scala/kafka/log/LogValidator.scala
+++ b/core/src/main/scala/kafka/log/LogValidator.scala
@@ -19,17 +19,46 @@ package kafka.log
 import java.nio.ByteBuffer
 
 import kafka.api.{ApiVersion, KAFKA_2_1_IV0}
-import kafka.common.LongRef
+import kafka.common.{LongRef, RecordValidationException}
 import kafka.message.{CompressionCodec, NoCompressionCodec, ZStdCompressionCodec}
+import kafka.server.BrokerTopicStats
 import kafka.utils.Logging
-import org.apache.kafka.common.errors.{InvalidTimestampException, UnsupportedCompressionTypeException, UnsupportedForMessageFormatException}
-import org.apache.kafka.common.record.{AbstractRecords, CompressionType, InvalidRecordException, MemoryRecords, Record, RecordBatch, RecordConversionStats, TimestampType}
+import org.apache.kafka.common.errors.{CorruptRecordException, InvalidTimestampException, UnsupportedCompressionTypeException, UnsupportedForMessageFormatException}
+import org.apache.kafka.common.record.{AbstractRecords, BufferSupplier, CompressionType, MemoryRecords, Record, RecordBatch, RecordConversionStats, TimestampType}
+import org.apache.kafka.common.InvalidRecordException
+import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.requests.ProduceResponse.RecordError
 import org.apache.kafka.common.utils.Time
 
-import scala.collection.mutable
+import scala.collection.{Seq, mutable}
 import scala.collection.JavaConverters._
 
-private[kafka] object LogValidator extends Logging {
+/**
+ * The source of an append to the log. This is used when determining required validations.
+ */
+private[kafka] sealed trait AppendOrigin
+private[kafka] object AppendOrigin {
+
+  /**
+   * The log append came through replication from the leader. This typically implies minimal validation.
+   * Particularly, we do not decompress record batches in order to validate records individually.
+   */
+  case object Replication extends AppendOrigin
+
+  /**
+   * The log append came from either the group coordinator or the transaction coordinator. We validate
+   * producer epochs for normal log entries (specifically offset commits from the group coordinator) and
+   * we validate coordinate end transaction markers from the transaction coordinator.
+   */
+  case object Coordinator extends AppendOrigin
+
+  /**
+   * The log append came from the client, which implies full validation.
+   */
+  case object Client extends AppendOrigin
+}
+
+private[log] object LogValidator extends Logging {
 
   /**
    * Update the offsets for this message set and do further validation on messages including:
@@ -46,57 +75,100 @@ private[kafka] object LogValidator extends Logging {
    * 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.
    */
-  private[kafka] def validateMessagesAndAssignOffsets(records: MemoryRecords,
-                                                      offsetCounter: LongRef,
-                                                      time: Time,
-                                                      now: Long,
-                                                      sourceCodec: CompressionCodec,
-                                                      targetCodec: CompressionCodec,
-                                                      compactedTopic: Boolean,
-                                                      magic: Byte,
-                                                      timestampType: TimestampType,
-                                                      timestampDiffMaxMs: Long,
-                                                      partitionLeaderEpoch: Int,
-                                                      isFromClient: Boolean,
-                                                      interBrokerProtocolVersion: ApiVersion): ValidationAndOffsetAssignResult = {
+  private[log] def validateMessagesAndAssignOffsets(records: MemoryRecords,
+                                                    topicPartition: TopicPartition,
+                                                    offsetCounter: LongRef,
+                                                    time: Time,
+                                                    now: Long,
+                                                    sourceCodec: CompressionCodec,
+                                                    targetCodec: CompressionCodec,
+                                                    compactedTopic: Boolean,
+                                                    magic: Byte,
+                                                    timestampType: TimestampType,
+                                                    timestampDiffMaxMs: Long,
+                                                    partitionLeaderEpoch: Int,
+                                                    origin: AppendOrigin,
+                                                    interBrokerProtocolVersion: ApiVersion,
+                                                    brokerTopicStats: BrokerTopicStats,
+                                                    decompressionEnable: Boolean = true): ValidationAndOffsetAssignResult = {
     if (sourceCodec == NoCompressionCodec && targetCodec == NoCompressionCodec) {
       // check the magic value
       if (!records.hasMatchingMagic(magic))
-        convertAndAssignOffsetsNonCompressed(records, offsetCounter, compactedTopic, time, now, timestampType,
-          timestampDiffMaxMs, magic, partitionLeaderEpoch, isFromClient)
+        convertAndAssignOffsetsNonCompressed(records, topicPartition, offsetCounter, compactedTopic, time, now, timestampType,
+          timestampDiffMaxMs, magic, partitionLeaderEpoch, origin, brokerTopicStats)
       else
         // Do in-place validation, offset assignment and maybe set timestamp
-        assignOffsetsNonCompressed(records, offsetCounter, now, compactedTopic, timestampType, timestampDiffMaxMs,
-          partitionLeaderEpoch, isFromClient, magic)
+        assignOffsetsNonCompressed(records, topicPartition, offsetCounter, now, compactedTopic, timestampType, timestampDiffMaxMs,
+          partitionLeaderEpoch, origin, magic, brokerTopicStats)
     } else {
-      validateMessagesAndAssignOffsetsCompressed(records, offsetCounter, time, now, sourceCodec, targetCodec, compactedTopic,
-        magic, timestampType, timestampDiffMaxMs, partitionLeaderEpoch, isFromClient, interBrokerProtocolVersion)
+      validateMessagesAndAssignOffsetsCompressed(records, topicPartition, offsetCounter, time, now, sourceCodec, targetCodec, compactedTopic,
+        magic, timestampType, timestampDiffMaxMs, partitionLeaderEpoch, origin, interBrokerProtocolVersion, brokerTopicStats, decompressionEnable)
+    }
+  }
+
+  private def getFirstBatchAndMaybeValidateNoMoreBatches(records: MemoryRecords, sourceCodec: CompressionCodec): RecordBatch = {
+    val batchIterator = records.batches.iterator
+
+    if (!batchIterator.hasNext) {
+      throw new InvalidRecordException("Record batch has no batches at all")
+    }
+
+    val batch = batchIterator.next()
+
+    // if the format is v2 and beyond, or if the messages are compressed, we should check there's only one batch.
+    if (batch.magic() >= RecordBatch.MAGIC_VALUE_V2 || sourceCodec != NoCompressionCodec) {
+      if (batchIterator.hasNext) {
+        throw new InvalidRecordException("Compressed outer record has more than one batch")
+      }
     }
+
+    batch
   }
 
-  private def validateBatch(batch: RecordBatch, isFromClient: Boolean, toMagic: Byte): Unit = {
-    if (isFromClient) {
+  private def validateBatch(topicPartition: TopicPartition,
+                            firstBatch: RecordBatch,
+                            batch: RecordBatch,
+                            origin: AppendOrigin,
+                            toMagic: Byte,
+                            brokerTopicStats: BrokerTopicStats): Unit = {
+    // batch magic byte should have the same magic as the first batch
+    if (firstBatch.magic() != batch.magic()) {
+      brokerTopicStats.allTopicsStats.invalidMagicNumberRecordsPerSec.mark()
+      throw new InvalidRecordException(s"Batch magic ${batch.magic()} is not the same as the first batch'es magic byte ${firstBatch.magic()} in topic partition $topicPartition.")
+    }
+
+    if (origin == AppendOrigin.Client) {
       if (batch.magic >= RecordBatch.MAGIC_VALUE_V2) {
         val countFromOffsets = batch.lastOffset - batch.baseOffset + 1
-        if (countFromOffsets <= 0)
-          throw new InvalidRecordException(s"Batch has an invalid offset range: [${batch.baseOffset}, ${batch.lastOffset}]")
+        if (countFromOffsets <= 0) {
+          brokerTopicStats.allTopicsStats.invalidOffsetOrSequenceRecordsPerSec.mark()
+          throw new InvalidRecordException(s"Batch has an invalid offset range: [${batch.baseOffset}, ${batch.lastOffset}] in topic partition $topicPartition.")
+        }
 
         // v2 and above messages always have a non-null count
         val count = batch.countOrNull
-        if (count <= 0)
-          throw new InvalidRecordException(s"Invalid reported count for record batch: $count")
+        if (count <= 0) {
+          brokerTopicStats.allTopicsStats.invalidOffsetOrSequenceRecordsPerSec.mark()
+          throw new InvalidRecordException(s"Invalid reported count for record batch: $count in topic partition $topicPartition.")
+        }
 
-        if (countFromOffsets != batch.countOrNull)
+        if (countFromOffsets != batch.countOrNull) {
+          brokerTopicStats.allTopicsStats.invalidOffsetOrSequenceRecordsPerSec.mark()
           throw new InvalidRecordException(s"Inconsistent batch offset range [${batch.baseOffset}, ${batch.lastOffset}] " +
-            s"and count of records $count")
+            s"and count of records $count in topic partition $topicPartition.")
+        }
       }
 
-      if (batch.hasProducerId && batch.baseSequence < 0)
-        throw new InvalidRecordException(s"Invalid sequence number ${batch.baseSequence} in record batch " +
-          s"with producerId ${batch.producerId}")
+      if (batch.isControlBatch) {
+        brokerTopicStats.allTopicsStats.invalidOffsetOrSequenceRecordsPerSec.mark()
+        throw new InvalidRecordException(s"Clients are not allowed to write control records in topic partition $topicPartition.")
+      }
 
-      if (batch.isControlBatch)
-        throw new InvalidRecordException("Clients are not allowed to write control records")
+      if (batch.hasProducerId && batch.baseSequence < 0) {
+        brokerTopicStats.allTopicsStats.invalidOffsetOrSequenceRecordsPerSec.mark()
+        throw new InvalidRecordException(s"Invalid sequence number ${batch.baseSequence} in record batch " +
+          s"with producerId ${batch.producerId} in topic partition $topicPartition.")
+      }
     }
 
     if (batch.isTransactional && toMagic < RecordBatch.MAGIC_VALUE_V2)
@@ -106,23 +178,36 @@ private[kafka] object LogValidator extends Logging {
       throw new UnsupportedForMessageFormatException(s"Idempotent records cannot be used with magic version $toMagic")
   }
 
-  private def validateRecord(batch: RecordBatch, record: Record, now: Long, timestampType: TimestampType,
-                             timestampDiffMaxMs: Long, compactedTopic: Boolean): Unit = {
-    if (!record.hasMagic(batch.magic))
-      throw new InvalidRecordException(s"Log record magic does not match outer magic ${batch.magic}")
+  private def validateRecord(batch: RecordBatch, topicPartition: TopicPartition, record: Record, batchIndex: Int, now: Long,
+                             timestampType: TimestampType, timestampDiffMaxMs: Long, compactedTopic: Boolean,
+                             brokerTopicStats: BrokerTopicStats): Unit = {
+    if (!record.hasMagic(batch.magic)) {
+      brokerTopicStats.allTopicsStats.invalidMagicNumberRecordsPerSec.mark()
+      throw new RecordValidationException(
+        new InvalidRecordException(s"Log record $record's magic does not match outer magic ${batch.magic} in topic partition $topicPartition."),
+        List(new RecordError(batchIndex)))
+    }
 
     // verify the record-level CRC only if this is one of the deep entries of a compressed message
     // set for magic v0 and v1. For non-compressed messages, there is no inner record for magic v0 and v1,
     // so we depend on the batch-level CRC check in Log.analyzeAndValidateRecords(). For magic v2 and above,
     // there is no record-level CRC to check.
-    if (batch.magic <= RecordBatch.MAGIC_VALUE_V1 && batch.isCompressed)
-      record.ensureValid()
+    if (batch.magic <= RecordBatch.MAGIC_VALUE_V1 && batch.isCompressed) {
+      try {
+        record.ensureValid()
+      } catch {
+        case e: InvalidRecordException =>
+          brokerTopicStats.allTopicsStats.invalidMessageCrcRecordsPerSec.mark()
+          throw new CorruptRecordException(e.getMessage + s" in topic partition $topicPartition.")
+      }
+    }
 
-    validateKey(record, compactedTopic)
-    validateTimestamp(batch, record, now, timestampType, timestampDiffMaxMs)
+    validateKey(record, batchIndex, topicPartition, compactedTopic, brokerTopicStats)
+    validateTimestamp(batch, record, batchIndex, now, timestampType, timestampDiffMaxMs)
   }
 
   private def convertAndAssignOffsetsNonCompressed(records: MemoryRecords,
+                                                   topicPartition: TopicPartition,
                                                    offsetCounter: LongRef,
                                                    compactedTopic: Boolean,
                                                    time: Time,
@@ -131,7 +216,8 @@ private[kafka] object LogValidator extends Logging {
                                                    timestampDiffMaxMs: Long,
                                                    toMagicValue: Byte,
                                                    partitionLeaderEpoch: Int,
-                                                   isFromClient: Boolean): ValidationAndOffsetAssignResult = {
+                                                   origin: AppendOrigin,
+                                                   brokerTopicStats: BrokerTopicStats): ValidationAndOffsetAssignResult = {
     val startNanos = time.nanoseconds
     val sizeInBytesAfterConversion = AbstractRecords.estimateSizeInBytes(toMagicValue, offsetCounter.value,
       CompressionType.NONE, records.records)
@@ -145,11 +231,13 @@ private[kafka] object LogValidator extends Logging {
     val builder = MemoryRecords.builder(newBuffer, toMagicValue, CompressionType.NONE, timestampType,
       offsetCounter.value, now, producerId, producerEpoch, sequence, isTransactional, partitionLeaderEpoch)
 
+    val firstBatch = getFirstBatchAndMaybeValidateNoMoreBatches(records, NoCompressionCodec)
+
     for (batch <- records.batches.asScala) {
-      validateBatch(batch, isFromClient, toMagicValue)
+      validateBatch(topicPartition, firstBatch, batch, origin, toMagicValue, brokerTopicStats)
 
-      for (record <- batch.asScala) {
-        validateRecord(batch, record, now, timestampType, timestampDiffMaxMs, compactedTopic)
+      for ((record, batchIndex) <- batch.asScala.view.zipWithIndex) {
+        validateRecord(batch, topicPartition, record, batchIndex, now, timestampType, timestampDiffMaxMs, compactedTopic, brokerTopicStats)
         builder.appendWithOffset(offsetCounter.getAndIncrement(), record)
       }
     }
@@ -164,30 +252,36 @@ private[kafka] object LogValidator extends Logging {
       maxTimestamp = info.maxTimestamp,
       shallowOffsetOfMaxTimestamp = info.shallowOffsetOfMaxTimestamp,
       messageSizeMaybeChanged = true,
-      recordConversionStats = recordConversionStats)
+      recordConversionStats = recordConversionStats,
+      recompressApplied = false)
   }
 
   private def assignOffsetsNonCompressed(records: MemoryRecords,
+                                         topicPartition: TopicPartition,
                                          offsetCounter: LongRef,
                                          now: Long,
                                          compactedTopic: Boolean,
                                          timestampType: TimestampType,
                                          timestampDiffMaxMs: Long,
                                          partitionLeaderEpoch: Int,
-                                         isFromClient: Boolean,
-                                         magic: Byte): ValidationAndOffsetAssignResult = {
+                                         origin: AppendOrigin,
+                                         magic: Byte,
+                                         brokerTopicStats: BrokerTopicStats): ValidationAndOffsetAssignResult = {
     var maxTimestamp = RecordBatch.NO_TIMESTAMP
     var offsetOfMaxTimestamp = -1L
     val initialOffset = offsetCounter.value
 
+    val firstBatch = getFirstBatchAndMaybeValidateNoMoreBatches(records, NoCompressionCodec)
+
     for (batch <- records.batches.asScala) {
-      validateBatch(batch, isFromClient, magic)
+      validateBatch(topicPartition, firstBatch, batch, origin, magic, brokerTopicStats)
 
       var maxBatchTimestamp = RecordBatch.NO_TIMESTAMP
       var offsetOfMaxBatchTimestamp = -1L
 
-      for (record <- batch.asScala) {
-        validateRecord(batch, record, now, timestampType, timestampDiffMaxMs, compactedTopic)
+      for ((record, batchIndex) <- batch.asScala.view.zipWithIndex) {
+        validateRecord(batch, topicPartition, record, batchIndex, now, timestampType, timestampDiffMaxMs, compactedTopic, brokerTopicStats)
+
         val offset = offsetCounter.getAndIncrement()
         if (batch.magic > RecordBatch.MAGIC_VALUE_V0 && record.timestamp > maxBatchTimestamp) {
           maxBatchTimestamp = record.timestamp
@@ -226,17 +320,18 @@ private[kafka] object LogValidator extends Logging {
       maxTimestamp = maxTimestamp,
       shallowOffsetOfMaxTimestamp = offsetOfMaxTimestamp,
       messageSizeMaybeChanged = false,
-      recordConversionStats = RecordConversionStats.EMPTY)
+      recordConversionStats = RecordConversionStats.EMPTY,
+      recompressApplied = false)
   }
 
   /**
    * We cannot do in place assignment in one of the following situations:
    * 1. Source and target compression codec are different
-   * 2. When magic value to use is 0 because offsets need to be overwritten
-   * 3. When magic value to use is above 0, but some fields of inner messages need to be overwritten.
-   * 4. Message format conversion is needed.
+   * 2. When the target magic is not equal to batches' magic, meaning format conversion is needed.
+   * 3. When the target magic is equal to V0, meaning absolute offsets need to be re-assigned.
    */
   def validateMessagesAndAssignOffsetsCompressed(records: MemoryRecords,
+                                                 topicPartition: TopicPartition,
                                                  offsetCounter: LongRef,
                                                  time: Time,
                                                  now: Long,
@@ -247,85 +342,134 @@ private[kafka] object LogValidator extends Logging {
                                                  timestampType: TimestampType,
                                                  timestampDiffMaxMs: Long,
                                                  partitionLeaderEpoch: Int,
-                                                 isFromClient: Boolean,
-                                                 interBrokerProtocolVersion: ApiVersion): ValidationAndOffsetAssignResult = {
-      // No in place assignment situation 1 and 2
-      var inPlaceAssignment = sourceCodec == targetCodec && toMagic > RecordBatch.MAGIC_VALUE_V0
+                                                 origin: AppendOrigin,
+                                                 interBrokerProtocolVersion: ApiVersion,
+                                                 brokerTopicStats: BrokerTopicStats,
+                                                 decompressionEnable: Boolean): ValidationAndOffsetAssignResult = {
 
-      var maxTimestamp = RecordBatch.NO_TIMESTAMP
-      val expectedInnerOffset = new LongRef(0)
-      val validatedRecords = new mutable.ArrayBuffer[Record]
+    if (targetCodec == ZStdCompressionCodec && interBrokerProtocolVersion < KAFKA_2_1_IV0)
+      throw new UnsupportedCompressionTypeException("Produce requests to inter.broker.protocol.version < 2.1 broker " +
+        "are not allowed to use ZStandard compression")
 
-      var uncompressedSizeInBytes = 0
+    // No in place assignment situation 1
+    var inPlaceAssignment = sourceCodec == targetCodec
 
-      for (batch <- records.batches.asScala) {
-        validateBatch(batch, isFromClient, toMagic)
-        uncompressedSizeInBytes += AbstractRecords.recordBatchHeaderSizeInBytes(toMagic, batch.compressionType())
-
-        // Do not compress control records unless they are written compressed
-        if (sourceCodec == NoCompressionCodec && batch.isControlBatch)
-          inPlaceAssignment = true
+    var maxTimestamp = RecordBatch.NO_TIMESTAMP
+    val expectedInnerOffset = new LongRef(0)
+    val validatedRecords = new mutable.ArrayBuffer[Record]
+
+    var uncompressedSizeInBytes = 0
+
+    // Assume there's only one batch with compressed memory records; otherwise, return InvalidRecordException
+    // One exception though is that with format smaller than v2, if sourceCodec is noCompression, then each batch is actually
+    // a single record so we'd need to special handle it by creating a single wrapper batch that includes all the records
+    val firstBatch = getFirstBatchAndMaybeValidateNoMoreBatches(records, sourceCodec)
+    // Check is made in {@link LogValidator#validateBatch} that the relative offset(s) of the record(s) in the record
+    // version V2 batch is monotonically increasing by 1, which is one of the requirement to avoid decompression.
+    val decompressBatch = decompressionEnable ||
+      !records.hasMatchingMagic(toMagic) ||
+      (firstBatch.magic < RecordBatch.MAGIC_VALUE_V2)
+
+    // No in place assignment situation 2 and 3: we only need to check for the first batch because:
+    //  1. For most cases (compressed records, v2, for example), there's only one batch anyways.
+    //  2. For cases that there may be multiple batches, all batches' magic should be the same.
+    if (firstBatch.magic != toMagic || toMagic == RecordBatch.MAGIC_VALUE_V0)
+      inPlaceAssignment = false
+
+    // Do not compress control records unless they are written compressed
+    if (sourceCodec == NoCompressionCodec && firstBatch.isControlBatch)
+      inPlaceAssignment = true
+
+    val batches = records.batches.asScala
+    for (batch <- batches) {
+      validateBatch(topicPartition, firstBatch, batch, origin, toMagic, brokerTopicStats)
+      uncompressedSizeInBytes += AbstractRecords.recordBatchHeaderSizeInBytes(toMagic, batch.compressionType())
+
+      // The following inner block are not indented to make the hotfix cherry-picking easier
+      if (decompressBatch || !inPlaceAssignment) {
+      // if we are on version 2 and beyond, and we know we are going for in place assignment,
+      // then we can optimize the iterator to skip key / value / headers since they would not be used at all
+      val recordsIterator = if (inPlaceAssignment && firstBatch.magic >= RecordBatch.MAGIC_VALUE_V2)
+        batch.skipKeyValueIterator(BufferSupplier.NO_CACHING)
+      else
+        batch.streamingIterator(BufferSupplier.NO_CACHING)
 
-        for (record <- batch.asScala) {
+      try {
+        for ((record, batchIndex) <- batch.asScala.view.zipWithIndex) {
           if (sourceCodec != NoCompressionCodec && record.isCompressed)
-            throw new InvalidRecordException("Compressed outer record should not have an inner record with a " +
-              s"compression attribute set: $record")
-          if (targetCodec == ZStdCompressionCodec && interBrokerProtocolVersion < KAFKA_2_1_IV0)
-            throw new UnsupportedCompressionTypeException("Produce requests to inter.broker.protocol.version < 2.1 broker " + "are not allowed to use ZStandard compression")
-          validateRecord(batch, record, now, timestampType, timestampDiffMaxMs, compactedTopic)
+            throw new RecordValidationException(
+              new InvalidRecordException(s"Compressed outer record should not have an inner record with a compression attribute set: $record"),
+              List(new RecordError(batchIndex)))
+
+          validateRecord(batch, topicPartition, record, batchIndex, now, timestampType, timestampDiffMaxMs, compactedTopic, brokerTopicStats)
 
           uncompressedSizeInBytes += record.sizeInBytes()
           if (batch.magic > RecordBatch.MAGIC_VALUE_V0 && toMagic > RecordBatch.MAGIC_VALUE_V0) {
-            // Check if we need to overwrite offset
-            // No in place assignment situation 3
-            if (record.offset != expectedInnerOffset.getAndIncrement())
+            // Some older clients do not implement the V1 internal offsets correctly.
+            // Historically the broker handled this by rewriting the batches rather
+            // than rejecting the request. We must continue this handling here to avoid
+            // breaking these clients.
+            val expectedOffset = expectedInnerOffset.getAndIncrement()
+            if (record.offset != expectedOffset)
               inPlaceAssignment = false
+
             if (record.timestamp > maxTimestamp)
               maxTimestamp = record.timestamp
           }
 
-          // No in place assignment situation 4
-          if (!record.hasMagic(toMagic))
-            inPlaceAssignment = false
-
           validatedRecords += record
         }
+      } finally {
+        recordsIterator.close()
+      }
       }
+    }
 
-      if (!inPlaceAssignment) {
-        val (producerId, producerEpoch, sequence, isTransactional) = {
-          // note that we only reassign offsets for requests coming straight from a producer. For records with magic V2,
-          // there should be exactly one RecordBatch per request, so the following is all we need to do. For Records
-          // with older magic versions, there will never be a producer id, etc.
-          val first = records.batches.asScala.head
-          (first.producerId, first.producerEpoch, first.baseSequence, first.isTransactional)
-        }
-        buildRecordsAndAssignOffsets(toMagic, offsetCounter, time, timestampType, CompressionType.forId(targetCodec.codec), now,
-          validatedRecords, producerId, producerEpoch, sequence, isTransactional, partitionLeaderEpoch, isFromClient,
-          uncompressedSizeInBytes)
-      } else {
-        // we can update the batch only and write the compressed payload as is
-        val batch = records.batches.iterator.next()
-        val lastOffset = offsetCounter.addAndGet(validatedRecords.size) - 1
+    if (!inPlaceAssignment) {
+      val (producerId, producerEpoch, sequence, isTransactional) = {
+        // note that we only reassign offsets for requests coming straight from a producer. For records with magic V2,
+        // there should be exactly one RecordBatch per request, so the following is all we need to do. For Records
+        // with older magic versions, there will never be a producer id, etc.
+        val first = records.batches.asScala.head
+        (first.producerId, first.producerEpoch, first.baseSequence, first.isTransactional)
+      }
+      buildRecordsAndAssignOffsets(toMagic, offsetCounter, time, timestampType, CompressionType.forId(targetCodec.codec),
+        now, validatedRecords, producerId, producerEpoch, sequence, isTransactional, partitionLeaderEpoch,
+        uncompressedSizeInBytes)
+    } else {
+      // we can update the batch only and write the compressed payload as is;
+      // again we assume only one record batch within the compressed set
+      val batch = records.batches.iterator.next()
+      val lastOffset = if (decompressBatch)
+        offsetCounter.addAndGet(validatedRecords.size) - 1
+      else {
+        // batch.countOrNull() will never be null as the following line is execteud
+        // for record format version V2.
+        offsetCounter.addAndGet(batch.countOrNull().longValue()) - 1
+      }
 
-        batch.setLastOffset(lastOffset)
+      batch.setLastOffset(lastOffset)
 
-        if (timestampType == TimestampType.LOG_APPEND_TIME)
-          maxTimestamp = now
+      if (timestampType == TimestampType.LOG_APPEND_TIME)
+        maxTimestamp = now
+      else if (!decompressBatch)
+        maxTimestamp = batch.maxTimestamp
 
-        if (toMagic >= RecordBatch.MAGIC_VALUE_V1)
-          batch.setMaxTimestamp(timestampType, maxTimestamp)
 
-        if (toMagic >= RecordBatch.MAGIC_VALUE_V2)
-          batch.setPartitionLeaderEpoch(partitionLeaderEpoch)
+      if (toMagic >= RecordBatch.MAGIC_VALUE_V1)
+        batch.setMaxTimestamp(timestampType, maxTimestamp)
 
-        val recordConversionStats = new RecordConversionStats(uncompressedSizeInBytes, 0, 0)
-        ValidationAndOffsetAssignResult(validatedRecords = records,
-          maxTimestamp = maxTimestamp,
-          shallowOffsetOfMaxTimestamp = lastOffset,
-          messageSizeMaybeChanged = false,
-          recordConversionStats = recordConversionStats)
-      }
+      if (toMagic >= RecordBatch.MAGIC_VALUE_V2)
+        batch.setPartitionLeaderEpoch(partitionLeaderEpoch)
+
+      val recordConversionStats = new RecordConversionStats(uncompressedSizeInBytes, 0, 0)
+      ValidationAndOffsetAssignResult(validatedRecords = records,
+        maxTimestamp = maxTimestamp,
+        shallowOffsetOfMaxTimestamp = lastOffset,
+        messageSizeMaybeChanged = false,
+        recordConversionStats = recordConversionStats,
+        recompressApplied = false)
+    }
   }
 
   private def buildRecordsAndAssignOffsets(magic: Byte,
@@ -340,8 +484,7 @@ private[kafka] object LogValidator extends Logging {
                                            baseSequence: Int,
                                            isTransactional: Boolean,
                                            partitionLeaderEpoch: Int,
-                                           isFromClient: Boolean,
-                                           uncompresssedSizeInBytes: Int): ValidationAndOffsetAssignResult = {
+                                           uncompressedSizeInBytes: Int): ValidationAndOffsetAssignResult = {
     val startNanos = time.nanoseconds
     val estimatedSize = AbstractRecords.estimateSizeInBytes(magic, offsetCounter.value, compressionType,
       validatedRecords.asJava)
@@ -362,7 +505,7 @@ private[kafka] object LogValidator extends Logging {
     // message format V0 or if the inner offsets are not consecutive. This is OK since the impact is the same: we have
     // to rebuild the records (including recompression if enabled).
     val conversionCount = builder.numRecords
-    val recordConversionStats = new RecordConversionStats(uncompresssedSizeInBytes + builder.uncompressedBytesWritten,
+    val recordConversionStats = new RecordConversionStats(uncompressedSizeInBytes + builder.uncompressedBytesWritten,
       conversionCount, time.nanoseconds - startNanos)
 
     ValidationAndOffsetAssignResult(
@@ -370,12 +513,17 @@ private[kafka] object LogValidator extends Logging {
       maxTimestamp = info.maxTimestamp,
       shallowOffsetOfMaxTimestamp = info.shallowOffsetOfMaxTimestamp,
       messageSizeMaybeChanged = true,
-      recordConversionStats = recordConversionStats)
+      recordConversionStats = recordConversionStats,
+      recompressApplied = true)
   }
 
-  private def validateKey(record: Record, compactedTopic: Boolean) {
-    if (compactedTopic && !record.hasKey)
-      throw new InvalidRecordException("Compacted topic cannot accept message without key.")
+  private def validateKey(record: Record, batchIndex: Int, topicPartition: TopicPartition, compactedTopic: Boolean, brokerTopicStats: BrokerTopicStats) {
+    if (compactedTopic && !record.hasKey) {
+      brokerTopicStats.allTopicsStats.noKeyCompactedTopicRecordsPerSec.mark()
+      throw new RecordValidationException(
+        new InvalidRecordException(s"Compacted topic cannot accept message without key in topic partition $topicPartition."),
+        List(new RecordError(batchIndex)))
+    }
   }
 
   /**
@@ -384,23 +532,29 @@ private[kafka] object LogValidator extends Logging {
    */
   private def validateTimestamp(batch: RecordBatch,
                                 record: Record,
+                                batchIndex: Int,
                                 now: Long,
                                 timestampType: TimestampType,
-                                timestampDiffMaxMs: Long) {
+                                timestampDiffMaxMs: Long): Unit = {
     if (timestampType == TimestampType.CREATE_TIME
       && record.timestamp != RecordBatch.NO_TIMESTAMP
       && math.abs(record.timestamp - now) > timestampDiffMaxMs)
-      throw new InvalidTimestampException(s"Timestamp ${record.timestamp} of message with offset ${record.offset} is " +
-        s"out of range. The timestamp should be within [${now - timestampDiffMaxMs}, ${now + timestampDiffMaxMs}]")
+      throw new RecordValidationException(
+        new InvalidTimestampException(s"Timestamp ${record.timestamp} of message with offset ${record.offset} is " +
+          s"out of range. The timestamp should be within [${now - timestampDiffMaxMs}, ${now + timestampDiffMaxMs}]"),
+        List(new RecordError(batchIndex)))
     if (batch.timestampType == TimestampType.LOG_APPEND_TIME)
-      throw new InvalidTimestampException(s"Invalid timestamp type in message $record. Producer should not set " +
-        s"timestamp type to LogAppendTime.")
+      throw new RecordValidationException(
+        new InvalidTimestampException(s"Invalid timestamp type in message $record. Producer should not set " +
+          s"timestamp type to LogAppendTime."),
+        List(new RecordError(batchIndex)))
   }
 
   case class ValidationAndOffsetAssignResult(validatedRecords: MemoryRecords,
                                              maxTimestamp: Long,
                                              shallowOffsetOfMaxTimestamp: Long,
                                              messageSizeMaybeChanged: Boolean,
-                                             recordConversionStats: RecordConversionStats)
+                                             recordConversionStats: RecordConversionStats,
+                                             recompressApplied: Boolean)
 
 }
diff --git a/core/src/main/scala/kafka/log/OffsetIndex.scala b/core/src/main/scala/kafka/log/OffsetIndex.scala
index 876895cd5cce4..78ac79572fa7b 100755
--- a/core/src/main/scala/kafka/log/OffsetIndex.scala
+++ b/core/src/main/scala/kafka/log/OffsetIndex.scala
@@ -51,7 +51,7 @@ import org.apache.kafka.common.errors.InvalidOffsetException
  */
 // Avoid shadowing mutable `file` in AbstractIndex
 class OffsetIndex(_file: File, baseOffset: Long, maxIndexSize: Int = -1, writable: Boolean = true)
-    extends AbstractIndex[Long, Int](_file, baseOffset, maxIndexSize, writable) {
+    extends AbstractIndex(_file, baseOffset, maxIndexSize, writable) {
   import OffsetIndex._
 
   override def entrySize = 8
@@ -59,8 +59,9 @@ class OffsetIndex(_file: File, baseOffset: Long, maxIndexSize: Int = -1, writabl
   /* the last offset in the index */
   private[this] var _lastOffset = lastEntry.offset
 
-  debug("Loaded index file %s with maxEntries = %d, maxIndexSize = %d, entries = %d, lastOffset = %d, file position = %d"
-    .format(file.getAbsolutePath, maxEntries, maxIndexSize, _entries, _lastOffset, mmap.position()))
+  if (isDebugEnabled)
+    debug(s"Loaded index file ${file.getAbsolutePath} with maxEntries = $maxEntries, " +
+      s"maxIndexSize = $maxIndexSize, entries = ${_entries}, lastOffset = ${_lastOffset}, file position = ${mmap.position()}")
 
   /**
    * The last entry in the index
@@ -128,35 +129,37 @@ class OffsetIndex(_file: File, baseOffset: Long, maxIndexSize: Int = -1, writabl
   def entry(n: Int): OffsetPosition = {
     maybeLock(lock) {
       if (n >= _entries)
-        throw new IllegalArgumentException(s"Attempt to fetch the ${n}th entry from an index of size ${_entries}.")
+        throw new IllegalArgumentException(s"Attempt to fetch the ${n}th entry from index ${file.getAbsolutePath}, " +
+          s"which has size ${_entries}.")
       parseEntry(mmap, n)
     }
   }
 
   /**
    * Append an entry for the given offset/location pair to the index. This entry must have a larger offset than all subsequent entries.
-   * @throws IndexOffsetOverflowException if the offset causes index offset to overflow
+   * @throws kafka.common.IndexOffsetOverflowException#IndexOffsetOverflowException if the offset causes index offset to overflow
    */
-  def append(offset: Long, position: Int) {
+  def append(offset: Long, position: Int): Unit = {
     inLock(lock) {
       require(!isFull, "Attempt to append to a full index (size = " + _entries + ").")
       if (_entries == 0 || offset > _lastOffset) {
-        debug("Adding index entry %d => %d to %s.".format(offset, position, file.getName))
+        if (isTraceEnabled)
+          trace(s"Adding index entry $offset => $position to ${file.getAbsolutePath}")
         mmap.putInt(relativeOffset(offset))
         mmap.putInt(position)
         _entries += 1
         _lastOffset = offset
         require(_entries * entrySize == mmap.position(), entries + " entries but file position in index is " + mmap.position() + ".")
       } else {
-        throw new InvalidOffsetException("Attempt to append an offset (%d) to position %d no larger than the last offset appended (%d) to %s."
-          .format(offset, entries, _lastOffset, file.getAbsolutePath))
+        throw new InvalidOffsetException(s"Attempt to append an offset ($offset) to position $entries no larger than" +
+          s" the last offset appended (${_lastOffset}) to ${file.getAbsolutePath}.")
       }
     }
   }
 
   override def truncate() = truncateToEntries(0)
 
-  override def truncateTo(offset: Long) {
+  override def truncateTo(offset: Long): Unit = {
     inLock(lock) {
       val idx = mmap.duplicate
       val slot = largestLowerBoundSlotFor(idx, offset, IndexSearchType.KEY)
@@ -180,15 +183,18 @@ class OffsetIndex(_file: File, baseOffset: Long, maxIndexSize: Int = -1, writabl
   /**
    * Truncates index to a known number of entries.
    */
-  private def truncateToEntries(entries: Int) {
+  private def truncateToEntries(entries: Int): Unit = {
     inLock(lock) {
       _entries = entries
       mmap.position(_entries * entrySize)
       _lastOffset = lastEntry.offset
+      if (isDebugEnabled)
+        debug(s"Truncated index ${file.getAbsolutePath} to $entries entries;" +
+              s" position is now ${mmap.position()} and last offset is now ${_lastOffset}")
     }
   }
 
-  override def sanityCheck() {
+  override def sanityCheck(): Unit = {
     if (_entries != 0 && _lastOffset < baseOffset)
       throw new CorruptIndexException(s"Corrupt index found, index file (${file.getAbsolutePath}) has non-zero size " +
         s"but the last offset is ${_lastOffset} which is less than the base offset $baseOffset.")
@@ -202,37 +208,3 @@ class OffsetIndex(_file: File, baseOffset: Long, maxIndexSize: Int = -1, writabl
 object OffsetIndex extends Logging {
   override val loggerName: String = classOf[OffsetIndex].getName
 }
-
-
-
-/**
-  * A thin wrapper on top of the raw OffsetIndex object to avoid initialization on construction. This defers the OffsetIndex
-  * initialization to the time it gets accessed so the cost of the heavy memory mapped operation gets amortized over time.
-  *
-  * Combining with skipping sanity check for safely flushed segments, the startup time of a broker can be reduced, especially
-  * for the the broker with a lot of log segments
-  *
-  */
-class LazyOffsetIndex(@volatile private var _file: File, baseOffset: Long, maxIndexSize: Int = -1, writable: Boolean = true) {
-  @volatile private var offsetIndex: Option[OffsetIndex] = None
-
-  def file: File = {
-    if (offsetIndex.isDefined)
-      offsetIndex.get.file
-    else
-      _file
-  }
-
-  def file_=(f: File) {
-    if (offsetIndex.isDefined)
-      offsetIndex.get.file = f
-    else
-      _file = f
-  }
-
-  def get: OffsetIndex = {
-    if (offsetIndex.isEmpty)
-      offsetIndex = Some(new OffsetIndex(_file, baseOffset, maxIndexSize, writable))
-    offsetIndex.get
-  }
-}
\ No newline at end of file
diff --git a/core/src/main/scala/kafka/log/OffsetMap.scala b/core/src/main/scala/kafka/log/OffsetMap.scala
index 219bed3c6d22a..22b5305203e90 100755
--- a/core/src/main/scala/kafka/log/OffsetMap.scala
+++ b/core/src/main/scala/kafka/log/OffsetMap.scala
@@ -25,10 +25,10 @@ import org.apache.kafka.common.utils.Utils
 
 trait OffsetMap {
   def slots: Int
-  def put(key: ByteBuffer, offset: Long)
+  def put(key: ByteBuffer, offset: Long): Unit
   def get(key: ByteBuffer): Long
-  def updateLatestOffset(offset: Long)
-  def clear()
+  def updateLatestOffset(offset: Long): Unit
+  def clear(): Unit
   def size: Int
   def utilization: Double = size.toDouble / slots
   def latestOffset: Long
@@ -81,7 +81,7 @@ class SkimpyOffsetMap(val memory: Int, val hashAlgorithm: String = "MD5") extend
    * @param key The key
    * @param offset The offset
    */
-  override def put(key: ByteBuffer, offset: Long) {
+  override def put(key: ByteBuffer, offset: Long): Unit = {
     require(entries < slots, "Attempt to add a new entry to a full offset map.")
     lookups += 1
     hashInto(key, hash1)
@@ -144,7 +144,7 @@ class SkimpyOffsetMap(val memory: Int, val hashAlgorithm: String = "MD5") extend
   /**
    * Change the salt used for key hashing making all existing keys unfindable.
    */
-  override def clear() {
+  override def clear(): Unit = {
     this.entries = 0
     this.lookups = 0L
     this.probes = 0L
@@ -191,7 +191,7 @@ class SkimpyOffsetMap(val memory: Int, val hashAlgorithm: String = "MD5") extend
    * @param key The key to hash
    * @param buffer The buffer to store the hash into
    */
-  private def hashInto(key: ByteBuffer, buffer: Array[Byte]) {
+  private def hashInto(key: ByteBuffer, buffer: Array[Byte]): Unit = {
     key.mark()
     digest.update(key)
     key.reset()
diff --git a/core/src/main/scala/kafka/log/ProducerStateManager.scala b/core/src/main/scala/kafka/log/ProducerStateManager.scala
index a3b03d063f256..f195906f6fcc9 100644
--- a/core/src/main/scala/kafka/log/ProducerStateManager.scala
+++ b/core/src/main/scala/kafka/log/ProducerStateManager.scala
@@ -26,38 +26,22 @@ import kafka.server.LogOffsetMetadata
 import kafka.utils.{Logging, nonthreadsafe, threadsafe}
 import org.apache.kafka.common.{KafkaException, TopicPartition}
 import org.apache.kafka.common.errors._
-import org.apache.kafka.common.internals.Topic
 import org.apache.kafka.common.protocol.types._
 import org.apache.kafka.common.record.{ControlRecordType, DefaultRecordBatch, EndTransactionMarker, RecordBatch}
 import org.apache.kafka.common.utils.{ByteUtils, Crc32C}
 
+import scala.collection.JavaConverters._
 import scala.collection.mutable.ListBuffer
 import scala.collection.{immutable, mutable}
 
 class CorruptSnapshotException(msg: String) extends KafkaException(msg)
 
+/**
+ * The last written record for a given producer. The last data offset may be undefined
+ * if the only log entry for a producer is a transaction marker.
+ */
+case class LastRecord(lastDataOffset: Option[Long], producerEpoch: Short)
 
-// ValidationType and its subtypes define the extent of the validation to perform on a given ProducerAppendInfo instance
-private[log] sealed trait ValidationType
-private[log] object ValidationType {
-
-  /**
-    * This indicates no validation should be performed on the incoming append. This is the case for all appends on
-    * a replica, as well as appends when the producer state is being built from the log.
-    */
-  case object None extends ValidationType
-
-  /**
-    * We only validate the epoch (and not the sequence numbers) for offset commit requests coming from the transactional
-    * producer. These appends will not have sequence numbers, so we can't validate them.
-    */
-  case object EpochOnly extends ValidationType
-
-  /**
-    * Perform the full validation. This should be used fo regular produce requests coming to the leader.
-    */
-  case object Full extends ValidationType
-}
 
 private[log] case class TxnMetadata(producerId: Long, var firstOffset: LogOffsetMetadata, var lastOffset: Option[Long] = None) {
   def this(producerId: Long, firstOffset: Long) = this(producerId, LogOffsetMetadata(firstOffset))
@@ -72,12 +56,18 @@ private[log] case class TxnMetadata(producerId: Long, var firstOffset: LogOffset
 
 private[log] object ProducerStateEntry {
   private[log] val NumBatchesToRetain = 5
-  def empty(producerId: Long) = new ProducerStateEntry(producerId, mutable.Queue[BatchMetadata](), RecordBatch.NO_PRODUCER_EPOCH, -1, None)
+
+  def empty(producerId: Long) = new ProducerStateEntry(producerId,
+    batchMetadata = mutable.Queue[BatchMetadata](),
+    producerEpoch = RecordBatch.NO_PRODUCER_EPOCH,
+    coordinatorEpoch = -1,
+    lastTimestamp = RecordBatch.NO_TIMESTAMP,
+    currentTxnFirstOffset = None)
 }
 
 private[log] case class BatchMetadata(lastSeq: Int, lastOffset: Long, offsetDelta: Int, timestamp: Long) {
-  def firstSeq =  DefaultRecordBatch.decrementSequence(lastSeq, offsetDelta)
-  def firstOffset = lastOffset - offsetDelta
+  def firstSeq: Int =  DefaultRecordBatch.decrementSequence(lastSeq, offsetDelta)
+  def firstOffset: Long = lastOffset - offsetDelta
 
   override def toString: String = {
     "BatchMetadata(" +
@@ -96,28 +86,28 @@ private[log] class ProducerStateEntry(val producerId: Long,
                                       val batchMetadata: mutable.Queue[BatchMetadata],
                                       var producerEpoch: Short,
                                       var coordinatorEpoch: Int,
+                                      var lastTimestamp: Long,
                                       var currentTxnFirstOffset: Option[Long]) {
 
   def firstSeq: Int = if (isEmpty) RecordBatch.NO_SEQUENCE else batchMetadata.front.firstSeq
 
-  def firstOffset: Long = if (isEmpty) -1L else batchMetadata.front.firstOffset
+  def firstDataOffset: Long = if (isEmpty) -1L else batchMetadata.front.firstOffset
 
   def lastSeq: Int = if (isEmpty) RecordBatch.NO_SEQUENCE else batchMetadata.last.lastSeq
 
   def lastDataOffset: Long = if (isEmpty) -1L else batchMetadata.last.lastOffset
 
-  def lastTimestamp = if (isEmpty) RecordBatch.NO_TIMESTAMP else batchMetadata.last.timestamp
-
   def lastOffsetDelta : Int = if (isEmpty) 0 else batchMetadata.last.offsetDelta
 
   def isEmpty: Boolean = batchMetadata.isEmpty
 
   def addBatch(producerEpoch: Short, lastSeq: Int, lastOffset: Long, offsetDelta: Int, timestamp: Long): Unit = {
-    maybeUpdateEpoch(producerEpoch)
+    maybeUpdateProducerEpoch(producerEpoch)
     addBatchMetadata(BatchMetadata(lastSeq, lastOffset, offsetDelta, timestamp))
+    this.lastTimestamp = timestamp
   }
 
-  def maybeUpdateEpoch(producerEpoch: Short): Boolean = {
+  def maybeUpdateProducerEpoch(producerEpoch: Short): Boolean = {
     if (this.producerEpoch != producerEpoch) {
       batchMetadata.clear()
       this.producerEpoch = producerEpoch
@@ -134,15 +124,14 @@ private[log] class ProducerStateEntry(val producerId: Long,
   }
 
   def update(nextEntry: ProducerStateEntry): Unit = {
-    maybeUpdateEpoch(nextEntry.producerEpoch)
+    maybeUpdateProducerEpoch(nextEntry.producerEpoch)
     while (nextEntry.batchMetadata.nonEmpty)
       addBatchMetadata(nextEntry.batchMetadata.dequeue())
     this.coordinatorEpoch = nextEntry.coordinatorEpoch
     this.currentTxnFirstOffset = nextEntry.currentTxnFirstOffset
+    this.lastTimestamp = nextEntry.lastTimestamp
   }
 
-  def removeBatchesOlderThan(offset: Long): Unit = batchMetadata.dropWhile(_.lastOffset < offset)
-
   def findDuplicateBatch(batch: RecordBatch): Option[BatchMetadata] = {
     if (batch.producerEpoch != producerEpoch)
        None
@@ -164,6 +153,7 @@ private[log] class ProducerStateEntry(val producerId: Long,
       s"producerEpoch=$producerEpoch, " +
       s"currentTxnFirstOffset=$currentTxnFirstOffset, " +
       s"coordinatorEpoch=$coordinatorEpoch, " +
+      s"lastTimestamp=$lastTimestamp, " +
       s"batchMetadata=$batchMetadata"
   }
 }
@@ -179,50 +169,54 @@ private[log] class ProducerStateEntry(val producerId: Long,
  *                      the most recent appends made by the producer. Validation of the first incoming append will
  *                      be made against the latest append in the current entry. New appends will replace older appends
  *                      in the current entry so that the space overhead is constant.
- * @param validationType Indicates the extent of validation to perform on the appends on this instance. Offset commits
- *                       coming from the producer should have ValidationType.EpochOnly. Appends which aren't from a client
- *                       should have ValidationType.None. Appends coming from a client for produce requests should have
- *                       ValidationType.Full.
+ * @param origin Indicates the origin of the append which implies the extent of validation. For example, offset
+ *               commits, which originate from the group coordinator, do not have sequence numbers and therefore
+ *               only producer epoch validation is done. Appends which come through replciation are not validated
+ *               (we assume the validation has already been done) and appends from clients require full validation.
  */
-private[log] class ProducerAppendInfo(val producerId: Long,
+private[log] class ProducerAppendInfo(val topicPartition: TopicPartition,
+                                      val producerId: Long,
                                       val currentEntry: ProducerStateEntry,
-                                      val validationType: ValidationType) {
+                                      val origin: AppendOrigin) extends Logging {
+
   private val transactions = ListBuffer.empty[TxnMetadata]
   private val updatedEntry = ProducerStateEntry.empty(producerId)
 
   updatedEntry.producerEpoch = currentEntry.producerEpoch
   updatedEntry.coordinatorEpoch = currentEntry.coordinatorEpoch
+  updatedEntry.lastTimestamp = currentEntry.lastTimestamp
   updatedEntry.currentTxnFirstOffset = currentEntry.currentTxnFirstOffset
 
-  private def maybeValidateAppend(producerEpoch: Short, firstSeq: Int) = {
-    validationType match {
-      case ValidationType.None =>
-
-      case ValidationType.EpochOnly =>
-        checkProducerEpoch(producerEpoch)
-
-      case ValidationType.Full =>
-        checkProducerEpoch(producerEpoch)
-        checkSequence(producerEpoch, firstSeq)
+  private def maybeValidateDataBatch(producerEpoch: Short, firstSeq: Int, offset: Long): Unit = {
+    checkProducerEpoch(producerEpoch, offset)
+    if (origin == AppendOrigin.Client) {
+      checkSequence(producerEpoch, firstSeq, offset)
     }
   }
 
-  private def checkProducerEpoch(producerEpoch: Short): Unit = {
+  private def checkProducerEpoch(producerEpoch: Short, offset: Long): Unit = {
     if (producerEpoch < updatedEntry.producerEpoch) {
-      throw new ProducerFencedException(s"Producer's epoch is no longer valid. There is probably another producer " +
-        s"with a newer epoch. $producerEpoch (request epoch), ${updatedEntry.producerEpoch} (server epoch)")
+      val message = s"Producer's epoch at offset $offset in $topicPartition is $producerEpoch, which is " +
+        s"smaller than the last seen epoch ${updatedEntry.producerEpoch}"
+
+      if (origin == AppendOrigin.Replication) {
+        warn(message)
+      } else {
+        throw new ProducerFencedException(message)
+      }
     }
   }
 
-  private def checkSequence(producerEpoch: Short, appendFirstSeq: Int): Unit = {
+  private def checkSequence(producerEpoch: Short, appendFirstSeq: Int, offset: Long): Unit = {
     if (producerEpoch != updatedEntry.producerEpoch) {
       if (appendFirstSeq != 0) {
         if (updatedEntry.producerEpoch != RecordBatch.NO_PRODUCER_EPOCH) {
-          throw new OutOfOrderSequenceException(s"Invalid sequence number for new epoch: $producerEpoch " +
-            s"(request epoch), $appendFirstSeq (seq. number)")
+          throw new OutOfOrderSequenceException(s"Invalid sequence number for new epoch at offset $offset in " +
+            s"partition $topicPartition: $producerEpoch (request epoch), $appendFirstSeq (seq. number)")
         } else {
-          throw new UnknownProducerIdException(s"Found no record of producerId=$producerId on the broker. It is possible " +
-            s"that the last message with the producerId=$producerId has been removed due to hitting the retention limit.")
+          throw new UnknownProducerIdException(s"Found no record of producerId=$producerId on the broker at offset $offset" +
+            s"in partition $topicPartition. It is possible that the last message with the producerId=$producerId has " +
+            "been removed due to hitting the retention limit.")
         }
       }
     } else {
@@ -240,10 +234,12 @@ private[log] class ProducerAppendInfo(val producerId: Long,
         // the sequence number. Note that this check follows the fencing check, so the marker still fences
         // old producers even if it cannot determine our next expected sequence number.
         throw new UnknownProducerIdException(s"Local producer state matches expected epoch $producerEpoch " +
-          s"for producerId=$producerId, but next expected sequence number is not known.")
+          s"for producerId=$producerId at offset $offset in partition $topicPartition, but the next expected " +
+          "sequence number is not known.")
       } else if (!inSequence(currentLastSeq, appendFirstSeq)) {
-        throw new OutOfOrderSequenceException(s"Out of order sequence number for producerId $producerId: $appendFirstSeq " +
-          s"(incoming seq. number), $currentLastSeq (current end sequence number)")
+        throw new OutOfOrderSequenceException(s"Out of order sequence number for producerId $producerId at " +
+          s"offset $offset in partition $topicPartition: $appendFirstSeq (incoming seq. number), " +
+          s"$currentLastSeq (current end sequence number)")
       }
     }
   }
@@ -252,7 +248,7 @@ private[log] class ProducerAppendInfo(val producerId: Long,
     nextSeq == lastSeq + 1L || (nextSeq == 0 && lastSeq == Int.MaxValue)
   }
 
-  def append(batch: RecordBatch): Option[CompletedTxn] = {
+  def append(batch: RecordBatch, firstOffsetMetadataOpt: Option[LogOffsetMetadata]): Option[CompletedTxn] = {
     if (batch.isControlBatch) {
       val recordIterator = batch.iterator
       if (recordIterator.hasNext) {
@@ -265,47 +261,59 @@ private[log] class ProducerAppendInfo(val producerId: Long,
         None
       }
     } else {
-      append(batch.producerEpoch, batch.baseSequence, batch.lastSequence, batch.maxTimestamp, batch.baseOffset, batch.lastOffset,
-        batch.isTransactional)
+      val firstOffsetMetadata = firstOffsetMetadataOpt.getOrElse(LogOffsetMetadata(batch.baseOffset))
+      appendDataBatch(batch.producerEpoch, batch.baseSequence, batch.lastSequence, batch.maxTimestamp,
+        firstOffsetMetadata, batch.lastOffset, batch.isTransactional)
       None
     }
   }
 
-  def append(epoch: Short,
-             firstSeq: Int,
-             lastSeq: Int,
-             lastTimestamp: Long,
-             firstOffset: Long,
-             lastOffset: Long,
-             isTransactional: Boolean): Unit = {
-    maybeValidateAppend(epoch, firstSeq)
+  def appendDataBatch(epoch: Short,
+                      firstSeq: Int,
+                      lastSeq: Int,
+                      lastTimestamp: Long,
+                      firstOffsetMetadata: LogOffsetMetadata,
+                      lastOffset: Long,
+                      isTransactional: Boolean): Unit = {
+    val firstOffset = firstOffsetMetadata.messageOffset
+    maybeValidateDataBatch(epoch, firstSeq, firstOffset)
     updatedEntry.addBatch(epoch, lastSeq, lastOffset, (lastOffset - firstOffset).toInt, lastTimestamp)
 
     updatedEntry.currentTxnFirstOffset match {
       case Some(_) if !isTransactional =>
         // Received a non-transactional message while a transaction is active
-        throw new InvalidTxnStateException(s"Expected transactional write from producer $producerId")
+        throw new InvalidTxnStateException(s"Expected transactional write from producer $producerId at " +
+          s"offset $firstOffsetMetadata in partition $topicPartition")
 
       case None if isTransactional =>
         // Began a new transaction
         updatedEntry.currentTxnFirstOffset = Some(firstOffset)
-        transactions += new TxnMetadata(producerId, firstOffset)
+        transactions += TxnMetadata(producerId, firstOffsetMetadata)
 
       case _ => // nothing to do
     }
   }
 
+  private def checkCoordinatorEpoch(endTxnMarker: EndTransactionMarker, offset: Long): Unit = {
+    if (updatedEntry.coordinatorEpoch > endTxnMarker.coordinatorEpoch) {
+      if (origin == AppendOrigin.Replication) {
+        info(s"Detected invalid coordinator epoch for producerId $producerId at " +
+          s"offset $offset in partition $topicPartition: ${endTxnMarker.coordinatorEpoch} " +
+          s"is older than previously known coordinator epoch ${updatedEntry.coordinatorEpoch}")
+      } else {
+        throw new TransactionCoordinatorFencedException(s"Invalid coordinator epoch for producerId $producerId at " +
+          s"offset $offset in partition $topicPartition: ${endTxnMarker.coordinatorEpoch} " +
+          s"(zombie), ${updatedEntry.coordinatorEpoch} (current)")
+      }
+    }
+  }
+
   def appendEndTxnMarker(endTxnMarker: EndTransactionMarker,
                          producerEpoch: Short,
                          offset: Long,
                          timestamp: Long): CompletedTxn = {
-    checkProducerEpoch(producerEpoch)
-
-    if (updatedEntry.coordinatorEpoch > endTxnMarker.coordinatorEpoch)
-      throw new TransactionCoordinatorFencedException(s"Invalid coordinator epoch: ${endTxnMarker.coordinatorEpoch} " +
-        s"(zombie), ${updatedEntry.coordinatorEpoch} (current)")
-
-    updatedEntry.maybeUpdateEpoch(producerEpoch)
+    checkProducerEpoch(producerEpoch, offset)
+    checkCoordinatorEpoch(endTxnMarker, offset)
 
     val firstOffset = updatedEntry.currentTxnFirstOffset match {
       case Some(txnFirstOffset) => txnFirstOffset
@@ -314,8 +322,11 @@ private[log] class ProducerAppendInfo(val producerId: Long,
         offset
     }
 
+    updatedEntry.maybeUpdateProducerEpoch(producerEpoch)
     updatedEntry.currentTxnFirstOffset = None
     updatedEntry.coordinatorEpoch = endTxnMarker.coordinatorEpoch
+    updatedEntry.lastTimestamp = timestamp
+
     CompletedTxn(producerId, firstOffset, offset, endTxnMarker.controlType == ControlRecordType.ABORT)
   }
 
@@ -323,18 +334,6 @@ private[log] class ProducerAppendInfo(val producerId: Long,
 
   def startedTransactions: List[TxnMetadata] = transactions.toList
 
-  def maybeCacheTxnFirstOffsetMetadata(logOffsetMetadata: LogOffsetMetadata): Unit = {
-    // we will cache the log offset metadata if it corresponds to the starting offset of
-    // the last transaction that was started. This is optimized for leader appends where it
-    // is only possible to have one transaction started for each log append, and the log
-    // offset metadata will always match in that case since no data from other producers
-    // is mixed into the append
-    transactions.headOption.foreach { txn =>
-      if (txn.firstOffset.messageOffset == logOffsetMetadata.messageOffset)
-        txn.firstOffset = logOffsetMetadata
-    }
-  }
-
   override def toString: String = {
     "ProducerAppendInfo(" +
       s"producerId=$producerId, " +
@@ -343,6 +342,7 @@ private[log] class ProducerAppendInfo(val producerId: Long,
       s"lastSequence=${updatedEntry.lastSeq}, " +
       s"currentTxnFirstOffset=${updatedEntry.currentTxnFirstOffset}, " +
       s"coordinatorEpoch=${updatedEntry.coordinatorEpoch}, " +
+      s"lastTimestamp=${updatedEntry.lastTimestamp}, " +
       s"startedTransactions=$transactions)"
   }
 }
@@ -396,7 +396,7 @@ object ProducerStateManager {
 
       struct.getArray(ProducerEntriesField).map { producerEntryObj =>
         val producerEntryStruct = producerEntryObj.asInstanceOf[Struct]
-        val producerId: Long = producerEntryStruct.getLong(ProducerIdField)
+        val producerId = producerEntryStruct.getLong(ProducerIdField)
         val producerEpoch = producerEntryStruct.getShort(ProducerEpochField)
         val seq = producerEntryStruct.getInt(LastSequenceField)
         val offset = producerEntryStruct.getLong(LastOffsetField)
@@ -404,8 +404,12 @@ object ProducerStateManager {
         val offsetDelta = producerEntryStruct.getInt(OffsetDeltaField)
         val coordinatorEpoch = producerEntryStruct.getInt(CoordinatorEpochField)
         val currentTxnFirstOffset = producerEntryStruct.getLong(CurrentTxnFirstOffsetField)
-        val newEntry = new ProducerStateEntry(producerId, mutable.Queue[BatchMetadata](BatchMetadata(seq, offset, offsetDelta, timestamp)), producerEpoch,
-          coordinatorEpoch, if (currentTxnFirstOffset >= 0) Some(currentTxnFirstOffset) else None)
+        val lastAppendedDataBatches = mutable.Queue.empty[BatchMetadata]
+        if (offset >= 0)
+          lastAppendedDataBatches += BatchMetadata(seq, offset, offsetDelta, timestamp)
+
+        val newEntry = new ProducerStateEntry(producerId, lastAppendedDataBatches, producerEpoch,
+          coordinatorEpoch, timestamp, if (currentTxnFirstOffset >= 0) Some(currentTxnFirstOffset) else None)
         newEntry
       }
     } catch {
@@ -414,7 +418,7 @@ object ProducerStateManager {
     }
   }
 
-  private def writeSnapshot(file: File, entries: mutable.Map[Long, ProducerStateEntry]) {
+  private def writeSnapshot(file: File, entries: mutable.Map[Long, ProducerStateEntry]): Unit = {
     val struct = new Struct(PidSnapshotMapSchema)
     struct.set(VersionField, ProducerSnapshotVersion)
     struct.set(CrcField, 0L) // we'll fill this after writing the entries
@@ -460,7 +464,7 @@ object ProducerStateManager {
   // visible for testing
   private[log] def deleteSnapshotsBefore(dir: File, offset: Long): Unit = deleteSnapshotFiles(dir, _ < offset)
 
-  private def deleteSnapshotFiles(dir: File, predicate: Long => Boolean = _ => true) {
+  private def deleteSnapshotFiles(dir: File, predicate: Long => Boolean = _ => true): Unit = {
     listSnapshotFiles(dir).filter(file => predicate(offsetFromFile(file))).foreach { file =>
       Files.deleteIfExists(file.toPath)
     }
@@ -538,7 +542,7 @@ class ProducerStateManager(val topicPartition: TopicPartition,
   /**
    * Returns the last offset of this map
    */
-  def mapEndOffset = lastMapOffset
+  def mapEndOffset: Long = lastMapOffset
 
   /**
    * Get a copy of the active producers
@@ -547,15 +551,13 @@ class ProducerStateManager(val topicPartition: TopicPartition,
 
   def isEmpty: Boolean = producers.isEmpty && unreplicatedTxns.isEmpty
 
-  private def loadFromSnapshot(logStartOffset: Long, currentTime: Long) {
+  private def loadFromSnapshot(logStartOffset: Long, currentTime: Long): Unit = {
     while (true) {
       latestSnapshotFile match {
         case Some(file) =>
           try {
             info(s"Loading producer state from snapshot file '$file'")
-            val loadedProducers = readSnapshot(file).filter { producerEntry =>
-              isProducerRetained(producerEntry, logStartOffset) && !isProducerExpired(currentTime, producerEntry)
-            }
+            val loadedProducers = readSnapshot(file).filter { producerEntry => !isProducerExpired(currentTime, producerEntry) }
             loadedProducers.foreach(loadProducerEntry)
             lastSnapOffset = offsetFromFile(file)
             lastMapOffset = lastSnapOffset
@@ -588,7 +590,7 @@ class ProducerStateManager(val topicPartition: TopicPartition,
   /**
    * Expire any producer ids which have been idle longer than the configured maximum expiration timeout.
    */
-  def removeExpiredProducers(currentTimeMs: Long) {
+  def removeExpiredProducers(currentTimeMs: Long): Unit = {
     producers.retain { case (_, lastEntry) =>
       !isProducerExpired(currentTimeMs, lastEntry)
     }
@@ -596,10 +598,13 @@ class ProducerStateManager(val topicPartition: TopicPartition,
 
   /**
    * Truncate the producer id mapping to the given offset range and reload the entries from the most recent
-   * snapshot in range (if there is one). Note that the log end offset is assumed to be less than
-   * or equal to the high watermark.
+   * snapshot in range (if there is one). We delete snapshot files prior to the logStartOffset but do not remove
+   * producer state from the map. This means that in-memory and on-disk state can diverge, and in the case of
+   * broker failover or unclean shutdown, any in-memory state not persisted in the snapshots will be lost, which
+   * would lead to UNKNOWN_PRODUCER_ID errors. Note that the log end offset is assumed to be less than or equal
+   * to the high watermark.
    */
-  def truncateAndReload(logStartOffset: Long, logEndOffset: Long, currentTimeMs: Long) {
+  def truncateAndReload(logStartOffset: Long, logEndOffset: Long, currentTimeMs: Long): Unit = {
     // remove all out of range snapshots
     deleteSnapshotFiles(logDir, { snapOffset =>
       snapOffset > logEndOffset || snapOffset <= logStartOffset
@@ -618,17 +623,9 @@ class ProducerStateManager(val topicPartition: TopicPartition,
     }
   }
 
-  def prepareUpdate(producerId: Long, isFromClient: Boolean): ProducerAppendInfo = {
-    val validationToPerform =
-      if (!isFromClient)
-        ValidationType.None
-      else if (topicPartition.topic == Topic.GROUP_METADATA_TOPIC_NAME)
-        ValidationType.EpochOnly
-      else
-        ValidationType.Full
-
+  def prepareUpdate(producerId: Long, origin: AppendOrigin): ProducerAppendInfo = {
     val currentEntry = lastEntry(producerId).getOrElse(ProducerStateEntry.empty(producerId))
-    new ProducerAppendInfo(producerId, currentEntry, validationToPerform)
+    new ProducerAppendInfo(topicPartition, producerId, currentEntry, origin)
   }
 
   /**
@@ -688,28 +685,11 @@ class ProducerStateManager(val topicPartition: TopicPartition,
    */
   def oldestSnapshotOffset: Option[Long] = oldestSnapshotFile.map(file => offsetFromFile(file))
 
-  private def isProducerRetained(producerStateEntry: ProducerStateEntry, logStartOffset: Long): Boolean = {
-    producerStateEntry.removeBatchesOlderThan(logStartOffset)
-    producerStateEntry.lastDataOffset >= logStartOffset
-  }
-
   /**
-   * When we remove the head of the log due to retention, we need to clean up the id map. This method takes
-   * the new start offset and removes all producerIds which have a smaller last written offset. Additionally,
-   * we remove snapshots older than the new log start offset.
-   *
-   * Note that snapshots from offsets greater than the log start offset may have producers included which
-   * should no longer be retained: these producers will be removed if and when we need to load state from
-   * the snapshot.
+   * When we remove the head of the log due to retention, we need to remove snapshots older than
+   * the new log start offset.
    */
-  def truncateHead(logStartOffset: Long) {
-    val evictedProducerEntries = producers.filter { case (_, producerState) =>
-      !isProducerRetained(producerState, logStartOffset)
-    }
-    val evictedProducerIds = evictedProducerEntries.keySet
-
-    producers --= evictedProducerIds
-    removeEvictedOngoingTransactions(evictedProducerIds)
+  def truncateHead(logStartOffset: Long): Unit = {
     removeUnreplicatedTransactions(logStartOffset)
 
     if (lastMapOffset < logStartOffset)
@@ -719,15 +699,6 @@ class ProducerStateManager(val topicPartition: TopicPartition,
     lastSnapOffset = latestSnapshotOffset.getOrElse(logStartOffset)
   }
 
-  private def removeEvictedOngoingTransactions(expiredProducerIds: collection.Set[Long]): Unit = {
-    val iterator = ongoingTxns.entrySet.iterator
-    while (iterator.hasNext) {
-      val txnEntry = iterator.next()
-      if (expiredProducerIds.contains(txnEntry.getValue.producerId))
-        iterator.remove()
-    }
-  }
-
   private def removeUnreplicatedTransactions(offset: Long): Unit = {
     val iterator = unreplicatedTxns.entrySet.iterator
     while (iterator.hasNext) {
@@ -741,7 +712,7 @@ class ProducerStateManager(val topicPartition: TopicPartition,
   /**
    * Truncate the producer id mapping and remove all snapshots. This resets the state of the mapping.
    */
-  def truncate() {
+  def truncate(): Unit = {
     producers.clear()
     ongoingTxns.clear()
     unreplicatedTxns.clear()
@@ -751,9 +722,20 @@ class ProducerStateManager(val topicPartition: TopicPartition,
   }
 
   /**
-   * Complete the transaction and return the last stable offset.
+   * Compute the last stable offset of a completed transaction, but do not yet mark the transaction complete.
+   * That will be done in `completeTxn` below. This is used to compute the LSO that will be appended to the
+   * transaction index, but the completion must be done only after successfully appending to the index.
    */
-  def completeTxn(completedTxn: CompletedTxn): Long = {
+  def lastStableOffset(completedTxn: CompletedTxn): Long = {
+    val nextIncompleteTxn = ongoingTxns.values.asScala.find(_.producerId != completedTxn.producerId)
+    nextIncompleteTxn.map(_.firstOffset.messageOffset).getOrElse(completedTxn.lastOffset  + 1)
+  }
+
+  /**
+   * Mark a transaction as completed. We will still await advancement of the high watermark before
+   * advancing the first unstable offset.
+   */
+  def completeTxn(completedTxn: CompletedTxn): Unit = {
     val txnMetadata = ongoingTxns.remove(completedTxn.firstOffset)
     if (txnMetadata == null)
       throw new IllegalArgumentException(s"Attempted to complete transaction $completedTxn on partition $topicPartition " +
@@ -761,9 +743,6 @@ class ProducerStateManager(val topicPartition: TopicPartition,
 
     txnMetadata.lastOffset = Some(completedTxn.lastOffset)
     unreplicatedTxns.put(completedTxn.firstOffset, txnMetadata)
-
-    val lastStableOffset = firstUndecidedOffset.getOrElse(completedTxn.lastOffset + 1)
-    lastStableOffset
   }
 
   @threadsafe
diff --git a/core/src/main/scala/kafka/log/TimeIndex.scala b/core/src/main/scala/kafka/log/TimeIndex.scala
index f5c4d90aac5ad..79516ba0a04d8 100644
--- a/core/src/main/scala/kafka/log/TimeIndex.scala
+++ b/core/src/main/scala/kafka/log/TimeIndex.scala
@@ -51,13 +51,17 @@ import org.apache.kafka.common.record.RecordBatch
  */
 // Avoid shadowing mutable file in AbstractIndex
 class TimeIndex(_file: File, baseOffset: Long, maxIndexSize: Int = -1, writable: Boolean = true)
-    extends AbstractIndex[Long, Long](_file, baseOffset, maxIndexSize, writable) {
+    extends AbstractIndex(_file, baseOffset, maxIndexSize, writable) {
   import TimeIndex._
 
   @volatile private var _lastEntry = lastEntryFromIndexFile
 
   override def entrySize = 12
 
+  if (isDebugEnabled)
+    debug(s"Loaded index file ${file.getAbsolutePath} with maxEntries = $maxEntries, maxIndexSize = $maxIndexSize," +
+          s" entries = ${_entries}, lastOffset = ${_lastEntry}, file position = ${mmap.position()}")
+
   // We override the full check to reserve the last time index entry slot for the on roll call.
   override def isFull: Boolean = entries >= maxEntries - 1
 
@@ -87,7 +91,8 @@ class TimeIndex(_file: File, baseOffset: Long, maxIndexSize: Int = -1, writable:
   def entry(n: Int): TimestampOffset = {
     maybeLock(lock) {
       if(n >= _entries)
-        throw new IllegalArgumentException("Attempt to fetch the %dth entry from a time index of size %d.".format(n, _entries))
+        throw new IllegalArgumentException(s"Attempt to fetch the ${n}th entry from  time index ${file.getAbsolutePath} " +
+          s"which has size ${_entries}.")
       parseEntry(mmap, n)
     }
   }
@@ -106,7 +111,7 @@ class TimeIndex(_file: File, baseOffset: Long, maxIndexSize: Int = -1, writable:
    * @param skipFullCheck To skip checking whether the segment is full or not. We only skip the check when the segment
    *                      gets rolled or the segment is closed.
    */
-  def maybeAppend(timestamp: Long, offset: Long, skipFullCheck: Boolean = false) {
+  def maybeAppend(timestamp: Long, offset: Long, skipFullCheck: Boolean = false): Unit = {
     inLock(lock) {
       if (!skipFullCheck)
         require(!isFull, "Attempt to append to a full time index (size = " + _entries + ").")
@@ -116,17 +121,20 @@ class TimeIndex(_file: File, baseOffset: Long, maxIndexSize: Int = -1, writable:
       // because that could happen in the following two scenarios:
       // 1. A log segment is closed.
       // 2. LogSegment.onBecomeInactiveSegment() is called when an active log segment is rolled.
-      if (_entries != 0 && offset < lastEntry.offset)
-        throw new InvalidOffsetException("Attempt to append an offset (%d) to slot %d no larger than the last offset appended (%d) to %s."
-          .format(offset, _entries, lastEntry.offset, file.getAbsolutePath))
-      if (_entries != 0 && timestamp < lastEntry.timestamp)
-        throw new IllegalStateException("Attempt to append a timestamp (%d) to slot %d no larger than the last timestamp appended (%d) to %s."
-            .format(timestamp, _entries, lastEntry.timestamp, file.getAbsolutePath))
+      if (_entries != 0) {
+        if (offset < lastEntry.offset)
+          throw new InvalidOffsetException(s"Attempt to append an offset ($offset) to slot ${_entries} no larger than" +
+                                           s" the last offset appended (${lastEntry.offset}) to ${file.getAbsolutePath}.")
+        if (timestamp < lastEntry.timestamp)
+          throw new IllegalStateException(s"Attempt to append a timestamp ($timestamp) to slot ${_entries} no larger" +
+                                          s" than the last timestamp appended (${lastEntry.timestamp}) to ${file.getAbsolutePath}.")
+      }
       // We only append to the time index when the timestamp is greater than the last inserted timestamp.
       // If all the messages are in message format v0, the timestamp will always be NoTimestamp. In that case, the time
       // index will be empty.
       if (timestamp > lastEntry.timestamp) {
-        debug("Adding index entry %d => %d to %s.".format(timestamp, offset, file.getName))
+        if (isTraceEnabled)
+          trace(s"Adding index entry $timestamp => $offset to ${file.getAbsolutePath}.")
         mmap.putLong(timestamp)
         mmap.putInt(relativeOffset(offset))
         _entries += 1
@@ -161,7 +169,7 @@ class TimeIndex(_file: File, baseOffset: Long, maxIndexSize: Int = -1, writable:
    * Remove all entries from the index which have an offset greater than or equal to the given offset.
    * Truncating to an offset larger than the largest in the index has no effect.
    */
-  override def truncateTo(offset: Long) {
+  override def truncateTo(offset: Long): Unit = {
     inLock(lock) {
       val idx = mmap.duplicate
       val slot = largestLowerBoundSlotFor(idx, offset, IndexSearchType.VALUE)
@@ -195,15 +203,17 @@ class TimeIndex(_file: File, baseOffset: Long, maxIndexSize: Int = -1, writable:
   /**
    * Truncates index to a known number of entries.
    */
-  private def truncateToEntries(entries: Int) {
+  private def truncateToEntries(entries: Int): Unit = {
     inLock(lock) {
       _entries = entries
       mmap.position(_entries * entrySize)
       _lastEntry = lastEntryFromIndexFile
+      if (isDebugEnabled)
+        debug(s"Truncated index ${file.getAbsolutePath} to $entries entries; position is now ${mmap.position()} and last entry is now ${_lastEntry}")
     }
   }
 
-  override def sanityCheck() {
+  override def sanityCheck(): Unit = {
     val lastTimestamp = lastEntry.timestamp
     val lastOffset = lastEntry.offset
     if (_entries != 0 && lastTimestamp < timestamp(mmap, 0))
@@ -222,37 +232,3 @@ class TimeIndex(_file: File, baseOffset: Long, maxIndexSize: Int = -1, writable:
 object TimeIndex extends Logging {
   override val loggerName: String = classOf[TimeIndex].getName
 }
-
-
-
-/**
-  * A thin wrapper on top of the raw TimeIndex object to avoid initialization on construction. This defers the TimeIndex
-  * initialization to the time it gets accessed so the cost of the heavy memory mapped operation gets amortized over time.
-  *
-  * Combining with skipping sanity check for safely flushed segments, the startup time of a broker can be reduced, especially
-  * for the the broker with a lot of log segments
-  *
-  */
-class LazyTimeIndex(@volatile private var _file: File, baseOffset: Long, maxIndexSize: Int = -1, writable: Boolean = true) {
-  @volatile private var timeIndex: Option[TimeIndex] = None
-
-  def file: File = {
-    if (timeIndex.isDefined)
-      timeIndex.get.file
-    else
-      _file
-  }
-
-  def file_=(f: File) {
-    if (timeIndex.isDefined)
-      timeIndex.get.file = f
-    else
-      _file = f
-  }
-
-  def get: TimeIndex = {
-    if (timeIndex.isEmpty)
-      timeIndex = Some(new TimeIndex(_file, baseOffset, maxIndexSize, writable))
-    timeIndex.get
-  }
-}
\ No newline at end of file
diff --git a/core/src/main/scala/kafka/log/TransactionIndex.scala b/core/src/main/scala/kafka/log/TransactionIndex.scala
index da7fce815dcf9..9152bc41ab353 100644
--- a/core/src/main/scala/kafka/log/TransactionIndex.scala
+++ b/core/src/main/scala/kafka/log/TransactionIndex.scala
@@ -42,25 +42,31 @@ private[log] case class TxnIndexSearchResult(abortedTransactions: List[AbortedTx
  * order to find the start of the transactions.
  */
 @nonthreadsafe
-class TransactionIndex(val startOffset: Long, @volatile var file: File) extends Logging {
+class TransactionIndex(val startOffset: Long, @volatile private var _file: File) extends Logging {
+
   // note that the file is not created until we need it
   @volatile private var maybeChannel: Option[FileChannel] = None
   private var lastOffset: Option[Long] = None
 
-  if (file.exists)
+  if (_file.exists)
     openChannel()
 
   def append(abortedTxn: AbortedTxn): Unit = {
     lastOffset.foreach { offset =>
       if (offset >= abortedTxn.lastOffset)
-        throw new IllegalArgumentException("The last offset of appended transactions must increase sequentially")
+        throw new IllegalArgumentException(s"The last offset of appended transactions must increase sequentially, but " +
+          s"${abortedTxn.lastOffset} is not greater than current last offset $offset of index ${file.getAbsolutePath}")
     }
     lastOffset = Some(abortedTxn.lastOffset)
-    Utils.writeFully(channel, abortedTxn.buffer.duplicate())
+    Utils.writeFully(channel(), abortedTxn.buffer.duplicate())
   }
 
   def flush(): Unit = maybeChannel.foreach(_.force(true))
 
+  def file: File = _file
+
+  def updateParentDir(parentDir: File): Unit = _file = new File(parentDir, file.getName)
+
   /**
    * Delete this index.
    *
@@ -73,7 +79,7 @@ class TransactionIndex(val startOffset: Long, @volatile var file: File) extends
     Files.deleteIfExists(file.toPath)
   }
 
-  private def channel: FileChannel = {
+  private def channel(): FileChannel = {
     maybeChannel match {
       case Some(channel) => channel
       case None => openChannel()
@@ -105,7 +111,7 @@ class TransactionIndex(val startOffset: Long, @volatile var file: File) extends
     try {
       if (file.exists)
         Utils.atomicMoveWithFallback(file.toPath, f.toPath)
-    } finally file = f
+    } finally _file = f
   }
 
   def truncateTo(offset: Long): Unit = {
@@ -113,7 +119,7 @@ class TransactionIndex(val startOffset: Long, @volatile var file: File) extends
     var newLastOffset: Option[Long] = None
     for ((abortedTxn, position) <- iterator(() => buffer)) {
       if (abortedTxn.lastOffset >= offset) {
-        channel.truncate(position)
+        channel().truncate(position)
         lastOffset = newLastOffset
         return
       }
@@ -138,8 +144,8 @@ class TransactionIndex(val startOffset: Long, @volatile var file: File) extends
 
               val abortedTxn = new AbortedTxn(buffer)
               if (abortedTxn.version > AbortedTxn.CurrentVersion)
-                throw new KafkaException(s"Unexpected aborted transaction version ${abortedTxn.version}, " +
-                  s"current version is ${AbortedTxn.CurrentVersion}")
+                throw new KafkaException(s"Unexpected aborted transaction version ${abortedTxn.version} " +
+                  s"in transaction index ${file.getAbsolutePath}, current version is ${AbortedTxn.CurrentVersion}")
               val nextEntry = (abortedTxn, position)
               position += AbortedTxn.TotalSize
               nextEntry
@@ -147,7 +153,7 @@ class TransactionIndex(val startOffset: Long, @volatile var file: File) extends
               case e: IOException =>
                 // We received an unexpected error reading from the index file. We propagate this as an
                 // UNKNOWN error to the consumer, which will cause it to retry the fetch.
-                throw new KafkaException(s"Failed to read from the transaction index $file", e)
+                throw new KafkaException(s"Failed to read from the transaction index ${file.getAbsolutePath}", e)
             }
           }
         }
@@ -187,8 +193,8 @@ class TransactionIndex(val startOffset: Long, @volatile var file: File) extends
     val buffer = ByteBuffer.allocate(AbortedTxn.TotalSize)
     for ((abortedTxn, _) <- iterator(() => buffer)) {
       if (abortedTxn.lastOffset < startOffset)
-        throw new CorruptIndexException(s"Last offset of aborted transaction $abortedTxn is less than start offset " +
-          s"$startOffset")
+        throw new CorruptIndexException(s"Last offset of aborted transaction $abortedTxn in index " +
+          s"${file.getAbsolutePath} is less than start offset $startOffset")
     }
   }
 
diff --git a/core/src/main/scala/kafka/metrics/KafkaCSVMetricsReporter.scala b/core/src/main/scala/kafka/metrics/KafkaCSVMetricsReporter.scala
index 81c20a7e117bc..82317277814c5 100755
--- a/core/src/main/scala/kafka/metrics/KafkaCSVMetricsReporter.scala
+++ b/core/src/main/scala/kafka/metrics/KafkaCSVMetricsReporter.scala
@@ -45,7 +45,7 @@ private class KafkaCSVMetricsReporter extends KafkaMetricsReporter
   override def getMBeanName = "kafka:type=kafka.metrics.KafkaCSVMetricsReporter"
 
 
-  override def init(props: VerifiableProperties) {
+  override def init(props: VerifiableProperties): Unit = {
     synchronized {
       if (!initialized) {
         val metricsConfig = new KafkaMetricsConfig(props)
@@ -62,7 +62,7 @@ private class KafkaCSVMetricsReporter extends KafkaMetricsReporter
   }
 
 
-  override def startReporter(pollingPeriodSecs: Long) {
+  override def startReporter(pollingPeriodSecs: Long): Unit = {
     synchronized {
       if (initialized && !running) {
         underlying.start(pollingPeriodSecs, TimeUnit.SECONDS)
@@ -73,7 +73,7 @@ private class KafkaCSVMetricsReporter extends KafkaMetricsReporter
   }
 
 
-  override def stopReporter() {
+  override def stopReporter(): Unit = {
     synchronized {
       if (initialized && running) {
         underlying.shutdown()
diff --git a/core/src/main/scala/kafka/metrics/KafkaMetricsGroup.scala b/core/src/main/scala/kafka/metrics/KafkaMetricsGroup.scala
index 03a4f7ccd4c03..a5f9911a2b077 100644
--- a/core/src/main/scala/kafka/metrics/KafkaMetricsGroup.scala
+++ b/core/src/main/scala/kafka/metrics/KafkaMetricsGroup.scala
@@ -71,6 +71,9 @@ trait KafkaMetricsGroup extends Logging {
   def newMeter(name: String, eventType: String, timeUnit: TimeUnit, tags: scala.collection.Map[String, String] = Map.empty) =
     Metrics.defaultRegistry().newMeter(metricName(name, tags), eventType, timeUnit)
 
+  def newCounter(name: String, tags: scala.collection.Map[String, String] = Map.empty) =
+    Metrics.defaultRegistry().newCounter(metricName(name, tags))
+
   def newHistogram(name: String, biased: Boolean = true, tags: scala.collection.Map[String, String] = Map.empty) =
     Metrics.defaultRegistry().newHistogram(metricName(name, tags), biased)
 
diff --git a/core/src/main/scala/kafka/metrics/KafkaMetricsReporter.scala b/core/src/main/scala/kafka/metrics/KafkaMetricsReporter.scala
index 6d35539915eeb..814cdb2b91b7b 100755
--- a/core/src/main/scala/kafka/metrics/KafkaMetricsReporter.scala
+++ b/core/src/main/scala/kafka/metrics/KafkaMetricsReporter.scala
@@ -23,6 +23,7 @@ package kafka.metrics
 import kafka.utils.{CoreUtils, VerifiableProperties}
 import java.util.concurrent.atomic.AtomicBoolean
 
+import scala.collection.Seq
 import scala.collection.mutable.ArrayBuffer
 
 
@@ -34,9 +35,8 @@ import scala.collection.mutable.ArrayBuffer
  * registered MBean is compliant with the standard MBean convention.
  */
 trait KafkaMetricsReporterMBean {
-  def startReporter(pollingPeriodInSeconds: Long)
-  def stopReporter()
-
+  def startReporter(pollingPeriodInSeconds: Long): Unit
+  def stopReporter(): Unit
   /**
    *
    * @return The name with which the MBean will be registered.
@@ -48,7 +48,7 @@ trait KafkaMetricsReporterMBean {
   * Implement {@link org.apache.kafka.common.ClusterResourceListener} to receive cluster metadata once it's available. Please see the class documentation for ClusterResourceListener for more information.
   */
 trait KafkaMetricsReporter {
-  def init(props: VerifiableProperties)
+  def init(props: VerifiableProperties): Unit
 }
 
 object KafkaMetricsReporter {
diff --git a/core/src/main/scala/kafka/network/RequestChannel.scala b/core/src/main/scala/kafka/network/RequestChannel.scala
index 35d64f0e2901b..542b82e00d4b5 100644
--- a/core/src/main/scala/kafka/network/RequestChannel.scala
+++ b/core/src/main/scala/kafka/network/RequestChannel.scala
@@ -115,7 +115,7 @@ object RequestChannel extends Logging {
       math.max(apiLocalCompleteTimeNanos - requestDequeueTimeNanos, 0L)
     }
 
-    def updateRequestMetrics(networkThreadTimeNanos: Long, response: Response) {
+    def updateRequestMetrics(networkThreadTimeNanos: Long, response: Response): Unit = {
       val endTimeNanos = Time.SYSTEM.nanoseconds
       // In some corner cases, apiLocalCompleteTimeNanos may not be set when the request completes if the remote
       // processing time is really small. This value is set in KafkaApis from a request handling thread.
@@ -272,7 +272,7 @@ object RequestChannel extends Logging {
   }
 }
 
-class RequestChannel(val queueSize: Int, val metricNamePrefix : String) extends KafkaMetricsGroup {
+class RequestChannel(val queueSize: Int, val metricNamePrefix : String, val time: Time) extends KafkaMetricsGroup {
   import RequestChannel._
   val metrics = new RequestChannel.Metrics
   private val requestQueue = new ArrayBlockingQueue[BaseRequest](queueSize)
@@ -280,6 +280,13 @@ class RequestChannel(val queueSize: Int, val metricNamePrefix : String) extends
   val requestQueueSizeMetricName = metricNamePrefix.concat(RequestQueueSizeMetric)
   val responseQueueSizeMetricName = metricNamePrefix.concat(ResponseQueueSizeMetric)
 
+  // Set this to Long.Maxvalue so that KafkaHealthCheck will not shutdown broker if it
+  // reads lastDequeueTimeMs before lastDequeueTimeMs is updated by any KafkaRequestHandler thread.
+  @volatile var lastDequeueTimeMs = Long.MaxValue
+  // This metric can help user select a suitable threshold for requestMaxLocalTimeMs so that broker can shutdown itself only when it
+  // is stuck or too slow. A suggested value of requestMaxLocalTimeMs could be twice the 999'th percentile of the RequestDequeuePollIntervalMs.
+  private val requestDequeuePollIntervalMs = newHistogram("RequestDequeuePollIntervalMs")
+
   newGauge(requestQueueSizeMetricName, new Gauge[Int] {
       def value = requestQueue.size
   })
@@ -308,12 +315,12 @@ class RequestChannel(val queueSize: Int, val metricNamePrefix : String) extends
   }
 
   /** Send a request to be handled, potentially blocking until there is room in the queue for the request */
-  def sendRequest(request: RequestChannel.Request) {
+  def sendRequest(request: RequestChannel.Request): Unit = {
     requestQueue.put(request)
   }
 
   /** Send a response back to the socket server to be sent over the network */
-  def sendResponse(response: RequestChannel.Response) {
+  def sendResponse(response: RequestChannel.Response): Unit = {
     if (isTraceEnabled) {
       val requestHeader = response.request.header
       val message = response match {
@@ -340,24 +347,32 @@ class RequestChannel(val queueSize: Int, val metricNamePrefix : String) extends
   }
 
   /** Get the next request or block until specified time has elapsed */
-  def receiveRequest(timeout: Long): RequestChannel.BaseRequest =
+  def receiveRequest(timeout: Long): RequestChannel.BaseRequest = {
+    val curTime = time.milliseconds
+    requestDequeuePollIntervalMs.update(curTime - lastDequeueTimeMs)
+    lastDequeueTimeMs = curTime
     requestQueue.poll(timeout, TimeUnit.MILLISECONDS)
+  }
 
   /** Get the next request or block until there is one */
-  def receiveRequest(): RequestChannel.BaseRequest =
+  def receiveRequest(): RequestChannel.BaseRequest = {
+    val curTime = time.milliseconds
+    requestDequeuePollIntervalMs.update(curTime - lastDequeueTimeMs)
+    lastDequeueTimeMs = curTime
     requestQueue.take()
+  }
 
-  def updateErrorMetrics(apiKey: ApiKeys, errors: collection.Map[Errors, Integer]) {
+  def updateErrorMetrics(apiKey: ApiKeys, errors: collection.Map[Errors, Integer]): Unit = {
     errors.foreach { case (error, count) =>
       metrics(apiKey.name).markErrorMeter(error, count)
     }
   }
 
-  def clear() {
+  def clear(): Unit = {
     requestQueue.clear()
   }
 
-  def shutdown() {
+  def shutdown(): Unit = {
     clear()
     metrics.close()
   }
@@ -453,7 +468,7 @@ class RequestMetrics(name: String) extends KafkaMetricsGroup {
     }
   }
 
-  def markErrorMeter(error: Errors, count: Int) {
+  def markErrorMeter(error: Errors, count: Int): Unit = {
     errorMeters(error).getOrCreateMeter().mark(count.toLong)
   }
 
diff --git a/core/src/main/scala/kafka/network/SocketServer.scala b/core/src/main/scala/kafka/network/SocketServer.scala
index d8c4a74cdd25c..48d2a5000e337 100644
--- a/core/src/main/scala/kafka/network/SocketServer.scala
+++ b/core/src/main/scala/kafka/network/SocketServer.scala
@@ -21,6 +21,8 @@ import java.io.IOException
 import java.net._
 import java.nio.channels._
 import java.nio.channels.{Selector => NSelector}
+import java.util
+import java.util.Optional
 import java.util.concurrent._
 import java.util.concurrent.atomic._
 import java.util.function.Supplier
@@ -32,15 +34,16 @@ import kafka.network.RequestChannel.{CloseConnectionResponse, EndThrottlingRespo
 import kafka.network.Processor._
 import kafka.network.SocketServer._
 import kafka.security.CredentialProvider
-import kafka.server.KafkaConfig
+import kafka.server.{BrokerReconfigurable, KafkaConfig}
 import kafka.utils._
-import org.apache.kafka.common.{KafkaException, Reconfigurable}
-import org.apache.kafka.common.memory.{MemoryPool, SimpleMemoryPool}
+import org.apache.kafka.common.config.ConfigException
+import org.apache.kafka.common.{Endpoint, KafkaException, Reconfigurable}
+import org.apache.kafka.common.memory.{MemoryPool, RecyclingMemoryPool, SimpleMemoryPool}
 import org.apache.kafka.common.metrics._
-import org.apache.kafka.common.metrics.stats.Meter
-import org.apache.kafka.common.metrics.stats.Total
+import org.apache.kafka.common.metrics.stats.Percentiles.BucketSizing
+import org.apache.kafka.common.metrics.stats.{CumulativeSum, Max, Meter, Percentile, Percentiles}
 import org.apache.kafka.common.network.KafkaChannel.ChannelMuteEvent
-import org.apache.kafka.common.network.{ChannelBuilder, ChannelBuilders, KafkaChannel, ListenerName, Selectable, Send, Selector => KSelector}
+import org.apache.kafka.common.network.{ChannelBuilder, ChannelBuilders, KafkaChannel, ListenerName, ListenerReconfigurable, Selectable, Send, Selector => KSelector}
 import org.apache.kafka.common.protocol.ApiKeys
 import org.apache.kafka.common.requests.{RequestContext, RequestHeader}
 import org.apache.kafka.common.security.auth.SecurityProtocol
@@ -70,26 +73,38 @@ import scala.util.control.ControlThrowable
  *      Acceptor has 1 Processor thread that has its own selector and read requests from the socket.
  *      1 Handler thread that handles requests and produce responses back to the processor thread for writing.
  */
-class SocketServer(val config: KafkaConfig, val metrics: Metrics, val time: Time, val credentialProvider: CredentialProvider) extends Logging with KafkaMetricsGroup {
+class SocketServer(val config: KafkaConfig,
+                   val metrics: Metrics,
+                   val time: Time,
+                   val credentialProvider: CredentialProvider)
+  extends Logging with KafkaMetricsGroup with BrokerReconfigurable {
 
   private val maxQueuedRequests = config.queuedMaxRequests
 
   private val logContext = new LogContext(s"[SocketServer brokerId=${config.brokerId}] ")
   this.logIdent = logContext.logPrefix
 
-  private val memoryPoolSensor = metrics.sensor("MemoryPoolUtilization")
+  private val memoryPoolUsageSensor = metrics.sensor("MemoryPoolUtilization")
   private val memoryPoolDepletedPercentMetricName = metrics.metricName("MemoryPoolAvgDepletedPercent", MetricsGroup)
   private val memoryPoolDepletedTimeMetricName = metrics.metricName("MemoryPoolDepletedTimeTotal", MetricsGroup)
-  memoryPoolSensor.add(new Meter(TimeUnit.MILLISECONDS, memoryPoolDepletedPercentMetricName, memoryPoolDepletedTimeMetricName))
-  private val memoryPool = if (config.queuedMaxBytes > 0) new SimpleMemoryPool(config.queuedMaxBytes, config.socketRequestMaxBytes, false, memoryPoolSensor) else MemoryPool.NONE
+  memoryPoolUsageSensor.add(new Meter(TimeUnit.MILLISECONDS, memoryPoolDepletedPercentMetricName, memoryPoolDepletedTimeMetricName))
+  private val memoryPoolAllocationSensor = metrics.sensor("MemoryPoolAllocation")
+  private val memoryPoolMaxAllocateSizeMetricName = metrics.metricName("MemoryPoolMaxAllocateSize", MetricsGroup)
+  memoryPoolAllocationSensor.add(memoryPoolMaxAllocateSizeMetricName, new Max())
+  private val percentiles = (1 to 9).map( i => new Percentile(metrics.metricName("MemoryPoolAllocateSize%dPercentile".format(i * 10), MetricsGroup), i * 10))
+  // At current stage, we do not know the max decrypted request size, temporarily set it to 10MB.
+  memoryPoolAllocationSensor.add(new Percentiles(400, 0.0, 10485760, BucketSizing.CONSTANT, percentiles:_*))
+  private val memoryPool = if (config.queuedMaxBytes > 0) new SimpleMemoryPool(config.queuedMaxBytes, config.socketRequestMaxBytes, false, memoryPoolUsageSensor, memoryPoolAllocationSensor)
+                           else if (config.socketRequestCommonBytes > 0) new RecyclingMemoryPool(config.socketRequestCommonBytes, config.socketRequestBufferCacheSize, memoryPoolAllocationSensor)
+                           else MemoryPool.NONE
   // data-plane
   private val dataPlaneProcessors = new ConcurrentHashMap[Int, Processor]()
   private[network] val dataPlaneAcceptors = new ConcurrentHashMap[EndPoint, Acceptor]()
-  val dataPlaneRequestChannel = new RequestChannel(maxQueuedRequests, DataPlaneMetricPrefix)
+  val dataPlaneRequestChannel = new RequestChannel(maxQueuedRequests, DataPlaneMetricPrefix, time)
   // control-plane
   private var controlPlaneProcessorOpt : Option[Processor] = None
   private[network] var controlPlaneAcceptorOpt : Option[Acceptor] = None
-  val controlPlaneRequestChannelOpt: Option[RequestChannel] = config.controlPlaneListenerName.map(_ => new RequestChannel(20, ControlPlaneMetricPrefix))
+  val controlPlaneRequestChannelOpt: Option[RequestChannel] = config.controlPlaneListenerName.map(_ => new RequestChannel(20, ControlPlaneMetricPrefix, time))
 
   private var nextProcessorId = 0
   private var connectionQuotas: ConnectionQuotas = _
@@ -98,7 +113,8 @@ class SocketServer(val config: KafkaConfig, val metrics: Metrics, val time: Time
   /**
    * Start the socket server. Acceptors for all the listeners are started. Processors
    * are started if `startupProcessors` is true. If not, processors are only started when
-   * [[kafka.network.SocketServer#startProcessors()]] is invoked. Delayed starting of processors
+   * [[kafka.network.SocketServer#startDataPlaneProcessors()]] or
+   * [[kafka.network.SocketServer#startControlPlaneProcessor()]] is invoked. Delayed starting of processors
    * is used to delay processing client connections until server is fully initialized, e.g.
    * to ensure that all credentials have been loaded before authentications are performed.
    * Acceptors are always started during `startup` so that the bound port is known when this
@@ -107,14 +123,14 @@ class SocketServer(val config: KafkaConfig, val metrics: Metrics, val time: Time
    *
    * @param startupProcessors Flag indicating whether `Processor`s must be started.
    */
-  def startup(startupProcessors: Boolean = true) {
+  def startup(startupProcessors: Boolean = true): Unit = {
     this.synchronized {
-      connectionQuotas = new ConnectionQuotas(config.maxConnectionsPerIp, config.maxConnectionsPerIpOverrides)
+      connectionQuotas = new ConnectionQuotas(config, time)
       createControlPlaneAcceptorAndProcessor(config.controlPlaneListener)
       createDataPlaneAcceptorsAndProcessors(config.numNetworkThreads, config.dataPlaneListeners)
       if (startupProcessors) {
-        startControlPlaneProcessor()
-        startDataPlaneProcessors()
+        startControlPlaneProcessor(Map.empty)
+        startDataPlaneProcessors(Map.empty)
       }
     }
 
@@ -189,9 +205,25 @@ class SocketServer(val config: KafkaConfig, val metrics: Metrics, val time: Time
    * Starts processors of all the data-plane acceptors of this server if they have not already been started.
    * This method is used for delayed starting of data-plane processors if [[kafka.network.SocketServer#startup]]
    * was invoked with `startupProcessors=false`.
+   *
+   * Before starting processors for each endpoint, we ensure that authorizer has all the metadata
+   * to authorize requests on that endpoint by waiting on the provided future. We start inter-broker listener
+   * before other listeners. This allows authorization metadata for other listeners to be stored in Kafka topics
+   * in this cluster.
    */
-  def startDataPlaneProcessors(): Unit = synchronized {
-    dataPlaneAcceptors.values.asScala.foreach { _.startProcessors(DataPlaneThreadPrefix) }
+  def startDataPlaneProcessors(authorizerFutures: Map[Endpoint, CompletableFuture[Void]] = Map.empty): Unit = synchronized {
+    val interBrokerListener = dataPlaneAcceptors.asScala.keySet
+      .find(_.listenerName == config.interBrokerListenerName)
+      .getOrElse(throw new IllegalStateException(s"Inter-broker listener ${config.interBrokerListenerName} not found, endpoints=${dataPlaneAcceptors.keySet}"))
+    val orderedAcceptors = List(dataPlaneAcceptors.get(interBrokerListener)) ++
+      dataPlaneAcceptors.asScala.filterKeys(_ != interBrokerListener).values
+    orderedAcceptors.foreach { acceptor =>
+      val endpoint = acceptor.endPoint
+      debug(s"Wait for authorizer to complete start up on listener ${endpoint.listenerName}")
+      waitForAuthorizerFuture(acceptor, authorizerFutures)
+      debug(s"Start processors on listener ${endpoint.listenerName}")
+      acceptor.startProcessors(DataPlaneThreadPrefix)
+    }
     info(s"Started data-plane processors for ${dataPlaneAcceptors.size} acceptors")
   }
 
@@ -200,8 +232,9 @@ class SocketServer(val config: KafkaConfig, val metrics: Metrics, val time: Time
    * This method is used for delayed starting of control-plane processor if [[kafka.network.SocketServer#startup]]
    * was invoked with `startupProcessors=false`.
    */
-  def startControlPlaneProcessor(): Unit = synchronized {
+  def startControlPlaneProcessor(authorizerFutures: Map[Endpoint, CompletableFuture[Void]] = Map.empty): Unit = synchronized {
     controlPlaneAcceptorOpt.foreach { controlPlaneAcceptor =>
+      waitForAuthorizerFuture(controlPlaneAcceptor, authorizerFutures)
       controlPlaneAcceptor.startProcessors(ControlPlaneThreadPrefix)
       info(s"Started control-plane processor for the control-plane acceptor")
     }
@@ -212,6 +245,7 @@ class SocketServer(val config: KafkaConfig, val metrics: Metrics, val time: Time
   private def createDataPlaneAcceptorsAndProcessors(dataProcessorsPerListener: Int,
                                                     endpoints: Seq[EndPoint]): Unit = synchronized {
     endpoints.foreach { endpoint =>
+      connectionQuotas.addListener(config, endpoint.listenerName)
       val dataPlaneAcceptor = createAcceptor(endpoint, DataPlaneMetricPrefix)
       addDataPlaneProcessors(dataPlaneAcceptor, endpoint, dataProcessorsPerListener)
       KafkaThread.nonDaemon(s"data-plane-kafka-socket-acceptor-${endpoint.listenerName}-${endpoint.securityProtocol}-${endpoint.port}", dataPlaneAcceptor).start()
@@ -223,6 +257,7 @@ class SocketServer(val config: KafkaConfig, val metrics: Metrics, val time: Time
 
   private def createControlPlaneAcceptorAndProcessor(endpointOpt: Option[EndPoint]): Unit = synchronized {
     endpointOpt.foreach { endpoint =>
+      connectionQuotas.addListener(config, endpoint.listenerName)
       val controlPlaneAcceptor = createAcceptor(endpoint, ControlPlaneMetricPrefix)
       val controlPlaneProcessor = newProcessor(nextProcessorId, controlPlaneRequestChannelOpt.get, connectionQuotas, endpoint.listenerName, endpoint.securityProtocol, memoryPool)
       controlPlaneAcceptorOpt = Some(controlPlaneAcceptor)
@@ -324,18 +359,42 @@ class SocketServer(val config: KafkaConfig, val metrics: Metrics, val time: Time
   def removeListeners(listenersRemoved: Seq[EndPoint]): Unit = synchronized {
     info(s"Removing data-plane listeners for endpoints $listenersRemoved")
     listenersRemoved.foreach { endpoint =>
+      connectionQuotas.removeListener(config, endpoint.listenerName)
       dataPlaneAcceptors.asScala.remove(endpoint).foreach(_.shutdown())
     }
   }
 
-  def updateMaxConnectionsPerIp(maxConnectionsPerIp: Int): Unit = {
-    info(s"Updating maxConnectionsPerIp: $maxConnectionsPerIp")
-    connectionQuotas.updateMaxConnectionsPerIp(maxConnectionsPerIp)
+  override def reconfigurableConfigs: Set[String] = SocketServer.ReconfigurableConfigs
+
+  override def validateReconfiguration(newConfig: KafkaConfig): Unit = {
+
   }
 
-  def updateMaxConnectionsPerIpOverride(maxConnectionsPerIpOverrides: Map[String, Int]): Unit = {
-    info(s"Updating maxConnectionsPerIpOverrides: ${maxConnectionsPerIpOverrides.map { case (k, v) => s"$k=$v" }.mkString(",")}")
-    connectionQuotas.updateMaxConnectionsPerIpOverride(maxConnectionsPerIpOverrides)
+  override def reconfigure(oldConfig: KafkaConfig, newConfig: KafkaConfig): Unit = {
+    val maxConnectionsPerIp = newConfig.maxConnectionsPerIp
+    if (maxConnectionsPerIp != oldConfig.maxConnectionsPerIp) {
+      info(s"Updating maxConnectionsPerIp: $maxConnectionsPerIp")
+      connectionQuotas.updateMaxConnectionsPerIp(maxConnectionsPerIp)
+    }
+    val maxConnectionsPerIpOverrides = newConfig.maxConnectionsPerIpOverrides
+    if (maxConnectionsPerIpOverrides != oldConfig.maxConnectionsPerIpOverrides) {
+      info(s"Updating maxConnectionsPerIpOverrides: ${maxConnectionsPerIpOverrides.map { case (k, v) => s"$k=$v" }.mkString(",")}")
+      connectionQuotas.updateMaxConnectionsPerIpOverride(maxConnectionsPerIpOverrides)
+    }
+    val maxConnections = newConfig.maxConnections
+    if (maxConnections != oldConfig.maxConnections) {
+      info(s"Updating broker-wide maxConnections: $maxConnections")
+      connectionQuotas.updateBrokerMaxConnections(maxConnections)
+    }
+  }
+
+  private def waitForAuthorizerFuture(acceptor: Acceptor,
+                                      authorizerFutures: Map[Endpoint, CompletableFuture[Void]]): Unit = {
+    //we can't rely on authorizerFutures.get() due to ephemeral ports. Get the future using listener name
+    authorizerFutures.foreach { case (endpoint, future) =>
+      if (endpoint.listenerName == Optional.of(acceptor.endPoint.listenerName.value))
+        future.join()
+    }
   }
 
   // `protected` for test usage
@@ -373,6 +432,13 @@ object SocketServer {
   val ControlPlaneThreadPrefix = "control-plane"
   val DataPlaneMetricPrefix = ""
   val ControlPlaneMetricPrefix = "ControlPlane"
+
+  val ReconfigurableConfigs = Set(
+    KafkaConfig.MaxConnectionsPerIpProp,
+    KafkaConfig.MaxConnectionsPerIpOverridesProp,
+    KafkaConfig.MaxConnectionsProp)
+
+  val ListenerReconfigurableConfigs = Set(KafkaConfig.MaxConnectionsProp)
 }
 
 /**
@@ -427,10 +493,10 @@ private[kafka] abstract class AbstractServerThread(connectionQuotas: ConnectionQ
   /**
    * Close `channel` and decrement the connection count.
    */
-  def close(channel: SocketChannel): Unit = {
+  def close(listenerName: ListenerName, channel: SocketChannel): Unit = {
     if (channel != null) {
       debug(s"Closing connection from ${channel.socket.getRemoteSocketAddress()}")
-      connectionQuotas.dec(channel.socket.getInetAddress)
+      connectionQuotas.dec(listenerName, channel.socket.getInetAddress)
       CoreUtils.swallow(channel.socket().close(), this, Level.ERROR)
       CoreUtils.swallow(channel.close(), this, Level.ERROR)
     }
@@ -493,13 +559,14 @@ private[kafka] class Acceptor(val endPoint: EndPoint,
   /**
    * Accept loop that checks for new connection attempts
    */
-  def run() {
+  def run(): Unit = {
     serverChannel.register(nioSelector, SelectionKey.OP_ACCEPT)
     startupComplete()
     try {
       var currentProcessorIndex = 0
       while (isRunning) {
         try {
+
           val ready = nioSelector.select(500)
           if (ready > 0) {
             val keys = nioSelector.selectedKeys()
@@ -508,6 +575,7 @@ private[kafka] class Acceptor(val endPoint: EndPoint,
               try {
                 val key = iter.next
                 iter.remove()
+
                 if (key.isAcceptable) {
                   accept(key).foreach { socketChannel =>
 
@@ -567,7 +635,7 @@ private[kafka] class Acceptor(val endPoint: EndPoint,
 
     try {
       serverChannel.socket.bind(socketAddress)
-      info(s"Awaiting socket connections on s${socketAddress.getHostString}:${serverChannel.socket.getLocalPort}.")
+      info(s"Awaiting socket connections on ${socketAddress.getHostString}:${serverChannel.socket.getLocalPort}.")
     } catch {
       case e: SocketException =>
         throw new KafkaException(s"Socket server failed to bind to ${socketAddress.getHostString}:$port: ${e.getMessage}.", e)
@@ -582,7 +650,7 @@ private[kafka] class Acceptor(val endPoint: EndPoint,
     val serverSocketChannel = key.channel().asInstanceOf[ServerSocketChannel]
     val socketChannel = serverSocketChannel.accept()
     try {
-      connectionQuotas.inc(socketChannel.socket().getInetAddress)
+      connectionQuotas.inc(endPoint.listenerName, socketChannel.socket.getInetAddress, blockedPercentMeter)
       socketChannel.configureBlocking(false)
       socketChannel.socket().setTcpNoDelay(true)
       socketChannel.socket().setKeepAlive(true)
@@ -592,7 +660,7 @@ private[kafka] class Acceptor(val endPoint: EndPoint,
     } catch {
       case e: TooManyConnectionsException =>
         info(s"Rejected connection from ${e.ip}, address already has the configured maximum of ${e.count} connections.")
-        close(socketChannel)
+        close(endPoint.listenerName, socketChannel)
         None
     }
   }
@@ -680,7 +748,7 @@ private[kafka] class Processor(val id: Int,
     Map(NetworkProcessorMetricTag -> id.toString)
   )
 
-  val expiredConnectionsKilledCount = new Total()
+  val expiredConnectionsKilledCount = new CumulativeSum()
   private val expiredConnectionsKilledCountMetricName = metrics.metricName("expired-connections-killed-count", "socket-server-metrics", metricTags)
   metrics.addMetric(expiredConnectionsKilledCountMetricName, expiredConnectionsKilledCount)
 
@@ -718,7 +786,7 @@ private[kafka] class Processor(val id: Int,
   // closed, connection ids are not reused while requests from the closed connection are being processed.
   private var nextConnectionIndex = 0
 
-  override def run() {
+  override def run(): Unit = {
     startupComplete()
     try {
       while (isRunning) {
@@ -731,6 +799,7 @@ private[kafka] class Processor(val id: Int,
           processCompletedReceives()
           processCompletedSends()
           processDisconnected()
+          closeExcessConnections()
         } catch {
           // We catch all the throwables here to prevent the processor thread from exiting. We do this because
           // letting a processor exit might cause a bigger impact on the broker. This behavior might need to be
@@ -748,14 +817,14 @@ private[kafka] class Processor(val id: Int,
     }
   }
 
-  private def processException(errorMessage: String, throwable: Throwable) {
+  private def processException(errorMessage: String, throwable: Throwable): Unit = {
     throwable match {
       case e: ControlThrowable => throw e
       case e => error(errorMessage, e)
     }
   }
 
-  private def processChannelException(channelId: String, errorMessage: String, throwable: Throwable) {
+  private def processChannelException(channelId: String, errorMessage: String, throwable: Throwable): Unit = {
     if (openOrClosingChannel(channelId).isDefined) {
       error(s"Closing socket for $channelId because of error", throwable)
       close(channelId)
@@ -763,7 +832,7 @@ private[kafka] class Processor(val id: Int,
     processException(errorMessage, throwable)
   }
 
-  private def processNewResponses() {
+  private def processNewResponses(): Unit = {
     var currentResponse: RequestChannel.Response = null
     while ({currentResponse = dequeueResponse(); currentResponse != null}) {
       val channelId = currentResponse.request.context.connectionId
@@ -804,7 +873,7 @@ private[kafka] class Processor(val id: Int,
   }
 
   // `protected` for test usage
-  protected[network] def sendResponse(response: RequestChannel.Response, responseSend: Send) {
+  protected[network] def sendResponse(response: RequestChannel.Response, responseSend: Send): Unit = {
     val connectionId = response.request.context.connectionId
     trace(s"Socket server received response to send to $connectionId, registering for write and sending data: $response")
     // `channel` can be None if the connection was closed remotely or if selector closed it for being idle for too long
@@ -825,7 +894,7 @@ private[kafka] class Processor(val id: Int,
     override def get(): java.lang.Long = time.nanoseconds()
   }
 
-  private def poll() {
+  private def poll(): Unit = {
     val pollTimeout = if (newConnections.isEmpty) 300 else 0
     try selector.poll(pollTimeout)
     catch {
@@ -836,7 +905,7 @@ private[kafka] class Processor(val id: Int,
     }
   }
 
-  private def processCompletedReceives() {
+  private def processCompletedReceives(): Unit = {
     selector.completedReceives.asScala.foreach { receive =>
       try {
         openOrClosingChannel(receive.source) match {
@@ -847,8 +916,9 @@ private[kafka] class Processor(val id: Int,
             else {
               val nowNanos = time.nanoseconds()
               if (channel.serverAuthenticationSessionExpired(nowNanos)) {
-                channel.disconnect()
-                debug(s"Disconnected expired channel: $channel : $header")
+                // be sure to decrease connection count and drop any in-flight responses
+                debug(s"Disconnecting expired channel: $channel : $header")
+                close(channel.id)
                 expiredConnectionsKilledCount.record(null, 1, 0)
               } else {
                 val connectionId = receive.source
@@ -874,7 +944,7 @@ private[kafka] class Processor(val id: Int,
     }
   }
 
-  private def processCompletedSends() {
+  private def processCompletedSends(): Unit = {
     selector.completedSends.asScala.foreach { send =>
       try {
         val response = inflightResponses.remove(send.destination).getOrElse {
@@ -903,7 +973,7 @@ private[kafka] class Processor(val id: Int,
     request.updateRequestMetrics(networkThreadTimeNanos, response)
   }
 
-  private def processDisconnected() {
+  private def processDisconnected(): Unit = {
     selector.disconnected.keySet.asScala.foreach { connectionId =>
       try {
         val remoteHost = ConnectionId.fromString(connectionId).getOrElse {
@@ -911,13 +981,21 @@ private[kafka] class Processor(val id: Int,
         }.remoteHost
         inflightResponses.remove(connectionId).foreach(updateRequestMetrics)
         // the channel has been closed by the selector but the quotas still need to be updated
-        connectionQuotas.dec(InetAddress.getByName(remoteHost))
+        connectionQuotas.dec(listenerName, InetAddress.getByName(remoteHost))
       } catch {
         case e: Throwable => processException(s"Exception while processing disconnection of $connectionId", e)
       }
     }
   }
 
+  private def closeExcessConnections(): Unit = {
+    if (connectionQuotas.maxConnectionsExceeded(listenerName)) {
+      val channel = selector.lowestPriorityChannel()
+      if (channel != null)
+        close(channel.id)
+    }
+  }
+
   /**
    * Close the connection identified by `connectionId` and decrement the connection count.
    * The channel will be immediately removed from the selector's `channels` or `closingChannels`
@@ -930,7 +1008,7 @@ private[kafka] class Processor(val id: Int,
       debug(s"Closing selector connection $connectionId")
       val address = channel.socketAddress
       if (address != null)
-        connectionQuotas.dec(address)
+        connectionQuotas.dec(listenerName, address)
       selector.close(connectionId)
 
       inflightResponses.remove(connectionId).foreach(response => updateRequestMetrics(response))
@@ -964,7 +1042,7 @@ private[kafka] class Processor(val id: Int,
    * in each iteration is limited to ensure that traffic and connection close notifications of
    * existing channels are handled promptly.
    */
-  private def configureNewConnections() {
+  private def configureNewConnections(): Unit = {
     var connectionsProcessed = 0
     while (connectionsProcessed < connectionQueueSize && !newConnections.isEmpty) {
       val channel = newConnections.poll()
@@ -977,7 +1055,7 @@ private[kafka] class Processor(val id: Int,
         case e: Throwable =>
           val remoteAddress = channel.socket.getRemoteSocketAddress
           // need to close the channel here to avoid a socket leak.
-          close(channel)
+          close(listenerName, channel)
           processException(s"Processor $id closed connection from $remoteAddress", e)
       }
     }
@@ -986,7 +1064,7 @@ private[kafka] class Processor(val id: Int,
   /**
    * Close the selector and all open connections
    */
-  private def closeAll() {
+  private def closeAll(): Unit = {
     selector.channels.asScala.foreach { channel =>
       close(channel.id)
     }
@@ -1058,31 +1136,72 @@ private[kafka] class Processor(val id: Int,
 
 }
 
-class ConnectionQuotas(val defaultMax: Int, overrideQuotas: Map[String, Int]) {
+class ConnectionQuotas(config: KafkaConfig, time: Time) extends Logging {
 
-  @volatile private var defaultMaxConnectionsPerIp = defaultMax
-  @volatile private var maxConnectionsPerIpOverrides = overrideQuotas.map { case (host, count) => (InetAddress.getByName(host), count) }
+  @volatile private var defaultMaxConnectionsPerIp: Int = config.maxConnectionsPerIp
+  @volatile private var maxConnectionsPerIpOverrides = config.maxConnectionsPerIpOverrides.map { case (host, count) => (InetAddress.getByName(host), count) }
+  @volatile private var brokerMaxConnections = config.maxConnections
   private val counts = mutable.Map[InetAddress, Int]()
 
-  def inc(address: InetAddress) {
+  // Listener counts and configs are synchronized on `counts`
+  private val listenerCounts = mutable.Map[ListenerName, Int]()
+  private val maxConnectionsPerListener = mutable.Map[ListenerName, ListenerConnectionQuota]()
+  @volatile private var totalCount = 0
+
+  def inc(listenerName: ListenerName, address: InetAddress, acceptorBlockedPercentMeter: com.yammer.metrics.core.Meter): Unit = {
     counts.synchronized {
+      waitForConnectionSlot(listenerName, acceptorBlockedPercentMeter)
+
       val count = counts.getOrElseUpdate(address, 0)
       counts.put(address, count + 1)
+      totalCount += 1
+      if (listenerCounts.contains(listenerName)) {
+        listenerCounts.put(listenerName, listenerCounts(listenerName) + 1)
+      }
       val max = maxConnectionsPerIpOverrides.getOrElse(address, defaultMaxConnectionsPerIp)
       if (count >= max)
         throw new TooManyConnectionsException(address, max)
     }
   }
 
-  def updateMaxConnectionsPerIp(maxConnectionsPerIp: Int): Unit = {
+  private[network] def updateMaxConnectionsPerIp(maxConnectionsPerIp: Int): Unit = {
     defaultMaxConnectionsPerIp = maxConnectionsPerIp
   }
 
-  def updateMaxConnectionsPerIpOverride(overrideQuotas: Map[String, Int]): Unit = {
+  private[network] def updateMaxConnectionsPerIpOverride(overrideQuotas: Map[String, Int]): Unit = {
     maxConnectionsPerIpOverrides = overrideQuotas.map { case (host, count) => (InetAddress.getByName(host), count) }
   }
 
-  def dec(address: InetAddress) {
+  private[network] def updateBrokerMaxConnections(maxConnections: Int): Unit = {
+    counts.synchronized {
+      brokerMaxConnections = maxConnections
+      counts.notifyAll()
+    }
+  }
+
+  private[network] def addListener(config: KafkaConfig, listenerName: ListenerName): Unit = {
+    counts.synchronized {
+      if (!maxConnectionsPerListener.contains(listenerName)) {
+        val newListenerQuota = new ListenerConnectionQuota(counts, listenerName)
+        maxConnectionsPerListener.put(listenerName, newListenerQuota)
+        listenerCounts.put(listenerName, 0)
+        config.addReconfigurable(newListenerQuota)
+      }
+      counts.notifyAll()
+    }
+  }
+
+  private[network] def removeListener(config: KafkaConfig, listenerName: ListenerName): Unit = {
+    counts.synchronized {
+      maxConnectionsPerListener.remove(listenerName).foreach { listenerQuota =>
+        listenerCounts.remove(listenerName)
+        counts.notifyAll() // wake up any waiting acceptors to close cleanly
+        config.removeReconfigurable(listenerQuota)
+      }
+    }
+  }
+
+  def dec(listenerName: ListenerName, address: InetAddress): Unit = {
     counts.synchronized {
       val count = counts.getOrElse(address,
         throw new IllegalArgumentException(s"Attempted to decrease connection count for address with no connections, address: $address"))
@@ -1090,6 +1209,19 @@ class ConnectionQuotas(val defaultMax: Int, overrideQuotas: Map[String, Int]) {
         counts.remove(address)
       else
         counts.put(address, count - 1)
+
+      if (totalCount <= 0)
+        error(s"Attempted to decrease total connection count for broker with no connections")
+      totalCount -= 1
+
+      if (maxConnectionsPerListener.contains(listenerName)) {
+        val listenerCount = listenerCounts(listenerName)
+        if (listenerCount == 0)
+          error(s"Attempted to decrease connection count for listener $listenerName with no connections")
+        else
+          listenerCounts.put(listenerName, listenerCount - 1)
+      }
+      counts.notifyAll() // wake up any acceptors waiting to process a new connection since listener connection limit was reached
     }
   }
 
@@ -1097,6 +1229,72 @@ class ConnectionQuotas(val defaultMax: Int, overrideQuotas: Map[String, Int]) {
     counts.getOrElse(address, 0)
   }
 
+  private def waitForConnectionSlot(listenerName: ListenerName,
+                                    acceptorBlockedPercentMeter: com.yammer.metrics.core.Meter): Unit = {
+    counts.synchronized {
+      if (!connectionSlotAvailable(listenerName)) {
+        val startNs = time.nanoseconds
+        do {
+          counts.wait()
+        } while (!connectionSlotAvailable(listenerName))
+        acceptorBlockedPercentMeter.mark(time.nanoseconds - startNs)
+      }
+    }
+  }
+
+  // This is invoked in every poll iteration and we close one LRU connection in an iteration
+  // if necessary
+  def maxConnectionsExceeded(listenerName: ListenerName): Boolean = {
+    totalCount > brokerMaxConnections && !protectedListener(listenerName)
+  }
+
+  private def connectionSlotAvailable(listenerName: ListenerName): Boolean = {
+    if (listenerCounts(listenerName) >= maxListenerConnections(listenerName))
+      false
+    else if (protectedListener(listenerName))
+      true
+    else
+      totalCount < brokerMaxConnections
+  }
+
+  private def protectedListener(listenerName: ListenerName): Boolean =
+    config.interBrokerListenerName == listenerName && config.listeners.size > 1
+
+  private def maxListenerConnections(listenerName: ListenerName): Int =
+    maxConnectionsPerListener.get(listenerName).map(_.maxConnections).getOrElse(Int.MaxValue)
+
+  class ListenerConnectionQuota(lock: Object, listener: ListenerName) extends ListenerReconfigurable {
+    @volatile private var _maxConnections = Int.MaxValue
+
+    def maxConnections: Int = _maxConnections
+
+    override def listenerName(): ListenerName = listener
+
+    override def configure(configs: util.Map[String, _]): Unit = {
+      _maxConnections = maxConnections(configs)
+    }
+
+    override def reconfigurableConfigs(): util.Set[String] = {
+      SocketServer.ListenerReconfigurableConfigs.asJava
+    }
+
+    override def validateReconfiguration(configs: util.Map[String, _]): Unit = {
+      val value = maxConnections(configs)
+      if (value <= 0)
+        throw new ConfigException("Invalid max.connections $listenerMax")
+    }
+
+    override def reconfigure(configs: util.Map[String, _]): Unit = {
+      lock.synchronized {
+        _maxConnections = maxConnections(configs)
+        lock.notifyAll()
+      }
+    }
+
+    private def maxConnections(configs: util.Map[String, _]): Int = {
+      Option(configs.get(KafkaConfig.MaxConnectionsProp)).map(_.toString.toInt).getOrElse(Int.MaxValue)
+    }
+  }
 }
 
 class TooManyConnectionsException(val ip: InetAddress, val count: Int) extends KafkaException(s"Too many connections from $ip (maximum = $count)")
diff --git a/core/src/main/scala/kafka/security/CredentialProvider.scala b/core/src/main/scala/kafka/security/CredentialProvider.scala
index f761b8b281a8d..9aa8bc915d4a7 100644
--- a/core/src/main/scala/kafka/security/CredentialProvider.scala
+++ b/core/src/main/scala/kafka/security/CredentialProvider.scala
@@ -31,7 +31,7 @@ class CredentialProvider(scramMechanisms: Collection[String], val tokenCache: De
   val credentialCache = new CredentialCache
   ScramCredentialUtils.createCache(credentialCache, scramMechanisms)
 
-  def updateCredentials(username: String, config: Properties) {
+  def updateCredentials(username: String, config: Properties): Unit = {
     for (mechanism <- ScramMechanism.values()) {
       val cache = credentialCache.cache(mechanism.mechanismName, classOf[ScramCredential])
       if (cache != null) {
diff --git a/core/src/main/scala/kafka/security/auth/Authorizer.scala b/core/src/main/scala/kafka/security/auth/Authorizer.scala
index 9be8e6c704930..7509171313a53 100644
--- a/core/src/main/scala/kafka/security/auth/Authorizer.scala
+++ b/core/src/main/scala/kafka/security/auth/Authorizer.scala
@@ -32,6 +32,7 @@ import org.apache.kafka.common.security.auth.KafkaPrincipal
  * If `authorizer.class.name` has no value specified, then no authorization will be performed, and all operations are
  * permitted.
  */
+@deprecated("Use org.apache.kafka.server.authorizer.Authorizer", "Since 2.4")
 trait Authorizer extends Configurable {
 
   /**
diff --git a/core/src/main/scala/kafka/security/auth/SimpleAclAuthorizer.scala b/core/src/main/scala/kafka/security/auth/SimpleAclAuthorizer.scala
index 8a0b4a072e40a..56935d9b12e68 100644
--- a/core/src/main/scala/kafka/security/auth/SimpleAclAuthorizer.scala
+++ b/core/src/main/scala/kafka/security/auth/SimpleAclAuthorizer.scala
@@ -17,392 +17,160 @@
 package kafka.security.auth
 
 import java.util
-import java.util.concurrent.locks.ReentrantReadWriteLock
 
-import com.typesafe.scalalogging.Logger
-import kafka.api.KAFKA_2_0_IV1
 import kafka.network.RequestChannel.Session
-import kafka.security.auth.SimpleAclAuthorizer.{NoAcls, VersionedAcls}
-import kafka.server.KafkaConfig
-import kafka.utils.CoreUtils.{inReadLock, inWriteLock}
+import kafka.security.auth.SimpleAclAuthorizer.BaseAuthorizer
+import kafka.security.authorizer.{AclAuthorizer, AuthorizerUtils}
 import kafka.utils._
-import kafka.zk.{AclChangeNotificationHandler, AclChangeSubscription, KafkaZkClient, ZkAclChangeStore, ZkAclStore, ZkVersion}
-import org.apache.kafka.common.errors.UnsupportedVersionException
-import org.apache.kafka.common.resource.PatternType
+import kafka.zk.ZkVersion
+import org.apache.kafka.common.acl.{AccessControlEntryFilter, AclBinding, AclBindingFilter, AclOperation, AclPermissionType}
+import org.apache.kafka.common.errors.ApiException
+import org.apache.kafka.common.resource.{PatternType, ResourcePatternFilter}
 import org.apache.kafka.common.security.auth.KafkaPrincipal
-import org.apache.kafka.common.utils.{SecurityUtils, Time}
+import org.apache.kafka.server.authorizer.{Action, AuthorizableRequestContext, AuthorizationResult}
 
+import scala.collection.mutable
 import scala.collection.JavaConverters._
-import scala.util.{Failure, Random, Success, Try}
+import scala.compat.java8.OptionConverters._
 
+@deprecated("Use kafka.security.authorizer.AclAuthorizer", "Since 2.4")
 object SimpleAclAuthorizer {
   //optional override zookeeper cluster configuration where acls will be stored, if not specified acls will be stored in
   //same zookeeper where all other kafka broker info is stored.
-  val ZkUrlProp = "authorizer.zookeeper.url"
-  val ZkConnectionTimeOutProp = "authorizer.zookeeper.connection.timeout.ms"
-  val ZkSessionTimeOutProp = "authorizer.zookeeper.session.timeout.ms"
-  val ZkMaxInFlightRequests = "authorizer.zookeeper.max.in.flight.requests"
+  val ZkUrlProp = AclAuthorizer.ZkUrlProp
+  val ZkConnectionTimeOutProp = AclAuthorizer.ZkConnectionTimeOutProp
+  val ZkSessionTimeOutProp = AclAuthorizer.ZkSessionTimeOutProp
+  val ZkMaxInFlightRequests = AclAuthorizer.ZkMaxInFlightRequests
 
   //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 = AclAuthorizer.SuperUsersProp
   //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 = AclAuthorizer.AllowEveryoneIfNoAclIsFoundProp
 
   case class VersionedAcls(acls: Set[Acl], zkVersion: Int) {
     def exists: Boolean = zkVersion != ZkVersion.UnknownVersion
   }
   val NoAcls = VersionedAcls(Set.empty, ZkVersion.UnknownVersion)
+
+  private[auth] class BaseAuthorizer extends AclAuthorizer {
+    override def logAuditMessage(requestContext: AuthorizableRequestContext, action: Action, authorized: Boolean): Unit = {
+      val principal = requestContext.principal
+      val host = requestContext.clientAddress.getHostAddress
+      val operation = Operation.fromJava(action.operation)
+      val resource = AuthorizerUtils.convertToResource(action.resourcePattern)
+      def logMessage: String = {
+        val authResult = if (authorized) "Allowed" else "Denied"
+        s"Principal = $principal is $authResult Operation = $operation from host = $host on resource = $resource"
+      }
+
+      if (authorized) authorizerLogger.debug(logMessage)
+      else authorizerLogger.info(logMessage)
+    }
+  }
 }
 
+@deprecated("Use kafka.security.authorizer.AclAuthorizer", "Since 2.4")
 class SimpleAclAuthorizer extends Authorizer with Logging {
-  private val authorizerLogger = Logger("kafka.authorizer.logger")
-  private var superUsers = Set.empty[KafkaPrincipal]
-  private var shouldAllowEveryoneIfNoAclIsFound = false
-  private var zkClient: KafkaZkClient = _
-  private var aclChangeListeners: Iterable[AclChangeSubscription] = Iterable.empty
-  private var extendedAclSupport: Boolean = _
 
-  @volatile
-  private var aclCache = new scala.collection.immutable.TreeMap[Resource, VersionedAcls]()(ResourceOrdering)
-  private val lock = new ReentrantReadWriteLock()
+  private val aclAuthorizer = new BaseAuthorizer
 
   // The maximum number of times we should try to update the resource acls in zookeeper before failing;
   // This should never occur, but is a safeguard just in case.
   protected[auth] var maxUpdateRetries = 10
 
-  private val retryBackoffMs = 100
-  private val retryBackoffJitterMs = 50
 
   /**
    * Guaranteed to be called before any authorize call is made.
    */
-  override def configure(javaConfigs: util.Map[String, _]) {
-    val configs = javaConfigs.asScala
-    val props = new java.util.Properties()
-    configs.foreach { case (key, value) => props.put(key, value.toString) }
-
-    superUsers = configs.get(SimpleAclAuthorizer.SuperUsersProp).collect {
-      case str: String if str.nonEmpty => str.split(";").map(s => SecurityUtils.parseKafkaPrincipal(s.trim)).toSet
-    }.getOrElse(Set.empty[KafkaPrincipal])
-
-    shouldAllowEveryoneIfNoAclIsFound = configs.get(SimpleAclAuthorizer.AllowEveryoneIfNoAclIsFoundProp).exists(_.toString.toBoolean)
-
-    // Use `KafkaConfig` in order to get the default ZK config values if not present in `javaConfigs`. Note that this
-    // means that `KafkaConfig.zkConnect` must always be set by the user (even if `SimpleAclAuthorizer.ZkUrlProp` is also
-    // set).
-    val kafkaConfig = KafkaConfig.fromProps(props, doLog = false)
-    val zkUrl = configs.get(SimpleAclAuthorizer.ZkUrlProp).map(_.toString).getOrElse(kafkaConfig.zkConnect)
-    val zkConnectionTimeoutMs = configs.get(SimpleAclAuthorizer.ZkConnectionTimeOutProp).map(_.toString.toInt).getOrElse(kafkaConfig.zkConnectionTimeoutMs)
-    val zkSessionTimeOutMs = configs.get(SimpleAclAuthorizer.ZkSessionTimeOutProp).map(_.toString.toInt).getOrElse(kafkaConfig.zkSessionTimeoutMs)
-    val zkMaxInFlightRequests = configs.get(SimpleAclAuthorizer.ZkMaxInFlightRequests).map(_.toString.toInt).getOrElse(kafkaConfig.zkMaxInFlightRequests)
-
-    val time = Time.SYSTEM
-    zkClient = KafkaZkClient(zkUrl, kafkaConfig.zkEnableSecureAcls, zkSessionTimeOutMs, zkConnectionTimeoutMs,
-      zkMaxInFlightRequests, time, "kafka.security", "SimpleAclAuthorizer")
-    zkClient.createAclPaths()
-
-    extendedAclSupport = kafkaConfig.interBrokerProtocolVersion >= KAFKA_2_0_IV1
-
-    // Start change listeners first and then populate the cache so that there is no timing window
-    // between loading cache and processing change notifications.
-    startZkChangeListeners()
-    loadCache()
+  override def configure(javaConfigs: util.Map[String, _]): Unit = {
+    aclAuthorizer.configure(javaConfigs)
   }
 
   override def authorize(session: Session, operation: Operation, resource: Resource): Boolean = {
-    if (resource.patternType != PatternType.LITERAL) {
-      throw new IllegalArgumentException("Only literal resources are supported. Got: " + resource.patternType)
-    }
-
-    // ensure we compare identical classes
-    val sessionPrincipal = session.principal
-    val principal = if (classOf[KafkaPrincipal] != sessionPrincipal.getClass)
-      new KafkaPrincipal(sessionPrincipal.getPrincipalType, sessionPrincipal.getName)
-    else
-      sessionPrincipal
-
-    val host = session.clientAddress.getHostAddress
-
-    def isEmptyAclAndAuthorized(acls: Set[Acl]): Boolean = {
-      if (acls.isEmpty) {
-        // No ACLs found for this resource, permission is determined by value of config allow.everyone.if.no.acl.found
-        authorizerLogger.debug(s"No acl found for resource $resource, authorized = $shouldAllowEveryoneIfNoAclIsFound")
-        shouldAllowEveryoneIfNoAclIsFound
-      } else false
-    }
-
-    def denyAclExists(acls: Set[Acl]): Boolean = {
-      // Check if there are any Deny ACLs which would forbid this operation.
-      aclMatch(operation, resource, principal, host, Deny, acls)
-    }
-
-    def allowAclExists(acls: Set[Acl]): Boolean = {
-      // Check if there are any Allow ACLs which would allow this operation.
-      // Allowing read, write, delete, or alter implies allowing describe.
-      // See #{org.apache.kafka.common.acl.AclOperation} for more details about ACL inheritance.
-      val allowOps = operation match {
-        case Describe => Set[Operation](Describe, Read, Write, Delete, Alter)
-        case DescribeConfigs => Set[Operation](DescribeConfigs, AlterConfigs)
-        case _ => Set[Operation](operation)
-      }
-      allowOps.exists(operation => aclMatch(operation, resource, principal, host, Allow, acls))
-    }
-
-    def aclsAllowAccess = {
-      //we allow an operation if no acls are found and user has configured to allow all users
-      //when no acls are found or if no deny acls are found and at least one allow acls matches.
-      val acls = getMatchingAcls(resource.resourceType, resource.name)
-      isEmptyAclAndAuthorized(acls) || (!denyAclExists(acls) && allowAclExists(acls))
-    }
-
-    // Evaluate if operation is allowed
-    val authorized = isSuperUser(operation, resource, principal, host) || aclsAllowAccess
-
-    logAuditMessage(principal, authorized, operation, resource, host)
-    authorized
+    val requestContext = AuthorizerUtils.sessionToRequestContext(session)
+    val action = new Action(operation.toJava, resource.toPattern, 1, true, true)
+    aclAuthorizer.authorize(requestContext, List(action).asJava).asScala.head == AuthorizationResult.ALLOWED
   }
 
   def isSuperUser(operation: Operation, resource: Resource, principal: KafkaPrincipal, host: String): Boolean = {
-    if (superUsers.contains(principal)) {
-      authorizerLogger.debug(s"principal = $principal is a super user, allowing operation without checking acls.")
-      true
-    } else false
-  }
-
-  private def aclMatch(operation: Operation, resource: Resource, principal: KafkaPrincipal, host: String, permissionType: PermissionType, acls: Set[Acl]): Boolean = {
-    acls.find { acl =>
-      acl.permissionType == permissionType &&
-        (acl.principal == principal || acl.principal == Acl.WildCardPrincipal) &&
-        (operation == acl.operation || acl.operation == All) &&
-        (acl.host == host || acl.host == Acl.WildCardHost)
-    }.exists { acl =>
-      authorizerLogger.debug(s"operation = $operation on resource = $resource from host = $host is $permissionType based on acl = $acl")
-      true
-    }
+    aclAuthorizer.isSuperUser(principal)
   }
 
-  override def addAcls(acls: Set[Acl], resource: Resource) {
+  override def addAcls(acls: Set[Acl], resource: Resource): Unit = {
+    aclAuthorizer.maxUpdateRetries = maxUpdateRetries
     if (acls != null && acls.nonEmpty) {
-      if (!extendedAclSupport && resource.patternType == PatternType.PREFIXED) {
-        throw new UnsupportedVersionException(s"Adding ACLs on prefixed resource patterns requires " +
-          s"${KafkaConfig.InterBrokerProtocolVersionProp} of $KAFKA_2_0_IV1 or greater")
-      }
-
-      inWriteLock(lock) {
-        updateResourceAcls(resource) { currentAcls =>
-          currentAcls ++ acls
-        }
-      }
+      val bindings = acls.map { acl => AuthorizerUtils.convertToAclBinding(resource, acl) }
+      createAcls(bindings)
     }
   }
 
   override def removeAcls(aclsTobeRemoved: Set[Acl], resource: Resource): Boolean = {
-    inWriteLock(lock) {
-      updateResourceAcls(resource) { currentAcls =>
-        currentAcls -- aclsTobeRemoved
-      }
+    val filters = aclsTobeRemoved.map { acl =>
+      new AclBindingFilter(resource.toPattern.toFilter, AuthorizerUtils.convertToAccessControlEntry(acl).toFilter)
     }
+    deleteAcls(filters)
   }
 
   override def removeAcls(resource: Resource): Boolean = {
-    inWriteLock(lock) {
-      val result = zkClient.deleteResource(resource)
-      updateCache(resource, NoAcls)
-      updateAclChangedFlag(resource)
-      result
-    }
+    val filter = new AclBindingFilter(resource.toPattern.toFilter, AccessControlEntryFilter.ANY)
+    deleteAcls(Set(filter))
   }
 
   override def getAcls(resource: Resource): Set[Acl] = {
-    inReadLock(lock) {
-      aclCache.get(resource).map(_.acls).getOrElse(Set.empty[Acl])
-    }
+    val filter = new AclBindingFilter(resource.toPattern.toFilter, AccessControlEntryFilter.ANY)
+    acls(filter).getOrElse(resource, Set.empty)
   }
 
   override def getAcls(principal: KafkaPrincipal): Map[Resource, Set[Acl]] = {
-    inReadLock(lock) {
-      aclCache.mapValues { versionedAcls =>
-        versionedAcls.acls.filter(_.principal == principal)
-      }.filter { case (_, acls) =>
-        acls.nonEmpty
-      }
-    }
+    val filter = new AclBindingFilter(ResourcePatternFilter.ANY,
+      new AccessControlEntryFilter(principal.toString, null, AclOperation.ANY, AclPermissionType.ANY))
+    acls(filter)
   }
 
   def getMatchingAcls(resourceType: ResourceType, resourceName: String): Set[Acl] = {
-    inReadLock(lock) {
-      val wildcard = aclCache.get(Resource(resourceType, Acl.WildCardResource, PatternType.LITERAL))
-        .map(_.acls)
-        .getOrElse(Set.empty[Acl])
-
-      val literal = aclCache.get(Resource(resourceType, resourceName, PatternType.LITERAL))
-        .map(_.acls)
-        .getOrElse(Set.empty[Acl])
-
-      val prefixed = aclCache
-        .from(Resource(resourceType, resourceName, PatternType.PREFIXED))
-        .to(Resource(resourceType, resourceName.take(1), PatternType.PREFIXED))
-        .filterKeys(resource => resourceName.startsWith(resource.name))
-        .flatMap { case (resource, versionedAcls) => versionedAcls.acls }
-        .toSet
-
-      prefixed ++ wildcard ++ literal
-    }
+    val filter = new AclBindingFilter(new ResourcePatternFilter(resourceType.toJava, resourceName, PatternType.MATCH),
+      AccessControlEntryFilter.ANY)
+    acls(filter).flatMap(_._2).toSet
   }
 
   override def getAcls(): Map[Resource, Set[Acl]] = {
-    inReadLock(lock) {
-      aclCache.mapValues(_.acls)
-    }
-  }
-
-  def close() {
-    aclChangeListeners.foreach(listener => listener.close())
-    if (zkClient != null) zkClient.close()
-  }
-
-  private def loadCache() {
-    inWriteLock(lock) {
-      ZkAclStore.stores.foreach(store => {
-        val resourceTypes = zkClient.getResourceTypes(store.patternType)
-        for (rType <- resourceTypes) {
-          val resourceType = Try(ResourceType.fromString(rType))
-          resourceType match {
-            case Success(resourceTypeObj) => {
-              val resourceNames = zkClient.getResourceNames(store.patternType, resourceTypeObj)
-              for (resourceName <- resourceNames) {
-                val resource = new Resource(resourceTypeObj, resourceName, store.patternType)
-                val versionedAcls = getAclsFromZk(resource)
-                updateCache(resource, versionedAcls)
-              }
-            }
-            case Failure(f) => warn(s"Ignoring unknown ResourceType: $rType")
-          }
-        }
-      })
-    }
-  }
-
-  private[auth] def startZkChangeListeners(): Unit = {
-    aclChangeListeners = ZkAclChangeStore.stores
-      .map(store => store.createListener(AclChangedNotificationHandler, zkClient))
-  }
-
-  private def logAuditMessage(principal: KafkaPrincipal, authorized: Boolean, operation: Operation, resource: Resource, host: String) {
-    def logMessage: String = {
-      val authResult = if (authorized) "Allowed" else "Denied"
-      s"Principal = $principal is $authResult Operation = $operation from host = $host on resource = $resource"
-    }
-
-    if (authorized) authorizerLogger.debug(logMessage)
-    else authorizerLogger.info(logMessage)
-  }
-
-  /**
-    * Safely updates the resources ACLs by ensuring reads and writes respect the expected zookeeper version.
-    * Continues to retry until it successfully updates zookeeper.
-    *
-    * Returns a boolean indicating if the content of the ACLs was actually changed.
-    *
-    * @param resource the resource to change ACLs for
-    * @param getNewAcls function to transform existing acls to new ACLs
-    * @return boolean indicating if a change was made
-    */
-  private def updateResourceAcls(resource: Resource)(getNewAcls: Set[Acl] => Set[Acl]): Boolean = {
-    var currentVersionedAcls =
-      if (aclCache.contains(resource))
-        getAclsFromCache(resource)
-      else
-        getAclsFromZk(resource)
-    var newVersionedAcls: VersionedAcls = null
-    var writeComplete = false
-    var retries = 0
-    while (!writeComplete && retries <= maxUpdateRetries) {
-      val newAcls = getNewAcls(currentVersionedAcls.acls)
-      val (updateSucceeded, updateVersion) =
-        if (newAcls.nonEmpty) {
-          if (currentVersionedAcls.exists)
-            zkClient.conditionalSetAclsForResource(resource, newAcls, currentVersionedAcls.zkVersion)
-          else
-            zkClient.createAclsForResourceIfNotExists(resource, newAcls)
-        } else {
-          trace(s"Deleting path for $resource because it had no ACLs remaining")
-          (zkClient.conditionalDelete(resource, currentVersionedAcls.zkVersion), 0)
-        }
-
-      if (!updateSucceeded) {
-        trace(s"Failed to update ACLs for $resource. Used version ${currentVersionedAcls.zkVersion}. Reading data and retrying update.")
-        Thread.sleep(backoffTime)
-        currentVersionedAcls = getAclsFromZk(resource)
-        retries += 1
-      } else {
-        newVersionedAcls = VersionedAcls(newAcls, updateVersion)
-        writeComplete = updateSucceeded
-      }
-    }
-
-    if(!writeComplete)
-      throw new IllegalStateException(s"Failed to update ACLs for $resource after trying a maximum of $maxUpdateRetries times")
-
-    if (newVersionedAcls.acls != currentVersionedAcls.acls) {
-      debug(s"Updated ACLs for $resource to ${newVersionedAcls.acls} with version ${newVersionedAcls.zkVersion}")
-      updateCache(resource, newVersionedAcls)
-      updateAclChangedFlag(resource)
-      true
-    } else {
-      debug(s"Updated ACLs for $resource, no change was made")
-      updateCache(resource, newVersionedAcls) // Even if no change, update the version
-      false
-    }
-  }
-
-  private def getAclsFromCache(resource: Resource): VersionedAcls = {
-    aclCache.getOrElse(resource, throw new IllegalArgumentException(s"ACLs do not exist in the cache for resource $resource"))
+    acls(AclBindingFilter.ANY)
   }
 
-  private def getAclsFromZk(resource: Resource): VersionedAcls = {
-    zkClient.getVersionedAclsForResource(resource)
+  def close(): Unit = {
+    aclAuthorizer.close()
   }
 
-  private def updateCache(resource: Resource, versionedAcls: VersionedAcls) {
-    if (versionedAcls.acls.nonEmpty) {
-      aclCache = aclCache + (resource -> versionedAcls)
-    } else {
-      aclCache = aclCache - resource
-    }
+  private def createAcls(bindings: Set[AclBinding]): Unit = {
+    aclAuthorizer.maxUpdateRetries = maxUpdateRetries
+    val results = aclAuthorizer.createAcls(null, bindings.toList.asJava).asScala.map(_.toCompletableFuture.get)
+    results.foreach { result => result.exception.asScala.foreach(throwException) }
   }
 
-  private def updateAclChangedFlag(resource: Resource) {
-      zkClient.createAclChangeNotification(resource)
+  private def deleteAcls(filters: Set[AclBindingFilter]): Boolean = {
+    aclAuthorizer.maxUpdateRetries = maxUpdateRetries
+    val results = aclAuthorizer.deleteAcls(null, filters.toList.asJava).asScala.map(_.toCompletableFuture.get)
+    results.foreach { result => result.exception.asScala.foreach(throwException) }
+    results.flatMap(_.aclBindingDeleteResults.asScala).foreach { result => result.exception.asScala.foreach(e => throw e) }
+    results.exists(r => r.aclBindingDeleteResults.asScala.exists(d => !d.exception.isPresent))
   }
 
-  private def backoffTime = {
-    retryBackoffMs + Random.nextInt(retryBackoffJitterMs)
-  }
-
-  object AclChangedNotificationHandler extends AclChangeNotificationHandler {
-    override def processNotification(resource: Resource) {
-      inWriteLock(lock) {
-        val versionedAcls = getAclsFromZk(resource)
-        updateCache(resource, versionedAcls)
-      }
+  private def acls(filter: AclBindingFilter): Map[Resource, Set[Acl]] = {
+    val result = mutable.Map[Resource, mutable.Set[Acl]]()
+    aclAuthorizer.acls(filter).asScala.foreach { binding =>
+      val resource = AuthorizerUtils.convertToResource(binding.pattern)
+      val acl = AuthorizerUtils.convertToAcl(binding.entry)
+      result.getOrElseUpdate(resource, mutable.Set()).add(acl)
     }
+    result.mapValues(_.toSet).toMap
   }
 
-  // Orders by resource type, then resource pattern type and finally reverse ordering by name.
-  private object ResourceOrdering extends Ordering[Resource] {
-
-    def compare(a: Resource, b: Resource): Int = {
-      val rt = a.resourceType compare b.resourceType
-      if (rt != 0)
-        rt
-      else {
-        val rnt = a.patternType compareTo b.patternType
-        if (rnt != 0)
-          rnt
-        else
-          (a.name compare b.name) * -1
-      }
-    }
+  // To retain the same exceptions as in previous versions, throw the underlying exception when the exception
+  // was wrapped by AclAuthorizer in an ApiException
+  private def throwException(e: ApiException): Unit = {
+    if (e.getCause != null)
+      throw e.getCause
+    else
+      throw e
   }
 }
diff --git a/core/src/main/scala/kafka/security/authorizer/AclAuthorizer.scala b/core/src/main/scala/kafka/security/authorizer/AclAuthorizer.scala
new file mode 100644
index 0000000000000..8ba060d814921
--- /dev/null
+++ b/core/src/main/scala/kafka/security/authorizer/AclAuthorizer.scala
@@ -0,0 +1,499 @@
+/**
+ * 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 java.{lang, util}
+import java.util.concurrent.{CompletableFuture, CompletionStage}
+import java.util.concurrent.locks.ReentrantReadWriteLock
+
+import com.typesafe.scalalogging.Logger
+import kafka.api.KAFKA_2_0_IV1
+import kafka.security.authorizer.AclAuthorizer.VersionedAcls
+import kafka.security.auth.{Acl, Operation, PermissionType, Resource, ResourceType}
+import kafka.security.auth.{All, Allow, Alter, AlterConfigs, Delete, Deny, Describe, DescribeConfigs, Read, Write}
+import kafka.server.KafkaConfig
+import kafka.utils.CoreUtils.{inReadLock, inWriteLock}
+import kafka.utils._
+import kafka.zk._
+import org.apache.kafka.common.Endpoint
+import org.apache.kafka.common.acl._
+import org.apache.kafka.common.errors.{ApiException, InvalidRequestException, UnsupportedVersionException}
+import org.apache.kafka.common.protocol.ApiKeys
+import org.apache.kafka.common.resource.PatternType
+import org.apache.kafka.common.security.auth.KafkaPrincipal
+import org.apache.kafka.common.utils.{Time, SecurityUtils => JSecurityUtils}
+import org.apache.kafka.server.authorizer.AclDeleteResult.AclBindingDeleteResult
+import org.apache.kafka.server.authorizer.{AclCreateResult, AclDeleteResult, Action, AuthorizableRequestContext, AuthorizationResult, Authorizer, AuthorizerServerInfo}
+
+import scala.collection.mutable
+import scala.collection.JavaConverters._
+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 ZkUrlProp = "authorizer.zookeeper.url"
+  val ZkConnectionTimeOutProp = "authorizer.zookeeper.connection.timeout.ms"
+  val ZkSessionTimeOutProp = "authorizer.zookeeper.session.timeout.ms"
+  val ZkMaxInFlightRequests = "authorizer.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"
+  // 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"
+
+  case class VersionedAcls(acls: Set[Acl], zkVersion: Int) {
+    def exists: Boolean = zkVersion != ZkVersion.UnknownVersion
+  }
+  val NoAcls = VersionedAcls(Set.empty, ZkVersion.UnknownVersion)
+
+  // Orders by resource type, then resource pattern type and finally reverse ordering by name.
+  private object ResourceOrdering extends Ordering[Resource] {
+
+    def compare(a: Resource, b: Resource): Int = {
+      val rt = a.resourceType compare b.resourceType
+      if (rt != 0)
+        rt
+      else {
+        val rnt = a.patternType compareTo b.patternType
+        if (rnt != 0)
+          rnt
+        else
+          (a.name compare b.name) * -1
+      }
+    }
+  }
+}
+
+class AclAuthorizer extends Authorizer with Logging {
+  private[security] val authorizerLogger = Logger("kafka.authorizer.logger")
+  private var superUsers = Set.empty[KafkaPrincipal]
+  private var shouldAllowEveryoneIfNoAclIsFound = false
+  private var zkClient: KafkaZkClient = _
+  private var aclChangeListeners: Iterable[AclChangeSubscription] = Iterable.empty
+  private var extendedAclSupport: Boolean = _
+
+  @volatile
+  private var aclCache = new scala.collection.immutable.TreeMap[Resource, VersionedAcls]()(AclAuthorizer.ResourceOrdering)
+  private val lock = new ReentrantReadWriteLock()
+
+  // The maximum number of times we should try to update the resource acls in zookeeper before failing;
+  // This should never occur, but is a safeguard just in case.
+  protected[security] var maxUpdateRetries = 10
+
+  private val retryBackoffMs = 100
+  private val retryBackoffJitterMs = 50
+
+  /**
+   * Guaranteed to be called before any authorize call is made.
+   */
+  override def configure(javaConfigs: util.Map[String, _]): Unit = {
+    val configs = javaConfigs.asScala
+    val props = new java.util.Properties()
+    configs.foreach { case (key, value) => props.put(key, value.toString) }
+
+    superUsers = configs.get(AclAuthorizer.SuperUsersProp).collect {
+      case str: String if str.nonEmpty => str.split(";").map(s => JSecurityUtils.parseKafkaPrincipal(s.trim)).toSet
+    }.getOrElse(Set.empty[KafkaPrincipal])
+
+    shouldAllowEveryoneIfNoAclIsFound = configs.get(AclAuthorizer.AllowEveryoneIfNoAclIsFoundProp).exists(_.toString.toBoolean)
+
+    // Use `KafkaConfig` in order to get the default ZK config values if not present in `javaConfigs`. Note that this
+    // means that `KafkaConfig.zkConnect` must always be set by the user (even if `AclAuthorizer.ZkUrlProp` is also
+    // set).
+    val kafkaConfig = KafkaConfig.fromProps(props, doLog = false)
+    val zkUrl = configs.get(AclAuthorizer.ZkUrlProp).map(_.toString).getOrElse(kafkaConfig.zkConnect)
+    val zkConnectionTimeoutMs = configs.get(AclAuthorizer.ZkConnectionTimeOutProp).map(_.toString.toInt).getOrElse(kafkaConfig.zkConnectionTimeoutMs)
+    val zkSessionTimeOutMs = configs.get(AclAuthorizer.ZkSessionTimeOutProp).map(_.toString.toInt).getOrElse(kafkaConfig.zkSessionTimeoutMs)
+    val zkMaxInFlightRequests = configs.get(AclAuthorizer.ZkMaxInFlightRequests).map(_.toString.toInt).getOrElse(kafkaConfig.zkMaxInFlightRequests)
+
+    val time = Time.SYSTEM
+    zkClient = KafkaZkClient(zkUrl, kafkaConfig.zkEnableSecureAcls, zkSessionTimeOutMs, zkConnectionTimeoutMs,
+      zkMaxInFlightRequests, time, "kafka.security", "AclAuthorizer", name=Some("ACL authorizer"))
+    zkClient.createAclPaths()
+
+    extendedAclSupport = kafkaConfig.interBrokerProtocolVersion >= KAFKA_2_0_IV1
+
+    // Start change listeners first and then populate the cache so that there is no timing window
+    // between loading cache and processing change notifications.
+    startZkChangeListeners()
+    loadCache()
+  }
+
+  override def start(serverInfo: AuthorizerServerInfo): util.Map[Endpoint, _ <: CompletionStage[Void]] = {
+    serverInfo.endpoints.asScala.map { endpoint =>
+      endpoint -> CompletableFuture.completedFuture[Void](null) }.toMap.asJava
+  }
+
+  override def authorize(requestContext: AuthorizableRequestContext, actions: util.List[Action]): util.List[AuthorizationResult] = {
+    actions.asScala.map { action => authorizeAction(requestContext, action) }.asJava
+  }
+
+  override def createAcls(requestContext: AuthorizableRequestContext,
+                          aclBindings: util.List[AclBinding]): util.List[_ <: CompletionStage[AclCreateResult]] = {
+    val results = new Array[AclCreateResult](aclBindings.size)
+    val aclsToCreate = aclBindings.asScala.zipWithIndex
+      .filter { case (aclBinding, i) =>
+        try {
+          if (!extendedAclSupport && aclBinding.pattern.patternType == PatternType.PREFIXED) {
+            throw new UnsupportedVersionException(s"Adding ACLs on prefixed resource patterns requires " +
+              s"${KafkaConfig.InterBrokerProtocolVersionProp} of $KAFKA_2_0_IV1 or greater")
+          }
+          AuthorizerUtils.validateAclBinding(aclBinding)
+          true
+        } catch {
+          case e: Throwable =>
+            results(i) = new AclCreateResult(new InvalidRequestException("Failed to create ACL", apiException(e)))
+            false
+        }
+      }.groupBy(_._1.pattern)
+
+    if (aclsToCreate.nonEmpty) {
+      inWriteLock(lock) {
+        aclsToCreate.foreach { case (resource, aclsWithIndex) =>
+          try {
+            updateResourceAcls(AuthorizerUtils.convertToResource(resource)) { currentAcls =>
+              val newAcls = aclsWithIndex.map { case (acl, index) => AuthorizerUtils.convertToAcl(acl.entry) }
+              currentAcls ++ newAcls
+            }
+            aclsWithIndex.foreach { case (_, index) => results(index) = AclCreateResult.SUCCESS }
+          } catch {
+            case e: Throwable =>
+              aclsWithIndex.foreach { case (_, index) => results(index) = new AclCreateResult(apiException(e)) }
+          }
+        }
+      }
+    }
+    results.toList.map(CompletableFuture.completedFuture[AclCreateResult]).asJava
+  }
+
+  override def deleteAcls(requestContext: AuthorizableRequestContext,
+                          aclBindingFilters: util.List[AclBindingFilter]): util.List[_ <: CompletionStage[AclDeleteResult]] = {
+    val deletedBindings = new mutable.HashMap[AclBinding, Int]()
+    val deleteExceptions = new mutable.HashMap[AclBinding, ApiException]()
+    val filters = aclBindingFilters.asScala.zipWithIndex
+    inWriteLock(lock) {
+      val resourcesToUpdate = aclCache.keys.map { resource =>
+        val matchingFilters = filters.filter { case (filter, _) =>
+          filter.patternFilter.matches(resource.toPattern)
+        }
+        resource -> matchingFilters
+      }.toMap.filter(_._2.nonEmpty)
+
+      resourcesToUpdate.foreach { case (resource, matchingFilters) =>
+        val resourceBindingsBeingDeleted = new mutable.HashMap[AclBinding, Int]()
+        try {
+          updateResourceAcls(resource) { currentAcls =>
+            val aclsToRemove = currentAcls.filter { acl =>
+              matchingFilters.exists { case (filter, index) =>
+                val matches = filter.entryFilter.matches(AuthorizerUtils.convertToAccessControlEntry(acl))
+                if (matches) {
+                  val binding = AuthorizerUtils.convertToAclBinding(resource, acl)
+                  deletedBindings.getOrElseUpdate(binding, index)
+                  resourceBindingsBeingDeleted.getOrElseUpdate(binding, index)
+                }
+                matches
+              }
+            }
+            currentAcls -- aclsToRemove
+          }
+        } catch {
+          case e: Exception =>
+            resourceBindingsBeingDeleted.foreach { case (binding, index) =>
+                deleteExceptions.getOrElseUpdate(binding, apiException(e))
+            }
+        }
+      }
+    }
+    val deletedResult = deletedBindings.groupBy(_._2)
+      .mapValues(_.map{ case (binding, _) => new AclBindingDeleteResult(binding, deleteExceptions.getOrElse(binding, null)) })
+    (0 until aclBindingFilters.size).map { i =>
+      new AclDeleteResult(deletedResult.getOrElse(i, Set.empty[AclBindingDeleteResult]).toSet.asJava)
+    }.map(CompletableFuture.completedFuture[AclDeleteResult]).asJava
+  }
+
+  override def acls(filter: AclBindingFilter): lang.Iterable[AclBinding] = {
+    inReadLock(lock) {
+      unorderedAcls.flatMap { case (resource, versionedAcls) =>
+        versionedAcls.acls.map(acl => AuthorizerUtils.convertToAclBinding(resource, acl))
+            .filter(filter.matches)
+      }.asJava
+    }
+  }
+
+  override def close(): Unit = {
+    aclChangeListeners.foreach(listener => listener.close())
+    if (zkClient != null) zkClient.close()
+  }
+
+  private def authorizeAction(requestContext: AuthorizableRequestContext, action: Action): AuthorizationResult = {
+    val resource = AuthorizerUtils.convertToResource(action.resourcePattern)
+    if (resource.patternType != PatternType.LITERAL) {
+      throw new IllegalArgumentException("Only literal resources are supported. Got: " + resource.patternType)
+    }
+
+    // ensure we compare identical classes
+    val sessionPrincipal = requestContext.principal
+    val principal = if (classOf[KafkaPrincipal] != sessionPrincipal.getClass)
+      new KafkaPrincipal(sessionPrincipal.getPrincipalType, sessionPrincipal.getName)
+    else
+      sessionPrincipal
+
+    val host = requestContext.clientAddress.getHostAddress
+    val operation = Operation.fromJava(action.operation)
+
+    def isEmptyAclAndAuthorized(acls: Set[Acl]): Boolean = {
+      if (acls.isEmpty) {
+        // No ACLs found for this resource, permission is determined by value of config allow.everyone.if.no.acl.found
+        authorizerLogger.debug(s"No acl found for resource $resource, authorized = $shouldAllowEveryoneIfNoAclIsFound")
+        shouldAllowEveryoneIfNoAclIsFound
+      } else false
+    }
+
+    def denyAclExists(acls: Set[Acl]): Boolean = {
+      // Check if there are any Deny ACLs which would forbid this operation.
+      matchingAclExists(operation, resource, principal, host, Deny, acls)
+    }
+
+    def allowAclExists(acls: Set[Acl]): Boolean = {
+      // Check if there are any Allow ACLs which would allow this operation.
+      // Allowing read, write, delete, or alter implies allowing describe.
+      // See #{org.apache.kafka.common.acl.AclOperation} for more details about ACL inheritance.
+      val allowOps = operation match {
+        case Describe => Set[Operation](Describe, Read, Write, Delete, Alter)
+        case DescribeConfigs => Set[Operation](DescribeConfigs, AlterConfigs)
+        case _ => Set[Operation](operation)
+      }
+      allowOps.exists(operation => matchingAclExists(operation, resource, principal, host, Allow, acls))
+    }
+
+    def aclsAllowAccess = {
+      //we allow an operation if no acls are found and user has configured to allow all users
+      //when no acls are found or if no deny acls are found and at least one allow acls matches.
+      val acls = matchingAcls(resource.resourceType, resource.name)
+      isEmptyAclAndAuthorized(acls) || (!denyAclExists(acls) && allowAclExists(acls))
+    }
+
+    // Evaluate if operation is allowed
+    val authorized = isSuperUser(principal) || aclsAllowAccess
+
+    logAuditMessage(requestContext, action, authorized)
+    if (authorized) AuthorizationResult.ALLOWED else AuthorizationResult.DENIED
+  }
+
+  def isSuperUser(principal: KafkaPrincipal): Boolean = {
+    if (superUsers.contains(principal)) {
+      authorizerLogger.debug(s"principal = $principal is a super user, allowing operation without checking acls.")
+      true
+    } else false
+  }
+
+  private def matchingAcls(resourceType: ResourceType, resourceName: String): Set[Acl] = {
+    inReadLock(lock) {
+      val wildcard = aclCache.get(Resource(resourceType, Acl.WildCardResource, PatternType.LITERAL))
+        .map(_.acls)
+        .getOrElse(Set.empty[Acl])
+
+      val literal = aclCache.get(Resource(resourceType, resourceName, PatternType.LITERAL))
+        .map(_.acls)
+        .getOrElse(Set.empty[Acl])
+
+      val prefixed = aclCache
+        .from(Resource(resourceType, resourceName, PatternType.PREFIXED))
+        .to(Resource(resourceType, resourceName.take(1), PatternType.PREFIXED))
+        .filterKeys(resource => resourceName.startsWith(resource.name))
+        .values
+        .flatMap { _.acls }
+        .toSet
+
+      prefixed ++ wildcard ++ literal
+    }
+  }
+
+  private def matchingAclExists(operation: Operation, resource: Resource, principal: KafkaPrincipal, host: String, permissionType: PermissionType, acls: Set[Acl]): Boolean = {
+    acls.find { acl =>
+      acl.permissionType == permissionType &&
+        (acl.principal == principal || acl.principal == Acl.WildCardPrincipal) &&
+        (operation == acl.operation || acl.operation == All) &&
+        (acl.host == host || acl.host == Acl.WildCardHost)
+    }.exists { acl =>
+      authorizerLogger.debug(s"operation = $operation on resource = $resource from host = $host is $permissionType based on acl = $acl")
+      true
+    }
+  }
+
+  private def loadCache(): Unit = {
+    inWriteLock(lock) {
+      ZkAclStore.stores.foreach(store => {
+        val resourceTypes = zkClient.getResourceTypes(store.patternType)
+        for (rType <- resourceTypes) {
+          val resourceType = Try(ResourceType.fromString(rType))
+          resourceType match {
+            case Success(resourceTypeObj) =>
+              val resourceNames = zkClient.getResourceNames(store.patternType, resourceTypeObj)
+              for (resourceName <- resourceNames) {
+                val resource = new Resource(resourceTypeObj, resourceName, store.patternType)
+                val versionedAcls = getAclsFromZk(resource)
+                updateCache(resource, versionedAcls)
+              }
+            case Failure(_) => warn(s"Ignoring unknown ResourceType: $rType")
+          }
+        }
+      })
+    }
+  }
+
+  private[authorizer] def startZkChangeListeners(): Unit = {
+    aclChangeListeners = ZkAclChangeStore.stores
+      .map(store => store.createListener(AclChangedNotificationHandler, zkClient))
+  }
+
+  def logAuditMessage(requestContext: AuthorizableRequestContext, action: Action, authorized: Boolean): Unit = {
+    def logMessage: String = {
+      val principal = requestContext.principal
+      val operation = Operation.fromJava(action.operation)
+      val host = requestContext.clientAddress.getHostAddress
+      val resource = AuthorizerUtils.convertToResource(action.resourcePattern)
+      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
+      s"Principal = $principal is $authResult Operation = $operation from host = $host on resource = $resource for request = $apiKey with resourceRefCount = $refCount"
+    }
+
+    if (authorized) {
+      // logIfAllowed is true if access is granted to the resource as a result of this authorization.
+      // In this case, log at debug level. If false, no access is actually granted, the result is used
+      // only to determine authorized operations. So log only at trace level.
+      if (action.logIfAllowed)
+        authorizerLogger.debug(logMessage)
+      else
+        authorizerLogger.trace(logMessage)
+    } else {
+      // logIfDenied is true if access to the resource was explicitly requested. Since this is an attempt
+      // to access unauthorized resources, log at info level. If false, this is either a request to determine
+      // authorized operations or a filter (e.g for regex subscriptions) to filter out authorized resources.
+      // In this case, log only at trace level.
+      if (action.logIfDenied)
+        authorizerLogger.info(logMessage)
+      else
+        authorizerLogger.trace(logMessage)
+    }
+  }
+
+  /**
+    * Safely updates the resources ACLs by ensuring reads and writes respect the expected zookeeper version.
+    * Continues to retry until it successfully updates zookeeper.
+    *
+    * Returns a boolean indicating if the content of the ACLs was actually changed.
+    *
+    * @param resource the resource to change ACLs for
+    * @param getNewAcls function to transform existing acls to new ACLs
+    * @return boolean indicating if a change was made
+    */
+  private def updateResourceAcls(resource: Resource)(getNewAcls: Set[Acl] => Set[Acl]): Boolean = {
+    var currentVersionedAcls =
+      if (aclCache.contains(resource))
+        getAclsFromCache(resource)
+      else
+        getAclsFromZk(resource)
+    var newVersionedAcls: VersionedAcls = null
+    var writeComplete = false
+    var retries = 0
+    while (!writeComplete && retries <= maxUpdateRetries) {
+      val newAcls = getNewAcls(currentVersionedAcls.acls)
+      val (updateSucceeded, updateVersion) =
+        if (newAcls.nonEmpty) {
+          if (currentVersionedAcls.exists)
+            zkClient.conditionalSetAclsForResource(resource, newAcls, currentVersionedAcls.zkVersion)
+          else
+            zkClient.createAclsForResourceIfNotExists(resource, newAcls)
+        } else {
+          trace(s"Deleting path for $resource because it had no ACLs remaining")
+          (zkClient.conditionalDelete(resource, currentVersionedAcls.zkVersion), 0)
+        }
+
+      if (!updateSucceeded) {
+        trace(s"Failed to update ACLs for $resource. Used version ${currentVersionedAcls.zkVersion}. Reading data and retrying update.")
+        Thread.sleep(backoffTime)
+        currentVersionedAcls = getAclsFromZk(resource)
+        retries += 1
+      } else {
+        newVersionedAcls = VersionedAcls(newAcls, updateVersion)
+        writeComplete = updateSucceeded
+      }
+    }
+
+    if(!writeComplete)
+      throw new IllegalStateException(s"Failed to update ACLs for $resource after trying a maximum of $maxUpdateRetries times")
+
+    if (newVersionedAcls.acls != currentVersionedAcls.acls) {
+      debug(s"Updated ACLs for $resource to ${newVersionedAcls.acls} with version ${newVersionedAcls.zkVersion}")
+      updateCache(resource, newVersionedAcls)
+      updateAclChangedFlag(resource)
+      true
+    } else {
+      debug(s"Updated ACLs for $resource, no change was made")
+      updateCache(resource, newVersionedAcls) // Even if no change, update the version
+      false
+    }
+  }
+
+  // Returns Map instead of SortedMap since most callers don't care about ordering. In Scala 2.13, mapping from SortedMap
+  // to Map is restricted by default
+  private def unorderedAcls: Map[Resource, VersionedAcls] = aclCache
+
+  private def getAclsFromCache(resource: Resource): VersionedAcls = {
+    aclCache.getOrElse(resource, throw new IllegalArgumentException(s"ACLs do not exist in the cache for resource $resource"))
+  }
+
+  private def getAclsFromZk(resource: Resource): VersionedAcls = {
+    zkClient.getVersionedAclsForResource(resource)
+  }
+
+  private def updateCache(resource: Resource, versionedAcls: VersionedAcls): Unit = {
+    if (versionedAcls.acls.nonEmpty) {
+      aclCache = aclCache + (resource -> versionedAcls)
+    } else {
+      aclCache = aclCache - resource
+    }
+  }
+
+  private def updateAclChangedFlag(resource: Resource): Unit = {
+      zkClient.createAclChangeNotification(resource)
+  }
+
+  private def backoffTime = {
+    retryBackoffMs + Random.nextInt(retryBackoffJitterMs)
+  }
+
+  private def apiException(e: Throwable): ApiException = {
+    e match {
+      case e1: ApiException => e1
+      case e1 => new ApiException(e1)
+    }
+  }
+
+  object AclChangedNotificationHandler extends AclChangeNotificationHandler {
+    override def processNotification(resource: Resource): Unit = {
+      inWriteLock(lock) {
+        val versionedAcls = getAclsFromZk(resource)
+        updateCache(resource, versionedAcls)
+      }
+    }
+  }
+}
diff --git a/core/src/main/scala/kafka/security/SecurityUtils.scala b/core/src/main/scala/kafka/security/authorizer/AuthorizerUtils.scala
similarity index 52%
rename from core/src/main/scala/kafka/security/SecurityUtils.scala
rename to core/src/main/scala/kafka/security/authorizer/AuthorizerUtils.scala
index 311e195795d7e..909f3a78bbb7e 100644
--- a/core/src/main/scala/kafka/security/SecurityUtils.scala
+++ b/core/src/main/scala/kafka/security/authorizer/AuthorizerUtils.scala
@@ -15,18 +15,26 @@
   * limitations under the License.
   */
 
-package kafka.security
+package kafka.security.authorizer
 
-import kafka.security.auth.{Acl, Operation, PermissionType, Resource, ResourceType}
-import org.apache.kafka.common.acl.{AccessControlEntry, AclBinding, AclBindingFilter}
+import java.net.InetAddress
+
+import kafka.network.RequestChannel.Session
+import kafka.security.auth._
+import org.apache.kafka.common.acl.{AccessControlEntry, AclBinding, AclBindingFilter, AclOperation}
 import org.apache.kafka.common.protocol.Errors
 import org.apache.kafka.common.requests.ApiError
-import org.apache.kafka.common.resource.ResourcePattern
+import org.apache.kafka.common.resource.{ResourcePattern, ResourceType => JResourceType}
+import org.apache.kafka.common.security.auth.{KafkaPrincipal, SecurityProtocol}
 import org.apache.kafka.common.utils.SecurityUtils._
+import org.apache.kafka.server.authorizer.AuthorizableRequestContext
+
 import scala.util.{Failure, Success, Try}
 
 
-object SecurityUtils {
+object AuthorizerUtils {
+  val WildcardPrincipal = "User:*"
+  val WildcardHost = "*"
 
   def convertToResourceAndAcl(filter: AclBindingFilter): Either[ApiError, (Resource, Acl)] = {
     (for {
@@ -44,10 +52,44 @@ object SecurityUtils {
 
   def convertToAclBinding(resource: Resource, acl: Acl): AclBinding = {
     val resourcePattern = new ResourcePattern(resource.resourceType.toJava, resource.name, resource.patternType)
-    val entry = new AccessControlEntry(acl.principal.toString, acl.host.toString,
+    new AclBinding(resourcePattern, convertToAccessControlEntry(acl))
+  }
+
+  def convertToAccessControlEntry(acl: Acl): AccessControlEntry = {
+    new AccessControlEntry(acl.principal.toString, acl.host.toString,
       acl.operation.toJava, acl.permissionType.toJava)
-    new AclBinding(resourcePattern, entry)
+  }
+
+  def convertToAcl(ace: AccessControlEntry): Acl = {
+    new Acl(parseKafkaPrincipal(ace.principal), PermissionType.fromJava(ace.permissionType), ace.host,
+      Operation.fromJava(ace.operation))
+  }
+
+  def convertToResource(resourcePattern: ResourcePattern): Resource = {
+    Resource(ResourceType.fromJava(resourcePattern.resourceType), resourcePattern.name, resourcePattern.patternType)
+  }
+
+  def validateAclBinding(aclBinding: AclBinding): Unit = {
+    if (aclBinding.isUnknown)
+      throw new IllegalArgumentException("ACL binding contains unknown elements")
+  }
+
+  def supportedOperations(resourceType: JResourceType): Set[AclOperation] = {
+    ResourceType.fromJava(resourceType).supportedOperations.map(_.toJava)
   }
 
   def isClusterResource(name: String): Boolean = name.equals(Resource.ClusterResourceName)
+
+  def sessionToRequestContext(session: Session): AuthorizableRequestContext = {
+    new AuthorizableRequestContext {
+      override def clientId(): String = ""
+      override def requestType(): Int = -1
+      override def listenerName(): String = ""
+      override def clientAddress(): InetAddress = session.clientAddress
+      override def principal(): KafkaPrincipal = session.principal
+      override def securityProtocol(): SecurityProtocol = null
+      override def correlationId(): Int = -1
+      override def requestVersion(): Int = -1
+    }
+  }
 }
diff --git a/core/src/main/scala/kafka/security/authorizer/AuthorizerWrapper.scala b/core/src/main/scala/kafka/security/authorizer/AuthorizerWrapper.scala
new file mode 100644
index 0000000000000..9556c8c6b9136
--- /dev/null
+++ b/core/src/main/scala/kafka/security/authorizer/AuthorizerWrapper.scala
@@ -0,0 +1,136 @@
+/**
+ * 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 java.util.concurrent.{CompletableFuture, CompletionStage}
+import java.{lang, util}
+
+import kafka.network.RequestChannel.Session
+import kafka.security.auth.{Acl, Operation, Resource}
+import org.apache.kafka.common.Endpoint
+import org.apache.kafka.common.acl.{AccessControlEntry, AclBinding, AclBindingFilter}
+import org.apache.kafka.common.errors.{ApiException, InvalidRequestException}
+import org.apache.kafka.common.requests.ApiError
+import org.apache.kafka.common.resource.ResourcePattern
+import org.apache.kafka.server.authorizer.AclDeleteResult.AclBindingDeleteResult
+import org.apache.kafka.server.authorizer.{AuthorizableRequestContext, AuthorizerServerInfo, _}
+
+import scala.collection.JavaConverters._
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.{Seq, immutable, mutable}
+
+class AuthorizerWrapper(private[kafka] val baseAuthorizer: kafka.security.auth.Authorizer) extends Authorizer {
+
+  override def configure(configs: util.Map[String, _]): Unit = {
+    baseAuthorizer.configure(configs)
+  }
+
+  override def start(serverInfo: AuthorizerServerInfo): util.Map[Endpoint, _ <: CompletionStage[Void]] = {
+    serverInfo.endpoints.asScala.map { endpoint =>
+      endpoint -> CompletableFuture.completedFuture[Void](null) }.toMap.asJava
+  }
+
+  override def authorize(requestContext: AuthorizableRequestContext, actions: util.List[Action]): util.List[AuthorizationResult] = {
+    val session = Session(requestContext.principal, requestContext.clientAddress)
+    actions.asScala.map { action =>
+      val operation = Operation.fromJava(action.operation)
+      if (baseAuthorizer.authorize(session, operation, AuthorizerUtils.convertToResource(action.resourcePattern)))
+        AuthorizationResult.ALLOWED
+      else
+        AuthorizationResult.DENIED
+    }.asJava
+  }
+
+  override def createAcls(requestContext: AuthorizableRequestContext,
+                          aclBindings: util.List[AclBinding]): util.List[_ <: CompletionStage[AclCreateResult]] = {
+    aclBindings.asScala
+      .map { aclBinding =>
+        AuthorizerUtils.convertToResourceAndAcl(aclBinding.toFilter) match {
+          case Left(apiError) => new AclCreateResult(apiError.exception)
+          case Right((resource, acl)) =>
+            try {
+              baseAuthorizer.addAcls(Set(acl), resource)
+              AclCreateResult.SUCCESS
+            } catch {
+              case e: ApiException => new AclCreateResult(e)
+              case e: Throwable => new AclCreateResult(new InvalidRequestException("Failed to create ACL", e))
+            }
+        }
+      }.toList.map(CompletableFuture.completedFuture[AclCreateResult]).asJava
+  }
+
+  override def deleteAcls(requestContext: AuthorizableRequestContext,
+                          aclBindingFilters: util.List[AclBindingFilter]): util.List[_ <: CompletionStage[AclDeleteResult]] = {
+    val filters = aclBindingFilters.asScala
+    val results = mutable.Map[Int, AclDeleteResult]()
+    val toDelete = mutable.Map[Int, ArrayBuffer[(Resource, Acl)]]()
+
+    if (filters.forall(_.matchesAtMostOne)) {
+      // Delete based on a list of ACL fixtures.
+      for ((filter, i) <- filters.zipWithIndex) {
+        AuthorizerUtils.convertToResourceAndAcl(filter) match {
+          case Left(apiError) => results.put(i, new AclDeleteResult(apiError.exception))
+          case Right(binding) => toDelete.put(i, ArrayBuffer(binding))
+        }
+      }
+    } else {
+      // Delete based on filters that may match more than one ACL.
+      val aclMap = baseAuthorizer.getAcls()
+      val filtersWithIndex = filters.zipWithIndex
+      for ((resource, acls) <- aclMap; acl <- acls) {
+        val binding = new AclBinding(
+          new ResourcePattern(resource.resourceType.toJava, resource.name, resource.patternType),
+          new AccessControlEntry(acl.principal.toString, acl.host.toString, acl.operation.toJava,
+            acl.permissionType.toJava))
+
+        for ((filter, i) <- filtersWithIndex if filter.matches(binding))
+          toDelete.getOrElseUpdate(i, ArrayBuffer.empty) += ((resource, acl))
+      }
+    }
+
+    for ((i, acls) <- toDelete) {
+      val deletionResults = acls.flatMap { case (resource, acl) =>
+        val aclBinding = AuthorizerUtils.convertToAclBinding(resource, acl)
+        try {
+          if (baseAuthorizer.removeAcls(immutable.Set(acl), resource))
+            Some(new AclBindingDeleteResult(aclBinding))
+          else None
+        } catch {
+          case throwable: Throwable =>
+            Some(new AclBindingDeleteResult(aclBinding, ApiError.fromThrowable(throwable).exception))
+        }
+      }.asJava
+
+      results.put(i, new AclDeleteResult(deletionResults))
+    }
+
+    filters.indices.map { i =>
+      results.getOrElse(i, new AclDeleteResult(Seq.empty[AclBindingDeleteResult].asJava))
+    }.map(CompletableFuture.completedFuture[AclDeleteResult]).asJava
+  }
+
+  override def acls(filter: AclBindingFilter): lang.Iterable[AclBinding] = {
+    baseAuthorizer.getAcls().flatMap { case (resource, acls) =>
+      acls.map(acl => AuthorizerUtils.convertToAclBinding(resource, acl)).filter(filter.matches)
+    }.asJava
+  }
+
+  override def close(): Unit = {
+    baseAuthorizer.close()
+  }
+}
\ No newline at end of file
diff --git a/core/src/main/scala/kafka/server/AbstractFetcherManager.scala b/core/src/main/scala/kafka/server/AbstractFetcherManager.scala
index 10ae8df132952..fbb49c975b271 100755
--- a/core/src/main/scala/kafka/server/AbstractFetcherManager.scala
+++ b/core/src/main/scala/kafka/server/AbstractFetcherManager.scala
@@ -34,13 +34,14 @@ abstract class AbstractFetcherManager[T <: AbstractFetcherThread](val name: Stri
   private[server] val fetcherThreadMap = new mutable.HashMap[BrokerIdAndFetcherId, T]
   private val lock = new Object
   private var numFetchersPerBroker = numFetchers
+  val failedPartitions = new FailedPartitions
   this.logIdent = "[" + name + "] "
 
   newGauge(
     "MaxLag",
     new Gauge[Long] {
       // current max lag across all fetchers/topics/partitions
-      def value = fetcherThreadMap.foldLeft(0L)((curMaxAll, fetcherThreadMapEntry) => {
+      def value: Long = fetcherThreadMap.foldLeft(0L)((curMaxAll, fetcherThreadMapEntry) => {
         fetcherThreadMapEntry._2.fetcherLagStats.stats.foldLeft(0L)((curMaxThread, fetcherLagStatsEntry) => {
           curMaxThread.max(fetcherLagStatsEntry._2.lag)
         }).max(curMaxAll)
@@ -53,7 +54,7 @@ abstract class AbstractFetcherManager[T <: AbstractFetcherThread](val name: Stri
   "MinFetchRate", {
     new Gauge[Double] {
       // current min fetch rate across all fetchers/topics/partitions
-      def value = {
+      def value: Double = {
         val headRate: Double =
           fetcherThreadMap.headOption.map(_._2.fetcherStats.requestRate.oneMinuteRate).getOrElse(0)
 
@@ -66,6 +67,42 @@ abstract class AbstractFetcherManager[T <: AbstractFetcherThread](val name: Stri
   Map("clientId" -> clientId)
   )
 
+  newGauge(
+    "FetchFailureRate", {
+      new Gauge[Double] {
+        // fetch failure rate sum across all fetchers/topics/partitions
+        def value: Double = {
+          val headRate: Double =
+            fetcherThreadMap.headOption.map(_._2.fetcherStats.requestFailureRate.oneMinuteRate).getOrElse(0)
+
+          fetcherThreadMap.foldLeft(headRate)((curSum, fetcherThreadMapEntry) => {
+            fetcherThreadMapEntry._2.fetcherStats.requestRate.oneMinuteRate + curSum
+          })
+        }
+      }
+    },
+    Map("clientId" -> clientId)
+  )
+
+  val failedPartitionsCount = newGauge(
+    "FailedPartitionsCount", {
+      new Gauge[Int] {
+        def value: Int = failedPartitions.size
+      }
+    },
+    Map("clientId" -> clientId)
+  )
+
+  newGauge("DeadThreadCount", {
+    new Gauge[Int] {
+      def value: Int = {
+        deadThreadCount
+      }
+    }
+  }, Map("clientId" -> clientId))
+
+  private[server] def deadThreadCount: Int = lock synchronized { fetcherThreadMap.values.count(_.isThreadFailed) }
+
   def resizeThreadPool(newSize: Int): Unit = {
     def migratePartitions(newSize: Int): Unit = {
       fetcherThreadMap.foreach { case (id, thread) =>
@@ -107,7 +144,7 @@ abstract class AbstractFetcherManager[T <: AbstractFetcherThread](val name: Stri
   }
 
   // This method is only needed by ReplicaAlterDirManager
-  def markPartitionsForTruncation(brokerId: Int, topicPartition: TopicPartition, truncationOffset: Long) {
+  def markPartitionsForTruncation(brokerId: Int, topicPartition: TopicPartition, truncationOffset: Long): Unit = {
     lock synchronized {
       val fetcherId = getFetcherId(topicPartition)
       val brokerIdAndFetcherId = BrokerIdAndFetcherId(brokerId, fetcherId)
@@ -120,13 +157,14 @@ abstract class AbstractFetcherManager[T <: AbstractFetcherThread](val name: Stri
   // to be defined in subclass to create a specific fetcher
   def createFetcherThread(fetcherId: Int, sourceBroker: BrokerEndPoint): T
 
-  def addFetcherForPartitions(partitionAndOffsets: Map[TopicPartition, InitialFetchState]) {
+  def addFetcherForPartitions(partitionAndOffsets: Map[TopicPartition, InitialFetchState]): Unit = {
     lock synchronized {
       val partitionsPerFetcher = partitionAndOffsets.groupBy { case (topicPartition, brokerAndInitialFetchOffset) =>
         BrokerAndFetcherId(brokerAndInitialFetchOffset.leader, getFetcherId(topicPartition))
       }
 
-      def addAndStartFetcherThread(brokerAndFetcherId: BrokerAndFetcherId, brokerIdAndFetcherId: BrokerIdAndFetcherId): AbstractFetcherThread = {
+      def addAndStartFetcherThread(brokerAndFetcherId: BrokerAndFetcherId,
+                                   brokerIdAndFetcherId: BrokerIdAndFetcherId): T = {
         val fetcherThread = createFetcherThread(brokerAndFetcherId.fetcherId, brokerAndFetcherId.broker)
         fetcherThreadMap.put(brokerIdAndFetcherId, fetcherThread)
         fetcherThread.start()
@@ -150,25 +188,32 @@ abstract class AbstractFetcherManager[T <: AbstractFetcherThread](val name: Stri
           tp -> OffsetAndEpoch(brokerAndInitOffset.initOffset, brokerAndInitOffset.currentLeaderEpoch)
         }
 
-        fetcherThread.addPartitions(initialOffsetAndEpochs)
-        info(s"Added fetcher to broker ${brokerAndFetcherId.broker} for partitions $initialOffsetAndEpochs")
+        addPartitionsToFetcherThread(fetcherThread, initialOffsetAndEpochs)
       }
     }
   }
 
-  def removeFetcherForPartitions(partitions: Set[TopicPartition]) {
+  protected def addPartitionsToFetcherThread(fetcherThread: T,
+                                             initialOffsetAndEpochs: collection.Map[TopicPartition, OffsetAndEpoch]): Unit = {
+    fetcherThread.addPartitions(initialOffsetAndEpochs)
+    info(s"Added fetcher to broker ${fetcherThread.sourceBroker.id} for partitions $initialOffsetAndEpochs")
+  }
+
+  def removeFetcherForPartitions(partitions: Set[TopicPartition]): Unit = {
     lock synchronized {
       for (fetcher <- fetcherThreadMap.values)
         fetcher.removePartitions(partitions)
+      failedPartitions.removeAll(partitions)
     }
-    info(s"Removed fetcher for partitions $partitions")
+    if (partitions.nonEmpty)
+      info(s"Removed fetcher for partitions $partitions")
   }
 
-  def shutdownIdleFetcherThreads() {
+  def shutdownIdleFetcherThreads(): Unit = {
     lock synchronized {
       val keysToBeRemoved = new mutable.HashSet[BrokerIdAndFetcherId]
       for ((key, fetcher) <- fetcherThreadMap) {
-        if (fetcher.partitionCount <= 0) {
+        if (fetcher.idle) {
           fetcher.shutdown()
           keysToBeRemoved += key
         }
@@ -177,7 +222,7 @@ abstract class AbstractFetcherManager[T <: AbstractFetcherThread](val name: Stri
     }
   }
 
-  def closeAllFetchers() {
+  def closeAllFetchers(): Unit = {
     lock synchronized {
       for ( (_, fetcher) <- fetcherThreadMap) {
         fetcher.initiateShutdown()
@@ -191,6 +236,37 @@ abstract class AbstractFetcherManager[T <: AbstractFetcherThread](val name: Stri
   }
 }
 
+/**
+  * The class FailedPartitions would keep a track of partitions marked as failed either during truncation or appending
+  * resulting from one of the following errors -
+  * 
        
+  *   
      1.  Storage exception
+  *   
      2.  Fenced epoch
+  *   
      3.  Unexpected errors
+  * 
      
+  * The partitions which fail due to storage error are eventually removed from this set after the log directory is
+  * taken offline.
+  */
+class FailedPartitions {
+  private val failedPartitionsSet = new mutable.HashSet[TopicPartition]
+
+  def size: Int = synchronized {
+    failedPartitionsSet.size
+  }
+
+  def add(topicPartition: TopicPartition): Unit = synchronized {
+    failedPartitionsSet += topicPartition
+  }
+
+  def removeAll(topicPartitions: Set[TopicPartition]): Unit = synchronized {
+    failedPartitionsSet --= topicPartitions
+  }
+
+  def contains(topicPartition: TopicPartition): Boolean = synchronized {
+    failedPartitionsSet.contains(topicPartition)
+  }
+}
+
 case class BrokerAndFetcherId(broker: BrokerEndPoint, fetcherId: Int)
 
 case class InitialFetchState(leader: BrokerEndPoint, currentLeaderEpoch: Int, initOffset: Long)
diff --git a/core/src/main/scala/kafka/server/AbstractFetcherThread.scala b/core/src/main/scala/kafka/server/AbstractFetcherThread.scala
index 3cc6137fcb732..c8d59f142ab89 100755
--- a/core/src/main/scala/kafka/server/AbstractFetcherThread.scala
+++ b/core/src/main/scala/kafka/server/AbstractFetcherThread.scala
@@ -31,7 +31,7 @@ import kafka.utils.CoreUtils.inLock
 import org.apache.kafka.common.protocol.Errors
 import AbstractFetcherThread._
 
-import scala.collection.{Map, Set, mutable}
+import scala.collection.{Map, Seq, Set, mutable}
 import scala.collection.JavaConverters._
 import java.util.concurrent.TimeUnit
 import java.util.concurrent.atomic.AtomicLong
@@ -39,7 +39,7 @@ import java.util.function.Consumer
 
 import com.yammer.metrics.core.Gauge
 import kafka.log.LogAppendInfo
-import org.apache.kafka.common.{KafkaException, TopicPartition}
+import org.apache.kafka.common.{InvalidRecordException, TopicPartition}
 import org.apache.kafka.common.internals.PartitionStates
 import org.apache.kafka.common.record.{FileRecords, MemoryRecords, Records}
 import org.apache.kafka.common.requests._
@@ -52,6 +52,7 @@ import scala.math._
 abstract class AbstractFetcherThread(name: String,
                                      clientId: String,
                                      val sourceBroker: BrokerEndPoint,
+                                     failedPartitions: FailedPartitions,
                                      fetchBackOffMs: Int = 0,
                                      isInterruptible: Boolean = true)
   extends ShutdownableThread(name, isInterruptible) {
@@ -60,13 +61,15 @@ abstract class AbstractFetcherThread(name: String,
   type EpochData = OffsetsForLeaderEpochRequest.PartitionData
 
   private val partitionStates = new PartitionStates[PartitionFetchState]
-  private val partitionMapLock = new ReentrantLock
+  protected val partitionMapLock = new ReentrantLock
   private val partitionMapCond = partitionMapLock.newCondition()
 
   private val metricId = ClientIdAndBroker(clientId, sourceBroker.host, sourceBroker.port)
   val fetcherStats = new FetcherStats(metricId)
   val fetcherLagStats = new FetcherLagStats(metricId)
 
+  @volatile var idle = false
+
   /* callbacks to be defined in subclass */
 
   // process fetched data
@@ -96,7 +99,7 @@ abstract class AbstractFetcherThread(name: String,
 
   protected def isOffsetForLeaderEpochSupported: Boolean
 
-  override def shutdown() {
+  override def shutdown(): Unit = {
     initiateShutdown()
     inLock(partitionMapLock) {
       partitionMapCond.signalAll()
@@ -108,7 +111,7 @@ abstract class AbstractFetcherThread(name: String,
     fetcherLagStats.unregister()
   }
 
-  override def doWork() {
+  override def doWork(): Unit = {
     maybeTruncate()
     maybeFetch()
   }
@@ -134,10 +137,11 @@ abstract class AbstractFetcherThread(name: String,
   }
 
   // deal with partitions with errors, potentially due to leadership changes
-  private def handlePartitionsWithErrors(partitions: Iterable[TopicPartition], methodName: String) {
-    if (partitions.nonEmpty)
+  private def handlePartitionsWithErrors(partitions: Iterable[TopicPartition], methodName: String): Unit = {
+    if (partitions.nonEmpty) {
       debug(s"Handling errors in $methodName for partitions $partitions")
       delayPartitions(partitions, fetchBackOffMs)
+    }
   }
 
   /**
@@ -175,12 +179,30 @@ abstract class AbstractFetcherThread(name: String,
     }
   }
 
+  private def doTruncate(topicPartition: TopicPartition, truncationState: OffsetTruncationState): Boolean = {
+    try {
+      truncate(topicPartition, truncationState)
+      true
+    }
+    catch {
+      case e: KafkaStorageException =>
+        error(s"Failed to truncate $topicPartition at offset ${truncationState.offset}", e)
+        markPartitionFailed(topicPartition)
+        false
+      case t: Throwable =>
+        error(s"Unexpected error occurred during truncation for $topicPartition "
+          + s"at offset ${truncationState.offset}", t)
+        markPartitionFailed(topicPartition)
+        false
+    }
+  }
+
   /**
     * - Build a leader epoch fetch based on partitions that are in the Truncating phase
     * - Send OffsetsForLeaderEpochRequest, retrieving the latest offset for each partition's
     *   leader epoch. This is the offset the follower should truncate to ensure
     *   accurate log replication.
-    * - Finally truncate the logs for partitions in the truncating phase and mark them
+    * - Finally truncate the logs for partitions in the truncating phase and mark the
     *   truncation complete. Do this within a lock to ensure no leadership changes can
     *   occur during truncation.
     */
@@ -199,7 +221,7 @@ abstract class AbstractFetcherThread(name: String,
         curPartitionState != null && leaderEpochInRequest == curPartitionState.currentLeaderEpoch
       }
 
-      val ResultWithPartitions(fetchOffsets, partitionsWithError) = maybeTruncateToEpochEndOffsets(epochEndOffsets)
+      val ResultWithPartitions(fetchOffsets, partitionsWithError) = maybeTruncateToEpochEndOffsets(epochEndOffsets, latestEpochsForPartitions)
       handlePartitionsWithErrors(partitionsWithError, "truncateToEpochEndOffsets")
       updateFetchOffsetAndMaybeMarkTruncationComplete(fetchOffsets)
     }
@@ -208,53 +230,43 @@ abstract class AbstractFetcherThread(name: String,
   // Visible for testing
   private[server] def truncateToHighWatermark(partitions: Set[TopicPartition]): Unit = inLock(partitionMapLock) {
     val fetchOffsets = mutable.HashMap.empty[TopicPartition, OffsetTruncationState]
-    val partitionsWithError = mutable.HashSet.empty[TopicPartition]
 
     for (tp <- partitions) {
       val partitionState = partitionStates.stateValue(tp)
       if (partitionState != null) {
-        try {
-          val highWatermark = partitionState.fetchOffset
-          val truncationState = OffsetTruncationState(highWatermark, truncationCompleted = true)
-
-          info(s"Truncating partition $tp to local high watermark $highWatermark")
-          truncate(tp, truncationState)
+        val highWatermark = partitionState.fetchOffset
+        val truncationState = OffsetTruncationState(highWatermark, truncationCompleted = true)
 
+        info(s"Truncating partition $tp to local high watermark $highWatermark")
+        if (doTruncate(tp, truncationState))
           fetchOffsets.put(tp, truncationState)
-        } catch {
-          case e: KafkaStorageException =>
-            info(s"Failed to truncate $tp", e)
-            partitionsWithError += tp
-        }
       }
     }
 
-    handlePartitionsWithErrors(partitionsWithError, "truncateToHighWatermark")
     updateFetchOffsetAndMaybeMarkTruncationComplete(fetchOffsets)
   }
 
-  private def maybeTruncateToEpochEndOffsets(fetchedEpochs: Map[TopicPartition, EpochEndOffset]): ResultWithPartitions[Map[TopicPartition, OffsetTruncationState]] = {
+  private def maybeTruncateToEpochEndOffsets(fetchedEpochs: Map[TopicPartition, EpochEndOffset],
+                                             latestEpochsForPartitions: Map[TopicPartition, EpochData]): ResultWithPartitions[Map[TopicPartition, OffsetTruncationState]] = {
     val fetchOffsets = mutable.HashMap.empty[TopicPartition, OffsetTruncationState]
     val partitionsWithError = mutable.HashSet.empty[TopicPartition]
 
     fetchedEpochs.foreach { case (tp, leaderEpochOffset) =>
-      try {
-        leaderEpochOffset.error match {
-          case Errors.NONE =>
-            val offsetTruncationState = getOffsetTruncationState(tp, leaderEpochOffset)
-            truncate(tp, offsetTruncationState)
+      leaderEpochOffset.error match {
+        case Errors.NONE =>
+          val offsetTruncationState = getOffsetTruncationState(tp, leaderEpochOffset)
+          if(doTruncate(tp, offsetTruncationState))
             fetchOffsets.put(tp, offsetTruncationState)
 
-          case Errors.FENCED_LEADER_EPOCH =>
-            onPartitionFenced(tp)
+        case Errors.FENCED_LEADER_EPOCH =>
+          if (onPartitionFenced(tp, latestEpochsForPartitions.get(tp).flatMap {
+            p =>
+              if (p.currentLeaderEpoch.isPresent) Some(p.currentLeaderEpoch.get())
+              else None
+          })) partitionsWithError += tp
 
-          case error =>
-            info(s"Retrying leaderEpoch request for partition $tp as the leader reported an error: $error")
-            partitionsWithError += tp
-        }
-      } catch {
-        case e: KafkaStorageException =>
-          info(s"Failed to truncate $tp", e)
+        case error =>
+          info(s"Retrying leaderEpoch request for partition $tp as the leader reported an error: $error")
           partitionsWithError += tp
       }
     }
@@ -262,12 +274,22 @@ abstract class AbstractFetcherThread(name: String,
     ResultWithPartitions(fetchOffsets, partitionsWithError)
   }
 
-  private def onPartitionFenced(tp: TopicPartition): Unit = inLock(partitionMapLock) {
-    Option(partitionStates.stateValue(tp)).foreach { currentFetchState =>
+  /**
+   * remove the partition if the partition state is NOT updated. Otherwise, keep the partition active.
+   * @return true if the epoch in this thread is updated. otherwise, false
+   */
+  private def onPartitionFenced(tp: TopicPartition, requestEpoch: Option[Int]): Boolean = inLock(partitionMapLock) {
+    Option(partitionStates.stateValue(tp)).exists { currentFetchState =>
       val currentLeaderEpoch = currentFetchState.currentLeaderEpoch
-      info(s"Partition $tp has an older epoch ($currentLeaderEpoch) than the current leader. Will await " +
-        s"the new LeaderAndIsr state before resuming fetching.")
-      partitionStates.remove(tp)
+      if (requestEpoch.contains(currentLeaderEpoch)) {
+        info(s"Partition $tp has an older epoch ($currentLeaderEpoch) than the current leader. Will await " +
+          s"the new LeaderAndIsr state before resuming fetching.")
+        markPartitionFailed(tp)
+        false
+      } else {
+        info(s"Partition $tp has an new epoch ($currentLeaderEpoch) than the current leader. retry the partition later")
+        true
+      }
     }
   }
 
@@ -283,6 +305,7 @@ abstract class AbstractFetcherThread(name: String,
       case t: Throwable =>
         if (isRunning) {
           warn(s"Error in response for fetch request $fetchRequest", t)
+          fetcherStats.requestFailureRate.mark()
           inLock(partitionMapLock) {
             partitionsWithError ++= partitionStates.partitionSet.asScala
             // there is an error occurred while fetching partitions, sleep a while
@@ -304,6 +327,10 @@ abstract class AbstractFetcherThread(name: String,
             // the current offset is the same as the offset requested.
             val fetchState = fetchStates(topicPartition)
             if (fetchState.fetchOffset == currentFetchState.fetchOffset && currentFetchState.isReadyForFetch) {
+              val requestEpoch = if (fetchState.currentLeaderEpoch >= 0)
+                Some(fetchState.currentLeaderEpoch)
+              else
+                None
               partitionData.error match {
                 case Errors.NONE =>
                   try {
@@ -326,7 +353,7 @@ abstract class AbstractFetcherThread(name: String,
                       }
                     }
                   } catch {
-                    case ime: CorruptRecordException =>
+                    case ime@( _: CorruptRecordException | _: InvalidRecordException) =>
                       // we log the error and continue. This ensures two things
                       // 1. If there is a corrupt message in a topic partition, it does not bring the fetcher thread
                       //    down and cause other topic partition to also lag
@@ -336,14 +363,17 @@ abstract class AbstractFetcherThread(name: String,
                         s"offset ${currentFetchState.fetchOffset}", ime)
                       partitionsWithError += topicPartition
                     case e: KafkaStorageException =>
-                      error(s"Error while processing data for partition $topicPartition", e)
-                      partitionsWithError += topicPartition
-                    case e: Throwable =>
-                      throw new KafkaException(s"Error processing data for partition $topicPartition " +
-                        s"offset ${currentFetchState.fetchOffset}", e)
+                      error(s"Error while processing data for partition $topicPartition " +
+                        s"at offset ${currentFetchState.fetchOffset}", e)
+                      markPartitionFailed(topicPartition)
+                    case t: Throwable =>
+                      // stop monitoring this partition and add it to the set of failed partitions
+                      error(s"Unexpected error occurred while processing data for partition $topicPartition " +
+                        s"at offset ${currentFetchState.fetchOffset}", t)
+                      markPartitionFailed(topicPartition)
                   }
                 case Errors.OFFSET_OUT_OF_RANGE =>
-                  if (!handleOutOfRangeError(topicPartition, currentFetchState))
+                  if (handleOutOfRangeError(topicPartition, currentFetchState, requestEpoch))
                     partitionsWithError += topicPartition
 
                 case Errors.UNKNOWN_LEADER_EPOCH =>
@@ -352,13 +382,18 @@ abstract class AbstractFetcherThread(name: String,
                   partitionsWithError += topicPartition
 
                 case Errors.FENCED_LEADER_EPOCH =>
-                  onPartitionFenced(topicPartition)
+                  if (onPartitionFenced(topicPartition, requestEpoch)) partitionsWithError += topicPartition
 
                 case Errors.NOT_LEADER_FOR_PARTITION =>
                   debug(s"Remote broker is not the leader for partition $topicPartition, which could indicate " +
                     "that the partition is being moved")
                   partitionsWithError += topicPartition
 
+                case Errors.UNKNOWN_TOPIC_OR_PARTITION =>
+                  warn(s"Remote broker does not host the partition $topicPartition, which could indicate " +
+                    "that the partition is being created or deleted.")
+                  partitionsWithError += topicPartition
+
                 case _ =>
                   error(s"Error for partition $topicPartition at offset ${currentFetchState.fetchOffset}",
                     partitionData.error.exception)
@@ -375,7 +410,7 @@ abstract class AbstractFetcherThread(name: String,
     }
   }
 
-  def markPartitionsForTruncation(topicPartition: TopicPartition, truncationOffset: Long) {
+  def markPartitionsForTruncation(topicPartition: TopicPartition, truncationOffset: Long): Unit = {
     partitionMapLock.lockInterruptibly()
     try {
       Option(partitionStates.stateValue(topicPartition)).foreach { state =>
@@ -387,9 +422,20 @@ abstract class AbstractFetcherThread(name: String,
     } finally partitionMapLock.unlock()
   }
 
-  def addPartitions(initialFetchStates: Map[TopicPartition, OffsetAndEpoch]) {
+  private def markPartitionFailed(topicPartition: TopicPartition): Unit = {
+    partitionMapLock.lock()
+    try {
+      failedPartitions.add(topicPartition)
+      removePartitions(Set(topicPartition))
+    } finally partitionMapLock.unlock()
+    warn(s"Partition $topicPartition marked as failed")
+  }
+
+  def addPartitions(initialFetchStates: Map[TopicPartition, OffsetAndEpoch]): Set[TopicPartition] = {
     partitionMapLock.lockInterruptibly()
     try {
+      failedPartitions.removeAll(initialFetchStates.keySet)
+
       initialFetchStates.foreach { case (tp, initialFetchState) =>
         // We can skip the truncation step iff the leader epoch matches the existing epoch
         val currentState = partitionStates.stateValue(tp)
@@ -405,7 +451,9 @@ abstract class AbstractFetcherThread(name: String,
         partitionStates.updateAndMoveToEnd(tp, updatedState)
       }
 
+      maybeUpdateIdleFlag()
       partitionMapCond.signalAll()
+      initialFetchStates.keySet
     } finally partitionMapLock.unlock()
   }
 
@@ -415,7 +463,7 @@ abstract class AbstractFetcherThread(name: String,
     *
     * @param fetchOffsets the partitions to update fetch offset and maybe mark truncation complete
     */
-  private def updateFetchOffsetAndMaybeMarkTruncationComplete(fetchOffsets: Map[TopicPartition, OffsetTruncationState]) {
+  private def updateFetchOffsetAndMaybeMarkTruncationComplete(fetchOffsets: Map[TopicPartition, OffsetTruncationState]): Unit = {
     val newStates: Map[TopicPartition, PartitionFetchState] = partitionStates.partitionStates.asScala
       .map { state =>
         val currentFetchState = state.value
@@ -429,6 +477,7 @@ abstract class AbstractFetcherThread(name: String,
         (state.topicPartition, maybeTruncationComplete)
       }.toMap
     partitionStates.set(newStates.asJava)
+    maybeUpdateIdleFlag()
   }
 
   /**
@@ -504,32 +553,34 @@ abstract class AbstractFetcherThread(name: String,
   }
 
   /**
-   * Handle the out of range error. Return true if the request succeeded or was fenced, which means we need
-   * not backoff and retry. False if there was a retriable error.
+   * Handle the out of range error. Return false if
+   * 1) the request succeeded or
+   * 2) was fenced and this thread haven't received new epoch,
+   * which means we need not backoff and retry. True if there was a retriable error.
    */
   private def handleOutOfRangeError(topicPartition: TopicPartition,
-                                    fetchState: PartitionFetchState): Boolean = {
+                                    fetchState: PartitionFetchState,
+                                    requestEpoch: Option[Int]): Boolean = {
     try {
       val newOffset = fetchOffsetAndTruncate(topicPartition, fetchState.currentLeaderEpoch)
       val newFetchState = PartitionFetchState(newOffset, fetchState.currentLeaderEpoch, state = Fetching)
       partitionStates.updateAndMoveToEnd(topicPartition, newFetchState)
       info(s"Current offset ${fetchState.fetchOffset} for partition $topicPartition is " +
         s"out of range, which typically implies a leader change. Reset fetch offset to $newOffset")
-      true
+      false
     } catch {
       case _: FencedLeaderEpochException =>
-        onPartitionFenced(topicPartition)
-        true
+        onPartitionFenced(topicPartition, requestEpoch)
 
       case e @ (_ : UnknownTopicOrPartitionException |
                 _ : UnknownLeaderEpochException |
                 _ : NotLeaderForPartitionException) =>
         info(s"Could not fetch offset for $topicPartition due to error: ${e.getMessage}")
-        false
+        true
 
       case e: Throwable =>
         error(s"Error getting offset for partition $topicPartition", e)
-        false
+        true
     }
   }
 
@@ -588,7 +639,7 @@ abstract class AbstractFetcherThread(name: String,
     }
   }
 
-  def delayPartitions(partitions: Iterable[TopicPartition], delay: Long) {
+  def delayPartitions(partitions: Iterable[TopicPartition], delay: Long): Unit = {
     partitionMapLock.lockInterruptibly()
     try {
       for (partition <- partitions) {
@@ -603,17 +654,18 @@ abstract class AbstractFetcherThread(name: String,
     } finally partitionMapLock.unlock()
   }
 
-  def removePartitions(topicPartitions: Set[TopicPartition]) {
+  def removePartitions(topicPartitions: Set[TopicPartition]): Unit = {
     partitionMapLock.lockInterruptibly()
     try {
       topicPartitions.foreach { topicPartition =>
         partitionStates.remove(topicPartition)
         fetcherLagStats.unregister(topicPartition)
       }
+      maybeUpdateIdleFlag()
     } finally partitionMapLock.unlock()
   }
 
-  def partitionCount() = {
+  def partitionCount(): Int = {
     partitionMapLock.lockInterruptibly()
     try partitionStates.size
     finally partitionMapLock.unlock()
@@ -643,6 +695,11 @@ abstract class AbstractFetcherThread(name: String,
     }
   }
 
+  // This method should only be called when holding the partitionMapLock
+  private def maybeUpdateIdleFlag(): Unit = {
+    idle = partitionStates.size() <= 0
+  }
+
 }
 
 object AbstractFetcherThread {
@@ -654,6 +711,7 @@ object AbstractFetcherThread {
 object FetcherMetrics {
   val ConsumerLag = "ConsumerLag"
   val RequestsPerSec = "RequestsPerSec"
+  val RequestFailuresPerSec = "RequestFailuresPerSec"
   val BytesPerSec = "BytesPerSec"
 }
 
@@ -672,13 +730,13 @@ class FetcherLagMetrics(metricId: ClientIdTopicPartition) extends KafkaMetricsGr
     tags
   )
 
-  def lag_=(newLag: Long) {
+  def lag_=(newLag: Long): Unit = {
     lagVal.set(newLag)
   }
 
   def lag = lagVal.get
 
-  def unregister() {
+  def unregister(): Unit = {
     removeMetric(FetcherMetrics.ConsumerLag, tags)
   }
 }
@@ -699,12 +757,12 @@ class FetcherLagStats(metricId: ClientIdAndBroker) {
       false
   }
 
-  def unregister(topicPartition: TopicPartition) {
+  def unregister(topicPartition: TopicPartition): Unit = {
     val lagMetrics = stats.remove(ClientIdTopicPartition(metricId.clientId, topicPartition))
     if (lagMetrics != null) lagMetrics.unregister()
   }
 
-  def unregister() {
+  def unregister(): Unit = {
     stats.keys.toBuffer.foreach { key: ClientIdTopicPartition =>
       unregister(key.topicPartition)
     }
@@ -718,10 +776,13 @@ class FetcherStats(metricId: ClientIdAndBroker) extends KafkaMetricsGroup {
 
   val requestRate = newMeter(FetcherMetrics.RequestsPerSec, "requests", TimeUnit.SECONDS, tags)
 
+  val requestFailureRate = newMeter(FetcherMetrics.RequestFailuresPerSec, "requestFailures", TimeUnit.SECONDS, tags)
+
   val byteRate = newMeter(FetcherMetrics.BytesPerSec, "bytes", TimeUnit.SECONDS, tags)
 
-  def unregister() {
+  def unregister(): Unit = {
     removeMetric(FetcherMetrics.RequestsPerSec, tags)
+    removeMetric(FetcherMetrics.RequestFailuresPerSec, tags)
     removeMetric(FetcherMetrics.BytesPerSec, tags)
   }
 
diff --git a/core/src/main/scala/kafka/server/AdminManager.scala b/core/src/main/scala/kafka/server/AdminManager.scala
index 0cdaad6965ba6..8da9e3195cd38 100644
--- a/core/src/main/scala/kafka/server/AdminManager.scala
+++ b/core/src/main/scala/kafka/server/AdminManager.scala
@@ -18,16 +18,22 @@ package kafka.server
 
 import java.util.{Collections, Properties}
 
-import kafka.admin.{AdminOperationException, AdminUtils}
+import kafka.admin.AdminOperationException
 import kafka.common.TopicAlreadyMarkedForDeletionException
+import kafka.controller.KafkaController
 import kafka.log.LogConfig
+import kafka.utils.Log4jController
 import kafka.metrics.KafkaMetricsGroup
 import kafka.utils._
 import kafka.zk.{AdminZkClient, KafkaZkClient}
-import org.apache.kafka.common.config.{AbstractConfig, ConfigDef, ConfigException, ConfigResource}
-import org.apache.kafka.common.errors.{ApiException, InvalidConfigurationException, InvalidPartitionsException, InvalidReplicaAssignmentException, InvalidRequestException, ReassignmentInProgressException, UnknownTopicOrPartitionException}
+import org.apache.kafka.clients.admin.AlterConfigOp
+import org.apache.kafka.clients.admin.AlterConfigOp.OpType
+import org.apache.kafka.common.config.ConfigDef.ConfigKey
+import org.apache.kafka.common.config.{AbstractConfig, ConfigDef, ConfigException, ConfigResource, LogLevelConfig}
+import org.apache.kafka.common.errors.{ApiException, InvalidConfigurationException, InvalidPartitionsException, InvalidReplicaAssignmentException, InvalidRequestException, ReassignmentInProgressException, TopicExistsException, UnknownTopicOrPartitionException}
 import org.apache.kafka.common.internals.Topic
 import org.apache.kafka.common.message.CreateTopicsRequestData.CreatableTopic
+import org.apache.kafka.common.message.CreateTopicsResponseData.{CreatableTopicConfigs, CreatableTopicResult}
 import org.apache.kafka.common.metrics.Metrics
 import org.apache.kafka.common.network.ListenerName
 import org.apache.kafka.common.protocol.Errors
@@ -38,13 +44,14 @@ import org.apache.kafka.common.requests.{AlterConfigsRequest, ApiError, Describe
 import org.apache.kafka.server.policy.{AlterConfigPolicy, CreateTopicPolicy}
 import org.apache.kafka.server.policy.CreateTopicPolicy.RequestMetadata
 
-import scala.collection.{mutable, _}
+import scala.collection.{Map, mutable, _}
 import scala.collection.JavaConverters._
 
 class AdminManager(val config: KafkaConfig,
                    val metrics: Metrics,
                    val metadataCache: MetadataCache,
-                   val zkClient: KafkaZkClient) extends Logging with KafkaMetricsGroup {
+                   val zkClient: KafkaZkClient,
+                   val controller: KafkaController) extends Logging with KafkaMetricsGroup {
 
   this.logIdent = "[Admin Manager on Broker " + config.brokerId + "]: "
 
@@ -57,12 +64,15 @@ class AdminManager(val config: KafkaConfig,
   private val alterConfigPolicy =
     Option(config.getConfiguredInstance(KafkaConfig.AlterConfigPolicyClassNameProp, classOf[AlterConfigPolicy]))
 
-  def hasDelayedTopicOperations = topicPurgatory.delayed != 0
+  def hasDelayedTopicOperations = topicPurgatory.numDelayed != 0
+
+  private val defaultNumPartitions = config.numPartitions.intValue()
+  private val defaultReplicationFactor = config.defaultReplicationFactor.shortValue()
 
   /**
     * Try to complete delayed topic operations with the request key
     */
-  def tryCompleteDelayedTopicOperations(topic: String) {
+  def tryCompleteDelayedTopicOperations(topic: String): Unit = {
     val key = TopicKey(topic)
     val completed = topicPurgatory.checkAndComplete(key)
     debug(s"Request key ${key.keyLabel} unblocked $completed topic requests.")
@@ -75,15 +85,19 @@ class AdminManager(val config: KafkaConfig,
   def createTopics(timeout: Int,
                    validateOnly: Boolean,
                    toCreate: Map[String, CreatableTopic],
-                   responseCallback: Map[String, ApiError] => Unit) {
+                   includeConfigsAndMetatadata: Map[String, CreatableTopicResult],
+                   responseCallback: Map[String, ApiError] => Unit): Unit = {
 
     // 1. map over topics creating assignment and calling zookeeper
     val brokers = metadataCache.getAliveBrokers.map { b => kafka.admin.BrokerMetadata(b.id, b.rack) }
     val metadata = toCreate.values.map(topic =>
       try {
+        if (metadataCache.contains(topic.name))
+          throw new TopicExistsException(s"Topic '${topic.name}' already exists.")
+
         val configs = new Properties()
-        topic.configs().asScala.foreach { case entry =>
-          configs.setProperty(entry.name(), entry.value())
+        topic.configs.asScala.foreach { entry =>
+          configs.setProperty(entry.name, entry.value)
         }
         LogConfig.validate(configs)
 
@@ -92,8 +106,15 @@ class AdminManager(val config: KafkaConfig,
           throw new InvalidRequestException("Both numPartitions or replicationFactor and replicasAssignments were set. " +
             "Both cannot be used at the same time.")
         }
+
+        val resolvedNumPartitions = if (topic.numPartitions == NO_NUM_PARTITIONS)
+          defaultNumPartitions else topic.numPartitions
+        val resolvedReplicationFactor = if (topic.replicationFactor == NO_REPLICATION_FACTOR)
+          defaultReplicationFactor else topic.replicationFactor
+
         val assignments = if (topic.assignments().isEmpty) {
-          AdminUtils.assignReplicasToBrokers(brokers, topic.numPartitions, topic.replicationFactor)
+          adminZkClient.assignReplicasToAvailableBrokers(brokers, controller.partitionUnassignableBrokerIds.toSet,
+            resolvedNumPartitions, resolvedReplicationFactor)
         } else {
           val assignments = new mutable.HashMap[Int, Seq[Int]]
           // Note: we don't check that replicaAssignment contains unknown brokers - unlike in add-partitions case,
@@ -112,23 +133,15 @@ class AdminManager(val config: KafkaConfig,
 
             // Use `null` for unset fields in the public API
             val numPartitions: java.lang.Integer =
-              if (topic.numPartitions == NO_NUM_PARTITIONS) null else topic.numPartitions
+              if (topic.assignments().isEmpty) resolvedNumPartitions else null
             val replicationFactor: java.lang.Short =
-              if (topic.replicationFactor == NO_REPLICATION_FACTOR) null else topic.replicationFactor
+              if (topic.assignments().isEmpty) resolvedReplicationFactor else null
             val javaAssignments = if (topic.assignments().isEmpty) {
               null
             } else {
-              val map = new java.util.HashMap[Integer, java.util.List[Integer]]
-              assignments.foreach {
-                case (k, v) => {
-                  val list = new java.util.ArrayList[Integer]
-                  v.foreach {
-                    case i => list.add(Integer.valueOf(i))
-                  }
-                  map.put(k, list)
-                }
-              }
-              map
+              assignments.map { case (k, v) =>
+                (k: java.lang.Integer) -> v.map(i => i: java.lang.Integer).asJava
+              }.asJava
             }
             val javaConfigs = new java.util.HashMap[String, String]
             topic.configs().asScala.foreach(config => javaConfigs.put(config.name(), config.value()))
@@ -144,9 +157,33 @@ class AdminManager(val config: KafkaConfig,
             else
               adminZkClient.createTopicWithAssignment(topic.name, configs, assignments)
         }
+
+        // For responses with DescribeConfigs permission, populate metadata and configs
+        includeConfigsAndMetatadata.get(topic.name).foreach { result =>
+          val logConfig = LogConfig.fromProps(KafkaServer.copyKafkaConfigToLog(config), configs)
+          val createEntry = createTopicConfigEntry(logConfig, configs, includeSynonyms = false)(_, _)
+          val topicConfigs = logConfig.values.asScala.map { case (k, v) =>
+            val entry = createEntry(k, v)
+            val source = ConfigSource.values.indices.map(_.toByte)
+              .find(i => ConfigSource.forId(i.toByte) == entry.source)
+              .getOrElse(0.toByte)
+            new CreatableTopicConfigs()
+                .setName(k)
+                .setValue(entry.value)
+                .setIsSensitive(entry.isSensitive)
+                .setReadOnly(entry.isReadOnly)
+                .setConfigSource(source)
+          }.toList.asJava
+          result.setConfigs(topicConfigs)
+          result.setNumPartitions(assignments.size)
+          result.setReplicationFactor(assignments(0).size.toShort)
+        }
         CreatePartitionsMetadata(topic.name, assignments, ApiError.NONE)
       } catch {
         // Log client errors at a lower level than unexpected exceptions
+        case e: TopicExistsException =>
+          debug(s"Topic creation failed since topic '${topic.name}' already exists.", e)
+          CreatePartitionsMetadata(topic.name, Map(), ApiError.fromThrowable(e))
         case e: ApiException =>
           info(s"Error processing create topic request $topic", e)
           CreatePartitionsMetadata(topic.name, Map(), ApiError.fromThrowable(e))
@@ -156,7 +193,7 @@ class AdminManager(val config: KafkaConfig,
         case e: Throwable =>
           error(s"Error processing create topic request $topic", e)
           CreatePartitionsMetadata(topic.name, Map(), ApiError.fromThrowable(e))
-      })
+      }).toBuffer
 
     // 2. if timeout <= 0, validateOnly or no topics can proceed return immediately
     if (timeout <= 0 || validateOnly || !metadata.exists(_.error.is(Errors.NONE))) {
@@ -171,9 +208,8 @@ class AdminManager(val config: KafkaConfig,
       responseCallback(results)
     } else {
       // 3. else pass the assignments and errors to the delayed operation and set the keys
-      val delayedCreate = new DelayedCreatePartitions(timeout, metadata.toSeq, this, responseCallback)
-      val delayedCreateKeys = toCreate.values.map(
-        topic => new TopicKey(topic.name())).toSeq
+      val delayedCreate = new DelayedCreatePartitions(timeout, metadata, this, responseCallback)
+      val delayedCreateKeys = toCreate.values.map(topic => new TopicKey(topic.name)).toBuffer
       // try to complete the request immediately, otherwise put it into the purgatory
       topicPurgatory.tryCompleteElseWatch(delayedCreate, delayedCreateKeys)
     }
@@ -185,7 +221,7 @@ class AdminManager(val config: KafkaConfig,
     */
   def deleteTopics(timeout: Int,
                    topics: Set[String],
-                   responseCallback: Map[String, Errors] => Unit) {
+                   responseCallback: Map[String, Errors] => Unit): Unit = {
 
     // 1. map over topics calling the asynchronous delete
     val metadata = topics.map { topic =>
@@ -228,20 +264,20 @@ class AdminManager(val config: KafkaConfig,
                        listenerName: ListenerName,
                        callback: Map[String, ApiError] => Unit): Unit = {
 
-    val reassignPartitionsInProgress = zkClient.reassignPartitionsInProgress
     val allBrokers = adminZkClient.getBrokerMetadatas()
     val allBrokerIds = allBrokers.map(_.id)
 
     // 1. map over topics creating assignment and calling AdminUtils
     val metadata = newPartitions.map { case (topic, newPartition) =>
       try {
-        // We prevent addition partitions while a reassignment is in progress, since
-        // during reassignment there is no meaningful notion of replication factor
-        if (reassignPartitionsInProgress)
-          throw new ReassignmentInProgressException("A partition reassignment is in progress.")
-
-        val existingAssignment = zkClient.getReplicaAssignmentForTopics(immutable.Set(topic)).map {
-          case (topicPartition, replicas) => topicPartition.partition -> replicas
+        val existingAssignment = zkClient.getFullReplicaAssignmentForTopics(immutable.Set(topic)).map {
+          case (topicPartition, assignment) =>
+            if (assignment.isBeingReassigned) {
+              // We prevent adding partitions while topic reassignment is in progress, to protect from a race condition
+              // between the controller thread processing reassignment update and createPartitions(this) request.
+              throw new ReassignmentInProgressException(s"A partition reassignment is in progress for the topic '$topic'.")
+            }
+            topicPartition.partition -> assignment
         }
         if (existingAssignment.isEmpty)
           throw new UnknownTopicOrPartitionException(s"The topic '$topic' does not exist.")
@@ -256,7 +292,7 @@ class AdminManager(val config: KafkaConfig,
           throw new InvalidPartitionsException(s"Topic already has $oldNumPartitions partitions.")
         }
 
-        val reassignment = Option(newPartition.newAssignments).map(_.asScala.map(_.asScala.map(_.toInt))).map { assignments =>
+        val newPartitionsAssignment = Option(newPartition.newAssignments).map(_.asScala.map(_.asScala.map(_.toInt))).map { assignments =>
           val unknownBrokers = assignments.flatten.toSet -- allBrokerIds
           if (unknownBrokers.nonEmpty)
             throw new InvalidReplicaAssignmentException(
@@ -273,7 +309,8 @@ class AdminManager(val config: KafkaConfig,
         }
 
         val updatedReplicaAssignment = adminZkClient.addPartitions(topic, existingAssignment, allBrokers,
-          newPartition.totalCount, reassignment, validateOnly = validateOnly)
+          newPartition.totalCount, newPartitionsAssignment, validateOnly = validateOnly,
+          noNewPartitionBrokerIds = controller.partitionUnassignableBrokerIds.toSet)
         CreatePartitionsMetadata(topic, updatedReplicaAssignment, ApiError.NONE)
       } catch {
         case e: AdminOperationException =>
@@ -314,7 +351,7 @@ class AdminManager(val config: KafkaConfig,
         val filteredConfigPairs = configs.filter { case (configName, _) =>
           /* Always returns true if configNames is None */
           configNames.forall(_.contains(configName))
-        }.toIndexedSeq
+        }.toBuffer
 
         val configEntries = filteredConfigPairs.map { case (name, value) => createConfigEntry(name, value) }
         new DescribeConfigsResponse.Config(ApiError.NONE, configEntries.asJava)
@@ -343,8 +380,16 @@ class AdminManager(val config: KafkaConfig,
               createResponseConfig(allConfigs(config),
                 createBrokerConfigEntry(perBrokerConfig = true, includeSynonyms))
             else
-              throw new InvalidRequestException(s"Unexpected broker id, expected ${config.brokerId} or empty string, but received $resource.name")
+              throw new InvalidRequestException(s"Unexpected broker id, expected ${config.brokerId} or empty string, but received ${resource.name}")
 
+          case ConfigResource.Type.BROKER_LOGGER =>
+            if (resource.name == null || resource.name.isEmpty)
+              throw new InvalidRequestException("Broker id must not be empty")
+            else if (resourceNameToBrokerId(resource.name) != config.brokerId)
+              throw new InvalidRequestException(s"Unexpected broker id, expected ${config.brokerId} but received ${resource.name}")
+            else
+              createResponseConfig(Log4jController.loggers,
+                (name, value) => new DescribeConfigsResponse.ConfigEntry(name, value.toString, ConfigSource.DYNAMIC_BROKER_LOGGER_CONFIG, false, false, List.empty.asJava))
           case resourceType => throw new InvalidRequestException(s"Unsupported resource type: $resourceType")
         }
         resource -> resourceConfig
@@ -364,58 +409,131 @@ class AdminManager(val config: KafkaConfig,
   def alterConfigs(configs: Map[ConfigResource, AlterConfigsRequest.Config], validateOnly: Boolean): Map[ConfigResource, ApiError] = {
     configs.map { case (resource, config) =>
 
-      def validateConfigPolicy(resourceType: ConfigResource.Type): Unit = {
-        alterConfigPolicy match {
-          case Some(policy) =>
-            val configEntriesMap = config.entries.asScala.map(entry => (entry.name, entry.value)).toMap
-            policy.validate(new AlterConfigPolicy.RequestMetadata(
-              new ConfigResource(resourceType, resource.name), configEntriesMap.asJava))
-          case None =>
+      try {
+        val configEntriesMap = config.entries.asScala.map(entry => (entry.name, entry.value)).toMap
+
+        val configProps = new Properties
+        config.entries.asScala.foreach { configEntry =>
+          configProps.setProperty(configEntry.name, configEntry.value)
         }
+
+        resource.`type` match {
+          case ConfigResource.Type.TOPIC => alterTopicConfigs(resource, validateOnly, configProps, configEntriesMap)
+          case ConfigResource.Type.BROKER => alterBrokerConfigs(resource, validateOnly, configProps, configEntriesMap)
+          case resourceType =>
+            throw new InvalidRequestException(s"AlterConfigs is only supported for topics and brokers, but resource type is $resourceType")
+        }
+      } catch {
+        case e @ (_: ConfigException | _: IllegalArgumentException) =>
+          val message = s"Invalid config value for resource $resource: ${e.getMessage}"
+          info(message)
+          resource -> ApiError.fromThrowable(new InvalidRequestException(message, e))
+        case e: Throwable =>
+          // Log client errors at a lower level than unexpected exceptions
+          val message = s"Error processing alter configs request for resource $resource, config $config"
+          if (e.isInstanceOf[ApiException])
+            info(message, e)
+          else
+            error(message, e)
+          resource -> ApiError.fromThrowable(e)
       }
+    }.toMap
+  }
+
+  private def alterTopicConfigs(resource: ConfigResource, validateOnly: Boolean,
+                                configProps: Properties, configEntriesMap: Map[String, String]): (ConfigResource, ApiError) = {
+    val topic = resource.name
+    adminZkClient.validateTopicConfig(topic, configProps)
+    validateConfigPolicy(resource, configEntriesMap)
+    if (!validateOnly) {
+      info(s"Updating topic $topic with new configuration $config")
+      adminZkClient.changeTopicConfig(topic, configProps)
+    }
+
+    resource -> ApiError.NONE
+  }
+
+  private def alterBrokerConfigs(resource: ConfigResource, validateOnly: Boolean,
+                                 configProps: Properties, configEntriesMap: Map[String, String]): (ConfigResource, ApiError) = {
+    val brokerId = getBrokerId(resource)
+    val perBrokerConfig = brokerId.nonEmpty
+    this.config.dynamicConfig.validate(configProps, perBrokerConfig)
+    validateConfigPolicy(resource, configEntriesMap)
+    if (!validateOnly) {
+      if (perBrokerConfig)
+        this.config.dynamicConfig.reloadUpdatedFilesWithoutConfigChange(configProps)
+      adminZkClient.changeBrokerConfig(brokerId,
+        this.config.dynamicConfig.toPersistentProps(configProps, perBrokerConfig))
+    }
+
+    resource -> ApiError.NONE
+  }
+
+  private def alterLogLevelConfigs(alterConfigOps: List[AlterConfigOp]): Unit = {
+    alterConfigOps.foreach { alterConfigOp =>
+      val loggerName = alterConfigOp.configEntry().name()
+      val logLevel = alterConfigOp.configEntry().value()
+      alterConfigOp.opType() match {
+        case OpType.SET => Log4jController.logLevel(loggerName, logLevel)
+        case OpType.DELETE => Log4jController.unsetLogLevel(loggerName)
+      }
+    }
+  }
+
+  private def getBrokerId(resource: ConfigResource) = {
+    if (resource.name == null || resource.name.isEmpty)
+      None
+    else {
+      val id = resourceNameToBrokerId(resource.name)
+      if (id != this.config.brokerId)
+        throw new InvalidRequestException(s"Unexpected broker id, expected ${this.config.brokerId}, but received ${resource.name}")
+      Some(id)
+    }
+  }
+
+  private def validateConfigPolicy(resource: ConfigResource, configEntriesMap: Map[String, String]): Unit = {
+    alterConfigPolicy match {
+      case Some(policy) =>
+        policy.validate(new AlterConfigPolicy.RequestMetadata(
+          new ConfigResource(resource.`type`(), resource.name), configEntriesMap.asJava))
+      case None =>
+    }
+  }
+
+  def incrementalAlterConfigs(configs: Map[ConfigResource, List[AlterConfigOp]], validateOnly: Boolean): Map[ConfigResource, ApiError] = {
+    configs.map { case (resource, alterConfigOps) =>
       try {
+        // throw InvalidRequestException if any duplicate keys
+        val duplicateKeys = alterConfigOps.groupBy(config => config.configEntry().name())
+          .mapValues(_.size).filter(_._2 > 1).keys.toSet
+        if (duplicateKeys.nonEmpty)
+          throw new InvalidRequestException(s"Error due to duplicate config keys : ${duplicateKeys.mkString(",")}")
+
+        val configEntriesMap = alterConfigOps.map(entry => (entry.configEntry().name(), entry.configEntry().value())).toMap
+
         resource.`type` match {
           case ConfigResource.Type.TOPIC =>
-            val topic = resource.name
+            val configProps = adminZkClient.fetchEntityConfig(ConfigType.Topic, resource.name)
+            prepareIncrementalConfigs(alterConfigOps, configProps, LogConfig.configKeys)
+            alterTopicConfigs(resource, validateOnly, configProps, configEntriesMap)
 
-            val properties = new Properties
-            config.entries.asScala.foreach { configEntry =>
-              properties.setProperty(configEntry.name, configEntry.value)
-            }
+          case ConfigResource.Type.BROKER =>
+            val brokerId = getBrokerId(resource)
+            val perBrokerConfig = brokerId.nonEmpty
 
-            adminZkClient.validateTopicConfig(topic, properties)
-            validateConfigPolicy(ConfigResource.Type.TOPIC)
-            if (!validateOnly) {
-              info(s"Updating topic $topic with new configuration $config")
-              adminZkClient.changeTopicConfig(topic, properties)
-            }
+            val persistentProps = if (perBrokerConfig) adminZkClient.fetchEntityConfig(ConfigType.Broker, brokerId.get.toString)
+            else adminZkClient.fetchEntityConfig(ConfigType.Broker, ConfigEntityName.Default)
 
-            resource -> ApiError.NONE
+            val configProps = this.config.dynamicConfig.fromPersistentProps(persistentProps, perBrokerConfig)
+            prepareIncrementalConfigs(alterConfigOps, configProps, KafkaConfig.configKeys)
+            alterBrokerConfigs(resource, validateOnly, configProps, configEntriesMap)
 
-          case ConfigResource.Type.BROKER =>
-            val brokerId = if (resource.name == null || resource.name.isEmpty)
-              None
-            else {
-              val id = resourceNameToBrokerId(resource.name)
-              if (id != this.config.brokerId)
-                throw new InvalidRequestException(s"Unexpected broker id, expected ${this.config.brokerId}, but received $resource.name")
-              Some(id)
-            }
-            val configProps = new Properties
-            config.entries.asScala.foreach { configEntry =>
-              configProps.setProperty(configEntry.name, configEntry.value)
-            }
-
-            val perBrokerConfig = brokerId.nonEmpty
-            this.config.dynamicConfig.validate(configProps, perBrokerConfig)
-            validateConfigPolicy(ConfigResource.Type.BROKER)
-            if (!validateOnly) {
-              if (perBrokerConfig)
-                this.config.dynamicConfig.reloadUpdatedFilesWithoutConfigChange(configProps)
-              adminZkClient.changeBrokerConfig(brokerId,
-                this.config.dynamicConfig.toPersistentProps(configProps, perBrokerConfig))
-            }
+          case ConfigResource.Type.BROKER_LOGGER =>
+            getBrokerId(resource)
+            validateLogLevelConfigs(alterConfigOps)
 
+            if (!validateOnly)
+              alterLogLevelConfigs(alterConfigOps)
             resource -> ApiError.NONE
           case resourceType =>
             throw new InvalidRequestException(s"AlterConfigs is only supported for topics and brokers, but resource type is $resourceType")
@@ -427,7 +545,7 @@ class AdminManager(val config: KafkaConfig,
           resource -> ApiError.fromThrowable(new InvalidRequestException(message, e))
         case e: Throwable =>
           // Log client errors at a lower level than unexpected exceptions
-          val message = s"Error processing alter configs request for resource $resource, config $config"
+          val message = s"Error processing alter configs request for resource $resource, config $alterConfigOps"
           if (e.isInstanceOf[ApiException])
             info(message, e)
           else
@@ -437,7 +555,67 @@ class AdminManager(val config: KafkaConfig,
     }.toMap
   }
 
-  def shutdown() {
+  private def validateLogLevelConfigs(alterConfigOps: List[AlterConfigOp]): Unit = {
+    def validateLoggerNameExists(loggerName: String): Unit = {
+      if (!Log4jController.loggerExists(loggerName))
+        throw new ConfigException(s"Logger $loggerName does not exist!")
+    }
+
+    alterConfigOps.foreach { alterConfigOp =>
+      val loggerName = alterConfigOp.configEntry().name()
+      alterConfigOp.opType() match {
+        case OpType.SET =>
+          validateLoggerNameExists(loggerName)
+          val logLevel = alterConfigOp.configEntry().value()
+          if (!LogLevelConfig.VALID_LOG_LEVELS.contains(logLevel)) {
+            val validLevelsStr = LogLevelConfig.VALID_LOG_LEVELS.asScala.mkString(", ")
+            throw new ConfigException(
+              s"Cannot set the log level of $loggerName to $logLevel as it is not a supported log level. " +
+              s"Valid log levels are $validLevelsStr"
+            )
+          }
+        case OpType.DELETE =>
+          validateLoggerNameExists(loggerName)
+          if (loggerName == Log4jController.ROOT_LOGGER)
+            throw new InvalidRequestException(s"Removing the log level of the ${Log4jController.ROOT_LOGGER} logger is not allowed")
+        case OpType.APPEND => throw new InvalidRequestException(s"${OpType.APPEND} operation is not allowed for the ${ConfigResource.Type.BROKER_LOGGER} resource")
+        case OpType.SUBTRACT => throw new InvalidRequestException(s"${OpType.SUBTRACT} operation is not allowed for the ${ConfigResource.Type.BROKER_LOGGER} resource")
+      }
+    }
+  }
+
+  private def prepareIncrementalConfigs(alterConfigOps: List[AlterConfigOp], configProps: Properties, configKeys: Map[String, ConfigKey]): Unit = {
+
+    def listType(configName: String, configKeys: Map[String, ConfigKey]): Boolean = {
+      val configKey = configKeys(configName)
+      if (configKey == null)
+        throw new InvalidConfigurationException(s"Unknown topic config name: $configName")
+      configKey.`type` == ConfigDef.Type.LIST
+    }
+
+    alterConfigOps.foreach { alterConfigOp =>
+      alterConfigOp.opType() match {
+        case OpType.SET => configProps.setProperty(alterConfigOp.configEntry().name(), alterConfigOp.configEntry().value())
+        case OpType.DELETE => configProps.remove(alterConfigOp.configEntry().name())
+        case OpType.APPEND => {
+          if (!listType(alterConfigOp.configEntry().name(), configKeys))
+            throw new InvalidRequestException(s"Config value append is not allowed for config key: ${alterConfigOp.configEntry().name()}")
+          val oldValueList = configProps.getProperty(alterConfigOp.configEntry().name()).split(",").toList
+          val newValueList = oldValueList ::: alterConfigOp.configEntry().value().split(",").toList
+          configProps.setProperty(alterConfigOp.configEntry().name(), newValueList.mkString(","))
+        }
+        case OpType.SUBTRACT => {
+          if (!listType(alterConfigOp.configEntry().name(), configKeys))
+            throw new InvalidRequestException(s"Config value subtract is not allowed for config key: ${alterConfigOp.configEntry().name()}")
+          val oldValueList = configProps.getProperty(alterConfigOp.configEntry().name()).split(",").toList
+          val newValueList = oldValueList.diff(alterConfigOp.configEntry().value().split(",").toList)
+          configProps.setProperty(alterConfigOp.configEntry().name(), newValueList.mkString(","))
+        }
+      }
+    }
+  }
+
+  def shutdown(): Unit = {
     topicPurgatory.shutdown()
     CoreUtils.swallow(createTopicPolicy.foreach(_.close()), this)
     CoreUtils.swallow(alterConfigPolicy.foreach(_.close()), this)
@@ -455,7 +633,10 @@ class AdminManager(val config: KafkaConfig,
   }
 
   private def configType(name: String, synonyms: List[String]): ConfigDef.Type = {
-    val configType = config.typeOf(name)
+    var configType = config.typeOf(name)
+    if (configType == null)
+      configType = DynamicConfig.Broker.typeOf(name)
+
     if (configType != null)
       configType
     else
@@ -515,7 +696,7 @@ class AdminManager(val config: KafkaConfig,
         .filter(perBrokerConfig || _.source == ConfigSource.DYNAMIC_DEFAULT_BROKER_CONFIG)
     val synonyms = if (!includeSynonyms) List.empty else allSynonyms
     val source = if (allSynonyms.isEmpty) ConfigSource.DEFAULT_CONFIG else allSynonyms.head.source
-    val readOnly = !allNames.exists(DynamicBrokerConfig.AllDynamicConfigs.contains)
+    val readOnly = !DynamicBrokerConfig.AllDynamicConfigs.contains(name)
     new DescribeConfigsResponse.ConfigEntry(name, valueAsString, source, isSensitive, readOnly, synonyms.asJava)
   }
 }
diff --git a/core/src/main/scala/kafka/server/BrokerMetadataCheckpoint.scala b/core/src/main/scala/kafka/server/BrokerMetadataCheckpoint.scala
index 2b915a0f44592..ffbcb5df05968 100755
--- a/core/src/main/scala/kafka/server/BrokerMetadataCheckpoint.scala
+++ b/core/src/main/scala/kafka/server/BrokerMetadataCheckpoint.scala
@@ -24,7 +24,13 @@ import java.util.Properties
 import kafka.utils._
 import org.apache.kafka.common.utils.Utils
 
-case class BrokerMetadata(brokerId: Int)
+case class BrokerMetadata(brokerId: Int,
+                          clusterId: Option[String]) {
+
+  override def toString: String  = {
+    s"BrokerMetadata(brokerId=$brokerId, clusterId=${clusterId.map(_.toString).getOrElse("None")})"
+  }
+}
 
 /**
   * This class saves broker's metadata to a file
@@ -38,6 +44,9 @@ class BrokerMetadataCheckpoint(val file: File) extends Logging {
         val brokerMetaProps = new Properties()
         brokerMetaProps.setProperty("version", 0.toString)
         brokerMetaProps.setProperty("broker.id", brokerMetadata.brokerId.toString)
+        brokerMetadata.clusterId.foreach { clusterId =>
+          brokerMetaProps.setProperty("cluster.id", clusterId)
+        }
         val temp = new File(file.getAbsolutePath + ".tmp")
         val fileOutputStream = new FileOutputStream(temp)
         try {
@@ -66,7 +75,8 @@ class BrokerMetadataCheckpoint(val file: File) extends Logging {
         version match {
           case 0 =>
             val brokerId = brokerMetaProps.getIntInRange("broker.id", (0, Int.MaxValue))
-            return Some(BrokerMetadata(brokerId))
+            val clusterId = Option(brokerMetaProps.getString("cluster.id", null))
+            return Some(BrokerMetadata(brokerId, clusterId))
           case _ =>
             throw new IOException("Unrecognized version of the server meta.properties file: " + version)
         }
diff --git a/core/src/main/scala/kafka/server/BrokerStates.scala b/core/src/main/scala/kafka/server/BrokerStates.scala
index 2b66beb6328a5..f53ed833c22e9 100644
--- a/core/src/main/scala/kafka/server/BrokerStates.scala
+++ b/core/src/main/scala/kafka/server/BrokerStates.scala
@@ -67,12 +67,12 @@ case object BrokerShuttingDown extends BrokerStates { val state: Byte = 7 }
 case class BrokerState() {
   @volatile var currentState: Byte = NotRunning.state
 
-  def newState(newState: BrokerStates) {
+  def newState(newState: BrokerStates): Unit = {
     this.newState(newState.state)
   }
 
   // Allowing undefined custom state
-  def newState(newState: Byte) {
+  def newState(newState: Byte): Unit = {
     currentState = newState
   }
 }
diff --git a/core/src/main/scala/kafka/server/ClientQuotaManager.scala b/core/src/main/scala/kafka/server/ClientQuotaManager.scala
index 9817f2104c8b8..eda7ac3f49a72 100644
--- a/core/src/main/scala/kafka/server/ClientQuotaManager.scala
+++ b/core/src/main/scala/kafka/server/ClientQuotaManager.scala
@@ -23,11 +23,11 @@ import java.util.concurrent.locks.ReentrantReadWriteLock
 import kafka.network.RequestChannel
 import kafka.network.RequestChannel._
 import kafka.server.ClientQuotaManager._
-import kafka.utils.{Logging, ShutdownableThread}
+import kafka.utils.{KafkaScheduler, Logging, ShutdownableThread}
 import org.apache.kafka.common.{Cluster, MetricName}
 import org.apache.kafka.common.metrics._
 import org.apache.kafka.common.metrics.Metrics
-import org.apache.kafka.common.metrics.stats.{Avg, Rate, Total}
+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.quota.{ClientQuotaCallback, ClientQuotaEntity, ClientQuotaType}
@@ -62,8 +62,8 @@ object ClientQuotaManagerConfig {
   // Always have 10 whole windows + 1 current window
   val DefaultNumQuotaSamples = 11
   val DefaultQuotaWindowSizeSeconds = 1
-  // Purge sensors after 1 hour of inactivity
-  val InactiveSensorExpirationTimeSeconds  = 3600
+  // Do not expire quota sensors
+  val InactiveSensorExpirationTimeSeconds  = Int.MaxValue
   val QuotaRequestPercentDefault = Int.MaxValue.toDouble
   val NanosToPercentagePerSecond = 100.0 / TimeUnit.SECONDS.toNanos(1)
 
@@ -160,6 +160,7 @@ class ClientQuotaManager(private val config: ClientQuotaManagerConfig,
                          private val metrics: Metrics,
                          private val quotaType: QuotaType,
                          private val time: Time,
+                         private val schedulerOpt: Option[KafkaScheduler],
                          threadNamePrefix: String,
                          clientQuotaCallback: Option[ClientQuotaCallback] = None) extends Logging {
   private val staticConfigClientIdQuota = Quota.upperBound(config.quotaBytesPerSecondDefault)
@@ -176,15 +177,25 @@ class ClientQuotaManager(private val config: ClientQuotaManagerConfig,
   private[server] val throttledChannelReaper = new ThrottledChannelReaper(delayQueue, threadNamePrefix)
   private val quotaCallback = clientQuotaCallback.getOrElse(new DefaultQuotaCallback)
 
-  private val delayQueueSensor = metrics.sensor(quotaType + "-delayQueue")
+  private val delayQueueSensor = metrics.sensor(quotaType + "-delayQueueSize")
   delayQueueSensor.add(metrics.metricName("queue-size",
     quotaType.toString,
-    "Tracks the size of the delay queue"), new Total())
+    "Tracks the size of the delay queue"), new CumulativeSum())
   start() // Use start method to keep spotbugs happy
-  private def start() {
+  private def start(): Unit = {
     throttledChannelReaper.start()
+    schedulerOpt match {
+      case Some(scheduler) =>
+        scheduler.schedule("quota-metrics-logger-%s".format(quotaType), logQuotaMetrics, 60, 60, TimeUnit.SECONDS)
+      case _ =>
+    }
   }
 
+  private val throttledRequestCountSensor = metrics.sensor(quotaType + "-throttledRequestCount")
+  throttledRequestCountSensor.add(metrics.metricName("throttle-count",
+                                                     quotaType.toString,
+                                                     "Tracks the number of requests that have been throttled"), new CumulativeSum())
+
   /**
    * Reaper thread that triggers channel unmute callbacks on all throttled channels
    * @param delayQueue DelayQueue to dequeue from
@@ -203,6 +214,30 @@ class ClientQuotaManager(private val config: ClientQuotaManagerConfig,
     }
   }
 
+  private def isQuotaMetric(metricName : MetricName) : Boolean = {
+    metricName.group().equals(quotaType.toString) &&
+    (metricName.name().equals("byte-rate") || metricName.name().equals("throttle-time") ||
+    metricName.name().equals("request-time")) &&
+    (metricName.tags().containsKey("client-id") || metricName.tags().containsKey("user"))
+  }
+
+  def logQuotaMetrics() : Unit = {
+    import scala.collection.JavaConverters._
+    val metricsMap = metrics.metrics().asScala;
+    metricsMap.foreach {
+      case (metricName: MetricName, kafkaMetric: KafkaMetric) =>
+       if (isQuotaMetric(metricName)) {
+         val quotaValueStr = kafkaMetric.metricValue().toString()
+         if (!quotaValueStr.equals("NaN") && !quotaValueStr.equals("0.0")) {
+           val metricConfig = kafkaMetric.config()
+           val quota = metricConfig.quota();
+           val quotaBoundsStr = if (quota == null) "not configured" else quota.bound().toString()
+           info("Metric name (" + metricName + ") has value (" + quotaValueStr + ") with bound (" + quotaBoundsStr + ")")
+         }
+       }
+    }
+  }
+
   /**
    * Returns true if any quotas are enabled for this quota manager. This is used
    * to determine if quota related metrics should be created.
@@ -233,6 +268,22 @@ class ClientQuotaManager(private val config: ClientQuotaManagerConfig,
     }
   }
 
+  /**
+   * Returns maximum value (produced/consume bytes or request processing time) that could be recorded without guaranteed throttling.
+   * Recording any larger value will always be throttled, even if no other values were recorded in the quota window.
+   * This is used for deciding the maximum bytes that can be fetched at once
+   */
+  def getMaxValueInQuotaWindow(session: Session, clientId: String): Double = {
+    if (quotasEnabled) {
+      val clientSensors = getOrCreateQuotaSensors(session, clientId)
+      Option(quotaCallback.quotaLimit(clientQuotaType, clientSensors.metricTags.asJava))
+        .map(_.toDouble * (config.numQuotaSamples - 1) * config.quotaWindowSizeSeconds)
+        .getOrElse(Double.MaxValue)
+    } else {
+      Double.MaxValue
+    }
+  }
+
   def recordAndGetThrottleTimeMs(session: Session, clientId: String, value: Double, timeMs: Long): Int = {
     var throttleTimeMs = 0
     val clientSensors = getOrCreateQuotaSensors(session, clientId)
@@ -243,7 +294,7 @@ class ClientQuotaManager(private val config: ClientQuotaManagerConfig,
         // Compute the delay
         val clientMetric = metrics.metrics().get(clientRateMetricName(clientSensors.metricTags))
         throttleTimeMs = throttleTime(clientMetric).toInt
-        debug("Quota violated for sensor (%s). Delay time: (%d)".format(clientSensors.quotaSensor.name(), throttleTimeMs))
+        info("Quota violated for sensor (%s). Delay time: (%d)".format(clientSensors.quotaSensor.name(), throttleTimeMs))
     }
     throttleTimeMs
   }
@@ -276,6 +327,7 @@ class ClientQuotaManager(private val config: ClientQuotaManagerConfig,
       val throttledChannel = new ThrottledChannel(request, time, throttleTimeMs, channelThrottlingCallback)
       delayQueue.add(throttledChannel)
       delayQueueSensor.record()
+      throttledRequestCountSensor.record()
       debug("Channel throttled for sensor (%s). Delay time: (%d)".format(clientSensors.quotaSensor.name(), throttleTimeMs))
     }
   }
@@ -285,7 +337,7 @@ class ClientQuotaManager(private val config: ClientQuotaManagerConfig,
    * quota violation. The aggregate value will subsequently be used for throttling when the
    * next request is processed.
    */
-  def recordNoThrottle(clientSensors: ClientSensors, value: Double) {
+  def recordNoThrottle(clientSensors: ClientSensors, value: Double): Unit = {
     clientSensors.quotaSensor.record(value, time.milliseconds(), false)
   }
 
@@ -310,7 +362,7 @@ class ClientQuotaManager(private val config: ClientQuotaManagerConfig,
   }
 
   private def quotaLimit(metricTags: util.Map[String, String]): Double = {
-    Option(quotaCallback.quotaLimit(clientQuotaType, metricTags)).map(_.toDouble)getOrElse(Long.MaxValue)
+    Option(quotaCallback.quotaLimit(clientQuotaType, metricTags)).map(_.toDouble).getOrElse(Long.MaxValue)
   }
 
   /*
@@ -410,7 +462,7 @@ class ClientQuotaManager(private val config: ClientQuotaManagerConfig,
    * @param sanitizedClientId sanitized client ID to override if quota applies to  or 
    * @param quota             custom quota to apply or None if quota override is being removed
    */
-  def updateQuota(sanitizedUser: Option[String], clientId: Option[String], sanitizedClientId: Option[String], quota: Option[Quota]) {
+  def updateQuota(sanitizedUser: Option[String], clientId: Option[String], sanitizedClientId: Option[String], quota: Option[Quota]): Unit = {
     /*
      * Acquire the write lock to apply changes in the quota objects.
      * This method changes the quota in the overriddenQuota map and applies the update on the actual KafkaMetric object (if it exists).
@@ -483,7 +535,7 @@ class ClientQuotaManager(private val config: ClientQuotaManagerConfig,
       // Change the underlying metric config if the sensor has been created
       val metric = allMetrics.get(quotaMetricName)
       if (metric != null) {
-        Option(quotaCallback.quotaLimit(clientQuotaType, metricTags.asJava)).foreach { newQuota =>
+        Option(quotaLimit(metricTags.asJava)).foreach { newQuota =>
           info(s"Sensor for $quotaEntity already exists. Changing quota to $newQuota in MetricConfig")
           metric.config(getQuotaMetricConfig(newQuota))
         }
@@ -493,8 +545,7 @@ class ClientQuotaManager(private val config: ClientQuotaManagerConfig,
       allMetrics.asScala.filterKeys(n => n.name == quotaMetricName.name && n.group == quotaMetricName.group).foreach {
         case (metricName, metric) =>
           val metricTags = metricName.tags
-          Option(quotaCallback.quotaLimit(clientQuotaType, metricTags)).foreach { quota =>
-            val newQuota = quota.asInstanceOf[Double]
+          Option(quotaLimit(metricTags)).foreach { newQuota =>
             if (newQuota != metric.config.quota.bound) {
               info(s"Sensor for quota-id $metricTags already exists. Setting quota to $newQuota in MetricConfig")
               metric.config(getQuotaMetricConfig(newQuota))
diff --git a/core/src/main/scala/kafka/server/ClientRequestQuotaManager.scala b/core/src/main/scala/kafka/server/ClientRequestQuotaManager.scala
index 7bab61605bddb..d2e655189535b 100644
--- a/core/src/main/scala/kafka/server/ClientRequestQuotaManager.scala
+++ b/core/src/main/scala/kafka/server/ClientRequestQuotaManager.scala
@@ -19,6 +19,7 @@ package kafka.server
 import java.util.concurrent.TimeUnit
 
 import kafka.network.RequestChannel
+import kafka.utils.KafkaScheduler
 import org.apache.kafka.common.MetricName
 import org.apache.kafka.common.metrics._
 import org.apache.kafka.common.utils.Time
@@ -30,13 +31,14 @@ import scala.collection.JavaConverters._
 class ClientRequestQuotaManager(private val config: ClientQuotaManagerConfig,
                                 private val metrics: Metrics,
                                 private val time: Time,
+                                private val schedulerOpt: Option[KafkaScheduler],
                                 threadNamePrefix: String,
                                 quotaCallback: Option[ClientQuotaCallback])
-                                extends ClientQuotaManager(config, metrics, QuotaType.Request, time, threadNamePrefix, quotaCallback) {
+                                extends ClientQuotaManager(config, metrics, QuotaType.Request, time, schedulerOpt, threadNamePrefix, quotaCallback) {
   val maxThrottleTimeMs = TimeUnit.SECONDS.toMillis(this.config.quotaWindowSizeSeconds)
   def exemptSensor = getOrCreateSensor(exemptSensorName, exemptMetricName)
 
-  def recordExempt(value: Double) {
+  def recordExempt(value: Double): Unit = {
     exemptSensor.record(value)
   }
 
diff --git a/core/src/main/scala/kafka/server/ConfigHandler.scala b/core/src/main/scala/kafka/server/ConfigHandler.scala
index 5593225f0e6eb..2395dcde79db8 100644
--- a/core/src/main/scala/kafka/server/ConfigHandler.scala
+++ b/core/src/main/scala/kafka/server/ConfigHandler.scala
@@ -34,13 +34,14 @@ import org.apache.kafka.common.metrics.Quota._
 import org.apache.kafka.common.utils.Sanitizer
 
 import scala.collection.JavaConverters._
+import scala.collection.Seq
 import scala.util.Try
 
 /**
   * The ConfigHandler is used to process config change notifications received by the DynamicConfigManager
   */
 trait ConfigHandler {
-  def processConfigChanges(entityName: String, value: Properties)
+  def processConfigChanges(entityName: String, value: Properties): Unit
 }
 
 /**
@@ -49,11 +50,11 @@ trait ConfigHandler {
   */
 class TopicConfigHandler(private val logManager: LogManager, kafkaConfig: KafkaConfig, val quotas: QuotaManagers, kafkaController: KafkaController) extends ConfigHandler with Logging  {
 
-  def processConfigChanges(topic: String, topicConfig: Properties) {
-    // Validate the configurations.
-    val configNamesToExclude = excludedConfigs(topic, topicConfig)
-
-    val logs = logManager.logsByTopic(topic).toBuffer
+  private def updateLogConfig(topic: String,
+                              topicConfig: Properties,
+                              configNamesToExclude: Set[String]): Unit = {
+    logManager.topicConfigUpdated(topic)
+    val logs = logManager.logsByTopic(topic)
     if (logs.nonEmpty) {
       /* combine the default properties with the overrides in zk to create the new LogConfig */
       val props = new Properties()
@@ -61,8 +62,15 @@ class TopicConfigHandler(private val logManager: LogManager, kafkaConfig: KafkaC
         if (!configNamesToExclude.contains(key)) props.put(key, value)
       }
       val logConfig = LogConfig.fromProps(logManager.currentDefaultConfig.originals, props)
-      logs.foreach(_.updateConfig(topicConfig.asScala.keySet, logConfig))
+      logs.foreach(_.updateConfig(logConfig))
     }
+  }
+
+  def processConfigChanges(topic: String, topicConfig: Properties): Unit = {
+    // Validate the configurations.
+    val configNamesToExclude = excludedConfigs(topic, topicConfig)
+
+    updateLogConfig(topic, topicConfig, configNamesToExclude)
 
     def updateThrottledList(prop: String, quotaManager: ReplicationQuotaManager) = {
       if (topicConfig.containsKey(prop) && topicConfig.getProperty(prop).length > 0) {
@@ -116,7 +124,7 @@ class TopicConfigHandler(private val logManager: LogManager, kafkaConfig: KafkaC
  */
 class QuotaConfigHandler(private val quotaManagers: QuotaManagers) {
 
-  def updateQuotaConfig(sanitizedUser: Option[String], sanitizedClientId: Option[String], config: Properties) {
+  def updateQuotaConfig(sanitizedUser: Option[String], sanitizedClientId: Option[String], config: Properties): Unit = {
     val clientId = sanitizedClientId.map(Sanitizer.desanitize)
     val producerQuota =
       if (config.containsKey(DynamicConfig.Client.ProducerByteRateOverrideProp))
@@ -145,7 +153,7 @@ class QuotaConfigHandler(private val quotaManagers: QuotaManagers) {
  */
 class ClientIdConfigHandler(private val quotaManagers: QuotaManagers) extends QuotaConfigHandler(quotaManagers) with ConfigHandler {
 
-  def processConfigChanges(sanitizedClientId: String, clientConfig: Properties) {
+  def processConfigChanges(sanitizedClientId: String, clientConfig: Properties): Unit = {
     updateQuotaConfig(None, Some(sanitizedClientId), clientConfig)
   }
 }
@@ -157,7 +165,7 @@ class ClientIdConfigHandler(private val quotaManagers: QuotaManagers) extends Qu
  */
 class UserConfigHandler(private val quotaManagers: QuotaManagers, val credentialProvider: CredentialProvider) extends QuotaConfigHandler(quotaManagers) with ConfigHandler {
 
-  def processConfigChanges(quotaEntityPath: String, config: Properties) {
+  def processConfigChanges(quotaEntityPath: String, config: Properties): Unit = {
     // Entity path is  or /clients/
     val entities = quotaEntityPath.split("/")
     if (entities.length != 1 && entities.length != 3)
@@ -178,7 +186,7 @@ class UserConfigHandler(private val quotaManagers: QuotaManagers, val credential
 class BrokerConfigHandler(private val brokerConfig: KafkaConfig,
                           private val quotaManagers: QuotaManagers) extends ConfigHandler with Logging {
 
-  def processConfigChanges(brokerId: String, properties: Properties) {
+  def processConfigChanges(brokerId: String, properties: Properties): Unit = {
     def getOrDefault(prop: String): Long = {
       if (properties.containsKey(prop))
         properties.getProperty(prop).toLong
@@ -205,7 +213,7 @@ object ThrottledReplicaListValidator extends Validator {
       if (!(proposed.forall(_.toString.trim.matches("([0-9]+:[0-9]+)?"))
         || proposed.headOption.exists(_.toString.trim.equals("*"))))
         throw new ConfigException(name, value,
-          s"$name must be the literal '*' or a list of replicas in the following format: [partitionId],[brokerId]:[partitionId],[brokerId]:...")
+          s"$name must be the literal '*' or a list of replicas in the following format: [partitionId]:[brokerId],[partitionId]:[brokerId],...")
     }
     value match {
       case scalaSeq: Seq[_] => check(scalaSeq)
@@ -214,6 +222,6 @@ object ThrottledReplicaListValidator extends Validator {
     }
   }
 
-  override def toString: String = "[partitionId],[brokerId]:[partitionId],[brokerId]:..."
+  override def toString: String = "[partitionId]:[brokerId],[partitionId]:[brokerId],..."
 
 }
diff --git a/core/src/main/scala/kafka/server/DelayedCreatePartitions.scala b/core/src/main/scala/kafka/server/DelayedCreatePartitions.scala
index 0a9948311af0c..8542bc2a2ecee 100644
--- a/core/src/main/scala/kafka/server/DelayedCreatePartitions.scala
+++ b/core/src/main/scala/kafka/server/DelayedCreatePartitions.scala
@@ -61,7 +61,7 @@ class DelayedCreatePartitions(delayMs: Long,
   /**
     * Check for partitions that are still missing a leader, update their error code and call the responseCallback
     */
-  override def onComplete() {
+  override def onComplete(): Unit = {
     trace(s"Completing operation for $createMetadata")
     val results = createMetadata.map { metadata =>
       // ignore topics that already have errors
@@ -83,6 +83,6 @@ class DelayedCreatePartitions(delayMs: Long,
 
   private def isMissingLeader(topic: String, partition: Int): Boolean = {
     val partitionInfo = adminManager.metadataCache.getPartitionInfo(topic, partition)
-    partitionInfo.isEmpty || partitionInfo.get.basePartitionState.leader == LeaderAndIsr.NoLeader
+    partitionInfo.forall(_.leader == LeaderAndIsr.NoLeader)
   }
 }
diff --git a/core/src/main/scala/kafka/server/DelayedDeleteRecords.scala b/core/src/main/scala/kafka/server/DelayedDeleteRecords.scala
index a977d9a371d01..e83ce0a8af44d 100644
--- a/core/src/main/scala/kafka/server/DelayedDeleteRecords.scala
+++ b/core/src/main/scala/kafka/server/DelayedDeleteRecords.scala
@@ -21,8 +21,8 @@ package kafka.server
 import java.util.concurrent.TimeUnit
 
 import kafka.metrics.KafkaMetricsGroup
-import org.apache.kafka.common.protocol.Errors
 import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.protocol.Errors
 import org.apache.kafka.common.requests.DeleteRecordsResponse
 
 import scala.collection._
@@ -73,19 +73,19 @@ class DelayedDeleteRecords(delayMs: Long,
       // skip those partitions that have already been satisfied
       if (status.acksPending) {
         val (lowWatermarkReached, error, lw) = replicaManager.getPartition(topicPartition) match {
-          case Some(partition) =>
-            if (partition eq ReplicaManager.OfflinePartition) {
-              (false, Errors.KAFKA_STORAGE_ERROR, DeleteRecordsResponse.INVALID_LOW_WATERMARK)
-            } else {
-              partition.leaderReplicaIfLocal match {
-                case Some(_) =>
-                  val leaderLW = partition.lowWatermarkIfLeader
-                  (leaderLW >= status.requiredOffset, Errors.NONE, leaderLW)
-                case None =>
-                  (false, Errors.NOT_LEADER_FOR_PARTITION, DeleteRecordsResponse.INVALID_LOW_WATERMARK)
-              }
+          case HostedPartition.Online(partition) =>
+            partition.leaderLogIfLocal match {
+              case Some(_) =>
+                val leaderLW = partition.lowWatermarkIfLeader
+                (leaderLW >= status.requiredOffset, Errors.NONE, leaderLW)
+              case None =>
+                (false, Errors.NOT_LEADER_FOR_PARTITION, DeleteRecordsResponse.INVALID_LOW_WATERMARK)
             }
-          case None =>
+
+          case HostedPartition.Offline =>
+            (false, Errors.KAFKA_STORAGE_ERROR, DeleteRecordsResponse.INVALID_LOW_WATERMARK)
+
+          case HostedPartition.None =>
             (false, Errors.UNKNOWN_TOPIC_OR_PARTITION, DeleteRecordsResponse.INVALID_LOW_WATERMARK)
         }
         if (error != Errors.NONE || lowWatermarkReached) {
@@ -103,7 +103,7 @@ class DelayedDeleteRecords(delayMs: Long,
       false
   }
 
-  override def onExpiration() {
+  override def onExpiration(): Unit = {
     deleteRecordsStatus.foreach { case (topicPartition, status) =>
       if (status.acksPending) {
         DelayedDeleteRecordsMetrics.recordExpiration(topicPartition)
@@ -114,8 +114,8 @@ class DelayedDeleteRecords(delayMs: Long,
   /**
    * Upon completion, return the current response status along with the error code per partition
    */
-  override def onComplete() {
-    val responseStatus = deleteRecordsStatus.mapValues(status => status.responseStatus)
+  override def onComplete(): Unit = {
+    val responseStatus = deleteRecordsStatus.map { case (k, status) => k -> status.responseStatus }
     responseCallback(responseStatus)
   }
 }
@@ -124,7 +124,7 @@ object DelayedDeleteRecordsMetrics extends KafkaMetricsGroup {
 
   private val aggregateExpirationMeter = newMeter("ExpiresPerSec", "requests", TimeUnit.SECONDS)
 
-  def recordExpiration(partition: TopicPartition) {
+  def recordExpiration(partition: TopicPartition): Unit = {
     aggregateExpirationMeter.mark()
   }
 }
diff --git a/core/src/main/scala/kafka/server/DelayedDeleteTopics.scala b/core/src/main/scala/kafka/server/DelayedDeleteTopics.scala
index 95d6f50515c62..bb15a27948c73 100644
--- a/core/src/main/scala/kafka/server/DelayedDeleteTopics.scala
+++ b/core/src/main/scala/kafka/server/DelayedDeleteTopics.scala
@@ -57,7 +57,7 @@ class DelayedDeleteTopics(delayMs: Long,
   /**
     * Check for partitions that still exist, update their error code and call the responseCallback
     */
-  override def onComplete() {
+  override def onComplete(): Unit = {
     trace(s"Completing operation for $deleteMetadata")
     val results = deleteMetadata.map { metadata =>
       // ignore topics that already have errors
diff --git a/core/src/main/scala/kafka/server/DelayedElectPreferredLeader.scala b/core/src/main/scala/kafka/server/DelayedElectLeader.scala
similarity index 63%
rename from core/src/main/scala/kafka/server/DelayedElectPreferredLeader.scala
rename to core/src/main/scala/kafka/server/DelayedElectLeader.scala
index f3543a89518ea..d599f98411b3c 100644
--- a/core/src/main/scala/kafka/server/DelayedElectPreferredLeader.scala
+++ b/core/src/main/scala/kafka/server/DelayedElectLeader.scala
@@ -23,18 +23,19 @@ import org.apache.kafka.common.requests.ApiError
 
 import scala.collection.{Map, mutable}
 
-/** A delayed elect preferred leader operation that can be created by the replica manager and watched
-  * in the elect preferred leader purgatory
+/** A delayed elect leader operation that can be created by the replica manager and watched
+  * in the elect leader purgatory
   */
-class DelayedElectPreferredLeader(delayMs: Long,
-                                  expectedLeaders: Map[TopicPartition, Int],
-                                  results: Map[TopicPartition, ApiError],
-                                  replicaManager: ReplicaManager,
-                                  responseCallback: Map[TopicPartition, ApiError] => Unit)
-    extends DelayedOperation(delayMs) {
+class DelayedElectLeader(
+  delayMs: Long,
+  expectedLeaders: Map[TopicPartition, Int],
+  results: Map[TopicPartition, ApiError],
+  replicaManager: ReplicaManager,
+  responseCallback: Map[TopicPartition, ApiError] => Unit
+) extends DelayedOperation(delayMs) {
 
-  var waitingPartitions = expectedLeaders
-  val fullResults = results.to[mutable.Set]
+  private var waitingPartitions = expectedLeaders
+  private val fullResults = mutable.Map() ++= results
 
 
   /**
@@ -49,10 +50,10 @@ class DelayedElectPreferredLeader(delayMs: Long,
   override def onComplete(): Unit = {
     // This could be called to force complete, so I need the full list of partitions, so I can time them all out.
     updateWaiting()
-    val timedout = waitingPartitions.map{
-      case (tp, leader) => tp -> new ApiError(Errors.REQUEST_TIMED_OUT, null)
+    val timedOut = waitingPartitions.map {
+      case (tp, _) => tp -> new ApiError(Errors.REQUEST_TIMED_OUT, null)
     }.toMap
-    responseCallback(timedout ++ fullResults)
+    responseCallback(timedOut ++ fullResults)
   }
 
   /**
@@ -68,17 +69,14 @@ class DelayedElectPreferredLeader(delayMs: Long,
     waitingPartitions.isEmpty && forceComplete()
   }
 
-  private def updateWaiting() = {
-    waitingPartitions.foreach{case (tp, leader) =>
-      val ps = replicaManager.metadataCache.getPartitionInfo(tp.topic, tp.partition)
-      ps match {
-        case Some(ps) =>
-          if (leader == ps.basePartitionState.leader) {
-            waitingPartitions -= tp
-            fullResults += tp -> ApiError.NONE
-          }
-        case None =>
-      }
+  private def updateWaiting(): Unit = {
+    val metadataCache = replicaManager.metadataCache
+    val completedPartitions = waitingPartitions.collect {
+      case (tp, leader) if metadataCache.getPartitionInfo(tp.topic, tp.partition).exists(_.leader == leader) => tp
+    }
+    completedPartitions.foreach  { tp =>
+      waitingPartitions -= tp
+      fullResults += tp -> ApiError.NONE
     }
   }
 
diff --git a/core/src/main/scala/kafka/server/DelayedFetch.scala b/core/src/main/scala/kafka/server/DelayedFetch.scala
index 902009917592f..3cf55d7b8f242 100644
--- a/core/src/main/scala/kafka/server/DelayedFetch.scala
+++ b/core/src/main/scala/kafka/server/DelayedFetch.scala
@@ -22,14 +22,18 @@ import java.util.concurrent.TimeUnit
 import kafka.metrics.KafkaMetricsGroup
 import org.apache.kafka.common.TopicPartition
 import org.apache.kafka.common.errors._
+import org.apache.kafka.common.replica.ClientMetadata
 import org.apache.kafka.common.requests.FetchRequest.PartitionData
 
 import scala.collection._
 
 case class FetchPartitionStatus(startOffsetMetadata: LogOffsetMetadata, fetchInfo: PartitionData) {
 
-  override def toString = "[startOffsetMetadata: " + startOffsetMetadata + ", " +
-                          "fetchInfo: " + fetchInfo + "]"
+  override def toString: String = {
+    "[startOffsetMetadata: " + startOffsetMetadata +
+      ", fetchInfo: " + fetchInfo +
+      "]"
+  }
 }
 
 /**
@@ -59,6 +63,7 @@ class DelayedFetch(delayMs: Long,
                    fetchMetadata: FetchMetadata,
                    replicaManager: ReplicaManager,
                    quota: ReplicaQuota,
+                   clientMetadata: Option[ClientMetadata],
                    responseCallback: Seq[(TopicPartition, FetchPartitionData)] => Unit)
   extends DelayedOperation(delayMs) {
 
@@ -66,12 +71,13 @@ class DelayedFetch(delayMs: Long,
    * The operation can be completed if:
    *
    * Case A: This broker is no longer the leader for some partitions it tries to fetch
-   * Case B: This broker does not know of some partitions it tries to fetch
-   * Case C: The fetch offset locates not on the last segment of the log
-   * Case D: The accumulated bytes from all the fetching partitions exceeds the minimum bytes
-   * Case E: The partition is in an offline log directory on this broker
-   * Case F: This broker is the leader, but the requested epoch is now fenced
-   *
+   * Case B: The replica is no longer available on this broker
+   * Case C: This broker does not know of some partitions it tries to fetch
+   * Case D: The partition is in an offline log directory on this broker
+   * Case E: This broker is the leader, but the requested epoch is now fenced
+   * Case F: The fetch offset locates not on the last segment of the log
+   * Case G: The accumulated bytes from all the fetching partitions exceeds the minimum bytes
+   * Case H: The high watermark on this broker has changed within a FetchSession, need to propagate to follower (KIP-392)
    * Upon completion, should return whatever data is available for each valid partition
    */
   override def tryComplete(): Boolean = {
@@ -92,16 +98,16 @@ class DelayedFetch(delayMs: Long,
               case FetchTxnCommitted => offsetSnapshot.lastStableOffset
             }
 
-            // Go directly to the check for Case D if the message offsets are the same. If the log segment
+            // Go directly to the check for Case G if the message offsets are the same. If the log segment
             // has just rolled, then the high watermark offset will remain the same but be on the old segment,
-            // which would incorrectly be seen as an instance of Case C.
+            // which would incorrectly be seen as an instance of Case F.
             if (endOffset.messageOffset != fetchOffset.messageOffset) {
               if (endOffset.onOlderSegment(fetchOffset)) {
-                // Case C, this can happen when the new fetch operation is on a truncated leader
+                // Case F, this can happen when the new fetch operation is on a truncated leader
                 debug(s"Satisfying fetch $fetchMetadata since it is fetching later segments of partition $topicPartition.")
                 return forceComplete()
               } else if (fetchOffset.onOlderSegment(endOffset)) {
-                // Case C, this can happen when the fetch operation is falling behind the current segment
+                // Case F, this can happen when the fetch operation is falling behind the current segment
                 // or the partition has just rolled a new segment
                 debug(s"Satisfying fetch $fetchMetadata immediately since it is fetching older segments.")
                 // We will not force complete the fetch request if a replica should be throttled.
@@ -114,32 +120,43 @@ class DelayedFetch(delayMs: Long,
                   accumulatedSize += bytesAvailable
               }
             }
+
+            if (fetchMetadata.isFromFollower) {
+              // Case H check if the follower has the latest HW from the leader
+              if (partition.getReplica(fetchMetadata.replicaId)
+                .exists(r => offsetSnapshot.highWatermark.messageOffset > r.lastSentHighWatermark)) {
+                return forceComplete()
+              }
+            }
           }
         } catch {
-          case _: KafkaStorageException => // Case E
-            debug(s"Partition $topicPartition is in an offline log directory, satisfy $fetchMetadata immediately")
+          case _: NotLeaderForPartitionException =>  // Case A
+            debug(s"Broker is no longer the leader of $topicPartition, satisfy $fetchMetadata immediately")
+            return forceComplete()
+          case _: ReplicaNotAvailableException =>  // Case B
+            debug(s"Broker no longer has a replica of $topicPartition, satisfy $fetchMetadata immediately")
             return forceComplete()
-          case _: UnknownTopicOrPartitionException => // Case B
+          case _: UnknownTopicOrPartitionException => // Case C
             debug(s"Broker no longer knows of partition $topicPartition, satisfy $fetchMetadata immediately")
             return forceComplete()
-          case _: FencedLeaderEpochException => // Case F
+          case _: KafkaStorageException => // Case D
+            debug(s"Partition $topicPartition is in an offline log directory, satisfy $fetchMetadata immediately")
+            return forceComplete()
+          case _: FencedLeaderEpochException => // Case E
             debug(s"Broker is the leader of partition $topicPartition, but the requested epoch " +
               s"$fetchLeaderEpoch is fenced by the latest leader epoch, satisfy $fetchMetadata immediately")
             return forceComplete()
-          case _: NotLeaderForPartitionException =>  // Case A
-            debug("Broker is no longer the leader of %s, satisfy %s immediately".format(topicPartition, fetchMetadata))
-            return forceComplete()
         }
     }
 
-    // Case D
+    // Case G
     if (accumulatedSize >= fetchMetadata.fetchMinBytes)
        forceComplete()
     else
       false
   }
 
-  override def onExpiration() {
+  override def onExpiration(): Unit = {
     if (fetchMetadata.isFromFollower)
       DelayedFetchMetrics.followerExpiredRequestMeter.mark()
     else
@@ -149,7 +166,7 @@ class DelayedFetch(delayMs: Long,
   /**
    * Upon completion, read whatever data is available and pass to the complete callback
    */
-  override def onComplete() {
+  override def onComplete(): Unit = {
     val logReadResults = replicaManager.readFromLocalLog(
       replicaId = fetchMetadata.replicaId,
       fetchOnlyFromLeader = fetchMetadata.fetchOnlyLeader,
@@ -157,11 +174,12 @@ class DelayedFetch(delayMs: Long,
       fetchMaxBytes = fetchMetadata.fetchMaxBytes,
       hardMaxBytesLimit = fetchMetadata.hardMaxBytesLimit,
       readPartitionInfo = fetchMetadata.fetchPartitionStatus.map { case (tp, status) => tp -> status.fetchInfo },
+      clientMetadata = clientMetadata,
       quota = quota)
 
     val fetchPartitionData = logReadResults.map { case (tp, result) =>
       tp -> FetchPartitionData(result.error, result.highWatermark, result.leaderLogStartOffset, result.info.records,
-        result.lastStableOffset, result.info.abortedTransactions)
+        result.lastStableOffset, result.info.abortedTransactions, result.preferredReadReplica)
     }
 
     responseCallback(fetchPartitionData)
diff --git a/core/src/main/scala/kafka/server/DelayedFuture.scala b/core/src/main/scala/kafka/server/DelayedFuture.scala
new file mode 100644
index 0000000000000..f7e1e3fbf3d98
--- /dev/null
+++ b/core/src/main/scala/kafka/server/DelayedFuture.scala
@@ -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 kafka.server
+
+import java.util.concurrent._
+import java.util.function.BiConsumer
+
+import org.apache.kafka.common.errors.TimeoutException
+import org.apache.kafka.common.utils.KafkaThread
+
+import scala.collection.Seq
+
+/**
+  * A delayed operation using CompletionFutures that can be created by KafkaApis and watched
+  * in a DelayedFuturePurgatory purgatory. This is used for ACL updates using async Authorizers.
+  */
+class DelayedFuture[T](timeoutMs: Long,
+                       futures: List[CompletableFuture[T]],
+                       responseCallback: () => Unit)
+  extends DelayedOperation(timeoutMs) {
+
+  /**
+   * The operation can be completed if all the futures have completed successfully
+   * or failed with exceptions.
+   */
+  override def tryComplete() : Boolean = {
+    trace(s"Trying to complete operation for ${futures.size} futures")
+
+    val pending = futures.count(future => !future.isDone)
+    if (pending == 0) {
+      trace("All futures have been completed or have errors, completing the delayed operation")
+      forceComplete()
+    } else {
+      trace(s"$pending future still pending, not completing the delayed operation")
+      false
+    }
+  }
+
+  /**
+   * Timeout any pending futures and invoke responseCallback. This is invoked when all
+   * futures have completed or the operation has timed out.
+   */
+  override def onComplete(): Unit = {
+    val pendingFutures = futures.filterNot(_.isDone)
+    trace(s"Completing operation for ${futures.size} futures, expired ${pendingFutures.size}")
+    pendingFutures.foreach(_.completeExceptionally(new TimeoutException(s"Request has been timed out after $timeoutMs ms")))
+    responseCallback.apply()
+  }
+
+  /**
+   * This is invoked after onComplete(), so no actions required.
+   */
+  override def onExpiration(): Unit = {
+  }
+}
+
+class DelayedFuturePurgatory(purgatoryName: String, brokerId: Int) {
+  private val purgatory = DelayedOperationPurgatory[DelayedFuture[_]](purgatoryName, brokerId)
+  private val executor = new ThreadPoolExecutor(1, 1, 0, TimeUnit.MILLISECONDS,
+    new LinkedBlockingQueue[Runnable](),
+    new ThreadFactory {
+      override def newThread(r: Runnable): Thread = new KafkaThread(s"DelayedExecutor-$purgatoryName", r, true)
+    })
+  val purgatoryKey = new Object
+
+  def tryCompleteElseWatch[T](timeoutMs: Long,
+                              futures: List[CompletableFuture[T]],
+                              responseCallback: () => Unit): DelayedFuture[T] = {
+    val delayedFuture = new DelayedFuture[T](timeoutMs, futures, responseCallback)
+    val done = purgatory.tryCompleteElseWatch(delayedFuture, Seq(purgatoryKey))
+    if (!done) {
+      val callbackAction = new BiConsumer[Void, Throwable]() {
+        override def accept(result: Void, exception: Throwable): Unit = delayedFuture.forceComplete()
+      }
+      CompletableFuture.allOf(futures.toArray: _*).whenCompleteAsync(callbackAction, executor)
+    }
+    delayedFuture
+  }
+
+  def shutdown() {
+    executor.shutdownNow()
+    executor.awaitTermination(60, TimeUnit.SECONDS)
+    purgatory.shutdown()
+  }
+}
diff --git a/core/src/main/scala/kafka/server/DelayedOperation.scala b/core/src/main/scala/kafka/server/DelayedOperation.scala
index eb20e6d847da6..51c0e1a891913 100644
--- a/core/src/main/scala/kafka/server/DelayedOperation.scala
+++ b/core/src/main/scala/kafka/server/DelayedOperation.scala
@@ -44,7 +44,8 @@ import scala.collection.mutable.ListBuffer
  * A subclass of DelayedOperation needs to provide an implementation of both onComplete() and tryComplete().
  */
 abstract class DelayedOperation(override val delayMs: Long,
-    lockOpt: Option[Lock] = None) extends TimerTask with Logging {
+                                lockOpt: Option[Lock] = None)
+  extends TimerTask with Logging {
 
   private val completed = new AtomicBoolean(false)
   private val tryCompletePending = new AtomicBoolean(false)
@@ -209,7 +210,7 @@ final class DelayedOperationPurgatory[T <: DelayedOperation](purgatoryName: Stri
   newGauge(
     "NumDelayedOperations",
     new Gauge[Int] {
-      def value: Int = delayed
+      def value: Int = numDelayed
     },
     metricsTags
   )
@@ -288,10 +289,12 @@ final class DelayedOperationPurgatory[T <: DelayedOperation](purgatoryName: Stri
   def checkAndComplete(key: Any): Int = {
     val wl = watcherList(key)
     val watchers = inLock(wl.watchersLock) { wl.watchersByKey.get(key) }
-    if(watchers == null)
+    val numCompleted = if (watchers == null)
       0
     else
       watchers.tryCompleteWatched()
+    debug(s"Request key $key unblocked $numCompleted $purgatoryName operations")
+    numCompleted
   }
 
   /**
@@ -306,7 +309,7 @@ final class DelayedOperationPurgatory[T <: DelayedOperation](purgatoryName: Stri
   /**
    * Return the number of delayed operations in the expiry queue
    */
-  def delayed: Int = timeoutTimer.size
+  def numDelayed: Int = timeoutTimer.size
 
   /**
     * Cancel watching on any delayed operations for the given key. Note the operation will not be completed
@@ -326,7 +329,7 @@ final class DelayedOperationPurgatory[T <: DelayedOperation](purgatoryName: Stri
    * Return the watch list of the given key, note that we need to
    * grab the removeWatchersLock to avoid the operation being added to a removed watcher list
    */
-  private def watchForOperation(key: Any, operation: T) {
+  private def watchForOperation(key: Any, operation: T): Unit = {
     val wl = watcherList(key)
     inLock(wl.watchersLock) {
       val watcher = wl.watchersByKey.getAndMaybePut(key)
@@ -337,7 +340,7 @@ final class DelayedOperationPurgatory[T <: DelayedOperation](purgatoryName: Stri
   /*
    * Remove the key from watcher lists if its list is empty
    */
-  private def removeKeyIfEmpty(key: Any, watchers: Watchers) {
+  private def removeKeyIfEmpty(key: Any, watchers: Watchers): Unit = {
     val wl = watcherList(key)
     inLock(wl.watchersLock) {
       // if the current key is no longer correlated to the watchers to remove, skip
@@ -353,10 +356,12 @@ final class DelayedOperationPurgatory[T <: DelayedOperation](purgatoryName: Stri
   /**
    * Shutdown the expire reaper thread
    */
-  def shutdown() {
+  def shutdown(): Unit = {
     if (reaperEnabled)
       expirationReaper.shutdown()
     timeoutTimer.shutdown()
+    removeMetric("PurgatorySize", metricsTags)
+    removeMetric("NumDelayedOperations", metricsTags)
   }
 
   /**
@@ -371,7 +376,7 @@ final class DelayedOperationPurgatory[T <: DelayedOperation](purgatoryName: Stri
     def isEmpty: Boolean = operations.isEmpty
 
     // add the element to watch
-    def watch(t: T) {
+    def watch(t: T): Unit = {
       operations.add(t)
     }
 
@@ -429,17 +434,17 @@ final class DelayedOperationPurgatory[T <: DelayedOperation](purgatoryName: Stri
     }
   }
 
-  def advanceClock(timeoutMs: Long) {
+  def advanceClock(timeoutMs: Long): Unit = {
     timeoutTimer.advanceClock(timeoutMs)
 
     // Trigger a purge if the number of completed but still being watched operations is larger than
     // the purge threshold. That number is computed by the difference btw the estimated total number of
     // operations and the number of pending delayed operations.
-    if (estimatedTotalOperations.get - delayed > purgeInterval) {
+    if (estimatedTotalOperations.get - numDelayed > purgeInterval) {
       // now set estimatedTotalOperations to delayed (the number of pending operations) since we are going to
       // clean up watchers. Note that, if more operations are completed during the clean up, we may end up with
       // a little overestimated total number of operations.
-      estimatedTotalOperations.getAndSet(delayed)
+      estimatedTotalOperations.getAndSet(numDelayed)
       debug("Begin purging watch lists")
       val purged = watcherLists.foldLeft(0) {
         case (sum, watcherList) => sum + watcherList.allWatchers.map(_.purgeCompleted()).sum
@@ -455,7 +460,7 @@ final class DelayedOperationPurgatory[T <: DelayedOperation](purgatoryName: Stri
     "ExpirationReaper-%d-%s".format(brokerId, purgatoryName),
     false) {
 
-    override def doWork() {
+    override def doWork(): Unit = {
       advanceClock(200L)
     }
   }
diff --git a/core/src/main/scala/kafka/server/DelayedOperationKey.scala b/core/src/main/scala/kafka/server/DelayedOperationKey.scala
index bfa7fc29ea5cf..3be412b04a919 100644
--- a/core/src/main/scala/kafka/server/DelayedOperationKey.scala
+++ b/core/src/main/scala/kafka/server/DelayedOperationKey.scala
@@ -33,11 +33,16 @@ object DelayedOperationKey {
 /* used by delayed-produce and delayed-fetch operations */
 case class TopicPartitionOperationKey(topic: String, partition: Int) extends DelayedOperationKey {
 
-  def this(topicPartition: TopicPartition) = this(topicPartition.topic, topicPartition.partition)
 
   override def keyLabel = "%s-%d".format(topic, partition)
 }
 
+object TopicPartitionOperationKey {
+  def apply(topicPartition: TopicPartition): TopicPartitionOperationKey = {
+    apply(topicPartition.topic, topicPartition.partition)
+  }
+}
+
 /* used by delayed-join-group operations */
 case class MemberKey(groupId: String, consumerId: String) extends DelayedOperationKey {
 
diff --git a/core/src/main/scala/kafka/server/DelayedProduce.scala b/core/src/main/scala/kafka/server/DelayedProduce.scala
index dbecba4e018cf..83e6142008205 100644
--- a/core/src/main/scala/kafka/server/DelayedProduce.scala
+++ b/core/src/main/scala/kafka/server/DelayedProduce.scala
@@ -17,16 +17,14 @@
 
 package kafka.server
 
-
 import java.util.concurrent.TimeUnit
 import java.util.concurrent.locks.Lock
 
 import com.yammer.metrics.core.Meter
 import kafka.metrics.KafkaMetricsGroup
 import kafka.utils.Pool
-
-import org.apache.kafka.common.protocol.Errors
 import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.protocol.Errors
 import org.apache.kafka.common.requests.ProduceResponse.PartitionResponse
 
 import scala.collection._
@@ -87,16 +85,15 @@ class DelayedProduce(delayMs: Long,
       trace(s"Checking produce satisfaction for $topicPartition, current status $status")
       // skip those partitions that have already been satisfied
       if (status.acksPending) {
-        val (hasEnough, error) = replicaManager.getPartition(topicPartition) match {
-          case Some(partition) =>
-            if (partition eq ReplicaManager.OfflinePartition)
-              (false, Errors.KAFKA_STORAGE_ERROR)
-            else
-              partition.checkEnoughReplicasReachOffset(status.requiredOffset)
-          case None =>
+        val (hasEnough, error) = replicaManager.getPartitionOrError(topicPartition, expectLeader = true) match {
+          case Left(err) =>
             // Case A
-            (false, Errors.UNKNOWN_TOPIC_OR_PARTITION)
+            (false, err)
+
+          case Right(partition) =>
+            partition.checkEnoughReplicasReachOffset(status.requiredOffset)
         }
+
         // Case B.1 || B.2
         if (error != Errors.NONE || hasEnough) {
           status.acksPending = false
@@ -112,7 +109,7 @@ class DelayedProduce(delayMs: Long,
       false
   }
 
-  override def onExpiration() {
+  override def onExpiration(): Unit = {
     produceMetadata.produceStatus.foreach { case (topicPartition, status) =>
       if (status.acksPending) {
         debug(s"Expiring produce request for partition $topicPartition with status $status")
@@ -124,8 +121,8 @@ class DelayedProduce(delayMs: Long,
   /**
    * Upon completion, return the current response status along with the error code per partition
    */
-  override def onComplete() {
-    val responseStatus = produceMetadata.produceStatus.mapValues(status => status.responseStatus)
+  override def onComplete(): Unit = {
+    val responseStatus = produceMetadata.produceStatus.map { case (k, status) => k -> status.responseStatus }
     responseCallback(responseStatus)
   }
 }
@@ -141,7 +138,7 @@ object DelayedProduceMetrics extends KafkaMetricsGroup {
              tags = Map("topic" -> key.topic, "partition" -> key.partition.toString))
   private val partitionExpirationMeters = new Pool[TopicPartition, Meter](valueFactory = Some(partitionExpirationMeterFactory))
 
-  def recordExpiration(partition: TopicPartition) {
+  def recordExpiration(partition: TopicPartition): Unit = {
     aggregateExpirationMeter.mark()
     partitionExpirationMeters.getAndMaybePut(partition).mark()
   }
diff --git a/core/src/main/scala/kafka/server/DelegationTokenManager.scala b/core/src/main/scala/kafka/server/DelegationTokenManager.scala
index 90025f429898f..93053e23570cd 100644
--- a/core/src/main/scala/kafka/server/DelegationTokenManager.scala
+++ b/core/src/main/scala/kafka/server/DelegationTokenManager.scala
@@ -198,7 +198,7 @@ class DelegationTokenManager(val config: KafkaConfig,
     }
   }
 
-  private def loadCache() {
+  private def loadCache(): Unit = {
     lock.synchronized {
       val tokens = zkClient.getChildren(DelegationTokensZNode.path)
       info(s"Loading the token cache. Total token count: ${tokens.size}")
@@ -261,7 +261,7 @@ class DelegationTokenManager(val config: KafkaConfig,
   def createToken(owner: KafkaPrincipal,
                   renewers: List[KafkaPrincipal],
                   maxLifeTimeMs: Long,
-                  responseCallback: CreateResponseCallback) {
+                  responseCallback: CreateResponseCallback): Unit = {
 
     if (!config.tokenAuthEnabled) {
       responseCallback(CreateTokenResult(-1, -1, -1, "", Array[Byte](), Errors.DELEGATION_TOKEN_AUTH_DISABLED))
@@ -295,7 +295,7 @@ class DelegationTokenManager(val config: KafkaConfig,
   def renewToken(principal: KafkaPrincipal,
                  hmac: ByteBuffer,
                  renewLifeTimeMs: Long,
-                 renewCallback: RenewResponseCallback) {
+                 renewCallback: RenewResponseCallback): Unit = {
 
     if (!config.tokenAuthEnabled) {
       renewCallback(Errors.DELEGATION_TOKEN_AUTH_DISABLED, -1)
@@ -395,7 +395,7 @@ class DelegationTokenManager(val config: KafkaConfig,
   def expireToken(principal: KafkaPrincipal,
                   hmac: ByteBuffer,
                   expireLifeTimeMs: Long,
-                  expireResponseCallback: ExpireResponseCallback) {
+                  expireResponseCallback: ExpireResponseCallback): Unit = {
 
     if (!config.tokenAuthEnabled) {
       expireResponseCallback(Errors.DELEGATION_TOKEN_AUTH_DISABLED, -1)
@@ -477,7 +477,7 @@ class DelegationTokenManager(val config: KafkaConfig,
   }
 
   object TokenChangedNotificationHandler extends NotificationHandler {
-    override def processNotification(tokenIdBytes: Array[Byte]) {
+    override def processNotification(tokenIdBytes: Array[Byte]): Unit = {
       lock.synchronized {
         val tokenId = new String(tokenIdBytes, StandardCharsets.UTF_8)
         info(s"Processing Token Notification for tokenId: $tokenId")
diff --git a/core/src/main/scala/kafka/server/DynamicBrokerConfig.scala b/core/src/main/scala/kafka/server/DynamicBrokerConfig.scala
index 1c5657263c8d9..563fb87af6538 100755
--- a/core/src/main/scala/kafka/server/DynamicBrokerConfig.scala
+++ b/core/src/main/scala/kafka/server/DynamicBrokerConfig.scala
@@ -23,11 +23,12 @@ import java.util.concurrent.locks.ReentrantReadWriteLock
 
 import kafka.cluster.EndPoint
 import kafka.log.{LogCleaner, LogConfig, LogManager}
+import kafka.network.SocketServer
 import kafka.server.DynamicBrokerConfig._
 import kafka.utils.{CoreUtils, Logging, PasswordEncoder}
 import kafka.zk.{AdminZkClient, KafkaZkClient}
 import org.apache.kafka.common.Reconfigurable
-import org.apache.kafka.common.config.{ConfigDef, ConfigException, SslConfigs}
+import org.apache.kafka.common.config.{ConfigDef, ConfigException, SslConfigs, AbstractConfig}
 import org.apache.kafka.common.metrics.MetricsReporter
 import org.apache.kafka.common.config.types.Password
 import org.apache.kafka.common.network.{ListenerName, ListenerReconfigurable}
@@ -80,11 +81,16 @@ object DynamicBrokerConfig {
     DynamicLogConfig.ReconfigurableConfigs ++
     DynamicThreadPool.ReconfigurableConfigs ++
     Set(KafkaConfig.MetricReporterClassesProp) ++
+    Set(KafkaConfig.AutoCreateTopicsEnableProp) ++
+    Set(KafkaConfig.AllowPreferredControllerFallbackProp) ++
     DynamicListenerConfig.ReconfigurableConfigs ++
-    DynamicConnectionQuota.ReconfigurableConfigs
+    SocketServer.ReconfigurableConfigs
 
+  private val ClusterLevelListenerConfigs = Set(KafkaConfig.MaxConnectionsProp)
+  private val ClusterLevelConfigs = Set(KafkaConfig.AllowPreferredControllerFallbackProp) ++
+    DynamicConfig.Broker.ClusterLevelConfigs
   private val PerBrokerConfigs = DynamicSecurityConfigs  ++
-    DynamicListenerConfig.ReconfigurableConfigs
+    DynamicListenerConfig.ReconfigurableConfigs -- ClusterLevelListenerConfigs
   private val ListenerMechanismConfigs = Set(KafkaConfig.SaslJaasConfigProp)
 
   private val ReloadableFileConfigs = Set(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG, SslConfigs.SSL_TRUSTSTORE_LOCATION_CONFIG)
@@ -118,7 +124,7 @@ object DynamicBrokerConfig {
     }
   }
 
-  def validateConfigs(props: Properties, perBrokerConfig: Boolean): Unit =  {
+  def validateConfigs(props: Properties, perBrokerConfig: Boolean): Unit = {
     def checkInvalidProps(invalidPropNames: Set[String], errorMessage: String): Unit = {
       if (invalidPropNames.nonEmpty)
         throw new ConfigException(s"$errorMessage: $invalidPropNames")
@@ -127,7 +133,10 @@ object DynamicBrokerConfig {
     checkInvalidProps(securityConfigsWithoutListenerPrefix(props),
       "These security configs can be dynamically updated only per-listener using the listener prefix")
     validateConfigTypes(props)
-    if (!perBrokerConfig) {
+    if (perBrokerConfig) {
+      checkInvalidProps(clusterLevelConfigs(props),
+        "Cannot update these configs at per broker level, broker id must not be specified")
+    } else {
       checkInvalidProps(perBrokerConfigs(props),
         "Cannot update these configs at default cluster level, broker id must be specified")
     }
@@ -135,7 +144,18 @@ object DynamicBrokerConfig {
 
   private def perBrokerConfigs(props: Properties): Set[String] = {
     val configNames = props.asScala.keySet
-    configNames.intersect(PerBrokerConfigs) ++ configNames.filter(ListenerConfigRegex.findFirstIn(_).nonEmpty)
+    def perBrokerListenerConfig(name: String): Boolean = {
+      name match {
+        case ListenerConfigRegex(baseName) => !ClusterLevelListenerConfigs.contains(baseName)
+        case _ => false
+      }
+    }
+    configNames.intersect(PerBrokerConfigs) ++ configNames.filter(perBrokerListenerConfig)
+  }
+
+  private def clusterLevelConfigs(props: Properties): Set[String] = {
+    val configNames = props.asScala.keySet
+    configNames.intersect(ClusterLevelConfigs)
   }
 
   private def nonDynamicConfigs(props: Properties): Set[String] = {
@@ -169,6 +189,15 @@ object DynamicBrokerConfig {
       (name -> mode)
     }.toMap.asJava
   }
+
+  private[server] def resolveVariableConfigs(propsOriginal: Properties): Properties = {
+    val props = new Properties
+    val config = new AbstractConfig(new ConfigDef(), propsOriginal, false)
+    config.originals.asScala.filter(!_._1.startsWith(AbstractConfig.CONFIG_PROVIDERS_CONFIG)).foreach {case (key: String, value: Object) => {
+      props.put(key, value)
+    }}
+    props
+  }
 }
 
 class DynamicBrokerConfig(private val kafkaConfig: KafkaConfig) extends Logging {
@@ -204,15 +233,32 @@ class DynamicBrokerConfig(private val kafkaConfig: KafkaConfig) extends Logging
     brokerReconfigurables.clear()
   }
 
+  /**
+   * Add reconfigurables to be notified when a dynamic broker config is updated.
+   *
+   * `Reconfigurable` is the public API used by configurable plugins like metrics reporter
+   * and quota callbacks. These are reconfigured before `KafkaConfig` is updated so that
+   * the update can be aborted if `reconfigure()` fails with an exception.
+   *
+   * `BrokerReconfigurable` is used for internal reconfigurable classes. These are
+   * reconfigured after `KafkaConfig` is updated so that they can access `KafkaConfig`
+   * directly. They are provided both old and new configs.
+   */
   def addReconfigurables(kafkaServer: KafkaServer): Unit = {
+    kafkaServer.authorizer match {
+      case Some(authz: Reconfigurable) => addReconfigurable(authz)
+      case _ =>
+    }
+    addReconfigurable(new DynamicMetricsReporters(kafkaConfig.brokerId, kafkaServer))
+    addReconfigurable(new DynamicClientQuotaCallback(kafkaConfig.brokerId, kafkaServer))
+
     addBrokerReconfigurable(new DynamicThreadPool(kafkaServer))
     if (kafkaServer.logManager.cleaner != null)
       addBrokerReconfigurable(kafkaServer.logManager.cleaner)
-    addReconfigurable(new DynamicLogConfig(kafkaServer.logManager, kafkaServer))
-    addReconfigurable(new DynamicMetricsReporters(kafkaConfig.brokerId, kafkaServer))
-    addReconfigurable(new DynamicClientQuotaCallback(kafkaConfig.brokerId, kafkaServer))
+    addBrokerReconfigurable(new DynamicLogConfig(kafkaServer.logManager, kafkaServer))
     addBrokerReconfigurable(new DynamicListenerConfig(kafkaServer))
-    addBrokerReconfigurable(new DynamicConnectionQuota(kafkaServer))
+    addBrokerReconfigurable(kafkaServer.socketServer)
+    addBrokerReconfigurable(new DynamicAllowPreferredControllerFallback(kafkaServer))
   }
 
   def addReconfigurable(reconfigurable: Reconfigurable): Unit = CoreUtils.inWriteLock(lock) {
@@ -290,6 +336,11 @@ class DynamicBrokerConfig(private val kafkaConfig: KafkaConfig) extends Logging
       }
   }
 
+  private[server] def getMaintenanceBrokerList: Seq[Int] = CoreUtils.inReadLock(lock) {
+    DynamicConfig.Broker.getMaintenanceBrokerListFromString(dynamicDefaultConfigs.getOrElse(DynamicConfig.Broker.MaintenanceBrokerListProp,
+      DynamicConfig.Broker.DefaultMaintenanceBrokerList).toString)
+  }
+
   private def maybeCreatePasswordEncoder(secret: Option[Password]): Option[PasswordEncoder] = {
    secret.map { secret =>
       new PasswordEncoder(secret,
@@ -385,14 +436,15 @@ class DynamicBrokerConfig(private val kafkaConfig: KafkaConfig) extends Logging
    * Note: The caller must acquire the read or write lock before invoking this method.
    */
   private def validatedKafkaProps(propsOverride: Properties, perBrokerConfig: Boolean): Map[String, String] = {
-    validateConfigs(propsOverride, perBrokerConfig)
+    val propsResolved = DynamicBrokerConfig.resolveVariableConfigs(propsOverride)
+    validateConfigs(propsResolved, perBrokerConfig)
     val newProps = mutable.Map[String, String]()
     newProps ++= staticBrokerConfigs
     if (perBrokerConfig) {
       overrideProps(newProps, dynamicDefaultConfigs)
-      overrideProps(newProps, propsOverride.asScala)
+      overrideProps(newProps, propsResolved.asScala)
     } else {
-      overrideProps(newProps, propsOverride.asScala)
+      overrideProps(newProps, propsResolved.asScala)
       overrideProps(newProps, dynamicBrokerConfigs)
     }
     newProps
@@ -565,25 +617,36 @@ object DynamicLogConfig {
   val ReconfigurableConfigs = LogConfig.TopicConfigSynonyms.values.toSet -- ExcludedConfigs
   val KafkaConfigToLogConfigName = LogConfig.TopicConfigSynonyms.map { case (k, v) => (v, k) }
 }
-class DynamicLogConfig(logManager: LogManager, server: KafkaServer) extends Reconfigurable with Logging {
 
-  override def configure(configs: util.Map[String, _]): Unit = {}
+class DynamicLogConfig(logManager: LogManager, server: KafkaServer) extends BrokerReconfigurable with Logging {
 
-  override def reconfigurableConfigs(): util.Set[String] = {
-    DynamicLogConfig.ReconfigurableConfigs.asJava
+  override def reconfigurableConfigs: Set[String] = {
+    DynamicLogConfig.ReconfigurableConfigs
   }
 
-  override def validateReconfiguration(configs: util.Map[String, _]): Unit = {
+  override def validateReconfiguration(newConfig: KafkaConfig): Unit = {
     // For update of topic config overrides, only config names and types are validated
     // Names and types have already been validated. For consistency with topic config
     // validation, no additional validation is performed.
   }
 
-  override def reconfigure(configs: util.Map[String, _]): Unit = {
+  private def updateLogsConfig(newBrokerDefaults: Map[String, Object]): Unit = {
+    logManager.brokerConfigUpdated()
+    logManager.allLogs.foreach { log =>
+      val props = mutable.Map.empty[Any, Any]
+      props ++= newBrokerDefaults
+      props ++= log.config.originals.asScala.filterKeys(log.config.overriddenConfigs.contains)
+
+      val logConfig = LogConfig(props.asJava, log.config.overriddenConfigs)
+      log.updateConfig(logConfig)
+    }
+  }
+
+  override def reconfigure(oldConfig: KafkaConfig, newConfig: KafkaConfig): Unit = {
     val currentLogConfig = logManager.currentDefaultConfig
     val origUncleanLeaderElectionEnable = logManager.currentDefaultConfig.uncleanLeaderElectionEnable
     val newBrokerDefaults = new util.HashMap[String, Object](currentLogConfig.originals)
-    configs.asScala.filterKeys(DynamicLogConfig.ReconfigurableConfigs.contains).foreach { case (k, v) =>
+    newConfig.valuesFromThisConfig.asScala.filterKeys(DynamicLogConfig.ReconfigurableConfigs.contains).foreach { case (k, v) =>
       if (v != null) {
         DynamicLogConfig.KafkaConfigToLogConfigName.get(k).foreach { configName =>
           newBrokerDefaults.put(configName, v.asInstanceOf[AnyRef])
@@ -593,14 +656,8 @@ class DynamicLogConfig(logManager: LogManager, server: KafkaServer) extends Reco
 
     logManager.reconfigureDefaultLogConfig(LogConfig(newBrokerDefaults))
 
-    logManager.allLogs.foreach { log =>
-      val props = mutable.Map.empty[Any, Any]
-      props ++= newBrokerDefaults.asScala
-      props ++= log.config.originals.asScala.filterKeys(log.config.overriddenConfigs.contains)
+    updateLogsConfig(newBrokerDefaults.asScala)
 
-      val logConfig = LogConfig(props.asJava)
-      log.updateConfig(newBrokerDefaults.asScala.keySet, logConfig)
-    }
     if (logManager.currentDefaultConfig.uncleanLeaderElectionEnable && !origUncleanLeaderElectionEnable) {
       server.kafkaController.enableDefaultUncleanLeaderElection()
     }
@@ -616,6 +673,23 @@ object DynamicThreadPool {
     KafkaConfig.BackgroundThreadsProp)
 }
 
+class DynamicAllowPreferredControllerFallback(server: KafkaServer) extends BrokerReconfigurable {
+
+  override def reconfigurableConfigs: Set[String] = {
+    Set(KafkaConfig.AllowPreferredControllerFallbackProp)
+  }
+
+  override def validateReconfiguration(newConfig: KafkaConfig): Unit = {
+    // no additional validation is needed.
+  }
+
+  override def reconfigure(oldConfig: KafkaConfig, newConfig: KafkaConfig): Unit = {
+    if (newConfig.allowPreferredControllerFallback) {
+      server.kafkaController.enablePreferredControllerFallback
+    }
+  }
+}
+
 class DynamicThreadPool(server: KafkaServer) extends BrokerReconfigurable {
 
   override def reconfigurableConfigs: Set[String] = {
@@ -713,7 +787,7 @@ class DynamicMetricsReporters(brokerId: Int, server: KafkaServer) extends Reconf
       case reporter: Reconfigurable => dynamicConfig.maybeReconfigure(reporter, dynamicConfig.currentKafkaConfig, configs)
       case _ =>
     }
-    val added = updatedMetricsReporters -- currentReporters.keySet
+    val added = updatedMetricsReporters.filterNot(currentReporters.keySet)
     createReporters(added.asJava, configs)
   }
 
@@ -778,7 +852,10 @@ object DynamicListenerConfig {
     KafkaConfig.SaslLoginRefreshWindowFactorProp,
     KafkaConfig.SaslLoginRefreshWindowJitterProp,
     KafkaConfig.SaslLoginRefreshMinPeriodSecondsProp,
-    KafkaConfig.SaslLoginRefreshBufferSecondsProp
+    KafkaConfig.SaslLoginRefreshBufferSecondsProp,
+
+    // Connection limit configs
+    KafkaConfig.MaxConnectionsProp
   )
 }
 
@@ -822,9 +899,9 @@ class DynamicListenerConfig(server: KafkaServer) extends BrokerReconfigurable wi
   def validateReconfiguration(newConfig: KafkaConfig): Unit = {
 
     def immutableListenerConfigs(kafkaConfig: KafkaConfig, prefix: String): Map[String, AnyRef] = {
-      newConfig.originals.asScala
-        .filterKeys(_.startsWith(prefix))
-        .filterKeys(k => !DynamicSecurityConfigs.contains(k))
+      newConfig.originals.asScala.filter { case (key, _) =>
+        key.startsWith(prefix) && !DynamicSecurityConfigs.contains(key)
+      }
     }
 
     val oldConfig = server.config
@@ -873,24 +950,5 @@ class DynamicListenerConfig(server: KafkaServer) extends BrokerReconfigurable wi
 
 }
 
-object DynamicConnectionQuota {
-  val ReconfigurableConfigs = Set(KafkaConfig.MaxConnectionsPerIpProp, KafkaConfig.MaxConnectionsPerIpOverridesProp)
-}
-
-class DynamicConnectionQuota(server: KafkaServer) extends BrokerReconfigurable {
 
-  override def reconfigurableConfigs: Set[String] = {
-    DynamicConnectionQuota.ReconfigurableConfigs
-  }
-
-  override def validateReconfiguration(newConfig: KafkaConfig): Unit = {
-  }
-
-  override def reconfigure(oldConfig: KafkaConfig, newConfig: KafkaConfig): Unit = {
-    server.socketServer.updateMaxConnectionsPerIpOverride(newConfig.maxConnectionsPerIpOverrides)
-
-    if (newConfig.maxConnectionsPerIp != oldConfig.maxConnectionsPerIp)
-      server.socketServer.updateMaxConnectionsPerIp(newConfig.maxConnectionsPerIp)
-  }
-}
 
diff --git a/core/src/main/scala/kafka/server/DynamicConfig.scala b/core/src/main/scala/kafka/server/DynamicConfig.scala
index 1a401d2c326d9..695b1569fb4e0 100644
--- a/core/src/main/scala/kafka/server/DynamicConfig.scala
+++ b/core/src/main/scala/kafka/server/DynamicConfig.scala
@@ -21,10 +21,11 @@ import java.util.Properties
 
 import kafka.log.LogConfig
 import kafka.security.CredentialProvider
-import org.apache.kafka.common.config.ConfigDef
+import org.apache.kafka.common.config.{ConfigDef, ConfigException}
 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.common.config.ConfigDef.Validator
 
 import scala.collection.JavaConverters._
 
@@ -39,9 +40,11 @@ object DynamicConfig {
     val LeaderReplicationThrottledRateProp = "leader.replication.throttled.rate"
     val FollowerReplicationThrottledRateProp = "follower.replication.throttled.rate"
     val ReplicaAlterLogDirsIoMaxBytesPerSecondProp = "replica.alter.log.dirs.io.max.bytes.per.second"
+    val MaintenanceBrokerListProp = "maintenance.broker.list"
 
     //Defaults
     val DefaultReplicationThrottledRate = ReplicationQuotaManagerConfig.QuotaBytesPerSecondDefault
+    val DefaultMaintenanceBrokerList: String = ""
 
     //Documentation
     val LeaderReplicationThrottledRateDoc = "A long representing the upper bound (bytes/sec) on replication traffic for leaders enumerated in the " +
@@ -52,6 +55,10 @@ object DynamicConfig {
       s"limit be kept above 1MB/s for accurate behaviour."
     val ReplicaAlterLogDirsIoMaxBytesPerSecondDoc = "A long representing the upper bound (bytes/sec) on disk IO used for moving replica between log directories on the same broker. " +
       s"This property can be only set dynamically. It is suggested that the limit be kept above 1MB/s for accurate behaviour."
+    val MaintenanceBrokerListDoc = "A list containing maintenance broker Ids, separated by comma"
+
+    //cluster level only configs
+    val ClusterLevelConfigs = Set(MaintenanceBrokerListProp)
 
     //Definitions
     private val brokerConfigDef = new ConfigDef()
@@ -59,12 +66,35 @@ object DynamicConfig {
       .define(LeaderReplicationThrottledRateProp, LONG, DefaultReplicationThrottledRate, atLeast(0), MEDIUM, LeaderReplicationThrottledRateDoc)
       .define(FollowerReplicationThrottledRateProp, LONG, DefaultReplicationThrottledRate, atLeast(0), MEDIUM, FollowerReplicationThrottledRateDoc)
       .define(ReplicaAlterLogDirsIoMaxBytesPerSecondProp, LONG, DefaultReplicationThrottledRate, atLeast(0), MEDIUM, ReplicaAlterLogDirsIoMaxBytesPerSecondDoc)
+      .define(MaintenanceBrokerListProp, STRING, DefaultMaintenanceBrokerList, MaintenanceBrokerListValidator, MEDIUM, MaintenanceBrokerListDoc)
     DynamicBrokerConfig.addDynamicConfigs(brokerConfigDef)
     val nonDynamicProps = KafkaConfig.configNames.toSet -- brokerConfigDef.names.asScala
 
+    def typeOf(key: String): ConfigDef.Type = {
+      val configKey: ConfigDef.ConfigKey = brokerConfigDef.configKeys().get(key)
+      if (configKey == null)
+        null
+      else
+        configKey.`type`
+    }
+
     def names = brokerConfigDef.names
 
     def validate(props: Properties) = DynamicConfig.validate(brokerConfigDef, props, customPropsAllowed = true)
+
+    def getMaintenanceBrokerListFromString(brokerListStr: String): Seq[Int] = {
+      brokerListStr.split(",").map(_.trim).filter(_.nonEmpty).map(_.toInt)
+    }
+
+    object MaintenanceBrokerListValidator extends Validator {
+      override def ensureValid(name: String, value: Any): Unit = {
+        try {
+          getMaintenanceBrokerListFromString(value.toString)
+        } catch {
+          case e: NumberFormatException => throw new ConfigException(name, value.toString, e.getMessage)
+        }
+      }
+    }
   }
 
   object Client {
@@ -115,7 +145,8 @@ object DynamicConfig {
       val unknownKeys = propKeys.filter(!names.contains(_))
       require(unknownKeys.isEmpty, s"Unknown Dynamic Configuration: $unknownKeys.")
     }
+    val propResolved = DynamicBrokerConfig.resolveVariableConfigs(props)
     //ValidateValues
-    configDef.parse(props)
+    configDef.parse(propResolved)
   }
 }
diff --git a/core/src/main/scala/kafka/server/DynamicConfigManager.scala b/core/src/main/scala/kafka/server/DynamicConfigManager.scala
index be2edcfb55fba..4a80a6f15ea85 100644
--- a/core/src/main/scala/kafka/server/DynamicConfigManager.scala
+++ b/core/src/main/scala/kafka/server/DynamicConfigManager.scala
@@ -108,7 +108,7 @@ class DynamicConfigManager(private val zkClient: KafkaZkClient,
       }
     }
 
-    private def processEntityConfigChangeVersion1(jsonBytes: Array[Byte], js: JsonObject) {
+    private def processEntityConfigChangeVersion1(jsonBytes: Array[Byte], js: JsonObject): Unit = {
       val validConfigTypes = Set(ConfigType.Topic, ConfigType.Client)
       val entityType = js.get("entity_type").flatMap(_.to[Option[String]]).filter(validConfigTypes).getOrElse {
         throw new IllegalArgumentException("Version 1 config change notification must have 'entity_type' set to " +
@@ -126,7 +126,7 @@ class DynamicConfigManager(private val zkClient: KafkaZkClient,
 
     }
 
-    private def processEntityConfigChangeVersion2(jsonBytes: Array[Byte], js: JsonObject) {
+    private def processEntityConfigChangeVersion2(jsonBytes: Array[Byte], js: JsonObject): Unit = {
 
       val entityPath = js.get("entity_path").flatMap(_.to[Option[String]]).getOrElse {
         throw new IllegalArgumentException(s"Version 2 config change notification must specify 'entity_path'. " +
diff --git a/core/src/main/scala/kafka/server/FetchSession.scala b/core/src/main/scala/kafka/server/FetchSession.scala
index 16ee872824610..aa482eca7b66b 100644
--- a/core/src/main/scala/kafka/server/FetchSession.scala
+++ b/core/src/main/scala/kafka/server/FetchSession.scala
@@ -29,7 +29,7 @@ import org.apache.kafka.common.protocol.Errors
 import org.apache.kafka.common.record.Records
 import org.apache.kafka.common.requests.FetchMetadata.{FINAL_EPOCH, INITIAL_EPOCH, INVALID_SESSION_ID}
 import org.apache.kafka.common.requests.{FetchRequest, FetchResponse, FetchMetadata => JFetchMetadata}
-import org.apache.kafka.common.utils.{ImplicitLinkedHashSet, Time, Utils}
+import org.apache.kafka.common.utils.{ImplicitLinkedHashCollection, Time, Utils}
 
 import scala.math.Ordered.orderingToOrdered
 import scala.collection.{mutable, _}
@@ -38,7 +38,7 @@ import scala.collection.JavaConverters._
 object FetchSession {
   type REQ_MAP = util.Map[TopicPartition, FetchRequest.PartitionData]
   type RESP_MAP = util.LinkedHashMap[TopicPartition, FetchResponse.PartitionData[Records]]
-  type CACHE_MAP = ImplicitLinkedHashSet[CachedPartition]
+  type CACHE_MAP = ImplicitLinkedHashCollection[CachedPartition]
   type RESP_MAP_ITER = util.Iterator[util.Map.Entry[TopicPartition, FetchResponse.PartitionData[Records]]]
 
   val NUM_INCREMENTAL_FETCH_SESSISONS = "NumIncrementalFetchSessions"
@@ -79,10 +79,10 @@ class CachedPartition(val topic: String,
                       var highWatermark: Long,
                       var fetcherLogStartOffset: Long,
                       var localLogStartOffset: Long)
-    extends ImplicitLinkedHashSet.Element {
+    extends ImplicitLinkedHashCollection.Element {
 
-  var cachedNext: Int = ImplicitLinkedHashSet.INVALID_INDEX
-  var cachedPrev: Int = ImplicitLinkedHashSet.INVALID_INDEX
+  var cachedNext: Int = ImplicitLinkedHashCollection.INVALID_INDEX
+  var cachedPrev: Int = ImplicitLinkedHashCollection.INVALID_INDEX
 
   override def next = cachedNext
   override def setNext(next: Int) = this.cachedNext = next
@@ -144,6 +144,10 @@ class CachedPartition(val topic: String,
       if (updateResponseData)
         localLogStartOffset = respData.logStartOffset
     }
+    if (respData.preferredReadReplica.isPresent) {
+      // If the broker computed a preferred read replica, we need to include it in the response
+      mustRespond = true
+    }
     if (respData.error.code != 0) {
       // Partitions with errors are always included in the response.
       // We also set the cached highWatermark to an invalid offset, -1.
diff --git a/core/src/main/scala/kafka/server/KafkaApis.scala b/core/src/main/scala/kafka/server/KafkaApis.scala
index faf338e91fc46..2e7b4d375890d 100644
--- a/core/src/main/scala/kafka/server/KafkaApis.scala
+++ b/core/src/main/scala/kafka/server/KafkaApis.scala
@@ -17,61 +17,74 @@
 
 package kafka.server
 
-import java.lang.{Long => JLong}
 import java.lang.{Byte => JByte}
+import java.lang.{Long => JLong}
 import java.nio.ByteBuffer
 import java.util
-import java.util.concurrent.ConcurrentHashMap
+import java.util.{Collections, Optional}
+import java.util.concurrent.{CompletableFuture, ConcurrentHashMap}
 import java.util.concurrent.atomic.AtomicInteger
-import java.util.{Collections, Optional, Properties}
 
 import kafka.admin.{AdminUtils, RackAwareMode}
-import kafka.api.{ApiVersion, KAFKA_0_11_0_IV0}
+import kafka.api.ElectLeadersRequestOps
+import kafka.api.{ApiVersion, KAFKA_0_11_0_IV0, KAFKA_2_3_IV0}
 import kafka.cluster.Partition
 import kafka.common.OffsetAndMetadata
-import kafka.controller.KafkaController
-import kafka.coordinator.group.{GroupCoordinator, JoinGroupResult}
+import kafka.controller.{KafkaController, ReplicaAssignment}
+import kafka.coordinator.group.{GroupCoordinator, JoinGroupResult, LeaveGroupResult, SyncGroupResult}
 import kafka.coordinator.transaction.{InitProducerIdResult, TransactionCoordinator}
+import kafka.log.AppendOrigin
 import kafka.message.ZStdCompressionCodec
 import kafka.network.RequestChannel
-import kafka.security.SecurityUtils
-import kafka.security.auth.{Resource, _}
+import kafka.security.authorizer.AuthorizerUtils
 import kafka.server.QuotaFactory.{QuotaManagers, UnboundedQuota}
 import kafka.utils.{CoreUtils, Logging}
 import kafka.zk.{AdminZkClient, KafkaZkClient}
-import org.apache.kafka.common.acl.{AccessControlEntry, AclBinding}
+import org.apache.kafka.clients.admin.{AlterConfigOp, ConfigEntry}
+import org.apache.kafka.clients.admin.AlterConfigOp.OpType
+import org.apache.kafka.common.acl.{AclBinding, AclOperation}
+import org.apache.kafka.common.acl.AclOperation._
 import org.apache.kafka.common.config.ConfigResource
 import org.apache.kafka.common.errors._
 import org.apache.kafka.common.internals.FatalExitError
 import org.apache.kafka.common.internals.Topic.{GROUP_METADATA_TOPIC_NAME, TRANSACTION_STATE_TOPIC_NAME, isInternal}
 import org.apache.kafka.common.message.CreateTopicsRequestData.CreatableTopic
-import org.apache.kafka.common.message.{CreateTopicsResponseData, DescribeGroupsResponseData}
-import org.apache.kafka.common.message.CreateTopicsResponseData.{CreatableTopicResult, CreatableTopicResultSet}
-import org.apache.kafka.common.message.ElectPreferredLeadersResponseData
-import org.apache.kafka.common.message.LeaveGroupResponseData
-import org.apache.kafka.common.message.SaslAuthenticateResponseData
-import org.apache.kafka.common.message.SaslHandshakeResponseData
+import org.apache.kafka.common.message.{AlterPartitionReassignmentsResponseData, ApiVersionsResponseData, CreateTopicsResponseData, DeleteGroupsResponseData, DeleteTopicsResponseData, DescribeGroupsResponseData, ExpireDelegationTokenResponseData, FindCoordinatorResponseData, HeartbeatResponseData, InitProducerIdResponseData, JoinGroupResponseData, LeaveGroupResponseData, ListGroupsResponseData, ListPartitionReassignmentsResponseData, OffsetCommitRequestData, OffsetCommitResponseData, OffsetDeleteResponseData, RenewDelegationTokenResponseData, SaslAuthenticateResponseData, SaslHandshakeResponseData, StopReplicaResponseData, SyncGroupResponseData, UpdateMetadataResponseData}
+import org.apache.kafka.common.message.CreateTopicsResponseData.{CreatableTopicResult, CreatableTopicResultCollection}
+import org.apache.kafka.common.message.DeleteGroupsResponseData.{DeletableGroupResult, DeletableGroupResultCollection}
+import org.apache.kafka.common.message.AlterPartitionReassignmentsResponseData.{ReassignablePartitionResponse, ReassignableTopicResponse}
+import org.apache.kafka.common.message.ApiVersionsResponseData.{ApiVersionsResponseKey, ApiVersionsResponseKeyCollection}
+import org.apache.kafka.common.message.DeleteTopicsResponseData.{DeletableTopicResult, DeletableTopicResultCollection}
+import org.apache.kafka.common.message.ElectLeadersResponseData.PartitionResult
+import org.apache.kafka.common.message.ElectLeadersResponseData.ReplicaElectionResult
+import org.apache.kafka.common.message.LeaveGroupResponseData.MemberResponse
 import org.apache.kafka.common.metrics.Metrics
 import org.apache.kafka.common.network.{ListenerName, Send}
 import org.apache.kafka.common.protocol.{ApiKeys, Errors}
 import org.apache.kafka.common.record._
+import org.apache.kafka.common.replica.ClientMetadata
+import org.apache.kafka.common.replica.ClientMetadata.DefaultClientMetadata
 import org.apache.kafka.common.requests.CreateAclsResponse.AclCreationResponse
 import org.apache.kafka.common.requests.DeleteAclsResponse.{AclDeletionResult, AclFilterResponse}
 import org.apache.kafka.common.requests.DescribeLogDirsResponse.LogDirInfo
+import org.apache.kafka.common.requests.FindCoordinatorRequest.CoordinatorType
 import org.apache.kafka.common.requests.ProduceResponse.PartitionResponse
 import org.apache.kafka.common.requests._
-import org.apache.kafka.common.resource.PatternType.LITERAL
-import org.apache.kafka.common.resource.{PatternType, ResourcePattern}
+import org.apache.kafka.common.resource.Resource.CLUSTER_NAME
+import org.apache.kafka.common.resource.ResourceType._
+import org.apache.kafka.common.resource.{PatternType, Resource, ResourcePattern, ResourceType}
 import org.apache.kafka.common.security.auth.{KafkaPrincipal, SecurityProtocol}
 import org.apache.kafka.common.security.token.delegation.{DelegationToken, TokenInformation}
 import org.apache.kafka.common.utils.{Time, Utils}
 import org.apache.kafka.common.{Node, TopicPartition}
+import org.apache.kafka.server.authorizer._
 
+import scala.compat.java8.OptionConverters._
 import scala.collection.JavaConverters._
-import scala.collection._
-import scala.collection.mutable.ArrayBuffer
+import scala.collection.{Map, Seq, Set, immutable, mutable}
 import scala.util.{Failure, Success, Try}
 
+
 /**
  * Logic to handle the various Kafka requests
  */
@@ -87,6 +100,7 @@ class KafkaApis(val requestChannel: RequestChannel,
                 val metadataCache: MetadataCache,
                 val metrics: Metrics,
                 val authorizer: Option[Authorizer],
+                val observer: Observer,
                 val quotas: QuotaManagers,
                 val fetchManager: FetchManager,
                 brokerTopicStats: BrokerTopicStats,
@@ -97,15 +111,17 @@ class KafkaApis(val requestChannel: RequestChannel,
   type FetchResponseStats = Map[TopicPartition, RecordConversionStats]
   this.logIdent = "[KafkaApi-%d] ".format(brokerId)
   val adminZkClient = new AdminZkClient(zkClient)
+  private val alterAclsPurgatory = new DelayedFuturePurgatory(purgatoryName = "AlterAcls", brokerId = config.brokerId)
 
-  def close() {
+  def close(): Unit = {
+    alterAclsPurgatory.shutdown()
     info("Shutdown complete.")
   }
 
   /**
    * Top-level method that handles all requests and multiplexes to the right api
    */
-  def handle(request: RequestChannel.Request) {
+  def handle(request: RequestChannel.Request): Unit = {
     try {
       trace(s"Handling request:${request.requestDesc(true)} from connection ${request.context.connectionId};" +
         s"securityProtocol:${request.context.securityProtocol},principal:${request.context.principal}")
@@ -153,7 +169,11 @@ class KafkaApis(val requestChannel: RequestChannel,
         case ApiKeys.EXPIRE_DELEGATION_TOKEN => handleExpireTokenRequest(request)
         case ApiKeys.DESCRIBE_DELEGATION_TOKEN => handleDescribeTokensRequest(request)
         case ApiKeys.DELETE_GROUPS => handleDeleteGroupsRequest(request)
-        case ApiKeys.ELECT_PREFERRED_LEADERS => handleElectPreferredReplicaLeader(request)
+        case ApiKeys.ELECT_LEADERS => handleElectReplicaLeader(request)
+        case ApiKeys.INCREMENTAL_ALTER_CONFIGS => handleIncrementalAlterConfigsRequest(request)
+        case ApiKeys.ALTER_PARTITION_REASSIGNMENTS => handleAlterPartitionReassignmentsRequest(request)
+        case ApiKeys.LIST_PARTITION_REASSIGNMENTS => handleListPartitionReassignmentsRequest(request)
+        case ApiKeys.OFFSET_DELETE => handleOffsetDeleteRequest(request)
       }
     } catch {
       case e: FatalExitError => throw e
@@ -163,38 +183,39 @@ class KafkaApis(val requestChannel: RequestChannel,
     }
   }
 
-  def handleLeaderAndIsrRequest(request: RequestChannel.Request) {
+  def handleLeaderAndIsrRequest(request: RequestChannel.Request): Unit = {
     // ensureTopicExists is only for client facing requests
     // We can't have the ensureTopicExists check here since the controller sends it as an advisory to all brokers so they
     // stop serving data to clients for the topic being deleted
     val correlationId = request.header.correlationId
     val leaderAndIsrRequest = request.body[LeaderAndIsrRequest]
 
-    def onLeadershipChange(updatedLeaders: Iterable[Partition], updatedFollowers: Iterable[Partition]) {
+    def onLeadershipChange(updatedLeaders: Iterable[Partition], updatedFollowers: Iterable[Partition]): Unit = {
       // for each new leader or follower, call coordinator to handle consumer group migration.
       // this callback is invoked under the replica state change lock to ensure proper order of
       // leadership changes
       updatedLeaders.foreach { partition =>
         if (partition.topic == GROUP_METADATA_TOPIC_NAME)
-          groupCoordinator.handleGroupImmigration(partition.partitionId)
+          groupCoordinator.onElection(partition.partitionId)
         else if (partition.topic == TRANSACTION_STATE_TOPIC_NAME)
-          txnCoordinator.handleTxnImmigration(partition.partitionId, partition.getLeaderEpoch)
+          txnCoordinator.onElection(partition.partitionId, partition.getLeaderEpoch)
       }
 
       updatedFollowers.foreach { partition =>
         if (partition.topic == GROUP_METADATA_TOPIC_NAME)
-          groupCoordinator.handleGroupEmigration(partition.partitionId)
+          groupCoordinator.onResignation(partition.partitionId)
         else if (partition.topic == TRANSACTION_STATE_TOPIC_NAME)
-          txnCoordinator.handleTxnEmigration(partition.partitionId, partition.getLeaderEpoch)
+          txnCoordinator.onResignation(partition.partitionId, Some(partition.getLeaderEpoch))
       }
     }
 
-    authorizeClusterAction(request)
-    if (isBrokerEpochStale(leaderAndIsrRequest.brokerEpoch())) {
+    authorizeClusterOperation(request, CLUSTER_ACTION)
+    if (isBrokerEpochStale(leaderAndIsrRequest.brokerEpoch(), leaderAndIsrRequest.maxBrokerEpoch())) {
       // When the broker restarts very quickly, it is possible for this broker to receive request intended
       // for its previous generation so the broker should skip the stale request.
-      info("Received LeaderAndIsr request with broker epoch " +
-        s"${leaderAndIsrRequest.brokerEpoch()} smaller than the current broker epoch ${controller.brokerEpoch}")
+      info("Received LeaderAndIsr request with stale broker epoch info " +
+        "(broker epoch:" + leaderAndIsrRequest.brokerEpoch() +"},max broker epoch:" + leaderAndIsrRequest.maxBrokerEpoch() + ") " +
+        "when the current broker epoch is " + controller.brokerEpoch)
       sendResponseExemptThrottle(request, leaderAndIsrRequest.getErrorResponse(0, Errors.STALE_BROKER_EPOCH.exception))
     } else {
       val response = replicaManager.becomeLeaderOrFollower(correlationId, leaderAndIsrRequest, onLeadershipChange)
@@ -202,56 +223,69 @@ class KafkaApis(val requestChannel: RequestChannel,
     }
   }
 
-  def handleStopReplicaRequest(request: RequestChannel.Request) {
+  def handleStopReplicaRequest(request: RequestChannel.Request): Unit = {
     // ensureTopicExists is only for client facing requests
     // We can't have the ensureTopicExists check here since the controller sends it as an advisory to all brokers so they
     // stop serving data to clients for the topic being deleted
     val stopReplicaRequest = request.body[StopReplicaRequest]
-    authorizeClusterAction(request)
-    if (isBrokerEpochStale(stopReplicaRequest.brokerEpoch())) {
+    authorizeClusterOperation(request, CLUSTER_ACTION)
+    if (isBrokerEpochStale(stopReplicaRequest.brokerEpoch(), stopReplicaRequest.maxBrokerEpoch())) {
       // When the broker restarts very quickly, it is possible for this broker to receive request intended
       // for its previous generation so the broker should skip the stale request.
-      info("Received stop replica request with broker epoch " +
-        s"${stopReplicaRequest.brokerEpoch()} smaller than the current broker epoch ${controller.brokerEpoch}")
-      sendResponseExemptThrottle(request, new StopReplicaResponse(Errors.STALE_BROKER_EPOCH, Map.empty[TopicPartition, Errors].asJava))
+      info("Received stop replica request with stale broker epoch info " +
+        s"(broker epoch:${stopReplicaRequest.brokerEpoch()},max broker epoch:${stopReplicaRequest.maxBrokerEpoch()})" +
+        s"when the current broker epoch is ${controller.brokerEpoch}")
+      sendResponseExemptThrottle(request, new StopReplicaResponse(new StopReplicaResponseData().setErrorCode(Errors.STALE_BROKER_EPOCH.code)))
     } else {
       val (result, error) = replicaManager.stopReplicas(stopReplicaRequest)
-      // Clearing out the cache for groups that belong to an offsets topic partition for which this broker was the leader,
-      // since this broker is no longer a replica for that offsets topic partition.
-      // This is required to handle the following scenario :
-      // Consider old replicas : {[1,2,3], Leader = 1} is reassigned to new replicas : {[2,3,4], Leader = 2}, broker 1 does not receive a LeaderAndIsr
-      // request to become a follower due to which cache for groups that belong to an offsets topic partition for which broker 1 was the leader,
-      // is not cleared.
+      // Clear the coordinator caches in case we were the leader. In the case of a reassignment, we
+      // cannot rely on the LeaderAndIsr API for this since it is only sent to active replicas.
       result.foreach { case (topicPartition, error) =>
-        if (error == Errors.NONE && stopReplicaRequest.deletePartitions && topicPartition.topic == GROUP_METADATA_TOPIC_NAME) {
-          groupCoordinator.handleGroupEmigration(topicPartition.partition)
+        if (error == Errors.NONE && stopReplicaRequest.deletePartitions) {
+          if (topicPartition.topic == GROUP_METADATA_TOPIC_NAME) {
+            groupCoordinator.onResignation(topicPartition.partition)
+          } else if (topicPartition.topic == TRANSACTION_STATE_TOPIC_NAME) {
+            // The StopReplica API does not pass through the leader epoch
+            txnCoordinator.onResignation(topicPartition.partition, coordinatorEpoch = None)
+          }
         }
       }
-      sendResponseExemptThrottle(request, new StopReplicaResponse(error, result.asJava))
+
+      def toStopReplicaPartition(tp: TopicPartition, error: Errors) =
+        new StopReplicaResponseData.StopReplicaPartitionError()
+          .setTopicName(tp.topic)
+          .setPartitionIndex(tp.partition)
+          .setErrorCode(error.code)
+
+      sendResponseExemptThrottle(request, new StopReplicaResponse(new StopReplicaResponseData()
+        .setErrorCode(error.code)
+        .setPartitionErrors(result.map { case (tp, error) => toStopReplicaPartition(tp, error) }.toBuffer.asJava)))
     }
 
     CoreUtils.swallow(replicaManager.replicaFetcherManager.shutdownIdleFetcherThreads(), this)
   }
 
-  def handleUpdateMetadataRequest(request: RequestChannel.Request) {
+  def handleUpdateMetadataRequest(request: RequestChannel.Request): Unit = {
     val correlationId = request.header.correlationId
     val updateMetadataRequest = request.body[UpdateMetadataRequest]
 
-    authorizeClusterAction(request)
-    if (isBrokerEpochStale(updateMetadataRequest.brokerEpoch())) {
+    authorizeClusterOperation(request, CLUSTER_ACTION)
+    if (isBrokerEpochStale(updateMetadataRequest.brokerEpoch, updateMetadataRequest.maxBrokerEpoch)) {
       // When the broker restarts very quickly, it is possible for this broker to receive request intended
       // for its previous generation so the broker should skip the stale request.
-      info("Received update metadata request with broker epoch " +
-        s"${updateMetadataRequest.brokerEpoch()} smaller than the current broker epoch ${controller.brokerEpoch}")
-      sendResponseExemptThrottle(request, new UpdateMetadataResponse(Errors.STALE_BROKER_EPOCH))
+      info("Received update metadata request with stale broker epoch info " +
+        s"(broker epoch:${updateMetadataRequest.brokerEpoch},max broker epoch: ${updateMetadataRequest.maxBrokerEpoch}) " +
+        s"when the current broker epoch ${controller.brokerEpoch}")
+      sendResponseExemptThrottle(request,
+        new UpdateMetadataResponse(new UpdateMetadataResponseData().setErrorCode(Errors.STALE_BROKER_EPOCH.code)))
     } else {
       val deletedPartitions = replicaManager.maybeUpdateMetadataCache(correlationId, updateMetadataRequest)
       if (deletedPartitions.nonEmpty)
         groupCoordinator.handleDeletedPartitions(deletedPartitions)
 
       if (adminManager.hasDelayedTopicOperations) {
-        updateMetadataRequest.partitionStates.keySet.asScala.map(_.topic).foreach { topic =>
-          adminManager.tryCompleteDelayedTopicOperations(topic)
+        updateMetadataRequest.partitionStates.asScala.foreach { partitionState =>
+          adminManager.tryCompleteDelayedTopicOperations(partitionState.topicName)
         }
       }
       quotas.clientQuotaCallback.foreach { callback =>
@@ -262,79 +296,101 @@ class KafkaApis(val requestChannel: RequestChannel,
         }
       }
       if (replicaManager.hasDelayedElectionOperations) {
-        updateMetadataRequest.partitionStates.asScala.foreach { case (tp, ps) =>
-          replicaManager.tryCompleteElection(new TopicPartitionOperationKey(tp.topic(), tp.partition()))
+        updateMetadataRequest.partitionStates.asScala.foreach { partitionState =>
+          val tp = new TopicPartition(partitionState.topicName, partitionState.partitionIndex)
+          replicaManager.tryCompleteElection(TopicPartitionOperationKey(tp))
         }
       }
-      sendResponseExemptThrottle(request, new UpdateMetadataResponse(Errors.NONE))
+      sendResponseExemptThrottle(request, new UpdateMetadataResponse(
+        new UpdateMetadataResponseData().setErrorCode(Errors.NONE.code)))
     }
   }
 
-  def handleControlledShutdownRequest(request: RequestChannel.Request) {
+  def handleControlledShutdownRequest(request: RequestChannel.Request): Unit = {
     // ensureTopicExists is only for client facing requests
     // We can't have the ensureTopicExists check here since the controller sends it as an advisory to all brokers so they
     // stop serving data to clients for the topic being deleted
     val controlledShutdownRequest = request.body[ControlledShutdownRequest]
-    authorizeClusterAction(request)
+    authorizeClusterOperation(request, CLUSTER_ACTION)
 
     def controlledShutdownCallback(controlledShutdownResult: Try[Set[TopicPartition]]): Unit = {
       val response = controlledShutdownResult match {
         case Success(partitionsRemaining) =>
-          new ControlledShutdownResponse(Errors.NONE, partitionsRemaining.asJava)
+         ControlledShutdownResponse.prepareResponse(Errors.NONE, partitionsRemaining.asJava)
 
         case Failure(throwable) =>
           controlledShutdownRequest.getErrorResponse(throwable)
       }
       sendResponseExemptThrottle(request, response)
     }
-    controller.controlledShutdown(controlledShutdownRequest.brokerId, controlledShutdownRequest.brokerEpoch, controlledShutdownCallback)
+    controller.controlledShutdown(controlledShutdownRequest.data.brokerId, controlledShutdownRequest.data.brokerEpoch, controlledShutdownCallback)
   }
 
   /**
    * Handle an offset commit request
    */
-  def handleOffsetCommitRequest(request: RequestChannel.Request) {
+  def handleOffsetCommitRequest(request: RequestChannel.Request): Unit = {
     val header = request.header
     val offsetCommitRequest = request.body[OffsetCommitRequest]
 
+    val unauthorizedTopicErrors = mutable.Map[TopicPartition, Errors]()
+    val nonExistingTopicErrors = mutable.Map[TopicPartition, Errors]()
+    // the callback for sending an offset commit response
+    def sendResponseCallback(commitStatus: Map[TopicPartition, Errors]): Unit = {
+      val combinedCommitStatus = commitStatus ++ unauthorizedTopicErrors ++ nonExistingTopicErrors
+      if (isDebugEnabled)
+        combinedCommitStatus.foreach { case (topicPartition, error) =>
+          if (error != Errors.NONE) {
+            debug(s"Offset commit request with correlation id ${header.correlationId} from client ${header.clientId} " +
+              s"on partition $topicPartition failed due to ${error.exceptionName}")
+          }
+        }
+      sendResponseMaybeThrottle(request, requestThrottleMs =>
+        new OffsetCommitResponse(requestThrottleMs, combinedCommitStatus.asJava))
+    }
+
     // reject the request if not authorized to the group
-    if (!authorize(request.session, Read, Resource(Group, offsetCommitRequest.groupId, LITERAL))) {
+    if (!authorize(request, READ, GROUP, offsetCommitRequest.data().groupId)) {
       val error = Errors.GROUP_AUTHORIZATION_FAILED
-      val results = offsetCommitRequest.offsetData.keySet.asScala.map { topicPartition =>
-        (topicPartition, error)
-      }.toMap
-      sendResponseMaybeThrottle(request, requestThrottleMs => new OffsetCommitResponse(requestThrottleMs, results.asJava))
+      val responseTopicList = OffsetCommitRequest.getErrorResponseTopics(
+        offsetCommitRequest.data().topics(),
+        error)
+
+      sendResponseMaybeThrottle(request, requestThrottleMs => new OffsetCommitResponse(
+        new OffsetCommitResponseData()
+            .setTopics(responseTopicList)
+            .setThrottleTimeMs(requestThrottleMs)
+      ))
+    } else if (offsetCommitRequest.data.groupInstanceId != null && config.interBrokerProtocolVersion < KAFKA_2_3_IV0) {
+      // Only enable static membership when IBP >= 2.3, because it is not safe for the broker to use the static member logic
+      // until we are sure that all brokers support it. If static group being loaded by an older coordinator, it will discard
+      // the group.instance.id field, so static members could accidentally become "dynamic", which leads to wrong states.
+      val errorMap = new mutable.HashMap[TopicPartition, Errors]
+      for (topicData <- offsetCommitRequest.data().topics().asScala) {
+        for (partitionData <- topicData.partitions().asScala) {
+          val topicPartition = new TopicPartition(topicData.name(), partitionData.partitionIndex())
+          errorMap += topicPartition -> Errors.UNSUPPORTED_VERSION
+        }
+      }
+      sendResponseCallback(errorMap.toMap)
     } else {
-
-      val unauthorizedTopicErrors = mutable.Map[TopicPartition, Errors]()
-      val nonExistingTopicErrors = mutable.Map[TopicPartition, Errors]()
-      val authorizedTopicRequestInfoBldr = immutable.Map.newBuilder[TopicPartition, OffsetCommitRequest.PartitionData]
-
-      for ((topicPartition, partitionData) <- offsetCommitRequest.offsetData.asScala) {
-        if (!authorize(request.session, Read, Resource(Topic, topicPartition.topic, LITERAL)))
-          unauthorizedTopicErrors += (topicPartition -> Errors.TOPIC_AUTHORIZATION_FAILED)
-        else if (!metadataCache.contains(topicPartition))
-          nonExistingTopicErrors += (topicPartition -> Errors.UNKNOWN_TOPIC_OR_PARTITION)
-        else
-          authorizedTopicRequestInfoBldr += (topicPartition -> partitionData)
+      val authorizedTopicRequestInfoBldr = immutable.Map.newBuilder[TopicPartition, OffsetCommitRequestData.OffsetCommitRequestPartition]
+
+      val authorizedTopics = filterAuthorized(request, READ, TOPIC, offsetCommitRequest.data.topics.asScala.map(_.name))
+      for (topicData <- offsetCommitRequest.data().topics().asScala) {
+        for (partitionData <- topicData.partitions().asScala) {
+          val topicPartition = new TopicPartition(topicData.name(), partitionData.partitionIndex())
+          if (!authorizedTopics.contains(topicData.name()))
+            unauthorizedTopicErrors += (topicPartition -> Errors.TOPIC_AUTHORIZATION_FAILED)
+          else if (!metadataCache.contains(topicPartition))
+            nonExistingTopicErrors += (topicPartition -> Errors.UNKNOWN_TOPIC_OR_PARTITION)
+          else
+            authorizedTopicRequestInfoBldr += (topicPartition -> partitionData)
+        }
       }
 
       val authorizedTopicRequestInfo = authorizedTopicRequestInfoBldr.result()
 
-      // the callback for sending an offset commit response
-      def sendResponseCallback(commitStatus: immutable.Map[TopicPartition, Errors]) {
-        val combinedCommitStatus = commitStatus ++ unauthorizedTopicErrors ++ nonExistingTopicErrors
-        if (isDebugEnabled)
-          combinedCommitStatus.foreach { case (topicPartition, error) =>
-            if (error != Errors.NONE) {
-              debug(s"Offset commit request with correlation id ${header.correlationId} from client ${header.clientId} " +
-                s"on partition $topicPartition failed due to ${error.exceptionName}")
-            }
-          }
-        sendResponseMaybeThrottle(request, requestThrottleMs =>
-          new OffsetCommitResponse(requestThrottleMs, combinedCommitStatus.asJava))
-      }
-
       if (authorizedTopicRequestInfo.isEmpty)
         sendResponseCallback(Map.empty)
       else if (header.apiVersion == 0) {
@@ -342,10 +398,14 @@ class KafkaApis(val requestChannel: RequestChannel,
         val responseInfo = authorizedTopicRequestInfo.map {
           case (topicPartition, partitionData) =>
             try {
-              if (partitionData.metadata != null && partitionData.metadata.length > config.offsetMetadataMaxSize)
+              if (partitionData.committedMetadata() != null
+                && partitionData.committedMetadata().length > config.offsetMetadataMaxSize)
                 (topicPartition, Errors.OFFSET_METADATA_TOO_LARGE)
               else {
-                zkClient.setOrCreateConsumerOffset(offsetCommitRequest.groupId, topicPartition, partitionData.offset)
+                zkClient.setOrCreateConsumerOffset(
+                  offsetCommitRequest.data().groupId(),
+                  topicPartition,
+                  partitionData.committedOffset())
                 (topicPartition, Errors.NONE)
               }
             } catch {
@@ -363,17 +423,26 @@ class KafkaApis(val requestChannel: RequestChannel,
         //   - If v2/v3/v4 (no explicit commit timestamp) we treat it the same as v5.
         //   - For v5 and beyond there is no per partition expiration timestamp, so this field is no longer in effect
         val currentTimestamp = time.milliseconds
-        val partitionData = authorizedTopicRequestInfo.mapValues { partitionData =>
-          val metadata = if (partitionData.metadata == null) OffsetAndMetadata.NoMetadata else partitionData.metadata
-          new OffsetAndMetadata(
-            offset = partitionData.offset,
-            leaderEpoch = partitionData.leaderEpoch,
+        val partitionData = authorizedTopicRequestInfo.map { case (k, partitionData) =>
+          val metadata = if (partitionData.committedMetadata() == null)
+            OffsetAndMetadata.NoMetadata
+          else
+            partitionData.committedMetadata
+
+          val leaderEpochOpt = if (partitionData.committedLeaderEpoch == RecordBatch.NO_PARTITION_LEADER_EPOCH)
+            Optional.empty[Integer]
+          else
+            Optional.of[Integer](partitionData.committedLeaderEpoch)
+
+          k -> new OffsetAndMetadata(
+            offset = partitionData.committedOffset(),
+            leaderEpoch = leaderEpochOpt,
             metadata = metadata,
-            commitTimestamp = partitionData.timestamp match {
+            commitTimestamp = partitionData.commitTimestamp() match {
               case OffsetCommitRequest.DEFAULT_TIMESTAMP => currentTimestamp
               case customTimestamp => customTimestamp
             },
-            expireTimestamp = offsetCommitRequest.retentionTime match {
+            expireTimestamp = offsetCommitRequest.data().retentionTimeMs() match {
               case OffsetCommitRequest.DEFAULT_RETENTION_TIME => None
               case retentionTime => Some(currentTimestamp + retentionTime)
             }
@@ -382,35 +451,33 @@ class KafkaApis(val requestChannel: RequestChannel,
 
         // call coordinator to handle commit offset
         groupCoordinator.handleCommitOffsets(
-          offsetCommitRequest.groupId,
-          offsetCommitRequest.memberId,
-          offsetCommitRequest.generationId,
+          offsetCommitRequest.data.groupId,
+          offsetCommitRequest.data.memberId,
+          Option(offsetCommitRequest.data.groupInstanceId),
+          offsetCommitRequest.data.generationId,
           partitionData,
           sendResponseCallback)
       }
     }
   }
 
-  private def authorize(session: RequestChannel.Session, operation: Operation, resource: Resource): Boolean =
-    authorizer.forall(_.authorize(session, operation, resource))
-
   /**
    * Handle a produce request
    */
-  def handleProduceRequest(request: RequestChannel.Request) {
+  def handleProduceRequest(request: RequestChannel.Request): Unit = {
     val produceRequest = request.body[ProduceRequest]
     val numBytesAppended = request.header.toStruct.sizeOf + request.sizeOfBodyInBytes
 
     if (produceRequest.hasTransactionalRecords) {
       val isAuthorizedTransactional = produceRequest.transactionalId != null &&
-        authorize(request.session, Write, Resource(TransactionalId, produceRequest.transactionalId, LITERAL))
+        authorize(request, WRITE, TRANSACTIONAL_ID, produceRequest.transactionalId)
       if (!isAuthorizedTransactional) {
         sendErrorResponseMaybeThrottle(request, Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED.exception)
         return
       }
       // Note that authorization to a transactionalId implies ProducerId authorization
 
-    } else if (produceRequest.hasIdempotentRecords && !authorize(request.session, IdempotentWrite, Resource.ClusterResource)) {
+    } else if (produceRequest.hasIdempotentRecords && !authorize(request, IDEMPOTENT_WRITE, CLUSTER, CLUSTER_NAME)) {
       sendErrorResponseMaybeThrottle(request, Errors.CLUSTER_AUTHORIZATION_FAILED.exception)
       return
     }
@@ -419,9 +486,11 @@ class KafkaApis(val requestChannel: RequestChannel,
     val nonExistingTopicResponses = mutable.Map[TopicPartition, PartitionResponse]()
     val invalidRequestResponses = mutable.Map[TopicPartition, PartitionResponse]()
     val authorizedRequestInfo = mutable.Map[TopicPartition, MemoryRecords]()
+    val authorizedTopics = filterAuthorized(request, WRITE, TOPIC,
+      produceRequest.partitionRecordsOrFail.asScala.toSeq.map(_._1.topic))
 
     for ((topicPartition, memoryRecords) <- produceRequest.partitionRecordsOrFail.asScala) {
-      if (!authorize(request.session, Write, Resource(Topic, topicPartition.topic, LITERAL)))
+      if (!authorizedTopics.contains(topicPartition.topic))
         unauthorizedTopicResponses += topicPartition -> new PartitionResponse(Errors.TOPIC_AUTHORIZATION_FAILED)
       else if (!metadataCache.contains(topicPartition))
         nonExistingTopicResponses += topicPartition -> new PartitionResponse(Errors.UNKNOWN_TOPIC_OR_PARTITION)
@@ -436,7 +505,7 @@ class KafkaApis(val requestChannel: RequestChannel,
     }
 
     // the callback for sending a produce response
-    def sendResponseCallback(responseStatus: Map[TopicPartition, PartitionResponse]) {
+    def sendResponseCallback(responseStatus: Map[TopicPartition, PartitionResponse]): Unit = {
       val mergedResponseStatus = responseStatus ++ unauthorizedTopicResponses ++ nonExistingTopicResponses ++ invalidRequestResponses
       var errorInResponse = false
 
@@ -502,6 +571,14 @@ class KafkaApis(val requestChannel: RequestChannel,
     if (authorizedRequestInfo.isEmpty)
       sendResponseCallback(Map.empty)
     else {
+
+      try
+        observer.observeProduceRequest(request.context, request.body[ProduceRequest])
+      catch {
+        case e: Exception => error(s"Observer failed to observe the produce request " +
+          s"${Observer.describeRequestAndResponse(request, null)}", e)
+      }
+
       val internalTopicsAllowed = request.header.clientId == AdminUtils.AdminClientId
 
       // call the replica manager to append messages to the replicas
@@ -509,7 +586,7 @@ class KafkaApis(val requestChannel: RequestChannel,
         timeout = produceRequest.timeout.toLong,
         requiredAcks = produceRequest.acks,
         internalTopicsAllowed = internalTopicsAllowed,
-        isFromClient = true,
+        origin = AppendOrigin.Client,
         entriesPerPartition = authorizedRequestInfo,
         responseCallback = sendResponseCallback,
         recordConversionStatsCallback = processingStatsCallback)
@@ -523,7 +600,7 @@ class KafkaApis(val requestChannel: RequestChannel,
   /**
    * Handle a fetch request
    */
-  def handleFetchRequest(request: RequestChannel.Request) {
+  def handleFetchRequest(request: RequestChannel.Request): Unit = {
     val versionId = request.header.apiVersion
     val clientId = request.header.clientId
     val fetchRequest = request.body[FetchRequest]
@@ -533,6 +610,18 @@ class KafkaApis(val requestChannel: RequestChannel,
       fetchRequest.toForget,
       fetchRequest.isFromFollower)
 
+    val clientMetadata: Option[ClientMetadata] = if (versionId >= 11) {
+      // Fetch API version 11 added preferred replica logic
+      Some(new DefaultClientMetadata(
+        fetchRequest.rackId,
+        clientId,
+        request.context.clientAddress,
+        request.context.principal,
+        request.context.listenerName.value))
+    } else {
+      None
+    }
+
     def errorResponse[T >: MemoryRecords <: BaseRecords](error: Errors): FetchResponse.PartitionData[T] = {
       new FetchResponse.PartitionData[T](error, FetchResponse.INVALID_HIGHWATERMARK, FetchResponse.INVALID_LAST_STABLE_OFFSET,
         FetchResponse.INVALID_LOG_START_OFFSET, null, MemoryRecords.EMPTY)
@@ -542,7 +631,7 @@ class KafkaApis(val requestChannel: RequestChannel,
     val interesting = mutable.ArrayBuffer[(TopicPartition, FetchRequest.PartitionData)]()
     if (fetchRequest.isFromFollower) {
       // The follower must have ClusterAction on ClusterResource in order to fetch partition data.
-      if (authorize(request.session, ClusterAction, Resource.ClusterResource)) {
+      if (authorize(request, CLUSTER_ACTION, CLUSTER, CLUSTER_NAME)) {
         fetchContext.foreachPartition { (topicPartition, data) =>
           if (!metadataCache.contains(topicPartition))
             erroneous += topicPartition -> errorResponse(Errors.UNKNOWN_TOPIC_OR_PARTITION)
@@ -556,8 +645,11 @@ class KafkaApis(val requestChannel: RequestChannel,
       }
     } else {
       // Regular Kafka consumers need READ permission on each partition they are fetching.
+      val fetchTopics = new mutable.ArrayBuffer[String]
+      fetchContext.foreachPartition { (topicPartition, _) => fetchTopics += topicPartition.topic }
+      val authorizedTopics = filterAuthorized(request, READ, TOPIC, fetchTopics)
       fetchContext.foreachPartition { (topicPartition, data) =>
-        if (!authorize(request.session, Read, Resource(Topic, topicPartition.topic, LITERAL)))
+        if (!authorizedTopics.contains(topicPartition.topic))
           erroneous += topicPartition -> errorResponse(Errors.TOPIC_AUTHORIZATION_FAILED)
         else if (!metadataCache.contains(topicPartition))
           erroneous += topicPartition -> errorResponse(Errors.UNKNOWN_TOPIC_OR_PARTITION)
@@ -605,8 +697,10 @@ class KafkaApis(val requestChannel: RequestChannel,
                 // as possible. With KIP-283, we have the ability to lazily down-convert in a chunked manner. The lazy, chunked
                 // down-conversion always guarantees that at least one batch of messages is down-converted and sent out to the
                 // client.
+                replicaManager.incrementRecompressionCount()
                 new FetchResponse.PartitionData[BaseRecords](partitionData.error, partitionData.highWatermark,
-                  partitionData.lastStableOffset, partitionData.logStartOffset, partitionData.abortedTransactions,
+                  partitionData.lastStableOffset, partitionData.logStartOffset,
+                  partitionData.preferredReadReplica, partitionData.abortedTransactions,
                   new LazyDownConversionRecords(tp, unconvertedRecords, magic, fetchContext.getFetchOffset(tp).get, time))
               } catch {
                 case e: UnsupportedCompressionTypeException =>
@@ -615,7 +709,8 @@ class KafkaApis(val requestChannel: RequestChannel,
               }
             }
           case None => new FetchResponse.PartitionData[BaseRecords](partitionData.error, partitionData.highWatermark,
-            partitionData.lastStableOffset, partitionData.logStartOffset, partitionData.abortedTransactions,
+            partitionData.lastStableOffset, partitionData.logStartOffset,
+            partitionData.preferredReadReplica, partitionData.abortedTransactions,
             unconvertedRecords)
         }
       }
@@ -628,7 +723,8 @@ class KafkaApis(val requestChannel: RequestChannel,
         val abortedTransactions = data.abortedTransactions.map(_.asJava).orNull
         val lastStableOffset = data.lastStableOffset.getOrElse(FetchResponse.INVALID_LAST_STABLE_OFFSET)
         partitions.put(tp, new FetchResponse.PartitionData(data.error, data.highWatermark, lastStableOffset,
-          data.logStartOffset, abortedTransactions, data.records))
+          data.logStartOffset, data.preferredReadReplica.map(int2Integer).asJava,
+          abortedTransactions, data.records))
       }
       erroneous.foreach { case (tp, data) => partitions.put(tp, data) }
 
@@ -712,6 +808,16 @@ class KafkaApis(val requestChannel: RequestChannel,
       }
     }
 
+    // for fetch from consumer, cap fetchMaxBytes to the maximum bytes that could be fetched without being throttled given
+    // no bytes were recorded in the recent quota window
+    // trying to fetch more bytes would result in a guaranteed throttling potentially blocking consumer progress
+    val maxQuotaWindowBytes = if (fetchRequest.isFromFollower)
+      Int.MaxValue
+    else
+      quotas.fetch.getMaxValueInQuotaWindow(request.session, clientId).toInt
+
+    val fetchMaxBytes = Math.min(fetchRequest.maxBytes, maxQuotaWindowBytes)
+    val fetchMinBytes = Math.min(fetchRequest.minBytes, fetchMaxBytes)
     if (interesting.isEmpty)
       processResponseCallback(Seq.empty)
     else {
@@ -719,13 +825,14 @@ class KafkaApis(val requestChannel: RequestChannel,
       replicaManager.fetchMessages(
         fetchRequest.maxWait.toLong,
         fetchRequest.replicaId,
-        fetchRequest.minBytes,
-        fetchRequest.maxBytes,
+        fetchMinBytes,
+        fetchMaxBytes,
         versionId <= 2,
         interesting,
         replicationQuota(fetchRequest),
         processResponseCallback,
-        fetchRequest.isolationLevel)
+        fetchRequest.isolationLevel,
+        clientMetadata)
     }
   }
 
@@ -768,7 +875,7 @@ class KafkaApis(val requestChannel: RequestChannel,
   def replicationQuota(fetchRequest: FetchRequest): ReplicaQuota =
     if (fetchRequest.isFromFollower) quotas.leader else UnboundedQuota
 
-  def handleListOffsetRequest(request: RequestChannel.Request) {
+  def handleListOffsetRequest(request: RequestChannel.Request): Unit = {
     val version = request.header.apiVersion()
 
     val mergedResponseMap = if (version == 0)
@@ -784,8 +891,9 @@ class KafkaApis(val requestChannel: RequestChannel,
     val clientId = request.header.clientId
     val offsetRequest = request.body[ListOffsetRequest]
 
+    val authorizedTopics = filterAuthorized(request, DESCRIBE, TOPIC, offsetRequest.partitionTimestamps.asScala.toSeq.map(_._1.topic))
     val (authorizedRequestInfo, unauthorizedRequestInfo) = offsetRequest.partitionTimestamps.asScala.partition {
-      case (topicPartition, _) => authorize(request.session, Describe, Resource(Topic, topicPartition.topic, LITERAL))
+      case (topicPartition, _) => authorizedTopics.contains(topicPartition.topic)
     }
 
     val unauthorizedResponseStatus = unauthorizedRequestInfo.mapValues(_ =>
@@ -823,8 +931,9 @@ class KafkaApis(val requestChannel: RequestChannel,
     val clientId = request.header.clientId
     val offsetRequest = request.body[ListOffsetRequest]
 
+    val authorizedTopics = filterAuthorized(request, DESCRIBE, TOPIC, offsetRequest.partitionTimestamps.asScala.toSeq.map(_._1.topic))
     val (authorizedRequestInfo, unauthorizedRequestInfo) = offsetRequest.partitionTimestamps.asScala.partition {
-      case (topicPartition, _) => authorize(request.session, Describe, Resource(Topic, topicPartition.topic, LITERAL))
+      case (topicPartition, _) => authorizedTopics.contains(topicPartition.topic)
     }
 
     val unauthorizedResponseStatus = unauthorizedRequestInfo.mapValues(_ => {
@@ -907,20 +1016,20 @@ class KafkaApis(val requestChannel: RequestChannel,
   private def createTopic(topic: String,
                           numPartitions: Int,
                           replicationFactor: Int,
-                          properties: Properties = new Properties()): MetadataResponse.TopicMetadata = {
+                          properties: util.Properties = new util.Properties()): MetadataResponse.TopicMetadata = {
     try {
       adminZkClient.createTopic(topic, numPartitions, replicationFactor, properties, RackAwareMode.Safe)
       info("Auto creation of topic %s with %d partitions and replication factor %d is successful"
         .format(topic, numPartitions, replicationFactor))
       new MetadataResponse.TopicMetadata(Errors.LEADER_NOT_AVAILABLE, topic, isInternal(topic),
-        java.util.Collections.emptyList())
+        util.Collections.emptyList())
     } catch {
       case _: TopicExistsException => // let it go, possibly another broker created this topic
         new MetadataResponse.TopicMetadata(Errors.LEADER_NOT_AVAILABLE, topic, isInternal(topic),
-          java.util.Collections.emptyList())
+          util.Collections.emptyList())
       case ex: Throwable  => // Catch all to prevent unhandled errors
         new MetadataResponse.TopicMetadata(Errors.forException(ex), topic, isInternal(topic),
-          java.util.Collections.emptyList())
+          util.Collections.emptyList())
     }
   }
 
@@ -937,7 +1046,7 @@ class KafkaApis(val requestChannel: RequestChannel,
             s"'${config.offsetsTopicReplicationFactor}' for the offsets topic (configured via " +
             s"'${KafkaConfig.OffsetsTopicReplicationFactorProp}'). This error can be ignored if the cluster is starting up " +
             s"and not all brokers are up yet.")
-          new MetadataResponse.TopicMetadata(Errors.COORDINATOR_NOT_AVAILABLE, topic, true, java.util.Collections.emptyList())
+          new MetadataResponse.TopicMetadata(Errors.COORDINATOR_NOT_AVAILABLE, topic, true, util.Collections.emptyList())
         } else {
           createTopic(topic, config.offsetsTopicPartitions, config.offsetsTopicReplicationFactor.toInt,
             groupCoordinator.offsetsTopicConfigs)
@@ -948,7 +1057,7 @@ class KafkaApis(val requestChannel: RequestChannel,
             s"'${config.transactionTopicReplicationFactor}' for the transactions state topic (configured via " +
             s"'${KafkaConfig.TransactionsTopicReplicationFactorProp}'). This error can be ignored if the cluster is starting up " +
             s"and not all brokers are up yet.")
-          new MetadataResponse.TopicMetadata(Errors.COORDINATOR_NOT_AVAILABLE, topic, true, java.util.Collections.emptyList())
+          new MetadataResponse.TopicMetadata(Errors.COORDINATOR_NOT_AVAILABLE, topic, true, util.Collections.emptyList())
         } else {
           createTopic(topic, config.transactionTopicPartitions, config.transactionTopicReplicationFactor.toInt,
             txnCoordinator.transactionTopicConfigs)
@@ -962,10 +1071,11 @@ class KafkaApis(val requestChannel: RequestChannel,
     topicMetadata.headOption.getOrElse(createInternalTopic(topic))
   }
 
-  private def getTopicMetadata(allowAutoTopicCreation: Boolean, topics: Set[String], listenerName: ListenerName,
+  private def getTopicMetadata(allowAutoTopicCreation: Boolean, topics: Set[String], requestContext: RequestContext,
                                errorUnavailableEndpoints: Boolean,
                                errorUnavailableListeners: Boolean): Seq[MetadataResponse.TopicMetadata] = {
-    val topicResponses = metadataCache.getTopicMetadata(topics, listenerName,
+
+    val topicResponses = metadataCache.getTopicMetadata(topics, requestContext.listenerName,
         errorUnavailableEndpoints, errorUnavailableListeners)
     if (topics.isEmpty || topicResponses.size == topics.size) {
       topicResponses
@@ -975,13 +1085,19 @@ class KafkaApis(val requestChannel: RequestChannel,
         if (isInternal(topic)) {
           val topicMetadata = createInternalTopic(topic)
           if (topicMetadata.error == Errors.COORDINATOR_NOT_AVAILABLE)
-            new MetadataResponse.TopicMetadata(Errors.INVALID_REPLICATION_FACTOR, topic, true, java.util.Collections.emptyList())
+            new MetadataResponse.TopicMetadata(Errors.INVALID_REPLICATION_FACTOR, topic, true, util.Collections.emptyList())
           else
             topicMetadata
         } else if (allowAutoTopicCreation && config.autoCreateTopicsEnable) {
+          val msg =  "Automatically creating topic: " + topic + " with " + config.numPartitions +
+                     " partitions and replication factor " + config.defaultReplicationFactor +
+                     " due to request from " + requestContext.principal + " at IP address " +
+            requestContext.clientAddress + " and header " + requestContext.header
+
+          info(msg);
           createTopic(topic, config.numPartitions, config.defaultReplicationFactor)
         } else {
-          new MetadataResponse.TopicMetadata(Errors.UNKNOWN_TOPIC_OR_PARTITION, topic, false, java.util.Collections.emptyList())
+          new MetadataResponse.TopicMetadata(Errors.UNKNOWN_TOPIC_OR_PARTITION, topic, false, util.Collections.emptyList())
         }
       }
       topicResponses ++ responsesForNonExistentTopics
@@ -991,7 +1107,7 @@ class KafkaApis(val requestChannel: RequestChannel,
   /**
    * Handle a topic metadata request
    */
-  def handleTopicMetadataRequest(request: RequestChannel.Request) {
+  def handleTopicMetadataRequest(request: RequestChannel.Request): Unit = {
     val metadataRequest = request.body[MetadataRequest]
     val requestVersion = request.header.apiVersion
 
@@ -1000,18 +1116,16 @@ class KafkaApis(val requestChannel: RequestChannel,
     else
       metadataRequest.topics.asScala.toSet
 
-    var (authorizedTopics, unauthorizedForDescribeTopics) =
-      topics.partition(topic => authorize(request.session, Describe, Resource(Topic, topic, LITERAL)))
-
+    val authorizedForDescribeTopics = filterAuthorized(request, DESCRIBE, TOPIC, topics.toSeq, logIfDenied = !metadataRequest.isAllTopics)
+    var (authorizedTopics, unauthorizedForDescribeTopics) = topics.partition(authorizedForDescribeTopics.contains)
     var unauthorizedForCreateTopics = Set[String]()
 
     if (authorizedTopics.nonEmpty) {
       val nonExistingTopics = metadataCache.getNonExistingTopics(authorizedTopics)
       if (metadataRequest.allowAutoTopicCreation && config.autoCreateTopicsEnable && nonExistingTopics.nonEmpty) {
-        if (!authorize(request.session, Create, Resource.ClusterResource)) {
-          unauthorizedForCreateTopics = nonExistingTopics.filter { topic =>
-            !authorize(request.session, Create, new Resource(Topic, topic, PatternType.LITERAL))
-          }
+        if (!authorize(request, CREATE, CLUSTER, CLUSTER_NAME, logIfDenied = false)) {
+          val authorizedForCreateTopics = filterAuthorized(request, CREATE, TOPIC, nonExistingTopics.toSeq)
+          unauthorizedForCreateTopics = nonExistingTopics -- authorizedForCreateTopics
           authorizedTopics --= unauthorizedForCreateTopics
         }
       }
@@ -1019,7 +1133,7 @@ class KafkaApis(val requestChannel: RequestChannel,
 
     val unauthorizedForCreateTopicMetadata = unauthorizedForCreateTopics.map(topic =>
       new MetadataResponse.TopicMetadata(Errors.TOPIC_AUTHORIZATION_FAILED, topic, isInternal(topic),
-        java.util.Collections.emptyList()))
+        util.Collections.emptyList()))
 
     // do not disclose the existence of topics unauthorized for Describe, so we've not even checked if they exist or not
     val unauthorizedForDescribeTopicMetadata =
@@ -1028,7 +1142,7 @@ class KafkaApis(val requestChannel: RequestChannel,
         Set.empty[MetadataResponse.TopicMetadata]
       else
         unauthorizedForDescribeTopics.map(topic =>
-          new MetadataResponse.TopicMetadata(Errors.TOPIC_AUTHORIZATION_FAILED, topic, false, java.util.Collections.emptyList()))
+          new MetadataResponse.TopicMetadata(Errors.TOPIC_AUTHORIZATION_FAILED, topic, false, util.Collections.emptyList()))
 
     // In version 0, we returned an error when brokers with replicas were unavailable,
     // while in higher versions we simply don't include the broker in the returned broker list
@@ -1039,9 +1153,27 @@ class KafkaApis(val requestChannel: RequestChannel,
     val topicMetadata =
       if (authorizedTopics.isEmpty)
         Seq.empty[MetadataResponse.TopicMetadata]
-      else
-        getTopicMetadata(metadataRequest.allowAutoTopicCreation, authorizedTopics, request.context.listenerName,
+      else {
+        // If the metadata request is fetching metadata for all topics, auto topic creation is NOT allowed. This is an
+        // internal hotfix for issue KAFKA-10606
+        val allowAutoTopicCreation = (!metadataRequest.isAllTopics) && metadataRequest.allowAutoTopicCreation
+        getTopicMetadata(allowAutoTopicCreation, authorizedTopics, request.context,
           errorUnavailableEndpoints, errorUnavailableListeners)
+      }
+
+    var clusterAuthorizedOperations = 0
+
+    if (request.header.apiVersion >= 8) {
+      // get cluster authorized operations
+      if (metadataRequest.data().includeClusterAuthorizedOperations() &&
+        authorize(request, DESCRIBE, CLUSTER, CLUSTER_NAME))
+        clusterAuthorizedOperations = authorizedOperations(request, Resource.CLUSTER)
+      // get topic authorized operations
+      if (metadataRequest.data().includeTopicAuthorizedOperations())
+        topicMetadata.foreach(topicData => {
+          topicData.authorizedOperations(authorizedOperations(request, new Resource(ResourceType.TOPIC, topicData.topic())))
+        })
+    }
 
     val completeTopicMetadata = topicMetadata ++ unauthorizedForCreateTopicMetadata ++ unauthorizedForDescribeTopicMetadata
 
@@ -1051,34 +1183,37 @@ class KafkaApis(val requestChannel: RequestChannel,
       brokers.mkString(","), request.header.correlationId, request.header.clientId))
 
     sendResponseMaybeThrottle(request, requestThrottleMs =>
-      new MetadataResponse(
-        requestThrottleMs,
-        brokers.flatMap(_.getNode(request.context.listenerName)).asJava,
-        clusterId,
-        metadataCache.getControllerId.getOrElse(MetadataResponse.NO_CONTROLLER_ID),
-        completeTopicMetadata.asJava
+       MetadataResponse.prepareResponse(
+         requestThrottleMs,
+         brokers.flatMap(_.getNode(request.context.listenerName)).asJava,
+         clusterId,
+         metadataCache.getControllerId.getOrElse(MetadataResponse.NO_CONTROLLER_ID),
+         completeTopicMetadata.asJava,
+         clusterAuthorizedOperations
       ))
   }
 
   /**
    * Handle an offset fetch request
    */
-  def handleOffsetFetchRequest(request: RequestChannel.Request) {
+  def handleOffsetFetchRequest(request: RequestChannel.Request): Unit = {
     val header = request.header
     val offsetFetchRequest = request.body[OffsetFetchRequest]
 
-    def authorizeTopicDescribe(partition: TopicPartition) =
-      authorize(request.session, Describe, Resource(Topic, partition.topic, LITERAL))
+    def partitionAuthorized[T](elements: List[T], topic: T => String): (Seq[T], Seq[T]) = {
+      val authorizedTopics = filterAuthorized(request, DESCRIBE, TOPIC, elements.map(topic))
+      elements.partition(element => authorizedTopics.contains(topic.apply(element)))
+    }
 
     def createResponse(requestThrottleMs: Int): AbstractResponse = {
       val offsetFetchResponse =
         // reject the request if not authorized to the group
-        if (!authorize(request.session, Describe, Resource(Group, offsetFetchRequest.groupId, LITERAL)))
+        if (!authorize(request, DESCRIBE, GROUP, offsetFetchRequest.groupId))
           offsetFetchRequest.getErrorResponse(requestThrottleMs, Errors.GROUP_AUTHORIZATION_FAILED)
         else {
           if (header.apiVersion == 0) {
-            val (authorizedPartitions, unauthorizedPartitions) = offsetFetchRequest.partitions.asScala
-              .partition(authorizeTopicDescribe)
+            val (authorizedPartitions, unauthorizedPartitions) =
+              partitionAuthorized[TopicPartition](offsetFetchRequest.partitions.asScala.toList, tp => tp.topic)
 
             // version 0 reads offsets from ZK
             val authorizedPartitionData = authorizedPartitions.map { topicPartition =>
@@ -1112,12 +1247,12 @@ class KafkaApis(val requestChannel: RequestChannel,
                 offsetFetchRequest.getErrorResponse(requestThrottleMs, error)
               else {
                 // clients are not allowed to see offsets for topics that are not authorized for Describe
-                val authorizedPartitionData = allPartitionData.filter { case (topicPartition, _) => authorizeTopicDescribe(topicPartition) }
-                new OffsetFetchResponse(requestThrottleMs, Errors.NONE, authorizedPartitionData.asJava)
+                val (authorizedPartitionData, _) = partitionAuthorized[(TopicPartition, OffsetFetchResponse.PartitionData)](allPartitionData.toList, e => e._1.topic)
+                new OffsetFetchResponse(requestThrottleMs, Errors.NONE, authorizedPartitionData.toMap.asJava)
               }
             } else {
-              val (authorizedPartitions, unauthorizedPartitions) = offsetFetchRequest.partitions.asScala
-                .partition(authorizeTopicDescribe)
+              val (authorizedPartitions, unauthorizedPartitions) =
+                partitionAuthorized[TopicPartition](offsetFetchRequest.partitions.asScala.toList, tp => tp.topic)
               val (error, authorizedPartitionData) = groupCoordinator.handleFetchOffsets(offsetFetchRequest.groupId,
                 Some(authorizedPartitions))
               if (error != Errors.NONE)
@@ -1137,25 +1272,25 @@ class KafkaApis(val requestChannel: RequestChannel,
     sendResponseMaybeThrottle(request, createResponse)
   }
 
-  def handleFindCoordinatorRequest(request: RequestChannel.Request) {
+  def handleFindCoordinatorRequest(request: RequestChannel.Request): Unit = {
     val findCoordinatorRequest = request.body[FindCoordinatorRequest]
 
-    if (findCoordinatorRequest.coordinatorType == FindCoordinatorRequest.CoordinatorType.GROUP &&
-        !authorize(request.session, Describe, Resource(Group, findCoordinatorRequest.coordinatorKey, LITERAL)))
+    if (findCoordinatorRequest.data.keyType == CoordinatorType.GROUP.id &&
+        !authorize(request, DESCRIBE, GROUP, findCoordinatorRequest.data.key))
       sendErrorResponseMaybeThrottle(request, Errors.GROUP_AUTHORIZATION_FAILED.exception)
-    else if (findCoordinatorRequest.coordinatorType == FindCoordinatorRequest.CoordinatorType.TRANSACTION &&
-        !authorize(request.session, Describe, Resource(TransactionalId, findCoordinatorRequest.coordinatorKey, LITERAL)))
+    else if (findCoordinatorRequest.data.keyType == CoordinatorType.TRANSACTION.id &&
+        !authorize(request, DESCRIBE, TRANSACTIONAL_ID, findCoordinatorRequest.data.key))
       sendErrorResponseMaybeThrottle(request, Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED.exception)
     else {
       // get metadata (and create the topic if necessary)
-      val (partition, topicMetadata) = findCoordinatorRequest.coordinatorType match {
-        case FindCoordinatorRequest.CoordinatorType.GROUP =>
-          val partition = groupCoordinator.partitionFor(findCoordinatorRequest.coordinatorKey)
+      val (partition, topicMetadata) = CoordinatorType.forId(findCoordinatorRequest.data.keyType) match {
+        case CoordinatorType.GROUP =>
+          val partition = groupCoordinator.partitionFor(findCoordinatorRequest.data.key)
           val metadata = getOrCreateInternalTopic(GROUP_METADATA_TOPIC_NAME, request.context.listenerName)
           (partition, metadata)
 
-        case FindCoordinatorRequest.CoordinatorType.TRANSACTION =>
-          val partition = txnCoordinator.partitionFor(findCoordinatorRequest.coordinatorKey)
+        case CoordinatorType.TRANSACTION =>
+          val partition = txnCoordinator.partitionFor(findCoordinatorRequest.data.key)
           val metadata = getOrCreateInternalTopic(TRANSACTION_STATE_TOPIC_NAME, request.context.listenerName)
           (partition, metadata)
 
@@ -1164,8 +1299,18 @@ class KafkaApis(val requestChannel: RequestChannel,
       }
 
       def createResponse(requestThrottleMs: Int): AbstractResponse = {
+        def createFindCoordinatorResponse(error: Errors, node: Node): FindCoordinatorResponse = {
+          new FindCoordinatorResponse(
+              new FindCoordinatorResponseData()
+                .setErrorCode(error.code)
+                .setErrorMessage(error.message)
+                .setNodeId(node.id)
+                .setHost(node.host)
+                .setPort(node.port)
+                .setThrottleTimeMs(requestThrottleMs))
+        }
         val responseBody = if (topicMetadata.error != Errors.NONE) {
-          new FindCoordinatorResponse(requestThrottleMs, Errors.COORDINATOR_NOT_AVAILABLE, Node.noNode)
+          createFindCoordinatorResponse(Errors.COORDINATOR_NOT_AVAILABLE, Node.noNode)
         } else {
           val coordinatorEndpoint = topicMetadata.partitionMetadata.asScala
             .find(_.partition == partition)
@@ -1174,9 +1319,9 @@ class KafkaApis(val requestChannel: RequestChannel,
 
           coordinatorEndpoint match {
             case Some(endpoint) if !endpoint.isEmpty =>
-              new FindCoordinatorResponse(requestThrottleMs, Errors.NONE, endpoint)
+              createFindCoordinatorResponse(Errors.NONE, endpoint)
             case _ =>
-              new FindCoordinatorResponse(requestThrottleMs, Errors.COORDINATOR_NOT_AVAILABLE, Node.noNode)
+              createFindCoordinatorResponse(Errors.COORDINATOR_NOT_AVAILABLE, Node.noNode)
           }
         }
         trace("Sending FindCoordinator response %s for correlation id %d to client %s."
@@ -1187,7 +1332,7 @@ class KafkaApis(val requestChannel: RequestChannel,
     }
   }
 
-  def handleDescribeGroupRequest(request: RequestChannel.Request) {
+  def handleDescribeGroupRequest(request: RequestChannel.Request): Unit = {
 
     def sendResponseCallback(describeGroupsResponseData: DescribeGroupsResponseData): Unit = {
       def createResponse(requestThrottleMs: Int): AbstractResponse = {
@@ -1201,18 +1346,18 @@ class KafkaApis(val requestChannel: RequestChannel,
     val describeGroupsResponseData = new DescribeGroupsResponseData()
 
     describeRequest.data().groups().asScala.foreach { groupId =>
-      val resource = Resource(Group, groupId, LITERAL)
-      if (!authorize(request.session, Describe, resource)) {
+      if (!authorize(request, DESCRIBE, GROUP, groupId)) {
         describeGroupsResponseData.groups().add(DescribeGroupsResponse.forError(groupId, Errors.GROUP_AUTHORIZATION_FAILED))
       } else {
         val (error, summary) = groupCoordinator.handleDescribeGroup(groupId)
         val members = summary.members.map { member =>
           new DescribeGroupsResponseData.DescribedGroupMember()
             .setMemberId(member.memberId)
+            .setGroupInstanceId(member.groupInstanceId.orNull)
             .setClientId(member.clientId)
             .setClientHost(member.clientHost)
             .setMemberAssignment(member.assignment)
-            .setMemberMetadata(member.assignment)
+            .setMemberMetadata(member.metadata)
         }
 
         val describedGroup = new DescribeGroupsResponseData.DescribedGroup()
@@ -1225,7 +1370,7 @@ class KafkaApis(val requestChannel: RequestChannel,
 
         if (request.header.apiVersion >= 3) {
           if (error == Errors.NONE && describeRequest.data().includeAuthorizedOperations()) {
-            describedGroup.setAuthorizedOperations(authorizedOperations(request.session, resource))
+            describedGroup.setAuthorizedOperations(authorizedOperations(request, new Resource(ResourceType.GROUP, groupId)))
           } else {
             describedGroup.setAuthorizedOperations(0)
           }
@@ -1238,39 +1383,49 @@ class KafkaApis(val requestChannel: RequestChannel,
     sendResponseCallback(describeGroupsResponseData)
   }
 
-
-  private def authorizedOperations(session: RequestChannel.Session, resource: Resource): Int = {
-    val authorizedOps = authorizer match {
-      case None => resource.resourceType.supportedOperations
-      case Some(auth) => resource.resourceType.supportedOperations
-        .filter(operation => authorize(session, operation, resource))
-    }
-
-    Utils.to32BitField(authorizedOps.map(operation => operation.toJava.code().asInstanceOf[JByte]).asJava)
-  }
-
-  def handleListGroupsRequest(request: RequestChannel.Request) {
+  def handleListGroupsRequest(request: RequestChannel.Request): Unit = {
     val (error, groups) = groupCoordinator.handleListGroups()
-    if (authorize(request.session, Describe, Resource.ClusterResource))
+    if (authorize(request, DESCRIBE, CLUSTER, CLUSTER_NAME))
       // With describe cluster access all groups are returned. We keep this alternative for backward compatibility.
       sendResponseMaybeThrottle(request, requestThrottleMs =>
-        new ListGroupsResponse(requestThrottleMs, error, groups.map { group => new ListGroupsResponse.Group(group.groupId, group.protocolType) }.asJava))
+        new ListGroupsResponse(new ListGroupsResponseData()
+            .setErrorCode(error.code())
+            .setGroups(groups.map { group => new ListGroupsResponseData.ListedGroup()
+              .setGroupId(group.groupId)
+              .setProtocolType(group.protocolType)}.asJava
+            )
+            .setThrottleTimeMs(requestThrottleMs)
+        ))
     else {
-      val filteredGroups = groups.filter(group => authorize(request.session, Describe, new Resource(Group, group.groupId, LITERAL)))
+      val filteredGroups = groups.filter(group => authorize(request, DESCRIBE, GROUP, group.groupId))
       sendResponseMaybeThrottle(request, requestThrottleMs =>
-        new ListGroupsResponse(requestThrottleMs, error, filteredGroups.map { group => new ListGroupsResponse.Group(group.groupId, group.protocolType) }.asJava))
+        new ListGroupsResponse(new ListGroupsResponseData()
+          .setErrorCode(error.code())
+          .setGroups(filteredGroups.map { group => new ListGroupsResponseData.ListedGroup()
+            .setGroupId(group.groupId)
+            .setProtocolType(group.protocolType)}.asJava
+          )
+          .setThrottleTimeMs(requestThrottleMs)
+        ))
     }
   }
 
-  def handleJoinGroupRequest(request: RequestChannel.Request) {
+  def handleJoinGroupRequest(request: RequestChannel.Request): Unit = {
     val joinGroupRequest = request.body[JoinGroupRequest]
 
     // the callback for sending a join-group response
-    def sendResponseCallback(joinResult: JoinGroupResult) {
-      val members = joinResult.members map { case (memberId, metadataArray) => (memberId, ByteBuffer.wrap(metadataArray)) }
+    def sendResponseCallback(joinResult: JoinGroupResult): Unit = {
       def createResponse(requestThrottleMs: Int): AbstractResponse = {
-        val responseBody = new JoinGroupResponse(requestThrottleMs, joinResult.error, joinResult.generationId,
-          joinResult.subProtocol, joinResult.memberId, joinResult.leaderId, members.asJava)
+        val responseBody = new JoinGroupResponse(
+          new JoinGroupResponseData()
+            .setThrottleTimeMs(requestThrottleMs)
+            .setErrorCode(joinResult.error.code())
+            .setGenerationId(joinResult.generationId)
+            .setProtocolName(joinResult.subProtocol)
+            .setLeader(joinResult.leaderId)
+            .setMemberId(joinResult.memberId)
+            .setMembers(joinResult.members.asJava)
+        )
 
         trace("Sending join group response %s for correlation id %d to client %s."
           .format(responseBody, request.header.correlationId, request.header.clientId))
@@ -1279,54 +1434,88 @@ class KafkaApis(val requestChannel: RequestChannel,
       sendResponseMaybeThrottle(request, createResponse)
     }
 
-    if (!authorize(request.session, Read, Resource(Group, joinGroupRequest.groupId(), LITERAL))) {
+    if (joinGroupRequest.data.groupInstanceId != null && config.interBrokerProtocolVersion < KAFKA_2_3_IV0) {
+      // Only enable static membership when IBP >= 2.3, because it is not safe for the broker to use the static member logic
+      // until we are sure that all brokers support it. If static group being loaded by an older coordinator, it will discard
+      // the group.instance.id field, so static members could accidentally become "dynamic", which leads to wrong states.
+      sendResponseCallback(JoinGroupResult(
+        List.empty,
+        JoinGroupResponse.UNKNOWN_MEMBER_ID,
+        JoinGroupResponse.UNKNOWN_GENERATION_ID,
+        JoinGroupResponse.UNKNOWN_PROTOCOL,
+        JoinGroupResponse.UNKNOWN_MEMBER_ID,
+        Errors.UNSUPPORTED_VERSION
+      ))
+    } else if (!authorize(request, READ, GROUP, joinGroupRequest.data().groupId())) {
       sendResponseMaybeThrottle(request, requestThrottleMs =>
         new JoinGroupResponse(
-          requestThrottleMs,
-          Errors.GROUP_AUTHORIZATION_FAILED,
-          JoinGroupResponse.UNKNOWN_GENERATION_ID,
-          JoinGroupResponse.UNKNOWN_PROTOCOL,
-          JoinGroupResponse.UNKNOWN_MEMBER_ID, // memberId
-          JoinGroupResponse.UNKNOWN_MEMBER_ID, // leaderId
-          Collections.emptyMap())
+          new JoinGroupResponseData()
+            .setThrottleTimeMs(requestThrottleMs)
+            .setErrorCode(Errors.GROUP_AUTHORIZATION_FAILED.code)
+            .setGenerationId(JoinGroupResponse.UNKNOWN_GENERATION_ID)
+            .setProtocolName(JoinGroupResponse.UNKNOWN_PROTOCOL)
+            .setLeader(JoinGroupResponse.UNKNOWN_MEMBER_ID)
+            .setMemberId(JoinGroupResponse.UNKNOWN_MEMBER_ID)
+            .setMembers(util.Collections.emptyList())
+        )
       )
-    } else {
+    }  else {
+      val groupInstanceId = Option(joinGroupRequest.data.groupInstanceId)
+
       // Only return MEMBER_ID_REQUIRED error if joinGroupRequest version is >= 4
-      val requireKnownMemberId = joinGroupRequest.version >= 4
+      // and groupInstanceId is configured to unknown.
+      val requireKnownMemberId = joinGroupRequest.version >= 4 && groupInstanceId.isEmpty
 
       // let the coordinator handle join-group
-      val protocols = joinGroupRequest.groupProtocols().asScala.map(protocol =>
-        (protocol.name, Utils.toArray(protocol.metadata))).toList
+      val protocols = joinGroupRequest.data.protocols.valuesList.asScala.map(protocol =>
+        (protocol.name, protocol.metadata)).toList
       groupCoordinator.handleJoinGroup(
-        joinGroupRequest.groupId,
-        joinGroupRequest.memberId,
+        joinGroupRequest.data.groupId,
+        joinGroupRequest.data.memberId,
+        groupInstanceId,
         requireKnownMemberId,
         request.header.clientId,
-        request.session.clientAddress.toString,
-        joinGroupRequest.rebalanceTimeout,
-        joinGroupRequest.sessionTimeout,
-        joinGroupRequest.protocolType,
+        request.context.clientAddress.toString,
+        joinGroupRequest.data.rebalanceTimeoutMs,
+        joinGroupRequest.data.sessionTimeoutMs,
+        joinGroupRequest.data.protocolType,
         protocols,
         sendResponseCallback)
     }
   }
 
-  def handleSyncGroupRequest(request: RequestChannel.Request) {
+  def handleSyncGroupRequest(request: RequestChannel.Request): Unit = {
     val syncGroupRequest = request.body[SyncGroupRequest]
 
-    def sendResponseCallback(memberState: Array[Byte], error: Errors) {
+    def sendResponseCallback(syncGroupResult: SyncGroupResult): Unit = {
       sendResponseMaybeThrottle(request, requestThrottleMs =>
-        new SyncGroupResponse(requestThrottleMs, error, ByteBuffer.wrap(memberState)))
-    }
-
-    if (!authorize(request.session, Read, Resource(Group, syncGroupRequest.groupId(), LITERAL))) {
-      sendResponseCallback(Array[Byte](), Errors.GROUP_AUTHORIZATION_FAILED)
+        new SyncGroupResponse(
+          new SyncGroupResponseData()
+            .setErrorCode(syncGroupResult.error.code)
+            .setAssignment(syncGroupResult.memberAssignment)
+            .setThrottleTimeMs(requestThrottleMs)
+        ))
+    }
+
+    if (syncGroupRequest.data.groupInstanceId != null && config.interBrokerProtocolVersion < KAFKA_2_3_IV0) {
+      // Only enable static membership when IBP >= 2.3, because it is not safe for the broker to use the static member logic
+      // until we are sure that all brokers support it. If static group being loaded by an older coordinator, it will discard
+      // the group.instance.id field, so static members could accidentally become "dynamic", which leads to wrong states.
+      sendResponseCallback(SyncGroupResult(Array[Byte](), Errors.UNSUPPORTED_VERSION))
+    } else if (!authorize(request, READ, GROUP, syncGroupRequest.data.groupId)) {
+      sendResponseCallback(SyncGroupResult(Array[Byte](), Errors.GROUP_AUTHORIZATION_FAILED))
     } else {
+      val assignmentMap = immutable.Map.newBuilder[String, Array[Byte]]
+      syncGroupRequest.data.assignments.asScala.foreach { assignment =>
+        assignmentMap += (assignment.memberId -> assignment.assignment)
+      }
+
       groupCoordinator.handleSyncGroup(
-        syncGroupRequest.groupId,
-        syncGroupRequest.generationId,
-        syncGroupRequest.memberId,
-        syncGroupRequest.groupAssignment().asScala.mapValues(Utils.toArray),
+        syncGroupRequest.data.groupId,
+        syncGroupRequest.data.generationId,
+        syncGroupRequest.data.memberId,
+        Option(syncGroupRequest.data.groupInstanceId),
+        assignmentMap.result,
         sendResponseCallback
       )
     }
@@ -1334,26 +1523,41 @@ class KafkaApis(val requestChannel: RequestChannel,
 
   def handleDeleteGroupsRequest(request: RequestChannel.Request): Unit = {
     val deleteGroupsRequest = request.body[DeleteGroupsRequest]
-    val groups = deleteGroupsRequest.groups.asScala.toSet
+    val groups = deleteGroupsRequest.data.groupsNames.asScala.toSet
 
     val (authorizedGroups, unauthorizedGroups) = groups.partition { group =>
-      authorize(request.session, Delete, Resource(Group, group, LITERAL))
+      authorize(request, DELETE, GROUP, group)
     }
 
     val groupDeletionResult = groupCoordinator.handleDeleteGroups(authorizedGroups) ++
       unauthorizedGroups.map(_ -> Errors.GROUP_AUTHORIZATION_FAILED)
 
-    sendResponseMaybeThrottle(request, requestThrottleMs =>
-      new DeleteGroupsResponse(requestThrottleMs, groupDeletionResult.asJava))
+    sendResponseMaybeThrottle(request, requestThrottleMs => {
+      val deletionCollections = new DeletableGroupResultCollection()
+      groupDeletionResult.foreach { case (groupId, error) =>
+        deletionCollections.add(new DeletableGroupResult()
+          .setGroupId(groupId)
+          .setErrorCode(error.code)
+        )
+      }
+
+      new DeleteGroupsResponse(new DeleteGroupsResponseData()
+        .setResults(deletionCollections)
+        .setThrottleTimeMs(requestThrottleMs)
+      )
+    })
   }
 
-  def handleHeartbeatRequest(request: RequestChannel.Request) {
+  def handleHeartbeatRequest(request: RequestChannel.Request): Unit = {
     val heartbeatRequest = request.body[HeartbeatRequest]
 
     // the callback for sending a heartbeat response
-    def sendResponseCallback(error: Errors) {
+    def sendResponseCallback(error: Errors): Unit = {
       def createResponse(requestThrottleMs: Int): AbstractResponse = {
-        val response = new HeartbeatResponse(requestThrottleMs, error)
+        val response = new HeartbeatResponse(
+          new HeartbeatResponseData()
+            .setThrottleTimeMs(requestThrottleMs)
+            .setErrorCode(error.code))
         trace("Sending heartbeat response %s for correlation id %d to client %s."
           .format(response, request.header.correlationId, request.header.clientId))
         response
@@ -1361,63 +1565,82 @@ class KafkaApis(val requestChannel: RequestChannel,
       sendResponseMaybeThrottle(request, createResponse)
     }
 
-    if (!authorize(request.session, Read, Resource(Group, heartbeatRequest.groupId, LITERAL))) {
+    if (heartbeatRequest.data.groupInstanceId != null && config.interBrokerProtocolVersion < KAFKA_2_3_IV0) {
+      // Only enable static membership when IBP >= 2.3, because it is not safe for the broker to use the static member logic
+      // until we are sure that all brokers support it. If static group being loaded by an older coordinator, it will discard
+      // the group.instance.id field, so static members could accidentally become "dynamic", which leads to wrong states.
+      sendResponseCallback(Errors.UNSUPPORTED_VERSION)
+    } else if (!authorize(request, READ, GROUP, heartbeatRequest.data.groupId)) {
       sendResponseMaybeThrottle(request, requestThrottleMs =>
-        new HeartbeatResponse(requestThrottleMs, Errors.GROUP_AUTHORIZATION_FAILED))
+        new HeartbeatResponse(
+            new HeartbeatResponseData()
+              .setThrottleTimeMs(requestThrottleMs)
+              .setErrorCode(Errors.GROUP_AUTHORIZATION_FAILED.code)))
     } else {
       // let the coordinator to handle heartbeat
       groupCoordinator.handleHeartbeat(
-        heartbeatRequest.groupId,
-        heartbeatRequest.memberId,
-        heartbeatRequest.groupGenerationId,
+        heartbeatRequest.data.groupId,
+        heartbeatRequest.data.memberId,
+        Option(heartbeatRequest.data.groupInstanceId),
+        heartbeatRequest.data.generationId,
         sendResponseCallback)
     }
   }
 
-  def handleLeaveGroupRequest(request: RequestChannel.Request) {
+  def handleLeaveGroupRequest(request: RequestChannel.Request): Unit = {
     val leaveGroupRequest = request.body[LeaveGroupRequest]
 
-    // the callback for sending a leave-group response
-    def sendResponseCallback(error: Errors) {
-      def createResponse(requestThrottleMs: Int): AbstractResponse = {
-        val response = new LeaveGroupResponse(new LeaveGroupResponseData()
-          .setThrottleTimeMs(requestThrottleMs).setErrorCode(error.code()))
-        trace("Sending leave group response %s for correlation id %d to client %s."
-                    .format(response, request.header.correlationId, request.header.clientId))
-        response
-      }
-      sendResponseMaybeThrottle(request, createResponse)
-    }
+    val members = leaveGroupRequest.members().asScala.toList
 
-    if (!authorize(request.session, Read, Resource(Group, leaveGroupRequest.data().groupId(), LITERAL))) {
-      sendResponseMaybeThrottle(request, requestThrottleMs =>
+    if (!authorize(request, READ, GROUP, leaveGroupRequest.data().groupId())) {
+      sendResponseMaybeThrottle(request, requestThrottleMs => {
         new LeaveGroupResponse(new LeaveGroupResponseData()
-          .setThrottleTimeMs(requestThrottleMs).setErrorCode(Errors.GROUP_AUTHORIZATION_FAILED.code())))
+          .setThrottleTimeMs(requestThrottleMs)
+          .setErrorCode(Errors.GROUP_AUTHORIZATION_FAILED.code)
+        )
+      })
     } else {
-      // let the coordinator to handle leave-group
+      def sendResponseCallback(leaveGroupResult : LeaveGroupResult): Unit = {
+        val memberResponses = leaveGroupResult.memberResponses.map(
+          leaveGroupResult =>
+            new MemberResponse()
+              .setErrorCode(leaveGroupResult.error.code)
+              .setMemberId(leaveGroupResult.memberId)
+              .setGroupInstanceId(leaveGroupResult.groupInstanceId.orNull)
+        )
+        def createResponse(requestThrottleMs: Int): AbstractResponse = {
+          new LeaveGroupResponse(
+            memberResponses.asJava,
+            leaveGroupResult.topLevelError,
+            requestThrottleMs,
+            leaveGroupRequest.version)
+        }
+        sendResponseMaybeThrottle(request, createResponse)
+      }
+
       groupCoordinator.handleLeaveGroup(
-        leaveGroupRequest.data().groupId(),
-        leaveGroupRequest.data().memberId(),
+        leaveGroupRequest.data.groupId,
+        members,
         sendResponseCallback)
     }
   }
 
-  def handleSaslHandshakeRequest(request: RequestChannel.Request) {
+  def handleSaslHandshakeRequest(request: RequestChannel.Request): Unit = {
     val responseData = new SaslHandshakeResponseData().setErrorCode(Errors.ILLEGAL_SASL_STATE.code)
     sendResponseMaybeThrottle(request, _ => new SaslHandshakeResponse(responseData))
   }
 
-  def handleSaslAuthenticateRequest(request: RequestChannel.Request) {
+  def handleSaslAuthenticateRequest(request: RequestChannel.Request): Unit = {
     val responseData = new SaslAuthenticateResponseData()
       .setErrorCode(Errors.ILLEGAL_SASL_STATE.code)
       .setErrorMessage("SaslAuthenticate request received after successful authentication")
     sendResponseMaybeThrottle(request, _ => new SaslAuthenticateResponse(responseData))
   }
 
-  def handleApiVersionsRequest(request: RequestChannel.Request) {
+  def handleApiVersionsRequest(request: RequestChannel.Request): Unit = {
     // Note that broker returns its full list of supported ApiKeys and versions regardless of current
     // authentication state (e.g., before SASL authentication on an SASL listener, do note that no
-    // Kafka protocol requests may take place on a SSL listener before the SSL handshake is finished).
+    // Kafka protocol requests may take place on an SSL listener before the SSL handshake is finished).
     // If this is considered to leak information about the broker version a workaround is to use SSL
     // with client authentication which is performed at an earlier stage of the connection where the
     // ApiVersionRequest is not available.
@@ -1425,15 +1648,46 @@ class KafkaApis(val requestChannel: RequestChannel,
       val apiVersionRequest = request.body[ApiVersionsRequest]
       if (apiVersionRequest.hasUnsupportedRequestVersion)
         apiVersionRequest.getErrorResponse(requestThrottleMs, Errors.UNSUPPORTED_VERSION.exception)
-      else
-        ApiVersionsResponse.apiVersionsResponse(requestThrottleMs,
-          config.interBrokerProtocolVersion.recordVersion.value)
+      else if (!apiVersionRequest.isValid)
+        apiVersionRequest.getErrorResponse(requestThrottleMs, Errors.INVALID_REQUEST.exception)
+      else {
+        /**
+          * HOTFIX LIKAFKA-16384: brokers should suggest the max ProduceRequest ApiVersion that supports the broker default
+          * configured message format version instead of always suggesting the maximum ApiVersion.
+          */
+        val maxProduceApiVersion: Short = config.logMessageFormatVersion.recordVersion match {
+          case RecordVersion.V0 => 1
+          case RecordVersion.V1 => 2
+          case RecordVersion.V2 => 8
+        }
+        val response = ApiVersionsResponse.apiVersionsResponse(requestThrottleMs, config.interBrokerProtocolVersion.recordVersion.value)
+        val apiVersions = response.data.apiKeys().asScala.toList.map { apiVersionResponseKey =>
+          if (apiVersionResponseKey.apiKey == ApiKeys.PRODUCE.id) {
+            new ApiVersionsResponseKey()
+              .setApiKey(ApiKeys.PRODUCE.id)
+              .setMinVersion(apiVersionResponseKey.minVersion())
+              .setMaxVersion(maxProduceApiVersion)
+          } else {
+            // Since the apiVersionResponseKey variable is a slice of the ImplicitLinkedHashCollection#Element object,
+            // we cannot insert the same object into the new list.
+            // Otherwise the prev and next pointers will be changed, and the original collection will be corrupted.
+            new ApiVersionsResponseKey()
+              .setApiKey(apiVersionResponseKey.apiKey())
+              .setMinVersion(apiVersionResponseKey.minVersion())
+              .setMaxVersion(apiVersionResponseKey.maxVersion())
+          }
+        }.asJava
+        new ApiVersionsResponse(new ApiVersionsResponseData()
+          .setApiKeys(new ApiVersionsResponseKeyCollection(apiVersions.iterator()))
+          .setErrorCode(Errors.NONE.code())
+          .setThrottleTimeMs(requestThrottleMs))
+      }
     }
     sendResponseMaybeThrottle(request, createResponseCallback)
   }
 
-  def handleCreateTopicsRequest(request: RequestChannel.Request) {
-    def sendResponseCallback(results: CreatableTopicResultSet): Unit = {
+  def handleCreateTopicsRequest(request: RequestChannel.Request): Unit = {
+    def sendResponseCallback(results: CreatableTopicResultCollection): Unit = {
       def createResponse(requestThrottleMs: Int): AbstractResponse = {
         val responseData = new CreateTopicsResponseData().
           setThrottleTimeMs(requestThrottleMs).
@@ -1447,7 +1701,7 @@ class KafkaApis(val requestChannel: RequestChannel,
     }
 
     val createTopicsRequest = request.body[CreateTopicsRequest]
-    val results = new CreatableTopicResultSet(createTopicsRequest.data().topics().size())
+    val results = new CreatableTopicResultCollection(createTopicsRequest.data().topics().size())
     if (!controller.isActive) {
       createTopicsRequest.data.topics.asScala.foreach { case topic =>
         results.add(new CreatableTopicResult().setName(topic.name()).
@@ -1458,35 +1712,49 @@ class KafkaApis(val requestChannel: RequestChannel,
       createTopicsRequest.data.topics.asScala.foreach { case topic =>
         results.add(new CreatableTopicResult().setName(topic.name()))
       }
-      val hasClusterAuthorization = authorize(request.session, Create, Resource.ClusterResource)
+      val hasClusterAuthorization = authorize(request, CREATE, CLUSTER, CLUSTER_NAME, logIfDenied = false)
+      val topics = createTopicsRequest.data.topics.asScala.map(_.name)
+      val authorizedTopics = if (hasClusterAuthorization) topics.toSet else filterAuthorized(request, CREATE, TOPIC, topics.toSeq)
+      val authorizedForDescribeConfigs = filterAuthorized(request, DESCRIBE_CONFIGS, TOPIC, topics.toSeq)
+        .map(name => name -> results.find(name)).toMap
+
       results.asScala.foreach(topic => {
         if (results.findAll(topic.name()).size() > 1) {
           topic.setErrorCode(Errors.INVALID_REQUEST.code())
           topic.setErrorMessage("Found multiple entries for this topic.")
-        } else if ((!hasClusterAuthorization) && (!authorize(request.session, Create,
-              new Resource(Topic, topic.name(), PatternType.LITERAL)))) {
+        } else if (!authorizedTopics.contains(topic.name)) {
           topic.setErrorCode(Errors.TOPIC_AUTHORIZATION_FAILED.code())
           topic.setErrorMessage("Authorization failed.")
         }
+        if (!authorizedForDescribeConfigs.contains(topic.name)) {
+          topic.setTopicConfigErrorCode(Errors.TOPIC_AUTHORIZATION_FAILED.code)
+        }
       })
       val toCreate = mutable.Map[String, CreatableTopic]()
-      createTopicsRequest.data.topics.asScala.foreach { case topic =>
-        if (results.find(topic.name()).errorCode() == 0) {
-          toCreate += topic.name() -> topic
+      createTopicsRequest.data.topics.asScala.foreach { topic =>
+        if (results.find(topic.name).errorCode == Errors.NONE.code) {
+          toCreate += topic.name -> topic
         }
       }
       def handleCreateTopicsResults(errors: Map[String, ApiError]): Unit = {
-        errors.foreach {
-          case (topicName, error) =>
-            results.find(topicName).
-              setErrorCode(error.error().code()).
-              setErrorMessage(error.message())
+        errors.foreach { case (topicName, error) =>
+          val result = results.find(topicName)
+          result.setErrorCode(error.error.code)
+            .setErrorMessage(error.message)
+          // Reset any configs in the response if Create failed
+          if (error != ApiError.NONE) {
+            result.setConfigs(List.empty.asJava)
+              .setNumPartitions(-1)
+              .setReplicationFactor(-1)
+              .setTopicConfigErrorCode(0.toShort)
+          }
         }
         sendResponseCallback(results)
       }
-      adminManager.createTopics(createTopicsRequest.data.timeoutMs(),
-          createTopicsRequest.data.validateOnly(),
+      adminManager.createTopics(createTopicsRequest.data.timeoutMs,
+          createTopicsRequest.data.validateOnly,
           toCreate,
+          authorizedForDescribeConfigs,
           handleCreateTopicsResults)
     }
   }
@@ -1513,8 +1781,9 @@ class KafkaApis(val requestChannel: RequestChannel,
       // Special handling to add duplicate topics to the response
       val dupes = createPartitionsRequest.duplicates.asScala
       val notDuped = createPartitionsRequest.newPartitions.asScala -- dupes
+      val authorizedTopics = filterAuthorized(request, ALTER, TOPIC, notDuped.toSeq.map(_._1))
       val (authorized, unauthorized) = notDuped.partition { case (topic, _) =>
-        authorize(request.session, Alter, Resource(Topic, topic, LITERAL))
+        authorizedTopics.contains(topic)
       }
 
       val (queuedForDeletion, valid) = authorized.partition { case (topic, _) =>
@@ -1530,66 +1799,84 @@ class KafkaApis(val requestChannel: RequestChannel,
     }
   }
 
-  def handleDeleteTopicsRequest(request: RequestChannel.Request) {
-    val deleteTopicRequest = request.body[DeleteTopicsRequest]
-
-    val unauthorizedTopicErrors = mutable.Map[String, Errors]()
-    val nonExistingTopicErrors = mutable.Map[String, Errors]()
-    val authorizedForDeleteTopics =  mutable.Set[String]()
-
-    for (topic <- deleteTopicRequest.topics.asScala) {
-      if (!authorize(request.session, Delete, Resource(Topic, topic, LITERAL)))
-        unauthorizedTopicErrors += topic -> Errors.TOPIC_AUTHORIZATION_FAILED
-      else if (!metadataCache.contains(topic))
-        nonExistingTopicErrors += topic -> Errors.UNKNOWN_TOPIC_OR_PARTITION
-      else
-        authorizedForDeleteTopics.add(topic)
-    }
-
-    def sendResponseCallback(authorizedTopicErrors: Map[String, Errors]): Unit = {
+  def handleDeleteTopicsRequest(request: RequestChannel.Request): Unit = {
+    def sendResponseCallback(results: DeletableTopicResultCollection): Unit = {
       def createResponse(requestThrottleMs: Int): AbstractResponse = {
-        val completeResults = unauthorizedTopicErrors ++ nonExistingTopicErrors ++ authorizedTopicErrors
-        val responseBody = new DeleteTopicsResponse(requestThrottleMs, completeResults.asJava)
+        val responseData = new DeleteTopicsResponseData()
+          .setThrottleTimeMs(requestThrottleMs)
+          .setResponses(results)
+        val responseBody = new DeleteTopicsResponse(responseData)
         trace(s"Sending delete topics response $responseBody for correlation id ${request.header.correlationId} to client ${request.header.clientId}.")
         responseBody
       }
       sendResponseMaybeThrottle(request, createResponse)
     }
 
+    val deleteTopicRequest = request.body[DeleteTopicsRequest]
+    val results = new DeletableTopicResultCollection(deleteTopicRequest.data.topicNames.size)
+    val toDelete = mutable.Set[String]()
     if (!controller.isActive) {
-      val results = deleteTopicRequest.topics.asScala.map { topic =>
-        (topic, Errors.NOT_CONTROLLER)
-      }.toMap
+      deleteTopicRequest.data.topicNames.asScala.foreach { case topic =>
+        results.add(new DeletableTopicResult()
+          .setName(topic)
+          .setErrorCode(Errors.NOT_CONTROLLER.code))
+      }
       sendResponseCallback(results)
     } else if (!config.deleteTopicEnable) {
       val error = if (request.context.apiVersion < 3) Errors.INVALID_REQUEST else Errors.TOPIC_DELETION_DISABLED
-      val results = deleteTopicRequest.topics.asScala.map { topic =>
-        (topic, error)
-      }.toMap
+      deleteTopicRequest.data.topicNames.asScala.foreach { case topic =>
+        results.add(new DeletableTopicResult()
+          .setName(topic)
+          .setErrorCode(error.code))
+      }
       sendResponseCallback(results)
     } else {
+      deleteTopicRequest.data.topicNames.asScala.foreach { case topic =>
+        results.add(new DeletableTopicResult()
+          .setName(topic))
+      }
+      val authorizedTopics = filterAuthorized(request, DELETE, TOPIC, results.asScala.toSeq.map(_.name))
+      results.asScala.foreach(topic => {
+         if (!authorizedTopics.contains(topic.name))
+           topic.setErrorCode(Errors.TOPIC_AUTHORIZATION_FAILED.code)
+         else if (!metadataCache.contains(topic.name))
+           topic.setErrorCode(Errors.UNKNOWN_TOPIC_OR_PARTITION.code)
+         else
+           toDelete += topic.name
+      })
       // If no authorized topics return immediately
-      if (authorizedForDeleteTopics.isEmpty)
-        sendResponseCallback(Map())
+      if (toDelete.isEmpty)
+        sendResponseCallback(results)
       else {
+        def handleDeleteTopicsResults(errors: Map[String, Errors]): Unit = {
+          errors.foreach {
+            case (topicName, error) =>
+              results.find(topicName)
+                .setErrorCode(error.code)
+          }
+          sendResponseCallback(results)
+        }
+
         adminManager.deleteTopics(
-          deleteTopicRequest.timeout.toInt,
-          authorizedForDeleteTopics,
-          sendResponseCallback
+          deleteTopicRequest.data.timeoutMs.toInt,
+          toDelete,
+          handleDeleteTopicsResults
         )
       }
     }
   }
 
-  def handleDeleteRecordsRequest(request: RequestChannel.Request) {
+  def handleDeleteRecordsRequest(request: RequestChannel.Request): Unit = {
     val deleteRecordsRequest = request.body[DeleteRecordsRequest]
 
     val unauthorizedTopicResponses = mutable.Map[TopicPartition, DeleteRecordsResponse.PartitionResponse]()
     val nonExistingTopicResponses = mutable.Map[TopicPartition, DeleteRecordsResponse.PartitionResponse]()
     val authorizedForDeleteTopicOffsets = mutable.Map[TopicPartition, Long]()
 
+    val authorizedTopics = filterAuthorized(request, DELETE, TOPIC,
+      deleteRecordsRequest.partitionOffsets.asScala.toSeq.map(_._1.topic))
     for ((topicPartition, offset) <- deleteRecordsRequest.partitionOffsets.asScala) {
-      if (!authorize(request.session, Delete, Resource(Topic, topicPartition.topic, LITERAL)))
+      if (!authorizedTopics.contains(topicPartition.topic))
         unauthorizedTopicResponses += topicPartition -> new DeleteRecordsResponse.PartitionResponse(
           DeleteRecordsResponse.INVALID_LOW_WATERMARK, Errors.TOPIC_AUTHORIZATION_FAILED)
       else if (!metadataCache.contains(topicPartition))
@@ -1600,7 +1887,7 @@ class KafkaApis(val requestChannel: RequestChannel,
     }
 
     // the callback for sending a DeleteRecordsResponse
-    def sendResponseCallback(authorizedTopicResponses: Map[TopicPartition, DeleteRecordsResponse.PartitionResponse]) {
+    def sendResponseCallback(authorizedTopicResponses: Map[TopicPartition, DeleteRecordsResponse.PartitionResponse]): Unit = {
       val mergedResponseStatus = authorizedTopicResponses ++ unauthorizedTopicResponses ++ nonExistingTopicResponses
       mergedResponseStatus.foreach { case (topicPartition, status) =>
         if (status.error != Errors.NONE) {
@@ -1629,27 +1916,32 @@ class KafkaApis(val requestChannel: RequestChannel,
 
   def handleInitProducerIdRequest(request: RequestChannel.Request): Unit = {
     val initProducerIdRequest = request.body[InitProducerIdRequest]
-    val transactionalId = initProducerIdRequest.transactionalId
+    val transactionalId = initProducerIdRequest.data.transactionalId
 
     if (transactionalId != null) {
-      if (!authorize(request.session, Write, Resource(TransactionalId, transactionalId, LITERAL))) {
+      if (!authorize(request, WRITE, TRANSACTIONAL_ID, transactionalId)) {
         sendErrorResponseMaybeThrottle(request, Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED.exception)
         return
       }
-    } else if (!authorize(request.session, IdempotentWrite, Resource.ClusterResource)) {
+    } else if (!authorize(request, IDEMPOTENT_WRITE, CLUSTER, CLUSTER_NAME)) {
       sendErrorResponseMaybeThrottle(request, Errors.CLUSTER_AUTHORIZATION_FAILED.exception)
       return
     }
 
     def sendResponseCallback(result: InitProducerIdResult): Unit = {
       def createResponse(requestThrottleMs: Int): AbstractResponse = {
-        val responseBody = new InitProducerIdResponse(requestThrottleMs, result.error, result.producerId, result.producerEpoch)
+        val responseData = new InitProducerIdResponseData()
+          .setProducerId(result.producerId)
+          .setProducerEpoch(result.producerEpoch)
+          .setThrottleTimeMs(requestThrottleMs)
+          .setErrorCode(result.error.code)
+        val responseBody = new InitProducerIdResponse(responseData)
         trace(s"Completed $transactionalId's InitProducerIdRequest with result $result from client ${request.header.clientId}.")
         responseBody
       }
       sendResponseMaybeThrottle(request, createResponse)
     }
-    txnCoordinator.handleInitProducerId(transactionalId, initProducerIdRequest.transactionTimeoutMs, sendResponseCallback)
+    txnCoordinator.handleInitProducerId(transactionalId, initProducerIdRequest.data.transactionTimeoutMs, sendResponseCallback)
   }
 
   def handleEndTxnRequest(request: RequestChannel.Request): Unit = {
@@ -1657,8 +1949,8 @@ class KafkaApis(val requestChannel: RequestChannel,
     val endTxnRequest = request.body[EndTxnRequest]
     val transactionalId = endTxnRequest.transactionalId
 
-    if (authorize(request.session, Write, Resource(TransactionalId, transactionalId, LITERAL))) {
-      def sendResponseCallback(error: Errors) {
+    if (authorize(request, WRITE, TRANSACTIONAL_ID, transactionalId)) {
+      def sendResponseCallback(error: Errors): Unit = {
         def createResponse(requestThrottleMs: Int): AbstractResponse = {
           val responseBody = new EndTxnResponse(requestThrottleMs, error)
           trace(s"Completed ${endTxnRequest.transactionalId}'s EndTxnRequest with command: ${endTxnRequest.command}, errors: $error from client ${request.header.clientId}.")
@@ -1679,7 +1971,7 @@ class KafkaApis(val requestChannel: RequestChannel,
 
   def handleWriteTxnMarkersRequest(request: RequestChannel.Request): Unit = {
     ensureInterBrokerVersion(KAFKA_0_11_0_IV0)
-    authorizeClusterAction(request)
+    authorizeClusterOperation(request, CLUSTER_ACTION)
     val writeTxnMarkersRequest = request.body[WriteTxnMarkersRequest]
     val errors = new ConcurrentHashMap[java.lang.Long, util.Map[TopicPartition, Errors]]()
     val markers = writeTxnMarkersRequest.markers
@@ -1704,7 +1996,7 @@ class KafkaApis(val requestChannel: RequestChannel,
       */
     def maybeSendResponseCallback(producerId: Long, result: TransactionResult)(responseStatus: Map[TopicPartition, PartitionResponse]): Unit = {
       trace(s"End transaction marker append for producer id $producerId completed with status: $responseStatus")
-      val currentErrors = new ConcurrentHashMap[TopicPartition, Errors](responseStatus.mapValues(_.error).asJava)
+      val currentErrors = new ConcurrentHashMap[TopicPartition, Errors](responseStatus.map { case (k, v) => k -> v.error }.asJava)
       updateErrors(producerId, currentErrors)
       val successfulOffsetsPartitions = responseStatus.filter { case (topicPartition, partitionResponse) =>
         topicPartition.topic == GROUP_METADATA_TOPIC_NAME && partitionResponse.error == Errors.NONE
@@ -1770,7 +2062,7 @@ class KafkaApis(val requestChannel: RequestChannel,
           timeout = config.requestTimeoutMs.toLong,
           requiredAcks = -1,
           internalTopicsAllowed = true,
-          isFromClient = false,
+          origin = AppendOrigin.Coordinator,
           entriesPerPartition = controlRecords,
           responseCallback = maybeSendResponseCallback(producerId, marker.transactionResult))
       }
@@ -1792,7 +2084,7 @@ class KafkaApis(val requestChannel: RequestChannel,
     val addPartitionsToTxnRequest = request.body[AddPartitionsToTxnRequest]
     val transactionalId = addPartitionsToTxnRequest.transactionalId
     val partitionsToAdd = addPartitionsToTxnRequest.partitions.asScala
-    if (!authorize(request.session, Write, Resource(TransactionalId, transactionalId, LITERAL)))
+    if (!authorize(request, WRITE, TRANSACTIONAL_ID, transactionalId))
       sendResponseMaybeThrottle(request, requestThrottleMs =>
         addPartitionsToTxnRequest.getErrorResponse(requestThrottleMs, Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED.exception))
     else {
@@ -1800,9 +2092,10 @@ class KafkaApis(val requestChannel: RequestChannel,
       val nonExistingTopicErrors = mutable.Map[TopicPartition, Errors]()
       val authorizedPartitions = mutable.Set[TopicPartition]()
 
+      val authorizedTopics = filterAuthorized(request, WRITE, TOPIC,
+        partitionsToAdd.toSeq.map(_.topic).filterNot(org.apache.kafka.common.internals.Topic.isInternal))
       for (topicPartition <- partitionsToAdd) {
-        if (org.apache.kafka.common.internals.Topic.isInternal(topicPartition.topic) ||
-            !authorize(request.session, Write, Resource(Topic, topicPartition.topic, LITERAL)))
+        if (!authorizedTopics.contains(topicPartition.topic))
           unauthorizedTopicErrors += topicPartition -> Errors.TOPIC_AUTHORIZATION_FAILED
         else if (!metadataCache.contains(topicPartition))
           nonExistingTopicErrors += topicPartition -> Errors.UNKNOWN_TOPIC_OR_PARTITION
@@ -1846,10 +2139,10 @@ class KafkaApis(val requestChannel: RequestChannel,
     val groupId = addOffsetsToTxnRequest.consumerGroupId
     val offsetTopicPartition = new TopicPartition(GROUP_METADATA_TOPIC_NAME, groupCoordinator.partitionFor(groupId))
 
-    if (!authorize(request.session, Write, Resource(TransactionalId, transactionalId, LITERAL)))
+    if (!authorize(request, WRITE, TRANSACTIONAL_ID, transactionalId))
       sendResponseMaybeThrottle(request, requestThrottleMs =>
         new AddOffsetsToTxnResponse(requestThrottleMs, Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED))
-    else if (!authorize(request.session, Read, Resource(Group, groupId, LITERAL)))
+    else if (!authorize(request, READ, GROUP, groupId))
       sendResponseMaybeThrottle(request, requestThrottleMs =>
         new AddOffsetsToTxnResponse(requestThrottleMs, Errors.GROUP_AUTHORIZATION_FAILED))
     else {
@@ -1878,17 +2171,18 @@ class KafkaApis(val requestChannel: RequestChannel,
 
     // authorize for the transactionalId and the consumer group. Note that we skip producerId authorization
     // since it is implied by transactionalId authorization
-    if (!authorize(request.session, Write, Resource(TransactionalId, txnOffsetCommitRequest.transactionalId, LITERAL)))
+    if (!authorize(request, WRITE, TRANSACTIONAL_ID, txnOffsetCommitRequest.data.transactionalId))
       sendErrorResponseMaybeThrottle(request, Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED.exception)
-    else if (!authorize(request.session, Read, Resource(Group, txnOffsetCommitRequest.consumerGroupId, LITERAL)))
+    else if (!authorize(request, READ, GROUP, txnOffsetCommitRequest.data.groupId))
       sendErrorResponseMaybeThrottle(request, Errors.GROUP_AUTHORIZATION_FAILED.exception)
     else {
       val unauthorizedTopicErrors = mutable.Map[TopicPartition, Errors]()
       val nonExistingTopicErrors = mutable.Map[TopicPartition, Errors]()
       val authorizedTopicCommittedOffsets = mutable.Map[TopicPartition, TxnOffsetCommitRequest.CommittedOffset]()
+      val authorizedTopics = filterAuthorized(request, READ, TOPIC, txnOffsetCommitRequest.offsets.keySet.asScala.toSeq.map(_.topic))
 
       for ((topicPartition, commitedOffset) <- txnOffsetCommitRequest.offsets.asScala) {
-        if (!authorize(request.session, Read, Resource(Topic, topicPartition.topic, LITERAL)))
+        if (!authorizedTopics.contains(topicPartition.topic))
           unauthorizedTopicErrors += topicPartition -> Errors.TOPIC_AUTHORIZATION_FAILED
         else if (!metadataCache.contains(topicPartition))
           nonExistingTopicErrors += topicPartition -> Errors.UNKNOWN_TOPIC_OR_PARTITION
@@ -1897,7 +2191,7 @@ class KafkaApis(val requestChannel: RequestChannel,
       }
 
       // the callback for sending an offset commit response
-      def sendResponseCallback(authorizedTopicErrors: Map[TopicPartition, Errors]) {
+      def sendResponseCallback(authorizedTopicErrors: Map[TopicPartition, Errors]): Unit = {
         val combinedCommitStatus = authorizedTopicErrors ++ unauthorizedTopicErrors ++ nonExistingTopicErrors
         if (isDebugEnabled)
           combinedCommitStatus.foreach { case (topicPartition, error) =>
@@ -1915,9 +2209,9 @@ class KafkaApis(val requestChannel: RequestChannel,
       else {
         val offsetMetadata = convertTxnOffsets(authorizedTopicCommittedOffsets.toMap)
         groupCoordinator.handleTxnCommitOffsets(
-          txnOffsetCommitRequest.consumerGroupId,
-          txnOffsetCommitRequest.producerId,
-          txnOffsetCommitRequest.producerEpoch,
+          txnOffsetCommitRequest.data.groupId,
+          txnOffsetCommitRequest.data.producerId,
+          txnOffsetCommitRequest.data.producerEpoch,
           offsetMetadata,
           sendResponseCallback)
       }
@@ -1938,64 +2232,69 @@ class KafkaApis(val requestChannel: RequestChannel,
   }
 
   def handleDescribeAcls(request: RequestChannel.Request): Unit = {
-    authorizeClusterDescribe(request)
+    authorizeClusterOperation(request, DESCRIBE)
     val describeAclsRequest = request.body[DescribeAclsRequest]
     authorizer match {
       case None =>
         sendResponseMaybeThrottle(request, requestThrottleMs =>
           new DescribeAclsResponse(requestThrottleMs,
-            new ApiError(Errors.SECURITY_DISABLED, "No Authorizer is configured on the broker"), Collections.emptySet()))
+            new ApiError(Errors.SECURITY_DISABLED, "No Authorizer is configured on the broker"), util.Collections.emptySet()))
       case Some(auth) =>
-        val filter = describeAclsRequest.filter()
-        val returnedAcls = auth.getAcls.toSeq.flatMap { case (resource, acls) =>
-          acls.flatMap { acl =>
-            val fixture = new AclBinding(new ResourcePattern(resource.resourceType.toJava, resource.name, resource.patternType),
-                new AccessControlEntry(acl.principal.toString, acl.host.toString, acl.operation.toJava, acl.permissionType.toJava))
-            Some(fixture).filter(filter.matches)
-          }
-        }
+        val filter = describeAclsRequest.filter
+        val returnedAcls = new util.HashSet[AclBinding]()
+        auth.acls(filter).asScala.foreach(returnedAcls.add)
         sendResponseMaybeThrottle(request, requestThrottleMs =>
-          new DescribeAclsResponse(requestThrottleMs, ApiError.NONE, returnedAcls.asJava))
+          new DescribeAclsResponse(requestThrottleMs, ApiError.NONE, returnedAcls))
     }
   }
 
   def handleCreateAcls(request: RequestChannel.Request): Unit = {
-    authorizeClusterAlter(request)
+    authorizeClusterOperation(request, ALTER)
     val createAclsRequest = request.body[CreateAclsRequest]
+
     authorizer match {
       case None =>
         sendResponseMaybeThrottle(request, requestThrottleMs =>
           createAclsRequest.getErrorResponse(requestThrottleMs,
             new SecurityDisabledException("No Authorizer is configured on the broker.")))
       case Some(auth) =>
-        val aclCreationResults = createAclsRequest.aclCreations.asScala.map { aclCreation =>
-          SecurityUtils.convertToResourceAndAcl(aclCreation.acl.toFilter) match {
-            case Left(apiError) => new AclCreationResponse(apiError)
-            case Right((resource, acl)) => try {
-                if (resource.resourceType.equals(Cluster) && !SecurityUtils.isClusterResource(resource.name))
-                  throw new InvalidRequestException("The only valid name for the CLUSTER resource is " +
-                      Resource.ClusterResourceName)
-                if (resource.name.isEmpty)
-                  throw new InvalidRequestException("Invalid empty resource name")
-                auth.addAcls(immutable.Set(acl), resource)
-
-                debug(s"Added acl $acl to $resource")
-
-                new AclCreationResponse(ApiError.NONE)
-              } catch {
-                case throwable: Throwable =>
-                  debug(s"Failed to add acl $acl to $resource", throwable)
-                  new AclCreationResponse(ApiError.fromThrowable(throwable))
-              }
+
+        val errorResults = mutable.Map[AclBinding, AclCreateResult]()
+        val aclBindings = createAclsRequest.aclCreations.asScala.map(_.acl)
+        val validBindings = aclBindings
+          .filter { acl =>
+            val resource = acl.pattern
+            val throwable = if (resource.resourceType == ResourceType.CLUSTER && !AuthorizerUtils.isClusterResource(resource.name))
+                new InvalidRequestException("The only valid name for the CLUSTER resource is " + CLUSTER_NAME)
+            else if (resource.name.isEmpty)
+              new InvalidRequestException("Invalid empty resource name")
+            else
+              null
+            if (throwable != null) {
+              debug(s"Failed to add acl $acl to $resource", throwable)
+              errorResults(acl) = new AclCreateResult(throwable)
+              false
+            } else
+              true
+          }
+
+        val createResults = auth.createAcls(request.context, validBindings.asJava)
+          .asScala.map(_.toCompletableFuture).toList
+        def sendResponseCallback(): Unit = {
+          val aclCreationResults = aclBindings.map { acl =>
+            val result = errorResults.getOrElse(acl, createResults(validBindings.indexOf(acl)).get)
+            new AclCreationResponse(result.exception.asScala.map(ApiError.fromThrowable).getOrElse(ApiError.NONE))
           }
+          sendResponseMaybeThrottle(request, requestThrottleMs =>
+            new CreateAclsResponse(requestThrottleMs, aclCreationResults.asJava))
         }
-        sendResponseMaybeThrottle(request, requestThrottleMs =>
-          new CreateAclsResponse(requestThrottleMs, aclCreationResults.asJava))
+
+        alterAclsPurgatory.tryCompleteElseWatch(config.connectionsMaxIdleMs, createResults, sendResponseCallback)
     }
   }
 
   def handleDeleteAcls(request: RequestChannel.Request): Unit = {
-    authorizeClusterAlter(request)
+    authorizeClusterOperation(request, ALTER)
     val deleteAclsRequest = request.body[DeleteAclsRequest]
     authorizer match {
       case None =>
@@ -2003,53 +2302,25 @@ class KafkaApis(val requestChannel: RequestChannel,
           deleteAclsRequest.getErrorResponse(requestThrottleMs,
             new SecurityDisabledException("No Authorizer is configured on the broker.")))
       case Some(auth) =>
-        val filters = deleteAclsRequest.filters.asScala
-        val filterResponseMap = mutable.Map[Int, AclFilterResponse]()
-        val toDelete = mutable.Map[Int, ArrayBuffer[(Resource, Acl)]]()
-
-        if (filters.forall(_.matchesAtMostOne)) {
-          // Delete based on a list of ACL fixtures.
-          for ((filter, i) <- filters.zipWithIndex) {
-            SecurityUtils.convertToResourceAndAcl(filter) match {
-              case Left(apiError) => filterResponseMap.put(i, new AclFilterResponse(apiError, Seq.empty.asJava))
-              case Right(binding) => toDelete.put(i, ArrayBuffer(binding))
-            }
-          }
-        } else {
-          // Delete based on filters that may match more than one ACL.
-          val aclMap = auth.getAcls()
-          val filtersWithIndex = filters.zipWithIndex
-          for ((resource, acls) <- aclMap; acl <- acls) {
-            val binding = new AclBinding(
-              new ResourcePattern(resource.resourceType.toJava, resource.name, resource.patternType),
-              new AccessControlEntry(acl.principal.toString, acl.host.toString, acl.operation.toJava,
-                acl.permissionType.toJava))
-
-            for ((filter, i) <- filtersWithIndex if filter.matches(binding))
-              toDelete.getOrElseUpdate(i, ArrayBuffer.empty) += ((resource, acl))
-          }
-        }
 
-        for ((i, acls) <- toDelete) {
-          val deletionResults = acls.flatMap { case (resource, acl) =>
-            val aclBinding = SecurityUtils.convertToAclBinding(resource, acl)
-            try {
-              if (auth.removeAcls(immutable.Set(acl), resource))
-                Some(new AclDeletionResult(aclBinding))
-              else None
-            } catch {
-              case throwable: Throwable =>
-                Some(new AclDeletionResult(ApiError.fromThrowable(throwable), aclBinding))
-            }
-          }.asJava
+        val deleteResults = auth.deleteAcls(request.context, deleteAclsRequest.filters)
+          .asScala.map(_.toCompletableFuture).toList
 
-          filterResponseMap.put(i, new AclFilterResponse(deletionResults))
+        def toErrorCode(exception: Optional[ApiException]): ApiError = {
+          exception.asScala.map(ApiError.fromThrowable).getOrElse(ApiError.NONE)
         }
 
-        val filterResponses = filters.indices.map { i =>
-          filterResponseMap.getOrElse(i, new AclFilterResponse(Seq.empty.asJava))
-        }.asJava
-        sendResponseMaybeThrottle(request, requestThrottleMs => new DeleteAclsResponse(requestThrottleMs, filterResponses))
+        def sendResponseCallback(): Unit = {
+          val filterResponses = deleteResults.map(_.get).map { result =>
+            val deletions = result.aclBindingDeleteResults().asScala.toList.map { deletionResult =>
+              new AclDeletionResult(toErrorCode(deletionResult.exception), deletionResult.aclBinding)
+            }.asJava
+            new AclFilterResponse(toErrorCode(result.exception), deletions)
+          }.asJava
+          sendResponseMaybeThrottle(request, requestThrottleMs =>
+            new DeleteAclsResponse(requestThrottleMs, filterResponses))
+        }
+        alterAclsPurgatory.tryCompleteElseWatch(config.connectionsMaxIdleMs, deleteResults, sendResponseCallback)
     }
   }
 
@@ -2060,11 +2331,12 @@ class KafkaApis(val requestChannel: RequestChannel,
     // The OffsetsForLeaderEpoch API was initially only used for inter-broker communication and required
     // cluster permission. With KIP-320, the consumer now also uses this API to check for log truncation
     // following a leader change, so we also allow topic describe permission.
-    val (authorizedPartitions, unauthorizedPartitions) = if (isAuthorizedClusterAction(request)) {
+    val (authorizedPartitions, unauthorizedPartitions) = if (authorize(request, CLUSTER_ACTION, CLUSTER, CLUSTER_NAME, logIfDenied = false)) {
       (requestInfo, Map.empty[TopicPartition, OffsetsForLeaderEpochRequest.PartitionData])
     } else {
+      val authorizedTopics = filterAuthorized(request, DESCRIBE, TOPIC, requestInfo.keySet.toSeq.map(_.topic))
       requestInfo.partition {
-        case (tp, _) => authorize(request.session, Describe, Resource(Topic, tp.topic, LITERAL))
+        case (tp, _) => authorizedTopics.contains(tp.topic)
       }
     }
 
@@ -2082,37 +2354,153 @@ class KafkaApis(val requestChannel: RequestChannel,
     val alterConfigsRequest = request.body[AlterConfigsRequest]
     val (authorizedResources, unauthorizedResources) = alterConfigsRequest.configs.asScala.partition { case (resource, _) =>
       resource.`type` match {
+        case ConfigResource.Type.BROKER_LOGGER =>
+          throw new InvalidRequestException(s"AlterConfigs is deprecated and does not support the resource type ${ConfigResource.Type.BROKER_LOGGER}")
         case ConfigResource.Type.BROKER =>
-          authorize(request.session, AlterConfigs, Resource.ClusterResource)
+          authorize(request, ALTER_CONFIGS, CLUSTER, CLUSTER_NAME)
         case ConfigResource.Type.TOPIC =>
-          authorize(request.session, AlterConfigs, Resource(Topic, resource.name, LITERAL))
+          authorize(request, ALTER_CONFIGS, TOPIC, resource.name)
         case rt => throw new InvalidRequestException(s"Unexpected resource type $rt")
       }
     }
     val authorizedResult = adminManager.alterConfigs(authorizedResources, alterConfigsRequest.validateOnly)
     val unauthorizedResult = unauthorizedResources.keys.map { resource =>
-      resource -> configsAuthorizationApiError(request.session, resource)
+      resource -> configsAuthorizationApiError(resource)
     }
     sendResponseMaybeThrottle(request, requestThrottleMs =>
       new AlterConfigsResponse(requestThrottleMs, (authorizedResult ++ unauthorizedResult).asJava))
   }
 
-  private def configsAuthorizationApiError(session: RequestChannel.Session, resource: ConfigResource): ApiError = {
+  def handleAlterPartitionReassignmentsRequest(request: RequestChannel.Request): Unit = {
+    authorizeClusterOperation(request, ALTER)
+    val alterPartitionReassignmentsRequest = request.body[AlterPartitionReassignmentsRequest]
+
+    def sendResponseCallback(result: Either[Map[TopicPartition, ApiError], ApiError]): Unit = {
+      val responseData = result match {
+        case Right(topLevelError) =>
+          new AlterPartitionReassignmentsResponseData().setErrorMessage(topLevelError.message()).setErrorCode(topLevelError.error().code())
+
+        case Left(assignments) =>
+          val topicResponses = assignments.groupBy(_._1.topic()).map {
+            case (topic, reassignmentsByTp) =>
+              val partitionResponses = reassignmentsByTp.map {
+                case (topicPartition, error) =>
+                  new ReassignablePartitionResponse().setPartitionIndex(topicPartition.partition())
+                    .setErrorCode(error.error().code()).setErrorMessage(error.message())
+              }
+              new ReassignableTopicResponse().setName(topic).setPartitions(partitionResponses.toList.asJava)
+          }
+          new AlterPartitionReassignmentsResponseData().setResponses(topicResponses.toList.asJava)
+      }
+
+      sendResponseMaybeThrottle(request, requestThrottleMs =>
+        new AlterPartitionReassignmentsResponse(responseData.setThrottleTimeMs(requestThrottleMs))
+      )
+    }
+
+    val reassignments = alterPartitionReassignmentsRequest.data().topics().asScala.flatMap {
+      reassignableTopic => reassignableTopic.partitions().asScala.map {
+        reassignablePartition =>
+          val tp = new TopicPartition(reassignableTopic.name(), reassignablePartition.partitionIndex())
+          if (reassignablePartition.replicas() == null)
+            tp -> None // revert call
+          else
+            tp -> Some(reassignablePartition.replicas().asScala.map(_.toInt))
+      }
+    }.toMap
+
+    controller.alterPartitionReassignments(reassignments, sendResponseCallback)
+  }
+
+  def handleListPartitionReassignmentsRequest(request: RequestChannel.Request): Unit = {
+    authorizeClusterOperation(request, DESCRIBE)
+    val listPartitionReassignmentsRequest = request.body[ListPartitionReassignmentsRequest]
+
+    def sendResponseCallback(result: Either[Map[TopicPartition, ReplicaAssignment], ApiError]): Unit = {
+      val responseData = result match {
+        case Right(error) => new ListPartitionReassignmentsResponseData().setErrorMessage(error.message()).setErrorCode(error.error().code())
+
+        case Left(assignments) =>
+          val topicReassignments = assignments.groupBy(_._1.topic()).map {
+            case (topic, reassignmentsByTp) =>
+              val partitionReassignments = reassignmentsByTp.map {
+                case (topicPartition, assignment) =>
+                  new ListPartitionReassignmentsResponseData.OngoingPartitionReassignment()
+                    .setPartitionIndex(topicPartition.partition())
+                    .setAddingReplicas(assignment.addingReplicas.toList.asJava.asInstanceOf[java.util.List[java.lang.Integer]])
+                    .setRemovingReplicas(assignment.removingReplicas.toList.asJava.asInstanceOf[java.util.List[java.lang.Integer]])
+                    .setReplicas(assignment.replicas.toList.asJava.asInstanceOf[java.util.List[java.lang.Integer]])
+              }.toList
+
+              new ListPartitionReassignmentsResponseData.OngoingTopicReassignment().setName(topic)
+                .setPartitions(partitionReassignments.asJava)
+          }.toList
+
+          new ListPartitionReassignmentsResponseData().setTopics(topicReassignments.asJava)
+      }
+
+      sendResponseMaybeThrottle(request, requestThrottleMs =>
+        new ListPartitionReassignmentsResponse(responseData.setThrottleTimeMs(requestThrottleMs))
+      )
+    }
+
+    val partitionsOpt = listPartitionReassignmentsRequest.data().topics() match {
+      case topics: Any =>
+        Some(topics.iterator().asScala.flatMap { topic =>
+          topic.partitionIndexes().iterator().asScala
+            .map { tp => new TopicPartition(topic.name(), tp) }
+        }.toSet)
+      case _ => None
+    }
+
+    controller.listPartitionReassignments(partitionsOpt, sendResponseCallback)
+  }
+
+  private def configsAuthorizationApiError(resource: ConfigResource): ApiError = {
     val error = resource.`type` match {
-      case ConfigResource.Type.BROKER => Errors.CLUSTER_AUTHORIZATION_FAILED
+      case ConfigResource.Type.BROKER | ConfigResource.Type.BROKER_LOGGER => Errors.CLUSTER_AUTHORIZATION_FAILED
       case ConfigResource.Type.TOPIC => Errors.TOPIC_AUTHORIZATION_FAILED
       case rt => throw new InvalidRequestException(s"Unexpected resource type $rt for resource ${resource.name}")
     }
     new ApiError(error, null)
   }
 
+  def handleIncrementalAlterConfigsRequest(request: RequestChannel.Request): Unit = {
+    val alterConfigsRequest = request.body[IncrementalAlterConfigsRequest]
+
+    val configs = alterConfigsRequest.data().resources().iterator().asScala.map { alterConfigResource =>
+      val configResource = new ConfigResource(ConfigResource.Type.forId(alterConfigResource.resourceType()), alterConfigResource.resourceName())
+      configResource -> alterConfigResource.configs().iterator().asScala.map {
+        alterConfig => new AlterConfigOp(new ConfigEntry(alterConfig.name(), alterConfig.value()), OpType.forId(alterConfig.configOperation())) }.toList
+    }.toMap
+
+    val (authorizedResources, unauthorizedResources) = configs.partition { case (resource, _) =>
+      resource.`type` match {
+        case ConfigResource.Type.BROKER | ConfigResource.Type.BROKER_LOGGER =>
+          authorize(request, ALTER_CONFIGS, CLUSTER, CLUSTER_NAME)
+        case ConfigResource.Type.TOPIC =>
+          authorize(request, ALTER_CONFIGS, TOPIC, resource.name)
+        case rt => throw new InvalidRequestException(s"Unexpected resource type $rt")
+      }
+    }
+
+    val authorizedResult = adminManager.incrementalAlterConfigs(authorizedResources, alterConfigsRequest.data().validateOnly())
+    val unauthorizedResult = unauthorizedResources.keys.map { resource =>
+      resource -> configsAuthorizationApiError(resource)
+    }
+    sendResponseMaybeThrottle(request, requestThrottleMs =>
+      new IncrementalAlterConfigsResponse(IncrementalAlterConfigsResponse.toResponseData(requestThrottleMs,
+        (authorizedResult ++ unauthorizedResult).asJava)))
+  }
+
   def handleDescribeConfigsRequest(request: RequestChannel.Request): Unit = {
     val describeConfigsRequest = request.body[DescribeConfigsRequest]
     val (authorizedResources, unauthorizedResources) = describeConfigsRequest.resources.asScala.partition { resource =>
       resource.`type` match {
-        case ConfigResource.Type.BROKER => authorize(request.session, DescribeConfigs, Resource.ClusterResource)
+        case ConfigResource.Type.BROKER | ConfigResource.Type.BROKER_LOGGER =>
+          authorize(request, DESCRIBE_CONFIGS, CLUSTER, CLUSTER_NAME)
         case ConfigResource.Type.TOPIC =>
-          authorize(request.session, DescribeConfigs, Resource(Topic, resource.name, LITERAL))
+          authorize(request, DESCRIBE_CONFIGS, TOPIC, resource.name)
         case rt => throw new InvalidRequestException(s"Unexpected resource type $rt for resource ${resource.name}")
       }
     }
@@ -2120,8 +2508,8 @@ class KafkaApis(val requestChannel: RequestChannel,
       resource -> Option(describeConfigsRequest.configNames(resource)).map(_.asScala.toSet)
     }.toMap, describeConfigsRequest.includeSynonyms)
     val unauthorizedConfigs = unauthorizedResources.map { resource =>
-      val error = configsAuthorizationApiError(request.session, resource)
-      resource -> new DescribeConfigsResponse.Config(error, Collections.emptyList[DescribeConfigsResponse.ConfigEntry])
+      val error = configsAuthorizationApiError(resource)
+      resource -> new DescribeConfigsResponse.Config(error, util.Collections.emptyList[DescribeConfigsResponse.ConfigEntry])
     }
 
     sendResponseMaybeThrottle(request, requestThrottleMs =>
@@ -2131,7 +2519,7 @@ class KafkaApis(val requestChannel: RequestChannel,
   def handleAlterReplicaLogDirsRequest(request: RequestChannel.Request): Unit = {
     val alterReplicaDirsRequest = request.body[AlterReplicaLogDirsRequest]
     val responseMap = {
-      if (authorize(request.session, Alter, Resource.ClusterResource))
+      if (authorize(request, ALTER, CLUSTER, CLUSTER_NAME))
         replicaManager.alterReplicaLogDirs(alterReplicaDirsRequest.partitionDirs.asScala)
       else
         alterReplicaDirsRequest.partitionDirs.asScala.keys.map((_, Errors.CLUSTER_AUTHORIZATION_FAILED)).toMap
@@ -2142,7 +2530,7 @@ class KafkaApis(val requestChannel: RequestChannel,
   def handleDescribeLogDirsRequest(request: RequestChannel.Request): Unit = {
     val describeLogDirsDirRequest = request.body[DescribeLogDirsRequest]
     val logDirInfos = {
-      if (authorize(request.session, Describe, Resource.ClusterResource)) {
+      if (authorize(request, DESCRIBE, CLUSTER, CLUSTER_NAME)) {
         val partitions =
           if (describeLogDirsDirRequest.isAllTopicPartitions)
             replicaManager.logManager.allLogs.map(_.topicPartition).toSet
@@ -2157,78 +2545,87 @@ class KafkaApis(val requestChannel: RequestChannel,
     sendResponseMaybeThrottle(request, throttleTimeMs => new DescribeLogDirsResponse(throttleTimeMs, logDirInfos.asJava))
   }
 
-  def handleCreateTokenRequest(request: RequestChannel.Request) {
+  def handleCreateTokenRequest(request: RequestChannel.Request): Unit = {
     val createTokenRequest = request.body[CreateDelegationTokenRequest]
 
     // the callback for sending a create token response
-    def sendResponseCallback(createResult: CreateTokenResult) {
+    def sendResponseCallback(createResult: CreateTokenResult): Unit = {
       trace("Sending create token response for correlation id %d to client %s."
         .format(request.header.correlationId, request.header.clientId))
       sendResponseMaybeThrottle(request, requestThrottleMs =>
-        new CreateDelegationTokenResponse(requestThrottleMs, createResult.error, request.session.principal, createResult.issueTimestamp,
+        CreateDelegationTokenResponse.prepareResponse(requestThrottleMs, createResult.error, request.context.principal, createResult.issueTimestamp,
           createResult.expiryTimestamp, createResult.maxTimestamp, createResult.tokenId, ByteBuffer.wrap(createResult.hmac)))
     }
 
     if (!allowTokenRequests(request))
       sendResponseMaybeThrottle(request, requestThrottleMs =>
-        new CreateDelegationTokenResponse(requestThrottleMs, Errors.DELEGATION_TOKEN_REQUEST_NOT_ALLOWED, request.session.principal))
+        CreateDelegationTokenResponse.prepareResponse(requestThrottleMs, Errors.DELEGATION_TOKEN_REQUEST_NOT_ALLOWED, request.context.principal))
     else {
-      val renewerList = createTokenRequest.renewers().asScala.toList
+      val renewerList = createTokenRequest.data.renewers.asScala.toList.map(entry =>
+        new KafkaPrincipal(entry.principalType, entry.principalName))
 
-      if (renewerList.exists(principal =>  principal.getPrincipalType != KafkaPrincipal.USER_TYPE)) {
+      if (renewerList.exists(principal => principal.getPrincipalType != KafkaPrincipal.USER_TYPE)) {
         sendResponseMaybeThrottle(request, requestThrottleMs =>
-          new CreateDelegationTokenResponse(requestThrottleMs, Errors.INVALID_PRINCIPAL_TYPE, request.session.principal))
+          CreateDelegationTokenResponse.prepareResponse(requestThrottleMs, Errors.INVALID_PRINCIPAL_TYPE, request.context.principal))
       }
       else {
         tokenManager.createToken(
-          request.session.principal,
-          createTokenRequest.renewers().asScala.toList,
-          createTokenRequest.maxLifeTime(),
+          request.context.principal,
+          renewerList,
+          createTokenRequest.data.maxLifetimeMs,
           sendResponseCallback
         )
       }
     }
   }
 
-  def handleRenewTokenRequest(request: RequestChannel.Request) {
+  def handleRenewTokenRequest(request: RequestChannel.Request): Unit = {
     val renewTokenRequest = request.body[RenewDelegationTokenRequest]
 
     // the callback for sending a renew token response
-    def sendResponseCallback(error: Errors, expiryTimestamp: Long) {
+    def sendResponseCallback(error: Errors, expiryTimestamp: Long): Unit = {
       trace("Sending renew token response %s for correlation id %d to client %s."
         .format(request.header.correlationId, request.header.clientId))
       sendResponseMaybeThrottle(request, requestThrottleMs =>
-        new RenewDelegationTokenResponse(requestThrottleMs, error, expiryTimestamp))
+        new RenewDelegationTokenResponse(
+             new RenewDelegationTokenResponseData()
+               .setThrottleTimeMs(requestThrottleMs)
+               .setErrorCode(error.code)
+               .setExpiryTimestampMs(expiryTimestamp)))
     }
 
     if (!allowTokenRequests(request))
       sendResponseCallback(Errors.DELEGATION_TOKEN_REQUEST_NOT_ALLOWED, DelegationTokenManager.ErrorTimestamp)
     else {
       tokenManager.renewToken(
-        request.session.principal,
-        renewTokenRequest.hmac,
-        renewTokenRequest.renewTimePeriod(),
+        request.context.principal,
+        ByteBuffer.wrap(renewTokenRequest.data.hmac),
+        renewTokenRequest.data.renewPeriodMs,
         sendResponseCallback
       )
     }
   }
 
-  def handleExpireTokenRequest(request: RequestChannel.Request) {
+  def handleExpireTokenRequest(request: RequestChannel.Request): Unit = {
     val expireTokenRequest = request.body[ExpireDelegationTokenRequest]
 
     // the callback for sending a expire token response
-    def sendResponseCallback(error: Errors, expiryTimestamp: Long) {
+    def sendResponseCallback(error: Errors, expiryTimestamp: Long): Unit = {
       trace("Sending expire token response for correlation id %d to client %s."
         .format(request.header.correlationId, request.header.clientId))
       sendResponseMaybeThrottle(request, requestThrottleMs =>
-        new ExpireDelegationTokenResponse(requestThrottleMs, error, expiryTimestamp))
+        new ExpireDelegationTokenResponse(
+            new ExpireDelegationTokenResponseData()
+              .setThrottleTimeMs(requestThrottleMs)
+              .setErrorCode(error.code)
+              .setExpiryTimestampMs(expiryTimestamp)))
     }
 
     if (!allowTokenRequests(request))
       sendResponseCallback(Errors.DELEGATION_TOKEN_REQUEST_NOT_ALLOWED, DelegationTokenManager.ErrorTimestamp)
     else {
       tokenManager.expireToken(
-        request.session.principal,
+        request.context.principal,
         expireTokenRequest.hmac(),
         expireTokenRequest.expiryTimePeriod(),
         sendResponseCallback
@@ -2236,11 +2633,11 @@ class KafkaApis(val requestChannel: RequestChannel,
     }
   }
 
-  def handleDescribeTokensRequest(request: RequestChannel.Request) {
+  def handleDescribeTokensRequest(request: RequestChannel.Request): Unit = {
     val describeTokenRequest = request.body[DescribeDelegationTokenRequest]
 
     // the callback for sending a describe token response
-    def sendResponseCallback(error: Errors, tokenDetails: List[DelegationToken]) {
+    def sendResponseCallback(error: Errors, tokenDetails: List[DelegationToken]): Unit = {
       sendResponseMaybeThrottle(request, requestThrottleMs =>
         new DescribeDelegationTokenResponse(requestThrottleMs, error, tokenDetails.asJava))
       trace("Sending describe token response for correlation id %d to client %s."
@@ -2252,14 +2649,17 @@ class KafkaApis(val requestChannel: RequestChannel,
     else if (!config.tokenAuthEnabled)
       sendResponseCallback(Errors.DELEGATION_TOKEN_AUTH_DISABLED, List.empty)
     else {
-      val requestPrincipal = request.session.principal
+      val requestPrincipal = request.context.principal
 
       if (describeTokenRequest.ownersListEmpty()) {
         sendResponseCallback(Errors.NONE, List())
       }
       else {
-        val owners = if (describeTokenRequest.owners == null) None else Some(describeTokenRequest.owners.asScala.toList)
-        def authorizeToken(tokenId: String) = authorize(request.session, Describe, Resource(kafka.security.auth.DelegationToken, tokenId, LITERAL))
+        val owners = if (describeTokenRequest.data.owners == null)
+          None
+        else
+          Some(describeTokenRequest.data.owners.asScala.map(p => new KafkaPrincipal(p.principalType(), p.principalName())).toList)
+        def authorizeToken(tokenId: String) = authorize(request, DESCRIBE, DELEGATION_TOKEN, tokenId)
         def eligible(token: TokenInformation) = DelegationTokenManager.filterToken(requestPrincipal, owners, token, authorizeToken)
         val tokens =  tokenManager.getTokens(eligible)
         sendResponseCallback(Errors.NONE, tokens)
@@ -2269,74 +2669,199 @@ class KafkaApis(val requestChannel: RequestChannel,
 
   def allowTokenRequests(request: RequestChannel.Request): Boolean = {
     val protocol = request.context.securityProtocol
-    if (request.session.principal.tokenAuthenticated ||
+    if (request.context.principal.tokenAuthenticated ||
       protocol == SecurityProtocol.PLAINTEXT ||
       // disallow requests from 1-way SSL
-      (protocol == SecurityProtocol.SSL && request.session.principal == KafkaPrincipal.ANONYMOUS))
+      (protocol == SecurityProtocol.SSL && request.context.principal == KafkaPrincipal.ANONYMOUS))
       false
     else
       true
   }
 
-  def handleElectPreferredReplicaLeader(request: RequestChannel.Request): Unit = {
+  def handleElectReplicaLeader(request: RequestChannel.Request): Unit = {
+
+    val electionRequest = request.body[ElectLeadersRequest]
+
+    def sendResponseCallback(
+      error: ApiError
+    )(
+      results: Map[TopicPartition, ApiError]
+    ): Unit = {
+      sendResponseMaybeThrottle(request, requestThrottleMs => {
+        val adjustedResults = if (electionRequest.data().topicPartitions() == null) {
+          /* When performing elections across all of the partitions we should only return
+           * partitions for which there was an eleciton or resulted in an error. In other
+           * words, partitions that didn't need election because they ready have the correct
+           * leader are not returned to the client.
+           */
+          results.filter { case (_, error) =>
+            error.error != Errors.ELECTION_NOT_NEEDED
+          }
+        } else results
+
+        val electionResults = new util.ArrayList[ReplicaElectionResult]()
+        adjustedResults
+          .groupBy { case (tp, _) => tp.topic }
+          .foreach { case (topic, ps) =>
+            val electionResult = new ReplicaElectionResult()
+
+            electionResult.setTopic(topic)
+            ps.foreach { case (topicPartition, error) =>
+              val partitionResult = new PartitionResult()
+              partitionResult.setPartitionId(topicPartition.partition)
+              partitionResult.setErrorCode(error.error.code)
+              partitionResult.setErrorMessage(error.message)
+              electionResult.partitionResult.add(partitionResult)
+            }
+
+            electionResults.add(electionResult)
+          }
+
+        new ElectLeadersResponse(
+          requestThrottleMs,
+          error.error.code,
+          electionResults,
+          electionRequest.version
+        )
+      })
+    }
+
+    if (!authorize(request, ALTER, CLUSTER, CLUSTER_NAME)) {
+      val error = new ApiError(Errors.CLUSTER_AUTHORIZATION_FAILED, null)
+      val partitionErrors: Map[TopicPartition, ApiError] =
+        electionRequest.topicPartitions.iterator.map(partition => partition -> error).toMap
 
-    val electionRequest = request.body[ElectPreferredLeadersRequest]
-    val partitions =
-      if (electionRequest.data().topicPartitions() == null) {
+      sendResponseCallback(error)(partitionErrors)
+    } else {
+      val partitions = if (electionRequest.data().topicPartitions() == null) {
         metadataCache.getAllPartitions()
       } else {
-        electionRequest.data().topicPartitions().asScala.flatMap{tp =>
-          tp.partitionId().asScala.map(partitionId => new TopicPartition(tp.topic, partitionId))}.toSet
+        electionRequest.topicPartitions
       }
-    def sendResponseCallback(result: Map[TopicPartition, ApiError]): Unit = {
-      sendResponseMaybeThrottle(request, requestThrottleMs => {
-        val results = result.
-          groupBy{case (tp, error) => tp.topic}.
-          map{case (topic, ps) => new ElectPreferredLeadersResponseData.ReplicaElectionResult()
-            .setTopic(topic)
-            .setPartitionResult(ps.map{
-            case (tp, error) =>
-              new ElectPreferredLeadersResponseData.PartitionResult()
-                .setErrorCode(error.error.code)
-                .setErrorMessage(error.message())
-                .setPartitionId(tp.partition)}.toList.asJava)}
-        val data = new ElectPreferredLeadersResponseData()
-          .setThrottleTimeMs(requestThrottleMs)
-          .setReplicaElectionResults(results.toList.asJava)
-        new ElectPreferredLeadersResponse(data)})
-    }
-    if (!authorize(request.session, Alter, Resource.ClusterResource)) {
-      val error = new ApiError(Errors.CLUSTER_AUTHORIZATION_FAILED, null);
-      val partitionErrors =
-      if (electionRequest.data().topicPartitions() == null) {
-        // Don't leak the set of partitions if the client lack authz
-        Map.empty[TopicPartition, ApiError]
-      } else {
-        partitions.map(partition => partition -> error).toMap
+
+      replicaManager.electLeaders(
+        controller,
+        partitions,
+        electionRequest.electionType,
+        sendResponseCallback(ApiError.NONE),
+        electionRequest.data().timeoutMs()
+      )
+    }
+  }
+
+  def handleOffsetDeleteRequest(request: RequestChannel.Request): Unit = {
+    val offsetDeleteRequest = request.body[OffsetDeleteRequest]
+    val groupId = offsetDeleteRequest.data.groupId
+
+    if (authorize(request, DELETE, GROUP, groupId)) {
+      val authorizedTopics = filterAuthorized(request, READ, TOPIC,
+        offsetDeleteRequest.data.topics.asScala.map(_.name).toSeq)
+
+      val topicPartitionErrors = mutable.Map[TopicPartition, Errors]()
+      val topicPartitions = mutable.ArrayBuffer[TopicPartition]()
+
+      for (topic <- offsetDeleteRequest.data.topics.asScala) {
+        for (partition <- topic.partitions.asScala) {
+          val tp = new TopicPartition(topic.name, partition.partitionIndex)
+          if (!authorizedTopics.contains(topic.name))
+            topicPartitionErrors(tp) = Errors.TOPIC_AUTHORIZATION_FAILED
+          else if (!metadataCache.contains(tp))
+            topicPartitionErrors(tp) = Errors.UNKNOWN_TOPIC_OR_PARTITION
+          else
+            topicPartitions += tp
+        }
       }
-      sendResponseCallback(partitionErrors)
+
+      val (groupError, authorizedTopicPartitionsErrors) = groupCoordinator.handleDeleteOffsets(
+        groupId, topicPartitions)
+
+      topicPartitionErrors ++= authorizedTopicPartitionsErrors
+
+      sendResponseMaybeThrottle(request, requestThrottleMs => {
+        if (groupError != Errors.NONE)
+          offsetDeleteRequest.getErrorResponse(requestThrottleMs, groupError)
+        else {
+          val topics = new OffsetDeleteResponseData.OffsetDeleteResponseTopicCollection
+          topicPartitionErrors.groupBy(_._1.topic).map { case (topic, topicPartitions) =>
+            val partitions = new OffsetDeleteResponseData.OffsetDeleteResponsePartitionCollection
+            topicPartitions.map { case (topicPartition, error) =>
+              partitions.add(
+                new OffsetDeleteResponseData.OffsetDeleteResponsePartition()
+                  .setPartitionIndex(topicPartition.partition)
+                  .setErrorCode(error.code)
+              )
+              topics.add(new OffsetDeleteResponseData.OffsetDeleteResponseTopic()
+                .setName(topic)
+                .setPartitions(partitions))
+            }
+          }
+
+          new OffsetDeleteResponse(new OffsetDeleteResponseData()
+            .setTopics(topics)
+            .setThrottleTimeMs(requestThrottleMs))
+        }
+      })
     } else {
-      replicaManager.electPreferredLeaders(controller, partitions, sendResponseCallback, electionRequest.data().timeoutMs())
+      sendResponseMaybeThrottle(request, requestThrottleMs =>
+        offsetDeleteRequest.getErrorResponse(requestThrottleMs, Errors.GROUP_AUTHORIZATION_FAILED))
     }
   }
 
-  def authorizeClusterAction(request: RequestChannel.Request): Unit = {
-    if (!isAuthorizedClusterAction(request))
-      throw new ClusterAuthorizationException(s"Request $request is not authorized.")
+  private def authorize(request: RequestChannel.Request,
+                        operation: AclOperation,
+                        resourceType: ResourceType,
+                        resourceName: String,
+                        logIfAllowed: Boolean = true,
+                        logIfDenied: Boolean = true,
+                        refCount: Int = 1): Boolean = {
+    authorizer.forall { authZ =>
+      val resource = new ResourcePattern(resourceType, resourceName, PatternType.LITERAL)
+      val actions = Collections.singletonList(new Action(operation, resource, refCount, logIfAllowed, logIfDenied))
+      authZ.authorize(request.context, actions).asScala.head == AuthorizationResult.ALLOWED
+    }
   }
 
-  private def isAuthorizedClusterAction(request: RequestChannel.Request): Boolean = {
-    authorize(request.session, ClusterAction, Resource.ClusterResource)
+  private def filterAuthorized(request: RequestChannel.Request,
+                               operation: AclOperation,
+                               resourceType: ResourceType,
+                               resourceNames: Seq[String],
+                               logIfAllowed: Boolean = true,
+                               logIfDenied: Boolean = true): Set[String] = {
+    authorizer match {
+      case Some(authZ) =>
+        val resources = resourceNames.groupBy(identity).mapValues(_.size).toList
+        val actions = resources.map { case (resourceName, count) =>
+          val resource = new ResourcePattern(resourceType, resourceName, PatternType.LITERAL)
+          new Action(operation, resource, count, logIfAllowed, logIfDenied)
+        }
+        authZ.authorize(request.context, actions.asJava).asScala
+          .zip(resources.map(_._1)) // zip with resource name
+          .filter(_._1 == AuthorizationResult.ALLOWED) // filter authorized resources
+          .map(_._2).toSet
+      case None =>
+        resourceNames.toSet
+    }
   }
 
-  def authorizeClusterAlter(request: RequestChannel.Request): Unit = {
-    if (!authorize(request.session, Alter, Resource.ClusterResource))
+  private def authorizeClusterOperation(request: RequestChannel.Request, operation: AclOperation): Unit = {
+    if (!authorize(request, operation, CLUSTER, CLUSTER_NAME))
       throw new ClusterAuthorizationException(s"Request $request is not authorized.")
   }
 
-  def authorizeClusterDescribe(request: RequestChannel.Request): Unit = {
-    if (!authorize(request.session, Describe, Resource.ClusterResource))
-      throw new ClusterAuthorizationException(s"Request $request is not authorized.")
+  private def authorizedOperations(request: RequestChannel.Request, resource: Resource): Int = {
+    val supportedOps = AuthorizerUtils.supportedOperations(resource.resourceType).toList
+    val authorizedOps = authorizer match {
+      case Some(authZ) =>
+        val resourcePattern = new ResourcePattern(resource.resourceType, resource.name, PatternType.LITERAL)
+        val actions = supportedOps.map { op => new Action(op, resourcePattern, 1, false, false) }
+        authZ.authorize(request.context, actions.asJava).asScala
+          .zip(supportedOps)
+          .filter(_._1 == AuthorizationResult.ALLOWED)
+          .map(_._2).toSet
+      case None =>
+        supportedOps.toSet
+    }
+    Utils.to32BitField(authorizedOps.map(operation => operation.code().asInstanceOf[JByte]).asJava)
   }
 
   private def updateRecordConversionStats(request: RequestChannel.Request,
@@ -2359,12 +2884,13 @@ class KafkaApis(val requestChannel: RequestChannel,
     request.temporaryMemoryBytes = conversionStats.temporaryMemoryBytes
   }
 
-  private def handleError(request: RequestChannel.Request, e: Throwable) {
+  private def handleError(request: RequestChannel.Request, e: Throwable): Unit = {
     val mayThrottle = e.isInstanceOf[ClusterAuthorizationException] || !request.header.apiKey.clusterAction
     error("Error when handling request: " +
       s"clientId=${request.header.clientId}, " +
       s"correlationId=${request.header.correlationId}, " +
       s"api=${request.header.apiKey}, " +
+      s"version=${request.header.apiVersion}, " +
       s"body=${request.body[AbstractRequest]}", e)
     if (mayThrottle)
       sendErrorResponseMaybeThrottle(request, e)
@@ -2382,7 +2908,7 @@ class KafkaApis(val requestChannel: RequestChannel,
     sendResponse(request, Some(createResponse(throttleTimeMs)), onComplete)
   }
 
-  private def sendErrorResponseMaybeThrottle(request: RequestChannel.Request, error: Throwable) {
+  private def sendErrorResponseMaybeThrottle(request: RequestChannel.Request, error: Throwable): Unit = {
     val throttleTimeMs = quotas.request.maybeRecordAndGetThrottleTimeMs(request)
     quotas.request.throttle(request, throttleTimeMs, sendResponse)
     sendErrorOrCloseConnection(request, error, throttleTimeMs)
@@ -2433,8 +2959,11 @@ class KafkaApis(val requestChannel: RequestChannel,
         val responseString =
           if (RequestChannel.isRequestLoggingEnabled) Some(response.toString(request.context.apiVersion))
           else None
+        observeRequestResponse(request, response)
+
         new RequestChannel.SendResponse(request, responseSend, responseString, onComplete)
       case None =>
+        observeRequestResponse(request, null)
         new RequestChannel.NoOpResponse(request)
     }
     sendResponse(response)
@@ -2444,16 +2973,23 @@ class KafkaApis(val requestChannel: RequestChannel,
     requestChannel.sendResponse(response)
   }
 
-  private def isBrokerEpochStale(brokerEpochInRequest: Long): Boolean = {
-    // Broker epoch in LeaderAndIsr/UpdateMetadata/StopReplica request is unknown
-    // if the controller hasn't been upgraded to use KIP-380
-    if (brokerEpochInRequest == AbstractControlRequest.UNKNOWN_BROKER_EPOCH) false
-    else {
+  private def isBrokerEpochStale(brokerEpochInRequest: Long, maxBrokerEpochInRequest: Long): Boolean = {
+    if (maxBrokerEpochInRequest != AbstractControlRequest.UNKNOWN_BROKER_EPOCH)  {
+      maxBrokerEpochInRequest < controller.brokerEpoch
+    }
+    else if (brokerEpochInRequest != AbstractControlRequest.UNKNOWN_BROKER_EPOCH) {
       val curBrokerEpoch = controller.brokerEpoch
       if (brokerEpochInRequest < curBrokerEpoch) true
       else if (brokerEpochInRequest == curBrokerEpoch) false
       else throw new IllegalStateException(s"Epoch $brokerEpochInRequest larger than current broker epoch $curBrokerEpoch")
-    }
+    } else false
   }
 
+  private def observeRequestResponse(request: RequestChannel.Request, response: AbstractResponse): Unit = {
+    try {
+      observer.observe(request.context, request.body[AbstractRequest], response)
+    } catch {
+      case e: Exception => error(s"Observer failed to observe ${Observer.describeRequestAndResponse(request, response)}", e)
+    }
+  }
 }
diff --git a/core/src/main/scala/kafka/server/KafkaConfig.scala b/core/src/main/scala/kafka/server/KafkaConfig.scala
index 02e171527e501..f82792474caf1 100755
--- a/core/src/main/scala/kafka/server/KafkaConfig.scala
+++ b/core/src/main/scala/kafka/server/KafkaConfig.scala
@@ -17,29 +17,33 @@
 
 package kafka.server
 
+import java.io.File
 import java.util
-import java.util.{Collections, Properties}
+import java.util.{Collections, Locale, Properties}
 
 import kafka.api.{ApiVersion, ApiVersionValidator, KAFKA_0_10_0_IV1, KAFKA_2_1_IV0}
 import kafka.cluster.EndPoint
 import kafka.coordinator.group.OffsetConfig
 import kafka.coordinator.transaction.{TransactionLog, TransactionStateManager}
 import kafka.message.{BrokerCompressionCodec, CompressionCodec, ZStdCompressionCodec}
+import kafka.security.authorizer.AuthorizerWrapper
 import kafka.utils.CoreUtils
 import kafka.utils.Implicits._
 import org.apache.kafka.clients.CommonClientConfigs
 import org.apache.kafka.common.Reconfigurable
+import org.apache.kafka.common.config.SecurityConfig
 import org.apache.kafka.common.config.ConfigDef.{ConfigKey, ValidList}
 import org.apache.kafka.common.config.internals.BrokerSecurityConfigs
-import org.apache.kafka.common.config.{AbstractConfig, ConfigDef, ConfigException, SaslConfigs, SslConfigs, TopicConfig}
+import org.apache.kafka.common.config.{AbstractConfig, ConfigDef, ConfigException, SaslConfigs, SslClientAuth, SslConfigs, TopicConfig}
 import org.apache.kafka.common.metrics.Sensor
 import org.apache.kafka.common.network.ListenerName
 import org.apache.kafka.common.record.{LegacyRecord, Records, TimestampType}
 import org.apache.kafka.common.security.auth.SecurityProtocol
 import org.apache.kafka.common.utils.Utils
+import org.apache.kafka.server.authorizer.Authorizer
 
 import scala.collection.JavaConverters._
-import scala.collection.Map
+import scala.collection.{Map, Seq}
 
 object Defaults {
   /** ********* Zookeeper Configuration ***********/
@@ -58,10 +62,17 @@ object Defaults {
   val BackgroundThreads = 10
   val QueuedMaxRequests = 500
   val QueuedMaxRequestBytes = -1
+  val ProducerBatchDecompressionEnable = true
+  val PreferredController = false
+  val AllowPreferredControllerFallback = true
 
   /************* Authorizer Configuration ***********/
   val AuthorizerClassName = ""
 
+  /** ********* Broker-side configuration ***********/
+  val ObserverClassName = "kafka.server.NoOpObserver"
+  val ObserverShutdownTimeoutMs = 2000
+
   /** ********* Socket Server Configuration ***********/
   val Port = 9092
   val HostName: String = new String("")
@@ -73,12 +84,19 @@ object Defaults {
   val SocketSendBufferBytes: Int = 100 * 1024
   val SocketReceiveBufferBytes: Int = 100 * 1024
   val SocketRequestMaxBytes: Int = 100 * 1024 * 1024
+  val SocketRequestCommonBytes: Int = -1
+  val SocketRequestBufferCacheSize: Int = 0
+  val RequestMaxLocalTimeMs = Long.MaxValue
   val MaxConnectionsPerIp: Int = Int.MaxValue
   val MaxConnectionsPerIpOverrides: String = ""
+  val MaxConnections: Int = Int.MaxValue
   val ConnectionsMaxIdleMs = 10 * 60 * 1000L
   val RequestTimeoutMs = 30000
   val FailedAuthenticationDelayMs = 100
 
+  val HeapDumpFolder = "."
+  val HeapDumpTimeout = 30000
+
   /** ********* Log Configuration ***********/
   val NumPartitions = 1
   val LogDir = "/tmp/kafka-logs"
@@ -102,6 +120,7 @@ object Defaults {
   val LogCleanerEnable = true
   val LogCleanerDeleteRetentionMs = 24 * 60 * 60 * 1000L
   val LogCleanerMinCompactionLagMs = 0L
+  val LogCleanerMaxCompactionLagMs = Long.MaxValue
   val LogIndexSizeMaxBytes = 10 * 1024 * 1024
   val LogIndexIntervalBytes = 4096
   val LogFlushIntervalMessages = Long.MaxValue
@@ -150,7 +169,7 @@ object Defaults {
 
   /** ********* Group coordinator configuration ***********/
   val GroupMinSessionTimeoutMs = 6000
-  val GroupMaxSessionTimeoutMs = 300000
+  val GroupMaxSessionTimeoutMs = 1800000
   val GroupInitialRebalanceDelayMs = 3000
   val GroupMaxSize: Int = Int.MaxValue
 
@@ -165,6 +184,9 @@ object Defaults {
   val OffsetsRetentionCheckIntervalMs: Long = OffsetConfig.DefaultOffsetsRetentionCheckIntervalMs
   val OffsetCommitTimeoutMs = OffsetConfig.DefaultOffsetCommitTimeoutMs
   val OffsetCommitRequiredAcks = OffsetConfig.DefaultOffsetCommitRequiredAcks
+  val OffsetsTopicMaxMessageBytes = OffsetConfig.DefaultOffsetsTopicMaxMessageBytes
+  val OffsetsTopicMinInSyncReplicas = OffsetConfig.DefaultOffsetsTopicMinInSyncReplicas
+  val OffsetsTopicMinCompactionLagMs = OffsetConfig.DefaultOffsetsTopicMinCompactionLagMs
 
   /** ********* Transaction management configuration ***********/
   val TransactionalIdExpirationMs = TransactionStateManager.DefaultTransactionalIdExpirationMs
@@ -203,6 +225,7 @@ object Defaults {
   val MetricSampleWindowMs = 30000
   val MetricReporterClasses = ""
   val MetricRecordingLevel = Sensor.RecordingLevel.INFO.toString()
+  val MetricReplaceOnDuplicate = false
 
 
   /** ********* Kafka Yammer Metrics Reporter Configuration ***********/
@@ -211,16 +234,15 @@ object Defaults {
 
   /** ********* SSL configuration ***********/
   val SslProtocol = SslConfigs.DEFAULT_SSL_PROTOCOL
+  val SslContextProviderClass = SslConfigs.DEFAULT_SSL_CONTEXT_PROVIDER_CLASS
   val SslEnabledProtocols = SslConfigs.DEFAULT_SSL_ENABLED_PROTOCOLS
   val SslKeystoreType = SslConfigs.DEFAULT_SSL_KEYSTORE_TYPE
   val SslTruststoreType = SslConfigs.DEFAULT_SSL_TRUSTSTORE_TYPE
   val SslKeyManagerAlgorithm = SslConfigs.DEFAULT_SSL_KEYMANGER_ALGORITHM
   val SslTrustManagerAlgorithm = SslConfigs.DEFAULT_SSL_TRUSTMANAGER_ALGORITHM
   val SslEndpointIdentificationAlgorithm = SslConfigs.DEFAULT_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM
-  val SslClientAuthRequired = "required"
-  val SslClientAuthRequested = "requested"
-  val SslClientAuthNone = "none"
-  val SslClientAuth = SslClientAuthNone
+  val SslClientAuthentication = SslClientAuth.NONE.name().toLowerCase(Locale.ROOT)
+  val SslClientAuthenticationValidValues = SslClientAuth.VALUES.asScala.map(v => v.toString().toLowerCase(Locale.ROOT)).asJava.toArray(new Array[String](0))
   val SslPrincipalMappingRules = BrokerSecurityConfigs.DEFAULT_SSL_PRINCIPAL_MAPPING_RULES
 
     /** ********* General Security configuration ***********/
@@ -254,8 +276,8 @@ object KafkaConfig {
 
   private val LogConfigPrefix = "log."
 
-  def main(args: Array[String]) {
-    System.out.println(configDef.toHtmlTable(DynamicBrokerConfig.dynamicConfigUpdateModes))
+  def main(args: Array[String]): Unit = {
+    System.out.println(configDef.toHtml(DynamicBrokerConfig.dynamicConfigUpdateModes))
   }
 
   /** ********* Zookeeper Configuration ***********/
@@ -277,8 +299,19 @@ object KafkaConfig {
   val QueuedMaxRequestsProp = "queued.max.requests"
   val QueuedMaxBytesProp = "queued.max.request.bytes"
   val RequestTimeoutMsProp = CommonClientConfigs.REQUEST_TIMEOUT_MS_CONFIG
+  val HeapDumpFolderProp = "heap.dump.folder"
+  val HeapDumpTimeoutProp = "heap.dump.timeout"
+  val ProducerBatchDecompressionEnableProp = "producer.batch.decompression.enable"
+  val PreferredControllerProp = "preferred.controller"
+  val AllowPreferredControllerFallbackProp = "allow.preferred.controller.fallback"
+
   /************* Authorizer Configuration ***********/
   val AuthorizerClassNameProp = "authorizer.class.name"
+
+  /** ********* Broker-side observer Configuration ****************/
+  val ObserverClassNameProp = "observer.class.name"
+  val ObserverShutdownTimeoutMsProp = "observer.shutdown.timeout"
+
   /** ********* Socket Server Configuration ***********/
   val PortProp = "port"
   val HostNameProp = "host.name"
@@ -290,9 +323,13 @@ object KafkaConfig {
   val ControlPlaneListenerNameProp = "control.plane.listener.name"
   val SocketSendBufferBytesProp = "socket.send.buffer.bytes"
   val SocketReceiveBufferBytesProp = "socket.receive.buffer.bytes"
+  val RequestMaxLocalTimeMsProp = "request.max.local.time.ms"
   val SocketRequestMaxBytesProp = "socket.request.max.bytes"
+  val SocketRequestCommonBytesProp = "socket.request.common.bytes"
+  val SocketRequestBufferCacheSizeProp = "socket.request.buffer.cache.size"
   val MaxConnectionsPerIpProp = "max.connections.per.ip"
   val MaxConnectionsPerIpOverridesProp = "max.connections.per.ip.overrides"
+  val MaxConnectionsProp = "max.connections"
   val ConnectionsMaxIdleMsProp = "connections.max.idle.ms"
   val FailedAuthenticationDelayMsProp = "connection.failed.authentication.delay.ms"
   /***************** rack configuration *************/
@@ -326,6 +363,7 @@ object KafkaConfig {
   val LogCleanerEnableProp = "log.cleaner.enable"
   val LogCleanerDeleteRetentionMsProp = "log.cleaner.delete.retention.ms"
   val LogCleanerMinCompactionLagMsProp = "log.cleaner.min.compaction.lag.ms"
+  val LogCleanerMaxCompactionLagMsProp = "log.cleaner.max.compaction.lag.ms"
   val LogIndexSizeMaxBytesProp = "log.index.size.max.bytes"
   val LogIndexIntervalBytesProp = "log.index.interval.bytes"
   val LogFlushIntervalMessagesProp = "log.flush.interval.messages"
@@ -368,6 +406,7 @@ object KafkaConfig {
   val InterBrokerSecurityProtocolProp = "security.inter.broker.protocol"
   val InterBrokerProtocolVersionProp = "inter.broker.protocol.version"
   val InterBrokerListenerNameProp = "inter.broker.listener.name"
+  val ReplicaSelectorClassProp = "replica.selector.class"
   /** ********* Controlled shutdown configuration ***********/
   val ControlledShutdownMaxRetriesProp = "controlled.shutdown.max.retries"
   val ControlledShutdownRetryBackoffMsProp = "controlled.shutdown.retry.backoff.ms"
@@ -388,6 +427,9 @@ object KafkaConfig {
   val OffsetsRetentionCheckIntervalMsProp = "offsets.retention.check.interval.ms"
   val OffsetCommitTimeoutMsProp = "offsets.commit.timeout.ms"
   val OffsetCommitRequiredAcksProp = "offsets.commit.required.acks"
+  val OffsetsTopicMaxMessageBytesProp = "offsets.topic.max.message.bytes"
+  val OffsetsTopicMinInSyncReplicasProp = "offsets.topic.min.insync.replicas"
+  val OffsetsTopicMinCompactionLagMsProp = "offsets.topic.min.compaction.lag.ms"
   /** ********* Transaction management configuration ***********/
   val TransactionalIdExpirationMsProp = "transactional.id.expiration.ms"
   val TransactionsMaxTimeoutMsProp = "transaction.max.timeout.ms"
@@ -421,6 +463,7 @@ object KafkaConfig {
   val MetricNumSamplesProp: String = CommonClientConfigs.METRICS_NUM_SAMPLES_CONFIG
   val MetricReporterClassesProp: String = CommonClientConfigs.METRIC_REPORTER_CLASSES_CONFIG
   val MetricRecordingLevelProp: String = CommonClientConfigs.METRICS_RECORDING_LEVEL_CONFIG
+  val MetricReplaceOnDuplicateProp: String = CommonClientConfigs.METRICS_REPLACE_ON_DUPLICATE_CONFIG
 
   /** ********* Kafka Yammer Metrics Reporters Configuration ***********/
   val KafkaMetricsReporterClassesProp = "kafka.metrics.reporters"
@@ -429,10 +472,12 @@ object KafkaConfig {
   /** ******** Common Security Configuration *************/
   val PrincipalBuilderClassProp = BrokerSecurityConfigs.PRINCIPAL_BUILDER_CLASS_CONFIG
   val ConnectionsMaxReauthMsProp = BrokerSecurityConfigs.CONNECTIONS_MAX_REAUTH_MS
+  val securityProviderClassProp = SecurityConfig.SECURITY_PROVIDERS_CONFIG
 
   /** ********* SSL Configuration ****************/
   val SslProtocolProp = SslConfigs.SSL_PROTOCOL_CONFIG
   val SslProviderProp = SslConfigs.SSL_PROVIDER_CONFIG
+  val SslContextProviderClassProp = SslConfigs.SSL_CONTEXT_PROVIDER_CLASS_CONFIG
   val SslCipherSuitesProp = SslConfigs.SSL_CIPHER_SUITES_CONFIG
   val SslEnabledProtocolsProp = SslConfigs.SSL_ENABLED_PROTOCOLS_CONFIG
   val SslKeystoreTypeProp = SslConfigs.SSL_KEYSTORE_TYPE_CONFIG
@@ -501,7 +546,7 @@ object KafkaConfig {
   "To avoid conflicts between zookeeper generated broker id's and user configured broker id's, generated broker ids " +
   "start from " + MaxReservedBrokerIdProp + " + 1."
   val MessageMaxBytesDoc = TopicConfig.MAX_MESSAGE_BYTES_DOC +
-    s"
      This can be set per topic with the topic level ${TopicConfig.MAX_MESSAGE_BYTES_CONFIG} config.
      "
+    s"This can be set per topic with the topic level ${TopicConfig.MAX_MESSAGE_BYTES_CONFIG} config."
   val NumNetworkThreadsDoc = "The number of threads that the server uses for receiving requests from the network and sending responses to the network"
   val NumIoThreadsDoc = "The number of threads that the server uses for processing requests, which may include disk I/O"
   val NumReplicaAlterLogDirsThreadsDoc = "The number of threads that can move replicas between log directories, which may include disk I/O"
@@ -509,8 +554,19 @@ object KafkaConfig {
   val QueuedMaxRequestsDoc = "The number of queued requests allowed for data-plane, before blocking the network threads"
   val QueuedMaxRequestBytesDoc = "The number of queued bytes allowed before no more requests are read"
   val RequestTimeoutMsDoc = CommonClientConfigs.REQUEST_TIMEOUT_MS_DOC
+  val HeapDumpFolderDoc = "The Folder under which heap dumps will be written by the watchdog"
+  val HeapDumpTimeoutDoc = "The max amount of time (in millis) to wait for heap dump to complete before halting regardless"
+  val ProducerBatchDecompressionEnableDoc = "Decompress batch sent by producer to perform verification of individual records inside the batch"
+  val PreferredControllerDoc = "Specifies whether the broker is a dedicated controller node. If set to true, the broker is a preferred controller node."
+  // Although AllowPreferredControllerFallback is expected to be configured dynamically at per cluster level, providing a static configuration entry
+  // here allows its value to be obtained without holding the dynamic broker configuration lock.
+  val AllowPreferredControllerFallbackDoc = "Specifies whether a non-preferred controller node (broker) is allowed to become the controller." +
+  " This configuration is expected to be configured at cluster level via dynamic broker configuration to provide a consistent configuration among all brokers." +
+  " If AllowPreferredControllerFallback is dynamically set to false and there is no preferred controllers, the non-preferred active controller does not resign."
   /************* Authorizer Configuration ***********/
-  val AuthorizerClassNameDoc = "The authorizer class that should be used for authorization"
+  val AuthorizerClassNameDoc = s"The fully qualified name of a class that implements s${classOf[Authorizer].getName}" +
+  " interface, which is used by the broker for authorization. This config also supports authorizers that implement the deprecated" +
+  " kafka.security.auth.Authorizer trait which was previously used for authorization."
   /** ********* Socket Server Configuration ***********/
   val PortDoc = "DEPRECATED: only used when listeners is not set. " +
   "Use listeners instead. \n" +
@@ -568,10 +624,22 @@ object KafkaConfig {
 
   val SocketSendBufferBytesDoc = "The SO_SNDBUF buffer of the socket server sockets. If the value is -1, the OS default will be used."
   val SocketReceiveBufferBytesDoc = "The SO_RCVBUF buffer of the socket server sockets. If the value is -1, the OS default will be used."
+  val RequestMaxLocalTimeMsDoc = "The maximum allowable request local processing time. If a request's local processing " +
+    "takes longer than this time the broker will kill itself as violating this timeout is a symptom of a more serious broker zombie state." +
+    " It is useful to observe the RequestDequeuePollIntervalMs metric to find a suitable setting for this configuration."
   val SocketRequestMaxBytesDoc = "The maximum number of bytes in a socket request"
+  val SocketRequestCommonBytesDoc = "The common size in bytes of a socket request"
+  val SocketRequestBufferCacheSizeDoc = "The maximal number of cache slot recycling memory pool will keep"
   val MaxConnectionsPerIpDoc = "The maximum number of connections we allow from each ip address. This can be set to 0 if there are overrides " +
-    "configured using " + MaxConnectionsPerIpOverridesProp + " property"
-  val MaxConnectionsPerIpOverridesDoc = "A comma-separated list of per-ip or hostname overrides to the default maximum number of connections. An example value is \"hostName:100,127.0.0.1:200\""
+    s"configured using $MaxConnectionsPerIpOverridesProp property. New connections from the ip address are dropped if the limit is reached."
+  val MaxConnectionsPerIpOverridesDoc = "A comma-separated list of per-ip or hostname overrides to the default maximum number of connections. " +
+    "An example value is \"hostName:100,127.0.0.1:200\""
+  val MaxConnectionsDoc = "The maximum number of connections we allow in the broker at any time. This limit is applied in addition " +
+    s"to any per-ip limits configured using $MaxConnectionsPerIpProp. Listener-level limits may also be configured by prefixing the " +
+    s"config name with the listener prefix, for example, listener.name.internal.$MaxConnectionsProp. Broker-wide limit " +
+    "should be configured based on broker capacity while listener limits should be configured based on application requirements. " +
+    "New connections are blocked if either the listener or broker limit is reached. Connections on the inter-broker listener are " +
+    "permitted even if broker-wide limit is reached. The least recently used connection on another listener will be closed in this case."
   val ConnectionsMaxIdleMsDoc = "Idle connections timeout: the server socket processor threads close the connections that idle more than this"
   val FailedAuthenticationDelayMsDoc = "Connection close delay on failed authentication: this is the time (in milliseconds) by which connection close will be delayed on authentication failure. " +
     s"This must be configured to be less than $ConnectionsMaxIdleMsProp to prevent connection timeout."
@@ -588,7 +656,7 @@ object KafkaConfig {
   val LogRollTimeJitterMillisDoc = "The maximum jitter to subtract from logRollTimeMillis (in milliseconds). If not set, the value in " + LogRollTimeJitterHoursProp + " is used"
   val LogRollTimeJitterHoursDoc = "The maximum jitter to subtract from logRollTimeMillis (in hours), secondary to " + LogRollTimeJitterMillisProp + " property"
 
-  val LogRetentionTimeMillisDoc = "The number of milliseconds to keep a log file before deleting it (in milliseconds), If not set, the value in " + LogRetentionTimeMinutesProp + " is used"
+  val LogRetentionTimeMillisDoc = "The number of milliseconds to keep a log file before deleting it (in milliseconds), If not set, the value in " + LogRetentionTimeMinutesProp + " is used. If set to -1, no time limit is applied."
   val LogRetentionTimeMinsDoc = "The number of minutes to keep a log file before deleting it (in minutes), secondary to " + LogRetentionTimeMillisProp + " property. If not set, the value in " + LogRetentionTimeHoursProp + " is used"
   val LogRetentionTimeHoursDoc = "The number of hours to keep a log file before deleting it (in hours), tertiary to " + LogRetentionTimeMillisProp + " property"
 
@@ -602,10 +670,16 @@ object KafkaConfig {
   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"
+  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 = "How long are delete records retained?"
   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 "
@@ -672,7 +746,7 @@ object KafkaConfig {
   val FetchPurgatoryPurgeIntervalRequestsDoc = "The purge interval (in number of requests) of the fetch request purgatory"
   val ProducerPurgatoryPurgeIntervalRequestsDoc = "The purge interval (in number of requests) of the producer request purgatory"
   val DeleteRecordsPurgatoryPurgeIntervalRequestsDoc = "The purge interval (in number of requests) of the delete records request purgatory"
-  val AutoLeaderRebalanceEnableDoc = "Enables auto leader balancing. A background thread checks and triggers leader balance if required at regular intervals"
+  val AutoLeaderRebalanceEnableDoc = "Enables auto leader balancing. A background thread checks the distribution of partition leaders at regular intervals, configurable by `leader.imbalance.check.interval.seconds`. If the leader imbalance exceeds `leader.imbalance.per.broker.percentage`, leader rebalance to the preferred leader for partitions is triggered."
   val LeaderImbalancePerBrokerPercentageDoc = "The ratio of leader imbalance allowed per broker. The controller would trigger a leader balance if it goes above this value per broker. The value is specified in percentage."
   val LeaderImbalanceCheckIntervalSecondsDoc = "The frequency with which the partition rebalance check is triggered by the controller"
   val UncleanLeaderElectionEnableDoc = "Indicates whether to enable replicas not in the ISR set to be elected as leader as a last resort, even though doing so may result in data loss"
@@ -684,6 +758,7 @@ object KafkaConfig {
   " Example of some valid values are: 0.8.0, 0.8.1, 0.8.1.1, 0.8.2, 0.8.2.0, 0.8.2.1, 0.9.0.0, 0.9.0.1 Check ApiVersion for the full list."
   val InterBrokerListenerNameDoc = s"Name of listener used for communication between brokers. If this is unset, the listener name is defined by $InterBrokerSecurityProtocolProp. " +
     s"It is an error to set this and $InterBrokerSecurityProtocolProp properties at the same time."
+  val ReplicaSelectorClassDoc = "The fully qualified class name that implements ReplicaSelector. This is used by the broker to find the preferred read replica. By default, we use an implementation that returns the leader."
   /** ********* Controlled shutdown configuration ***********/
   val ControlledShutdownMaxRetriesDoc = "Controlled shutdown can fail for multiple reasons. This determines the number of retries when such failure happens"
   val ControlledShutdownRetryBackoffMsDoc = "Before each retry, the system needs time to recover from the state that caused the previous failure (Controller fail over, replica lag etc). This config determines the amount of time to wait before retrying."
@@ -707,8 +782,13 @@ object KafkaConfig {
   val OffsetCommitTimeoutMsDoc = "Offset commit will be delayed until all replicas for the offsets topic receive the commit " +
   "or this timeout is reached. This is similar to the producer request timeout."
   val OffsetCommitRequiredAcksDoc = "The required acks before the commit can be accepted. In general, the default (-1) should not be overridden"
+  val OffsetsTopicMaxMessageBytesDoc = "Overriden " + MessageMaxBytesProp + " config for the consumer_offset topic."
+  val OffsetsTopicMinInSyncReplicasDoc = "Overridden " + MinInSyncReplicasProp + " config for the consumer_offset topic."
+  val OffsetsTopicMinCompactionLagMsDoc = "Overridden " + LogCleanerMinCompactionLagMsProp + " config for the consumer_offset topic."
   /** ********* Transaction management configuration ***********/
-  val TransactionalIdExpirationMsDoc = "The maximum amount of time in ms that the transaction coordinator will wait before proactively expire a producer's transactional id without receiving any transaction status updates from it."
+  val TransactionalIdExpirationMsDoc = "The time in ms that the transaction coordinator will wait without receiving any transaction status updates " +
+    "for the current transaction before expiring its transactional id. This setting also influences producer id expiration - producer ids are expired " + 
+    "once this time has elapsed after the last write with the given producer id. Note that producer ids may expire sooner if the last write from the producer id is deleted due to the topic's retention settings."
   val TransactionsMaxTimeoutMsDoc = "The maximum allowed timeout for transactions. " +
     "If a client’s requested transaction time exceed this, then the broker will return an error in InitProducerIdRequest. This prevents a client from too large of a timeout, which can stall consumers reading from topics included in the transaction."
   val TransactionsTopicMinISRDoc = "Overridden " + MinInSyncReplicasProp + " config for the transaction topic."
@@ -718,7 +798,7 @@ object KafkaConfig {
   val TransactionsTopicPartitionsDoc = "The number of partitions for the transaction topic (should not change after deployment)."
   val TransactionsTopicSegmentBytesDoc = "The transaction topic segment bytes should be kept relatively small in order to facilitate faster log compaction and cache loads"
   val TransactionsAbortTimedOutTransactionsIntervalMsDoc = "The interval at which to rollback transactions that have timed out"
-  val TransactionsRemoveExpiredTransactionsIntervalMsDoc = "The interval at which to remove transactions that have expired due to transactional.id.expiration.ms passing"
+  val TransactionsRemoveExpiredTransactionsIntervalMsDoc = "The interval at which to remove transactions that have expired due to transactional.id.expiration.ms passing"
 
   /** ********* Fetch Session Configuration **************/
   val MaxIncrementalFetchSessionCacheSlotsDoc = "The maximum number of incremental fetch sessions that we will maintain."
@@ -738,11 +818,6 @@ object KafkaConfig {
     "which is used to determine quota limits applied to client requests. By default, ,  or  " +
     "quotas stored in ZooKeeper are applied. For any given request, the most specific quota that matches the user principal " +
     "of the session and the client-id of the request is applied."
-  /** ********* Transaction Configuration ***********/
-  val TransactionIdExpirationMsDoc = "The maximum time of inactivity before a transactional id is expired by the " +
-    "transaction coordinator. Note that this also influences producer id expiration: Producer ids are guaranteed to expire " +
-    "after expiration of this timeout from the last write by the producer id (they may expire sooner if the last write " +
-    "from the producer id is deleted due to the topic's retention settings)."
 
   val DeleteTopicEnableDoc = "Enables delete topic. Delete topic through the admin tool will have no effect if this config is turned off"
   val CompressionTypeDoc = "Specify the final compression type for a given topic. This configuration accepts the standard compression codecs " +
@@ -754,6 +829,7 @@ object KafkaConfig {
   val MetricNumSamplesDoc = CommonClientConfigs.METRICS_NUM_SAMPLES_DOC
   val MetricReporterClassesDoc = CommonClientConfigs.METRIC_REPORTER_CLASSES_DOC
   val MetricRecordingLevelDoc = CommonClientConfigs.METRICS_RECORDING_LEVEL_DOC
+  val MetricReplaceOnDuplicateDoc = CommonClientConfigs.METRICS_REPLACE_ON_DUPLICATE_DOC
 
 
   /** ********* Kafka Yammer Metrics Reporter Configuration ***********/
@@ -769,10 +845,12 @@ object KafkaConfig {
   /** ******** Common Security Configuration *************/
   val PrincipalBuilderClassDoc = BrokerSecurityConfigs.PRINCIPAL_BUILDER_CLASS_DOC
   val ConnectionsMaxReauthMsDoc = BrokerSecurityConfigs.CONNECTIONS_MAX_REAUTH_MS_DOC
+  val securityProviderClassDoc = SecurityConfig.SECURITY_PROVIDERS_DOC
 
   /** ********* SSL Configuration ****************/
   val SslProtocolDoc = SslConfigs.SSL_PROTOCOL_DOC
   val SslProviderDoc = SslConfigs.SSL_PROVIDER_DOC
+  val SslContextProviderClassDoc = SslConfigs.SSL_CONTEXT_PROVIDER_CLASS_DOC
   val SslCipherSuitesDoc = SslConfigs.SSL_CIPHER_SUITES_DOC
   val SslEnabledProtocolsDoc = SslConfigs.SSL_ENABLED_PROTOCOLS_DOC
   val SslKeystoreTypeDoc = SslConfigs.SSL_KEYSTORE_TYPE_DOC
@@ -826,6 +904,12 @@ object KafkaConfig {
   val PasswordEncoderKeyLengthDoc =  "The key length used for encoding dynamically configured passwords."
   val PasswordEncoderIterationsDoc =  "The iteration count used for encoding dynamically configured passwords."
 
+  /** *********  Broker-side Observer Configuration *********/
+  val ObserverClassNameDoc = "The name of the observer class that is used to observe requests and/or response on broker."
+  val ObserverShutdownTimeoutMsDoc = "The maximum time of closing/shutting down an observer. This property can not be less than or equal to " +
+    "zero. When closing/shutting down an observer, most time is spent on flushing the observed stats. The reasonable timeout should be close to " +
+    "the time it takes to flush the stats."
+
   private val configDef = {
     import ConfigDef.Importance._
     import ConfigDef.Range._
@@ -854,10 +938,19 @@ object KafkaConfig {
       .define(QueuedMaxRequestsProp, INT, Defaults.QueuedMaxRequests, atLeast(1), HIGH, QueuedMaxRequestsDoc)
       .define(QueuedMaxBytesProp, LONG, Defaults.QueuedMaxRequestBytes, MEDIUM, QueuedMaxRequestBytesDoc)
       .define(RequestTimeoutMsProp, INT, Defaults.RequestTimeoutMs, HIGH, RequestTimeoutMsDoc)
+      .define(HeapDumpFolderProp, STRING, Defaults.HeapDumpFolder, LOW, HeapDumpFolderDoc)
+      .define(HeapDumpTimeoutProp, LONG, Defaults.HeapDumpTimeout, LOW, HeapDumpTimeoutDoc)
+      .define(ProducerBatchDecompressionEnableProp, BOOLEAN, Defaults.ProducerBatchDecompressionEnable, LOW, ProducerBatchDecompressionEnableDoc)
+      .define(PreferredControllerProp, BOOLEAN, Defaults.PreferredController, HIGH, PreferredControllerDoc)
+      .define(AllowPreferredControllerFallbackProp, BOOLEAN, Defaults.AllowPreferredControllerFallback, HIGH, AllowPreferredControllerFallbackDoc)
 
       /************* Authorizer Configuration ***********/
       .define(AuthorizerClassNameProp, STRING, Defaults.AuthorizerClassName, LOW, AuthorizerClassNameDoc)
 
+      /************* Broker-side Observer Configuration ***********/
+      .define(ObserverClassNameProp, STRING, Defaults.ObserverClassName, MEDIUM, ObserverClassNameDoc)
+      .define(ObserverShutdownTimeoutMsProp, LONG, Defaults.ObserverShutdownTimeoutMs, atLeast(1), MEDIUM, ObserverShutdownTimeoutMsDoc)
+
       /** ********* Socket Server Configuration ***********/
       .define(PortProp, INT, Defaults.Port, HIGH, PortDoc)
       .define(HostNameProp, STRING, Defaults.HostName, HIGH, HostNameDoc)
@@ -868,10 +961,14 @@ object KafkaConfig {
       .define(ListenerSecurityProtocolMapProp, STRING, Defaults.ListenerSecurityProtocolMap, LOW, ListenerSecurityProtocolMapDoc)
       .define(ControlPlaneListenerNameProp, STRING, null, HIGH, controlPlaneListenerNameDoc)
       .define(SocketSendBufferBytesProp, INT, Defaults.SocketSendBufferBytes, HIGH, SocketSendBufferBytesDoc)
+      .define(RequestMaxLocalTimeMsProp, LONG, Defaults.RequestMaxLocalTimeMs, atLeast(1), MEDIUM, RequestMaxLocalTimeMsDoc)
       .define(SocketReceiveBufferBytesProp, INT, Defaults.SocketReceiveBufferBytes, HIGH, SocketReceiveBufferBytesDoc)
       .define(SocketRequestMaxBytesProp, INT, Defaults.SocketRequestMaxBytes, atLeast(1), HIGH, SocketRequestMaxBytesDoc)
+      .define(SocketRequestCommonBytesProp, INT, Defaults.SocketRequestCommonBytes, MEDIUM, SocketRequestCommonBytesDoc)
+      .define(SocketRequestBufferCacheSizeProp, INT, Defaults.SocketRequestBufferCacheSize, atLeast(0), MEDIUM, SocketRequestBufferCacheSizeDoc)
       .define(MaxConnectionsPerIpProp, INT, Defaults.MaxConnectionsPerIp, atLeast(0), MEDIUM, MaxConnectionsPerIpDoc)
       .define(MaxConnectionsPerIpOverridesProp, STRING, Defaults.MaxConnectionsPerIpOverrides, MEDIUM, MaxConnectionsPerIpOverridesDoc)
+      .define(MaxConnectionsProp, INT, Defaults.MaxConnections, atLeast(0), MEDIUM, MaxConnectionsDoc)
       .define(ConnectionsMaxIdleMsProp, LONG, Defaults.ConnectionsMaxIdleMs, MEDIUM, ConnectionsMaxIdleMsDoc)
       .define(FailedAuthenticationDelayMsProp, INT, Defaults.FailedAuthenticationDelayMs, atLeast(0), LOW, FailedAuthenticationDelayMsDoc)
 
@@ -907,6 +1004,7 @@ object KafkaConfig {
       .define(LogCleanerEnableProp, BOOLEAN, Defaults.LogCleanerEnable, MEDIUM, LogCleanerEnableDoc)
       .define(LogCleanerDeleteRetentionMsProp, LONG, Defaults.LogCleanerDeleteRetentionMs, MEDIUM, LogCleanerDeleteRetentionMsDoc)
       .define(LogCleanerMinCompactionLagMsProp, LONG, Defaults.LogCleanerMinCompactionLagMs, MEDIUM, LogCleanerMinCompactionLagMsDoc)
+      .define(LogCleanerMaxCompactionLagMsProp, LONG, Defaults.LogCleanerMaxCompactionLagMs, MEDIUM, LogCleanerMaxCompactionLagMsDoc)
       .define(LogIndexSizeMaxBytesProp, INT, Defaults.LogIndexSizeMaxBytes, atLeast(4), MEDIUM, LogIndexSizeMaxBytesDoc)
       .define(LogIndexIntervalBytesProp, INT, Defaults.LogIndexIntervalBytes, atLeast(0), MEDIUM, LogIndexIntervalBytesDoc)
       .define(LogFlushIntervalMessagesProp, LONG, Defaults.LogFlushIntervalMessages, atLeast(1), HIGH, LogFlushIntervalMessagesDoc)
@@ -949,6 +1047,7 @@ object KafkaConfig {
       .define(InterBrokerSecurityProtocolProp, STRING, Defaults.InterBrokerSecurityProtocol, MEDIUM, InterBrokerSecurityProtocolDoc)
       .define(InterBrokerProtocolVersionProp, STRING, Defaults.InterBrokerProtocolVersion, ApiVersionValidator, MEDIUM, InterBrokerProtocolVersionDoc)
       .define(InterBrokerListenerNameProp, STRING, null, MEDIUM, InterBrokerListenerNameDoc)
+      .define(ReplicaSelectorClassProp, STRING, null, MEDIUM, ReplicaSelectorClassDoc)
 
       /** ********* Controlled shutdown configuration ***********/
       .define(ControlledShutdownMaxRetriesProp, INT, Defaults.ControlledShutdownMaxRetries, MEDIUM, ControlledShutdownMaxRetriesDoc)
@@ -974,6 +1073,9 @@ object KafkaConfig {
       .define(OffsetCommitRequiredAcksProp, SHORT, Defaults.OffsetCommitRequiredAcks, HIGH, OffsetCommitRequiredAcksDoc)
       .define(DeleteTopicEnableProp, BOOLEAN, Defaults.DeleteTopicEnable, HIGH, DeleteTopicEnableDoc)
       .define(CompressionTypeProp, STRING, Defaults.CompressionType, HIGH, CompressionTypeDoc)
+      .define(OffsetsTopicMaxMessageBytesProp, INT, Defaults.OffsetsTopicMaxMessageBytes, atLeast(0), HIGH, OffsetsTopicMaxMessageBytesDoc)
+      .define(OffsetsTopicMinInSyncReplicasProp, INT, Defaults.OffsetsTopicMinInSyncReplicas, atLeast(1), HIGH, OffsetsTopicMinInSyncReplicasDoc)
+      .define(OffsetsTopicMinCompactionLagMsProp, LONG, Defaults.OffsetsTopicMinCompactionLagMs, atLeast(0), HIGH, OffsetsTopicMinCompactionLagMsDoc)
 
       /** ********* Transaction management configuration ***********/
       .define(TransactionalIdExpirationMsProp, INT, Defaults.TransactionalIdExpirationMs, atLeast(1), HIGH, TransactionalIdExpirationMsDoc)
@@ -994,6 +1096,7 @@ object KafkaConfig {
       .define(MetricSampleWindowMsProp, LONG, Defaults.MetricSampleWindowMs, atLeast(1), LOW, MetricSampleWindowMsDoc)
       .define(MetricReporterClassesProp, LIST, Defaults.MetricReporterClasses, LOW, MetricReporterClassesDoc)
       .define(MetricRecordingLevelProp, STRING, Defaults.MetricRecordingLevel, LOW, MetricRecordingLevelDoc)
+      .define(MetricReplaceOnDuplicateProp, BOOLEAN, Defaults.MetricReplaceOnDuplicate, LOW, MetricReplaceOnDuplicateDoc)
 
       /** ********* Kafka Yammer Metrics Reporter Configuration for docs ***********/
       .define(KafkaMetricsReporterClassesProp, LIST, Defaults.KafkaMetricReporterClasses, LOW, KafkaMetricsReporterClassesDoc)
@@ -1012,11 +1115,13 @@ object KafkaConfig {
 
       /** ********* General Security Configuration ****************/
       .define(ConnectionsMaxReauthMsProp, LONG, Defaults.ConnectionsMaxReauthMsDefault, MEDIUM, ConnectionsMaxReauthMsDoc)
+      .define(securityProviderClassProp, STRING, null, LOW, securityProviderClassDoc)
 
       /** ********* SSL Configuration ****************/
       .define(PrincipalBuilderClassProp, CLASS, null, MEDIUM, PrincipalBuilderClassDoc)
       .define(SslProtocolProp, STRING, Defaults.SslProtocol, MEDIUM, SslProtocolDoc)
       .define(SslProviderProp, STRING, null, MEDIUM, SslProviderDoc)
+      .define(SslContextProviderClassProp, STRING, Defaults.SslContextProviderClass, MEDIUM, SslContextProviderClassDoc)
       .define(SslEnabledProtocolsProp, LIST, Defaults.SslEnabledProtocols, MEDIUM, SslEnabledProtocolsDoc)
       .define(SslKeystoreTypeProp, STRING, Defaults.SslKeystoreType, MEDIUM, SslKeystoreTypeDoc)
       .define(SslKeystoreLocationProp, STRING, null, MEDIUM, SslKeystoreLocationDoc)
@@ -1029,9 +1134,9 @@ object KafkaConfig {
       .define(SslTrustManagerAlgorithmProp, STRING, Defaults.SslTrustManagerAlgorithm, MEDIUM, SslTrustManagerAlgorithmDoc)
       .define(SslEndpointIdentificationAlgorithmProp, STRING, Defaults.SslEndpointIdentificationAlgorithm, LOW, SslEndpointIdentificationAlgorithmDoc)
       .define(SslSecureRandomImplementationProp, STRING, null, LOW, SslSecureRandomImplementationDoc)
-      .define(SslClientAuthProp, STRING, Defaults.SslClientAuth, in(Defaults.SslClientAuthRequired, Defaults.SslClientAuthRequested, Defaults.SslClientAuthNone), MEDIUM, SslClientAuthDoc)
+      .define(SslClientAuthProp, STRING, Defaults.SslClientAuthentication, in(Defaults.SslClientAuthenticationValidValues:_*), MEDIUM, SslClientAuthDoc)
       .define(SslCipherSuitesProp, LIST, Collections.emptyList(), MEDIUM, SslCipherSuitesDoc)
-      .define(SslPrincipalMappingRulesProp, LIST, Defaults.SslPrincipalMappingRules, LOW, SslPrincipalMappingRulesDoc)
+      .define(SslPrincipalMappingRulesProp, STRING, Defaults.SslPrincipalMappingRules, LOW, SslPrincipalMappingRulesDoc)
 
       /** ********* Sasl Configuration ****************/
       .define(SaslMechanismInterBrokerProtocolProp, STRING, Defaults.SaslMechanismInterBrokerProtocol, MEDIUM, SaslMechanismInterBrokerProtocolDoc)
@@ -1142,6 +1247,12 @@ class KafkaConfig(val props: java.util.Map[_, _], doLog: Boolean, dynamicConfigO
   def numIoThreads = getInt(KafkaConfig.NumIoThreadsProp)
   def messageMaxBytes = getInt(KafkaConfig.MessageMaxBytesProp)
   val requestTimeoutMs = getInt(KafkaConfig.RequestTimeoutMsProp)
+  val heapDumpFolder = new File(getString(KafkaConfig.HeapDumpFolderProp))
+  val heapDumpTimeout = getLong(KafkaConfig.HeapDumpTimeoutProp)
+  val producerBatchDecompressionEnable = getBoolean(KafkaConfig.ProducerBatchDecompressionEnableProp)
+
+  var preferredController = getBoolean(KafkaConfig.PreferredControllerProp)
+  def allowPreferredControllerFallback: Boolean = getBoolean(KafkaConfig.AllowPreferredControllerFallbackProp)
 
   def getNumReplicaAlterLogDirsThreads: Int = {
     val numThreads: Integer = Option(getInt(KafkaConfig.NumReplicaAlterLogDirsThreadsProp)).getOrElse(logDirs.size)
@@ -1149,7 +1260,23 @@ class KafkaConfig(val props: java.util.Map[_, _], doLog: Boolean, dynamicConfigO
   }
 
   /************* Authorizer Configuration ***********/
-  val authorizerClassName: String = getString(KafkaConfig.AuthorizerClassNameProp)
+  val authorizer: Option[Authorizer] = {
+    val className = getString(KafkaConfig.AuthorizerClassNameProp)
+    if (className == null || className.isEmpty)
+      None
+    else {
+      val authZ = Utils.newInstance(className, classOf[Object]) match {
+        case auth: Authorizer => auth
+        case auth: kafka.security.auth.Authorizer => new AuthorizerWrapper(auth)
+        case auth => throw new ConfigException(s"Authorizer does not implement ${classOf[Authorizer].getName} or kafka.security.auth.Authorizer .")
+      }
+      Some(authZ)
+    }
+  }
+
+  /************* Broker-side Observer Configuration ********/
+  val ObserverClassName: String = getString(KafkaConfig.ObserverClassNameProp)
+  val ObserverShutdownTimeoutMs: Long = getLong(KafkaConfig.ObserverShutdownTimeoutMsProp)
 
   /** ********* Socket Server Configuration ***********/
   val hostName = getString(KafkaConfig.HostNameProp)
@@ -1160,17 +1287,22 @@ class KafkaConfig(val props: java.util.Map[_, _], doLog: Boolean, dynamicConfigO
   val socketSendBufferBytes = getInt(KafkaConfig.SocketSendBufferBytesProp)
   val socketReceiveBufferBytes = getInt(KafkaConfig.SocketReceiveBufferBytesProp)
   val socketRequestMaxBytes = getInt(KafkaConfig.SocketRequestMaxBytesProp)
+  val socketRequestCommonBytes = getInt(KafkaConfig.SocketRequestCommonBytesProp)
+  val socketRequestBufferCacheSize = getInt(KafkaConfig.SocketRequestBufferCacheSizeProp)
+  val requestMaxLocalTimeMs = getLong(KafkaConfig.RequestMaxLocalTimeMsProp)
   val maxConnectionsPerIp = getInt(KafkaConfig.MaxConnectionsPerIpProp)
   val maxConnectionsPerIpOverrides: Map[String, Int] =
     getMap(KafkaConfig.MaxConnectionsPerIpOverridesProp, getString(KafkaConfig.MaxConnectionsPerIpOverridesProp)).map { case (k, v) => (k, v.toInt)}
+  def maxConnections = getInt(KafkaConfig.MaxConnectionsProp)
   val connectionsMaxIdleMs = getLong(KafkaConfig.ConnectionsMaxIdleMsProp)
   val failedAuthenticationDelayMs = getInt(KafkaConfig.FailedAuthenticationDelayMsProp)
 
   /***************** rack configuration **************/
   val rack = Option(getString(KafkaConfig.RackProp))
+  val replicaSelectorClassName = Option(getString(KafkaConfig.ReplicaSelectorClassProp))
 
   /** ********* Log Configuration ***********/
-  val autoCreateTopicsEnable = getBoolean(KafkaConfig.AutoCreateTopicsEnableProp)
+  def autoCreateTopicsEnable: java.lang.Boolean = getBoolean(KafkaConfig.AutoCreateTopicsEnableProp)
   val numPartitions = getInt(KafkaConfig.NumPartitionsProp)
   val logDirs = CoreUtils.parseCsvList(Option(getString(KafkaConfig.LogDirsProp)).getOrElse(getString(KafkaConfig.LogDirProp)))
   def logSegmentBytes = getInt(KafkaConfig.LogSegmentBytesProp)
@@ -1191,6 +1323,7 @@ class KafkaConfig(val props: java.util.Map[_, _], doLog: Boolean, dynamicConfigO
   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)
@@ -1256,6 +1389,9 @@ class KafkaConfig(val props: java.util.Map[_, _], doLog: Boolean, dynamicConfigO
   val offsetCommitRequiredAcks = getShort(KafkaConfig.OffsetCommitRequiredAcksProp)
   val offsetsTopicSegmentBytes = getInt(KafkaConfig.OffsetsTopicSegmentBytesProp)
   val offsetsTopicCompressionCodec = Option(getInt(KafkaConfig.OffsetsTopicCompressionCodecProp)).map(value => CompressionCodec.getCompressionCodec(value)).orNull
+  val offsetsTopicMaxMessageBytes = getInt(KafkaConfig.OffsetsTopicMaxMessageBytesProp)
+  val offsetsTopicMinInSyncReplicas = getInt(KafkaConfig.OffsetsTopicMinInSyncReplicasProp)
+  val offsetsTopicMinCompactionLagMs = getLong(KafkaConfig.OffsetsTopicMinCompactionLagMsProp)
 
   /** ********* Transaction management configuration ***********/
   val transactionalIdExpirationMs = getInt(KafkaConfig.TransactionalIdExpirationMsProp)
@@ -1273,6 +1409,7 @@ class KafkaConfig(val props: java.util.Map[_, _], doLog: Boolean, dynamicConfigO
   val metricNumSamples = getInt(KafkaConfig.MetricNumSamplesProp)
   val metricSampleWindowMs = getLong(KafkaConfig.MetricSampleWindowMsProp)
   val metricRecordingLevel = getString(KafkaConfig.MetricRecordingLevelProp)
+  val metricReplaceOnDuplicate = getBoolean(KafkaConfig.MetricReplaceOnDuplicateProp)
 
   /** ********* SSL/SASL Configuration **************/
   // Security configs may be overridden for listeners, so it is not safe to use the base values
@@ -1288,8 +1425,8 @@ class KafkaConfig(val props: java.util.Map[_, _], doLog: Boolean, dynamicConfigO
 
   def interBrokerListenerName = getInterBrokerListenerNameAndSecurityProtocol._1
   def interBrokerSecurityProtocol = getInterBrokerListenerNameAndSecurityProtocol._2
-  def controlPlaneListenerName = getControlPlaneListenerNameAndSecurityProtocol.map { case (listenerName, securityProtocol) => listenerName }
-  def controlPlaneSecurityProtocol = getControlPlaneListenerNameAndSecurityProtocol.map { case (listenerName, securityProtocol) => securityProtocol }
+  def controlPlaneListenerName = getControlPlaneListenerNameAndSecurityProtocol.map { case (listenerName, _) => listenerName }
+  def controlPlaneSecurityProtocol = getControlPlaneListenerNameAndSecurityProtocol.map { case (_, securityProtocol) => securityProtocol }
   def saslMechanismInterBrokerProtocol = getString(KafkaConfig.SaslMechanismInterBrokerProtocolProp)
   val saslInterBrokerHandshakeRequestEnable = interBrokerProtocolVersion >= KAFKA_0_10_0_IV1
 
@@ -1318,9 +1455,6 @@ class KafkaConfig(val props: java.util.Map[_, _], doLog: Boolean, dynamicConfigO
   val numAlterLogDirsReplicationQuotaSamples = getInt(KafkaConfig.NumAlterLogDirsReplicationQuotaSamplesProp)
   val alterLogDirsReplicationQuotaWindowSizeSeconds = getInt(KafkaConfig.AlterLogDirsReplicationQuotaWindowSizeSecondsProp)
 
-  /** ********* Transaction Configuration **************/
-  val transactionIdExpirationMs = getInt(KafkaConfig.TransactionalIdExpirationMsProp)
-
   /** ********* Fetch Session Configuration **************/
   val maxIncrementalFetchSessionCacheSlots = getInt(KafkaConfig.MaxIncrementalFetchSessionCacheSlots)
 
@@ -1331,6 +1465,10 @@ class KafkaConfig(val props: java.util.Map[_, _], doLog: Boolean, dynamicConfigO
     dynamicConfig.addReconfigurable(reconfigurable)
   }
 
+  def removeReconfigurable(reconfigurable: Reconfigurable): Unit = {
+    dynamicConfig.removeReconfigurable(reconfigurable)
+  }
+
   def logRetentionTimeMillis: Long = {
     val millisInMinute = 60L * 1000L
     val millisInHour = 60L * millisInMinute
@@ -1346,6 +1484,10 @@ class KafkaConfig(val props: java.util.Map[_, _], doLog: Boolean, dynamicConfigO
     millis
   }
 
+  def getMaintenanceBrokerList: Seq[Int] = {
+    dynamicConfig.getMaintenanceBrokerList
+  }
+
   private def getMap(propName: String, propValue: String): Map[String, String] = {
     try {
       CoreUtils.parseCsvMap(propValue)
@@ -1436,7 +1578,7 @@ class KafkaConfig(val props: java.util.Map[_, _], doLog: Boolean, dynamicConfigO
 
   validateValues()
 
-  private def validateValues() {
+  private def validateValues(): Unit = {
     if(brokerIdGenerationEnable) {
       require(brokerId >= -1 && brokerId <= maxReservedBrokerId, "broker.id must be equal or greater than -1 and not greater than reserved.broker.max.id")
     } else {
diff --git a/core/src/main/scala/kafka/server/KafkaRequestHandler.scala b/core/src/main/scala/kafka/server/KafkaRequestHandler.scala
index e0ad1b60cacdf..d43b09256ab64 100755
--- a/core/src/main/scala/kafka/server/KafkaRequestHandler.scala
+++ b/core/src/main/scala/kafka/server/KafkaRequestHandler.scala
@@ -23,11 +23,12 @@ import kafka.metrics.KafkaMetricsGroup
 import java.util.concurrent.{CountDownLatch, TimeUnit}
 import java.util.concurrent.atomic.AtomicInteger
 
-import com.yammer.metrics.core.Meter
+import com.yammer.metrics.core.{Counter, Meter}
 import org.apache.kafka.common.internals.FatalExitError
 import org.apache.kafka.common.utils.{KafkaThread, Time}
 
 import scala.collection.mutable
+import scala.collection.JavaConverters._
 
 /**
  * A thread that answers kafka requests.
@@ -43,7 +44,7 @@ class KafkaRequestHandler(id: Int,
   private val shutdownComplete = new CountDownLatch(1)
   @volatile private var stopped = false
 
-  def run() {
+  def run(): Unit = {
     while (!stopped) {
       // We use a single meter for aggregate idle percentage for the thread pool.
       // Since meter is calculated as total_recorded_value / time_window and
@@ -146,44 +147,155 @@ class BrokerTopicMetrics(name: Option[String]) extends KafkaMetricsGroup {
     case Some(topic) => Map("topic" -> topic)
   }
 
-  val messagesInRate = newMeter(BrokerTopicStats.MessagesInPerSec, "messages", TimeUnit.SECONDS, tags)
-  val bytesInRate = newMeter(BrokerTopicStats.BytesInPerSec, "bytes", TimeUnit.SECONDS, tags)
-  val bytesOutRate = newMeter(BrokerTopicStats.BytesOutPerSec, "bytes", TimeUnit.SECONDS, tags)
-  val bytesRejectedRate = newMeter(BrokerTopicStats.BytesRejectedPerSec, "bytes", TimeUnit.SECONDS, tags)
-  private[server] val replicationBytesInRate =
-    if (name.isEmpty) Some(newMeter(BrokerTopicStats.ReplicationBytesInPerSec, "bytes", TimeUnit.SECONDS, tags))
+  case class MeterWrapper(metricType: String, eventType: String) {
+    @volatile private var lazyMeter: Meter = _
+    private val meterLock = new Object
+
+    def meter(): Meter = {
+      var meter = lazyMeter
+      if (meter == null) {
+        meterLock synchronized {
+          meter = lazyMeter
+          if (meter == null) {
+            meter = newMeter(metricType, eventType, TimeUnit.SECONDS, tags)
+            lazyMeter = meter
+          }
+        }
+      }
+      meter
+    }
+
+    def close(): Unit = meterLock synchronized {
+      if (lazyMeter != null) {
+        removeMetric(metricType, tags)
+        lazyMeter = null
+      }
+    }
+
+    if (tags.isEmpty) // greedily initialize the general topic metrics
+      meter()
+  }
+
+  case class CounterWrapper(metricType: String) {
+    @volatile private var lazyCounter: Counter = _
+    private val counterLock = new Object
+
+    def counter(): Counter = {
+      var counter = lazyCounter
+      if (counter == null) {
+        counterLock synchronized {
+          counter = lazyCounter
+          if (counter == null) {
+            counter = newCounter(metricType, tags)
+            lazyCounter = counter
+          }
+        }
+      }
+      counter
+    }
+
+    def close(): Unit = counterLock synchronized {
+      if (lazyCounter != null) {
+        removeMetric(metricType, tags)
+        lazyCounter = null
+      }
+    }
+
+    if (tags.isEmpty) // greedily initialize the general topic metrics
+    counter()
+  }
+
+  // an internal map for "lazy initialization" of certain metrics
+  private val metricTypeMap = new Pool[String, MeterWrapper]()
+  metricTypeMap.putAll(Map(
+    BrokerTopicStats.MessagesInPerSec -> MeterWrapper(BrokerTopicStats.MessagesInPerSec, "messages"),
+    BrokerTopicStats.BytesInPerSec -> MeterWrapper(BrokerTopicStats.BytesInPerSec, "bytes"),
+    BrokerTopicStats.BytesOutPerSec -> MeterWrapper(BrokerTopicStats.BytesOutPerSec, "bytes"),
+    BrokerTopicStats.BytesRejectedPerSec -> MeterWrapper(BrokerTopicStats.BytesRejectedPerSec, "bytes"),
+    BrokerTopicStats.FailedProduceRequestsPerSec -> MeterWrapper(BrokerTopicStats.FailedProduceRequestsPerSec, "requests"),
+    BrokerTopicStats.FailedFetchRequestsPerSec -> MeterWrapper(BrokerTopicStats.FailedFetchRequestsPerSec, "requests"),
+    BrokerTopicStats.TotalProduceRequestsPerSec -> MeterWrapper(BrokerTopicStats.TotalProduceRequestsPerSec, "requests"),
+    BrokerTopicStats.TotalFetchRequestsPerSec -> MeterWrapper(BrokerTopicStats.TotalFetchRequestsPerSec, "requests"),
+    BrokerTopicStats.FetchMessageConversionsPerSec -> MeterWrapper(BrokerTopicStats.FetchMessageConversionsPerSec, "requests"),
+    BrokerTopicStats.ProduceMessageConversionsPerSec -> MeterWrapper(BrokerTopicStats.ProduceMessageConversionsPerSec, "requests"),
+    BrokerTopicStats.NoKeyCompactedTopicRecordsPerSec -> MeterWrapper(BrokerTopicStats.NoKeyCompactedTopicRecordsPerSec, "requests"),
+    BrokerTopicStats.InvalidMagicNumberRecordsPerSec -> MeterWrapper(BrokerTopicStats.InvalidMagicNumberRecordsPerSec, "requests"),
+    BrokerTopicStats.InvalidMessageCrcRecordsPerSec -> MeterWrapper(BrokerTopicStats.InvalidMessageCrcRecordsPerSec, "requests"),
+    BrokerTopicStats.InvalidOffsetOrSequenceRecordsPerSec -> MeterWrapper(BrokerTopicStats.InvalidOffsetOrSequenceRecordsPerSec, "requests")
+  ).asJava)
+  if (name.isEmpty) {
+    metricTypeMap.put(BrokerTopicStats.ReplicationBytesInPerSec, MeterWrapper(BrokerTopicStats.ReplicationBytesInPerSec, "bytes"))
+    metricTypeMap.put(BrokerTopicStats.ReplicationBytesOutPerSec, MeterWrapper(BrokerTopicStats.ReplicationBytesOutPerSec, "bytes"))
+  }
+
+  private val counterMetricTypeMap = new Pool[String, CounterWrapper]()
+  counterMetricTypeMap.putAll(Map(
+    BrokerTopicStats.MessagesInTotal-> CounterWrapper(BrokerTopicStats.MessagesInTotal),
+    BrokerTopicStats.BytesInTotal -> CounterWrapper(BrokerTopicStats.BytesInTotal)
+  ).asJava)
+  // used for testing only
+  def metricMap: Map[String, MeterWrapper] = metricTypeMap.toMap
+
+  def counterMetricMap: Map[String, CounterWrapper] = counterMetricTypeMap.toMap
+
+  def messagesInRate = metricTypeMap.get(BrokerTopicStats.MessagesInPerSec).meter()
+
+  def messagesInTotal = counterMetricTypeMap.get(BrokerTopicStats.MessagesInTotal).counter()
+
+  def bytesInRate = metricTypeMap.get(BrokerTopicStats.BytesInPerSec).meter()
+
+  def bytesInTotal = counterMetricTypeMap.get(BrokerTopicStats.BytesInTotal).counter()
+
+  def bytesOutRate = metricTypeMap.get(BrokerTopicStats.BytesOutPerSec).meter()
+
+  def bytesRejectedRate = metricTypeMap.get(BrokerTopicStats.BytesRejectedPerSec).meter()
+
+  private[server] def replicationBytesInRate =
+    if (name.isEmpty) Some(metricTypeMap.get(BrokerTopicStats.ReplicationBytesInPerSec).meter())
     else None
-  private[server] val replicationBytesOutRate =
-    if (name.isEmpty) Some(newMeter(BrokerTopicStats.ReplicationBytesOutPerSec, "bytes", TimeUnit.SECONDS, tags))
+
+  private[server] def replicationBytesOutRate =
+    if (name.isEmpty) Some(metricTypeMap.get(BrokerTopicStats.ReplicationBytesOutPerSec).meter())
     else None
-  val failedProduceRequestRate = newMeter(BrokerTopicStats.FailedProduceRequestsPerSec, "requests", TimeUnit.SECONDS, tags)
-  val failedFetchRequestRate = newMeter(BrokerTopicStats.FailedFetchRequestsPerSec, "requests", TimeUnit.SECONDS, tags)
-  val totalProduceRequestRate = newMeter(BrokerTopicStats.TotalProduceRequestsPerSec, "requests", TimeUnit.SECONDS, tags)
-  val totalFetchRequestRate = newMeter(BrokerTopicStats.TotalFetchRequestsPerSec, "requests", TimeUnit.SECONDS, tags)
-  val fetchMessageConversionsRate = newMeter(BrokerTopicStats.FetchMessageConversionsPerSec, "requests", TimeUnit.SECONDS, tags)
-  val produceMessageConversionsRate = newMeter(BrokerTopicStats.ProduceMessageConversionsPerSec, "requests", TimeUnit.SECONDS, tags)
-
-  def close() {
-    removeMetric(BrokerTopicStats.MessagesInPerSec, tags)
-    removeMetric(BrokerTopicStats.BytesInPerSec, tags)
-    removeMetric(BrokerTopicStats.BytesOutPerSec, tags)
-    removeMetric(BrokerTopicStats.BytesRejectedPerSec, tags)
-    if (replicationBytesInRate.isDefined)
-      removeMetric(BrokerTopicStats.ReplicationBytesInPerSec, tags)
-    if (replicationBytesOutRate.isDefined)
-      removeMetric(BrokerTopicStats.ReplicationBytesOutPerSec, tags)
-    removeMetric(BrokerTopicStats.FailedProduceRequestsPerSec, tags)
-    removeMetric(BrokerTopicStats.FailedFetchRequestsPerSec, tags)
-    removeMetric(BrokerTopicStats.TotalProduceRequestsPerSec, tags)
-    removeMetric(BrokerTopicStats.TotalFetchRequestsPerSec, tags)
-    removeMetric(BrokerTopicStats.FetchMessageConversionsPerSec, tags)
-    removeMetric(BrokerTopicStats.ProduceMessageConversionsPerSec, tags)
+
+  def failedProduceRequestRate = metricTypeMap.get(BrokerTopicStats.FailedProduceRequestsPerSec).meter()
+
+  def failedFetchRequestRate = metricTypeMap.get(BrokerTopicStats.FailedFetchRequestsPerSec).meter()
+
+  def totalProduceRequestRate = metricTypeMap.get(BrokerTopicStats.TotalProduceRequestsPerSec).meter()
+
+  def totalFetchRequestRate = metricTypeMap.get(BrokerTopicStats.TotalFetchRequestsPerSec).meter()
+
+  def fetchMessageConversionsRate = metricTypeMap.get(BrokerTopicStats.FetchMessageConversionsPerSec).meter()
+
+  def produceMessageConversionsRate = metricTypeMap.get(BrokerTopicStats.ProduceMessageConversionsPerSec).meter()
+
+  def noKeyCompactedTopicRecordsPerSec = metricTypeMap.get(BrokerTopicStats.NoKeyCompactedTopicRecordsPerSec).meter()
+
+  def invalidMagicNumberRecordsPerSec = metricTypeMap.get(BrokerTopicStats.InvalidMagicNumberRecordsPerSec).meter()
+
+  def invalidMessageCrcRecordsPerSec = metricTypeMap.get(BrokerTopicStats.InvalidMessageCrcRecordsPerSec).meter()
+
+  def invalidOffsetOrSequenceRecordsPerSec = metricTypeMap.get(BrokerTopicStats.InvalidOffsetOrSequenceRecordsPerSec).meter()
+
+  def closeMetric(metricType: String): Unit = {
+    val meter = metricTypeMap.get(metricType)
+    if (meter != null)
+      meter.close()
+  }
+
+  def close(): Unit = {
+    metricTypeMap.values.foreach(_.close())
+    removeMetric(BrokerTopicStats.MessagesInTotal, tags)
+    removeMetric(BrokerTopicStats.BytesInTotal, tags)
   }
 }
 
 object BrokerTopicStats {
   val MessagesInPerSec = "MessagesInPerSec"
+  val MessagesInTotal = "MessagesInTotal"
   val BytesInPerSec = "BytesInPerSec"
+  val BytesInTotal = "BytesInTotal"
   val BytesOutPerSec = "BytesOutPerSec"
   val BytesRejectedPerSec = "BytesRejectedPerSec"
   val ReplicationBytesInPerSec = "ReplicationBytesInPerSec"
@@ -194,6 +306,13 @@ object BrokerTopicStats {
   val TotalFetchRequestsPerSec = "TotalFetchRequestsPerSec"
   val FetchMessageConversionsPerSec = "FetchMessageConversionsPerSec"
   val ProduceMessageConversionsPerSec = "ProduceMessageConversionsPerSec"
+
+  // These following topics are for LogValidator for better debugging on failed records
+  val NoKeyCompactedTopicRecordsPerSec = "NoKeyCompactedTopicRecordsPerSec"
+  val InvalidMagicNumberRecordsPerSec = "InvalidMagicNumberRecordsPerSec"
+  val InvalidMessageCrcRecordsPerSec = "InvalidMessageCrcRecordsPerSec"
+  val InvalidOffsetOrSequenceRecordsPerSec = "InvalidOffsetOrSequenceRecordsPerSec"
+
   private val valueFactory = (k: String) => new BrokerTopicMetrics(Some(k))
 }
 
@@ -206,25 +325,46 @@ class BrokerTopicStats {
   def topicStats(topic: String): BrokerTopicMetrics =
     stats.getAndMaybePut(topic)
 
-  def updateReplicationBytesIn(value: Long) {
+  def updateReplicationBytesIn(value: Long): Unit = {
     allTopicsStats.replicationBytesInRate.foreach { metric =>
       metric.mark(value)
     }
   }
 
-  private def updateReplicationBytesOut(value: Long) {
+  private def updateReplicationBytesOut(value: Long): Unit = {
     allTopicsStats.replicationBytesOutRate.foreach { metric =>
       metric.mark(value)
     }
   }
 
-  def removeMetrics(topic: String) {
+  // This method only removes metrics only used for leader
+  def removeOldLeaderMetrics(topic: String): Unit = {
+    val topicMetrics = topicStats(topic)
+    if (topicMetrics != null) {
+      topicMetrics.closeMetric(BrokerTopicStats.MessagesInPerSec)
+      topicMetrics.closeMetric(BrokerTopicStats.BytesInPerSec)
+      topicMetrics.closeMetric(BrokerTopicStats.BytesRejectedPerSec)
+      topicMetrics.closeMetric(BrokerTopicStats.FailedProduceRequestsPerSec)
+      topicMetrics.closeMetric(BrokerTopicStats.TotalProduceRequestsPerSec)
+      topicMetrics.closeMetric(BrokerTopicStats.ProduceMessageConversionsPerSec)
+      topicMetrics.closeMetric(BrokerTopicStats.ReplicationBytesOutPerSec)
+    }
+  }
+
+  // This method only removes metrics only used for follower
+  def removeOldFollowerMetrics(topic: String): Unit = {
+    val topicMetrics = topicStats(topic)
+    if (topicMetrics != null)
+      topicMetrics.closeMetric(BrokerTopicStats.ReplicationBytesInPerSec)
+  }
+
+  def removeMetrics(topic: String): Unit = {
     val metrics = stats.remove(topic)
     if (metrics != null)
       metrics.close()
   }
 
-  def updateBytesOut(topic: String, isFollower: Boolean, value: Long) {
+  def updateBytesOut(topic: String, isFollower: Boolean, value: Long): Unit = {
     if (isFollower) {
       updateReplicationBytesOut(value)
     } else {
@@ -233,7 +373,6 @@ class BrokerTopicStats {
     }
   }
 
-
   def close(): Unit = {
     allTopicsStats.close()
     stats.values.foreach(_.close())
diff --git a/core/src/main/scala/kafka/server/KafkaServer.scala b/core/src/main/scala/kafka/server/KafkaServer.scala
index 8fc5197b7e19d..155ea5fde5330 100755
--- a/core/src/main/scala/kafka/server/KafkaServer.scala
+++ b/core/src/main/scala/kafka/server/KafkaServer.scala
@@ -24,9 +24,9 @@ import java.util.concurrent._
 import java.util.concurrent.atomic.{AtomicBoolean, AtomicInteger}
 
 import com.yammer.metrics.core.Gauge
-import kafka.api.{KAFKA_0_9_0, KAFKA_2_2_IV0}
+import kafka.api.{KAFKA_0_9_0, KAFKA_2_2_IV0, KAFKA_2_4_IV1}
 import kafka.cluster.Broker
-import kafka.common.{GenerateBrokerIdException, InconsistentBrokerIdException}
+import kafka.common.{GenerateBrokerIdException, InconsistentBrokerIdException, InconsistentBrokerMetadataException, InconsistentClusterIdException}
 import kafka.controller.KafkaController
 import kafka.coordinator.group.GroupCoordinator
 import kafka.coordinator.transaction.TransactionCoordinator
@@ -34,11 +34,11 @@ import kafka.log.{LogConfig, LogManager}
 import kafka.metrics.{KafkaMetricsGroup, KafkaMetricsReporter}
 import kafka.network.SocketServer
 import kafka.security.CredentialProvider
-import kafka.security.auth.Authorizer
 import kafka.utils._
 import kafka.zk.{BrokerInfo, KafkaZkClient}
 import org.apache.kafka.clients.{ApiVersions, ClientDnsLookup, ManualMetadataUpdater, NetworkClient, NetworkClientUtils}
 import org.apache.kafka.common.internals.ClusterResourceListeners
+import org.apache.kafka.common.message.ControlledShutdownRequestData
 import org.apache.kafka.common.metrics.{JmxReporter, Metrics, _}
 import org.apache.kafka.common.network._
 import org.apache.kafka.common.protocol.Errors
@@ -46,8 +46,9 @@ import org.apache.kafka.common.requests.{ControlledShutdownRequest, ControlledSh
 import org.apache.kafka.common.security.scram.internals.ScramMechanism
 import org.apache.kafka.common.security.token.delegation.internals.DelegationTokenCache
 import org.apache.kafka.common.security.{JaasContext, JaasUtils}
-import org.apache.kafka.common.utils.{AppInfoParser, LogContext, Time}
-import org.apache.kafka.common.{ClusterResource, Node}
+import org.apache.kafka.common.utils.{AppInfoParser, LogContext, PoisonPill, Time}
+import org.apache.kafka.common.{ClusterResource, Endpoint, Node}
+import org.apache.kafka.server.authorizer.Authorizer
 
 import scala.collection.JavaConverters._
 import scala.collection.{Map, Seq, mutable}
@@ -69,6 +70,7 @@ object KafkaServer {
     logProps.put(LogConfig.IndexIntervalBytesProp, kafkaConfig.logIndexIntervalBytes)
     logProps.put(LogConfig.DeleteRetentionMsProp, kafkaConfig.logCleanerDeleteRetentionMs)
     logProps.put(LogConfig.MinCompactionLagMsProp, kafkaConfig.logCleanerMinCompactionLagMs)
+    logProps.put(LogConfig.MaxCompactionLagMsProp, kafkaConfig.logCleanerMaxCompactionLagMs)
     logProps.put(LogConfig.FileDeleteDelayMsProp, kafkaConfig.logDeleteDelayMs)
     logProps.put(LogConfig.MinCleanableDirtyRatioProp, kafkaConfig.logCleanerMinCleanRatio)
     logProps.put(LogConfig.CleanupPolicyProp, kafkaConfig.logCleanupPolicy)
@@ -80,6 +82,7 @@ object KafkaServer {
     logProps.put(LogConfig.MessageTimestampTypeProp, kafkaConfig.logMessageTimestampType.name)
     logProps.put(LogConfig.MessageTimestampDifferenceMaxMsProp, kafkaConfig.logMessageTimestampDifferenceMaxMs: java.lang.Long)
     logProps.put(LogConfig.MessageDownConversionEnableProp, kafkaConfig.logMessageDownConversionEnable: java.lang.Boolean)
+    logProps.put(LogConfig.ProducerBatchDecompressionEnableProp, kafkaConfig.producerBatchDecompressionEnable: java.lang.Boolean)
     logProps
   }
 
@@ -117,6 +120,7 @@ class KafkaServer(val config: KafkaConfig, time: Time = Time.SYSTEM, threadNameP
   var controlPlaneRequestProcessor: KafkaApis = null
 
   var authorizer: Option[Authorizer] = None
+  var observer: Observer = null
   var socketServer: SocketServer = null
   var dataPlaneRequestHandlerPool: KafkaRequestHandlerPool = null
   var controlPlaneRequestHandlerPool: KafkaRequestHandlerPool = null
@@ -152,6 +156,16 @@ class KafkaServer(val config: KafkaConfig, time: Time = Time.SYSTEM, threadNameP
   private var _clusterId: String = null
   private var _brokerTopicStats: BrokerTopicStats = null
 
+  private var healthCheckScheduler: KafkaScheduler = null
+
+  private def haltIfNotHealthy() {
+    // This relies on io-thread to receive request from RequestChannel with 300 ms timeout, so that lastDequeueTimeMs
+    // will keep increasing even if there is no incoming request
+    if (time.milliseconds - socketServer.dataPlaneRequestChannel.lastDequeueTimeMs > config.requestMaxLocalTimeMs) {
+      fatal(s"It has been more than ${config.requestMaxLocalTimeMs} ms since the last time any io-thread reads from RequestChannel. Shutdown broker now.")
+      PoisonPill.die(config.heapDumpFolder, config.heapDumpTimeout)
+    }
+  }
 
   def clusterId: String = _clusterId
 
@@ -187,7 +201,7 @@ class KafkaServer(val config: KafkaConfig, time: Time = Time.SYSTEM, threadNameP
    * Start up API for bringing up a single instance of the Kafka server.
    * Instantiates the LogManager, the SocketServer and the request handlers - KafkaRequestHandlers
    */
-  def startup() {
+  def startup(): Unit = {
     try {
       info("starting")
 
@@ -208,9 +222,17 @@ class KafkaServer(val config: KafkaConfig, time: Time = Time.SYSTEM, threadNameP
         _clusterId = getOrGenerateClusterId(zkClient)
         info(s"Cluster ID = $clusterId")
 
+        /* load metadata */
+        val (preloadedBrokerMetadataCheckpoint, initialOfflineDirs) = getBrokerMetadataAndOfflineDirs
+
+        /* check cluster id */
+        if (preloadedBrokerMetadataCheckpoint.clusterId.isDefined && preloadedBrokerMetadataCheckpoint.clusterId.get != clusterId)
+          throw new InconsistentClusterIdException(
+            s"The Cluster ID ${clusterId} doesn't match stored clusterId ${preloadedBrokerMetadataCheckpoint.clusterId} in meta.properties. " +
+            s"The broker is trying to join the wrong cluster. Configured zookeeper.connect may be wrong.")
+
         /* generate brokerId */
-        val (brokerId, initialOfflineDirs) = getBrokerIdAndOfflineDirs
-        config.brokerId = brokerId
+        config.brokerId = getOrGenerateBrokerId(preloadedBrokerMetadataCheckpoint)
         logContext = new LogContext(s"[KafkaServer id=${config.brokerId}] ")
         this.logIdent = logContext.logPrefix
 
@@ -227,11 +249,12 @@ class KafkaServer(val config: KafkaConfig, time: Time = Time.SYSTEM, threadNameP
         reporters.add(new JmxReporter(jmxPrefix))
         val metricConfig = KafkaServer.metricConfig(config)
         metrics = new Metrics(metricConfig, reporters, time, true)
+        metrics.setReplaceOnDuplicateMetric(config.metricReplaceOnDuplicate)
 
         /* register broker metrics */
         _brokerTopicStats = new BrokerTopicStats
 
-        quotaManagers = QuotaFactory.instantiate(config, metrics, time, threadNamePrefix.getOrElse(""))
+        quotaManagers = QuotaFactory.instantiate(config, metrics, time, Some(kafkaScheduler), threadNamePrefix.getOrElse(""))
         notifyClusterListeners(kafkaMetricsReporters ++ metrics.reporters.asScala)
 
         logDirFailureChannel = new LogDirFailureChannel(config.logDirs.size)
@@ -257,10 +280,20 @@ class KafkaServer(val config: KafkaConfig, time: Time = Time.SYSTEM, threadNameP
         replicaManager.startup()
 
         val brokerInfo = createBrokerInfo
+        if (config.preferredController) {
+          zkClient.registerPreferredControllerId(brokerInfo.broker.id)
+        }
         val brokerEpoch = zkClient.registerBroker(brokerInfo)
 
-        // Now that the broker id is successfully registered, checkpoint it
-        checkpointBrokerId(config.brokerId)
+        healthCheckScheduler = new KafkaScheduler(threads = 1, threadNamePrefix = "kafka-healthcheck-scheduler-")
+        healthCheckScheduler.startup()
+        healthCheckScheduler.schedule(name = "halt-broker-if-not-healthy",
+                                      fun = haltIfNotHealthy,
+                                      period = 10000,
+                                      unit = TimeUnit.MILLISECONDS)
+
+        // Now that the broker is successfully registered, checkpoint its metadata
+        checkpointBrokerMetadata(BrokerMetadata(config.brokerId, Some(clusterId)))
 
         /* start token manager */
         tokenManager = new DelegationTokenManager(config, tokenCache, time , zkClient)
@@ -270,11 +303,11 @@ class KafkaServer(val config: KafkaConfig, time: Time = Time.SYSTEM, threadNameP
         kafkaController = new KafkaController(config, zkClient, time, metrics, brokerInfo, brokerEpoch, tokenManager, threadNamePrefix)
         kafkaController.startup()
 
-        adminManager = new AdminManager(config, metrics, metadataCache, zkClient)
+        adminManager = new AdminManager(config, metrics, metadataCache, zkClient, kafkaController)
 
         /* start group coordinator */
         // Hardcode Time.SYSTEM for now as some Streams tests fail otherwise, it would be good to fix the underlying issue
-        groupCoordinator = GroupCoordinator(config, zkClient, replicaManager, Time.SYSTEM)
+        groupCoordinator = GroupCoordinator(config, zkClient, replicaManager, Time.SYSTEM, metrics)
         groupCoordinator.startup()
 
         /* start transaction coordinator, with a separate background thread scheduler for transaction expiration and log loading */
@@ -283,11 +316,23 @@ class KafkaServer(val config: KafkaConfig, time: Time = Time.SYSTEM, threadNameP
         transactionCoordinator.startup()
 
         /* Get the authorizer and initialize it if one is specified.*/
-        authorizer = Option(config.authorizerClassName).filter(_.nonEmpty).map { authorizerClassName =>
-          val authZ = CoreUtils.createObject[Authorizer](authorizerClassName)
-          authZ.configure(config.originals())
-          authZ
+        authorizer = config.authorizer
+        authorizer.foreach(_.configure(config.originals))
+        val authorizerFutures: Map[Endpoint, CompletableFuture[Void]] = authorizer match {
+          case Some(authZ) =>
+            authZ.start(brokerInfo.broker.toServerInfo(clusterId, config)).asScala.mapValues(_.toCompletableFuture).toMap
+          case None =>
+            brokerInfo.broker.endPoints.map { ep => ep.toJava -> CompletableFuture.completedFuture[Void](null) }.toMap
+        }
+
+        observer = try {
+          CoreUtils.createObject[Observer](config.ObserverClassName)
+        } catch {
+          case e: Exception =>
+            error(s"Creating observer instance from the given class name ${config.ObserverClassName} failed.", e)
+            new NoOpObserver
         }
+        observer.configure(config.originals())
 
         val fetchManager = new FetchManager(Time.SYSTEM,
           new FetchSessionCache(config.maxIncrementalFetchSessionCacheSlots,
@@ -295,7 +340,7 @@ class KafkaServer(val config: KafkaConfig, time: Time = Time.SYSTEM, threadNameP
 
         /* start processing requests */
         dataPlaneRequestProcessor = new KafkaApis(socketServer.dataPlaneRequestChannel, replicaManager, adminManager, groupCoordinator, transactionCoordinator,
-          kafkaController, zkClient, config.brokerId, config, metadataCache, metrics, authorizer, quotaManagers,
+          kafkaController, zkClient, config.brokerId, config, metadataCache, metrics, authorizer, observer, quotaManagers,
           fetchManager, brokerTopicStats, clusterId, time, tokenManager)
 
         dataPlaneRequestHandlerPool = new KafkaRequestHandlerPool(config.brokerId, socketServer.dataPlaneRequestChannel, dataPlaneRequestProcessor, time,
@@ -303,7 +348,7 @@ class KafkaServer(val config: KafkaConfig, time: Time = Time.SYSTEM, threadNameP
 
         socketServer.controlPlaneRequestChannelOpt.foreach { controlPlaneRequestChannel =>
           controlPlaneRequestProcessor = new KafkaApis(controlPlaneRequestChannel, replicaManager, adminManager, groupCoordinator, transactionCoordinator,
-            kafkaController, zkClient, config.brokerId, config, metadataCache, metrics, authorizer, quotaManagers,
+            kafkaController, zkClient, config.brokerId, config, metadataCache, metrics, authorizer, observer, quotaManagers,
             fetchManager, brokerTopicStats, clusterId, time, tokenManager)
 
           controlPlaneRequestHandlerPool = new KafkaRequestHandlerPool(config.brokerId, socketServer.controlPlaneRequestChannelOpt.get, controlPlaneRequestProcessor, time,
@@ -325,13 +370,13 @@ class KafkaServer(val config: KafkaConfig, time: Time = Time.SYSTEM, threadNameP
         dynamicConfigManager = new DynamicConfigManager(zkClient, dynamicConfigHandlers)
         dynamicConfigManager.startup()
 
-        socketServer.startDataPlaneProcessors()
-        socketServer.startControlPlaneProcessor()
+        socketServer.startControlPlaneProcessor(authorizerFutures)
+        socketServer.startDataPlaneProcessors(authorizerFutures)
         brokerState.newState(RunningAsBroker)
         shutdownLatch = new CountDownLatch(1)
         startupComplete.set(true)
         isStartingUp.set(false)
-        AppInfoParser.registerAppInfo(jmxPrefix, config.brokerId.toString, metrics)
+        AppInfoParser.registerAppInfo(jmxPrefix, config.brokerId.toString, metrics, time.milliseconds())
         info("started")
       }
     }
@@ -359,7 +404,7 @@ class KafkaServer(val config: KafkaConfig, time: Time = Time.SYSTEM, threadNameP
 
     def createZkClient(zkConnect: String, isSecure: Boolean) =
       KafkaZkClient(zkConnect, isSecure, config.zkSessionTimeoutMs, config.zkConnectionTimeoutMs,
-        config.zkMaxInFlightRequests, time)
+        config.zkMaxInFlightRequests, time, name = Some("Kafka server"))
 
     val chrootIndex = config.zkConnect.indexOf("/")
     val chrootOption = {
@@ -371,7 +416,8 @@ class KafkaServer(val config: KafkaConfig, time: Time = Time.SYSTEM, threadNameP
     val isZkSecurityEnabled = JaasUtils.isZkSecurityEnabled()
 
     if (secureAclsEnabled && !isZkSecurityEnabled)
-      throw new java.lang.SecurityException(s"${KafkaConfig.ZkEnableSecureAclsProp} is true, but the verification of the JAAS login file failed.")
+      throw new java.lang.SecurityException(s"${KafkaConfig.ZkEnableSecureAclsProp} is true, but the " +
+        s"verification of the JAAS login file failed ${JaasUtils.zkSecuritySysConfigString}")
 
     // make sure chroot path exists
     chrootOption.foreach { chroot =>
@@ -419,7 +465,7 @@ class KafkaServer(val config: KafkaConfig, time: Time = Time.SYSTEM, threadNameP
   /**
    * Performs controlled shutdown
    */
-  private def controlledShutdown() {
+  private def controlledShutdown(): Unit = {
 
     def node(broker: Broker): Node = broker.node(config.interBrokerListenerName)
 
@@ -514,21 +560,25 @@ class KafkaServer(val config: KafkaConfig, time: Time = Time.SYSTEM, threadNameP
               val controlledShutdownApiVersion: Short =
                 if (config.interBrokerProtocolVersion < KAFKA_0_9_0) 0
                 else if (config.interBrokerProtocolVersion < KAFKA_2_2_IV0) 1
-                else 2
-
-              val controlledShutdownRequest = new ControlledShutdownRequest.Builder(config.brokerId,
-                kafkaController.brokerEpoch, controlledShutdownApiVersion)
+                else if (config.interBrokerProtocolVersion < KAFKA_2_4_IV1) 2
+                else 3
+
+              val controlledShutdownRequest = new ControlledShutdownRequest.Builder(
+                  new ControlledShutdownRequestData()
+                    .setBrokerId(config.brokerId)
+                    .setBrokerEpoch(kafkaController.brokerEpoch),
+                    controlledShutdownApiVersion)
               val request = networkClient.newClientRequest(node(prevController).idString, controlledShutdownRequest,
                 time.milliseconds(), true)
               val clientResponse = NetworkClientUtils.sendAndReceive(networkClient, request, time)
 
               val shutdownResponse = clientResponse.responseBody.asInstanceOf[ControlledShutdownResponse]
-              if (shutdownResponse.error == Errors.NONE && shutdownResponse.partitionsRemaining.isEmpty) {
+              if (shutdownResponse.error == Errors.NONE && shutdownResponse.data.remainingPartitions.isEmpty) {
                 shutdownSucceeded = true
                 info("Controlled shutdown succeeded")
               }
               else {
-                info(s"Remaining partitions to move: ${shutdownResponse.partitionsRemaining.asScala.mkString(",")}")
+                info(s"Remaining partitions to move: ${shutdownResponse.data.remainingPartitions}")
                 info(s"Error from controller: ${shutdownResponse.error}")
               }
             }
@@ -571,7 +621,7 @@ class KafkaServer(val config: KafkaConfig, time: Time = Time.SYSTEM, threadNameP
    * Shutdown API for shutting down a single instance of the Kafka server.
    * Shuts down the LogManager, the SocketServer and the log cleaner scheduler thread
    */
-  def shutdown() {
+  def shutdown(): Unit = {
     try {
       info("shutting down")
 
@@ -585,6 +635,9 @@ class KafkaServer(val config: KafkaConfig, time: Time = Time.SYSTEM, threadNameP
         CoreUtils.swallow(controlledShutdown(), this)
         brokerState.newState(BrokerShuttingDown)
 
+        if (healthCheckScheduler != null)
+          healthCheckScheduler.shutdown()
+
         if (dynamicConfigManager != null)
           CoreUtils.swallow(dynamicConfigManager.shutdown(), this)
 
@@ -604,6 +657,9 @@ class KafkaServer(val config: KafkaConfig, time: Time = Time.SYSTEM, threadNameP
         if (controlPlaneRequestProcessor != null)
           CoreUtils.swallow(controlPlaneRequestProcessor.close(), this)
         CoreUtils.swallow(authorizer.foreach(_.close()), this)
+
+        CoreUtils.swallow(observer.close(config.ObserverShutdownTimeoutMs, TimeUnit.MILLISECONDS), this)
+
         if (adminManager != null)
           CoreUtils.swallow(adminManager.shutdown(), this)
 
@@ -669,28 +725,24 @@ class KafkaServer(val config: KafkaConfig, time: Time = Time.SYSTEM, threadNameP
   def boundPort(listenerName: ListenerName): Int = socketServer.boundPort(listenerName)
 
   /**
-    * Generates new brokerId if enabled or reads from meta.properties based on following conditions
-    * 
        
-    * 
      1.  config has no broker.id provided and broker id generation is enabled, generates a broker.id based on Zookeeper's sequence
-    * 
      2.  stored broker.id in meta.properties doesn't match in all the log.dirs throws InconsistentBrokerIdException
-    * 
      3.  config has broker.id and meta.properties contains broker.id if they don't match throws InconsistentBrokerIdException
-    * 
      4.  config has broker.id and there is no meta.properties file, creates new meta.properties and stores broker.id
-    * 
          
-    *
-    * The log directories whose meta.properties can not be accessed due to IOException will be returned to the caller
-    *
-    * @return A 2-tuple containing the brokerId and a sequence of offline log directories.
-    */
-  private def getBrokerIdAndOfflineDirs: (Int, Seq[String]) = {
-    var brokerId = config.brokerId
-    val brokerIdSet = mutable.HashSet[Int]()
+   * Reads the BrokerMetadata. If the BrokerMetadata doesn't match in all the log.dirs, InconsistentBrokerMetadataException is
+   * thrown.
+   *
+   * The log directories whose meta.properties can not be accessed due to IOException will be returned to the caller
+   *
+   * @return A 2-tuple containing the brokerMetadata and a sequence of offline log directories.
+   */
+  private def getBrokerMetadataAndOfflineDirs: (BrokerMetadata, Seq[String]) = {
+    val brokerMetadataMap = mutable.HashMap[String, BrokerMetadata]()
+    val brokerMetadataSet = mutable.HashSet[BrokerMetadata]()
     val offlineDirs = mutable.ArrayBuffer.empty[String]
 
     for (logDir <- config.logDirs) {
       try {
         val brokerMetadataOpt = brokerMetadataCheckpoints(logDir).read()
         brokerMetadataOpt.foreach { brokerMetadata =>
-          brokerIdSet.add(brokerMetadata.brokerId)
+          brokerMetadataMap += (logDir -> brokerMetadata)
+          brokerMetadataSet += brokerMetadata
         }
       } catch {
         case e: IOException =>
@@ -699,39 +751,61 @@ class KafkaServer(val config: KafkaConfig, time: Time = Time.SYSTEM, threadNameP
       }
     }
 
-    if (brokerIdSet.size > 1)
-      throw new InconsistentBrokerIdException(
-        s"Failed to match broker.id across log.dirs. This could happen if multiple brokers shared a log directory (log.dirs) " +
-        s"or partial data was manually copied from another broker. Found $brokerIdSet")
-    else if (brokerId >= 0 && brokerIdSet.size == 1 && brokerIdSet.last != brokerId)
-      throw new InconsistentBrokerIdException(
-        s"Configured broker.id $brokerId doesn't match stored broker.id ${brokerIdSet.last} in meta.properties. " +
-        s"If you moved your data, make sure your configured broker.id matches. " +
-        s"If you intend to create a new broker, you should remove all data in your data directories (log.dirs).")
-    else if (brokerIdSet.isEmpty && brokerId < 0 && config.brokerIdGenerationEnable) // generate a new brokerId from Zookeeper
-      brokerId = generateBrokerId
-    else if (brokerIdSet.size == 1) // pick broker.id from meta.properties
-      brokerId = brokerIdSet.last
+    if (brokerMetadataSet.size > 1) {
+      val builder = StringBuilder.newBuilder
 
+      for ((logDir, brokerMetadata) <- brokerMetadataMap)
+        builder ++= s"- $logDir -> $brokerMetadata\n"
 
-    (brokerId, offlineDirs)
+      throw new InconsistentBrokerMetadataException(
+        s"BrokerMetadata is not consistent across log.dirs. This could happen if multiple brokers shared a log directory (log.dirs) " +
+        s"or partial data was manually copied from another broker. Found:\n${builder.toString()}"
+      )
+    } else if (brokerMetadataSet.size == 1)
+      (brokerMetadataSet.last, offlineDirs)
+    else
+      (BrokerMetadata(-1, None), offlineDirs)
   }
 
-  private def checkpointBrokerId(brokerId: Int) {
-    var logDirsWithoutMetaProps: List[String] = List()
 
+  /**
+   * Checkpoint the BrokerMetadata to all the online log.dirs
+   *
+   * @param brokerMetadata
+   */
+  private def checkpointBrokerMetadata(brokerMetadata: BrokerMetadata) = {
     for (logDir <- config.logDirs if logManager.isLogDirOnline(new File(logDir).getAbsolutePath)) {
-      val brokerMetadataOpt = brokerMetadataCheckpoints(logDir).read()
-      if (brokerMetadataOpt.isEmpty)
-        logDirsWithoutMetaProps ++= List(logDir)
-    }
-
-    for (logDir <- logDirsWithoutMetaProps) {
       val checkpoint = brokerMetadataCheckpoints(logDir)
-      checkpoint.write(BrokerMetadata(brokerId))
+      checkpoint.write(brokerMetadata)
     }
   }
 
+  /**
+   * Generates new brokerId if enabled or reads from meta.properties based on following conditions
+   * 
            
+   * 
          1.  config has no broker.id provided and broker id generation is enabled, generates a broker.id based on Zookeeper's sequence
+   * 
          2.  config has broker.id and meta.properties contains broker.id if they don't match throws InconsistentBrokerIdException
+   * 
          3.  config has broker.id and there is no meta.properties file, creates new meta.properties and stores broker.id
+   * 
              
+   *
+   * @return The brokerId.
+   */
+  private def getOrGenerateBrokerId(brokerMetadata: BrokerMetadata): Int = {
+    val brokerId = config.brokerId
+
+    if (brokerId >= 0 && brokerMetadata.brokerId >= 0 && brokerMetadata.brokerId != brokerId)
+      throw new InconsistentBrokerIdException(
+        s"Configured broker.id $brokerId doesn't match stored broker.id ${brokerMetadata.brokerId} in meta.properties. " +
+        s"If you moved your data, make sure your configured broker.id matches. " +
+        s"If you intend to create a new broker, you should remove all data in your data directories (log.dirs).")
+    else if (brokerMetadata.brokerId < 0 && brokerId < 0 && config.brokerIdGenerationEnable) // generate a new brokerId from Zookeeper
+      generateBrokerId
+    else if (brokerMetadata.brokerId >= 0) // pick broker.id from meta.properties
+      brokerMetadata.brokerId
+    else
+      brokerId
+  }
+
   /**
     * Return a sequence id generated by updating the broker sequence id path in ZK.
     * Users can provide brokerId in the config. To avoid conflicts between ZK generated
diff --git a/core/src/main/scala/kafka/server/KafkaServerStartable.scala b/core/src/main/scala/kafka/server/KafkaServerStartable.scala
index bbe9abae63fc3..d4f4c152db0a8 100644
--- a/core/src/main/scala/kafka/server/KafkaServerStartable.scala
+++ b/core/src/main/scala/kafka/server/KafkaServerStartable.scala
@@ -22,19 +22,25 @@ import java.util.Properties
 import kafka.metrics.KafkaMetricsReporter
 import kafka.utils.{Exit, Logging, VerifiableProperties}
 
+import scala.collection.Seq
+
 object KafkaServerStartable {
-  def fromProps(serverProps: Properties) = {
+  def fromProps(serverProps: Properties): KafkaServerStartable = {
+    fromProps(serverProps, None)
+  }
+
+  def fromProps(serverProps: Properties, threadNamePrefix: Option[String]): KafkaServerStartable = {
     val reporters = KafkaMetricsReporter.startReporters(new VerifiableProperties(serverProps))
-    new KafkaServerStartable(KafkaConfig.fromProps(serverProps, false), reporters)
+    new KafkaServerStartable(KafkaConfig.fromProps(serverProps, false), reporters, threadNamePrefix)
   }
 }
 
-class KafkaServerStartable(val staticServerConfig: KafkaConfig, reporters: Seq[KafkaMetricsReporter]) extends Logging {
-  private val server = new KafkaServer(staticServerConfig, kafkaMetricsReporters = reporters)
+class KafkaServerStartable(val staticServerConfig: KafkaConfig, reporters: Seq[KafkaMetricsReporter], threadNamePrefix: Option[String] = None) extends Logging {
+  private val server = new KafkaServer(staticServerConfig, kafkaMetricsReporters = reporters, threadNamePrefix = threadNamePrefix)
 
   def this(serverConfig: KafkaConfig) = this(serverConfig, Seq.empty)
 
-  def startup() {
+  def startup(): Unit = {
     try server.startup()
     catch {
       case _: Throwable =>
@@ -44,7 +50,7 @@ class KafkaServerStartable(val staticServerConfig: KafkaConfig, reporters: Seq[K
     }
   }
 
-  def shutdown() {
+  def shutdown(): Unit = {
     try server.shutdown()
     catch {
       case _: Throwable =>
@@ -58,7 +64,7 @@ class KafkaServerStartable(val staticServerConfig: KafkaConfig, reporters: Seq[K
    * Allow setting broker state from the startable.
    * This is needed when a custom kafka server startable want to emit new states that it introduces.
    */
-  def setServerState(newState: Byte) {
+  def setServerState(newState: Byte): Unit = {
     server.brokerState.newState(newState)
   }
 
diff --git a/core/src/main/scala/kafka/server/LiCreateTopicPolicy.scala b/core/src/main/scala/kafka/server/LiCreateTopicPolicy.scala
new file mode 100644
index 0000000000000..56c0f46e5da89
--- /dev/null
+++ b/core/src/main/scala/kafka/server/LiCreateTopicPolicy.scala
@@ -0,0 +1,61 @@
+/**
+ * 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 org.apache.kafka.common.errors.PolicyViolationException
+import org.apache.kafka.server.policy.CreateTopicPolicy
+
+class LiCreateTopicPolicy extends CreateTopicPolicy {
+  var minRf = 0
+
+  @throws[PolicyViolationException]
+  override def validate(requestMetadata: CreateTopicPolicy.RequestMetadata): Unit = {
+    val requestTopic = requestMetadata.topic
+    val requestRF = requestMetadata.replicationFactor()
+
+    // For scala 2.11.x, a scala object converted from a null java object via asScala is NOT null. This could lead to
+    // NPE, thus we check the java object directly. AdminClientIntegrationTest verifies the behavior.
+    if (requestMetadata.replicasAssignments() == null && requestRF == null)
+      throw new PolicyViolationException(s"Topic [$requestTopic] is missing both replica assignment and " +
+        s"replication factor.")
+
+    // In createTopics() in AdminManager, replicationFactor and replicasAssignments are not both set at same time. We
+    // follow the same rationale here and prioritize replicasAssignments over replicationFactor
+    if (requestMetadata.replicasAssignments() != null) {
+      import collection.JavaConverters._
+      requestMetadata.replicasAssignments().asScala.foreach { case (p, assignment) =>
+          if (assignment.size() < minRf)
+            throw new PolicyViolationException(s"Topic [$requestTopic] fails RF requirement. Received RF for " +
+              s"[partition-$p]: ${assignment.size()}, min required RF: $minRf.")
+      }
+    } else if (requestRF < minRf) {
+      throw new PolicyViolationException(s"Topic [$requestTopic] fails RF requirement. " +
+        s"Received RF: ${requestMetadata.replicationFactor}, min required RF: $minRf.")
+    }
+  }
+
+  /**
+    * Configure this class with the given key-value pairs
+    */
+  override def configure(configs: java.util.Map[String, _]): Unit = {
+    minRf = configs.get(KafkaConfig.DefaultReplicationFactorProp).asInstanceOf[String].toInt
+  }
+
+  @throws[Exception]
+  override def close(): Unit = {
+  }
+}
diff --git a/core/src/main/scala/kafka/server/LogOffsetMetadata.scala b/core/src/main/scala/kafka/server/LogOffsetMetadata.scala
index effbaa04ea081..6423cfcf9eabc 100644
--- a/core/src/main/scala/kafka/server/LogOffsetMetadata.scala
+++ b/core/src/main/scala/kafka/server/LogOffsetMetadata.scala
@@ -17,11 +17,11 @@
 
 package kafka.server
 
+import kafka.log.Log
 import org.apache.kafka.common.KafkaException
 
 object LogOffsetMetadata {
-  val UnknownOffsetMetadata = new LogOffsetMetadata(-1, 0, 0)
-  val UnknownSegBaseOffset = -1L
+  val UnknownOffsetMetadata = LogOffsetMetadata(-1, 0, 0)
   val UnknownFilePosition = -1
 
   class OffsetOrdering extends Ordering[LogOffsetMetadata] {
@@ -39,7 +39,7 @@ object LogOffsetMetadata {
  *  3. the physical position on the located segment
  */
 case class LogOffsetMetadata(messageOffset: Long,
-                             segmentBaseOffset: Long = LogOffsetMetadata.UnknownSegBaseOffset,
+                             segmentBaseOffset: Long = Log.UnknownOffset,
                              relativePositionInSegment: Int = LogOffsetMetadata.UnknownFilePosition) {
 
   // check if this offset is already on an older segment compared with the given offset
@@ -76,7 +76,7 @@ case class LogOffsetMetadata(messageOffset: Long,
 
   // decide if the offset metadata only contains message offset info
   def messageOffsetOnly: Boolean = {
-    segmentBaseOffset == LogOffsetMetadata.UnknownSegBaseOffset && relativePositionInSegment == LogOffsetMetadata.UnknownFilePosition
+    segmentBaseOffset == Log.UnknownOffset && relativePositionInSegment == LogOffsetMetadata.UnknownFilePosition
   }
 
   override def toString = s"(offset=$messageOffset segment=[$segmentBaseOffset:$relativePositionInSegment])"
diff --git a/core/src/main/scala/kafka/server/MetadataCache.scala b/core/src/main/scala/kafka/server/MetadataCache.scala
index ec5a2b9ed38b1..1ad0e41352e5a 100755
--- a/core/src/main/scala/kafka/server/MetadataCache.scala
+++ b/core/src/main/scala/kafka/server/MetadataCache.scala
@@ -28,10 +28,13 @@ import kafka.controller.StateChangeLogger
 import kafka.utils.CoreUtils._
 import kafka.utils.Logging
 import org.apache.kafka.common.internals.Topic
+import org.apache.kafka.common.message.UpdateMetadataRequestData.UpdateMetadataPartitionState
 import org.apache.kafka.common.{Cluster, Node, PartitionInfo, TopicPartition}
 import org.apache.kafka.common.network.ListenerName
 import org.apache.kafka.common.protocol.Errors
 import org.apache.kafka.common.requests.{MetadataResponse, UpdateMetadataRequest}
+import org.apache.kafka.common.security.auth.SecurityProtocol
+
 
 /**
  *  A cache for the state (e.g., current leader) of each partition. This cache is updated through
@@ -51,9 +54,10 @@ class MetadataCache(brokerId: Int) extends Logging {
   private val stateChangeLogger = new StateChangeLogger(brokerId, inControllerContext = false, None)
 
   // This method is the main hotspot when it comes to the performance of metadata requests,
-  // we should be careful about adding additional logic here.
+  // we should be careful about adding additional logic here. Relatedly, `brokers` is
+  // `Iterable[Integer]` instead of `Iterable[Int]` to avoid a collection copy.
   // filterUnavailableEndpoints exists to support v0 MetadataResponses
-  private def getEndpoints(snapshot: MetadataSnapshot, brokers: Iterable[Int], listenerName: ListenerName, filterUnavailableEndpoints: Boolean): Seq[Node] = {
+  private def getEndpoints(snapshot: MetadataSnapshot, brokers: Iterable[java.lang.Integer], listenerName: ListenerName, filterUnavailableEndpoints: Boolean): Seq[Node] = {
     val result = new mutable.ArrayBuffer[Node](math.min(snapshot.aliveBrokers.size, brokers.size))
     brokers.foreach { brokerId =>
       val endpoint = getAliveEndpoint(snapshot, brokerId, listenerName) match {
@@ -73,13 +77,15 @@ class MetadataCache(brokerId: Int) extends Logging {
     snapshot.partitionStates.get(topic).map { partitions =>
       partitions.map { case (partitionId, partitionState) =>
         val topicPartition = new TopicPartition(topic, partitionId.toInt)
-        val leaderBrokerId = partitionState.basePartitionState.leader
-        val leaderEpoch = partitionState.basePartitionState.leaderEpoch
+        val leaderBrokerId = partitionState.leader
+        val leaderEpoch = partitionState.leaderEpoch
         val maybeLeader = getAliveEndpoint(snapshot, leaderBrokerId, listenerName)
-        val replicas = partitionState.basePartitionState.replicas.asScala.map(_.toInt)
+        val replicas = partitionState.replicas.asScala
         val replicaInfo = getEndpoints(snapshot, replicas, listenerName, errorUnavailableEndpoints)
-        val offlineReplicaInfo = getEndpoints(snapshot, partitionState.offlineReplicas.asScala.map(_.toInt), listenerName, errorUnavailableEndpoints)
+        val offlineReplicaInfo = getEndpoints(snapshot, partitionState.offlineReplicas.asScala, listenerName, errorUnavailableEndpoints)
 
+        val isr = partitionState.isr.asScala
+        val isrInfo = getEndpoints(snapshot, isr, listenerName, errorUnavailableEndpoints)
         maybeLeader match {
           case None =>
             val error = if (!snapshot.aliveBrokers.contains(brokerId)) { // we are already holding the read lock
@@ -90,13 +96,10 @@ class MetadataCache(brokerId: Int) extends Logging {
               if (errorUnavailableListeners) Errors.LISTENER_NOT_FOUND else Errors.LEADER_NOT_AVAILABLE
             }
             new MetadataResponse.PartitionMetadata(error, partitionId.toInt, Node.noNode(),
-              Optional.empty(), replicaInfo.asJava, java.util.Collections.emptyList(),
+              Optional.empty(), replicaInfo.asJava, isrInfo.asJava,
               offlineReplicaInfo.asJava)
 
           case Some(leader) =>
-            val isr = partitionState.basePartitionState.isr.asScala.map(_.toInt)
-            val isrInfo = getEndpoints(snapshot, isr, listenerName, errorUnavailableEndpoints)
-
             if (replicaInfo.size < replicas.size) {
               debug(s"Error while fetching metadata for $topicPartition: replica information not available for " +
                 s"following brokers ${replicas.filterNot(replicaInfo.map(_.id).contains).mkString(",")}")
@@ -147,7 +150,7 @@ class MetadataCache(brokerId: Int) extends Logging {
     snapshot.partitionStates.keySet
   }
 
-  private def getAllPartitions(snapshot: MetadataSnapshot): Map[TopicPartition, UpdateMetadataRequest.PartitionState] = {
+  private def getAllPartitions(snapshot: MetadataSnapshot): Map[TopicPartition, UpdateMetadataPartitionState] = {
     snapshot.partitionStates.flatMap { case (topic, partitionStates) =>
       partitionStates.map { case (partition, state ) => (new TopicPartition(topic, partition.toInt), state) }
     }.toMap
@@ -157,23 +160,23 @@ class MetadataCache(brokerId: Int) extends Logging {
     topics -- metadataSnapshot.partitionStates.keySet
   }
 
-  def isBrokerAlive(brokerId: Int): Boolean = {
-    metadataSnapshot.aliveBrokers.contains(brokerId)
+  def getAliveBroker(brokerId: Int): Option[Broker] = {
+    metadataSnapshot.aliveBrokers.get(brokerId)
   }
 
   def getAliveBrokers: Seq[Broker] = {
     metadataSnapshot.aliveBrokers.values.toBuffer
   }
 
-  private def addOrUpdatePartitionInfo(partitionStates: mutable.AnyRefMap[String, mutable.LongMap[UpdateMetadataRequest.PartitionState]],
+  private def addOrUpdatePartitionInfo(partitionStates: mutable.AnyRefMap[String, mutable.LongMap[UpdateMetadataPartitionState]],
                                        topic: String,
                                        partitionId: Int,
-                                       stateInfo: UpdateMetadataRequest.PartitionState) {
+                                       stateInfo: UpdateMetadataPartitionState): Unit = {
     val infos = partitionStates.getOrElseUpdate(topic, mutable.LongMap())
     infos(partitionId) = stateInfo
   }
 
-  def getPartitionInfo(topic: String, partitionId: Int): Option[UpdateMetadataRequest.PartitionState] = {
+  def getPartitionInfo(topic: String, partitionId: Int): Option[UpdateMetadataPartitionState] = {
     metadataSnapshot.partitionStates.get(topic).flatMap(_.get(partitionId))
   }
 
@@ -183,7 +186,7 @@ class MetadataCache(brokerId: Int) extends Logging {
   def getPartitionLeaderEndpoint(topic: String, partitionId: Int, listenerName: ListenerName): Option[Node] = {
     val snapshot = metadataSnapshot
     snapshot.partitionStates.get(topic).flatMap(_.get(partitionId)) map { partitionInfo =>
-      val leaderId = partitionInfo.basePartitionState.leader
+      val leaderId = partitionInfo.leader
 
       snapshot.aliveNodes.get(leaderId) match {
         case Some(nodeMap) =>
@@ -194,6 +197,24 @@ class MetadataCache(brokerId: Int) extends Logging {
     }
   }
 
+  def getPartitionReplicaEndpoints(tp: TopicPartition, listenerName: ListenerName): Map[Int, Node] = {
+    val snapshot = metadataSnapshot
+    snapshot.partitionStates.get(tp.topic).flatMap(_.get(tp.partition)).map { partitionInfo =>
+      val replicaIds = partitionInfo.replicas
+      replicaIds.asScala
+        .map(replicaId => replicaId.intValue() -> {
+          snapshot.aliveBrokers.get(replicaId.longValue()) match {
+            case Some(broker) =>
+              broker.getNode(listenerName).getOrElse(Node.noNode())
+            case None =>
+              Node.noNode()
+          }}).toMap
+        .filter(pair => pair match {
+          case (_, node) => !node.isEmpty
+        })
+    }.getOrElse(Map.empty[Int, Node])
+  }
+
   def getControllerId: Option[Int] = metadataSnapshot.controllerId
 
   def getClusterMetadata(clusterId: String, listenerName: ListenerName): Cluster = {
@@ -201,17 +222,17 @@ class MetadataCache(brokerId: Int) extends Logging {
     val nodes = snapshot.aliveNodes.map { case (id, nodes) => (id, nodes.get(listenerName).orNull) }
     def node(id: Integer): Node = nodes.get(id.toLong).orNull
     val partitions = getAllPartitions(snapshot)
-      .filter { case (_, state) => state.basePartitionState.leader != LeaderAndIsr.LeaderDuringDelete }
+      .filter { case (_, state) => state.leader != LeaderAndIsr.LeaderDuringDelete }
       .map { case (tp, state) =>
-        new PartitionInfo(tp.topic, tp.partition, node(state.basePartitionState.leader),
-          state.basePartitionState.replicas.asScala.map(node).toArray,
-          state.basePartitionState.isr.asScala.map(node).toArray,
+        new PartitionInfo(tp.topic, tp.partition, node(state.leader),
+          state.replicas.asScala.map(node).toArray,
+          state.isr.asScala.map(node).toArray,
           state.offlineReplicas.asScala.map(node).toArray)
       }
     val unauthorizedTopics = Collections.emptySet[String]
     val internalTopics = getAllTopics(snapshot).filter(Topic.isInternal).asJava
-    new Cluster(clusterId, nodes.values.filter(_ != null).toList.asJava,
-      partitions.toList.asJava,
+    new Cluster(clusterId, nodes.values.filter(_ != null).toBuffer.asJava,
+      partitions.toBuffer.asJava,
       unauthorizedTopics, internalTopics,
       snapshot.controllerId.map(id => node(id)).orNull)
   }
@@ -233,9 +254,10 @@ class MetadataCache(brokerId: Int) extends Logging {
         // move to `AnyRefMap`, which has comparable performance.
         val nodes = new java.util.HashMap[ListenerName, Node]
         val endPoints = new mutable.ArrayBuffer[EndPoint]
-        broker.endPoints.asScala.foreach { ep =>
-          endPoints += EndPoint(ep.host, ep.port, ep.listenerName, ep.securityProtocol)
-          nodes.put(ep.listenerName, new Node(broker.id, ep.host, ep.port))
+        broker.endpoints.asScala.foreach { ep =>
+          val listenerName = new ListenerName(ep.listener)
+          endPoints += new EndPoint(ep.host, ep.port, listenerName, SecurityProtocol.forId(ep.securityProtocol))
+          nodes.put(listenerName, new Node(broker.id, ep.host, ep.port))
         }
         aliveBrokers(broker.id) = Broker(broker.id, endPoints, Option(broker.rack))
         aliveNodes(broker.id) = nodes.asScala
@@ -247,20 +269,21 @@ class MetadataCache(brokerId: Int) extends Logging {
       }
 
       val deletedPartitions = new mutable.ArrayBuffer[TopicPartition]
-      if (updateMetadataRequest.partitionStates().isEmpty) {
+      if (!updateMetadataRequest.partitionStates.iterator.hasNext) {
         metadataSnapshot = MetadataSnapshot(metadataSnapshot.partitionStates, controllerId, aliveBrokers, aliveNodes)
       } else {
         //since kafka may do partial metadata updates, we start by copying the previous state
-        val partitionStates = new mutable.AnyRefMap[String, mutable.LongMap[UpdateMetadataRequest.PartitionState]](metadataSnapshot.partitionStates.size)
+        val partitionStates = new mutable.AnyRefMap[String, mutable.LongMap[UpdateMetadataPartitionState]](metadataSnapshot.partitionStates.size)
         metadataSnapshot.partitionStates.foreach { case (topic, oldPartitionStates) =>
-          val copy = new mutable.LongMap[UpdateMetadataRequest.PartitionState](oldPartitionStates.size)
+          val copy = new mutable.LongMap[UpdateMetadataPartitionState](oldPartitionStates.size)
           copy ++= oldPartitionStates
           partitionStates += (topic -> copy)
         }
-        updateMetadataRequest.partitionStates.asScala.foreach { case (tp, info) =>
+        updateMetadataRequest.partitionStates.asScala.foreach { info =>
           val controllerId = updateMetadataRequest.controllerId
           val controllerEpoch = updateMetadataRequest.controllerEpoch
-          if (info.basePartitionState.leader == LeaderAndIsr.LeaderDuringDelete) {
+          val tp = new TopicPartition(info.topicName, info.partitionIndex)
+          if (info.leader == LeaderAndIsr.LeaderDuringDelete) {
             removePartitionInfo(partitionStates, tp.topic, tp.partition)
             stateChangeLogger.trace(s"Deleted partition $tp from metadata cache in response to UpdateMetadata " +
               s"request sent by controller $controllerId epoch $controllerEpoch with correlation id $correlationId")
@@ -283,7 +306,8 @@ class MetadataCache(brokerId: Int) extends Logging {
 
   def contains(tp: TopicPartition): Boolean = getPartitionInfo(tp.topic, tp.partition).isDefined
 
-  private def removePartitionInfo(partitionStates: mutable.AnyRefMap[String, mutable.LongMap[UpdateMetadataRequest.PartitionState]], topic: String, partitionId: Int): Boolean = {
+  private def removePartitionInfo(partitionStates: mutable.AnyRefMap[String, mutable.LongMap[UpdateMetadataPartitionState]],
+                                  topic: String, partitionId: Int): Boolean = {
     partitionStates.get(topic).exists { infos =>
       infos.remove(partitionId)
       if (infos.isEmpty) partitionStates.remove(topic)
@@ -291,7 +315,7 @@ class MetadataCache(brokerId: Int) extends Logging {
     }
   }
 
-  case class MetadataSnapshot(partitionStates: mutable.AnyRefMap[String, mutable.LongMap[UpdateMetadataRequest.PartitionState]],
+  case class MetadataSnapshot(partitionStates: mutable.AnyRefMap[String, mutable.LongMap[UpdateMetadataPartitionState]],
                               controllerId: Option[Int],
                               aliveBrokers: mutable.LongMap[Broker],
                               aliveNodes: mutable.LongMap[collection.Map[ListenerName, Node]])
diff --git a/core/src/main/scala/kafka/server/NoOpObserver.scala b/core/src/main/scala/kafka/server/NoOpObserver.scala
new file mode 100644
index 0000000000000..3fea54a7c8528
--- /dev/null
+++ b/core/src/main/scala/kafka/server/NoOpObserver.scala
@@ -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 kafka.server
+
+import java.util.Map
+import java.util.concurrent.TimeUnit
+import org.apache.kafka.common.requests.{AbstractRequest, AbstractResponse, ProduceRequest, RequestContext}
+
+/**
+  * An observer implementation that has no operation and serves as a place holder.
+  */
+class NoOpObserver extends Observer {
+
+  def configure(configs: Map[String, _]): Unit = {}
+
+  /**
+    * Observe a request and its corresponding response.
+    */
+  def observe(requestContext: RequestContext, request: AbstractRequest, response: AbstractResponse): Unit = {}
+
+  /**
+    * Observe a produce request
+    */
+  def observeProduceRequest(requestContext: RequestContext, produceRequest: ProduceRequest): Unit = {}
+
+  /**
+    * Close the observer with timeout.
+    */
+  def close(timeout: Long, unit: TimeUnit): Unit = {}
+
+}
diff --git a/core/src/main/scala/kafka/server/Observer.scala b/core/src/main/scala/kafka/server/Observer.scala
new file mode 100644
index 0000000000000..591fa05189ebb
--- /dev/null
+++ b/core/src/main/scala/kafka/server/Observer.scala
@@ -0,0 +1,106 @@
+/**
+  * 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 java.util.concurrent.TimeUnit
+import kafka.network.RequestChannel
+import org.apache.kafka.common.requests.{AbstractRequest, AbstractResponse, ProduceRequest, RequestContext}
+import org.apache.kafka.common.Configurable
+
+/**
+  * Top level interface that all pluggable observer must implement. Kafka will read the 'observer.class.name' config
+  * value at startup time, create an instance of the specificed class using the default constructor, and call its
+  * 'configure' method.
+  *
+  * From that point onwards, every pair of request and response will be routed to the 'record' method.
+  *
+  * If 'observer.class.name' has no value specified or the specified class does not exist, the NoOpObserver
+  * will be used as a place holder.
+  */
+trait Observer extends Configurable {
+
+  /**
+    * Observe a request and its corresponding response
+    *
+    * @param requestContext the context information about the request
+    * @param request  the request being observed for a various purpose(s)
+    * @param response the response to the request
+    */
+  def observe(requestContext: RequestContext, request: AbstractRequest, response: AbstractResponse): Unit
+
+  /**
+    * Observe a produce request. This method handles only the produce request since produce request is special in
+    * two ways. Firstly, if ACK is set to be 0, there is no produce response associated with the produce request.
+    * Secondly, the lifecycle of some inner fields in a ProduceRequest is shorter than the lifecycle of the produce
+    * request itself. That means in some situations, when observe is called on a produce request and
+    * response pair, some fields in the produce request has been null-ed already so that the produce request and
+    * response is not observable (or no useful information). Therefore this method exists for the purpose of allowing
+    * users to observe on the produce request before its corresponding response is created.
+    *
+    * @param requestContext the context information about the request
+    * @param produceRequest  the produce request being observed for a various purpose(s)
+    */
+  def observeProduceRequest(requestContext: RequestContext, produceRequest: ProduceRequest): Unit
+
+  /**
+    * Close the observer with timeout.
+    *
+    * @param timeout the maximum time to wait to close the observer.
+    * @param unit    the time unit.
+    */
+  def close(timeout: Long, unit: TimeUnit): Unit
+}
+
+object Observer {
+
+  /**
+    * Generates a description of the given request and response. It could be used mostly for debugging purpose.
+    *
+    * @param request  the request being described
+    * @param response the response to the request
+    */
+  def describeRequestAndResponse(request: RequestChannel.Request, response: AbstractResponse): String = {
+    var requestDesc = "Request"
+    var responseDesc = "Response"
+    try {
+      if (request == null) {
+        requestDesc += " null"
+      } else {
+        requestDesc += (" header: " + request.header)
+        requestDesc += (" connection ID: " + request.context.connectionId)
+        requestDesc += (" from service with principal: " +
+          request.session.sanitizedUser +
+          " IP address: " + request.session.clientAddress)
+      }
+      requestDesc += " | " // Separate the response description from the request description
+
+      if (response == null) {
+        responseDesc += " null"
+      } else {
+        responseDesc += (if (response.errorCounts == null || response.errorCounts.size == 0) {
+          " with no error"
+        } else {
+          " with errors: " + response.errorCounts
+        })
+      }
+    } catch {
+      case e: Exception => return e.toString // If describing fails, return the exception message directly
+    }
+    requestDesc + responseDesc
+  }
+}
\ No newline at end of file
diff --git a/core/src/main/scala/kafka/server/QuotaFactory.scala b/core/src/main/scala/kafka/server/QuotaFactory.scala
index ed04dcf096fb8..00951c6fabcb8 100644
--- a/core/src/main/scala/kafka/server/QuotaFactory.scala
+++ b/core/src/main/scala/kafka/server/QuotaFactory.scala
@@ -19,6 +19,7 @@ package kafka.server
 import kafka.server.QuotaType._
 import kafka.utils.Logging
 import org.apache.kafka.common.TopicPartition
+import kafka.utils.KafkaScheduler
 import org.apache.kafka.common.metrics.Metrics
 import org.apache.kafka.server.quota.ClientQuotaCallback
 import org.apache.kafka.common.utils.Time
@@ -48,7 +49,7 @@ object QuotaFactory extends Logging {
                            follower: ReplicationQuotaManager,
                            alterLogDirs: ReplicationQuotaManager,
                            clientQuotaCallback: Option[ClientQuotaCallback]) {
-    def shutdown() {
+    def shutdown(): Unit = {
       fetch.shutdown
       produce.shutdown
       request.shutdown
@@ -57,13 +58,16 @@ object QuotaFactory extends Logging {
   }
 
   def instantiate(cfg: KafkaConfig, metrics: Metrics, time: Time, threadNamePrefix: String): QuotaManagers = {
+    instantiate(cfg, metrics, time, None, threadNamePrefix)
+  }
 
+  def instantiate(cfg: KafkaConfig, metrics: Metrics, time: Time, schedulerOpt: Option[KafkaScheduler], threadNamePrefix: String): QuotaManagers = {
     val clientQuotaCallback = Option(cfg.getConfiguredInstance(KafkaConfig.ClientQuotaCallbackClassProp,
       classOf[ClientQuotaCallback]))
     QuotaManagers(
-      new ClientQuotaManager(clientFetchConfig(cfg), metrics, Fetch, time, threadNamePrefix, clientQuotaCallback),
-      new ClientQuotaManager(clientProduceConfig(cfg), metrics, Produce, time, threadNamePrefix, clientQuotaCallback),
-      new ClientRequestQuotaManager(clientRequestConfig(cfg), metrics, time, threadNamePrefix, clientQuotaCallback),
+      new ClientQuotaManager(clientFetchConfig(cfg), metrics, Fetch, time, schedulerOpt, threadNamePrefix, clientQuotaCallback),
+      new ClientQuotaManager(clientProduceConfig(cfg), metrics, Produce, time, schedulerOpt, threadNamePrefix, clientQuotaCallback),
+      new ClientRequestQuotaManager(clientRequestConfig(cfg), metrics, time, schedulerOpt, threadNamePrefix, clientQuotaCallback),
       new ReplicationQuotaManager(replicationConfig(cfg), metrics, LeaderReplication, time),
       new ReplicationQuotaManager(replicationConfig(cfg), metrics, FollowerReplication, time),
       new ReplicationQuotaManager(alterLogDirsReplicationConfig(cfg), metrics, AlterLogDirsReplication, time),
diff --git a/core/src/main/scala/kafka/server/ReplicaAlterLogDirsManager.scala b/core/src/main/scala/kafka/server/ReplicaAlterLogDirsManager.scala
index 1616b84e946eb..85eaeaa89caf3 100644
--- a/core/src/main/scala/kafka/server/ReplicaAlterLogDirsManager.scala
+++ b/core/src/main/scala/kafka/server/ReplicaAlterLogDirsManager.scala
@@ -18,6 +18,7 @@
 package kafka.server
 
 import kafka.cluster.BrokerEndPoint
+import org.apache.kafka.common.TopicPartition
 
 class ReplicaAlterLogDirsManager(brokerConfig: KafkaConfig,
                                  replicaManager: ReplicaManager,
@@ -30,11 +31,26 @@ class ReplicaAlterLogDirsManager(brokerConfig: KafkaConfig,
 
   override def createFetcherThread(fetcherId: Int, sourceBroker: BrokerEndPoint): ReplicaAlterLogDirsThread = {
     val threadName = s"ReplicaAlterLogDirsThread-$fetcherId"
-    new ReplicaAlterLogDirsThread(threadName, sourceBroker, brokerConfig, replicaManager,
+    new ReplicaAlterLogDirsThread(threadName, sourceBroker, brokerConfig, failedPartitions, replicaManager,
       quotaManager, brokerTopicStats)
   }
 
-  def shutdown() {
+  override protected def addPartitionsToFetcherThread(fetcherThread: ReplicaAlterLogDirsThread,
+                                                      initialOffsetAndEpochs: collection.Map[TopicPartition, OffsetAndEpoch]): Unit = {
+    val addedPartitions = fetcherThread.addPartitions(initialOffsetAndEpochs)
+    val (addedInitialOffsets, notAddedInitialOffsets) = initialOffsetAndEpochs.partition { case (tp, _) =>
+      addedPartitions.contains(tp)
+    }
+
+    if (addedInitialOffsets.nonEmpty)
+      info(s"Added log dir fetcher for partitions with initial offsets $addedInitialOffsets")
+
+    if (notAddedInitialOffsets.nonEmpty)
+      info(s"Failed to add log dir fetch for partitions ${notAddedInitialOffsets.keySet} " +
+        s"since the log dir reassignment has already completed")
+  }
+
+  def shutdown(): Unit = {
     info("shutting down")
     closeAllFetchers()
     info("shutdown completed")
diff --git a/core/src/main/scala/kafka/server/ReplicaAlterLogDirsThread.scala b/core/src/main/scala/kafka/server/ReplicaAlterLogDirsThread.scala
index a1f0134fffec0..d0dc9efb969bd 100644
--- a/core/src/main/scala/kafka/server/ReplicaAlterLogDirsThread.scala
+++ b/core/src/main/scala/kafka/server/ReplicaAlterLogDirsThread.scala
@@ -22,7 +22,7 @@ import java.util.Optional
 
 import kafka.api.Request
 import kafka.cluster.BrokerEndPoint
-import kafka.log.{LogAppendInfo, LogOffsetSnapshot}
+import kafka.log.LogAppendInfo
 import kafka.server.AbstractFetcherThread.ResultWithPartitions
 import kafka.server.QuotaFactory.UnboundedQuota
 import org.apache.kafka.common.TopicPartition
@@ -39,12 +39,14 @@ import scala.collection.{Map, Seq, Set, mutable}
 class ReplicaAlterLogDirsThread(name: String,
                                 sourceBroker: BrokerEndPoint,
                                 brokerConfig: KafkaConfig,
+                                failedPartitions: FailedPartitions,
                                 replicaMgr: ReplicaManager,
                                 quota: ReplicationQuotaManager,
                                 brokerTopicStats: BrokerTopicStats)
   extends AbstractFetcherThread(name = name,
                                 clientId = name,
                                 sourceBroker = sourceBroker,
+                                failedPartitions,
                                 fetchBackOffMs = brokerConfig.replicaFetchBackoffMs,
                                 isInterruptible = false) {
 
@@ -54,22 +56,22 @@ class ReplicaAlterLogDirsThread(name: String,
   private var inProgressPartition: Option[TopicPartition] = None
 
   override protected def latestEpoch(topicPartition: TopicPartition): Option[Int] = {
-    replicaMgr.futureLocalReplicaOrException(topicPartition).latestEpoch
+    replicaMgr.futureLocalLogOrException(topicPartition).latestEpoch
   }
 
   override protected def logEndOffset(topicPartition: TopicPartition): Long = {
-    replicaMgr.futureLocalReplicaOrException(topicPartition).logEndOffset
+    replicaMgr.futureLocalLogOrException(topicPartition).logEndOffset
   }
 
   override protected def endOffsetForEpoch(topicPartition: TopicPartition, epoch: Int): Option[OffsetAndEpoch] = {
-    replicaMgr.futureLocalReplicaOrException(topicPartition).endOffsetForEpoch(epoch)
+    replicaMgr.futureLocalLogOrException(topicPartition).endOffsetForEpoch(epoch)
   }
 
   def fetchFromLeader(fetchRequest: FetchRequest.Builder): Seq[(TopicPartition, FetchData)] = {
     var partitionData: Seq[(TopicPartition, FetchResponse.PartitionData[Records])] = null
     val request = fetchRequest.build()
 
-    def processResponseCallback(responsePartitionData: Seq[(TopicPartition, FetchPartitionData)]) {
+    def processResponseCallback(responsePartitionData: Seq[(TopicPartition, FetchPartitionData)]): Unit = {
       partitionData = responsePartitionData.map { case (tp, data) =>
         val abortedTransactions = data.abortedTransactions.map(_.asJava).orNull
         val lastStableOffset = data.lastStableOffset.getOrElse(FetchResponse.INVALID_LAST_STABLE_OFFSET)
@@ -83,11 +85,12 @@ class ReplicaAlterLogDirsThread(name: String,
       Request.FutureLocalReplicaId,
       request.minBytes,
       request.maxBytes,
-      request.version <= 2,
+      false,
       request.fetchData.asScala.toSeq,
       UnboundedQuota,
       processResponseCallback,
-      request.isolationLevel)
+      request.isolationLevel,
+      None)
 
     if (partitionData == null)
       throw new IllegalStateException(s"Failed to fetch data for partitions ${request.fetchData.keySet().toArray.mkString(",")}")
@@ -99,18 +102,21 @@ class ReplicaAlterLogDirsThread(name: String,
   override def processPartitionData(topicPartition: TopicPartition,
                                     fetchOffset: Long,
                                     partitionData: PartitionData[Records]): Option[LogAppendInfo] = {
-    val futureReplica = replicaMgr.futureLocalReplicaOrException(topicPartition)
-    val partition = replicaMgr.getPartition(topicPartition).get
+    val partition = replicaMgr.nonOfflinePartition(topicPartition).get
+    val futureLog = partition.futureLocalLogOrException
     val records = toMemoryRecords(partitionData.records)
 
-    if (fetchOffset != futureReplica.logEndOffset)
+    if (fetchOffset != futureLog.logEndOffset)
       throw new IllegalStateException("Offset mismatch for the future replica %s: fetched offset = %d, log end offset = %d.".format(
-        topicPartition, fetchOffset, futureReplica.logEndOffset))
+        topicPartition, fetchOffset, futureLog.logEndOffset))
 
-    val logAppendInfo = partition.appendRecordsToFollowerOrFutureReplica(records, isFuture = true)
-    val futureReplicaHighWatermark = futureReplica.logEndOffset.min(partitionData.highWatermark)
-    futureReplica.highWatermark = new LogOffsetMetadata(futureReplicaHighWatermark)
-    futureReplica.maybeIncrementLogStartOffset(partitionData.logStartOffset)
+    val logAppendInfo = if (records.sizeInBytes() > 0)
+      partition.appendRecordsToFollowerOrFutureReplica(records, isFuture = true)
+    else
+      None
+
+    futureLog.updateHighWatermark(partitionData.highWatermark)
+    futureLog.maybeIncrementLogStartOffset(partitionData.logStartOffset)
 
     if (partition.maybeReplaceCurrentWithFutureReplica())
       removePartitions(Set(topicPartition))
@@ -119,19 +125,28 @@ class ReplicaAlterLogDirsThread(name: String,
     logAppendInfo
   }
 
-  override protected def fetchEarliestOffsetFromLeader(topicPartition: TopicPartition, leaderEpoch: Int): Long = {
-    val offsetSnapshot = offsetSnapshotFromCurrentReplica(topicPartition, leaderEpoch)
-    offsetSnapshot.logStartOffset
+  override def addPartitions(initialFetchStates: Map[TopicPartition, OffsetAndEpoch]): Set[TopicPartition] = {
+    partitionMapLock.lockInterruptibly()
+    try {
+      // It is possible that the log dir fetcher completed just before this call, so we
+      // filter only the partitions which still have a future log dir.
+      val filteredFetchStates = initialFetchStates.filter { case (tp, _) =>
+        replicaMgr.futureLogExists(tp)
+      }
+      super.addPartitions(filteredFetchStates)
+    } finally {
+      partitionMapLock.unlock()
+    }
   }
 
-  override protected def fetchLatestOffsetFromLeader(topicPartition: TopicPartition, leaderEpoch: Int): Long = {
-    val offsetSnapshot = offsetSnapshotFromCurrentReplica(topicPartition, leaderEpoch)
-    offsetSnapshot.logEndOffset.messageOffset
+  override protected def fetchEarliestOffsetFromLeader(topicPartition: TopicPartition, leaderEpoch: Int): Long = {
+    val partition = replicaMgr.getPartitionOrException(topicPartition, expectLeader = false)
+    partition.localLogOrException.logStartOffset
   }
 
-  private def offsetSnapshotFromCurrentReplica(topicPartition: TopicPartition, leaderEpoch: Int): LogOffsetSnapshot = {
+  override protected def fetchLatestOffsetFromLeader(topicPartition: TopicPartition, leaderEpoch: Int): Long = {
     val partition = replicaMgr.getPartitionOrException(topicPartition, expectLeader = false)
-    partition.fetchOffsetSnapshot(Optional.of[Integer](leaderEpoch), fetchOnlyFromLeader = false)
+    partition.localLogOrException.logEndOffset
   }
 
   /**
@@ -226,7 +241,7 @@ class ReplicaAlterLogDirsThread(name: String,
     val partitionsWithError = mutable.Set[TopicPartition]()
 
     try {
-      val logStartOffset = replicaMgr.futureLocalReplicaOrException(tp).logStartOffset
+      val logStartOffset = replicaMgr.futureLocalLogOrException(tp).logStartOffset
       requestMap.put(tp, new FetchRequest.PartitionData(fetchState.fetchOffset, logStartOffset,
         fetchSize, Optional.of(fetchState.currentLeaderEpoch)))
     } catch {
diff --git a/core/src/main/scala/kafka/server/ReplicaFetcherBlockingSend.scala b/core/src/main/scala/kafka/server/ReplicaFetcherBlockingSend.scala
index 924111c4a862f..ee2dd608c0a4d 100644
--- a/core/src/main/scala/kafka/server/ReplicaFetcherBlockingSend.scala
+++ b/core/src/main/scala/kafka/server/ReplicaFetcherBlockingSend.scala
@@ -26,7 +26,7 @@ import org.apache.kafka.common.requests.AbstractRequest
 import org.apache.kafka.common.security.JaasContext
 import org.apache.kafka.common.utils.{LogContext, Time}
 import org.apache.kafka.clients.{ApiVersions, ClientResponse, ManualMetadataUpdater, NetworkClient}
-import org.apache.kafka.common.Node
+import org.apache.kafka.common.{Node, Reconfigurable}
 import org.apache.kafka.common.requests.AbstractRequest.Builder
 
 import scala.collection.JavaConverters._
@@ -35,9 +35,9 @@ trait BlockingSend {
 
   def sendRequest(requestBuilder: AbstractRequest.Builder[_ <: AbstractRequest]): ClientResponse
 
-  def initiateClose()
+  def initiateClose(): Unit
 
-  def close()
+  def close(): Unit
 }
 
 class ReplicaFetcherBlockingSend(sourceBroker: BrokerEndPoint,
@@ -51,7 +51,7 @@ class ReplicaFetcherBlockingSend(sourceBroker: BrokerEndPoint,
   private val sourceNode = new Node(sourceBroker.id, sourceBroker.host, sourceBroker.port)
   private val socketTimeout: Int = brokerConfig.replicaSocketTimeoutMs
 
-  private val networkClient = {
+  private val (networkClient, reconfigurableChannelBuilder) = {
     val channelBuilder = ChannelBuilders.clientChannelBuilder(
       brokerConfig.interBrokerSecurityProtocol,
       JaasContext.Type.SERVER,
@@ -61,6 +61,12 @@ class ReplicaFetcherBlockingSend(sourceBroker: BrokerEndPoint,
       time,
       brokerConfig.saslInterBrokerHandshakeRequestEnable
     )
+    val reconfigurableChannelBuilder = channelBuilder match {
+      case reconfigurable: Reconfigurable =>
+        brokerConfig.addReconfigurable(reconfigurable)
+        Some(reconfigurable)
+      case _ => None
+    }
     val selector = new Selector(
       NetworkReceive.UNLIMITED,
       brokerConfig.connectionsMaxIdleMs,
@@ -72,7 +78,7 @@ class ReplicaFetcherBlockingSend(sourceBroker: BrokerEndPoint,
       channelBuilder,
       logContext
     )
-    new NetworkClient(
+    val networkClient = new NetworkClient(
       selector,
       new ManualMetadataUpdater(),
       clientId,
@@ -88,6 +94,7 @@ class ReplicaFetcherBlockingSend(sourceBroker: BrokerEndPoint,
       new ApiVersions,
       logContext
     )
+    (networkClient, reconfigurableChannelBuilder)
   }
 
   override def sendRequest(requestBuilder: Builder[_ <: AbstractRequest]): ClientResponse = {
@@ -108,6 +115,7 @@ class ReplicaFetcherBlockingSend(sourceBroker: BrokerEndPoint,
   }
 
   override def initiateClose(): Unit = {
+    reconfigurableChannelBuilder.foreach(brokerConfig.removeReconfigurable)
     networkClient.initiateClose()
   }
 
diff --git a/core/src/main/scala/kafka/server/ReplicaFetcherManager.scala b/core/src/main/scala/kafka/server/ReplicaFetcherManager.scala
index fa902b9d9efb0..d547e1b5d769b 100644
--- a/core/src/main/scala/kafka/server/ReplicaFetcherManager.scala
+++ b/core/src/main/scala/kafka/server/ReplicaFetcherManager.scala
@@ -35,11 +35,11 @@ class ReplicaFetcherManager(brokerConfig: KafkaConfig,
   override def createFetcherThread(fetcherId: Int, sourceBroker: BrokerEndPoint): ReplicaFetcherThread = {
     val prefix = threadNamePrefix.map(tp => s"$tp:").getOrElse("")
     val threadName = s"${prefix}ReplicaFetcherThread-$fetcherId-${sourceBroker.id}"
-    new ReplicaFetcherThread(threadName, fetcherId, sourceBroker, brokerConfig, replicaManager,
+    new ReplicaFetcherThread(threadName, fetcherId, sourceBroker, brokerConfig, failedPartitions, replicaManager,
       metrics, time, quotaManager)
   }
 
-  def shutdown() {
+  def shutdown(): Unit = {
     info("shutting down")
     closeAllFetchers()
     info("shutdown completed")
diff --git a/core/src/main/scala/kafka/server/ReplicaFetcherThread.scala b/core/src/main/scala/kafka/server/ReplicaFetcherThread.scala
index 3e22653ff9b88..1d2fdeed2d65c 100644
--- a/core/src/main/scala/kafka/server/ReplicaFetcherThread.scala
+++ b/core/src/main/scala/kafka/server/ReplicaFetcherThread.scala
@@ -40,6 +40,7 @@ class ReplicaFetcherThread(name: String,
                            fetcherId: Int,
                            sourceBroker: BrokerEndPoint,
                            brokerConfig: KafkaConfig,
+                           failedPartitions: FailedPartitions,
                            replicaMgr: ReplicaManager,
                            metrics: Metrics,
                            time: Time,
@@ -48,6 +49,7 @@ class ReplicaFetcherThread(name: String,
   extends AbstractFetcherThread(name = name,
                                 clientId = name,
                                 sourceBroker = sourceBroker,
+                                failedPartitions,
                                 fetchBackOffMs = brokerConfig.replicaFetchBackoffMs,
                                 isInterruptible = false) {
 
@@ -62,7 +64,8 @@ class ReplicaFetcherThread(name: String,
 
   // Visible for testing
   private[server] val fetchRequestVersion: Short =
-    if (brokerConfig.interBrokerProtocolVersion >= KAFKA_2_1_IV2) 10
+    if (brokerConfig.interBrokerProtocolVersion >= KAFKA_2_3_IV1) 11
+    else if (brokerConfig.interBrokerProtocolVersion >= KAFKA_2_1_IV2) 10
     else if (brokerConfig.interBrokerProtocolVersion >= KAFKA_2_0_IV1) 8
     else if (brokerConfig.interBrokerProtocolVersion >= KAFKA_1_1_IV0) 7
     else if (brokerConfig.interBrokerProtocolVersion >= KAFKA_0_11_0_IV1) 5
@@ -74,7 +77,8 @@ class ReplicaFetcherThread(name: String,
 
   // Visible for testing
   private[server] val offsetForLeaderEpochRequestVersion: Short =
-    if (brokerConfig.interBrokerProtocolVersion >= KAFKA_2_1_IV1) 2
+    if (brokerConfig.interBrokerProtocolVersion >= KAFKA_2_3_IV1) 3
+    else if (brokerConfig.interBrokerProtocolVersion >= KAFKA_2_1_IV1) 2
     else if (brokerConfig.interBrokerProtocolVersion >= KAFKA_2_0_IV0) 1
     else 0
 
@@ -95,15 +99,15 @@ class ReplicaFetcherThread(name: String,
   private val fetchSessionHandler = new FetchSessionHandler(logContext, sourceBroker.id)
 
   override protected def latestEpoch(topicPartition: TopicPartition): Option[Int] = {
-    replicaMgr.localReplicaOrException(topicPartition).latestEpoch
+    replicaMgr.localLogOrException(topicPartition).latestEpoch
   }
 
   override protected def logEndOffset(topicPartition: TopicPartition): Long = {
-    replicaMgr.localReplicaOrException(topicPartition).logEndOffset
+    replicaMgr.localLogOrException(topicPartition).logEndOffset
   }
 
   override protected def endOffsetForEpoch(topicPartition: TopicPartition, epoch: Int): Option[OffsetAndEpoch] = {
-    replicaMgr.localReplicaOrException(topicPartition).endOffsetForEpoch(epoch)
+    replicaMgr.localLogOrException(topicPartition).endOffsetForEpoch(epoch)
   }
 
   override def initiateShutdown(): Boolean = {
@@ -139,33 +143,32 @@ class ReplicaFetcherThread(name: String,
   override def processPartitionData(topicPartition: TopicPartition,
                                     fetchOffset: Long,
                                     partitionData: FetchData): Option[LogAppendInfo] = {
-    val replica = replicaMgr.localReplicaOrException(topicPartition)
-    val partition = replicaMgr.getPartition(topicPartition).get
+    val partition = replicaMgr.nonOfflinePartition(topicPartition).get
+    val log = partition.localLogOrException
     val records = toMemoryRecords(partitionData.records)
 
     maybeWarnIfOversizedRecords(records, topicPartition)
 
-    if (fetchOffset != replica.logEndOffset)
+    if (fetchOffset != log.logEndOffset)
       throw new IllegalStateException("Offset mismatch for partition %s: fetched offset = %d, log end offset = %d.".format(
-        topicPartition, fetchOffset, replica.logEndOffset))
+        topicPartition, fetchOffset, log.logEndOffset))
 
     if (isTraceEnabled)
       trace("Follower has replica log end offset %d for partition %s. Received %d messages and leader hw %d"
-        .format(replica.logEndOffset, topicPartition, records.sizeInBytes, partitionData.highWatermark))
+        .format(log.logEndOffset, topicPartition, records.sizeInBytes, partitionData.highWatermark))
 
     // Append the leader's messages to the log
     val logAppendInfo = partition.appendRecordsToFollowerOrFutureReplica(records, isFuture = false)
 
     if (isTraceEnabled)
       trace("Follower has replica log end offset %d after appending %d bytes of messages for partition %s"
-        .format(replica.logEndOffset, records.sizeInBytes, topicPartition))
-    val followerHighWatermark = replica.logEndOffset.min(partitionData.highWatermark)
+        .format(log.logEndOffset, records.sizeInBytes, topicPartition))
     val leaderLogStartOffset = partitionData.logStartOffset
-    // for the follower replica, we do not need to keep
-    // its segment base offset the physical position,
-    // these values will be computed upon making the leader
-    replica.highWatermark = new LogOffsetMetadata(followerHighWatermark)
-    replica.maybeIncrementLogStartOffset(leaderLogStartOffset)
+
+    // For the follower replica, we do not need to keep its segment base offset and physical position.
+    // These values will be computed upon becoming leader or handling a preferred read replica fetch.
+    val followerHighWatermark = log.updateHighWatermark(partitionData.highWatermark)
+    log.maybeIncrementLogStartOffset(leaderLogStartOffset)
     if (isTraceEnabled)
       trace(s"Follower set replica high watermark for partition $topicPartition to $followerHighWatermark")
 
@@ -241,7 +244,7 @@ class ReplicaFetcherThread(name: String,
       // We will not include a replica in the fetch request if it should be throttled.
       if (fetchState.isReadyForFetch && !shouldFollowerThrottle(quota, topicPartition)) {
         try {
-          val logStartOffset = replicaMgr.localReplicaOrException(topicPartition).logStartOffset
+          val logStartOffset = replicaMgr.localLogOrException(topicPartition).logStartOffset
           builder.add(topicPartition, new FetchRequest.PartitionData(
             fetchState.fetchOffset, logStartOffset, fetchSize, Optional.of(fetchState.currentLeaderEpoch)))
         } catch {
@@ -273,13 +276,14 @@ class ReplicaFetcherThread(name: String,
    * The logic for finding the truncation offset is implemented in AbstractFetcherThread.getOffsetTruncationState
    */
   override def truncate(tp: TopicPartition, offsetTruncationState: OffsetTruncationState): Unit = {
-    val replica = replicaMgr.localReplicaOrException(tp)
-    val partition = replicaMgr.getPartition(tp).get
+    val partition = replicaMgr.nonOfflinePartition(tp).get
+    val log = partition.localLogOrException
+
     partition.truncateTo(offsetTruncationState.offset, isFuture = false)
 
-    if (offsetTruncationState.offset < replica.highWatermark.messageOffset)
+    if (offsetTruncationState.offset < log.highWatermark)
       warn(s"Truncating $tp to offset ${offsetTruncationState.offset} below high watermark " +
-        s"${replica.highWatermark.messageOffset}")
+        s"${log.highWatermark}")
 
     // mark the future replica for truncation only when we do last truncation
     if (offsetTruncationState.truncationCompleted)
@@ -288,7 +292,7 @@ class ReplicaFetcherThread(name: String,
   }
 
   override protected def truncateFullyAndStartAt(topicPartition: TopicPartition, offset: Long): Unit = {
-    val partition = replicaMgr.getPartition(topicPartition).get
+    val partition = replicaMgr.nonOfflinePartition(topicPartition).get
     partition.truncateFullyAndStartAt(offset, isFuture = false)
   }
 
@@ -299,7 +303,7 @@ class ReplicaFetcherThread(name: String,
       return Map.empty
     }
 
-    val epochRequest = new OffsetsForLeaderEpochRequest.Builder(offsetForLeaderEpochRequestVersion, partitions.asJava)
+    val epochRequest = OffsetsForLeaderEpochRequest.Builder.forFollower(offsetForLeaderEpochRequestVersion, partitions.asJava, brokerConfig.brokerId)
     debug(s"Sending offset for leader epoch request $epochRequest")
 
     try {
diff --git a/core/src/main/scala/kafka/server/ReplicaManager.scala b/core/src/main/scala/kafka/server/ReplicaManager.scala
index c194bd56abfab..5e0c994c94f9c 100644
--- a/core/src/main/scala/kafka/server/ReplicaManager.scala
+++ b/core/src/main/scala/kafka/server/ReplicaManager.scala
@@ -22,33 +22,42 @@ import java.util.concurrent.TimeUnit
 import java.util.concurrent.atomic.{AtomicBoolean, AtomicLong}
 import java.util.concurrent.locks.Lock
 
-import com.yammer.metrics.core.Gauge
+import com.yammer.metrics.core.{Gauge, Meter}
 import kafka.api._
-import kafka.cluster.{BrokerEndPoint, Partition, Replica}
+import kafka.cluster.{BrokerEndPoint, Partition}
+import kafka.common.RecordValidationException
 import kafka.controller.{KafkaController, StateChangeLogger}
 import kafka.log._
 import kafka.metrics.KafkaMetricsGroup
-import kafka.server.QuotaFactory.{QuotaManagers, UnboundedQuota}
-import kafka.server.checkpoints.OffsetCheckpointFile
+import kafka.server.QuotaFactory.QuotaManagers
+import kafka.server.checkpoints.{LazyOffsetCheckpoints, OffsetCheckpointFile, OffsetCheckpoints}
 import kafka.utils._
 import kafka.zk.KafkaZkClient
+import org.apache.kafka.common.{ElectionType, Node, TopicPartition}
 import org.apache.kafka.common.errors._
 import org.apache.kafka.common.internals.Topic
+import org.apache.kafka.common.message.LeaderAndIsrRequestData.LeaderAndIsrPartitionState
+import org.apache.kafka.common.message.LeaderAndIsrResponseData
+import org.apache.kafka.common.message.LeaderAndIsrResponseData.LeaderAndIsrPartitionError
 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.record.FileRecords.TimestampAndOffset
 import org.apache.kafka.common.record._
-import org.apache.kafka.common.requests.FetchResponse.AbortedTransaction
+import org.apache.kafka.common.replica.PartitionView.DefaultPartitionView
 import org.apache.kafka.common.requests.DescribeLogDirsResponse.{LogDirInfo, ReplicaInfo}
 import org.apache.kafka.common.requests.EpochEndOffset._
 import org.apache.kafka.common.requests.FetchRequest.PartitionData
+import org.apache.kafka.common.requests.FetchResponse.AbortedTransaction
 import org.apache.kafka.common.requests.ProduceResponse.PartitionResponse
-import org.apache.kafka.common.requests._
+import org.apache.kafka.common.requests.{ApiError, DeleteRecordsResponse, DescribeLogDirsResponse, EpochEndOffset, IsolationLevel, LeaderAndIsrRequest, LeaderAndIsrResponse, OffsetsForLeaderEpochRequest, StopReplicaRequest, UpdateMetadataRequest}
 import org.apache.kafka.common.utils.Time
-import org.apache.kafka.common.TopicPartition
-import org.apache.kafka.common.record.FileRecords.TimestampAndOffset
+import org.apache.kafka.common.replica.ReplicaView.DefaultReplicaView
+import org.apache.kafka.common.replica.{ClientMetadata, _}
 
+import scala.compat.java8.OptionConverters._
 import scala.collection.JavaConverters._
-import scala.collection._
+import scala.collection.{Map, Seq, Set, mutable}
 
 /*
  * Result metadata of a log append operation on the log
@@ -74,7 +83,8 @@ case class LogDeleteRecordsResult(requestedOffset: Long, lowWatermark: Long, exc
  * @param readSize amount of data that was read from the log i.e. size of the fetch
  * @param isReadFromLogEnd true if the request read up to the log end offset snapshot
  *                         when the read was initiated, false otherwise
- * @param error Exception if error encountered while reading from the log
+ * @param preferredReadReplica the preferred read replica to be used for future fetches
+ * @param exception Exception if error encountered while reading from the log
  */
 case class LogReadResult(info: FetchDataInfo,
                          highWatermark: Long,
@@ -84,6 +94,8 @@ case class LogReadResult(info: FetchDataInfo,
                          fetchTimeMs: Long,
                          readSize: Int,
                          lastStableOffset: Option[Long],
+                         preferredReadReplica: Option[Int] = None,
+                         followerNeedsHwUpdate: Boolean = false,
                          exception: Option[Throwable] = None) {
 
   def error: Errors = exception match {
@@ -105,32 +117,37 @@ case class FetchPartitionData(error: Errors = Errors.NONE,
                               logStartOffset: Long,
                               records: Records,
                               lastStableOffset: Option[Long],
-                              abortedTransactions: Option[List[AbortedTransaction]])
-
-object LogReadResult {
-  val UnknownLogReadResult = LogReadResult(info = FetchDataInfo(LogOffsetMetadata.UnknownOffsetMetadata, MemoryRecords.EMPTY),
-                                           highWatermark = -1L,
-                                           leaderLogStartOffset = -1L,
-                                           leaderLogEndOffset = -1L,
-                                           followerLogStartOffset = -1L,
-                                           fetchTimeMs = -1L,
-                                           readSize = -1,
-                                           lastStableOffset = None)
+                              abortedTransactions: Option[List[AbortedTransaction]],
+                              preferredReadReplica: Option[Int])
+
+
+/**
+ * Trait to represent the state of hosted partitions. We create a concrete (active) Partition
+ * instance when the broker receives a LeaderAndIsr request from the controller indicating
+ * that it should be either a leader or follower of a partition.
+ */
+sealed trait HostedPartition
+object HostedPartition {
+  /**
+   * This broker does not have any state for this partition locally.
+   */
+  final object None extends HostedPartition
+
+  /**
+   * This broker hosts the partition and it is online.
+   */
+  final case class Online(partition: Partition) extends HostedPartition
+
+  /**
+   * This broker hosts the partition, but it is in an offline log directory.
+   */
+  final object Offline extends HostedPartition
 }
 
 object ReplicaManager {
   val HighWatermarkFilename = "replication-offset-checkpoint"
   val IsrChangePropagationBlackOut = 5000L
   val IsrChangePropagationInterval = 60000L
-  val OfflinePartition: Partition = new Partition(new TopicPartition("", -1),
-    isOffline = true,
-    replicaLagTimeMaxMs = 0L,
-    interBrokerProtocolVersion = ApiVersion.latestVersion,
-    localBrokerId = -1,
-    time = null,
-    replicaManager = null,
-    logManager = null,
-    zkClient = null)
 }
 
 class ReplicaManager(val config: KafkaConfig,
@@ -147,7 +164,7 @@ class ReplicaManager(val config: KafkaConfig,
                      val delayedProducePurgatory: DelayedOperationPurgatory[DelayedProduce],
                      val delayedFetchPurgatory: DelayedOperationPurgatory[DelayedFetch],
                      val delayedDeleteRecordsPurgatory: DelayedOperationPurgatory[DelayedDeleteRecords],
-                     val delayedElectPreferredLeaderPurgatory: DelayedOperationPurgatory[DelayedElectPreferredLeader],
+                     val delayedElectLeaderPurgatory: DelayedOperationPurgatory[DelayedElectLeader],
                      threadNamePrefix: Option[String]) extends Logging with KafkaMetricsGroup {
 
   def this(config: KafkaConfig,
@@ -173,35 +190,37 @@ class ReplicaManager(val config: KafkaConfig,
       DelayedOperationPurgatory[DelayedDeleteRecords](
         purgatoryName = "DeleteRecords", brokerId = config.brokerId,
         purgeInterval = config.deleteRecordsPurgatoryPurgeIntervalRequests),
-      DelayedOperationPurgatory[DelayedElectPreferredLeader](
-        purgatoryName = "ElectPreferredLeader", brokerId = config.brokerId),
+      DelayedOperationPurgatory[DelayedElectLeader](
+        purgatoryName = "ElectLeader", brokerId = config.brokerId),
       threadNamePrefix)
   }
 
   /* epoch of the controller that last changed the leader */
   @volatile var controllerEpoch: Int = KafkaController.InitialControllerEpoch
   private val localBrokerId = config.brokerId
-  private val allPartitions = new Pool[TopicPartition, Partition](valueFactory = Some(tp =>
-    Partition(tp, time, this)))
+  private val allPartitions = new Pool[TopicPartition, HostedPartition](
+    valueFactory = Some(tp => HostedPartition.Online(Partition(tp, time, this)))
+  )
   private val replicaStateChangeLock = new Object
   val replicaFetcherManager = createReplicaFetcherManager(metrics, time, threadNamePrefix, quotaManagers.follower)
   val replicaAlterLogDirsManager = createReplicaAlterLogDirsManager(quotaManagers.alterLogDirs, brokerTopicStats)
   private val highWatermarkCheckPointThreadStarted = new AtomicBoolean(false)
-  @volatile var highWatermarkCheckpoints = logManager.liveLogDirs.map(dir =>
+  @volatile var highWatermarkCheckpoints: Map[String, OffsetCheckpointFile] = logManager.liveLogDirs.map(dir =>
     (dir.getAbsolutePath, new OffsetCheckpointFile(new File(dir, ReplicaManager.HighWatermarkFilename), logDirFailureChannel))).toMap
 
-  private var hwThreadInitialized = false
   this.logIdent = s"[ReplicaManager broker=$localBrokerId] "
   private val stateChangeLogger = new StateChangeLogger(localBrokerId, inControllerContext = false, None)
 
+  private val isrChangeNotificationBatchSize = 3000
   private val isrChangeSet: mutable.Set[TopicPartition] = new mutable.HashSet[TopicPartition]()
   private val lastIsrChangeMs = new AtomicLong(System.currentTimeMillis())
   private val lastIsrPropagationMs = new AtomicLong(System.currentTimeMillis())
 
   private var logDirFailureHandler: LogDirFailureHandler = null
+  val recompressedBatchCount: AtomicLong = new AtomicLong(0)
 
   private class LogDirFailureHandler(name: String, haltBrokerOnDirFailure: Boolean) extends ShutdownableThread(name) {
-    override def doWork() {
+    override def doWork(): Unit = {
       val newOfflineLogDir = logDirFailureChannel.takeNextOfflineLogDir()
       if (haltBrokerOnDirFailure) {
         fatal(s"Halting broker because dir $newOfflineLogDir is offline")
@@ -211,6 +230,8 @@ class ReplicaManager(val config: KafkaConfig,
     }
   }
 
+  val replicaSelectorOpt: Option[ReplicaSelector] = createReplicaSelector()
+
   val leaderCount = newGauge(
     "LeaderCount",
     new Gauge[Int] {
@@ -226,7 +247,7 @@ class ReplicaManager(val config: KafkaConfig,
   val offlineReplicaCount = newGauge(
     "OfflineReplicaCount",
     new Gauge[Int] {
-      def value = offlinePartitionsIterator.size
+      def value = offlinePartitionCount
     }
   )
   val underReplicatedPartitions = newGauge(
@@ -241,38 +262,67 @@ class ReplicaManager(val config: KafkaConfig,
       def value = leaderPartitionsIterator.count(_.isUnderMinIsr)
     }
   )
+  val atMinIsrPartitionCount = newGauge(
+    "AtMinIsrPartitionCount",
+    new Gauge[Int] {
+      def value = leaderPartitionsIterator.count(_.isAtMinIsr)
+    }
+  )
+  val oneAboveMinIsrPartitionCount = newGauge(
+    "OneAboveMinIsrPartitionCount",
+    new Gauge[Int] {
+      def value = leaderPartitionsIterator.count(_.isOneAboveMinIsr)
+    }
+  )
+
+  val isrExpandRate: Meter = newMeter("IsrExpandsPerSec", "expands", TimeUnit.SECONDS)
+  val isrShrinkRate: Meter = newMeter("IsrShrinksPerSec", "shrinks", TimeUnit.SECONDS)
+  val failedIsrUpdatesRate: Meter = newMeter("FailedIsrUpdatesPerSec", "failedUpdates", TimeUnit.SECONDS)
+  val recompressionCount = newGauge(
+    "recompressionCount",
+    new Gauge[Long] {
+      def value = recompressedBatchCount.get()
+    }
+  )
 
-  val isrExpandRate = newMeter("IsrExpandsPerSec", "expands", TimeUnit.SECONDS)
-  val isrShrinkRate = newMeter("IsrShrinksPerSec", "shrinks", TimeUnit.SECONDS)
-  val failedIsrUpdatesRate = newMeter("FailedIsrUpdatesPerSec", "failedUpdates", TimeUnit.SECONDS)
 
   def underReplicatedPartitionCount: Int = leaderPartitionsIterator.count(_.isUnderReplicated)
 
-  def startHighWaterMarksCheckPointThread() = {
-    if(highWatermarkCheckPointThreadStarted.compareAndSet(false, true))
+  def startHighWatermarkCheckPointThread() = {
+    if (highWatermarkCheckPointThreadStarted.compareAndSet(false, true))
       scheduler.schedule("highwatermark-checkpoint", checkpointHighWatermarks _, period = config.replicaHighWatermarkCheckpointIntervalMs, unit = TimeUnit.MILLISECONDS)
   }
 
-  def recordIsrChange(topicPartition: TopicPartition) {
+  def recordIsrChange(topicPartition: TopicPartition): Unit = {
     isrChangeSet synchronized {
       isrChangeSet += topicPartition
-      lastIsrChangeMs.set(System.currentTimeMillis())
+      lastIsrChangeMs.set(time.milliseconds())
     }
   }
+
+  def incrementRecompressionCount() {
+    recompressedBatchCount.incrementAndGet()
+  }
+
   /**
    * This function periodically runs to see if ISR needs to be propagated. It propagates ISR when:
    * 1. There is ISR change not propagated yet.
    * 2. There is no ISR Change in the last five seconds, or it has been more than 60 seconds since the last ISR propagation.
    * This allows an occasional ISR change to be propagated within a few seconds, and avoids overwhelming controller and
    * other brokers when large amount of ISR change occurs.
+   *
+   * ISR changes are batched into chunks of isrChangeNotificationBatchSize TopicPartitions to avoid writing ZNodes which
+   * exceed the default jute.maxbuffer size of 1MB in ZooKeeper. With this batch size, the worst case data size is about
+   * 820k, which gives about 203k of head room for the rest of the ZooKeeper CreateRequest. Since we have no idea what
+   * ACLs or other metadata could be included we need to leave lots of extra room. See KAFKA-6469.
    */
-  def maybePropagateIsrChanges() {
-    val now = System.currentTimeMillis()
+  def maybePropagateIsrChanges(): Unit = {
+    val now = time.milliseconds()
     isrChangeSet synchronized {
       if (isrChangeSet.nonEmpty &&
         (lastIsrChangeMs.get() + ReplicaManager.IsrChangePropagationBlackOut < now ||
           lastIsrPropagationMs.get() + ReplicaManager.IsrChangePropagationInterval < now)) {
-        zkClient.propagateIsrChanges(isrChangeSet)
+        isrChangeSet.grouped(isrChangeNotificationBatchSize).foreach(zkClient.propagateIsrChanges)
         isrChangeSet.clear()
         lastIsrPropagationMs.set(now)
       }
@@ -289,47 +339,14 @@ class ReplicaManager(val config: KafkaConfig,
 
   def getLog(topicPartition: TopicPartition): Option[Log] = logManager.getLog(topicPartition)
 
-  /**
-   * Try to complete some delayed produce requests with the request key;
-   * this can be triggered when:
-   *
-   * 1. The partition HW has changed (for acks = -1)
-   * 2. A follower replica's fetch operation is received (for acks > 1)
-   */
-  def tryCompleteDelayedProduce(key: DelayedOperationKey) {
-    val completed = delayedProducePurgatory.checkAndComplete(key)
-    debug("Request key %s unblocked %d producer requests.".format(key.keyLabel, completed))
-  }
-
-  /**
-   * Try to complete some delayed fetch requests with the request key;
-   * this can be triggered when:
-   *
-   * 1. The partition HW has changed (for regular fetch)
-   * 2. A new message set is appended to the local log (for follower fetch)
-   */
-  def tryCompleteDelayedFetch(key: DelayedOperationKey) {
-    val completed = delayedFetchPurgatory.checkAndComplete(key)
-    debug("Request key %s unblocked %d fetch requests.".format(key.keyLabel, completed))
-  }
-
-  /**
-   * Try to complete some delayed DeleteRecordsRequest with the request key;
-   * this needs to be triggered when the partition low watermark has changed
-   */
-  def tryCompleteDelayedDeleteRecords(key: DelayedOperationKey) {
-    val completed = delayedDeleteRecordsPurgatory.checkAndComplete(key)
-    debug("Request key %s unblocked %d DeleteRecordsRequest.".format(key.keyLabel, completed))
-  }
-
-  def hasDelayedElectionOperations = delayedElectPreferredLeaderPurgatory.delayed != 0
+  def hasDelayedElectionOperations: Boolean = delayedElectLeaderPurgatory.numDelayed != 0
 
   def tryCompleteElection(key: DelayedOperationKey): Unit = {
-    val completed = delayedElectPreferredLeaderPurgatory.checkAndComplete(key)
-    debug("Request key %s unblocked %d ElectPreferredLeader.".format(key.keyLabel, completed))
+    val completed = delayedElectLeaderPurgatory.checkAndComplete(key)
+    debug("Request key %s unblocked %d ElectLeader.".format(key.keyLabel, completed))
   }
 
-  def startup() {
+  def startup(): Unit = {
     // start ISR expiration thread
     // A follower can lag behind leader for up to config.replicaLagTimeMaxMs x 1.5 before it is removed from ISR
     scheduler.schedule("isr-expiration", maybeShrinkIsr _, period = config.replicaLagTimeMaxMs / 2, unit = TimeUnit.MILLISECONDS)
@@ -344,24 +361,33 @@ class ReplicaManager(val config: KafkaConfig,
     logDirFailureHandler.start()
   }
 
+  private def maybeRemoveTopicMetrics(topic: String): Unit = {
+    val topicHasOnlinePartition = allPartitions.values.exists {
+      case HostedPartition.Online(partition) => topic == partition.topic
+      case HostedPartition.None | HostedPartition.Offline => false
+    }
+    if (!topicHasOnlinePartition)
+      brokerTopicStats.removeMetrics(topic)
+  }
+
   def stopReplica(topicPartition: TopicPartition, deletePartition: Boolean)  = {
     stateChangeLogger.trace(s"Handling stop replica (delete=$deletePartition) for partition $topicPartition")
 
     if (deletePartition) {
-      val removedPartition = allPartitions.remove(topicPartition)
-      if (removedPartition eq ReplicaManager.OfflinePartition) {
-        allPartitions.put(topicPartition, ReplicaManager.OfflinePartition)
-        throw new KafkaStorageException(s"Partition $topicPartition is on an offline disk")
-      }
+      getPartition(topicPartition) match {
+        case HostedPartition.Offline =>
+          throw new KafkaStorageException(s"Partition $topicPartition is on an offline disk")
+
+        case hostedPartition @ HostedPartition.Online(removedPartition) =>
+          if (allPartitions.remove(topicPartition, hostedPartition)) {
+            maybeRemoveTopicMetrics(topicPartition.topic)
+            // this will delete the local log. This call may throw exception if the log is on offline directory
+            removedPartition.delete()
+          }
 
-      if (removedPartition != null) {
-        val topicHasPartitions = allPartitions.values.exists(partition => topicPartition.topic == partition.topic)
-        if (!topicHasPartitions)
-          brokerTopicStats.removeMetrics(topicPartition.topic)
-        // this will delete the local log. This call may throw exception if the log is on offline directory
-        removedPartition.delete()
-      } else {
-        stateChangeLogger.trace(s"Ignoring stop replica (delete=$deletePartition) for partition $topicPartition as replica doesn't exist on broker")
+        case HostedPartition.None =>
+          stateChangeLogger.trace(s"Ignoring stop replica (delete=$deletePartition) for partition " +
+            s"$topicPartition as replica doesn't exist on broker")
       }
 
       // Delete log and corresponding folders in case replica manager doesn't hold them anymore.
@@ -371,9 +397,20 @@ class ReplicaManager(val config: KafkaConfig,
       if (logManager.getLog(topicPartition, isFuture = true).isDefined)
         logManager.asyncDelete(topicPartition, isFuture = true)
     }
+
+    // If we were the leader, we may have some operations still waiting for completion.
+    // We force completion to prevent them from timing out.
+    completeDelayedFetchOrProduceRequests(topicPartition)
+
     stateChangeLogger.trace(s"Finished handling stop replica (delete=$deletePartition) for partition $topicPartition")
   }
 
+  private def completeDelayedFetchOrProduceRequests(topicPartition: TopicPartition): Unit = {
+    val topicPartitionOperationKey = TopicPartitionOperationKey(topicPartition)
+    delayedProducePurgatory.checkAndComplete(topicPartitionOperationKey)
+    delayedFetchPurgatory.checkAndComplete(topicPartitionOperationKey)
+  }
+
   def stopReplicas(stopReplicaRequest: StopReplicaRequest): (mutable.Map[TopicPartition, Errors], Errors) = {
     replicaStateChangeLock synchronized {
       val responseMap = new collection.mutable.HashMap[TopicPartition, Errors]
@@ -382,7 +419,7 @@ class ReplicaManager(val config: KafkaConfig,
           s"${stopReplicaRequest.controllerEpoch}. Latest known controller epoch is $controllerEpoch")
         (responseMap, Errors.STALE_CONTROLLER_EPOCH)
       } else {
-        val partitions = stopReplicaRequest.partitions.asScala
+        val partitions = stopReplicaRequest.partitions.asScala.toSet
         controllerEpoch = stopReplicaRequest.controllerEpoch
         // First stop fetchers for all partitions, then stop the corresponding replicas
         replicaFetcherManager.removeFetcherForPartitions(partitions)
@@ -403,68 +440,91 @@ class ReplicaManager(val config: KafkaConfig,
     }
   }
 
-  def getOrCreatePartition(topicPartition: TopicPartition): Partition =
-    allPartitions.getAndMaybePut(topicPartition)
+  def getPartition(topicPartition: TopicPartition): HostedPartition = {
+    Option(allPartitions.get(topicPartition)).getOrElse(HostedPartition.None)
+  }
 
-  def getPartition(topicPartition: TopicPartition): Option[Partition] =
-    Option(allPartitions.get(topicPartition))
+  // Visible for testing
+  def createPartition(topicPartition: TopicPartition): Partition = {
+    val partition = Partition(topicPartition, time, this)
+    allPartitions.put(topicPartition, HostedPartition.Online(partition))
+    partition
+  }
 
-  def nonOfflinePartition(topicPartition: TopicPartition): Option[Partition] =
-    getPartition(topicPartition).filter(_ ne ReplicaManager.OfflinePartition)
+  def nonOfflinePartition(topicPartition: TopicPartition): Option[Partition] = {
+    getPartition(topicPartition) match {
+      case HostedPartition.Online(partition) => Some(partition)
+      case HostedPartition.None | HostedPartition.Offline => None
+    }
+  }
 
   // An iterator over all non offline partitions. This is a weakly consistent iterator; a partition made offline after
   // the iterator has been constructed could still be returned by this iterator.
-  private def nonOfflinePartitionsIterator: Iterator[Partition] =
-    allPartitions.values.iterator.filter(_ ne ReplicaManager.OfflinePartition)
-
-  // An iterator over all offline partitions. This is a weakly consistent iterator; a partition made offline after the
-  // iterator has been constructed may not be visible.
-  private def offlinePartitionsIterator: Iterator[Partition] =
-    allPartitions.values.iterator.filter(_ eq ReplicaManager.OfflinePartition)
+  private def nonOfflinePartitionsIterator: Iterator[Partition] = {
+    allPartitions.values.iterator.flatMap {
+      case HostedPartition.Online(partition) => Some(partition)
+      case HostedPartition.None | HostedPartition.Offline => None
+    }
+  }
 
+  private def offlinePartitionCount: Int = {
+    allPartitions.values.iterator.count(_ == HostedPartition.Offline)
+  }
 
   def getPartitionOrException(topicPartition: TopicPartition, expectLeader: Boolean): Partition = {
+    getPartitionOrError(topicPartition, expectLeader) match {
+      case Left(Errors.KAFKA_STORAGE_ERROR) =>
+        throw new KafkaStorageException(s"Partition $topicPartition is in an offline log directory")
+
+      case Left(error) =>
+        throw error.exception(s"Error while fetching partition state for $topicPartition")
+
+      case Right(partition) => partition
+    }
+  }
+
+  def getPartitionOrError(topicPartition: TopicPartition, expectLeader: Boolean): Either[Errors, Partition] = {
     getPartition(topicPartition) match {
-      case Some(partition) =>
-        if (partition eq ReplicaManager.OfflinePartition)
-          throw new KafkaStorageException(s"Partition $topicPartition is in an offline log directory")
-        else
-          partition
+      case HostedPartition.Online(partition) =>
+        Right(partition)
 
-      case None if metadataCache.contains(topicPartition) =>
+      case HostedPartition.Offline =>
+        Left(Errors.KAFKA_STORAGE_ERROR)
+
+      case HostedPartition.None if metadataCache.contains(topicPartition) =>
         if (expectLeader) {
           // The topic exists, but this broker is no longer a replica of it, so we return NOT_LEADER which
           // forces clients to refresh metadata to find the new location. This can happen, for example,
           // during a partition reassignment if a produce request from the client is sent to a broker after
           // the local replica has been deleted.
-          throw new NotLeaderForPartitionException(s"Broker $localBrokerId is not a replica of $topicPartition")
+          Left(Errors.NOT_LEADER_FOR_PARTITION)
         } else {
-          throw new ReplicaNotAvailableException(s"Partition $topicPartition is not available")
+          Left(Errors.REPLICA_NOT_AVAILABLE)
         }
 
-      case None =>
-        throw new UnknownTopicOrPartitionException(s"Partition $topicPartition doesn't exist")
+      case HostedPartition.None =>
+        Left(Errors.UNKNOWN_TOPIC_OR_PARTITION)
     }
   }
 
-  def localReplicaOrException(topicPartition: TopicPartition): Replica = {
-    getPartitionOrException(topicPartition, expectLeader = false).localReplicaOrException
+  def localLogOrException(topicPartition: TopicPartition): Log = {
+    getPartitionOrException(topicPartition, expectLeader = false).localLogOrException
   }
 
-  def futureLocalReplicaOrException(topicPartition: TopicPartition): Replica = {
-    getPartitionOrException(topicPartition, expectLeader = false).futureLocalReplicaOrException
+  def futureLocalLogOrException(topicPartition: TopicPartition): Log = {
+    getPartitionOrException(topicPartition, expectLeader = false).futureLocalLogOrException
   }
 
-  def futureLocalReplica(topicPartition: TopicPartition): Option[Replica] = {
-    nonOfflinePartition(topicPartition).flatMap(_.futureLocalReplica)
+  def futureLogExists(topicPartition: TopicPartition): Boolean = {
+    getPartitionOrException(topicPartition, expectLeader = false).futureLog.isDefined
   }
 
-  def localReplica(topicPartition: TopicPartition): Option[Replica] = {
-    nonOfflinePartition(topicPartition).flatMap(_.localReplica)
+  def localLog(topicPartition: TopicPartition): Option[Log] = {
+    nonOfflinePartition(topicPartition).flatMap(_.log)
   }
 
   def getLogDir(topicPartition: TopicPartition): Option[String] = {
-    localReplica(topicPartition).flatMap(_.log).map(_.dir.getParent)
+    localLog(topicPartition).map(_.parentDir)
   }
 
   /**
@@ -475,25 +535,26 @@ class ReplicaManager(val config: KafkaConfig,
   def appendRecords(timeout: Long,
                     requiredAcks: Short,
                     internalTopicsAllowed: Boolean,
-                    isFromClient: Boolean,
+                    origin: AppendOrigin,
                     entriesPerPartition: Map[TopicPartition, MemoryRecords],
                     responseCallback: Map[TopicPartition, PartitionResponse] => Unit,
                     delayedProduceLock: Option[Lock] = None,
-                    recordConversionStatsCallback: Map[TopicPartition, RecordConversionStats] => Unit = _ => ()) {
+                    recordConversionStatsCallback: Map[TopicPartition, RecordConversionStats] => Unit = _ => ()): Unit = {
     if (isValidRequiredAcks(requiredAcks)) {
       val sTime = time.milliseconds
       val localProduceResults = appendToLocalLog(internalTopicsAllowed = internalTopicsAllowed,
-        isFromClient = isFromClient, entriesPerPartition, requiredAcks)
+        origin, entriesPerPartition, requiredAcks)
       debug("Produce to local log in %d ms".format(time.milliseconds - sTime))
 
       val produceStatus = localProduceResults.map { case (topicPartition, result) =>
         topicPartition ->
                 ProducePartitionStatus(
                   result.info.lastOffset + 1, // required offset
-                  new PartitionResponse(result.error, result.info.firstOffset.getOrElse(-1), result.info.logAppendTime, result.info.logStartOffset)) // response status
+                  new PartitionResponse(result.error, result.info.firstOffset.getOrElse(-1), result.info.logAppendTime,
+                    result.info.logStartOffset, result.info.recordErrors.asJava, result.info.errorMessage)) // response status
       }
 
-      recordConversionStatsCallback(localProduceResults.mapValues(_.info.recordConversionStats))
+      recordConversionStatsCallback(localProduceResults.map { case (k, v) => k -> v.info.recordConversionStats })
 
       if (delayedProduceRequestRequired(requiredAcks, entriesPerPartition, localProduceResults)) {
         // create delayed produce operation
@@ -501,7 +562,7 @@ class ReplicaManager(val config: KafkaConfig,
         val delayedProduce = new DelayedProduce(timeout, produceMetadata, this, responseCallback, delayedProduceLock)
 
         // create a list of (topic, partition) pairs to use as keys for this delayed produce operation
-        val producerRequestKeys = entriesPerPartition.keys.map(new TopicPartitionOperationKey(_)).toSeq
+        val producerRequestKeys = entriesPerPartition.keys.map(TopicPartitionOperationKey(_)).toSeq
 
         // try to complete the request immediately, otherwise put it into the purgatory
         // this is because while the delayed produce operation is being created, new
@@ -510,7 +571,7 @@ class ReplicaManager(val config: KafkaConfig,
 
       } else {
         // we can respond immediately
-        val produceResponseStatus = produceStatus.mapValues(status => status.responseStatus)
+        val produceResponseStatus = produceStatus.map { case (k, status) => k -> status.responseStatus }
         responseCallback(produceResponseStatus)
       }
     } else {
@@ -543,7 +604,6 @@ class ReplicaManager(val config: KafkaConfig,
           case e@ (_: UnknownTopicOrPartitionException |
                    _: NotLeaderForPartitionException |
                    _: OffsetOutOfRangeException |
-                   _: PolicyViolationException |
                    _: KafkaStorageException) =>
             (topicPartition, LogDeleteRecordsResult(-1L, -1L, Some(e)))
           case t: Throwable =>
@@ -574,18 +634,25 @@ class ReplicaManager(val config: KafkaConfig,
     replicaStateChangeLock synchronized {
       partitionDirs.map { case (topicPartition, destinationDir) =>
         try {
+          /* If the topic name is exceptionally long, we can't support altering the log directory.
+           * See KAFKA-4893 for details.
+           * TODO: fix this by implementing topic IDs. */
+          if (Log.logFutureDirName(topicPartition).size > 255)
+            throw new InvalidTopicException("The topic name is too long.")
           if (!logManager.isLogDirOnline(destinationDir))
             throw new KafkaStorageException(s"Log directory $destinationDir is offline")
 
-          getPartition(topicPartition).foreach { partition =>
-            if (partition eq ReplicaManager.OfflinePartition)
+          getPartition(topicPartition) match {
+            case HostedPartition.Online(partition) =>
+              // Stop current replica movement if the destinationDir is different from the existing destination log directory
+              if (partition.futureReplicaDirChanged(destinationDir)) {
+                replicaAlterLogDirsManager.removeFetcherForPartitions(Set(topicPartition))
+                partition.removeFutureLocalReplica()
+              }
+            case HostedPartition.Offline =>
               throw new KafkaStorageException(s"Partition $topicPartition is offline")
 
-            // Stop current replica movement if the destinationDir is different from the existing destination log directory
-            if (partition.futureReplicaDirChanged(destinationDir)) {
-              replicaAlterLogDirsManager.removeFetcherForPartitions(Set(topicPartition))
-              partition.removeFutureLocalReplica()
-            }
+            case HostedPartition.None => // Do nothing
           }
 
           // If the log for this partition has not been created yet:
@@ -594,27 +661,29 @@ class ReplicaManager(val config: KafkaConfig,
           // 2) Respond with ReplicaNotAvailableException for this partition in the AlterReplicaLogDirsResponse
           logManager.maybeUpdatePreferredLogDir(topicPartition, destinationDir)
 
-          // throw ReplicaNotAvailableException if replica does not exit for the given partition
+          // throw ReplicaNotAvailableException if replica does not exist for the given partition
           val partition = getPartitionOrException(topicPartition, expectLeader = false)
-          partition.localReplicaOrException
+          partition.localLogOrException
 
           // If the destinationLDir is different from the current log directory of the replica:
           // - If there is no offline log directory, create the future log in the destinationDir (if it does not exist) and
           //   start ReplicaAlterDirThread to move data of this partition from the current log to the future log
           // - Otherwise, return KafkaStorageException. We do not create the future log while there is offline log directory
           //   so that we can avoid creating future log for the same partition in multiple log directories.
-          if (partition.maybeCreateFutureReplica(destinationDir)) {
-            val futureReplica = futureLocalReplicaOrException(topicPartition)
+          val highWatermarkCheckpoints = new LazyOffsetCheckpoints(this.highWatermarkCheckpoints)
+          if (partition.maybeCreateFutureReplica(destinationDir, highWatermarkCheckpoints)) {
+            val futureLog = futureLocalLogOrException(topicPartition)
             logManager.abortAndPauseCleaning(topicPartition)
 
             val initialFetchState = InitialFetchState(BrokerEndPoint(config.brokerId, "localhost", -1),
-              partition.getLeaderEpoch, futureReplica.highWatermark.messageOffset)
+              partition.getLeaderEpoch, futureLog.highWatermark)
             replicaAlterLogDirsManager.addFetcherForPartitions(Map(topicPartition -> initialFetchState))
           }
 
           (topicPartition, Errors.NONE)
         } catch {
-          case e@(_: LogDirNotFoundException |
+          case e@(_: InvalidTopicException |
+                  _: LogDirNotFoundException |
                   _: ReplicaNotAvailableException |
                   _: KafkaStorageException) =>
             (topicPartition, Errors.forException(e))
@@ -635,7 +704,7 @@ class ReplicaManager(val config: KafkaConfig,
    *    are included. There may be future logs (which will replace the current logs of the partition in the future) on the broker after KIP-113 is implemented.
    */
   def describeLogDirs(partitions: Set[TopicPartition]): Map[String, LogDirInfo] = {
-    val logsByDir = logManager.allLogs.groupBy(log => log.dir.getParent)
+    val logsByDir = logManager.allLogs.groupBy(log => log.parentDir)
 
     config.logDirs.toSet.map { logDir: String =>
       val absolutePath = new File(logDir).getAbsolutePath
@@ -667,12 +736,12 @@ class ReplicaManager(val config: KafkaConfig,
   }
 
   def getLogEndOffsetLag(topicPartition: TopicPartition, logEndOffset: Long, isFuture: Boolean): Long = {
-    localReplica(topicPartition) match {
-      case Some(replica) =>
+    localLog(topicPartition) match {
+      case Some(log) =>
         if (isFuture)
-          replica.logEndOffset - logEndOffset
+          log.logEndOffset - logEndOffset
         else
-          math.max(replica.highWatermark.messageOffset - logEndOffset, 0)
+          math.max(log.highWatermark - logEndOffset, 0)
       case None =>
         // return -1L to indicate that the LEO lag is not available if the replica is not created or is offline
         DescribeLogDirsResponse.INVALID_OFFSET_LAG
@@ -681,7 +750,7 @@ class ReplicaManager(val config: KafkaConfig,
 
   def deleteRecords(timeout: Long,
                     offsetPerPartition: Map[TopicPartition, Long],
-                    responseCallback: Map[TopicPartition, DeleteRecordsResponse.PartitionResponse] => Unit) {
+                    responseCallback: Map[TopicPartition, DeleteRecordsResponse.PartitionResponse] => Unit): Unit = {
     val timeBeforeLocalDeleteRecords = time.milliseconds
     val localDeleteRecordsResults = deleteRecordsOnLocalLog(offsetPerPartition)
     debug("Delete records on local log in %d ms".format(time.milliseconds - timeBeforeLocalDeleteRecords))
@@ -698,7 +767,7 @@ class ReplicaManager(val config: KafkaConfig,
       val delayedDeleteRecords = new DelayedDeleteRecords(timeout, deleteRecordsStatus, this, responseCallback)
 
       // create a list of (topic, partition) pairs to use as keys for this delayed delete records operation
-      val deleteRecordsRequestKeys = offsetPerPartition.keys.map(new TopicPartitionOperationKey(_)).toSeq
+      val deleteRecordsRequestKeys = offsetPerPartition.keys.map(TopicPartitionOperationKey(_)).toSeq
 
       // try to complete the request immediately, otherwise put it into the purgatory
       // this is because while the delayed delete records operation is being created, new
@@ -706,7 +775,7 @@ class ReplicaManager(val config: KafkaConfig,
       delayedDeleteRecordsPurgatory.tryCompleteElseWatch(delayedDeleteRecords, deleteRecordsRequestKeys)
     } else {
       // we can respond immediately
-      val deleteRecordsResponseStatus = deleteRecordsStatus.mapValues(status => status.responseStatus)
+      val deleteRecordsResponseStatus = deleteRecordsStatus.map { case (k, status) => k -> status.responseStatus }
       responseCallback(deleteRecordsResponseStatus)
     }
   }
@@ -732,9 +801,22 @@ class ReplicaManager(val config: KafkaConfig,
    * Append the messages to the local replica logs
    */
   private def appendToLocalLog(internalTopicsAllowed: Boolean,
-                               isFromClient: Boolean,
+                               origin: AppendOrigin,
                                entriesPerPartition: Map[TopicPartition, MemoryRecords],
                                requiredAcks: Short): Map[TopicPartition, LogAppendResult] = {
+
+    def processFailedRecord(topicPartition: TopicPartition, t: Throwable) = {
+      val logStartOffset = getPartition(topicPartition) match {
+        case HostedPartition.Online(partition) => partition.logStartOffset
+        case HostedPartition.None | HostedPartition.Offline => -1L
+      }
+      brokerTopicStats.topicStats(topicPartition.topic).failedProduceRequestRate.mark()
+      brokerTopicStats.allTopicsStats.failedProduceRequestRate.mark()
+      error(s"Error processing append operation on partition $topicPartition", t)
+
+      logStartOffset
+    }
+
     trace(s"Append [$entriesPerPartition] to local log")
     entriesPerPartition.map { case (topicPartition, records) =>
       brokerTopicStats.topicStats(topicPartition.topic).totalProduceRequestRate.mark()
@@ -748,14 +830,24 @@ class ReplicaManager(val config: KafkaConfig,
       } else {
         try {
           val partition = getPartitionOrException(topicPartition, expectLeader = true)
-          val info = partition.appendRecordsToLeader(records, isFromClient, requiredAcks)
+          val info = partition.appendRecordsToLeader(records, origin, requiredAcks)
           val numAppendedMessages = info.numMessages
 
+          // update stats for compressed or decompressed batches on broker
+          recompressedBatchCount.addAndGet(info.recompressedBatchCount)
+
           // update stats for successfully appended bytes and messages as bytesInRate and messageInRate
           brokerTopicStats.topicStats(topicPartition.topic).bytesInRate.mark(records.sizeInBytes)
+          brokerTopicStats.topicStats(topicPartition.topic).bytesInTotal.inc(records.sizeInBytes)
+
           brokerTopicStats.allTopicsStats.bytesInRate.mark(records.sizeInBytes)
+          brokerTopicStats.allTopicsStats.bytesInTotal.inc(records.sizeInBytes)
+
           brokerTopicStats.topicStats(topicPartition.topic).messagesInRate.mark(numAppendedMessages)
+          brokerTopicStats.topicStats(topicPartition.topic).messagesInTotal.inc(numAppendedMessages)
+
           brokerTopicStats.allTopicsStats.messagesInRate.mark(numAppendedMessages)
+          brokerTopicStats.allTopicsStats.messagesInTotal.inc(numAppendedMessages)
 
           trace(s"${records.sizeInBytes} written to log $topicPartition beginning at offset " +
             s"${info.firstOffset.getOrElse(-1)} and ending at offset ${info.lastOffset}")
@@ -768,19 +860,15 @@ class ReplicaManager(val config: KafkaConfig,
                    _: RecordTooLargeException |
                    _: RecordBatchTooLargeException |
                    _: CorruptRecordException |
-                   _: KafkaStorageException |
-                   _: InvalidTimestampException) =>
+                   _: KafkaStorageException) =>
             (topicPartition, LogAppendResult(LogAppendInfo.UnknownLogAppendInfo, Some(e)))
+          case rve: RecordValidationException =>
+            val logStartOffset = processFailedRecord(topicPartition, rve.invalidException)
+            val recordErrors = rve.recordErrors
+            (topicPartition, LogAppendResult(LogAppendInfo.unknownLogAppendInfoWithAdditionalInfo(
+              logStartOffset, recordErrors, rve.invalidException.getMessage), Some(rve.invalidException)))
           case t: Throwable =>
-            val logStartOffset = getPartition(topicPartition) match {
-              case Some(partition) =>
-                partition.logStartOffset
-              case _ =>
-                -1
-            }
-            brokerTopicStats.topicStats(topicPartition.topic).failedProduceRequestRate.mark()
-            brokerTopicStats.allTopicsStats.failedProduceRequestRate.mark()
-            error(s"Error processing append operation on partition $topicPartition", t)
+            val logStartOffset = processFailedRecord(topicPartition, t)
             (topicPartition, LogAppendResult(LogAppendInfo.unknownLogAppendInfoWithLogStartOffset(logStartOffset), Some(t)))
         }
       }
@@ -806,8 +894,9 @@ class ReplicaManager(val config: KafkaConfig,
   }
 
   /**
-   * Fetch messages from the leader replica, and wait until enough data can be fetched and return;
-   * the callback function will be triggered either when timeout or required fetch info is satisfied
+   * Fetch messages from a replica, and wait until enough data can be fetched and return;
+   * the callback function will be triggered either when timeout or required fetch info is satisfied.
+   * Consumers may fetch from any replica, but followers can only fetch from the leader.
    */
   def fetchMessages(timeout: Long,
                     replicaId: Int,
@@ -815,20 +904,21 @@ class ReplicaManager(val config: KafkaConfig,
                     fetchMaxBytes: Int,
                     hardMaxBytesLimit: Boolean,
                     fetchInfos: Seq[(TopicPartition, PartitionData)],
-                    quota: ReplicaQuota = UnboundedQuota,
+                    quota: ReplicaQuota,
                     responseCallback: Seq[(TopicPartition, FetchPartitionData)] => Unit,
-                    isolationLevel: IsolationLevel) {
+                    isolationLevel: IsolationLevel,
+                    clientMetadata: Option[ClientMetadata]): Unit = {
     val isFromFollower = Request.isValidBrokerId(replicaId)
-    val fetchOnlyFromLeader = replicaId != Request.DebuggingConsumerId && replicaId != Request.FutureLocalReplicaId
-
-    val fetchIsolation = if (isFromFollower || replicaId == Request.FutureLocalReplicaId)
+    val isFromConsumer = !(isFromFollower || replicaId == Request.FutureLocalReplicaId)
+    val fetchIsolation = if (!isFromConsumer)
       FetchLogEnd
     else if (isolationLevel == IsolationLevel.READ_COMMITTED)
       FetchTxnCommitted
     else
       FetchHighWatermark
 
-
+    // Restrict fetching to leader if request is from follower or from a client with older version (no ClientMetadata)
+    val fetchOnlyFromLeader = isFromFollower || (isFromConsumer && clientMetadata.isEmpty)
     def readFromLog(): Seq[(TopicPartition, LogReadResult)] = {
       val result = readFromLocalLog(
         replicaId = replicaId,
@@ -837,8 +927,9 @@ class ReplicaManager(val config: KafkaConfig,
         fetchMaxBytes = fetchMaxBytes,
         hardMaxBytesLimit = hardMaxBytesLimit,
         readPartitionInfo = fetchInfos,
-        quota = quota)
-      if (isFromFollower) updateFollowerLogReadResults(replicaId, result)
+        quota = quota,
+        clientMetadata = clientMetadata)
+      if (isFromFollower) updateFollowerFetchState(replicaId, result)
       else result
     }
 
@@ -848,21 +939,26 @@ class ReplicaManager(val config: KafkaConfig,
     var bytesReadable: Long = 0
     var errorReadingData = false
     val logReadResultMap = new mutable.HashMap[TopicPartition, LogReadResult]
+    var anyPartitionsNeedHwUpdate = false
     logReadResults.foreach { case (topicPartition, logReadResult) =>
       if (logReadResult.error != Errors.NONE)
         errorReadingData = true
       bytesReadable = bytesReadable + logReadResult.info.records.sizeInBytes
       logReadResultMap.put(topicPartition, logReadResult)
+      if (isFromFollower && logReadResult.followerNeedsHwUpdate) {
+        anyPartitionsNeedHwUpdate = true
+      }
     }
 
     // respond immediately if 1) fetch request does not want to wait
     //                        2) fetch request does not require any data
     //                        3) has enough data to respond
     //                        4) some error happens while reading data
-    if (timeout <= 0 || fetchInfos.isEmpty || bytesReadable >= fetchMinBytes || errorReadingData) {
+    //                        5) any of the requested partitions need HW update
+    if (timeout <= 0 || fetchInfos.isEmpty || bytesReadable >= fetchMinBytes || errorReadingData || anyPartitionsNeedHwUpdate) {
       val fetchPartitionData = logReadResults.map { case (tp, result) =>
         tp -> FetchPartitionData(result.error, result.highWatermark, result.leaderLogStartOffset, result.info.records,
-          result.lastStableOffset, result.info.abortedTransactions)
+          result.lastStableOffset, result.info.abortedTransactions, result.preferredReadReplica)
       }
       responseCallback(fetchPartitionData)
     } else {
@@ -876,10 +972,11 @@ class ReplicaManager(val config: KafkaConfig,
       }
       val fetchMetadata = FetchMetadata(fetchMinBytes, fetchMaxBytes, hardMaxBytesLimit, fetchOnlyFromLeader,
         fetchIsolation, isFromFollower, replicaId, fetchPartitionStatus)
-      val delayedFetch = new DelayedFetch(timeout, fetchMetadata, this, quota, responseCallback)
+      val delayedFetch = new DelayedFetch(timeout, fetchMetadata, this, quota, clientMetadata,
+        responseCallback)
 
       // create a list of (topic, partition) pairs to use as keys for this delayed fetch operation
-      val delayedFetchKeys = fetchPartitionStatus.map { case (tp, _) => new TopicPartitionOperationKey(tp) }
+      val delayedFetchKeys = fetchPartitionStatus.map { case (tp, _) => TopicPartitionOperationKey(tp) }
 
       // try to complete the request immediately, otherwise put it into the purgatory;
       // this is because while the delayed fetch operation is being created, new requests
@@ -897,7 +994,8 @@ class ReplicaManager(val config: KafkaConfig,
                        fetchMaxBytes: Int,
                        hardMaxBytesLimit: Boolean,
                        readPartitionInfo: Seq[(TopicPartition, PartitionData)],
-                       quota: ReplicaQuota): Seq[(TopicPartition, LogReadResult)] = {
+                       quota: ReplicaQuota,
+                       clientMetadata: Option[ClientMetadata]): Seq[(TopicPartition, LogReadResult)] = {
 
     def read(tp: TopicPartition, fetchInfo: PartitionData, limitBytes: Int, minOneMessage: Boolean): LogReadResult = {
       val offset = fetchInfo.fetchOffset
@@ -916,35 +1014,64 @@ class ReplicaManager(val config: KafkaConfig,
         val partition = getPartitionOrException(tp, expectLeader = fetchOnlyFromLeader)
         val fetchTimeMs = time.milliseconds
 
-        // Try the read first, this tells us whether we need all of adjustedFetchSize for this partition
-        val readInfo = partition.readRecords(
-          fetchOffset = fetchInfo.fetchOffset,
-          currentLeaderEpoch = fetchInfo.currentLeaderEpoch,
-          maxBytes = adjustedMaxBytes,
-          fetchIsolation = fetchIsolation,
-          fetchOnlyFromLeader = fetchOnlyFromLeader,
-          minOneMessage = minOneMessage)
-
-        val fetchDataInfo = if (shouldLeaderThrottle(quota, tp, replicaId)) {
-          // If the partition is being throttled, simply return an empty set.
-          FetchDataInfo(readInfo.fetchedData.fetchOffsetMetadata, MemoryRecords.EMPTY)
-        } else if (!hardMaxBytesLimit && readInfo.fetchedData.firstEntryIncomplete) {
-          // For FetchRequest version 3, we replace incomplete message sets with an empty one as consumers can make
-          // progress in such cases and don't need to report a `RecordTooLargeException`
-          FetchDataInfo(readInfo.fetchedData.fetchOffsetMetadata, MemoryRecords.EMPTY)
+        // If we are the leader, determine the preferred read-replica
+        val preferredReadReplica = clientMetadata.flatMap(
+          metadata => findPreferredReadReplica(tp, metadata, replicaId, fetchInfo.fetchOffset, fetchTimeMs))
+
+        if (preferredReadReplica.isDefined) {
+          replicaSelectorOpt.foreach{ selector =>
+            debug(s"Replica selector ${selector.getClass.getSimpleName} returned preferred replica " +
+              s"${preferredReadReplica.get} for $clientMetadata")
+          }
+          // If a preferred read-replica is set, skip the read
+          val offsetSnapshot = partition.fetchOffsetSnapshot(fetchInfo.currentLeaderEpoch, fetchOnlyFromLeader = false)
+          LogReadResult(info = FetchDataInfo(LogOffsetMetadata.UnknownOffsetMetadata, MemoryRecords.EMPTY),
+            highWatermark = offsetSnapshot.highWatermark.messageOffset,
+            leaderLogStartOffset = offsetSnapshot.logStartOffset,
+            leaderLogEndOffset = offsetSnapshot.logEndOffset.messageOffset,
+            followerLogStartOffset = followerLogStartOffset,
+            fetchTimeMs = -1L,
+            readSize = 0,
+            lastStableOffset = Some(offsetSnapshot.lastStableOffset.messageOffset),
+            preferredReadReplica = preferredReadReplica,
+            exception = None)
         } else {
-          readInfo.fetchedData
-        }
+          // Try the read first, this tells us whether we need all of adjustedFetchSize for this partition
+          val readInfo: LogReadInfo = partition.readRecords(
+            fetchOffset = fetchInfo.fetchOffset,
+            currentLeaderEpoch = fetchInfo.currentLeaderEpoch,
+            maxBytes = adjustedMaxBytes,
+            fetchIsolation = fetchIsolation,
+            fetchOnlyFromLeader = fetchOnlyFromLeader,
+            minOneMessage = minOneMessage)
+
+          // Check if the HW known to the follower is behind the actual HW
+          val followerNeedsHwUpdate: Boolean = partition.getReplica(replicaId)
+            .exists(replica => replica.lastSentHighWatermark < readInfo.highWatermark)
+
+          val fetchDataInfo = if (shouldLeaderThrottle(quota, tp, replicaId)) {
+            // If the partition is being throttled, simply return an empty set.
+            FetchDataInfo(readInfo.fetchedData.fetchOffsetMetadata, MemoryRecords.EMPTY)
+          } else if (!hardMaxBytesLimit && readInfo.fetchedData.firstEntryIncomplete) {
+            // For FetchRequest version 3, we replace incomplete message sets with an empty one as consumers can make
+            // progress in such cases and don't need to report a `RecordTooLargeException`
+            FetchDataInfo(readInfo.fetchedData.fetchOffsetMetadata, MemoryRecords.EMPTY)
+          } else {
+            readInfo.fetchedData
+          }
 
-        LogReadResult(info = fetchDataInfo,
-                      highWatermark = readInfo.highWatermark,
-                      leaderLogStartOffset = readInfo.logStartOffset,
-                      leaderLogEndOffset = readInfo.logEndOffset,
-                      followerLogStartOffset = followerLogStartOffset,
-                      fetchTimeMs = fetchTimeMs,
-                      readSize = adjustedMaxBytes,
-                      lastStableOffset = Some(readInfo.lastStableOffset),
-                      exception = None)
+          LogReadResult(info = fetchDataInfo,
+            highWatermark = readInfo.highWatermark,
+            leaderLogStartOffset = readInfo.logStartOffset,
+            leaderLogEndOffset = readInfo.logEndOffset,
+            followerLogStartOffset = followerLogStartOffset,
+            fetchTimeMs = fetchTimeMs,
+            readSize = adjustedMaxBytes,
+            lastStableOffset = Some(readInfo.lastStableOffset),
+            preferredReadReplica = preferredReadReplica,
+            followerNeedsHwUpdate = followerNeedsHwUpdate,
+            exception = None)
+        }
       } catch {
         // NOTE: Failed fetch requests metric is not incremented for known exceptions since it
         // is supposed to indicate un-expected failure of a broker in handling a fetch request
@@ -956,14 +1083,14 @@ class ReplicaManager(val config: KafkaConfig,
                  _: KafkaStorageException |
                  _: OffsetOutOfRangeException) =>
           LogReadResult(info = FetchDataInfo(LogOffsetMetadata.UnknownOffsetMetadata, MemoryRecords.EMPTY),
-                        highWatermark = -1L,
-                        leaderLogStartOffset = -1L,
-                        leaderLogEndOffset = -1L,
-                        followerLogStartOffset = -1L,
-                        fetchTimeMs = -1L,
-                        readSize = 0,
-                        lastStableOffset = None,
-                        exception = Some(e))
+            highWatermark = Log.UnknownOffset,
+            leaderLogStartOffset = Log.UnknownOffset,
+            leaderLogEndOffset = Log.UnknownOffset,
+            followerLogStartOffset = Log.UnknownOffset,
+            fetchTimeMs = -1L,
+            readSize = 0,
+            lastStableOffset = None,
+            exception = Some(e))
         case e: Throwable =>
           brokerTopicStats.topicStats(tp.topic).failedFetchRequestRate.mark()
           brokerTopicStats.allTopicsStats.failedFetchRequestRate.mark()
@@ -973,14 +1100,14 @@ class ReplicaManager(val config: KafkaConfig,
             s"on partition $tp: $fetchInfo", e)
 
           LogReadResult(info = FetchDataInfo(LogOffsetMetadata.UnknownOffsetMetadata, MemoryRecords.EMPTY),
-                        highWatermark = -1L,
-                        leaderLogStartOffset = -1L,
-                        leaderLogEndOffset = -1L,
-                        followerLogStartOffset = -1L,
-                        fetchTimeMs = -1L,
-                        readSize = 0,
-                        lastStableOffset = None,
-                        exception = Some(e))
+            highWatermark = Log.UnknownOffset,
+            leaderLogStartOffset = Log.UnknownOffset,
+            leaderLogEndOffset = Log.UnknownOffset,
+            followerLogStartOffset = Log.UnknownOffset,
+            fetchTimeMs = -1L,
+            readSize = 0,
+            lastStableOffset = None,
+            exception = Some(e))
       }
     }
 
@@ -999,18 +1126,69 @@ class ReplicaManager(val config: KafkaConfig,
     result
   }
 
+  /**
+    * Using the configured [[ReplicaSelector]], determine the preferred read replica for a partition given the
+    * client metadata, the requested offset, and the current set of replicas. If the preferred read replica is the
+    * leader, return None
+    */
+  def findPreferredReadReplica(tp: TopicPartition,
+                               clientMetadata: ClientMetadata,
+                               replicaId: Int,
+                               fetchOffset: Long,
+                               currentTimeMs: Long): Option[Int] = {
+    val partition = getPartitionOrException(tp, expectLeader = false)
+
+    if (partition.isLeader) {
+      if (Request.isValidBrokerId(replicaId)) {
+        // Don't look up preferred for follower fetches via normal replication
+        Option.empty
+      } else {
+        replicaSelectorOpt.flatMap { replicaSelector =>
+          val replicaEndpoints = metadataCache.getPartitionReplicaEndpoints(tp, new ListenerName(clientMetadata.listenerName))
+          var replicaInfoSet: Set[ReplicaView] = partition.remoteReplicas
+            // Exclude replicas that don't have the requested offset (whether or not if they're in the ISR)
+            .filter(replica => replica.logEndOffset >= fetchOffset)
+            .filter(replica => replica.logStartOffset <= fetchOffset)
+            .map(replica => new DefaultReplicaView(
+              replicaEndpoints.getOrElse(replica.brokerId, Node.noNode()),
+              replica.logEndOffset,
+              currentTimeMs - replica.lastCaughtUpTimeMs))
+            .toSet
+
+          if (partition.leaderReplicaIdOpt.isDefined) {
+            val leaderReplica: ReplicaView = partition.leaderReplicaIdOpt
+              .map(replicaId => replicaEndpoints.getOrElse(replicaId, Node.noNode()))
+              .map(leaderNode => new DefaultReplicaView(leaderNode, partition.localLogOrException.logEndOffset, 0L))
+              .get
+            replicaInfoSet ++= Set(leaderReplica)
+
+            val partitionInfo = new DefaultPartitionView(replicaInfoSet.asJava, leaderReplica)
+            replicaSelector.select(tp, clientMetadata, partitionInfo).asScala
+              .filter(!_.endpoint.isEmpty)
+              // Even though the replica selector can return the leader, we don't want to send it out with the
+              // FetchResponse, so we exclude it here
+              .filter(!_.equals(leaderReplica))
+              .map(_.endpoint.id)
+          } else {
+            None
+          }
+        }
+      }
+    } else {
+      None
+    }
+  }
+
   /**
    *  To avoid ISR thrashing, we only throttle a replica on the leader if it's in the throttled replica list,
    *  the quota is exceeded and the replica is not in sync.
    */
   def shouldLeaderThrottle(quota: ReplicaQuota, topicPartition: TopicPartition, replicaId: Int): Boolean = {
-    val isReplicaInSync = nonOfflinePartition(topicPartition).exists { partition =>
-      partition.getReplica(replicaId).exists(partition.inSyncReplicas.contains)
-    }
+    val isReplicaInSync = nonOfflinePartition(topicPartition).exists(_.inSyncReplicaIds.contains(replicaId))
     !isReplicaInSync && quota.isThrottled(topicPartition) && quota.isQuotaExceeded
   }
 
-  def getLogConfig(topicPartition: TopicPartition): Option[LogConfig] = localReplica(topicPartition).flatMap(_.log.map(_.config))
+  def getLogConfig(topicPartition: TopicPartition): Option[LogConfig] = localLog(topicPartition).map(_.config)
 
   def getMagic(topicPartition: TopicPartition): Option[Byte] = getLogConfig(topicPartition).map(_.messageFormatVersion.recordVersion.value)
 
@@ -1033,10 +1211,12 @@ class ReplicaManager(val config: KafkaConfig,
   def becomeLeaderOrFollower(correlationId: Int,
                              leaderAndIsrRequest: LeaderAndIsrRequest,
                              onLeadershipChange: (Iterable[Partition], Iterable[Partition]) => Unit): LeaderAndIsrResponse = {
-    leaderAndIsrRequest.partitionStates.asScala.foreach { case (topicPartition, stateInfo) =>
-      stateChangeLogger.trace(s"Received LeaderAndIsr request $stateInfo " +
-        s"correlation id $correlationId from controller ${leaderAndIsrRequest.controllerId} " +
-        s"epoch ${leaderAndIsrRequest.controllerEpoch} for partition $topicPartition")
+    if (stateChangeLogger.isTraceEnabled) {
+      leaderAndIsrRequest.partitionStates.asScala.foreach { partitionState =>
+        stateChangeLogger.trace(s"Received LeaderAndIsr request $partitionState " +
+          s"correlation id $correlationId from controller ${leaderAndIsrRequest.controllerId} " +
+          s"epoch ${leaderAndIsrRequest.controllerEpoch}")
+      }
     }
     replicaStateChangeLock synchronized {
       if (leaderAndIsrRequest.controllerEpoch < controllerEpoch) {
@@ -1050,93 +1230,125 @@ class ReplicaManager(val config: KafkaConfig,
         controllerEpoch = leaderAndIsrRequest.controllerEpoch
 
         // First check partition's leader epoch
-        val partitionState = new mutable.HashMap[Partition, LeaderAndIsrRequest.PartitionState]()
-        val newPartitions = new mutable.HashSet[Partition]
-
-        leaderAndIsrRequest.partitionStates.asScala.foreach { case (topicPartition, stateInfo) =>
-          val partition = getPartition(topicPartition).getOrElse {
-            val createdPartition = getOrCreatePartition(topicPartition)
-            newPartitions.add(createdPartition)
-            createdPartition
+        val partitionStates = new mutable.HashMap[Partition, LeaderAndIsrPartitionState]()
+        val updatedPartitions = new mutable.HashSet[Partition]
+
+        leaderAndIsrRequest.partitionStates.asScala.foreach { partitionState =>
+          val topicPartition = new TopicPartition(partitionState.topicName, partitionState.partitionIndex)
+          val partitionOpt = getPartition(topicPartition) match {
+            case HostedPartition.Offline =>
+              stateChangeLogger.warn(s"Ignoring LeaderAndIsr request from " +
+                s"controller $controllerId with correlation id $correlationId " +
+                s"epoch $controllerEpoch for partition $topicPartition as the local replica for the " +
+                "partition is in an offline log directory")
+              responseMap.put(topicPartition, Errors.KAFKA_STORAGE_ERROR)
+              None
+
+            case HostedPartition.Online(partition) =>
+              updatedPartitions.add(partition)
+              Some(partition)
+
+            case HostedPartition.None =>
+              val partition = Partition(topicPartition, time, this)
+              allPartitions.putIfNotExists(topicPartition, HostedPartition.Online(partition))
+              updatedPartitions.add(partition)
+              Some(partition)
           }
-          val currentLeaderEpoch = partition.getLeaderEpoch
-          val requestLeaderEpoch = stateInfo.basePartitionState.leaderEpoch
-          if (partition eq ReplicaManager.OfflinePartition) {
-            stateChangeLogger.warn(s"Ignoring LeaderAndIsr request from " +
-              s"controller $controllerId with correlation id $correlationId " +
-              s"epoch $controllerEpoch for partition $topicPartition as the local replica for the " +
-              "partition is in an offline log directory")
-            responseMap.put(topicPartition, Errors.KAFKA_STORAGE_ERROR)
-          } else if (requestLeaderEpoch > currentLeaderEpoch) {
-            // If the leader epoch is valid record the epoch of the controller that made the leadership decision.
-            // This is useful while updating the isr to maintain the decision maker controller's epoch in the zookeeper path
-            if(stateInfo.basePartitionState.replicas.contains(localBrokerId))
-              partitionState.put(partition, stateInfo)
-            else {
-              stateChangeLogger.warn(s"Ignoring LeaderAndIsr request from controller $controllerId with " +
-                s"correlation id $correlationId epoch $controllerEpoch for partition $topicPartition as itself is not " +
-                s"in assigned replica list ${stateInfo.basePartitionState.replicas.asScala.mkString(",")}")
-              responseMap.put(topicPartition, Errors.UNKNOWN_TOPIC_OR_PARTITION)
+
+          partitionOpt.foreach { partition =>
+            val currentLeaderEpoch = partition.getLeaderEpoch
+            val requestLeaderEpoch = partitionState.leaderEpoch
+            if (requestLeaderEpoch > currentLeaderEpoch) {
+              // If the leader epoch is valid record the epoch of the controller that made the leadership decision.
+              // This is useful while updating the isr to maintain the decision maker controller's epoch in the zookeeper path
+              if (partitionState.replicas.contains(localBrokerId))
+                partitionStates.put(partition, partitionState)
+              else {
+                stateChangeLogger.warn(s"Ignoring LeaderAndIsr request from controller $controllerId with " +
+                  s"correlation id $correlationId epoch $controllerEpoch for partition $topicPartition as itself is not " +
+                  s"in assigned replica list ${partitionState.replicas.asScala.mkString(",")}")
+                responseMap.put(topicPartition, Errors.UNKNOWN_TOPIC_OR_PARTITION)
+              }
+            } else if (requestLeaderEpoch < currentLeaderEpoch) {
+              stateChangeLogger.warn(s"Ignoring LeaderAndIsr request from " +
+                s"controller $controllerId with correlation id $correlationId " +
+                s"epoch $controllerEpoch for partition $topicPartition since its associated " +
+                s"leader epoch $requestLeaderEpoch is smaller than the current " +
+                s"leader epoch $currentLeaderEpoch")
+              responseMap.put(topicPartition, Errors.STALE_CONTROLLER_EPOCH)
+            } else {
+              stateChangeLogger.debug(s"Ignoring LeaderAndIsr request from " +
+                s"controller $controllerId with correlation id $correlationId " +
+                s"epoch $controllerEpoch for partition $topicPartition since its associated " +
+                s"leader epoch $requestLeaderEpoch matches the current leader epoch")
+              responseMap.put(topicPartition, Errors.STALE_CONTROLLER_EPOCH)
             }
-          } else {
-            // Otherwise record the error code in response
-            stateChangeLogger.warn(s"Ignoring LeaderAndIsr request from " +
-              s"controller $controllerId with correlation id $correlationId " +
-              s"epoch $controllerEpoch for partition $topicPartition since its associated " +
-              s"leader epoch $requestLeaderEpoch is not higher than the current " +
-              s"leader epoch $currentLeaderEpoch")
-            responseMap.put(topicPartition, Errors.STALE_CONTROLLER_EPOCH)
           }
         }
 
-        val partitionsTobeLeader = partitionState.filter { case (_, stateInfo) =>
-          stateInfo.basePartitionState.leader == localBrokerId
+        val partitionsTobeLeader = partitionStates.filter { case (_, partitionState) =>
+          partitionState.leader == localBrokerId
         }
-        val partitionsToBeFollower = partitionState -- partitionsTobeLeader.keys
+        val partitionsToBeFollower = partitionStates -- partitionsTobeLeader.keys
 
+        val highWatermarkCheckpoints = new LazyOffsetCheckpoints(this.highWatermarkCheckpoints)
         val partitionsBecomeLeader = if (partitionsTobeLeader.nonEmpty)
-          makeLeaders(controllerId, controllerEpoch, partitionsTobeLeader, correlationId, responseMap)
+          makeLeaders(controllerId, controllerEpoch, partitionsTobeLeader, correlationId, responseMap,
+            highWatermarkCheckpoints)
         else
           Set.empty[Partition]
         val partitionsBecomeFollower = if (partitionsToBeFollower.nonEmpty)
-          makeFollowers(controllerId, controllerEpoch, partitionsToBeFollower, correlationId, responseMap)
+          makeFollowers(controllerId, controllerEpoch, partitionsToBeFollower, correlationId, responseMap,
+            highWatermarkCheckpoints)
         else
           Set.empty[Partition]
 
-        leaderAndIsrRequest.partitionStates.asScala.keys.foreach { topicPartition =>
+        /*
+         * KAFKA-8392
+         * For topic partitions of which the broker is no longer a leader, delete metrics related to
+         * those topics. Note that this means the broker stops being either a replica or a leader of
+         * partitions of said topics
+         */
+        val leaderTopicSet = leaderPartitionsIterator.map(_.topic).toSet
+        val followerTopicSet = partitionsBecomeFollower.map(_.topic).toSet
+        followerTopicSet.diff(leaderTopicSet).foreach(brokerTopicStats.removeOldLeaderMetrics)
+
+        // remove metrics for brokers which are not followers of a topic
+        leaderTopicSet.diff(followerTopicSet).foreach(brokerTopicStats.removeOldFollowerMetrics)
+
+        leaderAndIsrRequest.partitionStates.asScala.foreach { partitionState =>
+          val topicPartition = new TopicPartition(partitionState.topicName, partitionState.partitionIndex)
           /*
            * If there is offline log directory, a Partition object may have been created by getOrCreatePartition()
            * before getOrCreateReplica() failed to create local replica due to KafkaStorageException.
            * In this case ReplicaManager.allPartitions will map this topic-partition to an empty Partition object.
            * we need to map this topic-partition to OfflinePartition instead.
            */
-          if (localReplica(topicPartition).isEmpty && (allPartitions.get(topicPartition) ne ReplicaManager.OfflinePartition))
-            allPartitions.put(topicPartition, ReplicaManager.OfflinePartition)
+          if (localLog(topicPartition).isEmpty)
+            markPartitionOffline(topicPartition)
         }
 
         // we initialize highwatermark thread after the first leaderisrrequest. This ensures that all the partitions
         // have been completely populated before starting the checkpointing there by avoiding weird race conditions
-        if (!hwThreadInitialized) {
-          startHighWaterMarksCheckPointThread()
-          hwThreadInitialized = true
-        }
+        startHighWatermarkCheckPointThread()
 
         val futureReplicasAndInitialOffset = new mutable.HashMap[TopicPartition, InitialFetchState]
-        for (partition <- newPartitions) {
+        for (partition <- updatedPartitions) {
           val topicPartition = partition.topicPartition
           if (logManager.getLog(topicPartition, isFuture = true).isDefined) {
-            partition.localReplica.foreach { replica =>
+            partition.log.foreach { log =>
               val leader = BrokerEndPoint(config.brokerId, "localhost", -1)
 
               // Add future replica to partition's map
-              partition.getOrCreateReplica(Request.FutureLocalReplicaId, isNew = false)
+              partition.createLogIfNotExists(isNew = false, isFutureReplica = true,
+                highWatermarkCheckpoints)
 
               // pause cleaning for partitions that are being moved and start ReplicaAlterDirThread to move
               // replica from source dir to destination dir
               logManager.abortAndPauseCleaning(topicPartition)
 
               futureReplicasAndInitialOffset.put(topicPartition, InitialFetchState(leader,
-                partition.getLeaderEpoch, replica.highWatermark.messageOffset))
+                partition.getLeaderEpoch, log.highWatermark))
             }
           }
         }
@@ -1145,7 +1357,15 @@ class ReplicaManager(val config: KafkaConfig,
         replicaFetcherManager.shutdownIdleFetcherThreads()
         replicaAlterLogDirsManager.shutdownIdleFetcherThreads()
         onLeadershipChange(partitionsBecomeLeader, partitionsBecomeFollower)
-        new LeaderAndIsrResponse(Errors.NONE, responseMap.asJava)
+        val responsePartitions = responseMap.iterator.map { case (tp, error) =>
+          new LeaderAndIsrPartitionError()
+            .setTopicName(tp.topic)
+            .setPartitionIndex(tp.partition)
+            .setErrorCode(error.code)
+        }.toBuffer
+        new LeaderAndIsrResponse(new LeaderAndIsrResponseData()
+          .setErrorCode(Errors.NONE.code)
+          .setPartitionErrors(responsePartitions.asJava))
       }
     }
   }
@@ -1164,42 +1384,43 @@ class ReplicaManager(val config: KafkaConfig,
    *  TODO: the above may need to be fixed later
    */
   private def makeLeaders(controllerId: Int,
-                          epoch: Int,
-                          partitionState: Map[Partition, LeaderAndIsrRequest.PartitionState],
+                          controllerEpoch: Int,
+                          partitionStates: Map[Partition, LeaderAndIsrPartitionState],
                           correlationId: Int,
-                          responseMap: mutable.Map[TopicPartition, Errors]): Set[Partition] = {
-    partitionState.keys.foreach { partition =>
+                          responseMap: mutable.Map[TopicPartition, Errors],
+                          highWatermarkCheckpoints: OffsetCheckpoints): Set[Partition] = {
+    partitionStates.keys.foreach { partition =>
       stateChangeLogger.trace(s"Handling LeaderAndIsr request correlationId $correlationId from " +
-        s"controller $controllerId epoch $epoch starting the become-leader transition for " +
+        s"controller $controllerId epoch $controllerEpoch starting the become-leader transition for " +
         s"partition ${partition.topicPartition}")
     }
 
-    for (partition <- partitionState.keys)
+    for (partition <- partitionStates.keys)
       responseMap.put(partition.topicPartition, Errors.NONE)
 
     val partitionsToMakeLeaders = mutable.Set[Partition]()
 
     try {
       // First stop fetchers for all the partitions
-      replicaFetcherManager.removeFetcherForPartitions(partitionState.keySet.map(_.topicPartition))
+      replicaFetcherManager.removeFetcherForPartitions(partitionStates.keySet.map(_.topicPartition))
       // Update the partition information to be the leader
-      partitionState.foreach{ case (partition, partitionStateInfo) =>
+      partitionStates.foreach { case (partition, partitionState) =>
         try {
-          if (partition.makeLeader(controllerId, partitionStateInfo, correlationId)) {
+          if (partition.makeLeader(partitionState, highWatermarkCheckpoints)) {
             partitionsToMakeLeaders += partition
             stateChangeLogger.trace(s"Stopped fetchers as part of become-leader request from " +
-              s"controller $controllerId epoch $epoch with correlation id $correlationId for partition ${partition.topicPartition} " +
-              s"(last update controller epoch ${partitionStateInfo.basePartitionState.controllerEpoch})")
+              s"controller $controllerId epoch $controllerEpoch with correlation id $correlationId for partition ${partition.topicPartition} " +
+              s"(last update controller epoch ${partitionState.controllerEpoch})")
           } else
             stateChangeLogger.info(s"Skipped the become-leader state change after marking its " +
-              s"partition as leader with correlation id $correlationId from controller $controllerId epoch $epoch for " +
-              s"partition ${partition.topicPartition} (last update controller epoch ${partitionStateInfo.basePartitionState.controllerEpoch}) " +
+              s"partition as leader with correlation id $correlationId from controller $controllerId epoch $controllerEpoch for " +
+              s"partition ${partition.topicPartition} (last update controller epoch ${partitionState.controllerEpoch}) " +
               s"since it is already the leader for the partition.")
         } catch {
           case e: KafkaStorageException =>
             stateChangeLogger.error(s"Skipped the become-leader state change with " +
-              s"correlation id $correlationId from controller $controllerId epoch $epoch for partition ${partition.topicPartition} " +
-              s"(last update controller epoch ${partitionStateInfo.basePartitionState.controllerEpoch}) since " +
+              s"correlation id $correlationId from controller $controllerId epoch $controllerEpoch for partition ${partition.topicPartition} " +
+              s"(last update controller epoch ${partitionState.controllerEpoch}) since " +
               s"the replica for the partition is offline due to disk error $e")
             val dirOpt = getLogDir(partition.topicPartition)
             error(s"Error while making broker the leader for partition $partition in dir $dirOpt", e)
@@ -1209,17 +1430,17 @@ class ReplicaManager(val config: KafkaConfig,
 
     } catch {
       case e: Throwable =>
-        partitionState.keys.foreach { partition =>
+        partitionStates.keys.foreach { partition =>
           stateChangeLogger.error(s"Error while processing LeaderAndIsr request correlationId $correlationId received " +
-            s"from controller $controllerId epoch $epoch for partition ${partition.topicPartition}", e)
+            s"from controller $controllerId epoch $controllerEpoch for partition ${partition.topicPartition}", e)
         }
         // Re-throw the exception for it to be caught in KafkaApis
         throw e
     }
 
-    partitionState.keys.foreach { partition =>
+    partitionStates.keys.foreach { partition =>
       stateChangeLogger.trace(s"Completed LeaderAndIsr request correlationId $correlationId from controller $controllerId " +
-        s"epoch $epoch for the become-leader transition for partition ${partition.topicPartition}")
+        s"epoch $controllerEpoch for the become-leader transition for partition ${partition.topicPartition}")
     }
 
     partitionsToMakeLeaders
@@ -1244,53 +1465,54 @@ class ReplicaManager(val config: KafkaConfig,
    * return the set of partitions that are made follower due to this method
    */
   private def makeFollowers(controllerId: Int,
-                            epoch: Int,
-                            partitionStates: Map[Partition, LeaderAndIsrRequest.PartitionState],
+                            controllerEpoch: Int,
+                            partitionStates: Map[Partition, LeaderAndIsrPartitionState],
                             correlationId: Int,
-                            responseMap: mutable.Map[TopicPartition, Errors]) : Set[Partition] = {
+                            responseMap: mutable.Map[TopicPartition, Errors],
+                            highWatermarkCheckpoints: OffsetCheckpoints) : Set[Partition] = {
     partitionStates.foreach { case (partition, partitionState) =>
       stateChangeLogger.trace(s"Handling LeaderAndIsr request correlationId $correlationId from controller $controllerId " +
-        s"epoch $epoch starting the become-follower transition for partition ${partition.topicPartition} with leader " +
-        s"${partitionState.basePartitionState.leader}")
+        s"epoch $controllerEpoch starting the become-follower transition for partition ${partition.topicPartition} with leader " +
+        s"${partitionState.leader}")
     }
 
     for (partition <- partitionStates.keys)
       responseMap.put(partition.topicPartition, Errors.NONE)
 
     val partitionsToMakeFollower: mutable.Set[Partition] = mutable.Set()
-
     try {
       // TODO: Delete leaders from LeaderAndIsrRequest
-      partitionStates.foreach { case (partition, partitionStateInfo) =>
-        val newLeaderBrokerId = partitionStateInfo.basePartitionState.leader
+      partitionStates.foreach { case (partition, partitionState) =>
+        val newLeaderBrokerId = partitionState.leader
         try {
           metadataCache.getAliveBrokers.find(_.id == newLeaderBrokerId) match {
             // Only change partition state when the leader is available
             case Some(_) =>
-              if (partition.makeFollower(controllerId, partitionStateInfo, correlationId))
+              if (partition.makeFollower(partitionState, highWatermarkCheckpoints))
                 partitionsToMakeFollower += partition
               else
                 stateChangeLogger.info(s"Skipped the become-follower state change after marking its partition as " +
-                  s"follower with correlation id $correlationId from controller $controllerId epoch $epoch " +
+                  s"follower with correlation id $correlationId from controller $controllerId epoch $controllerEpoch " +
                   s"for partition ${partition.topicPartition} (last update " +
-                  s"controller epoch ${partitionStateInfo.basePartitionState.controllerEpoch}) " +
+                  s"controller epoch ${partitionState.controllerEpoch}) " +
                   s"since the new leader $newLeaderBrokerId is the same as the old leader")
             case None =>
               // The leader broker should always be present in the metadata cache.
               // If not, we should record the error message and abort the transition process for this partition
               stateChangeLogger.error(s"Received LeaderAndIsrRequest with correlation id $correlationId from " +
-                s"controller $controllerId epoch $epoch for partition ${partition.topicPartition} " +
-                s"(last update controller epoch ${partitionStateInfo.basePartitionState.controllerEpoch}) " +
+                s"controller $controllerId epoch $controllerEpoch for partition ${partition.topicPartition} " +
+                s"(last update controller epoch ${partitionState.controllerEpoch}) " +
                 s"but cannot become follower since the new leader $newLeaderBrokerId is unavailable.")
               // Create the local replica even if the leader is unavailable. This is required to ensure that we include
               // the partition's high watermark in the checkpoint file (see KAFKA-1647)
-              partition.getOrCreateReplica(localBrokerId, isNew = partitionStateInfo.isNew)
+              partition.createLogIfNotExists(isNew = partitionState.isNew, isFutureReplica = false,
+                highWatermarkCheckpoints)
           }
         } catch {
           case e: KafkaStorageException =>
             stateChangeLogger.error(s"Skipped the become-follower state change with correlation id $correlationId from " +
-              s"controller $controllerId epoch $epoch for partition ${partition.topicPartition} " +
-              s"(last update controller epoch ${partitionStateInfo.basePartitionState.controllerEpoch}) with leader " +
+              s"controller $controllerId epoch $controllerEpoch for partition ${partition.topicPartition} " +
+              s"(last update controller epoch ${partitionState.controllerEpoch}) with leader " +
               s"$newLeaderBrokerId since the replica for the partition is offline due to disk error $e")
             val dirOpt = getLogDir(partition.topicPartition)
             error(s"Error while making broker the follower for partition $partition with leader " +
@@ -1302,58 +1524,55 @@ class ReplicaManager(val config: KafkaConfig,
       replicaFetcherManager.removeFetcherForPartitions(partitionsToMakeFollower.map(_.topicPartition))
       partitionsToMakeFollower.foreach { partition =>
         stateChangeLogger.trace(s"Stopped fetchers as part of become-follower request from controller $controllerId " +
-          s"epoch $epoch with correlation id $correlationId for partition ${partition.topicPartition} with leader " +
-          s"${partitionStates(partition).basePartitionState.leader}")
+          s"epoch $controllerEpoch with correlation id $correlationId for partition ${partition.topicPartition} with leader " +
+          s"${partitionStates(partition).leader}")
       }
 
       partitionsToMakeFollower.foreach { partition =>
-        val topicPartitionOperationKey = new TopicPartitionOperationKey(partition.topicPartition)
-        tryCompleteDelayedProduce(topicPartitionOperationKey)
-        tryCompleteDelayedFetch(topicPartitionOperationKey)
+        completeDelayedFetchOrProduceRequests(partition.topicPartition)
       }
 
       partitionsToMakeFollower.foreach { partition =>
         stateChangeLogger.trace(s"Truncated logs and checkpointed recovery boundaries for partition " +
           s"${partition.topicPartition} as part of become-follower request with correlation id $correlationId from " +
-          s"controller $controllerId epoch $epoch with leader ${partitionStates(partition).basePartitionState.leader}")
+          s"controller $controllerId epoch $controllerEpoch with leader ${partitionStates(partition).leader}")
       }
 
       if (isShuttingDown.get()) {
         partitionsToMakeFollower.foreach { partition =>
           stateChangeLogger.trace(s"Skipped the adding-fetcher step of the become-follower state " +
-            s"change with correlation id $correlationId from controller $controllerId epoch $epoch for " +
-            s"partition ${partition.topicPartition} with leader ${partitionStates(partition).basePartitionState.leader} " +
+            s"change with correlation id $correlationId from controller $controllerId epoch $controllerEpoch for " +
+            s"partition ${partition.topicPartition} with leader ${partitionStates(partition).leader} " +
             "since it is shutting down")
         }
-      }
-      else {
+      } else {
         // we do not need to check if the leader exists again since this has been done at the beginning of this process
         val partitionsToMakeFollowerWithLeaderAndOffset = partitionsToMakeFollower.map { partition =>
           val leader = metadataCache.getAliveBrokers.find(_.id == partition.leaderReplicaIdOpt.get).get
             .brokerEndPoint(config.interBrokerListenerName)
-          val fetchOffset = partition.localReplicaOrException.highWatermark.messageOffset
+          val fetchOffset = partition.localLogOrException.highWatermark
           partition.topicPartition -> InitialFetchState(leader, partition.getLeaderEpoch, fetchOffset)
-        }.toMap
-        replicaFetcherManager.addFetcherForPartitions(partitionsToMakeFollowerWithLeaderAndOffset)
+       }.toMap
 
-        partitionsToMakeFollower.foreach { partition =>
+        replicaFetcherManager.addFetcherForPartitions(partitionsToMakeFollowerWithLeaderAndOffset)
+        partitionsToMakeFollowerWithLeaderAndOffset.foreach { case (partition, initialFetchState) =>
           stateChangeLogger.trace(s"Started fetcher to new leader as part of become-follower " +
-            s"request from controller $controllerId epoch $epoch with correlation id $correlationId for " +
-            s"partition ${partition.topicPartition} with leader ${partitionStates(partition).basePartitionState.leader}")
+            s"request from controller $controllerId epoch $controllerEpoch with correlation id $correlationId for " +
+            s"partition $partition with leader ${initialFetchState.leader}")
         }
       }
     } catch {
       case e: Throwable =>
         stateChangeLogger.error(s"Error while processing LeaderAndIsr request with correlationId $correlationId " +
-          s"received from controller $controllerId epoch $epoch", e)
+          s"received from controller $controllerId epoch $controllerEpoch", e)
         // Re-throw the exception for it to be caught in KafkaApis
         throw e
     }
 
     partitionStates.keys.foreach { partition =>
       stateChangeLogger.trace(s"Completed LeaderAndIsr request correlationId $correlationId from controller $controllerId " +
-        s"epoch $epoch for the become-follower transition for partition ${partition.topicPartition} with leader " +
-        s"${partitionStates(partition).basePartitionState.leader}")
+        s"epoch $controllerEpoch for the become-follower transition for partition ${partition.topicPartition} with leader " +
+        s"${partitionStates(partition).leader}")
     }
 
     partitionsToMakeFollower
@@ -1369,86 +1588,98 @@ class ReplicaManager(val config: KafkaConfig,
   }
 
   /**
-   * Update the follower's fetch state in the leader based on the last fetch request and update `readResult`,
-   * if the follower replica is not recognized to be one of the assigned replicas. Do not update
-   * `readResult` otherwise, so that log start/end offset and high watermark is consistent with
+   * Update the follower's fetch state on the leader based on the last fetch request and update `readResult`.
+   * If the follower replica is not recognized to be one of the assigned replicas, do not update
+   * `readResult` so that log start/end offset and high watermark is consistent with
    * records in fetch response. Log start/end offset and high watermark may change not only due to
    * this fetch request, e.g., rolling new log segment and removing old log segment may move log
    * start offset further than the last offset in the fetched records. The followers will get the
    * updated leader's state in the next fetch response.
    */
-  private def updateFollowerLogReadResults(replicaId: Int,
-                                           readResults: Seq[(TopicPartition, LogReadResult)]): Seq[(TopicPartition, LogReadResult)] = {
-    debug(s"Recording follower broker $replicaId log end offsets: $readResults")
+  private def updateFollowerFetchState(followerId: Int,
+                                       readResults: Seq[(TopicPartition, LogReadResult)]): Seq[(TopicPartition, LogReadResult)] = {
     readResults.map { case (topicPartition, readResult) =>
-      var updatedReadResult = readResult
-      nonOfflinePartition(topicPartition) match {
-        case Some(partition) =>
-          partition.getReplica(replicaId) match {
-            case Some(replica) =>
-              partition.updateReplicaLogReadResult(replica, readResult)
-            case None =>
-              warn(s"Leader $localBrokerId failed to record follower $replicaId's position " +
-                s"${readResult.info.fetchOffsetMetadata.messageOffset} since the replica is not recognized to be " +
-                s"one of the assigned replicas ${partition.assignedReplicas.map(_.brokerId).mkString(",")} " +
+      val updatedReadResult = if (readResult.error != Errors.NONE) {
+        debug(s"Skipping update of fetch state for follower $followerId since the " +
+          s"log read returned error ${readResult.error}")
+        readResult
+      } else {
+        nonOfflinePartition(topicPartition) match {
+          case Some(partition) =>
+            if (partition.updateFollowerFetchState(followerId,
+              followerFetchOffsetMetadata = readResult.info.fetchOffsetMetadata,
+              followerStartOffset = readResult.followerLogStartOffset,
+              followerFetchTimeMs = readResult.fetchTimeMs,
+              leaderEndOffset = readResult.leaderLogEndOffset,
+              lastSentHighwatermark = readResult.highWatermark)) {
+              readResult
+            } else {
+              warn(s"Leader $localBrokerId failed to record follower $followerId's position " +
+                s"${readResult.info.fetchOffsetMetadata.messageOffset}, and last sent HW since the replica " +
+                s"is not recognized to be one of the assigned replicas ${partition.allReplicaIds.mkString(",")} " +
                 s"for partition $topicPartition. Empty records will be returned for this partition.")
-              updatedReadResult = readResult.withEmptyFetchInfo
-          }
-        case None =>
-          warn(s"While recording the replica LEO, the partition $topicPartition hasn't been created.")
+              readResult.withEmptyFetchInfo
+            }
+          case None =>
+            warn(s"While recording the replica LEO, the partition $topicPartition hasn't been created.")
+            readResult
+        }
       }
       topicPartition -> updatedReadResult
     }
   }
 
   private def leaderPartitionsIterator: Iterator[Partition] =
-    nonOfflinePartitionsIterator.filter(_.leaderReplicaIfLocal.isDefined)
+    nonOfflinePartitionsIterator.filter(_.leaderLogIfLocal.isDefined)
 
   def getLogEndOffset(topicPartition: TopicPartition): Option[Long] =
-    nonOfflinePartition(topicPartition).flatMap(_.leaderReplicaIfLocal.map(_.logEndOffset))
+    nonOfflinePartition(topicPartition).flatMap(_.leaderLogIfLocal.map(_.logEndOffset))
 
   // Flushes the highwatermark value for all partitions to the highwatermark file
-  def checkpointHighWatermarks() {
-    val replicas = nonOfflinePartitionsIterator.flatMap { partition =>
-      val replicasList: mutable.Set[Replica] = mutable.Set()
-      partition.localReplica.foreach(replicasList.add)
-      partition.futureLocalReplica.foreach(replicasList.add)
-      replicasList
-    }.filter(_.log.isDefined).toBuffer
-    val replicasByDir = replicas.groupBy(_.log.get.dir.getParent)
-    for ((dir, reps) <- replicasByDir) {
-      val hwms = reps.map(r => r.topicPartition -> r.highWatermark.messageOffset).toMap
+  def checkpointHighWatermarks(): Unit = {
+    def putHw(logDirToCheckpoints: mutable.AnyRefMap[String, mutable.AnyRefMap[TopicPartition, Long]],
+              log: Log): Unit = {
+      val checkpoints = logDirToCheckpoints.getOrElseUpdate(log.parentDir,
+        new mutable.AnyRefMap[TopicPartition, Long]())
+      checkpoints.put(log.topicPartition, log.highWatermark)
+    }
+
+    val logDirToHws = new mutable.AnyRefMap[String, mutable.AnyRefMap[TopicPartition, Long]](
+      allPartitions.size)
+    nonOfflinePartitionsIterator.foreach { partition =>
+      partition.log.foreach(putHw(logDirToHws, _))
+      partition.futureLog.foreach(putHw(logDirToHws, _))
+    }
+
+    for ((logDir, hws) <- logDirToHws) {
       try {
-        highWatermarkCheckpoints.get(dir).foreach(_.write(hwms))
+        highWatermarkCheckpoints.get(logDir).foreach(_.write(hws))
       } catch {
         case e: KafkaStorageException =>
-          error(s"Error while writing to highwatermark file in directory $dir", e)
+          error(s"Error while writing to highwatermark file in directory $logDir", e)
       }
     }
   }
 
   // Used only by test
-  def markPartitionOffline(tp: TopicPartition) {
-    allPartitions.put(tp, ReplicaManager.OfflinePartition)
+  def markPartitionOffline(tp: TopicPartition): Unit = replicaStateChangeLock synchronized {
+    allPartitions.put(tp, HostedPartition.Offline)
+    Partition.removeMetrics(tp)
   }
 
   // logDir should be an absolute path
   // sendZkNotification is needed for unit test
-  def handleLogDirFailure(dir: String, sendZkNotification: Boolean = true) {
+  def handleLogDirFailure(dir: String, sendZkNotification: Boolean = true): Unit = {
     if (!logManager.isLogDirOnline(dir))
       return
     info(s"Stopping serving replicas in dir $dir")
     replicaStateChangeLock synchronized {
       val newOfflinePartitions = nonOfflinePartitionsIterator.filter { partition =>
-        partition.localReplica.exists { replica =>
-          replica.log.isDefined && replica.log.get.dir.getParent == dir
-        }
+        partition.log.exists { _.parentDir == dir }
       }.map(_.topicPartition).toSet
 
       val partitionsWithOfflineFutureReplica = nonOfflinePartitionsIterator.filter { partition =>
-        partition.futureLocalReplica.exists { replica =>
-          replica.log.isDefined && replica.log.get.dir.getParent == dir
-        }
+        partition.futureLog.exists { _.parentDir == dir }
       }.toSet
 
       replicaFetcherManager.removeFetcherForPartitions(newOfflinePartitions)
@@ -1456,15 +1687,12 @@ class ReplicaManager(val config: KafkaConfig,
 
       partitionsWithOfflineFutureReplica.foreach(partition => partition.removeFutureLocalReplica(deleteFromLogDir = false))
       newOfflinePartitions.foreach { topicPartition =>
-        val partition = allPartitions.put(topicPartition, ReplicaManager.OfflinePartition)
-        partition.removePartitionMetrics()
+        markPartitionOffline(topicPartition)
       }
       newOfflinePartitions.map(_.topic).foreach { topic: String =>
-        val topicHasPartitions = allPartitions.values.exists(partition => topic == partition.topic)
-        if (!topicHasPartitions)
-          brokerTopicStats.removeMetrics(topic)
+        maybeRemoveTopicMetrics(topic)
       }
-      highWatermarkCheckpoints = highWatermarkCheckpoints.filterKeys(_ != dir)
+      highWatermarkCheckpoints = highWatermarkCheckpoints.filter { case (checkpointDir, _) => checkpointDir != dir }
 
       info(s"Broker $localBrokerId stopped fetcher for partitions ${newOfflinePartitions.mkString(",")} and stopped moving logs " +
            s"for partitions ${partitionsWithOfflineFutureReplica.mkString(",")} because they are in the failed log directory $dir.")
@@ -1476,16 +1704,18 @@ class ReplicaManager(val config: KafkaConfig,
     info(s"Stopped serving replicas in dir $dir")
   }
 
-  def removeMetrics() {
+  def removeMetrics(): Unit = {
     removeMetric("LeaderCount")
     removeMetric("PartitionCount")
     removeMetric("OfflineReplicaCount")
     removeMetric("UnderReplicatedPartitions")
     removeMetric("UnderMinIsrPartitionCount")
+    removeMetric("AtMinIsrPartitionCount")
+    removeMetric("recompressedBatch")
   }
 
   // High watermark do not need to be checkpointed only when under unit tests
-  def shutdown(checkpointHW: Boolean = true) {
+  def shutdown(checkpointHW: Boolean = true): Unit = {
     info("Shutting down")
     removeMetrics()
     if (logDirFailureHandler != null)
@@ -1495,9 +1725,10 @@ class ReplicaManager(val config: KafkaConfig,
     delayedFetchPurgatory.shutdown()
     delayedProducePurgatory.shutdown()
     delayedDeleteRecordsPurgatory.shutdown()
-    delayedElectPreferredLeaderPurgatory.shutdown()
+    delayedElectLeaderPurgatory.shutdown()
     if (checkpointHW)
       checkpointHighWatermarks()
+    replicaSelectorOpt.foreach(_.close)
     info("Shut down completely")
   }
 
@@ -1509,49 +1740,73 @@ class ReplicaManager(val config: KafkaConfig,
     new ReplicaAlterLogDirsManager(config, this, quotaManager, brokerTopicStats)
   }
 
+  protected def createReplicaSelector(): Option[ReplicaSelector] = {
+    config.replicaSelectorClassName.map { className =>
+      val tmpReplicaSelector: ReplicaSelector = CoreUtils.createObject[ReplicaSelector](className)
+      tmpReplicaSelector.configure(config.originals())
+      tmpReplicaSelector
+    }
+  }
+
   def lastOffsetForLeaderEpoch(requestedEpochInfo: Map[TopicPartition, OffsetsForLeaderEpochRequest.PartitionData]): Map[TopicPartition, EpochEndOffset] = {
     requestedEpochInfo.map { case (tp, partitionData) =>
       val epochEndOffset = getPartition(tp) match {
-        case Some(partition) =>
-          if (partition eq ReplicaManager.OfflinePartition)
-            new EpochEndOffset(Errors.KAFKA_STORAGE_ERROR, UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET)
-          else
-            partition.lastOffsetForLeaderEpoch(partitionData.currentLeaderEpoch, partitionData.leaderEpoch,
-              fetchOnlyFromLeader = true)
-
-        case None if metadataCache.contains(tp) =>
+        case HostedPartition.Online(partition) =>
+          partition.lastOffsetForLeaderEpoch(partitionData.currentLeaderEpoch, partitionData.leaderEpoch,
+            fetchOnlyFromLeader = true)
+
+        case HostedPartition.Offline =>
+          new EpochEndOffset(Errors.KAFKA_STORAGE_ERROR, UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET)
+
+        case HostedPartition.None if metadataCache.contains(tp) =>
           new EpochEndOffset(Errors.NOT_LEADER_FOR_PARTITION, UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET)
 
-        case None =>
+        case HostedPartition.None =>
           new EpochEndOffset(Errors.UNKNOWN_TOPIC_OR_PARTITION, UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET)
       }
       tp -> epochEndOffset
     }
   }
 
-  def electPreferredLeaders(controller: KafkaController,
-                            partitions: Set[TopicPartition],
-                            responseCallback: Map[TopicPartition, ApiError] => Unit,
-                            requestTimeout: Long): Unit = {
+  def electLeaders(
+    controller: KafkaController,
+    partitions: Set[TopicPartition],
+    electionType: ElectionType,
+    responseCallback: Map[TopicPartition, ApiError] => Unit,
+    requestTimeout: Int
+  ): Unit = {
 
     val deadline = time.milliseconds() + requestTimeout
 
-    def electionCallback(expectedLeaders: Map[TopicPartition, Int],
-                         results: Map[TopicPartition, ApiError]): Unit = {
+    def electionCallback(results: Map[TopicPartition, Either[ApiError, Int]]): Unit = {
+      val expectedLeaders = mutable.Map.empty[TopicPartition, Int]
+      val failures = mutable.Map.empty[TopicPartition, ApiError]
+      results.foreach {
+        case (partition, Right(leader)) => expectedLeaders += partition -> leader
+        case (partition, Left(error)) => failures += partition -> error
+      }
+
       if (expectedLeaders.nonEmpty) {
-        val watchKeys = expectedLeaders.map{
-          case (tp, leader) => new TopicPartitionOperationKey(tp.topic, tp.partition)
-        }.toSeq
-        delayedElectPreferredLeaderPurgatory.tryCompleteElseWatch(
-          new DelayedElectPreferredLeader(deadline - time.milliseconds(), expectedLeaders, results,
-            this, responseCallback),
-          watchKeys)
+        val watchKeys = expectedLeaders.iterator.map {
+          case (tp, _) => TopicPartitionOperationKey(tp)
+        }.toBuffer
+
+        delayedElectLeaderPurgatory.tryCompleteElseWatch(
+          new DelayedElectLeader(
+            math.max(0, deadline - time.milliseconds()),
+            expectedLeaders,
+            failures,
+            this,
+            responseCallback
+          ),
+          watchKeys
+        )
       } else {
           // There are no partitions actually being elected, so return immediately
-          responseCallback(results)
+          responseCallback(failures)
       }
     }
 
-    controller.electPreferredLeaders(partitions, electionCallback)
+    controller.electLeaders(partitions, electionType, electionCallback)
   }
 }
diff --git a/core/src/main/scala/kafka/server/ReplicationQuotaManager.scala b/core/src/main/scala/kafka/server/ReplicationQuotaManager.scala
index 7835c9dbafe4e..c4024645564cf 100644
--- a/core/src/main/scala/kafka/server/ReplicationQuotaManager.scala
+++ b/core/src/main/scala/kafka/server/ReplicationQuotaManager.scala
@@ -17,13 +17,15 @@
 package kafka.server
 
 import java.util.concurrent.{ConcurrentHashMap, TimeUnit}
+import java.util.concurrent.locks.ReentrantReadWriteLock
+
+import scala.collection.Seq
 
 import kafka.server.Constants._
 import kafka.server.ReplicationQuotaManagerConfig._
 import kafka.utils.CoreUtils._
 import kafka.utils.Logging
 import org.apache.kafka.common.metrics._
-import java.util.concurrent.locks.ReentrantReadWriteLock
 
 import org.apache.kafka.common.TopicPartition
 import org.apache.kafka.common.metrics.stats.SimpleRate
@@ -84,7 +86,7 @@ class ReplicationQuotaManager(val config: ReplicationQuotaManagerConfig,
     *
     * @param quota
     */
-  def updateQuota(quota: Quota) {
+  def updateQuota(quota: Quota): Unit = {
     inWriteLock(lock) {
       this.quota = quota
       //The metric could be expired by another thread, so use a local variable and null check.
@@ -130,7 +132,7 @@ class ReplicationQuotaManager(val config: ReplicationQuotaManagerConfig,
     *
     * @param value
     */
-  def record(value: Long) {
+  def record(value: Long): Unit = {
     try {
       sensor().record(value)
     } catch {
@@ -147,7 +149,7 @@ class ReplicationQuotaManager(val config: ReplicationQuotaManagerConfig,
     * @param partitions the set of throttled partitions
     * @return
     */
-  def markThrottled(topic: String, partitions: Seq[Int]) {
+  def markThrottled(topic: String, partitions: Seq[Int]): Unit = {
     throttledPartitions.put(topic, partitions)
   }
 
@@ -157,7 +159,7 @@ class ReplicationQuotaManager(val config: ReplicationQuotaManagerConfig,
     * @param topic
     * @return
     */
-  def markThrottled(topic: String) {
+  def markThrottled(topic: String): Unit = {
     markThrottled(topic, AllReplicas)
   }
 
@@ -167,7 +169,7 @@ class ReplicationQuotaManager(val config: ReplicationQuotaManagerConfig,
     * @param topic
     * @return
     */
-  def removeThrottle(topic: String) {
+  def removeThrottle(topic: String): Unit = {
     throttledPartitions.remove(topic)
   }
 
diff --git a/core/src/main/scala/kafka/server/checkpoints/CheckpointFile.scala b/core/src/main/scala/kafka/server/checkpoints/CheckpointFile.scala
index 1878ae263f034..8a05775e0869c 100644
--- a/core/src/main/scala/kafka/server/checkpoints/CheckpointFile.scala
+++ b/core/src/main/scala/kafka/server/checkpoints/CheckpointFile.scala
@@ -45,7 +45,7 @@ class CheckpointFile[T](val file: File,
   try Files.createFile(file.toPath) // create the file if it doesn't exist
   catch { case _: FileAlreadyExistsException => }
 
-  def write(entries: Seq[T]) {
+  def write(entries: Iterable[T]): Unit = {
     lock synchronized {
       try {
         // write to temp file and then swap with the existing file
diff --git a/core/src/main/scala/kafka/server/checkpoints/LeaderEpochCheckpointFile.scala b/core/src/main/scala/kafka/server/checkpoints/LeaderEpochCheckpointFile.scala
index 79a11c7dcc893..6b88bebf6b4a4 100644
--- a/core/src/main/scala/kafka/server/checkpoints/LeaderEpochCheckpointFile.scala
+++ b/core/src/main/scala/kafka/server/checkpoints/LeaderEpochCheckpointFile.scala
@@ -25,7 +25,7 @@ import kafka.server.epoch.EpochEntry
 import scala.collection._
 
 trait LeaderEpochCheckpoint {
-  def write(epochs: Seq[EpochEntry])
+  def write(epochs: Seq[EpochEntry]): Unit
   def read(): Seq[EpochEntry]
 }
 
diff --git a/core/src/main/scala/kafka/server/checkpoints/OffsetCheckpointFile.scala b/core/src/main/scala/kafka/server/checkpoints/OffsetCheckpointFile.scala
index 2769cb4240a52..fbe2f6e2bb036 100644
--- a/core/src/main/scala/kafka/server/checkpoints/OffsetCheckpointFile.scala
+++ b/core/src/main/scala/kafka/server/checkpoints/OffsetCheckpointFile.scala
@@ -45,7 +45,7 @@ object OffsetCheckpointFile {
 }
 
 trait OffsetCheckpoint {
-  def write(epochs: Seq[EpochEntry])
+  def write(epochs: Seq[EpochEntry]): Unit
   def read(): Seq[EpochEntry]
 }
 
@@ -56,8 +56,36 @@ class OffsetCheckpointFile(val file: File, logDirFailureChannel: LogDirFailureCh
   val checkpoint = new CheckpointFile[(TopicPartition, Long)](file, OffsetCheckpointFile.CurrentVersion,
     OffsetCheckpointFile.Formatter, logDirFailureChannel, file.getParent)
 
-  def write(offsets: Map[TopicPartition, Long]): Unit = checkpoint.write(offsets.toSeq)
+  def write(offsets: Map[TopicPartition, Long]): Unit = checkpoint.write(offsets)
 
   def read(): Map[TopicPartition, Long] = checkpoint.read().toMap
 
 }
+
+trait OffsetCheckpoints {
+  def fetch(logDir: String, topicPartition: TopicPartition): Option[Long]
+}
+
+/**
+ * Loads checkpoint files on demand and caches the offsets for reuse.
+ */
+class LazyOffsetCheckpoints(checkpointsByLogDir: Map[String, OffsetCheckpointFile]) extends OffsetCheckpoints {
+  private val lazyCheckpointsByLogDir = checkpointsByLogDir.map { case (logDir, checkpointFile) =>
+    logDir -> new LazyOffsetCheckpointMap(checkpointFile)
+  }.toMap
+
+  override def fetch(logDir: String, topicPartition: TopicPartition): Option[Long] = {
+    val offsetCheckpointFile = lazyCheckpointsByLogDir.getOrElse(logDir,
+      throw new IllegalArgumentException(s"No checkpoint file for log dir $logDir"))
+    offsetCheckpointFile.fetch(topicPartition)
+  }
+}
+
+class LazyOffsetCheckpointMap(checkpoint: OffsetCheckpointFile) {
+  private lazy val offsets: Map[TopicPartition, Long] = checkpoint.read()
+
+  def fetch(topicPartition: TopicPartition): Option[Long] = {
+    offsets.get(topicPartition)
+  }
+
+}
diff --git a/core/src/main/scala/kafka/server/epoch/LeaderEpochFileCache.scala b/core/src/main/scala/kafka/server/epoch/LeaderEpochFileCache.scala
index 0c885b71c7c7b..9d14ae6f2bd09 100644
--- a/core/src/main/scala/kafka/server/epoch/LeaderEpochFileCache.scala
+++ b/core/src/main/scala/kafka/server/epoch/LeaderEpochFileCache.scala
@@ -24,7 +24,8 @@ import kafka.utils.CoreUtils._
 import kafka.utils.Logging
 import org.apache.kafka.common.TopicPartition
 
-import scala.collection.mutable.ListBuffer
+import scala.collection.Seq
+import scala.collection.mutable.ArrayBuffer
 
 /**
  * Represents a cache of (LeaderEpoch => Offset) mappings for a particular replica.
@@ -42,7 +43,10 @@ class LeaderEpochFileCache(topicPartition: TopicPartition,
   this.logIdent = s"[LeaderEpochCache $topicPartition] "
 
   private val lock = new ReentrantReadWriteLock()
-  private var epochs: ListBuffer[EpochEntry] = inWriteLock(lock) { ListBuffer(checkpoint.read(): _*) }
+  private var epochs: ArrayBuffer[EpochEntry] = inWriteLock(lock) {
+    val read = checkpoint.read()
+    new ArrayBuffer(read.size) ++= read
+  }
 
   /**
     * Assigns the supplied Leader Epoch to the supplied Offset
@@ -78,7 +82,10 @@ class LeaderEpochFileCache(topicPartition: TopicPartition,
 
     if (removedEpochs.isEmpty) {
       debug(s"Appended new epoch entry $entryToAppend. Cache now contains ${epochs.size} entries.")
-    } else {
+    } else if (removedEpochs.size > 1 || removedEpochs.head.startOffset != entryToAppend.startOffset) {
+      // Only log a warning if there were non-trivial removals. If the start offset of the new entry
+      // matches the start offset of the removed epoch, then no data has been written and the truncation
+      // is expected.
       warn(s"New epoch entry $entryToAppend caused truncation of conflicting entries $removedEpochs. " +
         s"Cache now contains ${epochs.size} entries.")
     }
@@ -116,7 +123,7 @@ class LeaderEpochFileCache(topicPartition: TopicPartition,
     * Offset if the latest epoch was requested.
     *
     * During the upgrade phase, where there are existing messages may not have a leader epoch,
-    * if requestedEpoch is < the first epoch cached, UNSUPPORTED_EPOCH_OFFSET will be returned
+    * if requestedEpoch is < the first epoch cached, UNDEFINED_EPOCH_OFFSET will be returned
     * so that the follower falls back to High Water Mark.
     *
     * @param requestedEpoch requested leader epoch
@@ -220,7 +227,7 @@ class LeaderEpochFileCache(topicPartition: TopicPartition,
   }
 
   // Visible for testing
-  def epochEntries: ListBuffer[EpochEntry] = {
+  def epochEntries: Seq[EpochEntry] = {
     epochs
   }
 
diff --git a/core/src/main/scala/kafka/tools/ConsoleConsumer.scala b/core/src/main/scala/kafka/tools/ConsoleConsumer.scala
index d246e7b9007b3..22917c9594146 100755
--- a/core/src/main/scala/kafka/tools/ConsoleConsumer.scala
+++ b/core/src/main/scala/kafka/tools/ConsoleConsumer.scala
@@ -48,7 +48,7 @@ object ConsoleConsumer extends Logging {
 
   private val shutdownLatch = new CountDownLatch(1)
 
-  def main(args: Array[String]) {
+  def main(args: Array[String]): Unit = {
     val conf = new ConsumerConfig(args)
     try {
       run(conf)
@@ -62,7 +62,7 @@ object ConsoleConsumer extends Logging {
     }
   }
 
-  def run(conf: ConsumerConfig) {
+  def run(conf: ConsumerConfig): Unit = {
     val timeoutMs = if (conf.timeoutMs >= 0) conf.timeoutMs else Long.MaxValue
     val consumer = new KafkaConsumer(consumerProps(conf), new ByteArrayDeserializer, new ByteArrayDeserializer)
 
@@ -84,9 +84,9 @@ object ConsoleConsumer extends Logging {
     }
   }
 
-  def addShutdownHook(consumer: ConsumerWrapper, conf: ConsumerConfig) {
+  def addShutdownHook(consumer: ConsumerWrapper, conf: ConsumerConfig): Unit = {
     Runtime.getRuntime.addShutdownHook(new Thread() {
-      override def run() {
+      override def run(): Unit = {
         consumer.wakeup()
 
         shutdownLatch.await()
@@ -99,7 +99,7 @@ object ConsoleConsumer extends Logging {
   }
 
   def process(maxMessages: Integer, formatter: MessageFormatter, consumer: ConsumerWrapper, output: PrintStream,
-              skipMessageOnError: Boolean) {
+              skipMessageOnError: Boolean): Unit = {
     while (messageCount < maxMessages || maxMessages == -1) {
       val msg: ConsumerRecord[Array[Byte], Array[Byte]] = try {
         consumer.receive()
@@ -133,7 +133,7 @@ object ConsoleConsumer extends Logging {
     }
   }
 
-  def reportRecordCount() {
+  def reportRecordCount(): Unit = {
     System.err.println(s"Processed a total of $messageCount messages")
   }
 
@@ -168,7 +168,7 @@ object ConsoleConsumer extends Logging {
     * 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) {
+  def setAutoOffsetResetValue(config: ConsumerConfig, props: Properties): Unit = {
     val (earliestConfigValue, latestConfigValue) = ("earliest", "latest")
 
     if (props.containsKey(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG)) {
@@ -394,7 +394,7 @@ object ConsoleConsumer extends Logging {
     consumerInit()
     var recordIter = Collections.emptyList[ConsumerRecord[Array[Byte], Array[Byte]]]().iterator()
 
-    def consumerInit() {
+    def consumerInit(): Unit = {
       (topic, partitionId, offset, whitelist) match {
         case (Some(topic), Some(partitionId), Some(offset), None) =>
           seek(topic, partitionId, offset)
@@ -413,7 +413,7 @@ object ConsoleConsumer extends Logging {
       }
     }
 
-    def seek(topic: String, partitionId: Int, offset: Long) {
+    def seek(topic: String, partitionId: Int, offset: Long): Unit = {
       val topicPartition = new TopicPartition(topic, partitionId)
       consumer.assign(Collections.singletonList(topicPartition))
       offset match {
@@ -423,7 +423,7 @@ object ConsoleConsumer extends Logging {
       }
     }
 
-    def resetUnconsumedOffsets() {
+    def resetUnconsumedOffsets(): Unit = {
       val smallestUnconsumedOffsets = collection.mutable.Map[TopicPartition, Long]()
       while (recordIter.hasNext) {
         val record = recordIter.next()
@@ -448,7 +448,7 @@ object ConsoleConsumer extends Logging {
       this.consumer.wakeup()
     }
 
-    def cleanup() {
+    def cleanup(): Unit = {
       resetUnconsumedOffsets()
       this.consumer.close()
     }
@@ -466,7 +466,7 @@ class DefaultMessageFormatter extends MessageFormatter {
   var keyDeserializer: Option[Deserializer[_]] = None
   var valueDeserializer: Option[Deserializer[_]] = None
 
-  override def init(props: Properties) {
+  override def init(props: Properties): Unit = {
     if (props.containsKey("print.timestamp"))
       printTimestamp = props.getProperty("print.timestamp").trim.equalsIgnoreCase("true")
     if (props.containsKey("print.key"))
@@ -500,7 +500,7 @@ class DefaultMessageFormatter extends MessageFormatter {
     newProps
   }
 
-  def writeTo(consumerRecord: ConsumerRecord[Array[Byte], Array[Byte]], output: PrintStream) {
+  def writeTo(consumerRecord: ConsumerRecord[Array[Byte], Array[Byte]], output: PrintStream): Unit = {
 
     def writeSeparator(columnSeparator: Boolean): Unit = {
       if (columnSeparator)
@@ -509,7 +509,7 @@ class DefaultMessageFormatter extends MessageFormatter {
         output.write(lineSeparator)
     }
 
-    def write(deserializer: Option[Deserializer[_]], sourceBytes: Array[Byte], topic: String) {
+    def write(deserializer: Option[Deserializer[_]], sourceBytes: Array[Byte], topic: String): Unit = {
       val nonNullBytes = Option(sourceBytes).getOrElse("null".getBytes(StandardCharsets.UTF_8))
       val convertedBytes = deserializer.map(_.deserialize(topic, nonNullBytes).toString.
         getBytes(StandardCharsets.UTF_8)).getOrElse(nonNullBytes)
@@ -553,15 +553,15 @@ class LoggingMessageFormatter extends MessageFormatter with LazyLogging {
 }
 
 class NoOpMessageFormatter extends MessageFormatter {
-  override def init(props: Properties) {}
+  override def init(props: Properties): Unit = {}
 
-  def writeTo(consumerRecord: ConsumerRecord[Array[Byte], Array[Byte]], output: PrintStream){}
+  def writeTo(consumerRecord: ConsumerRecord[Array[Byte], Array[Byte]], output: PrintStream): Unit = {}
 }
 
 class ChecksumMessageFormatter extends MessageFormatter {
   private var topicStr: String = _
 
-  override def init(props: Properties) {
+  override def init(props: Properties): Unit = {
     topicStr = props.getProperty("topic")
     if (topicStr != null)
       topicStr = topicStr + ":"
@@ -569,7 +569,7 @@ class ChecksumMessageFormatter extends MessageFormatter {
       topicStr = ""
   }
 
-  def writeTo(consumerRecord: ConsumerRecord[Array[Byte], Array[Byte]], output: PrintStream) {
+  def writeTo(consumerRecord: ConsumerRecord[Array[Byte], Array[Byte]], output: PrintStream): Unit = {
     output.println(topicStr + "checksum:" + consumerRecord.checksum)
   }
 }
diff --git a/core/src/main/scala/kafka/tools/ConsoleProducer.scala b/core/src/main/scala/kafka/tools/ConsoleProducer.scala
index 829e2717813bb..16f6af89b65ad 100644
--- a/core/src/main/scala/kafka/tools/ConsoleProducer.scala
+++ b/core/src/main/scala/kafka/tools/ConsoleProducer.scala
@@ -21,6 +21,7 @@ import java.io._
 import java.nio.charset.StandardCharsets
 import java.util.Properties
 
+import joptsimple.{OptionException, OptionParser, OptionSet}
 import kafka.common._
 import kafka.message._
 import kafka.utils.Implicits._
@@ -44,7 +45,7 @@ object ConsoleProducer {
         val producer = new KafkaProducer[Array[Byte], Array[Byte]](producerProps(config))
 
         Runtime.getRuntime.addShutdownHook(new Thread() {
-          override def run() {
+          override def run(): Unit = {
             producer.close()
           }
         })
@@ -213,7 +214,7 @@ object ConsoleProducer {
       .describedAs("config file")
       .ofType(classOf[String])
 
-    options = parser.parse(args : _*)
+    options = tryParse(parser, args)
 
     CommandLineUtils.printHelpAndExitIfNeeded(this, "This tool helps to read data from standard input and publish it to Kafka.")
     CommandLineUtils.checkRequiredArgs(parser, options, topicOpt, brokerListOpt)
@@ -232,6 +233,15 @@ object ConsoleProducer {
     val readerClass = options.valueOf(messageReaderOpt)
     val cmdLineProps = CommandLineUtils.parseKeyValueArgs(options.valuesOf(propertyOpt).asScala)
     val extraProducerProps = CommandLineUtils.parseKeyValueArgs(options.valuesOf(producerPropertyOpt).asScala)
+
+    def tryParse(parser: OptionParser, args: Array[String]): OptionSet = {
+      try
+        parser.parse(args: _*)
+      catch {
+        case e: OptionException =>
+          CommandLineUtils.printUsageAndDie(parser, e.getMessage)
+      }
+    }
   }
 
   class LineMessageReader extends MessageReader {
@@ -242,7 +252,7 @@ object ConsoleProducer {
     var ignoreError = false
     var lineNumber = 0
 
-    override def init(inputStream: InputStream, props: Properties) {
+    override def init(inputStream: InputStream, props: Properties): Unit = {
       topic = props.getProperty("topic")
       if (props.containsKey("parse.key"))
         parseKey = props.getProperty("parse.key").trim.equalsIgnoreCase("true")
diff --git a/core/src/main/scala/kafka/tools/ConsumerPerformance.scala b/core/src/main/scala/kafka/tools/ConsumerPerformance.scala
index 226341e4122bc..278d69486b012 100644
--- a/core/src/main/scala/kafka/tools/ConsumerPerformance.scala
+++ b/core/src/main/scala/kafka/tools/ConsumerPerformance.scala
@@ -100,7 +100,7 @@ object ConsumerPerformance extends LazyLogging {
               totalMessagesRead: AtomicLong,
               totalBytesRead: AtomicLong,
               joinTime: AtomicLong,
-              testStartTime: Long) {
+              testStartTime: Long): Unit = {
     var bytesRead = 0L
     var messagesRead = 0L
     var lastBytesRead = 0L
@@ -109,11 +109,11 @@ object ConsumerPerformance extends LazyLogging {
     var joinTimeMsInSingleRound = 0L
 
     consumer.subscribe(topics.asJava, new ConsumerRebalanceListener {
-      def onPartitionsAssigned(partitions: util.Collection[TopicPartition]) {
+      def onPartitionsAssigned(partitions: util.Collection[TopicPartition]): Unit = {
         joinTime.addAndGet(System.currentTimeMillis - joinStart)
         joinTimeMsInSingleRound += System.currentTimeMillis - joinStart
       }
-      def onPartitionsRevoked(partitions: util.Collection[TopicPartition]) {
+      def onPartitionsRevoked(partitions: util.Collection[TopicPartition]): Unit = {
         joinStart = System.currentTimeMillis
       }})
 
diff --git a/core/src/main/scala/kafka/tools/DumpLogSegments.scala b/core/src/main/scala/kafka/tools/DumpLogSegments.scala
index 5998902245ce3..5af4a473346dc 100755
--- a/core/src/main/scala/kafka/tools/DumpLogSegments.scala
+++ b/core/src/main/scala/kafka/tools/DumpLogSegments.scala
@@ -18,28 +18,25 @@
 package kafka.tools
 
 import java.io._
-import java.nio.ByteBuffer
 
-import kafka.coordinator.group.{GroupMetadataKey, GroupMetadataManager, OffsetKey}
+import kafka.coordinator.group.GroupMetadataManager
 import kafka.coordinator.transaction.TransactionLog
 import kafka.log._
 import kafka.serializer.Decoder
 import kafka.utils._
-import org.apache.kafka.clients.consumer.internals.ConsumerProtocol
-import org.apache.kafka.common.KafkaException
 import org.apache.kafka.common.record._
-import org.apache.kafka.common.utils.{Time, Utils}
+import org.apache.kafka.common.utils.Utils
 
-import scala.collection.{Map, mutable}
-import scala.collection.mutable.ArrayBuffer
 import scala.collection.JavaConverters._
+import scala.collection.mutable
+import scala.collection.mutable.ArrayBuffer
 
 object DumpLogSegments {
 
   // visible for testing
   private[tools] val RecordIndent = "|"
 
-  def main(args: Array[String]) {
+  def main(args: Array[String]): Unit = {
     val opts = new DumpLogSegmentsOptions(args)
     CommandLineUtils.printHelpAndExitIfNeeded(opts, "This tool helps to parse a log file and dump its contents to the console, useful for debugging a seemingly corrupt log segment.")
     opts.checkArgs()
@@ -119,7 +116,7 @@ object DumpLogSegments {
                                indexSanityOnly: Boolean,
                                verifyOnly: Boolean,
                                misMatchesForIndexFilesMap: mutable.Map[String, List[(Long, Long)]],
-                               maxMessageSize: Int) {
+                               maxMessageSize: Int): Unit = {
     val startOffset = file.getName.split("\\.")(0).toLong
     val logFile = new File(file.getAbsoluteFile.getParent, file.getName.split("\\.")(0) + Log.LogFileSuffix)
     val fileRecords = FileRecords.open(logFile, false)
@@ -156,7 +153,7 @@ object DumpLogSegments {
                                    indexSanityOnly: Boolean,
                                    verifyOnly: Boolean,
                                    timeIndexDumpErrors: TimeIndexDumpErrors,
-                                   maxMessageSize: Int) {
+                                   maxMessageSize: Int): Unit = {
     val startOffset = file.getName.split("\\.")(0).toLong
     val logFile = new File(file.getAbsoluteFile.getParent, file.getName.split("\\.")(0) + Log.LogFileSuffix)
     val fileRecords = FileRecords.open(logFile, false)
@@ -212,20 +209,20 @@ object DumpLogSegments {
     }
   }
 
-  private trait MessageParser[K, V] {
+  private[kafka] trait MessageParser[K, V] {
     def parse(record: Record): (Option[K], Option[V])
   }
 
   private class DecoderMessageParser[K, V](keyDecoder: Decoder[K], valueDecoder: Decoder[V]) extends MessageParser[K, V] {
     override def parse(record: Record): (Option[K], Option[V]) = {
+      val key = if (record.hasKey)
+        Some(keyDecoder.fromBytes(Utils.readBytes(record.key)))
+      else
+        None
+
       if (!record.hasValue) {
-        (None, None)
+        (key, None)
       } else {
-        val key = if (record.hasKey)
-          Some(keyDecoder.fromBytes(Utils.readBytes(record.key)))
-        else
-          None
-
         val payload = Some(valueDecoder.fromBytes(Utils.readBytes(record.value)))
 
         (key, payload)
@@ -233,100 +230,13 @@ object DumpLogSegments {
     }
   }
 
-  private class TransactionLogMessageParser extends MessageParser[String, String] {
-
-    override def parse(record: Record): (Option[String], Option[String]) = {
-      val txnKey = TransactionLog.readTxnRecordKey(record.key)
-      val txnMetadata = TransactionLog.readTxnRecordValue(txnKey.transactionalId, record.value)
-
-      val keyString = s"transactionalId=${txnKey.transactionalId}"
-      val valueString = s"producerId:${txnMetadata.producerId}," +
-        s"producerEpoch:${txnMetadata.producerEpoch}," +
-        s"state=${txnMetadata.state}," +
-        s"partitions=${txnMetadata.topicPartitions}," +
-        s"txnLastUpdateTimestamp=${txnMetadata.txnLastUpdateTimestamp}," +
-        s"txnTimeoutMs=${txnMetadata.txnTimeoutMs}"
-
-      (Some(keyString), Some(valueString))
-    }
-
-  }
-
-  private class OffsetsMessageParser extends MessageParser[String, String] {
-    private def hex(bytes: Array[Byte]): String = {
-      if (bytes.isEmpty)
-        ""
-      else
-        "%X".format(BigInt(1, bytes))
-    }
-
-    private def parseOffsets(offsetKey: OffsetKey, payload: ByteBuffer) = {
-      val group = offsetKey.key.group
-      val topicPartition = offsetKey.key.topicPartition
-      val offset = GroupMetadataManager.readOffsetMessageValue(payload)
-
-      val keyString = s"offset::$group:${topicPartition.topic}:${topicPartition.partition}"
-      val valueString = if (offset.metadata.isEmpty)
-        String.valueOf(offset.offset)
-      else
-        s"${offset.offset}:${offset.metadata}"
-
-      (Some(keyString), Some(valueString))
-    }
-
-    private def parseGroupMetadata(groupMetadataKey: GroupMetadataKey, payload: ByteBuffer) = {
-      val groupId = groupMetadataKey.key
-      val group = GroupMetadataManager.readGroupMessageValue(groupId, payload, Time.SYSTEM)
-      val protocolType = group.protocolType.getOrElse("")
-
-      val assignment = group.allMemberMetadata.map { member =>
-        if (protocolType == ConsumerProtocol.PROTOCOL_TYPE) {
-          val partitionAssignment = ConsumerProtocol.deserializeAssignment(ByteBuffer.wrap(member.assignment))
-          val userData = hex(Utils.toArray(partitionAssignment.userData()))
-
-          if (userData.isEmpty)
-            s"${member.memberId}=${partitionAssignment.partitions()}"
-          else
-            s"${member.memberId}=${partitionAssignment.partitions()}:$userData"
-        } else {
-          s"${member.memberId}=${hex(member.assignment)}"
-        }
-      }.mkString("{", ",", "}")
-
-      val keyString = Json.encodeAsString(Map("metadata" -> groupId).asJava)
-
-      val valueString = Json.encodeAsString(Map(
-        "protocolType" -> protocolType,
-        "protocol" -> group.protocolOrNull,
-        "generationId" -> group.generationId,
-        "assignment" -> assignment
-      ).asJava)
-
-      (Some(keyString), Some(valueString))
-    }
-
-    override def parse(record: Record): (Option[String], Option[String]) = {
-      if (!record.hasValue)
-        (None, None)
-      else if (!record.hasKey) {
-        throw new KafkaException("Failed to decode message using offset topic decoder (message had a missing key)")
-      } else {
-        GroupMetadataManager.readMessageKey(record.key) match {
-          case offsetKey: OffsetKey => parseOffsets(offsetKey, record.value)
-          case groupMetadataKey: GroupMetadataKey => parseGroupMetadata(groupMetadataKey, record.value)
-          case _ => throw new KafkaException("Failed to decode message using offset topic decoder (message had an invalid key)")
-        }
-      }
-    }
-  }
-
   /* print out the contents of the log */
   private def dumpLog(file: File,
                       printContents: Boolean,
                       nonConsecutivePairsForLogFilesMap: mutable.Map[String, List[(Long, Long)]],
                       isDeepIteration: Boolean,
                       maxMessageSize: Int,
-                      parser: MessageParser[_, _]) {
+                      parser: MessageParser[_, _]): Unit = {
     val startOffset = file.getName.split("\\.")(0).toLong
     println("Starting offset: " + startOffset)
     val fileRecords = FileRecords.open(file, false)
@@ -401,28 +311,28 @@ object DumpLogSegments {
     val outOfOrderTimestamp = mutable.Map[String, ArrayBuffer[(Long, Long)]]()
     val shallowOffsetNotFound = mutable.Map[String, ArrayBuffer[(Long, Long)]]()
 
-    def recordMismatchTimeIndex(file: File, indexTimestamp: Long, logTimestamp: Long) {
+    def recordMismatchTimeIndex(file: File, indexTimestamp: Long, logTimestamp: Long): Unit = {
       val misMatchesSeq = misMatchesForTimeIndexFilesMap.getOrElse(file.getAbsolutePath, new ArrayBuffer[(Long, Long)]())
       if (misMatchesSeq.isEmpty)
         misMatchesForTimeIndexFilesMap.put(file.getAbsolutePath, misMatchesSeq)
       misMatchesSeq += ((indexTimestamp, logTimestamp))
     }
 
-    def recordOutOfOrderIndexTimestamp(file: File, indexTimestamp: Long, prevIndexTimestamp: Long) {
+    def recordOutOfOrderIndexTimestamp(file: File, indexTimestamp: Long, prevIndexTimestamp: Long): Unit = {
       val outOfOrderSeq = outOfOrderTimestamp.getOrElse(file.getAbsolutePath, new ArrayBuffer[(Long, Long)]())
       if (outOfOrderSeq.isEmpty)
         outOfOrderTimestamp.put(file.getAbsolutePath, outOfOrderSeq)
       outOfOrderSeq += ((indexTimestamp, prevIndexTimestamp))
     }
 
-    def recordShallowOffsetNotFound(file: File, indexOffset: Long, logOffset: Long) {
+    def recordShallowOffsetNotFound(file: File, indexOffset: Long, logOffset: Long): Unit = {
       val shallowOffsetNotFoundSeq = shallowOffsetNotFound.getOrElse(file.getAbsolutePath, new ArrayBuffer[(Long, Long)]())
       if (shallowOffsetNotFoundSeq.isEmpty)
         shallowOffsetNotFound.put(file.getAbsolutePath, shallowOffsetNotFoundSeq)
       shallowOffsetNotFoundSeq += ((indexOffset, logOffset))
     }
 
-    def printErrors() {
+    def printErrors(): Unit = {
       misMatchesForTimeIndexFilesMap.foreach {
         case (fileName, listOfMismatches) => {
           System.err.println("Found timestamp mismatch in :" + fileName)
@@ -450,6 +360,18 @@ object DumpLogSegments {
     }
   }
 
+  private class OffsetsMessageParser extends MessageParser[String, String] {
+    override def parse(record: Record): (Option[String], Option[String]) = {
+      GroupMetadataManager.formatRecordKeyAndValue(record)
+    }
+  }
+
+  private class TransactionLogMessageParser extends MessageParser[String, String] {
+    override def parse(record: Record): (Option[String], Option[String]) = {
+      TransactionLog.formatRecordKeyAndValue(record)
+    }
+  }
+
   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.")
diff --git a/core/src/main/scala/kafka/tools/EndToEndLatency.scala b/core/src/main/scala/kafka/tools/EndToEndLatency.scala
index 810758446af53..e97be2f30c62b 100755
--- a/core/src/main/scala/kafka/tools/EndToEndLatency.scala
+++ b/core/src/main/scala/kafka/tools/EndToEndLatency.scala
@@ -49,7 +49,7 @@ object EndToEndLatency {
   private val defaultReplicationFactor: Short = 1
   private val defaultNumPartitions: Int = 1
 
-  def main(args: Array[String]) {
+  def main(args: Array[String]): Unit = {
     if (args.length != 5 && args.length != 6) {
       System.err.println("USAGE: java " + getClass.getName + " broker_list topic num_messages producer_acks message_size_bytes [optional] properties_file")
       Exit.exit(1)
@@ -88,7 +88,7 @@ object EndToEndLatency {
     producerProps.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.ByteArraySerializer")
     val producer = new KafkaProducer[Array[Byte], Array[Byte]](producerProps)
 
-    def finalise() {
+    def finalise(): Unit = {
       consumer.commitSync()
       producer.close()
       consumer.close()
diff --git a/core/src/main/scala/kafka/tools/GetOffsetShell.scala b/core/src/main/scala/kafka/tools/GetOffsetShell.scala
index 3474974f477a9..b8b2ad60e3776 100644
--- a/core/src/main/scala/kafka/tools/GetOffsetShell.scala
+++ b/core/src/main/scala/kafka/tools/GetOffsetShell.scala
@@ -28,6 +28,7 @@ import org.apache.kafka.common.requests.ListOffsetRequest
 import org.apache.kafka.common.serialization.ByteArrayDeserializer
 
 import scala.collection.JavaConverters._
+import scala.collection.Seq
 
 object GetOffsetShell {
 
@@ -127,7 +128,9 @@ object GetOffsetShell {
       case ListOffsetRequest.LATEST_TIMESTAMP => consumer.endOffsets(topicPartitions.asJava).asScala
       case _ =>
         val timestampsToSearch = topicPartitions.map(tp => tp -> (listOffsetsTimestamp: java.lang.Long)).toMap.asJava
-        consumer.offsetsForTimes(timestampsToSearch).asScala.mapValues(x => if (x == null) null else x.offset)
+        consumer.offsetsForTimes(timestampsToSearch).asScala.map { case (k, x) =>
+          if (x == null) (k, null) else (k, x.offset: java.lang.Long)
+        }
     }
 
     partitionOffsets.toSeq.sortBy { case (tp, _) => tp.partition }.foreach { case (tp, offset) =>
diff --git a/core/src/main/scala/kafka/tools/JmxTool.scala b/core/src/main/scala/kafka/tools/JmxTool.scala
index c5303a9d96123..7b677019a04ef 100644
--- a/core/src/main/scala/kafka/tools/JmxTool.scala
+++ b/core/src/main/scala/kafka/tools/JmxTool.scala
@@ -18,17 +18,18 @@
  */
 package kafka.tools
 
-import java.util.Date
+import java.util.{Date, Objects}
 import java.text.SimpleDateFormat
 import javax.management._
 import javax.management.remote._
+import javax.rmi.ssl.SslRMIClientSocketFactory
 
 import joptsimple.OptionParser
 
 import scala.collection.JavaConverters._
 import scala.collection.mutable
 import scala.math._
-import kafka.utils.{CommandLineUtils , Exit, Logging}
+import kafka.utils.{CommandLineUtils, Exit, Logging}
 
 
 /**
@@ -39,7 +40,7 @@ import kafka.utils.{CommandLineUtils , Exit, Logging}
   */
 object JmxTool extends Logging {
 
-  def main(args: Array[String]) {
+  def main(args: Array[String]): Unit = {
     // Parse command line
     val parser = new OptionParser(false)
     val objectNameOpt =
@@ -82,6 +83,16 @@ object JmxTool extends Logging {
       .describedAs("report-format")
       .ofType(classOf[java.lang.String])
       .defaultsTo("original")
+    val jmxAuthPropOpt = parser.accepts("jmx-auth-prop", "A mechanism to pass property in the form 'username=password' " +
+      "when enabling remote JMX with password authentication.")
+      .withRequiredArg
+      .describedAs("jmx-auth-prop")
+      .ofType(classOf[String])
+    val jmxSslEnableOpt = parser.accepts("jmx-ssl-enable", "Flag to enable remote JMX with SSL.")
+      .withRequiredArg
+      .describedAs("ssl-enable")
+      .ofType(classOf[java.lang.Boolean])
+      .defaultsTo(false)
     val waitOpt = parser.accepts("wait", "Wait for requested JMX objects to become available before starting output. " +
       "Only supported when the list of objects is non-empty and contains no object name patterns.")
     val helpOpt = parser.accepts("help", "Print usage information.")
@@ -109,6 +120,9 @@ object JmxTool extends Logging {
     val reportFormat = parseFormat(options.valueOf(reportFormatOpt).toLowerCase)
     val reportFormatOriginal = reportFormat.equals("original")
 
+    val enablePasswordAuth = options.has(jmxAuthPropOpt)
+    val enableSsl = options.has(jmxSslEnableOpt)
+
     var jmxc: JMXConnector = null
     var mbsc: MBeanServerConnection = null
     var connected = false
@@ -117,7 +131,18 @@ object JmxTool extends Logging {
     do {
       try {
         System.err.println(s"Trying to connect to JMX url: $url.")
-        jmxc = JMXConnectorFactory.connect(url, null)
+        val env = new java.util.HashMap[String, AnyRef]
+        // ssl enable
+        if (enableSsl) {
+          val csf = new SslRMIClientSocketFactory
+          env.put("com.sun.jndi.rmi.factory.socket", csf)
+        }
+        // password authentication enable
+        if (enablePasswordAuth) {
+          val credentials = options.valueOf(jmxAuthPropOpt).split("=", 2)
+          env.put(JMXConnector.CREDENTIALS, credentials)
+        }
+        jmxc = JMXConnectorFactory.connect(url, env)
         mbsc = jmxc.getMBeanServerConnection
         connected = true
       } catch {
@@ -140,7 +165,7 @@ object JmxTool extends Logging {
       else
         List(null)
 
-    val hasPatternQueries = queries.exists((name: ObjectName) => name.isPattern)
+    val hasPatternQueries = queries.filterNot(Objects.isNull).exists((name: ObjectName) => name.isPattern)
 
     var names: Iterable[ObjectName] = null
     def namesSet = Option(names).toSet.flatten
@@ -165,12 +190,20 @@ object JmxTool extends Logging {
     }
 
     val numExpectedAttributes: Map[ObjectName, Int] =
-      if (attributesWhitelistExists)
-        queries.map((_, attributesWhitelist.get.length)).toMap
-      else {
-        names.map{(name: ObjectName) =>
+      if (!attributesWhitelistExists)
+        names.map{name: ObjectName =>
           val mbean = mbsc.getMBeanInfo(name)
           (name, mbsc.getAttributes(name, mbean.getAttributes.map(_.getName)).size)}.toMap
+      else {
+        if (!hasPatternQueries)
+          names.map{name: ObjectName =>
+            val mbean = mbsc.getMBeanInfo(name)
+            val attributes = mbsc.getAttributes(name, mbean.getAttributes.map(_.getName))
+            val expectedAttributes = attributes.asScala.asInstanceOf[mutable.Buffer[Attribute]]
+              .filter(attr => attributesWhitelist.get.contains(attr.getName))
+            (name, expectedAttributes.size)}.toMap.filter(_._2 > 0)
+        else
+          queries.map((_, attributesWhitelist.get.length)).toMap
       }
 
     if(numExpectedAttributes.isEmpty) {
diff --git a/core/src/main/scala/kafka/tools/MirrorMaker.scala b/core/src/main/scala/kafka/tools/MirrorMaker.scala
index 823f0fd283845..5d6ea9ab0d288 100755
--- a/core/src/main/scala/kafka/tools/MirrorMaker.scala
+++ b/core/src/main/scala/kafka/tools/MirrorMaker.scala
@@ -20,7 +20,7 @@ package kafka.tools
 import java.time.Duration
 import java.util
 import java.util.concurrent.atomic.{AtomicBoolean, AtomicInteger}
-import java.util.concurrent.{CountDownLatch, TimeUnit}
+import java.util.concurrent.CountDownLatch
 import java.util.regex.Pattern
 import java.util.{Collections, Properties}
 
@@ -81,7 +81,7 @@ object MirrorMaker extends Logging with KafkaMetricsGroup {
       def value = numDroppedMessages.get()
     })
 
-  def main(args: Array[String]) {
+  def main(args: Array[String]): Unit = {
 
     info("Starting mirror maker")
     try {
@@ -133,6 +133,7 @@ object MirrorMaker extends Logging with KafkaMetricsGroup {
             // so if we catch it in commit we can safely retry
             // and re-throw to break the loop
             commitOffsets(consumerWrapper)
+            consumerWrapper.clearOffsetMap()
             throw e
 
           case _: TimeoutException =>
@@ -151,10 +152,15 @@ object MirrorMaker extends Logging with KafkaMetricsGroup {
 
           case _: CommitFailedException =>
             retryNeeded = false
+            consumerWrapper.clearOffsetMap()
             warn("Failed to commit offsets because the consumer group has rebalanced and assigned partitions to " +
               "another instance. If you see this regularly, it could indicate that you need to either increase " +
               s"the consumer's ${ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG} or reduce the number of records " +
               s"handled on each iteration with ${ConsumerConfig.MAX_POLL_RECORDS_CONFIG}")
+          // HOTFIX LIKAFKA-12852
+          case e: KafkaException if e.getMessage != null && e.getMessage.contains("may not exist or user may not have Describe access to topic") =>
+            consumerWrapper.clearOffsetMap()
+            error("Failed to commit offsets due to an unrecoverable error such as committing to a deleted topic", e)
         }
       }
     } else {
@@ -162,7 +168,7 @@ object MirrorMaker extends Logging with KafkaMetricsGroup {
     }
   }
 
-  def cleanShutdown() {
+  def cleanShutdown(): Unit = {
     if (isShuttingDown.compareAndSet(false, true)) {
       info("Start clean shutdown.")
       // Shutdown consumer threads.
@@ -177,7 +183,7 @@ object MirrorMaker extends Logging with KafkaMetricsGroup {
     }
   }
 
-  private def maybeSetDefaultProperty(properties: Properties, propertyName: String, defaultValue: String) {
+  private def maybeSetDefaultProperty(properties: Properties, propertyName: String, defaultValue: String): Unit = {
     val propertyValue = properties.getProperty(propertyName)
     properties.setProperty(propertyName, Option(propertyValue).getOrElse(defaultValue))
     if (properties.getProperty(propertyName) != defaultValue)
@@ -204,7 +210,7 @@ object MirrorMaker extends Logging with KafkaMetricsGroup {
         record.value,
         record.headers)
 
-    override def run() {
+    override def run(): Unit = {
       info(s"Starting mirror maker thread $threadName")
       try {
         consumerWrapper.init()
@@ -260,7 +266,7 @@ object MirrorMaker extends Logging with KafkaMetricsGroup {
       }
     }
 
-    def maybeFlushAndCommitOffsets() {
+    def maybeFlushAndCommitOffsets(): Unit = {
       if (System.currentTimeMillis() - lastOffsetCommitMs > offsetCommitIntervalMs) {
         debug("Committing MirrorMaker state.")
         producer.flush()
@@ -269,7 +275,7 @@ object MirrorMaker extends Logging with KafkaMetricsGroup {
       }
     }
 
-    def shutdown() {
+    def shutdown(): Unit = {
       try {
         info(s"$threadName shutting down")
         shuttingDown = true
@@ -281,7 +287,7 @@ object MirrorMaker extends Logging with KafkaMetricsGroup {
       }
     }
 
-    def awaitShutdown() {
+    def awaitShutdown(): Unit = {
       try {
         shutdownLatch.await()
         info("Mirror maker thread shutdown complete")
@@ -303,7 +309,7 @@ object MirrorMaker extends Logging with KafkaMetricsGroup {
     // Visible for testing
     private[tools] val offsets = new HashMap[TopicPartition, Long]()
 
-    def init() {
+    def init(): Unit = {
       debug("Initiating consumer")
       val consumerRebalanceListener = new InternalRebalanceListener(this, customRebalanceListener)
       whitelistOpt.foreach { whitelist =>
@@ -324,7 +330,7 @@ object MirrorMaker extends Logging with KafkaMetricsGroup {
         // uncommitted record since last poll. Using one second as poll's timeout ensures that
         // offsetCommitIntervalMs, of value greater than 1 second, does not see delays in offset
         // commit.
-        recordIter = consumer.poll(Duration.ofSeconds(1)).iterator
+        recordIter = consumer.poll(Duration.ofSeconds(1L)).iterator
         if (!recordIter.hasNext)
           throw new NoRecordsException
       }
@@ -336,31 +342,39 @@ object MirrorMaker extends Logging with KafkaMetricsGroup {
       record
     }
 
-    def wakeup() {
+    def wakeup(): Unit = {
       consumer.wakeup()
     }
 
-    def close() {
+    def close(): Unit = {
       consumer.close()
     }
 
-    def commit() {
+    def commit(): Unit = {
       consumer.commitSync(offsets.map { case (tp, offset) => (tp, new OffsetAndMetadata(offset)) }.asJava)
       offsets.clear()
     }
+
+    // HOTFIX LIKAFKA-12852, need to clear the offsets map when a KafkaException has been thrown due to topic
+    // deletion during offset commet.
+    def clearOffsetMap(): Unit = {
+      offsets.clear()
+    }
   }
 
   private class InternalRebalanceListener(consumerWrapper: ConsumerWrapper,
                                           customRebalanceListener: Option[ConsumerRebalanceListener])
     extends ConsumerRebalanceListener {
 
-    override def onPartitionsRevoked(partitions: util.Collection[TopicPartition]) {
+    override def onPartitionsLost(partitions: util.Collection[TopicPartition]): Unit = {}
+
+    override def onPartitionsRevoked(partitions: util.Collection[TopicPartition]): Unit = {
       producer.flush()
       commitOffsets(consumerWrapper)
       customRebalanceListener.foreach(_.onPartitionsRevoked(partitions))
     }
 
-    override def onPartitionsAssigned(partitions: util.Collection[TopicPartition]) {
+    override def onPartitionsAssigned(partitions: util.Collection[TopicPartition]): Unit = {
       customRebalanceListener.foreach(_.onPartitionsAssigned(partitions))
     }
   }
@@ -369,7 +383,7 @@ object MirrorMaker extends Logging with KafkaMetricsGroup {
 
     val producer = new KafkaProducer[Array[Byte], Array[Byte]](producerProps)
 
-    def send(record: ProducerRecord[Array[Byte], Array[Byte]]) {
+    def send(record: ProducerRecord[Array[Byte], Array[Byte]]): Unit = {
       if (sync) {
         this.producer.send(record).get()
       } else {
@@ -378,23 +392,23 @@ object MirrorMaker extends Logging with KafkaMetricsGroup {
       }
     }
 
-    def flush() {
+    def flush(): Unit = {
       this.producer.flush()
     }
 
-    def close() {
+    def close(): Unit = {
       this.producer.close()
     }
 
-    def close(timeout: Long) {
-      this.producer.close(timeout, TimeUnit.MILLISECONDS)
+    def close(timeout: Long): Unit = {
+      this.producer.close(Duration.ofMillis(timeout))
     }
   }
 
   private class MirrorMakerProducerCallback (topic: String, key: Array[Byte], value: Array[Byte])
     extends ErrorLoggingCallback(topic, key, value, false) {
 
-    override def onCompletion(metadata: RecordMetadata, exception: Exception) {
+    override def onCompletion(metadata: RecordMetadata, exception: Exception): Unit = {
       if (exception != null) {
         // Use default call back to log error. This means the max retries of producer has reached and message
         // still could not be sent.
@@ -496,6 +510,9 @@ object MirrorMaker extends Logging with KafkaMetricsGroup {
       .ofType(classOf[String])
       .defaultsTo("true")
 
+    val passthroughCompressionOpt = parser.accepts("enable.passthrough",
+      "When enabled, it avoids decompressing consumed record batches and doesn't re-compress in the producer")
+
     options = parser.parse(args: _*)
 
     def checkArgs() = {
@@ -518,7 +535,7 @@ object MirrorMaker extends Logging with KafkaMetricsGroup {
       val numStreams = options.valueOf(numStreamsOpt).intValue()
 
       Runtime.getRuntime.addShutdownHook(new Thread("MirrorMakerShutdownHook") {
-        override def run() {
+        override def run(): Unit = {
           cleanShutdown()
         }
       })
@@ -532,6 +549,10 @@ object MirrorMaker extends Logging with KafkaMetricsGroup {
       maybeSetDefaultProperty(producerProps, ProducerConfig.MAX_BLOCK_MS_CONFIG, Long.MaxValue.toString)
       maybeSetDefaultProperty(producerProps, ProducerConfig.ACKS_CONFIG, "all")
       maybeSetDefaultProperty(producerProps, ProducerConfig.MAX_IN_FLIGHT_REQUESTS_PER_CONNECTION, "1")
+      if (options.has(passthroughCompressionOpt)) {
+        consumerProps.setProperty("enable.shallow.iterator", "true")
+        producerProps.setProperty(ProducerConfig.COMPRESSION_TYPE_CONFIG, "passthrough")
+      }
       // Always set producer key and value serializer to ByteArraySerializer.
       producerProps.setProperty(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, classOf[ByteArraySerializer].getName)
       producerProps.setProperty(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, classOf[ByteArraySerializer].getName)
diff --git a/core/src/main/scala/kafka/tools/ReplicaVerificationTool.scala b/core/src/main/scala/kafka/tools/ReplicaVerificationTool.scala
index 6edd31557d1ee..8455b0d97e45f 100644
--- a/core/src/main/scala/kafka/tools/ReplicaVerificationTool.scala
+++ b/core/src/main/scala/kafka/tools/ReplicaVerificationTool.scala
@@ -30,7 +30,7 @@ import kafka.api._
 import kafka.utils.Whitelist
 import kafka.utils._
 import org.apache.kafka.clients._
-import org.apache.kafka.clients.admin.{ListTopicsOptions, TopicDescription}
+import org.apache.kafka.clients.admin.{Admin, ListTopicsOptions, TopicDescription}
 import org.apache.kafka.clients.consumer.{ConsumerConfig, KafkaConsumer}
 import org.apache.kafka.common.metrics.Metrics
 import org.apache.kafka.common.network.{NetworkReceive, Selectable, Selector}
@@ -43,6 +43,7 @@ import org.apache.kafka.common.utils.{LogContext, Time}
 import org.apache.kafka.common.{Node, TopicPartition}
 
 import scala.collection.JavaConverters._
+import scala.collection.Seq
 
 /**
  * For verifying the consistency among replicas.
@@ -105,11 +106,18 @@ object ReplicaVerificationTool extends Logging {
                          .describedAs("ms")
                          .ofType(classOf[java.lang.Long])
                          .defaultsTo(30 * 1000L)
+    val helpOpt = parser.accepts("help", "Print usage information.").forHelp()
+    val versionOpt = parser.accepts("version", "Print version information and exit.").forHelp()
 
-    if (args.length == 0)
+    val options = parser.parse(args: _*)
+
+    if (args.length == 0 || options.has(helpOpt)) {
       CommandLineUtils.printUsageAndDie(parser, "Validate that all replicas for a set of topics have the same data.")
+    }
 
-    val options = parser.parse(args: _*)
+    if (options.has(versionOpt)) {
+      CommandLineUtils.printVersionAndDie()
+    }
     CommandLineUtils.checkRequiredArgs(parser, options, brokerListOpt)
 
     val regex = options.valueOf(topicWhiteListOpt)
@@ -177,7 +185,7 @@ object ReplicaVerificationTool extends Logging {
     // create all replica fetcher threads
     val verificationBrokerId = brokerToTopicPartitions.head._1
     val counter = new AtomicInteger(0)
-    val fetcherThreads: Iterable[ReplicaFetcher] = brokerToTopicPartitions.map { case (brokerId, topicPartitions) =>
+    val fetcherThreads = brokerToTopicPartitions.map { case (brokerId, topicPartitions) =>
       new ReplicaFetcher(name = s"ReplicaFetcher-$brokerId",
         sourceBroker = brokerInfo(brokerId),
         topicPartitions = topicPartitions,
@@ -193,7 +201,7 @@ object ReplicaVerificationTool extends Logging {
     }
 
     Runtime.getRuntime.addShutdownHook(new Thread() {
-      override def run() {
+      override def run(): Unit = {
         info("Stopping all fetchers")
         fetcherThreads.foreach(_.shutdown())
       }
@@ -203,16 +211,16 @@ object ReplicaVerificationTool extends Logging {
 
   }
 
-  private def listTopicsMetadata(adminClient: admin.AdminClient): Seq[TopicDescription] = {
+  private def listTopicsMetadata(adminClient: Admin): Seq[TopicDescription] = {
     val topics = adminClient.listTopics(new ListTopicsOptions().listInternal(true)).names.get
     adminClient.describeTopics(topics).all.get.values.asScala.toBuffer
   }
 
-  private def brokerDetails(adminClient: admin.AdminClient): Map[Int, Node] = {
+  private def brokerDetails(adminClient: Admin): Map[Int, Node] = {
     adminClient.describeCluster.nodes.get.asScala.map(n => (n.id, n)).toMap
   }
 
-  private def createAdminClient(brokerUrl: String): admin.AdminClient = {
+  private def createAdminClient(brokerUrl: String): Admin = {
     val props = new Properties()
     props.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, brokerUrl)
     admin.AdminClient.create(props)
@@ -264,31 +272,31 @@ private class ReplicaBuffer(expectedReplicasPerTopicPartition: Map[TopicPartitio
   private var maxLagTopicAndPartition: TopicPartition = null
   initialize()
 
-  def createNewFetcherBarrier() {
+  def createNewFetcherBarrier(): Unit = {
     fetcherBarrier.set(new CountDownLatch(expectedNumFetchers))
   }
 
   def getFetcherBarrier() = fetcherBarrier.get
 
-  def createNewVerificationBarrier() {
+  def createNewVerificationBarrier(): Unit = {
     verificationBarrier.set(new CountDownLatch(1))
   }
 
   def getVerificationBarrier() = verificationBarrier.get
 
-  private def initialize() {
+  private def initialize(): Unit = {
     for (topicPartition <- expectedReplicasPerTopicPartition.keySet)
       recordsCache.put(topicPartition, new Pool[Int, FetchResponse.PartitionData[MemoryRecords]])
     setInitialOffsets()
   }
 
 
-  private def setInitialOffsets() {
+  private def setInitialOffsets(): Unit = {
     for ((tp, offset) <- initialOffsets)
       fetchOffsetMap.put(tp, offset)
   }
 
-  def addFetchedData(topicAndPartition: TopicPartition, replicaId: Int, partitionData: FetchResponse.PartitionData[MemoryRecords]) {
+  def addFetchedData(topicAndPartition: TopicPartition, replicaId: Int, partitionData: FetchResponse.PartitionData[MemoryRecords]): Unit = {
     recordsCache.get(topicAndPartition).put(replicaId, partitionData)
   }
 
@@ -296,7 +304,7 @@ private class ReplicaBuffer(expectedReplicasPerTopicPartition: Map[TopicPartitio
     fetchOffsetMap.get(topicAndPartition)
   }
 
-  def verifyCheckSum(println: String => Unit) {
+  def verifyCheckSum(println: String => Unit): Unit = {
     debug("Begin verification")
     maxLag = -1L
     for ((topicPartition, fetchResponsePerReplica) <- recordsCache) {
@@ -382,7 +390,7 @@ private class ReplicaFetcher(name: String, sourceBroker: Node, topicPartitions:
   private val fetchEndpoint = new ReplicaFetcherBlockingSend(sourceBroker, new ConsumerConfig(consumerConfig), new Metrics(), Time.SYSTEM, fetcherId,
     s"broker-${Request.DebuggingConsumerId}-fetcher-$fetcherId")
 
-  override def doWork() {
+  override def doWork(): Unit = {
 
     val fetcherBarrier = replicaBuffer.getFetcherBarrier()
     val verificationBarrier = replicaBuffer.getVerificationBarrier()
diff --git a/core/src/main/scala/kafka/tools/StateChangeLogMerger.scala b/core/src/main/scala/kafka/tools/StateChangeLogMerger.scala
index a3c80d17bb0a4..49eaccbe537d5 100755
--- a/core/src/main/scala/kafka/tools/StateChangeLogMerger.scala
+++ b/core/src/main/scala/kafka/tools/StateChangeLogMerger.scala
@@ -57,7 +57,7 @@ object StateChangeLogMerger extends Logging {
   var startDate: Date = null
   var endDate: Date = null
 
-  def main(args: Array[String]) {
+  def main(args: Array[String]): Unit = {
 
     // Parse input arguments.
     val parser = new OptionParser(false)
@@ -137,7 +137,7 @@ object StateChangeLogMerger extends Logging {
      */
     val pqueue = new mutable.PriorityQueue[LineIterator]()(dateBasedOrdering)
     val output: OutputStream = new BufferedOutputStream(System.out, 1024*1024)
-    val lineIterators = files.map(io.Source.fromFile(_).getLines)
+    val lineIterators = files.map(scala.io.Source.fromFile(_).getLines)
     var lines: List[LineIterator] = List()
 
     for (itr <- lineIterators) {
diff --git a/core/src/main/scala/kafka/tools/StreamsResetter.java b/core/src/main/scala/kafka/tools/StreamsResetter.java
index 3666f6740e37b..574e9c66e291a 100644
--- a/core/src/main/scala/kafka/tools/StreamsResetter.java
+++ b/core/src/main/scala/kafka/tools/StreamsResetter.java
@@ -22,11 +22,10 @@
 import joptsimple.OptionSpec;
 import joptsimple.OptionSpecBuilder;
 import kafka.utils.CommandLineUtils;
-import org.apache.kafka.clients.admin.AdminClient;
+import org.apache.kafka.clients.admin.Admin;
 import org.apache.kafka.clients.admin.DeleteTopicsResult;
 import org.apache.kafka.clients.admin.DescribeConsumerGroupsOptions;
 import org.apache.kafka.clients.admin.DescribeConsumerGroupsResult;
-import org.apache.kafka.clients.admin.KafkaAdminClient;
 import org.apache.kafka.clients.admin.MemberDescription;
 import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.clients.consumer.ConsumerConfig;
@@ -38,6 +37,7 @@
 import org.apache.kafka.common.serialization.ByteArrayDeserializer;
 import org.apache.kafka.common.utils.Exit;
 import org.apache.kafka.common.utils.Utils;
+import scala.collection.JavaConverters;
 
 import java.io.IOException;
 import java.text.ParseException;
@@ -102,7 +102,8 @@ public class StreamsResetter {
     private static OptionSpec fromFileOption;
     private static OptionSpec shiftByOption;
     private static OptionSpecBuilder dryRunOption;
-    private static OptionSpecBuilder helpOption;
+    private static OptionSpec helpOption;
+    private static OptionSpec versionOption;
     private static OptionSpecBuilder executeOption;
     private static OptionSpec commandConfigOption;
 
@@ -134,8 +135,7 @@ public int run(final String[] args,
                    final Properties config) {
         int exitCode;
 
-        KafkaAdminClient kafkaAdminClient = null;
-
+        Admin adminClient = null;
         try {
             parseArguments(args);
 
@@ -148,11 +148,11 @@ public int run(final String[] args,
             }
             properties.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, options.valueOf(bootstrapServerOption));
 
-            kafkaAdminClient = (KafkaAdminClient) AdminClient.create(properties);
-            validateNoActiveConsumers(groupId, kafkaAdminClient);
+            adminClient = Admin.create(properties);
+            validateNoActiveConsumers(groupId, adminClient);
 
             allTopics.clear();
-            allTopics.addAll(kafkaAdminClient.listTopics().names().get(60, TimeUnit.SECONDS));
+            allTopics.addAll(adminClient.listTopics().names().get(60, TimeUnit.SECONDS));
 
             if (dryRun) {
                 System.out.println("----Dry run displays the actions which will be performed when running Streams Reset Tool----");
@@ -161,15 +161,15 @@ public int run(final String[] args,
             final HashMap consumerConfig = new HashMap<>(config);
             consumerConfig.putAll(properties);
             exitCode = maybeResetInputAndSeekToEndIntermediateTopicOffsets(consumerConfig, dryRun);
-            maybeDeleteInternalTopics(kafkaAdminClient, dryRun);
+            maybeDeleteInternalTopics(adminClient, dryRun);
 
         } catch (final Throwable e) {
             exitCode = EXIT_CODE_ERROR;
             System.err.println("ERROR: " + e);
             e.printStackTrace(System.err);
         } finally {
-            if (kafkaAdminClient != null) {
-                kafkaAdminClient.close(Duration.ofSeconds(60));
+            if (adminClient != null) {
+                adminClient.close(Duration.ofSeconds(60));
             }
         }
 
@@ -177,11 +177,12 @@ public int run(final String[] args,
     }
 
     private void validateNoActiveConsumers(final String groupId,
-                                           final AdminClient adminClient)
+                                           final Admin adminClient)
         throws ExecutionException, InterruptedException {
 
-        final DescribeConsumerGroupsResult describeResult = adminClient.describeConsumerGroups(Collections.singleton(groupId),
-                (new DescribeConsumerGroupsOptions()).timeoutMs(10 * 1000));
+        final DescribeConsumerGroupsResult describeResult = adminClient.describeConsumerGroups(
+            Collections.singleton(groupId),
+            new DescribeConsumerGroupsOptions().timeoutMs(10 * 1000));
         final List members =
             new ArrayList<>(describeResult.describedGroups().get(groupId).get().members());
         if (!members.isEmpty()) {
@@ -191,8 +192,7 @@ private void validateNoActiveConsumers(final String groupId,
         }
     }
 
-    private void parseArguments(final String[] args) throws IOException {
-
+    private void parseArguments(final String[] args) {
         final OptionParser optionParser = new OptionParser(false);
         applicationIdOption = optionParser.accepts("application-id", "The Kafka Streams application ID (application.id).")
             .withRequiredArg()
@@ -238,7 +238,8 @@ private void parseArguments(final String[] args) throws IOException {
             .describedAs("file name");
         executeOption = optionParser.accepts("execute", "Execute the command.");
         dryRunOption = optionParser.accepts("dry-run", "Display the actions that would be performed without executing the reset commands.");
-        helpOption = optionParser.accepts("help", "Print usage information.");
+        helpOption = optionParser.accepts("help", "Print usage information.").forHelp();
+        versionOption = optionParser.accepts("version", "Print version information and exit.").forHelp();
 
         // TODO: deprecated in 1.0; can be removed eventually: https://issues.apache.org/jira/browse/KAFKA-7606
         optionParser.accepts("zookeeper", "Zookeeper option is deprecated by bootstrap.servers, as the reset tool would no longer access Zookeeper directly.");
@@ -248,31 +249,46 @@ private void parseArguments(final String[] args) throws IOException {
             if (args.length == 0 || options.has(helpOption)) {
                 CommandLineUtils.printUsageAndDie(optionParser, usage);
             }
+            if (options.has(versionOption)) {
+                CommandLineUtils.printVersionAndDie();
+            }
         } catch (final OptionException e) {
-            printHelp(optionParser);
-            throw e;
+            CommandLineUtils.printUsageAndDie(optionParser, e.getMessage());
         }
 
         if (options.has(executeOption) && options.has(dryRunOption)) {
             CommandLineUtils.printUsageAndDie(optionParser, "Only one of --dry-run and --execute can be specified");
         }
 
-        final scala.collection.immutable.HashSet> allScenarioOptions = new scala.collection.immutable.HashSet>();
-        allScenarioOptions.$plus(toOffsetOption);
-        allScenarioOptions.$plus(toDatetimeOption);
-        allScenarioOptions.$plus(byDurationOption);
-        allScenarioOptions.$plus(toEarliestOption);
-        allScenarioOptions.$plus(toLatestOption);
-        allScenarioOptions.$plus(fromFileOption);
-        allScenarioOptions.$plus(shiftByOption);
-
-        CommandLineUtils.checkInvalidArgs(optionParser, options, toOffsetOption, allScenarioOptions.$minus(toOffsetOption));
-        CommandLineUtils.checkInvalidArgs(optionParser, options, toDatetimeOption, allScenarioOptions.$minus(toDatetimeOption));
-        CommandLineUtils.checkInvalidArgs(optionParser, options, byDurationOption, allScenarioOptions.$minus(byDurationOption));
-        CommandLineUtils.checkInvalidArgs(optionParser, options, toEarliestOption, allScenarioOptions.$minus(toEarliestOption));
-        CommandLineUtils.checkInvalidArgs(optionParser, options, toLatestOption, allScenarioOptions.$minus(toLatestOption));
-        CommandLineUtils.checkInvalidArgs(optionParser, options, fromFileOption, allScenarioOptions.$minus(fromFileOption));
-        CommandLineUtils.checkInvalidArgs(optionParser, options, shiftByOption, allScenarioOptions.$minus(shiftByOption));
+        final Set> allScenarioOptions = new HashSet<>();
+        allScenarioOptions.add(toOffsetOption);
+        allScenarioOptions.add(toDatetimeOption);
+        allScenarioOptions.add(byDurationOption);
+        allScenarioOptions.add(toEarliestOption);
+        allScenarioOptions.add(toLatestOption);
+        allScenarioOptions.add(fromFileOption);
+        allScenarioOptions.add(shiftByOption);
+
+        checkInvalidArgs(optionParser, options, allScenarioOptions, toOffsetOption);
+        checkInvalidArgs(optionParser, options, allScenarioOptions, toDatetimeOption);
+        checkInvalidArgs(optionParser, options, allScenarioOptions, byDurationOption);
+        checkInvalidArgs(optionParser, options, allScenarioOptions, toEarliestOption);
+        checkInvalidArgs(optionParser, options, allScenarioOptions, toLatestOption);
+        checkInvalidArgs(optionParser, options, allScenarioOptions, fromFileOption);
+        checkInvalidArgs(optionParser, options, allScenarioOptions, shiftByOption);
+    }
+
+    private  void checkInvalidArgs(final OptionParser optionParser,
+                                      final OptionSet options,
+                                      final Set> allOptions,
+                                      final OptionSpec option) {
+        final Set> invalidOptions = new HashSet<>(allOptions);
+        invalidOptions.remove(option);
+        CommandLineUtils.checkInvalidArgs(
+            optionParser,
+            options,
+            option,
+            JavaConverters.asScalaSetConverter(invalidOptions).asScala());
     }
 
     private int maybeResetInputAndSeekToEndIntermediateTopicOffsets(final Map consumerConfig,
@@ -396,7 +412,6 @@ public void maybeSeekToEnd(final String groupId,
                     System.out.println("Topic: " + topicPartition.topic());
                 }
             }
-
             client.seekToEnd(intermediateTopicPartitions);
         }
     }
@@ -614,8 +629,7 @@ private boolean isIntermediateTopic(final String topic) {
         return options.valuesOf(intermediateTopicsOption).contains(topic);
     }
 
-    private void maybeDeleteInternalTopics(final KafkaAdminClient adminClient, final boolean dryRun) {
-
+    private void maybeDeleteInternalTopics(final Admin adminClient, final boolean dryRun) {
         System.out.println("Deleting all internal/auto-created topics for application " + options.valueOf(applicationIdOption));
         final List topicsToDelete = new ArrayList<>();
         for (final String listing : allTopics) {
@@ -635,7 +649,7 @@ private void maybeDeleteInternalTopics(final KafkaAdminClient adminClient, final
 
     // visible for testing
     public void doDelete(final List topicsToDelete,
-                          final AdminClient adminClient) {
+                          final Admin adminClient) {
         boolean hasDeleteErrors = false;
         final DeleteTopicsResult deleteTopicsResult = adminClient.deleteTopics(topicsToDelete);
         final Map> results = deleteTopicsResult.values();
@@ -654,7 +668,6 @@ public void doDelete(final List topicsToDelete,
         }
     }
 
-
     private boolean isInternalTopic(final String topicName) {
         // Specified input/intermediate topics might be named like internal topics (by chance).
         // Even is this is not expected in general, we need to exclude those topics here
@@ -662,12 +675,9 @@ private boolean isInternalTopic(final String topicName) {
         // Cf. https://issues.apache.org/jira/browse/KAFKA-7930
         return !isInputTopic(topicName) && !isIntermediateTopic(topicName)
             && topicName.startsWith(options.valueOf(applicationIdOption) + "-")
-            && (topicName.endsWith("-changelog") || topicName.endsWith("-repartition"));
-    }
-
-    private void printHelp(final OptionParser parser) throws IOException {
-        System.err.println(usage);
-        parser.printHelpOn(System.err);
+            && (topicName.endsWith("-changelog") || topicName.endsWith("-repartition")
+                || topicName.endsWith("-subscription-registration-topic")
+                || topicName.endsWith("-subscription-response-topic"));
     }
 
     public static void main(final String[] args) {
diff --git a/core/src/main/scala/kafka/utils/CommandDefaultOptions.scala b/core/src/main/scala/kafka/utils/CommandDefaultOptions.scala
index 096fa958cac2f..2cdb408b4bb8c 100644
--- a/core/src/main/scala/kafka/utils/CommandDefaultOptions.scala
+++ b/core/src/main/scala/kafka/utils/CommandDefaultOptions.scala
@@ -21,6 +21,7 @@ import joptsimple.{OptionParser, OptionSet}
 
 abstract class CommandDefaultOptions(val args: Array[String], allowCommandOptionAbbreviation: Boolean = false) {
   val parser = new OptionParser(allowCommandOptionAbbreviation)
-  val helpOpt = parser.accepts("help", "Print usage information.")
+  val helpOpt = parser.accepts("help", "Print usage information.").forHelp()
+  val versionOpt = parser.accepts("version", "Display Kafka version.").forHelp()
   var options: OptionSet = _
 }
diff --git a/core/src/main/scala/kafka/utils/CommandLineUtils.scala b/core/src/main/scala/kafka/utils/CommandLineUtils.scala
index 2c2157259b087..3576f4990f2f6 100644
--- a/core/src/main/scala/kafka/utils/CommandLineUtils.scala
+++ b/core/src/main/scala/kafka/utils/CommandLineUtils.scala
@@ -33,12 +33,21 @@ object CommandLineUtils extends Logging {
     * @return true on matching the help check condition
     */
   def isPrintHelpNeeded(commandOpts: CommandDefaultOptions): Boolean = {
-    return commandOpts.args.length == 0 || commandOpts.options.has(commandOpts.helpOpt)
+    commandOpts.args.length == 0 || commandOpts.options.has(commandOpts.helpOpt)
+  }
+
+  def isPrintVersionNeeded(commandOpts: CommandDefaultOptions): Boolean = {
+    commandOpts.options.has(commandOpts.versionOpt)
   }
 
   /**
     * Check and print help message if there is no options or `--help` option
-    * from command line
+    * from command line, if `--version` is specified on the command line
+    * print version information and exit.
+    * NOTE: The function name is not strictly speaking correct anymore
+    * as it also checks whether the version needs to be printed, but
+    * refactoring this would have meant changing all command line tools
+    * and unnecessarily increased the blast radius of this change.
     *
     * @param commandOpts Acceptable options for a command
     * @param message     Message to display on successful check
@@ -46,12 +55,14 @@ object CommandLineUtils extends Logging {
   def printHelpAndExitIfNeeded(commandOpts: CommandDefaultOptions, message: String) = {
     if (isPrintHelpNeeded(commandOpts))
       printUsageAndDie(commandOpts.parser, message)
+    if (isPrintVersionNeeded(commandOpts))
+      printVersionAndDie()
   }
 
   /**
    * Check that all the listed options are present
    */
-  def checkRequiredArgs(parser: OptionParser, options: OptionSet, required: OptionSpec[_]*) {
+  def checkRequiredArgs(parser: OptionParser, options: OptionSet, required: OptionSpec[_]*): Unit = {
     for (arg <- required) {
       if (!options.has(arg))
         printUsageAndDie(parser, "Missing required argument \"" + arg + "\"")
@@ -61,7 +72,7 @@ object CommandLineUtils extends Logging {
   /**
    * Check that none of the listed options are present
    */
-  def checkInvalidArgs(parser: OptionParser, options: OptionSet, usedOption: OptionSpec[_], invalidOptions: Set[OptionSpec[_]]) {
+  def checkInvalidArgs(parser: OptionParser, options: OptionSet, usedOption: OptionSpec[_], invalidOptions: Set[OptionSpec[_]]): Unit = {
     if (options.has(usedOption)) {
       for (arg <- invalidOptions) {
         if (options.has(arg))
@@ -73,7 +84,7 @@ object CommandLineUtils extends Logging {
   /**
     * Check that none of the listed options are present with the combination of used options
     */
-  def checkInvalidArgsSet(parser: OptionParser, options: OptionSet, usedOptions: Set[OptionSpec[_]], invalidOptions: Set[OptionSpec[_]]) {
+  def checkInvalidArgsSet(parser: OptionParser, options: OptionSet, usedOptions: Set[OptionSpec[_]], invalidOptions: Set[OptionSpec[_]]): Unit = {
     if (usedOptions.count(options.has) == usedOptions.size) {
       for (arg <- invalidOptions) {
         if (options.has(arg))
@@ -91,6 +102,11 @@ object CommandLineUtils extends Logging {
     Exit.exit(1, Some(message))
   }
 
+  def printVersionAndDie(): Nothing = {
+    System.out.println(VersionInfo.getVersionString)
+    Exit.exit(0)
+  }
+
   /**
    * Parse key-value pairs in the form key=value
    * value may contain equals sign
@@ -116,7 +132,7 @@ object CommandLineUtils extends Logging {
     * 3) otherwise, use the default value of {@code spec}.
     * A {@code null} value means to remove {@code key} from the {@code props}.
     */
-  def maybeMergeOptions[V](props: Properties, key: String, options: OptionSet, spec: OptionSpec[V]) {
+  def maybeMergeOptions[V](props: Properties, key: String, options: OptionSet, spec: OptionSpec[V]): Unit = {
     if (options.has(spec) || !props.containsKey(key)) {
       val value = options.valueOf(spec)
       if (value == null)
diff --git a/core/src/main/scala/kafka/utils/CoreUtils.scala b/core/src/main/scala/kafka/utils/CoreUtils.scala
index 778cc0dd67b80..2abac1bcbcba7 100755
--- a/core/src/main/scala/kafka/utils/CoreUtils.scala
+++ b/core/src/main/scala/kafka/utils/CoreUtils.scala
@@ -28,7 +28,7 @@ import com.typesafe.scalalogging.Logger
 import javax.management._
 
 import scala.collection._
-import scala.collection.mutable
+import scala.collection.{Seq, mutable}
 import kafka.cluster.EndPoint
 import org.apache.kafka.common.network.ListenerName
 import org.apache.kafka.common.security.auth.SecurityProtocol
@@ -50,10 +50,10 @@ object CoreUtils {
   private val logger = Logger(getClass)
 
   /**
-   * Return the smallest element in `traversable` if it is not empty. Otherwise return `ifEmpty`.
+   * Return the smallest element in `iterable` if it is not empty. Otherwise return `ifEmpty`.
    */
-  def min[A, B >: A](traversable: TraversableOnce[A], ifEmpty: A)(implicit cmp: Ordering[B]): A =
-    if (traversable.isEmpty) ifEmpty else traversable.min(cmp)
+  def min[A, B >: A](iterable: Iterable[A], ifEmpty: A)(implicit cmp: Ordering[B]): A =
+    if (iterable.isEmpty) ifEmpty else iterable.min(cmp)
 
   /**
    * Wrap the given function in a java.lang.Runnable
@@ -83,7 +83,7 @@ object CoreUtils {
     * @param logging The logging instance to use for logging the thrown exception.
     * @param logLevel The log level to use for logging.
     */
-  def swallow(action: => Unit, logging: Logging, logLevel: Level = Level.WARN) {
+  def swallow(action: => Unit, logging: Logging, logLevel: Level = Level.WARN): Unit = {
     try {
       action
     } catch {
@@ -159,7 +159,7 @@ object CoreUtils {
    * Unregister the mbean with the given name, if there is one registered
    * @param name The mbean name to unregister
    */
-  def unregisterMBean(name: String) {
+  def unregisterMBean(name: String): Unit = {
     val mbs = ManagementFactory.getPlatformMBeanServer()
     mbs synchronized {
       val objName = new ObjectName(name)
@@ -260,31 +260,6 @@ object CoreUtils {
 
   def inWriteLock[T](lock: ReadWriteLock)(fun: => T): T = inLock[T](lock.writeLock)(fun)
 
-
-  //JSON strings need to be escaped based on ECMA-404 standard http://json.org
-  def JSONEscapeString (s : String) : String = {
-    s.map {
-      case '"'  => "\\\""
-      case '\\' => "\\\\"
-      case '/'  => "\\/"
-      case '\b' => "\\b"
-      case '\f' => "\\f"
-      case '\n' => "\\n"
-      case '\r' => "\\r"
-      case '\t' => "\\t"
-      /* We'll unicode escape any control characters. These include:
-       * 0x0 -> 0x1f  : ASCII Control (C0 Control Codes)
-       * 0x7f         : ASCII DELETE
-       * 0x80 -> 0x9f : C1 Control Codes
-       *
-       * Per RFC4627, section 2.5, we're not technically required to
-       * encode the C1 codes, but we do to be safe.
-       */
-      case c if (c >= '\u0000' && c <= '\u001f') || (c >= '\u007f' && c <= '\u009f') => "\\u%04x".format(c: Int)
-      case c => c
-    }.mkString
-  }
-
   /**
    * Returns a list of duplicated items
    */
diff --git a/core/src/main/scala/kafka/utils/FileLock.scala b/core/src/main/scala/kafka/utils/FileLock.scala
index fd31b12a2b0d1..c635f76dff3a2 100644
--- a/core/src/main/scala/kafka/utils/FileLock.scala
+++ b/core/src/main/scala/kafka/utils/FileLock.scala
@@ -34,7 +34,7 @@ class FileLock(val file: File) extends Logging {
   /**
    * Lock the file or throw an exception if the lock is already held
    */
-  def lock() {
+  def lock(): Unit = {
     this synchronized {
       trace(s"Acquiring lock on ${file.getAbsolutePath}")
       flock = channel.lock()
@@ -62,7 +62,7 @@ class FileLock(val file: File) extends Logging {
   /**
    * Unlock the lock if it is held
    */
-  def unlock() {
+  def unlock(): Unit = {
     this synchronized {
       trace(s"Releasing lock on ${file.getAbsolutePath}")
       if(flock != null)
diff --git a/core/src/main/scala/kafka/utils/KafkaScheduler.scala b/core/src/main/scala/kafka/utils/KafkaScheduler.scala
index b4fae0bc85a7c..a175fde04c4e9 100755
--- a/core/src/main/scala/kafka/utils/KafkaScheduler.scala
+++ b/core/src/main/scala/kafka/utils/KafkaScheduler.scala
@@ -32,13 +32,13 @@ trait Scheduler {
   /**
    * Initialize this scheduler so it is ready to accept scheduling of tasks
    */
-  def startup()
+  def startup(): Unit
   
   /**
    * Shutdown this scheduler. When this method is complete no more executions of background tasks will occur. 
    * This includes tasks scheduled with a delayed execution.
    */
-  def shutdown()
+  def shutdown(): Unit
   
   /**
    * Check if the scheduler has been started
@@ -51,8 +51,9 @@ trait Scheduler {
    * @param delay The amount of time to wait before the first execution
    * @param period The period with which to execute the task. If < 0 the task will execute only once.
    * @param unit The unit for the preceding times.
+   * @return A Future object to manage the task scheduled.
    */
-  def schedule(name: String, fun: ()=>Unit, delay: Long = 0, period: Long = -1, unit: TimeUnit = TimeUnit.MILLISECONDS)
+  def schedule(name: String, fun: ()=>Unit, delay: Long = 0, period: Long = -1, unit: TimeUnit = TimeUnit.MILLISECONDS) : ScheduledFuture[_]
 }
 
 /**
@@ -71,7 +72,7 @@ class KafkaScheduler(val threads: Int,
   private var executor: ScheduledThreadPoolExecutor = null
   private val schedulerThreadId = new AtomicInteger(0)
 
-  override def startup() {
+  override def startup(): Unit = {
     debug("Initializing task scheduler.")
     this synchronized {
       if(isStarted)
@@ -79,6 +80,7 @@ class KafkaScheduler(val threads: Int,
       executor = new ScheduledThreadPoolExecutor(threads)
       executor.setContinueExistingPeriodicTasksAfterShutdownPolicy(false)
       executor.setExecuteExistingDelayedTasksAfterShutdownPolicy(false)
+      executor.setRemoveOnCancelPolicy(true)
       executor.setThreadFactory(new ThreadFactory() {
                                   def newThread(runnable: Runnable): Thread = 
                                     new KafkaThread(threadNamePrefix + schedulerThreadId.getAndIncrement(), runnable, daemon)
@@ -86,7 +88,7 @@ class KafkaScheduler(val threads: Int,
     }
   }
   
-  override def shutdown() {
+  override def shutdown(): Unit = {
     debug("Shutting down task scheduler.")
     // We use the local variable to avoid NullPointerException if another thread shuts down scheduler at same time.
     val cachedExecutor = this.executor
@@ -103,7 +105,7 @@ class KafkaScheduler(val threads: Int,
     schedule(name, fun, delay = 0L, period = -1L, unit = TimeUnit.MILLISECONDS)
   }
 
-  def schedule(name: String, fun: () => Unit, delay: Long, period: Long, unit: TimeUnit) {
+  def schedule(name: String, fun: () => Unit, delay: Long, period: Long, unit: TimeUnit): ScheduledFuture[_] = {
     debug("Scheduling task %s with initial delay %d ms and period %d ms."
         .format(name, TimeUnit.MILLISECONDS.convert(delay, unit), TimeUnit.MILLISECONDS.convert(period, unit)))
     this synchronized {
@@ -125,6 +127,13 @@ class KafkaScheduler(val threads: Int,
     }
   }
 
+  /**
+   * Package private for testing.
+   */
+  private[kafka] def taskRunning(task: ScheduledFuture[_]): Boolean = {
+    executor.getQueue().contains(task)
+  }
+
   def resizeThreadPool(newSize: Int): Unit = {
     executor.setCorePoolSize(newSize)
   }
diff --git a/core/src/main/scala/kafka/utils/Log4jController.scala b/core/src/main/scala/kafka/utils/Log4jController.scala
index 95d07334a61d5..ba0649c0ca2db 100755
--- a/core/src/main/scala/kafka/utils/Log4jController.scala
+++ b/core/src/main/scala/kafka/utils/Log4jController.scala
@@ -22,69 +22,95 @@ import java.util.Locale
 
 import org.apache.log4j.{Level, LogManager, Logger}
 
+import scala.collection.mutable
+import scala.collection.JavaConverters._
 
-/**
- * An MBean that allows the user to dynamically alter log4j levels at runtime.
- * The companion object contains the singleton instance of this class and
- * registers the MBean. The [[kafka.utils.Logging]] trait forces initialization
- * of the companion object.
- */
-private class Log4jController extends Log4jControllerMBean {
 
-  def getLoggers = {
-    val lst = new util.ArrayList[String]()
-    lst.add("root=" + existingLogger("root").getLevel.toString)
+object Log4jController {
+  val ROOT_LOGGER = "root"
+
+  /**
+    * Returns a map of the log4j loggers and their assigned log level.
+    * If a logger does not have a log level assigned, we return the root logger's log level
+    */
+  def loggers: mutable.Map[String, String] = {
+    val logs = new mutable.HashMap[String, String]()
+    val rootLoggerLvl = existingLogger(ROOT_LOGGER).getLevel.toString
+    logs.put(ROOT_LOGGER, rootLoggerLvl)
+
     val loggers = LogManager.getCurrentLoggers
     while (loggers.hasMoreElements) {
       val logger = loggers.nextElement().asInstanceOf[Logger]
       if (logger != null) {
-        val level =  if (logger != null) logger.getLevel else null
-        lst.add("%s=%s".format(logger.getName, if (level != null) level.toString else "null"))
+        val level = if (logger.getLevel != null) logger.getLevel.toString else rootLoggerLvl
+        logs.put(logger.getName, level)
       }
     }
-    lst
+    logs
   }
 
+  /**
+    * Sets the log level of a particular logger
+    */
+  def logLevel(loggerName: String, logLevel: String): Boolean = {
+    val log = existingLogger(loggerName)
+    if (!loggerName.trim.isEmpty && !logLevel.trim.isEmpty && log != null) {
+      log.setLevel(Level.toLevel(logLevel.toUpperCase(Locale.ROOT)))
+      true
+    }
+    else false
+  }
 
-  private def newLogger(loggerName: String) =
-    if (loggerName == "root")
-      LogManager.getRootLogger
-    else LogManager.getLogger(loggerName)
+  def unsetLogLevel(loggerName: String): Boolean = {
+    val log = existingLogger(loggerName)
+    if (!loggerName.trim.isEmpty && log != null) {
+      log.setLevel(null)
+      true
+    }
+    else false
+  }
 
+  def loggerExists(loggerName: String): Boolean = existingLogger(loggerName) != null
 
   private def existingLogger(loggerName: String) =
-    if (loggerName == "root")
+    if (loggerName == ROOT_LOGGER)
       LogManager.getRootLogger
     else LogManager.exists(loggerName)
+}
 
+/**
+ * An MBean that allows the user to dynamically alter log4j levels at runtime.
+ * The companion object contains the singleton instance of this class and
+ * registers the MBean. The [[kafka.utils.Logging]] trait forces initialization
+ * of the companion object.
+ */
+class Log4jController extends Log4jControllerMBean {
 
-  def getLogLevel(loggerName: String) = {
-    val log = existingLogger(loggerName)
+  def getLoggers: util.List[String] = {
+    Log4jController.loggers.map {
+      case (logger, level) => s"$logger=$level"
+    }.toList.asJava
+  }
+
+
+  def getLogLevel(loggerName: String): String = {
+    val log = Log4jController.existingLogger(loggerName)
     if (log != null) {
       val level = log.getLevel
       if (level != null)
         log.getLevel.toString
-      else "Null log level."
+      else
+        Log4jController.existingLogger(Log4jController.ROOT_LOGGER).getLevel.toString
     }
     else "No such logger."
   }
 
-
-  def setLogLevel(loggerName: String, level: String) = {
-    val log = newLogger(loggerName)
-    if (!loggerName.trim.isEmpty && !level.trim.isEmpty && log != null) {
-      log.setLevel(Level.toLevel(level.toUpperCase(Locale.ROOT)))
-      true
-    }
-    else false
-  }
-
+  def setLogLevel(loggerName: String, level: String): Boolean = Log4jController.logLevel(loggerName, level)
 }
 
 
-private trait Log4jControllerMBean {
+trait Log4jControllerMBean {
   def getLoggers: java.util.List[String]
   def getLogLevel(logger: String): String
   def setLogLevel(logger: String, level: String): Boolean
 }
-
diff --git a/core/src/main/scala/kafka/utils/Pool.scala b/core/src/main/scala/kafka/utils/Pool.scala
index 5dd61946f0d29..9a81ce3f9c50e 100644
--- a/core/src/main/scala/kafka/utils/Pool.scala
+++ b/core/src/main/scala/kafka/utils/Pool.scala
@@ -27,9 +27,11 @@ import collection.JavaConverters._
 class Pool[K,V](valueFactory: Option[K => V] = None) extends Iterable[(K, V)] {
 
   private val pool: ConcurrentMap[K, V] = new ConcurrentHashMap[K, V]
-  
+
   def put(k: K, v: V): V = pool.put(k, v)
-  
+
+  def putAll(map: java.util.Map[K, V]): Unit = pool.putAll(map)
+
   def putIfNotExists(k: K, v: V): V = pool.putIfAbsent(k, v)
 
   /**
@@ -73,7 +75,7 @@ class Pool[K,V](valueFactory: Option[K => V] = None) extends Iterable[(K, V)] {
 
   def values: Iterable[V] = pool.values.asScala
 
-  def clear() { pool.clear() }
+  def clear(): Unit = { pool.clear() }
   
   override def size: Int = pool.size
   
diff --git a/core/src/main/scala/kafka/utils/ReplicationUtils.scala b/core/src/main/scala/kafka/utils/ReplicationUtils.scala
index 33de22ba82521..e2733b8936fd9 100644
--- a/core/src/main/scala/kafka/utils/ReplicationUtils.scala
+++ b/core/src/main/scala/kafka/utils/ReplicationUtils.scala
@@ -39,13 +39,13 @@ object ReplicationUtils extends Logging {
     try {
       val (writtenLeaderOpt, writtenStat) = zkClient.getDataAndStat(path)
       val expectedLeaderOpt = TopicPartitionStateZNode.decode(expectedLeaderAndIsrInfo, writtenStat)
-      val succeeded = writtenLeaderOpt.map { writtenData =>
+      val succeeded = writtenLeaderOpt.exists { writtenData =>
         val writtenLeaderOpt = TopicPartitionStateZNode.decode(writtenData, writtenStat)
         (expectedLeaderOpt, writtenLeaderOpt) match {
           case (Some(expectedLeader), Some(writtenLeader)) if expectedLeader == writtenLeader => true
           case _ => false
         }
-      }.getOrElse(false)
+      }
       if (succeeded) (true, writtenStat.getVersion)
       else (false, -1)
     } catch {
diff --git a/core/src/main/scala/kafka/utils/ShutdownableThread.scala b/core/src/main/scala/kafka/utils/ShutdownableThread.scala
index 02d09dab96d82..0ca21c4ab166d 100644
--- a/core/src/main/scala/kafka/utils/ShutdownableThread.scala
+++ b/core/src/main/scala/kafka/utils/ShutdownableThread.scala
@@ -34,9 +34,16 @@ abstract class ShutdownableThread(val name: String, val isInterruptible: Boolean
     awaitShutdown()
   }
 
-  def isShutdownComplete: Boolean = {
-    shutdownComplete.getCount == 0
-  }
+  def isShutdownInitiated: Boolean = shutdownInitiated.getCount == 0
+
+  def isShutdownComplete: Boolean = shutdownComplete.getCount == 0
+
+  /**
+    * @return true if there has been an unexpected error and the thread shut down
+    */
+  // mind that run() might set both when we're shutting down the broker
+  // but the return value of this function at that point wouldn't matter
+  def isThreadFailed: Boolean = isShutdownComplete && !isShutdownInitiated
 
   def initiateShutdown(): Boolean = {
     this.synchronized {
@@ -55,7 +62,7 @@ abstract class ShutdownableThread(val name: String, val isInterruptible: Boolean
    * After calling initiateShutdown(), use this API to wait until the shutdown is complete
    */
   def awaitShutdown(): Unit = {
-    if (shutdownInitiated.getCount != 0)
+    if (!isShutdownInitiated)
       throw new IllegalStateException("initiateShutdown() was not called before awaitShutdown()")
     else {
       if (isStarted)
@@ -102,7 +109,5 @@ abstract class ShutdownableThread(val name: String, val isInterruptible: Boolean
     info("Stopped")
   }
 
-  def isRunning: Boolean = {
-    shutdownInitiated.getCount() != 0
-  }
+  def isRunning: Boolean = !isShutdownInitiated
 }
diff --git a/core/src/main/scala/kafka/utils/Throttler.scala b/core/src/main/scala/kafka/utils/Throttler.scala
index 9fe3cdcf13f44..cce6270cf02e8 100644
--- a/core/src/main/scala/kafka/utils/Throttler.scala
+++ b/core/src/main/scala/kafka/utils/Throttler.scala
@@ -49,7 +49,7 @@ class Throttler(desiredRatePerSec: Double,
   private var periodStartNs: Long = time.nanoseconds
   private var observedSoFar: Double = 0.0
   
-  def maybeThrottle(observed: Double) {
+  def maybeThrottle(observed: Double): Unit = {
     val msPerSec = TimeUnit.SECONDS.toMillis(1)
     val nsPerSec = TimeUnit.SECONDS.toNanos(1)
 
@@ -83,7 +83,7 @@ class Throttler(desiredRatePerSec: Double,
 
 object Throttler {
   
-  def main(args: Array[String]) {
+  def main(args: Array[String]): Unit = {
     val rand = new Random()
     val throttler = new Throttler(100000, 100, true, time = Time.SYSTEM)
     val interval = 30000
diff --git a/core/src/main/scala/kafka/utils/VerifiableProperties.scala b/core/src/main/scala/kafka/utils/VerifiableProperties.scala
index 5d70db5e11d06..8a0f33837c1d8 100755
--- a/core/src/main/scala/kafka/utils/VerifiableProperties.scala
+++ b/core/src/main/scala/kafka/utils/VerifiableProperties.scala
@@ -214,7 +214,7 @@ class VerifiableProperties(val props: Properties) extends Logging {
     }
   }
 
-  def verify() {
+  def verify(): Unit = {
     info("Verifying properties")
     val propNames = Collections.list(props.propertyNames).asScala.map(_.toString).sorted
     for(key <- propNames) {
diff --git a/core/src/main/scala/kafka/utils/VersionInfo.scala b/core/src/main/scala/kafka/utils/VersionInfo.scala
index 3d42d4d2b67e4..9d3130e6685d3 100644
--- a/core/src/main/scala/kafka/utils/VersionInfo.scala
+++ b/core/src/main/scala/kafka/utils/VersionInfo.scala
@@ -21,10 +21,20 @@ import org.apache.kafka.common.utils.AppInfoParser
 
 object VersionInfo {
 
-  def main(args: Array[String]) {
-    val version = AppInfoParser.getVersion
-    val commitId = AppInfoParser.getCommitId
-    System.out.println(s"${version} (Commit:${commitId})")
+  def main(args: Array[String]): Unit = {
+    System.out.println(getVersionString)
     System.exit(0)
   }
+
+  def getVersion: String = {
+    AppInfoParser.getVersion
+  }
+
+  def getCommit: String = {
+    AppInfoParser.getCommitId
+  }
+
+  def getVersionString: String = {
+    s"${getVersion} (Commit:${getCommit})"
+  }
 }
diff --git a/core/src/main/scala/kafka/utils/ZkUtils.scala b/core/src/main/scala/kafka/utils/ZkUtils.scala
deleted file mode 100644
index dd850906b42d8..0000000000000
--- a/core/src/main/scala/kafka/utils/ZkUtils.scala
+++ /dev/null
@@ -1,890 +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.nio.charset.StandardCharsets
-
-import kafka.admin._
-import kafka.api.LeaderAndIsr
-import kafka.cluster._
-import kafka.common.{KafkaException, NoEpochForPartitionException, TopicAndPartition}
-import kafka.controller.{LeaderIsrAndControllerEpoch, ReassignedPartitionsContext}
-import kafka.zk.{BrokerIdZNode, ReassignPartitionsZNode, ZkData}
-import org.I0Itec.zkclient.exception.{ZkBadVersionException, ZkMarshallingError, ZkNoNodeException, ZkNodeExistsException}
-import org.I0Itec.zkclient.serialize.ZkSerializer
-import org.I0Itec.zkclient.{IZkChildListener, IZkDataListener, IZkStateListener, ZkClient, ZkConnection}
-import org.apache.kafka.common.config.ConfigException
-import org.apache.zookeeper.data.{ACL, Stat}
-import org.apache.zookeeper.ZooDefs
-
-import scala.collection._
-import scala.collection.JavaConverters._
-import org.apache.kafka.common.TopicPartition
-
-@deprecated("This is an internal class that is no longer used by Kafka and will be removed in a future release. Please " +
-  "use org.apache.kafka.clients.admin.AdminClient instead.", since = "2.0.0")
-object ZkUtils {
-
-  private val UseDefaultAcls = new java.util.ArrayList[ACL]
-
-  // Important: it is necessary to add any new top level Zookeeper path here
-  val AdminPath = "/admin"
-  val BrokersPath = "/brokers"
-  val ClusterPath = "/cluster"
-  val ConfigPath = "/config"
-  val ControllerPath = "/controller"
-  val ControllerEpochPath = "/controller_epoch"
-  val IsrChangeNotificationPath = "/isr_change_notification"
-  val LogDirEventNotificationPath = "/log_dir_event_notification"
-  val KafkaAclPath = "/kafka-acl"
-  val KafkaAclChangesPath = "/kafka-acl-changes"
-
-  val ConsumersPath = "/consumers"
-  val ClusterIdPath = s"$ClusterPath/id"
-  val BrokerIdsPath = s"$BrokersPath/ids"
-  val BrokerTopicsPath = s"$BrokersPath/topics"
-  val ReassignPartitionsPath = s"$AdminPath/reassign_partitions"
-  val DeleteTopicsPath = s"$AdminPath/delete_topics"
-  val PreferredReplicaLeaderElectionPath = s"$AdminPath/preferred_replica_election"
-  val BrokerSequenceIdPath = s"$BrokersPath/seqid"
-  val ConfigChangesPath = s"$ConfigPath/changes"
-  val ConfigUsersPath = s"$ConfigPath/users"
-  val ConfigBrokersPath = s"$ConfigPath/brokers"
-  val ProducerIdBlockPath = "/latest_producer_id_block"
-
-  val SecureZkRootPaths = ZkData.SecureRootPaths
-
-  val SensitiveZkRootPaths = ZkData.SensitiveRootPaths
-
-  def apply(zkUrl: String, sessionTimeout: Int, connectionTimeout: Int, isZkSecurityEnabled: Boolean): ZkUtils = {
-    val (zkClient, zkConnection) = createZkClientAndConnection(zkUrl, sessionTimeout, connectionTimeout)
-    new ZkUtils(zkClient, zkConnection, isZkSecurityEnabled)
-  }
-
-  /*
-   * Used in tests
-   */
-  def apply(zkClient: ZkClient, isZkSecurityEnabled: Boolean): ZkUtils = {
-    new ZkUtils(zkClient, null, isZkSecurityEnabled)
-  }
-
-  def createZkClient(zkUrl: String, sessionTimeout: Int, connectionTimeout: Int): ZkClient = {
-    val zkClient = new ZkClient(zkUrl, sessionTimeout, connectionTimeout, ZKStringSerializer)
-    zkClient
-  }
-
-  def createZkClientAndConnection(zkUrl: String, sessionTimeout: Int, connectionTimeout: Int): (ZkClient, ZkConnection) = {
-    val zkConnection = new ZkConnection(zkUrl, sessionTimeout)
-    val zkClient = new ZkClient(zkConnection, connectionTimeout, ZKStringSerializer)
-    (zkClient, zkConnection)
-  }
-
-  def sensitivePath(path: String): Boolean = ZkData.sensitivePath(path)
-
-  def defaultAcls(isSecure: Boolean, path: String): java.util.List[ACL] = ZkData.defaultAcls(isSecure, path).asJava
-
-  def maybeDeletePath(zkUrl: String, dir: String) {
-    try {
-      val zk = createZkClient(zkUrl, 30*1000, 30*1000)
-      zk.deleteRecursive(dir)
-      zk.close()
-    } catch {
-      case _: Throwable => // swallow
-    }
-  }
-
-  /*
-   * Get calls that only depend on static paths
-   */
-  def getTopicPath(topic: String): String = {
-    ZkUtils.BrokerTopicsPath + "/" + topic
-  }
-
-  def getTopicPartitionsPath(topic: String): String = {
-    getTopicPath(topic) + "/partitions"
-  }
-
-  def getTopicPartitionPath(topic: String, partitionId: Int): String =
-    getTopicPartitionsPath(topic) + "/" + partitionId
-
-  def getTopicPartitionLeaderAndIsrPath(topic: String, partitionId: Int): String =
-    getTopicPartitionPath(topic, partitionId) + "/" + "state"
-
-  def getEntityConfigRootPath(entityType: String): String =
-    ZkUtils.ConfigPath + "/" + entityType
-
-  def getEntityConfigPath(entityType: String, entity: String): String =
-    getEntityConfigRootPath(entityType) + "/" + entity
-
-  def getEntityConfigPath(entityPath: String): String =
-    ZkUtils.ConfigPath + "/" + entityPath
-
-  def getDeleteTopicPath(topic: String): String =
-    DeleteTopicsPath + "/" + topic
-
-  def parsePartitionReassignmentData(jsonData: String): Map[TopicAndPartition, Seq[Int]] = {
-    val utf8Bytes = jsonData.getBytes(StandardCharsets.UTF_8)
-    val assignments = ReassignPartitionsZNode.decode(utf8Bytes) match {
-      case Left(e) => throw e
-      case Right(result) => result
-    }
-
-    assignments.map { case (tp, p) => (new TopicAndPartition(tp), p) }
-  }
-
-  def controllerZkData(brokerId: Int, timestamp: Long): String = {
-    Json.legacyEncodeAsString(Map("version" -> 1, "brokerid" -> brokerId, "timestamp" -> timestamp.toString))
-  }
-
-  def preferredReplicaLeaderElectionZkData(partitions: scala.collection.Set[TopicAndPartition]): String = {
-    Json.legacyEncodeAsString(Map("version" -> 1, "partitions" -> partitions.map(tp => Map("topic" -> tp.topic, "partition" -> tp.partition))))
-  }
-
-  def formatAsReassignmentJson(partitionsToBeReassigned: Map[TopicAndPartition, Seq[Int]]): String = {
-    Json.legacyEncodeAsString(Map(
-      "version" -> 1,
-      "partitions" -> partitionsToBeReassigned.map { case (TopicAndPartition(topic, partition), replicas) =>
-        Map(
-          "topic" -> topic,
-          "partition" -> partition,
-          "replicas" -> replicas
-        )
-      }
-    ))
-  }
-
-  def getReassignmentJson(partitionsToBeReassigned: Map[TopicPartition, Seq[Int]]): String = {
-    Json.encodeAsString(Map(
-      "version" -> 1,
-      "partitions" -> partitionsToBeReassigned.map { case (tp, replicas) =>
-        Map(
-          "topic" -> tp.topic,
-          "partition" -> tp.partition,
-          "replicas" -> replicas.asJava
-        ).asJava
-      }.asJava
-    ).asJava)
-  }
-}
-
-/**
- * Legacy class for interacting with ZooKeeper. Whenever possible, ``KafkaZkClient`` should be used instead.
- */
-@deprecated("This is an internal class that is no longer used by Kafka and will be removed in a future release. Please " +
-  "use org.apache.kafka.clients.admin.AdminClient instead.", since = "2.0.0")
-class ZkUtils(val zkClient: ZkClient,
-              val zkConnection: ZkConnection,
-              val isSecure: Boolean) extends Logging {
-  import ZkUtils._
-
-  // These are persistent ZK paths that should exist on kafka broker startup.
-  val persistentZkPaths = ZkData.PersistentZkPaths
-
-  // Visible for testing
-  val zkPath = new ZkPath(zkClient)
-
-  def defaultAcls(path: String): java.util.List[ACL] = ZkUtils.defaultAcls(isSecure, path)
-
-  def getController(): Int = {
-    readDataMaybeNull(ControllerPath)._1 match {
-      case Some(controller) => parseControllerId(controller)
-      case None => throw new KafkaException("Controller doesn't exist")
-    }
-  }
-
-  def parseControllerId(controllerInfoString: String): Int = {
-    try {
-      Json.parseFull(controllerInfoString) match {
-        case Some(js) => js.asJsonObject("brokerid").to[Int]
-        case None => throw new KafkaException("Failed to parse the controller info json [%s].".format(controllerInfoString))
-      }
-    } catch {
-      case _: Throwable =>
-        // It may be due to an incompatible controller register version
-        warn("Failed to parse the controller info as json. "
-          + "Probably this controller is still using the old format [%s] to store the broker id in zookeeper".format(controllerInfoString))
-        try controllerInfoString.toInt
-        catch {
-          case t: Throwable => throw new KafkaException(s"Failed to parse the controller info: $controllerInfoString. This is neither the new or the old format.", t)
-        }
-    }
-  }
-
-  /* Represents a cluster identifier. Stored in Zookeeper in JSON format: {"version" -> "1", "id" -> id } */
-  object ClusterId {
-
-    def toJson(id: String) = {
-      Json.legacyEncodeAsString(Map("version" -> "1", "id" -> id))
-    }
-
-    def fromJson(clusterIdJson: String): String = {
-      Json.parseFull(clusterIdJson).map(_.asJsonObject("id").to[String]).getOrElse {
-        throw new KafkaException(s"Failed to parse the cluster id json $clusterIdJson")
-      }
-    }
-  }
-
-  def getClusterId: Option[String] =
-    readDataMaybeNull(ClusterIdPath)._1.map(ClusterId.fromJson)
-
-  def createOrGetClusterId(proposedClusterId: String): String = {
-    try {
-      createPersistentPath(ClusterIdPath, ClusterId.toJson(proposedClusterId))
-      proposedClusterId
-    } catch {
-      case _: ZkNodeExistsException =>
-        getClusterId.getOrElse(throw new KafkaException("Failed to get cluster id from Zookeeper. This can only happen if /cluster/id is deleted from Zookeeper."))
-    }
-  }
-
-  def getSortedBrokerList(): Seq[Int] =
-    getChildren(BrokerIdsPath).map(_.toInt).sorted
-
-  def getAllBrokersInCluster(): Seq[Broker] = {
-    val brokerIds = getChildrenParentMayNotExist(BrokerIdsPath).sorted
-    brokerIds.map(_.toInt).map(getBrokerInfo(_)).filter(_.isDefined).map(_.get)
-  }
-
-  def getLeaderAndIsrForPartition(topic: String, partition: Int): Option[LeaderAndIsr] = {
-    val leaderAndIsrPath = getTopicPartitionLeaderAndIsrPath(topic, partition)
-    val (leaderAndIsrOpt, stat) = readDataMaybeNull(leaderAndIsrPath)
-    debug(s"Read leaderISR $leaderAndIsrOpt for $topic-$partition")
-    leaderAndIsrOpt.flatMap(leaderAndIsrStr => parseLeaderAndIsr(leaderAndIsrStr, leaderAndIsrPath, stat).map(_.leaderAndIsr))
-  }
-
-  private def parseLeaderAndIsr(leaderAndIsrStr: String, path: String, stat: Stat): Option[LeaderIsrAndControllerEpoch] = {
-    Json.parseFull(leaderAndIsrStr).flatMap { js =>
-      val leaderIsrAndEpochInfo = js.asJsonObject
-      val leader = leaderIsrAndEpochInfo("leader").to[Int]
-      val epoch = leaderIsrAndEpochInfo("leader_epoch").to[Int]
-      val isr = leaderIsrAndEpochInfo("isr").to[List[Int]]
-      val controllerEpoch = leaderIsrAndEpochInfo("controller_epoch").to[Int]
-      val zkPathVersion = stat.getVersion
-      trace(s"Leader $leader, Epoch $epoch, Isr $isr, Zk path version $zkPathVersion for leaderAndIsrPath $path")
-      Some(LeaderIsrAndControllerEpoch(LeaderAndIsr(leader, epoch, isr, zkPathVersion), controllerEpoch))}
-  }
-
-  def setupCommonPaths() {
-    for(path <- persistentZkPaths)
-      makeSurePersistentPathExists(path)
-  }
-
-  def getLeaderForPartition(topic: String, partition: Int): Option[Int] = {
-    readDataMaybeNull(getTopicPartitionLeaderAndIsrPath(topic, partition))._1.flatMap { leaderAndIsr =>
-      Json.parseFull(leaderAndIsr).map(_.asJsonObject("leader").to[Int])
-    }
-  }
-
-  /**
-   * This API should read the epoch in the ISR path. It is sufficient to read the epoch in the ISR path, since if the
-   * leader fails after updating epoch in the leader path and before updating epoch in the ISR path, effectively some
-   * other broker will retry becoming leader with the same new epoch value.
-   */
-  def getEpochForPartition(topic: String, partition: Int): Int = {
-    readDataMaybeNull(getTopicPartitionLeaderAndIsrPath(topic, partition))._1 match {
-      case Some(leaderAndIsr) =>
-        Json.parseFull(leaderAndIsr) match {
-          case None => throw new NoEpochForPartitionException("No epoch, leaderAndISR data for partition [%s,%d] is invalid".format(topic, partition))
-          case Some(js) => js.asJsonObject("leader_epoch").to[Int]
-        }
-      case None => throw new NoEpochForPartitionException("No epoch, ISR path for partition [%s,%d] is empty"
-        .format(topic, partition))
-    }
-  }
-
-  /** returns a sequence id generated by updating BrokerSequenceIdPath in Zk.
-    * users can provide brokerId in the config , inorder to avoid conflicts between zk generated
-    * seqId and config.brokerId we increment zk seqId by KafkaConfig.MaxReservedBrokerId.
-    */
-  def getBrokerSequenceId(MaxReservedBrokerId: Int): Int = {
-    getSequenceId(BrokerSequenceIdPath) + MaxReservedBrokerId
-  }
-
-  /**
-   * Gets the in-sync replicas (ISR) for a specific topic and partition
-   */
-  def getInSyncReplicasForPartition(topic: String, partition: Int): Seq[Int] = {
-    val leaderAndIsrOpt = readDataMaybeNull(getTopicPartitionLeaderAndIsrPath(topic, partition))._1
-    leaderAndIsrOpt match {
-      case Some(leaderAndIsr) =>
-        Json.parseFull(leaderAndIsr) match {
-          case Some(js) => js.asJsonObject("isr").to[Seq[Int]]
-          case None => Seq.empty[Int]
-        }
-      case None => Seq.empty[Int]
-    }
-  }
-
-  /**
-   * Gets the assigned replicas (AR) for a specific topic and partition
-   */
-  def getReplicasForPartition(topic: String, partition: Int): Seq[Int] = {
-    val seqOpt = for {
-      jsonPartitionMap <- readDataMaybeNull(getTopicPath(topic))._1
-      js <- Json.parseFull(jsonPartitionMap)
-      replicaMap <- js.asJsonObject.get("partitions")
-      seq <- replicaMap.asJsonObject.get(partition.toString)
-    } yield seq.to[Seq[Int]]
-    seqOpt.getOrElse(Seq.empty)
-  }
-
-  def leaderAndIsrZkData(leaderAndIsr: LeaderAndIsr, controllerEpoch: Int): String = {
-    Json.legacyEncodeAsString(Map("version" -> 1, "leader" -> leaderAndIsr.leader, "leader_epoch" -> leaderAndIsr.leaderEpoch,
-                    "controller_epoch" -> controllerEpoch, "isr" -> leaderAndIsr.isr))
-  }
-
-  /**
-   * Get JSON partition to replica map from zookeeper.
-   */
-  def replicaAssignmentZkData(map: Map[String, Seq[Int]]): String = {
-    Json.legacyEncodeAsString(Map("version" -> 1, "partitions" -> map))
-  }
-
-  /**
-   *  make sure a persistent path exists in ZK. Create the path if not exist.
-   */
-  def makeSurePersistentPathExists(path: String, acls: java.util.List[ACL] = UseDefaultAcls) {
-    //Consumer path is kept open as different consumers will write under this node.
-    val acl = if (path == null || path.isEmpty || path.equals(ConsumersPath)) {
-      ZooDefs.Ids.OPEN_ACL_UNSAFE
-    } else if (acls eq UseDefaultAcls) {
-      ZkUtils.defaultAcls(isSecure, path)
-    } else {
-      acls
-    }
-
-    if (!zkClient.exists(path))
-      zkPath.createPersistent(path, createParents = true, acl) //won't throw NoNodeException or NodeExistsException
-  }
-
-  /**
-   *  create the parent path
-   */
-  private def createParentPath(path: String, acls: java.util.List[ACL] = UseDefaultAcls): Unit = {
-    val acl = if (acls eq UseDefaultAcls) ZkUtils.defaultAcls(isSecure, path) else acls
-    val parentDir = path.substring(0, path.lastIndexOf('/'))
-    if (parentDir.length != 0) {
-      zkPath.createPersistent(parentDir, createParents = true, acl)
-    }
-  }
-
-  /**
-   * Create an ephemeral node with the given path and data. Create parents if necessary.
-   */
-  private def createEphemeralPath(path: String, data: String, acls: java.util.List[ACL]): Unit = {
-    val acl = if (acls eq UseDefaultAcls) ZkUtils.defaultAcls(isSecure, path) else acls
-    try {
-      zkPath.createEphemeral(path, data, acl)
-    } catch {
-      case _: ZkNoNodeException =>
-        createParentPath(path)
-        zkPath.createEphemeral(path, data, acl)
-    }
-  }
-
-  /**
-   * Create an ephemeral node with the given path and data.
-   * Throw NodeExistException if node already exists.
-   */
-  def createEphemeralPathExpectConflict(path: String, data: String, acls: java.util.List[ACL] = UseDefaultAcls): Unit = {
-    val acl = if (acls eq UseDefaultAcls) ZkUtils.defaultAcls(isSecure, path) else acls
-    try {
-      createEphemeralPath(path, data, acl)
-    } catch {
-      case e: ZkNodeExistsException =>
-        // this can happen when there is connection loss; make sure the data is what we intend to write
-        var storedData: String = null
-        try {
-          storedData = readData(path)._1
-        } catch {
-          case _: ZkNoNodeException => // the node disappeared; treat as if node existed and let caller handles this
-        }
-        if (storedData == null || storedData != data) {
-          info(s"conflict in $path data: $data stored data: $storedData")
-          throw e
-        } else {
-          // otherwise, the creation succeeded, return normally
-          info(s"$path exists with value $data during connection loss; this is ok")
-        }
-    }
-  }
-
-  /**
-   * Create a persistent node with the given path and data. Create parents if necessary.
-   */
-  def createPersistentPath(path: String, data: String = "", acls: java.util.List[ACL] = UseDefaultAcls): Unit = {
-    val acl = if (acls eq UseDefaultAcls) ZkUtils.defaultAcls(isSecure, path) else acls
-    try {
-      zkPath.createPersistent(path, data, acl)
-    } catch {
-      case _: ZkNoNodeException =>
-        createParentPath(path)
-        zkPath.createPersistent(path, data, acl)
-    }
-  }
-
-  def createSequentialPersistentPath(path: String, data: String = "", acls: java.util.List[ACL] = UseDefaultAcls): String = {
-    val acl = if (acls eq UseDefaultAcls) ZkUtils.defaultAcls(isSecure, path) else acls
-    zkPath.createPersistentSequential(path, data, acl)
-  }
-
-  /**
-   * Update the value of a persistent node with the given path and data.
-   * create parent directory if necessary. Never throw NodeExistException.
-   * Return the updated path zkVersion
-   */
-  def updatePersistentPath(path: String, data: String, acls: java.util.List[ACL] = UseDefaultAcls) = {
-    val acl = if (acls eq UseDefaultAcls) ZkUtils.defaultAcls(isSecure, path) else acls
-    try {
-      zkClient.writeData(path, data)
-    } catch {
-      case _: ZkNoNodeException =>
-        createParentPath(path)
-        try {
-          zkPath.createPersistent(path, data, acl)
-        } catch {
-          case _: ZkNodeExistsException =>
-            zkClient.writeData(path, data)
-        }
-    }
-  }
-
-  /**
-   * Conditional update the persistent path data, return (true, newVersion) if it succeeds, otherwise (the path doesn't
-   * exist, the current version is not the expected version, etc.) return (false, -1)
-   *
-   * When there is a ConnectionLossException during the conditional update, zkClient will retry the update and may fail
-   * since the previous update may have succeeded (but the stored zkVersion no longer matches the expected one).
-   * In this case, we will run the optionalChecker to further check if the previous write did indeed succeeded.
-   */
-  def conditionalUpdatePersistentPath(path: String, data: String, expectVersion: Int,
-    optionalChecker:Option[(ZkUtils, String, String) => (Boolean,Int)] = None): (Boolean, Int) = {
-    try {
-      val stat = zkClient.writeDataReturnStat(path, data, expectVersion)
-      debug("Conditional update of path %s with value %s and expected version %d succeeded, returning the new version: %d"
-        .format(path, data, expectVersion, stat.getVersion))
-      (true, stat.getVersion)
-    } catch {
-      case e1: ZkBadVersionException =>
-        optionalChecker match {
-          case Some(checker) => checker(this, path, data)
-          case _ =>
-            debug("Checker method is not passed skipping zkData match")
-            debug("Conditional update of path %s with data %s and expected version %d failed due to %s"
-              .format(path, data,expectVersion, e1.getMessage))
-            (false, -1)
-        }
-      case e2: Exception =>
-        debug("Conditional update of path %s with data %s and expected version %d failed due to %s".format(path, data,
-          expectVersion, e2.getMessage))
-        (false, -1)
-    }
-  }
-
-  /**
-   * Conditional update the persistent path data, return (true, newVersion) if it succeeds, otherwise (the current
-   * version is not the expected version, etc.) return (false, -1). If path doesn't exist, throws ZkNoNodeException
-   */
-  def conditionalUpdatePersistentPathIfExists(path: String, data: String, expectVersion: Int): (Boolean, Int) = {
-    try {
-      val stat = zkClient.writeDataReturnStat(path, data, expectVersion)
-      debug("Conditional update of path %s with value %s and expected version %d succeeded, returning the new version: %d"
-        .format(path, data, expectVersion, stat.getVersion))
-      (true, stat.getVersion)
-    } catch {
-      case nne: ZkNoNodeException => throw nne
-      case e: Exception =>
-        error("Conditional update of path %s with data %s and expected version %d failed due to %s".format(path, data,
-          expectVersion, e.getMessage))
-        (false, -1)
-    }
-  }
-
-  /**
-   * Update the value of a ephemeral node with the given path and data.
-   * create parent directory if necessary. Never throw NodeExistException.
-   */
-  def updateEphemeralPath(path: String, data: String, acls: java.util.List[ACL] = UseDefaultAcls): Unit = {
-    val acl = if (acls eq UseDefaultAcls) ZkUtils.defaultAcls(isSecure, path) else acls
-    try {
-      zkClient.writeData(path, data)
-    } catch {
-      case _: ZkNoNodeException =>
-        createParentPath(path)
-        zkPath.createEphemeral(path, data, acl)
-    }
-  }
-
-  def deletePath(path: String): Boolean = {
-    zkClient.delete(path)
-  }
-
-  /**
-    * Conditional delete the persistent path data, return true if it succeeds,
-    * false otherwise (the current version is not the expected version)
-    */
-   def conditionalDeletePath(path: String, expectedVersion: Int): Boolean = {
-    try {
-      zkClient.delete(path, expectedVersion)
-      true
-    } catch {
-      case _: ZkBadVersionException => false
-    }
-  }
-
-  def deletePathRecursive(path: String) {
-    zkClient.deleteRecursive(path)
-  }
-
-  def subscribeDataChanges(path: String, listener: IZkDataListener): Unit =
-    zkClient.subscribeDataChanges(path, listener)
-
-  def unsubscribeDataChanges(path: String, dataListener: IZkDataListener): Unit =
-    zkClient.unsubscribeDataChanges(path, dataListener)
-
-  def subscribeStateChanges(listener: IZkStateListener): Unit =
-    zkClient.subscribeStateChanges(listener)
-
-  def subscribeChildChanges(path: String, listener: IZkChildListener): Option[Seq[String]] =
-    Option(zkClient.subscribeChildChanges(path, listener)).map(_.asScala)
-
-  def unsubscribeChildChanges(path: String, childListener: IZkChildListener): Unit =
-    zkClient.unsubscribeChildChanges(path, childListener)
-
-  def unsubscribeAll(): Unit =
-    zkClient.unsubscribeAll()
-
-  def readData(path: String): (String, Stat) = {
-    val stat: Stat = new Stat()
-    val dataStr: String = zkClient.readData[String](path, stat)
-    (dataStr, stat)
-  }
-
-  def readDataMaybeNull(path: String): (Option[String], Stat) = {
-    val stat = new Stat()
-    val dataAndStat = try {
-                        val dataStr = zkClient.readData[String](path, stat)
-                        (Some(dataStr), stat)
-                      } catch {
-                        case _: ZkNoNodeException =>
-                          (None, stat)
-                      }
-    dataAndStat
-  }
-
-  def readDataAndVersionMaybeNull(path: String): (Option[String], Int) = {
-    val stat = new Stat()
-    try {
-      val data = zkClient.readData[String](path, stat)
-      (Option(data), stat.getVersion)
-    } catch {
-      case _: ZkNoNodeException => (None, stat.getVersion)
-    }
-  }
-
-  def getChildren(path: String): Seq[String] = zkClient.getChildren(path).asScala
-
-  def getChildrenParentMayNotExist(path: String): Seq[String] = {
-    try {
-      zkClient.getChildren(path).asScala
-    } catch {
-      case _: ZkNoNodeException => Nil
-    }
-  }
-
-  /**
-   * Check if the given path exists
-   */
-  def pathExists(path: String): Boolean = {
-    zkClient.exists(path)
-  }
-
-  def isTopicMarkedForDeletion(topic: String): Boolean = {
-    pathExists(getDeleteTopicPath(topic))
-  }
-
-  def getCluster(): Cluster = {
-    val cluster = new Cluster
-    val nodes = getChildrenParentMayNotExist(BrokerIdsPath)
-    for (node <- nodes) {
-      val brokerZKString = readData(BrokerIdsPath + "/" + node)._1
-      cluster.add(parseBrokerJson(node.toInt, brokerZKString))
-    }
-    cluster
-  }
-
-  private def parseBrokerJson(id: Int, jsonString: String): Broker = {
-    BrokerIdZNode.decode(id, jsonString.getBytes(StandardCharsets.UTF_8)).broker
-  }
-
-  def getPartitionLeaderAndIsrForTopics(topicAndPartitions: Set[TopicAndPartition]): mutable.Map[TopicAndPartition, LeaderIsrAndControllerEpoch] = {
-    val ret = new mutable.HashMap[TopicAndPartition, LeaderIsrAndControllerEpoch]
-    for(topicAndPartition <- topicAndPartitions) {
-      getLeaderIsrAndEpochForPartition(topicAndPartition.topic, topicAndPartition.partition).foreach { leaderIsrAndControllerEpoch =>
-        ret.put(topicAndPartition, leaderIsrAndControllerEpoch)
-      }
-    }
-    ret
-  }
-
-  private[utils] def getLeaderIsrAndEpochForPartition(topic: String, partition: Int): Option[LeaderIsrAndControllerEpoch] = {
-    val leaderAndIsrPath = getTopicPartitionLeaderAndIsrPath(topic, partition)
-    val (leaderAndIsrOpt, stat) = readDataMaybeNull(leaderAndIsrPath)
-    debug(s"Read leaderISR $leaderAndIsrOpt for $topic-$partition")
-    leaderAndIsrOpt.flatMap(leaderAndIsrStr => parseLeaderAndIsr(leaderAndIsrStr, leaderAndIsrPath, stat))
-  }
-
-  def getReplicaAssignmentForTopics(topics: Seq[String]): mutable.Map[TopicAndPartition, Seq[Int]] = {
-    val ret = new mutable.HashMap[TopicAndPartition, Seq[Int]]
-    topics.foreach { topic =>
-      readDataMaybeNull(getTopicPath(topic))._1.foreach { jsonPartitionMap =>
-        Json.parseFull(jsonPartitionMap).foreach { js =>
-          js.asJsonObject.get("partitions").foreach { partitionsJs =>
-            partitionsJs.asJsonObject.iterator.foreach { case (partition, replicas) =>
-              ret.put(TopicAndPartition(topic, partition.toInt), replicas.to[Seq[Int]])
-              debug("Replicas assigned to topic [%s], partition [%s] are [%s]".format(topic, partition, replicas))
-            }
-          }
-        }
-      }
-    }
-    ret
-  }
-
-  def getPartitionAssignmentForTopics(topics: Seq[String]): mutable.Map[String, collection.Map[Int, Seq[Int]]] = {
-    val ret = new mutable.HashMap[String, Map[Int, Seq[Int]]]()
-    topics.foreach { topic =>
-      val partitionMapOpt = for {
-        jsonPartitionMap <- readDataMaybeNull(getTopicPath(topic))._1
-        js <- Json.parseFull(jsonPartitionMap)
-        replicaMap <- js.asJsonObject.get("partitions")
-      } yield replicaMap.asJsonObject.iterator.map { case (k, v) => (k.toInt, v.to[Seq[Int]]) }.toMap
-      val partitionMap = partitionMapOpt.getOrElse(Map.empty)
-      debug("Partition map for /brokers/topics/%s is %s".format(topic, partitionMap))
-      ret += (topic -> partitionMap)
-    }
-    ret
-  }
-
-  def getPartitionsForTopics(topics: Seq[String]): mutable.Map[String, Seq[Int]] = {
-    getPartitionAssignmentForTopics(topics).map { topicAndPartitionMap =>
-      val topic = topicAndPartitionMap._1
-      val partitionMap = topicAndPartitionMap._2
-      debug("partition assignment of /brokers/topics/%s is %s".format(topic, partitionMap))
-      topic -> partitionMap.keys.toSeq.sortWith((s, t) => s < t)
-    }
-  }
-
-  def getTopicPartitionCount(topic: String): Option[Int] = {
-    val topicData = getPartitionAssignmentForTopics(Seq(topic))
-    if (topicData(topic).nonEmpty)
-      Some(topicData(topic).size)
-    else
-      None
-  }
-
-  def getPartitionsBeingReassigned(): Map[TopicAndPartition, ReassignedPartitionsContext] = {
-    // read the partitions and their new replica list
-    val jsonPartitionMapOpt = readDataMaybeNull(ReassignPartitionsPath)._1
-    jsonPartitionMapOpt match {
-      case Some(jsonPartitionMap) =>
-        val reassignedPartitions = parsePartitionReassignmentData(jsonPartitionMap)
-        reassignedPartitions.map { case (tp, newReplicas) =>
-          tp -> new ReassignedPartitionsContext(newReplicas, null)
-        }
-      case None => Map.empty[TopicAndPartition, ReassignedPartitionsContext]
-    }
-  }
-
-  def updatePartitionReassignmentData(partitionsToBeReassigned: Map[TopicAndPartition, Seq[Int]]) {
-    val zkPath = ZkUtils.ReassignPartitionsPath
-    partitionsToBeReassigned.size match {
-      case 0 => // need to delete the /admin/reassign_partitions path
-        deletePath(zkPath)
-        info("No more partitions need to be reassigned. Deleting zk path %s".format(zkPath))
-      case _ =>
-        val jsonData = formatAsReassignmentJson(partitionsToBeReassigned)
-        try {
-          updatePersistentPath(zkPath, jsonData)
-          debug("Updated partition reassignment path with %s".format(jsonData))
-        } catch {
-          case _: ZkNoNodeException =>
-            createPersistentPath(zkPath, jsonData)
-            debug("Created path %s with %s for partition reassignment".format(zkPath, jsonData))
-          case e2: Throwable => throw new AdminOperationException(e2.toString)
-        }
-    }
-  }
-
-  def getPartitionsUndergoingPreferredReplicaElection(): Set[TopicAndPartition] = {
-    // read the partitions and their new replica list
-    val jsonPartitionListOpt = readDataMaybeNull(PreferredReplicaLeaderElectionPath)._1
-    jsonPartitionListOpt match {
-      case Some(jsonPartitionList) => PreferredReplicaLeaderElectionCommand.parsePreferredReplicaElectionData(jsonPartitionList).map(tp => new TopicAndPartition(tp))
-      case None => Set.empty[TopicAndPartition]
-    }
-  }
-
-  def deletePartition(brokerId: Int, topic: String) {
-    val brokerIdPath = BrokerIdsPath + "/" + brokerId
-    zkClient.delete(brokerIdPath)
-    val brokerPartTopicPath = ZkUtils.BrokerTopicsPath + "/" + topic + "/" + brokerId
-    zkClient.delete(brokerPartTopicPath)
-  }
-
-  /**
-   * This API takes in a broker id, queries zookeeper for the broker metadata and returns the metadata for that broker
-   * or throws an exception if the broker dies before the query to zookeeper finishes
-   *
-   * @param brokerId The broker id
-   * @return An optional Broker object encapsulating the broker metadata
-   */
-  def getBrokerInfo(brokerId: Int): Option[Broker] = {
-    readDataMaybeNull(BrokerIdsPath + "/" + brokerId)._1 match {
-      case Some(brokerInfo) => Some(parseBrokerJson(brokerId, brokerInfo))
-      case None => None
-    }
-  }
-
-  /**
-    * This API produces a sequence number by creating / updating given path in zookeeper
-    * It uses the stat returned by the zookeeper and return the version. Every time
-    * client updates the path stat.version gets incremented. Starting value of sequence number is 1.
-    */
-  def getSequenceId(path: String, acls: java.util.List[ACL] = UseDefaultAcls): Int = {
-    val acl = if (acls == UseDefaultAcls) ZkUtils.defaultAcls(isSecure, path) else acls
-    def writeToZk: Int = zkClient.writeDataReturnStat(path, "", -1).getVersion
-    try {
-      writeToZk
-    } catch {
-      case _: ZkNoNodeException =>
-        makeSurePersistentPathExists(path, acl)
-        writeToZk
-    }
-  }
-
-  def getAllTopics(): Seq[String] = {
-    val topics = getChildrenParentMayNotExist(BrokerTopicsPath)
-    if(topics == null)
-      Seq.empty[String]
-    else
-      topics
-  }
-
-  /**
-   * Returns all the entities whose configs have been overridden.
-   */
-  def getAllEntitiesWithConfig(entityType: String): Seq[String] = {
-    val entities = getChildrenParentMayNotExist(getEntityConfigRootPath(entityType))
-    if(entities == null)
-      Seq.empty[String]
-    else
-      entities
-  }
-
-  def getAllPartitions(): Set[TopicAndPartition] = {
-    val topics = getChildrenParentMayNotExist(BrokerTopicsPath)
-    if (topics == null) Set.empty[TopicAndPartition]
-    else {
-      topics.flatMap { topic =>
-        // The partitions path may not exist if the topic is in the process of being deleted
-        getChildrenParentMayNotExist(getTopicPartitionsPath(topic)).map(_.toInt).map(TopicAndPartition(topic, _))
-      }.toSet
-    }
-  }
-
-  def close() {
-    zkClient.close()
-  }
-}
-
-private object ZKStringSerializer extends ZkSerializer {
-
-  @throws(classOf[ZkMarshallingError])
-  def serialize(data : Object): Array[Byte] = data.asInstanceOf[String].getBytes("UTF-8")
-
-  @throws(classOf[ZkMarshallingError])
-  def deserialize(bytes : Array[Byte]): Object = {
-    if (bytes == null)
-      null
-    else
-      new String(bytes, "UTF-8")
-  }
-}
-
-object ZKConfig {
-  val ZkConnectProp = "zookeeper.connect"
-  val ZkSessionTimeoutMsProp = "zookeeper.session.timeout.ms"
-  val ZkConnectionTimeoutMsProp = "zookeeper.connection.timeout.ms"
-  val ZkSyncTimeMsProp = "zookeeper.sync.time.ms"
-}
-
-class ZKConfig(props: VerifiableProperties) {
-  import ZKConfig._
-
-  /** ZK host string */
-  val zkConnect = props.getString(ZkConnectProp)
-
-  /** zookeeper session timeout */
-  val zkSessionTimeoutMs = props.getInt(ZkSessionTimeoutMsProp, 6000)
-
-  /** the max time that the client waits to establish a connection to zookeeper */
-  val zkConnectionTimeoutMs = props.getInt(ZkConnectionTimeoutMsProp, zkSessionTimeoutMs)
-
-  /** how far a ZK follower can be behind a ZK leader */
-  val zkSyncTimeMs = props.getInt(ZkSyncTimeMsProp, 2000)
-}
-
-class ZkPath(zkClient: ZkClient) {
-
-  @volatile private var isNamespacePresent: Boolean = false
-
-  def checkNamespace() {
-    if (isNamespacePresent)
-      return
-
-    if (!zkClient.exists("/")) {
-      throw new ConfigException("Zookeeper namespace does not exist")
-    }
-    isNamespacePresent = true
-  }
-
-  def resetNamespaceCheckedState() {
-    isNamespacePresent = false
-  }
-
-  def createPersistent(path: String, data: Object, acls: java.util.List[ACL]) {
-    checkNamespace()
-    zkClient.createPersistent(path, data, acls)
-  }
-
-  def createPersistent(path: String, createParents: Boolean, acls: java.util.List[ACL]) {
-    checkNamespace()
-    zkClient.createPersistent(path, createParents, acls)
-  }
-
-  def createEphemeral(path: String, data: Object, acls: java.util.List[ACL]) {
-    checkNamespace()
-    zkClient.createEphemeral(path, data, acls)
-  }
-
-  def createPersistentSequential(path: String, data: Object, acls: java.util.List[ACL]): String = {
-    checkNamespace()
-    zkClient.createPersistentSequential(path, data, acls)
-  }
-}
diff --git a/core/src/main/scala/kafka/utils/json/DecodeJson.scala b/core/src/main/scala/kafka/utils/json/DecodeJson.scala
index 0b57fa7b94fe2..6962c03af2f49 100644
--- a/core/src/main/scala/kafka/utils/json/DecodeJson.scala
+++ b/core/src/main/scala/kafka/utils/json/DecodeJson.scala
@@ -17,10 +17,10 @@
 
 package kafka.utils.json
 
-import scala.collection._
+import scala.collection.{Map, Seq}
+import scala.collection.compat._
 import scala.language.higherKinds
-import JavaConverters._
-import generic.CanBuildFrom
+import scala.collection.JavaConverters._
 
 import com.fasterxml.jackson.databind.{JsonMappingException, JsonNode}
 
@@ -88,7 +88,7 @@ object DecodeJson {
     }
   }
 
-  implicit def decodeSeq[E, S[+T] <: Seq[E]](implicit decodeJson: DecodeJson[E], cbf: CanBuildFrom[Nothing, E, S[E]]): DecodeJson[S[E]] = new DecodeJson[S[E]] {
+  implicit def decodeSeq[E, S[+T] <: Seq[E]](implicit decodeJson: DecodeJson[E], factory: Factory[E, S[E]]): DecodeJson[S[E]] = new DecodeJson[S[E]] {
     def decodeEither(node: JsonNode): Either[String, S[E]] = {
       if (node.isArray)
         decodeIterator(node.elements.asScala)(decodeJson.decodeEither)
@@ -96,16 +96,16 @@ object DecodeJson {
     }
   }
 
-  implicit def decodeMap[V, M[K, +V] <: Map[K, V]](implicit decodeJson: DecodeJson[V], cbf: CanBuildFrom[Nothing, (String, V), M[String, V]]): DecodeJson[M[String, V]] = new DecodeJson[M[String, V]] {
+  implicit def decodeMap[V, M[K, +V] <: Map[K, V]](implicit decodeJson: DecodeJson[V], factory: Factory[(String, V), M[String, V]]): DecodeJson[M[String, V]] = new DecodeJson[M[String, V]] {
     def decodeEither(node: JsonNode): Either[String, M[String, V]] = {
       if (node.isObject)
-        decodeIterator(node.fields.asScala)(e => decodeJson.decodeEither(e.getValue).right.map(v => (e.getKey, v)))(cbf)
+        decodeIterator(node.fields.asScala)(e => decodeJson.decodeEither(e.getValue).right.map(v => (e.getKey, v)))
       else Left(s"Expected JSON object, received $node")
     }
   }
 
-  private def decodeIterator[S, T, C](it: Iterator[S])(f: S => Either[String, T])(implicit cbf: CanBuildFrom[Nothing, T, C]): Either[String, C] = {
-    val result = cbf()
+  private def decodeIterator[S, T, C](it: Iterator[S])(f: S => Either[String, T])(implicit factory: Factory[T, C]): Either[String, C] = {
+    val result = factory.newBuilder
     while (it.hasNext) {
       f(it.next) match {
         case Right(x) => result += x
diff --git a/core/src/main/scala/kafka/utils/timer/Timer.scala b/core/src/main/scala/kafka/utils/timer/Timer.scala
index ae8caff863658..79994a5af9b70 100644
--- a/core/src/main/scala/kafka/utils/timer/Timer.scala
+++ b/core/src/main/scala/kafka/utils/timer/Timer.scala
@@ -122,7 +122,7 @@ class SystemTimer(executorName: String,
 
   def size: Int = taskCounter.get
 
-  override def shutdown() {
+  override def shutdown(): Unit = {
     taskExecutor.shutdown()
   }
 
diff --git a/core/src/main/scala/kafka/zk/AdminZkClient.scala b/core/src/main/scala/kafka/zk/AdminZkClient.scala
index b10a089f058a6..f3a50ea483952 100644
--- a/core/src/main/scala/kafka/zk/AdminZkClient.scala
+++ b/core/src/main/scala/kafka/zk/AdminZkClient.scala
@@ -20,6 +20,7 @@ import java.util.Properties
 
 import kafka.admin.{AdminOperationException, AdminUtils, BrokerMetadata, RackAwareMode}
 import kafka.common.TopicAlreadyMarkedForDeletionException
+import kafka.controller.ReplicaAssignment
 import kafka.log.LogConfig
 import kafka.server.{ConfigEntityName, ConfigType, DynamicConfig}
 import kafka.utils._
@@ -50,9 +51,10 @@ class AdminZkClient(zkClient: KafkaZkClient) extends Logging {
                   partitions: Int,
                   replicationFactor: Int,
                   topicConfig: Properties = new Properties,
-                  rackAwareMode: RackAwareMode = RackAwareMode.Enforced) {
+                  rackAwareMode: RackAwareMode = RackAwareMode.Enforced): Unit = {
     val brokerMetadatas = getBrokerMetadatas(rackAwareMode)
-    val replicaAssignment = AdminUtils.assignReplicasToBrokers(brokerMetadatas, partitions, replicationFactor)
+    val noNewPartitionBrokerIds = getMaintenanceBrokerList()
+    val replicaAssignment = assignReplicasToAvailableBrokers(brokerMetadatas, noNewPartitionBrokerIds.toSet, partitions, replicationFactor)
     createTopicWithAssignment(topic, topicConfig, replicaAssignment)
   }
 
@@ -80,6 +82,16 @@ class AdminZkClient(zkClient: KafkaZkClient) extends Logging {
     brokerMetadatas.sortBy(_.id)
   }
 
+  /**
+    * fetch maintenance broker list from zk
+    */
+  def getMaintenanceBrokerList(): Seq[Int] = {
+    val maintenanceBrokerConfig = fetchEntityConfig(ConfigType.Broker, ConfigEntityName.Default)
+      .getProperty(DynamicConfig.Broker.MaintenanceBrokerListProp, DynamicConfig.Broker.DefaultMaintenanceBrokerList)
+
+    DynamicConfig.Broker.getMaintenanceBrokerListFromString(maintenanceBrokerConfig)
+  }
+
   def createTopicWithAssignment(topic: String,
                                 config: Properties,
                                 partitionReplicaAssignment: Map[Int, Seq[Int]]): Unit = {
@@ -92,7 +104,7 @@ class AdminZkClient(zkClient: KafkaZkClient) extends Logging {
     zkClient.setOrCreateEntityConfigs(ConfigType.Topic, topic, config)
 
     // create the partition assignment
-    writeTopicPartitionAssignment(topic, partitionReplicaAssignment, isUpdate = false)
+    writeTopicPartitionAssignment(topic, partitionReplicaAssignment.mapValues(ReplicaAssignment(_)).toMap, isUpdate = false)
   }
 
   /**
@@ -134,12 +146,13 @@ class AdminZkClient(zkClient: KafkaZkClient) extends Logging {
     LogConfig.validate(config)
   }
 
-  private def writeTopicPartitionAssignment(topic: String, replicaAssignment: Map[Int, Seq[Int]], isUpdate: Boolean) {
+
+  def writeTopicPartitionAssignment(topic: String, replicaAssignment: Map[Int, ReplicaAssignment], isUpdate: Boolean): Unit = {
     try {
       val assignment = replicaAssignment.map { case (partitionId, replicas) => (new TopicPartition(topic,partitionId), replicas) }.toMap
 
       if (!isUpdate) {
-        zkClient.createTopicAssignment(topic, assignment)
+        zkClient.createTopicAssignment(topic, assignment.mapValues(_.replicas).toMap)
       } else {
         zkClient.setTopicAssignment(topic, assignment)
       }
@@ -154,7 +167,7 @@ class AdminZkClient(zkClient: KafkaZkClient) extends Logging {
    * Creates a delete path for a given topic
    * @param topic
    */
-  def deleteTopic(topic: String) {
+  def deleteTopic(topic: String): Unit = {
     if (zkClient.topicExists(topic)) {
       try {
         zkClient.createDeleteTopicPath(topic)
@@ -168,27 +181,86 @@ class AdminZkClient(zkClient: KafkaZkClient) extends Logging {
     }
   }
 
+  /**
+    * Assign replicas to brokers that take new partitions.
+    * If the number of replicationFactor is greater than the number of brokers that take new partitions,
+    * all brokers are used for assignment.
+    * @param brokerMetadatas
+    * @param noNewPartitionBrokerIds
+    * @param nPartitions
+    * @param replicationFactor
+    * @param fixedStartIndex
+    * @param startPartitionId
+    * @return
+    */
+  def assignReplicasToAvailableBrokers(brokerMetadatas: Seq[BrokerMetadata],
+    noNewPartitionBrokerIds: Set[Int],
+    nPartitions: Int,
+    replicationFactor: Int,
+    fixedStartIndex: Int = -1,
+    startPartitionId: Int = -1): Map[Int, Seq[Int]] = {
+
+    val availableBrokerMetadata = brokerMetadatas.filter {
+      brokerMetadata =>
+        if (noNewPartitionBrokerIds.contains(brokerMetadata.id)) false
+        else true
+    }
+
+    if (replicationFactor > availableBrokerMetadata.size) {
+      info(s"Using all brokers for replica assignment since replicationFactor[$replicationFactor] " +
+        s"is larger than the number of nonMaintenanceBroker[${availableBrokerMetadata.size}]")
+      AdminUtils.assignReplicasToBrokers(brokerMetadatas, nPartitions, replicationFactor, fixedStartIndex, startPartitionId)
+    } else
+      AdminUtils.assignReplicasToBrokers(availableBrokerMetadata, nPartitions, replicationFactor, fixedStartIndex, startPartitionId)
+  }
+
+  /**
+   * Add partitions to existing topic with optional replica assignment and no constraints on whether brokers can be
+   * assigned new replicas.
+   *
+   * This method ensures that the API consistency is maintained with the upstream Kafka -- i.e. Originally this API
+   * exists in Apache Kafka 2.3, but has been removed from LinkedIn Kafka 2.3.
+   *
+   * @param topic Topic for adding partitions to
+   * @param existingAssignment A map from partition id to its assigned replicas
+   * @param allBrokers All brokers in the cluster
+   * @param numPartitions Number of partitions to be set
+   * @param replicaAssignment Manual replica assignment, or none
+   * @param validateOnly If true, validate the parameters without actually adding the partitions
+   * @return the updated replica assignment
+   */
+  def addPartitions(topic: String,
+                    existingAssignment: Map[Int, ReplicaAssignment],
+                    allBrokers: Seq[BrokerMetadata],
+                    numPartitions: Int = 1,
+                    replicaAssignment: Option[Map[Int, Seq[Int]]] = None,
+                    validateOnly: Boolean = false): Map[Int, Seq[Int]] = {
+    addPartitions(topic, existingAssignment, allBrokers, numPartitions, replicaAssignment, validateOnly, Set.empty[Int])
+  }
+
   /**
   * Add partitions to existing topic with optional replica assignment
   *
   * @param topic Topic for adding partitions to
-  * @param existingAssignment A map from partition id to its assigned replicas
+  * @param existingAssignment A map from partition id to its assignment
   * @param allBrokers All brokers in the cluster
   * @param numPartitions Number of partitions to be set
   * @param replicaAssignment Manual replica assignment, or none
   * @param validateOnly If true, validate the parameters without actually adding the partitions
+  * @param noNewPartitionBrokerIds Brokers that do not take new partitions
   * @return the updated replica assignment
   */
   def addPartitions(topic: String,
-                    existingAssignment: Map[Int, Seq[Int]],
+                    existingAssignment: Map[Int, ReplicaAssignment],
                     allBrokers: Seq[BrokerMetadata],
-                    numPartitions: Int = 1,
-                    replicaAssignment: Option[Map[Int, Seq[Int]]] = None,
-                    validateOnly: Boolean = false): Map[Int, Seq[Int]] = {
+                    numPartitions: Int,
+                    replicaAssignment: Option[Map[Int, Seq[Int]]],
+                    validateOnly: Boolean,
+                    noNewPartitionBrokerIds: Set[Int]): Map[Int, Seq[Int]] = {
     val existingAssignmentPartition0 = existingAssignment.getOrElse(0,
       throw new AdminOperationException(
         s"Unexpected existing replica assignment for topic '$topic', partition id 0 is missing. " +
-          s"Assignment: $existingAssignment"))
+          s"Assignment: $existingAssignment")).replicas
 
     val partitionsToAdd = numPartitions - existingAssignment.size
     if (partitionsToAdd <= 0)
@@ -204,18 +276,20 @@ class AdminZkClient(zkClient: KafkaZkClient) extends Logging {
 
     val proposedAssignmentForNewPartitions = replicaAssignment.getOrElse {
       val startIndex = math.max(0, allBrokers.indexWhere(_.id >= existingAssignmentPartition0.head))
-      AdminUtils.assignReplicasToBrokers(allBrokers, partitionsToAdd, existingAssignmentPartition0.size,
-        startIndex, existingAssignment.size)
+      assignReplicasToAvailableBrokers(allBrokers, noNewPartitionBrokerIds, partitionsToAdd,
+        existingAssignmentPartition0.size, startIndex, existingAssignment.size)
     }
 
-    val proposedAssignment = existingAssignment ++ proposedAssignmentForNewPartitions
+    val proposedAssignment = existingAssignment ++ proposedAssignmentForNewPartitions.map { case (tp, replicas) =>
+      tp -> ReplicaAssignment(replicas, List(), List())
+    }
     if (!validateOnly) {
       info(s"Creating $partitionsToAdd partitions for '$topic' with the following replica assignment: " +
         s"$proposedAssignmentForNewPartitions.")
 
       writeTopicPartitionAssignment(topic, proposedAssignment, isUpdate = true)
     }
-    proposedAssignment
+    proposedAssignment.mapValues(_.replicas).toMap
   }
 
   private def validateReplicaAssignment(replicaAssignment: Map[Int, Seq[Int]],
@@ -289,7 +363,7 @@ class AdminZkClient(zkClient: KafkaZkClient) extends Logging {
    *                 existing configs need to be deleted, it should be done prior to invoking this API
    *
    */
-  def changeClientIdConfig(sanitizedClientId: String, configs: Properties) {
+  def changeClientIdConfig(sanitizedClientId: String, configs: Properties): Unit = {
     DynamicConfig.Client.validate(configs)
     changeEntityConfig(ConfigType.Client, sanitizedClientId, configs)
   }
@@ -304,7 +378,7 @@ class AdminZkClient(zkClient: KafkaZkClient) extends Logging {
    *                 existing configs need to be deleted, it should be done prior to invoking this API
    *
    */
-  def changeUserOrUserClientIdConfig(sanitizedEntityName: String, configs: Properties) {
+  def changeUserOrUserClientIdConfig(sanitizedEntityName: String, configs: Properties): Unit = {
     if (sanitizedEntityName == ConfigEntityName.Default || sanitizedEntityName.contains("/clients"))
       DynamicConfig.Client.validate(configs)
     else
@@ -373,7 +447,7 @@ class AdminZkClient(zkClient: KafkaZkClient) extends Logging {
     DynamicConfig.Broker.validate(configs)
   }
 
-  private def changeEntityConfig(rootEntityType: String, fullSanitizedEntityName: String, configs: Properties) {
+  private def changeEntityConfig(rootEntityType: String, fullSanitizedEntityName: String, configs: Properties): Unit = {
     val sanitizedEntityPath = rootEntityType + '/' + fullSanitizedEntityName
     zkClient.setOrCreateEntityConfigs(rootEntityType, fullSanitizedEntityName, configs)
 
diff --git a/core/src/main/scala/kafka/zk/KafkaZkClient.scala b/core/src/main/scala/kafka/zk/KafkaZkClient.scala
index 6d8d50443e9cd..4f4eedc97a3c8 100644
--- a/core/src/main/scala/kafka/zk/KafkaZkClient.scala
+++ b/core/src/main/scala/kafka/zk/KafkaZkClient.scala
@@ -21,36 +21,32 @@ import java.util.Properties
 import com.yammer.metrics.core.MetricName
 import kafka.api.LeaderAndIsr
 import kafka.cluster.Broker
-import kafka.controller.{KafkaController, LeaderIsrAndControllerEpoch}
+import kafka.controller.{KafkaController, LeaderIsrAndControllerEpoch, ReplicaAssignment}
 import kafka.log.LogConfig
 import kafka.metrics.KafkaMetricsGroup
-import kafka.security.auth.SimpleAclAuthorizer.{NoAcls, VersionedAcls}
+import kafka.security.authorizer.AclAuthorizer.{NoAcls, VersionedAcls}
 import kafka.security.auth.{Acl, Resource, ResourceType}
 import kafka.server.ConfigType
 import kafka.utils.Logging
 import kafka.zookeeper._
-import org.apache.kafka.common.{KafkaException, TopicPartition}
-import org.apache.kafka.common.resource.PatternType
 import org.apache.kafka.common.errors.ControllerMovedException
+import org.apache.kafka.common.resource.PatternType
 import org.apache.kafka.common.security.token.delegation.{DelegationToken, TokenInformation}
 import org.apache.kafka.common.utils.{Time, Utils}
+import org.apache.kafka.common.{KafkaException, TopicPartition}
 import org.apache.zookeeper.KeeperException.{Code, NodeExistsException}
 import org.apache.zookeeper.OpResult.{CreateResult, ErrorResult, SetDataResult}
 import org.apache.zookeeper.data.{ACL, Stat}
 import org.apache.zookeeper.{CreateMode, KeeperException, ZooKeeper}
-
-import scala.collection.mutable.ArrayBuffer
-import scala.collection.{Seq, mutable}
+import scala.collection.{Map, Seq, mutable}
 
 /**
  * Provides higher level Kafka-specific operations on top of the pipelined [[kafka.zookeeper.ZooKeeperClient]].
  *
- * This performs better than [[kafka.utils.ZkUtils]] and should replace it completely, eventually.
- *
  * Implementation note: this class includes methods for various components (Controller, Configs, Old Consumer, etc.)
- * and returns instances of classes from the calling packages in some cases. This is not ideal, but it makes it
- * easier to quickly migrate away from `ZkUtils`. We should revisit this once the migration is completed and tests are
- * in place. We should also consider whether a monolithic [[kafka.zk.ZkData]] is the way to go.
+ * and returns instances of classes from the calling packages in some cases. This is not ideal, but it made it
+ * easier to migrate away from `ZkUtils` (since removed). We should revisit this. We should also consider whether a
+ * monolithic [[kafka.zk.ZkData]] is the way to go.
  */
 class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boolean, time: Time) extends AutoCloseable with
   Logging with KafkaMetricsGroup {
@@ -99,6 +95,16 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
     stat.getCzxid
   }
 
+  /**
+    * Registers the preferred controller id in zookeeper.
+    * @param id controller id
+    */
+  def registerPreferredControllerId(id: Int): Unit = {
+    val path = PreferredControllerIdZNode.path(id)
+    val stat = checkedEphemeralCreate(path, null)
+    info(s"Registered preferred controller ${id} at path $path with czxid (preferred controller epoch): ${stat.getCzxid}")
+  }
+
   /**
    * Registers a given broker in zookeeper as the controller and increments controller epoch.
    * @param controllerId the id of the broker that is to be registered as the controller.
@@ -250,10 +256,11 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
    * @param expectedControllerEpochZkVersion expected controller epoch zkVersion.
    * @return UpdateLeaderAndIsrResult instance containing per partition results.
    */
-  def updateLeaderAndIsr(leaderAndIsrs: Map[TopicPartition, LeaderAndIsr], controllerEpoch: Int, expectedControllerEpochZkVersion: Int): UpdateLeaderAndIsrResult = {
-    val successfulUpdates = mutable.Map.empty[TopicPartition, LeaderAndIsr]
-    val updatesToRetry = mutable.Buffer.empty[TopicPartition]
-    val failed = mutable.Map.empty[TopicPartition, Exception]
+  def updateLeaderAndIsr(
+    leaderAndIsrs: Map[TopicPartition, LeaderAndIsr],
+    controllerEpoch: Int,
+    expectedControllerEpochZkVersion: Int
+  ): UpdateLeaderAndIsrResult = {
     val leaderIsrAndControllerEpochs = leaderAndIsrs.map { case (partition, leaderAndIsr) =>
       partition -> LeaderIsrAndControllerEpoch(leaderAndIsr, controllerEpoch)
     }
@@ -262,20 +269,26 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
     } catch {
       case e: ControllerMovedException => throw e
       case e: Exception =>
-        leaderAndIsrs.keys.foreach(partition => failed.put(partition, e))
-        return UpdateLeaderAndIsrResult(successfulUpdates.toMap, updatesToRetry, failed.toMap)
+        return UpdateLeaderAndIsrResult(leaderAndIsrs.keys.iterator.map(_ -> Left(e)).toMap, Seq.empty)
     }
-    setDataResponses.foreach { setDataResponse =>
+
+    val updatesToRetry = mutable.Buffer.empty[TopicPartition]
+    val finished = setDataResponses.iterator.flatMap { setDataResponse =>
       val partition = setDataResponse.ctx.get.asInstanceOf[TopicPartition]
       setDataResponse.resultCode match {
         case Code.OK =>
           val updatedLeaderAndIsr = leaderAndIsrs(partition).withZkVersion(setDataResponse.stat.getVersion)
-          successfulUpdates.put(partition, updatedLeaderAndIsr)
-        case Code.BADVERSION => updatesToRetry += partition
-        case _ => failed.put(partition, setDataResponse.resultException.get)
+          Some(partition -> Right(updatedLeaderAndIsr))
+        case Code.BADVERSION =>
+          // Update the buffer for partitions to retry
+          updatesToRetry += partition
+          None
+        case _ =>
+          Some(partition -> Left(setDataResponse.resultException.get))
       }
-    }
-    UpdateLeaderAndIsrResult(successfulUpdates.toMap, updatesToRetry, failed.toMap)
+    }.toMap
+
+    UpdateLeaderAndIsrResult(finished, updatesToRetry)
   }
 
   /**
@@ -286,8 +299,10 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
    *         1. The successfully gathered log configs
    *         2. Exceptions corresponding to failed log config lookups.
    */
-  def getLogConfigs(topics: Seq[String], config: java.util.Map[String, AnyRef]):
-  (Map[String, LogConfig], Map[String, Exception]) = {
+  def getLogConfigs(
+    topics: Set[String],
+    config: java.util.Map[String, AnyRef]
+  ): (Map[String, LogConfig], Map[String, Exception]) = {
     val logConfigs = mutable.Map.empty[String, LogConfig]
     val failed = mutable.Map.empty[String, Exception]
     val configResponses = try {
@@ -449,15 +464,20 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
    */
   def getSortedBrokerList: Seq[Int] = getChildren(BrokerIdsZNode.path).map(_.toInt).sorted
 
+  /**
+    * Gets the list of preferred controller Ids
+    */
+  def getPreferredControllerList: Seq[Int] = getChildren(PreferredControllersZNode.path).map(_.toInt)
+
   /**
    * Gets all topics in the cluster.
    * @return sequence of topics in the cluster.
    */
-  def getAllTopicsInCluster: Seq[String] = {
+  def getAllTopicsInCluster: Set[String] = {
     val getChildrenResponse = retryRequestUntilConnected(GetChildrenRequest(TopicsZNode.path))
     getChildrenResponse.resultCode match {
-      case Code.OK => getChildrenResponse.children
-      case Code.NONODE => Seq.empty
+      case Code.OK => getChildrenResponse.children.toSet
+      case Code.NONODE => Set.empty
       case _ => throw getChildrenResponse.resultException.get
     }
 
@@ -479,7 +499,9 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
    * @param expectedControllerEpochZkVersion expected controller epoch zkVersion.
    * @return SetDataResponse
    */
-  def setTopicAssignmentRaw(topic: String, assignment: collection.Map[TopicPartition, Seq[Int]], expectedControllerEpochZkVersion: Int): SetDataResponse = {
+  def setTopicAssignmentRaw(topic: String,
+                            assignment: collection.Map[TopicPartition, ReplicaAssignment],
+                            expectedControllerEpochZkVersion: Int): SetDataResponse = {
     val setDataRequest = SetDataRequest(TopicZNode.path(topic), TopicZNode.encode(assignment), ZkVersion.MatchAnyVersion)
     retryRequestUntilConnected(setDataRequest, expectedControllerEpochZkVersion)
   }
@@ -491,7 +513,9 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
    * @param expectedControllerEpochZkVersion expected controller epoch zkVersion.
    * @throws KeeperException if there is an error while setting assignment
    */
-  def setTopicAssignment(topic: String, assignment: Map[TopicPartition, Seq[Int]], expectedControllerEpochZkVersion: Int = ZkVersion.MatchAnyVersion) = {
+  def setTopicAssignment(topic: String,
+                         assignment: Map[TopicPartition, ReplicaAssignment],
+                         expectedControllerEpochZkVersion: Int = ZkVersion.MatchAnyVersion) = {
     val setDataResponse = setTopicAssignmentRaw(topic, assignment, expectedControllerEpochZkVersion)
     setDataResponse.maybeThrow
   }
@@ -503,7 +527,8 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
    * @throws KeeperException if there is an error while creating assignment
    */
   def createTopicAssignment(topic: String, assignment: Map[TopicPartition, Seq[Int]]) = {
-    createRecursive(TopicZNode.path(topic), TopicZNode.encode(assignment))
+    val persistedAssignments = assignment.mapValues(ReplicaAssignment(_)).toMap
+    createRecursive(TopicZNode.path(topic), TopicZNode.encode(persistedAssignments))
   }
 
   /**
@@ -564,18 +589,28 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
   }
 
   /**
-   * Gets the assignments for the given topics.
+   * Gets the replica assignments for the given topics.
+   * This function does not return information about which replicas are being added or removed from the assignment.
    * @param topics the topics whose partitions we wish to get the assignments for.
    * @return the replica assignment for each partition from the given topics.
    */
   def getReplicaAssignmentForTopics(topics: Set[String]): Map[TopicPartition, Seq[Int]] = {
+    getFullReplicaAssignmentForTopics(topics).mapValues(_.replicas).toMap
+  }
+
+  /**
+    * Gets the replica assignments for the given topics.
+    * @param topics the topics whose partitions we wish to get the assignments for.
+    * @return the full replica assignment for each partition from the given topics.
+    */
+  def getFullReplicaAssignmentForTopics(topics: Set[String]): Map[TopicPartition, ReplicaAssignment] = {
     val getDataRequests = topics.map(topic => GetDataRequest(TopicZNode.path(topic), ctx = Some(topic)))
     val getDataResponses = retryRequestsUntilConnected(getDataRequests.toSeq)
     getDataResponses.flatMap { getDataResponse =>
       val topic = getDataResponse.ctx.get.asInstanceOf[String]
       getDataResponse.resultCode match {
         case Code.OK => TopicZNode.decode(topic, getDataResponse.data)
-        case Code.NONODE => Map.empty[TopicPartition, Seq[Int]]
+        case Code.NONODE => Map.empty[TopicPartition, ReplicaAssignment]
         case _ => throw getDataResponse.resultException.get
       }
     }.toMap
@@ -586,7 +621,7 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
    * @param topics the topics whose partitions we wish to get the assignments for.
    * @return the partition assignment for each partition from the given topics.
    */
-  def getPartitionAssignmentForTopics(topics: Set[String]): Map[String, Map[Int, Seq[Int]]] = {
+  def getPartitionAssignmentForTopics(topics: Set[String]): Map[String, Map[Int, ReplicaAssignment]] = {
     val getDataRequests = topics.map(topic => GetDataRequest(TopicZNode.path(topic), ctx = Some(topic)))
     val getDataResponses = retryRequestsUntilConnected(getDataRequests.toSeq)
     getDataResponses.flatMap { getDataResponse =>
@@ -595,13 +630,30 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
         val partitionMap = TopicZNode.decode(topic, getDataResponse.data).map { case (k, v) => (k.partition, v) }
         Map(topic -> partitionMap)
       } else if (getDataResponse.resultCode == Code.NONODE) {
-        Map.empty[String, Map[Int, Seq[Int]]]
+        Map.empty[String, Map[Int, ReplicaAssignment]]
       } else {
         throw getDataResponse.resultException.get
       }
     }.toMap
   }
 
+  def getPartitionNodeNonExistsTopics(topics: Set[String]): Set[String] = {
+    val existsRequests = topics.map(topic => ExistsRequest(TopicPartitionsZNode.path(topic), ctx = Some(topic)))
+    val existsResponses = retryRequestsUntilConnected(existsRequests.toSeq)
+    val newTopics = scala.collection.mutable.Set.empty[String]
+
+    existsResponses.foreach {
+      existsResponse =>
+        val topic = existsResponse.ctx.get.asInstanceOf[String]
+        existsResponse.resultCode match {
+          case Code.OK =>
+          case Code.NONODE => newTopics.add(topic)
+          case _ => throw existsResponse.resultException.get
+        }
+    }
+    newTopics.toSet
+  }
+
   /**
    * Gets the partition numbers for the given topics
    * @param topics the topics whose partitions we wish to get.
@@ -653,6 +705,15 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
     }
   }
 
+  /**
+    * Gets all partitions in for a topic
+    * @param topic topic name
+    * @return all partitions for the topic
+    */
+  def getTopicPartitions(topic: String) : Seq[String] = {
+    getChildren(TopicPartitionsZNode.path(topic))
+  }
+
   /**
    * Gets the data and version at the given zk path
    * @param path zk node path
@@ -773,6 +834,37 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
     }
   }
 
+  /**
+   * Creates the delete topic flag znode.
+   * @throws KeeperException if there is an error while setting or creating the znode
+   */
+  def createDeleteTopicFlagPath(): Unit = {
+    createRecursive(DeleteTopicFlagZNode.path)
+  }
+
+  /**
+   * Get topic deletion flag in zookeeper.
+   * @return topic deletion flag in zookeeper.
+   */
+  def getTopicDeletionFlag: String = {
+    val getDataResponse = retryRequestUntilConnected(GetDataRequest(DeleteTopicFlagZNode.path))
+    getDataResponse.resultCode match {
+      case Code.OK => DeleteTopicFlagZNode.decode(getDataResponse.data)
+      case _ => throw getDataResponse.resultException.get
+    }
+  }
+
+  /**
+   * Set topic deletion flag in zookeeper.
+   */
+  def setTopicDeletionFlag(flag: String): Unit = {
+    val setDataResponse = retryRequestUntilConnected(SetDataRequest(DeleteTopicFlagZNode.path, DeleteTopicFlagZNode.encode(flag), -1))
+    setDataResponse.resultCode match {
+      case Code.OK =>
+      case _ => throw setDataResponse.resultException.get
+    }
+  }
+
   /**
    * Remove the given topics from the topics marked for deletion.
    * @param topics the topics to remove.
@@ -786,8 +878,10 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
   /**
    * Returns all reassignments.
    * @return the reassignments for each partition.
+   * @deprecated Use the PartitionReassignment Kafka API instead
    */
-  def getPartitionReassignment: collection.Map[TopicPartition, Seq[Int]] = {
+  @Deprecated
+  def getPartitionReassignment(): collection.Map[TopicPartition, Seq[Int]] = {
     val getDataRequest = GetDataRequest(ReassignPartitionsZNode.path)
     val getDataResponse = retryRequestUntilConnected(getDataRequest)
     getDataResponse.resultCode match {
@@ -810,7 +904,9 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
    * @param reassignment the reassignment to set on the reassignment znode
    * @param expectedControllerEpochZkVersion expected controller epoch zkVersion.
    * @throws KeeperException if there is an error while setting or creating the znode
+   * @deprecated Use the PartitionReassignment Kafka API instead
    */
+  @Deprecated
   def setOrCreatePartitionReassignment(reassignment: collection.Map[TopicPartition, Seq[Int]], expectedControllerEpochZkVersion: Int): Unit = {
 
     def set(reassignmentData: Array[Byte]): SetDataResponse = {
@@ -838,15 +934,23 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
    * Creates the partition reassignment znode with the given reassignment.
    * @param reassignment the reassignment to set on the reassignment znode.
    * @throws KeeperException if there is an error while creating the znode
+   * @deprecated Use the PartitionReassignment Kafka API instead
    */
+  @Deprecated
   def createPartitionReassignment(reassignment: Map[TopicPartition, Seq[Int]])  = {
     createRecursive(ReassignPartitionsZNode.path, ReassignPartitionsZNode.encode(reassignment))
   }
 
+  def deletePartitionReassignment(): Unit = {
+    deletePath(ReassignPartitionsZNode.path, ZkVersion.MatchAnyVersion)
+  }
+
   /**
    * Deletes the partition reassignment znode.
    * @param expectedControllerEpochZkVersion expected controller epoch zkVersion.
+   * @deprecated Use the PartitionReassignment Kafka API instead
    */
+  @Deprecated
   def deletePartitionReassignment(expectedControllerEpochZkVersion: Int): Unit = {
     deletePath(ReassignPartitionsZNode.path, expectedControllerEpochZkVersion)
   }
@@ -854,7 +958,9 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
   /**
    * Checks if reassign partitions is in progress
    * @return true if reassign partitions is in progress, else false
+   * @deprecated Use the PartitionReassignment Kafka API instead
    */
+  @Deprecated
   def reassignPartitionsInProgress(): Boolean = {
     pathExists(ReassignPartitionsZNode.path)
   }
@@ -980,7 +1086,7 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
    * @param partitions
    * @throws KeeperException if there is an error while creating the znode
    */
-  def createPreferredReplicaElection(partitions: Set[TopicPartition]): Unit  = {
+  def createPreferredReplicaElection(partitions: Set[TopicPartition]): Unit = {
     createRecursive(PreferredReplicaElectionZNode.path, PreferredReplicaElectionZNode.encode(partitions))
   }
 
@@ -1150,7 +1256,7 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
     createResponse.maybeThrow
   }
 
-  def propagateLogDirEvent(brokerId: Int) {
+  def propagateLogDirEvent(brokerId: Int): Unit = {
     val logDirEventNotificationPath: String = createSequentialPersistentPath(
       LogDirEventNotificationZNode.path + "/" + LogDirEventNotificationSequenceZNode.SequenceNumberPrefix,
       LogDirEventNotificationSequenceZNode.encode(brokerId))
@@ -1619,7 +1725,7 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
     retryRequestsUntilConnected(createRequests, expectedControllerEpochZkVersion)
   }
 
-  private def createTopicPartitions(topics: Seq[String], expectedControllerEpochZkVersion: Int):Seq[CreateResponse] = {
+  private def createTopicPartitions(topics: Seq[String], expectedControllerEpochZkVersion: Int): Seq[CreateResponse] = {
     val createRequests = topics.map { topic =>
       val path = TopicPartitionsZNode.path(topic)
       CreateRequest(path, null, defaultAcls(path), CreateMode.PERSISTENT, Some(topic))
@@ -1627,10 +1733,11 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
     retryRequestsUntilConnected(createRequests, expectedControllerEpochZkVersion)
   }
 
-  private def getTopicConfigs(topics: Seq[String]): Seq[GetDataResponse] = {
-    val getDataRequests = topics.map { topic =>
+  private def getTopicConfigs(topics: Set[String]): Seq[GetDataResponse] = {
+    val getDataRequests: Seq[GetDataRequest] = topics.iterator.map { topic =>
       GetDataRequest(ConfigEntityZNode.path(ConfigType.Topic, topic), ctx = Some(topic))
-    }
+    }.toBuffer
+
     retryRequestsUntilConnected(getDataRequests)
   }
 
@@ -1654,8 +1761,8 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
   }
 
   private def retryRequestsUntilConnected[Req <: AsyncRequest](requests: Seq[Req]): Seq[Req#Response] = {
-    val remainingRequests = ArrayBuffer(requests: _*)
-    val responses = new ArrayBuffer[Req#Response]
+    val remainingRequests = new mutable.ArrayBuffer(requests.size) ++= requests
+    val responses = new mutable.ArrayBuffer[Req#Response]
     while (remainingRequests.nonEmpty) {
       val batchResponses = zooKeeperClient.handleRequests(remainingRequests)
 
@@ -1715,7 +1822,7 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
           SetDataOp(path, data, 0)))
       )
       val stat = response.resultCode match {
-        case code@ Code.OK =>
+        case Code.OK =>
           val setDataResult = response.zkOpResults(1).rawOpResult.asInstanceOf[SetDataResult]
           setDataResult.getStat
         case Code.NODEEXISTS =>
@@ -1798,15 +1905,16 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
 object KafkaZkClient {
 
   /**
-   * @param successfulPartitions The successfully updated partition states with adjusted znode versions.
+   * @param finishedPartitions Partitions that finished either in successfully
+   *                      updated partition states or failed with an exception.
    * @param partitionsToRetry The partitions that we should retry due to a zookeeper BADVERSION conflict. Version conflicts
    *                      can occur if the partition leader updated partition state while the controller attempted to
    *                      update partition state.
-   * @param failedPartitions Exceptions corresponding to failed partition state updates.
    */
-  case class UpdateLeaderAndIsrResult(successfulPartitions: Map[TopicPartition, LeaderAndIsr],
-                                      partitionsToRetry: Seq[TopicPartition],
-                                      failedPartitions: Map[TopicPartition, Exception])
+  case class UpdateLeaderAndIsrResult(
+    finishedPartitions: Map[TopicPartition, Either[Exception, LeaderAndIsr]],
+    partitionsToRetry: Seq[TopicPartition]
+  )
 
   /**
    * Create an instance of this class with the provided parameters.
@@ -1820,9 +1928,10 @@ object KafkaZkClient {
             maxInFlightRequests: Int,
             time: Time,
             metricGroup: String = "kafka.server",
-            metricType: String = "SessionExpireListener") = {
+            metricType: String = "SessionExpireListener",
+            name: Option[String] = None) = {
     val zooKeeperClient = new ZooKeeperClient(connectString, sessionTimeoutMs, connectionTimeoutMs, maxInFlightRequests,
-      time, metricGroup, metricType)
+      time, metricGroup, metricType, name)
     new KafkaZkClient(zooKeeperClient, isSecure, time)
   }
 
diff --git a/core/src/main/scala/kafka/zk/ZkData.scala b/core/src/main/scala/kafka/zk/ZkData.scala
index f8ad5ea333e68..601ea53a8ddc9 100644
--- a/core/src/main/scala/kafka/zk/ZkData.scala
+++ b/core/src/main/scala/kafka/zk/ZkData.scala
@@ -17,6 +17,7 @@
 package kafka.zk
 
 import java.nio.charset.StandardCharsets.UTF_8
+import java.util
 import java.util.Properties
 
 import com.fasterxml.jackson.annotation.JsonProperty
@@ -24,12 +25,13 @@ import com.fasterxml.jackson.core.JsonProcessingException
 import kafka.api.{ApiVersion, KAFKA_0_10_0_IV1, LeaderAndIsr}
 import kafka.cluster.{Broker, EndPoint}
 import kafka.common.{NotificationHandler, ZkNodeChangeNotificationListener}
-import kafka.controller.{IsrChangeNotificationHandler, LeaderIsrAndControllerEpoch}
+import kafka.controller.{IsrChangeNotificationHandler, LeaderIsrAndControllerEpoch, ReplicaAssignment}
 import kafka.security.auth.Resource.Separator
-import kafka.security.auth.SimpleAclAuthorizer.VersionedAcls
+import kafka.security.authorizer.AclAuthorizer.VersionedAcls
 import kafka.security.auth.{Acl, Resource, ResourceType}
 import kafka.server.{ConfigType, DelegationTokenManager}
 import kafka.utils.Json
+import kafka.utils.json.JsonObject
 import org.apache.kafka.common.{KafkaException, TopicPartition}
 import org.apache.kafka.common.errors.UnsupportedVersionException
 import org.apache.kafka.common.network.ListenerName
@@ -43,7 +45,7 @@ import org.apache.zookeeper.data.{ACL, Stat}
 import scala.beans.BeanProperty
 import scala.collection.JavaConverters._
 import scala.collection.mutable.ArrayBuffer
-import scala.collection.{Seq, breakOut}
+import scala.collection.{Map, Seq, mutable}
 import scala.util.{Failure, Success, Try}
 
 // This file contains objects for encoding/decoding data stored in ZooKeeper nodes (znodes).
@@ -72,6 +74,15 @@ object BrokersZNode {
   def path = "/brokers"
 }
 
+object PreferredControllersZNode {
+  def path = s"${BrokersZNode.path}/preferred_controllers"
+  def encode: Array[Byte] = null
+}
+
+object PreferredControllerIdZNode {
+  def path(id: Int) = s"${BrokersZNode.path}/preferred_controllers/$id"
+}
+
 object BrokerIdsZNode {
   def path = s"${BrokersZNode.path}/ids"
   def encode: Array[Byte] = null
@@ -238,19 +249,49 @@ object TopicsZNode {
 
 object TopicZNode {
   def path(topic: String) = s"${TopicsZNode.path}/$topic"
-  def encode(assignment: collection.Map[TopicPartition, Seq[Int]]): Array[Byte] = {
-    val assignmentJson = assignment.map { case (partition, replicas) =>
-      partition.partition.toString -> replicas.asJava
+  def encode(assignment: collection.Map[TopicPartition, ReplicaAssignment]): Array[Byte] = {
+    val replicaAssignmentJson = mutable.Map[String, util.List[Int]]()
+    val addingReplicasAssignmentJson = mutable.Map[String, util.List[Int]]()
+    val removingReplicasAssignmentJson = mutable.Map[String, util.List[Int]]()
+
+    for ((partition, replicaAssignment) <- assignment) {
+      replicaAssignmentJson += (partition.partition.toString -> replicaAssignment.replicas.asJava)
+      if (replicaAssignment.addingReplicas.nonEmpty)
+        addingReplicasAssignmentJson += (partition.partition.toString -> replicaAssignment.addingReplicas.asJava)
+      if (replicaAssignment.removingReplicas.nonEmpty)
+        removingReplicasAssignmentJson += (partition.partition.toString -> replicaAssignment.removingReplicas.asJava)
     }
-    Json.encodeAsBytes(Map("version" -> 1, "partitions" -> assignmentJson.asJava).asJava)
+
+    Json.encodeAsBytes(Map(
+      "version" -> 2,
+      "partitions" -> replicaAssignmentJson.asJava,
+      "adding_replicas" -> addingReplicasAssignmentJson.asJava,
+      "removing_replicas" -> removingReplicasAssignmentJson.asJava
+    ).asJava)
   }
-  def decode(topic: String, bytes: Array[Byte]): Map[TopicPartition, Seq[Int]] = {
+  def decode(topic: String, bytes: Array[Byte]): Map[TopicPartition, ReplicaAssignment] = {
+    def getReplicas(replicasJsonOpt: Option[JsonObject], partition: String): Seq[Int] = {
+      replicasJsonOpt match {
+        case Some(replicasJson) => replicasJson.get(partition) match {
+          case Some(ar) => ar.to[Seq[Int]]
+          case None => Seq.empty[Int]
+        }
+        case None => Seq.empty[Int]
+      }
+    }
+
     Json.parseBytes(bytes).flatMap { js =>
       val assignmentJson = js.asJsonObject
       val partitionsJsonOpt = assignmentJson.get("partitions").map(_.asJsonObject)
+      val addingReplicasJsonOpt = assignmentJson.get("adding_replicas").map(_.asJsonObject)
+      val removingReplicasJsonOpt = assignmentJson.get("removing_replicas").map(_.asJsonObject)
       partitionsJsonOpt.map { partitionsJson =>
         partitionsJson.iterator.map { case (partition, replicas) =>
-          new TopicPartition(topic, partition.toInt) -> replicas.to[Seq[Int]]
+          new TopicPartition(topic, partition.toInt) -> ReplicaAssignment(
+            replicas.to[Seq[Int]],
+            getReplicas(addingReplicasJsonOpt, partition),
+            getReplicas(removingReplicasJsonOpt, partition)
+          )
         }
       }
     }.map(_.toMap).getOrElse(Map.empty)
@@ -373,6 +414,16 @@ object DeleteTopicsTopicZNode {
   def path(topic: String) = s"${DeleteTopicsZNode.path}/$topic"
 }
 
+object DeleteTopicFlagZNode {
+  def path = "/topic_deletion_flag"
+  def encode(topicDeletionFlag: String): Array[Byte] = topicDeletionFlag.getBytes(UTF_8)
+  def decode(bytes: Array[Byte]): String = if (bytes != null) new String(bytes, UTF_8) else ""
+}
+
+/**
+ * The znode for initiating a partition reassignment.
+ * @deprecated Since 2.4, use the PartitionReassignment Kafka API instead.
+ */
 object ReassignPartitionsZNode {
 
   /**
@@ -388,14 +439,16 @@ object ReassignPartitionsZNode {
   /**
     * An assignment consists of a `version` and a list of `partitions`, which represent the
     * assignment of topic-partitions to brokers.
+    * @deprecated Use the PartitionReassignment Kafka API instead
     */
-  case class PartitionAssignment(@BeanProperty @JsonProperty("version") version: Int,
-                                 @BeanProperty @JsonProperty("partitions") partitions: java.util.List[ReplicaAssignment])
+  @Deprecated
+  case class LegacyPartitionAssignment(@BeanProperty @JsonProperty("version") version: Int,
+                                       @BeanProperty @JsonProperty("partitions") partitions: java.util.List[ReplicaAssignment])
 
   def path = s"${AdminZNode.path}/reassign_partitions"
 
   def encode(reassignmentMap: collection.Map[TopicPartition, Seq[Int]]): Array[Byte] = {
-    val reassignment = PartitionAssignment(1,
+    val reassignment = LegacyPartitionAssignment(1,
       reassignmentMap.toSeq.map { case (tp, replicas) =>
         ReplicaAssignment(tp.topic, tp.partition, replicas.asJava)
       }.asJava
@@ -404,10 +457,10 @@ object ReassignPartitionsZNode {
   }
 
   def decode(bytes: Array[Byte]): Either[JsonProcessingException, collection.Map[TopicPartition, Seq[Int]]] =
-    Json.parseBytesAs[PartitionAssignment](bytes).right.map { partitionAssignment =>
-      partitionAssignment.partitions.asScala.map { replicaAssignment =>
+    Json.parseBytesAs[LegacyPartitionAssignment](bytes).right.map { partitionAssignment =>
+      partitionAssignment.partitions.asScala.iterator.map { replicaAssignment =>
         new TopicPartition(replicaAssignment.topic, replicaAssignment.partition) -> replicaAssignment.replicas.asScala
-      }(breakOut)
+      }.toMap
     }
 }
 
@@ -731,6 +784,7 @@ object ZkData {
   val PersistentZkPaths = Seq(
     ConsumerPathZNode.path, // old consumer path
     BrokerIdsZNode.path,
+    PreferredControllersZNode.path,
     TopicsZNode.path,
     ConfigEntityChangeNotificationZNode.path,
     DeleteTopicsZNode.path,
diff --git a/core/src/main/scala/kafka/zookeeper/ZooKeeperClient.scala b/core/src/main/scala/kafka/zookeeper/ZooKeeperClient.scala
index 4d5ef96f4659b..8ae8e17de4c93 100755
--- a/core/src/main/scala/kafka/zookeeper/ZooKeeperClient.scala
+++ b/core/src/main/scala/kafka/zookeeper/ZooKeeperClient.scala
@@ -35,6 +35,7 @@ import org.apache.zookeeper.data.{ACL, Stat}
 import org.apache.zookeeper._
 
 import scala.collection.JavaConverters._
+import scala.collection.Seq
 import scala.collection.mutable.Set
 
 /**
@@ -44,6 +45,7 @@ import scala.collection.mutable.Set
  * @param sessionTimeoutMs session timeout in milliseconds
  * @param connectionTimeoutMs connection timeout in milliseconds
  * @param maxInFlightRequests maximum number of unacknowledged requests the client will send before blocking.
+ * @param name name of the client instance
  */
 class ZooKeeperClient(connectString: String,
                       sessionTimeoutMs: Int,
@@ -51,8 +53,23 @@ class ZooKeeperClient(connectString: String,
                       maxInFlightRequests: Int,
                       time: Time,
                       metricGroup: String,
-                      metricType: String) extends Logging with KafkaMetricsGroup {
-  this.logIdent = "[ZooKeeperClient] "
+                      metricType: String,
+                      name: Option[String]) extends Logging with KafkaMetricsGroup {
+
+  def this(connectString: String,
+           sessionTimeoutMs: Int,
+           connectionTimeoutMs: Int,
+           maxInFlightRequests: Int,
+           time: Time,
+           metricGroup: String,
+           metricType: String) = {
+    this(connectString, sessionTimeoutMs, connectionTimeoutMs, maxInFlightRequests, time, metricGroup, metricType, None)
+  }
+
+  this.logIdent = name match {
+    case Some(n) => s"[ZooKeeperClient $n] "
+    case _ => "[ZooKeeperClient] "
+  }
   private val initializationLock = new ReentrantReadWriteLock()
   private val isConnectedOrExpiredLock = new ReentrantLock()
   private val isConnectedOrExpiredCondition = isConnectedOrExpiredLock.newCondition()
@@ -320,6 +337,12 @@ class ZooKeeperClient(connectString: String,
 
   def close(): Unit = {
     info("Closing.")
+
+    // Shutdown scheduler outside of lock to avoid deadlock if scheduler
+    // is waiting for lock to process session expiry. Close expiry thread
+    // first to ensure that new clients are not created during close().
+    expiryScheduler.shutdown()
+
     inWriteLock(initializationLock) {
       zNodeChangeHandlers.clear()
       zNodeChildChangeHandlers.clear()
@@ -327,9 +350,6 @@ class ZooKeeperClient(connectString: String,
       zooKeeper.close()
       metricNames.foreach(removeMetric(_))
     }
-    // Shutdown scheduler outside of lock to avoid deadlock if scheduler
-    // is waiting for lock to process session expiry
-    expiryScheduler.shutdown()
     info("Closed.")
   }
 
diff --git a/core/src/test/resources/log4j.properties b/core/src/test/resources/log4j.properties
index d394a2abc9495..0648d7f3fc00e 100644
--- a/core/src/test/resources/log4j.properties
+++ b/core/src/test/resources/log4j.properties
@@ -22,5 +22,4 @@ log4j.logger.kafka=ERROR
 log4j.logger.org.apache.kafka=ERROR
 
 # zkclient can be verbose, during debugging it is common to adjust it separately
-log4j.logger.org.I0Itec.zkclient.ZkClient=WARN
 log4j.logger.org.apache.zookeeper=WARN
diff --git a/core/src/test/scala/integration/kafka/admin/BrokerApiVersionsCommandTest.scala b/core/src/test/scala/integration/kafka/admin/BrokerApiVersionsCommandTest.scala
index 1f4b5e291d38e..28b1f069a2710 100644
--- a/core/src/test/scala/integration/kafka/admin/BrokerApiVersionsCommandTest.scala
+++ b/core/src/test/scala/integration/kafka/admin/BrokerApiVersionsCommandTest.scala
@@ -20,6 +20,8 @@ package kafka.admin
 import java.io.{ByteArrayOutputStream, PrintStream}
 import java.nio.charset.StandardCharsets
 
+import scala.collection.Seq
+
 import kafka.integration.KafkaServerTestHarness
 import kafka.server.KafkaConfig
 import kafka.utils.TestUtils
@@ -33,7 +35,7 @@ class BrokerApiVersionsCommandTest extends KafkaServerTestHarness {
   def generateConfigs: Seq[KafkaConfig] = TestUtils.createBrokerConfigs(1, zkConnect).map(KafkaConfig.fromProps)
 
   @Test(timeout=120000)
-  def checkBrokerApiVersionCommandOutput() {
+  def checkBrokerApiVersionCommandOutput(): Unit = {
     val byteArrayOutputStream = new ByteArrayOutputStream
     val printStream = new PrintStream(byteArrayOutputStream, false, StandardCharsets.UTF_8.name())
     BrokerApiVersionsCommand.execute(Array("--bootstrap-server", brokerList), printStream)
diff --git a/core/src/test/scala/integration/kafka/admin/ReassignPartitionsIntegrationTest.scala b/core/src/test/scala/integration/kafka/admin/ReassignPartitionsIntegrationTest.scala
index 19a6b47da86bc..c4733860ff006 100644
--- a/core/src/test/scala/integration/kafka/admin/ReassignPartitionsIntegrationTest.scala
+++ b/core/src/test/scala/integration/kafka/admin/ReassignPartitionsIntegrationTest.scala
@@ -20,7 +20,7 @@ import org.junit.Test
 class ReassignPartitionsIntegrationTest extends ZooKeeperTestHarness with RackAwareTest {
 
   @Test
-  def testRackAwareReassign() {
+  def testRackAwareReassign(): Unit = {
     val rackInfo = Map(0 -> "rack1", 1 -> "rack2", 2 -> "rack2", 3 -> "rack1", 4 -> "rack3", 5 -> "rack3")
     TestUtils.createBrokersInZk(toBrokerMetadata(rackInfo), zkClient)
 
diff --git a/core/src/test/scala/integration/kafka/api/AbstractConsumerTest.scala b/core/src/test/scala/integration/kafka/api/AbstractConsumerTest.scala
new file mode 100644
index 0000000000000..fcd1402126ce9
--- /dev/null
+++ b/core/src/test/scala/integration/kafka/api/AbstractConsumerTest.scala
@@ -0,0 +1,441 @@
+/*
+ * 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 java.time.Duration
+import java.util
+import java.util.Properties
+
+import org.apache.kafka.clients.consumer._
+import org.apache.kafka.clients.producer.{ProducerConfig, ProducerRecord}
+import org.apache.kafka.common.record.TimestampType
+import org.apache.kafka.common.TopicPartition
+import kafka.utils.{ShutdownableThread, TestUtils}
+import kafka.server.{BaseRequestTest, KafkaConfig}
+import org.junit.Assert._
+import org.junit.Before
+
+import scala.collection.JavaConverters._
+import scala.collection.mutable.{ArrayBuffer, Buffer}
+import org.apache.kafka.clients.producer.KafkaProducer
+import org.apache.kafka.common.errors.WakeupException
+
+import scala.collection.mutable
+
+/**
+ * Extension point for consumer integration tests.
+ */
+abstract class AbstractConsumerTest extends BaseRequestTest {
+
+  val epsilon = 0.1
+  override def brokerCount: Int = 3
+
+  val topic = "topic"
+  val part = 0
+  val tp = new TopicPartition(topic, part)
+  val part2 = 1
+  val tp2 = new TopicPartition(topic, part2)
+  val group = "my-test"
+  val producerClientId = "ConsumerTestProducer"
+  val consumerClientId = "ConsumerTestConsumer"
+  val groupMaxSessionTimeoutMs = 30000L
+
+  this.producerConfig.setProperty(ProducerConfig.ACKS_CONFIG, "all")
+  this.producerConfig.setProperty(ProducerConfig.CLIENT_ID_CONFIG, producerClientId)
+  this.consumerConfig.setProperty(ConsumerConfig.CLIENT_ID_CONFIG, consumerClientId)
+  this.consumerConfig.setProperty(ConsumerConfig.GROUP_ID_CONFIG, group)
+  this.consumerConfig.setProperty(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest")
+  this.consumerConfig.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "false")
+  this.consumerConfig.setProperty(ConsumerConfig.METADATA_MAX_AGE_CONFIG, "100")
+  this.consumerConfig.setProperty(ConsumerConfig.MAX_POLL_INTERVAL_MS_CONFIG, "6000")
+
+
+  override protected def brokerPropertyOverrides(properties: Properties): Unit = {
+    properties.setProperty(KafkaConfig.ControlledShutdownEnableProp, "false") // speed up shutdown
+    properties.setProperty(KafkaConfig.OffsetsTopicReplicationFactorProp, "3") // don't want to lose offset
+    properties.setProperty(KafkaConfig.OffsetsTopicPartitionsProp, "1")
+    properties.setProperty(KafkaConfig.GroupMinSessionTimeoutMsProp, "100") // set small enough session timeout
+    properties.setProperty(KafkaConfig.GroupMaxSessionTimeoutMsProp, groupMaxSessionTimeoutMs.toString)
+    properties.setProperty(KafkaConfig.GroupInitialRebalanceDelayMsProp, "10")
+  }
+
+  @Before
+  override def setUp(): Unit = {
+    super.setUp()
+
+    // create the test topic with all the brokers as replicas
+    createTopic(topic, 2, brokerCount)
+  }
+
+  protected class TestConsumerReassignmentListener extends ConsumerRebalanceListener {
+    var callsToAssigned = 0
+    var callsToRevoked = 0
+
+    def onPartitionsAssigned(partitions: java.util.Collection[TopicPartition]): Unit = {
+      info("onPartitionsAssigned called.")
+      callsToAssigned += 1
+    }
+
+    def onPartitionsRevoked(partitions: java.util.Collection[TopicPartition]): Unit = {
+      info("onPartitionsRevoked called.")
+      callsToRevoked += 1
+    }
+  }
+
+  protected def createConsumerWithGroupId(groupId: String): KafkaConsumer[Array[Byte], Array[Byte]] = {
+    val groupOverrideConfig = new Properties
+    groupOverrideConfig.setProperty(ConsumerConfig.GROUP_ID_CONFIG, groupId)
+    createConsumer(configOverrides = groupOverrideConfig)
+  }
+
+  protected def sendRecords(producer: KafkaProducer[Array[Byte], Array[Byte]], numRecords: Int,
+                            tp: TopicPartition): Seq[ProducerRecord[Array[Byte], Array[Byte]]] = {
+    val records = (0 until numRecords).map { i =>
+      val record = new ProducerRecord(tp.topic(), tp.partition(), i.toLong, s"key $i".getBytes, s"value $i".getBytes)
+      producer.send(record)
+      record
+    }
+    producer.flush()
+
+    records
+  }
+
+  protected def consumeAndVerifyRecords(consumer: Consumer[Array[Byte], Array[Byte]],
+                                        numRecords: Int,
+                                        startingOffset: Int,
+                                        startingKeyAndValueIndex: Int = 0,
+                                        startingTimestamp: Long = 0L,
+                                        timestampType: TimestampType = TimestampType.CREATE_TIME,
+                                        tp: TopicPartition = tp,
+                                        maxPollRecords: Int = Int.MaxValue): Unit = {
+    val records = consumeRecords(consumer, numRecords, maxPollRecords = maxPollRecords)
+    val now = System.currentTimeMillis()
+    for (i <- 0 until numRecords) {
+      val record = records(i)
+      val offset = startingOffset + i
+      assertEquals(tp.topic, record.topic)
+      assertEquals(tp.partition, record.partition)
+      if (timestampType == TimestampType.CREATE_TIME) {
+        assertEquals(timestampType, record.timestampType)
+        val timestamp = startingTimestamp + i
+        assertEquals(timestamp.toLong, record.timestamp)
+      } else
+        assertTrue(s"Got unexpected timestamp ${record.timestamp}. Timestamp should be between [$startingTimestamp, $now}]",
+          record.timestamp >= startingTimestamp && record.timestamp <= now)
+      assertEquals(offset.toLong, record.offset)
+      val keyAndValueIndex = startingKeyAndValueIndex + i
+      assertEquals(s"key $keyAndValueIndex", new String(record.key))
+      assertEquals(s"value $keyAndValueIndex", new String(record.value))
+      // this is true only because K and V are byte arrays
+      assertEquals(s"key $keyAndValueIndex".length, record.serializedKeySize)
+      assertEquals(s"value $keyAndValueIndex".length, record.serializedValueSize)
+    }
+  }
+
+  protected def consumeRecords[K, V](consumer: Consumer[K, V],
+                                     numRecords: Int,
+                                     maxPollRecords: Int = Int.MaxValue): ArrayBuffer[ConsumerRecord[K, V]] = {
+    val records = new ArrayBuffer[ConsumerRecord[K, V]]
+    def pollAction(polledRecords: ConsumerRecords[K, V]): Boolean = {
+      assertTrue(polledRecords.asScala.size <= maxPollRecords)
+      records ++= polledRecords.asScala
+      records.size >= numRecords
+    }
+    TestUtils.pollRecordsUntilTrue(consumer, pollAction, waitTimeMs = 60000,
+      msg = s"Timed out before consuming expected $numRecords records. " +
+        s"The number consumed was ${records.size}.")
+    records
+  }
+
+  protected def sendAndAwaitAsyncCommit[K, V](consumer: Consumer[K, V],
+                                              offsetsOpt: Option[Map[TopicPartition, OffsetAndMetadata]] = None): Unit = {
+
+    def sendAsyncCommit(callback: OffsetCommitCallback) = {
+      offsetsOpt match {
+        case Some(offsets) => consumer.commitAsync(offsets.asJava, callback)
+        case None => consumer.commitAsync(callback)
+      }
+    }
+
+    class RetryCommitCallback extends OffsetCommitCallback {
+      var isComplete = false
+      var error: Option[Exception] = None
+
+      override def onComplete(offsets: util.Map[TopicPartition, OffsetAndMetadata], exception: Exception): Unit = {
+        exception match {
+          case e: RetriableCommitFailedException =>
+            sendAsyncCommit(this)
+          case e =>
+            isComplete = true
+            error = Option(e)
+        }
+      }
+    }
+
+    val commitCallback = new RetryCommitCallback
+
+    sendAsyncCommit(commitCallback)
+    TestUtils.pollUntilTrue(consumer, () => commitCallback.isComplete,
+      "Failed to observe commit callback before timeout", waitTimeMs = 10000)
+
+    assertEquals(None, commitCallback.error)
+  }
+
+  /**
+    * Create 'numOfConsumersToAdd' consumers add then to the consumer group 'consumerGroup', and create corresponding
+    * pollers for these consumers. Wait for partition re-assignment and validate.
+    *
+    * Currently, assignment validation requires that total number of partitions is greater or equal to
+    * number of consumers, so subscriptions.size must be greater or equal the resulting number of consumers in the group
+    *
+    * @param numOfConsumersToAdd number of consumers to create and add to the consumer group
+    * @param consumerGroup current consumer group
+    * @param consumerPollers current consumer pollers
+    * @param topicsToSubscribe topics to which new consumers will subscribe to
+    * @param subscriptions set of all topic partitions
+    */
+  def addConsumersToGroupAndWaitForGroupAssignment(numOfConsumersToAdd: Int,
+                                                   consumerGroup: mutable.Buffer[KafkaConsumer[Array[Byte], Array[Byte]]],
+                                                   consumerPollers: mutable.Buffer[ConsumerAssignmentPoller],
+                                                   topicsToSubscribe: List[String],
+                                                   subscriptions: Set[TopicPartition],
+                                                   group: String = group): (mutable.Buffer[KafkaConsumer[Array[Byte], Array[Byte]]], mutable.Buffer[ConsumerAssignmentPoller]) = {
+    assertTrue(consumerGroup.size + numOfConsumersToAdd <= subscriptions.size)
+    addConsumersToGroup(numOfConsumersToAdd, consumerGroup, consumerPollers, topicsToSubscribe, subscriptions, group)
+    // wait until topics get re-assigned and validate assignment
+    validateGroupAssignment(consumerPollers, subscriptions)
+
+    (consumerGroup, consumerPollers)
+  }
+
+  /**
+    * Create 'numOfConsumersToAdd' consumers add then to the consumer group 'consumerGroup', and create corresponding
+    * pollers for these consumers.
+    *
+    *
+    * @param numOfConsumersToAdd number of consumers to create and add to the consumer group
+    * @param consumerGroup current consumer group
+    * @param consumerPollers current consumer pollers
+    * @param topicsToSubscribe topics to which new consumers will subscribe to
+    * @param subscriptions set of all topic partitions
+    */
+  def addConsumersToGroup(numOfConsumersToAdd: Int,
+                          consumerGroup: mutable.Buffer[KafkaConsumer[Array[Byte], Array[Byte]]],
+                          consumerPollers: mutable.Buffer[ConsumerAssignmentPoller],
+                          topicsToSubscribe: List[String],
+                          subscriptions: Set[TopicPartition],
+                          group: String = group): (mutable.Buffer[KafkaConsumer[Array[Byte], Array[Byte]]], mutable.Buffer[ConsumerAssignmentPoller]) = {
+    for (_ <- 0 until numOfConsumersToAdd) {
+      val consumer = createConsumerWithGroupId(group)
+      consumerGroup += consumer
+      consumerPollers += subscribeConsumerAndStartPolling(consumer, topicsToSubscribe)
+    }
+
+    (consumerGroup, consumerPollers)
+  }
+
+  /**
+    * Wait for consumers to get partition assignment and validate it.
+    *
+    * @param consumerPollers consumer pollers corresponding to the consumer group we are testing
+    * @param subscriptions set of all topic partitions
+    * @param msg message to print when waiting for/validating assignment fails
+    */
+  def validateGroupAssignment(consumerPollers: mutable.Buffer[ConsumerAssignmentPoller],
+                              subscriptions: Set[TopicPartition],
+                              msg: Option[String] = None,
+                              waitTime: Long = 10000L): Unit = {
+    val assignments = mutable.Buffer[Set[TopicPartition]]()
+    TestUtils.waitUntilTrue(() => {
+      assignments.clear()
+      consumerPollers.foreach(assignments += _.consumerAssignment())
+      isPartitionAssignmentValid(assignments, subscriptions)
+    }, msg.getOrElse(s"Did not get valid assignment for partitions $subscriptions. Instead, got $assignments"), waitTime)
+  }
+
+  /**
+    * Subscribes consumer 'consumer' to a given list of topics 'topicsToSubscribe', creates
+    * consumer poller and starts polling.
+    * Assumes that the consumer is not subscribed to any topics yet
+    *
+    * @param consumer consumer
+    * @param topicsToSubscribe topics that this consumer will subscribe to
+    * @return consumer poller for the given consumer
+    */
+  def subscribeConsumerAndStartPolling(consumer: Consumer[Array[Byte], Array[Byte]],
+                                       topicsToSubscribe: List[String],
+                                       partitionsToAssign: Set[TopicPartition] = Set.empty[TopicPartition]): ConsumerAssignmentPoller = {
+    assertEquals(0, consumer.assignment().size)
+    val consumerPoller = if (topicsToSubscribe.nonEmpty)
+      new ConsumerAssignmentPoller(consumer, topicsToSubscribe)
+    else
+      new ConsumerAssignmentPoller(consumer, partitionsToAssign)
+
+    consumerPoller.start()
+    consumerPoller
+  }
+
+  protected def awaitRebalance(consumer: Consumer[_, _], rebalanceListener: TestConsumerReassignmentListener): Unit = {
+    val numReassignments = rebalanceListener.callsToAssigned
+    TestUtils.pollUntilTrue(consumer, () => rebalanceListener.callsToAssigned > numReassignments,
+      "Timed out before expected rebalance completed")
+  }
+
+  protected def ensureNoRebalance(consumer: Consumer[_, _], rebalanceListener: TestConsumerReassignmentListener): Unit = {
+    // The best way to verify that the current membership is still active is to commit offsets.
+    // This would fail if the group had rebalanced.
+    val initialRevokeCalls = rebalanceListener.callsToRevoked
+    sendAndAwaitAsyncCommit(consumer)
+    assertEquals(initialRevokeCalls, rebalanceListener.callsToRevoked)
+  }
+
+  protected class CountConsumerCommitCallback extends OffsetCommitCallback {
+    var successCount = 0
+    var failCount = 0
+    var lastError: Option[Exception] = None
+
+    override def onComplete(offsets: util.Map[TopicPartition, OffsetAndMetadata], exception: Exception): Unit = {
+      if (exception == null) {
+        successCount += 1
+      } else {
+        failCount += 1
+        lastError = Some(exception)
+      }
+    }
+  }
+
+  protected class ConsumerAssignmentPoller(consumer: Consumer[Array[Byte], Array[Byte]],
+                                           topicsToSubscribe: List[String],
+                                           partitionsToAssign: Set[TopicPartition])
+    extends ShutdownableThread("daemon-consumer-assignment", false) {
+
+    def this(consumer: Consumer[Array[Byte], Array[Byte]], topicsToSubscribe: List[String]) {
+      this(consumer, topicsToSubscribe, Set.empty[TopicPartition])
+    }
+
+    def this(consumer: Consumer[Array[Byte], Array[Byte]], partitionsToAssign: Set[TopicPartition]) {
+      this(consumer, List.empty[String], partitionsToAssign)
+    }
+
+    @volatile var thrownException: Option[Throwable] = None
+    @volatile var receivedMessages = 0
+
+    private val partitionAssignment = mutable.Set[TopicPartition]()
+    @volatile private var subscriptionChanged = false
+    private var topicsSubscription = topicsToSubscribe
+
+    val rebalanceListener: ConsumerRebalanceListener = new ConsumerRebalanceListener {
+      override def onPartitionsAssigned(partitions: util.Collection[TopicPartition]) = {
+        partitionAssignment ++= partitions.toArray(new Array[TopicPartition](0))
+      }
+
+      override def onPartitionsRevoked(partitions: util.Collection[TopicPartition]) = {
+        partitionAssignment --= partitions.toArray(new Array[TopicPartition](0))
+      }
+    }
+
+    if (partitionsToAssign.isEmpty) {
+      consumer.subscribe(topicsToSubscribe.asJava, rebalanceListener)
+    } else {
+      consumer.assign(partitionsToAssign.asJava)
+    }
+
+    def consumerAssignment(): Set[TopicPartition] = {
+      partitionAssignment.toSet
+    }
+
+    /**
+     * Subscribe consumer to a new set of topics.
+     * Since this method most likely be called from a different thread, this function
+     * just "schedules" the subscription change, and actual call to consumer.subscribe is done
+     * in the doWork() method
+     *
+     * This method does not allow to change subscription until doWork processes the previous call
+     * to this method. This is just to avoid race conditions and enough functionality for testing purposes
+     * @param newTopicsToSubscribe
+     */
+    def subscribe(newTopicsToSubscribe: List[String]): Unit = {
+      if (subscriptionChanged)
+        throw new IllegalStateException("Do not call subscribe until the previous subscribe request is processed.")
+      if (partitionsToAssign.nonEmpty)
+        throw new IllegalStateException("Cannot call subscribe when configured to use manual partition assignment")
+
+      topicsSubscription = newTopicsToSubscribe
+      subscriptionChanged = true
+    }
+
+    def isSubscribeRequestProcessed: Boolean = {
+      !subscriptionChanged
+    }
+
+    override def initiateShutdown(): Boolean = {
+      val res = super.initiateShutdown()
+      consumer.wakeup()
+      res
+    }
+
+    override def doWork(): Unit = {
+      if (subscriptionChanged) {
+        consumer.subscribe(topicsSubscription.asJava, rebalanceListener)
+        subscriptionChanged = false
+      }
+      try {
+        receivedMessages += consumer.poll(Duration.ofMillis(50)).count()
+      } catch {
+        case _: WakeupException => // ignore for shutdown
+        case e: Throwable =>
+          thrownException = Some(e)
+          throw e
+      }
+    }
+  }
+
+  /**
+   * Check whether partition assignment is valid
+   * Assumes partition assignment is valid iff
+   * 1. Every consumer got assigned at least one partition
+   * 2. Each partition is assigned to only one consumer
+   * 3. Every partition is assigned to one of the consumers
+   *
+   * @param assignments set of consumer assignments; one per each consumer
+   * @param partitions set of partitions that consumers subscribed to
+   * @return true if partition assignment is valid
+   */
+  def isPartitionAssignmentValid(assignments: Buffer[Set[TopicPartition]],
+                                 partitions: Set[TopicPartition]): Boolean = {
+    val allNonEmptyAssignments = assignments.forall(assignment => assignment.nonEmpty)
+    if (!allNonEmptyAssignments) {
+      // at least one consumer got empty assignment
+      return false
+    }
+
+    // make sure that sum of all partitions to all consumers equals total number of partitions
+    val totalPartitionsInAssignments = assignments.foldLeft(0)(_ + _.size)
+    if (totalPartitionsInAssignments != partitions.size) {
+      // either same partitions got assigned to more than one consumer or some
+      // partitions were not assigned
+      return false
+    }
+
+    // The above checks could miss the case where one or more partitions were assigned to more
+    // than one consumer and the same number of partitions were missing from assignments.
+    // Make sure that all unique assignments are the same as 'partitions'
+    val uniqueAssignedPartitions = assignments.foldLeft(Set.empty[TopicPartition])(_ ++ _)
+    uniqueAssignedPartitions == partitions
+  }
+
+}
diff --git a/core/src/test/scala/integration/kafka/api/AdminClientIntegrationTest.scala b/core/src/test/scala/integration/kafka/api/AdminClientIntegrationTest.scala
index 1ee22346f9d77..b3afa1bc1265f 100644
--- a/core/src/test/scala/integration/kafka/api/AdminClientIntegrationTest.scala
+++ b/core/src/test/scala/integration/kafka/api/AdminClientIntegrationTest.scala
@@ -16,45 +16,44 @@
  */
 package kafka.api
 
-import java.util
-import java.util.{Collections, Properties}
-import java.util.Arrays.asList
-import java.util.concurrent.{ExecutionException, TimeUnit}
 import java.io.File
+import java.lang.{Long => JLong}
+import java.time.{Duration => JDuration}
+import java.util.Arrays.asList
 import java.util.concurrent.atomic.{AtomicBoolean, AtomicInteger}
+import java.util.concurrent.{CountDownLatch, ExecutionException, TimeUnit}
+import java.util.{Collections, Optional, Properties}
+import java.{time, util}
 
-import org.apache.kafka.clients.admin.KafkaAdminClientTest
-import org.apache.kafka.common.utils.{Time, Utils}
 import kafka.log.LogConfig
+import kafka.security.auth.{Cluster, Group, Topic}
 import kafka.server.{Defaults, KafkaConfig, KafkaServer}
-import org.apache.kafka.clients.admin._
-import kafka.utils.{Logging, TestUtils}
-import kafka.utils.TestUtils._
 import kafka.utils.Implicits._
-import org.apache.kafka.clients.admin.NewTopic
+import kafka.utils.TestUtils._
+import kafka.utils.{Log4jController, Logging, TestUtils}
+import kafka.zk.KafkaZkClient
+import org.apache.kafka.clients.admin._
 import org.apache.kafka.clients.consumer.{ConsumerConfig, KafkaConsumer}
 import org.apache.kafka.clients.producer.KafkaProducer
 import org.apache.kafka.clients.producer.ProducerRecord
-import org.apache.kafka.common.{ConsumerGroupState, TopicPartition, TopicPartitionReplica}
+import org.apache.kafka.common.{ConsumerGroupState, ElectionType, TopicPartition, TopicPartitionInfo, TopicPartitionReplica}
 import org.apache.kafka.common.acl._
-import org.apache.kafka.common.config.ConfigResource
+import org.apache.kafka.common.config.{ConfigResource, LogLevelConfig}
 import org.apache.kafka.common.errors._
-import org.junit.{After, Before, Rule, Test}
 import org.apache.kafka.common.requests.{DeleteRecordsRequest, MetadataResponse}
-import org.apache.kafka.common.resource.{PatternType, ResourcePattern, ResourceType}
-import org.junit.rules.Timeout
+import org.apache.kafka.common.resource.{PatternType, Resource, ResourcePattern, ResourceType}
+import org.apache.kafka.common.utils.{Time, Utils}
 import org.junit.Assert._
+import org.junit.rules.Timeout
+import org.junit.{After, Before, Ignore, Rule, Test}
+import org.scalatest.Assertions.intercept
 
-import scala.util.Random
 import scala.collection.JavaConverters._
-import kafka.zk.KafkaZkClient
-
+import scala.collection.Seq
+import scala.compat.java8.OptionConverters._
 import scala.concurrent.duration.Duration
 import scala.concurrent.{Await, Future}
-import java.lang.{Long => JLong}
-
-import kafka.security.auth.Group
-
+import scala.util.Random
 /**
  * An integration test of the KafkaAdminClient.
  *
@@ -67,31 +66,37 @@ class AdminClientIntegrationTest extends IntegrationTestHarness with Logging {
   @Rule
   def globalTimeout = Timeout.millis(120000)
 
-  var client: AdminClient = null
+  var client: Admin = null
+  var brokerLoggerConfigResource: ConfigResource = null
+  var changedBrokerLoggers = scala.collection.mutable.Set[String]()
 
   val topic = "topic"
   val partition = 0
   val topicPartition = new TopicPartition(topic, partition)
+  val clusterResourcePattern = new ResourcePattern(ResourceType.CLUSTER, Resource.CLUSTER_NAME, PatternType.LITERAL)
+
 
   @Before
   override def setUp(): Unit = {
     super.setUp
     TestUtils.waitUntilBrokerMetadataIsPropagated(servers)
+    brokerLoggerConfigResource = new ConfigResource(ConfigResource.Type.BROKER_LOGGER, servers.head.config.brokerId.toString)
   }
 
   @After
   override def tearDown(): Unit = {
+    teardownBrokerLoggers()
     if (client != null)
       Utils.closeQuietly(client, "AdminClient")
     super.tearDown()
   }
 
-  val serverCount = 3
+  val brokerCount = 3
   val consumerCount = 1
   val producerCount = 1
 
   override def generateConfigs = {
-    val cfgs = TestUtils.createBrokerConfigs(serverCount, zkConnect, interBrokerSecurityProtocol = Some(securityProtocol),
+    val cfgs = TestUtils.createBrokerConfigs(brokerCount, zkConnect, interBrokerSecurityProtocol = Some(securityProtocol),
       trustStoreFile = trustStoreFile, saslProperties = serverSaslProperties, logDirCount = 2)
     cfgs.foreach { config =>
       config.setProperty(KafkaConfig.ListenersProp, s"${listenerName.value}://localhost:${TestUtils.RandomPort}")
@@ -102,6 +107,9 @@ class AdminClientIntegrationTest extends IntegrationTestHarness with Logging {
       config.setProperty(KafkaConfig.GroupInitialRebalanceDelayMsProp, "0")
       config.setProperty(KafkaConfig.AutoLeaderRebalanceEnableProp, "false")
       config.setProperty(KafkaConfig.ControlledShutdownEnableProp, "false")
+      // Set up CreateTopicPolicy to be included in test.
+      config.setProperty(KafkaConfig.DefaultReplicationFactorProp, "1");
+      config.setProperty(KafkaConfig.CreateTopicPolicyClassNameProp, "kafka.server.LiCreateTopicPolicy")
       // We set this in order to test that we don't expose sensitive data via describe configs. This will already be
       // set for subclasses with security enabled and we don't want to overwrite it.
       if (!config.containsKey(KafkaConfig.SslTruststorePasswordProp))
@@ -121,7 +129,7 @@ class AdminClientIntegrationTest extends IntegrationTestHarness with Logging {
     config
   }
 
-  def waitForTopics(client: AdminClient, expectedPresent: Seq[String], expectedMissing: Seq[String]): Unit = {
+  def waitForTopics(client: Admin, expectedPresent: Seq[String], expectedMissing: Seq[String]): Unit = {
     TestUtils.waitUntilTrue(() => {
         val topics = client.listTopics.names.get()
         expectedPresent.forall(topicName => topics.contains(topicName)) &&
@@ -151,22 +159,43 @@ class AdminClientIntegrationTest extends IntegrationTestHarness with Logging {
   @Test
   def testCreateDeleteTopics(): Unit = {
     client = AdminClient.create(createConfig())
-    val topics = Seq("mytopic", "mytopic2")
+    val topics = Seq("mytopic", "mytopic2", "mytopic3")
     val newTopics = Seq(
       new NewTopic("mytopic", Map((0: Integer) -> Seq[Integer](1, 2).asJava, (1: Integer) -> Seq[Integer](2, 0).asJava).asJava),
-      new NewTopic("mytopic2", 3, 3)
+      new NewTopic("mytopic2", 3, 3.toShort),
+      new NewTopic("mytopic3", Option.empty[Integer].asJava, Option.empty[java.lang.Short].asJava)
     )
-    client.createTopics(newTopics.asJava, new CreateTopicsOptions().validateOnly(true)).all.get()
+    val validateResult = client.createTopics(newTopics.asJava, new CreateTopicsOptions().validateOnly(true))
+    validateResult.all.get()
     waitForTopics(client, List(), topics)
 
-    client.createTopics(newTopics.asJava).all.get()
+    def validateMetadataAndConfigs(result: CreateTopicsResult): Unit = {
+      assertEquals(2, result.numPartitions("mytopic").get())
+      assertEquals(2, result.replicationFactor("mytopic").get())
+      assertEquals(3, result.numPartitions("mytopic2").get())
+      assertEquals(3, result.replicationFactor("mytopic2").get())
+      assertEquals(configs.head.numPartitions, result.numPartitions("mytopic3").get())
+      assertEquals(configs.head.defaultReplicationFactor, result.replicationFactor("mytopic3").get())
+      assertFalse(result.config("mytopic").get().entries.isEmpty)
+    }
+    validateMetadataAndConfigs(validateResult)
+
+    val createResult = client.createTopics(newTopics.asJava)
+    createResult.all.get()
     waitForTopics(client, topics, List())
+    validateMetadataAndConfigs(createResult)
 
-    val results = client.createTopics(newTopics.asJava).values()
+    val failedCreateResult = client.createTopics(newTopics.asJava)
+    val results = failedCreateResult.values()
     assertTrue(results.containsKey("mytopic"))
     assertFutureExceptionTypeEquals(results.get("mytopic"), classOf[TopicExistsException])
     assertTrue(results.containsKey("mytopic2"))
     assertFutureExceptionTypeEquals(results.get("mytopic2"), classOf[TopicExistsException])
+    assertTrue(results.containsKey("mytopic3"))
+    assertFutureExceptionTypeEquals(results.get("mytopic3"), classOf[TopicExistsException])
+    assertFutureExceptionTypeEquals(failedCreateResult.numPartitions("mytopic3"), classOf[TopicExistsException])
+    assertFutureExceptionTypeEquals(failedCreateResult.replicationFactor("mytopic3"), classOf[TopicExistsException])
+    assertFutureExceptionTypeEquals(failedCreateResult.config("mytopic3"), classOf[TopicExistsException])
 
     val topicToDescription = client.describeTopics(topics.asJava).all.get()
     assertEquals(topics.toSet, topicToDescription.keySet.asScala)
@@ -196,7 +225,7 @@ class AdminClientIntegrationTest extends IntegrationTestHarness with Logging {
       assertEquals(3, partition.replicas.size)
       partition.replicas.asScala.foreach { replica =>
         assertTrue(replica.id >= 0)
-        assertTrue(replica.id < serverCount)
+        assertTrue(replica.id < brokerCount)
       }
       assertEquals("No duplicate replica ids", partition.replicas.size, partition.replicas.asScala.map(_.id).distinct.size)
 
@@ -205,15 +234,37 @@ class AdminClientIntegrationTest extends IntegrationTestHarness with Logging {
       assertTrue(partition.replicas.contains(partition.leader))
     }
 
+    val topic3 = topicToDescription.get("mytopic3")
+    assertEquals("mytopic3", topic3.name)
+    assertEquals(configs.head.numPartitions, topic3.partitions.size)
+    assertEquals(configs.head.defaultReplicationFactor, topic3.partitions.get(0).replicas().size())
+
     client.deleteTopics(topics.asJava).all.get()
     waitForTopics(client, List(), topics)
   }
 
+  @Test
+  def testCreateExistingTopicsThrowTopicExistsException(): Unit = {
+    client = AdminClient.create(createConfig())
+    val topic = "mytopic"
+    val topics = Seq(topic)
+    val newTopics = Seq(new NewTopic(topic, 1, 1.toShort))
+
+    client.createTopics(newTopics.asJava).all.get()
+    waitForTopics(client, topics, List())
+
+    val newTopicsWithInvalidRF = Seq(new NewTopic(topic, 1, (servers.size + 1).toShort))
+    val e = intercept[ExecutionException] {
+      client.createTopics(newTopicsWithInvalidRF.asJava, new CreateTopicsOptions().validateOnly(true)).all.get()
+    }
+    assertTrue(e.getCause.isInstanceOf[TopicExistsException])
+  }
+
   @Test
   def testMetadataRefresh(): Unit = {
     client = AdminClient.create(createConfig())
     val topics = Seq("mytopic")
-    val newTopics = Seq(new NewTopic("mytopic", 3, 3))
+    val newTopics = Seq(new NewTopic("mytopic", 3, 3.toShort))
     client.createTopics(newTopics.asJava).all.get()
     waitForTopics(client, expectedPresent = topics, expectedMissing = List())
 
@@ -224,6 +275,39 @@ class AdminClientIntegrationTest extends IntegrationTestHarness with Logging {
     assertEquals(topics.toSet, topicDesc.keySet.asScala)
   }
 
+  @Test
+  def testAuthorizedOperations(): Unit = {
+    client = AdminClient.create(createConfig())
+
+    // without includeAuthorizedOperations flag
+    var result = client.describeCluster
+    assertEquals(Set().asJava, result.authorizedOperations().get())
+
+    //with includeAuthorizedOperations flag
+    result = client.describeCluster(new DescribeClusterOptions().includeAuthorizedOperations(true))
+    var expectedOperations = configuredClusterPermissions.asJava
+    assertEquals(expectedOperations, result.authorizedOperations().get())
+
+    val topic = "mytopic"
+    val newTopics = Seq(new NewTopic(topic, 3, 3.toShort))
+    client.createTopics(newTopics.asJava).all.get()
+    waitForTopics(client, expectedPresent = Seq(topic), expectedMissing = List())
+
+    // without includeAuthorizedOperations flag
+    var topicResult = getTopicMetadata(client, topic)
+    assertEquals(Set().asJava, topicResult.authorizedOperations)
+
+    //with includeAuthorizedOperations flag
+    topicResult = getTopicMetadata(client, topic, new DescribeTopicsOptions().includeAuthorizedOperations(true))
+    expectedOperations = Topic.supportedOperations
+      .map(operation => operation.toJava).asJava
+    assertEquals(expectedOperations, topicResult.authorizedOperations)
+  }
+
+  def configuredClusterPermissions() : Set[AclOperation] = {
+    Cluster.supportedOperations.map(operation => operation.toJava)
+  }
+
   /**
     * describe should not auto create topics
     */
@@ -232,7 +316,7 @@ class AdminClientIntegrationTest extends IntegrationTestHarness with Logging {
     client = AdminClient.create(createConfig())
 
     val existingTopic = "existing-topic"
-    client.createTopics(Seq(existingTopic).map(new NewTopic(_, 1, 1)).asJava).all.get()
+    client.createTopics(Seq(existingTopic).map(new NewTopic(_, 1, 1.toShort)).asJava).all.get()
     waitForTopics(client, Seq(existingTopic), List())
 
     val nonExistingTopic = "non-existing"
@@ -245,10 +329,11 @@ class AdminClientIntegrationTest extends IntegrationTestHarness with Logging {
   @Test
   def testDescribeCluster(): Unit = {
     client = AdminClient.create(createConfig())
-    val nodes = client.describeCluster.nodes.get()
-    val clusterId = client.describeCluster().clusterId().get()
+    val result = client.describeCluster
+    val nodes = result.nodes.get()
+    val clusterId = result.clusterId().get()
     assertEquals(servers.head.dataPlaneRequestProcessor.clusterId, clusterId)
-    val controller = client.describeCluster().controller().get()
+    val controller = result.controller().get()
     assertEquals(servers.head.dataPlaneRequestProcessor.metadataCache.getControllerId.
       getOrElse(MetadataResponse.NO_CONTROLLER_ID), controller.id())
     val brokers = brokerList.split(",")
@@ -265,10 +350,10 @@ class AdminClientIntegrationTest extends IntegrationTestHarness with Logging {
     val topic = "topic"
     val leaderByPartition = createTopic(topic, numPartitions = 10, replicationFactor = 1)
     val partitionsByBroker = leaderByPartition.groupBy { case (partitionId, leaderId) => leaderId }.mapValues(_.keys.toSeq)
-    val brokers = (0 until serverCount).map(Integer.valueOf)
+    val brokers = (0 until brokerCount).map(Integer.valueOf)
     val logDirInfosByBroker = client.describeLogDirs(brokers.asJava).all.get
 
-    (0 until serverCount).foreach { brokerId =>
+    (0 until brokerCount).foreach { brokerId =>
       val server = servers.find(_.config.brokerId == brokerId).get
       val expectedPartitions = partitionsByBroker(brokerId)
       val logDirInfos = logDirInfosByBroker.get(brokerId)
@@ -325,7 +410,7 @@ class AdminClientIntegrationTest extends IntegrationTestHarness with Logging {
       assertTrue(exception.getCause.isInstanceOf[UnknownTopicOrPartitionException])
     }
 
-    createTopic(topic, numPartitions = 1, replicationFactor = serverCount)
+    createTopic(topic, numPartitions = 1, replicationFactor = brokerCount)
     servers.foreach { server =>
       val logDir = server.logManager.getLog(tp).get.dir.getParent
       assertEquals(firstReplicaAssignment(new TopicPartitionReplica(topic, 0, server.config.brokerId)), logDir)
@@ -471,17 +556,19 @@ class AdminClientIntegrationTest extends IntegrationTestHarness with Logging {
     createTopic(topic2, numPartitions = 1, replicationFactor = 2)
 
     // assert that both the topics have 1 partition
-    assertEquals(1, client.describeTopics(Set(topic1).asJava).values.get(topic1).get.partitions.size)
-    assertEquals(1, client.describeTopics(Set(topic2).asJava).values.get(topic2).get.partitions.size)
+    val topic1_metadata = getTopicMetadata(client, topic1)
+    val topic2_metadata = getTopicMetadata(client, topic2)
+    assertEquals(1, topic1_metadata.partitions.size)
+    assertEquals(1, topic2_metadata.partitions.size)
 
     val validateOnly = new CreatePartitionsOptions().validateOnly(true)
     val actuallyDoIt = new CreatePartitionsOptions().validateOnly(false)
 
-    def partitions(topic: String) =
-      client.describeTopics(Set(topic).asJava).values.get(topic).get.partitions
+    def partitions(topic: String, expectedNumPartitionsOpt: Option[Int] = None): util.List[TopicPartitionInfo] = {
+      getTopicMetadata(client, topic, expectedNumPartitionsOpt = expectedNumPartitionsOpt).partitions
+    }
 
-    def numPartitions(topic: String) =
-      partitions(topic).size
+    def numPartitions(topic: String, expectedNumPartitionsOpt: Option[Int] = None): Int = partitions(topic, expectedNumPartitionsOpt).size
 
     // validateOnly: try creating a new partition (no assignments), to bring the total to 3 partitions
     var alterResult = client.createPartitions(Map(topic1 ->
@@ -493,7 +580,7 @@ class AdminClientIntegrationTest extends IntegrationTestHarness with Logging {
     alterResult = client.createPartitions(Map(topic1 ->
       NewPartitions.increaseTo(3)).asJava, actuallyDoIt)
     altered = alterResult.values.get(topic1).get
-    assertEquals(3, numPartitions(topic1))
+    TestUtils.waitUntilTrue(() => numPartitions(topic1) == 3, "Timed out waiting for new partitions to appear")
 
     // validateOnly: now try creating a new partition (with assignments), to bring the total to 3 partitions
     val newPartition2Assignments = asList[util.List[Integer]](asList(0, 1), asList(1, 2))
@@ -506,7 +593,7 @@ class AdminClientIntegrationTest extends IntegrationTestHarness with Logging {
     alterResult = client.createPartitions(Map(topic2 ->
       NewPartitions.increaseTo(3, newPartition2Assignments)).asJava, actuallyDoIt)
     altered = alterResult.values.get(topic2).get
-    val actualPartitions2 = partitions(topic2)
+    val actualPartitions2 = partitions(topic2, expectedNumPartitionsOpt = Some(3))
     assertEquals(3, actualPartitions2.size)
     assertEquals(Seq(0, 1), actualPartitions2.get(1).replicas.asScala.map(_.id).toList)
     assertEquals(Seq(1, 2), actualPartitions2.get(2).replicas.asScala.map(_.id).toList)
@@ -538,7 +625,7 @@ class AdminClientIntegrationTest extends IntegrationTestHarness with Logging {
         case e: ExecutionException =>
           assertTrue(desc, e.getCause.isInstanceOf[InvalidPartitionsException])
           assertEquals(desc, "Topic already has 3 partitions.", e.getCause.getMessage)
-          assertEquals(desc, 3, numPartitions(topic2))
+          assertEquals(desc, 3, numPartitions(topic2, Some(3)))
       }
 
       // try a newCount which would be a noop (where the assignment matches current state)
@@ -550,7 +637,7 @@ class AdminClientIntegrationTest extends IntegrationTestHarness with Logging {
         case e: ExecutionException =>
           assertTrue(desc, e.getCause.isInstanceOf[InvalidPartitionsException])
           assertEquals(desc, "Topic already has 3 partitions.", e.getCause.getMessage)
-          assertEquals(desc, 3, numPartitions(topic2))
+          assertEquals(desc, 3, numPartitions(topic2, Some(3)))
       }
 
       // try a newCount which would be a noop (where the assignment doesn't match current state)
@@ -562,7 +649,7 @@ class AdminClientIntegrationTest extends IntegrationTestHarness with Logging {
         case e: ExecutionException =>
           assertTrue(desc, e.getCause.isInstanceOf[InvalidPartitionsException])
           assertEquals(desc, "Topic already has 3 partitions.", e.getCause.getMessage)
-          assertEquals(desc, 3, numPartitions(topic2))
+          assertEquals(desc, 3, numPartitions(topic2, Some(3)))
       }
 
       // try a bad topic name
@@ -694,11 +781,10 @@ class AdminClientIntegrationTest extends IntegrationTestHarness with Logging {
       topic2 -> NewPartitions.increaseTo(2)).asJava, actuallyDoIt)
     // assert that the topic1 now has 4 partitions
     altered = alterResult.values.get(topic1).get
-    assertEquals(4, numPartitions(topic1))
+    TestUtils.waitUntilTrue(() => numPartitions(topic1) == 4, "Timed out waiting for new partitions to appear")
     try {
       altered = alterResult.values.get(topic2).get
     } catch {
-      case e: ExecutionException =>
       case e: ExecutionException =>
         assertTrue(e.getCause.isInstanceOf[InvalidPartitionsException])
         assertEquals("Topic currently has 3 partitions, which is higher than the requested 2.", e.getCause.getMessage)
@@ -723,7 +809,7 @@ class AdminClientIntegrationTest extends IntegrationTestHarness with Logging {
 
   @Test
   def testSeekAfterDeleteRecords(): Unit = {
-    createTopic(topic, numPartitions = 2, replicationFactor = serverCount)
+    createTopic(topic, numPartitions = 2, replicationFactor = brokerCount)
 
     client = AdminClient.create(createConfig)
 
@@ -752,7 +838,7 @@ class AdminClientIntegrationTest extends IntegrationTestHarness with Logging {
 
   @Test
   def testLogStartOffsetCheckpoint(): Unit = {
-    createTopic(topic, numPartitions = 2, replicationFactor = serverCount)
+    createTopic(topic, numPartitions = 2, replicationFactor = brokerCount)
 
     client = AdminClient.create(createConfig)
 
@@ -765,7 +851,7 @@ class AdminClientIntegrationTest extends IntegrationTestHarness with Logging {
     var lowWatermark: Option[Long] = Some(result.lowWatermarks.get(topicPartition).get.lowWatermark)
     assertEquals(Some(5), lowWatermark)
 
-    for (i <- 0 until serverCount) {
+    for (i <- 0 until brokerCount) {
       killBroker(i)
     }
     restartDeadBrokers()
@@ -792,7 +878,7 @@ class AdminClientIntegrationTest extends IntegrationTestHarness with Logging {
 
   @Test
   def testLogStartOffsetAfterDeleteRecords(): Unit = {
-    createTopic(topic, numPartitions = 2, replicationFactor = serverCount)
+    createTopic(topic, numPartitions = 2, replicationFactor = brokerCount)
 
     client = AdminClient.create(createConfig)
 
@@ -806,26 +892,26 @@ class AdminClientIntegrationTest extends IntegrationTestHarness with Logging {
     val lowWatermark = result.lowWatermarks.get(topicPartition).get.lowWatermark
     assertEquals(3L, lowWatermark)
 
-    for (i <- 0 until serverCount)
-      assertEquals(3, servers(i).replicaManager.localReplica(topicPartition).get.logStartOffset)
+    for (i <- 0 until brokerCount)
+      assertEquals(3, servers(i).replicaManager.localLog(topicPartition).get.logStartOffset)
   }
 
   @Test
   def testReplicaCanFetchFromLogStartOffsetAfterDeleteRecords(): Unit = {
-    val leaders = createTopic(topic, numPartitions = 1, replicationFactor = serverCount)
+    val leaders = createTopic(topic, numPartitions = 1, replicationFactor = brokerCount)
     val followerIndex = if (leaders(0) != servers(0).config.brokerId) 0 else 1
 
     def waitForFollowerLog(expectedStartOffset: Long, expectedEndOffset: Long): Unit = {
-      TestUtils.waitUntilTrue(() => servers(followerIndex).replicaManager.localReplica(topicPartition) != None,
+      TestUtils.waitUntilTrue(() => servers(followerIndex).replicaManager.localLog(topicPartition) != None,
                               "Expected follower to create replica for partition")
 
       // wait until the follower discovers that log start offset moved beyond its HW
       TestUtils.waitUntilTrue(() => {
-        servers(followerIndex).replicaManager.localReplica(topicPartition).get.logStartOffset == expectedStartOffset
+        servers(followerIndex).replicaManager.localLog(topicPartition).get.logStartOffset == expectedStartOffset
       }, s"Expected follower to discover new log start offset $expectedStartOffset")
 
       TestUtils.waitUntilTrue(() => {
-        servers(followerIndex).replicaManager.localReplica(topicPartition).get.logEndOffset == expectedEndOffset
+        servers(followerIndex).replicaManager.localLog(topicPartition).get.logEndOffset == expectedEndOffset
       }, s"Expected follower to catch up to log end offset $expectedEndOffset")
     }
 
@@ -845,8 +931,8 @@ class AdminClientIntegrationTest extends IntegrationTestHarness with Logging {
     waitForFollowerLog(expectedStartOffset=3L, expectedEndOffset=100L)
 
     // after the new replica caught up, all replicas should have same log start offset
-    for (i <- 0 until serverCount)
-      assertEquals(3, servers(i).replicaManager.localReplica(topicPartition).get.logStartOffset)
+    for (i <- 0 until brokerCount)
+      assertEquals(3, servers(i).replicaManager.localLog(topicPartition).get.logStartOffset)
 
     // kill the same follower again, produce more records, and delete records beyond follower's LOE
     killBroker(followerIndex)
@@ -860,7 +946,7 @@ class AdminClientIntegrationTest extends IntegrationTestHarness with Logging {
   @Test
   def testAlterLogDirsAfterDeleteRecords(): Unit = {
     client = AdminClient.create(createConfig)
-    createTopic(topic, numPartitions = 1, replicationFactor = serverCount)
+    createTopic(topic, numPartitions = 1, replicationFactor = brokerCount)
     val expectedLEO = 100
     val producer = createProducer()
     sendRecords(producer, expectedLEO, topicPartition)
@@ -869,9 +955,9 @@ class AdminClientIntegrationTest extends IntegrationTestHarness with Logging {
     val result = client.deleteRecords(Map(topicPartition -> RecordsToDelete.beforeOffset(3L)).asJava)
     result.all().get()
     // make sure we are in the expected state after delete records
-    for (i <- 0 until serverCount) {
-      assertEquals(3, servers(i).replicaManager.localReplica(topicPartition).get.logStartOffset)
-      assertEquals(expectedLEO, servers(i).replicaManager.localReplica(topicPartition).get.logEndOffset)
+    for (i <- 0 until brokerCount) {
+      assertEquals(3, servers(i).replicaManager.localLog(topicPartition).get.logStartOffset)
+      assertEquals(expectedLEO, servers(i).replicaManager.localLog(topicPartition).get.logEndOffset)
     }
 
     // we will create another dir just for one server
@@ -885,13 +971,13 @@ class AdminClientIntegrationTest extends IntegrationTestHarness with Logging {
     }, "timed out waiting for replica movement")
 
     // once replica moved, its LSO and LEO should match other replicas
-    assertEquals(3, servers(0).replicaManager.localReplica(topicPartition).get.logStartOffset)
-    assertEquals(expectedLEO, servers(0).replicaManager.localReplica(topicPartition).get.logEndOffset)
+    assertEquals(3, servers.head.replicaManager.localLog(topicPartition).get.logStartOffset)
+    assertEquals(expectedLEO, servers.head.replicaManager.localLog(topicPartition).get.logEndOffset)
   }
 
   @Test
   def testOffsetsForTimesAfterDeleteRecords(): Unit = {
-    createTopic(topic, numPartitions = 2, replicationFactor = serverCount)
+    createTopic(topic, numPartitions = 2, replicationFactor = brokerCount)
 
     client = AdminClient.create(createConfig)
 
@@ -963,7 +1049,7 @@ class AdminClientIntegrationTest extends IntegrationTestHarness with Logging {
 
   @Test
   def testDescribeConfigsForTopic(): Unit = {
-    createTopic(topic, numPartitions = 2, replicationFactor = serverCount)
+    createTopic(topic, numPartitions = 2, replicationFactor = brokerCount)
     client = AdminClient.create(createConfig)
 
     val existingTopic = new ConfigResource(ConfigResource.Type.TOPIC, topic)
@@ -1029,13 +1115,13 @@ class AdminClientIntegrationTest extends IntegrationTestHarness with Logging {
   def testDelayedClose(): Unit = {
     client = AdminClient.create(createConfig())
     val topics = Seq("mytopic", "mytopic2")
-    val newTopics = topics.map(new NewTopic(_, 1, 1))
+    val newTopics = topics.map(new NewTopic(_, 1, 1.toShort))
     val future = client.createTopics(newTopics.asJava, new CreateTopicsOptions().validateOnly(true)).all()
-    client.close(2, TimeUnit.HOURS)
+    client.close(time.Duration.ofHours(2))
     val future2 = client.createTopics(newTopics.asJava, new CreateTopicsOptions().validateOnly(true)).all()
     assertFutureExceptionTypeEquals(future2, classOf[TimeoutException])
     future.get
-    client.close(30, TimeUnit.MINUTES) // multiple close-with-timeout should have no effect
+    client.close(time.Duration.ofMinutes(30)) // multiple close-with-timeout should have no effect
   }
 
   /**
@@ -1049,9 +1135,9 @@ class AdminClientIntegrationTest extends IntegrationTestHarness with Logging {
     client = AdminClient.create(config)
     // Because the bootstrap servers are set up incorrectly, this call will not complete, but must be
     // cancelled by the close operation.
-    val future = client.createTopics(Seq("mytopic", "mytopic2").map(new NewTopic(_, 1, 1)).asJava,
+    val future = client.createTopics(Seq("mytopic", "mytopic2").map(new NewTopic(_, 1, 1.toShort)).asJava,
       new CreateTopicsOptions().timeoutMs(900000)).all()
-    client.close(0, TimeUnit.MILLISECONDS)
+    client.close(time.Duration.ZERO)
     assertFutureExceptionTypeEquals(future, classOf[TimeoutException])
   }
 
@@ -1066,7 +1152,7 @@ class AdminClientIntegrationTest extends IntegrationTestHarness with Logging {
     config.put(AdminClientConfig.REQUEST_TIMEOUT_MS_CONFIG, "0")
     client = AdminClient.create(config)
     val startTimeMs = Time.SYSTEM.milliseconds()
-    val future = client.createTopics(Seq("mytopic", "mytopic2").map(new NewTopic(_, 1, 1)).asJava,
+    val future = client.createTopics(Seq("mytopic", "mytopic2").map(new NewTopic(_, 1, 1.toShort)).asJava,
       new CreateTopicsOptions().timeoutMs(2)).all()
     assertFutureExceptionTypeEquals(future, classOf[TimeoutException])
     val endTimeMs = Time.SYSTEM.milliseconds()
@@ -1082,10 +1168,10 @@ class AdminClientIntegrationTest extends IntegrationTestHarness with Logging {
     config.put(AdminClientConfig.REQUEST_TIMEOUT_MS_CONFIG, "100000000")
     val factory = new KafkaAdminClientTest.FailureInjectingTimeoutProcessorFactory()
     client = KafkaAdminClientTest.createInternal(new AdminClientConfig(config), factory)
-    val future = client.createTopics(Seq("mytopic", "mytopic2").map(new NewTopic(_, 1, 1)).asJava,
+    val future = client.createTopics(Seq("mytopic", "mytopic2").map(new NewTopic(_, 1, 1.toShort)).asJava,
         new CreateTopicsOptions().validateOnly(true)).all()
     assertFutureExceptionTypeEquals(future, classOf[TimeoutException])
-    val future2 = client.createTopics(Seq("mytopic3", "mytopic4").map(new NewTopic(_, 1, 1)).asJava,
+    val future2 = client.createTopics(Seq("mytopic3", "mytopic4").map(new NewTopic(_, 1, 1.toShort)).asJava,
       new CreateTopicsOptions().validateOnly(true)).all()
     future2.get
     assertEquals(1, factory.failuresInjected)
@@ -1107,7 +1193,7 @@ class AdminClientIntegrationTest extends IntegrationTestHarness with Logging {
       val testTopicName = "test_topic"
       val testNumPartitions = 2
       client.createTopics(Collections.singleton(
-        new NewTopic(testTopicName, testNumPartitions, 1))).all().get()
+        new NewTopic(testTopicName, testNumPartitions, 1.toShort))).all().get()
       waitForTopics(client, List(testTopicName), List())
 
       val producer = createProducer()
@@ -1118,40 +1204,51 @@ class AdminClientIntegrationTest extends IntegrationTestHarness with Logging {
       }
       val testGroupId = "test_group_id"
       val testClientId = "test_client_id"
+      val testInstanceId = "test_instance_id"
       val fakeGroupId = "fake_group_id"
       val newConsumerConfig = new Properties(consumerConfig)
       newConsumerConfig.setProperty(ConsumerConfig.GROUP_ID_CONFIG, testGroupId)
       newConsumerConfig.setProperty(ConsumerConfig.CLIENT_ID_CONFIG, testClientId)
+      newConsumerConfig.setProperty(ConsumerConfig.GROUP_INSTANCE_ID_CONFIG, testInstanceId)
       val consumer = createConsumer(configOverrides = newConsumerConfig)
+      val latch = new CountDownLatch(1)
       try {
         // Start a consumer in a thread that will subscribe to a new group.
         val consumerThread = new Thread {
-          override def run {
+          override def run : Unit = {
             consumer.subscribe(Collections.singleton(testTopicName))
-            while (true) {
-              consumer.poll(5000)
-              consumer.commitSync()
+
+            try {
+              while (true) {
+                consumer.poll(JDuration.ofSeconds(5))
+                if (!consumer.assignment.isEmpty && latch.getCount > 0L)
+                  latch.countDown()
+                consumer.commitSync()
+              }
+            } catch {
+              case _: InterruptException => // Suppress the output to stderr
             }
           }
         }
         try {
           consumerThread.start
+          assertTrue(latch.await(30000, TimeUnit.MILLISECONDS))
           // Test that we can list the new group.
           TestUtils.waitUntilTrue(() => {
             val matching = client.listConsumerGroups.all.get().asScala.filter(_.groupId == testGroupId)
-            !matching.isEmpty
+            matching.nonEmpty
           }, s"Expected to be able to list $testGroupId")
 
-          val result = client.describeConsumerGroups(Seq(testGroupId, fakeGroupId).asJava,
+          val describeWithFakeGroupResult = client.describeConsumerGroups(Seq(testGroupId, fakeGroupId).asJava,
             new DescribeConsumerGroupsOptions().includeAuthorizedOperations(true))
-          assertEquals(2, result.describedGroups().size())
+          assertEquals(2, describeWithFakeGroupResult.describedGroups().size())
 
           // Test that we can get information about the test consumer group.
-          assertTrue(result.describedGroups().containsKey(testGroupId))
-          val testGroupDescription = result.describedGroups().get(testGroupId).get()
+          assertTrue(describeWithFakeGroupResult.describedGroups().containsKey(testGroupId))
+          var testGroupDescription = describeWithFakeGroupResult.describedGroups().get(testGroupId).get()
 
           assertEquals(testGroupId, testGroupDescription.groupId())
-          assertFalse(testGroupDescription.isSimpleConsumerGroup())
+          assertFalse(testGroupDescription.isSimpleConsumerGroup)
           assertEquals(1, testGroupDescription.members().size())
           val member = testGroupDescription.members().iterator().next()
           assertEquals(testClientId, member.clientId())
@@ -1164,8 +1261,8 @@ class AdminClientIntegrationTest extends IntegrationTestHarness with Logging {
           assertEquals(expectedOperations, testGroupDescription.authorizedOperations())
 
           // Test that the fake group is listed as dead.
-          assertTrue(result.describedGroups().containsKey(fakeGroupId))
-          val fakeGroupDescription = result.describedGroups().get(fakeGroupId).get()
+          assertTrue(describeWithFakeGroupResult.describedGroups().containsKey(fakeGroupId))
+          val fakeGroupDescription = describeWithFakeGroupResult.describedGroups().get(fakeGroupId).get()
 
           assertEquals(fakeGroupId, fakeGroupDescription.groupId())
           assertEquals(0, fakeGroupDescription.members().size())
@@ -1174,7 +1271,7 @@ class AdminClientIntegrationTest extends IntegrationTestHarness with Logging {
           assertEquals(expectedOperations, fakeGroupDescription.authorizedOperations())
 
           // Test that all() returns 2 results
-          assertEquals(2, result.all().get().size())
+          assertEquals(2, describeWithFakeGroupResult.all().get().size())
 
           // Test listConsumerGroupOffsets
           TestUtils.waitUntilTrue(() => {
@@ -1182,9 +1279,19 @@ class AdminClientIntegrationTest extends IntegrationTestHarness with Logging {
             val part = new TopicPartition(testTopicName, 0)
             parts.containsKey(part) && (parts.get(part).offset() == 1)
           }, s"Expected the offset for partition 0 to eventually become 1.")
+          
+          // Test delete non-exist consumer instance
+          val invalidInstanceId = "invalid-instance-id"
+          var removeMembersResult = client.removeMembersFromConsumerGroup(testGroupId, new RemoveMembersFromConsumerGroupOptions(
+            Collections.singleton(new MemberToRemove(invalidInstanceId))
+          ))
+
+          TestUtils.assertFutureExceptionTypeEquals(removeMembersResult.all, classOf[UnknownMemberIdException])
+          val firstMemberFuture = removeMembersResult.memberResult(new MemberToRemove(invalidInstanceId))
+          TestUtils.assertFutureExceptionTypeEquals(firstMemberFuture, classOf[UnknownMemberIdException])
 
           // Test consumer group deletion
-          val deleteResult = client.deleteConsumerGroups(Seq(testGroupId, fakeGroupId).asJava)
+          var deleteResult = client.deleteConsumerGroups(Seq(testGroupId, fakeGroupId).asJava)
           assertEquals(2, deleteResult.deletedGroups().size())
 
           // Deleting the fake group ID should get GroupIdNotFoundException.
@@ -1196,6 +1303,33 @@ class AdminClientIntegrationTest extends IntegrationTestHarness with Logging {
           assertTrue(deleteResult.deletedGroups().containsKey(testGroupId))
           assertFutureExceptionTypeEquals(deleteResult.deletedGroups().get(testGroupId),
             classOf[GroupNotEmptyException])
+
+          // Test delete correct member
+          removeMembersResult = client.removeMembersFromConsumerGroup(testGroupId, new RemoveMembersFromConsumerGroupOptions(
+            Collections.singleton(new MemberToRemove(testInstanceId))
+          ))
+
+          assertNull(removeMembersResult.all().get())
+          val validMemberFuture = removeMembersResult.memberResult(new MemberToRemove(testInstanceId))
+          assertNull(validMemberFuture.get())
+
+          // The group should contain no member now.
+          val describeTestGroupResult = client.describeConsumerGroups(Seq(testGroupId).asJava,
+            new DescribeConsumerGroupsOptions().includeAuthorizedOperations(true))
+          assertEquals(1, describeTestGroupResult.describedGroups().size())
+
+          testGroupDescription = describeTestGroupResult.describedGroups().get(testGroupId).get()
+
+          assertEquals(testGroupId, testGroupDescription.groupId)
+          assertFalse(testGroupDescription.isSimpleConsumerGroup)
+          assertTrue(testGroupDescription.members().isEmpty)
+
+          // Consumer group deletion on empty group should succeed
+          deleteResult = client.deleteConsumerGroups(Seq(testGroupId).asJava)
+          assertEquals(1, deleteResult.deletedGroups().size())
+
+          assertTrue(deleteResult.deletedGroups().containsKey(testGroupId))
+          assertNull(deleteResult.deletedGroups().get(testGroupId).get())
         } finally {
           consumerThread.interrupt()
           consumerThread.join()
@@ -1208,6 +1342,78 @@ class AdminClientIntegrationTest extends IntegrationTestHarness with Logging {
     }
   }
 
+  @Test
+  def testDeleteConsumerGroupOffsets(): Unit = {
+    val config = createConfig()
+    client = AdminClient.create(config)
+    try {
+      val testTopicName = "test_topic"
+      val testGroupId = "test_group_id"
+      val testClientId = "test_client_id"
+      val fakeGroupId = "fake_group_id"
+
+      val tp1 = new TopicPartition(testTopicName, 0)
+      val tp2 = new TopicPartition("foo", 0)
+
+      client.createTopics(Collections.singleton(
+        new NewTopic(testTopicName, 1, 1.toShort))).all().get()
+      waitForTopics(client, List(testTopicName), List())
+
+      val producer = createProducer()
+      try {
+        producer.send(new ProducerRecord(testTopicName, 0, null, null)).get()
+      } finally {
+        Utils.closeQuietly(producer, "producer")
+      }
+
+      val newConsumerConfig = new Properties(consumerConfig)
+      newConsumerConfig.setProperty(ConsumerConfig.GROUP_ID_CONFIG, testGroupId)
+      newConsumerConfig.setProperty(ConsumerConfig.CLIENT_ID_CONFIG, testClientId)
+      // Increase timeouts to avoid having a rebalance during the test
+      newConsumerConfig.setProperty(ConsumerConfig.MAX_POLL_INTERVAL_MS_CONFIG, Integer.MAX_VALUE.toString)
+      newConsumerConfig.setProperty(ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG, Defaults.GroupMaxSessionTimeoutMs.toString)
+      val consumer = createConsumer(configOverrides = newConsumerConfig)
+
+      try {
+        TestUtils.subscribeAndWaitForRecords(testTopicName, consumer)
+        consumer.commitSync()
+
+        // Test offset deletion while consuming
+        val offsetDeleteResult = client.deleteConsumerGroupOffsets(testGroupId, Set(tp1, tp2).asJava)
+
+        // Top level error will equal to the first partition level error
+        assertFutureExceptionTypeEquals(offsetDeleteResult.all(), classOf[GroupSubscribedToTopicException])
+        assertFutureExceptionTypeEquals(offsetDeleteResult.partitionResult(tp1),
+          classOf[GroupSubscribedToTopicException])
+        assertFutureExceptionTypeEquals(offsetDeleteResult.partitionResult(tp2),
+          classOf[UnknownTopicOrPartitionException])
+
+        // Test the fake group ID
+        val fakeDeleteResult = client.deleteConsumerGroupOffsets(fakeGroupId, Set(tp1, tp2).asJava)
+
+        assertFutureExceptionTypeEquals(fakeDeleteResult.all(), classOf[GroupIdNotFoundException])
+        assertFutureExceptionTypeEquals(fakeDeleteResult.partitionResult(tp1),
+          classOf[GroupIdNotFoundException])
+        assertFutureExceptionTypeEquals(fakeDeleteResult.partitionResult(tp2),
+          classOf[GroupIdNotFoundException])
+
+      } finally {
+        Utils.closeQuietly(consumer, "consumer")
+      }
+
+      // Test offset deletion when group is empty
+      val offsetDeleteResult = client.deleteConsumerGroupOffsets(testGroupId, Set(tp1, tp2).asJava)
+
+      assertFutureExceptionTypeEquals(offsetDeleteResult.all(),
+        classOf[UnknownTopicOrPartitionException])
+      assertNull(offsetDeleteResult.partitionResult(tp1).get())
+      assertFutureExceptionTypeEquals(offsetDeleteResult.partitionResult(tp2),
+        classOf[UnknownTopicOrPartitionException])
+    } finally {
+      Utils.closeQuietly(client, "adminClient")
+    }
+  }
+
   @Test
   def testElectPreferredLeaders(): Unit = {
     client = AdminClient.create(createConfig)
@@ -1222,22 +1428,17 @@ class AdminClientIntegrationTest extends IntegrationTestHarness with Logging {
     val partition2 = new TopicPartition("elect-preferred-leaders-topic-2", 0)
     TestUtils.createTopic(zkClient, partition2.topic, Map[Int, Seq[Int]](partition2.partition -> prefer0), servers)
 
-    def currentLeader(topicPartition: TopicPartition) =
-      client.describeTopics(asList(topicPartition.topic)).values.get(topicPartition.topic).
-        get.partitions.get(topicPartition.partition).leader.id
-
-    def preferredLeader(topicPartition: TopicPartition) =
-      client.describeTopics(asList(topicPartition.topic)).values.get(topicPartition.topic).
-        get.partitions.get(topicPartition.partition).replicas.get(0).id
-
-    def waitForLeaderToBecome(topicPartition: TopicPartition, leader: Int) =
-      TestUtils.waitUntilTrue(() => currentLeader(topicPartition) == leader, s"Expected leader to become $leader", 10000)
+    def preferredLeader(topicPartition: TopicPartition): Int = {
+      val partitionMetadata = getTopicMetadata(client, topicPartition.topic).partitions.get(topicPartition.partition)
+      val preferredLeaderMetadata = partitionMetadata.replicas.get(0)
+      preferredLeaderMetadata.id
+    }
 
     /** Changes the preferred leader without changing the current leader. */
     def changePreferredLeader(newAssignment: Seq[Int]) = {
       val preferred = newAssignment.head
-      val prior1 = currentLeader(partition1)
-      val prior2 = currentLeader(partition2)
+      val prior1 = zkClient.getLeaderForPartition(partition1).get
+      val prior2 = zkClient.getLeaderForPartition(partition2).get
 
       var m = Map.empty[TopicPartition, Seq[Int]]
 
@@ -1249,125 +1450,830 @@ class AdminClientIntegrationTest extends IntegrationTestHarness with Logging {
       zkClient.createPartitionReassignment(m)
       TestUtils.waitUntilTrue(
         () => preferredLeader(partition1) == preferred && preferredLeader(partition2) == preferred,
-        s"Expected preferred leader to become $preferred, but is ${preferredLeader(partition1)} and ${preferredLeader(partition2)}", 10000)
+        s"Expected preferred leader to become $preferred, but is ${preferredLeader(partition1)} and ${preferredLeader(partition2)}",
+        10000)
       // Check the leader hasn't moved
-      assertEquals(prior1, currentLeader(partition1))
-      assertEquals(prior2, currentLeader(partition2))
+      TestUtils.assertLeader(client, partition1, prior1)
+      TestUtils.assertLeader(client, partition2, prior2)
     }
 
     // Check current leaders are 0
-    assertEquals(0, currentLeader(partition1))
-    assertEquals(0, currentLeader(partition2))
+    TestUtils.assertLeader(client, partition1, 0)
+    TestUtils.assertLeader(client, partition2, 0)
 
     // Noop election
-    var electResult = client.electPreferredLeaders(asList(partition1))
-    electResult.partitionResult(partition1).get()
-    assertEquals(0, currentLeader(partition1))
+    var electResult = client.electLeaders(ElectionType.PREFERRED, Set(partition1).asJava)
+    var exception = electResult.partitions.get.get(partition1).get
+    assertEquals(classOf[ElectionNotNeededException], exception.getClass)
+    assertEquals("Leader election not needed for topic partition", exception.getMessage)
+    TestUtils.assertLeader(client, partition1, 0)
 
     // Noop election with null partitions
-    electResult = client.electPreferredLeaders(null)
-    electResult.partitionResult(partition1).get()
-    assertEquals(0, currentLeader(partition1))
-    electResult.partitionResult(partition2).get()
-    assertEquals(0, currentLeader(partition2))
+    electResult = client.electLeaders(ElectionType.PREFERRED, null)
+    assertTrue(electResult.partitions.get.isEmpty)
+    TestUtils.assertLeader(client, partition1, 0)
+    TestUtils.assertLeader(client, partition2, 0)
 
     // Now change the preferred leader to 1
     changePreferredLeader(prefer1)
 
     // meaningful election
-    electResult = client.electPreferredLeaders(asList(partition1))
-    assertEquals(Set(partition1).asJava, electResult.partitions.get)
-    electResult.partitionResult(partition1).get()
-    waitForLeaderToBecome(partition1, 1)
+    electResult = client.electLeaders(ElectionType.PREFERRED, Set(partition1).asJava)
+    assertEquals(Set(partition1).asJava, electResult.partitions.get.keySet)
+    assertFalse(electResult.partitions.get.get(partition1).isPresent)
+    TestUtils.assertLeader(client, partition1, 1)
 
     // topic 2 unchanged
-    var e = intercept[ExecutionException](electResult.partitionResult(partition2).get()).getCause
-    assertEquals(classOf[UnknownTopicOrPartitionException], e.getClass)
-    assertEquals("Preferred leader election for partition \"elect-preferred-leaders-topic-2-0\" was not attempted",
-      e.getMessage)
-    assertEquals(0, currentLeader(partition2))
+    assertFalse(electResult.partitions.get.containsKey(partition2))
+    TestUtils.assertLeader(client, partition2, 0)
 
     // meaningful election with null partitions
-    electResult = client.electPreferredLeaders(null)
-    assertEquals(Set(partition1, partition2), electResult.partitions.get.asScala.filterNot(_.topic == "__consumer_offsets"))
-    electResult.partitionResult(partition1).get()
-    waitForLeaderToBecome(partition1, 1)
-    electResult.partitionResult(partition2).get()
-    waitForLeaderToBecome(partition2, 1)
+    electResult = client.electLeaders(ElectionType.PREFERRED, null)
+    assertEquals(Set(partition2), electResult.partitions.get.keySet.asScala)
+    assertFalse(electResult.partitions.get.get(partition2).isPresent)
+    TestUtils.assertLeader(client, partition2, 1)
 
     // unknown topic
     val unknownPartition = new TopicPartition("topic-does-not-exist", 0)
-    electResult = client.electPreferredLeaders(asList(unknownPartition))
-    assertEquals(Set(unknownPartition).asJava, electResult.partitions.get)
-    e = intercept[ExecutionException](electResult.partitionResult(unknownPartition).get()).getCause
-    assertEquals(classOf[UnknownTopicOrPartitionException], e.getClass)
-    assertEquals("The partition does not exist.", e.getMessage)
-    assertEquals(1, currentLeader(partition1))
-    assertEquals(1, currentLeader(partition2))
+    electResult = client.electLeaders(ElectionType.PREFERRED, Set(unknownPartition).asJava)
+    assertEquals(Set(unknownPartition).asJava, electResult.partitions.get.keySet)
+    exception = electResult.partitions.get.get(unknownPartition).get
+    assertEquals(classOf[UnknownTopicOrPartitionException], exception.getClass)
+    assertEquals("The partition does not exist.", exception.getMessage)
+    TestUtils.assertLeader(client, partition1, 1)
+    TestUtils.assertLeader(client, partition2, 1)
 
     // Now change the preferred leader to 2
     changePreferredLeader(prefer2)
 
     // mixed results
-    electResult = client.electPreferredLeaders(asList(unknownPartition, partition1))
-    assertEquals(Set(unknownPartition, partition1).asJava, electResult.partitions.get)
-    waitForLeaderToBecome(partition1, 2)
-    assertEquals(1, currentLeader(partition2))
-    e = intercept[ExecutionException](electResult.partitionResult(unknownPartition).get()).getCause
-    assertEquals(classOf[UnknownTopicOrPartitionException], e.getClass)
-    assertEquals("The partition does not exist.", e.getMessage)
-
-    // dupe partitions
-    electResult = client.electPreferredLeaders(asList(partition2, partition2))
-    assertEquals(Set(partition2).asJava, electResult.partitions.get)
-    electResult.partitionResult(partition2).get()
-    waitForLeaderToBecome(partition2, 2)
+    electResult = client.electLeaders(ElectionType.PREFERRED, Set(unknownPartition, partition1).asJava)
+    assertEquals(Set(unknownPartition, partition1).asJava, electResult.partitions.get.keySet)
+    TestUtils.assertLeader(client, partition1, 2)
+    TestUtils.assertLeader(client, partition2, 1)
+    exception = electResult.partitions.get.get(unknownPartition).get
+    assertEquals(classOf[UnknownTopicOrPartitionException], exception.getClass)
+    assertEquals("The partition does not exist.", exception.getMessage)
+
+    // elect preferred leader for partition 2
+    electResult = client.electLeaders(ElectionType.PREFERRED, Set(partition2).asJava)
+    assertEquals(Set(partition2).asJava, electResult.partitions.get.keySet)
+    assertFalse(electResult.partitions.get.get(partition2).isPresent)
+    TestUtils.assertLeader(client, partition2, 2)
 
     // Now change the preferred leader to 1
     changePreferredLeader(prefer1)
     // but shut it down...
     servers(1).shutdown()
-    waitUntilTrue (() => {
-      val description = client.describeTopics(Set(partition1.topic, partition2.topic).asJava).all().get()
-      val isr = description.asScala.values.flatMap(_.partitions.asScala.flatMap(_.isr.asScala))
-      !isr.exists(_.id == 1)
-    }, "Expect broker 1 to no longer be in any ISR")
+    TestUtils.waitForBrokersOutOfIsr(client, Set(partition1, partition2), Set(1))
 
     // ... now what happens if we try to elect the preferred leader and it's down?
-    val shortTimeout = new ElectPreferredLeadersOptions().timeoutMs(10000)
-    electResult = client.electPreferredLeaders(asList(partition1), shortTimeout)
-    assertEquals(Set(partition1).asJava, electResult.partitions.get)
-    e = intercept[ExecutionException](electResult.partitionResult(partition1).get()).getCause
-    assertEquals(classOf[PreferredLeaderNotAvailableException], e.getClass)
-    assertTrue(s"Wrong message ${e.getMessage}", e.getMessage.contains(
+    val shortTimeout = new ElectLeadersOptions().timeoutMs(10000)
+    electResult = client.electLeaders(ElectionType.PREFERRED, Set(partition1).asJava, shortTimeout)
+    assertEquals(Set(partition1).asJava, electResult.partitions.get.keySet)
+    exception = electResult.partitions.get.get(partition1).get
+    assertEquals(classOf[PreferredLeaderNotAvailableException], exception.getClass)
+    assertTrue(s"Wrong message ${exception.getMessage}", exception.getMessage.contains(
       "Failed to elect leader for partition elect-preferred-leaders-topic-1-0 under strategy PreferredReplicaPartitionLeaderElectionStrategy"))
-    assertEquals(2, currentLeader(partition1))
+    TestUtils.assertLeader(client, partition1, 2)
 
     // preferred leader unavailable with null argument
-    electResult = client.electPreferredLeaders(null, shortTimeout)
-    e = intercept[ExecutionException](electResult.partitions.get()).getCause
-    assertEquals(classOf[PreferredLeaderNotAvailableException], e.getClass)
+    electResult = client.electLeaders(ElectionType.PREFERRED, null, shortTimeout)
 
-    e = intercept[ExecutionException](electResult.partitionResult(partition1).get()).getCause
-    assertEquals(classOf[PreferredLeaderNotAvailableException], e.getClass)
-    assertTrue(s"Wrong message ${e.getMessage}", e.getMessage.contains(
+    exception = electResult.partitions.get.get(partition1).get
+    assertEquals(classOf[PreferredLeaderNotAvailableException], exception.getClass)
+    assertTrue(s"Wrong message ${exception.getMessage}", exception.getMessage.contains(
       "Failed to elect leader for partition elect-preferred-leaders-topic-1-0 under strategy PreferredReplicaPartitionLeaderElectionStrategy"))
 
-    e = intercept[ExecutionException](electResult.partitionResult(partition2).get()).getCause
-    assertEquals(classOf[PreferredLeaderNotAvailableException], e.getClass)
-    assertTrue(s"Wrong message ${e.getMessage}", e.getMessage.contains(
+    exception = electResult.partitions.get.get(partition2).get
+    assertEquals(classOf[PreferredLeaderNotAvailableException], exception.getClass)
+    assertTrue(s"Wrong message ${exception.getMessage}", exception.getMessage.contains(
       "Failed to elect leader for partition elect-preferred-leaders-topic-2-0 under strategy PreferredReplicaPartitionLeaderElectionStrategy"))
 
-    assertEquals(2, currentLeader(partition1))
-    assertEquals(2, currentLeader(partition2))
+    TestUtils.assertLeader(client, partition1, 2)
+    TestUtils.assertLeader(client, partition2, 2)
+  }
+
+  @Test
+  def testElectUncleanLeadersForOnePartition(): Unit = {
+    // Case: unclean leader election with one topic partition
+    client = AdminClient.create(createConfig)
+
+    val broker1 = 1
+    val broker2 = 2
+    val assignment1 = Seq(broker1, broker2)
+
+    val partition1 = new TopicPartition("unclean-test-topic-1", 0)
+    TestUtils.createTopic(zkClient, partition1.topic, Map[Int, Seq[Int]](partition1.partition -> assignment1), servers)
+
+    TestUtils.assertLeader(client, partition1, broker1)
+
+    servers(broker2).shutdown()
+    TestUtils.waitForBrokersOutOfIsr(client, Set(partition1), Set(broker2))
+    servers(broker1).shutdown()
+    TestUtils.assertNoLeader(client, partition1)
+    servers(broker2).startup()
+
+    val electResult = client.electLeaders(ElectionType.UNCLEAN, Set(partition1).asJava)
+    assertFalse(electResult.partitions.get.get(partition1).isPresent)
+    TestUtils.assertLeader(client, partition1, broker2)
+  }
+
+  @Test
+  def testElectUncleanLeadersForManyPartitions(): Unit = {
+    // Case: unclean leader election with many topic partitions
+    client = AdminClient.create(createConfig)
+
+    val broker1 = 1
+    val broker2 = 2
+    val assignment1 = Seq(broker1, broker2)
+    val assignment2 = Seq(broker1, broker2)
+
+    val topic = "unclean-test-topic-1"
+    val partition1 = new TopicPartition(topic, 0)
+    val partition2 = new TopicPartition(topic, 1)
+
+    TestUtils.createTopic(
+      zkClient,
+      topic,
+      Map(partition1.partition -> assignment1, partition2.partition -> assignment2),
+      servers
+    )
+
+    TestUtils.assertLeader(client, partition1, broker1)
+    TestUtils.assertLeader(client, partition2, broker1)
+
+    servers(broker2).shutdown()
+    TestUtils.waitForBrokersOutOfIsr(client, Set(partition1, partition2), Set(broker2))
+    servers(broker1).shutdown()
+    TestUtils.assertNoLeader(client, partition1)
+    TestUtils.assertNoLeader(client, partition2)
+    servers(broker2).startup()
+
+    val electResult = client.electLeaders(ElectionType.UNCLEAN, Set(partition1, partition2).asJava)
+    assertFalse(electResult.partitions.get.get(partition1).isPresent)
+    assertFalse(electResult.partitions.get.get(partition2).isPresent)
+    TestUtils.assertLeader(client, partition1, broker2)
+    TestUtils.assertLeader(client, partition2, broker2)
+  }
+
+  @Test
+  def testElectUncleanLeadersForAllPartitions(): Unit = {
+    // Case: noop unclean leader election and valid unclean leader election for all partitions
+    client = AdminClient.create(createConfig)
+
+    val broker1 = 1
+    val broker2 = 2
+    val broker3 = 0
+    val assignment1 = Seq(broker1, broker2)
+    val assignment2 = Seq(broker1, broker3)
+
+    val topic = "unclean-test-topic-1"
+    val partition1 = new TopicPartition(topic, 0)
+    val partition2 = new TopicPartition(topic, 1)
+
+    TestUtils.createTopic(
+      zkClient,
+      topic,
+      Map(partition1.partition -> assignment1, partition2.partition -> assignment2),
+      servers
+    )
+
+    TestUtils.assertLeader(client, partition1, broker1)
+    TestUtils.assertLeader(client, partition2, broker1)
+
+    servers(broker2).shutdown()
+    TestUtils.waitForBrokersOutOfIsr(client, Set(partition1), Set(broker2))
+    servers(broker1).shutdown()
+    TestUtils.assertNoLeader(client, partition1)
+    TestUtils.assertLeader(client, partition2, broker3)
+    servers(broker2).startup()
+
+    val electResult = client.electLeaders(ElectionType.UNCLEAN, null)
+    assertFalse(electResult.partitions.get.get(partition1).isPresent)
+    assertFalse(electResult.partitions.get.containsKey(partition2))
+    TestUtils.assertLeader(client, partition1, broker2)
+    TestUtils.assertLeader(client, partition2, broker3)
+  }
+
+  @Test
+  def testElectUncleanLeadersForUnknownPartitions(): Unit = {
+    // Case: unclean leader election for unknown topic
+    client = AdminClient.create(createConfig)
+
+    val broker1 = 1
+    val broker2 = 2
+    val assignment1 = Seq(broker1, broker2)
+
+    val topic = "unclean-test-topic-1"
+    val unknownPartition = new TopicPartition(topic, 1)
+    val unknownTopic = new TopicPartition("unknown-topic", 0)
+
+    TestUtils.createTopic(
+      zkClient,
+      topic,
+      Map(0 -> assignment1),
+      servers
+    )
+
+    TestUtils.assertLeader(client, new TopicPartition(topic, 0), broker1)
+
+    val electResult = client.electLeaders(ElectionType.UNCLEAN, Set(unknownPartition, unknownTopic).asJava)
+    assertTrue(electResult.partitions.get.get(unknownPartition).get.isInstanceOf[UnknownTopicOrPartitionException])
+    assertTrue(electResult.partitions.get.get(unknownTopic).get.isInstanceOf[UnknownTopicOrPartitionException])
+  }
+
+  @Test
+  def testElectUncleanLeadersWhenNoLiveBrokers(): Unit = {
+    // Case: unclean leader election with no live brokers
+    client = AdminClient.create(createConfig)
+
+    val broker1 = 1
+    val broker2 = 2
+    val assignment1 = Seq(broker1, broker2)
+
+    val topic = "unclean-test-topic-1"
+    val partition1 = new TopicPartition(topic, 0)
+
+    TestUtils.createTopic(
+      zkClient,
+      topic,
+      Map(partition1.partition -> assignment1),
+      servers
+    )
+
+    TestUtils.assertLeader(client, partition1, broker1)
+
+    servers(broker2).shutdown()
+    TestUtils.waitForBrokersOutOfIsr(client, Set(partition1), Set(broker2))
+    servers(broker1).shutdown()
+    TestUtils.assertNoLeader(client, partition1)
+
+    val electResult = client.electLeaders(ElectionType.UNCLEAN, Set(partition1).asJava)
+    assertTrue(electResult.partitions.get.get(partition1).get.isInstanceOf[EligibleLeadersNotAvailableException])
+  }
+
+  @Test
+  def testElectUncleanLeadersNoop(): Unit = {
+    // Case: noop unclean leader election with explicit topic partitions
+    client = AdminClient.create(createConfig)
+
+    val broker1 = 1
+    val broker2 = 2
+    val assignment1 = Seq(broker1, broker2)
+
+    val topic = "unclean-test-topic-1"
+    val partition1 = new TopicPartition(topic, 0)
+
+    TestUtils.createTopic(
+      zkClient,
+      topic,
+      Map(partition1.partition -> assignment1),
+      servers
+    )
+
+    TestUtils.assertLeader(client, partition1, broker1)
+
+    servers(broker1).shutdown()
+    TestUtils.assertLeader(client, partition1, broker2)
+    servers(broker1).startup()
+
+    val electResult = client.electLeaders(ElectionType.UNCLEAN, Set(partition1).asJava)
+    assertTrue(electResult.partitions.get.get(partition1).get.isInstanceOf[ElectionNotNeededException])
+  }
+
+  @Test
+  def testElectUncleanLeadersAndNoop(): Unit = {
+    // Case: one noop unclean leader election and one valid unclean leader election
+    client = AdminClient.create(createConfig)
+
+    val broker1 = 1
+    val broker2 = 2
+    val broker3 = 0
+    val assignment1 = Seq(broker1, broker2)
+    val assignment2 = Seq(broker1, broker3)
+
+    val topic = "unclean-test-topic-1"
+    val partition1 = new TopicPartition(topic, 0)
+    val partition2 = new TopicPartition(topic, 1)
+
+    TestUtils.createTopic(
+      zkClient,
+      topic,
+      Map(partition1.partition -> assignment1, partition2.partition -> assignment2),
+      servers
+    )
+
+    TestUtils.assertLeader(client, partition1, broker1)
+    TestUtils.assertLeader(client, partition2, broker1)
+
+    servers(broker2).shutdown()
+    TestUtils.waitForBrokersOutOfIsr(client, Set(partition1), Set(broker2))
+    servers(broker1).shutdown()
+    TestUtils.assertNoLeader(client, partition1)
+    TestUtils.assertLeader(client, partition2, broker3)
+    servers(broker2).startup()
+
+    val electResult = client.electLeaders(ElectionType.UNCLEAN, Set(partition1, partition2).asJava)
+    assertFalse(electResult.partitions.get.get(partition1).isPresent)
+    assertTrue(electResult.partitions.get.get(partition2).get.isInstanceOf[ElectionNotNeededException])
+    TestUtils.assertLeader(client, partition1, broker2)
+    TestUtils.assertLeader(client, partition2, broker3)
+  }
+
+  @Test
+  def testListReassignmentsDoesNotShowNonReassigningPartitions(): Unit = {
+    client = AdminClient.create(createConfig())
+
+    // Create topics
+    val topic = "list-reassignments-no-reassignments"
+    createTopic(topic, numPartitions = 1, replicationFactor = 3)
+    val tp = new TopicPartition(topic, 0)
+
+    val reassignmentsMap = client.listPartitionReassignments(Set(tp).asJava).reassignments().get()
+    assertEquals(0, reassignmentsMap.size())
+
+    val allReassignmentsMap = client.listPartitionReassignments().reassignments().get()
+    assertEquals(0, allReassignmentsMap.size())
+  }
+
+  @Test
+  def testListReassignmentsDoesNotShowDeletedPartitions(): Unit = {
+    client = AdminClient.create(createConfig())
+    
+    val topic = "list-reassignments-no-reassignments"
+    val tp = new TopicPartition(topic, 0)
+
+    val reassignmentsMap = client.listPartitionReassignments(Set(tp).asJava).reassignments().get()
+    assertEquals(0, reassignmentsMap.size())
+
+    val allReassignmentsMap = client.listPartitionReassignments().reassignments().get()
+    assertEquals(0, allReassignmentsMap.size())
+  }
+
+  @Test
+  def testValidIncrementalAlterConfigs(): Unit = {
+    client = AdminClient.create(createConfig)
+
+    // Create topics
+    val topic1 = "incremental-alter-configs-topic-1"
+    val topic1Resource = new ConfigResource(ConfigResource.Type.TOPIC, topic1)
+    val topic1CreateConfigs = new Properties
+    topic1CreateConfigs.setProperty(LogConfig.RetentionMsProp, "60000000")
+    topic1CreateConfigs.setProperty(LogConfig.CleanupPolicyProp, LogConfig.Compact)
+    createTopic(topic1, numPartitions = 1, replicationFactor = 1, topic1CreateConfigs)
+
+    val topic2 = "incremental-alter-configs-topic-2"
+    val topic2Resource = new ConfigResource(ConfigResource.Type.TOPIC, topic2)
+    createTopic(topic2)
+
+    // Alter topic configs
+    var topic1AlterConfigs = Seq(
+      new AlterConfigOp(new ConfigEntry(LogConfig.FlushMsProp, "1000"), AlterConfigOp.OpType.SET),
+      new AlterConfigOp(new ConfigEntry(LogConfig.CleanupPolicyProp, LogConfig.Delete), AlterConfigOp.OpType.APPEND),
+      new AlterConfigOp(new ConfigEntry(LogConfig.RetentionMsProp, ""), AlterConfigOp.OpType.DELETE)
+    ).asJavaCollection
+
+    val topic2AlterConfigs = Seq(
+      new AlterConfigOp(new ConfigEntry(LogConfig.MinCleanableDirtyRatioProp, "0.9"), AlterConfigOp.OpType.SET),
+      new AlterConfigOp(new ConfigEntry(LogConfig.CompressionTypeProp, "lz4"), AlterConfigOp.OpType.SET)
+    ).asJavaCollection
+
+    var alterResult = client.incrementalAlterConfigs(Map(
+      topic1Resource -> topic1AlterConfigs,
+      topic2Resource -> topic2AlterConfigs
+    ).asJava)
+
+    assertEquals(Set(topic1Resource, topic2Resource).asJava, alterResult.values.keySet)
+    alterResult.all.get
+
+    // Verify that topics were updated correctly
+    var describeResult = client.describeConfigs(Seq(topic1Resource, topic2Resource).asJava)
+    var configs = describeResult.all.get
+
+    assertEquals(2, configs.size)
+
+    assertEquals("1000", configs.get(topic1Resource).get(LogConfig.FlushMsProp).value)
+    assertEquals("compact,delete", configs.get(topic1Resource).get(LogConfig.CleanupPolicyProp).value)
+    assertEquals((Defaults.LogRetentionHours * 60 * 60 * 1000).toString, configs.get(topic1Resource).get(LogConfig.RetentionMsProp).value)
+
+    assertEquals("0.9", configs.get(topic2Resource).get(LogConfig.MinCleanableDirtyRatioProp).value)
+    assertEquals("lz4", configs.get(topic2Resource).get(LogConfig.CompressionTypeProp).value)
+
+    //verify subtract operation
+    topic1AlterConfigs = Seq(
+      new AlterConfigOp(new ConfigEntry(LogConfig.CleanupPolicyProp, LogConfig.Compact), AlterConfigOp.OpType.SUBTRACT)
+    ).asJava
+
+   alterResult = client.incrementalAlterConfigs(Map(
+      topic1Resource -> topic1AlterConfigs
+    ).asJava)
+    alterResult.all.get
+
+    // Verify that topics were updated correctly
+    describeResult = client.describeConfigs(Seq(topic1Resource).asJava)
+    configs = describeResult.all.get
+
+    assertEquals("delete", configs.get(topic1Resource).get(LogConfig.CleanupPolicyProp).value)
+    assertEquals("1000", configs.get(topic1Resource).get(LogConfig.FlushMsProp).value) // verify previous change is still intact
+
+    // Alter topics with validateOnly=true
+    topic1AlterConfigs = Seq(
+      new AlterConfigOp(new ConfigEntry(LogConfig.CleanupPolicyProp, LogConfig.Compact), AlterConfigOp.OpType.APPEND)
+    ).asJava
+
+    alterResult = client.incrementalAlterConfigs(Map(
+      topic1Resource -> topic1AlterConfigs
+    ).asJava, new AlterConfigsOptions().validateOnly(true))
+    alterResult.all.get
+
+    // Verify that topics were not updated due to validateOnly = true
+    describeResult = client.describeConfigs(Seq(topic1Resource).asJava)
+    configs = describeResult.all.get
+
+    assertEquals("delete", configs.get(topic1Resource).get(LogConfig.CleanupPolicyProp).value)
+
+    //Alter topics with validateOnly=true with invalid configs
+    topic1AlterConfigs = Seq(
+      new AlterConfigOp(new ConfigEntry(LogConfig.CompressionTypeProp, "zip"), AlterConfigOp.OpType.SET)
+    ).asJava
+
+    alterResult = client.incrementalAlterConfigs(Map(
+      topic1Resource -> topic1AlterConfigs
+    ).asJava, new AlterConfigsOptions().validateOnly(true))
+
+    assertFutureExceptionTypeEquals(alterResult.values().get(topic1Resource), classOf[InvalidRequestException],
+      Some("Invalid config value for resource"))
+  }
+
+  @Test
+  def testInvalidIncrementalAlterConfigs(): Unit = {
+    client = AdminClient.create(createConfig)
+
+    // Create topics
+    val topic1 = "incremental-alter-configs-topic-1"
+    val topic1Resource = new ConfigResource(ConfigResource.Type.TOPIC, topic1)
+    createTopic(topic1)
+
+    val topic2 = "incremental-alter-configs-topic-2"
+    val topic2Resource = new ConfigResource(ConfigResource.Type.TOPIC, topic2)
+    createTopic(topic2)
+
+    //Add duplicate Keys for topic1
+    var topic1AlterConfigs = Seq(
+      new AlterConfigOp(new ConfigEntry(LogConfig.MinCleanableDirtyRatioProp, "0.75"), AlterConfigOp.OpType.SET),
+      new AlterConfigOp(new ConfigEntry(LogConfig.MinCleanableDirtyRatioProp, "0.65"), AlterConfigOp.OpType.SET),
+      new AlterConfigOp(new ConfigEntry(LogConfig.CompressionTypeProp, "gzip"), AlterConfigOp.OpType.SET) // valid entry
+    ).asJavaCollection
+
+    //Add valid config for topic2
+    var topic2AlterConfigs = Seq(
+      new AlterConfigOp(new ConfigEntry(LogConfig.MinCleanableDirtyRatioProp, "0.9"), AlterConfigOp.OpType.SET)
+    ).asJavaCollection
+
+    var alterResult = client.incrementalAlterConfigs(Map(
+      topic1Resource -> topic1AlterConfigs,
+      topic2Resource -> topic2AlterConfigs
+    ).asJava)
+    assertEquals(Set(topic1Resource, topic2Resource).asJava, alterResult.values.keySet)
+
+    //InvalidRequestException error for topic1
+    assertFutureExceptionTypeEquals(alterResult.values().get(topic1Resource), classOf[InvalidRequestException],
+      Some("Error due to duplicate config keys"))
+
+    //operation should succeed for topic2
+    alterResult.values().get(topic2Resource).get()
+
+    // Verify that topic1 is not config not updated, and topic2 config is updated
+    val describeResult = client.describeConfigs(Seq(topic1Resource, topic2Resource).asJava)
+    val configs = describeResult.all.get
+    assertEquals(2, configs.size)
+
+    assertEquals(Defaults.LogCleanerMinCleanRatio.toString, configs.get(topic1Resource).get(LogConfig.MinCleanableDirtyRatioProp).value)
+    assertEquals(Defaults.CompressionType.toString, configs.get(topic1Resource).get(LogConfig.CompressionTypeProp).value)
+    assertEquals("0.9", configs.get(topic2Resource).get(LogConfig.MinCleanableDirtyRatioProp).value)
+
+    //check invalid use of append/subtract operation types
+    topic1AlterConfigs = Seq(
+      new AlterConfigOp(new ConfigEntry(LogConfig.CompressionTypeProp, "gzip"), AlterConfigOp.OpType.APPEND)
+    ).asJavaCollection
+
+    topic2AlterConfigs = Seq(
+      new AlterConfigOp(new ConfigEntry(LogConfig.CompressionTypeProp, "snappy"), AlterConfigOp.OpType.SUBTRACT)
+    ).asJavaCollection
+
+    alterResult = client.incrementalAlterConfigs(Map(
+      topic1Resource -> topic1AlterConfigs,
+      topic2Resource -> topic2AlterConfigs
+    ).asJava)
+    assertEquals(Set(topic1Resource, topic2Resource).asJava, alterResult.values.keySet)
+
+    assertFutureExceptionTypeEquals(alterResult.values().get(topic1Resource), classOf[InvalidRequestException],
+      Some("Config value append is not allowed for config"))
+
+    assertFutureExceptionTypeEquals(alterResult.values().get(topic2Resource), classOf[InvalidRequestException],
+      Some("Config value subtract is not allowed for config"))
+
+
+    //try to add invalid config
+    topic1AlterConfigs = Seq(
+      new AlterConfigOp(new ConfigEntry(LogConfig.MinCleanableDirtyRatioProp, "1.1"), AlterConfigOp.OpType.SET)
+    ).asJavaCollection
+
+    alterResult = client.incrementalAlterConfigs(Map(
+      topic1Resource -> topic1AlterConfigs
+    ).asJava)
+    assertEquals(Set(topic1Resource).asJava, alterResult.values.keySet)
+
+    assertFutureExceptionTypeEquals(alterResult.values().get(topic1Resource), classOf[InvalidRequestException],
+      Some("Invalid config value for resource"))
+  }
+
+  @Test
+  def testInvalidAlterPartitionReassignments(): Unit = {
+    client = AdminClient.create(createConfig)
+    val topic = "alter-reassignments-topic-1"
+    val tp1 = new TopicPartition(topic, 0)
+    val tp2 = new TopicPartition(topic, 1)
+    val tp3 = new TopicPartition(topic, 2)
+    createTopic(topic, numPartitions = 4)
+
+
+    val validAssignment = Optional.of(new NewPartitionReassignment(
+      (0 until brokerCount).map(_.asInstanceOf[Integer]).asJava
+    ))
+
+    val nonExistentTp1 = new TopicPartition("topicA", 0)
+    val nonExistentTp2 = new TopicPartition(topic, 4)
+    val nonExistentPartitionsResult = client.alterPartitionReassignments(Map(
+      tp1 -> validAssignment,
+      tp2 -> validAssignment,
+      tp3 -> validAssignment,
+      nonExistentTp1 -> validAssignment,
+      nonExistentTp2 -> validAssignment
+    ).asJava).values()
+    assertFutureExceptionTypeEquals(nonExistentPartitionsResult.get(nonExistentTp1), classOf[UnknownTopicOrPartitionException])
+    assertFutureExceptionTypeEquals(nonExistentPartitionsResult.get(nonExistentTp2), classOf[UnknownTopicOrPartitionException])
+
+    val extraNonExistentReplica = Optional.of(new NewPartitionReassignment((0 until brokerCount + 1).map(_.asInstanceOf[Integer]).asJava))
+    val negativeIdReplica = Optional.of(new NewPartitionReassignment(Seq(-3, -2, -1).map(_.asInstanceOf[Integer]).asJava))
+    val duplicateReplica = Optional.of(new NewPartitionReassignment(Seq(0, 1, 1).map(_.asInstanceOf[Integer]).asJava))
+    val invalidReplicaResult = client.alterPartitionReassignments(Map(
+      tp1 -> extraNonExistentReplica,
+      tp2 -> negativeIdReplica,
+      tp3 -> duplicateReplica
+    ).asJava).values()
+    assertFutureExceptionTypeEquals(invalidReplicaResult.get(tp1), classOf[InvalidReplicaAssignmentException])
+    assertFutureExceptionTypeEquals(invalidReplicaResult.get(tp2), classOf[InvalidReplicaAssignmentException])
+    assertFutureExceptionTypeEquals(invalidReplicaResult.get(tp3), classOf[InvalidReplicaAssignmentException])
+  }
+
+  @Test
+  def testLongTopicNames(): Unit = {
+    val client = AdminClient.create(createConfig)
+    val longTopicName = String.join("", Collections.nCopies(249, "x"));
+    val invalidTopicName = String.join("", Collections.nCopies(250, "x"));
+    val newTopics2 = Seq(new NewTopic(invalidTopicName, 3, 3.toShort),
+                         new NewTopic(longTopicName, 3, 3.toShort))
+    val results = client.createTopics(newTopics2.asJava).values()
+    assertTrue(results.containsKey(longTopicName))
+    results.get(longTopicName).get()
+    assertTrue(results.containsKey(invalidTopicName))
+    assertFutureExceptionTypeEquals(results.get(invalidTopicName), classOf[InvalidTopicException])
+    assertFutureExceptionTypeEquals(client.alterReplicaLogDirs(
+      Map(new TopicPartitionReplica(longTopicName, 0, 0) -> servers(0).config.logDirs(0)).asJava).all(),
+      classOf[InvalidTopicException])
+    client.close()
+  }
+
+  @Test
+  def testDescribeConfigsForLog4jLogLevels(): Unit = {
+    client = AdminClient.create(createConfig())
+
+    val loggerConfig = describeBrokerLoggers()
+    val rootLogLevel = loggerConfig.get(Log4jController.ROOT_LOGGER).value()
+    val logCleanerLogLevelConfig = loggerConfig.get("kafka.cluster.Replica")
+    assertEquals(rootLogLevel, logCleanerLogLevelConfig.value()) // we expect an undefined log level to be the same as the root logger
+    assertEquals("kafka.cluster.Replica", logCleanerLogLevelConfig.name())
+    assertEquals(ConfigEntry.ConfigSource.DYNAMIC_BROKER_LOGGER_CONFIG, logCleanerLogLevelConfig.source())
+    assertEquals(false, logCleanerLogLevelConfig.isReadOnly)
+    assertEquals(false, logCleanerLogLevelConfig.isSensitive)
+    assertTrue(logCleanerLogLevelConfig.synonyms().isEmpty)
+  }
+
+  @Test
+  @Ignore // To be re-enabled once KAFKA-8779 is resolved
+  def testIncrementalAlterConfigsForLog4jLogLevels(): Unit = {
+    client = AdminClient.create(createConfig())
+
+    val initialLoggerConfig = describeBrokerLoggers()
+    val initialRootLogLevel = initialLoggerConfig.get(Log4jController.ROOT_LOGGER).value()
+    assertEquals(initialRootLogLevel, initialLoggerConfig.get("kafka.controller.KafkaController").value())
+    assertEquals(initialRootLogLevel, initialLoggerConfig.get("kafka.log.LogCleaner").value())
+    assertEquals(initialRootLogLevel, initialLoggerConfig.get("kafka.server.ReplicaManager").value())
+
+    val newRootLogLevel = LogLevelConfig.DEBUG_LOG_LEVEL
+    val alterRootLoggerEntry = Seq(
+      new AlterConfigOp(new ConfigEntry(Log4jController.ROOT_LOGGER, newRootLogLevel), AlterConfigOp.OpType.SET)
+    ).asJavaCollection
+    // Test validateOnly does not change anything
+    alterBrokerLoggers(alterRootLoggerEntry, validateOnly = true)
+    val validatedLoggerConfig = describeBrokerLoggers()
+    assertEquals(initialRootLogLevel, validatedLoggerConfig.get(Log4jController.ROOT_LOGGER).value())
+    assertEquals(initialRootLogLevel, validatedLoggerConfig.get("kafka.controller.KafkaController").value())
+    assertEquals(initialRootLogLevel, validatedLoggerConfig.get("kafka.log.LogCleaner").value())
+    assertEquals(initialRootLogLevel, validatedLoggerConfig.get("kafka.server.ReplicaManager").value())
+    assertEquals(initialRootLogLevel, validatedLoggerConfig.get("kafka.zookeeper.ZooKeeperClient").value())
+
+    // test that we can change them and unset loggers still use the root's log level
+    alterBrokerLoggers(alterRootLoggerEntry)
+    val changedRootLoggerConfig = describeBrokerLoggers()
+    assertEquals(newRootLogLevel, changedRootLoggerConfig.get(Log4jController.ROOT_LOGGER).value())
+    assertEquals(newRootLogLevel, changedRootLoggerConfig.get("kafka.controller.KafkaController").value())
+    assertEquals(newRootLogLevel, changedRootLoggerConfig.get("kafka.log.LogCleaner").value())
+    assertEquals(newRootLogLevel, changedRootLoggerConfig.get("kafka.server.ReplicaManager").value())
+    assertEquals(newRootLogLevel, changedRootLoggerConfig.get("kafka.zookeeper.ZooKeeperClient").value())
+
+    // alter the ZK client's logger so we can later test resetting it
+    val alterZKLoggerEntry = Seq(
+      new AlterConfigOp(new ConfigEntry("kafka.zookeeper.ZooKeeperClient", LogLevelConfig.ERROR_LOG_LEVEL), AlterConfigOp.OpType.SET)
+    ).asJavaCollection
+    alterBrokerLoggers(alterZKLoggerEntry)
+    val changedZKLoggerConfig = describeBrokerLoggers()
+    assertEquals(LogLevelConfig.ERROR_LOG_LEVEL, changedZKLoggerConfig.get("kafka.zookeeper.ZooKeeperClient").value())
+
+    // properly test various set operations and one delete
+    val alterLogLevelsEntries = Seq(
+      new AlterConfigOp(new ConfigEntry("kafka.controller.KafkaController", LogLevelConfig.INFO_LOG_LEVEL), AlterConfigOp.OpType.SET),
+      new AlterConfigOp(new ConfigEntry("kafka.log.LogCleaner", LogLevelConfig.ERROR_LOG_LEVEL), AlterConfigOp.OpType.SET),
+      new AlterConfigOp(new ConfigEntry("kafka.server.ReplicaManager", LogLevelConfig.TRACE_LOG_LEVEL), AlterConfigOp.OpType.SET),
+      new AlterConfigOp(new ConfigEntry("kafka.zookeeper.ZooKeeperClient", ""), AlterConfigOp.OpType.DELETE) // should reset to the root logger level
+    ).asJavaCollection
+    alterBrokerLoggers(alterLogLevelsEntries)
+    val alteredLoggerConfig = describeBrokerLoggers()
+    assertEquals(newRootLogLevel, alteredLoggerConfig.get(Log4jController.ROOT_LOGGER).value())
+    assertEquals(LogLevelConfig.INFO_LOG_LEVEL, alteredLoggerConfig.get("kafka.controller.KafkaController").value())
+    assertEquals(LogLevelConfig.ERROR_LOG_LEVEL, alteredLoggerConfig.get("kafka.log.LogCleaner").value())
+    assertEquals(LogLevelConfig.TRACE_LOG_LEVEL, alteredLoggerConfig.get("kafka.server.ReplicaManager").value())
+    assertEquals(newRootLogLevel, alteredLoggerConfig.get("kafka.zookeeper.ZooKeeperClient").value())
+  }
+
+  /**
+    * 1. Assume ROOT logger == TRACE
+    * 2. Change kafka.controller.KafkaController logger to INFO
+    * 3. Unset kafka.controller.KafkaController via AlterConfigOp.OpType.DELETE (resets it to the root logger - TRACE)
+    * 4. Change ROOT logger to ERROR
+    * 5. Ensure the kafka.controller.KafkaController logger's level is ERROR (the curent root logger level)
+    */
+  @Test
+  @Ignore // To be re-enabled once KAFKA-8779 is resolved
+  def testIncrementalAlterConfigsForLog4jLogLevelsCanResetLoggerToCurrentRoot(): Unit = {
+    client = AdminClient.create(createConfig())
+    // step 1 - configure root logger
+    val initialRootLogLevel = LogLevelConfig.TRACE_LOG_LEVEL
+    val alterRootLoggerEntry = Seq(
+      new AlterConfigOp(new ConfigEntry(Log4jController.ROOT_LOGGER, initialRootLogLevel), AlterConfigOp.OpType.SET)
+    ).asJavaCollection
+    alterBrokerLoggers(alterRootLoggerEntry)
+    val initialLoggerConfig = describeBrokerLoggers()
+    assertEquals(initialRootLogLevel, initialLoggerConfig.get(Log4jController.ROOT_LOGGER).value())
+    assertEquals(initialRootLogLevel, initialLoggerConfig.get("kafka.controller.KafkaController").value())
+
+    // step 2 - change KafkaController logger to INFO
+    val alterControllerLoggerEntry = Seq(
+      new AlterConfigOp(new ConfigEntry("kafka.controller.KafkaController", LogLevelConfig.INFO_LOG_LEVEL), AlterConfigOp.OpType.SET)
+    ).asJavaCollection
+    alterBrokerLoggers(alterControllerLoggerEntry)
+    val changedControllerLoggerConfig = describeBrokerLoggers()
+    assertEquals(initialRootLogLevel, changedControllerLoggerConfig.get(Log4jController.ROOT_LOGGER).value())
+    assertEquals(LogLevelConfig.INFO_LOG_LEVEL, changedControllerLoggerConfig.get("kafka.controller.KafkaController").value())
+
+    // step 3 - unset KafkaController logger
+    val deleteControllerLoggerEntry = Seq(
+      new AlterConfigOp(new ConfigEntry("kafka.controller.KafkaController", ""), AlterConfigOp.OpType.DELETE)
+    ).asJavaCollection
+    alterBrokerLoggers(deleteControllerLoggerEntry)
+    val deletedControllerLoggerConfig = describeBrokerLoggers()
+    assertEquals(initialRootLogLevel, deletedControllerLoggerConfig.get(Log4jController.ROOT_LOGGER).value())
+    assertEquals(initialRootLogLevel, deletedControllerLoggerConfig.get("kafka.controller.KafkaController").value())
+
+    val newRootLogLevel = LogLevelConfig.ERROR_LOG_LEVEL
+    val newAlterRootLoggerEntry = Seq(
+      new AlterConfigOp(new ConfigEntry(Log4jController.ROOT_LOGGER, newRootLogLevel), AlterConfigOp.OpType.SET)
+    ).asJavaCollection
+    alterBrokerLoggers(newAlterRootLoggerEntry)
+    val newRootLoggerConfig = describeBrokerLoggers()
+    assertEquals(newRootLogLevel, newRootLoggerConfig.get(Log4jController.ROOT_LOGGER).value())
+    assertEquals(newRootLogLevel, newRootLoggerConfig.get("kafka.controller.KafkaController").value())
+  }
+
+  @Test
+  @Ignore // To be re-enabled once KAFKA-8779 is resolved
+  def testIncrementalAlterConfigsForLog4jLogLevelsCannotResetRootLogger(): Unit = {
+    client = AdminClient.create(createConfig())
+    val deleteRootLoggerEntry = Seq(
+      new AlterConfigOp(new ConfigEntry(Log4jController.ROOT_LOGGER, ""), AlterConfigOp.OpType.DELETE)
+    ).asJavaCollection
+
+    assertTrue(intercept[ExecutionException](alterBrokerLoggers(deleteRootLoggerEntry)).getCause.isInstanceOf[InvalidRequestException])
+  }
+
+  @Test
+  @Ignore // To be re-enabled once KAFKA-8779 is resolved
+  def testIncrementalAlterConfigsForLog4jLogLevelsDoesNotWorkWithInvalidConfigs(): Unit = {
+    client = AdminClient.create(createConfig())
+    val validLoggerName = "kafka.server.KafkaRequestHandler"
+    val expectedValidLoggerLogLevel = describeBrokerLoggers().get(validLoggerName)
+    def assertLogLevelDidNotChange(): Unit = {
+      assertEquals(
+        expectedValidLoggerLogLevel,
+        describeBrokerLoggers().get(validLoggerName)
+      )
+    }
+
+    val appendLogLevelEntries = Seq(
+      new AlterConfigOp(new ConfigEntry("kafka.server.KafkaRequestHandler", LogLevelConfig.INFO_LOG_LEVEL), AlterConfigOp.OpType.SET), // valid
+      new AlterConfigOp(new ConfigEntry("kafka.network.SocketServer", LogLevelConfig.ERROR_LOG_LEVEL), AlterConfigOp.OpType.APPEND) // append is not supported
+    ).asJavaCollection
+    assertTrue(intercept[ExecutionException](alterBrokerLoggers(appendLogLevelEntries)).getCause.isInstanceOf[InvalidRequestException])
+    assertLogLevelDidNotChange()
+
+    val subtractLogLevelEntries = Seq(
+      new AlterConfigOp(new ConfigEntry("kafka.server.KafkaRequestHandler", LogLevelConfig.INFO_LOG_LEVEL), AlterConfigOp.OpType.SET), // valid
+      new AlterConfigOp(new ConfigEntry("kafka.network.SocketServer", LogLevelConfig.ERROR_LOG_LEVEL), AlterConfigOp.OpType.SUBTRACT) // subtract is not supported
+    ).asJavaCollection
+    assertTrue(intercept[ExecutionException](alterBrokerLoggers(subtractLogLevelEntries)).getCause.isInstanceOf[InvalidRequestException])
+    assertLogLevelDidNotChange()
+
+    val invalidLogLevelLogLevelEntries = Seq(
+      new AlterConfigOp(new ConfigEntry("kafka.server.KafkaRequestHandler", LogLevelConfig.INFO_LOG_LEVEL), AlterConfigOp.OpType.SET), // valid
+      new AlterConfigOp(new ConfigEntry("kafka.network.SocketServer", "OFF"), AlterConfigOp.OpType.SET) // OFF is not a valid log level
+    ).asJavaCollection
+    assertTrue(intercept[ExecutionException](alterBrokerLoggers(invalidLogLevelLogLevelEntries)).getCause.isInstanceOf[InvalidRequestException])
+    assertLogLevelDidNotChange()
+
+    val invalidLoggerNameLogLevelEntries = Seq(
+      new AlterConfigOp(new ConfigEntry("kafka.server.KafkaRequestHandler", LogLevelConfig.INFO_LOG_LEVEL), AlterConfigOp.OpType.SET), // valid
+      new AlterConfigOp(new ConfigEntry("Some Other LogCleaner", LogLevelConfig.ERROR_LOG_LEVEL), AlterConfigOp.OpType.SET) // invalid logger name is not supported
+    ).asJavaCollection
+    assertTrue(intercept[ExecutionException](alterBrokerLoggers(invalidLoggerNameLogLevelEntries)).getCause.isInstanceOf[InvalidRequestException])
+    assertLogLevelDidNotChange()
+  }
+
+  /**
+    * The AlterConfigs API is deprecated and should not support altering log levels
+    */
+  @Test
+  @Ignore // To be re-enabled once KAFKA-8779 is resolved
+  def testAlterConfigsForLog4jLogLevelsDoesNotWork(): Unit = {
+    client = AdminClient.create(createConfig())
+
+    val alterLogLevelsEntries = Seq(
+      new ConfigEntry("kafka.controller.KafkaController", LogLevelConfig.INFO_LOG_LEVEL)
+    ).asJavaCollection
+    val alterResult = client.alterConfigs(Map(brokerLoggerConfigResource -> new Config(alterLogLevelsEntries)).asJava)
+    assertTrue(intercept[ExecutionException](alterResult.values.get(brokerLoggerConfigResource).get).getCause.isInstanceOf[InvalidRequestException])
+  }
+
+  def alterBrokerLoggers(entries: util.Collection[AlterConfigOp], validateOnly: Boolean = false): Unit = {
+    if (!validateOnly) {
+      for (entry <- entries.asScala)
+        changedBrokerLoggers.add(entry.configEntry().name())
+    }
+
+    client.incrementalAlterConfigs(Map(brokerLoggerConfigResource -> entries).asJava, new AlterConfigsOptions().validateOnly(validateOnly))
+      .values.get(brokerLoggerConfigResource).get()
+  }
+
+  def describeBrokerLoggers(): Config =
+    client.describeConfigs(Collections.singletonList(brokerLoggerConfigResource)).values.get(brokerLoggerConfigResource).get()
+
+  /**
+    * Due to the fact that log4j is not re-initialized across tests, changing a logger's log level persists across test classes.
+    * We need to clean up the changes done while testing.
+    */
+  def teardownBrokerLoggers(): Unit = {
+    if (changedBrokerLoggers.nonEmpty) {
+      val validLoggers = describeBrokerLoggers().entries().asScala.filterNot(_.name().equals(Log4jController.ROOT_LOGGER)).map(_.name).toSet
+      val unsetBrokerLoggersEntries = changedBrokerLoggers
+        .intersect(validLoggers)
+        .map { logger => new AlterConfigOp(new ConfigEntry(logger, ""), AlterConfigOp.OpType.DELETE) }
+        .asJavaCollection
+
+      // ensure that we first reset the root logger to an arbitrary log level. Note that we cannot reset it to its original value
+      alterBrokerLoggers(List(
+        new AlterConfigOp(new ConfigEntry(Log4jController.ROOT_LOGGER, LogLevelConfig.FATAL_LOG_LEVEL), AlterConfigOp.OpType.SET)
+      ).asJavaCollection)
+      alterBrokerLoggers(unsetBrokerLoggersEntries)
+
+      changedBrokerLoggers.clear()
+    }
   }
 }
 
 object AdminClientIntegrationTest {
 
-  import org.scalatest.Assertions._
-
-  def checkValidAlterConfigs(client: AdminClient, topicResource1: ConfigResource, topicResource2: ConfigResource): Unit = {
+  def checkValidAlterConfigs(client: Admin, topicResource1: ConfigResource, topicResource2: ConfigResource): Unit = {
     // Alter topics
     var topicConfigEntries1 = Seq(
       new ConfigEntry(LogConfig.FlushMsProp, "1000")
@@ -1429,7 +2335,7 @@ object AdminClientIntegrationTest {
     assertEquals("0.9", configs.get(topicResource2).get(LogConfig.MinCleanableDirtyRatioProp).value)
   }
 
-  def checkInvalidAlterConfigs(zkClient: KafkaZkClient, servers: Seq[KafkaServer], client: AdminClient): Unit = {
+  def checkInvalidAlterConfigs(zkClient: KafkaZkClient, servers: Seq[KafkaServer], client: Admin): Unit = {
     // Create topics
     val topic1 = "invalid-alter-configs-topic-1"
     val topicResource1 = new ConfigResource(ConfigResource.Type.TOPIC, topic1)
@@ -1504,4 +2410,22 @@ object AdminClientIntegrationTest {
     assertEquals(Defaults.CompressionType.toString, configs.get(brokerResource).get(KafkaConfig.CompressionTypeProp).value)
   }
 
+  private def getTopicMetadata(client: Admin,
+                               topic: String,
+                               describeOptions: DescribeTopicsOptions = new DescribeTopicsOptions,
+                               expectedNumPartitionsOpt: Option[Int] = None): TopicDescription = {
+    var result: TopicDescription = null
+
+    TestUtils.waitUntilTrue(() => {
+      val topicResult = client.describeTopics(Set(topic).asJava, describeOptions).values.get(topic)
+      try {
+        result = topicResult.get
+        expectedNumPartitionsOpt.map(_ == result.partitions.size).getOrElse(true)
+      } catch {
+        case e: ExecutionException if e.getCause.isInstanceOf[UnknownTopicOrPartitionException] => false  // metadata may not have propagated yet, so retry
+      }
+    }, s"Timed out waiting for metadata for $topic")
+
+    result
+  }
 }
diff --git a/core/src/test/scala/integration/kafka/api/AdminClientWithPoliciesIntegrationTest.scala b/core/src/test/scala/integration/kafka/api/AdminClientWithPoliciesIntegrationTest.scala
index 1bea039e20bbf..0da0829ad628e 100644
--- a/core/src/test/scala/integration/kafka/api/AdminClientWithPoliciesIntegrationTest.scala
+++ b/core/src/test/scala/integration/kafka/api/AdminClientWithPoliciesIntegrationTest.scala
@@ -21,7 +21,7 @@ import kafka.integration.KafkaServerTestHarness
 import kafka.log.LogConfig
 import kafka.server.{Defaults, KafkaConfig}
 import kafka.utils.{Logging, TestUtils}
-import org.apache.kafka.clients.admin.{AdminClient, AdminClientConfig, AlterConfigsOptions, Config, ConfigEntry}
+import org.apache.kafka.clients.admin.{Admin, AdminClient, AdminClientConfig, AlterConfigsOptions, Config, ConfigEntry}
 import org.apache.kafka.common.config.{ConfigResource, TopicConfig}
 import org.apache.kafka.common.errors.{InvalidRequestException, PolicyViolationException}
 import org.apache.kafka.common.utils.Utils
@@ -29,6 +29,7 @@ import org.apache.kafka.server.policy.AlterConfigPolicy
 import org.junit.Assert.{assertEquals, assertNull, assertTrue}
 import org.junit.{After, Before, Rule, Test}
 import org.junit.rules.Timeout
+import org.scalatest.Assertions.intercept
 
 import scala.collection.JavaConverters._
 
@@ -39,7 +40,7 @@ class AdminClientWithPoliciesIntegrationTest extends KafkaServerTestHarness with
 
   import AdminClientWithPoliciesIntegrationTest._
 
-  var client: AdminClient = null
+  var client: Admin = null
   val brokerCount = 3
 
   @Rule
diff --git a/core/src/test/scala/integration/kafka/api/AuthorizerIntegrationTest.scala b/core/src/test/scala/integration/kafka/api/AuthorizerIntegrationTest.scala
index 044f5952670f5..b25e81d3a1364 100644
--- a/core/src/test/scala/integration/kafka/api/AuthorizerIntegrationTest.scala
+++ b/core/src/test/scala/integration/kafka/api/AuthorizerIntegrationTest.scala
@@ -22,39 +22,65 @@ import java.time.Duration
 import kafka.admin.ConsumerGroupCommand.{ConsumerGroupCommandOptions, ConsumerGroupService}
 import kafka.log.LogConfig
 import kafka.network.SocketServer
-import kafka.security.auth._
+import kafka.security.auth.{SimpleAclAuthorizer, Topic, ResourceType => AuthResourceType}
+import kafka.security.authorizer.AuthorizerUtils.WildcardHost
 import kafka.server.{BaseRequestTest, KafkaConfig}
 import kafka.utils.TestUtils
-import org.apache.kafka.clients.admin.{AdminClient, AdminClientConfig}
+import org.apache.kafka.clients.admin.{Admin, AdminClient, AdminClientConfig, AlterConfigOp}
 import org.apache.kafka.clients.consumer._
 import org.apache.kafka.clients.consumer.internals.NoOpConsumerRebalanceListener
 import org.apache.kafka.clients.producer._
-import org.apache.kafka.common.acl.{AccessControlEntry, AccessControlEntryFilter, AclBinding, AclBindingFilter, AclOperation, AclPermissionType}
-import org.apache.kafka.common.config.ConfigResource
+import org.apache.kafka.common.ElectionType
+import org.apache.kafka.common.acl.{AccessControlEntry, AccessControlEntryFilter, AclBinding, AclBindingFilter, AclOperation}
+import org.apache.kafka.common.acl.AclOperation._
+import org.apache.kafka.common.acl.AclPermissionType.{ALLOW, DENY}
+import org.apache.kafka.common.config.{ConfigResource, LogLevelConfig}
 import org.apache.kafka.common.errors._
 import org.apache.kafka.common.internals.Topic.GROUP_METADATA_TOPIC_NAME
-import org.apache.kafka.common.message.{CreateTopicsRequestData, DescribeGroupsRequestData, LeaveGroupRequestData}
-import org.apache.kafka.common.message.CreateTopicsRequestData.{CreatableTopic, CreatableTopicSet}
+import org.apache.kafka.common.message.AlterPartitionReassignmentsRequestData
+import org.apache.kafka.common.message.ControlledShutdownRequestData
+import org.apache.kafka.common.message.CreateTopicsRequestData
+import org.apache.kafka.common.message.CreateTopicsRequestData.{CreatableTopic, CreatableTopicCollection}
+import org.apache.kafka.common.message.DeleteGroupsRequestData
+import org.apache.kafka.common.message.DeleteTopicsRequestData
+import org.apache.kafka.common.message.DescribeGroupsRequestData
+import org.apache.kafka.common.message.FindCoordinatorRequestData
+import org.apache.kafka.common.message.HeartbeatRequestData
+import org.apache.kafka.common.message.IncrementalAlterConfigsRequestData
+import org.apache.kafka.common.message.IncrementalAlterConfigsRequestData.{AlterConfigsResource, AlterableConfig, AlterableConfigCollection}
+import org.apache.kafka.common.message.JoinGroupRequestData
+import org.apache.kafka.common.message.JoinGroupRequestData.JoinGroupRequestProtocolCollection
+import org.apache.kafka.common.message.LeaderAndIsrRequestData.LeaderAndIsrPartitionState
+import org.apache.kafka.common.message.LeaveGroupRequestData.MemberIdentity
+import org.apache.kafka.common.message.ListPartitionReassignmentsRequestData
+import org.apache.kafka.common.message.OffsetCommitRequestData
+import org.apache.kafka.common.message.SyncGroupRequestData
+import org.apache.kafka.common.message.JoinGroupRequestData.JoinGroupRequestProtocolCollection
+import org.apache.kafka.common.message.LeaveGroupRequestData.MemberIdentity
+import org.apache.kafka.common.message.UpdateMetadataRequestData.{UpdateMetadataBroker, UpdateMetadataEndpoint, UpdateMetadataPartitionState}
 import org.apache.kafka.common.network.ListenerName
 import org.apache.kafka.common.protocol.{ApiKeys, Errors}
-import org.apache.kafka.common.record.{CompressionType, MemoryRecords, Records, SimpleRecord}
+import org.apache.kafka.common.record.{CompressionType, MemoryRecords, RecordBatch, Records, SimpleRecord}
 import org.apache.kafka.common.requests.CreateAclsRequest.AclCreation
 import org.apache.kafka.common.requests._
 import org.apache.kafka.common.resource.PatternType.LITERAL
-import org.apache.kafka.common.resource.{ResourcePattern, ResourcePatternFilter, ResourceType => AdminResourceType}
+import org.apache.kafka.common.resource.{Resource, ResourcePattern, ResourcePatternFilter, ResourceType}
+import org.apache.kafka.common.resource.ResourceType._
 import org.apache.kafka.common.security.auth.{KafkaPrincipal, SecurityProtocol}
 import org.apache.kafka.common.{KafkaException, Node, TopicPartition, requests}
 import org.apache.kafka.test.{TestUtils => JTestUtils}
 import org.junit.Assert._
 import org.junit.{After, Assert, Before, Test}
+import org.scalatest.Assertions.intercept
 
 import scala.collection.JavaConverters._
+import scala.compat.java8.OptionConverters._
 import scala.collection.mutable
 import scala.collection.mutable.Buffer
 
 class AuthorizerIntegrationTest extends BaseRequestTest {
 
-  override def numBrokers: Int = 1
+  override def brokerCount: Int = 1
   val brokerId: Integer = 0
   def userPrincipal = KafkaPrincipal.ANONYMOUS
 
@@ -71,42 +97,46 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
   val logDir = "logDir"
   val deleteRecordsPartition = new TopicPartition(deleteTopic, part)
   val group = "my-group"
-  val topicResource = Resource(Topic, topic, LITERAL)
-  val groupResource = Resource(Group, group, LITERAL)
-  val deleteTopicResource = Resource(Topic, deleteTopic, LITERAL)
-  val transactionalIdResource = Resource(TransactionalId, transactionalId, LITERAL)
-  val createTopicResource = Resource(Topic, createTopic, LITERAL)
-
-  val groupReadAcl = Map(groupResource -> Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Read)))
-  val groupDescribeAcl = Map(groupResource -> Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Describe)))
-  val groupDeleteAcl = Map(groupResource -> Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Delete)))
-  val clusterAcl = Map(Resource.ClusterResource -> Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, ClusterAction)))
-  val clusterCreateAcl = Map(Resource.ClusterResource -> Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Create)))
-  val clusterAlterAcl = Map(Resource.ClusterResource -> Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Alter)))
-  val clusterDescribeAcl = Map(Resource.ClusterResource -> Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Describe)))
-  val clusterIdempotentWriteAcl = Map(Resource.ClusterResource -> Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, IdempotentWrite)))
-  val topicCreateAcl = Map(createTopicResource -> Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Create)))
-  val topicReadAcl = Map(topicResource -> Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Read)))
-  val topicWriteAcl = Map(topicResource -> Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Write)))
-  val topicDescribeAcl = Map(topicResource -> Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Describe)))
-  val topicAlterAcl = Map(topicResource -> Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Alter)))
-  val topicDeleteAcl = Map(deleteTopicResource -> Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Delete)))
-  val topicDescribeConfigsAcl = Map(topicResource -> Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, DescribeConfigs)))
-  val topicAlterConfigsAcl = Map(topicResource -> Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, AlterConfigs)))
-  val transactionIdWriteAcl = Map(transactionalIdResource -> Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Write)))
-  val transactionalIdDescribeAcl = Map(transactionalIdResource -> Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Describe)))
+  val clusterResource = new ResourcePattern(CLUSTER, Resource.CLUSTER_NAME, LITERAL)
+  val topicResource = new ResourcePattern(TOPIC, topic, LITERAL)
+  val groupResource = new ResourcePattern(GROUP, group, LITERAL)
+  val deleteTopicResource = new ResourcePattern(TOPIC, deleteTopic, LITERAL)
+  val transactionalIdResource = new ResourcePattern(TRANSACTIONAL_ID, transactionalId, LITERAL)
+  val createTopicResource = new ResourcePattern(TOPIC, createTopic, LITERAL)
+  val userPrincipalStr = userPrincipal.toString
+
+  val groupReadAcl = Map(groupResource -> Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)))
+  val groupDescribeAcl = Map(groupResource -> Set(new AccessControlEntry(userPrincipalStr, WildcardHost, DESCRIBE, ALLOW)))
+  val groupDeleteAcl = Map(groupResource -> Set(new AccessControlEntry(userPrincipalStr, WildcardHost, DELETE, ALLOW)))
+  val clusterAcl = Map(clusterResource -> Set(new AccessControlEntry(userPrincipalStr, WildcardHost, CLUSTER_ACTION, ALLOW)))
+  val clusterCreateAcl = Map(clusterResource -> Set(new AccessControlEntry(userPrincipalStr, WildcardHost, CREATE, ALLOW)))
+  val clusterAlterAcl = Map(clusterResource -> Set(new AccessControlEntry(userPrincipalStr, WildcardHost, ALTER, ALLOW)))
+  val clusterDescribeAcl = Map(clusterResource -> Set(new AccessControlEntry(userPrincipalStr, WildcardHost, DESCRIBE, ALLOW)))
+  val clusterAlterConfigsAcl = Map(clusterResource -> Set(new AccessControlEntry(userPrincipalStr, WildcardHost, ALTER_CONFIGS, ALLOW)))
+  val clusterIdempotentWriteAcl = Map(clusterResource -> Set(new AccessControlEntry(userPrincipalStr, WildcardHost, IDEMPOTENT_WRITE, ALLOW)))
+  val topicCreateAcl = Map(createTopicResource -> Set(new AccessControlEntry(userPrincipalStr, WildcardHost, CREATE, ALLOW)))
+  val topicReadAcl = Map(topicResource -> Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)))
+  val topicWriteAcl = Map(topicResource -> Set(new AccessControlEntry(userPrincipalStr, WildcardHost, WRITE, ALLOW)))
+  val topicDescribeAcl = Map(topicResource -> Set(new AccessControlEntry(userPrincipalStr, WildcardHost, DESCRIBE, ALLOW)))
+  val topicAlterAcl = Map(topicResource -> Set(new AccessControlEntry(userPrincipalStr, WildcardHost, ALTER, ALLOW)))
+  val topicDeleteAcl = Map(deleteTopicResource -> Set(new AccessControlEntry(userPrincipalStr, WildcardHost, DELETE, ALLOW)))
+  val topicDescribeConfigsAcl = Map(topicResource -> Set(new AccessControlEntry(userPrincipalStr, WildcardHost, DESCRIBE_CONFIGS, ALLOW)))
+  val topicAlterConfigsAcl = Map(topicResource -> Set(new AccessControlEntry(userPrincipalStr, WildcardHost, ALTER_CONFIGS, ALLOW)))
+  val transactionIdWriteAcl = Map(transactionalIdResource -> Set(new AccessControlEntry(userPrincipalStr, WildcardHost, WRITE, ALLOW)))
+  val transactionalIdDescribeAcl = Map(transactionalIdResource -> Set(new AccessControlEntry(userPrincipalStr, WildcardHost, DESCRIBE, ALLOW)))
 
   val numRecords = 1
-  val adminClients = Buffer[AdminClient]()
+  val adminClients = Buffer[Admin]()
 
   producerConfig.setProperty(ProducerConfig.ACKS_CONFIG, "1")
   producerConfig.setProperty(ProducerConfig.MAX_BLOCK_MS_CONFIG, "50000")
   consumerConfig.setProperty(ConsumerConfig.GROUP_ID_CONFIG, group)
 
-  override def propertyOverrides(properties: Properties): Unit = {
+  override def brokerPropertyOverrides(properties: Properties): Unit = {
     properties.put(KafkaConfig.AuthorizerClassNameProp, classOf[SimpleAclAuthorizer].getName)
     properties.put(KafkaConfig.BrokerIdProp, brokerId.toString)
     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")
@@ -148,20 +178,25 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
       ApiKeys.ALTER_REPLICA_LOG_DIRS -> classOf[AlterReplicaLogDirsResponse],
       ApiKeys.DESCRIBE_LOG_DIRS -> classOf[DescribeLogDirsResponse],
       ApiKeys.CREATE_PARTITIONS -> classOf[CreatePartitionsResponse],
-      ApiKeys.ELECT_PREFERRED_LEADERS -> classOf[ElectPreferredLeadersResponse]
-  )
+      ApiKeys.ELECT_LEADERS -> classOf[ElectLeadersResponse],
+      ApiKeys.INCREMENTAL_ALTER_CONFIGS -> classOf[IncrementalAlterConfigsResponse],
+      ApiKeys.ALTER_PARTITION_REASSIGNMENTS -> classOf[AlterPartitionReassignmentsResponse],
+      ApiKeys.LIST_PARTITION_REASSIGNMENTS -> classOf[ListPartitionReassignmentsResponse],
+      ApiKeys.OFFSET_DELETE -> classOf[OffsetDeleteResponse]
+    )
 
   val requestKeyToError = 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) => resp.responses.asScala.find(_._1 == tp).get._2.error),
     ApiKeys.FETCH -> ((resp: requests.FetchResponse[Records]) => resp.responseData.asScala.find(_._1 == tp).get._2.error),
     ApiKeys.LIST_OFFSETS -> ((resp: requests.ListOffsetResponse) => resp.responseData.asScala.find(_._1 == tp).get._2.error),
-    ApiKeys.OFFSET_COMMIT -> ((resp: requests.OffsetCommitResponse) => resp.responseData.asScala.find(_._1 == tp).get._2),
+    ApiKeys.OFFSET_COMMIT -> ((resp: requests.OffsetCommitResponse) => Errors.forCode(
+      resp.data().topics().get(0).partitions().get(0).errorCode())),
     ApiKeys.OFFSET_FETCH -> ((resp: requests.OffsetFetchResponse) => resp.error),
     ApiKeys.FIND_COORDINATOR -> ((resp: FindCoordinatorResponse) => resp.error),
     ApiKeys.UPDATE_METADATA -> ((resp: requests.UpdateMetadataResponse) => resp.error),
     ApiKeys.JOIN_GROUP -> ((resp: JoinGroupResponse) => resp.error),
-    ApiKeys.SYNC_GROUP -> ((resp: SyncGroupResponse) => resp.error),
+    ApiKeys.SYNC_GROUP -> ((resp: SyncGroupResponse) => Errors.forCode(resp.data.errorCode())),
     ApiKeys.DESCRIBE_GROUPS -> ((resp: DescribeGroupsResponse) => {
       val errorCode = resp.data().groups().asScala.find(g => group.equals(g.groupId())).head.errorCode()
       Errors.forCode(errorCode)
@@ -169,11 +204,13 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
     ApiKeys.HEARTBEAT -> ((resp: HeartbeatResponse) => resp.error),
     ApiKeys.LEAVE_GROUP -> ((resp: LeaveGroupResponse) => resp.error),
     ApiKeys.DELETE_GROUPS -> ((resp: DeleteGroupsResponse) => resp.get(group)),
-    ApiKeys.LEADER_AND_ISR -> ((resp: requests.LeaderAndIsrResponse) => resp.responses.asScala.find(_._1 == tp).get._2),
-    ApiKeys.STOP_REPLICA -> ((resp: requests.StopReplicaResponse) => resp.responses.asScala.find(_._1 == tp).get._2),
+    ApiKeys.LEADER_AND_ISR -> ((resp: requests.LeaderAndIsrResponse) => Errors.forCode(
+      resp.partitions.asScala.find(p => p.topicName == tp.topic && p.partitionIndex == tp.partition).get.errorCode)),
+    ApiKeys.STOP_REPLICA -> ((resp: requests.StopReplicaResponse) => Errors.forCode(
+      resp.partitionErrors.asScala.find(pe => pe.topicName == tp.topic && pe.partitionIndex == tp.partition).get.errorCode)),
     ApiKeys.CONTROLLED_SHUTDOWN -> ((resp: requests.ControlledShutdownResponse) => resp.error),
-    ApiKeys.CREATE_TOPICS -> ((resp: CreateTopicsResponse) => Errors.forCode(resp.data().topics().find(createTopic).errorCode())),
-    ApiKeys.DELETE_TOPICS -> ((resp: requests.DeleteTopicsResponse) => resp.errors.asScala.find(_._1 == deleteTopic).get._2),
+    ApiKeys.CREATE_TOPICS -> ((resp: CreateTopicsResponse) => Errors.forCode(resp.data.topics.find(createTopic).errorCode())),
+    ApiKeys.DELETE_TOPICS -> ((resp: requests.DeleteTopicsResponse) => Errors.forCode(resp.data.responses.find(deleteTopic).errorCode())),
     ApiKeys.DELETE_RECORDS -> ((resp: requests.DeleteRecordsResponse) => resp.responses.get(deleteRecordsPartition).error),
     ApiKeys.OFFSET_FOR_LEADER_EPOCH -> ((resp: OffsetsForLeaderEpochResponse) => resp.responses.get(tp).error),
     ApiKeys.DESCRIBE_CONFIGS -> ((resp: DescribeConfigsResponse) =>
@@ -193,11 +230,27 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
     ApiKeys.DESCRIBE_LOG_DIRS -> ((resp: DescribeLogDirsResponse) =>
       if (resp.logDirInfos.size() > 0) resp.logDirInfos.asScala.head._2.error else Errors.CLUSTER_AUTHORIZATION_FAILED),
     ApiKeys.CREATE_PARTITIONS -> ((resp: CreatePartitionsResponse) => resp.errors.asScala.find(_._1 == topic).get._2.error),
-    ApiKeys.ELECT_PREFERRED_LEADERS -> ((resp: ElectPreferredLeadersResponse) =>
-      ElectPreferredLeadersRequest.fromResponseData(resp.data()).get(tp).error())
+    ApiKeys.ELECT_LEADERS -> ((resp: ElectLeadersResponse) => Errors.forCode(resp.data().errorCode())),
+    ApiKeys.INCREMENTAL_ALTER_CONFIGS -> ((resp: IncrementalAlterConfigsResponse) => {
+      val topicResourceError = IncrementalAlterConfigsResponse.fromResponseData(resp.data()).get(new ConfigResource(ConfigResource.Type.TOPIC, tp.topic))
+      if (topicResourceError == null)
+        IncrementalAlterConfigsResponse.fromResponseData(resp.data()).get(new ConfigResource(ConfigResource.Type.BROKER_LOGGER, brokerId.toString)).error
+      else
+        topicResourceError.error()
+    }),
+    ApiKeys.ALTER_PARTITION_REASSIGNMENTS -> ((resp: AlterPartitionReassignmentsResponse) => Errors.forCode(resp.data().errorCode())),
+    ApiKeys.LIST_PARTITION_REASSIGNMENTS -> ((resp: ListPartitionReassignmentsResponse) => Errors.forCode(resp.data().errorCode())),
+    ApiKeys.OFFSET_DELETE -> ((resp: OffsetDeleteResponse) => {
+      Errors.forCode(
+        resp.data
+          .topics().asScala.find(_.name() == topic).get
+          .partitions().asScala.find(_.partitionIndex() == part).get
+          .errorCode()
+      )
+    })
   )
 
-  val requestKeysToAcls = Map[ApiKeys, Map[Resource, Set[Acl]]](
+  val requestKeysToAcls = Map[ApiKeys, Map[ResourcePattern, Set[AccessControlEntry]]](
     ApiKeys.METADATA -> topicDescribeAcl,
     ApiKeys.PRODUCE -> (topicWriteAcl ++ transactionIdWriteAcl ++ clusterIdempotentWriteAcl),
     ApiKeys.FETCH -> topicReadAcl,
@@ -233,14 +286,18 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
     ApiKeys.ALTER_REPLICA_LOG_DIRS -> clusterAlterAcl,
     ApiKeys.DESCRIBE_LOG_DIRS -> clusterDescribeAcl,
     ApiKeys.CREATE_PARTITIONS -> topicAlterAcl,
-    ApiKeys.ELECT_PREFERRED_LEADERS -> clusterAlterAcl
+    ApiKeys.ELECT_LEADERS -> clusterAlterAcl,
+    ApiKeys.INCREMENTAL_ALTER_CONFIGS -> topicAlterConfigsAcl,
+    ApiKeys.ALTER_PARTITION_REASSIGNMENTS -> clusterAlterAcl,
+    ApiKeys.LIST_PARTITION_REASSIGNMENTS -> clusterDescribeAcl,
+    ApiKeys.OFFSET_DELETE -> groupReadAcl
   )
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     doSetup(createOffsetsTopic = false)
 
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, ClusterAction)), Resource.ClusterResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, CLUSTER_ACTION, ALLOW)), clusterResource)
 
     TestUtils.createOffsetsTopic(zkClient, servers)
 
@@ -287,7 +344,7 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
 
   private def offsetsForLeaderEpochRequest: OffsetsForLeaderEpochRequest = {
     val epochs = Map(tp -> new OffsetsForLeaderEpochRequest.PartitionData(Optional.of(27), 7))
-    new OffsetsForLeaderEpochRequest.Builder(ApiKeys.OFFSET_FOR_LEADER_EPOCH.latestVersion, epochs.asJava).build()
+    OffsetsForLeaderEpochRequest.Builder.forConsumer(epochs.asJava).build()
   }
 
   private def createOffsetFetchRequest = {
@@ -295,28 +352,61 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
   }
 
   private def createFindCoordinatorRequest = {
-    new FindCoordinatorRequest.Builder(FindCoordinatorRequest.CoordinatorType.GROUP, group).build()
+    new FindCoordinatorRequest.Builder(
+        new FindCoordinatorRequestData()
+          .setKeyType(FindCoordinatorRequest.CoordinatorType.GROUP.id)
+          .setKey(group)).build()
   }
 
   private def createUpdateMetadataRequest = {
-    val partitionState = Map(tp -> new UpdateMetadataRequest.PartitionState(
-      Int.MaxValue, brokerId, Int.MaxValue, List(brokerId).asJava, 2, Seq(brokerId).asJava, Seq.empty[Integer].asJava)).asJava
+    val partitionStates = Seq(new UpdateMetadataPartitionState()
+      .setTopicName(tp.topic)
+      .setPartitionIndex(tp.partition)
+      .setControllerEpoch(Int.MaxValue)
+      .setLeader(brokerId)
+      .setLeaderEpoch(Int.MaxValue)
+      .setIsr(List(brokerId).asJava)
+      .setZkVersion(2)
+      .setReplicas(Seq(brokerId).asJava)).asJava
     val securityProtocol = SecurityProtocol.PLAINTEXT
-    val brokers = Set(new requests.UpdateMetadataRequest.Broker(brokerId,
-      Seq(new requests.UpdateMetadataRequest.EndPoint("localhost", 0, securityProtocol,
-        ListenerName.forSecurityProtocol(securityProtocol))).asJava, null)).asJava
+    val brokers = Seq(new UpdateMetadataBroker()
+      .setId(brokerId)
+      .setEndpoints(Seq(new UpdateMetadataEndpoint()
+        .setHost("localhost")
+        .setPort(0)
+        .setSecurityProtocol(securityProtocol.id)
+        .setListener(ListenerName.forSecurityProtocol(securityProtocol).value)).asJava)).asJava
     val version = ApiKeys.UPDATE_METADATA.latestVersion
-    new requests.UpdateMetadataRequest.Builder(version, brokerId, Int.MaxValue, Long.MaxValue, partitionState, brokers).build()
+    new requests.UpdateMetadataRequest.Builder(version, brokerId, Int.MaxValue, Long.MaxValue, Long.MaxValue, partitionStates, brokers).build()
   }
 
   private def createJoinGroupRequest = {
-    new JoinGroupRequest.Builder(group, 10000, "", "consumer",
-      List( new JoinGroupRequest.ProtocolMetadata("consumer-range",ByteBuffer.wrap("test".getBytes()))).asJava)
-      .setRebalanceTimeout(60000).build()
+    val protocolSet = new JoinGroupRequestProtocolCollection(
+      Collections.singletonList(new JoinGroupRequestData.JoinGroupRequestProtocol()
+        .setName("consumer-range")
+        .setMetadata("test".getBytes())
+    ).iterator())
+
+    new JoinGroupRequest.Builder(
+      new JoinGroupRequestData()
+        .setGroupId(group)
+        .setSessionTimeoutMs(10000)
+        .setMemberId(JoinGroupRequest.UNKNOWN_MEMBER_ID)
+        .setGroupInstanceId(null)
+        .setProtocolType("consumer")
+        .setProtocols(protocolSet)
+        .setRebalanceTimeoutMs(60000)
+    ).build()
   }
 
   private def createSyncGroupRequest = {
-    new SyncGroupRequest.Builder(group, 1, "", Map[String, ByteBuffer]().asJava).build()
+    new SyncGroupRequest.Builder(
+      new SyncGroupRequestData()
+        .setGroupId(group)
+        .setGenerationId(1)
+        .setMemberId(JoinGroupRequest.UNKNOWN_MEMBER_ID)
+        .setAssignments(Collections.emptyList())
+    ).build()
   }
 
   private def createDescribeGroupsRequest = {
@@ -325,9 +415,23 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
 
   private def createOffsetCommitRequest = {
     new requests.OffsetCommitRequest.Builder(
-      group, Map(tp -> new requests.OffsetCommitRequest.PartitionData(0L, Optional.empty[Integer](), "metadata")).asJava).
-      setMemberId("").setGenerationId(1).
-      build()
+        new OffsetCommitRequestData()
+          .setGroupId(group)
+          .setMemberId(JoinGroupRequest.UNKNOWN_MEMBER_ID)
+          .setGenerationId(1)
+          .setTopics(Collections.singletonList(
+            new OffsetCommitRequestData.OffsetCommitRequestTopic()
+              .setName(topic)
+              .setPartitions(Collections.singletonList(
+                new OffsetCommitRequestData.OffsetCommitRequestPartition()
+                  .setPartitionIndex(part)
+                  .setCommittedOffset(0)
+                  .setCommittedLeaderEpoch(RecordBatch.NO_PARTITION_LEADER_EPOCH)
+                  .setCommitTimestamp(OffsetCommitRequest.DEFAULT_TIMESTAMP)
+                  .setCommittedMetadata("metadata")
+              )))
+          )
+    ).build()
   }
 
   private def createPartitionsRequest = {
@@ -336,30 +440,57 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
     ).build()
   }
 
-  private def heartbeatRequest = new HeartbeatRequest.Builder(group, 1, "").build()
+  private def heartbeatRequest = new HeartbeatRequest.Builder(
+    new HeartbeatRequestData()
+      .setGroupId(group)
+      .setGenerationId(1)
+      .setMemberId(JoinGroupRequest.UNKNOWN_MEMBER_ID)).build()
 
-  private def leaveGroupRequest = new LeaveGroupRequest.Builder(new LeaveGroupRequestData().setGroupId(group).setMemberId(JoinGroupRequest.UNKNOWN_MEMBER_ID)).build()
+  private def leaveGroupRequest = new LeaveGroupRequest.Builder(
+    group, Collections.singletonList(
+      new MemberIdentity()
+        .setMemberId(JoinGroupRequest.UNKNOWN_MEMBER_ID)
+    )).build()
 
-  private def deleteGroupsRequest = new DeleteGroupsRequest.Builder(Set(group).asJava).build()
+  private def deleteGroupsRequest = new DeleteGroupsRequest.Builder(
+    new DeleteGroupsRequestData()
+      .setGroupsNames(Collections.singletonList(group))
+  ).build()
 
   private def leaderAndIsrRequest = {
-    new requests.LeaderAndIsrRequest.Builder(ApiKeys.LEADER_AND_ISR.latestVersion, brokerId, Int.MaxValue, Long.MaxValue,
-      Map(tp -> new LeaderAndIsrRequest.PartitionState(Int.MaxValue, brokerId, Int.MaxValue, List(brokerId).asJava, 2, Seq(brokerId).asJava, false)).asJava,
+    new requests.LeaderAndIsrRequest.Builder(ApiKeys.LEADER_AND_ISR.latestVersion, brokerId, Int.MaxValue, Long.MaxValue, Long.MaxValue,
+      Seq(new LeaderAndIsrPartitionState()
+        .setTopicName(tp.topic)
+        .setPartitionIndex(tp.partition)
+        .setControllerEpoch(Int.MaxValue)
+        .setLeader(brokerId)
+        .setLeaderEpoch(Int.MaxValue)
+        .setIsr(List(brokerId).asJava)
+        .setZkVersion(2)
+        .setReplicas(Seq(brokerId).asJava)
+        .setIsNew(false)).asJava,
       Set(new Node(brokerId, "localhost", 0)).asJava).build()
   }
 
-  private def stopReplicaRequest = new StopReplicaRequest.Builder(ApiKeys.STOP_REPLICA.latestVersion, brokerId, Int.MaxValue, Long.MaxValue, true, Set(tp).asJava).build()
+  private def stopReplicaRequest = new StopReplicaRequest.Builder(ApiKeys.STOP_REPLICA.latestVersion, brokerId, Int.MaxValue, Long.MaxValue, Long.MaxValue, true, Set(tp).asJava).build()
 
-  private def controlledShutdownRequest = new requests.ControlledShutdownRequest.Builder(brokerId, Long.MaxValue,
-    ApiKeys.CONTROLLED_SHUTDOWN.latestVersion).build()
+  private def controlledShutdownRequest = new ControlledShutdownRequest.Builder(
+      new ControlledShutdownRequestData()
+        .setBrokerId(brokerId)
+        .setBrokerEpoch(Long.MaxValue),
+      ApiKeys.CONTROLLED_SHUTDOWN.latestVersion).build()
 
   private def createTopicsRequest =
     new CreateTopicsRequest.Builder(new CreateTopicsRequestData().setTopics(
-      new CreatableTopicSet(Collections.singleton(new CreatableTopic().
+      new CreatableTopicCollection(Collections.singleton(new CreatableTopic().
         setName(createTopic).setNumPartitions(1).
           setReplicationFactor(1.toShort)).iterator))).build()
 
-  private def deleteTopicsRequest = new DeleteTopicsRequest.Builder(Set(deleteTopic).asJava, 5000).build()
+  private def deleteTopicsRequest =
+    new DeleteTopicsRequest.Builder(
+      new DeleteTopicsRequestData()
+        .setTopicNames(Collections.singletonList(deleteTopic))
+        .setTimeoutMs(5000)).build()
 
   private def deleteRecordsRequest = new DeleteRecordsRequest.Builder(5000, Collections.singletonMap(deleteRecordsPartition, 0L)).build()
 
@@ -373,17 +504,30 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
           new AlterConfigsRequest.ConfigEntry(LogConfig.MaxMessageBytesProp, "1000000")
         ))), true).build()
 
+  private def incrementalAlterConfigsRequest = {
+    val data = new IncrementalAlterConfigsRequestData
+    val alterableConfig = new AlterableConfig
+    alterableConfig.setName(LogConfig.MaxMessageBytesProp).
+      setValue("1000000").setConfigOperation(AlterConfigOp.OpType.SET.id())
+    val alterableConfigSet = new AlterableConfigCollection
+    alterableConfigSet.add(alterableConfig)
+    data.resources().add(new AlterConfigsResource().
+      setResourceName(tp.topic).setResourceType(ConfigResource.Type.TOPIC.id()).
+      setConfigs(alterableConfigSet))
+    new IncrementalAlterConfigsRequest.Builder(data).build()
+  }
+
   private def describeAclsRequest = new DescribeAclsRequest.Builder(AclBindingFilter.ANY).build()
 
   private def createAclsRequest = new CreateAclsRequest.Builder(
     Collections.singletonList(new AclCreation(new AclBinding(
-      new ResourcePattern(AdminResourceType.TOPIC, "mytopic", LITERAL),
-      new AccessControlEntry(userPrincipal.toString, "*", AclOperation.WRITE, AclPermissionType.DENY))))).build()
+      new ResourcePattern(ResourceType.TOPIC, "mytopic", LITERAL),
+      new AccessControlEntry(userPrincipalStr, "*", AclOperation.WRITE, DENY))))).build()
 
   private def deleteAclsRequest = new DeleteAclsRequest.Builder(
     Collections.singletonList(new AclBindingFilter(
-      new ResourcePatternFilter(AdminResourceType.TOPIC, null, LITERAL),
-      new AccessControlEntryFilter(userPrincipal.toString, "*", AclOperation.ANY, AclPermissionType.DENY)))).build()
+      new ResourcePatternFilter(ResourceType.TOPIC, null, LITERAL),
+      new AccessControlEntryFilter(userPrincipalStr, "*", AclOperation.ANY, DENY)))).build()
 
   private def alterReplicaLogDirsRequest = new AlterReplicaLogDirsRequest.Builder(Collections.singletonMap(tp, logDir)).build()
 
@@ -393,11 +537,34 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
 
   private def addOffsetsToTxnRequest = new AddOffsetsToTxnRequest.Builder(transactionalId, 1, 1, group).build()
 
-  private def electPreferredLeadersRequest = new ElectPreferredLeadersRequest.Builder(
-    ElectPreferredLeadersRequest.toRequestData(Collections.singleton(tp), 10000)).build()
+  private def electLeadersRequest = new ElectLeadersRequest.Builder(
+    ElectionType.PREFERRED,
+    Collections.singleton(tp),
+    10000
+  ).build()
+
+  private def alterPartitionReassignmentsRequest = new AlterPartitionReassignmentsRequest.Builder(
+    new AlterPartitionReassignmentsRequestData().setTopics(
+      List(new AlterPartitionReassignmentsRequestData.ReassignableTopic()
+        .setName(topic)
+        .setPartitions(
+          List(new AlterPartitionReassignmentsRequestData.ReassignablePartition().setPartitionIndex(tp.partition())).asJava
+        )).asJava
+    )
+  ).build()
+
+  private def listPartitionReassignmentsRequest = new ListPartitionReassignmentsRequest.Builder(
+    new ListPartitionReassignmentsRequestData().setTopics(
+      List(new ListPartitionReassignmentsRequestData.ListPartitionReassignmentsTopics()
+        .setName(topic)
+        .setPartitionIndexes(
+          List(Integer.valueOf(tp.partition)).asJava
+        )).asJava
+    )
+  ).build()
 
   @Test
-  def testAuthorizationWithTopicExisting() {
+  def testAuthorizationWithTopicExisting(): Unit = {
     val requestKeyToRequest = mutable.LinkedHashMap[ApiKeys, AbstractRequest](
       ApiKeys.METADATA -> createMetadataRequest(allowAutoTopicCreation = true),
       ApiKeys.PRODUCE -> createProduceRequest,
@@ -428,9 +595,12 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
       ApiKeys.CREATE_PARTITIONS -> createPartitionsRequest,
       ApiKeys.ADD_PARTITIONS_TO_TXN -> addPartitionsToTxnRequest,
       ApiKeys.ADD_OFFSETS_TO_TXN -> addOffsetsToTxnRequest,
+      ApiKeys.ELECT_LEADERS -> electLeadersRequest,
+      ApiKeys.INCREMENTAL_ALTER_CONFIGS -> incrementalAlterConfigsRequest,
+      ApiKeys.ALTER_PARTITION_REASSIGNMENTS -> alterPartitionReassignmentsRequest,
+      ApiKeys.LIST_PARTITION_REASSIGNMENTS -> listPartitionReassignmentsRequest,
       // Check StopReplica last since some APIs depend on replica availability
-      ApiKeys.STOP_REPLICA -> stopReplicaRequest,
-      ApiKeys.ELECT_PREFERRED_LEADERS -> electPreferredLeadersRequest
+      ApiKeys.STOP_REPLICA -> stopReplicaRequest
     )
 
     for ((key, request) <- requestKeyToRequest) {
@@ -457,7 +627,7 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
    * even if the topic doesn't exist, request APIs should not leak the topic name
    */
   @Test
-  def testAuthorizationWithTopicNotExisting() {
+  def testAuthorizationWithTopicNotExisting(): Unit = {
     adminZkClient.deleteTopic(topic)
     TestUtils.verifyTopicDeletion(zkClient, topic, 1, servers)
     adminZkClient.deleteTopic(deleteTopic)
@@ -477,7 +647,7 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
       ApiKeys.CREATE_PARTITIONS -> createPartitionsRequest,
       ApiKeys.DELETE_GROUPS -> deleteGroupsRequest,
       ApiKeys.OFFSET_FOR_LEADER_EPOCH -> offsetsForLeaderEpochRequest,
-      ApiKeys.ELECT_PREFERRED_LEADERS -> electPreferredLeadersRequest
+      ApiKeys.ELECT_LEADERS -> electLeadersRequest
     )
 
     for ((key, request) <- requestKeyToRequest) {
@@ -501,9 +671,9 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
   }
 
   @Test
-  def testCreateTopicAuthorizationWithClusterCreate() {
+  def testCreateTopicAuthorizationWithClusterCreate(): Unit = {
     removeAllAcls()
-    val resources = Set[ResourceType](Topic)
+    val resources = Set[ResourceType](TOPIC)
 
     sendRequestAndVerifyResponseError(ApiKeys.CREATE_TOPICS, createTopicsRequest, resources, isAuthorized = false)
 
@@ -513,20 +683,42 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
   }
 
   @Test
-  def testFetchFollowerRequest() {
+  def testFetchFollowerRequest(): Unit = {
     val key = ApiKeys.FETCH
     val request = createFetchFollowerRequest
 
     removeAllAcls()
-    val resources = Set(topicResource.resourceType, Resource.ClusterResource.resourceType)
+    val resources = Set(topicResource.resourceType, clusterResource.resourceType)
     sendRequestAndVerifyResponseError(key, request, resources, isAuthorized = false)
 
     val readAcls = topicReadAcl(topicResource)
     addAndVerifyAcls(readAcls, topicResource)
     sendRequestAndVerifyResponseError(key, request, resources, isAuthorized = false)
 
-    val clusterAcls = clusterAcl(Resource.ClusterResource)
-    addAndVerifyAcls(clusterAcls, Resource.ClusterResource)
+    val clusterAcls = clusterAcl(clusterResource)
+    addAndVerifyAcls(clusterAcls, clusterResource)
+    sendRequestAndVerifyResponseError(key, request, resources, isAuthorized = true)
+  }
+
+  @Test
+  def testIncrementalAlterConfigsRequestRequiresClusterPermissionForBrokerLogger(): Unit = {
+    val data = new IncrementalAlterConfigsRequestData
+    val alterableConfig = new AlterableConfig().setName("kafka.controller.KafkaController").
+      setValue(LogLevelConfig.DEBUG_LOG_LEVEL).setConfigOperation(AlterConfigOp.OpType.DELETE.id())
+    val alterableConfigSet = new AlterableConfigCollection
+    alterableConfigSet.add(alterableConfig)
+    data.resources().add(new AlterConfigsResource().
+      setResourceName(brokerId.toString).setResourceType(ConfigResource.Type.BROKER_LOGGER.id()).
+      setConfigs(alterableConfigSet))
+    val key = ApiKeys.INCREMENTAL_ALTER_CONFIGS
+    val request = new IncrementalAlterConfigsRequest.Builder(data).build()
+
+    removeAllAcls()
+    val resources = Set(topicResource.resourceType, clusterResource.resourceType)
+    sendRequestAndVerifyResponseError(key, request, resources, isAuthorized = false)
+
+    val clusterAcls = clusterAlterConfigsAcl(clusterResource)
+    addAndVerifyAcls(clusterAcls, clusterResource)
     sendRequestAndVerifyResponseError(key, request, resources, isAuthorized = true)
   }
 
@@ -537,18 +729,18 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
 
     removeAllAcls()
 
-    val resources = Set(topicResource.resourceType, Resource.ClusterResource.resourceType)
+    val resources = Set(topicResource.resourceType, clusterResource.resourceType)
     sendRequestAndVerifyResponseError(key, request, resources, isAuthorized = false)
 
     // Although the OffsetsForLeaderEpoch API now accepts topic describe, we should continue
     // allowing cluster action for backwards compatibility
-    val clusterAcls = clusterAcl(Resource.ClusterResource)
-    addAndVerifyAcls(clusterAcls, Resource.ClusterResource)
+    val clusterAcls = clusterAcl(clusterResource)
+    addAndVerifyAcls(clusterAcls, clusterResource)
     sendRequestAndVerifyResponseError(key, request, resources, isAuthorized = true)
   }
 
   @Test
-  def testProduceWithNoTopicAccess() {
+  def testProduceWithNoTopicAccess(): Unit = {
     try {
       val producer = createProducer()
       sendRecords(producer, numRecords, tp)
@@ -559,8 +751,8 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
   }
 
   @Test
-  def testProduceWithTopicDescribe() {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Describe)), topicResource)
+  def testProduceWithTopicDescribe(): Unit = {
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, DESCRIBE, ALLOW)), topicResource)
     try {
       val producer = createProducer()
       sendRecords(producer, numRecords, tp)
@@ -572,8 +764,8 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
   }
 
   @Test
-  def testProduceWithTopicRead() {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Read)), topicResource)
+  def testProduceWithTopicRead(): Unit = {
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)), topicResource)
     try {
       val producer = createProducer()
       sendRecords(producer, numRecords, tp)
@@ -585,26 +777,26 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
   }
 
   @Test
-  def testProduceWithTopicWrite() {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Write)), topicResource)
+  def testProduceWithTopicWrite(): Unit = {
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, WRITE, ALLOW)), topicResource)
     val producer = createProducer()
     sendRecords(producer, numRecords, tp)
   }
 
   @Test
-  def testCreatePermissionOnTopicToWriteToNonExistentTopic() {
-    testCreatePermissionNeededToWriteToNonExistentTopic(Topic)
+  def testCreatePermissionOnTopicToWriteToNonExistentTopic(): Unit = {
+    testCreatePermissionNeededToWriteToNonExistentTopic(TOPIC)
   }
 
   @Test
-  def testCreatePermissionOnClusterToWriteToNonExistentTopic() {
-    testCreatePermissionNeededToWriteToNonExistentTopic(Cluster)
+  def testCreatePermissionOnClusterToWriteToNonExistentTopic(): Unit = {
+    testCreatePermissionNeededToWriteToNonExistentTopic(CLUSTER)
   }
 
-  private def testCreatePermissionNeededToWriteToNonExistentTopic(resType: ResourceType) {
+  private def testCreatePermissionNeededToWriteToNonExistentTopic(resType: ResourceType): Unit = {
     val topicPartition = new TopicPartition(createTopic, 0)
-    val newTopicResource = Resource(Topic, createTopic, LITERAL)
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Write)), newTopicResource)
+    val newTopicResource = new ResourcePattern(TOPIC, createTopic, LITERAL)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, WRITE, ALLOW)), newTopicResource)
     val producer = createProducer()
     try {
       sendRecords(producer, numRecords, topicPartition)
@@ -614,15 +806,15 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
         assertEquals(Collections.singleton(createTopic), e.unauthorizedTopics())
     }
 
-    val resource = if (resType == Topic) newTopicResource else Resource.ClusterResource
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Create)), resource)
+    val resource = if (resType == Topic) newTopicResource else clusterResource
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, CREATE, ALLOW)), resource)
 
     sendRecords(producer, numRecords, topicPartition)
   }
 
   @Test(expected = classOf[TopicAuthorizationException])
   def testConsumeUsingAssignWithNoAccess(): Unit = {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Write)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, WRITE, ALLOW)), topicResource)
     val producer = createProducer()
     sendRecords(producer, 1, tp)
     removeAllAcls()
@@ -634,12 +826,12 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
 
   @Test
   def testSimpleConsumeWithOffsetLookupAndNoGroupAccess(): Unit = {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Write)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, WRITE, ALLOW)), topicResource)
     val producer = createProducer()
     sendRecords(producer, 1, tp)
     removeAllAcls()
 
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Read)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)), topicResource)
     try {
       // 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)
@@ -654,12 +846,12 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
 
   @Test
   def testSimpleConsumeWithExplicitSeekAndNoGroupAccess(): Unit = {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Write)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, WRITE, ALLOW)), topicResource)
     val producer = createProducer()
     sendRecords(producer, 1, tp)
     removeAllAcls()
 
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Read)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)), topicResource)
 
     // in this case, we do an explicit seek, so there should be no need to query the coordinator at all
     val consumer = createConsumer()
@@ -669,13 +861,13 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
   }
 
   @Test(expected = classOf[KafkaException])
-  def testConsumeWithoutTopicDescribeAccess() {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Write)), topicResource)
+  def testConsumeWithoutTopicDescribeAccess(): Unit = {
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, WRITE, ALLOW)), topicResource)
     val producer = createProducer()
     sendRecords(producer, 1, tp)
     removeAllAcls()
 
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Read)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)), groupResource)
     val consumer = createConsumer()
     consumer.assign(List(tp).asJava)
 
@@ -685,14 +877,14 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
   }
 
   @Test
-  def testConsumeWithTopicDescribe() {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Write)), topicResource)
+  def testConsumeWithTopicDescribe(): Unit = {
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, WRITE, ALLOW)), topicResource)
     val producer = createProducer()
     sendRecords(producer, 1, tp)
     removeAllAcls()
 
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Describe)), topicResource)
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Read)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, DESCRIBE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)), groupResource)
     try {
       val consumer = createConsumer()
       consumer.assign(List(tp).asJava)
@@ -704,14 +896,14 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
   }
 
   @Test
-  def testConsumeWithTopicWrite() {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Write)), topicResource)
+  def testConsumeWithTopicWrite(): Unit = {
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, WRITE, ALLOW)), topicResource)
     val producer = createProducer()
     sendRecords(producer, 1, tp)
     removeAllAcls()
 
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Write)), topicResource)
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Read)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, WRITE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)), groupResource)
     try {
       val consumer = createConsumer()
       consumer.assign(List(tp).asJava)
@@ -724,14 +916,14 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
   }
 
   @Test
-  def testConsumeWithTopicAndGroupRead() {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Write)), topicResource)
+  def testConsumeWithTopicAndGroupRead(): Unit = {
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, WRITE, ALLOW)), topicResource)
     val producer = createProducer()
     sendRecords(producer, 1, tp)
     removeAllAcls()
 
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Read)), topicResource)
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Read)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)), groupResource)
 
     val consumer = createConsumer()
     consumer.assign(List(tp).asJava)
@@ -739,13 +931,13 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
   }
 
   @Test
-  def testPatternSubscriptionWithNoTopicAccess() {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Write)), topicResource)
+  def testPatternSubscriptionWithNoTopicAccess(): Unit = {
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, WRITE, ALLOW)), topicResource)
     val producer = createProducer()
     sendRecords(producer, 1, tp)
     removeAllAcls()
 
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Read)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)), groupResource)
 
     val consumer = createConsumer()
     consumer.subscribe(Pattern.compile(topicPattern), new NoOpConsumerRebalanceListener)
@@ -754,14 +946,14 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
   }
 
   @Test
-  def testPatternSubscriptionWithTopicDescribeOnlyAndGroupRead() {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Write)), topicResource)
+  def testPatternSubscriptionWithTopicDescribeOnlyAndGroupRead(): Unit = {
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, WRITE, ALLOW)), topicResource)
     val producer = createProducer()
     sendRecords(producer, 1, tp)
     removeAllAcls()
 
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Describe)), topicResource)
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Read)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, DESCRIBE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)), groupResource)
     val consumer = createConsumer()
     consumer.subscribe(Pattern.compile(topicPattern))
     try {
@@ -773,27 +965,27 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
   }
 
   @Test
-  def testPatternSubscriptionWithTopicAndGroupRead() {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Write)), topicResource)
+  def testPatternSubscriptionWithTopicAndGroupRead(): Unit = {
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, WRITE, ALLOW)), topicResource)
     val producer = createProducer()
     sendRecords(producer, 1, tp)
 
     // create an unmatched topic
     val unmatchedTopic = "unmatched"
     createTopic(unmatchedTopic)
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Write)),  Resource(Topic, unmatchedTopic, LITERAL))
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, WRITE, ALLOW)),  new ResourcePattern(TOPIC, unmatchedTopic, LITERAL))
     sendRecords(producer, 1, new TopicPartition(unmatchedTopic, part))
     removeAllAcls()
 
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Read)), topicResource)
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Read)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, 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 Acl(userPrincipal, Allow, Acl.WildCardHost, Read)),  Resource(Topic,
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)),  new ResourcePattern(TOPIC,
       GROUP_METADATA_TOPIC_NAME, LITERAL))
     consumer.subscribe(Pattern.compile(GROUP_METADATA_TOPIC_NAME))
     consumer.poll(0)
@@ -802,14 +994,14 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
   }
 
   @Test
-  def testPatternSubscriptionMatchingInternalTopic() {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Write)), topicResource)
+  def testPatternSubscriptionMatchingInternalTopic(): Unit = {
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, WRITE, ALLOW)), topicResource)
     val producer = createProducer()
     sendRecords(producer, 1, tp)
     removeAllAcls()
 
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Read)), topicResource)
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Read)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)), groupResource)
 
     consumerConfig.put(ConsumerConfig.EXCLUDE_INTERNAL_TOPICS_CONFIG, "false")
     val consumer = createConsumer()
@@ -819,7 +1011,7 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
     assertEquals(Set(topic).asJava, consumer.subscription)
 
     // now authorize the user for the internal topic and verify that we can subscribe
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Read)), Resource(Topic,
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)), new ResourcePattern(TOPIC,
       GROUP_METADATA_TOPIC_NAME, LITERAL))
     consumer.subscribe(Pattern.compile(GROUP_METADATA_TOPIC_NAME))
     consumer.poll(0)
@@ -827,16 +1019,16 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
   }
 
   @Test
-  def testPatternSubscriptionMatchingInternalTopicWithDescribeOnlyPermission() {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Write)), topicResource)
+  def testPatternSubscriptionMatchingInternalTopicWithDescribeOnlyPermission(): Unit = {
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, WRITE, ALLOW)), topicResource)
     val producer = createProducer()
     sendRecords(producer, 1, tp)
     removeAllAcls()
 
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Read)), topicResource)
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Read)), groupResource)
-    val internalTopicResource = Resource(Topic, GROUP_METADATA_TOPIC_NAME, LITERAL)
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Describe)), internalTopicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)), groupResource)
+    val internalTopicResource = new ResourcePattern(TOPIC, GROUP_METADATA_TOPIC_NAME, LITERAL)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, DESCRIBE, ALLOW)), internalTopicResource)
 
     consumerConfig.put(ConsumerConfig.EXCLUDE_INTERNAL_TOPICS_CONFIG, "false")
     val consumer = createConsumer()
@@ -852,14 +1044,14 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
   }
 
   @Test
-  def testPatternSubscriptionNotMatchingInternalTopic() {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Write)), topicResource)
+  def testPatternSubscriptionNotMatchingInternalTopic(): Unit = {
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, WRITE, ALLOW)), topicResource)
     val producer = createProducer()
     sendRecords(producer, 1, tp)
     removeAllAcls()
 
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Read)), topicResource)
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Read)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)), groupResource)
 
     consumerConfig.put(ConsumerConfig.EXCLUDE_INTERNAL_TOPICS_CONFIG, "false")
     val consumer = createConsumer()
@@ -867,26 +1059,26 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
       consumer.subscribe(Pattern.compile(topicPattern))
       consumeRecords(consumer)
     } finally consumer.close()
-}
+  }
 
   @Test
-  def testCreatePermissionOnTopicToReadFromNonExistentTopic() {
+  def testCreatePermissionOnTopicToReadFromNonExistentTopic(): Unit = {
     testCreatePermissionNeededToReadFromNonExistentTopic("newTopic",
-      Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Create)),
-      Topic)
+      Set(new AccessControlEntry(userPrincipalStr, WildcardHost, CREATE, ALLOW)),
+      TOPIC)
   }
 
   @Test
-  def testCreatePermissionOnClusterToReadFromNonExistentTopic() {
+  def testCreatePermissionOnClusterToReadFromNonExistentTopic(): Unit = {
     testCreatePermissionNeededToReadFromNonExistentTopic("newTopic",
-      Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Create)),
-      Cluster)
+      Set(new AccessControlEntry(userPrincipalStr, WildcardHost, CREATE, ALLOW)),
+      CLUSTER)
   }
 
-  private def testCreatePermissionNeededToReadFromNonExistentTopic(newTopic: String, acls: Set[Acl], resType: ResourceType) {
+  private def testCreatePermissionNeededToReadFromNonExistentTopic(newTopic: String, acls: Set[AccessControlEntry], resType: ResourceType): Unit = {
     val topicPartition = new TopicPartition(newTopic, 0)
-    val newTopicResource = Resource(Topic, newTopic, LITERAL)
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Read)), newTopicResource)
+    val newTopicResource = new ResourcePattern(TOPIC, newTopic, LITERAL)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)), newTopicResource)
     addAndVerifyAcls(groupReadAcl(groupResource), groupResource)
     val consumer = createConsumer()
     consumer.assign(List(topicPartition).asJava)
@@ -895,7 +1087,7 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
     }.unauthorizedTopics
     assertEquals(Collections.singleton(newTopic), unauthorizedTopics)
 
-    val resource = if (resType == Topic) newTopicResource else Resource.ClusterResource
+    val resource = if (resType == TOPIC) newTopicResource else clusterResource
     addAndVerifyAcls(acls, resource)
 
     TestUtils.waitUntilTrue(() => {
@@ -905,7 +1097,7 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
   }
 
   @Test
-  def testCreatePermissionMetadataRequestAutoCreate() {
+  def testCreatePermissionMetadataRequestAutoCreate(): Unit = {
     val readAcls = topicReadAcl.get(topicResource).get
     addAndVerifyAcls(readAcls, topicResource)
     assertTrue(zkClient.topicExists(topicResource.name))
@@ -930,95 +1122,95 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
   }
 
   @Test(expected = classOf[AuthorizationException])
-  def testCommitWithNoAccess() {
+  def testCommitWithNoAccess(): Unit = {
     val consumer = createConsumer()
     consumer.commitSync(Map(tp -> new OffsetAndMetadata(5)).asJava)
   }
 
   @Test(expected = classOf[KafkaException])
-  def testCommitWithNoTopicAccess() {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Read)), groupResource)
+  def testCommitWithNoTopicAccess(): Unit = {
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)), groupResource)
     val consumer = createConsumer()
     consumer.commitSync(Map(tp -> new OffsetAndMetadata(5)).asJava)
   }
 
   @Test(expected = classOf[TopicAuthorizationException])
-  def testCommitWithTopicWrite() {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Read)), groupResource)
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Write)), topicResource)
+  def testCommitWithTopicWrite(): Unit = {
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, WRITE, ALLOW)), topicResource)
     val consumer = createConsumer()
     consumer.commitSync(Map(tp -> new OffsetAndMetadata(5)).asJava)
   }
 
   @Test(expected = classOf[TopicAuthorizationException])
-  def testCommitWithTopicDescribe() {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Read)), groupResource)
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Describe)), topicResource)
+  def testCommitWithTopicDescribe(): Unit = {
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, DESCRIBE, ALLOW)), topicResource)
     val consumer = createConsumer()
     consumer.commitSync(Map(tp -> new OffsetAndMetadata(5)).asJava)
   }
 
   @Test(expected = classOf[GroupAuthorizationException])
-  def testCommitWithNoGroupAccess() {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Read)), topicResource)
+  def testCommitWithNoGroupAccess(): Unit = {
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)), topicResource)
     val consumer = createConsumer()
     consumer.commitSync(Map(tp -> new OffsetAndMetadata(5)).asJava)
   }
 
   @Test
-  def testCommitWithTopicAndGroupRead() {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Read)), groupResource)
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Read)), topicResource)
+  def testCommitWithTopicAndGroupRead(): Unit = {
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)), topicResource)
     val consumer = createConsumer()
     consumer.commitSync(Map(tp -> new OffsetAndMetadata(5)).asJava)
   }
 
   @Test(expected = classOf[AuthorizationException])
-  def testOffsetFetchWithNoAccess() {
+  def testOffsetFetchWithNoAccess(): Unit = {
     val consumer = createConsumer()
     consumer.assign(List(tp).asJava)
     consumer.position(tp)
   }
 
   @Test(expected = classOf[GroupAuthorizationException])
-  def testOffsetFetchWithNoGroupAccess() {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Read)), topicResource)
+  def testOffsetFetchWithNoGroupAccess(): Unit = {
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)), topicResource)
     val consumer = createConsumer()
     consumer.assign(List(tp).asJava)
     consumer.position(tp)
   }
 
   @Test(expected = classOf[KafkaException])
-  def testOffsetFetchWithNoTopicAccess() {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Read)), groupResource)
+  def testOffsetFetchWithNoTopicAccess(): Unit = {
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)), groupResource)
     val consumer = createConsumer()
     consumer.assign(List(tp).asJava)
     consumer.position(tp)
   }
 
   @Test
-  def testFetchAllOffsetsTopicAuthorization() {
+  def testFetchAllOffsetsTopicAuthorization(): Unit = {
     val offset = 15L
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Read)), groupResource)
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Read)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)), topicResource)
     val consumer = createConsumer()
     consumer.assign(List(tp).asJava)
     consumer.commitSync(Map(tp -> new OffsetAndMetadata(offset)).asJava)
 
     removeAllAcls()
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Read)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, 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
 
     // without describe permission on the topic, we shouldn't be able to fetch offsets
-    val offsetFetchRequest = requests.OffsetFetchRequest.forAllPartitions(group)
+    val offsetFetchRequest = requests.OffsetFetchRequest.Builder.allTopicPartitions(group).build()
     var offsetFetchResponse = sendOffsetFetchRequest(offsetFetchRequest, anySocketServer)
     assertEquals(Errors.NONE, offsetFetchResponse.error)
     assertTrue(offsetFetchResponse.responseData.isEmpty)
 
     // now add describe permission on the topic and verify that the offset can be fetched
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Describe)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, DESCRIBE, ALLOW)), topicResource)
     offsetFetchResponse = sendOffsetFetchRequest(offsetFetchRequest, anySocketServer)
     assertEquals(Errors.NONE, offsetFetchResponse.error)
     assertTrue(offsetFetchResponse.responseData.containsKey(tp))
@@ -1026,87 +1218,87 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
   }
 
   @Test
-  def testOffsetFetchTopicDescribe() {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Describe)), groupResource)
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Describe)), topicResource)
+  def testOffsetFetchTopicDescribe(): Unit = {
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, DESCRIBE, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, DESCRIBE, ALLOW)), topicResource)
     val consumer = createConsumer()
     consumer.assign(List(tp).asJava)
     consumer.position(tp)
   }
 
   @Test
-  def testOffsetFetchWithTopicAndGroupRead() {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Read)), groupResource)
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Read)), topicResource)
+  def testOffsetFetchWithTopicAndGroupRead(): Unit = {
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)), topicResource)
     val consumer = createConsumer()
     consumer.assign(List(tp).asJava)
     consumer.position(tp)
   }
 
   @Test(expected = classOf[TopicAuthorizationException])
-  def testMetadataWithNoTopicAccess() {
+  def testMetadataWithNoTopicAccess(): Unit = {
     val consumer = createConsumer()
     consumer.partitionsFor(topic)
   }
 
   @Test
-  def testMetadataWithTopicDescribe() {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Describe)), topicResource)
+  def testMetadataWithTopicDescribe(): Unit = {
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, DESCRIBE, ALLOW)), topicResource)
     val consumer = createConsumer()
     consumer.partitionsFor(topic)
   }
 
   @Test(expected = classOf[TopicAuthorizationException])
-  def testListOffsetsWithNoTopicAccess() {
+  def testListOffsetsWithNoTopicAccess(): Unit = {
     val consumer = createConsumer()
     consumer.endOffsets(Set(tp).asJava)
   }
 
   @Test
-  def testListOffsetsWithTopicDescribe() {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Describe)), topicResource)
+  def testListOffsetsWithTopicDescribe(): Unit = {
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, DESCRIBE, ALLOW)), topicResource)
     val consumer = createConsumer()
     consumer.endOffsets(Set(tp).asJava)
   }
 
   @Test
-  def testDescribeGroupApiWithNoGroupAcl() {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Describe)), topicResource)
+  def testDescribeGroupApiWithNoGroupAcl(): Unit = {
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, DESCRIBE, ALLOW)), topicResource)
     val result = createAdminClient().describeConsumerGroups(Seq(group).asJava)
     TestUtils.assertFutureExceptionTypeEquals(result.describedGroups().get(group), classOf[GroupAuthorizationException])
   }
 
   @Test
-  def testDescribeGroupApiWithGroupDescribe() {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Describe)), groupResource)
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Describe)), topicResource)
+  def testDescribeGroupApiWithGroupDescribe(): Unit = {
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, DESCRIBE, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, DESCRIBE, ALLOW)), topicResource)
     createAdminClient().describeConsumerGroups(Seq(group).asJava).describedGroups().get(group).get()
   }
 
   @Test
-  def testDescribeGroupCliWithGroupDescribe() {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Describe)), groupResource)
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Describe)), topicResource)
+  def testDescribeGroupCliWithGroupDescribe(): Unit = {
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, DESCRIBE, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, DESCRIBE, ALLOW)), topicResource)
 
     val cgcArgs = Array("--bootstrap-server", brokerList, "--describe", "--group", group)
     val opts = new ConsumerGroupCommandOptions(cgcArgs)
     val consumerGroupService = new ConsumerGroupService(opts)
-    consumerGroupService.describeGroup()
+    consumerGroupService.describeGroups()
     consumerGroupService.close()
   }
 
   @Test
-  def testListGroupApiWithAndWithoutListGroupAcls() {
+  def testListGroupApiWithAndWithoutListGroupAcls(): Unit = {
     // write some record to the topic
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Write)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, 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 Acl(userPrincipal, Allow, Acl.WildCardHost, Read)), groupResource)
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Read)), Resource(Group, group2, LITERAL))
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Read)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)), new ResourcePattern(GROUP, group2, LITERAL))
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)), topicResource)
     val consumer = createConsumer()
     consumer.subscribe(Collections.singleton(topic))
     consumeRecords(consumer)
@@ -1119,13 +1311,13 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
 
     // first use cluster describe permission
     removeAllAcls()
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Describe)), Resource.ClusterResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, 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
     removeAllAcls()
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Read)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)), groupResource)
     // it should list only one group now
     val groupList = adminClient.listConsumerGroups().all().get().asScala.toList
     assertEquals(1, groupList.length)
@@ -1140,10 +1332,10 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
   }
 
   @Test
-  def testDeleteGroupApiWithDeleteGroupAcl() {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Read)), groupResource)
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Read)), topicResource)
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Delete)), groupResource)
+  def testDeleteGroupApiWithDeleteGroupAcl(): Unit = {
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, DELETE, ALLOW)), groupResource)
     val consumer = createConsumer()
     consumer.assign(List(tp).asJava)
     consumer.commitSync(Map(tp -> new OffsetAndMetadata(5, "")).asJava)
@@ -1151,9 +1343,9 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
   }
 
   @Test
-  def testDeleteGroupApiWithNoDeleteGroupAcl() {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Read)), groupResource)
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Read)), topicResource)
+  def testDeleteGroupApiWithNoDeleteGroupAcl(): Unit = {
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)), topicResource)
     val consumer = createConsumer()
     consumer.assign(List(tp).asJava)
     consumer.commitSync(Map(tp -> new OffsetAndMetadata(5, "")).asJava)
@@ -1162,41 +1354,91 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
   }
 
   @Test
-  def testDeleteGroupApiWithNoDeleteGroupAcl2() {
+  def testDeleteGroupApiWithNoDeleteGroupAcl2(): Unit = {
     val result = createAdminClient().deleteConsumerGroups(Seq(group).asJava)
     TestUtils.assertFutureExceptionTypeEquals(result.deletedGroups().get(group), classOf[GroupAuthorizationException])
   }
 
   @Test
-  def testUnauthorizedDeleteTopicsWithoutDescribe() {
+  def testDeleteGroupOffsetsWithAcl(): Unit = {
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, DELETE, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)), topicResource)
+    val consumer = createConsumer()
+    consumer.assign(List(tp).asJava)
+    consumer.commitSync(Map(tp -> new OffsetAndMetadata(5, "")).asJava)
+    consumer.close()
+    val result = createAdminClient().deleteConsumerGroupOffsets(group, Set(tp).asJava)
+    assertNull(result.partitionResult(tp).get())
+  }
+
+  @Test
+  def testDeleteGroupOffsetsWithoutDeleteAcl(): Unit = {
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)), topicResource)
+    val consumer = createConsumer()
+    consumer.assign(List(tp).asJava)
+    consumer.commitSync(Map(tp -> new OffsetAndMetadata(5, "")).asJava)
+    consumer.close()
+    val result = createAdminClient().deleteConsumerGroupOffsets(group, Set(tp).asJava)
+    TestUtils.assertFutureExceptionTypeEquals(result.all(), classOf[GroupAuthorizationException])
+  }
+
+  @Test
+  def testDeleteGroupOffsetsWithDeleteAclWithoutTopicAcl(): Unit = {
+    // Create the consumer group
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)), topicResource)
+    val consumer = createConsumer()
+    consumer.assign(List(tp).asJava)
+    consumer.commitSync(Map(tp -> new OffsetAndMetadata(5, "")).asJava)
+    consumer.close()
+
+    // Remove the topic ACL & Check that it does not work without it
+    removeAllAcls()
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, DELETE, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, READ, ALLOW)), groupResource)
+    val result = createAdminClient().deleteConsumerGroupOffsets(group, Set(tp).asJava)
+    TestUtils.assertFutureExceptionTypeEquals(result.all(), classOf[TopicAuthorizationException])
+    TestUtils.assertFutureExceptionTypeEquals(result.partitionResult(tp), classOf[TopicAuthorizationException])
+  }
+
+  @Test
+  def testDeleteGroupOffsetsWithNoAcl(): Unit = {
+    val result = createAdminClient().deleteConsumerGroupOffsets(group, Set(tp).asJava)
+    TestUtils.assertFutureExceptionTypeEquals(result.all(), classOf[GroupAuthorizationException])
+  }
+
+  @Test
+  def testUnauthorizedDeleteTopicsWithoutDescribe(): Unit = {
     val response = connectAndSend(deleteTopicsRequest, ApiKeys.DELETE_TOPICS)
     val version = ApiKeys.DELETE_TOPICS.latestVersion
     val deleteResponse = DeleteTopicsResponse.parse(response, version)
-    assertEquals(Errors.TOPIC_AUTHORIZATION_FAILED, deleteResponse.errors.asScala.head._2)
+    assertEquals(Errors.TOPIC_AUTHORIZATION_FAILED.code, deleteResponse.data.responses.find(deleteTopic).errorCode)
   }
 
   @Test
-  def testUnauthorizedDeleteTopicsWithDescribe() {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Describe)), deleteTopicResource)
+  def testUnauthorizedDeleteTopicsWithDescribe(): Unit = {
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, DESCRIBE, ALLOW)), deleteTopicResource)
     val response = connectAndSend(deleteTopicsRequest, ApiKeys.DELETE_TOPICS)
     val version = ApiKeys.DELETE_TOPICS.latestVersion
     val deleteResponse = DeleteTopicsResponse.parse(response, version)
 
-    assertEquals(Errors.TOPIC_AUTHORIZATION_FAILED, deleteResponse.errors.asScala.head._2)
+    assertEquals(Errors.TOPIC_AUTHORIZATION_FAILED.code, deleteResponse.data.responses.find(deleteTopic).errorCode)
   }
 
   @Test
-  def testDeleteTopicsWithWildCardAuth() {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Delete)), Resource(Topic, "*", LITERAL))
+  def testDeleteTopicsWithWildCardAuth(): Unit = {
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, DELETE, ALLOW)), new ResourcePattern(TOPIC, "*", LITERAL))
     val response = connectAndSend(deleteTopicsRequest, ApiKeys.DELETE_TOPICS)
     val version = ApiKeys.DELETE_TOPICS.latestVersion
     val deleteResponse = DeleteTopicsResponse.parse(response, version)
 
-    assertEquals(Errors.NONE, deleteResponse.errors.asScala.head._2)
+    assertEquals(Errors.NONE.code, deleteResponse.data.responses.find(deleteTopic).errorCode)
   }
 
   @Test
-  def testUnauthorizedDeleteRecordsWithoutDescribe() {
+  def testUnauthorizedDeleteRecordsWithoutDescribe(): Unit = {
     val response = connectAndSend(deleteRecordsRequest, ApiKeys.DELETE_RECORDS)
     val version = ApiKeys.DELETE_RECORDS.latestVersion
     val deleteRecordsResponse = DeleteRecordsResponse.parse(response, version)
@@ -1204,8 +1446,8 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
   }
 
   @Test
-  def testUnauthorizedDeleteRecordsWithDescribe() {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Describe)), deleteTopicResource)
+  def testUnauthorizedDeleteRecordsWithDescribe(): Unit = {
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, DESCRIBE, ALLOW)), deleteTopicResource)
     val response = connectAndSend(deleteRecordsRequest, ApiKeys.DELETE_RECORDS)
     val version = ApiKeys.DELETE_RECORDS.latestVersion
     val deleteRecordsResponse = DeleteRecordsResponse.parse(response, version)
@@ -1213,8 +1455,8 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
   }
 
   @Test
-  def testDeleteRecordsWithWildCardAuth() {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Delete)), Resource(Topic, "*", LITERAL))
+  def testDeleteRecordsWithWildCardAuth(): Unit = {
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, DELETE, ALLOW)), new ResourcePattern(TOPIC, "*", LITERAL))
     val response = connectAndSend(deleteRecordsRequest, ApiKeys.DELETE_RECORDS)
     val version = ApiKeys.DELETE_RECORDS.latestVersion
     val deleteRecordsResponse = DeleteRecordsResponse.parse(response, version)
@@ -1223,7 +1465,7 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
   }
 
   @Test
-  def testUnauthorizedCreatePartitions() {
+  def testUnauthorizedCreatePartitions(): Unit = {
     val response = connectAndSend(createPartitionsRequest, ApiKeys.CREATE_PARTITIONS)
     val version = ApiKeys.CREATE_PARTITIONS.latestVersion
     val createPartitionsResponse = CreatePartitionsResponse.parse(response, version)
@@ -1231,8 +1473,8 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
   }
 
   @Test
-  def testCreatePartitionsWithWildCardAuth() {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Alter)), Resource(Topic, "*", LITERAL))
+  def testCreatePartitionsWithWildCardAuth(): Unit = {
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, ALTER, ALLOW)), new ResourcePattern(TOPIC, "*", LITERAL))
     val response = connectAndSend(createPartitionsRequest, ApiKeys.CREATE_PARTITIONS)
     val version = ApiKeys.CREATE_PARTITIONS.latestVersion
     val createPartitionsResponse = CreatePartitionsResponse.parse(response, version)
@@ -1241,7 +1483,7 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
 
   @Test(expected = classOf[TransactionalIdAuthorizationException])
   def testTransactionalProducerInitTransactionsNoWriteTransactionalIdAcl(): Unit = {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Describe)), transactionalIdResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, DESCRIBE, ALLOW)), transactionalIdResource)
     val producer = buildTransactionalProducer()
     producer.initTransactions()
   }
@@ -1254,9 +1496,9 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
 
   @Test
   def testSendOffsetsWithNoConsumerGroupDescribeAccess(): Unit = {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, ClusterAction)), Resource.ClusterResource)
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Write)), topicResource)
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Write)), transactionalIdResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, CLUSTER_ACTION, ALLOW)), clusterResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, WRITE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, WRITE, ALLOW)), transactionalIdResource)
     val producer = buildTransactionalProducer()
     producer.initTransactions()
     producer.beginTransaction()
@@ -1270,8 +1512,8 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
 
   @Test
   def testSendOffsetsWithNoConsumerGroupWriteAccess(): Unit = {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Write)), transactionalIdResource)
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Describe)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, WRITE, ALLOW)), transactionalIdResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, DESCRIBE, ALLOW)), groupResource)
     val producer = buildTransactionalProducer()
     producer.initTransactions()
     producer.beginTransaction()
@@ -1285,7 +1527,7 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
 
   @Test
   def testIdempotentProducerNoIdempotentWriteAclInInitProducerId(): Unit = {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Write)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, WRITE, ALLOW)), topicResource)
     val producer = buildIdempotentProducer()
     try {
       // the InitProducerId is sent asynchronously, so we expect the error either in the callback
@@ -1309,8 +1551,8 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
 
   @Test
   def testIdempotentProducerNoIdempotentWriteAclInProduce(): Unit = {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Write)), topicResource)
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, IdempotentWrite)), Resource.ClusterResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, WRITE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, IDEMPOTENT_WRITE, ALLOW)), clusterResource)
 
     val producer = buildIdempotentProducer()
 
@@ -1319,7 +1561,7 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
 
     // revoke the IdempotentWrite permission
     removeAllAcls()
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Write)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, WRITE, ALLOW)), topicResource)
 
     try {
       // the send should now fail with a cluster auth error
@@ -1342,16 +1584,16 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
 
   @Test
   def shouldInitTransactionsWhenAclSet(): Unit = {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Write)), transactionalIdResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, WRITE, ALLOW)), transactionalIdResource)
     val producer = buildTransactionalProducer()
     producer.initTransactions()
   }
 
   @Test
   def testTransactionalProducerTopicAuthorizationExceptionInSendCallback(): Unit = {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Write)), transactionalIdResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, WRITE, ALLOW)), transactionalIdResource)
     // add describe access so that we can fetch metadata
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Describe)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, DESCRIBE, ALLOW)), topicResource)
     val producer = buildTransactionalProducer()
     producer.initTransactions()
     producer.beginTransaction()
@@ -1371,9 +1613,9 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
 
   @Test
   def testTransactionalProducerTopicAuthorizationExceptionInCommit(): Unit = {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Write)), transactionalIdResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, WRITE, ALLOW)), transactionalIdResource)
     // add describe access so that we can fetch metadata
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Describe)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, DESCRIBE, ALLOW)), topicResource)
     val producer = buildTransactionalProducer()
     producer.initTransactions()
     producer.beginTransaction()
@@ -1389,11 +1631,11 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
 
   @Test
   def shouldThrowTransactionalIdAuthorizationExceptionWhenNoTransactionAccessDuringSend(): Unit = {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Write)), transactionalIdResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, WRITE, ALLOW)), transactionalIdResource)
     val producer = buildTransactionalProducer()
     producer.initTransactions()
     removeAllAcls()
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Write)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, WRITE, ALLOW)), topicResource)
     try {
       producer.beginTransaction()
       producer.send(new ProducerRecord(tp.topic, tp.partition, "1".getBytes, "1".getBytes)).get
@@ -1406,8 +1648,8 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
 
   @Test
   def shouldThrowTransactionalIdAuthorizationExceptionWhenNoTransactionAccessOnEndTransaction(): Unit = {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Write)), transactionalIdResource)
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Write)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, WRITE, ALLOW)), transactionalIdResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, WRITE, ALLOW)), topicResource)
     val producer = buildTransactionalProducer()
     producer.initTransactions()
     producer.beginTransaction()
@@ -1423,9 +1665,9 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
 
   @Test
   def shouldSuccessfullyAbortTransactionAfterTopicAuthorizationException(): Unit = {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Write)), transactionalIdResource)
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Write)), topicResource)
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Describe)), Resource(Topic, deleteTopic, LITERAL))
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, WRITE, ALLOW)), transactionalIdResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, WRITE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, DESCRIBE, ALLOW)), new ResourcePattern(TOPIC, deleteTopic, LITERAL))
     val producer = buildTransactionalProducer()
     producer.initTransactions()
     producer.beginTransaction()
@@ -1443,8 +1685,8 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
 
   @Test
   def shouldThrowTransactionalIdAuthorizationExceptionWhenNoTransactionAccessOnSendOffsetsToTxn(): Unit = {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Write)), transactionalIdResource)
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Write)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, WRITE, ALLOW)), transactionalIdResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, WRITE, ALLOW)), groupResource)
     val producer = buildTransactionalProducer()
     producer.initTransactions()
     producer.beginTransaction()
@@ -1460,19 +1702,35 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
 
   @Test
   def shouldSendSuccessfullyWhenIdempotentAndHasCorrectACL(): Unit = {
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, IdempotentWrite)), Resource.ClusterResource)
-    addAndVerifyAcls(Set(new Acl(userPrincipal, Allow, Acl.WildCardHost, Write)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, IDEMPOTENT_WRITE, ALLOW)), clusterResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(userPrincipalStr, WildcardHost, WRITE, ALLOW)), topicResource)
     val producer = buildIdempotentProducer()
     producer.send(new ProducerRecord(tp.topic, tp.partition, "1".getBytes, "1".getBytes)).get
   }
 
-  def removeAllAcls() = {
-    servers.head.dataPlaneRequestProcessor.authorizer.get.getAcls().keys.foreach { resource =>
-      servers.head.dataPlaneRequestProcessor.authorizer.get.removeAcls(resource)
-      TestUtils.waitAndVerifyAcls(Set.empty[Acl], servers.head.dataPlaneRequestProcessor.authorizer.get, resource)
+  // Verify that metadata request without topics works without any ACLs and returns cluster id
+  @Test
+  def testClusterId(): Unit = {
+    val request = new requests.MetadataRequest.Builder(List.empty.asJava, false).build()
+    val apiKey = ApiKeys.METADATA
+    val resp = connectAndSend(request, apiKey)
+    val response = requestKeyToResponseDeserializer(apiKey).getMethod("parse", classOf[ByteBuffer], classOf[Short]).invoke(
+      null, resp, request.version: java.lang.Short).asInstanceOf[MetadataResponse]
+    assertEquals(Collections.emptyMap, response.errorCounts)
+    assertFalse("Cluster id not returned", response.clusterId.isEmpty)
+  }
+
+  def removeAllAcls(): Unit = {
+    val authorizer = servers.head.dataPlaneRequestProcessor.authorizer.get
+    val aclFilter = AclBindingFilter.ANY
+    authorizer.deleteAcls(null, List(aclFilter).asJava).asScala.map(_.toCompletableFuture.get).flatMap { deletion =>
+      deletion.aclBindingDeleteResults().asScala.map(_.aclBinding.pattern).toSet
+    }.foreach { resource =>
+      TestUtils.waitAndVerifyAcls(Set.empty[AccessControlEntry], authorizer, resource)
     }
   }
 
+
   def sendRequestAndVerifyResponseError(apiKey: ApiKeys,
                                         request: AbstractRequest,
                                         resources: Set[ResourceType],
@@ -1484,13 +1742,13 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
     val error = requestKeyToError(apiKey).asInstanceOf[AbstractResponse => Errors](response)
 
     val authorizationErrors = resources.flatMap { resourceType =>
-      if (resourceType == Topic) {
+      if (resourceType == TOPIC) {
         if (isAuthorized)
           Set(Errors.UNKNOWN_TOPIC_OR_PARTITION, Topic.error)
         else
           Set(Topic.error)
       } else {
-        Set(resourceType.error)
+        Set(AuthResourceType.fromJava(resourceType).error)
       }
     }
 
@@ -1499,7 +1757,7 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
         assertFalse(s"$apiKey should be allowed. Found unexpected authorization error $error", authorizationErrors.contains(error))
       else
         assertTrue(s"$apiKey should be forbidden. Found error $error but expected one of $authorizationErrors", authorizationErrors.contains(error))
-    else if (resources == Set(Topic))
+    else if (resources == Set(TOPIC))
       if (isAuthorized)
         assertEquals(s"$apiKey had an unexpected error", Errors.UNKNOWN_TOPIC_OR_PARTITION, error)
       else
@@ -1510,7 +1768,7 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
 
   private def sendRecords(producer: KafkaProducer[Array[Byte], Array[Byte]],
                           numRecords: Int,
-                          tp: TopicPartition) {
+                          tp: TopicPartition): Unit = {
     val futures = (0 until numRecords).map { i =>
       producer.send(new ProducerRecord(tp.topic(), tp.partition(), i.toString.getBytes, i.toString.getBytes))
     }
@@ -1521,16 +1779,23 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
     }
   }
 
-  private def addAndVerifyAcls(acls: Set[Acl], resource: Resource) = {
-    servers.head.dataPlaneRequestProcessor.authorizer.get.addAcls(acls, resource)
-    TestUtils.waitAndVerifyAcls(servers.head.dataPlaneRequestProcessor.authorizer.get.getAcls(resource) ++ acls, servers.head.dataPlaneRequestProcessor.authorizer.get, resource)
+  private def addAndVerifyAcls(acls: Set[AccessControlEntry], resource: ResourcePattern): Unit = {
+    val aclBindings = acls.map { acl => new AclBinding(resource, acl) }
+    servers.head.dataPlaneRequestProcessor.authorizer.get
+      .createAcls(null, aclBindings.toList.asJava).asScala.map(_.toCompletableFuture.get).foreach {result =>
+        result.exception.asScala.foreach { e => throw e }
+      }
+    val aclFilter = new AclBindingFilter(resource.toFilter, AccessControlEntryFilter.ANY)
+    TestUtils.waitAndVerifyAcls(
+      servers.head.dataPlaneRequestProcessor.authorizer.get.acls(aclFilter).asScala.map(_.entry).toSet ++ acls,
+      servers.head.dataPlaneRequestProcessor.authorizer.get, resource)
   }
 
   private def consumeRecords(consumer: Consumer[Array[Byte], Array[Byte]],
                              numRecords: Int = 1,
                              startingOffset: Int = 0,
                              topic: String = topic,
-                             part: Int = part) {
+                             part: Int = part): Unit = {
     val records = TestUtils.consumeRecords(consumer, numRecords)
 
     for (i <- 0 until numRecords) {
@@ -1560,7 +1825,7 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
     createProducer()
   }
 
-  private def createAdminClient(): AdminClient = {
+  private def createAdminClient(): Admin = {
     val props = new Properties()
     props.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, brokerList)
     val adminClient = AdminClient.create(props)
diff --git a/core/src/test/scala/integration/kafka/api/BaseConsumerTest.scala b/core/src/test/scala/integration/kafka/api/BaseConsumerTest.scala
index 1645b5badd63b..2b0700b38f869 100644
--- a/core/src/test/scala/integration/kafka/api/BaseConsumerTest.scala
+++ b/core/src/test/scala/integration/kafka/api/BaseConsumerTest.scala
@@ -1,76 +1,37 @@
-/**
- * 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
+/*
+ * 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
+ *    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.
+ * 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 java.time.Duration
-import java.util
-
-import org.apache.kafka.clients.consumer._
-import org.apache.kafka.clients.producer.{ProducerConfig, ProducerRecord}
-import org.apache.kafka.common.record.TimestampType
-import org.apache.kafka.common.{PartitionInfo, TopicPartition}
-import kafka.utils.{ShutdownableThread, TestUtils}
-import kafka.server.KafkaConfig
+import org.apache.kafka.clients.consumer.ConsumerConfig
+import org.apache.kafka.common.PartitionInfo
+import org.apache.kafka.common.internals.Topic
+import org.junit.Test
 import org.junit.Assert._
-import org.junit.{Before, Test}
 
 import scala.collection.JavaConverters._
-import scala.collection.mutable.{ArrayBuffer, Buffer}
-import org.apache.kafka.clients.producer.KafkaProducer
-import org.apache.kafka.common.errors.WakeupException
-import org.apache.kafka.common.internals.Topic
+import scala.collection.Seq
 
 /**
- * Integration tests for the consumer that cover basic usage as well as server failures
+ * Integration tests for the consumer that cover basic usage as well as coordinator failure
  */
-abstract class BaseConsumerTest extends IntegrationTestHarness {
-
-  val epsilon = 0.1
-  val serverCount = 3
-
-  val topic = "topic"
-  val part = 0
-  val tp = new TopicPartition(topic, part)
-  val part2 = 1
-  val tp2 = new TopicPartition(topic, part2)
-  val producerClientId = "ConsumerTestProducer"
-  val consumerClientId = "ConsumerTestConsumer"
-
-  // configure the servers and clients
-  this.serverConfig.setProperty(KafkaConfig.ControlledShutdownEnableProp, "false") // speed up shutdown
-  this.serverConfig.setProperty(KafkaConfig.OffsetsTopicReplicationFactorProp, "3") // don't want to lose offset
-  this.serverConfig.setProperty(KafkaConfig.OffsetsTopicPartitionsProp, "1")
-  this.serverConfig.setProperty(KafkaConfig.GroupMinSessionTimeoutMsProp, "100") // set small enough session timeout
-  this.serverConfig.setProperty(KafkaConfig.GroupMaxSessionTimeoutMsProp, "30000")
-  this.serverConfig.setProperty(KafkaConfig.GroupInitialRebalanceDelayMsProp, "10")
-  this.producerConfig.setProperty(ProducerConfig.ACKS_CONFIG, "all")
-  this.producerConfig.setProperty(ProducerConfig.CLIENT_ID_CONFIG, producerClientId)
-  this.consumerConfig.setProperty(ConsumerConfig.CLIENT_ID_CONFIG, consumerClientId)
-  this.consumerConfig.setProperty(ConsumerConfig.GROUP_ID_CONFIG, "my-test")
-  this.consumerConfig.setProperty(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest")
-  this.consumerConfig.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "false")
-  this.consumerConfig.setProperty(ConsumerConfig.METADATA_MAX_AGE_CONFIG, "100")
-
-  @Before
-  override def setUp() {
-    super.setUp()
-
-    // create the test topic with all the brokers as replicas
-    createTopic(topic, 2, serverCount)
-  }
+abstract class BaseConsumerTest extends AbstractConsumerTest {
 
   @Test
-  def testSimpleConsumption() {
+  def testSimpleConsumption(): Unit = {
     val numRecords = 10000
     val producer = createProducer()
     sendRecords(producer, numRecords, tp)
@@ -88,9 +49,9 @@ abstract class BaseConsumerTest extends IntegrationTestHarness {
   }
 
   @Test
-  def testCoordinatorFailover() {
+  def testCoordinatorFailover(): Unit = {
     val listener = new TestConsumerReassignmentListener()
-    this.consumerConfig.setProperty(ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG, "5000")
+    this.consumerConfig.setProperty(ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG, "5001")
     this.consumerConfig.setProperty(ConsumerConfig.HEARTBEAT_INTERVAL_MS_CONFIG, "2000")
     val consumer = createConsumer()
 
@@ -114,239 +75,4 @@ abstract class BaseConsumerTest extends IntegrationTestHarness {
     // the failover should not cause a rebalance
     ensureNoRebalance(consumer, listener)
   }
-
-  protected class TestConsumerReassignmentListener extends ConsumerRebalanceListener {
-    var callsToAssigned = 0
-    var callsToRevoked = 0
-
-    def onPartitionsAssigned(partitions: java.util.Collection[TopicPartition]) {
-      info("onPartitionsAssigned called.")
-      callsToAssigned += 1
-    }
-
-    def onPartitionsRevoked(partitions: java.util.Collection[TopicPartition]) {
-      info("onPartitionsRevoked called.")
-      callsToRevoked += 1
-    }
-  }
-
-  protected def sendRecords(producer: KafkaProducer[Array[Byte], Array[Byte]], numRecords: Int,
-                            tp: TopicPartition): Seq[ProducerRecord[Array[Byte], Array[Byte]]] = {
-    val records = (0 until numRecords).map { i =>
-      val record = new ProducerRecord(tp.topic(), tp.partition(), i.toLong, s"key $i".getBytes, s"value $i".getBytes)
-      producer.send(record)
-      record
-    }
-    producer.flush()
-
-    records
-  }
-
-  protected def consumeAndVerifyRecords(consumer: Consumer[Array[Byte], Array[Byte]],
-                                        numRecords: Int,
-                                        startingOffset: Int,
-                                        startingKeyAndValueIndex: Int = 0,
-                                        startingTimestamp: Long = 0L,
-                                        timestampType: TimestampType = TimestampType.CREATE_TIME,
-                                        tp: TopicPartition = tp,
-                                        maxPollRecords: Int = Int.MaxValue) {
-    val records = consumeRecords(consumer, numRecords, maxPollRecords = maxPollRecords)
-    val now = System.currentTimeMillis()
-    for (i <- 0 until numRecords) {
-      val record = records(i)
-      val offset = startingOffset + i
-      assertEquals(tp.topic, record.topic)
-      assertEquals(tp.partition, record.partition)
-      if (timestampType == TimestampType.CREATE_TIME) {
-        assertEquals(timestampType, record.timestampType)
-        val timestamp = startingTimestamp + i
-        assertEquals(timestamp.toLong, record.timestamp)
-      } else
-        assertTrue(s"Got unexpected timestamp ${record.timestamp}. Timestamp should be between [$startingTimestamp, $now}]",
-          record.timestamp >= startingTimestamp && record.timestamp <= now)
-      assertEquals(offset.toLong, record.offset)
-      val keyAndValueIndex = startingKeyAndValueIndex + i
-      assertEquals(s"key $keyAndValueIndex", new String(record.key))
-      assertEquals(s"value $keyAndValueIndex", new String(record.value))
-      // this is true only because K and V are byte arrays
-      assertEquals(s"key $keyAndValueIndex".length, record.serializedKeySize)
-      assertEquals(s"value $keyAndValueIndex".length, record.serializedValueSize)
-    }
-  }
-
-  protected def consumeRecords[K, V](consumer: Consumer[K, V],
-                                     numRecords: Int,
-                                     maxPollRecords: Int = Int.MaxValue): ArrayBuffer[ConsumerRecord[K, V]] = {
-    val records = new ArrayBuffer[ConsumerRecord[K, V]]
-    def pollAction(polledRecords: ConsumerRecords[K, V]): Boolean = {
-      assertTrue(polledRecords.asScala.size <= maxPollRecords)
-      records ++= polledRecords.asScala
-      records.size >= numRecords
-    }
-    TestUtils.pollRecordsUntilTrue(consumer, pollAction, waitTimeMs = 60000,
-      msg = s"Timed out before consuming expected $numRecords records. " +
-        s"The number consumed was ${records.size}.")
-    records
-  }
-
-  protected def sendAndAwaitAsyncCommit[K, V](consumer: Consumer[K, V],
-                                              offsetsOpt: Option[Map[TopicPartition, OffsetAndMetadata]] = None): Unit = {
-
-    def sendAsyncCommit(callback: OffsetCommitCallback) = {
-      offsetsOpt match {
-        case Some(offsets) => consumer.commitAsync(offsets.asJava, callback)
-        case None => consumer.commitAsync(callback)
-      }
-    }
-
-    class RetryCommitCallback extends OffsetCommitCallback {
-      var isComplete = false
-      var error: Option[Exception] = None
-
-      override def onComplete(offsets: util.Map[TopicPartition, OffsetAndMetadata], exception: Exception): Unit = {
-        exception match {
-          case e: RetriableCommitFailedException =>
-            sendAsyncCommit(this)
-          case e =>
-            isComplete = true
-            error = Option(e)
-        }
-      }
-    }
-
-    val commitCallback = new RetryCommitCallback
-
-    sendAsyncCommit(commitCallback)
-    TestUtils.pollUntilTrue(consumer, () => commitCallback.isComplete,
-      "Failed to observe commit callback before timeout", waitTimeMs = 10000)
-
-    assertEquals(None, commitCallback.error)
-  }
-
-  protected def awaitRebalance(consumer: Consumer[_, _], rebalanceListener: TestConsumerReassignmentListener): Unit = {
-    val numReassignments = rebalanceListener.callsToAssigned
-    TestUtils.pollUntilTrue(consumer, () => rebalanceListener.callsToAssigned > numReassignments,
-      "Timed out before expected rebalance completed")
-  }
-
-  protected def ensureNoRebalance(consumer: Consumer[_, _], rebalanceListener: TestConsumerReassignmentListener): Unit = {
-    // The best way to verify that the current membership is still active is to commit offsets.
-    // This would fail if the group had rebalanced.
-    val initialRevokeCalls = rebalanceListener.callsToRevoked
-    sendAndAwaitAsyncCommit(consumer)
-    assertEquals(initialRevokeCalls, rebalanceListener.callsToRevoked)
-  }
-
-  protected class CountConsumerCommitCallback extends OffsetCommitCallback {
-    var successCount = 0
-    var failCount = 0
-    var lastError: Option[Exception] = None
-
-    override def onComplete(offsets: util.Map[TopicPartition, OffsetAndMetadata], exception: Exception): Unit = {
-      if (exception == null) {
-        successCount += 1
-      } else {
-        failCount += 1
-        lastError = Some(exception)
-      }
-    }
-  }
-
-  protected class ConsumerAssignmentPoller(consumer: Consumer[Array[Byte], Array[Byte]],
-                                           topicsToSubscribe: List[String]) extends ShutdownableThread("daemon-consumer-assignment", false)
-  {
-    @volatile private var partitionAssignment: Set[TopicPartition] = Set.empty[TopicPartition]
-    private var topicsSubscription = topicsToSubscribe
-    @volatile private var subscriptionChanged = false
-
-    val rebalanceListener = new ConsumerRebalanceListener {
-      override def onPartitionsAssigned(partitions: util.Collection[TopicPartition]) = {
-        partitionAssignment = collection.immutable.Set(consumer.assignment().asScala.toArray: _*)
-      }
-
-      override def onPartitionsRevoked(partitions: util.Collection[TopicPartition]) = {
-        partitionAssignment = Set.empty[TopicPartition]
-      }
-    }
-    consumer.subscribe(topicsToSubscribe.asJava, rebalanceListener)
-
-    def consumerAssignment(): Set[TopicPartition] = {
-      partitionAssignment
-    }
-
-    /**
-     * Subscribe consumer to a new set of topics.
-     * Since this method most likely be called from a different thread, this function
-     * just "schedules" the subscription change, and actual call to consumer.subscribe is done
-     * in the doWork() method
-     *
-     * This method does not allow to change subscription until doWork processes the previous call
-     * to this method. This is just to avoid race conditions and enough functionality for testing purposes
-     * @param newTopicsToSubscribe
-     */
-    def subscribe(newTopicsToSubscribe: List[String]): Unit = {
-      if (subscriptionChanged) {
-        throw new IllegalStateException("Do not call subscribe until the previous subscribe request is processed.")
-      }
-      topicsSubscription = newTopicsToSubscribe
-      subscriptionChanged = true
-    }
-
-    def isSubscribeRequestProcessed(): Boolean = {
-      !subscriptionChanged
-    }
-
-    override def initiateShutdown(): Boolean = {
-      val res = super.initiateShutdown()
-      consumer.wakeup()
-      res
-    }
-
-    override def doWork(): Unit = {
-      if (subscriptionChanged) {
-        consumer.subscribe(topicsSubscription.asJava, rebalanceListener)
-        subscriptionChanged = false
-      }
-      try {
-        consumer.poll(Duration.ofMillis(50))
-      } catch {
-        case _: WakeupException => // ignore for shutdown
-      }
-    }
-  }
-
-  /**
-   * Check whether partition assignment is valid
-   * Assumes partition assignment is valid iff
-   * 1. Every consumer got assigned at least one partition
-   * 2. Each partition is assigned to only one consumer
-   * 3. Every partition is assigned to one of the consumers
-   *
-   * @param assignments set of consumer assignments; one per each consumer
-   * @param partitions set of partitions that consumers subscribed to
-   * @return true if partition assignment is valid
-   */
-  def isPartitionAssignmentValid(assignments: Buffer[Set[TopicPartition]],
-                                 partitions: Set[TopicPartition]): Boolean = {
-    val allNonEmptyAssignments = assignments.forall(assignment => assignment.nonEmpty)
-    if (!allNonEmptyAssignments) {
-      // at least one consumer got empty assignment
-      return false
-    }
-
-    // make sure that sum of all partitions to all consumers equals total number of partitions
-    val totalPartitionsInAssignments = (0 /: assignments) (_ + _.size)
-    if (totalPartitionsInAssignments != partitions.size) {
-      // either same partitions got assigned to more than one consumer or some
-      // partitions were not assigned
-      return false
-    }
-
-    // The above checks could miss the case where one or more partitions were assigned to more
-    // than one consumer and the same number of partitions were missing from assignments.
-    // Make sure that all unique assignments are the same as 'partitions'
-    val uniqueAssignedPartitions = (Set[TopicPartition]() /: assignments) (_ ++ _)
-    uniqueAssignedPartitions == partitions
-  }
-
 }
diff --git a/core/src/test/scala/integration/kafka/api/BaseProducerSendTest.scala b/core/src/test/scala/integration/kafka/api/BaseProducerSendTest.scala
index 9d454e93bf2b7..b318f8bef88ae 100644
--- a/core/src/test/scala/integration/kafka/api/BaseProducerSendTest.scala
+++ b/core/src/test/scala/integration/kafka/api/BaseProducerSendTest.scala
@@ -17,6 +17,7 @@
 
 package kafka.api
 
+import java.time.Duration
 import java.nio.charset.StandardCharsets
 import java.util.Properties
 import java.util.concurrent.TimeUnit
@@ -33,6 +34,7 @@ import org.apache.kafka.common.security.auth.SecurityProtocol
 import org.apache.kafka.common.{KafkaException, TopicPartition}
 import org.junit.Assert._
 import org.junit.{After, Before, Test}
+import org.scalatest.Assertions.fail
 
 import scala.collection.JavaConverters._
 import scala.collection.mutable.Buffer
@@ -40,8 +42,10 @@ import scala.concurrent.ExecutionException
 
 abstract class BaseProducerSendTest extends KafkaServerTestHarness {
 
+  def overrideConfigs() : Properties = {new Properties()}
+
   def generateConfigs = {
-    val overridingProps = new Properties()
+    val overridingProps = overrideConfigs()
     val numServers = 2
     overridingProps.put(KafkaConfig.NumPartitionsProp, 4.toString)
     TestUtils.createBrokerConfigs(numServers, zkConnect, false, interBrokerSecurityProtocol = Some(securityProtocol),
@@ -55,13 +59,13 @@ abstract class BaseProducerSendTest extends KafkaServerTestHarness {
   private val numRecords = 100
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     super.setUp()
     consumer = TestUtils.createConsumer(TestUtils.getBrokerListStrFromServers(servers), securityProtocol = SecurityProtocol.PLAINTEXT)
   }
 
   @After
-  override def tearDown() {
+  override def tearDown(): Unit = {
     consumer.close()
     // Ensure that all producers are closed since unclosed producers impact other tests when Kafka server ports are reused
     producers.foreach(_.close())
@@ -71,6 +75,7 @@ abstract class BaseProducerSendTest extends KafkaServerTestHarness {
 
   protected def createProducer(brokerList: String,
                                lingerMs: Int = 0,
+                               deliveryTimeoutMs: Int = 2 * 60 * 1000,
                                batchSize: Int = 16384,
                                compressionType: String = "none",
                                maxBlockMs: Long = 60 * 1000L): KafkaProducer[Array[Byte],Array[Byte]] = {
@@ -80,6 +85,7 @@ abstract class BaseProducerSendTest extends KafkaServerTestHarness {
       trustStoreFile = trustStoreFile,
       saslProperties = clientSaslProperties,
       lingerMs = lingerMs,
+      deliveryTimeoutMs = deliveryTimeoutMs,
       maxBlockMs = maxBlockMs)
     registerProducer(producer)
   }
@@ -96,14 +102,14 @@ abstract class BaseProducerSendTest extends KafkaServerTestHarness {
    * 2. Last message of the non-blocking send should return the correct offset metadata
    */
   @Test
-  def testSendOffset() {
+  def testSendOffset(): Unit = {
     val producer = createProducer(brokerList)
     val partition = 0
 
     object callback extends Callback {
       var offset = 0L
 
-      def onCompletion(metadata: RecordMetadata, exception: Exception) {
+      def onCompletion(metadata: RecordMetadata, exception: Exception): Unit = {
         if (exception == null) {
           assertEquals(offset, metadata.offset())
           assertEquals(topic, metadata.topic())
@@ -168,22 +174,23 @@ abstract class BaseProducerSendTest extends KafkaServerTestHarness {
   }
 
   @Test
-  def testSendCompressedMessageWithCreateTime() {
+  def testSendCompressedMessageWithCreateTime(): Unit = {
     val producer = createProducer(brokerList = brokerList,
       compressionType = "gzip",
-      lingerMs = Int.MaxValue)
+      lingerMs = Int.MaxValue,
+      deliveryTimeoutMs = Int.MaxValue)
     sendAndVerifyTimestamp(producer, TimestampType.CREATE_TIME)
   }
 
   @Test
-  def testSendNonCompressedMessageWithCreateTime() {
-    val producer = createProducer(brokerList = brokerList, lingerMs = Int.MaxValue)
+  def testSendNonCompressedMessageWithCreateTime(): Unit = {
+    val producer = createProducer(brokerList = brokerList, lingerMs = Int.MaxValue, deliveryTimeoutMs = Int.MaxValue)
     sendAndVerifyTimestamp(producer, TimestampType.CREATE_TIME)
   }
 
   protected def sendAndVerify(producer: KafkaProducer[Array[Byte], Array[Byte]],
                               numRecords: Int = numRecords,
-                              timeoutMs: Long = 20000L) {
+                              timeoutMs: Long = 20000L): Unit = {
     val partition = 0
     try {
       createTopic(topic, 1, 2)
@@ -193,7 +200,7 @@ abstract class BaseProducerSendTest extends KafkaServerTestHarness {
           s"value$i".getBytes(StandardCharsets.UTF_8))
         producer.send(record)
       }
-      producer.close(timeoutMs, TimeUnit.MILLISECONDS)
+      producer.close(Duration.ofMillis(timeoutMs))
       val lastOffset = futures.foldLeft(0) { (offset, future) =>
         val recordMetadata = future.get
         assertEquals(topic, recordMetadata.topic)
@@ -207,7 +214,7 @@ abstract class BaseProducerSendTest extends KafkaServerTestHarness {
     }
   }
 
-  protected def sendAndVerifyTimestamp(producer: KafkaProducer[Array[Byte], Array[Byte]], timestampType: TimestampType) {
+  protected def sendAndVerifyTimestamp(producer: KafkaProducer[Array[Byte], Array[Byte]], timestampType: TimestampType): Unit = {
     val partition = 0
 
     val baseTimestamp = 123456L
@@ -217,7 +224,7 @@ abstract class BaseProducerSendTest extends KafkaServerTestHarness {
       var offset = 0L
       var timestampDiff = 1L
 
-      def onCompletion(metadata: RecordMetadata, exception: Exception) {
+      def onCompletion(metadata: RecordMetadata, exception: Exception): Unit = {
         if (exception == null) {
           assertEquals(offset, metadata.offset)
           assertEquals(topic, metadata.topic)
@@ -248,7 +255,7 @@ abstract class BaseProducerSendTest extends KafkaServerTestHarness {
           s"value$i".getBytes(StandardCharsets.UTF_8))
         (record, producer.send(record, callback))
       }
-      producer.close(20000L, TimeUnit.MILLISECONDS)
+      producer.close(Duration.ofSeconds(20L))
       recordAndFutures.foreach { case (record, future) =>
         val recordMetadata = future.get
         if (timestampType == TimestampType.LOG_APPEND_TIME)
@@ -268,7 +275,7 @@ abstract class BaseProducerSendTest extends KafkaServerTestHarness {
    * After close() returns, all messages should be sent with correct returned offset metadata
    */
   @Test
-  def testClose() {
+  def testClose(): Unit = {
     val producer = createProducer(brokerList)
 
     try {
@@ -301,7 +308,7 @@ abstract class BaseProducerSendTest extends KafkaServerTestHarness {
    * The specified partition-id should be respected
    */
   @Test
-  def testSendToPartition() {
+  def testSendToPartition(): Unit = {
     val producer = createProducer(brokerList)
 
     try {
@@ -346,7 +353,7 @@ abstract class BaseProducerSendTest extends KafkaServerTestHarness {
     * succeed as long as the partition is included in the metadata.
     */
   @Test
-  def testSendBeforeAndAfterPartitionExpansion() {
+  def testSendBeforeAndAfterPartitionExpansion(): Unit = {
     val producer = createProducer(brokerList, maxBlockMs = 5 * 1000L)
 
     // create topic
@@ -376,8 +383,8 @@ abstract class BaseProducerSendTest extends KafkaServerTestHarness {
       }
     }
 
-    val existingAssignment = zkClient.getReplicaAssignmentForTopics(Set(topic)).map {
-      case (topicPartition, replicas) => topicPartition.partition -> replicas
+    val existingAssignment = zkClient.getFullReplicaAssignmentForTopics(Set(topic)).map {
+      case (topicPartition, assignment) => topicPartition.partition -> assignment
     }
     adminZkClient.addPartitions(topic, existingAssignment, adminZkClient.getBrokerMetadatas(), 2)
     // read metadata from a broker and verify the new topic partitions exist
@@ -412,8 +419,8 @@ abstract class BaseProducerSendTest extends KafkaServerTestHarness {
    * Test that flush immediately sends all accumulated requests.
    */
   @Test
-  def testFlush() {
-    val producer = createProducer(brokerList, lingerMs = Int.MaxValue)
+  def testFlush(): Unit = {
+    val producer = createProducer(brokerList, lingerMs = Int.MaxValue, deliveryTimeoutMs = Int.MaxValue)
     try {
       createTopic(topic, 2, 2)
       val record = new ProducerRecord[Array[Byte], Array[Byte]](topic,
@@ -428,12 +435,42 @@ abstract class BaseProducerSendTest extends KafkaServerTestHarness {
       producer.close()
     }
   }
+  /**
+  * Test that flush return with TimeoutException when producer is unable to finish sending buffered records in time.
+*/
+  @Test
+  def testBoundedFlush() {
+    val producer = createProducer(brokerList)
+    createTopic(topic, 2, 2)
+    try {
+      producer.send(new ProducerRecord(topic, null, "value1".getBytes())).get()
+      try {
+        killBroker(0)
+        } catch {
+        case _: Throwable =>
+        }
+      try {
+        killBroker(1)
+        } catch {
+        case _: Throwable =>
+        }
+      producer.send(new ProducerRecord(topic, null, "value2".getBytes()))
+      try {
+        producer.flush(1000, TimeUnit.MILLISECONDS)
+        fail("TimeoutException should have thrown")
+        } catch {
+        case _: TimeoutException =>
+        }
+      }finally {
+      producer.close(1000, TimeUnit.MILLISECONDS)
+      }
+    }
 
   /**
    * Test close with zero timeout from caller thread
    */
   @Test
-  def testCloseWithZeroTimeoutFromCallerThread() {
+  def testCloseWithZeroTimeoutFromCallerThread(): Unit = {
     createTopic(topic, 2, 2)
     val partition = 0
     consumer.assign(List(new TopicPartition(topic, partition)).asJava)
@@ -442,10 +479,10 @@ abstract class BaseProducerSendTest extends KafkaServerTestHarness {
 
     // Test closing from caller thread.
     for (_ <- 0 until 50) {
-      val producer = createProducer(brokerList, lingerMs = Int.MaxValue)
+      val producer = createProducer(brokerList, lingerMs = Int.MaxValue, deliveryTimeoutMs = Int.MaxValue)
       val responses = (0 until numRecords) map (_ => producer.send(record0))
       assertTrue("No request is complete.", responses.forall(!_.isDone()))
-      producer.close(0, TimeUnit.MILLISECONDS)
+      producer.close(Duration.ZERO)
       responses.foreach { future =>
         try {
           future.get()
@@ -454,7 +491,7 @@ abstract class BaseProducerSendTest extends KafkaServerTestHarness {
           case e: ExecutionException => assertEquals(classOf[KafkaException], e.getCause.getClass)
         }
       }
-      assertEquals("Fetch response should have no message returned.", 0, consumer.poll(50).count)
+      assertEquals("Fetch response should have no message returned.", 0, consumer.poll(Duration.ofMillis(50L)).count)
     }
   }
 
@@ -462,7 +499,7 @@ abstract class BaseProducerSendTest extends KafkaServerTestHarness {
    * Test close with zero and non-zero timeout from sender thread
    */
   @Test
-  def testCloseWithZeroTimeoutFromSenderThread() {
+  def testCloseWithZeroTimeoutFromSenderThread(): Unit = {
     createTopic(topic, 1, 2)
     val partition = 0
     consumer.assign(List(new TopicPartition(topic, partition)).asJava)
@@ -470,19 +507,19 @@ abstract class BaseProducerSendTest extends KafkaServerTestHarness {
 
     // Test closing from sender thread.
     class CloseCallback(producer: KafkaProducer[Array[Byte], Array[Byte]], sendRecords: Boolean) extends Callback {
-      override def onCompletion(metadata: RecordMetadata, exception: Exception) {
+      override def onCompletion(metadata: RecordMetadata, exception: Exception): Unit = {
         // Trigger another batch in accumulator before close the producer. These messages should
         // not be sent.
         if (sendRecords)
           (0 until numRecords) foreach (_ => producer.send(record))
         // The close call will be called by all the message callbacks. This tests idempotence of the close call.
-        producer.close(0, TimeUnit.MILLISECONDS)
+        producer.close(Duration.ZERO)
         // Test close with non zero timeout. Should not block at all.
-        producer.close(Long.MaxValue, TimeUnit.MICROSECONDS)
+        producer.close()
       }
     }
     for (i <- 0 until 50) {
-      val producer = createProducer(brokerList, lingerMs = Int.MaxValue)
+      val producer = createProducer(brokerList, lingerMs = Int.MaxValue, deliveryTimeoutMs = Int.MaxValue)
       try {
         // send message to partition 0
         // Only send the records in the first callback since we close the producer in the callback and no records
diff --git a/core/src/test/scala/integration/kafka/api/BaseQuotaTest.scala b/core/src/test/scala/integration/kafka/api/BaseQuotaTest.scala
index b28a40f890e00..4d9d23e70f90c 100644
--- a/core/src/test/scala/integration/kafka/api/BaseQuotaTest.scala
+++ b/core/src/test/scala/integration/kafka/api/BaseQuotaTest.scala
@@ -14,6 +14,8 @@
 
 package kafka.api
 
+import java.time.Duration
+import java.util.concurrent.TimeUnit
 import java.util.{Collections, HashMap, Properties}
 
 import kafka.api.QuotaTestClients._
@@ -32,7 +34,7 @@ import scala.collection.JavaConverters._
 
 abstract class BaseQuotaTest extends IntegrationTestHarness {
 
-  override val serverCount = 2
+  override val brokerCount = 2
 
   protected def producerClientId = "QuotasTestProducer-1"
   protected def consumerClientId = "QuotasTestConsumer-1"
@@ -65,29 +67,29 @@ abstract class BaseQuotaTest extends IntegrationTestHarness {
   var quotaTestClients: QuotaTestClients = _
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     super.setUp()
 
     val numPartitions = 1
-    val leaders = createTopic(topic1, numPartitions, serverCount)
+    val leaders = createTopic(topic1, numPartitions, brokerCount)
     leaderNode = if (leaders(0) == servers.head.config.brokerId) servers.head else servers(1)
     followerNode = if (leaders(0) != servers.head.config.brokerId) servers.head else servers(1)
     quotaTestClients = createQuotaTestClients(topic1, leaderNode)
   }
 
   @Test
-  def testThrottledProducerConsumer() {
+  def testThrottledProducerConsumer(): Unit = {
     val numRecords = 1000
     val produced = quotaTestClients.produceUntilThrottled(numRecords)
     quotaTestClients.verifyProduceThrottle(expectThrottle = true)
 
     // Consumer should read in a bursty manner and get throttled immediately
-    quotaTestClients.consumeUntilThrottled(produced)
+    assertTrue("Should have consumed at least one record", quotaTestClients.consumeUntilThrottled(produced) > 0)
     quotaTestClients.verifyConsumeThrottle(expectThrottle = true)
   }
 
   @Test
-  def testProducerConsumerOverrideUnthrottled() {
+  def testProducerConsumerOverrideUnthrottled(): Unit = {
     // Give effectively unlimited quota for producer and consumer
     val props = new Properties()
     props.put(DynamicConfig.Client.ProducerByteRateOverrideProp, Long.MaxValue.toString)
@@ -106,7 +108,24 @@ abstract class BaseQuotaTest extends IntegrationTestHarness {
   }
 
   @Test
-  def testQuotaOverrideDelete() {
+  def testProducerConsumerOverrideLowerQuota(): Unit = {
+    // consumer quota is set such that consumer quota * default quota window (10 seconds) is less than
+    // MAX_PARTITION_FETCH_BYTES_CONFIG, so that we can test consumer ability to fetch in this case
+    // In this case, 250 * 10 < 4096
+    quotaTestClients.overrideQuotas(2000, 250, Int.MaxValue)
+    quotaTestClients.waitForQuotaUpdate(2000, 250, Int.MaxValue)
+
+    val numRecords = 1000
+    val produced = quotaTestClients.produceUntilThrottled(numRecords)
+    quotaTestClients.verifyProduceThrottle(expectThrottle = true)
+
+    // Consumer should be able to consume at least one record, even when throttled
+    assertTrue("Should have consumed at least one record", quotaTestClients.consumeUntilThrottled(produced) > 0)
+    quotaTestClients.verifyConsumeThrottle(expectThrottle = true)
+  }
+
+  @Test
+  def testQuotaOverrideDelete(): Unit = {
     // Override producer and consumer quotas to unlimited
     quotaTestClients.overrideQuotas(Long.MaxValue, Long.MaxValue, Int.MaxValue)
     quotaTestClients.waitForQuotaUpdate(Long.MaxValue, Long.MaxValue, Int.MaxValue)
@@ -131,7 +150,7 @@ abstract class BaseQuotaTest extends IntegrationTestHarness {
   }
 
   @Test
-  def testThrottledRequest() {
+  def testThrottledRequest(): Unit = {
     quotaTestClients.overrideQuotas(Long.MaxValue, Long.MaxValue, 0.1)
     quotaTestClients.waitForQuotaUpdate(Long.MaxValue, Long.MaxValue, 0.1)
 
@@ -140,7 +159,7 @@ abstract class BaseQuotaTest extends IntegrationTestHarness {
     val endTimeMs = System.currentTimeMillis + 10000
     var throttled = false
     while ((!throttled || quotaTestClients.exemptRequestMetric == null) && System.currentTimeMillis < endTimeMs) {
-      consumer.poll(100)
+      consumer.poll(Duration.ofMillis(100L))
       val throttleMetric = quotaTestClients.throttleMetric(QuotaType.Request, consumerClientId)
       throttled = throttleMetric != null && metricValue(throttleMetric) > 0
     }
@@ -167,8 +186,8 @@ abstract class QuotaTestClients(topic: String,
                                 val consumer: KafkaConsumer[Array[Byte], Array[Byte]]) {
 
   def userPrincipal: KafkaPrincipal
-  def overrideQuotas(producerQuota: Long, consumerQuota: Long, requestQuota: Double)
-  def removeQuotaOverrides()
+  def overrideQuotas(producerQuota: Long, consumerQuota: Long, requestQuota: Double): Unit
+  def removeQuotaOverrides(): Unit
 
   def quotaMetricTags(clientId: String): Map[String, String]
 
@@ -193,20 +212,24 @@ abstract class QuotaTestClients(topic: String,
   }
 
   def consumeUntilThrottled(maxRecords: Int, waitForRequestCompletion: Boolean = true): Int = {
+    val timeoutMs = TimeUnit.MINUTES.toMillis(1)
+
     consumer.subscribe(Collections.singleton(topic))
     var numConsumed = 0
     var throttled = false
+    val startMs = System.currentTimeMillis
     do {
-      numConsumed += consumer.poll(100).count
+      numConsumed += consumer.poll(Duration.ofMillis(100L)).count
       val metric = throttleMetric(QuotaType.Fetch, consumerClientId)
       throttled = metric != null && metricValue(metric) > 0
-    }  while (numConsumed < maxRecords && !throttled)
+    }  while (numConsumed < maxRecords && !throttled && System.currentTimeMillis < startMs + timeoutMs)
 
     // If throttled, wait for the records from the last fetch to be received
     if (throttled && numConsumed < maxRecords && waitForRequestCompletion) {
       val minRecords = numConsumed + 1
-      while (numConsumed < minRecords)
-        numConsumed += consumer.poll(100).count
+      val startMs = System.currentTimeMillis
+      while (numConsumed < minRecords && System.currentTimeMillis < startMs + timeoutMs)
+        numConsumed += consumer.poll(Duration.ofMillis(100L)).count
     }
     numConsumed
   }
@@ -247,7 +270,7 @@ abstract class QuotaTestClients(topic: String,
     leaderNode.metrics.metrics.get(metricName)
   }
 
-  def verifyProducerClientThrottleTimeMetric(expectThrottle: Boolean) {
+  def verifyProducerClientThrottleTimeMetric(expectThrottle: Boolean): Unit = {
     val tags = new HashMap[String, String]
     tags.put("client-id", producerClientId)
     val avgMetric = producer.metrics.get(new MetricName("produce-throttle-time-avg", "producer-metrics", "", tags))
@@ -260,7 +283,7 @@ abstract class QuotaTestClients(topic: String,
       assertEquals("Should not have been throttled", 0.0, metricValue(maxMetric), 0.0)
   }
 
-  def verifyConsumerClientThrottleTimeMetric(expectThrottle: Boolean, maxThrottleTime: Option[Double] = None) {
+  def verifyConsumerClientThrottleTimeMetric(expectThrottle: Boolean, maxThrottleTime: Option[Double] = None): Unit = {
     val tags = new HashMap[String, String]
     tags.put("client-id", consumerClientId)
     val avgMetric = consumer.metrics.get(new MetricName("fetch-throttle-time-avg", "consumer-fetch-manager-metrics", "", tags))
@@ -283,7 +306,7 @@ abstract class QuotaTestClients(topic: String,
     props
   }
 
-  def waitForQuotaUpdate(producerQuota: Long, consumerQuota: Long, requestQuota: Double, server: KafkaServer = leaderNode) {
+  def waitForQuotaUpdate(producerQuota: Long, consumerQuota: Long, requestQuota: Double, server: KafkaServer = leaderNode): Unit = {
     TestUtils.retry(10000) {
       val quotaManagers = server.dataPlaneRequestProcessor.quotas
       val overrideProducerQuota = quota(quotaManagers.produce, userPrincipal, producerClientId)
diff --git a/core/src/test/scala/integration/kafka/api/ClientIdQuotaTest.scala b/core/src/test/scala/integration/kafka/api/ClientIdQuotaTest.scala
index ef8fb413e5262..386c39d969655 100644
--- a/core/src/test/scala/integration/kafka/api/ClientIdQuotaTest.scala
+++ b/core/src/test/scala/integration/kafka/api/ClientIdQuotaTest.scala
@@ -27,7 +27,7 @@ class ClientIdQuotaTest extends BaseQuotaTest {
   override def consumerClientId = "QuotasTestConsumer-!@#$%^&*()"
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     this.serverConfig.setProperty(KafkaConfig.ProducerQuotaBytesPerSecondDefaultProp, defaultProducerQuota.toString)
     this.serverConfig.setProperty(KafkaConfig.ConsumerQuotaBytesPerSecondDefaultProp, defaultConsumerQuota.toString)
     super.setUp()
@@ -43,7 +43,7 @@ class ClientIdQuotaTest extends BaseQuotaTest {
         Map("user" -> "", "client-id" -> clientId)
       }
 
-      override def overrideQuotas(producerQuota: Long, consumerQuota: Long, requestQuota: Double) {
+      override def overrideQuotas(producerQuota: Long, consumerQuota: Long, requestQuota: Double): Unit = {
         val producerProps = new Properties()
         producerProps.put(DynamicConfig.Client.ProducerByteRateOverrideProp, producerQuota.toString)
         producerProps.put(DynamicConfig.Client.RequestPercentageOverrideProp, requestQuota.toString)
@@ -55,13 +55,13 @@ class ClientIdQuotaTest extends BaseQuotaTest {
         updateQuotaOverride(consumerClientId, consumerProps)
       }
 
-      override def removeQuotaOverrides() {
+      override def removeQuotaOverrides(): Unit = {
         val emptyProps = new Properties
         updateQuotaOverride(producerClientId, emptyProps)
         updateQuotaOverride(consumerClientId, emptyProps)
       }
 
-      private def updateQuotaOverride(clientId: String, properties: Properties) {
+      private def updateQuotaOverride(clientId: String, properties: Properties): Unit = {
         adminZkClient.changeClientIdConfig(Sanitizer.sanitize(clientId), properties)
       }
     }
diff --git a/core/src/test/scala/integration/kafka/api/ConsumerBounceTest.scala b/core/src/test/scala/integration/kafka/api/ConsumerBounceTest.scala
index e535104306fa4..6c3201b5dc0a4 100644
--- a/core/src/test/scala/integration/kafka/api/ConsumerBounceTest.scala
+++ b/core/src/test/scala/integration/kafka/api/ConsumerBounceTest.scala
@@ -15,39 +15,36 @@ package kafka.api
 
 import java.time
 import java.util.concurrent._
-import java.util.concurrent.atomic.AtomicBoolean
-import java.util.concurrent.locks.ReentrantLock
 import java.util.{Collection, Collections, Properties}
 
-import util.control.Breaks._
-import kafka.server.{BaseRequestTest, KafkaConfig}
+import kafka.server.KafkaConfig
 import kafka.utils.{CoreUtils, Logging, ShutdownableThread, TestUtils}
 import org.apache.kafka.clients.consumer._
 import org.apache.kafka.clients.producer.{KafkaProducer, ProducerRecord}
 import org.apache.kafka.common.TopicPartition
 import org.apache.kafka.common.errors.GroupMaxSizeReachedException
+import org.apache.kafka.common.message.FindCoordinatorRequestData
 import org.apache.kafka.common.protocol.{ApiKeys, Errors}
 import org.apache.kafka.common.requests.{FindCoordinatorRequest, FindCoordinatorResponse}
 import org.junit.Assert._
-import org.junit.{After, Before, Ignore, Test}
+import org.junit.{After, Ignore, Test}
 
 import scala.collection.JavaConverters._
-import scala.collection.mutable.ArrayBuffer
-import scala.concurrent.duration.Duration
-import scala.concurrent.{Await, ExecutionContext, ExecutionContextExecutor, Future => SFuture}
+import scala.collection.mutable
+import scala.collection.Seq
 
 /**
  * Integration tests for the consumer that cover basic usage as well as server failures
  */
-class ConsumerBounceTest extends BaseRequestTest with Logging {
-  val topic = "topic"
-  val part = 0
-  val tp = new TopicPartition(topic, part)
+class ConsumerBounceTest extends AbstractConsumerTest with Logging {
   val maxGroupSize = 5
 
   // Time to process commit and leave group requests in tests when brokers are available
-  val gracefulCloseTimeMs = 1000
-  val executor = Executors.newScheduledThreadPool(2)
+  val gracefulCloseTimeMs = Some(1000L)
+  val executor: ScheduledExecutorService = Executors.newScheduledThreadPool(2)
+  val consumerPollers: mutable.Buffer[ConsumerAssignmentPoller] = mutable.Buffer[ConsumerAssignmentPoller]()
+
+  this.consumerConfig.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "true")
 
   override def generateConfigs = {
     generateKafkaConfigs()
@@ -63,21 +60,14 @@ class ConsumerBounceTest extends BaseRequestTest with Logging {
     properties.put(KafkaConfig.UncleanLeaderElectionEnableProp, "true")
     properties.put(KafkaConfig.AutoCreateTopicsEnableProp, "false")
 
-    FixedPortTestUtils.createBrokerConfigs(numBrokers, zkConnect, enableControlledShutdown = false)
+    FixedPortTestUtils.createBrokerConfigs(brokerCount, zkConnect, enableControlledShutdown = false)
       .map(KafkaConfig.fromProps(_, properties))
   }
 
-  @Before
-  override def setUp() {
-    super.setUp()
-
-    // create the test topic with all the brokers as replicas
-    createTopic(topic, 1, numBrokers)
-  }
-
   @After
-  override def tearDown() {
+  override def tearDown(): Unit = {
     try {
+      consumerPollers.foreach(_.shutdown())
       executor.shutdownNow()
       // Wait for any active tasks to terminate to ensure consumer is not closed while being used from another thread
       assertTrue("Executor did not terminate", executor.awaitTermination(5000, TimeUnit.MILLISECONDS))
@@ -94,7 +84,7 @@ class ConsumerBounceTest extends BaseRequestTest with Logging {
    * 1. Produce a bunch of messages
    * 2. Then consume the messages while killing and restarting brokers at random
    */
-  def consumeWithBrokerFailures(numIters: Int) {
+  def consumeWithBrokerFailures(numIters: Int): Unit = {
     val numRecords = 1000
     val producer = createProducer()
     sendRecords(producer, numRecords)
@@ -118,7 +108,7 @@ class ConsumerBounceTest extends BaseRequestTest with Logging {
 
       if (records.nonEmpty) {
         consumer.commitSync()
-        assertEquals(consumer.position(tp), consumer.committed(tp).offset)
+        assertEquals(consumer.position(tp), consumer.committed(Set(tp).asJava).get(tp).offset)
 
         if (consumer.position(tp) == numRecords) {
           consumer.seekToBeginning(Collections.emptyList())
@@ -132,7 +122,7 @@ class ConsumerBounceTest extends BaseRequestTest with Logging {
   @Test
   def testSeekAndCommitWithBrokerFailures() = seekAndCommitWithBrokerFailures(5)
 
-  def seekAndCommitWithBrokerFailures(numIters: Int) {
+  def seekAndCommitWithBrokerFailures(numIters: Int): Unit = {
     val numRecords = 1000
     val producer = createProducer()
     sendRecords(producer, numRecords)
@@ -143,7 +133,7 @@ class ConsumerBounceTest extends BaseRequestTest with Logging {
 
     // wait until all the followers have synced the last HW with leader
     TestUtils.waitUntilTrue(() => servers.forall(server =>
-      server.replicaManager.localReplica(tp).get.highWatermark.messageOffset == numRecords
+      server.replicaManager.localLog(tp).get.highWatermark == numRecords
     ), "Failed to update high watermark for followers after timeout")
 
     val scheduler = new BounceBrokerScheduler(numIters)
@@ -163,26 +153,26 @@ class ConsumerBounceTest extends BaseRequestTest with Logging {
       } else if (coin == 2) {
         info("Committing offset.")
         consumer.commitSync()
-        assertEquals(consumer.position(tp), consumer.committed(tp).offset)
+        assertEquals(consumer.position(tp), consumer.committed(Set(tp).asJava).get(tp).offset)
       }
     }
   }
 
   @Test
-  def testSubscribeWhenTopicUnavailable() {
+  def testSubscribeWhenTopicUnavailable(): Unit = {
     val numRecords = 1000
     val newtopic = "newtopic"
 
     val consumer = createConsumer()
     consumer.subscribe(Collections.singleton(newtopic))
     executor.schedule(new Runnable {
-        def run() = createTopic(newtopic, numPartitions = numBrokers, replicationFactor = numBrokers)
+        def run() = createTopic(newtopic, numPartitions = brokerCount, replicationFactor = brokerCount)
       }, 2, TimeUnit.SECONDS)
-    consumer.poll(0)
+    consumer.poll(time.Duration.ZERO)
 
     val producer = createProducer()
 
-    def sendRecords(numRecords: Int, topic: String) {
+    def sendRecords(numRecords: Int, topic: String): Unit = {
       var remainingRecords = numRecords
       val endTimeMs = System.currentTimeMillis + 20000
       while (remainingRecords > 0 && System.currentTimeMillis < endTimeMs) {
@@ -201,22 +191,26 @@ class ConsumerBounceTest extends BaseRequestTest with Logging {
       assertEquals(0, remainingRecords)
     }
 
+    val poller = new ConsumerAssignmentPoller(consumer, List(newtopic))
+    consumerPollers += poller
+    poller.start()
     sendRecords(numRecords, newtopic)
-    receiveRecords(consumer, numRecords, 10000)
+    receiveExactRecords(poller, numRecords, 10000)
+    poller.shutdown()
 
     servers.foreach(server => killBroker(server.config.brokerId))
     Thread.sleep(500)
     restartDeadBrokers()
 
-    val future = executor.submit(new Runnable {
-      def run() = receiveRecords(consumer, numRecords, 10000)
-    })
+    val poller2 = new ConsumerAssignmentPoller(consumer, List(newtopic))
+    consumerPollers += poller2
+    poller2.start()
     sendRecords(numRecords, newtopic)
-    future.get
+    receiveExactRecords(poller, numRecords, 10000L)
   }
 
   @Test
-  def testClose() {
+  def testClose(): Unit = {
     val numRecords = 10
     val producer = createProducer()
     sendRecords(producer, numRecords)
@@ -231,9 +225,9 @@ class ConsumerBounceTest extends BaseRequestTest with Logging {
    * and leave group. New consumer instance should be able join group and start consuming from
    * last committed offset.
    */
-  private def checkCloseGoodPath(numRecords: Int, groupId: String) {
+  private def checkCloseGoodPath(numRecords: Int, groupId: String): Unit = {
     val consumer = createConsumerAndReceive(groupId, false, numRecords)
-    val future = submitCloseAndValidate(consumer, Long.MaxValue, None, Some(gracefulCloseTimeMs))
+    val future = submitCloseAndValidate(consumer, Long.MaxValue, None, gracefulCloseTimeMs)
     future.get
     checkClosedState(groupId, numRecords)
   }
@@ -244,15 +238,17 @@ class ConsumerBounceTest extends BaseRequestTest with Logging {
    * Close of consumers using manual assignment should complete with successful commits since a
    * broker is available.
    */
-  private def checkCloseWithCoordinatorFailure(numRecords: Int, dynamicGroup: String, manualGroup: String) {
+  private def checkCloseWithCoordinatorFailure(numRecords: Int, dynamicGroup: String, manualGroup: String): Unit = {
     val consumer1 = createConsumerAndReceive(dynamicGroup, false, numRecords)
     val consumer2 = createConsumerAndReceive(manualGroup, true, numRecords)
 
     killBroker(findCoordinator(dynamicGroup))
     killBroker(findCoordinator(manualGroup))
 
-    val future1 = submitCloseAndValidate(consumer1, Long.MaxValue, None, Some(gracefulCloseTimeMs))
-    val future2 = submitCloseAndValidate(consumer2, Long.MaxValue, None, Some(gracefulCloseTimeMs))
+    val future1 = submitCloseAndValidate(consumer1, Long.MaxValue, None, gracefulCloseTimeMs)
+
+    val future2 = submitCloseAndValidate(consumer2, Long.MaxValue, None, gracefulCloseTimeMs)
+
     future1.get
     future2.get
 
@@ -262,7 +258,10 @@ class ConsumerBounceTest extends BaseRequestTest with Logging {
   }
 
   private def findCoordinator(group: String): Int = {
-    val request = new FindCoordinatorRequest.Builder(FindCoordinatorRequest.CoordinatorType.GROUP, group).build()
+    val request = new FindCoordinatorRequest.Builder(
+        new FindCoordinatorRequestData()
+          .setKeyType(FindCoordinatorRequest.CoordinatorType.GROUP.id)
+          .setKey(group)).build()
     var nodeId = -1
     TestUtils.waitUntilTrue(() => {
       val resp = connectAndSend(request, ApiKeys.FIND_COORDINATOR)
@@ -278,7 +277,7 @@ class ConsumerBounceTest extends BaseRequestTest with Logging {
    * there is no coordinator, but close should timeout and return. If close is invoked with a very
    * large timeout, close should timeout after request timeout.
    */
-  private def checkCloseWithClusterFailure(numRecords: Int, group1: String, group2: String) {
+  private def checkCloseWithClusterFailure(numRecords: Int, group1: String, group2: String): Unit = {
     val consumer1 = createConsumerAndReceive(group1, false, numRecords)
 
     val requestTimeout = 6000
@@ -289,7 +288,7 @@ class ConsumerBounceTest extends BaseRequestTest with Logging {
 
     servers.foreach(server => killBroker(server.config.brokerId))
     val closeTimeout = 2000
-    val future1 = submitCloseAndValidate(consumer1, closeTimeout, Some(closeTimeout), Some(closeTimeout))
+    val future1 = submitCloseAndValidate(consumer1, closeTimeout, None, Some(closeTimeout))
     val future2 = submitCloseAndValidate(consumer2, Long.MaxValue, Some(requestTimeout), Some(requestTimeout))
     future1.get
     future2.get
@@ -302,86 +301,39 @@ class ConsumerBounceTest extends BaseRequestTest with Logging {
     */
   @Test
   def testRollingBrokerRestartsWithSmallerMaxGroupSizeConfigDisruptsBigGroup(): Unit = {
+    val group = "group-max-size-test"
     val topic = "group-max-size-test"
     val maxGroupSize = 2
     val consumerCount = maxGroupSize + 1
-    var recordsProduced = maxGroupSize * 100
     val partitionCount = consumerCount * 2
-    if (recordsProduced % partitionCount != 0) {
-      // ensure even record distribution per partition
-      recordsProduced += partitionCount - recordsProduced % partitionCount
-    }
-    val executor = Executors.newScheduledThreadPool(consumerCount * 2)
+
     this.consumerConfig.setProperty(ConsumerConfig.MAX_POLL_INTERVAL_MS_CONFIG, "60000")
     this.consumerConfig.setProperty(ConsumerConfig.HEARTBEAT_INTERVAL_MS_CONFIG, "1000")
     this.consumerConfig.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "false")
-    val producer = createProducer()
-    createTopic(topic, numPartitions = partitionCount, replicationFactor = numBrokers)
-    val stableConsumers = createConsumersWithGroupId("group2", consumerCount, executor, topic = topic)
+    val partitions = createTopicPartitions(topic, numPartitions = partitionCount, replicationFactor = brokerCount)
 
-    // assert group is stable and working
-    sendRecords(producer, recordsProduced, topic, numPartitions = Some(partitionCount))
-    stableConsumers.foreach { cons => {
-      receiveAndCommit(cons, recordsProduced / consumerCount, 10000)
-    }}
+    addConsumersToGroupAndWaitForGroupAssignment(consumerCount, mutable.Buffer[KafkaConsumer[Array[Byte], Array[Byte]]](),
+      consumerPollers, List[String](topic), partitions, group)
 
     // roll all brokers with a lesser max group size to make sure coordinator has the new config
     val newConfigs = generateKafkaConfigs(maxGroupSize.toString)
-    val kickedConsumerOut = new AtomicBoolean(false)
-    var kickedOutConsumerIdx: Option[Int] = None
-    val lock = new ReentrantLock
-    // restart brokers until the group moves to a Coordinator with the new config
-    breakable { for (broker <- servers.indices) {
-      killBroker(broker)
-      sendRecords(producer, recordsProduced, topic, numPartitions = Some(partitionCount))
-
-      var successfulConsumes = 0
-
-      // compute consumptions in a non-blocking way in order to account for the rebalance once the group.size takes effect
-      val consumeFutures = new ArrayBuffer[SFuture[Any]]
-      implicit val executorContext: ExecutionContextExecutor = ExecutionContext.fromExecutor(executor)
-      stableConsumers.indices.foreach(idx => {
-        val currentConsumer = stableConsumers(idx)
-        val consumeFuture = SFuture {
-          try {
-            receiveAndCommit(currentConsumer, recordsProduced / consumerCount, 10000)
-            CoreUtils.inLock(lock) { successfulConsumes += 1 }
-          } catch {
-            case e: Throwable =>
-              if (!e.isInstanceOf[GroupMaxSizeReachedException]) {
-                throw e
-              }
-              if (!kickedConsumerOut.compareAndSet(false, true)) {
-                fail(s"Received more than one ${classOf[GroupMaxSizeReachedException]}")
-              }
-              kickedOutConsumerIdx = Some(idx)
-          }
-        }
+    for (serverIdx <- servers.indices) {
+      killBroker(serverIdx)
+      val config = newConfigs(serverIdx)
+      servers(serverIdx) = TestUtils.createServer(config, time = brokerTime(config.brokerId))
+      restartDeadBrokers()
+    }
 
-        consumeFutures += consumeFuture
-      })
-      Await.result(SFuture.sequence(consumeFutures), Duration("12sec"))
+    def raisedExceptions: Seq[Throwable] = {
+      consumerPollers.flatten(_.thrownException)
+    }
 
-      if (kickedConsumerOut.get()) {
-        // validate the rest N-1 consumers consumed successfully
-        assertEquals(maxGroupSize, successfulConsumes)
-        break
-      }
+    // we are waiting for the group to rebalance and one member to get kicked
+    TestUtils.waitUntilTrue(() => raisedExceptions.nonEmpty,
+      msg = "The remaining consumers in the group could not fetch the expected records", 10000L)
 
-      val config = newConfigs(broker)
-      servers(broker) = TestUtils.createServer(config, time = brokerTime(config.brokerId))
-      restartDeadBrokers()
-    }}
-    if (!kickedConsumerOut.get())
-      fail(s"Should have received an ${classOf[GroupMaxSizeReachedException]} during the cluster roll")
-
-    // assert that the group has gone through a rebalance and shed off one consumer
-    stableConsumers.remove(kickedOutConsumerIdx.get)
-    sendRecords(producer, recordsProduced, topic, numPartitions = Some(partitionCount))
-    // should be only maxGroupSize consumers left in the group
-    stableConsumers.foreach { cons => {
-      receiveAndCommit(cons, recordsProduced / maxGroupSize, 10000)
-    }}
+    assertEquals(1, raisedExceptions.size)
+    assertTrue(raisedExceptions.head.isInstanceOf[GroupMaxSizeReachedException])
   }
 
   /**
@@ -389,70 +341,30 @@ class ConsumerBounceTest extends BaseRequestTest with Logging {
     */
   @Test
   def testConsumerReceivesFatalExceptionWhenGroupPassesMaxSize(): Unit = {
-    val topic = "group-max-size-test"
-    val groupId = "group1"
-    val executor = Executors.newScheduledThreadPool(maxGroupSize * 2)
-    createTopic(topic, maxGroupSize, numBrokers)
+    val group = "fatal-exception-test"
+    val topic = "fatal-exception-test"
     this.consumerConfig.setProperty(ConsumerConfig.MAX_POLL_INTERVAL_MS_CONFIG, "60000")
     this.consumerConfig.setProperty(ConsumerConfig.HEARTBEAT_INTERVAL_MS_CONFIG, "1000")
     this.consumerConfig.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "false")
 
+    val partitions = createTopicPartitions(topic, numPartitions = maxGroupSize, replicationFactor = brokerCount)
+
     // Create N+1 consumers in the same consumer group and assert that the N+1th consumer receives a fatal error when it tries to join the group
-    val stableConsumers = createConsumersWithGroupId(groupId, maxGroupSize, executor, topic)
-    val newConsumer = createConsumerWithGroupId(groupId)
-    var failedRebalance = false
-    var exception: Exception = null
-    waitForRebalance(5000, subscribeAndPoll(newConsumer, executor = executor, onException = e => {failedRebalance = true; exception = e}),
-      executor = executor, stableConsumers:_*)
-    assertTrue("Rebalance did not fail as expected", failedRebalance)
-    assertTrue(exception.isInstanceOf[GroupMaxSizeReachedException])
+    addConsumersToGroupAndWaitForGroupAssignment(maxGroupSize, mutable.Buffer[KafkaConsumer[Array[Byte], Array[Byte]]](),
+      consumerPollers, List[String](topic), partitions, group)
+    val (_, rejectedConsumerPollers) = addConsumersToGroup(1,
+      mutable.Buffer[KafkaConsumer[Array[Byte], Array[Byte]]](), mutable.Buffer[ConsumerAssignmentPoller](), List[String](topic), partitions, group)
+    val rejectedConsumer = rejectedConsumerPollers.head
+    TestUtils.waitUntilTrue(() => {
+      rejectedConsumer.thrownException.isDefined
+    }, "Extra consumer did not throw an exception")
+    assertTrue(rejectedConsumer.thrownException.get.isInstanceOf[GroupMaxSizeReachedException])
 
     // assert group continues to live
-    val producer = createProducer()
-    sendRecords(producer, maxGroupSize * 100, topic, numPartitions = Some(maxGroupSize))
-    stableConsumers.foreach { cons => {
-        receiveExactRecords(cons, 100, 10000)
-    }}
-  }
-
-  /**
-    * Creates N consumers with the same group ID and ensures the group rebalances properly at each step
-    */
-  private def createConsumersWithGroupId(groupId: String, consumerCount: Int, executor: ExecutorService, topic: String): ArrayBuffer[KafkaConsumer[Array[Byte], Array[Byte]]] = {
-    val stableConsumers = ArrayBuffer[KafkaConsumer[Array[Byte], Array[Byte]]]()
-    for (_ <- 1.to(consumerCount)) {
-      val newConsumer = createConsumerWithGroupId(groupId)
-      waitForRebalance(5000, subscribeAndPoll(newConsumer, executor = executor, topic = topic),
-        executor = executor, stableConsumers:_*)
-      stableConsumers += newConsumer
-    }
-    stableConsumers
-  }
-
-  def subscribeAndPoll(consumer: KafkaConsumer[Array[Byte], Array[Byte]], executor: ExecutorService, revokeSemaphore: Option[Semaphore] = None,
-                       onException: Exception => Unit = e => { throw e }, topic: String = topic, pollTimeout: Int = 1000): Future[Any] = {
-    executor.submit(CoreUtils.runnable {
-      try {
-        consumer.subscribe(Collections.singletonList(topic))
-        consumer.poll(java.time.Duration.ofMillis(pollTimeout))
-      } catch {
-        case e: Exception => onException.apply(e)
-      }
-    }, 0)
-  }
-
-  def waitForRebalance(timeoutMs: Long, future: Future[Any], executor: ExecutorService, otherConsumers: KafkaConsumer[Array[Byte], Array[Byte]]*) {
-    val startMs = System.currentTimeMillis
-    implicit val executorContext: ExecutionContextExecutor = ExecutionContext.fromExecutor(executor)
-
-    while (System.currentTimeMillis < startMs + timeoutMs && !future.isDone) {
-      val consumeFutures = otherConsumers.map(consumer => SFuture {
-        consumer.poll(time.Duration.ofMillis(1000))
-      })
-      Await.result(SFuture.sequence(consumeFutures), Duration("1500ms"))
-    }
-
-    assertTrue("Rebalance did not complete in time", future.isDone)
+    sendRecords(createProducer(), maxGroupSize * 100, topic, numPartitions = Some(partitions.size))
+    TestUtils.waitUntilTrue(() => {
+      consumerPollers.forall(p => p.receivedMessages >= 100)
+    }, "The consumers in the group could not fetch the expected records", 10000L)
   }
 
   /**
@@ -461,34 +373,30 @@ class ConsumerBounceTest extends BaseRequestTest with Logging {
    * close should terminate immediately without sending leave group.
    */
   @Test
-  def testCloseDuringRebalance() {
+  def testCloseDuringRebalance(): Unit = {
     val topic = "closetest"
-    createTopic(topic, 10, numBrokers)
+    createTopic(topic, 10, brokerCount)
     this.consumerConfig.setProperty(ConsumerConfig.MAX_POLL_INTERVAL_MS_CONFIG, "60000")
     this.consumerConfig.setProperty(ConsumerConfig.HEARTBEAT_INTERVAL_MS_CONFIG, "1000")
     this.consumerConfig.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "false")
     checkCloseDuringRebalance("group1", topic, executor, true)
   }
 
-  private def checkCloseDuringRebalance(groupId: String, topic: String, executor: ExecutorService, brokersAvailableDuringClose: Boolean) {
+  private def checkCloseDuringRebalance(groupId: String, topic: String, executor: ExecutorService, brokersAvailableDuringClose: Boolean): Unit = {
 
     def subscribeAndPoll(consumer: KafkaConsumer[Array[Byte], Array[Byte]], revokeSemaphore: Option[Semaphore] = None): Future[Any] = {
       executor.submit(CoreUtils.runnable {
-          consumer.subscribe(Collections.singletonList(topic), new ConsumerRebalanceListener {
-            def onPartitionsAssigned(partitions: Collection[TopicPartition]) {
-            }
-            def onPartitionsRevoked(partitions: Collection[TopicPartition]) {
-              revokeSemaphore.foreach(s => s.release())
-            }
-          })
-          consumer.poll(0)
+        consumer.subscribe(Collections.singletonList(topic))
+        revokeSemaphore.foreach(s => s.release())
+        // requires to used deprecated `poll(long)` to trigger metadata update
+          consumer.poll(0L)
         }, 0)
     }
 
-    def waitForRebalance(timeoutMs: Long, future: Future[Any], otherConsumers: KafkaConsumer[Array[Byte], Array[Byte]]*) {
+    def waitForRebalance(timeoutMs: Long, future: Future[Any], otherConsumers: KafkaConsumer[Array[Byte], Array[Byte]]*): Unit = {
       val startMs = System.currentTimeMillis
       while (System.currentTimeMillis < startMs + timeoutMs && !future.isDone)
-          otherConsumers.foreach(consumer => consumer.poll(100))
+          otherConsumers.foreach(consumer => consumer.poll(time.Duration.ofMillis(100L)))
       assertTrue("Rebalance did not complete in time", future.isDone)
     }
 
@@ -509,7 +417,7 @@ class ConsumerBounceTest extends BaseRequestTest with Logging {
     val rebalanceFuture = createConsumerToRebalance()
 
     // consumer1 should leave group and close immediately even though rebalance is in progress
-    val closeFuture1 = submitCloseAndValidate(consumer1, Long.MaxValue, None, Some(gracefulCloseTimeMs))
+    val closeFuture1 = submitCloseAndValidate(consumer1, Long.MaxValue, None, gracefulCloseTimeMs)
 
     // Rebalance should complete without waiting for consumer1 to timeout since consumer1 has left the group
     waitForRebalance(2000, rebalanceFuture, consumer2)
@@ -527,50 +435,32 @@ class ConsumerBounceTest extends BaseRequestTest with Logging {
     closeFuture2.get(2000, TimeUnit.MILLISECONDS)
   }
 
-  private def createConsumerWithGroupId(groupId: String): KafkaConsumer[Array[Byte], Array[Byte]] = {
-    consumerConfig.setProperty(ConsumerConfig.GROUP_ID_CONFIG, groupId)
-    createConsumer()
-  }
-
   private def createConsumerAndReceive(groupId: String, manualAssign: Boolean, numRecords: Int): KafkaConsumer[Array[Byte], Array[Byte]] = {
     val consumer = createConsumerWithGroupId(groupId)
-    if (manualAssign)
-      consumer.assign(Collections.singleton(tp))
-    else
-      consumer.subscribe(Collections.singleton(topic))
-    receiveExactRecords(consumer, numRecords)
-    consumer
-  }
-
-  private def receiveRecords(consumer: KafkaConsumer[Array[Byte], Array[Byte]], numRecords: Int, timeoutMs: Long = 60000): Long = {
-    var received = 0L
-    val endTimeMs = System.currentTimeMillis + timeoutMs
-    while (received < numRecords && System.currentTimeMillis < endTimeMs)
-      received += consumer.poll(time.Duration.ofMillis(100)).count()
-
-    received
-  }
+    val consumerPoller = if (manualAssign)
+        subscribeConsumerAndStartPolling(consumer, List(), Set(tp))
+      else
+        subscribeConsumerAndStartPolling(consumer, List(topic))
 
-  private def receiveExactRecords(consumer: KafkaConsumer[Array[Byte], Array[Byte]], numRecords: Int, timeoutMs: Long = 60000): Unit = {
-    val received = receiveRecords(consumer, numRecords, timeoutMs)
-    assertEquals(numRecords, received)
+    receiveExactRecords(consumerPoller, numRecords)
+    consumerPoller.shutdown()
+    consumer
   }
 
-  @throws(classOf[CommitFailedException])
-  private def receiveAndCommit(consumer: KafkaConsumer[Array[Byte], Array[Byte]], numRecords: Int, timeoutMs: Long): Unit = {
-    val received = receiveRecords(consumer, numRecords, timeoutMs)
-    assertTrue(s"Received $received, expected at least $numRecords", numRecords <= received)
-    consumer.commitSync()
+  private def receiveExactRecords(consumer: ConsumerAssignmentPoller, numRecords: Int, timeoutMs: Long = 60000): Unit = {
+    TestUtils.waitUntilTrue(() => {
+      consumer.receivedMessages == numRecords
+    }, s"Consumer did not receive expected $numRecords. It received ${consumer.receivedMessages}", timeoutMs)
   }
 
   private def submitCloseAndValidate(consumer: KafkaConsumer[Array[Byte], Array[Byte]],
       closeTimeoutMs: Long, minCloseTimeMs: Option[Long], maxCloseTimeMs: Option[Long]): Future[Any] = {
     executor.submit(CoreUtils.runnable {
       val closeGraceTimeMs = 2000
-      val startNanos = System.nanoTime
+      val startMs = System.currentTimeMillis()
       info("Closing consumer with timeout " + closeTimeoutMs + " ms.")
-      consumer.close(closeTimeoutMs, TimeUnit.MILLISECONDS)
-      val timeTakenMs = TimeUnit.NANOSECONDS.toMillis(System.nanoTime - startNanos)
+      consumer.close(time.Duration.ofMillis(closeTimeoutMs))
+      val timeTakenMs = System.currentTimeMillis() - startMs
       maxCloseTimeMs.foreach { ms =>
         assertTrue("Close took too long " + timeTakenMs, timeTakenMs < ms + closeGraceTimeMs)
       }
@@ -581,21 +471,21 @@ class ConsumerBounceTest extends BaseRequestTest with Logging {
     }, 0)
   }
 
-  private def checkClosedState(groupId: String, committedRecords: Int) {
+  private def checkClosedState(groupId: String, committedRecords: Int): Unit = {
     // Check that close was graceful with offsets committed and leave group sent.
     // New instance of consumer should be assigned partitions immediately and should see committed offsets.
     val assignSemaphore = new Semaphore(0)
     val consumer = createConsumerWithGroupId(groupId)
     consumer.subscribe(Collections.singletonList(topic),  new ConsumerRebalanceListener {
-      def onPartitionsAssigned(partitions: Collection[TopicPartition]) {
+      def onPartitionsAssigned(partitions: Collection[TopicPartition]): Unit = {
         assignSemaphore.release()
       }
-      def onPartitionsRevoked(partitions: Collection[TopicPartition]) {
+      def onPartitionsRevoked(partitions: Collection[TopicPartition]): Unit = {
       }})
-    consumer.poll(3000)
+    consumer.poll(time.Duration.ofSeconds(3L))
     assertTrue("Assignment did not complete on time", assignSemaphore.tryAcquire(1, TimeUnit.SECONDS))
     if (committedRecords > 0)
-      assertEquals(committedRecords, consumer.committed(tp).offset)
+      assertEquals(committedRecords, consumer.committed(Set(tp).asJava).get(tp).offset)
     consumer.close()
   }
 
@@ -616,10 +506,16 @@ class ConsumerBounceTest extends BaseRequestTest with Logging {
     }
   }
 
+  private def createTopicPartitions(topic: String, numPartitions: Int, replicationFactor: Int,
+                                    topicConfig: Properties = new Properties): Set[TopicPartition] = {
+    createTopic(topic, numPartitions = numPartitions, replicationFactor = replicationFactor, topicConfig = topicConfig)
+    Range(0, numPartitions).map(part => new TopicPartition(topic, part)).toSet
+  }
+
   private def sendRecords(producer: KafkaProducer[Array[Byte], Array[Byte]],
                           numRecords: Int,
                           topic: String = this.topic,
-                          numPartitions: Option[Int] = None) {
+                          numPartitions: Option[Int] = None): Unit = {
     var partitionIndex = 0
     def getPartition: Int = {
       numPartitions match {
diff --git a/core/src/test/scala/integration/kafka/api/ConsumerTopicCreationTest.scala b/core/src/test/scala/integration/kafka/api/ConsumerTopicCreationTest.scala
new file mode 100644
index 0000000000000..a126c8397cbf5
--- /dev/null
+++ b/core/src/test/scala/integration/kafka/api/ConsumerTopicCreationTest.scala
@@ -0,0 +1,98 @@
+/**
+ * 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 java.lang.{Boolean => JBoolean}
+import java.time.Duration
+import java.util
+import java.util.Collections
+
+import kafka.server.KafkaConfig
+import kafka.utils.TestUtils
+import org.apache.kafka.clients.admin.NewTopic
+import org.apache.kafka.clients.consumer.ConsumerConfig
+import org.apache.kafka.clients.producer.{ProducerConfig, ProducerRecord}
+import org.junit.Assert._
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.Parameterized
+import org.junit.runners.Parameterized.Parameters
+
+/**
+ * Tests behavior of specifying auto topic creation configuration for the consumer and broker
+ */
+@RunWith(value = classOf[Parameterized])
+class ConsumerTopicCreationTest(brokerAutoTopicCreationEnable: JBoolean, consumerAllowAutoCreateTopics: JBoolean) extends IntegrationTestHarness {
+  override protected def brokerCount: Int = 1
+
+  val topic_1 = "topic-1"
+  val topic_2 = "topic-2"
+  val producerClientId = "ConsumerTestProducer"
+  val consumerClientId = "ConsumerTestConsumer"
+
+  // configure server properties
+  this.serverConfig.setProperty(KafkaConfig.ControlledShutdownEnableProp, "false") // speed up shutdown
+  this.serverConfig.setProperty(KafkaConfig.AutoCreateTopicsEnableProp, brokerAutoTopicCreationEnable.toString)
+
+  // configure client properties
+  this.producerConfig.setProperty(ProducerConfig.CLIENT_ID_CONFIG, producerClientId)
+  this.consumerConfig.setProperty(ConsumerConfig.CLIENT_ID_CONFIG, consumerClientId)
+  this.consumerConfig.setProperty(ConsumerConfig.GROUP_ID_CONFIG, "my-test")
+  this.consumerConfig.setProperty(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest")
+  this.consumerConfig.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "false")
+  this.consumerConfig.setProperty(ConsumerConfig.METADATA_MAX_AGE_CONFIG, "100")
+  this.consumerConfig.setProperty(ConsumerConfig.ALLOW_AUTO_CREATE_TOPICS_CONFIG, consumerAllowAutoCreateTopics.toString)
+
+  @Test
+  def testAutoTopicCreation(): Unit = {
+    val consumer = createConsumer()
+    val producer = createProducer()
+    val adminClient = createAdminClient()
+    val record = new ProducerRecord(topic_1, 0, "key".getBytes, "value".getBytes)
+
+    // create `topic_1` and produce a record to it
+    adminClient.createTopics(Collections.singleton(new NewTopic(topic_1, 1, 1.toShort))).all.get
+    producer.send(record).get
+
+    consumer.subscribe(util.Arrays.asList(topic_1, topic_2))
+
+    // Wait until the produced record was consumed. This guarantees that metadata request for `topic_2` was sent to the
+    // broker.
+    TestUtils.waitUntilTrue(() => {
+      consumer.poll(Duration.ofMillis(100)).count > 0
+    }, "Timed out waiting to consume")
+
+    // MetadataRequest is guaranteed to create the topic znode if creation was required
+    val topicCreated = zkClient.getAllTopicsInCluster.contains(topic_2)
+    if (brokerAutoTopicCreationEnable && consumerAllowAutoCreateTopics)
+      assertTrue(topicCreated)
+    else
+      assertFalse(topicCreated)
+  }
+}
+
+object ConsumerTopicCreationTest {
+  @Parameters(name = "brokerTopicCreation={0}, consumerTopicCreation={1}")
+  def parameters: java.util.Collection[Array[Object]] = {
+    val data = new java.util.ArrayList[Array[Object]]()
+    for (brokerAutoTopicCreationEnable <- Array(JBoolean.TRUE, JBoolean.FALSE))
+      for (consumerAutoCreateTopicsPolicy <- Array(JBoolean.TRUE, JBoolean.FALSE))
+        data.add(Array(brokerAutoTopicCreationEnable, consumerAutoCreateTopicsPolicy))
+    data
+  }
+}
diff --git a/core/src/test/scala/integration/kafka/api/CustomQuotaCallbackTest.scala b/core/src/test/scala/integration/kafka/api/CustomQuotaCallbackTest.scala
index 8c1d34ddc2c98..09b163a2c3973 100644
--- a/core/src/test/scala/integration/kafka/api/CustomQuotaCallbackTest.scala
+++ b/core/src/test/scala/integration/kafka/api/CustomQuotaCallbackTest.scala
@@ -18,7 +18,7 @@ import java.io.File
 import java.{lang, util}
 import java.util.concurrent.ConcurrentHashMap
 import java.util.concurrent.atomic.{AtomicBoolean, AtomicInteger}
-import java.util.{Collections, Properties}
+import java.util.Properties
 
 import kafka.api.GroupedUserPrincipalBuilder._
 import kafka.api.GroupedUserQuotaCallback._
@@ -26,7 +26,7 @@ import kafka.server._
 import kafka.utils.JaasTestUtils.ScramLoginModule
 import kafka.utils.{JaasTestUtils, Logging, TestUtils}
 import kafka.zk.ConfigEntityChangeNotificationZNode
-import org.apache.kafka.clients.admin.{AdminClient, AdminClientConfig}
+import org.apache.kafka.clients.admin.{Admin, AdminClient, AdminClientConfig}
 import org.apache.kafka.clients.consumer.{ConsumerConfig, KafkaConsumer}
 import org.apache.kafka.clients.producer.{KafkaProducer, ProducerConfig, ProducerRecord}
 import org.apache.kafka.common.{Cluster, Reconfigurable}
@@ -48,13 +48,13 @@ class CustomQuotaCallbackTest extends IntegrationTestHarness with SaslSetup {
   override protected def interBrokerListenerName: ListenerName = new ListenerName("BROKER")
 
   override protected lazy val trustStoreFile = Some(File.createTempFile("truststore", ".jks"))
-  override val serverCount: Int = 2
+  override val brokerCount: Int = 2
 
   private val kafkaServerSaslMechanisms = Seq("SCRAM-SHA-256")
   private val kafkaClientSaslMechanism = "SCRAM-SHA-256"
   override protected val serverSaslProperties = Some(kafkaServerSaslProperties(kafkaServerSaslMechanisms, kafkaClientSaslMechanism))
   override protected val clientSaslProperties = Some(kafkaClientSaslProperties(kafkaClientSaslMechanism))
-  private val adminClients = new ArrayBuffer[AdminClient]()
+  private val adminClients = new ArrayBuffer[Admin]()
   private var producerWithoutQuota: KafkaProducer[Array[Byte], Array[Byte]] = _
 
   val defaultRequestQuota = 1000
@@ -62,7 +62,7 @@ class CustomQuotaCallbackTest extends IntegrationTestHarness with SaslSetup {
   val defaultConsumeQuota = 1000 * 1000 * 1000
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     startSasl(jaasSections(kafkaServerSaslMechanisms, Some("SCRAM-SHA-256"), KafkaSasl, JaasTestUtils.KafkaServerContextName))
     this.serverConfig.setProperty(KafkaConfig.ProducerQuotaBytesPerSecondDefaultProp, Long.MaxValue.toString)
     this.serverConfig.setProperty(KafkaConfig.ConsumerQuotaBytesPerSecondDefaultProp, Long.MaxValue.toString)
@@ -85,14 +85,14 @@ class CustomQuotaCallbackTest extends IntegrationTestHarness with SaslSetup {
     super.tearDown()
   }
 
-  override def configureSecurityBeforeServersStart() {
+  override def configureSecurityBeforeServersStart(): Unit = {
     super.configureSecurityBeforeServersStart()
     zkClient.makeSurePersistentPathExists(ConfigEntityChangeNotificationZNode.path)
     createScramCredentials(zkConnect, JaasTestUtils.KafkaScramAdmin, JaasTestUtils.KafkaScramAdminPassword)
   }
 
   @Test
-  def testCustomQuotaCallback() {
+  def testCustomQuotaCallback(): Unit = {
     // Large quota override, should not throttle
     var brokerId = 0
     var user = createGroupWithOneUser("group0_user1", brokerId)
@@ -100,9 +100,11 @@ class CustomQuotaCallbackTest extends IntegrationTestHarness with SaslSetup {
     quotaLimitCalls.values.foreach(_.set(0))
     user.produceConsume(expectProduceThrottle = false, expectConsumeThrottle = false)
 
-    // ClientQuotaCallback#quotaLimit is invoked by each quota manager once for each new client
+    // ClientQuotaCallback#quotaLimit is invoked by each quota manager once per throttled produce request for each client
     assertEquals(1, quotaLimitCalls(ClientQuotaType.PRODUCE).get)
-    assertEquals(1, quotaLimitCalls(ClientQuotaType.FETCH).get)
+    // ClientQuotaCallback#quotaLimit is invoked once per each unthrottled and two for each throttled request
+    // since we don't know the total number of requests, we verify it was called at least twice (at least one throttled request)
+    assertTrue("quotaLimit must be called at least twice", quotaLimitCalls(ClientQuotaType.FETCH).get > 2)
     assertTrue(s"Too many quotaLimit calls $quotaLimitCalls", quotaLimitCalls(ClientQuotaType.REQUEST).get <= 10) // sanity check
     // Large quota updated to small quota, should throttle
     user.configureAndWaitForQuota(9000, 3000)
@@ -161,7 +163,7 @@ class CustomQuotaCallbackTest extends IntegrationTestHarness with SaslSetup {
     user.waitForQuotaUpdate(8000, 2500, defaultRequestQuota)
     user.produceConsume(expectProduceThrottle = true, expectConsumeThrottle = true)
 
-    assertEquals(serverCount, callbackInstances.get)
+    assertEquals(brokerCount, callbackInstances.get)
   }
 
   /**
@@ -181,7 +183,7 @@ class CustomQuotaCallbackTest extends IntegrationTestHarness with SaslSetup {
     TestUtils.createTopic(zkClient, topic, assignment, servers)
   }
 
-  private def createAdminClient(): AdminClient = {
+  private def createAdminClient(): Admin = {
     val config = new util.HashMap[String, Object]
     config.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG,
       TestUtils.bootstrapServers(servers, new ListenerName("BROKER")))
@@ -315,7 +317,7 @@ object GroupedUserQuotaCallback {
   val QuotaGroupTag = "group"
   val DefaultProduceQuotaProp = "default.produce.quota"
   val DefaultFetchQuotaProp = "default.fetch.quota"
-  val UnlimitedQuotaMetricTags = Collections.emptyMap[String, String]
+  val UnlimitedQuotaMetricTags = new util.HashMap[String, String]
   val quotaLimitCalls = Map(
     ClientQuotaType.PRODUCE -> new AtomicInteger,
     ClientQuotaType.FETCH -> new AtomicInteger,
@@ -344,10 +346,8 @@ object GroupedUserQuotaCallback {
 class GroupedUserQuotaCallback extends ClientQuotaCallback with Reconfigurable with Logging {
 
   var brokerId: Int = -1
-  val customQuotasUpdated = ClientQuotaType.values.toList
-    .map(quotaType =>(quotaType -> new AtomicBoolean)).toMap
-  val quotas = ClientQuotaType.values.toList
-    .map(quotaType => (quotaType -> new ConcurrentHashMap[String, Double])).toMap
+  val customQuotasUpdated = ClientQuotaType.values.map(quotaType => quotaType -> new AtomicBoolean).toMap
+  val quotas = ClientQuotaType.values.map(quotaType => quotaType -> new ConcurrentHashMap[String, Double]).toMap
 
   val partitionRatio = new ConcurrentHashMap[String, Double]()
 
@@ -398,7 +398,7 @@ class GroupedUserQuotaCallback extends ClientQuotaCallback with Reconfigurable w
   override def updateClusterMetadata(cluster: Cluster): Boolean = {
     val topicsByGroup = cluster.topics.asScala.groupBy(group)
 
-    !topicsByGroup.forall { case (group, groupTopics) =>
+    topicsByGroup.map { case (group, groupTopics) =>
       val groupPartitions = groupTopics.flatMap(topic => cluster.partitionsForTopic(topic).asScala)
       val totalPartitions = groupPartitions.size
       val partitionsOnThisBroker = groupPartitions.count { p => p.leader != null && p.leader.id == brokerId }
@@ -409,7 +409,7 @@ class GroupedUserQuotaCallback extends ClientQuotaCallback with Reconfigurable w
       else
         partitionsOnThisBroker.toDouble / totalPartitions
       partitionRatio.put(group, multiplier) != multiplier
-    }
+    }.exists(identity)
   }
 
   override def updateQuota(quotaType: ClientQuotaType, quotaEntity: ClientQuotaEntity, newValue: Double): Unit = {
diff --git a/core/src/test/scala/integration/kafka/api/DelegationTokenEndToEndAuthorizationTest.scala b/core/src/test/scala/integration/kafka/api/DelegationTokenEndToEndAuthorizationTest.scala
index 29c7507449be3..2f1d51cc34f8b 100644
--- a/core/src/test/scala/integration/kafka/api/DelegationTokenEndToEndAuthorizationTest.scala
+++ b/core/src/test/scala/integration/kafka/api/DelegationTokenEndToEndAuthorizationTest.scala
@@ -24,6 +24,7 @@ import kafka.zk.ConfigEntityChangeNotificationZNode
 import org.apache.kafka.clients.admin.{AdminClient, AdminClientConfig}
 import org.apache.kafka.common.config.SaslConfigs
 import org.apache.kafka.common.security.auth.SecurityProtocol
+import org.apache.kafka.common.security.scram.ScramCredential
 import org.apache.kafka.common.security.scram.internals.ScramMechanism
 import org.apache.kafka.common.security.token.delegation.DelegationToken
 import org.junit.Before
@@ -46,18 +47,19 @@ class DelegationTokenEndToEndAuthorizationTest extends EndToEndAuthorizationTest
 
   this.serverConfig.setProperty(KafkaConfig.DelegationTokenMasterKeyProp, "testKey")
 
-  override def configureSecurityBeforeServersStart() {
+  override def configureSecurityBeforeServersStart(): Unit = {
     super.configureSecurityBeforeServersStart()
     zkClient.makeSurePersistentPathExists(ConfigEntityChangeNotificationZNode.path)
     // Create broker admin credentials before starting brokers
     createScramCredentials(zkConnect, kafkaPrincipal, kafkaPassword)
   }
 
-  override def configureSecurityAfterServersStart() {
+  override def configureSecurityAfterServersStart(): Unit = {
     super.configureSecurityAfterServersStart()
 
     // create scram credential for user "scram-user"
     createScramCredentials(zkConnect, clientPrincipal, clientPassword)
+    waitForScramCredentials(clientPrincipal)
 
     //create a token with "scram-user" credentials
     val token = createDelegationToken()
@@ -66,10 +68,18 @@ class DelegationTokenEndToEndAuthorizationTest extends EndToEndAuthorizationTest
     val clientLoginContext = JaasTestUtils.tokenClientLoginModule(token.tokenInfo().tokenId(), token.hmacAsBase64String())
     producerConfig.put(SaslConfigs.SASL_JAAS_CONFIG, clientLoginContext)
     consumerConfig.put(SaslConfigs.SASL_JAAS_CONFIG, clientLoginContext)
+    adminClientConfig.put(SaslConfigs.SASL_JAAS_CONFIG, clientLoginContext)
+  }
+
+  private def waitForScramCredentials(clientPrincipal: String): Unit = {
+    servers.foreach { server =>
+      val cache = server.credentialProvider.credentialCache.cache(kafkaClientSaslMechanism, classOf[ScramCredential])
+      TestUtils.waitUntilTrue(() => cache.get(clientPrincipal) != null, s"SCRAM credentials not created for $clientPrincipal")
+    }
   }
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     startSasl(jaasSections(kafkaServerSaslMechanisms, Option(kafkaClientSaslMechanism), Both))
     super.setUp()
   }
diff --git a/core/src/test/scala/integration/kafka/api/DescribeAuthorizedOperationsTest.scala b/core/src/test/scala/integration/kafka/api/DescribeAuthorizedOperationsTest.scala
index 5e5359209c6a0..bda827127dbbc 100644
--- a/core/src/test/scala/integration/kafka/api/DescribeAuthorizedOperationsTest.scala
+++ b/core/src/test/scala/integration/kafka/api/DescribeAuthorizedOperationsTest.scala
@@ -16,40 +16,52 @@ import java.io.File
 import java.util
 import java.util.Properties
 
-import kafka.security.auth.{Allow, Alter, Authorizer, ClusterAction, Group, Operation, PermissionType, SimpleAclAuthorizer, Acl => AuthAcl, Resource => AuthResource}
+import kafka.security.auth.{Cluster, Group, Resource, Topic, Acl => AuthAcl}
+import kafka.security.authorizer.AclAuthorizer
 import kafka.server.KafkaConfig
 import kafka.utils.{CoreUtils, JaasTestUtils, TestUtils}
-import org.apache.kafka.clients.admin.{AdminClient, AdminClientConfig, DescribeConsumerGroupsOptions}
+import org.apache.kafka.clients.admin._
 import org.apache.kafka.common.acl._
+import org.apache.kafka.common.acl.AclOperation.{ALTER, DESCRIBE, CLUSTER_ACTION}
+import org.apache.kafka.common.acl.AclPermissionType.ALLOW
 import org.apache.kafka.common.resource.{PatternType, ResourcePattern, ResourceType}
 import org.apache.kafka.common.security.auth.{KafkaPrincipal, SecurityProtocol}
 import org.apache.kafka.common.utils.Utils
-import org.junit.Assert.assertEquals
+import org.apache.kafka.server.authorizer.Authorizer
+import org.junit.Assert.{assertEquals, assertFalse}
 import org.junit.{After, Before, Test}
 
 import scala.collection.JavaConverters._
 
 class DescribeAuthorizedOperationsTest extends IntegrationTestHarness with SaslSetup {
-  override val serverCount = 1
+  override val brokerCount = 1
   this.serverConfig.setProperty(KafkaConfig.ZkEnableSecureAclsProp, "true")
-  this.serverConfig.setProperty(KafkaConfig.AuthorizerClassNameProp, classOf[SimpleAclAuthorizer].getName)
+  this.serverConfig.setProperty(KafkaConfig.AuthorizerClassNameProp, classOf[AclAuthorizer].getName)
 
-  var client: AdminClient = _
+  var client: Admin = _
   val group1 = "group1"
   val group2 = "group2"
   val group3 = "group3"
+  val topic1 = "topic1"
+  val topic2 = "topic2"
 
   override protected def securityProtocol = SecurityProtocol.SASL_SSL
 
   override protected lazy val trustStoreFile = Some(File.createTempFile("truststore", ".jks"))
 
-  override def configureSecurityBeforeServersStart() {
-    val authorizer = CoreUtils.createObject[Authorizer](classOf[SimpleAclAuthorizer].getName)
+  override def configureSecurityBeforeServersStart(): Unit = {
+    val authorizer = CoreUtils.createObject[Authorizer](classOf[AclAuthorizer].getName)
+    val clusterResource = new ResourcePattern(ResourceType.CLUSTER, Resource.ClusterResource.name, PatternType.LITERAL)
+    val topicResource = new ResourcePattern(ResourceType.TOPIC, Resource.WildCardResource, PatternType.LITERAL)
+
     try {
       authorizer.configure(this.configs.head.originals())
-      authorizer.addAcls(Set(clusterAcl(JaasTestUtils.KafkaServerPrincipalUnqualifiedName, Allow, ClusterAction),
-        clusterAcl(JaasTestUtils.KafkaClientPrincipalUnqualifiedName2, Allow, Alter)),
-        AuthResource.ClusterResource)
+      val result = authorizer.createAcls(null, List(
+        new AclBinding(clusterResource, accessControlEntry(JaasTestUtils.KafkaServerPrincipalUnqualifiedName.toString, ALLOW, CLUSTER_ACTION)),
+        new AclBinding(clusterResource, accessControlEntry(JaasTestUtils.KafkaClientPrincipalUnqualifiedName2.toString, ALLOW, ALTER)),
+        new AclBinding(topicResource, accessControlEntry(JaasTestUtils.KafkaClientPrincipalUnqualifiedName2.toString, ALLOW, DESCRIBE))).asJava)
+      result.asScala.map(_.toCompletableFuture.get).foreach { result => assertFalse(result.exception.isPresent) }
+
     } finally {
       authorizer.close()
     }
@@ -62,9 +74,9 @@ class DescribeAuthorizedOperationsTest extends IntegrationTestHarness with SaslS
     TestUtils.waitUntilBrokerMetadataIsPropagated(servers)
   }
 
-  private def clusterAcl(userName: String, permissionType: PermissionType, operation: Operation): AuthAcl = {
-    new AuthAcl(new KafkaPrincipal(KafkaPrincipal.USER_TYPE, userName), permissionType,
-      AuthAcl.WildCardHost, operation)
+  private def accessControlEntry(userName: String, permissionType: AclPermissionType, operation: AclOperation): AccessControlEntry = {
+    new AccessControlEntry(new KafkaPrincipal(KafkaPrincipal.USER_TYPE, userName).toString,
+      AuthAcl.WildCardHost, operation, permissionType)
   }
 
   @After
@@ -84,6 +96,15 @@ class DescribeAuthorizedOperationsTest extends IntegrationTestHarness with SaslS
   val group3Acl = new AclBinding(new ResourcePattern(ResourceType.GROUP, group3, PatternType.LITERAL),
     new AccessControlEntry("User:" + JaasTestUtils.KafkaClientPrincipalUnqualifiedName2, "*", AclOperation.DELETE, AclPermissionType.ALLOW))
 
+  val clusteAllAcl = new AclBinding(Resource.ClusterResource.toPattern,
+    new AccessControlEntry("User:" + JaasTestUtils.KafkaClientPrincipalUnqualifiedName2, "*", AclOperation.ALL, AclPermissionType.ALLOW))
+
+  val topic1Acl = new AclBinding(new ResourcePattern(ResourceType.TOPIC, topic1, PatternType.LITERAL),
+    new AccessControlEntry("User:" + JaasTestUtils.KafkaClientPrincipalUnqualifiedName2, "*", AclOperation.ALL, AclPermissionType.ALLOW))
+
+  val topic2All = new AclBinding(new ResourcePattern(ResourceType.TOPIC, topic2, PatternType.LITERAL),
+    new AccessControlEntry("User:" + JaasTestUtils.KafkaClientPrincipalUnqualifiedName2, "*", AclOperation.DELETE, AclPermissionType.ALLOW))
+
   def createConfig(): Properties = {
     val adminClientConfig = new Properties()
     adminClientConfig.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, brokerList)
@@ -118,4 +139,63 @@ class DescribeAuthorizedOperationsTest extends IntegrationTestHarness with SaslS
     assertEquals(Set(AclOperation.DESCRIBE, AclOperation.DELETE), group3Description.authorizedOperations().asScala.toSet)
   }
 
+  @Test
+  def testClusterAuthorizedOperations(): Unit = {
+    client = AdminClient.create(createConfig())
+
+    // test without includeAuthorizedOperations flag
+    var clusterDescribeResult = client.describeCluster()
+    assertEquals(Set(), clusterDescribeResult.authorizedOperations().get().asScala.toSet)
+
+    //test with includeAuthorizedOperations flag, we have give Alter permission
+    // in configureSecurityBeforeServersStart()
+    clusterDescribeResult = client.describeCluster(new DescribeClusterOptions().
+      includeAuthorizedOperations(true))
+    assertEquals(Set(AclOperation.DESCRIBE, AclOperation.ALTER),
+      clusterDescribeResult.authorizedOperations().get().asScala.toSet)
+
+    // enable all operations for cluster resource
+    val results = client.createAcls(List(clusteAllAcl).asJava)
+    assertEquals(Set(clusteAllAcl), results.values.keySet.asScala)
+    results.all.get
+
+    val expectedOperations = Cluster.supportedOperations
+      .map(operation => operation.toJava).asJava
+
+    clusterDescribeResult = client.describeCluster(new DescribeClusterOptions().
+      includeAuthorizedOperations(true))
+    assertEquals(expectedOperations, clusterDescribeResult.authorizedOperations().get())
+  }
+
+  @Test
+  def testTopicAuthorizedOperations(): Unit = {
+    client = AdminClient.create(createConfig())
+    createTopic(topic1)
+    createTopic(topic2)
+
+    // test without includeAuthorizedOperations flag
+    var describeTopicsResult = client.describeTopics(Set(topic1, topic2).asJava).all.get()
+    assertEquals(Set(), describeTopicsResult.get(topic1).authorizedOperations().asScala.toSet)
+    assertEquals(Set(), describeTopicsResult.get(topic2).authorizedOperations().asScala.toSet)
+
+    //test with includeAuthorizedOperations flag
+    describeTopicsResult = client.describeTopics(Set(topic1, topic2).asJava,
+      new DescribeTopicsOptions().includeAuthorizedOperations(true)).all.get()
+    assertEquals(Set(AclOperation.DESCRIBE), describeTopicsResult.get(topic1).authorizedOperations().asScala.toSet)
+    assertEquals(Set(AclOperation.DESCRIBE), describeTopicsResult.get(topic2).authorizedOperations().asScala.toSet)
+
+    //add few permissions
+    val results = client.createAcls(List(topic1Acl, topic2All).asJava)
+    assertEquals(Set(topic1Acl, topic2All), results.values.keySet.asScala)
+    results.all.get
+
+    val expectedOperations = Topic.supportedOperations
+      .map(operation => operation.toJava).asJava
+
+    describeTopicsResult = client.describeTopics(Set(topic1, topic2).asJava,
+      new DescribeTopicsOptions().includeAuthorizedOperations(true)).all.get()
+    assertEquals(expectedOperations, describeTopicsResult.get(topic1).authorizedOperations())
+    assertEquals(Set(AclOperation.DESCRIBE, AclOperation.DELETE),
+      describeTopicsResult.get(topic2).authorizedOperations().asScala.toSet)
+  }
 }
diff --git a/core/src/test/scala/integration/kafka/api/EndToEndAuthorizationTest.scala b/core/src/test/scala/integration/kafka/api/EndToEndAuthorizationTest.scala
index 0587a6de16021..9a40d69335b63 100644
--- a/core/src/test/scala/integration/kafka/api/EndToEndAuthorizationTest.scala
+++ b/core/src/test/scala/integration/kafka/api/EndToEndAuthorizationTest.scala
@@ -19,23 +19,27 @@ package kafka.api
 
 import com.yammer.metrics.Metrics
 import com.yammer.metrics.core.Gauge
-
 import java.io.File
 import java.util.concurrent.ExecutionException
 
 import kafka.admin.AclCommand
-import kafka.security.auth._
+import kafka.security.authorizer.AclAuthorizer
+import kafka.security.authorizer.AuthorizerUtils.WildcardHost
 import kafka.server._
 import kafka.utils._
-import org.apache.kafka.clients.consumer.{Consumer, ConsumerConfig}
+import org.apache.kafka.clients.consumer.{Consumer, ConsumerConfig, ConsumerRecords}
 import org.apache.kafka.clients.producer.{KafkaProducer, ProducerRecord}
-import org.apache.kafka.common.security.auth.KafkaPrincipal
+import org.apache.kafka.common.acl._
+import org.apache.kafka.common.acl.AclOperation._
+import org.apache.kafka.common.acl.AclPermissionType._
 import org.apache.kafka.common.{KafkaException, TopicPartition}
 import org.apache.kafka.common.errors.{GroupAuthorizationException, TopicAuthorizationException}
-import org.apache.kafka.common.resource.PatternType
+import org.apache.kafka.common.resource._
+import org.apache.kafka.common.resource.ResourceType._
 import org.apache.kafka.common.resource.PatternType.{LITERAL, PREFIXED}
 import org.junit.Assert._
 import org.junit.{After, Before, Test}
+import org.scalatest.Assertions.{assertThrows, fail, intercept}
 
 import scala.collection.JavaConverters._
 
@@ -58,9 +62,9 @@ import scala.collection.JavaConverters._
   * would end up with ZooKeeperTestHarness twice.
   */
 abstract class EndToEndAuthorizationTest extends IntegrationTestHarness with SaslSetup {
-  override val serverCount = 3
+  override val brokerCount = 3
 
-  override def configureSecurityBeforeServersStart() {
+  override def configureSecurityBeforeServersStart(): Unit = {
     AclCommand.main(clusterActionArgs)
     AclCommand.main(topicBrokerReadAclArgs)
   }
@@ -77,14 +81,17 @@ abstract class EndToEndAuthorizationTest extends IntegrationTestHarness with Sas
   val kafkaPrincipal: String
 
   override protected lazy val trustStoreFile = Some(File.createTempFile("truststore", ".jks"))
-
-  val topicResource = Resource(Topic, topic, LITERAL)
-  val groupResource = Resource(Group, group, LITERAL)
-  val clusterResource = Resource.ClusterResource
-  val prefixedTopicResource = Resource(Topic, topicPrefix, PREFIXED)
-  val prefixedGroupResource = Resource(Group, groupPrefix, PREFIXED)
-  val wildcardTopicResource = Resource(Topic, wildcard, LITERAL)
-  val wildcardGroupResource = Resource(Group, wildcard, LITERAL)
+  protected def authorizerClass: Class[_] = classOf[AclAuthorizer]
+
+  val topicResource = new ResourcePattern(TOPIC, topic, LITERAL)
+  val groupResource =  new ResourcePattern(GROUP, group, LITERAL)
+  val clusterResource = new ResourcePattern(CLUSTER, Resource.CLUSTER_NAME, LITERAL)
+  val prefixedTopicResource =  new ResourcePattern(TOPIC, topicPrefix, PREFIXED)
+  val prefixedGroupResource =  new ResourcePattern(GROUP, groupPrefix, PREFIXED)
+  val wildcardTopicResource =  new ResourcePattern(TOPIC, wildcard, LITERAL)
+  val wildcardGroupResource =  new ResourcePattern(GROUP, wildcard, LITERAL)
+  def kafkaPrincipalStr = s"$kafkaPrincipalType:$kafkaPrincipal"
+  def clientPrincipalStr = s"$kafkaPrincipalType:$clientPrincipal"
 
   // Arguments to AclCommand to set ACLs.
   def clusterActionArgs: Array[String] = Array("--authorizer-properties",
@@ -92,52 +99,52 @@ abstract class EndToEndAuthorizationTest extends IntegrationTestHarness with Sas
                                           s"--add",
                                           s"--cluster",
                                           s"--operation=ClusterAction",
-                                          s"--allow-principal=$kafkaPrincipalType:$kafkaPrincipal")
+                                          s"--allow-principal=$kafkaPrincipalStr")
   def topicBrokerReadAclArgs: Array[String] = Array("--authorizer-properties",
                                           s"zookeeper.connect=$zkConnect",
                                           s"--add",
                                           s"--topic=$wildcard",
                                           s"--operation=Read",
-                                          s"--allow-principal=$kafkaPrincipalType:$kafkaPrincipal")
+                                          s"--allow-principal=$kafkaPrincipalStr")
   def produceAclArgs(topic: String): Array[String] = Array("--authorizer-properties",
                                           s"zookeeper.connect=$zkConnect",
                                           s"--add",
                                           s"--topic=$topic",
                                           s"--producer",
-                                          s"--allow-principal=$kafkaPrincipalType:$clientPrincipal")
+                                          s"--allow-principal=$clientPrincipalStr")
   def describeAclArgs: Array[String] = Array("--authorizer-properties",
                                           s"zookeeper.connect=$zkConnect",
                                           s"--add",
                                           s"--topic=$topic",
                                           s"--operation=Describe",
-                                          s"--allow-principal=$kafkaPrincipalType:$clientPrincipal")
+                                          s"--allow-principal=$clientPrincipalStr")
   def deleteDescribeAclArgs: Array[String] = Array("--authorizer-properties",
                                           s"zookeeper.connect=$zkConnect",
                                           s"--remove",
                                           s"--force",
                                           s"--topic=$topic",
                                           s"--operation=Describe",
-                                          s"--allow-principal=$kafkaPrincipalType:$clientPrincipal")
+                                          s"--allow-principal=$clientPrincipalStr")
   def deleteWriteAclArgs: Array[String] = Array("--authorizer-properties",
                                           s"zookeeper.connect=$zkConnect",
                                           s"--remove",
                                           s"--force",
                                           s"--topic=$topic",
                                           s"--operation=Write",
-                                          s"--allow-principal=$kafkaPrincipalType:$clientPrincipal")
+                                          s"--allow-principal=$clientPrincipalStr")
   def consumeAclArgs(topic: String): Array[String] = Array("--authorizer-properties",
                                           s"zookeeper.connect=$zkConnect",
                                           s"--add",
                                           s"--topic=$topic",
                                           s"--group=$group",
                                           s"--consumer",
-                                          s"--allow-principal=$kafkaPrincipalType:$clientPrincipal")
+                                          s"--allow-principal=$clientPrincipalStr")
   def groupAclArgs: Array[String] = Array("--authorizer-properties",
                                           s"zookeeper.connect=$zkConnect",
                                           s"--add",
                                           s"--group=$group",
                                           s"--operation=Read",
-                                          s"--allow-principal=$kafkaPrincipalType:$clientPrincipal")
+                                          s"--allow-principal=$clientPrincipalStr")
   def produceConsumeWildcardAclArgs: Array[String] = Array("--authorizer-properties",
                                           s"zookeeper.connect=$zkConnect",
                                           s"--add",
@@ -145,7 +152,7 @@ abstract class EndToEndAuthorizationTest extends IntegrationTestHarness with Sas
                                           s"--group=$wildcard",
                                           s"--consumer",
                                           s"--producer",
-                                          s"--allow-principal=$kafkaPrincipalType:$clientPrincipal")
+                                          s"--allow-principal=$clientPrincipalStr")
   def produceConsumePrefixedAclsArgs: Array[String] = Array("--authorizer-properties",
                                           s"zookeeper.connect=$zkConnect",
                                           s"--add",
@@ -154,19 +161,19 @@ abstract class EndToEndAuthorizationTest extends IntegrationTestHarness with Sas
                                           s"--resource-pattern-type=prefixed",
                                           s"--consumer",
                                           s"--producer",
-                                          s"--allow-principal=$kafkaPrincipalType:$clientPrincipal")
-
-  def ClusterActionAcl = Set(new Acl(new KafkaPrincipal(kafkaPrincipalType, kafkaPrincipal), Allow, Acl.WildCardHost, ClusterAction))
-  def TopicBrokerReadAcl = Set(new Acl(new KafkaPrincipal(kafkaPrincipalType, kafkaPrincipal), Allow, Acl.WildCardHost, Read))
-  def GroupReadAcl = Set(new Acl(new KafkaPrincipal(kafkaPrincipalType, clientPrincipal), Allow, Acl.WildCardHost, Read))
-  def TopicReadAcl = Set(new Acl(new KafkaPrincipal(kafkaPrincipalType, clientPrincipal), Allow, Acl.WildCardHost, Read))
-  def TopicWriteAcl = Set(new Acl(new KafkaPrincipal(kafkaPrincipalType, clientPrincipal), Allow, Acl.WildCardHost, Write))
-  def TopicDescribeAcl = Set(new Acl(new KafkaPrincipal(kafkaPrincipalType, clientPrincipal), Allow, Acl.WildCardHost, Describe))
-  def TopicCreateAcl = Set(new Acl(new KafkaPrincipal(kafkaPrincipalType, clientPrincipal), Allow, Acl.WildCardHost, Create))
+                                          s"--allow-principal=$clientPrincipalStr")
+
+  def ClusterActionAcl = Set(new AccessControlEntry(kafkaPrincipalStr, WildcardHost, CLUSTER_ACTION, ALLOW))
+  def TopicBrokerReadAcl = Set(new AccessControlEntry(kafkaPrincipalStr, WildcardHost, READ, ALLOW))
+  def GroupReadAcl = Set(new AccessControlEntry(clientPrincipalStr, WildcardHost, READ, ALLOW))
+  def TopicReadAcl = Set(new AccessControlEntry(clientPrincipalStr, WildcardHost, READ, ALLOW))
+  def TopicWriteAcl = Set(new AccessControlEntry(clientPrincipalStr, WildcardHost, WRITE, ALLOW))
+  def TopicDescribeAcl = Set(new AccessControlEntry(clientPrincipalStr, WildcardHost, DESCRIBE, ALLOW))
+  def TopicCreateAcl = Set(new AccessControlEntry(clientPrincipalStr, WildcardHost, CREATE, ALLOW))
   // 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(KafkaConfig.AuthorizerClassNameProp, classOf[SimpleAclAuthorizer].getName)
+  this.serverConfig.setProperty(KafkaConfig.AuthorizerClassNameProp, authorizerClass.getName)
   // Some needed configuration for brokers, producers, and consumers
   this.serverConfig.setProperty(KafkaConfig.OffsetsTopicPartitionsProp, "1")
   this.serverConfig.setProperty(KafkaConfig.OffsetsTopicReplicationFactorProp, "3")
@@ -174,16 +181,17 @@ abstract class EndToEndAuthorizationTest extends IntegrationTestHarness with Sas
   this.serverConfig.setProperty(KafkaConfig.DefaultReplicationFactorProp, "3")
   this.serverConfig.setProperty(KafkaConfig.ConnectionsMaxReauthMsProp, "1500")
   this.consumerConfig.setProperty(ConsumerConfig.GROUP_ID_CONFIG, "group")
+  this.consumerConfig.setProperty(ConsumerConfig.METADATA_MAX_AGE_CONFIG, "1500")
 
   /**
     * Starts MiniKDC and only then sets up the parent trait.
     */
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     super.setUp()
     servers.foreach { s =>
-      TestUtils.waitAndVerifyAcls(ClusterActionAcl, s.dataPlaneRequestProcessor.authorizer.get, Resource.ClusterResource)
-      TestUtils.waitAndVerifyAcls(TopicBrokerReadAcl, s.dataPlaneRequestProcessor.authorizer.get, Resource(Topic, "*", LITERAL))
+      TestUtils.waitAndVerifyAcls(ClusterActionAcl, s.dataPlaneRequestProcessor.authorizer.get, clusterResource)
+      TestUtils.waitAndVerifyAcls(TopicBrokerReadAcl, s.dataPlaneRequestProcessor.authorizer.get, new ResourcePattern(TOPIC, "*", LITERAL))
     }
     // create the test topic with all the brokers as replicas
     createTopic(topic, 1, 3)
@@ -193,7 +201,7 @@ abstract class EndToEndAuthorizationTest extends IntegrationTestHarness with Sas
     * Closes MiniKDC last when tearing down.
     */
   @After
-  override def tearDown() {
+  override def tearDown(): Unit = {
     super.tearDown()
     closeSasl()
   }
@@ -210,7 +218,7 @@ abstract class EndToEndAuthorizationTest extends IntegrationTestHarness with Sas
     confirmReauthenticationMetrics
   }
 
-  protected def confirmReauthenticationMetrics() : Unit = {
+  protected def confirmReauthenticationMetrics(): Unit = {
     val expiredConnectionsKilledCountTotal = getGauge("ExpiredConnectionsKilledCount").value()
     servers.foreach { s =>
         val numExpiredKilled = TestUtils.totalMetricValue(s, "expired-connections-killed-count")
@@ -272,7 +280,7 @@ abstract class EndToEndAuthorizationTest extends IntegrationTestHarness with Sas
     confirmReauthenticationMetrics
   }
 
-  private def setWildcardResourceAcls() {
+  private def setWildcardResourceAcls(): Unit = {
     AclCommand.main(produceConsumeWildcardAclArgs)
     servers.foreach { s =>
       TestUtils.waitAndVerifyAcls(TopicReadAcl ++ TopicWriteAcl ++ TopicDescribeAcl ++ TopicCreateAcl ++ TopicBrokerReadAcl, s.dataPlaneRequestProcessor.authorizer.get, wildcardTopicResource)
@@ -280,7 +288,7 @@ abstract class EndToEndAuthorizationTest extends IntegrationTestHarness with Sas
     }
   }
 
-  private def setPrefixedResourceAcls() {
+  private def setPrefixedResourceAcls(): Unit = {
     AclCommand.main(produceConsumePrefixedAclsArgs)
     servers.foreach { s =>
       TestUtils.waitAndVerifyAcls(TopicReadAcl ++ TopicWriteAcl ++ TopicDescribeAcl ++ TopicCreateAcl, s.dataPlaneRequestProcessor.authorizer.get, prefixedTopicResource)
@@ -288,27 +296,80 @@ abstract class EndToEndAuthorizationTest extends IntegrationTestHarness with Sas
     }
   }
 
-  protected def setAclsAndProduce(tp: TopicPartition) {
+  private def setReadAndWriteAcls(tp: TopicPartition): Unit = {
     AclCommand.main(produceAclArgs(tp.topic))
     AclCommand.main(consumeAclArgs(tp.topic))
     servers.foreach { s =>
       TestUtils.waitAndVerifyAcls(TopicReadAcl ++ TopicWriteAcl ++ TopicDescribeAcl ++ TopicCreateAcl, s.dataPlaneRequestProcessor.authorizer.get,
-        new Resource(Topic, tp.topic, PatternType.LITERAL))
+        new ResourcePattern(TOPIC, tp.topic, LITERAL))
       TestUtils.waitAndVerifyAcls(GroupReadAcl, s.dataPlaneRequestProcessor.authorizer.get, groupResource)
     }
+  }
+
+  protected def setAclsAndProduce(tp: TopicPartition): Unit = {
+    setReadAndWriteAcls(tp)
     val producer = createProducer()
     sendRecords(producer, numRecords, tp)
   }
 
+  private def setConsumerGroupAcls(): Unit = {
+    AclCommand.main(groupAclArgs)
+    servers.foreach { s =>
+      TestUtils.waitAndVerifyAcls(GroupReadAcl, s.dataPlaneRequestProcessor.authorizer.get, groupResource)
+    }
+  }
+
   /**
-    * Tests that a producer fails to publish messages when the appropriate ACL
-    * isn't set.
+    * Tests that producer, consumer and adminClient fail to publish messages, consume
+    * messages and describe topics respectively when the describe ACL isn't set.
+    * Also verifies that subsequent publish, consume and describe to authorized topic succeeds.
     */
-  @Test(expected = classOf[TopicAuthorizationException])
-  def testNoProduceWithoutDescribeAcl(): Unit = {
+  @Test
+  def testNoDescribeProduceOrConsumeWithoutTopicDescribeAcl(): Unit = {
+    // Set consumer group acls since we are testing topic authorization
+    setConsumerGroupAcls()
+
+    // Verify produce/consume/describe throw TopicAuthorizationException
     val producer = createProducer()
+    assertThrows[TopicAuthorizationException] { sendRecords(producer, numRecords, tp) }
+    val consumer = createConsumer()
+    consumer.assign(List(tp).asJava)
+    assertThrows[TopicAuthorizationException] { consumeRecords(consumer, numRecords, topic = tp.topic) }
+    val adminClient = createAdminClient()
+    val e1 = intercept[ExecutionException] { adminClient.describeTopics(Set(topic).asJava).all().get() }
+    assertTrue("Unexpected exception " + e1.getCause, e1.getCause.isInstanceOf[TopicAuthorizationException])
+
+    // Verify successful produce/consume/describe on another topic using the same producer, consumer and adminClient
+    val topic2 = "topic2"
+    val tp2 = new TopicPartition(topic2, 0)
+    setReadAndWriteAcls(tp2)
+    sendRecords(producer, numRecords, tp2)
+    consumer.assign(List(tp2).asJava)
+    consumeRecords(consumer, numRecords, topic = topic2)
+    val describeResults = adminClient.describeTopics(Set(topic, topic2).asJava).values
+    assertEquals(1, describeResults.get(topic2).get().partitions().size())
+    val e2 = intercept[ExecutionException] { adminClient.describeTopics(Set(topic).asJava).all().get() }
+    assertTrue("Unexpected exception " + e2.getCause, e2.getCause.isInstanceOf[TopicAuthorizationException])
+
+    // Verify that consumer manually assigning both authorized and unauthorized topic doesn't consume from either
+    consumer.assign(List(tp, tp2).asJava)
+    sendRecords(producer, numRecords, tp2)
+    def verifyNoRecords(records: ConsumerRecords[Array[Byte], Array[Byte]]): Boolean = {
+      assertTrue("Consumed records: " + records, records.isEmpty)
+      !records.isEmpty
+    }
+    assertThrows[TopicAuthorizationException] {
+      TestUtils.pollRecordsUntilTrue(consumer, verifyNoRecords, "Consumer didn't fail with authorization exception within timeout")
+    }
+
+    // Add ACLs and verify successful produce/consume/describe on first topic
+    setReadAndWriteAcls(tp)
+    consumeRecordsIgnoreOneAuthorizationException(consumer, numRecords, startingOffset = 1, topic2)
     sendRecords(producer, numRecords, tp)
-    confirmReauthenticationMetrics
+    consumeRecords(consumer, numRecords, topic = topic)
+    val describeResults2 = adminClient.describeTopics(Set(topic, topic2).asJava).values
+    assertEquals(1, describeResults2.get(topic).get().partitions().size())
+    assertEquals(1, describeResults2.get(topic2).get().partitions().size())
   }
 
   @Test
@@ -342,13 +403,22 @@ abstract class EndToEndAuthorizationTest extends IntegrationTestHarness with Sas
     confirmReauthenticationMetrics
   }
 
-  @Test(expected = classOf[TopicAuthorizationException])
+  @Test
   def testNoConsumeWithoutDescribeAclViaSubscribe(): Unit = {
     noConsumeWithoutDescribeAclSetup()
     val consumer = createConsumer()
     consumer.subscribe(List(topic).asJava)
     // this should timeout since the consumer will not be able to fetch any metadata for the topic
-    consumeRecords(consumer, timeout = 3000)
+    assertThrows[TopicAuthorizationException] { consumeRecords(consumer, timeout = 3000) }
+
+    // Verify that no records are consumed even if one of the requested topics is authorized
+    setReadAndWriteAcls(tp)
+    consumer.subscribe(List(topic, "topic2").asJava)
+    assertThrows[TopicAuthorizationException] { consumeRecords(consumer, timeout = 3000) }
+
+    // Verify that records are consumed if all topics are authorized
+    consumer.subscribe(List(topic).asJava)
+    consumeRecordsIgnoreOneAuthorizationException(consumer)
   }
 
   private def noConsumeWithoutDescribeAclSetup(): Unit = {
@@ -438,7 +508,7 @@ abstract class EndToEndAuthorizationTest extends IntegrationTestHarness with Sas
   }
 
   protected final def sendRecords(producer: KafkaProducer[Array[Byte], Array[Byte]],
-                                  numRecords: Int, tp: TopicPartition) {
+                                  numRecords: Int, tp: TopicPartition): Unit = {
     val futures = (0 until numRecords).map { i =>
       val record = new ProducerRecord(tp.topic(), tp.partition(), s"$i".getBytes, s"$i".getBytes)
       debug(s"Sending this record: $record")
@@ -456,7 +526,7 @@ abstract class EndToEndAuthorizationTest extends IntegrationTestHarness with Sas
                                      startingOffset: Int = 0,
                                      topic: String = topic,
                                      part: Int = part,
-                                     timeout: Long = 10000) {
+                                     timeout: Long = 10000): Unit = {
     val records = TestUtils.consumeRecords(consumer, numRecords, timeout)
 
     for (i <- 0 until numRecords) {
@@ -467,5 +537,17 @@ abstract class EndToEndAuthorizationTest extends IntegrationTestHarness with Sas
       assertEquals(offset.toLong, record.offset)
     }
   }
+
+  // Consume records, ignoring at most one TopicAuthorization exception from previously sent request
+  private def consumeRecordsIgnoreOneAuthorizationException(consumer: Consumer[Array[Byte], Array[Byte]],
+                                                            numRecords: Int = 1,
+                                                            startingOffset: Int = 0,
+                                                            topic: String = topic): Unit = {
+    try {
+      consumeRecords(consumer, numRecords, startingOffset, topic)
+    } catch {
+      case _: TopicAuthorizationException => consumeRecords(consumer, numRecords, startingOffset, topic)
+    }
+  }
 }
 
diff --git a/core/src/test/scala/integration/kafka/api/EndToEndClusterIdTest.scala b/core/src/test/scala/integration/kafka/api/EndToEndClusterIdTest.scala
index 04dcf28a68094..072c44d1a4d90 100644
--- a/core/src/test/scala/integration/kafka/api/EndToEndClusterIdTest.scala
+++ b/core/src/test/scala/integration/kafka/api/EndToEndClusterIdTest.scala
@@ -55,7 +55,7 @@ object EndToEndClusterIdTest {
 
   class MockConsumerMetricsReporter extends MockMetricsReporter with ClusterResourceListener {
 
-    override def onUpdate(clusterMetadata: ClusterResource) {
+    override def onUpdate(clusterMetadata: ClusterResource): Unit = {
       MockConsumerMetricsReporter.CLUSTER_META.set(clusterMetadata)
     }
   }
@@ -66,7 +66,7 @@ object EndToEndClusterIdTest {
 
   class MockProducerMetricsReporter extends MockMetricsReporter with ClusterResourceListener {
 
-    override def onUpdate(clusterMetadata: ClusterResource) {
+    override def onUpdate(clusterMetadata: ClusterResource): Unit = {
       MockProducerMetricsReporter.CLUSTER_META.set(clusterMetadata)
     }
   }
@@ -77,7 +77,7 @@ object EndToEndClusterIdTest {
 
   class MockBrokerMetricsReporter extends MockMetricsReporter with ClusterResourceListener {
 
-    override def onUpdate(clusterMetadata: ClusterResource) {
+    override def onUpdate(clusterMetadata: ClusterResource): Unit = {
       MockBrokerMetricsReporter.CLUSTER_META.set(clusterMetadata)
     }
   }
@@ -107,7 +107,7 @@ class EndToEndClusterIdTest extends KafkaServerTestHarness {
   }
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     super.setUp()
     MockDeserializer.resetStaticVariables
     // create the consumer offset topic
@@ -115,7 +115,7 @@ class EndToEndClusterIdTest extends KafkaServerTestHarness {
   }
 
   @Test
-  def testEndToEnd() {
+  def testEndToEnd(): Unit = {
     val appendStr = "mock"
     MockConsumerInterceptor.resetCounters()
     MockProducerInterceptor.resetCounters()
@@ -183,7 +183,7 @@ class EndToEndClusterIdTest extends KafkaServerTestHarness {
     MockProducerInterceptor.resetCounters()
   }
 
-  private def sendRecords(producer: KafkaProducer[Array[Byte], Array[Byte]], numRecords: Int, tp: TopicPartition) {
+  private def sendRecords(producer: KafkaProducer[Array[Byte], Array[Byte]], numRecords: Int, tp: TopicPartition): Unit = {
     val futures = (0 until numRecords).map { i =>
       val record = new ProducerRecord(tp.topic(), tp.partition(), s"$i".getBytes, s"$i".getBytes)
       debug(s"Sending this record: $record")
@@ -200,7 +200,7 @@ class EndToEndClusterIdTest extends KafkaServerTestHarness {
                              numRecords: Int,
                              startingOffset: Int = 0,
                              topic: String = topic,
-                             part: Int = part) {
+                             part: Int = part): Unit = {
     val records = TestUtils.consumeRecords(consumer, numRecords)
 
     for (i <- 0 until numRecords) {
diff --git a/core/src/test/scala/integration/kafka/api/GroupAuthorizerIntegrationTest.scala b/core/src/test/scala/integration/kafka/api/GroupAuthorizerIntegrationTest.scala
index 27c3f31c3ff82..6d2161d52e35f 100644
--- a/core/src/test/scala/integration/kafka/api/GroupAuthorizerIntegrationTest.scala
+++ b/core/src/test/scala/integration/kafka/api/GroupAuthorizerIntegrationTest.scala
@@ -33,9 +33,9 @@ class GroupAuthorizerIntegrationTest extends AuthorizerIntegrationTest {
   override val kafkaPrincipalType = GroupPrincipalType
   override def userPrincipal = TestGroupPrincipal
 
-  override def propertyOverrides(properties: Properties): Unit = {
+  override def brokerPropertyOverrides(properties: Properties): Unit = {
     properties.setProperty(BrokerSecurityConfigs.PRINCIPAL_BUILDER_CLASS_CONFIG,
       classOf[GroupPrincipalBuilder].getName)
-    super.propertyOverrides(properties)
+    super.brokerPropertyOverrides(properties)
   }
 }
\ No newline at end of file
diff --git a/core/src/test/scala/integration/kafka/api/GroupCoordinatorIntegrationTest.scala b/core/src/test/scala/integration/kafka/api/GroupCoordinatorIntegrationTest.scala
index 61845cf194505..e29867b886900 100644
--- a/core/src/test/scala/integration/kafka/api/GroupCoordinatorIntegrationTest.scala
+++ b/core/src/test/scala/integration/kafka/api/GroupCoordinatorIntegrationTest.scala
@@ -38,7 +38,7 @@ class GroupCoordinatorIntegrationTest extends KafkaServerTestHarness {
   }
 
   @Test
-  def testGroupCoordinatorPropagatesOffsetsTopicCompressionCodec() {
+  def testGroupCoordinatorPropagatesOffsetsTopicCompressionCodec(): Unit = {
     val consumer = TestUtils.createConsumer(TestUtils.getBrokerListStrFromServers(servers))
     val offsetMap = Map(
       new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, 0) -> new OffsetAndMetadata(10, "")
diff --git a/core/src/test/scala/integration/kafka/api/IntegrationTestHarness.scala b/core/src/test/scala/integration/kafka/api/IntegrationTestHarness.scala
index 5a2000557424f..187766d7c2ecd 100644
--- a/core/src/test/scala/integration/kafka/api/IntegrationTestHarness.scala
+++ b/core/src/test/scala/integration/kafka/api/IntegrationTestHarness.scala
@@ -23,37 +23,50 @@ import org.apache.kafka.clients.consumer.{ConsumerConfig, KafkaConsumer}
 import kafka.utils.TestUtils
 import kafka.utils.Implicits._
 import java.util.Properties
-import java.util.concurrent.TimeUnit
 
 import org.apache.kafka.clients.producer.{KafkaProducer, ProducerConfig}
 import kafka.server.KafkaConfig
 import kafka.integration.KafkaServerTestHarness
+import org.apache.kafka.clients.admin.{Admin, AdminClient, AdminClientConfig}
 import org.apache.kafka.common.network.{ListenerName, Mode}
 import org.apache.kafka.common.serialization.{ByteArrayDeserializer, ByteArraySerializer, Deserializer, Serializer}
 import org.junit.{After, Before}
 
 import scala.collection.mutable
+import scala.collection.Seq
 
 /**
  * A helper class for writing integration tests that involve producers, consumers, and servers
  */
 abstract class IntegrationTestHarness extends KafkaServerTestHarness {
-  protected def serverCount: Int
+  protected def brokerCount: Int
   protected def logDirCount: Int = 1
 
   val producerConfig = new Properties
   val consumerConfig = new Properties
+  val adminClientConfig = new Properties
   val serverConfig = new Properties
 
   private val consumers = mutable.Buffer[KafkaConsumer[_, _]]()
   private val producers = mutable.Buffer[KafkaProducer[_, _]]()
+  private val adminClients = mutable.Buffer[Admin]()
 
   protected def interBrokerListenerName: ListenerName = listenerName
 
+  protected def modifyConfigs(props: Seq[Properties]): Unit = {
+    configureListeners(props)
+    props.foreach(_ ++= serverConfig)
+  }
+
   override def generateConfigs: Seq[KafkaConfig] = {
-    val cfgs = TestUtils.createBrokerConfigs(serverCount, zkConnect, interBrokerSecurityProtocol = Some(securityProtocol),
+    val cfgs = TestUtils.createBrokerConfigs(brokerCount, zkConnect, interBrokerSecurityProtocol = Some(securityProtocol),
       trustStoreFile = trustStoreFile, saslProperties = serverSaslProperties, logDirCount = logDirCount)
-    cfgs.foreach { config =>
+    modifyConfigs(cfgs)
+    cfgs.map(KafkaConfig.fromProps)
+  }
+
+  protected def configureListeners(props: Seq[Properties]): Unit = {
+    props.foreach { config =>
       config.remove(KafkaConfig.InterBrokerSecurityProtocolProp)
       config.setProperty(KafkaConfig.InterBrokerListenerNameProp, interBrokerListenerName.value)
 
@@ -64,12 +77,10 @@ abstract class IntegrationTestHarness extends KafkaServerTestHarness {
       config.setProperty(KafkaConfig.ListenersProp, listeners)
       config.setProperty(KafkaConfig.ListenerSecurityProtocolMapProp, listenerSecurityMap)
     }
-    cfgs.foreach(_ ++= serverConfig)
-    cfgs.map(KafkaConfig.fromProps)
   }
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     doSetup(createOffsetsTopic = true)
   }
 
@@ -77,6 +88,7 @@ abstract class IntegrationTestHarness extends KafkaServerTestHarness {
     // Generate client security properties before starting the brokers in case certs are needed
     producerConfig ++= clientSecurityProps("producer")
     consumerConfig ++= clientSecurityProps("consumer")
+    adminClientConfig ++= clientSecurityProps("adminClient")
 
     super.setUp()
 
@@ -84,12 +96,16 @@ abstract class IntegrationTestHarness extends KafkaServerTestHarness {
     producerConfig.putIfAbsent(ProducerConfig.ACKS_CONFIG, "-1")
     producerConfig.putIfAbsent(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, classOf[ByteArraySerializer].getName)
     producerConfig.putIfAbsent(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, classOf[ByteArraySerializer].getName)
+    producerConfig.putIfAbsent(ProducerConfig.ALLOW_AUTO_CREATE_TOPICS_CONFIG, boolean2Boolean(true))
 
     consumerConfig.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, brokerList)
     consumerConfig.putIfAbsent(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest")
     consumerConfig.putIfAbsent(ConsumerConfig.GROUP_ID_CONFIG, "group")
     consumerConfig.putIfAbsent(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, classOf[ByteArrayDeserializer].getName)
     consumerConfig.putIfAbsent(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, classOf[ByteArrayDeserializer].getName)
+    consumerConfig.putIfAbsent(ConsumerConfig.ALLOW_AUTO_CREATE_TOPICS_CONFIG, boolean2Boolean(true))
+
+    adminClientConfig.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, brokerList)
 
     if (createOffsetsTopic)
       TestUtils.createOffsetsTopic(zkClient, servers)
@@ -124,13 +140,26 @@ abstract class IntegrationTestHarness extends KafkaServerTestHarness {
     consumer
   }
 
+  def createAdminClient(configOverrides: Properties = new Properties): Admin = {
+    val props = new Properties
+    props ++= adminClientConfig
+    props ++= configOverrides
+    val adminClient = AdminClient.create(props)
+    adminClients += adminClient
+    adminClient
+  }
+
   @After
-  override def tearDown() {
-    producers.foreach(_.close(0, TimeUnit.MILLISECONDS))
+  override def tearDown(): Unit = {
+    producers.foreach(_.close(Duration.ZERO))
     consumers.foreach(_.wakeup())
     consumers.foreach(_.close(Duration.ZERO))
+    adminClients.foreach(_.close(Duration.ZERO))
+
     producers.clear()
     consumers.clear()
+    adminClients.clear()
+
     super.tearDown()
   }
 
diff --git a/core/src/test/scala/integration/kafka/api/LegacyAdminClientTest.scala b/core/src/test/scala/integration/kafka/api/LegacyAdminClientTest.scala
deleted file mode 100644
index 08a0224704d1b..0000000000000
--- a/core/src/test/scala/integration/kafka/api/LegacyAdminClientTest.scala
+++ /dev/null
@@ -1,156 +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.api
-
-import java.util.Collections
-
-import kafka.admin.AdminClient
-import kafka.server.KafkaConfig
-import java.lang.{Long => JLong}
-
-import kafka.utils.{Logging, TestUtils}
-import org.apache.kafka.clients.consumer.{ConsumerConfig, KafkaConsumer}
-import org.apache.kafka.clients.producer.ProducerConfig
-import org.apache.kafka.common.TopicPartition
-import org.apache.kafka.common.protocol.ApiKeys
-import org.junit.{After, Before, Test}
-import org.junit.Assert._
-
-import scala.collection.JavaConverters._
-
-/**
-  * Tests for the deprecated Scala AdminClient.
-  */
-@deprecated("The Scala AdminClient has been deprecated in favour of org.apache.kafka.clients.admin.AdminClient",
-  since = "0.11.0")
-class LegacyAdminClientTest extends IntegrationTestHarness with Logging {
-
-  val producerCount = 1
-  val consumerCount = 2
-  val serverCount = 3
-  val groupId = "my-test"
-  val clientId = "consumer-498"
-
-  val topic = "topic"
-  val part = 0
-  val tp = new TopicPartition(topic, part)
-  val part2 = 1
-  val tp2 = new TopicPartition(topic, part2)
-
-  var client: AdminClient = null
-
-  // configure the servers and clients
-  this.serverConfig.setProperty(KafkaConfig.ControlledShutdownEnableProp, "false") // speed up shutdown
-  this.serverConfig.setProperty(KafkaConfig.OffsetsTopicReplicationFactorProp, "3") // don't want to lose offset
-  this.serverConfig.setProperty(KafkaConfig.OffsetsTopicPartitionsProp, "1")
-  this.serverConfig.setProperty(KafkaConfig.GroupMinSessionTimeoutMsProp, "100") // set small enough session timeout
-  this.serverConfig.setProperty(KafkaConfig.GroupInitialRebalanceDelayMsProp, "0") // do initial rebalance immediately
-  this.producerConfig.setProperty(ProducerConfig.ACKS_CONFIG, "all")
-  this.consumerConfig.setProperty(ConsumerConfig.GROUP_ID_CONFIG, groupId)
-  this.consumerConfig.setProperty(ConsumerConfig.CLIENT_ID_CONFIG, clientId)
-  this.consumerConfig.setProperty(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest")
-  this.consumerConfig.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "false")
-  this.consumerConfig.setProperty(ConsumerConfig.METADATA_MAX_AGE_CONFIG, "100")
-
-  @Before
-  override def setUp() {
-    super.setUp()
-    client = AdminClient.createSimplePlaintext(this.brokerList)
-    createTopic(topic, 2, serverCount)
-  }
-
-  @After
-  override def tearDown() {
-    client.close()
-    super.tearDown()
-  }
-
-  @Test
-  def testOffsetsForTimesWhenOffsetNotFound() {
-    val consumer = createConsumer()
-    assertNull(consumer.offsetsForTimes(Map(tp -> JLong.valueOf(0L)).asJava).get(tp))
-  }
-
-  @Test
-  def testListGroups() {
-    val consumer = createConsumer()
-    subscribeAndWaitForAssignment(topic, consumer)
-
-    val groups = client.listAllGroupsFlattened
-    assertFalse(groups.isEmpty)
-    val group = groups.head
-    assertEquals(groupId, group.groupId)
-    assertEquals("consumer", group.protocolType)
-  }
-
-  @Test
-  def testListAllBrokerVersionInfo() {
-    val consumer = createConsumer()
-    subscribeAndWaitForAssignment(topic, consumer)
-
-    val brokerVersionInfos = client.listAllBrokerVersionInfo
-    val brokers = brokerList.split(",")
-    assertEquals(brokers.size, brokerVersionInfos.size)
-    for ((node, tryBrokerVersionInfo) <- brokerVersionInfos) {
-      val hostStr = s"${node.host}:${node.port}"
-      assertTrue(s"Unknown host:port pair $hostStr in brokerVersionInfos", brokers.contains(hostStr))
-      val brokerVersionInfo = tryBrokerVersionInfo.get
-      assertEquals(2, brokerVersionInfo.latestUsableVersion(ApiKeys.API_VERSIONS))
-    }
-  }
-
-  @Test
-  def testGetConsumerGroupSummary() {
-    val consumer = createConsumer()
-    subscribeAndWaitForAssignment(topic, consumer)
-
-    val group = client.describeConsumerGroup(groupId)
-    assertEquals("range", group.assignmentStrategy)
-    assertEquals("Stable", group.state)
-    assertFalse(group.consumers.isEmpty)
-
-    val member = group.consumers.get.head
-    assertEquals(clientId, member.clientId)
-    assertFalse(member.host.isEmpty)
-    assertFalse(member.consumerId.isEmpty)
-  }
-
-  @Test
-  def testDescribeConsumerGroup() {
-    val consumer = createConsumer()
-    subscribeAndWaitForAssignment(topic, consumer)
-
-    val consumerGroupSummary = client.describeConsumerGroup(groupId)
-    assertEquals(1, consumerGroupSummary.consumers.get.size)
-    assertEquals(List(tp, tp2), consumerGroupSummary.consumers.get.flatMap(_.assignment))
-  }
-
-  @Test
-  def testDescribeConsumerGroupForNonExistentGroup() {
-    val nonExistentGroup = "non" + groupId
-    assertTrue("Expected empty ConsumerSummary list", client.describeConsumerGroup(nonExistentGroup).consumers.get.isEmpty)
-  }
-
-  private def subscribeAndWaitForAssignment(topic: String, consumer: KafkaConsumer[Array[Byte], Array[Byte]]) {
-    consumer.subscribe(Collections.singletonList(topic))
-    TestUtils.waitUntilTrue(() => {
-      consumer.poll(0)
-      !consumer.assignment.isEmpty
-    }, "Expected non-empty assignment")
-  }
-
-}
diff --git a/core/src/test/scala/integration/kafka/api/LogAppendTimeTest.scala b/core/src/test/scala/integration/kafka/api/LogAppendTimeTest.scala
index 795f954a32d12..74a774c682f67 100644
--- a/core/src/test/scala/integration/kafka/api/LogAppendTimeTest.scala
+++ b/core/src/test/scala/integration/kafka/api/LogAppendTimeTest.scala
@@ -33,7 +33,7 @@ import org.junit.Assert.{assertEquals, assertNotEquals, assertTrue}
 class LogAppendTimeTest extends IntegrationTestHarness {
   val producerCount: Int = 1
   val consumerCount: Int = 1
-  val serverCount: Int = 2
+  val brokerCount: Int = 2
 
   // This will be used for the offsets topic as well
   serverConfig.put(KafkaConfig.LogMessageTimestampTypeProp, TimestampType.LOG_APPEND_TIME.name)
@@ -42,13 +42,13 @@ class LogAppendTimeTest extends IntegrationTestHarness {
   private val topic = "topic"
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     super.setUp()
     createTopic(topic)
   }
 
   @Test
-  def testProduceConsume() {
+  def testProduceConsume(): Unit = {
     val producer = createProducer()
     val now = System.currentTimeMillis()
     val createTime = now - TimeUnit.DAYS.toMillis(1)
diff --git a/core/src/test/scala/integration/kafka/api/MetricsTest.scala b/core/src/test/scala/integration/kafka/api/MetricsTest.scala
index b2f5a7ea3b949..83c75dfb77f2b 100644
--- a/core/src/test/scala/integration/kafka/api/MetricsTest.scala
+++ b/core/src/test/scala/integration/kafka/api/MetricsTest.scala
@@ -28,12 +28,13 @@ import org.apache.kafka.common.network.ListenerName
 import org.apache.kafka.common.security.auth.SecurityProtocol
 import org.junit.{After, Before, Test}
 import org.junit.Assert._
+import org.scalatest.Assertions.fail
 
 import scala.collection.JavaConverters._
 
 class MetricsTest extends IntegrationTestHarness with SaslSetup {
 
-  override val serverCount = 1
+  override val brokerCount = 1
 
   override protected def listenerName = new ListenerName("CLIENT")
   private val kafkaClientSaslMechanism = "PLAIN"
@@ -133,7 +134,7 @@ class MetricsTest extends IntegrationTestHarness with SaslSetup {
   }
 
   private def verifyKafkaRateMetricsHaveCumulativeCount(producer: KafkaProducer[Array[Byte], Array[Byte]],
-                                                        consumer: KafkaConsumer[Array[Byte], Array[Byte]]): Unit =  {
+                                                        consumer: KafkaConsumer[Array[Byte], Array[Byte]]): Unit = {
 
     def exists(name: String, rateMetricName: MetricName, allMetricNames: Set[MetricName]): Boolean = {
       allMetricNames.contains(new MetricName(name, rateMetricName.group, "", rateMetricName.tags))
diff --git a/core/src/test/scala/integration/kafka/api/PlaintextConsumerTest.scala b/core/src/test/scala/integration/kafka/api/PlaintextConsumerTest.scala
index c11fc12d16d4a..9df468a34c57a 100644
--- a/core/src/test/scala/integration/kafka/api/PlaintextConsumerTest.scala
+++ b/core/src/test/scala/integration/kafka/api/PlaintextConsumerTest.scala
@@ -19,7 +19,6 @@ import java.util.regex.Pattern
 import java.util.{Collections, Locale, Optional, Properties}
 
 import kafka.log.LogConfig
-import kafka.server.KafkaConfig
 import kafka.utils.TestUtils
 import org.apache.kafka.clients.consumer._
 import org.apache.kafka.clients.producer.{KafkaProducer, ProducerConfig, ProducerRecord}
@@ -32,17 +31,20 @@ import org.apache.kafka.common.utils.Utils
 import org.apache.kafka.test.{MockConsumerInterceptor, MockProducerInterceptor}
 import org.junit.Assert._
 import org.junit.Test
+import org.scalatest.Assertions.intercept
 
 import scala.collection.JavaConverters._
 import scala.collection.mutable.Buffer
 import kafka.server.QuotaType
 import kafka.server.KafkaServer
 
+import scala.collection.mutable
+
 /* We have some tests in this class instead of `BaseConsumerTest` in order to keep the build time under control. */
 class PlaintextConsumerTest extends BaseConsumerTest {
 
   @Test
-  def testHeaders() {
+  def testHeaders(): Unit = {
     val numRecords = 1
     val record = new ProducerRecord(tp.topic, tp.partition, null, "key".getBytes, "value".getBytes)
 
@@ -156,7 +158,7 @@ class PlaintextConsumerTest extends BaseConsumerTest {
   }
 
   @Test
-  def testMaxPollRecords() {
+  def testMaxPollRecords(): Unit = {
     val maxPollRecords = 2
     val numRecords = 10000
 
@@ -170,7 +172,7 @@ class PlaintextConsumerTest extends BaseConsumerTest {
   }
 
   @Test
-  def testMaxPollIntervalMs() {
+  def testMaxPollIntervalMs(): Unit = {
     this.consumerConfig.setProperty(ConsumerConfig.MAX_POLL_INTERVAL_MS_CONFIG, 3000.toString)
     this.consumerConfig.setProperty(ConsumerConfig.HEARTBEAT_INTERVAL_MS_CONFIG, 500.toString)
     this.consumerConfig.setProperty(ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG, 2000.toString)
@@ -183,18 +185,18 @@ class PlaintextConsumerTest extends BaseConsumerTest {
     // rebalance to get the initial assignment
     awaitRebalance(consumer, listener)
     assertEquals(1, listener.callsToAssigned)
-    assertEquals(1, listener.callsToRevoked)
+    assertEquals(0, listener.callsToRevoked)
 
     Thread.sleep(3500)
 
     // we should fall out of the group and need to rebalance
     awaitRebalance(consumer, listener)
     assertEquals(2, listener.callsToAssigned)
-    assertEquals(2, listener.callsToRevoked)
+    assertEquals(1, listener.callsToRevoked)
   }
 
   @Test
-  def testMaxPollIntervalMsDelayInRevocation() {
+  def testMaxPollIntervalMsDelayInRevocation(): 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)
@@ -205,8 +207,9 @@ class PlaintextConsumerTest extends BaseConsumerTest {
     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 (callsToRevoked > 0) {
+        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
@@ -233,7 +236,7 @@ class PlaintextConsumerTest extends BaseConsumerTest {
   }
 
   @Test
-  def testMaxPollIntervalMsDelayInAssignment() {
+  def testMaxPollIntervalMsDelayInAssignment(): 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)
@@ -257,7 +260,7 @@ class PlaintextConsumerTest extends BaseConsumerTest {
   }
 
   @Test
-  def testAutoCommitOnClose() {
+  def testAutoCommitOnClose(): Unit = {
     this.consumerConfig.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "true")
     val consumer = createConsumer()
 
@@ -275,12 +278,12 @@ class PlaintextConsumerTest extends BaseConsumerTest {
 
     // now we should see the committed positions from another consumer
     val anotherConsumer = createConsumer()
-    assertEquals(300, anotherConsumer.committed(tp).offset)
-    assertEquals(500, anotherConsumer.committed(tp2).offset)
+    assertEquals(300, anotherConsumer.committed(Set(tp).asJava).get(tp).offset)
+    assertEquals(500, anotherConsumer.committed(Set(tp2).asJava).get(tp2).offset)
   }
 
   @Test
-  def testAutoCommitOnCloseAfterWakeup() {
+  def testAutoCommitOnCloseAfterWakeup(): Unit = {
     this.consumerConfig.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "true")
     val consumer = createConsumer()
 
@@ -302,12 +305,12 @@ class PlaintextConsumerTest extends BaseConsumerTest {
 
     // now we should see the committed positions from another consumer
     val anotherConsumer = createConsumer()
-    assertEquals(300, anotherConsumer.committed(tp).offset)
-    assertEquals(500, anotherConsumer.committed(tp2).offset)
+    assertEquals(300, anotherConsumer.committed(Set(tp).asJava).get(tp).offset)
+    assertEquals(500, anotherConsumer.committed(Set(tp2).asJava).get(tp2).offset)
   }
 
   @Test
-  def testAutoOffsetReset() {
+  def testAutoOffsetReset(): Unit = {
     val producer = createProducer()
     sendRecords(producer, numRecords = 1, tp)
 
@@ -317,7 +320,7 @@ class PlaintextConsumerTest extends BaseConsumerTest {
   }
 
   @Test
-  def testGroupConsumption() {
+  def testGroupConsumption(): Unit = {
     val producer = createProducer()
     sendRecords(producer, numRecords = 10, tp)
 
@@ -336,23 +339,23 @@ class PlaintextConsumerTest extends BaseConsumerTest {
    * of that topic are assigned to it.
    */
   @Test
-  def testPatternSubscription() {
+  def testPatternSubscription(): Unit = {
     val numRecords = 10000
     val producer = createProducer()
     sendRecords(producer, numRecords, tp)
 
     val topic1 = "tblablac" // matches subscribed pattern
-    createTopic(topic1, 2, serverCount)
+    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, serverCount)
+    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, serverCount)
+    createTopic(topic3, 2, brokerCount)
     sendRecords(producer, numRecords = 1000, new TopicPartition(topic3, 0))
     sendRecords(producer, numRecords = 1000, new TopicPartition(topic3, 1))
 
@@ -370,7 +373,7 @@ class PlaintextConsumerTest extends BaseConsumerTest {
     awaitAssignment(consumer, assignment)
 
     val topic4 = "tsomec" // matches subscribed pattern
-    createTopic(topic4, 2, serverCount)
+    createTopic(topic4, 2, brokerCount)
     sendRecords(producer, numRecords = 1000, new TopicPartition(topic4, 0))
     sendRecords(producer, numRecords = 1000, new TopicPartition(topic4, 1))
 
@@ -393,7 +396,7 @@ class PlaintextConsumerTest extends BaseConsumerTest {
    * that it is the subscription call that triggers a metadata refresh, and not the timeout.
    */
   @Test
-  def testSubsequentPatternSubscription() {
+  def testSubsequentPatternSubscription(): Unit = {
     this.consumerConfig.setProperty(ConsumerConfig.METADATA_MAX_AGE_CONFIG, "30000")
     val consumer = createConsumer()
 
@@ -404,7 +407,7 @@ class PlaintextConsumerTest extends BaseConsumerTest {
     // the first topic ('topic')  matches first subscription pattern only
 
     val fooTopic = "foo" // matches both subscription patterns
-    createTopic(fooTopic, 1, serverCount)
+    createTopic(fooTopic, 1, brokerCount)
     sendRecords(producer, numRecords = 1000, new TopicPartition(fooTopic, 0))
 
     assertEquals(0, consumer.assignment().size)
@@ -419,7 +422,7 @@ class PlaintextConsumerTest extends BaseConsumerTest {
     awaitAssignment(consumer, assignment)
 
     val barTopic = "bar" // matches the next subscription pattern
-    createTopic(barTopic, 1, serverCount)
+    createTopic(barTopic, 1, brokerCount)
     sendRecords(producer, numRecords = 1000, new TopicPartition(barTopic, 0))
 
     val pattern2 = Pattern.compile("...") // only 'foo' and 'bar' match this
@@ -444,13 +447,13 @@ class PlaintextConsumerTest extends BaseConsumerTest {
    * assignments are cleared right away.
    */
   @Test
-  def testPatternUnsubscription() {
+  def testPatternUnsubscription(): Unit = {
     val numRecords = 10000
     val producer = createProducer()
     sendRecords(producer, numRecords, tp)
 
     val topic1 = "tblablac" // matches the subscription pattern
-    createTopic(topic1, 2, serverCount)
+    createTopic(topic1, 2, brokerCount)
     sendRecords(producer, numRecords = 1000, new TopicPartition(topic1, 0))
     sendRecords(producer, numRecords = 1000, new TopicPartition(topic1, 1))
 
@@ -470,28 +473,28 @@ class PlaintextConsumerTest extends BaseConsumerTest {
   }
 
   @Test
-  def testCommitMetadata() {
+  def testCommitMetadata(): 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(tp))
+    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(tp))
+    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(tp))
+    assertEquals(nullMetadata, consumer.committed(Set(tp).asJava).get(tp))
   }
 
   @Test
-  def testAsyncCommit() {
+  def testAsyncCommit(): Unit = {
     val consumer = createConsumer()
     consumer.assign(List(tp).asJava)
 
@@ -506,27 +509,27 @@ class PlaintextConsumerTest extends BaseConsumerTest {
 
     assertEquals(None, callback.lastError)
     assertEquals(count, callback.successCount)
-    assertEquals(new OffsetAndMetadata(count), consumer.committed(tp))
+    assertEquals(new OffsetAndMetadata(count), consumer.committed(Set(tp).asJava).get(tp))
   }
 
   @Test
-  def testExpandingTopicSubscriptions() {
+  def testExpandingTopicSubscriptions(): 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, serverCount)
+    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)
   }
 
   @Test
-  def testShrinkingTopicSubscriptions() {
+  def testShrinkingTopicSubscriptions(): Unit = {
     val otherTopic = "other"
-    createTopic(otherTopic, 2, serverCount)
+    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)
@@ -538,7 +541,7 @@ class PlaintextConsumerTest extends BaseConsumerTest {
   }
 
   @Test
-  def testPartitionsFor() {
+  def testPartitionsFor(): Unit = {
     val numParts = 2
     createTopic("part-test", numParts, 1)
     val consumer = createConsumer()
@@ -548,20 +551,20 @@ class PlaintextConsumerTest extends BaseConsumerTest {
   }
 
   @Test
-  def testPartitionsForAutoCreate() {
+  def testPartitionsForAutoCreate(): Unit = {
     val consumer = createConsumer()
     val partitions = consumer.partitionsFor("non-exist-topic")
     assertFalse(partitions.isEmpty)
   }
 
   @Test(expected = classOf[InvalidTopicException])
-  def testPartitionsForInvalidTopic() {
+  def testPartitionsForInvalidTopic(): Unit = {
     val consumer = createConsumer()
     consumer.partitionsFor(";3# ads,{234")
   }
 
   @Test
-  def testSeek() {
+  def testSeek(): Unit = {
     val consumer = createConsumer()
     val totalRecords = 50L
     val mid = totalRecords / 2
@@ -603,7 +606,7 @@ class PlaintextConsumerTest extends BaseConsumerTest {
       startingTimestamp = mid.toLong, tp = tp2)
   }
 
-  private def sendCompressedMessages(numRecords: Int, tp: TopicPartition) {
+  private def sendCompressedMessages(numRecords: Int, tp: TopicPartition): Unit = {
     val producerProps = new Properties()
     producerProps.setProperty(ProducerConfig.COMPRESSION_TYPE_CONFIG, CompressionType.GZIP.name)
     producerProps.setProperty(ProducerConfig.LINGER_MS_CONFIG, Int.MaxValue.toString)
@@ -615,28 +618,29 @@ class PlaintextConsumerTest extends BaseConsumerTest {
   }
 
   @Test
-  def testPositionAndCommit() {
+  def testPositionAndCommit(): Unit = {
     val producer = createProducer()
     sendRecords(producer, numRecords = 5, tp)
 
+    val topicPartition = new TopicPartition(topic, 15)
     val consumer = createConsumer()
-    assertNull(consumer.committed(new TopicPartition(topic, 15)))
+    assertNull(consumer.committed(Set(topicPartition).asJava).get(topicPartition))
 
     // position() on a partition that we aren't subscribed to throws an exception
     intercept[IllegalStateException] {
-      consumer.position(new TopicPartition(topic, 15))
+      consumer.position(topicPartition)
     }
 
     consumer.assign(List(tp).asJava)
 
     assertEquals("position() on a partition that we are subscribed to should reset the offset", 0L, consumer.position(tp))
     consumer.commitSync()
-    assertEquals(0L, consumer.committed(tp).offset)
+    assertEquals(0L, consumer.committed(Set(tp).asJava).get(tp).offset)
 
     consumeAndVerifyRecords(consumer = consumer, numRecords = 5, startingOffset = 0)
     assertEquals("After consuming 5 records, position should be 5", 5L, consumer.position(tp))
     consumer.commitSync()
-    assertEquals("Committed offset should be returned", 5L, consumer.committed(tp).offset)
+    assertEquals("Committed offset should be returned", 5L, consumer.committed(Set(tp).asJava).get(tp).offset)
 
     sendRecords(producer, numRecords = 1, tp)
 
@@ -647,7 +651,7 @@ class PlaintextConsumerTest extends BaseConsumerTest {
   }
 
   @Test
-  def testPartitionPauseAndResume() {
+  def testPartitionPauseAndResume(): Unit = {
     val partitions = List(tp).asJava
     val producer = createProducer()
     sendRecords(producer, numRecords = 5, tp)
@@ -663,7 +667,7 @@ class PlaintextConsumerTest extends BaseConsumerTest {
   }
 
   @Test
-  def testFetchInvalidOffset() {
+  def testFetchInvalidOffset(): Unit = {
     this.consumerConfig.setProperty(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "none")
     val consumer = createConsumer()
 
@@ -691,7 +695,7 @@ class PlaintextConsumerTest extends BaseConsumerTest {
   }
 
   @Test
-  def testFetchRecordLargerThanFetchMaxBytes() {
+  def testFetchRecordLargerThanFetchMaxBytes(): Unit = {
     val maxFetchBytes = 10 * 1024
     this.consumerConfig.setProperty(ConsumerConfig.FETCH_MAX_BYTES_CONFIG, maxFetchBytes.toString)
     checkLargeRecord(maxFetchBytes + 1)
@@ -759,7 +763,7 @@ class PlaintextConsumerTest extends BaseConsumerTest {
   }
 
   @Test
-  def testFetchRecordLargerThanMaxPartitionFetchBytes() {
+  def testFetchRecordLargerThanMaxPartitionFetchBytes(): Unit = {
     val maxPartitionFetchBytes = 10 * 1024
     this.consumerConfig.setProperty(ConsumerConfig.MAX_PARTITION_FETCH_BYTES_CONFIG, maxPartitionFetchBytes.toString)
     checkLargeRecord(maxPartitionFetchBytes + 1)
@@ -781,7 +785,7 @@ class PlaintextConsumerTest extends BaseConsumerTest {
     val partitionCount = 30
     val topics = Seq(topic1, topic2, topic3)
     topics.foreach { topicName =>
-      createTopic(topicName, partitionCount, serverCount)
+      createTopic(topicName, partitionCount, brokerCount)
     }
 
     val partitions = topics.flatMap { topic =>
@@ -810,7 +814,7 @@ class PlaintextConsumerTest extends BaseConsumerTest {
   }
 
   @Test
-  def testRoundRobinAssignment() {
+  def testRoundRobinAssignment(): 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)
@@ -846,7 +850,7 @@ class PlaintextConsumerTest extends BaseConsumerTest {
   }
 
   @Test
-  def testMultiConsumerRoundRobinAssignment() {
+  def testMultiConsumerRoundRobinAssignment(): Unit = {
     this.consumerConfig.setProperty(ConsumerConfig.GROUP_ID_CONFIG, "roundrobin-group")
     this.consumerConfig.setProperty(ConsumerConfig.PARTITION_ASSIGNMENT_STRATEGY_CONFIG, classOf[RoundRobinAssignor].getName)
 
@@ -861,11 +865,11 @@ class PlaintextConsumerTest extends BaseConsumerTest {
     // for the topic partition assignment
     val (consumerGroup, consumerPollers) = createConsumerGroupAndWaitForAssignment(10, List(topic1, topic2), subscriptions)
     try {
-      validateGroupAssignment(consumerPollers, subscriptions, s"Did not get valid initial assignment for partitions ${subscriptions.asJava}")
+      validateGroupAssignment(consumerPollers, subscriptions)
 
       // add one more consumer and validate re-assignment
       addConsumersToGroupAndWaitForGroupAssignment(1, consumerGroup, consumerPollers,
-        List(topic1, topic2), subscriptions)
+        List(topic1, topic2), subscriptions, "roundrobin-group")
     } finally {
       consumerPollers.foreach(_.shutdown())
     }
@@ -883,7 +887,7 @@ class PlaintextConsumerTest extends BaseConsumerTest {
    *    will move to consumer #10, leading to a total of (#par mod 9) partition movement
    */
   @Test
-  def testMultiConsumerStickyAssignment() {
+  def testMultiConsumerStickyAssignment(): Unit = {
 
     def reverse(m: Map[Long, Set[TopicPartition]]) =
       m.values.toSet.flatten.map(v => (v, m.keys.filter(m(_).contains(v)).head)).toMap
@@ -900,11 +904,11 @@ class PlaintextConsumerTest extends BaseConsumerTest {
     // 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, s"Did not get valid initial assignment for partitions ${partitions.asJava}")
+    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)
+    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)
@@ -929,7 +933,7 @@ class PlaintextConsumerTest extends BaseConsumerTest {
    * As a result, it is testing the default assignment strategy set by BaseConsumerTest
    */
   @Test
-  def testMultiConsumerDefaultAssignment() {
+  def testMultiConsumerDefaultAssignment(): Unit = {
     // use consumers and topics defined in this class + one more topic
     val producer = createProducer()
     sendRecords(producer, numRecords = 100, tp)
@@ -945,7 +949,7 @@ class PlaintextConsumerTest extends BaseConsumerTest {
 
     val consumerPollers = subscribeConsumers(consumersInGroup, List(topic, topic1))
     try {
-      validateGroupAssignment(consumerPollers, subscriptions, s"Did not get valid initial assignment for partitions ${subscriptions.asJava}")
+      validateGroupAssignment(consumerPollers, subscriptions)
 
       // add 2 more consumers and validate re-assignment
       addConsumersToGroupAndWaitForGroupAssignment(2, consumersInGroup, consumerPollers, List(topic, topic1), subscriptions)
@@ -974,7 +978,7 @@ class PlaintextConsumerTest extends BaseConsumerTest {
   }
 
   @Test
-  def testInterceptors() {
+  def testInterceptors(): Unit = {
     val appendStr = "mock"
     MockConsumerInterceptor.resetCounters()
     MockProducerInterceptor.resetCounters()
@@ -1021,12 +1025,12 @@ class PlaintextConsumerTest extends BaseConsumerTest {
     // commit sync and verify onCommit is called
     val commitCountBefore = MockConsumerInterceptor.ON_COMMIT_COUNT.intValue
     testConsumer.commitSync(Map[TopicPartition, OffsetAndMetadata]((tp, new OffsetAndMetadata(2L))).asJava)
-    assertEquals(2, testConsumer.committed(tp).offset)
+    assertEquals(2, testConsumer.committed(Set(tp).asJava).get(tp).offset)
     assertEquals(commitCountBefore + 1, MockConsumerInterceptor.ON_COMMIT_COUNT.intValue)
 
     // commit async and verify onCommit is called
     sendAndAwaitAsyncCommit(testConsumer, Some(Map(tp -> new OffsetAndMetadata(5L))))
-    assertEquals(5, testConsumer.committed(tp).offset)
+    assertEquals(5, testConsumer.committed(Set(tp).asJava).get(tp).offset)
     assertEquals(commitCountBefore + 2, MockConsumerInterceptor.ON_COMMIT_COUNT.intValue)
 
     testConsumer.close()
@@ -1038,9 +1042,9 @@ class PlaintextConsumerTest extends BaseConsumerTest {
   }
 
   @Test
-  def testAutoCommitIntercept() {
+  def testAutoCommitIntercept(): Unit = {
     val topic2 = "topic2"
-    createTopic(topic2, 2, serverCount)
+    createTopic(topic2, 2, brokerCount)
 
     // produce records
     val numRecords = 100
@@ -1073,8 +1077,8 @@ class PlaintextConsumerTest extends BaseConsumerTest {
                                                     rebalanceListener)
 
     // after rebalancing, we should have reset to the committed positions
-    assertEquals(10, testConsumer.committed(tp).offset)
-    assertEquals(20, testConsumer.committed(tp2).offset)
+    assertEquals(10, testConsumer.committed(Set(tp).asJava).get(tp).offset)
+    assertEquals(20, testConsumer.committed(Set(tp2).asJava).get(tp2).offset)
     assertTrue(MockConsumerInterceptor.ON_COMMIT_COUNT.intValue() > commitCountBeforeRebalance)
 
     // verify commits are intercepted on close
@@ -1088,7 +1092,7 @@ class PlaintextConsumerTest extends BaseConsumerTest {
   }
 
   @Test
-  def testInterceptorsWithWrongKeyValue() {
+  def testInterceptorsWithWrongKeyValue(): Unit = {
     val appendStr = "mock"
     // create producer with interceptor that has different key and value types from the producer
     val producerProps = new Properties()
@@ -1114,7 +1118,7 @@ class PlaintextConsumerTest extends BaseConsumerTest {
   }
 
   @Test
-  def testConsumeMessagesWithCreateTime() {
+  def testConsumeMessagesWithCreateTime(): Unit = {
     val numRecords = 50
     // Test non-compressed messages
     val producer = createProducer()
@@ -1132,7 +1136,7 @@ class PlaintextConsumerTest extends BaseConsumerTest {
   }
 
   @Test
-  def testConsumeMessagesWithLogAppendTime() {
+  def testConsumeMessagesWithLogAppendTime(): Unit = {
     val topicName = "testConsumeMessagesWithLogAppendTime"
     val topicProps = new Properties()
     topicProps.setProperty(LogConfig.MessageTimestampTypeProp, "LogAppendTime")
@@ -1160,7 +1164,7 @@ class PlaintextConsumerTest extends BaseConsumerTest {
   }
 
   @Test
-  def testListTopics() {
+  def testListTopics(): Unit = {
     val numParts = 2
     val topic1 = "part-test-topic-1"
     val topic2 = "part-test-topic-2"
@@ -1180,7 +1184,7 @@ class PlaintextConsumerTest extends BaseConsumerTest {
   }
 
   @Test
-  def testOffsetsForTimes() {
+  def testOffsetsForTimes(): Unit = {
     val numParts = 2
     val topic1 = "part-test-topic-1"
     val topic2 = "part-test-topic-2"
@@ -1243,7 +1247,7 @@ class PlaintextConsumerTest extends BaseConsumerTest {
   }
 
   @Test
-  def testEarliestOrLatestOffsets() {
+  def testEarliestOrLatestOffsets(): Unit = {
     val topic0 = "topicWithNewMessageFormat"
     val topic1 = "topicWithOldMessageFormat"
     val producer = createProducer()
@@ -1271,7 +1275,7 @@ class PlaintextConsumerTest extends BaseConsumerTest {
   }
 
   @Test
-  def testUnsubscribeTopic() {
+  def testUnsubscribeTopic(): 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()
@@ -1287,7 +1291,7 @@ class PlaintextConsumerTest extends BaseConsumerTest {
   }
 
   @Test
-  def testPauseStateNotPreservedByRebalance() {
+  def testPauseStateNotPreservedByRebalance(): 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()
@@ -1307,7 +1311,7 @@ class PlaintextConsumerTest extends BaseConsumerTest {
   }
 
   @Test
-  def testCommitSpecifiedOffsets() {
+  def testCommitSpecifiedOffsets(): Unit = {
     val producer = createProducer()
     sendRecords(producer, numRecords = 5, tp)
     sendRecords(producer, numRecords = 7, tp2)
@@ -1318,25 +1322,25 @@ class PlaintextConsumerTest extends BaseConsumerTest {
     val pos1 = consumer.position(tp)
     val pos2 = consumer.position(tp2)
     consumer.commitSync(Map[TopicPartition, OffsetAndMetadata]((tp, new OffsetAndMetadata(3L))).asJava)
-    assertEquals(3, consumer.committed(tp).offset)
-    assertNull(consumer.committed(tp2))
+    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(tp).offset)
-    assertEquals(5, consumer.committed(tp2).offset)
+    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(tp2).offset)
+    assertEquals(7, consumer.committed(Set(tp2).asJava).get(tp2).offset)
   }
 
   @Test
-  def testAutoCommitOnRebalance() {
+  def testAutoCommitOnRebalance(): Unit = {
     val topic2 = "topic2"
-    createTopic(topic2, 2, serverCount)
+    createTopic(topic2, 2, brokerCount)
 
     this.consumerConfig.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "true")
     val consumer = createConsumer()
@@ -1368,15 +1372,15 @@ class PlaintextConsumerTest extends BaseConsumerTest {
     awaitAssignment(consumer, newAssignment)
 
     // after rebalancing, we should have reset to the committed positions
-    assertEquals(300, consumer.committed(tp).offset)
-    assertEquals(500, consumer.committed(tp2).offset)
+    assertEquals(300, consumer.committed(Set(tp).asJava).get(tp).offset)
+    assertEquals(500, consumer.committed(Set(tp2).asJava).get(tp2).offset)
   }
 
   @Test
-  def testPerPartitionLeadMetricsCleanUpWithSubscribe() {
+  def testPerPartitionLeadMetricsCleanUpWithSubscribe(): Unit = {
     val numMessages = 1000
     val topic2 = "topic2"
-    createTopic(topic2, 2, serverCount)
+    createTopic(topic2, 2, brokerCount)
     // send some messages.
     val producer = createProducer()
     sendRecords(producer, numMessages, tp)
@@ -1412,10 +1416,10 @@ class PlaintextConsumerTest extends BaseConsumerTest {
   }
 
   @Test
-  def testPerPartitionLagMetricsCleanUpWithSubscribe() {
+  def testPerPartitionLagMetricsCleanUpWithSubscribe(): Unit = {
     val numMessages = 1000
     val topic2 = "topic2"
-    createTopic(topic2, 2, serverCount)
+    createTopic(topic2, 2, brokerCount)
     // send some messages.
     val producer = createProducer()
     sendRecords(producer, numMessages, tp)
@@ -1452,7 +1456,7 @@ class PlaintextConsumerTest extends BaseConsumerTest {
   }
 
   @Test
-  def testPerPartitionLeadMetricsCleanUpWithAssign() {
+  def testPerPartitionLeadMetricsCleanUpWithAssign(): Unit = {
     val numMessages = 1000
     // Test assign
     // send some messages.
@@ -1481,7 +1485,7 @@ class PlaintextConsumerTest extends BaseConsumerTest {
   }
 
   @Test
-  def testPerPartitionLagMetricsCleanUpWithAssign() {
+  def testPerPartitionLagMetricsCleanUpWithAssign(): Unit = {
     val numMessages = 1000
     // Test assign
     // send some messages.
@@ -1512,7 +1516,7 @@ class PlaintextConsumerTest extends BaseConsumerTest {
   }
 
   @Test
-  def testPerPartitionLagMetricsWhenReadCommitted() {
+  def testPerPartitionLagMetricsWhenReadCommitted(): Unit = {
     val numMessages = 1000
     // send some messages.
     val producer = createProducer()
@@ -1535,7 +1539,7 @@ class PlaintextConsumerTest extends BaseConsumerTest {
   }
 
   @Test
-  def testPerPartitionLeadWithMaxPollRecords() {
+  def testPerPartitionLeadWithMaxPollRecords(): Unit = {
     val numMessages = 1000
     val maxPollRecords = 10
     val producer = createProducer()
@@ -1557,7 +1561,7 @@ class PlaintextConsumerTest extends BaseConsumerTest {
   }
 
   @Test
-  def testPerPartitionLagWithMaxPollRecords() {
+  def testPerPartitionLagWithMaxPollRecords(): Unit = {
     val numMessages = 1000
     val maxPollRecords = 10
     val producer = createProducer()
@@ -1580,7 +1584,7 @@ class PlaintextConsumerTest extends BaseConsumerTest {
   }
 
   @Test
-  def testQuotaMetricsNotCreatedIfNoQuotasConfigured() {
+  def testQuotaMetricsNotCreatedIfNoQuotasConfigured(): Unit = {
     val numRecords = 1000
     val producer = createProducer()
     sendRecords(producer, numRecords, tp)
@@ -1590,7 +1594,7 @@ class PlaintextConsumerTest extends BaseConsumerTest {
     consumer.seek(tp, 0)
     consumeAndVerifyRecords(consumer = consumer, numRecords = numRecords, startingOffset = 0)
 
-    def assertNoMetric(broker: KafkaServer, name: String, quotaType: QuotaType, clientId: String) {
+    def assertNoMetric(broker: KafkaServer, name: String, quotaType: QuotaType, clientId: String): Unit = {
         val metricName = broker.metrics.metricName("throttle-time",
                                   quotaType.toString,
                                   "",
@@ -1608,7 +1612,7 @@ class PlaintextConsumerTest extends BaseConsumerTest {
     servers.foreach(assertNoMetric(_, "request-time", QuotaType.Request, consumerClientId))
     servers.foreach(assertNoMetric(_, "throttle-time", QuotaType.Request, consumerClientId))
 
-    def assertNoExemptRequestMetric(broker: KafkaServer) {
+    def assertNoExemptRequestMetric(broker: KafkaServer): Unit = {
         val metricName = broker.metrics.metricName("exempt-request-time", QuotaType.Request.toString, "")
         assertNull("Metric should not hanve been created " + metricName, broker.metrics.metric(metricName))
     }
@@ -1635,22 +1639,40 @@ class PlaintextConsumerTest extends BaseConsumerTest {
     consumerPollers += timeoutPoller
 
     // validate the initial assignment
-    validateGroupAssignment(consumerPollers, subscriptions, s"Did not get valid initial assignment for partitions ${subscriptions.asJava}")
+    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()
 
-    val maxSessionTimeout = this.serverConfig.getProperty(KafkaConfig.GroupMaxSessionTimeoutMsProp).toLong
     validateGroupAssignment(consumerPollers, subscriptions,
-      s"Did not get valid assignment for partitions ${subscriptions.asJava} after one consumer left", 3 * maxSessionTimeout)
+      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[KafkaConsumer[Array[Byte], Array[Byte]]],
+                         topicsToSubscribe: List[String]): mutable.Buffer[ConsumerAssignmentPoller] = {
+    val consumerPollers = mutable.Buffer[ConsumerAssignmentPoller]()
+    for (consumer <- consumerGroup)
+      consumerPollers += subscribeConsumerAndStartPolling(consumer, topicsToSubscribe)
+    consumerPollers
+  }
+
   /**
    * Creates topic 'topicName' with 'numPartitions' partitions and produces 'recordsPerPartition'
    * records to each partition
@@ -1659,7 +1681,7 @@ class PlaintextConsumerTest extends BaseConsumerTest {
                                 topicName: String,
                                 numPartitions: Int,
                                 recordsPerPartition: Int): Set[TopicPartition] = {
-    createTopic(topicName, numPartitions, serverCount)
+    createTopic(topicName, numPartitions, brokerCount)
     var parts = Set[TopicPartition]()
     for (partition <- 0 until numPartitions) {
       val tp = new TopicPartition(topicName, partition)
@@ -1670,51 +1692,16 @@ class PlaintextConsumerTest extends BaseConsumerTest {
   }
 
   /**
-   * Subscribes consumer 'consumer' to a given list of topics 'topicsToSubscribe', creates
-   * consumer poller and starts polling.
-   * Assumes that the consumer is not subscribed to any topics yet
-   *
-   * @param consumer consumer
-   * @param topicsToSubscribe topics that this consumer will subscribe to
-   * @return consumer poller for the given consumer
-   */
-  def subscribeConsumerAndStartPolling(consumer: Consumer[Array[Byte], Array[Byte]],
-                                       topicsToSubscribe: List[String]): ConsumerAssignmentPoller = {
-    assertEquals(0, consumer.assignment().size)
-    val consumerPoller = new ConsumerAssignmentPoller(consumer, topicsToSubscribe)
-    consumerPoller.start()
-    consumerPoller
-  }
-
-  /**
-   * 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: Buffer[KafkaConsumer[Array[Byte], Array[Byte]]],
-                         topicsToSubscribe: List[String]): Buffer[ConsumerAssignmentPoller] = {
-    val consumerPollers = 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
-   */
+    * 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[KafkaConsumer[Array[Byte], Array[Byte]]], Buffer[ConsumerAssignmentPoller]) = {
@@ -1728,54 +1715,6 @@ class PlaintextConsumerTest extends BaseConsumerTest {
     (consumerGroup, consumerPollers)
   }
 
-  /**
-   * Create 'numOfConsumersToAdd' consumers add then to the consumer group 'consumerGroup', and create corresponding
-   * pollers for these consumers. Wait for partition re-assignment and validate.
-   *
-   * Currently, assignment validation requires that total number of partitions is greater or equal to
-   * number of consumers, so subscriptions.size must be greater or equal the resulting number of consumers in the group
-   *
-   * @param numOfConsumersToAdd number of consumers to create and add to the consumer group
-   * @param consumerGroup current consumer group
-   * @param consumerPollers current consumer pollers
-   * @param topicsToSubscribe topics to which new consumers will subscribe to
-   * @param subscriptions set of all topic partitions
-   */
-  def addConsumersToGroupAndWaitForGroupAssignment(numOfConsumersToAdd: Int,
-                                                   consumerGroup: Buffer[KafkaConsumer[Array[Byte], Array[Byte]]],
-                                                   consumerPollers: Buffer[ConsumerAssignmentPoller],
-                                                   topicsToSubscribe: List[String],
-                                                   subscriptions: Set[TopicPartition]): Unit = {
-    assertTrue(consumerGroup.size + numOfConsumersToAdd <= subscriptions.size)
-    for (_ <- 0 until numOfConsumersToAdd) {
-      val consumer = createConsumer()
-      consumerGroup += consumer
-      consumerPollers += subscribeConsumerAndStartPolling(consumer, topicsToSubscribe)
-    }
-
-    // wait until topics get re-assigned and validate assignment
-    validateGroupAssignment(consumerPollers, subscriptions,
-      s"Did not get valid assignment for partitions ${subscriptions.asJava} after we added $numOfConsumersToAdd consumer(s)")
-  }
-
-  /**
-   * Wait for consumers to get partition assignment and validate it.
-   *
-   * @param consumerPollers consumer pollers corresponding to the consumer group we are testing
-   * @param subscriptions set of all topic partitions
-   * @param msg message to print when waiting for/validating assignment fails
-   */
-  def validateGroupAssignment(consumerPollers: Buffer[ConsumerAssignmentPoller],
-                              subscriptions: Set[TopicPartition],
-                              msg: String,
-                              waitTime: Long = 10000L): Unit = {
-    TestUtils.waitUntilTrue(() => {
-      val assignments = Buffer[Set[TopicPartition]]()
-      consumerPollers.foreach(assignments += _.consumerAssignment())
-      isPartitionAssignmentValid(assignments, subscriptions)
-    }, msg, waitTime)
-  }
-
   def changeConsumerGroupSubscriptionAndValidateAssignment(consumerPollers: Buffer[ConsumerAssignmentPoller],
                                                            topicsToSubscribe: List[String],
                                                            subscriptions: Set[TopicPartition]): Unit = {
@@ -1785,11 +1724,11 @@ class PlaintextConsumerTest extends BaseConsumerTest {
     // 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)
+      consumerPollers.forall { poller => poller.isSubscribeRequestProcessed }
+    }, s"Failed to call subscribe on all consumers in the group for subscription $subscriptions", 1000L)
 
     validateGroupAssignment(consumerPollers, subscriptions,
-      s"Did not get valid assignment for partitions ${subscriptions.asJava} after we changed subscription")
+      Some(s"Did not get valid assignment for partitions ${subscriptions.asJava} after we changed subscription"))
   }
 
   def changeConsumerSubscriptionAndValidateAssignment[K, V](consumer: Consumer[K, V],
@@ -1870,7 +1809,7 @@ class PlaintextConsumerTest extends BaseConsumerTest {
     }
 
     try {
-      consumer2.committed(tp)
+      consumer2.committed(Set(tp).asJava)
       fail("Expected committed offset fetch to fail due to null group id")
     } catch {
       case e: InvalidGroupIdException => // OK
diff --git a/core/src/test/scala/integration/kafka/api/PlaintextEndToEndAuthorizationTest.scala b/core/src/test/scala/integration/kafka/api/PlaintextEndToEndAuthorizationTest.scala
index fc1853bf018f2..b8549c577c69a 100644
--- a/core/src/test/scala/integration/kafka/api/PlaintextEndToEndAuthorizationTest.scala
+++ b/core/src/test/scala/integration/kafka/api/PlaintextEndToEndAuthorizationTest.scala
@@ -22,6 +22,7 @@ import org.apache.kafka.common.security.auth._
 import org.junit.{Before, Test}
 import org.junit.Assert._
 import org.apache.kafka.common.errors.TopicAuthorizationException
+import org.scalatest.Assertions.intercept
 
 // This test case uses a separate listener for client and inter-broker communication, from
 // which we derive corresponding principals
@@ -70,13 +71,13 @@ class PlaintextEndToEndAuthorizationTest extends EndToEndAuthorizationTest {
   override val kafkaPrincipal = "server"
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     startSasl(jaasSections(List.empty, None, ZkSasl))
     super.setUp()
   }
 
   @Test
-  def testListenerName() {
+  def testListenerName(): Unit = {
     // To check the client listener name, establish a session on the server by sending any request eg sendRecords
     val producer = createProducer()
     intercept[TopicAuthorizationException](sendRecords(producer, numRecords = 1, tp))
diff --git a/core/src/test/scala/integration/kafka/api/PlaintextProducerSendTest.scala b/core/src/test/scala/integration/kafka/api/PlaintextProducerSendTest.scala
index 8ab32af2ed2fb..717154c7d2749 100644
--- a/core/src/test/scala/integration/kafka/api/PlaintextProducerSendTest.scala
+++ b/core/src/test/scala/integration/kafka/api/PlaintextProducerSendTest.scala
@@ -31,35 +31,38 @@ import org.junit.Test
 class PlaintextProducerSendTest extends BaseProducerSendTest {
 
   @Test(expected = classOf[SerializationException])
-  def testWrongSerializer() {
+  def testWrongSerializer(): Unit = {
     val producerProps = new Properties()
     producerProps.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, brokerList)
     producerProps.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.StringSerializer")
     producerProps.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.StringSerializer")
+    producerProps.put(ProducerConfig.ALLOW_AUTO_CREATE_TOPICS_CONFIG, boolean2Boolean(true))
     val producer = registerProducer(new KafkaProducer(producerProps))
     val record = new ProducerRecord[Array[Byte], Array[Byte]](topic, 0, "key".getBytes, "value".getBytes)
     producer.send(record)
   }
 
   @Test
-  def testBatchSizeZero() {
+  def testBatchSizeZero(): Unit = {
     val producer = createProducer(brokerList = brokerList,
       lingerMs = Int.MaxValue,
+      deliveryTimeoutMs = Int.MaxValue,
       batchSize = 0)
     sendAndVerify(producer)
   }
 
   @Test
-  def testSendCompressedMessageWithLogAppendTime() {
+  def testSendCompressedMessageWithLogAppendTime(): Unit = {
     val producer = createProducer(brokerList = brokerList,
       compressionType = "gzip",
-      lingerMs = Int.MaxValue)
+      lingerMs = Int.MaxValue,
+      deliveryTimeoutMs = Int.MaxValue)
     sendAndVerifyTimestamp(producer, TimestampType.LOG_APPEND_TIME)
   }
 
   @Test
-  def testSendNonCompressedMessageWithLogAppendTime() {
-    val producer = createProducer(brokerList = brokerList, lingerMs = Int.MaxValue)
+  def testSendNonCompressedMessageWithLogAppendTime(): Unit = {
+    val producer = createProducer(brokerList = brokerList, lingerMs = Int.MaxValue, deliveryTimeoutMs = Int.MaxValue)
     sendAndVerifyTimestamp(producer, TimestampType.LOG_APPEND_TIME)
   }
 
@@ -69,7 +72,7 @@ class PlaintextProducerSendTest extends BaseProducerSendTest {
    * The topic should be created upon sending the first message
    */
   @Test
-  def testAutoCreateTopic() {
+  def testAutoCreateTopic(): Unit = {
     val producer = createProducer(brokerList)
     try {
       // Send a message to auto-create the topic
@@ -85,7 +88,7 @@ class PlaintextProducerSendTest extends BaseProducerSendTest {
   }
 
   @Test
-  def testSendWithInvalidCreateTime() {
+  def testSendWithInvalidCreateTime(): Unit = {
     val topicProps = new Properties()
     topicProps.setProperty(LogConfig.MessageTimestampDifferenceMaxMsProp, "1000")
     createTopic(topic, 1, 2, topicProps)
@@ -111,5 +114,4 @@ class PlaintextProducerSendTest extends BaseProducerSendTest {
       compressedProducer.close()
     }
   }
-
 }
diff --git a/core/src/test/scala/integration/kafka/api/ProducerCompressionTest.scala b/core/src/test/scala/integration/kafka/api/ProducerCompressionTest.scala
index 24193521f91ee..d65c43e71b556 100755
--- a/core/src/test/scala/integration/kafka/api/ProducerCompressionTest.scala
+++ b/core/src/test/scala/integration/kafka/api/ProducerCompressionTest.scala
@@ -42,14 +42,14 @@ class ProducerCompressionTest(compression: String) extends ZooKeeperTestHarness
   private var server: KafkaServer = null
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     super.setUp()
     val props = TestUtils.createBrokerConfig(brokerId, zkConnect)
     server = TestUtils.createServer(KafkaConfig.fromProps(props))
   }
 
   @After
-  override def tearDown() {
+  override def tearDown(): Unit = {
     TestUtils.shutdownServers(Seq(server))
     super.tearDown()
   }
@@ -60,7 +60,7 @@ class ProducerCompressionTest(compression: String) extends ZooKeeperTestHarness
    * Compressed messages should be able to sent and consumed correctly
    */
   @Test
-  def testCompression() {
+  def testCompression(): Unit = {
 
     val producerProps = new Properties()
     val bootstrapServers = TestUtils.getBrokerListStrFromServers(Seq(server))
diff --git a/core/src/test/scala/integration/kafka/api/ProducerFailureHandlingTest.scala b/core/src/test/scala/integration/kafka/api/ProducerFailureHandlingTest.scala
index 17d68d119cfe1..beb52323e4689 100644
--- a/core/src/test/scala/integration/kafka/api/ProducerFailureHandlingTest.scala
+++ b/core/src/test/scala/integration/kafka/api/ProducerFailureHandlingTest.scala
@@ -30,6 +30,7 @@ import org.apache.kafka.common.internals.Topic
 import org.apache.kafka.common.record.{DefaultRecord, DefaultRecordBatch}
 import org.junit.Assert._
 import org.junit.{After, Before, Test}
+import org.scalatest.Assertions.intercept
 
 class ProducerFailureHandlingTest extends KafkaServerTestHarness {
   private val producerBufferSize = 30000
@@ -60,7 +61,7 @@ class ProducerFailureHandlingTest extends KafkaServerTestHarness {
   private val topic2 = "topic-2"
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     super.setUp()
 
     producer1 = TestUtils.createProducer(brokerList, acks = 0, retries = 0, requestTimeoutMs = 30000, maxBlockMs = 10000L,
@@ -72,7 +73,7 @@ class ProducerFailureHandlingTest extends KafkaServerTestHarness {
   }
 
   @After
-  override def tearDown() {
+  override def tearDown(): Unit = {
     if (producer1 != null) producer1.close()
     if (producer2 != null) producer2.close()
     if (producer3 != null) producer3.close()
@@ -85,7 +86,7 @@ class ProducerFailureHandlingTest extends KafkaServerTestHarness {
    * With ack == 0 the future metadata will have no exceptions with offset -1
    */
   @Test
-  def testTooLargeRecordWithAckZero() {
+  def testTooLargeRecordWithAckZero(): Unit = {
     // create topic
     createTopic(topic1, replicationFactor = numServers)
 
@@ -102,7 +103,7 @@ class ProducerFailureHandlingTest extends KafkaServerTestHarness {
    * With ack == 1 the future metadata will throw ExecutionException caused by RecordTooLargeException
    */
   @Test
-  def testTooLargeRecordWithAckOne() {
+  def testTooLargeRecordWithAckOne(): Unit = {
     // create topic
     createTopic(topic1, replicationFactor = numServers)
 
@@ -113,7 +114,7 @@ class ProducerFailureHandlingTest extends KafkaServerTestHarness {
     }
   }
 
-  private def checkTooLargeRecordForReplicationWithAckAll(maxFetchSize: Int) {
+  private def checkTooLargeRecordForReplicationWithAckAll(maxFetchSize: Int): Unit = {
     val maxMessageSize = maxFetchSize + 100
     val topicConfig = new Properties
     topicConfig.setProperty(LogConfig.MinInSyncReplicasProp, numServers.toString)
@@ -133,13 +134,13 @@ class ProducerFailureHandlingTest extends KafkaServerTestHarness {
 
   /** This should succeed as the replica fetcher thread can handle oversized messages since KIP-74 */
   @Test
-  def testPartitionTooLargeForReplicationWithAckAll() {
+  def testPartitionTooLargeForReplicationWithAckAll(): Unit = {
     checkTooLargeRecordForReplicationWithAckAll(replicaFetchMaxPartitionBytes)
   }
 
   /** This should succeed as the replica fetcher thread can handle oversized messages since KIP-74 */
   @Test
-  def testResponseTooLargeForReplicationWithAckAll() {
+  def testResponseTooLargeForReplicationWithAckAll(): Unit = {
     checkTooLargeRecordForReplicationWithAckAll(replicaFetchMaxResponseBytes)
   }
 
@@ -147,7 +148,7 @@ class ProducerFailureHandlingTest extends KafkaServerTestHarness {
    * With non-exist-topic the future metadata should return ExecutionException caused by TimeoutException
    */
   @Test
-  def testNonExistentTopic() {
+  def testNonExistentTopic(): Unit = {
     // send a record with non-exist topic
     val record = new ProducerRecord(topic2, null, "key".getBytes, "value".getBytes)
     intercept[ExecutionException] {
@@ -166,7 +167,7 @@ class ProducerFailureHandlingTest extends KafkaServerTestHarness {
    *    TimeoutException
    */
   @Test
-  def testWrongBrokerList() {
+  def testWrongBrokerList(): Unit = {
     // create topic
     createTopic(topic1, replicationFactor = numServers)
 
@@ -185,7 +186,7 @@ class ProducerFailureHandlingTest extends KafkaServerTestHarness {
     * when partition is higher than the upper bound of partitions.
     */
   @Test
-  def testInvalidPartition() {
+  def testInvalidPartition(): Unit = {
     // create topic with a single partition
     createTopic(topic1, numPartitions = 1, replicationFactor = numServers)
 
@@ -203,7 +204,7 @@ class ProducerFailureHandlingTest extends KafkaServerTestHarness {
    * The send call after producer closed should throw IllegalStateException
    */
   @Test
-  def testSendAfterClosed() {
+  def testSendAfterClosed(): Unit = {
     // create topic
     createTopic(topic1, replicationFactor = numServers)
 
@@ -229,7 +230,7 @@ class ProducerFailureHandlingTest extends KafkaServerTestHarness {
   }
 
   @Test
-  def testCannotSendToInternalTopic() {
+  def testCannotSendToInternalTopic(): Unit = {
     TestUtils.createOffsetsTopic(zkClient, servers)
     val thrown = intercept[ExecutionException] {
       producer2.send(new ProducerRecord(Topic.GROUP_METADATA_TOPIC_NAME, "test".getBytes, "test".getBytes)).get
@@ -238,7 +239,7 @@ class ProducerFailureHandlingTest extends KafkaServerTestHarness {
   }
 
   @Test
-  def testNotEnoughReplicas() {
+  def testNotEnoughReplicas(): Unit = {
     val topicName = "minisrtest"
     val topicProps = new Properties()
     topicProps.put("min.insync.replicas",(numServers+1).toString)
@@ -258,7 +259,7 @@ class ProducerFailureHandlingTest extends KafkaServerTestHarness {
   }
 
   @Test
-  def testNotEnoughReplicasAfterBrokerShutdown() {
+  def testNotEnoughReplicasAfterBrokerShutdown(): Unit = {
     val topicName = "minisrtest2"
     val topicProps = new Properties()
     topicProps.put("min.insync.replicas", numServers.toString)
diff --git a/core/src/test/scala/integration/kafka/api/RackAwareAutoTopicCreationTest.scala b/core/src/test/scala/integration/kafka/api/RackAwareAutoTopicCreationTest.scala
index 5fc626bb0bbe5..4ba3fd5eac945 100644
--- a/core/src/test/scala/integration/kafka/api/RackAwareAutoTopicCreationTest.scala
+++ b/core/src/test/scala/integration/kafka/api/RackAwareAutoTopicCreationTest.scala
@@ -43,7 +43,7 @@ class RackAwareAutoTopicCreationTest extends KafkaServerTestHarness with RackAwa
   private val topic = "topic"
 
   @Test
-  def testAutoCreateTopic() {
+  def testAutoCreateTopic(): Unit = {
     val producer = TestUtils.createProducer(brokerList)
     try {
       // Send a message to auto-create the topic
diff --git a/core/src/test/scala/integration/kafka/api/RecordHeaderProducerSendTest.scala b/core/src/test/scala/integration/kafka/api/RecordHeaderProducerSendTest.scala
new file mode 100644
index 0000000000000..6475d1620f418
--- /dev/null
+++ b/core/src/test/scala/integration/kafka/api/RecordHeaderProducerSendTest.scala
@@ -0,0 +1,79 @@
+/**
+ * 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 integration.kafka.api
+
+import java.util.Properties
+
+import kafka.api.BaseProducerSendTest
+import kafka.server.KafkaConfig
+import org.apache.kafka.clients.producer.{KafkaProducer, ProducerConfig, ProducerRecord}
+import org.apache.kafka.common.header.internals.RecordHeaders
+import org.junit.Test
+
+class RecordHeaderProducerSendTest extends BaseProducerSendTest {
+  @Test
+  def testRecordHeaders(): Unit = {
+    val producerProps = new Properties()
+    producerProps.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, brokerList)
+    producerProps.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.StringSerializer")
+    producerProps.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.StringSerializer")
+    producerProps.put(ProducerConfig.ALLOW_AUTO_CREATE_TOPICS_CONFIG, "true")
+    val producer = new KafkaProducer[String, String](producerProps)
+    try {
+      var invalidRecordHeaders = new RecordHeaders()
+      invalidRecordHeaders.add("RecordHeaderKey", "RecordHeaderValue".getBytes)
+      // The record is invalid because it contain header with keys that doesn't start with a "_"
+      // Keys that do start with a "_" are internal to the clients and are dropped at the producer.
+      val invalidRecord = new ProducerRecord[String, String](
+        topic,
+        new Integer(0),
+        "RecordKey",
+        "RecordValue",
+        invalidRecordHeaders
+      )
+      try {
+        producer.send(invalidRecord).get
+        throw new IllegalStateException("The invalid record should have thrown an exception")
+      } catch {
+        // Ignore the exception because a non internal header was introduced into the producer record
+        case ignored: IllegalArgumentException =>
+      }
+  
+      val validRecordHeaders = new RecordHeaders()
+      validRecordHeaders.add("_RecordHeaderKey", "RecordHeaderValue".getBytes)
+      val record = new ProducerRecord[String, String](
+        topic,
+        new Integer(0),
+        "RecordKey",
+        "RecordValue",
+        validRecordHeaders
+      )
+      producer.send(record).get
+    } finally {
+      producer.close()
+    }
+  }
+
+
+  override def overrideConfigs(): Properties = {
+    val properties = new Properties()
+    properties.put(KafkaConfig.InterBrokerProtocolVersionProp, "0.10.2")
+    properties.put(KafkaConfig.LogMessageFormatVersionProp, "0.10.2")
+    properties
+  }
+}
diff --git a/core/src/test/scala/integration/kafka/api/SaslClientsWithInvalidCredentialsTest.scala b/core/src/test/scala/integration/kafka/api/SaslClientsWithInvalidCredentialsTest.scala
index c353b526dddcc..83d32fa27337b 100644
--- a/core/src/test/scala/integration/kafka/api/SaslClientsWithInvalidCredentialsTest.scala
+++ b/core/src/test/scala/integration/kafka/api/SaslClientsWithInvalidCredentialsTest.scala
@@ -39,7 +39,7 @@ class SaslClientsWithInvalidCredentialsTest extends IntegrationTestHarness with
   override protected val clientSaslProperties = Some(kafkaClientSaslProperties(kafkaClientSaslMechanism))
   val consumerCount = 1
   val producerCount = 1
-  val serverCount = 1
+  val brokerCount = 1
 
   this.serverConfig.setProperty(KafkaConfig.OffsetsTopicReplicationFactorProp, "1")
   this.serverConfig.setProperty(KafkaConfig.TransactionsTopicReplicationFactorProp, "1")
@@ -50,7 +50,7 @@ class SaslClientsWithInvalidCredentialsTest extends IntegrationTestHarness with
   val numPartitions = 1
   val tp = new TopicPartition(topic, 0)
 
-  override def configureSecurityBeforeServersStart() {
+  override def configureSecurityBeforeServersStart(): Unit = {
     super.configureSecurityBeforeServersStart()
     zkClient.makeSurePersistentPathExists(ConfigEntityChangeNotificationZNode.path)
     // Create broker credentials before starting brokers
@@ -62,7 +62,7 @@ class SaslClientsWithInvalidCredentialsTest extends IntegrationTestHarness with
     startSasl(jaasSections(kafkaServerSaslMechanisms, Some(kafkaClientSaslMechanism), Both,
       JaasTestUtils.KafkaServerContextName))
     super.setUp()
-    createTopic(topic, numPartitions, serverCount)
+    createTopic(topic, numPartitions, brokerCount)
   }
 
   @After
@@ -72,7 +72,7 @@ class SaslClientsWithInvalidCredentialsTest extends IntegrationTestHarness with
   }
 
   @Test
-  def testProducerWithAuthenticationFailure() {
+  def testProducerWithAuthenticationFailure(): Unit = {
     val producer = createProducer()
     verifyAuthenticationException(sendOneRecord(producer, maxWaitMs = 10000))
     verifyAuthenticationException(producer.partitionsFor(topic))
@@ -82,7 +82,7 @@ class SaslClientsWithInvalidCredentialsTest extends IntegrationTestHarness with
   }
 
   @Test
-  def testTransactionalProducerWithAuthenticationFailure() {
+  def testTransactionalProducerWithAuthenticationFailure(): Unit = {
     val txProducer = createTransactionalProducer()
     verifyAuthenticationException(txProducer.initTransactions())
 
@@ -96,21 +96,21 @@ class SaslClientsWithInvalidCredentialsTest extends IntegrationTestHarness with
   }
 
   @Test
-  def testConsumerWithAuthenticationFailure() {
+  def testConsumerWithAuthenticationFailure(): Unit = {
     val consumer = createConsumer()
     consumer.subscribe(List(topic).asJava)
     verifyConsumerWithAuthenticationFailure(consumer)
   }
 
   @Test
-  def testManualAssignmentConsumerWithAuthenticationFailure() {
+  def testManualAssignmentConsumerWithAuthenticationFailure(): Unit = {
     val consumer = createConsumer()
     consumer.assign(List(tp).asJava)
     verifyConsumerWithAuthenticationFailure(consumer)
   }
 
   @Test
-  def testManualAssignmentConsumerWithAutoCommitDisabledWithAuthenticationFailure() {
+  def testManualAssignmentConsumerWithAutoCommitDisabledWithAuthenticationFailure(): Unit = {
     this.consumerConfig.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, false.toString)
     val consumer = createConsumer()
     consumer.assign(List(tp).asJava)
@@ -118,7 +118,7 @@ class SaslClientsWithInvalidCredentialsTest extends IntegrationTestHarness with
     verifyConsumerWithAuthenticationFailure(consumer)
   }
 
-  private def verifyConsumerWithAuthenticationFailure(consumer: KafkaConsumer[Array[Byte], Array[Byte]]) {
+  private def verifyConsumerWithAuthenticationFailure(consumer: KafkaConsumer[Array[Byte], Array[Byte]]): Unit = {
     verifyAuthenticationException(consumer.poll(Duration.ofMillis(1000)))
     verifyAuthenticationException(consumer.partitionsFor(topic))
 
@@ -129,7 +129,7 @@ class SaslClientsWithInvalidCredentialsTest extends IntegrationTestHarness with
   }
 
   @Test
-  def testKafkaAdminClientWithAuthenticationFailure() {
+  def testKafkaAdminClientWithAuthenticationFailure(): Unit = {
     val props = TestUtils.adminClientSecurityConfigs(securityProtocol, trustStoreFile, clientSaslProperties)
     props.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, brokerList)
     val adminClient = AdminClient.create(props)
@@ -157,7 +157,7 @@ class SaslClientsWithInvalidCredentialsTest extends IntegrationTestHarness with
   }
 
   @Test
-  def testConsumerGroupServiceWithAuthenticationFailure() {
+  def testConsumerGroupServiceWithAuthenticationFailure(): Unit = {
     val consumerGroupService: ConsumerGroupService = prepareConsumerGroupService
 
     val consumer = createConsumer()
@@ -168,7 +168,7 @@ class SaslClientsWithInvalidCredentialsTest extends IntegrationTestHarness with
   }
 
   @Test
-  def testConsumerGroupServiceWithAuthenticationSuccess() {
+  def testConsumerGroupServiceWithAuthenticationSuccess(): Unit = {
     createClientCredential()
     val consumerGroupService: ConsumerGroupService = prepareConsumerGroupService
 
diff --git a/core/src/test/scala/integration/kafka/api/SaslEndToEndAuthorizationTest.scala b/core/src/test/scala/integration/kafka/api/SaslEndToEndAuthorizationTest.scala
index c8521f67ed868..a1b867ef753dd 100644
--- a/core/src/test/scala/integration/kafka/api/SaslEndToEndAuthorizationTest.scala
+++ b/core/src/test/scala/integration/kafka/api/SaslEndToEndAuthorizationTest.scala
@@ -21,6 +21,7 @@ import org.apache.kafka.common.security.auth.SecurityProtocol
 import org.apache.kafka.common.errors.{GroupAuthorizationException, TopicAuthorizationException}
 import org.junit.{Before, Test}
 import org.junit.Assert.{assertEquals, assertTrue}
+import org.scalatest.Assertions.fail
 
 import scala.collection.immutable.List
 import scala.collection.JavaConverters._
@@ -34,7 +35,7 @@ abstract class SaslEndToEndAuthorizationTest extends EndToEndAuthorizationTest {
   protected def kafkaServerSaslMechanisms: List[String]
   
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     // create static config including client login context with credentials for JaasTestUtils 'client2'
     startSasl(jaasSections(kafkaServerSaslMechanisms, Option(kafkaClientSaslMechanism), Both))
     // set dynamic properties with credentials for JaasTestUtils 'client1' so that dynamic JAAS configuration is also
@@ -42,6 +43,7 @@ abstract class SaslEndToEndAuthorizationTest extends EndToEndAuthorizationTest {
     val clientLoginContext = jaasClientLoginModule(kafkaClientSaslMechanism)
     producerConfig.put(SaslConfigs.SASL_JAAS_CONFIG, clientLoginContext)
     consumerConfig.put(SaslConfigs.SASL_JAAS_CONFIG, clientLoginContext)
+    adminClientConfig.put(SaslConfigs.SASL_JAAS_CONFIG, clientLoginContext)
     super.setUp()
   }
 
diff --git a/core/src/test/scala/integration/kafka/api/SaslGssapiSslEndToEndAuthorizationTest.scala b/core/src/test/scala/integration/kafka/api/SaslGssapiSslEndToEndAuthorizationTest.scala
index a630293bb6625..1142bf7a6a029 100644
--- a/core/src/test/scala/integration/kafka/api/SaslGssapiSslEndToEndAuthorizationTest.scala
+++ b/core/src/test/scala/integration/kafka/api/SaslGssapiSslEndToEndAuthorizationTest.scala
@@ -16,6 +16,7 @@
   */
 package kafka.api
 
+import kafka.security.auth.SimpleAclAuthorizer
 import kafka.server.KafkaConfig
 import kafka.utils.JaasTestUtils
 
@@ -27,6 +28,7 @@ class SaslGssapiSslEndToEndAuthorizationTest extends SaslEndToEndAuthorizationTe
 
   override protected def kafkaClientSaslMechanism = "GSSAPI"
   override protected def kafkaServerSaslMechanisms = List("GSSAPI")
+  override protected def authorizerClass = classOf[SimpleAclAuthorizer]
 
   // Configure brokers to require SSL client authentication in order to verify that SASL_SSL works correctly even if the
   // client doesn't have a keystore. We want to cover the scenario where a broker requires either SSL client
diff --git a/core/src/test/scala/integration/kafka/api/SaslMultiMechanismConsumerTest.scala b/core/src/test/scala/integration/kafka/api/SaslMultiMechanismConsumerTest.scala
index 94b5e6f673255..b24937a039432 100644
--- a/core/src/test/scala/integration/kafka/api/SaslMultiMechanismConsumerTest.scala
+++ b/core/src/test/scala/integration/kafka/api/SaslMultiMechanismConsumerTest.scala
@@ -44,7 +44,7 @@ class SaslMultiMechanismConsumerTest extends BaseConsumerTest with SaslSetup {
   }
 
   @Test
-  def testMultipleBrokerMechanisms() {
+  def testMultipleBrokerMechanisms(): Unit = {
     val plainSaslProducer = createProducer()
     val plainSaslConsumer = createConsumer()
 
diff --git a/core/src/test/scala/integration/kafka/api/SaslPlainPlaintextConsumerTest.scala b/core/src/test/scala/integration/kafka/api/SaslPlainPlaintextConsumerTest.scala
index c15a5080bf2ec..f31575323a06c 100644
--- a/core/src/test/scala/integration/kafka/api/SaslPlainPlaintextConsumerTest.scala
+++ b/core/src/test/scala/integration/kafka/api/SaslPlainPlaintextConsumerTest.scala
@@ -13,7 +13,9 @@
 package kafka.api
 
 import java.io.File
-import java.util.Locale
+import java.util.{Locale, Properties}
+
+import scala.collection.Seq
 
 import kafka.server.KafkaConfig
 import kafka.utils.{JaasTestUtils, TestUtils}
@@ -47,12 +49,17 @@ class SaslPlainPlaintextConsumerTest extends BaseConsumerTest with SaslSetup {
     closeSasl()
   }
 
+  override def modifyConfigs(props: Seq[Properties]): Unit = {
+    super.modifyConfigs(props)
+    configureListeners(props)
+  }
+
   /**
-   * Checks that everyone can access ZkUtils.SecureZkRootPaths and ZkUtils.SensitiveZkRootPaths
+   * Checks that everyone can access ZkData.SecureZkRootPaths and ZkData.SensitiveZkRootPaths
    * when zookeeper.set.acl=false, even if ZooKeeper is SASL-enabled.
    */
   @Test
-  def testZkAclsDisabled() {
+  def testZkAclsDisabled(): Unit = {
     TestUtils.verifyUnsecureZkAcls(zkClient)
   }
 }
diff --git a/core/src/test/scala/integration/kafka/api/SaslPlainSslEndToEndAuthorizationTest.scala b/core/src/test/scala/integration/kafka/api/SaslPlainSslEndToEndAuthorizationTest.scala
index bbe0dd8356d84..0c1559106c955 100644
--- a/core/src/test/scala/integration/kafka/api/SaslPlainSslEndToEndAuthorizationTest.scala
+++ b/core/src/test/scala/integration/kafka/api/SaslPlainSslEndToEndAuthorizationTest.scala
@@ -21,6 +21,8 @@ import javax.security.auth.callback._
 import javax.security.auth.Subject
 import javax.security.auth.login.AppConfigurationEntry
 
+import scala.collection.Seq
+
 import kafka.server.KafkaConfig
 import kafka.utils.{TestUtils}
 import kafka.utils.JaasTestUtils._
@@ -57,8 +59,8 @@ object SaslPlainSslEndToEndAuthorizationTest {
   }
 
   class TestServerCallbackHandler extends AuthenticateCallbackHandler {
-    def configure(configs: java.util.Map[String, _], saslMechanism: String, jaasConfigEntries: java.util.List[AppConfigurationEntry]) {}
-    def handle(callbacks: Array[Callback]) {
+    def configure(configs: java.util.Map[String, _], saslMechanism: String, jaasConfigEntries: java.util.List[AppConfigurationEntry]): Unit = {}
+    def handle(callbacks: Array[Callback]): Unit = {
       var username: String = null
       for (callback <- callbacks) {
         if (callback.isInstanceOf[NameCallback])
@@ -70,12 +72,12 @@ object SaslPlainSslEndToEndAuthorizationTest {
           throw new UnsupportedCallbackException(callback)
       }
     }
-    def close() {}
+    def close(): Unit = {}
   }
 
   class TestClientCallbackHandler extends AuthenticateCallbackHandler {
-    def configure(configs: java.util.Map[String, _], saslMechanism: String, jaasConfigEntries: java.util.List[AppConfigurationEntry]) {}
-    def handle(callbacks: Array[Callback]) {
+    def configure(configs: java.util.Map[String, _], saslMechanism: String, jaasConfigEntries: java.util.List[AppConfigurationEntry]): Unit = {}
+    def handle(callbacks: Array[Callback]): Unit = {
       val subject = Subject.getSubject(AccessController.getContext())
       val username = subject.getPublicCredentials(classOf[String]).iterator().next()
       for (callback <- callbacks) {
@@ -88,7 +90,7 @@ object SaslPlainSslEndToEndAuthorizationTest {
           throw new UnsupportedCallbackException(callback)
       }
     }
-    def close() {}
+    def close(): Unit = {}
   }
 }
 
@@ -107,9 +109,11 @@ class SaslPlainSslEndToEndAuthorizationTest extends SaslEndToEndAuthorizationTes
   this.serverConfig.put(s"$mechanismPrefix${KafkaConfig.SaslServerCallbackHandlerClassProp}", classOf[TestServerCallbackHandler].getName)
   this.producerConfig.put(SaslConfigs.SASL_CLIENT_CALLBACK_HANDLER_CLASS, classOf[TestClientCallbackHandler].getName)
   this.consumerConfig.put(SaslConfigs.SASL_CLIENT_CALLBACK_HANDLER_CLASS, classOf[TestClientCallbackHandler].getName)
+  this.adminClientConfig.put(SaslConfigs.SASL_CLIENT_CALLBACK_HANDLER_CLASS, classOf[TestClientCallbackHandler].getName)
   private val plainLogin = s"org.apache.kafka.common.security.plain.PlainLoginModule username=$KafkaPlainUser required;"
   this.producerConfig.put(SaslConfigs.SASL_JAAS_CONFIG, plainLogin)
   this.consumerConfig.put(SaslConfigs.SASL_JAAS_CONFIG, plainLogin)
+  this.adminClientConfig.put(SaslConfigs.SASL_JAAS_CONFIG, plainLogin)
 
   override protected def kafkaClientSaslMechanism = "PLAIN"
   override protected def kafkaServerSaslMechanisms = List("PLAIN")
@@ -132,7 +136,7 @@ class SaslPlainSslEndToEndAuthorizationTest extends SaslEndToEndAuthorizationTes
    * have expected ACLs.
    */
   @Test
-  def testAcls() {
+  def testAcls(): Unit = {
     TestUtils.verifySecureZkAcls(zkClient, 1)
   }
 }
diff --git a/core/src/test/scala/integration/kafka/api/SaslScramSslEndToEndAuthorizationTest.scala b/core/src/test/scala/integration/kafka/api/SaslScramSslEndToEndAuthorizationTest.scala
index 8652a6fedbcd1..bc425911baa83 100644
--- a/core/src/test/scala/integration/kafka/api/SaslScramSslEndToEndAuthorizationTest.scala
+++ b/core/src/test/scala/integration/kafka/api/SaslScramSslEndToEndAuthorizationTest.scala
@@ -30,7 +30,7 @@ class SaslScramSslEndToEndAuthorizationTest extends SaslEndToEndAuthorizationTes
   override val kafkaPrincipal = JaasTestUtils.KafkaScramAdmin
   private val kafkaPassword = JaasTestUtils.KafkaScramAdminPassword
 
-  override def configureSecurityBeforeServersStart() {
+  override def configureSecurityBeforeServersStart(): Unit = {
     super.configureSecurityBeforeServersStart()
     zkClient.makeSurePersistentPathExists(ConfigEntityChangeNotificationZNode.path)
     // Create broker credentials before starting brokers
@@ -38,7 +38,7 @@ class SaslScramSslEndToEndAuthorizationTest extends SaslEndToEndAuthorizationTes
   }
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     super.setUp()
     // Create client credentials after starting brokers so that dynamic credential creation is also tested
     createScramCredentials(zkConnect, JaasTestUtils.KafkaScramUser, JaasTestUtils.KafkaScramPassword)
diff --git a/core/src/test/scala/integration/kafka/api/SaslSetup.scala b/core/src/test/scala/integration/kafka/api/SaslSetup.scala
index 81de1059068a8..7b06eb40bdf9d 100644
--- a/core/src/test/scala/integration/kafka/api/SaslSetup.scala
+++ b/core/src/test/scala/integration/kafka/api/SaslSetup.scala
@@ -21,6 +21,8 @@ import java.io.File
 import java.util.Properties
 import javax.security.auth.login.Configuration
 
+import scala.collection.Seq
+
 import kafka.admin.ConfigCommand
 import kafka.security.minikdc.MiniKdc
 import kafka.server.KafkaConfig
@@ -51,7 +53,7 @@ trait SaslSetup {
   private var serverKeytabFile: Option[File] = None
   private var clientKeytabFile: Option[File] = None
 
-  def startSasl(jaasSections: Seq[JaasSection]) {
+  def startSasl(jaasSections: Seq[JaasSection]): Unit = {
     // Important if tests leak consumers, producers or brokers
     LoginManager.closeAll()
     val hasKerberos = jaasSections.exists(_.modules.exists {
@@ -106,14 +108,14 @@ trait SaslSetup {
     }
   }
 
-  private def writeJaasConfigurationToFile(jaasSections: Seq[JaasSection]) {
+  private def writeJaasConfigurationToFile(jaasSections: Seq[JaasSection]): Unit = {
     val file = JaasTestUtils.writeJaasContextsToFile(jaasSections)
     System.setProperty(JaasUtils.JAVA_LOGIN_CONFIG_PARAM, file.getAbsolutePath)
     // This will cause a reload of the Configuration singleton when `getConfiguration` is called
     Configuration.setConfiguration(null)
   }
 
-  def closeSasl() {
+  def closeSasl(): Unit = {
     if (kdc != null)
       kdc.stop()
     // Important if tests leak consumers, producers or brokers
diff --git a/core/src/test/scala/integration/kafka/api/SaslSslAdminClientIntegrationTest.scala b/core/src/test/scala/integration/kafka/api/SaslSslAdminClientIntegrationTest.scala
index cb2186c42271b..fd4b29f3ff484 100644
--- a/core/src/test/scala/integration/kafka/api/SaslSslAdminClientIntegrationTest.scala
+++ b/core/src/test/scala/integration/kafka/api/SaslSslAdminClientIntegrationTest.scala
@@ -15,30 +15,36 @@ package kafka.api
 import java.io.File
 import java.util
 
+import kafka.log.LogConfig
 import kafka.security.auth.{All, Allow, Alter, AlterConfigs, Authorizer, ClusterAction, Create, Delete, Deny, Describe, Group, Operation, PermissionType, SimpleAclAuthorizer, Topic, Acl => AuthAcl, Resource => AuthResource}
-import kafka.server.KafkaConfig
+import kafka.security.authorizer.AuthorizerWrapper
+import kafka.server.{Defaults, KafkaConfig}
 import kafka.utils.{CoreUtils, JaasTestUtils, TestUtils}
 import kafka.utils.TestUtils._
-
-import org.apache.kafka.clients.admin.{AdminClient, CreateAclsOptions, DeleteAclsOptions}
+import org.apache.kafka.clients.admin._
 import org.apache.kafka.common.acl._
-import org.apache.kafka.common.errors.{ClusterAuthorizationException, InvalidRequestException}
+import org.apache.kafka.common.config.ConfigResource
+import org.apache.kafka.common.errors.{ClusterAuthorizationException, InvalidRequestException, TopicAuthorizationException, UnknownTopicOrPartitionException}
 import org.apache.kafka.common.resource.{PatternType, ResourcePattern, ResourcePatternFilter, ResourceType}
 import org.apache.kafka.common.security.auth.{KafkaPrincipal, SecurityProtocol}
-import org.junit.Assert.assertEquals
+import org.junit.Assert.{assertEquals, assertTrue}
 import org.junit.{After, Assert, Before, Test}
 
 import scala.collection.JavaConverters._
+import scala.collection.Seq
+import scala.compat.java8.OptionConverters._
+import scala.concurrent.ExecutionException
 import scala.util.{Failure, Success, Try}
 
 class SaslSslAdminClientIntegrationTest extends AdminClientIntegrationTest with SaslSetup {
   this.serverConfig.setProperty(KafkaConfig.ZkEnableSecureAclsProp, "true")
+  // This tests the old SimpleAclAuthorizer, we have another test for the new AclAuthorizer
   this.serverConfig.setProperty(KafkaConfig.AuthorizerClassNameProp, classOf[SimpleAclAuthorizer].getName)
 
   override protected def securityProtocol = SecurityProtocol.SASL_SSL
   override protected lazy val trustStoreFile = Some(File.createTempFile("truststore", ".jks"))
 
-  override def configureSecurityBeforeServersStart() {
+  override def configureSecurityBeforeServersStart(): Unit = {
     val authorizer = CoreUtils.createObject[Authorizer](classOf[SimpleAclAuthorizer].getName)
     try {
       authorizer.configure(this.configs.head.originals())
@@ -60,26 +66,30 @@ class SaslSslAdminClientIntegrationTest extends AdminClientIntegrationTest with
 
   @Before
   override def setUp(): Unit = {
-    startSasl(jaasSections(Seq("GSSAPI"), Some("GSSAPI"), Both, JaasTestUtils.KafkaServerContextName))
+    setUpSasl()
     super.setUp()
   }
 
+  def setUpSasl(): Unit = {
+    startSasl(jaasSections(Seq("GSSAPI"), Some("GSSAPI"), Both, JaasTestUtils.KafkaServerContextName))
+  }
+
   private def clusterAcl(permissionType: PermissionType, operation: Operation): AuthAcl = {
     new AuthAcl(new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "*"), permissionType,
       AuthAcl.WildCardHost, operation)
   }
 
-  private def addClusterAcl(permissionType: PermissionType, operation: Operation): Unit = {
+  def addClusterAcl(permissionType: PermissionType, operation: Operation): Unit = {
     val acls = Set(clusterAcl(permissionType, operation))
-    val authorizer = servers.head.dataPlaneRequestProcessor.authorizer.get
+    val authorizer = simpleAclAuthorizer
     val prevAcls = authorizer.getAcls(AuthResource.ClusterResource)
     authorizer.addAcls(acls, AuthResource.ClusterResource)
     TestUtils.waitAndVerifyAcls(prevAcls ++ acls, authorizer, AuthResource.ClusterResource)
   }
 
-  private def removeClusterAcl(permissionType: PermissionType, operation: Operation): Unit = {
+  def removeClusterAcl(permissionType: PermissionType, operation: Operation): Unit = {
     val acls = Set(clusterAcl(permissionType, operation))
-    val authorizer = servers.head.dataPlaneRequestProcessor.authorizer.get
+    val authorizer = simpleAclAuthorizer
     val prevAcls = authorizer.getAcls(AuthResource.ClusterResource)
     Assert.assertTrue(authorizer.removeAcls(acls, AuthResource.ClusterResource))
     TestUtils.waitAndVerifyAcls(prevAcls -- acls, authorizer, AuthResource.ClusterResource)
@@ -278,6 +288,11 @@ class SaslSslAdminClientIntegrationTest extends AdminClientIntegrationTest with
     assertFutureExceptionTypeEquals(results.values.get(emptyResourceNameAcl), classOf[InvalidRequestException])
   }
 
+  override def configuredClusterPermissions(): Set[AclOperation] = {
+    Set(AclOperation.ALTER, AclOperation.CREATE, AclOperation.CLUSTER_ACTION, AclOperation.ALTER_CONFIGS,
+      AclOperation.DESCRIBE, AclOperation.DESCRIBE_CONFIGS)
+  }
+
   private def verifyCauseIsClusterAuth(e: Throwable): Unit = {
     if (!e.getCause.isInstanceOf[ClusterAuthorizationException]) {
       throw e.getCause
@@ -391,7 +406,79 @@ class SaslSslAdminClientIntegrationTest extends AdminClientIntegrationTest with
     testAclCreateGetDelete(expectAuth = false)
   }
 
-  private def waitForDescribeAcls(client: AdminClient, filter: AclBindingFilter, acls: Set[AclBinding]): Unit = {
+  @Test
+  def testCreateTopicsResponseMetadataAndConfig(): Unit = {
+    val topic1 = "mytopic1"
+    val topic2 = "mytopic2"
+    val denyAcl = new AclBinding(new ResourcePattern(ResourceType.TOPIC, topic2, PatternType.LITERAL),
+      new AccessControlEntry("User:*", "*", AclOperation.DESCRIBE_CONFIGS, AclPermissionType.DENY))
+
+    client = AdminClient.create(createConfig())
+    client.createAcls(List(denyAcl).asJava, new CreateAclsOptions()).all().get()
+
+    val topics = Seq(topic1, topic2)
+    val configsOverride = Map(LogConfig.SegmentBytesProp -> "100000").asJava
+    val newTopics = Seq(
+      new NewTopic(topic1, 2, 3.toShort).configs(configsOverride),
+      new NewTopic(topic2, Option.empty[Integer].asJava, Option.empty[java.lang.Short].asJava).configs(configsOverride))
+    val validateResult = client.createTopics(newTopics.asJava, new CreateTopicsOptions().validateOnly(true))
+    validateResult.all.get()
+    waitForTopics(client, List(), topics)
+
+    def validateMetadataAndConfigs(result: CreateTopicsResult): Unit = {
+      assertEquals(2, result.numPartitions(topic1).get())
+      assertEquals(3, result.replicationFactor(topic1).get())
+      val topicConfigs = result.config(topic1).get().entries.asScala
+      assertTrue(topicConfigs.nonEmpty)
+      val segmentBytesConfig = topicConfigs.find(_.name == LogConfig.SegmentBytesProp).get
+      assertEquals(100000, segmentBytesConfig.value.toLong)
+      assertEquals(ConfigEntry.ConfigSource.DYNAMIC_TOPIC_CONFIG, segmentBytesConfig.source)
+      val compressionConfig = topicConfigs.find(_.name == LogConfig.CompressionTypeProp).get
+      assertEquals(Defaults.CompressionType, compressionConfig.value)
+      assertEquals(ConfigEntry.ConfigSource.DEFAULT_CONFIG, compressionConfig.source)
+
+      assertFutureExceptionTypeEquals(result.numPartitions(topic2), classOf[TopicAuthorizationException])
+      assertFutureExceptionTypeEquals(result.replicationFactor(topic2), classOf[TopicAuthorizationException])
+      assertFutureExceptionTypeEquals(result.config(topic2), classOf[TopicAuthorizationException])
+    }
+    validateMetadataAndConfigs(validateResult)
+
+    val createResult = client.createTopics(newTopics.asJava, new CreateTopicsOptions())
+    createResult.all.get()
+    waitForTopics(client, topics, List())
+    validateMetadataAndConfigs(createResult)
+    val createResponseConfig = createResult.config(topic1).get().entries.asScala
+
+    val describeResponseConfig = describeConfigs(topic1)
+    assertEquals(describeResponseConfig.map(_.name).toSet, createResponseConfig.map(_.name).toSet)
+    describeResponseConfig.foreach { describeEntry =>
+      val name = describeEntry.name
+      val createEntry = createResponseConfig.find(_.name == name).get
+      assertEquals(s"Value mismatch for $name", describeEntry.value, createEntry.value)
+      assertEquals(s"isReadOnly mismatch for $name", describeEntry.isReadOnly, createEntry.isReadOnly)
+      assertEquals(s"isSensitive mismatch for $name", describeEntry.isSensitive, createEntry.isSensitive)
+      assertEquals(s"Source mismatch for $name", describeEntry.source, createEntry.source)
+    }
+  }
+
+  private def describeConfigs(topic: String): Iterable[ConfigEntry] = {
+    val topicResource = new ConfigResource(ConfigResource.Type.TOPIC, topic)
+    var configEntries: Iterable[ConfigEntry] = null
+
+    TestUtils.waitUntilTrue(() => {
+      try {
+        val topicResponse = client.describeConfigs(List(topicResource).asJava).all.get.get(topicResource)
+        configEntries = topicResponse.entries.asScala
+        true
+      } catch {
+        case e: ExecutionException if e.getCause.isInstanceOf[UnknownTopicOrPartitionException] => false
+      }
+    }, "Timed out waiting for describeConfigs")
+
+    configEntries
+  }
+
+  private def waitForDescribeAcls(client: Admin, filter: AclBindingFilter, acls: Set[AclBinding]): Unit = {
     var lastResults: util.Collection[AclBinding] = null
     TestUtils.waitUntilTrue(() => {
       lastResults = client.describeAcls(filter).values.get()
@@ -408,4 +495,9 @@ class SaslSslAdminClientIntegrationTest extends AdminClientIntegrationTest with
   private def getAcls(allTopicAcls: AclBindingFilter) = {
     client.describeAcls(allTopicAcls).values.get().asScala.toSet
   }
+
+  private def simpleAclAuthorizer: Authorizer = {
+    val authorizerWrapper = servers.head.dataPlaneRequestProcessor.authorizer.get.asInstanceOf[AuthorizerWrapper]
+    authorizerWrapper.baseAuthorizer
+  }
 }
diff --git a/core/src/test/scala/integration/kafka/api/SslAdminClientIntegrationTest.scala b/core/src/test/scala/integration/kafka/api/SslAdminClientIntegrationTest.scala
new file mode 100644
index 0000000000000..f5fb42133e52a
--- /dev/null
+++ b/core/src/test/scala/integration/kafka/api/SslAdminClientIntegrationTest.scala
@@ -0,0 +1,313 @@
+/**
+  * 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 java.io.File
+import java.util
+import java.util.Collections
+import java.util.concurrent._
+import java.util.function.BiConsumer
+
+import com.yammer.metrics.Metrics
+import com.yammer.metrics.core.Gauge
+import kafka.security.authorizer.AclAuthorizer
+import kafka.security.authorizer.AuthorizerUtils.{WildcardHost, WildcardPrincipal}
+import kafka.security.auth.{Operation, PermissionType}
+import kafka.server.KafkaConfig
+import kafka.utils.{CoreUtils, TestUtils}
+import org.apache.kafka.clients.admin.{AdminClient, AdminClientConfig, CreateAclsResult}
+import org.apache.kafka.common.acl._
+import org.apache.kafka.common.acl.AclOperation._
+import org.apache.kafka.common.acl.AclPermissionType._
+import org.apache.kafka.common.protocol.ApiKeys
+import org.apache.kafka.common.resource.ResourcePattern
+import org.apache.kafka.common.resource.PatternType._
+import org.apache.kafka.common.resource.ResourceType._
+import org.apache.kafka.common.security.auth.{KafkaPrincipal, SecurityProtocol}
+import org.apache.kafka.server.authorizer._
+import org.junit.Assert.{assertEquals, assertFalse, assertNotNull, assertTrue}
+import org.junit.{Assert, Test}
+
+import scala.collection.JavaConverters._
+import scala.collection.mutable
+
+object SslAdminClientIntegrationTest {
+  @volatile var semaphore: Option[Semaphore] = None
+  @volatile var executor: Option[ExecutorService] = None
+  @volatile var lastUpdateRequestContext: Option[AuthorizableRequestContext] = None
+  class TestableAclAuthorizer extends AclAuthorizer {
+    override def createAcls(requestContext: AuthorizableRequestContext,
+                            aclBindings: util.List[AclBinding]): util.List[_ <: CompletionStage[AclCreateResult]] = {
+      lastUpdateRequestContext = Some(requestContext)
+      execute[AclCreateResult](aclBindings.size, () => super.createAcls(requestContext, aclBindings))
+    }
+
+    override def deleteAcls(requestContext: AuthorizableRequestContext,
+                            aclBindingFilters: util.List[AclBindingFilter]): util.List[_ <: CompletionStage[AclDeleteResult]] = {
+      lastUpdateRequestContext = Some(requestContext)
+      execute[AclDeleteResult](aclBindingFilters.size, () => super.deleteAcls(requestContext, aclBindingFilters))
+    }
+
+    private def execute[T](batchSize: Int, action: () => util.List[_ <: CompletionStage[T]]): util.List[CompletableFuture[T]] = {
+      val futures = (0 until batchSize).map(_ => new CompletableFuture[T]).toList
+      val runnable = new Runnable {
+        override def run(): Unit = {
+          semaphore.foreach(_.acquire())
+          try {
+            action.apply().asScala.zip(futures).foreach { case (baseFuture, resultFuture) =>
+              baseFuture.whenComplete(new BiConsumer[T, Throwable]() {
+                override def accept(result: T, exception: Throwable): Unit = {
+                  if (exception != null)
+                    resultFuture.completeExceptionally(exception)
+                  else
+                    resultFuture.complete(result)
+                }
+              })
+            }
+          } finally {
+            semaphore.foreach(_.release())
+          }
+        }
+      }
+      executor match {
+        case Some(executorService) => executorService.submit(runnable)
+        case None => runnable.run()
+      }
+      futures.asJava
+    }
+  }
+}
+
+class SslAdminClientIntegrationTest extends SaslSslAdminClientIntegrationTest {
+  this.serverConfig.setProperty(KafkaConfig.ZkEnableSecureAclsProp, "true")
+  this.serverConfig.setProperty(KafkaConfig.AuthorizerClassNameProp, classOf[SslAdminClientIntegrationTest.TestableAclAuthorizer].getName)
+
+  override protected def securityProtocol = SecurityProtocol.SSL
+  override protected lazy val trustStoreFile = Some(File.createTempFile("truststore", ".jks"))
+  private val adminClients = mutable.Buffer.empty[AdminClient]
+
+  override def configureSecurityBeforeServersStart(): Unit = {
+    val authorizer = CoreUtils.createObject[Authorizer](classOf[AclAuthorizer].getName)
+    try {
+      authorizer.configure(this.configs.head.originals())
+      val ace = new AccessControlEntry(WildcardPrincipal, WildcardHost, 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)
+
+      authorizer.createAcls(null, List(clusterAcl(ALLOW, CREATE),
+                             clusterAcl(ALLOW, DELETE),
+                             clusterAcl(ALLOW, CLUSTER_ACTION),
+                             clusterAcl(ALLOW, ALTER_CONFIGS),
+                             clusterAcl(ALLOW, ALTER))
+        .map(ace => new AclBinding(clusterResourcePattern, ace)).asJava)
+    } finally {
+      authorizer.close()
+    }
+  }
+
+  override def setUpSasl(): Unit = {
+    SslAdminClientIntegrationTest.semaphore = None
+    SslAdminClientIntegrationTest.executor = None
+    SslAdminClientIntegrationTest.lastUpdateRequestContext = None
+
+    startSasl(jaasSections(List.empty, None, ZkSasl))
+  }
+
+  override def tearDown(): Unit = {
+    // Ensure semaphore doesn't block shutdown even if test has failed
+    val semaphore = SslAdminClientIntegrationTest.semaphore
+    SslAdminClientIntegrationTest.semaphore = None
+    semaphore.foreach(s => s.release(s.getQueueLength))
+
+    adminClients.foreach(_.close())
+    super.tearDown()
+  }
+
+  override def addClusterAcl(permissionType: PermissionType, operation: Operation): Unit = {
+    val ace = clusterAcl(permissionType.toJava, operation.toJava)
+    val aclBinding = new AclBinding(clusterResourcePattern, ace)
+    val authorizer = servers.head.dataPlaneRequestProcessor.authorizer.get
+    val prevAcls = authorizer.acls(new AclBindingFilter(clusterResourcePattern.toFilter, AccessControlEntryFilter.ANY))
+      .asScala.map(_.entry).toSet
+    authorizer.createAcls(null, Collections.singletonList(aclBinding))
+    TestUtils.waitAndVerifyAcls(prevAcls ++ Set(ace), authorizer, clusterResourcePattern)
+  }
+
+  override def removeClusterAcl(permissionType: PermissionType, operation: Operation): Unit = {
+    val ace = clusterAcl(permissionType.toJava, operation.toJava)
+    val authorizer = servers.head.dataPlaneRequestProcessor.authorizer.get
+    val clusterFilter = new AclBindingFilter(clusterResourcePattern.toFilter, AccessControlEntryFilter.ANY)
+    val prevAcls = authorizer.acls(clusterFilter).asScala.map(_.entry).toSet
+    val deleteFilter = new AclBindingFilter(clusterResourcePattern.toFilter, ace.toFilter)
+    Assert.assertFalse(authorizer.deleteAcls(null, Collections.singletonList(deleteFilter))
+      .get(0).toCompletableFuture.get.aclBindingDeleteResults().asScala.head.exception.isPresent)
+    TestUtils.waitAndVerifyAcls(prevAcls -- Set(ace), authorizer, clusterResourcePattern)
+  }
+
+  private def clusterAcl(permissionType: AclPermissionType, operation: AclOperation): AccessControlEntry = {
+    new AccessControlEntry(new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "*").toString,
+      WildcardHost, operation, permissionType)
+  }
+
+  @Test
+  def testAclUpdatesUsingSynchronousAuthorizer(): Unit = {
+    verifyAclUpdates()
+  }
+
+  @Test
+  def testAclUpdatesUsingAsynchronousAuthorizer(): Unit = {
+    SslAdminClientIntegrationTest.executor = Some(Executors.newSingleThreadExecutor)
+    verifyAclUpdates()
+  }
+
+  /**
+   * Verify that ACL updates using synchronous authorizer are performed synchronously
+   * on request threads without any performance overhead introduced by a purgatory.
+   */
+  @Test
+  def testSynchronousAuthorizerAclUpdatesBlockRequestThreads(): Unit = {
+    val testSemaphore = new Semaphore(0)
+    SslAdminClientIntegrationTest.semaphore = Some(testSemaphore)
+    waitForNoBlockedRequestThreads()
+
+    // Queue requests until all threads are blocked. ACL create requests are sent to least loaded
+    // node, so we may need more than `numRequestThreads` requests to block all threads.
+    val aclFutures = mutable.Buffer[CreateAclsResult]()
+    while (blockedRequestThreads.size < numRequestThreads) {
+      aclFutures += createAdminClient.createAcls(List(acl2).asJava)
+      assertTrue(s"Request threads not blocked numRequestThreads=$numRequestThreads blocked=$blockedRequestThreads",
+        aclFutures.size < numRequestThreads * 10)
+    }
+    assertEquals(0, purgatoryMetric("NumDelayedOperations"))
+    assertEquals(0, purgatoryMetric("PurgatorySize"))
+
+    // Verify that operations on other clients are blocked
+    val describeFuture = createAdminClient.describeCluster().clusterId()
+    assertFalse(describeFuture.isDone)
+
+    // Release the semaphore and verify that all requests complete
+    testSemaphore.release(aclFutures.size)
+    waitForNoBlockedRequestThreads()
+    assertNotNull(describeFuture.get(10, TimeUnit.SECONDS))
+    // If any of the requests time out since we were blocking the threads earlier, retry the request.
+    val numTimedOut = aclFutures.count { future =>
+      try {
+        future.all().get()
+        false
+      } catch {
+        case e: ExecutionException =>
+          if (e.getCause.isInstanceOf[org.apache.kafka.common.errors.TimeoutException])
+            true
+          else
+            throw e.getCause
+      }
+    }
+    (0 until numTimedOut)
+      .map(_ => createAdminClient.createAcls(List(acl2).asJava))
+      .foreach(_.all().get(30, TimeUnit.SECONDS))
+  }
+
+  /**
+   * Verify that ACL updates using an asynchronous authorizer are completed asynchronously
+   * using a purgatory, enabling other requests to be processed even when ACL updates are blocked.
+   */
+  @Test
+  def testAsynchronousAuthorizerAclUpdatesDontBlockRequestThreads(): Unit = {
+    SslAdminClientIntegrationTest.executor = Some(Executors.newSingleThreadExecutor)
+    val testSemaphore = new Semaphore(0)
+    SslAdminClientIntegrationTest.semaphore = Some(testSemaphore)
+
+    waitForNoBlockedRequestThreads()
+
+    val aclFutures = (0 until numRequestThreads).map(_ => createAdminClient.createAcls(List(acl2).asJava))
+    waitForNoBlockedRequestThreads()
+    assertTrue(aclFutures.forall(future => !future.all.isDone))
+    // Other requests should succeed even though ACL updates are blocked
+    assertNotNull(createAdminClient.describeCluster().clusterId().get(10, TimeUnit.SECONDS))
+    TestUtils.waitUntilTrue(() => purgatoryMetric("PurgatorySize") > 0, "PurgatorySize metrics not updated")
+    TestUtils.waitUntilTrue(() => purgatoryMetric("NumDelayedOperations") > 0, "NumDelayedOperations metrics not updated")
+
+    // Release the semaphore and verify that ACL update requests complete
+    testSemaphore.release(aclFutures.size)
+    aclFutures.foreach(_.all.get())
+    assertEquals(0, purgatoryMetric("NumDelayedOperations"))
+  }
+
+  private def verifyAclUpdates(): Unit = {
+    def validateRequestContext(context: AuthorizableRequestContext, apiKey: ApiKeys): Unit = {
+      assertEquals(SecurityProtocol.SSL, context.securityProtocol)
+      assertEquals("SSL", context.listenerName)
+      assertEquals(KafkaPrincipal.ANONYMOUS, context.principal)
+      assertEquals(apiKey.id.toInt, context.requestType)
+      assertEquals(apiKey.latestVersion.toInt, context.requestVersion)
+      assertTrue(s"Invalid correlation id: ${context.correlationId}", context.correlationId > 0)
+      assertTrue(s"Invalid client id: ${context.clientId}", context.clientId.startsWith("adminclient"))
+      assertTrue(s"Invalid host address: ${context.clientAddress}", context.clientAddress.isLoopbackAddress)
+    }
+
+    val testSemaphore = new Semaphore(0)
+    SslAdminClientIntegrationTest.semaphore = Some(testSemaphore)
+
+    client = AdminClient.create(createConfig())
+    val results = client.createAcls(List(acl2, acl3).asJava).values
+    assertEquals(Set(acl2, acl3), results.keySet().asScala)
+    assertFalse(results.values().asScala.exists(_.isDone))
+    TestUtils.waitUntilTrue(() => testSemaphore.hasQueuedThreads, "Authorizer not blocked in createAcls")
+    testSemaphore.release()
+    results.values().asScala.foreach(_.get)
+    validateRequestContext(SslAdminClientIntegrationTest.lastUpdateRequestContext.get, ApiKeys.CREATE_ACLS)
+
+    testSemaphore.acquire()
+    val results2 = client.deleteAcls(List(ACL1.toFilter, acl2.toFilter, acl3.toFilter).asJava).values
+    assertEquals(Set(ACL1.toFilter, acl2.toFilter, acl3.toFilter), results2.keySet.asScala)
+    assertFalse(results2.values().asScala.exists(_.isDone))
+    TestUtils.waitUntilTrue(() => testSemaphore.hasQueuedThreads, "Authorizer not blocked in deleteAcls")
+    testSemaphore.release()
+    results.values().asScala.foreach(_.get)
+    assertEquals(0, results2.get(ACL1.toFilter).get.values.size())
+    assertEquals(Set(acl2), results2.get(acl2.toFilter).get.values.asScala.map(_.binding).toSet)
+    assertEquals(Set(acl3), results2.get(acl3.toFilter).get.values.asScala.map(_.binding).toSet)
+    validateRequestContext(SslAdminClientIntegrationTest.lastUpdateRequestContext.get, ApiKeys.DELETE_ACLS)
+  }
+
+  private def createAdminClient: AdminClient = {
+    val config = createConfig()
+    config.put(AdminClientConfig.REQUEST_TIMEOUT_MS_CONFIG, "40000")
+    val client = AdminClient.create(config)
+    adminClients += client
+    client
+  }
+
+  private def blockedRequestThreads: List[Thread] = {
+    val requestThreads = Thread.getAllStackTraces.keySet.asScala
+      .filter(_.getName.contains("data-plane-kafka-request-handler"))
+    assertEquals(numRequestThreads, requestThreads.size)
+    requestThreads.filter(_.getState == Thread.State.WAITING).toList
+  }
+
+  private def numRequestThreads = servers.head.config.numIoThreads * servers.size
+
+  private def waitForNoBlockedRequestThreads(): Unit = {
+    val (blockedThreads, _) = TestUtils.computeUntilTrue(blockedRequestThreads)(_.isEmpty)
+    assertEquals(List.empty, blockedThreads)
+  }
+
+  private def purgatoryMetric(name: String): Int = {
+    val allMetrics = Metrics.defaultRegistry.allMetrics.asScala
+    val metrics = allMetrics.filter { case (metricName, _) =>
+      metricName.getMBeanName.contains("delayedOperation=AlterAcls") && metricName.getMBeanName.contains(s"name=$name")
+    }.values.toList
+    assertTrue(s"Unable to find metric $name: allMetrics: ${allMetrics.keySet.map(_.getMBeanName)}", metrics.nonEmpty)
+    metrics.map(_.asInstanceOf[Gauge[Int]].value).sum
+  }
+}
diff --git a/core/src/test/scala/integration/kafka/api/SslEndToEndAuthorizationTest.scala b/core/src/test/scala/integration/kafka/api/SslEndToEndAuthorizationTest.scala
index 8354ee06a4d5a..fe087ec864037 100644
--- a/core/src/test/scala/integration/kafka/api/SslEndToEndAuthorizationTest.scala
+++ b/core/src/test/scala/integration/kafka/api/SslEndToEndAuthorizationTest.scala
@@ -67,7 +67,7 @@ class SslEndToEndAuthorizationTest extends EndToEndAuthorizationTest {
   override val clientPrincipal = s"O=A client,CN=$clientCn"
   override val kafkaPrincipal = "server"
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     startSasl(jaasSections(List.empty, None, ZkSasl))
     super.setUp()
   }
diff --git a/core/src/test/scala/integration/kafka/api/TransactionsBounceTest.scala b/core/src/test/scala/integration/kafka/api/TransactionsBounceTest.scala
index eabafccffb2b8..b24698bed6713 100644
--- a/core/src/test/scala/integration/kafka/api/TransactionsBounceTest.scala
+++ b/core/src/test/scala/integration/kafka/api/TransactionsBounceTest.scala
@@ -71,7 +71,7 @@ class TransactionsBounceTest extends KafkaServerTestHarness {
   }
 
   @Test
-  def testBrokerFailure() {
+  def testBrokerFailure(): Unit = {
     // basic idea is to seed a topic with 10000 records, and copy it transactionally while bouncing brokers
     // constantly through the period.
     val consumerGroup = "myGroup"
@@ -181,7 +181,7 @@ class TransactionsBounceTest extends KafkaServerTestHarness {
       (0 until numPartitions).foreach(partition => TestUtils.waitUntilLeaderIsElectedOrChanged(zkClient, outputTopic, partition))
     }
 
-    override def shutdown(){
+    override def shutdown(): Unit = {
       super.shutdown()
    }
   }
diff --git a/core/src/test/scala/integration/kafka/api/TransactionsTest.scala b/core/src/test/scala/integration/kafka/api/TransactionsTest.scala
index 34dea70ca8c6e..63bfc9f4639c4 100644
--- a/core/src/test/scala/integration/kafka/api/TransactionsTest.scala
+++ b/core/src/test/scala/integration/kafka/api/TransactionsTest.scala
@@ -18,6 +18,7 @@
 package kafka.api
 
 import java.lang.{Long => JLong}
+import java.time.Duration
 import java.util.{Optional, Properties}
 import java.util.concurrent.TimeUnit
 
@@ -31,9 +32,11 @@ import org.apache.kafka.common.{KafkaException, TopicPartition}
 import org.apache.kafka.common.errors.{ProducerFencedException, TimeoutException}
 import org.junit.{After, Before, Test}
 import org.junit.Assert._
+import org.scalatest.Assertions.fail
 
 import scala.collection.JavaConverters._
 import scala.collection.mutable.Buffer
+import scala.collection.Seq
 import scala.concurrent.ExecutionException
 
 class TransactionsTest extends KafkaServerTestHarness {
@@ -386,11 +389,11 @@ class TransactionsTest extends KafkaServerTestHarness {
     val producer2 = transactionalProducers(1)
     producer2.initTransactions()
 
-    assertEquals(offsetAndMetadata, consumer.committed(tp))
+    TestUtils.waitUntilTrue(() => offsetAndMetadata.equals(consumer.committed(Set(tp).asJava).get(tp)), "cannot read committed offset")
   }
 
   @Test
-  def testFencingOnSend() {
+  def testFencingOnSend(): Unit = {
     val producer1 = transactionalProducers(0)
     val producer2 = transactionalProducers(1)
     val consumer = transactionalConsumers(0)
@@ -495,7 +498,7 @@ class TransactionsTest extends KafkaServerTestHarness {
 
     try {
       // Now that the transaction has expired, the second send should fail with a ProducerFencedException.
-      producer.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic2, "2", "2", willBeCommitted = false)).get()
+      producer.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic1, "2", "2", willBeCommitted = false)).get()
       fail("should have raised a ProducerFencedException since the transaction has expired")
     } catch {
       case _: ProducerFencedException =>
@@ -504,9 +507,13 @@ class TransactionsTest extends KafkaServerTestHarness {
     }
 
     // Verify that the first message was aborted and the second one was never written at all.
-    val nonTransactionalConsumer = nonTransactionalConsumers(0)
+    val nonTransactionalConsumer = nonTransactionalConsumers.head
     nonTransactionalConsumer.subscribe(List(topic1).asJava)
-    val records = TestUtils.consumeRecordsFor(nonTransactionalConsumer, 1000)
+
+    // Attempt to consume the one written record. We should not see the second. The
+    // assertion does not strictly guarantee that the record wasn't written, but the
+    // data is small enough that had it been written, it would have been in the first fetch.
+    val records = TestUtils.consumeRecords(nonTransactionalConsumer, numRecords = 1)
     assertEquals(1, records.size)
     assertEquals("1", TestUtils.recordValueAsString(records.head))
 
@@ -578,7 +585,7 @@ class TransactionsTest extends KafkaServerTestHarness {
     try {
       producer.commitTransaction()
     } finally {
-      producer.close(0, TimeUnit.MILLISECONDS)
+      producer.close(Duration.ZERO)
     }
   }
 
diff --git a/core/src/test/scala/integration/kafka/api/UserClientIdQuotaTest.scala b/core/src/test/scala/integration/kafka/api/UserClientIdQuotaTest.scala
index 828c98b414bdf..9a3c4df0c6ad8 100644
--- a/core/src/test/scala/integration/kafka/api/UserClientIdQuotaTest.scala
+++ b/core/src/test/scala/integration/kafka/api/UserClientIdQuotaTest.scala
@@ -31,7 +31,7 @@ class UserClientIdQuotaTest extends BaseQuotaTest {
   override def consumerClientId = "QuotasTestConsumer-!@#$%^&*()"
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     this.serverConfig.setProperty(KafkaConfig.SslClientAuthProp, "required")
     this.serverConfig.setProperty(KafkaConfig.ProducerQuotaBytesPerSecondDefaultProp, Long.MaxValue.toString)
     this.serverConfig.setProperty(KafkaConfig.ConsumerQuotaBytesPerSecondDefaultProp, Long.MaxValue.toString)
@@ -51,7 +51,7 @@ class UserClientIdQuotaTest extends BaseQuotaTest {
         Map("user" -> Sanitizer.sanitize(userPrincipal.getName), "client-id" -> clientId)
       }
 
-      override def overrideQuotas(producerQuota: Long, consumerQuota: Long, requestQuota: Double) {
+      override def overrideQuotas(producerQuota: Long, consumerQuota: Long, requestQuota: Double): Unit = {
         val producerProps = new Properties()
         producerProps.setProperty(DynamicConfig.Client.ProducerByteRateOverrideProp, producerQuota.toString)
         producerProps.setProperty(DynamicConfig.Client.RequestPercentageOverrideProp, requestQuota.toString)
@@ -63,7 +63,7 @@ class UserClientIdQuotaTest extends BaseQuotaTest {
         updateQuotaOverride(userPrincipal.getName, consumerClientId, consumerProps)
       }
 
-      override def removeQuotaOverrides() {
+      override def removeQuotaOverrides(): Unit = {
         val emptyProps = new Properties
         adminZkClient.changeUserOrUserClientIdConfig(Sanitizer.sanitize(userPrincipal.getName) +
           "/clients/" + Sanitizer.sanitize(producerClientId), emptyProps)
@@ -71,7 +71,7 @@ class UserClientIdQuotaTest extends BaseQuotaTest {
           "/clients/" + Sanitizer.sanitize(consumerClientId), emptyProps)
       }
 
-      private def updateQuotaOverride(userPrincipal: String, clientId: String, properties: Properties) {
+      private def updateQuotaOverride(userPrincipal: String, clientId: String, properties: Properties): Unit = {
         adminZkClient.changeUserOrUserClientIdConfig(Sanitizer.sanitize(userPrincipal) + "/clients/" + Sanitizer.sanitize(clientId), properties)
       }
     }
diff --git a/core/src/test/scala/integration/kafka/api/UserQuotaTest.scala b/core/src/test/scala/integration/kafka/api/UserQuotaTest.scala
index d60b8751c7048..37870d9121b0c 100644
--- a/core/src/test/scala/integration/kafka/api/UserQuotaTest.scala
+++ b/core/src/test/scala/integration/kafka/api/UserQuotaTest.scala
@@ -33,7 +33,7 @@ class UserQuotaTest extends BaseQuotaTest with SaslSetup {
   override protected val clientSaslProperties = Some(kafkaClientSaslProperties(kafkaClientSaslMechanism))
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     startSasl(jaasSections(kafkaServerSaslMechanisms, Some("GSSAPI"), KafkaSasl, JaasTestUtils.KafkaServerContextName))
     this.serverConfig.setProperty(KafkaConfig.ProducerQuotaBytesPerSecondDefaultProp, Long.MaxValue.toString)
     this.serverConfig.setProperty(KafkaConfig.ConsumerQuotaBytesPerSecondDefaultProp, Long.MaxValue.toString)
@@ -59,18 +59,18 @@ class UserQuotaTest extends BaseQuotaTest with SaslSetup {
         Map("user" -> userPrincipal.getName, "client-id" -> "")
       }
 
-      override def overrideQuotas(producerQuota: Long, consumerQuota: Long, requestQuota: Double) {
+      override def overrideQuotas(producerQuota: Long, consumerQuota: Long, requestQuota: Double): Unit = {
         val props = quotaProperties(producerQuota, consumerQuota, requestQuota)
         updateQuotaOverride(props)
       }
 
-      override def removeQuotaOverrides() {
+      override def removeQuotaOverrides(): Unit = {
         val emptyProps = new Properties
         updateQuotaOverride(emptyProps)
         updateQuotaOverride(emptyProps)
       }
 
-      private def updateQuotaOverride(properties: Properties) {
+      private def updateQuotaOverride(properties: Properties): Unit = {
         adminZkClient.changeUserOrUserClientIdConfig(Sanitizer.sanitize(userPrincipal.getName), properties)
       }
     }
diff --git a/core/src/test/scala/integration/kafka/network/DynamicConnectionQuotaTest.scala b/core/src/test/scala/integration/kafka/network/DynamicConnectionQuotaTest.scala
index 78d9af7bdfca3..16e3cbc9260c7 100644
--- a/core/src/test/scala/integration/kafka/network/DynamicConnectionQuotaTest.scala
+++ b/core/src/test/scala/integration/kafka/network/DynamicConnectionQuotaTest.scala
@@ -21,106 +21,167 @@ package kafka.network
 import java.io.IOException
 import java.net.{InetAddress, Socket}
 import java.util.Properties
+import java.util.concurrent._
 
 import kafka.server.{BaseRequestTest, KafkaConfig}
-import kafka.utils.TestUtils
-import org.apache.kafka.clients.admin.{AdminClient, AdminClientConfig}
-import org.apache.kafka.common.TopicPartition
+import kafka.utils.{CoreUtils, TestUtils}
+import org.apache.kafka.clients.admin.{Admin, AdminClient, AdminClientConfig}
+import org.apache.kafka.common.{KafkaException, TopicPartition}
 import org.apache.kafka.common.network.ListenerName
 import org.apache.kafka.common.protocol.{ApiKeys, Errors}
 import org.apache.kafka.common.record.{CompressionType, MemoryRecords, SimpleRecord}
 import org.apache.kafka.common.requests.{ProduceRequest, ProduceResponse}
 import org.apache.kafka.common.security.auth.SecurityProtocol
-import org.junit.Assert.assertEquals
-import org.junit.{Before, Test}
+import org.junit.Assert._
+import org.junit.{After, Before, Test}
+import org.scalatest.Assertions.intercept
 
 import scala.collection.JavaConverters._
 
 class DynamicConnectionQuotaTest extends BaseRequestTest {
 
-  override def numBrokers = 1
+  override def brokerCount = 1
 
   val topic = "test"
+  val listener = ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT)
+  val localAddress = InetAddress.getByName("127.0.0.1")
+  var executor: ExecutorService = _
 
   @Before
   override def setUp(): Unit = {
     super.setUp()
-    TestUtils.createTopic(zkClient, topic, numBrokers, numBrokers, servers)
+    TestUtils.createTopic(zkClient, topic, brokerCount, brokerCount, servers)
   }
 
-  @Test
-  def testDynamicConnectionQuota(): Unit = {
-    def connect(socketServer: SocketServer, protocol: SecurityProtocol = SecurityProtocol.PLAINTEXT, localAddr: InetAddress = null) = {
-      new Socket("localhost", socketServer.boundPort(ListenerName.forSecurityProtocol(protocol)), localAddr, 0)
+  @After
+  override def tearDown(): Unit = {
+    try {
+      if (executor != null) {
+        executor.shutdownNow()
+        assertTrue(executor.awaitTermination(10, TimeUnit.SECONDS))
+      }
+    } finally {
+      super.tearDown()
     }
+  }
 
-    val socketServer = servers.head.socketServer
-    val localAddress = InetAddress.getByName("127.0.0.1")
-    def connectionCount = socketServer.connectionCount(localAddress)
-    val initialConnectionCount = connectionCount
+  @Test
+  def testDynamicConnectionQuota(): Unit = {
     val maxConnectionsPerIP = 5
 
+    def connectAndVerify(): Unit = {
+      val socket = connect()
+      try {
+        sendAndReceive(produceRequest, ApiKeys.PRODUCE, socket)
+      } finally {
+        socket.close()
+      }
+    }
+
     val props = new Properties
     props.put(KafkaConfig.MaxConnectionsPerIpProp, maxConnectionsPerIP.toString)
     reconfigureServers(props, perBrokerConfig = false, (KafkaConfig.MaxConnectionsPerIpProp, maxConnectionsPerIP.toString))
 
-    //wait for adminClient connections to close
-    TestUtils.waitUntilTrue(() => initialConnectionCount == connectionCount, "Connection count mismatch")
-
-    //create connections up to maxConnectionsPerIP - 1, leave space for one connection
-    var conns = (connectionCount until (maxConnectionsPerIP - 1)).map(_ => connect(socketServer))
-
-    // produce should succeed
-    var produceResponse = sendProduceRequest()
-    assertEquals(1, produceResponse.responses.size)
-    val (tp, partitionResponse) = produceResponse.responses.asScala.head
-    assertEquals(Errors.NONE, partitionResponse.error)
-
-    TestUtils.waitUntilTrue(() => connectionCount == (maxConnectionsPerIP - 1), "produce request connection is not closed")
-    conns = conns :+ connect(socketServer)
-    // now try one more (should fail)
-    intercept[IOException](sendProduceRequest())
-
-    conns.foreach(conn => conn.close())
-    TestUtils.waitUntilTrue(() => initialConnectionCount == connectionCount, "Connection count mismatch")
+    verifyMaxConnections(maxConnectionsPerIP, () => connectAndVerify)
 
     // Increase MaxConnectionsPerIpOverrides for localhost to 7
     val maxConnectionsPerIPOverride = 7
     props.put(KafkaConfig.MaxConnectionsPerIpOverridesProp, s"localhost:$maxConnectionsPerIPOverride")
     reconfigureServers(props, perBrokerConfig = false, (KafkaConfig.MaxConnectionsPerIpOverridesProp, s"localhost:$maxConnectionsPerIPOverride"))
 
-    //wait for adminClient connections to close
-    TestUtils.waitUntilTrue(() => initialConnectionCount == connectionCount, "Connection count mismatch")
-
-    //create connections up to maxConnectionsPerIPOverride - 1, leave space for one connection
-    conns = (connectionCount until maxConnectionsPerIPOverride - 1).map(_ => connect(socketServer))
+    verifyMaxConnections(maxConnectionsPerIPOverride, () => connectAndVerify)
+  }
 
-    // send should succeed
-    produceResponse = sendProduceRequest()
-    assertEquals(1, produceResponse.responses.size)
-    val (tp1, partitionResponse1) = produceResponse.responses.asScala.head
-    assertEquals(Errors.NONE, partitionResponse1.error)
+  @Test
+  def testDynamicListenerConnectionQuota(): Unit = {
+    val initialConnectionCount = connectionCount
 
-    TestUtils.waitUntilTrue(() => connectionCount == (maxConnectionsPerIPOverride - 1), "produce request connection is not closed")
-    conns = conns :+ connect(socketServer)
-    // now try one more (should fail)
-    intercept[IOException](sendProduceRequest())
+    def connectAndVerify(): Unit = {
+      val socket = connect("PLAINTEXT")
+      socket.setSoTimeout(1000)
+      try {
+        val response = sendAndReceive(produceRequest, ApiKeys.PRODUCE, socket)
+        assertEquals(0, response.remaining)
+      } finally {
+        socket.close()
+      }
+    }
 
-    //close one connection
-    conns.head.close()
-    TestUtils.waitUntilTrue(() => connectionCount == (maxConnectionsPerIPOverride - 1), "connection is not closed")
-    // send should succeed
-    sendProduceRequest()
+    // Reduce total broker MaxConnections to 5 at the cluster level
+    val props = new Properties
+    props.put(KafkaConfig.MaxConnectionsProp, "5")
+    reconfigureServers(props, perBrokerConfig = false, (KafkaConfig.MaxConnectionsProp, "5"))
+    verifyMaxConnections(5, () => connectAndVerify)
+
+    // Create another listener and verify listener connection limit of 5 for each listener
+    val newListeners = "PLAINTEXT://localhost:0,INTERNAL://localhost:0"
+    props.put(KafkaConfig.ListenersProp, newListeners)
+    props.put(KafkaConfig.ListenerSecurityProtocolMapProp, "PLAINTEXT:PLAINTEXT,INTERNAL:PLAINTEXT")
+    props.put(KafkaConfig.MaxConnectionsProp, "10")
+    props.put("listener.name.internal.max.connections", "5")
+    props.put("listener.name.plaintext.max.connections", "5")
+    reconfigureServers(props, perBrokerConfig = true, (KafkaConfig.ListenersProp, newListeners))
+    waitForListener("INTERNAL")
+
+    var conns = (connectionCount until 5).map(_ => connect("PLAINTEXT"))
+    conns ++= (5 until 10).map(_ => connect("INTERNAL"))
+    conns.foreach(verifyConnection)
+    conns.foreach(_.close())
+    TestUtils.waitUntilTrue(() => initialConnectionCount == connectionCount, "Connections not closed")
+
+    // Increase MaxConnections for PLAINTEXT listener to 7 at the broker level
+    val maxConnectionsPlaintext = 7
+    val listenerProp = s"${listener.configPrefix}${KafkaConfig.MaxConnectionsProp}"
+    props.put(listenerProp, maxConnectionsPlaintext.toString)
+    reconfigureServers(props, perBrokerConfig = true, (listenerProp, maxConnectionsPlaintext.toString))
+    verifyMaxConnections(maxConnectionsPlaintext, () => connectAndVerify)
+
+    // Verify that connection blocked on the limit connects successfully when an existing connection is closed
+    val plaintextConnections = (connectionCount until maxConnectionsPlaintext).map(_ => connect("PLAINTEXT"))
+    executor = Executors.newSingleThreadExecutor
+    val future = executor.submit(CoreUtils.runnable { createAndVerifyConnection() })
+    Thread.sleep(100)
+    assertFalse(future.isDone)
+    plaintextConnections.head.close()
+    future.get(30, TimeUnit.SECONDS)
+    plaintextConnections.foreach(_.close())
+    TestUtils.waitUntilTrue(() => initialConnectionCount == connectionCount, "Connections not closed")
+
+    // Verify that connections on inter-broker listener succeed even if broker max connections has been
+    // reached by closing connections on another listener
+    var plaintextConns = (connectionCount until 5).map(_ => connect("PLAINTEXT"))
+    val internalConns = (5 until 10).map(_ => connect("INTERNAL"))
+    plaintextConns.foreach(verifyConnection)
+    internalConns.foreach(verifyConnection)
+    plaintextConns ++= (0 until 2).map(_ => connect("PLAINTEXT"))
+    TestUtils.waitUntilTrue(() => connectionCount <= 10, "Internal connections not closed")
+    plaintextConns.foreach(verifyConnection)
+    intercept[IOException](internalConns.foreach { socket => sendAndReceive(produceRequest, ApiKeys.PRODUCE, socket) })
+    plaintextConns.foreach(_.close())
+    internalConns.foreach(_.close())
+    TestUtils.waitUntilTrue(() => initialConnectionCount == connectionCount, "Connections not closed")
   }
 
   private def reconfigureServers(newProps: Properties, perBrokerConfig: Boolean, aPropToVerify: (String, String)): Unit = {
+    val initialConnectionCount = connectionCount
     val adminClient = createAdminClient()
     TestUtils.alterConfigs(servers, adminClient, newProps, perBrokerConfig).all.get()
     waitForConfigOnServer(aPropToVerify._1, aPropToVerify._2)
     adminClient.close()
+    TestUtils.waitUntilTrue(() => initialConnectionCount == connectionCount, "Admin client connection not closed")
   }
 
-  private def createAdminClient(): AdminClient = {
+  private def waitForListener(listenerName: String): Unit = {
+    TestUtils.retry(maxWaitMs = 10000) {
+      try {
+        assertTrue(servers.head.socketServer.boundPort(ListenerName.normalised(listenerName)) > 0)
+      } catch {
+        case e: KafkaException => throw new AssertionError(e)
+      }
+    }
+  }
+
+  private def createAdminClient(): Admin = {
     val bootstrapServers = TestUtils.bootstrapServers(servers, new ListenerName(securityProtocol.name))
     val config = new Properties()
     config.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers)
@@ -135,12 +196,59 @@ class DynamicConnectionQuotaTest extends BaseRequestTest {
     }
   }
 
-  private def sendProduceRequest(): ProduceResponse = {
+  private def produceRequest: ProduceRequest = {
     val topicPartition = new TopicPartition(topic, 0)
     val memoryRecords = MemoryRecords.withRecords(CompressionType.NONE, new SimpleRecord(System.currentTimeMillis(), "key".getBytes, "value".getBytes))
     val partitionRecords = Map(topicPartition -> memoryRecords)
-    val request = ProduceRequest.Builder.forCurrentMagic(-1, 3000, partitionRecords.asJava).build()
-    val response = connectAndSend(request, ApiKeys.PRODUCE, servers.head.socketServer)
-    ProduceResponse.parse(response, request.version)
+    ProduceRequest.Builder.forCurrentMagic(-1, 3000, partitionRecords.asJava).build()
+  }
+
+  def connectionCount: Int = servers.head.socketServer.connectionCount(localAddress)
+
+  def connect(listener: String): Socket = {
+    val listenerName = ListenerName.normalised(listener)
+    new Socket("localhost", servers.head.socketServer.boundPort(listenerName))
+  }
+
+  private def createAndVerifyConnection(listener: String = "PLAINTEXT"): Unit = {
+    val socket = connect(listener)
+    try {
+      verifyConnection(socket)
+    } finally {
+      socket.close()
+    }
+  }
+
+  private def verifyConnection(socket: Socket): Unit = {
+    val request = produceRequest
+    val response = sendAndReceive(request, ApiKeys.PRODUCE, socket)
+    val produceResponse = ProduceResponse.parse(response, request.version)
+    assertEquals(1, produceResponse.responses.size)
+    val (_, partitionResponse) = produceResponse.responses.asScala.head
+    assertEquals(Errors.NONE, partitionResponse.error)
+  }
+
+  private def verifyMaxConnections(maxConnections: Int, connectWithFailure: () => Unit): Unit = {
+    val initialConnectionCount = connectionCount
+
+    //create connections up to maxConnectionsPerIP - 1, leave space for one connection
+    var conns = (connectionCount until (maxConnections - 1)).map(_ => connect("PLAINTEXT"))
+
+    // produce should succeed on a new connection
+    createAndVerifyConnection()
+
+    TestUtils.waitUntilTrue(() => connectionCount == (maxConnections - 1), "produce request connection is not closed")
+    conns = conns :+ connect("PLAINTEXT")
+
+    // now try one more (should fail)
+    intercept[IOException](connectWithFailure.apply())
+
+    //close one connection
+    conns.head.close()
+    TestUtils.waitUntilTrue(() => connectionCount == (maxConnections - 1), "connection is not closed")
+    createAndVerifyConnection()
+
+    conns.foreach(_.close())
+    TestUtils.waitUntilTrue(() => initialConnectionCount == connectionCount, "Connections not closed")
   }
 }
diff --git a/core/src/test/scala/integration/kafka/server/DelayedFetchTest.scala b/core/src/test/scala/integration/kafka/server/DelayedFetchTest.scala
index 6a5a33d2702a9..317d8f78d4a36 100644
--- a/core/src/test/scala/integration/kafka/server/DelayedFetchTest.scala
+++ b/core/src/test/scala/integration/kafka/server/DelayedFetchTest.scala
@@ -18,9 +18,11 @@ package kafka.server
 
 import java.util.Optional
 
-import kafka.cluster.Partition
+import scala.collection.Seq
+import kafka.cluster.{Partition, Replica}
+import kafka.log.LogOffsetSnapshot
 import org.apache.kafka.common.TopicPartition
-import org.apache.kafka.common.errors.FencedLeaderEpochException
+import org.apache.kafka.common.errors.{FencedLeaderEpochException, ReplicaNotAvailableException}
 import org.apache.kafka.common.protocol.Errors
 import org.apache.kafka.common.record.MemoryRecords
 import org.apache.kafka.common.requests.FetchRequest
@@ -56,6 +58,7 @@ class DelayedFetchTest extends EasyMockSupport {
       fetchMetadata = fetchMetadata,
       replicaManager = replicaManager,
       quota = replicaQuota,
+      clientMetadata = None,
       responseCallback = callback)
 
     val partition: Partition = mock(classOf[Partition])
@@ -77,6 +80,127 @@ class DelayedFetchTest extends EasyMockSupport {
     assertEquals(Errors.FENCED_LEADER_EPOCH, fetchResult.error)
   }
 
+  @Test
+  def testReplicaNotAvailable(): Unit = {
+    val topicPartition = new TopicPartition("topic", 0)
+    val fetchOffset = 500L
+    val logStartOffset = 0L
+    val currentLeaderEpoch = Optional.of[Integer](10)
+    val replicaId = 1
+
+    val fetchStatus = FetchPartitionStatus(
+      startOffsetMetadata = LogOffsetMetadata(fetchOffset),
+      fetchInfo = new FetchRequest.PartitionData(fetchOffset, logStartOffset, maxBytes, currentLeaderEpoch))
+    val fetchMetadata = buildFetchMetadata(replicaId, topicPartition, fetchStatus)
+
+    var fetchResultOpt: Option[FetchPartitionData] = None
+    def callback(responses: Seq[(TopicPartition, FetchPartitionData)]): Unit = {
+      fetchResultOpt = Some(responses.head._2)
+    }
+
+    val delayedFetch = new DelayedFetch(
+      delayMs = 500,
+      fetchMetadata = fetchMetadata,
+      replicaManager = replicaManager,
+      quota = replicaQuota,
+      clientMetadata = None,
+      responseCallback = callback)
+
+    val partition: Partition = mock(classOf[Partition])
+
+    EasyMock.expect(replicaManager.getPartitionOrException(topicPartition, expectLeader = true))
+      .andThrow(new ReplicaNotAvailableException(s"Replica for $topicPartition not available"))
+    expectReadFromReplicaWithError(replicaId, topicPartition, fetchStatus.fetchInfo, Errors.REPLICA_NOT_AVAILABLE)
+
+    replayAll()
+
+    assertTrue(delayedFetch.tryComplete())
+    assertTrue(delayedFetch.isCompleted)
+    assertTrue(fetchResultOpt.isDefined)
+  }
+
+  def checkCompleteWhenFollowerLaggingHW(followerHW: Option[Long], checkResult: DelayedFetch => Unit): Unit = {
+    val topicPartition = new TopicPartition("topic", 0)
+    val fetchOffset = 500L
+    val logStartOffset = 0L
+    val currentLeaderEpoch = Optional.of[Integer](10)
+    val replicaId = 1
+
+    val fetchStatus = FetchPartitionStatus(
+      startOffsetMetadata = LogOffsetMetadata(fetchOffset),
+      fetchInfo = new FetchRequest.PartitionData(fetchOffset, logStartOffset, maxBytes, currentLeaderEpoch))
+    val fetchMetadata = buildFetchMetadata(replicaId, topicPartition, fetchStatus)
+
+    var fetchResultOpt: Option[FetchPartitionData] = None
+    def callback(responses: Seq[(TopicPartition, FetchPartitionData)]): Unit = {
+      fetchResultOpt = Some(responses.head._2)
+    }
+
+    val delayedFetch = new DelayedFetch(
+      delayMs = 500,
+      fetchMetadata = fetchMetadata,
+      replicaManager = replicaManager,
+      quota = replicaQuota,
+      clientMetadata = None,
+      responseCallback = callback
+    )
+
+    val partition: Partition = mock(classOf[Partition])
+
+    EasyMock.expect(replicaManager.getPartitionOrException(topicPartition, expectLeader = true))
+      .andReturn(partition)
+    EasyMock.expect(partition.fetchOffsetSnapshot(currentLeaderEpoch, fetchOnlyFromLeader = true))
+      .andReturn(
+        LogOffsetSnapshot(
+          logStartOffset = 0,
+          logEndOffset = new LogOffsetMetadata(500L),
+          highWatermark = new LogOffsetMetadata(480L),
+          lastStableOffset = new LogOffsetMetadata(400L)))
+
+    expectReadFromReplica(replicaId, topicPartition, fetchStatus.fetchInfo)
+
+    val follower = new Replica(replicaId, topicPartition)
+    followerHW.foreach(hw => {
+      follower.updateFetchState(LogOffsetMetadata.UnknownOffsetMetadata, 0L, 0L, 0L, hw)
+    })
+    EasyMock.expect(partition.getReplica(replicaId))
+        .andReturn(Some(follower))
+
+    replayAll()
+    checkResult.apply(delayedFetch)
+  }
+
+  @Test
+  def testCompleteWhenFollowerLaggingHW(): Unit = {
+    // No HW from the follower, should complete
+    resetAll
+    checkCompleteWhenFollowerLaggingHW(None, delayedFetch => {
+      assertTrue(delayedFetch.tryComplete())
+      assertTrue(delayedFetch.isCompleted)
+    })
+
+    // A higher HW from the follower (shouldn't actually be possible)
+    resetAll
+    checkCompleteWhenFollowerLaggingHW(Some(500), delayedFetch => {
+      assertFalse(delayedFetch.tryComplete())
+      assertFalse(delayedFetch.isCompleted)
+    })
+
+    // An equal HW from follower
+    resetAll
+    checkCompleteWhenFollowerLaggingHW(Some(480), delayedFetch => {
+      assertFalse(delayedFetch.tryComplete())
+      assertFalse(delayedFetch.isCompleted)
+    })
+
+    // A lower HW from follower, should complete the fetch
+    resetAll
+    checkCompleteWhenFollowerLaggingHW(Some(470), delayedFetch => {
+      assertTrue(delayedFetch.tryComplete())
+      assertTrue(delayedFetch.isCompleted)
+    })
+  }
+
   private def buildFetchMetadata(replicaId: Int,
                                  topicPartition: TopicPartition,
                                  fetchStatus: FetchPartitionStatus): FetchMetadata = {
@@ -101,10 +225,38 @@ class DelayedFetchTest extends EasyMockSupport {
       fetchMaxBytes = maxBytes,
       hardMaxBytesLimit = false,
       readPartitionInfo = Seq((topicPartition, fetchPartitionData)),
+      clientMetadata = None,
       quota = replicaQuota))
       .andReturn(Seq((topicPartition, buildReadResultWithError(error))))
   }
 
+  private def expectReadFromReplica(replicaId: Int,
+                                    topicPartition: TopicPartition,
+                                    fetchPartitionData: FetchRequest.PartitionData): Unit = {
+    val result = LogReadResult(
+      exception = None,
+      info = FetchDataInfo(LogOffsetMetadata.UnknownOffsetMetadata, MemoryRecords.EMPTY),
+      highWatermark = -1L,
+      leaderLogStartOffset = -1L,
+      leaderLogEndOffset = -1L,
+      followerLogStartOffset = -1L,
+      fetchTimeMs = -1L,
+      readSize = -1,
+      lastStableOffset = None)
+
+
+    EasyMock.expect(replicaManager.readFromLocalLog(
+      replicaId = replicaId,
+      fetchOnlyFromLeader = true,
+      fetchIsolation = FetchLogEnd,
+      fetchMaxBytes = maxBytes,
+      hardMaxBytesLimit = false,
+      readPartitionInfo = Seq((topicPartition, fetchPartitionData)),
+      clientMetadata = None,
+      quota = replicaQuota))
+      .andReturn(Seq((topicPartition, result))).anyTimes()
+  }
+
   private def buildReadResultWithError(error: Errors): LogReadResult = {
     LogReadResult(
       exception = Some(error.exception),
diff --git a/core/src/test/scala/integration/kafka/server/DynamicBrokerReconfigurationTest.scala b/core/src/test/scala/integration/kafka/server/DynamicBrokerReconfigurationTest.scala
index 80ed131ec244d..715ff20c9134a 100644
--- a/core/src/test/scala/integration/kafka/server/DynamicBrokerReconfigurationTest.scala
+++ b/core/src/test/scala/integration/kafka/server/DynamicBrokerReconfigurationTest.scala
@@ -18,7 +18,7 @@
 
 package kafka.server
 
-import java.io.{Closeable, File, FileWriter}
+import java.io.{Closeable, File, FileWriter, IOException, Reader, StringReader}
 import java.nio.file.{Files, Paths, StandardCopyOption}
 import java.lang.management.ManagementFactory
 import java.security.KeyStore
@@ -26,12 +26,13 @@ import java.time.Duration
 import java.util
 import java.util.{Collections, Properties}
 import java.util.concurrent._
-
 import javax.management.ObjectName
+
 import com.yammer.metrics.Metrics
 import com.yammer.metrics.core.MetricName
 import kafka.admin.ConfigCommand
 import kafka.api.{KafkaSasl, SaslSetup}
+import kafka.controller.{ControllerBrokerStateInfo, ControllerChannelManager}
 import kafka.log.LogConfig
 import kafka.message.ProducerCompressionCodec
 import kafka.network.{Processor, RequestChannel}
@@ -39,6 +40,7 @@ import kafka.utils._
 import kafka.utils.Implicits._
 import kafka.zk.{ConfigEntityChangeNotificationZNode, ZooKeeperTestHarness}
 import org.apache.kafka.clients.CommonClientConfigs
+import org.apache.kafka.clients.admin.AlterConfigOp.OpType
 import org.apache.kafka.clients.admin.ConfigEntry.{ConfigSource, ConfigSynonym}
 import org.apache.kafka.clients.admin._
 import org.apache.kafka.clients.consumer.{ConsumerConfig, ConsumerRecord, ConsumerRecords, KafkaConsumer}
@@ -47,6 +49,7 @@ import org.apache.kafka.common.{ClusterResource, ClusterResourceListener, Reconf
 import org.apache.kafka.common.config.{ConfigException, ConfigResource}
 import org.apache.kafka.common.config.SslConfigs._
 import org.apache.kafka.common.config.types.Password
+import org.apache.kafka.common.config.provider.FileConfigProvider
 import org.apache.kafka.common.errors.{AuthenticationException, InvalidRequestException}
 import org.apache.kafka.common.internals.Topic
 import org.apache.kafka.common.metrics.{KafkaMetric, MetricsReporter}
@@ -55,9 +58,10 @@ import org.apache.kafka.common.network.CertStores.{KEYSTORE_PROPS, TRUSTSTORE_PR
 import org.apache.kafka.common.record.TimestampType
 import org.apache.kafka.common.security.auth.SecurityProtocol
 import org.apache.kafka.common.serialization.{StringDeserializer, StringSerializer}
-import org.apache.kafka.test.TestSslUtils
+import org.apache.kafka.test.{TestSslUtils, TestUtils => JTestUtils}
 import org.junit.Assert._
 import org.junit.{After, Before, Ignore, Test}
+import org.scalatest.Assertions.intercept
 
 import scala.collection._
 import scala.collection.mutable.ArrayBuffer
@@ -78,7 +82,7 @@ class DynamicBrokerReconfigurationTest extends ZooKeeperTestHarness with SaslSet
   private val numPartitions = 10
   private val producers = new ArrayBuffer[KafkaProducer[String, String]]
   private val consumers = new ArrayBuffer[KafkaConsumer[String, String]]
-  private val adminClients = new ArrayBuffer[AdminClient]()
+  private val adminClients = new ArrayBuffer[Admin]()
   private val clientThreads = new ArrayBuffer[ShutdownableThread]()
   private val executors = new ArrayBuffer[ExecutorService]
   private val topic = "testtopic"
@@ -115,6 +119,8 @@ class DynamicBrokerReconfigurationTest extends ZooKeeperTestHarness with SaslSet
       props.put(KafkaConfig.NumReplicaFetchersProp, "2") // greater than one to test reducing threads
       props.put(KafkaConfig.ProducerQuotaBytesPerSecondDefaultProp, "10000000") // non-default value to trigger a new metric
       props.put(KafkaConfig.PasswordEncoderSecretProp, "dynamic-config-secret")
+      props.put(KafkaConfig.LogRetentionTimeMillisProp, 1680000000.toString)
+      props.put(KafkaConfig.LogRetentionTimeHoursProp, 168.toString)
 
       props ++= sslProperties1
       props ++= securityProps(sslProperties1, KEYSTORE_PROPS, listenerPrefix(SecureInternal))
@@ -141,12 +147,12 @@ class DynamicBrokerReconfigurationTest extends ZooKeeperTestHarness with SaslSet
   }
 
   @After
-  override def tearDown() {
+  override def tearDown(): Unit = {
     clientThreads.foreach(_.interrupt())
     clientThreads.foreach(_.initiateShutdown())
     clientThreads.foreach(_.join(5 * 1000))
     executors.foreach(_.shutdownNow())
-    producers.foreach(_.close(0, TimeUnit.MILLISECONDS))
+    producers.foreach(_.close(Duration.ZERO))
     consumers.foreach(_.close(Duration.ofMillis(0)))
     adminClients.foreach(_.close())
     TestUtils.shutdownServers(servers)
@@ -155,9 +161,10 @@ class DynamicBrokerReconfigurationTest extends ZooKeeperTestHarness with SaslSet
   }
 
   @Test
-  def testKeyStoreDescribeUsingAdminClient(): Unit = {
+  def testConfigDescribeUsingAdminClient(): Unit = {
 
-    def verifyConfig(configName: String, configEntry: ConfigEntry, isSensitive: Boolean, expectedProps: Properties): Unit = {
+    def verifyConfig(configName: String, configEntry: ConfigEntry, isSensitive: Boolean, isReadOnly: Boolean,
+                     expectedProps: Properties): Unit = {
       if (isSensitive) {
         assertTrue(s"Value is sensitive: $configName", configEntry.isSensitive)
         assertNull(s"Sensitive value returned for $configName", configEntry.value)
@@ -165,6 +172,7 @@ class DynamicBrokerReconfigurationTest extends ZooKeeperTestHarness with SaslSet
         assertFalse(s"Config is not sensitive: $configName", configEntry.isSensitive)
         assertEquals(expectedProps.getProperty(configName), configEntry.value)
       }
+      assertEquals(s"isReadOnly incorrect for $configName: $configEntry", isReadOnly, configEntry.isReadOnly)
     }
 
     def verifySynonym(configName: String, synonym: ConfigSynonym, isSensitive: Boolean,
@@ -199,7 +207,7 @@ class DynamicBrokerReconfigurationTest extends ZooKeeperTestHarness with SaslSet
       KEYSTORE_PROPS.asScala.foreach { configName =>
         val desc = configEntry(configDesc, s"$prefix$configName")
         val isSensitive = configName.contains("password")
-        verifyConfig(configName, desc, isSensitive, if (prefix.isEmpty) invalidSslProperties else sslProperties1)
+        verifyConfig(configName, desc, isSensitive, isReadOnly = prefix.nonEmpty, if (prefix.isEmpty) invalidSslProperties else sslProperties1)
         val defaultValue = if (configName == SSL_KEYSTORE_TYPE_CONFIG) Some("JKS") else None
         verifySynonyms(configName, desc.synonyms, isSensitive, prefix, defaultValue)
       }
@@ -211,6 +219,107 @@ class DynamicBrokerReconfigurationTest extends ZooKeeperTestHarness with SaslSet
     val configDesc = describeConfig(adminClient)
     verifySslConfig("listener.name.external.", sslProperties1, configDesc)
     verifySslConfig("", invalidSslProperties, configDesc)
+
+    // Verify a few log configs with and without synonyms
+    val expectedProps = new Properties
+    expectedProps.setProperty(KafkaConfig.LogRetentionTimeMillisProp, "1680000000")
+    expectedProps.setProperty(KafkaConfig.LogRetentionTimeHoursProp, "168")
+    expectedProps.setProperty(KafkaConfig.LogRollTimeHoursProp, "168")
+    expectedProps.setProperty(KafkaConfig.LogCleanerThreadsProp, "1")
+    val logRetentionMs = configEntry(configDesc, KafkaConfig.LogRetentionTimeMillisProp)
+    verifyConfig(KafkaConfig.LogRetentionTimeMillisProp, logRetentionMs,
+      isSensitive = false, isReadOnly = false, expectedProps)
+    val logRetentionHours = configEntry(configDesc, KafkaConfig.LogRetentionTimeHoursProp)
+    verifyConfig(KafkaConfig.LogRetentionTimeHoursProp, logRetentionHours,
+      isSensitive = false, isReadOnly = true, expectedProps)
+    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,
+      isSensitive = false, isReadOnly = false, expectedProps)
+
+    def synonymsList(configEntry: ConfigEntry): List[(String, ConfigSource)] =
+      configEntry.synonyms.asScala.map(s => (s.name, s.source)).toList
+    assertEquals(List((KafkaConfig.LogRetentionTimeMillisProp, ConfigSource.STATIC_BROKER_CONFIG),
+      (KafkaConfig.LogRetentionTimeHoursProp, ConfigSource.STATIC_BROKER_CONFIG),
+      (KafkaConfig.LogRetentionTimeHoursProp, ConfigSource.DEFAULT_CONFIG)),
+      synonymsList(logRetentionMs))
+    assertEquals(List((KafkaConfig.LogRetentionTimeHoursProp, ConfigSource.STATIC_BROKER_CONFIG),
+      (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))
+  }
+
+  @Test
+  def testUpdatesUsingConfigProvider(): Unit = {
+    val PollingIntervalVal = "${file:polling.interval:interval}"
+    val PollingIntervalUpdateVal = "${file:polling.interval:updinterval}"
+    val SslTruststoreTypeVal = "${file:ssl.truststore.type:storetype}"
+    val SslKeystorePasswordVal = "${file:ssl.keystore.password:password}"
+
+    val configPrefix = listenerPrefix(SecureExternal)
+    val brokerConfigs = describeConfig(adminClients.head, servers).entries.asScala
+    // the following are values before updated
+    assertTrue("Initial value of polling interval", brokerConfigs.find(_.name == TestMetricsReporter.PollingIntervalProp) == None)
+    assertTrue("Initial value of ssl truststore type", brokerConfigs.find(_.name == configPrefix+KafkaConfig.SslTruststoreTypeProp) == None)
+    assertTrue("Initial value of ssl keystore password", brokerConfigs.find(_.name == configPrefix+KafkaConfig.SslKeystorePasswordProp).get.value == null)
+
+    // setup ssl properties
+    val secProps = securityProps(sslProperties1, KEYSTORE_PROPS, configPrefix)
+
+    // configure config providers and properties need be updated
+    val updatedProps = new Properties
+    updatedProps.setProperty("config.providers", "file")
+    updatedProps.setProperty("config.providers.file.class", "kafka.server.MockFileConfigProvider")
+    updatedProps.put(KafkaConfig.MetricReporterClassesProp, classOf[TestMetricsReporter].getName)
+
+    // 1. update Integer property using config provider
+    updatedProps.put(TestMetricsReporter.PollingIntervalProp, PollingIntervalVal)
+
+    // 2. update String property using config provider
+    updatedProps.put(configPrefix+KafkaConfig.SslTruststoreTypeProp, SslTruststoreTypeVal)
+
+    // merge two properties
+    updatedProps ++= secProps
+
+    // 3. update password property using config provider
+    updatedProps.put(configPrefix+KafkaConfig.SslKeystorePasswordProp, SslKeystorePasswordVal)
+
+    alterConfigsUsingConfigCommand(updatedProps)
+    waitForConfig(TestMetricsReporter.PollingIntervalProp, "1000")
+    waitForConfig(configPrefix+KafkaConfig.SslTruststoreTypeProp, "JKS")
+    waitForConfig(configPrefix+KafkaConfig.SslKeystorePasswordProp, "ServerPassword")
+
+    // wait for MetricsReporter
+    val reporters = TestMetricsReporter.waitForReporters(servers.size)
+    reporters.foreach { reporter =>
+      reporter.verifyState(reconfigureCount = 0, deleteCount = 0, pollingInterval = 1000)
+      assertFalse("No metrics found", reporter.kafkaMetrics.isEmpty)
+    }
+
+    // fetch from ZK, values should be unresolved
+    val props = fetchBrokerConfigsFromZooKeeper(servers.head)
+    assertTrue("polling interval is not updated in ZK", props.getProperty(TestMetricsReporter.PollingIntervalProp) == PollingIntervalVal)
+    assertTrue("store type is not updated in ZK", props.getProperty(configPrefix+KafkaConfig.SslTruststoreTypeProp) == SslTruststoreTypeVal)
+    assertTrue("keystore password is not updated in ZK", props.getProperty(configPrefix+KafkaConfig.SslKeystorePasswordProp) == SslKeystorePasswordVal)
+
+    // verify the update
+    // 1. verify update not occurring if the value of property is same.
+    alterConfigsUsingConfigCommand(updatedProps)
+    waitForConfig(TestMetricsReporter.PollingIntervalProp, "1000")
+    reporters.foreach { reporter =>
+      reporter.verifyState(reconfigureCount = 0, deleteCount = 0, pollingInterval = 1000)
+    }
+
+    // 2. verify update occurring if the value of property changed.
+    updatedProps.put(TestMetricsReporter.PollingIntervalProp, PollingIntervalUpdateVal)
+    alterConfigsUsingConfigCommand(updatedProps)
+    waitForConfig(TestMetricsReporter.PollingIntervalProp, "2000")
+    reporters.foreach { reporter =>
+      reporter.verifyState(reconfigureCount = 1, deleteCount = 0, pollingInterval = 2000)
+    }
   }
 
   @Test
@@ -253,6 +362,29 @@ class DynamicBrokerReconfigurationTest extends ZooKeeperTestHarness with SaslSet
     alterSslKeystore(adminClient, sslPropertiesCopy, SecureInternal)
     verifyProduceConsume(producer, consumer, 10, topic2)
 
+    // Verify that keystores can be updated using same file name.
+    val reusableProps = sslProperties2.clone().asInstanceOf[Properties]
+    val reusableFile = File.createTempFile("keystore", ".jks")
+    reusableProps.setProperty(SSL_KEYSTORE_LOCATION_CONFIG, reusableFile.getPath)
+    Files.copy(new File(sslProperties1.getProperty(SSL_KEYSTORE_LOCATION_CONFIG)).toPath,
+      reusableFile.toPath, StandardCopyOption.REPLACE_EXISTING)
+    alterSslKeystore(adminClient, reusableProps, SecureExternal)
+    val producer3 = ProducerBuilder().trustStoreProps(sslProperties2).maxRetries(0).build()
+    verifyAuthenticationFailure(producer3)
+    // Now alter using same file name. We can't check if the update has completed by comparing config on
+    // the broker, so we wait for producer operation to succeed to verify that the update has been performed.
+    Files.copy(new File(sslProperties2.getProperty(SSL_KEYSTORE_LOCATION_CONFIG)).toPath,
+      reusableFile.toPath, StandardCopyOption.REPLACE_EXISTING)
+    reusableFile.setLastModified(System.currentTimeMillis() + 1000)
+    alterSslKeystore(adminClient, reusableProps, SecureExternal)
+    TestUtils.waitUntilTrue(() => {
+      try {
+        producer3.partitionsFor(topic).size() == numPartitions
+      } catch {
+        case _: Exception  => false
+      }
+    }, "Keystore not updated")
+
     // Verify that all messages sent with retries=0 while keystores were being altered were consumed
     stopAndVerifyProduceConsume(producerThread, consumerThread)
   }
@@ -306,9 +438,28 @@ class DynamicBrokerReconfigurationTest extends ZooKeeperTestHarness with SaslSet
     Files.copy(Paths.get(combinedStoreProps.getProperty(SSL_TRUSTSTORE_LOCATION_CONFIG)),
       Paths.get(sslProperties1.getProperty(SSL_TRUSTSTORE_LOCATION_CONFIG)),
       StandardCopyOption.REPLACE_EXISTING)
-    TestUtils.alterConfigs(servers, adminClients.head, oldTruststoreProps, perBrokerConfig = true).all.get()
+    TestUtils.incrementalAlterConfigs(servers, adminClients.head, oldTruststoreProps, perBrokerConfig = true).all.get()
     verifySslProduceConsume(sslProperties1, "alter-truststore-4")
     verifySslProduceConsume(sslProperties2, "alter-truststore-5")
+
+    // Update internal keystore/truststore and validate new client connections from broker (e.g. controller).
+    // Alter internal keystore from `sslProperties1` to `sslProperties2`, force disconnect of a controller connection
+    // and verify that metadata is propagated for new topic.
+    val props2 = securityProps(sslProperties2, KEYSTORE_PROPS, prefix)
+    props2 ++= securityProps(combinedStoreProps, TRUSTSTORE_PROPS, prefix)
+    TestUtils.incrementalAlterConfigs(servers, adminClients.head, props2, perBrokerConfig = true).all.get(15, TimeUnit.SECONDS)
+    verifySslProduceConsume(sslProperties2, "alter-truststore-6")
+    props2 ++= securityProps(sslProperties2, TRUSTSTORE_PROPS, prefix)
+    TestUtils.incrementalAlterConfigs(servers, adminClients.head, props2, perBrokerConfig = true).all.get(15, TimeUnit.SECONDS)
+    verifySslProduceConsume(sslProperties2, "alter-truststore-7")
+    waitForAuthenticationFailure(producerBuilder.keyStoreProps(sslProperties1))
+
+    val controller = servers.find(_.config.brokerId == TestUtils.waitUntilControllerElected(zkClient)).get
+    val controllerChannelManager = controller.kafkaController.controllerChannelManager
+    val brokerStateInfo: mutable.HashMap[Int, ControllerBrokerStateInfo] =
+      JTestUtils.fieldValue(controllerChannelManager, classOf[ControllerChannelManager], "brokerStateInfo")
+    brokerStateInfo(0).networkClient.disconnect("0")
+    TestUtils.createTopic(zkClient, "testtopic2", numPartitions, replicationFactor = numServers, servers)
   }
 
   @Test
@@ -353,6 +504,40 @@ class DynamicBrokerReconfigurationTest extends ZooKeeperTestHarness with SaslSet
     stopAndVerifyProduceConsume(producerThread, consumerThread)
   }
 
+  @Test
+  def testConsecutiveConfigChange(): Unit = {
+    val topic2 = "testtopic2"
+    val topicProps = new Properties
+    topicProps.put(KafkaConfig.MinInSyncReplicasProp, "2")
+    TestUtils.createTopic(zkClient, topic2, 1, replicationFactor = numServers, servers, topicProps)
+    var log = servers.head.logManager.getLog(new TopicPartition(topic2, 0)).getOrElse(throw new IllegalStateException("Log not found"))
+    assertTrue(log.config.overriddenConfigs.contains(KafkaConfig.MinInSyncReplicasProp))
+    assertEquals("2", log.config.originals().get(KafkaConfig.MinInSyncReplicasProp).toString)
+
+    val props = new Properties
+    props.put(KafkaConfig.MinInSyncReplicasProp, "3")
+    // Make a broker-default config
+    reconfigureServers(props, perBrokerConfig = false, (KafkaConfig.MinInSyncReplicasProp, "3"))
+    // Verify that all broker defaults have been updated again
+    servers.foreach { server =>
+      props.asScala.foreach { case (k, v) =>
+        assertEquals(s"Not reconfigured $k", v, server.config.originals.get(k).toString)
+      }
+    }
+
+    log = servers.head.logManager.getLog(new TopicPartition(topic2, 0)).getOrElse(throw new IllegalStateException("Log not found"))
+    assertTrue(log.config.overriddenConfigs.contains(KafkaConfig.MinInSyncReplicasProp))
+    assertEquals("2", log.config.originals().get(KafkaConfig.MinInSyncReplicasProp).toString) // Verify topic-level config survives
+
+    // Make a second broker-default change
+    props.clear()
+    props.put(KafkaConfig.LogRetentionTimeMillisProp, "604800000")
+    reconfigureServers(props, perBrokerConfig = false, (KafkaConfig.LogRetentionTimeMillisProp, "604800000"))
+    log = servers.head.logManager.getLog(new TopicPartition(topic2, 0)).getOrElse(throw new IllegalStateException("Log not found"))
+    assertTrue(log.config.overriddenConfigs.contains(KafkaConfig.MinInSyncReplicasProp))
+    assertEquals("2", log.config.originals().get(KafkaConfig.MinInSyncReplicasProp).toString) // Verify topic-level config still survives
+  }
+
   @Test
   def testDefaultTopicConfig(): Unit = {
     val (producerThread, consumerThread) = startProduceConsume(retries = 0)
@@ -391,7 +576,9 @@ class DynamicBrokerReconfigurationTest extends ZooKeeperTestHarness with SaslSet
 
     // Verify that configs of existing logs have been updated
     val newLogConfig = LogConfig(KafkaServer.copyKafkaConfigToLog(servers.head.config))
-    assertEquals(newLogConfig, servers.head.logManager.currentDefaultConfig)
+    TestUtils.waitUntilTrue(() => servers.head.logManager.currentDefaultConfig == newLogConfig,
+      "Config not updated in LogManager")
+
     val log = servers.head.logManager.getLog(new TopicPartition(topic, 0)).getOrElse(throw new IllegalStateException("Log not found"))
     TestUtils.waitUntilTrue(() => log.config.segmentSize == 4000, "Existing topic config using defaults not updated")
     props.asScala.foreach { case (k, v) =>
@@ -484,7 +671,7 @@ class DynamicBrokerReconfigurationTest extends ZooKeeperTestHarness with SaslSet
     // Enable unclean leader election
     val newProps = new Properties
     newProps.put(KafkaConfig.UncleanLeaderElectionEnableProp, "true")
-    TestUtils.alterConfigs(servers, adminClients.head, newProps, perBrokerConfig = false).all.get
+    TestUtils.incrementalAlterConfigs(servers, adminClients.head, newProps, perBrokerConfig = false).all.get
     waitForConfigOnServer(controller, KafkaConfig.UncleanLeaderElectionEnableProp, "true")
 
     // Verify that the old follower with missing records is elected as the new leader
@@ -621,7 +808,7 @@ class DynamicBrokerReconfigurationTest extends ZooKeeperTestHarness with SaslSet
     val partitions = (0 until numPartitions).map(i => new TopicPartition(topic, i)).filter { tp =>
       zkClient.getLeaderForPartition(tp).contains(leaderId)
     }
-    assertTrue(s"Partitons not found with leader $leaderId", partitions.nonEmpty)
+    assertTrue(s"Partitions not found with leader $leaderId", partitions.nonEmpty)
     partitions.foreach { tp =>
       (1 to 2).foreach { i =>
         val replicaFetcherManager = servers(i).replicaManager.replicaFetcherManager
@@ -686,12 +873,12 @@ class DynamicBrokerReconfigurationTest extends ZooKeeperTestHarness with SaslSet
     servers.foreach { server =>
       assertEquals(classOf[TestMetricsReporter].getName, server.config.originals.get(KafkaConfig.MetricReporterClassesProp))
     }
-    newReporters.foreach(_.verifyState(reconfigureCount = 1, deleteCount = 0, pollingInterval = 2000))
+    newReporters.foreach(_.verifyState(reconfigureCount = 0, deleteCount = 0, pollingInterval = 2000))
 
     // Verify that validation failure of custom config fails reconfiguration and leaves config unchanged
     newProps.put(TestMetricsReporter.PollingIntervalProp, "invalid")
     reconfigureServers(newProps, perBrokerConfig = false, (TestMetricsReporter.PollingIntervalProp, "2000"), expectFailure = true)
-    newReporters.foreach(_.verifyState(reconfigureCount = 1, deleteCount = 0, pollingInterval = 2000))
+    newReporters.foreach(_.verifyState(reconfigureCount = 0, deleteCount = 0, pollingInterval = 2000))
 
     // Delete reporters
     configureMetricsReporters(Seq.empty[Class[_]], newProps)
@@ -704,7 +891,13 @@ class DynamicBrokerReconfigurationTest extends ZooKeeperTestHarness with SaslSet
     alterConfigsOnServer(servers.head, newProps)
     TestUtils.waitUntilTrue(() => !TestMetricsReporter.testReporters.isEmpty, "Metrics reporter not created")
     val perBrokerReporter = TestMetricsReporter.waitForReporters(1).head
-    perBrokerReporter.verifyState(reconfigureCount = 1, deleteCount = 0, pollingInterval = 4000)
+    perBrokerReporter.verifyState(reconfigureCount = 0, deleteCount = 0, pollingInterval = 4000)
+
+    // update TestMetricsReporter.PollingIntervalProp to 3000
+    newProps.put(TestMetricsReporter.PollingIntervalProp, "3000")
+    alterConfigsOnServer(servers.head, newProps)
+    perBrokerReporter.verifyState(reconfigureCount = 1, deleteCount = 0, pollingInterval = 3000)
+
     servers.tail.foreach { server => assertEquals("", server.config.originals.get(KafkaConfig.MetricReporterClassesProp)) }
 
     // Verify that produce/consume worked throughout this test without any retries in producer
@@ -883,7 +1076,7 @@ class DynamicBrokerReconfigurationTest extends ZooKeeperTestHarness with SaslSet
     val unknownConfig = "some.config"
     props.put(unknownConfig, "some.config.value")
 
-    TestUtils.alterConfigs(servers, adminClients.head, props, perBrokerConfig = true).all.get
+    TestUtils.incrementalAlterConfigs(servers, adminClients.head, props, perBrokerConfig = true).all.get
 
     TestUtils.waitUntilTrue(() => servers.forall(server => server.config.listeners.size == existingListenerCount + 1),
       "Listener config not updated")
@@ -908,6 +1101,8 @@ class DynamicBrokerReconfigurationTest extends ZooKeeperTestHarness with SaslSet
   private def verifyAddListener(listenerName: String, securityProtocol: SecurityProtocol,
                                 saslMechanisms: Seq[String]): Unit = {
     addListener(servers, listenerName, securityProtocol, saslMechanisms)
+    TestUtils.waitUntilTrue(() => servers.forall(hasListenerMetric(_, listenerName)),
+      "Processors not started for new listener")
     if (saslMechanisms.nonEmpty)
       saslMechanisms.foreach { mechanism =>
         verifyListener(securityProtocol, Some(mechanism), s"add-listener-group-$securityProtocol-$mechanism")
@@ -931,8 +1126,6 @@ class DynamicBrokerReconfigurationTest extends ZooKeeperTestHarness with SaslSet
       .autoOffsetReset("latest")
       .build()
     verifyProduceConsume(producer1, consumer1, numRecords = 10, topic)
-    // send another message to check consumer later
-    producer1.send(new ProducerRecord(topic, "key", "value")).get(1, TimeUnit.SECONDS)
 
     val config = servers.head.config
     val existingListenerCount = config.listeners.size
@@ -946,14 +1139,21 @@ class DynamicBrokerReconfigurationTest extends ZooKeeperTestHarness with SaslSet
       .mkString(",")
 
     val props = fetchBrokerConfigsFromZooKeeper(servers.head)
-    val listenerProps = props.asScala.keySet.filter(_.startsWith(listenerPrefix(listenerName)))
-    listenerProps.foreach(props.remove)
+    val deleteListenerProps = new Properties()
+    deleteListenerProps ++= props.asScala.filter(entry => entry._1.startsWith(listenerPrefix(listenerName)))
+    TestUtils.incrementalAlterConfigs(servers, adminClients.head, deleteListenerProps, perBrokerConfig = true, opType = OpType.DELETE).all.get
+
+    props.clear()
     props.put(KafkaConfig.ListenersProp, listeners)
     props.put(KafkaConfig.ListenerSecurityProtocolMapProp, listenerMap)
-    TestUtils.alterConfigs(servers, adminClients.head, props, perBrokerConfig = true).all.get
+    TestUtils.incrementalAlterConfigs(servers, adminClients.head, props, perBrokerConfig = true).all.get
 
     TestUtils.waitUntilTrue(() => servers.forall(server => server.config.listeners.size == existingListenerCount - 1),
       "Listeners not updated")
+    // Wait until metrics of the listener have been removed to ensure that processors have been shutdown before
+    // verifying that connections to the removed listener fail.
+    TestUtils.waitUntilTrue(() => !servers.exists(hasListenerMetric(_, listenerName)),
+      "Processors not shutdown for removed listener")
 
     // Test that connections using deleted listener don't work
     val producerFuture = verifyConnectionFailure(producer1)
@@ -992,6 +1192,10 @@ class DynamicBrokerReconfigurationTest extends ZooKeeperTestHarness with SaslSet
     verifyProduceConsume(producer, consumer, numRecords = 10, topic)
   }
 
+  private def hasListenerMetric(server: KafkaServer, listenerName: String): Boolean = {
+    server.socketServer.metrics.metrics.keySet.asScala.exists(_.tags.get("listener") == listenerName)
+  }
+
   private def fetchBrokerConfigsFromZooKeeper(server: KafkaServer): Properties = {
     val props = adminZkClient.fetchEntityConfig(ConfigType.Broker, server.config.brokerId.toString)
     server.config.dynamicConfig.fromPersistentProps(props, perBrokerConfig = true)
@@ -1012,7 +1216,7 @@ class DynamicBrokerReconfigurationTest extends ZooKeeperTestHarness with SaslSet
     securityProps(props, props.keySet)
   }
 
-  private def createAdminClient(securityProtocol: SecurityProtocol, listenerName: String): AdminClient = {
+  private def createAdminClient(securityProtocol: SecurityProtocol, listenerName: String): Admin = {
     val config = clientProps(securityProtocol)
     val bootstrapServers = TestUtils.bootstrapServers(servers, new ListenerName(listenerName))
     config.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers)
@@ -1040,7 +1244,18 @@ class DynamicBrokerReconfigurationTest extends ZooKeeperTestHarness with SaslSet
     }
   }
 
-  private def describeConfig(adminClient: AdminClient, servers: Seq[KafkaServer] = this.servers): Config = {
+  private def waitForAuthenticationFailure(producerBuilder: ProducerBuilder): Unit = {
+    TestUtils.waitUntilTrue(() => {
+      try {
+        verifyAuthenticationFailure(producerBuilder.build())
+        true
+      } catch {
+        case e: Error => false
+      }
+    }, "Did not fail authentication with invalid config")
+  }
+
+  private def describeConfig(adminClient: Admin, servers: Seq[KafkaServer] = this.servers): Config = {
     val configResources = servers.map { server =>
       new ConfigResource(ConfigResource.Type.BROKER, server.config.brokerId.toString)
     }
@@ -1088,7 +1303,7 @@ class DynamicBrokerReconfigurationTest extends ZooKeeperTestHarness with SaslSet
     newStoreProps
   }
 
-  private def alterSslKeystore(adminClient: AdminClient, props: Properties, listener: String, expectFailure: Boolean  = false): Unit = {
+  private def alterSslKeystore(adminClient: Admin, props: Properties, listener: String, expectFailure: Boolean  = false): Unit = {
     val configPrefix = listenerPrefix(listener)
     val newProps = securityProps(props, KEYSTORE_PROPS, configPrefix)
     reconfigureServers(newProps, perBrokerConfig = true,
@@ -1098,36 +1313,18 @@ class DynamicBrokerReconfigurationTest extends ZooKeeperTestHarness with SaslSet
   private def alterSslKeystoreUsingConfigCommand(props: Properties, listener: String): Unit = {
     val configPrefix = listenerPrefix(listener)
     val newProps = securityProps(props, KEYSTORE_PROPS, configPrefix)
-
-    val propsFile = TestUtils.tempFile()
-    val propsWriter = new FileWriter(propsFile)
-    try {
-      clientProps(SecurityProtocol.SSL).asScala.foreach {
-        case (k, v) => propsWriter.write(s"$k=$v\n")
-      }
-    } finally {
-      propsWriter.close()
-    }
-
-    servers.foreach { server =>
-      val args = Array("--bootstrap-server", TestUtils.bootstrapServers(servers, new ListenerName(SecureInternal)),
-        "--command-config", propsFile.getAbsolutePath,
-        "--alter", "--add-config", newProps.asScala.map { case (k, v) => s"$k=$v" }.mkString(","),
-        "--entity-type", "brokers",
-        "--entity-name", server.config.brokerId.toString)
-      ConfigCommand.main(args)
-    }
+    alterConfigsUsingConfigCommand(newProps)
     waitForConfig(s"$configPrefix$SSL_KEYSTORE_LOCATION_CONFIG", props.getProperty(SSL_KEYSTORE_LOCATION_CONFIG))
   }
 
-  private def serverEndpoints(adminClient: AdminClient): String = {
+  private def serverEndpoints(adminClient: Admin): String = {
     val nodes = adminClient.describeCluster().nodes().get
     nodes.asScala.map { node =>
       s"${node.host}:${node.port}"
     }.mkString(",")
   }
 
-  private def alterAdvertisedListener(adminClient: AdminClient, externalAdminClient: AdminClient, oldHost: String, newHost: String): Unit = {
+  private def alterAdvertisedListener(adminClient: Admin, externalAdminClient: Admin, oldHost: String, newHost: String): Unit = {
     val configs = servers.map { server =>
       val resource = new ConfigResource(ConfigResource.Type.BROKER, server.config.brokerId.toString)
       val newListeners = server.config.advertisedListeners.map { e =>
@@ -1164,7 +1361,7 @@ class DynamicBrokerReconfigurationTest extends ZooKeeperTestHarness with SaslSet
   private def reconfigureServers(newProps: Properties, perBrokerConfig: Boolean, aPropToVerify: (String, String), expectFailure: Boolean = false): Unit = {
     val alterResult = TestUtils.alterConfigs(servers, adminClients.head, newProps, perBrokerConfig)
     if (expectFailure) {
-      val oldProps = servers.head.config.values.asScala.filterKeys(newProps.containsKey)
+      val oldProps = servers.head.config.values.asScala.filter { case (k, _) => newProps.containsKey(k) }
       val brokerResources = if (perBrokerConfig)
         servers.map(server => new ConfigResource(ConfigResource.Type.BROKER, server.config.brokerId.toString))
       else
@@ -1173,7 +1370,9 @@ class DynamicBrokerReconfigurationTest extends ZooKeeperTestHarness with SaslSet
         val exception = intercept[ExecutionException](alterResult.values.get(brokerResource).get)
         assertTrue(exception.getCause.isInstanceOf[InvalidRequestException])
       }
-      servers.foreach { server => assertEquals(oldProps, server.config.values.asScala.filterKeys(newProps.containsKey)) }
+      servers.foreach { server =>
+        assertEquals(oldProps, server.config.values.asScala.filter { case (k, _) => newProps.containsKey(k) })
+      }
     } else {
       alterResult.all.get
       waitForConfig(aPropToVerify._1, aPropToVerify._2)
@@ -1233,7 +1432,7 @@ class DynamicBrokerReconfigurationTest extends ZooKeeperTestHarness with SaslSet
   }
 
   private def configureMetricsReporters(reporters: Seq[Class[_]], props: Properties,
-                                       perBrokerConfig: Boolean = false): Unit = {
+                                        perBrokerConfig: Boolean = false): Unit = {
     val reporterStr = reporters.map(_.getName).mkString(",")
     props.put(KafkaConfig.MetricReporterClassesProp, reporterStr)
     reconfigureServers(props, perBrokerConfig, (KafkaConfig.MetricReporterClassesProp, reporterStr))
@@ -1292,7 +1491,7 @@ class DynamicBrokerReconfigurationTest extends ZooKeeperTestHarness with SaslSet
     val executor = Executors.newSingleThreadExecutor
     executors += executor
     val future = executor.submit(new Runnable() {
-      def run() {
+      def run(): Unit = {
         producer.send(new ProducerRecord(topic, "key", "value")).get
       }
     })
@@ -1304,8 +1503,8 @@ class DynamicBrokerReconfigurationTest extends ZooKeeperTestHarness with SaslSet
     val executor = Executors.newSingleThreadExecutor
     executors += executor
     val future = executor.submit(new Runnable() {
-      def run() {
-        assertEquals(0, consumer.poll(100).count)
+      def run(): Unit = {
+        consumer.commitSync()
       }
     })
     verifyTimeout(future)
@@ -1346,6 +1545,27 @@ class DynamicBrokerReconfigurationTest extends ZooKeeperTestHarness with SaslSet
     }
   }
 
+  private def alterConfigsUsingConfigCommand(props: Properties): Unit = {
+    val propsFile = TestUtils.tempFile()
+    val propsWriter = new FileWriter(propsFile)
+    try {
+      clientProps(SecurityProtocol.SSL).asScala.foreach {
+        case (k, v) => propsWriter.write(s"$k=$v\n")
+      }
+    } finally {
+      propsWriter.close()
+    }
+
+    servers.foreach { server =>
+      val args = Array("--bootstrap-server", TestUtils.bootstrapServers(servers, new ListenerName(SecureInternal)),
+        "--command-config", propsFile.getAbsolutePath,
+        "--alter", "--add-config", props.asScala.map { case (k, v) => s"$k=$v" }.mkString(","),
+        "--entity-type", "brokers",
+        "--entity-name", server.config.brokerId.toString)
+      ConfigCommand.main(args)
+    }
+  }
+
   private abstract class ClientBuilder[T]() {
     protected var _bootstrapServers: Option[String] = None
     protected var _listenerName = SecureExternal
@@ -1460,7 +1680,7 @@ class DynamicBrokerReconfigurationTest extends ZooKeeperTestHarness with SaslSet
     override def doWork(): Unit = {
       try {
         while (isRunning || (lastReceived != producerThread.lastSent && System.currentTimeMillis < endTimeMs)) {
-          val records = consumer.poll(50)
+          val records = consumer.poll(Duration.ofMillis(50L))
           received += records.count
           if (!records.isEmpty) {
             lastBatch = records
@@ -1558,8 +1778,8 @@ class TestMetricsReporter extends MetricsReporter with Reconfigurable with Close
   }
 
   override def validateReconfiguration(configs: util.Map[String, _]): Unit = {
-    val pollingInterval = configs.get(PollingIntervalProp).toString
-    if (configs.get(PollingIntervalProp).toString.toInt <= 0)
+    val pollingInterval = configs.get(PollingIntervalProp).toString.toInt
+    if (pollingInterval <= 0)
       throw new ConfigException(s"Invalid polling interval $pollingInterval")
   }
 
@@ -1575,7 +1795,7 @@ class TestMetricsReporter extends MetricsReporter with Reconfigurable with Close
   def verifyState(reconfigureCount: Int, deleteCount: Int, pollingInterval: Int): Unit = {
     assertEquals(1, initializeCount)
     assertEquals(1, configureCount)
-    assertEquals(reconfigureCount, reconfigureCount)
+    assertEquals(reconfigureCount, this.reconfigureCount)
     assertEquals(deleteCount, closeCount)
     assertEquals(1, clusterUpdateCount)
     assertEquals(pollingInterval, this.pollingInterval)
@@ -1588,3 +1808,11 @@ class TestMetricsReporter extends MetricsReporter with Reconfigurable with Close
     assertTrue("Invalid metric value", total > 0.0)
   }
 }
+
+
+class MockFileConfigProvider extends FileConfigProvider {
+  @throws(classOf[IOException])
+  override def reader(path: String): Reader = {
+    new StringReader("key=testKey\npassword=ServerPassword\ninterval=1000\nupdinterval=2000\nstoretype=JKS");
+  }
+}
diff --git a/core/src/test/scala/integration/kafka/server/GssapiAuthenticationTest.scala b/core/src/test/scala/integration/kafka/server/GssapiAuthenticationTest.scala
index 07885cdbb0062..f5fed6c608c44 100644
--- a/core/src/test/scala/integration/kafka/server/GssapiAuthenticationTest.scala
+++ b/core/src/test/scala/integration/kafka/server/GssapiAuthenticationTest.scala
@@ -39,7 +39,7 @@ import org.junit.{After, Before, Test}
 import scala.collection.JavaConverters._
 
 class GssapiAuthenticationTest extends IntegrationTestHarness with SaslSetup {
-  override val serverCount = 1
+  override val brokerCount = 1
   override protected def securityProtocol = SecurityProtocol.SASL_PLAINTEXT
 
   private val kafkaClientSaslMechanism = "GSSAPI"
@@ -56,7 +56,7 @@ class GssapiAuthenticationTest extends IntegrationTestHarness with SaslSetup {
   private val failedAuthenticationDelayMs = 2000
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     startSasl(jaasSections(kafkaServerSaslMechanisms, Option(kafkaClientSaslMechanism), Both))
     serverConfig.put(KafkaConfig.SslClientAuthProp, "required")
     serverConfig.put(KafkaConfig.FailedAuthenticationDelayMsProp, failedAuthenticationDelayMs.toString)
@@ -70,7 +70,7 @@ class GssapiAuthenticationTest extends IntegrationTestHarness with SaslSetup {
     clientConfig.put(CommonClientConfigs.CONNECTIONS_MAX_IDLE_MS_CONFIG, "5000")
 
     // create the test topic with all the brokers as replicas
-    createTopic(topic, 2, serverCount)
+    createTopic(topic, 2, brokerCount)
   }
 
   @After
diff --git a/core/src/test/scala/integration/kafka/server/MultipleListenersWithAdditionalJaasContextTest.scala b/core/src/test/scala/integration/kafka/server/MultipleListenersWithAdditionalJaasContextTest.scala
index f2c6753786e00..39aa3c3639ca8 100644
--- a/core/src/test/scala/integration/kafka/server/MultipleListenersWithAdditionalJaasContextTest.scala
+++ b/core/src/test/scala/integration/kafka/server/MultipleListenersWithAdditionalJaasContextTest.scala
@@ -19,6 +19,8 @@ package kafka.server
 
 import java.util.Properties
 
+import scala.collection.Seq
+
 import kafka.utils.JaasTestUtils
 import kafka.utils.JaasTestUtils.JaasSection
 
diff --git a/core/src/test/scala/integration/kafka/server/MultipleListenersWithDefaultJaasContextTest.scala b/core/src/test/scala/integration/kafka/server/MultipleListenersWithDefaultJaasContextTest.scala
index 00695ed2bf6ff..23746c495579b 100644
--- a/core/src/test/scala/integration/kafka/server/MultipleListenersWithDefaultJaasContextTest.scala
+++ b/core/src/test/scala/integration/kafka/server/MultipleListenersWithDefaultJaasContextTest.scala
@@ -19,6 +19,8 @@ package kafka.server
 
 import java.util.Properties
 
+import scala.collection.Seq
+
 import kafka.api.Both
 import kafka.utils.JaasTestUtils.JaasSection
 
diff --git a/core/src/test/scala/integration/kafka/server/MultipleListenersWithSameSecurityProtocolBaseTest.scala b/core/src/test/scala/integration/kafka/server/MultipleListenersWithSameSecurityProtocolBaseTest.scala
index 95a4843eb650e..6e811886ac40a 100644
--- a/core/src/test/scala/integration/kafka/server/MultipleListenersWithSameSecurityProtocolBaseTest.scala
+++ b/core/src/test/scala/integration/kafka/server/MultipleListenersWithSameSecurityProtocolBaseTest.scala
@@ -38,6 +38,7 @@ import org.junit.{After, Before, Test}
 
 import scala.collection.mutable
 import scala.collection.mutable.ArrayBuffer
+import scala.collection.Seq
 
 object MultipleListenersWithSameSecurityProtocolBaseTest {
   val SecureInternal = "SECURE_INTERNAL"
@@ -147,7 +148,7 @@ abstract class MultipleListenersWithSameSecurityProtocolBaseTest extends ZooKeep
   }
 
   @After
-  override def tearDown() {
+  override def tearDown(): Unit = {
     producers.values.foreach(_.close())
     consumers.values.foreach(_.close())
     TestUtils.shutdownServers(servers)
diff --git a/core/src/test/scala/integration/kafka/server/ScramServerStartupTest.scala b/core/src/test/scala/integration/kafka/server/ScramServerStartupTest.scala
index 0f4065080ce76..70c30ee7dcbf4 100644
--- a/core/src/test/scala/integration/kafka/server/ScramServerStartupTest.scala
+++ b/core/src/test/scala/integration/kafka/server/ScramServerStartupTest.scala
@@ -36,7 +36,7 @@ import scala.collection.JavaConverters._
  */
 class ScramServerStartupTest extends IntegrationTestHarness with SaslSetup {
 
-  override val serverCount = 1
+  override val brokerCount = 1
 
   private val kafkaClientSaslMechanism = "SCRAM-SHA-256"
   private val kafkaServerSaslMechanisms = Collections.singletonList("SCRAM-SHA-256").asScala
@@ -46,7 +46,7 @@ class ScramServerStartupTest extends IntegrationTestHarness with SaslSetup {
   override protected val serverSaslProperties = Some(kafkaServerSaslProperties(kafkaServerSaslMechanisms, kafkaClientSaslMechanism))
   override protected val clientSaslProperties = Some(kafkaClientSaslProperties(kafkaClientSaslMechanism))
 
-  override def configureSecurityBeforeServersStart() {
+  override def configureSecurityBeforeServersStart(): Unit = {
     super.configureSecurityBeforeServersStart()
     zkClient.makeSurePersistentPathExists(ConfigEntityChangeNotificationZNode.path)
     // Create credentials before starting brokers
diff --git a/core/src/test/scala/integration/kafka/tools/MirrorMakerIntegrationTest.scala b/core/src/test/scala/integration/kafka/tools/MirrorMakerIntegrationTest.scala
index 0329f5275d706..a5a6caffda2e2 100644
--- a/core/src/test/scala/integration/kafka/tools/MirrorMakerIntegrationTest.scala
+++ b/core/src/test/scala/integration/kafka/tools/MirrorMakerIntegrationTest.scala
@@ -18,6 +18,8 @@ package kafka.tools
 
 import java.util.Properties
 
+import scala.collection.Seq
+
 import kafka.integration.KafkaServerTestHarness
 import kafka.server.KafkaConfig
 import kafka.tools.MirrorMaker.{ConsumerWrapper, MirrorMakerProducer, NoRecordsException}
@@ -73,6 +75,7 @@ class MirrorMakerIntegrationTest extends KafkaServerTestHarness {
     producerProps.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, brokerList)
     producerProps.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, classOf[ByteArraySerializer])
     producerProps.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, classOf[ByteArraySerializer])
+    producerProps.put(ProducerConfig.ALLOW_AUTO_CREATE_TOPICS_CONFIG, boolean2Boolean(true))
     val producer = new MirrorMakerProducer(true, producerProps)
     MirrorMaker.producer = producer
     MirrorMaker.producer.send(new ProducerRecord(topic, msg.getBytes()))
diff --git a/core/src/test/scala/kafka/security/minikdc/MiniKdc.scala b/core/src/test/scala/kafka/security/minikdc/MiniKdc.scala
index 9b3058122aaa2..256091ced732a 100644
--- a/core/src/test/scala/kafka/security/minikdc/MiniKdc.scala
+++ b/core/src/test/scala/kafka/security/minikdc/MiniKdc.scala
@@ -113,7 +113,7 @@ class MiniKdc(config: Properties, workDir: File) extends Logging {
 
   def host: String = config.getProperty(MiniKdc.KdcBindAddress)
 
-  def start() {
+  def start(): Unit = {
     if (kdc != null)
       throw new RuntimeException("KDC already started")
     if (closed)
@@ -123,7 +123,7 @@ class MiniKdc(config: Properties, workDir: File) extends Logging {
     initJvmKerberosConfig()
   }
 
-  private def initDirectoryService() {
+  private def initDirectoryService(): Unit = {
     ds = new DefaultDirectoryService
     ds.setInstanceLayout(new InstanceLayout(workDir))
     ds.setCacheService(new CacheService)
@@ -187,9 +187,9 @@ class MiniKdc(config: Properties, workDir: File) extends Logging {
     ds.getAdminSession.add(entry)
   }
 
-  private def initKdcServer() {
+  private def initKdcServer(): Unit = {
 
-    def addInitialEntriesToDirectoryService(bindAddress: String) {
+    def addInitialEntriesToDirectoryService(bindAddress: String): Unit = {
       val map = Map (
         "0" -> orgName.toLowerCase(Locale.ENGLISH),
         "1" -> orgDomain.toLowerCase(Locale.ENGLISH),
@@ -245,7 +245,7 @@ class MiniKdc(config: Properties, workDir: File) extends Logging {
     refreshJvmKerberosConfig()
   }
 
-  private def writeKrb5Conf() {
+  private def writeKrb5Conf(): Unit = {
     val stringBuilder = new StringBuilder
     val reader = new BufferedReader(
       new InputStreamReader(MiniKdc.getResourceAsStream("minikdc-krb5.conf"), StandardCharsets.UTF_8))
@@ -268,7 +268,7 @@ class MiniKdc(config: Properties, workDir: File) extends Logging {
     klass.getMethod("refresh").invoke(klass)
   }
 
-  def stop() {
+  def stop(): Unit = {
     if (!closed) {
       closed = true
       if (kdc != null) {
@@ -291,7 +291,7 @@ class MiniKdc(config: Properties, workDir: File) extends Logging {
     * @param principal principal name, do not include the domain.
     * @param password  password.
     */
-  private def createPrincipal(principal: String, password: String) {
+  private def createPrincipal(principal: String, password: String): Unit = {
     val ldifContent = s"""
       |dn: uid=$principal,ou=users,dc=${orgName.toLowerCase(Locale.ENGLISH)},dc=${orgDomain.toLowerCase(Locale.ENGLISH)}
       |objectClass: top
@@ -316,7 +316,7 @@ class MiniKdc(config: Properties, workDir: File) extends Logging {
     * @param keytabFile keytab file to add the created principals
     * @param principals principals to add to the KDC, do not include the domain.
     */
-  def createPrincipal(keytabFile: File, principals: String*) {
+  def createPrincipal(keytabFile: File, principals: String*): Unit = {
     val generatedPassword = UUID.randomUUID.toString
     val keytab = new Keytab
     val entries = principals.flatMap { principal =>
@@ -347,7 +347,7 @@ object MiniKdc {
   val JavaSecurityKrb5Conf = "java.security.krb5.conf"
   val SunSecurityKrb5Debug = "sun.security.krb5.debug"
 
-  def main(args: Array[String]) {
+  def main(args: Array[String]): Unit = {
     args match {
       case Array(workDirPath, configPath, keytabPath, principals@ _*) if principals.nonEmpty =>
         val workDir = new File(workDirPath)
@@ -370,7 +370,7 @@ object MiniKdc {
     }
   }
 
-  private def start(workDir: File, config: Properties, keytabFile: File, principals: Seq[String]) {
+  private def start(workDir: File, config: Properties, keytabFile: File, principals: Seq[String]): Unit = {
     val miniKdc = new MiniKdc(config, workDir)
     miniKdc.start()
     miniKdc.createPrincipal(keytabFile, principals: _*)
diff --git a/core/src/test/scala/kafka/tools/CustomDeserializerTest.scala b/core/src/test/scala/kafka/tools/CustomDeserializerTest.scala
index 7fb3cf30cc9d1..d6b0100bf303d 100644
--- a/core/src/test/scala/kafka/tools/CustomDeserializerTest.scala
+++ b/core/src/test/scala/kafka/tools/CustomDeserializerTest.scala
@@ -24,7 +24,7 @@ import org.apache.kafka.common.serialization.Deserializer
 import org.hamcrest.CoreMatchers
 import org.hamcrest.MatcherAssert._
 import org.junit.Test
-import org.scalatest.mockito.MockitoSugar
+import org.mockito.Mockito._
 
 class CustomDeserializer extends Deserializer[String] {
 
@@ -34,14 +34,15 @@ class CustomDeserializer extends Deserializer[String] {
   }
 }
 
-class CustomDeserializerTest extends MockitoSugar {
+class CustomDeserializerTest {
 
   @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[PrintStream])
+    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/LogCompactionTester.scala b/core/src/test/scala/kafka/tools/LogCompactionTester.scala
index 3690856cd90ab..9edac52bffcfc 100755
--- a/core/src/test/scala/kafka/tools/LogCompactionTester.scala
+++ b/core/src/test/scala/kafka/tools/LogCompactionTester.scala
@@ -58,7 +58,7 @@ object LogCompactionTester {
   //maximum line size while reading produced/consumed record text file
   private val ReadAheadLimit = 4906
 
-  def main(args: Array[String]) {
+  def main(args: Array[String]): Unit = {
     val parser = new OptionParser(false)
     val numMessagesOpt = parser.accepts("messages", "The number of messages to send or consume.")
       .withRequiredArg
@@ -142,7 +142,7 @@ object LogCompactionTester {
 
     try {
       val topicConfigs = Map(TopicConfig.CLEANUP_POLICY_CONFIG -> TopicConfig.CLEANUP_POLICY_COMPACT)
-      val newTopics = topics.map(name => new NewTopic(name, 1, 1).configs(topicConfigs.asJava)).asJava
+      val newTopics = topics.map(name => new NewTopic(name, 1, 1.toShort).configs(topicConfigs.asJava)).asJava
       adminClient.createTopics(newTopics).all.get
 
       var pendingTopics: Seq[String] = Seq()
@@ -157,7 +157,7 @@ object LogCompactionTester {
 
   def lineCount(filPath: Path): Int = Files.readAllLines(filPath).size
 
-  def validateOutput(producedDataFile: File, consumedDataFile: File) {
+  def validateOutput(producedDataFile: File, consumedDataFile: File): Unit = {
     val producedReader = externalSort(producedDataFile)
     val consumedReader = externalSort(consumedDataFile)
     val produced = valuesIterator(producedReader)
@@ -192,7 +192,7 @@ object LogCompactionTester {
     Utils.delete(consumedDedupedFile)
   }
 
-  def require(requirement: Boolean, message: => Any) {
+  def require(requirement: Boolean, message: => Any): Unit = {
     if (!requirement) {
       System.err.println(s"Data validation failed : $message")
       Exit.exit(1)
@@ -242,7 +242,7 @@ object LogCompactionTester {
     val builder = new ProcessBuilder("sort", "--key=1,2", "--stable", "--buffer-size=20%", "--temporary-directory=" + Files.createTempDirectory("log_compaction_test"), file.getAbsolutePath)
     val process = builder.start
     new Thread() {
-      override def run() {
+      override def run(): Unit = {
         val exitCode = process.waitFor()
         if (exitCode != 0) {
           System.err.println("Process exited abnormally.")
diff --git a/core/src/test/scala/other/kafka/ReplicationQuotasTestRig.scala b/core/src/test/scala/other/kafka/ReplicationQuotasTestRig.scala
index 20c28e76a2015..5f31818da182b 100644
--- a/core/src/test/scala/other/kafka/ReplicationQuotasTestRig.scala
+++ b/core/src/test/scala/other/kafka/ReplicationQuotasTestRig.scala
@@ -18,6 +18,7 @@
 package kafka
 
 import java.io.{File, PrintWriter}
+import java.nio.charset.StandardCharsets
 import java.nio.file.{Files, StandardOpenOption}
 import javax.imageio.ImageIO
 
@@ -25,7 +26,7 @@ import kafka.admin.ReassignPartitionsCommand
 import kafka.admin.ReassignPartitionsCommand.Throttle
 import kafka.server.{KafkaConfig, KafkaServer, QuotaType}
 import kafka.utils.TestUtils._
-import kafka.utils.{Exit, Logging, TestUtils, ZkUtils}
+import kafka.utils.{Exit, Logging, TestUtils}
 import kafka.zk.{ReassignPartitionsZNode, ZooKeeperTestHarness}
 import org.apache.kafka.clients.producer.ProducerRecord
 import org.apache.kafka.common.TopicPartition
@@ -73,7 +74,7 @@ object ReplicationQuotasTestRig {
       Exit.exit(0)
   }
 
-  def run(config: ExperimentDef, journal: Journal, displayChartsOnScreen: Boolean) {
+  def run(config: ExperimentDef, journal: Journal, displayChartsOnScreen: Boolean): Unit = {
     val experiment = new Experiment()
     try {
       experiment.setUp
@@ -100,18 +101,18 @@ object ReplicationQuotasTestRig {
     val leaderRates = mutable.Map[Int, Array[Double]]()
     val followerRates = mutable.Map[Int, Array[Double]]()
 
-    def startBrokers(brokerIds: Seq[Int]) {
+    def startBrokers(brokerIds: Seq[Int]): Unit = {
       println("Starting Brokers")
       servers = brokerIds.map(i => createBrokerConfig(i, zkConnect))
         .map(c => createServer(KafkaConfig.fromProps(c)))
     }
 
-    override def tearDown() {
+    override def tearDown(): Unit = {
       TestUtils.shutdownServers(servers)
       super.tearDown()
     }
 
-    def run(config: ExperimentDef, journal: Journal, displayChartsOnScreen: Boolean) {
+    def run(config: ExperimentDef, journal: Journal, displayChartsOnScreen: Boolean): Unit = {
       experimentName = config.name
       val brokers = (100 to 100 + config.brokers)
       var count = 0
@@ -138,7 +139,8 @@ object ReplicationQuotasTestRig {
       val newAssignment = ReassignPartitionsCommand.generateAssignment(zkClient, brokers, json(topicName), true)._1
 
       val start = System.currentTimeMillis()
-      ReassignPartitionsCommand.executeAssignment(zkClient, None, ZkUtils.getReassignmentJson(newAssignment), Throttle(config.throttle))
+      ReassignPartitionsCommand.executeAssignment(zkClient, None,
+        new String(ReassignPartitionsZNode.encode(newAssignment), StandardCharsets.UTF_8), Throttle(config.throttle))
 
       //Await completion
       waitForReassignmentToComplete()
@@ -171,9 +173,9 @@ object ReplicationQuotasTestRig {
 
       //Long stats
       println("The replicas are " + replicas.toSeq.sortBy(_._1).map("\n" + _))
-      println("This is the current replica assignment:\n" + actual.toSeq)
+      println("This is the current replica assignment:\n" + actual.mapValues(_.replicas).toMap.toSeq)
       println("proposed assignment is: \n" + newAssignment)
-      println("This is the assignment we ended up with" + actual)
+      println("This is the assignment we ended up with" + actual.mapValues(_.replicas).toMap)
 
       //Test Stats
       println(s"numBrokers: ${config.brokers}")
@@ -185,7 +187,7 @@ object ReplicationQuotasTestRig {
       println(s"Worst case duration is ${config.targetBytesPerBrokerMB * 1000 * 1000/ config.throttle}")
     }
 
-    def waitForReassignmentToComplete() {
+    def waitForReassignmentToComplete(): Unit = {
       waitUntilTrue(() => {
         printRateMetrics()
         !zkClient.reassignPartitionsInProgress()
@@ -246,7 +248,7 @@ object ReplicationQuotasTestRig {
       rates.put(brokerId, leaderRatesBroker)
     }
 
-    def printRateMetrics() {
+    def printRateMetrics(): Unit = {
       for (broker <- servers) {
         val leaderRate: Double = measuredRate(broker, QuotaType.LeaderReplication)
         if (broker.config.brokerId == 100)
diff --git a/core/src/test/scala/other/kafka/StressTestLog.scala b/core/src/test/scala/other/kafka/StressTestLog.scala
index 7821a0e007d3c..8a5eefb859aad 100755
--- a/core/src/test/scala/other/kafka/StressTestLog.scala
+++ b/core/src/test/scala/other/kafka/StressTestLog.scala
@@ -21,7 +21,7 @@ import java.util.Properties
 import java.util.concurrent.atomic._
 
 import kafka.log._
-import kafka.server.{BrokerTopicStats, LogDirFailureChannel}
+import kafka.server.{BrokerTopicStats, FetchLogEnd, LogDirFailureChannel}
 import kafka.utils._
 import org.apache.kafka.clients.consumer.OffsetOutOfRangeException
 import org.apache.kafka.common.record.FileRecords
@@ -34,7 +34,7 @@ import org.apache.kafka.common.utils.Utils
 object StressTestLog {
   val running = new AtomicBoolean(true)
 
-  def main(args: Array[String]) {
+  def main(args: Array[String]): Unit = {
     val dir = TestUtils.randomPartitionLogDir(TestUtils.tempDir())
     val time = new MockTime
     val logProperties = new Properties()
@@ -77,7 +77,7 @@ object StressTestLog {
   abstract class WorkerThread extends Thread {
     val threadInfo = "Thread: " + Thread.currentThread.getName + " Class: " + getClass.getName
 
-    override def run() {
+    override def run(): Unit = {
       try {
         while(running.get)
           work()
@@ -90,7 +90,7 @@ object StressTestLog {
       }
     }
 
-    def work()
+    def work(): Unit
     def isMakingProgress(): Boolean
   }
 
@@ -108,7 +108,7 @@ object StressTestLog {
       false
     }
 
-    def checkProgress() {
+    def checkProgress(): Unit = {
       // Check if we are making progress every 500ms
       val curTime = System.currentTimeMillis
       if ((curTime - lastProgressCheckTime) > 500) {
@@ -119,7 +119,7 @@ object StressTestLog {
   }
 
   class WriterThread(val log: Log) extends WorkerThread with LogProgress {
-    override def work() {
+    override def work(): Unit = {
       val logAppendInfo = log.appendAsLeader(TestUtils.singletonRecords(currentOffset.toString.getBytes), 0)
       require(logAppendInfo.firstOffset.forall(_ == currentOffset) && logAppendInfo.lastOffset == currentOffset)
       currentOffset += 1
@@ -129,13 +129,12 @@ object StressTestLog {
   }
 
   class ReaderThread(val log: Log) extends WorkerThread with LogProgress {
-    override def work() {
+    override def work(): Unit = {
       try {
         log.read(currentOffset,
-          maxLength = 1024,
-          maxOffset = Some(currentOffset + 1),
-          minOneMessage = true,
-          includeAbortedTxns = false).records match {
+          maxLength = 1,
+          isolation = FetchLogEnd,
+          minOneMessage = true).records match {
           case read: FileRecords if read.sizeInBytes > 0 => {
             val first = read.batches.iterator.next()
             require(first.lastOffset == currentOffset, "We should either read nothing or the message we asked for.")
diff --git a/core/src/test/scala/other/kafka/TestLinearWriteSpeed.scala b/core/src/test/scala/other/kafka/TestLinearWriteSpeed.scala
index 3c8dfa2b13e78..bf21deeae357e 100755
--- a/core/src/test/scala/other/kafka/TestLinearWriteSpeed.scala
+++ b/core/src/test/scala/other/kafka/TestLinearWriteSpeed.scala
@@ -174,7 +174,7 @@ object TestLinearWriteSpeed {
 
   trait Writable {
     def write(): Int
-    def close()
+    def close(): Unit
   }
 
   class MmapWritable(val file: File, size: Long, val content: ByteBuffer) extends Writable {
@@ -187,7 +187,7 @@ object TestLinearWriteSpeed {
       content.rewind()
       content.limit()
     }
-    def close() {
+    def close(): Unit = {
       raf.close()
       Utils.delete(file)
     }
@@ -202,7 +202,7 @@ object TestLinearWriteSpeed {
       content.rewind()
       content.limit()
     }
-    def close() {
+    def close(): Unit = {
       channel.close()
       Utils.delete(file)
     }
@@ -216,7 +216,7 @@ object TestLinearWriteSpeed {
       log.appendAsLeader(messages, leaderEpoch = 0)
       messages.sizeInBytes
     }
-    def close() {
+    def close(): Unit = {
       log.close()
       Utils.delete(log.dir)
     }
diff --git a/core/src/test/scala/unit/kafka/admin/AclCommandTest.scala b/core/src/test/scala/unit/kafka/admin/AclCommandTest.scala
index 7d7b8e492566c..a2b68a6b57f16 100644
--- a/core/src/test/scala/unit/kafka/admin/AclCommandTest.scala
+++ b/core/src/test/scala/unit/kafka/admin/AclCommandTest.scala
@@ -19,17 +19,23 @@ package kafka.admin
 import java.util.Properties
 
 import kafka.admin.AclCommand.AclCommandOptions
-import kafka.security.auth._
+import kafka.security.auth.Authorizer
+import kafka.security.authorizer.{AclAuthorizer, AuthorizerUtils}
 import kafka.server.{KafkaConfig, KafkaServer}
 import kafka.utils.{Exit, Logging, TestUtils}
 import kafka.zk.ZooKeeperTestHarness
-import org.apache.kafka.common.resource.PatternType
+import org.apache.kafka.common.acl.{AccessControlEntry, AclOperation, AclPermissionType}
+import org.apache.kafka.common.acl.AclOperation._
+import org.apache.kafka.common.acl.AclPermissionType._
+import org.apache.kafka.common.resource.{PatternType, Resource, ResourcePattern}
+import org.apache.kafka.common.resource.ResourceType._
 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.SecurityUtils
+import org.apache.kafka.server.authorizer.{Authorizer => JAuthorizer}
 import org.junit.{After, Before, Test}
+import org.scalatest.Assertions.intercept
 
 class AclCommandTest extends ZooKeeperTestHarness with Logging {
 
@@ -42,51 +48,52 @@ class AclCommandTest extends ZooKeeperTestHarness with Logging {
   private val AllowHostCommand = Array("--allow-host", "host1", "--allow-host", "host2")
   private val DenyHostCommand = Array("--deny-host", "host1", "--deny-host", "host2")
 
-  private val TopicResources = Set(Resource(Topic, "test-1", LITERAL), Resource(Topic, "test-2", LITERAL))
-  private val GroupResources = Set(Resource(Group, "testGroup-1", LITERAL), Resource(Group, "testGroup-2", LITERAL))
-  private val TransactionalIdResources = Set(Resource(TransactionalId, "t0", LITERAL), Resource(TransactionalId, "t1", LITERAL))
-  private val TokenResources = Set(Resource(DelegationToken, "token1", LITERAL), Resource(DelegationToken, "token2", LITERAL))
+  private val ClusterResource = new ResourcePattern(CLUSTER, Resource.CLUSTER_NAME, LITERAL)
+  private val TopicResources = Set(new ResourcePattern(TOPIC, "test-1", LITERAL), new ResourcePattern(TOPIC, "test-2", LITERAL))
+  private val GroupResources = Set(new ResourcePattern(GROUP, "testGroup-1", LITERAL), new ResourcePattern(GROUP, "testGroup-2", LITERAL))
+  private val TransactionalIdResources = Set(new ResourcePattern(TRANSACTIONAL_ID, "t0", LITERAL), new ResourcePattern(TRANSACTIONAL_ID, "t1", LITERAL))
+  private val TokenResources = Set(new ResourcePattern(DELEGATION_TOKEN, "token1", LITERAL), new ResourcePattern(DELEGATION_TOKEN, "token2", LITERAL))
 
-  private val ResourceToCommand = Map[Set[Resource], Array[String]](
+  private val ResourceToCommand = Map[Set[ResourcePattern], Array[String]](
     TopicResources -> Array("--topic", "test-1", "--topic", "test-2"),
-    Set(Resource.ClusterResource) -> Array("--cluster"),
+    Set(ClusterResource) -> Array("--cluster"),
     GroupResources -> Array("--group", "testGroup-1", "--group", "testGroup-2"),
     TransactionalIdResources -> Array("--transactional-id", "t0", "--transactional-id", "t1"),
     TokenResources -> Array("--delegation-token", "token1", "--delegation-token", "token2")
   )
 
-  private val ResourceToOperations = Map[Set[Resource], (Set[Operation], Array[String])](
-    TopicResources -> (Set(Read, Write, Create, Describe, Delete, DescribeConfigs, AlterConfigs, Alter),
-      Array("--operation", "Read" , "--operation", "Write", "--operation", "Create", "--operation", "Describe", "--operation", "Delete",
+  private val ResourceToOperations = Map[Set[ResourcePattern], (Set[AclOperation], Array[String])](
+    TopicResources -> (Set(READ, WRITE, CREATE, DESCRIBE, DELETE, DESCRIBE_CONFIGS, ALTER_CONFIGS, ALTER),
+      Array("--operation", "Read" , "--operation", "WRITE", "--operation", "Create", "--operation", "Describe", "--operation", "Delete",
         "--operation", "DescribeConfigs", "--operation", "AlterConfigs", "--operation", "Alter")),
-    Set(Resource.ClusterResource) -> (Set(Create, ClusterAction, DescribeConfigs, AlterConfigs, IdempotentWrite, Alter, Describe),
+    Set(ClusterResource) -> (Set(CREATE, CLUSTER_ACTION, DESCRIBE_CONFIGS, ALTER_CONFIGS, IDEMPOTENT_WRITE, ALTER, DESCRIBE),
       Array("--operation", "Create", "--operation", "ClusterAction", "--operation", "DescribeConfigs",
         "--operation", "AlterConfigs", "--operation", "IdempotentWrite", "--operation", "Alter", "--operation", "Describe")),
-    GroupResources -> (Set(Read, Describe, Delete), Array("--operation", "Read", "--operation", "Describe", "--operation", "Delete")),
-    TransactionalIdResources -> (Set(Describe, Write), Array("--operation", "Describe", "--operation", "Write")),
-    TokenResources -> (Set(Describe), Array("--operation", "Describe"))
+    GroupResources -> (Set(READ, DESCRIBE, DELETE), Array("--operation", "Read", "--operation", "Describe", "--operation", "Delete")),
+    TransactionalIdResources -> (Set(DESCRIBE, WRITE), Array("--operation", "Describe", "--operation", "WRITE")),
+    TokenResources -> (Set(DESCRIBE), Array("--operation", "Describe"))
   )
 
-  private def ProducerResourceToAcls(enableIdempotence: Boolean = false) = Map[Set[Resource], Set[Acl]](
-    TopicResources -> AclCommand.getAcls(Users, Allow, Set(Write, Describe, Create), Hosts),
-    TransactionalIdResources -> AclCommand.getAcls(Users, Allow, Set(Write, Describe), Hosts),
-    Set(Resource.ClusterResource) -> AclCommand.getAcls(Users, Allow,
-      Set(if (enableIdempotence) Some(IdempotentWrite) else None).flatten, Hosts)
+  private def ProducerResourceToAcls(enableIdempotence: Boolean = false) = Map[Set[ResourcePattern], Set[AccessControlEntry]](
+    TopicResources -> AclCommand.getAcls(Users, ALLOW, Set(WRITE, DESCRIBE, CREATE), Hosts),
+    TransactionalIdResources -> AclCommand.getAcls(Users, ALLOW, Set(WRITE, DESCRIBE), Hosts),
+    Set(ClusterResource) -> AclCommand.getAcls(Users, ALLOW,
+      Set(if (enableIdempotence) Some(IDEMPOTENT_WRITE) else None).flatten, Hosts)
   )
 
-  private val ConsumerResourceToAcls = Map[Set[Resource], Set[Acl]](
-    TopicResources -> AclCommand.getAcls(Users, Allow, Set(Read, Describe), Hosts),
-    GroupResources -> AclCommand.getAcls(Users, Allow, Set(Read), Hosts)
+  private val ConsumerResourceToAcls = Map[Set[ResourcePattern], Set[AccessControlEntry]](
+    TopicResources -> AclCommand.getAcls(Users, ALLOW, Set(READ, DESCRIBE), Hosts),
+    GroupResources -> AclCommand.getAcls(Users, ALLOW, Set(READ), Hosts)
   )
 
-  private val CmdToResourcesToAcl = Map[Array[String], Map[Set[Resource], Set[Acl]]](
+  private val CmdToResourcesToAcl = Map[Array[String], Map[Set[ResourcePattern], Set[AccessControlEntry]]](
     Array[String]("--producer") -> ProducerResourceToAcls(),
     Array[String]("--producer", "--idempotent") -> ProducerResourceToAcls(enableIdempotence = true),
     Array[String]("--consumer") -> ConsumerResourceToAcls,
     Array[String]("--producer", "--consumer") -> ConsumerResourceToAcls.map { case (k, v) => k -> (v ++
-      ProducerResourceToAcls().getOrElse(k, Set.empty[Acl])) },
+      ProducerResourceToAcls().getOrElse(k, Set.empty[AccessControlEntry])) },
     Array[String]("--producer", "--idempotent", "--consumer") -> ConsumerResourceToAcls.map { case (k, v) => k -> (v ++
-      ProducerResourceToAcls(enableIdempotence = true).getOrElse(k, Set.empty[Acl])) }
+      ProducerResourceToAcls(enableIdempotence = true).getOrElse(k, Set.empty[AccessControlEntry])) }
   )
 
   private var brokerProps: Properties = _
@@ -98,14 +105,14 @@ class AclCommandTest extends ZooKeeperTestHarness with Logging {
     super.setUp()
 
     brokerProps = TestUtils.createBrokerConfig(0, zkConnect)
-    brokerProps.put(KafkaConfig.AuthorizerClassNameProp, "kafka.security.auth.SimpleAclAuthorizer")
-    brokerProps.put(SimpleAclAuthorizer.SuperUsersProp, "User:ANONYMOUS")
+    brokerProps.put(KafkaConfig.AuthorizerClassNameProp, classOf[AclAuthorizer].getName)
+    brokerProps.put(AclAuthorizer.SuperUsersProp, "User:ANONYMOUS")
 
     zkArgs = Array("--authorizer-properties", "zookeeper.connect=" + zkConnect)
   }
 
   @After
-  override def tearDown() {
+  override def tearDown(): Unit = {
     TestUtils.shutdownServers(servers)
     super.tearDown()
   }
@@ -127,9 +134,9 @@ class AclCommandTest extends ZooKeeperTestHarness with Logging {
     adminArgs = Array("--bootstrap-server", TestUtils.bootstrapServers(servers, listenerName))
   }
 
-  private def testAclCli(cmdArgs: Array[String]) {
+  private def testAclCli(cmdArgs: Array[String]): Unit = {
     for ((resources, resourceCmd) <- ResourceToCommand) {
-      for (permissionType <- PermissionType.values) {
+      for (permissionType <- Set(ALLOW, DENY)) {
         val operationToCmd = ResourceToOperations(resources)
         val (acls, cmd) = getAclToCommand(permissionType, operationToCmd._1)
           AclCommand.main(cmdArgs ++ cmd ++ resourceCmd ++ operationToCmd._2 :+ "--add")
@@ -155,10 +162,10 @@ class AclCommandTest extends ZooKeeperTestHarness with Logging {
     testProducerConsumerCli(adminArgs)
   }
 
-  private def testProducerConsumerCli(cmdArgs: Array[String]) {
+  private def testProducerConsumerCli(cmdArgs: Array[String]): Unit = {
     for ((cmd, resourcesToAcls) <- CmdToResourcesToAcl) {
       val resourceCommand: Array[String] = resourcesToAcls.keys.map(ResourceToCommand).foldLeft(Array[String]())(_ ++ _)
-      AclCommand.main(cmdArgs ++ getCmd(Allow) ++ resourceCommand ++ cmd :+ "--add")
+      AclCommand.main(cmdArgs ++ getCmd(ALLOW) ++ resourceCommand ++ cmd :+ "--add")
       for ((resources, acls) <- resourcesToAcls) {
         for (resource <- resources) {
           withAuthorizer() { authorizer =>
@@ -187,29 +194,37 @@ class AclCommandTest extends ZooKeeperTestHarness with Logging {
     AclCommand.main(cmdArgs ++ cmd :+ "--add")
 
     withAuthorizer() { authorizer =>
-      val writeAcl = Acl(principal, Allow, Acl.WildCardHost, Write)
-      val describeAcl = Acl(principal, Allow, Acl.WildCardHost, Describe)
-      val createAcl = Acl(principal, Allow, Acl.WildCardHost, Create)
-      TestUtils.waitAndVerifyAcls(Set(writeAcl, describeAcl, createAcl), authorizer, Resource(Topic, "Test-", PREFIXED))
+      val writeAcl = new AccessControlEntry(principal.toString, AuthorizerUtils.WildcardHost, WRITE, ALLOW)
+      val describeAcl = new AccessControlEntry(principal.toString, AuthorizerUtils.WildcardHost, DESCRIBE, ALLOW)
+      val createAcl = new AccessControlEntry(principal.toString, AuthorizerUtils.WildcardHost, CREATE, ALLOW)
+      TestUtils.waitAndVerifyAcls(Set(writeAcl, describeAcl, createAcl), authorizer,
+        new ResourcePattern(TOPIC, "Test-", PREFIXED))
     }
 
     AclCommand.main(cmdArgs ++ cmd :+ "--remove" :+ "--force")
 
     withAuthorizer() { authorizer =>
-      TestUtils.waitAndVerifyAcls(Set.empty[Acl], authorizer, Resource(Cluster, "kafka-cluster", LITERAL))
-      TestUtils.waitAndVerifyAcls(Set.empty[Acl], authorizer, Resource(Topic, "Test-", PREFIXED))
+      TestUtils.waitAndVerifyAcls(Set.empty[AccessControlEntry], authorizer, new ResourcePattern(CLUSTER, "kafka-cluster", LITERAL))
+      TestUtils.waitAndVerifyAcls(Set.empty[AccessControlEntry], authorizer, new ResourcePattern(TOPIC, "Test-", PREFIXED))
     }
   }
 
   @Test(expected = classOf[IllegalArgumentException])
-  def testInvalidAuthorizerProperty() {
+  def testInvalidAuthorizerProperty(): Unit = {
     val args = Array("--authorizer-properties", "zookeeper.connect " + zkConnect)
-    val aclCommandService = new AclCommand.AuthorizerService(new AclCommandOptions(args))
+    val aclCommandService = new AclCommand.AuthorizerService(classOf[Authorizer], new AclCommandOptions(args))
+    aclCommandService.listAcls()
+  }
+
+  @Test(expected = classOf[IllegalArgumentException])
+  def testInvalidJAuthorizerProperty() {
+    val args = Array("--authorizer-properties", "zookeeper.connect " + zkConnect)
+    val aclCommandService = new AclCommand.JAuthorizerService(classOf[JAuthorizer], new AclCommandOptions(args))
     aclCommandService.listAcls()
   }
 
   @Test
-  def testPatternTypes() {
+  def testPatternTypes(): Unit = {
     Exit.setExitProcedure { (status, _) =>
       if (status == 1)
         throw new RuntimeException("Exiting command")
@@ -231,35 +246,36 @@ class AclCommandTest extends ZooKeeperTestHarness with Logging {
         verifyPatternType(listCmd, isValid = patternType != PatternType.UNKNOWN)
         val removeCmd = zkArgs ++ Array("--topic", "Test", "--force", "--remove", "--resource-pattern-type", patternType.toString)
         verifyPatternType(removeCmd, isValid = patternType != PatternType.UNKNOWN)
+
       }
     } finally {
       Exit.resetExitProcedure()
     }
   }
 
-  private def testRemove(cmdArgs: Array[String], resources: Set[Resource], resourceCmd: Array[String]) {
+  private def testRemove(cmdArgs: Array[String], resources: Set[ResourcePattern], resourceCmd: Array[String]): Unit = {
     for (resource <- resources) {
       AclCommand.main(cmdArgs ++ resourceCmd :+ "--remove" :+ "--force")
       withAuthorizer() { authorizer =>
-        TestUtils.waitAndVerifyAcls(Set.empty[Acl], authorizer, resource)
+        TestUtils.waitAndVerifyAcls(Set.empty[AccessControlEntry], authorizer, resource)
       }
     }
   }
 
-  private def getAclToCommand(permissionType: PermissionType, operations: Set[Operation]): (Set[Acl], Array[String]) = {
+  private def getAclToCommand(permissionType: AclPermissionType, operations: Set[AclOperation]): (Set[AccessControlEntry], Array[String]) = {
     (AclCommand.getAcls(Users, permissionType, operations, Hosts), getCmd(permissionType))
   }
 
-  private def getCmd(permissionType: PermissionType): Array[String] = {
-    val principalCmd = if (permissionType == Allow) "--allow-principal" else "--deny-principal"
-    val cmd = if (permissionType == Allow) AllowHostCommand else DenyHostCommand
+  private def getCmd(permissionType: AclPermissionType): Array[String] = {
+    val principalCmd = if (permissionType == ALLOW) "--allow-principal" else "--deny-principal"
+    val cmd = if (permissionType == ALLOW) AllowHostCommand else DenyHostCommand
 
     Users.foldLeft(cmd) ((cmd, user) => cmd ++ Array(principalCmd, user.toString))
   }
 
-  private def withAuthorizer()(f: Authorizer => Unit) {
+  private def withAuthorizer()(f: JAuthorizer => Unit): Unit = {
     val kafkaConfig = KafkaConfig.fromProps(brokerProps, doLog = false)
-    val authZ = new SimpleAclAuthorizer
+    val authZ = new AclAuthorizer
     try {
       authZ.configure(kafkaConfig.originals)
       f(authZ)
diff --git a/core/src/test/scala/unit/kafka/admin/AddPartitionsTest.scala b/core/src/test/scala/unit/kafka/admin/AddPartitionsTest.scala
index 31283468aac4b..115f091acb007 100755
--- a/core/src/test/scala/unit/kafka/admin/AddPartitionsTest.scala
+++ b/core/src/test/scala/unit/kafka/admin/AddPartitionsTest.scala
@@ -17,6 +17,7 @@
 
 package kafka.admin
 
+import kafka.controller.ReplicaAssignment
 import kafka.network.SocketServer
 import org.junit.Assert._
 import kafka.utils.TestUtils._
@@ -33,29 +34,29 @@ import scala.collection.JavaConverters._
 
 class AddPartitionsTest extends BaseRequestTest {
 
-  protected override def numBrokers: Int = 4
+  override def brokerCount: Int = 4
 
   val partitionId = 0
 
   val topic1 = "new-topic1"
-  val topic1Assignment = Map(0->Seq(0,1))
+  val topic1Assignment = Map(0 -> ReplicaAssignment(Seq(0,1), List(), List()))
   val topic2 = "new-topic2"
-  val topic2Assignment = Map(0->Seq(1,2))
+  val topic2Assignment = Map(0 -> ReplicaAssignment(Seq(1,2), List(), List()))
   val topic3 = "new-topic3"
-  val topic3Assignment = Map(0->Seq(2,3,0,1))
+  val topic3Assignment = Map(0 -> ReplicaAssignment(Seq(2,3,0,1), List(), List()))
   val topic4 = "new-topic4"
-  val topic4Assignment = Map(0->Seq(0,3))
+  val topic4Assignment = Map(0 -> ReplicaAssignment(Seq(0,3), List(), List()))
   val topic5 = "new-topic5"
-  val topic5Assignment = Map(1->Seq(0,1))
+  val topic5Assignment = Map(1 -> ReplicaAssignment(Seq(0,1), List(), List()))
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     super.setUp()
 
-    createTopic(topic1, partitionReplicaAssignment = topic1Assignment)
-    createTopic(topic2, partitionReplicaAssignment = topic2Assignment)
-    createTopic(topic3, partitionReplicaAssignment = topic3Assignment)
-    createTopic(topic4, partitionReplicaAssignment = topic4Assignment)
+    createTopic(topic1, partitionReplicaAssignment = topic1Assignment.mapValues(_.replicas).toMap)
+    createTopic(topic2, partitionReplicaAssignment = topic2Assignment.mapValues(_.replicas).toMap)
+    createTopic(topic3, partitionReplicaAssignment = topic3Assignment.mapValues(_.replicas).toMap)
+    createTopic(topic4, partitionReplicaAssignment = topic4Assignment.mapValues(_.replicas).toMap)
   }
 
   @Test
diff --git a/core/src/test/scala/unit/kafka/admin/AdminRackAwareTest.scala b/core/src/test/scala/unit/kafka/admin/AdminRackAwareTest.scala
index bece037fa812c..48bddd453f1f4 100644
--- a/core/src/test/scala/unit/kafka/admin/AdminRackAwareTest.scala
+++ b/core/src/test/scala/unit/kafka/admin/AdminRackAwareTest.scala
@@ -27,7 +27,7 @@ import scala.collection.Map
 class AdminRackAwareTest extends RackAwareTest with Logging {
 
   @Test
-  def testGetRackAlternatedBrokerListAndAssignReplicasToBrokers() {
+  def testGetRackAlternatedBrokerListAndAssignReplicasToBrokers(): Unit = {
     val rackMap = Map(0 -> "rack1", 1 -> "rack3", 2 -> "rack3", 3 -> "rack2", 4 -> "rack2", 5 -> "rack1")
     val newList = AdminUtils.getRackAlternatedBrokerList(rackMap)
     assertEquals(List(0, 3, 1, 5, 4, 2), newList)
@@ -45,7 +45,7 @@ class AdminRackAwareTest extends RackAwareTest with Logging {
   }
 
   @Test
-  def testAssignmentWithRackAware() {
+  def testAssignmentWithRackAware(): Unit = {
     val brokerRackMapping = Map(0 -> "rack1", 1 -> "rack2", 2 -> "rack2", 3 -> "rack3", 4 -> "rack3", 5 -> "rack1")
     val numPartitions = 6
     val replicationFactor = 3
@@ -56,7 +56,7 @@ class AdminRackAwareTest extends RackAwareTest with Logging {
   }
 
   @Test
-  def testAssignmentWithRackAwareWithRandomStartIndex() {
+  def testAssignmentWithRackAwareWithRandomStartIndex(): Unit = {
     val brokerRackMapping = Map(0 -> "rack1", 1 -> "rack2", 2 -> "rack2", 3 -> "rack3", 4 -> "rack3", 5 -> "rack1")
     val numPartitions = 6
     val replicationFactor = 3
@@ -67,7 +67,7 @@ class AdminRackAwareTest extends RackAwareTest with Logging {
   }
 
   @Test
-  def testAssignmentWithRackAwareWithUnevenReplicas() {
+  def testAssignmentWithRackAwareWithUnevenReplicas(): Unit = {
     val brokerRackMapping = Map(0 -> "rack1", 1 -> "rack2", 2 -> "rack2", 3 -> "rack3", 4 -> "rack3", 5 -> "rack1")
     val numPartitions = 13
     val replicationFactor = 3
@@ -78,7 +78,7 @@ class AdminRackAwareTest extends RackAwareTest with Logging {
   }
 
   @Test
-  def testAssignmentWithRackAwareWithUnevenRacks() {
+  def testAssignmentWithRackAwareWithUnevenRacks(): Unit = {
     val brokerRackMapping = Map(0 -> "rack1", 1 -> "rack1", 2 -> "rack2", 3 -> "rack3", 4 -> "rack3", 5 -> "rack1")
     val numPartitions = 12
     val replicationFactor = 3
@@ -89,7 +89,7 @@ class AdminRackAwareTest extends RackAwareTest with Logging {
   }
 
   @Test
-  def testAssignmentWith2ReplicasRackAware() {
+  def testAssignmentWith2ReplicasRackAware(): Unit = {
     val brokerRackMapping = Map(0 -> "rack1", 1 -> "rack2", 2 -> "rack2", 3 -> "rack3", 4 -> "rack3", 5 -> "rack1")
     val numPartitions = 12
     val replicationFactor = 2
@@ -100,7 +100,7 @@ class AdminRackAwareTest extends RackAwareTest with Logging {
   }
 
   @Test
-  def testRackAwareExpansion() {
+  def testRackAwareExpansion(): Unit = {
     val brokerRackMapping = Map(6 -> "rack1", 7 -> "rack2", 8 -> "rack2", 9 -> "rack3", 10 -> "rack3", 11 -> "rack1")
     val numPartitions = 12
     val replicationFactor = 2
@@ -111,7 +111,7 @@ class AdminRackAwareTest extends RackAwareTest with Logging {
   }
 
   @Test
-  def testAssignmentWith2ReplicasRackAwareWith6Partitions() {
+  def testAssignmentWith2ReplicasRackAwareWith6Partitions(): Unit = {
     val brokerRackMapping = Map(0 -> "rack1", 1 -> "rack2", 2 -> "rack2", 3 -> "rack3", 4 -> "rack3", 5 -> "rack1")
     val numPartitions = 6
     val replicationFactor = 2
@@ -122,7 +122,7 @@ class AdminRackAwareTest extends RackAwareTest with Logging {
   }
 
   @Test
-  def testAssignmentWith2ReplicasRackAwareWith6PartitionsAnd3Brokers() {
+  def testAssignmentWith2ReplicasRackAwareWith6PartitionsAnd3Brokers(): Unit = {
     val brokerRackMapping = Map(0 -> "rack1", 1 -> "rack2", 4 -> "rack3")
     val numPartitions = 3
     val replicationFactor = 2
@@ -131,7 +131,7 @@ class AdminRackAwareTest extends RackAwareTest with Logging {
   }
 
   @Test
-  def testLargeNumberPartitionsAssignment() {
+  def testLargeNumberPartitionsAssignment(): Unit = {
     val numPartitions = 96
     val replicationFactor = 3
     val brokerRackMapping = Map(0 -> "rack1", 1 -> "rack2", 2 -> "rack2", 3 -> "rack3", 4 -> "rack3", 5 -> "rack1",
@@ -143,37 +143,37 @@ class AdminRackAwareTest extends RackAwareTest with Logging {
   }
 
   @Test
-  def testMoreReplicasThanRacks() {
+  def testMoreReplicasThanRacks(): Unit = {
     val numPartitions = 6
     val replicationFactor = 5
     val brokerRackMapping = Map(0 -> "rack1", 1 -> "rack2", 2 -> "rack2", 3 -> "rack3", 4 -> "rack3", 5 -> "rack2")
     val assignment = AdminUtils.assignReplicasToBrokers(toBrokerMetadata(brokerRackMapping), numPartitions, replicationFactor)
-    assertEquals(List.fill(assignment.size)(replicationFactor), assignment.values.map(_.size))
+    assertEquals(List.fill(assignment.size)(replicationFactor), assignment.values.toIndexedSeq.map(_.size))
     val distribution = getReplicaDistribution(assignment, brokerRackMapping)
     for (partition <- 0 until numPartitions)
       assertEquals(3, distribution.partitionRacks(partition).toSet.size)
   }
 
   @Test
-  def testLessReplicasThanRacks() {
+  def testLessReplicasThanRacks(): Unit = {
     val numPartitions = 6
     val replicationFactor = 2
     val brokerRackMapping = Map(0 -> "rack1", 1 -> "rack2", 2 -> "rack2", 3 -> "rack3", 4 -> "rack3", 5 -> "rack2")
     val assignment = AdminUtils.assignReplicasToBrokers(toBrokerMetadata(brokerRackMapping), numPartitions,
       replicationFactor)
-    assertEquals(List.fill(assignment.size)(replicationFactor), assignment.values.map(_.size))
+    assertEquals(List.fill(assignment.size)(replicationFactor), assignment.values.toIndexedSeq.map(_.size))
     val distribution = getReplicaDistribution(assignment, brokerRackMapping)
     for (partition <- 0 to 5)
       assertEquals(2, distribution.partitionRacks(partition).toSet.size)
   }
 
   @Test
-  def testSingleRack() {
+  def testSingleRack(): Unit = {
     val numPartitions = 6
     val replicationFactor = 3
     val brokerRackMapping = Map(0 -> "rack1", 1 -> "rack1", 2 -> "rack1", 3 -> "rack1", 4 -> "rack1", 5 -> "rack1")
     val assignment = AdminUtils.assignReplicasToBrokers(toBrokerMetadata(brokerRackMapping), numPartitions, replicationFactor)
-    assertEquals(List.fill(assignment.size)(replicationFactor), assignment.values.map(_.size))
+    assertEquals(List.fill(assignment.size)(replicationFactor), assignment.values.toIndexedSeq.map(_.size))
     val distribution = getReplicaDistribution(assignment, brokerRackMapping)
     for (partition <- 0 until numPartitions)
       assertEquals(1, distribution.partitionRacks(partition).toSet.size)
@@ -182,7 +182,7 @@ class AdminRackAwareTest extends RackAwareTest with Logging {
   }
 
   @Test
-  def testSkipBrokerWithReplicaAlreadyAssigned() {
+  def testSkipBrokerWithReplicaAlreadyAssigned(): Unit = {
     val rackInfo = Map(0 -> "a", 1 -> "b", 2 -> "c", 3 -> "a", 4 -> "a")
     val brokerList = 0 to 4
     val numPartitions = 6
@@ -196,7 +196,7 @@ class AdminRackAwareTest extends RackAwareTest with Logging {
   }
 
   @Test
-  def testReplicaAssignment() {
+  def testReplicaAssignment(): Unit = {
     val brokerMetadatas = (0 to 4).map(new BrokerMetadata(_, None))
 
     // test 0 replication factor
diff --git a/core/src/test/scala/unit/kafka/admin/AdminTest.scala b/core/src/test/scala/unit/kafka/admin/AdminTest.scala
deleted file mode 100755
index a64f6e7ab9217..0000000000000
--- a/core/src/test/scala/unit/kafka/admin/AdminTest.scala
+++ /dev/null
@@ -1,212 +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 org.apache.kafka.common.errors.{InvalidReplicaAssignmentException, InvalidTopicException, TopicExistsException}
-import org.apache.kafka.common.metrics.Quota
-import org.easymock.EasyMock
-import org.junit.Assert._
-import org.junit.{After, Before, Test}
-import java.util.Properties
-
-import kafka.utils._
-import kafka.zk.{ConfigEntityZNode, ZooKeeperTestHarness}
-import kafka.utils.{Logging, TestUtils, ZkUtils}
-import kafka.server.{ConfigType, KafkaConfig, KafkaServer}
-
-import scala.collection.{Map, immutable}
-import org.apache.kafka.common.security.JaasUtils
-
-import scala.collection.JavaConverters._
-
-@deprecated("This test has been deprecated and will be removed in a future release.", "1.1.0")
-class AdminTest extends ZooKeeperTestHarness with Logging with RackAwareTest {
-
-  var servers: Seq[KafkaServer] = Seq()
-  var zkUtils: ZkUtils = null
-
-  @Before
-  override def setUp() {
-    super.setUp()
-    zkUtils = ZkUtils(zkConnect, zkSessionTimeout, zkConnectionTimeout, zkAclsEnabled.getOrElse(JaasUtils.isZkSecurityEnabled))
-  }
-
-  @After
-  override def tearDown() {
-    if (zkUtils != null)
-     CoreUtils.swallow(zkUtils.close(), this)
-    TestUtils.shutdownServers(servers)
-    super.tearDown()
-  }
-
-  @Test
-  def testManualReplicaAssignment() {
-    val brokers = List(0, 1, 2, 3, 4)
-    TestUtils.createBrokersInZk(zkClient, brokers)
-
-    // duplicate brokers
-    intercept[InvalidReplicaAssignmentException] {
-      AdminUtils.createOrUpdateTopicPartitionAssignmentPathInZK(zkUtils, "test", Map(0->Seq(0,0)))
-    }
-
-    // inconsistent replication factor
-    intercept[InvalidReplicaAssignmentException] {
-      AdminUtils.createOrUpdateTopicPartitionAssignmentPathInZK(zkUtils, "test", Map(0->Seq(0,1), 1->Seq(0)))
-    }
-
-    // good assignment
-    val assignment = Map(0 -> List(0, 1, 2),
-                         1 -> List(1, 2, 3))
-    AdminUtils.createOrUpdateTopicPartitionAssignmentPathInZK(zkUtils, "test", assignment)
-    val found = zkUtils.getPartitionAssignmentForTopics(Seq("test"))
-    assertEquals(assignment, found("test"))
-  }
-
-  @Test
-  def testTopicCreationInZK() {
-    val expectedReplicaAssignment = Map(
-      0  -> List(0, 1, 2),
-      1  -> List(1, 2, 3),
-      2  -> List(2, 3, 4),
-      3  -> List(3, 4, 0),
-      4  -> List(4, 0, 1),
-      5  -> List(0, 2, 3),
-      6  -> List(1, 3, 4),
-      7  -> List(2, 4, 0),
-      8  -> List(3, 0, 1),
-      9  -> List(4, 1, 2),
-      10 -> List(1, 2, 3),
-      11 -> List(1, 3, 4)
-    )
-    val leaderForPartitionMap = immutable.Map(
-      0 -> 0,
-      1 -> 1,
-      2 -> 2,
-      3 -> 3,
-      4 -> 4,
-      5 -> 0,
-      6 -> 1,
-      7 -> 2,
-      8 -> 3,
-      9 -> 4,
-      10 -> 1,
-      11 -> 1
-    )
-    val topic = "test"
-    TestUtils.createBrokersInZk(zkClient, List(0, 1, 2, 3, 4))
-    // create the topic
-    AdminUtils.createOrUpdateTopicPartitionAssignmentPathInZK(zkUtils, topic, expectedReplicaAssignment)
-    // create leaders for all partitions
-    TestUtils.makeLeaderForPartition(zkClient, topic, leaderForPartitionMap, 1)
-    val actualReplicaList = leaderForPartitionMap.keys.toArray.map(p => p -> zkUtils.getReplicasForPartition(topic, p)).toMap
-    assertEquals(expectedReplicaAssignment.size, actualReplicaList.size)
-    for(i <- 0 until actualReplicaList.size)
-      assertEquals(expectedReplicaAssignment.get(i).get, actualReplicaList(i))
-
-    intercept[TopicExistsException] {
-      // shouldn't be able to create a topic that already exists
-      AdminUtils.createOrUpdateTopicPartitionAssignmentPathInZK(zkUtils, topic, expectedReplicaAssignment)
-    }
-  }
-
-  @Test
-  def testTopicCreationWithCollision() {
-    val topic = "test.topic"
-    val collidingTopic = "test_topic"
-    TestUtils.createBrokersInZk(zkClient, List(0, 1, 2, 3, 4))
-    // create the topic
-    AdminUtils.createTopic(zkUtils, topic, 3, 1)
-
-    intercept[InvalidTopicException] {
-      // shouldn't be able to create a topic that collides
-      AdminUtils.createTopic(zkUtils, collidingTopic, 3, 1)
-    }
-  }
-
-  @Test
-  def testConcurrentTopicCreation() {
-    val topic = "test.topic"
-
-    // simulate the ZK interactions that can happen when a topic is concurrently created by multiple processes
-    val zkMock: ZkUtils = EasyMock.createNiceMock(classOf[ZkUtils])
-    EasyMock.expect(zkMock.pathExists(s"/brokers/topics/$topic")).andReturn(false)
-    EasyMock.expect(zkMock.getAllTopics).andReturn(Seq("some.topic", topic, "some.other.topic"))
-    EasyMock.replay(zkMock)
-
-    intercept[TopicExistsException] {
-      AdminUtils.validateCreateOrUpdateTopic(zkMock, topic, Map.empty, new Properties, update = false)
-    }
-  }
-
-  /**
-   * This test simulates a client config change in ZK whose notification has been purged.
-   * Basically, it asserts that notifications are bootstrapped from ZK
-   */
-  @Test
-  def testBootstrapClientIdConfig() {
-    val clientId = "my-client"
-    val props = new Properties()
-    props.setProperty("producer_byte_rate", "1000")
-    props.setProperty("consumer_byte_rate", "2000")
-
-    // Write config without notification to ZK.
-    val configMap = Map[String, String] ("producer_byte_rate" -> "1000", "consumer_byte_rate" -> "2000")
-    val map = Map("version" -> 1, "config" -> configMap.asJava)
-    zkUtils.updatePersistentPath(ConfigEntityZNode.path(ConfigType.Client, clientId), Json.encodeAsString(map.asJava))
-
-    val configInZk: Map[String, Properties] = AdminUtils.fetchAllEntityConfigs(zkUtils, ConfigType.Client)
-    assertEquals("Must have 1 overridden client config", 1, configInZk.size)
-    assertEquals(props, configInZk(clientId))
-
-    // Test that the existing clientId overrides are read
-    val server = TestUtils.createServer(KafkaConfig.fromProps(TestUtils.createBrokerConfig(0, zkConnect)))
-    servers = Seq(server)
-    assertEquals(new Quota(1000, true), server.dataPlaneRequestProcessor.quotas.produce.quota("ANONYMOUS", clientId))
-    assertEquals(new Quota(2000, true), server.dataPlaneRequestProcessor.quotas.fetch.quota("ANONYMOUS", clientId))
-  }
-
-  @Test
-  def testGetBrokerMetadatas() {
-    // broker 4 has no rack information
-    val brokerList = 0 to 5
-    val rackInfo = Map(0 -> "rack1", 1 -> "rack2", 2 -> "rack2", 3 -> "rack1", 5 -> "rack3")
-    val brokerMetadatas = toBrokerMetadata(rackInfo, brokersWithoutRack = brokerList.filterNot(rackInfo.keySet))
-    TestUtils.createBrokersInZk(brokerMetadatas, zkClient)
-
-    val processedMetadatas1 = AdminUtils.getBrokerMetadatas(zkUtils, RackAwareMode.Disabled)
-    assertEquals(brokerList, processedMetadatas1.map(_.id))
-    assertEquals(List.fill(brokerList.size)(None), processedMetadatas1.map(_.rack))
-
-    val processedMetadatas2 = AdminUtils.getBrokerMetadatas(zkUtils, RackAwareMode.Safe)
-    assertEquals(brokerList, processedMetadatas2.map(_.id))
-    assertEquals(List.fill(brokerList.size)(None), processedMetadatas2.map(_.rack))
-
-    intercept[AdminOperationException] {
-      AdminUtils.getBrokerMetadatas(zkUtils, RackAwareMode.Enforced)
-    }
-
-    val partialList = List(0, 1, 2, 3, 5)
-    val processedMetadatas3 = AdminUtils.getBrokerMetadatas(zkUtils, RackAwareMode.Enforced, Some(partialList))
-    assertEquals(partialList, processedMetadatas3.map(_.id))
-    assertEquals(partialList.map(rackInfo), processedMetadatas3.flatMap(_.rack))
-
-    val numPartitions = 3
-    AdminUtils.createTopic(zkUtils, "foo", numPartitions, 2, rackAwareMode = RackAwareMode.Safe)
-    val assignment = zkUtils.getReplicaAssignmentForTopics(Seq("foo"))
-    assertEquals(numPartitions, assignment.size)
-  }
-}
diff --git a/core/src/test/scala/unit/kafka/admin/ConfigCommandTest.scala b/core/src/test/scala/unit/kafka/admin/ConfigCommandTest.scala
index ee4a6ef84211d..4fd110aab6a13 100644
--- a/core/src/test/scala/unit/kafka/admin/ConfigCommandTest.scala
+++ b/core/src/test/scala/unit/kafka/admin/ConfigCommandTest.scala
@@ -22,7 +22,7 @@ import java.util.Properties
 import kafka.admin.ConfigCommand.ConfigCommandOptions
 import kafka.api.ApiVersion
 import kafka.cluster.{Broker, EndPoint}
-import kafka.server.{ConfigEntityName, KafkaConfig}
+import kafka.server.{ConfigEntityName, ConfigType, KafkaConfig}
 import kafka.utils.{Exit, Logging}
 import kafka.zk.{AdminZkClient, BrokerInfo, KafkaZkClient, ZooKeeperTestHarness}
 import org.apache.kafka.clients.admin._
@@ -37,6 +37,7 @@ import org.apache.kafka.common.utils.Sanitizer
 import org.easymock.EasyMock
 import org.junit.Assert._
 import org.junit.Test
+import org.scalatest.Assertions.intercept
 
 import scala.collection.{Seq, mutable}
 import scala.collection.JavaConverters._
@@ -86,47 +87,58 @@ class ConfigCommandTest extends ZooKeeperTestHarness with Logging {
   }
 
   @Test
-  def shouldParseArgumentsForClientsEntityType() {
+  def shouldParseArgumentsForClientsEntityType(): Unit = {
     testArgumentParse("clients")
   }
 
   @Test
-  def shouldParseArgumentsForTopicsEntityType() {
+  def shouldParseArgumentsForTopicsEntityType(): Unit = {
     testArgumentParse("topics")
   }
 
   @Test
-  def shouldParseArgumentsForBrokersEntityType() {
+  def shouldParseArgumentsForBrokersEntityType(): Unit = {
     testArgumentParse("brokers")
   }
 
-  def testArgumentParse(entityType: String) = {
+  @Test
+  def shouldParseArgumentsForBrokerLoggersEntityType(): Unit = {
+    testArgumentParse("broker-loggers",
+      zkConfig = false)
+  }
+
+  def testArgumentParse(entityType: String, zkConfig: Boolean=true): Unit = {
+    val connectOpts = if (zkConfig)
+      ("--zookeeper", zkConnect)
+    else
+      ("--bootstrap-server", "localhost:9092")
+
     // Should parse correctly
-    var createOpts = new ConfigCommandOptions(Array("--zookeeper", zkConnect,
-      "--entity-name", "x",
+    var createOpts = new ConfigCommandOptions(Array(connectOpts._1, connectOpts._2,
+      "--entity-name", "1",
       "--entity-type", entityType,
       "--describe"))
     createOpts.checkArgs()
 
     // For --alter and added config
-    createOpts = new ConfigCommandOptions(Array("--zookeeper", zkConnect,
-      "--entity-name", "x",
+    createOpts = new ConfigCommandOptions(Array(connectOpts._1, connectOpts._2,
+      "--entity-name", "1",
       "--entity-type", entityType,
       "--alter",
       "--add-config", "a=b,c=d"))
     createOpts.checkArgs()
 
     // For alter and deleted config
-    createOpts = new ConfigCommandOptions(Array("--zookeeper", zkConnect,
-      "--entity-name", "x",
+    createOpts = new ConfigCommandOptions(Array(connectOpts._1, connectOpts._2,
+      "--entity-name", "1",
       "--entity-type", entityType,
       "--alter",
       "--delete-config", "a,b,c"))
     createOpts.checkArgs()
 
     // For alter and both added, deleted config
-    createOpts = new ConfigCommandOptions(Array("--zookeeper", zkConnect,
-      "--entity-name", "x",
+    createOpts = new ConfigCommandOptions(Array(connectOpts._1, connectOpts._2,
+      "--entity-name", "1",
       "--entity-type", entityType,
       "--alter",
       "--add-config", "a=b,c=d",
@@ -142,8 +154,8 @@ class ConfigCommandTest extends ZooKeeperTestHarness with Logging {
     assertEquals(1, deletedProps.size)
     assertEquals("a", deletedProps.head)
 
-    createOpts = new ConfigCommandOptions(Array("--zookeeper", zkConnect,
-      "--entity-name", "x",
+    createOpts = new ConfigCommandOptions(Array(connectOpts._1, connectOpts._2,
+      "--entity-name", "1",
       "--entity-type", entityType,
       "--alter",
       "--add-config", "a=b,c=,d=e,f="))
@@ -164,6 +176,13 @@ class ConfigCommandTest extends ZooKeeperTestHarness with Logging {
     ConfigCommand.alterConfig(null, createOpts, new DummyAdminZkClient(zkClient))
   }
 
+  @Test(expected = classOf[IllegalArgumentException])
+  def shouldFailIfBrokerEntityTypeIsNotAnInteger(): Unit = {
+    val createOpts = new ConfigCommandOptions(Array("--zookeeper", zkConnect,
+      "--entity-name", "A", "--entity-type", "brokers", "--alter", "--add-config", "a=b,c=d"))
+    ConfigCommand.alterConfig(null, createOpts, new DummyAdminZkClient(zkClient))
+  }
+
   @Test
   def shouldAddClientConfig(): Unit = {
     val createOpts = new ConfigCommandOptions(Array("--zookeeper", zkConnect,
@@ -227,6 +246,69 @@ class ConfigCommandTest extends ZooKeeperTestHarness with Logging {
     verifyAlterBrokerConfig(node, "1", List("--entity-name", "1"))
   }
 
+  @Test
+  def shouldAddBrokerLoggerConfig(): Unit = {
+    val node = new Node(1, "localhost", 9092)
+    verifyAlterBrokerLoggerConfig(node, "1", "1", List(
+      new ConfigEntry("kafka.log.LogCleaner", "INFO"),
+      new ConfigEntry("kafka.server.ReplicaManager", "INFO"),
+      new ConfigEntry("kafka.server.KafkaApi", "INFO")
+    ))
+  }
+
+  @Test
+  def testNoSpecifiedEntityOptionWithDescribeBrokersInZKIsAllowed(): Unit = {
+    val optsList = List("--zookeeper", "localhost:9092",
+      "--entity-type", ConfigType.Broker,
+      "--describe"
+    )
+
+    new ConfigCommandOptions(optsList.toArray).checkArgs()
+  }
+
+  @Test(expected = classOf[IllegalArgumentException])
+  def testNoSpecifiedEntityOptionWithDescribeBrokersInBootstrapServerIsNotAllowed(): Unit = {
+    val optsList = List("--bootstrap-server", "localhost:9092",
+      "--entity-type", ConfigType.Broker,
+      "--describe"
+    )
+
+    new ConfigCommandOptions(optsList.toArray).checkArgs()
+  }
+
+  @Test(expected = classOf[IllegalArgumentException])
+  def testEntityDefaultOptionWithDescribeBrokerLoggerIsNotAllowed(): Unit = {
+    val optsList = List("--bootstrap-server", "localhost:9092",
+      "--entity-type", ConfigCommand.BrokerLoggerConfigType,
+      "--entity-default",
+      "--describe"
+    )
+
+    new ConfigCommandOptions(optsList.toArray).checkArgs()
+  }
+
+  @Test(expected = classOf[IllegalArgumentException])
+  def testEntityDefaultOptionWithAlterBrokerLoggerIsNotAllowed(): Unit = {
+    val optsList = List("--bootstrap-server", "localhost:9092",
+      "--entity-type", ConfigCommand.BrokerLoggerConfigType,
+      "--entity-default",
+      "--alter",
+      "--add-config", "kafka.log.LogCleaner=DEBUG"
+    )
+
+    new ConfigCommandOptions(optsList.toArray).checkArgs()
+  }
+
+  @Test(expected = classOf[InvalidConfigurationException])
+  def shouldRaiseInvalidConfigurationExceptionWhenAddingInvalidBrokerLoggerConfig(): Unit = {
+    val node = new Node(1, "localhost", 9092)
+    // verifyAlterBrokerLoggerConfig tries to alter kafka.log.LogCleaner, kafka.server.ReplicaManager and kafka.server.KafkaApi
+    // yet, we make it so DescribeConfigs returns only one logger, implying that kafka.server.ReplicaManager and kafka.log.LogCleaner are invalid
+    verifyAlterBrokerLoggerConfig(node, "1", "1", List(
+      new ConfigEntry("kafka.server.KafkaApi", "INFO")
+    ))
+  }
+
   @Test
   def shouldAddDefaultBrokerDynamicConfig(): Unit = {
     val node = new Node(1, "localhost", 9092)
@@ -273,11 +355,66 @@ class ConfigCommandTest extends ZooKeeperTestHarness with Logging {
       }
     }
     EasyMock.replay(alterResult, describeResult)
-    ConfigCommand.alterBrokerConfig(mockAdminClient, alterOpts, resourceName)
+    ConfigCommand.alterBrokerConfig(mockAdminClient, alterOpts, ConfigType.Broker, resourceName)
     assertEquals(Map("message.max.bytes" -> "10", "num.io.threads" -> "5"), brokerConfigs.toMap)
     EasyMock.reset(alterResult, describeResult)
   }
 
+  def verifyAlterBrokerLoggerConfig(node: Node, resourceName: String, entityName: String,
+                                    describeConfigEntries: List[ConfigEntry]): Unit = {
+    val optsList = List("--bootstrap-server", "localhost:9092",
+      "--entity-type", ConfigCommand.BrokerLoggerConfigType,
+      "--alter",
+      "--entity-name", entityName,
+      "--add-config", "kafka.log.LogCleaner=DEBUG",
+      "--delete-config", "kafka.server.ReplicaManager,kafka.server.KafkaApi")
+    val alterOpts = new ConfigCommandOptions(optsList.toArray)
+    var alteredConfigs = false
+
+    val resource = new ConfigResource(ConfigResource.Type.BROKER_LOGGER, resourceName)
+    val future = new KafkaFutureImpl[util.Map[ConfigResource, Config]]
+    future.complete(util.Collections.singletonMap(resource, new Config(describeConfigEntries.asJava)))
+    val describeResult: DescribeConfigsResult = EasyMock.createNiceMock(classOf[DescribeConfigsResult])
+    EasyMock.expect(describeResult.all()).andReturn(future).once()
+
+    val alterFuture = new KafkaFutureImpl[Void]
+    alterFuture.complete(null)
+    val alterResult: AlterConfigsResult = EasyMock.createNiceMock(classOf[AlterConfigsResult])
+    EasyMock.expect(alterResult.all()).andReturn(alterFuture)
+
+    val mockAdminClient = new MockAdminClient(util.Collections.singletonList(node), node) {
+      override def describeConfigs(resources: util.Collection[ConfigResource]): DescribeConfigsResult = {
+        assertEquals(1, resources.size)
+        val resource = resources.iterator.next
+        assertEquals(ConfigResource.Type.BROKER_LOGGER, resource.`type`)
+        assertEquals(resourceName, resource.name)
+        describeResult
+      }
+
+      override def incrementalAlterConfigs(configs: util.Map[ConfigResource, util.Collection[AlterConfigOp]], options: AlterConfigsOptions): AlterConfigsResult = {
+        assertEquals(1, configs.size)
+        val entry = configs.entrySet.iterator.next
+        val resource = entry.getKey
+        val alterConfigOps = entry.getValue
+        assertEquals(ConfigResource.Type.BROKER_LOGGER, resource.`type`)
+        assertEquals(3, alterConfigOps.size)
+
+        val expectedConfigOps = List(
+          new AlterConfigOp(new ConfigEntry("kafka.log.LogCleaner", "DEBUG"), AlterConfigOp.OpType.SET),
+          new AlterConfigOp(new ConfigEntry("kafka.server.ReplicaManager", ""), AlterConfigOp.OpType.DELETE),
+          new AlterConfigOp(new ConfigEntry("kafka.server.KafkaApi", ""), AlterConfigOp.OpType.DELETE)
+        )
+        assertEquals(expectedConfigOps, alterConfigOps.asScala.toList)
+        alteredConfigs = true
+        alterResult
+      }
+    }
+    EasyMock.replay(alterResult, describeResult)
+    ConfigCommand.alterBrokerConfig(mockAdminClient, alterOpts, ConfigCommand.BrokerLoggerConfigType, resourceName)
+    assertTrue(alteredConfigs)
+    EasyMock.reset(alterResult, describeResult)
+  }
+
   @Test
   def shouldSupportCommaSeparatedValues(): Unit = {
     val createOpts = new ConfigCommandOptions(Array("--zookeeper", zkConnect,
@@ -507,7 +644,7 @@ class ConfigCommandTest extends ZooKeeperTestHarness with Logging {
   }
 
   @Test
-  def testQuotaConfigEntity() {
+  def testQuotaConfigEntity(): Unit = {
 
     def createOpts(entityType: String, entityName: Option[String], otherArgs: Array[String]) : ConfigCommandOptions = {
       val optArray = Array("--zookeeper", zkConnect,
@@ -519,7 +656,7 @@ class ConfigCommandTest extends ZooKeeperTestHarness with Logging {
       new ConfigCommandOptions(optArray ++ nameArray ++ otherArgs)
     }
 
-    def checkEntity(entityType: String, entityName: Option[String], expectedEntityName: String, otherArgs: Array[String]) {
+    def checkEntity(entityType: String, entityName: Option[String], expectedEntityName: String, otherArgs: Array[String]): Unit = {
       val opts = createOpts(entityType, entityName, otherArgs)
       opts.checkArgs()
       val entity = ConfigCommand.parseEntity(opts)
@@ -527,7 +664,7 @@ class ConfigCommandTest extends ZooKeeperTestHarness with Logging {
       assertEquals(expectedEntityName, entity.fullSanitizedName)
     }
 
-    def checkInvalidEntity(entityType: String, entityName: Option[String], otherArgs: Array[String]) {
+    def checkInvalidEntity(entityType: String, entityName: Option[String], otherArgs: Array[String]): Unit = {
       val opts = createOpts(entityType, entityName, otherArgs)
       try {
         opts.checkArgs()
@@ -579,8 +716,8 @@ class ConfigCommandTest extends ZooKeeperTestHarness with Logging {
   }
 
   @Test
-  def testUserClientQuotaOpts() {
-    def checkEntity(expectedEntityType: String, expectedEntityName: String, args: String*) {
+  def testUserClientQuotaOpts(): Unit = {
+    def checkEntity(expectedEntityType: String, expectedEntityName: String, args: String*): Unit = {
       val opts = new ConfigCommandOptions(Array("--zookeeper", zkConnect) ++ args)
       opts.checkArgs()
       val entity = ConfigCommand.parseEntity(opts)
@@ -621,10 +758,10 @@ class ConfigCommandTest extends ZooKeeperTestHarness with Logging {
   }
 
   @Test
-  def testQuotaDescribeEntities() {
+  def testQuotaDescribeEntities(): Unit = {
     val zkClient: KafkaZkClient = EasyMock.createNiceMock(classOf[KafkaZkClient])
 
-    def checkEntities(opts: Array[String], expectedFetches: Map[String, Seq[String]], expectedEntityNames: Seq[String]) {
+    def checkEntities(opts: Array[String], expectedFetches: Map[String, Seq[String]], expectedEntityNames: Seq[String]): Unit = {
       val entity = ConfigCommand.parseEntity(new ConfigCommandOptions(opts :+ "--describe"))
       expectedFetches.foreach {
         case (name, values) => EasyMock.expect(zkClient.getAllEntitiesWithConfig(name)).andReturn(values)
diff --git a/core/src/test/scala/unit/kafka/admin/ConsumerGroupCommandTest.scala b/core/src/test/scala/unit/kafka/admin/ConsumerGroupCommandTest.scala
index d5eea981635c6..1fbdb0811c505 100644
--- a/core/src/test/scala/unit/kafka/admin/ConsumerGroupCommandTest.scala
+++ b/core/src/test/scala/unit/kafka/admin/ConsumerGroupCommandTest.scala
@@ -25,14 +25,15 @@ import kafka.admin.ConsumerGroupCommand.{ConsumerGroupCommandOptions, ConsumerGr
 import kafka.integration.KafkaServerTestHarness
 import kafka.server.KafkaConfig
 import kafka.utils.TestUtils
-import org.apache.kafka.clients.consumer.{KafkaConsumer, RangeAssignor}
-import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.clients.admin.AdminClientConfig
+import org.apache.kafka.clients.consumer.{ConsumerConfig, 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.{After, Before}
 
-import scala.collection.mutable.ArrayBuffer
 import scala.collection.JavaConverters._
+import scala.collection.mutable.ArrayBuffer
 
 class ConsumerGroupCommandTest extends KafkaServerTestHarness {
   import ConsumerGroupCommandTest._
@@ -51,7 +52,7 @@ class ConsumerGroupCommandTest extends KafkaServerTestHarness {
   }
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     super.setUp()
     createTopic(topic, 1, 1)
   }
@@ -64,27 +65,27 @@ class ConsumerGroupCommandTest extends KafkaServerTestHarness {
   }
 
   def committedOffsets(topic: String = topic, group: String = group): Map[TopicPartition, Long] = {
-    val props = new Properties
-    props.put("bootstrap.servers", brokerList)
-    props.put("group.id", group)
-    val consumer = new KafkaConsumer(props, new StringDeserializer, new StringDeserializer)
+    val consumer = createNoAutoCommitConsumer(group)
     try {
-      consumer.partitionsFor(topic).asScala.flatMap { partitionInfo =>
-        val tp = new TopicPartition(partitionInfo.topic, partitionInfo.partition)
-        val committed = consumer.committed(tp)
-        if (committed == null)
-          None
-        else
-          Some(tp -> committed.offset)
-      }.toMap
+      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).mapValues(_.offset()).toMap
     } finally {
       consumer.close()
     }
   }
 
+  def createNoAutoCommitConsumer(group: String): KafkaConsumer[String, String] = {
+    val props = new Properties
+    props.put("bootstrap.servers", brokerList)
+    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)
+    val service = new ConsumerGroupService(opts, Map(AdminClientConfig.RETRIES_CONFIG -> Int.MaxValue.toString))
     consumerGroupService = service :: consumerGroupService
     service
   }
@@ -93,8 +94,9 @@ class ConsumerGroupCommandTest extends KafkaServerTestHarness {
                                topic: String = topic,
                                group: String = group,
                                strategy: String = classOf[RangeAssignor].getName,
-                               customPropsOpt: Option[Properties] = None): ConsumerGroupExecutor = {
-    val executor = new ConsumerGroupExecutor(brokerList, numConsumers, group, topic, strategy, customPropsOpt)
+                               customPropsOpt: Option[Properties] = None,
+                               syncCommit: Boolean = false): ConsumerGroupExecutor = {
+    val executor = new ConsumerGroupExecutor(brokerList, numConsumers, group, topic, strategy, customPropsOpt, syncCommit)
     addExecutor(executor)
     executor
   }
@@ -115,7 +117,8 @@ class ConsumerGroupCommandTest extends KafkaServerTestHarness {
 
 object ConsumerGroupCommandTest {
 
-  abstract class AbstractConsumerRunnable(broker: String, groupId: String, customPropsOpt: Option[Properties] = None) extends Runnable {
+  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)
@@ -126,15 +129,19 @@ object ConsumerGroupCommandTest {
       props.put("group.id", groupId)
       props.put("key.deserializer", classOf[StringDeserializer].getName)
       props.put("value.deserializer", classOf[StringDeserializer].getName)
+      props.put(ConsumerConfig.ALLOW_AUTO_CREATE_TOPICS_CONFIG, boolean2Boolean(true))
     }
 
     def subscribe(): Unit
 
-    def run() {
+    def run(): Unit = {
       try {
         subscribe()
-        while (true)
+        while (true) {
           consumer.poll(Duration.ofMillis(Long.MaxValue))
+          if (syncCommit)
+            consumer.commitSync()
+        }
       } catch {
         case _: WakeupException => // OK
       } finally {
@@ -142,13 +149,14 @@ object ConsumerGroupCommandTest {
       }
     }
 
-    def shutdown() {
+    def shutdown(): Unit = {
       consumer.wakeup()
     }
   }
 
-  class ConsumerRunnable(broker: String, groupId: String, topic: String, strategy: String, customPropsOpt: Option[Properties] = None)
-    extends AbstractConsumerRunnable(broker, groupId, customPropsOpt) {
+  class ConsumerRunnable(broker: String, groupId: String, topic: String, strategy: String,
+                         customPropsOpt: Option[Properties] = None, syncCommit: Boolean = false)
+    extends AbstractConsumerRunnable(broker, groupId, customPropsOpt, syncCommit) {
 
     override def configure(props: Properties): Unit = {
       super.configure(props)
@@ -172,12 +180,12 @@ object ConsumerGroupCommandTest {
     private val executor: ExecutorService = Executors.newFixedThreadPool(numThreads)
     private val consumers = new ArrayBuffer[AbstractConsumerRunnable]()
 
-    def submit(consumerThread: AbstractConsumerRunnable) {
+    def submit(consumerThread: AbstractConsumerRunnable): Unit = {
       consumers += consumerThread
       executor.submit(consumerThread)
     }
 
-    def shutdown() {
+    def shutdown(): Unit = {
       consumers.foreach(_.shutdown())
       executor.shutdown()
       executor.awaitTermination(5000, TimeUnit.MILLISECONDS)
@@ -185,11 +193,11 @@ object ConsumerGroupCommandTest {
   }
 
   class ConsumerGroupExecutor(broker: String, numConsumers: Int, groupId: String, topic: String, strategy: String,
-                              customPropsOpt: Option[Properties] = None)
+                              customPropsOpt: Option[Properties] = None, syncCommit: Boolean = false)
     extends AbstractConsumerGroupExecutor(numConsumers) {
 
     for (_ <- 1 to numConsumers) {
-      submit(new ConsumerRunnable(broker, groupId, topic, strategy, customPropsOpt))
+      submit(new ConsumerRunnable(broker, groupId, topic, strategy, customPropsOpt, syncCommit))
     }
 
   }
diff --git a/core/src/test/scala/unit/kafka/admin/DelegationTokenCommandTest.scala b/core/src/test/scala/unit/kafka/admin/DelegationTokenCommandTest.scala
index a4c27e502883f..cae0902de4c44 100644
--- a/core/src/test/scala/unit/kafka/admin/DelegationTokenCommandTest.scala
+++ b/core/src/test/scala/unit/kafka/admin/DelegationTokenCommandTest.scala
@@ -22,10 +22,11 @@ import kafka.admin.DelegationTokenCommand.DelegationTokenCommandOptions
 import kafka.api.{KafkaSasl, SaslSetup}
 import kafka.server.{BaseRequestTest, KafkaConfig}
 import kafka.utils.{JaasTestUtils, TestUtils}
-import org.apache.kafka.clients.admin.AdminClientConfig
+import org.apache.kafka.clients.admin.{Admin, AdminClient, AdminClientConfig}
 import org.apache.kafka.common.security.auth.SecurityProtocol
 import org.junit.Assert._
 import org.junit.{After, Before, Test}
+import org.scalatest.Assertions.intercept
 
 import scala.collection.JavaConverters._
 import scala.collection.mutable.ListBuffer
@@ -37,9 +38,9 @@ class DelegationTokenCommandTest extends BaseRequestTest with SaslSetup {
   private val kafkaServerSaslMechanisms = List("PLAIN")
   protected override val serverSaslProperties = Some(kafkaServerSaslProperties(kafkaServerSaslMechanisms, kafkaClientSaslMechanism))
   protected override val clientSaslProperties = Some(kafkaClientSaslProperties(kafkaClientSaslMechanism))
-  var adminClient: org.apache.kafka.clients.admin.AdminClient = null
+  var adminClient: Admin = null
 
-  override def numBrokers = 1
+  override def brokerCount = 1
 
   @Before
   override def setUp(): Unit = {
@@ -48,11 +49,11 @@ class DelegationTokenCommandTest extends BaseRequestTest with SaslSetup {
   }
 
   override def generateConfigs = {
-    val props = TestUtils.createBrokerConfigs(numBrokers, zkConnect,
+    val props = TestUtils.createBrokerConfigs(brokerCount, zkConnect,
       enableControlledShutdown = false,
       interBrokerSecurityProtocol = Some(securityProtocol),
       trustStoreFile = trustStoreFile, saslProperties = serverSaslProperties, enableToken = true)
-    props.foreach(propertyOverrides)
+    props.foreach(brokerPropertyOverrides)
     props.map(KafkaConfig.fromProps)
   }
 
@@ -67,7 +68,7 @@ class DelegationTokenCommandTest extends BaseRequestTest with SaslSetup {
 
   @Test
   def testDelegationTokenRequests(): Unit = {
-    adminClient = org.apache.kafka.clients.admin.AdminClient.create(createAdminConfig)
+    adminClient = AdminClient.create(createAdminConfig)
     val renewer1 = "User:renewer1"
     val renewer2 = "User:renewer2"
 
diff --git a/core/src/test/scala/unit/kafka/admin/DeleteConsumerGroupsTest.scala b/core/src/test/scala/unit/kafka/admin/DeleteConsumerGroupsTest.scala
index 1bcc316a09c04..23de7904d5eb6 100644
--- a/core/src/test/scala/unit/kafka/admin/DeleteConsumerGroupsTest.scala
+++ b/core/src/test/scala/unit/kafka/admin/DeleteConsumerGroupsTest.scala
@@ -26,14 +26,14 @@ import org.junit.Test
 class DeleteConsumerGroupsTest extends ConsumerGroupCommandTest {
 
   @Test(expected = classOf[OptionException])
-  def testDeleteWithTopicOption() {
+  def testDeleteWithTopicOption(): Unit = {
     TestUtils.createOffsetsTopic(zkClient, servers)
     val cgcArgs = Array("--bootstrap-server", brokerList, "--delete", "--group", group, "--topic")
     getConsumerGroupService(cgcArgs)
   }
 
   @Test
-  def testDeleteCmdNonExistingGroup() {
+  def testDeleteCmdNonExistingGroup(): Unit = {
     TestUtils.createOffsetsTopic(zkClient, servers)
     val missingGroup = "missing.group"
 
@@ -42,11 +42,11 @@ class DeleteConsumerGroupsTest extends ConsumerGroupCommandTest {
 
     val output = TestUtils.grabConsoleOutput(service.deleteGroups())
     assertTrue(s"The expected error (${Errors.GROUP_ID_NOT_FOUND}) was not detected while deleting consumer group",
-        output.contains(s"Group '$missingGroup' could not be deleted due to:") && output.contains(Errors.GROUP_ID_NOT_FOUND.message))
+      output.contains(s"Group '$missingGroup' could not be deleted due to:") && output.contains(Errors.GROUP_ID_NOT_FOUND.message))
   }
 
   @Test
-  def testDeleteNonExistingGroup() {
+  def testDeleteNonExistingGroup(): Unit = {
     TestUtils.createOffsetsTopic(zkClient, servers)
     val missingGroup = "missing.group"
 
@@ -56,12 +56,12 @@ class DeleteConsumerGroupsTest extends ConsumerGroupCommandTest {
 
     val result = service.deleteGroups()
     assertTrue(s"The expected error (${Errors.GROUP_ID_NOT_FOUND}) was not detected while deleting consumer group",
-      result.size == 1 && result.keySet.contains(missingGroup) && result.get(missingGroup).get.getCause
+      result.size == 1 && result.keySet.contains(missingGroup) && result(missingGroup).getCause
         .isInstanceOf[GroupIdNotFoundException])
   }
 
   @Test
-  def testDeleteCmdNonEmptyGroup() {
+  def testDeleteCmdNonEmptyGroup(): Unit = {
     TestUtils.createOffsetsTopic(zkClient, servers)
 
     // run one consumer in the group
@@ -70,8 +70,8 @@ class DeleteConsumerGroupsTest extends ConsumerGroupCommandTest {
     val service = getConsumerGroupService(cgcArgs)
 
     TestUtils.waitUntilTrue(() => {
-      service.collectGroupMembers(false)._2.get.size == 1
-    }, "The group did not initialize as expected.", maxRetries = 3)
+      service.collectGroupMembers(group, false)._2.get.size == 1
+    }, "The group did not initialize as expected.")
 
     val output = TestUtils.grabConsoleOutput(service.deleteGroups())
     assertTrue(s"The expected error (${Errors.NON_EMPTY_GROUP}) was not detected while deleting consumer group. Output was: (${output})",
@@ -79,7 +79,7 @@ class DeleteConsumerGroupsTest extends ConsumerGroupCommandTest {
   }
 
   @Test
-  def testDeleteNonEmptyGroup() {
+  def testDeleteNonEmptyGroup(): Unit = {
     TestUtils.createOffsetsTopic(zkClient, servers)
 
     // run one consumer in the group
@@ -88,17 +88,17 @@ class DeleteConsumerGroupsTest extends ConsumerGroupCommandTest {
     val service = getConsumerGroupService(cgcArgs)
 
     TestUtils.waitUntilTrue(() => {
-      service.collectGroupMembers(false)._2.get.size == 1
-    }, "The group did not initialize as expected.", maxRetries = 3)
+      service.collectGroupMembers(group, false)._2.get.size == 1
+    }, "The group did not initialize as expected.")
 
     val result = service.deleteGroups()
-    assertNotNull(s"Group was deleted successfully, but it shouldn't have been. Result was:(${result})", result.get(group).get)
+    assertNotNull(s"Group was deleted successfully, but it shouldn't have been. Result was:(${result})", result(group))
     assertTrue(s"The expected error (${Errors.NON_EMPTY_GROUP}) was not detected while deleting consumer group. Result was:(${result})",
-      result.size == 1 && result.keySet.contains(group) && result.get(group).get.getCause.isInstanceOf[GroupNotEmptyException])
+      result.size == 1 && result.keySet.contains(group) && result(group).getCause.isInstanceOf[GroupNotEmptyException])
   }
 
   @Test
-  def testDeleteCmdEmptyGroup() {
+  def testDeleteCmdEmptyGroup(): Unit = {
     TestUtils.createOffsetsTopic(zkClient, servers)
 
     // run one consumer in the group
@@ -107,14 +107,14 @@ class DeleteConsumerGroupsTest extends ConsumerGroupCommandTest {
     val service = getConsumerGroupService(cgcArgs)
 
     TestUtils.waitUntilTrue(() => {
-      service.listGroups().contains(group)
-    }, "The group did not initialize as expected.", maxRetries = 3)
+      service.listGroups().contains(group) && service.collectGroupState(group).state == "Stable"
+    }, "The group did not initialize as expected.")
 
     executor.shutdown()
 
     TestUtils.waitUntilTrue(() => {
-      service.collectGroupState().state == "Empty"
-    }, "The group did become empty as expected.", maxRetries = 3)
+      service.collectGroupState(group).state == "Empty"
+    }, "The group did not become empty as expected.")
 
     val output = TestUtils.grabConsoleOutput(service.deleteGroups())
     assertTrue(s"The consumer group could not be deleted as expected",
@@ -122,7 +122,45 @@ class DeleteConsumerGroupsTest extends ConsumerGroupCommandTest {
   }
 
   @Test
-  def testDeleteEmptyGroup() {
+  def testDeleteCmdAllGroups(): Unit = {
+    TestUtils.createOffsetsTopic(zkClient, servers)
+
+    // 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", brokerList, "--delete", "--all-groups")
+    val service = getConsumerGroupService(cgcArgs)
+
+    TestUtils.waitUntilTrue(() => {
+      service.listGroups().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(s"The consumer group(s) could not be deleted as expected",
+      output.matches(s"Deletion of requested consumer groups (.*) was successful.")
+        && deletedGroupsGrepped == expectedGroupsForDeletion
+    )
+  }
+
+  @Test
+  def testDeleteEmptyGroup(): Unit = {
     TestUtils.createOffsetsTopic(zkClient, servers)
 
     // run one consumer in the group
@@ -131,22 +169,22 @@ class DeleteConsumerGroupsTest extends ConsumerGroupCommandTest {
     val service = getConsumerGroupService(cgcArgs)
 
     TestUtils.waitUntilTrue(() => {
-      service.listGroups().contains(group)
-    }, "The group did not initialize as expected.", maxRetries = 3)
+      service.listGroups().contains(group) && service.collectGroupState(group).state == "Stable"
+    }, "The group did not initialize as expected.")
 
     executor.shutdown()
 
     TestUtils.waitUntilTrue(() => {
-      service.collectGroupState().state == "Empty"
-    }, "The group did become empty as expected.", maxRetries = 3)
+      service.collectGroupState(group).state == "Empty"
+    }, "The group did not become empty as expected.")
 
     val result = service.deleteGroups()
     assertTrue(s"The consumer group could not be deleted as expected",
-      result.size == 1 && result.keySet.contains(group) && result.get(group).get == null)
+      result.size == 1 && result.keySet.contains(group) && result(group) == null)
   }
 
   @Test
-  def testDeleteCmdWithMixOfSuccessAndError() {
+  def testDeleteCmdWithMixOfSuccessAndError(): Unit = {
     TestUtils.createOffsetsTopic(zkClient, servers)
     val missingGroup = "missing.group"
 
@@ -156,24 +194,24 @@ class DeleteConsumerGroupsTest extends ConsumerGroupCommandTest {
     val service = getConsumerGroupService(cgcArgs)
 
     TestUtils.waitUntilTrue(() => {
-      service.listGroups().contains(group)
-    }, "The group did not initialize as expected.", maxRetries = 3)
+      service.listGroups().contains(group) && service.collectGroupState(group).state == "Stable"
+    }, "The group did not initialize as expected.")
 
     executor.shutdown()
 
     TestUtils.waitUntilTrue(() => {
-      service.collectGroupState().state == "Empty"
-    }, "The group did become empty as expected.", maxRetries = 3)
+      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(s"The consumer group deletion did not work as expected",
       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'"))
+        output.contains(s"These consumer groups were deleted successfully: '$group'"))
   }
 
   @Test
-  def testDeleteWithMixOfSuccessAndError() {
+  def testDeleteWithMixOfSuccessAndError(): Unit = {
     TestUtils.createOffsetsTopic(zkClient, servers)
     val missingGroup = "missing.group"
 
@@ -183,26 +221,27 @@ class DeleteConsumerGroupsTest extends ConsumerGroupCommandTest {
     val service = getConsumerGroupService(cgcArgs)
 
     TestUtils.waitUntilTrue(() => {
-      service.listGroups().contains(group)
-    }, "The group did not initialize as expected.", maxRetries = 3)
+      service.listGroups().contains(group) && service.collectGroupState(group).state == "Stable"
+    }, "The group did not initialize as expected.")
 
     executor.shutdown()
 
     TestUtils.waitUntilTrue(() => {
-      service.collectGroupState().state == "Empty"
-    }, "The group did become empty as expected.", maxRetries = 3)
+      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(s"The consumer group deletion did not work as expected",
       result.size == 2 &&
-        result.keySet.contains(group) && result.get(group).get == null &&
-        result.keySet.contains(missingGroup) && result.get(missingGroup).get.getMessage.contains(Errors.GROUP_ID_NOT_FOUND.message))
+        result.keySet.contains(group) && result(group) == null &&
+        result.keySet.contains(missingGroup) &&
+        result(missingGroup).getMessage.contains(Errors.GROUP_ID_NOT_FOUND.message))
   }
 
 
   @Test(expected = classOf[OptionException])
-  def testDeleteWithUnrecognizedNewConsumerOption() {
+  def testDeleteWithUnrecognizedNewConsumerOption(): Unit = {
     val cgcArgs = Array("--new-consumer", "--bootstrap-server", brokerList, "--delete", "--group", group)
     getConsumerGroupService(cgcArgs)
   }
diff --git a/core/src/test/scala/unit/kafka/admin/DeleteOffsetsConsumerGroupCommandIntegrationTest.scala b/core/src/test/scala/unit/kafka/admin/DeleteOffsetsConsumerGroupCommandIntegrationTest.scala
new file mode 100644
index 0000000000000..0e2885231fe41
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/admin/DeleteOffsetsConsumerGroupCommandIntegrationTest.scala
@@ -0,0 +1,197 @@
+/**
+ * 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.server.Defaults
+import kafka.utils.TestUtils
+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.junit.Test
+import org.junit.Assert._
+
+class DeleteOffsetsConsumerGroupCommandIntegrationTest extends ConsumerGroupCommandTest {
+
+  def getArgs(group: String, topic: String): Array[String] = {
+    Array(
+      "--bootstrap-server", brokerList,
+      "--delete-offsets",
+      "--group", group,
+      "--topic", topic
+    )
+  }
+
+  @Test
+  def testDeleteOffsetsNonExistingGroup(): 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)
+  }
+
+  @Test
+  def testDeleteOffsetsOfStableConsumerGroupWithTopicPartition(): Unit = {
+    testWithStableConsumerGroup(topic, 0, 0, Errors.GROUP_SUBSCRIBED_TO_TOPIC)
+  }
+
+  @Test
+  def testDeleteOffsetsOfStableConsumerGroupWithTopicOnly(): Unit = {
+    testWithStableConsumerGroup(topic, -1, 0, Errors.GROUP_SUBSCRIBED_TO_TOPIC)
+  }
+
+  @Test
+  def testDeleteOffsetsOfStableConsumerGroupWithUnknownTopicPartition(): Unit = {
+    testWithStableConsumerGroup("foobar", 0, 0, Errors.UNKNOWN_TOPIC_OR_PARTITION)
+  }
+
+  @Test
+  def testDeleteOffsetsOfStableConsumerGroupWithUnknownTopicOnly(): Unit = {
+    testWithStableConsumerGroup("foobar", -1, -1, Errors.UNKNOWN_TOPIC_OR_PARTITION)
+  }
+
+  @Test
+  def testDeleteOffsetsOfEmptyConsumerGroupWithTopicPartition(): Unit = {
+    testWithEmptyConsumerGroup(topic, 0, 0, Errors.NONE)
+  }
+
+  @Test
+  def testDeleteOffsetsOfEmptyConsumerGroupWithTopicOnly(): Unit = {
+    testWithEmptyConsumerGroup(topic, -1, 0, Errors.NONE)
+  }
+
+  @Test
+  def testDeleteOffsetsOfEmptyConsumerGroupWithUnknownTopicPartition(): Unit = {
+    testWithEmptyConsumerGroup("foobar", 0, 0, Errors.UNKNOWN_TOPIC_OR_PARTITION)
+  }
+
+  @Test
+  def testDeleteOffsetsOfEmptyConsumerGroupWithUnknownTopicOnly(): 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, brokerList)
+    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()): KafkaConsumer[Array[Byte], Array[Byte]] = {
+    config.putIfAbsent(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, brokerList)
+    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.GroupMaxSessionTimeoutMs.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 3f3e75464bc07..d2374cccc2094 100644
--- a/core/src/test/scala/unit/kafka/admin/DeleteTopicTest.scala
+++ b/core/src/test/scala/unit/kafka/admin/DeleteTopicTest.scala
@@ -16,34 +16,39 @@
  */
 package kafka.admin
 
+import java.util.Properties
+
+import scala.collection.Seq
+
 import kafka.log.Log
-import kafka.zk.{TopicPartitionZNode, ZooKeeperTestHarness}
+import kafka.zk.{DeleteTopicFlagZNode, TopicPartitionZNode, TopicZNode, ZooKeeperTestHarness}
 import kafka.utils.TestUtils
 import kafka.server.{KafkaConfig, KafkaServer}
 import org.junit.Assert._
 import org.junit.{After, Test}
-import java.util.Properties
 
 import kafka.admin.TopicCommand.ZookeeperTopicService
 import kafka.common.TopicAlreadyMarkedForDeletionException
-import kafka.controller.{OfflineReplica, PartitionAndReplica, ReplicaDeletionSuccessful}
+import kafka.controller.{OfflineReplica, PartitionAndReplica, ReplicaAssignment, ReplicaDeletionSuccessful}
 import org.apache.kafka.common.TopicPartition
 import org.apache.kafka.common.errors.UnknownTopicOrPartitionException
+import org.scalatest.Assertions.fail
 
 class DeleteTopicTest extends ZooKeeperTestHarness {
 
   var servers: Seq[KafkaServer] = Seq()
 
   val expectedReplicaAssignment = Map(0 -> List(0, 1, 2))
+  val expectedReplicaFullAssignment = expectedReplicaAssignment.mapValues(ReplicaAssignment(_, List(), List())).toMap
 
   @After
-  override def tearDown() {
+  override def tearDown(): Unit = {
     TestUtils.shutdownServers(servers)
     super.tearDown()
   }
 
   @Test
-  def testDeleteTopicWithAllAliveReplicas() {
+  def testDeleteTopicWithAllAliveReplicas(): Unit = {
     val topic = "test"
     servers = createTestTopicAndCluster(topic)
     // start topic deletion
@@ -52,7 +57,7 @@ class DeleteTopicTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  def testResumeDeleteTopicWithRecoveredFollower() {
+  def testResumeDeleteTopicWithRecoveredFollower(): Unit = {
     val topicPartition = new TopicPartition("test", 0)
     val topic = topicPartition.topic
     servers = createTestTopicAndCluster(topic)
@@ -69,14 +74,14 @@ class DeleteTopicTest extends ZooKeeperTestHarness {
         .forall(_.getLogManager().getLog(topicPartition).isEmpty), "Replicas 0,1 have not deleted log.")
     // ensure topic deletion is halted
     TestUtils.waitUntilTrue(() => zkClient.isTopicMarkedForDeletion(topic),
-      "Admin path /admin/delete_topic/test path deleted even when a follower replica is down")
+      "Admin path /admin/delete_topics/test path deleted even when a follower replica is down")
     // restart follower replica
     follower.startup()
     TestUtils.verifyTopicDeletion(zkClient, topic, 1, servers)
   }
 
   @Test
-  def testResumeDeleteTopicOnControllerFailover() {
+  def testResumeDeleteTopicOnControllerFailover(): Unit = {
     val topicPartition = new TopicPartition("test", 0)
     val topic = topicPartition.topic
     servers = createTestTopicAndCluster(topic)
@@ -93,7 +98,7 @@ class DeleteTopicTest extends ZooKeeperTestHarness {
 
     // ensure topic deletion is halted
     TestUtils.waitUntilTrue(() => zkClient.isTopicMarkedForDeletion(topic),
-      "Admin path /admin/delete_topic/test path deleted even when a replica is down")
+      "Admin path /admin/delete_topics/test path deleted even when a replica is down")
 
     controller.startup()
     follower.startup()
@@ -102,8 +107,7 @@ class DeleteTopicTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  def testPartitionReassignmentDuringDeleteTopic() {
-    val expectedReplicaAssignment = Map(0 -> List(0, 1, 2))
+  def testPartitionReassignmentDuringDeleteTopic(): Unit = {
     val topic = "test"
     val topicPartition = new TopicPartition(topic, 0)
     val brokerConfigs = TestUtils.createBrokerConfigs(4, zkConnect, false)
@@ -138,8 +142,11 @@ class DeleteTopicTest extends ZooKeeperTestHarness {
     }, "Partition reassignment shouldn't complete.")
     val controllerId = zkClient.getControllerId.getOrElse(fail("Controller doesn't exist"))
     val controller = servers.filter(s => s.config.brokerId == controllerId).head
-    assertFalse("Partition reassignment should fail",
-      controller.kafkaController.controllerContext.partitionsBeingReassigned.contains(topicPartition))
+
+    // partitionsBeingReassigned is updated after re-assignment znode is removed, so wait again
+    TestUtils.waitUntilTrue(() => {
+      !controller.kafkaController.controllerContext.partitionsBeingReassigned.contains(topicPartition)
+    }, "Partition should be removed from partitionsBeingReassigned.")
     val assignedReplicas = zkClient.getReplicasForPartition(new TopicPartition(topic, 0))
     assertEquals("Partition should not be reassigned to 0, 1, 2", oldAssignedReplicas, assignedReplicas)
     follower.startup()
@@ -170,8 +177,7 @@ class DeleteTopicTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  def testIncreasePartitionCountDuringDeleteTopic() {
-    val expectedReplicaAssignment = Map(0 -> List(0, 1, 2))
+  def testIncreasePartitionCountDuringDeleteTopic(): Unit = {
     val topic = "test"
     val topicPartition = new TopicPartition(topic, 0)
     val brokerConfigs = TestUtils.createBrokerConfigs(4, zkConnect, false)
@@ -198,8 +204,8 @@ class DeleteTopicTest extends ZooKeeperTestHarness {
     val (controller, controllerId) = getController()
     val allReplicasForTopic = getAllReplicasFromAssignment(topic, expectedReplicaAssignment)
     TestUtils.waitUntilTrue(() => {
-      val replicasInDeletionSuccessful = controller.kafkaController.replicaStateMachine.replicasInState(topic, ReplicaDeletionSuccessful)
-      val offlineReplicas = controller.kafkaController.replicaStateMachine.replicasInState(topic, OfflineReplica)
+      val replicasInDeletionSuccessful = controller.kafkaController.controllerContext.replicasInState(topic, ReplicaDeletionSuccessful)
+      val offlineReplicas = controller.kafkaController.controllerContext.replicasInState(topic, OfflineReplica)
       allReplicasForTopic == (replicasInDeletionSuccessful union offlineReplicas)
     }, s"Not all replicas for topic $topic are in states of either ReplicaDeletionSuccessful or OfflineReplica")
 
@@ -227,7 +233,7 @@ class DeleteTopicTest extends ZooKeeperTestHarness {
 
 
   @Test
-  def testDeleteTopicDuringAddPartition() {
+  def testDeleteTopicDuringAddPartition(): Unit = {
     val topic = "test"
     servers = createTestTopicAndCluster(topic)
     val leaderIdOpt = zkClient.getLeaderForPartition(new TopicPartition(topic, 0))
@@ -241,7 +247,7 @@ class DeleteTopicTest extends ZooKeeperTestHarness {
     TestUtils.waitUntilTrue(() => zkClient.getBroker(follower.config.brokerId).isEmpty,
       s"Follower ${follower.config.brokerId} was not removed from ZK")
     // add partitions to topic
-    adminZkClient.addPartitions(topic, expectedReplicaAssignment, brokers, 2,
+    adminZkClient.addPartitions(topic, expectedReplicaFullAssignment, brokers, 2,
       Some(Map(1 -> Seq(0, 1, 2), 2 -> Seq(0, 1, 2))))
     // start topic deletion
     adminZkClient.deleteTopic(topic)
@@ -255,7 +261,7 @@ class DeleteTopicTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  def testAddPartitionDuringDeleteTopic() {
+  def testAddPartitionDuringDeleteTopic(): Unit = {
     zkClient.createTopLevelPaths()
     val topic = "test"
     servers = createTestTopicAndCluster(topic)
@@ -264,7 +270,7 @@ class DeleteTopicTest extends ZooKeeperTestHarness {
     adminZkClient.deleteTopic(topic)
     // add partitions to topic
     val newPartition = new TopicPartition(topic, 1)
-    adminZkClient.addPartitions(topic, expectedReplicaAssignment, brokers, 2,
+    adminZkClient.addPartitions(topic, expectedReplicaFullAssignment, brokers, 2,
       Some(Map(1 -> Seq(0, 1, 2), 2 -> Seq(0, 1, 2))))
     TestUtils.verifyTopicDeletion(zkClient, topic, 1, servers)
     // verify that new partition doesn't exist on any broker either
@@ -273,7 +279,7 @@ class DeleteTopicTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  def testRecreateTopicAfterDeletion() {
+  def testRecreateTopicAfterDeletion(): Unit = {
     val expectedReplicaAssignment = Map(0 -> List(0, 1, 2))
     val topic = "test"
     val topicPartition = new TopicPartition(topic, 0)
@@ -289,7 +295,7 @@ class DeleteTopicTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  def testDeleteNonExistingTopic() {
+  def testDeleteNonExistingTopic(): Unit = {
     val topicPartition = new TopicPartition("test", 0)
     val topic = topicPartition.topic
     servers = createTestTopicAndCluster(topic)
@@ -312,7 +318,7 @@ class DeleteTopicTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  def testDeleteTopicWithCleaner() {
+  def testDeleteTopicWithCleaner(): Unit = {
     val topicName = "test"
     val topicPartition = new TopicPartition(topicName, 0)
     val topic = topicPartition.topic
@@ -342,7 +348,7 @@ class DeleteTopicTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  def testDeleteTopicAlreadyMarkedAsDeleted() {
+  def testDeleteTopicAlreadyMarkedAsDeleted(): Unit = {
     val topicPartition = new TopicPartition("test", 0)
     val topic = topicPartition.topic
     servers = createTestTopicAndCluster(topic)
@@ -390,14 +396,14 @@ class DeleteTopicTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  def testDisableDeleteTopic() {
+  def testDisableDeleteTopic(): Unit = {
     val topicPartition = new TopicPartition("test", 0)
     val topic = topicPartition.topic
     servers = createTestTopicAndCluster(topic, deleteTopicEnabled = false)
     // mark the topic for deletion
     adminZkClient.deleteTopic("test")
     TestUtils.waitUntilTrue(() => !zkClient.isTopicMarkedForDeletion(topic),
-      "Admin path /admin/delete_topic/%s path not deleted even if deleteTopic is disabled".format(topic))
+      "Admin path /admin/delete_topics/%s path not deleted even if deleteTopic is disabled".format(topic))
     // verify that topic test is untouched
     assertTrue(servers.forall(_.getLogManager().getLog(topicPartition).isDefined))
     // test the topic path exists
@@ -407,8 +413,64 @@ class DeleteTopicTest extends ZooKeeperTestHarness {
     assertTrue("Leader should exist for topic test", leaderIdOpt.isDefined)
   }
 
+  def testDeleteTopicAfterEnableZkDeleteTopicFlag() {
+    val topicPartition = new TopicPartition("test", 0)
+    val topic = topicPartition.topic
+    servers = createTestTopicAndCluster(topic, deleteTopicEnabled = false)
+    // mark the topic for deletion
+    adminZkClient.deleteTopic("test")
+    TestUtils.waitUntilTrue(() => !zkClient.isTopicMarkedForDeletion(topic),
+      "Admin path /admin/delete_topic/%s path not deleted even if deleteTopic is disabled".format(topic))
+    // verify that topic test is untouched
+    assertTrue(servers.forall(_.getLogManager().getLog(topicPartition).isDefined))
+    // test the topic path exists
+    assertTrue("Topic path disappeared even when topic deletion is disabled", zkClient.pathExists(TopicZNode.path(topic)))
+    // topic test should have a leader
+    val leaderIdOpt = zkClient.getLeaderForPartition(new TopicPartition(topic, 0))
+    assertTrue("Leader should exist for topic test", leaderIdOpt.isDefined)
+
+    // Set TopicDeletionFlag to true in zk and try delete topic again
+    zkClient.setTopicDeletionFlag("true")
+    TestUtils.waitUntilTrue(() =>
+      try {
+        zkClient.getTopicDeletionFlag.equalsIgnoreCase("true")
+      } catch {
+        case _: Throwable => false
+      },
+      "TopicDeletionFlag is not set")
+    TestUtils.waitUntilTrue( () => getController()._1.kafkaController.topicDeletionManager.isDeleteTopicEnabled,
+      "Delete topic is not enabled")
+    // mark the topic for deletion
+    adminZkClient.deleteTopic("test")
+    TestUtils.verifyTopicDeletion(zkClient, "test", 1, servers)
+
+    // Set TopicDeletionFlag to invalid value in zk
+    zkClient.setTopicDeletionFlag("flase")
+    TestUtils.waitUntilTrue(() =>
+      try {
+        zkClient.getTopicDeletionFlag.equalsIgnoreCase("true")
+      } catch {
+        case _: Throwable => false
+      },
+      "TopicDeletionFlag is not overwritten")
+
+    // delete TopicDeletionFlagPath in zk
+    zkClient.deletePath(DeleteTopicFlagZNode.path)
+    TestUtils.waitUntilTrue(() =>
+      try {
+        !zkClient.pathExists(DeleteTopicFlagZNode.path)
+      } catch {
+        case _: Throwable => false
+      },
+      "TopicDeletionFlagPath is not deleted")
+    TestUtils.waitUntilTrue(() =>
+      getController()._1.kafkaController.topicDeletionManager.isDeleteTopicEnabled == false,
+      "Topic deletion flag is not rest"
+    )
+  }
+
   @Test
-  def testDeletingPartiallyDeletedTopic() {
+  def testDeletingPartiallyDeletedTopic(): Unit = {
     /**
       * A previous controller could have deleted some partitions of a topic from ZK, but not all partitions, and then crashed.
       * In that case, the new controller should be able to handle the partially deleted topic, and finish the deletion.
diff --git a/core/src/test/scala/unit/kafka/admin/DescribeConsumerGroupTest.scala b/core/src/test/scala/unit/kafka/admin/DescribeConsumerGroupTest.scala
index 88f9f4a686293..e4d167cd3b6b9 100644
--- a/core/src/test/scala/unit/kafka/admin/DescribeConsumerGroupTest.scala
+++ b/core/src/test/scala/unit/kafka/admin/DescribeConsumerGroupTest.scala
@@ -22,7 +22,7 @@ import joptsimple.OptionException
 import kafka.utils.TestUtils
 import org.apache.kafka.clients.consumer.{ConsumerConfig, RoundRobinAssignor}
 import org.apache.kafka.common.TopicPartition
-import org.apache.kafka.common.errors.{TimeoutException}
+import org.apache.kafka.common.errors.TimeoutException
 import org.junit.Assert._
 import org.junit.Test
 
@@ -37,7 +37,7 @@ class DescribeConsumerGroupTest extends ConsumerGroupCommandTest {
   private val describeTypes = describeTypeOffsets ++ describeTypeMembers ++ describeTypeState
 
   @Test
-  def testDescribeNonExistingGroup() {
+  def testDescribeNonExistingGroup(): Unit = {
     TestUtils.createOffsetsTopic(zkClient, servers)
     val missingGroup = "missing.group"
 
@@ -46,64 +46,67 @@ class DescribeConsumerGroupTest extends ConsumerGroupCommandTest {
       val cgcArgs = Array("--bootstrap-server", brokerList, "--describe", "--group", missingGroup) ++ describeType
       val service = getConsumerGroupService(cgcArgs)
 
-      val output = TestUtils.grabConsoleOutput(service.describeGroup())
+      val output = TestUtils.grabConsoleOutput(service.describeGroups())
       assertTrue(s"Expected error was not detected for describe option '${describeType.mkString(" ")}'",
           output.contains(s"Consumer group '$missingGroup' does not exist."))
     }
   }
 
   @Test(expected = classOf[OptionException])
-  def testDescribeWithMultipleSubActions() {
+  def testDescribeWithMultipleSubActions(): Unit = {
     TestUtils.createOffsetsTopic(zkClient, servers)
     val cgcArgs = Array("--bootstrap-server", brokerList, "--describe", "--group", group, "--members", "--state")
     getConsumerGroupService(cgcArgs)
   }
 
   @Test
-  def testDescribeOffsetsOfNonExistingGroup() {
+  def testDescribeOffsetsOfNonExistingGroup(): Unit = {
+    val group = "missing.group"
     TestUtils.createOffsetsTopic(zkClient, servers)
 
     // run one consumer in the group consuming from a single-partition topic
     addConsumerGroupExecutor(numConsumers = 1)
     // note the group to be queried is a different (non-existing) group
-    val cgcArgs = Array("--bootstrap-server", brokerList, "--describe", "--group", "missing.group")
+    val cgcArgs = Array("--bootstrap-server", brokerList, "--describe", "--group", group)
     val service = getConsumerGroupService(cgcArgs)
 
-    val (state, assignments) = service.collectGroupOffsets()
+    val (state, assignments) = service.collectGroupOffsets(group)
     assertTrue(s"Expected the state to be 'Dead', with no members in the group '$group'.",
         state.contains("Dead") && assignments.contains(List()))
   }
 
   @Test
-  def testDescribeMembersOfNonExistingGroup() {
+  def testDescribeMembersOfNonExistingGroup(): Unit = {
+    val group = "missing.group"
     TestUtils.createOffsetsTopic(zkClient, servers)
 
     // run one consumer in the group consuming from a single-partition topic
     addConsumerGroupExecutor(numConsumers = 1)
     // note the group to be queried is a different (non-existing) group
-    val cgcArgs = Array("--bootstrap-server", brokerList, "--describe", "--group", "missing.group")
+    val cgcArgs = Array("--bootstrap-server", brokerList, "--describe", "--group", group)
     val service = getConsumerGroupService(cgcArgs)
 
-    val (state, assignments) = service.collectGroupMembers(false)
+    val (state, assignments) = service.collectGroupMembers(group, false)
     assertTrue(s"Expected the state to be 'Dead', with no members in the group '$group'.",
         state.contains("Dead") && assignments.contains(List()))
 
-    val (state2, assignments2) = service.collectGroupMembers(true)
+    val (state2, assignments2) = service.collectGroupMembers(group, true)
     assertTrue(s"Expected the state to be 'Dead', with no members in the group '$group' (verbose option).",
         state2.contains("Dead") && assignments2.contains(List()))
   }
 
   @Test
-  def testDescribeStateOfNonExistingGroup() {
+  def testDescribeStateOfNonExistingGroup(): Unit = {
+    val group = "missing.group"
     TestUtils.createOffsetsTopic(zkClient, servers)
 
     // run one consumer in the group consuming from a single-partition topic
     addConsumerGroupExecutor(numConsumers = 1)
     // note the group to be queried is a different (non-existing) group
-    val cgcArgs = Array("--bootstrap-server", brokerList, "--describe", "--group", "missing.group")
+    val cgcArgs = Array("--bootstrap-server", brokerList, "--describe", "--group", group)
     val service = getConsumerGroupService(cgcArgs)
 
-    val state = service.collectGroupState()
+    val state = service.collectGroupState(group)
     assertTrue(s"Expected the state to be 'Dead', with no members in the group '$group'.",
         state.state == "Dead" && state.numMembers == 0 &&
         state.coordinator != null && servers.map(_.config.brokerId).toList.contains(state.coordinator.id)
@@ -111,7 +114,7 @@ class DescribeConsumerGroupTest extends ConsumerGroupCommandTest {
   }
 
   @Test
-  def testDescribeExistingGroup() {
+  def testDescribeExistingGroup(): Unit = {
     TestUtils.createOffsetsTopic(zkClient, servers)
 
     for (describeType <- describeTypes) {
@@ -122,14 +125,63 @@ class DescribeConsumerGroupTest extends ConsumerGroupCommandTest {
       val service = getConsumerGroupService(cgcArgs)
 
       TestUtils.waitUntilTrue(() => {
-        val (output, error) = TestUtils.grabConsoleOutputAndError(service.describeGroup())
+        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(" ")}.", maxRetries = 3)
+      }, s"Expected a data row and no error in describe results with describe type ${describeType.mkString(" ")}.")
+    }
+  }
+
+  @Test
+  def testDescribeExistingGroups(): Unit = {
+    TestUtils.createOffsetsTopic(zkClient, servers)
+
+    // 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)
+      Array("--group", group)
+    }).flatten
+
+    val expectedNumLines = describeTypes.length * 2
+
+    for (describeType <- describeTypes) {
+      val cgcArgs = Array("--bootstrap-server", brokerList, "--describe") ++ groups ++ describeType
+      val service = getConsumerGroupService(cgcArgs)
+
+      TestUtils.waitUntilTrue(() => {
+        val (output, error) = TestUtils.grabConsoleOutputAndError(service.describeGroups())
+        val numLines = output.trim.split("\n").filterNot(line => line.isEmpty).length
+        (numLines == expectedNumLines) && error.isEmpty
+      }, s"Expected a data row and no error in describe results with describe type ${describeType.mkString(" ")}.")
+    }
+  }
+
+  @Test
+  def testDescribeAllExistingGroups(): Unit = {
+    TestUtils.createOffsetsTopic(zkClient, servers)
+
+    // 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)
+    }
+
+    val expectedNumLines = describeTypes.length * 2
+
+    for (describeType <- describeTypes) {
+      val cgcArgs = Array("--bootstrap-server", brokerList, "--describe", "--all-groups") ++ describeType
+      val service = getConsumerGroupService(cgcArgs)
+
+      TestUtils.waitUntilTrue(() => {
+        val (output, error) = TestUtils.grabConsoleOutputAndError(service.describeGroups())
+        val numLines = output.trim.split("\n").filterNot(line => line.isEmpty).length
+        (numLines == expectedNumLines) && error.isEmpty
+      }, s"Expected a data row and no error in describe results with describe type ${describeType.mkString(" ")}.")
     }
   }
 
   @Test
-  def testDescribeOffsetsOfExistingGroup() {
+  def testDescribeOffsetsOfExistingGroup(): Unit = {
     TestUtils.createOffsetsTopic(zkClient, servers)
 
     // run one consumer in the group consuming from a single-partition topic
@@ -139,18 +191,18 @@ class DescribeConsumerGroupTest extends ConsumerGroupCommandTest {
     val service = getConsumerGroupService(cgcArgs)
 
     TestUtils.waitUntilTrue(() => {
-      val (state, assignments) = service.collectGroupOffsets()
+      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.", maxRetries = 3)
+    }, s"Expected a 'Stable' group status, rows and valid values for consumer id / client id / host columns in describe results for group $group.")
   }
 
   @Test
-  def testDescribeMembersOfExistingGroup() {
+  def testDescribeMembersOfExistingGroup(): Unit = {
     TestUtils.createOffsetsTopic(zkClient, servers)
 
     // run one consumer in the group consuming from a single-partition topic
@@ -159,7 +211,7 @@ class DescribeConsumerGroupTest extends ConsumerGroupCommandTest {
     val service = getConsumerGroupService(cgcArgs)
 
     TestUtils.waitUntilTrue(() => {
-      val (state, assignments) = service.collectGroupMembers(false)
+      val (state, assignments) = service.collectGroupMembers(group, false)
       state.contains("Stable") &&
         (assignments match {
           case Some(memberAssignments) =>
@@ -170,9 +222,9 @@ class DescribeConsumerGroupTest extends ConsumerGroupCommandTest {
           case None =>
             false
         })
-    }, s"Expected a 'Stable' group status, rows and valid member information for group $group.", maxRetries = 3)
+    }, s"Expected a 'Stable' group status, rows and valid member information for group $group.")
 
-    val (_, assignments) = service.collectGroupMembers(true)
+    val (_, assignments) = service.collectGroupMembers(group, true)
     assignments match {
       case None =>
         fail(s"Expected partition assignments for members of group $group")
@@ -184,7 +236,7 @@ class DescribeConsumerGroupTest extends ConsumerGroupCommandTest {
   }
 
   @Test
-  def testDescribeStateOfExistingGroup() {
+  def testDescribeStateOfExistingGroup(): Unit = {
     TestUtils.createOffsetsTopic(zkClient, servers)
 
     // run one consumer in the group consuming from a single-partition topic
@@ -193,17 +245,17 @@ class DescribeConsumerGroupTest extends ConsumerGroupCommandTest {
     val service = getConsumerGroupService(cgcArgs)
 
     TestUtils.waitUntilTrue(() => {
-      val state = service.collectGroupState()
+      val state = service.collectGroupState(group)
       state.state == "Stable" &&
         state.numMembers == 1 &&
         state.assignmentStrategy == "range" &&
         state.coordinator != null &&
         servers.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.", maxRetries = 3)
+    }, s"Expected a 'Stable' group status, with one member and round robin assignment strategy for group $group.")
   }
 
   @Test
-  def testDescribeStateOfExistingGroupWithRoundRobinAssignor() {
+  def testDescribeStateOfExistingGroupWithRoundRobinAssignor(): Unit = {
     TestUtils.createOffsetsTopic(zkClient, servers)
 
     // run one consumer in the group consuming from a single-partition topic
@@ -212,17 +264,17 @@ class DescribeConsumerGroupTest extends ConsumerGroupCommandTest {
     val service = getConsumerGroupService(cgcArgs)
 
     TestUtils.waitUntilTrue(() => {
-      val state = service.collectGroupState()
+      val state = service.collectGroupState(group)
       state.state == "Stable" &&
         state.numMembers == 1 &&
         state.assignmentStrategy == "roundrobin" &&
         state.coordinator != null &&
         servers.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.", maxRetries = 3)
+    }, s"Expected a 'Stable' group status, with one member and round robin assignment strategy for group $group.")
   }
 
   @Test
-  def testDescribeExistingGroupWithNoMembers() {
+  def testDescribeExistingGroupWithNoMembers(): Unit = {
     TestUtils.createOffsetsTopic(zkClient, servers)
 
     for (describeType <- describeTypes) {
@@ -233,20 +285,20 @@ class DescribeConsumerGroupTest extends ConsumerGroupCommandTest {
       val service = getConsumerGroupService(cgcArgs)
 
       TestUtils.waitUntilTrue(() => {
-        val (output, error) = TestUtils.grabConsoleOutputAndError(service.describeGroup())
+        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(" ")}'", maxRetries = 3)
+      }, 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.describeGroup()).contains(s"Consumer group '$group' has no active members.")
-      }, s"Expected no active member in describe group results with describe type ${describeType.mkString(" ")}", maxRetries = 3)
+        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(" ")}")
     }
   }
 
   @Test
-  def testDescribeOffsetsOfExistingGroupWithNoMembers() {
+  def testDescribeOffsetsOfExistingGroupWithNoMembers(): Unit = {
     TestUtils.createOffsetsTopic(zkClient, servers)
 
     // run one consumer in the group consuming from a single-partition topic
@@ -256,14 +308,14 @@ class DescribeConsumerGroupTest extends ConsumerGroupCommandTest {
     val service = getConsumerGroupService(cgcArgs)
 
     TestUtils.waitUntilTrue(() => {
-      val (state, assignments) = service.collectGroupOffsets()
+      val (state, assignments) = service.collectGroupOffsets(group)
       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.", maxRetries = 3)
+    }, "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()) {
+    val (result, succeeded) = TestUtils.computeUntilTrue(service.collectGroupOffsets(group)) {
       case (state, assignments) =>
         val testGroupAssignments = assignments.toSeq.flatMap(_.filter(_.group == group))
         def assignment = testGroupAssignments.head
@@ -279,7 +331,7 @@ class DescribeConsumerGroupTest extends ConsumerGroupCommandTest {
   }
 
   @Test
-  def testDescribeMembersOfExistingGroupWithNoMembers() {
+  def testDescribeMembersOfExistingGroupWithNoMembers(): Unit = {
     TestUtils.createOffsetsTopic(zkClient, servers)
 
     // run one consumer in the group consuming from a single-partition topic
@@ -289,21 +341,21 @@ class DescribeConsumerGroupTest extends ConsumerGroupCommandTest {
     val service = getConsumerGroupService(cgcArgs)
 
     TestUtils.waitUntilTrue(() => {
-      val (state, assignments) = service.collectGroupMembers(false)
+      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.", maxRetries = 3)
+    }, "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(false)
+      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'", maxRetries = 3)
+    }, s"Expected no member in describe group members results for group '$group'")
   }
 
   @Test
-  def testDescribeStateOfExistingGroupWithNoMembers() {
+  def testDescribeStateOfExistingGroupWithNoMembers(): Unit = {
     TestUtils.createOffsetsTopic(zkClient, servers)
 
     // run one consumer in the group consuming from a single-partition topic
@@ -313,24 +365,24 @@ class DescribeConsumerGroupTest extends ConsumerGroupCommandTest {
     val service = getConsumerGroupService(cgcArgs)
 
     TestUtils.waitUntilTrue(() => {
-      val state = service.collectGroupState()
+      val state = service.collectGroupState(group)
       state.state == "Stable" &&
         state.numMembers == 1 &&
         state.coordinator != null &&
         servers.map(_.config.brokerId).toList.contains(state.coordinator.id)
-    }, s"Expected the group '$group' to initially become stable, and have a single member.", maxRetries = 3)
+    }, 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()
+      val state = service.collectGroupState(group)
       state.state == "Empty" && state.numMembers == 0 && state.assignmentStrategy == ""
-    }, s"Expected the group '$group' to become empty after the only member leaving.", maxRetries = 3)
+    }, s"Expected the group '$group' to become empty after the only member leaving.")
   }
 
   @Test
-  def testDescribeWithConsumersWithoutAssignedPartitions() {
+  def testDescribeWithConsumersWithoutAssignedPartitions(): Unit = {
     TestUtils.createOffsetsTopic(zkClient, servers)
 
     for (describeType <- describeTypes) {
@@ -341,15 +393,15 @@ class DescribeConsumerGroupTest extends ConsumerGroupCommandTest {
       val service = getConsumerGroupService(cgcArgs)
 
       TestUtils.waitUntilTrue(() => {
-        val (output, error) = TestUtils.grabConsoleOutputAndError(service.describeGroup())
+        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(" ")}'", maxRetries = 3)
+      }, s"Expected a single data row in describe group result with describe type '${describeType.mkString(" ")}'")
     }
   }
 
   @Test
-  def testDescribeOffsetsWithConsumersWithoutAssignedPartitions() {
+  def testDescribeOffsetsWithConsumersWithoutAssignedPartitions(): Unit = {
     TestUtils.createOffsetsTopic(zkClient, servers)
 
     // run two consumers in the group consuming from a single-partition topic
@@ -359,16 +411,16 @@ class DescribeConsumerGroupTest extends ConsumerGroupCommandTest {
     val service = getConsumerGroupService(cgcArgs)
 
     TestUtils.waitUntilTrue(() => {
-      val (state, assignments) = service.collectGroupOffsets()
+      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", maxRetries = 3)
+    }, "Expected rows for consumers with no assigned partitions in describe group results")
   }
 
   @Test
-  def testDescribeMembersWithConsumersWithoutAssignedPartitions() {
+  def testDescribeMembersWithConsumersWithoutAssignedPartitions(): Unit = {
     TestUtils.createOffsetsTopic(zkClient, servers)
 
     // run two consumers in the group consuming from a single-partition topic
@@ -378,22 +430,22 @@ class DescribeConsumerGroupTest extends ConsumerGroupCommandTest {
     val service = getConsumerGroupService(cgcArgs)
 
     TestUtils.waitUntilTrue(() => {
-      val (state, assignments) = service.collectGroupMembers(false)
+      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.count(_.assignment.nonEmpty) == 0
-    }, "Expected rows for consumers with no assigned partitions in describe group results", maxRetries = 3)
+    }, "Expected rows for consumers with no assigned partitions in describe group results")
 
-    val (state, assignments) = service.collectGroupMembers(true)
+    val (state, assignments) = service.collectGroupMembers(group, true)
     assertTrue("Expected additional columns in verbose version of describe members",
         state.contains("Stable") && assignments.get.count(_.assignment.nonEmpty) > 0)
   }
 
   @Test
-  def testDescribeStateWithConsumersWithoutAssignedPartitions() {
+  def testDescribeStateWithConsumersWithoutAssignedPartitions(): Unit = {
     TestUtils.createOffsetsTopic(zkClient, servers)
 
     // run two consumers in the group consuming from a single-partition topic
@@ -403,13 +455,13 @@ class DescribeConsumerGroupTest extends ConsumerGroupCommandTest {
     val service = getConsumerGroupService(cgcArgs)
 
     TestUtils.waitUntilTrue(() => {
-      val state = service.collectGroupState()
+      val state = service.collectGroupState(group)
       state.state == "Stable" && state.numMembers == 2
-    }, "Expected two consumers in describe group results", maxRetries = 3)
+    }, "Expected two consumers in describe group results")
   }
 
   @Test
-  def testDescribeWithMultiPartitionTopicAndMultipleConsumers() {
+  def testDescribeWithMultiPartitionTopicAndMultipleConsumers(): Unit = {
     TestUtils.createOffsetsTopic(zkClient, servers)
     val topic2 = "foo2"
     createTopic(topic2, 2, 1)
@@ -422,15 +474,15 @@ class DescribeConsumerGroupTest extends ConsumerGroupCommandTest {
       val service = getConsumerGroupService(cgcArgs)
 
       TestUtils.waitUntilTrue(() => {
-        val (output, error) = TestUtils.grabConsoleOutputAndError(service.describeGroup())
+        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(" ")}'", maxRetries = 3)
+      }, s"Expected a single data row in describe group result with describe type '${describeType.mkString(" ")}'")
     }
   }
 
   @Test
-  def testDescribeOffsetsWithMultiPartitionTopicAndMultipleConsumers() {
+  def testDescribeOffsetsWithMultiPartitionTopicAndMultipleConsumers(): Unit = {
     TestUtils.createOffsetsTopic(zkClient, servers)
     val topic2 = "foo2"
     createTopic(topic2, 2, 1)
@@ -442,17 +494,17 @@ class DescribeConsumerGroupTest extends ConsumerGroupCommandTest {
     val service = getConsumerGroupService(cgcArgs)
 
     TestUtils.waitUntilTrue(() => {
-      val (state, assignments) = service.collectGroupOffsets()
+      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.", maxRetries = 3)
+    }, "Expected two rows (one row per consumer) in describe group results.")
   }
 
   @Test
-  def testDescribeMembersWithMultiPartitionTopicAndMultipleConsumers() {
+  def testDescribeMembersWithMultiPartitionTopicAndMultipleConsumers(): Unit = {
     TestUtils.createOffsetsTopic(zkClient, servers)
     val topic2 = "foo2"
     createTopic(topic2, 2, 1)
@@ -464,21 +516,21 @@ class DescribeConsumerGroupTest extends ConsumerGroupCommandTest {
     val service = getConsumerGroupService(cgcArgs)
 
     TestUtils.waitUntilTrue(() => {
-      val (state, assignments) = service.collectGroupMembers(false)
+      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.", maxRetries = 3)
+    }, "Expected two rows (one row per consumer) in describe group members results.")
 
-    val (state, assignments) = service.collectGroupMembers(true)
+    val (state, assignments) = service.collectGroupMembers(group, true)
     assertTrue("Expected additional columns in verbose version of describe members",
         state.contains("Stable") && assignments.get.count(_.assignment.isEmpty) == 0)
   }
 
   @Test
-  def testDescribeStateWithMultiPartitionTopicAndMultipleConsumers() {
+  def testDescribeStateWithMultiPartitionTopicAndMultipleConsumers(): Unit = {
     TestUtils.createOffsetsTopic(zkClient, servers)
     val topic2 = "foo2"
     createTopic(topic2, 2, 1)
@@ -490,13 +542,13 @@ class DescribeConsumerGroupTest extends ConsumerGroupCommandTest {
     val service = getConsumerGroupService(cgcArgs)
 
     TestUtils.waitUntilTrue(() => {
-      val state = service.collectGroupState()
+      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.", maxRetries = 3)
+    }, "Expected a stable group with two members in describe group state result.")
   }
 
   @Test
-  def testDescribeSimpleConsumerGroup() {
+  def testDescribeSimpleConsumerGroup(): Unit = {
     // Ensure that the offsets of consumers which don't use group management are still displayed
 
     TestUtils.createOffsetsTopic(zkClient, servers)
@@ -508,13 +560,13 @@ class DescribeConsumerGroupTest extends ConsumerGroupCommandTest {
     val service = getConsumerGroupService(cgcArgs)
 
     TestUtils.waitUntilTrue(() => {
-      val (state, assignments) = service.collectGroupOffsets()
+      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.", maxRetries = 3)
+    }, "Expected a stable group with two members in describe group state result.")
   }
 
   @Test
-  def testDescribeGroupWithShortInitializationTimeout() {
+  def testDescribeGroupWithShortInitializationTimeout(): Unit = {
     // Let creation of the offsets topic happen during group initialization to ensure that initialization doesn't
     // complete before the timeout expires
 
@@ -527,7 +579,7 @@ class DescribeConsumerGroupTest extends ConsumerGroupCommandTest {
     val service = getConsumerGroupService(cgcArgs)
 
     try {
-      TestUtils.grabConsoleOutputAndError(service.describeGroup())
+      TestUtils.grabConsoleOutputAndError(service.describeGroups())
       fail(s"The consumer group command should have failed due to low initialization timeout (describe type: ${describeType.mkString(" ")})")
     } catch {
       case e: ExecutionException => assert(e.getCause.isInstanceOf[TimeoutException]) // OK
@@ -535,7 +587,7 @@ class DescribeConsumerGroupTest extends ConsumerGroupCommandTest {
   }
 
   @Test
-  def testDescribeGroupOffsetsWithShortInitializationTimeout() {
+  def testDescribeGroupOffsetsWithShortInitializationTimeout(): Unit = {
     // Let creation of the offsets topic happen during group initialization to ensure that initialization doesn't
     // complete before the timeout expires
 
@@ -547,7 +599,7 @@ class DescribeConsumerGroupTest extends ConsumerGroupCommandTest {
     val service = getConsumerGroupService(cgcArgs)
 
     try {
-      service.collectGroupOffsets()
+      service.collectGroupOffsets(group)
       fail("The consumer group command should fail due to low initialization timeout")
     } catch {
       case e : ExecutionException => assert(e.getCause.isInstanceOf[TimeoutException]) // OK
@@ -555,7 +607,7 @@ class DescribeConsumerGroupTest extends ConsumerGroupCommandTest {
   }
 
   @Test
-  def testDescribeGroupMembersWithShortInitializationTimeout() {
+  def testDescribeGroupMembersWithShortInitializationTimeout(): Unit = {
     // Let creation of the offsets topic happen during group initialization to ensure that initialization doesn't
     // complete before the timeout expires
 
@@ -567,12 +619,12 @@ class DescribeConsumerGroupTest extends ConsumerGroupCommandTest {
     val service = getConsumerGroupService(cgcArgs)
 
     try {
-      service.collectGroupMembers(false)
+      service.collectGroupMembers(group, false)
       fail("The consumer group command should fail due to low initialization timeout")
     } catch {
       case e: ExecutionException => assert(e.getCause.isInstanceOf[TimeoutException])// OK
         try {
-          service.collectGroupMembers(true)
+          service.collectGroupMembers(group, true)
           fail("The consumer group command should fail due to low initialization timeout (verbose)")
         } catch {
           case e: ExecutionException => assert(e.getCause.isInstanceOf[TimeoutException]) // OK
@@ -581,7 +633,7 @@ class DescribeConsumerGroupTest extends ConsumerGroupCommandTest {
   }
 
   @Test
-  def testDescribeGroupStateWithShortInitializationTimeout() {
+  def testDescribeGroupStateWithShortInitializationTimeout(): Unit = {
     // Let creation of the offsets topic happen during group initialization to ensure that initialization doesn't
     // complete before the timeout expires
 
@@ -593,7 +645,7 @@ class DescribeConsumerGroupTest extends ConsumerGroupCommandTest {
     val service = getConsumerGroupService(cgcArgs)
 
     try {
-      service.collectGroupState()
+      service.collectGroupState(group)
       fail("The consumer group command should fail due to low initialization timeout")
     } catch {
       case e: ExecutionException => assert(e.getCause.isInstanceOf[TimeoutException]) // OK
@@ -601,14 +653,14 @@ class DescribeConsumerGroupTest extends ConsumerGroupCommandTest {
   }
 
   @Test(expected = classOf[joptsimple.OptionException])
-  def testDescribeWithUnrecognizedNewConsumerOption() {
+  def testDescribeWithUnrecognizedNewConsumerOption(): Unit = {
     val cgcArgs = Array("--new-consumer", "--bootstrap-server", brokerList, "--describe", "--group", group)
     getConsumerGroupService(cgcArgs)
     fail("Expected an error due to presence of unrecognized --new-consumer option")
   }
 
   @Test
-  def testDescribeNonOffsetCommitGroup() {
+  def testDescribeNonOffsetCommitGroup(): Unit = {
     TestUtils.createOffsetsTopic(zkClient, servers)
 
     val customProps = new Properties
@@ -621,7 +673,7 @@ class DescribeConsumerGroupTest extends ConsumerGroupCommandTest {
     val service = getConsumerGroupService(cgcArgs)
 
     TestUtils.waitUntilTrue(() => {
-      val (state, assignments) = service.collectGroupOffsets()
+      val (state, assignments) = service.collectGroupOffsets(group)
       state.contains("Stable") &&
         assignments.isDefined &&
         assignments.get.count(_.group == group) == 1 &&
diff --git a/core/src/test/scala/unit/kafka/admin/LeaderElectionCommandTest.scala b/core/src/test/scala/unit/kafka/admin/LeaderElectionCommandTest.scala
new file mode 100644
index 0000000000000..45fb7b9ebc3e8
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/admin/LeaderElectionCommandTest.scala
@@ -0,0 +1,333 @@
+/**
+ * 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.File
+import java.nio.charset.StandardCharsets
+import java.nio.file.Files
+import java.nio.file.Path
+
+import kafka.common.AdminCommandFailedException
+import kafka.server.KafkaConfig
+import kafka.server.KafkaServer
+import kafka.utils.TestUtils
+import kafka.zk.ZooKeeperTestHarness
+import org.apache.kafka.clients.admin.{AdminClient, AdminClientConfig}
+import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.errors.TimeoutException
+import org.apache.kafka.common.errors.UnknownTopicOrPartitionException
+import org.apache.kafka.common.network.ListenerName
+import org.junit.After
+import org.junit.Assert._
+import org.junit.Before
+import org.junit.Test
+
+import scala.collection.JavaConverters._
+import scala.collection.Seq
+import scala.concurrent.duration._
+
+final class LeaderElectionCommandTest extends ZooKeeperTestHarness {
+  import LeaderElectionCommandTest._
+
+  var servers = Seq.empty[KafkaServer]
+  val broker1 = 0
+  val broker2 = 1
+  val broker3 = 2
+
+  @Before
+  override def setUp(): Unit = {
+    super.setUp()
+
+    val brokerConfigs = TestUtils.createBrokerConfigs(3, zkConnect, false)
+    servers = brokerConfigs.map { config =>
+      config.setProperty("auto.leader.rebalance.enable", "false")
+      config.setProperty("controlled.shutdown.enable", "true")
+      config.setProperty("controlled.shutdown.max.retries", "1")
+      config.setProperty("controlled.shutdown.retry.backoff.ms", "1000")
+      TestUtils.createServer(KafkaConfig.fromProps(config))
+    }
+  }
+
+  @After
+  override def tearDown(): Unit = {
+    TestUtils.shutdownServers(servers)
+
+    super.tearDown()
+  }
+
+  @Test
+  def testAllTopicPartition(): Unit = {
+    TestUtils.resource(AdminClient.create(createConfig(servers).asJava)) { client =>
+      val topic = "unclean-topic"
+      val partition = 0
+      val assignment = Seq(broker2, broker3)
+
+      TestUtils.createTopic(zkClient, topic, Map(partition -> assignment), servers)
+
+      val topicPartition = new TopicPartition(topic, partition)
+
+      TestUtils.assertLeader(client, topicPartition, broker2)
+
+      servers(broker3).shutdown()
+      TestUtils.waitForBrokersOutOfIsr(client, Set(topicPartition), Set(broker3))
+      servers(broker2).shutdown()
+      TestUtils.assertNoLeader(client, topicPartition)
+      servers(broker3).startup()
+
+      LeaderElectionCommand.main(
+        Array(
+          "--bootstrap-server", bootstrapServers(servers),
+          "--election-type", "unclean",
+          "--all-topic-partitions"
+        )
+      )
+
+      TestUtils.assertLeader(client, topicPartition, broker3)
+    }
+  }
+
+  @Test
+  def testTopicPartition(): Unit = {
+    TestUtils.resource(AdminClient.create(createConfig(servers).asJava)) { client =>
+      val topic = "unclean-topic"
+      val partition = 0
+      val assignment = Seq(broker2, broker3)
+
+      TestUtils.createTopic(zkClient, topic, Map(partition -> assignment), servers)
+
+      val topicPartition = new TopicPartition(topic, partition)
+
+      TestUtils.assertLeader(client, topicPartition, broker2)
+
+      servers(broker3).shutdown()
+      TestUtils.waitForBrokersOutOfIsr(client, Set(topicPartition), Set(broker3))
+      servers(broker2).shutdown()
+      TestUtils.assertNoLeader(client, topicPartition)
+      servers(broker3).startup()
+
+      LeaderElectionCommand.main(
+        Array(
+          "--bootstrap-server", bootstrapServers(servers),
+          "--election-type", "unclean",
+          "--topic", topic,
+          "--partition", partition.toString
+        )
+      )
+
+      TestUtils.assertLeader(client, topicPartition, broker3)
+    }
+  }
+
+  @Test
+  def testPathToJsonFile(): Unit = {
+    TestUtils.resource(AdminClient.create(createConfig(servers).asJava)) { client =>
+      val topic = "unclean-topic"
+      val partition = 0
+      val assignment = Seq(broker2, broker3)
+
+      TestUtils.createTopic(zkClient, topic, Map(partition -> assignment), servers)
+
+      val topicPartition = new TopicPartition(topic, partition)
+
+      TestUtils.assertLeader(client, topicPartition, broker2)
+
+      servers(broker3).shutdown()
+      TestUtils.waitForBrokersOutOfIsr(client, Set(topicPartition), Set(broker3))
+      servers(broker2).shutdown()
+      TestUtils.assertNoLeader(client, topicPartition)
+      servers(broker3).startup()
+
+      val topicPartitionPath = tempTopicPartitionFile(Set(topicPartition))
+
+      LeaderElectionCommand.main(
+        Array(
+          "--bootstrap-server", bootstrapServers(servers),
+          "--election-type", "unclean",
+          "--path-to-json-file", topicPartitionPath.toString
+        )
+      )
+
+      TestUtils.assertLeader(client, topicPartition, broker3)
+    }
+  }
+
+  @Test
+  def testPreferredReplicaElection(): Unit = {
+    TestUtils.resource(AdminClient.create(createConfig(servers).asJava)) { client =>
+      val topic = "unclean-topic"
+      val partition = 0
+      val assignment = Seq(broker2, broker3)
+
+      TestUtils.createTopic(zkClient, topic, Map(partition -> assignment), servers)
+
+      val topicPartition = new TopicPartition(topic, partition)
+
+      TestUtils.assertLeader(client, topicPartition, broker2)
+
+      servers(broker2).shutdown()
+      TestUtils.assertLeader(client, topicPartition, broker3)
+      servers(broker2).startup()
+      TestUtils.waitForBrokersInIsr(client, topicPartition, Set(broker2))
+
+      LeaderElectionCommand.main(
+        Array(
+          "--bootstrap-server", bootstrapServers(servers),
+          "--election-type", "preferred",
+          "--all-topic-partitions"
+        )
+      )
+
+      TestUtils.assertLeader(client, topicPartition, broker2)
+    }
+  }
+
+  @Test
+  def testTopicWithoutPartition(): Unit = {
+    try {
+      LeaderElectionCommand.main(
+        Array(
+          "--bootstrap-server", bootstrapServers(servers),
+          "--election-type", "unclean",
+          "--topic", "some-topic"
+        )
+      )
+      fail()
+    } catch {
+      case e: Throwable =>
+        assertTrue(e.getMessage.startsWith("Missing required option(s)"))
+        assertTrue(e.getMessage.contains(" partition"))
+    }
+  }
+
+  @Test
+  def testPartitionWithoutTopic(): Unit = {
+    try {
+      LeaderElectionCommand.main(
+        Array(
+          "--bootstrap-server", bootstrapServers(servers),
+          "--election-type", "unclean",
+          "--all-topic-partitions",
+          "--partition", "0"
+        )
+      )
+      fail()
+    } catch {
+      case e: Throwable =>
+        assertEquals("Option partition is only allowed if topic is used", e.getMessage)
+    }
+  }
+
+  @Test
+  def testTopicDoesNotExist(): Unit = {
+    try {
+      LeaderElectionCommand.main(
+        Array(
+          "--bootstrap-server", bootstrapServers(servers),
+          "--election-type", "preferred",
+          "--topic", "unknown-topic-name",
+          "--partition", "0"
+        )
+      )
+      fail()
+    } catch {
+      case e: AdminCommandFailedException =>
+        assertTrue(e.getSuppressed()(0).isInstanceOf[UnknownTopicOrPartitionException])
+    }
+  }
+
+  @Test
+  def testMissingElectionType(): Unit = {
+    try {
+      LeaderElectionCommand.main(
+        Array(
+          "--bootstrap-server", bootstrapServers(servers),
+          "--topic", "some-topic",
+          "--partition", "0"
+        )
+      )
+      fail()
+    } catch {
+      case e: Throwable =>
+        assertTrue(e.getMessage.startsWith("Missing required option(s)"))
+        assertTrue(e.getMessage.contains(" election-type"))
+    }
+  }
+
+  @Test
+  def testMissingTopicPartitionSelection(): Unit = {
+    try {
+      LeaderElectionCommand.main(
+        Array(
+          "--bootstrap-server", bootstrapServers(servers),
+          "--election-type", "preferred"
+        )
+      )
+      fail()
+    } catch {
+      case e: Throwable =>
+        assertTrue(e.getMessage.startsWith("One and only one of the following options is required: "))
+        assertTrue(e.getMessage.contains(" all-topic-partitions"))
+        assertTrue(e.getMessage.contains(" topic"))
+        assertTrue(e.getMessage.contains(" path-to-json-file"))
+    }
+  }
+
+  @Test
+  def testInvalidBroker(): Unit = {
+    try {
+      LeaderElectionCommand.run(
+        Array(
+          "--bootstrap-server", "example.com:1234",
+          "--election-type", "unclean",
+          "--all-topic-partitions"
+        ),
+        1.seconds
+      )
+      fail()
+    } catch {
+      case e: AdminCommandFailedException =>
+        assertTrue(e.getCause.isInstanceOf[TimeoutException])
+    }
+  }
+}
+
+object LeaderElectionCommandTest {
+  def createConfig(servers: Seq[KafkaServer]): Map[String, Object] = {
+    Map(
+      AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG -> bootstrapServers(servers),
+      AdminClientConfig.REQUEST_TIMEOUT_MS_CONFIG -> "20000"
+    )
+  }
+
+  def bootstrapServers(servers: Seq[KafkaServer]): String = {
+    servers.map { server =>
+      val port = server.socketServer.boundPort(ListenerName.normalised("PLAINTEXT"))
+      s"localhost:$port"
+    }.headOption.mkString(",")
+  }
+
+  def tempTopicPartitionFile(partitions: Set[TopicPartition]): Path = {
+    val file = File.createTempFile("leader-election-command", ".json")
+    file.deleteOnExit()
+
+    val jsonString = TestUtils.stringifyTopicPartitions(partitions)
+
+    Files.write(file.toPath, jsonString.getBytes(StandardCharsets.UTF_8))
+
+    file.toPath
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/admin/ListConsumerGroupTest.scala b/core/src/test/scala/unit/kafka/admin/ListConsumerGroupTest.scala
index 1a35c4ccd502e..7429a43a1f102 100644
--- a/core/src/test/scala/unit/kafka/admin/ListConsumerGroupTest.scala
+++ b/core/src/test/scala/unit/kafka/admin/ListConsumerGroupTest.scala
@@ -23,7 +23,7 @@ import kafka.utils.TestUtils
 class ListConsumerGroupTest extends ConsumerGroupCommandTest {
 
   @Test
-  def testListConsumerGroups() {
+  def testListConsumerGroups(): Unit = {
     val simpleGroup = "simple-group"
     addSimpleGroupExecutor(group = simpleGroup)
     addConsumerGroupExecutor(numConsumers = 1)
@@ -36,11 +36,11 @@ class ListConsumerGroupTest extends ConsumerGroupCommandTest {
     TestUtils.waitUntilTrue(() => {
       foundGroups = service.listGroups().toSet
       expectedGroups == foundGroups
-    }, s"Expected --list to show groups $expectedGroups, but found $foundGroups.", maxRetries = 3)
+    }, s"Expected --list to show groups $expectedGroups, but found $foundGroups.")
   }
 
   @Test(expected = classOf[OptionException])
-  def testListWithUnrecognizedNewConsumerOption() {
+  def testListWithUnrecognizedNewConsumerOption(): Unit = {
     val cgcArgs = Array("--new-consumer", "--bootstrap-server", brokerList, "--list")
     getConsumerGroupService(cgcArgs)
   }
diff --git a/core/src/test/scala/unit/kafka/admin/PreferredReplicaElectionCommandTest.scala b/core/src/test/scala/unit/kafka/admin/PreferredReplicaElectionCommandTest.scala
deleted file mode 100644
index cf752b8b42f76..0000000000000
--- a/core/src/test/scala/unit/kafka/admin/PreferredReplicaElectionCommandTest.scala
+++ /dev/null
@@ -1,70 +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.server.{KafkaConfig, KafkaServer}
-import kafka.utils.{Logging, TestUtils}
-import kafka.zk.ZooKeeperTestHarness
-import org.apache.kafka.common.TopicPartition
-import org.junit.Assert.assertEquals
-import org.junit.{After, Test}
-
-import scala.collection.{Map, Set}
-
-class PreferredReplicaElectionCommandTest extends ZooKeeperTestHarness with Logging {
-  var servers: Seq[KafkaServer] = Seq()
-
-  @After
-  override def tearDown() {
-    TestUtils.shutdownServers(servers)
-    super.tearDown()
-  }
-
-  @Test
-  def testPreferredReplicaJsonData() {
-    // write preferred replica json data to zk path
-    val partitionsForPreferredReplicaElection = Set(new TopicPartition("test", 1), new TopicPartition("test2", 1))
-    PreferredReplicaLeaderElectionCommand.writePreferredReplicaElectionData(zkClient, partitionsForPreferredReplicaElection)
-    // try to read it back and compare with what was written
-    val partitionsUndergoingPreferredReplicaElection = zkClient.getPreferredReplicaElection
-    assertEquals("Preferred replica election ser-de failed", partitionsForPreferredReplicaElection,
-      partitionsUndergoingPreferredReplicaElection)
-  }
-
-  @Test
-  def testBasicPreferredReplicaElection() {
-    val expectedReplicaAssignment = Map(0  -> List(0, 1, 2))
-    val topic = "test"
-    val partition = 0
-    val preferredReplica = 0
-    // create brokers
-    val brokerRack = Map(0 -> "rack0", 1 -> "rack1", 2 -> "rack2")
-    val serverConfigs = TestUtils.createBrokerConfigs(3, zkConnect, false, rackInfo = brokerRack).map(KafkaConfig.fromProps)
-    // create the topic
-    adminZkClient.createTopicWithAssignment(topic, config = new Properties, expectedReplicaAssignment)
-    servers = serverConfigs.reverseMap(s => TestUtils.createServer(s))
-    // broker 2 should be the leader since it was started first
-    val currentLeader = TestUtils.waitUntilLeaderIsElectedOrChanged(zkClient, topic, partition, oldLeaderOpt = None)
-    // trigger preferred replica election
-    val preferredReplicaElection = new PreferredReplicaLeaderElectionCommand(zkClient, Set(new TopicPartition(topic, partition)))
-    preferredReplicaElection.moveLeaderToPreferredReplica()
-    val newLeader = TestUtils.waitUntilLeaderIsElectedOrChanged(zkClient, topic, partition, oldLeaderOpt = Some(currentLeader))
-    assertEquals("Preferred replica election failed", preferredReplica, newLeader)
-  }
-}
diff --git a/core/src/test/scala/unit/kafka/admin/PreferredReplicaLeaderElectionCommandTest.scala b/core/src/test/scala/unit/kafka/admin/PreferredReplicaLeaderElectionCommandTest.scala
index 96e7dace4e7b2..37032f17501a8 100644
--- a/core/src/test/scala/unit/kafka/admin/PreferredReplicaLeaderElectionCommandTest.scala
+++ b/core/src/test/scala/unit/kafka/admin/PreferredReplicaLeaderElectionCommandTest.scala
@@ -19,31 +19,41 @@ package kafka.admin
 import java.io.File
 import java.nio.charset.StandardCharsets
 import java.nio.file.{Files, Paths}
+import java.util
 import java.util.Properties
 
-import kafka.common.{AdminCommandFailedException, TopicAndPartition}
-import kafka.network.RequestChannel
-import kafka.security.auth._
+import scala.collection.Seq
+import kafka.common.AdminCommandFailedException
+import kafka.security.authorizer.AclAuthorizer
 import kafka.server.{KafkaConfig, KafkaServer}
-import kafka.utils.{Logging, TestUtils, ZkUtils}
+import kafka.utils.{Logging, TestUtils}
 import kafka.zk.ZooKeeperTestHarness
 import org.apache.kafka.common.TopicPartition
-import org.apache.kafka.common.errors.{ClusterAuthorizationException, PreferredLeaderNotAvailableException, TimeoutException, UnknownTopicOrPartitionException}
+import org.apache.kafka.common.acl.AclOperation
+import org.apache.kafka.common.errors.ClusterAuthorizationException
+import org.apache.kafka.common.errors.PreferredLeaderNotAvailableException
+import org.apache.kafka.common.errors.TimeoutException
+import org.apache.kafka.common.errors.UnknownTopicOrPartitionException
 import org.apache.kafka.common.network.ListenerName
+import org.apache.kafka.common.resource.ResourceType
+import org.apache.kafka.server.authorizer.{Action, AuthorizableRequestContext, AuthorizationResult}
+import org.apache.kafka.test
 import org.junit.Assert._
 import org.junit.{After, Test}
 
+import scala.collection.JavaConverters._
+
 class PreferredReplicaLeaderElectionCommandTest extends ZooKeeperTestHarness with Logging {
   var servers: Seq[KafkaServer] = Seq()
 
   @After
-  override def tearDown() {
+  override def tearDown(): Unit = {
     TestUtils.shutdownServers(servers)
     super.tearDown()
   }
 
   private def createTestTopicAndCluster(topicPartition: Map[TopicPartition, List[Int]],
-                                        authorizer: Option[String] = None) {
+                                        authorizer: Option[String] = None): Unit = {
 
     val brokerConfigs = TestUtils.createBrokerConfigs(3, zkConnect, false)
     brokerConfigs.foreach(p => p.setProperty("auto.leader.rebalance.enable", "false"))
@@ -52,32 +62,39 @@ class PreferredReplicaLeaderElectionCommandTest extends ZooKeeperTestHarness wit
         brokerConfigs.foreach(p => p.setProperty("authorizer.class.name", className))
       case None =>
     }
-    createTestTopicAndCluster(topicPartition,brokerConfigs)
+    createTestTopicAndCluster(topicPartition, brokerConfigs)
   }
 
   private def createTestTopicAndCluster(partitionsAndAssignments: Map[TopicPartition, List[Int]],
-                                        brokerConfigs: Seq[Properties]) {
+                                        brokerConfigs: Seq[Properties]): Unit = {
     // create brokers
     servers = brokerConfigs.map(b => TestUtils.createServer(KafkaConfig.fromProps(b)))
     // create the topic
-    partitionsAndAssignments.foreach { case (tp, assigment) =>
-      zkClient.createTopicAssignment(tp.topic(),
-      Map(tp -> assigment))
+    partitionsAndAssignments.foreach { case (tp, assignment) =>
+      zkClient.createTopicAssignment(tp.topic,
+      Map(tp -> assignment))
     }
     // wait until replica log is created on every broker
-    TestUtils.waitUntilTrue(() => servers.forall(server => partitionsAndAssignments.forall(partitionAndAssignment => server.getLogManager().getLog(partitionAndAssignment._1).isDefined)),
-      "Replicas for topic test not created")
+    TestUtils.waitUntilTrue(
+      () =>
+        servers.forall { server =>
+          partitionsAndAssignments.forall { partitionAndAssignment =>
+            server.getLogManager().getLog(partitionAndAssignment._1).isDefined
+          }
+        },
+      "Replicas for topic test not created"
+    )
   }
 
   /** Bounce the given targetServer and wait for all servers to get metadata for the given partition */
-  private def bounceServer(targetServer: Int, partition: TopicPartition) {
+  private def bounceServer(targetServer: Int, partition: TopicPartition): Unit = {
     debug(s"Shutting down server $targetServer so a non-preferred replica becomes leader")
     servers(targetServer).shutdown()
     debug(s"Starting server $targetServer now that a non-preferred replica is leader")
     servers(targetServer).startup()
     TestUtils.waitUntilTrue(() => servers.forall { server =>
-      server.metadataCache.getPartitionInfo(partition.topic(), partition.partition()).exists { partitionState =>
-        partitionState.basePartitionState.isr.contains(targetServer)
+      server.metadataCache.getPartitionInfo(partition.topic, partition.partition).exists { partitionState =>
+        partitionState.isr.contains(targetServer)
       }
     },
       s"Replicas for partition $partition not created")
@@ -87,8 +104,10 @@ class PreferredReplicaLeaderElectionCommandTest extends ZooKeeperTestHarness wit
     servers.find(p => p.kafkaController.isActive)
   }
 
-  private def getLeader(topicPartition: TopicPartition) = {
-    servers(0).metadataCache.getPartitionInfo(topicPartition.topic(), topicPartition.partition()).get.basePartitionState.leader
+  private def awaitLeader(topicPartition: TopicPartition, timeoutMs: Long = test.TestUtils.DEFAULT_MAX_WAIT_MS): Int = {
+    TestUtils.awaitValue(() => {
+      servers.head.metadataCache.getPartitionInfo(topicPartition.topic, topicPartition.partition).map(_.leader)
+    }, s"Timed out waiting to find current leader of $topicPartition", timeoutMs)
   }
 
   private def bootstrapServer(broker: Int = 0): String = {
@@ -104,20 +123,20 @@ class PreferredReplicaLeaderElectionCommandTest extends ZooKeeperTestHarness wit
 
   /** Test the case multiple values are given for --bootstrap-broker */
   @Test
-  def testMultipleBrokersGiven() {
+  def testMultipleBrokersGiven(): Unit = {
     createTestTopicAndCluster(testPartitionAndAssignment)
     bounceServer(testPartitionPreferredLeader, testPartition)
     // Check the leader for the partition is not the preferred one
-    assertNotEquals(testPartitionPreferredLeader, getLeader(testPartition))
+    assertNotEquals(testPartitionPreferredLeader, awaitLeader(testPartition))
     PreferredReplicaLeaderElectionCommand.run(Array(
       "--bootstrap-server", s"${bootstrapServer(1)},${bootstrapServer(0)}"))
     // Check the leader for the partition IS the preferred one
-    assertEquals(testPartitionPreferredLeader, getLeader(testPartition))
+    assertEquals(testPartitionPreferredLeader, awaitLeader(testPartition))
   }
 
   /** Test the case when an invalid broker is given for --bootstrap-broker */
   @Test
-  def testInvalidBrokerGiven() {
+  def testInvalidBrokerGiven(): Unit = {
     try {
       PreferredReplicaLeaderElectionCommand.run(Array(
         "--bootstrap-server", "example.com:1234"),
@@ -131,21 +150,21 @@ class PreferredReplicaLeaderElectionCommandTest extends ZooKeeperTestHarness wit
 
   /** Test the case where no partitions are given (=> elect all partitions) */
   @Test
-  def testNoPartitionsGiven() {
+  def testNoPartitionsGiven(): Unit = {
     createTestTopicAndCluster(testPartitionAndAssignment)
     bounceServer(testPartitionPreferredLeader, testPartition)
     // Check the leader for the partition is not the preferred one
-    assertNotEquals(testPartitionPreferredLeader, getLeader(testPartition))
+    assertNotEquals(testPartitionPreferredLeader, awaitLeader(testPartition))
     PreferredReplicaLeaderElectionCommand.run(Array(
       "--bootstrap-server", bootstrapServer()))
     // Check the leader for the partition IS the preferred one
-    assertEquals(testPartitionPreferredLeader, getLeader(testPartition))
+    assertEquals(testPartitionPreferredLeader, awaitLeader(testPartition))
   }
 
-  private def toJsonFile(partitions: scala.collection.Set[TopicPartition]): File = {
+  private def toJsonFile(partitions: Set[TopicPartition]): File = {
     val jsonFile = File.createTempFile("preferredreplicaelection", ".js")
     jsonFile.deleteOnExit()
-    val jsonString = ZkUtils.preferredReplicaLeaderElectionZkData(partitions.map(new TopicAndPartition(_)))
+    val jsonString = TestUtils.stringifyTopicPartitions(partitions)
     debug("Using json: "+jsonString)
     Files.write(Paths.get(jsonFile.getAbsolutePath), jsonString.getBytes(StandardCharsets.UTF_8))
     jsonFile
@@ -153,11 +172,11 @@ class PreferredReplicaLeaderElectionCommandTest extends ZooKeeperTestHarness wit
 
   /** Test the case where a list of partitions is given */
   @Test
-  def testSingletonPartitionGiven() {
+  def testSingletonPartitionGiven(): Unit = {
     createTestTopicAndCluster(testPartitionAndAssignment)
     bounceServer(testPartitionPreferredLeader, testPartition)
     // Check the leader for the partition is not the preferred one
-    assertNotEquals(testPartitionPreferredLeader, getLeader(testPartition))
+    assertNotEquals(testPartitionPreferredLeader, awaitLeader(testPartition))
     val jsonFile = toJsonFile(testPartitionAndAssignment.keySet)
     try {
       PreferredReplicaLeaderElectionCommand.run(Array(
@@ -167,12 +186,12 @@ class PreferredReplicaLeaderElectionCommandTest extends ZooKeeperTestHarness wit
       jsonFile.delete()
     }
     // Check the leader for the partition IS the preferred one
-    assertEquals(testPartitionPreferredLeader, getLeader(testPartition))
+    assertEquals(testPartitionPreferredLeader, awaitLeader(testPartition))
   }
 
   /** Test the case where a topic does not exist */
   @Test
-  def testTopicDoesNotExist() {
+  def testTopicDoesNotExist(): Unit = {
     val nonExistentPartition = new TopicPartition("does.not.exist", 0)
     val nonExistentPartitionAssignment = List(1, 2, 0)
     val nonExistentPartitionAndAssignment = Map(nonExistentPartition -> nonExistentPartitionAssignment)
@@ -197,7 +216,7 @@ class PreferredReplicaLeaderElectionCommandTest extends ZooKeeperTestHarness wit
 
   /** Test the case where several partitions are given */
   @Test
-  def testMultiplePartitionsSameAssignment() {
+  def testMultiplePartitionsSameAssignment(): Unit = {
     val testPartitionA = new TopicPartition("testA", 0)
     val testPartitionB = new TopicPartition("testB", 0)
     val testPartitionAssignment = List(1, 2, 0)
@@ -207,8 +226,8 @@ class PreferredReplicaLeaderElectionCommandTest extends ZooKeeperTestHarness wit
     createTestTopicAndCluster(testPartitionAndAssignment)
     bounceServer(testPartitionPreferredLeader, testPartitionA)
     // Check the leader for the partition is not the preferred one
-    assertNotEquals(testPartitionPreferredLeader, getLeader(testPartitionA))
-    assertNotEquals(testPartitionPreferredLeader, getLeader(testPartitionB))
+    assertNotEquals(testPartitionPreferredLeader, awaitLeader(testPartitionA))
+    assertNotEquals(testPartitionPreferredLeader, awaitLeader(testPartitionB))
     val jsonFile = toJsonFile(testPartitionAndAssignment.keySet)
     try {
       PreferredReplicaLeaderElectionCommand.run(Array(
@@ -218,16 +237,16 @@ class PreferredReplicaLeaderElectionCommandTest extends ZooKeeperTestHarness wit
       jsonFile.delete()
     }
     // Check the leader for the partition IS the preferred one
-    assertEquals(testPartitionPreferredLeader, getLeader(testPartitionA))
-    assertEquals(testPartitionPreferredLeader, getLeader(testPartitionB))
+    assertEquals(testPartitionPreferredLeader, awaitLeader(testPartitionA))
+    assertEquals(testPartitionPreferredLeader, awaitLeader(testPartitionB))
   }
 
   /** What happens when the preferred replica is already the leader? */
   @Test
-  def testNoopElection() {
+  def testNoopElection(): Unit = {
     createTestTopicAndCluster(testPartitionAndAssignment)
-    // Don't bounce the server. Doublec heck the leader for the partition is the preferred one
-    assertEquals(testPartitionPreferredLeader, getLeader(testPartition))
+    // Don't bounce the server. Doublecheck the leader for the partition is the preferred one
+    assertEquals(testPartitionPreferredLeader, awaitLeader(testPartition))
     val jsonFile = toJsonFile(testPartitionAndAssignment.keySet)
     try {
       // Now do the election, even though the preferred replica is *already* the leader
@@ -235,7 +254,7 @@ class PreferredReplicaLeaderElectionCommandTest extends ZooKeeperTestHarness wit
         "--bootstrap-server", bootstrapServer(),
         "--path-to-json-file", jsonFile.getAbsolutePath))
       // Check the leader for the partition still is the preferred one
-      assertEquals(testPartitionPreferredLeader, getLeader(testPartition))
+      assertEquals(testPartitionPreferredLeader, awaitLeader(testPartition))
     } finally {
       jsonFile.delete()
     }
@@ -243,11 +262,11 @@ class PreferredReplicaLeaderElectionCommandTest extends ZooKeeperTestHarness wit
 
   /** What happens if the preferred replica is offline? */
   @Test
-  def testWithOfflinePreferredReplica() {
+  def testWithOfflinePreferredReplica(): Unit = {
     createTestTopicAndCluster(testPartitionAndAssignment)
     bounceServer(testPartitionPreferredLeader, testPartition)
     // Check the leader for the partition is not the preferred one
-    val leader = getLeader(testPartition)
+    val leader = awaitLeader(testPartition)
     assertNotEquals(testPartitionPreferredLeader, leader)
     // Now kill the preferred one
     servers(testPartitionPreferredLeader).shutdown()
@@ -265,7 +284,7 @@ class PreferredReplicaLeaderElectionCommandTest extends ZooKeeperTestHarness wit
         assertTrue(suppressed.isInstanceOf[PreferredLeaderNotAvailableException])
         assertTrue(suppressed.getMessage, suppressed.getMessage.contains("Failed to elect leader for partition test-0 under strategy PreferredReplicaPartitionLeaderElectionStrategy"))
         // Check we still have the same leader
-        assertEquals(leader, getLeader(testPartition))
+        assertEquals(leader, awaitLeader(testPartition))
     } finally {
       jsonFile.delete()
     }
@@ -273,11 +292,11 @@ class PreferredReplicaLeaderElectionCommandTest extends ZooKeeperTestHarness wit
 
   /** What happens if the controller gets killed just before an election? */
   @Test
-  def testTimeout() {
+  def testTimeout(): Unit = {
     createTestTopicAndCluster(testPartitionAndAssignment)
     bounceServer(testPartitionPreferredLeader, testPartition)
     // Check the leader for the partition is not the preferred one
-    val leader = getLeader(testPartition)
+    val leader = awaitLeader(testPartition)
     assertNotEquals(testPartitionPreferredLeader, leader)
     // Now kill the controller just before we trigger the election
     val controller = getController().get.config.brokerId
@@ -291,10 +310,9 @@ class PreferredReplicaLeaderElectionCommandTest extends ZooKeeperTestHarness wit
       fail();
     } catch {
       case e: AdminCommandFailedException =>
-        assertEquals("1 preferred replica(s) could not be elected", e.getMessage)
-        assertTrue(e.getSuppressed()(0).getMessage.contains("Timed out waiting for a node assignment"))
+        assertEquals("Timeout waiting for election results", e.getMessage)
         // Check we still have the same leader
-        assertEquals(leader, getLeader(testPartition))
+        assertEquals(leader, awaitLeader(testPartition))
     } finally {
       jsonFile.delete()
     }
@@ -302,14 +320,14 @@ class PreferredReplicaLeaderElectionCommandTest extends ZooKeeperTestHarness wit
 
   /** Test the case where client is not authorized */
   @Test
-  def testAuthzFailure() {
+  def testAuthzFailure(): Unit = {
     createTestTopicAndCluster(testPartitionAndAssignment, Some(classOf[PreferredReplicaLeaderElectionCommandTestAuthorizer].getName))
     bounceServer(testPartitionPreferredLeader, testPartition)
     // Check the leader for the partition is not the preferred one
-    val leader = getLeader(testPartition)
+    val leader = awaitLeader(testPartition)
     assertNotEquals(testPartitionPreferredLeader, leader)
     // Check the leader for the partition is not the preferred one
-    assertNotEquals(testPartitionPreferredLeader, getLeader(testPartition))
+    assertNotEquals(testPartitionPreferredLeader, awaitLeader(testPartition))
     val jsonFile = toJsonFile(testPartitionAndAssignment.keySet)
     try {
       PreferredReplicaLeaderElectionCommand.run(Array(
@@ -318,18 +336,55 @@ class PreferredReplicaLeaderElectionCommandTest extends ZooKeeperTestHarness wit
       fail();
     } catch {
       case e: AdminCommandFailedException =>
-        assertEquals("1 preferred replica(s) could not be elected", e.getMessage)
-        assertTrue(e.getSuppressed()(0).isInstanceOf[ClusterAuthorizationException])
+        assertEquals("Not authorized to perform leader election", e.getMessage)
+        assertTrue(e.getCause().isInstanceOf[ClusterAuthorizationException])
         // Check we still have the same leader
-        assertEquals(leader, getLeader(testPartition))
+        assertEquals(leader, awaitLeader(testPartition))
     } finally {
       jsonFile.delete()
     }
   }
 
+  @Test
+  def testPreferredReplicaJsonData(): Unit = {
+    // write preferred replica json data to zk path
+    val partitionsForPreferredReplicaElection = Set(new TopicPartition("test", 1), new TopicPartition("test2", 1))
+    PreferredReplicaLeaderElectionCommand.writePreferredReplicaElectionData(zkClient, partitionsForPreferredReplicaElection)
+    // try to read it back and compare with what was written
+    val partitionsUndergoingPreferredReplicaElection = zkClient.getPreferredReplicaElection
+    assertEquals("Preferred replica election ser-de failed", partitionsForPreferredReplicaElection,
+      partitionsUndergoingPreferredReplicaElection)
+  }
+
+  @Test
+  def testBasicPreferredReplicaElection(): Unit = {
+    val expectedReplicaAssignment = Map(0  -> List(0, 1, 2))
+    val topic = "test"
+    val partition = 0
+    val preferredReplica = 0
+    // create brokers
+    val brokerRack = Map(0 -> "rack0", 1 -> "rack1", 2 -> "rack2")
+    val serverConfigs = TestUtils.createBrokerConfigs(3, zkConnect, false, rackInfo = brokerRack).map(KafkaConfig.fromProps)
+    // create the topic
+    adminZkClient.createTopicWithAssignment(topic, config = new Properties, expectedReplicaAssignment)
+    servers = serverConfigs.reverseMap(s => TestUtils.createServer(s))
+    // broker 2 should be the leader since it was started first
+    val currentLeader = TestUtils.waitUntilLeaderIsElectedOrChanged(zkClient, topic, partition, oldLeaderOpt = None)
+    // trigger preferred replica election
+    val preferredReplicaElection = new PreferredReplicaLeaderElectionCommand(zkClient, Set(new TopicPartition(topic, partition)))
+    preferredReplicaElection.moveLeaderToPreferredReplica()
+    val newLeader = TestUtils.waitUntilLeaderIsElectedOrChanged(zkClient, topic, partition, oldLeaderOpt = Some(currentLeader))
+    assertEquals("Preferred replica election failed", preferredReplica, newLeader)
+  }
 }
 
-class PreferredReplicaLeaderElectionCommandTestAuthorizer extends SimpleAclAuthorizer {
-  override def authorize(session: RequestChannel.Session, operation: Operation, resource: Resource): Boolean =
-    operation != Alter || resource.resourceType != Cluster
+class PreferredReplicaLeaderElectionCommandTestAuthorizer extends AclAuthorizer {
+  override def authorize(requestContext: AuthorizableRequestContext, actions: util.List[Action]): util.List[AuthorizationResult] = {
+    actions.asScala.map { action =>
+      if (action.operation != AclOperation.ALTER || action.resourcePattern.resourceType != ResourceType.CLUSTER)
+        AuthorizationResult.ALLOWED
+      else
+        AuthorizationResult.DENIED
+    }.asJava
+  }
 }
diff --git a/core/src/test/scala/unit/kafka/admin/RackAwareTest.scala b/core/src/test/scala/unit/kafka/admin/RackAwareTest.scala
index facc7458333dc..4f4c92e0406bf 100644
--- a/core/src/test/scala/unit/kafka/admin/RackAwareTest.scala
+++ b/core/src/test/scala/unit/kafka/admin/RackAwareTest.scala
@@ -28,7 +28,7 @@ trait RackAwareTest {
                                replicationFactor: Int,
                                verifyRackAware: Boolean = true,
                                verifyLeaderDistribution: Boolean = true,
-                               verifyReplicasDistribution: Boolean = true) {
+                               verifyReplicasDistribution: Boolean = true): Unit = {
     // always verify that no broker will be assigned for more than one replica
     for ((_, brokerList) <- assignment) {
       assertEquals("More than one replica is assigned to same broker for the same partition", brokerList.toSet.size, brokerList.size)
diff --git a/core/src/test/scala/unit/kafka/admin/ReassignPartitionsClusterTest.scala b/core/src/test/scala/unit/kafka/admin/ReassignPartitionsClusterTest.scala
index 654a92ee0ac41..5933127157973 100644
--- a/core/src/test/scala/unit/kafka/admin/ReassignPartitionsClusterTest.scala
+++ b/core/src/test/scala/unit/kafka/admin/ReassignPartitionsClusterTest.scala
@@ -12,46 +12,51 @@
   */
 package kafka.admin
 
-import java.util.Collections
-import java.util.Properties
-
 import kafka.admin.ReassignPartitionsCommand._
 import kafka.common.AdminCommandFailedException
-import kafka.server.{KafkaConfig, KafkaServer}
+import kafka.server.{DynamicConfig, KafkaConfig, KafkaServer}
 import kafka.utils.TestUtils._
 import kafka.utils.{Logging, TestUtils}
 import kafka.zk.{ReassignPartitionsZNode, ZkVersion, ZooKeeperTestHarness}
-import org.junit.Assert.{assertEquals, assertTrue}
+import org.junit.Assert.{assertEquals, assertFalse, assertTrue}
 import org.junit.{After, Before, Test}
 import kafka.admin.ReplicationQuotaUtils._
-import org.apache.kafka.clients.admin.AdminClientConfig
-import org.apache.kafka.clients.admin.{AdminClient => JAdminClient}
+import org.apache.kafka.clients.admin.{Admin, AdminClientConfig, AlterConfigOp, ConfigEntry, NewPartitionReassignment, NewPartitions, PartitionReassignment, AdminClient => JAdminClient}
 import org.apache.kafka.common.{TopicPartition, TopicPartitionReplica}
 
 import scala.collection.JavaConverters._
-import scala.collection.Map
-import scala.collection.Seq
+import scala.collection.{Map, Seq}
 import scala.util.Random
 import java.io.File
+import java.util.{Collections, Optional, Properties}
+import java.util.concurrent.ExecutionException
 
+import kafka.controller.ReplicaAssignment
+import kafka.log.LogConfig
+import org.apache.kafka.common.config.ConfigResource
 import org.apache.kafka.clients.producer.ProducerRecord
+import org.apache.kafka.common.errors.{NoReassignmentInProgressException, ReassignmentInProgressException}
+import org.scalatest.Assertions.intercept
 
 class ReassignPartitionsClusterTest extends ZooKeeperTestHarness with Logging {
-  val partitionId = 0
   var servers: Seq[KafkaServer] = null
+  var brokerIds: Seq[Int] = null
   val topicName = "my-topic"
+  val tp0 = new TopicPartition(topicName, 0)
+  val tp1 = new TopicPartition(topicName, 1)
   val delayMs = 1000
-  var adminClient: JAdminClient = null
+  var adminClient: Admin = null
 
   def zkUpdateDelay(): Unit = Thread.sleep(delayMs)
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     super.setUp()
   }
 
-  def startBrokers(brokerIds: Seq[Int]) {
-    servers = brokerIds.map { i =>
+  def startBrokers(ids: Seq[Int]): Unit = {
+    brokerIds = ids
+    servers = ids.map { i =>
       val props = createBrokerConfig(i, zkConnect, enableControlledShutdown = false, logDirCount = 3)
       // shorter backoff to reduce test durations when no active partitions are eligible for fetching due to throttling
       props.put(KafkaConfig.ReplicaFetchBackoffMsProp, "100")
@@ -59,21 +64,21 @@ class ReassignPartitionsClusterTest extends ZooKeeperTestHarness with Logging {
     }.map(c => createServer(KafkaConfig.fromProps(c)))
   }
 
-  def createAdminClient(servers: Seq[KafkaServer]): JAdminClient = {
+  def createAdminClient(servers: Seq[KafkaServer]): Admin = {
     val props = new Properties()
     props.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, TestUtils.getBrokerListStrFromServers(servers))
     props.put(AdminClientConfig.REQUEST_TIMEOUT_MS_CONFIG, "10000")
     JAdminClient.create(props)
   }
 
-  def getRandomLogDirAssignment(brokerId: Int): String = {
+  def getRandomLogDirAssignment(brokerId: Int, excluded: Option[String] = None): String = {
     val server = servers.find(_.config.brokerId == brokerId).get
-    val logDirs = server.config.logDirs
+    val logDirs = server.config.logDirs.filterNot(excluded.contains)
     new File(logDirs(Random.nextInt(logDirs.size))).getAbsolutePath
   }
 
   @After
-  override def tearDown() {
+  override def tearDown(): Unit = {
     if (adminClient != null) {
       adminClient.close()
       adminClient = null
@@ -87,27 +92,28 @@ class ReassignPartitionsClusterTest extends ZooKeeperTestHarness with Logging {
     //Given a single replica on server 100
     startBrokers(Seq(100, 101, 102))
     adminClient = createAdminClient(servers)
-    createTopic(zkClient, topicName, Map(0 -> Seq(100)), servers = servers)
+    createTopic(zkClient, topicName, Map(tp0.partition() -> Seq(100)), servers = servers)
 
-    val topicPartition = new TopicPartition(topicName, 0)
     val leaderServer = servers.find(_.config.brokerId == 100).get
-    leaderServer.replicaManager.logManager.truncateFullyAndStartAt(topicPartition, 100L, false)
+    leaderServer.replicaManager.logManager.truncateFullyAndStartAt(tp0, 100L, false)
 
-    val topicJson: String = s"""{"version":1,"partitions":[{"topic":"$topicName","partition":0,"replicas":[101, 102]}]}"""
+    val topicJson = executeAssignmentJson(Seq(
+      PartitionAssignmentJson(tp0, replicas=Seq(101, 102))
+    ))
     ReassignPartitionsCommand.executeAssignment(zkClient, Some(adminClient), topicJson, NoThrottle)
 
     val newLeaderServer = servers.find(_.config.brokerId == 101).get
 
-    TestUtils.waitUntilTrue (
-      () => newLeaderServer.replicaManager.getPartition(topicPartition).flatMap(_.leaderReplicaIfLocal).isDefined,
+    waitUntilTrue (
+      () => newLeaderServer.replicaManager.nonOfflinePartition(tp0).flatMap(_.leaderLogIfLocal).isDefined,
       "broker 101 should be the new leader", pause = 1L
     )
 
-    assertEquals(100, newLeaderServer.replicaManager.localReplicaOrException(topicPartition)
-      .highWatermark.messageOffset)
+    assertEquals(100, newLeaderServer.replicaManager.localLogOrException(tp0)
+      .highWatermark)
     val newFollowerServer = servers.find(_.config.brokerId == 102).get
-    TestUtils.waitUntilTrue(() => newFollowerServer.replicaManager.localReplicaOrException(topicPartition)
-      .highWatermark.messageOffset == 100,
+    waitUntilTrue(() => newFollowerServer.replicaManager.localLogOrException(tp0)
+      .highWatermark == 100,
       "partition follower's highWatermark should be 100")
   }
 
@@ -116,42 +122,77 @@ class ReassignPartitionsClusterTest extends ZooKeeperTestHarness with Logging {
     //Given a single replica on server 100
     startBrokers(Seq(100, 101))
     adminClient = createAdminClient(servers)
-    val partition = 0
     // Get a random log directory on broker 101
     val expectedLogDir = getRandomLogDirAssignment(101)
-    createTopic(zkClient, topicName, Map(partition -> Seq(100)), servers = servers)
+    createTopic(zkClient, topicName, Map(tp0.partition() -> Seq(100)), servers = servers)
 
     //When we move the replica on 100 to broker 101
-    val topicJson: String = s"""{"version":1,"partitions":[{"topic":"$topicName","partition":0,"replicas":[101],"log_dirs":["$expectedLogDir"]}]}"""
+    val topicJson = executeAssignmentJson(Seq(
+      PartitionAssignmentJson(tp0, replicas = Seq(101), logDirectories = Some(Seq(expectedLogDir)))
+    ))
     ReassignPartitionsCommand.executeAssignment(zkClient, Some(adminClient), topicJson, NoThrottle)
-    waitForReassignmentToComplete()
+    waitForZkReassignmentToComplete()
 
     //Then the replica should be on 101
-    assertEquals(Seq(101), zkClient.getPartitionAssignmentForTopics(Set(topicName)).get(topicName).get(partition))
+    val partitionAssignment = zkClient.getPartitionAssignmentForTopics(Set(topicName)).get(topicName).get(tp0.partition())
+    assertMoveForPartitionOccurred(Seq(101), partitionAssignment)
     // The replica should be in the expected log directory on broker 101
     val replica = new TopicPartitionReplica(topicName, 0, 101)
     assertEquals(expectedLogDir, adminClient.describeReplicaLogDirs(Collections.singleton(replica)).all().get.get(replica).getCurrentReplicaLogDir)
   }
 
   @Test
-  def shouldMoveSinglePartitionWithinBroker() {
+  def testReassignmentMatchesCurrentAssignment(): Unit = {
+    // Given a single replica on server 100
+    startBrokers(Seq(100))
+    adminClient = createAdminClient(servers)
+    createTopic(zkClient, topicName, Map(tp0.partition() -> Seq(100)), servers = servers)
+
+    // Execute no-op reassignment
+    val topicJson = executeAssignmentJson(Seq(
+      PartitionAssignmentJson(tp0, replicas = Seq(100), logDirectories = None)
+    ))
+    ReassignPartitionsCommand.executeAssignment(zkClient, Some(adminClient), topicJson, NoThrottle)
+    waitForZkReassignmentToComplete()
+
+    // The replica should remain on 100
+    val partitionAssignment = zkClient.getPartitionAssignmentForTopics(Set(topicName))(topicName)(tp0.partition)
+    assertMoveForPartitionOccurred(Seq(100), partitionAssignment)
+  }
+
+  @Test
+  def shouldMoveSinglePartitionToSameFolderWithinBroker(): Unit = shouldMoveSinglePartitionWithinBroker(true)
+
+  @Test
+  def shouldMoveSinglePartitionToDifferentFolderWithinBroker(): Unit = shouldMoveSinglePartitionWithinBroker(false)
+
+  private[this] def shouldMoveSinglePartitionWithinBroker(moveToSameFolder: Boolean): Unit = {
     // Given a single replica on server 100
     startBrokers(Seq(100, 101))
     adminClient = createAdminClient(servers)
-    val expectedLogDir = getRandomLogDirAssignment(100)
-    createTopic(zkClient, topicName, Map(0 -> Seq(100)), servers = servers)
+    createTopic(zkClient, topicName, Map(tp0.partition() -> Seq(100)), servers = servers)
+
+    val replica = new TopicPartitionReplica(topicName, 0, 100)
+    val currentLogDir = adminClient.describeReplicaLogDirs(java.util.Collections.singleton(replica))
+      .all()
+      .get()
+      .get(replica)
+      .getCurrentReplicaLogDir
+
+    val expectedLogDir = if (moveToSameFolder) currentLogDir else getRandomLogDirAssignment(100, excluded = Some(currentLogDir))
 
     // When we execute an assignment that moves an existing replica to another log directory on the same broker
-    val topicJson: String = s"""{"version":1,"partitions":[{"topic":"$topicName","partition":0,"replicas":[100],"log_dirs":["$expectedLogDir"]}]}"""
+    val topicJson = executeAssignmentJson(Seq(
+      PartitionAssignmentJson(tp0, replicas = Seq(100), logDirectories = Some(Seq(expectedLogDir)))
+    ))
     ReassignPartitionsCommand.executeAssignment(zkClient, Some(adminClient), topicJson, NoThrottle)
-    val replica = new TopicPartitionReplica(topicName, 0, 100)
-    TestUtils.waitUntilTrue(() => {
+    waitUntilTrue(() => {
       expectedLogDir == adminClient.describeReplicaLogDirs(Collections.singleton(replica)).all().get.get(replica).getCurrentReplicaLogDir
     }, "Partition should have been moved to the expected log directory", 1000)
   }
 
   @Test
-  def shouldExpandCluster() {
+  def shouldExpandCluster(): Unit = {
     val brokers = Array(100, 101, 102)
     startBrokers(brokers)
     adminClient = createAdminClient(servers)
@@ -162,7 +203,7 @@ class ReassignPartitionsClusterTest extends ZooKeeperTestHarness with Logging {
     ), servers = servers)
 
     //When rebalancing
-    val newAssignment = generateAssignment(zkClient, brokers, json(topicName), true)._1
+    val newAssignment = generateAssignment(zkClient, brokers, generateAssignmentJson(topicName), true)._1
     // Find a partition in the new assignment on broker 102 and a random log directory on broker 102,
     // which currently does not have any partition for this topic
     val partition1 = newAssignment.find { case (_, brokerIds) => brokerIds.contains(102) }.get._1.partition
@@ -178,11 +219,11 @@ class ReassignPartitionsClusterTest extends ZooKeeperTestHarness with Logging {
     val newReplicaAssignment = Map(replica1 -> expectedLogDir1, replica2 -> expectedLogDir2)
     ReassignPartitionsCommand.executeAssignment(zkClient, Some(adminClient),
       ReassignPartitionsCommand.formatAsReassignmentJson(newAssignment, newReplicaAssignment), NoThrottle)
-    waitForReassignmentToComplete()
+    waitForZkReassignmentToComplete()
 
     // Then the replicas should span all three brokers
     val actual = zkClient.getPartitionAssignmentForTopics(Set(topicName))(topicName)
-    assertEquals(Seq(100, 101, 102), actual.values.flatten.toSeq.distinct.sorted)
+    assertMoveForTopicOccurred(Seq(100, 101, 102), actual)
     // The replica should be in the expected log directory on broker 102 and 100
     waitUntilTrue(() => {
       expectedLogDir1 == adminClient.describeReplicaLogDirs(Collections.singleton(replica1)).all().get.get(replica1).getCurrentReplicaLogDir
@@ -193,7 +234,7 @@ class ReassignPartitionsClusterTest extends ZooKeeperTestHarness with Logging {
   }
 
   @Test
-  def shouldShrinkCluster() {
+  def shouldShrinkCluster(): Unit = {
     //Given partitions on 3 of 3 brokers
     val brokers = Array(100, 101, 102)
     startBrokers(brokers)
@@ -204,18 +245,18 @@ class ReassignPartitionsClusterTest extends ZooKeeperTestHarness with Logging {
     ), servers = servers)
 
     //When rebalancing
-    val newAssignment = generateAssignment(zkClient, Array(100, 101), json(topicName), true)._1
+    val newAssignment = generateAssignment(zkClient, Array(100, 101), generateAssignmentJson(topicName), true)._1
     ReassignPartitionsCommand.executeAssignment(zkClient, None,
       ReassignPartitionsCommand.formatAsReassignmentJson(newAssignment, Map.empty), NoThrottle)
-    waitForReassignmentToComplete()
+    waitForZkReassignmentToComplete()
 
     //Then replicas should only span the first two brokers
     val actual = zkClient.getPartitionAssignmentForTopics(Set(topicName))(topicName)
-    assertEquals(Seq(100, 101), actual.values.flatten.toSeq.distinct.sorted)
+    assertMoveForTopicOccurred(Seq(100, 101), actual)
   }
 
   @Test
-  def shouldMoveSubsetOfPartitions() {
+  def shouldMoveSubsetOfPartitions(): Unit = {
     //Given partitions on 3 of 3 brokers
     val brokers = Array(100, 101, 102)
     startBrokers(brokers)
@@ -248,16 +289,16 @@ class ReassignPartitionsClusterTest extends ZooKeeperTestHarness with Logging {
     //When rebalancing
     ReassignPartitionsCommand.executeAssignment(zkClient, Some(adminClient),
       ReassignPartitionsCommand.formatAsReassignmentJson(proposed, proposedReplicaAssignment), NoThrottle)
-    waitForReassignmentToComplete()
+    waitForZkReassignmentToComplete()
 
     //Then the proposed changes should have been made
     val actual = zkClient.getPartitionAssignmentForTopics(Set("topic1", "topic2"))
-    assertEquals(Seq(100, 102), actual("topic1")(0))//changed
-    assertEquals(Seq(101, 102), actual("topic1")(1))
-    assertEquals(Seq(100, 102), actual("topic1")(2))//changed
-    assertEquals(Seq(100, 101), actual("topic2")(0))
-    assertEquals(Seq(101, 100), actual("topic2")(1))//changed
-    assertEquals(Seq(100, 102), actual("topic2")(2))//changed
+    assertMoveForPartitionOccurred(Seq(100, 102), actual("topic1")(0)) //changed
+    assertMoveForPartitionOccurred(Seq(101, 102), actual("topic1")(1))
+    assertMoveForPartitionOccurred(Seq(100, 102), actual("topic1")(2)) //changed
+    assertMoveForPartitionOccurred(Seq(100, 101), actual("topic2")(0))
+    assertMoveForPartitionOccurred(Seq(101, 100), actual("topic2")(1)) //changed
+    assertMoveForPartitionOccurred(Seq(100, 102), actual("topic2")(2)) //changed
 
     // The replicas should be in the expected log directories
     val replicaDirs = adminClient.describeReplicaLogDirs(List(replica1, replica2).asJava).all().get()
@@ -266,7 +307,7 @@ class ReassignPartitionsClusterTest extends ZooKeeperTestHarness with Logging {
   }
 
   @Test
-  def shouldExecuteThrottledReassignment() {
+  def shouldExecuteThrottledReassignment(): Unit = {
 
     //Given partitions on 3 of 3 brokers
     val brokers = Array(100, 101, 102)
@@ -284,7 +325,7 @@ class ReassignPartitionsClusterTest extends ZooKeeperTestHarness with Logging {
     assertEquals(expectedDurationSecs, numMessages * msgSize / initialThrottle.interBrokerLimit)
 
     //Start rebalance which will move replica on 100 -> replica on 102
-    val newAssignment = generateAssignment(zkClient, Array(101, 102), json(topicName), true)._1
+    val newAssignment = generateAssignment(zkClient, Array(101, 102), generateAssignmentJson(topicName), true)._1
 
     val start = System.currentTimeMillis()
     ReassignPartitionsCommand.executeAssignment(zkClient, None,
@@ -294,12 +335,12 @@ class ReassignPartitionsClusterTest extends ZooKeeperTestHarness with Logging {
     checkThrottleConfigAddedToZK(adminZkClient, initialThrottle.interBrokerLimit, servers, topicName, Set("0:100","0:101"), Set("0:102"))
 
     //Await completion
-    waitForReassignmentToComplete()
+    waitForZkReassignmentToComplete()
     val took = System.currentTimeMillis() - start - delayMs
 
     //Check move occurred
     val actual = zkClient.getPartitionAssignmentForTopics(Set(topicName))(topicName)
-    assertEquals(Seq(101, 102), actual.values.flatten.toSeq.distinct.sorted)
+    assertMoveForTopicOccurred(Seq(101, 102), actual)
 
     //Then command should have taken longer than the throttle rate
     assertTrue(s"Expected replication to be > ${expectedDurationSecs * 0.9 * 1000} but was $took",
@@ -310,7 +351,7 @@ class ReassignPartitionsClusterTest extends ZooKeeperTestHarness with Logging {
 
 
   @Test
-  def shouldOnlyThrottleMovingReplicas() {
+  def shouldOnlyThrottleMovingReplicas(): Unit = {
     //Given 6 brokers, two topics
     val brokers = Array(100, 101, 102, 103, 104, 105)
     startBrokers(brokers)
@@ -355,7 +396,7 @@ class ReassignPartitionsClusterTest extends ZooKeeperTestHarness with Logging {
   }
 
   @Test
-  def shouldChangeThrottleOnRerunAndRemoveOnVerify() {
+  def shouldChangeThrottleOnRerunAndRemoveOnVerify(): Unit = {
     //Given partitions on 3 of 3 brokers
     val brokers = Array(100, 101, 102)
     startBrokers(brokers)
@@ -368,7 +409,7 @@ class ReassignPartitionsClusterTest extends ZooKeeperTestHarness with Logging {
     produceMessages(topicName, numMessages = 200, acks = 0, valueLength = 100 * 1000)
 
     //Start rebalance
-    val newAssignment = generateAssignment(zkClient, Array(101, 102), json(topicName), true)._1
+    val newAssignment = generateAssignment(zkClient, Array(101, 102), generateAssignmentJson(topicName), true)._1
 
     ReassignPartitionsCommand.executeAssignment(zkClient, None,
       ReassignPartitionsCommand.formatAsReassignmentJson(newAssignment, Map.empty), Throttle(initialThrottle))
@@ -392,7 +433,7 @@ class ReassignPartitionsClusterTest extends ZooKeeperTestHarness with Logging {
     checkThrottleConfigAddedToZK(adminZkClient, newThrottle, servers, topicName, Set("0:100","0:101"), Set("0:102"))
 
     //Await completion
-    waitForReassignmentToComplete()
+    waitForZkReassignmentToComplete()
 
     //Verify should remove the throttle
     verifyAssignment(zkClient, None, ReassignPartitionsCommand.formatAsReassignmentJson(newAssignment, Map.empty))
@@ -402,69 +443,69 @@ class ReassignPartitionsClusterTest extends ZooKeeperTestHarness with Logging {
 
     //Check move occurred
     val actual = zkClient.getPartitionAssignmentForTopics(Set(topicName))(topicName)
-    assertEquals(Seq(101, 102), actual.values.flatten.toSeq.distinct.sorted)
+    assertMoveForTopicOccurred(Seq(101, 102), actual)
   }
 
   @Test(expected = classOf[AdminCommandFailedException])
-  def shouldFailIfProposedDoesNotMatchExisting() {
+  def shouldFailIfProposedDoesNotMatchExisting(): Unit = {
     //Given a single replica on server 100
     startBrokers(Seq(100, 101))
-    createTopic(zkClient, topicName, Map(0 -> Seq(100)), servers = servers)
+    createTopic(zkClient, topicName, Map(tp0.partition() -> Seq(100)), servers = servers)
 
     //When we execute an assignment that includes an invalid partition (1:101 in this case)
-    val topicJson = s"""{"version":1,"partitions":[{"topic":"$topicName","partition":1,"replicas":[101]}]}"""
+    val topicJson = executeAssignmentJson(Seq(PartitionAssignmentJson(tp1, Seq(101))))
     ReassignPartitionsCommand.executeAssignment(zkClient, None, topicJson, NoThrottle)
   }
 
   @Test(expected = classOf[AdminCommandFailedException])
-  def shouldFailIfProposedHasEmptyReplicaList() {
+  def shouldFailIfProposedHasEmptyReplicaList(): Unit = {
     //Given a single replica on server 100
     startBrokers(Seq(100, 101))
-    createTopic(zkClient, topicName, Map(0 -> Seq(100)), servers = servers)
+    createTopic(zkClient, topicName, Map(tp0.partition() -> Seq(100)), servers = servers)
 
     //When we execute an assignment that specifies an empty replica list (0: empty list in this case)
-    val topicJson = s"""{"version":1,"partitions":[{"topic":"$topicName","partition":0,"replicas":[]}]}"""
+    val topicJson = executeAssignmentJson(Seq(PartitionAssignmentJson(tp0, Seq())))
     ReassignPartitionsCommand.executeAssignment(zkClient, None, topicJson, NoThrottle)
   }
 
   @Test(expected = classOf[AdminCommandFailedException])
-  def shouldFailIfProposedHasInvalidBrokerID() {
+  def shouldFailIfProposedHasInvalidBrokerID(): Unit = {
     //Given a single replica on server 100
     startBrokers(Seq(100, 101))
-    createTopic(zkClient, topicName, Map(0 -> Seq(100)), servers = servers)
+    createTopic(zkClient, topicName, Map(tp0.partition() -> Seq(100)), servers = servers)
 
     //When we execute an assignment that specifies an invalid brokerID (102: invalid broker ID in this case)
-    val topicJson = s"""{"version":1,"partitions":[{"topic":"$topicName","partition":0,"replicas":[101, 102]}]}"""
+    val topicJson = executeAssignmentJson(Seq(PartitionAssignmentJson(tp0, Seq(101, 102))))
     ReassignPartitionsCommand.executeAssignment(zkClient, None, topicJson, NoThrottle)
   }
 
   @Test(expected = classOf[AdminCommandFailedException])
-  def shouldFailIfProposedHasInvalidLogDir() {
+  def shouldFailIfProposedHasInvalidLogDir(): Unit = {
     // Given a single replica on server 100
     startBrokers(Seq(100, 101))
     adminClient = createAdminClient(servers)
-    createTopic(zkClient, topicName, Map(0 -> Seq(100)), servers = servers)
+    createTopic(zkClient, topicName, Map(tp0.partition() -> Seq(100)), servers = servers)
 
     // When we execute an assignment that specifies an invalid log directory
-    val topicJson: String = s"""{"version":1,"partitions":[{"topic":"$topicName","partition":0,"replicas":[101],"log_dirs":["invalidDir"]}]}"""
+    val topicJson = executeAssignmentJson(Seq(PartitionAssignmentJson(tp0, Seq(101), logDirectories = Some(Seq("invalidDir")))))
     ReassignPartitionsCommand.executeAssignment(zkClient, Some(adminClient), topicJson, NoThrottle)
   }
 
   @Test(expected = classOf[AdminCommandFailedException])
-  def shouldFailIfProposedHasInconsistentReplicasAndLogDirs() {
+  def shouldFailIfProposedHasInconsistentReplicasAndLogDirs(): Unit = {
     // Given a single replica on server 100
     startBrokers(Seq(100, 101))
     adminClient = createAdminClient(servers)
     val logDir = getRandomLogDirAssignment(100)
-    createTopic(zkClient, topicName, Map(0 -> Seq(100)), servers = servers)
+    createTopic(zkClient, topicName, Map(tp0.partition() -> Seq(100)), servers = servers)
 
-    // When we execute an assignment whose length of replicas doesn't match that of replicas
-    val topicJson: String = s"""{"version":1,"partitions":[{"topic":"$topicName","partition":0,"replicas":[101],"log_dirs":["$logDir", "$logDir"]}]}"""
+    // When we execute an assignment whose length of replicas doesn't match that of log dirs
+    val topicJson = executeAssignmentJson(Seq(PartitionAssignmentJson(tp0, Seq(101), logDirectories = Some(Seq(logDir, logDir)))))
     ReassignPartitionsCommand.executeAssignment(zkClient, Some(adminClient), topicJson, NoThrottle)
   }
 
   @Test
-  def shouldPerformThrottledReassignmentOverVariousTopics() {
+  def shouldPerformThrottledReassignmentOverVariousTopics(): Unit = {
     val throttle = Throttle(1000L)
 
     startBrokers(Seq(0, 1, 2, 3))
@@ -486,7 +527,7 @@ class ReassignPartitionsClusterTest extends ZooKeeperTestHarness with Logging {
     //When we run a throttled reassignment
     new ReassignPartitionsCommand(zkClient, None, move, adminZkClient = adminZkClient).reassignPartitions(throttle)
 
-    waitForReassignmentToComplete()
+    waitForZkReassignmentToComplete()
 
     //Check moved replicas did move
     assertEquals(Seq(0, 2, 3), zkClient.getReplicasForPartition(new TopicPartition("orders", 0)))
@@ -509,7 +550,7 @@ class ReassignPartitionsClusterTest extends ZooKeeperTestHarness with Logging {
    * often enough to detect a regression.
    */
   @Test
-  def shouldPerformMultipleReassignmentOperationsOverVariousTopics() {
+  def shouldPerformMultipleReassignmentOperationsOverVariousTopics(): Unit = {
     startBrokers(Seq(0, 1, 2, 3))
 
     createTopic(zkClient, "orders", Map(0 -> List(0, 1, 2), 1 -> List(0, 1, 2)), servers)
@@ -526,7 +567,7 @@ class ReassignPartitionsClusterTest extends ZooKeeperTestHarness with Logging {
 
     new ReassignPartitionsCommand(zkClient, None, firstMove, adminZkClient = adminZkClient).reassignPartitions()
     // Low pause to detect deletion of the reassign_partitions znode before the reassignment is complete
-    waitForReassignmentToComplete(pause = 1L)
+    waitForZkReassignmentToComplete(pause = 1L)
 
     // Check moved replicas did move
     assertEquals(Seq(0, 2, 3), zkClient.getReplicasForPartition(new TopicPartition("orders", 0)))
@@ -551,7 +592,7 @@ class ReassignPartitionsClusterTest extends ZooKeeperTestHarness with Logging {
 
     new ReassignPartitionsCommand(zkClient, None, secondMove, adminZkClient = adminZkClient).reassignPartitions()
     // Low pause to detect deletion of the reassign_partitions znode before the reassignment is complete
-    waitForReassignmentToComplete(pause = 1L)
+    waitForZkReassignmentToComplete(pause = 1L)
 
     // Check moved replicas did move
     assertEquals(Seq(0, 2, 3), zkClient.getReplicasForPartition(new TopicPartition("orders", 0)))
@@ -585,7 +626,7 @@ class ReassignPartitionsClusterTest extends ZooKeeperTestHarness with Logging {
     }.exists(identity)
 
     // Low pause to detect deletion of the reassign_partitions znode before the reassignment is complete
-    waitForReassignmentToComplete(pause = 1L)
+    waitForZkReassignmentToComplete(pause = 1L)
 
     // Check moved replicas for thirdMove and fourthMove
     assertEquals(Seq(1, 2, 3), zkClient.getReplicasForPartition(new TopicPartition("orders", 0)))
@@ -621,23 +662,757 @@ class ReassignPartitionsClusterTest extends ZooKeeperTestHarness with Logging {
     zkClient.setOrCreatePartitionReassignment(firstMove, ZkVersion.MatchAnyVersion)
 
     servers.foreach(_.startup())
-    waitForReassignmentToComplete()
+    waitForZkReassignmentToComplete()
 
     assertEquals(Seq(2, 1), zkClient.getReplicasForPartition(new TopicPartition("orders", 0)))
     assertEquals(Seq(1, 2), zkClient.getReplicasForPartition(new TopicPartition("orders", 1)))
     assertEquals(Seq.empty, zkClient.getReplicasForPartition(new TopicPartition("customers", 0)))
   }
 
-  def waitForReassignmentToComplete(pause: Long = 100L) {
+  /**
+    * Set a reassignment through the `/topics/` znode and set the `reassign_partitions` znode while the brokers are down.
+    * Verify that the reassignment is triggered by the Controller during start-up with the `reassign_partitions` znode taking precedence
+    */
+  @Test
+  def shouldTriggerReassignmentWithZnodePrecedenceOnControllerStartup(): Unit = {
+    startBrokers(Seq(0, 1, 2))
+    adminClient = createAdminClient(servers)
+    createTopic(zkClient, "orders", Map(0 -> List(0, 1), 1 -> List(1, 2), 2 -> List(0, 1), 3 -> List(0, 1)), servers)
+    val sameMoveTp = new TopicPartition("orders", 2)
+
+    // Throttle to ensure we minimize race conditions and test flakiness
+    throttle(Seq("orders"), throttleSettingForSeconds(10), Map(
+      sameMoveTp -> Seq(0, 1, 2)
+    ))
+
+    servers.foreach(_.shutdown())
+    adminClient.close()
+
+    zkClient.setTopicAssignment("orders", Map(
+      new TopicPartition("orders", 0) -> ReplicaAssignment(List(0, 1), List(2), List(0)), // should be overwritten
+      new TopicPartition("orders", 1) -> ReplicaAssignment(List(1, 2), List(3), List(1)), // should be overwritten
+      // should be overwritten (so we know to remove it from ZK) even though we do the exact same move
+      sameMoveTp -> ReplicaAssignment(List(0, 1, 2), List(2), List(0)),
+      new TopicPartition("orders", 3) -> ReplicaAssignment(List(0, 1, 2), List(2), List(0)) // moves
+    ))
+    val move = Map(
+      new TopicPartition("orders", 0) -> Seq(2, 1), // moves
+      new TopicPartition("orders", 1) -> Seq(1, 2), // stays
+      sameMoveTp -> Seq(1, 2), // same reassignment
+      // orders-3 intentionally left for API
+      new TopicPartition("customers", 0) -> Seq(1, 2) // non-existent topic, triggers topic deleted path
+    )
+
+    // Set znode directly to avoid non-existent topic validation
+    zkClient.setOrCreatePartitionReassignment(move, ZkVersion.MatchAnyVersion)
+
+    servers.foreach(_.startup())
+    TestUtils.waitUntilBrokerMetadataIsPropagated(servers)
+    adminClient = createAdminClient(servers)
+    resetBrokersThrottle()
+
+    waitForZkReassignmentToComplete()
+
+    assertEquals(Seq(2, 1), zkClient.getReplicasForPartition(new TopicPartition("orders", 0)))
+    assertEquals(Seq(1, 2), zkClient.getReplicasForPartition(new TopicPartition("orders", 1)))
+    assertEquals(Seq(1, 2), zkClient.getReplicasForPartition(sameMoveTp))
+    assertEquals(Seq(1, 2), zkClient.getReplicasForPartition(new TopicPartition("orders", 3)))
+    assertEquals(Seq.empty, zkClient.getReplicasForPartition(new TopicPartition("customers", 0)))
+  }
+
+  @Test
+  def shouldListReassignmentsTriggeredByZk(): Unit = {
+    // Given a single replica on server 100
+    startBrokers(Seq(100, 101))
+    adminClient = createAdminClient(servers)
+    // Get a random log directory on broker 101
+    val expectedLogDir = getRandomLogDirAssignment(101)
+    createTopic(zkClient, topicName, Map(tp0.partition() -> Seq(100)), servers = servers)
+    // Given throttle set so replication will take at least 2 sec (to ensure we don't minimize race condition and test flakiness
+    val throttle: Long = 1000 * 1000
+    produceMessages(topicName, numMessages = 20, acks = 0, valueLength = 100 * 1000)
+
+    // When we move the replica on 100 to broker 101
+    val topicJson = executeAssignmentJson(Seq(
+      PartitionAssignmentJson(tp0, replicas = Seq(101), Some(Seq(expectedLogDir)))))
+    ReassignPartitionsCommand.executeAssignment(zkClient, Some(adminClient), topicJson, Throttle(throttle))
+    // Then the replica should be removing
+    val reassigningPartitionsResult = adminClient.listPartitionReassignments(Set(tp0).asJava).reassignments().get().get(tp0)
+    assertIsReassigning(from = Seq(100), to = Seq(101), reassigningPartitionsResult)
+
+    waitForZkReassignmentToComplete()
+
+    // Then the replica should be on 101
+    val partitionAssignment = zkClient.getPartitionAssignmentForTopics(Set(topicName)).get(topicName).get(tp0.partition())
+    assertMoveForPartitionOccurred(Seq(101), partitionAssignment)
+  }
+
+  @Test
+  def shouldReassignThroughApi(): Unit = {
+    startBrokers(Seq(100, 101))
+    adminClient = createAdminClient(servers)
+    createTopic(zkClient, topicName, Map(tp0.partition() -> Seq(100)), servers = servers)
+
+    assertTrue(adminClient.listPartitionReassignments(Set(tp0).asJava).reassignments().get().isEmpty)
+    assertEquals(Seq(100), zkClient.getReplicasForPartition(tp0))
+    adminClient.alterPartitionReassignments(
+      Map(reassignmentEntry(tp0, Seq(101))).asJava
+    ).all().get()
+
+    waitForAllReassignmentsToComplete()
+    assertEquals(Seq(101), zkClient.getReplicasForPartition(tp0))
+  }
+
+  @Test
+  def testProduceAndConsumeWithReassignmentInProgress(): Unit = {
+    startBrokers(Seq(100, 101))
+    adminClient = createAdminClient(servers)
+    createTopic(zkClient, topicName, Map(tp0.partition() -> Seq(100)), servers = servers)
+
+    produceMessages(tp0.topic, 500, acks = -1, valueLength = 100 * 1000)
+
+    TestUtils.throttleAllBrokersReplication(adminClient, Seq(101), throttleBytes = 1)
+    TestUtils.assignThrottledPartitionReplicas(adminClient, Map(tp0 -> Seq(101)))
+
+    adminClient.alterPartitionReassignments(
+      Map(reassignmentEntry(tp0, Seq(100, 101))).asJava
+    ).all().get()
+
+    awaitReassignmentInProgress(tp0)
+
+    produceMessages(tp0.topic, 500, acks = -1, valueLength = 64)
+    val consumer = TestUtils.createConsumer(TestUtils.getBrokerListStrFromServers(servers))
+    try {
+      consumer.assign(Seq(tp0).asJava)
+      pollUntilAtLeastNumRecords(consumer, numRecords = 1000)
+    } finally {
+      consumer.close()
+    }
+
+    assertTrue(isAssignmentInProgress(tp0))
+
+    TestUtils.resetBrokersThrottle(adminClient, Seq(101))
+    TestUtils.removePartitionReplicaThrottles(adminClient, Set(tp0))
+
+    waitForAllReassignmentsToComplete()
+    assertEquals(Seq(100, 101), zkClient.getReplicasForPartition(tp0))
+  }
+
+  @Test
+  def shouldListMovingPartitionsThroughApi(): Unit = {
+    startBrokers(Seq(100, 101))
+    adminClient = createAdminClient(servers)
+    val topic2 = "topic2"
+    val tp2 = new TopicPartition(topic2, 0)
+
+    createTopic(zkClient, topicName,
+      Map(tp0.partition() -> Seq(100),
+          tp1.partition() -> Seq(101)),
+      servers = servers)
+    createTopic(zkClient, topic2,
+      Map(tp2.partition() -> Seq(100)),
+      servers = servers)
+    assertTrue(adminClient.listPartitionReassignments().reassignments().get().isEmpty)
+
+    // Throttle to ensure we minimize race conditions and test flakiness
+    throttle(Seq(topicName), throttleSettingForSeconds(10), Map(
+      tp0 -> Seq(100, 101),
+      tp2 -> Seq(100, 101)
+    ))
+    adminClient.alterPartitionReassignments(
+      Map(reassignmentEntry(tp0, Seq(101)),
+          reassignmentEntry(tp2, Seq(101))).asJava
+    ).all().get()
+
+    val reassignmentsInProgress = adminClient.listPartitionReassignments(Set(tp0, tp1, tp2).asJava).reassignments().get()
+    assertFalse(reassignmentsInProgress.containsKey(tp1)) // tp1 is not reassigning
+    assertIsReassigning(from = Seq(100), to = Seq(101), reassignmentsInProgress.get(tp0))
+    assertIsReassigning(from = Seq(100), to = Seq(101), reassignmentsInProgress.get(tp2))
+
+    resetBrokersThrottle()
+    waitForAllReassignmentsToComplete()
+    assertEquals(Seq(101), zkClient.getReplicasForPartition(tp0))
+    assertEquals(Seq(101), zkClient.getReplicasForPartition(tp2))
+  }
+
+  @Test
+  def shouldUseLatestOrderingIfTwoConsecutiveReassignmentsHaveSameSetButDifferentOrdering(): Unit = {
+    startBrokers(Seq(100, 101, 102))
+    adminClient = createAdminClient(servers)
+    createTopic(zkClient, topicName,
+      Map(tp0.partition() -> Seq(100, 101),
+          tp1.partition() -> Seq(100, 101)),
+      servers = servers)
+
+    // Throttle to ensure we minimize race conditions and test flakiness
+    throttle(Seq(topicName), throttleSettingForSeconds(10), Map(
+      tp0 -> Seq(100, 101, 102),
+      tp1 -> Seq(100, 101, 102)
+    ))
+
+    adminClient.alterPartitionReassignments(
+      Map(reassignmentEntry(tp0, Seq(100, 101, 102)),
+          reassignmentEntry(tp1, Seq(100, 101, 102))).asJava
+    ).all().get()
+    val apiReassignmentsInProgress = adminClient.listPartitionReassignments(Set(tp0, tp1).asJava).reassignments().get()
+    assertIsReassigning(from = Seq(100, 101), to = Seq(100, 101, 102), apiReassignmentsInProgress.get(tp0))
+    assertIsReassigning(from = Seq(100, 101), to = Seq(100, 101, 102), apiReassignmentsInProgress.get(tp1))
+
+    // API reassignment to the same replicas but a different order
+    adminClient.alterPartitionReassignments(
+      Map(reassignmentEntry(tp0, Seq(102, 101, 100)),
+          reassignmentEntry(tp1, Seq(102, 101, 100))).asJava
+    ).all().get()
+    val apiReassignmentsInProgress2 = adminClient.listPartitionReassignments(Set(tp0, tp1).asJava).reassignments().get()
+    // assert same replicas, ignoring ordering
+    assertIsReassigning(from = Seq(100, 101), to = Seq(100, 101, 102), apiReassignmentsInProgress2.get(tp0))
+    assertIsReassigning(from = Seq(100, 101), to = Seq(100, 101, 102), apiReassignmentsInProgress2.get(tp1))
+
+    resetBrokersThrottle()
+    waitForAllReassignmentsToComplete()
+
+    //Check move occurred
+    val actual = zkClient.getPartitionAssignmentForTopics(Set(topicName))(topicName)
+    assertMoveForPartitionOccurred(Seq(102, 101, 100), actual(tp0.partition()))
+    assertMoveForPartitionOccurred(Seq(102, 101, 100), actual(tp1.partition()))
+  }
+
+  /**
+    * 1. Trigger API reassignment for partitions
+    * 2. Trigger ZK reassignment for partitions
+    * Ensure ZK reassignment overrides API reassignment and znode is deleted
+    */
+  @Test
+  def znodeReassignmentShouldOverrideApiTriggeredReassignment(): Unit = {
+    startBrokers(Seq(100, 101, 102))
+    adminClient = createAdminClient(servers)
+    createTopic(zkClient, topicName,
+      Map(tp0.partition() -> Seq(100),
+          tp1.partition() -> Seq(100)),
+      servers = servers)
+
+    // Throttle to avoid race conditions
+    val throttleSetting = throttleSettingForSeconds(10)
+    throttle(Seq(topicName), throttleSetting, Map(
+      tp0 -> Seq(100, 101, 102),
+      tp1 -> Seq(100, 101, 102)
+    ))
+
+    // API reassignment to 101 for both partitions
+    adminClient.alterPartitionReassignments(
+      Map(reassignmentEntry(tp0, Seq(101)),
+          reassignmentEntry(tp1, Seq(101))).asJava
+    ).all().get()
+    val apiReassignmentsInProgress = adminClient.listPartitionReassignments(Set(tp0, tp1).asJava).reassignments().get()
+    assertIsReassigning(from = Seq(100), to = Seq(101), apiReassignmentsInProgress.get(tp0))
+    assertIsReassigning(from = Seq(100), to = Seq(101), apiReassignmentsInProgress.get(tp1))
+
+    // znode reassignment to 102 for both partitions
+    val topicJson = executeAssignmentJson(Seq(
+      PartitionAssignmentJson(tp0, Seq(102)),
+      PartitionAssignmentJson(tp1, Seq(102))
+    ))
+    ReassignPartitionsCommand.executeAssignment(zkClient, Some(adminClient), topicJson, Throttle(throttleSetting.throttleBytes.toLong))
+    waitUntilTrue(() => {
+      !adminClient.listPartitionReassignments().reassignments().get().isEmpty
+    }, "Controller should have picked up on znode creation", 1000)
+
+    val zkReassignmentsInProgress = adminClient.listPartitionReassignments(Set(tp0, tp1).asJava).reassignments().get()
+    assertIsReassigning(from = Seq(100), to = Seq(102), zkReassignmentsInProgress.get(tp0))
+    assertIsReassigning(from = Seq(100), to = Seq(102), zkReassignmentsInProgress.get(tp1))
+
+    resetBrokersThrottle()
+    waitForZkReassignmentToComplete()
+    assertTrue(adminClient.listPartitionReassignments(Set(tp0, tp1).asJava).reassignments().get().isEmpty)
+    assertEquals(Seq(102), zkClient.getReplicasForPartition(tp0))
+    assertEquals(Seq(102), zkClient.getReplicasForPartition(tp1))
+  }
+
+  /**
+    * 1. Trigger ZK reassignment for TP A-0, A-1
+    * 2. Trigger API reassignment for partitions TP A-1, B-0
+    * 3. Unthrottle A-0, A-1 so the ZK reassignment finishes quickly
+    * 4. Ensure ZK node is emptied out after the API reassignment of 1 finishes
+    */
+  @Test
+  def shouldDeleteReassignmentZnodeAfterApiReassignmentForPartitionCompletes(): Unit = {
+    startBrokers(Seq(100, 101, 102))
+    adminClient = createAdminClient(servers)
+    val tpA0 = new TopicPartition("A", 0)
+    val tpA1 = new TopicPartition("A", 1)
+    val tpB0 = new TopicPartition("B", 0)
+
+    createTopic(zkClient, "A",
+      Map(tpA0.partition() -> Seq(100),
+          tpA1.partition() -> Seq(100)),
+      servers = servers)
+    createTopic(zkClient, "B",
+      Map(tpB0.partition() -> Seq(100)),
+      servers = servers)
+
+    // Throttle to avoid race conditions
+    throttle(Seq("A", "B"), throttleSettingForSeconds(10), Map(
+      tpA0 -> Seq(100, 101, 102),
+      tpA1 -> Seq(100, 101, 102),
+      tpB0 -> Seq(100, 101, 102)
+    ))
+
+    // 1. znode reassignment to 101 for TP A-0, A-1
+    val topicJson = executeAssignmentJson(Seq(
+      PartitionAssignmentJson(tpA0, replicas=Seq(101)),
+      PartitionAssignmentJson(tpA1, replicas=Seq(101))
+    ))
+    ReassignPartitionsCommand.executeAssignment(zkClient, Some(adminClient), topicJson, NoThrottle)
+    waitUntilTrue(() => {
+      !adminClient.listPartitionReassignments().reassignments().get().isEmpty
+    }, "Controller should have picked up on znode creation", 1000)
+    val zkReassignmentsInProgress = adminClient.listPartitionReassignments(Set(tpA0, tpA1).asJava).reassignments().get()
+    assertIsReassigning(from = Seq(100), to = Seq(101), zkReassignmentsInProgress.get(tpA0))
+    assertIsReassigning(from = Seq(100), to = Seq(101), zkReassignmentsInProgress.get(tpA1))
+
+    // 2. API reassignment to 102 for TP A-1, B-0
+    adminClient.alterPartitionReassignments(
+      Map(reassignmentEntry(tpA1, Seq(102)), reassignmentEntry(tpB0, Seq(102))).asJava
+    ).all().get()
+    val apiReassignmentsInProgress = adminClient.listPartitionReassignments(Set(tpA1, tpB0).asJava).reassignments().get()
+    assertIsReassigning(from = Seq(100), to = Seq(102), apiReassignmentsInProgress.get(tpA1))
+    assertIsReassigning(from = Seq(100), to = Seq(102), apiReassignmentsInProgress.get(tpB0))
+
+    // 3. Unthrottle topic A
+    removePartitionReplicaThrottles(Set(tpA0, tpA1))
+    waitForZkReassignmentToComplete()
+    // 4. Ensure the API reassignment not part of the znode is still in progress
+    val leftoverReassignments = adminClient.listPartitionReassignments(Set(tpA0, tpA1, tpB0).asJava).reassignments().get()
+    assertTrue(leftoverReassignments.keySet().asScala.subsetOf(Set(tpA1, tpB0)))
+
+    resetBrokersThrottle()
+    waitForAllReassignmentsToComplete()
+    assertEquals(Seq(101), zkClient.getReplicasForPartition(tpA0))
+    assertEquals(Seq(102), zkClient.getReplicasForPartition(tpA1))
+    assertEquals(Seq(102), zkClient.getReplicasForPartition(tpB0))
+  }
+
+  @Test
+  def shouldBeAbleToCancelThroughApi(): Unit = {
+    startBrokers(Seq(100, 101, 102))
+    adminClient = createAdminClient(servers)
+    createTopic(zkClient, topicName, Map(tp0.partition() -> Seq(100, 101)), servers = servers)
+    // Throttle to ensure we minimize race conditions and test flakiness
+    throttle(Seq(topicName), throttleSettingForSeconds(10), Map(
+      tp0 -> Seq(100, 101, 102)
+    ))
+
+    // move to [102, 101]
+    adminClient.alterPartitionReassignments(
+      Map(reassignmentEntry(tp0, Seq(102, 101))).asJava
+    ).all().get()
+    val apiReassignmentsInProgress = adminClient.listPartitionReassignments().reassignments().get()
+    val tpReassignment = apiReassignmentsInProgress.get(tp0)
+    assertIsReassigning(from = Seq(100, 101), to = Seq(101, 102), tpReassignment)
+
+    adminClient.alterPartitionReassignments(
+      Map(cancelReassignmentEntry(tp0)).asJava
+    ).all().get()
+
+    resetBrokersThrottle()
+    waitForAllReassignmentsToComplete()
+    assertEquals(Seq(100, 101), zkClient.getReplicasForPartition(tp0).sorted) // revert ordering is not guaranteed
+  }
+
+  @Test
+  def shouldBeAbleToCancelZkTriggeredReassignmentThroughApi(): Unit = {
+    startBrokers(Seq(100, 101))
+    adminClient = createAdminClient(servers)
+    createTopic(zkClient, topicName,
+      Map(tp0.partition() -> Seq(100),
+          tp1.partition() -> Seq(100)),
+      servers = servers)
+
+    // Throttle to avoid race conditions
+    throttle(Seq(topicName), throttleSettingForSeconds(10), Map(
+      tp0 -> Seq(100, 101),
+      tp1 -> Seq(100, 101)
+    ))
+
+    val move = Map(
+      tp0 -> Seq(101),
+      tp1 -> Seq(101)
+    )
+    zkClient.setOrCreatePartitionReassignment(move, ZkVersion.MatchAnyVersion)
+    waitUntilTrue(() => {
+      !adminClient.listPartitionReassignments().reassignments().get().isEmpty
+    }, "Controller should have picked up on znode creation", 1000)
+    var reassignmentIsOngoing = adminClient.listPartitionReassignments().reassignments().get().size() > 0
+    assertTrue(reassignmentIsOngoing)
+
+    adminClient.alterPartitionReassignments(
+      Map(cancelReassignmentEntry(tp0), cancelReassignmentEntry(tp1)).asJava
+    ).all().get()
+
+    resetBrokersThrottle()
+    waitForZkReassignmentToComplete()
+    reassignmentIsOngoing = adminClient.listPartitionReassignments().reassignments().get().size() > 0
+    assertFalse(reassignmentIsOngoing)
+    assertEquals(Seq(100), zkClient.getReplicasForPartition(tp0))
+    assertEquals(Seq(100), zkClient.getReplicasForPartition(tp1))
+  }
+
+  /**
+    * Cancel and set reassignments in the same API call.
+    * Even though one cancellation is invalid, ensure the other entries in the request pass
+    */
+  @Test
+  def testCancelAndSetSomeReassignments(): Unit = {
+    startBrokers(Seq(100, 101, 102))
+    adminClient = createAdminClient(servers)
+    val tp2 = new TopicPartition(topicName, 2)
+    val tp3 = new TopicPartition(topicName, 3)
+
+    createTopic(zkClient, topicName,
+      Map(tp0.partition() -> Seq(100), tp1.partition() -> Seq(100), tp2.partition() -> Seq(100), tp3.partition() -> Seq(100)),
+      servers = servers)
+
+    // Throttle to avoid race conditions
+    throttle(Seq(topicName), throttleSettingForSeconds(10), Map(
+      tp0 -> Seq(100, 101, 102),
+      tp1 -> Seq(100, 101, 102),
+      tp2 -> Seq(100, 101, 102),
+      tp3 -> Seq(100, 101, 102)
+    ))
+
+    // API reassignment to 101 for tp0 and tp1
+    adminClient.alterPartitionReassignments(
+      Map(reassignmentEntry(tp0, Seq(101)), reassignmentEntry(tp1, Seq(101))).asJava
+    ).all().get()
+
+    // cancel tp0, reassign tp1 to 102 (override), assign tp2 to 101 (new reassignment) and cancel tp3 (it is not moving)
+    val alterResults = adminClient.alterPartitionReassignments(
+      Map(cancelReassignmentEntry(tp0), reassignmentEntry(tp1, Seq(102)),
+          reassignmentEntry(tp2, Seq(101)), cancelReassignmentEntry(tp3)).asJava
+    ).values()
+    alterResults.get(tp0).get()
+    alterResults.get(tp1).get()
+    alterResults.get(tp2).get()
+    try {
+      alterResults.get(tp3).get()
+    } catch {
+      case exception: Exception =>
+        assertEquals(exception.getCause.getClass, classOf[NoReassignmentInProgressException])
+    }
+
+    resetBrokersThrottle()
+    waitForAllReassignmentsToComplete()
+    assertEquals(Seq(100), zkClient.getReplicasForPartition(tp0))
+    assertEquals(Seq(102), zkClient.getReplicasForPartition(tp1))
+    assertEquals(Seq(101), zkClient.getReplicasForPartition(tp2))
+    assertEquals(Seq(100), zkClient.getReplicasForPartition(tp3))
+  }
+
+  /**
+    * Three different Alter Reassignment calls should all create reassignments
+    */
+  @Test
+  def shouldBeAbleToIncrementallyStackDifferentReassignments(): Unit = {
+    startBrokers(Seq(100, 101))
+    adminClient = createAdminClient(servers)
+    val tpA0 = new TopicPartition("A", 0)
+    val tpA1 = new TopicPartition("A", 1)
+    val tpB0 = new TopicPartition("B", 0)
+
+    createTopic(zkClient, "A",
+      Map(tpA0.partition() -> Seq(100),
+          tpA1.partition() -> Seq(100)),
+      servers = servers)
+    createTopic(zkClient, "B",
+      Map(tpB0.partition() -> Seq(100)),
+      servers = servers)
+
+    // Throttle to avoid race conditions
+    throttle(Seq("A", "B"), throttleSettingForSeconds(10), Map(
+      tpA0 -> Seq(100, 101, 102),
+      tpA1 -> Seq(100, 101, 102),
+      tpB0 -> Seq(100, 101, 102)
+    ))
+
+    adminClient.alterPartitionReassignments(Map(reassignmentEntry(tpA0, Seq(101))).asJava).all().get()
+    val apiReassignmentsInProgress1 = adminClient.listPartitionReassignments().reassignments().get()
+    assertEquals(1, apiReassignmentsInProgress1.size())
+    assertIsReassigning(
+      from = Seq(100), to = Seq(101),
+      apiReassignmentsInProgress1.get(tpA0)
+    )
+
+    adminClient.alterPartitionReassignments(Map(reassignmentEntry(tpA1, Seq(101))).asJava).all().get()
+    val apiReassignmentsInProgress2 = adminClient.listPartitionReassignments().reassignments().get()
+    assertEquals(2, apiReassignmentsInProgress2.size())
+    assertIsReassigning(from = Seq(100), to = Seq(101), apiReassignmentsInProgress2.get(tpA0))
+    assertIsReassigning(
+      from = Seq(100), to = Seq(101),
+      apiReassignmentsInProgress2.get(tpA1)
+    )
+
+    adminClient.alterPartitionReassignments(Map(reassignmentEntry(tpB0, Seq(101))).asJava).all().get()
+    val apiReassignmentsInProgress3 = adminClient.listPartitionReassignments().reassignments().get()
+    assertEquals(s"${apiReassignmentsInProgress3}", 3, apiReassignmentsInProgress3.size())
+    assertIsReassigning(from = Seq(100), to = Seq(101), apiReassignmentsInProgress3.get(tpA0))
+    assertIsReassigning(from = Seq(100), to = Seq(101), apiReassignmentsInProgress3.get(tpA1))
+    assertIsReassigning(
+      from = Seq(100), to = Seq(101),
+      apiReassignmentsInProgress3.get(tpB0)
+    )
+
+    resetBrokersThrottle()
+    waitForAllReassignmentsToComplete()
+    assertEquals(Seq(101), zkClient.getReplicasForPartition(tpA0))
+    assertEquals(Seq(101), zkClient.getReplicasForPartition(tpA1))
+    assertEquals(Seq(101), zkClient.getReplicasForPartition(tpB0))
+  }
+
+  /**
+   * Verifies that partitions can be created for topics not in reassignment and for the topics that are in reassignment
+   * an ReassignmentInProgressException should be thrown. The test creates two topics `topicName` and `otherTopicName`,
+   * the `topicName` topic undergoes partition reassignment and the test validates that during reassignment createPartitions
+   * call throws ReassignmentInProgressException `topicName` topic and for topic `otherTopicName` which is not being reassigned
+   * successfully creates partitions. Further validates that after the reassignment is complete for topic `topicName`
+   * createPartition is successful for that topic.
+   */
+  @Test
+  def shouldCreatePartitionsForTopicNotInReassignment(): Unit = {
+    startBrokers(Seq(100, 101))
+    val otherTopicName = "anyTopic"
+    val otp0 = new TopicPartition(otherTopicName, 0)
+    val otp1 = new TopicPartition(otherTopicName, 1)
+    adminClient = createAdminClient(servers)
+    createTopic(zkClient, topicName,
+      Map(otp0.partition() -> Seq(100),
+          otp1.partition() -> Seq(100)),
+      servers = servers)
+    createTopic(zkClient, otherTopicName,
+      Map(tp0.partition() -> Seq(100),
+          tp1.partition() -> Seq(100)),
+      servers = servers)
+
+    // Throttle to avoid race conditions
+    throttle(Seq(topicName), throttleSettingForSeconds(10), Map(
+      tp0 -> Seq(100, 101),
+      tp1 -> Seq(100, 101)
+    ))
+
+    // Alter `topicName` partition reassignment
+    adminClient.alterPartitionReassignments(
+      Map(reassignmentEntry(tp0, Seq(101)),
+        reassignmentEntry(tp1, Seq(101))).asJava
+    ).all().get()
+    waitUntilTrue(() => {
+      !adminClient.listPartitionReassignments().reassignments().get().isEmpty
+    }, "Controller should have picked up reassignment", 1000)
+
+    def testCreatePartitions(topicName: String, isTopicBeingReassigned: Boolean): Unit = {
+      if (isTopicBeingReassigned)
+        assertTrue("createPartitions for topic under reassignment should throw an exception", intercept[ExecutionException](
+          adminClient.createPartitions(Map(topicName -> NewPartitions.increaseTo(4)).asJava).values.get(topicName).get()).
+          getCause.isInstanceOf[ReassignmentInProgressException])
+      else
+        adminClient.createPartitions(Map(topicName -> NewPartitions.increaseTo(4)).asJava).values.get(topicName).get()
+    }
+
+    // Test case: createPartitions throws ReassignmentInProgressException Topics with partitions in reassignment.
+    testCreatePartitions(topicName, true)
+    // Test case: createPartitions is successful for Topics with partitions NOT in reassignment.
+    testCreatePartitions(otherTopicName, false)
+
+    // complete reassignment
+    resetBrokersThrottle()
+    waitForAllReassignmentsToComplete()
+
+    // Test case: createPartitions is successful for Topics with partitions after reassignment has completed.
+    testCreatePartitions(topicName, false)
+  }
+
+  @Test
+  def shouldFinishReassignmentAfterTopicDeletion() {
+
+    //Given partitions on 3 of 3 brokers
+    val brokers = Array(100, 101, 102)
+
+    // start servers with topic deletion enabled
+    servers = brokers.map(i => createBrokerConfig(i, zkConnect, enableDeleteTopic = true, enableControlledShutdown = false, logDirCount = 3))
+      .map(c => createServer(KafkaConfig.fromProps(c)))
+
+    createTopic(zkClient, topicName, Map(
+      0 -> Seq(100, 101)
+    ), servers = servers)
+
+    val topicPartition = new TopicPartition(topicName, 0)
+    //Given a small throttle to delay partition reassignment
+    val initialThrottle = Throttle(1)
+
+    val numMessages = 1000
+    val msgSize = 100 * 1000
+    produceMessages(topicName, numMessages, acks = 0, msgSize)
+
+    //Start rebalance which will move replica on 100 -> replica on 102
+    val newAssignment = generateAssignment(zkClient, Array(101, 102), generateAssignmentJson(topicName), true)._1
+
+    // Start reassignment
+    ReassignPartitionsCommand.executeAssignment(zkClient, None,
+      ReassignPartitionsCommand.formatAsReassignmentJson(newAssignment, Map.empty), initialThrottle)
+
+    // make sure all brokers have received partition reassignment request
+    TestUtils.waitUntilTrue(() => servers.forall(_.getLogManager().getLog(topicPartition).isDefined),
+      "reassignment for topic not started.")
+
+    // delete the topic
+    adminZkClient.deleteTopic(topicName)
+    TestUtils.verifyTopicDeletion(zkClient, topicName, 1, servers)
+
+    //Await completion
+    waitForZkReassignmentToComplete()
+  }
+
+  /**
+    * Asserts that a replica is being reassigned from the given replicas to the target replicas
+    */
+  def assertIsReassigning(from: Seq[Int], to: Seq[Int], reassignment: PartitionReassignment): Unit = {
+    assertReplicas((from ++ to).distinct, reassignment.replicas())
+    assertReplicas(to.filterNot(from.contains(_)), reassignment.addingReplicas())
+    assertReplicas(from.filterNot(to.contains(_)), reassignment.removingReplicas())
+  }
+
+  /**
+   * Asserts that a topic's reassignments completed and span across the expected replicas
+   */
+  def assertMoveForTopicOccurred(expectedReplicas: Seq[Int],
+                                 partitionAssignments: Map[Int, ReplicaAssignment]): Unit = {
+    assertEquals(expectedReplicas, partitionAssignments.values.flatMap(_.replicas).toSeq.distinct.sorted)
+    assertTrue(partitionAssignments.values.flatMap(_.addingReplicas).isEmpty)
+    assertTrue(partitionAssignments.values.flatMap(_.removingReplicas).isEmpty)
+  }
+
+  /**
+   * Asserts that a partition moved to the exact expected replicas in the specific order
+   */
+  def assertMoveForPartitionOccurred(expectedReplicas: Seq[Int],
+                                     partitionAssignment: ReplicaAssignment): Unit = {
+    assertEquals(expectedReplicas, partitionAssignment.replicas)
+    assertTrue(partitionAssignment.addingReplicas.isEmpty)
+    assertTrue(partitionAssignment.removingReplicas.isEmpty)
+  }
+
+  /**
+   * Asserts that two replica sets are equal, ignoring ordering
+   */
+  def assertReplicas(expectedReplicas: Seq[Int], receivedReplicas: java.util.List[Integer]): Unit = {
+    assertEquals(expectedReplicas.sorted, receivedReplicas.asScala.map(_.toInt).sorted)
+  }
+
+  def throttleAllBrokersReplication(throttleBytes: String): Unit = {
+    val throttleConfigs = Seq(
+      new AlterConfigOp(new ConfigEntry(DynamicConfig.Broker.LeaderReplicationThrottledRateProp, throttleBytes), AlterConfigOp.OpType.SET),
+      new AlterConfigOp(new ConfigEntry(DynamicConfig.Broker.FollowerReplicationThrottledRateProp, throttleBytes), AlterConfigOp.OpType.SET)
+    ).asJavaCollection
+
+    adminClient.incrementalAlterConfigs(
+      brokerIds.map { brokerId =>
+        new ConfigResource(ConfigResource.Type.BROKER, brokerId.toString) -> throttleConfigs
+      }.toMap.asJava
+    ).all().get()
+  }
+
+  def resetBrokersThrottle(): Unit = throttleAllBrokersReplication(Int.MaxValue.toString)
+
+  def assignThrottledPartitionReplicas(allReplicasByPartition: Map[TopicPartition, Seq[Int]]): Unit = {
+    val throttles = allReplicasByPartition.groupBy(_._1.topic()).map {
+      case (topic, replicasByPartition) =>
+        new ConfigResource(ConfigResource.Type.TOPIC, topic) -> Seq(
+          new AlterConfigOp(new ConfigEntry(LogConfig.LeaderReplicationThrottledReplicasProp, formatReplicaThrottles(replicasByPartition)), AlterConfigOp.OpType.SET),
+          new AlterConfigOp(new ConfigEntry(LogConfig.FollowerReplicationThrottledReplicasProp, formatReplicaThrottles(replicasByPartition)), AlterConfigOp.OpType.SET)
+        ).asJavaCollection
+    }
+    adminClient.incrementalAlterConfigs(throttles.asJava).all().get()
+  }
+
+  def removePartitionReplicaThrottles(partitions: Set[TopicPartition]): Unit = {
+    val throttles = partitions.map {
+      tp =>
+        new ConfigResource(ConfigResource.Type.TOPIC, tp.topic()) -> Seq(
+          new AlterConfigOp(new ConfigEntry(LogConfig.LeaderReplicationThrottledReplicasProp, ""), AlterConfigOp.OpType.DELETE),
+          new AlterConfigOp(new ConfigEntry(LogConfig.FollowerReplicationThrottledReplicasProp, ""), AlterConfigOp.OpType.DELETE)
+        ).asJavaCollection
+    }.toMap
+    adminClient.incrementalAlterConfigs(throttles.asJava).all().get()
+  }
+
+  def formatReplicaThrottles(moves: Map[TopicPartition, Seq[Int]]): String =
+    moves.flatMap { case (tp, assignment) =>
+      assignment.map(replicaId => s"${tp.partition}:$replicaId")
+    }.mkString(",")
+
+  def reassignmentEntry(tp: TopicPartition, replicas: Seq[Int]): (TopicPartition, java.util.Optional[NewPartitionReassignment]) =
+    tp -> Optional.of(new NewPartitionReassignment((replicas.map(_.asInstanceOf[Integer]).asJava)))
+
+  def cancelReassignmentEntry(tp: TopicPartition): (TopicPartition, java.util.Optional[NewPartitionReassignment]) =
+    tp -> java.util.Optional.empty()
+
+  def waitForZkReassignmentToComplete(pause: Long = 100L): Unit = {
     waitUntilTrue(() => !zkClient.reassignPartitionsInProgress,
       s"Znode ${ReassignPartitionsZNode.path} wasn't deleted", pause = pause)
   }
 
-  def json(topic: String*): String = {
+  def awaitReassignmentInProgress(topicPartition: TopicPartition): Unit = {
+    waitUntilTrue(() => isAssignmentInProgress(topicPartition),
+      "Timed out waiting for expected reassignment to begin")
+  }
+
+  def isAssignmentInProgress(topicPartition: TopicPartition): Boolean = {
+    val reassignments = adminClient.listPartitionReassignments().reassignments().get()
+    reassignments.asScala.get(topicPartition).isDefined
+  }
+
+  def waitForAllReassignmentsToComplete(pause: Long = 100L): Unit = {
+    waitUntilTrue(() => adminClient.listPartitionReassignments().reassignments().get().isEmpty,
+      s"There still are ongoing reassignments", pause = pause)
+  }
+
+  def generateAssignmentJson(topic: String*): String = {
     val topicStr = topic.map { t => "{\"topic\": \"" + t + "\"}" }.mkString(",")
     s"""{"topics": [$topicStr],"version":1}"""
   }
 
+  def executeAssignmentJson(partitions: Seq[PartitionAssignmentJson]): String =
+    s"""{"version":1,"partitions":[${partitions.map(_.toJson).mkString(",")}]}"""
+
+  case class PartitionAssignmentJson(topicPartition: TopicPartition, replicas: Seq[Int],
+                                     logDirectories: Option[Seq[String]] = None) {
+    def toJson: String = {
+      val logDirsSuffix = logDirectories match {
+        case Some(dirs) => s""","log_dirs":[${dirs.map("\"" + _ + "\"").mkString(",")}]"""
+        case None => ""
+      }
+      s"""{"topic":"${topicPartition.topic()}","partition":${topicPartition.partition()}""" +
+        s""","replicas":[${replicas.mkString(",")}]""" +
+        s"$logDirsSuffix}"
+    }
+  }
+
+  case class ThrottleSetting(throttleBytes: String, numMessages: Int, messageSizeBytes: Int)
+
+  def throttleSettingForSeconds(secondsDuration: Int): ThrottleSetting = {
+    val throttle = 1000 * 1000 // 1 MB/s throttle
+    val messageSize = 100 * 100 // 0.01 MB message size
+    val messagesPerSecond = throttle / messageSize
+    ThrottleSetting(throttle.toString, messagesPerSecond * secondsDuration, messageSize)
+  }
+
+  def throttle(topics: Seq[String], throttle: ThrottleSetting, replicasToThrottle: Map[TopicPartition, Seq[Int]]): Unit = {
+    val messagesPerTopic = throttle.numMessages / topics.size
+    for (topic <- topics) {
+      produceMessages(topic, numMessages = messagesPerTopic, acks = 0, valueLength = throttle.messageSizeBytes)
+    }
+    throttleAllBrokersReplication(throttle.throttleBytes)
+    assignThrottledPartitionReplicas(replicasToThrottle)
+  }
+
   private def produceMessages(topic: String, numMessages: Int, acks: Int, valueLength: Int): Unit = {
     val records = (0 until numMessages).map(_ => new ProducerRecord[Array[Byte], Array[Byte]](topic,
       new Array[Byte](valueLength)))
diff --git a/core/src/test/scala/unit/kafka/admin/ReassignPartitionsCommandArgsTest.scala b/core/src/test/scala/unit/kafka/admin/ReassignPartitionsCommandArgsTest.scala
index 58768c637c26f..7037f2980ce8e 100644
--- a/core/src/test/scala/unit/kafka/admin/ReassignPartitionsCommandArgsTest.scala
+++ b/core/src/test/scala/unit/kafka/admin/ReassignPartitionsCommandArgsTest.scala
@@ -19,17 +19,16 @@ package kafka.admin
 import kafka.utils.Exit
 import org.junit.Assert._
 import org.junit.{After, Before, Test}
-import org.scalatest.junit.JUnitSuite
 
-class ReassignPartitionsCommandArgsTest extends JUnitSuite {
+class ReassignPartitionsCommandArgsTest {
 
   @Before
-  def setUp() {
+  def setUp(): Unit = {
     Exit.setExitProcedure((_, message) => throw new IllegalArgumentException(message.orNull))
   }
 
   @After
-  def tearDown() {
+  def tearDown(): Unit = {
     Exit.resetExitProcedure()
   }
 
diff --git a/core/src/test/scala/unit/kafka/admin/ReassignPartitionsCommandTest.scala b/core/src/test/scala/unit/kafka/admin/ReassignPartitionsCommandTest.scala
index 128919f83ebbb..790e3d8076af3 100644
--- a/core/src/test/scala/unit/kafka/admin/ReassignPartitionsCommandTest.scala
+++ b/core/src/test/scala/unit/kafka/admin/ReassignPartitionsCommandTest.scala
@@ -31,24 +31,26 @@ import org.easymock.EasyMock._
 import org.easymock.{Capture, CaptureType, EasyMock}
 import org.junit.{After, Before, Test}
 import org.junit.Assert.{assertEquals, assertFalse, assertNull, assertTrue}
+import org.scalatest.Assertions.fail
 
 import scala.collection.JavaConverters._
 import org.apache.kafka.common.TopicPartition
 
 import scala.collection.mutable
+import scala.collection.Seq
 
 class ReassignPartitionsCommandTest extends ZooKeeperTestHarness with Logging {
   var servers: Seq[KafkaServer] = Seq()
   var calls = 0
 
   @After
-  override def tearDown() {
+  override def tearDown(): Unit = {
     TestUtils.shutdownServers(servers)
     super.tearDown()
   }
 
   @Test
-  def shouldFindMovingReplicas() {
+  def shouldFindMovingReplicas(): Unit = {
     val control = new TopicPartition("topic1", 1) -> Seq(100, 102)
     val assigner = new ReassignPartitionsCommand(null, null, null, null, null)
 
@@ -71,7 +73,7 @@ class ReassignPartitionsCommandTest extends ZooKeeperTestHarness with Logging {
   }
 
   @Test
-  def shouldFindMovingReplicasWhenProposedIsSubsetOfExisting() {
+  def shouldFindMovingReplicasWhenProposedIsSubsetOfExisting(): Unit = {
     val assigner = new ReassignPartitionsCommand(null, null, null, null, null)
 
     //Given we have more existing partitions than we are proposing
@@ -106,7 +108,7 @@ class ReassignPartitionsCommandTest extends ZooKeeperTestHarness with Logging {
   }
 
   @Test
-  def shouldFindMovingReplicasMultiplePartitions() {
+  def shouldFindMovingReplicasMultiplePartitions(): Unit = {
     val control = new TopicPartition("topic1", 2) -> Seq(100, 102)
     val assigner = new ReassignPartitionsCommand(null, null, null, null, null)
 
@@ -131,7 +133,7 @@ class ReassignPartitionsCommandTest extends ZooKeeperTestHarness with Logging {
   }
 
   @Test
-  def shouldFindMovingReplicasMultipleTopics() {
+  def shouldFindMovingReplicasMultipleTopics(): Unit = {
     val control = new TopicPartition("topic1", 1) -> Seq(100, 102)
     val assigner = new ReassignPartitionsCommand(null, null, null, null, null)
 
@@ -163,7 +165,7 @@ class ReassignPartitionsCommandTest extends ZooKeeperTestHarness with Logging {
   }
 
   @Test
-  def shouldFindMovingReplicasMultipleTopicsAndPartitions() {
+  def shouldFindMovingReplicasMultipleTopicsAndPartitions(): Unit = {
     val assigner = new ReassignPartitionsCommand(null, null, null, null, null)
 
     //Given
@@ -206,7 +208,7 @@ class ReassignPartitionsCommandTest extends ZooKeeperTestHarness with Logging {
   }
 
   @Test
-  def shouldFindTwoMovingReplicasInSamePartition() {
+  def shouldFindTwoMovingReplicasInSamePartition(): Unit = {
     val control = new TopicPartition("topic1", 1) -> Seq(100, 102)
     val assigner = new ReassignPartitionsCommand(null, null, null, null, null)
 
@@ -424,7 +426,7 @@ class ReassignPartitionsCommandTest extends ZooKeeperTestHarness with Logging {
   }
 
   @Test
-  def testPartitionReassignmentWithLeaderInNewReplicas() {
+  def testPartitionReassignmentWithLeaderInNewReplicas(): Unit = {
     val expectedReplicaAssignment = Map(0  -> List(0, 1, 2))
     val topic = "test"
     // create brokers
@@ -453,7 +455,7 @@ class ReassignPartitionsCommandTest extends ZooKeeperTestHarness with Logging {
   }
 
   @Test
-  def testPartitionReassignmentWithLeaderNotInNewReplicas() {
+  def testPartitionReassignmentWithLeaderNotInNewReplicas(): Unit = {
     val expectedReplicaAssignment = Map(0  -> List(0, 1, 2))
     val topic = "test"
     // create brokers
@@ -481,7 +483,7 @@ class ReassignPartitionsCommandTest extends ZooKeeperTestHarness with Logging {
   }
 
   @Test
-  def testPartitionReassignmentNonOverlappingReplicas() {
+  def testPartitionReassignmentNonOverlappingReplicas(): Unit = {
     val expectedReplicaAssignment = Map(0  -> List(0, 1))
     val topic = "test"
     // create brokers
@@ -509,7 +511,7 @@ class ReassignPartitionsCommandTest extends ZooKeeperTestHarness with Logging {
   }
 
   @Test
-  def testReassigningNonExistingPartition() {
+  def testReassigningNonExistingPartition(): Unit = {
     val topic = "test"
     // create brokers
     servers = TestUtils.createBrokerConfigs(4, zkConnect, false).map(b => TestUtils.createServer(KafkaConfig.fromProps(b)))
@@ -524,11 +526,11 @@ class ReassignPartitionsCommandTest extends ZooKeeperTestHarness with Logging {
   }
 
   @Test
-  def testResumePartitionReassignmentThatWasCompleted() {
-    val expectedReplicaAssignment = Map(0  -> List(0, 1))
+  def testResumePartitionReassignmentThatWasCompleted(): Unit = {
+    val initialAssignment = Map(0  -> List(0, 2))
     val topic = "test"
     // create the topic
-    adminZkClient.createTopicWithAssignment(topic, config = new Properties, expectedReplicaAssignment)
+    adminZkClient.createTopicWithAssignment(topic, config = new Properties, initialAssignment)
     // put the partition in the reassigned path as well
     // reassign partition 0
     val newReplicas = Seq(0, 1)
diff --git a/core/src/test/scala/unit/kafka/admin/ResetConsumerGroupOffsetTest.scala b/core/src/test/scala/unit/kafka/admin/ResetConsumerGroupOffsetTest.scala
index 007edd9a9730a..f9322b36e76f1 100644
--- a/core/src/test/scala/unit/kafka/admin/ResetConsumerGroupOffsetTest.scala
+++ b/core/src/test/scala/unit/kafka/admin/ResetConsumerGroupOffsetTest.scala
@@ -22,13 +22,17 @@ 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.Assert._
 import org.junit.Test
 
+import scala.collection.JavaConverters._
+import scala.collection.Seq
+
 class TimeConversionTests {
 
   @Test
-  def testDateTimeFormats() {
+  def testDateTimeFormats(): Unit = {
     //check valid formats
     invokeGetDateTimeMethod(new SimpleDateFormat("yyyy-MM-dd'T'HH:mm:ss.SSS"))
     invokeGetDateTimeMethod(new SimpleDateFormat("yyyy-MM-dd'T'HH:mm:ss.SSSZ"))
@@ -52,7 +56,7 @@ class TimeConversionTests {
     }
   }
 
-  private def invokeGetDateTimeMethod(format: SimpleDateFormat) {
+  private def invokeGetDateTimeMethod(format: SimpleDateFormat): Unit = {
     val checkpoint = new Date()
     val timestampString = format.format(checkpoint)
     ConsumerGroupCommand.convertTimestamp(timestampString)
@@ -71,42 +75,111 @@ class TimeConversionTests {
   * - 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))  
-  } 
+      .map(KafkaConfig.fromProps(_, overridingProps))
+  }
+
+  private def basicArgs: Array[String] = {
+    Array("--reset-offsets",
+      "--bootstrap-server", brokerList,
+      "--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() {
-    val args = Array("--bootstrap-server", brokerList, "--reset-offsets", "--group", "missing.group", "--all-topics",
-      "--to-current", "--execute")
+  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().coordinator.host() == "localhost"
-    }, "Can't find a coordinator.", maxRetries = 3)
-    val resetOffsets = consumerGroupCommand.resetOffsets()
+      consumerGroupCommand.collectGroupState(group).coordinator.host() == "localhost"
+    }, "Can't find a coordinator")
+    val resetOffsets = consumerGroupCommand.resetOffsets()(group)
     assertEquals(Map.empty, resetOffsets)
-    assertEquals(resetOffsets, committedOffsets(group = "missing.group"))
+    assertEquals(resetOffsets, committedOffsets(group = group))
   }
 
   @Test
   def testResetOffsetsExistingTopic(): Unit = {
-    val args = Array("--bootstrap-server", brokerList, "--reset-offsets", "--group", "new.group", "--topic", topic,
-      "--to-offset", "50")
+    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, group = "new.group")
+    resetAndAssertOffsets(args ++ Array("--execute"), expectedOffset = 50)
   }
 
   @Test
-  def testResetOffsetsToLocalDateTime() {
+  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)
@@ -117,13 +190,12 @@ class ResetConsumerGroupOffsetTest extends ConsumerGroupCommandTest {
     awaitConsumerProgress(count = 100L)
     executor.shutdown()
 
-    val args = Array("--bootstrap-server", brokerList, "--reset-offsets", "--group", group, "--all-topics",
-      "--to-datetime", format.format(calendar.getTime), "--execute")
+    val args = buildArgsForGroup(group, "--all-topics", "--to-datetime", format.format(calendar.getTime), "--execute")
     resetAndAssertOffsets(args, expectedOffset = 0)
   }
 
   @Test
-  def testResetOffsetsToZonedDateTime() {
+  def testResetOffsetsToZonedDateTime(): Unit = {
     val format = new SimpleDateFormat("yyyy-MM-dd'T'HH:mm:ss.SSSXXX")
 
     produceMessages(topic, 50)
@@ -134,115 +206,102 @@ class ResetConsumerGroupOffsetTest extends ConsumerGroupCommandTest {
     awaitConsumerProgress(count = 100L)
     executor.shutdown()
 
-    val args = Array("--bootstrap-server", brokerList, "--reset-offsets", "--group", group, "--all-topics",
-      "--to-datetime", format.format(checkpoint), "--execute")
+    val args = buildArgsForGroup(group, "--all-topics", "--to-datetime", format.format(checkpoint), "--execute")
     resetAndAssertOffsets(args, expectedOffset = 50)
   }
 
   @Test
-  def testResetOffsetsByDuration() {
-    val args = Array("--bootstrap-server", brokerList, "--reset-offsets", "--group", group, "--all-topics",
-      "--by-duration", "PT1M", "--execute")
-    produceConsumeAndShutdown(topic, totalMessages = 100)
+  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() {
-    val args = Array("--bootstrap-server", brokerList, "--reset-offsets", "--group", group, "--all-topics",
-      "--by-duration", "PT0.1S", "--execute")
-    produceConsumeAndShutdown(topic, totalMessages = 100)
+  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 testResetOffsetsToEarliest() {
-    val args = Array("--bootstrap-server", brokerList, "--reset-offsets", "--group", group, "--all-topics",
-      "--to-earliest", "--execute")
-    produceConsumeAndShutdown(topic, totalMessages = 100)
+  def testResetOffsetsToEarliest(): Unit = {
+    val args = buildArgsForGroup(group, "--all-topics", "--to-earliest", "--execute")
+    produceConsumeAndShutdown(topic, group, totalMessages = 100)
     resetAndAssertOffsets(args, expectedOffset = 0)
   }
 
   @Test
-  def testResetOffsetsToLatest() {
-    val args = Array("--bootstrap-server", brokerList, "--reset-offsets", "--group", group, "--all-topics",
-      "--to-latest", "--execute")
-    produceConsumeAndShutdown(topic, totalMessages = 100)
+  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() {
-    val args = Array("--bootstrap-server", brokerList, "--reset-offsets", "--group", group, "--all-topics",
-      "--to-current", "--execute")
-    produceConsumeAndShutdown(topic, totalMessages = 100)
+  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() {
-    val args = Array("--bootstrap-server", brokerList, "--reset-offsets", "--group", group, "--all-topics",
-      "--to-offset", "1", "--execute")
-    produceConsumeAndShutdown(topic, totalMessages = 100)
+  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() {
-    val args = Array("--bootstrap-server", brokerList, "--reset-offsets", "--group", group, "--all-topics",
-      "--shift-by", "50", "--execute")
-    produceConsumeAndShutdown(topic, totalMessages = 100)
+  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() {
-    val args = Array("--bootstrap-server", brokerList, "--reset-offsets", "--group", group, "--all-topics",
-      "--shift-by", "-50", "--execute")
-    produceConsumeAndShutdown(topic, totalMessages = 100)
+  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() {
-    val args = Array("--bootstrap-server", brokerList, "--reset-offsets", "--group", group, "--all-topics",
-      "--shift-by", "-150", "--execute")
-    produceConsumeAndShutdown(topic, totalMessages = 100)
+  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() {
-    val args = Array("--bootstrap-server", brokerList, "--reset-offsets", "--group", group, "--all-topics",
-      "--shift-by", "150", "--execute")
-    produceConsumeAndShutdown(topic, totalMessages = 100)
+  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() {
-    val args = Array("--bootstrap-server", brokerList, "--reset-offsets", "--group", group, "--topic", topic,
-      "--to-earliest", "--execute")
-    produceConsumeAndShutdown(topic, totalMessages = 100)
+  def testResetOffsetsToEarliestOnOneTopic(): Unit = {
+    val args = buildArgsForGroup(group, "--topic", topic, "--to-earliest", "--execute")
+    produceConsumeAndShutdown(topic, group, totalMessages = 100)
     resetAndAssertOffsets(args, expectedOffset = 0)
   }
 
   @Test
-  def testResetOffsetsToEarliestOnOneTopicAndPartition() {
+  def testResetOffsetsToEarliestOnOneTopicAndPartition(): Unit = {
     val topic = "bar"
     createTopic(topic, 2, 1)
 
-    val args = Array("--bootstrap-server", brokerList, "--reset-offsets", "--group", group, "--topic",
-      s"$topic:1", "--to-earliest", "--execute")
+    val args = buildArgsForGroup(group, "--topic", s"$topic:1", "--to-earliest", "--execute")
     val consumerGroupCommand = getConsumerGroupService(args)
 
-    produceConsumeAndShutdown(topic, totalMessages = 100, numConsumers = 2)
+    produceConsumeAndShutdown(topic, group, totalMessages = 100, numConsumers = 2)
     val priorCommittedOffsets = committedOffsets(topic = topic)
 
     val tp0 = new TopicPartition(topic, 0)
@@ -254,23 +313,22 @@ class ResetConsumerGroupOffsetTest extends ConsumerGroupCommandTest {
   }
 
   @Test
-  def testResetOffsetsToEarliestOnTopics() {
+  def testResetOffsetsToEarliestOnTopics(): Unit = {
     val topic1 = "topic1"
     val topic2 = "topic2"
     createTopic(topic1, 1, 1)
     createTopic(topic2, 1, 1)
 
-    val args = Array("--bootstrap-server", brokerList, "--reset-offsets", "--group", group, "--topic", topic1,
-      "--topic", topic2, "--to-earliest", "--execute")
+    val args = buildArgsForGroup(group, "--topic", topic1, "--topic", topic2, "--to-earliest", "--execute")
     val consumerGroupCommand = getConsumerGroupService(args)
 
-    produceConsumeAndShutdown(topic1, 100, 1)
-    produceConsumeAndShutdown(topic2, 100, 1)
+    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)
+    val allResetOffsets = resetOffsets(consumerGroupCommand)(group).mapValues(_.offset).toMap
     assertEquals(Map(tp1 -> 0L, tp2 -> 0L), allResetOffsets)
     assertEquals(Map(tp1 -> 0L), committedOffsets(topic1))
     assertEquals(Map(tp2 -> 0L), committedOffsets(topic2))
@@ -280,26 +338,25 @@ class ResetConsumerGroupOffsetTest extends ConsumerGroupCommandTest {
   }
 
   @Test
-  def testResetOffsetsToEarliestOnTopicsAndPartitions() {
+  def testResetOffsetsToEarliestOnTopicsAndPartitions(): Unit = {
     val topic1 = "topic1"
     val topic2 = "topic2"
 
     createTopic(topic1, 2, 1)
     createTopic(topic2, 2, 1)
 
-    val args = Array("--bootstrap-server", brokerList, "--reset-offsets", "--group", group, "--topic",
-      s"$topic1:1", "--topic", s"$topic2:1", "--to-earliest", "--execute")
+    val args = buildArgsForGroup(group, "--topic", s"$topic1:1", "--topic", s"$topic2:1", "--to-earliest", "--execute")
     val consumerGroupCommand = getConsumerGroupService(args)
 
-    produceConsumeAndShutdown(topic1, 100, 2)
-    produceConsumeAndShutdown(topic2, 100, 2)
+    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)
+    val allResetOffsets = resetOffsets(consumerGroupCommand)(group).mapValues(_.offset).toMap
     assertEquals(Map(tp1 -> 0, tp2 -> 0), allResetOffsets)
 
     assertEquals(priorCommittedOffsets1 + (tp1 -> 0L), committedOffsets(topic1))
@@ -310,38 +367,84 @@ class ResetConsumerGroupOffsetTest extends ConsumerGroupCommandTest {
   }
 
   @Test
-  def testResetOffsetsExportImportPlan() {
+  // 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 = Array("--bootstrap-server", brokerList, "--reset-offsets", "--group", group, "--all-topics",
-      "--to-offset", "2", "--export")
+    val cgcArgs = buildArgsForGroup(group, "--all-topics", "--to-offset", "2", "--export")
     val consumerGroupCommand = getConsumerGroupService(cgcArgs)
 
-    produceConsumeAndShutdown(topic, 100, 2)
+    produceConsumeAndShutdown(topic = topic, group = group, totalMessages = 100, numConsumers = 2)
 
     val file = File.createTempFile("reset", ".csv")
     file.deleteOnExit()
 
     val exportedOffsets = consumerGroupCommand.resetOffsets()
     val bw = new BufferedWriter(new FileWriter(file))
-    bw.write(consumerGroupCommand.exportOffsetsToReset(exportedOffsets))
+    bw.write(consumerGroupCommand.exportOffsetsToCsv(exportedOffsets))
     bw.close()
-    assertEquals(Map(tp0 -> 2L, tp1 -> 2L), exportedOffsets.mapValues(_.offset))
+    assertEquals(Map(tp0 -> 2L, tp1 -> 2L), exportedOffsets(group).mapValues(_.offset).toMap)
 
-    val cgcArgsExec = Array("--bootstrap-server", brokerList, "--reset-offsets", "--group", group, "--all-topics",
-      "--from-file", file.getCanonicalPath, "--dry-run")
+    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.mapValues(_.offset))
+    assertEquals(Map(tp0 -> 2L, tp1 -> 2L), importedOffsets(group).mapValues(_.offset).toMap)
+
+    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, numConsumers = 2)
+    produceConsumeAndShutdown(topic = topic2, group = group2, totalMessages = 100, numConsumers = 5)
+
+    val file = File.createTempFile("reset", ".csv")
+    file.deleteOnExit()
+
+    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).mapValues(_.offset).toMap)
+    assertEquals(Map(t2p0 -> 2L, t2p1 -> 2L), exportedOffsets(group2).mapValues(_.offset).toMap)
+
+    // 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).mapValues(_.offset).toMap)
+    assertEquals(Map(t2p0 -> 2L, t2p1 -> 2L), importedOffsets(group2).mapValues(_.offset).toMap)
+
+    // 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).mapValues(_.offset).toMap)
 
     adminZkClient.deleteTopic(topic)
   }
 
   @Test(expected = classOf[OptionException])
-  def testResetWithUnrecognizedNewConsumerOption() {
+  def testResetWithUnrecognizedNewConsumerOption(): Unit = {
     val cgcArgs = Array("--new-consumer", "--bootstrap-server", brokerList, "--reset-offsets", "--group", group, "--all-topics",
       "--to-offset", "2", "--export")
     getConsumerGroupService(cgcArgs)
@@ -353,30 +456,54 @@ class ResetConsumerGroupOffsetTest extends ConsumerGroupCommandTest {
     TestUtils.produceMessages(servers, records, acks = 1)
   }
 
-  private def produceConsumeAndShutdown(topic: String, totalMessages: Int, numConsumers: Int = 1) {
+  private def produceConsumeAndShutdown(topic: String, group: String, totalMessages: Int, numConsumers: Int = 1): Unit = {
     produceMessages(topic, totalMessages)
-    val executor =  addConsumerGroupExecutor(numConsumers, topic)
-    awaitConsumerProgress(topic, totalMessages)
+    val executor = addConsumerGroupExecutor(numConsumers = numConsumers, topic = topic, group = group)
+    awaitConsumerProgress(topic, group, totalMessages)
     executor.shutdown()
   }
 
-  private def awaitConsumerProgress(topic: String = topic, count: Long): Unit = {
-    TestUtils.waitUntilTrue(() => {
-      val offsets = committedOffsets(topic).values
-      count == offsets.sum
-    }, "Expected that consumer group has consumed all messages from topic/partition.")
+  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 resetAndAssertOffsets(args: Array[String],
-                                           expectedOffset: Long,
-                                           group: String = group,
-                                           dryRun: Boolean = false): Unit = {
+                                    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 {
-      val priorOffsets = committedOffsets(group = group)
-      val expectedOffsets = Map(new TopicPartition(topic, 0) -> expectedOffset)
-      assertEquals(expectedOffsets, resetOffsets(consumerGroupCommand))
-      assertEquals(if (dryRun) priorOffsets else expectedOffsets, committedOffsets(group = group))
+      for {
+        topic <- topics
+        (group, partitionInfo) <- resetOffsetsResultByGroup
+      } {
+        val priorOffsets = committedOffsets(topic = topic, group = group)
+        assertEquals(expectedOffsets(topic), partitionInfo.filter(partitionInfo => partitionInfo._1.topic() == topic).mapValues(_.offset).toMap)
+        assertEquals(if (dryRun) priorOffsets else expectedOffsets(topic), committedOffsets(topic = topic, group = group))
+      }
     } finally {
       consumerGroupCommand.close()
     }
@@ -386,14 +513,16 @@ class ResetConsumerGroupOffsetTest extends ConsumerGroupCommandTest {
                                              expectedOffsets: Map[TopicPartition, Long],
                                              topic: String): Unit = {
     val allResetOffsets = resetOffsets(consumerGroupService)
-    allResetOffsets.foreach { case (tp, offset) =>
-      assertEquals(offset, expectedOffsets(tp))
+    for {
+      (group, offsetsInfo) <- allResetOffsets
+      (tp, offsetMetadata) <- offsetsInfo
+    } {
+      assertEquals(offsetMetadata.offset(), expectedOffsets(tp))
+      assertEquals(expectedOffsets, committedOffsets(topic, group))
     }
-    assertEquals(expectedOffsets, committedOffsets(topic))
   }
 
-  private def resetOffsets(consumerGroupService: ConsumerGroupService): Map[TopicPartition, Long] = {
-    consumerGroupService.resetOffsets().mapValues(_.offset)
+  private def resetOffsets(consumerGroupService: ConsumerGroupService) = {
+    consumerGroupService.resetOffsets()
   }
-
 }
diff --git a/core/src/test/scala/unit/kafka/admin/TestAdminUtils.scala b/core/src/test/scala/unit/kafka/admin/TestAdminUtils.scala
deleted file mode 100644
index 7cbad05bbe51a..0000000000000
--- a/core/src/test/scala/unit/kafka/admin/TestAdminUtils.scala
+++ /dev/null
@@ -1,29 +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.ZkUtils
-
-@deprecated("This class is deprecated since AdminUtilities will be replaced by kafka.zk.AdminZkClient.", "1.1.0")
-class TestAdminUtils extends AdminUtilities {
-  override def changeBrokerConfig(zkUtils: ZkUtils, brokerIds: Seq[Int], configs: Properties): Unit = {}
-  override def fetchEntityConfig(zkUtils: ZkUtils, entityType: String, entityName: String): Properties = {new Properties}
-  override def changeClientIdConfig(zkUtils: ZkUtils, clientId: String, configs: Properties): Unit = {}
-  override def changeUserOrUserClientIdConfig(zkUtils: ZkUtils, sanitizedEntityName: String, configs: Properties): Unit = {}
-  override def changeTopicConfig(zkUtils: ZkUtils, topic: String, configs: Properties): Unit = {}
-}
diff --git a/core/src/test/scala/unit/kafka/admin/TopicCommandTest.scala b/core/src/test/scala/unit/kafka/admin/TopicCommandTest.scala
index 008fd662b79d2..652bf95593b47 100644
--- a/core/src/test/scala/unit/kafka/admin/TopicCommandTest.scala
+++ b/core/src/test/scala/unit/kafka/admin/TopicCommandTest.scala
@@ -27,6 +27,7 @@ import org.apache.kafka.common.internals.Topic
 import org.junit.Assert._
 import org.junit.rules.TestName
 import org.junit.{After, Before, Rule, Test}
+import org.scalatest.Assertions.intercept
 
 import scala.util.Random
 
@@ -75,7 +76,7 @@ class TopicCommandTest extends ZooKeeperTestHarness with Logging with RackAwareT
   }
 
   @Test
-  def testCreateIfNotExists() {
+  def testCreateIfNotExists(): Unit = {
     // create brokers
     val brokers = List(0, 1, 2)
     TestUtils.createBrokersInZk(zkClient, brokers)
@@ -114,7 +115,7 @@ class TopicCommandTest extends ZooKeeperTestHarness with Logging with RackAwareT
   }
 
   @Test
-  def testCreateWithInvalidReplicationFactor() {
+  def testCreateWithInvalidReplicationFactor(): Unit = {
     val brokers = List(0)
     TestUtils.createBrokersInZk(zkClient, brokers)
 
@@ -259,7 +260,7 @@ class TopicCommandTest extends ZooKeeperTestHarness with Logging with RackAwareT
   }
 
   @Test
-  def testAlterIfExists() {
+  def testAlterIfExists(): Unit = {
     // create brokers
     val brokers = List(0, 1, 2)
     TestUtils.createBrokersInZk(zkClient, brokers)
@@ -276,7 +277,7 @@ class TopicCommandTest extends ZooKeeperTestHarness with Logging with RackAwareT
   }
 
   @Test
-  def testAlterConfigs() {
+  def testAlterConfigs(): Unit = {
     // create brokers
     val brokers = List(0, 1, 2)
     TestUtils.createBrokersInZk(zkClient, brokers)
@@ -289,18 +290,18 @@ class TopicCommandTest extends ZooKeeperTestHarness with Logging with RackAwareT
 
     val output = TestUtils.grabConsoleOutput(
       topicService.describeTopic(new TopicCommandOptions(Array("--topic", testTopicName))))
-    assertTrue("The output should contain the modified config", output.contains("Configs:cleanup.policy=compact"))
+    assertTrue("The output should contain the modified config", output.contains("Configs: cleanup.policy=compact"))
 
     topicService.alterTopic(new TopicCommandOptions(
       Array("--topic", testTopicName, "--config", "cleanup.policy=delete")))
 
     val output2 = TestUtils.grabConsoleOutput(
       topicService.describeTopic(new TopicCommandOptions(Array("--topic", testTopicName))))
-    assertTrue("The output should contain the modified config", output2.contains("Configs:cleanup.policy=delete"))
+    assertTrue("The output should contain the modified config", output2.contains("Configs: cleanup.policy=delete"))
   }
 
   @Test
-  def testConfigPreservationAcrossPartitionAlteration() {
+  def testConfigPreservationAcrossPartitionAlteration(): Unit = {
     val numPartitionsOriginal = 1
     val cleanupKey = "cleanup.policy"
     val cleanupVal = "compact"
@@ -330,7 +331,7 @@ class TopicCommandTest extends ZooKeeperTestHarness with Logging with RackAwareT
   }
 
   @Test
-  def testTopicDeletion() {
+  def testTopicDeletion(): Unit = {
 
     val numPartitionsOriginal = 1
 
@@ -368,7 +369,7 @@ class TopicCommandTest extends ZooKeeperTestHarness with Logging with RackAwareT
   }
 
   @Test
-  def testDeleteIfExists() {
+  def testDeleteIfExists(): Unit = {
     // create brokers
     val brokers = List(0, 1, 2)
     TestUtils.createBrokersInZk(zkClient, brokers)
@@ -404,7 +405,7 @@ class TopicCommandTest extends ZooKeeperTestHarness with Logging with RackAwareT
   }
 
   @Test
-  def testDescribeIfTopicNotExists() {
+  def testDescribeIfTopicNotExists(): Unit = {
     // create brokers
     val brokers = List(0, 1, 2)
     TestUtils.createBrokersInZk(zkClient, brokers)
@@ -415,6 +416,11 @@ class TopicCommandTest extends ZooKeeperTestHarness with Logging with RackAwareT
       topicService.describeTopic(describeOpts)
     }
 
+    // describe all topics
+    val describeOptsAllTopics = new TopicCommandOptions(Array())
+    // should not throw any error
+    topicService.describeTopic(describeOptsAllTopics)
+
     // describe topic that does not exist with --if-exists
     val describeOptsWithExists = new TopicCommandOptions(Array("--topic", testTopicName, "--if-exists"))
     // should not throw any error
@@ -422,7 +428,7 @@ class TopicCommandTest extends ZooKeeperTestHarness with Logging with RackAwareT
   }
 
   @Test
-  def testCreateAlterTopicWithRackAware() {
+  def testCreateAlterTopicWithRackAware(): Unit = {
     val rackInfo = Map(0 -> "rack1", 1 -> "rack2", 2 -> "rack2", 3 -> "rack1", 4 -> "rack3", 5 -> "rack3")
     TestUtils.createBrokersInZk(toBrokerMetadata(rackInfo), zkClient)
 
@@ -479,7 +485,7 @@ class TopicCommandTest extends ZooKeeperTestHarness with Logging with RackAwareT
   }
 
   @Test
-  def testDescribeAndListTopicsMarkedForDeletion() {
+  def testDescribeAndListTopicsMarkedForDeletion(): Unit = {
     val brokers = List(0)
     val markedForDeletionDescribe = "MarkedForDeletion"
     val markedForDeletionList = "marked for deletion"
@@ -493,20 +499,20 @@ class TopicCommandTest extends ZooKeeperTestHarness with Logging with RackAwareT
     topicService.deleteTopic(new TopicCommandOptions(Array("--topic", testTopicName)))
 
     // Test describe topics
-    def describeTopicsWithConfig() {
+    def describeTopicsWithConfig(): Unit = {
       topicService.describeTopic(new TopicCommandOptions(Array("--describe")))
     }
     val outputWithConfig = TestUtils.grabConsoleOutput(describeTopicsWithConfig())
     assertTrue(outputWithConfig.contains(testTopicName) && outputWithConfig.contains(markedForDeletionDescribe))
 
-    def describeTopicsNoConfig() {
+    def describeTopicsNoConfig(): Unit = {
       topicService.describeTopic(new TopicCommandOptions(Array("--describe", "--unavailable-partitions")))
     }
     val outputNoConfig = TestUtils.grabConsoleOutput(describeTopicsNoConfig())
     assertTrue(outputNoConfig.contains(testTopicName) && outputNoConfig.contains(markedForDeletionDescribe))
 
     // Test list topics
-    def listTopics() {
+    def listTopics(): Unit = {
       topicService.listTopics(new TopicCommandOptions(Array("--list")))
     }
     val output = TestUtils.grabConsoleOutput(listTopics())
@@ -514,7 +520,7 @@ class TopicCommandTest extends ZooKeeperTestHarness with Logging with RackAwareT
   }
 
   @Test
-  def testDescribeAndListTopicsWithoutInternalTopics() {
+  def testDescribeAndListTopicsWithoutInternalTopics(): Unit = {
     val brokers = List(0)
     TestUtils.createBrokersInZk(zkClient, brokers)
 
@@ -569,4 +575,26 @@ class TopicCommandTest extends ZooKeeperTestHarness with Logging with RackAwareT
     assertFalse(TestUtils.grabConsoleOutput(topicService.deleteTopic(escapedCommandOpts)).contains(topic2))
     assertTrue(TestUtils.grabConsoleOutput(topicService.deleteTopic(unescapedCommandOpts)).contains(topic2))
   }
+
+  @Test
+  def testAlterInternalTopicPartitionCount(): Unit = {
+    val brokers = List(0)
+    TestUtils.createBrokersInZk(zkClient, brokers)
+    
+    // create internal topics
+    adminZkClient.createTopic(Topic.GROUP_METADATA_TOPIC_NAME, 1, 1)
+    adminZkClient.createTopic(Topic.TRANSACTION_STATE_TOPIC_NAME, 1, 1)
+
+    def expectAlterInternalTopicPartitionCountFailed(topic: String): Unit = {
+      try {
+        topicService.alterTopic(new TopicCommandOptions(
+          Array("--topic", topic, "--partitions", "2")))
+        fail("Should have thrown an IllegalArgumentException")
+      } catch {
+        case _: IllegalArgumentException => // expected
+      }
+    }
+    expectAlterInternalTopicPartitionCountFailed(Topic.GROUP_METADATA_TOPIC_NAME)
+    expectAlterInternalTopicPartitionCountFailed(Topic.TRANSACTION_STATE_TOPIC_NAME)
+  }
 }
diff --git a/core/src/test/scala/unit/kafka/admin/TopicCommandWithAdminClientTest.scala b/core/src/test/scala/unit/kafka/admin/TopicCommandWithAdminClientTest.scala
index e0597de748751..13d68e5026308 100644
--- a/core/src/test/scala/unit/kafka/admin/TopicCommandWithAdminClientTest.scala
+++ b/core/src/test/scala/unit/kafka/admin/TopicCommandWithAdminClientTest.scala
@@ -25,14 +25,19 @@ import kafka.server.{ConfigType, KafkaConfig}
 import kafka.utils.{Exit, Logging, TestUtils}
 import kafka.zk.{ConfigEntityChangeNotificationZNode, DeleteTopicsTopicZNode}
 import org.apache.kafka.clients.CommonClientConfigs
-import org.apache.kafka.clients.admin.{ListTopicsOptions, NewTopic, AdminClient => JAdminClient}
+import org.apache.kafka.clients.admin.{Admin, AdminClient, ListTopicsOptions, NewTopic}
 import org.apache.kafka.common.config.{ConfigException, ConfigResource, TopicConfig}
 import org.apache.kafka.common.internals.Topic
+import org.apache.kafka.common.network.ListenerName
+import org.apache.kafka.common.protocol.Errors
+import org.apache.kafka.common.security.auth.SecurityProtocol
 import org.junit.Assert.{assertEquals, assertFalse, assertTrue}
-import org.junit.{After, Before, Rule, Test}
 import org.junit.rules.TestName
+import org.junit.{After, Before, Rule, Test}
+import org.scalatest.Assertions.{fail, intercept}
 
 import scala.collection.JavaConverters._
+import scala.collection.Seq
 import scala.concurrent.ExecutionException
 import scala.util.Random
 
@@ -45,17 +50,22 @@ class TopicCommandWithAdminClientTest extends KafkaServerTestHarness with Loggin
   override def generateConfigs: Seq[KafkaConfig] = TestUtils.createBrokerConfigs(
     numConfigs = 6,
     zkConnect = zkConnect,
-    rackInfo = Map(0 -> "rack1", 1 -> "rack2", 2 -> "rack2", 3 -> "rack1", 4 -> "rack3", 5 -> "rack3"
-    )).map(KafkaConfig.fromProps)
+    rackInfo = Map(0 -> "rack1", 1 -> "rack2", 2 -> "rack2", 3 -> "rack1", 4 -> "rack3", 5 -> "rack3"),
+    numPartitions = numPartitions,
+    defaultReplicationFactor = defaultReplicationFactor
+    ).map(KafkaConfig.fromProps)
+
+  private val numPartitions = 1
+  private val defaultReplicationFactor = 1.toShort
 
   private var topicService: AdminClientTopicService = _
-  private var adminClient: JAdminClient = _
+  private var adminClient: Admin = _
   private var testTopicName: String = _
 
   private val _testName = new TestName
   @Rule def testName = _testName
 
-  def assertExitCode(expected: Int, method: () => Unit) {
+  def assertExitCode(expected: Int, method: () => Unit): Unit = {
     def mockExitProcedure(exitCode: Int, exitMessage: Option[String]): Nothing = {
       assertEquals(expected, exitCode)
       throw new RuntimeException
@@ -70,7 +80,7 @@ class TopicCommandWithAdminClientTest extends KafkaServerTestHarness with Loggin
     }
   }
 
-  def assertCheckArgsExitCode(expected: Int, options: TopicCommandOptions) {
+  def assertCheckArgsExitCode(expected: Int, options: TopicCommandOptions): Unit = {
     assertExitCode(expected, options.checkArgs _)
   }
 
@@ -91,11 +101,11 @@ class TopicCommandWithAdminClientTest extends KafkaServerTestHarness with Loggin
   }
 
   @Before
-  def setup() {
+  def setup(): Unit = {
     // create adminClient
     val props = new Properties()
     props.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, brokerList)
-    adminClient = JAdminClient.create(props)
+    adminClient = AdminClient.create(props)
     topicService = AdminClientTopicService(adminClient)
     testTopicName = s"${testName.getMethodName}-${Random.alphanumeric.take(10).mkString}"
   }
@@ -151,6 +161,52 @@ class TopicCommandWithAdminClientTest extends KafkaServerTestHarness with Loggin
     adminClient.listTopics().names().get().contains(testTopicName)
   }
 
+  @Test
+  def testCreateWithDefaults(): Unit = {
+    createAndWaitTopic(new TopicCommandOptions(
+      Array("--topic", testTopicName)))
+
+    val partitions = adminClient
+      .describeTopics(Collections.singletonList(testTopicName))
+      .all()
+      .get()
+      .get(testTopicName)
+      .partitions()
+    assertEquals(partitions.size(), numPartitions)
+    assertEquals(partitions.get(0).replicas().size(), defaultReplicationFactor)
+  }
+
+  @Test
+  def testCreateWithDefaultReplication(): Unit = {
+    createAndWaitTopic(new TopicCommandOptions(
+      Array("--topic", testTopicName, "--partitions", "2")))
+
+    val partitions = adminClient
+      .describeTopics(Collections.singletonList(testTopicName))
+      .all()
+      .get()
+      .get(testTopicName)
+      .partitions()
+    assertEquals(partitions.size(), 2)
+    assertEquals(partitions.get(0).replicas().size(), defaultReplicationFactor)
+  }
+
+  @Test
+  def testCreateWithDefaultPartitions(): Unit = {
+    createAndWaitTopic(new TopicCommandOptions(
+      Array("--topic", testTopicName, "--replication-factor", "2")))
+
+    val partitions = adminClient
+      .describeTopics(Collections.singletonList(testTopicName))
+      .all()
+      .get()
+      .get(testTopicName)
+      .partitions()
+
+    assertEquals(partitions.size(), numPartitions)
+    assertEquals(partitions.get(0).replicas().size(), 2)
+  }
+
   @Test
   def testCreateWithConfigs(): Unit = {
     val configResource = new ConfigResource(ConfigResource.Type.TOPIC, testTopicName)
@@ -198,7 +254,7 @@ class TopicCommandWithAdminClientTest extends KafkaServerTestHarness with Loggin
   }
 
   @Test
-  def testCreateWithInvalidReplicationFactor() {
+  def testCreateWithInvalidReplicationFactor(): Unit = {
     intercept[IllegalArgumentException] {
       topicService.createTopic(new TopicCommandOptions(
         Array("--partitions", "2", "--replication-factor", (Short.MaxValue+1).toString, "--topic", testTopicName)))
@@ -207,7 +263,7 @@ class TopicCommandWithAdminClientTest extends KafkaServerTestHarness with Loggin
 
   @Test
   def testCreateWithNegativeReplicationFactor(): Unit = {
-    intercept[ExecutionException] {
+    intercept[IllegalArgumentException] {
       topicService.createTopic(new TopicCommandOptions(
         Array("--partitions", "2", "--replication-factor", "-1", "--topic", testTopicName)))
     }
@@ -237,17 +293,17 @@ class TopicCommandWithAdminClientTest extends KafkaServerTestHarness with Loggin
 
   @Test
   def testCreateWithNegativePartitionCount(): Unit = {
-    intercept[ExecutionException] {
+    intercept[IllegalArgumentException] {
       topicService.createTopic(new TopicCommandOptions(
         Array("--partitions", "-1", "--replication-factor", "1", "--topic", testTopicName)))
     }
   }
 
   @Test
-  def testCreateWithUnspecifiedPartitionCount(): Unit = {
-    assertExitCode(1,
-      () => topicService.createTopic(new TopicCommandOptions(
-        Array("--replication-factor", "1", "--topic", testTopicName))))
+  def testCreateWithUnspecifiedReplicationFactorAndPartitionsWithZkClient(): Unit = {
+    assertExitCode(1, () =>
+      new TopicCommandOptions(Array("--create", "--zookeeper", "zk", "--topic", testTopicName)).checkArgs()
+    )
   }
 
   @Test
@@ -277,9 +333,9 @@ class TopicCommandWithAdminClientTest extends KafkaServerTestHarness with Loggin
     val topic2 = "kafka.testTopic2"
     val topic3 = "oooof.testTopic1"
     adminClient.createTopics(
-      List(new NewTopic(topic1, 2, 2),
-        new NewTopic(topic2, 2, 2),
-        new NewTopic(topic3, 2, 2)).asJavaCollection)
+      List(new NewTopic(topic1, 2, 2.toShort),
+        new NewTopic(topic2, 2, 2.toShort),
+        new NewTopic(topic3, 2, 2.toShort)).asJavaCollection)
       .all().get()
     waitForTopicCreated(topic1)
     waitForTopicCreated(topic2)
@@ -297,8 +353,8 @@ class TopicCommandWithAdminClientTest extends KafkaServerTestHarness with Loggin
   def testListTopicsWithExcludeInternal(): Unit = {
     val topic1 = "kafka.testTopic1"
     adminClient.createTopics(
-      List(new NewTopic(topic1, 2, 2),
-        new NewTopic(Topic.GROUP_METADATA_TOPIC_NAME, 2, 2)).asJavaCollection)
+      List(new NewTopic(topic1, 2, 2.toShort),
+        new NewTopic(Topic.GROUP_METADATA_TOPIC_NAME, 2, 2.toShort)).asJavaCollection)
       .all().get()
     waitForTopicCreated(topic1)
 
@@ -312,7 +368,7 @@ class TopicCommandWithAdminClientTest extends KafkaServerTestHarness with Loggin
   @Test
   def testAlterPartitionCount(): Unit = {
     adminClient.createTopics(
-      List(new NewTopic(testTopicName, 2, 2)).asJavaCollection).all().get()
+      List(new NewTopic(testTopicName, 2, 2.toShort)).asJavaCollection).all().get()
     waitForTopicCreated(testTopicName)
 
     topicService.alterTopic(new TopicCommandOptions(
@@ -325,7 +381,7 @@ class TopicCommandWithAdminClientTest extends KafkaServerTestHarness with Loggin
   @Test
   def testAlterAssignment(): Unit = {
     adminClient.createTopics(
-      Collections.singletonList(new NewTopic(testTopicName, 2, 2))).all().get()
+      Collections.singletonList(new NewTopic(testTopicName, 2, 2.toShort))).all().get()
     waitForTopicCreated(testTopicName)
 
     topicService.alterTopic(new TopicCommandOptions(
@@ -339,7 +395,7 @@ class TopicCommandWithAdminClientTest extends KafkaServerTestHarness with Loggin
   @Test
   def testAlterAssignmentWithMoreAssignmentThanPartitions(): Unit = {
     adminClient.createTopics(
-      List(new NewTopic(testTopicName, 2, 2)).asJavaCollection).all().get()
+      List(new NewTopic(testTopicName, 2, 2.toShort)).asJavaCollection).all().get()
     waitForTopicCreated(testTopicName)
 
     intercept[ExecutionException] {
@@ -351,7 +407,7 @@ class TopicCommandWithAdminClientTest extends KafkaServerTestHarness with Loggin
   @Test
   def testAlterAssignmentWithMorePartitionsThanAssignment(): Unit = {
     adminClient.createTopics(
-      List(new NewTopic(testTopicName, 2, 2)).asJavaCollection).all().get()
+      List(new NewTopic(testTopicName, 2, 2.toShort)).asJavaCollection).all().get()
     waitForTopicCreated(testTopicName)
 
     intercept[ExecutionException] {
@@ -465,7 +521,7 @@ class TopicCommandWithAdminClientTest extends KafkaServerTestHarness with Loggin
     val deletePath = DeleteTopicsTopicZNode.path(testTopicName)
     assertFalse("Delete path for topic shouldn't exist before deletion.", zkClient.pathExists(deletePath))
     topicService.deleteTopic(deleteOpts)
-    assertTrue("Delete path for topic should exist after deletion.", zkClient.pathExists(deletePath))
+    TestUtils.verifyTopicDeletion(zkClient, testTopicName, 1, servers)
   }
 
   @Test
@@ -483,7 +539,7 @@ class TopicCommandWithAdminClientTest extends KafkaServerTestHarness with Loggin
     val deleteOffsetTopicPath = DeleteTopicsTopicZNode.path(Topic.GROUP_METADATA_TOPIC_NAME)
     assertFalse("Delete path for topic shouldn't exist before deletion.", zkClient.pathExists(deleteOffsetTopicPath))
     topicService.deleteTopic(deleteOffsetTopicOpts)
-    assertTrue("Delete path for topic should exist after deletion.", zkClient.pathExists(deleteOffsetTopicPath))
+    TestUtils.verifyTopicDeletion(zkClient, Topic.GROUP_METADATA_TOPIC_NAME, 1, servers)
   }
 
   @Test
@@ -498,7 +554,7 @@ class TopicCommandWithAdminClientTest extends KafkaServerTestHarness with Loggin
   @Test
   def testDescribe(): Unit = {
     adminClient.createTopics(
-      Collections.singletonList(new NewTopic(testTopicName, 2, 2))).all().get()
+      Collections.singletonList(new NewTopic(testTopicName, 2, 2.toShort))).all().get()
     waitForTopicCreated(testTopicName)
 
     val output = TestUtils.grabConsoleOutput(
@@ -511,16 +567,42 @@ class TopicCommandWithAdminClientTest extends KafkaServerTestHarness with Loggin
   @Test
   def testDescribeUnavailablePartitions(): Unit = {
     adminClient.createTopics(
-      Collections.singletonList(new NewTopic(testTopicName, 6, 1))).all().get()
+      Collections.singletonList(new NewTopic(testTopicName, 6, 1.toShort))).all().get()
     waitForTopicCreated(testTopicName)
 
     try {
+      // check which partition is on broker 0 which we'll kill
+      val testTopicDescription = adminClient.describeTopics(Collections.singletonList(testTopicName))
+        .all().get().asScala(testTopicName)
+      val partitionOnBroker0 = testTopicDescription.partitions().asScala.find(_.leader().id() == 0).get.partition()
+
       killBroker(0)
+
+      // wait until the topic metadata for the test topic is propagated to each alive broker
+      TestUtils.waitUntilTrue(() => {
+        servers
+          .filterNot(_.config.brokerId == 0)
+          .foldLeft(true) {
+            (result, server) => {
+              val topicMetadatas = server.dataPlaneRequestProcessor.metadataCache
+                .getTopicMetadata(Set(testTopicName), ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT))
+              val testPartitionMetadata = topicMetadatas.find(_.topic().equals(testTopicName)).get.partitionMetadata().asScala.find(_.partition() == partitionOnBroker0)
+              testPartitionMetadata match {
+                case None => fail(s"Partition metadata is not found in metadata cache")
+                case Some(metadata) => {
+                  result && metadata.error() == Errors.LEADER_NOT_AVAILABLE
+                }
+              }
+            }
+          }
+      }, s"Partition metadata for $testTopicName is not propagated")
+
+      // grab the console output and assert
       val output = TestUtils.grabConsoleOutput(
-        topicService.describeTopic(new TopicCommandOptions(Array("--topic", testTopicName, "--unavailable-partitions"))))
+          topicService.describeTopic(new TopicCommandOptions(Array("--topic", testTopicName, "--unavailable-partitions"))))
       val rows = output.split("\n")
       assertTrue(rows(0).startsWith(s"\tTopic: $testTopicName"))
-      assertTrue(rows(0).endsWith("Leader: none\tReplicas: 0\tIsr: "))
+      assertTrue(rows(0).contains("Leader: none\tReplicas: 0\tIsr:"))
     } finally {
       restartDeadBrokers()
     }
@@ -529,11 +611,13 @@ class TopicCommandWithAdminClientTest extends KafkaServerTestHarness with Loggin
   @Test
   def testDescribeUnderReplicatedPartitions(): Unit = {
     adminClient.createTopics(
-      Collections.singletonList(new NewTopic(testTopicName, 1, 6))).all().get()
+      Collections.singletonList(new NewTopic(testTopicName, 1, 6.toShort))).all().get()
     waitForTopicCreated(testTopicName)
 
     try {
       killBroker(0)
+      val aliveServers = servers.filterNot(_.config.brokerId == 0)
+      TestUtils.waitUntilMetadataIsPropagated(aliveServers, testTopicName, 0)
       val output = TestUtils.grabConsoleOutput(
         topicService.describeTopic(new TopicCommandOptions(Array("--under-replicated-partitions"))))
       val rows = output.split("\n")
@@ -549,11 +633,13 @@ class TopicCommandWithAdminClientTest extends KafkaServerTestHarness with Loggin
     configMap.put(TopicConfig.MIN_IN_SYNC_REPLICAS_CONFIG, "6")
 
     adminClient.createTopics(
-      Collections.singletonList(new NewTopic(testTopicName, 1, 6).configs(configMap))).all().get()
+      Collections.singletonList(new NewTopic(testTopicName, 1, 6.toShort).configs(configMap))).all().get()
     waitForTopicCreated(testTopicName)
 
     try {
       killBroker(0)
+      val aliveServers = servers.filterNot(_.config.brokerId == 0)
+      TestUtils.waitUntilMetadataIsPropagated(aliveServers, testTopicName, 0)
       val output = TestUtils.grabConsoleOutput(
         topicService.describeTopic(new TopicCommandOptions(Array("--under-min-isr-partitions"))))
       val rows = output.split("\n")
@@ -563,6 +649,28 @@ class TopicCommandWithAdminClientTest extends KafkaServerTestHarness with Loggin
     }
   }
 
+  @Test
+  def testDescribeAtMinIsrPartitions(): Unit = {
+    val configMap = new java.util.HashMap[String, String]()
+    configMap.put(TopicConfig.MIN_IN_SYNC_REPLICAS_CONFIG, "4")
+
+    adminClient.createTopics(
+      Collections.singletonList(new NewTopic(testTopicName, 1, 6.toShort).configs(configMap))).all().get()
+    waitForTopicCreated(testTopicName)
+
+    try {
+      killBroker(0)
+      killBroker(1)
+      val output = TestUtils.grabConsoleOutput(
+        topicService.describeTopic(new TopicCommandOptions(Array("--at-min-isr-partitions"))))
+      val rows = output.split("\n")
+      assertTrue(rows(0).startsWith(s"\tTopic: $testTopicName"))
+      assertEquals(1, rows.length);
+    } finally {
+      restartDeadBrokers()
+    }
+  }
+
   /**
     * Test describe --under-min-isr-partitions option with four topics:
     *   (1) topic with partition under the configured min ISR count
@@ -584,8 +692,8 @@ class TopicCommandWithAdminClientTest extends KafkaServerTestHarness with Loggin
 
     adminClient.createTopics(
       java.util.Arrays.asList(
-        new NewTopic(underMinIsrTopic, 1, 6).configs(configMap),
-        new NewTopic(notUnderMinIsrTopic, 1, 6),
+        new NewTopic(underMinIsrTopic, 1, 6.toShort).configs(configMap),
+        new NewTopic(notUnderMinIsrTopic, 1, 6.toShort),
         new NewTopic(offlineTopic, Collections.singletonMap(0, Collections.singletonList(0))),
         new NewTopic(fullyReplicatedTopic, Collections.singletonMap(0, java.util.Arrays.asList(1, 2, 3))))).all().get()
 
@@ -596,6 +704,8 @@ class TopicCommandWithAdminClientTest extends KafkaServerTestHarness with Loggin
 
     try {
       killBroker(0)
+      val aliveServers = servers.filterNot(_.config.brokerId == 0)
+      TestUtils.waitUntilMetadataIsPropagated(aliveServers, underMinIsrTopic, 0)
       val output = TestUtils.grabConsoleOutput(
         topicService.describeTopic(new TopicCommandOptions(Array("--under-min-isr-partitions"))))
       val rows = output.split("\n")
diff --git a/core/src/test/scala/unit/kafka/api/ApiUtilsTest.scala b/core/src/test/scala/unit/kafka/api/ApiUtilsTest.scala
index 9fe4cbf1965b4..08994c87ad01a 100644
--- a/core/src/test/scala/unit/kafka/api/ApiUtilsTest.scala
+++ b/core/src/test/scala/unit/kafka/api/ApiUtilsTest.scala
@@ -18,7 +18,6 @@
 package kafka.api
 
 import org.junit._
-import org.scalatest.junit.JUnitSuite
 import org.junit.Assert._
 
 import scala.util.Random
@@ -31,10 +30,10 @@ object ApiUtilsTest {
   val rnd: Random = new Random()
 }
 
-class ApiUtilsTest extends JUnitSuite {
+class ApiUtilsTest {
 
   @Test
-  def testShortStringNonASCII() {
+  def testShortStringNonASCII(): Unit = {
     // Random-length strings
     for(_ <- 0 to 100) {
       // Since we're using UTF-8 encoding, each encoded byte will be one to four bytes long 
@@ -47,7 +46,7 @@ class ApiUtilsTest extends JUnitSuite {
   }
 
   @Test
-  def testShortStringASCII() {
+  def testShortStringASCII(): Unit = {
     // Random-length strings
     for(_ <- 0 to 100) {
       val s: String = TestUtils.randomString(math.abs(ApiUtilsTest.rnd.nextInt()) % Short.MaxValue)  
diff --git a/core/src/test/scala/unit/kafka/api/ApiVersionTest.scala b/core/src/test/scala/unit/kafka/api/ApiVersionTest.scala
index 571a0775bd3a4..caff74ce2d666 100644
--- a/core/src/test/scala/unit/kafka/api/ApiVersionTest.scala
+++ b/core/src/test/scala/unit/kafka/api/ApiVersionTest.scala
@@ -88,6 +88,15 @@ class ApiVersionTest {
     assertEquals(KAFKA_2_2_IV1, ApiVersion("2.2"))
     assertEquals(KAFKA_2_2_IV0, ApiVersion("2.2-IV0"))
     assertEquals(KAFKA_2_2_IV1, ApiVersion("2.2-IV1"))
+
+    assertEquals(KAFKA_2_3_IV2, ApiVersion("2.3"))
+    assertEquals(KAFKA_2_3_IV0, ApiVersion("2.3-IV0"))
+    assertEquals(KAFKA_2_3_IV1, ApiVersion("2.3-IV1"))
+    assertEquals(KAFKA_2_3_IV2, ApiVersion("2.3-IV2"))
+
+    assertEquals(KAFKA_2_4_IV1, ApiVersion("2.4"))
+    assertEquals(KAFKA_2_4_IV0, ApiVersion("2.4-IV0"))
+    assertEquals(KAFKA_2_4_IV1, ApiVersion("2.4-IV1"))
   }
 
   @Test
@@ -116,7 +125,22 @@ class ApiVersionTest {
   def testShortVersion(): Unit = {
     assertEquals("0.8.0", KAFKA_0_8_0.shortVersion)
     assertEquals("0.10.0", KAFKA_0_10_0_IV0.shortVersion)
+    assertEquals("0.10.0", KAFKA_0_10_0_IV1.shortVersion)
     assertEquals("0.11.0", KAFKA_0_11_0_IV0.shortVersion)
+    assertEquals("0.11.0", KAFKA_0_11_0_IV1.shortVersion)
+    assertEquals("0.11.0", KAFKA_0_11_0_IV2.shortVersion)
+    assertEquals("1.0", KAFKA_1_0_IV0.shortVersion)
+    assertEquals("1.1", KAFKA_1_1_IV0.shortVersion)
+    assertEquals("2.0", KAFKA_2_0_IV0.shortVersion)
+    assertEquals("2.0", KAFKA_2_0_IV1.shortVersion)
+    assertEquals("2.1", KAFKA_2_1_IV0.shortVersion)
+    assertEquals("2.1", KAFKA_2_1_IV1.shortVersion)
+    assertEquals("2.1", KAFKA_2_1_IV2.shortVersion)
+    assertEquals("2.2", KAFKA_2_2_IV0.shortVersion)
+    assertEquals("2.2", KAFKA_2_2_IV1.shortVersion)
+    assertEquals("2.3", KAFKA_2_3_IV0.shortVersion)
+    assertEquals("2.3", KAFKA_2_3_IV1.shortVersion)
+    assertEquals("2.4", KAFKA_2_4_IV0.shortVersion)
   }
 
   @Test
diff --git a/core/src/test/scala/unit/kafka/cluster/PartitionTest.scala b/core/src/test/scala/unit/kafka/cluster/PartitionTest.scala
index ba53cd02e6365..800f672f659f4 100644
--- a/core/src/test/scala/unit/kafka/cluster/PartitionTest.scala
+++ b/core/src/test/scala/unit/kafka/cluster/PartitionTest.scala
@@ -22,46 +22,53 @@ import java.util.{Optional, Properties}
 import java.util.concurrent.{CountDownLatch, Executors, TimeUnit, TimeoutException}
 import java.util.concurrent.atomic.AtomicBoolean
 
-import kafka.api.{ApiVersion, Request}
+import com.yammer.metrics.Metrics
+import com.yammer.metrics.core.Metric
+import kafka.api.{ApiVersion, LeaderAndIsr}
 import kafka.common.UnexpectedAppendOffsetException
 import kafka.log.{Defaults => _, _}
-import kafka.server.QuotaFactory.QuotaManagers
 import kafka.server._
+import kafka.server.checkpoints.OffsetCheckpoints
 import kafka.utils._
-import kafka.zk.KafkaZkClient
 import org.apache.kafka.common.TopicPartition
 import org.apache.kafka.common.errors.{ApiException, OffsetNotAvailableException, ReplicaNotAvailableException}
-import org.apache.kafka.common.metrics.Metrics
+import org.apache.kafka.common.message.LeaderAndIsrRequestData.LeaderAndIsrPartitionState
 import org.apache.kafka.common.protocol.Errors
 import org.apache.kafka.common.record.FileRecords.TimestampAndOffset
 import org.apache.kafka.common.utils.Utils
 import org.apache.kafka.common.record._
-import org.apache.kafka.common.requests.{EpochEndOffset, IsolationLevel, LeaderAndIsrRequest, ListOffsetRequest}
+import org.apache.kafka.common.requests.{EpochEndOffset, IsolationLevel, ListOffsetRequest}
 import org.junit.{After, Before, Test}
 import org.junit.Assert._
+import org.mockito.Mockito.{doAnswer, doNothing, mock, spy, times, verify, when}
 import org.scalatest.Assertions.assertThrows
-import org.easymock.{Capture, EasyMock, IAnswer}
+import org.mockito.ArgumentMatchers
+import org.mockito.invocation.InvocationOnMock
+import org.mockito.stubbing.Answer
 
 import scala.collection.JavaConverters._
+import scala.collection.mutable.ListBuffer
 
 class PartitionTest {
 
   val brokerId = 101
   val topicPartition = new TopicPartition("test-topic", 0)
   val time = new MockTime()
-  val brokerTopicStats = new BrokerTopicStats
-  val metrics = new Metrics
-
   var tmpDir: File = _
   var logDir1: File = _
   var logDir2: File = _
-  var replicaManager: ReplicaManager = _
   var logManager: LogManager = _
   var logConfig: LogConfig = _
-  var quotaManagers: QuotaManagers = _
+  val stateStore: PartitionStateStore = mock(classOf[PartitionStateStore])
+  val delayedOperations: DelayedOperations = mock(classOf[DelayedOperations])
+  val metadataCache: MetadataCache = mock(classOf[MetadataCache])
+  val offsetCheckpoints: OffsetCheckpoints = mock(classOf[OffsetCheckpoints])
+  var partition: Partition = _
 
   @Before
   def setup(): Unit = {
+    TestUtils.clearYammerMetrics()
+
     val logProps = createLogProperties(Map.empty)
     logConfig = LogConfig(logProps)
 
@@ -72,20 +79,19 @@ class PartitionTest {
       logDirs = Seq(logDir1, logDir2), defaultConfig = logConfig, CleanerConfig(enableCleaner = false), time)
     logManager.startup()
 
-    val brokerProps = TestUtils.createBrokerConfig(brokerId, TestUtils.MockZkConnect)
-    brokerProps.put(KafkaConfig.LogDirsProp, Seq(logDir1, logDir2).map(_.getAbsolutePath).mkString(","))
-    val brokerConfig = KafkaConfig.fromProps(brokerProps)
-    val kafkaZkClient: KafkaZkClient = EasyMock.createMock(classOf[KafkaZkClient])
-    quotaManagers = QuotaFactory.instantiate(brokerConfig, metrics, time, "")
-    replicaManager = new ReplicaManager(
-      config = brokerConfig, metrics, time, zkClient = kafkaZkClient, new MockScheduler(time),
-      logManager, new AtomicBoolean(false), quotaManagers,
-      brokerTopicStats, new MetadataCache(brokerId), new LogDirFailureChannel(brokerConfig.logDirs.size))
-
-    EasyMock.expect(kafkaZkClient.getEntityConfigs(EasyMock.anyString(), EasyMock.anyString())).andReturn(logProps).anyTimes()
-    EasyMock.expect(kafkaZkClient.conditionalUpdatePath(EasyMock.anyObject(), EasyMock.anyObject(), EasyMock.anyObject(), EasyMock.anyObject()))
-      .andReturn((true, 0)).anyTimes()
-    EasyMock.replay(kafkaZkClient)
+    partition = new Partition(topicPartition,
+      replicaLagTimeMaxMs = Defaults.ReplicaLagTimeMaxMs,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      localBrokerId = brokerId,
+      time,
+      stateStore,
+      delayedOperations,
+      metadataCache,
+      logManager)
+
+    when(stateStore.fetchTopicConfig()).thenReturn(createLogProperties(Map.empty))
+    when(offsetCheckpoints.fetch(ArgumentMatchers.anyString, ArgumentMatchers.eq(topicPartition)))
+      .thenReturn(None)
   }
 
   private def createLogProperties(overrides: Map[String, String]): Properties = {
@@ -99,21 +105,16 @@ class PartitionTest {
 
   @After
   def tearDown(): Unit = {
-    brokerTopicStats.close()
-    metrics.close()
-
     logManager.shutdown()
     Utils.delete(tmpDir)
-    logManager.liveLogDirs.foreach(Utils.delete)
-    replicaManager.shutdown(checkpointHW = false)
-    quotaManagers.shutdown()
+    TestUtils.clearYammerMetrics()
   }
 
   @Test
   def testMakeLeaderUpdatesEpochCache(): Unit = {
     val leaderEpoch = 8
 
-    val log = logManager.getOrCreateLog(topicPartition, logConfig)
+    val log = logManager.getOrCreateLog(topicPartition, () => logConfig)
     log.appendAsLeader(MemoryRecords.withRecords(0L, CompressionType.NONE, 0,
       new SimpleRecord("k1".getBytes, "v1".getBytes),
       new SimpleRecord("k2".getBytes, "v2".getBytes)
@@ -125,7 +126,7 @@ class PartitionTest {
     assertEquals(4, log.logEndOffset)
 
     val partition = setupPartitionWithMocks(leaderEpoch = leaderEpoch, isLeader = true, log = log)
-    assertEquals(Some(4), partition.leaderReplicaIfLocal.map(_.logEndOffset))
+    assertEquals(Some(4), partition.leaderLogIfLocal.map(_.logEndOffset))
 
     val epochEndOffset = partition.lastOffsetForLeaderEpoch(currentLeaderEpoch = Optional.of[Integer](leaderEpoch),
       leaderEpoch = leaderEpoch, fetchOnlyFromLeader = true)
@@ -139,7 +140,7 @@ class PartitionTest {
 
     val logConfig = LogConfig(createLogProperties(Map(
       LogConfig.MessageFormatVersionProp -> kafka.api.KAFKA_0_10_2_IV0.shortVersion)))
-    val log = logManager.getOrCreateLog(topicPartition, logConfig)
+    val log = logManager.getOrCreateLog(topicPartition, () => logConfig)
     log.appendAsLeader(TestUtils.records(List(
       new SimpleRecord("k1".getBytes, "v1".getBytes),
       new SimpleRecord("k2".getBytes, "v2".getBytes)),
@@ -153,7 +154,7 @@ class PartitionTest {
     assertEquals(4, log.logEndOffset)
 
     val partition = setupPartitionWithMocks(leaderEpoch = leaderEpoch, isLeader = true, log = log)
-    assertEquals(Some(4), partition.leaderReplicaIfLocal.map(_.logEndOffset))
+    assertEquals(Some(4), partition.leaderLogIfLocal.map(_.logEndOffset))
     assertEquals(None, log.latestEpoch)
 
     val epochEndOffset = partition.lastOffsetForLeaderEpoch(currentLeaderEpoch = Optional.of[Integer](leaderEpoch),
@@ -168,17 +169,9 @@ class PartitionTest {
     val latch = new CountDownLatch(1)
 
     logManager.maybeUpdatePreferredLogDir(topicPartition, logDir1.getAbsolutePath)
-    val log1 = logManager.getOrCreateLog(topicPartition, logConfig)
+    partition.createLogIfNotExists(isNew = true, isFutureReplica = false, offsetCheckpoints)
     logManager.maybeUpdatePreferredLogDir(topicPartition, logDir2.getAbsolutePath)
-    val log2 = logManager.getOrCreateLog(topicPartition, logConfig, isFuture = true)
-    val currentReplica = new Replica(brokerId, topicPartition, time, log = Some(log1))
-    val futureReplica = new Replica(Request.FutureLocalReplicaId, topicPartition, time, log = Some(log2))
-    val partition = Partition(topicPartition, time, replicaManager)
-
-    partition.addReplicaIfNotExists(futureReplica)
-    partition.addReplicaIfNotExists(currentReplica)
-    assertEquals(Some(currentReplica), partition.localReplica)
-    assertEquals(Some(futureReplica), partition.futureLocalReplica)
+    partition.maybeCreateFutureReplica(logDir2.getAbsolutePath, offsetCheckpoints)
 
     val thread1 = new Thread {
       override def run(): Unit = {
@@ -200,17 +193,21 @@ class PartitionTest {
     latch.countDown()
     thread1.join()
     thread2.join()
-    assertEquals(None, partition.futureLocalReplica)
+    assertEquals(None, partition.futureLog)
   }
 
   @Test
   // Verify that replacement works when the replicas have the same log end offset but different base offsets in the
   // active segment
   def testMaybeReplaceCurrentWithFutureReplicaDifferentBaseOffsets(): Unit = {
-    // Write records with duplicate keys to current replica and roll at offset 6
     logManager.maybeUpdatePreferredLogDir(topicPartition, logDir1.getAbsolutePath)
-    val log1 = logManager.getOrCreateLog(topicPartition, logConfig)
-    log1.appendAsLeader(MemoryRecords.withRecords(0L, CompressionType.NONE, 0,
+    partition.createLogIfNotExists(isNew = true, isFutureReplica = false, offsetCheckpoints)
+    logManager.maybeUpdatePreferredLogDir(topicPartition, logDir2.getAbsolutePath)
+    partition.maybeCreateFutureReplica(logDir2.getAbsolutePath, offsetCheckpoints)
+
+    // Write records with duplicate keys to current replica and roll at offset 6
+    val currentLog = partition.log.get
+    currentLog.appendAsLeader(MemoryRecords.withRecords(0L, CompressionType.NONE, 0,
       new SimpleRecord("k1".getBytes, "v1".getBytes),
       new SimpleRecord("k1".getBytes, "v2".getBytes),
       new SimpleRecord("k1".getBytes, "v3".getBytes),
@@ -218,33 +215,24 @@ class PartitionTest {
       new SimpleRecord("k2".getBytes, "v5".getBytes),
       new SimpleRecord("k2".getBytes, "v6".getBytes)
     ), leaderEpoch = 0)
-    log1.roll()
-    log1.appendAsLeader(MemoryRecords.withRecords(0L, CompressionType.NONE, 0,
+    currentLog.roll()
+    currentLog.appendAsLeader(MemoryRecords.withRecords(0L, CompressionType.NONE, 0,
       new SimpleRecord("k3".getBytes, "v7".getBytes),
       new SimpleRecord("k4".getBytes, "v8".getBytes)
     ), leaderEpoch = 0)
 
     // Write to the future replica as if the log had been compacted, and do not roll the segment
-    logManager.maybeUpdatePreferredLogDir(topicPartition, logDir2.getAbsolutePath)
-    val log2 = logManager.getOrCreateLog(topicPartition, logConfig, isFuture = true)
+
     val buffer = ByteBuffer.allocate(1024)
-    var builder = MemoryRecords.builder(buffer, RecordBatch.CURRENT_MAGIC_VALUE, CompressionType.NONE,
+    val builder = MemoryRecords.builder(buffer, RecordBatch.CURRENT_MAGIC_VALUE, CompressionType.NONE,
       TimestampType.CREATE_TIME, 0L, RecordBatch.NO_TIMESTAMP, 0)
     builder.appendWithOffset(2L, new SimpleRecord("k1".getBytes, "v3".getBytes))
     builder.appendWithOffset(5L, new SimpleRecord("k2".getBytes, "v6".getBytes))
     builder.appendWithOffset(6L, new SimpleRecord("k3".getBytes, "v7".getBytes))
     builder.appendWithOffset(7L, new SimpleRecord("k4".getBytes, "v8".getBytes))
 
-    log2.appendAsFollower(builder.build())
-
-    val currentReplica = new Replica(brokerId, topicPartition, time, log = Some(log1))
-    val futureReplica = new Replica(Request.FutureLocalReplicaId, topicPartition, time, log = Some(log2))
-    val partition = Partition(topicPartition, time, replicaManager)
-
-    partition.addReplicaIfNotExists(futureReplica)
-    partition.addReplicaIfNotExists(currentReplica)
-    assertEquals(Some(currentReplica), partition.localReplica)
-    assertEquals(Some(futureReplica), partition.futureLocalReplica)
+    val futureLog = partition.futureLocalLogOrException
+    futureLog.appendAsFollower(builder.build())
 
     assertTrue(partition.maybeReplaceCurrentWithFutureReplica())
   }
@@ -255,9 +243,11 @@ class PartitionTest {
     val partition = setupPartitionWithMocks(leaderEpoch, isLeader = true)
 
     def assertSnapshotError(expectedError: Errors, currentLeaderEpoch: Optional[Integer]): Unit = {
-      partition.fetchOffsetSnapshotOrError(currentLeaderEpoch, fetchOnlyFromLeader = true) match {
-        case Left(_) => assertEquals(Errors.NONE, expectedError)
-        case Right(error) => assertEquals(expectedError, error)
+      try {
+        partition.fetchOffsetSnapshot(currentLeaderEpoch, fetchOnlyFromLeader = true)
+        assertEquals(Errors.NONE, expectedError)
+      } catch {
+        case error: ApiException => assertEquals(expectedError, Errors.forException(error))
       }
     }
 
@@ -275,9 +265,11 @@ class PartitionTest {
     def assertSnapshotError(expectedError: Errors,
                             currentLeaderEpoch: Optional[Integer],
                             fetchOnlyLeader: Boolean): Unit = {
-      partition.fetchOffsetSnapshotOrError(currentLeaderEpoch, fetchOnlyFromLeader = fetchOnlyLeader) match {
-        case Left(_) => assertEquals(expectedError, Errors.NONE)
-        case Right(error) => assertEquals(expectedError, error)
+      try {
+        partition.fetchOffsetSnapshot(currentLeaderEpoch, fetchOnlyFromLeader = fetchOnlyLeader)
+        assertEquals(Errors.NONE, expectedError)
+      } catch {
+        case error: ApiException => assertEquals(expectedError, Errors.forException(error))
       }
     }
 
@@ -480,8 +472,7 @@ class PartitionTest {
     val leader = brokerId
     val follower1 = brokerId + 1
     val follower2 = brokerId + 2
-    val controllerId = brokerId + 3
-    val replicas = List[Integer](leader, follower1, follower2).asJava
+    val replicas = List(leader, follower1, follower2)
     val isr = List[Integer](leader, follower2).asJava
     val leaderEpoch = 8
     val batch1 = TestUtils.records(records = List(
@@ -491,32 +482,36 @@ class PartitionTest {
       new SimpleRecord(20,"k4".getBytes, "v2".getBytes),
       new SimpleRecord(21,"k5".getBytes, "v3".getBytes)))
 
-    val partition = Partition(topicPartition, time, replicaManager)
+    val leaderState = new LeaderAndIsrPartitionState()
+      .setControllerEpoch(controllerEpoch)
+      .setLeader(leader)
+      .setLeaderEpoch(leaderEpoch)
+      .setIsr(isr)
+      .setZkVersion(1)
+      .setReplicas(replicas.map(Int.box).asJava)
+      .setIsNew(true)
+
     assertTrue("Expected first makeLeader() to return 'leader changed'",
-      partition.makeLeader(controllerId, new LeaderAndIsrRequest.PartitionState(controllerEpoch, leader, leaderEpoch, isr, 1, replicas, true), 0))
+      partition.makeLeader(leaderState, offsetCheckpoints))
     assertEquals("Current leader epoch", leaderEpoch, partition.getLeaderEpoch)
-    assertEquals("ISR", Set[Integer](leader, follower2), partition.inSyncReplicas.map(_.brokerId))
+    assertEquals("ISR", Set[Integer](leader, follower2), partition.inSyncReplicaIds)
 
     // after makeLeader(() call, partition should know about all the replicas
-    val leaderReplica = partition.getReplica(leader).get
-    val follower1Replica = partition.getReplica(follower1).get
-    val follower2Replica = partition.getReplica(follower2).get
-
     // append records with initial leader epoch
-    partition.appendRecordsToLeader(batch1, isFromClient = true)
-    partition.appendRecordsToLeader(batch2, isFromClient = true)
-    assertEquals("Expected leader's HW not move", leaderReplica.logStartOffset, leaderReplica.highWatermark.messageOffset)
+    partition.appendRecordsToLeader(batch1, origin = AppendOrigin.Client, requiredAcks = 0)
+    partition.appendRecordsToLeader(batch2, origin = AppendOrigin.Client, requiredAcks = 0)
+    assertEquals("Expected leader's HW not move", partition.localLogOrException.logStartOffset,
+      partition.localLogOrException.highWatermark)
 
     // let the follower in ISR move leader's HW to move further but below LEO
-    def readResult(fetchInfo: FetchDataInfo, leaderReplica: Replica): LogReadResult = {
-      LogReadResult(info = fetchInfo,
-        highWatermark = leaderReplica.highWatermark.messageOffset,
-        leaderLogStartOffset = leaderReplica.logStartOffset,
-        leaderLogEndOffset = leaderReplica.logEndOffset,
-        followerLogStartOffset = 0,
-        fetchTimeMs = time.milliseconds,
-        readSize = 10240,
-        lastStableOffset = None)
+    def updateFollowerFetchState(followerId: Int, fetchOffsetMetadata: LogOffsetMetadata): Unit = {
+      partition.updateFollowerFetchState(
+        followerId,
+        followerFetchOffsetMetadata = fetchOffsetMetadata,
+        followerStartOffset = 0L,
+        followerFetchTimeMs = time.milliseconds(),
+        leaderEndOffset = partition.localLogOrException.logEndOffset,
+        lastSentHighwatermark = partition.localLogOrException.highWatermark)
     }
 
     def fetchOffsetsForTimestamp(timestamp: Long, isolation: Option[IsolationLevel]): Either[ApiException, Option[TimestampAndOffset]] = {
@@ -532,20 +527,18 @@ class PartitionTest {
       }
     }
 
-    // Update follower 1
-    partition.updateReplicaLogReadResult(
-      follower1Replica, readResult(FetchDataInfo(LogOffsetMetadata(0), batch1), leaderReplica))
-    partition.updateReplicaLogReadResult(
-      follower1Replica, readResult(FetchDataInfo(LogOffsetMetadata(2), batch2), leaderReplica))
+    when(stateStore.expandIsr(controllerEpoch, new LeaderAndIsr(leader, leaderEpoch,
+      List(leader, follower2, follower1), 1)))
+      .thenReturn(Some(2))
+
+    updateFollowerFetchState(follower1, LogOffsetMetadata(0))
+    updateFollowerFetchState(follower1, LogOffsetMetadata(2))
 
-    // Update follower 2
-    partition.updateReplicaLogReadResult(
-      follower2Replica, readResult(FetchDataInfo(LogOffsetMetadata(0), batch1), leaderReplica))
-    partition.updateReplicaLogReadResult(
-      follower2Replica, readResult(FetchDataInfo(LogOffsetMetadata(2), batch2), leaderReplica))
+    updateFollowerFetchState(follower2, LogOffsetMetadata(0))
+    updateFollowerFetchState(follower2, LogOffsetMetadata(2))
 
     // At this point, the leader has gotten 5 writes, but followers have only fetched two
-    assertEquals(2, partition.localReplica.get.highWatermark.messageOffset)
+    assertEquals(2, partition.localLogOrException.highWatermark)
 
     // Get the LEO
     fetchOffsetsForTimestamp(ListOffsetRequest.LATEST_TIMESTAMP, None) match {
@@ -565,12 +558,28 @@ class PartitionTest {
     assertEquals(Right(None), fetchOffsetsForTimestamp(30, Some(IsolationLevel.READ_UNCOMMITTED)))
 
     // Make into a follower
-    assertTrue(partition.makeFollower(controllerId,
-      new LeaderAndIsrRequest.PartitionState(controllerEpoch, follower2, leaderEpoch + 1, isr, 1, replicas, false), 1))
+    val followerState = new LeaderAndIsrPartitionState()
+      .setControllerEpoch(controllerEpoch)
+      .setLeader(follower2)
+      .setLeaderEpoch(leaderEpoch + 1)
+      .setIsr(isr)
+      .setZkVersion(4)
+      .setReplicas(replicas.map(Int.box).asJava)
+      .setIsNew(false)
+
+    assertTrue(partition.makeFollower(followerState, offsetCheckpoints))
 
     // Back to leader, this resets the startLogOffset for this epoch (to 2), we're now in the fault condition
-    assertTrue(partition.makeLeader(controllerId,
-      new LeaderAndIsrRequest.PartitionState(controllerEpoch, leader, leaderEpoch + 2, isr, 1, replicas, false), 2))
+    val newLeaderState = new LeaderAndIsrPartitionState()
+      .setControllerEpoch(controllerEpoch)
+      .setLeader(leader)
+      .setLeaderEpoch(leaderEpoch + 2)
+      .setIsr(isr)
+      .setZkVersion(5)
+      .setReplicas(replicas.map(Int.box).asJava)
+      .setIsNew(false)
+
+    assertTrue(partition.makeLeader(newLeaderState, offsetCheckpoints))
 
     // Try to get offsets as a client
     fetchOffsetsForTimestamp(ListOffsetRequest.LATEST_TIMESTAMP, Some(IsolationLevel.READ_UNCOMMITTED)) match {
@@ -611,12 +620,13 @@ class PartitionTest {
       case Left(e: ApiException) => fail(s"Should have seen OffsetNotAvailableException, saw $e")
     }
 
+    when(stateStore.expandIsr(controllerEpoch, new LeaderAndIsr(leader, leaderEpoch + 2,
+      List(leader, follower2, follower1), 5)))
+      .thenReturn(Some(2))
 
     // Next fetch from replicas, HW is moved up to 5 (ahead of the LEO)
-    partition.updateReplicaLogReadResult(
-      follower1Replica, readResult(FetchDataInfo(LogOffsetMetadata(5), MemoryRecords.EMPTY), leaderReplica))
-    partition.updateReplicaLogReadResult(
-      follower2Replica, readResult(FetchDataInfo(LogOffsetMetadata(5), MemoryRecords.EMPTY), leaderReplica))
+    updateFollowerFetchState(follower1, LogOffsetMetadata(5))
+    updateFollowerFetchState(follower2, LogOffsetMetadata(5))
 
     // Error goes away
     fetchOffsetsForTimestamp(ListOffsetRequest.LATEST_TIMESTAMP, Some(IsolationLevel.READ_UNCOMMITTED)) match {
@@ -629,45 +639,38 @@ class PartitionTest {
     assertEquals(Right(None), fetchOffsetsForTimestamp(100, Some(IsolationLevel.READ_UNCOMMITTED)))
   }
 
-
   private def setupPartitionWithMocks(leaderEpoch: Int,
                                       isLeader: Boolean,
-                                      log: Log = logManager.getOrCreateLog(topicPartition, logConfig)): Partition = {
-    val replica = new Replica(brokerId, topicPartition, time, log = Some(log))
-    val replicaManager: ReplicaManager = EasyMock.mock(classOf[ReplicaManager])
-    val zkClient: KafkaZkClient = EasyMock.mock(classOf[KafkaZkClient])
-
-    val partition = new Partition(topicPartition,
-      isOffline = false,
-      replicaLagTimeMaxMs = Defaults.ReplicaLagTimeMaxMs,
-      interBrokerProtocolVersion = ApiVersion.latestVersion,
-      localBrokerId = brokerId,
-      time,
-      replicaManager,
-      logManager,
-      zkClient)
+                                      log: Log = logManager.getOrCreateLog(topicPartition, () => logConfig)): Partition = {
+    partition.createLogIfNotExists(isNew = false, isFutureReplica = false, offsetCheckpoints)
 
-    EasyMock.replay(replicaManager, zkClient)
-
-    partition.addReplicaIfNotExists(replica)
-
-    val controllerId = 0
     val controllerEpoch = 0
     val replicas = List[Integer](brokerId, brokerId + 1).asJava
     val isr = replicas
 
     if (isLeader) {
       assertTrue("Expected become leader transition to succeed",
-        partition.makeLeader(controllerId, new LeaderAndIsrRequest.PartitionState(controllerEpoch, brokerId,
-          leaderEpoch, isr, 1, replicas, true), 0))
+        partition.makeLeader(new LeaderAndIsrPartitionState()
+          .setControllerEpoch(controllerEpoch)
+          .setLeader(brokerId)
+          .setLeaderEpoch(leaderEpoch)
+          .setIsr(isr)
+          .setZkVersion(1)
+          .setReplicas(replicas)
+          .setIsNew(true), offsetCheckpoints))
       assertEquals(leaderEpoch, partition.getLeaderEpoch)
-      assertEquals(Some(replica), partition.leaderReplicaIfLocal)
     } else {
       assertTrue("Expected become follower transition to succeed",
-        partition.makeFollower(controllerId, new LeaderAndIsrRequest.PartitionState(controllerEpoch, brokerId + 1,
-          leaderEpoch, isr, 1, replicas, true), 0))
+        partition.makeFollower(new LeaderAndIsrPartitionState()
+          .setControllerEpoch(controllerEpoch)
+          .setLeader(brokerId + 1)
+          .setLeaderEpoch(leaderEpoch)
+          .setIsr(isr)
+          .setZkVersion(1)
+          .setReplicas(replicas)
+          .setIsNew(true), offsetCheckpoints))
       assertEquals(leaderEpoch, partition.getLeaderEpoch)
-      assertEquals(None, partition.leaderReplicaIfLocal)
+      assertEquals(None, partition.leaderLogIfLocal)
     }
 
     partition
@@ -675,25 +678,22 @@ class PartitionTest {
 
   @Test
   def testAppendRecordsAsFollowerBelowLogStartOffset(): Unit = {
-    val log = logManager.getOrCreateLog(topicPartition, logConfig)
-    val replica = new Replica(brokerId, topicPartition, time, log = Some(log))
-    val partition = Partition(topicPartition, time, replicaManager)
-    partition.addReplicaIfNotExists(replica)
-    assertEquals(Some(replica), partition.localReplica)
+    partition.createLogIfNotExists(isNew = false, isFutureReplica = false, offsetCheckpoints)
+    val log = partition.localLogOrException
 
     val initialLogStartOffset = 5L
     partition.truncateFullyAndStartAt(initialLogStartOffset, isFuture = false)
     assertEquals(s"Log end offset after truncate fully and start at $initialLogStartOffset:",
-                 initialLogStartOffset, replica.logEndOffset)
+                 initialLogStartOffset, log.logEndOffset)
     assertEquals(s"Log start offset after truncate fully and start at $initialLogStartOffset:",
-                 initialLogStartOffset, replica.logStartOffset)
+                 initialLogStartOffset, log.logStartOffset)
 
     // verify that we cannot append records that do not contain log start offset even if the log is empty
     assertThrows[UnexpectedAppendOffsetException] {
       // append one record with offset = 3
       partition.appendRecordsToFollowerOrFutureReplica(createRecords(List(new SimpleRecord("k1".getBytes, "v1".getBytes)), baseOffset = 3L), isFuture = false)
     }
-    assertEquals(s"Log end offset should not change after failure to append", initialLogStartOffset, replica.logEndOffset)
+    assertEquals(s"Log end offset should not change after failure to append", initialLogStartOffset, log.logEndOffset)
 
     // verify that we can append records that contain log start offset, even when first
     // offset < log start offset if the log is empty
@@ -703,13 +703,13 @@ class PartitionTest {
                                      new SimpleRecord("k3".getBytes, "v3".getBytes)),
                                 baseOffset = newLogStartOffset)
     partition.appendRecordsToFollowerOrFutureReplica(records, isFuture = false)
-    assertEquals(s"Log end offset after append of 3 records with base offset $newLogStartOffset:", 7L, replica.logEndOffset)
-    assertEquals(s"Log start offset after append of 3 records with base offset $newLogStartOffset:", newLogStartOffset, replica.logStartOffset)
+    assertEquals(s"Log end offset after append of 3 records with base offset $newLogStartOffset:", 7L, log.logEndOffset)
+    assertEquals(s"Log start offset after append of 3 records with base offset $newLogStartOffset:", newLogStartOffset, log.logStartOffset)
 
     // and we can append more records after that
     partition.appendRecordsToFollowerOrFutureReplica(createRecords(List(new SimpleRecord("k1".getBytes, "v1".getBytes)), baseOffset = 7L), isFuture = false)
-    assertEquals(s"Log end offset after append of 1 record at offset 7:", 8L, replica.logEndOffset)
-    assertEquals(s"Log start offset not expected to change:", newLogStartOffset, replica.logStartOffset)
+    assertEquals(s"Log end offset after append of 1 record at offset 7:", 8L, log.logEndOffset)
+    assertEquals(s"Log start offset not expected to change:", newLogStartOffset, log.logStartOffset)
 
     // but we cannot append to offset < log start if the log is not empty
     assertThrows[UnexpectedAppendOffsetException] {
@@ -718,56 +718,42 @@ class PartitionTest {
                                    baseOffset = 3L)
       partition.appendRecordsToFollowerOrFutureReplica(records2, isFuture = false)
     }
-    assertEquals(s"Log end offset should not change after failure to append", 8L, replica.logEndOffset)
+    assertEquals(s"Log end offset should not change after failure to append", 8L, log.logEndOffset)
 
     // we still can append to next offset
     partition.appendRecordsToFollowerOrFutureReplica(createRecords(List(new SimpleRecord("k1".getBytes, "v1".getBytes)), baseOffset = 8L), isFuture = false)
-    assertEquals(s"Log end offset after append of 1 record at offset 8:", 9L, replica.logEndOffset)
-    assertEquals(s"Log start offset not expected to change:", newLogStartOffset, replica.logStartOffset)
+    assertEquals(s"Log end offset after append of 1 record at offset 8:", 9L, log.logEndOffset)
+    assertEquals(s"Log start offset not expected to change:", newLogStartOffset, log.logStartOffset)
   }
 
   @Test
   def testListOffsetIsolationLevels(): Unit = {
-    val log = logManager.getOrCreateLog(topicPartition, logConfig)
-    val replica = new Replica(brokerId, topicPartition, time, log = Some(log))
-    val replicaManager: ReplicaManager = EasyMock.mock(classOf[ReplicaManager])
-    val zkClient: KafkaZkClient = EasyMock.mock(classOf[KafkaZkClient])
-
-    val partition = new Partition(topicPartition,
-      isOffline = false,
-      replicaLagTimeMaxMs = Defaults.ReplicaLagTimeMaxMs,
-      interBrokerProtocolVersion = ApiVersion.latestVersion,
-      localBrokerId = brokerId,
-      time,
-      replicaManager,
-      logManager,
-      zkClient)
-
-    val controllerId = 0
     val controllerEpoch = 0
     val leaderEpoch = 5
     val replicas = List[Integer](brokerId, brokerId + 1).asJava
     val isr = replicas
 
-    EasyMock.expect(replicaManager.tryCompleteDelayedFetch(EasyMock.anyObject[TopicPartitionOperationKey]))
-        .andVoid()
-
-    EasyMock.replay(replicaManager, zkClient)
+    doNothing().when(delayedOperations).checkAndCompleteFetch()
 
-    partition.addReplicaIfNotExists(replica)
+    partition.createLogIfNotExists(isNew = false, isFutureReplica = false, offsetCheckpoints)
 
     assertTrue("Expected become leader transition to succeed",
-      partition.makeLeader(controllerId, new LeaderAndIsrRequest.PartitionState(controllerEpoch, brokerId,
-        leaderEpoch, isr, 1, replicas, true), 0))
+      partition.makeLeader(new LeaderAndIsrPartitionState()
+        .setControllerEpoch(controllerEpoch)
+        .setLeader(brokerId)
+        .setLeaderEpoch(leaderEpoch)
+        .setIsr(isr)
+        .setZkVersion(1)
+        .setReplicas(replicas)
+        .setIsNew(true), offsetCheckpoints))
     assertEquals(leaderEpoch, partition.getLeaderEpoch)
-    assertEquals(Some(replica), partition.leaderReplicaIfLocal)
 
     val records = createTransactionalRecords(List(
       new SimpleRecord("k1".getBytes, "v1".getBytes),
       new SimpleRecord("k2".getBytes, "v2".getBytes),
       new SimpleRecord("k3".getBytes, "v3".getBytes)),
       baseOffset = 0L)
-    partition.appendRecordsToLeader(records, isFromClient = true)
+    partition.appendRecordsToLeader(records, origin = AppendOrigin.Client, requiredAcks = 0)
 
     def fetchLatestOffset(isolationLevel: Option[IsolationLevel]): TimestampAndOffset = {
       val res = partition.fetchOffsetForTimestamp(ListOffsetRequest.LATEST_TIMESTAMP,
@@ -791,7 +777,7 @@ class PartitionTest {
     assertEquals(0L, fetchLatestOffset(isolationLevel = Some(IsolationLevel.READ_UNCOMMITTED)).offset)
     assertEquals(0L, fetchLatestOffset(isolationLevel = Some(IsolationLevel.READ_COMMITTED)).offset)
 
-    replica.highWatermark = LogOffsetMetadata(1L)
+    partition.log.get.updateHighWatermark(1L)
 
     assertEquals(3L, fetchLatestOffset(isolationLevel = None).offset)
     assertEquals(1L, fetchLatestOffset(isolationLevel = Some(IsolationLevel.READ_UNCOMMITTED)).offset)
@@ -804,23 +790,14 @@ class PartitionTest {
 
   @Test
   def testGetReplica(): Unit = {
-    val log = logManager.getOrCreateLog(topicPartition, logConfig)
-    val replica = new Replica(brokerId, topicPartition, time, log = Some(log))
-    val partition = Partition(topicPartition, time, replicaManager)
-
-    assertEquals(None, partition.localReplica)
+    assertEquals(None, partition.log)
     assertThrows[ReplicaNotAvailableException] {
-      partition.localReplicaOrException
+      partition.localLogOrException
     }
-
-    partition.addReplicaIfNotExists(replica)
-    assertEquals(Some(replica), partition.localReplica)
-    assertEquals(replica, partition.localReplicaOrException)
   }
 
   @Test
   def testAppendRecordsToFollowerWithNoReplicaThrowsException(): Unit = {
-    val partition = Partition(topicPartition, time, replicaManager)
     assertThrows[ReplicaNotAvailableException] {
       partition.appendRecordsToFollowerOrFutureReplica(
            createRecords(List(new SimpleRecord("k1".getBytes, "v1".getBytes)), baseOffset = 0L), isFuture = false)
@@ -829,22 +806,37 @@ class PartitionTest {
 
   @Test
   def testMakeFollowerWithNoLeaderIdChange(): Unit = {
-    val partition = Partition(topicPartition, time, replicaManager)
-
     // Start off as follower
-    var partitionStateInfo = new LeaderAndIsrRequest.PartitionState(0, 1, 1,
-      List[Integer](0, 1, 2).asJava, 1, List[Integer](0, 1, 2).asJava, false)
-    partition.makeFollower(0, partitionStateInfo, 0)
+    var partitionState = new LeaderAndIsrPartitionState()
+      .setControllerEpoch(0)
+      .setLeader(1)
+      .setLeaderEpoch(1)
+      .setIsr(List[Integer](0, 1, 2, brokerId).asJava)
+      .setZkVersion(1)
+      .setReplicas(List[Integer](0, 1, 2, brokerId).asJava)
+      .setIsNew(false)
+    partition.makeFollower(partitionState, offsetCheckpoints)
 
     // Request with same leader and epoch increases by only 1, do become-follower steps
-    partitionStateInfo = new LeaderAndIsrRequest.PartitionState(0, 1, 4,
-      List[Integer](0, 1, 2).asJava, 1, List[Integer](0, 1, 2).asJava, false)
-    assertTrue(partition.makeFollower(0, partitionStateInfo, 2))
+    partitionState = new LeaderAndIsrPartitionState()
+      .setControllerEpoch(0)
+      .setLeader(1)
+      .setLeaderEpoch(4)
+      .setIsr(List[Integer](0, 1, 2, brokerId).asJava)
+      .setZkVersion(1)
+      .setReplicas(List[Integer](0, 1, 2, brokerId).asJava)
+      .setIsNew(false)
+    assertTrue(partition.makeFollower(partitionState, offsetCheckpoints))
 
     // Request with same leader and same epoch, skip become-follower steps
-    partitionStateInfo = new LeaderAndIsrRequest.PartitionState(0, 1, 4,
-      List[Integer](0, 1, 2).asJava, 1, List[Integer](0, 1, 2).asJava, false)
-    assertFalse(partition.makeFollower(0, partitionStateInfo, 2))
+    partitionState = new LeaderAndIsrPartitionState()
+      .setControllerEpoch(0)
+      .setLeader(1)
+      .setLeaderEpoch(4)
+      .setIsr(List[Integer](0, 1, 2, brokerId).asJava)
+      .setZkVersion(1)
+      .setReplicas(List[Integer](0, 1, 2, brokerId).asJava)
+    assertFalse(partition.makeFollower(partitionState, offsetCheckpoints))
   }
 
   @Test
@@ -853,7 +845,6 @@ class PartitionTest {
     val leader = brokerId
     val follower1 = brokerId + 1
     val follower2 = brokerId + 2
-    val controllerId = brokerId + 3
     val replicas = List[Integer](leader, follower1, follower2).asJava
     val isr = List[Integer](leader, follower2).asJava
     val leaderEpoch = 8
@@ -865,62 +856,79 @@ class PartitionTest {
     val batch3 = TestUtils.records(records = List(new SimpleRecord("k6".getBytes, "v1".getBytes),
                                                   new SimpleRecord("k7".getBytes, "v2".getBytes)))
 
-    val partition = Partition(topicPartition, time, replicaManager)
+    val leaderState = new LeaderAndIsrPartitionState()
+      .setControllerEpoch(controllerEpoch)
+      .setLeader(leader)
+      .setLeaderEpoch(leaderEpoch)
+      .setIsr(isr)
+      .setZkVersion(1)
+      .setReplicas(replicas)
+      .setIsNew(true)
     assertTrue("Expected first makeLeader() to return 'leader changed'",
-               partition.makeLeader(controllerId, new LeaderAndIsrRequest.PartitionState(controllerEpoch, leader, leaderEpoch, isr, 1, replicas, true), 0))
+               partition.makeLeader(leaderState, offsetCheckpoints))
     assertEquals("Current leader epoch", leaderEpoch, partition.getLeaderEpoch)
-    assertEquals("ISR", Set[Integer](leader, follower2), partition.inSyncReplicas.map(_.brokerId))
+    assertEquals("ISR", Set[Integer](leader, follower2), partition.inSyncReplicaIds)
 
     // after makeLeader(() call, partition should know about all the replicas
-    val leaderReplica = partition.getReplica(leader).get
-    val follower1Replica = partition.getReplica(follower1).get
-    val follower2Replica = partition.getReplica(follower2).get
-
     // append records with initial leader epoch
-    val lastOffsetOfFirstBatch = partition.appendRecordsToLeader(batch1, isFromClient = true).lastOffset
-    partition.appendRecordsToLeader(batch2, isFromClient = true)
-    assertEquals("Expected leader's HW not move", leaderReplica.logStartOffset, leaderReplica.highWatermark.messageOffset)
+    val lastOffsetOfFirstBatch = partition.appendRecordsToLeader(batch1, origin = AppendOrigin.Client,
+      requiredAcks = 0).lastOffset
+    partition.appendRecordsToLeader(batch2, origin = AppendOrigin.Client, requiredAcks = 0)
+    assertEquals("Expected leader's HW not move", partition.localLogOrException.logStartOffset,
+      partition.log.get.highWatermark)
 
     // let the follower in ISR move leader's HW to move further but below LEO
-    def readResult(fetchInfo: FetchDataInfo, leaderReplica: Replica): LogReadResult = {
-      LogReadResult(info = fetchInfo,
-                    highWatermark = leaderReplica.highWatermark.messageOffset,
-                    leaderLogStartOffset = leaderReplica.logStartOffset,
-                    leaderLogEndOffset = leaderReplica.logEndOffset,
-                    followerLogStartOffset = 0,
-                    fetchTimeMs = time.milliseconds,
-                    readSize = 10240,
-                    lastStableOffset = None)
+    def updateFollowerFetchState(followerId: Int, fetchOffsetMetadata: LogOffsetMetadata): Unit = {
+      partition.updateFollowerFetchState(
+        followerId,
+        followerFetchOffsetMetadata = fetchOffsetMetadata,
+        followerStartOffset = 0L,
+        followerFetchTimeMs = time.milliseconds(),
+        leaderEndOffset = partition.localLogOrException.logEndOffset,
+        lastSentHighwatermark = partition.localLogOrException.highWatermark)
     }
-    partition.updateReplicaLogReadResult(
-      follower2Replica, readResult(FetchDataInfo(LogOffsetMetadata(0), batch1), leaderReplica))
-    partition.updateReplicaLogReadResult(
-      follower2Replica, readResult(FetchDataInfo(LogOffsetMetadata(lastOffsetOfFirstBatch), batch2), leaderReplica))
-    assertEquals("Expected leader's HW", lastOffsetOfFirstBatch, leaderReplica.highWatermark.messageOffset)
+
+    updateFollowerFetchState(follower2, LogOffsetMetadata(0))
+    updateFollowerFetchState(follower2, LogOffsetMetadata(lastOffsetOfFirstBatch))
+    assertEquals("Expected leader's HW", lastOffsetOfFirstBatch, partition.log.get.highWatermark)
 
     // current leader becomes follower and then leader again (without any new records appended)
-    partition.makeFollower(
-      controllerId, new LeaderAndIsrRequest.PartitionState(controllerEpoch, follower2, leaderEpoch + 1, isr, 1, replicas, false), 1)
+    val followerState = new LeaderAndIsrPartitionState()
+      .setControllerEpoch(controllerEpoch)
+      .setLeader(follower2)
+      .setLeaderEpoch(leaderEpoch + 1)
+      .setIsr(isr)
+      .setZkVersion(1)
+      .setReplicas(replicas)
+      .setIsNew(false)
+    partition.makeFollower(followerState, offsetCheckpoints)
+
+    val newLeaderState = new LeaderAndIsrPartitionState()
+      .setControllerEpoch(controllerEpoch)
+      .setLeader(leader)
+      .setLeaderEpoch(leaderEpoch + 2)
+      .setIsr(isr)
+      .setZkVersion(1)
+      .setReplicas(replicas)
+      .setIsNew(false)
     assertTrue("Expected makeLeader() to return 'leader changed' after makeFollower()",
-               partition.makeLeader(controllerEpoch, new LeaderAndIsrRequest.PartitionState(
-                 controllerEpoch, leader, leaderEpoch + 2, isr, 1, replicas, false), 2))
-    val currentLeaderEpochStartOffset = leaderReplica.logEndOffset
+               partition.makeLeader(newLeaderState, offsetCheckpoints))
+    val currentLeaderEpochStartOffset = partition.localLogOrException.logEndOffset
 
     // append records with the latest leader epoch
-    partition.appendRecordsToLeader(batch3, isFromClient = true)
+    partition.appendRecordsToLeader(batch3, origin = AppendOrigin.Client, requiredAcks = 0)
 
     // fetch from follower not in ISR from log start offset should not add this follower to ISR
-    partition.updateReplicaLogReadResult(follower1Replica,
-                                         readResult(FetchDataInfo(LogOffsetMetadata(0), batch1), leaderReplica))
-    partition.updateReplicaLogReadResult(follower1Replica,
-                                         readResult(FetchDataInfo(LogOffsetMetadata(lastOffsetOfFirstBatch), batch2), leaderReplica))
-    assertEquals("ISR", Set[Integer](leader, follower2), partition.inSyncReplicas.map(_.brokerId))
+    updateFollowerFetchState(follower1, LogOffsetMetadata(0))
+    updateFollowerFetchState(follower1, LogOffsetMetadata(lastOffsetOfFirstBatch))
+    assertEquals("ISR", Set[Integer](leader, follower2), partition.inSyncReplicaIds)
 
     // fetch from the follower not in ISR from start offset of the current leader epoch should
     // add this follower to ISR
-    partition.updateReplicaLogReadResult(follower1Replica,
-                                         readResult(FetchDataInfo(LogOffsetMetadata(currentLeaderEpochStartOffset), batch3), leaderReplica))
-    assertEquals("ISR", Set[Integer](leader, follower1, follower2), partition.inSyncReplicas.map(_.brokerId))
+    when(stateStore.expandIsr(controllerEpoch, new LeaderAndIsr(leader, leaderEpoch + 2,
+      List(leader, follower2, follower1), 1))).thenReturn(Some(2))
+    updateFollowerFetchState(follower1, LogOffsetMetadata(currentLeaderEpochStartOffset))
+    assertEquals("ISR", Set[Integer](leader, follower1, follower2), partition.inSyncReplicaIds)
   }
 
   /**
@@ -932,9 +940,6 @@ class PartitionTest {
    */
   @Test
   def testDelayedFetchAfterAppendRecords(): Unit = {
-    val replicaManager: ReplicaManager = EasyMock.mock(classOf[ReplicaManager])
-    val zkClient: KafkaZkClient = EasyMock.mock(classOf[KafkaZkClient])
-    val controllerId = 0
     val controllerEpoch = 0
     val leaderEpoch = 5
     val replicaIds = List[Integer](brokerId, brokerId + 1).asJava
@@ -942,36 +947,44 @@ class PartitionTest {
     val logConfig = LogConfig(new Properties)
 
     val topicPartitions = (0 until 5).map { i => new TopicPartition("test-topic", i) }
-    val logs = topicPartitions.map { tp => logManager.getOrCreateLog(tp, logConfig) }
-    val replicas = logs.map { log => new Replica(brokerId, log.topicPartition, time, log = Some(log)) }
-    val partitions = replicas.map { replica =>
-      val tp = replica.topicPartition
+    val logs = topicPartitions.map { tp => logManager.getOrCreateLog(tp, () => logConfig) }
+    val partitions = ListBuffer.empty[Partition]
+
+    logs.foreach { log =>
+      val tp = log.topicPartition
+      val delayedOperations: DelayedOperations = mock(classOf[DelayedOperations])
       val partition = new Partition(tp,
-        isOffline = false,
         replicaLagTimeMaxMs = Defaults.ReplicaLagTimeMaxMs,
         interBrokerProtocolVersion = ApiVersion.latestVersion,
         localBrokerId = brokerId,
         time,
-        replicaManager,
-        logManager,
-        zkClient)
-      partition.addReplicaIfNotExists(replica)
-      partition.makeLeader(controllerId, new LeaderAndIsrRequest.PartitionState(controllerEpoch, brokerId,
-        leaderEpoch, isr, 1, replicaIds, true), 0)
-      partition
-    }
+        stateStore,
+        delayedOperations,
+        metadataCache,
+        logManager)
+
+      when(delayedOperations.checkAndCompleteFetch())
+        .thenAnswer(new Answer[Unit] {
+          override def answer(invocation: InvocationOnMock): Unit = {
+            // Acquire leaderIsrUpdate read lock of a different partition when completing delayed fetch
+            val anotherPartition = (tp.partition + 1) % topicPartitions.size
+            val partition = partitions(anotherPartition)
+            partition.fetchOffsetSnapshot(Optional.of(leaderEpoch), fetchOnlyFromLeader = true)
+          }
+        })
 
-    // Acquire leaderIsrUpdate read lock of a different partition when completing delayed fetch
-    val tpKey: Capture[TopicPartitionOperationKey] = EasyMock.newCapture()
-    EasyMock.expect(replicaManager.tryCompleteDelayedFetch(EasyMock.capture(tpKey)))
-      .andAnswer(new IAnswer[Unit] {
-        override def answer(): Unit = {
-          val anotherPartition = (tpKey.getValue.partition + 1) % topicPartitions.size
-          val partition = partitions(anotherPartition)
-          partition.fetchOffsetSnapshot(Optional.of(leaderEpoch), fetchOnlyFromLeader = true)
-        }
-      }).anyTimes()
-    EasyMock.replay(replicaManager, zkClient)
+      partition.setLog(log, isFutureLog = false)
+      val leaderState = new LeaderAndIsrPartitionState()
+        .setControllerEpoch(controllerEpoch)
+        .setLeader(brokerId)
+        .setLeaderEpoch(leaderEpoch)
+        .setIsr(isr)
+        .setZkVersion(1)
+        .setReplicas(replicaIds)
+        .setIsNew(true)
+      partition.makeLeader(leaderState, offsetCheckpoints)
+      partitions += partition
+    }
 
     def createRecords(baseOffset: Long): MemoryRecords = {
       val records = List(
@@ -997,7 +1010,11 @@ class PartitionTest {
       // Append records to partitions, one partition-per-thread
       val futures = partitions.map { partition =>
         executor.submit(CoreUtils.runnable {
-          (1 to 10000).foreach { _ => partition.appendRecordsToLeader(createRecords(baseOffset = 0), isFromClient = true) }
+          (1 to 10000).foreach { _ =>
+            partition.appendRecordsToLeader(createRecords(baseOffset = 0),
+              origin = AppendOrigin.Client,
+              requiredAcks = 0)
+          }
         })
       }
       futures.foreach(_.get(15, TimeUnit.SECONDS))
@@ -1022,8 +1039,7 @@ class PartitionTest {
   }
 
   def createTransactionalRecords(records: Iterable[SimpleRecord],
-                                 baseOffset: Long,
-                                 partitionLeaderEpoch: Int = 0): MemoryRecords = {
+                                 baseOffset: Long): MemoryRecords = {
     val producerId = 1L
     val producerEpoch = 0.toShort
     val baseSequence = 0
@@ -1035,4 +1051,593 @@ class PartitionTest {
     builder.build()
   }
 
+  /**
+    * Test for AtMinIsr partition state. We set the partition replica set size as 3, but only set one replica as an ISR.
+    * As the default minIsr configuration is 1, then the partition should be at min ISR (isAtMinIsr = true).
+    */
+  @Test
+  def testAtMinIsr(): Unit = {
+    val controllerEpoch = 3
+    val leader = brokerId
+    val follower1 = brokerId + 1
+    val follower2 = brokerId + 2
+    val replicas = List[Integer](leader, follower1, follower2).asJava
+    val isr = List[Integer](leader).asJava
+    val leaderEpoch = 8
+
+    assertFalse(partition.isAtMinIsr)
+    // Make isr set to only have leader to trigger AtMinIsr (default min isr config is 1)
+    val leaderState = new LeaderAndIsrPartitionState()
+      .setControllerEpoch(controllerEpoch)
+      .setLeader(leader)
+      .setLeaderEpoch(leaderEpoch)
+      .setIsr(isr)
+      .setZkVersion(1)
+      .setReplicas(replicas)
+      .setIsNew(true)
+    partition.makeLeader(leaderState, offsetCheckpoints)
+    assertTrue(partition.isAtMinIsr)
+  }
+
+  @Test
+  def testUpdateFollowerFetchState(): Unit = {
+    val log = logManager.getOrCreateLog(topicPartition, () => logConfig)
+    seedLogData(log, numRecords = 6, leaderEpoch = 4)
+
+    val controllerEpoch = 0
+    val leaderEpoch = 5
+    val remoteBrokerId = brokerId + 1
+    val replicas = List[Integer](brokerId, remoteBrokerId).asJava
+    val isr = replicas
+
+    doNothing().when(delayedOperations).checkAndCompleteFetch()
+
+    partition.createLogIfNotExists(isNew = false, isFutureReplica = false, offsetCheckpoints)
+
+    val initializeTimeMs = time.milliseconds()
+    assertTrue("Expected become leader transition to succeed",
+      partition.makeLeader(
+        new LeaderAndIsrPartitionState()
+          .setControllerEpoch(controllerEpoch)
+          .setLeader(brokerId)
+          .setLeaderEpoch(leaderEpoch)
+          .setIsr(isr)
+          .setZkVersion(1)
+          .setReplicas(replicas)
+          .setIsNew(true),
+        offsetCheckpoints))
+
+    val remoteReplica = partition.getReplica(remoteBrokerId).get
+    assertEquals(initializeTimeMs, remoteReplica.lastCaughtUpTimeMs)
+    assertEquals(LogOffsetMetadata.UnknownOffsetMetadata.messageOffset, remoteReplica.logEndOffset)
+    assertEquals(Log.UnknownOffset, remoteReplica.logStartOffset)
+
+    time.sleep(500)
+
+    partition.updateFollowerFetchState(remoteBrokerId,
+      followerFetchOffsetMetadata = LogOffsetMetadata(3),
+      followerStartOffset = 0L,
+      followerFetchTimeMs = time.milliseconds(),
+      leaderEndOffset = 6L,
+      lastSentHighwatermark = partition.localLogOrException.highWatermark)
+
+    assertEquals(initializeTimeMs, remoteReplica.lastCaughtUpTimeMs)
+    assertEquals(3L, remoteReplica.logEndOffset)
+    assertEquals(0L, remoteReplica.logStartOffset)
+
+    time.sleep(500)
+
+    partition.updateFollowerFetchState(remoteBrokerId,
+      followerFetchOffsetMetadata = LogOffsetMetadata(6L),
+      followerStartOffset = 0L,
+      followerFetchTimeMs = time.milliseconds(),
+      leaderEndOffset = 6L,
+      lastSentHighwatermark = partition.localLogOrException.highWatermark)
+
+    assertEquals(time.milliseconds(), remoteReplica.lastCaughtUpTimeMs)
+    assertEquals(6L, remoteReplica.logEndOffset)
+    assertEquals(0L, remoteReplica.logStartOffset)
+
+  }
+
+  @Test
+  def testIsrExpansion(): Unit = {
+    val log = logManager.getOrCreateLog(topicPartition, () => logConfig)
+    seedLogData(log, numRecords = 10, leaderEpoch = 4)
+
+    val controllerEpoch = 0
+    val leaderEpoch = 5
+    val remoteBrokerId = brokerId + 1
+    val replicas = List(brokerId, remoteBrokerId)
+    val isr = List[Integer](brokerId).asJava
+
+    doNothing().when(delayedOperations).checkAndCompleteFetch()
+
+    partition.createLogIfNotExists(isNew = false, isFutureReplica = false, offsetCheckpoints)
+    assertTrue(
+      "Expected become leader transition to succeed",
+      partition.makeLeader(
+        new LeaderAndIsrPartitionState()
+          .setControllerEpoch(controllerEpoch)
+          .setLeader(brokerId)
+          .setLeaderEpoch(leaderEpoch)
+          .setIsr(isr)
+          .setZkVersion(1)
+          .setReplicas(replicas.map(Int.box).asJava)
+          .setIsNew(true),
+        offsetCheckpoints)
+    )
+    assertEquals(Set(brokerId), partition.inSyncReplicaIds)
+
+    val remoteReplica = partition.getReplica(remoteBrokerId).get
+    assertEquals(LogOffsetMetadata.UnknownOffsetMetadata.messageOffset, remoteReplica.logEndOffset)
+    assertEquals(Log.UnknownOffset, remoteReplica.logStartOffset)
+
+    partition.updateFollowerFetchState(remoteBrokerId,
+      followerFetchOffsetMetadata = LogOffsetMetadata(3),
+      followerStartOffset = 0L,
+      followerFetchTimeMs = time.milliseconds(),
+      leaderEndOffset = 6L,
+      lastSentHighwatermark = partition.localLogOrException.highWatermark)
+
+    assertEquals(Set(brokerId), partition.inSyncReplicaIds)
+    assertEquals(3L, remoteReplica.logEndOffset)
+    assertEquals(0L, remoteReplica.logStartOffset)
+
+    // The next update should bring the follower back into the ISR
+    val updatedLeaderAndIsr = LeaderAndIsr(
+      leader = brokerId,
+      leaderEpoch = leaderEpoch,
+      isr = List(brokerId, remoteBrokerId),
+      zkVersion = 1)
+    when(stateStore.expandIsr(controllerEpoch, updatedLeaderAndIsr)).thenReturn(Some(2))
+
+    partition.updateFollowerFetchState(remoteBrokerId,
+      followerFetchOffsetMetadata = LogOffsetMetadata(10),
+      followerStartOffset = 0L,
+      followerFetchTimeMs = time.milliseconds(),
+      leaderEndOffset = 6L,
+      lastSentHighwatermark = partition.localLogOrException.highWatermark)
+
+    assertEquals(Set(brokerId, remoteBrokerId), partition.inSyncReplicaIds)
+    assertEquals(10L, remoteReplica.logEndOffset)
+    assertEquals(0L, remoteReplica.logStartOffset)
+  }
+
+  @Test
+  def testIsrNotExpandedIfUpdateFails(): Unit = {
+    val log = logManager.getOrCreateLog(topicPartition, () => logConfig)
+    seedLogData(log, numRecords = 10, leaderEpoch = 4)
+
+    val controllerEpoch = 0
+    val leaderEpoch = 5
+    val remoteBrokerId = brokerId + 1
+    val replicas = List[Integer](brokerId, remoteBrokerId).asJava
+    val isr = List[Integer](brokerId).asJava
+
+    doNothing().when(delayedOperations).checkAndCompleteFetch()
+
+    partition.createLogIfNotExists(isNew = false, isFutureReplica = false, offsetCheckpoints)
+    assertTrue("Expected become leader transition to succeed",
+      partition.makeLeader(
+        new LeaderAndIsrPartitionState()
+          .setControllerEpoch(controllerEpoch)
+          .setLeader(brokerId)
+          .setLeaderEpoch(leaderEpoch)
+          .setIsr(isr)
+          .setZkVersion(1)
+          .setReplicas(replicas)
+          .setIsNew(true),
+        offsetCheckpoints))
+    assertEquals(Set(brokerId), partition.inSyncReplicaIds)
+
+    val remoteReplica = partition.getReplica(remoteBrokerId).get
+    assertEquals(LogOffsetMetadata.UnknownOffsetMetadata.messageOffset, remoteReplica.logEndOffset)
+    assertEquals(Log.UnknownOffset, remoteReplica.logStartOffset)
+
+    // Mock the expected ISR update failure
+    val updatedLeaderAndIsr = LeaderAndIsr(
+      leader = brokerId,
+      leaderEpoch = leaderEpoch,
+      isr = List(brokerId, remoteBrokerId),
+      zkVersion = 1)
+    when(stateStore.expandIsr(controllerEpoch, updatedLeaderAndIsr)).thenReturn(None)
+
+    partition.updateFollowerFetchState(remoteBrokerId,
+      followerFetchOffsetMetadata = LogOffsetMetadata(10),
+      followerStartOffset = 0L,
+      followerFetchTimeMs = time.milliseconds(),
+      leaderEndOffset = 10L,
+      lastSentHighwatermark = partition.localLogOrException.highWatermark)
+
+    // Follower state is updated, but the ISR has not expanded
+    assertEquals(Set(brokerId), partition.inSyncReplicaIds)
+    assertEquals(10L, remoteReplica.logEndOffset)
+    assertEquals(0L, remoteReplica.logStartOffset)
+  }
+
+  @Test
+  def testMaybeShrinkIsr(): Unit = {
+    val log = logManager.getOrCreateLog(topicPartition, () => logConfig)
+    seedLogData(log, numRecords = 10, leaderEpoch = 4)
+
+    val controllerEpoch = 0
+    val leaderEpoch = 5
+    val remoteBrokerId = brokerId + 1
+    val replicas = List(brokerId, remoteBrokerId)
+    val isr = List[Integer](brokerId, remoteBrokerId).asJava
+
+    doNothing().when(delayedOperations).checkAndCompleteFetch()
+
+    val initializeTimeMs = time.milliseconds()
+    partition.createLogIfNotExists(isNew = false, isFutureReplica = false, offsetCheckpoints)
+    assertTrue(
+      "Expected become leader transition to succeed",
+      partition.makeLeader(
+        new LeaderAndIsrPartitionState()
+          .setControllerEpoch(controllerEpoch)
+          .setLeader(brokerId)
+          .setLeaderEpoch(leaderEpoch)
+          .setIsr(isr)
+          .setZkVersion(1)
+          .setReplicas(replicas.map(Int.box).asJava)
+          .setIsNew(true),
+        offsetCheckpoints)
+    )
+    assertEquals(Set(brokerId, remoteBrokerId), partition.inSyncReplicaIds)
+    assertEquals(0L, partition.localLogOrException.highWatermark)
+
+    val remoteReplica = partition.getReplica(remoteBrokerId).get
+    assertEquals(initializeTimeMs, remoteReplica.lastCaughtUpTimeMs)
+    assertEquals(LogOffsetMetadata.UnknownOffsetMetadata.messageOffset, remoteReplica.logEndOffset)
+    assertEquals(Log.UnknownOffset, remoteReplica.logStartOffset)
+
+    // On initialization, the replica is considered caught up and should not be removed
+    partition.maybeShrinkIsr(10000)
+    assertEquals(Set(brokerId, remoteBrokerId), partition.inSyncReplicaIds)
+
+    // If enough time passes without a fetch update, the ISR should shrink
+    time.sleep(10001)
+    val updatedLeaderAndIsr = LeaderAndIsr(
+      leader = brokerId,
+      leaderEpoch = leaderEpoch,
+      isr = List(brokerId),
+      zkVersion = 1)
+    when(stateStore.shrinkIsr(controllerEpoch, updatedLeaderAndIsr)).thenReturn(Some(2))
+
+    partition.maybeShrinkIsr(10000)
+    assertEquals(Set(brokerId), partition.inSyncReplicaIds)
+    assertEquals(10L, partition.localLogOrException.highWatermark)
+  }
+
+  @Test
+  def testShouldNotShrinkIsrIfPreviousFetchIsCaughtUp(): Unit = {
+    val log = logManager.getOrCreateLog(topicPartition, () => logConfig)
+    seedLogData(log, numRecords = 10, leaderEpoch = 4)
+
+    val controllerEpoch = 0
+    val leaderEpoch = 5
+    val remoteBrokerId = brokerId + 1
+    val replicas = List(brokerId, remoteBrokerId)
+    val isr = List[Integer](brokerId, remoteBrokerId).asJava
+
+    doNothing().when(delayedOperations).checkAndCompleteFetch()
+
+    val initializeTimeMs = time.milliseconds()
+    partition.createLogIfNotExists(isNew = false, isFutureReplica = false, offsetCheckpoints)
+    assertTrue(
+      "Expected become leader transition to succeed",
+      partition.makeLeader(
+        new LeaderAndIsrPartitionState()
+          .setControllerEpoch(controllerEpoch)
+          .setLeader(brokerId)
+          .setLeaderEpoch(leaderEpoch)
+          .setIsr(isr)
+          .setZkVersion(1)
+          .setReplicas(replicas.map(Int.box).asJava)
+          .setIsNew(true),
+        offsetCheckpoints)
+    )
+    assertEquals(Set(brokerId, remoteBrokerId), partition.inSyncReplicaIds)
+    assertEquals(0L, partition.localLogOrException.highWatermark)
+
+    val remoteReplica = partition.getReplica(remoteBrokerId).get
+    assertEquals(initializeTimeMs, remoteReplica.lastCaughtUpTimeMs)
+    assertEquals(LogOffsetMetadata.UnknownOffsetMetadata.messageOffset, remoteReplica.logEndOffset)
+    assertEquals(Log.UnknownOffset, remoteReplica.logStartOffset)
+
+    // There is a short delay before the first fetch. The follower is not yet caught up to the log end.
+    time.sleep(5000)
+    val firstFetchTimeMs = time.milliseconds()
+    partition.updateFollowerFetchState(remoteBrokerId,
+      followerFetchOffsetMetadata = LogOffsetMetadata(5),
+      followerStartOffset = 0L,
+      followerFetchTimeMs = firstFetchTimeMs,
+      leaderEndOffset = 10L,
+      lastSentHighwatermark = partition.localLogOrException.highWatermark)
+    assertEquals(initializeTimeMs, remoteReplica.lastCaughtUpTimeMs)
+    assertEquals(5L, partition.localLogOrException.highWatermark)
+    assertEquals(5L, remoteReplica.logEndOffset)
+    assertEquals(0L, remoteReplica.logStartOffset)
+
+    // Some new data is appended, but the follower catches up to the old end offset.
+    // The total elapsed time from initialization is larger than the max allowed replica lag.
+    time.sleep(5001)
+    seedLogData(log, numRecords = 5, leaderEpoch = leaderEpoch)
+    partition.updateFollowerFetchState(remoteBrokerId,
+      followerFetchOffsetMetadata = LogOffsetMetadata(10),
+      followerStartOffset = 0L,
+      followerFetchTimeMs = time.milliseconds(),
+      leaderEndOffset = 15L,
+      lastSentHighwatermark = partition.localLogOrException.highWatermark)
+    assertEquals(firstFetchTimeMs, remoteReplica.lastCaughtUpTimeMs)
+    assertEquals(10L, partition.localLogOrException.highWatermark)
+    assertEquals(10L, remoteReplica.logEndOffset)
+    assertEquals(0L, remoteReplica.logStartOffset)
+
+    // The ISR should not be shrunk because the follower has caught up with the leader at the
+    // time of the first fetch.
+    partition.maybeShrinkIsr(10000)
+    assertEquals(Set(brokerId, remoteBrokerId), partition.inSyncReplicaIds)
+  }
+
+  @Test
+  def testShouldNotShrinkIsrIfFollowerCaughtUpToLogEnd(): Unit = {
+    val log = logManager.getOrCreateLog(topicPartition, () => logConfig)
+    seedLogData(log, numRecords = 10, leaderEpoch = 4)
+
+    val controllerEpoch = 0
+    val leaderEpoch = 5
+    val remoteBrokerId = brokerId + 1
+    val replicas = List(brokerId, remoteBrokerId)
+    val isr = List[Integer](brokerId, remoteBrokerId).asJava
+
+    doNothing().when(delayedOperations).checkAndCompleteFetch()
+
+    val initializeTimeMs = time.milliseconds()
+    partition.createLogIfNotExists(isNew = false, isFutureReplica = false, offsetCheckpoints)
+    assertTrue(
+      "Expected become leader transition to succeed",
+      partition.makeLeader(
+        new LeaderAndIsrPartitionState()
+          .setControllerEpoch(controllerEpoch)
+          .setLeader(brokerId)
+          .setLeaderEpoch(leaderEpoch)
+          .setIsr(isr)
+          .setZkVersion(1)
+          .setReplicas(replicas.map(Int.box).asJava)
+          .setIsNew(true),
+        offsetCheckpoints)
+    )
+    assertEquals(Set(brokerId, remoteBrokerId), partition.inSyncReplicaIds)
+    assertEquals(0L, partition.localLogOrException.highWatermark)
+
+    val remoteReplica = partition.getReplica(remoteBrokerId).get
+    assertEquals(initializeTimeMs, remoteReplica.lastCaughtUpTimeMs)
+    assertEquals(LogOffsetMetadata.UnknownOffsetMetadata.messageOffset, remoteReplica.logEndOffset)
+    assertEquals(Log.UnknownOffset, remoteReplica.logStartOffset)
+
+    // The follower catches up to the log end immediately.
+    partition.updateFollowerFetchState(remoteBrokerId,
+      followerFetchOffsetMetadata = LogOffsetMetadata(10),
+      followerStartOffset = 0L,
+      followerFetchTimeMs = time.milliseconds(),
+      leaderEndOffset = 10L,
+      lastSentHighwatermark = partition.localLogOrException.highWatermark)
+    assertEquals(initializeTimeMs, remoteReplica.lastCaughtUpTimeMs)
+    assertEquals(10L, partition.localLogOrException.highWatermark)
+    assertEquals(10L, remoteReplica.logEndOffset)
+    assertEquals(0L, remoteReplica.logStartOffset)
+
+    // Sleep longer than the max allowed follower lag
+    time.sleep(10001)
+
+    // The ISR should not be shrunk because the follower is caught up to the leader's log end
+    partition.maybeShrinkIsr(10000)
+    assertEquals(Set(brokerId, remoteBrokerId), partition.inSyncReplicaIds)
+  }
+
+  @Test
+  def testIsrNotShrunkIfUpdateFails(): Unit = {
+    val log = logManager.getOrCreateLog(topicPartition, () => logConfig)
+    seedLogData(log, numRecords = 10, leaderEpoch = 4)
+
+    val controllerEpoch = 0
+    val leaderEpoch = 5
+    val remoteBrokerId = brokerId + 1
+    val replicas = List[Integer](brokerId, remoteBrokerId).asJava
+    val isr = List[Integer](brokerId, remoteBrokerId).asJava
+
+    doNothing().when(delayedOperations).checkAndCompleteFetch()
+
+    val initializeTimeMs = time.milliseconds()
+    partition.createLogIfNotExists(isNew = false, isFutureReplica = false, offsetCheckpoints)
+    assertTrue("Expected become leader transition to succeed",
+      partition.makeLeader(
+        new LeaderAndIsrPartitionState()
+          .setControllerEpoch(controllerEpoch)
+          .setLeader(brokerId)
+          .setLeaderEpoch(leaderEpoch)
+          .setIsr(isr)
+          .setZkVersion(1)
+          .setReplicas(replicas)
+          .setIsNew(true),
+        offsetCheckpoints))
+    assertEquals(Set(brokerId, remoteBrokerId), partition.inSyncReplicaIds)
+    assertEquals(0L, partition.localLogOrException.highWatermark)
+
+    val remoteReplica = partition.getReplica(remoteBrokerId).get
+    assertEquals(initializeTimeMs, remoteReplica.lastCaughtUpTimeMs)
+    assertEquals(LogOffsetMetadata.UnknownOffsetMetadata.messageOffset, remoteReplica.logEndOffset)
+    assertEquals(Log.UnknownOffset, remoteReplica.logStartOffset)
+
+    time.sleep(10001)
+
+    // Mock the expected ISR update failure
+    val updatedLeaderAndIsr = LeaderAndIsr(
+      leader = brokerId,
+      leaderEpoch = leaderEpoch,
+      isr = List(brokerId),
+      zkVersion = 1)
+    when(stateStore.shrinkIsr(controllerEpoch, updatedLeaderAndIsr)).thenReturn(None)
+
+    partition.maybeShrinkIsr(10000)
+    assertEquals(Set(brokerId, remoteBrokerId), partition.inSyncReplicaIds)
+    assertEquals(0L, partition.localLogOrException.highWatermark)
+  }
+
+  @Test
+  def testUseCheckpointToInitializeHighWatermark(): Unit = {
+    val log = logManager.getOrCreateLog(topicPartition, () => logConfig)
+    seedLogData(log, numRecords = 6, leaderEpoch = 5)
+
+    when(offsetCheckpoints.fetch(logDir1.getAbsolutePath, topicPartition))
+      .thenReturn(Some(4L))
+
+    val controllerEpoch = 3
+    val replicas = List[Integer](brokerId, brokerId + 1).asJava
+    val leaderState = new LeaderAndIsrPartitionState()
+      .setControllerEpoch(controllerEpoch)
+      .setLeader(brokerId)
+      .setLeaderEpoch(6)
+      .setIsr(replicas)
+      .setZkVersion(1)
+      .setReplicas(replicas)
+      .setIsNew(false)
+    partition.makeLeader(leaderState, offsetCheckpoints)
+    assertEquals(4, partition.localLogOrException.highWatermark)
+  }
+
+  @Test
+  def testAddAndRemoveMetrics(): Unit = {
+    val metricsToCheck = List(
+      "UnderReplicated",
+      "UnderMinIsr",
+      "InSyncReplicasCount",
+      "ReplicasCount",
+      "LastStableOffsetLag",
+      "AtMinIsr")
+
+    def getMetric(metric: String): Option[Metric] = {
+      Metrics.defaultRegistry().allMetrics().asScala.filterKeys { metricName =>
+        metricName.getName == metric && metricName.getType == "Partition"
+      }.headOption.map(_._2)
+    }
+
+    assertTrue(metricsToCheck.forall(getMetric(_).isDefined))
+
+    Partition.removeMetrics(topicPartition)
+
+    assertEquals(Set(), Metrics.defaultRegistry().allMetrics().asScala.keySet.filter(_.getType == "Partition"))
+  }
+
+  /**
+   * Test when log is getting initialized, its config remains untouched after initialization is done.
+   */
+  @Test
+  def testLogConfigNotDirty(): Unit = {
+    val spyLogManager = spy(logManager)
+    val partition = new Partition(topicPartition,
+      replicaLagTimeMaxMs = Defaults.ReplicaLagTimeMaxMs,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      localBrokerId = brokerId,
+      time,
+      stateStore,
+      delayedOperations,
+      metadataCache,
+      spyLogManager)
+
+    partition.createLog(isNew = true, isFutureReplica = false, offsetCheckpoints)
+
+    // Validate that initializingLog and finishedInitializingLog was called
+    verify(spyLogManager).initializingLog(ArgumentMatchers.eq(topicPartition))
+    verify(spyLogManager).finishedInitializingLog(ArgumentMatchers.eq(topicPartition),
+      ArgumentMatchers.any(),
+      ArgumentMatchers.any()) // This doesn't get evaluated, but needed to satisfy compilation
+
+    // We should get config from ZK only once
+    verify(stateStore).fetchTopicConfig()
+  }
+
+  /**
+   * Test when log is getting initialized, its config remains gets reloaded if Topic config gets changed
+   * before initialization is done.
+   */
+  @Test
+  def testLogConfigDirtyAsTopicUpdated(): Unit = {
+    val spyLogManager = spy(logManager)
+    doAnswer(new Answer[Unit] {
+      def answer(invocation: InvocationOnMock): Unit = {
+        logManager.initializingLog(topicPartition)
+        logManager.topicConfigUpdated(topicPartition.topic())
+      }
+    }).when(spyLogManager).initializingLog(ArgumentMatchers.eq(topicPartition))
+
+    val partition = new Partition(topicPartition,
+      replicaLagTimeMaxMs = Defaults.ReplicaLagTimeMaxMs,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      localBrokerId = brokerId,
+      time,
+      stateStore,
+      delayedOperations,
+      metadataCache,
+      spyLogManager)
+
+    partition.createLog(isNew = true, isFutureReplica = false, offsetCheckpoints)
+
+    // Validate that initializingLog and finishedInitializingLog was called
+    verify(spyLogManager).initializingLog(ArgumentMatchers.eq(topicPartition))
+    verify(spyLogManager).finishedInitializingLog(ArgumentMatchers.eq(topicPartition),
+      ArgumentMatchers.any(),
+      ArgumentMatchers.any()) // This doesn't get evaluated, but needed to satisfy compilation
+
+    // We should get config from ZK twice, once before log is created, and second time once
+    // we find log config is dirty and refresh it.
+    verify(stateStore, times(2)).fetchTopicConfig()
+  }
+
+  /**
+   * Test when log is getting initialized, its config remains gets reloaded if Broker config gets changed
+   * before initialization is done.
+   */
+  @Test
+  def testLogConfigDirtyAsBrokerUpdated(): Unit = {
+    val spyLogManager = spy(logManager)
+    doAnswer(new Answer[Unit] {
+      def answer(invocation: InvocationOnMock): Unit = {
+        logManager.initializingLog(topicPartition)
+        logManager.brokerConfigUpdated()
+      }
+    }).when(spyLogManager).initializingLog(ArgumentMatchers.eq(topicPartition))
+
+    val partition = new Partition(topicPartition,
+      replicaLagTimeMaxMs = Defaults.ReplicaLagTimeMaxMs,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      localBrokerId = brokerId,
+      time,
+      stateStore,
+      delayedOperations,
+      metadataCache,
+      spyLogManager)
+
+    partition.createLog(isNew = true, isFutureReplica = false, offsetCheckpoints)
+
+    // Validate that initializingLog and finishedInitializingLog was called
+    verify(spyLogManager).initializingLog(ArgumentMatchers.eq(topicPartition))
+    verify(spyLogManager).finishedInitializingLog(ArgumentMatchers.eq(topicPartition),
+      ArgumentMatchers.any(),
+      ArgumentMatchers.any()) // This doesn't get evaluated, but needed to satisfy compilation
+
+    // We should get config from ZK twice, once before log is created, and second time once
+    // we find log config is dirty and refresh it.
+    verify(stateStore, times(2)).fetchTopicConfig()
+  }
+
+  private def seedLogData(log: Log, numRecords: Int, leaderEpoch: Int): Unit = {
+    for (i <- 0 until numRecords) {
+      val records = MemoryRecords.withRecords(0L, CompressionType.NONE, leaderEpoch,
+        new SimpleRecord(s"k$i".getBytes, s"v$i".getBytes))
+      log.appendAsLeader(records, leaderEpoch)
+    }
+  }
 }
diff --git a/core/src/test/scala/unit/kafka/cluster/ReplicaTest.scala b/core/src/test/scala/unit/kafka/cluster/ReplicaTest.scala
index d63901a3c3462..d6f10946fbf62 100644
--- a/core/src/test/scala/unit/kafka/cluster/ReplicaTest.scala
+++ b/core/src/test/scala/unit/kafka/cluster/ReplicaTest.scala
@@ -19,9 +19,8 @@ package kafka.cluster
 import java.util.Properties
 
 import kafka.log.{Log, LogConfig, LogManager}
-import kafka.server.{BrokerTopicStats, LogDirFailureChannel, LogOffsetMetadata}
+import kafka.server.{BrokerTopicStats, LogDirFailureChannel}
 import kafka.utils.{MockTime, TestUtils}
-import org.apache.kafka.common.TopicPartition
 import org.apache.kafka.common.errors.OffsetOutOfRangeException
 import org.apache.kafka.common.utils.Utils
 import org.junit.{After, Before, Test}
@@ -34,7 +33,6 @@ class ReplicaTest {
   val time = new MockTime()
   val brokerTopicStats = new BrokerTopicStats
   var log: Log = _
-  var replica: Replica = _
 
   @Before
   def setup(): Unit = {
@@ -53,11 +51,6 @@ class ReplicaTest {
       maxProducerIdExpirationMs = 60 * 60 * 1000,
       producerIdExpirationCheckIntervalMs = LogManager.ProducerIdExpirationCheckIntervalMs,
       logDirFailureChannel = new LogDirFailureChannel(10))
-
-    replica = new Replica(brokerId = 0,
-      topicPartition = new TopicPartition("foo", 0),
-      time = time,
-      log = Some(log))
   }
 
   @After
@@ -69,25 +62,19 @@ class ReplicaTest {
 
   @Test
   def testSegmentDeletionWithHighWatermarkInitialization(): Unit = {
-    val initialHighWatermark = 25L
-    replica = new Replica(brokerId = 0,
-      topicPartition = new TopicPartition("foo", 0),
-      time = time,
-      initialHighWatermarkValue = initialHighWatermark,
-      log = Some(log))
-
-    assertEquals(initialHighWatermark, replica.highWatermark.messageOffset)
-
     val expiredTimestamp = time.milliseconds() - 1000
     for (i <- 0 until 100) {
       val records = TestUtils.singletonRecords(value = s"test$i".getBytes, timestamp = expiredTimestamp)
       log.appendAsLeader(records, leaderEpoch = 0)
     }
 
+    val initialHighWatermark = log.updateHighWatermark(25L)
+    assertEquals(25L, initialHighWatermark)
+
     val initialNumSegments = log.numberOfSegments
     log.deleteOldSegments()
     assertTrue(log.numberOfSegments < initialNumSegments)
-    assertTrue(replica.logStartOffset <= initialHighWatermark)
+    assertTrue(log.logStartOffset <= initialHighWatermark)
   }
 
   @Test
@@ -100,25 +87,25 @@ class ReplicaTest {
 
     // ensure we have at least a few segments so the test case is not trivial
     assertTrue(log.numberOfSegments > 5)
-    assertEquals(0L, replica.highWatermark.messageOffset)
-    assertEquals(0L, replica.logStartOffset)
-    assertEquals(100L, replica.logEndOffset)
+    assertEquals(0L, log.highWatermark)
+    assertEquals(0L, log.logStartOffset)
+    assertEquals(100L, log.logEndOffset)
 
     for (hw <- 0 to 100) {
-      replica.highWatermark = new LogOffsetMetadata(hw)
-      assertEquals(hw, replica.highWatermark.messageOffset)
+      log.updateHighWatermark(hw)
+      assertEquals(hw, log.highWatermark)
       log.deleteOldSegments()
-      assertTrue(replica.logStartOffset <= hw)
+      assertTrue(log.logStartOffset <= hw)
 
       // verify that all segments up to the high watermark have been deleted
 
       log.logSegments.headOption.foreach { segment =>
         assertTrue(segment.baseOffset <= hw)
-        assertTrue(segment.baseOffset >= replica.logStartOffset)
+        assertTrue(segment.baseOffset >= log.logStartOffset)
       }
       log.logSegments.tail.foreach { segment =>
         assertTrue(segment.baseOffset > hw)
-        assertTrue(segment.baseOffset >= replica.logStartOffset)
+        assertTrue(segment.baseOffset >= log.logStartOffset)
       }
     }
 
@@ -134,8 +121,7 @@ class ReplicaTest {
       log.appendAsLeader(records, leaderEpoch = 0)
     }
 
-    replica.highWatermark = new LogOffsetMetadata(25L)
-    replica.maybeIncrementLogStartOffset(26L)
+    log.updateHighWatermark(25L)
+    log.maybeIncrementLogStartOffset(26L)
   }
-
 }
diff --git a/core/src/test/scala/unit/kafka/common/ZkNodeChangeNotificationListenerTest.scala b/core/src/test/scala/unit/kafka/common/ZkNodeChangeNotificationListenerTest.scala
index 0462300b248ac..05b56aae64232 100644
--- a/core/src/test/scala/unit/kafka/common/ZkNodeChangeNotificationListenerTest.scala
+++ b/core/src/test/scala/unit/kafka/common/ZkNodeChangeNotificationListenerTest.scala
@@ -23,6 +23,7 @@ import org.apache.kafka.common.resource.PatternType.LITERAL
 import org.junit.{After, Before, Test}
 
 import scala.collection.mutable.ArrayBuffer
+import scala.collection.Seq
 
 class ZkNodeChangeNotificationListenerTest extends ZooKeeperTestHarness {
 
@@ -46,7 +47,7 @@ class ZkNodeChangeNotificationListenerTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  def testProcessNotification() {
+  def testProcessNotification(): Unit = {
     val notificationMessage1 = Resource(Group, "messageA", LITERAL)
     val notificationMessage2 = Resource(Group, "messageB", LITERAL)
 
@@ -62,7 +63,7 @@ class ZkNodeChangeNotificationListenerTest extends ZooKeeperTestHarness {
      * There is no easy way to test purging. Even if we mock kafka time with MockTime, the purging compares kafka time
      * with the time stored in ZooKeeper stat and the embedded ZooKeeper server does not provide a way to mock time.
      * So to test purging we would have to use Time.SYSTEM.sleep(changeExpirationMs + 1) issue a write and check
-     * Assert.assertEquals(1, ZkUtils.getChildren(zkClient, seqNodeRoot).size). However even that the assertion
+     * Assert.assertEquals(1, KafkaZkClient.getChildren(seqNodeRoot).size). However even that the assertion
      * can fail as the second node can be deleted depending on how threads get scheduled.
      */
 
@@ -77,7 +78,7 @@ class ZkNodeChangeNotificationListenerTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  def testSwallowsProcessorException() : Unit = {
+  def testSwallowsProcessorException(): Unit = {
     notificationHandler.setThrowSize(2)
     notificationListener = new ZkNodeChangeNotificationListener(zkClient, LiteralAclChangeStore.aclChangePath,
       ZkAclChangeStore.SequenceNumberPrefix, notificationHandler, changeExpirationMs)
@@ -106,4 +107,4 @@ class ZkNodeChangeNotificationListenerTest extends ZooKeeperTestHarness {
 
     def setThrowSize(index: Int): Unit = throwSize = Option(index)
   }
-}
\ No newline at end of file
+}
diff --git a/core/src/test/scala/unit/kafka/controller/ControllerChannelManagerTest.scala b/core/src/test/scala/unit/kafka/controller/ControllerChannelManagerTest.scala
new file mode 100644
index 0000000000000..100014358ce66
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/controller/ControllerChannelManagerTest.scala
@@ -0,0 +1,753 @@
+/*
+ * 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.controller
+
+import java.util.Properties
+
+import kafka.api.{ApiVersion, KAFKA_0_10_0_IV1, KAFKA_0_10_2_IV0, KAFKA_0_9_0, KAFKA_1_0_IV0, KAFKA_2_2_IV0, KAFKA_2_3_IV2, KAFKA_2_4_IV0, KAFKA_2_4_IV1, LeaderAndIsr}
+import kafka.cluster.{Broker, EndPoint}
+import kafka.server.KafkaConfig
+import kafka.utils.TestUtils
+import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.message.{LeaderAndIsrResponseData, StopReplicaResponseData}
+import org.apache.kafka.common.message.LeaderAndIsrResponseData.LeaderAndIsrPartitionError
+import org.apache.kafka.common.message.StopReplicaResponseData.StopReplicaPartitionError
+import org.apache.kafka.common.network.ListenerName
+import org.apache.kafka.common.protocol.{ApiKeys, Errors}
+import org.apache.kafka.common.requests.{AbstractControlRequest, AbstractResponse, LeaderAndIsrRequest, LeaderAndIsrResponse, StopReplicaRequest, StopReplicaResponse, UpdateMetadataRequest}
+import org.apache.kafka.common.security.auth.SecurityProtocol
+import org.junit.Assert._
+import org.junit.Test
+import org.scalatest.Assertions
+
+import scala.collection.JavaConverters._
+import scala.collection.mutable
+import scala.collection.mutable.ListBuffer
+
+class ControllerChannelManagerTest {
+  private val controllerId = 1
+  private val controllerEpoch = 1
+  private val config = KafkaConfig.fromProps(TestUtils.createBrokerConfig(controllerId, "zkConnect"))
+  private val logger = new StateChangeLogger(controllerId, true, None)
+
+  type ControlRequest = AbstractControlRequest.Builder[_ <: AbstractControlRequest]
+
+  @Test
+  def testLeaderAndIsrRequestSent(): Unit = {
+    val context = initContext(Seq(1, 2, 3), Set("foo", "bar"), 2, 3)
+    val batch = new MockControllerBrokerRequestBatch(context)
+
+    val partitions = Map(
+      new TopicPartition("foo", 0) -> LeaderAndIsr(1, List(1, 2)),
+      new TopicPartition("foo", 1) -> LeaderAndIsr(2, List(2, 3)),
+      new TopicPartition("bar", 1) -> LeaderAndIsr(3, List(1, 3))
+    )
+
+    batch.newBatch()
+    partitions.foreach { case (partition, leaderAndIsr) =>
+      val leaderIsrAndControllerEpoch = LeaderIsrAndControllerEpoch(leaderAndIsr, controllerEpoch)
+      context.partitionLeadershipInfo.put(partition, leaderIsrAndControllerEpoch)
+      batch.addLeaderAndIsrRequestForBrokers(Seq(2), partition, leaderIsrAndControllerEpoch, replicaAssignment(Seq(1, 2, 3)), isNew = false)
+    }
+    batch.sendRequestsToBrokers(controllerEpoch)
+
+    val leaderAndIsrRequests = batch.collectLeaderAndIsrRequestsFor(2)
+    val updateMetadataRequests = batch.collectUpdateMetadataRequestsFor(2)
+    assertEquals(1, leaderAndIsrRequests.size)
+    assertEquals(1, updateMetadataRequests.size)
+
+    val leaderAndIsrRequest = leaderAndIsrRequests.head
+    assertEquals(controllerId, leaderAndIsrRequest.controllerId)
+    assertEquals(controllerEpoch, leaderAndIsrRequest.controllerEpoch)
+    assertEquals(partitions.keySet,
+      leaderAndIsrRequest.partitionStates.asScala.map(p => new TopicPartition(p.topicName, p.partitionIndex)).toSet)
+    assertEquals(partitions.map { case (k, v) => (k, v.leader) },
+      leaderAndIsrRequest.partitionStates.asScala.map(p => new TopicPartition(p.topicName, p.partitionIndex) -> p.leader).toMap)
+    assertEquals(partitions.map { case (k, v) => (k, v.isr) },
+      leaderAndIsrRequest.partitionStates.asScala.map(p => new TopicPartition(p.topicName, p.partitionIndex) -> p.isr.asScala).toMap)
+
+    applyLeaderAndIsrResponseCallbacks(Errors.NONE, batch.sentRequests(2).toList)
+    assertEquals(1, batch.sentEvents.size)
+
+    val LeaderAndIsrResponseReceived(response, brokerId) = batch.sentEvents.head
+    val leaderAndIsrResponse = response.asInstanceOf[LeaderAndIsrResponse]
+    assertEquals(2, brokerId)
+    assertEquals(partitions.keySet,
+      leaderAndIsrResponse.partitions.asScala.map(p => new TopicPartition(p.topicName, p.partitionIndex)).toSet)
+  }
+
+  @Test
+  def testLeaderAndIsrRequestIsNew(): Unit = {
+    val context = initContext(Seq(1, 2, 3), Set("foo", "bar"), 2, 3)
+    val batch = new MockControllerBrokerRequestBatch(context)
+
+    val partition = new TopicPartition("foo", 0)
+    val leaderAndIsr = LeaderAndIsr(1, List(1, 2))
+
+    val leaderIsrAndControllerEpoch = LeaderIsrAndControllerEpoch(leaderAndIsr, controllerEpoch)
+    context.partitionLeadershipInfo.put(partition, leaderIsrAndControllerEpoch)
+
+    batch.newBatch()
+    batch.addLeaderAndIsrRequestForBrokers(Seq(2), partition, leaderIsrAndControllerEpoch, replicaAssignment(Seq(1, 2, 3)), isNew = true)
+    batch.addLeaderAndIsrRequestForBrokers(Seq(2), partition, leaderIsrAndControllerEpoch, replicaAssignment(Seq(1, 2, 3)), isNew = false)
+    batch.sendRequestsToBrokers(controllerEpoch)
+
+    val leaderAndIsrRequests = batch.collectLeaderAndIsrRequestsFor(2)
+    val updateMetadataRequests = batch.collectUpdateMetadataRequestsFor(2)
+    assertEquals(1, leaderAndIsrRequests.size)
+    assertEquals(1, updateMetadataRequests.size)
+
+    val leaderAndIsrRequest = leaderAndIsrRequests.head
+    val partitionStates = leaderAndIsrRequest.partitionStates.asScala
+    assertEquals(Seq(partition), partitionStates.map(p => new TopicPartition(p.topicName, p.partitionIndex)))
+    val partitionState = partitionStates.find(p => p.topicName == partition.topic && p.partitionIndex == partition.partition)
+    assertEquals(Some(true), partitionState.map(_.isNew))
+  }
+
+  @Test
+  def testLeaderAndIsrRequestSentToLiveOrShuttingDownBrokers(): Unit = {
+    val context = initContext(Seq(1, 2, 3), Set("foo", "bar"), 2, 3)
+    val batch = new MockControllerBrokerRequestBatch(context)
+
+    // 2 is shutting down, 3 is dead
+    context.shuttingDownBrokerIds.add(2)
+    context.removeLiveBrokers(Set(3))
+
+    val partition = new TopicPartition("foo", 0)
+    val leaderAndIsr = LeaderAndIsr(1, List(1, 2))
+
+    val leaderIsrAndControllerEpoch = LeaderIsrAndControllerEpoch(leaderAndIsr, controllerEpoch)
+    context.partitionLeadershipInfo.put(partition, leaderIsrAndControllerEpoch)
+
+    batch.newBatch()
+    batch.addLeaderAndIsrRequestForBrokers(Seq(1, 2, 3), partition, leaderIsrAndControllerEpoch, replicaAssignment(Seq(1, 2, 3)), isNew = false)
+    batch.sendRequestsToBrokers(controllerEpoch)
+
+    assertEquals(0, batch.sentEvents.size)
+    assertEquals(2, batch.sentRequests.size)
+    assertEquals(Set(1, 2), batch.sentRequests.keySet)
+
+    for (brokerId <- Set(1, 2)) {
+      val leaderAndIsrRequests = batch.collectLeaderAndIsrRequestsFor(brokerId)
+      val updateMetadataRequests = batch.collectUpdateMetadataRequestsFor(brokerId)
+      assertEquals(1, leaderAndIsrRequests.size)
+      assertEquals(1, updateMetadataRequests.size)
+      val leaderAndIsrRequest = leaderAndIsrRequests.head
+      assertEquals(Seq(partition), leaderAndIsrRequest.partitionStates.asScala.map(p => new TopicPartition(p.topicName, p.partitionIndex)))
+    }
+  }
+
+  @Test
+  def testLeaderAndIsrInterBrokerProtocolVersion(): Unit = {
+    testLeaderAndIsrRequestFollowsInterBrokerProtocolVersion(ApiVersion.latestVersion, ApiKeys.LEADER_AND_ISR.latestVersion)
+
+    for (apiVersion <- ApiVersion.allVersions) {
+      val leaderAndIsrRequestVersion: Short =
+        if (apiVersion >= KAFKA_2_4_IV1) 5
+        else if (apiVersion >= KAFKA_2_4_IV0) 4
+        else if (apiVersion >= KAFKA_2_3_IV2) 3
+        else if (apiVersion >= KAFKA_2_2_IV0) 2
+        else if (apiVersion >= KAFKA_1_0_IV0) 1
+        else 0
+
+      testLeaderAndIsrRequestFollowsInterBrokerProtocolVersion(apiVersion, leaderAndIsrRequestVersion)
+    }
+  }
+
+  private def testLeaderAndIsrRequestFollowsInterBrokerProtocolVersion(interBrokerProtocolVersion: ApiVersion,
+                                                                       expectedLeaderAndIsrVersion: Short): Unit = {
+    val context = initContext(Seq(1, 2, 3), Set("foo", "bar"), 2, 3)
+    val config = createConfig(interBrokerProtocolVersion)
+    val batch = new MockControllerBrokerRequestBatch(context, config)
+
+    val partition = new TopicPartition("foo", 0)
+    val leaderAndIsr = LeaderAndIsr(1, List(1, 2))
+
+    val leaderIsrAndControllerEpoch = LeaderIsrAndControllerEpoch(leaderAndIsr, controllerEpoch)
+    context.partitionLeadershipInfo.put(partition, leaderIsrAndControllerEpoch)
+
+    batch.newBatch()
+    batch.addLeaderAndIsrRequestForBrokers(Seq(2), partition, leaderIsrAndControllerEpoch, replicaAssignment(Seq(1, 2, 3)), isNew = false)
+    batch.sendRequestsToBrokers(controllerEpoch)
+
+    val leaderAndIsrRequests = batch.collectLeaderAndIsrRequestsFor(2)
+    assertEquals(1, leaderAndIsrRequests.size)
+    assertEquals(s"IBP $interBrokerProtocolVersion should use version $expectedLeaderAndIsrVersion",
+      expectedLeaderAndIsrVersion, leaderAndIsrRequests.head.version)
+  }
+
+  @Test
+  def testUpdateMetadataRequestSent(): Unit = {
+    val context = initContext(Seq(1, 2, 3), Set("foo", "bar"), 2, 3)
+    val batch = new MockControllerBrokerRequestBatch(context)
+
+    val partitions = Map(
+      new TopicPartition("foo", 0) -> LeaderAndIsr(1, List(1, 2)),
+      new TopicPartition("foo", 1) -> LeaderAndIsr(2, List(2, 3)),
+      new TopicPartition("bar", 1) -> LeaderAndIsr(3, List(1, 3))
+    )
+
+    partitions.foreach { case (partition, leaderAndIsr) =>
+      context.partitionLeadershipInfo.put(partition, LeaderIsrAndControllerEpoch(leaderAndIsr, controllerEpoch))
+    }
+
+    batch.newBatch()
+    batch.addUpdateMetadataRequestForBrokers(Seq(2), partitions.keySet)
+    batch.sendRequestsToBrokers(controllerEpoch)
+
+    val updateMetadataRequests = batch.collectUpdateMetadataRequestsFor(2)
+    assertEquals(1, updateMetadataRequests.size)
+
+    val updateMetadataRequest = updateMetadataRequests.head
+    val partitionStates = updateMetadataRequest.partitionStates.asScala.toBuffer
+    assertEquals(3, partitionStates.size)
+    assertEquals(partitions.map { case (k, v) => (k, v.leader) },
+      partitionStates.map(ps => (new TopicPartition(ps.topicName, ps.partitionIndex), ps.leader)).toMap)
+    assertEquals(partitions.map { case (k, v) => (k, v.isr) },
+      partitionStates.map(ps => (new TopicPartition(ps.topicName, ps.partitionIndex), ps.isr.asScala)).toMap)
+
+    assertEquals(controllerId, updateMetadataRequest.controllerId)
+    assertEquals(controllerEpoch, updateMetadataRequest.controllerEpoch)
+    assertEquals(3, updateMetadataRequest.liveBrokers.size)
+    assertEquals(Set(1, 2, 3), updateMetadataRequest.liveBrokers.asScala.map(_.id).toSet)
+  }
+
+  @Test
+  def testUpdateMetadataDoesNotIncludePartitionsWithoutLeaderAndIsr(): Unit = {
+    val context = initContext(Seq(1, 2, 3), Set("foo", "bar"), 2, 3)
+    val batch = new MockControllerBrokerRequestBatch(context)
+
+    val partitions = Set(
+      new TopicPartition("foo", 0),
+      new TopicPartition("foo", 1),
+      new TopicPartition("bar", 1)
+    )
+
+    batch.newBatch()
+    batch.addUpdateMetadataRequestForBrokers(Seq(2), partitions)
+    batch.sendRequestsToBrokers(controllerEpoch)
+
+    assertEquals(0, batch.sentEvents.size)
+    assertEquals(1, batch.sentRequests.size)
+    assertTrue(batch.sentRequests.contains(2))
+
+    val updateMetadataRequests = batch.collectUpdateMetadataRequestsFor(2)
+    assertEquals(1, updateMetadataRequests.size)
+
+    val updateMetadataRequest = updateMetadataRequests.head
+    assertEquals(0, updateMetadataRequest.partitionStates.asScala.size)
+    assertEquals(3, updateMetadataRequest.liveBrokers.size)
+    assertEquals(Set(1, 2, 3), updateMetadataRequest.liveBrokers.asScala.map(_.id).toSet)
+  }
+
+  @Test
+  def testUpdateMetadataRequestDuringTopicDeletion(): Unit = {
+    val context = initContext(Seq(1, 2, 3), Set("foo", "bar"), 2, 3)
+    val batch = new MockControllerBrokerRequestBatch(context)
+
+    val partitions = Map(
+      new TopicPartition("foo", 0) -> LeaderAndIsr(1, List(1, 2)),
+      new TopicPartition("foo", 1) -> LeaderAndIsr(2, List(2, 3)),
+      new TopicPartition("bar", 1) -> LeaderAndIsr(3, List(1, 3))
+    )
+
+    partitions.foreach { case (partition, leaderAndIsr) =>
+      context.partitionLeadershipInfo.put(partition, LeaderIsrAndControllerEpoch(leaderAndIsr, controllerEpoch))
+    }
+
+    context.queueTopicDeletion(Set("foo"))
+
+    batch.newBatch()
+    batch.addUpdateMetadataRequestForBrokers(Seq(2), partitions.keySet)
+    batch.sendRequestsToBrokers(controllerEpoch)
+
+    val updateMetadataRequests = batch.collectUpdateMetadataRequestsFor(2)
+    assertEquals(1, updateMetadataRequests.size)
+
+    val updateMetadataRequest = updateMetadataRequests.head
+    assertEquals(3, updateMetadataRequest.partitionStates.asScala.size)
+
+    assertTrue(updateMetadataRequest.partitionStates.asScala
+      .filter(_.topicName == "foo")
+      .map(_.leader)
+      .forall(leaderId => leaderId == LeaderAndIsr.LeaderDuringDelete))
+
+    assertEquals(partitions.filter { case (k, _) => k.topic == "bar" }.map { case (k, v) => (k, v.leader) },
+      updateMetadataRequest.partitionStates.asScala.filter(ps => ps.topicName == "bar").map { ps =>
+        (new TopicPartition(ps.topicName, ps.partitionIndex), ps.leader) }.toMap)
+    assertEquals(partitions.map { case (k, v) => (k, v.isr) },
+      updateMetadataRequest.partitionStates.asScala.map(ps => (new TopicPartition(ps.topicName, ps.partitionIndex), ps.isr.asScala)).toMap)
+
+    assertEquals(3, updateMetadataRequest.liveBrokers.size)
+    assertEquals(Set(1, 2, 3), updateMetadataRequest.liveBrokers.asScala.map(_.id).toSet)
+  }
+
+  @Test
+  def testUpdateMetadataIncludesLiveOrShuttingDownBrokers(): Unit = {
+    val context = initContext(Seq(1, 2, 3), Set("foo", "bar"), 2, 3)
+    val batch = new MockControllerBrokerRequestBatch(context)
+
+    // 2 is shutting down, 3 is dead
+    context.shuttingDownBrokerIds.add(2)
+    context.removeLiveBrokers(Set(3))
+
+    batch.newBatch()
+    batch.addUpdateMetadataRequestForBrokers(Seq(1, 2, 3), Set.empty)
+    batch.sendRequestsToBrokers(controllerEpoch)
+
+    assertEquals(Set(1, 2), batch.sentRequests.keySet)
+
+    for (brokerId <- Set(1, 2)) {
+      val updateMetadataRequests = batch.collectUpdateMetadataRequestsFor(brokerId)
+      assertEquals(1, updateMetadataRequests.size)
+
+      val updateMetadataRequest = updateMetadataRequests.head
+      assertEquals(0, updateMetadataRequest.partitionStates.asScala.size)
+      assertEquals(2, updateMetadataRequest.liveBrokers.size)
+      assertEquals(Set(1, 2), updateMetadataRequest.liveBrokers.asScala.map(_.id).toSet)
+    }
+  }
+
+  @Test
+  def testUpdateMetadataInterBrokerProtocolVersion(): Unit = {
+    testUpdateMetadataFollowsInterBrokerProtocolVersion(ApiVersion.latestVersion, ApiKeys.UPDATE_METADATA.latestVersion)
+
+    for (apiVersion <- ApiVersion.allVersions) {
+      val updateMetadataRequestVersion: Short =
+        if (apiVersion >= KAFKA_2_4_IV1) 7
+        else if (apiVersion >= KAFKA_2_3_IV2) 6
+        else if (apiVersion >= KAFKA_2_2_IV0) 5
+        else if (apiVersion >= KAFKA_1_0_IV0) 4
+        else if (apiVersion >= KAFKA_0_10_2_IV0) 3
+        else if (apiVersion >= KAFKA_0_10_0_IV1) 2
+        else if (apiVersion >= KAFKA_0_9_0) 1
+        else 0
+
+      testUpdateMetadataFollowsInterBrokerProtocolVersion(apiVersion, updateMetadataRequestVersion)
+    }
+  }
+
+  private def testUpdateMetadataFollowsInterBrokerProtocolVersion(interBrokerProtocolVersion: ApiVersion,
+                                                          expectedUpdateMetadataVersion: Short): Unit = {
+    val context = initContext(Seq(1, 2, 3), Set("foo", "bar"), 2, 3)
+    val config = createConfig(interBrokerProtocolVersion)
+    val batch = new MockControllerBrokerRequestBatch(context, config)
+
+    batch.newBatch()
+    batch.addUpdateMetadataRequestForBrokers(Seq(2), Set.empty)
+    batch.sendRequestsToBrokers(controllerEpoch)
+
+    assertEquals(0, batch.sentEvents.size)
+    assertEquals(1, batch.sentRequests.size)
+    assertTrue(batch.sentRequests.contains(2))
+
+    val requests = batch.collectUpdateMetadataRequestsFor(2)
+    val allVersions = requests.map(_.version)
+    assertTrue(s"IBP $interBrokerProtocolVersion should use version $expectedUpdateMetadataVersion, " +
+      s"but found versions $allVersions",
+      allVersions.forall(_ == expectedUpdateMetadataVersion))
+  }
+
+  @Test
+  def testStopReplicaRequestSent(): Unit = {
+    val context = initContext(Seq(1, 2, 3), Set("foo", "bar"), 2, 3)
+    val batch = new MockControllerBrokerRequestBatch(context)
+
+    val partitions = Set(
+      new TopicPartition("foo", 0),
+      new TopicPartition("foo", 1),
+      new TopicPartition("bar", 1)
+    )
+
+    batch.newBatch()
+    partitions.foreach { partition =>
+      batch.addStopReplicaRequestForBrokers(Seq(2), partition, deletePartition = false)
+    }
+    batch.sendRequestsToBrokers(controllerEpoch)
+
+    assertEquals(0, batch.sentEvents.size)
+    assertEquals(1, batch.sentRequests.size)
+    assertTrue(batch.sentRequests.contains(2))
+
+    val sentRequests = batch.sentRequests(2)
+    assertEquals(1, sentRequests.size)
+
+    val sentStopReplicaRequests = batch.collectStopReplicaRequestsFor(2)
+    assertEquals(1, sentStopReplicaRequests.size)
+
+    val stopReplicaRequest = sentStopReplicaRequests.head
+    assertFalse(stopReplicaRequest.deletePartitions())
+    assertEquals(partitions, stopReplicaRequest.partitions.asScala.toSet)
+
+    applyStopReplicaResponseCallbacks(Errors.NONE, batch.sentRequests(2).toList)
+    assertEquals(0, batch.sentEvents.size)
+  }
+
+  @Test
+  def testStopReplicaRequestsWhileTopicQueuedForDeletion(): Unit = {
+    val context = initContext(Seq(1, 2, 3), Set("foo", "bar"), 2, 3)
+    val batch = new MockControllerBrokerRequestBatch(context)
+
+    val partitions = Set(
+      new TopicPartition("foo", 0),
+      new TopicPartition("foo", 1),
+      new TopicPartition("bar", 1)
+    )
+
+    // Topic deletion is queued, but has not begun
+    context.queueTopicDeletion(Set("foo"))
+
+    batch.newBatch()
+    partitions.foreach { partition =>
+      batch.addStopReplicaRequestForBrokers(Seq(2), partition, deletePartition = true)
+    }
+    batch.sendRequestsToBrokers(controllerEpoch)
+
+    assertEquals(0, batch.sentEvents.size)
+    assertEquals(1, batch.sentRequests.size)
+    assertTrue(batch.sentRequests.contains(2))
+
+    val sentRequests = batch.sentRequests(2)
+    assertEquals(1, sentRequests.size)
+
+    val sentStopReplicaRequests = batch.collectStopReplicaRequestsFor(2)
+    assertEquals(1, sentStopReplicaRequests.size)
+    assertEquals(partitions, sentStopReplicaRequests.flatMap(_.partitions.asScala).toSet)
+    assertTrue(sentStopReplicaRequests.forall(_.deletePartitions))
+
+    // No events will be sent after the response returns
+    applyStopReplicaResponseCallbacks(Errors.NONE, batch.sentRequests(2).toList)
+    assertEquals(0, batch.sentEvents.size)
+  }
+
+  @Test
+  def testStopReplicaRequestsWhileTopicDeletionStarted(): Unit = {
+    val context = initContext(Seq(1, 2, 3), Set("foo", "bar"), 2, 3)
+    val batch = new MockControllerBrokerRequestBatch(context)
+
+    val partitions = Set(
+      new TopicPartition("foo", 0),
+      new TopicPartition("foo", 1),
+      new TopicPartition("bar", 1)
+    )
+
+    context.queueTopicDeletion(Set("foo"))
+    context.beginTopicDeletion(Set("foo"))
+
+    batch.newBatch()
+    partitions.foreach { partition =>
+      batch.addStopReplicaRequestForBrokers(Seq(2), partition, deletePartition = true)
+    }
+    batch.sendRequestsToBrokers(controllerEpoch)
+
+    assertEquals(0, batch.sentEvents.size)
+    assertEquals(1, batch.sentRequests.size)
+    assertTrue(batch.sentRequests.contains(2))
+
+    val sentRequests = batch.sentRequests(2)
+    assertEquals(1, sentRequests.size)
+
+    val sentStopReplicaRequests = batch.collectStopReplicaRequestsFor(2)
+    assertEquals(1, sentStopReplicaRequests.size)
+    assertEquals(partitions, sentStopReplicaRequests.flatMap(_.partitions.asScala).toSet)
+    assertTrue(sentStopReplicaRequests.forall(_.deletePartitions()))
+
+    // When the topic is being deleted, we should provide a callback which sends
+    // the received event for the StopReplica response
+    applyStopReplicaResponseCallbacks(Errors.NONE, batch.sentRequests(2).toList)
+    assertEquals(1, batch.sentEvents.size)
+
+    // We should only receive events for the topic being deleted
+    val includedPartitions = batch.sentEvents.flatMap {
+      case event: TopicDeletionStopReplicaResponseReceived => event.partitionErrors.keySet
+      case otherEvent => Assertions.fail(s"Unexpected sent event: $otherEvent")
+    }.toSet
+    assertEquals(partitions.filter(_.topic == "foo"), includedPartitions)
+  }
+
+  @Test
+  def testMixedDeleteAndNotDeleteStopReplicaRequests(): Unit = {
+    val context = initContext(Seq(1, 2, 3), Set("foo", "bar"), 2, 3)
+    val batch = new MockControllerBrokerRequestBatch(context)
+
+    val deletePartitions = Set(
+      new TopicPartition("foo", 0),
+      new TopicPartition("foo", 1)
+    )
+
+    val nonDeletePartitions = Set(
+      new TopicPartition("bar", 0),
+      new TopicPartition("bar", 1)
+    )
+
+    batch.newBatch()
+    deletePartitions.foreach { partition =>
+      batch.addStopReplicaRequestForBrokers(Seq(2), partition, deletePartition = true)
+    }
+    nonDeletePartitions.foreach { partition =>
+      batch.addStopReplicaRequestForBrokers(Seq(2), partition, deletePartition = false)
+    }
+    batch.sendRequestsToBrokers(controllerEpoch)
+
+    assertEquals(0, batch.sentEvents.size)
+    assertEquals(1, batch.sentRequests.size)
+    assertTrue(batch.sentRequests.contains(2))
+
+    val sentRequests = batch.sentRequests(2)
+    assertEquals(2, sentRequests.size)
+
+    val sentStopReplicaRequests = batch.collectStopReplicaRequestsFor(2)
+    assertEquals(2, sentStopReplicaRequests.size)
+
+    val (deleteRequests, nonDeleteRequests) = sentStopReplicaRequests.partition(_.deletePartitions())
+    assertEquals(1, deleteRequests.size)
+    assertEquals(deletePartitions, deleteRequests.head.partitions.asScala.toSet)
+    assertEquals(1, nonDeleteRequests.size)
+    assertEquals(nonDeletePartitions, nonDeleteRequests.head.partitions.asScala.toSet)
+  }
+
+  @Test
+  def testStopReplicaGroupsByBroker(): Unit = {
+    val context = initContext(Seq(1, 2, 3), Set("foo", "bar"), 2, 3)
+    val batch = new MockControllerBrokerRequestBatch(context)
+
+    val partitions = Set(
+      new TopicPartition("foo", 0),
+      new TopicPartition("foo", 1),
+      new TopicPartition("bar", 1)
+    )
+
+    batch.newBatch()
+    partitions.foreach { partition =>
+      batch.addStopReplicaRequestForBrokers(Seq(2, 3), partition, deletePartition = false)
+    }
+    batch.sendRequestsToBrokers(controllerEpoch)
+
+    assertEquals(0, batch.sentEvents.size)
+    assertEquals(2, batch.sentRequests.size)
+    assertTrue(batch.sentRequests.contains(2))
+    assertTrue(batch.sentRequests.contains(3))
+
+    val sentRequests = batch.sentRequests(2)
+    assertEquals(1, sentRequests.size)
+
+    for (brokerId <- Set(2, 3)) {
+      val sentStopReplicaRequests = batch.collectStopReplicaRequestsFor(brokerId)
+      assertEquals(1, sentStopReplicaRequests.size)
+
+      val stopReplicaRequest = sentStopReplicaRequests.head
+      assertFalse(stopReplicaRequest.deletePartitions())
+      assertEquals(partitions, stopReplicaRequest.partitions.asScala.toSet)
+
+      applyStopReplicaResponseCallbacks(Errors.NONE, batch.sentRequests(2).toList)
+      assertEquals(0, batch.sentEvents.size)
+    }
+  }
+
+  @Test
+  def testStopReplicaSentOnlyToLiveAndShuttingDownBrokers(): Unit = {
+    val context = initContext(Seq(1, 2, 3), Set("foo", "bar"), 2, 3)
+    val batch = new MockControllerBrokerRequestBatch(context)
+
+    // 2 is shutting down, 3 is dead
+    context.shuttingDownBrokerIds.add(2)
+    context.removeLiveBrokers(Set(3))
+
+    val partitions = Set(
+      new TopicPartition("foo", 0),
+      new TopicPartition("foo", 1),
+      new TopicPartition("bar", 1)
+    )
+
+    batch.newBatch()
+    partitions.foreach { partition =>
+      batch.addStopReplicaRequestForBrokers(Seq(2, 3), partition, deletePartition = false)
+    }
+    batch.sendRequestsToBrokers(controllerEpoch)
+
+    assertEquals(0, batch.sentEvents.size)
+    assertEquals(1, batch.sentRequests.size)
+    assertTrue(batch.sentRequests.contains(2))
+
+    val sentRequests = batch.sentRequests(2)
+    assertEquals(1, sentRequests.size)
+
+    val sentStopReplicaRequests = batch.collectStopReplicaRequestsFor(2)
+    assertEquals(1, sentStopReplicaRequests.size)
+
+    val stopReplicaRequest = sentStopReplicaRequests.head
+    assertFalse(stopReplicaRequest.deletePartitions())
+    assertEquals(partitions, stopReplicaRequest.partitions.asScala.toSet)
+  }
+
+  @Test
+  def testStopReplicaInterBrokerProtocolVersion(): Unit = {
+    testStopReplicaFollowsInterBrokerProtocolVersion(ApiVersion.latestVersion, ApiKeys.STOP_REPLICA.latestVersion)
+
+    for (apiVersion <- ApiVersion.allVersions) {
+      if (apiVersion < KAFKA_2_2_IV0)
+        testStopReplicaFollowsInterBrokerProtocolVersion(apiVersion, 0.toShort)
+      else if (apiVersion < KAFKA_2_3_IV2)
+        testStopReplicaFollowsInterBrokerProtocolVersion(apiVersion, 1.toShort)
+      else if (apiVersion < KAFKA_2_4_IV1)
+        testStopReplicaFollowsInterBrokerProtocolVersion(apiVersion, 2.toShort)
+      else
+        testStopReplicaFollowsInterBrokerProtocolVersion(apiVersion, 3.toShort)
+    }
+  }
+
+  private def testStopReplicaFollowsInterBrokerProtocolVersion(interBrokerProtocolVersion: ApiVersion,
+                                                               expectedStopReplicaRequestVersion: Short): Unit = {
+    val context = initContext(Seq(1, 2, 3), Set("foo"), 2, 3)
+    val config = createConfig(interBrokerProtocolVersion)
+    val batch = new MockControllerBrokerRequestBatch(context, config)
+
+    val partition = new TopicPartition("foo", 0)
+
+    batch.newBatch()
+    batch.addStopReplicaRequestForBrokers(Seq(2), partition, deletePartition = false)
+    batch.sendRequestsToBrokers(controllerEpoch)
+
+    assertEquals(0, batch.sentEvents.size)
+    assertEquals(1, batch.sentRequests.size)
+    assertTrue(batch.sentRequests.contains(2))
+
+    val requests = batch.collectStopReplicaRequestsFor(2)
+    val allVersions = requests.map(_.version)
+    assertTrue(s"IBP $interBrokerProtocolVersion should use version $expectedStopReplicaRequestVersion, " +
+      s"but found versions $allVersions",
+      allVersions.forall(_ == expectedStopReplicaRequestVersion))
+  }
+
+  private def applyStopReplicaResponseCallbacks(error: Errors, sentRequests: List[SentRequest]): Unit = {
+    sentRequests.filter(_.responseCallback != null).foreach { sentRequest =>
+      val stopReplicaRequest = sentRequest.request.build().asInstanceOf[StopReplicaRequest]
+      val stopReplicaResponse =
+        if (error == Errors.NONE) {
+          val partitionErrors = stopReplicaRequest.partitions.asScala.map { tp =>
+            new StopReplicaPartitionError()
+              .setTopicName(tp.topic)
+              .setPartitionIndex(tp.partition)
+              .setErrorCode(error.code)
+          }.toBuffer.asJava
+          new StopReplicaResponse(new StopReplicaResponseData().setPartitionErrors(partitionErrors))
+        } else {
+          stopReplicaRequest.getErrorResponse(error.exception)
+        }
+      sentRequest.responseCallback.apply(stopReplicaResponse)
+    }
+  }
+
+  private def applyLeaderAndIsrResponseCallbacks(error: Errors, sentRequests: List[SentRequest]): Unit = {
+    sentRequests.filter(_.request.apiKey == ApiKeys.LEADER_AND_ISR).filter(_.responseCallback != null).foreach { sentRequest =>
+      val leaderAndIsrRequest = sentRequest.request.build().asInstanceOf[LeaderAndIsrRequest]
+      val partitionErrors = leaderAndIsrRequest.partitionStates.asScala.map(p =>
+        new LeaderAndIsrPartitionError()
+          .setTopicName(p.topicName)
+          .setPartitionIndex(p.partitionIndex)
+          .setErrorCode(error.code))
+      val leaderAndIsrResponse = new LeaderAndIsrResponse(
+        new LeaderAndIsrResponseData()
+          .setErrorCode(error.code)
+          .setPartitionErrors(partitionErrors.toBuffer.asJava))
+      sentRequest.responseCallback(leaderAndIsrResponse)
+    }
+  }
+
+  private def createConfig(interBrokerVersion: ApiVersion): KafkaConfig = {
+    val props = new Properties()
+    props.put(KafkaConfig.BrokerIdProp, controllerId.toString)
+    props.put(KafkaConfig.ZkConnectProp, "zkConnect")
+    props.put(KafkaConfig.InterBrokerProtocolVersionProp, interBrokerVersion.version)
+    props.put(KafkaConfig.LogMessageFormatVersionProp, interBrokerVersion.version)
+    KafkaConfig.fromProps(props)
+  }
+
+  private def replicaAssignment(replicas: Seq[Int]): ReplicaAssignment = ReplicaAssignment(replicas, Seq(), Seq())
+
+  private def initContext(brokers: Seq[Int],
+                          topics: Set[String],
+                          numPartitions: Int,
+                          replicationFactor: Int): ControllerContext = {
+    val context = new ControllerContext
+    val brokerEpochs = brokers.map { brokerId =>
+      val endpoint = new EndPoint("localhost", 9900 + brokerId, new ListenerName("PLAINTEXT"),
+        SecurityProtocol.PLAINTEXT)
+      Broker(brokerId, Seq(endpoint), rack = None) -> 1L
+    }.toMap
+
+    context.setLiveBrokerAndEpochs(brokerEpochs)
+
+    // Simple round-robin replica assignment
+    var leaderIndex = 0
+    for (topic <- topics; partitionId <- 0 until numPartitions) {
+      val partition = new TopicPartition(topic, partitionId)
+      val replicas = (0 until replicationFactor).map { i =>
+        val replica = brokers((i + leaderIndex) % brokers.size)
+        replica
+      }
+      context.updatePartitionReplicaAssignment(partition, replicas)
+      leaderIndex += 1
+    }
+    context
+  }
+
+  private case class SentRequest(request: ControlRequest, responseCallback: AbstractResponse => Unit)
+
+  private class MockControllerBrokerRequestBatch(context: ControllerContext, config: KafkaConfig = config)
+    extends AbstractControllerBrokerRequestBatch(config, context, logger) {
+
+    val sentEvents = ListBuffer.empty[ControllerEvent]
+    val sentRequests = mutable.Map.empty[Int, ListBuffer[SentRequest]]
+
+    override def sendEvent(event: ControllerEvent): Unit = {
+      sentEvents.append(event)
+    }
+    override def sendRequest(brokerId: Int, request: ControlRequest, callback: AbstractResponse => Unit): Unit = {
+      sentRequests.getOrElseUpdate(brokerId, ListBuffer.empty)
+      sentRequests(brokerId).append(SentRequest(request, callback))
+    }
+
+    def collectStopReplicaRequestsFor(brokerId: Int): List[StopReplicaRequest] = {
+      sentRequests.get(brokerId) match {
+        case Some(requests) => requests
+          .filter(_.request.apiKey == ApiKeys.STOP_REPLICA)
+          .map(_.request.build().asInstanceOf[StopReplicaRequest]).toList
+        case None => List.empty[StopReplicaRequest]
+      }
+    }
+
+    def collectUpdateMetadataRequestsFor(brokerId: Int): List[UpdateMetadataRequest] = {
+      sentRequests.get(brokerId) match {
+        case Some(requests) => requests
+          .filter(_.request.apiKey == ApiKeys.UPDATE_METADATA)
+          .map(_.request.build().asInstanceOf[UpdateMetadataRequest]).toList
+        case None => List.empty[UpdateMetadataRequest]
+      }
+    }
+
+    def collectLeaderAndIsrRequestsFor(brokerId: Int): List[LeaderAndIsrRequest] = {
+      sentRequests.get(brokerId) match {
+        case Some(requests) => requests
+          .filter(_.request.apiKey == ApiKeys.LEADER_AND_ISR)
+          .map(_.request.build().asInstanceOf[LeaderAndIsrRequest]).toList
+        case None => List.empty[LeaderAndIsrRequest]
+      }
+    }
+  }
+
+}
diff --git a/core/src/test/scala/unit/kafka/controller/ControllerContextTest.scala b/core/src/test/scala/unit/kafka/controller/ControllerContextTest.scala
new file mode 100644
index 0000000000000..a0121e26adfab
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/controller/ControllerContextTest.scala
@@ -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 unit.kafka.controller
+
+import kafka.cluster.{Broker, EndPoint}
+import kafka.controller.{ControllerContext, ReplicaAssignment}
+import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.network.ListenerName
+import org.apache.kafka.common.security.auth.SecurityProtocol
+import org.junit.{Before, Test}
+import org.junit.Assert.assertEquals
+import org.junit.Assert.assertTrue
+import org.junit.Assert.assertFalse
+
+
+class ControllerContextTest {
+
+  var context: ControllerContext = null
+  val brokers: Seq[Int] = Seq(1, 2, 3)
+  val tp1 = new TopicPartition("A", 0)
+  val tp2 = new TopicPartition("A", 1)
+  val tp3 = new TopicPartition("B", 0)
+
+  @Before
+  def setUp(): Unit = {
+    context = new ControllerContext
+
+    val brokerEpochs = Seq(1,2,3).map { brokerId =>
+      val endpoint = new EndPoint("localhost", 9900 + brokerId, new ListenerName("PLAINTEXT"),
+        SecurityProtocol.PLAINTEXT)
+      Broker(brokerId, Seq(endpoint), rack = None) -> 1L
+    }.toMap
+
+    context.setLiveBrokerAndEpochs(brokerEpochs)
+
+    // Simple round-robin replica assignment
+    var leaderIndex = 0
+    Seq(tp1, tp2, tp3).foreach {
+      partition =>
+        val replicas = brokers.indices.map { i =>
+          val replica = brokers((i + leaderIndex) % brokers.size)
+          replica
+        }
+        context.updatePartitionReplicaAssignment(partition, replicas)
+        leaderIndex += 1
+    }
+  }
+
+  @Test
+  def testUpdatePartitionReplicaAssignmentUpdatesReplicaAssignmentOnly(): Unit = {
+    val expectedReplicas = Seq(4)
+    context.updatePartitionReplicaAssignment(tp1, expectedReplicas)
+    val assignment = context.partitionReplicaAssignment(tp1)
+    val fullAssignment = context.partitionFullReplicaAssignment(tp1)
+
+    assertEquals(expectedReplicas, assignment)
+    assertEquals(expectedReplicas, fullAssignment.replicas)
+    assertEquals(Seq(), fullAssignment.addingReplicas)
+    assertEquals(Seq(), fullAssignment.removingReplicas)
+  }
+
+  @Test
+  def testUpdatePartitionReplicaAssignmentUpdatesReplicaAssignmentOnlyAndDoesNotOverwrite(): Unit = {
+    val expectedReplicas = Seq(4)
+    val expectedFullAssignment = ReplicaAssignment(Seq(3), Seq(1), Seq(2))
+    context.updatePartitionFullReplicaAssignment(tp1, expectedFullAssignment)
+
+    context.updatePartitionReplicaAssignment(tp1, expectedReplicas) // update only the replicas
+
+    val assignment = context.partitionReplicaAssignment(tp1)
+    val fullAssignment = context.partitionFullReplicaAssignment(tp1)
+    assertEquals(expectedReplicas, assignment)
+    assertEquals(expectedReplicas, fullAssignment.replicas)
+    // adding/removing replicas preserved
+    assertEquals(Seq(1), fullAssignment.addingReplicas)
+    assertEquals(Seq(2), fullAssignment.removingReplicas)
+  }
+
+  @Test
+  def testUpdatePartitionFullReplicaAssignmentUpdatesReplicaAssignment(): Unit = {
+    val initialReplicas = Seq(4)
+    context.updatePartitionReplicaAssignment(tp1, initialReplicas) // update only the replicas
+    val fullAssignment = context.partitionFullReplicaAssignment(tp1)
+    assertEquals(initialReplicas, fullAssignment.replicas)
+    assertEquals(Seq(), fullAssignment.addingReplicas)
+    assertEquals(Seq(), fullAssignment.removingReplicas)
+
+    val expectedFullAssignment = ReplicaAssignment(Seq(3), Seq(1), Seq(2))
+    context.updatePartitionFullReplicaAssignment(tp1, expectedFullAssignment)
+    val updatedFullAssignment = context.partitionFullReplicaAssignment(tp1)
+    assertEquals(expectedFullAssignment.replicas, updatedFullAssignment.replicas)
+    assertEquals(expectedFullAssignment.addingReplicas, updatedFullAssignment.addingReplicas)
+    assertEquals(expectedFullAssignment.removingReplicas, updatedFullAssignment.removingReplicas)
+  }
+
+  @Test
+  def testPartitionReplicaAssignmentReturnsEmptySeqIfTopicOrPartitionDoesNotExist(): Unit = {
+    val noTopicReplicas = context.partitionReplicaAssignment(new TopicPartition("NONEXISTENT", 0))
+    assertEquals(Seq.empty, noTopicReplicas)
+    val noPartitionReplicas = context.partitionReplicaAssignment(new TopicPartition("A", 100))
+    assertEquals(Seq.empty, noPartitionReplicas)
+  }
+
+  @Test
+  def testPartitionFullReplicaAssignmentReturnsEmptyAssignmentIfTopicOrPartitionDoesNotExist(): Unit = {
+    val expectedEmptyAssignment = ReplicaAssignment(Seq.empty, Seq.empty, Seq.empty)
+
+    val noTopicAssignment = context.partitionFullReplicaAssignment(new TopicPartition("NONEXISTENT", 0))
+    assertEquals(expectedEmptyAssignment, noTopicAssignment)
+    val noPartitionAssignment = context.partitionFullReplicaAssignment(new TopicPartition("A", 100))
+    assertEquals(expectedEmptyAssignment, noPartitionAssignment)
+  }
+
+  @Test
+  def testPartitionReplicaAssignmentForTopicReturnsEmptyMapIfTopicDoesNotExist(): Unit = {
+    assertEquals(Map.empty, context.partitionReplicaAssignmentForTopic("NONEXISTENT"))
+  }
+
+  @Test
+  def testPartitionReplicaAssignmentForTopicReturnsExpectedReplicaAssignments(): Unit = {
+    val expectedAssignments = Map(
+      tp1 -> context.partitionReplicaAssignment(tp1),
+      tp2 -> context.partitionReplicaAssignment(tp2)
+    )
+    val receivedAssignments = context.partitionReplicaAssignmentForTopic("A")
+    assertEquals(expectedAssignments, receivedAssignments)
+  }
+
+  @Test
+  def testPartitionReplicaAssignment(): Unit = {
+    val reassigningPartition = ReplicaAssignment(List(1, 2, 3, 4, 5, 6), List(2, 3, 4), List(1, 5, 6))
+    assertTrue(reassigningPartition.isBeingReassigned)
+    assertEquals(List(2, 3, 4), reassigningPartition.targetReplicas)
+
+    val reassigningPartition2 = ReplicaAssignment(List(1, 2, 3, 4), List(), List(1, 4))
+    assertTrue(reassigningPartition2.isBeingReassigned)
+    assertEquals(List(2, 3), reassigningPartition2.targetReplicas)
+
+    val reassigningPartition3 = ReplicaAssignment(List(1, 2, 3, 4), List(4), List(2))
+    assertTrue(reassigningPartition3.isBeingReassigned)
+    assertEquals(List(1, 3, 4), reassigningPartition3.targetReplicas)
+
+    val partition = ReplicaAssignment(List(1, 2, 3, 4, 5, 6), List(), List())
+    assertFalse(partition.isBeingReassigned)
+    assertEquals(List(1, 2, 3, 4, 5, 6), partition.targetReplicas)
+
+    val reassigningPartition4 = ReplicaAssignment.fromOldAndNewReplicas(
+      List(1, 2, 3, 4), List(4, 2, 5, 3)
+    )
+    assertEquals(List(4, 2, 5, 3, 1), reassigningPartition4.replicas)
+    assertEquals(List(4, 2, 5, 3), reassigningPartition4.targetReplicas)
+    assertEquals(List(5), reassigningPartition4.addingReplicas)
+    assertEquals(List(1), reassigningPartition4.removingReplicas)
+    assertTrue(reassigningPartition4.isBeingReassigned)
+
+    val reassigningPartition5 = ReplicaAssignment.fromOldAndNewReplicas(
+      List(1, 2, 3), List(4, 5, 6)
+    )
+    assertEquals(List(4, 5, 6, 1, 2, 3), reassigningPartition5.replicas)
+    assertEquals(List(4, 5, 6), reassigningPartition5.targetReplicas)
+    assertEquals(List(4, 5, 6), reassigningPartition5.addingReplicas)
+    assertEquals(List(1, 2, 3), reassigningPartition5.removingReplicas)
+    assertTrue(reassigningPartition5.isBeingReassigned)
+
+    val nonReassigningPartition = ReplicaAssignment.fromOldAndNewReplicas(
+      List(1, 2, 3), List(3, 1, 2)
+    )
+    assertEquals(List(3, 1, 2), nonReassigningPartition.replicas)
+    assertEquals(List(3, 1, 2), nonReassigningPartition.targetReplicas)
+    assertEquals(List(), nonReassigningPartition.addingReplicas)
+    assertEquals(List(), nonReassigningPartition.removingReplicas)
+    assertFalse(nonReassigningPartition.isBeingReassigned)
+  }
+
+  @Test
+  def testReassignTo(): Unit = {
+    val assignment = ReplicaAssignment(Seq(1, 2, 3))
+    val firstReassign = assignment.reassignTo(Seq(4, 5, 6))
+
+    assertEquals(ReplicaAssignment(Seq(4, 5, 6, 1, 2, 3), Seq(4, 5, 6), Seq(1, 2, 3)), firstReassign)
+    assertEquals(ReplicaAssignment(Seq(7, 8, 9, 1, 2, 3), Seq(7, 8, 9), Seq(1, 2, 3)), firstReassign.reassignTo(Seq(7, 8, 9)))
+    assertEquals(ReplicaAssignment(Seq(7, 8, 9, 1, 2, 3), Seq(7, 8, 9), Seq(1, 2, 3)), assignment.reassignTo(Seq(7, 8, 9)))
+    assertEquals(assignment, firstReassign.reassignTo(Seq(1,2,3)))
+  }
+
+}
diff --git a/core/src/test/scala/unit/kafka/controller/ControllerEventManagerTest.scala b/core/src/test/scala/unit/kafka/controller/ControllerEventManagerTest.scala
index e0a753cacaf74..83bc38fddcfee 100644
--- a/core/src/test/scala/unit/kafka/controller/ControllerEventManagerTest.scala
+++ b/core/src/test/scala/unit/kafka/controller/ControllerEventManagerTest.scala
@@ -21,10 +21,11 @@ import java.util.concurrent.CountDownLatch
 import java.util.concurrent.atomic.AtomicInteger
 
 import com.yammer.metrics.Metrics
-import com.yammer.metrics.core.Timer
+import com.yammer.metrics.core.{Histogram, MetricName, Timer}
 import kafka.utils.TestUtils
+import org.apache.kafka.common.utils.MockTime
+import org.junit.Assert.{assertEquals, assertTrue, fail}
 import org.junit.{After, Test}
-import org.junit.Assert.{assertEquals, fail}
 
 import scala.collection.JavaConverters._
 
@@ -38,37 +39,106 @@ class ControllerEventManagerTest {
       controllerEventManager.close()
   }
 
+  @Test
+  def testMetricsCleanedOnClose(): Unit = {
+    val time = new MockTime()
+    val controllerStats = new ControllerStats
+    val eventProcessor = new ControllerEventProcessor {
+      override def process(event: ControllerEvent): Unit = {}
+      override def preempt(event: ControllerEvent): Unit = {}
+    }
+
+    def allEventManagerMetrics: Set[MetricName] = {
+      Metrics.defaultRegistry.allMetrics.asScala.keySet
+        .filter(_.getMBeanName.startsWith("kafka.controller:type=ControllerEventManager"))
+        .toSet
+    }
+
+    controllerEventManager = new ControllerEventManager(0, eventProcessor,
+      time, controllerStats.rateAndTimeMetrics)
+    controllerEventManager.start()
+    assertTrue(allEventManagerMetrics.nonEmpty)
+
+    controllerEventManager.close()
+    assertTrue(allEventManagerMetrics.isEmpty)
+  }
+
+  @Test
+  def testEventQueueTime(): Unit = {
+    val metricName = "kafka.controller:type=ControllerEventManager,name=EventQueueTimeMs"
+    val controllerStats = new ControllerStats
+    val time = new MockTime()
+    val latch = new CountDownLatch(1)
+    val processedEvents = new AtomicInteger()
+
+    val eventProcessor = new ControllerEventProcessor {
+      override def process(event: ControllerEvent): Unit = {
+        latch.await()
+        time.sleep(500)
+        processedEvents.incrementAndGet()
+      }
+      override def preempt(event: ControllerEvent): Unit = {}
+    }
+
+    // The metric should not already exist
+    assertTrue(Metrics.defaultRegistry.allMetrics.asScala.filterKeys(_.getMBeanName == metricName).values.isEmpty)
+
+    controllerEventManager = new ControllerEventManager(0, eventProcessor,
+      time, controllerStats.rateAndTimeMetrics)
+    controllerEventManager.start()
+
+    controllerEventManager.put(TopicChange)
+    controllerEventManager.put(TopicChange)
+    latch.countDown()
+
+    TestUtils.waitUntilTrue(() => processedEvents.get() == 2,
+      "Timed out waiting for processing of all events")
+
+    val queueTimeHistogram = Metrics.defaultRegistry.allMetrics.asScala.filterKeys(_.getMBeanName == metricName).values.headOption
+      .getOrElse(fail(s"Unable to find metric $metricName")).asInstanceOf[Histogram]
+
+    assertEquals(2, queueTimeHistogram.count)
+    assertEquals(0, queueTimeHistogram.min, 0.01)
+    assertEquals(500, queueTimeHistogram.max, 0.01)
+  }
+
   @Test
   def testSuccessfulEvent(): Unit = {
-    check("kafka.controller:type=ControllerStats,name=AutoLeaderBalanceRateAndTimeMs", ControllerState.AutoLeaderBalance,
-      () => Unit)
+    check("kafka.controller:type=ControllerStats,name=AutoLeaderBalanceRateAndTimeMs",
+      AutoPreferredReplicaLeaderElection, () => ())
   }
 
   @Test
   def testEventThatThrowsException(): Unit = {
-    check("kafka.controller:type=ControllerStats,name=LeaderElectionRateAndTimeMs", ControllerState.BrokerChange,
-      () => throw new NullPointerException)
+    check("kafka.controller:type=ControllerStats,name=LeaderElectionRateAndTimeMs",
+      BrokerChange, () => throw new NullPointerException)
   }
 
-  private def check(metricName: String, controllerState: ControllerState, process: () => Unit): Unit = {
+  private def check(metricName: String,
+                    event: ControllerEvent,
+                    func: () => Unit): Unit = {
     val controllerStats = new ControllerStats
     val eventProcessedListenerCount = new AtomicInteger
-    controllerEventManager = new ControllerEventManager(0, controllerStats.rateAndTimeMetrics,
-      _ => eventProcessedListenerCount.incrementAndGet, () => ())
+    val latch = new CountDownLatch(1)
+    val eventProcessor = new ControllerEventProcessor {
+      override def process(event: ControllerEvent): Unit = {
+        // Only return from `process()` once we have checked `controllerEventManager.state`
+        latch.await()
+        eventProcessedListenerCount.incrementAndGet()
+        func()
+      }
+      override def preempt(event: ControllerEvent): Unit = {}
+    }
+
+    controllerEventManager = new ControllerEventManager(0, eventProcessor,
+      new MockTime(), controllerStats.rateAndTimeMetrics)
     controllerEventManager.start()
 
     val initialTimerCount = timer(metricName).count
 
-    // Only return from `process()` once we have checked `controllerEventManager.state`
-    val latch = new CountDownLatch(1)
-    val eventMock = ControllerTestUtils.createMockControllerEvent(controllerState, { () =>
-      latch.await()
-      process()
-    })
-
-    controllerEventManager.put(eventMock)
-    TestUtils.waitUntilTrue(() => controllerEventManager.state == controllerState,
-      s"Controller state is not $controllerState")
+    controllerEventManager.put(event)
+    TestUtils.waitUntilTrue(() => controllerEventManager.state == event.state,
+      s"Controller state is not ${event.state}")
     latch.countDown()
 
     TestUtils.waitUntilTrue(() => controllerEventManager.state == ControllerState.Idle,
diff --git a/core/src/test/scala/unit/kafka/controller/ControllerFailoverTest.scala b/core/src/test/scala/unit/kafka/controller/ControllerFailoverTest.scala
index 6cfa72cf3842f..abbf6721b176e 100644
--- a/core/src/test/scala/unit/kafka/controller/ControllerFailoverTest.scala
+++ b/core/src/test/scala/unit/kafka/controller/ControllerFailoverTest.scala
@@ -19,6 +19,7 @@ package kafka.controller
 
 import java.util.Properties
 import java.util.concurrent.CountDownLatch
+import java.util.concurrent.atomic.AtomicReference
 
 import kafka.integration.KafkaServerTestHarness
 import kafka.server.KafkaConfig
@@ -28,6 +29,7 @@ import org.apache.kafka.common.metrics.Metrics
 import org.apache.log4j.Logger
 import org.junit.{After, Test}
 import org.junit.Assert._
+import org.scalatest.Assertions.fail
 
 class ControllerFailoverTest extends KafkaServerTestHarness with Logging {
   val log = Logger.getLogger(classOf[ControllerFailoverTest])
@@ -43,7 +45,7 @@ class ControllerFailoverTest extends KafkaServerTestHarness with Logging {
     .map(KafkaConfig.fromProps(_, overridingProps))
 
   @After
-  override def tearDown() {
+  override def tearDown(): Unit = {
     super.tearDown()
     this.metrics.close()
   }
@@ -53,7 +55,7 @@ class ControllerFailoverTest extends KafkaServerTestHarness with Logging {
    * for the background of this test case
    */
   @Test
-  def testHandleIllegalStateException() {
+  def testHandleIllegalStateException(): Unit = {
     val initialController = servers.find(_.kafkaController.isActive).map(_.kafkaController).getOrElse {
       fail("Could not find controller")
     }
@@ -62,25 +64,27 @@ class ControllerFailoverTest extends KafkaServerTestHarness with Logging {
     createTopic(topic, 1, 1)
     val topicPartition = new TopicPartition("topic1", 0)
     TestUtils.waitUntilTrue(() =>
-      initialController.partitionStateMachine.partitionsInState(OnlinePartition).contains(topicPartition),
+      initialController.controllerContext.partitionsInState(OnlinePartition).contains(topicPartition),
       s"Partition $topicPartition did not transition to online state")
 
     // Wait until we have verified that we have resigned
     val latch = new CountDownLatch(1)
-    @volatile var exceptionThrown: Option[Throwable] = None
-    val illegalStateEvent = ControllerTestUtils.createMockControllerEvent(ControllerState.BrokerChange, { () =>
-      try initialController.handleIllegalState(new IllegalStateException("Thrown for test purposes"))
-      catch {
-        case t: Throwable => exceptionThrown = Some(t)
+    val exceptionThrown = new AtomicReference[Throwable]()
+    val illegalStateEvent = new MockEvent(ControllerState.BrokerChange) {
+      override def process(): Unit = {
+        try initialController.handleIllegalState(new IllegalStateException("Thrown for test purposes"))
+        catch {
+          case t: Throwable => exceptionThrown.set(t)
+        }
+        latch.await()
       }
-      latch.await()
-    })
+    }
     initialController.eventManager.put(illegalStateEvent)
     // Check that we have shutdown the scheduler (via onControllerResigned)
     TestUtils.waitUntilTrue(() => !initialController.kafkaScheduler.isStarted, "Scheduler was not shutdown")
     TestUtils.waitUntilTrue(() => !initialController.isActive, "Controller did not become inactive")
     latch.countDown()
-    TestUtils.waitUntilTrue(() => exceptionThrown.isDefined, "handleIllegalState did not throw an exception")
+    TestUtils.waitUntilTrue(() => Option(exceptionThrown.get()).isDefined, "handleIllegalState did not throw an exception")
     assertTrue(s"handleIllegalState should throw an IllegalStateException, but $exceptionThrown was thrown",
       exceptionThrown.get.isInstanceOf[IllegalStateException])
 
diff --git a/core/src/test/scala/unit/kafka/controller/ControllerIntegrationTest.scala b/core/src/test/scala/unit/kafka/controller/ControllerIntegrationTest.scala
index f167876644c29..78764f62e32b7 100644
--- a/core/src/test/scala/unit/kafka/controller/ControllerIntegrationTest.scala
+++ b/core/src/test/scala/unit/kafka/controller/ControllerIntegrationTest.scala
@@ -33,10 +33,12 @@ import org.apache.kafka.common.errors.{ControllerMovedException, StaleBrokerEpoc
 import org.apache.log4j.Level
 import kafka.utils.LogCaptureAppender
 import org.apache.kafka.common.metrics.KafkaMetric
+import org.scalatest.Assertions.fail
 
 import scala.collection.JavaConverters._
 import scala.collection.mutable
-import scala.util.Try
+import scala.collection.Seq
+import scala.util.{Failure, Success, Try}
 
 class ControllerIntegrationTest extends ZooKeeperTestHarness {
   var servers = Seq.empty[KafkaServer]
@@ -44,13 +46,13 @@ class ControllerIntegrationTest extends ZooKeeperTestHarness {
   val firstControllerEpochZkVersion = KafkaController.InitialControllerEpochZkVersion + 1
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     super.setUp
     servers = Seq.empty[KafkaServer]
   }
 
   @After
-  override def tearDown() {
+  override def tearDown(): Unit = {
     TestUtils.shutdownServers(servers)
     super.tearDown
   }
@@ -87,7 +89,9 @@ class ControllerIntegrationTest extends ZooKeeperTestHarness {
 
   @Test
   def testMetadataPropagationOnControlPlane(): Unit = {
-    servers = makeServers(1, listeners = Some("PLAINTEXT://localhost:0,CONTROLLER://localhost:5000"), listenerSecurityProtocolMap = Some("PLAINTEXT:PLAINTEXT,CONTROLLER:PLAINTEXT"),
+    servers = makeServers(1,
+      listeners = Some("PLAINTEXT://localhost:0,CONTROLLER://localhost:0"),
+      listenerSecurityProtocolMap = Some("PLAINTEXT:PLAINTEXT,CONTROLLER:PLAINTEXT"),
       controlPlaneListenerName = Some("CONTROLLER"))
     TestUtils.waitUntilBrokerMetadataIsPropagated(servers)
     val controlPlaneMetricMap = mutable.Map[String, KafkaMetric]()
@@ -100,14 +104,14 @@ class ControllerIntegrationTest extends ZooKeeperTestHarness {
         dataPlaneMetricMap.put(kafkaMetric.metricName().name(), kafkaMetric)
       }
     }
-    assertEquals(1e-0, controlPlaneMetricMap.get("response-total").get.metricValue().asInstanceOf[Double], 0)
-    assertEquals(0e-0, dataPlaneMetricMap.get("response-total").get.metricValue().asInstanceOf[Double], 0)
-    assertEquals(1e-0, controlPlaneMetricMap.get("request-total").get.metricValue().asInstanceOf[Double], 0)
-    assertEquals(0e-0, dataPlaneMetricMap.get("request-total").get.metricValue().asInstanceOf[Double], 0)
-    assertTrue(controlPlaneMetricMap.get("incoming-byte-total").get.metricValue().asInstanceOf[Double] > 1.0)
-    assertTrue(dataPlaneMetricMap.get("incoming-byte-total").get.metricValue().asInstanceOf[Double] == 0.0)
-    assertTrue(controlPlaneMetricMap.get("network-io-total").get.metricValue().asInstanceOf[Double] == 2.0)
-    assertTrue(dataPlaneMetricMap.get("network-io-total").get.metricValue().asInstanceOf[Double] == 0.0)
+    assertEquals(1e-0, controlPlaneMetricMap("response-total").metricValue().asInstanceOf[Double], 0)
+    assertEquals(0e-0, dataPlaneMetricMap("response-total").metricValue().asInstanceOf[Double], 0)
+    assertEquals(1e-0, controlPlaneMetricMap("request-total").metricValue().asInstanceOf[Double], 0)
+    assertEquals(0e-0, dataPlaneMetricMap("request-total").metricValue().asInstanceOf[Double], 0)
+    assertTrue(controlPlaneMetricMap("incoming-byte-total").metricValue().asInstanceOf[Double] > 1.0)
+    assertTrue(dataPlaneMetricMap("incoming-byte-total").metricValue().asInstanceOf[Double] == 0.0)
+    assertTrue(controlPlaneMetricMap("network-io-total").metricValue().asInstanceOf[Double] == 2.0)
+    assertTrue(dataPlaneMetricMap("network-io-total").metricValue().asInstanceOf[Double] == 0.0)
   }
 
   // This test case is used to ensure that there will be no correctness issue after we avoid sending out full
@@ -139,10 +143,10 @@ class ControllerIntegrationTest extends ZooKeeperTestHarness {
       val offlineReplicaPartitionInfo = server.metadataCache.getPartitionInfo(topic, 0).get
       assertEquals(1, offlineReplicaPartitionInfo.offlineReplicas.size())
       assertEquals(testBroker.config.brokerId, offlineReplicaPartitionInfo.offlineReplicas.get(0))
-      assertEquals(assignment(0).asJava, offlineReplicaPartitionInfo.basePartitionState.replicas)
-      assertEquals(Seq(remainingBrokers.head.config.brokerId).asJava, offlineReplicaPartitionInfo.basePartitionState.isr)
+      assertEquals(assignment(0).asJava, offlineReplicaPartitionInfo.replicas)
+      assertEquals(Seq(remainingBrokers.head.config.brokerId).asJava, offlineReplicaPartitionInfo.isr)
       val onlinePartitionInfo = server.metadataCache.getPartitionInfo(topic, 1).get
-      assertEquals(assignment(1).asJava, onlinePartitionInfo.basePartitionState.replicas)
+      assertEquals(assignment(1).asJava, onlinePartitionInfo.replicas)
       assertTrue(onlinePartitionInfo.offlineReplicas.isEmpty)
     }
 
@@ -154,7 +158,7 @@ class ControllerIntegrationTest extends ZooKeeperTestHarness {
           val partitionInfoOpt = server.metadataCache.getPartitionInfo(topic, partitionId)
           if (partitionInfoOpt.isDefined) {
             val partitionInfo = partitionInfoOpt.get
-            !partitionInfo.offlineReplicas.isEmpty || !partitionInfo.basePartitionState.replicas.asScala.equals(replicas)
+            !partitionInfo.offlineReplicas.isEmpty || !partitionInfo.replicas.asScala.equals(replicas)
           } else {
             true
           }
@@ -163,6 +167,58 @@ class ControllerIntegrationTest extends ZooKeeperTestHarness {
     }, "Inconsistent metadata after broker startup")
   }
 
+  @Test
+  def testMetadataPropagationForOfflineReplicas(): Unit = {
+    servers = makeServers(3)
+    TestUtils.waitUntilBrokerMetadataIsPropagated(servers)
+    val controllerId = TestUtils.waitUntilControllerElected(zkClient)
+
+    //get brokerId for topic creation with single partition and RF =1
+    val replicaBroker = servers.filter(e => e.config.brokerId != controllerId).head
+
+    val controllerBroker = servers.filter(e => e.config.brokerId == controllerId).head
+    val otherBroker = servers.filter(e => e.config.brokerId != controllerId &&
+      e.config.brokerId != replicaBroker.config.brokerId).head
+
+    val topic = "topic1"
+    val assignment = Map(0 -> Seq(replicaBroker.config.brokerId))
+
+    // Create topic
+    TestUtils.createTopic(zkClient, topic, assignment, servers)
+
+    // Shutdown the other broker
+    otherBroker.shutdown()
+    otherBroker.awaitShutdown()
+
+    // Shutdown the broker with replica
+    replicaBroker.shutdown()
+    replicaBroker.awaitShutdown()
+
+    //Shutdown controller broker
+    controllerBroker.shutdown()
+    controllerBroker.awaitShutdown()
+
+    def verifyMetadata(broker: KafkaServer): Unit = {
+      broker.startup()
+      TestUtils.waitUntilTrue(() => {
+        val partitionInfoOpt = broker.metadataCache.getPartitionInfo(topic, 0)
+        if (partitionInfoOpt.isDefined) {
+          val partitionInfo = partitionInfoOpt.get
+          (!partitionInfo.offlineReplicas.isEmpty && partitionInfo.leader == -1
+            && !partitionInfo.replicas.isEmpty && !partitionInfo.isr.isEmpty)
+        } else {
+          false
+        }
+      }, "Inconsistent metadata after broker startup")
+    }
+
+    //Start controller broker and check metadata
+    verifyMetadata(controllerBroker)
+
+    //Start other broker and check metadata
+    verifyMetadata(otherBroker)
+  }
+
   @Test
   def testTopicCreation(): Unit = {
     servers = makeServers(1)
@@ -187,13 +243,42 @@ class ControllerIntegrationTest extends ZooKeeperTestHarness {
       "failed to get expected partition state upon topic creation")
   }
 
+  /*
+  * Tests that controller will fix insufficient RF topic by assigning sufficient replicas
+  * */
+  @Test
+  def testTopicCreationWithFixingRF(): Unit = {
+    val topicRF1 = "test_topic_rf1"
+    val partition = 0
+    val defaultRF = 2
+    val serverConfigs = TestUtils.createBrokerConfigs(3, zkConnect, false)
+      .map(prop => {
+        prop.setProperty(KafkaConfig.DefaultReplicationFactorProp,defaultRF.toString)
+        prop.setProperty(KafkaConfig.CreateTopicPolicyClassNameProp, "kafka.server.LiCreateTopicPolicy")
+        prop
+      }).map(KafkaConfig.fromProps)
+    servers = serverConfigs.map(s => TestUtils.createServer(s))
+
+    val controllerId = TestUtils.waitUntilControllerElected(zkClient)
+    val assignment = Map(partition -> List(controllerId)) // topicRF1 original RF set to 1
+    TestUtils.createTopic(zkClient, topicRF1, partitionReplicaAssignment = assignment, servers = servers, new Properties())
+
+    waitForPartitionState(new TopicPartition(topicRF1, partition), firstControllerEpoch, zkClient.getAllBrokersInCluster.map(b => b.id),
+      LeaderAndIsr.initialLeaderEpoch, "failed to get expected partition state upon valid RF topic creation")
+
+    // Actual RF should be corrected to 2
+    assertEquals(defaultRF, zkClient.getReplicasForPartition(new TopicPartition(topicRF1, partition)).size)
+  }
+
   @Test
   def testTopicPartitionExpansion(): Unit = {
     servers = makeServers(1)
     val tp0 = new TopicPartition("t", 0)
     val tp1 = new TopicPartition("t", 1)
     val assignment = Map(tp0.partition -> Seq(0))
-    val expandedAssignment = Map(tp0 -> Seq(0), tp1 -> Seq(0))
+    val expandedAssignment = Map(
+      tp0 -> ReplicaAssignment(Seq(0), Seq(), Seq()),
+      tp1 -> ReplicaAssignment(Seq(0), Seq(), Seq()))
     TestUtils.createTopic(zkClient, tp0.topic, partitionReplicaAssignment = assignment, servers = servers)
     zkClient.setTopicAssignment(tp0.topic, expandedAssignment, firstControllerEpochZkVersion)
     waitForPartitionState(tp1, firstControllerEpoch, 0, LeaderAndIsr.initialLeaderEpoch,
@@ -209,7 +294,9 @@ class ControllerIntegrationTest extends ZooKeeperTestHarness {
     val tp0 = new TopicPartition("t", 0)
     val tp1 = new TopicPartition("t", 1)
     val assignment = Map(tp0.partition -> Seq(otherBrokerId, controllerId))
-    val expandedAssignment = Map(tp0 -> Seq(otherBrokerId, controllerId), tp1 -> Seq(otherBrokerId, controllerId))
+    val expandedAssignment = Map(
+      tp0 -> ReplicaAssignment(Seq(otherBrokerId, controllerId), Seq(), Seq()),
+      tp1 -> ReplicaAssignment(Seq(otherBrokerId, controllerId), Seq(), Seq()))
     TestUtils.createTopic(zkClient, tp0.topic, partitionReplicaAssignment = assignment, servers = servers)
     servers(otherBrokerId).shutdown()
     servers(otherBrokerId).awaitShutdown()
@@ -224,18 +311,18 @@ class ControllerIntegrationTest extends ZooKeeperTestHarness {
     servers = makeServers(2)
     val controllerId = TestUtils.waitUntilControllerElected(zkClient)
 
-    val metricName = s"kafka.controller:type=ControllerStats,name=${ControllerState.PartitionReassignment.rateAndTimeMetricName.get}"
+    val metricName = s"kafka.controller:type=ControllerStats,name=${ControllerState.AlterPartitionReassignment.rateAndTimeMetricName.get}"
     val timerCount = timer(metricName).count
 
     val otherBrokerId = servers.map(_.config.brokerId).filter(_ != controllerId).head
     val tp = new TopicPartition("t", 0)
     val assignment = Map(tp.partition -> Seq(controllerId))
-    val reassignment = Map(tp -> Seq(otherBrokerId))
+    val reassignment = Map(tp -> ReplicaAssignment(Seq(otherBrokerId), List(), List()))
     TestUtils.createTopic(zkClient, tp.topic, partitionReplicaAssignment = assignment, servers = servers)
-    zkClient.createPartitionReassignment(reassignment)
+    zkClient.createPartitionReassignment(reassignment.mapValues(_.replicas).toMap)
     waitForPartitionState(tp, firstControllerEpoch, otherBrokerId, LeaderAndIsr.initialLeaderEpoch + 3,
       "failed to get expected partition state after partition reassignment")
-    TestUtils.waitUntilTrue(() =>  zkClient.getReplicaAssignmentForTopics(Set(tp.topic)) == reassignment,
+    TestUtils.waitUntilTrue(() =>  zkClient.getFullReplicaAssignmentForTopics(Set(tp.topic)) == reassignment,
       "failed to get updated partition assignment on topic znode after partition reassignment")
     TestUtils.waitUntilTrue(() => !zkClient.reassignPartitionsInProgress(),
       "failed to remove reassign partitions path after completion")
@@ -270,17 +357,17 @@ class ControllerIntegrationTest extends ZooKeeperTestHarness {
     val otherBrokerId = servers.map(_.config.brokerId).filter(_ != controllerId).head
     val tp = new TopicPartition("t", 0)
     val assignment = Map(tp.partition -> Seq(controllerId))
-    val reassignment = Map(tp -> Seq(otherBrokerId))
+    val reassignment = Map(tp -> ReplicaAssignment(Seq(otherBrokerId), List(), List()))
     TestUtils.createTopic(zkClient, tp.topic, partitionReplicaAssignment = assignment, servers = servers)
     servers(otherBrokerId).shutdown()
     servers(otherBrokerId).awaitShutdown()
-    zkClient.createPartitionReassignment(reassignment)
+    zkClient.createPartitionReassignment(reassignment.mapValues(_.replicas).toMap)
     waitForPartitionState(tp, firstControllerEpoch, controllerId, LeaderAndIsr.initialLeaderEpoch + 1,
       "failed to get expected partition state during partition reassignment with offline replica")
     servers(otherBrokerId).startup()
     waitForPartitionState(tp, firstControllerEpoch, otherBrokerId, LeaderAndIsr.initialLeaderEpoch + 4,
       "failed to get expected partition state after partition reassignment")
-    TestUtils.waitUntilTrue(() => zkClient.getReplicaAssignmentForTopics(Set(tp.topic)) == reassignment,
+    TestUtils.waitUntilTrue(() => zkClient.getFullReplicaAssignmentForTopics(Set(tp.topic)) == reassignment,
       "failed to get updated partition assignment on topic znode after partition reassignment")
     TestUtils.waitUntilTrue(() => !zkClient.reassignPartitionsInProgress(),
       "failed to remove reassign partitions path after completion")
@@ -343,6 +430,24 @@ class ControllerIntegrationTest extends ZooKeeperTestHarness {
       "failed to get expected partition state upon broker startup")
   }
 
+  @Test
+  def testTopicDeletionCleanUpPartitionState(): Unit = {
+    servers = makeServers(3, enableDeleteTopic = true)
+    val controllerId = TestUtils.waitUntilControllerElected(zkClient)
+    val controller = servers.find(broker => broker.config.brokerId == controllerId).get.kafkaController
+    val tp = new TopicPartition("t", 0)
+    val assignment = Map(tp.partition -> Seq(1, 0, 2))
+    TestUtils.createTopic(zkClient, tp.topic, partitionReplicaAssignment = assignment, servers = servers)
+
+    // Make sure the partition state has been populated
+    assertTrue(controller.controllerContext.partitionStates.contains(tp))
+    adminZkClient.deleteTopic(tp.topic())
+    TestUtils.verifyTopicDeletion(zkClient, tp.topic(), 1, servers)
+
+    // Make sure the partition state has been removed
+    assertTrue(!controller.controllerContext.partitionStates.contains(tp))
+  }
+
   @Test
   def testLeaderAndIsrWhenEntireIsrOfflineAndUncleanLeaderElectionDisabled(): Unit = {
     servers = makeServers(2)
@@ -384,7 +489,7 @@ class ControllerIntegrationTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  def testControlledShutdown() {
+  def testControlledShutdown(): Unit = {
     val expectedReplicaAssignment = Map(0  -> List(0, 1, 2))
     val topic = "test"
     val partition = 0
@@ -403,28 +508,31 @@ class ControllerIntegrationTest extends ZooKeeperTestHarness {
     var activeServers = servers.filter(s => s.config.brokerId != 2)
     // wait for the update metadata request to trickle to the brokers
     TestUtils.waitUntilTrue(() =>
-      activeServers.forall(_.dataPlaneRequestProcessor.metadataCache.getPartitionInfo(topic,partition).get.basePartitionState.isr.size != 3),
+      activeServers.forall(_.dataPlaneRequestProcessor.metadataCache.getPartitionInfo(topic,partition).get.isr.size != 3),
       "Topic test not created after timeout")
     assertEquals(0, partitionsRemaining.size)
     var partitionStateInfo = activeServers.head.dataPlaneRequestProcessor.metadataCache.getPartitionInfo(topic,partition).get
-    var leaderAfterShutdown = partitionStateInfo.basePartitionState.leader
+    var leaderAfterShutdown = partitionStateInfo.leader
     assertEquals(0, leaderAfterShutdown)
-    assertEquals(2, partitionStateInfo.basePartitionState.isr.size)
-    assertEquals(List(0,1), partitionStateInfo.basePartitionState.isr.asScala)
+    assertEquals(2, partitionStateInfo.isr.size)
+    assertEquals(List(0,1), partitionStateInfo.isr.asScala)
     controller.controlledShutdown(1, servers.find(_.config.brokerId == 1).get.kafkaController.brokerEpoch, controlledShutdownCallback)
-    partitionsRemaining = resultQueue.take().get
+    partitionsRemaining = resultQueue.take() match {
+      case Success(partitions) => partitions
+      case Failure(exception) => fail("Controlled shutdown failed due to error", exception)
+    }
     assertEquals(0, partitionsRemaining.size)
     activeServers = servers.filter(s => s.config.brokerId == 0)
     partitionStateInfo = activeServers.head.dataPlaneRequestProcessor.metadataCache.getPartitionInfo(topic,partition).get
-    leaderAfterShutdown = partitionStateInfo.basePartitionState.leader
+    leaderAfterShutdown = partitionStateInfo.leader
     assertEquals(0, leaderAfterShutdown)
 
-    assertTrue(servers.forall(_.dataPlaneRequestProcessor.metadataCache.getPartitionInfo(topic,partition).get.basePartitionState.leader == 0))
+    assertTrue(servers.forall(_.dataPlaneRequestProcessor.metadataCache.getPartitionInfo(topic,partition).get.leader == 0))
     controller.controlledShutdown(0, servers.find(_.config.brokerId == 0).get.kafkaController.brokerEpoch, controlledShutdownCallback)
     partitionsRemaining = resultQueue.take().get
     assertEquals(1, partitionsRemaining.size)
     // leader doesn't change since all the replicas are shut down
-    assertTrue(servers.forall(_.dataPlaneRequestProcessor.metadataCache.getPartitionInfo(topic,partition).get.basePartitionState.leader == 0))
+    assertTrue(servers.forall(_.dataPlaneRequestProcessor.metadataCache.getPartitionInfo(topic,partition).get.leader == 0))
   }
 
   @Test
@@ -516,7 +624,10 @@ class ControllerIntegrationTest extends ZooKeeperTestHarness {
 
     // Let the controller event thread await on a latch until broker bounce finishes.
     // This is used to simulate fast broker bounce
-    controller.eventManager.put(KafkaController.AwaitOnLatch(latch))
+
+    controller.eventManager.put(new MockEvent(ControllerState.TopicChange) {
+      override def process(): Unit = latch.await()
+    })
 
     otherBroker.shutdown()
     otherBroker.startup()
@@ -535,7 +646,7 @@ class ControllerIntegrationTest extends ZooKeeperTestHarness {
   private def testControllerMove(fun: () => Unit): Unit = {
     val controller = getController().kafkaController
     val appender = LogCaptureAppender.createAndRegister()
-    val previousLevel = LogCaptureAppender.setClassLoggerLevel(controller.eventManager.thread.getClass, Level.INFO)
+    val previousLevel = LogCaptureAppender.setClassLoggerLevel(controller.getClass, Level.INFO)
 
     try {
       TestUtils.waitUntilTrue(() => {
@@ -546,7 +657,10 @@ class ControllerIntegrationTest extends ZooKeeperTestHarness {
 
       // Let the controller event thread await on a latch before the pre-defined logic is triggered.
       // This is used to make sure that when the event thread resumes and starts processing events, the controller has already moved.
-      controller.eventManager.put(KafkaController.AwaitOnLatch(latch))
+      controller.eventManager.put(new MockEvent(ControllerState.TopicChange) {
+        override def process(): Unit = latch.await()
+      })
+
       // Execute pre-defined logic. This can be topic creation/deletion, preferred leader election, etc.
       fun()
 
@@ -561,7 +675,7 @@ class ControllerIntegrationTest extends ZooKeeperTestHarness {
       // Expect to capture the ControllerMovedException in the log of ControllerEventThread
       val event = appender.getMessages.find(e => e.getLevel == Level.INFO
         && e.getThrowableInformation != null
-        && e.getThrowableInformation.getThrowable.getClass.getName.equals(new ControllerMovedException("").getClass.getName))
+        && e.getThrowableInformation.getThrowable.getClass.getName.equals(classOf[ControllerMovedException].getName))
       assertTrue(event.isDefined)
 
     } finally {
@@ -601,6 +715,18 @@ class ControllerIntegrationTest extends ZooKeeperTestHarness {
     }, message)
   }
 
+  private def waitForPartitionState(tp: TopicPartition,
+    controllerEpoch: Int,
+    leaders: Seq[Int],
+    leaderEpoch: Int,
+    message: String): Unit = {
+    TestUtils.waitUntilTrue(() => {
+      val leaderIsrAndControllerEpochMap = zkClient.getTopicPartitionStates(Seq(tp))
+      leaderIsrAndControllerEpochMap.contains(tp) && leaders.exists(
+        leader=> isExpectedPartitionState(leaderIsrAndControllerEpochMap(tp), controllerEpoch, leader, leaderEpoch))
+    }, message)
+  }
+
   private def isExpectedPartitionState(leaderIsrAndControllerEpoch: LeaderIsrAndControllerEpoch,
                                        controllerEpoch: Int,
                                        leader: Int,
@@ -615,7 +741,8 @@ class ControllerIntegrationTest extends ZooKeeperTestHarness {
                           enableControlledShutdown: Boolean = true,
                           listeners : Option[String] = None,
                           listenerSecurityProtocolMap : Option[String] = None,
-                          controlPlaneListenerName : Option[String] = None) = {
+                          controlPlaneListenerName : Option[String] = None,
+                          enableDeleteTopic: Boolean = false) = {
     val configs = TestUtils.createBrokerConfigs(numConfigs, zkConnect, enableControlledShutdown = enableControlledShutdown)
     configs.foreach { config =>
       config.setProperty(KafkaConfig.AutoLeaderRebalanceEnableProp, autoLeaderRebalanceEnable.toString)
@@ -624,6 +751,7 @@ class ControllerIntegrationTest extends ZooKeeperTestHarness {
       listeners.foreach(listener => config.setProperty(KafkaConfig.ListenersProp, listener))
       listenerSecurityProtocolMap.foreach(listenerMap => config.setProperty(KafkaConfig.ListenerSecurityProtocolMapProp, listenerMap))
       controlPlaneListenerName.foreach(controlPlaneListener => config.setProperty(KafkaConfig.ControlPlaneListenerNameProp, controlPlaneListener))
+      config.setProperty(KafkaConfig.DeleteTopicEnableProp, enableDeleteTopic.toString)
     }
     configs.map(config => TestUtils.createServer(KafkaConfig.fromProps(config)))
   }
diff --git a/core/src/test/scala/unit/kafka/controller/MockPartitionStateMachine.scala b/core/src/test/scala/unit/kafka/controller/MockPartitionStateMachine.scala
new file mode 100644
index 0000000000000..0c6c00dd63051
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/controller/MockPartitionStateMachine.scala
@@ -0,0 +1,119 @@
+/*
+ * 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.controller
+
+import kafka.api.LeaderAndIsr
+import kafka.common.StateChangeFailedException
+import kafka.controller.Election._
+import org.apache.kafka.common.TopicPartition
+
+import scala.collection.{Seq, mutable}
+
+class MockPartitionStateMachine(controllerContext: ControllerContext,
+                                uncleanLeaderElectionEnabled: Boolean)
+  extends PartitionStateMachine(controllerContext) {
+
+  override def handleStateChanges(
+    partitions: Seq[TopicPartition],
+    targetState: PartitionState,
+    leaderElectionStrategy: Option[PartitionLeaderElectionStrategy]
+  ): Map[TopicPartition, Either[Throwable, LeaderAndIsr]] = {
+    partitions.foreach(partition => controllerContext.putPartitionStateIfNotExists(partition, NonExistentPartition))
+    val (validPartitions, invalidPartitions) = controllerContext.checkValidPartitionStateChange(partitions, targetState)
+    if (invalidPartitions.nonEmpty) {
+      val currentStates = invalidPartitions.map(p => controllerContext.partitionStates.get(p))
+      throw new IllegalStateException(s"Invalid state transition to $targetState for partitions $currentStates")
+    }
+
+    if (targetState == OnlinePartition) {
+      val uninitializedPartitions = validPartitions.filter(partition => controllerContext.partitionState(partition) == NewPartition)
+      val partitionsToElectLeader = partitions.filter { partition =>
+        val currentState = controllerContext.partitionState(partition)
+        currentState == OfflinePartition || currentState == OnlinePartition
+      }
+
+      uninitializedPartitions.foreach { partition =>
+        controllerContext.putPartitionState(partition, targetState)
+      }
+
+      val electionResults = doLeaderElections(partitionsToElectLeader, leaderElectionStrategy.get)
+      electionResults.foreach {
+        case (partition, Right(_)) => controllerContext.putPartitionState(partition, targetState)
+        case (_, Left(_)) => // Ignore; No need to update the context if the election failed
+      }
+
+      electionResults
+    } else {
+      validPartitions.foreach { partition =>
+        controllerContext.putPartitionState(partition, targetState)
+      }
+      Map.empty
+    }
+  }
+
+  private def doLeaderElections(
+    partitions: Seq[TopicPartition],
+    leaderElectionStrategy: PartitionLeaderElectionStrategy
+  ): Map[TopicPartition, Either[Throwable, LeaderAndIsr]] = {
+    val failedElections = mutable.Map.empty[TopicPartition, Either[Throwable, LeaderAndIsr]]
+    val validLeaderAndIsrs = mutable.Buffer.empty[(TopicPartition, LeaderAndIsr)]
+
+    for (partition <- partitions) {
+      val leaderIsrAndControllerEpoch = controllerContext.partitionLeadershipInfo(partition)
+      if (leaderIsrAndControllerEpoch.controllerEpoch > controllerContext.epoch) {
+        val failMsg = s"Aborted leader election for partition $partition since the LeaderAndIsr path was " +
+          s"already written by another controller. This probably means that the current controller went through " +
+          s"a soft failure and another controller was elected with epoch ${leaderIsrAndControllerEpoch.controllerEpoch}."
+        failedElections.put(partition, Left(new StateChangeFailedException(failMsg)))
+      } else {
+        validLeaderAndIsrs.append((partition, leaderIsrAndControllerEpoch.leaderAndIsr))
+      }
+    }
+
+    val electionResults = leaderElectionStrategy match {
+      case OfflinePartitionLeaderElectionStrategy(isUnclean) =>
+        val partitionsWithUncleanLeaderElectionState = validLeaderAndIsrs.map { case (partition, leaderAndIsr) =>
+          (partition, Some(leaderAndIsr), isUnclean || uncleanLeaderElectionEnabled)
+        }
+        leaderForOffline(controllerContext, partitionsWithUncleanLeaderElectionState)
+      case ReassignPartitionLeaderElectionStrategy =>
+        leaderForReassign(controllerContext, validLeaderAndIsrs)
+      case PreferredReplicaPartitionLeaderElectionStrategy =>
+        leaderForPreferredReplica(controllerContext, validLeaderAndIsrs)
+      case ControlledShutdownPartitionLeaderElectionStrategy =>
+        leaderForControlledShutdown(controllerContext, validLeaderAndIsrs)
+    }
+
+    val results: Map[TopicPartition, Either[Exception, LeaderAndIsr]] = electionResults.map { electionResult =>
+      val partition = electionResult.topicPartition
+      val value = electionResult.leaderAndIsr match {
+        case None =>
+          val failMsg = s"Failed to elect leader for partition $partition under strategy $leaderElectionStrategy"
+          Left(new StateChangeFailedException(failMsg))
+        case Some(leaderAndIsr) =>
+          val leaderIsrAndControllerEpoch = LeaderIsrAndControllerEpoch(leaderAndIsr, controllerContext.epoch)
+          controllerContext.partitionLeadershipInfo.put(partition, leaderIsrAndControllerEpoch)
+          Right(leaderAndIsr)
+      }
+
+      partition -> value
+    }.toMap
+
+    results ++ failedElections
+  }
+
+}
diff --git a/core/src/test/scala/unit/kafka/controller/MockReplicaStateMachine.scala b/core/src/test/scala/unit/kafka/controller/MockReplicaStateMachine.scala
new file mode 100644
index 0000000000000..e5207bfddb39f
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/controller/MockReplicaStateMachine.scala
@@ -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 kafka.controller
+
+import scala.collection.Seq
+
+class MockReplicaStateMachine(controllerContext: ControllerContext) extends ReplicaStateMachine(controllerContext) {
+
+  override def handleStateChanges(replicas: Seq[PartitionAndReplica], targetState: ReplicaState): Unit = {
+    replicas.foreach(replica => controllerContext.putReplicaStateIfNotExists(replica, NonExistentReplica))
+    val (validReplicas, invalidReplicas) = controllerContext.checkValidReplicaStateChange(replicas, targetState)
+    if (invalidReplicas.nonEmpty) {
+      val currentStates = invalidReplicas.map(replica => replica -> controllerContext.replicaStates.get(replica)).toMap
+      throw new IllegalStateException(s"Invalid state transition to $targetState for replicas $currentStates")
+    }
+    validReplicas.foreach { replica =>
+      if (targetState == NonExistentReplica)
+        controllerContext.removeReplicaState(replica)
+      else
+        controllerContext.putReplicaState(replica, targetState)
+    }
+  }
+
+}
diff --git a/core/src/test/scala/unit/kafka/controller/PartitionLeaderElectionAlgorithmsTest.scala b/core/src/test/scala/unit/kafka/controller/PartitionLeaderElectionAlgorithmsTest.scala
index 113a39d54300e..41bc8833ac88a 100644
--- a/core/src/test/scala/unit/kafka/controller/PartitionLeaderElectionAlgorithmsTest.scala
+++ b/core/src/test/scala/unit/kafka/controller/PartitionLeaderElectionAlgorithmsTest.scala
@@ -18,16 +18,8 @@ package kafka.controller
 
 import org.junit.Assert._
 import org.junit.{Before, Test}
-import org.scalatest.junit.JUnitSuite
 
-class PartitionLeaderElectionAlgorithmsTest  extends JUnitSuite {
-  private var controllerContext: ControllerContext = null
-
-  @Before
-  def setUp(): Unit = {
-    controllerContext = new ControllerContext
-    controllerContext.stats.removeMetric("UncleanLeaderElectionsPerSec")
-  }
+class PartitionLeaderElectionAlgorithmsTest {
 
   @Test
   def testOfflinePartitionLeaderElection(): Unit = {
@@ -37,9 +29,8 @@ class PartitionLeaderElectionAlgorithmsTest  extends JUnitSuite {
     val leaderOpt = PartitionLeaderElectionAlgorithms.offlinePartitionLeaderElection(assignment,
       isr,
       liveReplicas,
-      uncleanLeaderElectionEnabled = false,
-      controllerContext)
-    assertEquals(Option(4), leaderOpt)
+      uncleanLeaderElectionEnabled = false)
+    assertEquals(Option(4, false), leaderOpt)
   }
 
   @Test
@@ -50,10 +41,8 @@ class PartitionLeaderElectionAlgorithmsTest  extends JUnitSuite {
     val leaderOpt = PartitionLeaderElectionAlgorithms.offlinePartitionLeaderElection(assignment,
       isr,
       liveReplicas,
-      uncleanLeaderElectionEnabled = false,
-      controllerContext)
+      uncleanLeaderElectionEnabled = false)
     assertEquals(None, leaderOpt)
-    assertEquals(0, controllerContext.stats.uncleanLeaderElectionRate.count())
   }
 
   @Test
@@ -64,10 +53,8 @@ class PartitionLeaderElectionAlgorithmsTest  extends JUnitSuite {
     val leaderOpt = PartitionLeaderElectionAlgorithms.offlinePartitionLeaderElection(assignment,
       isr,
       liveReplicas,
-      uncleanLeaderElectionEnabled = true,
-      controllerContext)
-    assertEquals(Option(4), leaderOpt)
-    assertEquals(1, controllerContext.stats.uncleanLeaderElectionRate.count())
+      uncleanLeaderElectionEnabled = true)
+    assertEquals(Option(4, true), leaderOpt)
   }
 
   @Test
diff --git a/core/src/test/scala/unit/kafka/controller/PartitionStateMachineTest.scala b/core/src/test/scala/unit/kafka/controller/PartitionStateMachineTest.scala
index d711ae0849172..ef496a146abdd 100644
--- a/core/src/test/scala/unit/kafka/controller/PartitionStateMachineTest.scala
+++ b/core/src/test/scala/unit/kafka/controller/PartitionStateMachineTest.scala
@@ -20,8 +20,8 @@ import kafka.api.LeaderAndIsr
 import kafka.log.LogConfig
 import kafka.server.KafkaConfig
 import kafka.utils.TestUtils
-import kafka.zk.{KafkaZkClient, TopicPartitionStateZNode}
 import kafka.zk.KafkaZkClient.UpdateLeaderAndIsrResult
+import kafka.zk.{KafkaZkClient, TopicPartitionStateZNode}
 import kafka.zookeeper._
 import org.apache.kafka.common.TopicPartition
 import org.apache.zookeeper.KeeperException.Code
@@ -29,16 +29,13 @@ import org.apache.zookeeper.data.Stat
 import org.easymock.EasyMock
 import org.junit.Assert._
 import org.junit.{Before, Test}
-import org.scalatest.junit.JUnitSuite
-
-import scala.collection.mutable
+import org.mockito.Mockito
+import scala.collection.JavaConverters._
 
-class PartitionStateMachineTest extends JUnitSuite {
+class PartitionStateMachineTest {
   private var controllerContext: ControllerContext = null
   private var mockZkClient: KafkaZkClient = null
   private var mockControllerBrokerRequestBatch: ControllerBrokerRequestBatch = null
-  private var mockTopicDeletionManager: TopicDeletionManager = null
-  private var partitionState: mutable.Map[TopicPartition, PartitionState] = null
   private var partitionStateMachine: PartitionStateMachine = null
 
   private val brokerId = 5
@@ -53,11 +50,12 @@ class PartitionStateMachineTest extends JUnitSuite {
     controllerContext.epoch = controllerEpoch
     mockZkClient = EasyMock.createMock(classOf[KafkaZkClient])
     mockControllerBrokerRequestBatch = EasyMock.createMock(classOf[ControllerBrokerRequestBatch])
-    mockTopicDeletionManager = EasyMock.createMock(classOf[TopicDeletionManager])
-    partitionState = mutable.Map.empty[TopicPartition, PartitionState]
-    partitionStateMachine = new PartitionStateMachine(config, new StateChangeLogger(brokerId, true, None), controllerContext,
-      mockZkClient, partitionState, mockControllerBrokerRequestBatch)
-    partitionStateMachine.setTopicDeletionManager(mockTopicDeletionManager)
+    partitionStateMachine = new ZkPartitionStateMachine(config, new StateChangeLogger(brokerId, true, None), controllerContext,
+      mockZkClient, mockControllerBrokerRequestBatch)
+  }
+
+  private def partitionState(partition: TopicPartition): PartitionState = {
+    controllerContext.partitionState(partition)
   }
 
   @Test
@@ -68,7 +66,11 @@ class PartitionStateMachineTest extends JUnitSuite {
 
   @Test
   def testInvalidNonexistentPartitionToOnlinePartitionTransition(): Unit = {
-    partitionStateMachine.handleStateChanges(partitions, OnlinePartition, Option(OfflinePartitionLeaderElectionStrategy))
+    partitionStateMachine.handleStateChanges(
+      partitions,
+      OnlinePartition,
+      Option(OfflinePartitionLeaderElectionStrategy(false))
+    )
     assertEquals(NonExistentPartition, partitionState(partition))
   }
 
@@ -82,32 +84,40 @@ class PartitionStateMachineTest extends JUnitSuite {
   def testNewPartitionToOnlinePartitionTransition(): Unit = {
     controllerContext.setLiveBrokerAndEpochs(Map(TestUtils.createBrokerAndEpoch(brokerId, "host", 0)))
     controllerContext.updatePartitionReplicaAssignment(partition, Seq(brokerId))
-    partitionState.put(partition, NewPartition)
+    controllerContext.putPartitionState(partition, NewPartition)
     val leaderIsrAndControllerEpoch = LeaderIsrAndControllerEpoch(LeaderAndIsr(brokerId, List(brokerId)), controllerEpoch)
     EasyMock.expect(mockControllerBrokerRequestBatch.newBatch())
     EasyMock.expect(mockZkClient.createTopicPartitionStatesRaw(Map(partition -> leaderIsrAndControllerEpoch), controllerContext.epochZkVersion))
       .andReturn(Seq(CreateResponse(Code.OK, null, Some(partition), null, ResponseMetadata(0, 0))))
     EasyMock.expect(mockControllerBrokerRequestBatch.addLeaderAndIsrRequestForBrokers(Seq(brokerId),
-      partition, leaderIsrAndControllerEpoch, Seq(brokerId), isNew = true))
+      partition, leaderIsrAndControllerEpoch, replicaAssignment(Seq(brokerId)), isNew = true))
     EasyMock.expect(mockControllerBrokerRequestBatch.sendRequestsToBrokers(controllerEpoch))
     EasyMock.replay(mockZkClient, mockControllerBrokerRequestBatch)
-    partitionStateMachine.handleStateChanges(partitions, OnlinePartition, Option(OfflinePartitionLeaderElectionStrategy))
+    partitionStateMachine.handleStateChanges(
+      partitions,
+      OnlinePartition,
+      Option(OfflinePartitionLeaderElectionStrategy(false))
+    )
     EasyMock.verify(mockZkClient, mockControllerBrokerRequestBatch)
     assertEquals(OnlinePartition, partitionState(partition))
   }
 
   @Test
-  def testNewPartitionToOnlinePartitionTransitionZkUtilsExceptionFromCreateStates(): Unit = {
+  def testNewPartitionToOnlinePartitionTransitionZooKeeperClientExceptionFromCreateStates(): Unit = {
     controllerContext.setLiveBrokerAndEpochs(Map(TestUtils.createBrokerAndEpoch(brokerId, "host", 0)))
     controllerContext.updatePartitionReplicaAssignment(partition, Seq(brokerId))
-    partitionState.put(partition, NewPartition)
+    controllerContext.putPartitionState(partition, NewPartition)
     val leaderIsrAndControllerEpoch = LeaderIsrAndControllerEpoch(LeaderAndIsr(brokerId, List(brokerId)), controllerEpoch)
     EasyMock.expect(mockControllerBrokerRequestBatch.newBatch())
     EasyMock.expect(mockZkClient.createTopicPartitionStatesRaw(Map(partition -> leaderIsrAndControllerEpoch), controllerContext.epochZkVersion))
       .andThrow(new ZooKeeperClientException("test"))
     EasyMock.expect(mockControllerBrokerRequestBatch.sendRequestsToBrokers(controllerEpoch))
     EasyMock.replay(mockZkClient, mockControllerBrokerRequestBatch)
-    partitionStateMachine.handleStateChanges(partitions, OnlinePartition, Option(OfflinePartitionLeaderElectionStrategy))
+    partitionStateMachine.handleStateChanges(
+      partitions,
+      OnlinePartition,
+      Option(OfflinePartitionLeaderElectionStrategy(false))
+    )
     EasyMock.verify(mockZkClient, mockControllerBrokerRequestBatch)
     assertEquals(NewPartition, partitionState(partition))
   }
@@ -116,28 +126,32 @@ class PartitionStateMachineTest extends JUnitSuite {
   def testNewPartitionToOnlinePartitionTransitionErrorCodeFromCreateStates(): Unit = {
     controllerContext.setLiveBrokerAndEpochs(Map(TestUtils.createBrokerAndEpoch(brokerId, "host", 0)))
     controllerContext.updatePartitionReplicaAssignment(partition, Seq(brokerId))
-    partitionState.put(partition, NewPartition)
+    controllerContext.putPartitionState(partition, NewPartition)
     val leaderIsrAndControllerEpoch = LeaderIsrAndControllerEpoch(LeaderAndIsr(brokerId, List(brokerId)), controllerEpoch)
     EasyMock.expect(mockControllerBrokerRequestBatch.newBatch())
     EasyMock.expect(mockZkClient.createTopicPartitionStatesRaw(Map(partition -> leaderIsrAndControllerEpoch), controllerContext.epochZkVersion))
       .andReturn(Seq(CreateResponse(Code.NODEEXISTS, null, Some(partition), null, ResponseMetadata(0, 0))))
     EasyMock.expect(mockControllerBrokerRequestBatch.sendRequestsToBrokers(controllerEpoch))
     EasyMock.replay(mockZkClient, mockControllerBrokerRequestBatch)
-    partitionStateMachine.handleStateChanges(partitions, OnlinePartition, Option(OfflinePartitionLeaderElectionStrategy))
+    partitionStateMachine.handleStateChanges(
+      partitions,
+      OnlinePartition,
+      Option(OfflinePartitionLeaderElectionStrategy(false))
+    )
     EasyMock.verify(mockZkClient, mockControllerBrokerRequestBatch)
     assertEquals(NewPartition, partitionState(partition))
   }
 
   @Test
   def testNewPartitionToOfflinePartitionTransition(): Unit = {
-    partitionState.put(partition, NewPartition)
+    controllerContext.putPartitionState(partition, NewPartition)
     partitionStateMachine.handleStateChanges(partitions, OfflinePartition)
     assertEquals(OfflinePartition, partitionState(partition))
   }
 
   @Test
   def testInvalidNewPartitionToNonexistentPartitionTransition(): Unit = {
-    partitionState.put(partition, NewPartition)
+    controllerContext.putPartitionState(partition, NewPartition)
     partitionStateMachine.handleStateChanges(partitions, NonExistentPartition)
     assertEquals(NewPartition, partitionState(partition))
   }
@@ -146,7 +160,7 @@ class PartitionStateMachineTest extends JUnitSuite {
   def testOnlinePartitionToOnlineTransition(): Unit = {
     controllerContext.setLiveBrokerAndEpochs(Map(TestUtils.createBrokerAndEpoch(brokerId, "host", 0)))
     controllerContext.updatePartitionReplicaAssignment(partition, Seq(brokerId))
-    partitionState.put(partition, OnlinePartition)
+    controllerContext.putPartitionState(partition, OnlinePartition)
     val leaderAndIsr = LeaderAndIsr(brokerId, List(brokerId))
     val leaderIsrAndControllerEpoch = LeaderIsrAndControllerEpoch(leaderAndIsr, controllerEpoch)
     controllerContext.partitionLeadershipInfo.put(partition, leaderIsrAndControllerEpoch)
@@ -160,9 +174,9 @@ class PartitionStateMachineTest extends JUnitSuite {
     val leaderAndIsrAfterElection = leaderAndIsr.newLeader(brokerId)
     val updatedLeaderAndIsr = leaderAndIsrAfterElection.withZkVersion(2)
     EasyMock.expect(mockZkClient.updateLeaderAndIsr(Map(partition -> leaderAndIsrAfterElection), controllerEpoch, controllerContext.epochZkVersion))
-      .andReturn(UpdateLeaderAndIsrResult(Map(partition -> updatedLeaderAndIsr), Seq.empty, Map.empty))
+      .andReturn(UpdateLeaderAndIsrResult(Map(partition -> Right(updatedLeaderAndIsr)), Seq.empty))
     EasyMock.expect(mockControllerBrokerRequestBatch.addLeaderAndIsrRequestForBrokers(Seq(brokerId),
-      partition, LeaderIsrAndControllerEpoch(updatedLeaderAndIsr, controllerEpoch), Seq(brokerId), isNew = false))
+      partition, LeaderIsrAndControllerEpoch(updatedLeaderAndIsr, controllerEpoch), replicaAssignment(Seq(brokerId)), isNew = false))
     EasyMock.expect(mockControllerBrokerRequestBatch.sendRequestsToBrokers(controllerEpoch))
     EasyMock.replay(mockZkClient, mockControllerBrokerRequestBatch)
 
@@ -179,7 +193,7 @@ class PartitionStateMachineTest extends JUnitSuite {
       TestUtils.createBrokerAndEpoch(otherBrokerId, "host", 0)))
     controllerContext.shuttingDownBrokerIds.add(brokerId)
     controllerContext.updatePartitionReplicaAssignment(partition, Seq(brokerId, otherBrokerId))
-    partitionState.put(partition, OnlinePartition)
+    controllerContext.putPartitionState(partition, OnlinePartition)
     val leaderAndIsr = LeaderAndIsr(brokerId, List(brokerId, otherBrokerId))
     val leaderIsrAndControllerEpoch = LeaderIsrAndControllerEpoch(leaderAndIsr, controllerEpoch)
     controllerContext.partitionLeadershipInfo.put(partition, leaderIsrAndControllerEpoch)
@@ -193,11 +207,11 @@ class PartitionStateMachineTest extends JUnitSuite {
     val leaderAndIsrAfterElection = leaderAndIsr.newLeaderAndIsr(otherBrokerId, List(otherBrokerId))
     val updatedLeaderAndIsr = leaderAndIsrAfterElection.withZkVersion(2)
     EasyMock.expect(mockZkClient.updateLeaderAndIsr(Map(partition -> leaderAndIsrAfterElection), controllerEpoch, controllerContext.epochZkVersion))
-      .andReturn(UpdateLeaderAndIsrResult(Map(partition -> updatedLeaderAndIsr), Seq.empty, Map.empty))
+      .andReturn(UpdateLeaderAndIsrResult(Map(partition -> Right(updatedLeaderAndIsr)), Seq.empty))
 
     // The leaderAndIsr request should be sent to both brokers, including the shutting down one
     EasyMock.expect(mockControllerBrokerRequestBatch.addLeaderAndIsrRequestForBrokers(Seq(brokerId, otherBrokerId),
-      partition, LeaderIsrAndControllerEpoch(updatedLeaderAndIsr, controllerEpoch), Seq(brokerId, otherBrokerId),
+      partition, LeaderIsrAndControllerEpoch(updatedLeaderAndIsr, controllerEpoch), replicaAssignment(Seq(brokerId, otherBrokerId)),
       isNew = false))
     EasyMock.expect(mockControllerBrokerRequestBatch.sendRequestsToBrokers(controllerEpoch))
     EasyMock.replay(mockZkClient, mockControllerBrokerRequestBatch)
@@ -209,31 +223,40 @@ class PartitionStateMachineTest extends JUnitSuite {
 
   @Test
   def testOnlinePartitionToOfflineTransition(): Unit = {
-    partitionState.put(partition, OnlinePartition)
+    controllerContext.putPartitionState(partition, OnlinePartition)
     partitionStateMachine.handleStateChanges(partitions, OfflinePartition)
     assertEquals(OfflinePartition, partitionState(partition))
   }
 
   @Test
   def testInvalidOnlinePartitionToNonexistentPartitionTransition(): Unit = {
-    partitionState.put(partition, OnlinePartition)
+    controllerContext.putPartitionState(partition, OnlinePartition)
     partitionStateMachine.handleStateChanges(partitions, NonExistentPartition)
     assertEquals(OnlinePartition, partitionState(partition))
   }
 
   @Test
   def testInvalidOnlinePartitionToNewPartitionTransition(): Unit = {
-    partitionState.put(partition, OnlinePartition)
+    controllerContext.putPartitionState(partition, OnlinePartition)
     partitionStateMachine.handleStateChanges(partitions, NewPartition)
     assertEquals(OnlinePartition, partitionState(partition))
   }
 
+  def testOfflinePartitionToOnlinePartitionTransitionInIsr(): Unit = {
+    testOfflinePartitionToOnlinePartitionTransition(testUncleanLeaderElection = false)
+  }
+
   @Test
-  def testOfflinePartitionToOnlinePartitionTransition(): Unit = {
+  def testOfflinePartitionToOnlinePartitionTransitionUncleanLeaderElection(): Unit = {
+    testOfflinePartitionToOnlinePartitionTransition(testUncleanLeaderElection = true)
+  }
+
+  private[this] def testOfflinePartitionToOnlinePartitionTransition(testUncleanLeaderElection : Boolean) : Unit = {
     controllerContext.setLiveBrokerAndEpochs(Map(TestUtils.createBrokerAndEpoch(brokerId, "host", 0)))
     controllerContext.updatePartitionReplicaAssignment(partition, Seq(brokerId))
-    partitionState.put(partition, OfflinePartition)
-    val leaderAndIsr = LeaderAndIsr(LeaderAndIsr.NoLeader, List(brokerId))
+    controllerContext.putPartitionState(partition, OfflinePartition)
+    val leaderAndIsr = LeaderAndIsr(LeaderAndIsr.NoLeader, if (testUncleanLeaderElection) List.empty else List(brokerId))
+
     val leaderIsrAndControllerEpoch = LeaderIsrAndControllerEpoch(leaderAndIsr, controllerEpoch)
     controllerContext.partitionLeadershipInfo.put(partition, leaderIsrAndControllerEpoch)
 
@@ -243,27 +266,108 @@ class PartitionStateMachineTest extends JUnitSuite {
       .andReturn(Seq(GetDataResponse(Code.OK, null, Some(partition),
         TopicPartitionStateZNode.encode(leaderIsrAndControllerEpoch), stat, ResponseMetadata(0, 0))))
 
-    EasyMock.expect(mockZkClient.getLogConfigs(Seq.empty, config.originals()))
-      .andReturn((Map(partition.topic -> LogConfig()), Map.empty))
-    val leaderAndIsrAfterElection = leaderAndIsr.newLeader(brokerId)
+    if (testUncleanLeaderElection) {
+      val uleLogConfig = LogConfig(Map(LogConfig.UncleanLeaderElectionEnableProp -> "true").asJava)
+      EasyMock.expect(mockZkClient.getLogConfigs(Set("t"), config.originals()))
+        .andReturn((Map(partition.topic -> uleLogConfig), Map.empty))
+    } else {
+      EasyMock.expect(mockZkClient.getLogConfigs(Set.empty, config.originals()))
+        .andReturn((Map(partition.topic -> LogConfig()), Map.empty))
+    }
+    val leaderAndIsrAfterElection = leaderAndIsr.newLeaderAndIsr(brokerId, List(brokerId))
+
     val updatedLeaderAndIsr = leaderAndIsrAfterElection.withZkVersion(2)
     EasyMock.expect(mockZkClient.updateLeaderAndIsr(Map(partition -> leaderAndIsrAfterElection), controllerEpoch, controllerContext.epochZkVersion))
-      .andReturn(UpdateLeaderAndIsrResult(Map(partition -> updatedLeaderAndIsr), Seq.empty, Map.empty))
+      .andReturn(UpdateLeaderAndIsrResult(Map(partition -> Right(updatedLeaderAndIsr)), Seq.empty))
     EasyMock.expect(mockControllerBrokerRequestBatch.addLeaderAndIsrRequestForBrokers(Seq(brokerId),
-      partition, LeaderIsrAndControllerEpoch(updatedLeaderAndIsr, controllerEpoch), Seq(brokerId), isNew = false))
+      partition, LeaderIsrAndControllerEpoch(updatedLeaderAndIsr, controllerEpoch), replicaAssignment(Seq(brokerId)), isNew = false))
     EasyMock.expect(mockControllerBrokerRequestBatch.sendRequestsToBrokers(controllerEpoch))
     EasyMock.replay(mockZkClient, mockControllerBrokerRequestBatch)
 
-    partitionStateMachine.handleStateChanges(partitions, OnlinePartition, Option(OfflinePartitionLeaderElectionStrategy))
+    val numUnCleanLeaderElectionsBefore = controllerContext.stats.uncleanLeaderElectionRate.count()
+    partitionStateMachine.handleStateChanges(
+      partitions,
+      OnlinePartition,
+      Option(OfflinePartitionLeaderElectionStrategy(false))
+    )
     EasyMock.verify(mockZkClient, mockControllerBrokerRequestBatch)
     assertEquals(OnlinePartition, partitionState(partition))
+    if (testUncleanLeaderElection)
+      assertEquals(1, controllerContext.stats.uncleanLeaderElectionRate.count() - numUnCleanLeaderElectionsBefore)
   }
 
   @Test
-  def testOfflinePartitionToOnlinePartitionTransitionZkUtilsExceptionFromStateLookup(): Unit = {
+  def testOfflinePartitionToUncleanOnlinePartitionTransition(): Unit = {
+    /* Starting scenario: Leader: X, Isr: [X], Replicas: [X, Y], LiveBrokers: [Y]
+     * Ending scenario: Leader: Y, Isr: [Y], Replicas: [X, Y], LiverBrokers: [Y]
+     *
+     * For the give staring scenario verify that performing an unclean leader
+     * election on the offline partition results on the first live broker getting
+     * elected.
+     */
+    val leaderBrokerId = brokerId + 1
+    controllerContext.setLiveBrokerAndEpochs(Map(TestUtils.createBrokerAndEpoch(brokerId, "host", 0)))
+    controllerContext.updatePartitionReplicaAssignment(partition, Seq(leaderBrokerId, brokerId))
+    controllerContext.putPartitionState(partition, OfflinePartition)
+
+    val leaderAndIsr = LeaderAndIsr(leaderBrokerId, List(leaderBrokerId))
+    val leaderIsrAndControllerEpoch = LeaderIsrAndControllerEpoch(leaderAndIsr, controllerEpoch)
+    controllerContext.partitionLeadershipInfo.put(partition, leaderIsrAndControllerEpoch)
+
+    EasyMock.expect(mockControllerBrokerRequestBatch.newBatch())
+    EasyMock
+      .expect(mockZkClient.getTopicPartitionStatesRaw(partitions))
+      .andReturn(
+        Seq(
+          GetDataResponse(
+            Code.OK,
+            null,
+            Option(partition),
+            TopicPartitionStateZNode.encode(leaderIsrAndControllerEpoch),
+            new Stat(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0),
+            ResponseMetadata(0, 0)
+          )
+        )
+      )
+
+    val leaderAndIsrAfterElection = leaderAndIsr.newLeaderAndIsr(brokerId, List(brokerId))
+    val updatedLeaderAndIsr = leaderAndIsrAfterElection.withZkVersion(2)
+
+    EasyMock
+      .expect(
+        mockZkClient.updateLeaderAndIsr(
+          Map(partition -> leaderAndIsrAfterElection),
+          controllerEpoch,
+          controllerContext.epochZkVersion
+        )
+      )
+      .andReturn(UpdateLeaderAndIsrResult(Map(partition -> Right(updatedLeaderAndIsr)), Seq.empty))
+    EasyMock.expect(
+      mockControllerBrokerRequestBatch.addLeaderAndIsrRequestForBrokers(
+        Seq(brokerId),
+        partition,
+        LeaderIsrAndControllerEpoch(updatedLeaderAndIsr, controllerEpoch),
+        replicaAssignment(Seq(leaderBrokerId, brokerId)),
+        false
+      )
+    )
+    EasyMock.expect(mockControllerBrokerRequestBatch.sendRequestsToBrokers(controllerEpoch))
+    EasyMock.replay(mockZkClient, mockControllerBrokerRequestBatch)
+
+    partitionStateMachine.handleStateChanges(
+      partitions,
+      OnlinePartition,
+      Option(OfflinePartitionLeaderElectionStrategy(true))
+    )
+    EasyMock.verify(mockZkClient, mockControllerBrokerRequestBatch)
+    assertEquals(OnlinePartition, partitionState(partition))
+  }
+
+  @Test
+  def testOfflinePartitionToOnlinePartitionTransitionZooKeeperClientExceptionFromStateLookup(): Unit = {
     controllerContext.setLiveBrokerAndEpochs(Map(TestUtils.createBrokerAndEpoch(brokerId, "host", 0)))
     controllerContext.updatePartitionReplicaAssignment(partition, Seq(brokerId))
-    partitionState.put(partition, OfflinePartition)
+    controllerContext.putPartitionState(partition, OfflinePartition)
     val leaderAndIsr = LeaderAndIsr(LeaderAndIsr.NoLeader, List(brokerId))
     val leaderIsrAndControllerEpoch = LeaderIsrAndControllerEpoch(leaderAndIsr, controllerEpoch)
     controllerContext.partitionLeadershipInfo.put(partition, leaderIsrAndControllerEpoch)
@@ -275,7 +379,11 @@ class PartitionStateMachineTest extends JUnitSuite {
     EasyMock.expect(mockControllerBrokerRequestBatch.sendRequestsToBrokers(controllerEpoch))
     EasyMock.replay(mockZkClient, mockControllerBrokerRequestBatch)
 
-    partitionStateMachine.handleStateChanges(partitions, OnlinePartition, Option(OfflinePartitionLeaderElectionStrategy))
+    partitionStateMachine.handleStateChanges(
+      partitions,
+      OnlinePartition,
+      Option(OfflinePartitionLeaderElectionStrategy(false))
+    )
     EasyMock.verify(mockZkClient, mockControllerBrokerRequestBatch)
     assertEquals(OfflinePartition, partitionState(partition))
   }
@@ -284,7 +392,7 @@ class PartitionStateMachineTest extends JUnitSuite {
   def testOfflinePartitionToOnlinePartitionTransitionErrorCodeFromStateLookup(): Unit = {
     controllerContext.setLiveBrokerAndEpochs(Map(TestUtils.createBrokerAndEpoch(brokerId, "host", 0)))
     controllerContext.updatePartitionReplicaAssignment(partition, Seq(brokerId))
-    partitionState.put(partition, OfflinePartition)
+    controllerContext.putPartitionState(partition, OfflinePartition)
     val leaderAndIsr = LeaderAndIsr(LeaderAndIsr.NoLeader, List(brokerId))
     val leaderIsrAndControllerEpoch = LeaderIsrAndControllerEpoch(leaderAndIsr, controllerEpoch)
     controllerContext.partitionLeadershipInfo.put(partition, leaderIsrAndControllerEpoch)
@@ -298,21 +406,25 @@ class PartitionStateMachineTest extends JUnitSuite {
     EasyMock.expect(mockControllerBrokerRequestBatch.sendRequestsToBrokers(controllerEpoch))
     EasyMock.replay(mockZkClient, mockControllerBrokerRequestBatch)
 
-    partitionStateMachine.handleStateChanges(partitions, OnlinePartition, Option(OfflinePartitionLeaderElectionStrategy))
+    partitionStateMachine.handleStateChanges(
+      partitions,
+      OnlinePartition,
+      Option(OfflinePartitionLeaderElectionStrategy(false))
+    )
     EasyMock.verify(mockZkClient, mockControllerBrokerRequestBatch)
     assertEquals(OfflinePartition, partitionState(partition))
   }
 
   @Test
   def testOfflinePartitionToNonexistentPartitionTransition(): Unit = {
-    partitionState.put(partition, OfflinePartition)
+    controllerContext.putPartitionState(partition, OfflinePartition)
     partitionStateMachine.handleStateChanges(partitions, NonExistentPartition)
     assertEquals(NonExistentPartition, partitionState(partition))
   }
 
   @Test
   def testInvalidOfflinePartitionToNewPartitionTransition(): Unit = {
-    partitionState.put(partition, OfflinePartition)
+    controllerContext.putPartitionState(partition, OfflinePartition)
     partitionStateMachine.handleStateChanges(partitions, NewPartition)
     assertEquals(OfflinePartition, partitionState(partition))
   }
@@ -330,18 +442,17 @@ class PartitionStateMachineTest extends JUnitSuite {
     prepareMockToGetTopicPartitionsStatesRaw()
 
     def prepareMockToGetLogConfigs(): Unit = {
-      val topicsForPartitionsWithNoLiveInSyncReplicas = Seq()
-      EasyMock.expect(mockZkClient.getLogConfigs(topicsForPartitionsWithNoLiveInSyncReplicas, config.originals()))
+      EasyMock.expect(mockZkClient.getLogConfigs(Set.empty, config.originals()))
         .andReturn(Map.empty, Map.empty)
     }
     prepareMockToGetLogConfigs()
 
     def prepareMockToUpdateLeaderAndIsr(): Unit = {
-      val updatedLeaderAndIsr = partitions.map { partition =>
+      val updatedLeaderAndIsr: Map[TopicPartition, LeaderAndIsr] = partitions.map { partition =>
         partition -> leaderAndIsr.newLeaderAndIsr(brokerId, List(brokerId))
       }.toMap
       EasyMock.expect(mockZkClient.updateLeaderAndIsr(updatedLeaderAndIsr, controllerEpoch, controllerContext.epochZkVersion))
-        .andReturn(UpdateLeaderAndIsrResult(updatedLeaderAndIsr, Seq.empty, Map.empty))
+        .andReturn(UpdateLeaderAndIsrResult(updatedLeaderAndIsr.mapValues(Right(_)).toMap, Seq.empty))
     }
     prepareMockToUpdateLeaderAndIsr()
   }
@@ -356,23 +467,21 @@ class PartitionStateMachineTest extends JUnitSuite {
 
     val partitionIds = Seq(0, 1, 2, 3)
     val topic = "test"
-    val partitions = partitionIds.map(new TopicPartition("test", _))
+    val partitions = partitionIds.map(new TopicPartition(topic, _))
 
     partitions.foreach { partition =>
       controllerContext.updatePartitionReplicaAssignment(partition, Seq(brokerId))
     }
 
-    EasyMock.expect(mockTopicDeletionManager.isTopicWithDeletionStarted(topic)).andReturn(false)
-    EasyMock.expectLastCall().anyTimes()
     prepareMockToElectLeaderForPartitions(partitions)
-    EasyMock.replay(mockZkClient, mockTopicDeletionManager)
+    EasyMock.replay(mockZkClient)
 
     partitionStateMachine.handleStateChanges(partitions, NewPartition)
     partitionStateMachine.handleStateChanges(partitions, OfflinePartition)
-    assertEquals(s"There should be ${partitions.size} offline partition(s)", partitions.size, partitionStateMachine.offlinePartitionCount)
+    assertEquals(s"There should be ${partitions.size} offline partition(s)", partitions.size, controllerContext.offlinePartitionCount)
 
-    partitionStateMachine.handleStateChanges(partitions, OnlinePartition, Some(OfflinePartitionLeaderElectionStrategy))
-    assertEquals(s"There should be no offline partition(s)", 0, partitionStateMachine.offlinePartitionCount)
+    partitionStateMachine.handleStateChanges(partitions, OnlinePartition, Some(OfflinePartitionLeaderElectionStrategy(false)))
+    assertEquals(s"There should be no offline partition(s)", 0, controllerContext.offlinePartitionCount)
   }
 
   /**
@@ -383,15 +492,14 @@ class PartitionStateMachineTest extends JUnitSuite {
   def testNoOfflinePartitionsChangeForTopicsBeingDeleted() = {
     val partitionIds = Seq(0, 1, 2, 3)
     val topic = "test"
-    val partitions = partitionIds.map(new TopicPartition("test", _))
+    val partitions = partitionIds.map(new TopicPartition(topic, _))
 
-    EasyMock.expect(mockTopicDeletionManager.isTopicWithDeletionStarted(topic)).andReturn(true)
-    EasyMock.expectLastCall().anyTimes()
-    EasyMock.replay(mockTopicDeletionManager)
+    controllerContext.topicsToBeDeleted.add(topic)
+    controllerContext.topicsWithDeletionStarted.add(topic)
 
     partitionStateMachine.handleStateChanges(partitions, NewPartition)
     partitionStateMachine.handleStateChanges(partitions, OfflinePartition)
-    assertEquals(s"There should be no offline partition(s)", 0, partitionStateMachine.offlinePartitionCount)
+    assertEquals(s"There should be no offline partition(s)", 0, controllerContext.offlinePartitionCount)
   }
 
   /**
@@ -411,52 +519,24 @@ class PartitionStateMachineTest extends JUnitSuite {
       controllerContext.updatePartitionReplicaAssignment(partition, Seq(brokerId))
     }
 
-    val props = TestUtils.createBrokerConfig(brokerId, "zkConnect")
-    props.put(KafkaConfig.DeleteTopicEnableProp, "true")
-
-    val customConfig = KafkaConfig.fromProps(props)
-
-    def createMockReplicaStateMachine() = {
-      val replicaStateMachine: ReplicaStateMachine = EasyMock.createMock(classOf[ReplicaStateMachine])
-      EasyMock.expect(replicaStateMachine.areAllReplicasForTopicDeleted(topic)).andReturn(false).anyTimes()
-      EasyMock.expect(replicaStateMachine.isAtLeastOneReplicaInDeletionStartedState(topic)).andReturn(false).anyTimes()
-      EasyMock.expect(replicaStateMachine.isAnyReplicaInState(topic, ReplicaDeletionIneligible)).andReturn(false).anyTimes()
-      EasyMock.expect(replicaStateMachine.replicasInState(topic, ReplicaDeletionIneligible)).andReturn(Set.empty).anyTimes()
-      EasyMock.expect(replicaStateMachine.replicasInState(topic, ReplicaDeletionStarted)).andReturn(Set.empty).anyTimes()
-      EasyMock.expect(replicaStateMachine.replicasInState(topic, ReplicaDeletionSuccessful)).andReturn(Set.empty).anyTimes()
-      EasyMock.expect(replicaStateMachine.handleStateChanges(EasyMock.anyObject[Seq[PartitionAndReplica]],
-        EasyMock.anyObject[ReplicaState], EasyMock.anyObject[Callbacks]))
-
-      EasyMock.expectLastCall().anyTimes()
-      replicaStateMachine
-    }
-    val replicaStateMachine = createMockReplicaStateMachine()
-    partitionStateMachine = new PartitionStateMachine(customConfig, new StateChangeLogger(brokerId, true, None), controllerContext,
-      mockZkClient, partitionState, mockControllerBrokerRequestBatch)
-
-    def createMockController() = {
-      val mockController: KafkaController = EasyMock.createMock(classOf[KafkaController])
-      EasyMock.expect(mockController.controllerContext).andReturn(controllerContext).anyTimes()
-      EasyMock.expect(mockController.config).andReturn(customConfig).anyTimes()
-      EasyMock.expect(mockController.partitionStateMachine).andReturn(partitionStateMachine).anyTimes()
-      EasyMock.expect(mockController.replicaStateMachine).andReturn(replicaStateMachine).anyTimes()
-      EasyMock.expect(mockController.sendUpdateMetadataRequest(Seq.empty, partitions.toSet))
-      EasyMock.expectLastCall().anyTimes()
-      mockController
-    }
-
-    val mockController = createMockController()
-    val mockEventManager: ControllerEventManager = EasyMock.createMock(classOf[ControllerEventManager])
-    EasyMock.replay(mockController, replicaStateMachine, mockEventManager)
-
-    val topicDeletionManager = new TopicDeletionManager(mockController, mockEventManager, mockZkClient)
-    partitionStateMachine.setTopicDeletionManager(topicDeletionManager)
+    val partitionStateMachine = new MockPartitionStateMachine(controllerContext, uncleanLeaderElectionEnabled = false)
+    val replicaStateMachine = new MockReplicaStateMachine(controllerContext)
+    val deletionClient = Mockito.mock(classOf[DeletionClient])
+    val topicDeletionManager = new TopicDeletionManager(config, controllerContext,
+      replicaStateMachine, partitionStateMachine, deletionClient)
 
     partitionStateMachine.handleStateChanges(partitions, NewPartition)
     partitionStateMachine.handleStateChanges(partitions, OfflinePartition)
-    assertEquals(s"There should be ${partitions.size} offline partition(s)", partitions.size, mockController.partitionStateMachine.offlinePartitionCount)
+    partitions.foreach { partition =>
+      val replica = PartitionAndReplica(partition, brokerId)
+      controllerContext.putReplicaState(replica, OfflineReplica)
+    }
 
+    assertEquals(s"There should be ${partitions.size} offline partition(s)", partitions.size, controllerContext.offlinePartitionCount)
     topicDeletionManager.enqueueTopicsForDeletion(Set(topic))
-    assertEquals(s"There should be no offline partition(s)", 0, partitionStateMachine.offlinePartitionCount)
+    assertEquals(s"There should be no offline partition(s)", 0, controllerContext.offlinePartitionCount)
   }
+
+  private def replicaAssignment(replicas: Seq[Int]): ReplicaAssignment = ReplicaAssignment(replicas, Seq(), Seq())
+
 }
diff --git a/core/src/test/scala/unit/kafka/controller/ReplicaStateMachineTest.scala b/core/src/test/scala/unit/kafka/controller/ReplicaStateMachineTest.scala
index ef274fa4fa787..7a76496f18e00 100644
--- a/core/src/test/scala/unit/kafka/controller/ReplicaStateMachineTest.scala
+++ b/core/src/test/scala/unit/kafka/controller/ReplicaStateMachineTest.scala
@@ -17,27 +17,25 @@
 package kafka.controller
 
 import kafka.api.LeaderAndIsr
+import kafka.cluster.{Broker, EndPoint}
 import kafka.server.KafkaConfig
 import kafka.utils.TestUtils
 import kafka.zk.KafkaZkClient.UpdateLeaderAndIsrResult
 import kafka.zk.{KafkaZkClient, TopicPartitionStateZNode}
 import kafka.zookeeper.{GetDataResponse, ResponseMetadata}
 import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.network.ListenerName
+import org.apache.kafka.common.security.auth.SecurityProtocol
 import org.apache.zookeeper.KeeperException.Code
 import org.apache.zookeeper.data.Stat
 import org.easymock.EasyMock
 import org.junit.Assert._
 import org.junit.{Before, Test}
-import org.scalatest.junit.JUnitSuite
 
-import scala.collection.mutable
-
-class ReplicaStateMachineTest extends JUnitSuite {
+class ReplicaStateMachineTest {
   private var controllerContext: ControllerContext = null
   private var mockZkClient: KafkaZkClient = null
   private var mockControllerBrokerRequestBatch: ControllerBrokerRequestBatch = null
-  private var mockTopicDeletionManager: TopicDeletionManager = null
-  private var replicaState: mutable.Map[PartitionAndReplica, ReplicaState] = null
   private var replicaStateMachine: ReplicaStateMachine = null
 
   private val brokerId = 5
@@ -54,10 +52,46 @@ class ReplicaStateMachineTest extends JUnitSuite {
     controllerContext.epoch = controllerEpoch
     mockZkClient = EasyMock.createMock(classOf[KafkaZkClient])
     mockControllerBrokerRequestBatch = EasyMock.createMock(classOf[ControllerBrokerRequestBatch])
-    mockTopicDeletionManager = EasyMock.createMock(classOf[TopicDeletionManager])
-    replicaState = mutable.Map.empty[PartitionAndReplica, ReplicaState]
-    replicaStateMachine = new ReplicaStateMachine(config, new StateChangeLogger(brokerId, true, None), controllerContext, mockTopicDeletionManager, mockZkClient,
-      replicaState, mockControllerBrokerRequestBatch)
+    replicaStateMachine = new ZkReplicaStateMachine(config, new StateChangeLogger(brokerId, true, None),
+      controllerContext, mockZkClient, mockControllerBrokerRequestBatch)
+  }
+
+  private def replicaState(replica: PartitionAndReplica): ReplicaState = {
+    controllerContext.replicaState(replica)
+  }
+
+  @Test
+  def testStartupOnlinePartition(): Unit = {
+    val endpoint1 = new EndPoint("localhost", 9997, new ListenerName("blah"),
+      SecurityProtocol.PLAINTEXT)
+    val liveBrokerEpochs = Map(Broker(brokerId, Seq(endpoint1), rack = None) -> 1L)
+    controllerContext.setLiveBrokerAndEpochs(liveBrokerEpochs)
+    controllerContext.updatePartitionReplicaAssignment(partition, Seq(brokerId))
+    assertEquals(None, controllerContext.replicaStates.get(replica))
+    replicaStateMachine.startup()
+    assertEquals(OnlineReplica, replicaState(replica))
+  }
+
+  @Test
+  def testStartupOfflinePartition(): Unit = {
+    controllerContext.updatePartitionReplicaAssignment(partition, Seq(brokerId))
+    assertEquals(None, controllerContext.replicaStates.get(replica))
+    replicaStateMachine.startup()
+    assertEquals(OfflineReplica, replicaState(replica))
+  }
+
+  @Test
+  def testStartupWithReplicaWithoutLeader(): Unit = {
+    val shutdownBrokerId = 100
+    val offlineReplica = PartitionAndReplica(partition, shutdownBrokerId)
+    val endpoint1 = new EndPoint("localhost", 9997, new ListenerName("blah"),
+      SecurityProtocol.PLAINTEXT)
+    val liveBrokerEpochs = Map(Broker(brokerId, Seq(endpoint1), rack = None) -> 1L)
+    controllerContext.setLiveBrokerAndEpochs(liveBrokerEpochs)
+    controllerContext.updatePartitionReplicaAssignment(partition, Seq(shutdownBrokerId))
+    assertEquals(None, controllerContext.replicaStates.get(offlineReplica))
+    replicaStateMachine.startup()
+    assertEquals(OfflineReplica, replicaState(offlineReplica))
   }
 
   @Test
@@ -103,7 +137,7 @@ class ReplicaStateMachineTest extends JUnitSuite {
 
   @Test
   def testNewReplicaToOnlineReplicaTransition(): Unit = {
-    replicaState.put(replica, NewReplica)
+    controllerContext.putReplicaState(replica, NewReplica)
     controllerContext.updatePartitionReplicaAssignment(partition, Seq(brokerId))
     replicaStateMachine.handleStateChanges(replicas, OnlineReplica)
     assertEquals(OnlineReplica, replicaState(replica))
@@ -111,11 +145,16 @@ class ReplicaStateMachineTest extends JUnitSuite {
 
   @Test
   def testNewReplicaToOfflineReplicaTransition(): Unit = {
-    replicaState.put(replica, NewReplica)
+    val endpoint1 = new EndPoint("localhost", 9997, new ListenerName("blah"),
+      SecurityProtocol.PLAINTEXT)
+    val liveBrokerEpochs = Map(Broker(brokerId, Seq(endpoint1), rack = None) -> 1L)
+    controllerContext.setLiveBrokerAndEpochs(liveBrokerEpochs)
+    controllerContext.putReplicaState(replica, NewReplica)
     EasyMock.expect(mockControllerBrokerRequestBatch.newBatch())
-    EasyMock.expect(mockControllerBrokerRequestBatch.addStopReplicaRequestForBrokers(EasyMock.eq(Seq(brokerId)),
-      EasyMock.eq(partition), EasyMock.eq(false), EasyMock.anyObject()))
+    EasyMock.expect(mockControllerBrokerRequestBatch.addStopReplicaRequestForBrokers(EasyMock.eq(Seq(brokerId)), EasyMock.eq(partition), EasyMock.eq(false)))
+    EasyMock.expect(mockControllerBrokerRequestBatch.addUpdateMetadataRequestForBrokers(EasyMock.eq(Seq(brokerId)),  EasyMock.eq(Set(partition))))
     EasyMock.expect(mockControllerBrokerRequestBatch.sendRequestsToBrokers(controllerEpoch))
+
     EasyMock.replay(mockControllerBrokerRequestBatch)
     replicaStateMachine.handleStateChanges(replicas, OfflineReplica)
     EasyMock.verify(mockControllerBrokerRequestBatch)
@@ -149,13 +188,13 @@ class ReplicaStateMachineTest extends JUnitSuite {
 
   @Test
   def testOnlineReplicaToOnlineReplicaTransition(): Unit = {
-    replicaState.put(replica, OnlineReplica)
+    controllerContext.putReplicaState(replica, OnlineReplica)
     controllerContext.updatePartitionReplicaAssignment(partition, Seq(brokerId))
     val leaderIsrAndControllerEpoch = LeaderIsrAndControllerEpoch(LeaderAndIsr(brokerId, List(brokerId)), controllerEpoch)
     controllerContext.partitionLeadershipInfo.put(partition, leaderIsrAndControllerEpoch)
     EasyMock.expect(mockControllerBrokerRequestBatch.newBatch())
     EasyMock.expect(mockControllerBrokerRequestBatch.addLeaderAndIsrRequestForBrokers(Seq(brokerId),
-      partition, leaderIsrAndControllerEpoch, Seq(brokerId), isNew = false))
+      partition, leaderIsrAndControllerEpoch, replicaAssignment(Seq(brokerId)), isNew = false))
     EasyMock.expect(mockControllerBrokerRequestBatch.sendRequestsToBrokers(controllerEpoch))
     EasyMock.replay(mockZkClient, mockControllerBrokerRequestBatch)
     replicaStateMachine.handleStateChanges(replicas, OnlineReplica)
@@ -167,7 +206,7 @@ class ReplicaStateMachineTest extends JUnitSuite {
   def testOnlineReplicaToOfflineReplicaTransition(): Unit = {
     val otherBrokerId = brokerId + 1
     val replicaIds = List(brokerId, otherBrokerId)
-    replicaState.put(replica, OnlineReplica)
+    controllerContext.putReplicaState(replica, OnlineReplica)
     controllerContext.updatePartitionReplicaAssignment(partition, replicaIds)
     val leaderAndIsr = LeaderAndIsr(brokerId, replicaIds)
     val leaderIsrAndControllerEpoch = LeaderIsrAndControllerEpoch(leaderAndIsr, controllerEpoch)
@@ -175,8 +214,7 @@ class ReplicaStateMachineTest extends JUnitSuite {
 
     val stat = new Stat(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0)
     EasyMock.expect(mockControllerBrokerRequestBatch.newBatch())
-    EasyMock.expect(mockControllerBrokerRequestBatch.addStopReplicaRequestForBrokers(EasyMock.eq(Seq(brokerId)),
-      EasyMock.eq(partition), EasyMock.eq(false), EasyMock.anyObject()))
+    EasyMock.expect(mockControllerBrokerRequestBatch.addStopReplicaRequestForBrokers(EasyMock.eq(Seq(brokerId)), EasyMock.eq(partition), EasyMock.eq(false)))
     val adjustedLeaderAndIsr = leaderAndIsr.newLeaderAndIsr(LeaderAndIsr.NoLeader, List(otherBrokerId))
     val updatedLeaderAndIsr = adjustedLeaderAndIsr.withZkVersion(adjustedLeaderAndIsr .zkVersion + 1)
     val updatedLeaderIsrAndControllerEpoch = LeaderIsrAndControllerEpoch(updatedLeaderAndIsr, controllerEpoch)
@@ -184,15 +222,14 @@ class ReplicaStateMachineTest extends JUnitSuite {
       Seq(GetDataResponse(Code.OK, null, Some(partition),
         TopicPartitionStateZNode.encode(leaderIsrAndControllerEpoch), stat, ResponseMetadata(0, 0))))
     EasyMock.expect(mockZkClient.updateLeaderAndIsr(Map(partition -> adjustedLeaderAndIsr), controllerEpoch, controllerContext.epochZkVersion))
-      .andReturn(UpdateLeaderAndIsrResult(Map(partition -> updatedLeaderAndIsr), Seq.empty, Map.empty))
-    EasyMock.expect(mockTopicDeletionManager.isTopicQueuedUpForDeletion(partition.topic)).andReturn(false)
+      .andReturn(UpdateLeaderAndIsrResult(Map(partition -> Right(updatedLeaderAndIsr)), Seq.empty))
     EasyMock.expect(mockControllerBrokerRequestBatch.addLeaderAndIsrRequestForBrokers(Seq(otherBrokerId),
-      partition, updatedLeaderIsrAndControllerEpoch, replicaIds, isNew = false))
+      partition, updatedLeaderIsrAndControllerEpoch, replicaAssignment(replicaIds), isNew = false))
     EasyMock.expect(mockControllerBrokerRequestBatch.sendRequestsToBrokers(controllerEpoch))
 
-    EasyMock.replay(mockZkClient, mockControllerBrokerRequestBatch, mockTopicDeletionManager)
+    EasyMock.replay(mockZkClient, mockControllerBrokerRequestBatch)
     replicaStateMachine.handleStateChanges(replicas, OfflineReplica)
-    EasyMock.verify(mockZkClient, mockControllerBrokerRequestBatch, mockTopicDeletionManager)
+    EasyMock.verify(mockZkClient, mockControllerBrokerRequestBatch)
     assertEquals(updatedLeaderIsrAndControllerEpoch, controllerContext.partitionLeadershipInfo(partition))
     assertEquals(OfflineReplica, replicaState(replica))
   }
@@ -224,13 +261,13 @@ class ReplicaStateMachineTest extends JUnitSuite {
 
   @Test
   def testOfflineReplicaToOnlineReplicaTransition(): Unit = {
-    replicaState.put(replica, OfflineReplica)
+    controllerContext.putReplicaState(replica, OfflineReplica)
     controllerContext.updatePartitionReplicaAssignment(partition, Seq(brokerId))
     val leaderIsrAndControllerEpoch = LeaderIsrAndControllerEpoch(LeaderAndIsr(brokerId, List(brokerId)), controllerEpoch)
     controllerContext.partitionLeadershipInfo.put(partition, leaderIsrAndControllerEpoch)
     EasyMock.expect(mockControllerBrokerRequestBatch.newBatch())
     EasyMock.expect(mockControllerBrokerRequestBatch.addLeaderAndIsrRequestForBrokers(Seq(brokerId),
-      partition, leaderIsrAndControllerEpoch, Seq(brokerId), isNew = false))
+      partition, leaderIsrAndControllerEpoch, replicaAssignment(Seq(brokerId)), isNew = false))
     EasyMock.expect(mockControllerBrokerRequestBatch.sendRequestsToBrokers(controllerEpoch))
     EasyMock.replay(mockZkClient, mockControllerBrokerRequestBatch)
     replicaStateMachine.handleStateChanges(replicas, OnlineReplica)
@@ -240,21 +277,21 @@ class ReplicaStateMachineTest extends JUnitSuite {
 
   @Test
   def testOfflineReplicaToReplicaDeletionStartedTransition(): Unit = {
-    val callbacks = new Callbacks()
-    replicaState.put(replica, OfflineReplica)
+    controllerContext.putReplicaState(replica, OfflineReplica)
     EasyMock.expect(mockControllerBrokerRequestBatch.newBatch())
-    EasyMock.expect(mockControllerBrokerRequestBatch.addStopReplicaRequestForBrokers(Seq(brokerId),
-      partition, true, callbacks.stopReplicaResponseCallback))
+    EasyMock.expect(mockControllerBrokerRequestBatch.addStopReplicaRequestForBrokers(Seq(brokerId), partition, true))
     EasyMock.expect(mockControllerBrokerRequestBatch.sendRequestsToBrokers(controllerEpoch))
     EasyMock.replay(mockZkClient, mockControllerBrokerRequestBatch)
-    replicaStateMachine.handleStateChanges(replicas, ReplicaDeletionStarted, callbacks)
+    replicaStateMachine.handleStateChanges(replicas, ReplicaDeletionStarted)
     EasyMock.verify(mockZkClient, mockControllerBrokerRequestBatch)
     assertEquals(ReplicaDeletionStarted, replicaState(replica))
   }
 
   @Test
-  def testInvalidOfflineReplicaToReplicaDeletionIneligibleTransition(): Unit = {
-    testInvalidTransition(OfflineReplica, ReplicaDeletionIneligible)
+  def testOfflineReplicaToReplicaDeletionIneligibleTransition(): Unit = {
+    controllerContext.putReplicaState(replica, OfflineReplica)
+    replicaStateMachine.handleStateChanges(replicas, ReplicaDeletionIneligible)
+    assertEquals(ReplicaDeletionIneligible, replicaState(replica))
   }
 
   @Test
@@ -284,25 +321,25 @@ class ReplicaStateMachineTest extends JUnitSuite {
 
   @Test
   def testReplicaDeletionStartedToReplicaDeletionIneligibleTransition(): Unit = {
-    replicaState.put(replica, ReplicaDeletionStarted)
+    controllerContext.putReplicaState(replica, ReplicaDeletionStarted)
     replicaStateMachine.handleStateChanges(replicas, ReplicaDeletionIneligible)
     assertEquals(ReplicaDeletionIneligible, replicaState(replica))
   }
 
   @Test
   def testReplicaDeletionStartedToReplicaDeletionSuccessfulTransition(): Unit = {
-    replicaState.put(replica, ReplicaDeletionStarted)
+    controllerContext.putReplicaState(replica, ReplicaDeletionStarted)
     replicaStateMachine.handleStateChanges(replicas, ReplicaDeletionSuccessful)
     assertEquals(ReplicaDeletionSuccessful, replicaState(replica))
   }
 
   @Test
   def testReplicaDeletionSuccessfulToNonexistentReplicaTransition(): Unit = {
-    replicaState.put(replica, ReplicaDeletionSuccessful)
+    controllerContext.putReplicaState(replica, ReplicaDeletionSuccessful)
     controllerContext.updatePartitionReplicaAssignment(partition, Seq(brokerId))
     replicaStateMachine.handleStateChanges(replicas, NonExistentReplica)
     assertEquals(Seq.empty, controllerContext.partitionReplicaAssignment(partition))
-    assertEquals(None, replicaState.get(replica))
+    assertEquals(None, controllerContext.replicaStates.get(replica))
   }
 
   @Test
@@ -342,13 +379,13 @@ class ReplicaStateMachineTest extends JUnitSuite {
 
   @Test
   def testReplicaDeletionIneligibleToOnlineReplicaTransition(): Unit = {
-    replicaState.put(replica, ReplicaDeletionIneligible)
+    controllerContext.putReplicaState(replica, ReplicaDeletionIneligible)
     controllerContext.updatePartitionReplicaAssignment(partition, Seq(brokerId))
     val leaderIsrAndControllerEpoch = LeaderIsrAndControllerEpoch(LeaderAndIsr(brokerId, List(brokerId)), controllerEpoch)
     controllerContext.partitionLeadershipInfo.put(partition, leaderIsrAndControllerEpoch)
     EasyMock.expect(mockControllerBrokerRequestBatch.newBatch())
     EasyMock.expect(mockControllerBrokerRequestBatch.addLeaderAndIsrRequestForBrokers(Seq(brokerId),
-      partition, leaderIsrAndControllerEpoch, Seq(brokerId), isNew = false))
+      partition, leaderIsrAndControllerEpoch, replicaAssignment(Seq(brokerId)), isNew = false))
     EasyMock.expect(mockControllerBrokerRequestBatch.sendRequestsToBrokers(controllerEpoch))
     EasyMock.replay(mockZkClient, mockControllerBrokerRequestBatch)
     replicaStateMachine.handleStateChanges(replicas, OnlineReplica)
@@ -367,8 +404,11 @@ class ReplicaStateMachineTest extends JUnitSuite {
   }
 
   private def testInvalidTransition(fromState: ReplicaState, toState: ReplicaState): Unit = {
-    replicaState.put(replica, fromState)
+    controllerContext.putReplicaState(replica, fromState)
     replicaStateMachine.handleStateChanges(replicas, toState)
     assertEquals(fromState, replicaState(replica))
   }
+
+  private def replicaAssignment(replicas: Seq[Int]): ReplicaAssignment = ReplicaAssignment(replicas, Seq(), Seq())
+
 }
diff --git a/core/src/test/scala/unit/kafka/controller/TopicDeletionManagerTest.scala b/core/src/test/scala/unit/kafka/controller/TopicDeletionManagerTest.scala
new file mode 100644
index 0000000000000..e3e1996cb646c
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/controller/TopicDeletionManagerTest.scala
@@ -0,0 +1,256 @@
+/*
+ * 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.controller
+
+import kafka.cluster.{Broker, EndPoint}
+import kafka.server.KafkaConfig
+import kafka.utils.TestUtils
+import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.network.ListenerName
+import org.apache.kafka.common.security.auth.SecurityProtocol
+import org.junit.Assert._
+import org.junit.Test
+import org.mockito.Mockito._
+
+class TopicDeletionManagerTest {
+
+  private val brokerId = 1
+  private val config = KafkaConfig.fromProps(TestUtils.createBrokerConfig(brokerId, "zkConnect"))
+  private val deletionClient = mock(classOf[DeletionClient])
+
+  @Test
+  def testInitialization(): Unit = {
+    val controllerContext = initContext(
+      brokers = Seq(1, 2, 3),
+      topics = Set("foo", "bar", "baz"),
+      numPartitions = 2,
+      replicationFactor = 3)
+
+    val replicaStateMachine = new MockReplicaStateMachine(controllerContext)
+    replicaStateMachine.startup()
+
+    val partitionStateMachine = new MockPartitionStateMachine(controllerContext, uncleanLeaderElectionEnabled = false)
+    partitionStateMachine.startup()
+
+    val deletionManager = new TopicDeletionManager(config, controllerContext, replicaStateMachine,
+      partitionStateMachine, deletionClient)
+
+    assertTrue(deletionManager.isDeleteTopicEnabled)
+    deletionManager.init(initialTopicsToBeDeleted = Set("foo", "bar"), initialTopicsIneligibleForDeletion = Set("bar", "baz"))
+
+    assertEquals(Set("foo", "bar"), controllerContext.topicsToBeDeleted.toSet)
+    assertEquals(Set("bar"), controllerContext.topicsIneligibleForDeletion.toSet)
+  }
+
+  @Test
+  def testBasicDeletion(): Unit = {
+    val controllerContext = initContext(
+      brokers = Seq(1, 2, 3),
+      topics = Set("foo", "bar"),
+      numPartitions = 2,
+      replicationFactor = 3)
+    val replicaStateMachine = new MockReplicaStateMachine(controllerContext)
+    replicaStateMachine.startup()
+
+    val partitionStateMachine = new MockPartitionStateMachine(controllerContext, uncleanLeaderElectionEnabled = false)
+    partitionStateMachine.startup()
+
+    val deletionManager = new TopicDeletionManager(config, controllerContext, replicaStateMachine,
+      partitionStateMachine, deletionClient)
+    assertTrue(deletionManager.isDeleteTopicEnabled)
+    deletionManager.init(Set.empty, Set.empty)
+
+    val fooPartitions = controllerContext.partitionsForTopic("foo")
+    val fooReplicas = controllerContext.replicasForPartition(fooPartitions).toSet
+
+    // Queue the topic for deletion
+    deletionManager.enqueueTopicsForDeletion(Set("foo"))
+
+    assertEquals(fooPartitions, controllerContext.partitionsInState("foo", NonExistentPartition))
+    assertEquals(fooReplicas, controllerContext.replicasInState("foo", ReplicaDeletionStarted))
+    verify(deletionClient).sendMetadataUpdate(fooPartitions)
+    assertEquals(Set("foo"), controllerContext.topicsToBeDeleted)
+    assertEquals(Set("foo"), controllerContext.topicsWithDeletionStarted)
+    assertEquals(Set(), controllerContext.topicsIneligibleForDeletion)
+
+    // Complete the deletion
+    deletionManager.completeReplicaDeletion(fooReplicas)
+
+    assertEquals(Set.empty, controllerContext.partitionsForTopic("foo"))
+    assertEquals(Set.empty[PartitionAndReplica], controllerContext.replicaStates.keySet.filter(_.topic == "foo"))
+    assertEquals(Set(), controllerContext.topicsToBeDeleted)
+    assertEquals(Set(), controllerContext.topicsWithDeletionStarted)
+    assertEquals(Set(), controllerContext.topicsIneligibleForDeletion)
+  }
+
+  @Test
+  def testDeletionWithBrokerOffline(): Unit = {
+    val controllerContext = initContext(
+      brokers = Seq(1, 2, 3),
+      topics = Set("foo", "bar"),
+      numPartitions = 2,
+      replicationFactor = 3)
+
+    val replicaStateMachine = new MockReplicaStateMachine(controllerContext)
+    replicaStateMachine.startup()
+
+    val partitionStateMachine = new MockPartitionStateMachine(controllerContext, uncleanLeaderElectionEnabled = false)
+    partitionStateMachine.startup()
+
+    val deletionManager = new TopicDeletionManager(config, controllerContext, replicaStateMachine,
+      partitionStateMachine, deletionClient)
+    assertTrue(deletionManager.isDeleteTopicEnabled)
+    deletionManager.init(Set.empty, Set.empty)
+
+    val fooPartitions = controllerContext.partitionsForTopic("foo")
+    val fooReplicas = controllerContext.replicasForPartition(fooPartitions).toSet
+
+    // Broker 2 is taken offline
+    val failedBrokerId = 2
+    val offlineBroker = controllerContext.liveOrShuttingDownBroker(failedBrokerId).get
+    val lastEpoch = controllerContext.liveBrokerIdAndEpochs(failedBrokerId)
+    controllerContext.removeLiveBrokers(Set(failedBrokerId))
+    assertEquals(Set(1, 3), controllerContext.liveBrokerIds)
+
+    val (offlineReplicas, onlineReplicas) = fooReplicas.partition(_.replica == failedBrokerId)
+    replicaStateMachine.handleStateChanges(offlineReplicas.toSeq, OfflineReplica)
+
+    // Start topic deletion
+    deletionManager.enqueueTopicsForDeletion(Set("foo"))
+    assertEquals(fooPartitions, controllerContext.partitionsInState("foo", NonExistentPartition))
+    verify(deletionClient).sendMetadataUpdate(fooPartitions)
+    assertEquals(onlineReplicas, controllerContext.replicasInState("foo", ReplicaDeletionStarted))
+    assertEquals(offlineReplicas, controllerContext.replicasInState("foo", ReplicaDeletionIneligible))
+
+    assertEquals(Set("foo"), controllerContext.topicsToBeDeleted)
+    assertEquals(Set("foo"), controllerContext.topicsWithDeletionStarted)
+    assertEquals(Set("foo"), controllerContext.topicsIneligibleForDeletion)
+
+    // Deletion succeeds for online replicas
+    deletionManager.completeReplicaDeletion(onlineReplicas)
+
+    assertEquals(fooPartitions, controllerContext.partitionsInState("foo", NonExistentPartition))
+    assertEquals(Set("foo"), controllerContext.topicsToBeDeleted)
+    assertEquals(Set("foo"), controllerContext.topicsWithDeletionStarted)
+    assertEquals(Set("foo"), controllerContext.topicsIneligibleForDeletion)
+    assertEquals(onlineReplicas, controllerContext.replicasInState("foo", ReplicaDeletionSuccessful))
+    assertEquals(offlineReplicas, controllerContext.replicasInState("foo", OfflineReplica))
+
+    // Broker 2 comes back online and deletion is resumed
+    controllerContext.addLiveBrokersAndEpochs(Map(offlineBroker -> (lastEpoch + 1L)))
+    deletionManager.resumeDeletionForTopics(Set("foo"))
+
+    assertEquals(onlineReplicas, controllerContext.replicasInState("foo", ReplicaDeletionSuccessful))
+    assertEquals(offlineReplicas, controllerContext.replicasInState("foo", ReplicaDeletionStarted))
+
+    deletionManager.completeReplicaDeletion(offlineReplicas)
+    assertEquals(Set.empty, controllerContext.partitionsForTopic("foo"))
+    assertEquals(Set.empty[PartitionAndReplica], controllerContext.replicaStates.keySet.filter(_.topic == "foo"))
+    assertEquals(Set(), controllerContext.topicsToBeDeleted)
+    assertEquals(Set(), controllerContext.topicsWithDeletionStarted)
+    assertEquals(Set(), controllerContext.topicsIneligibleForDeletion)
+  }
+
+  @Test
+  def testBrokerFailureAfterDeletionStarted(): Unit = {
+    val controllerContext = initContext(
+      brokers = Seq(1, 2, 3),
+      topics = Set("foo", "bar"),
+      numPartitions = 2,
+      replicationFactor = 3)
+
+    val replicaStateMachine = new MockReplicaStateMachine(controllerContext)
+    replicaStateMachine.startup()
+
+    val partitionStateMachine = new MockPartitionStateMachine(controllerContext, uncleanLeaderElectionEnabled = false)
+    partitionStateMachine.startup()
+
+    val deletionManager = new TopicDeletionManager(config, controllerContext, replicaStateMachine,
+      partitionStateMachine, deletionClient)
+    deletionManager.init(Set.empty, Set.empty)
+
+    val fooPartitions = controllerContext.partitionsForTopic("foo")
+    val fooReplicas = controllerContext.replicasForPartition(fooPartitions).toSet
+
+    // Queue the topic for deletion
+    deletionManager.enqueueTopicsForDeletion(Set("foo"))
+    assertEquals(fooPartitions, controllerContext.partitionsInState("foo", NonExistentPartition))
+    assertEquals(fooReplicas, controllerContext.replicasInState("foo", ReplicaDeletionStarted))
+
+    // Broker 2 fails
+    val failedBrokerId = 2
+    val offlineBroker = controllerContext.liveOrShuttingDownBroker(failedBrokerId).get
+    val lastEpoch = controllerContext.liveBrokerIdAndEpochs(failedBrokerId)
+    controllerContext.removeLiveBrokers(Set(failedBrokerId))
+    assertEquals(Set(1, 3), controllerContext.liveBrokerIds)
+    val (offlineReplicas, onlineReplicas) = fooReplicas.partition(_.replica == failedBrokerId)
+
+    // Fail replica deletion
+    deletionManager.failReplicaDeletion(offlineReplicas)
+    assertEquals(Set("foo"), controllerContext.topicsToBeDeleted)
+    assertEquals(Set("foo"), controllerContext.topicsWithDeletionStarted)
+    assertEquals(Set("foo"), controllerContext.topicsIneligibleForDeletion)
+    assertEquals(offlineReplicas, controllerContext.replicasInState("foo", ReplicaDeletionIneligible))
+    assertEquals(onlineReplicas, controllerContext.replicasInState("foo", ReplicaDeletionStarted))
+
+    // Broker 2 is restarted. The offline replicas remain ineligable
+    // (TODO: this is probably not desired)
+    controllerContext.addLiveBrokersAndEpochs(Map(offlineBroker -> (lastEpoch + 1L)))
+    deletionManager.resumeDeletionForTopics(Set("foo"))
+    assertEquals(Set("foo"), controllerContext.topicsToBeDeleted)
+    assertEquals(Set("foo"), controllerContext.topicsWithDeletionStarted)
+    assertEquals(Set(), controllerContext.topicsIneligibleForDeletion)
+    assertEquals(onlineReplicas, controllerContext.replicasInState("foo", ReplicaDeletionStarted))
+    assertEquals(offlineReplicas, controllerContext.replicasInState("foo", ReplicaDeletionIneligible))
+
+    // When deletion completes for the replicas which started, then deletion begins for the remaining ones
+    deletionManager.completeReplicaDeletion(onlineReplicas)
+    assertEquals(Set("foo"), controllerContext.topicsToBeDeleted)
+    assertEquals(Set("foo"), controllerContext.topicsWithDeletionStarted)
+    assertEquals(Set(), controllerContext.topicsIneligibleForDeletion)
+    assertEquals(onlineReplicas, controllerContext.replicasInState("foo", ReplicaDeletionSuccessful))
+    assertEquals(offlineReplicas, controllerContext.replicasInState("foo", ReplicaDeletionStarted))
+
+  }
+
+  def initContext(brokers: Seq[Int],
+                  topics: Set[String],
+                  numPartitions: Int,
+                  replicationFactor: Int): ControllerContext = {
+    val context = new ControllerContext
+    val brokerEpochs = brokers.map { brokerId =>
+      val endpoint = new EndPoint("localhost", 9900 + brokerId, new ListenerName("blah"),
+        SecurityProtocol.PLAINTEXT)
+      Broker(brokerId, Seq(endpoint), rack = None) -> 1L
+    }.toMap
+    context.setLiveBrokerAndEpochs(brokerEpochs)
+
+    // Simple round-robin replica assignment
+    var leaderIndex = 0
+    for (topic <- topics; partitionId <- 0 until numPartitions) {
+      val partition = new TopicPartition(topic, partitionId)
+      val replicas = (0 until replicationFactor).map { i =>
+        val replica = brokers((i + leaderIndex) % brokers.size)
+        replica
+      }
+      context.updatePartitionReplicaAssignment(partition, replicas)
+      leaderIndex += 1
+    }
+    context
+  }
+
+}
diff --git a/core/src/test/scala/unit/kafka/controller/TopicsSatisfyMinRFTest.scala b/core/src/test/scala/unit/kafka/controller/TopicsSatisfyMinRFTest.scala
new file mode 100644
index 0000000000000..2de930612bfba
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/controller/TopicsSatisfyMinRFTest.scala
@@ -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 unit.kafka.controller
+
+import kafka.controller.{KafkaController, ReplicaAssignment}
+import kafka.server.KafkaConfig
+import kafka.utils.TestUtils
+import kafka.zk.KafkaZkClient
+import org.apache.kafka.server.policy.CreateTopicPolicy
+import org.easymock.EasyMock
+import org.junit.{Before, Test}
+import org.junit.Assert.assertEquals
+import org.scalatest.junit.JUnitSuite
+
+class TopicsSatisfyMinRFTest extends JUnitSuite {
+  private var createTopicPolicy : Option[CreateTopicPolicy] = null
+  private var mockZkClient: KafkaZkClient = null
+  @Before
+  def setUp(): Unit = {
+    val config = TestUtils.createBrokerConfigs(1, "")
+      .map(prop => {
+        prop.setProperty(KafkaConfig.DefaultReplicationFactorProp,"2")  // Set min RF to be 2
+        prop.setProperty(KafkaConfig.CreateTopicPolicyClassNameProp, "kafka.server.LiCreateTopicPolicy")
+        prop
+      }).map(KafkaConfig.fromProps)
+    createTopicPolicy = Option(config.head.getConfiguredInstance(KafkaConfig.CreateTopicPolicyClassNameProp,
+      classOf[CreateTopicPolicy]))
+    mockZkClient = EasyMock.createMock(classOf[KafkaZkClient])
+  }
+
+  // Mock new topic with sufficient RF, should satisfy
+  @Test
+  def testNewTopicSucceed: Unit = {
+    val topicSucceed = "test_topic1"
+    assertEquals(true, KafkaController.satisfiesLiCreateTopicPolicy(createTopicPolicy, mockZkClient,
+      topicSucceed, Map(0 -> new ReplicaAssignment(Seq(0, 1), Seq.empty[Int], Seq.empty[Int]))))      // RF = 2
+  }
+
+  // Mock new topic with insufficient RF, should fail
+  @Test
+  def testNewTopicFail: Unit = {
+    val topicFail = "test_topic1"
+    EasyMock.expect(mockZkClient.getTopicPartitions(topicFail)).andReturn(Seq.empty)
+    EasyMock.replay(mockZkClient)
+    assertEquals(false, KafkaController.satisfiesLiCreateTopicPolicy(createTopicPolicy, mockZkClient,
+      topicFail, Map(0 -> new ReplicaAssignment(Seq(1), Seq.empty[Int], Seq.empty[Int]))))            // RF = 1
+    EasyMock.verify(mockZkClient)
+  }
+
+  // Mock existing topic with insufficient RF, because it's pre-existing, should still satisfy
+  @Test
+  def testExistingTopicSucceed: Unit = {
+    val topicExisting = "test_topic1"
+    EasyMock.expect(mockZkClient.getTopicPartitions(topicExisting)).andReturn(Seq("0"))  // Partition number = 1
+    EasyMock.replay(mockZkClient)
+    assertEquals(true, KafkaController.satisfiesLiCreateTopicPolicy(createTopicPolicy, mockZkClient,
+      topicExisting, Map(0 -> new ReplicaAssignment(Seq(1), Seq.empty[Int], Seq.empty[Int]))))        // RF = 1
+    EasyMock.verify(mockZkClient)
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/coordinator/AbstractCoordinatorConcurrencyTest.scala b/core/src/test/scala/unit/kafka/coordinator/AbstractCoordinatorConcurrencyTest.scala
index d5beceae8e574..c39a23f2b9686 100644
--- a/core/src/test/scala/unit/kafka/coordinator/AbstractCoordinatorConcurrencyTest.scala
+++ b/core/src/test/scala/unit/kafka/coordinator/AbstractCoordinatorConcurrencyTest.scala
@@ -23,7 +23,7 @@ import java.util.concurrent.atomic.AtomicInteger
 import java.util.concurrent.locks.Lock
 
 import kafka.coordinator.AbstractCoordinatorConcurrencyTest._
-import kafka.log.Log
+import kafka.log.{AppendOrigin, Log}
 import kafka.server._
 import kafka.utils._
 import kafka.utils.timer.MockTimer
@@ -52,7 +52,7 @@ abstract class AbstractCoordinatorConcurrencyTest[M <: CoordinatorMember] {
   val random = new Random
 
   @Before
-  def setUp() {
+  def setUp(): Unit = {
 
     replicaManager = EasyMock.partialMockBuilder(classOf[TestReplicaManager]).createMock()
     replicaManager.createDelayedProducePurgatory(timer)
@@ -61,7 +61,7 @@ abstract class AbstractCoordinatorConcurrencyTest[M <: CoordinatorMember] {
   }
 
   @After
-  def tearDown() {
+  def tearDown(): Unit = {
     EasyMock.reset(replicaManager)
     if (executor != null)
       executor.shutdownNow()
@@ -70,7 +70,7 @@ abstract class AbstractCoordinatorConcurrencyTest[M <: CoordinatorMember] {
   /**
     * Verify that concurrent operations run in the normal sequence produce the expected results.
     */
-  def verifyConcurrentOperations(createMembers: String => Set[M], operations: Seq[Operation]) {
+  def verifyConcurrentOperations(createMembers: String => Set[M], operations: Seq[Operation]): Unit = {
     OrderedOperationSequence(createMembers("verifyConcurrentOperations"), operations).run()
   }
 
@@ -78,7 +78,7 @@ abstract class AbstractCoordinatorConcurrencyTest[M <: CoordinatorMember] {
     * Verify that arbitrary operations run in some random sequence don't leave the coordinator
     * in a bad state. Operations in the normal sequence should continue to work as expected.
     */
-  def verifyConcurrentRandomSequences(createMembers: String => Set[M], operations: Seq[Operation]) {
+  def verifyConcurrentRandomSequences(createMembers: String => Set[M], operations: Seq[Operation]): Unit = {
     EasyMock.reset(replicaManager)
     for (i <- 0 to 10) {
       // Run some random operations
@@ -89,7 +89,7 @@ abstract class AbstractCoordinatorConcurrencyTest[M <: CoordinatorMember] {
     }
   }
 
-  def verifyConcurrentActions(actions: Set[Action]) {
+  def verifyConcurrentActions(actions: Set[Action]): Unit = {
     val futures = actions.map(executor.submit)
     futures.map(_.get)
     enableCompletion()
@@ -173,11 +173,11 @@ object AbstractCoordinatorConcurrencyTest {
     override def appendRecords(timeout: Long,
                                requiredAcks: Short,
                                internalTopicsAllowed: Boolean,
-                               isFromClient: Boolean,
+                               origin: AppendOrigin,
                                entriesPerPartition: Map[TopicPartition, MemoryRecords],
                                responseCallback: Map[TopicPartition, PartitionResponse] => Unit,
                                delayedProduceLock: Option[Lock] = None,
-                               processingStatsCallback: Map[TopicPartition, RecordConversionStats] => Unit = _ => ()) {
+                               processingStatsCallback: Map[TopicPartition, RecordConversionStats] => Unit = _ => ()): Unit = {
 
       if (entriesPerPartition.isEmpty)
         return
@@ -194,14 +194,14 @@ object AbstractCoordinatorConcurrencyTest {
           else
             false
         }
-        override def onComplete() {
+        override def onComplete(): Unit = {
           responseCallback(entriesPerPartition.map {
             case (tp, _) =>
               (tp, new PartitionResponse(Errors.NONE, 0L, RecordBatch.NO_TIMESTAMP, 0L))
           })
         }
       }
-      val producerRequestKeys = entriesPerPartition.keys.map(new TopicPartitionOperationKey(_)).toSeq
+      val producerRequestKeys = entriesPerPartition.keys.map(TopicPartitionOperationKey(_)).toSeq
       watchKeys ++= producerRequestKeys
       producePurgatory.tryCompleteElseWatch(delayedProduce, producerRequestKeys)
       tryCompleteDelayedRequests()
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 ba179a505b38d..207aab134d655 100644
--- a/core/src/test/scala/unit/kafka/coordinator/group/GroupCoordinatorConcurrencyTest.scala
+++ b/core/src/test/scala/unit/kafka/coordinator/group/GroupCoordinatorConcurrencyTest.scala
@@ -26,6 +26,8 @@ import kafka.coordinator.group.GroupCoordinatorConcurrencyTest._
 import kafka.server.{DelayedOperationPurgatory, KafkaConfig}
 import org.apache.kafka.common.TopicPartition
 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.requests.JoinGroupRequest
 import org.apache.kafka.common.utils.Time
@@ -66,7 +68,7 @@ class GroupCoordinatorConcurrencyTest extends AbstractCoordinatorConcurrencyTest
   var groupCoordinator: GroupCoordinator = _
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     super.setUp()
 
     EasyMock.expect(zkClient.getTopicPartitionCount(Topic.GROUP_METADATA_TOPIC_NAME))
@@ -83,12 +85,12 @@ class GroupCoordinatorConcurrencyTest extends AbstractCoordinatorConcurrencyTest
     val heartbeatPurgatory = new DelayedOperationPurgatory[DelayedHeartbeat]("Heartbeat", timer, config.brokerId, reaperEnabled = false)
     val joinPurgatory = new DelayedOperationPurgatory[DelayedJoin]("Rebalance", timer, config.brokerId, reaperEnabled = false)
 
-    groupCoordinator = GroupCoordinator(config, zkClient, replicaManager, heartbeatPurgatory, joinPurgatory, timer.time)
+    groupCoordinator = GroupCoordinator(config, zkClient, replicaManager, heartbeatPurgatory, joinPurgatory, timer.time, new Metrics())
     groupCoordinator.startup(false)
   }
 
   @After
-  override def tearDown() {
+  override def tearDown(): Unit = {
     try {
       if (groupCoordinator != null)
         groupCoordinator.shutdown()
@@ -104,17 +106,17 @@ class GroupCoordinatorConcurrencyTest extends AbstractCoordinatorConcurrencyTest
   }
 
   @Test
-  def testConcurrentGoodPathSequence() {
+  def testConcurrentGoodPathSequence(): Unit = {
     verifyConcurrentOperations(createGroupMembers, allOperations)
   }
 
   @Test
-  def testConcurrentTxnGoodPathSequence() {
+  def testConcurrentTxnGoodPathSequence(): Unit = {
     verifyConcurrentOperations(createGroupMembers, allOperationsWithTxn)
   }
 
   @Test
-  def testConcurrentRandomSequence() {
+  def testConcurrentRandomSequence(): Unit = {
     verifyConcurrentRandomSequences(createGroupMembers, allOperationsWithTxn)
   }
 
@@ -153,13 +155,13 @@ class GroupCoordinatorConcurrencyTest extends AbstractCoordinatorConcurrencyTest
   }
 
 
-  class JoinGroupOperation extends GroupOperation[JoinGroupResult, JoinGroupCallback] {
-    override def responseCallback(responsePromise: Promise[JoinGroupResult]): JoinGroupCallback = {
+  class JoinGroupOperation extends GroupOperation[JoinGroupCallbackParams, JoinGroupCallback] {
+    override def responseCallback(responsePromise: Promise[JoinGroupCallbackParams]): JoinGroupCallback = {
       val callback: JoinGroupCallback = responsePromise.success(_)
       callback
     }
     override def runWithCallback(member: GroupMember, responseCallback: JoinGroupCallback): Unit = {
-      groupCoordinator.handleJoinGroup(member.groupId, member.memberId, requireKnownMemberId = false, "clientId", "clientHost",
+      groupCoordinator.handleJoinGroup(member.groupId, member.memberId, None, requireKnownMemberId = false, "clientId", "clientHost",
        DefaultRebalanceTimeout, DefaultSessionTimeout,
        protocolType, protocols, responseCallback)
     }
@@ -173,17 +175,17 @@ class GroupCoordinatorConcurrencyTest extends AbstractCoordinatorConcurrencyTest
 
   class SyncGroupOperation extends GroupOperation[SyncGroupCallbackParams, SyncGroupCallback] {
     override def responseCallback(responsePromise: Promise[SyncGroupCallbackParams]): SyncGroupCallback = {
-      val callback: SyncGroupCallback = (assignment, error) =>
-        responsePromise.success((assignment, error))
+      val callback: SyncGroupCallback = syncGroupResult =>
+        responsePromise.success(syncGroupResult.memberAssignment, syncGroupResult.error)
       callback
     }
     override def runWithCallback(member: GroupMember, responseCallback: SyncGroupCallback): Unit = {
       if (member.leader) {
         groupCoordinator.handleSyncGroup(member.groupId, member.generationId, member.memberId,
-            member.group.assignment, responseCallback)
+          member.groupInstanceId, member.group.assignment, responseCallback)
       } else {
-         groupCoordinator.handleSyncGroup(member.groupId, member.generationId, member.memberId,
-             Map.empty[String, Array[Byte]], responseCallback)
+        groupCoordinator.handleSyncGroup(member.groupId, member.generationId, member.memberId,
+          member.groupInstanceId, Map.empty[String, Array[Byte]], responseCallback)
       }
     }
     override def awaitAndVerify(member: GroupMember): Unit = {
@@ -198,7 +200,8 @@ class GroupCoordinatorConcurrencyTest extends AbstractCoordinatorConcurrencyTest
       callback
     }
     override def runWithCallback(member: GroupMember, responseCallback: HeartbeatCallback): Unit = {
-      groupCoordinator.handleHeartbeat( member.groupId, member.memberId,  member.generationId, responseCallback)
+      groupCoordinator.handleHeartbeat(member.groupId, member.memberId,
+        member.groupInstanceId, member.generationId, responseCallback)
     }
     override def awaitAndVerify(member: GroupMember): Unit = {
        val error = await(member, DefaultSessionTimeout)
@@ -213,8 +216,8 @@ class GroupCoordinatorConcurrencyTest extends AbstractCoordinatorConcurrencyTest
     override def runWithCallback(member: GroupMember, responseCallback: CommitOffsetCallback): Unit = {
       val tp = new TopicPartition("topic", 0)
       val offsets = immutable.Map(tp -> OffsetAndMetadata(1, "", Time.SYSTEM.milliseconds()))
-      groupCoordinator.handleCommitOffsets(member.groupId, member.memberId, member.generationId,
-          offsets, responseCallback)
+      groupCoordinator.handleCommitOffsets(member.groupId, member.memberId,
+        member.groupInstanceId, member.generationId, offsets, responseCallback)
     }
     override def awaitAndVerify(member: GroupMember): Unit = {
        val offsets = await(member, 500)
@@ -262,30 +265,37 @@ class GroupCoordinatorConcurrencyTest extends AbstractCoordinatorConcurrencyTest
 
   class LeaveGroupOperation extends GroupOperation[LeaveGroupCallbackParams, LeaveGroupCallback] {
     override def responseCallback(responsePromise: Promise[LeaveGroupCallbackParams]): LeaveGroupCallback = {
-      val callback: LeaveGroupCallback = error => responsePromise.success(error)
+      val callback: LeaveGroupCallback = result => responsePromise.success(result)
       callback
     }
     override def runWithCallback(member: GroupMember, responseCallback: LeaveGroupCallback): Unit = {
-      groupCoordinator.handleLeaveGroup(member.group.groupId, member.memberId, responseCallback)
+      val memberIdentity = new MemberIdentity()
+          .setMemberId(member.memberId)
+      groupCoordinator.handleLeaveGroup(member.group.groupId, List(memberIdentity), responseCallback)
     }
     override def awaitAndVerify(member: GroupMember): Unit = {
-       val error = await(member, DefaultSessionTimeout)
-       assertEquals(Errors.NONE, error)
+      val leaveGroupResult = await(member, DefaultSessionTimeout)
+
+      val memberResponses = leaveGroupResult.memberResponses
+      GroupCoordinatorTest.verifyLeaveGroupResult(leaveGroupResult, Errors.NONE, List(Errors.NONE))
+      assertEquals(member.memberId, memberResponses.head.memberId)
+      assertEquals(None, memberResponses.head.groupInstanceId)
     }
   }
 }
 
 object GroupCoordinatorConcurrencyTest {
 
+  type JoinGroupCallbackParams = JoinGroupResult
   type JoinGroupCallback = JoinGroupResult => Unit
   type SyncGroupCallbackParams = (Array[Byte], Errors)
-  type SyncGroupCallback = (Array[Byte], Errors) => Unit
+  type SyncGroupCallback = SyncGroupResult => Unit
   type HeartbeatCallbackParams = Errors
   type HeartbeatCallback = Errors => Unit
   type CommitOffsetCallbackParams = Map[TopicPartition, Errors]
   type CommitOffsetCallback = Map[TopicPartition, Errors] => Unit
-  type LeaveGroupCallbackParams = Errors
-  type LeaveGroupCallback = Errors => Unit
+  type LeaveGroupCallbackParams = LeaveGroupResult
+  type LeaveGroupCallback = LeaveGroupResult => Unit
   type CompleteTxnCallbackParams = Errors
   type CompleteTxnCallback = Errors => Unit
 
@@ -307,6 +317,7 @@ object GroupCoordinatorConcurrencyTest {
 
   class GroupMember(val group: Group, val groupPartitionId: Int, val leader: Boolean) extends CoordinatorMember {
     @volatile var memberId: String = JoinGroupRequest.UNKNOWN_MEMBER_ID
+    @volatile var groupInstanceId: Option[String] = None
     @volatile var generationId: Int = -1
     def groupId: String = group.groupId
   }
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 a529750566b59..4d6d694f6d7f3 100644
--- a/core/src/test/scala/unit/kafka/coordinator/group/GroupCoordinatorTest.scala
+++ b/core/src/test/scala/unit/kafka/coordinator/group/GroupCoordinatorTest.scala
@@ -20,7 +20,7 @@ package kafka.coordinator.group
 import java.util.Optional
 
 import kafka.common.OffsetAndMetadata
-import kafka.server.{DelayedOperationPurgatory, KafkaConfig, ReplicaManager}
+import kafka.server.{DelayedOperationPurgatory, HostedPartition, KafkaConfig, ReplicaManager}
 import kafka.utils._
 import kafka.utils.timer.MockTimer
 import org.apache.kafka.common.TopicPartition
@@ -33,33 +33,40 @@ import java.util.concurrent.TimeUnit
 import java.util.concurrent.locks.ReentrantLock
 
 import kafka.cluster.Partition
+import kafka.log.AppendOrigin
 import kafka.zk.KafkaZkClient
+import org.apache.kafka.clients.consumer.ConsumerPartitionAssignor.Subscription
+import org.apache.kafka.clients.consumer.internals.ConsumerProtocol
 import org.apache.kafka.common.internals.Topic
+import org.apache.kafka.common.metrics.Metrics
+import org.apache.kafka.common.message.LeaveGroupRequestData.MemberIdentity
 import org.junit.Assert._
 import org.junit.{After, Assert, Before, Test}
-import org.scalatest.junit.JUnitSuite
+import org.scalatest.Assertions.intercept
 
-import scala.collection.mutable
+import scala.collection.JavaConverters._
+import scala.collection.{Seq, mutable}
 import scala.collection.mutable.ArrayBuffer
 import scala.concurrent.duration.Duration
 import scala.concurrent.{Await, Future, Promise, TimeoutException}
 
-class GroupCoordinatorTest extends JUnitSuite {
+class GroupCoordinatorTest {
+  import GroupCoordinatorTest._
+
   type JoinGroupCallback = JoinGroupResult => Unit
   type SyncGroupCallbackParams = (Array[Byte], Errors)
-  type SyncGroupCallback = (Array[Byte], Errors) => Unit
+  type SyncGroupCallback = SyncGroupResult => Unit
   type HeartbeatCallbackParams = Errors
   type HeartbeatCallback = Errors => Unit
   type CommitOffsetCallbackParams = Map[TopicPartition, Errors]
   type CommitOffsetCallback = Map[TopicPartition, Errors] => Unit
-  type LeaveGroupCallbackParams = Errors
-  type LeaveGroupCallback = Errors => Unit
+  type LeaveGroupCallback = LeaveGroupResult => Unit
 
   val ClientId = "consumer-test"
   val ClientHost = "localhost"
   val GroupMinSessionTimeout = 10
   val GroupMaxSessionTimeout = 10 * 60 * 1000
-  val GroupMaxSize = 3
+  val GroupMaxSize = 4
   val DefaultRebalanceTimeout = 500
   val DefaultSessionTimeout = 500
   val GroupInitialRebalanceDelay = 50
@@ -72,15 +79,20 @@ class GroupCoordinatorTest extends JUnitSuite {
   private val groupId = "groupId"
   private val protocolType = "consumer"
   private val memberId = "memberId"
+  private val groupInstanceId = Some("groupInstanceId")
+  private val leaderInstanceId = Some("leader")
+  private val followerInstanceId = Some("follower")
+  private val invalidMemberId = "invalidMember"
   private val metadata = Array[Byte]()
   private val protocols = List(("range", metadata))
+  private val protocolSuperset = List(("range", metadata), ("roundrobin", metadata))
   private var groupPartitionId: Int = -1
 
   // we use this string value since its hashcode % #.partitions is different
   private val otherGroupId = "otherGroup"
 
   @Before
-  def setUp() {
+  def setUp(): Unit = {
     val props = TestUtils.createBrokerConfig(nodeId = 0, zkConnect = "")
     props.setProperty(KafkaConfig.GroupMinSessionTimeoutMsProp, GroupMinSessionTimeout.toString)
     props.setProperty(KafkaConfig.GroupMaxSessionTimeoutMsProp, GroupMaxSessionTimeout.toString)
@@ -104,7 +116,7 @@ class GroupCoordinatorTest extends JUnitSuite {
     val heartbeatPurgatory = new DelayedOperationPurgatory[DelayedHeartbeat]("Heartbeat", timer, config.brokerId, reaperEnabled = false)
     val joinPurgatory = new DelayedOperationPurgatory[DelayedJoin]("Rebalance", timer, config.brokerId, reaperEnabled = false)
 
-    groupCoordinator = GroupCoordinator(config, zkClient, replicaManager, heartbeatPurgatory, joinPurgatory, timer.time)
+    groupCoordinator = GroupCoordinator(config, zkClient, replicaManager, heartbeatPurgatory, joinPurgatory, timer.time, new Metrics())
     groupCoordinator.startup(enableMetadataExpiration = false)
 
     // add the partition into the owned partition list
@@ -113,7 +125,7 @@ class GroupCoordinatorTest extends JUnitSuite {
   }
 
   @After
-  def tearDown() {
+  def tearDown(): Unit = {
     EasyMock.reset(replicaManager)
     if (groupCoordinator != null)
       groupCoordinator.shutdown()
@@ -127,28 +139,33 @@ class GroupCoordinatorTest extends JUnitSuite {
     groupCoordinator.groupManager.addLoadingPartition(otherGroupPartitionId)
     assertTrue(groupCoordinator.groupManager.isGroupLoading(otherGroupId))
 
-    // JoinGroup
+    // Dynamic Member JoinGroup
     var joinGroupResponse: Option[JoinGroupResult] = None
-    groupCoordinator.handleJoinGroup(otherGroupId, memberId, true, "clientId", "clientHost", 60000, 10000, "consumer",
+    groupCoordinator.handleJoinGroup(otherGroupId, memberId, None, true, "clientId", "clientHost", 60000, 10000, "consumer",
+      List("range" -> new Array[Byte](0)), result => { joinGroupResponse = Some(result)})
+    assertEquals(Some(Errors.COORDINATOR_LOAD_IN_PROGRESS), joinGroupResponse.map(_.error))
+
+    // Static Member JoinGroup
+    groupCoordinator.handleJoinGroup(otherGroupId, memberId, Some("groupInstanceId"), false, "clientId", "clientHost", 60000, 10000, "consumer",
       List("range" -> new Array[Byte](0)), result => { joinGroupResponse = Some(result)})
     assertEquals(Some(Errors.COORDINATOR_LOAD_IN_PROGRESS), joinGroupResponse.map(_.error))
 
     // SyncGroup
     var syncGroupResponse: Option[Errors] = None
-    groupCoordinator.handleSyncGroup(otherGroupId, 1, memberId, Map.empty[String, Array[Byte]],
-      (_, error)=> syncGroupResponse = Some(error))
+    groupCoordinator.handleSyncGroup(otherGroupId, 1, memberId, None, Map.empty[String, Array[Byte]],
+      syncGroupResult => syncGroupResponse = Some(syncGroupResult.error))
     assertEquals(Some(Errors.REBALANCE_IN_PROGRESS), syncGroupResponse)
 
     // OffsetCommit
     val topicPartition = new TopicPartition("foo", 0)
     var offsetCommitErrors = Map.empty[TopicPartition, Errors]
-    groupCoordinator.handleCommitOffsets(otherGroupId, memberId, 1,
+    groupCoordinator.handleCommitOffsets(otherGroupId, memberId, None, 1,
       Map(topicPartition -> offsetAndMetadata(15L)), result => { offsetCommitErrors = result })
     assertEquals(Some(Errors.COORDINATOR_LOAD_IN_PROGRESS), offsetCommitErrors.get(topicPartition))
 
     // Heartbeat
     var heartbeatError: Option[Errors] = None
-    groupCoordinator.handleHeartbeat(otherGroupId, memberId, 1, error => { heartbeatError = Some(error) })
+    groupCoordinator.handleHeartbeat(otherGroupId, memberId, None, 1, error => { heartbeatError = Some(error) })
     assertEquals(Some(Errors.NONE), heartbeatError)
 
     // DescribeGroups
@@ -176,7 +193,7 @@ class GroupCoordinatorTest extends JUnitSuite {
   }
 
   @Test
-  def testOffsetsRetentionMsIntegerOverflow() {
+  def testOffsetsRetentionMsIntegerOverflow(): Unit = {
     val props = TestUtils.createBrokerConfig(nodeId = 0, zkConnect = "")
     props.setProperty(KafkaConfig.OffsetsRetentionMinutesProp, Integer.MAX_VALUE.toString)
     val config = KafkaConfig.fromProps(props)
@@ -185,21 +202,24 @@ class GroupCoordinatorTest extends JUnitSuite {
   }
 
   @Test
-  def testJoinGroupWrongCoordinator() {
+  def testJoinGroupWrongCoordinator(): Unit = {
     val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
 
-    val joinGroupResult = joinGroup(otherGroupId, memberId, protocolType, protocols)
-    val joinGroupError = joinGroupResult.error
-    assertEquals(Errors.NOT_COORDINATOR, joinGroupError)
+    var joinGroupResult = dynamicJoinGroup(otherGroupId, memberId, protocolType, protocols)
+    assertEquals(Errors.NOT_COORDINATOR, joinGroupResult.error)
+
+    EasyMock.reset(replicaManager)
+    joinGroupResult = staticJoinGroup(otherGroupId, memberId, groupInstanceId, protocolType, protocols)
+    assertEquals(Errors.NOT_COORDINATOR, joinGroupResult.error)
   }
 
   @Test
-  def testJoinGroupShouldReceiveErrorIfGroupOverMaxSize() {
-    var futures = ArrayBuffer[Future[JoinGroupResult]]()
+  def testJoinGroupShouldReceiveErrorIfGroupOverMaxSize(): Unit = {
+    val futures = ArrayBuffer[Future[JoinGroupResult]]()
     val rebalanceTimeout = GroupInitialRebalanceDelay * 2
 
     for (i <- 1.to(GroupMaxSize)) {
-      futures += sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols, rebalanceTimeout)
+      futures += sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols, rebalanceTimeout = rebalanceTimeout)
       if (i != 1)
         timer.advanceClock(GroupInitialRebalanceDelay)
       EasyMock.reset(replicaManager)
@@ -211,59 +231,62 @@ class GroupCoordinatorTest extends JUnitSuite {
     }
 
     // Should receive an error since the group is full
-    val errorFuture = sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols, rebalanceTimeout)
+    val errorFuture = sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols, rebalanceTimeout = rebalanceTimeout)
     assertEquals(Errors.GROUP_MAX_SIZE_REACHED, await(errorFuture, 1).error)
   }
 
   @Test
-  def testJoinGroupSessionTimeoutTooSmall() {
+  def testJoinGroupSessionTimeoutTooSmall(): Unit = {
     val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
 
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols, sessionTimeout = GroupMinSessionTimeout - 1)
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols, sessionTimeout = GroupMinSessionTimeout - 1)
     val joinGroupError = joinGroupResult.error
     assertEquals(Errors.INVALID_SESSION_TIMEOUT, joinGroupError)
   }
 
   @Test
-  def testJoinGroupSessionTimeoutTooLarge() {
+  def testJoinGroupSessionTimeoutTooLarge(): Unit = {
     val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
 
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols, sessionTimeout = GroupMaxSessionTimeout + 1)
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols, sessionTimeout = GroupMaxSessionTimeout + 1)
     val joinGroupError = joinGroupResult.error
     assertEquals(Errors.INVALID_SESSION_TIMEOUT, joinGroupError)
   }
 
   @Test
-  def testJoinGroupUnknownConsumerNewGroup() {
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols)
-    val joinGroupError = joinGroupResult.error
-    assertEquals(Errors.UNKNOWN_MEMBER_ID, joinGroupError)
+  def testJoinGroupUnknownConsumerNewGroup(): Unit = {
+    var joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols)
+    assertEquals(Errors.UNKNOWN_MEMBER_ID, joinGroupResult.error)
+
+    EasyMock.reset(replicaManager)
+    joinGroupResult = staticJoinGroup(groupId, memberId, groupInstanceId, protocolType, protocols)
+    assertEquals(Errors.UNKNOWN_MEMBER_ID, joinGroupResult.error)
   }
 
   @Test
-  def testInvalidGroupId() {
+  def testInvalidGroupId(): Unit = {
     val groupId = ""
     val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
 
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols)
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols)
     assertEquals(Errors.INVALID_GROUP_ID, joinGroupResult.error)
   }
 
   @Test
-  def testValidJoinGroup() {
+  def testValidJoinGroup(): Unit = {
     val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
 
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols)
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols)
     val joinGroupError = joinGroupResult.error
     assertEquals(Errors.NONE, joinGroupError)
   }
 
   @Test
-  def testJoinGroupInconsistentProtocolType() {
+  def testJoinGroupInconsistentProtocolType(): Unit = {
     val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
     val otherMemberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
 
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols)
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols)
     assertEquals(Errors.NONE, joinGroupResult.error)
 
     EasyMock.reset(replicaManager)
@@ -272,21 +295,55 @@ class GroupCoordinatorTest extends JUnitSuite {
   }
 
   @Test
-  def testJoinGroupWithEmptyProtocolType() {
+  def testJoinGroupWithEmptyProtocolType(): Unit = {
     val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
 
-    val joinGroupResult = joinGroup(groupId, memberId, "", protocols)
+    var joinGroupResult = dynamicJoinGroup(groupId, memberId, "", protocols)
+    assertEquals(Errors.INCONSISTENT_GROUP_PROTOCOL, joinGroupResult.error)
+
+    EasyMock.reset(replicaManager)
+    joinGroupResult = staticJoinGroup(groupId, memberId, groupInstanceId, "", protocols)
     assertEquals(Errors.INCONSISTENT_GROUP_PROTOCOL, joinGroupResult.error)
   }
 
   @Test
-  def testJoinGroupWithEmptyGroupProtocol() {
+  def testJoinGroupWithEmptyGroupProtocol(): Unit = {
     val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
 
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, List())
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, List())
     assertEquals(Errors.INCONSISTENT_GROUP_PROTOCOL, joinGroupResult.error)
   }
 
+  @Test
+  def testNewMemberTimeoutCompletion(): Unit = {
+    val sessionTimeout = GroupCoordinator.NewMemberJoinTimeoutMs + 5000
+    val responseFuture = sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols, None, sessionTimeout, DefaultRebalanceTimeout, false)
+
+    timer.advanceClock(GroupInitialRebalanceDelay + 1)
+
+    val joinResult = Await.result(responseFuture, Duration(DefaultRebalanceTimeout + 100, TimeUnit.MILLISECONDS))
+    val group = groupCoordinator.groupManager.getGroup(groupId).get
+    val memberId = joinResult.memberId
+
+    assertEquals(Errors.NONE, joinResult.error)
+    assertEquals(0, group.allMemberMetadata.count(_.isNew))
+
+    EasyMock.reset(replicaManager)
+    val syncGroupResult = syncGroupLeader(groupId, joinResult.generationId, memberId, Map(memberId -> Array[Byte]()))
+    val syncGroupError = syncGroupResult._2
+
+    assertEquals(Errors.NONE, syncGroupError)
+    assertEquals(1, group.size)
+
+    timer.advanceClock(GroupCoordinator.NewMemberJoinTimeoutMs + 100)
+
+    // Make sure the NewMemberTimeout is not still in effect, and the member is not kicked
+    assertEquals(1, group.size)
+
+    timer.advanceClock(sessionTimeout + 100)
+    assertEquals(0, group.size)
+  }
+
   @Test
   def testNewMemberJoinExpiration(): Unit = {
     // This tests new member expiration during a protracted rebalance. We first create a
@@ -297,7 +354,7 @@ class GroupCoordinatorTest extends JUnitSuite {
     val sessionTimeout = GroupCoordinator.NewMemberJoinTimeoutMs + 5000
     val rebalanceTimeout = GroupCoordinator.NewMemberJoinTimeoutMs * 2
 
-    val firstJoinResult = joinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols,
+    val firstJoinResult = dynamicJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols,
       sessionTimeout, rebalanceTimeout)
     val firstMemberId = firstJoinResult.memberId
     assertEquals(firstMemberId, firstJoinResult.leaderId)
@@ -309,8 +366,7 @@ class GroupCoordinatorTest extends JUnitSuite {
     assertEquals(0, group.allMemberMetadata.count(_.isNew))
 
     EasyMock.reset(replicaManager)
-
-    val responseFuture = sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols, sessionTimeout, rebalanceTimeout)
+    val responseFuture = sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols, None, sessionTimeout, rebalanceTimeout)
     assertFalse(responseFuture.isCompleted)
 
     assertEquals(2, group.allMembers.size)
@@ -330,7 +386,96 @@ class GroupCoordinatorTest extends JUnitSuite {
   }
 
   @Test
-  def testJoinGroupInconsistentGroupProtocol() {
+  def testNewMemberFailureAfterJoinGroupCompletion(): Unit = {
+    // For old versions of the JoinGroup protocol, new members were subject
+    // to expiration if the rebalance took long enough. This test case ensures
+    // that following completion of the JoinGroup phase, new members follow
+    // normal heartbeat expiration logic.
+
+    val firstJoinResult = dynamicJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
+    val firstMemberId = firstJoinResult.memberId
+    val firstGenerationId = firstJoinResult.generationId
+    assertEquals(firstMemberId, firstJoinResult.leaderId)
+    assertEquals(Errors.NONE, firstJoinResult.error)
+
+    EasyMock.reset(replicaManager)
+    val firstSyncResult = syncGroupLeader(groupId, firstGenerationId, firstMemberId,
+      Map(firstMemberId -> Array[Byte]()))
+    assertEquals(Errors.NONE, firstSyncResult._2)
+
+    EasyMock.reset(replicaManager)
+    val otherJoinFuture = sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
+
+    EasyMock.reset(replicaManager)
+    val joinFuture = sendJoinGroup(groupId, firstMemberId, protocolType, protocols,
+      requireKnownMemberId = false)
+
+    val joinResult = await(joinFuture, DefaultSessionTimeout+100)
+    val otherJoinResult = await(otherJoinFuture, DefaultSessionTimeout+100)
+    assertEquals(Errors.NONE, joinResult.error)
+    assertEquals(Errors.NONE, otherJoinResult.error)
+
+    verifySessionExpiration(groupId)
+  }
+
+  @Test
+  def testNewMemberFailureAfterSyncGroupCompletion(): Unit = {
+    // For old versions of the JoinGroup protocol, new members were subject
+    // to expiration if the rebalance took long enough. This test case ensures
+    // that following completion of the SyncGroup phase, new members follow
+    // normal heartbeat expiration logic.
+
+    val firstJoinResult = dynamicJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
+    val firstMemberId = firstJoinResult.memberId
+    val firstGenerationId = firstJoinResult.generationId
+    assertEquals(firstMemberId, firstJoinResult.leaderId)
+    assertEquals(Errors.NONE, firstJoinResult.error)
+
+    EasyMock.reset(replicaManager)
+    val firstSyncResult = syncGroupLeader(groupId, firstGenerationId, firstMemberId,
+      Map(firstMemberId -> Array[Byte]()))
+    assertEquals(Errors.NONE, firstSyncResult._2)
+
+    EasyMock.reset(replicaManager)
+    val otherJoinFuture = sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
+
+    EasyMock.reset(replicaManager)
+    val joinFuture = sendJoinGroup(groupId, firstMemberId, protocolType, protocols,
+      requireKnownMemberId = false)
+
+    val joinResult = await(joinFuture, DefaultSessionTimeout+100)
+    val otherJoinResult = await(otherJoinFuture, DefaultSessionTimeout+100)
+    assertEquals(Errors.NONE, joinResult.error)
+    assertEquals(Errors.NONE, otherJoinResult.error)
+    val secondGenerationId = joinResult.generationId
+    val secondMemberId = otherJoinResult.memberId
+
+    EasyMock.reset(replicaManager)
+    sendSyncGroupFollower(groupId, secondGenerationId, secondMemberId)
+
+    EasyMock.reset(replicaManager)
+    val syncGroupResult = syncGroupLeader(groupId, secondGenerationId, firstMemberId,
+      Map(firstMemberId -> Array.emptyByteArray, secondMemberId -> Array.emptyByteArray))
+    assertEquals(Errors.NONE, syncGroupResult._2)
+
+    verifySessionExpiration(groupId)
+  }
+
+  private def verifySessionExpiration(groupId: String): Unit = {
+    EasyMock.reset(replicaManager)
+    EasyMock.expect(replicaManager.getMagic(EasyMock.anyObject()))
+      .andReturn(Some(RecordBatch.CURRENT_MAGIC_VALUE)).anyTimes()
+    EasyMock.replay(replicaManager)
+
+    timer.advanceClock(DefaultSessionTimeout + 1)
+
+    val groupMetadata = group(groupId)
+    assertEquals(Empty, groupMetadata.currentState)
+    assertEquals(0, groupMetadata.allMembers.size)
+  }
+
+  @Test
+  def testJoinGroupInconsistentGroupProtocol(): Unit = {
     val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
 
     val otherMemberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
@@ -338,7 +483,7 @@ class GroupCoordinatorTest extends JUnitSuite {
     val joinGroupFuture = sendJoinGroup(groupId, memberId, protocolType, List(("range", metadata)))
 
     EasyMock.reset(replicaManager)
-    val otherJoinGroupResult = joinGroup(groupId, otherMemberId, protocolType, List(("roundrobin", metadata)))
+    val otherJoinGroupResult = dynamicJoinGroup(groupId, otherMemberId, protocolType, List(("roundrobin", metadata)))
     timer.advanceClock(GroupInitialRebalanceDelay + 1)
 
     val joinGroupResult = await(joinGroupFuture, 1)
@@ -347,11 +492,11 @@ class GroupCoordinatorTest extends JUnitSuite {
   }
 
   @Test
-  def testJoinGroupUnknownConsumerExistingGroup() {
+  def testJoinGroupUnknownConsumerExistingGroup(): Unit = {
     val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
     val otherMemberId = "memberId"
 
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols)
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols)
     assertEquals(Errors.NONE, joinGroupResult.error)
 
     EasyMock.reset(replicaManager)
@@ -360,18 +505,29 @@ class GroupCoordinatorTest extends JUnitSuite {
   }
 
   @Test
-  def testJoinGroupUnknownConsumerDeadGroup() {
+  def testJoinGroupUnknownConsumerNewDeadGroup(): Unit = {
     val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
 
     val deadGroupId = "deadGroupId"
 
+    groupCoordinator.groupManager.addGroup(new  GroupMetadata(deadGroupId, Dead, new MockTime()))
+    val joinGroupResult = dynamicJoinGroup(deadGroupId, memberId, protocolType, protocols)
+    assertEquals(Errors.COORDINATOR_NOT_AVAILABLE, joinGroupResult.error)
+  }
+
+  @Test
+  def testSyncDeadGroup(): Unit = {
+    val memberId = "memberId"
+
+    val deadGroupId = "deadGroupId"
+
     groupCoordinator.groupManager.addGroup(new GroupMetadata(deadGroupId, Dead, new MockTime()))
-    val joinGroupResult = joinGroup(deadGroupId, memberId, protocolType, protocols)
-    assertEquals(Errors.UNKNOWN_MEMBER_ID, joinGroupResult.error)
+    val syncGroupResult = syncGroupFollower(deadGroupId, 1, memberId)
+    assertEquals(Errors.COORDINATOR_NOT_AVAILABLE, syncGroupResult._2)
   }
 
   @Test
-  def testJoinGroupSecondJoinInconsistentProtocol() {
+  def testJoinGroupSecondJoinInconsistentProtocol(): Unit = {
     var responseFuture = sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols, requireKnownMemberId = true)
     var joinGroupResult = Await.result(responseFuture, Duration(DefaultRebalanceTimeout + 1, TimeUnit.MILLISECONDS))
     assertEquals(Errors.MEMBER_ID_REQUIRED, joinGroupResult.error)
@@ -392,24 +548,846 @@ class GroupCoordinatorTest extends JUnitSuite {
   }
 
   @Test
-  def testHeartbeatWrongCoordinator() {
+  def staticMemberJoinAsFirstMember(): Unit = {
+    val joinGroupResult = staticJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, groupInstanceId, protocolType, protocols)
+    assertEquals(Errors.NONE, joinGroupResult.error)
+  }
+
+  @Test
+  def staticMemberReJoinWithExplicitUnknownMemberId(): Unit = {
+    var joinGroupResult = staticJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, groupInstanceId, protocolType, protocols)
+    assertEquals(Errors.NONE, joinGroupResult.error)
+
+    EasyMock.reset(replicaManager)
+    val unknownMemberId = "unknown_member"
+    joinGroupResult = staticJoinGroup(groupId, unknownMemberId, groupInstanceId, protocolType, protocols)
+    assertEquals(Errors.FENCED_INSTANCE_ID, joinGroupResult.error)
+  }
+
+  @Test
+  def staticMemberFenceDuplicateRejoinedFollower(): Unit = {
+    val rebalanceResult = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId)
+
+    EasyMock.reset(replicaManager)
+    // A third member joins will trigger rebalance.
+    sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
+    timer.advanceClock(1)
+    assertTrue(getGroup(groupId).is(PreparingRebalance))
+
+    EasyMock.reset(replicaManager)
+    timer.advanceClock(1)
+    // Old follower rejoins group will be matching current member.id.
+    val oldFollowerJoinGroupFuture =
+      sendJoinGroup(groupId, rebalanceResult.followerId, protocolType, protocols, groupInstanceId = followerInstanceId)
+
+    EasyMock.reset(replicaManager)
+    timer.advanceClock(1)
+    // Duplicate follower joins group with unknown member id will trigger member.id replacement.
+    val duplicateFollowerJoinFuture =
+      sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols, groupInstanceId = followerInstanceId)
+
+    timer.advanceClock(1)
+    // Old member shall be fenced immediately upon duplicate follower joins.
+    val oldFollowerJoinGroupResult = Await.result(oldFollowerJoinGroupFuture, Duration(1, TimeUnit.MILLISECONDS))
+    checkJoinGroupResult(oldFollowerJoinGroupResult,
+      Errors.FENCED_INSTANCE_ID,
+      -1,
+      Set.empty,
+      groupId,
+      PreparingRebalance)
+    verifyDelayedTaskNotCompleted(duplicateFollowerJoinFuture)
+  }
+
+  @Test
+  def staticMemberFenceDuplicateSyncingFollowerAfterMemberIdChanged(): Unit = {
+    val rebalanceResult = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId)
+
+    EasyMock.reset(replicaManager)
+    // Known leader rejoins will trigger rebalance.
+    val leaderJoinGroupFuture =
+      sendJoinGroup(groupId, rebalanceResult.leaderId, protocolType, protocols, groupInstanceId = leaderInstanceId)
+    timer.advanceClock(1)
+    assertTrue(getGroup(groupId).is(PreparingRebalance))
+
+    EasyMock.reset(replicaManager)
+    timer.advanceClock(1)
+    // Old follower rejoins group will match current member.id.
+    val oldFollowerJoinGroupFuture =
+      sendJoinGroup(groupId, rebalanceResult.followerId, protocolType, protocols, groupInstanceId = followerInstanceId)
+
+    timer.advanceClock(1)
+    val leaderJoinGroupResult = Await.result(leaderJoinGroupFuture, Duration(1, TimeUnit.MILLISECONDS))
+    checkJoinGroupResult(leaderJoinGroupResult,
+      Errors.NONE,
+      rebalanceResult.generation + 1,
+      Set(leaderInstanceId, followerInstanceId),
+      groupId,
+      CompletingRebalance)
+    assertEquals(leaderJoinGroupResult.leaderId, leaderJoinGroupResult.memberId)
+    assertEquals(rebalanceResult.leaderId, leaderJoinGroupResult.leaderId)
+
+    // Old member shall be getting a successful join group response.
+    val oldFollowerJoinGroupResult = Await.result(oldFollowerJoinGroupFuture, Duration(1, TimeUnit.MILLISECONDS))
+    checkJoinGroupResult(oldFollowerJoinGroupResult,
+      Errors.NONE,
+      rebalanceResult.generation + 1,
+      Set.empty,
+      groupId,
+      CompletingRebalance,
+      expectedLeaderId = leaderJoinGroupResult.memberId)
+
+    EasyMock.reset(replicaManager)
+    val oldFollowerSyncGroupFuture = sendSyncGroupFollower(groupId, oldFollowerJoinGroupResult.generationId,
+      oldFollowerJoinGroupResult.memberId, followerInstanceId)
+
+    // Duplicate follower joins group with unknown member id will trigger member.id replacement.
+    EasyMock.reset(replicaManager)
+    val duplicateFollowerJoinFuture =
+      sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols, groupInstanceId = followerInstanceId)
+    timer.advanceClock(1)
+
+    // Old follower sync callback will return fenced exception while broker replaces the member identity.
+    val oldFollowerSyncGroupResult = Await.result(oldFollowerSyncGroupFuture, Duration(1, TimeUnit.MILLISECONDS))
+    assertEquals(oldFollowerSyncGroupResult._2, Errors.FENCED_INSTANCE_ID)
+
+    // Duplicate follower will get the same response as old follower.
+    val duplicateFollowerJoinGroupResult = Await.result(duplicateFollowerJoinFuture, Duration(1, TimeUnit.MILLISECONDS))
+    checkJoinGroupResult(duplicateFollowerJoinGroupResult,
+      Errors.NONE,
+      rebalanceResult.generation + 1,
+      Set.empty,
+      groupId,
+      CompletingRebalance,
+      expectedLeaderId = leaderJoinGroupResult.memberId)
+  }
+
+  @Test
+  def staticMemberFenceDuplicateRejoiningFollowerAfterMemberIdChanged(): Unit = {
+    val rebalanceResult = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId)
+
+    EasyMock.reset(replicaManager)
+    // Known leader rejoins will trigger rebalance.
+    val leaderJoinGroupFuture =
+      sendJoinGroup(groupId, rebalanceResult.leaderId, protocolType, protocols, groupInstanceId = leaderInstanceId)
+    timer.advanceClock(1)
+    assertTrue(getGroup(groupId).is(PreparingRebalance))
+
+    EasyMock.reset(replicaManager)
+    // Duplicate follower joins group will trigger member.id replacement.
+    val duplicateFollowerJoinGroupFuture =
+      sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols, groupInstanceId = followerInstanceId)
+
+    EasyMock.reset(replicaManager)
+    timer.advanceClock(1)
+    // Old follower rejoins group will fail because member.id already updated.
+    val oldFollowerJoinGroupFuture =
+      sendJoinGroup(groupId, rebalanceResult.followerId, protocolType, protocols, groupInstanceId = followerInstanceId)
+
+    val leaderRejoinGroupResult = Await.result(leaderJoinGroupFuture, Duration(1, TimeUnit.MILLISECONDS))
+    checkJoinGroupResult(leaderRejoinGroupResult,
+      Errors.NONE,
+      rebalanceResult.generation + 1,
+      Set(leaderInstanceId, followerInstanceId),
+      groupId,
+      CompletingRebalance)
+
+    val duplicateFollowerJoinGroupResult = Await.result(duplicateFollowerJoinGroupFuture, Duration(1, TimeUnit.MILLISECONDS))
+    checkJoinGroupResult(duplicateFollowerJoinGroupResult,
+      Errors.NONE,
+      rebalanceResult.generation + 1,
+      Set.empty,
+      groupId,
+      CompletingRebalance)
+    assertNotEquals(rebalanceResult.followerId, duplicateFollowerJoinGroupResult.memberId)
+
+    val oldFollowerJoinGroupResult = Await.result(oldFollowerJoinGroupFuture, Duration(1, TimeUnit.MILLISECONDS))
+    checkJoinGroupResult(oldFollowerJoinGroupResult,
+      Errors.FENCED_INSTANCE_ID,
+      -1,
+      Set.empty,
+      groupId,
+      CompletingRebalance)
+  }
+
+  @Test
+  def staticMemberRejoinWithKnownMemberId(): Unit = {
+    var joinGroupResult = staticJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, groupInstanceId, protocolType, protocols)
+    assertEquals(Errors.NONE, joinGroupResult.error)
+
+    EasyMock.reset(replicaManager)
+    val assignedMemberId = joinGroupResult.memberId
+    // The second join group should return immediately since we are using the same metadata during CompletingRebalance.
+    val rejoinResponseFuture = sendJoinGroup(groupId, assignedMemberId, protocolType, protocols, groupInstanceId)
+    timer.advanceClock(1)
+    joinGroupResult = Await.result(rejoinResponseFuture, Duration(1, TimeUnit.MILLISECONDS))
+    assertEquals(Errors.NONE, joinGroupResult.error)
+    assertTrue(getGroup(groupId).is(CompletingRebalance))
+
+    EasyMock.reset(replicaManager)
+    val syncGroupFuture = sendSyncGroupLeader(groupId, joinGroupResult.generationId, assignedMemberId, groupInstanceId, Map(assignedMemberId -> Array[Byte]()))
+    timer.advanceClock(1)
+    val syncGroupResult = Await.result(syncGroupFuture, Duration(1, TimeUnit.MILLISECONDS))
+    assertEquals(Errors.NONE, syncGroupResult._2)
+    assertTrue(getGroup(groupId).is(Stable))
+  }
+
+  @Test
+  def staticMemberRejoinWithLeaderIdAndUnknownMemberId(): Unit = {
+    val rebalanceResult = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId)
+
+    // A static leader rejoin with unknown id will not trigger rebalance, and no assignment will be returned.
+    EasyMock.reset(replicaManager)
+    val joinGroupResult = staticJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, leaderInstanceId, protocolType, protocolSuperset, clockAdvance = 1)
+
+    checkJoinGroupResult(joinGroupResult,
+      Errors.NONE,
+      rebalanceResult.generation, // The group should be at the same generation
+      Set.empty,
+      groupId,
+      Stable,
+      rebalanceResult.leaderId)
+
+    EasyMock.reset(replicaManager)
+    val oldLeaderJoinGroupResult = staticJoinGroup(groupId, rebalanceResult.leaderId, leaderInstanceId, protocolType, protocolSuperset, clockAdvance = 1)
+    assertEquals(Errors.FENCED_INSTANCE_ID, oldLeaderJoinGroupResult.error)
+
+    EasyMock.reset(replicaManager)
+    // Old leader will get fenced.
+    val oldLeaderSyncGroupResult = syncGroupLeader(groupId, rebalanceResult.generation, rebalanceResult.leaderId, Map.empty, leaderInstanceId)
+    assertEquals(Errors.FENCED_INSTANCE_ID, oldLeaderSyncGroupResult._2)
+
+    // Calling sync on old leader.id will fail because that leader.id is no longer valid and replaced.
+    EasyMock.reset(replicaManager)
+    val newLeaderSyncGroupResult = syncGroupLeader(groupId, rebalanceResult.generation, joinGroupResult.leaderId, Map.empty)
+    assertEquals(Errors.UNKNOWN_MEMBER_ID, newLeaderSyncGroupResult._2)
+  }
+
+  @Test
+  def staticMemberRejoinWithLeaderIdAndKnownMemberId(): Unit = {
+    val rebalanceResult = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId,
+      sessionTimeout = DefaultRebalanceTimeout / 2)
+
+    // A static leader with known id rejoin will trigger rebalance.
+    EasyMock.reset(replicaManager)
+    val joinGroupResult = staticJoinGroup(groupId, rebalanceResult.leaderId, leaderInstanceId,
+      protocolType, protocolSuperset, clockAdvance = DefaultRebalanceTimeout + 1)
+    // Timeout follower in the meantime.
+    assertFalse(getGroup(groupId).hasStaticMember(followerInstanceId))
+    checkJoinGroupResult(joinGroupResult,
+      Errors.NONE,
+      rebalanceResult.generation + 1, // The group has promoted to the new generation.
+      Set(leaderInstanceId),
+      groupId,
+      CompletingRebalance,
+      rebalanceResult.leaderId,
+      rebalanceResult.leaderId)
+  }
+
+  @Test
+  def staticMemberRejoinWithLeaderIdAndUnexpectedDeadGroup(): Unit = {
+    val rebalanceResult = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId)
+
+    getGroup(groupId).transitionTo(Dead)
+
+    EasyMock.reset(replicaManager)
+    val joinGroupResult = staticJoinGroup(groupId, rebalanceResult.leaderId, leaderInstanceId, protocolType, protocols, clockAdvance = 1)
+    assertEquals(Errors.COORDINATOR_NOT_AVAILABLE, joinGroupResult.error)
+  }
+
+  @Test
+  def staticMemberRejoinWithLeaderIdAndUnexpectedEmptyGroup(): Unit = {
+    val rebalanceResult = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId)
+
+    getGroup(groupId).transitionTo(PreparingRebalance)
+    getGroup(groupId).transitionTo(Empty)
+
+    EasyMock.reset(replicaManager)
+    val joinGroupResult = staticJoinGroup(groupId, rebalanceResult.leaderId, leaderInstanceId, protocolType, protocols, clockAdvance = 1)
+    assertEquals(Errors.UNKNOWN_MEMBER_ID, joinGroupResult.error)
+  }
+
+  @Test
+  def staticMemberRejoinWithFollowerIdAndChangeOfProtocol(): Unit = {
+    val rebalanceResult = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId, sessionTimeout = DefaultSessionTimeout * 2)
+
+    // A static follower rejoin with changed protocol will trigger rebalance.
+    EasyMock.reset(replicaManager)
+    val newProtocols = List(("roundrobin", metadata))
+    // Old leader hasn't joined in the meantime, triggering a re-election.
+    val joinGroupResult = staticJoinGroup(groupId, rebalanceResult.followerId, followerInstanceId, protocolType, newProtocols, clockAdvance = DefaultSessionTimeout + 1)
+
+    assertEquals(rebalanceResult.followerId, joinGroupResult.memberId)
+    assertTrue(getGroup(groupId).hasStaticMember(leaderInstanceId))
+    assertTrue(getGroup(groupId).isLeader(rebalanceResult.followerId))
+    checkJoinGroupResult(joinGroupResult,
+      Errors.NONE,
+      rebalanceResult.generation + 1, // The group has promoted to the new generation, and leader has changed because old one times out.
+      Set(leaderInstanceId, followerInstanceId),
+      groupId,
+      CompletingRebalance,
+      rebalanceResult.followerId,
+      rebalanceResult.followerId)
+  }
+
+  @Test
+  def staticMemberRejoinWithUnknownMemberIdAndChangeOfProtocol(): Unit = {
+    val rebalanceResult = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId)
+
+    // A static follower rejoin with protocol changing to leader protocol subset won't trigger rebalance.
+    EasyMock.reset(replicaManager)
+    val newProtocols = List(("roundrobin", metadata))
+    // Timeout old leader in the meantime.
+    val joinGroupResult = staticJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, followerInstanceId, protocolType, newProtocols, clockAdvance = 1)
+
+    checkJoinGroupResult(joinGroupResult,
+      Errors.NONE,
+      rebalanceResult.generation,
+      Set.empty,
+      groupId,
+      Stable)
+
+    EasyMock.reset(replicaManager)
+    // Join with old member id will fail because the member id is updated
+    assertNotEquals(rebalanceResult.followerId, joinGroupResult.memberId)
+    val oldFollowerJoinGroupResult = staticJoinGroup(groupId, rebalanceResult.followerId, followerInstanceId, protocolType, newProtocols, clockAdvance = 1)
+    assertEquals(Errors.FENCED_INSTANCE_ID, oldFollowerJoinGroupResult.error)
+
+    EasyMock.reset(replicaManager)
+    // Sync with old member id will fail because the member id is updated
+    val syncGroupWithOldMemberIdResult = syncGroupFollower(groupId, rebalanceResult.generation, rebalanceResult.followerId, followerInstanceId)
+    assertEquals(Errors.FENCED_INSTANCE_ID, syncGroupWithOldMemberIdResult._2)
+
+    EasyMock.reset(replicaManager)
+    val syncGroupWithNewMemberIdResult = syncGroupFollower(groupId, rebalanceResult.generation, joinGroupResult.memberId, followerInstanceId)
+    assertEquals(Errors.NONE, syncGroupWithNewMemberIdResult._2)
+    assertEquals(rebalanceResult.followerAssignment, syncGroupWithNewMemberIdResult._1)
+  }
+
+  @Test
+  def staticMemberRejoinWithKnownLeaderIdToTriggerRebalanceAndFollowerWithChangeofProtocol(): Unit = {
+    val rebalanceResult = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId)
+
+    // A static leader rejoin with known member id will trigger rebalance.
+    EasyMock.reset(replicaManager)
+    val leaderRejoinGroupFuture = sendJoinGroup(groupId, rebalanceResult.leaderId, protocolType, protocolSuperset, leaderInstanceId)
+    // Rebalance complete immediately after follower rejoin.
+    EasyMock.reset(replicaManager)
+    val followerRejoinWithFuture = sendJoinGroup(groupId, rebalanceResult.followerId, protocolType, protocolSuperset, followerInstanceId)
+
+    timer.advanceClock(1)
+
+    // Leader should get the same assignment as last round.
+    checkJoinGroupResult(await(leaderRejoinGroupFuture, 1),
+      Errors.NONE,
+      rebalanceResult.generation + 1, // The group has promoted to the new generation.
+      Set(leaderInstanceId, followerInstanceId),
+      groupId,
+      CompletingRebalance,
+      rebalanceResult.leaderId,
+      rebalanceResult.leaderId)
+
+    checkJoinGroupResult(await(followerRejoinWithFuture, 1),
+      Errors.NONE,
+      rebalanceResult.generation + 1, // The group has promoted to the new generation.
+      Set.empty,
+      groupId,
+      CompletingRebalance,
+      rebalanceResult.leaderId,
+      rebalanceResult.followerId)
+
+    EasyMock.reset(replicaManager)
+    // The follower protocol changed from protocolSuperset to general protocols.
+    val followerRejoinWithProtocolChangeFuture = sendJoinGroup(groupId, rebalanceResult.followerId, protocolType, protocols, followerInstanceId)
+    // The group will transit to PreparingRebalance due to protocol change from follower.
+    assertTrue(getGroup(groupId).is(PreparingRebalance))
+
+    timer.advanceClock(DefaultRebalanceTimeout + 1)
+    checkJoinGroupResult(await(followerRejoinWithProtocolChangeFuture, 1),
+      Errors.NONE,
+      rebalanceResult.generation + 2, // The group has promoted to the new generation.
+      Set(followerInstanceId),
+      groupId,
+      CompletingRebalance,
+      rebalanceResult.followerId,
+      rebalanceResult.followerId)
+  }
+
+  @Test
+  def staticMemberRejoinAsFollowerWithUnknownMemberId(): Unit = {
+    val rebalanceResult = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId)
+
+    EasyMock.reset(replicaManager)
+    // A static follower rejoin with no protocol change will not trigger rebalance.
+    val joinGroupResult = staticJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, followerInstanceId, protocolType, protocolSuperset, clockAdvance = 1)
+
+    // Old leader shouldn't be timed out.
+    assertTrue(getGroup(groupId).hasStaticMember(leaderInstanceId))
+    checkJoinGroupResult(joinGroupResult,
+      Errors.NONE,
+      rebalanceResult.generation, // The group has no change.
+      Set.empty,
+      groupId,
+      Stable)
+
+    assertNotEquals(rebalanceResult.followerId, joinGroupResult.memberId)
+
+    EasyMock.reset(replicaManager)
+    val syncGroupResult = syncGroupFollower(groupId, rebalanceResult.generation, joinGroupResult.memberId)
+    assertEquals(Errors.NONE, syncGroupResult._2)
+    assertEquals(rebalanceResult.followerAssignment, syncGroupResult._1)
+  }
+
+  @Test
+  def staticMemberRejoinAsFollowerWithKnownMemberIdAndNoProtocolChange(): Unit = {
+    val rebalanceResult  = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId)
+
+    // A static follower rejoin with no protocol change will not trigger rebalance.
+    EasyMock.reset(replicaManager)
+    val joinGroupResult = staticJoinGroup(groupId, rebalanceResult.followerId, followerInstanceId, protocolType, protocolSuperset, clockAdvance = 1)
+
+    // Old leader shouldn't be timed out.
+    assertTrue(getGroup(groupId).hasStaticMember(leaderInstanceId))
+    checkJoinGroupResult(joinGroupResult,
+      Errors.NONE,
+      rebalanceResult.generation, // The group has no change.
+      Set.empty,
+      groupId,
+      Stable,
+      rebalanceResult.leaderId,
+      rebalanceResult.followerId)
+  }
+
+  @Test
+  def staticMemberRejoinAsFollowerWithMismatchedMemberId(): Unit = {
+    val rebalanceResult = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId)
+
+    EasyMock.reset(replicaManager)
+    val joinGroupResult = staticJoinGroup(groupId, rebalanceResult.followerId, leaderInstanceId, protocolType, protocolSuperset, clockAdvance = 1)
+    assertEquals(Errors.FENCED_INSTANCE_ID, joinGroupResult.error)
+  }
+
+  @Test
+  def staticMemberRejoinAsLeaderWithMismatchedMemberId(): Unit = {
+    val rebalanceResult = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId)
+
+    EasyMock.reset(replicaManager)
+    val joinGroupResult = staticJoinGroup(groupId, rebalanceResult.leaderId, followerInstanceId, protocolType, protocolSuperset, clockAdvance = 1)
+    assertEquals(Errors.FENCED_INSTANCE_ID, joinGroupResult.error)
+  }
+
+  @Test
+  def staticMemberSyncAsLeaderWithInvalidMemberId(): Unit = {
+    val rebalanceResult = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId)
+
+    EasyMock.reset(replicaManager)
+    val syncGroupResult = syncGroupLeader(groupId, rebalanceResult.generation, "invalid", Map.empty, leaderInstanceId)
+    assertEquals(Errors.FENCED_INSTANCE_ID, syncGroupResult._2)
+  }
+
+  @Test
+  def staticMemberHeartbeatLeaderWithInvalidMemberId(): Unit = {
+    val rebalanceResult = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId)
+
+    EasyMock.reset(replicaManager)
+    val syncGroupResult = syncGroupLeader(groupId, rebalanceResult.generation, rebalanceResult.leaderId, Map.empty)
+    assertEquals(Errors.NONE, syncGroupResult._2)
+
+    EasyMock.reset(replicaManager)
+    val validHeartbeatResult = heartbeat(groupId, rebalanceResult.leaderId, rebalanceResult.generation)
+    assertEquals(Errors.NONE, validHeartbeatResult)
+
+    EasyMock.reset(replicaManager)
+    val invalidHeartbeatResult = heartbeat(groupId, invalidMemberId, rebalanceResult.generation, leaderInstanceId)
+    assertEquals(Errors.FENCED_INSTANCE_ID, invalidHeartbeatResult)
+  }
+
+  @Test
+  def shouldGetDifferentStaticMemberIdAfterEachRejoin(): Unit = {
+    val initialResult = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId)
+
+    val timeAdvance = 1
+    var lastMemberId = initialResult.leaderId
+    for (_ <- 1 to 5) {
+      EasyMock.reset(replicaManager)
+
+      val joinGroupResult = staticJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID,
+        leaderInstanceId, protocolType, protocols, clockAdvance = timeAdvance)
+      assertTrue(joinGroupResult.memberId.startsWith(leaderInstanceId.get))
+      assertNotEquals(lastMemberId, joinGroupResult.memberId)
+      lastMemberId = joinGroupResult.memberId
+    }
+  }
+
+  @Test
+  def testOffsetCommitDeadGroup(): Unit = {
+    val memberId = "memberId"
+
+    val deadGroupId = "deadGroupId"
+    val tp = new TopicPartition("topic", 0)
+    val offset = offsetAndMetadata(0)
+
+    groupCoordinator.groupManager.addGroup(new GroupMetadata(deadGroupId, Dead, new MockTime()))
+    val offsetCommitResult = commitOffsets(deadGroupId, memberId, 1, Map(tp -> offset))
+    assertEquals(Errors.COORDINATOR_NOT_AVAILABLE, offsetCommitResult(tp))
+  }
+
+  @Test
+  def staticMemberCommitOffsetWithInvalidMemberId(): Unit = {
+    val rebalanceResult = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId)
+
+    EasyMock.reset(replicaManager)
+    val syncGroupResult = syncGroupLeader(groupId, rebalanceResult.generation, rebalanceResult.leaderId, Map.empty)
+    assertEquals(Errors.NONE, syncGroupResult._2)
+
+    val tp = new TopicPartition("topic", 0)
+    val offset = offsetAndMetadata(0)
+    EasyMock.reset(replicaManager)
+    val validOffsetCommitResult = commitOffsets(groupId, rebalanceResult.leaderId, rebalanceResult.generation, Map(tp -> offset))
+    assertEquals(Errors.NONE, validOffsetCommitResult(tp))
+
+    EasyMock.reset(replicaManager)
+    val invalidOffsetCommitResult = commitOffsets(groupId, invalidMemberId, rebalanceResult.generation, Map(tp -> offset), leaderInstanceId)
+    assertEquals(Errors.FENCED_INSTANCE_ID, invalidOffsetCommitResult(tp))
+  }
+
+  @Test
+  def staticMemberJoinWithUnknownInstanceIdAndKnownMemberId(): Unit = {
+    val rebalanceResult = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId)
+
+    EasyMock.reset(replicaManager)
+    val joinGroupResult = staticJoinGroup(groupId, rebalanceResult.leaderId, Some("unknown_instance"), protocolType, protocolSuperset, clockAdvance = 1)
+
+    assertEquals(Errors.UNKNOWN_MEMBER_ID, joinGroupResult.error)
+  }
+
+  @Test
+  def staticMemberJoinWithIllegalStateAsPendingMember(): Unit = {
+    val rebalanceResult = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId)
+    val group = groupCoordinator.groupManager.getGroup(groupId).get
+    group.addPendingMember(rebalanceResult.followerId)
+    group.remove(rebalanceResult.followerId)
+    EasyMock.reset(replicaManager)
+
+    // Illegal state exception shall trigger since follower id resides in pending member bucket.
+    val expectedException = intercept[IllegalStateException] {
+      staticJoinGroup(groupId, rebalanceResult.followerId, followerInstanceId, protocolType, protocolSuperset, clockAdvance = 1)
+    }
+
+    val message = expectedException.getMessage
+    assertTrue(message.contains(rebalanceResult.followerId))
+    assertTrue(message.contains(followerInstanceId.get))
+  }
+
+  @Test
+  def staticMemberLeaveWithIllegalStateAsPendingMember(): Unit = {
+    val rebalanceResult = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId)
+    val group = groupCoordinator.groupManager.getGroup(groupId).get
+    group.addPendingMember(rebalanceResult.followerId)
+    group.remove(rebalanceResult.followerId)
+    EasyMock.reset(replicaManager)
+
+    // Illegal state exception shall trigger since follower id resides in pending member bucket.
+    val expectedException = intercept[IllegalStateException] {
+      singleLeaveGroup(groupId, rebalanceResult.followerId, followerInstanceId)
+    }
+
+    val message = expectedException.getMessage
+    assertTrue(message.contains(rebalanceResult.followerId))
+    assertTrue(message.contains(followerInstanceId.get))
+  }
+
+  @Test
+  def staticMemberReJoinWithIllegalStateAsUnknownMember(): Unit = {
+    staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId)
+    val group = groupCoordinator.groupManager.getGroup(groupId).get
+    group.transitionTo(PreparingRebalance)
+    group.transitionTo(Empty)
+
+    EasyMock.reset(replicaManager)
+
+    // Illegal state exception shall trigger since follower id resides in pending member bucket.
+    val expectedException = intercept[IllegalStateException] {
+      staticJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, followerInstanceId, protocolType, protocolSuperset, clockAdvance = 1)
+    }
+
+    val message = expectedException.getMessage
+    assertTrue(message.contains(group.groupId))
+    assertTrue(message.contains(followerInstanceId.get))
+  }
+
+  @Test
+  def staticMemberReJoinWithIllegalArgumentAsMissingOldMember(): Unit = {
+    staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId)
+    val group = groupCoordinator.groupManager.getGroup(groupId).get
+    val invalidMemberId = "invalid_member_id"
+    group.addStaticMember(followerInstanceId, invalidMemberId)
+    EasyMock.reset(replicaManager)
+
+    // Illegal state exception shall trigger since follower corresponding id is not defined in member list.
+    val expectedException = intercept[IllegalArgumentException] {
+      staticJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, followerInstanceId, protocolType, protocolSuperset, clockAdvance = 1)
+    }
+
+    val message = expectedException.getMessage
+    assertTrue(message.contains(invalidMemberId))
+  }
+
+  @Test
+  def testLeaderFailToRejoinBeforeFinalRebalanceTimeoutWithLongSessionTimeout(): Unit = {
+    groupStuckInRebalanceTimeoutDueToNonjoinedStaticMember()
+
+    timer.advanceClock(DefaultRebalanceTimeout + 1)
+    // The static leader should already session timeout, moving group towards Empty
+    assertEquals(Set.empty, getGroup(groupId).allMembers)
+    assertEquals(null, getGroup(groupId).leaderOrNull)
+    assertEquals(3, getGroup(groupId).generationId)
+    assertGroupState(groupState = Empty)
+  }
+
+  @Test
+  def testLeaderRejoinBeforeFinalRebalanceTimeoutWithLongSessionTimeout(): Unit = {
+    groupStuckInRebalanceTimeoutDueToNonjoinedStaticMember()
+
+    EasyMock.reset(replicaManager)
+    // The static leader should be back now, moving group towards CompletingRebalance
+    val leaderRejoinGroupResult = staticJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, leaderInstanceId, protocolType, protocols)
+    checkJoinGroupResult(leaderRejoinGroupResult,
+      Errors.NONE,
+      3,
+      Set(leaderInstanceId),
+      groupId,
+      CompletingRebalance
+    )
+    assertEquals(1, getGroup(groupId).allMembers.size)
+    assertNotEquals(null, getGroup(groupId).leaderOrNull)
+    assertEquals(3, getGroup(groupId).generationId)
+  }
+
+  def groupStuckInRebalanceTimeoutDueToNonjoinedStaticMember(): Unit = {
+    val longSessionTimeout = DefaultSessionTimeout * 2
+    val rebalanceResult = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId, sessionTimeout = longSessionTimeout)
+
+    EasyMock.reset(replicaManager)
+
+    val dynamicJoinFuture = sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocolSuperset, sessionTimeout = longSessionTimeout)
+    timer.advanceClock(DefaultRebalanceTimeout + 1)
+
+    val dynamicJoinResult = await(dynamicJoinFuture, 100)
+    // The new dynamic member has been elected as leader
+    assertEquals(dynamicJoinResult.leaderId, dynamicJoinResult.memberId)
+    assertEquals(Errors.NONE, dynamicJoinResult.error)
+    assertEquals(3, dynamicJoinResult.members.size)
+    assertEquals(2, dynamicJoinResult.generationId)
+    assertGroupState(groupState = CompletingRebalance)
+
+    // Send a special leave group request from static follower, moving group towards PreparingRebalance
+    EasyMock.reset(replicaManager)
+    val followerLeaveGroupResults = singleLeaveGroup(groupId, rebalanceResult.followerId)
+    verifyLeaveGroupResult(followerLeaveGroupResults)
+    assertGroupState(groupState = PreparingRebalance)
+
+    timer.advanceClock(DefaultRebalanceTimeout + 1)
+    // Only static leader is maintained, and group is stuck at PreparingRebalance stage
+    assertEquals(1, getGroup(groupId).allMembers.size)
+    assertEquals(Set(rebalanceResult.leaderId), getGroup(groupId).allMembers)
+    assertEquals(2, getGroup(groupId).generationId)
+    assertGroupState(groupState = PreparingRebalance)
+  }
+
+  @Test
+  def testStaticMemberFollowerFailToRejoinBeforeRebalanceTimeout(): Unit = {
+    // Increase session timeout so that the follower won't be evicted when rebalance timeout is reached.
+    val initialRebalanceResult = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId, sessionTimeout = DefaultRebalanceTimeout * 2)
+
+    EasyMock.reset(replicaManager)
+    val newMemberInstanceId = Some("newMember")
+
+    val leaderId = initialRebalanceResult.leaderId
+
+    val newMemberJoinGroupFuture = sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocolSuperset, newMemberInstanceId)
+    assertGroupState(groupState = PreparingRebalance)
+
+    EasyMock.reset(replicaManager)
+    val leaderRejoinGroupResult = staticJoinGroup(groupId, leaderId, leaderInstanceId, protocolType, protocolSuperset, clockAdvance = DefaultRebalanceTimeout + 1)
+    checkJoinGroupResult(leaderRejoinGroupResult,
+      Errors.NONE,
+      initialRebalanceResult.generation + 1,
+      Set(leaderInstanceId, followerInstanceId, newMemberInstanceId),
+      groupId,
+      CompletingRebalance,
+      expectedLeaderId = leaderId,
+      expectedMemberId = leaderId)
+
+    val newMemberJoinGroupResult = Await.result(newMemberJoinGroupFuture, Duration(1, TimeUnit.MILLISECONDS))
+    assertEquals(Errors.NONE, newMemberJoinGroupResult.error)
+    checkJoinGroupResult(newMemberJoinGroupResult,
+      Errors.NONE,
+      initialRebalanceResult.generation + 1,
+      Set.empty,
+      groupId,
+      CompletingRebalance,
+      expectedLeaderId = leaderId)
+  }
+
+  @Test
+  def testStaticMemberLeaderFailToRejoinBeforeRebalanceTimeout(): Unit = {
+    // Increase session timeout so that the leader won't be evicted when rebalance timeout is reached.
+    val initialRebalanceResult = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId, sessionTimeout = DefaultRebalanceTimeout * 2)
+
+    EasyMock.reset(replicaManager)
+    val newMemberInstanceId = Some("newMember")
+
+    val newMemberJoinGroupFuture = sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocolSuperset, newMemberInstanceId)
+    timer.advanceClock(1)
+    assertGroupState(groupState = PreparingRebalance)
+
+    EasyMock.reset(replicaManager)
+    val oldFollowerRejoinGroupResult = staticJoinGroup(groupId, initialRebalanceResult.followerId, followerInstanceId, protocolType, protocolSuperset, clockAdvance = DefaultRebalanceTimeout + 1)
+    val newMemberJoinGroupResult = Await.result(newMemberJoinGroupFuture, Duration(1, TimeUnit.MILLISECONDS))
+
+    val (newLeaderResult, newFollowerResult) = if (oldFollowerRejoinGroupResult.leaderId == oldFollowerRejoinGroupResult.memberId)
+      (oldFollowerRejoinGroupResult, newMemberJoinGroupResult)
+    else
+      (newMemberJoinGroupResult, oldFollowerRejoinGroupResult)
+
+    checkJoinGroupResult(newLeaderResult,
+      Errors.NONE,
+      initialRebalanceResult.generation + 1,
+      Set(leaderInstanceId, followerInstanceId, newMemberInstanceId),
+      groupId,
+      CompletingRebalance)
+
+    checkJoinGroupResult(newFollowerResult,
+      Errors.NONE,
+      initialRebalanceResult.generation + 1,
+      Set.empty,
+      groupId,
+      CompletingRebalance,
+      expectedLeaderId = newLeaderResult.memberId)
+  }
+
+  private class RebalanceResult(val generation: Int,
+                                val leaderId: String,
+                                val leaderAssignment: Array[Byte],
+                                val followerId: String,
+                                val followerAssignment: Array[Byte])
+  /**
+    * Generate static member rebalance results, including:
+    *   - generation
+    *   - leader id
+    *   - leader assignment
+    *   - follower id
+    *   - follower assignment
+    */
+  private def staticMembersJoinAndRebalance(leaderInstanceId: Option[String],
+                                            followerInstanceId: Option[String],
+                                            sessionTimeout: Int = DefaultSessionTimeout): RebalanceResult = {
+    EasyMock.reset(replicaManager)
+    val leaderResponseFuture = sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocolSuperset, leaderInstanceId, sessionTimeout)
+
+    EasyMock.reset(replicaManager)
+    val followerResponseFuture = sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocolSuperset, followerInstanceId, sessionTimeout)
+    // 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.
+    timer.advanceClock(GroupInitialRebalanceDelay + 1)
+    timer.advanceClock(DefaultRebalanceTimeout + 1)
+    val newGeneration = 1
+
+    val leaderJoinGroupResult = await(leaderResponseFuture, 1)
+    assertEquals(Errors.NONE, leaderJoinGroupResult.error)
+    assertEquals(newGeneration, leaderJoinGroupResult.generationId)
+
+    val followerJoinGroupResult = await(followerResponseFuture, 1)
+    assertEquals(Errors.NONE, followerJoinGroupResult.error)
+    assertEquals(newGeneration, followerJoinGroupResult.generationId)
+
+    EasyMock.reset(replicaManager)
+    val leaderId = leaderJoinGroupResult.memberId
+    val leaderSyncGroupResult = syncGroupLeader(groupId, leaderJoinGroupResult.generationId, leaderId, Map(leaderId -> Array[Byte]()))
+    assertEquals(Errors.NONE, leaderSyncGroupResult._2)
+    assertTrue(getGroup(groupId).is(Stable))
+
+    EasyMock.reset(replicaManager)
+    val followerId = followerJoinGroupResult.memberId
+    val followerSyncGroupResult = syncGroupFollower(groupId, leaderJoinGroupResult.generationId, followerId)
+    assertEquals(Errors.NONE, followerSyncGroupResult._2)
+    assertTrue(getGroup(groupId).is(Stable))
+
+    new RebalanceResult(newGeneration,
+      leaderId,
+      leaderSyncGroupResult._1,
+      followerId,
+      followerSyncGroupResult._1)
+  }
+
+  private def checkJoinGroupResult(joinGroupResult: JoinGroupResult,
+                                   expectedError: Errors,
+                                   expectedGeneration: Int,
+                                   expectedGroupInstanceIds: Set[Option[String]],
+                                   groupId: String,
+                                   expectedGroupState: GroupState,
+                                   expectedLeaderId: String = JoinGroupRequest.UNKNOWN_MEMBER_ID,
+                                   expectedMemberId: String = JoinGroupRequest.UNKNOWN_MEMBER_ID): Unit = {
+    assertEquals(expectedError, joinGroupResult.error)
+    assertEquals(expectedGeneration, joinGroupResult.generationId)
+    assertEquals(expectedGroupInstanceIds.size, joinGroupResult.members.size)
+    val resultedGroupInstanceIds = joinGroupResult.members.map(member => Some(member.groupInstanceId())).toSet
+    assertEquals(expectedGroupInstanceIds, resultedGroupInstanceIds)
+    assertGroupState(groupState = expectedGroupState)
+
+    if (!expectedLeaderId.equals(JoinGroupRequest.UNKNOWN_MEMBER_ID)) {
+      assertEquals(expectedLeaderId, joinGroupResult.leaderId)
+    }
+    if (!expectedMemberId.equals(JoinGroupRequest.UNKNOWN_MEMBER_ID)) {
+      assertEquals(expectedMemberId, joinGroupResult.memberId)
+    }
+  }
+
+  @Test
+  def testHeartbeatWrongCoordinator(): Unit = {
     val heartbeatResult = heartbeat(otherGroupId, memberId, -1)
     assertEquals(Errors.NOT_COORDINATOR, heartbeatResult)
   }
 
   @Test
-  def testHeartbeatUnknownGroup() {
-
+  def testHeartbeatUnknownGroup(): Unit = {
     val heartbeatResult = heartbeat(groupId, memberId, -1)
     assertEquals(Errors.UNKNOWN_MEMBER_ID, heartbeatResult)
   }
 
   @Test
-  def testHeartbeatUnknownConsumerExistingGroup() {
+  def testheartbeatDeadGroup(): Unit = {
+    val memberId = "memberId"
+
+    val deadGroupId = "deadGroupId"
+
+    groupCoordinator.groupManager.addGroup(new GroupMetadata(deadGroupId, Dead, new MockTime()))
+    val heartbeatResult = heartbeat(deadGroupId, memberId, 1)
+    assertEquals(Errors.COORDINATOR_NOT_AVAILABLE, heartbeatResult)
+  }
+
+  @Test
+  def testheartbeatEmptyGroup(): Unit = {
+    val memberId = "memberId"
+
+    val group = new GroupMetadata(groupId, Empty, new MockTime())
+    val member = new MemberMetadata(memberId, groupId, groupInstanceId,
+      ClientId, ClientHost, DefaultRebalanceTimeout, DefaultSessionTimeout,
+      protocolType, List(("range", Array.empty[Byte]), ("roundrobin", Array.empty[Byte])))
+
+    group.add(member)
+    groupCoordinator.groupManager.addGroup(group)
+    val heartbeatResult = heartbeat(groupId, memberId, 0)
+    assertEquals(Errors.UNKNOWN_MEMBER_ID, heartbeatResult)
+  }
+
+  @Test
+  def testHeartbeatUnknownConsumerExistingGroup(): Unit = {
     val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
     val otherMemberId = "memberId"
 
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols)
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols)
     val assignedMemberId = joinGroupResult.memberId
     val joinGroupError = joinGroupResult.error
     assertEquals(Errors.NONE, joinGroupError)
@@ -425,24 +1403,24 @@ class GroupCoordinatorTest extends JUnitSuite {
   }
 
   @Test
-  def testHeartbeatRebalanceInProgress() {
+  def testHeartbeatRebalanceInProgress(): Unit = {
     val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
 
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols)
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols)
     val assignedMemberId = joinGroupResult.memberId
     val joinGroupError = joinGroupResult.error
     assertEquals(Errors.NONE, joinGroupError)
 
     EasyMock.reset(replicaManager)
-    val heartbeatResult = heartbeat(groupId, assignedMemberId, 2)
+    val heartbeatResult = heartbeat(groupId, assignedMemberId, 1)
     assertEquals(Errors.REBALANCE_IN_PROGRESS, heartbeatResult)
   }
 
   @Test
-  def testHeartbeatIllegalGeneration() {
+  def testHeartbeatIllegalGeneration(): Unit = {
     val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
 
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols)
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols)
     val assignedMemberId = joinGroupResult.memberId
     val joinGroupError = joinGroupResult.error
     assertEquals(Errors.NONE, joinGroupError)
@@ -458,10 +1436,10 @@ class GroupCoordinatorTest extends JUnitSuite {
   }
 
   @Test
-  def testValidHeartbeat() {
+  def testValidHeartbeat(): Unit = {
     val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
 
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols)
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols)
     val assignedConsumerId = joinGroupResult.memberId
     val generationId = joinGroupResult.generationId
     val joinGroupError = joinGroupResult.error
@@ -478,10 +1456,10 @@ class GroupCoordinatorTest extends JUnitSuite {
   }
 
   @Test
-  def testSessionTimeout() {
+  def testSessionTimeout(): Unit = {
     val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
 
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols)
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols)
     val assignedConsumerId = joinGroupResult.memberId
     val generationId = joinGroupResult.generationId
     val joinGroupError = joinGroupResult.error
@@ -492,7 +1470,8 @@ class GroupCoordinatorTest extends JUnitSuite {
     assertEquals(Errors.NONE, syncGroupError)
 
     EasyMock.reset(replicaManager)
-    EasyMock.expect(replicaManager.getPartition(new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, groupPartitionId))).andReturn(None)
+    EasyMock.expect(replicaManager.getPartition(new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, groupPartitionId)))
+      .andReturn(HostedPartition.None)
     EasyMock.expect(replicaManager.getMagic(EasyMock.anyObject())).andReturn(Some(RecordBatch.MAGIC_VALUE_V1)).anyTimes()
     EasyMock.replay(replicaManager)
 
@@ -504,11 +1483,11 @@ class GroupCoordinatorTest extends JUnitSuite {
   }
 
   @Test
-  def testHeartbeatMaintainsSession() {
+  def testHeartbeatMaintainsSession(): Unit = {
     val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
     val sessionTimeout = 1000
 
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols,
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols,
       rebalanceTimeout = sessionTimeout, sessionTimeout = sessionTimeout)
     val assignedConsumerId = joinGroupResult.memberId
     val generationId = joinGroupResult.generationId
@@ -533,40 +1512,40 @@ class GroupCoordinatorTest extends JUnitSuite {
   }
 
   @Test
-  def testCommitMaintainsSession() {
+  def testCommitMaintainsSession(): Unit = {
     val sessionTimeout = 1000
     val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
     val tp = new TopicPartition("topic", 0)
     val offset = offsetAndMetadata(0)
 
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols,
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols,
       rebalanceTimeout = sessionTimeout, sessionTimeout = sessionTimeout)
-    val assignedConsumerId = joinGroupResult.memberId
+    val assignedMemberId = joinGroupResult.memberId
     val generationId = joinGroupResult.generationId
     val joinGroupError = joinGroupResult.error
     assertEquals(Errors.NONE, joinGroupError)
 
     EasyMock.reset(replicaManager)
-    val (_, syncGroupError) = syncGroupLeader(groupId, generationId, assignedConsumerId, Map(assignedConsumerId -> Array[Byte]()))
+    val (_, syncGroupError) = syncGroupLeader(groupId, generationId, assignedMemberId, Map(assignedMemberId -> Array[Byte]()))
     assertEquals(Errors.NONE, syncGroupError)
 
     timer.advanceClock(sessionTimeout / 2)
 
     EasyMock.reset(replicaManager)
-    val commitOffsetResult = commitOffsets(groupId, assignedConsumerId, generationId, Map(tp -> offset))
+    val commitOffsetResult = commitOffsets(groupId, assignedMemberId, generationId, Map(tp -> offset))
     assertEquals(Errors.NONE, commitOffsetResult(tp))
 
     timer.advanceClock(sessionTimeout / 2 + 100)
 
     EasyMock.reset(replicaManager)
-    val heartbeatResult = heartbeat(groupId, assignedConsumerId, 1)
+    val heartbeatResult = heartbeat(groupId, assignedMemberId, 1)
     assertEquals(Errors.NONE, heartbeatResult)
   }
 
   @Test
-  def testSessionTimeoutDuringRebalance() {
+  def testSessionTimeoutDuringRebalance(): Unit = {
     // create a group with a single member
-    val firstJoinResult = joinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols,
+    val firstJoinResult = dynamicJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols,
       rebalanceTimeout = 2000, sessionTimeout = 1000)
     val firstMemberId = firstJoinResult.memberId
     val firstGenerationId = firstJoinResult.generationId
@@ -600,9 +1579,9 @@ class GroupCoordinatorTest extends JUnitSuite {
   }
 
   @Test
-  def testRebalanceCompletesBeforeMemberJoins() {
+  def testRebalanceCompletesBeforeMemberJoins(): Unit = {
     // create a group with a single member
-    val firstJoinResult = joinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols,
+    val firstJoinResult = staticJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, leaderInstanceId, protocolType, protocols,
       rebalanceTimeout = 1200, sessionTimeout = 1000)
     val firstMemberId = firstJoinResult.memberId
     val firstGenerationId = firstJoinResult.generationId
@@ -615,49 +1594,68 @@ class GroupCoordinatorTest extends JUnitSuite {
 
     // now have a new member join to trigger a rebalance
     EasyMock.reset(replicaManager)
+    val otherMemberSessionTimeout = DefaultSessionTimeout
     val otherJoinFuture = sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
 
     // send a couple heartbeats to keep the member alive while the rebalance finishes
-    timer.advanceClock(500)
-    EasyMock.reset(replicaManager)
-    var heartbeatResult = heartbeat(groupId, firstMemberId, firstGenerationId)
-    assertEquals(Errors.REBALANCE_IN_PROGRESS, heartbeatResult)
-
-    timer.advanceClock(500)
-    EasyMock.reset(replicaManager)
-    heartbeatResult = heartbeat(groupId, firstMemberId, firstGenerationId)
-    assertEquals(Errors.REBALANCE_IN_PROGRESS, heartbeatResult)
+    var expectedResultList = List(Errors.REBALANCE_IN_PROGRESS, Errors.REBALANCE_IN_PROGRESS)
+    for (expectedResult <- expectedResultList) {
+      timer.advanceClock(otherMemberSessionTimeout)
+      EasyMock.reset(replicaManager)
+      val heartbeatResult = heartbeat(groupId, firstMemberId, firstGenerationId)
+      assertEquals(expectedResult, heartbeatResult)
+    }
 
     // now timeout the rebalance
-    timer.advanceClock(500)
-    val otherJoinResult = await(otherJoinFuture, DefaultSessionTimeout+100)
+    timer.advanceClock(otherMemberSessionTimeout)
+    val otherJoinResult = await(otherJoinFuture, otherMemberSessionTimeout+100)
     val otherMemberId = otherJoinResult.memberId
     val otherGenerationId = otherJoinResult.generationId
     EasyMock.reset(replicaManager)
     val syncResult = syncGroupLeader(groupId, otherGenerationId, otherMemberId, Map(otherMemberId -> Array[Byte]()))
     assertEquals(Errors.NONE, syncResult._2)
 
-    // the unjoined member should be kicked out from the group
-    assertEquals(Errors.NONE, otherJoinResult.error)
+    // the unjoined static member should be remained in the group before session timeout.
+    assertEquals(Errors.NONE, otherJoinResult.error)
+    EasyMock.reset(replicaManager)
+    var heartbeatResult = heartbeat(groupId, firstMemberId, firstGenerationId)
+    assertEquals(Errors.ILLEGAL_GENERATION, heartbeatResult)
+
+    expectedResultList = List(Errors.NONE, Errors.NONE, Errors.REBALANCE_IN_PROGRESS)
+
+    // now session timeout the unjoined member. Still keeping the new member.
+    for (expectedResult <- expectedResultList) {
+      timer.advanceClock(otherMemberSessionTimeout)
+      EasyMock.reset(replicaManager)
+      heartbeatResult = heartbeat(groupId, otherMemberId, otherGenerationId)
+      assertEquals(expectedResult, heartbeatResult)
+    }
+
+    EasyMock.reset(replicaManager)
+    val otherRejoinGroupFuture = sendJoinGroup(groupId, otherMemberId, protocolType, protocols)
+    val otherReJoinResult = await(otherRejoinGroupFuture, otherMemberSessionTimeout+100)
+    assertEquals(Errors.NONE, otherReJoinResult.error)
+
     EasyMock.reset(replicaManager)
-    heartbeatResult = heartbeat(groupId, firstMemberId, firstGenerationId)
-    assertEquals(Errors.UNKNOWN_MEMBER_ID, heartbeatResult)
+    val otherRejoinGenerationId = otherReJoinResult.generationId
+    val reSyncResult = syncGroupLeader(groupId, otherRejoinGenerationId, otherMemberId, Map(otherMemberId -> Array[Byte]()))
+    assertEquals(Errors.NONE, reSyncResult._2)
 
     // 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 ( _ <-  1 to 20) {
       timer.advanceClock(500)
       EasyMock.reset(replicaManager)
-      heartbeatResult = heartbeat(groupId, otherMemberId, otherGenerationId)
+      heartbeatResult = heartbeat(groupId, otherMemberId, otherRejoinGenerationId)
       assertEquals(Errors.NONE, heartbeatResult)
     }
   }
 
   @Test
-  def testSyncGroupEmptyAssignment() {
+  def testSyncGroupEmptyAssignment(): Unit = {
     val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
 
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols)
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols)
     val assignedConsumerId = joinGroupResult.memberId
     val generationId = joinGroupResult.generationId
     val joinGroupError = joinGroupResult.error
@@ -675,7 +1673,7 @@ class GroupCoordinatorTest extends JUnitSuite {
   }
 
   @Test
-  def testSyncGroupNotCoordinator() {
+  def testSyncGroupNotCoordinator(): Unit = {
     val generation = 1
 
     val syncGroupResult = syncGroupFollower(otherGroupId, generation, memberId)
@@ -683,18 +1681,16 @@ class GroupCoordinatorTest extends JUnitSuite {
   }
 
   @Test
-  def testSyncGroupFromUnknownGroup() {
-    val generation = 1
-
-    val syncGroupResult = syncGroupFollower(groupId, generation, memberId)
+  def testSyncGroupFromUnknownGroup(): Unit = {
+    val syncGroupResult = syncGroupFollower(groupId, 1, memberId)
     assertEquals(Errors.UNKNOWN_MEMBER_ID, syncGroupResult._2)
   }
 
   @Test
-  def testSyncGroupFromUnknownMember() {
+  def testSyncGroupFromUnknownMember(): Unit = {
     val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
 
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols)
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols)
     val assignedConsumerId = joinGroupResult.memberId
     val generationId = joinGroupResult.generationId
     assertEquals(Errors.NONE, joinGroupResult.error)
@@ -711,10 +1707,10 @@ class GroupCoordinatorTest extends JUnitSuite {
   }
 
   @Test
-  def testSyncGroupFromIllegalGeneration() {
+  def testSyncGroupFromIllegalGeneration(): Unit = {
     val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
 
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols)
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols)
     val assignedConsumerId = joinGroupResult.memberId
     val generationId = joinGroupResult.generationId
     assertEquals(Errors.NONE, joinGroupResult.error)
@@ -726,12 +1722,12 @@ class GroupCoordinatorTest extends JUnitSuite {
   }
 
   @Test
-  def testJoinGroupFromUnchangedFollowerDoesNotRebalance() {
+  def testJoinGroupFromUnchangedFollowerDoesNotRebalance(): Unit = {
     // to get a group of two members:
     // 1. join and sync with a single member (because we can't immediately join with two members)
     // 2. join and sync with the first member and a new member
 
-    val firstJoinResult = joinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
+    val firstJoinResult = dynamicJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
     val firstMemberId = firstJoinResult.memberId
     val firstGenerationId = firstJoinResult.generationId
     assertEquals(firstMemberId, firstJoinResult.leaderId)
@@ -767,8 +1763,8 @@ class GroupCoordinatorTest extends JUnitSuite {
   }
 
   @Test
-  def testJoinGroupFromUnchangedLeaderShouldRebalance() {
-    val firstJoinResult = joinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
+  def testJoinGroupFromUnchangedLeaderShouldRebalance(): Unit = {
+    val firstJoinResult = dynamicJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
     val firstMemberId = firstJoinResult.memberId
     val firstGenerationId = firstJoinResult.generationId
     assertEquals(firstMemberId, firstJoinResult.leaderId)
@@ -795,8 +1791,8 @@ class GroupCoordinatorTest extends JUnitSuite {
     * which should remain stable throughout this test.
     */
   @Test
-  def testSecondMemberPartiallyJoinAndTimeout() {
-    val firstJoinResult = joinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
+  def testSecondMemberPartiallyJoinAndTimeout(): Unit = {
+    val firstJoinResult = dynamicJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
     val firstMemberId = firstJoinResult.memberId
     val firstGenerationId = firstJoinResult.generationId
     assertEquals(firstMemberId, firstJoinResult.leaderId)
@@ -852,7 +1848,7 @@ class GroupCoordinatorTest extends JUnitSuite {
     */
   private def setupGroupWithPendingMember(): JoinGroupResult = {
     // add the first member
-    val joinResult1 = joinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
+    val joinResult1 = dynamicJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
     assertGroupState(groupState = CompletingRebalance)
 
     // now the group is stable, with the one member that joined above
@@ -885,7 +1881,7 @@ class GroupCoordinatorTest extends JUnitSuite {
 
     // create a pending member in the group
     EasyMock.reset(replicaManager)
-    var pendingMember = joinGroupPartial(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols, sessionTimeout=100)
+    val pendingMember = joinGroupPartial(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols, sessionTimeout=100)
     assertEquals(1, groupCoordinator.groupManager.getGroup(groupId).get.numPending)
 
     // re-join the second existing member
@@ -902,7 +1898,7 @@ class GroupCoordinatorTest extends JUnitSuite {
     * operation
     */
   @Test
-  def testJoinGroupCompletionWhenPendingMemberJoins() {
+  def testJoinGroupCompletionWhenPendingMemberJoins(): Unit = {
     val pendingMember = setupGroupWithPendingMember()
 
     // compete join group for the pending member
@@ -920,7 +1916,7 @@ class GroupCoordinatorTest extends JUnitSuite {
     * cause the group to return to a CompletingRebalance state.
     */
   @Test
-  def testJoinGroupCompletionWhenPendingMemberTimesOut() {
+  def testJoinGroupCompletionWhenPendingMemberTimesOut(): Unit = {
     setupGroupWithPendingMember()
 
     // Advancing Clock by > 100 (session timeout for third and fourth member)
@@ -941,8 +1937,8 @@ class GroupCoordinatorTest extends JUnitSuite {
     val pending = setupGroupWithPendingMember()
 
     EasyMock.reset(replicaManager)
-    val leaveGroupResult = leaveGroup(groupId, pending.memberId)
-    assertEquals(Errors.NONE, leaveGroupResult)
+    val leaveGroupResults = singleLeaveGroup(groupId, pending.memberId)
+    verifyLeaveGroupResult(leaveGroupResults)
 
     assertGroupState(groupState = CompletingRebalance)
     assertEquals(2, group().allMembers.size)
@@ -971,17 +1967,17 @@ class GroupCoordinatorTest extends JUnitSuite {
                                sessionTimeout: Int = DefaultSessionTimeout,
                                rebalanceTimeout: Int = DefaultRebalanceTimeout): JoinGroupResult = {
     val requireKnownMemberId = true
-    val responseFuture = sendJoinGroup(groupId, memberId, protocolType, protocols, sessionTimeout, rebalanceTimeout, requireKnownMemberId)
+    val responseFuture = sendJoinGroup(groupId, memberId, protocolType, protocols, None, sessionTimeout, rebalanceTimeout, requireKnownMemberId)
     Await.result(responseFuture, Duration(rebalanceTimeout + 100, TimeUnit.MILLISECONDS))
   }
 
   @Test
-  def testLeaderFailureInSyncGroup() {
+  def testLeaderFailureInSyncGroup(): Unit = {
     // to get a group of two members:
     // 1. join and sync with a single member (because we can't immediately join with two members)
     // 2. join and sync with the first member and a new member
 
-    val firstJoinResult = joinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
+    val firstJoinResult = dynamicJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
     val firstMemberId = firstJoinResult.memberId
     val firstGenerationId = firstJoinResult.generationId
     assertEquals(firstMemberId, firstJoinResult.leaderId)
@@ -1008,10 +2004,10 @@ class GroupCoordinatorTest extends JUnitSuite {
 
     val nextGenerationId = joinResult.generationId
 
-    // with no leader SyncGroup, the follower's request should failure with an error indicating
+    // with no leader SyncGroup, the follower's request should fail with an error indicating
     // that it should rejoin
     EasyMock.reset(replicaManager)
-    val followerSyncFuture = sendSyncGroupFollower(groupId, nextGenerationId, otherJoinResult.memberId)
+    val followerSyncFuture = sendSyncGroupFollower(groupId, nextGenerationId, otherJoinResult.memberId, None)
 
     timer.advanceClock(DefaultSessionTimeout + 100)
 
@@ -1020,12 +2016,12 @@ class GroupCoordinatorTest extends JUnitSuite {
   }
 
   @Test
-  def testSyncGroupFollowerAfterLeader() {
+  def testSyncGroupFollowerAfterLeader(): Unit = {
     // to get a group of two members:
     // 1. join and sync with a single member (because we can't immediately join with two members)
     // 2. join and sync with the first member and a new member
 
-    val firstJoinResult = joinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
+    val firstJoinResult = dynamicJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
     val firstMemberId = firstJoinResult.memberId
     val firstGenerationId = firstJoinResult.generationId
     assertEquals(firstMemberId, firstJoinResult.leaderId)
@@ -1069,12 +2065,12 @@ class GroupCoordinatorTest extends JUnitSuite {
   }
 
   @Test
-  def testSyncGroupLeaderAfterFollower() {
+  def testSyncGroupLeaderAfterFollower(): Unit = {
     // to get a group of two members:
     // 1. join and sync with a single member (because we can't immediately join with two members)
     // 2. join and sync with the first member and a new member
 
-    val joinGroupResult = joinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
+    val joinGroupResult = dynamicJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
     val firstMemberId = joinGroupResult.memberId
     val firstGenerationId = joinGroupResult.generationId
     assertEquals(firstMemberId, joinGroupResult.leaderId)
@@ -1107,7 +2103,7 @@ class GroupCoordinatorTest extends JUnitSuite {
     assertEquals(firstMemberId, otherJoinResult.leaderId)
 
     EasyMock.reset(replicaManager)
-    val followerSyncFuture = sendSyncGroupFollower(groupId, nextGenerationId, followerId)
+    val followerSyncFuture = sendSyncGroupFollower(groupId, nextGenerationId, followerId, None)
 
     EasyMock.reset(replicaManager)
     val leaderSyncResult = syncGroupLeader(groupId, nextGenerationId, leaderId,
@@ -1121,7 +2117,7 @@ class GroupCoordinatorTest extends JUnitSuite {
   }
 
   @Test
-  def testCommitOffsetFromUnknownGroup() {
+  def testCommitOffsetFromUnknownGroup(): Unit = {
     val generationId = 1
     val tp = new TopicPartition("topic", 0)
     val offset = offsetAndMetadata(0)
@@ -1131,7 +2127,7 @@ class GroupCoordinatorTest extends JUnitSuite {
   }
 
   @Test
-  def testCommitOffsetWithDefaultGeneration() {
+  def testCommitOffsetWithDefaultGeneration(): Unit = {
     val tp = new TopicPartition("topic", 0)
     val offset = offsetAndMetadata(0)
 
@@ -1147,15 +2143,15 @@ class GroupCoordinatorTest extends JUnitSuite {
 
     // A group member joins
     val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols)
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols)
     val assignedMemberId = joinGroupResult.memberId
     val joinGroupError = joinGroupResult.error
     assertEquals(Errors.NONE, joinGroupError)
 
     // and leaves.
     EasyMock.reset(replicaManager)
-    val leaveGroupResult = leaveGroup(groupId, assignedMemberId)
-    assertEquals(Errors.NONE, leaveGroupResult)
+    val leaveGroupResults = singleLeaveGroup(groupId, assignedMemberId)
+    verifyLeaveGroupResult(leaveGroupResults)
 
     // The simple offset commit should now fail
     EasyMock.reset(replicaManager)
@@ -1193,7 +2189,7 @@ class GroupCoordinatorTest extends JUnitSuite {
   }
 
   @Test
-  def testCommitAndFetchOffsetsWithEmptyGroup() {
+  def testCommitAndFetchOffsetsWithEmptyGroup(): Unit = {
     // For backwards compatibility, the coordinator supports committing/fetching offsets with an empty groupId.
     // To allow inspection and removal of the empty group, we must also support DescribeGroups and DeleteGroups
 
@@ -1218,7 +2214,7 @@ class GroupCoordinatorTest extends JUnitSuite {
 
     EasyMock.reset(replicaManager)
     EasyMock.expect(replicaManager.getMagic(EasyMock.anyObject())).andStubReturn(Some(RecordBatch.CURRENT_MAGIC_VALUE))
-    EasyMock.expect(replicaManager.getPartition(groupTopicPartition)).andStubReturn(Some(partition))
+    EasyMock.expect(replicaManager.getPartition(groupTopicPartition)).andStubReturn(HostedPartition.Online(partition))
     EasyMock.expect(replicaManager.nonOfflinePartition(groupTopicPartition)).andStubReturn(Some(partition))
     EasyMock.replay(replicaManager, partition)
 
@@ -1231,7 +2227,7 @@ class GroupCoordinatorTest extends JUnitSuite {
   }
 
   @Test
-  def testBasicFetchTxnOffsets() {
+  def testBasicFetchTxnOffsets(): Unit = {
     val tp = new TopicPartition("topic", 0)
     val offset = offsetAndMetadata(0)
     val producerId = 1000L
@@ -1258,7 +2254,7 @@ class GroupCoordinatorTest extends JUnitSuite {
   }
 
   @Test
-  def testFetchTxnOffsetsWithAbort() {
+  def testFetchTxnOffsetsWithAbort(): Unit = {
     val tp = new TopicPartition("topic", 0)
     val offset = offsetAndMetadata(0)
     val producerId = 1000L
@@ -1282,7 +2278,7 @@ class GroupCoordinatorTest extends JUnitSuite {
   }
 
   @Test
-  def testFetchTxnOffsetsIgnoreSpuriousCommit() {
+  def testFetchTxnOffsetsIgnoreSpuriousCommit(): Unit = {
     val tp = new TopicPartition("topic", 0)
     val offset = offsetAndMetadata(0)
     val producerId = 1000L
@@ -1311,7 +2307,7 @@ class GroupCoordinatorTest extends JUnitSuite {
   }
 
   @Test
-  def testFetchTxnOffsetsOneProducerMultipleGroups() {
+  def testFetchTxnOffsetsOneProducerMultipleGroups(): Unit = {
     // One producer, two groups located on separate offsets topic partitions.
     // Both group have pending offset commits.
     // Marker for only one partition is received. That commit should be materialized while the other should not.
@@ -1390,7 +2386,7 @@ class GroupCoordinatorTest extends JUnitSuite {
   }
 
   @Test
-  def testFetchTxnOffsetsMultipleProducersOneGroup() {
+  def testFetchTxnOffsetsMultipleProducersOneGroup(): Unit = {
     // One group, two producers
     // Different producers will commit offsets for different partitions.
     // Each partition's offsets should be materialized when the corresponding producer's marker is received.
@@ -1464,7 +2460,7 @@ class GroupCoordinatorTest extends JUnitSuite {
   }
 
   @Test
-  def testFetchAllOffsets() {
+  def testFetchAllOffsets(): Unit = {
     val tp1 = new TopicPartition("topic", 0)
     val tp2 = new TopicPartition("topic", 1)
     val tp3 = new TopicPartition("other-topic", 0)
@@ -1490,12 +2486,12 @@ class GroupCoordinatorTest extends JUnitSuite {
   }
 
   @Test
-  def testCommitOffsetInCompletingRebalance() {
+  def testCommitOffsetInCompletingRebalance(): Unit = {
     val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
     val tp = new TopicPartition("topic", 0)
     val offset = offsetAndMetadata(0)
 
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols)
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols)
     val assignedMemberId = joinGroupResult.memberId
     val generationId = joinGroupResult.generationId
     val joinGroupError = joinGroupResult.error
@@ -1507,9 +2503,42 @@ class GroupCoordinatorTest extends JUnitSuite {
   }
 
   @Test
-  def testHeartbeatDuringRebalanceCausesRebalanceInProgress() {
+  def testCommitOffsetInCompletingRebalanceFromUnknownMemberId(): Unit = {
+    val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
+    val tp = new TopicPartition("topic", 0)
+    val offset = offsetAndMetadata(0)
+
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols)
+    val generationId = joinGroupResult.generationId
+    val joinGroupError = joinGroupResult.error
+    assertEquals(Errors.NONE, joinGroupError)
+
+    EasyMock.reset(replicaManager)
+    val commitOffsetResult = commitOffsets(groupId, memberId, generationId, Map(tp -> offset))
+    assertEquals(Errors.UNKNOWN_MEMBER_ID, commitOffsetResult(tp))
+  }
+
+  @Test
+  def testCommitOffsetInCompletingRebalanceFromIllegalGeneration(): Unit = {
+    val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
+    val tp = new TopicPartition("topic", 0)
+    val offset = offsetAndMetadata(0)
+
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols)
+    val assignedMemberId = joinGroupResult.memberId
+    val generationId = joinGroupResult.generationId
+    val joinGroupError = joinGroupResult.error
+    assertEquals(Errors.NONE, joinGroupError)
+
+    EasyMock.reset(replicaManager)
+    val commitOffsetResult = commitOffsets(groupId, assignedMemberId, generationId + 1, Map(tp -> offset))
+    assertEquals(Errors.ILLEGAL_GENERATION, commitOffsetResult(tp))
+  }
+
+  @Test
+  def testHeartbeatDuringRebalanceCausesRebalanceInProgress(): Unit = {
     // First start up a group (with a slightly larger timeout to give us time to heartbeat when the rebalance starts)
-    val joinGroupResult = joinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
+    val joinGroupResult = dynamicJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
     val assignedConsumerId = joinGroupResult.memberId
     val initialGenerationId = joinGroupResult.generationId
     val joinGroupError = joinGroupResult.error
@@ -1526,8 +2555,8 @@ class GroupCoordinatorTest extends JUnitSuite {
   }
 
   @Test
-  def testGenerationIdIncrementsOnRebalance() {
-    val joinGroupResult = joinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
+  def testGenerationIdIncrementsOnRebalance(): Unit = {
+    val joinGroupResult = dynamicJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
     val initialGenerationId = joinGroupResult.generationId
     val joinGroupError = joinGroupResult.error
     val memberId = joinGroupResult.memberId
@@ -1550,52 +2579,178 @@ class GroupCoordinatorTest extends JUnitSuite {
   }
 
   @Test
-  def testLeaveGroupWrongCoordinator() {
+  def testLeaveGroupWrongCoordinator(): Unit = {
     val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
 
-    val leaveGroupResult = leaveGroup(otherGroupId, memberId)
-    assertEquals(Errors.NOT_COORDINATOR, leaveGroupResult)
+    val leaveGroupResults = singleLeaveGroup(otherGroupId, memberId)
+    verifyLeaveGroupResult(leaveGroupResults, Errors.NOT_COORDINATOR)
   }
 
   @Test
-  def testLeaveGroupUnknownGroup() {
-
-    val leaveGroupResult = leaveGroup(groupId, memberId)
-    assertEquals(Errors.UNKNOWN_MEMBER_ID, leaveGroupResult)
+  def testLeaveGroupUnknownGroup(): Unit = {
+    val leaveGroupResults = singleLeaveGroup(groupId, memberId)
+    verifyLeaveGroupResult(leaveGroupResults, Errors.NONE, List(Errors.UNKNOWN_MEMBER_ID))
   }
 
   @Test
-  def testLeaveGroupUnknownConsumerExistingGroup() {
+  def testLeaveGroupUnknownConsumerExistingGroup(): Unit = {
     val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
     val otherMemberId = "consumerId"
 
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols)
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols)
     val joinGroupError = joinGroupResult.error
     assertEquals(Errors.NONE, joinGroupError)
 
     EasyMock.reset(replicaManager)
-    val leaveGroupResult = leaveGroup(groupId, otherMemberId)
-    assertEquals(Errors.UNKNOWN_MEMBER_ID, leaveGroupResult)
+    val leaveGroupResults = singleLeaveGroup(groupId, otherMemberId)
+    verifyLeaveGroupResult(leaveGroupResults, Errors.NONE, List(Errors.UNKNOWN_MEMBER_ID))
+  }
+
+  @Test
+  def testSingleLeaveDeadGroup(): Unit = {
+    val deadGroupId = "deadGroupId"
+
+    groupCoordinator.groupManager.addGroup(new GroupMetadata(deadGroupId, Dead, new MockTime()))
+    val leaveGroupResults = singleLeaveGroup(deadGroupId, memberId)
+    verifyLeaveGroupResult(leaveGroupResults, Errors.COORDINATOR_NOT_AVAILABLE)
   }
 
   @Test
-  def testValidLeaveGroup() {
+  def testBatchLeaveDeadGroup(): Unit = {
+    val deadGroupId = "deadGroupId"
+
+    groupCoordinator.groupManager.addGroup(new GroupMetadata(deadGroupId, Dead, new MockTime()))
+    val leaveGroupResults = batchLeaveGroup(deadGroupId,
+      List(new MemberIdentity().setMemberId(memberId), new MemberIdentity().setMemberId(memberId)))
+    verifyLeaveGroupResult(leaveGroupResults, Errors.COORDINATOR_NOT_AVAILABLE)
+  }
+
+  @Test
+  def testValidLeaveGroup(): Unit = {
     val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
 
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols)
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols)
     val assignedMemberId = joinGroupResult.memberId
     val joinGroupError = joinGroupResult.error
     assertEquals(Errors.NONE, joinGroupError)
 
     EasyMock.reset(replicaManager)
-    val leaveGroupResult = leaveGroup(groupId, assignedMemberId)
-    assertEquals(Errors.NONE, leaveGroupResult)
+    val leaveGroupResults = singleLeaveGroup(groupId, assignedMemberId)
+    verifyLeaveGroupResult(leaveGroupResults)
+  }
+
+  @Test
+  def testLeaveGroupWithFencedInstanceId(): Unit = {
+    val joinGroupResult = staticJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, leaderInstanceId, protocolType, protocolSuperset)
+    assertEquals(Errors.NONE, joinGroupResult.error)
+
+    EasyMock.reset(replicaManager)
+    val leaveGroupResults = singleLeaveGroup(groupId, "some_member", leaderInstanceId)
+    verifyLeaveGroupResult(leaveGroupResults, Errors.NONE, List(Errors.FENCED_INSTANCE_ID))
+  }
+
+  @Test
+  def testLeaveGroupStaticMemberWithUnknownMemberId(): Unit = {
+    val joinGroupResult = staticJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, leaderInstanceId, protocolType, protocolSuperset)
+    assertEquals(Errors.NONE, joinGroupResult.error)
+
+    EasyMock.reset(replicaManager)
+    // Having unknown member id will not affect the request processing.
+    val leaveGroupResults = singleLeaveGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, leaderInstanceId)
+    verifyLeaveGroupResult(leaveGroupResults, Errors.NONE, List(Errors.NONE))
+  }
+
+  @Test
+  def testStaticMembersValidBatchLeaveGroup(): Unit = {
+    staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId)
+
+    EasyMock.reset(replicaManager)
+    val leaveGroupResults = batchLeaveGroup(groupId, List(new MemberIdentity()
+      .setGroupInstanceId(leaderInstanceId.get), new MemberIdentity().setGroupInstanceId(followerInstanceId.get)))
+
+    verifyLeaveGroupResult(leaveGroupResults, Errors.NONE, List(Errors.NONE, Errors.NONE))
+  }
+
+  @Test
+  def testStaticMembersWrongCoordinatorBatchLeaveGroup(): Unit = {
+    staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId)
+
+    EasyMock.reset(replicaManager)
+    val leaveGroupResults = batchLeaveGroup("invalid-group", List(new MemberIdentity()
+      .setGroupInstanceId(leaderInstanceId.get), new MemberIdentity().setGroupInstanceId(followerInstanceId.get)))
+
+    verifyLeaveGroupResult(leaveGroupResults, Errors.NOT_COORDINATOR)
+  }
+
+  @Test
+  def testStaticMembersUnknownGroupBatchLeaveGroup(): Unit = {
+    val leaveGroupResults = batchLeaveGroup(groupId, List(new MemberIdentity()
+      .setGroupInstanceId(leaderInstanceId.get), new MemberIdentity().setGroupInstanceId(followerInstanceId.get)))
+
+    verifyLeaveGroupResult(leaveGroupResults, Errors.NONE, List(Errors.UNKNOWN_MEMBER_ID, Errors.UNKNOWN_MEMBER_ID))
+  }
+
+  @Test
+  def testStaticMembersFencedInstanceBatchLeaveGroup(): Unit = {
+    staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId)
+
+    EasyMock.reset(replicaManager)
+    val leaveGroupResults = batchLeaveGroup(groupId, List(new MemberIdentity()
+      .setGroupInstanceId(leaderInstanceId.get), new MemberIdentity()
+      .setGroupInstanceId(followerInstanceId.get)
+      .setMemberId("invalid-member")))
+
+    verifyLeaveGroupResult(leaveGroupResults, Errors.NONE, List(Errors.NONE, Errors.FENCED_INSTANCE_ID))
+  }
+
+  @Test
+  def testStaticMembersUnknownInstanceBatchLeaveGroup(): Unit = {
+    staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId)
+
+    EasyMock.reset(replicaManager)
+    val leaveGroupResults = batchLeaveGroup(groupId, List(new MemberIdentity()
+      .setGroupInstanceId("unknown-instance"), new MemberIdentity()
+      .setGroupInstanceId(followerInstanceId.get)))
+
+    verifyLeaveGroupResult(leaveGroupResults, Errors.NONE, List(Errors.UNKNOWN_MEMBER_ID, Errors.NONE))
+  }
+
+  @Test
+  def testPendingMemberBatchLeaveGroup(): Unit = {
+    val pendingMember = setupGroupWithPendingMember()
+
+    EasyMock.reset(replicaManager)
+    val leaveGroupResults = batchLeaveGroup(groupId, List(new MemberIdentity()
+      .setGroupInstanceId("unknown-instance"), new MemberIdentity()
+      .setMemberId(pendingMember.memberId)))
+
+    verifyLeaveGroupResult(leaveGroupResults, Errors.NONE, List(Errors.UNKNOWN_MEMBER_ID, Errors.NONE))
+  }
+
+  @Test
+  def testPendingMemberWithUnexpectedInstanceIdBatchLeaveGroup(): Unit = {
+    val pendingMember = setupGroupWithPendingMember()
+
+    EasyMock.reset(replicaManager)
+
+    // Bypass the FENCED_INSTANCE_ID check by defining pending member as a static member.
+    val instanceId = "instanceId"
+    val pendingMemberId = pendingMember.memberId
+    getGroup(groupId).addStaticMember(Option(instanceId), pendingMemberId)
+    val expectedException = intercept[IllegalStateException] {
+      batchLeaveGroup(groupId, List(new MemberIdentity().setGroupInstanceId("unknown-instance"),
+        new MemberIdentity().setGroupInstanceId(instanceId).setMemberId(pendingMemberId)))
+    }
+
+    val message = expectedException.getMessage
+    assertTrue(message.contains(instanceId))
+    assertTrue(message.contains(pendingMemberId))
   }
 
   @Test
-  def testListGroupsIncludesStableGroups() {
+  def testListGroupsIncludesStableGroups(): Unit = {
     val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols)
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols)
     val assignedMemberId = joinGroupResult.memberId
     val generationId = joinGroupResult.generationId
     assertEquals(Errors.NONE, joinGroupResult.error)
@@ -1612,9 +2767,9 @@ class GroupCoordinatorTest extends JUnitSuite {
   }
 
   @Test
-  def testListGroupsIncludesRebalancingGroups() {
+  def testListGroupsIncludesRebalancingGroups(): Unit = {
     val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols)
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols)
     assertEquals(Errors.NONE, joinGroupResult.error)
 
     val (error, groups) = groupCoordinator.handleListGroups()
@@ -1624,14 +2779,14 @@ class GroupCoordinatorTest extends JUnitSuite {
   }
 
   @Test
-  def testDescribeGroupWrongCoordinator() {
+  def testDescribeGroupWrongCoordinator(): Unit = {
     EasyMock.reset(replicaManager)
     val (error, _) = groupCoordinator.handleDescribeGroup(otherGroupId)
     assertEquals(Errors.NOT_COORDINATOR, error)
   }
 
   @Test
-  def testDescribeGroupInactiveGroup() {
+  def testDescribeGroupInactiveGroup(): Unit = {
     EasyMock.reset(replicaManager)
     val (error, summary) = groupCoordinator.handleDescribeGroup(groupId)
     assertEquals(Errors.NONE, error)
@@ -1639,9 +2794,30 @@ class GroupCoordinatorTest extends JUnitSuite {
   }
 
   @Test
-  def testDescribeGroupStable() {
-    val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols)
+  def testDescribeGroupStableForDynamicMember(): Unit = {
+    val joinGroupResult = dynamicJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
+    val assignedMemberId = joinGroupResult.memberId
+    val generationId = joinGroupResult.generationId
+    val joinGroupError = joinGroupResult.error
+    assertEquals(Errors.NONE, joinGroupError)
+
+    EasyMock.reset(replicaManager)
+    val syncGroupResult = syncGroupLeader(groupId, generationId, assignedMemberId, Map(assignedMemberId -> Array[Byte]()))
+
+    val syncGroupError = syncGroupResult._2
+    assertEquals(Errors.NONE, syncGroupError)
+
+    EasyMock.reset(replicaManager)
+    val (error, summary) = groupCoordinator.handleDescribeGroup(groupId)
+    assertEquals(Errors.NONE, error)
+    assertEquals(protocolType, summary.protocolType)
+    assertEquals("range", summary.protocol)
+    assertEquals(List(assignedMemberId), summary.members.map(_.memberId))
+  }
+
+  @Test
+  def testDescribeGroupStableForStaticMember(): Unit = {
+    val joinGroupResult = staticJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, leaderInstanceId, protocolType, protocols)
     val assignedMemberId = joinGroupResult.memberId
     val generationId = joinGroupResult.generationId
     val joinGroupError = joinGroupResult.error
@@ -1659,12 +2835,13 @@ class GroupCoordinatorTest extends JUnitSuite {
     assertEquals(protocolType, summary.protocolType)
     assertEquals("range", summary.protocol)
     assertEquals(List(assignedMemberId), summary.members.map(_.memberId))
+    assertEquals(List(leaderInstanceId), summary.members.map(_.groupInstanceId))
   }
 
   @Test
-  def testDescribeGroupRebalancing() {
+  def testDescribeGroupRebalancing(): Unit = {
     val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols)
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols)
     val joinGroupError = joinGroupResult.error
     assertEquals(Errors.NONE, joinGroupError)
 
@@ -1680,42 +2857,42 @@ class GroupCoordinatorTest extends JUnitSuite {
   }
 
   @Test
-  def testDeleteNonEmptyGroup() {
+  def testDeleteNonEmptyGroup(): Unit = {
     val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
-    joinGroup(groupId, memberId, protocolType, protocols)
+    dynamicJoinGroup(groupId, memberId, protocolType, protocols)
 
     val result = groupCoordinator.handleDeleteGroups(Set(groupId))
     assert(result.size == 1 && result.contains(groupId) && result.get(groupId).contains(Errors.NON_EMPTY_GROUP))
   }
 
   @Test
-  def testDeleteGroupWithInvalidGroupId() {
+  def testDeleteGroupWithInvalidGroupId(): Unit = {
     val invalidGroupId = null
     val result = groupCoordinator.handleDeleteGroups(Set(invalidGroupId))
     assert(result.size == 1 && result.contains(invalidGroupId) && result.get(invalidGroupId).contains(Errors.INVALID_GROUP_ID))
   }
 
   @Test
-  def testDeleteGroupWithWrongCoordinator() {
+  def testDeleteGroupWithWrongCoordinator(): Unit = {
     val result = groupCoordinator.handleDeleteGroups(Set(otherGroupId))
     assert(result.size == 1 && result.contains(otherGroupId) && result.get(otherGroupId).contains(Errors.NOT_COORDINATOR))
   }
 
   @Test
-  def testDeleteEmptyGroup() {
+  def testDeleteEmptyGroup(): Unit = {
     val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols)
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols)
 
     EasyMock.reset(replicaManager)
-    val leaveGroupResult = leaveGroup(groupId, joinGroupResult.memberId)
-    assertEquals(Errors.NONE, leaveGroupResult)
+    val leaveGroupResults = singleLeaveGroup(groupId, joinGroupResult.memberId)
+    verifyLeaveGroupResult(leaveGroupResults)
 
     val groupTopicPartition = new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, groupPartitionId)
     val partition: Partition = EasyMock.niceMock(classOf[Partition])
 
     EasyMock.reset(replicaManager)
     EasyMock.expect(replicaManager.getMagic(EasyMock.anyObject())).andStubReturn(Some(RecordBatch.CURRENT_MAGIC_VALUE))
-    EasyMock.expect(replicaManager.getPartition(groupTopicPartition)).andStubReturn(Some(partition))
+    EasyMock.expect(replicaManager.getPartition(groupTopicPartition)).andStubReturn(HostedPartition.Online(partition))
     EasyMock.expect(replicaManager.nonOfflinePartition(groupTopicPartition)).andStubReturn(Some(partition))
     EasyMock.replay(replicaManager, partition)
 
@@ -1724,10 +2901,10 @@ class GroupCoordinatorTest extends JUnitSuite {
   }
 
   @Test
-  def testDeleteEmptyGroupWithStoredOffsets() {
+  def testDeleteEmptyGroupWithStoredOffsets(): Unit = {
     val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
 
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols)
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols)
     val assignedMemberId = joinGroupResult.memberId
     val joinGroupError = joinGroupResult.error
     assertEquals(Errors.NONE, joinGroupError)
@@ -1748,15 +2925,15 @@ class GroupCoordinatorTest extends JUnitSuite {
     assertEquals(assignedMemberId, describeGroupResult._2.members.head.memberId)
 
     EasyMock.reset(replicaManager)
-    val leaveGroupResult = leaveGroup(groupId, assignedMemberId)
-    assertEquals(Errors.NONE, leaveGroupResult)
+    val leaveGroupResults = singleLeaveGroup(groupId, assignedMemberId)
+    verifyLeaveGroupResult(leaveGroupResults)
 
     val groupTopicPartition = new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, groupPartitionId)
     val partition: Partition = EasyMock.niceMock(classOf[Partition])
 
     EasyMock.reset(replicaManager)
     EasyMock.expect(replicaManager.getMagic(EasyMock.anyObject())).andStubReturn(Some(RecordBatch.CURRENT_MAGIC_VALUE))
-    EasyMock.expect(replicaManager.getPartition(groupTopicPartition)).andStubReturn(Some(partition))
+    EasyMock.expect(replicaManager.getPartition(groupTopicPartition)).andStubReturn(HostedPartition.Online(partition))
     EasyMock.expect(replicaManager.nonOfflinePartition(groupTopicPartition)).andStubReturn(Some(partition))
     EasyMock.replay(replicaManager, partition)
 
@@ -1767,7 +2944,211 @@ class GroupCoordinatorTest extends JUnitSuite {
   }
 
   @Test
-  def shouldDelayInitialRebalanceByGroupInitialRebalanceDelayOnEmptyGroup() {
+  def testDeleteOffsetOfNonExistingGroup(): Unit = {
+    val tp = new TopicPartition("foo", 0)
+    val (groupError, topics) = groupCoordinator.handleDeleteOffsets(groupId, Seq(tp))
+
+    assertEquals(Errors.GROUP_ID_NOT_FOUND, groupError)
+    assertTrue(topics.isEmpty)
+  }
+
+  @Test
+  def testDeleteOffsetOfNonEmptyNonConsumerGroup(): Unit = {
+    val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
+    dynamicJoinGroup(groupId, memberId, "My Protocol", protocols)
+    val tp = new TopicPartition("foo", 0)
+    val (groupError, topics) = groupCoordinator.handleDeleteOffsets(groupId, Seq(tp))
+
+    assertEquals(Errors.NON_EMPTY_GROUP, groupError)
+    assertTrue(topics.isEmpty)
+  }
+
+  @Test
+  def testDeleteOffsetOfEmptyNonConsumerGroup(): Unit = {
+    // join the group
+    val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
+
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, "My Protocol", protocols)
+    assertEquals(Errors.NONE, joinGroupResult.error)
+
+    EasyMock.reset(replicaManager)
+    val syncGroupResult = syncGroupLeader(groupId, joinGroupResult.generationId, joinGroupResult.leaderId, Map.empty)
+    assertEquals(Errors.NONE, syncGroupResult._2)
+
+    val t1p0 = new TopicPartition("foo", 0)
+    val t2p0 = new TopicPartition("bar", 0)
+    val offset = offsetAndMetadata(37)
+
+    EasyMock.reset(replicaManager)
+    val validOffsetCommitResult = commitOffsets(groupId, joinGroupResult.memberId, joinGroupResult.generationId,
+      Map(t1p0 -> offset, t2p0 -> offset))
+    assertEquals(Errors.NONE, validOffsetCommitResult(t1p0))
+    assertEquals(Errors.NONE, validOffsetCommitResult(t2p0))
+
+    // and leaves.
+    EasyMock.reset(replicaManager)
+    val leaveGroupResults = singleLeaveGroup(groupId, joinGroupResult.memberId)
+    verifyLeaveGroupResult(leaveGroupResults)
+
+    assertTrue(groupCoordinator.groupManager.getGroup(groupId).exists(_.is(Empty)))
+
+    val groupTopicPartition = new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, groupPartitionId)
+    val partition: Partition = EasyMock.niceMock(classOf[Partition])
+
+    EasyMock.reset(replicaManager)
+    EasyMock.expect(replicaManager.getMagic(EasyMock.anyObject())).andStubReturn(Some(RecordBatch.CURRENT_MAGIC_VALUE))
+    EasyMock.expect(replicaManager.getPartition(groupTopicPartition)).andStubReturn(HostedPartition.Online(partition))
+    EasyMock.expect(replicaManager.nonOfflinePartition(groupTopicPartition)).andStubReturn(Some(partition))
+    EasyMock.replay(replicaManager, partition)
+
+    val (groupError, topics) = groupCoordinator.handleDeleteOffsets(groupId, Seq(t1p0))
+
+    assertEquals(Errors.NONE, groupError)
+    assertEquals(1, topics.size)
+    assertEquals(Some(Errors.NONE), topics.get(t1p0))
+
+    val cachedOffsets = groupCoordinator.groupManager.getOffsets(groupId, Some(Seq(t1p0, t2p0)))
+
+    assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), cachedOffsets.get(t1p0).map(_.offset))
+    assertEquals(Some(offset.offset), cachedOffsets.get(t2p0).map(_.offset))
+  }
+
+  @Test
+  def testDeleteOffsetOfConsumerGroupWithUnparsableProtocol(): Unit = {
+    val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols)
+
+    EasyMock.reset(replicaManager)
+    val syncGroupResult = syncGroupLeader(groupId, joinGroupResult.generationId, joinGroupResult.leaderId, Map.empty)
+    assertEquals(Errors.NONE, syncGroupResult._2)
+
+    val tp = new TopicPartition("foo", 0)
+    val offset = offsetAndMetadata(37)
+
+    EasyMock.reset(replicaManager)
+    val validOffsetCommitResult = commitOffsets(groupId, joinGroupResult.memberId, joinGroupResult.generationId,
+      Map(tp -> offset))
+    assertEquals(Errors.NONE, validOffsetCommitResult(tp))
+
+    val (groupError, topics) = groupCoordinator.handleDeleteOffsets(groupId, Seq(tp))
+
+    assertEquals(Errors.NONE, groupError)
+    assertEquals(1, topics.size)
+    assertEquals(Some(Errors.GROUP_SUBSCRIBED_TO_TOPIC), topics.get(tp))
+  }
+
+  @Test
+  def testDeleteOffsetOfDeadConsumerGroup(): Unit = {
+    val group = new GroupMetadata(groupId, Dead, new MockTime())
+    group.protocolType = Some(protocolType)
+    groupCoordinator.groupManager.addGroup(group)
+
+    val tp = new TopicPartition("foo", 0)
+    val (groupError, topics) = groupCoordinator.handleDeleteOffsets(groupId, Seq(tp))
+
+    assertEquals(Errors.GROUP_ID_NOT_FOUND, groupError)
+    assertTrue(topics.isEmpty)
+  }
+
+  @Test
+  def testDeleteOffsetOfEmptyConsumerGroup(): Unit = {
+    // join the group
+    val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols)
+    assertEquals(Errors.NONE, joinGroupResult.error)
+
+    EasyMock.reset(replicaManager)
+    val syncGroupResult = syncGroupLeader(groupId, joinGroupResult.generationId, joinGroupResult.leaderId, Map.empty)
+    assertEquals(Errors.NONE, syncGroupResult._2)
+
+    val t1p0 = new TopicPartition("foo", 0)
+    val t2p0 = new TopicPartition("bar", 0)
+    val offset = offsetAndMetadata(37)
+
+    EasyMock.reset(replicaManager)
+    val validOffsetCommitResult = commitOffsets(groupId, joinGroupResult.memberId, joinGroupResult.generationId,
+      Map(t1p0 -> offset, t2p0 -> offset))
+    assertEquals(Errors.NONE, validOffsetCommitResult(t1p0))
+    assertEquals(Errors.NONE, validOffsetCommitResult(t2p0))
+
+    // and leaves.
+    EasyMock.reset(replicaManager)
+    val leaveGroupResults = singleLeaveGroup(groupId, joinGroupResult.memberId)
+    verifyLeaveGroupResult(leaveGroupResults)
+
+    assertTrue(groupCoordinator.groupManager.getGroup(groupId).exists(_.is(Empty)))
+
+    val groupTopicPartition = new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, groupPartitionId)
+    val partition: Partition = EasyMock.niceMock(classOf[Partition])
+
+    EasyMock.reset(replicaManager)
+    EasyMock.expect(replicaManager.getMagic(EasyMock.anyObject())).andStubReturn(Some(RecordBatch.CURRENT_MAGIC_VALUE))
+    EasyMock.expect(replicaManager.getPartition(groupTopicPartition)).andStubReturn(HostedPartition.Online(partition))
+    EasyMock.expect(replicaManager.nonOfflinePartition(groupTopicPartition)).andStubReturn(Some(partition))
+    EasyMock.replay(replicaManager, partition)
+
+    val (groupError, topics) = groupCoordinator.handleDeleteOffsets(groupId, Seq(t1p0))
+
+    assertEquals(Errors.NONE, groupError)
+    assertEquals(1, topics.size)
+    assertEquals(Some(Errors.NONE), topics.get(t1p0))
+
+    val cachedOffsets = groupCoordinator.groupManager.getOffsets(groupId, Some(Seq(t1p0, t2p0)))
+
+    assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), cachedOffsets.get(t1p0).map(_.offset))
+    assertEquals(Some(offset.offset), cachedOffsets.get(t2p0).map(_.offset))
+  }
+
+  @Test
+  def testDeleteOffsetOfStableConsumerGroup(): Unit = {
+    // join the group
+    val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
+    val subscription = new Subscription(List("bar").asJava)
+
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType,
+      List(("protocol", ConsumerProtocol.serializeSubscription(subscription).array())))
+    assertEquals(Errors.NONE, joinGroupResult.error)
+
+    EasyMock.reset(replicaManager)
+    val syncGroupResult = syncGroupLeader(groupId, joinGroupResult.generationId, joinGroupResult.leaderId, Map.empty)
+    assertEquals(Errors.NONE, syncGroupResult._2)
+
+    val t1p0 = new TopicPartition("foo", 0)
+    val t2p0 = new TopicPartition("bar", 0)
+    val offset = offsetAndMetadata(37)
+
+    EasyMock.reset(replicaManager)
+    val validOffsetCommitResult = commitOffsets(groupId, joinGroupResult.memberId, joinGroupResult.generationId,
+      Map(t1p0 -> offset, t2p0 -> offset))
+    assertEquals(Errors.NONE, validOffsetCommitResult(t1p0))
+    assertEquals(Errors.NONE, validOffsetCommitResult(t2p0))
+
+    assertTrue(groupCoordinator.groupManager.getGroup(groupId).exists(_.is(Stable)))
+
+    val groupTopicPartition = new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, groupPartitionId)
+    val partition: Partition = EasyMock.niceMock(classOf[Partition])
+
+    EasyMock.reset(replicaManager)
+    EasyMock.expect(replicaManager.getMagic(EasyMock.anyObject())).andStubReturn(Some(RecordBatch.CURRENT_MAGIC_VALUE))
+    EasyMock.expect(replicaManager.getPartition(groupTopicPartition)).andStubReturn(HostedPartition.Online(partition))
+    EasyMock.expect(replicaManager.nonOfflinePartition(groupTopicPartition)).andStubReturn(Some(partition))
+    EasyMock.replay(replicaManager, partition)
+
+    val (groupError, topics) = groupCoordinator.handleDeleteOffsets(groupId, Seq(t1p0, t2p0))
+
+    assertEquals(Errors.NONE, groupError)
+    assertEquals(2, topics.size)
+    assertEquals(Some(Errors.NONE), topics.get(t1p0))
+    assertEquals(Some(Errors.GROUP_SUBSCRIBED_TO_TOPIC), topics.get(t2p0))
+
+    val cachedOffsets = groupCoordinator.groupManager.getOffsets(groupId, Some(Seq(t1p0, t2p0)))
+
+    assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), cachedOffsets.get(t1p0).map(_.offset))
+    assertEquals(Some(offset.offset), cachedOffsets.get(t2p0).map(_.offset))
+  }
+
+  @Test
+  def shouldDelayInitialRebalanceByGroupInitialRebalanceDelayOnEmptyGroup(): Unit = {
     val firstJoinFuture = sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
     timer.advanceClock(GroupInitialRebalanceDelay / 2)
     verifyDelayedTaskNotCompleted(firstJoinFuture)
@@ -1786,7 +3167,7 @@ class GroupCoordinatorTest extends JUnitSuite {
   }
 
   @Test
-  def shouldResetRebalanceDelayWhenNewMemberJoinsGroupInInitialRebalance() {
+  def shouldResetRebalanceDelayWhenNewMemberJoinsGroupInInitialRebalance(): Unit = {
     val rebalanceTimeout = GroupInitialRebalanceDelay * 3
     val firstMemberJoinFuture = sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols, rebalanceTimeout = rebalanceTimeout)
     EasyMock.reset(replicaManager)
@@ -1810,7 +3191,7 @@ class GroupCoordinatorTest extends JUnitSuite {
   }
 
   @Test
-  def shouldDelayRebalanceUptoRebalanceTimeout() {
+  def shouldDelayRebalanceUptoRebalanceTimeout(): Unit = {
     val rebalanceTimeout = GroupInitialRebalanceDelay * 2
     val firstMemberJoinFuture = sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols, rebalanceTimeout = rebalanceTimeout)
     EasyMock.reset(replicaManager)
@@ -1836,6 +3217,35 @@ class GroupCoordinatorTest extends JUnitSuite {
     assertEquals(Errors.NONE, thirdResult.error)
   }
 
+  @Test
+  def testCompleteHeartbeatWithGroupDead(): Unit = {
+    val rebalanceResult = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId)
+    EasyMock.reset(replicaManager)
+    heartbeat(groupId, rebalanceResult.leaderId, rebalanceResult.generation)
+    val group = getGroup(groupId)
+    group.transitionTo(Dead)
+    val leaderMemberId = rebalanceResult.leaderId
+    assertTrue(groupCoordinator.tryCompleteHeartbeat(group, leaderMemberId, false, () => true))
+    groupCoordinator.onExpireHeartbeat(group, leaderMemberId, false)
+    assertTrue(group.has(leaderMemberId))
+  }
+
+  @Test
+  def testCompleteHeartbeatWithMemberAlreadyRemoved(): Unit = {
+    val rebalanceResult = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId)
+    EasyMock.reset(replicaManager)
+    heartbeat(groupId, rebalanceResult.leaderId, rebalanceResult.generation)
+    val group = getGroup(groupId)
+    val leaderMemberId = rebalanceResult.leaderId
+    group.remove(leaderMemberId)
+    assertTrue(groupCoordinator.tryCompleteHeartbeat(group, leaderMemberId, false, () => true))
+  }
+
+  private def getGroup(groupId: String): GroupMetadata = {
+    val groupOpt = groupCoordinator.groupManager.getGroup(groupId)
+    assertTrue(groupOpt.isDefined)
+    groupOpt.get
+  }
   private def setupJoinGroupCallback: (Future[JoinGroupResult], JoinGroupCallback) = {
     val responsePromise = Promise[JoinGroupResult]
     val responseFuture = responsePromise.future
@@ -1846,8 +3256,8 @@ class GroupCoordinatorTest extends JUnitSuite {
   private def setupSyncGroupCallback: (Future[SyncGroupCallbackParams], SyncGroupCallback) = {
     val responsePromise = Promise[SyncGroupCallbackParams]
     val responseFuture = responsePromise.future
-    val responseCallback: SyncGroupCallback = (assignment, error) =>
-      responsePromise.success((assignment, error))
+    val responseCallback: SyncGroupCallback = syncGroupResult =>
+      responsePromise.success(syncGroupResult.memberAssignment, syncGroupResult.error)
     (responseFuture, responseCallback)
   }
 
@@ -1865,10 +3275,18 @@ class GroupCoordinatorTest extends JUnitSuite {
     (responseFuture, responseCallback)
   }
 
+  private def setupLeaveGroupCallback: (Future[LeaveGroupResult], LeaveGroupCallback) = {
+    val responsePromise = Promise[LeaveGroupResult]
+    val responseFuture = responsePromise.future
+    val responseCallback: LeaveGroupCallback = result => responsePromise.success(result)
+    (responseFuture, responseCallback)
+  }
+
   private def sendJoinGroup(groupId: String,
                             memberId: String,
                             protocolType: String,
                             protocols: List[(String, Array[Byte])],
+                            groupInstanceId: Option[String] = None,
                             sessionTimeout: Int = DefaultSessionTimeout,
                             rebalanceTimeout: Int = DefaultRebalanceTimeout,
                             requireKnownMemberId: Boolean = false): Future[JoinGroupResult] = {
@@ -1876,8 +3294,8 @@ class GroupCoordinatorTest extends JUnitSuite {
 
     EasyMock.replay(replicaManager)
 
-    groupCoordinator.handleJoinGroup(groupId, memberId, requireKnownMemberId, "clientId", "clientHost", rebalanceTimeout, sessionTimeout,
-      protocolType, protocols, responseCallback)
+    groupCoordinator.handleJoinGroup(groupId, memberId, groupInstanceId,
+      requireKnownMemberId, "clientId", "clientHost", rebalanceTimeout, sessionTimeout, protocolType, protocols, responseCallback)
     responseFuture
   }
 
@@ -1885,6 +3303,7 @@ class GroupCoordinatorTest extends JUnitSuite {
   private def sendSyncGroupLeader(groupId: String,
                                   generation: Int,
                                   leaderId: String,
+                                  groupInstanceId: Option[String],
                                   assignment: Map[String, Array[Byte]]): Future[SyncGroupCallbackParams] = {
     val (responseFuture, responseCallback) = setupSyncGroupCallback
 
@@ -1893,7 +3312,7 @@ class GroupCoordinatorTest extends JUnitSuite {
     EasyMock.expect(replicaManager.appendRecords(EasyMock.anyLong(),
       EasyMock.anyShort(),
       internalTopicsAllowed = EasyMock.eq(true),
-      isFromClient = EasyMock.eq(false),
+      origin = EasyMock.eq(AppendOrigin.Coordinator),
       EasyMock.anyObject().asInstanceOf[Map[TopicPartition, MemoryRecords]],
       EasyMock.capture(capturedArgument),
       EasyMock.anyObject().asInstanceOf[Option[ReentrantLock]],
@@ -1906,29 +3325,31 @@ class GroupCoordinatorTest extends JUnitSuite {
     EasyMock.expect(replicaManager.getMagic(EasyMock.anyObject())).andReturn(Some(RecordBatch.MAGIC_VALUE_V1)).anyTimes()
     EasyMock.replay(replicaManager)
 
-    groupCoordinator.handleSyncGroup(groupId, generation, leaderId, assignment, responseCallback)
+    groupCoordinator.handleSyncGroup(groupId, generation, leaderId, groupInstanceId, assignment, responseCallback)
     responseFuture
   }
 
   private def sendSyncGroupFollower(groupId: String,
                                     generation: Int,
-                                    memberId: String): Future[SyncGroupCallbackParams] = {
+                                    memberId: String,
+                                    groupInstanceId: Option[String] = None): Future[SyncGroupCallbackParams] = {
     val (responseFuture, responseCallback) = setupSyncGroupCallback
 
     EasyMock.replay(replicaManager)
 
-    groupCoordinator.handleSyncGroup(groupId, generation, memberId, Map.empty[String, Array[Byte]], responseCallback)
+    groupCoordinator.handleSyncGroup(groupId, generation, memberId, groupInstanceId,
+      Map.empty[String, Array[Byte]], responseCallback)
     responseFuture
   }
 
-  private def joinGroup(groupId: String,
-                        memberId: String,
-                        protocolType: String,
-                        protocols: List[(String, Array[Byte])],
-                        sessionTimeout: Int = DefaultSessionTimeout,
-                        rebalanceTimeout: Int = DefaultRebalanceTimeout): JoinGroupResult = {
+  private def dynamicJoinGroup(groupId: String,
+                               memberId: String,
+                               protocolType: String,
+                               protocols: List[(String, Array[Byte])],
+                               sessionTimeout: Int = DefaultSessionTimeout,
+                               rebalanceTimeout: Int = DefaultRebalanceTimeout): JoinGroupResult = {
     val requireKnownMemberId = true
-    var responseFuture = sendJoinGroup(groupId, memberId, protocolType, protocols, sessionTimeout, rebalanceTimeout, requireKnownMemberId)
+    var responseFuture = sendJoinGroup(groupId, memberId, protocolType, protocols, None, sessionTimeout, rebalanceTimeout, requireKnownMemberId)
 
     // Since member id is required, we need another bounce to get the successful join group result.
     if (memberId == JoinGroupRequest.UNKNOWN_MEMBER_ID && requireKnownMemberId) {
@@ -1938,19 +3359,34 @@ class GroupCoordinatorTest extends JUnitSuite {
         return joinGroupResult
       }
       EasyMock.reset(replicaManager)
-      responseFuture = sendJoinGroup(groupId, joinGroupResult.memberId, protocolType, protocols, sessionTimeout, rebalanceTimeout, requireKnownMemberId)
+      responseFuture = sendJoinGroup(groupId, joinGroupResult.memberId, protocolType, protocols, None, sessionTimeout, rebalanceTimeout, requireKnownMemberId)
     }
     timer.advanceClock(GroupInitialRebalanceDelay + 1)
     // should only have to wait as long as session timeout, but allow some extra time in case of an unexpected delay
     Await.result(responseFuture, Duration(rebalanceTimeout + 100, TimeUnit.MILLISECONDS))
   }
 
+  private def staticJoinGroup(groupId: String,
+                              memberId: String,
+                              groupInstanceId: Option[String],
+                              protocolType: String,
+                              protocols: List[(String, Array[Byte])],
+                              clockAdvance: Int = GroupInitialRebalanceDelay + 1,
+                              sessionTimeout: Int = DefaultSessionTimeout,
+                              rebalanceTimeout: Int = DefaultRebalanceTimeout): JoinGroupResult = {
+    val responseFuture = sendJoinGroup(groupId, memberId, protocolType, protocols, groupInstanceId, sessionTimeout, rebalanceTimeout)
+
+    timer.advanceClock(clockAdvance)
+    // should only have to wait as long as session timeout, but allow some extra time in case of an unexpected delay
+    Await.result(responseFuture, Duration(rebalanceTimeout + 100, TimeUnit.MILLISECONDS))
+  }
 
   private def syncGroupFollower(groupId: String,
                                 generationId: Int,
                                 memberId: String,
+                                groupInstanceId: Option[String] = None,
                                 sessionTimeout: Int = DefaultSessionTimeout): SyncGroupCallbackParams = {
-    val responseFuture = sendSyncGroupFollower(groupId, generationId, memberId)
+    val responseFuture = sendSyncGroupFollower(groupId, generationId, memberId, groupInstanceId)
     Await.result(responseFuture, Duration(sessionTimeout + 100, TimeUnit.MILLISECONDS))
   }
 
@@ -1958,19 +3394,21 @@ class GroupCoordinatorTest extends JUnitSuite {
                               generationId: Int,
                               memberId: String,
                               assignment: Map[String, Array[Byte]],
+                              groupInstanceId: Option[String] = None,
                               sessionTimeout: Int = DefaultSessionTimeout): SyncGroupCallbackParams = {
-    val responseFuture = sendSyncGroupLeader(groupId, generationId, memberId, assignment)
+    val responseFuture = sendSyncGroupLeader(groupId, generationId, memberId, groupInstanceId, assignment)
     Await.result(responseFuture, Duration(sessionTimeout + 100, TimeUnit.MILLISECONDS))
   }
 
   private def heartbeat(groupId: String,
                         consumerId: String,
-                        generationId: Int): HeartbeatCallbackParams = {
+                        generationId: Int,
+                        groupInstanceId: Option[String] = None): HeartbeatCallbackParams = {
     val (responseFuture, responseCallback) = setupHeartbeatCallback
 
     EasyMock.replay(replicaManager)
 
-    groupCoordinator.handleHeartbeat(groupId, consumerId, generationId, responseCallback)
+    groupCoordinator.handleHeartbeat(groupId, consumerId, groupInstanceId, generationId, responseCallback)
     Await.result(responseFuture, Duration(40, TimeUnit.MILLISECONDS))
   }
 
@@ -1979,9 +3417,10 @@ class GroupCoordinatorTest extends JUnitSuite {
   }
 
   private def commitOffsets(groupId: String,
-                            consumerId: String,
+                            memberId: String,
                             generationId: Int,
-                            offsets: Map[TopicPartition, OffsetAndMetadata]): CommitOffsetCallbackParams = {
+                            offsets: Map[TopicPartition, OffsetAndMetadata],
+                            groupInstanceId: Option[String] = None): CommitOffsetCallbackParams = {
     val (responseFuture, responseCallback) = setupCommitOffsetsCallback
 
     val capturedArgument: Capture[scala.collection.Map[TopicPartition, PartitionResponse] => Unit] = EasyMock.newCapture()
@@ -1989,7 +3428,7 @@ class GroupCoordinatorTest extends JUnitSuite {
     EasyMock.expect(replicaManager.appendRecords(EasyMock.anyLong(),
       EasyMock.anyShort(),
       internalTopicsAllowed = EasyMock.eq(true),
-      isFromClient = EasyMock.eq(false),
+      origin = EasyMock.eq(AppendOrigin.Coordinator),
       EasyMock.anyObject().asInstanceOf[Map[TopicPartition, MemoryRecords]],
       EasyMock.capture(capturedArgument),
       EasyMock.anyObject().asInstanceOf[Option[ReentrantLock]],
@@ -2004,7 +3443,7 @@ class GroupCoordinatorTest extends JUnitSuite {
     EasyMock.expect(replicaManager.getMagic(EasyMock.anyObject())).andReturn(Some(RecordBatch.MAGIC_VALUE_V1)).anyTimes()
     EasyMock.replay(replicaManager)
 
-    groupCoordinator.handleCommitOffsets(groupId, consumerId, generationId, offsets, responseCallback)
+    groupCoordinator.handleCommitOffsets(groupId, memberId, groupInstanceId, generationId, offsets, responseCallback)
     Await.result(responseFuture, Duration(40, TimeUnit.MILLISECONDS))
   }
 
@@ -2019,7 +3458,7 @@ class GroupCoordinatorTest extends JUnitSuite {
     EasyMock.expect(replicaManager.appendRecords(EasyMock.anyLong(),
       EasyMock.anyShort(),
       internalTopicsAllowed = EasyMock.eq(true),
-      isFromClient = EasyMock.eq(false),
+      origin = EasyMock.eq(AppendOrigin.Coordinator),
       EasyMock.anyObject().asInstanceOf[Map[TopicPartition, MemoryRecords]],
       EasyMock.capture(capturedArgument),
       EasyMock.anyObject().asInstanceOf[Option[ReentrantLock]],
@@ -2039,14 +3478,26 @@ class GroupCoordinatorTest extends JUnitSuite {
     result
   }
 
-  private def leaveGroup(groupId: String, consumerId: String): LeaveGroupCallbackParams = {
-    val (responseFuture, responseCallback) = setupHeartbeatCallback
+  private def singleLeaveGroup(groupId: String,
+                               consumerId: String,
+                               groupInstanceId: Option[String] = None): LeaveGroupResult = {
+    val singleMemberIdentity = List(
+      new MemberIdentity()
+        .setMemberId(consumerId)
+        .setGroupInstanceId(groupInstanceId.orNull))
+    batchLeaveGroup(groupId, singleMemberIdentity)
+  }
 
-    EasyMock.expect(replicaManager.getPartition(new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, groupPartitionId))).andReturn(None)
+  private def batchLeaveGroup(groupId: String,
+                              memberIdentities: List[MemberIdentity]): LeaveGroupResult = {
+    val (responseFuture, responseCallback) = setupLeaveGroupCallback
+
+    EasyMock.expect(replicaManager.getPartition(new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, groupPartitionId)))
+      .andReturn(HostedPartition.None)
     EasyMock.expect(replicaManager.getMagic(EasyMock.anyObject())).andReturn(Some(RecordBatch.MAGIC_VALUE_V1)).anyTimes()
     EasyMock.replay(replicaManager)
 
-    groupCoordinator.handleLeaveGroup(groupId, consumerId, responseCallback)
+    groupCoordinator.handleLeaveGroup(groupId, memberIdentities, responseCallback)
     Await.result(responseFuture, Duration(40, TimeUnit.MILLISECONDS))
   }
 
@@ -2060,5 +3511,18 @@ class GroupCoordinatorTest extends JUnitSuite {
   private def offsetAndMetadata(offset: Long): OffsetAndMetadata = {
     OffsetAndMetadata(offset, "", timer.time.milliseconds())
   }
+}
 
+object GroupCoordinatorTest {
+  def verifyLeaveGroupResult(leaveGroupResult: LeaveGroupResult,
+                             expectedTopLevelError: Errors = Errors.NONE,
+                             expectedMemberLevelErrors: List[Errors] = List.empty): Unit = {
+    assertEquals(expectedTopLevelError, leaveGroupResult.topLevelError)
+    if (expectedMemberLevelErrors.nonEmpty) {
+      assertEquals(expectedMemberLevelErrors.size, leaveGroupResult.memberResponses.size)
+      for (i <- expectedMemberLevelErrors.indices) {
+            assertEquals(expectedMemberLevelErrors(i), leaveGroupResult.memberResponses(i).error)
+          }
+    }
+  }
 }
diff --git a/core/src/test/scala/unit/kafka/coordinator/group/GroupMetadataManagerTest.scala b/core/src/test/scala/unit/kafka/coordinator/group/GroupMetadataManagerTest.scala
index 0287888ad1f87..b3a7f661471ad 100644
--- a/core/src/test/scala/unit/kafka/coordinator/group/GroupMetadataManagerTest.scala
+++ b/core/src/test/scala/unit/kafka/coordinator/group/GroupMetadataManagerTest.scala
@@ -17,36 +17,39 @@
 
 package kafka.coordinator.group
 
+import java.lang.management.ManagementFactory
+import java.nio.ByteBuffer
+import java.util.concurrent.locks.ReentrantLock
+import java.util.{Collections, Optional}
+
+import com.yammer.metrics.Metrics
+import com.yammer.metrics.core.Gauge
+import javax.management.ObjectName
 import kafka.api._
 import kafka.cluster.Partition
 import kafka.common.OffsetAndMetadata
-import kafka.log.{Log, LogAppendInfo}
-import kafka.server.{FetchDataInfo, KafkaConfig, LogOffsetMetadata, ReplicaManager}
-import kafka.utils.TestUtils.fail
+import kafka.log.{AppendOrigin, Log, LogAppendInfo}
+import kafka.server.{FetchDataInfo, FetchLogEnd, HostedPartition, KafkaConfig, LogOffsetMetadata, ReplicaManager}
 import kafka.utils.{KafkaScheduler, MockTime, TestUtils}
+import kafka.zk.KafkaZkClient
+import org.apache.kafka.clients.consumer.ConsumerPartitionAssignor
+import org.apache.kafka.clients.consumer.ConsumerPartitionAssignor.Subscription
 import org.apache.kafka.clients.consumer.internals.ConsumerProtocol
-import org.apache.kafka.clients.consumer.internals.PartitionAssignor.Subscription
 import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.internals.Topic
+import org.apache.kafka.common.metrics.{JmxReporter, Metrics => kMetrics}
 import org.apache.kafka.common.protocol.Errors
 import org.apache.kafka.common.record._
 import org.apache.kafka.common.requests.OffsetFetchResponse
 import org.apache.kafka.common.requests.ProduceResponse.PartitionResponse
+import org.apache.kafka.common.utils.Utils
 import org.easymock.{Capture, EasyMock, IAnswer}
 import org.junit.Assert.{assertEquals, assertFalse, assertNull, assertTrue}
 import org.junit.{Before, Test}
-import java.nio.ByteBuffer
-import java.util.Collections
-import java.util.Optional
-
-import com.yammer.metrics.Metrics
-import com.yammer.metrics.core.Gauge
-import org.apache.kafka.common.internals.Topic
+import org.scalatest.Assertions.fail
 
 import scala.collection.JavaConverters._
 import scala.collection._
-import java.util.concurrent.locks.ReentrantLock
-
-import kafka.zk.KafkaZkClient
 
 class GroupMetadataManagerTest {
 
@@ -57,8 +60,10 @@ class GroupMetadataManagerTest {
   var zkClient: KafkaZkClient = null
   var partition: Partition = null
   var defaultOffsetRetentionMs = Long.MaxValue
+  var metrics: kMetrics = null
 
   val groupId = "foo"
+  val groupInstanceId = Some("bar")
   val groupPartitionId = 0
   val groupTopicPartition = new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, groupPartitionId)
   val protocolType = "protocolType"
@@ -66,7 +71,7 @@ class GroupMetadataManagerTest {
   val sessionTimeout = 10000
 
   @Before
-  def setUp() {
+  def setUp(): Unit = {
     val config = KafkaConfig.fromProps(TestUtils.createBrokerConfig(nodeId = 0, zkConnect = ""))
 
     val offsetConfig = OffsetConfig(maxMetadataSize = config.offsetMetadataMaxSize,
@@ -87,14 +92,15 @@ class GroupMetadataManagerTest {
     EasyMock.expect(zkClient.getTopicPartitionCount(Topic.GROUP_METADATA_TOPIC_NAME)).andReturn(Some(2))
     EasyMock.replay(zkClient)
 
+    metrics = new kMetrics()
     time = new MockTime
     replicaManager = EasyMock.createNiceMock(classOf[ReplicaManager])
-    groupMetadataManager = new GroupMetadataManager(0, ApiVersion.latestVersion, offsetConfig, replicaManager, zkClient, time)
+    groupMetadataManager = new GroupMetadataManager(0, ApiVersion.latestVersion, offsetConfig, replicaManager, zkClient, time, metrics)
     partition = EasyMock.niceMock(classOf[Partition])
   }
 
   @Test
-  def testLoadOffsetsWithoutGroup() {
+  def testLoadOffsetsWithoutGroup(): Unit = {
     val groupMetadataTopicPartition = groupTopicPartition
     val startOffset = 15L
 
@@ -105,7 +111,7 @@ class GroupMetadataManagerTest {
     )
 
     val offsetCommitRecords = createCommittedOffsetRecords(committedOffsets)
-    val records = MemoryRecords.withRecords(startOffset, CompressionType.NONE, offsetCommitRecords: _*)
+    val records = MemoryRecords.withRecords(startOffset, CompressionType.NONE, offsetCommitRecords.toArray: _*)
     expectGroupMetadataLoad(groupMetadataTopicPartition, startOffset, records)
 
     EasyMock.replay(replicaManager)
@@ -122,7 +128,7 @@ class GroupMetadataManagerTest {
   }
 
   @Test
-  def testLoadEmptyGroupWithOffsets() {
+  def testLoadEmptyGroupWithOffsets(): Unit = {
     val groupMetadataTopicPartition = groupTopicPartition
     val generation = 15
     val protocolType = "consumer"
@@ -136,7 +142,7 @@ class GroupMetadataManagerTest {
     val offsetCommitRecords = createCommittedOffsetRecords(committedOffsets)
     val groupMetadataRecord = buildEmptyGroupRecord(generation, protocolType)
     val records = MemoryRecords.withRecords(startOffset, CompressionType.NONE,
-      offsetCommitRecords ++ Seq(groupMetadataRecord): _*)
+      (offsetCommitRecords ++ Seq(groupMetadataRecord)).toArray: _*)
 
     expectGroupMetadataLoad(groupMetadataTopicPartition, startOffset, records)
 
@@ -157,7 +163,7 @@ class GroupMetadataManagerTest {
   }
 
   @Test
-  def testLoadTransactionalOffsetsWithoutGroup() {
+  def testLoadTransactionalOffsetsWithoutGroup(): Unit = {
     val groupMetadataTopicPartition = groupTopicPartition
     val producerId = 1000L
     val producerEpoch: Short = 2
@@ -191,7 +197,7 @@ class GroupMetadataManagerTest {
   }
 
   @Test
-  def testDoNotLoadAbortedTransactionalOffsetCommits() {
+  def testDoNotLoadAbortedTransactionalOffsetCommits(): Unit = {
     val groupMetadataTopicPartition = groupTopicPartition
     val producerId = 1000L
     val producerEpoch: Short = 2
@@ -255,7 +261,7 @@ class GroupMetadataManagerTest {
   }
 
   @Test
-  def testLoadWithCommittedAndAbortedTransactionalOffsetCommits() {
+  def testLoadWithCommittedAndAbortedTransactionalOffsetCommits(): Unit = {
     // A test which loads a log with a mix of committed and aborted transactional offset committed messages.
     val groupMetadataTopicPartition = groupTopicPartition
     val producerId = 1000L
@@ -301,7 +307,7 @@ class GroupMetadataManagerTest {
   }
 
   @Test
-  def testLoadWithCommittedAndAbortedAndPendingTransactionalOffsetCommits() {
+  def testLoadWithCommittedAndAbortedAndPendingTransactionalOffsetCommits(): Unit = {
     val groupMetadataTopicPartition = groupTopicPartition
     val producerId = 1000L
     val producerEpoch: Short = 2
@@ -505,7 +511,7 @@ class GroupMetadataManagerTest {
   }
 
   @Test
-  def testGroupNotExists() {
+  def testGroupNotExists(): Unit = {
     // group is not owned
     assertFalse(groupMetadataManager.groupNotExists(groupId))
 
@@ -553,7 +559,7 @@ class GroupMetadataManagerTest {
   }
 
   @Test
-  def testLoadOffsetsWithTombstones() {
+  def testLoadOffsetsWithTombstones(): Unit = {
     val groupMetadataTopicPartition = groupTopicPartition
     val startOffset = 15L
 
@@ -567,7 +573,7 @@ class GroupMetadataManagerTest {
     val offsetCommitRecords = createCommittedOffsetRecords(committedOffsets)
     val tombstone = new SimpleRecord(GroupMetadataManager.offsetCommitKey(groupId, tombstonePartition), null)
     val records = MemoryRecords.withRecords(startOffset, CompressionType.NONE,
-      offsetCommitRecords ++ Seq(tombstone): _*)
+      (offsetCommitRecords ++ Seq(tombstone)).toArray: _*)
 
     expectGroupMetadataLoad(groupMetadataTopicPartition, startOffset, records)
 
@@ -588,7 +594,7 @@ class GroupMetadataManagerTest {
   }
 
   @Test
-  def testLoadOffsetsAndGroup() {
+  def testLoadOffsetsAndGroup(): Unit = {
     val groupMetadataTopicPartition = groupTopicPartition
     val generation = 935
     val protocolType = "consumer"
@@ -604,7 +610,7 @@ class GroupMetadataManagerTest {
     val memberId = "98098230493"
     val groupMetadataRecord = buildStableGroupRecordWithMember(generation, protocolType, protocol, memberId)
     val records = MemoryRecords.withRecords(startOffset, CompressionType.NONE,
-      offsetCommitRecords ++ Seq(groupMetadataRecord): _*)
+      (offsetCommitRecords ++ Seq(groupMetadataRecord)).toArray: _*)
 
     expectGroupMetadataLoad(groupMetadataTopicPartition, startOffset, records)
 
@@ -628,7 +634,7 @@ class GroupMetadataManagerTest {
   }
 
   @Test
-  def testLoadGroupWithTombstone() {
+  def testLoadGroupWithTombstone(): Unit = {
     val groupMetadataTopicPartition = groupTopicPartition
     val startOffset = 15L
     val memberId = "98098230493"
@@ -636,7 +642,7 @@ class GroupMetadataManagerTest {
       protocolType = "consumer", protocol = "range", memberId)
     val groupMetadataTombstone = new SimpleRecord(GroupMetadataManager.groupMetadataKey(groupId), null)
     val records = MemoryRecords.withRecords(startOffset, CompressionType.NONE,
-      Seq(groupMetadataRecord, groupMetadataTombstone): _*)
+      Seq(groupMetadataRecord, groupMetadataTombstone).toArray: _*)
 
     expectGroupMetadataLoad(groupMetadataTopicPartition, startOffset, records)
 
@@ -648,7 +654,7 @@ class GroupMetadataManagerTest {
   }
 
   @Test
-  def testLoadGroupWithLargeGroupMetadataRecord() {
+  def testLoadGroupWithLargeGroupMetadataRecord(): Unit = {
     val groupMetadataTopicPartition = groupTopicPartition
     val startOffset = 15L
     val committedOffsets = Map(
@@ -663,9 +669,9 @@ class GroupMetadataManagerTest {
 
     val offsetCommitRecords = createCommittedOffsetRecords(committedOffsets)
     val groupMetadataRecord = buildStableGroupRecordWithMember(generation = 15,
-      protocolType = "consumer", protocol = "range", memberId, assignmentSize)
+      protocolType = "consumer", protocol = "range", memberId, new Array[Byte](assignmentSize))
     val records = MemoryRecords.withRecords(startOffset, CompressionType.NONE,
-      offsetCommitRecords ++ Seq(groupMetadataRecord): _*)
+      (offsetCommitRecords ++ Seq(groupMetadataRecord)).toArray: _*)
 
     expectGroupMetadataLoad(groupMetadataTopicPartition, startOffset, records)
 
@@ -679,6 +685,30 @@ class GroupMetadataManagerTest {
     }
   }
 
+  @Test
+  def testLoadGroupAndOffsetsWithCorruptedLog(): Unit = {
+    // Simulate a case where startOffset < endOffset but log is empty. This could theoretically happen
+    // when all the records are expired and the active segment is truncated or when the partition
+    // is accidentally corrupted.
+    val startOffset = 0L
+    val endOffset = 10L
+
+    val logMock: Log = EasyMock.mock(classOf[Log])
+    EasyMock.expect(replicaManager.getLog(groupTopicPartition)).andStubReturn(Some(logMock))
+    expectGroupMetadataLoad(logMock, startOffset, MemoryRecords.EMPTY)
+    EasyMock.expect(replicaManager.getLogEndOffset(groupTopicPartition)).andStubReturn(Some(endOffset))
+    EasyMock.replay(logMock)
+
+    EasyMock.replay(replicaManager)
+
+    groupMetadataManager.loadGroupsAndOffsets(groupTopicPartition, _ => ())
+
+    EasyMock.verify(logMock)
+    EasyMock.verify(replicaManager)
+
+    assertFalse(groupMetadataManager.isPartitionLoading(groupTopicPartition.partition()))
+  }
+
   @Test
   def testOffsetWriteAfterGroupRemoved(): Unit = {
     // this test case checks the following scenario:
@@ -701,7 +731,7 @@ class GroupMetadataManagerTest {
     val groupMetadataRecord = buildStableGroupRecordWithMember(generation, protocolType, protocol, memberId)
     val groupMetadataTombstone = new SimpleRecord(GroupMetadataManager.groupMetadataKey(groupId), null)
     val records = MemoryRecords.withRecords(startOffset, CompressionType.NONE,
-      Seq(groupMetadataRecord, groupMetadataTombstone) ++ offsetCommitRecords: _*)
+      (Seq(groupMetadataRecord, groupMetadataTombstone) ++ offsetCommitRecords).toArray: _*)
 
     expectGroupMetadataLoad(groupMetadataTopicPartition, startOffset, records)
 
@@ -709,7 +739,7 @@ class GroupMetadataManagerTest {
 
     groupMetadataManager.loadGroupsAndOffsets(groupMetadataTopicPartition, _ => ())
 
-    val group = groupMetadataManager.getGroup(groupId).getOrElse(TestUtils.fail("Group was not loaded into the cache"))
+    val group = groupMetadataManager.getGroup(groupId).getOrElse(fail("Group was not loaded into the cache"))
     assertEquals(groupId, group.groupId)
     assertEquals(Empty, group.currentState)
     assertEquals(committedOffsets.size, group.allOffsets.size)
@@ -735,15 +765,15 @@ class GroupMetadataManagerTest {
     val segment1MemberId = "a"
     val segment1Offsets = Map(tp0 -> 23L, tp1 -> 455L, tp3 -> 42L)
     val segment1Records = MemoryRecords.withRecords(startOffset, CompressionType.NONE,
-      createCommittedOffsetRecords(segment1Offsets) ++ Seq(buildStableGroupRecordWithMember(
-        generation, protocolType, protocol, segment1MemberId)): _*)
+      (createCommittedOffsetRecords(segment1Offsets) ++ Seq(buildStableGroupRecordWithMember(
+        generation, protocolType, protocol, segment1MemberId))).toArray: _*)
     val segment1End = expectGroupMetadataLoad(logMock, startOffset, segment1Records)
 
     val segment2MemberId = "b"
     val segment2Offsets = Map(tp0 -> 33L, tp2 -> 8992L, tp3 -> 10L)
     val segment2Records = MemoryRecords.withRecords(segment1End, CompressionType.NONE,
-      createCommittedOffsetRecords(segment2Offsets) ++ Seq(buildStableGroupRecordWithMember(
-        generation, protocolType, protocol, segment2MemberId)): _*)
+      (createCommittedOffsetRecords(segment2Offsets) ++ Seq(buildStableGroupRecordWithMember(
+        generation, protocolType, protocol, segment2MemberId))).toArray: _*)
     val segment2End = expectGroupMetadataLoad(logMock, segment1End, segment2Records)
 
     EasyMock.expect(replicaManager.getLogEndOffset(groupTopicPartition)).andStubReturn(Some(segment2End))
@@ -768,14 +798,129 @@ class GroupMetadataManagerTest {
   }
 
   @Test
-  def testAddGroup() {
+  def testAddGroup(): Unit = {
     val group = new GroupMetadata("foo", Empty, time)
     assertEquals(group, groupMetadataManager.addGroup(group))
     assertEquals(group, groupMetadataManager.addGroup(new GroupMetadata("foo", Empty, time)))
   }
 
   @Test
-  def testStoreEmptyGroup() {
+  def testloadGroupWithStaticMember(): Unit = {
+    val generation = 27
+    val protocolType = "consumer"
+    val staticMemberId = "staticMemberId"
+    val dynamicMemberId = "dynamicMemberId"
+
+    val staticMember = new MemberMetadata(staticMemberId, groupId, groupInstanceId, "", "", rebalanceTimeout, sessionTimeout,
+      protocolType, List(("protocol", Array[Byte]())))
+
+    val dynamicMember = new MemberMetadata(dynamicMemberId, groupId, None, "", "", rebalanceTimeout, sessionTimeout,
+      protocolType, List(("protocol", Array[Byte]())))
+
+    val members = Seq(staticMember, dynamicMember)
+
+    val group = GroupMetadata.loadGroup(groupId, Empty, generation, protocolType, null, null, None, members, time)
+
+    assertTrue(group.is(Empty))
+    assertEquals(generation, group.generationId)
+    assertEquals(Some(protocolType), group.protocolType)
+    assertTrue(group.has(staticMemberId))
+    assertTrue(group.has(dynamicMemberId))
+    assertTrue(group.hasStaticMember(groupInstanceId))
+    assertEquals(staticMemberId, group.getStaticMemberId(groupInstanceId))
+  }
+
+  @Test
+  def testLoadConsumerGroup(): Unit = {
+    val generation = 27
+    val protocolType = "consumer"
+    val protocol = "protocol"
+    val memberId = "member1"
+    val topic = "foo"
+
+    val subscriptions = List(
+      ("protocol", ConsumerProtocol.serializeSubscription(new Subscription(List(topic).asJava)).array())
+    )
+
+    val member = new MemberMetadata(memberId, groupId, groupInstanceId, "", "", rebalanceTimeout,
+      sessionTimeout, protocolType, subscriptions)
+
+    val members = Seq(member)
+
+    val group = GroupMetadata.loadGroup(groupId, Stable, generation, protocolType, protocol, null, None,
+      members, time)
+
+    assertTrue(group.is(Stable))
+    assertEquals(generation, group.generationId)
+    assertEquals(Some(protocolType), group.protocolType)
+    assertEquals(protocol, group.protocolOrNull)
+    assertEquals(Some(Set(topic)), group.getSubscribedTopics)
+    assertTrue(group.has(memberId))
+  }
+
+  @Test
+  def testLoadEmptyConsumerGroup(): Unit = {
+    val generation = 27
+    val protocolType = "consumer"
+
+    val group = GroupMetadata.loadGroup(groupId, Empty, generation, protocolType, null, null, None,
+      Seq(), time)
+
+    assertTrue(group.is(Empty))
+    assertEquals(generation, group.generationId)
+    assertEquals(Some(protocolType), group.protocolType)
+    assertNull(group.protocolOrNull)
+    assertEquals(Some(Set.empty), group.getSubscribedTopics)
+  }
+
+  @Test
+  def testLoadConsumerGroupWithFaultyConsumerProtocol(): Unit = {
+    val generation = 27
+    val protocolType = "consumer"
+    val protocol = "protocol"
+    val memberId = "member1"
+
+    val subscriptions = List(("protocol", Array[Byte]()))
+
+    val member = new MemberMetadata(memberId, groupId, groupInstanceId, "", "", rebalanceTimeout,
+      sessionTimeout, protocolType, subscriptions)
+
+    val members = Seq(member)
+
+    val group = GroupMetadata.loadGroup(groupId, Stable, generation, protocolType, protocol, null, None,
+      members, time)
+
+    assertTrue(group.is(Stable))
+    assertEquals(generation, group.generationId)
+    assertEquals(Some(protocolType), group.protocolType)
+    assertEquals(protocol, group.protocolOrNull)
+    assertEquals(None, group.getSubscribedTopics)
+    assertTrue(group.has(memberId))
+  }
+
+  @Test
+  def testReadFromOldGroupMetadata(): Unit = {
+    val generation = 1
+    val protocol = "range"
+    val memberId = "memberId"
+    val oldApiVersions = Array(KAFKA_0_9_0, KAFKA_0_10_1_IV0, KAFKA_2_1_IV0)
+
+    for (apiVersion <- oldApiVersions) {
+      val groupMetadataRecord = buildStableGroupRecordWithMember(generation, protocolType, protocol, memberId, apiVersion = apiVersion)
+
+      val deserializedGroupMetadata = GroupMetadataManager.readGroupMessageValue(groupId, groupMetadataRecord.value(), time)
+      assertEquals(groupId, deserializedGroupMetadata.groupId)
+      assertEquals(generation, deserializedGroupMetadata.generationId)
+      assertEquals(protocolType, deserializedGroupMetadata.protocolType.get)
+      assertEquals(protocol, deserializedGroupMetadata.protocolOrNull)
+      assertEquals(1, deserializedGroupMetadata.allMembers.size)
+      assertTrue(deserializedGroupMetadata.allMembers.contains(memberId))
+      assertTrue(deserializedGroupMetadata.allStaticMembers.isEmpty)
+    }
+  }
+
+  @Test
+  def testStoreEmptyGroup(): Unit = {
     val generation = 27
     val protocolType = "consumer"
 
@@ -786,7 +931,7 @@ class GroupMetadataManagerTest {
     EasyMock.replay(replicaManager)
 
     var maybeError: Option[Errors] = None
-    def callback(error: Errors) {
+    def callback(error: Errors): Unit = {
       maybeError = Some(error)
     }
 
@@ -805,7 +950,7 @@ class GroupMetadataManagerTest {
   }
 
   @Test
-  def testStoreEmptySimpleGroup() {
+  def testStoreEmptySimpleGroup(): Unit = {
     val group = new GroupMetadata(groupId, Empty, time)
     groupMetadataManager.addGroup(group)
 
@@ -813,7 +958,7 @@ class GroupMetadataManagerTest {
     EasyMock.replay(replicaManager)
 
     var maybeError: Option[Errors] = None
-    def callback(error: Errors) {
+    def callback(error: Errors): Unit = {
       maybeError = Some(error)
     }
 
@@ -834,7 +979,7 @@ class GroupMetadataManagerTest {
   }
 
   @Test
-  def testStoreGroupErrorMapping() {
+  def testStoreGroupErrorMapping(): Unit = {
     assertStoreGroupErrorMapping(Errors.NONE, Errors.NONE)
     assertStoreGroupErrorMapping(Errors.UNKNOWN_TOPIC_OR_PARTITION, Errors.COORDINATOR_NOT_AVAILABLE)
     assertStoreGroupErrorMapping(Errors.NOT_ENOUGH_REPLICAS, Errors.COORDINATOR_NOT_AVAILABLE)
@@ -846,7 +991,7 @@ class GroupMetadataManagerTest {
     assertStoreGroupErrorMapping(Errors.CORRUPT_MESSAGE, Errors.CORRUPT_MESSAGE)
   }
 
-  private def assertStoreGroupErrorMapping(appendError: Errors, expectedError: Errors) {
+  private def assertStoreGroupErrorMapping(appendError: Errors, expectedError: Errors): Unit = {
     EasyMock.reset(replicaManager)
 
     val group = new GroupMetadata(groupId, Empty, time)
@@ -856,7 +1001,7 @@ class GroupMetadataManagerTest {
     EasyMock.replay(replicaManager)
 
     var maybeError: Option[Errors] = None
-    def callback(error: Errors) {
+    def callback(error: Errors): Unit = {
       maybeError = Some(error)
     }
 
@@ -867,7 +1012,7 @@ class GroupMetadataManagerTest {
   }
 
   @Test
-  def testStoreNonEmptyGroup() {
+  def testStoreNonEmptyGroup(): Unit = {
     val memberId = "memberId"
     val clientId = "clientId"
     val clientHost = "localhost"
@@ -875,7 +1020,7 @@ class GroupMetadataManagerTest {
     val group = new GroupMetadata(groupId, Empty, time)
     groupMetadataManager.addGroup(group)
 
-    val member = new MemberMetadata(memberId, groupId, clientId, clientHost, rebalanceTimeout, sessionTimeout,
+    val member = new MemberMetadata(memberId, groupId, groupInstanceId, clientId, clientHost, rebalanceTimeout, sessionTimeout,
       protocolType, List(("protocol", Array[Byte]())))
     group.add(member, _ => ())
     group.transitionTo(PreparingRebalance)
@@ -885,7 +1030,7 @@ class GroupMetadataManagerTest {
     EasyMock.replay(replicaManager)
 
     var maybeError: Option[Errors] = None
-    def callback(error: Errors) {
+    def callback(error: Errors): Unit = {
       maybeError = Some(error)
     }
 
@@ -896,7 +1041,7 @@ class GroupMetadataManagerTest {
   }
 
   @Test
-  def testStoreNonEmptyGroupWhenCoordinatorHasMoved() {
+  def testStoreNonEmptyGroupWhenCoordinatorHasMoved(): Unit = {
     EasyMock.expect(replicaManager.getMagic(EasyMock.anyObject())).andReturn(None)
     val memberId = "memberId"
     val clientId = "clientId"
@@ -904,7 +1049,7 @@ class GroupMetadataManagerTest {
 
     val group = new GroupMetadata(groupId, Empty, time)
 
-    val member = new MemberMetadata(memberId, groupId, clientId, clientHost, rebalanceTimeout, sessionTimeout,
+    val member = new MemberMetadata(memberId, groupId, groupInstanceId, clientId, clientHost, rebalanceTimeout, sessionTimeout,
       protocolType, List(("protocol", Array[Byte]())))
     group.add(member, _ => ())
     group.transitionTo(PreparingRebalance)
@@ -913,7 +1058,7 @@ class GroupMetadataManagerTest {
     EasyMock.replay(replicaManager)
 
     var maybeError: Option[Errors] = None
-    def callback(error: Errors) {
+    def callback(error: Errors): Unit = {
       maybeError = Some(error)
     }
 
@@ -924,7 +1069,7 @@ class GroupMetadataManagerTest {
   }
 
   @Test
-  def testCommitOffset() {
+  def testCommitOffset(): Unit = {
     val memberId = ""
     val topicPartition = new TopicPartition("foo", 0)
     val offset = 37
@@ -940,7 +1085,7 @@ class GroupMetadataManagerTest {
     EasyMock.replay(replicaManager)
 
     var commitErrors: Option[immutable.Map[TopicPartition, Errors]] = None
-    def callback(errors: immutable.Map[TopicPartition, Errors]) {
+    def callback(errors: immutable.Map[TopicPartition, Errors]): Unit = {
       commitErrors = Some(errors)
     }
 
@@ -964,7 +1109,7 @@ class GroupMetadataManagerTest {
   }
 
   @Test
-  def testTransactionalCommitOffsetCommitted() {
+  def testTransactionalCommitOffsetCommitted(): Unit = {
     val memberId = ""
     val topicPartition = new TopicPartition("foo", 0)
     val offset = 37
@@ -983,7 +1128,7 @@ class GroupMetadataManagerTest {
     EasyMock.replay(replicaManager)
 
     var commitErrors: Option[immutable.Map[TopicPartition, Errors]] = None
-    def callback(errors: immutable.Map[TopicPartition, Errors]) {
+    def callback(errors: immutable.Map[TopicPartition, Errors]): Unit = {
       commitErrors = Some(errors)
     }
 
@@ -1005,7 +1150,7 @@ class GroupMetadataManagerTest {
   }
 
   @Test
-  def testTransactionalCommitOffsetAppendFailure() {
+  def testTransactionalCommitOffsetAppendFailure(): Unit = {
     val memberId = ""
     val topicPartition = new TopicPartition("foo", 0)
     val offset = 37
@@ -1023,7 +1168,7 @@ class GroupMetadataManagerTest {
     EasyMock.replay(replicaManager)
 
     var commitErrors: Option[immutable.Map[TopicPartition, Errors]] = None
-    def callback(errors: immutable.Map[TopicPartition, Errors]) {
+    def callback(errors: immutable.Map[TopicPartition, Errors]): Unit = {
       commitErrors = Some(errors)
     }
 
@@ -1044,7 +1189,7 @@ class GroupMetadataManagerTest {
   }
 
   @Test
-  def testTransactionalCommitOffsetAborted() {
+  def testTransactionalCommitOffsetAborted(): Unit = {
     val memberId = ""
     val topicPartition = new TopicPartition("foo", 0)
     val offset = 37
@@ -1062,7 +1207,7 @@ class GroupMetadataManagerTest {
     EasyMock.replay(replicaManager)
 
     var commitErrors: Option[immutable.Map[TopicPartition, Errors]] = None
-    def callback(errors: immutable.Map[TopicPartition, Errors]) {
+    def callback(errors: immutable.Map[TopicPartition, Errors]): Unit = {
       commitErrors = Some(errors)
     }
 
@@ -1083,7 +1228,7 @@ class GroupMetadataManagerTest {
   }
 
   @Test
-  def testCommitOffsetWhenCoordinatorHasMoved() {
+  def testCommitOffsetWhenCoordinatorHasMoved(): Unit = {
     EasyMock.expect(replicaManager.getMagic(EasyMock.anyObject())).andReturn(None)
     val memberId = ""
     val topicPartition = new TopicPartition("foo", 0)
@@ -1099,7 +1244,7 @@ class GroupMetadataManagerTest {
     EasyMock.replay(replicaManager)
 
     var commitErrors: Option[immutable.Map[TopicPartition, Errors]] = None
-    def callback(errors: immutable.Map[TopicPartition, Errors]) {
+    def callback(errors: immutable.Map[TopicPartition, Errors]): Unit = {
       commitErrors = Some(errors)
     }
 
@@ -1113,7 +1258,7 @@ class GroupMetadataManagerTest {
   }
 
   @Test
-  def testCommitOffsetFailure() {
+  def testCommitOffsetFailure(): Unit = {
     assertCommitOffsetErrorMapping(Errors.UNKNOWN_TOPIC_OR_PARTITION, Errors.COORDINATOR_NOT_AVAILABLE)
     assertCommitOffsetErrorMapping(Errors.NOT_ENOUGH_REPLICAS, Errors.COORDINATOR_NOT_AVAILABLE)
     assertCommitOffsetErrorMapping(Errors.NOT_ENOUGH_REPLICAS_AFTER_APPEND, Errors.COORDINATOR_NOT_AVAILABLE)
@@ -1142,7 +1287,7 @@ class GroupMetadataManagerTest {
     EasyMock.replay(replicaManager)
 
     var commitErrors: Option[immutable.Map[TopicPartition, Errors]] = None
-    def callback(errors: immutable.Map[TopicPartition, Errors]) {
+    def callback(errors: immutable.Map[TopicPartition, Errors]): Unit = {
       commitErrors = Some(errors)
     }
 
@@ -1163,7 +1308,7 @@ class GroupMetadataManagerTest {
   }
 
   @Test
-  def testExpireOffset() {
+  def testExpireOffset(): Unit = {
     val memberId = ""
     val topicPartition1 = new TopicPartition("foo", 0)
     val topicPartition2 = new TopicPartition("foo", 1)
@@ -1185,7 +1330,7 @@ class GroupMetadataManagerTest {
     EasyMock.replay(replicaManager)
 
     var commitErrors: Option[immutable.Map[TopicPartition, Errors]] = None
-    def callback(errors: immutable.Map[TopicPartition, Errors]) {
+    def callback(errors: immutable.Map[TopicPartition, Errors]): Unit = {
       commitErrors = Some(errors)
     }
 
@@ -1200,7 +1345,7 @@ class GroupMetadataManagerTest {
 
     EasyMock.reset(partition)
     EasyMock.expect(partition.appendRecordsToLeader(EasyMock.anyObject(classOf[MemoryRecords]),
-      isFromClient = EasyMock.eq(false), requiredAcks = EasyMock.anyInt()))
+      origin = EasyMock.eq(AppendOrigin.Coordinator), requiredAcks = EasyMock.anyInt()))
       .andReturn(LogAppendInfo.UnknownLogAppendInfo)
     EasyMock.replay(partition)
 
@@ -1218,7 +1363,7 @@ class GroupMetadataManagerTest {
   }
 
   @Test
-  def testGroupMetadataRemoval() {
+  def testGroupMetadataRemoval(): Unit = {
     val topicPartition1 = new TopicPartition("foo", 0)
     val topicPartition2 = new TopicPartition("foo", 1)
 
@@ -1235,7 +1380,7 @@ class GroupMetadataManagerTest {
     EasyMock.expect(replicaManager.getMagic(EasyMock.anyObject())).andStubReturn(Some(RecordBatch.CURRENT_MAGIC_VALUE))
     mockGetPartition()
     EasyMock.expect(partition.appendRecordsToLeader(EasyMock.capture(recordsCapture),
-      isFromClient = EasyMock.eq(false), requiredAcks = EasyMock.anyInt()))
+      origin = EasyMock.eq(AppendOrigin.Coordinator), requiredAcks = EasyMock.anyInt()))
       .andReturn(LogAppendInfo.UnknownLogAppendInfo)
     EasyMock.replay(replicaManager, partition)
 
@@ -1266,7 +1411,7 @@ class GroupMetadataManagerTest {
   }
 
   @Test
-  def testGroupMetadataRemovalWithLogAppendTime() {
+  def testGroupMetadataRemovalWithLogAppendTime(): Unit = {
     val topicPartition1 = new TopicPartition("foo", 0)
     val topicPartition2 = new TopicPartition("foo", 1)
 
@@ -1283,7 +1428,7 @@ class GroupMetadataManagerTest {
     EasyMock.expect(replicaManager.getMagic(EasyMock.anyObject())).andStubReturn(Some(RecordBatch.CURRENT_MAGIC_VALUE))
     mockGetPartition()
     EasyMock.expect(partition.appendRecordsToLeader(EasyMock.capture(recordsCapture),
-      isFromClient = EasyMock.eq(false), requiredAcks = EasyMock.anyInt()))
+      origin = EasyMock.eq(AppendOrigin.Coordinator), requiredAcks = EasyMock.anyInt()))
       .andReturn(LogAppendInfo.UnknownLogAppendInfo)
     EasyMock.replay(replicaManager, partition)
 
@@ -1315,7 +1460,7 @@ class GroupMetadataManagerTest {
   }
 
   @Test
-  def testExpireGroupWithOffsetsOnly() {
+  def testExpireGroupWithOffsetsOnly(): Unit = {
     // verify that the group is removed properly, but no tombstone is written if
     // this is a group which is only using kafka for offset storage
 
@@ -1340,7 +1485,7 @@ class GroupMetadataManagerTest {
     EasyMock.replay(replicaManager)
 
     var commitErrors: Option[immutable.Map[TopicPartition, Errors]] = None
-    def callback(errors: immutable.Map[TopicPartition, Errors]) {
+    def callback(errors: immutable.Map[TopicPartition, Errors]): Unit = {
       commitErrors = Some(errors)
     }
 
@@ -1358,7 +1503,7 @@ class GroupMetadataManagerTest {
     val recordsCapture: Capture[MemoryRecords] = EasyMock.newCapture()
 
     EasyMock.expect(partition.appendRecordsToLeader(EasyMock.capture(recordsCapture),
-      isFromClient = EasyMock.eq(false), requiredAcks = EasyMock.anyInt()))
+      origin = EasyMock.eq(AppendOrigin.Coordinator), requiredAcks = EasyMock.anyInt()))
       .andReturn(LogAppendInfo.UnknownLogAppendInfo)
     EasyMock.replay(partition)
 
@@ -1387,7 +1532,7 @@ class GroupMetadataManagerTest {
   }
 
   @Test
-  def testOffsetExpirationSemantics() {
+  def testOffsetExpirationSemantics(): Unit = {
     val memberId = "memberId"
     val clientId = "clientId"
     val clientHost = "localhost"
@@ -1403,7 +1548,7 @@ class GroupMetadataManagerTest {
     groupMetadataManager.addGroup(group)
 
     val subscription = new Subscription(List(topic).asJava)
-    val member = new MemberMetadata(memberId, groupId, clientId, clientHost, rebalanceTimeout, sessionTimeout,
+    val member = new MemberMetadata(memberId, groupId, groupInstanceId, clientId, clientHost, rebalanceTimeout, sessionTimeout,
       protocolType, List(("protocol", ConsumerProtocol.serializeSubscription(subscription).array())))
     group.add(member, _ => ())
     group.transitionTo(PreparingRebalance)
@@ -1425,7 +1570,7 @@ class GroupMetadataManagerTest {
     EasyMock.replay(replicaManager)
 
     var commitErrors: Option[immutable.Map[TopicPartition, Errors]] = None
-    def callback(errors: immutable.Map[TopicPartition, Errors]) {
+    def callback(errors: immutable.Map[TopicPartition, Errors]): Unit = {
       commitErrors = Some(errors)
     }
 
@@ -1459,7 +1604,7 @@ class GroupMetadataManagerTest {
     // expect the offset tombstone
     EasyMock.reset(partition)
     EasyMock.expect(partition.appendRecordsToLeader(EasyMock.anyObject(classOf[MemoryRecords]),
-      isFromClient = EasyMock.eq(false), requiredAcks = EasyMock.anyInt()))
+      origin = EasyMock.eq(AppendOrigin.Coordinator), requiredAcks = EasyMock.anyInt()))
       .andReturn(LogAppendInfo.UnknownLogAppendInfo)
     EasyMock.replay(partition)
 
@@ -1483,7 +1628,7 @@ class GroupMetadataManagerTest {
     // expect the offset tombstone
     EasyMock.reset(partition)
     EasyMock.expect(partition.appendRecordsToLeader(EasyMock.anyObject(classOf[MemoryRecords]),
-      isFromClient = EasyMock.eq(false), requiredAcks = EasyMock.anyInt()))
+      origin = EasyMock.eq(AppendOrigin.Coordinator), requiredAcks = EasyMock.anyInt()))
       .andReturn(LogAppendInfo.UnknownLogAppendInfo)
     EasyMock.replay(partition)
 
@@ -1526,7 +1671,7 @@ class GroupMetadataManagerTest {
     // expect the offset tombstone
     EasyMock.reset(partition)
     EasyMock.expect(partition.appendRecordsToLeader(EasyMock.anyObject(classOf[MemoryRecords]),
-      isFromClient = EasyMock.eq(false), requiredAcks = EasyMock.anyInt()))
+      origin = EasyMock.eq(AppendOrigin.Coordinator), requiredAcks = EasyMock.anyInt()))
       .andReturn(LogAppendInfo.UnknownLogAppendInfo)
     EasyMock.replay(partition)
 
@@ -1549,7 +1694,7 @@ class GroupMetadataManagerTest {
   }
 
   @Test
-  def testOffsetExpirationOfSimpleConsumer() {
+  def testOffsetExpirationOfSimpleConsumer(): Unit = {
     val memberId = "memberId"
     val topic = "foo"
     val topicPartition1 = new TopicPartition(topic, 0)
@@ -1573,7 +1718,7 @@ class GroupMetadataManagerTest {
     EasyMock.replay(replicaManager)
 
     var commitErrors: Option[immutable.Map[TopicPartition, Errors]] = None
-    def callback(errors: immutable.Map[TopicPartition, Errors]) {
+    def callback(errors: immutable.Map[TopicPartition, Errors]): Unit = {
       commitErrors = Some(errors)
     }
 
@@ -1604,7 +1749,7 @@ class GroupMetadataManagerTest {
     // expect the offset tombstone
     EasyMock.reset(partition)
     EasyMock.expect(partition.appendRecordsToLeader(EasyMock.anyObject(classOf[MemoryRecords]),
-      isFromClient = EasyMock.eq(false), requiredAcks = EasyMock.anyInt()))
+      origin = EasyMock.eq(AppendOrigin.Coordinator), requiredAcks = EasyMock.anyInt()))
       .andReturn(LogAppendInfo.UnknownLogAppendInfo)
     EasyMock.replay(partition)
 
@@ -1622,6 +1767,144 @@ class GroupMetadataManagerTest {
     assert(group.is(Dead))
   }
 
+  @Test
+  def testOffsetExpirationOfActiveGroupSemantics(): Unit = {
+    val memberId = "memberId"
+    val clientId = "clientId"
+    val clientHost = "localhost"
+
+    val topic1 = "foo"
+    val topic1Partition0 = new TopicPartition(topic1, 0)
+    val topic1Partition1 = new TopicPartition(topic1, 1)
+
+    val topic2 = "bar"
+    val topic2Partition0 = new TopicPartition(topic2, 0)
+    val topic2Partition1 = new TopicPartition(topic2, 1)
+
+    val offset = 37
+
+    groupMetadataManager.addPartitionOwnership(groupPartitionId)
+
+    val group = new GroupMetadata(groupId, Empty, time)
+    groupMetadataManager.addGroup(group)
+
+    // Subscribe to topic1 and topic2
+    val subscriptionTopic1AndTopic2 = new Subscription(List(topic1, topic2).asJava)
+
+    val member = new MemberMetadata(
+      memberId,
+      groupId,
+      groupInstanceId,
+      clientId,
+      clientHost,
+      rebalanceTimeout,
+      sessionTimeout,
+      ConsumerProtocol.PROTOCOL_TYPE,
+      List(("protocol", ConsumerProtocol.serializeSubscription(subscriptionTopic1AndTopic2).array()))
+    )
+
+    group.add(member, _ => ())
+    group.transitionTo(PreparingRebalance)
+    group.initNextGeneration()
+    group.transitionTo(Stable)
+
+    val startMs = time.milliseconds
+
+    val t1p0OffsetAndMetadata = OffsetAndMetadata(offset, "", startMs)
+    val t1p1OffsetAndMetadata = OffsetAndMetadata(offset, "", startMs)
+
+    val t2p0OffsetAndMetadata = OffsetAndMetadata(offset, "", startMs)
+    val t2p1OffsetAndMetadata = OffsetAndMetadata(offset, "", startMs)
+
+    val offsets = immutable.Map(
+      topic1Partition0 -> t1p0OffsetAndMetadata,
+      topic1Partition1 -> t1p1OffsetAndMetadata,
+      topic2Partition0 -> t2p0OffsetAndMetadata,
+      topic2Partition1 -> t2p1OffsetAndMetadata)
+
+    mockGetPartition()
+    expectAppendMessage(Errors.NONE)
+    EasyMock.replay(replicaManager)
+
+    var commitErrors: Option[immutable.Map[TopicPartition, Errors]] = None
+    def callback(errors: immutable.Map[TopicPartition, Errors]): Unit = {
+      commitErrors = Some(errors)
+    }
+
+    groupMetadataManager.storeOffsets(group, memberId, offsets, callback)
+    assertTrue(group.hasOffsets)
+
+    assertFalse(commitErrors.isEmpty)
+    assertEquals(Some(Errors.NONE), commitErrors.get.get(topic1Partition0))
+
+    // advance time to just after the offset of last partition is to be expired
+    time.sleep(defaultOffsetRetentionMs + 2)
+
+    // no offset should expire because all topics are actively consumed
+    groupMetadataManager.cleanupGroupMetadata()
+
+    assertEquals(Some(group), groupMetadataManager.getGroup(groupId))
+    assert(group.is(Stable))
+
+    assertEquals(Some(t1p0OffsetAndMetadata), group.offset(topic1Partition0))
+    assertEquals(Some(t1p1OffsetAndMetadata), group.offset(topic1Partition1))
+    assertEquals(Some(t2p0OffsetAndMetadata), group.offset(topic2Partition0))
+    assertEquals(Some(t2p1OffsetAndMetadata), group.offset(topic2Partition1))
+
+    var cachedOffsets = groupMetadataManager.getOffsets(groupId,
+      Some(Seq(topic1Partition0, topic1Partition1, topic2Partition0, topic2Partition1)))
+
+    assertEquals(Some(offset), cachedOffsets.get(topic1Partition0).map(_.offset))
+    assertEquals(Some(offset), cachedOffsets.get(topic1Partition1).map(_.offset))
+    assertEquals(Some(offset), cachedOffsets.get(topic2Partition0).map(_.offset))
+    assertEquals(Some(offset), cachedOffsets.get(topic2Partition1).map(_.offset))
+
+    EasyMock.verify(replicaManager)
+
+    group.transitionTo(PreparingRebalance)
+
+    // Subscribe to topic1, offsets of topic2 should be removed
+    val subscriptionTopic1 = new Subscription(List(topic1).asJava)
+
+    group.updateMember(
+      member,
+      List(("protocol", ConsumerProtocol.serializeSubscription(subscriptionTopic1).array())),
+      null
+    )
+
+    group.initNextGeneration()
+    group.transitionTo(Stable)
+
+    // expect the offset tombstone
+    EasyMock.expect(partition.appendRecordsToLeader(EasyMock.anyObject(classOf[MemoryRecords]),
+      origin = EasyMock.eq(AppendOrigin.Coordinator), requiredAcks = EasyMock.anyInt()))
+      .andReturn(LogAppendInfo.UnknownLogAppendInfo)
+    EasyMock.expectLastCall().times(1)
+
+    EasyMock.replay(partition)
+
+    groupMetadataManager.cleanupGroupMetadata()
+
+    EasyMock.verify(partition)
+    EasyMock.verify(replicaManager)
+
+    assertEquals(Some(group), groupMetadataManager.getGroup(groupId))
+    assert(group.is(Stable))
+
+    assertEquals(Some(t1p0OffsetAndMetadata), group.offset(topic1Partition0))
+    assertEquals(Some(t1p1OffsetAndMetadata), group.offset(topic1Partition1))
+    assertEquals(None, group.offset(topic2Partition0))
+    assertEquals(None, group.offset(topic2Partition1))
+
+    cachedOffsets = groupMetadataManager.getOffsets(groupId,
+      Some(Seq(topic1Partition0, topic1Partition1, topic2Partition0, topic2Partition1)))
+
+    assertEquals(Some(offset), cachedOffsets.get(topic1Partition0).map(_.offset))
+    assertEquals(Some(offset), cachedOffsets.get(topic1Partition1).map(_.offset))
+    assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), cachedOffsets.get(topic2Partition0).map(_.offset))
+    assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), cachedOffsets.get(topic2Partition1).map(_.offset))
+  }
+
   @Test
   def testLoadOffsetFromOldCommit() = {
     val groupMetadataTopicPartition = groupTopicPartition
@@ -1640,7 +1923,7 @@ class GroupMetadataManagerTest {
     val memberId = "98098230493"
     val groupMetadataRecord = buildStableGroupRecordWithMember(generation, protocolType, protocol, memberId, apiVersion = apiVersion)
     val records = MemoryRecords.withRecords(startOffset, CompressionType.NONE,
-      offsetCommitRecords ++ Seq(groupMetadataRecord): _*)
+      (offsetCommitRecords ++ Seq(groupMetadataRecord)).toArray: _*)
 
     expectGroupMetadataLoad(groupMetadataTopicPartition, startOffset, records)
 
@@ -1680,7 +1963,7 @@ class GroupMetadataManagerTest {
     val memberId = "98098230493"
     val groupMetadataRecord = buildStableGroupRecordWithMember(generation, protocolType, protocol, memberId)
     val records = MemoryRecords.withRecords(startOffset, CompressionType.NONE,
-      offsetCommitRecords ++ Seq(groupMetadataRecord): _*)
+      (offsetCommitRecords ++ Seq(groupMetadataRecord)).toArray: _*)
 
     expectGroupMetadataLoad(groupMetadataTopicPartition, startOffset, records)
 
@@ -1796,7 +2079,7 @@ class GroupMetadataManagerTest {
   }
 
   @Test
-  def testLoadOffsetsWithEmptyControlBatch() {
+  def testLoadOffsetsWithEmptyControlBatch(): Unit = {
     val groupMetadataTopicPartition = groupTopicPartition
     val startOffset = 15L
     val generation = 15
@@ -1810,7 +2093,7 @@ class GroupMetadataManagerTest {
     val offsetCommitRecords = createCommittedOffsetRecords(committedOffsets)
     val groupMetadataRecord = buildEmptyGroupRecord(generation, protocolType)
     val records = MemoryRecords.withRecords(startOffset, CompressionType.NONE,
-      offsetCommitRecords ++ Seq(groupMetadataRecord): _*)
+      (offsetCommitRecords ++ Seq(groupMetadataRecord)).toArray: _*)
 
     // Prepend empty control batch to valid records
     val mockBatch: MutableRecordBatch = EasyMock.createMock(classOf[MutableRecordBatch])
@@ -1829,9 +2112,8 @@ class GroupMetadataManagerTest {
     EasyMock.expect(logMock.logStartOffset).andReturn(startOffset).anyTimes()
     EasyMock.expect(logMock.read(EasyMock.eq(startOffset),
       maxLength = EasyMock.anyInt(),
-      maxOffset = EasyMock.eq(None),
-      minOneMessage = EasyMock.eq(true),
-      includeAbortedTxns = EasyMock.eq(false)))
+      isolation = EasyMock.eq(FetchLogEnd),
+      minOneMessage = EasyMock.eq(true)))
       .andReturn(FetchDataInfo(LogOffsetMetadata(startOffset), mockRecords))
     EasyMock.expect(replicaManager.getLog(groupMetadataTopicPartition)).andStubReturn(Some(logMock))
     EasyMock.expect(replicaManager.getLogEndOffset(groupMetadataTopicPartition)).andStubReturn(Some(18))
@@ -1853,12 +2135,68 @@ class GroupMetadataManagerTest {
     }
   }
 
+  @Test
+  def testCommittedOffsetParsing(): Unit = {
+    val groupId = "group"
+    val topicPartition = new TopicPartition("topic", 0)
+    val offsetCommitRecord = TestUtils.records(Seq(
+      new SimpleRecord(
+        GroupMetadataManager.offsetCommitKey(groupId, topicPartition),
+        GroupMetadataManager.offsetCommitValue(OffsetAndMetadata(35L, "", time.milliseconds()), ApiVersion.latestVersion)
+      )
+    )).records.asScala.head
+    val (keyStringOpt, valueStringOpt) = GroupMetadataManager.formatRecordKeyAndValue(offsetCommitRecord)
+    assertEquals(Some(s"offset_commit::group=$groupId,partition=$topicPartition"), keyStringOpt)
+    assertEquals(Some("offset=35"), valueStringOpt)
+  }
+
+  @Test
+  def testCommittedOffsetTombstoneParsing(): Unit = {
+    val groupId = "group"
+    val topicPartition = new TopicPartition("topic", 0)
+    val offsetCommitRecord = TestUtils.records(Seq(
+      new SimpleRecord(GroupMetadataManager.offsetCommitKey(groupId, topicPartition), null)
+    )).records.asScala.head
+    val (keyStringOpt, valueStringOpt) = GroupMetadataManager.formatRecordKeyAndValue(offsetCommitRecord)
+    assertEquals(Some(s"offset_commit::group=$groupId,partition=$topicPartition"), keyStringOpt)
+    assertEquals(Some(""), valueStringOpt)
+  }
+
+  @Test
+  def testGroupMetadataParsingWithNullUserData(): Unit = {
+    val generation = 935
+    val protocolType = "consumer"
+    val protocol = "range"
+    val memberId = "98098230493"
+    val assignmentBytes = Utils.toArray(ConsumerProtocol.serializeAssignment(
+      new ConsumerPartitionAssignor.Assignment(List(new TopicPartition("topic", 0)).asJava, null)
+    ))
+    val groupMetadataRecord = TestUtils.records(Seq(
+      buildStableGroupRecordWithMember(generation, protocolType, protocol, memberId, assignmentBytes)
+    )).records.asScala.head
+    val (keyStringOpt, valueStringOpt) = GroupMetadataManager.formatRecordKeyAndValue(groupMetadataRecord)
+    assertEquals(Some(s"group_metadata::group=$groupId"), keyStringOpt)
+    assertEquals(Some("{\"protocolType\":\"consumer\",\"protocol\":\"range\"," +
+      "\"generationId\":935,\"assignment\":\"{98098230493=[topic-0]}\"}"), valueStringOpt)
+  }
+
+  @Test
+  def testGroupMetadataTombstoneParsing(): Unit = {
+    val groupId = "group"
+    val groupMetadataRecord = TestUtils.records(Seq(
+      new SimpleRecord(GroupMetadataManager.groupMetadataKey(groupId), null)
+    )).records.asScala.head
+    val (keyStringOpt, valueStringOpt) = GroupMetadataManager.formatRecordKeyAndValue(groupMetadataRecord)
+    assertEquals(Some(s"group_metadata::group=$groupId"), keyStringOpt)
+    assertEquals(Some(""), valueStringOpt)
+  }
+
   private def appendAndCaptureCallback(): Capture[Map[TopicPartition, PartitionResponse] => Unit] = {
     val capturedArgument: Capture[Map[TopicPartition, PartitionResponse] => Unit] = EasyMock.newCapture()
     EasyMock.expect(replicaManager.appendRecords(EasyMock.anyLong(),
       EasyMock.anyShort(),
       internalTopicsAllowed = EasyMock.eq(true),
-      isFromClient = EasyMock.eq(false),
+      origin = EasyMock.eq(AppendOrigin.Coordinator),
       EasyMock.anyObject().asInstanceOf[Map[TopicPartition, MemoryRecords]],
       EasyMock.capture(capturedArgument),
       EasyMock.anyObject().asInstanceOf[Option[ReentrantLock]],
@@ -1874,7 +2212,7 @@ class GroupMetadataManagerTest {
     EasyMock.expect(replicaManager.appendRecords(EasyMock.anyLong(),
       EasyMock.anyShort(),
       internalTopicsAllowed = EasyMock.eq(true),
-      isFromClient = EasyMock.eq(false),
+      origin = EasyMock.eq(AppendOrigin.Coordinator),
       EasyMock.capture(capturedRecords),
       EasyMock.capture(capturedCallback),
       EasyMock.anyObject().asInstanceOf[Option[ReentrantLock]],
@@ -1893,14 +2231,14 @@ class GroupMetadataManagerTest {
                                                protocolType: String,
                                                protocol: String,
                                                memberId: String,
-                                               assignmentSize: Int = 0,
+                                               assignmentBytes: Array[Byte] = Array.emptyByteArray,
                                                apiVersion: ApiVersion = ApiVersion.latestVersion): SimpleRecord = {
     val memberProtocols = List((protocol, Array.emptyByteArray))
-    val member = new MemberMetadata(memberId, groupId, "clientId", "clientHost", 30000, 10000, protocolType, memberProtocols)
+    val member = new MemberMetadata(memberId, groupId, groupInstanceId, "clientId", "clientHost", 30000, 10000, protocolType, memberProtocols)
     val group = GroupMetadata.loadGroup(groupId, Stable, generation, protocolType, protocol, memberId,
       if (apiVersion >= KAFKA_2_1_IV0) Some(time.milliseconds()) else None, Seq(member), time)
     val groupMetadataKey = GroupMetadataManager.groupMetadataKey(groupId)
-    val groupMetadataValue = GroupMetadataManager.groupMetadataValue(group, Map(memberId -> new Array[Byte](assignmentSize)), apiVersion)
+    val groupMetadataValue = GroupMetadataManager.groupMetadataValue(group, Map(memberId -> assignmentBytes), apiVersion)
     new SimpleRecord(groupMetadataKey, groupMetadataValue)
   }
 
@@ -1935,9 +2273,8 @@ class GroupMetadataManagerTest {
     EasyMock.expect(logMock.logStartOffset).andStubReturn(startOffset)
     EasyMock.expect(logMock.read(EasyMock.eq(startOffset),
       maxLength = EasyMock.anyInt(),
-      maxOffset = EasyMock.eq(None),
-      minOneMessage = EasyMock.eq(true),
-      includeAbortedTxns = EasyMock.eq(false)))
+      isolation = EasyMock.eq(FetchLogEnd),
+      minOneMessage = EasyMock.eq(true)))
       .andReturn(FetchDataInfo(LogOffsetMetadata(startOffset), fileRecordsMock))
 
     EasyMock.expect(fileRecordsMock.sizeInBytes()).andStubReturn(records.sizeInBytes)
@@ -1977,7 +2314,7 @@ class GroupMetadataManagerTest {
   }
 
   private def mockGetPartition(): Unit = {
-    EasyMock.expect(replicaManager.getPartition(groupTopicPartition)).andStubReturn(Some(partition))
+    EasyMock.expect(replicaManager.getPartition(groupTopicPartition)).andStubReturn(HostedPartition.Online(partition))
     EasyMock.expect(replicaManager.nonOfflinePartition(groupTopicPartition)).andStubReturn(Some(partition))
   }
 
@@ -1995,7 +2332,7 @@ class GroupMetadataManagerTest {
   }
 
   @Test
-  def testMetrics() {
+  def testMetrics(): Unit = {
     groupMetadataManager.cleanupGroupMetadata()
     expectMetrics(groupMetadataManager, 0, 0, 0)
     val group = new GroupMetadata("foo2", Stable, time)
@@ -2006,4 +2343,43 @@ class GroupMetadataManagerTest {
     group.transitionTo(CompletingRebalance)
     expectMetrics(groupMetadataManager, 1, 0, 1)
   }
+
+  @Test
+  def testPartitionLoadMetric(): Unit = {
+    val server = ManagementFactory.getPlatformMBeanServer
+    val mBeanName = "kafka.server:type=group-coordinator-metrics"
+    val reporter = new JmxReporter("kafka.server")
+    metrics.addReporter(reporter)
+
+    def partitionLoadTime(attribute: String): Double = {
+      server.getAttribute(new ObjectName(mBeanName), attribute).asInstanceOf[Double]
+    }
+
+    assertTrue(server.isRegistered(new ObjectName(mBeanName)))
+    assertEquals(Double.NaN, partitionLoadTime( "partition-load-time-max"), 0)
+    assertEquals(Double.NaN, partitionLoadTime("partition-load-time-avg"), 0)
+    assertTrue(reporter.containsMbean(mBeanName))
+
+    val groupMetadataTopicPartition = groupTopicPartition
+    val startOffset = 15L
+    val memberId = "98098230493"
+    val committedOffsets = Map(
+      new TopicPartition("foo", 0) -> 23L,
+      new TopicPartition("foo", 1) -> 455L,
+      new TopicPartition("bar", 0) -> 8992L
+    )
+
+    val offsetCommitRecords = createCommittedOffsetRecords(committedOffsets)
+    val groupMetadataRecord = buildStableGroupRecordWithMember(generation = 15,
+      protocolType = "consumer", protocol = "range", memberId)
+    val records = MemoryRecords.withRecords(startOffset, CompressionType.NONE,
+      (offsetCommitRecords ++ Seq(groupMetadataRecord)).toArray: _*)
+
+    expectGroupMetadataLoad(groupMetadataTopicPartition, startOffset, records)
+    EasyMock.replay(replicaManager)
+    groupMetadataManager.loadGroupsAndOffsets(groupMetadataTopicPartition, _ => ())
+
+    assertTrue(partitionLoadTime("partition-load-time-max") >= 0.0)
+    assertTrue(partitionLoadTime( "partition-load-time-avg") >= 0.0)
+  }
 }
diff --git a/core/src/test/scala/unit/kafka/coordinator/group/GroupMetadataTest.scala b/core/src/test/scala/unit/kafka/coordinator/group/GroupMetadataTest.scala
index 3108b150f16ce..82f151c361739 100644
--- a/core/src/test/scala/unit/kafka/coordinator/group/GroupMetadataTest.scala
+++ b/core/src/test/scala/unit/kafka/coordinator/group/GroupMetadataTest.scala
@@ -18,50 +18,59 @@
 package kafka.coordinator.group
 
 import kafka.common.OffsetAndMetadata
+import org.apache.kafka.clients.consumer.ConsumerPartitionAssignor.Subscription
+import org.apache.kafka.clients.consumer.internals.ConsumerProtocol
 import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.protocol.Errors
 import org.apache.kafka.common.utils.Time
 import org.junit.Assert._
 import org.junit.{Before, Test}
-import org.scalatest.junit.JUnitSuite
+
+import scala.collection.JavaConverters._
 
 /**
  * Test group state transitions and other GroupMetadata functionality
  */
-class GroupMetadataTest extends JUnitSuite {
+class GroupMetadataTest {
   private val protocolType = "consumer"
   private val groupId = "groupId"
+  private val groupInstanceId = Some("groupInstanceId")
+  private val memberId = "memberId"
   private val clientId = "clientId"
   private val clientHost = "clientHost"
   private val rebalanceTimeoutMs = 60000
   private val sessionTimeoutMs = 10000
 
   private var group: GroupMetadata = null
+  private var member: MemberMetadata = null
 
   @Before
-  def setUp() {
+  def setUp(): Unit = {
     group = new GroupMetadata("groupId", Empty, Time.SYSTEM)
+    member = new MemberMetadata(memberId, groupId, groupInstanceId, clientId, clientHost, rebalanceTimeoutMs, sessionTimeoutMs,
+      protocolType, List(("range", Array.empty[Byte]), ("roundrobin", Array.empty[Byte])))
   }
 
   @Test
-  def testCanRebalanceWhenStable() {
+  def testCanRebalanceWhenStable(): Unit = {
     assertTrue(group.canRebalance)
   }
 
   @Test
-  def testCanRebalanceWhenCompletingRebalance() {
+  def testCanRebalanceWhenCompletingRebalance(): Unit = {
     group.transitionTo(PreparingRebalance)
     group.transitionTo(CompletingRebalance)
     assertTrue(group.canRebalance)
   }
 
   @Test
-  def testCannotRebalanceWhenPreparingRebalance() {
+  def testCannotRebalanceWhenPreparingRebalance(): Unit = {
     group.transitionTo(PreparingRebalance)
     assertFalse(group.canRebalance)
   }
 
   @Test
-  def testCannotRebalanceWhenDead() {
+  def testCannotRebalanceWhenDead(): Unit = {
     group.transitionTo(PreparingRebalance)
     group.transitionTo(Empty)
     group.transitionTo(Dead)
@@ -69,19 +78,19 @@ class GroupMetadataTest extends JUnitSuite {
   }
 
   @Test
-  def testStableToPreparingRebalanceTransition() {
+  def testStableToPreparingRebalanceTransition(): Unit = {
     group.transitionTo(PreparingRebalance)
     assertState(group, PreparingRebalance)
   }
 
   @Test
-  def testStableToDeadTransition() {
+  def testStableToDeadTransition(): Unit = {
     group.transitionTo(Dead)
     assertState(group, Dead)
   }
 
   @Test
-  def testAwaitingRebalanceToPreparingRebalanceTransition() {
+  def testAwaitingRebalanceToPreparingRebalanceTransition(): Unit = {
     group.transitionTo(PreparingRebalance)
     group.transitionTo(CompletingRebalance)
     group.transitionTo(PreparingRebalance)
@@ -89,21 +98,21 @@ class GroupMetadataTest extends JUnitSuite {
   }
 
   @Test
-  def testPreparingRebalanceToDeadTransition() {
+  def testPreparingRebalanceToDeadTransition(): Unit = {
     group.transitionTo(PreparingRebalance)
     group.transitionTo(Dead)
     assertState(group, Dead)
   }
 
   @Test
-  def testPreparingRebalanceToEmptyTransition() {
+  def testPreparingRebalanceToEmptyTransition(): Unit = {
     group.transitionTo(PreparingRebalance)
     group.transitionTo(Empty)
     assertState(group, Empty)
   }
 
   @Test
-  def testEmptyToDeadTransition() {
+  def testEmptyToDeadTransition(): Unit = {
     group.transitionTo(PreparingRebalance)
     group.transitionTo(Empty)
     group.transitionTo(Dead)
@@ -111,7 +120,7 @@ class GroupMetadataTest extends JUnitSuite {
   }
 
   @Test
-  def testAwaitingRebalanceToStableTransition() {
+  def testAwaitingRebalanceToStableTransition(): Unit = {
     group.transitionTo(PreparingRebalance)
     group.transitionTo(CompletingRebalance)
     group.transitionTo(Stable)
@@ -119,12 +128,12 @@ class GroupMetadataTest extends JUnitSuite {
   }
 
   @Test(expected = classOf[IllegalStateException])
-  def testEmptyToStableIllegalTransition() {
+  def testEmptyToStableIllegalTransition(): Unit = {
     group.transitionTo(Stable)
   }
 
   @Test
-  def testStableToStableIllegalTransition() {
+  def testStableToStableIllegalTransition(): Unit = {
     group.transitionTo(PreparingRebalance)
     group.transitionTo(CompletingRebalance)
     group.transitionTo(Stable)
@@ -137,30 +146,30 @@ class GroupMetadataTest extends JUnitSuite {
   }
 
   @Test(expected = classOf[IllegalStateException])
-  def testEmptyToAwaitingRebalanceIllegalTransition() {
+  def testEmptyToAwaitingRebalanceIllegalTransition(): Unit = {
     group.transitionTo(CompletingRebalance)
   }
 
   @Test(expected = classOf[IllegalStateException])
-  def testPreparingRebalanceToPreparingRebalanceIllegalTransition() {
+  def testPreparingRebalanceToPreparingRebalanceIllegalTransition(): Unit = {
     group.transitionTo(PreparingRebalance)
     group.transitionTo(PreparingRebalance)
   }
 
   @Test(expected = classOf[IllegalStateException])
-  def testPreparingRebalanceToStableIllegalTransition() {
+  def testPreparingRebalanceToStableIllegalTransition(): Unit = {
     group.transitionTo(PreparingRebalance)
     group.transitionTo(Stable)
   }
 
   @Test(expected = classOf[IllegalStateException])
-  def testAwaitingRebalanceToAwaitingRebalanceIllegalTransition() {
+  def testAwaitingRebalanceToAwaitingRebalanceIllegalTransition(): Unit = {
     group.transitionTo(PreparingRebalance)
     group.transitionTo(CompletingRebalance)
     group.transitionTo(CompletingRebalance)
   }
 
-  def testDeadToDeadIllegalTransition() {
+  def testDeadToDeadIllegalTransition(): Unit = {
     group.transitionTo(PreparingRebalance)
     group.transitionTo(Dead)
     group.transitionTo(Dead)
@@ -168,37 +177,37 @@ class GroupMetadataTest extends JUnitSuite {
   }
 
   @Test(expected = classOf[IllegalStateException])
-  def testDeadToStableIllegalTransition() {
+  def testDeadToStableIllegalTransition(): Unit = {
     group.transitionTo(PreparingRebalance)
     group.transitionTo(Dead)
     group.transitionTo(Stable)
   }
 
   @Test(expected = classOf[IllegalStateException])
-  def testDeadToPreparingRebalanceIllegalTransition() {
+  def testDeadToPreparingRebalanceIllegalTransition(): Unit = {
     group.transitionTo(PreparingRebalance)
     group.transitionTo(Dead)
     group.transitionTo(PreparingRebalance)
   }
 
   @Test(expected = classOf[IllegalStateException])
-  def testDeadToAwaitingRebalanceIllegalTransition() {
+  def testDeadToAwaitingRebalanceIllegalTransition(): Unit = {
     group.transitionTo(PreparingRebalance)
     group.transitionTo(Dead)
     group.transitionTo(CompletingRebalance)
   }
 
   @Test
-  def testSelectProtocol() {
+  def testSelectProtocol(): Unit = {
     val memberId = "memberId"
-    val member = new MemberMetadata(memberId, groupId, clientId, clientHost, rebalanceTimeoutMs, sessionTimeoutMs,
+    val member = new MemberMetadata(memberId, groupId, groupInstanceId, clientId, clientHost, rebalanceTimeoutMs, sessionTimeoutMs,
       protocolType, List(("range", Array.empty[Byte]), ("roundrobin", Array.empty[Byte])))
 
     group.add(member)
     assertEquals("range", group.selectProtocol)
 
     val otherMemberId = "otherMemberId"
-    val otherMember = new MemberMetadata(otherMemberId, groupId, clientId, clientHost, rebalanceTimeoutMs,
+    val otherMember = new MemberMetadata(otherMemberId, groupId, groupInstanceId, clientId, clientHost, rebalanceTimeoutMs,
       sessionTimeoutMs, protocolType, List(("roundrobin", Array.empty[Byte]), ("range", Array.empty[Byte])))
 
     group.add(otherMember)
@@ -206,7 +215,7 @@ class GroupMetadataTest extends JUnitSuite {
     assertTrue(Set("range", "roundrobin")(group.selectProtocol))
 
     val lastMemberId = "lastMemberId"
-    val lastMember = new MemberMetadata(lastMemberId, groupId, clientId, clientHost, rebalanceTimeoutMs,
+    val lastMember = new MemberMetadata(lastMemberId, groupId, groupInstanceId, clientId, clientHost, rebalanceTimeoutMs,
       sessionTimeoutMs, protocolType, List(("roundrobin", Array.empty[Byte]), ("range", Array.empty[Byte])))
 
     group.add(lastMember)
@@ -215,19 +224,19 @@ class GroupMetadataTest extends JUnitSuite {
   }
 
   @Test(expected = classOf[IllegalStateException])
-  def testSelectProtocolRaisesIfNoMembers() {
+  def testSelectProtocolRaisesIfNoMembers(): Unit = {
     group.selectProtocol
     fail()
   }
 
   @Test
-  def testSelectProtocolChoosesCompatibleProtocol() {
+  def testSelectProtocolChoosesCompatibleProtocol(): Unit = {
     val memberId = "memberId"
-    val member = new MemberMetadata(memberId, groupId, clientId, clientHost, rebalanceTimeoutMs, sessionTimeoutMs,
+    val member = new MemberMetadata(memberId, groupId, groupInstanceId, clientId, clientHost, rebalanceTimeoutMs, sessionTimeoutMs,
       protocolType, List(("range", Array.empty[Byte]), ("roundrobin", Array.empty[Byte])))
 
     val otherMemberId = "otherMemberId"
-    val otherMember = new MemberMetadata(otherMemberId, groupId, clientId, clientHost, rebalanceTimeoutMs,
+    val otherMember = new MemberMetadata(otherMemberId, groupId, groupInstanceId, clientId, clientHost, rebalanceTimeoutMs,
       sessionTimeoutMs, protocolType, List(("roundrobin", Array.empty[Byte]), ("blah", Array.empty[Byte])))
 
     group.add(member)
@@ -236,14 +245,10 @@ class GroupMetadataTest extends JUnitSuite {
   }
 
   @Test
-  def testSupportsProtocols() {
+  def testSupportsProtocols(): Unit = {
     // by default, the group supports everything
     assertTrue(group.supportsProtocols(protocolType, Set("roundrobin", "range")))
 
-    val memberId = "memberId"
-    val member = new MemberMetadata(memberId, groupId, clientId, clientHost, rebalanceTimeoutMs,
-      sessionTimeoutMs, protocolType, List(("range", Array.empty[Byte]), ("roundrobin", Array.empty[Byte])))
-
     group.add(member)
     group.transitionTo(PreparingRebalance)
     assertTrue(group.supportsProtocols(protocolType, Set("roundrobin", "foo")))
@@ -251,7 +256,7 @@ class GroupMetadataTest extends JUnitSuite {
     assertFalse(group.supportsProtocols(protocolType, Set("foo", "bar")))
 
     val otherMemberId = "otherMemberId"
-    val otherMember = new MemberMetadata(otherMemberId, groupId, clientId, clientHost, rebalanceTimeoutMs,
+    val otherMember = new MemberMetadata(otherMemberId, groupId, groupInstanceId, clientId, clientHost, rebalanceTimeoutMs,
       sessionTimeoutMs, protocolType, List(("roundrobin", Array.empty[Byte]), ("blah", Array.empty[Byte])))
 
     group.add(otherMember)
@@ -262,10 +267,59 @@ class GroupMetadataTest extends JUnitSuite {
   }
 
   @Test
-  def testInitNextGeneration() {
+  def testSubscribedTopics(): Unit = {
+    // not able to compute it for a newly created group
+    assertEquals(None, group.getSubscribedTopics)
+
     val memberId = "memberId"
-    val member = new MemberMetadata(memberId, groupId, clientId, clientHost, rebalanceTimeoutMs, sessionTimeoutMs,
-      protocolType, List(("roundrobin", Array.empty[Byte])))
+    val member = new MemberMetadata(memberId, groupId, groupInstanceId, clientId, clientHost, rebalanceTimeoutMs,
+      sessionTimeoutMs, protocolType, List(("range", ConsumerProtocol.serializeSubscription(new Subscription(List("foo").asJava)).array())))
+
+    group.transitionTo(PreparingRebalance)
+    group.add(member)
+
+    group.initNextGeneration()
+
+    assertEquals(Some(Set("foo")), group.getSubscribedTopics)
+
+    group.transitionTo(PreparingRebalance)
+    group.remove(memberId)
+
+    group.initNextGeneration()
+
+    assertEquals(Some(Set.empty), group.getSubscribedTopics)
+
+    val memberWithFaultyProtocol  = new MemberMetadata(memberId, groupId, groupInstanceId, clientId, clientHost, rebalanceTimeoutMs,
+      sessionTimeoutMs, protocolType, List(("range", Array.empty[Byte])))
+
+    group.transitionTo(PreparingRebalance)
+    group.add(memberWithFaultyProtocol)
+
+    group.initNextGeneration()
+
+    assertEquals(None, group.getSubscribedTopics)
+  }
+
+  @Test
+  def testSubscribedTopicsNonConsumerGroup(): Unit = {
+    // not able to compute it for a newly created group
+    assertEquals(None, group.getSubscribedTopics)
+
+    val memberId = "memberId"
+    val member = new MemberMetadata(memberId, groupId, groupInstanceId, clientId, clientHost, rebalanceTimeoutMs,
+      sessionTimeoutMs, "My Protocol", List(("range", Array.empty[Byte])))
+
+    group.transitionTo(PreparingRebalance)
+    group.add(member)
+
+    group.initNextGeneration()
+
+    assertEquals(None, group.getSubscribedTopics)
+  }
+
+  @Test
+  def testInitNextGeneration(): Unit = {
+    member.supportedProtocols = List(("roundrobin", Array.empty[Byte]))
 
     group.transitionTo(PreparingRebalance)
     group.add(member, _ => ())
@@ -280,7 +334,7 @@ class GroupMetadataTest extends JUnitSuite {
   }
 
   @Test
-  def testInitNextGenerationEmptyGroup() {
+  def testInitNextGenerationEmptyGroup(): Unit = {
     assertEquals(Empty, group.currentState)
     assertEquals(0, group.generationId)
     assertNull(group.protocolOrNull)
@@ -465,7 +519,89 @@ class GroupMetadataTest extends JUnitSuite {
     assertFalse(group.hasPendingOffsetCommitsFromProducer(producerId))
   }
 
-  private def assertState(group: GroupMetadata, targetState: GroupState) {
+  @Test(expected = classOf[IllegalArgumentException])
+  def testReplaceGroupInstanceWithEmptyGroupInstanceId(): Unit = {
+    group.add(member)
+    group.addStaticMember(groupInstanceId, memberId)
+    assertTrue(group.isLeader(memberId))
+    assertEquals(memberId, group.getStaticMemberId(groupInstanceId))
+
+    val newMemberId = "newMemberId"
+    group.replaceGroupInstance(memberId, newMemberId, Option.empty)
+  }
+
+  @Test(expected = classOf[IllegalArgumentException])
+  def testReplaceGroupInstanceWithNonExistingMember(): Unit = {
+    val newMemberId = "newMemberId"
+    group.replaceGroupInstance(memberId, newMemberId, groupInstanceId)
+  }
+
+  @Test
+  def testReplaceGroupInstance(): Unit = {
+    var joinAwaitingMemberFenced = false
+    group.add(member, joinGroupResult => {
+      joinAwaitingMemberFenced = joinGroupResult.error == Errors.FENCED_INSTANCE_ID
+    })
+    var syncAwaitingMemberFenced = false
+    member.awaitingSyncCallback = syncGroupResult => {
+      syncAwaitingMemberFenced = syncGroupResult.error == Errors.FENCED_INSTANCE_ID
+    }
+    group.addStaticMember(groupInstanceId, memberId)
+    assertTrue(group.isLeader(memberId))
+    assertEquals(memberId, group.getStaticMemberId(groupInstanceId))
+
+    val newMemberId = "newMemberId"
+    group.replaceGroupInstance(memberId, newMemberId, groupInstanceId)
+    assertTrue(group.isLeader(newMemberId))
+    assertEquals(newMemberId, group.getStaticMemberId(groupInstanceId))
+    assertTrue(joinAwaitingMemberFenced)
+    assertTrue(syncAwaitingMemberFenced)
+    assertFalse(member.isAwaitingJoin)
+    assertFalse(member.isAwaitingSync)
+  }
+
+  @Test
+  def testInvokeJoinCallback(): Unit = {
+    var invoked = false
+    group.add(member, _ => {
+      invoked = true
+    })
+
+    assertTrue(group.hasAllMembersJoined)
+    group.maybeInvokeJoinCallback(member, GroupCoordinator.joinError(member.memberId, Errors.NONE))
+    assertTrue(invoked)
+    assertFalse(member.isAwaitingJoin)
+  }
+
+  @Test
+  def testNotInvokeJoinCallback(): Unit = {
+    group.add(member)
+
+    assertFalse(member.isAwaitingJoin)
+    group.maybeInvokeJoinCallback(member, GroupCoordinator.joinError(member.memberId, Errors.NONE))
+    assertFalse(member.isAwaitingJoin)
+  }
+
+  @Test
+  def testInvokeSyncCallback(): Unit = {
+    group.add(member)
+    member.awaitingSyncCallback = _ => {}
+
+    val invoked = group.maybeInvokeSyncCallback(member, SyncGroupResult(Array.empty, Errors.NONE))
+    assertTrue(invoked)
+    assertFalse(member.isAwaitingSync)
+  }
+
+  @Test
+  def testNotInvokeSyncCallback(): Unit = {
+    group.add(member)
+
+    val invoked = group.maybeInvokeSyncCallback(member, SyncGroupResult(Array.empty, Errors.NONE))
+    assertFalse(invoked)
+    assertFalse(member.isAwaitingSync)
+  }
+
+  private def assertState(group: GroupMetadata, targetState: GroupState): Unit = {
     val states: Set[GroupState] = Set(Stable, PreparingRebalance, CompletingRebalance, Dead)
     val otherStates = states - targetState
     otherStates.foreach { otherState =>
diff --git a/core/src/test/scala/unit/kafka/coordinator/group/MemberMetadataTest.scala b/core/src/test/scala/unit/kafka/coordinator/group/MemberMetadataTest.scala
index 48c92704ba4b7..d348cc8c931c3 100644
--- a/core/src/test/scala/unit/kafka/coordinator/group/MemberMetadataTest.scala
+++ b/core/src/test/scala/unit/kafka/coordinator/group/MemberMetadataTest.scala
@@ -20,10 +20,10 @@ import java.util.Arrays
 
 import org.junit.Assert._
 import org.junit.Test
-import org.scalatest.junit.JUnitSuite
 
-class MemberMetadataTest extends JUnitSuite {
+class MemberMetadataTest {
   val groupId = "groupId"
+  val groupInstanceId = Some("groupInstanceId")
   val clientId = "clientId"
   val clientHost = "clientHost"
   val memberId = "memberId"
@@ -33,10 +33,10 @@ class MemberMetadataTest extends JUnitSuite {
 
 
   @Test
-  def testMatchesSupportedProtocols() {
+  def testMatchesSupportedProtocols(): Unit = {
     val protocols = List(("range", Array.empty[Byte]))
 
-    val member = new MemberMetadata(memberId, groupId, clientId, clientHost, rebalanceTimeoutMs, sessionTimeoutMs,
+    val member = new MemberMetadata(memberId, groupId, groupInstanceId, clientId, clientHost, rebalanceTimeoutMs, sessionTimeoutMs,
       protocolType, protocols)
     assertTrue(member.matches(protocols))
     assertFalse(member.matches(List(("range", Array[Byte](0)))))
@@ -45,44 +45,52 @@ class MemberMetadataTest extends JUnitSuite {
   }
 
   @Test
-  def testVoteForPreferredProtocol() {
+  def testVoteForPreferredProtocol(): Unit = {
     val protocols = List(("range", Array.empty[Byte]), ("roundrobin", Array.empty[Byte]))
 
-    val member = new MemberMetadata(memberId, groupId, clientId, clientHost, rebalanceTimeoutMs, sessionTimeoutMs,
+    val member = new MemberMetadata(memberId, groupId, groupInstanceId, clientId, clientHost, rebalanceTimeoutMs, sessionTimeoutMs,
       protocolType, protocols)
     assertEquals("range", member.vote(Set("range", "roundrobin")))
     assertEquals("roundrobin", member.vote(Set("blah", "roundrobin")))
   }
 
   @Test
-  def testMetadata() {
+  def testMetadata(): Unit = {
     val protocols = List(("range", Array[Byte](0)), ("roundrobin", Array[Byte](1)))
 
-    val member = new MemberMetadata(memberId, groupId, clientId, clientHost, rebalanceTimeoutMs, sessionTimeoutMs,
+    val member = new MemberMetadata(memberId, groupId, groupInstanceId, clientId, clientHost, rebalanceTimeoutMs, sessionTimeoutMs,
       protocolType, protocols)
     assertTrue(Arrays.equals(Array[Byte](0), member.metadata("range")))
     assertTrue(Arrays.equals(Array[Byte](1), member.metadata("roundrobin")))
   }
 
   @Test(expected = classOf[IllegalArgumentException])
-  def testMetadataRaisesOnUnsupportedProtocol() {
+  def testMetadataRaisesOnUnsupportedProtocol(): Unit = {
     val protocols = List(("range", Array.empty[Byte]), ("roundrobin", Array.empty[Byte]))
 
-    val member = new MemberMetadata(memberId, groupId, clientId, clientHost, rebalanceTimeoutMs, sessionTimeoutMs,
+    val member = new MemberMetadata(memberId, groupId, groupInstanceId, clientId, clientHost, rebalanceTimeoutMs, sessionTimeoutMs,
       protocolType, protocols)
     member.metadata("blah")
     fail()
   }
 
   @Test(expected = classOf[IllegalArgumentException])
-  def testVoteRaisesOnNoSupportedProtocols() {
+  def testVoteRaisesOnNoSupportedProtocols(): Unit = {
     val protocols = List(("range", Array.empty[Byte]), ("roundrobin", Array.empty[Byte]))
 
-    val member = new MemberMetadata(memberId, groupId, clientId, clientHost, rebalanceTimeoutMs, sessionTimeoutMs,
+    val member = new MemberMetadata(memberId, groupId, groupInstanceId, clientId, clientHost, rebalanceTimeoutMs, sessionTimeoutMs,
       protocolType, protocols)
     member.vote(Set("blah"))
     fail()
   }
 
+  @Test
+  def testHasValidGroupInstanceId(): Unit = {
+    val protocols = List(("range", Array[Byte](0)), ("roundrobin", Array[Byte](1)))
 
+    val member = new MemberMetadata(memberId, groupId, groupInstanceId, clientId, clientHost, rebalanceTimeoutMs, sessionTimeoutMs,
+      protocolType, protocols)
+    assertTrue(member.isStaticMember)
+    assertEquals(groupInstanceId, member.groupInstanceId)
+  }
 }
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 b2cc4a51a592f..148d22b57f881 100644
--- a/core/src/test/scala/unit/kafka/coordinator/transaction/ProducerIdManagerTest.scala
+++ b/core/src/test/scala/unit/kafka/coordinator/transaction/ProducerIdManagerTest.scala
@@ -32,7 +32,7 @@ class ProducerIdManagerTest {
   }
 
   @Test
-  def testGetProducerId() {
+  def testGetProducerId(): Unit = {
     var zkVersion: Option[Int] = None
     var data: Array[Byte] = null
     EasyMock.expect(zkClient.getDataAndVersion(EasyMock.anyString)).andAnswer(new IAnswer[(Option[Array[Byte]], Int)] {
@@ -75,7 +75,7 @@ class ProducerIdManagerTest {
   }
 
   @Test(expected = classOf[KafkaException])
-  def testExceedProducerIdLimit() {
+  def testExceedProducerIdLimit(): Unit = {
     EasyMock.expect(zkClient.getDataAndVersion(EasyMock.anyString)).andAnswer(new IAnswer[(Option[Array[Byte]], Int)] {
       override def answer(): (Option[Array[Byte]], Int) = {
         val json = ProducerIdManager.generateProducerIdBlockJson(
diff --git a/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionCoordinatorConcurrencyTest.scala b/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionCoordinatorConcurrencyTest.scala
index 3cf956629b28a..79611cbf3a908 100644
--- a/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionCoordinatorConcurrencyTest.scala
+++ b/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionCoordinatorConcurrencyTest.scala
@@ -22,12 +22,13 @@ import kafka.coordinator.AbstractCoordinatorConcurrencyTest
 import kafka.coordinator.AbstractCoordinatorConcurrencyTest._
 import kafka.coordinator.transaction.TransactionCoordinatorConcurrencyTest._
 import kafka.log.Log
-import kafka.server.{DelayedOperationPurgatory, FetchDataInfo, KafkaConfig, LogOffsetMetadata, MetadataCache}
+import kafka.server.{DelayedOperationPurgatory, FetchDataInfo, FetchLogEnd, KafkaConfig, LogOffsetMetadata, MetadataCache}
 import kafka.utils.timer.MockTimer
 import kafka.utils.{Pool, TestUtils}
 import org.apache.kafka.clients.{ClientResponse, NetworkClient}
 import org.apache.kafka.common.{Node, TopicPartition}
 import org.apache.kafka.common.internals.Topic.TRANSACTION_STATE_TOPIC_NAME
+import org.apache.kafka.common.metrics.Metrics
 import org.apache.kafka.common.protocol.{ApiKeys, Errors}
 import org.apache.kafka.common.record.{CompressionType, FileRecords, MemoryRecords, SimpleRecord}
 import org.apache.kafka.common.requests._
@@ -60,7 +61,7 @@ class TransactionCoordinatorConcurrencyTest extends AbstractCoordinatorConcurren
     (0 until numPartitions).map { i => (i, mutable.ArrayBuffer[SimpleRecord]()) }.toMap
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     super.setUp()
 
     EasyMock.expect(zkClient.getTopicPartitionCount(TRANSACTION_STATE_TOPIC_NAME))
@@ -68,7 +69,7 @@ class TransactionCoordinatorConcurrencyTest extends AbstractCoordinatorConcurren
       .anyTimes()
     EasyMock.replay(zkClient)
 
-    txnStateManager = new TransactionStateManager(0, zkClient, scheduler, replicaManager, txnConfig, time)
+    txnStateManager = new TransactionStateManager(0, zkClient, scheduler, replicaManager, txnConfig, time, new Metrics())
     for (i <- 0 until numPartitions)
       txnStateManager.addLoadedTransactionsToCache(i, coordinatorEpoch, new Pool[String, TransactionMetadata]())
 
@@ -114,7 +115,7 @@ class TransactionCoordinatorConcurrencyTest extends AbstractCoordinatorConcurren
   }
 
   @After
-  override def tearDown() {
+  override def tearDown(): Unit = {
     try {
       EasyMock.reset(zkClient, replicaManager)
       transactionCoordinator.shutdown()
@@ -251,15 +252,14 @@ class TransactionCoordinatorConcurrencyTest extends AbstractCoordinatorConcurren
 
     val topicPartition = new TopicPartition(TRANSACTION_STATE_TOPIC_NAME, partitionId)
     val startOffset = replicaManager.getLogEndOffset(topicPartition).getOrElse(20L)
-    val records = MemoryRecords.withRecords(startOffset, CompressionType.NONE, txnRecordsByPartition(partitionId): _*)
+    val records = MemoryRecords.withRecords(startOffset, CompressionType.NONE, txnRecordsByPartition(partitionId).toArray: _*)
     val endOffset = startOffset + records.records.asScala.size
 
     EasyMock.expect(logMock.logStartOffset).andStubReturn(startOffset)
     EasyMock.expect(logMock.read(EasyMock.eq(startOffset),
       maxLength = EasyMock.anyInt(),
-      maxOffset = EasyMock.eq(None),
-      minOneMessage = EasyMock.eq(true),
-      includeAbortedTxns = EasyMock.eq(false)))
+      isolation = EasyMock.eq(FetchLogEnd),
+      minOneMessage = EasyMock.eq(true)))
       .andReturn(FetchDataInfo(LogOffsetMetadata(startOffset), fileRecordsMock))
 
     EasyMock.expect(fileRecordsMock.sizeInBytes()).andStubReturn(records.sizeInBytes)
@@ -344,7 +344,7 @@ class TransactionCoordinatorConcurrencyTest extends AbstractCoordinatorConcurren
 
   class LoadTxnPartitionAction(txnTopicPartitionId: Int) extends Action {
     override def run(): Unit = {
-      transactionCoordinator.handleTxnImmigration(txnTopicPartitionId, coordinatorEpoch)
+      transactionCoordinator.onElection(txnTopicPartitionId, coordinatorEpoch)
     }
     override def await(): Unit = {
       allTransactions.foreach { txn =>
@@ -358,7 +358,7 @@ class TransactionCoordinatorConcurrencyTest extends AbstractCoordinatorConcurren
   class UnloadTxnPartitionAction(txnTopicPartitionId: Int) extends Action {
     val txnRecords: mutable.ArrayBuffer[SimpleRecord] = mutable.ArrayBuffer[SimpleRecord]()
     override def run(): Unit = {
-      transactionCoordinator.handleTxnEmigration(txnTopicPartitionId, coordinatorEpoch)
+      transactionCoordinator.onResignation(txnTopicPartitionId, Some(coordinatorEpoch))
     }
     override def await(): Unit = {
       allTransactions.foreach { txn =>
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 75f06d5d58457..a637d16551e37 100644
--- a/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionCoordinatorTest.scala
+++ b/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionCoordinatorTest.scala
@@ -107,7 +107,7 @@ class TransactionCoordinatorTest {
       .andReturn(Right(None))
       .once()
 
-    EasyMock.expect(transactionManager.putTransactionStateIfNotExists(EasyMock.eq(transactionalId), EasyMock.capture(capturedTxn)))
+    EasyMock.expect(transactionManager.putTransactionStateIfNotExists(EasyMock.capture(capturedTxn)))
       .andAnswer(new IAnswer[Either[Errors, CoordinatorEpochAndTxnMetadata]] {
         override def answer(): Either[Errors, CoordinatorEpochAndTxnMetadata] = {
           assertTrue(capturedTxn.hasCaptured)
@@ -512,7 +512,7 @@ class TransactionCoordinatorTest {
     EasyMock.expect(transactionManager.validateTransactionTimeoutMs(EasyMock.anyInt()))
       .andReturn(true)
 
-    EasyMock.expect(transactionManager.putTransactionStateIfNotExists(EasyMock.eq(transactionalId), EasyMock.anyObject[TransactionMetadata]()))
+    EasyMock.expect(transactionManager.putTransactionStateIfNotExists(EasyMock.anyObject[TransactionMetadata]()))
       .andReturn(Right(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, txnMetadata)))
       .anyTimes()
 
@@ -551,7 +551,7 @@ class TransactionCoordinatorTest {
     EasyMock.expect(transactionManager.validateTransactionTimeoutMs(EasyMock.anyInt()))
       .andReturn(true)
 
-    EasyMock.expect(transactionManager.putTransactionStateIfNotExists(EasyMock.eq(transactionalId), EasyMock.anyObject[TransactionMetadata]()))
+    EasyMock.expect(transactionManager.putTransactionStateIfNotExists(EasyMock.anyObject[TransactionMetadata]()))
       .andReturn(Right(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, txnMetadata)))
       .anyTimes()
 
@@ -593,7 +593,7 @@ class TransactionCoordinatorTest {
     EasyMock.expect(transactionMarkerChannelManager.removeMarkersForTxnTopicPartition(0))
     EasyMock.replay(transactionManager, transactionMarkerChannelManager)
 
-    coordinator.handleTxnEmigration(0, coordinatorEpoch)
+    coordinator.onResignation(0, Some(coordinatorEpoch))
 
     EasyMock.verify(transactionManager, transactionMarkerChannelManager)
   }
diff --git a/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionLogTest.scala b/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionLogTest.scala
index c0edec7acb419..e9bdf3110a890 100644
--- a/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionLogTest.scala
+++ b/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionLogTest.scala
@@ -17,16 +17,16 @@
 package kafka.coordinator.transaction
 
 
+import kafka.utils.TestUtils
 import org.apache.kafka.common.TopicPartition
-import org.apache.kafka.common.record.{CompressionType, SimpleRecord, MemoryRecords}
-
+import org.apache.kafka.common.record.{CompressionType, MemoryRecords, SimpleRecord}
 import org.junit.Assert.assertEquals
 import org.junit.Test
-import org.scalatest.junit.JUnitSuite
+import org.scalatest.Assertions.intercept
 
 import scala.collection.JavaConverters._
 
-class TransactionLogTest extends JUnitSuite {
+class TransactionLogTest {
 
   val producerEpoch: Short = 0
   val transactionTimeoutMs: Int = 1000
@@ -38,7 +38,7 @@ class TransactionLogTest extends JUnitSuite {
     new TopicPartition("topic2", 2))
 
   @Test
-  def shouldThrowExceptionWriteInvalidTxn() {
+  def shouldThrowExceptionWriteInvalidTxn(): Unit = {
     val transactionalId = "transactionalId"
     val producerId = 23423L
 
@@ -51,7 +51,7 @@ class TransactionLogTest extends JUnitSuite {
   }
 
   @Test
-  def shouldReadWriteMessages() {
+  def shouldReadWriteMessages(): Unit = {
     val pidMappings = Map[String, Long]("zero" -> 0L,
       "one" -> 1L,
       "two" -> 2L,
@@ -86,7 +86,7 @@ class TransactionLogTest extends JUnitSuite {
     for (record <- records.records.asScala) {
       val txnKey = TransactionLog.readTxnRecordKey(record.key)
       val transactionalId = txnKey.transactionalId
-      val txnMetadata = TransactionLog.readTxnRecordValue(transactionalId, record.value)
+      val txnMetadata = TransactionLog.readTxnRecordValue(transactionalId, record.value).get
 
       assertEquals(pidMappings(transactionalId), txnMetadata.producerId)
       assertEquals(producerEpoch, txnMetadata.producerEpoch)
@@ -104,4 +104,38 @@ class TransactionLogTest extends JUnitSuite {
     assertEquals(pidMappings.size, count)
   }
 
+  @Test
+  def testTransactionMetadataParsing(): Unit = {
+    val transactionalId = "id"
+    val producerId = 1334L
+    val topicPartition = new TopicPartition("topic", 0)
+
+    val txnMetadata = TransactionMetadata(transactionalId, producerId, producerEpoch,
+      transactionTimeoutMs, Ongoing, 0)
+    txnMetadata.addPartitions(Set(topicPartition))
+
+    val keyBytes = TransactionLog.keyToBytes(transactionalId)
+    val valueBytes = TransactionLog.valueToBytes(txnMetadata.prepareNoTransit())
+    val transactionMetadataRecord = TestUtils.records(Seq(
+      new SimpleRecord(keyBytes, valueBytes)
+    )).records.asScala.head
+
+    val (keyStringOpt, valueStringOpt) = TransactionLog.formatRecordKeyAndValue(transactionMetadataRecord)
+    assertEquals(Some(s"transaction_metadata::transactionalId=$transactionalId"), keyStringOpt)
+    assertEquals(Some(s"producerId:$producerId,producerEpoch:$producerEpoch,state=Ongoing," +
+      s"partitions=[$topicPartition],txnLastUpdateTimestamp=0,txnTimeoutMs=$transactionTimeoutMs"), valueStringOpt)
+  }
+
+  @Test
+  def testTransactionMetadataTombstoneParsing(): Unit = {
+    val transactionalId = "id"
+    val transactionMetadataRecord = TestUtils.records(Seq(
+      new SimpleRecord(TransactionLog.keyToBytes(transactionalId), null)
+    )).records.asScala.head
+
+    val (keyStringOpt, valueStringOpt) = TransactionLog.formatRecordKeyAndValue(transactionMetadataRecord)
+    assertEquals(Some(s"transaction_metadata::transactionalId=$transactionalId"), keyStringOpt)
+    assertEquals(Some(""), valueStringOpt)
+  }
+
 }
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 84f3dffbfbd58..1db652acd7adc 100644
--- a/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionMarkerRequestCompletionHandlerTest.scala
+++ b/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionMarkerRequestCompletionHandlerTest.scala
@@ -189,6 +189,11 @@ class TransactionMarkerRequestCompletionHandlerTest {
     verifyRetriesPartitionOnError(Errors.NOT_ENOUGH_REPLICAS_AFTER_APPEND)
   }
 
+  @Test
+  def shouldRetryPartitionWhenKafkaStorageError(): Unit = {
+    verifyRetriesPartitionOnError(Errors.KAFKA_STORAGE_ERROR)
+  }
+
   @Test
   def shouldRemoveTopicPartitionFromWaitingSetOnUnsupportedForMessageFormat(): Unit = {
     mockCache()
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 34a37bebe4d74..b628534058f2b 100644
--- a/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionStateManagerTest.scala
+++ b/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionStateManagerTest.scala
@@ -16,28 +16,31 @@
  */
 package kafka.coordinator.transaction
 
+import java.lang.management.ManagementFactory
 import java.nio.ByteBuffer
+import java.util.concurrent.CountDownLatch
 import java.util.concurrent.locks.ReentrantLock
 
-import kafka.log.Log
-import kafka.server.{FetchDataInfo, LogOffsetMetadata, ReplicaManager}
-import kafka.utils.{MockScheduler, Pool}
-import kafka.utils.TestUtils.fail
+import javax.management.ObjectName
+import kafka.log.{AppendOrigin, Log}
+import kafka.server.{FetchDataInfo, FetchLogEnd, LogOffsetMetadata, ReplicaManager}
+import kafka.utils.{MockScheduler, Pool, TestUtils}
 import kafka.zk.KafkaZkClient
 import org.apache.kafka.common.TopicPartition
 import org.apache.kafka.common.internals.Topic.TRANSACTION_STATE_TOPIC_NAME
+import org.apache.kafka.common.metrics.{JmxReporter, Metrics}
 import org.apache.kafka.common.protocol.Errors
 import org.apache.kafka.common.record._
 import org.apache.kafka.common.requests.ProduceResponse.PartitionResponse
 import org.apache.kafka.common.requests.TransactionResult
 import org.apache.kafka.common.utils.MockTime
+import org.easymock.{Capture, EasyMock, IAnswer}
 import org.junit.Assert.{assertEquals, assertFalse, assertTrue}
 import org.junit.{After, Before, Test}
-import org.easymock.{Capture, EasyMock, IAnswer}
+import org.scalatest.Assertions.fail
 
-import scala.collection.Map
-import scala.collection.mutable
 import scala.collection.JavaConverters._
+import scala.collection.{Map, mutable}
 
 class TransactionStateManagerTest {
 
@@ -59,9 +62,10 @@ class TransactionStateManagerTest {
     .anyTimes()
 
   EasyMock.replay(zkClient)
+  val metrics = new Metrics()
 
   val txnConfig = TransactionConfig()
-  val transactionManager: TransactionStateManager = new TransactionStateManager(0, zkClient, scheduler, replicaManager, txnConfig, time)
+  val transactionManager: TransactionStateManager = new TransactionStateManager(0, zkClient, scheduler, replicaManager, txnConfig, time, metrics)
 
   val transactionalId1: String = "one"
   val transactionalId2: String = "two"
@@ -74,20 +78,20 @@ class TransactionStateManagerTest {
   var expectedError: Errors = Errors.NONE
 
   @Before
-  def setUp() {
+  def setUp(): Unit = {
     // make sure the transactional id hashes to the assigning partition id
     assertEquals(partitionId, transactionManager.partitionFor(transactionalId1))
     assertEquals(partitionId, transactionManager.partitionFor(transactionalId2))
   }
 
   @After
-  def tearDown() {
+  def tearDown(): Unit = {
     EasyMock.reset(zkClient, replicaManager)
     transactionManager.shutdown()
   }
 
   @Test
-  def testValidateTransactionTimeout() {
+  def testValidateTransactionTimeout(): Unit = {
     assertTrue(transactionManager.validateTransactionTimeoutMs(1))
     assertFalse(transactionManager.validateTransactionTimeoutMs(-1))
     assertFalse(transactionManager.validateTransactionTimeoutMs(0))
@@ -96,20 +100,114 @@ class TransactionStateManagerTest {
   }
 
   @Test
-  def testAddGetPids() {
+  def testAddGetPids(): Unit = {
     transactionManager.addLoadedTransactionsToCache(partitionId, coordinatorEpoch, new Pool[String, TransactionMetadata]())
 
     assertEquals(Right(None), transactionManager.getTransactionState(transactionalId1))
     assertEquals(Right(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, txnMetadata1)),
-      transactionManager.putTransactionStateIfNotExists(transactionalId1, txnMetadata1))
+      transactionManager.putTransactionStateIfNotExists(txnMetadata1))
     assertEquals(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, txnMetadata1))),
       transactionManager.getTransactionState(transactionalId1))
-    assertEquals(Right(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, txnMetadata1)),
-      transactionManager.putTransactionStateIfNotExists(transactionalId1, txnMetadata2))
+    assertEquals(Right(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, txnMetadata2)),
+      transactionManager.putTransactionStateIfNotExists(txnMetadata2))
+  }
+
+  @Test
+  def testDeletePartition(): Unit = {
+    val metadata1 = transactionMetadata("b", 5L)
+    val metadata2 = transactionMetadata("a", 10L)
+
+    assertEquals(0, transactionManager.partitionFor(metadata1.transactionalId))
+    assertEquals(1, transactionManager.partitionFor(metadata2.transactionalId))
+
+    transactionManager.addLoadedTransactionsToCache(0, coordinatorEpoch, new Pool[String, TransactionMetadata]())
+    transactionManager.putTransactionStateIfNotExists(metadata1)
+
+    transactionManager.addLoadedTransactionsToCache(1, coordinatorEpoch, new Pool[String, TransactionMetadata]())
+    transactionManager.putTransactionStateIfNotExists(metadata2)
+
+    def cachedProducerEpoch(transactionalId: String): Option[Short] = {
+      transactionManager.getTransactionState(transactionalId) match {
+        case Right(Some(epochAndTxnMetadata)) => Some(epochAndTxnMetadata.transactionMetadata.producerEpoch)
+        case _ => None
+      }
+    }
+
+    assertEquals(Some(metadata1.producerEpoch), cachedProducerEpoch(metadata1.transactionalId))
+    assertEquals(Some(metadata2.producerEpoch), cachedProducerEpoch(metadata2.transactionalId))
+
+    transactionManager.removeTransactionsForTxnTopicPartition(0)
+
+    assertEquals(None, cachedProducerEpoch(metadata1.transactionalId))
+    assertEquals(Some(metadata2.producerEpoch), cachedProducerEpoch(metadata2.transactionalId))
+  }
+
+  @Test
+  def testDeleteLoadingPartition(): Unit = {
+    // Verify the handling of a call to delete state for a partition while it is in the
+    // process of being loaded. Basically should be treated as a no-op.
+
+    val startOffset = 0L
+    val endOffset = 1L
+
+    val fileRecordsMock = EasyMock.mock[FileRecords](classOf[FileRecords])
+    val logMock = EasyMock.mock[Log](classOf[Log])
+    EasyMock.expect(replicaManager.getLog(topicPartition)).andStubReturn(Some(logMock))
+    EasyMock.expect(logMock.logStartOffset).andStubReturn(startOffset)
+    EasyMock.expect(logMock.read(EasyMock.eq(startOffset),
+      maxLength = EasyMock.anyInt(),
+      isolation = EasyMock.eq(FetchLogEnd),
+      minOneMessage = EasyMock.eq(true))
+    ).andReturn(FetchDataInfo(LogOffsetMetadata(startOffset), fileRecordsMock))
+    EasyMock.expect(replicaManager.getLogEndOffset(topicPartition)).andStubReturn(Some(endOffset))
+
+    txnMetadata1.state = PrepareCommit
+    txnMetadata1.addPartitions(Set[TopicPartition](
+      new TopicPartition("topic1", 0),
+      new TopicPartition("topic1", 1)))
+    val records = MemoryRecords.withRecords(startOffset, CompressionType.NONE,
+      new SimpleRecord(txnMessageKeyBytes1, TransactionLog.valueToBytes(txnMetadata1.prepareNoTransit())))
+
+    // We create a latch which is awaited while the log is loading. This ensures that the deletion
+    // is triggered before the loading returns
+    val latch = new CountDownLatch(1)
+
+    EasyMock.expect(fileRecordsMock.sizeInBytes()).andStubReturn(records.sizeInBytes)
+    val bufferCapture = EasyMock.newCapture[ByteBuffer]
+    fileRecordsMock.readInto(EasyMock.capture(bufferCapture), EasyMock.anyInt())
+    EasyMock.expectLastCall().andAnswer(new IAnswer[Unit] {
+      override def answer: Unit = {
+        latch.await()
+        val buffer = bufferCapture.getValue
+        buffer.put(records.buffer.duplicate)
+        buffer.flip()
+      }
+    })
+
+    EasyMock.replay(logMock, fileRecordsMock, replicaManager)
+
+    val coordinatorEpoch = 0
+    val partitionAndLeaderEpoch = TransactionPartitionAndLeaderEpoch(partitionId, coordinatorEpoch)
+
+    val loadingThread = new Thread(new Runnable {
+      override def run(): Unit = transactionManager.loadTransactionsForTxnTopicPartition(partitionId, coordinatorEpoch, (_, _, _, _, _) => ())
+    })
+    loadingThread.start()
+    TestUtils.waitUntilTrue(() => transactionManager.loadingPartitions.contains(partitionAndLeaderEpoch),
+      "Timed out waiting for loading partition", pause = 10)
+
+    transactionManager.removeTransactionsForTxnTopicPartition(partitionId)
+    assertFalse(transactionManager.loadingPartitions.contains(partitionAndLeaderEpoch))
+
+    latch.countDown()
+    loadingThread.join()
+
+    // Verify that transaction state was not loaded
+    assertEquals(Left(Errors.NOT_COORDINATOR), transactionManager.getTransactionState(txnMetadata1.transactionalId))
   }
 
   @Test
-  def testLoadAndRemoveTransactionsForPartition() {
+  def testLoadAndRemoveTransactionsForPartition(): Unit = {
     // generate transaction log messages for two pids traces:
 
     // pid1's transaction started with two partitions
@@ -156,7 +254,7 @@ class TransactionStateManagerTest {
     txnRecords += new SimpleRecord(txnMessageKeyBytes2, TransactionLog.valueToBytes(txnMetadata2.prepareNoTransit()))
 
     val startOffset = 15L   // it should work for any start offset
-    val records = MemoryRecords.withRecords(startOffset, CompressionType.NONE, txnRecords: _*)
+    val records = MemoryRecords.withRecords(startOffset, CompressionType.NONE, txnRecords.toArray: _*)
 
     prepareTxnLog(topicPartition, startOffset, records)
 
@@ -213,7 +311,7 @@ class TransactionStateManagerTest {
     transactionManager.addLoadedTransactionsToCache(partitionId, coordinatorEpoch, new Pool[String, TransactionMetadata]())
 
     // first insert the initial transaction metadata
-    transactionManager.putTransactionStateIfNotExists(transactionalId1, txnMetadata1)
+    transactionManager.putTransactionStateIfNotExists(txnMetadata1)
 
     prepareForTxnMessageAppend(Errors.NONE)
     expectedError = Errors.NONE
@@ -232,7 +330,7 @@ class TransactionStateManagerTest {
   @Test
   def testAppendFailToCoordinatorNotAvailableError(): Unit = {
     transactionManager.addLoadedTransactionsToCache(partitionId, coordinatorEpoch, new Pool[String, TransactionMetadata]())
-    transactionManager.putTransactionStateIfNotExists(transactionalId1, txnMetadata1)
+    transactionManager.putTransactionStateIfNotExists(txnMetadata1)
 
     expectedError = Errors.COORDINATOR_NOT_AVAILABLE
     var failedMetadata = txnMetadata1.prepareAddPartitions(Set[TopicPartition](new TopicPartition("topic2", 0)), time.milliseconds())
@@ -264,7 +362,7 @@ class TransactionStateManagerTest {
   @Test
   def testAppendFailToNotCoordinatorError(): Unit = {
     transactionManager.addLoadedTransactionsToCache(partitionId, coordinatorEpoch, new Pool[String, TransactionMetadata]())
-    transactionManager.putTransactionStateIfNotExists(transactionalId1, txnMetadata1)
+    transactionManager.putTransactionStateIfNotExists(txnMetadata1)
 
     expectedError = Errors.NOT_COORDINATOR
     var failedMetadata = txnMetadata1.prepareAddPartitions(Set[TopicPartition](new TopicPartition("topic2", 0)), time.milliseconds())
@@ -282,7 +380,7 @@ class TransactionStateManagerTest {
     prepareForTxnMessageAppend(Errors.NONE)
     transactionManager.removeTransactionsForTxnTopicPartition(partitionId, coordinatorEpoch)
     transactionManager.addLoadedTransactionsToCache(partitionId, coordinatorEpoch + 1, new Pool[String, TransactionMetadata]())
-    transactionManager.putTransactionStateIfNotExists(transactionalId1, txnMetadata1)
+    transactionManager.putTransactionStateIfNotExists(txnMetadata1)
     transactionManager.appendTransactionToLog(transactionalId1, coordinatorEpoch = 10, failedMetadata, assertCallback)
 
     prepareForTxnMessageAppend(Errors.NONE)
@@ -294,7 +392,7 @@ class TransactionStateManagerTest {
   @Test
   def testAppendFailToCoordinatorLoadingError(): Unit = {
     transactionManager.addLoadedTransactionsToCache(partitionId, coordinatorEpoch, new Pool[String, TransactionMetadata]())
-    transactionManager.putTransactionStateIfNotExists(transactionalId1, txnMetadata1)
+    transactionManager.putTransactionStateIfNotExists(txnMetadata1)
 
     expectedError = Errors.COORDINATOR_LOAD_IN_PROGRESS
     val failedMetadata = txnMetadata1.prepareAddPartitions(Set[TopicPartition](new TopicPartition("topic2", 0)), time.milliseconds())
@@ -306,9 +404,9 @@ class TransactionStateManagerTest {
   }
 
   @Test
-  def testAppendFailToUnknownError() {
+  def testAppendFailToUnknownError(): Unit = {
     transactionManager.addLoadedTransactionsToCache(partitionId, coordinatorEpoch, new Pool[String, TransactionMetadata]())
-    transactionManager.putTransactionStateIfNotExists(transactionalId1, txnMetadata1)
+    transactionManager.putTransactionStateIfNotExists(txnMetadata1)
 
     expectedError = Errors.UNKNOWN_SERVER_ERROR
     var failedMetadata = txnMetadata1.prepareAddPartitions(Set[TopicPartition](new TopicPartition("topic2", 0)), time.milliseconds())
@@ -326,9 +424,9 @@ class TransactionStateManagerTest {
   }
 
   @Test
-  def testPendingStateNotResetOnRetryAppend() {
+  def testPendingStateNotResetOnRetryAppend(): Unit = {
     transactionManager.addLoadedTransactionsToCache(partitionId, coordinatorEpoch, new Pool[String, TransactionMetadata]())
-    transactionManager.putTransactionStateIfNotExists(transactionalId1, txnMetadata1)
+    transactionManager.putTransactionStateIfNotExists(txnMetadata1)
 
     expectedError = Errors.COORDINATOR_NOT_AVAILABLE
     val failedMetadata = txnMetadata1.prepareAddPartitions(Set[TopicPartition](new TopicPartition("topic2", 0)), time.milliseconds())
@@ -344,7 +442,7 @@ class TransactionStateManagerTest {
     transactionManager.addLoadedTransactionsToCache(partitionId, 0, new Pool[String, TransactionMetadata]())
 
     // first insert the initial transaction metadata
-    transactionManager.putTransactionStateIfNotExists(transactionalId1, txnMetadata1)
+    transactionManager.putTransactionStateIfNotExists(txnMetadata1)
 
     prepareForTxnMessageAppend(Errors.NONE)
     expectedError = Errors.NOT_COORDINATOR
@@ -363,7 +461,7 @@ class TransactionStateManagerTest {
   def testAppendTransactionToLogWhilePendingStateChanged() = {
     // first insert the initial transaction metadata
     transactionManager.addLoadedTransactionsToCache(partitionId, coordinatorEpoch, new Pool[String, TransactionMetadata]())
-    transactionManager.putTransactionStateIfNotExists(transactionalId1, txnMetadata1)
+    transactionManager.putTransactionStateIfNotExists(txnMetadata1)
 
     prepareForTxnMessageAppend(Errors.NONE)
     expectedError = Errors.INVALID_PRODUCER_EPOCH
@@ -392,12 +490,12 @@ class TransactionStateManagerTest {
       transactionManager.addLoadedTransactionsToCache(partitionId, 0, new Pool[String, TransactionMetadata]())
     }
 
-    transactionManager.putTransactionStateIfNotExists("ongoing", transactionMetadata("ongoing", producerId = 0, state = Ongoing))
-    transactionManager.putTransactionStateIfNotExists("not-expiring", transactionMetadata("not-expiring", producerId = 1, state = Ongoing, txnTimeout = 10000))
-    transactionManager.putTransactionStateIfNotExists("prepare-commit", transactionMetadata("prepare-commit", producerId = 2, state = PrepareCommit))
-    transactionManager.putTransactionStateIfNotExists("prepare-abort", transactionMetadata("prepare-abort", producerId = 3, state = PrepareAbort))
-    transactionManager.putTransactionStateIfNotExists("complete-commit", transactionMetadata("complete-commit", producerId = 4, state = CompleteCommit))
-    transactionManager.putTransactionStateIfNotExists("complete-abort", transactionMetadata("complete-abort", producerId = 5, state = CompleteAbort))
+    transactionManager.putTransactionStateIfNotExists(transactionMetadata("ongoing", producerId = 0, state = Ongoing))
+    transactionManager.putTransactionStateIfNotExists(transactionMetadata("not-expiring", producerId = 1, state = Ongoing, txnTimeout = 10000))
+    transactionManager.putTransactionStateIfNotExists(transactionMetadata("prepare-commit", producerId = 2, state = PrepareCommit))
+    transactionManager.putTransactionStateIfNotExists(transactionMetadata("prepare-abort", producerId = 3, state = PrepareAbort))
+    transactionManager.putTransactionStateIfNotExists(transactionMetadata("complete-commit", producerId = 4, state = CompleteCommit))
+    transactionManager.putTransactionStateIfNotExists(transactionMetadata("complete-abort", producerId = 5, state = CompleteAbort))
 
     time.sleep(2000)
     val expiring = transactionManager.timedOutTransactions()
@@ -463,7 +561,62 @@ class TransactionStateManagerTest {
     verifyMetadataDoesExistAndIsUsable(transactionalId2)
   }
 
-  private def verifyMetadataDoesExistAndIsUsable(transactionalId: String) = {
+  @Test
+  def testSuccessfulReimmigration(): Unit = {
+    txnMetadata1.state = PrepareCommit
+    txnMetadata1.addPartitions(Set[TopicPartition](new TopicPartition("topic1", 0),
+      new TopicPartition("topic1", 1)))
+
+    txnRecords += new SimpleRecord(txnMessageKeyBytes1, TransactionLog.valueToBytes(txnMetadata1.prepareNoTransit()))
+    val startOffset = 0L
+    val records = MemoryRecords.withRecords(startOffset, CompressionType.NONE, txnRecords.toArray: _*)
+
+    prepareTxnLog(topicPartition, 0, records)
+
+    // immigrate partition at epoch 0
+    transactionManager.loadTransactionsForTxnTopicPartition(partitionId, coordinatorEpoch = 0, (_, _, _, _, _) => ())
+    assertEquals(0, transactionManager.loadingPartitions.size)
+
+    // Re-immigrate partition at epoch 1. This should be successful even though we didn't get to emigrate the partition.
+    prepareTxnLog(topicPartition, 0, records)
+    transactionManager.loadTransactionsForTxnTopicPartition(partitionId, coordinatorEpoch = 1, (_, _, _, _, _) => ())
+    assertEquals(0, transactionManager.loadingPartitions.size)
+    assertTrue(transactionManager.transactionMetadataCache.get(partitionId).isDefined)
+    assertEquals(1, transactionManager.transactionMetadataCache.get(partitionId).get.coordinatorEpoch)
+  }
+
+  @Test
+  def testLoadTransactionMetadataWithCorruptedLog(): Unit = {
+    // Simulate a case where startOffset < endOffset but log is empty. This could theoretically happen
+    // when all the records are expired and the active segment is truncated or when the partition
+    // is accidentally corrupted.
+    val startOffset = 0L
+    val endOffset = 10L
+
+    val logMock: Log = EasyMock.mock(classOf[Log])
+    EasyMock.expect(replicaManager.getLog(topicPartition)).andStubReturn(Some(logMock))
+    EasyMock.expect(logMock.logStartOffset).andStubReturn(startOffset)
+    EasyMock.expect(logMock.read(EasyMock.eq(startOffset),
+      maxLength = EasyMock.anyInt(),
+      isolation = EasyMock.eq(FetchLogEnd),
+      minOneMessage = EasyMock.eq(true))
+    ).andReturn(FetchDataInfo(LogOffsetMetadata(startOffset), MemoryRecords.EMPTY))
+    EasyMock.expect(replicaManager.getLogEndOffset(topicPartition)).andStubReturn(Some(endOffset))
+
+    EasyMock.replay(logMock)
+    EasyMock.replay(replicaManager)
+
+    transactionManager.loadTransactionsForTxnTopicPartition(partitionId, coordinatorEpoch = 0, (_, _, _, _, _) => ())
+
+    // let the time advance to trigger the background thread loading
+    scheduler.tick()
+
+    EasyMock.verify(logMock)
+    EasyMock.verify(replicaManager)
+    assertEquals(0, transactionManager.loadingPartitions.size)
+  }
+
+  private def verifyMetadataDoesExistAndIsUsable(transactionalId: String): Unit = {
     transactionManager.getTransactionState(transactionalId) match {
       case Left(errors) => fail("shouldn't have been any errors")
       case Right(None) => fail("metadata should have been removed")
@@ -497,7 +650,7 @@ class TransactionStateManagerTest {
         EasyMock.expect(replicaManager.appendRecords(EasyMock.anyLong(),
           EasyMock.eq((-1).toShort),
           EasyMock.eq(true),
-          EasyMock.eq(false),
+          EasyMock.eq(AppendOrigin.Coordinator),
           EasyMock.eq(recordsByPartition),
           EasyMock.capture(capturedArgument),
           EasyMock.anyObject().asInstanceOf[Option[ReentrantLock]],
@@ -518,10 +671,10 @@ class TransactionStateManagerTest {
 
     txnMetadata1.txnLastUpdateTimestamp = time.milliseconds() - txnConfig.transactionalIdExpirationMs
     txnMetadata1.state = txnState
-    transactionManager.putTransactionStateIfNotExists(transactionalId1, txnMetadata1)
+    transactionManager.putTransactionStateIfNotExists(txnMetadata1)
 
     txnMetadata2.txnLastUpdateTimestamp = time.milliseconds()
-    transactionManager.putTransactionStateIfNotExists(transactionalId2, txnMetadata2)
+    transactionManager.putTransactionStateIfNotExists(txnMetadata2)
 
     transactionManager.enableTransactionalIdExpiration()
     time.sleep(txnConfig.removeExpiredTransactionalIdsIntervalMs)
@@ -538,7 +691,7 @@ class TransactionStateManagerTest {
 
     txnRecords += new SimpleRecord(txnMessageKeyBytes1, TransactionLog.valueToBytes(txnMetadata1.prepareNoTransit()))
     val startOffset = 0L
-    val records = MemoryRecords.withRecords(startOffset, CompressionType.NONE, txnRecords: _*)
+    val records = MemoryRecords.withRecords(startOffset, CompressionType.NONE, txnRecords.toArray: _*)
 
     prepareTxnLog(topicPartition, 0, records)
 
@@ -584,9 +737,8 @@ class TransactionStateManagerTest {
     EasyMock.expect(logMock.logStartOffset).andStubReturn(startOffset)
     EasyMock.expect(logMock.read(EasyMock.eq(startOffset),
       maxLength = EasyMock.anyInt(),
-      maxOffset = EasyMock.eq(None),
-      minOneMessage = EasyMock.eq(true),
-      includeAbortedTxns = EasyMock.eq(false)))
+      isolation = EasyMock.eq(FetchLogEnd),
+      minOneMessage = EasyMock.eq(true)))
       .andReturn(FetchDataInfo(LogOffsetMetadata(startOffset), fileRecordsMock))
 
     EasyMock.expect(fileRecordsMock.sizeInBytes()).andStubReturn(records.sizeInBytes)
@@ -610,7 +762,7 @@ class TransactionStateManagerTest {
     EasyMock.expect(replicaManager.appendRecords(EasyMock.anyLong(),
       EasyMock.anyShort(),
       internalTopicsAllowed = EasyMock.eq(true),
-      isFromClient = EasyMock.eq(false),
+      origin = EasyMock.eq(AppendOrigin.Coordinator),
       EasyMock.anyObject().asInstanceOf[Map[TopicPartition, MemoryRecords]],
       EasyMock.capture(capturedArgument),
       EasyMock.anyObject().asInstanceOf[Option[ReentrantLock]],
@@ -628,4 +780,37 @@ class TransactionStateManagerTest {
 
     EasyMock.replay(replicaManager)
   }
+
+  @Test
+  def testPartitionLoadMetric(): Unit = {
+    val server = ManagementFactory.getPlatformMBeanServer
+    val mBeanName = "kafka.server:type=transaction-coordinator-metrics"
+    val reporter = new JmxReporter("kafka.server")
+    metrics.addReporter(reporter)
+
+    def partitionLoadTime(attribute: String): Double = {
+      server.getAttribute(new ObjectName(mBeanName), attribute).asInstanceOf[Double]
+    }
+
+    assertTrue(server.isRegistered(new ObjectName(mBeanName)))
+    assertEquals(Double.NaN, partitionLoadTime( "partition-load-time-max"), 0)
+    assertEquals(Double.NaN, partitionLoadTime("partition-load-time-avg"), 0)
+    assertTrue(reporter.containsMbean(mBeanName))
+
+    txnMetadata1.state = Ongoing
+    txnMetadata1.addPartitions(Set[TopicPartition](new TopicPartition("topic1", 1),
+      new TopicPartition("topic1", 1)))
+
+    txnRecords += new SimpleRecord(txnMessageKeyBytes1, TransactionLog.valueToBytes(txnMetadata1.prepareNoTransit()))
+
+    val startOffset = 15L
+    val records = MemoryRecords.withRecords(startOffset, CompressionType.NONE, txnRecords.toArray: _*)
+
+    prepareTxnLog(topicPartition, startOffset, records)
+    transactionManager.loadTransactionsForTxnTopicPartition(partitionId, 0, (_, _, _, _, _) => ())
+    scheduler.tick()
+
+    assertTrue(partitionLoadTime("partition-load-time-max") >= 0)
+    assertTrue(partitionLoadTime( "partition-load-time-avg") >= 0)
+  }
 }
diff --git a/core/src/test/scala/unit/kafka/integration/KafkaServerTestHarness.scala b/core/src/test/scala/unit/kafka/integration/KafkaServerTestHarness.scala
index 33001fd92a8a8..2b252503ed9e3 100755
--- a/core/src/test/scala/unit/kafka/integration/KafkaServerTestHarness.scala
+++ b/core/src/test/scala/unit/kafka/integration/KafkaServerTestHarness.scala
@@ -26,6 +26,7 @@ import kafka.zk.ZooKeeperTestHarness
 import org.apache.kafka.common.security.auth.{KafkaPrincipal, SecurityProtocol}
 import org.junit.{After, Before}
 
+import scala.collection.Seq
 import scala.collection.mutable.{ArrayBuffer, Buffer}
 import java.util.Properties
 
@@ -60,13 +61,13 @@ abstract class KafkaServerTestHarness extends ZooKeeperTestHarness {
    *
    * The default implementation of this method is a no-op.
    */
-  def configureSecurityBeforeServersStart() {}
+  def configureSecurityBeforeServersStart(): Unit = {}
 
   /**
    * Override this in case Tokens or security credentials needs to be created after `servers` are started.
    * The default implementation of this method is a no-op.
    */
-  def configureSecurityAfterServersStart() {}
+  def configureSecurityAfterServersStart(): Unit = {}
 
   def configs: Seq[KafkaConfig] = {
     if (instanceConfigs == null)
@@ -86,7 +87,7 @@ abstract class KafkaServerTestHarness extends ZooKeeperTestHarness {
   protected def brokerTime(brokerId: Int): Time = Time.SYSTEM
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     super.setUp()
 
     if (configs.isEmpty)
@@ -108,7 +109,7 @@ abstract class KafkaServerTestHarness extends ZooKeeperTestHarness {
   }
 
   @After
-  override def tearDown() {
+  override def tearDown(): Unit = {
     if (servers != null) {
       TestUtils.shutdownServers(servers)
     }
@@ -142,7 +143,7 @@ abstract class KafkaServerTestHarness extends ZooKeeperTestHarness {
     index
   }
 
-  def killBroker(index: Int) {
+  def killBroker(index: Int): Unit = {
     if(alive(index)) {
       servers(index).shutdown()
       servers(index).awaitShutdown()
@@ -153,7 +154,7 @@ abstract class KafkaServerTestHarness extends ZooKeeperTestHarness {
   /**
    * Restart any dead brokers
    */
-  def restartDeadBrokers() {
+  def restartDeadBrokers(): Unit = {
     for(i <- servers.indices if !alive(i)) {
       servers(i).startup()
       alive(i) = true
diff --git a/core/src/test/scala/unit/kafka/integration/MetricsDuringTopicCreationDeletionTest.scala b/core/src/test/scala/unit/kafka/integration/MetricsDuringTopicCreationDeletionTest.scala
index c1d310f93e835..c42d8eecfe071 100644
--- a/core/src/test/scala/unit/kafka/integration/MetricsDuringTopicCreationDeletionTest.scala
+++ b/core/src/test/scala/unit/kafka/integration/MetricsDuringTopicCreationDeletionTest.scala
@@ -21,6 +21,7 @@ import java.util.Properties
 import kafka.server.KafkaConfig
 import kafka.utils.{Logging, TestUtils}
 import scala.collection.JavaConverters.mapAsScalaMapConverter
+import org.scalatest.Assertions.fail
 
 import org.junit.{Before, Test}
 import com.yammer.metrics.Metrics
@@ -51,7 +52,7 @@ class MetricsDuringTopicCreationDeletionTest extends KafkaServerTestHarness with
     .map(KafkaConfig.fromProps(_, overridingProps))
 
   @Before
-  override def setUp {
+  override def setUp: Unit = {
     // Do some Metrics Registry cleanup by removing the metrics that this test checks.
     // This is a test workaround to the issue that prior harness runs may have left a populated registry.
     // see https://issues.apache.org/jira/browse/KAFKA-4605
@@ -67,7 +68,7 @@ class MetricsDuringTopicCreationDeletionTest extends KafkaServerTestHarness with
    * checking all metrics we care in a single test is faster though it would be more elegant to have 3 @Test methods
    */
   @Test
-  def testMetricsDuringTopicCreateDelete() {
+  def testMetricsDuringTopicCreateDelete(): Unit = {
 
     // For UnderReplicatedPartitions, because of https://issues.apache.org/jira/browse/KAFKA-4605
     // we can't access the metrics value of each server. So instead we directly invoke the method
@@ -80,14 +81,10 @@ class MetricsDuringTopicCreationDeletionTest extends KafkaServerTestHarness with
     @volatile var offlinePartitionsCount = offlinePartitionsCountGauge.value
     assert(offlinePartitionsCount == 0)
 
-    val preferredReplicaImbalanceCountGauge = getGauge("PreferredReplicaImbalanceCount")
-    @volatile var preferredReplicaImbalanceCount = preferredReplicaImbalanceCountGauge.value
-    assert(preferredReplicaImbalanceCount == 0)
-
     // Thread checking the metric continuously
     running = true
     val thread = new Thread(new Runnable {
-      def run() {
+      def run(): Unit = {
         while (running) {
           for ( s <- servers if running) {
             underReplicatedPartitionCount = s.replicaManager.underReplicatedPartitionCount
@@ -96,11 +93,6 @@ class MetricsDuringTopicCreationDeletionTest extends KafkaServerTestHarness with
             }
           }
 
-          preferredReplicaImbalanceCount = preferredReplicaImbalanceCountGauge.value
-          if (preferredReplicaImbalanceCount > 0) {
-             running = false
-          }
-
           offlinePartitionsCount = offlinePartitionsCountGauge.value
           if (offlinePartitionsCount > 0) {
              running = false
@@ -118,7 +110,6 @@ class MetricsDuringTopicCreationDeletionTest extends KafkaServerTestHarness with
     thread.join
 
     assert(offlinePartitionsCount==0, "OfflinePartitionCount not 0: "+ offlinePartitionsCount)
-    assert(preferredReplicaImbalanceCount==0, "PreferredReplicaImbalanceCount not 0: " + preferredReplicaImbalanceCount)
     assert(underReplicatedPartitionCount==0, "UnderReplicatedPartitionCount not 0: " + underReplicatedPartitionCount)
   }
 
@@ -130,7 +121,7 @@ class MetricsDuringTopicCreationDeletionTest extends KafkaServerTestHarness with
            ._2.asInstanceOf[Gauge[Int]]
   }
 
-  private def createDeleteTopics() {
+  private def createDeleteTopics(): Unit = {
     for (l <- 1 to createDeleteIterations if running) {
       // Create topics
       for (t <- topics if running) {
diff --git a/core/src/test/scala/unit/kafka/integration/MinIsrConfigTest.scala b/core/src/test/scala/unit/kafka/integration/MinIsrConfigTest.scala
index 1f01dc3be2b99..8ea5af789e8f6 100644
--- a/core/src/test/scala/unit/kafka/integration/MinIsrConfigTest.scala
+++ b/core/src/test/scala/unit/kafka/integration/MinIsrConfigTest.scala
@@ -30,7 +30,7 @@ class MinIsrConfigTest extends KafkaServerTestHarness {
   def generateConfigs = TestUtils.createBrokerConfigs(1, zkConnect).map(KafkaConfig.fromProps(_, overridingProps))
 
   @Test
-  def testDefaultKafkaConfig() {
+  def testDefaultKafkaConfig(): Unit = {
     assert(servers.head.getLogManager().initialDefaultConfig.minInSyncReplicas == 5)
   }
 
diff --git a/core/src/test/scala/unit/kafka/integration/UncleanLeaderElectionTest.scala b/core/src/test/scala/unit/kafka/integration/UncleanLeaderElectionTest.scala
index a13b570a01c2d..ac8b787a8331b 100755
--- a/core/src/test/scala/unit/kafka/integration/UncleanLeaderElectionTest.scala
+++ b/core/src/test/scala/unit/kafka/integration/UncleanLeaderElectionTest.scala
@@ -18,10 +18,11 @@
 package kafka.integration
 
 import org.apache.kafka.common.config.ConfigException
-import org.junit.{After, Before, Ignore, Test}
+import org.junit.{After, Before, Test}
 
 import scala.util.Random
 import scala.collection.JavaConverters._
+import scala.collection.Seq
 import org.apache.log4j.{Level, Logger}
 import java.util.Properties
 import java.util.concurrent.ExecutionException
@@ -35,8 +36,9 @@ import org.apache.kafka.common.errors.TimeoutException
 import org.apache.kafka.common.network.ListenerName
 import org.apache.kafka.common.security.auth.SecurityProtocol
 import org.apache.kafka.common.serialization.StringDeserializer
-import org.apache.kafka.clients.admin.{AdminClient, AdminClientConfig}
+import org.apache.kafka.clients.admin.{Admin, AdminClient, AdminClientConfig}
 import org.junit.Assert._
+import org.scalatest.Assertions.intercept
 
 class UncleanLeaderElectionTest extends ZooKeeperTestHarness {
   val brokerId1 = 0
@@ -60,7 +62,7 @@ class UncleanLeaderElectionTest extends ZooKeeperTestHarness {
   val networkProcessorLogger = Logger.getLogger(classOf[kafka.network.Processor])
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     super.setUp()
 
     configProps1 = createBrokerConfig(brokerId1, zkConnect)
@@ -78,7 +80,7 @@ class UncleanLeaderElectionTest extends ZooKeeperTestHarness {
   }
 
   @After
-  override def tearDown() {
+  override def tearDown(): Unit = {
     servers.foreach(server => shutdownServer(server))
     servers.foreach(server => CoreUtils.delete(server.config.logDirs))
 
@@ -89,7 +91,7 @@ class UncleanLeaderElectionTest extends ZooKeeperTestHarness {
     super.tearDown()
   }
 
-  private def startBrokers(cluster: Seq[Properties]) {
+  private def startBrokers(cluster: Seq[Properties]): Unit = {
     for (props <- cluster) {
       val config = KafkaConfig.fromProps(props)
       val server = createServer(config)
@@ -112,7 +114,6 @@ class UncleanLeaderElectionTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  @Ignore // Should be re-enabled after KAFKA-3096 is fixed
   def testUncleanLeaderElectionDisabled(): Unit = {
     // unclean leader election is disabled by default
     startBrokers(Seq(configProps1, configProps2))
@@ -139,8 +140,7 @@ class UncleanLeaderElectionTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  @Ignore // Should be re-enabled after KAFKA-3096 is fixed
-  def testCleanLeaderElectionDisabledByTopicOverride(): Unit = {
+  def testUncleanLeaderElectionDisabledByTopicOverride(): Unit = {
     // enable unclean leader election globally, but disable for our specific test topic
     configProps1.put("unclean.leader.election.enable", "true")
     configProps2.put("unclean.leader.election.enable", "true")
@@ -220,16 +220,16 @@ class UncleanLeaderElectionTest extends ZooKeeperTestHarness {
     assertEquals(List("first"), consumeAllMessages(topic, 1))
 
     // shutdown follower server
-    servers.filter(server => server.config.brokerId == followerId).map(server => shutdownServer(server))
+    servers.filter(server => server.config.brokerId == followerId).foreach(server => shutdownServer(server))
 
     produceMessage(servers, topic, "second")
     assertEquals(List("first", "second"), consumeAllMessages(topic, 2))
 
     //remove any previous unclean election metric
-    servers.map(server => server.kafkaController.controllerContext.stats.removeMetric("UncleanLeaderElectionsPerSec"))
+    servers.foreach(server => server.kafkaController.controllerContext.stats.removeMetric("UncleanLeaderElectionsPerSec"))
 
     // shutdown leader and then restart follower
-    servers.filter(server => server.config.brokerId == leaderId).map(server => shutdownServer(server))
+    servers.filter(server => server.config.brokerId == leaderId).foreach(server => shutdownServer(server))
     val followerServer = servers.find(_.config.brokerId == followerId).get
     followerServer.startup()
 
@@ -248,13 +248,17 @@ class UncleanLeaderElectionTest extends ZooKeeperTestHarness {
     assertEquals(List.empty[String], consumeAllMessages(topic, 0))
 
     // restart leader temporarily to send a successfully replicated message
-    servers.filter(server => server.config.brokerId == leaderId).map(server => server.startup())
+    servers.filter(server => server.config.brokerId == leaderId).foreach(server => server.startup())
     waitUntilLeaderIsElectedOrChanged(zkClient, topic, partitionId, newLeaderOpt = Some(leaderId))
 
     produceMessage(servers, topic, "third")
-    waitUntilMetadataIsPropagated(servers, topic, partitionId)
-    servers.filter(server => server.config.brokerId == leaderId).map(server => shutdownServer(server))
+    //make sure follower server joins the ISR
+    TestUtils.waitUntilTrue(() => {
+      val partitionInfoOpt = followerServer.metadataCache.getPartitionInfo(topic, partitionId)
+      partitionInfoOpt.isDefined && partitionInfoOpt.get.isr.contains(followerId)
+    }, "Inconsistent metadata after first server startup")
 
+    servers.filter(server => server.config.brokerId == leaderId).foreach(server => shutdownServer(server))
     // verify clean leader transition to ISR follower
     waitUntilLeaderIsElectedOrChanged(zkClient, topic, partitionId, newLeaderOpt = Some(followerId))
 
@@ -343,7 +347,7 @@ class UncleanLeaderElectionTest extends ZooKeeperTestHarness {
     assertEquals(List("first", "third"), consumeAllMessages(topic, 2))
   }
 
-  private def createAdminClient(): AdminClient = {
+  private def createAdminClient(): Admin = {
     val config = new Properties
     val bootstrapServers = TestUtils.bootstrapServers(servers, new ListenerName("PLAINTEXT"))
     config.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers)
diff --git a/core/src/test/scala/unit/kafka/log/AbstractLogCleanerIntegrationTest.scala b/core/src/test/scala/unit/kafka/log/AbstractLogCleanerIntegrationTest.scala
index e77833636e339..f2991afbc3368 100644
--- a/core/src/test/scala/unit/kafka/log/AbstractLogCleanerIntegrationTest.scala
+++ b/core/src/test/scala/unit/kafka/log/AbstractLogCleanerIntegrationTest.scala
@@ -43,9 +43,10 @@ abstract class AbstractLogCleanerIntegrationTest {
   private val logs = ListBuffer.empty[Log]
   private val defaultMaxMessageSize = 128
   private val defaultMinCleanableDirtyRatio = 0.0F
-  private val defaultCompactionLag = 0L
+  private val defaultMinCompactionLagMS = 0L
   private val defaultDeleteDelay = 1000
   private val defaultSegmentSize = 2048
+  private val defaultMaxCompactionLagMs = Long.MaxValue
 
   def time: MockTime
 
@@ -61,9 +62,10 @@ abstract class AbstractLogCleanerIntegrationTest {
   def logConfigProperties(propertyOverrides: Properties = new Properties(),
                           maxMessageSize: Int,
                           minCleanableDirtyRatio: Float = defaultMinCleanableDirtyRatio,
-                          compactionLag: Long = defaultCompactionLag,
+                          minCompactionLagMs: Long = defaultMinCompactionLagMS,
                           deleteDelay: Int = defaultDeleteDelay,
-                          segmentSize: Int = defaultSegmentSize): Properties = {
+                          segmentSize: Int = defaultSegmentSize,
+                          maxCompactionLagMs: Long = defaultMaxCompactionLagMs): Properties = {
     val props = new Properties()
     props.put(LogConfig.MaxMessageBytesProp, maxMessageSize: java.lang.Integer)
     props.put(LogConfig.SegmentBytesProp, segmentSize: java.lang.Integer)
@@ -72,7 +74,8 @@ abstract class AbstractLogCleanerIntegrationTest {
     props.put(LogConfig.CleanupPolicyProp, LogConfig.Compact)
     props.put(LogConfig.MinCleanableDirtyRatioProp, minCleanableDirtyRatio: java.lang.Float)
     props.put(LogConfig.MessageTimestampDifferenceMaxMsProp, Long.MaxValue.toString)
-    props.put(LogConfig.MinCompactionLagMsProp, compactionLag: java.lang.Long)
+    props.put(LogConfig.MinCompactionLagMsProp, minCompactionLagMs: java.lang.Long)
+    props.put(LogConfig.MaxCompactionLagMsProp, maxCompactionLagMs: java.lang.Long)
     props ++= propertyOverrides
     props
   }
@@ -82,9 +85,10 @@ abstract class AbstractLogCleanerIntegrationTest {
                   numThreads: Int = 1,
                   backOffMs: Long = 15000L,
                   maxMessageSize: Int = defaultMaxMessageSize,
-                  compactionLag: Long = defaultCompactionLag,
+                  minCompactionLagMs: Long = defaultMinCompactionLagMS,
                   deleteDelay: Int = defaultDeleteDelay,
                   segmentSize: Int = defaultSegmentSize,
+                  maxCompactionLagMs: Long = defaultMaxCompactionLagMs,
                   cleanerIoBufferSize: Option[Int] = None,
                   propertyOverrides: Properties = new Properties()): LogCleaner = {
 
@@ -96,9 +100,10 @@ abstract class AbstractLogCleanerIntegrationTest {
       val logConfig = LogConfig(logConfigProperties(propertyOverrides,
         maxMessageSize = maxMessageSize,
         minCleanableDirtyRatio = minCleanableDirtyRatio,
-        compactionLag = compactionLag,
+        minCompactionLagMs = minCompactionLagMs,
         deleteDelay = deleteDelay,
-        segmentSize = segmentSize))
+        segmentSize = segmentSize,
+        maxCompactionLagMs = maxCompactionLagMs))
       val log = Log(dir,
         logConfig,
         logStartOffset = 0L,
@@ -122,6 +127,7 @@ abstract class AbstractLogCleanerIntegrationTest {
       logDirs = Array(logDir),
       logs = logMap,
       logDirFailureChannel = new LogDirFailureChannel(1),
+      retentionCheckMs = 0L,
       time = time)
   }
 
diff --git a/core/src/test/scala/unit/kafka/log/BrokerCompressionTest.scala b/core/src/test/scala/unit/kafka/log/BrokerCompressionTest.scala
index 8372387d87b7a..717697a7c69d3 100755
--- a/core/src/test/scala/unit/kafka/log/BrokerCompressionTest.scala
+++ b/core/src/test/scala/unit/kafka/log/BrokerCompressionTest.scala
@@ -19,7 +19,6 @@ package kafka.log
 
 import kafka.utils._
 import kafka.message._
-import org.scalatest.junit.JUnitSuite
 import org.junit._
 import org.junit.Assert._
 import org.junit.runner.RunWith
@@ -29,12 +28,12 @@ import org.apache.kafka.common.record.{CompressionType, MemoryRecords, RecordBat
 import org.apache.kafka.common.utils.Utils
 import java.util.{Collection, Properties}
 
-import kafka.server.{BrokerTopicStats, LogDirFailureChannel}
+import kafka.server.{BrokerTopicStats, FetchLogEnd, LogDirFailureChannel}
 
 import scala.collection.JavaConverters._
 
 @RunWith(value = classOf[Parameterized])
-class BrokerCompressionTest(messageCompression: String, brokerCompression: String) extends JUnitSuite {
+class BrokerCompressionTest(messageCompression: String, brokerCompression: String) {
 
   val tmpDir = TestUtils.tempDir()
   val logDir = TestUtils.randomPartitionLogDir(tmpDir)
@@ -42,7 +41,7 @@ class BrokerCompressionTest(messageCompression: String, brokerCompression: Strin
   val logConfig = LogConfig()
 
   @After
-  def tearDown() {
+  def tearDown(): Unit = {
     Utils.delete(tmpDir)
   }
 
@@ -50,7 +49,7 @@ class BrokerCompressionTest(messageCompression: String, brokerCompression: Strin
    * Test broker-side compression configuration
    */
   @Test
-  def testBrokerSideCompression() {
+  def testBrokerSideCompression(): Unit = {
     val messageCompressionCode = CompressionCodec.getCompressionCodec(messageCompression)
     val logProps = new Properties()
     logProps.put(LogConfig.CompressionTypeProp, brokerCompression)
@@ -65,8 +64,10 @@ class BrokerCompressionTest(messageCompression: String, brokerCompression: Strin
           new SimpleRecord("hello".getBytes), new SimpleRecord("there".getBytes)), leaderEpoch = 0)
 
     def readBatch(offset: Int): RecordBatch = {
-      val fetchInfo = log.read(offset, 4096, maxOffset = None,
-        includeAbortedTxns = false, minOneMessage = true)
+      val fetchInfo = log.read(offset,
+        maxLength = 4096,
+        isolation = FetchLogEnd,
+        minOneMessage = true)
       fetchInfo.records.batches.iterator.next()
     }
 
@@ -85,6 +86,7 @@ object BrokerCompressionTest {
   def parameters: Collection[Array[String]] = {
     (for (brokerCompression <- BrokerCompressionCodec.brokerCompressionOptions;
          messageCompression <- CompressionType.values
+      if messageCompression != CompressionType.PASSTHROUGH
     ) yield Array(messageCompression.name, brokerCompression)).asJava
   }
 }
diff --git a/core/src/test/scala/unit/kafka/log/LogCleanerIntegrationTest.scala b/core/src/test/scala/unit/kafka/log/LogCleanerIntegrationTest.scala
index bfee811167c8f..7aea6c1fcdd80 100644
--- a/core/src/test/scala/unit/kafka/log/LogCleanerIntegrationTest.scala
+++ b/core/src/test/scala/unit/kafka/log/LogCleanerIntegrationTest.scala
@@ -20,28 +20,35 @@ package kafka.log
 import java.io.PrintWriter
 
 import com.yammer.metrics.Metrics
-import com.yammer.metrics.core.Gauge
+import com.yammer.metrics.core.{Gauge, MetricName}
+import kafka.metrics.KafkaMetricsGroup
 import kafka.utils.{MockTime, TestUtils}
 import org.apache.kafka.common.TopicPartition
-import org.apache.kafka.test.TestUtils.DEFAULT_MAX_WAIT_MS
 import org.apache.kafka.common.record.{CompressionType, RecordBatch}
-import org.junit.Assert.{assertFalse, assertTrue, fail}
-import org.junit.Test
+import org.apache.kafka.test.TestUtils.DEFAULT_MAX_WAIT_MS
+import org.junit.Assert._
+import org.junit.{After, Test}
 
 import scala.collection.JavaConverters.mapAsScalaMapConverter
+import scala.collection.{Iterable, JavaConverters, Seq}
 
 /**
   * This is an integration test that tests the fully integrated log cleaner
   */
-class LogCleanerIntegrationTest extends AbstractLogCleanerIntegrationTest {
+class LogCleanerIntegrationTest extends AbstractLogCleanerIntegrationTest with KafkaMetricsGroup {
 
   val codec: CompressionType = CompressionType.LZ4
 
   val time = new MockTime()
   val topicPartitions = Array(new TopicPartition("log", 0), new TopicPartition("log", 1), new TopicPartition("log", 2))
 
+  @After
+  def cleanup(): Unit = {
+    TestUtils.clearYammerMetrics()
+  }
+
   @Test(timeout = DEFAULT_MAX_WAIT_MS)
-  def testMarksPartitionsAsOfflineAndPopulatesUncleanableMetrics() {
+  def testMarksPartitionsAsOfflineAndPopulatesUncleanableMetrics(): Unit = {
     val largeMessageKey = 20
     val (_, largeMessageSet) = createLargeSingleMessageSet(largeMessageKey, RecordBatch.CURRENT_MAGIC_VALUE)
     val maxMessageSize = largeMessageSet.sizeInBytes
@@ -59,18 +66,6 @@ class LogCleanerIntegrationTest extends AbstractLogCleanerIntegrationTest {
       writeDups(numKeys = 20, numDups = 3, log = log, codec = codec)
     }
 
-    def getGauge[T](metricName: String, metricScope: String): Gauge[T] = {
-      Metrics.defaultRegistry.allMetrics.asScala
-        .filterKeys(k => {
-          k.getName.endsWith(metricName) && k.getScope.endsWith(metricScope)
-        })
-        .headOption
-        .getOrElse { fail(s"Unable to find metric $metricName") }
-        .asInstanceOf[(Any, Gauge[Any])]
-        ._2
-        .asInstanceOf[Gauge[T]]
-    }
-
     breakPartitionLog(topicPartitions(0))
     breakPartitionLog(topicPartitions(1))
 
@@ -83,8 +78,8 @@ class LogCleanerIntegrationTest extends AbstractLogCleanerIntegrationTest {
     val uncleanableBytesGauge = getGauge[Long]("uncleanable-bytes", uncleanableDirectory)
 
     TestUtils.waitUntilTrue(() => uncleanablePartitionsCountGauge.value() == 2, "There should be 2 uncleanable partitions", 2000L)
-    val expectedTotalUncleanableBytes = LogCleaner.calculateCleanableBytes(log, 0, log.logSegments.last.baseOffset)._2 +
-      LogCleaner.calculateCleanableBytes(log2, 0, log2.logSegments.last.baseOffset)._2
+    val expectedTotalUncleanableBytes = LogCleanerManager.calculateCleanableBytes(log, 0, log.logSegments.last.baseOffset)._2 +
+      LogCleanerManager.calculateCleanableBytes(log2, 0, log2.logSegments.last.baseOffset)._2
     TestUtils.waitUntilTrue(() => uncleanableBytesGauge.value() == expectedTotalUncleanableBytes,
       s"There should be $expectedTotalUncleanableBytes uncleanable bytes", 1000L)
 
@@ -93,4 +88,131 @@ class LogCleanerIntegrationTest extends AbstractLogCleanerIntegrationTest {
     assertTrue(uncleanablePartitions.contains(topicPartitions(1)))
     assertFalse(uncleanablePartitions.contains(topicPartitions(2)))
   }
+
+  private def getGauge[T](filter: MetricName => Boolean): Gauge[T] = {
+    Metrics.defaultRegistry.allMetrics.asScala
+      .filterKeys(filter(_))
+      .headOption
+      .getOrElse { fail(s"Unable to find metric") }
+      .asInstanceOf[(Any, Gauge[Any])]
+      ._2
+      .asInstanceOf[Gauge[T]]
+  }
+
+  private def getGauge[T](metricName: String): Gauge[T] = {
+    getGauge(mName => mName.getName.endsWith(metricName) && mName.getScope == null)
+  }
+
+  private def getGauge[T](metricName: String, metricScope: String): Gauge[T] = {
+    getGauge(k => k.getName.endsWith(metricName) && k.getScope.endsWith(metricScope))
+  }
+
+  @Test
+  def testMaxLogCompactionLag(): Unit = {
+    val msPerHour = 60 * 60 * 1000
+
+    val minCompactionLagMs = 1 * msPerHour
+    val maxCompactionLagMs = 6 * msPerHour
+
+    val cleanerBackOffMs = 200L
+    val segmentSize = 512
+    val topicPartitions = Array(new TopicPartition("log", 0), new TopicPartition("log", 1), new TopicPartition("log", 2))
+    val minCleanableDirtyRatio = 1.0F
+
+    cleaner = makeCleaner(partitions = topicPartitions,
+      backOffMs = cleanerBackOffMs,
+      minCompactionLagMs = minCompactionLagMs,
+      segmentSize = segmentSize,
+      maxCompactionLagMs= maxCompactionLagMs,
+      minCleanableDirtyRatio = minCleanableDirtyRatio)
+    val log = cleaner.logs.get(topicPartitions(0))
+
+    val T0 = time.milliseconds
+    writeKeyDups(numKeys = 100, numDups = 3, log, CompressionType.NONE, timestamp = T0, startValue = 0, step = 1)
+
+    val startSizeBlock0 = log.size
+
+    val activeSegAtT0 = log.activeSegment
+
+    cleaner.startup()
+
+    // advance to a time still less than maxCompactionLagMs from start
+    time.sleep(maxCompactionLagMs/2)
+    Thread.sleep(5 * cleanerBackOffMs) // give cleaning thread a chance to _not_ clean
+    assertEquals("There should be no cleaning until the max compaction lag has passed", startSizeBlock0, log.size)
+
+    // advance to time a bit more than one maxCompactionLagMs from start
+    time.sleep(maxCompactionLagMs/2 + 1)
+    val T1 = time.milliseconds
+
+    // write the second block of data: all zero keys
+    val appends1 = writeKeyDups(numKeys = 100, numDups = 1, log, CompressionType.NONE, timestamp = T1, startValue = 0, step = 0)
+
+    // roll the active segment
+    log.roll()
+    val activeSegAtT1 = log.activeSegment
+    val firstBlockCleanableSegmentOffset = activeSegAtT0.baseOffset
+
+    // the first block should get cleaned
+    cleaner.awaitCleaned(new TopicPartition("log", 0), firstBlockCleanableSegmentOffset)
+
+    val read1 = readFromLog(log)
+    val lastCleaned = cleaner.cleanerManager.allCleanerCheckpoints(new TopicPartition("log", 0))
+    assertTrue(s"log cleaner should have processed at least to offset $firstBlockCleanableSegmentOffset, " +
+      s"but lastCleaned=$lastCleaned", lastCleaned >= firstBlockCleanableSegmentOffset)
+
+    //minCleanableDirtyRatio  will prevent second block of data from compacting
+    assertNotEquals(s"log should still contain non-zero keys", appends1, read1)
+
+    time.sleep(maxCompactionLagMs + 1)
+    // the second block should get cleaned. only zero keys left
+    cleaner.awaitCleaned(new TopicPartition("log", 0), activeSegAtT1.baseOffset)
+
+    val read2 = readFromLog(log)
+
+    assertEquals(s"log should only contains zero keys now", appends1, read2)
+
+    val lastCleaned2 = cleaner.cleanerManager.allCleanerCheckpoints(new TopicPartition("log", 0))
+    val secondBlockCleanableSegmentOffset = activeSegAtT1.baseOffset
+    assertTrue(s"log cleaner should have processed at least to offset $secondBlockCleanableSegmentOffset, " +
+      s"but lastCleaned=$lastCleaned2", lastCleaned2 >= secondBlockCleanableSegmentOffset)
+  }
+
+  private def readFromLog(log: Log): Iterable[(Int, Int)] = {
+    import JavaConverters._
+    for (segment <- log.logSegments; record <- segment.log.records.asScala) yield {
+      val key = TestUtils.readString(record.key).toInt
+      val value = TestUtils.readString(record.value).toInt
+      key -> value
+    }
+  }
+
+  private def writeKeyDups(numKeys: Int, numDups: Int, log: Log, codec: CompressionType, timestamp: Long, startValue: Int, step: Int): Seq[(Int, Int)] = {
+    var valCounter = startValue
+    for (_ <- 0 until numDups; key <- 0 until numKeys) yield {
+      val curValue = valCounter
+      log.appendAsLeader(TestUtils.singletonRecords(value = curValue.toString.getBytes, codec = codec,
+        key = key.toString.getBytes, timestamp = timestamp), leaderEpoch = 0)
+      valCounter += step
+      (key, curValue)
+    }
+  }
+
+  @Test
+  def testIsThreadFailed(): Unit = {
+    val metricName = "DeadThreadCount"
+    cleaner = makeCleaner(partitions = topicPartitions, maxMessageSize = 100000, backOffMs = 100)
+    cleaner.startup()
+    assertEquals(0, cleaner.deadThreadCount)
+    // we simulate the unexpected error with an interrupt
+    cleaner.cleaners.foreach(_.interrupt())
+    // wait until interruption is propagated to all the threads
+    TestUtils.waitUntilTrue(
+      () => cleaner.cleaners.foldLeft(true)((result, thread) => {
+        thread.isThreadFailed && result
+      }), "Threads didn't terminate unexpectedly"
+    )
+    assertEquals(cleaner.cleaners.size, getGauge[Int](metricName).value())
+    assertEquals(cleaner.cleaners.size, cleaner.deadThreadCount)
+  }
 }
diff --git a/core/src/test/scala/unit/kafka/log/LogCleanerLagIntegrationTest.scala b/core/src/test/scala/unit/kafka/log/LogCleanerLagIntegrationTest.scala
index 6e8c9b9d33679..ffdf1ad0b63d3 100644
--- a/core/src/test/scala/unit/kafka/log/LogCleanerLagIntegrationTest.scala
+++ b/core/src/test/scala/unit/kafka/log/LogCleanerLagIntegrationTest.scala
@@ -35,8 +35,8 @@ import scala.collection._
 class LogCleanerLagIntegrationTest(compressionCodecName: String) extends AbstractLogCleanerIntegrationTest with Logging {
   val msPerHour = 60 * 60 * 1000
 
-  val compactionLag = 1 * msPerHour
-  assertTrue("compactionLag must be divisible by 2 for this test", compactionLag % 2 == 0)
+  val minCompactionLag = 1 * msPerHour
+  assertTrue("compactionLag must be divisible by 2 for this test", minCompactionLag % 2 == 0)
 
   val time = new MockTime(1400000000000L, 1000L)  // Tue May 13 16:53:20 UTC 2014 for `currentTimeMs`
   val cleanerBackOffMs = 200L
@@ -50,7 +50,7 @@ class LogCleanerLagIntegrationTest(compressionCodecName: String) extends Abstrac
   def cleanerTest(): Unit = {
     cleaner = makeCleaner(partitions = topicPartitions,
       backOffMs = cleanerBackOffMs,
-      compactionLag = compactionLag,
+      minCompactionLagMs = minCompactionLag,
       segmentSize = segmentSize)
     val log = cleaner.logs.get(topicPartitions(0))
 
@@ -69,13 +69,13 @@ class LogCleanerLagIntegrationTest(compressionCodecName: String) extends Abstrac
 
     // T0 < t < T1
     // advance to a time still less than one compaction lag from start
-    time.sleep(compactionLag/2)
+    time.sleep(minCompactionLag/2)
     Thread.sleep(5 * cleanerBackOffMs) // give cleaning thread a chance to _not_ clean
     assertEquals("There should be no cleaning until the compaction lag has passed", startSizeBlock0, log.size)
 
     // t = T1 > T0 + compactionLag
     // advance to time a bit more than one compaction lag from start
-    time.sleep(compactionLag/2 + 1)
+    time.sleep(minCompactionLag/2 + 1)
     val T1 = time.milliseconds
 
     // write another block of data
@@ -128,7 +128,7 @@ object LogCleanerLagIntegrationTest {
   @Parameters
   def parameters: java.util.Collection[Array[String]] = {
     val list = new java.util.ArrayList[Array[String]]()
-    for (codec <- CompressionType.values)
+    for (codec <- CompressionType.values.filter(_!=CompressionType.PASSTHROUGH))
       list.add(Array(codec.name))
     list
   }
diff --git a/core/src/test/scala/unit/kafka/log/LogCleanerManagerTest.scala b/core/src/test/scala/unit/kafka/log/LogCleanerManagerTest.scala
index 8ca26a86d16fe..0630fcc1d2fb2 100644
--- a/core/src/test/scala/unit/kafka/log/LogCleanerManagerTest.scala
+++ b/core/src/test/scala/unit/kafka/log/LogCleanerManagerTest.scala
@@ -27,14 +27,14 @@ import org.apache.kafka.common.record._
 import org.apache.kafka.common.utils.Utils
 import org.junit.Assert._
 import org.junit.{After, Test}
-import org.scalatest.junit.JUnitSuite
+import org.scalatest.Assertions.intercept
 
 import scala.collection.mutable
 
 /**
   * Unit tests for the log cleaning logic
   */
-class LogCleanerManagerTest extends JUnitSuite with Logging {
+class LogCleanerManagerTest extends Logging {
 
   val tmpDir = TestUtils.tempDir()
   val logDir = TestUtils.randomPartitionLogDir(tmpDir)
@@ -50,10 +50,15 @@ class LogCleanerManagerTest extends JUnitSuite with Logging {
 
   class LogCleanerManagerMock(logDirs: Seq[File],
                               logs: Pool[TopicPartition, Log],
-                              logDirFailureChannel: LogDirFailureChannel) extends LogCleanerManager(logDirs, logs, logDirFailureChannel) {
+                              logDirFailureChannel: LogDirFailureChannel) extends LogCleanerManager(logDirs, logs, logDirFailureChannel, 0L) {
     override def allCleanerCheckpoints: Map[TopicPartition, Long] = {
       cleanerCheckpoints.toMap
     }
+
+    override def updateCheckpoints(dataDir: File, update: Option[(TopicPartition,Long)]): Unit = {
+      val (tp, offset) = update.getOrElse(throw new IllegalArgumentException("update=None argument not yet handled"))
+      cleanerCheckpoints.put(tp, offset)
+    }
   }
 
   @After
@@ -61,107 +66,206 @@ class LogCleanerManagerTest extends JUnitSuite with Logging {
     Utils.delete(tmpDir)
   }
 
+  private def setupIncreasinglyFilthyLogs(partitions: Seq[TopicPartition],
+                                          startNumBatches: Int,
+                                          batchIncrement: Int): Pool[TopicPartition, Log] = {
+    val logs = new Pool[TopicPartition, Log]()
+    var numBatches = startNumBatches
+
+    for (tp <- partitions) {
+      val log = createLog(2048, LogConfig.Compact, topicPartition = tp)
+      logs.put(tp, log)
+
+      writeRecords(log, numBatches = numBatches, recordsPerBatch = 1, batchesPerSegment = 5)
+      numBatches += batchIncrement
+    }
+    logs
+  }
+
+  @Test
+  def testGrabFilthiestCompactedLogThrowsException(): Unit = {
+    val tp = new TopicPartition("A", 1)
+    val logSegmentSize = TestUtils.singletonRecords("test".getBytes).sizeInBytes * 10
+    val logSegmentsCount = 2
+    val tpDir = new File(logDir, "A-1")
+
+    // the exception should be catched and the partition that caused it marked as uncleanable
+    class LogMock(dir: File, config: LogConfig) extends Log(dir, config, 0L, 0L,
+      time.scheduler, new BrokerTopicStats, time, 60 * 60 * 1000, LogManager.ProducerIdExpirationCheckIntervalMs,
+      topicPartition, new ProducerStateManager(tp, tpDir, 60 * 60 * 1000), new LogDirFailureChannel(10)) {
+
+      // Throw an error in getFirstBatchTimestampForSegments since it is called in grabFilthiestLog()
+      override def getFirstBatchTimestampForSegments(segments: Iterable[LogSegment]): Iterable[Long] =
+        throw new IllegalStateException("Error!")
+    }
+
+    val log: Log = new LogMock(tpDir, createLowRetentionLogConfig(logSegmentSize, LogConfig.Compact))
+    writeRecords(log = log,
+      numBatches = logSegmentsCount * 2,
+      recordsPerBatch = 10,
+      batchesPerSegment = 2
+    )
+
+    val logsPool = new Pool[TopicPartition, Log]()
+    logsPool.put(tp, log)
+    val cleanerManager = createCleanerManagerMock(logsPool)
+    cleanerCheckpoints.put(tp, 1)
+
+    val thrownException = intercept[LogCleaningException] {
+      cleanerManager.grabFilthiestCompactedLog(time).get
+    }
+    assertEquals(log, thrownException.log)
+    assertTrue(thrownException.getCause.isInstanceOf[IllegalStateException])
+  }
+
   @Test
   def testGrabFilthiestCompactedLogReturnsLogWithDirtiestRatio(): Unit = {
-    val records = TestUtils.singletonRecords("test".getBytes)
-    val log1: Log = createLog(records.sizeInBytes * 5, LogConfig.Compact, 1)
-    val log2: Log = createLog(records.sizeInBytes * 10, LogConfig.Compact, 2)
-    val log3: Log = createLog(records.sizeInBytes * 15, LogConfig.Compact, 3)
+    val tp0 = new TopicPartition("wishing-well", 0)
+    val tp1 = new TopicPartition("wishing-well", 1)
+    val tp2 = new TopicPartition("wishing-well", 2)
+    val partitions = Seq(tp0, tp1, tp2)
 
-    val logs = new Pool[TopicPartition, Log]()
-    val tp1 = new TopicPartition("wishing well", 0) // active segment starts at 0
-    logs.put(tp1, log1)
-    val tp2 = new TopicPartition("wishing well", 1) // active segment starts at 10
-    logs.put(tp2, log2)
-    val tp3 = new TopicPartition("wishing well", 2) // // active segment starts at 20
-    logs.put(tp3, log3)
-    val cleanerManager: LogCleanerManagerMock = createCleanerManager(logs, toMock = true).asInstanceOf[LogCleanerManagerMock]
-    cleanerCheckpoints.put(tp1, 0) // all clean
-    cleanerCheckpoints.put(tp2, 1) // dirtiest - 9 dirty messages
-    cleanerCheckpoints.put(tp3, 15) // 5 dirty messages
+    // setup logs with cleanable range: [20, 20], [20, 25], [20, 30]
+    val logs = setupIncreasinglyFilthyLogs(partitions, startNumBatches = 20, batchIncrement = 5)
+    val cleanerManager = createCleanerManagerMock(logs)
+    partitions.foreach(partition => cleanerCheckpoints.put(partition, 20))
 
     val filthiestLog: LogToClean = cleanerManager.grabFilthiestCompactedLog(time).get
-
-    assertEquals(log2, filthiestLog.log)
     assertEquals(tp2, filthiestLog.topicPartition)
+    assertEquals(tp2, filthiestLog.log.topicPartition)
   }
 
   @Test
   def testGrabFilthiestCompactedLogIgnoresUncleanablePartitions(): Unit = {
-    val records = TestUtils.singletonRecords("test".getBytes)
-    val log1: Log = createLog(records.sizeInBytes * 5, LogConfig.Compact, 1)
-    val log2: Log = createLog(records.sizeInBytes * 10, LogConfig.Compact, 2)
-    val log3: Log = createLog(records.sizeInBytes * 15, LogConfig.Compact, 3)
+    val tp0 = new TopicPartition("wishing-well", 0)
+    val tp1 = new TopicPartition("wishing-well", 1)
+    val tp2 = new TopicPartition("wishing-well", 2)
+    val partitions = Seq(tp0, tp1, tp2)
 
-    val logs = new Pool[TopicPartition, Log]()
-    val tp1 = new TopicPartition("wishing well", 0) // active segment starts at 0
-    logs.put(tp1, log1)
-    val tp2 = new TopicPartition("wishing well", 1) // active segment starts at 10
-    logs.put(tp2, log2)
-    val tp3 = new TopicPartition("wishing well", 2) // // active segment starts at 20
-    logs.put(tp3, log3)
-    val cleanerManager: LogCleanerManagerMock = createCleanerManager(logs, toMock = true).asInstanceOf[LogCleanerManagerMock]
-    cleanerCheckpoints.put(tp1, 0) // all clean
-    cleanerCheckpoints.put(tp2, 1) // dirtiest - 9 dirty messages
-    cleanerCheckpoints.put(tp3, 15) // 5 dirty messages
-    cleanerManager.markPartitionUncleanable(log2.dir.getParent, tp2)
+    // setup logs with cleanable range: [20, 20], [20, 25], [20, 30]
+    val logs = setupIncreasinglyFilthyLogs(partitions, startNumBatches = 20, batchIncrement = 5)
+    val cleanerManager = createCleanerManagerMock(logs)
+    partitions.foreach(partition => cleanerCheckpoints.put(partition, 20))
 
-    val filthiestLog: LogToClean = cleanerManager.grabFilthiestCompactedLog(time).get
+    cleanerManager.markPartitionUncleanable(logs.get(tp2).dir.getParent, tp2)
 
-    assertEquals(log3, filthiestLog.log)
-    assertEquals(tp3, filthiestLog.topicPartition)
+    val filthiestLog: LogToClean = cleanerManager.grabFilthiestCompactedLog(time).get
+    assertEquals(tp1, filthiestLog.topicPartition)
+    assertEquals(tp1, filthiestLog.log.topicPartition)
   }
 
   @Test
   def testGrabFilthiestCompactedLogIgnoresInProgressPartitions(): Unit = {
-    val records = TestUtils.singletonRecords("test".getBytes)
-    val log1: Log = createLog(records.sizeInBytes * 5, LogConfig.Compact, 1)
-    val log2: Log = createLog(records.sizeInBytes * 10, LogConfig.Compact, 2)
-    val log3: Log = createLog(records.sizeInBytes * 15, LogConfig.Compact, 3)
+    val tp0 = new TopicPartition("wishing-well", 0)
+    val tp1 = new TopicPartition("wishing-well", 1)
+    val tp2 = new TopicPartition("wishing-well", 2)
+    val partitions = Seq(tp0, tp1, tp2)
+
+    // setup logs with cleanable range: [20, 20], [20, 25], [20, 30]
+    val logs = setupIncreasinglyFilthyLogs(partitions, startNumBatches = 20, batchIncrement = 5)
+    val cleanerManager = createCleanerManagerMock(logs)
+    partitions.foreach(partition => cleanerCheckpoints.put(partition, 20))
 
-    val logs = new Pool[TopicPartition, Log]()
-    val tp1 = new TopicPartition("wishing well", 0) // active segment starts at 0
-    logs.put(tp1, log1)
-    val tp2 = new TopicPartition("wishing well", 1) // active segment starts at 10
-    logs.put(tp2, log2)
-    val tp3 = new TopicPartition("wishing well", 2) // // active segment starts at 20
-    logs.put(tp3, log3)
-    val cleanerManager: LogCleanerManagerMock = createCleanerManager(logs, toMock = true).asInstanceOf[LogCleanerManagerMock]
-    cleanerCheckpoints.put(tp1, 0) // all clean
-    cleanerCheckpoints.put(tp2, 1) // dirtiest - 9 dirty messages
-    cleanerCheckpoints.put(tp3, 15) // 5 dirty messages
     cleanerManager.setCleaningState(tp2, LogCleaningInProgress)
 
     val filthiestLog: LogToClean = cleanerManager.grabFilthiestCompactedLog(time).get
 
-    assertEquals(log3, filthiestLog.log)
-    assertEquals(tp3, filthiestLog.topicPartition)
+    assertEquals(tp1, filthiestLog.topicPartition)
+    assertEquals(tp1, filthiestLog.log.topicPartition)
   }
 
   @Test
   def testGrabFilthiestCompactedLogIgnoresBothInProgressPartitionsAndUncleanablePartitions(): Unit = {
-    val records = TestUtils.singletonRecords("test".getBytes)
-    val log1: Log = createLog(records.sizeInBytes * 5, LogConfig.Compact, 1)
-    val log2: Log = createLog(records.sizeInBytes * 10, LogConfig.Compact, 2)
-    val log3: Log = createLog(records.sizeInBytes * 15, LogConfig.Compact, 3)
+    val tp0 = new TopicPartition("wishing-well", 0)
+    val tp1 = new TopicPartition("wishing-well", 1)
+    val tp2 = new TopicPartition("wishing-well", 2)
+    val partitions = Seq(tp0, tp1, tp2)
+
+    // setup logs with cleanable range: [20, 20], [20, 25], [20, 30]
+    val logs = setupIncreasinglyFilthyLogs(partitions, startNumBatches = 20, batchIncrement = 5)
+    val cleanerManager = createCleanerManagerMock(logs)
+    partitions.foreach(partition => cleanerCheckpoints.put(partition, 20))
 
-    val logs = new Pool[TopicPartition, Log]()
-    val tp1 = new TopicPartition("wishing well", 0) // active segment starts at 0
-    logs.put(tp1, log1)
-    val tp2 = new TopicPartition("wishing well", 1) // active segment starts at 10
-    logs.put(tp2, log2)
-    val tp3 = new TopicPartition("wishing well", 2) // // active segment starts at 20
-    logs.put(tp3, log3)
-    val cleanerManager: LogCleanerManagerMock = createCleanerManager(logs, toMock = true).asInstanceOf[LogCleanerManagerMock]
-    cleanerCheckpoints.put(tp1, 0) // all clean
-    cleanerCheckpoints.put(tp2, 1) // dirtiest - 9 dirty messages
-    cleanerCheckpoints.put(tp3, 15) // 5 dirty messages
     cleanerManager.setCleaningState(tp2, LogCleaningInProgress)
-    cleanerManager.markPartitionUncleanable(log3.dir.getParent, tp3)
+    cleanerManager.markPartitionUncleanable(logs.get(tp1).dir.getParent, tp1)
 
     val filthiestLog: Option[LogToClean] = cleanerManager.grabFilthiestCompactedLog(time)
+    assertEquals(None, filthiestLog)
+  }
 
-    assertTrue(filthiestLog.isEmpty)
+  @Test
+  def testDirtyOffsetResetIfLargerThanEndOffset(): Unit = {
+    val tp = new TopicPartition("foo", 0)
+    val logs = setupIncreasinglyFilthyLogs(Seq(tp), startNumBatches = 20, batchIncrement = 5)
+    val cleanerManager = createCleanerManagerMock(logs)
+    cleanerCheckpoints.put(tp, 200)
+
+    val filthiestLog = cleanerManager.grabFilthiestCompactedLog(time).get
+    assertEquals(0L, filthiestLog.firstDirtyOffset)
+  }
+
+  @Test
+  def testDirtyOffsetResetIfSmallerThanStartOffset(): Unit = {
+    val tp = new TopicPartition("foo", 0)
+    val logs = setupIncreasinglyFilthyLogs(Seq(tp), startNumBatches = 20, batchIncrement = 5)
+
+    logs.get(tp).maybeIncrementLogStartOffset(10L)
+
+    val cleanerManager = createCleanerManagerMock(logs)
+    cleanerCheckpoints.put(tp, 0L)
+
+    val filthiestLog = cleanerManager.grabFilthiestCompactedLog(time).get
+    assertEquals(10L, filthiestLog.firstDirtyOffset)
+  }
+
+  @Test
+  def testLogStartOffsetLargerThanActiveSegmentBaseOffset(): Unit = {
+    val tp = new TopicPartition("foo", 0)
+    val log = createLog(segmentSize = 2048, LogConfig.Compact, tp)
+
+    val logs = new Pool[TopicPartition, Log]()
+    logs.put(tp, log)
+
+    appendRecords(log, numRecords = 3)
+    appendRecords(log, numRecords = 3)
+    appendRecords(log, numRecords = 3)
+
+    assertEquals(1, log.logSegments.size)
+
+    log.maybeIncrementLogStartOffset(2L)
+
+    val cleanerManager = createCleanerManagerMock(logs)
+    cleanerCheckpoints.put(tp, 0L)
+
+    // The active segment is uncleanable and hence not filthy from the POV of the CleanerManager.
+    val filthiestLog = cleanerManager.grabFilthiestCompactedLog(time)
+    assertEquals(None, filthiestLog)
+  }
+
+  @Test
+  def testDirtyOffsetLargerThanActiveSegmentBaseOffset(): Unit = {
+    // It is possible in the case of an unclean leader election for the checkpoint
+    // dirty offset to get ahead of the active segment base offset, but still be
+    // within the range of the log.
+
+    val tp = new TopicPartition("foo", 0)
+
+    val logs = new Pool[TopicPartition, Log]()
+    val log = createLog(2048, LogConfig.Compact, topicPartition = tp)
+    logs.put(tp, log)
+
+    appendRecords(log, numRecords = 3)
+    appendRecords(log, numRecords = 3)
+
+    assertEquals(1, log.logSegments.size)
+    assertEquals(0L, log.activeSegment.baseOffset)
+
+    val cleanerManager = createCleanerManagerMock(logs)
+    cleanerCheckpoints.put(tp, 3L)
+
+    // These segments are uncleanable and hence not filthy
+    val filthiestLog = cleanerManager.grabFilthiestCompactedLog(time)
+    assertEquals(None, filthiestLog)
   }
 
   /**
@@ -218,7 +322,7 @@ class LogCleanerManagerTest extends JUnitSuite with Logging {
     log.appendAsLeader(records, leaderEpoch = 0)
     log.roll()
     log.appendAsLeader(records, leaderEpoch = 0)
-    log.onHighWatermarkIncremented(2L)
+    log.updateHighWatermark(2L)
 
     // simulate cleanup thread working on the log partition
     val deletableLog = cleanerManager.pauseCleaningForNonCompactedPartitions()
@@ -324,10 +428,10 @@ class LogCleanerManagerTest extends JUnitSuite with Logging {
     while(log.numberOfSegments < 8)
       log.appendAsLeader(records(log.logEndOffset.toInt, log.logEndOffset.toInt, time.milliseconds()), leaderEpoch = 0)
 
-    val lastClean = Map(topicPartition -> 0L)
-    val cleanableOffsets = LogCleanerManager.cleanableOffsets(log, topicPartition, lastClean, time.milliseconds)
-    assertEquals("The first cleanable offset starts at the beginning of the log.", 0L, cleanableOffsets._1)
-    assertEquals("The first uncleanable offset begins with the active segment.", log.activeSegment.baseOffset, cleanableOffsets._2)
+    val lastCleanOffset = Some(0L)
+    val cleanableOffsets = LogCleanerManager.cleanableOffsets(log, lastCleanOffset, time.milliseconds)
+    assertEquals("The first cleanable offset starts at the beginning of the log.", 0L, cleanableOffsets.firstDirtyOffset)
+    assertEquals("The first uncleanable offset begins with the active segment.", log.activeSegment.baseOffset, cleanableOffsets.firstUncleanableDirtyOffset)
   }
 
   /**
@@ -354,10 +458,10 @@ class LogCleanerManagerTest extends JUnitSuite with Logging {
     while (log.numberOfSegments < 8)
       log.appendAsLeader(records(log.logEndOffset.toInt, log.logEndOffset.toInt, t1), leaderEpoch = 0)
 
-    val lastClean = Map(topicPartition -> 0L)
-    val cleanableOffsets = LogCleanerManager.cleanableOffsets(log, topicPartition, lastClean, time.milliseconds)
-    assertEquals("The first cleanable offset starts at the beginning of the log.", 0L, cleanableOffsets._1)
-    assertEquals("The first uncleanable offset begins with the second block of log entries.", activeSegAtT0.baseOffset, cleanableOffsets._2)
+    val lastCleanOffset = Some(0L)
+    val cleanableOffsets = LogCleanerManager.cleanableOffsets(log, lastCleanOffset, time.milliseconds)
+    assertEquals("The first cleanable offset starts at the beginning of the log.", 0L, cleanableOffsets.firstDirtyOffset)
+    assertEquals("The first uncleanable offset begins with the second block of log entries.", activeSegAtT0.baseOffset, cleanableOffsets.firstUncleanableDirtyOffset)
   }
 
   /**
@@ -379,10 +483,29 @@ class LogCleanerManagerTest extends JUnitSuite with Logging {
 
     time.sleep(compactionLag + 1)
 
-    val lastClean = Map(topicPartition -> 0L)
-    val cleanableOffsets = LogCleanerManager.cleanableOffsets(log, topicPartition, lastClean, time.milliseconds)
-    assertEquals("The first cleanable offset starts at the beginning of the log.", 0L, cleanableOffsets._1)
-    assertEquals("The first uncleanable offset begins with active segment.", log.activeSegment.baseOffset, cleanableOffsets._2)
+    val lastCleanOffset = Some(0L)
+    val cleanableOffsets = LogCleanerManager.cleanableOffsets(log, lastCleanOffset, time.milliseconds)
+    assertEquals("The first cleanable offset starts at the beginning of the log.", 0L, cleanableOffsets.firstDirtyOffset)
+    assertEquals("The first uncleanable offset begins with active segment.", log.activeSegment.baseOffset, cleanableOffsets.firstUncleanableDirtyOffset)
+  }
+
+  @Test
+  def testCleanableOffsetsNeedsCheckpointReset(): Unit = {
+    val tp = new TopicPartition("foo", 0)
+    val logs = setupIncreasinglyFilthyLogs(Seq(tp), startNumBatches = 20, batchIncrement = 5)
+    logs.get(tp).maybeIncrementLogStartOffset(10L)
+
+    var lastCleanOffset = Some(15L)
+    var cleanableOffsets = LogCleanerManager.cleanableOffsets(logs.get(tp), lastCleanOffset, time.milliseconds)
+    assertFalse("Checkpoint offset should not be reset if valid", cleanableOffsets.forceUpdateCheckpoint)
+
+    logs.get(tp).maybeIncrementLogStartOffset(20L)
+    cleanableOffsets = LogCleanerManager.cleanableOffsets(logs.get(tp), lastCleanOffset, time.milliseconds)
+    assertTrue("Checkpoint offset needs to be reset if less than log start offset", cleanableOffsets.forceUpdateCheckpoint)
+
+    lastCleanOffset = Some(25L)
+    cleanableOffsets = LogCleanerManager.cleanableOffsets(logs.get(tp), lastCleanOffset, time.milliseconds)
+    assertTrue("Checkpoint offset needs to be reset if greater than log end offset", cleanableOffsets.forceUpdateCheckpoint)
   }
 
   @Test
@@ -403,33 +526,31 @@ class LogCleanerManagerTest extends JUnitSuite with Logging {
     log.appendAsLeader(MemoryRecords.withTransactionalRecords(CompressionType.NONE, producerId, producerEpoch, sequence + 2,
       new SimpleRecord(time.milliseconds(), "3".getBytes, "c".getBytes)), leaderEpoch = 0)
     log.roll()
-    log.onHighWatermarkIncremented(3L)
+    log.updateHighWatermark(3L)
 
     time.sleep(compactionLag + 1)
     // although the compaction lag has been exceeded, the undecided data should not be cleaned
-    var cleanableOffsets = LogCleanerManager.cleanableOffsets(log, topicPartition,
-      Map(topicPartition -> 0L), time.milliseconds())
-    assertEquals(0L, cleanableOffsets._1)
-    assertEquals(0L, cleanableOffsets._2)
+    var cleanableOffsets = LogCleanerManager.cleanableOffsets(log, Some(0L), time.milliseconds())
+    assertEquals(0L, cleanableOffsets.firstDirtyOffset)
+    assertEquals(0L, cleanableOffsets.firstUncleanableDirtyOffset)
 
     log.appendAsLeader(MemoryRecords.withEndTransactionMarker(time.milliseconds(), producerId, producerEpoch,
-      new EndTransactionMarker(ControlRecordType.ABORT, 15)), leaderEpoch = 0, isFromClient = false)
+      new EndTransactionMarker(ControlRecordType.ABORT, 15)), leaderEpoch = 0,
+      origin = AppendOrigin.Coordinator)
     log.roll()
-    log.onHighWatermarkIncremented(4L)
+    log.updateHighWatermark(4L)
 
     // the first segment should now become cleanable immediately
-    cleanableOffsets = LogCleanerManager.cleanableOffsets(log, topicPartition,
-      Map(topicPartition -> 0L), time.milliseconds())
-    assertEquals(0L, cleanableOffsets._1)
-    assertEquals(3L, cleanableOffsets._2)
+    cleanableOffsets = LogCleanerManager.cleanableOffsets(log, Some(0L), time.milliseconds())
+    assertEquals(0L, cleanableOffsets.firstDirtyOffset)
+    assertEquals(3L, cleanableOffsets.firstUncleanableDirtyOffset)
 
     time.sleep(compactionLag + 1)
 
     // the second segment becomes cleanable after the compaction lag
-    cleanableOffsets = LogCleanerManager.cleanableOffsets(log, topicPartition,
-      Map(topicPartition -> 0L), time.milliseconds())
-    assertEquals(0L, cleanableOffsets._1)
-    assertEquals(4L, cleanableOffsets._2)
+    cleanableOffsets = LogCleanerManager.cleanableOffsets(log, Some(0L), time.milliseconds())
+    assertEquals(0L, cleanableOffsets.firstDirtyOffset)
+    assertEquals(4L, cleanableOffsets.firstUncleanableDirtyOffset)
   }
 
   @Test
@@ -479,29 +600,62 @@ class LogCleanerManagerTest extends JUnitSuite with Logging {
     assertEquals(LogCleaningPaused(1), cleanerManager.cleaningState(tp).get)
   }
 
+  /**
+   * Logs with invalid checkpoint offsets should update their checkpoint offset even if the log doesn't need cleaning
+   */
+  @Test
+  def testCheckpointUpdatedForInvalidOffsetNoCleaning(): Unit = {
+    val tp = new TopicPartition("foo", 0)
+    val logs = setupIncreasinglyFilthyLogs(Seq(tp), startNumBatches = 20, batchIncrement = 5)
+
+    logs.get(tp).maybeIncrementLogStartOffset(20L)
+    val cleanerManager = createCleanerManagerMock(logs)
+    cleanerCheckpoints.put(tp, 15L)
+
+    val filthiestLog = cleanerManager.grabFilthiestCompactedLog(time)
+    assertEquals("Log should not be selected for cleaning", None, filthiestLog)
+    assertEquals("Unselected log should have checkpoint offset updated", 20L, cleanerCheckpoints.get(tp).get)
+  }
+
+  /**
+   * Logs with invalid checkpoint offsets should update their checkpoint offset even if they aren't selected
+   * for immediate cleaning
+   */
+  @Test
+  def testCheckpointUpdatedForInvalidOffsetNotSelected(): Unit = {
+    val tp0 = new TopicPartition("foo", 0)
+    val tp1 = new TopicPartition("foo", 1)
+    val partitions = Seq(tp0, tp1)
+
+    // create two logs, one with an invalid offset, and one that is dirtier than the log with an invalid offset
+    val logs = setupIncreasinglyFilthyLogs(partitions, startNumBatches = 20, batchIncrement = 5)
+    logs.get(tp0).maybeIncrementLogStartOffset(15L)
+    val cleanerManager = createCleanerManagerMock(logs)
+    cleanerCheckpoints.put(tp0, 10L)
+    cleanerCheckpoints.put(tp1, 5L)
+
+    val filthiestLog = cleanerManager.grabFilthiestCompactedLog(time).get
+    assertEquals("Dirtier log should be selected", tp1, filthiestLog.topicPartition)
+    assertEquals("Unselected log should have checkpoint offset updated", 15L, cleanerCheckpoints.get(tp0).get)
+  }
+
   private def createCleanerManager(log: Log): LogCleanerManager = {
     val logs = new Pool[TopicPartition, Log]()
     logs.put(topicPartition, log)
-    createCleanerManager(logs)
+    new LogCleanerManager(Array(logDir), logs, null, 0L)
   }
 
-  private def createCleanerManager(pool: Pool[TopicPartition, Log], toMock: Boolean = false): LogCleanerManager = {
-    if (toMock)
-      new LogCleanerManagerMock(Array(logDir), pool, null)
-    else
-      new LogCleanerManager(Array(logDir), pool, null)
+  private def createCleanerManagerMock(pool: Pool[TopicPartition, Log]): LogCleanerManagerMock = {
+    new LogCleanerManagerMock(Array(logDir), pool, null)
   }
 
-  private def createLog(segmentSize: Int, cleanupPolicy: String, segmentsCount: Int = 0): Log = {
-    val logProps = new Properties()
-    logProps.put(LogConfig.SegmentBytesProp, segmentSize: Integer)
-    logProps.put(LogConfig.RetentionMsProp, 1: Integer)
-    logProps.put(LogConfig.CleanupPolicyProp, cleanupPolicy)
-    logProps.put(LogConfig.MinCleanableDirtyRatioProp, 0.05: java.lang.Double) // small for easier and clearer tests
+  private def createLog(segmentSize: Int,
+                        cleanupPolicy: String,
+                        topicPartition: TopicPartition = new TopicPartition("log", 0)): Log = {
+    val config = createLowRetentionLogConfig(segmentSize, cleanupPolicy)
+    val partitionDir = new File(logDir, Log.logDirName(topicPartition))
 
-    val config = LogConfig(logProps)
-    val partitionDir = new File(logDir, "log-0")
-    val log = Log(partitionDir,
+    Log(partitionDir,
       config,
       logStartOffset = 0L,
       recoveryPoint = 0L,
@@ -511,23 +665,43 @@ class LogCleanerManagerTest extends JUnitSuite with Logging {
       maxProducerIdExpirationMs = 60 * 60 * 1000,
       producerIdExpirationCheckIntervalMs = LogManager.ProducerIdExpirationCheckIntervalMs,
       logDirFailureChannel = new LogDirFailureChannel(10))
-    for (i <- 0 until segmentsCount) {
-      val startOffset = i * 10
-      val endOffset = startOffset + 10
-      val segment = LogUtils.createSegment(startOffset, logDir)
-      var lastTimestamp = 0L
-      val records = (startOffset until endOffset).map { offset =>
-        val currentTimestamp = time.milliseconds()
-        if (offset == endOffset - 1)
-          lastTimestamp = currentTimestamp
-
-        new SimpleRecord(currentTimestamp, s"key-$offset".getBytes, s"value-$offset".getBytes)
-      }
-
-      segment.append(endOffset, lastTimestamp, endOffset, MemoryRecords.withRecords(CompressionType.NONE, records:_*))
-      log.addSegment(segment)
+  }
+
+  private def createLowRetentionLogConfig(segmentSize: Int, cleanupPolicy: String): LogConfig = {
+    val logProps = new Properties()
+    logProps.put(LogConfig.SegmentBytesProp, segmentSize: Integer)
+    logProps.put(LogConfig.RetentionMsProp, 1: Integer)
+    logProps.put(LogConfig.CleanupPolicyProp, cleanupPolicy)
+    logProps.put(LogConfig.MinCleanableDirtyRatioProp, 0.05: java.lang.Double) // small for easier and clearer tests
+
+    LogConfig(logProps)
+  }
+
+  private def writeRecords(log: Log,
+                           numBatches: Int,
+                           recordsPerBatch: Int,
+                           batchesPerSegment: Int): Unit = {
+    for (i <- 0 until numBatches) {
+      appendRecords(log, recordsPerBatch)
+      if (i % batchesPerSegment == 0)
+        log.roll()
     }
-    log
+    log.roll()
+  }
+
+  private def appendRecords(log: Log, numRecords: Int): Unit = {
+    val startOffset = log.logEndOffset
+    val endOffset = startOffset + numRecords
+    var lastTimestamp = 0L
+    val records = (startOffset until endOffset).map { offset =>
+      val currentTimestamp = time.milliseconds()
+      if (offset == endOffset - 1)
+        lastTimestamp = currentTimestamp
+      new SimpleRecord(currentTimestamp, s"key-$offset".getBytes, s"value-$offset".getBytes)
+    }
+
+    log.appendAsLeader(MemoryRecords.withRecords(CompressionType.NONE, records:_*), leaderEpoch = 1)
+    log.maybeIncrementHighWatermark(log.logEndOffsetMetadata)
   }
 
   private def makeLog(dir: File = logDir, config: LogConfig) =
diff --git a/core/src/test/scala/unit/kafka/log/LogCleanerParameterizedIntegrationTest.scala b/core/src/test/scala/unit/kafka/log/LogCleanerParameterizedIntegrationTest.scala
index 61e3ea57604d0..0d1e311659571 100755
--- a/core/src/test/scala/unit/kafka/log/LogCleanerParameterizedIntegrationTest.scala
+++ b/core/src/test/scala/unit/kafka/log/LogCleanerParameterizedIntegrationTest.scala
@@ -33,7 +33,6 @@ import org.junit.runner.RunWith
 import org.junit.runners.Parameterized
 import org.junit.runners.Parameterized.Parameters
 
-import scala.Seq
 import scala.collection._
 
 /**
@@ -49,7 +48,7 @@ class LogCleanerParameterizedIntegrationTest(compressionCodec: String) extends A
 
 
   @Test
-  def cleanerTest() {
+  def cleanerTest(): Unit = {
     val largeMessageKey = 20
     val (largeMessageValue, largeMessageSet) = createLargeSingleMessageSet(largeMessageKey, RecordBatch.CURRENT_MAGIC_VALUE)
     val maxMessageSize = largeMessageSet.sizeInBytes
@@ -102,7 +101,7 @@ class LogCleanerParameterizedIntegrationTest(compressionCodec: String) extends A
       val messages = writeDups(numKeys = numKeys, numDups = 3, log = log, codec = codec)
       val startSize = log.size
 
-      log.onHighWatermarkIncremented(log.logEndOffset)
+      log.updateHighWatermark(log.logEndOffset)
 
       val firstDirty = log.activeSegment.baseOffset
       cleaner.startup()
@@ -115,10 +114,12 @@ class LogCleanerParameterizedIntegrationTest(compressionCodec: String) extends A
     }
 
     val (log, _) = runCleanerAndCheckCompacted(100)
-    // should delete old segments
-    log.logSegments.foreach(_.lastModified = time.milliseconds - (2 * retentionMs))
 
-    TestUtils.waitUntilTrue(() => log.numberOfSegments == 1, "There should only be 1 segment remaining", 10000L)
+    // Set the last modified time to an old value to force deletion of old segments
+    val endOffset = log.logEndOffset
+    log.logSegments.foreach(_.lastModified = time.milliseconds - (2 * retentionMs))
+    TestUtils.waitUntilTrue(() => log.logStartOffset == endOffset,
+      "Timed out waiting for deletion of old segments")
     assertEquals(1, log.numberOfSegments)
 
     cleaner.shutdown()
@@ -225,7 +226,7 @@ class LogCleanerParameterizedIntegrationTest(compressionCodec: String) extends A
   }
 
   @Test
-  def cleanerConfigUpdateTest() {
+  def cleanerConfigUpdateTest(): Unit = {
     val largeMessageKey = 20
     val (largeMessageValue, largeMessageSet) = createLargeSingleMessageSet(largeMessageKey, RecordBatch.CURRENT_MAGIC_VALUE)
     val maxMessageSize = largeMessageSet.sizeInBytes
@@ -274,7 +275,7 @@ class LogCleanerParameterizedIntegrationTest(compressionCodec: String) extends A
     assertTrue(s"log should have been compacted: startSize=$startSize compactedSize=$compactedSize", startSize > compactedSize)
   }
 
-  private def checkLastCleaned(topic: String, partitionId: Int, firstDirty: Long) {
+  private def checkLastCleaned(topic: String, partitionId: Int, firstDirty: Long): Unit = {
     // wait until cleaning up to base_offset, note that cleaning happens only when "log dirty ratio" is higher than
     // LogConfig.MinCleanableDirtyRatioProp
     val topicPartition = new TopicPartition(topic, partitionId)
@@ -284,7 +285,7 @@ class LogCleanerParameterizedIntegrationTest(compressionCodec: String) extends A
       lastCleaned >= firstDirty)
   }
 
-  private def checkLogAfterAppendingDups(log: Log, startSize: Long, appends: Seq[(Int, String, Long)]) {
+  private def checkLogAfterAppendingDups(log: Log, startSize: Long, appends: Seq[(Int, String, Long)]): Unit = {
     val read = readFromLog(log)
     assertEquals("Contents of the map shouldn't change", toMap(appends), toMap(read))
     assertTrue(startSize > log.size)
@@ -327,7 +328,7 @@ object LogCleanerParameterizedIntegrationTest {
   @Parameters
   def parameters: java.util.Collection[Array[String]] = {
     val list = new java.util.ArrayList[Array[String]]()
-    for (codec <- CompressionType.values)
+    for (codec <- CompressionType.values.filter(_!=CompressionType.PASSTHROUGH))
       list.add(Array(codec.name))
     list
   }
diff --git a/core/src/test/scala/unit/kafka/log/LogCleanerTest.scala b/core/src/test/scala/unit/kafka/log/LogCleanerTest.scala
index 51477b627cbc5..644822cc4f45e 100755
--- a/core/src/test/scala/unit/kafka/log/LogCleanerTest.scala
+++ b/core/src/test/scala/unit/kafka/log/LogCleanerTest.scala
@@ -32,7 +32,7 @@ import org.apache.kafka.common.record._
 import org.apache.kafka.common.utils.Utils
 import org.junit.Assert._
 import org.junit.{After, Test}
-import org.scalatest.junit.JUnitSuite
+import org.scalatest.Assertions.{assertThrows, fail, intercept}
 
 import scala.collection.JavaConverters._
 import scala.collection._
@@ -40,7 +40,7 @@ import scala.collection._
 /**
  * Unit tests for the log cleaning logic
  */
-class LogCleanerTest extends JUnitSuite {
+class LogCleanerTest {
 
   val tmpdir = TestUtils.tempDir()
   val dir = TestUtils.randomPartitionLogDir(tmpdir)
@@ -84,7 +84,7 @@ class LogCleanerTest extends JUnitSuite {
     val segments = log.logSegments.take(3).toSeq
     val stats = new CleanerStats()
     val expectedBytesRead = segments.map(_.size).sum
-    cleaner.cleanSegments(log, segments, map, 0L, stats)
+    cleaner.cleanSegments(log, segments, map, 0L, stats, new CleanedTransactionMetadata)
     val shouldRemain = LogTest.keysInLog(log).filter(!keys.contains(_))
     assertEquals(shouldRemain, LogTest.keysInLog(log))
     assertEquals(expectedBytesRead, stats.bytesRead)
@@ -126,8 +126,9 @@ class LogCleanerTest extends JUnitSuite {
     val t = new Thread() {
       override def run(): Unit = {
         deleteStartLatch.await(5000, TimeUnit.MILLISECONDS)
+        log.updateHighWatermark(log.activeSegment.baseOffset)
         log.maybeIncrementLogStartOffset(log.activeSegment.baseOffset)
-        log.onHighWatermarkIncremented(log.activeSegment.baseOffset)
+        log.updateHighWatermark(log.activeSegment.baseOffset)
         log.deleteOldSegments()
         deleteCompleteLatch.countDown()
       }
@@ -151,7 +152,7 @@ class LogCleanerTest extends JUnitSuite {
     val segments = log.logSegments(0, log.activeSegment.baseOffset).toSeq
     val stats = new CleanerStats()
     cleaner.buildOffsetMap(log, 0, log.activeSegment.baseOffset, offsetMap, stats)
-    cleaner.cleanSegments(log, segments, offsetMap, 0L, stats)
+    cleaner.cleanSegments(log, segments, offsetMap, 0L, stats, new CleanedTransactionMetadata)
 
     // Validate based on the file name that log segment file is renamed exactly once for async deletion
     assertEquals(expectedFileName, firstLogFile.file().getPath)
@@ -264,10 +265,10 @@ class LogCleanerTest extends JUnitSuite {
     appendProducer1(Seq(1, 2))
     appendProducer2(Seq(2, 3))
     appendProducer1(Seq(3, 4))
-    log.appendAsLeader(abortMarker(pid1, producerEpoch), leaderEpoch = 0, isFromClient = false)
-    log.appendAsLeader(commitMarker(pid2, producerEpoch), leaderEpoch = 0, isFromClient = false)
+    log.appendAsLeader(abortMarker(pid1, producerEpoch), leaderEpoch = 0, origin = AppendOrigin.Coordinator)
+    log.appendAsLeader(commitMarker(pid2, producerEpoch), leaderEpoch = 0, origin = AppendOrigin.Coordinator)
     appendProducer1(Seq(2))
-    log.appendAsLeader(commitMarker(pid1, producerEpoch), leaderEpoch = 0, isFromClient = false)
+    log.appendAsLeader(commitMarker(pid1, producerEpoch), leaderEpoch = 0, origin = AppendOrigin.Coordinator)
 
     val abortedTransactions = log.collectAbortedTransactions(log.logStartOffset, log.logEndOffset)
 
@@ -305,11 +306,11 @@ class LogCleanerTest extends JUnitSuite {
     appendProducer2(Seq(5, 6))
     appendProducer3(Seq(6, 7))
     appendProducer1(Seq(7, 8))
-    log.appendAsLeader(abortMarker(pid2, producerEpoch), leaderEpoch = 0, isFromClient = false)
+    log.appendAsLeader(abortMarker(pid2, producerEpoch), leaderEpoch = 0, origin = AppendOrigin.Coordinator)
     appendProducer3(Seq(8, 9))
-    log.appendAsLeader(commitMarker(pid3, producerEpoch), leaderEpoch = 0, isFromClient = false)
+    log.appendAsLeader(commitMarker(pid3, producerEpoch), leaderEpoch = 0, origin = AppendOrigin.Coordinator)
     appendProducer1(Seq(9, 10))
-    log.appendAsLeader(abortMarker(pid1, producerEpoch), leaderEpoch = 0, isFromClient = false)
+    log.appendAsLeader(abortMarker(pid1, producerEpoch), leaderEpoch = 0, origin = AppendOrigin.Coordinator)
 
     // we have only cleaned the records in the first segment
     val dirtyOffset = cleaner.clean(LogToClean(new TopicPartition("test", 0), log, 0L, log.activeSegment.baseOffset))._1
@@ -340,32 +341,32 @@ class LogCleanerTest extends JUnitSuite {
 
     appendProducer(Seq(1))
     appendProducer(Seq(2, 3))
-    log.appendAsLeader(commitMarker(producerId, producerEpoch), leaderEpoch = 0, isFromClient = false)
+    log.appendAsLeader(commitMarker(producerId, producerEpoch), leaderEpoch = 0, origin = AppendOrigin.Coordinator)
     appendProducer(Seq(2))
-    log.appendAsLeader(commitMarker(producerId, producerEpoch), leaderEpoch = 0, isFromClient = false)
+    log.appendAsLeader(commitMarker(producerId, producerEpoch), leaderEpoch = 0, origin = AppendOrigin.Coordinator)
     log.roll()
 
     // cannot remove the marker in this pass because there are still valid records
-    var dirtyOffset = cleaner.doClean(LogToClean(tp, log, 0L, 100L), deleteHorizonMs = Long.MaxValue)._1
+    var dirtyOffset = cleaner.doClean(LogToClean(tp, log, 0L, log.activeSegment.baseOffset), deleteHorizonMs = Long.MaxValue)._1
     assertEquals(List(1, 3, 2), LogTest.keysInLog(log))
     assertEquals(List(0, 2, 3, 4, 5), offsetsInLog(log))
 
     appendProducer(Seq(1, 3))
-    log.appendAsLeader(commitMarker(producerId, producerEpoch), leaderEpoch = 0, isFromClient = false)
+    log.appendAsLeader(commitMarker(producerId, producerEpoch), leaderEpoch = 0, origin = AppendOrigin.Coordinator)
     log.roll()
 
     // the first cleaning preserves the commit marker (at offset 3) since there were still records for the transaction
-    dirtyOffset = cleaner.doClean(LogToClean(tp, log, dirtyOffset, 100L), deleteHorizonMs = Long.MaxValue)._1
+    dirtyOffset = cleaner.doClean(LogToClean(tp, log, dirtyOffset, log.activeSegment.baseOffset), deleteHorizonMs = Long.MaxValue)._1
     assertEquals(List(2, 1, 3), LogTest.keysInLog(log))
     assertEquals(List(3, 4, 5, 6, 7, 8), offsetsInLog(log))
 
     // delete horizon forced to 0 to verify marker is not removed early
-    dirtyOffset = cleaner.doClean(LogToClean(tp, log, dirtyOffset, 100L), deleteHorizonMs = 0L)._1
+    dirtyOffset = cleaner.doClean(LogToClean(tp, log, dirtyOffset, log.activeSegment.baseOffset), deleteHorizonMs = 0L)._1
     assertEquals(List(2, 1, 3), LogTest.keysInLog(log))
     assertEquals(List(3, 4, 5, 6, 7, 8), offsetsInLog(log))
 
     // clean again with large delete horizon and verify the marker is removed
-    dirtyOffset = cleaner.doClean(LogToClean(tp, log, dirtyOffset, 100L), deleteHorizonMs = Long.MaxValue)._1
+    dirtyOffset = cleaner.doClean(LogToClean(tp, log, dirtyOffset, log.activeSegment.baseOffset), deleteHorizonMs = Long.MaxValue)._1
     assertEquals(List(2, 1, 3), LogTest.keysInLog(log))
     assertEquals(List(4, 5, 6, 7, 8), offsetsInLog(log))
   }
@@ -388,17 +389,17 @@ class LogCleanerTest extends JUnitSuite {
     val appendProducer = appendTransactionalAsLeader(log, producerId, producerEpoch)
 
     appendProducer(Seq(1))
-    log.appendAsLeader(abortMarker(producerId, producerEpoch), leaderEpoch = 0, isFromClient = false)
+    log.appendAsLeader(abortMarker(producerId, producerEpoch), leaderEpoch = 0, origin = AppendOrigin.Coordinator)
     appendProducer(Seq(2))
     appendProducer(Seq(2))
-    log.appendAsLeader(commitMarker(producerId, producerEpoch), leaderEpoch = 0, isFromClient = false)
+    log.appendAsLeader(commitMarker(producerId, producerEpoch), leaderEpoch = 0, origin = AppendOrigin.Coordinator)
     log.roll()
 
-    cleaner.doClean(LogToClean(tp, log, 0L, 100L), deleteHorizonMs = Long.MaxValue)
+    cleaner.doClean(LogToClean(tp, log, 0L, log.activeSegment.baseOffset), deleteHorizonMs = Long.MaxValue)
     assertEquals(List(2), LogTest.keysInLog(log))
     assertEquals(List(1, 3, 4), offsetsInLog(log))
 
-    cleaner.doClean(LogToClean(tp, log, 0L, 100L), deleteHorizonMs = Long.MaxValue)
+    cleaner.doClean(LogToClean(tp, log, 0L, log.activeSegment.baseOffset), deleteHorizonMs = Long.MaxValue)
     assertEquals(List(2), LogTest.keysInLog(log))
     assertEquals(List(3, 4), offsetsInLog(log))
   }
@@ -421,48 +422,51 @@ class LogCleanerTest extends JUnitSuite {
 
     // [{Producer1: 2, 3}], [{Producer2: 2, 3}, {Producer2: Commit}]
     producer2(Seq(2, 3)) // offsets 2, 3
-    log.appendAsLeader(commitMarker(2L, producerEpoch), leaderEpoch = 0, isFromClient = false) // offset 4
+    log.appendAsLeader(commitMarker(2L, producerEpoch), leaderEpoch = 0,
+      origin = AppendOrigin.Coordinator) // offset 4
     log.roll()
 
     // [{Producer1: 2, 3}], [{Producer2: 2, 3}, {Producer2: Commit}], [{2}, {3}, {Producer1: Commit}]
     //  {0, 1},              {2, 3},            {4},                   {5}, {6}, {7} ==> Offsets
     log.appendAsLeader(record(2, 2), leaderEpoch = 0) // offset 5
     log.appendAsLeader(record(3, 3), leaderEpoch = 0) // offset 6
-    log.appendAsLeader(commitMarker(1L, producerEpoch), leaderEpoch = 0, isFromClient = false) // offset 7
+    log.appendAsLeader(commitMarker(1L, producerEpoch), leaderEpoch = 0,
+      origin = AppendOrigin.Coordinator) // offset 7
     log.roll()
 
     // first time through the records are removed
-    // Expected State: [{Producer1: EmptyBatch}, {Producer2: EmptyBatch}, {Producer2: Commit}, {2}, {3}]
-    var dirtyOffset = cleaner.doClean(LogToClean(tp, log, 0L, 100L), deleteHorizonMs = Long.MaxValue)._1
+    // Expected State: [{Producer1: EmptyBatch}, {Producer2: EmptyBatch}, {Producer2: Commit}, {2}, {3}, {Producer1: Commit}]
+    var dirtyOffset = cleaner.doClean(LogToClean(tp, log, 0L, log.activeSegment.baseOffset), deleteHorizonMs = Long.MaxValue)._1
     assertEquals(List(2, 3), LogTest.keysInLog(log))
-    assertEquals(List(4, 5, 6), offsetsInLog(log))
-    assertEquals(List(1, 3, 4, 5, 6), lastOffsetsPerBatchInLog(log))
+    assertEquals(List(4, 5, 6, 7), offsetsInLog(log))
+    assertEquals(List(1, 3, 4, 5, 6, 7), lastOffsetsPerBatchInLog(log))
 
     // the empty batch remains if cleaned again because it still holds the last sequence
-    // Expected State: [{Producer1: EmptyBatch}, {Producer2: EmptyBatch}, {Producer2: Commit}, {2}, {3}]
-    dirtyOffset = cleaner.doClean(LogToClean(tp, log, dirtyOffset, 100L), deleteHorizonMs = Long.MaxValue)._1
+    // Expected State: [{Producer1: EmptyBatch}, {Producer2: EmptyBatch}, {Producer2: Commit}, {2}, {3}, {Producer1: Commit}]
+    dirtyOffset = cleaner.doClean(LogToClean(tp, log, dirtyOffset, log.activeSegment.baseOffset), deleteHorizonMs = Long.MaxValue)._1
     assertEquals(List(2, 3), LogTest.keysInLog(log))
-    assertEquals(List(4, 5, 6), offsetsInLog(log))
-    assertEquals(List(1, 3, 4, 5, 6), lastOffsetsPerBatchInLog(log))
+    assertEquals(List(4, 5, 6, 7), offsetsInLog(log))
+    assertEquals(List(1, 3, 4, 5, 6, 7), lastOffsetsPerBatchInLog(log))
 
     // append a new record from the producer to allow cleaning of the empty batch
-    // [{Producer1: EmptyBatch}, {Producer2: EmptyBatch}, {Producer2: Commit}, {2}, {3}] [{Producer2: 1}, {Producer2: Commit}]
-    //  {1},                     {3},                     {4},                 {5}, {6},  {8},            {9} ==> Offsets
+    // [{Producer1: EmptyBatch}, {Producer2: EmptyBatch}, {Producer2: Commit}, {2}, {3}, {Producer1: Commit}, {Producer2: 1}, {Producer2: Commit}]
+    //  {1},                     {3},                     {4},                 {5}, {6}, {7},                 {8},            {9} ==> Offsets
     producer2(Seq(1)) // offset 8
-    log.appendAsLeader(commitMarker(2L, producerEpoch), leaderEpoch = 0, isFromClient = false) // offset 9
+    log.appendAsLeader(commitMarker(2L, producerEpoch), leaderEpoch = 0,
+      origin = AppendOrigin.Coordinator) // offset 9
     log.roll()
 
-    // Expected State: [{Producer1: EmptyBatch}, {Producer2: Commit}, {2}, {3}, {Producer2: 1}, {Producer2: Commit}]
-    dirtyOffset = cleaner.doClean(LogToClean(tp, log, dirtyOffset, 100L), deleteHorizonMs = Long.MaxValue)._1
+    // Expected State: [{Producer1: EmptyBatch}, {Producer2: Commit}, {2}, {3}, {Producer1: Commit}, {Producer2: 1}, {Producer2: Commit}]
+    dirtyOffset = cleaner.doClean(LogToClean(tp, log, dirtyOffset, log.activeSegment.baseOffset), deleteHorizonMs = Long.MaxValue)._1
     assertEquals(List(2, 3, 1), LogTest.keysInLog(log))
-    assertEquals(List(4, 5, 6, 8, 9), offsetsInLog(log))
-    assertEquals(List(1, 4, 5, 6, 8, 9), lastOffsetsPerBatchInLog(log))
+    assertEquals(List(4, 5, 6, 7, 8, 9), offsetsInLog(log))
+    assertEquals(List(1, 4, 5, 6, 7, 8, 9), lastOffsetsPerBatchInLog(log))
 
-    // Expected State: [{Producer1: EmptyBatch}, {2}, {3}, {Producer2: 1}, {Producer2: Commit}]
-    dirtyOffset = cleaner.doClean(LogToClean(tp, log, dirtyOffset, 100L), deleteHorizonMs = Long.MaxValue)._1
+    // Expected State: [{Producer1: EmptyBatch}, {2}, {3}, {Producer1: Commit}, {Producer2: 1}, {Producer2: Commit}]
+    dirtyOffset = cleaner.doClean(LogToClean(tp, log, dirtyOffset, log.activeSegment.baseOffset), deleteHorizonMs = Long.MaxValue)._1
     assertEquals(List(2, 3, 1), LogTest.keysInLog(log))
-    assertEquals(List(5, 6, 8, 9), offsetsInLog(log))
-    assertEquals(List(1, 5, 6, 8, 9), lastOffsetsPerBatchInLog(log))
+    assertEquals(List(5, 6, 7, 8, 9), offsetsInLog(log))
+    assertEquals(List(1, 5, 6, 7, 8, 9), lastOffsetsPerBatchInLog(log))
   }
 
   @Test
@@ -476,26 +480,80 @@ class LogCleanerTest extends JUnitSuite {
     val producerEpoch = 0.toShort
 
     // [{Producer1: Commit}, {2}, {3}]
-    log.appendAsLeader(commitMarker(1L, producerEpoch), leaderEpoch = 0, isFromClient = false) // offset 7
+    log.appendAsLeader(commitMarker(1L, producerEpoch), leaderEpoch = 0,
+      origin = AppendOrigin.Coordinator) // offset 1
     log.appendAsLeader(record(2, 2), leaderEpoch = 0) // offset 2
     log.appendAsLeader(record(3, 3), leaderEpoch = 0) // offset 3
     log.roll()
 
     // first time through the control batch is retained as an empty batch
     // Expected State: [{Producer1: EmptyBatch}], [{2}, {3}]
-    var dirtyOffset = cleaner.doClean(LogToClean(tp, log, 0L, 100L), deleteHorizonMs = Long.MaxValue)._1
+    var dirtyOffset = cleaner.doClean(LogToClean(tp, log, 0L, log.activeSegment.baseOffset), deleteHorizonMs = Long.MaxValue)._1
     assertEquals(List(2, 3), LogTest.keysInLog(log))
     assertEquals(List(1, 2), offsetsInLog(log))
     assertEquals(List(0, 1, 2), lastOffsetsPerBatchInLog(log))
 
     // the empty control batch does not cause an exception when cleaned
     // Expected State: [{Producer1: EmptyBatch}], [{2}, {3}]
-    dirtyOffset = cleaner.doClean(LogToClean(tp, log, dirtyOffset, 100L), deleteHorizonMs = Long.MaxValue)._1
+    dirtyOffset = cleaner.doClean(LogToClean(tp, log, dirtyOffset, log.activeSegment.baseOffset), deleteHorizonMs = Long.MaxValue)._1
     assertEquals(List(2, 3), LogTest.keysInLog(log))
     assertEquals(List(1, 2), offsetsInLog(log))
     assertEquals(List(0, 1, 2), lastOffsetsPerBatchInLog(log))
   }
 
+  @Test
+  def testCommittedTransactionSpanningSegments(): Unit = {
+    val tp = new TopicPartition("test", 0)
+    val cleaner = makeCleaner(Int.MaxValue)
+    val logProps = new Properties()
+    logProps.put(LogConfig.SegmentBytesProp, 128: java.lang.Integer)
+    val log = makeLog(config = LogConfig.fromProps(logConfig.originals, logProps))
+    val producerEpoch = 0.toShort
+    val producerId = 1L
+
+    val appendTransaction = appendTransactionalAsLeader(log, producerId, producerEpoch)
+    appendTransaction(Seq(1))
+    log.roll()
+
+    log.appendAsLeader(commitMarker(producerId, producerEpoch), leaderEpoch = 0, origin = AppendOrigin.Coordinator)
+    log.roll()
+
+    // Both the record and the marker should remain after cleaning
+    cleaner.doClean(LogToClean(tp, log, 0L, log.activeSegment.baseOffset), deleteHorizonMs = Long.MaxValue)
+    assertEquals(List(0, 1), offsetsInLog(log))
+    assertEquals(List(0, 1), lastOffsetsPerBatchInLog(log))
+  }
+
+  @Test
+  def testAbortedTransactionSpanningSegments(): Unit = {
+    val tp = new TopicPartition("test", 0)
+    val cleaner = makeCleaner(Int.MaxValue)
+    val logProps = new Properties()
+    logProps.put(LogConfig.SegmentBytesProp, 128: java.lang.Integer)
+    val log = makeLog(config = LogConfig.fromProps(logConfig.originals, logProps))
+    val producerEpoch = 0.toShort
+    val producerId = 1L
+
+    val appendTransaction = appendTransactionalAsLeader(log, producerId, producerEpoch)
+    appendTransaction(Seq(1))
+    log.roll()
+
+    log.appendAsLeader(abortMarker(producerId, producerEpoch), leaderEpoch = 0, origin = AppendOrigin.Coordinator)
+    log.roll()
+
+    // Both the batch and the marker should remain after cleaning. The batch is retained
+    // because it is the last entry for this producerId. The marker is retained because
+    // there are still batches remaining from this transaction.
+    cleaner.doClean(LogToClean(tp, log, 0L, log.activeSegment.baseOffset), deleteHorizonMs = Long.MaxValue)
+    assertEquals(List(1), offsetsInLog(log))
+    assertEquals(List(0, 1), lastOffsetsPerBatchInLog(log))
+
+    // The empty batch and the marker is still retained after a second cleaning.
+    cleaner.doClean(LogToClean(tp, log, 0L, log.activeSegment.baseOffset), deleteHorizonMs = Long.MaxValue)
+    assertEquals(List(1), offsetsInLog(log))
+    assertEquals(List(0, 1), lastOffsetsPerBatchInLog(log))
+  }
+
   @Test
   def testAbortMarkerRemoval(): Unit = {
     val tp = new TopicPartition("test", 0)
@@ -510,22 +568,63 @@ class LogCleanerTest extends JUnitSuite {
 
     appendProducer(Seq(1))
     appendProducer(Seq(2, 3))
-    log.appendAsLeader(abortMarker(producerId, producerEpoch), leaderEpoch = 0, isFromClient = false)
+    log.appendAsLeader(abortMarker(producerId, producerEpoch), leaderEpoch = 0, origin = AppendOrigin.Coordinator)
     appendProducer(Seq(3))
-    log.appendAsLeader(commitMarker(producerId, producerEpoch), leaderEpoch = 0, isFromClient = false)
+    log.appendAsLeader(commitMarker(producerId, producerEpoch), leaderEpoch = 0, origin = AppendOrigin.Coordinator)
     log.roll()
 
     // delete horizon set to 0 to verify marker is not removed early
-    val dirtyOffset = cleaner.doClean(LogToClean(tp, log, 0L, 100L), deleteHorizonMs = 0L)._1
+    val dirtyOffset = cleaner.doClean(LogToClean(tp, log, 0L, log.activeSegment.baseOffset), deleteHorizonMs = 0L)._1
     assertEquals(List(3), LogTest.keysInLog(log))
     assertEquals(List(3, 4, 5), offsetsInLog(log))
 
     // clean again with large delete horizon and verify the marker is removed
-    cleaner.doClean(LogToClean(tp, log, dirtyOffset, 100L), deleteHorizonMs = Long.MaxValue)
+    cleaner.doClean(LogToClean(tp, log, dirtyOffset, log.activeSegment.baseOffset), deleteHorizonMs = Long.MaxValue)
     assertEquals(List(3), LogTest.keysInLog(log))
     assertEquals(List(4, 5), offsetsInLog(log))
   }
 
+  @Test
+  def testEmptyBatchRemovalWithSequenceReuse(): Unit = {
+    // The group coordinator always writes batches beginning with sequence number 0. This test
+    // ensures that we still remove old empty batches and transaction markers under this expectation.
+
+    val producerEpoch = 0.toShort
+    val producerId = 1L
+    val tp = new TopicPartition("test", 0)
+    val cleaner = makeCleaner(Int.MaxValue)
+    val logProps = new Properties()
+    logProps.put(LogConfig.SegmentBytesProp, 2048: java.lang.Integer)
+    val log = makeLog(config = LogConfig.fromProps(logConfig.originals, logProps))
+
+    val appendFirstTransaction = appendTransactionalAsLeader(log, producerId, producerEpoch,
+      origin = AppendOrigin.Replication)
+    appendFirstTransaction(Seq(1))
+    log.appendAsLeader(commitMarker(producerId, producerEpoch), leaderEpoch = 0, origin = AppendOrigin.Coordinator)
+
+    val appendSecondTransaction = appendTransactionalAsLeader(log, producerId, producerEpoch,
+      origin = AppendOrigin.Replication)
+    appendSecondTransaction(Seq(2))
+    log.appendAsLeader(commitMarker(producerId, producerEpoch), leaderEpoch = 0, origin = AppendOrigin.Coordinator)
+
+    log.appendAsLeader(record(1, 1), leaderEpoch = 0)
+    log.appendAsLeader(record(2, 1), leaderEpoch = 0)
+
+    // Roll the log to ensure that the data is cleanable.
+    log.roll()
+
+    // Both transactional batches will be cleaned. The last one will remain in the log
+    // as an empty batch in order to preserve the producer sequence number and epoch
+    cleaner.doClean(LogToClean(tp, log, 0L, log.activeSegment.baseOffset), deleteHorizonMs = Long.MaxValue)
+    assertEquals(List(1, 3, 4, 5), offsetsInLog(log))
+    assertEquals(List(1, 2, 3, 4, 5), lastOffsetsPerBatchInLog(log))
+
+    // On the second round of cleaning, the marker from the first transaction should be removed.
+    cleaner.doClean(LogToClean(tp, log, 0L, log.activeSegment.baseOffset), deleteHorizonMs = Long.MaxValue)
+    assertEquals(List(3, 4, 5), offsetsInLog(log))
+    assertEquals(List(2, 3, 4, 5), lastOffsetsPerBatchInLog(log))
+  }
+
   @Test
   def testAbortMarkerRetentionWithEmptyBatch(): Unit = {
     val tp = new TopicPartition("test", 0)
@@ -539,7 +638,7 @@ class LogCleanerTest extends JUnitSuite {
     val appendProducer = appendTransactionalAsLeader(log, producerId, producerEpoch)
 
     appendProducer(Seq(2, 3)) // batch last offset is 1
-    log.appendAsLeader(abortMarker(producerId, producerEpoch), leaderEpoch = 0, isFromClient = false)
+    log.appendAsLeader(abortMarker(producerId, producerEpoch), leaderEpoch = 0, origin = AppendOrigin.Coordinator)
     log.roll()
 
     def assertAbortedTransactionIndexed(): Unit = {
@@ -553,14 +652,14 @@ class LogCleanerTest extends JUnitSuite {
     assertAbortedTransactionIndexed()
 
     // first time through the records are removed
-    var dirtyOffset = cleaner.doClean(LogToClean(tp, log, 0L, 100L), deleteHorizonMs = Long.MaxValue)._1
+    var dirtyOffset = cleaner.doClean(LogToClean(tp, log, 0L, log.activeSegment.baseOffset), deleteHorizonMs = Long.MaxValue)._1
     assertAbortedTransactionIndexed()
     assertEquals(List(), LogTest.keysInLog(log))
     assertEquals(List(2), offsetsInLog(log)) // abort marker is retained
     assertEquals(List(1, 2), lastOffsetsPerBatchInLog(log)) // empty batch is retained
 
     // the empty batch remains if cleaned again because it still holds the last sequence
-    dirtyOffset = cleaner.doClean(LogToClean(tp, log, dirtyOffset, 100L), deleteHorizonMs = Long.MaxValue)._1
+    dirtyOffset = cleaner.doClean(LogToClean(tp, log, dirtyOffset, log.activeSegment.baseOffset), deleteHorizonMs = Long.MaxValue)._1
     assertAbortedTransactionIndexed()
     assertEquals(List(), LogTest.keysInLog(log))
     assertEquals(List(2), offsetsInLog(log)) // abort marker is still retained
@@ -570,13 +669,13 @@ class LogCleanerTest extends JUnitSuite {
     appendProducer(Seq(1))
     log.roll()
 
-    dirtyOffset = cleaner.doClean(LogToClean(tp, log, dirtyOffset, 100L), deleteHorizonMs = Long.MaxValue)._1
+    dirtyOffset = cleaner.doClean(LogToClean(tp, log, dirtyOffset, log.activeSegment.baseOffset), deleteHorizonMs = Long.MaxValue)._1
     assertAbortedTransactionIndexed()
     assertEquals(List(1), LogTest.keysInLog(log))
     assertEquals(List(2, 3), offsetsInLog(log)) // abort marker is not yet gone because we read the empty batch
     assertEquals(List(2, 3), lastOffsetsPerBatchInLog(log)) // but we do not preserve the empty batch
 
-    dirtyOffset = cleaner.doClean(LogToClean(tp, log, dirtyOffset, 100L), deleteHorizonMs = Long.MaxValue)._1
+    dirtyOffset = cleaner.doClean(LogToClean(tp, log, dirtyOffset, log.activeSegment.baseOffset), deleteHorizonMs = Long.MaxValue)._1
     assertEquals(List(1), LogTest.keysInLog(log))
     assertEquals(List(3), offsetsInLog(log)) // abort marker is gone
     assertEquals(List(3), lastOffsetsPerBatchInLog(log))
@@ -589,7 +688,7 @@ class LogCleanerTest extends JUnitSuite {
    * Test log cleaning with logs containing messages larger than default message size
    */
   @Test
-  def testLargeMessage() {
+  def testLargeMessage(): Unit = {
     val largeMessageSize = 1024 * 1024
     // Create cleaner with very small default max message size
     val cleaner = makeCleaner(Int.MaxValue, maxMessageSize=1024)
@@ -611,7 +710,7 @@ class LogCleanerTest extends JUnitSuite {
 
     // clean the log
     val stats = new CleanerStats()
-    cleaner.cleanSegments(log, Seq(log.logSegments.head), map, 0L, stats)
+    cleaner.cleanSegments(log, Seq(log.logSegments.head), map, 0L, stats, new CleanedTransactionMetadata)
     val shouldRemain = LogTest.keysInLog(log).filter(!keys.contains(_))
     assertEquals(shouldRemain, LogTest.keysInLog(log))
   }
@@ -620,11 +719,11 @@ class LogCleanerTest extends JUnitSuite {
    * Test log cleaning with logs containing messages larger than topic's max message size
    */
   @Test
-  def testMessageLargerThanMaxMessageSize() {
+  def testMessageLargerThanMaxMessageSize(): Unit = {
     val (log, offsetMap) = createLogWithMessagesLargerThanMaxSize(largeMessageSize = 1024 * 1024)
 
     val cleaner = makeCleaner(Int.MaxValue, maxMessageSize=1024)
-    cleaner.cleanSegments(log, Seq(log.logSegments.head), offsetMap, 0L, new CleanerStats)
+    cleaner.cleanSegments(log, Seq(log.logSegments.head), offsetMap, 0L, new CleanerStats, new CleanedTransactionMetadata)
     val shouldRemain = LogTest.keysInLog(log).filter(k => !offsetMap.map.containsKey(k.toString))
     assertEquals(shouldRemain, LogTest.keysInLog(log))
   }
@@ -634,7 +733,7 @@ class LogCleanerTest extends JUnitSuite {
    * where header is corrupt
    */
   @Test
-  def testMessageLargerThanMaxMessageSizeWithCorruptHeader() {
+  def testMessageLargerThanMaxMessageSizeWithCorruptHeader(): Unit = {
     val (log, offsetMap) = createLogWithMessagesLargerThanMaxSize(largeMessageSize = 1024 * 1024)
     val file = new RandomAccessFile(log.logSegments.head.log.file, "rw")
     file.seek(Records.MAGIC_OFFSET)
@@ -643,7 +742,7 @@ class LogCleanerTest extends JUnitSuite {
 
     val cleaner = makeCleaner(Int.MaxValue, maxMessageSize=1024)
     intercept[CorruptRecordException] {
-      cleaner.cleanSegments(log, Seq(log.logSegments.head), offsetMap, 0L, new CleanerStats)
+      cleaner.cleanSegments(log, Seq(log.logSegments.head), offsetMap, 0L, new CleanerStats, new CleanedTransactionMetadata)
     }
   }
 
@@ -652,7 +751,7 @@ class LogCleanerTest extends JUnitSuite {
    * where message size is corrupt and larger than bytes available in log segment.
    */
   @Test
-  def testCorruptMessageSizeLargerThanBytesAvailable() {
+  def testCorruptMessageSizeLargerThanBytesAvailable(): Unit = {
     val (log, offsetMap) = createLogWithMessagesLargerThanMaxSize(largeMessageSize = 1024 * 1024)
     val file = new RandomAccessFile(log.logSegments.head.log.file, "rw")
     file.setLength(1024)
@@ -660,7 +759,7 @@ class LogCleanerTest extends JUnitSuite {
 
     val cleaner = makeCleaner(Int.MaxValue, maxMessageSize=1024)
     intercept[CorruptRecordException] {
-      cleaner.cleanSegments(log, Seq(log.logSegments.head), offsetMap, 0L, new CleanerStats)
+      cleaner.cleanSegments(log, Seq(log.logSegments.head), offsetMap, 0L, new CleanerStats, new CleanedTransactionMetadata)
     }
   }
 
@@ -779,7 +878,7 @@ class LogCleanerTest extends JUnitSuite {
 
     appendProducer(Seq(1))
     appendProducer(Seq(2, 3))
-    log.appendAsLeader(abortMarker(producerId, producerEpoch), leaderEpoch = 0, isFromClient = false)
+    log.appendAsLeader(abortMarker(producerId, producerEpoch), leaderEpoch = 0, origin = AppendOrigin.Coordinator)
     log.roll()
 
     cleaner.clean(LogToClean(new TopicPartition("test", 0), log, 0L, log.activeSegment.baseOffset))
@@ -991,7 +1090,8 @@ class LogCleanerTest extends JUnitSuite {
     val map = new FakeOffsetMap(Int.MaxValue)
     keys.foreach(k => map.put(key(k), Long.MaxValue))
     intercept[LogCleaningAbortedException] {
-      cleaner.cleanSegments(log, log.logSegments.take(3).toSeq, map, 0L, new CleanerStats())
+      cleaner.cleanSegments(log, log.logSegments.take(3).toSeq, map, 0L, new CleanerStats(),
+        new CleanedTransactionMetadata)
     }
   }
 
@@ -1153,7 +1253,7 @@ class LogCleanerTest extends JUnitSuite {
     val end = 500
     writeToLog(log, (start until end) zip (start until end))
 
-    def checkRange(map: FakeOffsetMap, start: Int, end: Int) {
+    def checkRange(map: FakeOffsetMap, start: Int, end: Int): Unit = {
       val stats = new CleanerStats()
       cleaner.buildOffsetMap(log, start, end, map, stats)
       val endOffset = map.latestOffset + 1
@@ -1189,7 +1289,7 @@ class LogCleanerTest extends JUnitSuite {
 
     val numSegmentsInitial = log.logSegments.size
     val allKeys = LogTest.keysInLog(log).toList
-    val expectedKeysAfterCleaning = mutable.MutableList[Long]()
+    val expectedKeysAfterCleaning = new mutable.ArrayBuffer[Long]()
 
     // pretend we want to clean every alternate key
     val offsetMap = new FakeOffsetMap(Int.MaxValue)
@@ -1200,7 +1300,8 @@ class LogCleanerTest extends JUnitSuite {
 
     // Try to clean segment with offset overflow. This will trigger log split and the cleaning itself must abort.
     assertThrows[LogCleaningAbortedException] {
-      cleaner.cleanSegments(log, List(segmentWithOverflow), offsetMap, 0L, new CleanerStats())
+      cleaner.cleanSegments(log, List(segmentWithOverflow), offsetMap, 0L, new CleanerStats(),
+        new CleanedTransactionMetadata)
     }
     assertEquals(numSegmentsInitial + 1, log.logSegments.size)
     assertEquals(allKeys, LogTest.keysInLog(log))
@@ -1208,7 +1309,8 @@ class LogCleanerTest extends JUnitSuite {
 
     // Clean each segment now that split is complete.
     for (segmentToClean <- log.logSegments)
-      cleaner.cleanSegments(log, List(segmentToClean), offsetMap, 0L, new CleanerStats())
+      cleaner.cleanSegments(log, List(segmentToClean), offsetMap, 0L, new CleanerStats(),
+        new CleanedTransactionMetadata)
     assertEquals(expectedKeysAfterCleaning, LogTest.keysInLog(log))
     assertFalse(LogTest.hasOffsetOverflow(log))
     log.close()
@@ -1248,7 +1350,8 @@ class LogCleanerTest extends JUnitSuite {
       offsetMap.put(key(k), Long.MaxValue)
 
     // clean the log
-    cleaner.cleanSegments(log, log.logSegments.take(9).toSeq, offsetMap, 0L, new CleanerStats())
+    cleaner.cleanSegments(log, log.logSegments.take(9).toSeq, offsetMap, 0L, new CleanerStats(),
+      new CleanedTransactionMetadata)
     // clear scheduler so that async deletes don't run
     time.scheduler.clear()
     var cleanedKeys = LogTest.keysInLog(log)
@@ -1263,7 +1366,8 @@ class LogCleanerTest extends JUnitSuite {
     log = recoverAndCheck(config, allKeys)
 
     // clean again
-    cleaner.cleanSegments(log, log.logSegments.take(9).toSeq, offsetMap, 0L, new CleanerStats())
+    cleaner.cleanSegments(log, log.logSegments.take(9).toSeq, offsetMap, 0L, new CleanerStats(),
+      new CleanedTransactionMetadata)
     // clear scheduler so that async deletes don't run
     time.scheduler.clear()
     cleanedKeys = LogTest.keysInLog(log)
@@ -1284,7 +1388,8 @@ class LogCleanerTest extends JUnitSuite {
     }
     for (k <- 1 until messageCount by 2)
       offsetMap.put(key(k), Long.MaxValue)
-    cleaner.cleanSegments(log, log.logSegments.take(9).toSeq, offsetMap, 0L, new CleanerStats())
+    cleaner.cleanSegments(log, log.logSegments.take(9).toSeq, offsetMap, 0L, new CleanerStats(),
+      new CleanedTransactionMetadata)
     // clear scheduler so that async deletes don't run
     time.scheduler.clear()
     cleanedKeys = LogTest.keysInLog(log)
@@ -1301,7 +1406,8 @@ class LogCleanerTest extends JUnitSuite {
     }
     for (k <- 1 until messageCount by 2)
       offsetMap.put(key(k), Long.MaxValue)
-    cleaner.cleanSegments(log, log.logSegments.take(9).toSeq, offsetMap, 0L, new CleanerStats())
+    cleaner.cleanSegments(log, log.logSegments.take(9).toSeq, offsetMap, 0L, new CleanerStats(),
+      new CleanedTransactionMetadata)
     // clear scheduler so that async deletes don't run
     time.scheduler.clear()
     cleanedKeys = LogTest.keysInLog(log)
@@ -1367,7 +1473,7 @@ class LogCleanerTest extends JUnitSuite {
    * This test verifies that messages corrupted by KAFKA-4298 are fixed by the cleaner
    */
   @Test
-  def testCleanCorruptMessageSet() {
+  def testCleanCorruptMessageSet(): Unit = {
     val codec = CompressionType.GZIP
 
     val logProps = new Properties()
@@ -1455,6 +1561,49 @@ class LogCleanerTest extends JUnitSuite {
     assertEquals("The tombstone should be retained.", 1, log.logSegments.head.log.batches.iterator.next().lastOffset)
   }
 
+  /**
+   * Verify that the clean is able to move beyond missing offsets records in dirty log
+   */
+  @Test
+  def testCleaningBeyondMissingOffsets(): Unit = {
+    val logProps = new Properties()
+    logProps.put(LogConfig.SegmentBytesProp, 1024*1024: java.lang.Integer)
+    logProps.put(LogConfig.CleanupPolicyProp, LogConfig.Compact)
+    val logConfig = LogConfig(logProps)
+    val cleaner = makeCleaner(Int.MaxValue)
+
+    {
+      val log = makeLog(dir = TestUtils.randomPartitionLogDir(tmpdir), config = logConfig)
+      writeToLog(log, (0 to 9) zip (0 to 9), (0L to 9L))
+      // roll new segment with baseOffset 11, leaving previous with holes in offset range [9,10]
+      log.roll(Some(11L))
+
+      // active segment record
+      log.appendAsFollower(messageWithOffset(1015, 1015, 11L))
+
+      val (nextDirtyOffset, _) = cleaner.clean(LogToClean(log.topicPartition, log, 0L, log.activeSegment.baseOffset, needCompactionNow = true))
+      assertEquals("Cleaning point should pass offset gap", log.activeSegment.baseOffset, nextDirtyOffset)
+    }
+
+
+    {
+      val log = makeLog(dir = TestUtils.randomPartitionLogDir(tmpdir), config = logConfig)
+      writeToLog(log, (0 to 9) zip (0 to 9), (0L to 9L))
+      // roll new segment with baseOffset 15, leaving previous with holes in offset rage [10, 14]
+      log.roll(Some(15L))
+
+      writeToLog(log, (15 to 24) zip (15 to 24), (15L to 24L))
+      // roll new segment with baseOffset 30, leaving previous with holes in offset range [25, 29]
+      log.roll(Some(30L))
+
+      // active segment record
+      log.appendAsFollower(messageWithOffset(1015, 1015, 30L))
+
+      val (nextDirtyOffset, _) = cleaner.clean(LogToClean(log.topicPartition, log, 0L, log.activeSegment.baseOffset, needCompactionNow = true))
+      assertEquals("Cleaning point should pass offset gap in multiple segments", log.activeSegment.baseOffset, nextDirtyOffset)
+    }
+  }
+
   private def writeToLog(log: Log, keysAndValues: Iterable[(Int, Int)], offsetSeq: Iterable[Long]): Iterable[Long] = {
     for(((key, value), offset) <- keysAndValues.zip(offsetSeq))
       yield log.appendAsFollower(messageWithOffset(key, value, offset)).lastOffset
@@ -1521,13 +1670,20 @@ class LogCleanerTest extends JUnitSuite {
       partitionLeaderEpoch, new SimpleRecord(key.toString.getBytes, value.toString.getBytes))
   }
 
-  private def appendTransactionalAsLeader(log: Log, producerId: Long, producerEpoch: Short): Seq[Int] => LogAppendInfo = {
-    appendIdempotentAsLeader(log, producerId, producerEpoch, isTransactional = true)
+  private def appendTransactionalAsLeader(log: Log,
+                                          producerId: Long,
+                                          producerEpoch: Short,
+                                          leaderEpoch: Int = 0,
+                                          origin: AppendOrigin = AppendOrigin.Client): Seq[Int] => LogAppendInfo = {
+    appendIdempotentAsLeader(log, producerId, producerEpoch, isTransactional = true, origin = origin)
   }
 
-  private def appendIdempotentAsLeader(log: Log, producerId: Long,
+  private def appendIdempotentAsLeader(log: Log,
+                                       producerId: Long,
                                        producerEpoch: Short,
-                                       isTransactional: Boolean = false): Seq[Int] => LogAppendInfo = {
+                                       isTransactional: Boolean = false,
+                                       leaderEpoch: Int = 0,
+                                       origin: AppendOrigin = AppendOrigin.Client): Seq[Int] => LogAppendInfo = {
     var sequence = 0
     keys: Seq[Int] => {
       val simpleRecords = keys.map { key =>
@@ -1535,11 +1691,11 @@ class LogCleanerTest extends JUnitSuite {
         new SimpleRecord(time.milliseconds(), keyBytes, keyBytes) // the value doesn't matter since we validate offsets
       }
       val records = if (isTransactional)
-        MemoryRecords.withTransactionalRecords(CompressionType.NONE, producerId, producerEpoch, sequence, simpleRecords: _*)
+        MemoryRecords.withTransactionalRecords(CompressionType.NONE, producerId, producerEpoch, sequence, simpleRecords.toArray: _*)
       else
-        MemoryRecords.withIdempotentRecords(CompressionType.NONE, producerId, producerEpoch, sequence, simpleRecords: _*)
+        MemoryRecords.withIdempotentRecords(CompressionType.NONE, producerId, producerEpoch, sequence, simpleRecords.toArray: _*)
       sequence += simpleRecords.size
-      log.appendAsLeader(records, leaderEpoch = 0)
+      log.appendAsLeader(records, leaderEpoch, origin)
     }
   }
 
diff --git a/core/src/test/scala/unit/kafka/log/LogConcurrencyTest.scala b/core/src/test/scala/unit/kafka/log/LogConcurrencyTest.scala
new file mode 100644
index 0000000000000..0670a4b1aec80
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/log/LogConcurrencyTest.scala
@@ -0,0 +1,184 @@
+/*
+ * 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.log
+
+import java.util.Properties
+import java.util.concurrent.{Callable, Executors}
+
+import kafka.server.{BrokerTopicStats, FetchHighWatermark, LogDirFailureChannel}
+import kafka.utils.{KafkaScheduler, TestUtils}
+import org.apache.kafka.common.record.SimpleRecord
+import org.apache.kafka.common.utils.{Time, Utils}
+import org.junit.Assert._
+import org.junit.{After, Before, Test}
+
+import scala.collection.JavaConverters._
+import scala.collection.mutable.ListBuffer
+import scala.util.Random
+
+class LogConcurrencyTest {
+  private val brokerTopicStats = new BrokerTopicStats
+  private val random = new Random()
+  private val scheduler = new KafkaScheduler(1)
+  private val tmpDir = TestUtils.tempDir()
+  private val logDir = TestUtils.randomPartitionLogDir(tmpDir)
+
+  @Before
+  def setup(): Unit = {
+    scheduler.startup()
+  }
+
+  @After
+  def shutdown(): Unit = {
+    scheduler.shutdown()
+    Utils.delete(tmpDir)
+  }
+
+  @Test
+  def testUncommittedDataNotConsumed(): Unit = {
+    testUncommittedDataNotConsumed(createLog())
+  }
+
+  @Test
+  def testUncommittedDataNotConsumedFrequentSegmentRolls(): Unit = {
+    val logProps = new Properties()
+    logProps.put(LogConfig.SegmentBytesProp, 237: Integer)
+    val logConfig = LogConfig(logProps)
+    testUncommittedDataNotConsumed(createLog(logConfig))
+  }
+
+  def testUncommittedDataNotConsumed(log: Log): Unit = {
+    val executor = Executors.newFixedThreadPool(2)
+    try {
+      val maxOffset = 5000
+      val consumer = new ConsumerTask(log, maxOffset)
+      val appendTask = new LogAppendTask(log, maxOffset)
+
+      val consumerFuture = executor.submit(consumer)
+      val fetcherTaskFuture = executor.submit(appendTask)
+
+      fetcherTaskFuture.get()
+      consumerFuture.get()
+
+      validateConsumedData(log, consumer.consumedBatches)
+    } finally executor.shutdownNow()
+  }
+
+  /**
+   * Simple consumption task which reads the log in ascending order and collects
+   * consumed batches for validation
+   */
+  private class ConsumerTask(log: Log, lastOffset: Int) extends Callable[Unit] {
+    val consumedBatches = ListBuffer.empty[FetchedBatch]
+
+    override def call(): Unit = {
+      var fetchOffset = 0L
+      while (log.highWatermark < lastOffset) {
+        val readInfo = log.read(
+          startOffset = fetchOffset,
+          maxLength = 1,
+          isolation = FetchHighWatermark,
+          minOneMessage = true
+        )
+        readInfo.records.batches().asScala.foreach { batch =>
+          consumedBatches += FetchedBatch(batch.baseOffset, batch.partitionLeaderEpoch)
+          fetchOffset = batch.lastOffset + 1
+        }
+      }
+    }
+  }
+
+  /**
+   * This class simulates basic leader/follower behavior.
+   */
+  private class LogAppendTask(log: Log, lastOffset: Long) extends Callable[Unit] {
+    override def call(): Unit = {
+      var leaderEpoch = 1
+      var isLeader = true
+
+      while (log.highWatermark < lastOffset) {
+        random.nextInt(2) match {
+          case 0 =>
+            val logEndOffsetMetadata = log.logEndOffsetMetadata
+            val logEndOffset = logEndOffsetMetadata.messageOffset
+            val batchSize = random.nextInt(9) + 1
+            val records = (0 to batchSize).map(i => new SimpleRecord(s"$i".getBytes))
+
+            if (isLeader) {
+              log.appendAsLeader(TestUtils.records(records), leaderEpoch)
+              log.maybeIncrementHighWatermark(logEndOffsetMetadata)
+            } else {
+              log.appendAsFollower(TestUtils.records(records,
+                baseOffset = logEndOffset,
+                partitionLeaderEpoch = leaderEpoch))
+              log.updateHighWatermark(logEndOffset)
+            }
+
+          case 1 =>
+            isLeader = !isLeader
+            leaderEpoch += 1
+
+            if (!isLeader) {
+              log.truncateTo(log.highWatermark)
+            }
+        }
+      }
+    }
+  }
+
+  private def createLog(config: LogConfig = LogConfig(new Properties())): Log = {
+    Log(dir = logDir,
+      config = config,
+      logStartOffset = 0L,
+      recoveryPoint = 0L,
+      scheduler = scheduler,
+      brokerTopicStats = brokerTopicStats,
+      time = Time.SYSTEM,
+      maxProducerIdExpirationMs = 60 * 60 * 1000,
+      producerIdExpirationCheckIntervalMs = LogManager.ProducerIdExpirationCheckIntervalMs,
+      logDirFailureChannel = new LogDirFailureChannel(10))
+  }
+
+  private def validateConsumedData(log: Log, consumedBatches: Iterable[FetchedBatch]): Unit = {
+    val iter = consumedBatches.iterator
+    log.logSegments.foreach { segment =>
+      segment.log.batches.asScala.foreach { batch =>
+        if (iter.hasNext) {
+          val consumedBatch = iter.next()
+          try {
+            assertEquals("Consumed batch with unexpected leader epoch",
+              batch.partitionLeaderEpoch, consumedBatch.epoch)
+            assertEquals("Consumed batch with unexpected base offset",
+              batch.baseOffset, consumedBatch.baseOffset)
+          } catch {
+            case t: Throwable =>
+              throw new AssertionError(s"Consumed batch $consumedBatch " +
+                s"does not match next expected batch in log $batch", t)
+          }
+        }
+      }
+    }
+  }
+
+  private case class FetchedBatch(baseOffset: Long, epoch: Int) {
+    override def toString: String = {
+      s"FetchedBatch(baseOffset=$baseOffset, epoch=$epoch)"
+    }
+  }
+
+}
diff --git a/core/src/test/scala/unit/kafka/log/LogConfigTest.scala b/core/src/test/scala/unit/kafka/log/LogConfigTest.scala
index 66702d683abcb..556810eff6fbf 100644
--- a/core/src/test/scala/unit/kafka/log/LogConfigTest.scala
+++ b/core/src/test/scala/unit/kafka/log/LogConfigTest.scala
@@ -19,24 +19,26 @@ package kafka.log
 
 import java.util.Properties
 
-import kafka.server.{ThrottledReplicaListValidator, KafkaConfig, KafkaServer}
+import kafka.server.{KafkaConfig, KafkaServer, ThrottledReplicaListValidator}
 import kafka.utils.TestUtils
-import org.apache.kafka.common.config.ConfigException
+import org.apache.kafka.common.config.ConfigDef.Importance.MEDIUM
+import org.apache.kafka.common.config.ConfigDef.Type.INT
+import org.apache.kafka.common.config.{ConfigException, TopicConfig}
 import org.junit.{Assert, Test}
 import org.junit.Assert._
 import org.scalatest.Assertions._
 
 class LogConfigTest {
 
-  /** 
-   * This test verifies that KafkaConfig object initialization does not depend on 
-   * LogConfig initialization. Bad things happen due to static initialization 
-   * order dependencies. For example, LogConfig.configDef ends up adding null 
-   * values in serverDefaultConfigNames. This test ensures that the mapping of 
+  /**
+   * This test verifies that KafkaConfig object initialization does not depend on
+   * LogConfig initialization. Bad things happen due to static initialization
+   * order dependencies. For example, LogConfig.configDef ends up adding null
+   * values in serverDefaultConfigNames. This test ensures that the mapping of
    * keys from LogConfig to KafkaConfig are not missing values.
    */
   @Test
-  def ensureNoStaticInitializationOrderDependency() {
+  def ensureNoStaticInitializationOrderDependency(): Unit = {
     // Access any KafkaConfig val to load KafkaConfig object before LogConfig.
     assertTrue(KafkaConfig.LogRetentionTimeMillisProp != null)
     assertTrue(LogConfig.configNames.forall { config =>
@@ -46,7 +48,7 @@ class LogConfigTest {
   }
 
   @Test
-  def testKafkaConfigToProps() {
+  def testKafkaConfigToProps(): Unit = {
     val millisInHour = 60L * 60L * 1000L
     val kafkaProps = TestUtils.createBrokerConfig(nodeId = 0, zkConnect = "")
     kafkaProps.put(KafkaConfig.LogRollTimeHoursProp, "2")
@@ -61,14 +63,14 @@ class LogConfigTest {
   }
 
   @Test
-  def testFromPropsEmpty() {
+  def testFromPropsEmpty(): Unit = {
     val p = new Properties()
     val config = LogConfig(p)
     Assert.assertEquals(LogConfig(), config)
   }
 
   @Test
-  def testFromPropsInvalid() {
+  def testFromPropsInvalid(): Unit = {
     LogConfig.configNames.foreach(name => name match {
       case LogConfig.UncleanLeaderElectionEnableProp => assertPropertyInvalid(name, "not a boolean")
       case LogConfig.RetentionBytesProp => assertPropertyInvalid(name, "not_a_number")
@@ -82,7 +84,17 @@ class LogConfigTest {
   }
 
   @Test
-  def shouldValidateThrottledReplicasConfig() {
+  def testInvalidCompactionLagConfig(): Unit = {
+    val props = new Properties
+    props.setProperty(LogConfig.MaxCompactionLagMsProp, "100")
+    props.setProperty(LogConfig.MinCompactionLagMsProp, "200")
+    intercept[Exception] {
+      LogConfig.validate(props)
+    }
+  }
+
+  @Test
+  def shouldValidateThrottledReplicasConfig(): Unit = {
     assertTrue(isValid("*"))
     assertTrue(isValid("* "))
     assertTrue(isValid(""))
@@ -103,6 +115,54 @@ class LogConfigTest {
     assertFalse(isValid("100:0,10 :   "))
   }
 
+  /* Sanity check that toHtmlTable produces one of the expected configs */
+  @Test
+  def testToHtmlTable(): Unit = {
+    val html = LogConfig.configDefCopy.toHtmlTable
+    val expectedConfig = "file.delete.delay.ms"
+    assertTrue(s"Could not find `$expectedConfig` in:\n $html", html.contains(expectedConfig))
+  }
+
+  /* Sanity check that toHtml produces one of the expected configs */
+  @Test
+  def testToHtml(): Unit = {
+    val html = LogConfig.configDefCopy.toHtml
+    val expectedConfig = "
            1. file.delete.delay.ms"
+    assertTrue(s"Could not find `$expectedConfig` in:\n $html", html.contains(expectedConfig))
+  }
+
+  /* Sanity check that toEnrichedRst produces one of the expected configs */
+  @Test
+  def testToEnrichedRst(): Unit = {
+    val rst = LogConfig.configDefCopy.toEnrichedRst
+    val expectedConfig = "``file.delete.delay.ms``"
+    assertTrue(s"Could not find `$expectedConfig` in:\n $rst", rst.contains(expectedConfig))
+  }
+
+  /* Sanity check that toEnrichedRst produces one of the expected configs */
+  @Test
+  def testToRst(): Unit = {
+    val rst = LogConfig.configDefCopy.toRst
+    val expectedConfig = "``file.delete.delay.ms``"
+    assertTrue(s"Could not find `$expectedConfig` in:\n $rst", rst.contains(expectedConfig))
+  }
+
+  @Test
+  def testGetConfigValue(): Unit = {
+    // Add a config that doesn't set the `serverDefaultConfigName`
+    val configDef = LogConfig.configDefCopy
+    val configNameWithNoServerMapping = "log.foo"
+    configDef.define(configNameWithNoServerMapping, INT, 1, MEDIUM, s"$configNameWithNoServerMapping doc")
+
+    val deleteDelayKey = configDef.configKeys.get(TopicConfig.FILE_DELETE_DELAY_MS_CONFIG)
+    val deleteDelayServerDefault = configDef.getConfigValue(deleteDelayKey, LogConfig.ServerDefaultHeaderName)
+    assertEquals(KafkaConfig.LogDeleteDelayMsProp, deleteDelayServerDefault)
+
+    val keyWithNoServerMapping = configDef.configKeys.get(configNameWithNoServerMapping)
+    val nullServerDefault = configDef.getConfigValue(keyWithNoServerMapping, LogConfig.ServerDefaultHeaderName)
+    assertNull(nullServerDefault)
+  }
+
   private def isValid(configValue: String): Boolean = {
     try {
       ThrottledReplicaListValidator.ensureValidString("", configValue)
@@ -112,7 +172,7 @@ class LogConfigTest {
     }
   }
 
-  private def assertPropertyInvalid(name: String, values: AnyRef*) {
+  private def assertPropertyInvalid(name: String, values: AnyRef*): Unit = {
     values.foreach((value) => {
       val props = new Properties
       props.setProperty(name, value.toString)
diff --git a/core/src/test/scala/unit/kafka/log/LogManagerTest.scala b/core/src/test/scala/unit/kafka/log/LogManagerTest.scala
index b3ecd23a9cf58..15139b5f22210 100755
--- a/core/src/test/scala/unit/kafka/log/LogManagerTest.scala
+++ b/core/src/test/scala/unit/kafka/log/LogManagerTest.scala
@@ -18,16 +18,24 @@
 package kafka.log
 
 import java.io._
-import java.util.Properties
+import java.util.{Collections, Properties}
 
-import kafka.server.FetchDataInfo
+import kafka.server.{FetchDataInfo, FetchLogEnd}
 import kafka.server.checkpoints.OffsetCheckpointFile
 import kafka.utils._
-import org.apache.kafka.common.{KafkaException, TopicPartition}
 import org.apache.kafka.common.errors.OffsetOutOfRangeException
 import org.apache.kafka.common.utils.Utils
+import org.apache.kafka.common.{KafkaException, TopicPartition}
+import org.easymock.EasyMock
 import org.junit.Assert._
 import org.junit.{After, Before, Test}
+import org.mockito.ArgumentMatchers.any
+import org.mockito.Mockito.{doAnswer, spy}
+import org.mockito.invocation.InvocationOnMock
+import org.mockito.stubbing.Answer
+
+import scala.collection.mutable
+import scala.util.{Failure, Try}
 
 class LogManagerTest {
 
@@ -46,14 +54,14 @@ class LogManagerTest {
   val veryLargeLogFlushInterval = 10000000L
 
   @Before
-  def setUp() {
+  def setUp(): Unit = {
     logDir = TestUtils.tempDir()
     logManager = createLogManager()
     logManager.startup()
   }
 
   @After
-  def tearDown() {
+  def tearDown(): Unit = {
     if (logManager != null)
       logManager.shutdown()
     Utils.delete(logDir)
@@ -65,8 +73,8 @@ class LogManagerTest {
    * Test that getOrCreateLog on a non-existent log creates a new log and that we can append to the new log.
    */
   @Test
-  def testCreateLog() {
-    val log = logManager.getOrCreateLog(new TopicPartition(name, 0), logConfig)
+  def testCreateLog(): Unit = {
+    val log = logManager.getOrCreateLog(new TopicPartition(name, 0), () => logConfig)
     assertEquals(1, logManager.liveLogDirs.size)
 
     val logFile = new File(logDir, name + "-0")
@@ -79,7 +87,7 @@ class LogManagerTest {
    * The LogManager is configured with one invalid log directory which should be marked as offline.
    */
   @Test
-  def testCreateLogWithInvalidLogDir() {
+  def testCreateLogWithInvalidLogDir(): Unit = {
     // Configure the log dir with the Nul character as the path, which causes dir.getCanonicalPath() to throw an
     // IOException. This simulates the scenario where the disk is not properly mounted (which is hard to achieve in
     // a unit test)
@@ -89,17 +97,58 @@ class LogManagerTest {
     logManager = createLogManager(dirs)
     logManager.startup()
 
-    val log = logManager.getOrCreateLog(new TopicPartition(name, 0), logConfig, isNew = true)
+    val log = logManager.getOrCreateLog(new TopicPartition(name, 0), () => logConfig, isNew = true)
     val logFile = new File(logDir, name + "-0")
     assertTrue(logFile.exists)
     log.appendAsLeader(TestUtils.singletonRecords("test".getBytes()), leaderEpoch = 0)
   }
 
+  @Test
+  def testCreateLogWithLogDirFallback(): Unit = {
+    // Configure a number of directories one level deeper in logDir,
+    // so they all get cleaned up in tearDown().
+    val dirs = (0 to 4)
+      .map(_.toString)
+      .map(logDir.toPath.resolve(_).toFile)
+
+    // Create a new LogManager with the configured directories and an overridden createLogDirectory.
+    logManager.shutdown()
+    logManager = spy(createLogManager(dirs))
+    val brokenDirs = mutable.Set[File]()
+    doAnswer(new Answer[Try[File]] {
+      override def answer(invocation: InvocationOnMock): Try[File] = {
+        // The first half of directories tried will fail, the rest goes through.
+        val logDir = invocation.getArgument[File](0)
+        if (brokenDirs.contains(logDir) || brokenDirs.size < dirs.length / 2) {
+          brokenDirs.add(logDir)
+          Failure(new Throwable("broken dir"))
+        } else {
+          invocation.callRealMethod().asInstanceOf[Try[File]]
+        }
+      }
+    }).when(logManager).createLogDirectory(any(), any())
+    logManager.startup()
+
+    // Request creating a new log.
+    // LogManager should try using all configured log directories until one succeeds.
+    logManager.getOrCreateLog(new TopicPartition(name, 0), () => logConfig, isNew = true)
+
+    // Verify that half the directories were considered broken,
+    assertEquals(dirs.length / 2, brokenDirs.size)
+
+    // and that exactly one log file was created,
+    val containsLogFile: File => Boolean = dir => new File(dir, name + "-0").exists()
+    assertEquals("More than one log file created", 1, dirs.count(containsLogFile))
+
+    // and that it wasn't created in one of the broken directories.
+    assertFalse(brokenDirs.exists(containsLogFile))
+  }
+
   /**
    * Test that get on a non-existent returns None and no log is created.
    */
   @Test
-  def testGetNonExistentLog() {
+  def testGetNonExistentLog(): Unit = {
     val log = logManager.getLog(new TopicPartition(name, 0))
     assertEquals("No log should be found.", None, log)
     val logFile = new File(logDir, name + "-0")
@@ -110,8 +159,8 @@ class LogManagerTest {
    * Test time-based log cleanup. First append messages, then set the time into the future and run cleanup.
    */
   @Test
-  def testCleanupExpiredSegments() {
-    val log = logManager.getOrCreateLog(new TopicPartition(name, 0), logConfig)
+  def testCleanupExpiredSegments(): Unit = {
+    val log = logManager.getOrCreateLog(new TopicPartition(name, 0), () => logConfig)
     var offset = 0L
     for(_ <- 0 until 200) {
       val set = TestUtils.singletonRecords("test".getBytes())
@@ -119,7 +168,7 @@ class LogManagerTest {
       offset = info.lastOffset
     }
     assertTrue("There should be more than one segment now.", log.numberOfSegments > 1)
-    log.onHighWatermarkIncremented(log.logEndOffset)
+    log.updateHighWatermark(log.logEndOffset)
 
     log.logSegments.foreach(_.log.file.setLastModified(time.milliseconds))
 
@@ -150,7 +199,7 @@ class LogManagerTest {
    * Test size-based cleanup. Append messages, then run cleanup and check that segments are deleted.
    */
   @Test
-  def testCleanupSegmentsToMaintainSize() {
+  def testCleanupSegmentsToMaintainSize(): Unit = {
     val setSize = TestUtils.singletonRecords("test".getBytes()).sizeInBytes
     logManager.shutdown()
     val logProps = new Properties()
@@ -162,7 +211,7 @@ class LogManagerTest {
     logManager.startup()
 
     // create a log
-    val log = logManager.getOrCreateLog(new TopicPartition(name, 0), config)
+    val log = logManager.getOrCreateLog(new TopicPartition(name, 0), () => config)
     var offset = 0L
 
     // add a bunch of messages that should be larger than the retentionSize
@@ -173,7 +222,7 @@ class LogManagerTest {
       offset = info.firstOffset.get
     }
 
-    log.onHighWatermarkIncremented(log.logEndOffset)
+    log.updateHighWatermark(log.logEndOffset)
     assertEquals("Check we have the expected number of segments.", numMessages * setSize / config.segmentSize, log.numberOfSegments)
 
     // this cleanup shouldn't find any expired segments but should delete some to reduce size
@@ -200,7 +249,7 @@ class LogManagerTest {
     * LogCleaner.CleanerThread handles all logs where compaction is enabled.
     */
   @Test
-  def testDoesntCleanLogsWithCompactDeletePolicy() {
+  def testDoesntCleanLogsWithCompactDeletePolicy(): Unit = {
     testDoesntCleanLogs(LogConfig.Compact + "," + LogConfig.Delete)
   }
 
@@ -209,14 +258,14 @@ class LogManagerTest {
     * LogCleaner.CleanerThread handles all logs where compaction is enabled.
     */
   @Test
-  def testDoesntCleanLogsWithCompactPolicy() {
+  def testDoesntCleanLogsWithCompactPolicy(): Unit = {
     testDoesntCleanLogs(LogConfig.Compact)
   }
 
-  private def testDoesntCleanLogs(policy: String) {
+  private def testDoesntCleanLogs(policy: String): Unit = {
     val logProps = new Properties()
     logProps.put(LogConfig.CleanupPolicyProp, policy)
-    val log = logManager.getOrCreateLog(new TopicPartition(name, 0), LogConfig.fromProps(logConfig.originals, logProps))
+    val log = logManager.getOrCreateLog(new TopicPartition(name, 0), () => LogConfig.fromProps(logConfig.originals, logProps))
     var offset = 0L
     for (_ <- 0 until 200) {
       val set = TestUtils.singletonRecords("test".getBytes(), key="test".getBytes())
@@ -237,7 +286,7 @@ class LogManagerTest {
    * Test that flush is invoked by the background scheduler thread.
    */
   @Test
-  def testTimeBasedFlush() {
+  def testTimeBasedFlush(): Unit = {
     logManager.shutdown()
     val logProps = new Properties()
     logProps.put(LogConfig.FlushMsProp, 1000: java.lang.Integer)
@@ -245,7 +294,7 @@ class LogManagerTest {
 
     logManager = createLogManager()
     logManager.startup()
-    val log = logManager.getOrCreateLog(new TopicPartition(name, 0), config)
+    val log = logManager.getOrCreateLog(new TopicPartition(name, 0), () => config)
     val lastFlush = log.lastFlushTime
     for (_ <- 0 until 200) {
       val set = TestUtils.singletonRecords("test".getBytes())
@@ -259,7 +308,7 @@ class LogManagerTest {
    * Test that new logs that are created are assigned to the least loaded log directory
    */
   @Test
-  def testLeastLoadedAssignment() {
+  def testLeastLoadedAssignment(): Unit = {
     // create a log manager with multiple data directories
     val dirs = Seq(TestUtils.tempDir(),
                      TestUtils.tempDir(),
@@ -269,7 +318,7 @@ class LogManagerTest {
 
     // verify that logs are always assigned to the least loaded partition
     for(partition <- 0 until 20) {
-      logManager.getOrCreateLog(new TopicPartition("test", partition), logConfig)
+      logManager.getOrCreateLog(new TopicPartition("test", partition), () => logConfig)
       assertEquals("We should have created the right number of logs", partition + 1, logManager.allLogs.size)
       val counts = logManager.allLogs.groupBy(_.dir.getParent).values.map(_.size)
       assertTrue("Load should balance evenly", counts.max <= counts.min + 1)
@@ -280,7 +329,7 @@ class LogManagerTest {
    * Test that it is not possible to open two log managers using the same data directory
    */
   @Test
-  def testTwoLogManagersUsingSameDirFails() {
+  def testTwoLogManagersUsingSameDirFails(): Unit = {
     try {
       createLogManager()
       fail("Should not be able to create a second log manager instance with the same data directory")
@@ -293,7 +342,7 @@ class LogManagerTest {
    * Test that recovery points are correctly written out to disk
    */
   @Test
-  def testCheckpointRecoveryPoints() {
+  def testCheckpointRecoveryPoints(): Unit = {
     verifyCheckpointRecovery(Seq(new TopicPartition("test-a", 1), new TopicPartition("test-b", 1)), logManager, logDir)
   }
 
@@ -301,7 +350,7 @@ class LogManagerTest {
    * Test that recovery points directory checking works with trailing slash
    */
   @Test
-  def testRecoveryDirectoryMappingWithTrailingSlash() {
+  def testRecoveryDirectoryMappingWithTrailingSlash(): Unit = {
     logManager.shutdown()
     logManager = TestUtils.createLogManager(logDirs = Seq(new File(TestUtils.tempDir().getAbsolutePath + File.separator)))
     logManager.startup()
@@ -312,15 +361,15 @@ class LogManagerTest {
    * Test that recovery points directory checking works with relative directory
    */
   @Test
-  def testRecoveryDirectoryMappingWithRelativeDirectory() {
+  def testRecoveryDirectoryMappingWithRelativeDirectory(): Unit = {
     logManager.shutdown()
     logManager = createLogManager(Seq(new File("data", logDir.getName).getAbsoluteFile))
     logManager.startup()
     verifyCheckpointRecovery(Seq(new TopicPartition("test-a", 1)), logManager, logManager.liveLogDirs.head)
   }
 
-  private def verifyCheckpointRecovery(topicPartitions: Seq[TopicPartition], logManager: LogManager, logDir: File) {
-    val logs = topicPartitions.map(logManager.getOrCreateLog(_, logConfig))
+  private def verifyCheckpointRecovery(topicPartitions: Seq[TopicPartition], logManager: LogManager, logDir: File): Unit = {
+    val logs = topicPartitions.map(logManager.getOrCreateLog(_, () => logConfig))
     logs.foreach { log =>
       for (_ <- 0 until 50)
         log.appendAsLeader(TestUtils.singletonRecords("test".getBytes()), leaderEpoch = 0)
@@ -345,8 +394,8 @@ class LogManagerTest {
   }
 
   @Test
-  def testFileReferencesAfterAsyncDelete() {
-    val log = logManager.getOrCreateLog(new TopicPartition(name, 0), logConfig)
+  def testFileReferencesAfterAsyncDelete(): Unit = {
+    val log = logManager.getOrCreateLog(new TopicPartition(name, 0), () => logConfig)
     val activeSegment = log.activeSegment
     val logName = activeSegment.log.file.getName
     val indexName = activeSegment.offsetIndex.file.getName
@@ -381,7 +430,14 @@ class LogManagerTest {
   }
 
   @Test
-  def testCheckpointForOnlyAffectedLogs() {
+  def testCreateAndDeleteOverlyLongTopic(): Unit = {
+    val invalidTopicName = String.join("", Collections.nCopies(253, "x"))
+    logManager.getOrCreateLog(new TopicPartition(invalidTopicName, 0), () => logConfig)
+    logManager.asyncDelete(new TopicPartition(invalidTopicName, 0))
+  }
+
+  @Test
+  def testCheckpointForOnlyAffectedLogs(): Unit = {
     val tps = Seq(
       new TopicPartition("test-a", 0),
       new TopicPartition("test-a", 1),
@@ -389,7 +445,7 @@ class LogManagerTest {
       new TopicPartition("test-b", 0),
       new TopicPartition("test-b", 1))
 
-    val allLogs = tps.map(logManager.getOrCreateLog(_, logConfig))
+    val allLogs = tps.map(logManager.getOrCreateLog(_, () => logConfig))
     allLogs.foreach { log =>
       for (_ <- 0 until 50)
         log.appendAsLeader(TestUtils.singletonRecords("test".getBytes), leaderEpoch = 0)
@@ -410,7 +466,104 @@ class LogManagerTest {
   }
 
   private def readLog(log: Log, offset: Long, maxLength: Int = 1024): FetchDataInfo = {
-    log.read(offset, maxLength, maxOffset = None, minOneMessage = true, includeAbortedTxns = false)
+    log.read(offset, maxLength, isolation = FetchLogEnd, minOneMessage = true)
+  }
+
+  /**
+   * Test when a configuration of a topic is updated while its log is getting initialized,
+   * the config is refreshed when log initialization is finished.
+   */
+  @Test
+  def testTopicConfigChangeUpdatesLogConfig(): Unit = {
+    val testTopicOne = "test-topic-one"
+    val testTopicTwo = "test-topic-two"
+    val testTopicOnePartition: TopicPartition = new TopicPartition(testTopicOne, 1)
+    val testTopicTwoPartition: TopicPartition = new TopicPartition(testTopicTwo, 1)
+    val mockLog: Log = EasyMock.mock(classOf[Log])
+
+    logManager.initializingLog(testTopicOnePartition)
+    logManager.initializingLog(testTopicTwoPartition)
+
+    logManager.topicConfigUpdated(testTopicOne)
+
+    val logConfig: LogConfig = null
+    var configUpdated = false
+    logManager.finishedInitializingLog(testTopicOnePartition, Some(mockLog), () => {
+      configUpdated = true
+      logConfig
+    })
+    assertTrue(configUpdated)
+
+    var configNotUpdated = true
+    logManager.finishedInitializingLog(testTopicTwoPartition, Some(mockLog), () => {
+      configNotUpdated = false
+      logConfig
+    })
+    assertTrue(configNotUpdated)
+  }
+
+  /**
+   * Test if an error occurs when creating log, log manager removes corresponding
+   * topic partition from the list of initializing partitions.
+   */
+  @Test
+  def testConfigChangeGetsCleanedUp(): Unit = {
+    val testTopicPartition: TopicPartition = new TopicPartition("test-topic", 1)
+
+    logManager.initializingLog(testTopicPartition)
+
+    val logConfig: LogConfig = null
+    var configUpdateNotCalled = true
+    logManager.finishedInitializingLog(testTopicPartition, None, () => {
+      configUpdateNotCalled = false
+      logConfig
+    })
+
+    assertTrue(logManager.partitionsInitializing.isEmpty)
+    assertTrue(configUpdateNotCalled)
+  }
+
+  /**
+   * Test when a broker configuration change happens all logs in process of initialization
+   * pick up latest config when finished with initialization.
+   */
+  @Test
+  def testBrokerConfigChangeDeliveredToAllLogs(): Unit = {
+    val testTopicOne = "test-topic-one"
+    val testTopicTwo = "test-topic-two"
+    val testTopicOnePartition: TopicPartition = new TopicPartition(testTopicOne, 1)
+    val testTopicTwoPartition: TopicPartition = new TopicPartition(testTopicTwo, 1)
+    val mockLog: Log = EasyMock.mock(classOf[Log])
+
+    logManager.initializingLog(testTopicOnePartition)
+    logManager.initializingLog(testTopicTwoPartition)
+
+    logManager.brokerConfigUpdated()
+
+    val logConfig: LogConfig = null
+    var totalChanges = 0
+    logManager.finishedInitializingLog(testTopicOnePartition, Some(mockLog), () => {
+      totalChanges += 1
+      logConfig
+    })
+    logManager.finishedInitializingLog(testTopicTwoPartition, Some(mockLog), () => {
+      totalChanges += 1
+      logConfig
+    })
+
+    assertEquals(2, totalChanges)
   }
 
+  /**
+   * Test even if no log is getting initialized, if config change events are delivered
+   * things continue to work correctly. This test should not throw.
+   *
+   * This makes sure that events can be delivered even when no log is getting initialized.
+   */
+  @Test
+  def testConfigChangesWithNoLogGettingInitialized(): Unit = {
+    logManager.brokerConfigUpdated()
+    logManager.topicConfigUpdated("test-topic")
+    assertTrue(logManager.partitionsInitializing.isEmpty)
+  }
 }
diff --git a/core/src/test/scala/unit/kafka/log/LogSegmentTest.scala b/core/src/test/scala/unit/kafka/log/LogSegmentTest.scala
index 0ebc0eff0b046..a29e7e5415c3a 100644
--- a/core/src/test/scala/unit/kafka/log/LogSegmentTest.scala
+++ b/core/src/test/scala/unit/kafka/log/LogSegmentTest.scala
@@ -56,7 +56,7 @@ class LogSegmentTest {
   }
 
   @After
-  def teardown() {
+  def teardown(): Unit = {
     segments.foreach(_.close())
     Utils.delete(logDir)
   }
@@ -65,9 +65,9 @@ class LogSegmentTest {
    * A read on an empty log segment should return null
    */
   @Test
-  def testReadOnEmptySegment() {
+  def testReadOnEmptySegment(): Unit = {
     val seg = createSegment(40)
-    val read = seg.read(startOffset = 40, maxSize = 300, maxOffset = None)
+    val read = seg.read(startOffset = 40, maxSize = 300)
     assertNull("Read beyond the last offset in the segment should be null", read)
   }
 
@@ -76,41 +76,23 @@ class LogSegmentTest {
    * beginning with the first message in the segment
    */
   @Test
-  def testReadBeforeFirstOffset() {
+  def testReadBeforeFirstOffset(): Unit = {
     val seg = createSegment(40)
     val ms = records(50, "hello", "there", "little", "bee")
     seg.append(53, RecordBatch.NO_TIMESTAMP, -1L, ms)
-    val read = seg.read(startOffset = 41, maxSize = 300, maxOffset = None).records
+    val read = seg.read(startOffset = 41, maxSize = 300).records
     checkEquals(ms.records.iterator, read.records.iterator)
   }
 
-  /**
-   * If we set the startOffset and maxOffset for the read to be the same value
-   * we should get only the first message in the log
-   */
-  @Test
-  def testMaxOffset() {
-    val baseOffset = 50
-    val seg = createSegment(baseOffset)
-    val ms = records(baseOffset, "hello", "there", "beautiful")
-    seg.append(52, RecordBatch.NO_TIMESTAMP, -1L, ms)
-    def validate(offset: Long) =
-      assertEquals(ms.records.asScala.filter(_.offset == offset).toList,
-                   seg.read(startOffset = offset, maxSize = 1024, maxOffset = Some(offset+1)).records.records.asScala.toList)
-    validate(50)
-    validate(51)
-    validate(52)
-  }
-
   /**
    * If we read from an offset beyond the last offset in the segment we should get null
    */
   @Test
-  def testReadAfterLast() {
+  def testReadAfterLast(): Unit = {
     val seg = createSegment(40)
     val ms = records(50, "hello", "there")
     seg.append(51, RecordBatch.NO_TIMESTAMP, -1L, ms)
-    val read = seg.read(startOffset = 52, maxSize = 200, maxOffset = None)
+    val read = seg.read(startOffset = 52, maxSize = 200)
     assertNull("Read beyond the last offset in the segment should give null", read)
   }
 
@@ -119,13 +101,13 @@ class LogSegmentTest {
    * with the least offset greater than the given startOffset.
    */
   @Test
-  def testReadFromGap() {
+  def testReadFromGap(): Unit = {
     val seg = createSegment(40)
     val ms = records(50, "hello", "there")
     seg.append(51, RecordBatch.NO_TIMESTAMP, -1L, ms)
     val ms2 = records(60, "alpha", "beta")
     seg.append(61, RecordBatch.NO_TIMESTAMP, -1L, ms2)
-    val read = seg.read(startOffset = 55, maxSize = 200, maxOffset = None)
+    val read = seg.read(startOffset = 55, maxSize = 200)
     checkEquals(ms2.records.iterator, read.records.records.iterator)
   }
 
@@ -134,7 +116,7 @@ class LogSegmentTest {
    * the first but not the second message.
    */
   @Test
-  def testTruncate() {
+  def testTruncate(): Unit = {
     val seg = createSegment(40)
     var offset = 40
     for (_ <- 0 until 30) {
@@ -143,11 +125,11 @@ class LogSegmentTest {
       val ms2 = records(offset + 1, "hello")
       seg.append(offset + 1, RecordBatch.NO_TIMESTAMP, -1L, ms2)
       // check that we can read back both messages
-      val read = seg.read(offset, None, 10000)
+      val read = seg.read(offset, 10000)
       assertEquals(List(ms1.records.iterator.next(), ms2.records.iterator.next()), read.records.records.asScala.toList)
       // now truncate off the last message
       seg.truncateTo(offset + 1)
-      val read2 = seg.read(offset, None, 10000)
+      val read2 = seg.read(offset, 10000)
       assertEquals(1, read2.records.records.asScala.size)
       checkEquals(ms1.records.iterator, read2.records.records.iterator)
       offset += 1
@@ -155,7 +137,7 @@ class LogSegmentTest {
   }
 
   @Test
-  def testTruncateEmptySegment() {
+  def testTruncateEmptySegment(): Unit = {
     // This tests the scenario in which the follower truncates to an empty segment. In this
     // case we must ensure that the index is resized so that the log segment is not mistakenly
     // rolled due to a full index
@@ -163,6 +145,9 @@ class LogSegmentTest {
     val maxSegmentMs = 300000
     val time = new MockTime
     val seg = createSegment(0, time = time)
+    // Force load indexes before closing the segment
+    seg.timeIndex
+    seg.offsetIndex
     seg.close()
 
     val reopened = createSegment(0, time = time)
@@ -193,7 +178,7 @@ class LogSegmentTest {
   }
 
   @Test
-  def testReloadLargestTimestampAndNextOffsetAfterTruncation() {
+  def testReloadLargestTimestampAndNextOffsetAfterTruncation(): Unit = {
     val numMessages = 30
     val seg = createSegment(40, 2 * records(0, "hello").sizeInBytes - 1)
     var offset = 40
@@ -216,7 +201,7 @@ class LogSegmentTest {
    * Test truncating the whole segment, and check that we can reappend with the original offset.
    */
   @Test
-  def testTruncateFull() {
+  def testTruncateFull(): Unit = {
     // test the case where we fully truncate the log
     val time = new MockTime
     val seg = createSegment(40, time = time)
@@ -230,7 +215,7 @@ class LogSegmentTest {
     assertEquals(0, seg.timeWaitedForRoll(time.milliseconds(), RecordBatch.NO_TIMESTAMP))
     assertFalse(seg.timeIndex.isFull)
     assertFalse(seg.offsetIndex.isFull)
-    assertNull("Segment should be empty.", seg.read(0, None, 1024))
+    assertNull("Segment should be empty.", seg.read(0, 1024))
 
     seg.append(41, RecordBatch.NO_TIMESTAMP, -1L, records(40, "hello", "there"))
   }
@@ -239,7 +224,7 @@ class LogSegmentTest {
    * Append messages with timestamp and search message by timestamp.
    */
   @Test
-  def testFindOffsetByTimestamp() {
+  def testFindOffsetByTimestamp(): Unit = {
     val messageSize = records(0, s"msg00").sizeInBytes
     val seg = createSegment(40, messageSize * 2 - 1)
     // Produce some messages
@@ -265,7 +250,7 @@ class LogSegmentTest {
    * Test that offsets are assigned sequentially and that the nextOffset variable is incremented
    */
   @Test
-  def testNextOffsetCalculation() {
+  def testNextOffsetCalculation(): Unit = {
     val seg = createSegment(40)
     assertEquals(40, seg.readNextOffset)
     seg.append(52, RecordBatch.NO_TIMESTAMP, -1L, records(50, "hello", "there", "you"))
@@ -276,15 +261,29 @@ class LogSegmentTest {
    * Test that we can change the file suffixes for the log and index files
    */
   @Test
-  def testChangeFileSuffixes() {
+  def testChangeFileSuffixes(): Unit = {
     val seg = createSegment(40)
     val logFile = seg.log.file
     val indexFile = seg.lazyOffsetIndex.file
+    val timeIndexFile = seg.lazyTimeIndex.file
+    // Ensure that files for offset and time indices have not been created eagerly.
+    assertFalse(seg.lazyOffsetIndex.file.exists)
+    assertFalse(seg.lazyTimeIndex.file.exists)
     seg.changeFileSuffixes("", ".deleted")
+    // Ensure that attempt to change suffixes for non-existing offset and time indices does not create new files.
+    assertFalse(seg.lazyOffsetIndex.file.exists)
+    assertFalse(seg.lazyTimeIndex.file.exists)
+    // Ensure that file names are updated accordingly.
     assertEquals(logFile.getAbsolutePath + ".deleted", seg.log.file.getAbsolutePath)
     assertEquals(indexFile.getAbsolutePath + ".deleted", seg.lazyOffsetIndex.file.getAbsolutePath)
+    assertEquals(timeIndexFile.getAbsolutePath + ".deleted", seg.lazyTimeIndex.file.getAbsolutePath)
     assertTrue(seg.log.file.exists)
+    // Ensure lazy creation of offset index file upon accessing it.
+    seg.lazyOffsetIndex.get
     assertTrue(seg.lazyOffsetIndex.file.exists)
+    // Ensure lazy creation of time index file upon accessing it.
+    seg.lazyTimeIndex.get
+    assertTrue(seg.lazyTimeIndex.file.exists)
   }
 
   /**
@@ -292,15 +291,17 @@ class LogSegmentTest {
    * and recover the segment, the entries should all be readable.
    */
   @Test
-  def testRecoveryFixesCorruptIndex() {
+  def testRecoveryFixesCorruptIndex(): Unit = {
     val seg = createSegment(0)
     for(i <- 0 until 100)
       seg.append(i, RecordBatch.NO_TIMESTAMP, -1L, records(i, i.toString))
     val indexFile = seg.lazyOffsetIndex.file
     TestUtils.writeNonsenseToFile(indexFile, 5, indexFile.length.toInt)
     seg.recover(new ProducerStateManager(topicPartition, logDir))
-    for(i <- 0 until 100)
-      assertEquals(i, seg.read(i, Some(i + 1), 1024).records.records.iterator.next().offset)
+    for(i <- 0 until 100) {
+      val records = seg.read(i, 1, minOneMessage = true).records.records
+      assertEquals(i, records.iterator.next().offset)
+    }
   }
 
   @Test
@@ -352,7 +353,8 @@ class LogSegmentTest {
     // recover again, but this time assuming the transaction from pid2 began on a previous segment
     stateManager = new ProducerStateManager(topicPartition, logDir)
     stateManager.loadProducerEntry(new ProducerStateEntry(pid2,
-      mutable.Queue[BatchMetadata](BatchMetadata(10, 10L, 5, RecordBatch.NO_TIMESTAMP)), producerEpoch, 0, Some(75L)))
+      mutable.Queue[BatchMetadata](BatchMetadata(10, 10L, 5, RecordBatch.NO_TIMESTAMP)), producerEpoch,
+      0, RecordBatch.NO_TIMESTAMP, Some(75L)))
     segment.recover(stateManager)
     assertEquals(108L, stateManager.mapEndOffset)
 
@@ -381,7 +383,7 @@ class LogSegmentTest {
    * and recover the segment, the entries should all be readable.
    */
   @Test
-  def testRecoveryFixesCorruptTimeIndex() {
+  def testRecoveryFixesCorruptTimeIndex(): Unit = {
     val seg = createSegment(0)
     for(i <- 0 until 100)
       seg.append(i, i * 10, i, records(i, i.toString))
@@ -399,7 +401,7 @@ class LogSegmentTest {
    * Randomly corrupt a log a number of times and attempt recovery.
    */
   @Test
-  def testRecoveryWithCorruptMessage() {
+  def testRecoveryWithCorruptMessage(): Unit = {
     val messagesAppended = 20
     for (_ <- 0 until 10) {
       val seg = createSegment(0)
@@ -433,19 +435,19 @@ class LogSegmentTest {
 
   /* create a segment with   pre allocate, put message to it and verify */
   @Test
-  def testCreateWithInitFileSizeAppendMessage() {
+  def testCreateWithInitFileSizeAppendMessage(): Unit = {
     val seg = createSegment(40, false, 512*1024*1024, true)
     val ms = records(50, "hello", "there")
     seg.append(51, RecordBatch.NO_TIMESTAMP, -1L, ms)
     val ms2 = records(60, "alpha", "beta")
     seg.append(61, RecordBatch.NO_TIMESTAMP, -1L, ms2)
-    val read = seg.read(startOffset = 55, maxSize = 200, maxOffset = None)
+    val read = seg.read(startOffset = 55, maxSize = 200)
     checkEquals(ms2.records.iterator, read.records.records.iterator)
   }
 
   /* create a segment with   pre allocate and clearly shut down*/
   @Test
-  def testCreateWithInitFileSizeClearShutdown() {
+  def testCreateWithInitFileSizeClearShutdown(): Unit = {
     val tempDir = TestUtils.tempDir()
     val logConfig = LogConfig(Map(
       LogConfig.IndexIntervalBytesProp -> 10,
@@ -460,7 +462,7 @@ class LogSegmentTest {
     seg.append(51, RecordBatch.NO_TIMESTAMP, -1L, ms)
     val ms2 = records(60, "alpha", "beta")
     seg.append(61, RecordBatch.NO_TIMESTAMP, -1L, ms2)
-    val read = seg.read(startOffset = 55, maxSize = 200, maxOffset = None)
+    val read = seg.read(startOffset = 55, maxSize = 200)
     checkEquals(ms2.records.iterator, read.records.records.iterator)
     val oldSize = seg.log.sizeInBytes()
     val oldPosition = seg.log.channel.position
@@ -474,7 +476,7 @@ class LogSegmentTest {
       initFileSize = 512 * 1024 * 1024, preallocate = true)
     segments += segReopen
 
-    val readAgain = segReopen.read(startOffset = 55, maxSize = 200, maxOffset = None)
+    val readAgain = segReopen.read(startOffset = 55, maxSize = 200)
     checkEquals(ms2.records.iterator, readAgain.records.records.iterator)
     val size = segReopen.log.sizeInBytes()
     val position = segReopen.log.channel.position
@@ -485,7 +487,7 @@ class LogSegmentTest {
   }
 
   @Test
-  def shouldTruncateEvenIfOffsetPointsToAGapInTheLog() {
+  def shouldTruncateEvenIfOffsetPointsToAGapInTheLog(): Unit = {
     val seg = createSegment(40)
     val offset = 40
 
@@ -503,7 +505,7 @@ class LogSegmentTest {
     seg.truncateTo(offset + 1)
 
     //Then we should still truncate the record that was present (i.e. offset + 3 is gone)
-    val log = seg.read(offset, None, 10000)
+    val log = seg.read(offset, 10000)
     assertEquals(offset, log.records.batches.iterator.next().baseOffset())
     assertEquals(1, log.records.batches.asScala.size)
   }
diff --git a/core/src/test/scala/unit/kafka/log/LogTest.scala b/core/src/test/scala/unit/kafka/log/LogTest.scala
index 6b6dd7c8a9de5..892997e09d088 100755
--- a/core/src/test/scala/unit/kafka/log/LogTest.scala
+++ b/core/src/test/scala/unit/kafka/log/LogTest.scala
@@ -20,16 +20,19 @@ package kafka.log
 import java.io._
 import java.nio.ByteBuffer
 import java.nio.file.{Files, Paths}
-import java.util.{Optional, Properties}
+import java.util.concurrent.{Callable, Executors}
+import java.util.regex.Pattern
+import java.util.{Collections, Optional, Properties}
 
+import com.yammer.metrics.Metrics
 import kafka.api.{ApiVersion, KAFKA_0_11_0_IV0}
-import kafka.common.{OffsetsOutOfOrderException, UnexpectedAppendOffsetException}
+import kafka.common.{OffsetsOutOfOrderException, RecordValidationException, UnexpectedAppendOffsetException}
 import kafka.log.Log.DeleteDirSuffix
 import kafka.server.checkpoints.LeaderEpochCheckpointFile
 import kafka.server.epoch.{EpochEntry, LeaderEpochFileCache}
-import kafka.server.{BrokerTopicStats, FetchDataInfo, KafkaConfig, LogDirFailureChannel}
+import kafka.server.{BrokerTopicStats, FetchDataInfo, FetchHighWatermark, FetchIsolation, FetchLogEnd, FetchTxnCommitted, KafkaConfig, LogDirFailureChannel, LogOffsetMetadata}
 import kafka.utils._
-import org.apache.kafka.common.{KafkaException, TopicPartition}
+import org.apache.kafka.common.{InvalidRecordException, KafkaException, TopicPartition}
 import org.apache.kafka.common.errors._
 import org.apache.kafka.common.record.FileRecords.TimestampAndOffset
 import org.apache.kafka.common.record.MemoryRecords.RecordFilter
@@ -43,31 +46,28 @@ import org.junit.Assert._
 import org.junit.{After, Before, Test}
 import org.scalatest.Assertions
 
-import scala.collection.Iterable
+import scala.collection.{Iterable, mutable}
 import scala.collection.JavaConverters._
-import scala.collection.mutable.{ArrayBuffer, ListBuffer}
+import scala.collection.mutable.ListBuffer
 import org.scalatest.Assertions.{assertThrows, intercept, withClue}
 
 class LogTest {
-  var config: KafkaConfig = null
   val brokerTopicStats = new BrokerTopicStats
   val tmpDir = TestUtils.tempDir()
   val logDir = TestUtils.randomPartitionLogDir(tmpDir)
   val mockTime = new MockTime()
+  def metricsKeySet = Metrics.defaultRegistry.allMetrics.keySet.asScala
 
   @Before
-  def setUp() {
-    val props = TestUtils.createBrokerConfig(0, "127.0.0.1:1", port = -1)
-    config = KafkaConfig.fromProps(props)
-  }
+  def setUp(): Unit = {}
 
   @After
-  def tearDown() {
+  def tearDown(): Unit = {
     brokerTopicStats.close()
     Utils.delete(tmpDir)
   }
 
-  def createEmptyLogs(dir: File, offsets: Int*) {
+  def createEmptyLogs(dir: File, offsets: Int*): Unit = {
     for(offset <- offsets) {
       Log.logFile(dir, offset).createNewFile()
       Log.offsetIndexFile(dir, offset).createNewFile()
@@ -75,7 +75,247 @@ class LogTest {
   }
 
   @Test
-  def testOffsetFromFile() {
+  def testHighWatermarkMetadataUpdatedAfterSegmentRoll(): Unit = {
+    val logConfig = LogTest.createLogConfig(segmentBytes = 1024 * 1024)
+    val log = createLog(logDir, logConfig)
+
+    def assertFetchSizeAndOffsets(fetchOffset: Long,
+                                  expectedSize: Int,
+                                  expectedOffsets: Seq[Long]): Unit = {
+      val readInfo = log.read(
+        startOffset = fetchOffset,
+        maxLength = 2048,
+        isolation = FetchHighWatermark,
+        minOneMessage = false)
+      assertEquals(expectedSize, readInfo.records.sizeInBytes)
+      assertEquals(expectedOffsets, readInfo.records.records.asScala.map(_.offset))
+    }
+
+    val records = TestUtils.records(List(
+      new SimpleRecord(mockTime.milliseconds, "a".getBytes, "value".getBytes),
+      new SimpleRecord(mockTime.milliseconds, "b".getBytes, "value".getBytes),
+      new SimpleRecord(mockTime.milliseconds, "c".getBytes, "value".getBytes)
+    ))
+
+    log.appendAsLeader(records, leaderEpoch = 0)
+    assertFetchSizeAndOffsets(fetchOffset = 0L, 0, Seq())
+
+    log.maybeIncrementHighWatermark(log.logEndOffsetMetadata)
+    assertFetchSizeAndOffsets(fetchOffset = 0L, records.sizeInBytes, Seq(0, 1, 2))
+
+    log.roll()
+    assertFetchSizeAndOffsets(fetchOffset = 0L, records.sizeInBytes, Seq(0, 1, 2))
+
+    log.appendAsLeader(records, leaderEpoch = 0)
+    assertFetchSizeAndOffsets(fetchOffset = 3L, 0, Seq())
+  }
+
+  @Test
+  def testHighWatermarkMaintenance(): Unit = {
+    val logConfig = LogTest.createLogConfig(segmentBytes = 1024 * 1024)
+    val log = createLog(logDir, logConfig)
+    val leaderEpoch = 0
+
+    def records(offset: Long): MemoryRecords = TestUtils.records(List(
+      new SimpleRecord(mockTime.milliseconds, "a".getBytes, "value".getBytes),
+      new SimpleRecord(mockTime.milliseconds, "b".getBytes, "value".getBytes),
+      new SimpleRecord(mockTime.milliseconds, "c".getBytes, "value".getBytes)
+    ), baseOffset = offset, partitionLeaderEpoch= leaderEpoch)
+
+    def assertHighWatermark(offset: Long): Unit = {
+      assertEquals(offset, log.highWatermark)
+      assertValidLogOffsetMetadata(log, log.fetchOffsetSnapshot.highWatermark)
+    }
+
+    // High watermark initialized to 0
+    assertHighWatermark(0L)
+
+    // High watermark not changed by append
+    log.appendAsLeader(records(0), leaderEpoch)
+    assertHighWatermark(0L)
+
+    // Update high watermark as leader
+    log.maybeIncrementHighWatermark(LogOffsetMetadata(1L))
+    assertHighWatermark(1L)
+
+    // Cannot update past the log end offset
+    log.updateHighWatermark(5L)
+    assertHighWatermark(3L)
+
+    // Update high watermark as follower
+    log.appendAsFollower(records(3L))
+    log.updateHighWatermark(6L)
+    assertHighWatermark(6L)
+
+    // High watermark should be adjusted by truncation
+    log.truncateTo(3L)
+    assertHighWatermark(3L)
+
+    log.appendAsLeader(records(0L), leaderEpoch = 0)
+    assertHighWatermark(3L)
+    assertEquals(6L, log.logEndOffset)
+    assertEquals(0L, log.logStartOffset)
+
+    // Full truncation should also reset high watermark
+    log.truncateFullyAndStartAt(4L)
+    assertEquals(4L, log.logEndOffset)
+    assertEquals(4L, log.logStartOffset)
+    assertHighWatermark(4L)
+  }
+
+  private def assertNonEmptyFetch(log: Log, offset: Long, isolation: FetchIsolation): Unit = {
+    val readInfo = log.read(startOffset = offset,
+      maxLength = Int.MaxValue,
+      isolation = isolation,
+      minOneMessage = true)
+
+    assertFalse(readInfo.firstEntryIncomplete)
+    assertTrue(readInfo.records.sizeInBytes > 0)
+
+    val upperBoundOffset = isolation match {
+      case FetchLogEnd => log.logEndOffset
+      case FetchHighWatermark => log.highWatermark
+      case FetchTxnCommitted => log.lastStableOffset
+    }
+
+    for (record <- readInfo.records.records.asScala)
+      assertTrue(record.offset < upperBoundOffset)
+
+    assertEquals(offset, readInfo.fetchOffsetMetadata.messageOffset)
+    assertValidLogOffsetMetadata(log, readInfo.fetchOffsetMetadata)
+  }
+
+  private def assertEmptyFetch(log: Log, offset: Long, isolation: FetchIsolation): Unit = {
+    val readInfo = log.read(startOffset = offset,
+      maxLength = Int.MaxValue,
+      isolation = isolation,
+      minOneMessage = true)
+    assertFalse(readInfo.firstEntryIncomplete)
+    assertEquals(0, readInfo.records.sizeInBytes)
+    assertEquals(offset, readInfo.fetchOffsetMetadata.messageOffset)
+    assertValidLogOffsetMetadata(log, readInfo.fetchOffsetMetadata)
+  }
+
+  @Test
+  def testFetchUpToLogEndOffset(): Unit = {
+    val logConfig = LogTest.createLogConfig(segmentBytes = 1024 * 1024)
+    val log = createLog(logDir, logConfig)
+
+    log.appendAsLeader(TestUtils.records(List(
+      new SimpleRecord("0".getBytes),
+      new SimpleRecord("1".getBytes),
+      new SimpleRecord("2".getBytes)
+    )), leaderEpoch = 0)
+    log.appendAsLeader(TestUtils.records(List(
+      new SimpleRecord("3".getBytes),
+      new SimpleRecord("4".getBytes)
+    )), leaderEpoch = 0)
+
+    (log.logStartOffset until log.logEndOffset).foreach { offset =>
+      assertNonEmptyFetch(log, offset, FetchLogEnd)
+    }
+  }
+
+  @Test
+  def testFetchUpToHighWatermark(): Unit = {
+    val logConfig = LogTest.createLogConfig(segmentBytes = 1024 * 1024)
+    val log = createLog(logDir, logConfig)
+
+    log.appendAsLeader(TestUtils.records(List(
+      new SimpleRecord("0".getBytes),
+      new SimpleRecord("1".getBytes),
+      new SimpleRecord("2".getBytes)
+    )), leaderEpoch = 0)
+    log.appendAsLeader(TestUtils.records(List(
+      new SimpleRecord("3".getBytes),
+      new SimpleRecord("4".getBytes)
+    )), leaderEpoch = 0)
+
+    def assertHighWatermarkBoundedFetches(): Unit = {
+      (log.logStartOffset until log.highWatermark).foreach { offset =>
+        assertNonEmptyFetch(log, offset, FetchHighWatermark)
+      }
+
+      (log.highWatermark to log.logEndOffset).foreach { offset =>
+        assertEmptyFetch(log, offset, FetchHighWatermark)
+      }
+    }
+
+    assertHighWatermarkBoundedFetches()
+
+    log.updateHighWatermark(3L)
+    assertHighWatermarkBoundedFetches()
+
+    log.updateHighWatermark(5L)
+    assertHighWatermarkBoundedFetches()
+  }
+
+  @Test
+  def testFetchUpToLastStableOffset(): Unit = {
+    val logConfig = LogTest.createLogConfig(segmentBytes = 1024 * 1024)
+    val log = createLog(logDir, logConfig)
+    val epoch = 0.toShort
+
+    val producerId1 = 1L
+    val producerId2 = 2L
+
+    val appendProducer1 = appendTransactionalAsLeader(log, producerId1, epoch)
+    val appendProducer2 = appendTransactionalAsLeader(log, producerId2, epoch)
+
+    appendProducer1(5)
+    appendNonTransactionalAsLeader(log, 3)
+    appendProducer2(2)
+    appendProducer1(4)
+    appendNonTransactionalAsLeader(log, 2)
+    appendProducer1(10)
+
+    def assertLsoBoundedFetches(): Unit = {
+      (log.logStartOffset until log.lastStableOffset).foreach { offset =>
+        assertNonEmptyFetch(log, offset, FetchTxnCommitted)
+      }
+
+      (log.lastStableOffset to log.logEndOffset).foreach { offset =>
+        assertEmptyFetch(log, offset, FetchTxnCommitted)
+      }
+    }
+
+    assertLsoBoundedFetches()
+
+    log.updateHighWatermark(log.logEndOffset)
+    assertLsoBoundedFetches()
+
+    appendEndTxnMarkerAsLeader(log, producerId1, epoch, ControlRecordType.COMMIT)
+    assertEquals(0L, log.lastStableOffset)
+
+    log.updateHighWatermark(log.logEndOffset)
+    assertEquals(8L, log.lastStableOffset)
+    assertLsoBoundedFetches()
+
+    appendEndTxnMarkerAsLeader(log, producerId2, epoch, ControlRecordType.ABORT)
+    assertEquals(8L, log.lastStableOffset)
+
+    log.updateHighWatermark(log.logEndOffset)
+    assertEquals(log.logEndOffset, log.lastStableOffset)
+    assertLsoBoundedFetches()
+  }
+
+  @Test
+  def testLogDeleteDirName(): Unit = {
+    val name1 = Log.logDeleteDirName(new TopicPartition("foo", 3))
+    assertTrue(name1.length <= 255)
+    assertTrue(Pattern.compile("foo-3\\.[0-9a-z]{32}-delete").matcher(name1).matches())
+    assertTrue(Log.DeleteDirPattern.matcher(name1).matches())
+    assertFalse(Log.FutureDirPattern.matcher(name1).matches())
+    val name2 = Log.logDeleteDirName(
+      new TopicPartition("n" + String.join("", Collections.nCopies(248, "o")), 5))
+    assertEquals(255, name2.length)
+    assertTrue(Pattern.compile("n[o]{212}-5\\.[0-9a-z]{32}-delete").matcher(name2).matches())
+    assertTrue(Log.DeleteDirPattern.matcher(name2).matches())
+    assertFalse(Log.FutureDirPattern.matcher(name2).matches())
+  }
+
+  @Test
+  def testOffsetFromFile(): Unit = {
     val offset = 23423423L
 
     val logFile = Log.logFile(tmpDir, offset)
@@ -96,7 +336,7 @@ class LogTest {
    * using the mock clock to force the log to roll and checks the number of segments.
    */
   @Test
-  def testTimeBasedLogRoll() {
+  def testTimeBasedLogRoll(): Unit = {
     def createRecords = TestUtils.singletonRecords("test".getBytes)
     val logConfig = LogTest.createLogConfig(segmentMs = 1 * 60 * 60L)
 
@@ -144,7 +384,7 @@ class LogTest {
   }
 
   @Test
-  def testRollSegmentThatAlreadyExists() {
+  def testRollSegmentThatAlreadyExists(): Unit = {
     val logConfig = LogTest.createLogConfig(segmentMs = 1 * 60 * 60L)
 
     // create a log
@@ -170,8 +410,8 @@ class LogTest {
     log.appendAsFollower(records2)
 
     assertEquals("Expect two records in the log", 2, log.logEndOffset)
-    assertEquals(0, readLog(log, 0, 100, Some(1)).records.batches.iterator.next().lastOffset)
-    assertEquals(1, readLog(log, 1, 100, Some(2)).records.batches.iterator.next().lastOffset)
+    assertEquals(0, readLog(log, 0, 1).records.batches.iterator.next().lastOffset)
+    assertEquals(1, readLog(log, 1, 1).records.batches.iterator.next().lastOffset)
 
     // roll so that active segment is empty
     log.roll()
@@ -185,7 +425,7 @@ class LogTest {
       baseOffset = 2L, partitionLeaderEpoch = 0)
     log.appendAsFollower(records3)
     assertTrue(log.activeSegment.offsetIndex.maxEntries > 1)
-    assertEquals(2, readLog(log, 2, 100, Some(3)).records.batches.iterator.next().lastOffset)
+    assertEquals(2, readLog(log, 2, 1).records.batches.iterator.next().lastOffset)
     assertEquals("Expect two segments.", 2, log.numberOfSegments)
   }
 
@@ -225,6 +465,41 @@ class LogTest {
     assertEquals(101L, log.logEndOffset)
   }
 
+  /**
+   * Test the values returned by the logSegments call
+   */
+  @Test
+  def testLogSegmentsCallCorrect(): Unit = {
+    // Create 3 segments and make sure we get the right values from various logSegments calls.
+    def createRecords = TestUtils.singletonRecords(value = "test".getBytes, timestamp = mockTime.milliseconds)
+    def getSegmentOffsets(log :Log, from: Long, to: Long) = log.logSegments(from, to).map { _.baseOffset }
+    val setSize = createRecords.sizeInBytes
+    val msgPerSeg = 10
+    val segmentSize = msgPerSeg * setSize  // each segment will be 10 messages
+    // create a log
+    val logConfig = LogTest.createLogConfig(segmentBytes = segmentSize)
+    val log = createLog(logDir, logConfig)
+    assertEquals("There should be exactly 1 segment.", 1, log.numberOfSegments)
+
+    // segments expire in size
+    for (_ <- 1 to (2 * msgPerSeg + 2))
+      log.appendAsLeader(createRecords, leaderEpoch = 0)
+    assertEquals("There should be exactly 3 segments.", 3, log.numberOfSegments)
+
+    // from == to should always be null
+    assertEquals(List.empty[LogSegment], getSegmentOffsets(log, 10, 10))
+    assertEquals(List.empty[LogSegment], getSegmentOffsets(log, 15, 15))
+
+    assertEquals(List[Long](0, 10, 20), getSegmentOffsets(log, 0, 21))
+
+    assertEquals(List[Long](0), getSegmentOffsets(log, 1, 5))
+    assertEquals(List[Long](10, 20), getSegmentOffsets(log, 13, 21))
+    assertEquals(List[Long](10), getSegmentOffsets(log, 13, 17))
+
+    // from < to is bad
+    assertThrows[IllegalArgumentException]({ log.logSegments(10, 0) })
+  }
+
   @Test
   def testInitializationOfProducerSnapshotsUpgradePath(): Unit = {
     // simulate the upgrade path by creating a new log with several segments, deleting the
@@ -268,6 +543,39 @@ class LogTest {
     log.close()
   }
 
+
+  @Test
+  def testRecoverAfterNonMonotonicCoordinatorEpochWrite(): Unit = {
+    // Due to KAFKA-9144, we may encounter a coordinator epoch which goes backwards.
+    // This test case verifies that recovery logic relaxes validation in this case and
+    // just takes the latest write.
+
+    val producerId = 1L
+    val coordinatorEpoch = 5
+    val logConfig = LogTest.createLogConfig(segmentBytes = 1024 * 1024 * 5)
+    var log = createLog(logDir, logConfig)
+    val epoch = 0.toShort
+
+    val firstAppendTimestamp = mockTime.milliseconds()
+    appendEndTxnMarkerAsLeader(log, producerId, epoch, ControlRecordType.ABORT,
+      timestamp = firstAppendTimestamp, coordinatorEpoch = coordinatorEpoch)
+    assertEquals(firstAppendTimestamp, log.producerStateManager.lastEntry(producerId).get.lastTimestamp)
+
+    mockTime.sleep(log.maxProducerIdExpirationMs)
+    assertEquals(None, log.producerStateManager.lastEntry(producerId))
+
+    val secondAppendTimestamp = mockTime.milliseconds()
+    appendEndTxnMarkerAsLeader(log, producerId, epoch, ControlRecordType.ABORT,
+      timestamp = secondAppendTimestamp, coordinatorEpoch = coordinatorEpoch - 1)
+
+    log.close()
+
+    // Force recovery by setting the recoveryPoint to the log start
+    log = createLog(logDir, logConfig, recoveryPoint = 0L)
+    assertEquals(secondAppendTimestamp, log.producerStateManager.lastEntry(producerId).get.lastTimestamp)
+    log.close()
+  }
+
   @Test
   def testProducerSnapshotsRecoveryAfterUncleanShutdownV1(): Unit = {
     testProducerSnapshotsRecoveryAfterUncleanShutdown(ApiVersion.minSupportedFor(RecordVersion.V1).version)
@@ -278,6 +586,131 @@ class LogTest {
     testProducerSnapshotsRecoveryAfterUncleanShutdown(ApiVersion.latestVersion.version)
   }
 
+  @Test
+  def testLogReinitializeAfterManualDelete(): Unit = {
+    val logConfig = LogTest.createLogConfig()
+    // simulate a case where log data does not exist but the start offset is non-zero
+    val log = createLog(logDir, logConfig, logStartOffset = 500)
+    assertEquals(500, log.logStartOffset)
+    assertEquals(500, log.logEndOffset)
+  }
+
+  @Test
+  def testLogEndLessThanStartAfterReopen(): Unit = {
+    val logConfig = LogTest.createLogConfig()
+    var log = createLog(logDir, logConfig)
+    for (i <- 0 until 5) {
+      val record = new SimpleRecord(mockTime.milliseconds, i.toString.getBytes)
+      log.appendAsLeader(TestUtils.records(List(record)), leaderEpoch = 0)
+      log.roll()
+    }
+    assertEquals(6, log.logSegments.size)
+
+    // Increment the log start offset
+    val startOffset = 4
+    log.updateHighWatermark(log.logEndOffset)
+    log.maybeIncrementLogStartOffset(startOffset)
+    assertTrue(log.logEndOffset > log.logStartOffset)
+
+    // Append garbage to a segment below the current log start offset
+    val segmentToForceTruncation = log.logSegments.take(2).last
+    val bw = new BufferedWriter(new FileWriter(segmentToForceTruncation.log.file))
+    bw.write("corruptRecord")
+    bw.close()
+    log.close()
+
+    // Reopen the log. This will cause truncate the segment to which we appended garbage and delete all other segments.
+    // All remaining segments will be lower than the current log start offset, which will force deletion of all segments
+    // and recreation of a single, active segment starting at logStartOffset.
+    log = createLog(logDir, logConfig, logStartOffset = startOffset)
+    assertEquals(1, log.logSegments.size)
+    assertEquals(startOffset, log.logStartOffset)
+    assertEquals(startOffset, log.logEndOffset)
+  }
+
+  @Test
+  def testNonActiveSegmentsFrom(): Unit = {
+    val logConfig = LogTest.createLogConfig()
+    val log = createLog(logDir, logConfig)
+
+    for (i <- 0 until 5) {
+      val record = new SimpleRecord(mockTime.milliseconds, i.toString.getBytes)
+      log.appendAsLeader(TestUtils.records(List(record)), leaderEpoch = 0)
+      log.roll()
+    }
+
+    def nonActiveBaseOffsetsFrom(startOffset: Long): Seq[Long] = {
+      log.nonActiveLogSegmentsFrom(startOffset).map(_.baseOffset).toSeq
+    }
+
+    assertEquals(5L, log.activeSegment.baseOffset)
+    assertEquals(0 until 5, nonActiveBaseOffsetsFrom(0L))
+    assertEquals(Seq.empty, nonActiveBaseOffsetsFrom(5L))
+    assertEquals(2 until 5, nonActiveBaseOffsetsFrom(2L))
+    assertEquals(Seq.empty, nonActiveBaseOffsetsFrom(6L))
+  }
+
+  @Test
+  def testInconsistentLogSegmentRange(): Unit = {
+    val logConfig = LogTest.createLogConfig()
+    val log = createLog(logDir, logConfig)
+
+    for (i <- 0 until 5) {
+      val record = new SimpleRecord(mockTime.milliseconds, i.toString.getBytes)
+      log.appendAsLeader(TestUtils.records(List(record)), leaderEpoch = 0)
+      log.roll()
+    }
+
+    assertThrows[IllegalArgumentException] {
+      log.logSegments(5, 1)
+    }
+  }
+
+  @Test
+  def testLogDelete(): Unit = {
+    val logConfig = LogTest.createLogConfig()
+    val log = createLog(logDir, logConfig)
+
+    for (i <- 0 to 100) {
+      val record = new SimpleRecord(mockTime.milliseconds, i.toString.getBytes)
+      log.appendAsLeader(TestUtils.records(List(record)), leaderEpoch = 0)
+      log.roll()
+    }
+
+    assertTrue(log.logSegments.size > 0)
+    assertFalse(logDir.listFiles.isEmpty)
+
+    // delete the log
+    log.delete()
+
+    assertEquals(0, log.logSegments.size)
+    assertFalse(logDir.exists)
+  }
+
+  /**
+   * Test that "PeriodicProducerExpirationCheck" scheduled task gets canceled after log
+   * is deleted.
+   */
+  @Test
+  def testProducerExpireCheckAfterDelete(): Unit = {
+    val scheduler = new KafkaScheduler(1)
+    try {
+      scheduler.startup()
+      val logConfig = LogTest.createLogConfig()
+      val log = createLog(logDir, logConfig, scheduler = scheduler)
+
+      val producerExpireCheck = log.producerExpireCheck
+      assertTrue("producerExpireCheck isn't as part of scheduled tasks",
+        scheduler.taskRunning(producerExpireCheck))
+
+      log.delete()
+      assertFalse("producerExpireCheck is part of scheduled tasks even after log deletion",
+        scheduler.taskRunning(producerExpireCheck))
+    } finally {
+      scheduler.shutdown();
+    }
+  }
+
   private def testProducerSnapshotsRecoveryAfterUncleanShutdown(messageFormatVersion: String): Unit = {
     val logConfig = LogTest.createLogConfig(segmentBytes = 64 * 10, messageFormatVersion = messageFormatVersion)
     var log = createLog(logDir, logConfig)
@@ -296,21 +729,21 @@ class LogTest {
     // 1 segment. We collect the data before closing the log.
     val offsetForSegmentAfterRecoveryPoint = segmentOffsets(segmentOffsets.size - 3)
     val offsetForRecoveryPointSegment = segmentOffsets(segmentOffsets.size - 4)
-    val (segOffsetsBeforeRecovery, segOffsetsAfterRecovery) = segmentOffsets.partition(_ < offsetForRecoveryPointSegment)
+    val (segOffsetsBeforeRecovery, segOffsetsAfterRecovery) = segmentOffsets.toSet.partition(_ < offsetForRecoveryPointSegment)
     val recoveryPoint = offsetForRecoveryPointSegment + 1
     assertTrue(recoveryPoint < offsetForSegmentAfterRecoveryPoint)
     log.close()
 
-    val segmentsWithReads = ArrayBuffer[LogSegment]()
-    val recoveredSegments = ArrayBuffer[LogSegment]()
-    val expectedSegmentsWithReads = ArrayBuffer[Long]()
-    val expectedSnapshotOffsets = ArrayBuffer[Long]()
+    val segmentsWithReads = mutable.Set[LogSegment]()
+    val recoveredSegments = mutable.Set[LogSegment]()
+    val expectedSegmentsWithReads = mutable.Set[Long]()
+    val expectedSnapshotOffsets = mutable.Set[Long]()
 
     if (logConfig.messageFormatVersion < KAFKA_0_11_0_IV0) {
       expectedSegmentsWithReads += activeSegmentOffset
       expectedSnapshotOffsets ++= log.logSegments.map(_.baseOffset).toVector.takeRight(2) :+ log.logEndOffset
     } else {
-      expectedSegmentsWithReads ++= segOffsetsBeforeRecovery ++ Seq(activeSegmentOffset)
+      expectedSegmentsWithReads ++= segOffsetsBeforeRecovery ++ Set(activeSegmentOffset)
       expectedSnapshotOffsets ++= log.logSegments.map(_.baseOffset).toVector.takeRight(4) :+ log.logEndOffset
     }
 
@@ -328,10 +761,9 @@ class LogTest {
           val wrapper = new LogSegment(segment.log, segment.lazyOffsetIndex, segment.lazyTimeIndex, segment.txnIndex, segment.baseOffset,
             segment.indexIntervalBytes, segment.rollJitterMs, mockTime) {
 
-            override def read(startOffset: Long, maxOffset: Option[Long], maxSize: Int, maxPosition: Long,
-                              minOneMessage: Boolean): FetchDataInfo = {
+            override def read(startOffset: Long, maxSize: Int, maxPosition: Long, minOneMessage: Boolean): FetchDataInfo = {
               segmentsWithReads += this
-              super.read(startOffset, maxOffset, maxSize, maxPosition, minOneMessage)
+              super.read(startOffset, maxSize, maxPosition, minOneMessage)
             }
 
             override def recover(producerStateManager: ProducerStateManager,
@@ -351,7 +783,7 @@ class LogTest {
     // We will reload all segments because the recovery point is behind the producer snapshot files (pre KAFKA-5829 behaviour)
     assertEquals(expectedSegmentsWithReads, segmentsWithReads.map(_.baseOffset))
     assertEquals(segOffsetsAfterRecovery, recoveredSegments.map(_.baseOffset))
-    assertEquals(expectedSnapshotOffsets, listProducerSnapshotOffsets)
+    assertEquals(expectedSnapshotOffsets, listProducerSnapshotOffsets.toSet)
     log.close()
     segmentsWithReads.clear()
     recoveredSegments.clear()
@@ -360,9 +792,9 @@ class LogTest {
     // avoid reading all segments
     ProducerStateManager.deleteSnapshotsBefore(logDir, offsetForRecoveryPointSegment)
     log = createLogWithInterceptedReads(recoveryPoint = recoveryPoint)
-    assertEquals(Seq(activeSegmentOffset), segmentsWithReads.map(_.baseOffset))
+    assertEquals(Set(activeSegmentOffset), segmentsWithReads.map(_.baseOffset))
     assertEquals(segOffsetsAfterRecovery, recoveredSegments.map(_.baseOffset))
-    assertEquals(expectedSnapshotOffsets, listProducerSnapshotOffsets)
+    assertEquals(expectedSnapshotOffsets, listProducerSnapshotOffsets.toSet)
 
     // Verify that we keep 2 snapshot files if we checkpoint the log end offset
     log.deleteSnapshotsAfterRecoveryPointCheckpoint()
@@ -385,7 +817,7 @@ class LogTest {
   }
 
   @Test
-  def testProducerIdMapOffsetUpdatedForNonIdempotentData() {
+  def testProducerIdMapOffsetUpdatedForNonIdempotentData(): Unit = {
     val logConfig = LogTest.createLogConfig(segmentBytes = 2048 * 5)
     val log = createLog(logDir, logConfig)
     val records = TestUtils.records(List(new SimpleRecord(mockTime.milliseconds, "key".getBytes, "value".getBytes)))
@@ -593,7 +1025,7 @@ class LogTest {
   }
 
   @Test
-  def testRebuildProducerIdMapWithCompactedData() {
+  def testRebuildProducerIdMapWithCompactedData(): Unit = {
     val logConfig = LogTest.createLogConfig(segmentBytes = 2048 * 5)
     val log = createLog(logDir, logConfig)
     val pid = 1L
@@ -603,10 +1035,10 @@ class LogTest {
 
     // create a batch with a couple gaps to simulate compaction
     val records = TestUtils.records(producerId = pid, producerEpoch = epoch, sequence = seq, baseOffset = baseOffset, records = List(
-      new SimpleRecord(System.currentTimeMillis(), "a".getBytes),
-      new SimpleRecord(System.currentTimeMillis(), "key".getBytes, "b".getBytes),
-      new SimpleRecord(System.currentTimeMillis(), "c".getBytes),
-      new SimpleRecord(System.currentTimeMillis(), "key".getBytes, "d".getBytes)))
+      new SimpleRecord(mockTime.milliseconds(), "a".getBytes),
+      new SimpleRecord(mockTime.milliseconds(), "key".getBytes, "b".getBytes),
+      new SimpleRecord(mockTime.milliseconds(), "c".getBytes),
+      new SimpleRecord(mockTime.milliseconds(), "key".getBytes, "d".getBytes)))
     records.batches.asScala.foreach(_.setPartitionLeaderEpoch(0))
 
     val filtered = ByteBuffer.allocate(2048)
@@ -621,8 +1053,8 @@ class LogTest {
 
     // append some more data and then truncate to force rebuilding of the PID map
     val moreRecords = TestUtils.records(baseOffset = baseOffset + 4, records = List(
-      new SimpleRecord(System.currentTimeMillis(), "e".getBytes),
-      new SimpleRecord(System.currentTimeMillis(), "f".getBytes)))
+      new SimpleRecord(mockTime.milliseconds(), "e".getBytes),
+      new SimpleRecord(mockTime.milliseconds(), "f".getBytes)))
     moreRecords.batches.asScala.foreach(_.setPartitionLeaderEpoch(0))
     log.appendAsFollower(moreRecords)
 
@@ -636,7 +1068,7 @@ class LogTest {
   }
 
   @Test
-  def testRebuildProducerStateWithEmptyCompactedBatch() {
+  def testRebuildProducerStateWithEmptyCompactedBatch(): Unit = {
     val logConfig = LogTest.createLogConfig(segmentBytes = 2048 * 5)
     val log = createLog(logDir, logConfig)
     val pid = 1L
@@ -646,8 +1078,8 @@ class LogTest {
 
     // create an empty batch
     val records = TestUtils.records(producerId = pid, producerEpoch = epoch, sequence = seq, baseOffset = baseOffset, records = List(
-      new SimpleRecord(System.currentTimeMillis(), "key".getBytes, "a".getBytes),
-      new SimpleRecord(System.currentTimeMillis(), "key".getBytes, "b".getBytes)))
+      new SimpleRecord(mockTime.milliseconds(), "key".getBytes, "a".getBytes),
+      new SimpleRecord(mockTime.milliseconds(), "key".getBytes, "b".getBytes)))
     records.batches.asScala.foreach(_.setPartitionLeaderEpoch(0))
 
     val filtered = ByteBuffer.allocate(2048)
@@ -662,8 +1094,8 @@ class LogTest {
 
     // append some more data and then truncate to force rebuilding of the PID map
     val moreRecords = TestUtils.records(baseOffset = baseOffset + 2, records = List(
-      new SimpleRecord(System.currentTimeMillis(), "e".getBytes),
-      new SimpleRecord(System.currentTimeMillis(), "f".getBytes)))
+      new SimpleRecord(mockTime.milliseconds(), "e".getBytes),
+      new SimpleRecord(mockTime.milliseconds(), "f".getBytes)))
     moreRecords.batches.asScala.foreach(_.setPartitionLeaderEpoch(0))
     log.appendAsFollower(moreRecords)
 
@@ -677,7 +1109,7 @@ class LogTest {
   }
 
   @Test
-  def testUpdateProducerIdMapWithCompactedData() {
+  def testUpdateProducerIdMapWithCompactedData(): Unit = {
     val logConfig = LogTest.createLogConfig(segmentBytes = 2048 * 5)
     val log = createLog(logDir, logConfig)
     val pid = 1L
@@ -687,10 +1119,10 @@ class LogTest {
 
     // create a batch with a couple gaps to simulate compaction
     val records = TestUtils.records(producerId = pid, producerEpoch = epoch, sequence = seq, baseOffset = baseOffset, records = List(
-      new SimpleRecord(System.currentTimeMillis(), "a".getBytes),
-      new SimpleRecord(System.currentTimeMillis(), "key".getBytes, "b".getBytes),
-      new SimpleRecord(System.currentTimeMillis(), "c".getBytes),
-      new SimpleRecord(System.currentTimeMillis(), "key".getBytes, "d".getBytes)))
+      new SimpleRecord(mockTime.milliseconds(), "a".getBytes),
+      new SimpleRecord(mockTime.milliseconds(), "key".getBytes, "b".getBytes),
+      new SimpleRecord(mockTime.milliseconds(), "c".getBytes),
+      new SimpleRecord(mockTime.milliseconds(), "key".getBytes, "d".getBytes)))
     records.batches.asScala.foreach(_.setPartitionLeaderEpoch(0))
 
     val filtered = ByteBuffer.allocate(2048)
@@ -710,7 +1142,7 @@ class LogTest {
   }
 
   @Test
-  def testProducerIdMapTruncateTo() {
+  def testProducerIdMapTruncateTo(): Unit = {
     val logConfig = LogTest.createLogConfig(segmentBytes = 2048 * 5)
     val log = createLog(logDir, logConfig)
     log.appendAsLeader(TestUtils.records(List(new SimpleRecord("a".getBytes))), leaderEpoch = 0)
@@ -734,7 +1166,7 @@ class LogTest {
   }
 
   @Test
-  def testProducerIdMapTruncateToWithNoSnapshots() {
+  def testProducerIdMapTruncateToWithNoSnapshots(): Unit = {
     // This ensures that the upgrade optimization path cannot be hit after initial loading
     val logConfig = LogTest.createLogConfig(segmentBytes = 2048 * 5)
     val log = createLog(logDir, logConfig)
@@ -772,17 +1204,20 @@ class LogTest {
       producerEpoch = epoch, sequence = 0), leaderEpoch = 0)
     assertEquals(2, log.activeProducersWithLastSequence.size)
 
+    log.updateHighWatermark(log.logEndOffset)
     log.maybeIncrementLogStartOffset(1L)
 
-    assertEquals(1, log.activeProducersWithLastSequence.size)
+    // Deleting records should not remove producer state
+    assertEquals(2, log.activeProducersWithLastSequence.size)
     val retainedLastSeqOpt = log.activeProducersWithLastSequence.get(pid2)
     assertTrue(retainedLastSeqOpt.isDefined)
     assertEquals(0, retainedLastSeqOpt.get)
 
     log.close()
 
+    // Because the log start offset did not advance, producer snapshots will still be present and the state will be rebuilt
     val reloadedLog = createLog(logDir, logConfig, logStartOffset = 1L)
-    assertEquals(1, reloadedLog.activeProducersWithLastSequence.size)
+    assertEquals(2, reloadedLog.activeProducersWithLastSequence.size)
     val reloadedLastSeqOpt = log.activeProducersWithLastSequence.get(pid2)
     assertEquals(retainedLastSeqOpt, reloadedLastSeqOpt)
   }
@@ -804,18 +1239,20 @@ class LogTest {
     assertEquals(2, log.logSegments.size)
     assertEquals(2, log.activeProducersWithLastSequence.size)
 
+    log.updateHighWatermark(log.logEndOffset)
     log.maybeIncrementLogStartOffset(1L)
-    log.onHighWatermarkIncremented(log.logEndOffset)
     log.deleteOldSegments()
 
+    // Deleting records should not remove producer state
     assertEquals(1, log.logSegments.size)
-    assertEquals(1, log.activeProducersWithLastSequence.size)
+    assertEquals(2, log.activeProducersWithLastSequence.size)
     val retainedLastSeqOpt = log.activeProducersWithLastSequence.get(pid2)
     assertTrue(retainedLastSeqOpt.isDefined)
     assertEquals(0, retainedLastSeqOpt.get)
 
     log.close()
 
+    // After reloading log, producer state should not be regenerated
     val reloadedLog = createLog(logDir, logConfig, logStartOffset = 1L)
     assertEquals(1, reloadedLog.activeProducersWithLastSequence.size)
     val reloadedEntryOpt = log.activeProducersWithLastSequence.get(pid2)
@@ -823,7 +1260,7 @@ class LogTest {
   }
 
   @Test
-  def testProducerIdMapTruncateFullyAndStartAt() {
+  def testProducerIdMapTruncateFullyAndStartAt(): Unit = {
     val records = TestUtils.singletonRecords("foo".getBytes)
     val logConfig = LogTest.createLogConfig(segmentBytes = records.sizeInBytes, retentionBytes = records.sizeInBytes * 2)
     val log = createLog(logDir, logConfig)
@@ -845,7 +1282,7 @@ class LogTest {
   }
 
   @Test
-  def testProducerIdExpirationOnSegmentDeletion() {
+  def testProducerIdExpirationOnSegmentDeletion(): Unit = {
     val pid1 = 1L
     val records = TestUtils.records(Seq(new SimpleRecord("foo".getBytes)), producerId = pid1, producerEpoch = 0, sequence = 0)
     val logConfig = LogTest.createLogConfig(segmentBytes = records.sizeInBytes, retentionBytes = records.sizeInBytes * 2)
@@ -863,15 +1300,16 @@ class LogTest {
     assertEquals(3, log.logSegments.size)
     assertEquals(Set(pid1, pid2), log.activeProducersWithLastSequence.keySet)
 
-    log.onHighWatermarkIncremented(log.logEndOffset)
+    log.updateHighWatermark(log.logEndOffset)
     log.deleteOldSegments()
 
+    // Producer state should not be removed when deleting log segment
     assertEquals(2, log.logSegments.size)
-    assertEquals(Set(pid2), log.activeProducersWithLastSequence.keySet)
+    assertEquals(Set(pid1, pid2), log.activeProducersWithLastSequence.keySet)
   }
 
   @Test
-  def testTakeSnapshotOnRollAndDeleteSnapshotOnRecoveryPointCheckpoint() {
+  def testTakeSnapshotOnRollAndDeleteSnapshotOnRecoveryPointCheckpoint(): Unit = {
     val logConfig = LogTest.createLogConfig(segmentBytes = 2048 * 5)
     val log = createLog(logDir, logConfig)
     log.appendAsLeader(TestUtils.singletonRecords("a".getBytes), leaderEpoch = 0)
@@ -931,7 +1369,7 @@ class LogTest {
 
     val lastEntry = log.producerStateManager.lastEntry(producerId)
     assertTrue(lastEntry.isDefined)
-    assertEquals(0L, lastEntry.get.firstOffset)
+    assertEquals(0L, lastEntry.get.firstDataOffset)
     assertEquals(0L, lastEntry.get.lastDataOffset)
   }
 
@@ -950,9 +1388,8 @@ class LogTest {
       new SimpleRecord("bar".getBytes),
       new SimpleRecord("baz".getBytes))
     log.appendAsLeader(records, leaderEpoch = 0)
-    val commitAppendInfo = log.appendAsLeader(endTxnRecords(ControlRecordType.ABORT, pid, epoch),
-      isFromClient = false, leaderEpoch = 0)
-    log.onHighWatermarkIncremented(commitAppendInfo.lastOffset + 1)
+    val abortAppendInfo = appendEndTxnMarkerAsLeader(log, pid, epoch, ControlRecordType.ABORT)
+    log.updateHighWatermark(abortAppendInfo.lastOffset + 1)
 
     // now there should be no first unstable offset
     assertEquals(None, log.firstUnstableOffset)
@@ -960,7 +1397,7 @@ class LogTest {
     log.close()
 
     val reopenedLog = createLog(logDir, logConfig)
-    reopenedLog.onHighWatermarkIncremented(commitAppendInfo.lastOffset + 1)
+    reopenedLog.updateHighWatermark(abortAppendInfo.lastOffset + 1)
     assertEquals(None, reopenedLog.firstUnstableOffset)
   }
 
@@ -969,13 +1406,14 @@ class LogTest {
                             epoch: Short,
                             offset: Long = 0L,
                             coordinatorEpoch: Int = 0,
-                            partitionLeaderEpoch: Int = 0): MemoryRecords = {
+                            partitionLeaderEpoch: Int = 0,
+                            timestamp: Long = mockTime.milliseconds()): MemoryRecords = {
     val marker = new EndTransactionMarker(controlRecordType, coordinatorEpoch)
-    MemoryRecords.withEndTransactionMarker(offset, mockTime.milliseconds(), partitionLeaderEpoch, producerId, epoch, marker)
+    MemoryRecords.withEndTransactionMarker(offset, timestamp, partitionLeaderEpoch, producerId, epoch, marker)
   }
 
   @Test
-  def testPeriodicProducerIdExpiration() {
+  def testPeriodicProducerIdExpiration(): Unit = {
     val maxProducerIdExpirationMs = 200
     val producerIdExpirationCheckIntervalMs = 100
 
@@ -1070,7 +1508,7 @@ class LogTest {
   }
 
   @Test
-  def testMultipleProducerIdsPerMemoryRecord() : Unit = {
+  def testMultipleProducerIdsPerMemoryRecord(): Unit = {
     // create a log
     val log = createLog(logDir, LogConfig())
 
@@ -1116,7 +1554,7 @@ class LogTest {
   }
 
   @Test
-  def testDuplicateAppendToFollower() : Unit = {
+  def testDuplicateAppendToFollower(): Unit = {
     val logConfig = LogTest.createLogConfig(segmentBytes = 1024 * 1024 * 5)
     val log = createLog(logDir, logConfig)
     val epoch: Short = 0
@@ -1137,7 +1575,7 @@ class LogTest {
   }
 
   @Test
-  def testMultipleProducersWithDuplicatesInSingleAppend() : Unit = {
+  def testMultipleProducersWithDuplicatesInSingleAppend(): Unit = {
     val logConfig = LogTest.createLogConfig(segmentBytes = 1024 * 1024 * 5)
     val log = createLog(logDir, logConfig)
 
@@ -1203,12 +1641,41 @@ class LogTest {
     log.appendAsLeader(nextRecords, leaderEpoch = 0)
   }
 
+  @Test
+  def testDeleteSnapshotsOnIncrementLogStartOffset(): Unit = {
+    val logConfig = LogTest.createLogConfig(segmentBytes = 2048 * 5)
+    val log = createLog(logDir, logConfig)
+    val pid1 = 1L
+    val pid2 = 2L
+    val epoch = 0.toShort
+
+    log.appendAsLeader(TestUtils.records(List(new SimpleRecord(mockTime.milliseconds(), "a".getBytes)), producerId = pid1,
+      producerEpoch = epoch, sequence = 0), leaderEpoch = 0)
+    log.roll()
+    log.appendAsLeader(TestUtils.records(List(new SimpleRecord(mockTime.milliseconds(), "b".getBytes)), producerId = pid2,
+      producerEpoch = epoch, sequence = 0), leaderEpoch = 0)
+    log.roll()
+
+    assertEquals(2, log.activeProducersWithLastSequence.size)
+    assertEquals(2, ProducerStateManager.listSnapshotFiles(log.producerStateManager.logDir).size)
+
+    log.updateHighWatermark(log.logEndOffset)
+    log.maybeIncrementLogStartOffset(2L)
+
+    // Deleting records should not remove producer state but should delete snapshots
+    assertEquals(2, log.activeProducersWithLastSequence.size)
+    assertEquals(1, ProducerStateManager.listSnapshotFiles(log.producerStateManager.logDir).size)
+    val retainedLastSeqOpt = log.activeProducersWithLastSequence.get(pid2)
+    assertTrue(retainedLastSeqOpt.isDefined)
+    assertEquals(0, retainedLastSeqOpt.get)
+  }
+
   /**
    * Test for jitter s for time based log roll. This test appends messages then changes the time
    * using the mock clock to force the log to roll and checks the number of segments.
    */
   @Test
-  def testTimeBasedLogRollJitter() {
+  def testTimeBasedLogRollJitter(): Unit = {
     var set = TestUtils.singletonRecords(value = "test".getBytes, timestamp = mockTime.milliseconds)
     val maxJitter = 20 * 60L
     // create a log
@@ -1231,7 +1698,7 @@ class LogTest {
    * Test that appending more than the maximum segment size rolls the log
    */
   @Test
-  def testSizeBasedLogRoll() {
+  def testSizeBasedLogRoll(): Unit = {
     def createRecords = TestUtils.singletonRecords(value = "test".getBytes, timestamp = mockTime.milliseconds)
     val setSize = createRecords.sizeInBytes
     val msgPerSeg = 10
@@ -1251,7 +1718,7 @@ class LogTest {
    * Test that we can open and append to an empty log
    */
   @Test
-  def testLoadEmptyLog() {
+  def testLoadEmptyLog(): Unit = {
     createEmptyLogs(logDir, 0)
     val log = createLog(logDir, LogConfig())
     log.appendAsLeader(TestUtils.singletonRecords(value = "test".getBytes, timestamp = mockTime.milliseconds), leaderEpoch = 0)
@@ -1261,7 +1728,7 @@ class LogTest {
    * This test case appends a bunch of messages and checks that we can read them all back using sequential offsets.
    */
   @Test
-  def testAppendAndReadWithSequentialOffsets() {
+  def testAppendAndReadWithSequentialOffsets(): Unit = {
     val logConfig = LogTest.createLogConfig(segmentBytes = 71)
     val log = createLog(logDir, logConfig)
     val values = (0 until 100 by 2).map(id => id.toString.getBytes).toArray
@@ -1270,7 +1737,7 @@ class LogTest {
       log.appendAsLeader(TestUtils.singletonRecords(value = value), leaderEpoch = 0)
 
     for(i <- values.indices) {
-      val read = readLog(log, i, 100, Some(i+1)).records.batches.iterator.next()
+      val read = readLog(log, i, 1).records.batches.iterator.next()
       assertEquals("Offset read should match order appended.", i, read.lastOffset)
       val actual = read.iterator.next()
       assertNull("Key should be null", actual.key)
@@ -1285,7 +1752,7 @@ class LogTest {
    * from any offset less than the logEndOffset including offsets not appended.
    */
   @Test
-  def testAppendAndReadWithNonSequentialOffsets() {
+  def testAppendAndReadWithNonSequentialOffsets(): Unit = {
     val logConfig = LogTest.createLogConfig(segmentBytes = 72)
     val log = createLog(logDir, logConfig)
     val messageIds = ((0 until 50) ++ (50 until 200 by 7)).toArray
@@ -1302,6 +1769,185 @@ class LogTest {
     }
   }
 
+  /**
+    * This test generates a single record set that is a concatenation of many valid record sets,
+    * and appends them to a log.  Verifies decompression was not performed by checking the size
+    * of buffer used to validate batch header.
+    */
+  @Test
+  def testAvoidDecompression() {
+    val logConfig = LogTest.createLogConfig(segmentBytes = 5000, producerBatchDecompressionEnable = false)
+    val log = createLog(TestUtils.randomPartitionLogDir(tmpDir), logConfig, logStartOffset = 1000)
+    val values = (0 until 5).map(id => TestUtils.randomBytes(10)).toArray
+
+    // Build a single buffer with all record appended
+    val recordBuffer = ByteBuffer.allocate(5000)
+    val builder = MemoryRecords.builder(recordBuffer, RecordBatch.MAGIC_VALUE_V2, CompressionType.GZIP, TimestampType.CREATE_TIME,
+      0, System.currentTimeMillis, RecordBatch.NO_PRODUCER_ID, RecordBatch.NO_PRODUCER_EPOCH,
+      RecordBatch.NO_SEQUENCE)
+    for(value <- values) {
+      builder.append(new SimpleRecord(RecordBatch.NO_TIMESTAMP, null, value))
+    }
+
+    val records = builder.build()
+    val logAppendInfo = log.appendAsLeader(records, leaderEpoch = 0)
+
+    assertEquals("Size of buffer used to validate batch should be equal to the size of record format version V2 batch header.",
+      DefaultRecordBatch.RECORD_BATCH_OVERHEAD, logAppendInfo.recordConversionStats.temporaryMemoryBytes)
+  }
+
+  /**
+    * This test generates a single record set that is a concatenation of many valid record sets,
+    * with non-monotonically increasing relative offset and appends them to a log.
+    */
+  @Test(expected = classOf[InvalidRecordException])
+  def testAppendBatchWithoutMonotonicallyIncreasingRelativeOffset() {
+    val logConfig = LogTest.createLogConfig(segmentBytes = 5000, producerBatchDecompressionEnable = false)
+    val log = createLog(TestUtils.randomPartitionLogDir(tmpDir), logConfig)
+    val values = (0 until 5).map(id => TestUtils.randomBytes(10)).toArray
+
+    // Build a single buffer with all record appended
+    val recordBuffer = ByteBuffer.allocate(5000)
+    val builder = MemoryRecords.builder(recordBuffer, RecordBatch.MAGIC_VALUE_V2, CompressionType.GZIP, TimestampType.CREATE_TIME,
+      10, System.currentTimeMillis, RecordBatch.NO_PRODUCER_ID, RecordBatch.NO_PRODUCER_EPOCH,
+      RecordBatch.NO_SEQUENCE)
+    var count = 0
+    for(value <- values) {
+      builder.appendWithOffset(10 + (count * 2), new SimpleRecord(RecordBatch.NO_TIMESTAMP, null, value))
+      count = count + 1
+    }
+
+    val records = builder.build()
+    log.appendAsLeader(records, leaderEpoch = 0)
+  }
+
+  /**
+   * This test generates a single record set that is a concatenation of many valid record sets, and appends them to
+   * a log. It then checks to make sure we can read them all back.
+   */
+  @Test
+  def testAppendManyRecordSets() {
+    val msgFormatSet = Set(RecordBatch.MAGIC_VALUE_V0, RecordBatch.MAGIC_VALUE_V1, RecordBatch.MAGIC_VALUE_V2)
+    msgFormatSet.foreach { magic =>
+      appendManyRecordSets(createLog(TestUtils.randomPartitionLogDir(tmpDir), LogTest.createLogConfig(segmentBytes = 100)),
+        CompressionType.NONE, magic)
+      appendManyRecordSets(createLog(TestUtils.randomPartitionLogDir(tmpDir), LogTest.createLogConfig(segmentBytes = 100)),
+        CompressionType.GZIP, magic)
+      appendManyRecordSets(createLog(TestUtils.randomPartitionLogDir(tmpDir), LogTest.createLogConfig(segmentBytes = 100)),
+        CompressionType.SNAPPY, magic)
+      appendManyRecordSets(createLog(TestUtils.randomPartitionLogDir(tmpDir), LogTest.createLogConfig(segmentBytes = 100)),
+        CompressionType.LZ4, magic)
+    }
+  }
+
+  def readUncommitted(log: Log, startOffset: Long, maxLength: Int, maxOffset: Option[Long] = None): FetchDataInfo = {
+    log.read(startOffset,
+      maxLength = maxLength,
+      isolation = FetchLogEnd,
+      minOneMessage = false)
+  }
+
+  def appendManyRecordSets(log: Log,
+                           compressionType: CompressionType, magicValue: Byte) {
+    val values = (0 until 50).map(id => TestUtils.randomBytes(10)).toArray
+
+    // Build a single buffer with all record sets appended
+    val recordBuffer = ByteBuffer.allocate(5000)
+    for(value <- values) {
+      recordBuffer.put(
+        TestUtils.singletonRecords(
+          value = value, codec = compressionType, magicValue = magicValue).buffer())
+    }
+
+    // Close the buffer and create a MemoryRecords wrapping it
+    recordBuffer.flip()
+    recordBuffer.position(0)
+    val records = MemoryRecords.readableRecords(recordBuffer.slice())
+
+    log.appendAsLeader(records, leaderEpoch = 0)
+
+
+    for(i <- values.indices) {
+      val read = readUncommitted(log, i, 100, Some(i+1)).records.batches.iterator.next()
+      assertEquals("Offset read should match order appended.", i, read.lastOffset)
+      val actual = read.iterator.next()
+      assertNull("Key should be null", actual.key)
+      assertEquals("Values not equal", ByteBuffer.wrap(values(i)), actual.value)
+    }
+    assertEquals("Reading beyond the last message returns nothing.", 0,
+      readUncommitted(log, values.length, 100, None).records.batches.asScala.size)
+  }
+
+  /**
+   * This test makes sure that appending an empty record set does not fail. There should be no exception raised
+   */
+  @Test
+  def testAppendEmptyRecordSet() {
+    val logConfig = LogTest.createLogConfig(segmentBytes = 71)
+    val log = createLog(logDir, logConfig)
+    log.appendAsFollower(MemoryRecords.withRecords(CompressionType.NONE))
+    log.appendAsFollower(MemoryRecords.withRecords(CompressionType.GZIP))
+    log.appendAsFollower(MemoryRecords.withRecords(CompressionType.LZ4))
+    log.appendAsFollower(MemoryRecords.withRecords(CompressionType.SNAPPY))
+  }
+
+  /**
+   * This test generates a single record set that is a concatenation of many valid record sets, with one empty record
+   * set in the middle, and appends them to a log. It then checks to make sure we can read them all back.
+   */
+  @Test
+  def testAppendManyRecordSetsWithEmpty() {
+    val msgFormatSet = Set(RecordBatch.MAGIC_VALUE_V0, RecordBatch.MAGIC_VALUE_V1, RecordBatch.MAGIC_VALUE_V2)
+    msgFormatSet.foreach { magic =>
+      appendManyRecordSets(createLog(TestUtils.randomPartitionLogDir(tmpDir), LogTest.createLogConfig(segmentBytes = 100)),
+        CompressionType.NONE, magic)
+      appendManyRecordSets(createLog(TestUtils.randomPartitionLogDir(tmpDir), LogTest.createLogConfig(segmentBytes = 100)),
+        CompressionType.GZIP, magic)
+      appendManyRecordSets(createLog(TestUtils.randomPartitionLogDir(tmpDir), LogTest.createLogConfig(segmentBytes = 100)),
+        CompressionType.SNAPPY, magic)
+      appendManyRecordSets(createLog(TestUtils.randomPartitionLogDir(tmpDir), LogTest.createLogConfig(segmentBytes = 100)),
+        CompressionType.LZ4, magic)
+    }
+
+  }
+
+  def appendManyRecordSetsWithEmpty(log: Log, compressionType: CompressionType, magicValue: Byte): Unit = {
+    val values1 = (0 until 50 by 2).map(id => TestUtils.randomBytes(10)).toArray
+    val values2 = (50 until 100 by 2).map(id => TestUtils.randomBytes(10)).toArray
+
+    // Build a single buffer with all record sets appended
+    val recordBuffer = ByteBuffer.allocate(4000)
+    for (v <- values1) {
+      recordBuffer.put(
+        TestUtils.singletonRecords(value = v, codec = compressionType, magicValue = magicValue).buffer()
+      )
+    }
+    recordBuffer.put(MemoryRecords.withRecords(magicValue, CompressionType.NONE).buffer())
+    for (v <- values2) {
+      recordBuffer.put(
+        TestUtils.singletonRecords(value = v, codec = compressionType, magicValue = magicValue).buffer()
+      )
+    }
+
+    // Close the buffer and create a MemoryRecords wrapping it
+    recordBuffer.flip()
+    recordBuffer.position(0)
+    val records = MemoryRecords.readableRecords(recordBuffer.slice())
+
+    log.appendAsLeader(records, leaderEpoch = 0)
+
+    val values = values1 ++ values2
+    for (i <- values.indices) {
+      val read = readUncommitted(log, i, 100, Some(i + 1)).records.batches.iterator.next()
+      assertEquals("Offset read should match order appended.", i, read.lastOffset)
+      val actual = read.iterator.next()
+      assertNull("Key should be null", actual.key)
+      assertEquals("Values not equal", ByteBuffer.wrap(values(i)), actual.value)
+    }
+    assertEquals("Reading beyond the last message returns nothing.", 0,
+      readUncommitted(log, values.length, 100, None).records.batches.asScala.size)
+  }
+
   /**
    * This test covers an odd case where we have a gap in the offsets that falls at the end of a log segment.
    * Specifically we create a log where the last message in the first segment has offset 0. If we
@@ -1309,7 +1955,7 @@ class LogTest {
    * first segment has the greatest lower bound on the offset.
    */
   @Test
-  def testReadAtLogGap() {
+  def testReadAtLogGap(): Unit = {
     val logConfig = LogTest.createLogConfig(segmentBytes = 300)
     val log = createLog(logDir, logConfig)
 
@@ -1325,7 +1971,7 @@ class LogTest {
   }
 
   @Test(expected = classOf[KafkaStorageException])
-  def testLogRollAfterLogHandlerClosed() {
+  def testLogRollAfterLogHandlerClosed(): Unit = {
     val logConfig = LogTest.createLogConfig()
     val log = createLog(logDir,  logConfig)
     log.closeHandlers()
@@ -1333,7 +1979,7 @@ class LogTest {
   }
 
   @Test
-  def testReadWithMinMessage() {
+  def testReadWithMinMessage(): Unit = {
     val logConfig = LogTest.createLogConfig(segmentBytes = 72)
     val log = createLog(logDir,  logConfig)
     val messageIds = ((0 until 50) ++ (50 until 200 by 7)).toArray
@@ -1347,21 +1993,18 @@ class LogTest {
       val idx = messageIds.indexWhere(_ >= i)
       val reads = Seq(
         readLog(log, i, 1),
-        readLog(log, i, 100),
-        readLog(log, i, 100, Some(10000))
+        readLog(log, i, 100000),
+        readLog(log, i, 100)
       ).map(_.records.records.iterator.next())
       reads.foreach { read =>
         assertEquals("Offset read should match message id.", messageIds(idx), read.offset)
         assertEquals("Message should match appended.", records(idx), new SimpleRecord(read))
       }
-
-      val fetchedData = readLog(log, i, 1, Some(1))
-      assertEquals(Seq.empty, fetchedData.records.batches.asScala.toIndexedSeq)
     }
   }
 
   @Test
-  def testReadWithTooSmallMaxLength() {
+  def testReadWithTooSmallMaxLength(): Unit = {
     val logConfig = LogTest.createLogConfig(segmentBytes = 72)
     val log = createLog(logDir,  logConfig)
     val messageIds = ((0 until 50) ++ (50 until 200 by 7)).toArray
@@ -1393,7 +2036,7 @@ class LogTest {
    * - reading beyond the log end offset should throw an OffsetOutOfRangeException
    */
   @Test
-  def testReadOutOfRange() {
+  def testReadOutOfRange(): Unit = {
     createEmptyLogs(logDir, 1024)
     // set up replica log starting with offset 1024 and with one message (at offset 1024)
     val logConfig = LogTest.createLogConfig(segmentBytes = 1024)
@@ -1416,9 +2059,6 @@ class LogTest {
     } catch {
       case _: OffsetOutOfRangeException => // This is good.
     }
-
-    assertEquals("Reading from below the specified maxOffset should produce 0 byte read.", 0,
-      readLog(log, 1025, 1000, Some(1024)).records.sizeInBytes)
   }
 
   /**
@@ -1426,7 +2066,7 @@ class LogTest {
    * and then reads them all back and checks that the message read and offset matches what was appended.
    */
   @Test
-  def testLogRolls() {
+  def testLogRolls(): Unit = {
     /* create a multipart log with 100 messages */
     val logConfig = LogTest.createLogConfig(segmentBytes = 100)
     val log = createLog(logDir, logConfig)
@@ -1450,8 +2090,7 @@ class LogTest {
       assertEquals(s"Timestamps not equal at offset $offset", expected.timestamp, actual.timestamp)
       offset = head.lastOffset + 1
     }
-    val lastRead = readLog(log, startOffset = numMessages, maxLength = 1024*1024,
-      maxOffset = Some(numMessages + 1)).records
+    val lastRead = readLog(log, startOffset = numMessages, maxLength = 1024*1024).records
     assertEquals("Should be no more messages", 0, lastRead.records.asScala.size)
 
     // check that rolling the log forced a flushed, the flush is async so retry in case of failure
@@ -1464,7 +2103,7 @@ class LogTest {
    * Test reads at offsets that fall within compressed message set boundaries.
    */
   @Test
-  def testCompressedMessages() {
+  def testCompressedMessages(): Unit = {
     /* this log should roll after every messageset */
     val logConfig = LogTest.createLogConfig(segmentBytes = 110)
     val log = createLog(logDir, logConfig)
@@ -1486,7 +2125,7 @@ class LogTest {
    * Test garbage collecting old segments
    */
   @Test
-  def testThatGarbageCollectingSegmentsDoesntChangeOffset() {
+  def testThatGarbageCollectingSegmentsDoesntChangeOffset(): Unit = {
     for(messagesToAppend <- List(0, 1, 25)) {
       logDir.mkdirs()
       // first test a log segment starting at 0
@@ -1499,7 +2138,7 @@ class LogTest {
       assertEquals(currOffset, messagesToAppend)
 
       // time goes by; the log file is deleted
-      log.onHighWatermarkIncremented(currOffset)
+      log.updateHighWatermark(currOffset)
       log.deleteOldSegments()
 
       assertEquals("Deleting segments shouldn't have changed the logEndOffset", currOffset, log.logEndOffset)
@@ -1520,7 +2159,7 @@ class LogTest {
    * appending a message set larger than the config.segmentSize setting and checking that an exception is thrown.
    */
   @Test
-  def testMessageSetSizeCheck() {
+  def testMessageSetSizeCheck(): Unit = {
     val messageSet = MemoryRecords.withRecords(CompressionType.NONE, new SimpleRecord("You".getBytes), new SimpleRecord("bethe".getBytes))
     // append messages to log
     val configSegmentSize = messageSet.sizeInBytes - 1
@@ -1536,7 +2175,7 @@ class LogTest {
   }
 
   @Test
-  def testCompactedTopicConstraints() {
+  def testCompactedTopicConstraints(): Unit = {
     val keyedMessage = new SimpleRecord("and here it is".getBytes, "this message has a key".getBytes)
     val anotherKeyedMessage = new SimpleRecord("another key".getBytes, "this message also has a key".getBytes)
     val unkeyedMessage = new SimpleRecord("this message does not have a key".getBytes)
@@ -1556,21 +2195,25 @@ class LogTest {
       log.appendAsLeader(messageSetWithUnkeyedMessage, leaderEpoch = 0)
       fail("Compacted topics cannot accept a message without a key.")
     } catch {
-      case _: CorruptRecordException => // this is good
+      case _: RecordValidationException => // this is good
     }
     try {
       log.appendAsLeader(messageSetWithOneUnkeyedMessage, leaderEpoch = 0)
       fail("Compacted topics cannot accept a message without a key.")
     } catch {
-      case _: CorruptRecordException => // this is good
+      case _: RecordValidationException => // this is good
     }
     try {
       log.appendAsLeader(messageSetWithCompressedUnkeyedMessage, leaderEpoch = 0)
       fail("Compacted topics cannot accept a message without a key.")
     } catch {
-      case _: CorruptRecordException => // this is good
+      case _: RecordValidationException => // this is good
     }
 
+    // check if metric for NoKeyCompactedTopicRecordsPerSec is logged
+    assertEquals(metricsKeySet.count(_.getMBeanName.endsWith(s"${BrokerTopicStats.NoKeyCompactedTopicRecordsPerSec}")), 1)
+    assertTrue(TestUtils.meterCount(s"${BrokerTopicStats.NoKeyCompactedTopicRecordsPerSec}") > 0)
+
     // the following should succeed without any InvalidMessageException
     log.appendAsLeader(messageSetWithKeyedMessage, leaderEpoch = 0)
     log.appendAsLeader(messageSetWithKeyedMessages, leaderEpoch = 0)
@@ -1582,7 +2225,7 @@ class LogTest {
    * setting and checking that an exception is thrown.
    */
   @Test
-  def testMessageSizeCheck() {
+  def testMessageSizeCheck(): Unit = {
     val first = MemoryRecords.withRecords(CompressionType.NONE, new SimpleRecord("You".getBytes), new SimpleRecord("bethe".getBytes))
     val second = MemoryRecords.withRecords(CompressionType.NONE,
       new SimpleRecord("change (I need more bytes)... blah blah blah.".getBytes),
@@ -1607,7 +2250,7 @@ class LogTest {
    * Append a bunch of messages to a log and then re-open it both with and without recovery and check that the log re-initializes correctly.
    */
   @Test
-  def testLogRecoversToCorrectOffset() {
+  def testLogRecoversToCorrectOffset(): Unit = {
     val numMessages = 100
     val messageSize = 100
     val segmentSize = 7 * messageSize
@@ -1630,7 +2273,7 @@ class LogTest {
     }
     log.close()
 
-    def verifyRecoveredLog(log: Log, expectedRecoveryPoint: Long) {
+    def verifyRecoveredLog(log: Log, expectedRecoveryPoint: Long): Unit = {
       assertEquals(s"Unexpected recovery point", expectedRecoveryPoint, log.recoveryPoint)
       assertEquals(s"Should have $numMessages messages when log is reopened w/o recovery", numMessages, log.logEndOffset)
       assertEquals("Should have same last index offset as before.", lastIndexOffset, log.activeSegment.offsetIndex.lastOffset)
@@ -1654,7 +2297,7 @@ class LogTest {
    * Test building the time index on the follower by setting assignOffsets to false.
    */
   @Test
-  def testBuildTimeIndexWhenNotAssigningOffsets() {
+  def testBuildTimeIndexWhenNotAssigningOffsets(): Unit = {
     val numMessages = 100
     val logConfig = LogTest.createLogConfig(segmentBytes = 10000, indexIntervalBytes = 1)
     val log = createLog(logDir, logConfig)
@@ -1673,7 +2316,7 @@ class LogTest {
    * Test that if we manually delete an index segment it is rebuilt when the log is re-opened
    */
   @Test
-  def testIndexRebuild() {
+  def testIndexRebuild(): Unit = {
     // publish the messages and close the log
     val numMessages = 200
     val logConfig = LogTest.createLogConfig(segmentBytes = 200, indexIntervalBytes = 1)
@@ -1746,7 +2389,7 @@ class LogTest {
    * Test that if messages format version of the messages in a segment is before 0.10.0, the time index should be empty.
    */
   @Test
-  def testRebuildTimeIndexForOldMessages() {
+  def testRebuildTimeIndexForOldMessages(): Unit = {
     val numMessages = 200
     val segmentSize = 200
     val logConfig = LogTest.createLogConfig(segmentBytes = segmentSize, indexIntervalBytes = 1, messageFormatVersion = "0.9.0")
@@ -1772,7 +2415,7 @@ class LogTest {
    * Test that if we have corrupted an index segment it is rebuilt when the log is re-opened
    */
   @Test
-  def testCorruptIndexRebuild() {
+  def testCorruptIndexRebuild(): Unit = {
     // publish the messages and close the log
     val numMessages = 200
     val logConfig = LogTest.createLogConfig(segmentBytes = 200, indexIntervalBytes = 1)
@@ -1814,7 +2457,7 @@ class LogTest {
    * Test the Log truncate operations
    */
   @Test
-  def testTruncateTo() {
+  def testTruncateTo(): Unit = {
     def createRecords = TestUtils.singletonRecords(value = "test".getBytes, timestamp = mockTime.milliseconds)
     val setSize = createRecords.sizeInBytes
     val msgPerSeg = 10
@@ -1869,7 +2512,7 @@ class LogTest {
    * Verify that when we truncate a log the index of the last segment is resized to the max index size to allow more appends
    */
   @Test
-  def testIndexResizingAtTruncation() {
+  def testIndexResizingAtTruncation(): Unit = {
     val setSize = TestUtils.singletonRecords(value = "test".getBytes, timestamp = mockTime.milliseconds).sizeInBytes
     val msgPerSeg = 10
     val segmentSize = msgPerSeg * setSize  // each segment will be 10 messages
@@ -1905,7 +2548,7 @@ class LogTest {
    * When we open a log any index segments without an associated log segment should be deleted.
    */
   @Test
-  def testBogusIndexSegmentsAreRemoved() {
+  def testBogusIndexSegmentsAreRemoved(): Unit = {
     val bogusIndex1 = Log.offsetIndexFile(logDir, 0)
     val bogusTimeIndex1 = Log.timeIndexFile(logDir, 0)
     val bogusIndex2 = Log.offsetIndexFile(logDir, 5)
@@ -1940,7 +2583,7 @@ class LogTest {
    * Verify that truncation works correctly after re-opening the log
    */
   @Test
-  def testReopenThenTruncate() {
+  def testReopenThenTruncate(): Unit = {
     def createRecords = TestUtils.singletonRecords(value = "test".getBytes, timestamp = mockTime.milliseconds)
     // create a log
     val logConfig = LogTest.createLogConfig(segmentBytes = createRecords.sizeInBytes * 5, segmentIndexBytes = 1000, indexIntervalBytes = 10000)
@@ -1960,7 +2603,7 @@ class LogTest {
    * Test that deleted files are deleted after the appropriate time.
    */
   @Test
-  def testAsyncDelete() {
+  def testAsyncDelete(): Unit = {
     def createRecords = TestUtils.singletonRecords(value = "test".getBytes, timestamp = mockTime.milliseconds - 1000L)
     val asyncDeleteMs = 1000
     val logConfig = LogTest.createLogConfig(segmentBytes = createRecords.sizeInBytes * 5, segmentIndexBytes = 1000, indexIntervalBytes = 10000,
@@ -1975,7 +2618,7 @@ class LogTest {
     val segments = log.logSegments.toArray
     val oldFiles = segments.map(_.log.file) ++ segments.map(_.lazyOffsetIndex.file)
 
-    log.onHighWatermarkIncremented(log.logEndOffset)
+    log.updateHighWatermark(log.logEndOffset)
     log.deleteOldSegments()
 
     assertEquals("Only one segment should remain.", 1, log.numberOfSegments)
@@ -1995,7 +2638,7 @@ class LogTest {
    * Any files ending in .deleted should be removed when the log is re-opened.
    */
   @Test
-  def testOpenDeletesObsoleteFiles() {
+  def testOpenDeletesObsoleteFiles(): Unit = {
     def createRecords = TestUtils.singletonRecords(value = "test".getBytes, timestamp = mockTime.milliseconds - 1000)
     val logConfig = LogTest.createLogConfig(segmentBytes = createRecords.sizeInBytes * 5, segmentIndexBytes = 1000, retentionMs = 999)
     var log = createLog(logDir, logConfig)
@@ -2005,7 +2648,7 @@ class LogTest {
       log.appendAsLeader(createRecords, leaderEpoch = 0)
 
     // expire all segments
-    log.onHighWatermarkIncremented(log.logEndOffset)
+    log.updateHighWatermark(log.logEndOffset)
     log.deleteOldSegments()
     log.close()
     log = createLog(logDir, logConfig)
@@ -2013,7 +2656,7 @@ class LogTest {
   }
 
   @Test
-  def testAppendMessageWithNullPayload() {
+  def testAppendMessageWithNullPayload(): Unit = {
     val log = createLog(logDir, LogConfig())
     log.appendAsLeader(TestUtils.singletonRecords(value = null), leaderEpoch = 0)
     val head = readLog(log, 0, 4096).records.records.iterator.next()
@@ -2022,7 +2665,7 @@ class LogTest {
   }
 
   @Test
-  def testAppendWithOutOfOrderOffsetsThrowsException() {
+  def testAppendWithOutOfOrderOffsetsThrowsException(): Unit = {
     val log = createLog(logDir, LogConfig())
 
     val appendOffsets = Seq(0L, 1L, 3L, 2L, 4L)
@@ -2037,13 +2680,13 @@ class LogTest {
     buffer.flip()
     val memoryRecords = MemoryRecords.readableRecords(buffer)
 
-    assertThrows[OffsetsOutOfOrderException] {
+    assertThrows[UnexpectedAppendOffsetException] {
       log.appendAsFollower(memoryRecords)
     }
   }
 
   @Test
-  def testAppendBelowExpectedOffsetThrowsException() {
+  def testAppendBelowExpectedOffsetThrowsException(): Unit = {
     val log = createLog(logDir, LogConfig())
     val records = (0 until 2).map(id => new SimpleRecord(id.toString.getBytes)).toArray
     records.foreach(record => log.appendAsLeader(MemoryRecords.withRecords(CompressionType.NONE, record), leaderEpoch = 0))
@@ -2067,26 +2710,33 @@ class LogTest {
     assertEquals(7L, log.logStartOffset)
     assertEquals(7L, log.logEndOffset)
 
-    val firstOffset = 4L
-    val magicVals = Seq(RecordBatch.MAGIC_VALUE_V0, RecordBatch.MAGIC_VALUE_V1, RecordBatch.MAGIC_VALUE_V2)
-    val compressionTypes = Seq(CompressionType.NONE, CompressionType.LZ4)
-    for (magic <- magicVals; compression <- compressionTypes) {
+    def testMain(magic: Byte, compression: CompressionType) {
+      val firstOffset = 4L
       val batch = TestUtils.records(List(new SimpleRecord("k1".getBytes, "v1".getBytes),
-                                         new SimpleRecord("k2".getBytes, "v2".getBytes),
-                                         new SimpleRecord("k3".getBytes, "v3".getBytes)),
-                                    magicValue = magic, codec = compression,
-                                    baseOffset = firstOffset)
+        new SimpleRecord("k2".getBytes, "v2".getBytes),
+        new SimpleRecord("k3".getBytes, "v3".getBytes)),
+        magicValue = magic, codec = compression,
+        baseOffset = firstOffset)
 
       withClue(s"Magic=$magic, compressionType=$compression") {
         val exception = intercept[UnexpectedAppendOffsetException] {
           log.appendAsFollower(records = batch)
         }
         assertEquals(s"Magic=$magic, compressionType=$compression, UnexpectedAppendOffsetException#firstOffset",
-                     firstOffset, exception.firstOffset)
+          firstOffset, exception.firstOffset)
         assertEquals(s"Magic=$magic, compressionType=$compression, UnexpectedAppendOffsetException#lastOffset",
-                     firstOffset + 2, exception.lastOffset)
+          firstOffset + 2, exception.lastOffset)
       }
     }
+
+    // For this test, we need a batch with multiple records. It appears that the above TestUtils.records() actually
+    // returns multiple batches with single record for magic values V0 and V1 when compression type is NONE. Exclude
+    // these combinations.
+    testMain(RecordBatch.MAGIC_VALUE_V2, CompressionType.NONE)
+    val magicVals = Seq(RecordBatch.MAGIC_VALUE_V0, RecordBatch.MAGIC_VALUE_V1, RecordBatch.MAGIC_VALUE_V2)
+    for (magic <- magicVals) {
+      testMain(magic, CompressionType.LZ4)
+    }
   }
 
   @Test
@@ -2097,7 +2747,7 @@ class LogTest {
   }
 
   @Test
-  def testCorruptLog() {
+  def testCorruptLog(): Unit = {
     // append some messages to create some segments
     val logConfig = LogTest.createLogConfig(segmentBytes = 1000, indexIntervalBytes = 1, maxMessageBytes = 64 * 1024)
     def createRecords = TestUtils.singletonRecords(value = "test".getBytes, timestamp = mockTime.milliseconds)
@@ -2184,6 +2834,19 @@ class LogTest {
     assertFalse(LeaderEpochCheckpointFile.newFile(this.logDir).exists())
   }
 
+  @Test
+  def testLeaderEpochCacheClearedAfterDowngradeInAppendedMessages(): Unit = {
+    val logConfig = LogTest.createLogConfig(segmentBytes = 1000, indexIntervalBytes = 1, maxMessageBytes = 64 * 1024)
+    val log = createLog(logDir, logConfig)
+    log.appendAsLeader(TestUtils.records(List(new SimpleRecord("foo".getBytes()))), leaderEpoch = 5)
+    assertEquals(Some(5), log.leaderEpochCache.flatMap(_.latestEpoch))
+
+    log.appendAsFollower(TestUtils.records(List(new SimpleRecord("foo".getBytes())),
+      baseOffset = 1L,
+      magicValue = RecordVersion.V1.value))
+    assertEquals(None, log.leaderEpochCache.flatMap(_.latestEpoch))
+  }
+
   @Test
   def testLeaderEpochCacheClearedAfterStaticMessageFormatDowngrade(): Unit = {
     val logConfig = LogTest.createLogConfig(segmentBytes = 1000, indexIntervalBytes = 1, maxMessageBytes = 64 * 1024)
@@ -2212,7 +2875,7 @@ class LogTest {
 
     val downgradedLogConfig = LogTest.createLogConfig(segmentBytes = 1000, indexIntervalBytes = 1,
       maxMessageBytes = 64 * 1024, messageFormatVersion = kafka.api.KAFKA_0_10_2_IV0.shortVersion)
-    log.updateConfig(Set(LogConfig.MessageFormatVersionProp), downgradedLogConfig)
+    log.updateConfig(downgradedLogConfig)
     assertLeaderEpochCacheEmpty(log)
 
     log.appendAsLeader(TestUtils.records(List(new SimpleRecord("bar".getBytes())),
@@ -2231,7 +2894,7 @@ class LogTest {
 
     val upgradedLogConfig = LogTest.createLogConfig(segmentBytes = 1000, indexIntervalBytes = 1,
       maxMessageBytes = 64 * 1024, messageFormatVersion = kafka.api.KAFKA_0_11_0_IV0.shortVersion)
-    log.updateConfig(Set(LogConfig.MessageFormatVersionProp), upgradedLogConfig)
+    log.updateConfig(upgradedLogConfig)
     log.appendAsLeader(TestUtils.records(List(new SimpleRecord("foo".getBytes()))), leaderEpoch = 5)
     assertEquals(Some(5), log.latestEpoch)
   }
@@ -2552,7 +3215,7 @@ class LogTest {
   }
 
   @Test
-  def testCleanShutdownFile() {
+  def testCleanShutdownFile(): Unit = {
     // append some messages to create some segments
     val logConfig = LogTest.createLogConfig(segmentBytes = 1000, indexIntervalBytes = 1, maxMessageBytes = 64 * 1024)
     def createRecords = TestUtils.singletonRecords(value = "test".getBytes, timestamp = mockTime.milliseconds)
@@ -2575,7 +3238,7 @@ class LogTest {
   }
 
   @Test
-  def testParseTopicPartitionName() {
+  def testParseTopicPartitionName(): Unit = {
     val topic = "test_topic"
     val partition = "143"
     val dir = new File(logDir, topicPartitionName(topic, partition))
@@ -2589,7 +3252,7 @@ class LogTest {
    * are parsed correctly by `Log.parseTopicPartitionName` (see KAFKA-5232 for details).
    */
   @Test
-  def testParseTopicPartitionNameWithPeriodForDeletedTopic() {
+  def testParseTopicPartitionNameWithPeriodForDeletedTopic(): Unit = {
     val topic = "foo.bar-testtopic"
     val partition = "42"
     val dir = new File(logDir, Log.logDeleteDirName(new TopicPartition(topic, partition.toInt)))
@@ -2599,7 +3262,7 @@ class LogTest {
   }
 
   @Test
-  def testParseTopicPartitionNameForEmptyName() {
+  def testParseTopicPartitionNameForEmptyName(): Unit = {
     try {
       val dir = new File("")
       Log.parseTopicPartitionName(dir)
@@ -2610,7 +3273,7 @@ class LogTest {
   }
 
   @Test
-  def testParseTopicPartitionNameForNull() {
+  def testParseTopicPartitionNameForNull(): Unit = {
     try {
       val dir: File = null
       Log.parseTopicPartitionName(dir)
@@ -2621,7 +3284,7 @@ class LogTest {
   }
 
   @Test
-  def testParseTopicPartitionNameForMissingSeparator() {
+  def testParseTopicPartitionNameForMissingSeparator(): Unit = {
     val topic = "test_topic"
     val partition = "1999"
     val dir = new File(logDir, topic + partition)
@@ -2642,7 +3305,7 @@ class LogTest {
   }
 
   @Test
-  def testParseTopicPartitionNameForMissingTopic() {
+  def testParseTopicPartitionNameForMissingTopic(): Unit = {
     val topic = ""
     val partition = "1999"
     val dir = new File(logDir, topicPartitionName(topic, partition))
@@ -2664,7 +3327,7 @@ class LogTest {
   }
 
   @Test
-  def testParseTopicPartitionNameForMissingPartition() {
+  def testParseTopicPartitionNameForMissingPartition(): Unit = {
     val topic = "test_topic"
     val partition = ""
     val dir = new File(logDir + topicPartitionName(topic, partition))
@@ -2685,7 +3348,7 @@ class LogTest {
   }
 
   @Test
-  def testParseTopicPartitionNameForInvalidPartition() {
+  def testParseTopicPartitionNameForInvalidPartition(): Unit = {
     val topic = "test_topic"
     val partition = "1999a"
     val dir = new File(logDir, topicPartitionName(topic, partition))
@@ -2706,7 +3369,7 @@ class LogTest {
   }
 
   @Test
-  def testParseTopicPartitionNameForExistingInvalidDir() {
+  def testParseTopicPartitionNameForExistingInvalidDir(): Unit = {
     val dir1 = new File(logDir + "/non_kafka_dir")
     try {
       Log.parseTopicPartitionName(dir1)
@@ -2727,7 +3390,7 @@ class LogTest {
     topic + "-" + partition
 
   @Test
-  def testDeleteOldSegments() {
+  def testDeleteOldSegments(): Unit = {
     def createRecords = TestUtils.singletonRecords(value = "test".getBytes, timestamp = mockTime.milliseconds - 1000)
     val logConfig = LogTest.createLogConfig(segmentBytes = createRecords.sizeInBytes * 5, segmentIndexBytes = 1000, retentionMs = 999)
     val log = createLog(logDir, logConfig)
@@ -2747,11 +3410,11 @@ class LogTest {
 
     // only segments with offset before the current high watermark are eligible for deletion
     for (hw <- 25 to 30) {
-      log.onHighWatermarkIncremented(hw)
+      log.updateHighWatermark(hw)
       log.deleteOldSegments()
       assertTrue(log.logStartOffset <= hw)
       log.logSegments.foreach { segment =>
-        val segmentFetchInfo = segment.read(startOffset = segment.baseOffset, maxOffset = None, maxSize = Int.MaxValue)
+        val segmentFetchInfo = segment.read(startOffset = segment.baseOffset, maxSize = Int.MaxValue)
         val segmentLastOffsetOpt = segmentFetchInfo.records.records.asScala.lastOption.map(_.offset)
         segmentLastOffsetOpt.foreach { lastOffset =>
           assertTrue(lastOffset >= hw)
@@ -2760,7 +3423,7 @@ class LogTest {
     }
 
     // expire all segments
-    log.onHighWatermarkIncremented(log.logEndOffset)
+    log.updateHighWatermark(log.logEndOffset)
     log.deleteOldSegments()
     assertEquals("The deleted segments should be gone.", 1, log.numberOfSegments)
     assertEquals("Epoch entries should have gone.", 1, epochCache(log).epochEntries.size)
@@ -2777,7 +3440,7 @@ class LogTest {
   }
 
   @Test
-  def testLogDeletionAfterClose() {
+  def testLogDeletionAfterClose(): Unit = {
     def createRecords = TestUtils.singletonRecords(value = "test".getBytes, timestamp = mockTime.milliseconds - 1000)
     val logConfig = LogTest.createLogConfig(segmentBytes = createRecords.sizeInBytes * 5, segmentIndexBytes = 1000, retentionMs = 999)
     val log = createLog(logDir, logConfig)
@@ -2795,7 +3458,7 @@ class LogTest {
   }
 
   @Test
-  def testLogDeletionAfterDeleteRecords() {
+  def testLogDeletionAfterDeleteRecords(): Unit = {
     def createRecords = TestUtils.singletonRecords("test".getBytes)
     val logConfig = LogTest.createLogConfig(segmentBytes = createRecords.sizeInBytes * 5)
     val log = createLog(logDir, logConfig)
@@ -2804,7 +3467,7 @@ class LogTest {
       log.appendAsLeader(createRecords, leaderEpoch = 0)
     assertEquals("should have 3 segments", 3, log.numberOfSegments)
     assertEquals(log.logStartOffset, 0)
-    log.onHighWatermarkIncremented(log.logEndOffset)
+    log.updateHighWatermark(log.logEndOffset)
 
     log.maybeIncrementLogStartOffset(1)
     log.deleteOldSegments()
@@ -2827,7 +3490,7 @@ class LogTest {
   }
 
   @Test
-  def shouldDeleteSizeBasedSegments() {
+  def shouldDeleteSizeBasedSegments(): Unit = {
     def createRecords = TestUtils.singletonRecords("test".getBytes)
     val logConfig = LogTest.createLogConfig(segmentBytes = createRecords.sizeInBytes * 5, retentionBytes = createRecords.sizeInBytes * 10)
     val log = createLog(logDir, logConfig)
@@ -2836,13 +3499,13 @@ class LogTest {
     for (_ <- 0 until 15)
       log.appendAsLeader(createRecords, leaderEpoch = 0)
 
-    log.onHighWatermarkIncremented(log.logEndOffset)
+    log.updateHighWatermark(log.logEndOffset)
     log.deleteOldSegments()
     assertEquals("should have 2 segments", 2,log.numberOfSegments)
   }
 
   @Test
-  def shouldNotDeleteSizeBasedSegmentsWhenUnderRetentionSize() {
+  def shouldNotDeleteSizeBasedSegmentsWhenUnderRetentionSize(): Unit = {
     def createRecords = TestUtils.singletonRecords("test".getBytes)
     val logConfig = LogTest.createLogConfig(segmentBytes = createRecords.sizeInBytes * 5, retentionBytes = createRecords.sizeInBytes * 15)
     val log = createLog(logDir, logConfig)
@@ -2851,13 +3514,13 @@ class LogTest {
     for (_ <- 0 until 15)
       log.appendAsLeader(createRecords, leaderEpoch = 0)
 
-    log.onHighWatermarkIncremented(log.logEndOffset)
+    log.updateHighWatermark(log.logEndOffset)
     log.deleteOldSegments()
     assertEquals("should have 3 segments", 3,log.numberOfSegments)
   }
 
   @Test
-  def shouldDeleteTimeBasedSegmentsReadyToBeDeleted() {
+  def shouldDeleteTimeBasedSegmentsReadyToBeDeleted(): Unit = {
     def createRecords = TestUtils.singletonRecords("test".getBytes, timestamp = 10)
     val logConfig = LogTest.createLogConfig(segmentBytes = createRecords.sizeInBytes * 5, retentionMs = 10000)
     val log = createLog(logDir, logConfig)
@@ -2866,13 +3529,13 @@ class LogTest {
     for (_ <- 0 until 15)
       log.appendAsLeader(createRecords, leaderEpoch = 0)
 
-    log.onHighWatermarkIncremented(log.logEndOffset)
+    log.updateHighWatermark(log.logEndOffset)
     log.deleteOldSegments()
     assertEquals("There should be 1 segment remaining", 1, log.numberOfSegments)
   }
 
   @Test
-  def shouldNotDeleteTimeBasedSegmentsWhenNoneReadyToBeDeleted() {
+  def shouldNotDeleteTimeBasedSegmentsWhenNoneReadyToBeDeleted(): Unit = {
     def createRecords = TestUtils.singletonRecords("test".getBytes, timestamp = mockTime.milliseconds)
     val logConfig = LogTest.createLogConfig(segmentBytes = createRecords.sizeInBytes * 5, retentionMs = 10000000)
     val log = createLog(logDir, logConfig)
@@ -2881,13 +3544,13 @@ class LogTest {
     for (_ <- 0 until 15)
       log.appendAsLeader(createRecords, leaderEpoch = 0)
 
-    log.onHighWatermarkIncremented(log.logEndOffset)
+    log.updateHighWatermark(log.logEndOffset)
     log.deleteOldSegments()
     assertEquals("There should be 3 segments remaining", 3, log.numberOfSegments)
   }
 
   @Test
-  def shouldNotDeleteSegmentsWhenPolicyDoesNotIncludeDelete() {
+  def shouldNotDeleteSegmentsWhenPolicyDoesNotIncludeDelete(): Unit = {
     def createRecords = TestUtils.singletonRecords("test".getBytes, key = "test".getBytes(), timestamp = 10L)
     val logConfig = LogTest.createLogConfig(segmentBytes = createRecords.sizeInBytes * 5, retentionMs = 10000, cleanupPolicy = "compact")
     val log = createLog(logDir, logConfig)
@@ -2900,13 +3563,13 @@ class LogTest {
     log.logSegments.head.lastModified = mockTime.milliseconds - 20000
 
     val segments = log.numberOfSegments
-    log.onHighWatermarkIncremented(log.logEndOffset)
+    log.updateHighWatermark(log.logEndOffset)
     log.deleteOldSegments()
     assertEquals("There should be 3 segments remaining", segments, log.numberOfSegments)
   }
 
   @Test
-  def shouldDeleteSegmentsReadyToBeDeletedWhenCleanupPolicyIsCompactAndDelete() {
+  def shouldDeleteSegmentsReadyToBeDeletedWhenCleanupPolicyIsCompactAndDelete(): Unit = {
     def createRecords = TestUtils.singletonRecords("test".getBytes, key = "test".getBytes, timestamp = 10L)
     val logConfig = LogTest.createLogConfig(segmentBytes = createRecords.sizeInBytes * 5, retentionMs = 10000, cleanupPolicy = "compact,delete")
     val log = createLog(logDir, logConfig)
@@ -2915,7 +3578,7 @@ class LogTest {
     for (_ <- 0 until 15)
       log.appendAsLeader(createRecords, leaderEpoch = 0)
 
-    log.onHighWatermarkIncremented(log.logEndOffset)
+    log.updateHighWatermark(log.logEndOffset)
     log.deleteOldSegments()
     assertEquals("There should be 1 segment remaining", 1, log.numberOfSegments)
   }
@@ -2923,23 +3586,23 @@ class LogTest {
   @Test
   def shouldDeleteStartOffsetBreachedSegmentsWhenPolicyDoesNotIncludeDelete(): Unit = {
     def createRecords = TestUtils.singletonRecords("test".getBytes, key = "test".getBytes, timestamp = 10L)
-    val logConfig = LogTest.createLogConfig(segmentBytes = createRecords.sizeInBytes * 5, retentionMs = 10000, cleanupPolicy = "compact")
-
-    // Create log with start offset ahead of the first log segment
-    val log = createLog(logDir, logConfig, brokerTopicStats, logStartOffset = 5L)
+    val recordsPerSegment = 5
+    val logConfig = LogTest.createLogConfig(segmentBytes = createRecords.sizeInBytes * recordsPerSegment, retentionMs = 10000, cleanupPolicy = "compact")
+    val log = createLog(logDir, logConfig, brokerTopicStats)
 
     // append some messages to create some segments
     for (_ <- 0 until 15)
       log.appendAsLeader(createRecords, leaderEpoch = 0)
 
-    // Three segments should be created, with the first one entirely preceding the log start offset
+    // Three segments should be created
     assertEquals(3, log.logSegments.count(_ => true))
-    assertTrue(log.logSegments.slice(1, 2).head.baseOffset <= log.logStartOffset)
+    log.updateHighWatermark(log.logEndOffset)
+    log.maybeIncrementLogStartOffset(recordsPerSegment)
 
     // The first segment, which is entirely before the log start offset, should be deleted
     // Of the remaining the segments, the first can overlap the log start offset and the rest must have a base offset
     // greater than the start offset
-    log.onHighWatermarkIncremented(log.logEndOffset)
+    log.updateHighWatermark(log.logEndOffset)
     log.deleteOldSegments()
     assertEquals("There should be 2 segments remaining", 2, log.numberOfSegments)
     assertTrue(log.logSegments.head.baseOffset <= log.logStartOffset)
@@ -2947,7 +3610,7 @@ class LogTest {
   }
 
   @Test
-  def shouldApplyEpochToMessageOnAppendIfLeader() {
+  def shouldApplyEpochToMessageOnAppendIfLeader(): Unit = {
     val records = (0 until 50).toArray.map(id => new SimpleRecord(id.toString.getBytes))
 
     //Given this partition is on leader epoch 72
@@ -2964,13 +3627,13 @@ class LogTest {
 
     //Then leader epoch should be set on messages
     for (i <- records.indices) {
-      val read = readLog(log, i, 100, Some(i+1)).records.batches.iterator.next()
+      val read = readLog(log, i, 1).records.batches.iterator.next()
       assertEquals("Should have set leader epoch", 72, read.partitionLeaderEpoch)
     }
   }
 
   @Test
-  def followerShouldSaveEpochInformationFromReplicatedMessagesToTheEpochCache() {
+  def followerShouldSaveEpochInformationFromReplicatedMessagesToTheEpochCache(): Unit = {
     val messageIds = (0 until 50).toArray
     val records = messageIds.map(id => new SimpleRecord(id.toString.getBytes))
 
@@ -2994,7 +3657,7 @@ class LogTest {
   }
 
   @Test
-  def shouldTruncateLeaderEpochsWhenDeletingSegments() {
+  def shouldTruncateLeaderEpochsWhenDeletingSegments(): Unit = {
     def createRecords = TestUtils.singletonRecords("test".getBytes)
     val logConfig = LogTest.createLogConfig(segmentBytes = createRecords.sizeInBytes * 5, retentionBytes = createRecords.sizeInBytes * 10)
     val log = createLog(logDir, logConfig)
@@ -3011,7 +3674,7 @@ class LogTest {
     cache.assign(2, 10)
 
     //When first segment is removed
-    log.onHighWatermarkIncremented(log.logEndOffset)
+    log.updateHighWatermark(log.logEndOffset)
     log.deleteOldSegments()
 
     //The oldest epoch entry should have been removed
@@ -3019,7 +3682,7 @@ class LogTest {
   }
 
   @Test
-  def shouldUpdateOffsetForLeaderEpochsWhenDeletingSegments() {
+  def shouldUpdateOffsetForLeaderEpochsWhenDeletingSegments(): Unit = {
     def createRecords = TestUtils.singletonRecords("test".getBytes)
     val logConfig = LogTest.createLogConfig(segmentBytes = createRecords.sizeInBytes * 5, retentionBytes = createRecords.sizeInBytes * 10)
     val log = createLog(logDir, logConfig)
@@ -3036,7 +3699,7 @@ class LogTest {
     cache.assign(2, 10)
 
     //When first segment removed (up to offset 5)
-    log.onHighWatermarkIncremented(log.logEndOffset)
+    log.updateHighWatermark(log.logEndOffset)
     log.deleteOldSegments()
 
     //The first entry should have gone from (0,0) => (0,5)
@@ -3044,7 +3707,7 @@ class LogTest {
   }
 
   @Test
-  def shouldTruncateLeaderEpochCheckpointFileWhenTruncatingLog() {
+  def shouldTruncateLeaderEpochCheckpointFileWhenTruncatingLog(): Unit = {
     def createRecords(startOffset: Long, epoch: Int): MemoryRecords = {
       TestUtils.records(Seq(new SimpleRecord("value".getBytes)),
         baseOffset = startOffset, partitionLeaderEpoch = epoch)
@@ -3096,7 +3759,7 @@ class LogTest {
    * Append a bunch of messages to a log and then re-open it with recovery and check that the leader epochs are recovered properly.
    */
   @Test
-  def testLogRecoversForLeaderEpoch() {
+  def testLogRecoversForLeaderEpoch(): Unit = {
     val log = createLog(logDir, LogConfig())
     val leaderEpochCache = epochCache(log)
     val firstBatch = singletonRecordsWithLeaderEpoch(value = "random".getBytes, leaderEpoch = 1, offset = 0)
@@ -3141,12 +3804,13 @@ class LogTest {
 
     val buf = ByteBuffer.allocate(DefaultRecordBatch.sizeInBytes(records.asJava))
     val builder = MemoryRecords.builder(buf, magicValue, codec, TimestampType.CREATE_TIME, offset,
-      System.currentTimeMillis, leaderEpoch)
+      mockTime.milliseconds, leaderEpoch)
     records.foreach(builder.append)
     builder.build()
   }
 
-  def testFirstUnstableOffsetNoTransactionalData() {
+  @Test
+  def testFirstUnstableOffsetNoTransactionalData(): Unit = {
     val logConfig = LogTest.createLogConfig(segmentBytes = 1024 * 1024 * 5)
     val log = createLog(logDir, logConfig)
 
@@ -3160,7 +3824,7 @@ class LogTest {
   }
 
   @Test
-  def testFirstUnstableOffsetWithTransactionalData() {
+  def testFirstUnstableOffsetWithTransactionalData(): Unit = {
     val logConfig = LogTest.createLogConfig(segmentBytes = 1024 * 1024 * 5)
     val log = createLog(logDir, logConfig)
 
@@ -3175,7 +3839,7 @@ class LogTest {
       new SimpleRecord("baz".getBytes))
 
     val firstAppendInfo = log.appendAsLeader(records, leaderEpoch = 0)
-    assertEquals(Some(firstAppendInfo.firstOffset.get), log.firstUnstableOffset.map(_.messageOffset))
+    assertEquals(firstAppendInfo.firstOffset, log.firstUnstableOffset)
 
     // add more transactional records
     seq += 3
@@ -3183,20 +3847,76 @@ class LogTest {
       new SimpleRecord("blah".getBytes)), leaderEpoch = 0)
 
     // LSO should not have changed
-    assertEquals(Some(firstAppendInfo.firstOffset.get), log.firstUnstableOffset.map(_.messageOffset))
+    assertEquals(firstAppendInfo.firstOffset, log.firstUnstableOffset)
 
     // now transaction is committed
-    val commitAppendInfo = log.appendAsLeader(endTxnRecords(ControlRecordType.COMMIT, pid, epoch),
-      isFromClient = false, leaderEpoch = 0)
+    val commitAppendInfo = appendEndTxnMarkerAsLeader(log, pid, epoch, ControlRecordType.COMMIT)
 
     // first unstable offset is not updated until the high watermark is advanced
-    assertEquals(Some(firstAppendInfo.firstOffset.get), log.firstUnstableOffset.map(_.messageOffset))
-    log.onHighWatermarkIncremented(commitAppendInfo.lastOffset + 1)
+    assertEquals(firstAppendInfo.firstOffset, log.firstUnstableOffset)
+    log.updateHighWatermark(commitAppendInfo.lastOffset + 1)
 
     // now there should be no first unstable offset
     assertEquals(None, log.firstUnstableOffset)
   }
 
+  @Test
+  def testReadCommittedWithConcurrentHighWatermarkUpdates(): Unit = {
+    val logConfig = LogTest.createLogConfig(segmentBytes = 1024 * 1024 * 5)
+    val log = createLog(logDir, logConfig)
+    val lastOffset = 50L
+
+    val producerEpoch = 0.toShort
+    val producerId = 15L
+    val appendProducer = appendTransactionalAsLeader(log, producerId, producerEpoch)
+
+    // Thread 1 writes single-record transactions and attempts to read them
+    // before they have been aborted, and then aborts them
+    val txnWriteAndReadLoop: Callable[Int] = new Callable[Int]() {
+      override def call(): Int = {
+        var nonEmptyReads = 0
+        while (log.logEndOffset < lastOffset) {
+          val currentLogEndOffset = log.logEndOffset
+
+          appendProducer(1)
+
+          val readInfo = log.read(
+            startOffset = currentLogEndOffset,
+            maxLength = Int.MaxValue,
+            isolation = FetchTxnCommitted,
+            minOneMessage = false)
+
+          if (readInfo.records.sizeInBytes() > 0)
+            nonEmptyReads += 1
+
+          appendEndTxnMarkerAsLeader(log, producerId, producerEpoch, ControlRecordType.ABORT)
+        }
+        nonEmptyReads
+      }
+    }
+
+    // Thread 2 watches the log and updates the high watermark
+    val hwUpdateLoop: Runnable = new Runnable() {
+      override def run(): Unit = {
+        while (log.logEndOffset < lastOffset) {
+          log.updateHighWatermark(log.logEndOffset)
+        }
+      }
+    }
+
+    val executor = Executors.newFixedThreadPool(2)
+    try {
+      executor.submit(hwUpdateLoop)
+
+      val future = executor.submit(txnWriteAndReadLoop)
+      val nonEmptyReads = future.get()
+
+      assertEquals(0, nonEmptyReads)
+    } finally {
+      executor.shutdownNow()
+    }
+  }
+
   @Test
   def testTransactionIndexUpdated(): Unit = {
     val logConfig = LogTest.createLogConfig(segmentBytes = 1024 * 1024 * 5)
@@ -3235,7 +3955,21 @@ class LogTest {
     appendEndTxnMarkerAsLeader(log, pid4, epoch, ControlRecordType.COMMIT) // 90
 
     val abortedTransactions = allAbortedTransactions(log)
-    assertEquals(List(new AbortedTxn(pid1, 0L, 29L, 8L), new AbortedTxn(pid2, 8L, 74L, 36L)), abortedTransactions)
+    val expectedTransactions = List(
+      new AbortedTxn(pid1, 0L, 29L, 8L),
+      new AbortedTxn(pid2, 8L, 74L, 36L)
+    )
+    assertEquals(expectedTransactions, abortedTransactions)
+
+    // Verify caching of the segment position of the first unstable offset
+    log.updateHighWatermark(30L)
+    assertCachedFirstUnstableOffset(log, expectedOffset = 8L)
+
+    log.updateHighWatermark(75L)
+    assertCachedFirstUnstableOffset(log, expectedOffset = 36L)
+
+    log.updateHighWatermark(log.logEndOffset)
+    assertEquals(None, log.firstUnstableOffset)
   }
 
   @Test
@@ -3436,7 +4170,52 @@ class LogTest {
     appendAsFollower(log, MemoryRecords.readableRecords(buffer))
 
     val abortedTransactions = allAbortedTransactions(log)
-    assertEquals(List(new AbortedTxn(pid1, 0L, 29L, 8L), new AbortedTxn(pid2, 8L, 74L, 36L)), abortedTransactions)
+    val expectedTransactions = List(
+      new AbortedTxn(pid1, 0L, 29L, 8L),
+      new AbortedTxn(pid2, 8L, 74L, 36L)
+    )
+
+    assertEquals(expectedTransactions, abortedTransactions)
+
+    // Verify caching of the segment position of the first unstable offset
+    log.updateHighWatermark(30L)
+    assertCachedFirstUnstableOffset(log, expectedOffset = 8L)
+
+    log.updateHighWatermark(75L)
+    assertCachedFirstUnstableOffset(log, expectedOffset = 36L)
+
+    log.updateHighWatermark(log.logEndOffset)
+    assertEquals(None, log.firstUnstableOffset)
+  }
+
+  private def assertCachedFirstUnstableOffset(log: Log, expectedOffset: Long): Unit = {
+    assertTrue(log.producerStateManager.firstUnstableOffset.isDefined)
+    val firstUnstableOffset = log.producerStateManager.firstUnstableOffset.get
+    assertEquals(expectedOffset, firstUnstableOffset.messageOffset)
+    assertFalse(firstUnstableOffset.messageOffsetOnly)
+    assertValidLogOffsetMetadata(log, firstUnstableOffset)
+  }
+
+  private def assertValidLogOffsetMetadata(log: Log, offsetMetadata: LogOffsetMetadata): Unit = {
+    assertFalse(offsetMetadata.messageOffsetOnly)
+
+    val segmentBaseOffset = offsetMetadata.segmentBaseOffset
+    val segmentOpt = log.logSegments(segmentBaseOffset, segmentBaseOffset + 1).headOption
+    assertTrue(segmentOpt.isDefined)
+
+    val segment = segmentOpt.get
+    assertEquals(segmentBaseOffset, segment.baseOffset)
+    assertTrue(offsetMetadata.relativePositionInSegment <= segment.size)
+
+    val readInfo = segment.read(offsetMetadata.messageOffset,
+      maxSize = 2048,
+      maxPosition = segment.size,
+      minOneMessage = false)
+
+    if (offsetMetadata.relativePositionInSegment < segment.size)
+      assertEquals(offsetMetadata, readInfo.fetchOffsetMetadata)
+    else
+      assertNull(readInfo)
   }
 
   @Test(expected = classOf[TransactionCoordinatorFencedException])
@@ -3458,7 +4237,43 @@ class LogTest {
   }
 
   @Test
-  def testFirstUnstableOffsetDoesNotExceedLogStartOffsetMidSegment(): Unit = {
+  def testZombieCoordinatorFencedEmptyTransaction(): Unit = {
+    val pid = 1L
+    val epoch = 0.toShort
+    val logConfig = LogTest.createLogConfig(segmentBytes = 1024 * 1024 * 5)
+    val log = createLog(logDir, logConfig)
+
+    val buffer = ByteBuffer.allocate(256)
+    val append = appendTransactionalToBuffer(buffer, pid, epoch, leaderEpoch = 1)
+    append(0, 10)
+    appendEndTxnMarkerToBuffer(buffer, pid, epoch, 10L, ControlRecordType.COMMIT,
+      coordinatorEpoch = 0, leaderEpoch = 1)
+
+    buffer.flip()
+    log.appendAsFollower(MemoryRecords.readableRecords(buffer))
+
+    appendEndTxnMarkerAsLeader(log, pid, epoch, ControlRecordType.ABORT, coordinatorEpoch = 2, leaderEpoch = 1)
+    appendEndTxnMarkerAsLeader(log, pid, epoch, ControlRecordType.ABORT, coordinatorEpoch = 2, leaderEpoch = 1)
+    assertThrows[TransactionCoordinatorFencedException] {
+      appendEndTxnMarkerAsLeader(log, pid, epoch, ControlRecordType.ABORT, coordinatorEpoch = 1, leaderEpoch = 1)
+    }
+  }
+
+  @Test
+  def testEndTxnWithFencedProducerEpoch(): Unit = {
+    val producerId = 1L
+    val epoch = 5.toShort
+    val logConfig = LogTest.createLogConfig(segmentBytes = 1024 * 1024 * 5)
+    val log = createLog(logDir, logConfig)
+    appendEndTxnMarkerAsLeader(log, producerId, epoch, ControlRecordType.ABORT, coordinatorEpoch = 1)
+
+    assertThrows[ProducerFencedException] {
+      appendEndTxnMarkerAsLeader(log, producerId, (epoch - 1).toShort, ControlRecordType.ABORT, coordinatorEpoch = 1)
+    }
+  }
+
+  @Test
+  def testLastStableOffsetDoesNotExceedLogStartOffsetMidSegment(): Unit = {
     val logConfig = LogTest.createLogConfig(segmentBytes = 1024 * 1024 * 5)
     val log = createLog(logDir, logConfig)
     val epoch = 0.toShort
@@ -3473,16 +4288,17 @@ class LogTest {
     assertEquals(2, log.logSegments.size)
     appendPid(5)
 
-    assertEquals(Some(0L), log.firstUnstableOffset.map(_.messageOffset))
+    assertEquals(Some(0L), log.firstUnstableOffset)
 
+    log.updateHighWatermark(log.logEndOffset)
     log.maybeIncrementLogStartOffset(5L)
 
     // the first unstable offset should be lower bounded by the log start offset
-    assertEquals(Some(5L), log.firstUnstableOffset.map(_.messageOffset))
+    assertEquals(Some(5L), log.firstUnstableOffset)
   }
 
   @Test
-  def testFirstUnstableOffsetDoesNotExceedLogStartOffsetAfterSegmentDeletion(): Unit = {
+  def testLastStableOffsetDoesNotExceedLogStartOffsetAfterSegmentDeletion(): Unit = {
     val logConfig = LogTest.createLogConfig(segmentBytes = 1024 * 1024 * 5)
     val log = createLog(logDir, logConfig)
     val epoch = 0.toShort
@@ -3497,19 +4313,83 @@ class LogTest {
     assertEquals(2, log.logSegments.size)
     appendPid(5)
 
-    assertEquals(Some(0L), log.firstUnstableOffset.map(_.messageOffset))
+    assertEquals(Some(0L), log.firstUnstableOffset)
 
+    log.updateHighWatermark(log.logEndOffset)
     log.maybeIncrementLogStartOffset(8L)
-    log.onHighWatermarkIncremented(log.logEndOffset)
+    log.updateHighWatermark(log.logEndOffset)
     log.deleteOldSegments()
     assertEquals(1, log.logSegments.size)
 
     // the first unstable offset should be lower bounded by the log start offset
-    assertEquals(Some(8L), log.firstUnstableOffset.map(_.messageOffset))
+    assertEquals(Some(8L), log.firstUnstableOffset)
+  }
+
+  @Test
+  def testAppendToTransactionIndexFailure(): Unit = {
+    val pid = 1L
+    val epoch = 0.toShort
+    val logConfig = LogTest.createLogConfig(segmentBytes = 1024 * 1024 * 5)
+    val log = createLog(logDir, logConfig)
+
+    val append = appendTransactionalAsLeader(log, pid, epoch)
+    append(10)
+
+    // Kind of a hack, but renaming the index to a directory ensures that the append
+    // to the index will fail.
+    log.activeSegment.txnIndex.renameTo(log.dir)
+
+    // The append will be written to the log successfully, but the write to the index will fail
+    assertThrows[KafkaStorageException] {
+      appendEndTxnMarkerAsLeader(log, pid, epoch, ControlRecordType.ABORT, coordinatorEpoch = 1)
+    }
+    assertEquals(11L, log.logEndOffset)
+    assertEquals(0L, log.lastStableOffset)
+
+    // Try the append a second time. The appended offset in the log should still increase.
+    assertThrows[KafkaStorageException] {
+      appendEndTxnMarkerAsLeader(log, pid, epoch, ControlRecordType.ABORT, coordinatorEpoch = 1)
+    }
+    assertEquals(12L, log.logEndOffset)
+    assertEquals(0L, log.lastStableOffset)
+
+    // Even if the high watermark is updated, the first unstable offset does not move
+    log.updateHighWatermark(12L)
+    assertEquals(0L, log.lastStableOffset)
+
+    log.close()
+
+    val reopenedLog = createLog(logDir, logConfig)
+    assertEquals(12L, reopenedLog.logEndOffset)
+    assertEquals(2, reopenedLog.activeSegment.txnIndex.allAbortedTxns.size)
+    reopenedLog.updateHighWatermark(12L)
+    assertEquals(None, reopenedLog.firstUnstableOffset)
   }
 
   @Test
-  def testLastStableOffsetWithMixedProducerData() {
+  def testOffsetSnapshot(): Unit = {
+    val logConfig = LogTest.createLogConfig(segmentBytes = 1024 * 1024 * 5)
+    val log = createLog(logDir, logConfig)
+
+    // append a few records
+    appendAsFollower(log, MemoryRecords.withRecords(CompressionType.NONE,
+      new SimpleRecord("a".getBytes),
+      new SimpleRecord("b".getBytes),
+      new SimpleRecord("c".getBytes)), 5)
+
+
+    log.updateHighWatermark(2L)
+    var offsets: LogOffsetSnapshot = log.fetchOffsetSnapshot
+    assertEquals(offsets.highWatermark.messageOffset, 2L)
+    assertFalse(offsets.highWatermark.messageOffsetOnly)
+
+    offsets = log.fetchOffsetSnapshot
+    assertEquals(offsets.highWatermark.messageOffset, 2L)
+    assertFalse(offsets.highWatermark.messageOffsetOnly)
+  }
+
+  @Test
+  def testLastStableOffsetWithMixedProducerData(): Unit = {
     val logConfig = LogTest.createLogConfig(segmentBytes = 1024 * 1024 * 5)
     val log = createLog(logDir, logConfig)
 
@@ -3526,7 +4406,7 @@ class LogTest {
       new SimpleRecord("a".getBytes),
       new SimpleRecord("b".getBytes),
       new SimpleRecord("c".getBytes)), leaderEpoch = 0)
-    assertEquals(Some(firstAppendInfo.firstOffset.get), log.firstUnstableOffset.map(_.messageOffset))
+    assertEquals(firstAppendInfo.firstOffset, log.firstUnstableOffset)
 
     // mix in some non-transactional data
     log.appendAsLeader(MemoryRecords.withRecords(CompressionType.NONE,
@@ -3541,27 +4421,25 @@ class LogTest {
       new SimpleRecord("f".getBytes)), leaderEpoch = 0)
 
     // LSO should not have changed
-    assertEquals(Some(firstAppendInfo.firstOffset.get), log.firstUnstableOffset.map(_.messageOffset))
+    assertEquals(firstAppendInfo.firstOffset, log.firstUnstableOffset)
 
     // now first producer's transaction is aborted
-    val abortAppendInfo = log.appendAsLeader(endTxnRecords(ControlRecordType.ABORT, pid1, epoch),
-      isFromClient = false, leaderEpoch = 0)
-    log.onHighWatermarkIncremented(abortAppendInfo.lastOffset + 1)
+    val abortAppendInfo = appendEndTxnMarkerAsLeader(log, pid1, epoch, ControlRecordType.ABORT)
+    log.updateHighWatermark(abortAppendInfo.lastOffset + 1)
 
     // LSO should now point to one less than the first offset of the second transaction
-    assertEquals(Some(secondAppendInfo.firstOffset.get), log.firstUnstableOffset.map(_.messageOffset))
+    assertEquals(secondAppendInfo.firstOffset, log.firstUnstableOffset)
 
     // commit the second transaction
-    val commitAppendInfo = log.appendAsLeader(endTxnRecords(ControlRecordType.COMMIT, pid2, epoch),
-      isFromClient = false, leaderEpoch = 0)
-    log.onHighWatermarkIncremented(commitAppendInfo.lastOffset + 1)
+    val commitAppendInfo = appendEndTxnMarkerAsLeader(log, pid2, epoch, ControlRecordType.COMMIT)
+    log.updateHighWatermark(commitAppendInfo.lastOffset + 1)
 
     // now there should be no first unstable offset
     assertEquals(None, log.firstUnstableOffset)
   }
 
   @Test
-  def testAbortedTransactionSpanningMultipleSegments() {
+  def testAbortedTransactionSpanningMultipleSegments(): Unit = {
     val pid = 137L
     val epoch = 5.toShort
     var seq = 0
@@ -3575,8 +4453,7 @@ class LogTest {
     val log = createLog(logDir, logConfig)
 
     val firstAppendInfo = log.appendAsLeader(records, leaderEpoch = 0)
-    assertEquals(Some(firstAppendInfo.firstOffset.get), log.firstUnstableOffset.map(_.messageOffset))
-    assertEquals(Some(0L), log.firstUnstableOffset.map(_.segmentBaseOffset))
+    assertEquals(firstAppendInfo.firstOffset, log.firstUnstableOffset)
 
     // this write should spill to the second segment
     seq = 3
@@ -3584,31 +4461,42 @@ class LogTest {
       new SimpleRecord("d".getBytes),
       new SimpleRecord("e".getBytes),
       new SimpleRecord("f".getBytes)), leaderEpoch = 0)
-    assertEquals(Some(firstAppendInfo.firstOffset.get), log.firstUnstableOffset.map(_.messageOffset))
-    assertEquals(Some(0L), log.firstUnstableOffset.map(_.segmentBaseOffset))
+    assertEquals(firstAppendInfo.firstOffset, log.firstUnstableOffset)
     assertEquals(3L, log.logEndOffsetMetadata.segmentBaseOffset)
 
     // now abort the transaction
-    val appendInfo = log.appendAsLeader(endTxnRecords(ControlRecordType.ABORT, pid, epoch),
-      isFromClient = false, leaderEpoch = 0)
-    log.onHighWatermarkIncremented(appendInfo.lastOffset + 1)
-    assertEquals(None, log.firstUnstableOffset.map(_.messageOffset))
+    val abortAppendInfo = appendEndTxnMarkerAsLeader(log, pid, epoch, ControlRecordType.ABORT)
+    log.updateHighWatermark(abortAppendInfo.lastOffset + 1)
+    assertEquals(None, log.firstUnstableOffset)
 
     // now check that a fetch includes the aborted transaction
-    val fetchDataInfo = log.read(0L, 2048, maxOffset = None, minOneMessage = true, includeAbortedTxns = true)
+    val fetchDataInfo = log.read(0L,
+      maxLength = 2048,
+      isolation = FetchTxnCommitted,
+      minOneMessage = true)
     assertEquals(1, fetchDataInfo.abortedTransactions.size)
 
     assertTrue(fetchDataInfo.abortedTransactions.isDefined)
     assertEquals(new AbortedTransaction(pid, 0), fetchDataInfo.abortedTransactions.get.head)
   }
 
- private def allAbortedTransactions(log: Log) = log.logSegments.flatMap(_.txnIndex.allAbortedTxns)
+  @Test
+  def testLoadPartitionDirWithNoSegmentsShouldNotThrow(): Unit = {
+    val dirName = Log.logDeleteDirName(new TopicPartition("foo", 3))
+    val logDir = new File(tmpDir, dirName)
+    logDir.mkdirs()
+    val logConfig = LogTest.createLogConfig()
+    val log = createLog(logDir, logConfig)
+    assertEquals(1, log.numberOfSegments)
+  }
+
+  private def allAbortedTransactions(log: Log) = log.logSegments.flatMap(_.txnIndex.allAbortedTxns)
 
   private def appendTransactionalAsLeader(log: Log, producerId: Long, producerEpoch: Short): Int => Unit = {
     var sequence = 0
     numRecords: Int => {
       val simpleRecords = (sequence until sequence + numRecords).map { seq =>
-        new SimpleRecord(s"$seq".getBytes)
+        new SimpleRecord(mockTime.milliseconds(), s"$seq".getBytes)
       }
       val records = MemoryRecords.withTransactionalRecords(CompressionType.NONE, producerId,
         producerEpoch, sequence, simpleRecords: _*)
@@ -3617,10 +4505,16 @@ class LogTest {
     }
   }
 
-  private def appendEndTxnMarkerAsLeader(log: Log, producerId: Long, producerEpoch: Short,
-                                         controlType: ControlRecordType, coordinatorEpoch: Int = 0): Unit = {
-    val records = endTxnRecords(controlType, producerId, producerEpoch, coordinatorEpoch = coordinatorEpoch)
-    log.appendAsLeader(records, isFromClient = false, leaderEpoch = 0)
+  private def appendEndTxnMarkerAsLeader(log: Log,
+                                         producerId: Long,
+                                         producerEpoch: Short,
+                                         controlType: ControlRecordType,
+                                         coordinatorEpoch: Int = 0,
+                                         leaderEpoch: Int = 0,
+                                         timestamp: Long = mockTime.milliseconds()): LogAppendInfo = {
+    val records = endTxnRecords(controlType, producerId, producerEpoch,
+      coordinatorEpoch = coordinatorEpoch, timestamp = timestamp)
+    log.appendAsLeader(records, origin = AppendOrigin.Coordinator, leaderEpoch = leaderEpoch)
   }
 
   private def appendNonTransactionalAsLeader(log: Log, numRecords: Int): Unit = {
@@ -3631,10 +4525,14 @@ class LogTest {
     log.appendAsLeader(records, leaderEpoch = 0)
   }
 
-  private def appendTransactionalToBuffer(buffer: ByteBuffer, producerId: Long, producerEpoch: Short): (Long, Int) => Unit = {
+  private def appendTransactionalToBuffer(buffer: ByteBuffer,
+                                          producerId: Long,
+                                          producerEpoch: Short,
+                                          leaderEpoch: Int = 0): (Long, Int) => Unit = {
     var sequence = 0
     (offset: Long, numRecords: Int) => {
-      val builder = MemoryRecords.builder(buffer, CompressionType.NONE, offset, producerId, producerEpoch, sequence, true)
+      val builder = MemoryRecords.builder(buffer, RecordBatch.CURRENT_MAGIC_VALUE, CompressionType.NONE, TimestampType.CREATE_TIME,
+        offset, mockTime.milliseconds(), producerId, producerEpoch, sequence, true, leaderEpoch)
       for (seq <- sequence until sequence + numRecords) {
         val record = new SimpleRecord(s"$seq".getBytes)
         builder.append(record)
@@ -3645,10 +4543,15 @@ class LogTest {
     }
   }
 
-  private def appendEndTxnMarkerToBuffer(buffer: ByteBuffer, producerId: Long, producerEpoch: Short, offset: Long,
-                                         controlType: ControlRecordType, coordinatorEpoch: Int = 0): Unit = {
+  private def appendEndTxnMarkerToBuffer(buffer: ByteBuffer,
+                                         producerId: Long,
+                                         producerEpoch: Short,
+                                         offset: Long,
+                                         controlType: ControlRecordType,
+                                         coordinatorEpoch: Int = 0,
+                                         leaderEpoch: Int = 0): Unit = {
     val marker = new EndTransactionMarker(controlType, coordinatorEpoch)
-    MemoryRecords.writeEndTransactionalMarker(buffer, offset, mockTime.milliseconds(), 0, producerId, producerEpoch, marker)
+    MemoryRecords.writeEndTransactionalMarker(buffer, offset, mockTime.milliseconds(), leaderEpoch, producerId, producerEpoch, marker)
   }
 
   private def appendNonTransactionalToBuffer(buffer: ByteBuffer, offset: Long, numRecords: Int): Unit = {
@@ -3712,10 +4615,12 @@ class LogTest {
       expectDeletedFiles)
   }
 
-  private def readLog(log: Log, startOffset: Long, maxLength: Int,
-                      maxOffset: Option[Long] = None,
+  private def readLog(log: Log,
+                      startOffset: Long,
+                      maxLength: Int,
+                      isolation: FetchIsolation = FetchLogEnd,
                       minOneMessage: Boolean = true): FetchDataInfo = {
-    log.read(startOffset, maxLength, maxOffset, minOneMessage, includeAbortedTxns = false)
+    log.read(startOffset, maxLength, isolation, minOneMessage)
   }
 
 }
@@ -3731,7 +4636,8 @@ object LogTest {
                       indexIntervalBytes: Int = Defaults.IndexInterval,
                       segmentIndexBytes: Int = Defaults.MaxIndexSize,
                       messageFormatVersion: String = Defaults.MessageFormatVersion,
-                      fileDeleteDelayMs: Long = Defaults.FileDeleteDelayMs): LogConfig = {
+                      fileDeleteDelayMs: Long = Defaults.FileDeleteDelayMs,
+                      producerBatchDecompressionEnable: Boolean = Defaults.ProducerBatchDecompressionEnable): LogConfig = {
     val logProps = new Properties()
 
     logProps.put(LogConfig.SegmentMsProp, segmentMs: java.lang.Long)
@@ -3745,6 +4651,7 @@ object LogTest {
     logProps.put(LogConfig.SegmentIndexBytesProp, segmentIndexBytes: Integer)
     logProps.put(LogConfig.MessageFormatVersionProp, messageFormatVersion)
     logProps.put(LogConfig.FileDeleteDelayMsProp, fileDeleteDelayMs: java.lang.Long)
+    logProps.put(LogConfig.ProducerBatchDecompressionEnableProp, producerBatchDecompressionEnable: java.lang.Boolean)
     LogConfig(logProps)
   }
 
diff --git a/core/src/test/scala/unit/kafka/log/LogUtils.scala b/core/src/test/scala/unit/kafka/log/LogUtils.scala
index 838a69974420e..8bf9812b86acd 100644
--- a/core/src/test/scala/unit/kafka/log/LogUtils.scala
+++ b/core/src/test/scala/unit/kafka/log/LogUtils.scala
@@ -31,8 +31,8 @@ object LogUtils {
                     indexIntervalBytes: Int = 10,
                     time: Time = Time.SYSTEM): LogSegment = {
     val ms = FileRecords.open(Log.logFile(logDir, offset))
-    val idx = new LazyOffsetIndex(Log.offsetIndexFile(logDir, offset), offset, maxIndexSize = 1000)
-    val timeIdx = new LazyTimeIndex(Log.timeIndexFile(logDir, offset), offset, maxIndexSize = 1500)
+    val idx = LazyIndex.forOffset(Log.offsetIndexFile(logDir, offset), offset, maxIndexSize = 1000)
+    val timeIdx = LazyIndex.forTime(Log.timeIndexFile(logDir, offset), offset, maxIndexSize = 1500)
     val txnIndex = new TransactionIndex(offset, Log.transactionIndexFile(logDir, offset))
 
     new LogSegment(ms, idx, timeIdx, txnIndex, offset, indexIntervalBytes, 0, time)
diff --git a/core/src/test/scala/unit/kafka/log/LogValidatorTest.scala b/core/src/test/scala/unit/kafka/log/LogValidatorTest.scala
index 37553b9376e1a..ba56f1449faee 100644
--- a/core/src/test/scala/unit/kafka/log/LogValidatorTest.scala
+++ b/core/src/test/scala/unit/kafka/log/LogValidatorTest.scala
@@ -19,33 +19,131 @@ package kafka.log
 import java.nio.ByteBuffer
 import java.util.concurrent.TimeUnit
 
-import kafka.api.{ApiVersion, KAFKA_2_0_IV1}
-import kafka.common.LongRef
+import com.yammer.metrics.Metrics
+import kafka.api.{ApiVersion, KAFKA_2_0_IV1, KAFKA_2_3_IV1}
+import kafka.common.{LongRef, RecordValidationException}
+import kafka.log.LogValidator.ValidationAndOffsetAssignResult
 import kafka.message._
+import kafka.server.BrokerTopicStats
+import kafka.utils.TestUtils.meterCount
 import org.apache.kafka.common.errors.{InvalidTimestampException, UnsupportedCompressionTypeException, UnsupportedForMessageFormatException}
 import org.apache.kafka.common.record._
 import org.apache.kafka.common.utils.Time
+import org.apache.kafka.common.{InvalidRecordException, TopicPartition}
 import org.apache.kafka.test.TestUtils
 import org.junit.Assert._
 import org.junit.Test
-import org.scalatest.Assertions.intercept
+import org.scalatest.Assertions.{assertThrows, intercept}
 
 import scala.collection.JavaConverters._
 
 class LogValidatorTest {
 
   val time = Time.SYSTEM
+  val topicPartition = new TopicPartition("topic", 0)
+  val brokerTopicStats = new BrokerTopicStats
+  val metricsKeySet = Metrics.defaultRegistry.allMetrics.keySet.asScala
 
   @Test
-  def testLogAppendTimeNonCompressedV1() {
+  def testOnlyOneBatch(): Unit = {
+    checkOnlyOneBatch(RecordBatch.MAGIC_VALUE_V0, CompressionType.GZIP, CompressionType.GZIP)
+    checkOnlyOneBatch(RecordBatch.MAGIC_VALUE_V1, CompressionType.GZIP, CompressionType.GZIP)
+    checkOnlyOneBatch(RecordBatch.MAGIC_VALUE_V2, CompressionType.GZIP, CompressionType.GZIP)
+    checkOnlyOneBatch(RecordBatch.MAGIC_VALUE_V0, CompressionType.GZIP, CompressionType.NONE)
+    checkOnlyOneBatch(RecordBatch.MAGIC_VALUE_V1, CompressionType.GZIP, CompressionType.NONE)
+    checkOnlyOneBatch(RecordBatch.MAGIC_VALUE_V2, CompressionType.GZIP, CompressionType.NONE)
+    checkOnlyOneBatch(RecordBatch.MAGIC_VALUE_V2, CompressionType.NONE, CompressionType.NONE)
+    checkOnlyOneBatch(RecordBatch.MAGIC_VALUE_V2, CompressionType.NONE, CompressionType.GZIP)
+  }
+
+  @Test
+  def testAllowMultiBatch(): Unit = {
+    checkAllowMultiBatch(RecordBatch.MAGIC_VALUE_V0, CompressionType.NONE, CompressionType.NONE)
+    checkAllowMultiBatch(RecordBatch.MAGIC_VALUE_V1, CompressionType.NONE, CompressionType.NONE)
+    checkAllowMultiBatch(RecordBatch.MAGIC_VALUE_V0, CompressionType.NONE, CompressionType.GZIP)
+    checkAllowMultiBatch(RecordBatch.MAGIC_VALUE_V1, CompressionType.NONE, CompressionType.GZIP)
+  }
+
+  @Test
+  def testValidationOfBatchesWithNonSequentialInnerOffsets(): Unit = {
+    def testMessageValidation(magicValue: Byte): Unit = {
+      val numRecords = 20
+      val invalidRecords = recordsWithNonSequentialInnerOffsets(magicValue, CompressionType.GZIP, numRecords)
+
+      // Validation for v2 and above is strict for this case. For older formats, we fix invalid
+      // internal offsets by rewriting the batch.
+      if (magicValue >= RecordBatch.MAGIC_VALUE_V2) {
+        assertThrows[InvalidRecordException] {
+          validateMessages(invalidRecords, magicValue, CompressionType.GZIP, CompressionType.GZIP)
+        }
+      } else {
+        val result = validateMessages(invalidRecords, magicValue, CompressionType.GZIP, CompressionType.GZIP)
+        assertEquals(0 until numRecords, result.validatedRecords.records.asScala.map(_.offset))
+      }
+    }
+
+    for (version <- RecordVersion.values) {
+      testMessageValidation(version.value)
+    }
+  }
+
+  @Test
+  def testMisMatchMagic(): Unit = {
+    checkMismatchMagic(RecordBatch.MAGIC_VALUE_V0, RecordBatch.MAGIC_VALUE_V1, CompressionType.GZIP)
+    checkMismatchMagic(RecordBatch.MAGIC_VALUE_V1, RecordBatch.MAGIC_VALUE_V0, CompressionType.GZIP)
+  }
+
+  private def checkOnlyOneBatch(magic: Byte, sourceCompressionType: CompressionType, targetCompressionType: CompressionType): Unit = {
+    assertThrows[InvalidRecordException] {
+      validateMessages(createTwoBatchedRecords(magic, 0L, sourceCompressionType), magic, sourceCompressionType, targetCompressionType)
+    }
+  }
+
+  private def checkAllowMultiBatch(magic: Byte, sourceCompressionType: CompressionType, targetCompressionType: CompressionType): Unit = {
+    validateMessages(createTwoBatchedRecords(magic, 0L, sourceCompressionType), magic, sourceCompressionType, targetCompressionType)
+  }
+
+  private def checkMismatchMagic(batchMagic: Byte, recordMagic: Byte, compressionType: CompressionType): Unit = {
+    assertThrows[RecordValidationException] {
+      validateMessages(recordsWithInvalidInnerMagic(batchMagic, recordMagic, compressionType), batchMagic, compressionType, compressionType)
+    }
+    assertEquals(metricsKeySet.count(_.getMBeanName.endsWith(s"${BrokerTopicStats.InvalidMagicNumberRecordsPerSec}")), 1)
+    assertTrue(meterCount(s"${BrokerTopicStats.InvalidMagicNumberRecordsPerSec}") > 0)
+  }
+
+  private def validateMessages(records: MemoryRecords,
+                               magic: Byte,
+                               sourceCompressionType: CompressionType,
+                               targetCompressionType: CompressionType): ValidationAndOffsetAssignResult = {
+    LogValidator.validateMessagesAndAssignOffsets(records,
+      topicPartition,
+      new LongRef(0L),
+      time,
+      now = 0L,
+      CompressionCodec.getCompressionCodec(sourceCompressionType.name),
+      CompressionCodec.getCompressionCodec(targetCompressionType.name),
+      compactedTopic = false,
+      magic,
+      TimestampType.CREATE_TIME,
+      1000L,
+      RecordBatch.NO_PRODUCER_EPOCH,
+      origin = AppendOrigin.Client,
+      KAFKA_2_3_IV1,
+      brokerTopicStats
+    )
+  }
+
+  @Test
+  def testLogAppendTimeNonCompressedV1(): Unit = {
     checkLogAppendTimeNonCompressed(RecordBatch.MAGIC_VALUE_V1)
   }
 
-  private def checkLogAppendTimeNonCompressed(magic: Byte) {
+  private def checkLogAppendTimeNonCompressed(magic: Byte): Unit = {
     val now = System.currentTimeMillis()
     // The timestamps should be overwritten
     val records = createRecords(magicValue = magic, timestamp = 1234L, codec = CompressionType.NONE)
     val validatedResults = LogValidator.validateMessagesAndAssignOffsets(records,
+      topicPartition,
       offsetCounter = new LongRef(0),
       time= time,
       now = now,
@@ -56,8 +154,9 @@ class LogValidatorTest {
       timestampType = TimestampType.LOG_APPEND_TIME,
       timestampDiffMaxMs = 1000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true,
-      interBrokerProtocolVersion = ApiVersion.latestVersion)
+      origin = AppendOrigin.Client,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      brokerTopicStats = brokerTopicStats)
     val validatedRecords = validatedResults.validatedRecords
     assertEquals("message set size should not change", records.records.asScala.size, validatedRecords.records.asScala.size)
     validatedRecords.batches.asScala.foreach(batch => validateLogAppendTime(now, 1234L, batch))
@@ -69,21 +168,22 @@ class LogValidatorTest {
       compressed = false)
   }
 
-  def testLogAppendTimeNonCompressedV2() {
+  def testLogAppendTimeNonCompressedV2(): Unit = {
     checkLogAppendTimeNonCompressed(RecordBatch.MAGIC_VALUE_V2)
   }
 
   @Test
-  def testLogAppendTimeWithRecompressionV1() {
+  def testLogAppendTimeWithRecompressionV1(): Unit = {
     checkLogAppendTimeWithRecompression(RecordBatch.MAGIC_VALUE_V1)
   }
 
-  private def checkLogAppendTimeWithRecompression(targetMagic: Byte) {
+  private def checkLogAppendTimeWithRecompression(targetMagic: Byte): Unit = {
     val now = System.currentTimeMillis()
     // The timestamps should be overwritten
     val records = createRecords(magicValue = RecordBatch.MAGIC_VALUE_V0, codec = CompressionType.GZIP)
     val validatedResults = LogValidator.validateMessagesAndAssignOffsets(
       records,
+      topicPartition,
       offsetCounter = new LongRef(0),
       time = time,
       now = now,
@@ -94,8 +194,9 @@ class LogValidatorTest {
       timestampType = TimestampType.LOG_APPEND_TIME,
       timestampDiffMaxMs = 1000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true,
-      interBrokerProtocolVersion = ApiVersion.latestVersion)
+      origin = AppendOrigin.Client,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      brokerTopicStats = brokerTopicStats)
     val validatedRecords = validatedResults.validatedRecords
 
     assertEquals("message set size should not change", records.records.asScala.size, validatedRecords.records.asScala.size)
@@ -111,21 +212,22 @@ class LogValidatorTest {
   }
 
   @Test
-  def testLogAppendTimeWithRecompressionV2() {
+  def testLogAppendTimeWithRecompressionV2(): Unit = {
     checkLogAppendTimeWithRecompression(RecordBatch.MAGIC_VALUE_V2)
   }
 
   @Test
-  def testLogAppendTimeWithoutRecompressionV1() {
+  def testLogAppendTimeWithoutRecompressionV1(): Unit = {
     checkLogAppendTimeWithoutRecompression(RecordBatch.MAGIC_VALUE_V1)
   }
 
-  private def checkLogAppendTimeWithoutRecompression(magic: Byte) {
+  private def checkLogAppendTimeWithoutRecompression(magic: Byte): Unit = {
     val now = System.currentTimeMillis()
     // The timestamps should be overwritten
     val records = createRecords(magicValue = magic, timestamp = 1234L, codec = CompressionType.GZIP)
     val validatedResults = LogValidator.validateMessagesAndAssignOffsets(
       records,
+      topicPartition,
       offsetCounter = new LongRef(0),
       time = time,
       now = now,
@@ -136,8 +238,9 @@ class LogValidatorTest {
       timestampType = TimestampType.LOG_APPEND_TIME,
       timestampDiffMaxMs = 1000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true,
-      interBrokerProtocolVersion = ApiVersion.latestVersion)
+      origin = AppendOrigin.Client,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      brokerTopicStats = brokerTopicStats)
     val validatedRecords = validatedResults.validatedRecords
 
     assertEquals("message set size should not change", records.records.asScala.size,
@@ -178,12 +281,13 @@ class LogValidatorTest {
     }
   }
 
-  private def validateRecordBatchWithCountOverrides(lastOffsetDelta: Int, count: Int) {
+  private def validateRecordBatchWithCountOverrides(lastOffsetDelta: Int, count: Int): Unit = {
     val records = createRecords(magicValue = RecordBatch.MAGIC_VALUE_V2, timestamp = 1234L, codec = CompressionType.NONE)
     records.buffer.putInt(DefaultRecordBatch.RECORDS_COUNT_OFFSET, count)
     records.buffer.putInt(DefaultRecordBatch.LAST_OFFSET_DELTA_OFFSET, lastOffsetDelta)
     LogValidator.validateMessagesAndAssignOffsets(
       records,
+      topicPartition,
       offsetCounter = new LongRef(0),
       time = time,
       now = time.milliseconds(),
@@ -194,21 +298,22 @@ class LogValidatorTest {
       timestampType = TimestampType.LOG_APPEND_TIME,
       timestampDiffMaxMs = 1000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true,
-      interBrokerProtocolVersion = ApiVersion.latestVersion)
+      origin = AppendOrigin.Client,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      brokerTopicStats = brokerTopicStats)
   }
 
   @Test
-  def testLogAppendTimeWithoutRecompressionV2() {
+  def testLogAppendTimeWithoutRecompressionV2(): Unit = {
     checkLogAppendTimeWithoutRecompression(RecordBatch.MAGIC_VALUE_V2)
   }
 
   @Test
-  def testNonCompressedV1() {
+  def testNonCompressedV1(): Unit = {
     checkNonCompressed(RecordBatch.MAGIC_VALUE_V1)
   }
 
-  private def checkNonCompressed(magic: Byte) {
+  private def checkNonCompressed(magic: Byte): Unit = {
     val now = System.currentTimeMillis()
     val timestampSeq = Seq(now - 1, now + 1, now)
 
@@ -226,6 +331,7 @@ class LogValidatorTest {
       new SimpleRecord(timestampSeq(2), "beautiful".getBytes))
 
     val validatingResults = LogValidator.validateMessagesAndAssignOffsets(records,
+      topicPartition,
       offsetCounter = new LongRef(0),
       time = time,
       now = System.currentTimeMillis(),
@@ -236,8 +342,9 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 1000L,
       partitionLeaderEpoch = partitionLeaderEpoch,
-      isFromClient = true,
-      interBrokerProtocolVersion = ApiVersion.latestVersion)
+      origin = AppendOrigin.Client,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      brokerTopicStats = brokerTopicStats)
     val validatedRecords = validatingResults.validatedRecords
 
     var i = 0
@@ -266,7 +373,7 @@ class LogValidatorTest {
   }
 
   @Test
-  def testNonCompressedV2() {
+  def testNonCompressedV2(): Unit = {
     checkNonCompressed(RecordBatch.MAGIC_VALUE_V2)
   }
 
@@ -293,6 +400,7 @@ class LogValidatorTest {
       new SimpleRecord(timestampSeq(2), "beautiful".getBytes))
 
     val validatingResults = LogValidator.validateMessagesAndAssignOffsets(records,
+      topicPartition,
       offsetCounter = new LongRef(0),
       time = time,
       now = System.currentTimeMillis(),
@@ -303,8 +411,9 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 1000L,
       partitionLeaderEpoch = partitionLeaderEpoch,
-      isFromClient = true,
-      interBrokerProtocolVersion = ApiVersion.latestVersion)
+      origin = AppendOrigin.Client,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      brokerTopicStats = brokerTopicStats)
     val validatedRecords = validatingResults.validatedRecords
 
     var i = 0
@@ -341,9 +450,10 @@ class LogValidatorTest {
     checkCreateTimeUpConversionFromV0(RecordBatch.MAGIC_VALUE_V1)
   }
 
-  private def checkCreateTimeUpConversionFromV0(toMagic: Byte) {
+  private def checkCreateTimeUpConversionFromV0(toMagic: Byte): Unit = {
     val records = createRecords(magicValue = RecordBatch.MAGIC_VALUE_V0, codec = CompressionType.GZIP)
     val validatedResults = LogValidator.validateMessagesAndAssignOffsets(records,
+      topicPartition,
       offsetCounter = new LongRef(0),
       time = time,
       now = System.currentTimeMillis(),
@@ -354,8 +464,9 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 1000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true,
-      interBrokerProtocolVersion = ApiVersion.latestVersion)
+      origin = AppendOrigin.Client,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      brokerTopicStats = brokerTopicStats)
     val validatedRecords = validatedResults.validatedRecords
 
     for (batch <- validatedRecords.batches.asScala) {
@@ -377,15 +488,16 @@ class LogValidatorTest {
   }
 
   @Test
-  def testCreateTimeUpConversionV0ToV2() {
+  def testCreateTimeUpConversionV0ToV2(): Unit = {
     checkCreateTimeUpConversionFromV0(RecordBatch.MAGIC_VALUE_V2)
   }
 
   @Test
-  def testCreateTimeUpConversionV1ToV2() {
+  def testCreateTimeUpConversionV1ToV2(): Unit = {
     val timestamp = System.currentTimeMillis()
     val records = createRecords(magicValue = RecordBatch.MAGIC_VALUE_V1, codec = CompressionType.GZIP, timestamp = timestamp)
     val validatedResults = LogValidator.validateMessagesAndAssignOffsets(records,
+      topicPartition,
       offsetCounter = new LongRef(0),
       time = time,
       now = timestamp,
@@ -396,8 +508,9 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 1000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true,
-      interBrokerProtocolVersion = ApiVersion.latestVersion)
+      origin = AppendOrigin.Client,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      brokerTopicStats = brokerTopicStats)
     val validatedRecords = validatedResults.validatedRecords
 
     for (batch <- validatedRecords.batches.asScala) {
@@ -419,11 +532,11 @@ class LogValidatorTest {
   }
 
   @Test
-  def testCompressedV1() {
+  def testCompressedV1(): Unit = {
     checkCompressed(RecordBatch.MAGIC_VALUE_V1)
   }
 
-  private def checkCompressed(magic: Byte) {
+  private def checkCompressed(magic: Byte): Unit = {
     val now = System.currentTimeMillis()
     val timestampSeq = Seq(now - 1, now + 1, now)
 
@@ -441,6 +554,7 @@ class LogValidatorTest {
       new SimpleRecord(timestampSeq(2), "beautiful".getBytes))
 
     val validatedResults = LogValidator.validateMessagesAndAssignOffsets(records,
+      topicPartition,
       offsetCounter = new LongRef(0),
       time = time,
       now = System.currentTimeMillis(),
@@ -451,8 +565,9 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 1000L,
       partitionLeaderEpoch = partitionLeaderEpoch,
-      isFromClient = true,
-      interBrokerProtocolVersion = ApiVersion.latestVersion)
+      origin = AppendOrigin.Client,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      brokerTopicStats = brokerTopicStats)
     val validatedRecords = validatedResults.validatedRecords
 
     var i = 0
@@ -481,17 +596,18 @@ class LogValidatorTest {
   }
 
   @Test
-  def testCompressedV2() {
+  def testCompressedV2(): Unit = {
     checkCompressed(RecordBatch.MAGIC_VALUE_V2)
   }
 
-  @Test(expected = classOf[InvalidTimestampException])
-  def testInvalidCreateTimeNonCompressedV1() {
+  @Test(expected = classOf[RecordValidationException])
+  def testInvalidCreateTimeNonCompressedV1(): Unit = {
     val now = System.currentTimeMillis()
     val records = createRecords(magicValue = RecordBatch.MAGIC_VALUE_V1, timestamp = now - 1001L,
       codec = CompressionType.NONE)
     LogValidator.validateMessagesAndAssignOffsets(
       records,
+      topicPartition,
       offsetCounter = new LongRef(0),
       time = time,
       now = System.currentTimeMillis(),
@@ -502,17 +618,19 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 1000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true,
-      interBrokerProtocolVersion = ApiVersion.latestVersion)
+      origin = AppendOrigin.Client,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      brokerTopicStats = brokerTopicStats)
   }
 
-  @Test(expected = classOf[InvalidTimestampException])
-  def testInvalidCreateTimeNonCompressedV2() {
+  @Test(expected = classOf[RecordValidationException])
+  def testInvalidCreateTimeNonCompressedV2(): Unit = {
     val now = System.currentTimeMillis()
     val records = createRecords(magicValue = RecordBatch.MAGIC_VALUE_V2, timestamp = now - 1001L,
       codec = CompressionType.NONE)
     LogValidator.validateMessagesAndAssignOffsets(
       records,
+      topicPartition,
       offsetCounter = new LongRef(0),
       time = time,
       now = System.currentTimeMillis(),
@@ -523,17 +641,19 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 1000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true,
-      interBrokerProtocolVersion = ApiVersion.latestVersion)
+      origin = AppendOrigin.Client,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      brokerTopicStats = brokerTopicStats)
   }
 
-  @Test(expected = classOf[InvalidTimestampException])
-  def testInvalidCreateTimeCompressedV1() {
+  @Test(expected = classOf[RecordValidationException])
+  def testInvalidCreateTimeCompressedV1(): Unit = {
     val now = System.currentTimeMillis()
     val records = createRecords(magicValue = RecordBatch.MAGIC_VALUE_V1, timestamp = now - 1001L,
       codec = CompressionType.GZIP)
     LogValidator.validateMessagesAndAssignOffsets(
       records,
+      topicPartition,
       offsetCounter = new LongRef(0),
       time = time,
       now = System.currentTimeMillis(),
@@ -544,17 +664,19 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 1000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true,
-      interBrokerProtocolVersion = ApiVersion.latestVersion)
+      origin = AppendOrigin.Client,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      brokerTopicStats = brokerTopicStats)
   }
 
-  @Test(expected = classOf[InvalidTimestampException])
-  def testInvalidCreateTimeCompressedV2() {
+  @Test(expected = classOf[RecordValidationException])
+  def testInvalidCreateTimeCompressedV2(): Unit = {
     val now = System.currentTimeMillis()
     val records = createRecords(magicValue = RecordBatch.MAGIC_VALUE_V2, timestamp = now - 1001L,
       codec = CompressionType.GZIP)
     LogValidator.validateMessagesAndAssignOffsets(
       records,
+      topicPartition,
       offsetCounter = new LongRef(0),
       time = time,
       now = System.currentTimeMillis(),
@@ -565,16 +687,18 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 1000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true,
-      interBrokerProtocolVersion = ApiVersion.latestVersion)
+      origin = AppendOrigin.Client,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      brokerTopicStats = brokerTopicStats)
   }
 
   @Test
-  def testAbsoluteOffsetAssignmentNonCompressed() {
+  def testAbsoluteOffsetAssignmentNonCompressed(): Unit = {
     val records = createRecords(magicValue = RecordBatch.MAGIC_VALUE_V0, codec = CompressionType.NONE)
     val offset = 1234567
     checkOffsets(records, 0)
     checkOffsets(LogValidator.validateMessagesAndAssignOffsets(records,
+      topicPartition,
       offsetCounter = new LongRef(offset),
       time = time,
       now = System.currentTimeMillis(),
@@ -585,16 +709,18 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 1000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true,
-      interBrokerProtocolVersion = ApiVersion.latestVersion).validatedRecords, offset)
+      origin = AppendOrigin.Client,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      brokerTopicStats = brokerTopicStats).validatedRecords, offset)
   }
 
   @Test
-  def testAbsoluteOffsetAssignmentCompressed() {
+  def testAbsoluteOffsetAssignmentCompressed(): Unit = {
     val records = createRecords(magicValue = RecordBatch.MAGIC_VALUE_V0, codec = CompressionType.GZIP)
     val offset = 1234567
     checkOffsets(records, 0)
     checkOffsets(LogValidator.validateMessagesAndAssignOffsets(records,
+      topicPartition,
       offsetCounter = new LongRef(offset),
       time = time,
       now = System.currentTimeMillis(),
@@ -605,17 +731,19 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 1000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true,
-      interBrokerProtocolVersion = ApiVersion.latestVersion).validatedRecords, offset)
+      origin = AppendOrigin.Client,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      brokerTopicStats = brokerTopicStats).validatedRecords, offset)
   }
 
   @Test
-  def testRelativeOffsetAssignmentNonCompressedV1() {
+  def testRelativeOffsetAssignmentNonCompressedV1(): Unit = {
     val now = System.currentTimeMillis()
     val records = createRecords(magicValue = RecordBatch.MAGIC_VALUE_V1, timestamp = now, codec = CompressionType.NONE)
     val offset = 1234567
     checkOffsets(records, 0)
     val messageWithOffset = LogValidator.validateMessagesAndAssignOffsets(records,
+      topicPartition,
       offsetCounter = new LongRef(offset),
       time = time,
       now = System.currentTimeMillis(),
@@ -626,18 +754,20 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 5000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true,
-      interBrokerProtocolVersion = ApiVersion.latestVersion).validatedRecords
+      origin = AppendOrigin.Client,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      brokerTopicStats = brokerTopicStats).validatedRecords
     checkOffsets(messageWithOffset, offset)
   }
 
   @Test
-  def testRelativeOffsetAssignmentNonCompressedV2() {
+  def testRelativeOffsetAssignmentNonCompressedV2(): Unit = {
     val now = System.currentTimeMillis()
     val records = createRecords(magicValue = RecordBatch.MAGIC_VALUE_V2, timestamp = now, codec = CompressionType.NONE)
     val offset = 1234567
     checkOffsets(records, 0)
     val messageWithOffset = LogValidator.validateMessagesAndAssignOffsets(records,
+      topicPartition,
       offsetCounter = new LongRef(offset),
       time = time,
       now = System.currentTimeMillis(),
@@ -648,19 +778,21 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 5000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true,
-      interBrokerProtocolVersion = ApiVersion.latestVersion).validatedRecords
+      origin = AppendOrigin.Client,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      brokerTopicStats = brokerTopicStats).validatedRecords
     checkOffsets(messageWithOffset, offset)
   }
 
   @Test
-  def testRelativeOffsetAssignmentCompressedV1() {
+  def testRelativeOffsetAssignmentCompressedV1(): Unit = {
     val now = System.currentTimeMillis()
     val records = createRecords(magicValue = RecordBatch.MAGIC_VALUE_V1, timestamp = now, codec = CompressionType.GZIP)
     val offset = 1234567
     checkOffsets(records, 0)
     val compressedMessagesWithOffset = LogValidator.validateMessagesAndAssignOffsets(
       records,
+      topicPartition,
       offsetCounter = new LongRef(offset),
       time = time,
       now = System.currentTimeMillis(),
@@ -671,19 +803,21 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 5000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true,
-      interBrokerProtocolVersion = ApiVersion.latestVersion).validatedRecords
+      origin = AppendOrigin.Client,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      brokerTopicStats = brokerTopicStats).validatedRecords
     checkOffsets(compressedMessagesWithOffset, offset)
   }
 
   @Test
-  def testRelativeOffsetAssignmentCompressedV2() {
+  def testRelativeOffsetAssignmentCompressedV2(): Unit = {
     val now = System.currentTimeMillis()
     val records = createRecords(magicValue = RecordBatch.MAGIC_VALUE_V2, timestamp = now, codec = CompressionType.GZIP)
     val offset = 1234567
     checkOffsets(records, 0)
     val compressedMessagesWithOffset = LogValidator.validateMessagesAndAssignOffsets(
       records,
+      topicPartition,
       offsetCounter = new LongRef(offset),
       time = time,
       now = System.currentTimeMillis(),
@@ -694,17 +828,19 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 5000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true,
-      interBrokerProtocolVersion = ApiVersion.latestVersion).validatedRecords
+      origin = AppendOrigin.Client,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      brokerTopicStats = brokerTopicStats).validatedRecords
     checkOffsets(compressedMessagesWithOffset, offset)
   }
 
   @Test
-  def testOffsetAssignmentAfterUpConversionV0ToV1NonCompressed() {
+  def testOffsetAssignmentAfterUpConversionV0ToV1NonCompressed(): Unit = {
     val records = createRecords(magicValue = RecordBatch.MAGIC_VALUE_V0, codec = CompressionType.NONE)
     checkOffsets(records, 0)
     val offset = 1234567
     val validatedResults = LogValidator.validateMessagesAndAssignOffsets(records,
+      topicPartition,
       offsetCounter = new LongRef(offset),
       time = time,
       now = System.currentTimeMillis(),
@@ -715,19 +851,21 @@ class LogValidatorTest {
       timestampType = TimestampType.LOG_APPEND_TIME,
       timestampDiffMaxMs = 1000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true,
-      interBrokerProtocolVersion = ApiVersion.latestVersion)
+      origin = AppendOrigin.Client,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      brokerTopicStats = brokerTopicStats)
     checkOffsets(validatedResults.validatedRecords, offset)
     verifyRecordConversionStats(validatedResults.recordConversionStats, numConvertedRecords = 3, records,
       compressed = false)
   }
 
   @Test
-  def testOffsetAssignmentAfterUpConversionV0ToV2NonCompressed() {
+  def testOffsetAssignmentAfterUpConversionV0ToV2NonCompressed(): Unit = {
     val records = createRecords(magicValue = RecordBatch.MAGIC_VALUE_V0, codec = CompressionType.NONE)
     checkOffsets(records, 0)
     val offset = 1234567
     val validatedResults = LogValidator.validateMessagesAndAssignOffsets(records,
+      topicPartition,
       offsetCounter = new LongRef(offset),
       time = time,
       now = System.currentTimeMillis(),
@@ -738,19 +876,21 @@ class LogValidatorTest {
       timestampType = TimestampType.LOG_APPEND_TIME,
       timestampDiffMaxMs = 1000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true,
-      interBrokerProtocolVersion = ApiVersion.latestVersion)
+      origin = AppendOrigin.Client,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      brokerTopicStats = brokerTopicStats)
     checkOffsets(validatedResults.validatedRecords, offset)
     verifyRecordConversionStats(validatedResults.recordConversionStats, numConvertedRecords = 3, records,
       compressed = false)
   }
 
   @Test
-  def testOffsetAssignmentAfterUpConversionV0ToV1Compressed() {
+  def testOffsetAssignmentAfterUpConversionV0ToV1Compressed(): Unit = {
     val records = createRecords(magicValue = RecordBatch.MAGIC_VALUE_V0, codec = CompressionType.GZIP)
     val offset = 1234567
     checkOffsets(records, 0)
     val validatedResults = LogValidator.validateMessagesAndAssignOffsets(records,
+      topicPartition,
       offsetCounter = new LongRef(offset),
       time = time,
       now = System.currentTimeMillis(),
@@ -761,19 +901,21 @@ class LogValidatorTest {
       timestampType = TimestampType.LOG_APPEND_TIME,
       timestampDiffMaxMs = 1000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true,
-      interBrokerProtocolVersion = ApiVersion.latestVersion)
+      origin = AppendOrigin.Client,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      brokerTopicStats = brokerTopicStats)
     checkOffsets(validatedResults.validatedRecords, offset)
     verifyRecordConversionStats(validatedResults.recordConversionStats, numConvertedRecords = 3, records,
       compressed = true)
   }
 
   @Test
-  def testOffsetAssignmentAfterUpConversionV0ToV2Compressed() {
+  def testOffsetAssignmentAfterUpConversionV0ToV2Compressed(): Unit = {
     val records = createRecords(magicValue = RecordBatch.MAGIC_VALUE_V0, codec = CompressionType.GZIP)
     val offset = 1234567
     checkOffsets(records, 0)
     val validatedResults = LogValidator.validateMessagesAndAssignOffsets(records,
+      topicPartition,
       offsetCounter = new LongRef(offset),
       time = time,
       now = System.currentTimeMillis(),
@@ -784,19 +926,21 @@ class LogValidatorTest {
       timestampType = TimestampType.LOG_APPEND_TIME,
       timestampDiffMaxMs = 1000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true,
-      interBrokerProtocolVersion = ApiVersion.latestVersion)
+      origin = AppendOrigin.Client,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      brokerTopicStats = brokerTopicStats)
     checkOffsets(validatedResults.validatedRecords, offset)
     verifyRecordConversionStats(validatedResults.recordConversionStats, numConvertedRecords = 3, records,
       compressed = true)
   }
 
   @Test(expected = classOf[InvalidRecordException])
-  def testControlRecordsNotAllowedFromClients() {
+  def testControlRecordsNotAllowedFromClients(): Unit = {
     val offset = 1234567
     val endTxnMarker = new EndTransactionMarker(ControlRecordType.COMMIT, 0)
     val records = MemoryRecords.withEndTransactionMarker(23423L, 5, endTxnMarker)
     LogValidator.validateMessagesAndAssignOffsets(records,
+      topicPartition,
       offsetCounter = new LongRef(offset),
       time = time,
       now = System.currentTimeMillis(),
@@ -807,16 +951,18 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 5000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true,
-      interBrokerProtocolVersion = ApiVersion.latestVersion)
+      origin = AppendOrigin.Client,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      brokerTopicStats = brokerTopicStats)
   }
 
   @Test
-  def testControlRecordsNotCompressed() {
+  def testControlRecordsNotCompressed(): Unit = {
     val offset = 1234567
     val endTxnMarker = new EndTransactionMarker(ControlRecordType.COMMIT, 0)
     val records = MemoryRecords.withEndTransactionMarker(23423L, 5, endTxnMarker)
     val result = LogValidator.validateMessagesAndAssignOffsets(records,
+      topicPartition,
       offsetCounter = new LongRef(offset),
       time = time,
       now = System.currentTimeMillis(),
@@ -827,8 +973,9 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 5000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = false,
-      interBrokerProtocolVersion = ApiVersion.latestVersion)
+      origin = AppendOrigin.Coordinator,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      brokerTopicStats = brokerTopicStats)
     val batches = TestUtils.toList(result.validatedRecords.batches)
     assertEquals(1, batches.size)
     val batch = batches.get(0)
@@ -836,12 +983,13 @@ class LogValidatorTest {
   }
 
   @Test
-  def testOffsetAssignmentAfterDownConversionV1ToV0NonCompressed() {
+  def testOffsetAssignmentAfterDownConversionV1ToV0NonCompressed(): Unit = {
     val offset = 1234567
     val now = System.currentTimeMillis()
     val records = createRecords(RecordBatch.MAGIC_VALUE_V1, now, codec = CompressionType.NONE)
     checkOffsets(records, 0)
     checkOffsets(LogValidator.validateMessagesAndAssignOffsets(records,
+      topicPartition,
       offsetCounter = new LongRef(offset),
       time = time,
       now = System.currentTimeMillis(),
@@ -852,17 +1000,19 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 5000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true,
-      interBrokerProtocolVersion = ApiVersion.latestVersion).validatedRecords, offset)
+      origin = AppendOrigin.Client,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      brokerTopicStats = brokerTopicStats).validatedRecords, offset)
   }
 
   @Test
-  def testOffsetAssignmentAfterDownConversionV1ToV0Compressed() {
+  def testOffsetAssignmentAfterDownConversionV1ToV0Compressed(): Unit = {
     val offset = 1234567
     val now = System.currentTimeMillis()
     val records = createRecords(RecordBatch.MAGIC_VALUE_V1, now, CompressionType.GZIP)
     checkOffsets(records, 0)
     checkOffsets(LogValidator.validateMessagesAndAssignOffsets(records,
+      topicPartition,
       offsetCounter = new LongRef(offset),
       time = time,
       now = System.currentTimeMillis(),
@@ -873,16 +1023,18 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 5000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true,
-      interBrokerProtocolVersion = ApiVersion.latestVersion).validatedRecords, offset)
+      origin = AppendOrigin.Client,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      brokerTopicStats = brokerTopicStats).validatedRecords, offset)
   }
 
   @Test
-  def testOffsetAssignmentAfterUpConversionV1ToV2NonCompressed() {
+  def testOffsetAssignmentAfterUpConversionV1ToV2NonCompressed(): Unit = {
     val records = createRecords(magicValue = RecordBatch.MAGIC_VALUE_V1, codec = CompressionType.NONE)
     checkOffsets(records, 0)
     val offset = 1234567
     checkOffsets(LogValidator.validateMessagesAndAssignOffsets(records,
+      topicPartition,
       offsetCounter = new LongRef(offset),
       time = time,
       now = System.currentTimeMillis(),
@@ -893,16 +1045,18 @@ class LogValidatorTest {
       timestampType = TimestampType.LOG_APPEND_TIME,
       timestampDiffMaxMs = 1000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true,
-      interBrokerProtocolVersion = ApiVersion.latestVersion).validatedRecords, offset)
+      origin = AppendOrigin.Client,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      brokerTopicStats = brokerTopicStats).validatedRecords, offset)
   }
 
   @Test
-  def testOffsetAssignmentAfterUpConversionV1ToV2Compressed() {
+  def testOffsetAssignmentAfterUpConversionV1ToV2Compressed(): Unit = {
     val records = createRecords(magicValue = RecordBatch.MAGIC_VALUE_V1, codec = CompressionType.GZIP)
     val offset = 1234567
     checkOffsets(records, 0)
     checkOffsets(LogValidator.validateMessagesAndAssignOffsets(records,
+      topicPartition,
       offsetCounter = new LongRef(offset),
       time = time,
       now = System.currentTimeMillis(),
@@ -913,17 +1067,19 @@ class LogValidatorTest {
       timestampType = TimestampType.LOG_APPEND_TIME,
       timestampDiffMaxMs = 1000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true,
-      interBrokerProtocolVersion = ApiVersion.latestVersion).validatedRecords, offset)
+      origin = AppendOrigin.Client,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      brokerTopicStats = brokerTopicStats).validatedRecords, offset)
   }
 
   @Test
-  def testOffsetAssignmentAfterDownConversionV2ToV1NonCompressed() {
+  def testOffsetAssignmentAfterDownConversionV2ToV1NonCompressed(): Unit = {
     val offset = 1234567
     val now = System.currentTimeMillis()
     val records = createRecords(RecordBatch.MAGIC_VALUE_V2, now, codec = CompressionType.NONE)
     checkOffsets(records, 0)
     checkOffsets(LogValidator.validateMessagesAndAssignOffsets(records,
+      topicPartition,
       offsetCounter = new LongRef(offset),
       time = time,
       now = System.currentTimeMillis(),
@@ -934,17 +1090,19 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 5000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true,
-      interBrokerProtocolVersion = ApiVersion.latestVersion).validatedRecords, offset)
+      origin = AppendOrigin.Client,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      brokerTopicStats = brokerTopicStats).validatedRecords, offset)
   }
 
   @Test
-  def testOffsetAssignmentAfterDownConversionV2ToV1Compressed() {
+  def testOffsetAssignmentAfterDownConversionV2ToV1Compressed(): Unit = {
     val offset = 1234567
     val now = System.currentTimeMillis()
     val records = createRecords(RecordBatch.MAGIC_VALUE_V2, now, CompressionType.GZIP)
     checkOffsets(records, 0)
     checkOffsets(LogValidator.validateMessagesAndAssignOffsets(records,
+      topicPartition,
       offsetCounter = new LongRef(offset),
       time = time,
       now = System.currentTimeMillis(),
@@ -955,12 +1113,13 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 5000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true,
-      interBrokerProtocolVersion = ApiVersion.latestVersion).validatedRecords, offset)
+      origin = AppendOrigin.Client,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      brokerTopicStats = brokerTopicStats).validatedRecords, offset)
   }
 
   @Test(expected = classOf[UnsupportedForMessageFormatException])
-  def testDownConversionOfTransactionalRecordsNotPermitted() {
+  def testDownConversionOfTransactionalRecordsNotPermitted(): Unit = {
     val offset = 1234567
     val producerId = 1344L
     val producerEpoch = 16.toShort
@@ -968,6 +1127,7 @@ class LogValidatorTest {
     val records = MemoryRecords.withTransactionalRecords(CompressionType.NONE, producerId, producerEpoch, sequence,
       new SimpleRecord("hello".getBytes), new SimpleRecord("there".getBytes), new SimpleRecord("beautiful".getBytes))
     checkOffsets(LogValidator.validateMessagesAndAssignOffsets(records,
+      topicPartition,
       offsetCounter = new LongRef(offset),
       time = time,
       now = System.currentTimeMillis(),
@@ -978,12 +1138,13 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 5000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true,
-      interBrokerProtocolVersion = ApiVersion.latestVersion).validatedRecords, offset)
+      origin = AppendOrigin.Client,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      brokerTopicStats = brokerTopicStats).validatedRecords, offset)
   }
 
   @Test(expected = classOf[UnsupportedForMessageFormatException])
-  def testDownConversionOfIdempotentRecordsNotPermitted() {
+  def testDownConversionOfIdempotentRecordsNotPermitted(): Unit = {
     val offset = 1234567
     val producerId = 1344L
     val producerEpoch = 16.toShort
@@ -991,6 +1152,7 @@ class LogValidatorTest {
     val records = MemoryRecords.withIdempotentRecords(CompressionType.NONE, producerId, producerEpoch, sequence,
       new SimpleRecord("hello".getBytes), new SimpleRecord("there".getBytes), new SimpleRecord("beautiful".getBytes))
     checkOffsets(LogValidator.validateMessagesAndAssignOffsets(records,
+      topicPartition,
       offsetCounter = new LongRef(offset),
       time = time,
       now = System.currentTimeMillis(),
@@ -1001,17 +1163,19 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 5000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true,
-      interBrokerProtocolVersion = ApiVersion.latestVersion).validatedRecords, offset)
+      origin = AppendOrigin.Client,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      brokerTopicStats = brokerTopicStats).validatedRecords, offset)
   }
 
   @Test
-  def testOffsetAssignmentAfterDownConversionV2ToV0NonCompressed() {
+  def testOffsetAssignmentAfterDownConversionV2ToV0NonCompressed(): Unit = {
     val offset = 1234567
     val now = System.currentTimeMillis()
     val records = createRecords(RecordBatch.MAGIC_VALUE_V2, now, codec = CompressionType.NONE)
     checkOffsets(records, 0)
     checkOffsets(LogValidator.validateMessagesAndAssignOffsets(records,
+      topicPartition,
       offsetCounter = new LongRef(offset),
       time = time,
       now = System.currentTimeMillis(),
@@ -1022,17 +1186,19 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 5000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true,
-      interBrokerProtocolVersion = ApiVersion.latestVersion).validatedRecords, offset)
+      origin = AppendOrigin.Client,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      brokerTopicStats = brokerTopicStats).validatedRecords, offset)
   }
 
   @Test
-  def testOffsetAssignmentAfterDownConversionV2ToV0Compressed() {
+  def testOffsetAssignmentAfterDownConversionV2ToV0Compressed(): Unit = {
     val offset = 1234567
     val now = System.currentTimeMillis()
     val records = createRecords(RecordBatch.MAGIC_VALUE_V2, now, CompressionType.GZIP)
     checkOffsets(records, 0)
     checkOffsets(LogValidator.validateMessagesAndAssignOffsets(records,
+      topicPartition,
       offsetCounter = new LongRef(offset),
       time = time,
       now = System.currentTimeMillis(),
@@ -1043,27 +1209,34 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 5000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true,
-      interBrokerProtocolVersion = ApiVersion.latestVersion).validatedRecords, offset)
+      origin = AppendOrigin.Client,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      brokerTopicStats = brokerTopicStats).validatedRecords, offset)
   }
 
-  @Test(expected = classOf[InvalidRecordException])
-  def testInvalidInnerMagicVersion(): Unit = {
-    val offset = 1234567
-    val records = recordsWithInvalidInnerMagic(offset)
-    LogValidator.validateMessagesAndAssignOffsets(records,
-      offsetCounter = new LongRef(offset),
-      time = time,
-      now = System.currentTimeMillis(),
-      sourceCodec = SnappyCompressionCodec,
-      targetCodec = SnappyCompressionCodec,
-      compactedTopic = false,
-      magic = RecordBatch.MAGIC_VALUE_V1,
-      timestampType = TimestampType.CREATE_TIME,
-      timestampDiffMaxMs = 5000L,
-      partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true,
-      interBrokerProtocolVersion = ApiVersion.latestVersion)
+  @Test
+  def testNonIncreasingOffsetRecordBatchHasMetricsLogged(): Unit = {
+    val records = createNonIncreasingOffsetRecords(RecordBatch.MAGIC_VALUE_V2)
+    records.batches().asScala.head.setLastOffset(2)
+    assertThrows[InvalidRecordException] {
+      LogValidator.validateMessagesAndAssignOffsets(records,
+        topicPartition,
+        offsetCounter = new LongRef(0L),
+        time = time,
+        now = System.currentTimeMillis(),
+        sourceCodec = DefaultCompressionCodec,
+        targetCodec = DefaultCompressionCodec,
+        compactedTopic = false,
+        magic = RecordBatch.MAGIC_VALUE_V0,
+        timestampType = TimestampType.CREATE_TIME,
+        timestampDiffMaxMs = 5000L,
+        partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
+        origin = AppendOrigin.Client,
+        interBrokerProtocolVersion = ApiVersion.latestVersion,
+        brokerTopicStats = brokerTopicStats)
+    }
+    assertEquals(metricsKeySet.count(_.getMBeanName.endsWith(s"${BrokerTopicStats.InvalidOffsetOrSequenceRecordsPerSec}")), 1)
+    assertTrue(meterCount(s"${BrokerTopicStats.InvalidOffsetOrSequenceRecordsPerSec}") > 0)
   }
 
   @Test(expected = classOf[InvalidRecordException])
@@ -1077,6 +1250,7 @@ class LogValidatorTest {
     // The timestamps should be overwritten
     val records = createRecords(magicValue = RecordBatch.MAGIC_VALUE_V2, timestamp = 1234L, codec = CompressionType.NONE)
     LogValidator.validateMessagesAndAssignOffsets(records,
+      topicPartition,
       offsetCounter = new LongRef(0),
       time= time,
       now = now,
@@ -1087,8 +1261,9 @@ class LogValidatorTest {
       timestampType = TimestampType.LOG_APPEND_TIME,
       timestampDiffMaxMs = 1000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true,
-      interBrokerProtocolVersion = KAFKA_2_0_IV1)
+      origin = AppendOrigin.Client,
+      interBrokerProtocolVersion = KAFKA_2_0_IV1,
+      brokerTopicStats = brokerTopicStats)
   }
 
   @Test(expected = classOf[InvalidRecordException])
@@ -1101,6 +1276,52 @@ class LogValidatorTest {
     testBatchWithoutRecordsNotAllowed(NoCompressionCodec, DefaultCompressionCodec)
   }
 
+  @Test
+  def testInvalidTimestampExceptionHasBatchIndex(): Unit = {
+    val now = System.currentTimeMillis()
+    val records = createRecords(magicValue = RecordBatch.MAGIC_VALUE_V2, timestamp = now - 1001L,
+      codec = CompressionType.GZIP)
+    val e = intercept[RecordValidationException] {
+      LogValidator.validateMessagesAndAssignOffsets(
+        records,
+        topicPartition,
+        offsetCounter = new LongRef(0),
+        time = time,
+        now = System.currentTimeMillis(),
+        sourceCodec = DefaultCompressionCodec,
+        targetCodec = DefaultCompressionCodec,
+        magic = RecordBatch.MAGIC_VALUE_V1,
+        compactedTopic = false,
+        timestampType = TimestampType.CREATE_TIME,
+        timestampDiffMaxMs = 1000L,
+        partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
+        origin = AppendOrigin.Client,
+        interBrokerProtocolVersion = ApiVersion.latestVersion,
+        brokerTopicStats = brokerTopicStats)
+    }
+
+    assertTrue(e.invalidException.isInstanceOf[InvalidTimestampException])
+    assertTrue(e.recordErrors.nonEmpty)
+    assertEquals(e.recordErrors.size, 1)
+    assertEquals(e.recordErrors.head.batchIndex, 0)
+    assertNull(e.recordErrors.head.message)
+  }
+
+  @Test
+  def testInvalidRecordExceptionHasBatchIndex(): Unit = {
+    val e = intercept[RecordValidationException] {
+      validateMessages(recordsWithInvalidInnerMagic(
+        RecordBatch.MAGIC_VALUE_V0, RecordBatch.MAGIC_VALUE_V1, CompressionType.GZIP),
+        RecordBatch.MAGIC_VALUE_V0, CompressionType.GZIP, CompressionType.GZIP)
+    }
+
+    assertTrue(e.invalidException.isInstanceOf[InvalidRecordException])
+    assertTrue(e.recordErrors.nonEmpty)
+    assertEquals(e.recordErrors.size, 1)
+    assertEquals(e.recordErrors.head.batchIndex, 0)
+    assertNull(e.recordErrors.head.message)
+  }
+
   private def testBatchWithoutRecordsNotAllowed(sourceCodec: CompressionCodec, targetCodec: CompressionCodec): Unit = {
     val offset = 1234567
     val (producerId, producerEpoch, baseSequence, isTransactional, partitionLeaderEpoch) =
@@ -1112,6 +1333,7 @@ class LogValidatorTest {
     buffer.flip()
     val records = MemoryRecords.readableRecords(buffer)
     LogValidator.validateMessagesAndAssignOffsets(records,
+      topicPartition,
       offsetCounter = new LongRef(offset),
       time = time,
       now = System.currentTimeMillis(),
@@ -1122,8 +1344,9 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 5000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true,
-      interBrokerProtocolVersion = ApiVersion.latestVersion)
+      origin = AppendOrigin.Client,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      brokerTopicStats = brokerTopicStats)
   }
 
   private def createRecords(magicValue: Byte,
@@ -1137,8 +1360,35 @@ class LogValidatorTest {
     builder.build()
   }
 
+  private def createNonIncreasingOffsetRecords(magicValue: Byte,
+                                               timestamp: Long = RecordBatch.NO_TIMESTAMP,
+                                               codec: CompressionType = CompressionType.NONE): MemoryRecords = {
+    val buf = ByteBuffer.allocate(512)
+    val builder = MemoryRecords.builder(buf, magicValue, codec, TimestampType.CREATE_TIME, 0L)
+    builder.appendWithOffset(0, timestamp, null, "hello".getBytes)
+    builder.appendWithOffset(2, timestamp, null, "there".getBytes)
+    builder.appendWithOffset(3, timestamp, null, "beautiful".getBytes)
+    builder.build()
+  }
+
+  private def createTwoBatchedRecords(magicValue: Byte,
+                                      timestamp: Long,
+                                      codec: CompressionType): MemoryRecords = {
+    val buf = ByteBuffer.allocate(2048)
+    var builder = MemoryRecords.builder(buf, magicValue, codec, TimestampType.CREATE_TIME, 0L)
+    builder.append(10L, "1".getBytes(), "a".getBytes())
+    builder.close()
+    builder = MemoryRecords.builder(buf, magicValue, codec, TimestampType.CREATE_TIME, 1L)
+    builder.append(11L, "2".getBytes(), "b".getBytes())
+    builder.append(12L, "3".getBytes(), "c".getBytes())
+    builder.close()
+
+    buf.flip()
+    MemoryRecords.readableRecords(buf.slice())
+  }
+
   /* check that offsets are assigned consecutively from the given base offset */
-  def checkOffsets(records: MemoryRecords, baseOffset: Long) {
+  def checkOffsets(records: MemoryRecords, baseOffset: Long): Unit = {
     assertTrue("Message set should not be empty", records.records.asScala.nonEmpty)
     var offset = baseOffset
     for (entry <- records.records.asScala) {
@@ -1147,18 +1397,37 @@ class LogValidatorTest {
     }
   }
 
-  private def recordsWithInvalidInnerMagic(initialOffset: Long): MemoryRecords = {
+  private def recordsWithNonSequentialInnerOffsets(magicValue: Byte,
+                                                   codec: CompressionType,
+                                                   numRecords: Int): MemoryRecords = {
+    val records = (0 until numRecords).map { id =>
+      new SimpleRecord(id.toString.getBytes)
+    }
+
+    val buffer = ByteBuffer.allocate(1024)
+    val builder = MemoryRecords.builder(buffer, magicValue, codec, TimestampType.CREATE_TIME, 0L)
+
+    records.foreach { record =>
+      builder.appendUncheckedWithOffset(0, record)
+    }
+
+    builder.build()
+  }
+
+  private def recordsWithInvalidInnerMagic(batchMagicValue: Byte,
+                                           recordMagicValue: Byte,
+                                           codec: CompressionType): MemoryRecords = {
     val records = (0 until 20).map(id =>
-      LegacyRecord.create(RecordBatch.MAGIC_VALUE_V0,
+      LegacyRecord.create(recordMagicValue,
         RecordBatch.NO_TIMESTAMP,
         id.toString.getBytes,
         id.toString.getBytes))
 
     val buffer = ByteBuffer.allocate(math.min(math.max(records.map(_.sizeInBytes()).sum / 2, 1024), 1 << 16))
-    val builder = MemoryRecords.builder(buffer, RecordBatch.MAGIC_VALUE_V1, CompressionType.GZIP,
+    val builder = MemoryRecords.builder(buffer, batchMagicValue, codec,
       TimestampType.CREATE_TIME, 0L)
 
-    var offset = initialOffset
+    var offset = 1234567
     records.foreach { record =>
       builder.appendUncheckedWithOffset(offset, record)
       offset += 1
@@ -1178,7 +1447,7 @@ class LogValidatorTest {
   /**
     * expectedLogAppendTime is only checked if batch.magic is V2 or higher
     */
-  def validateLogAppendTime(expectedLogAppendTime: Long, expectedBaseTimestamp: Long, batch: RecordBatch) {
+  def validateLogAppendTime(expectedLogAppendTime: Long, expectedBaseTimestamp: Long, batch: RecordBatch): Unit = {
     assertTrue(batch.isValid)
     assertTrue(batch.timestampType == TimestampType.LOG_APPEND_TIME)
     assertEquals(s"Unexpected max timestamp of batch $batch", expectedLogAppendTime, batch.maxTimestamp)
diff --git a/core/src/test/scala/unit/kafka/log/OffsetIndexTest.scala b/core/src/test/scala/unit/kafka/log/OffsetIndexTest.scala
index 4e2ab2faf696e..b8bc00631be09 100644
--- a/core/src/test/scala/unit/kafka/log/OffsetIndexTest.scala
+++ b/core/src/test/scala/unit/kafka/log/OffsetIndexTest.scala
@@ -24,32 +24,32 @@ import org.junit.Assert._
 import java.util.{Arrays, Collections}
 
 import org.junit._
-import org.scalatest.junit.JUnitSuite
+import org.scalatest.Assertions.intercept
 
 import scala.collection._
 import scala.util.Random
 import kafka.utils.TestUtils
 import org.apache.kafka.common.errors.InvalidOffsetException
 
-class OffsetIndexTest extends JUnitSuite {
+class OffsetIndexTest {
   
   var idx: OffsetIndex = null
   val maxEntries = 30
   val baseOffset = 45L
   
   @Before
-  def setup() {
+  def setup(): Unit = {
     this.idx = new OffsetIndex(nonExistentTempFile(), baseOffset, maxIndexSize = 30 * 8)
   }
   
   @After
-  def teardown() {
+  def teardown(): Unit = {
     if(this.idx != null)
       this.idx.file.delete()
   }
   
   @Test
-  def randomLookupTest() {
+  def randomLookupTest(): Unit = {
     assertEquals("Not present value should return physical offset 0.", OffsetPosition(idx.baseOffset, 0), idx.lookup(92L))
     
     // append some random values
@@ -77,7 +77,7 @@ class OffsetIndexTest extends JUnitSuite {
   }
   
   @Test
-  def lookupExtremeCases() {
+  def lookupExtremeCases(): Unit = {
     assertEquals("Lookup on empty file", OffsetPosition(idx.baseOffset, 0), idx.lookup(idx.baseOffset))
     for(i <- 0 until idx.maxEntries)
       idx.append(idx.baseOffset + i + 1, i)
@@ -100,7 +100,7 @@ class OffsetIndexTest extends JUnitSuite {
   }
   
   @Test
-  def appendTooMany() {
+  def appendTooMany(): Unit = {
     for(i <- 0 until idx.maxEntries) {
       val offset = idx.baseOffset + i + 1
       idx.append(offset, i)
@@ -109,13 +109,13 @@ class OffsetIndexTest extends JUnitSuite {
   }
   
   @Test(expected = classOf[InvalidOffsetException])
-  def appendOutOfOrder() {
+  def appendOutOfOrder(): Unit = {
     idx.append(51, 0)
     idx.append(50, 1)
   }
 
   @Test
-  def testFetchUpperBoundOffset() {
+  def testFetchUpperBoundOffset(): Unit = {
     val first = OffsetPosition(baseOffset + 0, 0)
     val second = OffsetPosition(baseOffset + 1, 10)
     val third = OffsetPosition(baseOffset + 2, 23)
@@ -137,7 +137,7 @@ class OffsetIndexTest extends JUnitSuite {
   }
 
   @Test
-  def testReopen() {
+  def testReopen(): Unit = {
     val first = OffsetPosition(51, 0)
     val sec = OffsetPosition(52, 1)
     idx.append(first.offset, first.position)
@@ -152,7 +152,7 @@ class OffsetIndexTest extends JUnitSuite {
   }
   
   @Test
-  def truncate() {
+  def truncate(): Unit = {
 	val idx = new OffsetIndex(nonExistentTempFile(), baseOffset = 0L, maxIndexSize = 10 * 8)
 	idx.truncate()
     for(i <- 1 until 10)
@@ -201,7 +201,7 @@ class OffsetIndexTest extends JUnitSuite {
     idx.sanityCheck()
   }
   
-  def assertWriteFails[T](message: String, idx: OffsetIndex, offset: Int, klass: Class[T]) {
+  def assertWriteFails[T](message: String, idx: OffsetIndex, offset: Int, klass: Class[T]): Unit = {
     try {
       idx.append(offset, 1)
       fail(message)
diff --git a/core/src/test/scala/unit/kafka/log/OffsetMapTest.scala b/core/src/test/scala/unit/kafka/log/OffsetMapTest.scala
index 405756288c5d0..5fa9389442d6a 100644
--- a/core/src/test/scala/unit/kafka/log/OffsetMapTest.scala
+++ b/core/src/test/scala/unit/kafka/log/OffsetMapTest.scala
@@ -21,13 +21,12 @@ import java.nio._
 
 import kafka.utils.Exit
 import org.junit._
-import org.scalatest.junit.JUnitSuite
 import org.junit.Assert._
 
-class OffsetMapTest extends JUnitSuite {
+class OffsetMapTest {
   
   @Test
-  def testBasicValidation() {
+  def testBasicValidation(): Unit = {
     validateMap(10)
     validateMap(100)
     validateMap(1000)
@@ -35,7 +34,7 @@ class OffsetMapTest extends JUnitSuite {
   }
   
   @Test
-  def testClear() {
+  def testClear(): Unit = {
     val map = new SkimpyOffsetMap(4000)
     for(i <- 0 until 10)
       map.put(key(i), i)
@@ -47,7 +46,7 @@ class OffsetMapTest extends JUnitSuite {
   }
   
   @Test
-  def testGetWhenFull() {
+  def testGetWhenFull(): Unit = {
     val map = new SkimpyOffsetMap(4096)
     var i = 37L  //any value would do
     while (map.size < map.slots) {
@@ -72,7 +71,7 @@ class OffsetMapTest extends JUnitSuite {
 }
 
 object OffsetMapTest {
-  def main(args: Array[String]) {
+  def main(args: Array[String]): Unit = {
     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/ProducerStateManagerTest.scala b/core/src/test/scala/unit/kafka/log/ProducerStateManagerTest.scala
index a2abf7b3f5aea..94f34283d89fe 100644
--- a/core/src/test/scala/unit/kafka/log/ProducerStateManagerTest.scala
+++ b/core/src/test/scala/unit/kafka/log/ProducerStateManagerTest.scala
@@ -33,9 +33,9 @@ import org.apache.kafka.common.utils.{MockTime, Utils}
 import org.easymock.EasyMock
 import org.junit.Assert._
 import org.junit.{After, Before, Test}
-import org.scalatest.junit.JUnitSuite
+import org.scalatest.Assertions.{assertThrows, fail}
 
-class ProducerStateManagerTest extends JUnitSuite {
+class ProducerStateManagerTest {
   var logDir: File = null
   var stateManager: ProducerStateManager = null
   val partition = new TopicPartition("test", 0)
@@ -117,7 +117,7 @@ class ProducerStateManagerTest extends JUnitSuite {
     val epoch = 15.toShort
     val sequence = Int.MaxValue
     val offset = 735L
-    append(stateManager, producerId, epoch, sequence, offset, isFromClient = false)
+    append(stateManager, producerId, epoch, sequence, offset, origin = AppendOrigin.Replication)
 
     append(stateManager, producerId, epoch, 0, offset + 500)
 
@@ -135,9 +135,10 @@ class ProducerStateManagerTest extends JUnitSuite {
   def testProducerSequenceWithWrapAroundBatchRecord(): Unit = {
     val epoch = 15.toShort
 
-    val appendInfo = stateManager.prepareUpdate(producerId, isFromClient = false)
+    val appendInfo = stateManager.prepareUpdate(producerId, origin = AppendOrigin.Replication)
     // Sequence number wrap around
-    appendInfo.append(epoch, Int.MaxValue-10, 9, time.milliseconds(), 2000L, 2020L, isTransactional = false)
+    appendInfo.appendDataBatch(epoch, Int.MaxValue - 10, 9, time.milliseconds(),
+      LogOffsetMetadata(2000L), 2020L, isTransactional = false)
     assertEquals(None, stateManager.lastEntry(producerId))
     stateManager.update(appendInfo)
     assertTrue(stateManager.lastEntry(producerId).isDefined)
@@ -145,7 +146,7 @@ class ProducerStateManagerTest extends JUnitSuite {
     val lastEntry = stateManager.lastEntry(producerId).get
     assertEquals(Int.MaxValue-10, lastEntry.firstSeq)
     assertEquals(9, lastEntry.lastSeq)
-    assertEquals(2000L, lastEntry.firstOffset)
+    assertEquals(2000L, lastEntry.firstDataOffset)
     assertEquals(2020L, lastEntry.lastDataOffset)
   }
 
@@ -154,7 +155,7 @@ class ProducerStateManagerTest extends JUnitSuite {
     val epoch = 15.toShort
     val sequence = Int.MaxValue
     val offset = 735L
-    append(stateManager, producerId, epoch, sequence, offset, isFromClient = false)
+    append(stateManager, producerId, epoch, sequence, offset, origin = AppendOrigin.Replication)
     append(stateManager, producerId, epoch, 1, offset + 500)
   }
 
@@ -163,7 +164,7 @@ class ProducerStateManagerTest extends JUnitSuite {
     val epoch = 5.toShort
     val sequence = 16
     val offset = 735L
-    append(stateManager, producerId, epoch, sequence, offset, isFromClient = false)
+    append(stateManager, producerId, epoch, sequence, offset, origin = AppendOrigin.Replication)
 
     val maybeLastEntry = stateManager.lastEntry(producerId)
     assertTrue(maybeLastEntry.isDefined)
@@ -173,7 +174,7 @@ class ProducerStateManagerTest extends JUnitSuite {
     assertEquals(sequence, lastEntry.firstSeq)
     assertEquals(sequence, lastEntry.lastSeq)
     assertEquals(offset, lastEntry.lastDataOffset)
-    assertEquals(offset, lastEntry.firstOffset)
+    assertEquals(offset, lastEntry.firstDataOffset)
   }
 
   @Test
@@ -208,47 +209,85 @@ class ProducerStateManagerTest extends JUnitSuite {
     val producerEpoch = 0.toShort
     val offset = 992342L
     val seq = 0
-    val producerAppendInfo = new ProducerAppendInfo(producerId, ProducerStateEntry.empty(producerId), ValidationType.Full)
-    producerAppendInfo.append(producerEpoch, seq, seq, time.milliseconds(), offset, offset, isTransactional = true)
+    val producerAppendInfo = new ProducerAppendInfo(partition, producerId, ProducerStateEntry.empty(producerId), AppendOrigin.Client)
 
-    val logOffsetMetadata = new LogOffsetMetadata(messageOffset = offset, segmentBaseOffset = 990000L,
+    val firstOffsetMetadata = LogOffsetMetadata(messageOffset = offset, segmentBaseOffset = 990000L,
       relativePositionInSegment = 234224)
-    producerAppendInfo.maybeCacheTxnFirstOffsetMetadata(logOffsetMetadata)
+    producerAppendInfo.appendDataBatch(producerEpoch, seq, seq, time.milliseconds(),
+      firstOffsetMetadata, offset, isTransactional = true)
     stateManager.update(producerAppendInfo)
 
-    assertEquals(Some(logOffsetMetadata), stateManager.firstUnstableOffset)
+    assertEquals(Some(firstOffsetMetadata), stateManager.firstUnstableOffset)
   }
 
   @Test
-  def testNonMatchingTxnFirstOffsetMetadataNotCached(): Unit = {
+  def testLastStableOffsetCompletedTxn(): Unit = {
     val producerEpoch = 0.toShort
-    val offset = 992342L
-    val seq = 0
-    val producerAppendInfo = new ProducerAppendInfo(producerId, ProducerStateEntry.empty(producerId), ValidationType.Full)
-    producerAppendInfo.append(producerEpoch, seq, seq, time.milliseconds(), offset, offset, isTransactional = true)
+    val segmentBaseOffset = 990000L
+
+    def beginTxn(producerId: Long, startOffset: Long): Unit = {
+      val relativeOffset = (startOffset - segmentBaseOffset).toInt
+      val producerAppendInfo = new ProducerAppendInfo(
+        partition,
+        producerId,
+        ProducerStateEntry.empty(producerId),
+        AppendOrigin.Client
+      )
+      val firstOffsetMetadata = LogOffsetMetadata(messageOffset = startOffset, segmentBaseOffset = segmentBaseOffset,
+        relativePositionInSegment = 50 * relativeOffset)
+      producerAppendInfo.appendDataBatch(producerEpoch, 0, 0, time.milliseconds(),
+        firstOffsetMetadata, startOffset, isTransactional = true)
+      stateManager.update(producerAppendInfo)
+    }
 
-    // use some other offset to simulate a follower append where the log offset metadata won't typically
-    // match any of the transaction first offsets
-    val logOffsetMetadata = new LogOffsetMetadata(messageOffset = offset - 23429, segmentBaseOffset = 990000L,
-      relativePositionInSegment = 234224)
-    producerAppendInfo.maybeCacheTxnFirstOffsetMetadata(logOffsetMetadata)
-    stateManager.update(producerAppendInfo)
+    val producerId1 = producerId
+    val startOffset1 = 992342L
+    beginTxn(producerId1, startOffset1)
+
+    val producerId2 = producerId + 1
+    val startOffset2 = startOffset1 + 25
+    beginTxn(producerId2, startOffset2)
+
+    val producerId3 = producerId + 2
+    val startOffset3 = startOffset1 + 57
+    beginTxn(producerId3, startOffset3)
+
+    val lastOffset1 = startOffset3 + 15
+    val completedTxn1 = CompletedTxn(producerId1, startOffset1, lastOffset1, isAborted = false)
+    assertEquals(startOffset2, stateManager.lastStableOffset(completedTxn1))
+    stateManager.completeTxn(completedTxn1)
+    stateManager.onHighWatermarkUpdated(lastOffset1 + 1)
+    assertEquals(Some(startOffset2), stateManager.firstUnstableOffset.map(_.messageOffset))
+
+    val lastOffset3 = lastOffset1 + 20
+    val completedTxn3 = CompletedTxn(producerId3, startOffset3, lastOffset3, isAborted = false)
+    assertEquals(startOffset2, stateManager.lastStableOffset(completedTxn3))
+    stateManager.completeTxn(completedTxn3)
+    stateManager.onHighWatermarkUpdated(lastOffset3 + 1)
+    assertEquals(Some(startOffset2), stateManager.firstUnstableOffset.map(_.messageOffset))
 
-    assertEquals(Some(LogOffsetMetadata(offset)), stateManager.firstUnstableOffset)
+    val lastOffset2 = lastOffset3 + 78
+    val completedTxn2 = CompletedTxn(producerId2, startOffset2, lastOffset2, isAborted = false)
+    assertEquals(lastOffset2 + 1, stateManager.lastStableOffset(completedTxn2))
+    stateManager.completeTxn(completedTxn2)
+    stateManager.onHighWatermarkUpdated(lastOffset2 + 1)
+    assertEquals(None, stateManager.firstUnstableOffset)
   }
 
   @Test
   def testPrepareUpdateDoesNotMutate(): Unit = {
     val producerEpoch = 0.toShort
 
-    val appendInfo = stateManager.prepareUpdate(producerId, isFromClient = true)
-    appendInfo.append(producerEpoch, 0, 5, time.milliseconds(), 15L, 20L, isTransactional = false)
+    val appendInfo = stateManager.prepareUpdate(producerId, origin = AppendOrigin.Client)
+    appendInfo.appendDataBatch(producerEpoch, 0, 5, time.milliseconds(),
+      LogOffsetMetadata(15L), 20L, isTransactional = false)
     assertEquals(None, stateManager.lastEntry(producerId))
     stateManager.update(appendInfo)
     assertTrue(stateManager.lastEntry(producerId).isDefined)
 
-    val nextAppendInfo = stateManager.prepareUpdate(producerId, isFromClient = true)
-    nextAppendInfo.append(producerEpoch, 6, 10, time.milliseconds(), 26L, 30L, isTransactional = false)
+    val nextAppendInfo = stateManager.prepareUpdate(producerId, origin = AppendOrigin.Client)
+    nextAppendInfo.appendDataBatch(producerEpoch, 6, 10, time.milliseconds(),
+      LogOffsetMetadata(26L), 30L, isTransactional = false)
     assertTrue(stateManager.lastEntry(producerId).isDefined)
 
     var lastEntry = stateManager.lastEntry(producerId).get
@@ -270,23 +309,25 @@ class ProducerStateManagerTest extends JUnitSuite {
     val offset = 9L
     append(stateManager, producerId, producerEpoch, 0, offset)
 
-    val appendInfo = stateManager.prepareUpdate(producerId, isFromClient = true)
-    appendInfo.append(producerEpoch, 1, 5, time.milliseconds(), 16L, 20L, isTransactional = true)
+    val appendInfo = stateManager.prepareUpdate(producerId, origin = AppendOrigin.Client)
+    appendInfo.appendDataBatch(producerEpoch, 1, 5, time.milliseconds(),
+      LogOffsetMetadata(16L), 20L, isTransactional = true)
     var lastEntry = appendInfo.toEntry
     assertEquals(producerEpoch, lastEntry.producerEpoch)
     assertEquals(1, lastEntry.firstSeq)
     assertEquals(5, lastEntry.lastSeq)
-    assertEquals(16L, lastEntry.firstOffset)
+    assertEquals(16L, lastEntry.firstDataOffset)
     assertEquals(20L, lastEntry.lastDataOffset)
     assertEquals(Some(16L), lastEntry.currentTxnFirstOffset)
     assertEquals(List(new TxnMetadata(producerId, 16L)), appendInfo.startedTransactions)
 
-    appendInfo.append(producerEpoch, 6, 10, time.milliseconds(), 26L, 30L, isTransactional = true)
+    appendInfo.appendDataBatch(producerEpoch, 6, 10, time.milliseconds(),
+      LogOffsetMetadata(26L), 30L, isTransactional = true)
     lastEntry = appendInfo.toEntry
     assertEquals(producerEpoch, lastEntry.producerEpoch)
     assertEquals(1, lastEntry.firstSeq)
     assertEquals(10, lastEntry.lastSeq)
-    assertEquals(16L, lastEntry.firstOffset)
+    assertEquals(16L, lastEntry.firstDataOffset)
     assertEquals(30L, lastEntry.lastDataOffset)
     assertEquals(Some(16L), lastEntry.currentTxnFirstOffset)
     assertEquals(List(new TxnMetadata(producerId, 16L)), appendInfo.startedTransactions)
@@ -303,7 +344,7 @@ class ProducerStateManagerTest extends JUnitSuite {
     // verify that appending the transaction marker doesn't affect the metadata of the cached record batches.
     assertEquals(1, lastEntry.firstSeq)
     assertEquals(10, lastEntry.lastSeq)
-    assertEquals(16L, lastEntry.firstOffset)
+    assertEquals(16L, lastEntry.firstDataOffset)
     assertEquals(30L, lastEntry.lastDataOffset)
     assertEquals(coordinatorEpoch, lastEntry.coordinatorEpoch)
     assertEquals(None, lastEntry.currentTxnFirstOffset)
@@ -376,17 +417,78 @@ class ProducerStateManagerTest extends JUnitSuite {
   }
 
   @Test
-  def testRecoverFromSnapshot(): Unit = {
+  def testRecoverFromSnapshotUnfinishedTransaction(): Unit = {
     val epoch = 0.toShort
-    append(stateManager, producerId, epoch, 0, 0L)
-    append(stateManager, producerId, epoch, 1, 1L)
+    append(stateManager, producerId, epoch, 0, 0L, isTransactional = true)
+    append(stateManager, producerId, epoch, 1, 1L, isTransactional = true)
 
     stateManager.takeSnapshot()
     val recoveredMapping = new ProducerStateManager(partition, logDir, maxPidExpirationMs)
     recoveredMapping.truncateAndReload(0L, 3L, time.milliseconds)
 
+    // The snapshot only persists the last appended batch metadata
+    val loadedEntry = recoveredMapping.lastEntry(producerId)
+    assertEquals(1, loadedEntry.get.firstDataOffset)
+    assertEquals(1, loadedEntry.get.firstSeq)
+    assertEquals(1, loadedEntry.get.lastDataOffset)
+    assertEquals(1, loadedEntry.get.lastSeq)
+    assertEquals(Some(0), loadedEntry.get.currentTxnFirstOffset)
+
     // entry added after recovery
-    append(recoveredMapping, producerId, epoch, 2, 2L)
+    append(recoveredMapping, producerId, epoch, 2, 2L, isTransactional = true)
+  }
+
+  @Test
+  def testRecoverFromSnapshotFinishedTransaction(): Unit = {
+    val epoch = 0.toShort
+    append(stateManager, producerId, epoch, 0, 0L, isTransactional = true)
+    append(stateManager, producerId, epoch, 1, 1L, isTransactional = true)
+    appendEndTxnMarker(stateManager, producerId, epoch, ControlRecordType.ABORT, offset = 2L)
+
+    stateManager.takeSnapshot()
+    val recoveredMapping = new ProducerStateManager(partition, logDir, maxPidExpirationMs)
+    recoveredMapping.truncateAndReload(0L, 3L, time.milliseconds)
+
+    // The snapshot only persists the last appended batch metadata
+    val loadedEntry = recoveredMapping.lastEntry(producerId)
+    assertEquals(1, loadedEntry.get.firstDataOffset)
+    assertEquals(1, loadedEntry.get.firstSeq)
+    assertEquals(1, loadedEntry.get.lastDataOffset)
+    assertEquals(1, loadedEntry.get.lastSeq)
+    assertEquals(None, loadedEntry.get.currentTxnFirstOffset)
+  }
+
+  @Test
+  def testRecoverFromSnapshotEmptyTransaction(): Unit = {
+    val epoch = 0.toShort
+    val appendTimestamp = time.milliseconds()
+    appendEndTxnMarker(stateManager, producerId, epoch, ControlRecordType.ABORT,
+      offset = 0L, timestamp = appendTimestamp)
+    stateManager.takeSnapshot()
+
+    val recoveredMapping = new ProducerStateManager(partition, logDir, maxPidExpirationMs)
+    recoveredMapping.truncateAndReload(logStartOffset = 0L, logEndOffset = 1L, time.milliseconds)
+
+    val lastEntry = recoveredMapping.lastEntry(producerId)
+    assertTrue(lastEntry.isDefined)
+    assertEquals(appendTimestamp, lastEntry.get.lastTimestamp)
+    assertEquals(None, lastEntry.get.currentTxnFirstOffset)
+  }
+
+  @Test
+  def testProducerStateAfterFencingAbortMarker(): Unit = {
+    val epoch = 0.toShort
+    append(stateManager, producerId, epoch, 0, 0L, isTransactional = true)
+    appendEndTxnMarker(stateManager, producerId, (epoch + 1).toShort, ControlRecordType.ABORT, offset = 1L)
+
+    val lastEntry = stateManager.lastEntry(producerId).get
+    assertEquals(None, lastEntry.currentTxnFirstOffset)
+    assertEquals(-1, lastEntry.lastDataOffset)
+    assertEquals(-1, lastEntry.firstDataOffset)
+
+    // The producer should not be expired because we want to preserve fencing epochs
+    stateManager.removeExpiredProducers(time.milliseconds())
+    assertTrue(stateManager.lastEntry(producerId).isDefined)
   }
 
   @Test(expected = classOf[UnknownProducerIdException])
@@ -418,7 +520,7 @@ class ProducerStateManagerTest extends JUnitSuite {
     // entry added after recovery. The pid should be expired now, and would not exist in the pid mapping. Nonetheless
     // the append on a replica should be accepted with the local producer state updated to the appended value.
     assertFalse(recoveredMapping.activeProducers.contains(producerId))
-    append(recoveredMapping, producerId, epoch, sequence, 2L, 70001, isFromClient = false)
+    append(recoveredMapping, producerId, epoch, sequence, 2L, 70001, origin = AppendOrigin.Replication)
     assertTrue(recoveredMapping.activeProducers.contains(producerId))
     val producerStateEntry = recoveredMapping.activeProducers.get(producerId).head
     assertEquals(epoch, producerStateEntry.producerEpoch)
@@ -434,7 +536,7 @@ class ProducerStateManagerTest extends JUnitSuite {
 
     // First we ensure that we raise an OutOfOrderSequenceException is raised when the append comes from a client.
     try {
-      append(stateManager, producerId, epoch, outOfOrderSequence, 1L, 1, isFromClient = true)
+      append(stateManager, producerId, epoch, outOfOrderSequence, 1L, 1, origin = AppendOrigin.Client)
       fail("Expected an OutOfOrderSequenceException to be raised.")
     } catch {
       case _ : OutOfOrderSequenceException =>
@@ -444,7 +546,7 @@ class ProducerStateManagerTest extends JUnitSuite {
     }
 
     assertEquals(0L, stateManager.activeProducers(producerId).lastSeq)
-    append(stateManager, producerId, epoch, outOfOrderSequence, 1L, 1, isFromClient = false)
+    append(stateManager, producerId, epoch, outOfOrderSequence, 1L, 1, origin = AppendOrigin.Replication)
     assertEquals(outOfOrderSequence, stateManager.activeProducers(producerId).lastSeq)
   }
 
@@ -520,52 +622,7 @@ class ProducerStateManagerTest extends JUnitSuite {
   }
 
   @Test
-  def testFirstUnstableOffsetAfterEviction(): Unit = {
-    val epoch = 0.toShort
-    val sequence = 0
-    append(stateManager, producerId, epoch, sequence, offset = 99, isTransactional = true)
-    assertEquals(Some(99), stateManager.firstUnstableOffset.map(_.messageOffset))
-    append(stateManager, 2L, epoch, 0, offset = 106, isTransactional = true)
-    stateManager.truncateHead(100)
-    assertEquals(Some(106), stateManager.firstUnstableOffset.map(_.messageOffset))
-  }
-
-  @Test
-  def testTruncateHead(): Unit = {
-    val epoch = 0.toShort
-
-    append(stateManager, producerId, epoch, 0, 0L)
-    append(stateManager, producerId, epoch, 1, 1L)
-    stateManager.takeSnapshot()
-
-    val anotherPid = 2L
-    append(stateManager, anotherPid, epoch, 0, 2L)
-    append(stateManager, anotherPid, epoch, 1, 3L)
-    stateManager.takeSnapshot()
-    assertEquals(Set(2, 4), currentSnapshotOffsets)
-
-    stateManager.truncateHead(2)
-    assertEquals(Set(2, 4), currentSnapshotOffsets)
-    assertEquals(Set(anotherPid), stateManager.activeProducers.keySet)
-    assertEquals(None, stateManager.lastEntry(producerId))
-
-    val maybeEntry = stateManager.lastEntry(anotherPid)
-    assertTrue(maybeEntry.isDefined)
-    assertEquals(3L, maybeEntry.get.lastDataOffset)
-
-    stateManager.truncateHead(3)
-    assertEquals(Set(anotherPid), stateManager.activeProducers.keySet)
-    assertEquals(Set(4), currentSnapshotOffsets)
-    assertEquals(4, stateManager.mapEndOffset)
-
-    stateManager.truncateHead(5)
-    assertEquals(Set(), stateManager.activeProducers.keySet)
-    assertEquals(Set(), currentSnapshotOffsets)
-    assertEquals(5, stateManager.mapEndOffset)
-  }
-
-  @Test
-  def testLoadFromSnapshotRemovesNonRetainedProducers(): Unit = {
+  def testLoadFromSnapshotRetainsNonExpiredProducers(): Unit = {
     val epoch = 0.toShort
     val pid1 = 1L
     val pid2 = 2L
@@ -576,13 +633,17 @@ class ProducerStateManagerTest extends JUnitSuite {
     assertEquals(2, stateManager.activeProducers.size)
 
     stateManager.truncateAndReload(1L, 2L, time.milliseconds())
-    assertEquals(1, stateManager.activeProducers.size)
-    assertEquals(None, stateManager.lastEntry(pid1))
+    assertEquals(2, stateManager.activeProducers.size)
 
-    val entry = stateManager.lastEntry(pid2)
-    assertTrue(entry.isDefined)
-    assertEquals(0, entry.get.lastSeq)
-    assertEquals(1L, entry.get.lastDataOffset)
+    val entry1 = stateManager.lastEntry(pid1)
+    assertTrue(entry1.isDefined)
+    assertEquals(0, entry1.get.lastSeq)
+    assertEquals(0L, entry1.get.lastDataOffset)
+
+    val entry2 = stateManager.lastEntry(pid2)
+    assertTrue(entry2.isDefined)
+    assertEquals(0, entry2.get.lastSeq)
+    assertEquals(1L, entry2.get.lastDataOffset)
   }
 
   @Test
@@ -618,7 +679,7 @@ class ProducerStateManagerTest extends JUnitSuite {
   }
 
   @Test(expected = classOf[UnknownProducerIdException])
-  def testPidExpirationTimeout() {
+  def testPidExpirationTimeout(): Unit = {
     val epoch = 5.toShort
     val sequence = 37
     append(stateManager, producerId, epoch, sequence, 1L)
@@ -628,7 +689,7 @@ class ProducerStateManagerTest extends JUnitSuite {
   }
 
   @Test
-  def testFirstUnstableOffset() {
+  def testFirstUnstableOffset(): Unit = {
     val epoch = 5.toShort
     val sequence = 0
 
@@ -662,7 +723,7 @@ class ProducerStateManagerTest extends JUnitSuite {
   }
 
   @Test
-  def testProducersWithOngoingTransactionsDontExpire() {
+  def testProducersWithOngoingTransactionsDontExpire(): Unit = {
     val epoch = 5.toShort
     val sequence = 0
 
@@ -685,9 +746,10 @@ class ProducerStateManagerTest extends JUnitSuite {
     val stateManager = new ProducerStateManager(partition, logDir, maxPidExpirationMs)
 
     val epoch = 0.toShort
-    append(stateManager, producerId, epoch, RecordBatch.NO_SEQUENCE, offset = 99, isTransactional = true)
-    append(stateManager, producerId, epoch, RecordBatch.NO_SEQUENCE, offset = 100, isTransactional = true)
-
+    append(stateManager, producerId, epoch, RecordBatch.NO_SEQUENCE, offset = 99,
+      isTransactional = true, origin = AppendOrigin.Coordinator)
+    append(stateManager, producerId, epoch, RecordBatch.NO_SEQUENCE, offset = 100,
+      isTransactional = true, origin = AppendOrigin.Coordinator)
   }
 
   @Test(expected = classOf[ProducerFencedException])
@@ -778,7 +840,7 @@ class ProducerStateManagerTest extends JUnitSuite {
     EasyMock.replay(batch)
 
     // Appending the empty control batch should not throw and a new transaction shouldn't be started
-    append(stateManager, producerId, producerEpoch, baseOffset, batch, isFromClient = true)
+    append(stateManager, producerId, producerEpoch, baseOffset, batch, origin = AppendOrigin.Client)
     assertEquals(None, stateManager.lastEntry(producerId).get.currentTxnFirstOffset)
   }
 
@@ -819,11 +881,12 @@ class ProducerStateManagerTest extends JUnitSuite {
                                  offset: Long,
                                  coordinatorEpoch: Int = 0,
                                  timestamp: Long = time.milliseconds()): (CompletedTxn, Long) = {
-    val producerAppendInfo = stateManager.prepareUpdate(producerId, isFromClient = true)
+    val producerAppendInfo = stateManager.prepareUpdate(producerId, origin = AppendOrigin.Coordinator)
     val endTxnMarker = new EndTransactionMarker(controlType, coordinatorEpoch)
     val completedTxn = producerAppendInfo.appendEndTxnMarker(endTxnMarker, producerEpoch, offset, timestamp)
     mapping.update(producerAppendInfo)
-    val lastStableOffset = mapping.completeTxn(completedTxn)
+    val lastStableOffset = mapping.lastStableOffset(completedTxn)
+    mapping.completeTxn(completedTxn)
     mapping.updateMapEndOffset(offset + 1)
     (completedTxn, lastStableOffset)
   }
@@ -835,9 +898,10 @@ class ProducerStateManagerTest extends JUnitSuite {
                      offset: Long,
                      timestamp: Long = time.milliseconds(),
                      isTransactional: Boolean = false,
-                     isFromClient : Boolean = true): Unit = {
-    val producerAppendInfo = stateManager.prepareUpdate(producerId, isFromClient)
-    producerAppendInfo.append(producerEpoch, seq, seq, timestamp, offset, offset, isTransactional)
+                     origin : AppendOrigin = AppendOrigin.Client): Unit = {
+    val producerAppendInfo = stateManager.prepareUpdate(producerId, origin)
+    producerAppendInfo.appendDataBatch(producerEpoch, seq, seq, timestamp,
+      LogOffsetMetadata(offset), offset, isTransactional)
     stateManager.update(producerAppendInfo)
     stateManager.updateMapEndOffset(offset + 1)
   }
@@ -847,14 +911,14 @@ class ProducerStateManagerTest extends JUnitSuite {
                      producerEpoch: Short,
                      offset: Long,
                      batch: RecordBatch,
-                     isFromClient : Boolean): Unit = {
-    val producerAppendInfo = stateManager.prepareUpdate(producerId, isFromClient)
-    producerAppendInfo.append(batch)
+                     origin: AppendOrigin): Unit = {
+    val producerAppendInfo = stateManager.prepareUpdate(producerId, origin)
+    producerAppendInfo.append(batch, firstOffsetMetadataOpt = None)
     stateManager.update(producerAppendInfo)
     stateManager.updateMapEndOffset(offset + 1)
   }
 
-  private def currentSnapshotOffsets =
+  private def currentSnapshotOffsets: Set[Long] =
     logDir.listFiles.map(Log.offsetFromFile).toSet
 
 }
diff --git a/core/src/test/scala/unit/kafka/log/TimeIndexTest.scala b/core/src/test/scala/unit/kafka/log/TimeIndexTest.scala
index 3653d3eceb14d..bfe8f56882f15 100644
--- a/core/src/test/scala/unit/kafka/log/TimeIndexTest.scala
+++ b/core/src/test/scala/unit/kafka/log/TimeIndexTest.scala
@@ -23,29 +23,29 @@ import kafka.utils.TestUtils
 import org.apache.kafka.common.errors.InvalidOffsetException
 import org.junit.{After, Before, Test}
 import org.junit.Assert.assertEquals
-import org.scalatest.junit.JUnitSuite
+import org.scalatest.Assertions.intercept
 
 /**
  * Unit test for time index.
  */
-class TimeIndexTest extends JUnitSuite {
+class TimeIndexTest {
   var idx: TimeIndex = null
   val maxEntries = 30
   val baseOffset = 45L
 
   @Before
-  def setup() {
+  def setup(): Unit = {
     this.idx = new TimeIndex(nonExistantTempFile(), baseOffset = baseOffset, maxIndexSize = maxEntries * 12)
   }
 
   @After
-  def teardown() {
+  def teardown(): Unit = {
     if(this.idx != null)
       this.idx.file.delete()
   }
 
   @Test
-  def testLookUp() {
+  def testLookUp(): Unit = {
     // Empty time index
     assertEquals(TimestampOffset(-1L, baseOffset), idx.lookup(100L))
 
@@ -75,7 +75,7 @@ class TimeIndexTest extends JUnitSuite {
   }
 
   @Test
-  def testTruncate() {
+  def testTruncate(): Unit = {
     appendEntries(maxEntries - 1)
     idx.truncate()
     assertEquals(0, idx.entries)
@@ -86,7 +86,7 @@ class TimeIndexTest extends JUnitSuite {
   }
 
   @Test
-  def testAppend() {
+  def testAppend(): Unit = {
     appendEntries(maxEntries - 1)
     intercept[IllegalArgumentException] {
       idx.maybeAppend(10000L, 1000L)
@@ -97,7 +97,7 @@ class TimeIndexTest extends JUnitSuite {
     idx.maybeAppend(10000L, 1000L, true)
   }
 
-  private def appendEntries(numEntries: Int) {
+  private def appendEntries(numEntries: Int): Unit = {
     for (i <- 1 to numEntries)
       idx.maybeAppend(i * 10, i * 10 + baseOffset)
   }
diff --git a/core/src/test/scala/unit/kafka/log/TransactionIndexTest.scala b/core/src/test/scala/unit/kafka/log/TransactionIndexTest.scala
index 574a8f586235a..0eb93e37ed804 100644
--- a/core/src/test/scala/unit/kafka/log/TransactionIndexTest.scala
+++ b/core/src/test/scala/unit/kafka/log/TransactionIndexTest.scala
@@ -22,7 +22,7 @@ import kafka.utils.TestUtils
 import org.apache.kafka.common.requests.FetchResponse.AbortedTransaction
 import org.junit.Assert._
 import org.junit.{After, Before, Test}
-import org.scalatest.junit.JUnitSuite
+import org.scalatestplus.junit.JUnitSuite
 
 class TransactionIndexTest extends JUnitSuite {
   var file: File = _
diff --git a/core/src/test/scala/unit/kafka/metrics/KafkaTimerTest.scala b/core/src/test/scala/unit/kafka/metrics/KafkaTimerTest.scala
index 3b3e4c39e491f..4ffa68aa07b1a 100644
--- a/core/src/test/scala/unit/kafka/metrics/KafkaTimerTest.scala
+++ b/core/src/test/scala/unit/kafka/metrics/KafkaTimerTest.scala
@@ -25,7 +25,7 @@ import com.yammer.metrics.core.{MetricsRegistry, Clock}
 class KafkaTimerTest {
 
   @Test
-  def testKafkaTimer() {
+  def testKafkaTimer(): Unit = {
     val clock = new ManualClock
     val testRegistry = new MetricsRegistry(clock)
     val metric = testRegistry.newTimer(this.getClass, "TestTimer")
@@ -52,7 +52,7 @@ class KafkaTimerTest {
       TimeUnit.NANOSECONDS.toMillis(ticksInNanos)
     }
 
-    def addMillis(millis: Long) {
+    def addMillis(millis: Long): Unit = {
       ticksInNanos += TimeUnit.MILLISECONDS.toNanos(millis)
     }
   }
diff --git a/core/src/test/scala/unit/kafka/metrics/MetricsTest.scala b/core/src/test/scala/unit/kafka/metrics/MetricsTest.scala
index 76cce0f1f6c42..3c2d30abb53e7 100644
--- a/core/src/test/scala/unit/kafka/metrics/MetricsTest.scala
+++ b/core/src/test/scala/unit/kafka/metrics/MetricsTest.scala
@@ -18,11 +18,11 @@
 package kafka.metrics
 
 import java.util.Properties
-import javax.management.ObjectName
 
+import javax.management.ObjectName
 import com.yammer.metrics.Metrics
-import com.yammer.metrics.core.{Meter, MetricPredicate}
-import org.junit.Test
+import com.yammer.metrics.core.MetricPredicate
+import org.junit.{Ignore, Test}
 import org.junit.Assert._
 import kafka.integration.KafkaServerTestHarness
 import kafka.server._
@@ -30,7 +30,6 @@ import kafka.utils._
 
 import scala.collection._
 import scala.collection.JavaConverters._
-import scala.util.matching.Regex
 import kafka.log.LogConfig
 import org.apache.kafka.common.TopicPartition
 
@@ -47,7 +46,7 @@ class MetricsTest extends KafkaServerTestHarness with Logging {
   val nMessages = 2
 
   @Test
-  def testMetricsReporterAfterDeletingTopic() {
+  def testMetricsReporterAfterDeletingTopic(): Unit = {
     val topic = "test-topic-metric"
     createTopic(topic, 1, 1)
     adminZkClient.deleteTopic(topic)
@@ -56,7 +55,7 @@ class MetricsTest extends KafkaServerTestHarness with Logging {
   }
 
   @Test
-  def testBrokerTopicMetricsUnregisteredAfterDeletingTopic() {
+  def testBrokerTopicMetricsUnregisteredAfterDeletingTopic(): Unit = {
     val topic = "test-broker-topic-metric"
     createTopic(topic, 2, 1)
     // Produce a few messages to create the metrics
@@ -76,6 +75,22 @@ class MetricsTest extends KafkaServerTestHarness with Logging {
     assertEquals(metrics.keySet.asScala.count(_.getMBeanName == "kafka.server:type=KafkaServer,name=ClusterId"), 1)
   }
 
+  @Test
+  @Ignore // reenable once it's not prone to the failure identified in KAFKA-9786
+  def testGeneralBrokerTopicMetricsAreGreedilyRegistered(): Unit = {
+    val topic = "test-broker-topic-metric"
+    createTopic(topic, 2, 1)
+
+    // The broker metrics for all topics should be greedily registered
+    assertTrue("General topic metrics don't exist", topicMetrics(None).nonEmpty)
+    assertEquals(servers.head.brokerTopicStats.allTopicsStats.metricMap.size +
+      servers.head.brokerTopicStats.allTopicsStats.counterMetricMap.size, topicMetrics(None).size)
+    // topic metrics should be lazily registered
+    assertTrue("Topic metrics aren't lazily registered", topicMetricGroups(topic).isEmpty)
+    TestUtils.generateAndProduceMessages(servers, topic, nMessages)
+    assertTrue("Topic metrics aren't registered", topicMetricGroups(topic).nonEmpty)
+  }
+
   @Test
   def testWindowsStyleTagNames(): Unit = {
     val path = "C:\\windows-path\\kafka-logs"
@@ -110,24 +125,27 @@ class MetricsTest extends KafkaServerTestHarness with Logging {
         logSize.map(_ > 0).getOrElse(false))
     }
 
-    val initialReplicationBytesIn = meterCount(replicationBytesIn)
-    val initialReplicationBytesOut = meterCount(replicationBytesOut)
-    val initialBytesIn = meterCount(bytesIn)
-    val initialBytesOut = meterCount(bytesOut)
+    // Consume messages to make bytesOut tick
+    TestUtils.consumeTopicRecords(servers, topic, nMessages)
+    val initialReplicationBytesIn = TestUtils.meterCount(replicationBytesIn)
+    val initialReplicationBytesOut = TestUtils.meterCount(replicationBytesOut)
+    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(servers, topic, nMessages)
 
-    assertTrue(meterCount(replicationBytesIn) > initialReplicationBytesIn)
-    assertTrue(meterCount(replicationBytesOut) > initialReplicationBytesOut)
-    assertTrue(meterCount(bytesIn) > initialBytesIn)
-    // BytesOut doesn't include replication, so it shouldn't have changed
-    assertEquals(initialBytesOut, meterCount(bytesOut))
+    assertTrue(TestUtils.meterCount(replicationBytesIn) > initialReplicationBytesIn)
+    assertTrue(TestUtils.meterCount(replicationBytesOut) > initialReplicationBytesOut)
+    assertTrue(TestUtils.meterCount(bytesIn) > initialBytesIn)
 
     // Consume messages to make bytesOut tick
-    TestUtils.consumeTopicRecords(servers, topic, nMessages * 2)
+    TestUtils.consumeTopicRecords(servers, topic, nMessages)
 
-    assertTrue(meterCount(bytesOut) > initialBytesOut)
+    assertTrue(TestUtils.meterCount(bytesOut) > initialBytesOut)
   }
 
   @Test
@@ -136,9 +154,12 @@ class MetricsTest extends KafkaServerTestHarness with Logging {
 
     assertEquals(metrics.keySet.asScala.count(_.getMBeanName == "kafka.controller:type=KafkaController,name=ActiveControllerCount"), 1)
     assertEquals(metrics.keySet.asScala.count(_.getMBeanName == "kafka.controller:type=KafkaController,name=OfflinePartitionsCount"), 1)
-    assertEquals(metrics.keySet.asScala.count(_.getMBeanName == "kafka.controller:type=KafkaController,name=PreferredReplicaImbalanceCount"), 1)
     assertEquals(metrics.keySet.asScala.count(_.getMBeanName == "kafka.controller:type=KafkaController,name=GlobalTopicCount"), 1)
     assertEquals(metrics.keySet.asScala.count(_.getMBeanName == "kafka.controller:type=KafkaController,name=GlobalPartitionCount"), 1)
+    assertEquals(metrics.keySet.asScala.count(_.getMBeanName == "kafka.controller:type=KafkaController,name=TopicsToDeleteCount"), 1)
+    assertEquals(metrics.keySet.asScala.count(_.getMBeanName == "kafka.controller:type=KafkaController,name=ReplicasToDeleteCount"), 1)
+    assertEquals(metrics.keySet.asScala.count(_.getMBeanName == "kafka.controller:type=KafkaController,name=TopicsIneligibleToDeleteCount"), 1)
+    assertEquals(metrics.keySet.asScala.count(_.getMBeanName == "kafka.controller:type=KafkaController,name=ReplicasIneligibleToDeleteCount"), 1)
   }
 
   /**
@@ -154,19 +175,18 @@ class MetricsTest extends KafkaServerTestHarness with Logging {
     assertEquals(metrics.keySet.asScala.count(_.getMBeanName == "kafka.server:type=SessionExpireListener,name=ZooKeeperDisconnectsPerSec"), 1)
   }
 
-  private def meterCount(metricName: String): Long = {
-    Metrics.defaultRegistry.allMetrics.asScala
-      .filterKeys(_.getMBeanName.endsWith(metricName))
-      .values
-      .headOption
-      .getOrElse(fail(s"Unable to find metric $metricName"))
-      .asInstanceOf[Meter]
-      .count
+  private def topicMetrics(topic: Option[String]): Set[String] = {
+    val metricNames = Metrics.defaultRegistry.allMetrics().keySet.asScala.map(_.getMBeanName)
+    filterByTopicMetricRegex(metricNames, topic)
   }
 
   private def topicMetricGroups(topic: String): Set[String] = {
-    val topicMetricRegex = new Regex(".*BrokerTopicMetrics.*("+topic+")$")
     val metricGroups = Metrics.defaultRegistry.groupedMetrics(MetricPredicate.ALL).keySet.asScala
-    metricGroups.filter(topicMetricRegex.pattern.matcher(_).matches)
+    filterByTopicMetricRegex(metricGroups, Some(topic))
+  }
+
+  private def filterByTopicMetricRegex(metrics: Set[String], topic: Option[String]): Set[String] = {
+    val pattern = (".*BrokerTopicMetrics.*" + topic.map(t => s"($t)$$").getOrElse("")).r.pattern
+    metrics.filter(pattern.matcher(_).matches())
   }
 }
diff --git a/core/src/test/scala/unit/kafka/network/SocketServerTest.scala b/core/src/test/scala/unit/kafka/network/SocketServerTest.scala
index 81a96a3801717..85e2fb7b6890e 100644
--- a/core/src/test/scala/unit/kafka/network/SocketServerTest.scala
+++ b/core/src/test/scala/unit/kafka/network/SocketServerTest.scala
@@ -21,38 +21,40 @@ import java.io._
 import java.net._
 import java.nio.ByteBuffer
 import java.nio.channels.SocketChannel
+import java.util.concurrent.{CompletableFuture, Executors}
 import java.util.{HashMap, Properties, Random}
 
 import com.yammer.metrics.core.{Gauge, Meter}
 import com.yammer.metrics.{Metrics => YammerMetrics}
 import javax.net.ssl._
-
 import kafka.security.CredentialProvider
 import kafka.server.{KafkaConfig, ThrottledChannel}
 import kafka.utils.Implicits._
-import kafka.utils.TestUtils
+import kafka.utils.{CoreUtils, TestUtils}
 import org.apache.kafka.common.TopicPartition
 import org.apache.kafka.common.memory.MemoryPool
+import org.apache.kafka.common.message.SaslAuthenticateRequestData
+import org.apache.kafka.common.message.SaslHandshakeRequestData
 import org.apache.kafka.common.metrics.Metrics
 import org.apache.kafka.common.network.KafkaChannel.ChannelMuteState
 import org.apache.kafka.common.network.{ChannelBuilder, ChannelState, KafkaChannel, ListenerName, NetworkReceive, NetworkSend, Selector, Send}
 import org.apache.kafka.common.protocol.{ApiKeys, Errors}
 import org.apache.kafka.common.record.MemoryRecords
-import org.apache.kafka.common.requests.{AbstractRequest, ProduceRequest, RequestHeader}
+import org.apache.kafka.common.requests.{AbstractRequest, ProduceRequest, RequestHeader, SaslAuthenticateRequest, SaslHandshakeRequest}
 import org.apache.kafka.common.security.auth.{KafkaPrincipal, SecurityProtocol}
 import org.apache.kafka.common.security.scram.internals.ScramMechanism
 import org.apache.kafka.common.utils.{LogContext, MockTime, Time}
 import org.apache.log4j.Level
 import org.junit.Assert._
 import org.junit._
-import org.scalatest.junit.JUnitSuite
+import org.scalatest.Assertions.fail
 
 import scala.collection.JavaConverters._
 import scala.collection.mutable
 import scala.collection.mutable.ArrayBuffer
 import scala.util.control.ControlThrowable
 
-class SocketServerTest extends JUnitSuite {
+class SocketServerTest {
   val props = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 0)
   props.put("listeners", "PLAINTEXT://localhost:0")
   props.put("num.network.threads", "1")
@@ -68,8 +70,7 @@ class SocketServerTest extends JUnitSuite {
   val localAddress = InetAddress.getLoopbackAddress
 
   // Clean-up any metrics left around by previous tests
-  for (metricName <- YammerMetrics.defaultRegistry.allMetrics.keySet.asScala)
-    YammerMetrics.defaultRegistry.removeMetric(metricName)
+  TestUtils.clearYammerMetrics()
 
   val server = new SocketServer(config, metrics, Time.SYSTEM, credentialProvider)
   server.startup()
@@ -86,14 +87,14 @@ class SocketServerTest extends JUnitSuite {
   }
 
   @After
-  def tearDown() {
+  def tearDown(): Unit = {
     shutdownServerAndMetrics(server)
     sockets.foreach(_.close())
     sockets.clear()
     kafkaLogger.setLevel(logLevelToRestore)
   }
 
-  def sendRequest(socket: Socket, request: Array[Byte], id: Option[Short] = None, flush: Boolean = true) {
+  def sendRequest(socket: Socket, request: Array[Byte], id: Option[Short] = None, flush: Boolean = true): Unit = {
     val outgoing = new DataOutputStream(socket.getOutputStream)
     id match {
       case Some(id) =>
@@ -107,6 +108,14 @@ class SocketServerTest extends JUnitSuite {
       outgoing.flush()
   }
 
+  def sendApiRequest(socket: Socket, request: AbstractRequest, header: RequestHeader) = {
+    val byteBuffer = request.serialize(header)
+    byteBuffer.rewind()
+    val serializedBytes = new Array[Byte](byteBuffer.remaining)
+    byteBuffer.get(serializedBytes)
+    sendRequest(socket, serializedBytes)
+  }
+
   def receiveResponse(socket: Socket): Array[Byte] = {
     val incoming = new DataInputStream(socket.getInputStream)
     val len = incoming.readInt()
@@ -124,11 +133,11 @@ class SocketServerTest extends JUnitSuite {
   }
 
   /* A simple request handler that just echos back the response */
-  def processRequest(channel: RequestChannel) {
+  def processRequest(channel: RequestChannel): Unit = {
     processRequest(channel, receiveRequest(channel))
   }
 
-  def processRequest(channel: RequestChannel, request: RequestChannel.Request) {
+  def processRequest(channel: RequestChannel, request: RequestChannel.Request): Unit = {
     val byteBuffer = request.body[AbstractRequest].serialize(request.header)
     byteBuffer.rewind()
 
@@ -136,7 +145,10 @@ class SocketServerTest extends JUnitSuite {
     channel.sendResponse(new RequestChannel.SendResponse(request, send, Some(request.header.toString), None))
   }
 
-  def connect(s: SocketServer = server, listenerName: ListenerName = ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT), localAddr: InetAddress = null, port: Int = 0) = {
+  def connect(s: SocketServer = server,
+              listenerName: ListenerName = ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT),
+              localAddr: InetAddress = null,
+              port: Int = 0) = {
     val socket = new Socket("localhost", s.boundPort(listenerName), localAddr, port)
     sockets += socket
     socket
@@ -177,7 +189,7 @@ class SocketServerTest extends JUnitSuite {
   }
 
   @Test
-  def simpleRequest() {
+  def simpleRequest(): Unit = {
     val plainSocket = connect()
     val serializedBytes = producerRequestBytes()
 
@@ -188,22 +200,73 @@ class SocketServerTest extends JUnitSuite {
     verifyAcceptorBlockedPercent("PLAINTEXT", expectBlocked = false)
   }
 
+  @Test
+  def testStagedListenerStartup(): Unit = {
+    val testProps = new Properties
+    testProps ++= props
+    testProps.put("listeners", "EXTERNAL://localhost:0,INTERNAL://localhost:0,CONTROLLER://localhost:0")
+    testProps.put("listener.security.protocol.map", "EXTERNAL:PLAINTEXT,INTERNAL:PLAINTEXT,CONTROLLER:PLAINTEXT")
+    testProps.put("control.plane.listener.name", "CONTROLLER")
+    testProps.put("inter.broker.listener.name", "INTERNAL")
+    val config = KafkaConfig.fromProps(testProps)
+    val testableServer = new TestableSocketServer(config)
+    testableServer.startup(startupProcessors = false)
+    val updatedEndPoints = config.advertisedListeners.map { endpoint =>
+      endpoint.copy(port = testableServer.boundPort(endpoint.listenerName))
+    }.map(_.toJava)
+
+    val externalReadyFuture = new CompletableFuture[Void]()
+    val executor = Executors.newSingleThreadExecutor()
+
+    def listenerStarted(listenerName: ListenerName) = {
+      try {
+        val socket = connect(testableServer, listenerName, localAddr = InetAddress.getLocalHost)
+        sendAndReceiveRequest(socket, testableServer)
+        true
+      } catch {
+        case _: Throwable => false
+      }
+    }
+
+    try {
+      testableServer.startControlPlaneProcessor()
+      val socket1 = connect(testableServer, config.controlPlaneListenerName.get, localAddr = InetAddress.getLocalHost)
+      sendAndReceiveControllerRequest(socket1, testableServer)
+
+      val externalListener = new ListenerName("EXTERNAL")
+      val externalEndpoint = updatedEndPoints.find(e => e.listenerName.get == externalListener.value).get
+      val futures =  Map(externalEndpoint -> externalReadyFuture)
+      val startFuture = executor.submit(CoreUtils.runnable(testableServer.startDataPlaneProcessors(futures)))
+      TestUtils.waitUntilTrue(() => listenerStarted(config.interBrokerListenerName), "Inter-broker listener not started")
+      assertFalse("Socket server startup did not wait for future to complete", startFuture.isDone)
+
+      assertFalse(listenerStarted(externalListener))
+
+      externalReadyFuture.complete(null)
+      TestUtils.waitUntilTrue(() => listenerStarted(externalListener), "External listener not started")
+    } finally {
+      executor.shutdownNow()
+      shutdownServerAndMetrics(testableServer)
+    }
+  }
+
   @Test
   def testControlPlaneRequest(): Unit = {
     val testProps = new Properties
     testProps ++= props
-    testProps.put("listeners", "PLAINTEXT://localhost:0,CONTROLLER://localhost:5000")
+    testProps.put("listeners", "PLAINTEXT://localhost:0,CONTROLLER://localhost:0")
     testProps.put("listener.security.protocol.map", "PLAINTEXT:PLAINTEXT,CONTROLLER:PLAINTEXT")
     testProps.put("control.plane.listener.name", "CONTROLLER")
     val config = KafkaConfig.fromProps(testProps)
     withTestableServer(config, { testableServer =>
-      val socket = connect(testableServer, config.controlPlaneListenerName.get, localAddr = InetAddress.getLocalHost, port = 5000)
+      val socket = connect(testableServer, config.controlPlaneListenerName.get,
+        localAddr = InetAddress.getLocalHost)
       sendAndReceiveControllerRequest(socket, testableServer)
     })
   }
 
   @Test
-  def tooBigRequestIsRejected() {
+  def tooBigRequestIsRejected(): Unit = {
     val tooManyBytes = new Array[Byte](server.config.socketRequestMaxBytes + 1)
     new Random().nextBytes(tooManyBytes)
     val socket = connect()
@@ -221,7 +284,7 @@ class SocketServerTest extends JUnitSuite {
   }
 
   @Test
-  def testGracefulClose() {
+  def testGracefulClose(): Unit = {
     val plainSocket = connect()
     val serializedBytes = producerRequestBytes()
 
@@ -250,7 +313,7 @@ class SocketServerTest extends JUnitSuite {
   }
 
   @Test
-  def testConnectionId() {
+  def testConnectionId(): Unit = {
     val sockets = (1 to 5).map(_ => connect())
     val serializedBytes = producerRequestBytes()
 
@@ -267,7 +330,7 @@ class SocketServerTest extends JUnitSuite {
   }
 
   @Test
-  def testIdleConnection() {
+  def testIdleConnection(): Unit = {
     val idleTimeMs = 60000
     val time = new MockTime()
     props.put(KafkaConfig.ConnectionsMaxIdleMsProp, idleTimeMs.toString)
@@ -312,7 +375,7 @@ class SocketServerTest extends JUnitSuite {
   }
 
   @Test
-  def testConnectionIdReuse() {
+  def testConnectionIdReuse(): Unit = {
     val idleTimeMs = 60000
     val time = new MockTime()
     props.put(KafkaConfig.ConnectionsMaxIdleMsProp, idleTimeMs.toString)
@@ -446,7 +509,7 @@ class SocketServerTest extends JUnitSuite {
     server.dataPlaneProcessor(0).openOrClosingChannel(request.context.connectionId)
 
   @Test
-  def testSendActionResponseWithThrottledChannelWhereThrottlingInProgress() {
+  def testSendActionResponseWithThrottledChannelWhereThrottlingInProgress(): Unit = {
     val socket = connect()
     val serializedBytes = producerRequestBytes()
     // SendAction with throttling in progress
@@ -460,7 +523,7 @@ class SocketServerTest extends JUnitSuite {
   }
 
   @Test
-  def testSendActionResponseWithThrottledChannelWhereThrottlingAlreadyDone() {
+  def testSendActionResponseWithThrottledChannelWhereThrottlingAlreadyDone(): Unit = {
     val socket = connect()
     val serializedBytes = producerRequestBytes()
     // SendAction with throttling in progress
@@ -475,7 +538,7 @@ class SocketServerTest extends JUnitSuite {
   }
 
   @Test
-  def testNoOpActionResponseWithThrottledChannelWhereThrottlingInProgress() {
+  def testNoOpActionResponseWithThrottledChannelWhereThrottlingInProgress(): Unit = {
     val socket = connect()
     val serializedBytes = producerRequestBytes()
     // SendAction with throttling in progress
@@ -487,7 +550,7 @@ class SocketServerTest extends JUnitSuite {
   }
 
   @Test
-  def testNoOpActionResponseWithThrottledChannelWhereThrottlingAlreadyDone() {
+  def testNoOpActionResponseWithThrottledChannelWhereThrottlingAlreadyDone(): Unit = {
     val socket = connect()
     val serializedBytes = producerRequestBytes()
     // SendAction with throttling in progress
@@ -500,7 +563,7 @@ class SocketServerTest extends JUnitSuite {
   }
 
   @Test
-  def testSocketsCloseOnShutdown() {
+  def testSocketsCloseOnShutdown(): Unit = {
     // open a connection
     val plainSocket = connect()
     plainSocket.setTcpNoDelay(true)
@@ -527,7 +590,7 @@ class SocketServerTest extends JUnitSuite {
   }
 
   @Test
-  def testMaxConnectionsPerIp() {
+  def testMaxConnectionsPerIp(): Unit = {
     // make the maximum allowable number of connections
     val conns = (0 until server.config.maxConnectionsPerIp).map(_ => connect())
     // now try one more (should fail)
@@ -549,7 +612,7 @@ class SocketServerTest extends JUnitSuite {
   }
 
   @Test
-  def testZeroMaxConnectionsPerIp() {
+  def testZeroMaxConnectionsPerIp(): Unit = {
     val newProps = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 0)
     newProps.setProperty(KafkaConfig.MaxConnectionsPerIpProp, "0")
     newProps.setProperty(KafkaConfig.MaxConnectionsPerIpOverridesProp, "%s:%s".format("127.0.0.1", "5"))
@@ -586,7 +649,7 @@ class SocketServerTest extends JUnitSuite {
   }
 
   @Test
-  def testMaxConnectionsPerIpOverrides() {
+  def testMaxConnectionsPerIpOverrides(): Unit = {
     val overrideNum = server.config.maxConnectionsPerIp + 1
     val overrideProps = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 0)
     overrideProps.put(KafkaConfig.MaxConnectionsPerIpOverridesProp, s"localhost:$overrideNum")
@@ -613,7 +676,7 @@ class SocketServerTest extends JUnitSuite {
   }
 
   @Test
-  def testSslSocketServer() {
+  def testSslSocketServer(): Unit = {
     val trustStoreFile = File.createTempFile("truststore", ".jks")
     val overrideProps = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, interBrokerSecurityProtocol = Some(SecurityProtocol.SSL),
       trustStoreFile = Some(trustStoreFile))
@@ -653,7 +716,72 @@ class SocketServerTest extends JUnitSuite {
   }
 
   @Test
-  def testSessionPrincipal() {
+  def testSaslReauthenticationFailure(): Unit = {
+    shutdownServerAndMetrics(server) // we will use our own instance because we require custom configs
+    val username = "admin"
+    val password = "admin-secret"
+    val reauthMs = 1500
+    val brokerProps = new Properties
+    brokerProps.setProperty("listeners", "SASL_PLAINTEXT://localhost:0")
+    brokerProps.setProperty("security.inter.broker.protocol", "SASL_PLAINTEXT")
+    brokerProps.setProperty("listener.name.sasl_plaintext.plain.sasl.jaas.config",
+      "org.apache.kafka.common.security.plain.PlainLoginModule required " +
+        "username=\"%s\" password=\"%s\" user_%s=\"%s\";".format(username, password, username, password))
+    brokerProps.setProperty("sasl.mechanism.inter.broker.protocol", "PLAIN")
+    brokerProps.setProperty("listener.name.sasl_plaintext.sasl.enabled.mechanisms", "PLAIN")
+    brokerProps.setProperty("num.network.threads", "1")
+    brokerProps.setProperty("connections.max.reauth.ms", reauthMs.toString)
+    val overrideProps = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect,
+      saslProperties = Some(brokerProps), enableSaslPlaintext = true)
+    val time = new MockTime()
+    val overrideServer = new TestableSocketServer(KafkaConfig.fromProps(overrideProps), time = time)
+    try {
+      overrideServer.startup()
+      val socket = connect(overrideServer, ListenerName.forSecurityProtocol(SecurityProtocol.SASL_PLAINTEXT))
+
+      val correlationId = -1
+      val clientId = ""
+      // send a SASL handshake request
+      val saslHandshakeRequest = new SaslHandshakeRequest.Builder(new SaslHandshakeRequestData().setMechanism("PLAIN"))
+        .build()
+      val saslHandshakeHeader = new RequestHeader(ApiKeys.SASL_HANDSHAKE, saslHandshakeRequest.version, clientId,
+        correlationId)
+      sendApiRequest(socket, saslHandshakeRequest, saslHandshakeHeader)
+      receiveResponse(socket)
+
+      // now send credentials within a SaslAuthenticateRequest
+      val authBytes = "admin\u0000admin\u0000admin-secret".getBytes("UTF-8")
+      val saslAuthenticateRequest = new SaslAuthenticateRequest.Builder(new SaslAuthenticateRequestData()
+        .setAuthBytes(authBytes)).build()
+      val saslAuthenticateHeader = new RequestHeader(ApiKeys.SASL_AUTHENTICATE, saslAuthenticateRequest.version,
+        clientId, correlationId)
+      sendApiRequest(socket, saslAuthenticateRequest, saslAuthenticateHeader)
+      receiveResponse(socket)
+      assertEquals(1, overrideServer.testableSelector.channels.size)
+
+      // advance the clock long enough to cause server-side disconnection upon next send...
+      time.sleep(reauthMs * 2)
+      // ...and now send something to trigger the disconnection
+      val ackTimeoutMs = 10000
+      val ack = 0: Short
+      val emptyRequest = ProduceRequest.Builder.forCurrentMagic(ack, ackTimeoutMs,
+        new HashMap[TopicPartition, MemoryRecords]()).build()
+      val emptyHeader = new RequestHeader(ApiKeys.PRODUCE, emptyRequest.version, clientId, correlationId)
+      sendApiRequest(socket, emptyRequest, emptyHeader)
+      // wait a little bit for the server-side disconnection to occur since it happens asynchronously
+      try {
+        TestUtils.waitUntilTrue(() => overrideServer.testableSelector.channels.isEmpty,
+          "Expired connection was not closed", 1000, 100)
+      } finally {
+        socket.close()
+      }
+    } finally {
+      shutdownServerAndMetrics(overrideServer)
+    }
+  }
+
+  @Test
+  def testSessionPrincipal(): Unit = {
     val socket = connect()
     val bytes = new Array[Byte](40)
     sendRequest(socket, bytes, Some(0))
@@ -662,7 +790,7 @@ class SocketServerTest extends JUnitSuite {
 
   /* Test that we update request metrics if the client closes the connection while the broker response is in flight. */
   @Test
-  def testClientDisconnectionUpdatesRequestMetrics() {
+  def testClientDisconnectionUpdatesRequestMetrics(): Unit = {
     val props = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 0)
     val serverMetrics = new Metrics
     var conn: Socket = null
@@ -672,7 +800,7 @@ class SocketServerTest extends JUnitSuite {
         new Processor(id, time, config.socketRequestMaxBytes, dataPlaneRequestChannel, connectionQuotas,
           config.connectionsMaxIdleMs, config.failedAuthenticationDelayMs, listenerName, protocol, config, metrics,
           credentialProvider, MemoryPool.NONE, new LogContext()) {
-          override protected[network] def sendResponse(response: RequestChannel.Response, responseSend: Send) {
+          override protected[network] def sendResponse(response: RequestChannel.Response, responseSend: Send): Unit = {
             conn.close()
             super.sendResponse(response, responseSend)
           }
@@ -708,7 +836,7 @@ class SocketServerTest extends JUnitSuite {
   }
 
   @Test
-  def testClientDisconnectionWithStagedReceivesFullyProcessed() {
+  def testClientDisconnectionWithStagedReceivesFullyProcessed(): Unit = {
     val serverMetrics = new Metrics
     @volatile var selector: TestableSelector = null
     val overrideConnectionId = "127.0.0.1:1-127.0.0.1:2-0"
@@ -761,7 +889,7 @@ class SocketServerTest extends JUnitSuite {
    * `selector.send` (selector closes old connections, for example).
    */
   @Test
-  def testBrokerSendAfterChannelClosedUpdatesRequestMetrics() {
+  def testBrokerSendAfterChannelClosedUpdatesRequestMetrics(): Unit = {
     val props = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 0)
     props.setProperty(KafkaConfig.ConnectionsMaxIdleMsProp, "110")
     val serverMetrics = new Metrics
@@ -808,8 +936,7 @@ class SocketServerTest extends JUnitSuite {
     def requestMetricMeters = YammerMetrics
       .defaultRegistry
       .allMetrics.asScala
-      .filterKeys(k => k.getType == "RequestMetrics")
-      .collect { case (k, metric: Meter) => (k.toString, metric.count) }
+      .collect { case (k, metric: Meter) if k.getType == "RequestMetrics" => (k.toString, metric.count) }
 
     assertEquals(nonZeroMeters, requestMetricMeters.filter { case (_, value) => value != 0 })
     server.shutdown()
@@ -823,7 +950,7 @@ class SocketServerTest extends JUnitSuite {
     val nonZeroMetricNamesAndValues = YammerMetrics
       .defaultRegistry
       .allMetrics.asScala
-      .filterKeys(k => k.getName.endsWith("IdlePercent") || k.getName.endsWith("NetworkProcessorAvgIdlePercent"))
+      .filter { case (k, _) => k.getName.endsWith("IdlePercent") || k.getName.endsWith("NetworkProcessorAvgIdlePercent") }
       .collect { case (k, metric: Gauge[_]) => (k, metric.value().asInstanceOf[Double]) }
       .filter { case (_, value) => value != 0.0 && !value.equals(Double.NaN) }
 
@@ -1069,7 +1196,7 @@ class SocketServerTest extends JUnitSuite {
 
       testableSelector.operationCounts.clear()
       testableSelector.addFailure(SelectorOperation.Poll,
-          Some(new RuntimeException("ControlThrowable exception during poll()") with ControlThrowable))
+          Some(new ControlThrowable() {}))
       testableSelector.waitForOperations(SelectorOperation.Poll, 1)
 
       testableSelector.waitForOperations(SelectorOperation.CloseSelector, 1)
@@ -1087,8 +1214,10 @@ class SocketServerTest extends JUnitSuite {
     val errors = new mutable.HashSet[String]
 
     def acceptorStackTraces: scala.collection.Map[Thread, String] = {
-      Thread.getAllStackTraces.asScala.filterKeys(_.getName.contains("kafka-socket-acceptor"))
-        .mapValues(_.toList.mkString("\n"))
+      Thread.getAllStackTraces.asScala.collect {
+        case (thread, stacktraceElement) if thread.getName.contains("kafka-socket-acceptor") =>
+          thread -> stacktraceElement.mkString("\n")
+      }
     }
 
     def acceptorBlocked: Boolean = {
@@ -1190,8 +1319,9 @@ class SocketServerTest extends JUnitSuite {
     }
   }
 
-  class TestableSocketServer(config : KafkaConfig = config, val connectionQueueSize: Int = 20) extends SocketServer(config,
-      new Metrics, Time.SYSTEM, credentialProvider) {
+  class TestableSocketServer(config : KafkaConfig = config, val connectionQueueSize: Int = 20,
+                             override val time: Time = Time.SYSTEM) extends SocketServer(config,
+      new Metrics, time, credentialProvider) {
 
     @volatile var selector: Option[TestableSelector] = None
 
@@ -1248,7 +1378,7 @@ class SocketServerTest extends JUnitSuite {
         extends Selector(config.socketRequestMaxBytes, config.connectionsMaxIdleMs, config.failedAuthenticationDelayMs,
             metrics, time, "socket-server", metricTags.asJava, false, true, channelBuilder, MemoryPool.NONE, new LogContext()) {
 
-    val failures = mutable.Map[SelectorOperation, Exception]()
+    val failures = mutable.Map[SelectorOperation, Throwable]()
     val operationCounts = mutable.Map[SelectorOperation, Int]().withDefaultValue(0)
     val allChannels = mutable.Set[String]()
     val allLocallyClosedChannels = mutable.Set[String]()
@@ -1284,7 +1414,7 @@ class SocketServerTest extends JUnitSuite {
     @volatile var pollTimeoutOverride: Option[Long] = None
     @volatile var pollCallback: () => Unit = () => {}
 
-    def addFailure(operation: SelectorOperation, exception: Option[Exception] = None) {
+    def addFailure(operation: SelectorOperation, exception: Option[Throwable] = None): Unit = {
       failures += operation ->
         exception.getOrElse(new IllegalStateException(s"Test exception during $operation"))
     }
diff --git a/core/src/test/scala/unit/kafka/security/auth/AclTest.scala b/core/src/test/scala/unit/kafka/security/auth/AclTest.scala
index beeac3766538b..a06d7a6d1b71c 100644
--- a/core/src/test/scala/unit/kafka/security/auth/AclTest.scala
+++ b/core/src/test/scala/unit/kafka/security/auth/AclTest.scala
@@ -21,7 +21,7 @@ import java.nio.charset.StandardCharsets.UTF_8
 import kafka.utils.Json
 import org.apache.kafka.common.security.auth.KafkaPrincipal
 import org.junit.{Assert, Test}
-import org.scalatest.junit.JUnitSuite
+import org.scalatestplus.junit.JUnitSuite
 import scala.collection.JavaConverters._
 
 class AclTest extends JUnitSuite {
diff --git a/core/src/test/scala/unit/kafka/security/auth/OperationTest.scala b/core/src/test/scala/unit/kafka/security/auth/OperationTest.scala
index 242c76809c7f0..6773096ed7f66 100644
--- a/core/src/test/scala/unit/kafka/security/auth/OperationTest.scala
+++ b/core/src/test/scala/unit/kafka/security/auth/OperationTest.scala
@@ -20,9 +20,8 @@ package kafka.security.auth
 import org.apache.kafka.common.acl.AclOperation
 import org.junit.Assert.assertEquals
 import org.junit.Test
-import org.scalatest.junit.JUnitSuite
 
-class OperationTest extends JUnitSuite {
+class OperationTest {
   /**
     * Test round trip conversions between org.apache.kafka.common.acl.AclOperation and
     * kafka.security.auth.Operation.
diff --git a/core/src/test/scala/unit/kafka/security/auth/PermissionTypeTest.scala b/core/src/test/scala/unit/kafka/security/auth/PermissionTypeTest.scala
index 0ee66e6bf6d82..8b2c5bc32cb8a 100644
--- a/core/src/test/scala/unit/kafka/security/auth/PermissionTypeTest.scala
+++ b/core/src/test/scala/unit/kafka/security/auth/PermissionTypeTest.scala
@@ -20,9 +20,9 @@ import kafka.common.KafkaException
 import org.apache.kafka.common.acl.AclPermissionType
 import org.junit.Assert.assertEquals
 import org.junit.Test
-import org.scalatest.junit.JUnitSuite
+import org.scalatest.Assertions.fail
 
-class PermissionTypeTest extends JUnitSuite {
+class PermissionTypeTest {
 
   @Test
   def testFromString(): Unit = {
diff --git a/core/src/test/scala/unit/kafka/security/auth/ResourceTypeTest.scala b/core/src/test/scala/unit/kafka/security/auth/ResourceTypeTest.scala
index 0d9937896efc8..bfefbba35c678 100644
--- a/core/src/test/scala/unit/kafka/security/auth/ResourceTypeTest.scala
+++ b/core/src/test/scala/unit/kafka/security/auth/ResourceTypeTest.scala
@@ -19,10 +19,10 @@ package kafka.security.auth
 import kafka.common.KafkaException
 import org.junit.Assert.assertEquals
 import org.junit.Test
-import org.scalatest.junit.JUnitSuite
+import org.scalatest.Assertions.fail
 import org.apache.kafka.common.resource.{ResourceType => JResourceType}
 
-class ResourceTypeTest extends JUnitSuite {
+class ResourceTypeTest {
 
   @Test
   def testFromString(): Unit = {
diff --git a/core/src/test/scala/unit/kafka/security/auth/SimpleAclAuthorizerTest.scala b/core/src/test/scala/unit/kafka/security/auth/SimpleAclAuthorizerTest.scala
index 1468003f3f011..d5169646e1616 100644
--- a/core/src/test/scala/unit/kafka/security/auth/SimpleAclAuthorizerTest.scala
+++ b/core/src/test/scala/unit/kafka/security/auth/SimpleAclAuthorizerTest.scala
@@ -19,13 +19,12 @@ package kafka.security.auth
 import java.net.InetAddress
 import java.nio.charset.StandardCharsets.UTF_8
 import java.util.UUID
-import java.util.concurrent.{Executors, Semaphore, TimeUnit}
 
 import kafka.api.{ApiVersion, KAFKA_2_0_IV0, KAFKA_2_0_IV1}
 import kafka.network.RequestChannel.Session
 import kafka.security.auth.Acl.{WildCardHost, WildCardResource}
 import kafka.server.KafkaConfig
-import kafka.utils.{CoreUtils, TestUtils}
+import kafka.utils.TestUtils
 import kafka.zk.{ZkAclStore, ZooKeeperTestHarness}
 import kafka.zookeeper.{GetChildrenRequest, GetDataRequest, ZooKeeperClient}
 import org.apache.kafka.common.errors.UnsupportedVersionException
@@ -36,6 +35,7 @@ import org.apache.kafka.common.utils.Time
 import org.junit.Assert._
 import org.junit.{After, Before, Test}
 
+@deprecated("Use AclAuthorizer", "Since 2.4")
 class SimpleAclAuthorizerTest extends ZooKeeperTestHarness {
 
   private val allowReadAcl = Acl(Acl.WildCardPrincipal, Allow, WildCardHost, Read)
@@ -60,7 +60,7 @@ class SimpleAclAuthorizerTest extends ZooKeeperTestHarness {
   }
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     super.setUp()
 
     // Increase maxUpdateRetries to avoid transient failures
@@ -88,7 +88,7 @@ class SimpleAclAuthorizerTest extends ZooKeeperTestHarness {
   }
 
   @Test(expected = classOf[IllegalArgumentException])
-  def testAuthorizeThrowsOnNoneLiteralResource() {
+  def testAuthorizeThrowsOnNonLiteralResource(): Unit = {
     simpleAclAuthorizer.authorize(session, Read, Resource(Topic, "something", PREFIXED))
   }
 
@@ -106,7 +106,7 @@ class SimpleAclAuthorizerTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  def testTopicAcl() {
+  def testTopicAcl(): Unit = {
     val user1 = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, username)
     val user2 = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "rob")
     val user3 = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "batman")
@@ -161,7 +161,7 @@ class SimpleAclAuthorizerTest extends ZooKeeperTestHarness {
     CustomPrincipals should be compared with their principal type and name
    */
   @Test
-  def testAllowAccessWithCustomPrincipal() {
+  def testAllowAccessWithCustomPrincipal(): Unit = {
     val user = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, username)
     val customUserPrincipal = new CustomPrincipal(KafkaPrincipal.USER_TYPE, username)
     val host1 = InetAddress.getByName("192.168.1.1")
@@ -181,7 +181,7 @@ class SimpleAclAuthorizerTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  def testDenyTakesPrecedence() {
+  def testDenyTakesPrecedence(): Unit = {
     val user = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, username)
     val host = InetAddress.getByName("192.168.2.1")
     val session = Session(user, host)
@@ -196,7 +196,7 @@ class SimpleAclAuthorizerTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  def testAllowAllAccess() {
+  def testAllowAllAccess(): Unit = {
     val allowAllAcl = Acl.AllowAllAcl
 
     changeAclAndVerify(Set.empty[Acl], Set[Acl](allowAllAcl), Set.empty[Acl])
@@ -206,7 +206,7 @@ class SimpleAclAuthorizerTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  def testSuperUserHasAccess() {
+  def testSuperUserHasAccess(): Unit = {
     val denyAllAcl = new Acl(Acl.WildCardPrincipal, Deny, WildCardHost, All)
 
     changeAclAndVerify(Set.empty[Acl], Set[Acl](denyAllAcl), Set.empty[Acl])
@@ -256,12 +256,12 @@ class SimpleAclAuthorizerTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  def testNoAclFound() {
+  def testNoAclFound(): Unit = {
     assertFalse("when acls = [],  authorizer should fail close.", simpleAclAuthorizer.authorize(session, Read, resource))
   }
 
   @Test
-  def testNoAclFoundOverride() {
+  def testNoAclFoundOverride(): Unit = {
     val props = TestUtils.createBrokerConfig(1, zkConnect)
     props.put(SimpleAclAuthorizer.AllowEveryoneIfNoAclIsFoundProp, "true")
 
@@ -276,7 +276,7 @@ class SimpleAclAuthorizerTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  def testAclManagementAPIs() {
+  def testAclManagementAPIs(): Unit = {
     val user1 = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, username)
     val user2 = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "bob")
     val host1 = "host1"
@@ -322,7 +322,7 @@ class SimpleAclAuthorizerTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  def testLoadCache() {
+  def testLoadCache(): Unit = {
     val user1 = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, username)
     val acl1 = new Acl(user1, Allow, "host-1", Read)
     val acls = Set[Acl](acl1)
@@ -346,42 +346,8 @@ class SimpleAclAuthorizerTest extends ZooKeeperTestHarness {
     }
   }
 
-  /**
-   * Verify that there is no timing window between loading ACL cache and setting
-   * up ZK change listener. Cache must be loaded before creating change listener
-   * in the authorizer to avoid the timing window.
-   */
-  @Test
-  def testChangeListenerTiming() {
-    val configureSemaphore = new Semaphore(0)
-    val listenerSemaphore = new Semaphore(0)
-    val executor = Executors.newSingleThreadExecutor
-    val simpleAclAuthorizer3 = new SimpleAclAuthorizer {
-      override private[auth] def startZkChangeListeners(): Unit = {
-        configureSemaphore.release()
-        listenerSemaphore.acquireUninterruptibly()
-        super.startZkChangeListeners()
-      }
-    }
-    try {
-      val future = executor.submit(CoreUtils.runnable(simpleAclAuthorizer3.configure(config.originals)))
-      configureSemaphore.acquire()
-      val user1 = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, username)
-      val acls = Set(new Acl(user1, Deny, "host-1", Read))
-      simpleAclAuthorizer.addAcls(acls, resource)
-
-      listenerSemaphore.release()
-      future.get(10, TimeUnit.SECONDS)
-
-      assertEquals(acls, simpleAclAuthorizer3.getAcls(resource))
-    } finally {
-      simpleAclAuthorizer3.close()
-      executor.shutdownNow()
-    }
-  }
-
   @Test
-  def testLocalConcurrentModificationOfResourceAcls() {
+  def testLocalConcurrentModificationOfResourceAcls(): Unit = {
     val commonResource = Resource(Topic, "test", LITERAL)
 
     val user1 = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, username)
@@ -397,7 +363,7 @@ class SimpleAclAuthorizerTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  def testDistributedConcurrentModificationOfResourceAcls() {
+  def testDistributedConcurrentModificationOfResourceAcls(): Unit = {
     val commonResource = Resource(Topic, "test", LITERAL)
 
     val user1 = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, username)
@@ -427,7 +393,7 @@ class SimpleAclAuthorizerTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  def testHighConcurrencyModificationOfResourceAcls() {
+  def testHighConcurrencyModificationOfResourceAcls(): Unit = {
     val commonResource = Resource(Topic, "test", LITERAL)
 
     val acls = (0 to 50).map { i =>
@@ -512,7 +478,7 @@ class SimpleAclAuthorizerTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  def testHighConcurrencyDeletionOfResourceAcls() {
+  def testHighConcurrencyDeletionOfResourceAcls(): Unit = {
     val acl = new Acl(new KafkaPrincipal(KafkaPrincipal.USER_TYPE, username), Allow, WildCardHost, All)
 
     // Alternate authorizer to keep adding and removing ZooKeeper path
@@ -721,7 +687,7 @@ class SimpleAclAuthorizerTest extends ZooKeeperTestHarness {
     assertEquals(expected, actual)
   }
 
-  private def givenAuthorizerWithProtocolVersion(protocolVersion: Option[ApiVersion]) {
+  private def givenAuthorizerWithProtocolVersion(protocolVersion: Option[ApiVersion]): Unit = {
     simpleAclAuthorizer.close()
 
     val props = TestUtils.createBrokerConfig(0, zkConnect)
diff --git a/core/src/test/scala/unit/kafka/security/auth/ZkAuthorizationTest.scala b/core/src/test/scala/unit/kafka/security/auth/ZkAuthorizationTest.scala
index cc845553181d1..6780d98a5367b 100644
--- a/core/src/test/scala/unit/kafka/security/auth/ZkAuthorizationTest.scala
+++ b/core/src/test/scala/unit/kafka/security/auth/ZkAuthorizationTest.scala
@@ -32,6 +32,7 @@ import scala.util.{Failure, Success, Try}
 import javax.security.auth.login.Configuration
 import kafka.api.ApiVersion
 import kafka.cluster.{Broker, EndPoint}
+import kafka.controller.ReplicaAssignment
 import org.apache.kafka.common.network.ListenerName
 import org.apache.kafka.common.security.auth.SecurityProtocol
 import org.apache.kafka.common.utils.Time
@@ -44,7 +45,7 @@ class ZkAuthorizationTest extends ZooKeeperTestHarness with Logging {
   val authProvider = "zookeeper.authProvider.1"
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     System.setProperty(JaasUtils.JAVA_LOGIN_CONFIG_PARAM, jaasFile.getAbsolutePath)
     Configuration.setConfiguration(null)
     System.setProperty(authProvider, "org.apache.zookeeper.server.auth.SASLAuthenticationProvider")
@@ -52,7 +53,7 @@ class ZkAuthorizationTest extends ZooKeeperTestHarness with Logging {
   }
 
   @After
-  override def tearDown() {
+  override def tearDown(): Unit = {
     super.tearDown()
     System.clearProperty(JaasUtils.JAVA_LOGIN_CONFIG_PARAM)
     System.clearProperty(authProvider)
@@ -130,7 +131,7 @@ class ZkAuthorizationTest extends ZooKeeperTestHarness with Logging {
 
     // Test that can update persistent nodes
     val updatedAssignment = assignment - new TopicPartition(topic1, 2)
-    zkClient.setTopicAssignment(topic1, updatedAssignment)
+    zkClient.setTopicAssignment(topic1, updatedAssignment.mapValues { case (v) => ReplicaAssignment(v, List(), List()) }.toMap)
     assertEquals(updatedAssignment.size, zkClient.getTopicPartitionCount(topic1).get)
   }
 
diff --git a/core/src/test/scala/unit/kafka/security/authorizer/AclAuthorizerTest.scala b/core/src/test/scala/unit/kafka/security/authorizer/AclAuthorizerTest.scala
new file mode 100644
index 0000000000000..6250958596620
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/security/authorizer/AclAuthorizerTest.scala
@@ -0,0 +1,884 @@
+/**
+ * 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 java.net.InetAddress
+import java.nio.charset.StandardCharsets.UTF_8
+import java.util.UUID
+import java.util.concurrent.{Executors, Semaphore, TimeUnit}
+
+import kafka.api.{ApiVersion, KAFKA_2_0_IV0, KAFKA_2_0_IV1}
+import kafka.security.auth.Resource
+import kafka.security.authorizer.AuthorizerUtils.{WildcardHost, WildcardPrincipal}
+import kafka.server.KafkaConfig
+import kafka.utils.{CoreUtils, TestUtils}
+import kafka.zk.{ZkAclStore, ZooKeeperTestHarness}
+import kafka.zookeeper.{GetChildrenRequest, GetDataRequest, ZooKeeperClient}
+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.common.errors.{ApiException, UnsupportedVersionException}
+import org.apache.kafka.common.network.ListenerName
+import org.apache.kafka.common.protocol.ApiKeys
+import org.apache.kafka.common.requests.{RequestContext, RequestHeader}
+import org.apache.kafka.common.resource.{PatternType, ResourcePattern, ResourcePatternFilter, ResourceType}
+import org.apache.kafka.common.resource.ResourcePattern.WILDCARD_RESOURCE
+import org.apache.kafka.common.resource.ResourceType._
+import org.apache.kafka.common.resource.PatternType.{LITERAL, MATCH, PREFIXED}
+import org.apache.kafka.common.security.auth.{KafkaPrincipal, SecurityProtocol}
+import org.apache.kafka.server.authorizer._
+import org.apache.kafka.common.utils.{Time, SecurityUtils => JSecurityUtils}
+import org.junit.Assert._
+import org.junit.{After, Before, Test}
+import org.scalatest.Assertions.intercept
+
+import scala.collection.JavaConverters._
+import scala.compat.java8.OptionConverters._
+
+class AclAuthorizerTest extends ZooKeeperTestHarness {
+
+  private val allowReadAcl = new AccessControlEntry(WildcardPrincipal, WildcardHost, READ, ALLOW)
+  private val allowWriteAcl = new AccessControlEntry(WildcardPrincipal, WildcardHost, WRITE, ALLOW)
+  private val denyReadAcl = new AccessControlEntry(WildcardPrincipal, WildcardHost, 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, Resource.ClusterResourceName, LITERAL)
+  private val wildcardPrincipal = JSecurityUtils.parseKafkaPrincipal(WildcardPrincipal)
+
+  private val aclAuthorizer = new AclAuthorizer
+  private val aclAuthorizer2 = new AclAuthorizer
+  private var resource: ResourcePattern = _
+  private val superUsers = "User:superuser1; User:superuser2"
+  private val username = "alice"
+  private val principal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, username)
+  private val requestContext = newRequestContext(principal, InetAddress.getByName("192.168.0.1"))
+  private var config: KafkaConfig = _
+  private var zooKeeperClient: ZooKeeperClient = _
+
+  class CustomPrincipal(principalType: String, name: String) extends KafkaPrincipal(principalType, name) {
+    override def equals(o: scala.Any): Boolean = false
+  }
+
+  @Before
+  override def setUp(): Unit = {
+    super.setUp()
+
+    // Increase maxUpdateRetries to avoid transient failures
+    aclAuthorizer.maxUpdateRetries = Int.MaxValue
+    aclAuthorizer2.maxUpdateRetries = Int.MaxValue
+
+    val props = TestUtils.createBrokerConfig(0, zkConnect)
+    props.put(AclAuthorizer.SuperUsersProp, superUsers)
+
+    config = KafkaConfig.fromProps(props)
+    aclAuthorizer.configure(config.originals)
+    aclAuthorizer2.configure(config.originals)
+    resource = new ResourcePattern(TOPIC, "foo-" + UUID.randomUUID(), LITERAL)
+
+    zooKeeperClient = new ZooKeeperClient(zkConnect, zkSessionTimeout, zkConnectionTimeout, zkMaxInFlightRequests,
+      Time.SYSTEM, "kafka.test", "AclAuthorizerTest")
+  }
+
+  @After
+  override def tearDown(): Unit = {
+    aclAuthorizer.close()
+    aclAuthorizer2.close()
+    zooKeeperClient.close()
+    super.tearDown()
+  }
+
+  @Test(expected = classOf[IllegalArgumentException])
+  def testAuthorizeThrowsOnNonLiteralResource(): Unit = {
+    authorize(aclAuthorizer, requestContext, READ, new ResourcePattern(TOPIC, "something", PREFIXED))
+  }
+
+  @Test
+  def testAuthorizeWithEmptyResourceName(): Unit = {
+    assertFalse(authorize(aclAuthorizer, requestContext, READ, new ResourcePattern(GROUP, "", LITERAL)))
+    addAcls(aclAuthorizer, Set(allowReadAcl), new ResourcePattern(GROUP, WILDCARD_RESOURCE, LITERAL))
+    assertTrue(authorize(aclAuthorizer, requestContext, READ, new ResourcePattern(GROUP, "", LITERAL)))
+  }
+
+  // Authorizing the empty resource is not supported because we create a znode with the resource name.
+  @Test
+  def testEmptyAclThrowsException(): Unit = {
+    val e = intercept[ApiException] {
+      addAcls(aclAuthorizer, Set(allowReadAcl), new ResourcePattern(GROUP, "", LITERAL))
+    }
+    assertTrue(s"Unexpected exception $e", e.getCause.isInstanceOf[IllegalArgumentException])
+  }
+
+  @Test
+  def testTopicAcl(): Unit = {
+    val user1 = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, username)
+    val user2 = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "rob")
+    val user3 = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "batman")
+    val host1 = InetAddress.getByName("192.168.1.1")
+    val host2 = InetAddress.getByName("192.168.1.2")
+
+    //user1 has READ access from host1 and host2.
+    val acl1 = new AccessControlEntry(user1.toString, host1.getHostAddress, READ, ALLOW)
+    val acl2 = new AccessControlEntry(user1.toString, host2.getHostAddress, READ, ALLOW)
+
+    //user1 does not have  READ access from host1.
+    val acl3 = new AccessControlEntry(user1.toString, host1.getHostAddress, READ, DENY)
+
+    //user1 has WRITE access from host1 only.
+    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)
+
+    //user2 has READ access from all hosts.
+    val acl6 = new AccessControlEntry(user2.toString, WildcardHost, READ, ALLOW)
+
+    //user3 has WRITE access from all hosts.
+    val acl7 = new AccessControlEntry(user3.toString, WildcardHost, WRITE, ALLOW)
+
+    val acls = Set(acl1, acl2, acl3, acl4, acl5, acl6, acl7)
+
+    changeAclAndVerify(Set.empty, acls, Set.empty)
+
+    val host1Context = newRequestContext(user1, host1)
+    val host2Context = newRequestContext(user1, host2)
+
+    assertTrue("User1 should have READ access from host2", authorize(aclAuthorizer, host2Context, READ, resource))
+    assertFalse("User1 should not have READ access from host1 due to denyAcl", authorize(aclAuthorizer, host1Context, READ, resource))
+    assertTrue("User1 should have WRITE access from host1", authorize(aclAuthorizer, host1Context, WRITE, resource))
+    assertFalse("User1 should not have WRITE access from host2 as no allow acl is defined", authorize(aclAuthorizer, host2Context, WRITE, resource))
+    assertTrue("User1 should not have DESCRIBE access from host1", authorize(aclAuthorizer, host1Context, DESCRIBE, resource))
+    assertTrue("User1 should have DESCRIBE access from host2", authorize(aclAuthorizer, host2Context, DESCRIBE, resource))
+    assertFalse("User1 should not have edit access from host1", authorize(aclAuthorizer, host1Context, ALTER, resource))
+    assertFalse("User1 should not have edit access from host2", authorize(aclAuthorizer, host2Context, ALTER, resource))
+
+    //test if user has READ and write access they also get describe access
+    val user2Context = newRequestContext(user2, host1)
+    val user3Context = newRequestContext(user3, host1)
+    assertTrue("User2 should have DESCRIBE access from host1", authorize(aclAuthorizer, user2Context, DESCRIBE, resource))
+    assertTrue("User3 should have DESCRIBE access from host2", authorize(aclAuthorizer, user3Context, DESCRIBE, resource))
+    assertTrue("User2 should have READ access from host1", authorize(aclAuthorizer, user2Context, READ, resource))
+    assertTrue("User3 should have WRITE access from host2", authorize(aclAuthorizer, user3Context, WRITE, resource))
+  }
+
+  /**
+    CustomPrincipals should be compared with their principal type and name
+   */
+  @Test
+  def testAllowAccessWithCustomPrincipal(): Unit = {
+    val user = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, username)
+    val customUserPrincipal = new CustomPrincipal(KafkaPrincipal.USER_TYPE, username)
+    val host1 = InetAddress.getByName("192.168.1.1")
+    val host2 = InetAddress.getByName("192.168.1.2")
+
+    // user has READ access from host2 but not from host1
+    val acl1 = new AccessControlEntry(user.toString, host1.getHostAddress, READ, DENY)
+    val acl2 = new AccessControlEntry(user.toString, host2.getHostAddress, READ, ALLOW)
+    val acls = Set(acl1, acl2)
+    changeAclAndVerify(Set.empty, acls, Set.empty)
+
+    val host1Context = newRequestContext(customUserPrincipal, host1)
+    val host2Context = newRequestContext(customUserPrincipal, host2)
+
+    assertTrue("User1 should have READ access from host2", authorize(aclAuthorizer, host2Context, READ, resource))
+    assertFalse("User1 should not have READ access from host1 due to denyAcl", authorize(aclAuthorizer, host1Context, READ, resource))
+  }
+
+  @Test
+  def testDenyTakesPrecedence(): Unit = {
+    val user = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, username)
+    val host = InetAddress.getByName("192.168.2.1")
+    val session = newRequestContext(user, host)
+
+    val allowAll = new AccessControlEntry(WildcardPrincipal, WildcardHost, AclOperation.ALL, ALLOW)
+    val denyAcl = new AccessControlEntry(user.toString, host.getHostAddress, AclOperation.ALL, DENY)
+    val acls = Set(allowAll, denyAcl)
+
+    changeAclAndVerify(Set.empty, acls, Set.empty)
+
+    assertFalse("deny should take precedence over allow.", authorize(aclAuthorizer, session, READ, resource))
+  }
+
+  @Test
+  def testAllowAllAccess(): Unit = {
+    val allowAllAcl = new AccessControlEntry(WildcardPrincipal, WildcardHost, AclOperation.ALL, ALLOW)
+
+    changeAclAndVerify(Set.empty, Set(allowAllAcl), Set.empty)
+
+    val context = newRequestContext(new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "random"), InetAddress.getByName("192.0.4.4"))
+    assertTrue("allow all acl should allow access to all.", authorize(aclAuthorizer, context, READ, resource))
+  }
+
+  @Test
+  def testSuperUserHasAccess(): Unit = {
+    val denyAllAcl = new AccessControlEntry(WildcardPrincipal, WildcardHost, AclOperation.ALL, DENY)
+
+    changeAclAndVerify(Set.empty, Set(denyAllAcl), Set.empty)
+
+    val session1 = newRequestContext(new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "superuser1"), InetAddress.getByName("192.0.4.4"))
+    val session2 = newRequestContext(new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "superuser2"), InetAddress.getByName("192.0.4.4"))
+
+    assertTrue("superuser always has access, no matter what acls.", authorize(aclAuthorizer, session1, READ, resource))
+    assertTrue("superuser always has access, no matter what acls.", authorize(aclAuthorizer, session2, READ, resource))
+  }
+
+  /**
+    CustomPrincipals should be compared with their principal type and name
+   */
+  @Test
+  def testSuperUserWithCustomPrincipalHasAccess(): Unit = {
+    val denyAllAcl = new AccessControlEntry(WildcardPrincipal, WildcardHost, 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"))
+
+    assertTrue("superuser with custom principal always has access, no matter what acls.", authorize(aclAuthorizer, session, READ, resource))
+  }
+
+  @Test
+  def testWildCardAcls(): Unit = {
+    assertFalse("when acls = [], authorizer should fail close.", authorize(aclAuthorizer, requestContext, READ, resource))
+
+    val user1 = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, username)
+    val host1 = InetAddress.getByName("192.168.3.1")
+    val readAcl = new AccessControlEntry(user1.toString, host1.getHostAddress, READ, ALLOW)
+
+    val acls = changeAclAndVerify(Set.empty, Set(readAcl), Set.empty, wildCardResource)
+
+    val host1Context = newRequestContext(user1, host1)
+    assertTrue("User1 should have READ access from host1", authorize(aclAuthorizer, host1Context, READ, resource))
+
+    //allow WRITE to specific topic.
+    val writeAcl = new AccessControlEntry(user1.toString, host1.getHostAddress, WRITE, ALLOW)
+    changeAclAndVerify(Set.empty, Set(writeAcl), Set.empty)
+
+    //deny WRITE to wild card topic.
+    val denyWriteOnWildCardResourceAcl = new AccessControlEntry(user1.toString, host1.getHostAddress, WRITE, DENY)
+    changeAclAndVerify(acls, Set(denyWriteOnWildCardResourceAcl), Set.empty, wildCardResource)
+
+    assertFalse("User1 should not have WRITE access from host1", authorize(aclAuthorizer, host1Context, WRITE, resource))
+  }
+
+  @Test
+  def testNoAclFound(): Unit = {
+    assertFalse("when acls = [], authorizer should deny op.", authorize(aclAuthorizer, requestContext, READ, resource))
+  }
+
+  @Test
+  def testNoAclFoundOverride(): Unit = {
+    val props = TestUtils.createBrokerConfig(1, zkConnect)
+    props.put(AclAuthorizer.AllowEveryoneIfNoAclIsFoundProp, "true")
+
+    val cfg = KafkaConfig.fromProps(props)
+    val testAuthorizer = new AclAuthorizer
+    try {
+      testAuthorizer.configure(cfg.originals)
+      assertTrue("when acls = null or [],  authorizer should allow op with allow.everyone = true.",
+        authorize(testAuthorizer, requestContext, READ, resource))
+    } finally {
+      testAuthorizer.close()
+    }
+  }
+
+  @Test
+  def testAclManagementAPIs(): Unit = {
+    val user1 = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, username)
+    val user2 = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "bob")
+    val host1 = "host1"
+    val host2 = "host2"
+
+    val acl1 = new AccessControlEntry(user1.toString, host1, READ, ALLOW)
+    val acl2 = new AccessControlEntry(user1.toString, host1, WRITE, ALLOW)
+    val acl3 = new AccessControlEntry(user2.toString, host2, READ, ALLOW)
+    val acl4 = new AccessControlEntry(user2.toString, host2, WRITE, ALLOW)
+
+    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)
+    acls = changeAclAndVerify(acls, Set(acl5), Set.empty)
+
+    //test get by principal name.
+    TestUtils.waitUntilTrue(() => Set(acl1, acl2).map(acl => new AclBinding(resource, acl)) == getAcls(aclAuthorizer, user1),
+      "changes not propagated in timeout period")
+    TestUtils.waitUntilTrue(() => Set(acl3, acl4, acl5).map(acl => new AclBinding(resource, acl)) == getAcls(aclAuthorizer, user2),
+      "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(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
+    )
+
+    resourceToAcls foreach { case (key, value) => changeAclAndVerify(Set.empty, value, Set.empty, key) }
+    val expectedAcls = (resourceToAcls + (resource -> acls)).flatMap {
+      case (res, resAcls) => resAcls.map { acl => new AclBinding(res, acl) }
+    }.toSet
+    TestUtils.waitUntilTrue(() => expectedAcls == getAcls(aclAuthorizer), "changes not propagated in timeout period.")
+
+    //test remove acl from existing acls.
+    acls = changeAclAndVerify(acls, Set.empty, Set(acl1, acl5))
+
+    //test remove all acls for resource
+    removeAcls(aclAuthorizer, Set.empty, resource)
+    TestUtils.waitAndVerifyAcls(Set.empty[AccessControlEntry], aclAuthorizer, resource)
+    assertTrue(!zkClient.resourceExists(AuthorizerUtils.convertToResource(resource)))
+
+    //test removing last acl also deletes ZooKeeper path
+    acls = changeAclAndVerify(Set.empty, Set(acl1), Set.empty)
+    changeAclAndVerify(acls, Set.empty, acls)
+    assertTrue(!zkClient.resourceExists(AuthorizerUtils.convertToResource(resource)))
+  }
+
+  @Test
+  def testLoadCache(): Unit = {
+    val user1 = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, username)
+    val acl1 = new AccessControlEntry(user1.toString, "host-1", READ, ALLOW)
+    val acls = Set(acl1)
+    addAcls(aclAuthorizer, acls, resource)
+
+    val user2 = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "bob")
+    val resource1 = new ResourcePattern(TOPIC, "test-2", LITERAL)
+    val acl2 = new AccessControlEntry(user2.toString, "host3", READ, DENY)
+    val acls1 = Set(acl2)
+    addAcls(aclAuthorizer, acls1, resource1)
+
+    zkClient.deleteAclChangeNotifications
+    val authorizer = new AclAuthorizer
+    try {
+      authorizer.configure(config.originals)
+
+      assertEquals(acls, getAcls(authorizer, resource))
+      assertEquals(acls1, getAcls(authorizer, resource1))
+    } finally {
+      authorizer.close()
+    }
+  }
+
+  /**
+   * Verify that there is no timing window between loading ACL cache and setting
+   * up ZK change listener. Cache must be loaded before creating change listener
+   * in the authorizer to avoid the timing window.
+   */
+  @Test
+  def testChangeListenerTiming(): Unit = {
+    val configureSemaphore = new Semaphore(0)
+    val listenerSemaphore = new Semaphore(0)
+    val executor = Executors.newSingleThreadExecutor
+    val aclAuthorizer3 = new AclAuthorizer {
+      override private[authorizer] def startZkChangeListeners(): Unit = {
+        configureSemaphore.release()
+        listenerSemaphore.acquireUninterruptibly()
+        super.startZkChangeListeners()
+      }
+    }
+    try {
+      val future = executor.submit(CoreUtils.runnable(aclAuthorizer3.configure(config.originals)))
+      configureSemaphore.acquire()
+      val user1 = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, username)
+      val acls = Set(new AccessControlEntry(user1.toString, "host-1", READ, DENY))
+      addAcls(aclAuthorizer, acls, resource)
+
+      listenerSemaphore.release()
+      future.get(10, TimeUnit.SECONDS)
+
+      assertEquals(acls, getAcls(aclAuthorizer3, resource))
+    } finally {
+      aclAuthorizer3.close()
+      executor.shutdownNow()
+    }
+  }
+
+  @Test
+  def testLocalConcurrentModificationOfResourceAcls(): Unit = {
+    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 user2 = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "bob")
+    val acl2 = new AccessControlEntry(user2.toString, WildcardHost, READ, DENY)
+
+    addAcls(aclAuthorizer, Set(acl1), commonResource)
+    addAcls(aclAuthorizer, Set(acl2), commonResource)
+
+    TestUtils.waitAndVerifyAcls(Set(acl1, acl2), aclAuthorizer, commonResource)
+  }
+
+  @Test
+  def testDistributedConcurrentModificationOfResourceAcls(): Unit = {
+    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 user2 = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "bob")
+    val acl2 = new AccessControlEntry(user2.toString, WildcardHost, READ, DENY)
+
+    // Add on each instance
+    addAcls(aclAuthorizer, Set(acl1), commonResource)
+    addAcls(aclAuthorizer2, Set(acl2), commonResource)
+
+    TestUtils.waitAndVerifyAcls(Set(acl1, acl2), aclAuthorizer, commonResource)
+    TestUtils.waitAndVerifyAcls(Set(acl1, acl2), aclAuthorizer2, commonResource)
+
+    val user3 = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "joe")
+    val acl3 = new AccessControlEntry(user3.toString, WildcardHost, READ, DENY)
+
+    // Add on one instance and delete on another
+    addAcls(aclAuthorizer, Set(acl3), commonResource)
+    val deleted = removeAcls(aclAuthorizer2, Set(acl3), commonResource)
+
+    assertTrue("The authorizer should see a value that needs to be deleted", deleted)
+
+    TestUtils.waitAndVerifyAcls(Set(acl1, acl2), aclAuthorizer, commonResource)
+    TestUtils.waitAndVerifyAcls(Set(acl1, acl2), aclAuthorizer2, commonResource)
+  }
+
+  @Test
+  def testHighConcurrencyModificationOfResourceAcls(): Unit = {
+    val commonResource = new ResourcePattern(TOPIC, "test", LITERAL)
+
+    val acls= (0 to 50).map { i =>
+      val useri = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, i.toString)
+      (new AccessControlEntry(useri.toString, WildcardHost, READ, ALLOW), i)
+    }
+
+    // Alternate authorizer, Remove all acls that end in 0
+    val concurrentFuctions = acls.map { case (acl, aclId) =>
+      () => {
+        if (aclId % 2 == 0) {
+          addAcls(aclAuthorizer, Set(acl), commonResource)
+        } else {
+          addAcls(aclAuthorizer2, Set(acl), commonResource)
+        }
+        if (aclId % 10 == 0) {
+          removeAcls(aclAuthorizer2, Set(acl), commonResource)
+        }
+      }
+    }
+
+    val expectedAcls = acls.filter { case (acl, aclId) =>
+      aclId % 10 != 0
+    }.map(_._1).toSet
+
+    TestUtils.assertConcurrent("Should support many concurrent calls", concurrentFuctions, 30 * 1000)
+
+    TestUtils.waitAndVerifyAcls(expectedAcls, aclAuthorizer, commonResource)
+    TestUtils.waitAndVerifyAcls(expectedAcls, aclAuthorizer2, commonResource)
+  }
+
+  /**
+    * Test ACL inheritance, as described in #{org.apache.kafka.common.acl.AclOperation}
+    */
+  @Test
+  def testAclInheritance(): Unit = {
+    testImplicationsOfAllow(AclOperation.ALL, Set(READ, WRITE, CREATE, DELETE, ALTER, DESCRIBE,
+      CLUSTER_ACTION, DESCRIBE_CONFIGS, ALTER_CONFIGS, IDEMPOTENT_WRITE))
+    testImplicationsOfDeny(AclOperation.ALL, Set(READ, WRITE, CREATE, DELETE, ALTER, DESCRIBE,
+      CLUSTER_ACTION, DESCRIBE_CONFIGS, ALTER_CONFIGS, IDEMPOTENT_WRITE))
+    testImplicationsOfAllow(READ, Set(DESCRIBE))
+    testImplicationsOfAllow(WRITE, Set(DESCRIBE))
+    testImplicationsOfAllow(DELETE, Set(DESCRIBE))
+    testImplicationsOfAllow(ALTER, Set(DESCRIBE))
+    testImplicationsOfDeny(DESCRIBE, Set())
+    testImplicationsOfAllow(ALTER_CONFIGS, Set(DESCRIBE_CONFIGS))
+    testImplicationsOfDeny(DESCRIBE_CONFIGS, Set())
+  }
+
+  private def testImplicationsOfAllow(parentOp: AclOperation, allowedOps: Set[AclOperation]): Unit = {
+    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)
+    addAcls(aclAuthorizer, Set(acl), clusterResource)
+    AclOperation.values.filter(validOp).foreach { op =>
+      val authorized = authorize(aclAuthorizer, hostContext, op, clusterResource)
+      if (allowedOps.contains(op) || op == parentOp)
+        assertTrue(s"ALLOW $parentOp should imply ALLOW $op", authorized)
+      else
+        assertFalse(s"ALLOW $parentOp should not imply ALLOW $op", authorized)
+    }
+    removeAcls(aclAuthorizer, Set(acl), clusterResource)
+  }
+
+  private def testImplicationsOfDeny(parentOp: AclOperation, deniedOps: Set[AclOperation]): Unit = {
+    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))
+    addAcls(aclAuthorizer, acls, clusterResource)
+    AclOperation.values.filter(validOp).foreach { op =>
+      val authorized = authorize(aclAuthorizer, host1Context, op, clusterResource)
+      if (deniedOps.contains(op) || op == parentOp)
+        assertFalse(s"DENY $parentOp should imply DENY $op", authorized)
+      else
+        assertTrue(s"DENY $parentOp should not imply DENY $op", authorized)
+    }
+    removeAcls(aclAuthorizer, acls, clusterResource)
+  }
+
+  @Test
+  def testHighConcurrencyDeletionOfResourceAcls(): Unit = {
+    val acl = new AccessControlEntry(new KafkaPrincipal(KafkaPrincipal.USER_TYPE, username).toString, WildcardHost, AclOperation.ALL, ALLOW)
+
+    // Alternate authorizer to keep adding and removing ZooKeeper path
+    val concurrentFuctions = (0 to 50).map { _ =>
+      () => {
+        addAcls(aclAuthorizer, Set(acl), resource)
+        removeAcls(aclAuthorizer2, Set(acl), resource)
+      }
+    }
+
+    TestUtils.assertConcurrent("Should support many concurrent calls", concurrentFuctions, 30 * 1000)
+
+    TestUtils.waitAndVerifyAcls(Set.empty[AccessControlEntry], aclAuthorizer, resource)
+    TestUtils.waitAndVerifyAcls(Set.empty[AccessControlEntry], aclAuthorizer2, resource)
+  }
+
+  @Test
+  def testAccessAllowedIfAllowAclExistsOnWildcardResource(): Unit = {
+    addAcls(aclAuthorizer, Set(allowReadAcl), wildCardResource)
+
+    assertTrue(authorize(aclAuthorizer, requestContext, READ, resource))
+  }
+
+  @Test
+  def testDeleteAclOnWildcardResource(): Unit = {
+    addAcls(aclAuthorizer, Set(allowReadAcl, allowWriteAcl), wildCardResource)
+
+    removeAcls(aclAuthorizer, Set(allowReadAcl), wildCardResource)
+
+    assertEquals(Set(allowWriteAcl), getAcls(aclAuthorizer, wildCardResource))
+  }
+
+  @Test
+  def testDeleteAllAclOnWildcardResource(): Unit = {
+    addAcls(aclAuthorizer, Set(allowReadAcl), wildCardResource)
+
+    removeAcls(aclAuthorizer, Set.empty, wildCardResource)
+
+    assertEquals(Set.empty, getAcls(aclAuthorizer))
+  }
+
+  @Test
+  def testAccessAllowedIfAllowAclExistsOnPrefixedResource(): Unit = {
+    addAcls(aclAuthorizer, Set(allowReadAcl), prefixedResource)
+
+    assertTrue(authorize(aclAuthorizer, requestContext, READ, resource))
+  }
+
+  @Test
+  def testDeleteAclOnPrefixedResource(): Unit = {
+    addAcls(aclAuthorizer, Set(allowReadAcl, allowWriteAcl), prefixedResource)
+
+    removeAcls(aclAuthorizer, Set(allowReadAcl), prefixedResource)
+
+    assertEquals(Set(allowWriteAcl), getAcls(aclAuthorizer, prefixedResource))
+  }
+
+  @Test
+  def testDeleteAllAclOnPrefixedResource(): Unit = {
+    addAcls(aclAuthorizer, Set(allowReadAcl, allowWriteAcl), prefixedResource)
+
+    removeAcls(aclAuthorizer, Set.empty, prefixedResource)
+
+    assertEquals(Set.empty, getAcls(aclAuthorizer))
+  }
+
+  @Test
+  def testAddAclsOnLiteralResource(): Unit = {
+    addAcls(aclAuthorizer, Set(allowReadAcl, allowWriteAcl), resource)
+    addAcls(aclAuthorizer, Set(allowWriteAcl, denyReadAcl), resource)
+
+    assertEquals(Set(allowReadAcl, allowWriteAcl, denyReadAcl), getAcls(aclAuthorizer, resource))
+    assertEquals(Set.empty, getAcls(aclAuthorizer, wildCardResource))
+    assertEquals(Set.empty, getAcls(aclAuthorizer, prefixedResource))
+  }
+
+  @Test
+  def testAddAclsOnWildcardResource(): Unit = {
+    addAcls(aclAuthorizer, Set(allowReadAcl, allowWriteAcl), wildCardResource)
+    addAcls(aclAuthorizer, Set(allowWriteAcl, denyReadAcl), wildCardResource)
+
+    assertEquals(Set(allowReadAcl, allowWriteAcl, denyReadAcl), getAcls(aclAuthorizer, wildCardResource))
+    assertEquals(Set.empty, getAcls(aclAuthorizer, resource))
+    assertEquals(Set.empty, getAcls(aclAuthorizer, prefixedResource))
+  }
+
+  @Test
+  def testAddAclsOnPrefixedResource(): Unit = {
+    addAcls(aclAuthorizer, Set(allowReadAcl, allowWriteAcl), prefixedResource)
+    addAcls(aclAuthorizer, Set(allowWriteAcl, denyReadAcl), prefixedResource)
+
+    assertEquals(Set(allowReadAcl, allowWriteAcl, denyReadAcl), getAcls(aclAuthorizer, prefixedResource))
+    assertEquals(Set.empty, getAcls(aclAuthorizer, wildCardResource))
+    assertEquals(Set.empty, getAcls(aclAuthorizer, resource))
+  }
+
+  @Test
+  def testAuthorizeWithPrefixedResource(): Unit = {
+    addAcls(aclAuthorizer, Set(denyReadAcl), new ResourcePattern(TOPIC, "a_other", LITERAL))
+    addAcls(aclAuthorizer, Set(denyReadAcl), new ResourcePattern(TOPIC, "a_other", PREFIXED))
+    addAcls(aclAuthorizer, Set(denyReadAcl), new ResourcePattern(TOPIC, "foo-" + UUID.randomUUID(), PREFIXED))
+    addAcls(aclAuthorizer, Set(denyReadAcl), new ResourcePattern(TOPIC, "foo-" + UUID.randomUUID(), PREFIXED))
+    addAcls(aclAuthorizer, Set(denyReadAcl), new ResourcePattern(TOPIC, "foo-" + UUID.randomUUID() + "-zzz", PREFIXED))
+    addAcls(aclAuthorizer, Set(denyReadAcl), new ResourcePattern(TOPIC, "fooo-" + UUID.randomUUID(), PREFIXED))
+    addAcls(aclAuthorizer, Set(denyReadAcl), new ResourcePattern(TOPIC, "fo-" + UUID.randomUUID(), PREFIXED))
+    addAcls(aclAuthorizer, Set(denyReadAcl), new ResourcePattern(TOPIC, "fop-" + UUID.randomUUID(), PREFIXED))
+    addAcls(aclAuthorizer, Set(denyReadAcl), new ResourcePattern(TOPIC, "fon-" + UUID.randomUUID(), PREFIXED))
+    addAcls(aclAuthorizer, Set(denyReadAcl), new ResourcePattern(TOPIC, "fon-", PREFIXED))
+    addAcls(aclAuthorizer, Set(denyReadAcl), new ResourcePattern(TOPIC, "z_other", PREFIXED))
+    addAcls(aclAuthorizer, Set(denyReadAcl), new ResourcePattern(TOPIC, "z_other", LITERAL))
+
+    addAcls(aclAuthorizer, Set(allowReadAcl), prefixedResource)
+
+    assertTrue(authorize(aclAuthorizer, requestContext, READ, resource))
+  }
+
+  @Test
+  def testSingleCharacterResourceAcls(): Unit = {
+    addAcls(aclAuthorizer, Set(allowReadAcl), new ResourcePattern(TOPIC, "f", LITERAL))
+    assertTrue(authorize(aclAuthorizer, requestContext, READ, new ResourcePattern(TOPIC, "f", LITERAL)))
+    assertFalse(authorize(aclAuthorizer, requestContext, READ, new ResourcePattern(TOPIC, "foo", LITERAL)))
+
+    addAcls(aclAuthorizer, Set(allowReadAcl), new ResourcePattern(TOPIC, "_", PREFIXED))
+    assertTrue(authorize(aclAuthorizer, requestContext, READ, new ResourcePattern(TOPIC, "_foo", LITERAL)))
+    assertTrue(authorize(aclAuthorizer, requestContext, READ, new ResourcePattern(TOPIC, "_", LITERAL)))
+    assertFalse(authorize(aclAuthorizer, requestContext, READ, new ResourcePattern(TOPIC, "foo_", LITERAL)))
+  }
+
+  @Test
+  def testGetAclsPrincipal(): Unit = {
+    val aclOnSpecificPrincipal = new AccessControlEntry(principal.toString, WildcardHost, WRITE, ALLOW)
+    addAcls(aclAuthorizer, Set(aclOnSpecificPrincipal), resource)
+
+    assertEquals("acl on specific should not be returned for wildcard request",
+      0, getAcls(aclAuthorizer, wildcardPrincipal).size)
+    assertEquals("acl on specific should be returned for specific request",
+      1, getAcls(aclAuthorizer, principal).size)
+    assertEquals("acl on specific should be returned for different principal instance",
+      1, getAcls(aclAuthorizer, new KafkaPrincipal(principal.getPrincipalType, principal.getName)).size)
+
+    removeAcls(aclAuthorizer, Set.empty, resource)
+    val aclOnWildcardPrincipal = new AccessControlEntry(WildcardPrincipal, WildcardHost, WRITE, ALLOW)
+    addAcls(aclAuthorizer, Set(aclOnWildcardPrincipal), resource)
+
+    assertEquals("acl on wildcard should be returned for wildcard request",
+      1, getAcls(aclAuthorizer, wildcardPrincipal).size)
+    assertEquals("acl on wildcard should not be returned for specific request",
+      0, getAcls(aclAuthorizer, principal).size)
+  }
+
+  @Test
+  def testAclsFilter(): Unit = {
+    val resource1 = new ResourcePattern(TOPIC, "foo-" + UUID.randomUUID(), LITERAL)
+    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 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))
+
+    aclAuthorizer.createAcls(requestContext, List(acl1, acl2, acl3, acl4).asJava)
+    assertEquals(Set(acl1, acl2, acl3, acl4), aclAuthorizer.acls(AclBindingFilter.ANY).asScala.toSet)
+    assertEquals(Set(acl1, acl2), aclAuthorizer.acls(new AclBindingFilter(resource1.toFilter, AccessControlEntryFilter.ANY)).asScala.toSet)
+    assertEquals(Set(acl4), aclAuthorizer.acls(new AclBindingFilter(prefixedResource.toFilter, AccessControlEntryFilter.ANY)).asScala.toSet)
+    val matchingFilter = new AclBindingFilter(new ResourcePatternFilter(ResourceType.ANY, resource2.name, MATCH), AccessControlEntryFilter.ANY)
+    assertEquals(Set(acl3, acl4), aclAuthorizer.acls(matchingFilter).asScala.toSet)
+
+    val filters = List(matchingFilter,
+      acl1.toFilter,
+      new AclBindingFilter(resource2.toFilter, AccessControlEntryFilter.ANY),
+      new AclBindingFilter(new ResourcePatternFilter(TOPIC, "baz", PatternType.ANY), AccessControlEntryFilter.ANY))
+    val deleteResults = aclAuthorizer.deleteAcls(requestContext, filters.asJava).asScala.map(_.toCompletableFuture.get)
+    assertEquals(List.empty, deleteResults.filter(_.exception.isPresent))
+    filters.indices.foreach { i =>
+      assertEquals(Set.empty, deleteResults(i).aclBindingDeleteResults.asScala.toSet.filter(_.exception.isPresent))
+    }
+    assertEquals(Set(acl3, acl4), deleteResults(0).aclBindingDeleteResults.asScala.map(_.aclBinding).toSet)
+    assertEquals(Set(acl1), deleteResults(1).aclBindingDeleteResults.asScala.map(_.aclBinding).toSet)
+    assertEquals(Set.empty, deleteResults(2).aclBindingDeleteResults.asScala.map(_.aclBinding).toSet)
+    assertEquals(Set.empty, deleteResults(3).aclBindingDeleteResults.asScala.map(_.aclBinding).toSet)
+  }
+
+  @Test
+  def testThrowsOnAddPrefixedAclIfInterBrokerProtocolVersionTooLow(): Unit = {
+    givenAuthorizerWithProtocolVersion(Option(KAFKA_2_0_IV0))
+    val e = intercept[ApiException] {
+      addAcls(aclAuthorizer, Set(denyReadAcl), new ResourcePattern(TOPIC, "z_other", PREFIXED))
+    }
+    assertTrue(s"Unexpected exception $e", e.getCause.isInstanceOf[UnsupportedVersionException])
+  }
+
+  @Test
+  def testWritesExtendedAclChangeEventIfInterBrokerProtocolNotSet(): Unit = {
+    givenAuthorizerWithProtocolVersion(Option.empty)
+    val resource = new ResourcePattern(TOPIC, "z_other", PREFIXED)
+    val expected = new String(ZkAclStore(PREFIXED).changeStore
+      .createChangeNode(AuthorizerUtils.convertToResource(resource)).bytes, UTF_8)
+
+    addAcls(aclAuthorizer, Set(denyReadAcl), resource)
+
+    val actual = getAclChangeEventAsString(PREFIXED)
+
+    assertEquals(expected, actual)
+  }
+
+  @Test
+  def testWritesExtendedAclChangeEventWhenInterBrokerProtocolAtLeastKafkaV2(): Unit = {
+    givenAuthorizerWithProtocolVersion(Option(KAFKA_2_0_IV1))
+    val resource = new ResourcePattern(TOPIC, "z_other", PREFIXED)
+    val expected = new String(ZkAclStore(PREFIXED).changeStore
+      .createChangeNode(AuthorizerUtils.convertToResource(resource)).bytes, UTF_8)
+
+    addAcls(aclAuthorizer, Set(denyReadAcl), resource)
+
+    val actual = getAclChangeEventAsString(PREFIXED)
+
+    assertEquals(expected, actual)
+  }
+
+  @Test
+  def testWritesLiteralWritesLiteralAclChangeEventWhenInterBrokerProtocolLessThanKafkaV2eralAclChangesForOlderProtocolVersions(): Unit = {
+    givenAuthorizerWithProtocolVersion(Option(KAFKA_2_0_IV0))
+    val resource = new ResourcePattern(TOPIC, "z_other", LITERAL)
+    val expected = new String(ZkAclStore(LITERAL).changeStore
+      .createChangeNode(AuthorizerUtils.convertToResource(resource)).bytes, UTF_8)
+
+    addAcls(aclAuthorizer, Set(denyReadAcl), resource)
+
+    val actual = getAclChangeEventAsString(LITERAL)
+
+    assertEquals(expected, actual)
+  }
+
+  @Test
+  def testWritesLiteralAclChangeEventWhenInterBrokerProtocolIsKafkaV2(): Unit = {
+    givenAuthorizerWithProtocolVersion(Option(KAFKA_2_0_IV1))
+    val resource = new ResourcePattern(TOPIC, "z_other", LITERAL)
+    val expected = new String(ZkAclStore(LITERAL).changeStore
+      .createChangeNode(AuthorizerUtils.convertToResource(resource)).bytes, UTF_8)
+
+    addAcls(aclAuthorizer, Set(denyReadAcl), resource)
+
+    val actual = getAclChangeEventAsString(LITERAL)
+
+    assertEquals(expected, actual)
+  }
+
+  private def givenAuthorizerWithProtocolVersion(protocolVersion: Option[ApiVersion]): Unit = {
+    aclAuthorizer.close()
+
+    val props = TestUtils.createBrokerConfig(0, zkConnect)
+    props.put(AclAuthorizer.SuperUsersProp, superUsers)
+    protocolVersion.foreach(version => props.put(KafkaConfig.InterBrokerProtocolVersionProp, version.toString))
+
+    config = KafkaConfig.fromProps(props)
+
+    aclAuthorizer.configure(config.originals)
+  }
+
+  private def getAclChangeEventAsString(patternType: PatternType) = {
+    val store = ZkAclStore(patternType)
+    val children = zooKeeperClient.handleRequest(GetChildrenRequest(store.changeStore.aclChangePath))
+    children.maybeThrow()
+    assertEquals("Expecting 1 change event", 1, children.children.size)
+
+    val data = zooKeeperClient.handleRequest(GetDataRequest(s"${store.changeStore.aclChangePath}/${children.children.head}"))
+    data.maybeThrow()
+
+    new String(data.data, UTF_8)
+  }
+
+  private def changeAclAndVerify(originalAcls: Set[AccessControlEntry],
+                                 addedAcls: Set[AccessControlEntry],
+                                 removedAcls: Set[AccessControlEntry],
+                                 resource: ResourcePattern = resource): Set[AccessControlEntry] = {
+    var acls = originalAcls
+
+    if(addedAcls.nonEmpty) {
+      addAcls(aclAuthorizer, addedAcls, resource)
+      acls ++= addedAcls
+    }
+
+    if(removedAcls.nonEmpty) {
+      removeAcls(aclAuthorizer, removedAcls, resource)
+      acls --=removedAcls
+    }
+
+    TestUtils.waitAndVerifyAcls(acls, aclAuthorizer, resource)
+
+    acls
+  }
+
+  private def newRequestContext(principal: KafkaPrincipal, clientAddress: InetAddress, apiKey: ApiKeys = ApiKeys.PRODUCE): RequestContext = {
+    val securityProtocol = SecurityProtocol.SASL_PLAINTEXT
+    val header = new RequestHeader(apiKey, 2, "", 1) //ApiKeys apiKey, short version, String clientId, int correlation
+    new RequestContext(header, "", clientAddress, principal, ListenerName.forSecurityProtocol(securityProtocol), securityProtocol)
+  }
+
+  private def authorize(authorizer: AclAuthorizer, requestContext: RequestContext, operation: AclOperation, resource: ResourcePattern): Boolean = {
+    val action = new Action(operation, resource, 1, true, true)
+    authorizer.authorize(requestContext, List(action).asJava).asScala.head == AuthorizationResult.ALLOWED
+  }
+
+  private def addAcls(authorizer: AclAuthorizer, aces: Set[AccessControlEntry], resourcePattern: ResourcePattern): Unit = {
+    val bindings = aces.map { ace => new AclBinding(resourcePattern, ace) }
+    authorizer.createAcls(requestContext, bindings.toList.asJava).asScala
+      .map(_.toCompletableFuture.get)
+      .foreach { result => result.exception.asScala.foreach { e => throw e } }
+  }
+
+  private def removeAcls(authorizer: AclAuthorizer, aces: Set[AccessControlEntry], resourcePattern: ResourcePattern): Boolean = {
+    val bindings = if (aces.isEmpty)
+      Set(new AclBindingFilter(resourcePattern.toFilter, AccessControlEntryFilter.ANY) )
+    else
+      aces.map { ace => new AclBinding(resourcePattern, ace).toFilter }
+    authorizer.deleteAcls(requestContext, bindings.toList.asJava).asScala
+      .map(_.toCompletableFuture.get)
+      .forall { result =>
+        result.exception.asScala.foreach { e => throw e }
+        result.aclBindingDeleteResults.asScala.foreach { r =>
+          r.exception.asScala.foreach { e => throw e }
+        }
+        result.aclBindingDeleteResults.asScala.exists(_.exception.asScala.isEmpty)
+      }
+  }
+
+  private def getAcls(authorizer: AclAuthorizer, resourcePattern: ResourcePattern): Set[AccessControlEntry] = {
+    val acls = authorizer.acls(new AclBindingFilter(resourcePattern.toFilter, AccessControlEntryFilter.ANY)).asScala.toSet
+    acls.map(_.entry)
+  }
+
+  private def getAcls(authorizer: AclAuthorizer, principal: KafkaPrincipal): Set[AclBinding] = {
+    val filter = new AclBindingFilter(ResourcePatternFilter.ANY,
+      new AccessControlEntryFilter(principal.toString, null, AclOperation.ANY, AclPermissionType.ANY))
+    authorizer.acls(filter).asScala.toSet
+  }
+
+  private def getAcls(authorizer: AclAuthorizer): Set[AclBinding] = {
+    authorizer.acls(AclBindingFilter.ANY).asScala.toSet
+  }
+
+  private def validOp(op: AclOperation): Boolean = {
+    op != AclOperation.ANY && op != AclOperation.UNKNOWN
+  }
+}
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 ed82f5ed4464c..90a3b93a941af 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
@@ -23,21 +23,28 @@ import java.util.{Base64, Properties}
 
 import kafka.network.RequestChannel.Session
 import kafka.security.auth.Acl.WildCardHost
-import kafka.security.auth._
+import kafka.security.authorizer.{AclAuthorizer, AuthorizerUtils}
 import kafka.server.{CreateTokenResult, Defaults, DelegationTokenManager, KafkaConfig}
 import kafka.utils.TestUtils
 import kafka.zk.{KafkaZkClient, ZooKeeperTestHarness}
+import org.apache.kafka.common.acl.{AccessControlEntry, AclBinding, AclOperation}
+import org.apache.kafka.common.acl.AclOperation._
+import org.apache.kafka.common.acl.AclPermissionType.ALLOW
 import org.apache.kafka.common.protocol.Errors
 import org.apache.kafka.common.resource.PatternType.LITERAL
+import org.apache.kafka.common.resource.ResourcePattern
+import org.apache.kafka.common.resource.ResourceType.DELEGATION_TOKEN
 import org.apache.kafka.common.security.auth.KafkaPrincipal
 import org.apache.kafka.common.security.scram.internals.ScramMechanism
 import org.apache.kafka.common.security.token.delegation.internals.DelegationTokenCache
 import org.apache.kafka.common.security.token.delegation.{DelegationToken, TokenInformation}
 import org.apache.kafka.common.utils.{MockTime, SecurityUtils, Time}
+import org.apache.kafka.server.authorizer._
 import org.junit.Assert._
 import org.junit.{After, Before, Test}
 
 import scala.collection.JavaConverters._
+import scala.compat.java8.OptionConverters._
 import scala.collection.mutable.Buffer
 
 class DelegationTokenManagerTest extends ZooKeeperTestHarness  {
@@ -58,7 +65,7 @@ class DelegationTokenManagerTest extends ZooKeeperTestHarness  {
   var expiryTimeStamp: Long = 0
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     super.setUp()
     props = TestUtils.createBrokerConfig(0, zkConnect, enableToken = true)
     props.put(KafkaConfig.SaslEnabledMechanismsProp, ScramMechanism.mechanismNames().asScala.mkString(","))
@@ -228,8 +235,8 @@ class DelegationTokenManagerTest extends ZooKeeperTestHarness  {
     val renewer3 = SecurityUtils.parseKafkaPrincipal("User:renewer3")
     val renewer4 = SecurityUtils.parseKafkaPrincipal("User:renewer4")
 
-    val simpleAclAuthorizer = new SimpleAclAuthorizer
-    simpleAclAuthorizer.configure(config.originals)
+    val aclAuthorizer = new AclAuthorizer
+    aclAuthorizer.configure(config.originals)
 
     var hostSession = new Session(owner1, InetAddress.getByName("192.168.1.1"))
 
@@ -250,61 +257,71 @@ class DelegationTokenManagerTest extends ZooKeeperTestHarness  {
     assert(tokenManager.getAllTokenInformation().size == 4 )
 
     //get tokens non-exiting owner
-    var  tokens = getTokens(tokenManager, simpleAclAuthorizer, hostSession, owner1, List(SecurityUtils.parseKafkaPrincipal("User:unknown")))
+    var  tokens = getTokens(tokenManager, aclAuthorizer, hostSession, owner1, List(SecurityUtils.parseKafkaPrincipal("User:unknown")))
     assert(tokens.size == 0)
 
     //get all tokens for  empty owner list
-    tokens = getTokens(tokenManager, simpleAclAuthorizer, hostSession, owner1, List())
+    tokens = getTokens(tokenManager, aclAuthorizer, hostSession, owner1, List())
     assert(tokens.size == 0)
 
     //get all tokens for owner1
-    tokens = getTokens(tokenManager, simpleAclAuthorizer, hostSession, owner1, List(owner1))
+    tokens = getTokens(tokenManager, aclAuthorizer, hostSession, owner1, List(owner1))
     assert(tokens.size == 2)
 
     //get all tokens for owner1
-    tokens = getTokens(tokenManager, simpleAclAuthorizer, hostSession, owner1, null)
+    tokens = getTokens(tokenManager, aclAuthorizer, hostSession, owner1, null)
     assert(tokens.size == 2)
 
     //get all tokens for unknown owner
-    tokens = getTokens(tokenManager, simpleAclAuthorizer, hostSession, SecurityUtils.parseKafkaPrincipal("User:unknown"), null)
+    tokens = getTokens(tokenManager, aclAuthorizer, hostSession, SecurityUtils.parseKafkaPrincipal("User:unknown"), null)
     assert(tokens.size == 0)
 
     //get all tokens for multiple owners (owner1, renewer4) and without permission for renewer4
-    tokens = getTokens(tokenManager, simpleAclAuthorizer, hostSession, owner1, List(owner1, renewer4))
+    tokens = getTokens(tokenManager, aclAuthorizer, hostSession, owner1, List(owner1, renewer4))
     assert(tokens.size == 2)
 
+    def createAcl(aclBinding: AclBinding): Unit = {
+      val result = aclAuthorizer.createAcls(null, List(aclBinding).asJava).get(0).toCompletableFuture.get
+      result.exception.asScala.foreach { e => throw e }
+    }
+
     //get all tokens for multiple owners (owner1, renewer4) and with permission
-    var acl = new Acl(owner1, Allow, WildCardHost, Describe)
-    simpleAclAuthorizer.addAcls(Set(acl), Resource(kafka.security.auth.DelegationToken, tokenId3, LITERAL))
-    tokens = getTokens(tokenManager, simpleAclAuthorizer, hostSession, owner1, List(owner1, renewer4))
+    createAcl(new AclBinding(new ResourcePattern(DELEGATION_TOKEN, tokenId3, LITERAL),
+      new AccessControlEntry(owner1.toString, WildCardHost, DESCRIBE, ALLOW)))
+    tokens = getTokens(tokenManager, aclAuthorizer, hostSession, owner1, List(owner1, renewer4))
     assert(tokens.size == 3)
 
     //get all tokens for renewer4 which is a renewer principal for some tokens
-    tokens = getTokens(tokenManager, simpleAclAuthorizer, hostSession,  renewer4, List(renewer4))
+    tokens = getTokens(tokenManager, aclAuthorizer, hostSession,  renewer4, List(renewer4))
     assert(tokens.size == 2)
 
     //get all tokens for multiple owners (renewer2, renewer3) which are token renewers principals and without permissions for renewer3
-    tokens = getTokens(tokenManager, simpleAclAuthorizer, hostSession,  renewer2, List(renewer2, renewer3))
+    tokens = getTokens(tokenManager, aclAuthorizer, hostSession,  renewer2, List(renewer2, renewer3))
     assert(tokens.size == 1)
 
     //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"))
-    acl = new Acl(renewer2, Allow, WildCardHost, Describe)
-    simpleAclAuthorizer.addAcls(Set(acl), Resource(kafka.security.auth.DelegationToken, tokenId2, LITERAL))
-    tokens = getTokens(tokenManager, simpleAclAuthorizer, hostSession,  renewer2, List(renewer2, renewer3))
+    hostSession = Session(renewer2, InetAddress.getByName("192.168.1.1"))
+    createAcl(new AclBinding(new ResourcePattern(DELEGATION_TOKEN, tokenId2, LITERAL),
+      new AccessControlEntry(renewer2.toString, WildCardHost, DESCRIBE, ALLOW)))
+    tokens = getTokens(tokenManager, aclAuthorizer, hostSession,  renewer2, List(renewer2, renewer3))
     assert(tokens.size == 2)
 
-    simpleAclAuthorizer.close()
+    aclAuthorizer.close()
   }
 
-  private def getTokens(tokenManager: DelegationTokenManager, simpleAclAuthorizer: SimpleAclAuthorizer, hostSession: Session,
+  private def getTokens(tokenManager: DelegationTokenManager, aclAuthorizer: AclAuthorizer, hostSession: Session,
                         requestPrincipal: KafkaPrincipal, requestedOwners: List[KafkaPrincipal]): List[DelegationToken] = {
 
     if (requestedOwners != null && requestedOwners.isEmpty) {
       List()
     }
     else {
-      def authorizeToken(tokenId: String) = simpleAclAuthorizer.authorize(hostSession, Describe, Resource(kafka.security.auth.DelegationToken, tokenId, LITERAL))
+      def authorizeToken(tokenId: String) = {
+        val requestContext = AuthorizerUtils.sessionToRequestContext(hostSession)
+        val action = new Action(AclOperation.DESCRIBE,
+          new ResourcePattern(DELEGATION_TOKEN, tokenId, LITERAL), 1, true, true)
+        aclAuthorizer.authorize(requestContext, List(action).asJava).asScala.head == AuthorizationResult.ALLOWED
+      }
       def eligible(token: TokenInformation) = DelegationTokenManager.filterToken(requestPrincipal, Option(requestedOwners), token, authorizeToken)
       tokenManager.getTokens(eligible)
     }
diff --git a/core/src/test/scala/unit/kafka/server/AbstractCreateTopicsRequestTest.scala b/core/src/test/scala/unit/kafka/server/AbstractCreateTopicsRequestTest.scala
index d6262b3da34e2..d0dbd8e5ab9c4 100644
--- a/core/src/test/scala/unit/kafka/server/AbstractCreateTopicsRequestTest.scala
+++ b/core/src/test/scala/unit/kafka/server/AbstractCreateTopicsRequestTest.scala
@@ -33,7 +33,7 @@ import scala.collection.JavaConverters._
 
 class AbstractCreateTopicsRequestTest extends BaseRequestTest {
 
-  override def propertyOverrides(properties: Properties): Unit =
+  override def brokerPropertyOverrides(properties: Properties): Unit =
     properties.put(KafkaConfig.AutoCreateTopicsEnableProp, false.toString)
 
   def topicsReq(topics: Seq[CreatableTopic],
@@ -41,7 +41,7 @@ class AbstractCreateTopicsRequestTest extends BaseRequestTest {
                 validateOnly: Boolean = false) = {
     val req = new CreateTopicsRequestData()
     req.setTimeoutMs(timeout)
-    req.setTopics(new CreatableTopicSet(topics.asJava.iterator()))
+    req.setTopics(new CreatableTopicCollection(topics.asJava.iterator()))
     req.setValidateOnly(validateOnly)
     new CreateTopicsRequest.Builder(req).build()
   }
@@ -68,7 +68,7 @@ class AbstractCreateTopicsRequestTest extends BaseRequestTest {
       topic.setReplicationFactor(1.toShort)
     }
     if (config != null) {
-      val effectiveConfigs = new CreateableTopicConfigSet()
+      val effectiveConfigs = new CreateableTopicConfigCollection()
       config.foreach {
         case (name, value) => {
           effectiveConfigs.add(new CreateableTopicConfig().setName(name).setValue(value))
@@ -77,7 +77,7 @@ class AbstractCreateTopicsRequestTest extends BaseRequestTest {
       topic.setConfigs(effectiveConfigs)
     }
     if (assignment != null) {
-      val effectiveAssignments = new CreatableReplicaAssignmentSet()
+      val effectiveAssignments = new CreatableReplicaAssignmentCollection()
       assignment.foreach {
         case (partitionIndex, brokerIdList) => {
           val effectiveAssignment = new CreatableReplicaAssignment()
@@ -124,8 +124,18 @@ class AbstractCreateTopicsRequestTest extends BaseRequestTest {
         else {
           assertNotNull("The topic should be created", metadataForTopic)
           assertEquals(Errors.NONE, metadataForTopic.error)
-          assertEquals("The topic should have the correct number of partitions", partitions, metadataForTopic.partitionMetadata.size)
-          assertEquals("The topic should have the correct replication factor", replication, metadataForTopic.partitionMetadata.asScala.head.replicas.size)
+          if (partitions == -1) {
+            assertEquals("The topic should have the default number of partitions", configs.head.numPartitions, metadataForTopic.partitionMetadata.size)
+          } else {
+            assertEquals("The topic should have the correct number of partitions", partitions, metadataForTopic.partitionMetadata.size)
+          }
+
+          if (replication == -1) {
+            assertEquals("The topic should have the default replication factor",
+              configs.head.defaultReplicationFactor, metadataForTopic.partitionMetadata.asScala.head.replicas.size)
+          } else {
+            assertEquals("The topic should have the correct replication factor", replication, metadataForTopic.partitionMetadata.asScala.head.replicas.size)
+          }
         }
       }
 
diff --git a/core/src/test/scala/unit/kafka/server/AbstractFetcherManagerTest.scala b/core/src/test/scala/unit/kafka/server/AbstractFetcherManagerTest.scala
index 0a4d7c11df36b..f940b836ba9b7 100644
--- a/core/src/test/scala/unit/kafka/server/AbstractFetcherManagerTest.scala
+++ b/core/src/test/scala/unit/kafka/server/AbstractFetcherManagerTest.scala
@@ -16,14 +16,29 @@
  */
 package kafka.server
 
+import com.yammer.metrics.Metrics
+import com.yammer.metrics.core.Gauge
 import kafka.cluster.BrokerEndPoint
+import kafka.utils.TestUtils
 import org.apache.kafka.common.TopicPartition
 import org.easymock.EasyMock
-import org.junit.Test
+import org.junit.{Before, Test}
 import org.junit.Assert._
 
+import scala.collection.JavaConverters._
+
 class AbstractFetcherManagerTest {
 
+  @Before
+  def cleanMetricRegistry(): Unit = {
+    TestUtils.clearYammerMetrics()
+  }
+
+  private def getMetricValue(name: String): Any = {
+    Metrics.defaultRegistry.allMetrics.asScala.filterKeys(_.getName == name).values.headOption.get.
+      asInstanceOf[Gauge[Int]].value()
+  }
+
   @Test
   def testAddAndRemovePartition(): Unit = {
     val fetcher: AbstractFetcherThread = EasyMock.mock(classOf[AbstractFetcherThread])
@@ -43,6 +58,7 @@ class AbstractFetcherManagerTest {
 
     EasyMock.expect(fetcher.start())
     EasyMock.expect(fetcher.addPartitions(Map(tp -> OffsetAndEpoch(fetchOffset, leaderEpoch))))
+        .andReturn(Set(tp))
     EasyMock.expect(fetcher.fetchState(tp))
       .andReturn(Some(PartitionFetchState(fetchOffset, leaderEpoch, Truncating)))
     EasyMock.expect(fetcher.removePartitions(Set(tp)))
@@ -58,4 +74,63 @@ class AbstractFetcherManagerTest {
     EasyMock.verify(fetcher)
   }
 
+  @Test
+  def testMetricFailedPartitionCount(): Unit = {
+    val fetcher: AbstractFetcherThread = EasyMock.mock(classOf[AbstractFetcherThread])
+    val fetcherManager = new AbstractFetcherManager[AbstractFetcherThread]("fetcher-manager", "fetcher-manager", 2) {
+      override def createFetcherThread(fetcherId: Int, sourceBroker: BrokerEndPoint): AbstractFetcherThread = {
+        fetcher
+      }
+    }
+
+    val tp = new TopicPartition("topic", 0)
+    val metricName = "FailedPartitionsCount"
+
+    // initial value for failed partition count
+    assertEquals(0, getMetricValue(metricName))
+
+    // partition marked as failed increments the count for failed partitions
+    fetcherManager.failedPartitions.add(tp)
+    assertEquals(1, getMetricValue(metricName))
+
+    // removing fetcher for the partition would remove the partition from set of failed partitions and decrement the
+    // count for failed partitions
+    fetcherManager.removeFetcherForPartitions(Set(tp))
+    assertEquals(0, getMetricValue(metricName))
+  }
+  @Test
+  def testDeadThreadCountMetric(): Unit = {
+    val fetcher: AbstractFetcherThread = EasyMock.mock(classOf[AbstractFetcherThread])
+    val fetcherManager = new AbstractFetcherManager[AbstractFetcherThread]("fetcher-manager", "fetcher-manager", 2) {
+      override def createFetcherThread(fetcherId: Int, sourceBroker: BrokerEndPoint): AbstractFetcherThread = {
+        fetcher
+      }
+    }
+
+    val fetchOffset = 10L
+    val leaderEpoch = 15
+    val tp = new TopicPartition("topic", 0)
+    val initialFetchState = InitialFetchState(
+      leader = new BrokerEndPoint(0, "localhost", 9092),
+      currentLeaderEpoch = leaderEpoch,
+      initOffset = fetchOffset)
+
+    EasyMock.expect(fetcher.start())
+    EasyMock.expect(fetcher.addPartitions(Map(tp -> OffsetAndEpoch(fetchOffset, leaderEpoch))))
+        .andReturn(Set(tp))
+    EasyMock.expect(fetcher.isThreadFailed).andReturn(true)
+    EasyMock.replay(fetcher)
+
+    fetcherManager.addFetcherForPartitions(Map(tp -> initialFetchState))
+
+    assertEquals(1, fetcherManager.deadThreadCount)
+    EasyMock.verify(fetcher)
+
+    EasyMock.reset(fetcher)
+    EasyMock.expect(fetcher.isThreadFailed).andReturn(false)
+    EasyMock.replay(fetcher)
+
+    assertEquals(0, fetcherManager.deadThreadCount)
+    EasyMock.verify(fetcher)
+  }
 }
diff --git a/core/src/test/scala/unit/kafka/server/AbstractFetcherThreadTest.scala b/core/src/test/scala/unit/kafka/server/AbstractFetcherThreadTest.scala
index 1fc079dde5387..7958fee626eb6 100644
--- a/core/src/test/scala/unit/kafka/server/AbstractFetcherThreadTest.scala
+++ b/core/src/test/scala/unit/kafka/server/AbstractFetcherThreadTest.scala
@@ -27,6 +27,7 @@ import kafka.log.LogAppendInfo
 import kafka.message.NoCompressionCodec
 import kafka.server.AbstractFetcherThread.ResultWithPartitions
 import kafka.utils.TestUtils
+import org.apache.kafka.common.KafkaException
 import org.apache.kafka.common.TopicPartition
 import org.apache.kafka.common.errors.{FencedLeaderEpochException, UnknownLeaderEpochException}
 import org.apache.kafka.common.protocol.{ApiKeys, Errors}
@@ -45,10 +46,13 @@ import scala.collection.mutable.ArrayBuffer
 
 class AbstractFetcherThreadTest {
 
+  private val partition1 = new TopicPartition("topic1", 0)
+  private val partition2 = new TopicPartition("topic2", 0)
+  private val failedPartitions = new FailedPartitions
+
   @Before
   def cleanMetricRegistry(): Unit = {
-    for (metricName <- Metrics.defaultRegistry().allMetrics().keySet().asScala)
-      Metrics.defaultRegistry().removeMetric(metricName)
+    TestUtils.clearYammerMetrics()
   }
 
   private def allMetricsNames: Set[String] = Metrics.defaultRegistry().allMetrics().asScala.keySet.map(_.getName)
@@ -76,7 +80,8 @@ class AbstractFetcherThreadTest {
 
     // wait until all fetcher metrics are present
     TestUtils.waitUntilTrue(() =>
-      allMetricsNames == Set(FetcherMetrics.BytesPerSec, FetcherMetrics.RequestsPerSec, FetcherMetrics.ConsumerLag),
+      allMetricsNames == Set(FetcherMetrics.BytesPerSec, FetcherMetrics.RequestFailuresPerSec, FetcherMetrics.RequestsPerSec,
+                             FetcherMetrics.ConsumerLag),
       "Failed waiting for all fetcher metrics to be registered")
 
     fetcher.shutdown()
@@ -147,8 +152,9 @@ class AbstractFetcherThreadTest {
     assertEquals(0L, replicaState.logEndOffset)
     assertEquals(0L, replicaState.highWatermark)
 
-    // After fencing, the fetcher should remove the partition from tracking
+    // After fencing, the fetcher should remove the partition from tracking and mark as failed
     assertTrue(fetcher.fetchState(partition).isEmpty)
+    assertTrue(failedPartitions.contains(partition))
   }
 
   @Test
@@ -176,8 +182,9 @@ class AbstractFetcherThreadTest {
 
     fetcher.doWork()
 
-    // After fencing, the fetcher should remove the partition from tracking
+    // After fencing, the fetcher should remove the partition from tracking and mark as failed
     assertTrue(fetcher.fetchState(partition).isEmpty)
+    assertTrue(failedPartitions.contains(partition))
   }
 
   @Test
@@ -480,11 +487,12 @@ class AbstractFetcherThreadTest {
     val leaderState = MockFetcherThread.PartitionState(leaderLog, leaderEpoch = 4, highWatermark = 2L)
     fetcher.setLeaderState(partition, leaderState)
 
-    // After the out of range error, we get a fenced error and remove the partition
+    // After the out of range error, we get a fenced error and remove the partition and mark as failed
     fetcher.doWork()
     assertEquals(0, replicaState.logEndOffset)
     assertTrue(fetchedEarliestOffset)
     assertTrue(fetcher.fetchState(partition).isEmpty)
+    assertTrue(failedPartitions.contains(partition))
   }
 
   @Test
@@ -722,6 +730,67 @@ class AbstractFetcherThreadTest {
     }
   }
 
+  @Test
+  def testFetcherThreadHandlingPartitionFailureDuringAppending(): Unit = {
+    val fetcherForAppend = new MockFetcherThread {
+      override def processPartitionData(topicPartition: TopicPartition, fetchOffset: Long, partitionData: FetchData): Option[LogAppendInfo] = {
+        if (topicPartition == partition1) {
+          throw new KafkaException()
+        } else {
+          super.processPartitionData(topicPartition, fetchOffset, partitionData)
+        }
+      }
+    }
+    verifyFetcherThreadHandlingPartitionFailure(fetcherForAppend)
+  }
+
+  @Test
+  def testFetcherThreadHandlingPartitionFailureDuringTruncation(): Unit = {
+    val fetcherForTruncation = new MockFetcherThread {
+      override def truncate(topicPartition: TopicPartition, truncationState: OffsetTruncationState): Unit = {
+        if(topicPartition == partition1)
+          throw new Exception()
+        else {
+          super.truncate(topicPartition: TopicPartition, truncationState: OffsetTruncationState)
+        }
+      }
+    }
+    verifyFetcherThreadHandlingPartitionFailure(fetcherForTruncation)
+  }
+
+  private def verifyFetcherThreadHandlingPartitionFailure(fetcher: MockFetcherThread): Unit = {
+
+    fetcher.setReplicaState(partition1, MockFetcherThread.PartitionState(leaderEpoch = 0))
+    fetcher.addPartitions(Map(partition1 -> offsetAndEpoch(0L, leaderEpoch = 0)))
+    fetcher.setLeaderState(partition1, MockFetcherThread.PartitionState(leaderEpoch = 0))
+
+    fetcher.setReplicaState(partition2, MockFetcherThread.PartitionState(leaderEpoch = 0))
+    fetcher.addPartitions(Map(partition2 -> offsetAndEpoch(0L, leaderEpoch = 0)))
+    fetcher.setLeaderState(partition2, MockFetcherThread.PartitionState(leaderEpoch = 0))
+
+    // processing data fails for partition1
+    fetcher.doWork()
+
+    // partition1 marked as failed
+    assertTrue(failedPartitions.contains(partition1))
+    assertEquals(None, fetcher.fetchState(partition1))
+
+    // make sure the fetcher continues to work with rest of the partitions
+    fetcher.doWork()
+    assertEquals(Some(Fetching), fetcher.fetchState(partition2).map(_.state))
+    assertFalse(failedPartitions.contains(partition2))
+
+    // simulate a leader change
+    fetcher.removePartitions(Set(partition1))
+    failedPartitions.removeAll(Set(partition1))
+    fetcher.addPartitions(Map(partition1 -> offsetAndEpoch(0L, leaderEpoch = 1)))
+
+    // partition1 added back
+    assertEquals(Some(Truncating), fetcher.fetchState(partition1).map(_.state))
+    assertFalse(failedPartitions.contains(partition1))
+
+  }
+
   object MockFetcherThread {
     class PartitionState(var log: mutable.Buffer[RecordBatch],
                          var leaderEpoch: Int,
@@ -745,7 +814,8 @@ class AbstractFetcherThreadTest {
   class MockFetcherThread(val replicaId: Int = 0, val leaderId: Int = 1)
     extends AbstractFetcherThread("mock-fetcher",
       clientId = "mock-fetcher",
-      sourceBroker = new BrokerEndPoint(leaderId, host = "localhost", port = Random.nextInt())) {
+      sourceBroker = new BrokerEndPoint(leaderId, host = "localhost", port = Random.nextInt()),
+      failedPartitions) {
 
     import MockFetcherThread.PartitionState
 
@@ -812,7 +882,8 @@ class AbstractFetcherThreadTest {
         shallowCount = batches.size,
         validBytes = partitionData.records.sizeInBytes,
         offsetsMonotonic = true,
-        lastOffsetOfFirstBatch = batches.headOption.map(_.lastOffset).getOrElse(-1)))
+        lastOffsetOfFirstBatch = batches.headOption.map(_.lastOffset).getOrElse(-1),
+        recompressedBatchCount = 0))
     }
 
     override def truncate(topicPartition: TopicPartition, truncationState: OffsetTruncationState): Unit = {
diff --git a/core/src/test/scala/unit/kafka/server/AddPartitionsToTxnRequestTest.scala b/core/src/test/scala/unit/kafka/server/AddPartitionsToTxnRequestTest.scala
index 9071f95e43856..d15409ced66f1 100644
--- a/core/src/test/scala/unit/kafka/server/AddPartitionsToTxnRequestTest.scala
+++ b/core/src/test/scala/unit/kafka/server/AddPartitionsToTxnRequestTest.scala
@@ -31,7 +31,7 @@ class AddPartitionsToTxnRequestTest extends BaseRequestTest {
   private val topic1 = "foobartopic"
   val numPartitions = 3
 
-  override def propertyOverrides(properties: Properties): Unit =
+  override def brokerPropertyOverrides(properties: Properties): Unit =
     properties.put(KafkaConfig.AutoCreateTopicsEnableProp, false.toString)
 
   @Before
diff --git a/core/src/test/scala/unit/kafka/server/AdvertiseBrokerTest.scala b/core/src/test/scala/unit/kafka/server/AdvertiseBrokerTest.scala
index 75038bf3bda0f..06004bcad56c6 100755
--- a/core/src/test/scala/unit/kafka/server/AdvertiseBrokerTest.scala
+++ b/core/src/test/scala/unit/kafka/server/AdvertiseBrokerTest.scala
@@ -31,7 +31,7 @@ class AdvertiseBrokerTest extends ZooKeeperTestHarness {
   val brokerId = 0
 
   @After
-  override def tearDown() {
+  override def tearDown(): Unit = {
     TestUtils.shutdownServers(servers)
     super.tearDown()
   }
diff --git a/core/src/test/scala/unit/kafka/server/AlterReplicaLogDirsRequestTest.scala b/core/src/test/scala/unit/kafka/server/AlterReplicaLogDirsRequestTest.scala
index 28ed81d060b28..eda0e5584c969 100644
--- a/core/src/test/scala/unit/kafka/server/AlterReplicaLogDirsRequestTest.scala
+++ b/core/src/test/scala/unit/kafka/server/AlterReplicaLogDirsRequestTest.scala
@@ -33,12 +33,12 @@ import scala.util.Random
 
 class AlterReplicaLogDirsRequestTest extends BaseRequestTest {
   override val logDirCount = 5
-  override val numBrokers = 1
+  override val brokerCount = 1
 
   val topic = "topic"
 
   @Test
-  def testAlterReplicaLogDirsRequest() {
+  def testAlterReplicaLogDirsRequest(): Unit = {
     val partitionNum = 5
 
     // Alter replica dir before topic creation
diff --git a/core/src/test/scala/unit/kafka/server/ApiVersionsRequestTest.scala b/core/src/test/scala/unit/kafka/server/ApiVersionsRequestTest.scala
index 83f7111bb2bbe..e915fcd357959 100644
--- a/core/src/test/scala/unit/kafka/server/ApiVersionsRequestTest.scala
+++ b/core/src/test/scala/unit/kafka/server/ApiVersionsRequestTest.scala
@@ -17,8 +17,9 @@
 
 package kafka.server
 
+import org.apache.kafka.common.message.ApiVersionsRequestData
+import org.apache.kafka.common.message.ApiVersionsResponseData.ApiVersionsResponseKey
 import org.apache.kafka.common.protocol.{ApiKeys, Errors}
-import org.apache.kafka.common.requests.ApiVersionsResponse.ApiVersion
 import org.apache.kafka.common.requests.{ApiVersionsRequest, ApiVersionsResponse}
 import org.junit.Assert._
 import org.junit.Test
@@ -26,9 +27,9 @@ import org.junit.Test
 import scala.collection.JavaConverters._
 
 object ApiVersionsRequestTest {
-  def validateApiVersionsResponse(apiVersionsResponse: ApiVersionsResponse) {
-    assertEquals("API keys in ApiVersionsResponse must match API keys supported by broker.", ApiKeys.values.length, apiVersionsResponse.apiVersions.size)
-    for (expectedApiVersion: ApiVersion <- ApiVersionsResponse.defaultApiVersionsResponse().apiVersions.asScala) {
+  def validateApiVersionsResponse(apiVersionsResponse: ApiVersionsResponse): Unit = {
+    assertEquals("API keys in ApiVersionsResponse must match API keys supported by broker.", ApiKeys.values.length, apiVersionsResponse.data.apiKeys().size())
+    for (expectedApiVersion: ApiVersionsResponseKey <- ApiVersionsResponse.DEFAULT_API_VERSIONS_RESPONSE.data.apiKeys().asScala) {
       val actualApiVersion = apiVersionsResponse.apiVersion(expectedApiVersion.apiKey)
       assertNotNull(s"API key ${actualApiVersion.apiKey} is supported by broker, but not received in ApiVersionsResponse.", actualApiVersion)
       assertEquals("API key must be supported by the broker.", expectedApiVersion.apiKey, actualApiVersion.apiKey)
@@ -40,22 +41,44 @@ object ApiVersionsRequestTest {
 
 class ApiVersionsRequestTest extends BaseRequestTest {
 
-  override def numBrokers: Int = 1
+  override def brokerCount: Int = 1
 
   @Test
-  def testApiVersionsRequest() {
-    val apiVersionsResponse = sendApiVersionsRequest(new ApiVersionsRequest.Builder().build())
+  def testApiVersionsRequest(): Unit = {
+    val request = new ApiVersionsRequest.Builder().build()
+    val apiVersionsResponse = sendApiVersionsRequest(request, None, request.version)
     ApiVersionsRequestTest.validateApiVersionsResponse(apiVersionsResponse)
   }
 
   @Test
-  def testApiVersionsRequestWithUnsupportedVersion() {
-    val apiVersionsRequest = new ApiVersionsRequest(0)
+  def testApiVersionsRequestWithUnsupportedVersion(): Unit = {
+    val apiVersionsRequest = new ApiVersionsRequest.Builder().build()
     val apiVersionsResponse = sendApiVersionsRequest(apiVersionsRequest, Some(Short.MaxValue), 0)
-    assertEquals(Errors.UNSUPPORTED_VERSION, apiVersionsResponse.error)
+    assertEquals(Errors.UNSUPPORTED_VERSION.code(), apiVersionsResponse.data.errorCode())
+    assertFalse(apiVersionsResponse.data.apiKeys().isEmpty)
+    val apiVersion = apiVersionsResponse.data.apiKeys().find(ApiKeys.API_VERSIONS.id)
+    assertEquals(ApiKeys.API_VERSIONS.id, apiVersion.apiKey())
+    assertEquals(ApiKeys.API_VERSIONS.oldestVersion(), apiVersion.minVersion())
+    assertEquals(ApiKeys.API_VERSIONS.latestVersion(), apiVersion.maxVersion())
   }
 
-  private def sendApiVersionsRequest(request: ApiVersionsRequest, apiVersion: Option[Short] = None, responseVersion: Short = 1): ApiVersionsResponse = {
+  @Test
+  def testApiVersionsRequestValidationV0(): Unit = {
+    val apiVersionsRequest = new ApiVersionsRequest.Builder().build( 0.asInstanceOf[Short])
+    val apiVersionsResponse = sendApiVersionsRequest(apiVersionsRequest, Some(0.asInstanceOf[Short]), 0)
+    ApiVersionsRequestTest.validateApiVersionsResponse(apiVersionsResponse)
+  }
+
+  @Test
+  def testApiVersionsRequestValidationV3(): Unit = {
+    // Invalid request because Name and Version are empty by default
+    val apiVersionsRequest = new ApiVersionsRequest(new ApiVersionsRequestData(), 3.asInstanceOf[Short])
+    val apiVersionsResponse = sendApiVersionsRequest(apiVersionsRequest, Some(3.asInstanceOf[Short]), 3)
+    assertEquals(Errors.INVALID_REQUEST.code(), apiVersionsResponse.data.errorCode())
+  }
+
+  private def sendApiVersionsRequest(request: ApiVersionsRequest, apiVersion: Option[Short] = None,
+                                     responseVersion: Short): ApiVersionsResponse = {
     val response = connectAndSend(request, ApiKeys.API_VERSIONS, apiVersion = apiVersion)
     ApiVersionsResponse.parse(response, responseVersion)
   }
diff --git a/core/src/test/scala/unit/kafka/server/BaseRequestTest.scala b/core/src/test/scala/unit/kafka/server/BaseRequestTest.scala
index 09ffe4fe83c24..2cc355395d4a6 100644
--- a/core/src/test/scala/unit/kafka/server/BaseRequestTest.scala
+++ b/core/src/test/scala/unit/kafka/server/BaseRequestTest.scala
@@ -22,32 +22,30 @@ import java.net.Socket
 import java.nio.ByteBuffer
 import java.util.Properties
 
+import scala.collection.Seq
+
 import kafka.api.IntegrationTestHarness
 import kafka.network.SocketServer
-import kafka.utils._
 import org.apache.kafka.common.network.ListenerName
 import org.apache.kafka.common.protocol.types.Struct
 import org.apache.kafka.common.protocol.ApiKeys
-import org.apache.kafka.common.requests.{AbstractRequest, AbstractRequestResponse, RequestHeader, ResponseHeader}
+import org.apache.kafka.common.requests.{AbstractRequest, RequestHeader, RequestUtils, ResponseHeader}
 import org.apache.kafka.common.security.auth.SecurityProtocol
 
 abstract class BaseRequestTest extends IntegrationTestHarness {
-  override val serverCount: Int = numBrokers
   private var correlationId = 0
 
   // If required, set number of brokers
-  protected def numBrokers: Int = 3
+  override def brokerCount: Int = 3
 
   // If required, override properties by mutating the passed Properties object
-  protected def propertyOverrides(properties: Properties) {}
+  protected def brokerPropertyOverrides(properties: Properties): Unit = {}
 
-  override def generateConfigs = {
-    val props = TestUtils.createBrokerConfigs(numBrokers, zkConnect,
-      enableControlledShutdown = false,
-      interBrokerSecurityProtocol = Some(securityProtocol),
-      trustStoreFile = trustStoreFile, saslProperties = serverSaslProperties, logDirCount = logDirCount)
-    props.foreach(propertyOverrides)
-    props.map(KafkaConfig.fromProps)
+  override def modifyConfigs(props: Seq[Properties]): Unit = {
+    props.foreach { p =>
+      p.put(KafkaConfig.ControlledShutdownEnableProp, "false")
+      brokerPropertyOverrides(p)
+    }
   }
 
   def anySocketServer = {
@@ -79,7 +77,7 @@ abstract class BaseRequestTest extends IntegrationTestHarness {
     new Socket("localhost", s.boundPort(ListenerName.forSecurityProtocol(protocol)))
   }
 
-  private def sendRequest(socket: Socket, request: Array[Byte]) {
+  private def sendRequest(socket: Socket, request: Array[Byte]): Unit = {
     val outgoing = new DataOutputStream(socket.getOutputStream)
     outgoing.writeInt(request.length)
     outgoing.write(request)
@@ -129,8 +127,8 @@ abstract class BaseRequestTest extends IntegrationTestHarness {
   /**
     * Serializes and sends the request to the given api.
     */
-  def send(request: AbstractRequest, apiKey: ApiKeys, socket: Socket, apiVersion: Option[Short] = None): Unit = {
-    val header = nextRequestHeader(apiKey, apiVersion.getOrElse(request.version))
+  def send(request: AbstractRequest, apiKey: ApiKeys, socket: Socket, apiVersion: Short): Unit = {
+    val header = nextRequestHeader(apiKey, apiVersion)
     val serializedBytes = request.serialize(header).array
     sendRequest(socket, serializedBytes)
   }
@@ -138,9 +136,9 @@ abstract class BaseRequestTest extends IntegrationTestHarness {
   /**
    * Receive response and return a ByteBuffer containing response without the header
    */
-  def receive(socket: Socket): ByteBuffer = {
+  def receive(socket: Socket, responseHeaderVersion: Short): ByteBuffer = {
     val response = receiveResponse(socket)
-    skipResponseHeader(response)
+    skipResponseHeader(response, responseHeaderVersion)
   }
 
   /**
@@ -148,9 +146,10 @@ abstract class BaseRequestTest extends IntegrationTestHarness {
     * A ByteBuffer containing the response is returned.
     */
   def sendAndReceive(request: AbstractRequest, apiKey: ApiKeys, socket: Socket, apiVersion: Option[Short] = None): ByteBuffer = {
-    send(request, apiKey, socket, apiVersion)
+    val version = apiVersion.getOrElse(request.version)
+    send(request, apiKey, socket, version)
     val response = receiveResponse(socket)
-    skipResponseHeader(response)
+    skipResponseHeader(response, apiKey.responseHeaderVersion(version))
   }
 
   /**
@@ -161,7 +160,7 @@ abstract class BaseRequestTest extends IntegrationTestHarness {
     val request = requestBuilder.build()
     val header = new RequestHeader(apiKey, request.version, clientId, correlationId)
     val response = requestAndReceive(socket, request.serialize(header).array)
-    val responseBuffer = skipResponseHeader(response)
+    val responseBuffer = skipResponseHeader(response, apiKey.responseHeaderVersion(request.version()))
     apiKey.parseResponse(request.version, responseBuffer)
   }
   
@@ -171,15 +170,15 @@ abstract class BaseRequestTest extends IntegrationTestHarness {
     */
   def sendStructAndReceive(requestStruct: Struct, apiKey: ApiKeys, socket: Socket, apiVersion: Short): ByteBuffer = {
     val header = nextRequestHeader(apiKey, apiVersion)
-    val serializedBytes = AbstractRequestResponse.serialize(header.toStruct, requestStruct).array
+    val serializedBytes = RequestUtils.serialize(header.toStruct, requestStruct).array
     val response = requestAndReceive(socket, serializedBytes)
-    skipResponseHeader(response)
+    skipResponseHeader(response, apiKey.responseHeaderVersion(apiVersion))
   }
 
-  protected def skipResponseHeader(response: Array[Byte]): ByteBuffer = {
+  protected def skipResponseHeader(response: Array[Byte], responseHeaderVersion: Short): ByteBuffer = {
     val responseBuffer = ByteBuffer.wrap(response)
     // Parse the header to ensure its valid and move the buffer forward
-    ResponseHeader.parse(responseBuffer)
+    ResponseHeader.parse(responseBuffer, responseHeaderVersion)
     responseBuffer
   }
 
diff --git a/core/src/test/scala/unit/kafka/server/BrokerEpochIntegrationTest.scala b/core/src/test/scala/unit/kafka/server/BrokerEpochIntegrationTest.scala
index d35b99712c494..dbeb70eb71a9c 100755
--- a/core/src/test/scala/unit/kafka/server/BrokerEpochIntegrationTest.scala
+++ b/core/src/test/scala/unit/kafka/server/BrokerEpochIntegrationTest.scala
@@ -24,10 +24,11 @@ import kafka.utils.TestUtils
 import kafka.utils.TestUtils.createTopic
 import kafka.zk.ZooKeeperTestHarness
 import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.message.LeaderAndIsrRequestData.LeaderAndIsrPartitionState
+import org.apache.kafka.common.message.UpdateMetadataRequestData.{UpdateMetadataBroker, UpdateMetadataEndpoint, UpdateMetadataPartitionState}
 import org.apache.kafka.common.metrics.Metrics
 import org.apache.kafka.common.network.ListenerName
 import org.apache.kafka.common.protocol.{ApiKeys, Errors}
-import org.apache.kafka.common.requests.UpdateMetadataRequest.EndPoint
 import org.apache.kafka.common.requests._
 import org.apache.kafka.common.security.auth.SecurityProtocol
 import org.apache.kafka.common.utils.Time
@@ -43,7 +44,7 @@ class BrokerEpochIntegrationTest extends ZooKeeperTestHarness {
   var servers: Seq[KafkaServer] = Seq.empty[KafkaServer]
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     super.setUp()
     val configs = Seq(
       TestUtils.createBrokerConfig(brokerId1, zkConnect),
@@ -57,7 +58,7 @@ class BrokerEpochIntegrationTest extends ZooKeeperTestHarness {
   }
 
   @After
-  override def tearDown() {
+  override def tearDown(): Unit = {
     TestUtils.shutdownServers(servers)
     super.tearDown()
   }
@@ -92,16 +93,16 @@ class BrokerEpochIntegrationTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  def testControlRequestWithCorrectBrokerEpoch() {
+  def testControlRequestWithCorrectBrokerEpoch(): Unit = {
     testControlRequestWithBrokerEpoch(false)
   }
 
   @Test
-  def testControlRequestWithStaleBrokerEpoch() {
+  def testControlRequestWithStaleBrokerEpoch(): Unit = {
     testControlRequestWithBrokerEpoch(true)
   }
 
-  private def testControlRequestWithBrokerEpoch(isEpochInRequestStale: Boolean) {
+  private def testControlRequestWithBrokerEpoch(isEpochInRequestStale: Boolean): Unit = {
     val tp = new TopicPartition("new-topic", 0)
 
     // create topic with 1 partition, 2 replicas, one on each broker
@@ -132,53 +133,72 @@ class BrokerEpochIntegrationTest extends ZooKeeperTestHarness {
     try {
       // Send LeaderAndIsr request with correct broker epoch
       {
-        val partitionStates = Map(
-          tp -> new LeaderAndIsrRequest.PartitionState(controllerEpoch, brokerId2, LeaderAndIsr.initialLeaderEpoch + 1,
-            Seq(brokerId1, brokerId2).map(Integer.valueOf).asJava, LeaderAndIsr.initialZKVersion,
-            Seq(0, 1).map(Integer.valueOf).asJava, false)
+        val partitionStates = Seq(
+          new LeaderAndIsrPartitionState()
+            .setTopicName(tp.topic)
+            .setPartitionIndex(tp.partition)
+            .setControllerEpoch(controllerEpoch)
+            .setLeader(brokerId2)
+            .setLeaderEpoch(LeaderAndIsr.initialLeaderEpoch + 1)
+            .setIsr(Seq(brokerId1, brokerId2).map(Integer.valueOf).asJava)
+            .setZkVersion(LeaderAndIsr.initialZKVersion)
+            .setReplicas(Seq(0, 1).map(Integer.valueOf).asJava)
+            .setIsNew(false)
         )
         val requestBuilder = new LeaderAndIsrRequest.Builder(
           ApiKeys.LEADER_AND_ISR.latestVersion, controllerId, controllerEpoch,
           epochInRequest,
+          epochInRequest,
           partitionStates.asJava, nodes.toSet.asJava)
 
         if (isEpochInRequestStale) {
-          sendAndVerifyStaleBrokerEpochInResponse(controllerChannelManager, ApiKeys.LEADER_AND_ISR, requestBuilder)
+          sendAndVerifyStaleBrokerEpochInResponse(controllerChannelManager, requestBuilder)
         }
         else {
-          sendAndVerifySuccessfulResponse(controllerChannelManager, ApiKeys.LEADER_AND_ISR, requestBuilder)
+          sendAndVerifySuccessfulResponse(controllerChannelManager, requestBuilder)
           TestUtils.waitUntilLeaderIsKnown(Seq(broker2), tp, 10000)
         }
       }
 
       // Send UpdateMetadata request with correct broker epoch
       {
-        val partitionStates = Map(
-          tp -> new UpdateMetadataRequest.PartitionState(controllerEpoch, brokerId2, LeaderAndIsr.initialLeaderEpoch + 1,
-            Seq(brokerId1, brokerId2).map(Integer.valueOf).asJava, LeaderAndIsr.initialZKVersion,
-            Seq(0, 1).map(Integer.valueOf).asJava, Seq.empty.asJava)
-        )
-        val liverBrokers = brokerAndEpochs.map { brokerAndEpoch =>
-          val broker = brokerAndEpoch._1
+        val partitionStates = Seq(
+          new UpdateMetadataPartitionState()
+            .setTopicName(tp.topic)
+            .setPartitionIndex(tp.partition)
+            .setControllerEpoch(controllerEpoch)
+            .setLeader(brokerId2)
+            .setLeaderEpoch(LeaderAndIsr.initialLeaderEpoch + 1)
+            .setIsr(Seq(brokerId1, brokerId2).map(Integer.valueOf).asJava)
+            .setZkVersion(LeaderAndIsr.initialZKVersion)
+            .setReplicas(Seq(0, 1).map(Integer.valueOf).asJava))
+        val liveBrokers = brokerAndEpochs.map { case (broker, _) =>
           val securityProtocol = SecurityProtocol.PLAINTEXT
           val listenerName = ListenerName.forSecurityProtocol(securityProtocol)
           val node = broker.node(listenerName)
-          val endPoints = Seq(new EndPoint(node.host, node.port, securityProtocol, listenerName))
-          new UpdateMetadataRequest.Broker(broker.id, endPoints.asJava, broker.rack.orNull)
-        }
+          val endpoints = Seq(new UpdateMetadataEndpoint()
+            .setHost(node.host)
+            .setPort(node.port)
+            .setSecurityProtocol(securityProtocol.id)
+            .setListener(listenerName.value))
+          new UpdateMetadataBroker()
+            .setId(broker.id)
+            .setEndpoints(endpoints.asJava)
+            .setRack(broker.rack.orNull)
+        }.toBuffer
         val requestBuilder = new UpdateMetadataRequest.Builder(
           ApiKeys.UPDATE_METADATA.latestVersion, controllerId, controllerEpoch,
-          epochInRequest,
-          partitionStates.asJava, liverBrokers.toSet.asJava)
+          epochInRequest, epochInRequest,
+          partitionStates.asJava, liveBrokers.asJava)
 
         if (isEpochInRequestStale) {
-          sendAndVerifyStaleBrokerEpochInResponse(controllerChannelManager, ApiKeys.UPDATE_METADATA, requestBuilder)
+          sendAndVerifyStaleBrokerEpochInResponse(controllerChannelManager, requestBuilder)
         }
         else {
-          sendAndVerifySuccessfulResponse(controllerChannelManager, ApiKeys.UPDATE_METADATA, requestBuilder)
-          TestUtils.waitUntilMetadataIsPropagated(Seq(broker2), tp.topic(), tp.partition(), 10000)
+          sendAndVerifySuccessfulResponse(controllerChannelManager, requestBuilder)
+          TestUtils.waitUntilMetadataIsPropagated(Seq(broker2), tp.topic, tp.partition, 10000)
           assertEquals(brokerId2,
-            broker2.metadataCache.getPartitionInfo(tp.topic(), tp.partition()).get.basePartitionState.leader)
+            broker2.metadataCache.getPartitionInfo(tp.topic, tp.partition).get.leader)
         }
       }
 
@@ -187,14 +207,14 @@ class BrokerEpochIntegrationTest extends ZooKeeperTestHarness {
         val requestBuilder = new StopReplicaRequest.Builder(
           ApiKeys.STOP_REPLICA.latestVersion, controllerId, controllerEpoch,
           epochInRequest, // Correct broker epoch
+          epochInRequest, // Correct broker epoch
           true, Set(tp).asJava)
 
         if (isEpochInRequestStale) {
-          sendAndVerifyStaleBrokerEpochInResponse(controllerChannelManager, ApiKeys.STOP_REPLICA, requestBuilder)
-        }
-        else {
-          sendAndVerifySuccessfulResponse(controllerChannelManager, ApiKeys.STOP_REPLICA, requestBuilder)
-          assertTrue(broker2.replicaManager.getPartition(tp).isEmpty)
+          sendAndVerifyStaleBrokerEpochInResponse(controllerChannelManager, requestBuilder)
+        } else {
+          sendAndVerifySuccessfulResponse(controllerChannelManager, requestBuilder)
+          assertEquals(HostedPartition.None, broker2.replicaManager.getPartition(tp))
         }
       }
     } finally {
@@ -221,22 +241,23 @@ class BrokerEpochIntegrationTest extends ZooKeeperTestHarness {
     }, "Broker epoch mismatches")
   }
 
-  private def sendAndVerifyStaleBrokerEpochInResponse(controllerChannelManager: ControllerChannelManager, apiKeys: ApiKeys,
-                                               builder: AbstractControlRequest.Builder[_ <: AbstractControlRequest]): Unit = {
+  private def sendAndVerifyStaleBrokerEpochInResponse(controllerChannelManager: ControllerChannelManager,
+                                                      builder: AbstractControlRequest.Builder[_ <: AbstractControlRequest]): Unit = {
     var staleBrokerEpochDetected = false
-    controllerChannelManager.sendRequest(brokerId2, apiKeys, builder,
-      response => {staleBrokerEpochDetected = response.errorCounts().containsKey(Errors.STALE_BROKER_EPOCH)})
+    controllerChannelManager.sendRequest(brokerId2, builder, response => {
+      staleBrokerEpochDetected = response.errorCounts().containsKey(Errors.STALE_BROKER_EPOCH)
+    })
     TestUtils.waitUntilTrue(() => staleBrokerEpochDetected, "Broker epoch should be stale")
     assertTrue("Stale broker epoch not detected by the broker", staleBrokerEpochDetected)
   }
 
-  private def sendAndVerifySuccessfulResponse(controllerChannelManager: ControllerChannelManager, apiKeys: ApiKeys,
+  private def sendAndVerifySuccessfulResponse(controllerChannelManager: ControllerChannelManager,
                                               builder: AbstractControlRequest.Builder[_ <: AbstractControlRequest]): Unit = {
     @volatile var succeed = false
-    controllerChannelManager.sendRequest(brokerId2, apiKeys, builder,
-      response => {
+    controllerChannelManager.sendRequest(brokerId2, builder, response => {
         succeed = response.errorCounts().isEmpty ||
-          (response.errorCounts().containsKey(Errors.NONE) && response.errorCounts().size() == 1)})
+          (response.errorCounts().containsKey(Errors.NONE) && response.errorCounts().size() == 1)
+    })
     TestUtils.waitUntilTrue(() => succeed, "Should receive response with no errors")
   }
 }
diff --git a/core/src/test/scala/unit/kafka/server/CacheableBrokerEpochIntegrationTest.scala b/core/src/test/scala/unit/kafka/server/CacheableBrokerEpochIntegrationTest.scala
new file mode 100644
index 0000000000000..b8ce0dea8d95b
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/server/CacheableBrokerEpochIntegrationTest.scala
@@ -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 kafka.server
+
+import java.util.Properties
+
+import kafka.api.KAFKA_2_3_IV1
+import kafka.server.KafkaConfig
+import kafka.utils.TestUtils
+import kafka.zk.ZooKeeperTestHarness
+import org.apache.kafka.common.TopicPartition
+import org.junit.Test
+
+class CacheableBrokerEpochIntegrationTest extends ZooKeeperTestHarness {
+  @Test
+  def testNewControllerConfig(): Unit = {
+    testControlRequests(true)
+  }
+
+  @Test
+  def testOldControllerConfig(): Unit = {
+    testControlRequests(false)
+  }
+
+  def testControlRequests(controllerUseNewConfig: Boolean): Unit = {
+    val controllerId = 0
+    val controllerConfig: Properties =
+      if (controllerUseNewConfig) {
+        TestUtils.createBrokerConfig(controllerId, zkConnect)
+      } else {
+        val oldConfig = TestUtils.createBrokerConfig(controllerId, zkConnect)
+        oldConfig.put(KafkaConfig.InterBrokerProtocolVersionProp, KAFKA_2_3_IV1.toString)
+        oldConfig
+      }
+    val controller = TestUtils.createServer(KafkaConfig.fromProps(controllerConfig))
+
+    // Note that broker side logic does not depend on the InterBrokerProtocolVersion config
+    val brokerId = 1
+    val brokerConfig = TestUtils.createBrokerConfig(brokerId, zkConnect)
+    val broker = TestUtils.createServer(KafkaConfig.fromProps(brokerConfig))
+    val servers = Seq(controller, broker)
+
+    val tp = new TopicPartition("new-topic", 0)
+
+    try {
+      // Use topic creation to test the LeaderAndIsr and UpdateMetadata requests
+      TestUtils.createTopic(zkClient, tp.topic(), partitionReplicaAssignment = Map(0 -> Seq(brokerId, controllerId)),
+        servers = servers)
+      TestUtils.waitUntilLeaderIsKnown(Seq(broker), tp, 10000)
+      TestUtils.waitUntilMetadataIsPropagated(Seq(broker), tp.topic(), tp.partition())
+
+      // Use topic deletion to test StopReplica requests
+      adminZkClient.deleteTopic(tp.topic())
+      TestUtils.verifyTopicDeletion(zkClient, tp.topic(), 1, servers)
+    } finally {
+      TestUtils.shutdownServers(servers)
+    }
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/server/ClientQuotaManagerTest.scala b/core/src/test/scala/unit/kafka/server/ClientQuotaManagerTest.scala
index e10d4b2d7c847..9b38ec2afeae8 100644
--- a/core/src/test/scala/unit/kafka/server/ClientQuotaManagerTest.scala
+++ b/core/src/test/scala/unit/kafka/server/ClientQuotaManagerTest.scala
@@ -23,6 +23,7 @@ import java.util.Collections
 import kafka.network.RequestChannel
 import kafka.network.RequestChannel.{EndThrottlingResponse, Session, StartThrottlingResponse}
 import kafka.server.QuotaType._
+import kafka.utils.KafkaScheduler
 import org.apache.kafka.common.TopicPartition
 import org.apache.kafka.common.memory.MemoryPool
 import org.apache.kafka.common.metrics.{MetricConfig, Metrics, Quota}
@@ -33,15 +34,22 @@ import org.apache.kafka.common.security.auth.{KafkaPrincipal, SecurityProtocol}
 import org.apache.kafka.common.utils.{MockTime, Sanitizer}
 import org.easymock.EasyMock
 import org.junit.Assert.{assertEquals, assertTrue}
-import org.junit.{Before, Test}
+import org.junit.{Before, After, Test}
 
 class ClientQuotaManagerTest {
   private val time = new MockTime
-
+  private val metrics = new Metrics(new MetricConfig(), Collections.emptyList(), time)
   private val config = ClientQuotaManagerConfig(quotaBytesPerSecondDefault = 500)
+  private val scheduler = new KafkaScheduler(1)
 
   var numCallbacks: Int = 0
-  def callback (response: RequestChannel.Response) {
+
+  @After
+  def tearDown(): Unit = {
+    metrics.close()
+  }
+
+  def callback (response: RequestChannel.Response): Unit = {
     // Count how many times this callback is called for notifyThrottlingDone().
     response match {
       case _: StartThrottlingResponse =>
@@ -52,6 +60,12 @@ class ClientQuotaManagerTest {
   @Before
   def beforeMethod() {
     numCallbacks = 0
+    scheduler.startup()
+  }
+
+  @After
+  def afterClass() {
+    scheduler.shutdown()
   }
 
   private def buildRequest[T <: AbstractRequest](builder: AbstractRequest.Builder[T],
@@ -75,13 +89,13 @@ class ClientQuotaManagerTest {
   }
 
   private def throttle(quotaManager: ClientQuotaManager, user: String, clientId: String, throttleTimeMs: Int,
-                       channelThrottlingCallback: (RequestChannel.Response) => Unit) {
+                       channelThrottlingCallback: (RequestChannel.Response) => Unit): Unit = {
     val (_, request) = buildRequest(FetchRequest.Builder.forConsumer(0, 1000, new util.HashMap[TopicPartition, PartitionData]))
     quotaManager.throttle(request, throttleTimeMs, channelThrottlingCallback)
   }
 
-  private def testQuotaParsing(config: ClientQuotaManagerConfig, client1: UserClient, client2: UserClient, randomClient: UserClient, defaultConfigClient: UserClient) {
-    val clientMetrics = new ClientQuotaManager(config, newMetrics, Produce, time, "")
+  private def testQuotaParsing(config: ClientQuotaManagerConfig, client1: UserClient, client2: UserClient, randomClient: UserClient, defaultConfigClient: UserClient): Unit = {
+    val clientMetrics = new ClientQuotaManager(config, metrics, Produce, time, Some(scheduler), "")
 
     try {
       // Case 1: Update the quota. Assert that the new quota value is returned
@@ -130,7 +144,7 @@ class ClientQuotaManagerTest {
    * Quota overrides persisted in ZooKeeper in /config/clients/, default persisted in /config/clients/
    */
   @Test
-  def testClientIdQuotaParsing() {
+  def testClientIdQuotaParsing(): Unit = {
     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)
@@ -143,7 +157,7 @@ class ClientQuotaManagerTest {
    * Quota overrides persisted in ZooKeeper in /config/users/, default persisted in /config/users/
    */
   @Test
-  def testUserQuotaParsing() {
+  def testUserQuotaParsing(): Unit = {
     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)
@@ -157,7 +171,7 @@ class ClientQuotaManagerTest {
    * Quotas persisted in ZooKeeper in /config/users//clients/, default in /config/users//clients/
    */
   @Test
-  def testUserClientIdQuotaParsing() {
+  def testUserClientIdQuotaParsing(): Unit = {
     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)
@@ -170,7 +184,7 @@ class ClientQuotaManagerTest {
    * Tests parsing for  quotas when client-id default quota properties are set.
    */
   @Test
-  def testUserQuotaParsingWithDefaultClientIdQuota() {
+  def testUserQuotaParsingWithDefaultClientIdQuota(): Unit = {
     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)
@@ -182,7 +196,7 @@ class ClientQuotaManagerTest {
    * Tests parsing for  quotas when client-id default quota properties are set.
    */
   @Test
-  def testUserClientQuotaParsingIdWithDefaultClientIdQuota() {
+  def testUserClientQuotaParsingIdWithDefaultClientIdQuota(): Unit = {
     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)
@@ -190,19 +204,103 @@ class ClientQuotaManagerTest {
     testQuotaParsing(config, client1, client2, randomClient, defaultConfigClient)
   }
 
+  private def checkQuota(quotaManager: ClientQuotaManager, user: String, clientId: String, expectedBound: Long, value: Int, expectThrottle: Boolean): Unit = {
+    assertEquals(expectedBound, quotaManager.quota(user, clientId).bound, 0.0)
+    val session = Session(new KafkaPrincipal(KafkaPrincipal.USER_TYPE, user), InetAddress.getLocalHost)
+    val expectedMaxValueInQuotaWindow =
+      if (expectedBound < Long.MaxValue) config.quotaWindowSizeSeconds * (config.numQuotaSamples - 1) * expectedBound else Double.MaxValue
+    assertEquals(expectedMaxValueInQuotaWindow, quotaManager.getMaxValueInQuotaWindow(session, clientId), 0.01)
+
+    val throttleTimeMs = maybeRecord(quotaManager, user, clientId, value * config.numQuotaSamples)
+    if (expectThrottle)
+      assertTrue(s"throttleTimeMs should be > 0. was $throttleTimeMs", throttleTimeMs > 0)
+    else
+      assertEquals(s"throttleTimeMs should be 0. was $throttleTimeMs", 0, throttleTimeMs)
+  }
+
   @Test
-  def testQuotaConfigPrecedence() {
-    val quotaManager = new ClientQuotaManager(ClientQuotaManagerConfig(quotaBytesPerSecondDefault=Long.MaxValue),
-        newMetrics, Produce, time, "")
-
-    def checkQuota(user: String, clientId: String, expectedBound: Int, value: Int, expectThrottle: Boolean) {
-      assertEquals(expectedBound, quotaManager.quota(user, clientId).bound, 0.0)
-      val throttleTimeMs = maybeRecord(quotaManager, user, clientId, value * config.numQuotaSamples)
-      if (expectThrottle)
-        assertTrue(s"throttleTimeMs should be > 0. was $throttleTimeMs", throttleTimeMs > 0)
-      else
-        assertEquals(s"throttleTimeMs should be 0. was $throttleTimeMs", 0, throttleTimeMs)
+  def testGetMaxValueInQuotaWindowWithNonDefaultQuotaWindow(): Unit = {
+    val numFullQuotaWindows = 3   // 3 seconds window (vs. 10 seconds default)
+    val nonDefaultConfig = ClientQuotaManagerConfig(quotaBytesPerSecondDefault = Long.MaxValue, numQuotaSamples = numFullQuotaWindows + 1)
+    val quotaManager = new ClientQuotaManager(nonDefaultConfig, metrics, Fetch, time, Some(scheduler), "")
+    val userSession = Session(new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "userA"), InetAddress.getLocalHost)
+
+    try {
+      // no quota set
+      assertEquals(Double.MaxValue, quotaManager.getMaxValueInQuotaWindow(userSession, "client1"), 0.01)
+
+      // Set default  quota config
+      quotaManager.updateQuota(Some(ConfigEntityName.Default), None, None, Some(new Quota(10, true)))
+      assertEquals(10 * numFullQuotaWindows, quotaManager.getMaxValueInQuotaWindow(userSession, "client1"), 0.01)
+    } finally {
+      quotaManager.shutdown()
     }
+  }
+
+  @Test
+  def testSetAndRemoveDefaultUserQuota(): Unit = {
+    // quotaTypesEnabled will be QuotaTypes.NoQuotas initially
+    val quotaManager = new ClientQuotaManager(ClientQuotaManagerConfig(quotaBytesPerSecondDefault = Long.MaxValue),
+      metrics, Produce, time, Some(scheduler), "")
+
+    try {
+      // no quota set yet, should not throttle
+      checkQuota(quotaManager, "userA", "client1", Long.MaxValue, 1000, false)
+
+      // Set default  quota config
+      quotaManager.updateQuota(Some(ConfigEntityName.Default), None, None, Some(new Quota(10, true)))
+      checkQuota(quotaManager, "userA", "client1", 10, 1000, true)
+
+      // Remove default  quota config, back to no quotas
+      quotaManager.updateQuota(Some(ConfigEntityName.Default), None, None, None)
+      checkQuota(quotaManager, "userA", "client1", Long.MaxValue, 1000, false)
+    } finally {
+      quotaManager.shutdown()
+    }
+  }
+
+  @Test
+  def testSetAndRemoveUserQuota(): Unit = {
+    // quotaTypesEnabled will be QuotaTypes.NoQuotas initially
+    val quotaManager = new ClientQuotaManager(ClientQuotaManagerConfig(quotaBytesPerSecondDefault = Long.MaxValue),
+      metrics, Produce, time, Some(scheduler), "")
+
+    try {
+      // Set  quota config
+      quotaManager.updateQuota(Some("userA"), None, None, Some(new Quota(10, true)))
+      checkQuota(quotaManager, "userA", "client1", 10, 1000, true)
+
+      // Remove  quota config, back to no quotas
+      quotaManager.updateQuota(Some("userA"), None, None, None)
+      checkQuota(quotaManager, "userA", "client1", Long.MaxValue, 1000, false)
+    } finally {
+      quotaManager.shutdown()
+    }
+  }
+
+  @Test
+  def testSetAndRemoveUserClientQuota(): Unit = {
+    // quotaTypesEnabled will be QuotaTypes.NoQuotas initially
+    val quotaManager = new ClientQuotaManager(ClientQuotaManagerConfig(quotaBytesPerSecondDefault = Long.MaxValue),
+      metrics, Produce, time, Some(scheduler), "")
+
+    try {
+      // Set  quota config
+      quotaManager.updateQuota(Some("userA"), Some("client1"), Some("client1"), Some(new Quota(10, true)))
+      checkQuota(quotaManager, "userA", "client1", 10, 1000, true)
+
+      // Remove  quota config, back to no quotas
+      quotaManager.updateQuota(Some("userA"), Some("client1"), Some("client1"), None)
+      checkQuota(quotaManager, "userA", "client1", Long.MaxValue, 1000, false)
+    } finally {
+      quotaManager.shutdown()
+    }
+  }
+
+  @Test
+  def testQuotaConfigPrecedence(): Unit = {
+    val quotaManager = new ClientQuotaManager(ClientQuotaManagerConfig(quotaBytesPerSecondDefault=Long.MaxValue),
+      metrics, Produce, time, Some(scheduler), "")
 
     try {
       quotaManager.updateQuota(Some(ConfigEntityName.Default), None, None, Some(new Quota(1000, true)))
@@ -216,47 +314,47 @@ class ClientQuotaManagerTest {
       quotaManager.updateQuota(Some("userC"), None, None, Some(new Quota(10000, true)))
       quotaManager.updateQuota(None, Some("client1"), Some("client1"), Some(new Quota(9000, true)))
 
-      checkQuota("userA", "client1", 5000, 4500, false) //  quota takes precedence over 
-      checkQuota("userA", "client2", 4000, 4500, true)  //  quota takes precedence over  and defaults
-      checkQuota("userA", "client3", 4000, 0, true)     //  quota is shared across clients of user
-      checkQuota("userA", "client1", 5000, 0, false)    //  is exclusive use, unaffected by other clients
+      checkQuota(quotaManager, "userA", "client1", 5000, 4500, false) //  quota takes precedence over 
+      checkQuota(quotaManager, "userA", "client2", 4000, 4500, true)  //  quota takes precedence over  and defaults
+      checkQuota(quotaManager, "userA", "client3", 4000, 0, true)     //  quota is shared across clients of user
+      checkQuota(quotaManager, "userA", "client1", 5000, 0, false)    //  is exclusive use, unaffected by other clients
 
-      checkQuota("userB", "client1", 7000, 8000, true)
-      checkQuota("userB", "client2", 8000, 7000, false) // Default per-client quota for exclusive use of 
-      checkQuota("userB", "client3", 8000, 7000, false)
+      checkQuota(quotaManager, "userB", "client1", 7000, 8000, true)
+      checkQuota(quotaManager, "userB", "client2", 8000, 7000, false) // Default per-client quota for exclusive use of 
+      checkQuota(quotaManager, "userB", "client3", 8000, 7000, false)
 
-      checkQuota("userD", "client1", 3000, 3500, true)  // Default  quota
-      checkQuota("userD", "client2", 3000, 2500, false)
-      checkQuota("userE", "client1", 3000, 2500, false)
+      checkQuota(quotaManager, "userD", "client1", 3000, 3500, true)  // Default  quota
+      checkQuota(quotaManager, "userD", "client2", 3000, 2500, false)
+      checkQuota(quotaManager, "userE", "client1", 3000, 2500, false)
 
       // Remove default  quota config, revert to  default
       quotaManager.updateQuota(Some(ConfigEntityName.Default), Some(ConfigEntityName.Default), Some(ConfigEntityName.Default), None)
-      checkQuota("userD", "client1", 1000, 0, false)    // Metrics tags changed, restart counter
-      checkQuota("userE", "client4", 1000, 1500, true)
-      checkQuota("userF", "client4", 1000, 800, false)  // Default  quota shared across clients of user
-      checkQuota("userF", "client5", 1000, 800, true)
+      checkQuota(quotaManager, "userD", "client1", 1000, 0, false)    // Metrics tags changed, restart counter
+      checkQuota(quotaManager, "userE", "client4", 1000, 1500, true)
+      checkQuota(quotaManager, "userF", "client4", 1000, 800, false)  // Default  quota shared across clients of user
+      checkQuota(quotaManager, "userF", "client5", 1000, 800, true)
 
       // Remove default  quota config, revert to  default
       quotaManager.updateQuota(Some(ConfigEntityName.Default), None, None, None)
-      checkQuota("userF", "client4", 2000, 0, false)  // Default  quota shared across client-id of all users
-      checkQuota("userF", "client5", 2000, 0, false)
-      checkQuota("userF", "client5", 2000, 2500, true)
-      checkQuota("userG", "client5", 2000, 0, true)
+      checkQuota(quotaManager, "userF", "client4", 2000, 0, false)  // Default  quota shared across client-id of all users
+      checkQuota(quotaManager, "userF", "client5", 2000, 0, false)
+      checkQuota(quotaManager, "userF", "client5", 2000, 2500, true)
+      checkQuota(quotaManager, "userG", "client5", 2000, 0, true)
 
       // Update quotas
       quotaManager.updateQuota(Some("userA"), None, None, Some(new Quota(8000, true)))
       quotaManager.updateQuota(Some("userA"), Some("client1"), Some("client1"), Some(new Quota(10000, true)))
-      checkQuota("userA", "client2", 8000, 0, false)
-      checkQuota("userA", "client2", 8000, 4500, true) // Throttled due to sum of new and earlier values
-      checkQuota("userA", "client1", 10000, 0, false)
-      checkQuota("userA", "client1", 10000, 6000, true)
+      checkQuota(quotaManager, "userA", "client2", 8000, 0, false)
+      checkQuota(quotaManager, "userA", "client2", 8000, 4500, true) // Throttled due to sum of new and earlier values
+      checkQuota(quotaManager, "userA", "client1", 10000, 0, false)
+      checkQuota(quotaManager, "userA", "client1", 10000, 6000, true)
       quotaManager.updateQuota(Some("userA"), Some("client1"), Some("client1"), None)
-      checkQuota("userA", "client6", 8000, 0, true)    // Throttled due to shared user quota
+      checkQuota(quotaManager, "userA", "client6", 8000, 0, true)    // Throttled due to shared user quota
       quotaManager.updateQuota(Some("userA"), Some("client6"), Some("client6"), Some(new Quota(11000, true)))
-      checkQuota("userA", "client6", 11000, 8500, false)
+      checkQuota(quotaManager, "userA", "client6", 11000, 8500, false)
       quotaManager.updateQuota(Some("userA"), Some(ConfigEntityName.Default), Some(ConfigEntityName.Default), Some(new Quota(12000, true)))
       quotaManager.updateQuota(Some("userA"), Some("client6"), Some("client6"), None)
-      checkQuota("userA", "client6", 12000, 4000, true) // Throttled due to sum of new and earlier values
+      checkQuota(quotaManager, "userA", "client6", 12000, 4000, true) // Throttled due to sum of new and earlier values
 
     } finally {
       quotaManager.shutdown()
@@ -264,10 +362,10 @@ class ClientQuotaManagerTest {
   }
 
   @Test
-  def testQuotaViolation() {
-    val metrics = newMetrics
-    val clientMetrics = new ClientQuotaManager(config, metrics, Produce, time, "")
+  def testQuotaViolation(): Unit = {
+    val clientMetrics = new ClientQuotaManager(config, metrics, Produce, time, Some(scheduler), "")
     val queueSizeMetric = metrics.metrics().get(metrics.metricName("queue-size", "Produce", ""))
+    val throttleCountMetric = metrics.metrics().get(metrics.metricName("throttle-count", "Produce", ""))
     try {
       /* We have 10 second windows. Make sure that there is no quota violation
        * if we produce under the quota
@@ -287,7 +385,10 @@ class ClientQuotaManagerTest {
 
       assertEquals("Should be throttled", 2100, sleepTime)
       throttle(clientMetrics, "ANONYMOYUS", "unknown", sleepTime, callback)
+
       assertEquals(1, queueSizeMetric.metricValue.asInstanceOf[Double].toInt)
+      assertEquals(1, throttleCountMetric.metricValue.asInstanceOf[Double].toInt)
+
       // After a request is delayed, the callback cannot be triggered immediately
       clientMetrics.throttledChannelReaper.doWork()
       assertEquals(0, numCallbacks)
@@ -295,7 +396,10 @@ class ClientQuotaManagerTest {
 
       // Callback can only be triggered after the delay time passes
       clientMetrics.throttledChannelReaper.doWork()
+
+      assertEquals(1, throttleCountMetric.metricValue.asInstanceOf[Double].toInt)
       assertEquals(0, queueSizeMetric.metricValue.asInstanceOf[Double].toInt)
+
       assertEquals(1, numCallbacks)
 
       // Could continue to see delays until the bursty sample disappears
@@ -312,9 +416,8 @@ class ClientQuotaManagerTest {
   }
 
   @Test
-  def testRequestPercentageQuotaViolation() {
-    val metrics = newMetrics
-    val quotaManager = new ClientRequestQuotaManager(config, metrics, time, "", None)
+  def testRequestPercentageQuotaViolation(): Unit = {
+    val quotaManager = new ClientRequestQuotaManager(config, metrics, time, None, "", None)
     quotaManager.updateQuota(Some("ANONYMOUS"), Some("test-client"), Some("test-client"), Some(Quota.upperBound(1)))
     val queueSizeMetric = metrics.metrics().get(metrics.metricName("queue-size", "Request", ""))
     def millisToPercent(millis: Double) = millis * 1000 * 1000 * ClientQuotaManagerConfig.NanosToPercentagePerSecond
@@ -375,9 +478,8 @@ class ClientQuotaManagerTest {
   }
 
   @Test
-  def testExpireThrottleTimeSensor() {
-    val metrics = newMetrics
-    val clientMetrics = new ClientQuotaManager(config, metrics, Produce, time, "")
+  def testExpireThrottleTimeSensor(): Unit = {
+    val clientMetrics = new ClientQuotaManager(config, metrics, Produce, time, Some(scheduler), "")
     try {
       maybeRecord(clientMetrics, "ANONYMOUS", "client1", 100)
       // remove the throttle time sensor
@@ -395,9 +497,8 @@ class ClientQuotaManagerTest {
   }
 
   @Test
-  def testExpireQuotaSensors() {
-    val metrics = newMetrics
-    val clientMetrics = new ClientQuotaManager(config, metrics, Produce, time, "")
+  def testExpireQuotaSensors(): Unit = {
+    val clientMetrics = new ClientQuotaManager(config, metrics, Produce, time, Some(scheduler), "")
     try {
       maybeRecord(clientMetrics, "ANONYMOUS", "client1", 100)
       // remove all the sensors
@@ -419,9 +520,8 @@ class ClientQuotaManagerTest {
   }
 
   @Test
-  def testClientIdNotSanitized() {
-    val metrics = newMetrics
-    val clientMetrics = new ClientQuotaManager(config, metrics, Produce, time, "")
+  def testClientIdNotSanitized(): Unit = {
+    val clientMetrics = new ClientQuotaManager(config, metrics, Produce, time, Some(scheduler), "")
     val clientId = "client@#$%"
     try {
       maybeRecord(clientMetrics, "ANONYMOUS", clientId, 100)
@@ -437,10 +537,6 @@ class ClientQuotaManagerTest {
     }
   }
 
-  def newMetrics: Metrics = {
-    new Metrics(new MetricConfig(), Collections.emptyList(), time)
-  }
-
   private case class UserClient(val user: String, val clientId: String, val configUser: Option[String] = None, val configClientId: Option[String] = None) {
     // 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
diff --git a/core/src/test/scala/unit/kafka/server/CreateTopicsRequestTest.scala b/core/src/test/scala/unit/kafka/server/CreateTopicsRequestTest.scala
index db2028c32b2a1..259fadcb7abb0 100644
--- a/core/src/test/scala/unit/kafka/server/CreateTopicsRequestTest.scala
+++ b/core/src/test/scala/unit/kafka/server/CreateTopicsRequestTest.scala
@@ -24,7 +24,7 @@ import org.junit.Test
 
 class CreateTopicsRequestTest extends AbstractCreateTopicsRequestTest {
   @Test
-  def testValidCreateTopicsRequests() {
+  def testValidCreateTopicsRequests(): Unit = {
     // Generated assignments
     validateValidCreateTopicsRequests(topicsReq(Seq(topicReq("topic1"))))
     validateValidCreateTopicsRequests(topicsReq(Seq(topicReq("topic2", replicationFactor = 3))))
@@ -43,19 +43,26 @@ class CreateTopicsRequestTest extends AbstractCreateTopicsRequestTest {
       topicReq("topic10", numPartitions = 5, replicationFactor = 2),
       topicReq("topic11", assignment = Map(0 -> List(0, 1), 1 -> List(1, 0), 2 -> List(1, 2)))),
       validateOnly = true))
+    // Defaults
+    validateValidCreateTopicsRequests(topicsReq(Seq(
+      topicReq("topic12", replicationFactor = -1, numPartitions = -1))))
+    validateValidCreateTopicsRequests(topicsReq(Seq(
+      topicReq("topic13", replicationFactor = 2, numPartitions = -1))))
+    validateValidCreateTopicsRequests(topicsReq(Seq(
+      topicReq("topic14", replicationFactor = -1, numPartitions = 2))))
   }
 
   @Test
-  def testErrorCreateTopicsRequests() {
+  def testErrorCreateTopicsRequests(): Unit = {
     val existingTopic = "existing-topic"
     createTopic(existingTopic, 1, 1)
     // Basic
     validateErrorCreateTopicsRequests(topicsReq(Seq(topicReq(existingTopic))),
       Map(existingTopic -> error(Errors.TOPIC_ALREADY_EXISTS, Some("Topic 'existing-topic' already exists."))))
-    validateErrorCreateTopicsRequests(topicsReq(Seq(topicReq("error-partitions", numPartitions = -1))),
+    validateErrorCreateTopicsRequests(topicsReq(Seq(topicReq("error-partitions", numPartitions = -2))),
       Map("error-partitions" -> error(Errors.INVALID_PARTITIONS)), checkErrorMessage = false)
     validateErrorCreateTopicsRequests(topicsReq(Seq(topicReq("error-replication",
-      replicationFactor = numBrokers + 1))),
+      replicationFactor = brokerCount + 1))),
       Map("error-replication" -> error(Errors.INVALID_REPLICATION_FACTOR)), checkErrorMessage = false)
     validateErrorCreateTopicsRequests(topicsReq(Seq(topicReq("error-config",
       config=Map("not.a.property" -> "error")))),
@@ -70,8 +77,8 @@ class CreateTopicsRequestTest extends AbstractCreateTopicsRequestTest {
     // Partial
     validateErrorCreateTopicsRequests(topicsReq(Seq(
       topicReq(existingTopic),
-      topicReq("partial-partitions", numPartitions = -1),
-      topicReq("partial-replication", replicationFactor=numBrokers + 1),
+      topicReq("partial-partitions", numPartitions = -2),
+      topicReq("partial-replication", replicationFactor=brokerCount + 1),
       topicReq("partial-assignment", assignment=Map(0 -> List(0, 1), 1 -> List(0))),
       topicReq("partial-none"))),
       Map(
@@ -106,7 +113,7 @@ class CreateTopicsRequestTest extends AbstractCreateTopicsRequestTest {
   }
 
   @Test
-  def testInvalidCreateTopicsRequests() {
+  def testInvalidCreateTopicsRequests(): Unit = {
     // Partitions/ReplicationFactor and ReplicaAssignment
     validateErrorCreateTopicsRequests(topicsReq(Seq(
       topicReq("bad-args-topic", numPartitions = 10, replicationFactor = 3,
@@ -120,7 +127,7 @@ class CreateTopicsRequestTest extends AbstractCreateTopicsRequestTest {
   }
 
   @Test
-  def testNotController() {
+  def testNotController(): Unit = {
     val req = topicsReq(Seq(topicReq("topic1")))
     val response = sendCreateTopicRequest(req, notControllerSocketServer)
     assertEquals(1, response.errorCounts().get(Errors.NOT_CONTROLLER))
diff --git a/core/src/test/scala/unit/kafka/server/CreateTopicsRequestWithPolicyTest.scala b/core/src/test/scala/unit/kafka/server/CreateTopicsRequestWithPolicyTest.scala
index 4fc3244053d41..aef408041c282 100644
--- a/core/src/test/scala/unit/kafka/server/CreateTopicsRequestWithPolicyTest.scala
+++ b/core/src/test/scala/unit/kafka/server/CreateTopicsRequestWithPolicyTest.scala
@@ -32,13 +32,13 @@ import scala.collection.JavaConverters._
 class CreateTopicsRequestWithPolicyTest extends AbstractCreateTopicsRequestTest {
   import CreateTopicsRequestWithPolicyTest._
 
-  override def propertyOverrides(properties: Properties): Unit = {
-    super.propertyOverrides(properties)
+  override def brokerPropertyOverrides(properties: Properties): Unit = {
+    super.brokerPropertyOverrides(properties)
     properties.put(KafkaConfig.CreateTopicPolicyClassNameProp, classOf[Policy].getName)
   }
 
   @Test
-  def testValidCreateTopicsRequests() {
+  def testValidCreateTopicsRequests(): Unit = {
     validateValidCreateTopicsRequests(topicsReq(Seq(topicReq("topic1",
       numPartitions = 5))))
 
@@ -56,7 +56,7 @@ class CreateTopicsRequestWithPolicyTest extends AbstractCreateTopicsRequestTest
   }
 
   @Test
-  def testErrorCreateTopicsRequests() {
+  def testErrorCreateTopicsRequests(): Unit = {
     val existingTopic = "existing-topic"
     createTopic(existingTopic, 1, 1)
 
@@ -94,14 +94,19 @@ class CreateTopicsRequestWithPolicyTest extends AbstractCreateTopicsRequestTest
         Some("Topic 'existing-topic' already exists."))))
 
     validateErrorCreateTopicsRequests(topicsReq(Seq(topicReq("error-replication",
-      numPartitions = 10, replicationFactor = numBrokers + 1)), validateOnly = true),
+      numPartitions = 10, replicationFactor = brokerCount + 1)), validateOnly = true),
       Map("error-replication" -> error(Errors.INVALID_REPLICATION_FACTOR,
         Some("Replication factor: 4 larger than available brokers: 3."))))
 
     validateErrorCreateTopicsRequests(topicsReq(Seq(topicReq("error-replication2",
-      numPartitions = 10, replicationFactor = -1)), validateOnly = true),
+      numPartitions = 10, replicationFactor = -2)), validateOnly = true),
       Map("error-replication2" -> error(Errors.INVALID_REPLICATION_FACTOR,
         Some("Replication factor must be larger than 0."))))
+
+    validateErrorCreateTopicsRequests(topicsReq(Seq(topicReq("error-partitions",
+      numPartitions = -2, replicationFactor = 1)), validateOnly = true),
+      Map("error-partitions" -> error(Errors.INVALID_PARTITIONS,
+        Some("Number of partitions must be larger than 0."))))
   }
 
 }
@@ -141,7 +146,7 @@ object CreateTopicsRequestWithPolicyTest {
         require(replicationFactor == null, s"replicationFactor should be null, but it is $replicationFactor")
         require(replicasAssignments != null, s"replicaAssigments should not be null, but it is $replicasAssignments")
 
-        replicasAssignments.asScala.foreach { case (partitionId, assignment) =>
+        replicasAssignments.asScala.toSeq.sortBy { case (tp, _) => tp }.foreach { case (partitionId, assignment) =>
           if (assignment.size < 2)
             throw new PolicyViolationException("Topic partitions should have at least 2 partitions, received " +
               s"${assignment.size} for partition $partitionId")
diff --git a/core/src/test/scala/unit/kafka/server/DelayedOperationTest.scala b/core/src/test/scala/unit/kafka/server/DelayedOperationTest.scala
index 3b077a0b438dc..212ef80e099dc 100644
--- a/core/src/test/scala/unit/kafka/server/DelayedOperationTest.scala
+++ b/core/src/test/scala/unit/kafka/server/DelayedOperationTest.scala
@@ -27,6 +27,9 @@ import kafka.utils.TestUtils
 import org.apache.kafka.common.utils.Time
 import org.junit.{After, Before, Test}
 import org.junit.Assert._
+import org.scalatest.Assertions.intercept
+
+import scala.collection.JavaConverters._
 
 class DelayedOperationTest {
 
@@ -34,19 +37,19 @@ class DelayedOperationTest {
   var executorService: ExecutorService = null
 
   @Before
-  def setUp() {
+  def setUp(): Unit = {
     purgatory = DelayedOperationPurgatory[MockDelayedOperation](purgatoryName = "mock")
   }
 
   @After
-  def tearDown() {
+  def tearDown(): Unit = {
     purgatory.shutdown()
     if (executorService != null)
       executorService.shutdown()
   }
 
   @Test
-  def testRequestSatisfaction() {
+  def testRequestSatisfaction(): Unit = {
     val r1 = new MockDelayedOperation(100000L)
     val r2 = new MockDelayedOperation(100000L)
     assertEquals("With no waiting requests, nothing should be satisfied", 0, purgatory.checkAndComplete("test1"))
@@ -63,7 +66,7 @@ class DelayedOperationTest {
   }
 
   @Test
-  def testRequestExpiry() {
+  def testRequestExpiry(): Unit = {
     val expiration = 20L
     val start = Time.SYSTEM.hiResClockMs
     val r1 = new MockDelayedOperation(expiration)
@@ -78,7 +81,69 @@ class DelayedOperationTest {
   }
 
   @Test
-  def testRequestPurge() {
+  def testDelayedFuture(): Unit = {
+    val purgatoryName = "testDelayedFuture"
+    val purgatory = new DelayedFuturePurgatory(purgatoryName, brokerId = 0)
+    val result = new AtomicInteger()
+
+    def hasExecutorThread: Boolean = Thread.getAllStackTraces.keySet.asScala.map(_.getName)
+      .exists(_.contains(s"DelayedExecutor-$purgatoryName"))
+    def updateResult(futures: List[CompletableFuture[Integer]]): Unit =
+      result.set(futures.filterNot(_.isCompletedExceptionally).map(_.get.intValue).sum)
+
+    assertFalse("Unnecessary thread created", hasExecutorThread)
+
+    // Two completed futures: callback should be executed immediately on the same thread
+    val futures1 = List(CompletableFuture.completedFuture(10.asInstanceOf[Integer]),
+      CompletableFuture.completedFuture(11.asInstanceOf[Integer]))
+    val r1 = purgatory.tryCompleteElseWatch[Integer](100000L, futures1, () => updateResult(futures1))
+    assertTrue("r1 not completed", r1.isCompleted)
+    assertEquals(21, result.get())
+    assertFalse("Unnecessary thread created", hasExecutorThread)
+
+    // Two delayed futures: callback should wait for both to complete
+    result.set(-1)
+    val futures2 = List(new CompletableFuture[Integer], new CompletableFuture[Integer])
+    val r2 = purgatory.tryCompleteElseWatch[Integer](100000L, futures2, () => updateResult(futures2))
+    assertFalse("r2 should be incomplete", r2.isCompleted)
+    futures2.head.complete(20)
+    assertFalse(r2.isCompleted)
+    assertEquals(-1, result.get())
+    futures2(1).complete(21)
+    TestUtils.waitUntilTrue(() => r2.isCompleted, "r2 not completed")
+    TestUtils.waitUntilTrue(() => result.get == 41, "callback not invoked")
+    assertTrue("Thread not created for executing delayed task", hasExecutorThread)
+
+    // One immediate and one delayed future: callback should wait for delayed task to complete
+    result.set(-1)
+    val futures3 = List(new CompletableFuture[Integer], CompletableFuture.completedFuture(31.asInstanceOf[Integer]))
+    val r3 = purgatory.tryCompleteElseWatch[Integer](100000L, futures3, () => updateResult(futures3))
+    assertFalse("r3 should be incomplete", r3.isCompleted)
+    assertEquals(-1, result.get())
+    futures3.head.complete(30)
+    TestUtils.waitUntilTrue(() => r3.isCompleted, "r3 not completed")
+    TestUtils.waitUntilTrue(() => result.get == 61, "callback not invoked")
+
+
+    // One future doesn't complete within timeout. Should expire and invoke callback after timeout.
+    result.set(-1)
+    val start = Time.SYSTEM.hiResClockMs
+    val expirationMs = 2000L
+    val futures4 = List(new CompletableFuture[Integer], new CompletableFuture[Integer])
+    val r4 = purgatory.tryCompleteElseWatch[Integer](expirationMs, futures4, () => updateResult(futures4))
+    futures4.head.complete(40)
+    TestUtils.waitUntilTrue(() => futures4(1).isDone, "r4 futures not expired")
+    assertTrue("r4 not completed after timeout", r4.isCompleted)
+    val elapsed = Time.SYSTEM.hiResClockMs - start
+    assertTrue(s"Time for expiration $elapsed should at least $expirationMs", elapsed >= expirationMs)
+    assertEquals(40, futures4.head.get)
+    assertEquals(classOf[org.apache.kafka.common.errors.TimeoutException],
+      intercept[ExecutionException](futures4(1).get).getCause.getClass)
+    assertEquals(40, result.get())
+  }
+
+  @Test
+  def testRequestPurge(): Unit = {
     val r1 = new MockDelayedOperation(100000L)
     val r2 = new MockDelayedOperation(100000L)
     val r3 = new MockDelayedOperation(100000L)
@@ -86,17 +151,17 @@ class DelayedOperationTest {
     purgatory.tryCompleteElseWatch(r2, Array("test1", "test2"))
     purgatory.tryCompleteElseWatch(r3, Array("test1", "test2", "test3"))
 
-    assertEquals("Purgatory should have 3 total delayed operations", 3, purgatory.delayed)
+    assertEquals("Purgatory should have 3 total delayed operations", 3, purgatory.numDelayed)
     assertEquals("Purgatory should have 6 watched elements", 6, purgatory.watched)
 
     // complete the operations, it should immediately be purged from the delayed operation
     r2.completable = true
     r2.tryComplete()
-    assertEquals("Purgatory should have 2 total delayed operations instead of " + purgatory.delayed, 2, purgatory.delayed)
+    assertEquals("Purgatory should have 2 total delayed operations instead of " + purgatory.numDelayed, 2, purgatory.numDelayed)
 
     r3.completable = true
     r3.tryComplete()
-    assertEquals("Purgatory should have 1 total delayed operations instead of " + purgatory.delayed, 1, purgatory.delayed)
+    assertEquals("Purgatory should have 1 total delayed operations instead of " + purgatory.numDelayed, 1, purgatory.numDelayed)
 
     // checking a watch should purge the watch list
     purgatory.checkAndComplete("test1")
@@ -110,19 +175,19 @@ class DelayedOperationTest {
   }
 
   @Test
-  def shouldCancelForKeyReturningCancelledOperations() {
+  def shouldCancelForKeyReturningCancelledOperations(): Unit = {
     purgatory.tryCompleteElseWatch(new MockDelayedOperation(10000L), Seq("key"))
     purgatory.tryCompleteElseWatch(new MockDelayedOperation(10000L), Seq("key"))
     purgatory.tryCompleteElseWatch(new MockDelayedOperation(10000L), Seq("key2"))
 
     val cancelledOperations = purgatory.cancelForKey("key")
     assertEquals(2, cancelledOperations.size)
-    assertEquals(1, purgatory.delayed)
+    assertEquals(1, purgatory.numDelayed)
     assertEquals(1, purgatory.watched)
   }
 
   @Test
-  def shouldReturnNilOperationsOnCancelForKeyWhenKeyDoesntExist() {
+  def shouldReturnNilOperationsOnCancelForKeyWhenKeyDoesntExist(): Unit = {
     val cancelledOperations = purgatory.cancelForKey("key")
     assertEquals(Nil, cancelledOperations)
   }
@@ -216,12 +281,12 @@ class DelayedOperationTest {
   }
 
   @Test
-  def testDelayedOperationLock() {
+  def testDelayedOperationLock(): Unit = {
     verifyDelayedOperationLock(new MockDelayedOperation(100000L), mismatchedLocks = false)
   }
 
   @Test
-  def testDelayedOperationLockOverride() {
+  def testDelayedOperationLockOverride(): Unit = {
     def newMockOperation = {
       val lock = new ReentrantLock
       new MockDelayedOperation(100000L, Some(lock), Some(lock))
@@ -232,7 +297,7 @@ class DelayedOperationTest {
         mismatchedLocks = true)
   }
 
-  def verifyDelayedOperationLock(mockDelayedOperation: => MockDelayedOperation, mismatchedLocks: Boolean) {
+  def verifyDelayedOperationLock(mockDelayedOperation: => MockDelayedOperation, mismatchedLocks: Boolean): Unit = {
     val key = "key"
     executorService = Executors.newSingleThreadExecutor
     def createDelayedOperations(count: Int): Seq[MockDelayedOperation] = {
@@ -327,7 +392,7 @@ class DelayedOperationTest {
                              extends DelayedOperation(delayMs, lockOpt) {
     var completable = false
 
-    def awaitExpiration() {
+    def awaitExpiration(): Unit = {
       synchronized {
         wait()
       }
@@ -340,11 +405,11 @@ class DelayedOperationTest {
         false
     }
 
-    override def onExpiration() {
+    override def onExpiration(): Unit = {
 
     }
 
-    override def onComplete() {
+    override def onComplete(): Unit = {
       responseLockOpt.foreach { lock =>
         if (!lock.tryLock())
           throw new IllegalStateException("Response callback lock could not be acquired in callback")
diff --git a/core/src/test/scala/unit/kafka/server/DelegationTokenRequestsOnPlainTextTest.scala b/core/src/test/scala/unit/kafka/server/DelegationTokenRequestsOnPlainTextTest.scala
index 6d56a023930d5..19ca1dd78341c 100644
--- a/core/src/test/scala/unit/kafka/server/DelegationTokenRequestsOnPlainTextTest.scala
+++ b/core/src/test/scala/unit/kafka/server/DelegationTokenRequestsOnPlainTextTest.scala
@@ -19,17 +19,18 @@ package kafka.server
 import java.util
 
 import kafka.utils.TestUtils
-import org.apache.kafka.clients.admin.{AdminClient, AdminClientConfig}
+import org.apache.kafka.clients.admin.{Admin, AdminClient, AdminClientConfig}
 import org.apache.kafka.common.errors.UnsupportedByAuthenticationException
 import org.junit.{After, Before, Test}
+import org.scalatest.Assertions.intercept
 
 import scala.collection.JavaConverters._
 import scala.concurrent.ExecutionException
 
 class DelegationTokenRequestsOnPlainTextTest extends BaseRequestTest {
-  var adminClient: AdminClient = null
+  var adminClient: Admin = null
 
-  override def numBrokers = 1
+  override def brokerCount = 1
 
   @Before
   override def setUp(): Unit = {
diff --git a/core/src/test/scala/unit/kafka/server/DelegationTokenRequestsTest.scala b/core/src/test/scala/unit/kafka/server/DelegationTokenRequestsTest.scala
index 6f79a9a20a6c6..52f78859a4c42 100644
--- a/core/src/test/scala/unit/kafka/server/DelegationTokenRequestsTest.scala
+++ b/core/src/test/scala/unit/kafka/server/DelegationTokenRequestsTest.scala
@@ -20,12 +20,13 @@ import java.util
 
 import kafka.api.{KafkaSasl, SaslSetup}
 import kafka.utils.{JaasTestUtils, TestUtils}
-import org.apache.kafka.clients.admin.{AdminClient, AdminClientConfig, CreateDelegationTokenOptions, DescribeDelegationTokenOptions}
+import org.apache.kafka.clients.admin.{Admin, AdminClient, AdminClientConfig, CreateDelegationTokenOptions, DescribeDelegationTokenOptions}
 import org.apache.kafka.common.errors.InvalidPrincipalTypeException
 import org.apache.kafka.common.security.auth.SecurityProtocol
 import org.apache.kafka.common.utils.SecurityUtils
 import org.junit.Assert._
 import org.junit.{After, Before, Test}
+import org.scalatest.Assertions.intercept
 
 import scala.collection.JavaConverters._
 import scala.concurrent.ExecutionException
@@ -36,9 +37,9 @@ class DelegationTokenRequestsTest extends BaseRequestTest with SaslSetup {
   private val kafkaServerSaslMechanisms = List("PLAIN")
   protected override val serverSaslProperties = Some(kafkaServerSaslProperties(kafkaServerSaslMechanisms, kafkaClientSaslMechanism))
   protected override val clientSaslProperties = Some(kafkaClientSaslProperties(kafkaClientSaslMechanism))
-  var adminClient: AdminClient = null
+  var adminClient: Admin = null
 
-  override def numBrokers = 1
+  override def brokerCount = 1
 
   @Before
   override def setUp(): Unit = {
@@ -47,11 +48,11 @@ class DelegationTokenRequestsTest extends BaseRequestTest with SaslSetup {
   }
 
   override def generateConfigs = {
-    val props = TestUtils.createBrokerConfigs(numBrokers, zkConnect,
+    val props = TestUtils.createBrokerConfigs(brokerCount, zkConnect,
       enableControlledShutdown = false,
       interBrokerSecurityProtocol = Some(securityProtocol),
       trustStoreFile = trustStoreFile, saslProperties = serverSaslProperties, enableToken = true)
-    props.foreach(propertyOverrides)
+    props.foreach(brokerPropertyOverrides)
     props.map(KafkaConfig.fromProps)
   }
 
diff --git a/core/src/test/scala/unit/kafka/server/DelegationTokenRequestsWithDisableTokenFeatureTest.scala b/core/src/test/scala/unit/kafka/server/DelegationTokenRequestsWithDisableTokenFeatureTest.scala
index 1203f839f895b..01c793c83b472 100644
--- a/core/src/test/scala/unit/kafka/server/DelegationTokenRequestsWithDisableTokenFeatureTest.scala
+++ b/core/src/test/scala/unit/kafka/server/DelegationTokenRequestsWithDisableTokenFeatureTest.scala
@@ -20,10 +20,11 @@ import java.util
 
 import kafka.api.{KafkaSasl, SaslSetup}
 import kafka.utils.{JaasTestUtils, TestUtils}
-import org.apache.kafka.clients.admin.{AdminClient, AdminClientConfig}
+import org.apache.kafka.clients.admin.{Admin, AdminClient, AdminClientConfig}
 import org.apache.kafka.common.errors.DelegationTokenDisabledException
 import org.apache.kafka.common.security.auth.SecurityProtocol
 import org.junit.{After, Before, Test}
+import org.scalatest.Assertions.intercept
 
 import scala.collection.JavaConverters._
 import scala.concurrent.ExecutionException
@@ -34,9 +35,9 @@ class DelegationTokenRequestsWithDisableTokenFeatureTest extends BaseRequestTest
   private val kafkaServerSaslMechanisms = List("PLAIN")
   protected override val serverSaslProperties = Some(kafkaServerSaslProperties(kafkaServerSaslMechanisms, kafkaClientSaslMechanism))
   protected override val clientSaslProperties = Some(kafkaClientSaslProperties(kafkaClientSaslMechanism))
-  var adminClient: AdminClient = null
+  var adminClient: Admin = null
 
-  override def numBrokers = 1
+  override def brokerCount = 1
 
   @Before
   override def setUp(): Unit = {
diff --git a/core/src/test/scala/unit/kafka/server/DeleteTopicsRequestTest.scala b/core/src/test/scala/unit/kafka/server/DeleteTopicsRequestTest.scala
index 4388e64631a5d..b9161a8d07b0c 100644
--- a/core/src/test/scala/unit/kafka/server/DeleteTopicsRequestTest.scala
+++ b/core/src/test/scala/unit/kafka/server/DeleteTopicsRequestTest.scala
@@ -19,52 +19,63 @@ package kafka.server
 
 import kafka.network.SocketServer
 import kafka.utils._
+import org.apache.kafka.common.message.DeleteTopicsRequestData
 import org.apache.kafka.common.protocol.{ApiKeys, Errors}
 import org.apache.kafka.common.requests.{DeleteTopicsRequest, DeleteTopicsResponse, MetadataRequest, MetadataResponse}
 import org.junit.Assert._
 import org.junit.Test
 
 import scala.collection.JavaConverters._
+import java.util.Collections
+import java.util.Arrays
 
 class DeleteTopicsRequestTest extends BaseRequestTest {
 
   @Test
-  def testValidDeleteTopicRequests() {
+  def testValidDeleteTopicRequests(): Unit = {
     val timeout = 10000
     // Single topic
     createTopic("topic-1", 1, 1)
-    validateValidDeleteTopicRequests(new DeleteTopicsRequest.Builder(Set("topic-1").asJava, timeout).build())
+    validateValidDeleteTopicRequests(new DeleteTopicsRequest.Builder(
+        new DeleteTopicsRequestData()
+          .setTopicNames(Arrays.asList("topic-1"))
+          .setTimeoutMs(timeout)).build())
     // Multi topic
     createTopic("topic-3", 5, 2)
     createTopic("topic-4", 1, 2)
-    validateValidDeleteTopicRequests(new DeleteTopicsRequest.Builder(Set("topic-3", "topic-4").asJava, timeout).build())
+    validateValidDeleteTopicRequests(new DeleteTopicsRequest.Builder(
+        new DeleteTopicsRequestData()
+          .setTopicNames(Arrays.asList("topic-3", "topic-4"))
+          .setTimeoutMs(timeout)).build())
   }
 
   private def validateValidDeleteTopicRequests(request: DeleteTopicsRequest): Unit = {
     val response = sendDeleteTopicsRequest(request)
-
-    val error = response.errors.values.asScala.find(_ != Errors.NONE)
-    assertTrue(s"There should be no errors, found ${response.errors.asScala}", error.isEmpty)
-
-    request.topics.asScala.foreach { topic =>
+    val error = response.errorCounts.asScala.find(_._1 != Errors.NONE)
+    assertTrue(s"There should be no errors, found ${response.data.responses.asScala}", error.isEmpty)
+    request.data.topicNames.asScala.foreach { topic =>
       validateTopicIsDeleted(topic)
     }
   }
 
   @Test
-  def testErrorDeleteTopicRequests() {
+  def testErrorDeleteTopicRequests(): Unit = {
     val timeout = 30000
     val timeoutTopic = "invalid-timeout"
 
     // Basic
-    validateErrorDeleteTopicRequests(new DeleteTopicsRequest.Builder(Set("invalid-topic").asJava, timeout).build(),
+    validateErrorDeleteTopicRequests(new DeleteTopicsRequest.Builder(
+        new DeleteTopicsRequestData()
+          .setTopicNames(Arrays.asList("invalid-topic"))
+          .setTimeoutMs(timeout)).build(),
       Map("invalid-topic" -> Errors.UNKNOWN_TOPIC_OR_PARTITION))
 
     // Partial
     createTopic("partial-topic-1", 1, 1)
-    validateErrorDeleteTopicRequests(new DeleteTopicsRequest.Builder(Set(
-      "partial-topic-1",
-      "partial-invalid-topic").asJava, timeout).build(),
+    validateErrorDeleteTopicRequests(new DeleteTopicsRequest.Builder(
+        new DeleteTopicsRequestData()
+          .setTopicNames(Arrays.asList("partial-topic-1", "partial-invalid-topic"))
+          .setTimeoutMs(timeout)).build(),
       Map(
         "partial-topic-1" -> Errors.NONE,
         "partial-invalid-topic" -> Errors.UNKNOWN_TOPIC_OR_PARTITION
@@ -74,7 +85,10 @@ class DeleteTopicsRequestTest extends BaseRequestTest {
     // Timeout
     createTopic(timeoutTopic, 5, 2)
     // Must be a 0ms timeout to avoid transient test failures. Even a timeout of 1ms has succeeded in the past.
-    validateErrorDeleteTopicRequests(new DeleteTopicsRequest.Builder(Set(timeoutTopic).asJava, 0).build(),
+    validateErrorDeleteTopicRequests(new DeleteTopicsRequest.Builder(
+        new DeleteTopicsRequestData()
+          .setTopicNames(Arrays.asList(timeoutTopic))
+          .setTimeoutMs(0)).build(),
       Map(timeoutTopic -> Errors.REQUEST_TIMED_OUT))
     // The topic should still get deleted eventually
     TestUtils.waitUntilTrue(() => !servers.head.metadataCache.contains(timeoutTopic), s"Topic $timeoutTopic is never deleted")
@@ -83,11 +97,13 @@ class DeleteTopicsRequestTest extends BaseRequestTest {
 
   private def validateErrorDeleteTopicRequests(request: DeleteTopicsRequest, expectedResponse: Map[String, Errors]): Unit = {
     val response = sendDeleteTopicsRequest(request)
-    val errors = response.errors.asScala
-    assertEquals("The response size should match", expectedResponse.size, response.errors.size)
+    val errors = response.data.responses
+
+    val errorCount = response.errorCounts().asScala.foldLeft(0)(_+_._2)
+    assertEquals("The response size should match", expectedResponse.size, errorCount)
 
     expectedResponse.foreach { case (topic, expectedError) =>
-      assertEquals("The response error should match", expectedResponse(topic), errors(topic))
+      assertEquals("The response error should match", expectedResponse(topic).code, errors.find(topic).errorCode)
       // If no error validate the topic was deleted
       if (expectedError == Errors.NONE) {
         validateTopicIsDeleted(topic)
@@ -96,12 +112,15 @@ class DeleteTopicsRequestTest extends BaseRequestTest {
   }
 
   @Test
-  def testNotController() {
-    val request = new DeleteTopicsRequest.Builder(Set("not-controller").asJava, 1000).build()
+  def testNotController(): Unit = {
+    val request = new DeleteTopicsRequest.Builder(
+        new DeleteTopicsRequestData()
+          .setTopicNames(Collections.singletonList("not-controller"))
+          .setTimeoutMs(1000)).build()
     val response = sendDeleteTopicsRequest(request, notControllerSocketServer)
 
-    val error = response.errors.asScala.head._2
-    assertEquals("Expected controller error when routed incorrectly",  Errors.NOT_CONTROLLER, error)
+    val error = response.data.responses().find("not-controller").errorCode()
+    assertEquals("Expected controller error when routed incorrectly",  Errors.NOT_CONTROLLER.code, error)
   }
 
   private def validateTopicIsDeleted(topic: String): Unit = {
diff --git a/core/src/test/scala/unit/kafka/server/DeleteTopicsRequestWithDeletionDisabledTest.scala b/core/src/test/scala/unit/kafka/server/DeleteTopicsRequestWithDeletionDisabledTest.scala
index 20e30c08a9f4e..ca0ccc0212aca 100644
--- a/core/src/test/scala/unit/kafka/server/DeleteTopicsRequestWithDeletionDisabledTest.scala
+++ b/core/src/test/scala/unit/kafka/server/DeleteTopicsRequestWithDeletionDisabledTest.scala
@@ -19,36 +19,43 @@ package kafka.server
 
 import kafka.network.SocketServer
 import kafka.utils._
+import org.apache.kafka.common.message.DeleteTopicsRequestData
 import org.apache.kafka.common.protocol.{ApiKeys, Errors}
 import org.apache.kafka.common.requests.{DeleteTopicsRequest, DeleteTopicsResponse}
 import org.junit.Assert._
 import org.junit.Test
 
-import scala.collection.JavaConverters._
+import java.util.Collections
 
 class DeleteTopicsRequestWithDeletionDisabledTest extends BaseRequestTest {
 
-  override def numBrokers: Int = 1
+  override def brokerCount: Int = 1
 
   override def generateConfigs = {
-    val props = TestUtils.createBrokerConfigs(numBrokers, zkConnect,
+    val props = TestUtils.createBrokerConfigs(brokerCount, zkConnect,
       enableControlledShutdown = false, enableDeleteTopic = false,
       interBrokerSecurityProtocol = Some(securityProtocol),
       trustStoreFile = trustStoreFile, saslProperties = serverSaslProperties, logDirCount = logDirCount)
-    props.foreach(propertyOverrides)
+    props.foreach(brokerPropertyOverrides)
     props.map(KafkaConfig.fromProps)
   }
 
   @Test
-  def testDeleteRecordsRequest() {
+  def testDeleteRecordsRequest(): Unit = {
     val topic = "topic-1"
-    val request = new DeleteTopicsRequest.Builder(Set(topic).asJava, 1000).build()
+    val request = new DeleteTopicsRequest.Builder(
+        new DeleteTopicsRequestData()
+          .setTopicNames(Collections.singletonList(topic))
+          .setTimeoutMs(1000)).build()
     val response = sendDeleteTopicsRequest(request)
-    assertEquals(Errors.TOPIC_DELETION_DISABLED, response.errors.get(topic))
+    assertEquals(Errors.TOPIC_DELETION_DISABLED.code, response.data.responses.find(topic).errorCode)
 
-    val v2request = new DeleteTopicsRequest.Builder(Set(topic).asJava, 1000).build(2)
+    val v2request = new DeleteTopicsRequest.Builder(
+        new DeleteTopicsRequestData()
+        .setTopicNames(Collections.singletonList(topic))
+        .setTimeoutMs(1000)).build(2)
     val v2response = sendDeleteTopicsRequest(v2request)
-    assertEquals(Errors.INVALID_REQUEST, v2response.errors.get(topic))
+    assertEquals(Errors.INVALID_REQUEST.code, v2response.data.responses.find(topic).errorCode)
   }
 
   private def sendDeleteTopicsRequest(request: DeleteTopicsRequest, socketServer: SocketServer = controllerSocketServer): DeleteTopicsResponse = {
diff --git a/core/src/test/scala/unit/kafka/server/DescribeLogDirsRequestTest.scala b/core/src/test/scala/unit/kafka/server/DescribeLogDirsRequestTest.scala
index 5a0244b75f38b..5e2707a646625 100644
--- a/core/src/test/scala/unit/kafka/server/DescribeLogDirsRequestTest.scala
+++ b/core/src/test/scala/unit/kafka/server/DescribeLogDirsRequestTest.scala
@@ -27,7 +27,7 @@ import java.io.File
 
 class DescribeLogDirsRequestTest extends BaseRequestTest {
   override val logDirCount = 2
-  override val numBrokers: Int = 1
+  override val brokerCount: Int = 1
 
   val topic = "topic"
   val partitionNum = 2
@@ -57,7 +57,8 @@ class DescribeLogDirsRequestTest extends BaseRequestTest {
     val log1 = servers.head.logManager.getLog(tp1).get
     assertEquals(log0.size, replicaInfo0.size)
     assertEquals(log1.size, replicaInfo1.size)
-    assertTrue(servers.head.logManager.getLog(tp0).get.logEndOffset > 0)
+    val logEndOffset = servers.head.logManager.getLog(tp0).get.logEndOffset
+    assertTrue(s"LogEndOffset '$logEndOffset' should be > 0", logEndOffset > 0)
     assertEquals(servers.head.replicaManager.getLogEndOffsetLag(tp0, log0.logEndOffset, false), replicaInfo0.offsetLag)
     assertEquals(servers.head.replicaManager.getLogEndOffsetLag(tp1, log1.logEndOffset, false), replicaInfo1.offsetLag)
   }
diff --git a/core/src/test/scala/unit/kafka/server/DynamicBrokerConfigTest.scala b/core/src/test/scala/unit/kafka/server/DynamicBrokerConfigTest.scala
index 45ef18f5187f0..31df6c772b322 100755
--- a/core/src/test/scala/unit/kafka/server/DynamicBrokerConfigTest.scala
+++ b/core/src/test/scala/unit/kafka/server/DynamicBrokerConfigTest.scala
@@ -17,23 +17,26 @@
 
 package kafka.server
 
-import java.util
+import java.{lang, util}
 import java.util.Properties
+import java.util.concurrent.CompletionStage
 
 import kafka.utils.TestUtils
 import kafka.zk.KafkaZkClient
-import org.apache.kafka.common.Reconfigurable
+import org.apache.kafka.common.{Endpoint, Reconfigurable}
+import org.apache.kafka.common.acl.{AclBinding, AclBindingFilter}
 import org.apache.kafka.common.config.types.Password
 import org.apache.kafka.common.config.{ConfigException, SslConfigs}
+import org.apache.kafka.server.authorizer._
 import org.easymock.EasyMock
 import org.junit.Assert._
 import org.junit.Test
-import org.scalatest.junit.JUnitSuite
+import org.scalatest.Assertions.intercept
 
 import scala.collection.JavaConverters._
 import scala.collection.Set
 
-class DynamicBrokerConfigTest extends JUnitSuite {
+class DynamicBrokerConfigTest {
 
   @Test
   def testConfigUpdate(): Unit = {
@@ -186,9 +189,15 @@ class DynamicBrokerConfigTest extends JUnitSuite {
     //test invalid address
     verifyConfigUpdate(KafkaConfig.MaxConnectionsPerIpOverridesProp, "hostName#:100", perBrokerConfig = true,
       expectFailure = true)
+
+    verifyConfigUpdate(KafkaConfig.MaxConnectionsProp, "100", perBrokerConfig = true, expectFailure = false)
+    verifyConfigUpdate(KafkaConfig.MaxConnectionsProp, "100", perBrokerConfig = false, expectFailure = false)
+    val listenerMaxConnectionsProp = s"listener.name.external.${KafkaConfig.MaxConnectionsProp}"
+    verifyConfigUpdate(listenerMaxConnectionsProp, "10", perBrokerConfig = true, expectFailure = false)
+    verifyConfigUpdate(listenerMaxConnectionsProp, "10", perBrokerConfig = false, expectFailure = false)
   }
 
-  private def verifyConfigUpdate(name: String, value: Object, perBrokerConfig: Boolean, expectFailure: Boolean) {
+  private def verifyConfigUpdate(name: String, value: Object, perBrokerConfig: Boolean, expectFailure: Boolean): Unit = {
     val configProps = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 8181)
     configProps.put(KafkaConfig.PasswordEncoderSecretProp, "broker.secret")
     val config = KafkaConfig(configProps)
@@ -315,6 +324,43 @@ class DynamicBrokerConfigTest extends JUnitSuite {
     dynamicListenerConfig.validateReconfiguration(newConfig)
   }
 
+  @Test
+  def testAuthorizerConfig(): Unit = {
+    val props = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 9092)
+    val oldConfig =  KafkaConfig.fromProps(props)
+    val kafkaServer: KafkaServer = EasyMock.createMock(classOf[kafka.server.KafkaServer])
+
+    class TestAuthorizer extends Authorizer with Reconfigurable {
+      @volatile var superUsers = ""
+      override def start(serverInfo: AuthorizerServerInfo): util.Map[Endpoint, _ <: CompletionStage[Void]] = Map.empty.asJava
+      override def authorize(requestContext: AuthorizableRequestContext, actions: util.List[Action]): util.List[AuthorizationResult] = null
+      override def createAcls(requestContext: AuthorizableRequestContext, aclBindings: util.List[AclBinding]): util.List[_ <: CompletionStage[AclCreateResult]] = null
+      override def deleteAcls(requestContext: AuthorizableRequestContext, aclBindingFilters: util.List[AclBindingFilter]): util.List[_ <: CompletionStage[AclDeleteResult]] = null
+      override def acls(filter: AclBindingFilter): lang.Iterable[AclBinding] = null
+      override def close(): Unit = {}
+      override def configure(configs: util.Map[String, _]): Unit = {}
+      override def reconfigurableConfigs(): util.Set[String] = Set("super.users").asJava
+      override def validateReconfiguration(configs: util.Map[String, _]): Unit = {}
+      override def reconfigure(configs: util.Map[String, _]): Unit = {
+        superUsers = configs.get("super.users").toString
+      }
+    }
+
+    val authorizer = new TestAuthorizer
+    EasyMock.expect(kafkaServer.config).andReturn(oldConfig).anyTimes()
+    EasyMock.expect(kafkaServer.authorizer).andReturn(Some(authorizer)).anyTimes()
+    EasyMock.replay(kafkaServer)
+    try {
+      kafkaServer.config.dynamicConfig.addReconfigurables(kafkaServer)
+    } catch {
+      case _: Throwable => // We are only testing authorizer reconfiguration, ignore any exceptions due to incomplete mock
+    }
+
+    props.put("super.users", "User:admin")
+    kafkaServer.config.dynamicConfig.updateBrokerConfig(0, props)
+    assertEquals("User:admin", authorizer.superUsers)
+  }
+
   @Test
   def testSynonyms(): Unit = {
     assertEquals(List("listener.name.secure.ssl.keystore.type", "ssl.keystore.type"),
diff --git a/core/src/test/scala/unit/kafka/server/DynamicConfigChangeTest.scala b/core/src/test/scala/unit/kafka/server/DynamicConfigChangeTest.scala
index f5cc134297024..04afa41f2dd8e 100644
--- a/core/src/test/scala/unit/kafka/server/DynamicConfigChangeTest.scala
+++ b/core/src/test/scala/unit/kafka/server/DynamicConfigChangeTest.scala
@@ -31,14 +31,14 @@ import kafka.admin.AdminOperationException
 import kafka.zk.ConfigEntityChangeNotificationZNode
 import org.apache.kafka.common.TopicPartition
 
-import scala.collection.Map
+import scala.collection.{Map, Seq}
 import scala.collection.JavaConverters._
 
 class DynamicConfigChangeTest extends KafkaServerTestHarness {
   def generateConfigs = List(KafkaConfig.fromProps(TestUtils.createBrokerConfig(0, zkConnect)))
 
   @Test
-  def testConfigChange() {
+  def testConfigChange(): Unit = {
     assertTrue("Should contain a ConfigHandler for topics",
       this.servers.head.dynamicConfigHandlers.contains(ConfigType.Topic))
     val oldVal: java.lang.Long = 100000L
@@ -60,7 +60,7 @@ class DynamicConfigChangeTest extends KafkaServerTestHarness {
   }
 
   @Test
-  def testDynamicTopicConfigChange() {
+  def testDynamicTopicConfigChange(): Unit = {
     val tp = new TopicPartition("test", 0)
     val oldSegmentSize = 1000
     val logProps = new Properties()
@@ -86,7 +86,7 @@ class DynamicConfigChangeTest extends KafkaServerTestHarness {
     assertTrue("Log segment size change not applied", log.logSegments.forall(_.size > 1000))
   }
 
-  private def testQuotaConfigChange(user: String, clientId: String, rootEntityType: String, configEntityName: String) {
+  private def testQuotaConfigChange(user: String, clientId: String, rootEntityType: String, configEntityName: String): Unit = {
     assertTrue("Should contain a ConfigHandler for " + rootEntityType ,
                this.servers.head.dynamicConfigHandlers.contains(rootEntityType))
     val props = new Properties()
@@ -129,37 +129,37 @@ class DynamicConfigChangeTest extends KafkaServerTestHarness {
   }
 
   @Test
-  def testClientIdQuotaConfigChange() {
+  def testClientIdQuotaConfigChange(): Unit = {
     testQuotaConfigChange("ANONYMOUS", "testClient", ConfigType.Client, "testClient")
   }
 
   @Test
-  def testUserQuotaConfigChange() {
+  def testUserQuotaConfigChange(): Unit = {
     testQuotaConfigChange("ANONYMOUS", "testClient", ConfigType.User, "ANONYMOUS")
   }
 
   @Test
-  def testUserClientIdQuotaChange() {
+  def testUserClientIdQuotaChange(): Unit = {
     testQuotaConfigChange("ANONYMOUS", "testClient", ConfigType.User, "ANONYMOUS/clients/testClient")
   }
 
   @Test
-  def testDefaultClientIdQuotaConfigChange() {
+  def testDefaultClientIdQuotaConfigChange(): Unit = {
     testQuotaConfigChange("ANONYMOUS", "testClient", ConfigType.Client, "")
   }
 
   @Test
-  def testDefaultUserQuotaConfigChange() {
+  def testDefaultUserQuotaConfigChange(): Unit = {
     testQuotaConfigChange("ANONYMOUS", "testClient", ConfigType.User, "")
   }
 
   @Test
-  def testDefaultUserClientIdQuotaConfigChange() {
+  def testDefaultUserClientIdQuotaConfigChange(): Unit = {
     testQuotaConfigChange("ANONYMOUS", "testClient", ConfigType.User, "/clients/")
   }
 
   @Test
-  def testQuotaInitialization() {
+  def testQuotaInitialization(): Unit = {
     val server = servers.head
     val clientIdProps = new Properties()
     server.shutdown()
@@ -190,7 +190,7 @@ class DynamicConfigChangeTest extends KafkaServerTestHarness {
   }
 
   @Test
-  def testConfigChangeOnNonExistingTopic() {
+  def testConfigChangeOnNonExistingTopic(): Unit = {
     val topic = TestUtils.tempTopic
     try {
       val logProps = new Properties()
@@ -302,7 +302,7 @@ class DynamicConfigChangeTest extends KafkaServerTestHarness {
   }
 
   @Test
-  def shouldParseRegardlessOfWhitespaceAroundValues() {
+  def shouldParseRegardlessOfWhitespaceAroundValues(): Unit = {
     val configHandler: TopicConfigHandler = new TopicConfigHandler(null, null, null, null)
     assertEquals(AllReplicas, parse(configHandler, "* "))
     assertEquals(Seq(), parse(configHandler, " "))
diff --git a/core/src/test/scala/unit/kafka/server/DynamicConfigTest.scala b/core/src/test/scala/unit/kafka/server/DynamicConfigTest.scala
index cf0984960f174..14e3ef5754e97 100644
--- a/core/src/test/scala/unit/kafka/server/DynamicConfigTest.scala
+++ b/core/src/test/scala/unit/kafka/server/DynamicConfigTest.scala
@@ -26,23 +26,23 @@ class DynamicConfigTest  extends ZooKeeperTestHarness {
   private final val someValue: String = "some interesting value"
 
   @Test(expected = classOf[IllegalArgumentException])
-  def shouldFailWhenChangingClientIdUnknownConfig() {
+  def shouldFailWhenChangingClientIdUnknownConfig(): Unit = {
     adminZkClient.changeClientIdConfig("ClientId", propsWith(nonExistentConfig, someValue))
   }
 
   @Test(expected = classOf[IllegalArgumentException])
-  def shouldFailWhenChangingUserUnknownConfig() {
+  def shouldFailWhenChangingUserUnknownConfig(): Unit = {
     adminZkClient.changeUserOrUserClientIdConfig("UserId", propsWith(nonExistentConfig, someValue))
   }
 
   @Test(expected = classOf[ConfigException])
-  def shouldFailLeaderConfigsWithInvalidValues() {
+  def shouldFailLeaderConfigsWithInvalidValues(): Unit = {
     adminZkClient.changeBrokerConfig(Seq(0),
       propsWith(DynamicConfig.Broker.LeaderReplicationThrottledRateProp, "-100"))
   }
 
   @Test(expected = classOf[ConfigException])
-  def shouldFailFollowerConfigsWithInvalidValues() {
+  def shouldFailFollowerConfigsWithInvalidValues(): Unit = {
     adminZkClient.changeBrokerConfig(Seq(0),
       propsWith(DynamicConfig.Broker.FollowerReplicationThrottledRateProp, "-100"))
   }
diff --git a/core/src/test/scala/unit/kafka/server/EdgeCaseRequestTest.scala b/core/src/test/scala/unit/kafka/server/EdgeCaseRequestTest.scala
index a42610831641e..b8374f2435d6c 100755
--- a/core/src/test/scala/unit/kafka/server/EdgeCaseRequestTest.scala
+++ b/core/src/test/scala/unit/kafka/server/EdgeCaseRequestTest.scala
@@ -50,7 +50,7 @@ class EdgeCaseRequestTest extends KafkaServerTestHarness {
     new Socket("localhost", s.boundPort(ListenerName.forSecurityProtocol(protocol)))
   }
 
-  private def sendRequest(socket: Socket, request: Array[Byte], id: Option[Short] = None) {
+  private def sendRequest(socket: Socket, request: Array[Byte], id: Option[Short] = None): Unit = {
     val outgoing = new DataOutputStream(socket.getOutputStream)
     id match {
       case Some(id) =>
@@ -98,7 +98,7 @@ class EdgeCaseRequestTest extends KafkaServerTestHarness {
     buffer.array()
   }
 
-  private def verifyDisconnect(request: Array[Byte]) {
+  private def verifyDisconnect(request: Array[Byte]): Unit = {
     val plainSocket = connect()
     try {
       sendRequest(plainSocket, requestHeaderBytes(-1, 0))
@@ -109,14 +109,14 @@ class EdgeCaseRequestTest extends KafkaServerTestHarness {
   }
 
   @Test
-  def testProduceRequestWithNullClientId() {
+  def testProduceRequestWithNullClientId(): Unit = {
     val topic = "topic"
     val topicPartition = new TopicPartition(topic, 0)
     val correlationId = -1
     createTopic(topic, numPartitions = 1, replicationFactor = 1)
 
     val version = ApiKeys.PRODUCE.latestVersion: Short
-    val serializedBytes = {
+    val (serializedBytes, responseHeaderVersion) = {
       val headerBytes = requestHeaderBytes(ApiKeys.PRODUCE.id, version, null,
         correlationId)
       val records = MemoryRecords.withRecords(CompressionType.NONE, new SimpleRecord("message".getBytes))
@@ -124,13 +124,13 @@ class EdgeCaseRequestTest extends KafkaServerTestHarness {
       val byteBuffer = ByteBuffer.allocate(headerBytes.length + request.toStruct.sizeOf)
       byteBuffer.put(headerBytes)
       request.toStruct.writeTo(byteBuffer)
-      byteBuffer.array()
+      (byteBuffer.array(), request.api.responseHeaderVersion(version))
     }
 
     val response = requestAndReceive(serializedBytes)
 
     val responseBuffer = ByteBuffer.wrap(response)
-    val responseHeader = ResponseHeader.parse(responseBuffer)
+    val responseHeader = ResponseHeader.parse(responseBuffer, responseHeaderVersion)
     val produceResponse = ProduceResponse.parse(responseBuffer, version)
 
     assertEquals("The response should parse completely", 0, responseBuffer.remaining)
@@ -143,22 +143,22 @@ class EdgeCaseRequestTest extends KafkaServerTestHarness {
   }
 
   @Test
-  def testHeaderOnlyRequest() {
+  def testHeaderOnlyRequest(): Unit = {
     verifyDisconnect(requestHeaderBytes(ApiKeys.PRODUCE.id, 1))
   }
 
   @Test
-  def testInvalidApiKeyRequest() {
+  def testInvalidApiKeyRequest(): Unit = {
     verifyDisconnect(requestHeaderBytes(-1, 0))
   }
 
   @Test
-  def testInvalidApiVersionRequest() {
+  def testInvalidApiVersionRequest(): Unit = {
     verifyDisconnect(requestHeaderBytes(ApiKeys.PRODUCE.id, -1))
   }
 
   @Test
-  def testMalformedHeaderRequest() {
+  def testMalformedHeaderRequest(): Unit = {
     val serializedBytes = {
       // Only send apiKey and apiVersion
       val buffer = ByteBuffer.allocate(
diff --git a/core/src/test/scala/unit/kafka/server/FetchRequestDownConversionConfigTest.scala b/core/src/test/scala/unit/kafka/server/FetchRequestDownConversionConfigTest.scala
index 673abe693c9c8..4ae7d936975d4 100644
--- a/core/src/test/scala/unit/kafka/server/FetchRequestDownConversionConfigTest.scala
+++ b/core/src/test/scala/unit/kafka/server/FetchRequestDownConversionConfigTest.scala
@@ -32,7 +32,7 @@ import org.junit.Test
 
 class FetchRequestDownConversionConfigTest extends BaseRequestTest {
   private var producer: KafkaProducer[String, String] = null
-  override def numBrokers: Int = 1
+  override def brokerCount: Int = 1
 
   override def setUp(): Unit = {
     super.setUp()
@@ -45,8 +45,8 @@ class FetchRequestDownConversionConfigTest extends BaseRequestTest {
     super.tearDown()
   }
 
-  override protected def propertyOverrides(properties: Properties): Unit = {
-    super.propertyOverrides(properties)
+  override protected def brokerPropertyOverrides(properties: Properties): Unit = {
+    super.brokerPropertyOverrides(properties)
     properties.put(KafkaConfig.LogMessageDownConversionEnableProp, "false")
   }
 
diff --git a/core/src/test/scala/unit/kafka/server/FetchRequestTest.scala b/core/src/test/scala/unit/kafka/server/FetchRequestTest.scala
index 86f3314d62401..895d9e55401e3 100644
--- a/core/src/test/scala/unit/kafka/server/FetchRequestTest.scala
+++ b/core/src/test/scala/unit/kafka/server/FetchRequestTest.scala
@@ -34,6 +34,7 @@ import org.junit.Assert._
 import org.junit.Test
 
 import scala.collection.JavaConverters._
+import scala.collection.Seq
 import scala.util.Random
 
 /**
@@ -44,7 +45,7 @@ class FetchRequestTest extends BaseRequestTest {
 
   private var producer: KafkaProducer[String, String] = null
 
-  override def tearDown() {
+  override def tearDown(): Unit = {
     if (producer != null)
       producer.close()
     super.tearDown()
@@ -205,7 +206,7 @@ class FetchRequestTest extends BaseRequestTest {
       Seq(topicPartition))).build()
     val fetchResponse = sendFetchRequest(nonReplicaId, fetchRequest)
     val partitionData = fetchResponse.responseData.get(topicPartition)
-    assertEquals(Errors.NOT_LEADER_FOR_PARTITION, partitionData.error)
+    assertEquals(Errors.REPLICA_NOT_AVAILABLE, partitionData.error)
   }
 
   @Test
@@ -238,8 +239,8 @@ class FetchRequestTest extends BaseRequestTest {
 
     // Check follower error codes
     val followerId = TestUtils.findFollowerId(topicPartition, servers)
-    assertResponseErrorForEpoch(Errors.NOT_LEADER_FOR_PARTITION, followerId, Optional.empty())
-    assertResponseErrorForEpoch(Errors.NOT_LEADER_FOR_PARTITION, followerId, Optional.of(secondLeaderEpoch))
+    assertResponseErrorForEpoch(Errors.NONE, followerId, Optional.empty())
+    assertResponseErrorForEpoch(Errors.NONE, followerId, Optional.of(secondLeaderEpoch))
     assertResponseErrorForEpoch(Errors.UNKNOWN_LEADER_EPOCH, followerId, Optional.of(secondLeaderEpoch + 1))
     assertResponseErrorForEpoch(Errors.FENCED_LEADER_EPOCH, followerId, Optional.of(secondLeaderEpoch - 1))
   }
@@ -257,6 +258,7 @@ class FetchRequestTest extends BaseRequestTest {
     val batchSize = 4 * msgValueLen
     val producer = TestUtils.createProducer(TestUtils.getBrokerListStrFromServers(servers),
       lingerMs = Int.MaxValue,
+      deliveryTimeoutMs = Int.MaxValue,
       batchSize = batchSize,
       keySerializer = new StringSerializer,
       valueSerializer = new ByteArraySerializer)
@@ -278,7 +280,7 @@ class FetchRequestTest extends BaseRequestTest {
 
       val socket = connect(brokerSocketServer(leaderId))
       try {
-        send(fetchRequest, ApiKeys.FETCH, socket)
+        send(fetchRequest, ApiKeys.FETCH, socket, fetchRequest.version())
         if (closeAfterPartialResponse) {
           // read some data to ensure broker has muted this channel and then close socket
           val size = new DataInputStream(socket.getInputStream).readInt()
@@ -289,7 +291,7 @@ class FetchRequestTest extends BaseRequestTest {
               size > maxPartitionBytes - batchSize)
           None
         } else {
-          Some(FetchResponse.parse(receive(socket), version))
+          Some(FetchResponse.parse(receive(socket, ApiKeys.FETCH.responseHeaderVersion(version)), version))
         }
       } finally {
         socket.close()
@@ -555,7 +557,7 @@ class FetchRequestTest extends BaseRequestTest {
   }
 
   private def records(partitionData: FetchResponse.PartitionData[MemoryRecords]): Seq[Record] = {
-    partitionData.records.records.asScala.toIndexedSeq
+    partitionData.records.records.asScala.toBuffer
   }
 
   private def checkFetchResponse(expectedPartitions: Seq[TopicPartition], fetchResponse: FetchResponse[MemoryRecords],
@@ -573,7 +575,7 @@ class FetchRequestTest extends BaseRequestTest {
       val records = partitionData.records
       responseBufferSize += records.sizeInBytes
 
-      val batches = records.batches.asScala.toIndexedSeq
+      val batches = records.batches.asScala.toBuffer
       assertTrue(batches.size < numMessagesPerPartition)
       val batchesSize = batches.map(_.sizeInBytes).sum
       responseSize += batchesSize
diff --git a/core/src/test/scala/unit/kafka/server/HighwatermarkPersistenceTest.scala b/core/src/test/scala/unit/kafka/server/HighwatermarkPersistenceTest.scala
index 61cbd2c94f488..e248f3b678ad2 100755
--- a/core/src/test/scala/unit/kafka/server/HighwatermarkPersistenceTest.scala
+++ b/core/src/test/scala/unit/kafka/server/HighwatermarkPersistenceTest.scala
@@ -24,12 +24,13 @@ import org.apache.kafka.common.utils.Utils
 import org.easymock.EasyMock
 import org.junit._
 import org.junit.Assert._
-import kafka.cluster.Replica
 import kafka.utils.{KafkaScheduler, MockTime, TestUtils}
 import kafka.zk.KafkaZkClient
 import java.util.concurrent.atomic.AtomicBoolean
 
+import kafka.cluster.Partition
 import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.record.SimpleRecord
 
 class HighwatermarkPersistenceTest {
 
@@ -47,13 +48,13 @@ class HighwatermarkPersistenceTest {
   }
 
   @After
-  def teardown() {
+  def teardown(): Unit = {
     for (manager <- logManagers; dir <- manager.liveLogDirs)
       Utils.delete(dir)
   }
 
   @Test
-  def testHighWatermarkPersistenceSinglePartition() {
+  def testHighWatermarkPersistenceSinglePartition(): Unit = {
     // mock zkclient
     EasyMock.replay(zkClient)
 
@@ -72,21 +73,24 @@ class HighwatermarkPersistenceTest {
       var fooPartition0Hw = hwmFor(replicaManager, topic, 0)
       assertEquals(0L, fooPartition0Hw)
       val tp0 = new TopicPartition(topic, 0)
-      val partition0 = replicaManager.getOrCreatePartition(tp0)
+      val partition0 = replicaManager.createPartition(tp0)
       // create leader and follower replicas
-      val log0 = logManagers.head.getOrCreateLog(new TopicPartition(topic, 0), LogConfig())
-      val leaderReplicaPartition0 = new Replica(configs.head.brokerId, tp0, time, 0, Some(log0))
-      partition0.addReplicaIfNotExists(leaderReplicaPartition0)
-      val followerReplicaPartition0 = new Replica(configs.last.brokerId, tp0, time)
-      partition0.addReplicaIfNotExists(followerReplicaPartition0)
+      val log0 = logManagers.head.getOrCreateLog(new TopicPartition(topic, 0), () => LogConfig())
+      partition0.setLog(log0, isFutureLog = false)
+
+      partition0.updateAssignmentAndIsr(
+        assignment = Seq(configs.head.brokerId, configs.last.brokerId),
+        isr = Set(configs.head.brokerId)
+      )
+
       replicaManager.checkpointHighWatermarks()
       fooPartition0Hw = hwmFor(replicaManager, topic, 0)
-      assertEquals(leaderReplicaPartition0.highWatermark.messageOffset, fooPartition0Hw)
+      assertEquals(log0.highWatermark, fooPartition0Hw)
       // set the high watermark for local replica
-      partition0.localReplica.get.highWatermark = new LogOffsetMetadata(5L)
+      partition0.localLogOrException.updateHighWatermark(5L)
       replicaManager.checkpointHighWatermarks()
       fooPartition0Hw = hwmFor(replicaManager, topic, 0)
-      assertEquals(leaderReplicaPartition0.highWatermark.messageOffset, fooPartition0Hw)
+      assertEquals(log0.highWatermark, fooPartition0Hw)
       EasyMock.verify(zkClient)
     } finally {
       // shutdown the replica manager upon test completion
@@ -97,7 +101,7 @@ class HighwatermarkPersistenceTest {
   }
 
   @Test
-  def testHighWatermarkPersistenceMultiplePartitions() {
+  def testHighWatermarkPersistenceMultiplePartitions(): Unit = {
     val topic1 = "foo1"
     val topic2 = "foo2"
     // mock zkclient
@@ -117,38 +121,39 @@ class HighwatermarkPersistenceTest {
       var topic1Partition0Hw = hwmFor(replicaManager, topic1, 0)
       assertEquals(0L, topic1Partition0Hw)
       val t1p0 = new TopicPartition(topic1, 0)
-      val topic1Partition0 = replicaManager.getOrCreatePartition(t1p0)
+      val topic1Partition0 = replicaManager.createPartition(t1p0)
       // create leader log
-      val topic1Log0 = logManagers.head.getOrCreateLog(t1p0, LogConfig())
+      val topic1Log0 = logManagers.head.getOrCreateLog(t1p0, () => LogConfig())
       // create a local replica for topic1
-      val leaderReplicaTopic1Partition0 = new Replica(configs.head.brokerId, t1p0, time, 0, Some(topic1Log0))
-      topic1Partition0.addReplicaIfNotExists(leaderReplicaTopic1Partition0)
+      topic1Partition0.setLog(topic1Log0, isFutureLog = false)
       replicaManager.checkpointHighWatermarks()
       topic1Partition0Hw = hwmFor(replicaManager, topic1, 0)
-      assertEquals(leaderReplicaTopic1Partition0.highWatermark.messageOffset, topic1Partition0Hw)
+      assertEquals(topic1Log0.highWatermark, topic1Partition0Hw)
       // set the high watermark for local replica
-      topic1Partition0.localReplica.get.highWatermark = new LogOffsetMetadata(5L)
+      append(topic1Partition0, count = 5)
+      topic1Partition0.localLogOrException.updateHighWatermark(5L)
       replicaManager.checkpointHighWatermarks()
       topic1Partition0Hw = hwmFor(replicaManager, topic1, 0)
-      assertEquals(5L, leaderReplicaTopic1Partition0.highWatermark.messageOffset)
+      assertEquals(5L, topic1Log0.highWatermark)
       assertEquals(5L, topic1Partition0Hw)
       // add another partition and set highwatermark
       val t2p0 = new TopicPartition(topic2, 0)
-      val topic2Partition0 = replicaManager.getOrCreatePartition(t2p0)
+      val topic2Partition0 = replicaManager.createPartition(t2p0)
       // create leader log
-      val topic2Log0 = logManagers.head.getOrCreateLog(t2p0, LogConfig())
+      val topic2Log0 = logManagers.head.getOrCreateLog(t2p0, () => LogConfig())
       // create a local replica for topic2
-      val leaderReplicaTopic2Partition0 =  new Replica(configs.head.brokerId, t2p0, time, 0, Some(topic2Log0))
-      topic2Partition0.addReplicaIfNotExists(leaderReplicaTopic2Partition0)
+      topic2Partition0.setLog(topic2Log0, isFutureLog = false)
       replicaManager.checkpointHighWatermarks()
       var topic2Partition0Hw = hwmFor(replicaManager, topic2, 0)
-      assertEquals(leaderReplicaTopic2Partition0.highWatermark.messageOffset, topic2Partition0Hw)
+      assertEquals(topic2Log0.highWatermark, topic2Partition0Hw)
       // set the highwatermark for local replica
-      topic2Partition0.localReplica.get.highWatermark = new LogOffsetMetadata(15L)
-      assertEquals(15L, leaderReplicaTopic2Partition0.highWatermark.messageOffset)
+      append(topic2Partition0, count = 15)
+      topic2Partition0.localLogOrException.updateHighWatermark(15L)
+      assertEquals(15L, topic2Log0.highWatermark)
       // change the highwatermark for topic1
-      topic1Partition0.localReplica.get.highWatermark = new LogOffsetMetadata(10L)
-      assertEquals(10L, leaderReplicaTopic1Partition0.highWatermark.messageOffset)
+      append(topic1Partition0, count = 5)
+      topic1Partition0.localLogOrException.updateHighWatermark(10L)
+      assertEquals(10L, topic1Log0.highWatermark)
       replicaManager.checkpointHighWatermarks()
       // verify checkpointed hw for topic 2
       topic2Partition0Hw = hwmFor(replicaManager, topic2, 0)
@@ -165,9 +170,13 @@ class HighwatermarkPersistenceTest {
     }
   }
 
-  def hwmFor(replicaManager: ReplicaManager, topic: String, partition: Int): Long = {
+  private def append(partition: Partition, count: Int): Unit = {
+    val records = TestUtils.records((0 to count).map(i => new SimpleRecord(s"$i".getBytes)))
+    partition.localLogOrException.appendAsLeader(records, leaderEpoch = 0)
+  }
+
+  private def hwmFor(replicaManager: ReplicaManager, topic: String, partition: Int): Long = {
     replicaManager.highWatermarkCheckpoints(new File(replicaManager.config.logDirs.head).getAbsolutePath).read.getOrElse(
       new TopicPartition(topic, partition), 0L)
   }
-
 }
diff --git a/core/src/test/scala/unit/kafka/server/ISRExpirationTest.scala b/core/src/test/scala/unit/kafka/server/ISRExpirationTest.scala
deleted file mode 100644
index 006067e4b35b2..0000000000000
--- a/core/src/test/scala/unit/kafka/server/ISRExpirationTest.scala
+++ /dev/null
@@ -1,266 +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.server
-
-import java.io.File
-import java.util.Properties
-import java.util.concurrent.atomic.AtomicBoolean
-
-import kafka.cluster.{Partition, Replica}
-import kafka.log.{Log, LogManager}
-import kafka.utils._
-import org.apache.kafka.common.TopicPartition
-import org.apache.kafka.common.metrics.Metrics
-import org.apache.kafka.common.record.MemoryRecords
-import org.apache.kafka.common.utils.Time
-import org.easymock.EasyMock
-import org.junit.Assert._
-import org.junit.{After, Before, Test}
-
-import scala.collection.mutable.{HashMap, Map}
-
-class IsrExpirationTest {
-
-  var topicPartitionIsr: Map[(String, Int), Seq[Int]] = new HashMap[(String, Int), Seq[Int]]()
-  val replicaLagTimeMaxMs = 100L
-  val replicaFetchWaitMaxMs = 100
-  val leaderLogEndOffset = 20
-
-  val overridingProps = new Properties()
-  overridingProps.put(KafkaConfig.ReplicaLagTimeMaxMsProp, replicaLagTimeMaxMs.toString)
-  overridingProps.put(KafkaConfig.ReplicaFetchWaitMaxMsProp, replicaFetchWaitMaxMs.toString)
-  val configs = TestUtils.createBrokerConfigs(2, TestUtils.MockZkConnect).map(KafkaConfig.fromProps(_, overridingProps))
-  val topic = "foo"
-
-  val time = new MockTime
-  val metrics = new Metrics
-
-  var replicaManager: ReplicaManager = null
-
-  @Before
-  def setUp() {
-    val logManager: LogManager = EasyMock.createMock(classOf[LogManager])
-    EasyMock.expect(logManager.liveLogDirs).andReturn(Array.empty[File]).anyTimes()
-    EasyMock.replay(logManager)
-
-    replicaManager = new ReplicaManager(configs.head, metrics, time, null, null, logManager, new AtomicBoolean(false),
-      QuotaFactory.instantiate(configs.head, metrics, time, ""), new BrokerTopicStats, new MetadataCache(configs.head.brokerId),
-      new LogDirFailureChannel(configs.head.logDirs.size))
-  }
-
-  @After
-  def tearDown() {
-    replicaManager.shutdown(false)
-    metrics.close()
-  }
-
-  /*
-   * Test the case where a follower is caught up but stops making requests to the leader. Once beyond the configured time limit, it should fall out of ISR
-   */
-  @Test
-  def testIsrExpirationForStuckFollowers() {
-    val log = logMock
-
-    // create one partition and all replicas
-    val partition0 = getPartitionWithAllReplicasInIsr(topic, 0, time, configs.head, log)
-    assertEquals("All replicas should be in ISR", configs.map(_.brokerId).toSet, partition0.inSyncReplicas.map(_.brokerId))
-    val leaderReplica = partition0.getReplica(configs.head.brokerId).get
-
-    // let the follower catch up to the Leader logEndOffset - 1
-    for (replica <- partition0.assignedReplicas - leaderReplica)
-      replica.updateLogReadResult(new LogReadResult(info = FetchDataInfo(new LogOffsetMetadata(leaderLogEndOffset - 1), MemoryRecords.EMPTY),
-                                                    highWatermark = leaderLogEndOffset - 1,
-                                                    leaderLogStartOffset = 0L,
-                                                    leaderLogEndOffset = leaderLogEndOffset,
-                                                    followerLogStartOffset = 0L,
-                                                    fetchTimeMs = time.milliseconds,
-                                                    readSize = -1,
-                                                    lastStableOffset = None))
-    var partition0OSR = partition0.getOutOfSyncReplicas(leaderReplica, configs.head.replicaLagTimeMaxMs)
-    assertEquals("No replica should be out of sync", Set.empty[Int], partition0OSR.map(_.brokerId))
-
-    // let some time pass
-    time.sleep(150)
-
-    // now follower hasn't pulled any data for > replicaMaxLagTimeMs ms. So it is stuck
-    partition0OSR = partition0.getOutOfSyncReplicas(leaderReplica, configs.head.replicaLagTimeMaxMs)
-    assertEquals("Replica 1 should be out of sync", Set(configs.last.brokerId), partition0OSR.map(_.brokerId))
-    EasyMock.verify(log)
-  }
-
-  /*
-   * Test the case where a follower never makes a fetch request. It should fall out of ISR because it will be declared stuck
-   */
-  @Test
-  def testIsrExpirationIfNoFetchRequestMade() {
-    val log = logMock
-
-    // create one partition and all replicas
-    val partition0 = getPartitionWithAllReplicasInIsr(topic, 0, time, configs.head, log)
-    assertEquals("All replicas should be in ISR", configs.map(_.brokerId).toSet, partition0.inSyncReplicas.map(_.brokerId))
-    val leaderReplica = partition0.getReplica(configs.head.brokerId).get
-
-    // Let enough time pass for the replica to be considered stuck
-    time.sleep(150)
-
-    val partition0OSR = partition0.getOutOfSyncReplicas(leaderReplica, configs.head.replicaLagTimeMaxMs)
-    assertEquals("Replica 1 should be out of sync", Set(configs.last.brokerId), partition0OSR.map(_.brokerId))
-    EasyMock.verify(log)
-  }
-
-  /*
-   * Test the case where a follower continually makes fetch requests but is unable to catch up. It should fall out of the ISR
-   * However, any time it makes a request to the LogEndOffset it should be back in the ISR
-   */
-  @Test
-  def testIsrExpirationForSlowFollowers() {
-    // create leader replica
-    val log = logMock
-    // add one partition
-    val partition0 = getPartitionWithAllReplicasInIsr(topic, 0, time, configs.head, log)
-    assertEquals("All replicas should be in ISR", configs.map(_.brokerId).toSet, partition0.inSyncReplicas.map(_.brokerId))
-    val leaderReplica = partition0.getReplica(configs.head.brokerId).get
-
-    // Make the remote replica not read to the end of log. It should be not be out of sync for at least 100 ms
-    for (replica <- partition0.assignedReplicas - leaderReplica)
-      replica.updateLogReadResult(new LogReadResult(info = FetchDataInfo(new LogOffsetMetadata(leaderLogEndOffset - 2), MemoryRecords.EMPTY),
-                                                    highWatermark = leaderLogEndOffset - 2,
-                                                    leaderLogStartOffset = 0L,
-                                                    leaderLogEndOffset = leaderLogEndOffset,
-                                                    followerLogStartOffset = 0L,
-                                                    fetchTimeMs = time.milliseconds,
-                                                    readSize = -1,
-                                                    lastStableOffset = None))
-
-    // Simulate 2 fetch requests spanning more than 100 ms which do not read to the end of the log.
-    // The replicas will no longer be in ISR. We do 2 fetches because we want to simulate the case where the replica is lagging but is not stuck
-    var partition0OSR = partition0.getOutOfSyncReplicas(leaderReplica, configs.head.replicaLagTimeMaxMs)
-    assertEquals("No replica should be out of sync", Set.empty[Int], partition0OSR.map(_.brokerId))
-
-    time.sleep(75)
-
-    (partition0.assignedReplicas - leaderReplica).foreach { r =>
-      r.updateLogReadResult(new LogReadResult(info = FetchDataInfo(new LogOffsetMetadata(leaderLogEndOffset - 1), MemoryRecords.EMPTY),
-                            highWatermark = leaderLogEndOffset - 1,
-                            leaderLogStartOffset = 0L,
-                            leaderLogEndOffset = leaderLogEndOffset,
-                            followerLogStartOffset = 0L,
-                            fetchTimeMs = time.milliseconds,
-                            readSize = -1,
-                            lastStableOffset = None))
-    }
-    partition0OSR = partition0.getOutOfSyncReplicas(leaderReplica, configs.head.replicaLagTimeMaxMs)
-    assertEquals("No replica should be out of sync", Set.empty[Int], partition0OSR.map(_.brokerId))
-
-    time.sleep(75)
-
-    // The replicas will no longer be in ISR
-    partition0OSR = partition0.getOutOfSyncReplicas(leaderReplica, configs.head.replicaLagTimeMaxMs)
-    assertEquals("Replica 1 should be out of sync", Set(configs.last.brokerId), partition0OSR.map(_.brokerId))
-
-    // Now actually make a fetch to the end of the log. The replicas should be back in ISR
-    (partition0.assignedReplicas - leaderReplica).foreach { r =>
-      r.updateLogReadResult(new LogReadResult(info = FetchDataInfo(new LogOffsetMetadata(leaderLogEndOffset), MemoryRecords.EMPTY),
-                            highWatermark = leaderLogEndOffset,
-                            leaderLogStartOffset = 0L,
-                            leaderLogEndOffset = leaderLogEndOffset,
-                            followerLogStartOffset = 0L,
-                            fetchTimeMs = time.milliseconds,
-                            readSize = -1,
-                            lastStableOffset = None))
-    }
-    partition0OSR = partition0.getOutOfSyncReplicas(leaderReplica, configs.head.replicaLagTimeMaxMs)
-    assertEquals("No replica should be out of sync", Set.empty[Int], partition0OSR.map(_.brokerId))
-
-    EasyMock.verify(log)
-  }
-
-  /*
-   * Test the case where a follower has already caught up with same log end offset with the leader. This follower should not be considered as out-of-sync
-   */
-  @Test
-  def testIsrExpirationForCaughtUpFollowers() {
-    val log = logMock
-
-    // create one partition and all replicas
-    val partition0 = getPartitionWithAllReplicasInIsr(topic, 0, time, configs.head, log)
-    assertEquals("All replicas should be in ISR", configs.map(_.brokerId).toSet, partition0.inSyncReplicas.map(_.brokerId))
-    val leaderReplica = partition0.getReplica(configs.head.brokerId).get
-
-    // let the follower catch up to the Leader logEndOffset
-    for (replica <- partition0.assignedReplicas - leaderReplica)
-      replica.updateLogReadResult(new LogReadResult(info = FetchDataInfo(new LogOffsetMetadata(leaderLogEndOffset), MemoryRecords.EMPTY),
-        highWatermark = leaderLogEndOffset,
-        leaderLogStartOffset = 0L,
-        leaderLogEndOffset = leaderLogEndOffset,
-        followerLogStartOffset = 0L,
-        fetchTimeMs = time.milliseconds,
-        readSize = -1,
-        lastStableOffset = None))
-    var partition0OSR = partition0.getOutOfSyncReplicas(leaderReplica, configs.head.replicaLagTimeMaxMs)
-    assertEquals("No replica should be out of sync", Set.empty[Int], partition0OSR.map(_.brokerId))
-
-    // let some time pass
-    time.sleep(150)
-
-    // even though follower hasn't pulled any data for > replicaMaxLagTimeMs ms, the follower has already caught up. So it is not out-of-sync.
-    partition0OSR = partition0.getOutOfSyncReplicas(leaderReplica, configs.head.replicaLagTimeMaxMs)
-    assertEquals("No replica should be out of sync", Set.empty[Int], partition0OSR.map(_.brokerId))
-    EasyMock.verify(log)
-  }
-
-  private def getPartitionWithAllReplicasInIsr(topic: String, partitionId: Int, time: Time, config: KafkaConfig,
-                                               localLog: Log): Partition = {
-    val leaderId = config.brokerId
-    val tp = new TopicPartition(topic, partitionId)
-    val partition = replicaManager.getOrCreatePartition(tp)
-    val leaderReplica = new Replica(leaderId, tp, time, 0, Some(localLog))
-
-    val allReplicas = getFollowerReplicas(partition, leaderId, time) :+ leaderReplica
-    allReplicas.foreach(r => partition.addReplicaIfNotExists(r))
-    // set in sync replicas for this partition to all the assigned replicas
-    partition.inSyncReplicas = allReplicas.toSet
-    // set lastCaughtUpTime to current time
-    for (replica <- partition.assignedReplicas - leaderReplica)
-      replica.updateLogReadResult(new LogReadResult(info = FetchDataInfo(new LogOffsetMetadata(0L), MemoryRecords.EMPTY),
-                                                    highWatermark = 0L,
-                                                    leaderLogStartOffset = 0L,
-                                                    leaderLogEndOffset = 0L,
-                                                    followerLogStartOffset = 0L,
-                                                    fetchTimeMs = time.milliseconds,
-                                                    readSize = -1,
-                                                    lastStableOffset = None))
-    // set the leader and its hw and the hw update time
-    partition.leaderReplicaIdOpt = Some(leaderId)
-    partition
-  }
-
-  private def logMock: Log = {
-    val log: Log = EasyMock.createMock(classOf[Log])
-    EasyMock.expect(log.dir).andReturn(TestUtils.tempDir()).anyTimes()
-    EasyMock.expect(log.onHighWatermarkIncremented(0L))
-    EasyMock.expect(log.logEndOffsetMetadata).andReturn(LogOffsetMetadata(leaderLogEndOffset)).anyTimes()
-    EasyMock.replay(log)
-    log
-  }
-
-  private def getFollowerReplicas(partition: Partition, leaderId: Int, time: Time): Seq[Replica] = {
-    configs.filter(_.brokerId != leaderId).map { config =>
-      new Replica(config.brokerId, partition.topicPartition, time)
-    }
-  }
-}
diff --git a/core/src/test/scala/unit/kafka/server/IsrExpirationTest.scala b/core/src/test/scala/unit/kafka/server/IsrExpirationTest.scala
new file mode 100644
index 0000000000000..a71b24cbf52bc
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/server/IsrExpirationTest.scala
@@ -0,0 +1,248 @@
+/**
+ * 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 java.io.File
+import java.util.Properties
+import java.util.concurrent.atomic.AtomicBoolean
+
+import kafka.cluster.Partition
+import kafka.log.{Log, LogManager}
+import kafka.utils._
+import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.metrics.Metrics
+import org.apache.kafka.common.utils.Time
+import org.easymock.EasyMock
+import org.junit.Assert._
+import org.junit.{After, Before, Test}
+
+import scala.collection.Seq
+import scala.collection.mutable.{HashMap, Map}
+
+class IsrExpirationTest {
+
+  var topicPartitionIsr: Map[(String, Int), Seq[Int]] = new HashMap[(String, Int), Seq[Int]]()
+  val replicaLagTimeMaxMs = 100L
+  val replicaFetchWaitMaxMs = 100
+  val leaderLogEndOffset = 20
+  val leaderLogHighWatermark = 20L
+
+  val overridingProps = new Properties()
+  overridingProps.put(KafkaConfig.ReplicaLagTimeMaxMsProp, replicaLagTimeMaxMs.toString)
+  overridingProps.put(KafkaConfig.ReplicaFetchWaitMaxMsProp, replicaFetchWaitMaxMs.toString)
+  val configs = TestUtils.createBrokerConfigs(2, TestUtils.MockZkConnect).map(KafkaConfig.fromProps(_, overridingProps))
+  val topic = "foo"
+
+  val time = new MockTime
+  val metrics = new Metrics
+
+  var replicaManager: ReplicaManager = null
+
+  @Before
+  def setUp(): Unit = {
+    val logManager: LogManager = EasyMock.createMock(classOf[LogManager])
+    EasyMock.expect(logManager.liveLogDirs).andReturn(Array.empty[File]).anyTimes()
+    EasyMock.replay(logManager)
+
+    replicaManager = new ReplicaManager(configs.head, metrics, time, null, null, logManager, new AtomicBoolean(false),
+      QuotaFactory.instantiate(configs.head, metrics, time, ""), new BrokerTopicStats, new MetadataCache(configs.head.brokerId),
+      new LogDirFailureChannel(configs.head.logDirs.size))
+  }
+
+  @After
+  def tearDown(): Unit = {
+    replicaManager.shutdown(false)
+    metrics.close()
+  }
+
+  /*
+   * Test the case where a follower is caught up but stops making requests to the leader. Once beyond the configured time limit, it should fall out of ISR
+   */
+  @Test
+  def testIsrExpirationForStuckFollowers(): Unit = {
+    val log = logMock
+
+    // create one partition and all replicas
+    val partition0 = getPartitionWithAllReplicasInIsr(topic, 0, time, configs.head, log)
+    assertEquals("All replicas should be in ISR", configs.map(_.brokerId).toSet, partition0.inSyncReplicaIds)
+
+    // let the follower catch up to the Leader logEndOffset - 1
+    for (replica <- partition0.remoteReplicas)
+      replica.updateFetchState(
+        followerFetchOffsetMetadata = LogOffsetMetadata(leaderLogEndOffset - 1),
+        followerStartOffset = 0L,
+        followerFetchTimeMs= time.milliseconds,
+        leaderEndOffset = leaderLogEndOffset,
+        lastSentHighwatermark = partition0.localLogOrException.highWatermark)
+    var partition0OSR = partition0.getOutOfSyncReplicas(configs.head.replicaLagTimeMaxMs)
+    assertEquals("No replica should be out of sync", Set.empty[Int], partition0OSR)
+
+    // let some time pass
+    time.sleep(150)
+
+    // now follower hasn't pulled any data for > replicaMaxLagTimeMs ms. So it is stuck
+    partition0OSR = partition0.getOutOfSyncReplicas(configs.head.replicaLagTimeMaxMs)
+    assertEquals("Replica 1 should be out of sync", Set(configs.last.brokerId), partition0OSR)
+    EasyMock.verify(log)
+  }
+
+  /*
+   * Test the case where a follower never makes a fetch request. It should fall out of ISR because it will be declared stuck
+   */
+  @Test
+  def testIsrExpirationIfNoFetchRequestMade(): Unit = {
+    val log = logMock
+
+    // create one partition and all replicas
+    val partition0 = getPartitionWithAllReplicasInIsr(topic, 0, time, configs.head, log)
+    assertEquals("All replicas should be in ISR", configs.map(_.brokerId).toSet, partition0.inSyncReplicaIds)
+
+    // Let enough time pass for the replica to be considered stuck
+    time.sleep(150)
+
+    val partition0OSR = partition0.getOutOfSyncReplicas(configs.head.replicaLagTimeMaxMs)
+    assertEquals("Replica 1 should be out of sync", Set(configs.last.brokerId), partition0OSR)
+    EasyMock.verify(log)
+  }
+
+  /*
+   * Test the case where a follower continually makes fetch requests but is unable to catch up. It should fall out of the ISR
+   * However, any time it makes a request to the LogEndOffset it should be back in the ISR
+   */
+  @Test
+  def testIsrExpirationForSlowFollowers(): Unit = {
+    // create leader replica
+    val log = logMock
+    // add one partition
+    val partition0 = getPartitionWithAllReplicasInIsr(topic, 0, time, configs.head, log)
+    assertEquals("All replicas should be in ISR", configs.map(_.brokerId).toSet, partition0.inSyncReplicaIds)
+    // Make the remote replica not read to the end of log. It should be not be out of sync for at least 100 ms
+    for (replica <- partition0.remoteReplicas)
+      replica.updateFetchState(
+        followerFetchOffsetMetadata = LogOffsetMetadata(leaderLogEndOffset - 2),
+        followerStartOffset = 0L,
+        followerFetchTimeMs= time.milliseconds,
+        leaderEndOffset = leaderLogEndOffset,
+        lastSentHighwatermark = partition0.localLogOrException.highWatermark)
+
+    // Simulate 2 fetch requests spanning more than 100 ms which do not read to the end of the log.
+    // The replicas will no longer be in ISR. We do 2 fetches because we want to simulate the case where the replica is lagging but is not stuck
+    var partition0OSR = partition0.getOutOfSyncReplicas(configs.head.replicaLagTimeMaxMs)
+    assertEquals("No replica should be out of sync", Set.empty[Int], partition0OSR)
+
+    time.sleep(75)
+
+    partition0.remoteReplicas.foreach { r =>
+      r.updateFetchState(
+        followerFetchOffsetMetadata = LogOffsetMetadata(leaderLogEndOffset - 1),
+        followerStartOffset = 0L,
+        followerFetchTimeMs= time.milliseconds,
+        leaderEndOffset = leaderLogEndOffset,
+        lastSentHighwatermark = partition0.localLogOrException.highWatermark)
+    }
+    partition0OSR = partition0.getOutOfSyncReplicas(configs.head.replicaLagTimeMaxMs)
+    assertEquals("No replica should be out of sync", Set.empty[Int], partition0OSR)
+
+    time.sleep(75)
+
+    // The replicas will no longer be in ISR
+    partition0OSR = partition0.getOutOfSyncReplicas(configs.head.replicaLagTimeMaxMs)
+    assertEquals("Replica 1 should be out of sync", Set(configs.last.brokerId), partition0OSR)
+
+    // Now actually make a fetch to the end of the log. The replicas should be back in ISR
+    partition0.remoteReplicas.foreach { r =>
+      r.updateFetchState(
+        followerFetchOffsetMetadata = LogOffsetMetadata(leaderLogEndOffset),
+        followerStartOffset = 0L,
+        followerFetchTimeMs= time.milliseconds,
+        leaderEndOffset = leaderLogEndOffset,
+        lastSentHighwatermark = partition0.localLogOrException.highWatermark)
+    }
+    partition0OSR = partition0.getOutOfSyncReplicas(configs.head.replicaLagTimeMaxMs)
+    assertEquals("No replica should be out of sync", Set.empty[Int], partition0OSR)
+
+    EasyMock.verify(log)
+  }
+
+  /*
+   * Test the case where a follower has already caught up with same log end offset with the leader. This follower should not be considered as out-of-sync
+   */
+  @Test
+  def testIsrExpirationForCaughtUpFollowers(): Unit = {
+    val log = logMock
+
+    // create one partition and all replicas
+    val partition0 = getPartitionWithAllReplicasInIsr(topic, 0, time, configs.head, log)
+    assertEquals("All replicas should be in ISR", configs.map(_.brokerId).toSet, partition0.inSyncReplicaIds)
+
+    // let the follower catch up to the Leader logEndOffset
+    for (replica <- partition0.remoteReplicas)
+      replica.updateFetchState(
+        followerFetchOffsetMetadata = LogOffsetMetadata(leaderLogEndOffset),
+        followerStartOffset = 0L,
+        followerFetchTimeMs= time.milliseconds,
+        leaderEndOffset = leaderLogEndOffset,
+        lastSentHighwatermark = partition0.localLogOrException.highWatermark)
+
+    var partition0OSR = partition0.getOutOfSyncReplicas(configs.head.replicaLagTimeMaxMs)
+    assertEquals("No replica should be out of sync", Set.empty[Int], partition0OSR)
+
+    // let some time pass
+    time.sleep(150)
+
+    // even though follower hasn't pulled any data for > replicaMaxLagTimeMs ms, the follower has already caught up. So it is not out-of-sync.
+    partition0OSR = partition0.getOutOfSyncReplicas(configs.head.replicaLagTimeMaxMs)
+    assertEquals("No replica should be out of sync", Set.empty[Int], partition0OSR)
+    EasyMock.verify(log)
+  }
+
+  private def getPartitionWithAllReplicasInIsr(topic: String, partitionId: Int, time: Time, config: KafkaConfig,
+                                               localLog: Log): Partition = {
+    val leaderId = config.brokerId
+    val tp = new TopicPartition(topic, partitionId)
+    val partition = replicaManager.createPartition(tp)
+    partition.setLog(localLog, isFutureLog = false)
+
+    partition.updateAssignmentAndIsr(
+      assignment = configs.map(_.brokerId),
+      isr = configs.map(_.brokerId).toSet
+    )
+
+    // set lastCaughtUpTime to current time
+    for (replica <- partition.remoteReplicas)
+      replica.updateFetchState(
+        followerFetchOffsetMetadata = LogOffsetMetadata(0L),
+        followerStartOffset = 0L,
+        followerFetchTimeMs= time.milliseconds,
+        leaderEndOffset = 0L,
+        lastSentHighwatermark = partition.localLogOrException.highWatermark)
+
+    // set the leader and its hw and the hw update time
+    partition.leaderReplicaIdOpt = Some(leaderId)
+    partition
+  }
+
+  private def logMock: Log = {
+    val log: Log = EasyMock.createMock(classOf[Log])
+    EasyMock.expect(log.dir).andReturn(TestUtils.tempDir()).anyTimes()
+    EasyMock.expect(log.logEndOffsetMetadata).andReturn(LogOffsetMetadata(leaderLogEndOffset)).anyTimes()
+    EasyMock.expect(log.logEndOffset).andReturn(leaderLogEndOffset).anyTimes()
+    EasyMock.expect(log.highWatermark).andReturn(leaderLogHighWatermark).anyTimes()
+    EasyMock.replay(log)
+    log
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/server/KafkaApisTest.scala b/core/src/test/scala/unit/kafka/server/KafkaApisTest.scala
index f3eae1ea8bf1e..d822ace68f3b9 100644
--- a/core/src/test/scala/unit/kafka/server/KafkaApisTest.scala
+++ b/core/src/test/scala/unit/kafka/server/KafkaApisTest.scala
@@ -20,21 +20,23 @@ package kafka.server
 import java.net.InetAddress
 import java.nio.charset.StandardCharsets
 import java.util
+import java.util.Random
 import java.util.{Collections, Optional}
 import java.util.Arrays.asList
 
-import kafka.api.{ApiVersion, KAFKA_0_10_2_IV0}
+import kafka.api.{ApiVersion, KAFKA_0_10_2_IV0, KAFKA_2_2_IV1}
 import kafka.controller.KafkaController
-import kafka.coordinator.group.GroupCoordinator
+import kafka.coordinator.group.{GroupCoordinator, GroupSummary, MemberSummary}
 import kafka.coordinator.transaction.TransactionCoordinator
+import kafka.log.AppendOrigin
 import kafka.network.RequestChannel
 import kafka.network.RequestChannel.SendResponse
-import kafka.security.auth.Authorizer
 import kafka.server.QuotaFactory.QuotaManagers
 import kafka.utils.{MockTime, TestUtils}
 import kafka.zk.KafkaZkClient
 import org.apache.kafka.common.TopicPartition
 import org.apache.kafka.common.errors.UnsupportedVersionException
+import org.apache.kafka.common.internals.Topic
 import org.apache.kafka.common.memory.MemoryPool
 import org.apache.kafka.common.metrics.Metrics
 import org.apache.kafka.common.network.ListenerName
@@ -42,17 +44,23 @@ import org.apache.kafka.common.protocol.{ApiKeys, Errors}
 import org.apache.kafka.common.record.FileRecords.TimestampAndOffset
 import org.apache.kafka.common.record._
 import org.apache.kafka.common.requests.ProduceResponse.PartitionResponse
-import org.apache.kafka.common.requests.UpdateMetadataRequest.{Broker, EndPoint}
 import org.apache.kafka.common.requests.WriteTxnMarkersRequest.TxnMarkerEntry
 import org.apache.kafka.common.requests.{FetchMetadata => JFetchMetadata, _}
 import org.apache.kafka.common.security.auth.{KafkaPrincipal, SecurityProtocol}
 import org.easymock.{Capture, EasyMock, IAnswer}
 import EasyMock._
-import org.junit.Assert.{assertEquals, assertNull, assertTrue}
+import org.apache.kafka.common.message.{DescribeGroupsRequestData, HeartbeatRequestData, JoinGroupRequestData, OffsetCommitRequestData, OffsetCommitResponseData, OffsetDeleteRequestData, SyncGroupRequestData, TxnOffsetCommitRequestData}
+import org.apache.kafka.common.message.JoinGroupRequestData.JoinGroupRequestProtocol
+import org.apache.kafka.common.message.LeaveGroupRequestData.MemberIdentity
+import org.apache.kafka.common.message.OffsetDeleteRequestData.{OffsetDeleteRequestPartition, OffsetDeleteRequestTopic, OffsetDeleteRequestTopicCollection}
+import org.apache.kafka.common.message.UpdateMetadataRequestData.{UpdateMetadataBroker, UpdateMetadataEndpoint, UpdateMetadataPartitionState}
+import org.apache.kafka.common.replica.ClientMetadata
+import org.apache.kafka.server.authorizer.Authorizer
+import org.junit.Assert.{assertArrayEquals, assertEquals, assertNull, assertTrue}
 import org.junit.{After, Test}
 
 import scala.collection.JavaConverters._
-import scala.collection.Map
+import scala.collection.{Map, Seq, mutable}
 
 class KafkaApisTest {
 
@@ -68,6 +76,7 @@ class KafkaApisTest {
   private val brokerId = 1
   private val metadataCache = new MetadataCache(brokerId)
   private val authorizer: Option[Authorizer] = None
+  private val observer: Observer = EasyMock.createNiceMock(classOf[Observer])
   private val clientQuotaManager: ClientQuotaManager = EasyMock.createNiceMock(classOf[ClientQuotaManager])
   private val clientRequestQuotaManager: ClientRequestQuotaManager = EasyMock.createNiceMock(classOf[ClientRequestQuotaManager])
   private val replicaQuotaManager: ReplicationQuotaManager = EasyMock.createNiceMock(classOf[ReplicationQuotaManager])
@@ -79,7 +88,7 @@ class KafkaApisTest {
   private val time = new MockTime
 
   @After
-  def tearDown() {
+  def tearDown(): Unit = {
     quotas.shutdown()
     metrics.close()
   }
@@ -100,6 +109,7 @@ class KafkaApisTest {
       metadataCache,
       metrics,
       authorizer,
+      observer,
       quotas,
       fetchManager,
       brokerTopicStats,
@@ -117,10 +127,21 @@ class KafkaApisTest {
     def checkInvalidPartition(invalidPartitionId: Int): Unit = {
       EasyMock.reset(replicaManager, clientRequestQuotaManager, requestChannel)
 
-      val invalidTopicPartition = new TopicPartition(topic, invalidPartitionId)
-      val partitionOffsetCommitData = new OffsetCommitRequest.PartitionData(15L, Optional.empty[Integer](), "")
-      val (offsetCommitRequest, request) = buildRequest(new OffsetCommitRequest.Builder("groupId",
-        Map(invalidTopicPartition -> partitionOffsetCommitData).asJava))
+      val (offsetCommitRequest, request) = buildRequest(new OffsetCommitRequest.Builder(
+        new OffsetCommitRequestData()
+          .setGroupId("groupId")
+          .setTopics(Collections.singletonList(
+            new OffsetCommitRequestData.OffsetCommitRequestTopic()
+              .setName(topic)
+              .setPartitions(Collections.singletonList(
+                new OffsetCommitRequestData.OffsetCommitRequestPartition()
+                  .setPartitionIndex(invalidPartitionId)
+                  .setCommittedOffset(15)
+                  .setCommittedLeaderEpoch(RecordBatch.NO_PARTITION_LEADER_EPOCH)
+                  .setCommittedMetadata(""))
+              )
+          ))
+      ))
 
       val capturedResponse = expectNoThrottling()
       EasyMock.replay(replicaManager, clientRequestQuotaManager, requestChannel)
@@ -128,7 +149,8 @@ class KafkaApisTest {
 
       val response = readResponse(ApiKeys.OFFSET_COMMIT, offsetCommitRequest, capturedResponse)
         .asInstanceOf[OffsetCommitResponse]
-      assertEquals(Errors.UNKNOWN_TOPIC_OR_PARTITION, response.responseData().get(invalidTopicPartition))
+      assertEquals(Errors.UNKNOWN_TOPIC_OR_PARTITION,
+        Errors.forCode(response.data().topics().get(0).partitions().get(0).errorCode()))
     }
 
     checkInvalidPartition(-1)
@@ -145,8 +167,15 @@ class KafkaApisTest {
 
       val invalidTopicPartition = new TopicPartition(topic, invalidPartitionId)
       val partitionOffsetCommitData = new TxnOffsetCommitRequest.CommittedOffset(15L, "", Optional.empty())
-      val (offsetCommitRequest, request) = buildRequest(new TxnOffsetCommitRequest.Builder("txnlId", "groupId",
-        15L, 0.toShort, Map(invalidTopicPartition -> partitionOffsetCommitData).asJava))
+      val (offsetCommitRequest, request) = buildRequest(new TxnOffsetCommitRequest.Builder(
+        new TxnOffsetCommitRequestData()
+          .setTransactionalId("txnlId")
+          .setGroupId("groupId")
+          .setProducerId(15L)
+          .setProducerEpoch(0.toShort)
+          .setTopics(TxnOffsetCommitRequest.getTopics(
+            Map(invalidTopicPartition -> partitionOffsetCommitData).asJava))
+      ))
 
       val capturedResponse = expectNoThrottling()
       EasyMock.replay(replicaManager, clientRequestQuotaManager, requestChannel)
@@ -268,7 +297,7 @@ class KafkaApisTest {
     EasyMock.expect(replicaManager.appendRecords(EasyMock.anyLong(),
       EasyMock.anyShort(),
       EasyMock.eq(true),
-      EasyMock.eq(false),
+      EasyMock.eq(AppendOrigin.Coordinator),
       EasyMock.anyObject(),
       EasyMock.capture(responseCallback),
       EasyMock.anyObject(),
@@ -289,6 +318,41 @@ class KafkaApisTest {
     EasyMock.verify(replicaManager)
   }
 
+  @Test
+  def shouldResignCoordinatorsIfStopReplicaReceivedWithDeleteFlag(): Unit = {
+    val controllerId = 0
+    val controllerEpoch = 5
+    val brokerEpoch = 230498320L
+
+    val groupMetadataPartition = new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, 0)
+    val txnStatePartition = new TopicPartition(Topic.TRANSACTION_STATE_TOPIC_NAME, 0)
+
+    val (_, request) = buildRequest(new StopReplicaRequest.Builder(
+      ApiKeys.STOP_REPLICA.latestVersion,
+      controllerId,
+      controllerEpoch,
+      brokerEpoch,
+      brokerEpoch,
+      true,
+      Set(groupMetadataPartition, txnStatePartition).asJava))
+
+    EasyMock.expect(replicaManager.stopReplicas(anyObject())).andReturn(
+      (mutable.Map(groupMetadataPartition -> Errors.NONE, txnStatePartition -> Errors.NONE), Errors.NONE))
+    EasyMock.expect(controller.brokerEpoch).andStubReturn(brokerEpoch)
+
+    txnCoordinator.onResignation(txnStatePartition.partition, None)
+    EasyMock.expectLastCall()
+
+    groupCoordinator.onResignation(groupMetadataPartition.partition)
+    EasyMock.expectLastCall()
+
+    EasyMock.replay(controller, replicaManager, txnCoordinator, groupCoordinator)
+
+    createKafkaApis().handleStopReplicaRequest(request)
+
+    EasyMock.verify(txnCoordinator, groupCoordinator)
+  }
+
   @Test
   def shouldRespondWithUnknownTopicOrPartitionForBadPartitionAndNoErrorsForGoodPartition(): Unit = {
     val tp1 = new TopicPartition("t", 0)
@@ -307,7 +371,7 @@ class KafkaApisTest {
     EasyMock.expect(replicaManager.appendRecords(EasyMock.anyLong(),
       EasyMock.anyShort(),
       EasyMock.eq(true),
-      EasyMock.eq(false),
+      EasyMock.eq(AppendOrigin.Coordinator),
       EasyMock.anyObject(),
       EasyMock.capture(responseCallback),
       EasyMock.anyObject(),
@@ -338,7 +402,7 @@ class KafkaApisTest {
     EasyMock.expect(replicaManager.appendRecords(EasyMock.anyLong(),
       EasyMock.anyShort(),
       EasyMock.eq(true),
-      EasyMock.eq(false),
+      EasyMock.eq(AppendOrigin.Coordinator),
       EasyMock.anyObject(),
       EasyMock.anyObject(),
       EasyMock.anyObject(),
@@ -370,6 +434,113 @@ class KafkaApisTest {
     testListOffsetFailedGetLeaderReplica(Errors.UNKNOWN_TOPIC_OR_PARTITION)
   }
 
+  @Test
+  def testDescribeGroups(): Unit = {
+    val groupId = "groupId"
+    val random = new Random()
+    val metadata = new Array[Byte](10)
+    random.nextBytes(metadata)
+    val assignment = new Array[Byte](10)
+    random.nextBytes(assignment)
+
+    val memberSummary = MemberSummary("memberid", Some("instanceid"), "clientid", "clienthost", metadata, assignment)
+    val groupSummary = GroupSummary("Stable", "consumer", "roundrobin", List(memberSummary))
+
+    EasyMock.reset(groupCoordinator, replicaManager, clientRequestQuotaManager, requestChannel)
+
+    val (describeGroupsRequest, request) = buildRequest(new DescribeGroupsRequest.Builder(
+      new DescribeGroupsRequestData().setGroups(List(groupId).asJava)
+    ))
+
+    val capturedResponse = expectNoThrottling()
+    EasyMock.expect(groupCoordinator.handleDescribeGroup(EasyMock.eq(groupId)))
+      .andReturn((Errors.NONE, groupSummary))
+    EasyMock.replay(groupCoordinator, replicaManager, clientRequestQuotaManager, requestChannel)
+
+    createKafkaApis().handleDescribeGroupRequest(request)
+
+    val response = readResponse(ApiKeys.DESCRIBE_GROUPS, describeGroupsRequest, capturedResponse)
+      .asInstanceOf[DescribeGroupsResponse]
+
+    val group = response.data().groups().get(0)
+    assertEquals(Errors.NONE, Errors.forCode(group.errorCode()))
+    assertEquals(groupId, group.groupId())
+    assertEquals(groupSummary.state, group.groupState())
+    assertEquals(groupSummary.protocolType, group.protocolType())
+    assertEquals(groupSummary.protocol, group.protocolData())
+    assertEquals(groupSummary.members.size, group.members().size())
+
+    val member = group.members().get(0)
+    assertEquals(memberSummary.memberId, member.memberId())
+    assertEquals(memberSummary.groupInstanceId.orNull, member.groupInstanceId())
+    assertEquals(memberSummary.clientId, member.clientId())
+    assertEquals(memberSummary.clientHost, member.clientHost())
+    assertArrayEquals(memberSummary.metadata, member.memberMetadata())
+    assertArrayEquals(memberSummary.assignment, member.memberAssignment())
+  }
+
+  @Test
+  def testOffsetDeleteWithInvalidPartition(): Unit = {
+    val group = "groupId"
+    val topic = "topic"
+    setupBasicMetadataCache(topic, numPartitions = 1)
+
+    def checkInvalidPartition(invalidPartitionId: Int): Unit = {
+      EasyMock.reset(groupCoordinator, replicaManager, clientRequestQuotaManager, requestChannel)
+
+      val topics = new OffsetDeleteRequestTopicCollection()
+      topics.add(new OffsetDeleteRequestTopic()
+        .setName(topic)
+        .setPartitions(Collections.singletonList(
+          new OffsetDeleteRequestPartition().setPartitionIndex(invalidPartitionId))))
+      val (offsetDeleteRequest, request) = buildRequest(new OffsetDeleteRequest.Builder(
+        new OffsetDeleteRequestData()
+          .setGroupId(group)
+          .setTopics(topics)
+      ))
+
+      val capturedResponse = expectNoThrottling()
+      EasyMock.expect(groupCoordinator.handleDeleteOffsets(EasyMock.eq(group), EasyMock.eq(Seq.empty)))
+        .andReturn((Errors.NONE, Map.empty))
+      EasyMock.replay(groupCoordinator, replicaManager, clientRequestQuotaManager, requestChannel)
+
+      createKafkaApis().handleOffsetDeleteRequest(request)
+
+      val response = readResponse(ApiKeys.OFFSET_DELETE, offsetDeleteRequest, capturedResponse)
+        .asInstanceOf[OffsetDeleteResponse]
+
+      assertEquals(Errors.UNKNOWN_TOPIC_OR_PARTITION,
+        Errors.forCode(response.data.topics.find(topic).partitions.find(invalidPartitionId).errorCode()))
+    }
+
+    checkInvalidPartition(-1)
+    checkInvalidPartition(1) // topic has only one partition
+  }
+
+  @Test
+  def testOffsetDeleteWithInvalidGroup(): Unit = {
+    val group = "groupId"
+
+    EasyMock.reset(groupCoordinator, replicaManager, clientRequestQuotaManager, requestChannel)
+
+    val (offsetDeleteRequest, request) = buildRequest(new OffsetDeleteRequest.Builder(
+      new OffsetDeleteRequestData()
+        .setGroupId(group)
+    ))
+
+    val capturedResponse = expectNoThrottling()
+    EasyMock.expect(groupCoordinator.handleDeleteOffsets(EasyMock.eq(group), EasyMock.eq(Seq.empty)))
+      .andReturn((Errors.GROUP_ID_NOT_FOUND, Map.empty))
+    EasyMock.replay(groupCoordinator, replicaManager, clientRequestQuotaManager, requestChannel)
+
+    createKafkaApis().handleOffsetDeleteRequest(request)
+
+    val response = readResponse(ApiKeys.OFFSET_DELETE, offsetDeleteRequest, capturedResponse)
+      .asInstanceOf[OffsetDeleteResponse]
+
+    assertEquals(Errors.GROUP_ID_NOT_FOUND, Errors.forCode(response.data.errorCode()))
+  }
+
   private def testListOffsetFailedGetLeaderReplica(error: Errors): Unit = {
     val tp = new TopicPartition("foo", 0)
     val isolationLevel = IsolationLevel.READ_UNCOMMITTED
@@ -450,14 +621,15 @@ class KafkaApisTest {
 
     replicaManager.fetchMessages(anyLong, anyInt, anyInt, anyInt, anyBoolean,
       anyObject[Seq[(TopicPartition, FetchRequest.PartitionData)]], anyObject[ReplicaQuota],
-      anyObject[Seq[(TopicPartition, FetchPartitionData)] => Unit](), anyObject[IsolationLevel])
+      anyObject[Seq[(TopicPartition, FetchPartitionData)] => Unit](), anyObject[IsolationLevel],
+      anyObject[Option[ClientMetadata]])
     expectLastCall[Unit].andAnswer(new IAnswer[Unit] {
       def answer: Unit = {
         val callback = getCurrentArguments.apply(7).asInstanceOf[(Seq[(TopicPartition, FetchPartitionData)] => Unit)]
         val records = MemoryRecords.withRecords(CompressionType.NONE,
           new SimpleRecord(timestamp, "foo".getBytes(StandardCharsets.UTF_8)))
         callback(Seq(tp -> new FetchPartitionData(Errors.NONE, hw, 0, records,
-          None, None)))
+          None, None, Option.empty)))
       }
     })
 
@@ -495,20 +667,229 @@ class KafkaApisTest {
     assertNull(partitionData.abortedTransactions)
   }
 
+  @Test
+  def testJoinGroupProtocolsOrder(): Unit = {
+    val protocols = List(
+      new JoinGroupRequestProtocol().setName("first").setMetadata("first".getBytes()),
+      new JoinGroupRequestProtocol().setName("second").setMetadata("second".getBytes())
+    )
+
+    EasyMock.expect(groupCoordinator.handleJoinGroup(
+      anyString,
+      anyString,
+      anyObject(classOf[Option[String]]),
+      anyBoolean,
+      anyString,
+      anyString,
+      anyInt,
+      anyInt,
+      anyString,
+      EasyMock.eq(protocols.map(protocol => (protocol.name, protocol.metadata))),
+      anyObject()
+    ))
+
+    createKafkaApis().handleJoinGroupRequest(
+      buildRequest(
+        new JoinGroupRequest.Builder(
+          new JoinGroupRequestData()
+            .setGroupId("test")
+            .setMemberId("test")
+            .setProtocolType("consumer")
+            .setProtocols(new JoinGroupRequestData.JoinGroupRequestProtocolCollection(protocols.iterator.asJava))
+        )
+      )._2)
+
+    EasyMock.replay(groupCoordinator)
+  }
+
+  @Test
+  def rejectJoinGroupRequestWhenStaticMembershipNotSupported(): Unit = {
+    val capturedResponse = expectNoThrottling()
+    EasyMock.replay(clientRequestQuotaManager, requestChannel)
+
+    val (joinGroupRequest, requestChannelRequest) = buildRequest(new JoinGroupRequest.Builder(
+      new JoinGroupRequestData()
+        .setGroupId("test")
+        .setMemberId("test")
+        .setGroupInstanceId("instanceId")
+        .setProtocolType("consumer")
+        .setProtocols(new JoinGroupRequestData.JoinGroupRequestProtocolCollection)
+    ))
+    createKafkaApis(KAFKA_2_2_IV1).handleJoinGroupRequest(requestChannelRequest)
+
+    val response = readResponse(ApiKeys.JOIN_GROUP, joinGroupRequest, capturedResponse).asInstanceOf[JoinGroupResponse]
+    assertEquals(Errors.UNSUPPORTED_VERSION, response.error())
+    EasyMock.replay(groupCoordinator)
+  }
+
+  @Test
+  def rejectSyncGroupRequestWhenStaticMembershipNotSupported(): Unit = {
+    val capturedResponse = expectNoThrottling()
+    EasyMock.replay(clientRequestQuotaManager, requestChannel)
+
+    val (syncGroupRequest, requestChannelRequest) = buildRequest(new SyncGroupRequest.Builder(
+      new SyncGroupRequestData()
+        .setGroupId("test")
+        .setMemberId("test")
+        .setGroupInstanceId("instanceId")
+        .setGenerationId(1)
+    ))
+    createKafkaApis(KAFKA_2_2_IV1).handleSyncGroupRequest(requestChannelRequest)
+
+    val response = readResponse(ApiKeys.SYNC_GROUP, syncGroupRequest, capturedResponse).asInstanceOf[SyncGroupResponse]
+    assertEquals(Errors.UNSUPPORTED_VERSION, response.error())
+    EasyMock.replay(groupCoordinator)
+  }
+
+  @Test
+  def rejectHeartbeatRequestWhenStaticMembershipNotSupported(): Unit = {
+    val capturedResponse = expectNoThrottling()
+    EasyMock.replay(clientRequestQuotaManager, requestChannel)
+
+    val (heartbeatRequest, requestChannelRequest) = buildRequest(new HeartbeatRequest.Builder(
+      new HeartbeatRequestData()
+        .setGroupId("test")
+        .setMemberId("test")
+        .setGroupInstanceId("instanceId")
+        .setGenerationId(1)
+    ))
+    createKafkaApis(KAFKA_2_2_IV1).handleHeartbeatRequest(requestChannelRequest)
+
+    val response = readResponse(ApiKeys.HEARTBEAT, heartbeatRequest, capturedResponse).asInstanceOf[HeartbeatResponse]
+    assertEquals(Errors.UNSUPPORTED_VERSION, response.error())
+    EasyMock.replay(groupCoordinator)
+  }
+
+  @Test
+  def rejectOffsetCommitRequestWhenStaticMembershipNotSupported(): Unit = {
+    val capturedResponse = expectNoThrottling()
+    EasyMock.replay(clientRequestQuotaManager, requestChannel)
+
+    val (offsetCommitRequest, requestChannelRequest) = buildRequest(new OffsetCommitRequest.Builder(
+      new OffsetCommitRequestData()
+        .setGroupId("test")
+        .setMemberId("test")
+        .setGroupInstanceId("instanceId")
+        .setGenerationId(100)
+        .setTopics(Collections.singletonList(
+          new OffsetCommitRequestData.OffsetCommitRequestTopic()
+            .setName("test")
+            .setPartitions(Collections.singletonList(
+              new OffsetCommitRequestData.OffsetCommitRequestPartition()
+                .setPartitionIndex(0)
+                .setCommittedOffset(100)
+                .setCommittedLeaderEpoch(RecordBatch.NO_PARTITION_LEADER_EPOCH)
+                .setCommittedMetadata("")
+            ))
+        ))
+    ))
+    createKafkaApis(KAFKA_2_2_IV1).handleOffsetCommitRequest(requestChannelRequest)
+
+    val expectedTopicErrors = Collections.singletonList(
+      new OffsetCommitResponseData.OffsetCommitResponseTopic()
+        .setName("test")
+        .setPartitions(Collections.singletonList(
+          new OffsetCommitResponseData.OffsetCommitResponsePartition()
+            .setPartitionIndex(0)
+            .setErrorCode(Errors.UNSUPPORTED_VERSION.code())
+        ))
+    )
+    val response = readResponse(ApiKeys.OFFSET_COMMIT, offsetCommitRequest, capturedResponse).asInstanceOf[OffsetCommitResponse]
+    assertEquals(expectedTopicErrors, response.data.topics())
+    EasyMock.replay(groupCoordinator)
+  }
+
+  @Test
+  def testMultipleLeaveGroup(): Unit = {
+    val groupId = "groupId"
+
+    val leaveMemberList = List(
+      new MemberIdentity()
+        .setMemberId("member-1")
+        .setGroupInstanceId("instance-1"),
+      new MemberIdentity()
+        .setMemberId("member-2")
+        .setGroupInstanceId("instance-2")
+    )
+
+    EasyMock.expect(groupCoordinator.handleLeaveGroup(
+      EasyMock.eq(groupId),
+      EasyMock.eq(leaveMemberList),
+      anyObject()
+    ))
+
+    val (_, leaveRequest) = buildRequest(
+      new LeaveGroupRequest.Builder(
+        groupId,
+        leaveMemberList.asJava)
+    )
+
+    createKafkaApis().handleLeaveGroupRequest(leaveRequest)
+
+    EasyMock.replay(groupCoordinator)
+  }
+
+  @Test
+  def testSingleLeaveGroup(): Unit = {
+    val groupId = "groupId"
+    val memberId = "member"
+
+    val singleLeaveMember = List(
+      new MemberIdentity()
+        .setMemberId(memberId)
+    )
+
+    EasyMock.expect(groupCoordinator.handleLeaveGroup(
+      EasyMock.eq(groupId),
+      EasyMock.eq(singleLeaveMember),
+      anyObject()
+    ))
+
+    val (_, leaveRequest) = buildRequest(
+      new LeaveGroupRequest.Builder(
+        groupId,
+        singleLeaveMember.asJava)
+    )
+
+    createKafkaApis().handleLeaveGroupRequest(leaveRequest)
+
+    EasyMock.replay(groupCoordinator)
+  }
+
   /**
    * Return pair of listener names in the metadataCache: PLAINTEXT and LISTENER2 respectively.
    */
   private def updateMetadataCacheWithInconsistentListeners(): (ListenerName, ListenerName) = {
     val plaintextListener = ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT)
     val anotherListener = new ListenerName("LISTENER2")
-    val brokers = Set(
-      new Broker(0, Seq(new EndPoint("broker0", 9092, SecurityProtocol.PLAINTEXT, plaintextListener),
-        new EndPoint("broker0", 9093, SecurityProtocol.PLAINTEXT, anotherListener)).asJava, "rack"),
-      new Broker(1, Seq(new EndPoint("broker1", 9092, SecurityProtocol.PLAINTEXT, plaintextListener)).asJava,
-        "rack")
+    val brokers = Seq(
+      new UpdateMetadataBroker()
+        .setId(0)
+        .setRack("rack")
+        .setEndpoints(Seq(
+          new UpdateMetadataEndpoint()
+            .setHost("broker0")
+            .setPort(9092)
+            .setSecurityProtocol(SecurityProtocol.PLAINTEXT.id)
+            .setListener(plaintextListener.value),
+          new UpdateMetadataEndpoint()
+            .setHost("broker0")
+            .setPort(9093)
+            .setSecurityProtocol(SecurityProtocol.PLAINTEXT.id)
+            .setListener(anotherListener.value)
+        ).asJava),
+      new UpdateMetadataBroker()
+        .setId(1)
+        .setRack("rack")
+        .setEndpoints(Seq(
+          new UpdateMetadataEndpoint()
+            .setHost("broker1")
+            .setPort(9092)
+            .setSecurityProtocol(SecurityProtocol.PLAINTEXT.id)
+            .setListener(plaintextListener.value)).asJava)
     )
     val updateMetadataRequest = new UpdateMetadataRequest.Builder(ApiKeys.UPDATE_METADATA.latestVersion, 0,
-      0, 0, Map.empty[TopicPartition, UpdateMetadataRequest.PartitionState].asJava, brokers.asJava).build()
+      0, 0, 0, Seq.empty[UpdateMetadataPartitionState].asJava, brokers.asJava).build()
     metadataCache.updateMetadata(correlationId = 0, updateMetadataRequest)
     (plaintextListener, anotherListener)
   }
@@ -571,7 +952,7 @@ class KafkaApisTest {
     val header = RequestHeader.parse(buffer)
     val context = new RequestContext(header, "1", InetAddress.getLocalHost, KafkaPrincipal.ANONYMOUS,
       listenerName, SecurityProtocol.PLAINTEXT)
-    (request, new RequestChannel.Request(processor = 1, context = context, startTimeNanos =  0, MemoryPool.NONE, buffer,
+    (request, new RequestChannel.Request(processor = 1, context = context, startTimeNanos = 0, MemoryPool.NONE, buffer,
       requestChannelMetrics))
   }
 
@@ -584,7 +965,7 @@ class KafkaApisTest {
     send.writeTo(channel)
     channel.close()
     channel.buffer.getInt() // read the size
-    ResponseHeader.parse(channel.buffer)
+    ResponseHeader.parse(channel.buffer, api.responseHeaderVersion(request.version()))
     val struct = api.responseSchema(request.version).read(channel.buffer)
     AbstractResponse.parseResponse(api, struct, request.version)
   }
@@ -602,12 +983,29 @@ class KafkaApisTest {
 
   private def setupBasicMetadataCache(topic: String, numPartitions: Int): Unit = {
     val replicas = List(0.asInstanceOf[Integer]).asJava
-    val partitionState = new UpdateMetadataRequest.PartitionState(1, 0, 1, replicas, 0, replicas, Collections.emptyList())
+
+    def createPartitionState(partition: Int) = new UpdateMetadataPartitionState()
+      .setTopicName(topic)
+      .setPartitionIndex(partition)
+      .setControllerEpoch(1)
+      .setLeader(0)
+      .setLeaderEpoch(1)
+      .setReplicas(replicas)
+      .setZkVersion(0)
+      .setReplicas(replicas)
+
     val plaintextListener = ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT)
-    val broker = new Broker(0, Seq(new EndPoint("broker0", 9092, SecurityProtocol.PLAINTEXT, plaintextListener)).asJava, "rack")
-    val partitions = (0 until numPartitions).map(new TopicPartition(topic, _) -> partitionState).toMap
+    val broker = new UpdateMetadataBroker()
+      .setId(0)
+      .setRack("rack")
+      .setEndpoints(Seq(new UpdateMetadataEndpoint()
+        .setHost("broker0")
+        .setPort(9092)
+        .setSecurityProtocol(SecurityProtocol.PLAINTEXT.id)
+        .setListener(plaintextListener.value)).asJava)
+    val partitionStates = (0 until numPartitions).map(createPartitionState)
     val updateMetadataRequest = new UpdateMetadataRequest.Builder(ApiKeys.UPDATE_METADATA.latestVersion, 0,
-      0, 0, partitions.asJava, Set(broker).asJava).build()
+      0, 0, 0, partitionStates.asJava, Seq(broker).asJava).build()
     metadataCache.updateMetadata(correlationId = 0, updateMetadataRequest)
   }
 }
diff --git a/core/src/test/scala/unit/kafka/server/KafkaConfigTest.scala b/core/src/test/scala/unit/kafka/server/KafkaConfigTest.scala
index 1d7e687a279f5..afbb7dcba0b40 100755
--- a/core/src/test/scala/unit/kafka/server/KafkaConfigTest.scala
+++ b/core/src/test/scala/unit/kafka/server/KafkaConfigTest.scala
@@ -35,7 +35,7 @@ import org.scalatest.Assertions.intercept
 class KafkaConfigTest {
 
   @Test
-  def testLogRetentionTimeHoursProvided() {
+  def testLogRetentionTimeHoursProvided(): Unit = {
     val props = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 8181)
     props.put(KafkaConfig.LogRetentionTimeHoursProp, "1")
 
@@ -44,7 +44,7 @@ class KafkaConfigTest {
   }
 
   @Test
-  def testLogRetentionTimeMinutesProvided() {
+  def testLogRetentionTimeMinutesProvided(): Unit = {
     val props = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 8181)
     props.put(KafkaConfig.LogRetentionTimeMinutesProp, "30")
 
@@ -53,7 +53,7 @@ class KafkaConfigTest {
   }
 
   @Test
-  def testLogRetentionTimeMsProvided() {
+  def testLogRetentionTimeMsProvided(): Unit = {
     val props = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 8181)
     props.put(KafkaConfig.LogRetentionTimeMillisProp, "1800000")
 
@@ -62,7 +62,7 @@ class KafkaConfigTest {
   }
 
   @Test
-  def testLogRetentionTimeNoConfigProvided() {
+  def testLogRetentionTimeNoConfigProvided(): Unit = {
     val props = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 8181)
 
     val cfg = KafkaConfig.fromProps(props)
@@ -70,7 +70,7 @@ class KafkaConfigTest {
   }
 
   @Test
-  def testLogRetentionTimeBothMinutesAndHoursProvided() {
+  def testLogRetentionTimeBothMinutesAndHoursProvided(): Unit = {
     val props = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 8181)
     props.put(KafkaConfig.LogRetentionTimeMinutesProp, "30")
     props.put(KafkaConfig.LogRetentionTimeHoursProp, "1")
@@ -80,7 +80,7 @@ class KafkaConfigTest {
   }
 
   @Test
-  def testLogRetentionTimeBothMinutesAndMsProvided() {
+  def testLogRetentionTimeBothMinutesAndMsProvided(): Unit = {
     val props = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 8181)
     props.put(KafkaConfig.LogRetentionTimeMillisProp, "1800000")
     props.put(KafkaConfig.LogRetentionTimeMinutesProp, "10")
@@ -90,7 +90,7 @@ class KafkaConfigTest {
   }
 
   @Test
-  def testLogRetentionUnlimited() {
+  def testLogRetentionUnlimited(): Unit = {
     val props1 = TestUtils.createBrokerConfig(0,TestUtils.MockZkConnect, port = 8181)
     val props2 = TestUtils.createBrokerConfig(0,TestUtils.MockZkConnect, port = 8181)
     val props3 = TestUtils.createBrokerConfig(0,TestUtils.MockZkConnect, port = 8181)
@@ -144,7 +144,7 @@ class KafkaConfigTest {
   }
 
   @Test
-  def testAdvertiseDefaults() {
+  def testAdvertiseDefaults(): Unit = {
     val port = "9999"
     val hostName = "fake-host"
 
@@ -160,7 +160,7 @@ class KafkaConfigTest {
   }
 
   @Test
-  def testAdvertiseConfigured() {
+  def testAdvertiseConfigured(): Unit = {
     val advertisedHostName = "routable-host"
     val advertisedPort = "1234"
 
@@ -177,7 +177,7 @@ class KafkaConfigTest {
   }
 
   @Test
-  def testAdvertisePortDefault() {
+  def testAdvertisePortDefault(): Unit = {
     val advertisedHostName = "routable-host"
     val port = "9999"
 
@@ -194,7 +194,7 @@ class KafkaConfigTest {
   }
 
   @Test
-  def testAdvertiseHostNameDefault() {
+  def testAdvertiseHostNameDefault(): Unit = {
     val hostName = "routable-host"
     val advertisedPort = "9999"
 
@@ -211,7 +211,7 @@ class KafkaConfigTest {
   }
 
   @Test
-  def testDuplicateListeners() {
+  def testDuplicateListeners(): Unit = {
     val props = new Properties()
     props.put(KafkaConfig.BrokerIdProp, "1")
     props.put(KafkaConfig.ZkConnectProp, "localhost:2181")
@@ -255,7 +255,7 @@ class KafkaConfigTest {
   }
 
   @Test
-  def testBadListenerProtocol() {
+  def testBadListenerProtocol(): Unit = {
     val props = new Properties()
     props.put(KafkaConfig.BrokerIdProp, "1")
     props.put(KafkaConfig.ZkConnectProp, "localhost:2181")
@@ -354,7 +354,7 @@ class KafkaConfigTest {
   }
 
   @Test
-  def testCaseInsensitiveListenerProtocol() {
+  def testCaseInsensitiveListenerProtocol(): Unit = {
     val props = new Properties()
     props.put(KafkaConfig.BrokerIdProp, "1")
     props.put(KafkaConfig.ZkConnectProp, "localhost:2181")
@@ -369,7 +369,7 @@ class KafkaConfigTest {
     CoreUtils.listenerListToEndPoints(listenerList, securityProtocolMap)
 
   @Test
-  def testListenerDefaults() {
+  def testListenerDefaults(): Unit = {
     val props = new Properties()
     props.put(KafkaConfig.BrokerIdProp, "1")
     props.put(KafkaConfig.ZkConnectProp, "localhost:2181")
@@ -398,7 +398,7 @@ class KafkaConfigTest {
   }
 
   @Test
-  def testVersionConfiguration() {
+  def testVersionConfiguration(): Unit = {
     val props = new Properties()
     props.put(KafkaConfig.BrokerIdProp, "1")
     props.put(KafkaConfig.ZkConnectProp, "localhost:2181")
@@ -432,7 +432,7 @@ class KafkaConfigTest {
   }
 
   @Test
-  def testUncleanLeaderElectionDefault() {
+  def testUncleanLeaderElectionDefault(): Unit = {
     val props = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 8181)
     val serverConfig = KafkaConfig.fromProps(props)
 
@@ -440,7 +440,7 @@ class KafkaConfigTest {
   }
 
   @Test
-  def testUncleanElectionDisabled() {
+  def testUncleanElectionDisabled(): Unit = {
     val props = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 8181)
     props.put(KafkaConfig.UncleanLeaderElectionEnableProp, String.valueOf(false))
     val serverConfig = KafkaConfig.fromProps(props)
@@ -449,7 +449,7 @@ class KafkaConfigTest {
   }
 
   @Test
-  def testUncleanElectionEnabled() {
+  def testUncleanElectionEnabled(): Unit = {
     val props = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 8181)
     props.put(KafkaConfig.UncleanLeaderElectionEnableProp, String.valueOf(true))
     val serverConfig = KafkaConfig.fromProps(props)
@@ -458,7 +458,7 @@ class KafkaConfigTest {
   }
 
   @Test
-  def testUncleanElectionInvalid() {
+  def testUncleanElectionInvalid(): Unit = {
     val props = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 8181)
     props.put(KafkaConfig.UncleanLeaderElectionEnableProp, "invalid")
 
@@ -468,7 +468,7 @@ class KafkaConfigTest {
   }
 
   @Test
-  def testLogRollTimeMsProvided() {
+  def testLogRollTimeMsProvided(): Unit = {
     val props = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 8181)
     props.put(KafkaConfig.LogRollTimeMillisProp, "1800000")
 
@@ -477,7 +477,7 @@ class KafkaConfigTest {
   }
 
   @Test
-  def testLogRollTimeBothMsAndHoursProvided() {
+  def testLogRollTimeBothMsAndHoursProvided(): Unit = {
     val props = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 8181)
     props.put(KafkaConfig.LogRollTimeMillisProp, "1800000")
     props.put(KafkaConfig.LogRollTimeHoursProp, "1")
@@ -487,7 +487,7 @@ class KafkaConfigTest {
   }
 
   @Test
-  def testLogRollTimeNoConfigProvided() {
+  def testLogRollTimeNoConfigProvided(): Unit = {
     val props = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 8181)
 
     val cfg = KafkaConfig.fromProps(props)
@@ -495,7 +495,7 @@ class KafkaConfigTest {
   }
 
   @Test
-  def testDefaultCompressionType() {
+  def testDefaultCompressionType(): Unit = {
     val props = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 8181)
     val serverConfig = KafkaConfig.fromProps(props)
 
@@ -503,7 +503,7 @@ class KafkaConfigTest {
   }
 
   @Test
-  def testValidCompressionType() {
+  def testValidCompressionType(): Unit = {
     val props = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 8181)
     props.put("compression.type", "gzip")
     val serverConfig = KafkaConfig.fromProps(props)
@@ -512,7 +512,7 @@ class KafkaConfigTest {
   }
 
   @Test
-  def testInvalidCompressionType() {
+  def testInvalidCompressionType(): Unit = {
     val props = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 8181)
     props.put(KafkaConfig.CompressionTypeProp, "abc")
     intercept[IllegalArgumentException] {
@@ -521,7 +521,7 @@ class KafkaConfigTest {
   }
 
   @Test
-  def testInvalidInterBrokerSecurityProtocol() {
+  def testInvalidInterBrokerSecurityProtocol(): Unit = {
     val props = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 8181)
     props.put(KafkaConfig.ListenersProp, "SSL://localhost:0")
     props.put(KafkaConfig.InterBrokerSecurityProtocolProp, SecurityProtocol.PLAINTEXT.toString)
@@ -531,7 +531,7 @@ class KafkaConfigTest {
   }
 
   @Test
-  def testEqualAdvertisedListenersProtocol() {
+  def testEqualAdvertisedListenersProtocol(): Unit = {
     val props = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 8181)
     props.put(KafkaConfig.ListenersProp, "PLAINTEXT://localhost:9092,SSL://localhost:9093")
     props.put(KafkaConfig.AdvertisedListenersProp, "PLAINTEXT://localhost:9092,SSL://localhost:9093")
@@ -539,7 +539,7 @@ class KafkaConfigTest {
   }
 
   @Test
-  def testInvalidAdvertisedListenersProtocol() {
+  def testInvalidAdvertisedListenersProtocol(): Unit = {
     val props = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 8181)
     props.put(KafkaConfig.ListenersProp, "TRACE://localhost:9091,SSL://localhost:9093")
     props.put(KafkaConfig.AdvertisedListenersProp, "PLAINTEXT://localhost:9092")
@@ -573,7 +573,7 @@ class KafkaConfigTest {
   }
 
   @Test
-  def testFromPropsInvalid() {
+  def testFromPropsInvalid(): Unit = {
     def getBaseProperties(): Properties = {
       val validRequiredProperties = new Properties()
       validRequiredProperties.put(KafkaConfig.ZkConnectProp, "127.0.0.1:2181")
@@ -599,6 +599,8 @@ class KafkaConfigTest {
         case KafkaConfig.NumReplicaAlterLogDirsThreadsProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number")
         case KafkaConfig.QueuedMaxBytesProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number")
         case KafkaConfig.RequestTimeoutMsProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number")
+        case KafkaConfig.HeapDumpFolderProp => //ignore string
+        case KafkaConfig.HeapDumpTimeoutProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number")
 
         case KafkaConfig.AuthorizerClassNameProp => //ignore string
         case KafkaConfig.CreateTopicPolicyClassNameProp => //ignore string
@@ -609,6 +611,8 @@ class KafkaConfigTest {
         case KafkaConfig.AdvertisedPortProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number")
         case KafkaConfig.SocketSendBufferBytesProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number")
         case KafkaConfig.SocketReceiveBufferBytesProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number")
+        case KafkaConfig.SocketRequestCommonBytesProp =>
+        case KafkaConfig.SocketRequestBufferCacheSizeProp =>
         case KafkaConfig.MaxConnectionsPerIpOverridesProp =>
           assertPropertyInvalid(getBaseProperties(), name, "127.0.0.1:not_a_number")
         case KafkaConfig.ConnectionsMaxIdleMsProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number")
@@ -635,6 +639,7 @@ class KafkaConfigTest {
         case KafkaConfig.LogCleanerEnableProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_boolean")
         case KafkaConfig.LogCleanerDeleteRetentionMsProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number")
         case KafkaConfig.LogCleanerMinCompactionLagMsProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number")
+        case KafkaConfig.LogCleanerMaxCompactionLagMsProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number")
         case KafkaConfig.LogCleanerMinCleanRatioProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number")
         case KafkaConfig.LogIndexSizeMaxBytesProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number", "3")
         case KafkaConfig.LogFlushIntervalMessagesProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number", "0")
@@ -654,6 +659,7 @@ class KafkaConfigTest {
         case KafkaConfig.ReplicaFetchWaitMaxMsProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number")
         case KafkaConfig.ReplicaFetchMinBytesProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number")
         case KafkaConfig.ReplicaFetchResponseMaxBytesProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number")
+        case KafkaConfig.ReplicaSelectorClassProp => // Ignore string
         case KafkaConfig.NumReplicaFetchersProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number")
         case KafkaConfig.ReplicaHighWatermarkCheckpointIntervalMsProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number")
         case KafkaConfig.FetchPurgatoryPurgeIntervalRequestsProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number")
@@ -680,6 +686,9 @@ class KafkaConfigTest {
         case KafkaConfig.OffsetsRetentionCheckIntervalMsProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number", "0")
         case KafkaConfig.OffsetCommitTimeoutMsProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number", "0")
         case KafkaConfig.OffsetCommitRequiredAcksProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number", "-2")
+        case KafkaConfig.OffsetsTopicMaxMessageBytesProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number", "-1")
+        case KafkaConfig.OffsetsTopicMinInSyncReplicasProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number", "0")
+        case KafkaConfig.OffsetsTopicMinCompactionLagMsProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number", "-1")
         case KafkaConfig.TransactionalIdExpirationMsProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number", "0", "-2")
         case KafkaConfig.TransactionsMaxTimeoutMsProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number", "0", "-2")
         case KafkaConfig.TransactionsTopicMinISRProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number", "0", "-2")
@@ -703,6 +712,7 @@ class KafkaConfigTest {
         case KafkaConfig.ConnectionsMaxReauthMsProp =>
         case KafkaConfig.SslProtocolProp => // ignore string
         case KafkaConfig.SslProviderProp => // ignore string
+        case KafkaConfig.SslContextProviderClassProp => // ignore string
         case KafkaConfig.SslEnabledProtocolsProp =>
         case KafkaConfig.SslKeystoreTypeProp => // ignore string
         case KafkaConfig.SslKeystoreLocationProp => // ignore string
@@ -738,6 +748,9 @@ class KafkaConfigTest {
         case KafkaConfig.SaslLoginRefreshMinPeriodSecondsProp =>
         case KafkaConfig.SaslLoginRefreshBufferSecondsProp =>
 
+        // Security config
+        case KafkaConfig.securityProviderClassProp =>
+
         // Password encoder configs
         case KafkaConfig.PasswordEncoderSecretProp =>
         case KafkaConfig.PasswordEncoderOldSecretProp =>
@@ -756,6 +769,10 @@ class KafkaConfigTest {
         case KafkaConfig.KafkaMetricsReporterClassesProp => // ignore
         case KafkaConfig.KafkaMetricsPollingIntervalSecondsProp => //ignore
 
+        // Broker-side observer configs
+        case KafkaConfig.ObserverClassNameProp => // ignore since even if the class name is invalid, a NoOpObserver class is used instead
+        case KafkaConfig.ObserverShutdownTimeoutMsProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number", "-1", "0")
+
         case _ => assertPropertyInvalid(getBaseProperties(), name, "not_a_number", "-1")
       }
     })
@@ -811,7 +828,7 @@ class KafkaConfigTest {
   }
 
   @Test
-  def testNonroutableAdvertisedListeners() {
+  def testNonroutableAdvertisedListeners(): Unit = {
     val props = new Properties()
     props.put(KafkaConfig.ZkConnectProp, "127.0.0.1:2181")
     props.put(KafkaConfig.ListenersProp, "PLAINTEXT://0.0.0.0:9092")
@@ -819,7 +836,7 @@ class KafkaConfigTest {
   }
 
   @Test
-  def testMaxConnectionsPerIpProp() {
+  def testMaxConnectionsPerIpProp(): Unit = {
     val props = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 8181)
     props.put(KafkaConfig.MaxConnectionsPerIpProp, "0")
     assertFalse(isValidKafkaConfig(props))
@@ -829,7 +846,7 @@ class KafkaConfigTest {
     assertFalse(isValidKafkaConfig(props))
   }
 
-  private def assertPropertyInvalid(validRequiredProps: => Properties, name: String, values: Any*) {
+  private def assertPropertyInvalid(validRequiredProps: => Properties, name: String, values: Any*): Unit = {
     values.foreach((value) => {
       val props = validRequiredProps
       props.setProperty(name, value.toString)
diff --git a/core/src/test/scala/unit/kafka/server/KafkaMetricReporterClusterIdTest.scala b/core/src/test/scala/unit/kafka/server/KafkaMetricReporterClusterIdTest.scala
index 4d514a0550df8..62d99c08b53e0 100755
--- a/core/src/test/scala/unit/kafka/server/KafkaMetricReporterClusterIdTest.scala
+++ b/core/src/test/scala/unit/kafka/server/KafkaMetricReporterClusterIdTest.scala
@@ -50,7 +50,7 @@ object KafkaMetricReporterClusterIdTest {
 
   class MockBrokerMetricsReporter extends MockMetricsReporter with ClusterResourceListener {
 
-    override def onUpdate(clusterMetadata: ClusterResource) {
+    override def onUpdate(clusterMetadata: ClusterResource): Unit = {
       MockBrokerMetricsReporter.CLUSTER_META.set(clusterMetadata)
     }
 
@@ -80,7 +80,7 @@ class KafkaMetricReporterClusterIdTest extends ZooKeeperTestHarness {
   var config: KafkaConfig = null
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     super.setUp()
     val props = TestUtils.createBrokerConfig(1, zkConnect)
     props.setProperty(KafkaConfig.KafkaMetricsReporterClassesProp, "kafka.server.KafkaMetricReporterClusterIdTest$MockKafkaMetricsReporter")
@@ -88,12 +88,12 @@ class KafkaMetricReporterClusterIdTest extends ZooKeeperTestHarness {
     props.setProperty(KafkaConfig.BrokerIdGenerationEnableProp, "true")
     props.setProperty(KafkaConfig.BrokerIdProp, "-1")
     config = KafkaConfig.fromProps(props)
-    server = KafkaServerStartable.fromProps(props)
+    server = KafkaServerStartable.fromProps(props, threadNamePrefix = Option(this.getClass.getName))
     server.startup()
   }
 
   @Test
-  def testClusterIdPresent() {
+  def testClusterIdPresent(): Unit = {
     assertEquals("", KafkaMetricReporterClusterIdTest.setupError.get())
 
     assertNotNull(KafkaMetricReporterClusterIdTest.MockKafkaMetricsReporter.CLUSTER_META)
@@ -104,13 +104,15 @@ class KafkaMetricReporterClusterIdTest extends ZooKeeperTestHarness {
 
     assertEquals(KafkaMetricReporterClusterIdTest.MockKafkaMetricsReporter.CLUSTER_META.get().clusterId(),
       KafkaMetricReporterClusterIdTest.MockBrokerMetricsReporter.CLUSTER_META.get().clusterId())
+
+    server.shutdown()
+    TestUtils.assertNoNonDaemonThreads(this.getClass.getName)
   }
 
   @After
-  override def tearDown() {
+  override def tearDown(): Unit = {
     server.shutdown()
     CoreUtils.delete(config.logDirs)
-    TestUtils.verifyNonDaemonThreadsStatus(this.getClass.getName)
     super.tearDown()
   }
 }
diff --git a/core/src/test/scala/unit/kafka/server/KafkaMetricReporterExceptionHandlingTest.scala b/core/src/test/scala/unit/kafka/server/KafkaMetricReporterExceptionHandlingTest.scala
index 30f3b234613ea..d6f37cd4f05bb 100644
--- a/core/src/test/scala/unit/kafka/server/KafkaMetricReporterExceptionHandlingTest.scala
+++ b/core/src/test/scala/unit/kafka/server/KafkaMetricReporterExceptionHandlingTest.scala
@@ -15,35 +15,36 @@
 package kafka.server
 
 import java.net.Socket
-import java.util.Properties
+import java.util.{Collections, Properties}
 
 import kafka.utils.TestUtils
 import org.apache.kafka.common.network.ListenerName
-import org.apache.kafka.common.requests.{ListGroupsRequest,ListGroupsResponse}
+import org.apache.kafka.common.requests.{ListGroupsRequest, ListGroupsResponse}
 import org.apache.kafka.common.metrics.MetricsReporter
 import org.apache.kafka.common.metrics.KafkaMetric
 import org.apache.kafka.common.security.auth.SecurityProtocol
-import org.apache.kafka.common.protocol.Errors
-
+import org.apache.kafka.common.protocol.{ApiKeys, Errors}
 import org.junit.Assert._
 import org.junit.{Before, Test}
 import org.junit.After
 import java.util.concurrent.atomic.AtomicInteger
 
+import org.apache.kafka.common.message.ListGroupsRequestData
+
 /*
  * this test checks that a reporter that throws an exception will not affect other reporters
  * and will not affect the broker's message handling
  */
 class KafkaMetricReporterExceptionHandlingTest extends BaseRequestTest {
 
-  override def numBrokers: Int = 1
+  override def brokerCount: Int = 1
 
-  override def propertyOverrides(properties: Properties): Unit = {
+  override def brokerPropertyOverrides(properties: Properties): Unit = {
     properties.put(KafkaConfig.MetricReporterClassesProp, classOf[KafkaMetricReporterExceptionHandlingTest.BadReporter].getName + "," + classOf[KafkaMetricReporterExceptionHandlingTest.GoodReporter].getName)
   }
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     super.setUp()
 
     // need a quota prop to register a "throttle-time" metrics after server startup
@@ -53,7 +54,7 @@ class KafkaMetricReporterExceptionHandlingTest extends BaseRequestTest {
   }
 
   @After
-  override def tearDown() {
+  override def tearDown(): Unit = {
     KafkaMetricReporterExceptionHandlingTest.goodReporterRegistered.set(0)
     KafkaMetricReporterExceptionHandlingTest.badReporterRegistered.set(0)
     
@@ -61,15 +62,17 @@ class KafkaMetricReporterExceptionHandlingTest extends BaseRequestTest {
   }
 
   @Test
-  def testBothReportersAreInvoked() {
+  def testBothReportersAreInvoked(): Unit = {
     val port = anySocketServer.boundPort(ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT))
     val socket = new Socket("localhost", port)
     socket.setSoTimeout(10000)
 
     try {
       TestUtils.retry(10000) {
-        val error = new ListGroupsResponse(requestResponse(socket, "clientId", 0, new ListGroupsRequest.Builder())).error()
-        assertEquals(Errors.NONE, error)
+        val error = new ListGroupsResponse(
+          requestResponse(socket, "clientId", 0, new ListGroupsRequest.Builder(new ListGroupsRequestData)), ApiKeys.LIST_GROUPS.latestVersion)
+          .errorCounts()
+        assertEquals(Collections.singletonMap(Errors.NONE, 1), error)
         assertEquals(KafkaMetricReporterExceptionHandlingTest.goodReporterRegistered.get, KafkaMetricReporterExceptionHandlingTest.badReporterRegistered.get)
         assertTrue(KafkaMetricReporterExceptionHandlingTest.goodReporterRegistered.get > 0)
       }
@@ -85,28 +88,28 @@ object KafkaMetricReporterExceptionHandlingTest {
 
   class GoodReporter extends MetricsReporter {
 
-    def configure(configs: java.util.Map[String, _]) {
+    def configure(configs: java.util.Map[String, _]): Unit = {
     }
 
-    def init(metrics: java.util.List[KafkaMetric]) {
+    def init(metrics: java.util.List[KafkaMetric]): Unit = {
     }
 
-    def metricChange(metric: KafkaMetric) {
+    def metricChange(metric: KafkaMetric): Unit = {
       if (metric.metricName.group == "Request") {
         goodReporterRegistered.incrementAndGet
       }
     }
 
-    def metricRemoval(metric: KafkaMetric) {
+    def metricRemoval(metric: KafkaMetric): Unit = {
     }
 
-    def close() {
+    def close(): Unit = {
     }
   }
 
   class BadReporter extends GoodReporter {
 
-    override def metricChange(metric: KafkaMetric) {
+    override def metricChange(metric: KafkaMetric): Unit = {
       if (metric.metricName.group == "Request") {
         badReporterRegistered.incrementAndGet
         throw new RuntimeException(metric.metricName.toString)
diff --git a/core/src/test/scala/unit/kafka/server/KafkaServerTest.scala b/core/src/test/scala/unit/kafka/server/KafkaServerTest.scala
index d78821a2ca5a6..a285fadf55b79 100755
--- a/core/src/test/scala/unit/kafka/server/KafkaServerTest.scala
+++ b/core/src/test/scala/unit/kafka/server/KafkaServerTest.scala
@@ -20,11 +20,12 @@ package kafka.server
 import kafka.utils.TestUtils
 import kafka.zk.ZooKeeperTestHarness
 import org.junit.Test
+import org.scalatest.Assertions.intercept
 
 class KafkaServerTest extends ZooKeeperTestHarness {
 
   @Test
-  def testAlreadyRegisteredAdvertisedListeners() {
+  def testAlreadyRegisteredAdvertisedListeners(): Unit = {
     //start a server with a advertised listener
     val server1 = createServer(1, "myhost", TestUtils.RandomPort)
 
diff --git a/core/src/test/scala/unit/kafka/server/LeaderElectionTest.scala b/core/src/test/scala/unit/kafka/server/LeaderElectionTest.scala
index 1772f324d2fb3..8b97d65e4d1bb 100755
--- a/core/src/test/scala/unit/kafka/server/LeaderElectionTest.scala
+++ b/core/src/test/scala/unit/kafka/server/LeaderElectionTest.scala
@@ -17,6 +17,8 @@
 
 package kafka.server
 
+import java.util.concurrent.{CountDownLatch, TimeUnit}
+
 import org.apache.kafka.common.TopicPartition
 
 import scala.collection.JavaConverters._
@@ -28,6 +30,8 @@ import kafka.cluster.Broker
 import kafka.controller.{ControllerChannelManager, ControllerContext, StateChangeLogger}
 import kafka.utils.TestUtils._
 import kafka.zk.ZooKeeperTestHarness
+import org.apache.kafka.common.message.LeaderAndIsrRequestData
+import org.apache.kafka.common.message.LeaderAndIsrRequestData.LeaderAndIsrPartitionState
 import org.apache.kafka.common.metrics.Metrics
 import org.apache.kafka.common.network.ListenerName
 import org.apache.kafka.common.protocol.{ApiKeys, Errors}
@@ -44,7 +48,7 @@ class LeaderElectionTest extends ZooKeeperTestHarness {
   var staleControllerEpochDetected = false
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     super.setUp()
 
     val configProps1 = TestUtils.createBrokerConfig(brokerId1, zkConnect, enableControlledShutdown = false)
@@ -60,7 +64,7 @@ class LeaderElectionTest extends ZooKeeperTestHarness {
   }
 
   @After
-  override def tearDown() {
+  override def tearDown(): Unit = {
     TestUtils.shutdownServers(servers)
     super.tearDown()
   }
@@ -71,47 +75,42 @@ class LeaderElectionTest extends ZooKeeperTestHarness {
     val topic = "new-topic"
     val partitionId = 0
 
+    TestUtils.waitUntilBrokerMetadataIsPropagated(servers)
+
     // create topic with 1 partition, 2 replicas, one on each broker
     val leader1 = createTopic(zkClient, topic, partitionReplicaAssignment = Map(0 -> Seq(0, 1)), servers = servers)(0)
 
     val leaderEpoch1 = zkClient.getEpochForPartition(new TopicPartition(topic, partitionId)).get
-    debug("leader Epoch: " + leaderEpoch1)
-    debug("Leader is elected to be: %s".format(leader1))
-    // NOTE: this is to avoid transient test failures
-    assertTrue("Leader could be broker 0 or broker 1", leader1 == 0 || leader1 == 1)
+    assertTrue("Leader should be broker 0", leader1 == 0)
     assertEquals("First epoch value should be 0", 0, leaderEpoch1)
 
-    // kill the server hosting the preferred replica
-    servers.last.shutdown()
+    // kill the server hosting the preferred replica/initial leader
+    servers.head.shutdown()
     // check if leader moves to the other server
-    val leader2 = waitUntilLeaderIsElectedOrChanged(zkClient, topic, partitionId,
-                                                    oldLeaderOpt = if (leader1 == 0) None else Some(leader1))
+    val leader2 = waitUntilLeaderIsElectedOrChanged(zkClient, topic, partitionId, oldLeaderOpt = Some(leader1))
     val leaderEpoch2 = zkClient.getEpochForPartition(new TopicPartition(topic, partitionId)).get
-    debug("Leader is elected to be: %s".format(leader1))
-    debug("leader Epoch: " + leaderEpoch2)
-    assertEquals("Leader must move to broker 0", 0, leader2)
-    if (leader1 == leader2)
-      assertEquals("Second epoch value should be " + leaderEpoch1+1, leaderEpoch1+1, leaderEpoch2)
-    else
-      assertEquals("Second epoch value should be %d".format(leaderEpoch1+1) , leaderEpoch1+1, leaderEpoch2)
-
-    servers.last.startup()
-    servers.head.shutdown()
+    assertEquals("Leader must move to broker 1", 1, leader2)
+    // new leaderEpoch will be leaderEpoch1+2, one increment during ReplicaStateMachine.startup()-> handleStateChanges
+    // for offline replica and one increment during PartitionStateMachine.triggerOnlinePartitionStateChange()
+    assertEquals("Second epoch value should be %d".format(leaderEpoch1 + 2) , leaderEpoch1 + 2, leaderEpoch2)
+
+    servers.head.startup()
+    //make sure second server joins the ISR
+    TestUtils.waitUntilTrue(() => {
+      servers.last.metadataCache.getPartitionInfo(topic, partitionId).exists(_.isr.size == 2)
+    }, "Inconsistent metadata after second broker startup")
+
+    servers.last.shutdown()
+
     Thread.sleep(zookeeper.tickTime)
-    val leader3 = waitUntilLeaderIsElectedOrChanged(zkClient, topic, partitionId,
-                                                    oldLeaderOpt = if (leader2 == 1) None else Some(leader2))
+    val leader3 = waitUntilLeaderIsElectedOrChanged(zkClient, topic, partitionId, oldLeaderOpt = Some(leader2))
     val leaderEpoch3 = zkClient.getEpochForPartition(new TopicPartition(topic, partitionId)).get
-    debug("leader Epoch: " + leaderEpoch3)
-    debug("Leader is elected to be: %s".format(leader3))
-    assertEquals("Leader must return to 1", 1, leader3)
-    if (leader2 == leader3)
-      assertEquals("Second epoch value should be " + leaderEpoch2, leaderEpoch2, leaderEpoch3)
-    else
-      assertEquals("Second epoch value should be %d".format(leaderEpoch2+1) , leaderEpoch2+1, leaderEpoch3)
+    assertEquals("Leader must return to 0", 0, leader3)
+    assertEquals("Second epoch value should be %d".format(leaderEpoch2 + 2) , leaderEpoch2 + 2, leaderEpoch3)
   }
 
   @Test
-  def testLeaderElectionWithStaleControllerEpoch() {
+  def testLeaderElectionWithStaleControllerEpoch(): Unit = {
     // start 2 brokers
     val topic = "new-topic"
     val partitionId = 0
@@ -145,16 +144,23 @@ class LeaderElectionTest extends ZooKeeperTestHarness {
     controllerChannelManager.startup()
     try {
       val staleControllerEpoch = 0
-      val partitionStates = Map(
-        new TopicPartition(topic, partitionId) -> new LeaderAndIsrRequest.PartitionState(2, brokerId2, LeaderAndIsr.initialLeaderEpoch,
-          Seq(brokerId1, brokerId2).map(Integer.valueOf).asJava, LeaderAndIsr.initialZKVersion,
-          Seq(0, 1).map(Integer.valueOf).asJava, false)
+      val partitionStates = Seq(
+        new LeaderAndIsrPartitionState()
+          .setTopicName(topic)
+          .setPartitionIndex(partitionId)
+          .setControllerEpoch(2)
+          .setLeader(brokerId2)
+          .setLeaderEpoch(LeaderAndIsr.initialLeaderEpoch)
+          .setIsr(Seq(brokerId1, brokerId2).map(Integer.valueOf).asJava)
+          .setZkVersion(LeaderAndIsr.initialZKVersion)
+          .setReplicas(Seq(0, 1).map(Integer.valueOf).asJava)
+          .setIsNew(false)
       )
       val requestBuilder = new LeaderAndIsrRequest.Builder(
-        ApiKeys.LEADER_AND_ISR.latestVersion, controllerId, staleControllerEpoch, servers(brokerId2).kafkaController.brokerEpoch ,partitionStates.asJava, nodes.toSet.asJava)
+        ApiKeys.LEADER_AND_ISR.latestVersion, controllerId, staleControllerEpoch, servers(brokerId2).kafkaController.brokerEpoch, 
+        servers(brokerId2).kafkaController.brokerEpoch, partitionStates.asJava, nodes.toSet.asJava)
 
-      controllerChannelManager.sendRequest(brokerId2, ApiKeys.LEADER_AND_ISR, requestBuilder,
-        staleControllerEpochCallback)
+      controllerChannelManager.sendRequest(brokerId2, requestBuilder, staleControllerEpochCallback)
       TestUtils.waitUntilTrue(() => staleControllerEpochDetected, "Controller epoch should be stale")
       assertTrue("Stale controller epoch not detected by the broker", staleControllerEpochDetected)
     } finally {
diff --git a/core/src/test/scala/unit/kafka/server/LogDirFailureTest.scala b/core/src/test/scala/unit/kafka/server/LogDirFailureTest.scala
index d56a9f0559c5e..ab22182ac60c1 100644
--- a/core/src/test/scala/unit/kafka/server/LogDirFailureTest.scala
+++ b/core/src/test/scala/unit/kafka/server/LogDirFailureTest.scala
@@ -20,17 +20,18 @@ import java.io.File
 import java.util.Collections
 import java.util.concurrent.{ExecutionException, TimeUnit}
 
-import kafka.server.LogDirFailureTest._
 import kafka.api.IntegrationTestHarness
 import kafka.controller.{OfflineReplica, PartitionAndReplica}
+import kafka.server.LogDirFailureTest._
 import kafka.utils.{CoreUtils, Exit, TestUtils}
 import org.apache.kafka.clients.consumer.KafkaConsumer
 import org.apache.kafka.clients.producer.{ProducerConfig, ProducerRecord}
 import org.apache.kafka.common.TopicPartition
-import org.apache.kafka.common.utils.Utils
 import org.apache.kafka.common.errors.{KafkaStorageException, NotLeaderForPartitionException}
-import org.junit.{Before, Test}
+import org.apache.kafka.common.utils.Utils
 import org.junit.Assert.{assertEquals, assertFalse, assertTrue}
+import org.junit.{Before, Test}
+import org.scalatest.Assertions.fail
 
 import scala.collection.JavaConverters._
 
@@ -41,7 +42,7 @@ class LogDirFailureTest extends IntegrationTestHarness {
 
   val producerCount: Int = 1
   val consumerCount: Int = 1
-  val serverCount: Int = 2
+  val brokerCount: Int = 2
   private val topic = "topic"
   private val partitionNum = 12
   override val logDirCount = 3
@@ -50,19 +51,24 @@ class LogDirFailureTest extends IntegrationTestHarness {
   this.serverConfig.setProperty(KafkaConfig.NumReplicaFetchersProp, "1")
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     super.setUp()
-    createTopic(topic, partitionNum, serverCount)
+    createTopic(topic, partitionNum, brokerCount)
   }
 
   @Test
-  def testIOExceptionDuringLogRoll() {
+  def testProduceErrorFromFailureOnLogRoll(): Unit = {
+    testProduceErrorsFromLogDirFailureOnLeader(Roll)
+  }
+
+  @Test
+  def testIOExceptionDuringLogRoll(): Unit = {
     testProduceAfterLogDirFailureOnLeader(Roll)
   }
 
   @Test
   // Broker should halt on any log directory failure if inter-broker protocol < 1.0
-  def brokerWithOldInterBrokerProtocolShouldHaltOnLogDirFailure() {
+  def brokerWithOldInterBrokerProtocolShouldHaltOnLogDirFailure(): Unit = {
     @volatile var statusCodeOption: Option[Int] = None
     Exit.setHaltProcedure { (statusCode, _) =>
       statusCodeOption = Some(statusCode)
@@ -71,7 +77,7 @@ class LogDirFailureTest extends IntegrationTestHarness {
 
     var server: KafkaServer = null
     try {
-      val props = TestUtils.createBrokerConfig(serverCount, zkConnect, logDirCount = 3)
+      val props = TestUtils.createBrokerConfig(brokerCount, zkConnect, logDirCount = 3)
       props.put(KafkaConfig.InterBrokerProtocolVersionProp, "0.11.0")
       props.put(KafkaConfig.LogMessageFormatVersionProp, "0.11.0")
       val kafkaConfig = KafkaConfig.fromProps(props)
@@ -91,12 +97,17 @@ class LogDirFailureTest extends IntegrationTestHarness {
   }
 
   @Test
-  def testIOExceptionDuringCheckpoint() {
+  def testProduceErrorFromFailureOnCheckpoint(): Unit = {
+    testProduceErrorsFromLogDirFailureOnLeader(Checkpoint)
+  }
+
+  @Test
+  def testIOExceptionDuringCheckpoint(): Unit = {
     testProduceAfterLogDirFailureOnLeader(Checkpoint)
   }
 
   @Test
-  def testReplicaFetcherThreadAfterLogDirFailureOnFollower() {
+  def testReplicaFetcherThreadAfterLogDirFailureOnFollower(): Unit = {
     this.producerConfig.setProperty(ProducerConfig.RETRIES_CONFIG, "0")
     val producer = createProducer()
     val partition = new TopicPartition(topic, 0)
@@ -111,23 +122,23 @@ class LogDirFailureTest extends IntegrationTestHarness {
     // Send a message to another partition whose leader is the same as partition 0
     // so that ReplicaFetcherThread on the follower will get response from leader immediately
     val anotherPartitionWithTheSameLeader = (1 until partitionNum).find { i =>
-      leaderServer.replicaManager.getPartition(new TopicPartition(topic, i)).flatMap(_.leaderReplicaIfLocal).isDefined
+      leaderServer.replicaManager.nonOfflinePartition(new TopicPartition(topic, i))
+        .flatMap(_.leaderLogIfLocal).isDefined
     }.get
     val record = new ProducerRecord[Array[Byte], Array[Byte]](topic, anotherPartitionWithTheSameLeader, topic.getBytes, "message".getBytes)
     // When producer.send(...).get returns, it is guaranteed that ReplicaFetcherThread on the follower
     // has fetched from the leader and attempts to append to the offline replica.
     producer.send(record).get
 
-    assertEquals(serverCount, leaderServer.replicaManager.getPartition(new TopicPartition(topic, anotherPartitionWithTheSameLeader)).get.inSyncReplicas.size)
+    assertEquals(brokerCount, leaderServer.replicaManager.nonOfflinePartition(new TopicPartition(topic, anotherPartitionWithTheSameLeader))
+      .get.inSyncReplicaIds.size)
     followerServer.replicaManager.replicaFetcherManager.fetcherThreadMap.values.foreach { thread =>
       assertFalse("ReplicaFetcherThread should still be working if its partition count > 0", thread.isShutdownComplete)
     }
   }
 
-  def testProduceAfterLogDirFailureOnLeader(failureType: LogDirFailureType) {
-    val consumer = createConsumer()
-    subscribeAndWaitForAssignment(topic, consumer)
-
+  def testProduceErrorsFromLogDirFailureOnLeader(failureType: LogDirFailureType): Unit = {
+    // Disable retries to allow exception to bubble up for validation
     this.producerConfig.setProperty(ProducerConfig.RETRIES_CONFIG, "0")
     val producer = createProducer()
 
@@ -137,33 +148,9 @@ class LogDirFailureTest extends IntegrationTestHarness {
     val leaderServerId = producer.partitionsFor(topic).asScala.find(_.partition() == 0).get.leader().id()
     val leaderServer = servers.find(_.config.brokerId == leaderServerId).get
 
-    // The first send() should succeed
-    producer.send(record).get()
-    TestUtils.consumeRecords(consumer, 1)
-
-    // Make log directory of the partition on the leader broker inaccessible by replacing it with a file
-    val replica = leaderServer.replicaManager.localReplicaOrException(partition)
-    val logDir = replica.log.get.dir.getParentFile
-    CoreUtils.swallow(Utils.delete(logDir), this)
-    logDir.createNewFile()
-    assertTrue(logDir.isFile)
-
-    if (failureType == Roll) {
-      try {
-        leaderServer.replicaManager.getLog(partition).get.roll()
-        fail("Log rolling should fail with KafkaStorageException")
-      } catch {
-        case e: KafkaStorageException => // This is expected
-      }
-    } else if (failureType == Checkpoint) {
-      leaderServer.replicaManager.checkpointHighWatermarks()
-    }
-
-    // Wait for ReplicaHighWatermarkCheckpoint to happen so that the log directory of the topic will be offline
-    TestUtils.waitUntilTrue(() => !leaderServer.logManager.isLogDirOnline(logDir.getAbsolutePath), "Expected log directory offline", 3000L)
-    assertTrue(leaderServer.replicaManager.localReplica(partition).isEmpty)
+    causeLogDirFailure(failureType, leaderServer, partition)
 
-    // The second send() should fail due to either KafkaStorageException or NotLeaderForPartitionException
+    // send() should fail due to either KafkaStorageException or NotLeaderForPartitionException
     try {
       producer.send(record).get(6000, TimeUnit.MILLISECONDS)
       fail("send() should fail with either KafkaStorageException or NotLeaderForPartitionException")
@@ -175,6 +162,25 @@ class LogDirFailureTest extends IntegrationTestHarness {
           case t: Throwable => fail(s"send() should fail with either KafkaStorageException or NotLeaderForPartitionException instead of ${t.toString}")
         }
     }
+  }
+
+  def testProduceAfterLogDirFailureOnLeader(failureType: LogDirFailureType) {
+    val consumer = createConsumer()
+    subscribeAndWaitForAssignment(topic, consumer)
+
+    val producer = createProducer()
+
+    val partition = new TopicPartition(topic, 0)
+    val record = new ProducerRecord(topic, 0, s"key".getBytes, s"value".getBytes)
+
+    val leaderServerId = producer.partitionsFor(topic).asScala.find(_.partition() == 0).get.leader().id()
+    val leaderServer = servers.find(_.config.brokerId == leaderServerId).get
+
+    // The first send() should succeed
+    producer.send(record).get()
+    TestUtils.consumeRecords(consumer, 1)
+
+    causeLogDirFailure(failureType, leaderServer, partition)
 
     TestUtils.waitUntilTrue(() => {
       // ProduceResponse may contain KafkaStorageException and trigger metadata update
@@ -193,11 +199,37 @@ class LogDirFailureTest extends IntegrationTestHarness {
 
     // The controller should have marked the replica on the original leader as offline
     val controllerServer = servers.find(_.kafkaController.isActive).get
-    val offlineReplicas = controllerServer.kafkaController.replicaStateMachine.replicasInState(topic, OfflineReplica)
+    val offlineReplicas = controllerServer.kafkaController.controllerContext.replicasInState(topic, OfflineReplica)
     assertTrue(offlineReplicas.contains(PartitionAndReplica(new TopicPartition(topic, 0), leaderServerId)))
   }
 
-  private def subscribeAndWaitForAssignment(topic: String, consumer: KafkaConsumer[Array[Byte], Array[Byte]]) {
+  private def causeLogDirFailure(failureType: LogDirFailureType,
+                                 leaderServer: KafkaServer,
+                                 partition: TopicPartition): Unit = {
+    // Make log directory of the partition on the leader broker inaccessible by replacing it with a file
+    val localLog = leaderServer.replicaManager.localLogOrException(partition)
+    val logDir = localLog.dir.getParentFile
+    CoreUtils.swallow(Utils.delete(logDir), this)
+    logDir.createNewFile()
+    assertTrue(logDir.isFile)
+
+    if (failureType == Roll) {
+      try {
+        leaderServer.replicaManager.getLog(partition).get.roll()
+        fail("Log rolling should fail with KafkaStorageException")
+      } catch {
+        case e: KafkaStorageException => // This is expected
+      }
+    } else if (failureType == Checkpoint) {
+      leaderServer.replicaManager.checkpointHighWatermarks()
+    }
+
+    // Wait for ReplicaHighWatermarkCheckpoint to happen so that the log directory of the topic will be offline
+    TestUtils.waitUntilTrue(() => !leaderServer.logManager.isLogDirOnline(logDir.getAbsolutePath), "Expected log directory offline", 3000L)
+    assertTrue(leaderServer.replicaManager.localLog(partition).isEmpty)
+  }
+
+  private def subscribeAndWaitForAssignment(topic: String, consumer: KafkaConsumer[Array[Byte], Array[Byte]]): Unit = {
     consumer.subscribe(Collections.singletonList(topic))
     TestUtils.pollUntilTrue(consumer, () => !consumer.assignment.isEmpty, "Expected non-empty assignment")
   }
diff --git a/core/src/test/scala/unit/kafka/server/LogOffsetTest.scala b/core/src/test/scala/unit/kafka/server/LogOffsetTest.scala
index 04b34675d86ed..abbb174e4665f 100755
--- a/core/src/test/scala/unit/kafka/server/LogOffsetTest.scala
+++ b/core/src/test/scala/unit/kafka/server/LogOffsetTest.scala
@@ -38,11 +38,11 @@ class LogOffsetTest extends BaseRequestTest {
 
   private lazy val time = new MockTime
 
-  protected override def numBrokers = 1
+  override def brokerCount = 1
 
   protected override def brokerTime(brokerId: Int) = time
 
-  protected override def propertyOverrides(props: Properties): Unit = {
+  protected override def brokerPropertyOverrides(props: Properties): Unit = {
     props.put("log.flush.interval.messages", "1")
     props.put("num.partitions", "20")
     props.put("log.retention.hours", "10")
@@ -52,7 +52,7 @@ class LogOffsetTest extends BaseRequestTest {
 
   @deprecated("ListOffsetsRequest V0", since = "")
   @Test
-  def testGetOffsetsForUnknownTopic() {
+  def testGetOffsetsForUnknownTopic(): Unit = {
     val topicPartition = new TopicPartition("foo", 0)
     val request = ListOffsetRequest.Builder.forConsumer(false, IsolationLevel.READ_UNCOMMITTED)
       .setTargetTimes(Map(topicPartition ->
@@ -63,7 +63,7 @@ class LogOffsetTest extends BaseRequestTest {
 
   @deprecated("ListOffsetsRequest V0", since = "")
   @Test
-  def testGetOffsetsAfterDeleteRecords() {
+  def testGetOffsetsAfterDeleteRecords(): Unit = {
     val topic = "kafka-"
     val topicPartition = new TopicPartition(topic, 0)
 
@@ -78,7 +78,7 @@ class LogOffsetTest extends BaseRequestTest {
       log.appendAsLeader(TestUtils.singletonRecords(value = Integer.toString(42).getBytes()), leaderEpoch = 0)
     log.flush()
 
-    log.onHighWatermarkIncremented(log.logEndOffset)
+    log.updateHighWatermark(log.logEndOffset)
     log.maybeIncrementLogStartOffset(3)
     log.deleteOldSegments()
 
@@ -95,7 +95,7 @@ class LogOffsetTest extends BaseRequestTest {
   }
 
   @Test
-  def testGetOffsetsBeforeLatestTime() {
+  def testGetOffsetsBeforeLatestTime(): Unit = {
     val topic = "kafka-"
     val topicPartition = new TopicPartition(topic, 0)
 
@@ -130,7 +130,7 @@ class LogOffsetTest extends BaseRequestTest {
   }
 
   @Test
-  def testEmptyLogsGetOffsets() {
+  def testEmptyLogsGetOffsets(): Unit = {
     val random = new Random
     val topic = "kafka-"
     val topicPartition = new TopicPartition(topic, random.nextInt(10))
@@ -155,7 +155,7 @@ class LogOffsetTest extends BaseRequestTest {
 
   @deprecated("legacyFetchOffsetsBefore", since = "")
   @Test
-  def testGetOffsetsBeforeNow() {
+  def testGetOffsetsBeforeNow(): Unit = {
     val random = new Random
     val topic = "kafka-"
     val topicPartition = new TopicPartition(topic, random.nextInt(3))
@@ -163,7 +163,7 @@ class LogOffsetTest extends BaseRequestTest {
     createTopic(topic, 3, 1)
 
     val logManager = server.getLogManager
-    val log = logManager.getOrCreateLog(topicPartition, logManager.initialDefaultConfig)
+    val log = logManager.getOrCreateLog(topicPartition, () => logManager.initialDefaultConfig)
 
     for (_ <- 0 until 20)
       log.appendAsLeader(TestUtils.singletonRecords(value = Integer.toString(42).getBytes()), leaderEpoch = 0)
@@ -185,7 +185,7 @@ class LogOffsetTest extends BaseRequestTest {
 
   @deprecated("legacyFetchOffsetsBefore", since = "")
   @Test
-  def testGetOffsetsBeforeEarliestTime() {
+  def testGetOffsetsBeforeEarliestTime(): Unit = {
     val random = new Random
     val topic = "kafka-"
     val topicPartition = new TopicPartition(topic, random.nextInt(3))
@@ -193,7 +193,7 @@ class LogOffsetTest extends BaseRequestTest {
     createTopic(topic, 3, 1)
 
     val logManager = server.getLogManager
-    val log = logManager.getOrCreateLog(topicPartition, logManager.initialDefaultConfig)
+    val log = logManager.getOrCreateLog(topicPartition, () => logManager.initialDefaultConfig)
     for (_ <- 0 until 20)
       log.appendAsLeader(TestUtils.singletonRecords(value = Integer.toString(42).getBytes()), leaderEpoch = 0)
     log.flush()
@@ -214,7 +214,7 @@ class LogOffsetTest extends BaseRequestTest {
   /* We test that `fetchOffsetsBefore` works correctly if `LogSegment.size` changes after each invocation (simulating
    * a race condition) */
   @Test
-  def testFetchOffsetsBeforeWithChangingSegmentSize() {
+  def testFetchOffsetsBeforeWithChangingSegmentSize(): Unit = {
     val log: Log = EasyMock.niceMock(classOf[Log])
     val logSegment: LogSegment = EasyMock.niceMock(classOf[LogSegment])
     EasyMock.expect(logSegment.size).andStubAnswer(new IAnswer[Int] {
@@ -231,7 +231,7 @@ class LogOffsetTest extends BaseRequestTest {
   /* We test that `fetchOffsetsBefore` works correctly if `Log.logSegments` content and size are
    * different (simulating a race condition) */
   @Test
-  def testFetchOffsetsBeforeWithChangingSegments() {
+  def testFetchOffsetsBeforeWithChangingSegments(): Unit = {
     val log: Log = EasyMock.niceMock(classOf[Log])
     val logSegment: LogSegment = EasyMock.niceMock(classOf[LogSegment])
     EasyMock.expect(log.logSegments).andStubAnswer {
diff --git a/core/src/test/scala/unit/kafka/server/LogRecoveryTest.scala b/core/src/test/scala/unit/kafka/server/LogRecoveryTest.scala
index 15f9a9b1d57ef..cd76bca0831d9 100755
--- a/core/src/test/scala/unit/kafka/server/LogRecoveryTest.scala
+++ b/core/src/test/scala/unit/kafka/server/LogRecoveryTest.scala
@@ -18,6 +18,8 @@ package kafka.server
 
 import java.util.Properties
 
+import scala.collection.Seq
+
 import kafka.utils.TestUtils
 import TestUtils._
 import kafka.zk.ZooKeeperTestHarness
@@ -73,7 +75,7 @@ class LogRecoveryTest extends ZooKeeperTestHarness {
   }
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     super.setUp()
 
     configs = TestUtils.createBrokerConfigs(2, zkConnect, enableControlledShutdown = false).map(KafkaConfig.fromProps(_, overridingProps))
@@ -91,7 +93,7 @@ class LogRecoveryTest extends ZooKeeperTestHarness {
   }
 
   @After
-  override def tearDown() {
+  override def tearDown(): Unit = {
     producer.close()
     TestUtils.shutdownServers(servers)
     super.tearDown()
@@ -104,7 +106,7 @@ class LogRecoveryTest extends ZooKeeperTestHarness {
 
     // give some time for the follower 1 to record leader HW
     TestUtils.waitUntilTrue(() =>
-      server2.replicaManager.localReplica(topicPartition).get.highWatermark.messageOffset == numMessages,
+      server2.replicaManager.localLogOrException(topicPartition).highWatermark == numMessages,
       "Failed to update high watermark for follower after timeout")
 
     servers.foreach(_.replicaManager.checkpointHighWatermarks())
@@ -147,7 +149,7 @@ class LogRecoveryTest extends ZooKeeperTestHarness {
       * is that server1 has caught up on the topicPartition, and has joined the ISR.
       * In the line below, we wait until the condition is met before shutting down server2
       */
-    waitUntilTrue(() => server2.replicaManager.getPartition(topicPartition).get.inSyncReplicas.size == 2,
+    waitUntilTrue(() => server2.replicaManager.nonOfflinePartition(topicPartition).get.inSyncReplicaIds.size == 2,
       "Server 1 is not able to join the ISR after restart")
 
 
@@ -166,7 +168,7 @@ class LogRecoveryTest extends ZooKeeperTestHarness {
 
     // give some time for follower 1 to record leader HW of 60
     TestUtils.waitUntilTrue(() =>
-      server2.replicaManager.localReplica(topicPartition).get.highWatermark.messageOffset == hw,
+      server2.replicaManager.localLogOrException(topicPartition).highWatermark == hw,
       "Failed to update high watermark for follower after timeout")
     // shutdown the servers to allow the hw to be checkpointed
     servers.foreach(_.shutdown())
@@ -180,7 +182,7 @@ class LogRecoveryTest extends ZooKeeperTestHarness {
     val hw = 20L
     // give some time for follower 1 to record leader HW of 600
     TestUtils.waitUntilTrue(() =>
-      server2.replicaManager.localReplica(topicPartition).get.highWatermark.messageOffset == hw,
+      server2.replicaManager.localLogOrException(topicPartition).highWatermark == hw,
       "Failed to update high watermark for follower after timeout")
     // shutdown the servers to allow the hw to be checkpointed
     servers.foreach(_.shutdown())
@@ -199,7 +201,7 @@ class LogRecoveryTest extends ZooKeeperTestHarness {
 
     // allow some time for the follower to get the leader HW
     TestUtils.waitUntilTrue(() =>
-      server2.replicaManager.localReplica(topicPartition).get.highWatermark.messageOffset == hw,
+      server2.replicaManager.localLogOrException(topicPartition).highWatermark == hw,
       "Failed to update high watermark for follower after timeout")
     // kill the server hosting the preferred replica
     server1.shutdown()
@@ -226,11 +228,11 @@ class LogRecoveryTest extends ZooKeeperTestHarness {
     hw += 2
 
     // allow some time for the follower to create replica
-    TestUtils.waitUntilTrue(() => server1.replicaManager.localReplica(topicPartition).nonEmpty,
+    TestUtils.waitUntilTrue(() => server1.replicaManager.localLog(topicPartition).nonEmpty,
       "Failed to create replica in follower after timeout")
     // allow some time for the follower to get the leader HW
     TestUtils.waitUntilTrue(() =>
-      server1.replicaManager.localReplica(topicPartition).get.highWatermark.messageOffset == hw,
+      server1.replicaManager.localLogOrException(topicPartition).highWatermark == hw,
       "Failed to update high watermark for follower after timeout")
     // shutdown the servers to allow the hw to be checkpointed
     servers.foreach(_.shutdown())
@@ -238,7 +240,7 @@ class LogRecoveryTest extends ZooKeeperTestHarness {
     assertEquals(hw, hwFile2.read.getOrElse(topicPartition, 0L))
   }
 
-  private def sendMessages(n: Int) {
+  private def sendMessages(n: Int): Unit = {
     (0 until n).map(_ => producer.send(new ProducerRecord(topic, 0, message))).foreach(_.get)
   }
 }
diff --git a/core/src/test/scala/unit/kafka/server/MaintenanceBrokerTest.scala b/core/src/test/scala/unit/kafka/server/MaintenanceBrokerTest.scala
new file mode 100644
index 0000000000000..c97740136623b
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/server/MaintenanceBrokerTest.scala
@@ -0,0 +1,215 @@
+/**
+  * 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 java.util.{Optional, Properties}
+
+import kafka.server.KafkaConfig.fromProps
+import kafka.utils.CoreUtils._
+import kafka.utils.TestUtils
+import kafka.utils.TestUtils._
+import kafka.zk.ZooKeeperTestHarness
+import org.apache.kafka.clients.admin._
+import org.apache.kafka.common.network.ListenerName
+import org.apache.kafka.common.security.auth.SecurityProtocol
+
+import scala.collection.JavaConverters._
+import org.junit.Assert._
+import org.junit.{After, Test}
+
+import scala.collection.Map
+
+/**
+  * This is the main test which ensure maintenance broker work correctly.
+  */
+class MaintenanceBrokerTest extends ZooKeeperTestHarness {
+
+  var brokers: Seq[KafkaServer] = null
+
+  @After
+  override def tearDown() {
+    shutdownServers(brokers)
+    super.tearDown()
+  }
+
+  @Test
+  def testTopicCreatedByZkclientShouldHonorMaintenanceBrokers(): Unit = {
+
+    brokers = (0 to 2).map { id => createServer(fromProps(createBrokerConfig(id, zkConnect))) }
+
+    // create topic using zkclient
+    TestUtils.createTopic(zkClient, "topic1", 3, 2, brokers)
+    (0 to 2).foreach {
+      brokerId =>
+        assertTrue("topic1 should be in broker " + brokerId, !ensureTopicNotInBrokers("topic1", Set(brokerId)))
+    }
+
+    // setting broker 1 to not take new topic partitions
+    setMaintenanceBrokers(Seq(1))
+
+    TestUtils.createTopic(zkClient, "topic2", 3, 2, brokers)
+    assertTrue("topic2 should not be in broker 1", ensureTopicNotInBrokers("topic2", Set(1)))
+
+    // setting broker 1 and 2 to not take new topic partitions
+    setMaintenanceBrokers(Seq(1, 2))
+
+    TestUtils.createTopic(zkClient, "topic3", 3, 1, brokers)
+
+    assertTrue("topic3 should not be in broker 1 and 2", ensureTopicNotInBrokers("topic3", Set(1, 2)))
+    assertTrue("topic3 should in broker 0", !ensureTopicNotInBrokers("topic3", Set(0)))
+  }
+
+  @Test
+  def testTopicCreatedByAdminClientShouldHonorMaintenanceBrokers(): Unit = {
+
+    brokers = (0 to 2).map { id => createServer(fromProps(createBrokerConfig(id, zkConnect))) }
+
+    val brokerList = TestUtils.bootstrapServers(brokers, ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT))
+    val adminClientConfig = new Properties
+    adminClientConfig.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, brokerList)
+    val client = AdminClient.create(adminClientConfig)
+
+    // create topic using admin client
+    val future1 = client.createTopics(Seq("topic1").map(new NewTopic(_, 3, 2.toShort)).asJava,
+      new CreateTopicsOptions()).all()
+    future1.get()
+
+    (0 to 2).foreach {
+      brokerId =>
+        assertTrue("topic1 should be in broker " + brokerId, !ensureTopicNotInBrokers("topic1", Set(brokerId)))
+    }
+
+    TestUtils.waitUntilControllerElected(zkClient)
+
+    // setting broker 1 to not take new topic partitions
+    setMaintenanceBrokers(Seq(1))
+
+    val future2 = client.createTopics(Seq("topic2").map(new NewTopic(_, 3, 2.toShort)).asJava,
+      new CreateTopicsOptions()).all()
+    future2.get()
+
+    assertTrue("topic2 should not be in broker 1", ensureTopicNotInBrokers("topic2", Set(1)))
+
+    // setting broker 1 and 2 to not take new topic partitions
+    setMaintenanceBrokers(Seq(1, 2))
+
+    val future3 = client.createTopics(Seq("topic3").map(new NewTopic(_, 3, 1.toShort)).asJava,
+      new CreateTopicsOptions()).all()
+    future3.get()
+
+    assertTrue("topic3 should not be in broker 1 and 2", ensureTopicNotInBrokers("topic3", Set(1, 2)))
+    assertTrue("topic3 should be in broker 0", !ensureTopicNotInBrokers("topic3", Set(0)))
+
+    // create topic with #replicas > #non-maintenance brokers
+    val future4 = client.createTopics(Seq("topic4").map(new NewTopic(_, 3, 3.toShort)).asJava,
+      new CreateTopicsOptions()).all()
+    future4.get()
+
+    (0 to 2).foreach {
+      brokerId =>
+        assertTrue("topic4 should be in broker " + brokerId + " because #replicas > #non-maintenance brokers",
+          !ensureTopicNotInBrokers("topic4", Set(brokerId)))
+    }
+
+    // clear maintenance broker
+    setMaintenanceBrokers(Seq.empty[Int])
+
+    // create topic using admin client
+    val future5 = client.createTopics(Seq("topic5").map(new NewTopic(_, 3, 1.toShort)).asJava,
+      new CreateTopicsOptions()).all()
+    future5.get()
+
+    (0 to 2).foreach {
+      brokerId =>
+        assertTrue("topic5 should be in broker " + brokerId, !ensureTopicNotInBrokers("topic5", Set(brokerId)))
+    }
+
+    client.close()
+  }
+
+  @Test
+  def testAddPartitionByAdminClientShouldHonorMaintenanceBrokers(): Unit = {
+
+    brokers = (0 to 2).map { id => createServer(fromProps(createBrokerConfig(id, zkConnect))) }
+
+    val brokerList = TestUtils.bootstrapServers(brokers, ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT))
+    val adminClientConfig = new Properties
+    adminClientConfig.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, brokerList)
+    val client = AdminClient.create(adminClientConfig)
+
+    TestUtils.waitUntilControllerElected(zkClient)
+
+    // setting broker 1 to not take new topic partitions
+    setMaintenanceBrokers(Seq(1))
+
+    val future1 = client.createTopics(Seq("topic1").map(new NewTopic(_, 3, 2.toShort)).asJava,
+      new CreateTopicsOptions()).all()
+    future1.get()
+
+    assertTrue("topic1 should not be in broker 1", ensureTopicNotInBrokers("topic1", Set(1)))
+
+    val future2 = client.createPartitions(Map("topic1" -> NewPartitions.increaseTo(5)).asJava).all()
+    future2.get()
+
+    assertTrue("topic1 should not be in broker 1 after increasing partition count",
+      ensureTopicNotInBrokers("topic1", Set(1)))
+
+    client.close()
+  }
+
+  @Test
+  def testTopicCreatedInZkShouldBeRearrangedForMaintenanceBrokers(): Unit = {
+
+    brokers = (0 to 2).map { id => createServer(fromProps(createBrokerConfig(id, zkConnect))) }
+    TestUtils.waitUntilControllerElected(zkClient)
+
+    TestUtils.createTopic(zkClient, "topic1", Map(0 -> List(0, 1), 1 -> List(1, 2)), brokers)
+
+    // setting broker 1 to not take new topic partitions
+    setMaintenanceBrokers(Seq(1))
+
+    TestUtils.createTopic(zkClient, "topic2", Map(0 -> List(0, 1), 1 -> List(1, 2)), brokers)
+
+    assertTrue("new topic topic2 should be rearranged and not be in broker 1", ensureTopicNotInBrokers("topic2", Set(1)))
+    assertTrue("old topic topic1 should still in broker 1", !ensureTopicNotInBrokers("topic1", Set(1)))
+  }
+
+  def ensureTopicNotInBrokers(topic: String, brokerIds: Set[Int]): Boolean = {
+    val topicAssignment = zkClient.getReplicaAssignmentForTopics(Set(topic))
+    topicAssignment.map(_._2).flatten.toSet.intersect(brokerIds).isEmpty
+  }
+
+  def createBrokers(brokerIds: Seq[Int]): Unit = {
+    brokerIds.foreach { id =>
+      brokers = brokers :+ createServer(fromProps(createBrokerConfig(id, zkConnect)))
+    }
+  }
+
+  def setMaintenanceBrokers(brokerIds: Seq[Int]): Unit = {
+    var propstring = brokerIds.mkString(",")
+    adminZkClient.changeBrokerConfig(None,
+      propsWith((DynamicConfig.Broker.MaintenanceBrokerListProp, propstring)))
+
+    val controllerId = TestUtils.waitUntilControllerElected(zkClient)
+
+    TestUtils.waitUntilTrue(() => brokers(controllerId).config.getMaintenanceBrokerList == brokerIds,
+      s"wait until broker $propstring is masked as maintenance broker not taking new partitions", 5000)
+
+  }
+
+}
diff --git a/core/src/test/scala/unit/kafka/server/MetadataCacheTest.scala b/core/src/test/scala/unit/kafka/server/MetadataCacheTest.scala
index 5ddabc0b6f7a0..0491dabb757ea 100644
--- a/core/src/test/scala/unit/kafka/server/MetadataCacheTest.scala
+++ b/core/src/test/scala/unit/kafka/server/MetadataCacheTest.scala
@@ -20,14 +20,14 @@ import java.util
 import java.util.Optional
 import util.Arrays.asList
 
-import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.message.UpdateMetadataRequestData.{UpdateMetadataBroker, UpdateMetadataEndpoint, UpdateMetadataPartitionState}
 import org.apache.kafka.common.network.ListenerName
 import org.apache.kafka.common.protocol.{ApiKeys, Errors}
 import org.apache.kafka.common.requests.UpdateMetadataRequest
-import org.apache.kafka.common.requests.UpdateMetadataRequest.{Broker, EndPoint}
 import org.apache.kafka.common.security.auth.SecurityProtocol
 import org.junit.Test
 import org.junit.Assert._
+import org.scalatest.Assertions
 
 import scala.collection.JavaConverters._
 
@@ -35,7 +35,7 @@ class MetadataCacheTest {
   val brokerEpoch = 0L
 
   @Test
-  def getTopicMetadataNonExistingTopics() {
+  def getTopicMetadataNonExistingTopics(): Unit = {
     val topic = "topic"
     val cache = new MetadataCache(1)
     val topicMetadata = cache.getTopicMetadata(Set(topic), ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT))
@@ -43,7 +43,7 @@ class MetadataCacheTest {
   }
 
   @Test
-  def getTopicMetadata() {
+  def getTopicMetadata(): Unit = {
     val topic0 = "topic-0"
     val topic1 = "topic-1"
 
@@ -53,25 +53,60 @@ class MetadataCacheTest {
     val controllerId = 2
     val controllerEpoch = 1
 
-    def endPoints(brokerId: Int): Seq[EndPoint] = {
+    def endpoints(brokerId: Int): Seq[UpdateMetadataEndpoint] = {
       val host = s"foo-$brokerId"
       Seq(
-        new EndPoint(host, 9092, SecurityProtocol.PLAINTEXT, ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT)),
-        new EndPoint(host, 9093, SecurityProtocol.SSL, ListenerName.forSecurityProtocol(SecurityProtocol.SSL))
+        new UpdateMetadataEndpoint()
+          .setHost(host)
+          .setPort(9092)
+          .setSecurityProtocol(SecurityProtocol.PLAINTEXT.id)
+          .setListener(ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT).value),
+        new UpdateMetadataEndpoint()
+          .setHost(host)
+          .setPort(9093)
+          .setSecurityProtocol(SecurityProtocol.SSL.id)
+          .setListener(ListenerName.forSecurityProtocol(SecurityProtocol.SSL).value)
       )
     }
 
     val brokers = (0 to 4).map { brokerId =>
-      new Broker(brokerId, endPoints(brokerId).asJava, "rack1")
-    }.toSet
+      new UpdateMetadataBroker()
+        .setId(brokerId)
+        .setEndpoints(endpoints(brokerId).asJava)
+        .setRack("rack1")
+    }
 
-    val partitionStates = Map(
-      new TopicPartition(topic0, 0) -> new UpdateMetadataRequest.PartitionState(controllerEpoch, 0, 0, asList(0, 1, 3), zkVersion, asList(0, 1, 3), asList()),
-      new TopicPartition(topic0, 1) -> new UpdateMetadataRequest.PartitionState(controllerEpoch, 1, 1, asList(1, 0), zkVersion, asList(1, 2, 0, 4), asList()),
-      new TopicPartition(topic1, 0) -> new UpdateMetadataRequest.PartitionState(controllerEpoch, 2, 2, asList(2, 1), zkVersion, asList(2, 1, 3), asList()))
+    val partitionStates = Seq(
+      new UpdateMetadataPartitionState()
+        .setTopicName(topic0)
+        .setPartitionIndex(0)
+        .setControllerEpoch(controllerEpoch)
+        .setLeader(0)
+        .setLeaderEpoch(0)
+        .setIsr(asList(0, 1, 3))
+        .setZkVersion(zkVersion)
+        .setReplicas(asList(0, 1, 3)),
+      new UpdateMetadataPartitionState()
+        .setTopicName(topic0)
+        .setPartitionIndex(1)
+        .setControllerEpoch(controllerEpoch)
+        .setLeader(1)
+        .setLeaderEpoch(1)
+        .setIsr(asList(1, 0))
+        .setZkVersion(zkVersion)
+        .setReplicas(asList(1, 2, 0, 4)),
+      new UpdateMetadataPartitionState()
+        .setTopicName(topic1)
+        .setPartitionIndex(0)
+        .setControllerEpoch(controllerEpoch)
+        .setLeader(2)
+        .setLeaderEpoch(2)
+        .setIsr(asList(2, 1))
+        .setZkVersion(zkVersion)
+        .setReplicas(asList(2, 1, 3)))
 
     val version = ApiKeys.UPDATE_METADATA.latestVersion
-    val updateMetadataRequest = new UpdateMetadataRequest.Builder(version, controllerId, controllerEpoch, brokerEpoch,
+    val updateMetadataRequest = new UpdateMetadataRequest.Builder(version, controllerId, controllerEpoch, brokerEpoch, brokerEpoch,
       partitionStates.asJava, brokers.asJava).build()
     cache.updateMetadata(15, updateMetadataRequest)
 
@@ -86,7 +121,7 @@ class MetadataCacheTest {
         assertEquals(Errors.NONE, topicMetadata.error)
         assertEquals(topic, topicMetadata.topic)
 
-        val topicPartitionStates = partitionStates.filter { case (tp, _) => tp.topic ==  topic }
+        val topicPartitionStates = partitionStates.filter { ps => ps.topicName == topic }
         val partitionMetadatas = topicMetadata.partitionMetadata.asScala.sortBy(_.partition)
         assertEquals(s"Unexpected partition count for topic $topic", topicPartitionStates.size, partitionMetadatas.size)
 
@@ -94,14 +129,15 @@ class MetadataCacheTest {
           assertEquals(Errors.NONE, partitionMetadata.error)
           assertEquals(partitionId, partitionMetadata.partition)
           val leader = partitionMetadata.leader
-          val partitionState = topicPartitionStates(new TopicPartition(topic, partitionId))
-          assertEquals(partitionState.basePartitionState.leader, leader.id)
-          assertEquals(Optional.of(partitionState.basePartitionState.leaderEpoch), partitionMetadata.leaderEpoch)
-          assertEquals(partitionState.basePartitionState.isr, partitionMetadata.isr.asScala.map(_.id).asJava)
-          assertEquals(partitionState.basePartitionState.replicas, partitionMetadata.replicas.asScala.map(_.id).asJava)
-          val endPoint = endPoints(partitionMetadata.leader.id).find(_.listenerName == listenerName).get
-          assertEquals(endPoint.host, leader.host)
-          assertEquals(endPoint.port, leader.port)
+          val partitionState = topicPartitionStates.find(_.partitionIndex == partitionId).getOrElse(
+            Assertions.fail(s"Unable to find partition state for partition $partitionId"))
+          assertEquals(partitionState.leader, leader.id)
+          assertEquals(Optional.of(partitionState.leaderEpoch), partitionMetadata.leaderEpoch)
+          assertEquals(partitionState.isr, partitionMetadata.isr.asScala.map(_.id).asJava)
+          assertEquals(partitionState.replicas, partitionMetadata.replicas.asScala.map(_.id).asJava)
+          val endpoint = endpoints(partitionMetadata.leader.id).find(_.listener == listenerName.value).get
+          assertEquals(endpoint.host, leader.host)
+          assertEquals(endpoint.port, leader.port)
         }
       }
 
@@ -115,7 +151,13 @@ class MetadataCacheTest {
   def getTopicMetadataPartitionLeaderNotAvailable(): Unit = {
     val securityProtocol = SecurityProtocol.PLAINTEXT
     val listenerName = ListenerName.forSecurityProtocol(securityProtocol)
-    val brokers = Set(new Broker(0, Seq(new EndPoint("foo", 9092, securityProtocol, listenerName)).asJava, null))
+    val brokers = Seq(new UpdateMetadataBroker()
+      .setId(0)
+      .setEndpoints(Seq(new UpdateMetadataEndpoint()
+        .setHost("foo")
+        .setPort(9092)
+        .setSecurityProtocol(securityProtocol.id)
+        .setListener(listenerName.value)).asJava))
     verifyTopicMetadataPartitionLeaderOrEndpointNotAvailable(brokers, listenerName,
       leader = 1, Errors.LEADER_NOT_AVAILABLE, errorUnavailableListeners = false)
     verifyTopicMetadataPartitionLeaderOrEndpointNotAvailable(brokers, listenerName,
@@ -127,10 +169,28 @@ class MetadataCacheTest {
     val plaintextListenerName = ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT)
     val sslListenerName = ListenerName.forSecurityProtocol(SecurityProtocol.SSL)
     val broker0Endpoints = Seq(
-      new EndPoint("host0", 9092, SecurityProtocol.PLAINTEXT, plaintextListenerName),
-      new EndPoint("host0", 9093, SecurityProtocol.SSL, sslListenerName))
-    val broker1Endpoints = Seq(new EndPoint("host1", 9092, SecurityProtocol.PLAINTEXT, plaintextListenerName))
-    val brokers = Set(new Broker(0, broker0Endpoints.asJava, null), new Broker(1, broker1Endpoints.asJava, null))
+      new UpdateMetadataEndpoint()
+        .setHost("host0")
+        .setPort(9092)
+        .setSecurityProtocol(SecurityProtocol.PLAINTEXT.id)
+        .setListener(plaintextListenerName.value),
+      new UpdateMetadataEndpoint()
+        .setHost("host0")
+        .setPort(9093)
+        .setSecurityProtocol(SecurityProtocol.SSL.id)
+        .setListener(sslListenerName.value))
+    val broker1Endpoints = Seq(new UpdateMetadataEndpoint()
+      .setHost("host1")
+      .setPort(9092)
+      .setSecurityProtocol(SecurityProtocol.PLAINTEXT.id)
+      .setListener(plaintextListenerName.value))
+    val brokers = Seq(
+      new UpdateMetadataBroker()
+        .setId(0)
+        .setEndpoints(broker0Endpoints.asJava),
+      new UpdateMetadataBroker()
+        .setId(1)
+        .setEndpoints(broker1Endpoints.asJava))
     verifyTopicMetadataPartitionLeaderOrEndpointNotAvailable(brokers, sslListenerName,
       leader = 1, Errors.LISTENER_NOT_FOUND, errorUnavailableListeners = true)
   }
@@ -140,15 +200,33 @@ class MetadataCacheTest {
     val plaintextListenerName = ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT)
     val sslListenerName = ListenerName.forSecurityProtocol(SecurityProtocol.SSL)
     val broker0Endpoints = Seq(
-      new EndPoint("host0", 9092, SecurityProtocol.PLAINTEXT, plaintextListenerName),
-      new EndPoint("host0", 9093, SecurityProtocol.SSL, sslListenerName))
-    val broker1Endpoints = Seq(new EndPoint("host1", 9092, SecurityProtocol.PLAINTEXT, plaintextListenerName))
-    val brokers = Set(new Broker(0, broker0Endpoints.asJava, null), new Broker(1, broker1Endpoints.asJava, null))
+      new UpdateMetadataEndpoint()
+        .setHost("host0")
+        .setPort(9092)
+        .setSecurityProtocol(SecurityProtocol.PLAINTEXT.id)
+        .setListener(plaintextListenerName.value),
+      new UpdateMetadataEndpoint()
+        .setHost("host0")
+        .setPort(9093)
+        .setSecurityProtocol(SecurityProtocol.SSL.id)
+        .setListener(sslListenerName.value))
+    val broker1Endpoints = Seq(new UpdateMetadataEndpoint()
+      .setHost("host1")
+      .setPort(9092)
+      .setSecurityProtocol(SecurityProtocol.PLAINTEXT.id)
+      .setListener(plaintextListenerName.value))
+    val brokers = Seq(
+      new UpdateMetadataBroker()
+        .setId(0)
+        .setEndpoints(broker0Endpoints.asJava),
+      new UpdateMetadataBroker()
+        .setId(1)
+        .setEndpoints(broker1Endpoints.asJava))
     verifyTopicMetadataPartitionLeaderOrEndpointNotAvailable(brokers, sslListenerName,
       leader = 1, Errors.LEADER_NOT_AVAILABLE, errorUnavailableListeners = false)
   }
 
-  private def verifyTopicMetadataPartitionLeaderOrEndpointNotAvailable(brokers: Set[Broker],
+  private def verifyTopicMetadataPartitionLeaderOrEndpointNotAvailable(brokers: Seq[UpdateMetadataBroker],
                                                                        listenerName: ListenerName,
                                                                        leader: Int,
                                                                        expectedError: Errors,
@@ -162,11 +240,18 @@ class MetadataCacheTest {
     val controllerEpoch = 1
 
     val leaderEpoch = 1
-    val partitionStates = Map(
-      new TopicPartition(topic, 0) -> new UpdateMetadataRequest.PartitionState(controllerEpoch, leader, leaderEpoch, asList(0), zkVersion, asList(0), asList()))
+    val partitionStates = Seq(new UpdateMetadataPartitionState()
+      .setTopicName(topic)
+      .setPartitionIndex(0)
+      .setControllerEpoch(controllerEpoch)
+      .setLeader(leader)
+      .setLeaderEpoch(leaderEpoch)
+      .setIsr(asList(0))
+      .setZkVersion(zkVersion)
+      .setReplicas(asList(0)))
 
     val version = ApiKeys.UPDATE_METADATA.latestVersion
-    val updateMetadataRequest = new UpdateMetadataRequest.Builder(version, controllerId, controllerEpoch, brokerEpoch,
+    val updateMetadataRequest = new UpdateMetadataRequest.Builder(version, controllerId, controllerEpoch, brokerEpoch, brokerEpoch,
       partitionStates.asJava, brokers.asJava).build()
     cache.updateMetadata(15, updateMetadataRequest)
 
@@ -182,13 +267,13 @@ class MetadataCacheTest {
     val partitionMetadata = partitionMetadatas.get(0)
     assertEquals(0, partitionMetadata.partition)
     assertEquals(expectedError, partitionMetadata.error)
-    assertTrue(partitionMetadata.isr.isEmpty)
+    assertFalse(partitionMetadata.isr.isEmpty)
     assertEquals(1, partitionMetadata.replicas.size)
     assertEquals(0, partitionMetadata.replicas.get(0).id)
   }
 
   @Test
-  def getTopicMetadataReplicaNotAvailable() {
+  def getTopicMetadataReplicaNotAvailable(): Unit = {
     val topic = "topic"
 
     val cache = new MetadataCache(1)
@@ -198,7 +283,13 @@ class MetadataCacheTest {
     val controllerEpoch = 1
     val securityProtocol = SecurityProtocol.PLAINTEXT
     val listenerName = ListenerName.forSecurityProtocol(securityProtocol)
-    val brokers = Set(new Broker(0, Seq(new EndPoint("foo", 9092, securityProtocol, listenerName)).asJava, null))
+    val brokers = Seq(new UpdateMetadataBroker()
+      .setId(0)
+      .setEndpoints(Seq(new UpdateMetadataEndpoint()
+        .setHost("foo")
+        .setPort(9092)
+        .setSecurityProtocol(securityProtocol.id)
+        .setListener(listenerName.value)).asJava))
 
     // replica 1 is not available
     val leader = 0
@@ -206,11 +297,19 @@ class MetadataCacheTest {
     val replicas = asList[Integer](0, 1)
     val isr = asList[Integer](0)
 
-    val partitionStates = Map(
-      new TopicPartition(topic, 0) -> new UpdateMetadataRequest.PartitionState(controllerEpoch, leader, leaderEpoch, isr, zkVersion, replicas, asList()))
+    val partitionStates = Seq(
+      new UpdateMetadataPartitionState()
+        .setTopicName(topic)
+        .setPartitionIndex(0)
+        .setControllerEpoch(controllerEpoch)
+        .setLeader(leader)
+        .setLeaderEpoch(leaderEpoch)
+        .setIsr(isr)
+        .setZkVersion(zkVersion)
+        .setReplicas(replicas))
 
     val version = ApiKeys.UPDATE_METADATA.latestVersion
-    val updateMetadataRequest = new UpdateMetadataRequest.Builder(version, controllerId, controllerEpoch, brokerEpoch,
+    val updateMetadataRequest = new UpdateMetadataRequest.Builder(version, controllerId, controllerEpoch, brokerEpoch, brokerEpoch,
       partitionStates.asJava, brokers.asJava).build()
     cache.updateMetadata(15, updateMetadataRequest)
 
@@ -248,7 +347,7 @@ class MetadataCacheTest {
   }
 
   @Test
-  def getTopicMetadataIsrNotAvailable() {
+  def getTopicMetadataIsrNotAvailable(): Unit = {
     val topic = "topic"
 
     val cache = new MetadataCache(1)
@@ -258,7 +357,14 @@ class MetadataCacheTest {
     val controllerEpoch = 1
     val securityProtocol = SecurityProtocol.PLAINTEXT
     val listenerName = ListenerName.forSecurityProtocol(securityProtocol)
-    val brokers = Set(new Broker(0, Seq(new EndPoint("foo", 9092, securityProtocol, listenerName)).asJava, "rack1"))
+    val brokers = Seq(new UpdateMetadataBroker()
+      .setId(0)
+      .setRack("rack1")
+      .setEndpoints(Seq(new UpdateMetadataEndpoint()
+        .setHost("foo")
+        .setPort(9092)
+        .setSecurityProtocol(securityProtocol.id)
+        .setListener(listenerName.value)).asJava))
 
     // replica 1 is not available
     val leader = 0
@@ -266,11 +372,18 @@ class MetadataCacheTest {
     val replicas = asList[Integer](0)
     val isr = asList[Integer](0, 1)
 
-    val partitionStates = Map(
-      new TopicPartition(topic, 0) -> new UpdateMetadataRequest.PartitionState(controllerEpoch, leader, leaderEpoch, isr, zkVersion, replicas, asList()))
+    val partitionStates = Seq(new UpdateMetadataPartitionState()
+      .setTopicName(topic)
+      .setPartitionIndex(0)
+      .setControllerEpoch(controllerEpoch)
+      .setLeader(leader)
+      .setLeaderEpoch(leaderEpoch)
+      .setIsr(isr)
+      .setZkVersion(zkVersion)
+      .setReplicas(replicas))
 
     val version = ApiKeys.UPDATE_METADATA.latestVersion
-    val updateMetadataRequest = new UpdateMetadataRequest.Builder(version, controllerId, controllerEpoch, brokerEpoch,
+    val updateMetadataRequest = new UpdateMetadataRequest.Builder(version, controllerId, controllerEpoch, brokerEpoch, brokerEpoch,
       partitionStates.asJava, brokers.asJava).build()
     cache.updateMetadata(15, updateMetadataRequest)
 
@@ -308,21 +421,34 @@ class MetadataCacheTest {
   }
 
   @Test
-  def getTopicMetadataWithNonSupportedSecurityProtocol() {
+  def getTopicMetadataWithNonSupportedSecurityProtocol(): Unit = {
     val topic = "topic"
     val cache = new MetadataCache(1)
     val securityProtocol = SecurityProtocol.PLAINTEXT
-    val brokers = Set(new Broker(0,
-      Seq(new EndPoint("foo", 9092, securityProtocol, ListenerName.forSecurityProtocol(securityProtocol))).asJava, ""))
+    val brokers = Seq(new UpdateMetadataBroker()
+      .setId(0)
+      .setRack("")
+      .setEndpoints(Seq(new UpdateMetadataEndpoint()
+        .setHost("foo")
+        .setPort(9092)
+        .setSecurityProtocol(securityProtocol.id)
+        .setListener(ListenerName.forSecurityProtocol(securityProtocol).value)).asJava))
     val controllerEpoch = 1
     val leader = 0
     val leaderEpoch = 0
     val replicas = asList[Integer](0)
     val isr = asList[Integer](0, 1)
-    val partitionStates = Map(
-      new TopicPartition(topic, 0) -> new UpdateMetadataRequest.PartitionState(controllerEpoch, leader, leaderEpoch, isr, 3, replicas, asList()))
+    val partitionStates = Seq(new UpdateMetadataPartitionState()
+      .setTopicName(topic)
+      .setPartitionIndex(0)
+      .setControllerEpoch(controllerEpoch)
+      .setLeader(leader)
+      .setLeaderEpoch(leaderEpoch)
+      .setIsr(isr)
+      .setZkVersion(3)
+      .setReplicas(replicas))
     val version = ApiKeys.UPDATE_METADATA.latestVersion
-    val updateMetadataRequest = new UpdateMetadataRequest.Builder(version, 2, controllerEpoch, brokerEpoch, partitionStates.asJava,
+    val updateMetadataRequest = new UpdateMetadataRequest.Builder(version, 2, controllerEpoch, brokerEpoch, brokerEpoch, partitionStates.asJava,
       brokers.asJava).build()
     cache.updateMetadata(15, updateMetadataRequest)
 
@@ -333,35 +459,48 @@ class MetadataCacheTest {
   }
 
   @Test
-  def getAliveBrokersShouldNotBeMutatedByUpdateCache() {
+  def getAliveBrokersShouldNotBeMutatedByUpdateCache(): Unit = {
     val topic = "topic"
     val cache = new MetadataCache(1)
 
-    def updateCache(brokerIds: Set[Int]) {
+    def updateCache(brokerIds: Seq[Int]): Unit = {
       val brokers = brokerIds.map { brokerId =>
         val securityProtocol = SecurityProtocol.PLAINTEXT
-        new Broker(brokerId, Seq(
-          new EndPoint("foo", 9092, securityProtocol, ListenerName.forSecurityProtocol(securityProtocol))).asJava, "")
+        new UpdateMetadataBroker()
+          .setId(brokerId)
+          .setRack("")
+          .setEndpoints(Seq(new UpdateMetadataEndpoint()
+            .setHost("foo")
+            .setPort(9092)
+            .setSecurityProtocol(securityProtocol.id)
+            .setListener(ListenerName.forSecurityProtocol(securityProtocol).value)).asJava)
       }
       val controllerEpoch = 1
       val leader = 0
       val leaderEpoch = 0
       val replicas = asList[Integer](0)
       val isr = asList[Integer](0, 1)
-      val partitionStates = Map(
-        new TopicPartition(topic, 0) -> new UpdateMetadataRequest.PartitionState(controllerEpoch, leader, leaderEpoch, isr, 3, replicas, asList()))
+      val partitionStates = Seq(new UpdateMetadataPartitionState()
+        .setTopicName(topic)
+        .setPartitionIndex(0)
+        .setControllerEpoch(controllerEpoch)
+        .setLeader(leader)
+        .setLeaderEpoch(leaderEpoch)
+        .setIsr(isr)
+        .setZkVersion(3)
+        .setReplicas(replicas))
       val version = ApiKeys.UPDATE_METADATA.latestVersion
-      val updateMetadataRequest = new UpdateMetadataRequest.Builder(version, 2, controllerEpoch, brokerEpoch, partitionStates.asJava,
+      val updateMetadataRequest = new UpdateMetadataRequest.Builder(version, 2, controllerEpoch, brokerEpoch, brokerEpoch, partitionStates.asJava,
         brokers.asJava).build()
       cache.updateMetadata(15, updateMetadataRequest)
     }
 
-    val initialBrokerIds = (0 to 2).toSet
+    val initialBrokerIds = (0 to 2)
     updateCache(initialBrokerIds)
     val aliveBrokersFromCache = cache.getAliveBrokers
     // This should not change `aliveBrokersFromCache`
-    updateCache((0 to 3).toSet)
-    assertEquals(initialBrokerIds, aliveBrokersFromCache.map(_.id).toSet)
+    updateCache((0 to 3))
+    assertEquals(initialBrokerIds.toSet, aliveBrokersFromCache.map(_.id).toSet)
   }
 
 }
diff --git a/core/src/test/scala/unit/kafka/server/MetadataRequestTest.scala b/core/src/test/scala/unit/kafka/server/MetadataRequestTest.scala
index ef3dece306325..2b3f321819085 100644
--- a/core/src/test/scala/unit/kafka/server/MetadataRequestTest.scala
+++ b/core/src/test/scala/unit/kafka/server/MetadataRequestTest.scala
@@ -23,6 +23,7 @@ import kafka.network.SocketServer
 import kafka.utils.TestUtils
 import org.apache.kafka.common.Node
 import org.apache.kafka.common.internals.Topic
+import org.apache.kafka.common.message.MetadataRequestData
 import org.apache.kafka.common.protocol.{ApiKeys, Errors}
 import org.apache.kafka.common.requests.{MetadataRequest, MetadataResponse}
 import org.junit.Assert._
@@ -30,10 +31,12 @@ import org.junit.{Before, Test}
 import org.apache.kafka.test.TestUtils.isValidClusterId
 
 import scala.collection.JavaConverters._
+import scala.collection.Seq
 
 class MetadataRequestTest extends BaseRequestTest {
 
-  override def propertyOverrides(properties: Properties) {
+  override def brokerPropertyOverrides(properties: Properties): Unit = {
+    properties.setProperty(KafkaConfig.OffsetsTopicPartitionsProp, "1")
     properties.setProperty(KafkaConfig.DefaultReplicationFactorProp, "2")
     properties.setProperty(KafkaConfig.RackProp, s"rack/${properties.getProperty(KafkaConfig.BrokerIdProp)}")
   }
@@ -44,20 +47,20 @@ class MetadataRequestTest extends BaseRequestTest {
   }
 
   @Test
-  def testClusterIdWithRequestVersion1() {
+  def testClusterIdWithRequestVersion1(): Unit = {
     val v1MetadataResponse = sendMetadataRequest(MetadataRequest.Builder.allTopics.build(1.toShort))
     val v1ClusterId = v1MetadataResponse.clusterId
     assertNull(s"v1 clusterId should be null", v1ClusterId)
   }
 
   @Test
-  def testClusterIdIsValid() {
+  def testClusterIdIsValid(): Unit = {
     val metadataResponse = sendMetadataRequest(MetadataRequest.Builder.allTopics.build(2.toShort))
     isValidClusterId(metadataResponse.clusterId)
   }
 
   @Test
-  def testControllerId() {
+  def testControllerId(): Unit = {
     val controllerServer = servers.find(_.kafkaController.isActive).get
     val controllerId = controllerServer.config.brokerId
     val metadataResponse = sendMetadataRequest(MetadataRequest.Builder.allTopics.build(1.toShort))
@@ -79,7 +82,7 @@ class MetadataRequestTest extends BaseRequestTest {
   }
 
   @Test
-  def testRack() {
+  def testRack(): Unit = {
     val metadataResponse = sendMetadataRequest(MetadataRequest.Builder.allTopics.build(1.toShort))
     // Validate rack matches what's set in generateConfigs() above
     metadataResponse.brokers.asScala.foreach { broker =>
@@ -88,7 +91,7 @@ class MetadataRequestTest extends BaseRequestTest {
   }
 
   @Test
-  def testIsInternal() {
+  def testIsInternal(): Unit = {
     val internalTopic = Topic.GROUP_METADATA_TOPIC_NAME
     val notInternalTopic = "notInternal"
     // create the topics
@@ -109,14 +112,14 @@ class MetadataRequestTest extends BaseRequestTest {
   }
 
   @Test
-  def testNoTopicsRequest() {
+  def testNoTopicsRequest(): Unit = {
     // create some topics
     createTopic("t1", 3, 2)
     createTopic("t2", 3, 2)
 
     // v0, Doesn't support a "no topics" request
     // v1, Empty list represents "no topics"
-    val metadataResponse = sendMetadataRequest(new MetadataRequest(List[String]().asJava, true, 1.toShort))
+    val metadataResponse = sendMetadataRequest(new MetadataRequest.Builder(List[String]().asJava, true, 1.toShort).build)
     assertTrue("Response should have no errors", metadataResponse.errors.isEmpty)
     assertTrue("Response should have no topics", metadataResponse.topicMetadata.isEmpty)
   }
@@ -137,15 +140,15 @@ class MetadataRequestTest extends BaseRequestTest {
     val topic4 = "t4"
     createTopic(topic1, 1, 1)
 
-    val response1 = sendMetadataRequest(new MetadataRequest(Seq(topic1, topic2).asJava, true, ApiKeys.METADATA.latestVersion))
+    val response1 = sendMetadataRequest(new MetadataRequest.Builder(Seq(topic1, topic2).asJava, true, ApiKeys.METADATA.latestVersion).build())
     checkAutoCreatedTopic(topic1, topic2, response1)
 
     // V3 doesn't support a configurable allowAutoTopicCreation, so the fact that we set it to `false` has no effect
-    val response2 = sendMetadataRequest(new MetadataRequest(Seq(topic2, topic3).asJava, false, 3))
+    val response2 = sendMetadataRequest(new MetadataRequest(requestData(List(topic2, topic3), false), 3.toShort))
     checkAutoCreatedTopic(topic2, topic3, response2)
 
     // V4 and higher support a configurable allowAutoTopicCreation
-    val response3 = sendMetadataRequest(new MetadataRequest(Seq(topic3, topic4).asJava, false, 4))
+    val response3 = sendMetadataRequest(new MetadataRequest.Builder(Seq(topic3, topic4).asJava, false, 4.toShort).build)
     assertNull(response3.errors.get(topic3))
     assertEquals(Errors.UNKNOWN_TOPIC_OR_PARTITION, response3.errors.get(topic4))
     assertEquals(None, zkClient.getTopicPartitionCount(topic4))
@@ -195,13 +198,13 @@ class MetadataRequestTest extends BaseRequestTest {
   }
 
   @Test
-  def testAllTopicsRequest() {
+  def testAllTopicsRequest(): Unit = {
     // create some topics
     createTopic("t1", 3, 2)
     createTopic("t2", 3, 2)
 
     // v0, Empty list represents all topics
-    val metadataResponseV0 = sendMetadataRequest(new MetadataRequest(List[String]().asJava, true, 0.toShort))
+    val metadataResponseV0 = sendMetadataRequest(new MetadataRequest(requestData(List(), true), 0.toShort))
     assertTrue("V0 Response should have no errors", metadataResponseV0.errors.isEmpty)
     assertEquals("V0 Response should have 2 (all) topics", 2, metadataResponseV0.topicMetadata.size())
 
@@ -238,8 +241,17 @@ class MetadataRequestTest extends BaseRequestTest {
     }
   }
 
+  def requestData(topics: List[String], allowAutoTopicCreation: Boolean): MetadataRequestData = {
+    val data = new MetadataRequestData
+    if (topics == null) data.setTopics(null)
+    else topics.foreach(topic => data.topics.add(new MetadataRequestData.MetadataRequestTopic().setName(topic)))
+
+    data.setAllowAutoTopicCreation(allowAutoTopicCreation)
+    data
+  }
+
   @Test
-  def testReplicaDownResponse() {
+  def testReplicaDownResponse(): Unit = {
     val replicaDownTopic = "replicaDown"
     val replicaCount = 3
 
@@ -247,7 +259,7 @@ class MetadataRequestTest extends BaseRequestTest {
     createTopic(replicaDownTopic, 1, replicaCount)
 
     // Kill a replica node that is not the leader
-    val metadataResponse = sendMetadataRequest(new MetadataRequest(List(replicaDownTopic).asJava, true, 1.toShort))
+    val metadataResponse = sendMetadataRequest(new MetadataRequest.Builder(List(replicaDownTopic).asJava, true, 1.toShort).build())
     val partitionMetadata = metadataResponse.topicMetadata.asScala.head.partitionMetadata.asScala.head
     val downNode = servers.find { server =>
       val serverId = server.dataPlaneRequestProcessor.brokerId
@@ -258,14 +270,14 @@ class MetadataRequestTest extends BaseRequestTest {
     downNode.shutdown()
 
     TestUtils.waitUntilTrue(() => {
-      val response = sendMetadataRequest(new MetadataRequest(List(replicaDownTopic).asJava, true, 1.toShort))
+      val response = sendMetadataRequest(new MetadataRequest.Builder(List(replicaDownTopic).asJava, true, 1.toShort).build())
       val metadata = response.topicMetadata.asScala.head.partitionMetadata.asScala.head
       val replica = metadata.replicas.asScala.find(_.id == downNode.dataPlaneRequestProcessor.brokerId).get
       replica.host == "" & replica.port == -1
     }, "Replica was not found down", 5000)
 
     // Validate version 0 still filters unavailable replicas and contains error
-    val v0MetadataResponse = sendMetadataRequest(new MetadataRequest(List(replicaDownTopic).asJava, true, 0.toShort))
+    val v0MetadataResponse = sendMetadataRequest(new MetadataRequest(requestData(List(replicaDownTopic), true), 0.toShort))
     val v0BrokerIds = v0MetadataResponse.brokers().asScala.map(_.id).toSeq
     assertTrue("Response should have no errors", v0MetadataResponse.errors.isEmpty)
     assertFalse(s"The downed broker should not be in the brokers list", v0BrokerIds.contains(downNode))
@@ -275,7 +287,7 @@ class MetadataRequestTest extends BaseRequestTest {
     assertTrue(s"Response should have ${replicaCount - 1} replicas", v0PartitionMetadata.replicas.size == replicaCount - 1)
 
     // Validate version 1 returns unavailable replicas with no error
-    val v1MetadataResponse = sendMetadataRequest(new MetadataRequest(List(replicaDownTopic).asJava, true, 1.toShort))
+    val v1MetadataResponse = sendMetadataRequest(new MetadataRequest.Builder(List(replicaDownTopic).asJava, true, 1.toShort).build())
     val v1BrokerIds = v1MetadataResponse.brokers().asScala.map(_.id).toSeq
     assertTrue("Response should have no errors", v1MetadataResponse.errors.isEmpty)
     assertFalse(s"The downed broker should not be in the brokers list", v1BrokerIds.contains(downNode))
diff --git a/core/src/test/scala/unit/kafka/server/OffsetsForLeaderEpochRequestTest.scala b/core/src/test/scala/unit/kafka/server/OffsetsForLeaderEpochRequestTest.scala
index 2cdd2e8d03490..b6f8a142012ab 100644
--- a/core/src/test/scala/unit/kafka/server/OffsetsForLeaderEpochRequestTest.scala
+++ b/core/src/test/scala/unit/kafka/server/OffsetsForLeaderEpochRequestTest.scala
@@ -35,7 +35,8 @@ class OffsetsForLeaderEpochRequestTest extends BaseRequestTest {
     val partition = new TopicPartition(topic, 0)
 
     val epochs = Map(partition -> new OffsetsForLeaderEpochRequest.PartitionData(Optional.empty[Integer], 0)).asJava
-    val request = new OffsetsForLeaderEpochRequest.Builder(ApiKeys.OFFSET_FOR_LEADER_EPOCH.latestVersion, epochs).build()
+    val request = OffsetsForLeaderEpochRequest.Builder.forFollower(
+      ApiKeys.OFFSET_FOR_LEADER_EPOCH.latestVersion, epochs, 1).build()
 
     // Unknown topic
     val randomBrokerId = servers.head.config.brokerId
@@ -61,8 +62,8 @@ class OffsetsForLeaderEpochRequestTest extends BaseRequestTest {
     def assertResponseErrorForEpoch(error: Errors, brokerId: Int, currentLeaderEpoch: Optional[Integer]): Unit = {
       val epochs = Map(topicPartition -> new OffsetsForLeaderEpochRequest.PartitionData(
         currentLeaderEpoch, 0)).asJava
-      val request = new OffsetsForLeaderEpochRequest.Builder(ApiKeys.OFFSET_FOR_LEADER_EPOCH.latestVersion, epochs)
-        .build()
+      val request = OffsetsForLeaderEpochRequest.Builder.forFollower(
+        ApiKeys.OFFSET_FOR_LEADER_EPOCH.latestVersion, epochs, 1).build()
       assertResponseError(error, brokerId, request)
     }
 
diff --git a/core/src/test/scala/unit/kafka/server/PreferredControllerTest.scala b/core/src/test/scala/unit/kafka/server/PreferredControllerTest.scala
new file mode 100644
index 0000000000000..3a130a22baafa
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/server/PreferredControllerTest.scala
@@ -0,0 +1,173 @@
+/**
+  * 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 java.util.Properties
+
+import kafka.server.KafkaConfig.fromProps
+import kafka.utils.CoreUtils._
+import kafka.utils.TestUtils
+import kafka.utils.TestUtils._
+import kafka.zk.ZooKeeperTestHarness
+import org.apache.kafka.clients.admin._
+import org.apache.kafka.common.network.ListenerName
+import org.apache.kafka.common.security.auth.SecurityProtocol
+import org.junit.Assert._
+import org.junit.{After, Test}
+import org.scalatest.Assertions.fail
+
+import scala.collection.JavaConverters._
+import scala.collection.Map
+
+class PreferredControllerTest extends ZooKeeperTestHarness {
+
+  var brokers: Seq[KafkaServer] = null
+
+  @After
+  override def tearDown() {
+    shutdownServers(brokers)
+    super.tearDown()
+  }
+
+  @Test
+  def testPartitionCreatedByAdminClientShouldNotBeAssignedToPreferredControllers(): Unit = {
+    val brokerConfigs = Seq((0, false), (1, true), (2, false))
+    createBrokersWithPreferredControllers(brokerConfigs, true)
+
+    val brokerList = TestUtils.bootstrapServers(brokers, ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT))
+    val adminClientConfig = new Properties
+    adminClientConfig.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, brokerList)
+    val client = AdminClient.create(adminClientConfig)
+
+    TestUtils.waitUntilControllerElected(zkClient)
+    // create topic using admin client
+    val future1 = client.createTopics(Seq("topic1").map(new NewTopic(_, 3, 2.toShort)).asJava,
+      new CreateTopicsOptions()).all()
+    future1.get()
+
+    assertTrue("topic1 should not be in broker 1", ensureTopicNotInBrokers("topic1", Set(1)))
+
+    val future2 = client.createPartitions(Map("topic1" -> NewPartitions.increaseTo(5)).asJava).all()
+    future2.get()
+
+    assertTrue("topic1 should not be in broker 1 after increasing partition count",
+      ensureTopicNotInBrokers("topic1", Set(1)))
+
+    client.close()
+  }
+
+  @Test
+  def testElectionWithoutPreferredControllersAndNoFallback(): Unit = {
+    val brokerConfigs = Seq((0, false), (1, false), (2, false))
+    createBrokersWithPreferredControllers(brokerConfigs, false)
+    // no broker can be elected as controller
+    ensureControllersInBrokers(Seq.empty, 5000L)
+  }
+
+  @Test
+  def testPreferredControllerElection(): Unit = {
+    val brokerConfigs = Seq((0, false), (1, true), (2, false))
+    createBrokersWithPreferredControllers(brokerConfigs, false)
+    // only broker 1 can be elected since it is the only preferred controller node
+    ensureControllersInBrokers(Seq(1))
+  }
+
+
+  @Test
+  def testNonPreferredControllerResignation(): Unit = {
+    val brokerConfigs = Seq((0, false), (1, true), (2, false))
+    createBrokersWithPreferredControllers(brokerConfigs, true)
+
+    // broker 1 should be elected since it is the only preferred controller node
+    ensureControllersInBrokers(Seq(1))
+    brokers(1).shutdown()
+
+    // broker 0 and broker 2 can become controller when broker 1 is offline
+    ensureControllersInBrokers(Seq(0, 2))
+    brokers(1).startup()
+    // broker 1 regains controllership
+    ensureControllersInBrokers(Seq(1))
+  }
+
+  @Test
+  def testDynamicAllowPreferredControllerFallback(): Unit = {
+    val brokerConfigs = Seq((0, false), (1, false), (2, false))
+    createBrokersWithPreferredControllers(brokerConfigs, false)
+
+    // non preferred controller nodes cannot be elected as the controller if fallback is not allowed
+    ensureControllersInBrokers(Seq.empty, 5000L)
+    setAllowPreferredControllerFallback(true)
+    // controller can be now elected among non preferred controller nodes
+    TestUtils.waitUntilControllerElected(zkClient)
+  }
+
+  @Test
+  def testCurrentControllerDoesNotResignWithoutPreferredControllersAndNoFallback(): Unit = {
+    val brokerConfigs = Seq((0, false), (1, false), (2, false))
+    createBrokersWithPreferredControllers(brokerConfigs, true)
+
+    val controllerId = TestUtils.waitUntilControllerElected(zkClient)
+
+    setAllowPreferredControllerFallback(false)
+
+    // current controller does not move
+    ensureControllersInBrokers(Seq(controllerId))
+  }
+
+  private def ensureControllersInBrokers(brokerIds: Seq[Int], timeout: Long = 15000L): Unit = {
+    val (controllerId, _) = TestUtils.computeUntilTrue(zkClient.getControllerId, waitTime = timeout) {
+      case controller =>
+        controller.isDefined && (brokerIds.isEmpty || brokerIds.contains(controller))
+    }
+    if (brokerIds.isEmpty) {
+      assertTrue("there should not be any controller", controllerId.isEmpty)
+    } else {
+      assertTrue(s"Controller should be elected in $brokerIds",
+        brokerIds.contains(controllerId.getOrElse(fail(s"Controller not elected after $timeout ms"))))
+    }
+  }
+
+  private def ensureTopicNotInBrokers(topic: String, brokerIds: Set[Int]): Boolean = {
+    val topicAssignment = zkClient.getReplicaAssignmentForTopics(Set(topic))
+    topicAssignment.map(_._2).flatten.toSet.intersect(brokerIds).isEmpty
+  }
+
+  /**
+    * @param brokerConfigs: a list of (brokerid, preferredController) configs
+    * @param allowFallback: "allow.preferred.controller.fallback" config
+    */
+  private def createBrokersWithPreferredControllers(brokerConfigs: Seq[(Int, Boolean)],  allowFallback: Boolean): Unit = {
+    brokers = brokerConfigs.map {
+      case (id, preferredController) =>
+        val props: Properties = createBrokerConfig(id, zkConnect)
+        props.put(KafkaConfig.PreferredControllerProp, preferredController.toString)
+        props.put(KafkaConfig.AllowPreferredControllerFallbackProp, allowFallback.toString)
+        createServer(fromProps(props))
+    }
+  }
+
+  private def setAllowPreferredControllerFallback(allowFallback: Boolean): Unit = {
+    adminZkClient.changeBrokerConfig(None,
+      propsWith((KafkaConfig.AllowPreferredControllerFallbackProp, allowFallback.toString)))
+
+    TestUtils.waitUntilTrue(() => {
+        brokers.forall(_.config.allowPreferredControllerFallback == allowFallback)
+      },
+      s"fail to set ${KafkaConfig.AllowPreferredControllerFallbackProp} to ${allowFallback}", 5000)
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/server/ProduceRequestTest.scala b/core/src/test/scala/unit/kafka/server/ProduceRequestTest.scala
index 906de71ecc6f9..3bc8d0aacb8f3 100644
--- a/core/src/test/scala/unit/kafka/server/ProduceRequestTest.scala
+++ b/core/src/test/scala/unit/kafka/server/ProduceRequestTest.scala
@@ -17,8 +17,10 @@
 
 package kafka.server
 
+import java.nio.ByteBuffer
 import java.util.Properties
 
+import com.yammer.metrics.Metrics
 import kafka.log.LogConfig
 import kafka.message.ZStdCompressionCodec
 import kafka.utils.TestUtils
@@ -28,6 +30,7 @@ import org.apache.kafka.common.record._
 import org.apache.kafka.common.requests.{ProduceRequest, ProduceResponse}
 import org.junit.Assert._
 import org.junit.Test
+import org.scalatest.Assertions.fail
 
 import scala.collection.JavaConverters._
 
@@ -37,8 +40,10 @@ import scala.collection.JavaConverters._
   */
 class ProduceRequestTest extends BaseRequestTest {
 
+  val metricsKeySet = Metrics.defaultRegistry.allMetrics.keySet.asScala
+
   @Test
-  def testSimpleProduceRequest() {
+  def testSimpleProduceRequest(): Unit = {
     val (partition, leader) = createTopicAndFindPartitionWithLeader("topic")
 
     def sendAndCheck(memoryRecords: MemoryRecords, expectedOffset: Long): ProduceResponse.PartitionResponse = {
@@ -52,6 +57,7 @@ class ProduceRequestTest extends BaseRequestTest {
       assertEquals(Errors.NONE, partitionResponse.error)
       assertEquals(expectedOffset, partitionResponse.baseOffset)
       assertEquals(-1, partitionResponse.logAppendTime)
+      assertTrue(partitionResponse.recordErrors.isEmpty)
       partitionResponse
     }
 
@@ -64,7 +70,40 @@ class ProduceRequestTest extends BaseRequestTest {
   }
 
   @Test
-  def testProduceToNonReplica() {
+  def testProduceWithInvalidTimestamp(): Unit = {
+    val topic = "topic"
+    val partition = 0
+    val topicConfig = new Properties
+    topicConfig.setProperty(LogConfig.MessageTimestampDifferenceMaxMsProp, "1000")
+    val partitionToLeader = TestUtils.createTopic(zkClient, topic, 1, 1, servers, topicConfig)
+    val leader = partitionToLeader(partition)
+
+    def createRecords(magicValue: Byte,
+                      timestamp: Long = RecordBatch.NO_TIMESTAMP,
+                      codec: CompressionType): MemoryRecords = {
+      val buf = ByteBuffer.allocate(512)
+      val builder = MemoryRecords.builder(buf, magicValue, codec, TimestampType.CREATE_TIME, 0L)
+      builder.appendWithOffset(0, timestamp, null, "hello".getBytes)
+      builder.appendWithOffset(1, timestamp, null, "there".getBytes)
+      builder.appendWithOffset(2, timestamp, null, "beautiful".getBytes)
+      builder.build()
+    }
+
+    val records = createRecords(RecordBatch.MAGIC_VALUE_V2, System.currentTimeMillis() - 1001L, CompressionType.GZIP)
+    val topicPartition = new TopicPartition("topic", partition)
+    val partitionRecords = Map(topicPartition -> records)
+    val produceResponse = sendProduceRequest(leader, ProduceRequest.Builder.forCurrentMagic(-1, 3000, partitionRecords.asJava).build())
+    val (tp, partitionResponse) = produceResponse.responses.asScala.head
+    assertEquals(topicPartition, tp)
+    assertEquals(Errors.INVALID_TIMESTAMP, partitionResponse.error)
+    assertEquals(1, partitionResponse.recordErrors.size())
+    assertEquals(0, partitionResponse.recordErrors.get(0).batchIndex)
+    assertNull(partitionResponse.recordErrors.get(0).message)
+    assertNotNull(partitionResponse.errorMessage)
+  }
+
+  @Test
+  def testProduceToNonReplica(): Unit = {
     val topic = "topic"
     val partition = 0
 
@@ -95,7 +134,7 @@ class ProduceRequestTest extends BaseRequestTest {
   }
 
   @Test
-  def testCorruptLz4ProduceRequest() {
+  def testCorruptLz4ProduceRequest(): Unit = {
     val (partition, leader) = createTopicAndFindPartitionWithLeader("topic")
     val timestamp = 1000000
     val memoryRecords = MemoryRecords.withRecords(CompressionType.LZ4,
@@ -113,6 +152,8 @@ class ProduceRequestTest extends BaseRequestTest {
     assertEquals(Errors.CORRUPT_MESSAGE, partitionResponse.error)
     assertEquals(-1, partitionResponse.baseOffset)
     assertEquals(-1, partitionResponse.logAppendTime)
+    assertEquals(metricsKeySet.count(_.getMBeanName.endsWith(s"${BrokerTopicStats.InvalidMessageCrcRecordsPerSec}")), 1)
+    assertTrue(TestUtils.meterCount(s"${BrokerTopicStats.InvalidMessageCrcRecordsPerSec}") > 0)
   }
 
   @Test
diff --git a/core/src/test/scala/unit/kafka/server/ReplicaAlterLogDirsThreadTest.scala b/core/src/test/scala/unit/kafka/server/ReplicaAlterLogDirsThreadTest.scala
index 9710cf22b7704..05015087304e7 100644
--- a/core/src/test/scala/unit/kafka/server/ReplicaAlterLogDirsThreadTest.scala
+++ b/core/src/test/scala/unit/kafka/server/ReplicaAlterLogDirsThreadTest.scala
@@ -18,32 +18,239 @@ package kafka.server
 
 import java.util.Optional
 
-import kafka.cluster.{BrokerEndPoint, Partition, Replica}
-import kafka.log.LogManager
+import kafka.api.Request
+import kafka.cluster.{BrokerEndPoint, Partition}
+import kafka.log.{Log, LogManager}
 import kafka.server.AbstractFetcherThread.ResultWithPartitions
+import kafka.server.QuotaFactory.UnboundedQuota
 import kafka.utils.{DelayedItem, TestUtils}
 import org.apache.kafka.common.TopicPartition
 import org.apache.kafka.common.errors.KafkaStorageException
 import org.apache.kafka.common.protocol.Errors
-import org.apache.kafka.common.requests.{EpochEndOffset, OffsetsForLeaderEpochRequest}
+import org.apache.kafka.common.record.MemoryRecords
 import org.apache.kafka.common.requests.EpochEndOffset.{UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET}
+import org.apache.kafka.common.requests.{EpochEndOffset, FetchRequest, IsolationLevel, OffsetsForLeaderEpochRequest}
 import org.easymock.EasyMock._
-import org.easymock.{Capture, CaptureType, EasyMock, IAnswer}
+import org.easymock.{Capture, CaptureType, EasyMock, IAnswer, IExpectationSetters}
 import org.junit.Assert._
 import org.junit.Test
+import org.mockito.Mockito.{doNothing, when}
+import org.mockito.invocation.InvocationOnMock
+import org.mockito.stubbing.Answer
+import org.mockito.{ArgumentCaptor, ArgumentMatchers, Mockito}
 
-import scala.collection.JavaConverters._
 import scala.collection.{Map, Seq}
+import scala.jdk.CollectionConverters._
 
 class ReplicaAlterLogDirsThreadTest {
 
   private val t1p0 = new TopicPartition("topic1", 0)
   private val t1p1 = new TopicPartition("topic1", 1)
+  private val failedPartitions = new FailedPartitions
 
   private def offsetAndEpoch(fetchOffset: Long, leaderEpoch: Int = 1): OffsetAndEpoch = {
     OffsetAndEpoch(offset = fetchOffset, leaderEpoch = leaderEpoch)
   }
 
+  @Test
+  def shouldNotAddPartitionIfFutureLogIsNotDefined(): Unit = {
+    val brokerId = 1
+    val config = KafkaConfig.fromProps(TestUtils.createBrokerConfig(brokerId, "localhost:1234"))
+
+    val replicaManager = Mockito.mock(classOf[ReplicaManager])
+    val quotaManager = Mockito.mock(classOf[ReplicationQuotaManager])
+
+    when(replicaManager.futureLogExists(t1p0)).thenReturn(false)
+
+    val endPoint = new BrokerEndPoint(0, "localhost", 1000)
+    val thread = new ReplicaAlterLogDirsThread(
+      "alter-logs-dirs-thread",
+      sourceBroker = endPoint,
+      brokerConfig = config,
+      failedPartitions = failedPartitions,
+      replicaMgr = replicaManager,
+      quota = quotaManager,
+      brokerTopicStats = new BrokerTopicStats)
+
+    val addedPartitions = thread.addPartitions(Map(t1p0 -> offsetAndEpoch(0L)))
+    assertEquals(Set.empty, addedPartitions)
+    assertEquals(0, thread.partitionCount())
+    assertEquals(None, thread.fetchState(t1p0))
+  }
+
+  @Test
+  def shouldUpdateLeaderEpochAfterFencedEpochError(): Unit = {
+    val brokerId = 1
+    val config = KafkaConfig.fromProps(TestUtils.createBrokerConfig(brokerId, "localhost:1234"))
+
+    val partition = Mockito.mock(classOf[Partition])
+    val replicaManager = Mockito.mock(classOf[ReplicaManager])
+    val quotaManager = Mockito.mock(classOf[ReplicationQuotaManager])
+    val futureLog = Mockito.mock(classOf[Log])
+
+    val leaderEpoch = 5
+    val logEndOffset = 0
+
+    when(replicaManager.futureLocalLogOrException(t1p0)).thenReturn(futureLog)
+    when(replicaManager.futureLogExists(t1p0)).thenReturn(true)
+    when(replicaManager.nonOfflinePartition(t1p0)).thenReturn(Some(partition))
+    when(replicaManager.getPartitionOrException(t1p0, expectLeader = false)).thenReturn(partition)
+
+    when(quotaManager.isQuotaExceeded).thenReturn(false)
+
+    when(partition.lastOffsetForLeaderEpoch(Optional.empty(), leaderEpoch, fetchOnlyFromLeader = false))
+      .thenReturn(new EpochEndOffset(leaderEpoch, logEndOffset))
+    when(partition.futureLocalLogOrException).thenReturn(futureLog)
+    doNothing().when(partition).truncateTo(offset = 0, isFuture = true)
+    when(partition.maybeReplaceCurrentWithFutureReplica()).thenReturn(true)
+
+    when(futureLog.logStartOffset).thenReturn(0L)
+    when(futureLog.logEndOffset).thenReturn(0L)
+    when(futureLog.latestEpoch).thenReturn(None)
+
+    val fencedRequestData = new FetchRequest.PartitionData(0L, 0L,
+      config.replicaFetchMaxBytes, Optional.of(leaderEpoch - 1))
+    val fencedResponseData = FetchPartitionData(
+      error = Errors.FENCED_LEADER_EPOCH,
+      highWatermark = -1,
+      logStartOffset = -1,
+      records = MemoryRecords.EMPTY,
+      lastStableOffset = None,
+      abortedTransactions = None,
+      preferredReadReplica = None)
+    mockFetchFromCurrentLog(t1p0, fencedRequestData, config, replicaManager, fencedResponseData)
+
+    val endPoint = new BrokerEndPoint(0, "localhost", 1000)
+    val thread = new ReplicaAlterLogDirsThread(
+      "alter-logs-dirs-thread",
+      sourceBroker = endPoint,
+      brokerConfig = config,
+      failedPartitions = failedPartitions,
+      replicaMgr = replicaManager,
+      quota = quotaManager,
+      brokerTopicStats = new BrokerTopicStats)
+
+    // Initially we add the partition with an older epoch which results in an error
+    thread.addPartitions(Map(t1p0 -> offsetAndEpoch(fetchOffset = 0L, leaderEpoch - 1)))
+    assertTrue(thread.fetchState(t1p0).isDefined)
+    assertEquals(1, thread.partitionCount())
+
+    thread.doWork()
+
+    assertTrue(failedPartitions.contains(t1p0))
+    assertEquals(None, thread.fetchState(t1p0))
+    assertEquals(0, thread.partitionCount())
+
+    // Next we update the epoch and assert that we can continue
+    thread.addPartitions(Map(t1p0 -> offsetAndEpoch(fetchOffset = 0L, leaderEpoch)))
+    assertEquals(Some(leaderEpoch), thread.fetchState(t1p0).map(_.currentLeaderEpoch))
+    assertEquals(1, thread.partitionCount())
+
+    val requestData = new FetchRequest.PartitionData(0L, 0L,
+      config.replicaFetchMaxBytes, Optional.of(leaderEpoch))
+    val responseData = FetchPartitionData(
+      error = Errors.NONE,
+      highWatermark = 0L,
+      logStartOffset = 0L,
+      records = MemoryRecords.EMPTY,
+      lastStableOffset = None,
+      abortedTransactions = None,
+      preferredReadReplica = None)
+    mockFetchFromCurrentLog(t1p0, requestData, config, replicaManager, responseData)
+
+    thread.doWork()
+
+    assertFalse(failedPartitions.contains(t1p0))
+    assertEquals(None, thread.fetchState(t1p0))
+    assertEquals(0, thread.partitionCount())
+  }
+
+  @Test
+  def shouldReplaceCurrentLogDirWhenCaughtUp(): Unit = {
+    val brokerId = 1
+    val config = KafkaConfig.fromProps(TestUtils.createBrokerConfig(brokerId, "localhost:1234"))
+
+    val partition = Mockito.mock(classOf[Partition])
+    val replicaManager = Mockito.mock(classOf[ReplicaManager])
+    val quotaManager = Mockito.mock(classOf[ReplicationQuotaManager])
+    val futureLog = Mockito.mock(classOf[Log])
+
+    val leaderEpoch = 5
+    val logEndOffset = 0
+
+    when(replicaManager.futureLocalLogOrException(t1p0)).thenReturn(futureLog)
+    when(replicaManager.futureLogExists(t1p0)).thenReturn(true)
+    when(replicaManager.nonOfflinePartition(t1p0)).thenReturn(Some(partition))
+    when(replicaManager.getPartitionOrException(t1p0, expectLeader = false)).thenReturn(partition)
+
+    when(quotaManager.isQuotaExceeded).thenReturn(false)
+
+    when(partition.lastOffsetForLeaderEpoch(Optional.empty(), leaderEpoch, fetchOnlyFromLeader = false))
+      .thenReturn(new EpochEndOffset(leaderEpoch, logEndOffset))
+    when(partition.futureLocalLogOrException).thenReturn(futureLog)
+    doNothing().when(partition).truncateTo(offset = 0, isFuture = true)
+    when(partition.maybeReplaceCurrentWithFutureReplica()).thenReturn(true)
+
+    when(futureLog.logStartOffset).thenReturn(0L)
+    when(futureLog.logEndOffset).thenReturn(0L)
+    when(futureLog.latestEpoch).thenReturn(None)
+
+    val requestData = new FetchRequest.PartitionData(0L, 0L,
+      config.replicaFetchMaxBytes, Optional.of(leaderEpoch))
+    val responseData = FetchPartitionData(
+      error = Errors.NONE,
+      highWatermark = 0L,
+      logStartOffset = 0L,
+      records = MemoryRecords.EMPTY,
+      lastStableOffset = None,
+      abortedTransactions = None,
+      preferredReadReplica = None)
+    mockFetchFromCurrentLog(t1p0, requestData, config, replicaManager, responseData)
+
+    val endPoint = new BrokerEndPoint(0, "localhost", 1000)
+    val thread = new ReplicaAlterLogDirsThread(
+      "alter-logs-dirs-thread",
+      sourceBroker = endPoint,
+      brokerConfig = config,
+      failedPartitions = failedPartitions,
+      replicaMgr = replicaManager,
+      quota = quotaManager,
+      brokerTopicStats = new BrokerTopicStats)
+
+    thread.addPartitions(Map(t1p0 -> offsetAndEpoch(fetchOffset = 0L, leaderEpoch)))
+    assertTrue(thread.fetchState(t1p0).isDefined)
+    assertEquals(1, thread.partitionCount())
+
+    thread.doWork()
+
+    assertEquals(None, thread.fetchState(t1p0))
+    assertEquals(0, thread.partitionCount())
+  }
+
+  private def mockFetchFromCurrentLog(topicPartition: TopicPartition,
+                                      requestData: FetchRequest.PartitionData,
+                                      config: KafkaConfig,
+                                      replicaManager: ReplicaManager,
+                                      responseData: FetchPartitionData): Unit = {
+    val callbackCaptor: ArgumentCaptor[Seq[(TopicPartition, FetchPartitionData)] => Unit] =
+      ArgumentCaptor.forClass(classOf[Seq[(TopicPartition, FetchPartitionData)] => Unit])
+    when(replicaManager.fetchMessages(
+      timeout = ArgumentMatchers.eq(0L),
+      replicaId = ArgumentMatchers.eq(Request.FutureLocalReplicaId),
+      fetchMinBytes = ArgumentMatchers.eq(0),
+      fetchMaxBytes = ArgumentMatchers.eq(config.replicaFetchResponseMaxBytes),
+      hardMaxBytesLimit = ArgumentMatchers.eq(false),
+      fetchInfos = ArgumentMatchers.eq(Seq(topicPartition -> requestData)),
+      quota = ArgumentMatchers.eq(UnboundedQuota),
+      responseCallback = callbackCaptor.capture(),
+      isolationLevel = ArgumentMatchers.eq(IsolationLevel.READ_UNCOMMITTED),
+      clientMetadata = ArgumentMatchers.eq(None)
+    ))
+    .thenAnswer(new Answer[Unit] {
+      override def answer(invocation: InvocationOnMock): Unit = callbackCaptor.getValue.apply(Seq((topicPartition, responseData)))
+    })
+  }
+
   @Test
   def issuesEpochRequestFromLocalReplica(): Unit = {
     val config = KafkaConfig.fromProps(TestUtils.createBrokerConfig(1, "localhost:1234"))
@@ -78,6 +285,7 @@ class ReplicaAlterLogDirsThreadTest {
       "alter-logs-dirs-thread-test1",
       sourceBroker = endPoint,
       brokerConfig = config,
+      failedPartitions = failedPartitions,
       replicaMgr = replicaManager,
       quota = null,
       brokerTopicStats = null)
@@ -123,6 +331,7 @@ class ReplicaAlterLogDirsThreadTest {
       "alter-logs-dirs-thread-test1",
       sourceBroker = endPoint,
       brokerConfig = config,
+      failedPartitions = failedPartitions,
       replicaMgr = replicaManager,
       quota = null,
       brokerTopicStats = null)
@@ -150,11 +359,11 @@ class ReplicaAlterLogDirsThreadTest {
     val config = KafkaConfig.fromProps(TestUtils.createBrokerConfig(1, "localhost:1234"))
     val quotaManager: ReplicationQuotaManager = createNiceMock(classOf[ReplicationQuotaManager])
     val logManager: LogManager = createMock(classOf[LogManager])
-    val replicaT1p0: Replica = createNiceMock(classOf[Replica])
-    val replicaT1p1: Replica = createNiceMock(classOf[Replica])
+    val logT1p0: Log = createNiceMock(classOf[Log])
+    val logT1p1: Log = createNiceMock(classOf[Log])
     // one future replica mock because our mocking methods return same values for both future replicas
-    val futureReplicaT1p0: Replica = createNiceMock(classOf[Replica])
-    val futureReplicaT1p1: Replica = createNiceMock(classOf[Replica])
+    val futureLogT1p0: Log = createNiceMock(classOf[Log])
+    val futureLogT1p1: Log = createNiceMock(classOf[Log])
     val partitionT1p0: Partition = createMock(classOf[Partition])
     val partitionT1p1: Partition = createMock(classOf[Partition])
     val replicaManager: ReplicaManager = createMock(classOf[ReplicaManager])
@@ -170,33 +379,34 @@ class ReplicaAlterLogDirsThreadTest {
       .andStubReturn(partitionT1p0)
     expect(replicaManager.getPartitionOrException(t1p1, expectLeader = false))
       .andStubReturn(partitionT1p1)
-    expect(replicaManager.futureLocalReplicaOrException(t1p0)).andStubReturn(futureReplicaT1p0)
-    expect(replicaManager.futureLocalReplicaOrException(t1p1)).andStubReturn(futureReplicaT1p1)
+    expect(replicaManager.futureLocalLogOrException(t1p0)).andStubReturn(futureLogT1p0)
+    expect(replicaManager.futureLogExists(t1p0)).andStubReturn(true)
+    expect(replicaManager.futureLocalLogOrException(t1p1)).andStubReturn(futureLogT1p1)
+    expect(replicaManager.futureLogExists(t1p1)).andStubReturn(true)
     expect(partitionT1p0.truncateTo(capture(truncateCaptureT1p0), anyBoolean())).anyTimes()
     expect(partitionT1p1.truncateTo(capture(truncateCaptureT1p1), anyBoolean())).anyTimes()
 
-    expect(futureReplicaT1p0.logEndOffset).andReturn(futureReplicaLEO).anyTimes()
-    expect(futureReplicaT1p1.logEndOffset).andReturn(futureReplicaLEO).anyTimes()
+    expect(futureLogT1p0.logEndOffset).andReturn(futureReplicaLEO).anyTimes()
+    expect(futureLogT1p1.logEndOffset).andReturn(futureReplicaLEO).anyTimes()
 
-    expect(futureReplicaT1p0.latestEpoch).andReturn(Some(leaderEpoch)).anyTimes()
-    expect(futureReplicaT1p0.endOffsetForEpoch(leaderEpoch)).andReturn(
+    expect(futureLogT1p0.latestEpoch).andReturn(Some(leaderEpoch)).anyTimes()
+    expect(futureLogT1p0.endOffsetForEpoch(leaderEpoch)).andReturn(
       Some(OffsetAndEpoch(futureReplicaLEO, leaderEpoch))).anyTimes()
     expect(partitionT1p0.lastOffsetForLeaderEpoch(Optional.of(1), leaderEpoch, fetchOnlyFromLeader = false))
       .andReturn(new EpochEndOffset(leaderEpoch, replicaT1p0LEO))
       .anyTimes()
 
-    expect(futureReplicaT1p1.latestEpoch).andReturn(Some(leaderEpoch)).anyTimes()
-    expect(futureReplicaT1p1.endOffsetForEpoch(leaderEpoch)).andReturn(
+    expect(futureLogT1p1.latestEpoch).andReturn(Some(leaderEpoch)).anyTimes()
+    expect(futureLogT1p1.endOffsetForEpoch(leaderEpoch)).andReturn(
       Some(OffsetAndEpoch(futureReplicaLEO, leaderEpoch))).anyTimes()
     expect(partitionT1p1.lastOffsetForLeaderEpoch(Optional.of(1), leaderEpoch, fetchOnlyFromLeader = false))
       .andReturn(new EpochEndOffset(leaderEpoch, replicaT1p1LEO))
       .anyTimes()
 
     expect(replicaManager.logManager).andReturn(logManager).anyTimes()
-    stubWithFetchMessages(replicaT1p0, replicaT1p1, futureReplicaT1p0, partitionT1p0, replicaManager, responseCallback)
+    stubWithFetchMessages(logT1p0, logT1p1, futureLogT1p0, partitionT1p0, replicaManager, responseCallback)
 
-    replay(replicaManager, logManager, quotaManager, replicaT1p0, replicaT1p1,
-      futureReplicaT1p0, partitionT1p0, partitionT1p1)
+    replay(replicaManager, logManager, quotaManager, partitionT1p0, partitionT1p1, logT1p0, logT1p1, futureLogT1p0, futureLogT1p1)
 
     //Create the thread
     val endPoint = new BrokerEndPoint(0, "localhost", 1000)
@@ -204,6 +414,7 @@ class ReplicaAlterLogDirsThreadTest {
       "alter-logs-dirs-thread-test1",
       sourceBroker = endPoint,
       brokerConfig = config,
+      failedPartitions = failedPartitions,
       replicaMgr = replicaManager,
       quota = quotaManager,
       brokerTopicStats = null)
@@ -227,9 +438,9 @@ class ReplicaAlterLogDirsThreadTest {
     val config = KafkaConfig.fromProps(TestUtils.createBrokerConfig(1, "localhost:1234"))
     val quotaManager: ReplicationQuotaManager = createNiceMock(classOf[ReplicationQuotaManager])
     val logManager: LogManager = createMock(classOf[LogManager])
-    val replica: Replica = createNiceMock(classOf[Replica])
+    val log: Log = createNiceMock(classOf[Log])
     // one future replica mock because our mocking methods return same values for both future replicas
-    val futureReplica: Replica = createNiceMock(classOf[Replica])
+    val futureLog: Log = createNiceMock(classOf[Log])
     val partition: Partition = createMock(classOf[Partition])
     val replicaManager: ReplicaManager = createMock(classOf[ReplicaManager])
     val responseCallback: Capture[Seq[(TopicPartition, FetchPartitionData)] => Unit]  = EasyMock.newCapture()
@@ -243,31 +454,32 @@ class ReplicaAlterLogDirsThreadTest {
     //Stubs
     expect(replicaManager.getPartitionOrException(t1p0, expectLeader = false))
       .andStubReturn(partition)
-    expect(replicaManager.futureLocalReplicaOrException(t1p0)).andStubReturn(futureReplica)
+    expect(replicaManager.futureLocalLogOrException(t1p0)).andStubReturn(futureLog)
+    expect(replicaManager.futureLogExists(t1p0)).andStubReturn(true)
 
     expect(partition.truncateTo(capture(truncateToCapture), EasyMock.eq(true))).anyTimes()
-    expect(futureReplica.logEndOffset).andReturn(futureReplicaLEO).anyTimes()
-    expect(futureReplica.latestEpoch).andReturn(Some(leaderEpoch)).once()
-    expect(futureReplica.latestEpoch).andReturn(Some(leaderEpoch - 2)).once()
+    expect(futureLog.logEndOffset).andReturn(futureReplicaLEO).anyTimes()
+    expect(futureLog.latestEpoch).andReturn(Some(leaderEpoch)).once()
+    expect(futureLog.latestEpoch).andReturn(Some(leaderEpoch - 2)).once()
 
     // leader replica truncated and fetched new offsets with new leader epoch
     expect(partition.lastOffsetForLeaderEpoch(Optional.of(1), leaderEpoch, fetchOnlyFromLeader = false))
       .andReturn(new EpochEndOffset(leaderEpoch - 1, replicaLEO))
       .anyTimes()
     // but future replica does not know about this leader epoch, so returns a smaller leader epoch
-    expect(futureReplica.endOffsetForEpoch(leaderEpoch - 1)).andReturn(
+    expect(futureLog.endOffsetForEpoch(leaderEpoch - 1)).andReturn(
       Some(OffsetAndEpoch(futureReplicaLEO, leaderEpoch - 2))).anyTimes()
     // finally, the leader replica knows about the leader epoch and returns end offset
     expect(partition.lastOffsetForLeaderEpoch(Optional.of(1), leaderEpoch - 2, fetchOnlyFromLeader = false))
       .andReturn(new EpochEndOffset(leaderEpoch - 2, replicaEpochEndOffset))
       .anyTimes()
-    expect(futureReplica.endOffsetForEpoch(leaderEpoch - 2)).andReturn(
+    expect(futureLog.endOffsetForEpoch(leaderEpoch - 2)).andReturn(
       Some(OffsetAndEpoch(futureReplicaEpochEndOffset, leaderEpoch - 2))).anyTimes()
 
     expect(replicaManager.logManager).andReturn(logManager).anyTimes()
-    stubWithFetchMessages(replica, replica, futureReplica, partition, replicaManager, responseCallback)
+    stubWithFetchMessages(log, null, futureLog, partition, replicaManager, responseCallback)
 
-    replay(replicaManager, logManager, quotaManager, replica, futureReplica, partition)
+    replay(replicaManager, logManager, quotaManager, partition, log, futureLog)
 
     //Create the thread
     val endPoint = new BrokerEndPoint(0, "localhost", 1000)
@@ -275,6 +487,7 @@ class ReplicaAlterLogDirsThreadTest {
       "alter-logs-dirs-thread-test1",
       sourceBroker = endPoint,
       brokerConfig = config,
+      failedPartitions = failedPartitions,
       replicaMgr = replicaManager,
       quota = quotaManager,
       brokerTopicStats = null)
@@ -300,29 +513,28 @@ class ReplicaAlterLogDirsThreadTest {
     val config = KafkaConfig.fromProps(TestUtils.createBrokerConfig(1, "localhost:1234"))
     val quotaManager: ReplicationQuotaManager = createNiceMock(classOf[ReplicationQuotaManager])
     val logManager: LogManager = createMock(classOf[LogManager])
-    val replica: Replica = createNiceMock(classOf[Replica])
-    val futureReplica: Replica = createNiceMock(classOf[Replica])
+    val log: Log = createNiceMock(classOf[Log])
+    val futureLog: Log = createNiceMock(classOf[Log])
     val partition: Partition = createMock(classOf[Partition])
     val replicaManager: ReplicaManager = createMock(classOf[ReplicaManager])
     val responseCallback: Capture[Seq[(TopicPartition, FetchPartitionData)] => Unit]  = EasyMock.newCapture()
 
     val initialFetchOffset = 100
-    val futureReplicaLEO = 111
 
     //Stubs
     expect(replicaManager.getPartitionOrException(t1p0, expectLeader = false))
       .andStubReturn(partition)
     expect(partition.truncateTo(capture(truncated), isFuture = EasyMock.eq(true))).anyTimes()
-    expect(replicaManager.futureLocalReplicaOrException(t1p0)).andStubReturn(futureReplica)
+    expect(replicaManager.futureLocalLogOrException(t1p0)).andStubReturn(futureLog)
+    expect(replicaManager.futureLogExists(t1p0)).andStubReturn(true)
 
-    expect(futureReplica.logEndOffset).andReturn(futureReplicaLEO).anyTimes()
     expect(replicaManager.logManager).andReturn(logManager).anyTimes()
 
     // pretend this is a completely new future replica, with no leader epochs recorded
-    expect(futureReplica.latestEpoch).andReturn(None).anyTimes()
+    expect(futureLog.latestEpoch).andReturn(None).anyTimes()
 
-    stubWithFetchMessages(replica, replica, futureReplica, partition, replicaManager, responseCallback)
-    replay(replicaManager, logManager, quotaManager, replica, futureReplica, partition)
+    stubWithFetchMessages(log, null, futureLog, partition, replicaManager, responseCallback)
+    replay(replicaManager, logManager, quotaManager, partition, log, futureLog)
 
     //Create the thread
     val endPoint = new BrokerEndPoint(0, "localhost", 1000)
@@ -330,6 +542,7 @@ class ReplicaAlterLogDirsThreadTest {
       "alter-logs-dirs-thread-test1",
       sourceBroker = endPoint,
       brokerConfig = config,
+      failedPartitions = failedPartitions,
       replicaMgr = replicaManager,
       quota = quotaManager,
       brokerTopicStats = null)
@@ -353,8 +566,8 @@ class ReplicaAlterLogDirsThreadTest {
     val config = KafkaConfig.fromProps(TestUtils.createBrokerConfig(1, "localhost:1234"))
     val quotaManager: ReplicationQuotaManager = createNiceMock(classOf[ReplicationQuotaManager])
     val logManager: LogManager = createMock(classOf[LogManager])
-    val replica: Replica = createNiceMock(classOf[Replica])
-    val futureReplica: Replica = createNiceMock(classOf[Replica])
+    val log: Log = createNiceMock(classOf[Log])
+    val futureLog: Log = createNiceMock(classOf[Log])
     val partition: Partition = createMock(classOf[Partition])
     val replicaManager: ReplicaManager = createMock(classOf[ReplicaManager])
     val responseCallback: Capture[Seq[(TopicPartition, FetchPartitionData)] => Unit]  = EasyMock.newCapture()
@@ -368,14 +581,13 @@ class ReplicaAlterLogDirsThreadTest {
       .andStubReturn(partition)
     expect(partition.truncateTo(capture(truncated), isFuture = EasyMock.eq(true))).once()
 
-    expect(replicaManager.futureLocalReplicaOrException(t1p0)).andStubReturn(futureReplica)
-    expect(futureReplica.latestEpoch).andStubReturn(Some(futureReplicaLeaderEpoch))
-    expect(futureReplica.endOffsetForEpoch(futureReplicaLeaderEpoch)).andReturn(
+    expect(replicaManager.futureLocalLogOrException(t1p0)).andStubReturn(futureLog)
+    expect(replicaManager.futureLogExists(t1p0)).andStubReturn(true)
+    expect(futureLog.logEndOffset).andReturn(futureReplicaLEO).anyTimes()
+    expect(futureLog.latestEpoch).andStubReturn(Some(futureReplicaLeaderEpoch))
+    expect(futureLog.endOffsetForEpoch(futureReplicaLeaderEpoch)).andReturn(
       Some(OffsetAndEpoch(futureReplicaLEO, futureReplicaLeaderEpoch)))
-    expect(futureReplica.logEndOffset).andReturn(futureReplicaLEO).anyTimes()
-    expect(replicaManager.localReplica(t1p0)).andReturn(Some(replica)).anyTimes()
-    expect(replicaManager.futureLocalReplica(t1p0)).andReturn(Some(futureReplica)).anyTimes()
-    expect(replicaManager.futureLocalReplicaOrException(t1p0)).andReturn(futureReplica).anyTimes()
+    expect(replicaManager.localLog(t1p0)).andReturn(Some(log)).anyTimes()
 
     // this will cause fetchEpochsFromLeader return an error with undefined offset
     expect(partition.lastOffsetForLeaderEpoch(Optional.of(1), futureReplicaLeaderEpoch, fetchOnlyFromLeader = false))
@@ -393,6 +605,7 @@ class ReplicaAlterLogDirsThreadTest {
       EasyMock.anyObject(),
       EasyMock.anyObject(),
       EasyMock.capture(responseCallback),
+      EasyMock.anyObject(),
       EasyMock.anyObject()))
       .andAnswer(new IAnswer[Unit] {
         override def answer(): Unit = {
@@ -400,7 +613,7 @@ class ReplicaAlterLogDirsThreadTest {
         }
       }).anyTimes()
 
-    replay(replicaManager, logManager, quotaManager, replica, futureReplica, partition)
+    replay(replicaManager, logManager, quotaManager, partition, log, futureLog)
 
     //Create the thread
     val endPoint = new BrokerEndPoint(0, "localhost", 1000)
@@ -408,6 +621,7 @@ class ReplicaAlterLogDirsThreadTest {
       "alter-logs-dirs-thread-test1",
       sourceBroker = endPoint,
       brokerConfig = config,
+      failedPartitions = failedPartitions,
       replicaMgr = replicaManager,
       quota = quotaManager,
       brokerTopicStats = null)
@@ -435,8 +649,8 @@ class ReplicaAlterLogDirsThreadTest {
     val config = KafkaConfig.fromProps(TestUtils.createBrokerConfig(1, "localhost:1234"))
     val quotaManager: ReplicationQuotaManager = createNiceMock(classOf[ReplicationQuotaManager])
     val logManager: LogManager = createMock(classOf[LogManager])
-    val replica: Replica = createNiceMock(classOf[Replica])
-    val futureReplica: Replica = createNiceMock(classOf[Replica])
+    val log: Log = createNiceMock(classOf[Log])
+    val futureLog: Log = createNiceMock(classOf[Log])
     val partition: Partition = createMock(classOf[Partition])
     val replicaManager: ReplicaManager = createMock(classOf[ReplicaManager])
     val responseCallback: Capture[Seq[(TopicPartition, FetchPartitionData)] => Unit]  = EasyMock.newCapture()
@@ -451,15 +665,16 @@ class ReplicaAlterLogDirsThreadTest {
         .andReturn(new EpochEndOffset(leaderEpoch, replicaLEO))
     expect(partition.truncateTo(futureReplicaLEO, isFuture = true)).once()
 
-    expect(replicaManager.futureLocalReplicaOrException(t1p0)).andStubReturn(futureReplica)
-    expect(futureReplica.latestEpoch).andStubReturn(Some(leaderEpoch))
-    expect(futureReplica.logEndOffset).andReturn(futureReplicaLEO).anyTimes()
-    expect(futureReplica.endOffsetForEpoch(leaderEpoch)).andReturn(
+    expect(replicaManager.futureLocalLogOrException(t1p0)).andStubReturn(futureLog)
+    expect(replicaManager.futureLogExists(t1p0)).andStubReturn(true)
+    expect(futureLog.latestEpoch).andStubReturn(Some(leaderEpoch))
+    expect(futureLog.logEndOffset).andStubReturn(futureReplicaLEO)
+    expect(futureLog.endOffsetForEpoch(leaderEpoch)).andReturn(
       Some(OffsetAndEpoch(futureReplicaLEO, leaderEpoch)))
     expect(replicaManager.logManager).andReturn(logManager).anyTimes()
-    stubWithFetchMessages(replica, replica, futureReplica, partition, replicaManager, responseCallback)
+    stubWithFetchMessages(log, null, futureLog, partition, replicaManager, responseCallback)
 
-    replay(replicaManager, logManager, quotaManager, replica, futureReplica, partition)
+    replay(replicaManager, logManager, quotaManager, partition, log, futureLog)
 
     //Create the fetcher thread
     val endPoint = new BrokerEndPoint(0, "localhost", 1000)
@@ -467,6 +682,7 @@ class ReplicaAlterLogDirsThreadTest {
       "alter-logs-dirs-thread-test1",
       sourceBroker = endPoint,
       brokerConfig = config,
+      failedPartitions = failedPartitions,
       replicaMgr = replicaManager,
       quota = quotaManager,
       brokerTopicStats = null)
@@ -488,17 +704,16 @@ class ReplicaAlterLogDirsThreadTest {
     val config = KafkaConfig.fromProps(TestUtils.createBrokerConfig(1, "localhost:1234"))
     val quotaManager: ReplicationQuotaManager = createNiceMock(classOf[ReplicationQuotaManager])
     val logManager: LogManager = createMock(classOf[LogManager])
-    val replica: Replica = createNiceMock(classOf[Replica])
-    val futureReplica: Replica = createNiceMock(classOf[Replica])
+    val log: Log = createNiceMock(classOf[Log])
+    val futureLog: Log = createNiceMock(classOf[Log])
     val partition: Partition = createMock(classOf[Partition])
     val replicaManager: ReplicaManager = createMock(classOf[ReplicaManager])
 
     //Stubs
-    expect(futureReplica.logStartOffset).andReturn(123).anyTimes()
     expect(replicaManager.logManager).andReturn(logManager).anyTimes()
-    stub(replica, replica, futureReplica, partition, replicaManager)
+    stub(log, null, futureLog, partition, replicaManager)
 
-    replay(replicaManager, logManager, quotaManager, replica, futureReplica, partition)
+    replay(replicaManager, logManager, quotaManager, partition, log)
 
     //Create the fetcher thread
     val endPoint = new BrokerEndPoint(0, "localhost", 1000)
@@ -507,6 +722,7 @@ class ReplicaAlterLogDirsThreadTest {
       "alter-logs-dirs-thread-test1",
       sourceBroker = endPoint,
       brokerConfig = config,
+      failedPartitions = failedPartitions,
       replicaMgr = replicaManager,
       quota = quotaManager,
       brokerTopicStats = null)
@@ -537,17 +753,18 @@ class ReplicaAlterLogDirsThreadTest {
     val config = KafkaConfig.fromProps(TestUtils.createBrokerConfig(1, "localhost:1234"))
     val quotaManager: ReplicationQuotaManager = createNiceMock(classOf[ReplicationQuotaManager])
     val logManager: LogManager = createMock(classOf[LogManager])
-    val replica: Replica = createNiceMock(classOf[Replica])
-    val futureReplica: Replica = createNiceMock(classOf[Replica])
+    val log: Log = createNiceMock(classOf[Log])
+    val futureLog: Log = createNiceMock(classOf[Log])
     val partition: Partition = createMock(classOf[Partition])
     val replicaManager: ReplicaManager = createMock(classOf[ReplicaManager])
 
     //Stubs
-    expect(futureReplica.logStartOffset).andReturn(123).anyTimes()
+    val startOffset = 123
+    expect(futureLog.logStartOffset).andReturn(startOffset).anyTimes()
     expect(replicaManager.logManager).andReturn(logManager).anyTimes()
-    stub(replica, replica, futureReplica, partition, replicaManager)
+    stub(log, null, futureLog, partition, replicaManager)
 
-    replay(replicaManager, logManager, quotaManager, replica, futureReplica, partition)
+    replay(replicaManager, logManager, quotaManager, partition, log, futureLog)
 
     //Create the fetcher thread
     val endPoint = new BrokerEndPoint(0, "localhost", 1000)
@@ -556,6 +773,7 @@ class ReplicaAlterLogDirsThreadTest {
       "alter-logs-dirs-thread-test1",
       sourceBroker = endPoint,
       brokerConfig = config,
+      failedPartitions = failedPartitions,
       replicaMgr = replicaManager,
       quota = quotaManager,
       brokerTopicStats = null)
@@ -599,23 +817,23 @@ class ReplicaAlterLogDirsThreadTest {
     assertFalse(partitionsWithError3.nonEmpty)
   }
 
-  def stub(replicaT1p0: Replica, replicaT1p1: Replica, futureReplica: Replica, partition: Partition, replicaManager: ReplicaManager) = {
-    expect(replicaManager.localReplica(t1p0)).andReturn(Some(replicaT1p0)).anyTimes()
-    expect(replicaManager.futureLocalReplica(t1p0)).andReturn(Some(futureReplica)).anyTimes()
-    expect(replicaManager.localReplicaOrException(t1p0)).andReturn(replicaT1p0).anyTimes()
-    expect(replicaManager.futureLocalReplicaOrException(t1p0)).andReturn(futureReplica).anyTimes()
-    expect(replicaManager.getPartition(t1p0)).andReturn(Some(partition)).anyTimes()
-    expect(replicaManager.localReplica(t1p1)).andReturn(Some(replicaT1p1)).anyTimes()
-    expect(replicaManager.futureLocalReplica(t1p1)).andReturn(Some(futureReplica)).anyTimes()
-    expect(replicaManager.localReplicaOrException(t1p1)).andReturn(replicaT1p1).anyTimes()
-    expect(replicaManager.futureLocalReplicaOrException(t1p1)).andReturn(futureReplica).anyTimes()
-    expect(replicaManager.getPartition(t1p1)).andReturn(Some(partition)).anyTimes()
+  def stub(logT1p0: Log, logT1p1: Log, futureLog: Log, partition: Partition,
+           replicaManager: ReplicaManager): IExpectationSetters[Option[Partition]] = {
+    expect(replicaManager.localLog(t1p0)).andReturn(Some(logT1p0)).anyTimes()
+    expect(replicaManager.localLogOrException(t1p0)).andReturn(logT1p0).anyTimes()
+    expect(replicaManager.futureLocalLogOrException(t1p0)).andReturn(futureLog).anyTimes()
+    expect(replicaManager.futureLogExists(t1p0)).andStubReturn(true)
+    expect(replicaManager.nonOfflinePartition(t1p0)).andReturn(Some(partition)).anyTimes()
+    expect(replicaManager.localLog(t1p1)).andReturn(Some(logT1p1)).anyTimes()
+    expect(replicaManager.localLogOrException(t1p1)).andReturn(logT1p1).anyTimes()
+    expect(replicaManager.futureLocalLogOrException(t1p1)).andReturn(futureLog).anyTimes()
+    expect(replicaManager.futureLogExists(t1p1)).andStubReturn(true)
+    expect(replicaManager.nonOfflinePartition(t1p1)).andReturn(Some(partition)).anyTimes()
   }
 
-  def stubWithFetchMessages(replicaT1p0: Replica, replicaT1p1: Replica, futureReplica: Replica,
-                            partition: Partition, replicaManager: ReplicaManager,
-                            responseCallback: Capture[Seq[(TopicPartition, FetchPartitionData)] => Unit]) = {
-    stub(replicaT1p0, replicaT1p1, futureReplica, partition, replicaManager)
+  def stubWithFetchMessages(logT1p0: Log, logT1p1: Log, futureLog: Log, partition: Partition, replicaManager: ReplicaManager,
+          responseCallback: Capture[Seq[(TopicPartition, FetchPartitionData)] => Unit]): IExpectationSetters[Unit] = {
+    stub(logT1p0, logT1p1, futureLog, partition, replicaManager)
     expect(replicaManager.fetchMessages(
       EasyMock.anyLong(),
       EasyMock.anyInt(),
@@ -625,6 +843,7 @@ class ReplicaAlterLogDirsThreadTest {
       EasyMock.anyObject(),
       EasyMock.anyObject(),
       EasyMock.capture(responseCallback),
+      EasyMock.anyObject(),
       EasyMock.anyObject()))
       .andAnswer(new IAnswer[Unit] {
         override def answer(): Unit = {
diff --git a/core/src/test/scala/unit/kafka/server/ReplicaFetchTest.scala b/core/src/test/scala/unit/kafka/server/ReplicaFetchTest.scala
index 8eba82471007e..f6f7ff129823e 100644
--- a/core/src/test/scala/unit/kafka/server/ReplicaFetchTest.scala
+++ b/core/src/test/scala/unit/kafka/server/ReplicaFetchTest.scala
@@ -17,6 +17,8 @@
 
 package kafka.server
 
+import scala.collection.Seq
+
 import org.junit.{After, Before, Test}
 import kafka.zk.ZooKeeperTestHarness
 import kafka.utils.TestUtils
@@ -31,20 +33,20 @@ class ReplicaFetchTest extends ZooKeeperTestHarness  {
   val topic2 = "bar"
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     super.setUp()
     val props = createBrokerConfigs(2, zkConnect)
     brokers = props.map(KafkaConfig.fromProps).map(TestUtils.createServer(_))
   }
 
   @After
-  override def tearDown() {
+  override def tearDown(): Unit = {
     TestUtils.shutdownServers(brokers)
     super.tearDown()
   }
 
   @Test
-  def testReplicaFetcherThread() {
+  def testReplicaFetcherThread(): Unit = {
     val partition = 0
     val testMessageList1 = List("test1", "test2", "test3", "test4")
     val testMessageList2 = List("test5", "test6", "test7", "test8")
diff --git a/core/src/test/scala/unit/kafka/server/ReplicaFetcherThreadTest.scala b/core/src/test/scala/unit/kafka/server/ReplicaFetcherThreadTest.scala
index 2a89a29484179..544c5991854d4 100644
--- a/core/src/test/scala/unit/kafka/server/ReplicaFetcherThreadTest.scala
+++ b/core/src/test/scala/unit/kafka/server/ReplicaFetcherThreadTest.scala
@@ -18,8 +18,8 @@ package kafka.server
 
 import java.util.Optional
 
-import kafka.cluster.{BrokerEndPoint, Partition, Replica}
-import kafka.log.LogManager
+import kafka.cluster.{BrokerEndPoint, Partition}
+import kafka.log.{Log, LogManager}
 import kafka.server.QuotaFactory.UnboundedQuota
 import kafka.server.epoch.util.ReplicaFetcherMockBlockingSend
 import kafka.utils.TestUtils
@@ -45,6 +45,7 @@ class ReplicaFetcherThreadTest {
   private val t2p1 = new TopicPartition("topic2", 1)
 
   private val brokerEndPoint = new BrokerEndPoint(0, "localhost", 1000)
+  private val failedPartitions = new FailedPartitions
 
   private def offsetAndEpoch(fetchOffset: Long, leaderEpoch: Int = 1): OffsetAndEpoch = {
     OffsetAndEpoch(offset = fetchOffset, leaderEpoch = leaderEpoch)
@@ -59,6 +60,7 @@ class ReplicaFetcherThreadTest {
       fetcherId = 0,
       sourceBroker = brokerEndPoint,
       brokerConfig = config,
+      failedPartitions: FailedPartitions,
       replicaMgr = null,
       metrics =  new Metrics(),
       time = new SystemTime(),
@@ -77,28 +79,29 @@ class ReplicaFetcherThreadTest {
     val quota: ReplicationQuotaManager = createNiceMock(classOf[ReplicationQuotaManager])
     val logManager: LogManager = createMock(classOf[LogManager])
     val replicaAlterLogDirsManager: ReplicaAlterLogDirsManager = createMock(classOf[ReplicaAlterLogDirsManager])
-    val replica: Replica = createNiceMock(classOf[Replica])
+    val log: Log = createNiceMock(classOf[Log])
     val partition: Partition = createMock(classOf[Partition])
     val replicaManager: ReplicaManager = createMock(classOf[ReplicaManager])
 
     val leaderEpoch = 5
 
     //Stubs
-    expect(replica.logEndOffset).andReturn(0).anyTimes()
-    expect(replica.highWatermark).andReturn(new LogOffsetMetadata(0)).anyTimes()
-    expect(replica.latestEpoch).andReturn(Some(leaderEpoch)).once()
-    expect(replica.latestEpoch).andReturn(Some(leaderEpoch)).once()
-    expect(replica.latestEpoch).andReturn(None).once()  // t2p1 doesnt support epochs
-    expect(replica.endOffsetForEpoch(leaderEpoch)).andReturn(
+    expect(partition.localLogOrException).andReturn(log).anyTimes()
+    expect(log.logEndOffset).andReturn(0).anyTimes()
+    expect(log.highWatermark).andReturn(0).anyTimes()
+    expect(log.latestEpoch).andReturn(Some(leaderEpoch)).once()
+    expect(log.latestEpoch).andReturn(Some(leaderEpoch)).once()
+    expect(log.latestEpoch).andReturn(None).once()  // t2p1 doesnt support epochs
+    expect(log.endOffsetForEpoch(leaderEpoch)).andReturn(
       Some(OffsetAndEpoch(0, leaderEpoch))).anyTimes()
     expect(replicaManager.logManager).andReturn(logManager).anyTimes()
     expect(replicaManager.replicaAlterLogDirsManager).andReturn(replicaAlterLogDirsManager).anyTimes()
-    stub(replica, partition, replicaManager)
+    stub(partition, replicaManager, log)
 
     //Expectations
-    expect(partition.truncateTo(anyLong(), anyBoolean())).once
+    expect(partition.truncateTo(anyLong(), anyBoolean())).times(3)
 
-    replay(replicaManager, logManager, quota, replica)
+    replay(replicaManager, logManager, quota, partition, log)
 
     //Define the offsets for the OffsetsForLeaderEpochResponse
     val offsets = Map(t1p0 -> new EpochEndOffset(leaderEpoch, 1),
@@ -107,7 +110,7 @@ class ReplicaFetcherThreadTest {
     //Create the fetcher thread
     val mockNetwork = new ReplicaFetcherMockBlockingSend(offsets, brokerEndPoint, new SystemTime())
 
-    val thread = new ReplicaFetcherThread("bob", 0, brokerEndPoint, config, replicaManager, new Metrics(), new SystemTime(), quota, Some(mockNetwork))
+    val thread = new ReplicaFetcherThread("bob", 0, brokerEndPoint, config, failedPartitions, replicaManager, new Metrics(), new SystemTime(), quota, Some(mockNetwork))
 
     // topic 1 supports epoch, t2 doesn't
     thread.addPartitions(Map(
@@ -181,6 +184,7 @@ class ReplicaFetcherThreadTest {
       fetcherId = 0,
       sourceBroker = brokerEndPoint,
       brokerConfig = config,
+      failedPartitions: FailedPartitions,
       replicaMgr = null,
       metrics =  new Metrics(),
       time = new SystemTime(),
@@ -207,33 +211,33 @@ class ReplicaFetcherThreadTest {
     //Setup all dependencies
     val logManager: LogManager = createMock(classOf[LogManager])
     val replicaAlterLogDirsManager: ReplicaAlterLogDirsManager = createMock(classOf[ReplicaAlterLogDirsManager])
-    val replica: Replica = createNiceMock(classOf[Replica])
+    val log: Log = createNiceMock(classOf[Log])
     val partition: Partition = createMock(classOf[Partition])
     val replicaManager: ReplicaManager = createMock(classOf[ReplicaManager])
 
     val leaderEpoch = 5
 
     //Stubs
-    expect(replica.logEndOffset).andReturn(0).anyTimes()
-    expect(replica.highWatermark).andReturn(new LogOffsetMetadata(0)).anyTimes()
-    expect(replica.latestEpoch).andReturn(Some(leaderEpoch)).anyTimes()
-    expect(replica.endOffsetForEpoch(leaderEpoch)).andReturn(
+    expect(partition.localLogOrException).andReturn(log).anyTimes()
+    expect(log.highWatermark).andReturn(0).anyTimes()
+    expect(log.latestEpoch).andReturn(Some(leaderEpoch)).anyTimes()
+    expect(log.endOffsetForEpoch(leaderEpoch)).andReturn(
       Some(OffsetAndEpoch(0, leaderEpoch))).anyTimes()
     expect(replicaManager.logManager).andReturn(logManager).anyTimes()
     expect(replicaManager.replicaAlterLogDirsManager).andReturn(replicaAlterLogDirsManager).anyTimes()
-    stub(replica, partition, replicaManager)
+    stub(partition, replicaManager, log)
 
     //Expectations
-    expect(partition.truncateTo(anyLong(), anyBoolean())).once
+    expect(partition.truncateTo(anyLong(), anyBoolean())).times(2)
 
-    replay(replicaManager, logManager, replica)
+    replay(replicaManager, logManager, partition, log)
 
     //Define the offsets for the OffsetsForLeaderEpochResponse
     val offsets = Map(t1p0 -> new EpochEndOffset(leaderEpoch, 1), t1p1 -> new EpochEndOffset(leaderEpoch, 1)).asJava
 
     //Create the fetcher thread
     val mockNetwork = new ReplicaFetcherMockBlockingSend(offsets, brokerEndPoint, new SystemTime())
-    val thread = new ReplicaFetcherThread("bob", 0, brokerEndPoint, config, replicaManager,
+    val thread = new ReplicaFetcherThread("bob", 0, brokerEndPoint, config, failedPartitions, replicaManager,
       new Metrics, new SystemTime, UnboundedQuota, Some(mockNetwork))
     thread.addPartitions(Map(t1p0 -> offsetAndEpoch(0L), t1p1 -> offsetAndEpoch(0L)))
 
@@ -267,7 +271,7 @@ class ReplicaFetcherThreadTest {
     val quota: ReplicationQuotaManager = createNiceMock(classOf[ReplicationQuotaManager])
     val logManager: LogManager = createMock(classOf[LogManager])
     val replicaAlterLogDirsManager: ReplicaAlterLogDirsManager = createMock(classOf[ReplicaAlterLogDirsManager])
-    val replica: Replica = createNiceMock(classOf[Replica])
+    val log: Log = createNiceMock(classOf[Log])
     val partition: Partition = createMock(classOf[Partition])
     val replicaManager: ReplicaManager = createMock(classOf[ReplicaManager])
 
@@ -276,23 +280,26 @@ class ReplicaFetcherThreadTest {
 
     //Stubs
     expect(partition.truncateTo(capture(truncateToCapture), anyBoolean())).anyTimes()
-    expect(replica.logEndOffset).andReturn(initialLEO).anyTimes()
-    expect(replica.highWatermark).andReturn(new LogOffsetMetadata(initialLEO - 1)).anyTimes()
-    expect(replica.latestEpoch).andReturn(Some(leaderEpoch)).anyTimes()
-    expect(replica.endOffsetForEpoch(leaderEpoch)).andReturn(
+    expect(partition.localLogOrException).andReturn(log).anyTimes()
+    expect(log.highWatermark).andReturn(initialLEO - 1).anyTimes()
+    expect(log.latestEpoch).andReturn(Some(leaderEpoch)).anyTimes()
+    expect(log.endOffsetForEpoch(leaderEpoch)).andReturn(
       Some(OffsetAndEpoch(initialLEO, leaderEpoch))).anyTimes()
+    expect(log.logEndOffset).andReturn(initialLEO).anyTimes()
+    expect(replicaManager.localLogOrException(anyObject(classOf[TopicPartition]))).andReturn(log).anyTimes()
     expect(replicaManager.logManager).andReturn(logManager).anyTimes()
     expect(replicaManager.replicaAlterLogDirsManager).andReturn(replicaAlterLogDirsManager).anyTimes()
-    stub(replica, partition, replicaManager)
+    stub(partition, replicaManager, log)
 
-    replay(replicaManager, logManager, quota, replica, partition)
+    replay(replicaManager, logManager, quota, partition, log)
 
     //Define the offsets for the OffsetsForLeaderEpochResponse, these are used for truncation
     val offsetsReply = Map(t1p0 -> new EpochEndOffset(leaderEpoch, 156), t2p1 -> new EpochEndOffset(leaderEpoch, 172)).asJava
 
     //Create the thread
     val mockNetwork = new ReplicaFetcherMockBlockingSend(offsetsReply, brokerEndPoint, new SystemTime())
-    val thread = new ReplicaFetcherThread("bob", 0, brokerEndPoint, configs(0), replicaManager, new Metrics(), new SystemTime(), quota, Some(mockNetwork))
+    val thread = new ReplicaFetcherThread("bob", 0, brokerEndPoint, configs.head, failedPartitions, replicaManager,
+      new Metrics(), new SystemTime(), quota, Some(mockNetwork))
     thread.addPartitions(Map(t1p0 -> offsetAndEpoch(0L), t2p1 -> offsetAndEpoch(0L)))
 
     //Run it
@@ -315,7 +322,7 @@ class ReplicaFetcherThreadTest {
     val quota: ReplicationQuotaManager = createNiceMock(classOf[ReplicationQuotaManager])
     val logManager: LogManager = createMock(classOf[LogManager])
     val replicaAlterLogDirsManager: ReplicaAlterLogDirsManager = createMock(classOf[ReplicaAlterLogDirsManager])
-    val replica: Replica = createNiceMock(classOf[Replica])
+    val log: Log = createNiceMock(classOf[Log])
     val partition: Partition = createMock(classOf[Partition])
     val replicaManager: ReplicaManager = createMock(classOf[ReplicaManager])
 
@@ -325,15 +332,17 @@ class ReplicaFetcherThreadTest {
 
     //Stubs
     expect(partition.truncateTo(capture(truncateToCapture), anyBoolean())).anyTimes()
-    expect(replica.logEndOffset).andReturn(initialLEO).anyTimes()
-    expect(replica.highWatermark).andReturn(new LogOffsetMetadata(initialLEO - 3)).anyTimes()
-    expect(replica.latestEpoch).andReturn(Some(leaderEpochAtFollower)).anyTimes()
-    expect(replica.endOffsetForEpoch(leaderEpochAtLeader)).andReturn(None).anyTimes()
+    expect(partition.localLogOrException).andReturn(log).anyTimes()
+    expect(log.highWatermark).andReturn(initialLEO - 3).anyTimes()
+    expect(log.latestEpoch).andReturn(Some(leaderEpochAtFollower)).anyTimes()
+    expect(log.endOffsetForEpoch(leaderEpochAtLeader)).andReturn(None).anyTimes()
+    expect(log.logEndOffset).andReturn(initialLEO).anyTimes()
+    expect(replicaManager.localLogOrException(anyObject(classOf[TopicPartition]))).andReturn(log).anyTimes()
     expect(replicaManager.logManager).andReturn(logManager).anyTimes()
     expect(replicaManager.replicaAlterLogDirsManager).andReturn(replicaAlterLogDirsManager).anyTimes()
-    stub(replica, partition, replicaManager)
+    stub(partition, replicaManager, log)
 
-    replay(replicaManager, logManager, quota, replica, partition)
+    replay(replicaManager, logManager, quota, partition, log)
 
     //Define the offsets for the OffsetsForLeaderEpochResponse, these are used for truncation
     val offsetsReply = Map(t1p0 -> new EpochEndOffset(leaderEpochAtLeader, 156),
@@ -341,7 +350,8 @@ class ReplicaFetcherThreadTest {
 
     //Create the thread
     val mockNetwork = new ReplicaFetcherMockBlockingSend(offsetsReply, brokerEndPoint, new SystemTime())
-    val thread = new ReplicaFetcherThread("bob", 0, brokerEndPoint, configs(0), replicaManager, new Metrics(), new SystemTime(), quota, Some(mockNetwork))
+    val thread = new ReplicaFetcherThread("bob", 0, brokerEndPoint, configs.head, failedPartitions,
+      replicaManager, new Metrics(), new SystemTime(), quota, Some(mockNetwork))
     thread.addPartitions(Map(t1p0 -> offsetAndEpoch(0L), t2p1 -> offsetAndEpoch(0L)))
 
     //Run it
@@ -367,7 +377,7 @@ class ReplicaFetcherThreadTest {
     val quota: ReplicationQuotaManager = createNiceMock(classOf[ReplicationQuotaManager])
     val logManager: LogManager = createMock(classOf[LogManager])
     val replicaAlterLogDirsManager: ReplicaAlterLogDirsManager = createMock(classOf[ReplicaAlterLogDirsManager])
-    val replica: Replica = createNiceMock(classOf[Replica])
+    val log: Log = createNiceMock(classOf[Log])
     val partition: Partition = createMock(classOf[Partition])
     val replicaManager: ReplicaManager = createMock(classOf[ReplicaManager])
 
@@ -375,25 +385,27 @@ class ReplicaFetcherThreadTest {
 
     // Stubs
     expect(partition.truncateTo(capture(truncateToCapture), anyBoolean())).anyTimes()
-    expect(replica.logEndOffset).andReturn(initialLEO).anyTimes()
-    expect(replica.highWatermark).andReturn(new LogOffsetMetadata(initialLEO - 2)).anyTimes()
-    expect(replica.latestEpoch).andReturn(Some(5)).anyTimes()
-    expect(replica.endOffsetForEpoch(4)).andReturn(
+    expect(partition.localLogOrException).andReturn(log).anyTimes()
+    expect(log.highWatermark).andReturn(initialLEO - 2).anyTimes()
+    expect(log.latestEpoch).andReturn(Some(5)).anyTimes()
+    expect(log.endOffsetForEpoch(4)).andReturn(
       Some(OffsetAndEpoch(120, 3))).anyTimes()
-    expect(replica.endOffsetForEpoch(3)).andReturn(
+    expect(log.endOffsetForEpoch(3)).andReturn(
       Some(OffsetAndEpoch(120, 3))).anyTimes()
+    expect(log.logEndOffset).andReturn(initialLEO).anyTimes()
+    expect(replicaManager.localLogOrException(anyObject(classOf[TopicPartition]))).andReturn(log).anyTimes()
     expect(replicaManager.logManager).andReturn(logManager).anyTimes()
     expect(replicaManager.replicaAlterLogDirsManager).andReturn(replicaAlterLogDirsManager).anyTimes()
-    stub(replica, partition, replicaManager)
+    stub(partition, replicaManager, log)
 
-    replay(replicaManager, logManager, quota, replica, partition)
+    replay(replicaManager, logManager, quota, partition, log)
 
     // Define the offsets for the OffsetsForLeaderEpochResponse
     val offsets = Map(t1p0 -> new EpochEndOffset(4, 155), t1p1 -> new EpochEndOffset(4, 143)).asJava
 
     // Create the fetcher thread
     val mockNetwork = new ReplicaFetcherMockBlockingSend(offsets, brokerEndPoint, new SystemTime())
-    val thread = new ReplicaFetcherThread("bob", 0, brokerEndPoint, config, replicaManager, new Metrics(), new SystemTime(), quota, Some(mockNetwork))
+    val thread = new ReplicaFetcherThread("bob", 0, brokerEndPoint, config, failedPartitions, replicaManager, new Metrics(), new SystemTime(), quota, Some(mockNetwork))
     thread.addPartitions(Map(t1p0 -> offsetAndEpoch(0L), t1p1 -> offsetAndEpoch(0L)))
 
     // Loop 1 -- both topic partitions will need to fetch another leader epoch
@@ -409,7 +421,7 @@ class ReplicaFetcherThreadTest {
     assertEquals(2, mockNetwork.epochFetchCount)
     assertEquals(1, mockNetwork.fetchCount)
     assertEquals("OffsetsForLeaderEpochRequest version.",
-      2, mockNetwork.lastUsedOffsetForLeaderEpochVersion)
+      3, mockNetwork.lastUsedOffsetForLeaderEpochVersion)
 
     //Loop 3 we should not fetch epochs
     thread.doWork()
@@ -438,7 +450,7 @@ class ReplicaFetcherThreadTest {
     val quota: ReplicationQuotaManager = createNiceMock(classOf[ReplicationQuotaManager])
     val logManager: LogManager = createMock(classOf[LogManager])
     val replicaAlterLogDirsManager: ReplicaAlterLogDirsManager = createMock(classOf[ReplicaAlterLogDirsManager])
-    val replica: Replica = createNiceMock(classOf[Replica])
+    val log: Log = createNiceMock(classOf[Log])
     val partition: Partition = createMock(classOf[Partition])
     val replicaManager: ReplicaManager = createMock(classOf[ReplicaManager])
 
@@ -446,18 +458,20 @@ class ReplicaFetcherThreadTest {
 
     // Stubs
     expect(partition.truncateTo(capture(truncateToCapture), anyBoolean())).anyTimes()
-    expect(replica.logEndOffset).andReturn(initialLEO).anyTimes()
-    expect(replica.highWatermark).andReturn(new LogOffsetMetadata(initialLEO - 2)).anyTimes()
-    expect(replica.latestEpoch).andReturn(Some(5)).anyTimes()
-    expect(replica.endOffsetForEpoch(4)).andReturn(
+    expect(partition.localLogOrException).andReturn(log).anyTimes()
+    expect(log.highWatermark).andReturn(initialLEO - 2).anyTimes()
+    expect(log.latestEpoch).andReturn(Some(5)).anyTimes()
+    expect(log.endOffsetForEpoch(4)).andReturn(
       Some(OffsetAndEpoch(120, 3))).anyTimes()
-    expect(replica.endOffsetForEpoch(3)).andReturn(
+    expect(log.endOffsetForEpoch(3)).andReturn(
       Some(OffsetAndEpoch(120, 3))).anyTimes()
+    expect(log.logEndOffset).andReturn(initialLEO).anyTimes()
+    expect(replicaManager.localLogOrException(anyObject(classOf[TopicPartition]))).andReturn(log).anyTimes()
     expect(replicaManager.logManager).andReturn(logManager).anyTimes()
     expect(replicaManager.replicaAlterLogDirsManager).andReturn(replicaAlterLogDirsManager).anyTimes()
-    stub(replica, partition, replicaManager)
+    stub(partition, replicaManager, log)
 
-    replay(replicaManager, logManager, quota, replica, partition)
+    replay(replicaManager, logManager, quota, partition, log)
 
     // Define the offsets for the OffsetsForLeaderEpochResponse with undefined epoch to simulate
     // older protocol version
@@ -465,7 +479,7 @@ class ReplicaFetcherThreadTest {
 
     // Create the fetcher thread
     val mockNetwork = new ReplicaFetcherMockBlockingSend(offsets, brokerEndPoint, new SystemTime())
-    val thread = new ReplicaFetcherThread("bob", 0, brokerEndPoint, config, replicaManager, new Metrics(), new SystemTime(), quota, Some(mockNetwork))
+    val thread = new ReplicaFetcherThread("bob", 0, brokerEndPoint, config, failedPartitions, replicaManager, new Metrics(), new SystemTime(), quota, Some(mockNetwork))
     thread.addPartitions(Map(t1p0 -> offsetAndEpoch(0L), t1p1 -> offsetAndEpoch(0L)))
 
     // Loop 1 -- both topic partitions will truncate to leader offset even though they don't know
@@ -499,29 +513,28 @@ class ReplicaFetcherThreadTest {
     val quota: ReplicationQuotaManager = createNiceMock(classOf[ReplicationQuotaManager])
     val logManager: LogManager = createMock(classOf[LogManager])
     val replicaAlterLogDirsManager: ReplicaAlterLogDirsManager = createMock(classOf[ReplicaAlterLogDirsManager])
-    val replica: Replica = createNiceMock(classOf[Replica])
+    val log: Log = createNiceMock(classOf[Log])
     val partition: Partition = createMock(classOf[Partition])
     val replicaManager: ReplicaManager = createMock(classOf[ReplicaManager])
 
     val initialFetchOffset = 100
-    val initialLeo = 300
 
     //Stubs
     expect(partition.truncateTo(capture(truncated), anyBoolean())).anyTimes()
-    expect(replica.logEndOffset).andReturn(initialLeo).anyTimes()
-    expect(replica.highWatermark).andReturn(new LogOffsetMetadata(initialFetchOffset)).anyTimes()
-    expect(replica.latestEpoch).andReturn(Some(5))
+    expect(partition.localLogOrException).andReturn(log).anyTimes()
+    expect(log.highWatermark).andReturn(initialFetchOffset).anyTimes()
+    expect(log.latestEpoch).andReturn(Some(5))
     expect(replicaManager.logManager).andReturn(logManager).anyTimes()
     expect(replicaManager.replicaAlterLogDirsManager).andReturn(replicaAlterLogDirsManager).anyTimes()
-    stub(replica, partition, replicaManager)
-    replay(replicaManager, logManager, quota, replica, partition)
+    stub(partition, replicaManager, log)
+    replay(replicaManager, logManager, quota, partition, log)
 
     //Define the offsets for the OffsetsForLeaderEpochResponse, these are used for truncation
     val offsetsReply = Map(t1p0 -> new EpochEndOffset(EpochEndOffset.UNDEFINED_EPOCH, EpochEndOffset.UNDEFINED_EPOCH_OFFSET)).asJava
 
     //Create the thread
     val mockNetwork = new ReplicaFetcherMockBlockingSend(offsetsReply, brokerEndPoint, new SystemTime())
-    val thread = new ReplicaFetcherThread("bob", 0, brokerEndPoint, configs(0), replicaManager, new Metrics(), new SystemTime(), quota, Some(mockNetwork))
+    val thread = new ReplicaFetcherThread("bob", 0, brokerEndPoint, configs.head, failedPartitions, replicaManager, new Metrics(), new SystemTime(), quota, Some(mockNetwork))
     thread.addPartitions(Map(t1p0 -> offsetAndEpoch(initialFetchOffset)))
 
     //Run it
@@ -542,7 +555,7 @@ class ReplicaFetcherThreadTest {
     val quota: ReplicationQuotaManager = createNiceMock(classOf[ReplicationQuotaManager])
     val logManager: LogManager = createMock(classOf[LogManager])
     val replicaAlterLogDirsManager: ReplicaAlterLogDirsManager = createMock(classOf[ReplicaAlterLogDirsManager])
-    val replica: Replica = createNiceMock(classOf[Replica])
+    val log: Log = createNiceMock(classOf[Log])
     val partition: Partition = createMock(classOf[Partition])
     val replicaManager: ReplicaManager = createMock(classOf[ReplicaManager])
 
@@ -551,17 +564,19 @@ class ReplicaFetcherThreadTest {
     val initialLeo = 300
 
     //Stubs
-    expect(replica.highWatermark).andReturn(new LogOffsetMetadata(highWaterMark)).anyTimes()
+    expect(log.highWatermark).andReturn(highWaterMark).anyTimes()
     expect(partition.truncateTo(capture(truncated), anyBoolean())).anyTimes()
-    expect(replica.logEndOffset).andReturn(initialLeo).anyTimes()
-    expect(replica.latestEpoch).andReturn(Some(leaderEpoch)).anyTimes()
+    expect(partition.localLogOrException).andReturn(log).anyTimes()
+    expect(log.latestEpoch).andReturn(Some(leaderEpoch)).anyTimes()
     // this is for the last reply with EpochEndOffset(5, 156)
-    expect(replica.endOffsetForEpoch(leaderEpoch)).andReturn(
+    expect(log.endOffsetForEpoch(leaderEpoch)).andReturn(
       Some(OffsetAndEpoch(initialLeo, leaderEpoch))).anyTimes()
+    expect(log.logEndOffset).andReturn(initialLeo).anyTimes()
+    expect(replicaManager.localLogOrException(anyObject(classOf[TopicPartition]))).andReturn(log).anyTimes()
     expect(replicaManager.logManager).andReturn(logManager).anyTimes()
     expect(replicaManager.replicaAlterLogDirsManager).andReturn(replicaAlterLogDirsManager).anyTimes()
-    stub(replica, partition, replicaManager)
-    replay(replicaManager, logManager, quota, replica, partition)
+    stub(partition, replicaManager, log)
+    replay(replicaManager, logManager, quota, partition, log)
 
     //Define the offsets for the OffsetsForLeaderEpochResponse, these are used for truncation
     val offsetsReply = mutable.Map(
@@ -571,7 +586,7 @@ class ReplicaFetcherThreadTest {
 
     //Create the thread
     val mockNetwork = new ReplicaFetcherMockBlockingSend(offsetsReply, brokerEndPoint, new SystemTime())
-    val thread = new ReplicaFetcherThread("bob", 0, brokerEndPoint, configs(0), replicaManager, new Metrics(), new SystemTime(), quota, Some(mockNetwork))
+    val thread = new ReplicaFetcherThread("bob", 0, brokerEndPoint, configs.head, failedPartitions, replicaManager, new Metrics(), new SystemTime(), quota, Some(mockNetwork))
     thread.addPartitions(Map(t1p0 -> offsetAndEpoch(0L), t1p1 -> offsetAndEpoch(0L)))
 
     //Run thread 3 times
@@ -599,23 +614,24 @@ class ReplicaFetcherThreadTest {
     val quota: ReplicationQuotaManager = createNiceMock(classOf[ReplicationQuotaManager])
     val logManager: LogManager = createNiceMock(classOf[LogManager])
     val replicaAlterLogDirsManager: ReplicaAlterLogDirsManager = createMock(classOf[ReplicaAlterLogDirsManager])
-    val replica: Replica = createNiceMock(classOf[Replica])
+    val log: Log = createNiceMock(classOf[Log])
     val partition: Partition = createMock(classOf[Partition])
     val replicaManager: ReplicaManager = createNiceMock(classOf[ReplicaManager])
 
     val leaderEpoch = 4
 
     //Stub return values
-    expect(replica.logEndOffset).andReturn(0).anyTimes()
-    expect(replica.highWatermark).andReturn(new LogOffsetMetadata(0)).anyTimes()
-    expect(replica.latestEpoch).andReturn(Some(leaderEpoch)).anyTimes()
-    expect(replica.endOffsetForEpoch(leaderEpoch)).andReturn(
+    expect(partition.truncateTo(0L, false)).times(2)
+    expect(partition.localLogOrException).andReturn(log).anyTimes()
+    expect(log.highWatermark).andReturn(0).anyTimes()
+    expect(log.latestEpoch).andReturn(Some(leaderEpoch)).anyTimes()
+    expect(log.endOffsetForEpoch(leaderEpoch)).andReturn(
       Some(OffsetAndEpoch(0, leaderEpoch))).anyTimes()
     expect(replicaManager.logManager).andReturn(logManager).anyTimes()
     expect(replicaManager.replicaAlterLogDirsManager).andReturn(replicaAlterLogDirsManager).anyTimes()
-    stub(replica, partition, replicaManager)
+    stub(partition, replicaManager, log)
 
-    replay(replicaManager, logManager, quota, replica)
+    replay(replicaManager, logManager, quota, partition, log)
 
     //Define the offsets for the OffsetsForLeaderEpochResponse
     val offsetsReply = Map(
@@ -624,7 +640,7 @@ class ReplicaFetcherThreadTest {
 
     //Create the fetcher thread
     val mockNetwork = new ReplicaFetcherMockBlockingSend(offsetsReply, brokerEndPoint, new SystemTime())
-    val thread = new ReplicaFetcherThread("bob", 0, brokerEndPoint, config, replicaManager, new Metrics(), new SystemTime(), quota, Some(mockNetwork))
+    val thread = new ReplicaFetcherThread("bob", 0, brokerEndPoint, config, failedPartitions, replicaManager, new Metrics(), new SystemTime(), quota, Some(mockNetwork))
 
     //When
     thread.addPartitions(Map(t1p0 -> offsetAndEpoch(0L), t1p1 -> offsetAndEpoch(0L)))
@@ -651,21 +667,23 @@ class ReplicaFetcherThreadTest {
     val quota: ReplicationQuotaManager = createNiceMock(classOf[ReplicationQuotaManager])
     val logManager: LogManager = createNiceMock(classOf[LogManager])
     val replicaAlterLogDirsManager: ReplicaAlterLogDirsManager = createMock(classOf[ReplicaAlterLogDirsManager])
-    val replica: Replica = createNiceMock(classOf[Replica])
+    val log: Log = createNiceMock(classOf[Log])
     val partition: Partition = createMock(classOf[Partition])
     val replicaManager: ReplicaManager = createNiceMock(classOf[ReplicaManager])
 
     //Stub return values
     expect(partition.truncateTo(capture(truncateToCapture), anyBoolean())).once
-    expect(replica.logEndOffset).andReturn(initialLEO).anyTimes()
-    expect(replica.highWatermark).andReturn(new LogOffsetMetadata(initialLEO - 2)).anyTimes()
-    expect(replica.latestEpoch).andReturn(Some(5)).anyTimes()
-    expect(replica.endOffsetForEpoch(5)).andReturn(Some(OffsetAndEpoch(initialLEO, 5))).anyTimes()
+    expect(partition.localLogOrException).andReturn(log).anyTimes()
+    expect(log.highWatermark).andReturn(initialLEO - 2).anyTimes()
+    expect(log.latestEpoch).andReturn(Some(5)).anyTimes()
+    expect(log.endOffsetForEpoch(5)).andReturn(Some(OffsetAndEpoch(initialLEO, 5))).anyTimes()
+    expect(log.logEndOffset).andReturn(initialLEO).anyTimes()
+    expect(replicaManager.localLogOrException(anyObject(classOf[TopicPartition]))).andReturn(log).anyTimes()
     expect(replicaManager.logManager).andReturn(logManager).anyTimes()
     expect(replicaManager.replicaAlterLogDirsManager).andReturn(replicaAlterLogDirsManager).anyTimes()
-    stub(replica, partition, replicaManager)
+    stub(partition, replicaManager, log)
 
-    replay(replicaManager, logManager, quota, replica, partition)
+    replay(replicaManager, logManager, quota, partition, log)
 
     //Define the offsets for the OffsetsForLeaderEpochResponse
     val offsetsReply = Map(
@@ -674,7 +692,8 @@ class ReplicaFetcherThreadTest {
 
     //Create the fetcher thread
     val mockNetwork = new ReplicaFetcherMockBlockingSend(offsetsReply, brokerEndPoint, new SystemTime())
-    val thread = new ReplicaFetcherThread("bob", 0, brokerEndPoint, config, replicaManager, new Metrics(), new SystemTime(), quota, Some(mockNetwork))
+    val thread = new ReplicaFetcherThread("bob", 0, brokerEndPoint, config, failedPartitions, replicaManager, new Metrics(),
+      new SystemTime(), quota, Some(mockNetwork))
 
     //When
     thread.addPartitions(Map(t1p0 -> offsetAndEpoch(0L), t1p1 -> offsetAndEpoch(0L)))
@@ -707,6 +726,7 @@ class ReplicaFetcherThreadTest {
       fetcherId = 0,
       sourceBroker = brokerEndPoint,
       brokerConfig = config,
+      failedPartitions = failedPartitions,
       replicaMgr = null,
       metrics =  new Metrics(),
       time = new SystemTime(),
@@ -723,15 +743,12 @@ class ReplicaFetcherThreadTest {
     verify(mockBlockingSend)
   }
 
-  def stub(replica: Replica, partition: Partition, replicaManager: ReplicaManager) = {
-    expect(replicaManager.localReplica(t1p0)).andReturn(Some(replica)).anyTimes()
-    expect(replicaManager.localReplicaOrException(t1p0)).andReturn(replica).anyTimes()
-    expect(replicaManager.getPartition(t1p0)).andReturn(Some(partition)).anyTimes()
-    expect(replicaManager.localReplica(t1p1)).andReturn(Some(replica)).anyTimes()
-    expect(replicaManager.localReplicaOrException(t1p1)).andReturn(replica).anyTimes()
-    expect(replicaManager.getPartition(t1p1)).andReturn(Some(partition)).anyTimes()
-    expect(replicaManager.localReplica(t2p1)).andReturn(Some(replica)).anyTimes()
-    expect(replicaManager.localReplicaOrException(t2p1)).andReturn(replica).anyTimes()
-    expect(replicaManager.getPartition(t2p1)).andReturn(Some(partition)).anyTimes()
+  def stub(partition: Partition, replicaManager: ReplicaManager, log: Log): Unit = {
+    expect(replicaManager.localLogOrException(t1p0)).andReturn(log).anyTimes()
+    expect(replicaManager.nonOfflinePartition(t1p0)).andReturn(Some(partition)).anyTimes()
+    expect(replicaManager.localLogOrException(t1p1)).andReturn(log).anyTimes()
+    expect(replicaManager.nonOfflinePartition(t1p1)).andReturn(Some(partition)).anyTimes()
+    expect(replicaManager.localLogOrException(t2p1)).andReturn(log).anyTimes()
+    expect(replicaManager.nonOfflinePartition(t2p1)).andReturn(Some(partition)).anyTimes()
   }
 }
diff --git a/core/src/test/scala/unit/kafka/server/ReplicaManagerQuotasTest.scala b/core/src/test/scala/unit/kafka/server/ReplicaManagerQuotasTest.scala
index 5b2f2aed77bef..e0dbf8cee4edd 100644
--- a/core/src/test/scala/unit/kafka/server/ReplicaManagerQuotasTest.scala
+++ b/core/src/test/scala/unit/kafka/server/ReplicaManagerQuotasTest.scala
@@ -20,7 +20,7 @@ import java.io.File
 import java.util.{Optional, Properties}
 import java.util.concurrent.atomic.AtomicBoolean
 
-import kafka.cluster.{Partition, Replica}
+import kafka.cluster.Partition
 import kafka.log.{Log, LogManager, LogOffsetSnapshot}
 import kafka.utils._
 import kafka.zk.KafkaZkClient
@@ -64,7 +64,8 @@ class ReplicaManagerQuotasTest {
       fetchMaxBytes = Int.MaxValue,
       hardMaxBytesLimit = false,
       readPartitionInfo = fetchInfo,
-      quota = quota)
+      quota = quota,
+      clientMetadata = None)
     assertEquals("Given two partitions, with only one throttled, we should get the first", 1,
       fetch.find(_._1 == topicPartition1).get._2.info.records.batches.asScala.size)
 
@@ -89,7 +90,8 @@ class ReplicaManagerQuotasTest {
       fetchMaxBytes = Int.MaxValue,
       hardMaxBytesLimit = false,
       readPartitionInfo = fetchInfo,
-      quota = quota)
+      quota = quota,
+      clientMetadata = None)
     assertEquals("Given two partitions, with both throttled, we should get no messages", 0,
       fetch.find(_._1 == topicPartition1).get._2.info.records.batches.asScala.size)
     assertEquals("Given two partitions, with both throttled, we should get no messages", 0,
@@ -113,7 +115,8 @@ class ReplicaManagerQuotasTest {
       fetchMaxBytes = Int.MaxValue,
       hardMaxBytesLimit = false,
       readPartitionInfo = fetchInfo,
-      quota = quota)
+      quota = quota,
+      clientMetadata = None)
     assertEquals("Given two partitions, with both non-throttled, we should get both messages", 1,
       fetch.find(_._1 == topicPartition1).get._2.info.records.batches.asScala.size)
     assertEquals("Given two partitions, with both non-throttled, we should get both messages", 1,
@@ -137,7 +140,8 @@ class ReplicaManagerQuotasTest {
       fetchMaxBytes = Int.MaxValue,
       hardMaxBytesLimit = false,
       readPartitionInfo = fetchInfo,
-      quota = quota)
+      quota = quota,
+      clientMetadata = None)
     assertEquals("Given two partitions, with only one throttled, we should get the first", 1,
       fetch.find(_._1 == topicPartition1).get._2.info.records.batches.asScala.size)
 
@@ -149,7 +153,7 @@ class ReplicaManagerQuotasTest {
   def testCompleteInDelayedFetchWithReplicaThrottling(): Unit = {
     // Set up DelayedFetch where there is data to return to a follower replica, either in-sync or out of sync
     def setupDelayedFetch(isReplicaInSync: Boolean): DelayedFetch = {
-      val endOffsetMetadata = new LogOffsetMetadata(messageOffset = 100L, segmentBaseOffset = 0L, relativePositionInSegment = 500)
+      val endOffsetMetadata = LogOffsetMetadata(messageOffset = 100L, segmentBaseOffset = 0L, relativePositionInSegment = 500)
       val partition: Partition = EasyMock.createMock(classOf[Partition])
 
       val offsetSnapshot = LogOffsetSnapshot(
@@ -167,10 +171,11 @@ class ReplicaManagerQuotasTest {
 
       EasyMock.expect(replicaManager.shouldLeaderThrottle(EasyMock.anyObject[ReplicaQuota], EasyMock.anyObject[TopicPartition], EasyMock.anyObject[Int]))
         .andReturn(!isReplicaInSync).anyTimes()
+      EasyMock.expect(partition.getReplica(1)).andReturn(None)
       EasyMock.replay(replicaManager, partition)
 
       val tp = new TopicPartition("t1", 0)
-      val fetchPartitionStatus = FetchPartitionStatus(new LogOffsetMetadata(messageOffset = 50L, segmentBaseOffset = 0L,
+      val fetchPartitionStatus = FetchPartitionStatus(LogOffsetMetadata(messageOffset = 50L, segmentBaseOffset = 0L,
          relativePositionInSegment = 250), new PartitionData(50, 0, 1, Optional.empty()))
       val fetchMetadata = FetchMetadata(fetchMinBytes = 1,
         fetchMaxBytes = 1000,
@@ -179,9 +184,10 @@ class ReplicaManagerQuotasTest {
         fetchIsolation = FetchLogEnd,
         isFromFollower = true,
         replicaId = 1,
-        fetchPartitionStatus = List((tp, fetchPartitionStatus)))
+        fetchPartitionStatus = List((tp, fetchPartitionStatus))
+      )
       new DelayedFetch(delayMs = 600, fetchMetadata = fetchMetadata, replicaManager = replicaManager,
-        quota = null, responseCallback = null) {
+        quota = null, clientMetadata = None, responseCallback = null) {
         override def forceComplete(): Boolean = true
       }
     }
@@ -190,7 +196,8 @@ class ReplicaManagerQuotasTest {
     assertFalse("Out of sync replica should not complete", setupDelayedFetch(isReplicaInSync = false).tryComplete())
   }
 
-  def setUpMocks(fetchInfo: Seq[(TopicPartition, PartitionData)], record: SimpleRecord = this.record, bothReplicasInSync: Boolean = false) {
+  def setUpMocks(fetchInfo: Seq[(TopicPartition, PartitionData)], record: SimpleRecord = this.record,
+                 bothReplicasInSync: Boolean = false): Unit = {
     val zkClient: KafkaZkClient = EasyMock.createMock(classOf[KafkaZkClient])
     val scheduler: KafkaScheduler = createNiceMock(classOf[KafkaScheduler])
 
@@ -198,27 +205,27 @@ class ReplicaManagerQuotasTest {
     val log: Log = createNiceMock(classOf[Log])
     expect(log.logStartOffset).andReturn(0L).anyTimes()
     expect(log.logEndOffset).andReturn(20L).anyTimes()
-    expect(log.logEndOffsetMetadata).andReturn(new LogOffsetMetadata(20L)).anyTimes()
+    expect(log.highWatermark).andReturn(5).anyTimes()
+    expect(log.lastStableOffset).andReturn(5).anyTimes()
+    expect(log.logEndOffsetMetadata).andReturn(LogOffsetMetadata(20L)).anyTimes()
 
     //if we ask for len 1 return a message
     expect(log.read(anyObject(),
       maxLength = geq(1),
-      maxOffset = anyObject(),
-      minOneMessage = anyBoolean(),
-      includeAbortedTxns = EasyMock.eq(false))).andReturn(
+      isolation = anyObject(),
+      minOneMessage = anyBoolean())).andReturn(
       FetchDataInfo(
-        new LogOffsetMetadata(0L, 0L, 0),
+        LogOffsetMetadata(0L, 0L, 0),
         MemoryRecords.withRecords(CompressionType.NONE, record)
       )).anyTimes()
 
     //if we ask for len = 0, return 0 messages
     expect(log.read(anyObject(),
       maxLength = EasyMock.eq(0),
-      maxOffset = anyObject(),
-      minOneMessage = anyBoolean(),
-      includeAbortedTxns = EasyMock.eq(false))).andReturn(
+      isolation = anyObject(),
+      minOneMessage = anyBoolean())).andReturn(
       FetchDataInfo(
-        new LogOffsetMetadata(0L, 0L, 0),
+        LogOffsetMetadata(0L, 0L, 0),
         MemoryRecords.EMPTY
       )).anyTimes()
     replay(log)
@@ -231,31 +238,27 @@ class ReplicaManagerQuotasTest {
     expect(logManager.liveLogDirs).andReturn(Array.empty[File]).anyTimes()
     replay(logManager)
 
+    val leaderBrokerId = configs.head.brokerId
     replicaManager = new ReplicaManager(configs.head, metrics, time, zkClient, scheduler, logManager,
       new AtomicBoolean(false), QuotaFactory.instantiate(configs.head, metrics, time, ""),
-      new BrokerTopicStats, new MetadataCache(configs.head.brokerId), new LogDirFailureChannel(configs.head.logDirs.size))
+      new BrokerTopicStats, new MetadataCache(leaderBrokerId), new LogDirFailureChannel(configs.head.logDirs.size))
 
     //create the two replicas
     for ((p, _) <- fetchInfo) {
-      val partition = replicaManager.getOrCreatePartition(p)
-      val leaderReplica = new Replica(configs.head.brokerId, p, time, 0, Some(log))
-      leaderReplica.highWatermark = new LogOffsetMetadata(5)
-      partition.leaderReplicaIdOpt = Some(leaderReplica.brokerId)
-      val followerReplica = new Replica(configs.last.brokerId, p, time, 0, Some(log))
-      val allReplicas = Set(leaderReplica, followerReplica)
-      allReplicas.foreach(partition.addReplicaIfNotExists)
-      if (bothReplicasInSync) {
-        partition.inSyncReplicas = allReplicas
-        followerReplica.highWatermark = new LogOffsetMetadata(5)
-      } else {
-        partition.inSyncReplicas = Set(leaderReplica)
-        followerReplica.highWatermark = new LogOffsetMetadata(0)
-      }
+      val partition = replicaManager.createPartition(p)
+      log.updateHighWatermark(5)
+      partition.leaderReplicaIdOpt = Some(leaderBrokerId)
+      partition.setLog(log, isFutureLog = false)
+
+      partition.updateAssignmentAndIsr(
+        assignment = Seq(leaderBrokerId, configs.last.brokerId),
+        isr = if (bothReplicasInSync) Set(leaderBrokerId, configs.last.brokerId) else Set(leaderBrokerId)
+      )
     }
   }
 
   @After
-  def tearDown() {
+  def tearDown(): Unit = {
     if (replicaManager != null)
       replicaManager.shutdown(false)
     metrics.close()
diff --git a/core/src/test/scala/unit/kafka/server/ReplicaManagerTest.scala b/core/src/test/scala/unit/kafka/server/ReplicaManagerTest.scala
index 65b62d0efa573..78d4782a7d553 100644
--- a/core/src/test/scala/unit/kafka/server/ReplicaManagerTest.scala
+++ b/core/src/test/scala/unit/kafka/server/ReplicaManagerTest.scala
@@ -18,41 +18,49 @@
 package kafka.server
 
 import java.io.File
-import java.util.{Optional, Properties}
+import java.net.InetAddress
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference}
 import java.util.concurrent.{CountDownLatch, TimeUnit}
-import java.util.concurrent.atomic.AtomicBoolean
+import java.util.{Optional, Properties}
 
-import kafka.log.{Log, LogConfig, LogManager, ProducerStateManager}
-import kafka.utils.{MockScheduler, MockTime, TestUtils}
-import TestUtils.createBroker
+import kafka.api.Request
+import kafka.log.{AppendOrigin, Log, LogConfig, LogManager, ProducerStateManager}
 import kafka.cluster.BrokerEndPoint
+import kafka.server.QuotaFactory.UnboundedQuota
+import kafka.server.checkpoints.LazyOffsetCheckpoints
 import kafka.server.epoch.util.ReplicaFetcherMockBlockingSend
+import kafka.utils.TestUtils.createBroker
 import kafka.utils.timer.MockTimer
-import kafka.zk.KafkaZkClient
-import org.I0Itec.zkclient.ZkClient
+import kafka.utils.{MockScheduler, MockTime, TestUtils}
+import kafka.zk.{IsrChangeNotificationSequenceZNode, KafkaZkClient}
+import org.apache.kafka.common.message.LeaderAndIsrRequestData.LeaderAndIsrPartitionState
 import org.apache.kafka.common.metrics.Metrics
 import org.apache.kafka.common.protocol.{ApiKeys, Errors}
 import org.apache.kafka.common.record._
-import org.apache.kafka.common.requests.{EpochEndOffset, IsolationLevel, LeaderAndIsrRequest}
-import org.apache.kafka.common.requests.ProduceResponse.PartitionResponse
+import org.apache.kafka.common.replica.ClientMetadata.DefaultClientMetadata
+import org.apache.kafka.common.replica.ClientMetadata
+import org.apache.kafka.common.requests.FetchRequest
 import org.apache.kafka.common.requests.FetchRequest.PartitionData
 import org.apache.kafka.common.requests.FetchResponse.AbortedTransaction
+import org.apache.kafka.common.requests.ProduceResponse.PartitionResponse
+import org.apache.kafka.common.requests.{EpochEndOffset, IsolationLevel, LeaderAndIsrRequest}
+import org.apache.kafka.common.security.auth.KafkaPrincipal
 import org.apache.kafka.common.utils.Time
+import org.apache.kafka.common.utils.Utils
 import org.apache.kafka.common.{Node, TopicPartition}
-import org.apache.zookeeper.data.Stat
-import org.easymock.EasyMock
+import org.easymock.{Capture, EasyMock}
 import org.junit.Assert._
-import org.junit.{After, Before, Test}
+import org.junit.{After, Before, Ignore, Test}
+import org.mockito.Mockito
 
 import scala.collection.JavaConverters._
-import scala.collection.Map
+import scala.collection.{Map, Seq}
 
 class ReplicaManagerTest {
 
   val topic = "test-topic"
   val time = new MockTime
   val metrics = new Metrics
-  var zkClient: ZkClient = _
   var kafkaZkClient: KafkaZkClient = _
 
   // Constants defined for readability
@@ -62,22 +70,57 @@ class ReplicaManagerTest {
   val brokerEpoch = 0L
 
   @Before
-  def setUp() {
-    zkClient = EasyMock.createMock(classOf[ZkClient])
+  def setUp(): Unit = {
     kafkaZkClient = EasyMock.createMock(classOf[KafkaZkClient])
     EasyMock.expect(kafkaZkClient.getEntityConfigs(EasyMock.anyString(), EasyMock.anyString())).andReturn(new Properties()).anyTimes()
     EasyMock.replay(kafkaZkClient)
-    EasyMock.expect(zkClient.readData(EasyMock.anyString(), EasyMock.anyObject[Stat])).andReturn(null).anyTimes()
-    EasyMock.replay(zkClient)
   }
 
   @After
-  def tearDown() {
+  def tearDown(): Unit = {
     metrics.close()
   }
 
   @Test
-  def testHighWaterMarkDirectoryMapping() {
+  def testBatchIsrChanges() : Unit = {
+    val time = new org.apache.kafka.common.utils.MockTime(ReplicaManager.IsrChangePropagationInterval)
+
+    kafkaZkClient = EasyMock.createMock(classOf[KafkaZkClient])
+    val capturedIsrChangeNotifications: Capture[collection.Set[TopicPartition]] = EasyMock.newCapture()
+    EasyMock.expect(kafkaZkClient.propagateIsrChanges(EasyMock.capture(capturedIsrChangeNotifications)))
+        .andAnswer(() => {
+          // Ensure that each serialized ISR change notification is well under the 1MiB limit in ZooKeeper.
+          val serializedSize = {
+            IsrChangeNotificationSequenceZNode.encode(capturedIsrChangeNotifications.getValue).length
+          }
+
+          // Just make sure the size never exceeds around 900KiB.
+          assertTrue(serializedSize < (900 * 1024));
+        }).times(2)
+    EasyMock.replay(kafkaZkClient)
+
+    val props = TestUtils.createBrokerConfig(1, TestUtils.MockZkConnect)
+    val config = KafkaConfig.fromProps(props)
+    val mockLogMgr = TestUtils.createLogManager(config.logDirs.map(new File(_)))
+    val rm = new ReplicaManager(config, metrics, time, kafkaZkClient, new MockScheduler(time), mockLogMgr,
+      new AtomicBoolean(false), QuotaFactory.instantiate(config, metrics, time, ""), new BrokerTopicStats,
+      new MetadataCache(config.brokerId), new LogDirFailureChannel(config.logDirs.size))
+
+    // Topic name which is the maximum length of 249.
+    val largeTopicName = "topic-".padTo(249, "x").mkString
+
+    // 5000 partitions. Large partition numbers chosen to make the serialized values as long as possible.
+    (10000 to 15000)
+      .map(n => new TopicPartition(largeTopicName, n))
+      .foreach(rm.recordIsrChange)
+
+    // Trigger sending of ISR notifications.
+    rm.maybePropagateIsrChanges()
+    EasyMock.verify(kafkaZkClient)
+  }
+
+  @Test
+  def testHighWaterMarkDirectoryMapping(): Unit = {
     val props = TestUtils.createBrokerConfig(1, TestUtils.MockZkConnect)
     val config = KafkaConfig.fromProps(props)
     val mockLogMgr = TestUtils.createLogManager(config.logDirs.map(new File(_)))
@@ -85,8 +128,9 @@ class ReplicaManagerTest {
       new AtomicBoolean(false), QuotaFactory.instantiate(config, metrics, time, ""), new BrokerTopicStats,
       new MetadataCache(config.brokerId), new LogDirFailureChannel(config.logDirs.size))
     try {
-      val partition = rm.getOrCreatePartition(new TopicPartition(topic, 1))
-      partition.getOrCreateReplica(1)
+      val partition = rm.createPartition(new TopicPartition(topic, 1))
+      partition.createLogIfNotExists(isNew = false, isFutureReplica = false,
+        new LazyOffsetCheckpoints(rm.highWatermarkCheckpoints))
       rm.checkpointHighWatermarks()
     } finally {
       // shutdown the replica manager upon test completion
@@ -95,7 +139,7 @@ class ReplicaManagerTest {
   }
 
   @Test
-  def testHighwaterMarkRelativeDirectoryMapping() {
+  def testHighwaterMarkRelativeDirectoryMapping(): Unit = {
     val props = TestUtils.createBrokerConfig(1, TestUtils.MockZkConnect)
     props.put("log.dir", TestUtils.tempRelativeDir("data").getAbsolutePath)
     val config = KafkaConfig.fromProps(props)
@@ -104,8 +148,9 @@ class ReplicaManagerTest {
       new AtomicBoolean(false), QuotaFactory.instantiate(config, metrics, time, ""), new BrokerTopicStats,
       new MetadataCache(config.brokerId), new LogDirFailureChannel(config.logDirs.size))
     try {
-      val partition = rm.getOrCreatePartition(new TopicPartition(topic, 1))
-      partition.getOrCreateReplica(1)
+      val partition = rm.createPartition(new TopicPartition(topic, 1))
+      partition.createLogIfNotExists(isNew = false, isFutureReplica = false,
+        new LazyOffsetCheckpoints(rm.highWatermarkCheckpoints))
       rm.checkpointHighWatermarks()
     } finally {
       // shutdown the replica manager upon test completion
@@ -114,7 +159,7 @@ class ReplicaManagerTest {
   }
 
   @Test
-  def testIllegalRequiredAcks() {
+  def testIllegalRequiredAcks(): Unit = {
     val props = TestUtils.createBrokerConfig(1, TestUtils.MockZkConnect)
     val config = KafkaConfig.fromProps(props)
     val mockLogMgr = TestUtils.createLogManager(config.logDirs.map(new File(_)))
@@ -129,7 +174,7 @@ class ReplicaManagerTest {
         timeout = 0,
         requiredAcks = 3,
         internalTopicsAllowed = false,
-        isFromClient = true,
+        origin = AppendOrigin.Client,
         entriesPerPartition = Map(new TopicPartition("test1", 0) -> MemoryRecords.withRecords(CompressionType.NONE,
           new SimpleRecord("first message".getBytes))),
         responseCallback = callback)
@@ -137,11 +182,11 @@ class ReplicaManagerTest {
       rm.shutdown(checkpointHW = false)
     }
 
-    TestUtils.verifyNonDaemonThreadsStatus(this.getClass.getName)
+    TestUtils.assertNoNonDaemonThreads(this.getClass.getName)
   }
 
   @Test
-  def testClearPurgatoryOnBecomingFollower() {
+  def testClearPurgatoryOnBecomingFollower(): Unit = {
     val props = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect)
     props.put("log.dir", TestUtils.tempRelativeDir("data").getAbsolutePath)
     val config = KafkaConfig.fromProps(props)
@@ -158,42 +203,115 @@ class ReplicaManagerTest {
     try {
       val brokerList = Seq[Integer](0, 1).asJava
 
-      val partition = rm.getOrCreatePartition(new TopicPartition(topic, 0))
-      partition.getOrCreateReplica(0)
+      val partition = rm.createPartition(new TopicPartition(topic, 0))
+      partition.createLogIfNotExists(isNew = false, isFutureReplica = false,
+        new LazyOffsetCheckpoints(rm.highWatermarkCheckpoints))
       // Make this replica the leader.
-      val leaderAndIsrRequest1 = new LeaderAndIsrRequest.Builder(ApiKeys.LEADER_AND_ISR.latestVersion, 0, 0, brokerEpoch,
-        collection.immutable.Map(new TopicPartition(topic, 0) ->
-          new LeaderAndIsrRequest.PartitionState(0, 0, 0, brokerList, 0, brokerList, false)).asJava,
+      val leaderAndIsrRequest1 = new LeaderAndIsrRequest.Builder(ApiKeys.LEADER_AND_ISR.latestVersion, 0, 0, brokerEpoch, brokerEpoch,
+        Seq(new LeaderAndIsrPartitionState()
+          .setTopicName(topic)
+          .setPartitionIndex(0)
+          .setControllerEpoch(0)
+          .setLeader(0)
+          .setLeaderEpoch(0)
+          .setIsr(brokerList)
+          .setZkVersion(0)
+          .setReplicas(brokerList)
+          .setIsNew(false)).asJava,
         Set(new Node(0, "host1", 0), new Node(1, "host2", 1)).asJava).build()
       rm.becomeLeaderOrFollower(0, leaderAndIsrRequest1, (_, _) => ())
       rm.getPartitionOrException(new TopicPartition(topic, 0), expectLeader = true)
-          .localReplicaOrException
+          .localLogOrException
 
       val records = MemoryRecords.withRecords(CompressionType.NONE, new SimpleRecord("first message".getBytes()))
       val appendResult = appendRecords(rm, new TopicPartition(topic, 0), records).onFire { response =>
         assertEquals(Errors.NOT_LEADER_FOR_PARTITION, response.error)
       }
 
-      // Fetch some messages
-      val fetchResult = fetchAsConsumer(rm, new TopicPartition(topic, 0),
-        new PartitionData(0, 0, 100000, Optional.empty()),
-        minBytes = 100000)
-      assertFalse(fetchResult.isFired)
-
       // Make this replica the follower
-      val leaderAndIsrRequest2 = new LeaderAndIsrRequest.Builder(ApiKeys.LEADER_AND_ISR.latestVersion, 0, 0, brokerEpoch,
-        collection.immutable.Map(new TopicPartition(topic, 0) ->
-          new LeaderAndIsrRequest.PartitionState(0, 1, 1, brokerList, 0, brokerList, false)).asJava,
+      val leaderAndIsrRequest2 = new LeaderAndIsrRequest.Builder(ApiKeys.LEADER_AND_ISR.latestVersion, 0, 0, brokerEpoch, brokerEpoch,
+        Seq(new LeaderAndIsrPartitionState()
+          .setTopicName(topic)
+          .setPartitionIndex(0)
+          .setControllerEpoch(0)
+          .setLeader(1)
+          .setLeaderEpoch(1)
+          .setIsr(brokerList)
+          .setZkVersion(0)
+          .setReplicas(brokerList)
+          .setIsNew(false)).asJava,
         Set(new Node(0, "host1", 0), new Node(1, "host2", 1)).asJava).build()
       rm.becomeLeaderOrFollower(1, leaderAndIsrRequest2, (_, _) => ())
 
       assertTrue(appendResult.isFired)
-      assertTrue(fetchResult.isFired)
     } finally {
       rm.shutdown(checkpointHW = false)
     }
   }
 
+  @Test
+  @Ignore // reenable once we upgrade to 2.5.0 KAFKA-9750
+  def testFencedErrorCausedByBecomeLeader(): Unit = {
+    testFencedErrorCausedByBecomeLeader(0)
+    testFencedErrorCausedByBecomeLeader(1)
+    testFencedErrorCausedByBecomeLeader(10)
+  }
+
+  private[this] def testFencedErrorCausedByBecomeLeader(loopEpochChange: Int): Unit = {
+    val replicaManager = setupReplicaManagerWithMockedPurgatories(new MockTimer)
+    try {
+      val brokerList = Seq[Integer](0, 1).asJava
+      val topicPartition = new TopicPartition(topic, 0)
+      replicaManager.createPartition(topicPartition)
+        .createLogIfNotExists(isNew = false, isFutureReplica = false,
+          new LazyOffsetCheckpoints(replicaManager.highWatermarkCheckpoints))
+
+      def leaderAndIsrRequest(epoch: Int): LeaderAndIsrRequest = new LeaderAndIsrRequest.Builder(ApiKeys.LEADER_AND_ISR.latestVersion, 0, 0, brokerEpoch, brokerEpoch,
+        Seq(new LeaderAndIsrPartitionState()
+          .setTopicName(topic)
+          .setPartitionIndex(0)
+          .setControllerEpoch(0)
+          .setLeader(0)
+          .setLeaderEpoch(epoch)
+          .setIsr(brokerList)
+          .setZkVersion(0)
+          .setReplicas(brokerList)
+          .setIsNew(true)).asJava,
+        Set(new Node(0, "host1", 0), new Node(1, "host2", 1)).asJava).build()
+
+      replicaManager.becomeLeaderOrFollower(0, leaderAndIsrRequest(0), (_, _) => ())
+      val partition = replicaManager.getPartitionOrException(new TopicPartition(topic, 0), expectLeader = true)
+      assertEquals(1, replicaManager.logManager.liveLogDirs.filterNot(_ == partition.log.get.dir.getParentFile).size)
+
+      val previousReplicaFolder = partition.log.get.dir.getParentFile
+      // find the live and different folder
+      val newReplicaFolder = replicaManager.logManager.liveLogDirs.filterNot(_ == partition.log.get.dir.getParentFile).head
+      assertEquals(0, replicaManager.replicaAlterLogDirsManager.fetcherThreadMap.size)
+      replicaManager.alterReplicaLogDirs(Map(topicPartition -> newReplicaFolder.getAbsolutePath))
+      // make sure the future log is created
+      replicaManager.futureLocalLogOrException(topicPartition)
+      assertEquals(1, replicaManager.replicaAlterLogDirsManager.fetcherThreadMap.size)
+      (1 to loopEpochChange).foreach(epoch => replicaManager.becomeLeaderOrFollower(0, leaderAndIsrRequest(epoch), (_, _) => ()))
+      // wait for the ReplicaAlterLogDirsThread to complete
+      TestUtils.waitUntilTrue(() => {
+        replicaManager.replicaAlterLogDirsManager.shutdownIdleFetcherThreads()
+        replicaManager.replicaAlterLogDirsManager.fetcherThreadMap.isEmpty
+      }, s"ReplicaAlterLogDirsThread should be gone")
+
+      // the fenced error should be recoverable
+      assertEquals(0, replicaManager.replicaAlterLogDirsManager.failedPartitions.size)
+      // the replica change is completed after retrying
+      assertTrue(partition.futureLog.isEmpty)
+      assertEquals(newReplicaFolder.getAbsolutePath, partition.log.get.dir.getParent)
+      // change the replica folder again
+      val response = replicaManager.alterReplicaLogDirs(Map(topicPartition -> previousReplicaFolder.getAbsolutePath))
+      assertNotEquals(0, response.size)
+      response.values.foreach(assertEquals(Errors.NONE, _))
+      // should succeed to invoke ReplicaAlterLogDirsThread again
+      assertEquals(1, replicaManager.replicaAlterLogDirsManager.fetcherThreadMap.size)
+    } finally replicaManager.shutdown(checkpointHW = false)
+  }
+
   @Test
   def testReceiveOutOfOrderSequenceExceptionWithLogStartOffset(): Unit = {
     val timer = new MockTimer
@@ -202,17 +320,26 @@ class ReplicaManagerTest {
     try {
       val brokerList = Seq[Integer](0, 1).asJava
 
-      val partition = replicaManager.getOrCreatePartition(new TopicPartition(topic, 0))
-      partition.getOrCreateReplica(0)
+      val partition = replicaManager.createPartition(new TopicPartition(topic, 0))
+      partition.createLogIfNotExists(isNew = false, isFutureReplica = false,
+        new LazyOffsetCheckpoints(replicaManager.highWatermarkCheckpoints))
 
       // Make this replica the leader.
-      val leaderAndIsrRequest1 = new LeaderAndIsrRequest.Builder(ApiKeys.LEADER_AND_ISR.latestVersion, 0, 0, brokerEpoch,
-        collection.immutable.Map(new TopicPartition(topic, 0) ->
-          new LeaderAndIsrRequest.PartitionState(0, 0, 0, brokerList, 0, brokerList, true)).asJava,
+      val leaderAndIsrRequest1 = new LeaderAndIsrRequest.Builder(ApiKeys.LEADER_AND_ISR.latestVersion, 0, 0, brokerEpoch, brokerEpoch,
+        Seq(new LeaderAndIsrPartitionState()
+          .setTopicName(topic)
+          .setPartitionIndex(0)
+          .setControllerEpoch(0)
+          .setLeader(0)
+          .setLeaderEpoch(0)
+          .setIsr(brokerList)
+          .setZkVersion(0)
+          .setReplicas(brokerList)
+          .setIsNew(true)).asJava,
         Set(new Node(0, "host1", 0), new Node(1, "host2", 1)).asJava).build()
       replicaManager.becomeLeaderOrFollower(0, leaderAndIsrRequest1, (_, _) => ())
       replicaManager.getPartitionOrException(new TopicPartition(topic, 0), expectLeader = true)
-        .localReplicaOrException
+        .localLogOrException
 
       val producerId = 234L
       val epoch = 5.toShort
@@ -253,17 +380,26 @@ class ReplicaManagerTest {
     try {
       val brokerList = Seq[Integer](0, 1).asJava
 
-      val partition = replicaManager.getOrCreatePartition(new TopicPartition(topic, 0))
-      partition.getOrCreateReplica(0)
+      val partition = replicaManager.createPartition(new TopicPartition(topic, 0))
+      partition.createLogIfNotExists(isNew = false, isFutureReplica = false,
+        new LazyOffsetCheckpoints(replicaManager.highWatermarkCheckpoints))
 
       // Make this replica the leader.
-      val leaderAndIsrRequest1 = new LeaderAndIsrRequest.Builder(ApiKeys.LEADER_AND_ISR.latestVersion, 0, 0, brokerEpoch,
-        collection.immutable.Map(new TopicPartition(topic, 0) ->
-          new LeaderAndIsrRequest.PartitionState(0, 0, 0, brokerList, 0, brokerList, true)).asJava,
+      val leaderAndIsrRequest1 = new LeaderAndIsrRequest.Builder(ApiKeys.LEADER_AND_ISR.latestVersion, 0, 0, brokerEpoch, brokerEpoch,
+        Seq(new LeaderAndIsrPartitionState()
+          .setTopicName(topic)
+          .setPartitionIndex(0)
+          .setControllerEpoch(0)
+          .setLeader(0)
+          .setLeaderEpoch(0)
+          .setIsr(brokerList)
+          .setZkVersion(0)
+          .setReplicas(brokerList)
+          .setIsNew(true)).asJava,
         Set(new Node(0, "host1", 0), new Node(1, "host2", 1)).asJava).build()
       replicaManager.becomeLeaderOrFollower(0, leaderAndIsrRequest1, (_, _) => ())
       replicaManager.getPartitionOrException(new TopicPartition(topic, 0), expectLeader = true)
-        .localReplicaOrException
+        .localLogOrException
 
       val producerId = 234L
       val epoch = 5.toShort
@@ -309,7 +445,8 @@ class ReplicaManagerTest {
       // now commit the transaction
       val endTxnMarker = new EndTransactionMarker(ControlRecordType.COMMIT, 0)
       val commitRecordBatch = MemoryRecords.withEndTransactionMarker(producerId, epoch, endTxnMarker)
-      appendRecords(replicaManager, new TopicPartition(topic, 0), commitRecordBatch, isFromClient = false)
+      appendRecords(replicaManager, new TopicPartition(topic, 0), commitRecordBatch,
+        origin = AppendOrigin.Coordinator)
         .onFire { response => assertEquals(Errors.NONE, response.error) }
 
       // the LSO has advanced, but the appended commit marker has not been replicated, so
@@ -349,17 +486,26 @@ class ReplicaManagerTest {
 
     try {
       val brokerList = Seq[Integer](0, 1).asJava
-      val partition = replicaManager.getOrCreatePartition(new TopicPartition(topic, 0))
-      partition.getOrCreateReplica(0)
+      val partition = replicaManager.createPartition(new TopicPartition(topic, 0))
+      partition.createLogIfNotExists(isNew = false, isFutureReplica = false,
+        new LazyOffsetCheckpoints(replicaManager.highWatermarkCheckpoints))
 
       // Make this replica the leader.
-      val leaderAndIsrRequest1 = new LeaderAndIsrRequest.Builder(ApiKeys.LEADER_AND_ISR.latestVersion, 0, 0, brokerEpoch,
-        collection.immutable.Map(new TopicPartition(topic, 0) ->
-          new LeaderAndIsrRequest.PartitionState(0, 0, 0, brokerList, 0, brokerList, true)).asJava,
+      val leaderAndIsrRequest1 = new LeaderAndIsrRequest.Builder(ApiKeys.LEADER_AND_ISR.latestVersion, 0, 0, brokerEpoch, brokerEpoch,
+        Seq(new LeaderAndIsrPartitionState()
+          .setTopicName(topic)
+          .setPartitionIndex(0)
+          .setControllerEpoch(0)
+          .setLeader(0)
+          .setLeaderEpoch(0)
+          .setIsr(brokerList)
+          .setZkVersion(0)
+          .setReplicas(brokerList)
+          .setIsNew(true)).asJava,
         Set(new Node(0, "host1", 0), new Node(1, "host2", 1)).asJava).build()
       replicaManager.becomeLeaderOrFollower(0, leaderAndIsrRequest1, (_, _) => ())
       replicaManager.getPartitionOrException(new TopicPartition(topic, 0), expectLeader = true)
-        .localReplicaOrException
+        .localLogOrException
 
       val producerId = 234L
       val epoch = 5.toShort
@@ -377,7 +523,8 @@ class ReplicaManagerTest {
       // now abort the transaction
       val endTxnMarker = new EndTransactionMarker(ControlRecordType.ABORT, 0)
       val abortRecordBatch = MemoryRecords.withEndTransactionMarker(producerId, epoch, endTxnMarker)
-      appendRecords(replicaManager, new TopicPartition(topic, 0), abortRecordBatch, isFromClient = false)
+      appendRecords(replicaManager, new TopicPartition(topic, 0), abortRecordBatch,
+        origin = AppendOrigin.Coordinator)
         .onFire { response => assertEquals(Errors.NONE, response.error) }
 
       // fetch as follower to advance the high watermark
@@ -410,22 +557,31 @@ class ReplicaManagerTest {
   }
 
   @Test
-  def testFetchBeyondHighWatermarkReturnEmptyResponse() {
+  def testFetchBeyondHighWatermarkReturnEmptyResponse(): Unit = {
     val rm = setupReplicaManagerWithMockedPurgatories(new MockTimer, aliveBrokerIds = Seq(0, 1, 2))
     try {
       val brokerList = Seq[Integer](0, 1, 2).asJava
 
-      val partition = rm.getOrCreatePartition(new TopicPartition(topic, 0))
-      partition.getOrCreateReplica(0)
+      val partition = rm.createPartition(new TopicPartition(topic, 0))
+      partition.createLogIfNotExists(isNew = false, isFutureReplica = false,
+        new LazyOffsetCheckpoints(rm.highWatermarkCheckpoints))
 
       // Make this replica the leader.
-      val leaderAndIsrRequest1 = new LeaderAndIsrRequest.Builder(ApiKeys.LEADER_AND_ISR.latestVersion, 0, 0, brokerEpoch,
-        collection.immutable.Map(new TopicPartition(topic, 0) ->
-          new LeaderAndIsrRequest.PartitionState(0, 0, 0, brokerList, 0, brokerList, false)).asJava,
+      val leaderAndIsrRequest1 = new LeaderAndIsrRequest.Builder(ApiKeys.LEADER_AND_ISR.latestVersion, 0, 0, brokerEpoch, brokerEpoch,
+        Seq(new LeaderAndIsrPartitionState()
+          .setTopicName(topic)
+          .setPartitionIndex(0)
+          .setControllerEpoch(0)
+          .setLeader(0)
+          .setLeaderEpoch(0)
+          .setIsr(brokerList)
+          .setZkVersion(0)
+          .setReplicas(brokerList)
+          .setIsNew(false)).asJava,
         Set(new Node(0, "host1", 0), new Node(1, "host2", 1), new Node(2, "host2", 2)).asJava).build()
       rm.becomeLeaderOrFollower(0, leaderAndIsrRequest1, (_, _) => ())
       rm.getPartitionOrException(new TopicPartition(topic, 0), expectLeader = true)
-        .localReplicaOrException
+        .localLogOrException
 
       // Append a couple of messages.
       for(i <- 1 to 2) {
@@ -453,26 +609,142 @@ class ReplicaManagerTest {
     }
   }
 
+  @Test
+  def testFollowerStateNotUpdatedIfLogReadFails(): Unit = {
+    val maxFetchBytes = 1024 * 1024
+    val aliveBrokersIds = Seq(0, 1)
+    val leaderEpoch = 5
+    val replicaManager = setupReplicaManagerWithMockedPurgatories(new MockTimer, aliveBrokersIds)
+    try {
+      val tp = new TopicPartition(topic, 0)
+      val replicas = aliveBrokersIds.toList.map(Int.box).asJava
+
+      // Broker 0 becomes leader of the partition
+      val leaderAndIsrPartitionState = new LeaderAndIsrPartitionState()
+        .setTopicName(topic)
+        .setPartitionIndex(0)
+        .setControllerEpoch(0)
+        .setLeader(0)
+        .setLeaderEpoch(leaderEpoch)
+        .setIsr(replicas)
+        .setZkVersion(0)
+        .setReplicas(replicas)
+        .setIsNew(true)
+      val leaderAndIsrRequest = new LeaderAndIsrRequest.Builder(ApiKeys.LEADER_AND_ISR.latestVersion, 0, 0, brokerEpoch, brokerEpoch,
+        Seq(leaderAndIsrPartitionState).asJava,
+        Set(new Node(0, "host1", 0), new Node(1, "host2", 1)).asJava).build()
+      val leaderAndIsrResponse = replicaManager.becomeLeaderOrFollower(0, leaderAndIsrRequest, (_, _) => ())
+      assertEquals(Errors.NONE, leaderAndIsrResponse.error)
+
+      // Follower replica state is initialized, but initial state is not known
+      assertTrue(replicaManager.nonOfflinePartition(tp).isDefined)
+      val partition = replicaManager.nonOfflinePartition(tp).get
+
+      assertTrue(partition.getReplica(1).isDefined)
+      val followerReplica = partition.getReplica(1).get
+      assertEquals(-1L, followerReplica.logStartOffset)
+      assertEquals(-1L, followerReplica.logEndOffset)
+
+      // Leader appends some data
+      for (i <- 1 to 5) {
+        appendRecords(replicaManager, tp, TestUtils.singletonRecords(s"message $i".getBytes)).onFire { response =>
+          assertEquals(Errors.NONE, response.error)
+        }
+      }
+
+      // We receive one valid request from the follower and replica state is updated
+      var successfulFetch: Option[FetchPartitionData] = None
+      def callback(response: Seq[(TopicPartition, FetchPartitionData)]): Unit = {
+        successfulFetch = response.headOption.filter(_._1 == tp).map(_._2)
+      }
+
+      val validFetchPartitionData = new FetchRequest.PartitionData(0L, 0L, maxFetchBytes,
+        Optional.of(leaderEpoch))
+
+      replicaManager.fetchMessages(
+        timeout = 0L,
+        replicaId = 1,
+        fetchMinBytes = 1,
+        fetchMaxBytes = maxFetchBytes,
+        hardMaxBytesLimit = false,
+        fetchInfos = Seq(tp -> validFetchPartitionData),
+        quota = UnboundedQuota,
+        isolationLevel = IsolationLevel.READ_UNCOMMITTED,
+        responseCallback = callback,
+        clientMetadata = None
+      )
+
+      assertTrue(successfulFetch.isDefined)
+      assertEquals(0L, followerReplica.logStartOffset)
+      assertEquals(0L, followerReplica.logEndOffset)
+
+
+      // Next we receive an invalid request with a higher fetch offset, but an old epoch.
+      // We expect that the replica state does not get updated.
+      val invalidFetchPartitionData = new FetchRequest.PartitionData(3L, 0L, maxFetchBytes,
+        Optional.of(leaderEpoch - 1))
+
+      replicaManager.fetchMessages(
+        timeout = 0L,
+        replicaId = 1,
+        fetchMinBytes = 1,
+        fetchMaxBytes = maxFetchBytes,
+        hardMaxBytesLimit = false,
+        fetchInfos = Seq(tp -> invalidFetchPartitionData),
+        quota = UnboundedQuota,
+        isolationLevel = IsolationLevel.READ_UNCOMMITTED,
+        responseCallback = callback,
+        clientMetadata = None
+      )
+
+      assertTrue(successfulFetch.isDefined)
+      assertEquals(0L, followerReplica.logStartOffset)
+      assertEquals(0L, followerReplica.logEndOffset)
+
+    } finally {
+      replicaManager.shutdown(checkpointHW = false)
+    }
+  }
+
   /**
    * If a follower sends a fetch request for 2 partitions and it's no longer the follower for one of them, the other
    * partition should not be affected.
    */
   @Test
-  def testFetchMessagesWhenNotFollowerForOnePartition() {
+  def testFetchMessagesWhenNotFollowerForOnePartition(): Unit = {
     val replicaManager = setupReplicaManagerWithMockedPurgatories(new MockTimer, aliveBrokerIds = Seq(0, 1, 2))
 
     try {
       // Create 2 partitions, assign replica 0 as the leader for both a different follower (1 and 2) for each
       val tp0 = new TopicPartition(topic, 0)
       val tp1 = new TopicPartition(topic, 1)
-      replicaManager.getOrCreatePartition(tp0).getOrCreateReplica(0)
-      replicaManager.getOrCreatePartition(tp1).getOrCreateReplica(0)
+      val offsetCheckpoints = new LazyOffsetCheckpoints(replicaManager.highWatermarkCheckpoints)
+      replicaManager.createPartition(tp0).createLogIfNotExists(isNew = false, isFutureReplica = false, offsetCheckpoints)
+      replicaManager.createPartition(tp1).createLogIfNotExists(isNew = false, isFutureReplica = false, offsetCheckpoints)
       val partition0Replicas = Seq[Integer](0, 1).asJava
       val partition1Replicas = Seq[Integer](0, 2).asJava
-      val leaderAndIsrRequest = new LeaderAndIsrRequest.Builder(ApiKeys.LEADER_AND_ISR.latestVersion, 0, 0, brokerEpoch,
-        collection.immutable.Map(
-          tp0 -> new LeaderAndIsrRequest.PartitionState(0, 0, 0, partition0Replicas, 0, partition0Replicas, true),
-          tp1 -> new LeaderAndIsrRequest.PartitionState(0, 0, 0, partition1Replicas, 0, partition1Replicas, true)
+      val leaderAndIsrRequest = new LeaderAndIsrRequest.Builder(ApiKeys.LEADER_AND_ISR.latestVersion, 0, 0, brokerEpoch, brokerEpoch,
+        Seq(
+          new LeaderAndIsrPartitionState()
+            .setTopicName(tp0.topic)
+            .setPartitionIndex(tp0.partition)
+            .setControllerEpoch(0)
+            .setLeader(0)
+            .setLeaderEpoch(0)
+            .setIsr(partition0Replicas)
+            .setZkVersion(0)
+            .setReplicas(partition0Replicas)
+            .setIsNew(true),
+          new LeaderAndIsrPartitionState()
+            .setTopicName(tp1.topic)
+            .setPartitionIndex(tp1.partition)
+            .setControllerEpoch(0)
+            .setLeader(0)
+            .setLeaderEpoch(0)
+            .setIsr(partition1Replicas)
+            .setZkVersion(0)
+            .setReplicas(partition1Replicas)
+            .setIsNew(true)
         ).asJava,
         Set(new Node(0, "host1", 0), new Node(1, "host2", 1)).asJava).build()
       replicaManager.becomeLeaderOrFollower(0, leaderAndIsrRequest, (_, _) => ())
@@ -518,17 +790,19 @@ class ReplicaManagerTest {
         fetchInfos = Seq(
           tp0 -> new PartitionData(1, 0, 100000, Optional.empty()),
           tp1 -> new PartitionData(1, 0, 100000, Optional.empty())),
+        quota = UnboundedQuota,
         responseCallback = fetchCallback,
-        isolationLevel = IsolationLevel.READ_UNCOMMITTED
+        isolationLevel = IsolationLevel.READ_UNCOMMITTED,
+        clientMetadata = None
       )
-      val tp0Replica = replicaManager.localReplica(tp0)
-      assertTrue(tp0Replica.isDefined)
-      assertEquals("hw should be incremented", 1, tp0Replica.get.highWatermark.messageOffset)
+      val tp0Log = replicaManager.localLog(tp0)
+      assertTrue(tp0Log.isDefined)
+      assertEquals("hw should be incremented", 1, tp0Log.get.highWatermark)
 
-      replicaManager.localReplica(tp1)
-      val tp1Replica = replicaManager.localReplica(tp1)
+      replicaManager.localLog(tp1)
+      val tp1Replica = replicaManager.localLog(tp1)
       assertTrue(tp1Replica.isDefined)
-      assertEquals("hw should not be incremented", 0, tp1Replica.get.highWatermark.messageOffset)
+      assertEquals("hw should not be incremented", 0, tp1Replica.get.highWatermark)
 
     } finally {
       replicaManager.shutdown(checkpointHW = false)
@@ -540,7 +814,7 @@ class ReplicaManagerTest {
     * if the epoch has increased by more than one (which suggests it has missed an update)
     */
   @Test
-  def testBecomeFollowerWhenLeaderIsUnchangedButMissedLeaderUpdate() {
+  def testBecomeFollowerWhenLeaderIsUnchangedButMissedLeaderUpdate(): Unit = {
     val topicPartition = 0
     val followerBrokerId = 0
     val leaderBrokerId = 1
@@ -556,19 +830,20 @@ class ReplicaManagerTest {
       topicPartition, leaderEpoch + leaderEpochIncrement, followerBrokerId, leaderBrokerId, countDownLatch, expectTruncation = true)
 
     // Initialize partition state to follower, with leader = 1, leaderEpoch = 1
-    val partition = replicaManager.getOrCreatePartition(new TopicPartition(topic, topicPartition))
-    partition.getOrCreateReplica(followerBrokerId)
-    partition.makeFollower(controllerId,
-      leaderAndIsrPartitionState(leaderEpoch, leaderBrokerId, aliveBrokerIds),
-      correlationId)
+    val tp = new TopicPartition(topic, topicPartition)
+    val partition = replicaManager.createPartition(tp)
+    val offsetCheckpoints = new LazyOffsetCheckpoints(replicaManager.highWatermarkCheckpoints)
+    partition.createLogIfNotExists(isNew = false, isFutureReplica = false, offsetCheckpoints)
+    partition.makeFollower(
+      leaderAndIsrPartitionState(tp, leaderEpoch, leaderBrokerId, aliveBrokerIds),
+      offsetCheckpoints)
 
     // Make local partition a follower - because epoch increased by more than 1, truncation should
     // trigger even though leader does not change
     leaderEpoch += leaderEpochIncrement
     val leaderAndIsrRequest0 = new LeaderAndIsrRequest.Builder(ApiKeys.LEADER_AND_ISR.latestVersion,
-      controllerId, controllerEpoch, brokerEpoch,
-      collection.immutable.Map(new TopicPartition(topic, topicPartition) ->
-        leaderAndIsrPartitionState(leaderEpoch, leaderBrokerId, aliveBrokerIds)).asJava,
+      controllerId, controllerEpoch, brokerEpoch, brokerEpoch,
+      Seq(leaderAndIsrPartitionState(tp, leaderEpoch, leaderBrokerId, aliveBrokerIds)).asJava,
       Set(new Node(followerBrokerId, "host1", 0),
         new Node(leaderBrokerId, "host2", 1)).asJava).build()
     replicaManager.becomeLeaderOrFollower(correlationId, leaderAndIsrRequest0,
@@ -579,6 +854,468 @@ class ReplicaManagerTest {
     EasyMock.verify(mockLogMgr)
   }
 
+  @Test
+  def testReplicaSelector(): Unit = {
+    val topicPartition = 0
+    val followerBrokerId = 0
+    val leaderBrokerId = 1
+    val leaderEpoch = 1
+    val leaderEpochIncrement = 2
+    val aliveBrokerIds = Seq[Integer] (followerBrokerId, leaderBrokerId)
+    val countDownLatch = new CountDownLatch(1)
+
+    // Prepare the mocked components for the test
+    val (replicaManager, _) = prepareReplicaManagerAndLogManager(
+      topicPartition, leaderEpoch + leaderEpochIncrement, followerBrokerId,
+      leaderBrokerId, countDownLatch, expectTruncation = true)
+
+    val tp = new TopicPartition(topic, topicPartition)
+    val partition = replicaManager.createPartition(tp)
+
+    val offsetCheckpoints = new LazyOffsetCheckpoints(replicaManager.highWatermarkCheckpoints)
+    partition.createLogIfNotExists(isNew = false, isFutureReplica = false, offsetCheckpoints)
+    partition.makeLeader(
+      leaderAndIsrPartitionState(tp, leaderEpoch, leaderBrokerId, aliveBrokerIds),
+      offsetCheckpoints
+    )
+
+    val tp0 = new TopicPartition(topic, 0)
+
+    val metadata: ClientMetadata = new DefaultClientMetadata("rack-a", "client-id",
+      InetAddress.getByName("localhost"), KafkaPrincipal.ANONYMOUS, "default")
+
+    // We expect to select the leader, which means we return None
+    val preferredReadReplica: Option[Int] = replicaManager.findPreferredReadReplica(
+      tp0, metadata, Request.OrdinaryConsumerId, 1L, System.currentTimeMillis)
+    assertFalse(preferredReadReplica.isDefined)
+  }
+
+  @Test
+  def testPreferredReplicaAsFollower(): Unit = {
+    val topicPartition = 0
+    val followerBrokerId = 0
+    val leaderBrokerId = 1
+    val leaderEpoch = 1
+    val leaderEpochIncrement = 2
+    val countDownLatch = new CountDownLatch(1)
+
+    // Prepare the mocked components for the test
+    val (replicaManager, _) = prepareReplicaManagerAndLogManager(
+      topicPartition, leaderEpoch + leaderEpochIncrement, followerBrokerId,
+      leaderBrokerId, countDownLatch, expectTruncation = true)
+
+    val brokerList = Seq[Integer](0, 1).asJava
+
+    val tp0 = new TopicPartition(topic, 0)
+
+    replicaManager.createPartition(new TopicPartition(topic, 0))
+
+    // Make this replica the follower
+    val leaderAndIsrRequest2 = new LeaderAndIsrRequest.Builder(ApiKeys.LEADER_AND_ISR.latestVersion, 0, 0, brokerEpoch, brokerEpoch,
+      Seq(new LeaderAndIsrPartitionState()
+        .setTopicName(topic)
+        .setPartitionIndex(0)
+        .setControllerEpoch(0)
+        .setLeader(1)
+        .setLeaderEpoch(1)
+        .setIsr(brokerList)
+        .setZkVersion(0)
+        .setReplicas(brokerList)
+        .setIsNew(false)).asJava,
+      Set(new Node(0, "host1", 0), new Node(1, "host2", 1)).asJava).build()
+    replicaManager.becomeLeaderOrFollower(1, leaderAndIsrRequest2, (_, _) => ())
+
+    val metadata: ClientMetadata = new DefaultClientMetadata("rack-a", "client-id",
+      InetAddress.getByName("localhost"), KafkaPrincipal.ANONYMOUS, "default")
+
+    val consumerResult = fetchAsConsumer(replicaManager, tp0,
+      new PartitionData(0, 0, 100000, Optional.empty()),
+      clientMetadata = Some(metadata))
+
+    // Fetch from follower succeeds
+    assertTrue(consumerResult.isFired)
+
+    // But only leader will compute preferred replica
+    assertTrue(consumerResult.assertFired.preferredReadReplica.isEmpty)
+  }
+
+  @Test
+  def testPreferredReplicaAsLeader(): Unit = {
+    val topicPartition = 0
+    val followerBrokerId = 0
+    val leaderBrokerId = 1
+    val leaderEpoch = 1
+    val leaderEpochIncrement = 2
+    val countDownLatch = new CountDownLatch(1)
+
+    // Prepare the mocked components for the test
+    val (replicaManager, _) = prepareReplicaManagerAndLogManager(
+      topicPartition, leaderEpoch + leaderEpochIncrement, followerBrokerId,
+      leaderBrokerId, countDownLatch, expectTruncation = true)
+
+    val brokerList = Seq[Integer](0, 1).asJava
+
+    val tp0 = new TopicPartition(topic, 0)
+
+    replicaManager.createPartition(new TopicPartition(topic, 0))
+
+    // Make this replica the follower
+    val leaderAndIsrRequest2 = new LeaderAndIsrRequest.Builder(ApiKeys.LEADER_AND_ISR.latestVersion, 0, 0, brokerEpoch, brokerEpoch,
+      Seq(new LeaderAndIsrPartitionState()
+        .setTopicName(topic)
+        .setPartitionIndex(0)
+        .setControllerEpoch(0)
+        .setLeader(0)
+        .setLeaderEpoch(1)
+        .setIsr(brokerList)
+        .setZkVersion(0)
+        .setReplicas(brokerList)
+        .setIsNew(false)).asJava,
+      Set(new Node(0, "host1", 0), new Node(1, "host2", 1)).asJava).build()
+    replicaManager.becomeLeaderOrFollower(1, leaderAndIsrRequest2, (_, _) => ())
+
+    val metadata: ClientMetadata = new DefaultClientMetadata("rack-a", "client-id",
+      InetAddress.getByName("localhost"), KafkaPrincipal.ANONYMOUS, "default")
+
+    val consumerResult = fetchAsConsumer(replicaManager, tp0,
+      new PartitionData(0, 0, 100000, Optional.empty()),
+      clientMetadata = Some(metadata))
+
+    // Fetch from follower succeeds
+    assertTrue(consumerResult.isFired)
+
+    // Returns a preferred replica (should just be the leader, which is None)
+    assertFalse(consumerResult.assertFired.preferredReadReplica.isDefined)
+  }
+
+  @Test(expected = classOf[ClassNotFoundException])
+  def testUnknownReplicaSelector(): Unit = {
+    val topicPartition = 0
+    val followerBrokerId = 0
+    val leaderBrokerId = 1
+    val leaderEpoch = 1
+    val leaderEpochIncrement = 2
+    val countDownLatch = new CountDownLatch(1)
+
+    val props = new Properties()
+    props.put(KafkaConfig.ReplicaSelectorClassProp, "non-a-class")
+    prepareReplicaManagerAndLogManager(
+      topicPartition, leaderEpoch + leaderEpochIncrement, followerBrokerId,
+      leaderBrokerId, countDownLatch, expectTruncation = true, extraProps = props)
+  }
+
+  @Test
+  def testDefaultReplicaSelector(): Unit = {
+    val topicPartition = 0
+    val followerBrokerId = 0
+    val leaderBrokerId = 1
+    val leaderEpoch = 1
+    val leaderEpochIncrement = 2
+    val countDownLatch = new CountDownLatch(1)
+
+    val (replicaManager, _) = prepareReplicaManagerAndLogManager(
+      topicPartition, leaderEpoch + leaderEpochIncrement, followerBrokerId,
+      leaderBrokerId, countDownLatch, expectTruncation = true)
+    assertFalse(replicaManager.replicaSelectorOpt.isDefined)
+  }
+
+  @Test
+  def testFetchFollowerNotAllowedForOlderClients(): Unit = {
+    val replicaManager = setupReplicaManagerWithMockedPurgatories(new MockTimer, aliveBrokerIds = Seq(0, 1))
+
+    val tp0 = new TopicPartition(topic, 0)
+    val offsetCheckpoints = new LazyOffsetCheckpoints(replicaManager.highWatermarkCheckpoints)
+    replicaManager.createPartition(tp0).createLogIfNotExists(isNew = false, isFutureReplica = false, offsetCheckpoints)
+    val partition0Replicas = Seq[Integer](0, 1).asJava
+    val becomeFollowerRequest = new LeaderAndIsrRequest.Builder(ApiKeys.LEADER_AND_ISR.latestVersion, 0, 0, brokerEpoch, brokerEpoch,
+      Seq(new LeaderAndIsrPartitionState()
+        .setTopicName(tp0.topic)
+        .setPartitionIndex(tp0.partition)
+        .setControllerEpoch(0)
+        .setLeader(1)
+        .setLeaderEpoch(0)
+        .setIsr(partition0Replicas)
+        .setZkVersion(0)
+        .setReplicas(partition0Replicas)
+        .setIsNew(true)).asJava,
+      Set(new Node(0, "host1", 0), new Node(1, "host2", 1)).asJava).build()
+    replicaManager.becomeLeaderOrFollower(0, becomeFollowerRequest, (_, _) => ())
+
+    // Fetch from follower, with non-empty ClientMetadata (FetchRequest v11+)
+    val clientMetadata = new DefaultClientMetadata("", "", null, KafkaPrincipal.ANONYMOUS, "")
+    var partitionData = new FetchRequest.PartitionData(0L, 0L, 100,
+      Optional.of(0))
+    var fetchResult = sendConsumerFetch(replicaManager, tp0, partitionData, Some(clientMetadata))
+    assertNotNull(fetchResult.get)
+    assertEquals(Errors.NONE, fetchResult.get.error)
+
+    // Fetch from follower, with empty ClientMetadata (which implies an older version)
+    partitionData = new FetchRequest.PartitionData(0L, 0L, 100,
+      Optional.of(0))
+    fetchResult = sendConsumerFetch(replicaManager, tp0, partitionData, None)
+    assertNotNull(fetchResult.get)
+    assertEquals(Errors.NOT_LEADER_FOR_PARTITION, fetchResult.get.error)
+  }
+
+  @Test
+  def testBecomeFollowerWhileOldClientFetchInPurgatory(): Unit = {
+    val mockTimer = new MockTimer
+    val replicaManager = setupReplicaManagerWithMockedPurgatories(mockTimer, aliveBrokerIds = Seq(0, 1))
+
+    val tp0 = new TopicPartition(topic, 0)
+    val offsetCheckpoints = new LazyOffsetCheckpoints(replicaManager.highWatermarkCheckpoints)
+    replicaManager.createPartition(tp0).createLogIfNotExists(isNew = false, isFutureReplica = false, offsetCheckpoints)
+    val partition0Replicas = Seq[Integer](0, 1).asJava
+
+    val becomeLeaderRequest = new LeaderAndIsrRequest.Builder(ApiKeys.LEADER_AND_ISR.latestVersion, 0, 0, brokerEpoch, brokerEpoch,
+      Seq(new LeaderAndIsrPartitionState()
+        .setTopicName(tp0.topic)
+        .setPartitionIndex(tp0.partition)
+        .setControllerEpoch(0)
+        .setLeader(0)
+        .setLeaderEpoch(1)
+        .setIsr(partition0Replicas)
+        .setZkVersion(0)
+        .setReplicas(partition0Replicas)
+        .setIsNew(true)).asJava,
+      Set(new Node(0, "host1", 0), new Node(1, "host2", 1)).asJava).build()
+    replicaManager.becomeLeaderOrFollower(1, becomeLeaderRequest, (_, _) => ())
+
+    val partitionData = new FetchRequest.PartitionData(0L, 0L, 100,
+      Optional.empty())
+    val fetchResult = sendConsumerFetch(replicaManager, tp0, partitionData, None, timeout = 10)
+    assertNull(fetchResult.get)
+
+    // Become a follower and ensure that the delayed fetch returns immediately
+    val becomeFollowerRequest = new LeaderAndIsrRequest.Builder(ApiKeys.LEADER_AND_ISR.latestVersion, 0, 0, brokerEpoch, brokerEpoch,
+      Seq(new LeaderAndIsrPartitionState()
+        .setTopicName(tp0.topic)
+        .setPartitionIndex(tp0.partition)
+        .setControllerEpoch(0)
+        .setLeader(1)
+        .setLeaderEpoch(2)
+        .setIsr(partition0Replicas)
+        .setZkVersion(0)
+        .setReplicas(partition0Replicas)
+        .setIsNew(true)).asJava,
+      Set(new Node(0, "host1", 0), new Node(1, "host2", 1)).asJava).build()
+    replicaManager.becomeLeaderOrFollower(0, becomeFollowerRequest, (_, _) => ())
+
+    assertNotNull(fetchResult.get)
+    assertEquals(Errors.NOT_LEADER_FOR_PARTITION, fetchResult.get.error)
+  }
+
+  @Test
+  def testBecomeFollowerWhileNewClientFetchInPurgatory(): Unit = {
+    val mockTimer = new MockTimer
+    val replicaManager = setupReplicaManagerWithMockedPurgatories(mockTimer, aliveBrokerIds = Seq(0, 1))
+
+    val tp0 = new TopicPartition(topic, 0)
+    val offsetCheckpoints = new LazyOffsetCheckpoints(replicaManager.highWatermarkCheckpoints)
+    replicaManager.createPartition(tp0).createLogIfNotExists(isNew = false, isFutureReplica = false, offsetCheckpoints)
+    val partition0Replicas = Seq[Integer](0, 1).asJava
+
+    val becomeLeaderRequest = new LeaderAndIsrRequest.Builder(ApiKeys.LEADER_AND_ISR.latestVersion, 0, 0, brokerEpoch, brokerEpoch,
+      Seq(new LeaderAndIsrPartitionState()
+        .setTopicName(tp0.topic)
+        .setPartitionIndex(tp0.partition)
+        .setControllerEpoch(0)
+        .setLeader(0)
+        .setLeaderEpoch(1)
+        .setIsr(partition0Replicas)
+        .setZkVersion(0)
+        .setReplicas(partition0Replicas)
+        .setIsNew(true)).asJava,
+      Set(new Node(0, "host1", 0), new Node(1, "host2", 1)).asJava).build()
+    replicaManager.becomeLeaderOrFollower(1, becomeLeaderRequest, (_, _) => ())
+
+    val clientMetadata = new DefaultClientMetadata("", "", null, KafkaPrincipal.ANONYMOUS, "")
+    val partitionData = new FetchRequest.PartitionData(0L, 0L, 100,
+      Optional.of(1))
+    val fetchResult = sendConsumerFetch(replicaManager, tp0, partitionData, Some(clientMetadata), timeout = 10)
+    assertNull(fetchResult.get)
+
+    // Become a follower and ensure that the delayed fetch returns immediately
+    val becomeFollowerRequest = new LeaderAndIsrRequest.Builder(ApiKeys.LEADER_AND_ISR.latestVersion, 0, 0, brokerEpoch, brokerEpoch,
+      Seq(new LeaderAndIsrPartitionState()
+        .setTopicName(tp0.topic)
+        .setPartitionIndex(tp0.partition)
+        .setControllerEpoch(0)
+        .setLeader(1)
+        .setLeaderEpoch(2)
+        .setIsr(partition0Replicas)
+        .setZkVersion(0)
+        .setReplicas(partition0Replicas)
+        .setIsNew(true)).asJava,
+      Set(new Node(0, "host1", 0), new Node(1, "host2", 1)).asJava).build()
+    replicaManager.becomeLeaderOrFollower(0, becomeFollowerRequest, (_, _) => ())
+
+    assertNotNull(fetchResult.get)
+    assertEquals(Errors.FENCED_LEADER_EPOCH, fetchResult.get.error)
+  }
+
+  @Test
+  def testFetchFromLeaderAlwaysAllowed(): Unit = {
+    val replicaManager = setupReplicaManagerWithMockedPurgatories(new MockTimer, aliveBrokerIds = Seq(0, 1))
+
+    val tp0 = new TopicPartition(topic, 0)
+    val offsetCheckpoints = new LazyOffsetCheckpoints(replicaManager.highWatermarkCheckpoints)
+    replicaManager.createPartition(tp0).createLogIfNotExists(isNew = false, isFutureReplica = false, offsetCheckpoints)
+    val partition0Replicas = Seq[Integer](0, 1).asJava
+
+    val becomeLeaderRequest = new LeaderAndIsrRequest.Builder(ApiKeys.LEADER_AND_ISR.latestVersion, 0, 0, brokerEpoch, brokerEpoch,
+      Seq(new LeaderAndIsrPartitionState()
+        .setTopicName(tp0.topic)
+        .setPartitionIndex(tp0.partition)
+        .setControllerEpoch(0)
+        .setLeader(0)
+        .setLeaderEpoch(1)
+        .setIsr(partition0Replicas)
+        .setZkVersion(0)
+        .setReplicas(partition0Replicas)
+        .setIsNew(true)).asJava,
+      Set(new Node(0, "host1", 0), new Node(1, "host2", 1)).asJava).build()
+    replicaManager.becomeLeaderOrFollower(1, becomeLeaderRequest, (_, _) => ())
+
+    val clientMetadata = new DefaultClientMetadata("", "", null, KafkaPrincipal.ANONYMOUS, "")
+    var partitionData = new FetchRequest.PartitionData(0L, 0L, 100,
+      Optional.of(1))
+    var fetchResult = sendConsumerFetch(replicaManager, tp0, partitionData, Some(clientMetadata))
+    assertNotNull(fetchResult.get)
+    assertEquals(Errors.NONE, fetchResult.get.error)
+
+    partitionData = new FetchRequest.PartitionData(0L, 0L, 100,
+      Optional.empty())
+    fetchResult = sendConsumerFetch(replicaManager, tp0, partitionData, Some(clientMetadata))
+    assertNotNull(fetchResult.get)
+    assertEquals(Errors.NONE, fetchResult.get.error)
+  }
+
+  @Test
+  def testClearFetchPurgatoryOnStopReplica(): Unit = {
+    // As part of a reassignment, we may send StopReplica to the old leader.
+    // In this case, we should ensure that pending purgatory operations are cancelled
+    // immediately rather than sitting around to timeout.
+
+    val mockTimer = new MockTimer
+    val replicaManager = setupReplicaManagerWithMockedPurgatories(mockTimer, aliveBrokerIds = Seq(0, 1))
+
+    val tp0 = new TopicPartition(topic, 0)
+    val offsetCheckpoints = new LazyOffsetCheckpoints(replicaManager.highWatermarkCheckpoints)
+    replicaManager.createPartition(tp0).createLogIfNotExists(isNew = false, isFutureReplica = false, offsetCheckpoints)
+    val partition0Replicas = Seq[Integer](0, 1).asJava
+
+    val becomeLeaderRequest = new LeaderAndIsrRequest.Builder(ApiKeys.LEADER_AND_ISR.latestVersion, 0, 0, brokerEpoch, brokerEpoch,
+      Seq(new LeaderAndIsrPartitionState()
+        .setTopicName(tp0.topic)
+        .setPartitionIndex(tp0.partition)
+        .setControllerEpoch(0)
+        .setLeader(0)
+        .setLeaderEpoch(1)
+        .setIsr(partition0Replicas)
+        .setZkVersion(0)
+        .setReplicas(partition0Replicas)
+        .setIsNew(true)).asJava,
+      Set(new Node(0, "host1", 0), new Node(1, "host2", 1)).asJava).build()
+    replicaManager.becomeLeaderOrFollower(1, becomeLeaderRequest, (_, _) => ())
+
+    val partitionData = new FetchRequest.PartitionData(0L, 0L, 100,
+      Optional.of(1))
+    val fetchResult = sendConsumerFetch(replicaManager, tp0, partitionData, None, timeout = 10)
+    assertNull(fetchResult.get)
+
+    Mockito.when(replicaManager.metadataCache.contains(tp0)).thenReturn(true)
+
+    // We have a fetch in purgatory, now receive a stop replica request and
+    // assert that the fetch returns with a NOT_LEADER error
+
+    replicaManager.stopReplica(tp0, deletePartition = true)
+    assertNotNull(fetchResult.get)
+    assertEquals(Errors.NOT_LEADER_FOR_PARTITION, fetchResult.get.error)
+  }
+
+  @Test
+  def testClearProducePurgatoryOnStopReplica(): Unit = {
+    val mockTimer = new MockTimer
+    val replicaManager = setupReplicaManagerWithMockedPurgatories(mockTimer, aliveBrokerIds = Seq(0, 1))
+
+    val tp0 = new TopicPartition(topic, 0)
+    val offsetCheckpoints = new LazyOffsetCheckpoints(replicaManager.highWatermarkCheckpoints)
+    replicaManager.createPartition(tp0).createLogIfNotExists(isNew = false, isFutureReplica = false, offsetCheckpoints)
+    val partition0Replicas = Seq[Integer](0, 1).asJava
+
+    val becomeLeaderRequest = new LeaderAndIsrRequest.Builder(ApiKeys.LEADER_AND_ISR.latestVersion, 0, 0, brokerEpoch, brokerEpoch,
+      Seq(new LeaderAndIsrPartitionState()
+        .setTopicName(tp0.topic)
+        .setPartitionIndex(tp0.partition)
+        .setControllerEpoch(0)
+        .setLeader(0)
+        .setLeaderEpoch(1)
+        .setIsr(partition0Replicas)
+        .setZkVersion(0)
+        .setReplicas(partition0Replicas)
+        .setIsNew(true)).asJava,
+      Set(new Node(0, "host1", 0), new Node(1, "host2", 1)).asJava).build()
+    replicaManager.becomeLeaderOrFollower(1, becomeLeaderRequest, (_, _) => ())
+
+    val produceResult = sendProducerAppend(replicaManager, tp0)
+    assertNull(produceResult.get)
+
+    Mockito.when(replicaManager.metadataCache.contains(tp0)).thenReturn(true)
+
+    replicaManager.stopReplica(tp0, deletePartition = true)
+    assertNotNull(produceResult.get)
+    assertEquals(Errors.NOT_LEADER_FOR_PARTITION, produceResult.get.error)
+  }
+
+  private def sendProducerAppend(replicaManager: ReplicaManager,
+                                 topicPartition: TopicPartition): AtomicReference[PartitionResponse] = {
+    val produceResult = new AtomicReference[PartitionResponse]()
+    def callback(response: Map[TopicPartition, PartitionResponse]): Unit = {
+      produceResult.set(response(topicPartition))
+    }
+
+    val records = MemoryRecords.withRecords(CompressionType.NONE,
+      new SimpleRecord("a".getBytes()),
+      new SimpleRecord("b".getBytes()),
+      new SimpleRecord("c".getBytes())
+    )
+
+    replicaManager.appendRecords(
+      timeout = 10,
+      requiredAcks = -1,
+      internalTopicsAllowed = false,
+      origin = AppendOrigin.Client,
+      entriesPerPartition = Map(topicPartition -> records),
+      responseCallback = callback
+    )
+    produceResult
+  }
+
+  private def sendConsumerFetch(replicaManager: ReplicaManager,
+                                topicPartition: TopicPartition,
+                                partitionData: FetchRequest.PartitionData,
+                                clientMetadataOpt: Option[ClientMetadata],
+                                timeout: Long = 0L): AtomicReference[FetchPartitionData] = {
+    val fetchResult = new AtomicReference[FetchPartitionData]()
+    def callback(response: Seq[(TopicPartition, FetchPartitionData)]): Unit = {
+      fetchResult.set(response.toMap.apply(topicPartition))
+    }
+    replicaManager.fetchMessages(
+      timeout = timeout,
+      replicaId = Request.OrdinaryConsumerId,
+      fetchMinBytes = 1,
+      fetchMaxBytes = 100,
+      hardMaxBytesLimit = false,
+      fetchInfos = Seq(topicPartition -> partitionData),
+      quota = UnboundedQuota,
+      isolationLevel = IsolationLevel.READ_UNCOMMITTED,
+      responseCallback = callback,
+      clientMetadata = clientMetadataOpt
+    )
+    fetchResult
+  }
+
   /**
    * This method assumes that the test using created ReplicaManager calls
    * ReplicaManager.becomeLeaderOrFollower() once with LeaderAndIsrRequest containing
@@ -589,9 +1326,11 @@ class ReplicaManagerTest {
                                                  followerBrokerId: Int,
                                                  leaderBrokerId: Int,
                                                  countDownLatch: CountDownLatch,
-                                                 expectTruncation: Boolean) : (ReplicaManager, LogManager) = {
+                                                 expectTruncation: Boolean,
+                                                 extraProps: Properties = new Properties()) : (ReplicaManager, LogManager) = {
     val props = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect)
     props.put("log.dir", TestUtils.tempRelativeDir("data").getAbsolutePath)
+    props.asScala ++= extraProps.asScala
     val config = KafkaConfig.fromProps(props)
 
     // Setup mock local log to have leader epoch of 3 and offset of 10
@@ -602,7 +1341,7 @@ class ReplicaManagerTest {
     val mockBrokerTopicStats = new BrokerTopicStats
     val mockLogDirFailureChannel = new LogDirFailureChannel(config.logDirs.size)
     val mockLog = new Log(
-      dir = new File(new File(config.logDirs.head), s"$topic-0"),
+      _dir = new File(new File(config.logDirs.head), s"$topic-0"),
       config = LogConfig(),
       logStartOffset = 0L,
       recoveryPoint = 0L,
@@ -624,18 +1363,27 @@ class ReplicaManagerTest {
       override def latestEpoch: Option[Int] = Some(leaderEpochFromLeader)
 
       override def logEndOffsetMetadata = LogOffsetMetadata(localLogOffset)
+
+      override def logEndOffset: Long = localLogOffset
     }
 
     // Expect to call LogManager.truncateTo exactly once
+    val topicPartitionObj = new TopicPartition(topic, topicPartition)
     val mockLogMgr: LogManager = EasyMock.createMock(classOf[LogManager])
     EasyMock.expect(mockLogMgr.liveLogDirs).andReturn(config.logDirs.map(new File(_).getAbsoluteFile)).anyTimes
-    EasyMock.expect(mockLogMgr.currentDefaultConfig).andReturn(LogConfig())
-    EasyMock.expect(mockLogMgr.getOrCreateLog(new TopicPartition(topic, topicPartition),
-      LogConfig(), isNew = false, isFuture = false)).andReturn(mockLog).anyTimes
+    EasyMock.expect(mockLogMgr.getOrCreateLog(EasyMock.eq(topicPartitionObj),
+      EasyMock.anyObject[() => LogConfig](), isNew = EasyMock.eq(false),
+      isFuture = EasyMock.eq(false))).andReturn(mockLog).anyTimes
     if (expectTruncation) {
-      EasyMock.expect(mockLogMgr.truncateTo(Map(new TopicPartition(topic, topicPartition) -> offsetFromLeader),
+      EasyMock.expect(mockLogMgr.truncateTo(Map(topicPartitionObj -> offsetFromLeader),
         isFuture = false)).once
     }
+    EasyMock.expect(mockLogMgr.initializingLog(topicPartitionObj)).anyTimes
+    EasyMock.expect(mockLogMgr.getLog(topicPartitionObj, isFuture = true)).andReturn(None)
+
+    EasyMock.expect(mockLogMgr.finishedInitializingLog(
+      EasyMock.eq(topicPartitionObj), EasyMock.anyObject(), EasyMock.anyObject())).anyTimes
+
     EasyMock.replay(mockLogMgr)
 
     val aliveBrokerIds = Seq[Integer](followerBrokerId, leaderBrokerId)
@@ -644,11 +1392,20 @@ class ReplicaManagerTest {
     val metadataCache: MetadataCache = EasyMock.createMock(classOf[MetadataCache])
     EasyMock.expect(metadataCache.getAliveBrokers).andReturn(aliveBrokers).anyTimes
     aliveBrokerIds.foreach { brokerId =>
-      EasyMock.expect(metadataCache.isBrokerAlive(EasyMock.eq(brokerId))).andReturn(true).anyTimes
+      EasyMock.expect(metadataCache.getAliveBroker(EasyMock.eq(brokerId)))
+        .andReturn(Option(createBroker(brokerId, s"host$brokerId", brokerId)))
+        .anyTimes
     }
+    EasyMock
+      .expect(metadataCache.getPartitionReplicaEndpoints(
+        EasyMock.anyObject(), EasyMock.anyObject()))
+      .andReturn(Map(
+        leaderBrokerId -> new Node(leaderBrokerId, "host1", 9092, "rack-a"),
+        followerBrokerId -> new Node(followerBrokerId, "host2", 9092, "rack-b")).toMap
+      )
+      .anyTimes()
     EasyMock.replay(metadataCache)
 
-
     val timer = new MockTimer
     val mockProducePurgatory = new DelayedOperationPurgatory[DelayedProduce](
       purgatoryName = "Produce", timer, reaperEnabled = false)
@@ -656,17 +1413,17 @@ class ReplicaManagerTest {
       purgatoryName = "Fetch", timer, reaperEnabled = false)
     val mockDeleteRecordsPurgatory = new DelayedOperationPurgatory[DelayedDeleteRecords](
       purgatoryName = "DeleteRecords", timer, reaperEnabled = false)
-    val mockElectPreferredLeaderPurgatory = new DelayedOperationPurgatory[DelayedElectPreferredLeader](
-      purgatoryName = "ElectPreferredLeader", timer, reaperEnabled = false)
+    val mockElectLeaderPurgatory = new DelayedOperationPurgatory[DelayedElectLeader](
+      purgatoryName = "ElectLeader", timer, reaperEnabled = false)
 
     // Mock network client to show leader offset of 5
     val quota = QuotaFactory.instantiate(config, metrics, time, "")
-    val blockingSend = new ReplicaFetcherMockBlockingSend(Map(new TopicPartition(topic, topicPartition) ->
+    val blockingSend = new ReplicaFetcherMockBlockingSend(Map(topicPartitionObj ->
       new EpochEndOffset(leaderEpochFromLeader, offsetFromLeader)).asJava, BrokerEndPoint(1, "host1" ,1), time)
     val replicaManager = new ReplicaManager(config, metrics, time, kafkaZkClient, mockScheduler, mockLogMgr,
       new AtomicBoolean(false), quota, mockBrokerTopicStats,
       metadataCache, mockLogDirFailureChannel, mockProducePurgatory, mockFetchPurgatory,
-      mockDeleteRecordsPurgatory, mockElectPreferredLeaderPurgatory, Option(this.getClass.getName)) {
+      mockDeleteRecordsPurgatory, mockElectLeaderPurgatory, Option(this.getClass.getName)) {
 
       override protected def createReplicaFetcherManager(metrics: Metrics,
                                                      time: Time,
@@ -676,7 +1433,7 @@ class ReplicaManagerTest {
 
           override def createFetcherThread(fetcherId: Int, sourceBroker: BrokerEndPoint): ReplicaFetcherThread = {
             new ReplicaFetcherThread(s"ReplicaFetcherThread-$fetcherId", fetcherId,
-              sourceBroker, config, replicaManager, metrics, time, quota.follower, Some(blockingSend)) {
+              sourceBroker, config, failedPartitions, replicaManager, metrics, time, quota.follower, Some(blockingSend)) {
 
               override def doWork() = {
                 // In case the thread starts before the partition is added by AbstractFetcherManager,
@@ -698,11 +1455,20 @@ class ReplicaManagerTest {
     (replicaManager, mockLogMgr)
   }
 
-  private def leaderAndIsrPartitionState(leaderEpoch: Int,
+  private def leaderAndIsrPartitionState(topicPartition: TopicPartition,
+                                         leaderEpoch: Int,
                                          leaderBrokerId: Int,
-                                         aliveBrokerIds: Seq[Integer]) : LeaderAndIsrRequest.PartitionState = {
-    new LeaderAndIsrRequest.PartitionState(controllerEpoch, leaderBrokerId, leaderEpoch, aliveBrokerIds.asJava,
-      zkVersion, aliveBrokerIds.asJava, false)
+                                         aliveBrokerIds: Seq[Integer]): LeaderAndIsrPartitionState = {
+    new LeaderAndIsrPartitionState()
+      .setTopicName(topic)
+      .setPartitionIndex(topicPartition.partition)
+      .setControllerEpoch(controllerEpoch)
+      .setLeader(leaderBrokerId)
+      .setLeaderEpoch(leaderEpoch)
+      .setIsr(aliveBrokerIds.asJava)
+      .setZkVersion(zkVersion)
+      .setReplicas(aliveBrokerIds.asJava)
+      .setIsNew(false)
   }
 
   private class CallbackResult[T] {
@@ -733,7 +1499,7 @@ class ReplicaManagerTest {
   private def appendRecords(replicaManager: ReplicaManager,
                             partition: TopicPartition,
                             records: MemoryRecords,
-                            isFromClient: Boolean = true,
+                            origin: AppendOrigin = AppendOrigin.Client,
                             requiredAcks: Short = -1): CallbackResult[PartitionResponse] = {
     val result = new CallbackResult[PartitionResponse]()
     def appendCallback(responses: Map[TopicPartition, PartitionResponse]): Unit = {
@@ -746,7 +1512,7 @@ class ReplicaManagerTest {
       timeout = 1000,
       requiredAcks = requiredAcks,
       internalTopicsAllowed = false,
-      isFromClient = isFromClient,
+      origin = origin,
       entriesPerPartition = Map(partition -> records),
       responseCallback = appendCallback)
 
@@ -757,16 +1523,18 @@ class ReplicaManagerTest {
                               partition: TopicPartition,
                               partitionData: PartitionData,
                               minBytes: Int = 0,
-                              isolationLevel: IsolationLevel = IsolationLevel.READ_UNCOMMITTED): CallbackResult[FetchPartitionData] = {
-    fetchMessages(replicaManager, replicaId = -1, partition, partitionData, minBytes, isolationLevel)
+                              isolationLevel: IsolationLevel = IsolationLevel.READ_UNCOMMITTED,
+                              clientMetadata: Option[ClientMetadata] = None): CallbackResult[FetchPartitionData] = {
+    fetchMessages(replicaManager, replicaId = -1, partition, partitionData, minBytes, isolationLevel, clientMetadata)
   }
 
   private def fetchAsFollower(replicaManager: ReplicaManager,
                               partition: TopicPartition,
                               partitionData: PartitionData,
                               minBytes: Int = 0,
-                              isolationLevel: IsolationLevel = IsolationLevel.READ_UNCOMMITTED): CallbackResult[FetchPartitionData] = {
-    fetchMessages(replicaManager, replicaId = 1, partition, partitionData, minBytes, isolationLevel)
+                              isolationLevel: IsolationLevel = IsolationLevel.READ_UNCOMMITTED,
+                              clientMetadata: Option[ClientMetadata] = None): CallbackResult[FetchPartitionData] = {
+    fetchMessages(replicaManager, replicaId = 1, partition, partitionData, minBytes, isolationLevel, clientMetadata)
   }
 
   private def fetchMessages(replicaManager: ReplicaManager,
@@ -774,7 +1542,8 @@ class ReplicaManagerTest {
                             partition: TopicPartition,
                             partitionData: PartitionData,
                             minBytes: Int,
-                            isolationLevel: IsolationLevel): CallbackResult[FetchPartitionData] = {
+                            isolationLevel: IsolationLevel,
+                            clientMetadata: Option[ClientMetadata]): CallbackResult[FetchPartitionData] = {
     val result = new CallbackResult[FetchPartitionData]()
     def fetchCallback(responseStatus: Seq[(TopicPartition, FetchPartitionData)]) = {
       assertEquals(1, responseStatus.size)
@@ -790,25 +1559,30 @@ class ReplicaManagerTest {
       fetchMaxBytes = Int.MaxValue,
       hardMaxBytesLimit = false,
       fetchInfos = Seq(partition -> partitionData),
+      quota = UnboundedQuota,
       responseCallback = fetchCallback,
-      isolationLevel = isolationLevel)
+      isolationLevel = isolationLevel,
+      clientMetadata = clientMetadata
+    )
 
     result
   }
 
   private def setupReplicaManagerWithMockedPurgatories(timer: MockTimer, aliveBrokerIds: Seq[Int] = Seq(0, 1)): ReplicaManager = {
     val props = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect)
-    props.put("log.dir", TestUtils.tempRelativeDir("data").getAbsolutePath)
+    props.put("log.dirs", TestUtils.tempRelativeDir("data").getAbsolutePath + "," + TestUtils.tempRelativeDir("data2").getAbsolutePath)
     val config = KafkaConfig.fromProps(props)
     val logProps = new Properties()
     val mockLogMgr = TestUtils.createLogManager(config.logDirs.map(new File(_)), LogConfig(logProps))
     val aliveBrokers = aliveBrokerIds.map(brokerId => createBroker(brokerId, s"host$brokerId", brokerId))
-    val metadataCache: MetadataCache = EasyMock.createMock(classOf[MetadataCache])
-    EasyMock.expect(metadataCache.getAliveBrokers).andReturn(aliveBrokers).anyTimes()
+
+    val metadataCache: MetadataCache = Mockito.mock(classOf[MetadataCache])
+    Mockito.when(metadataCache.getAliveBrokers).thenReturn(aliveBrokers)
+
     aliveBrokerIds.foreach { brokerId =>
-      EasyMock.expect(metadataCache.isBrokerAlive(EasyMock.eq(brokerId))).andReturn(true).anyTimes()
+      Mockito.when(metadataCache.getAliveBroker(brokerId))
+        .thenReturn(Option(createBroker(brokerId, s"host$brokerId", brokerId)))
     }
-    EasyMock.replay(metadataCache)
 
     val mockProducePurgatory = new DelayedOperationPurgatory[DelayedProduce](
       purgatoryName = "Produce", timer, reaperEnabled = false)
@@ -816,13 +1590,225 @@ class ReplicaManagerTest {
       purgatoryName = "Fetch", timer, reaperEnabled = false)
     val mockDeleteRecordsPurgatory = new DelayedOperationPurgatory[DelayedDeleteRecords](
       purgatoryName = "DeleteRecords", timer, reaperEnabled = false)
-    val mockDelayedElectPreferredLeaderPurgatory = new DelayedOperationPurgatory[DelayedElectPreferredLeader](
-      purgatoryName = "DelayedElectPreferredLeader", timer, reaperEnabled = false)
+    val mockDelayedElectLeaderPurgatory = new DelayedOperationPurgatory[DelayedElectLeader](
+      purgatoryName = "DelayedElectLeader", timer, reaperEnabled = false)
 
     new ReplicaManager(config, metrics, time, kafkaZkClient, new MockScheduler(time), mockLogMgr,
       new AtomicBoolean(false), QuotaFactory.instantiate(config, metrics, time, ""), new BrokerTopicStats,
       metadataCache, new LogDirFailureChannel(config.logDirs.size), mockProducePurgatory, mockFetchPurgatory,
-      mockDeleteRecordsPurgatory, mockDelayedElectPreferredLeaderPurgatory, Option(this.getClass.getName))
+      mockDeleteRecordsPurgatory, mockDelayedElectLeaderPurgatory, Option(this.getClass.getName))
+  }
+
+  @Test
+  def testOldLeaderLosesMetricsWhenReassignPartitions(): Unit = {
+    val controllerEpoch = 0
+    val leaderEpoch = 0
+    val leaderEpochIncrement = 1
+    val correlationId = 0
+    val controllerId = 0
+    val (rm0, rm1, _, mockTopicStats1) = prepareDifferentReplicaManagersWithMockedBrokerTopicStats()
+
+    EasyMock.expect(mockTopicStats1.removeOldLeaderMetrics(topic)).andVoid.once
+    EasyMock.replay(mockTopicStats1)
+
+    try {
+      // make broker 0 the leader of partition 0 and
+      // make broker 1 the leader of partition 1
+      val tp0 = new TopicPartition(topic, 0)
+      val tp1 = new TopicPartition(topic, 1)
+      val partition0Replicas = Seq[Integer](0, 1).asJava
+      val partition1Replicas = Seq[Integer](1, 0).asJava
+
+      val leaderAndIsrRequest1 = new LeaderAndIsrRequest.Builder(ApiKeys.LEADER_AND_ISR.latestVersion,
+        controllerId, 0, brokerEpoch, brokerEpoch,
+        Seq(
+          new LeaderAndIsrPartitionState()
+            .setTopicName(tp0.topic)
+            .setPartitionIndex(tp0.partition)
+            .setControllerEpoch(controllerEpoch)
+            .setLeader(0)
+            .setLeaderEpoch(leaderEpoch)
+            .setIsr(partition0Replicas)
+            .setZkVersion(0)
+            .setReplicas(partition0Replicas)
+            .setIsNew(true),
+          new LeaderAndIsrPartitionState()
+            .setTopicName(tp1.topic)
+            .setPartitionIndex(tp1.partition)
+            .setControllerEpoch(controllerEpoch)
+            .setLeader(1)
+            .setLeaderEpoch(leaderEpoch)
+            .setIsr(partition1Replicas)
+            .setZkVersion(0)
+            .setReplicas(partition1Replicas)
+            .setIsNew(true)
+        ).asJava,
+        Set(new Node(0, "host0", 0), new Node(1, "host1", 1)).asJava).build()
+
+      rm0.becomeLeaderOrFollower(correlationId, leaderAndIsrRequest1, (_, _) => ())
+      rm1.becomeLeaderOrFollower(correlationId, leaderAndIsrRequest1, (_, _) => ())
+
+      // make broker 0 the leader of partition 1 so broker 1 loses its leadership position
+      val leaderAndIsrRequest2 = new LeaderAndIsrRequest.Builder(ApiKeys.LEADER_AND_ISR.latestVersion, controllerId,
+        controllerEpoch, brokerEpoch, brokerEpoch,
+        Seq(
+          new LeaderAndIsrPartitionState()
+            .setTopicName(tp0.topic)
+            .setPartitionIndex(tp0.partition)
+            .setControllerEpoch(controllerEpoch)
+            .setLeader(0)
+            .setLeaderEpoch(leaderEpoch + leaderEpochIncrement)
+            .setIsr(partition0Replicas)
+            .setZkVersion(0)
+            .setReplicas(partition0Replicas)
+            .setIsNew(true),
+          new LeaderAndIsrPartitionState()
+            .setTopicName(tp1.topic)
+            .setPartitionIndex(tp1.partition)
+            .setControllerEpoch(controllerEpoch)
+            .setLeader(0)
+            .setLeaderEpoch(leaderEpoch + leaderEpochIncrement)
+            .setIsr(partition1Replicas)
+            .setZkVersion(0)
+            .setReplicas(partition1Replicas)
+            .setIsNew(true)
+        ).asJava,
+        Set(new Node(0, "host0", 0), new Node(1, "host1", 1)).asJava).build()
+
+      rm0.becomeLeaderOrFollower(correlationId, leaderAndIsrRequest2, (_, _) => ())
+      rm1.becomeLeaderOrFollower(correlationId, leaderAndIsrRequest2, (_, _) => ())
+    } finally {
+      rm0.shutdown()
+      rm1.shutdown()
+    }
+
+    // verify that broker 1 did remove its metrics when no longer being the leader of partition 1
+    EasyMock.verify(mockTopicStats1)
+  }
+
+  @Test
+  def testOldFollowerLosesMetricsWhenReassignPartitions(): Unit = {
+    val controllerEpoch = 0
+    val leaderEpoch = 0
+    val leaderEpochIncrement = 1
+    val correlationId = 0
+    val controllerId = 0
+    val (rm0, rm1, _, mockTopicStats1) = prepareDifferentReplicaManagersWithMockedBrokerTopicStats()
+
+    EasyMock.expect(mockTopicStats1.removeOldLeaderMetrics(topic)).andVoid.once
+    EasyMock.expect(mockTopicStats1.removeOldFollowerMetrics(topic)).andVoid.once
+    EasyMock.replay(mockTopicStats1)
+
+    try {
+      // make broker 0 the leader of partition 0 and
+      // make broker 1 the leader of partition 1
+      val tp0 = new TopicPartition(topic, 0)
+      val tp1 = new TopicPartition(topic, 1)
+      val partition0Replicas = Seq[Integer](1, 0).asJava
+      val partition1Replicas = Seq[Integer](1, 0).asJava
+
+      val leaderAndIsrRequest1 = new LeaderAndIsrRequest.Builder(ApiKeys.LEADER_AND_ISR.latestVersion,
+        controllerId, 0, brokerEpoch, brokerEpoch,
+        Seq(
+          new LeaderAndIsrPartitionState()
+            .setTopicName(tp0.topic)
+            .setPartitionIndex(tp0.partition)
+            .setControllerEpoch(controllerEpoch)
+            .setLeader(1)
+            .setLeaderEpoch(leaderEpoch)
+            .setIsr(partition0Replicas)
+            .setZkVersion(0)
+            .setReplicas(partition0Replicas)
+            .setIsNew(true),
+          new LeaderAndIsrPartitionState()
+            .setTopicName(tp1.topic)
+            .setPartitionIndex(tp1.partition)
+            .setControllerEpoch(controllerEpoch)
+            .setLeader(1)
+            .setLeaderEpoch(leaderEpoch)
+            .setIsr(partition1Replicas)
+            .setZkVersion(0)
+            .setReplicas(partition1Replicas)
+            .setIsNew(true)
+        ).asJava,
+        Set(new Node(0, "host0", 0), new Node(1, "host1", 1)).asJava).build()
+
+      rm0.becomeLeaderOrFollower(correlationId, leaderAndIsrRequest1, (_, _) => ())
+      rm1.becomeLeaderOrFollower(correlationId, leaderAndIsrRequest1, (_, _) => ())
+
+      // make broker 0 the leader of partition 1 so broker 1 loses its leadership position
+      val leaderAndIsrRequest2 = new LeaderAndIsrRequest.Builder(ApiKeys.LEADER_AND_ISR.latestVersion, controllerId,
+        controllerEpoch, brokerEpoch, brokerEpoch,
+        Seq(
+          new LeaderAndIsrPartitionState()
+            .setTopicName(tp0.topic)
+            .setPartitionIndex(tp0.partition)
+            .setControllerEpoch(controllerEpoch)
+            .setLeader(0)
+            .setLeaderEpoch(leaderEpoch + leaderEpochIncrement)
+            .setIsr(partition0Replicas)
+            .setZkVersion(0)
+            .setReplicas(partition0Replicas)
+            .setIsNew(true),
+          new LeaderAndIsrPartitionState()
+            .setTopicName(tp1.topic)
+            .setPartitionIndex(tp1.partition)
+            .setControllerEpoch(controllerEpoch)
+            .setLeader(0)
+            .setLeaderEpoch(leaderEpoch + leaderEpochIncrement)
+            .setIsr(partition1Replicas)
+            .setZkVersion(0)
+            .setReplicas(partition1Replicas)
+            .setIsNew(true)
+        ).asJava,
+        Set(new Node(0, "host0", 0), new Node(1, "host1", 1)).asJava).build()
+
+      rm0.becomeLeaderOrFollower(correlationId, leaderAndIsrRequest2, (_, _) => ())
+      rm1.becomeLeaderOrFollower(correlationId, leaderAndIsrRequest2, (_, _) => ())
+    } finally {
+      rm0.shutdown()
+      rm1.shutdown()
+    }
+
+    // verify that broker 1 did remove its metrics when no longer being the leader of partition 1
+    EasyMock.verify(mockTopicStats1)
+  }
+
+  private def prepareDifferentReplicaManagersWithMockedBrokerTopicStats(): (ReplicaManager, ReplicaManager, BrokerTopicStats, BrokerTopicStats) = {
+    val props0 = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect)
+    val props1 = TestUtils.createBrokerConfig(1, TestUtils.MockZkConnect)
+
+    props0.put("log0.dir", TestUtils.tempRelativeDir("data").getAbsolutePath)
+    props1.put("log1.dir", TestUtils.tempRelativeDir("data").getAbsolutePath)
+
+    val config0 = KafkaConfig.fromProps(props0)
+    val config1 = KafkaConfig.fromProps(props1)
+
+    val mockLogMgr0 = TestUtils.createLogManager(config0.logDirs.map(new File(_)))
+    val mockLogMgr1 = TestUtils.createLogManager(config1.logDirs.map(new File(_)))
+
+    val mockTopicStats0: BrokerTopicStats = EasyMock.createMock(classOf[BrokerTopicStats])
+    val mockTopicStats1: BrokerTopicStats = EasyMock.createMock(classOf[BrokerTopicStats])
+
+    val metadataCache0: MetadataCache = EasyMock.createMock(classOf[MetadataCache])
+    val metadataCache1: MetadataCache = EasyMock.createMock(classOf[MetadataCache])
+
+    val aliveBrokers = Seq(createBroker(0, "host0", 0), createBroker(1, "host1", 1))
+
+    EasyMock.expect(metadataCache0.getAliveBrokers).andReturn(aliveBrokers).anyTimes()
+    EasyMock.replay(metadataCache0)
+    EasyMock.expect(metadataCache1.getAliveBrokers).andReturn(aliveBrokers).anyTimes()
+    EasyMock.replay(metadataCache1)
+
+    // each replica manager is for a broker
+    val rm0 = new ReplicaManager(config0, metrics, time, kafkaZkClient, new MockScheduler(time), mockLogMgr0,
+      new AtomicBoolean(false), QuotaFactory.instantiate(config0, metrics, time, ""),
+      mockTopicStats0, metadataCache0, new LogDirFailureChannel(config0.logDirs.size))
+    val rm1 = new ReplicaManager(config1, metrics, time, kafkaZkClient, new MockScheduler(time), mockLogMgr1,
+      new AtomicBoolean(false), QuotaFactory.instantiate(config1, metrics, time, ""),
+      mockTopicStats1, metadataCache1, new LogDirFailureChannel(config1.logDirs.size))
+
+    (rm0, rm1, mockTopicStats0, mockTopicStats1)
   }
 
 }
diff --git a/core/src/test/scala/unit/kafka/server/ReplicationQuotaManagerTest.scala b/core/src/test/scala/unit/kafka/server/ReplicationQuotaManagerTest.scala
index dbb52e7f036f2..16c1ab7e2c708 100644
--- a/core/src/test/scala/unit/kafka/server/ReplicationQuotaManagerTest.scala
+++ b/core/src/test/scala/unit/kafka/server/ReplicationQuotaManagerTest.scala
@@ -23,16 +23,22 @@ import org.apache.kafka.common.TopicPartition
 import org.apache.kafka.common.metrics.{MetricConfig, Metrics, Quota}
 import org.apache.kafka.common.utils.MockTime
 import org.junit.Assert.{assertEquals, assertFalse, assertTrue}
-import org.junit.Test
+import org.junit.{After, Test}
 
 import scala.collection.JavaConverters._
 
 class ReplicationQuotaManagerTest {
   private val time = new MockTime
+  private val metrics = new Metrics(new MetricConfig(), Collections.emptyList(), time)
+
+  @After
+  def tearDown(): Unit = {
+    metrics.close()
+  }
 
   @Test
-  def shouldThrottleOnlyDefinedReplicas() {
-    val quota = new ReplicationQuotaManager(ReplicationQuotaManagerConfig(), newMetrics, QuotaType.Fetch, time)
+  def shouldThrottleOnlyDefinedReplicas(): Unit = {
+    val quota = new ReplicationQuotaManager(ReplicationQuotaManagerConfig(), metrics, QuotaType.Fetch, time)
     quota.markThrottled("topic1", Seq(1, 2, 3))
 
     assertTrue(quota.isThrottled(tp1(1)))
@@ -43,7 +49,6 @@ class ReplicationQuotaManagerTest {
 
   @Test
   def shouldExceedQuotaThenReturnBackBelowBoundAsTimePasses(): Unit = {
-    val metrics = newMetrics()
     val quota = new ReplicationQuotaManager(ReplicationQuotaManagerConfig(numQuotaSamples = 10, quotaWindowSizeSeconds = 1), metrics, LeaderReplication, time)
 
     //Given
@@ -105,7 +110,7 @@ class ReplicationQuotaManagerTest {
 
   @Test
   def shouldSupportWildcardThrottledReplicas(): Unit = {
-    val quota = new ReplicationQuotaManager(ReplicationQuotaManagerConfig(), newMetrics, LeaderReplication, time)
+    val quota = new ReplicationQuotaManager(ReplicationQuotaManagerConfig(), metrics, LeaderReplication, time)
 
     //When
     quota.markThrottled("MyTopic")
@@ -116,8 +121,4 @@ class ReplicationQuotaManagerTest {
   }
 
   private def tp1(id: Int): TopicPartition = new TopicPartition("topic1", id)
-
-  private def newMetrics(): Metrics = {
-    new Metrics(new MetricConfig(), Collections.emptyList(), time)
-  }
 }
diff --git a/core/src/test/scala/unit/kafka/server/ReplicationQuotasTest.scala b/core/src/test/scala/unit/kafka/server/ReplicationQuotasTest.scala
index 45c21d2096bdc..45836651927c4 100644
--- a/core/src/test/scala/unit/kafka/server/ReplicationQuotasTest.scala
+++ b/core/src/test/scala/unit/kafka/server/ReplicationQuotasTest.scala
@@ -51,7 +51,7 @@ class ReplicationQuotasTest extends ZooKeeperTestHarness {
   var producer: KafkaProducer[Array[Byte], Array[Byte]] = null
 
   @After
-  override def tearDown() {
+  override def tearDown(): Unit = {
     producer.close()
     shutdownServers(brokers)
     super.tearDown()
diff --git a/core/src/test/scala/unit/kafka/server/RequestQuotaTest.scala b/core/src/test/scala/unit/kafka/server/RequestQuotaTest.scala
index 1c8656de17fe7..79431bee503cd 100644
--- a/core/src/test/scala/unit/kafka/server/RequestQuotaTest.scala
+++ b/core/src/test/scala/unit/kafka/server/RequestQuotaTest.scala
@@ -14,22 +14,23 @@
 
 package kafka.server
 
-import java.nio.ByteBuffer
+import java.util
 import java.util.{Collections, LinkedHashMap, Optional, Properties}
 import java.util.concurrent.{Executors, Future, TimeUnit}
 
 import kafka.log.LogConfig
 import kafka.network.RequestChannel.Session
-import kafka.security.auth._
+import kafka.security.authorizer.AclAuthorizer
 import kafka.utils.TestUtils
-import org.apache.kafka.common.acl.{AccessControlEntry, AccessControlEntryFilter, AclBinding, AclBindingFilter, AclOperation, AclPermissionType}
+import org.apache.kafka.common.{ElectionType, Node, TopicPartition}
+import org.apache.kafka.common.acl._
 import org.apache.kafka.common.config.ConfigResource
-import org.apache.kafka.common.message.{CreateTopicsRequestData, DescribeGroupsRequestData, ElectPreferredLeadersRequestData, LeaveGroupRequestData}
-import org.apache.kafka.common.resource.{PatternType, ResourcePattern, ResourcePatternFilter, ResourceType => AdminResourceType}
-import org.apache.kafka.common.{Node, TopicPartition}
-import org.apache.kafka.common.message.CreateTopicsRequestData.{CreatableTopic, CreatableTopicSet}
-import org.apache.kafka.common.message.SaslAuthenticateRequestData
-import org.apache.kafka.common.message.SaslHandshakeRequestData
+import org.apache.kafka.common.message.CreateTopicsRequestData.{CreatableTopic, CreatableTopicCollection}
+import org.apache.kafka.common.message.JoinGroupRequestData.JoinGroupRequestProtocolCollection
+import org.apache.kafka.common.message.LeaderAndIsrRequestData.LeaderAndIsrPartitionState
+import org.apache.kafka.common.message.LeaveGroupRequestData.MemberIdentity
+import org.apache.kafka.common.message.UpdateMetadataRequestData.{UpdateMetadataBroker, UpdateMetadataEndpoint, UpdateMetadataPartitionState}
+import org.apache.kafka.common.message._
 import org.apache.kafka.common.metrics.{KafkaMetric, Quota, Sensor}
 import org.apache.kafka.common.network.ListenerName
 import org.apache.kafka.common.protocol.ApiKeys
@@ -37,9 +38,10 @@ import org.apache.kafka.common.protocol.types.Struct
 import org.apache.kafka.common.record._
 import org.apache.kafka.common.requests.CreateAclsRequest.AclCreation
 import org.apache.kafka.common.requests._
+import org.apache.kafka.common.resource.{PatternType, ResourcePattern, ResourcePatternFilter, ResourceType => AdminResourceType}
 import org.apache.kafka.common.security.auth.{AuthenticationContext, KafkaPrincipal, KafkaPrincipalBuilder, SecurityProtocol}
-import org.apache.kafka.common.utils.Sanitizer
-import org.apache.kafka.common.utils.SecurityUtils
+import org.apache.kafka.common.utils.{Sanitizer, SecurityUtils}
+import org.apache.kafka.server.authorizer.{Action, AuthorizableRequestContext, AuthorizationResult}
 import org.junit.Assert._
 import org.junit.{After, Before, Test}
 
@@ -48,7 +50,7 @@ import scala.collection.mutable.ListBuffer
 
 class RequestQuotaTest extends BaseRequestTest {
 
-  override def numBrokers: Int = 1
+  override def brokerCount: Int = 1
 
   private val topic = "topic-1"
   private val numPartitions = 1
@@ -65,7 +67,7 @@ class RequestQuotaTest extends BaseRequestTest {
   private val executor = Executors.newCachedThreadPool
   private val tasks = new ListBuffer[Task]
 
-  override def propertyOverrides(properties: Properties): Unit = {
+  override def brokerPropertyOverrides(properties: Properties): Unit = {
     properties.put(KafkaConfig.ControlledShutdownEnableProp, "false")
     properties.put(KafkaConfig.OffsetsTopicReplicationFactorProp, "1")
     properties.put(KafkaConfig.OffsetsTopicPartitionsProp, "1")
@@ -76,7 +78,7 @@ class RequestQuotaTest extends BaseRequestTest {
   }
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     RequestQuotaTest.principal = KafkaPrincipal.ANONYMOUS
     super.setUp()
 
@@ -102,21 +104,25 @@ class RequestQuotaTest extends BaseRequestTest {
     quotaProps.put(DynamicConfig.Client.RequestPercentageOverrideProp, "0.01")
     adminZkClient.changeClientIdConfig(Sanitizer.sanitize(smallQuotaConsumerClientId), quotaProps)
 
-    TestUtils.retry(10000) {
+    TestUtils.retry(20000) {
       val quotaManager = servers.head.dataPlaneRequestProcessor.quotas.request
       assertEquals(s"Default request quota not set", Quota.upperBound(0.01), quotaManager.quota("some-user", "some-client"))
       assertEquals(s"Request quota override not set", Quota.upperBound(2000), quotaManager.quota("some-user", unthrottledClientId))
+      val produceQuotaManager = servers.head.dataPlaneRequestProcessor.quotas.produce
+      assertEquals(s"Produce quota override not set", Quota.upperBound(1), produceQuotaManager.quota("some-user", smallQuotaProducerClientId))
+      val consumeQuotaManager = servers.head.dataPlaneRequestProcessor.quotas.fetch
+      assertEquals(s"Consume quota override not set", Quota.upperBound(1), consumeQuotaManager.quota("some-user", smallQuotaConsumerClientId))
     }
   }
 
   @After
-  override def tearDown() {
+  override def tearDown(): Unit = {
     try executor.shutdownNow()
     finally super.tearDown()
   }
 
   @Test
-  def testResponseThrottleTime() {
+  def testResponseThrottleTime(): Unit = {
     for (apiKey <- RequestQuotaTest.ClientActions)
       submitTest(apiKey, () => checkRequestThrottleTime(apiKey))
 
@@ -124,21 +130,21 @@ class RequestQuotaTest extends BaseRequestTest {
   }
 
   @Test
-  def testResponseThrottleTimeWhenBothProduceAndRequestQuotasViolated() {
+  def testResponseThrottleTimeWhenBothProduceAndRequestQuotasViolated(): Unit = {
     val apiKey = ApiKeys.PRODUCE
     submitTest(apiKey, () => checkSmallQuotaProducerRequestThrottleTime(apiKey))
     waitAndCheckResults()
   }
 
   @Test
-  def testResponseThrottleTimeWhenBothFetchAndRequestQuotasViolated() {
+  def testResponseThrottleTimeWhenBothFetchAndRequestQuotasViolated(): Unit = {
     val apiKey = ApiKeys.FETCH
     submitTest(apiKey, () => checkSmallQuotaConsumerRequestThrottleTime(apiKey))
     waitAndCheckResults()
   }
 
   @Test
-  def testUnthrottledClient() {
+  def testUnthrottledClient(): Unit = {
     for (apiKey <- RequestQuotaTest.ClientActions)
       submitTest(apiKey, () => checkUnthrottledClient(apiKey))
 
@@ -146,7 +152,7 @@ class RequestQuotaTest extends BaseRequestTest {
   }
 
   @Test
-  def testExemptRequestTime() {
+  def testExemptRequestTime(): Unit = {
     for (apiKey <- RequestQuotaTest.ClusterActions)
       submitTest(apiKey, () => checkExemptRequestMetric(apiKey))
 
@@ -154,7 +160,7 @@ class RequestQuotaTest extends BaseRequestTest {
   }
 
   @Test
-  def testUnauthorizedThrottle() {
+  def testUnauthorizedThrottle(): Unit = {
     RequestQuotaTest.principal = RequestQuotaTest.UnauthorizedPrincipal
 
     for (apiKey <- ApiKeys.values)
@@ -222,56 +228,128 @@ class RequestQuotaTest extends BaseRequestTest {
               0L, Optional.of[Integer](15))).asJava)
 
         case ApiKeys.LEADER_AND_ISR =>
-          new LeaderAndIsrRequest.Builder(ApiKeys.LEADER_AND_ISR.latestVersion, brokerId, Int.MaxValue, Long.MaxValue,
-            Map(tp -> new LeaderAndIsrRequest.PartitionState(Int.MaxValue, brokerId, Int.MaxValue, List(brokerId).asJava,
-              2, Seq(brokerId).asJava, true)).asJava,
+          new LeaderAndIsrRequest.Builder(ApiKeys.LEADER_AND_ISR.latestVersion, brokerId, Int.MaxValue, Long.MaxValue, Long.MaxValue,
+            Seq(new LeaderAndIsrPartitionState()
+              .setTopicName(tp.topic)
+              .setPartitionIndex(tp.partition)
+              .setControllerEpoch(Int.MaxValue)
+              .setLeader(brokerId)
+              .setLeaderEpoch(Int.MaxValue)
+              .setIsr(List(brokerId).asJava)
+              .setZkVersion(2)
+              .setReplicas(Seq(brokerId).asJava)
+              .setIsNew(true)).asJava,
             Set(new Node(brokerId, "localhost", 0)).asJava)
 
         case ApiKeys.STOP_REPLICA =>
-          new StopReplicaRequest.Builder(ApiKeys.STOP_REPLICA.latestVersion, brokerId, Int.MaxValue, Long.MaxValue, true, Set(tp).asJava)
+          new StopReplicaRequest.Builder(ApiKeys.STOP_REPLICA.latestVersion, brokerId, Int.MaxValue, Long.MaxValue, Long.MaxValue, true, Set(tp).asJava)
 
         case ApiKeys.UPDATE_METADATA =>
-          val partitionState = Map(tp -> new UpdateMetadataRequest.PartitionState(
-            Int.MaxValue, brokerId, Int.MaxValue, List(brokerId).asJava, 2, Seq(brokerId).asJava, Seq.empty[Integer].asJava)).asJava
+          val partitionState = Seq(new UpdateMetadataPartitionState()
+            .setTopicName(tp.topic)
+            .setPartitionIndex(tp.partition)
+            .setControllerEpoch(Int.MaxValue)
+            .setLeader(brokerId)
+            .setLeaderEpoch(Int.MaxValue)
+            .setIsr(List(brokerId).asJava)
+            .setZkVersion(2)
+            .setReplicas(Seq(brokerId).asJava)).asJava
           val securityProtocol = SecurityProtocol.PLAINTEXT
-          val brokers = Set(new UpdateMetadataRequest.Broker(brokerId,
-            Seq(new UpdateMetadataRequest.EndPoint("localhost", 0, securityProtocol,
-            ListenerName.forSecurityProtocol(securityProtocol))).asJava, null)).asJava
-          new UpdateMetadataRequest.Builder(ApiKeys.UPDATE_METADATA.latestVersion, brokerId, Int.MaxValue, Long.MaxValue, partitionState, brokers)
+          val brokers = Seq(new UpdateMetadataBroker()
+            .setId(brokerId)
+            .setEndpoints(Seq(new UpdateMetadataEndpoint()
+              .setHost("localhost")
+              .setPort(0)
+              .setSecurityProtocol(securityProtocol.id)
+              .setListener(ListenerName.forSecurityProtocol(securityProtocol).value)).asJava)).asJava
+          new UpdateMetadataRequest.Builder(ApiKeys.UPDATE_METADATA.latestVersion, brokerId, Int.MaxValue, Long.MaxValue, Long.MaxValue, partitionState, brokers)
 
         case ApiKeys.CONTROLLED_SHUTDOWN =>
-          new ControlledShutdownRequest.Builder(brokerId, Long.MaxValue, ApiKeys.CONTROLLED_SHUTDOWN.latestVersion)
+          new ControlledShutdownRequest.Builder(
+              new ControlledShutdownRequestData()
+                .setBrokerId(brokerId)
+                .setBrokerEpoch(Long.MaxValue),
+              ApiKeys.CONTROLLED_SHUTDOWN.latestVersion)
 
         case ApiKeys.OFFSET_COMMIT =>
-          new OffsetCommitRequest.Builder("test-group",
-            Map(tp -> new OffsetCommitRequest.PartitionData(0, Optional.empty[Integer](), "metadata")).asJava).
-            setMemberId("").setGenerationId(1)
-
+          new OffsetCommitRequest.Builder(
+            new OffsetCommitRequestData()
+              .setGroupId("test-group")
+              .setGenerationId(1)
+              .setMemberId(JoinGroupRequest.UNKNOWN_MEMBER_ID)
+              .setTopics(
+                Collections.singletonList(
+                  new OffsetCommitRequestData.OffsetCommitRequestTopic()
+                    .setName(topic)
+                    .setPartitions(
+                      Collections.singletonList(
+                        new OffsetCommitRequestData.OffsetCommitRequestPartition()
+                          .setPartitionIndex(0)
+                          .setCommittedLeaderEpoch(RecordBatch.NO_PARTITION_LEADER_EPOCH)
+                          .setCommittedOffset(0)
+                          .setCommittedMetadata("metadata")
+                      )
+                    )
+                )
+              )
+          )
         case ApiKeys.OFFSET_FETCH =>
           new OffsetFetchRequest.Builder("test-group", List(tp).asJava)
 
         case ApiKeys.FIND_COORDINATOR =>
-          new FindCoordinatorRequest.Builder(FindCoordinatorRequest.CoordinatorType.GROUP, "test-group")
+          new FindCoordinatorRequest.Builder(
+              new FindCoordinatorRequestData()
+                .setKeyType(FindCoordinatorRequest.CoordinatorType.GROUP.id)
+                .setKey("test-group"))
 
         case ApiKeys.JOIN_GROUP =>
-          new JoinGroupRequest.Builder("test-join-group", 200, "", "consumer",
-            List(new JoinGroupRequest.ProtocolMetadata("consumer-range", ByteBuffer.wrap("test".getBytes()))).asJava)
-           .setRebalanceTimeout(100)
+          new JoinGroupRequest.Builder(
+            new JoinGroupRequestData()
+              .setGroupId("test-join-group")
+              .setSessionTimeoutMs(200)
+              .setMemberId(JoinGroupRequest.UNKNOWN_MEMBER_ID)
+              .setGroupInstanceId(null)
+              .setProtocolType("consumer")
+              .setProtocols(
+                new JoinGroupRequestProtocolCollection(
+                  Collections.singletonList(new JoinGroupRequestData.JoinGroupRequestProtocol()
+                    .setName("consumer-range")
+                    .setMetadata("test".getBytes())).iterator()
+                )
+              )
+              .setRebalanceTimeoutMs(100)
+          )
 
         case ApiKeys.HEARTBEAT =>
-          new HeartbeatRequest.Builder("test-group", 1, "")
+          new HeartbeatRequest.Builder(
+            new HeartbeatRequestData()
+              .setGroupId("test-group")
+              .setGenerationId(1)
+              .setMemberId(JoinGroupRequest.UNKNOWN_MEMBER_ID)
+          )
 
         case ApiKeys.LEAVE_GROUP =>
-          new LeaveGroupRequest.Builder(new LeaveGroupRequestData().setGroupId("test-leave-group").setMemberId(JoinGroupRequest.UNKNOWN_MEMBER_ID))
+          new LeaveGroupRequest.Builder(
+            "test-leave-group",
+            Collections.singletonList(
+              new MemberIdentity()
+                .setMemberId(JoinGroupRequest.UNKNOWN_MEMBER_ID))
+          )
 
         case ApiKeys.SYNC_GROUP =>
-          new SyncGroupRequest.Builder("test-sync-group", 1, "", Map[String, ByteBuffer]().asJava)
+          new SyncGroupRequest.Builder(
+            new SyncGroupRequestData()
+              .setGroupId("test-sync-group")
+              .setGenerationId(1)
+              .setMemberId(JoinGroupRequest.UNKNOWN_MEMBER_ID)
+              .setAssignments(Collections.emptyList())
+          )
 
         case ApiKeys.DESCRIBE_GROUPS =>
           new DescribeGroupsRequest.Builder(new DescribeGroupsRequestData().setGroups(List("test-group").asJava))
 
         case ApiKeys.LIST_GROUPS =>
-          new ListGroupsRequest.Builder()
+          new ListGroupsRequest.Builder(new ListGroupsRequestData())
 
         case ApiKeys.SASL_HANDSHAKE =>
           new SaslHandshakeRequest.Builder(new SaslHandshakeRequestData().setMechanism("PLAIN"))
@@ -280,28 +358,34 @@ class RequestQuotaTest extends BaseRequestTest {
           new SaslAuthenticateRequest.Builder(new SaslAuthenticateRequestData().setAuthBytes(new Array[Byte](0)))
 
         case ApiKeys.API_VERSIONS =>
-          new ApiVersionsRequest.Builder
+          new ApiVersionsRequest.Builder()
 
         case ApiKeys.CREATE_TOPICS => {
           new CreateTopicsRequest.Builder(
             new CreateTopicsRequestData().setTopics(
-              new CreatableTopicSet(Collections.singleton(
+              new CreatableTopicCollection(Collections.singleton(
                 new CreatableTopic().setName("topic-2").setNumPartitions(1).
                   setReplicationFactor(1.toShort)).iterator())))
         }
 
         case ApiKeys.DELETE_TOPICS =>
-          new DeleteTopicsRequest.Builder(Set("topic-2").asJava, 5000)
+          new DeleteTopicsRequest.Builder(
+              new DeleteTopicsRequestData()
+              .setTopicNames(Collections.singletonList("topic-2"))
+              .setTimeoutMs(5000))
 
         case ApiKeys.DELETE_RECORDS =>
           new DeleteRecordsRequest.Builder(5000, Map(tp -> (0L: java.lang.Long)).asJava)
 
         case ApiKeys.INIT_PRODUCER_ID =>
-          new InitProducerIdRequest.Builder("abc")
+          val requestData = new InitProducerIdRequestData()
+            .setTransactionalId("test-transactional-id")
+            .setTransactionTimeoutMs(5000)
+          new InitProducerIdRequest.Builder(requestData)
 
         case ApiKeys.OFFSET_FOR_LEADER_EPOCH =>
-          new OffsetsForLeaderEpochRequest.Builder(ApiKeys.OFFSET_FOR_LEADER_EPOCH.latestVersion,
-            Map(tp -> new OffsetsForLeaderEpochRequest.PartitionData(Optional.of(15), 0)).asJava)
+          OffsetsForLeaderEpochRequest.Builder.forConsumer(Map(tp ->
+            new OffsetsForLeaderEpochRequest.PartitionData(Optional.of(15), 0)).asJava)
 
         case ApiKeys.ADD_PARTITIONS_TO_TXN =>
           new AddPartitionsToTxnRequest.Builder("test-transactional-id", 1, 0, List(tp).asJava)
@@ -316,8 +400,16 @@ class RequestQuotaTest extends BaseRequestTest {
           new WriteTxnMarkersRequest.Builder(List.empty.asJava)
 
         case ApiKeys.TXN_OFFSET_COMMIT =>
-          new TxnOffsetCommitRequest.Builder("test-transactional-id", "test-txn-group", 2, 0,
-            Map.empty[TopicPartition, TxnOffsetCommitRequest.CommittedOffset].asJava)
+          new TxnOffsetCommitRequest.Builder(
+            new TxnOffsetCommitRequestData()
+              .setTransactionalId("test-transactional-id")
+              .setGroupId("test-txn-group")
+              .setProducerId(2)
+              .setProducerEpoch(0)
+              .setTopics(TxnOffsetCommitRequest.getTopics(
+                Map.empty[TopicPartition, TxnOffsetCommitRequest.CommittedOffset].asJava
+              ))
+          )
 
         case ApiKeys.DESCRIBE_ACLS =>
           new DescribeAclsRequest.Builder(AclBindingFilter.ANY)
@@ -354,28 +446,64 @@ class RequestQuotaTest extends BaseRequestTest {
           )
 
         case ApiKeys.CREATE_DELEGATION_TOKEN =>
-          new CreateDelegationTokenRequest.Builder(Collections.singletonList(SecurityUtils.parseKafkaPrincipal("User:test")), 1000)
+          new CreateDelegationTokenRequest.Builder(
+              new CreateDelegationTokenRequestData()
+                .setRenewers(Collections.singletonList(new CreateDelegationTokenRequestData.CreatableRenewers()
+                .setPrincipalType("User")
+                .setPrincipalName("test")))
+                .setMaxLifetimeMs(1000)
+          )
 
         case ApiKeys.EXPIRE_DELEGATION_TOKEN =>
-          new ExpireDelegationTokenRequest.Builder("".getBytes, 1000)
+          new ExpireDelegationTokenRequest.Builder(
+              new ExpireDelegationTokenRequestData()
+                .setHmac("".getBytes)
+                .setExpiryTimePeriodMs(1000L))
 
         case ApiKeys.DESCRIBE_DELEGATION_TOKEN =>
           new DescribeDelegationTokenRequest.Builder(Collections.singletonList(SecurityUtils.parseKafkaPrincipal("User:test")))
 
         case ApiKeys.RENEW_DELEGATION_TOKEN =>
-          new RenewDelegationTokenRequest.Builder("".getBytes, 1000)
+          new RenewDelegationTokenRequest.Builder(
+              new RenewDelegationTokenRequestData()
+                .setHmac("".getBytes)
+                .setRenewPeriodMs(1000L))
 
         case ApiKeys.DELETE_GROUPS =>
-          new DeleteGroupsRequest.Builder(Collections.singleton("test-group"))
+          new DeleteGroupsRequest.Builder(new DeleteGroupsRequestData()
+            .setGroupsNames(Collections.singletonList("test-group")))
+
+        case ApiKeys.ELECT_LEADERS =>
+          new ElectLeadersRequest.Builder(
+            ElectionType.PREFERRED,
+            Collections.singletonList(new TopicPartition("my_topic", 0)),
+            0
+          )
 
-        case ApiKeys.ELECT_PREFERRED_LEADERS =>
-          val partition = new ElectPreferredLeadersRequestData.TopicPartitions()
-            .setPartitionId(Collections.singletonList(0))
-            .setTopic("my_topic")
-          new ElectPreferredLeadersRequest.Builder(
-            new ElectPreferredLeadersRequestData()
-                .setTimeoutMs(0)
-                .setTopicPartitions(Collections.singletonList(partition)))
+        case ApiKeys.INCREMENTAL_ALTER_CONFIGS =>
+          new IncrementalAlterConfigsRequest.Builder(
+            new IncrementalAlterConfigsRequestData())
+
+        case ApiKeys.ALTER_PARTITION_REASSIGNMENTS =>
+          new AlterPartitionReassignmentsRequest.Builder(
+            new AlterPartitionReassignmentsRequestData()
+          )
+
+        case ApiKeys.LIST_PARTITION_REASSIGNMENTS =>
+          new ListPartitionReassignmentsRequest.Builder(
+            new ListPartitionReassignmentsRequestData()
+          )
+
+        case ApiKeys.OFFSET_DELETE =>
+          new OffsetDeleteRequest.Builder(
+            new OffsetDeleteRequestData()
+              .setGroupId("test-group")
+              .setTopics(new OffsetDeleteRequestData.OffsetDeleteRequestTopicCollection(
+                Collections.singletonList(new OffsetDeleteRequestData.OffsetDeleteRequestTopic()
+                  .setName("test-topic")
+                  .setPartitions(Collections.singletonList(
+                    new OffsetDeleteRequestData.OffsetDeleteRequestPartition()
+                      .setPartitionIndex(0)))).iterator())))
 
         case _ =>
           throw new IllegalArgumentException("Unsupported API key " + apiKey)
@@ -404,20 +532,23 @@ class RequestQuotaTest extends BaseRequestTest {
     override def toString: String = {
       val requestTime = requestTimeMetricValue(clientId)
       val throttleTime = throttleTimeMetricValue(clientId)
-      s"Client $clientId apiKey $apiKey requests $correlationId requestTime $requestTime throttleTime $throttleTime"
+      val produceThrottleTime = throttleTimeMetricValueForQuotaType(clientId, QuotaType.Produce)
+      val consumeThrottleTime = throttleTimeMetricValueForQuotaType(clientId, QuotaType.Fetch)
+      s"Client $clientId apiKey $apiKey requests $correlationId requestTime $requestTime " +
+      s"throttleTime $throttleTime produceThrottleTime $produceThrottleTime consumeThrottleTime $consumeThrottleTime"
     }
   }
 
-  private def submitTest(apiKey: ApiKeys, test: () => Unit) {
+  private def submitTest(apiKey: ApiKeys, test: () => Unit): Unit = {
     val future = executor.submit(new Runnable() {
-      def run() {
+      def run(): Unit = {
         test.apply()
       }
     })
     tasks += Task(apiKey, future)
   }
 
-  private def waitAndCheckResults() {
+  private def waitAndCheckResults(): Unit = {
     for (task <- tasks) {
       try {
         task.future.get(15, TimeUnit.SECONDS)
@@ -435,27 +566,31 @@ class RequestQuotaTest extends BaseRequestTest {
       case ApiKeys.PRODUCE => new ProduceResponse(response).throttleTimeMs
       case ApiKeys.FETCH => FetchResponse.parse(response).throttleTimeMs
       case ApiKeys.LIST_OFFSETS => new ListOffsetResponse(response).throttleTimeMs
-      case ApiKeys.METADATA => new MetadataResponse(response).throttleTimeMs
-      case ApiKeys.OFFSET_COMMIT => new OffsetCommitResponse(response).throttleTimeMs
-      case ApiKeys.OFFSET_FETCH => new OffsetFetchResponse(response).throttleTimeMs
-      case ApiKeys.FIND_COORDINATOR => new FindCoordinatorResponse(response).throttleTimeMs
+      case ApiKeys.METADATA =>
+        new MetadataResponse(response, ApiKeys.DESCRIBE_GROUPS.latestVersion).throttleTimeMs
+      case ApiKeys.OFFSET_COMMIT =>
+        new OffsetCommitResponse(response, ApiKeys.OFFSET_COMMIT.latestVersion).throttleTimeMs
+      case ApiKeys.OFFSET_FETCH => new OffsetFetchResponse(response, ApiKeys.OFFSET_FETCH.latestVersion).throttleTimeMs
+      case ApiKeys.FIND_COORDINATOR =>
+        new FindCoordinatorResponse(response, ApiKeys.FIND_COORDINATOR.latestVersion).throttleTimeMs
       case ApiKeys.JOIN_GROUP => new JoinGroupResponse(response).throttleTimeMs
-      case ApiKeys.HEARTBEAT => new HeartbeatResponse(response).throttleTimeMs
+      case ApiKeys.HEARTBEAT => new HeartbeatResponse(response, ApiKeys.HEARTBEAT.latestVersion).throttleTimeMs
       case ApiKeys.LEAVE_GROUP => new LeaveGroupResponse(response).throttleTimeMs
       case ApiKeys.SYNC_GROUP => new SyncGroupResponse(response).throttleTimeMs
       case ApiKeys.DESCRIBE_GROUPS =>
-        new DescribeGroupsResponse(response, ApiKeys.DESCRIBE_GROUPS.latestVersion()).throttleTimeMs
-      case ApiKeys.LIST_GROUPS => new ListGroupsResponse(response).throttleTimeMs
+        new DescribeGroupsResponse(response, ApiKeys.DESCRIBE_GROUPS.latestVersion).throttleTimeMs
+      case ApiKeys.LIST_GROUPS => new ListGroupsResponse(response, ApiKeys.LIST_GROUPS.latestVersion).throttleTimeMs
       case ApiKeys.API_VERSIONS => new ApiVersionsResponse(response).throttleTimeMs
       case ApiKeys.CREATE_TOPICS =>
-        new CreateTopicsResponse(response, ApiKeys.CREATE_TOPICS.latestVersion()).throttleTimeMs
-      case ApiKeys.DELETE_TOPICS => new DeleteTopicsResponse(response).throttleTimeMs
+        new CreateTopicsResponse(response, ApiKeys.CREATE_TOPICS.latestVersion).throttleTimeMs
+      case ApiKeys.DELETE_TOPICS =>
+        new DeleteTopicsResponse(response, ApiKeys.DELETE_TOPICS.latestVersion).throttleTimeMs
       case ApiKeys.DELETE_RECORDS => new DeleteRecordsResponse(response).throttleTimeMs
-      case ApiKeys.INIT_PRODUCER_ID => new InitProducerIdResponse(response).throttleTimeMs
+      case ApiKeys.INIT_PRODUCER_ID => new InitProducerIdResponse(response, ApiKeys.INIT_PRODUCER_ID.latestVersion).throttleTimeMs
       case ApiKeys.ADD_PARTITIONS_TO_TXN => new AddPartitionsToTxnResponse(response).throttleTimeMs
       case ApiKeys.ADD_OFFSETS_TO_TXN => new AddOffsetsToTxnResponse(response).throttleTimeMs
       case ApiKeys.END_TXN => new EndTxnResponse(response).throttleTimeMs
-      case ApiKeys.TXN_OFFSET_COMMIT => new TxnOffsetCommitResponse(response).throttleTimeMs
+      case ApiKeys.TXN_OFFSET_COMMIT => new TxnOffsetCommitResponse(response, ApiKeys.TXN_OFFSET_COMMIT.latestVersion).throttleTimeMs
       case ApiKeys.DESCRIBE_ACLS => new DescribeAclsResponse(response).throttleTimeMs
       case ApiKeys.CREATE_ACLS => new CreateAclsResponse(response).throttleTimeMs
       case ApiKeys.DELETE_ACLS => new DeleteAclsResponse(response).throttleTimeMs
@@ -464,29 +599,38 @@ class RequestQuotaTest extends BaseRequestTest {
       case ApiKeys.ALTER_REPLICA_LOG_DIRS => new AlterReplicaLogDirsResponse(response).throttleTimeMs
       case ApiKeys.DESCRIBE_LOG_DIRS => new DescribeLogDirsResponse(response).throttleTimeMs
       case ApiKeys.CREATE_PARTITIONS => new CreatePartitionsResponse(response).throttleTimeMs
-      case ApiKeys.CREATE_DELEGATION_TOKEN => new CreateDelegationTokenResponse(response).throttleTimeMs
-      case ApiKeys.DESCRIBE_DELEGATION_TOKEN=> new DescribeDelegationTokenResponse(response).throttleTimeMs
-      case ApiKeys.EXPIRE_DELEGATION_TOKEN => new ExpireDelegationTokenResponse(response).throttleTimeMs
-      case ApiKeys.RENEW_DELEGATION_TOKEN => new RenewDelegationTokenResponse(response).throttleTimeMs
+      case ApiKeys.CREATE_DELEGATION_TOKEN => new CreateDelegationTokenResponse(response, ApiKeys.CREATE_DELEGATION_TOKEN.latestVersion).throttleTimeMs
+      case ApiKeys.DESCRIBE_DELEGATION_TOKEN=> new DescribeDelegationTokenResponse(response, ApiKeys.DESCRIBE_DELEGATION_TOKEN.latestVersion).throttleTimeMs
+      case ApiKeys.RENEW_DELEGATION_TOKEN => new RenewDelegationTokenResponse(response, ApiKeys.RENEW_DELEGATION_TOKEN.latestVersion).throttleTimeMs
+      case ApiKeys.EXPIRE_DELEGATION_TOKEN => new ExpireDelegationTokenResponse(response, ApiKeys.EXPIRE_DELEGATION_TOKEN.latestVersion).throttleTimeMs
       case ApiKeys.DELETE_GROUPS => new DeleteGroupsResponse(response).throttleTimeMs
       case ApiKeys.OFFSET_FOR_LEADER_EPOCH => new OffsetsForLeaderEpochResponse(response).throttleTimeMs
-      case ApiKeys.ELECT_PREFERRED_LEADERS => new ElectPreferredLeadersResponse(response).throttleTimeMs
+      case ApiKeys.ELECT_LEADERS => new ElectLeadersResponse(response).throttleTimeMs
+      case ApiKeys.INCREMENTAL_ALTER_CONFIGS =>
+        new IncrementalAlterConfigsResponse(response, ApiKeys.INCREMENTAL_ALTER_CONFIGS.latestVersion()).throttleTimeMs
+      case ApiKeys.ALTER_PARTITION_REASSIGNMENTS => new AlterPartitionReassignmentsResponse(response).throttleTimeMs
+      case ApiKeys.LIST_PARTITION_REASSIGNMENTS => new ListPartitionReassignmentsResponse(response).throttleTimeMs
+      case ApiKeys.OFFSET_DELETE => new OffsetDeleteResponse(response).throttleTimeMs()
       case requestId => throw new IllegalArgumentException(s"No throttle time for $requestId")
     }
   }
 
-  private def checkRequestThrottleTime(apiKey: ApiKeys) {
+  private def checkRequestThrottleTime(apiKey: ApiKeys): Unit = {
 
     // Request until throttled using client-id with default small quota
     val clientId = apiKey.toString
     val client = Client(clientId, apiKey)
-    val throttled = client.runUntil(response => responseThrottleTime(apiKey, response) > 0)
+
+    val throttled = client.runUntil(response => {
+      responseThrottleTime(apiKey, response) > 0
+    }
+    )
 
     assertTrue(s"Response not throttled: $client", throttled)
     assertTrue(s"Throttle time metrics not updated: $client" , throttleTimeMetricValue(clientId) > 0)
   }
 
-  private def checkSmallQuotaProducerRequestThrottleTime(apiKey: ApiKeys) {
+  private def checkSmallQuotaProducerRequestThrottleTime(apiKey: ApiKeys): Unit = {
 
     // Request until throttled using client-id with default small producer quota
     val smallQuotaProducerClient = Client(smallQuotaProducerClientId, apiKey)
@@ -499,20 +643,20 @@ class RequestQuotaTest extends BaseRequestTest {
       throttleTimeMetricValueForQuotaType(smallQuotaProducerClientId, QuotaType.Request).isNaN)
   }
 
-  private def checkSmallQuotaConsumerRequestThrottleTime(apiKey: ApiKeys) {
+  private def checkSmallQuotaConsumerRequestThrottleTime(apiKey: ApiKeys): Unit = {
 
     // Request until throttled using client-id with default small consumer quota
     val smallQuotaConsumerClient =   Client(smallQuotaConsumerClientId, apiKey)
     val throttled = smallQuotaConsumerClient.runUntil(response => responseThrottleTime(apiKey, response) > 0)
 
     assertTrue(s"Response not throttled: $smallQuotaConsumerClientId", throttled)
-    assertTrue(s"Throttle time metrics for consumer quota not updated: $smallQuotaConsumerClientId",
+    assertTrue(s"Throttle time metrics for consumer quota not updated: $smallQuotaConsumerClient",
       throttleTimeMetricValueForQuotaType(smallQuotaConsumerClientId, QuotaType.Fetch) > 0)
     assertTrue(s"Throttle time metrics for request quota updated: $smallQuotaConsumerClient",
       throttleTimeMetricValueForQuotaType(smallQuotaConsumerClientId, QuotaType.Request).isNaN)
   }
 
-  private def checkUnthrottledClient(apiKey: ApiKeys) {
+  private def checkUnthrottledClient(apiKey: ApiKeys): Unit = {
 
     // Test that request from client with large quota is not throttled
     val unthrottledClient = Client(unthrottledClientId, apiKey)
@@ -521,7 +665,7 @@ class RequestQuotaTest extends BaseRequestTest {
     assertTrue(s"Client should not have been throttled: $unthrottledClient", throttleTimeMetricValue(unthrottledClientId).isNaN)
   }
 
-  private def checkExemptRequestMetric(apiKey: ApiKeys) {
+  private def checkExemptRequestMetric(apiKey: ApiKeys): Unit = {
     val exemptTarget = exemptRequestMetricValue + 0.02
     val clientId = apiKey.toString
     val client = Client(clientId, apiKey)
@@ -531,7 +675,7 @@ class RequestQuotaTest extends BaseRequestTest {
     assertTrue(s"Client should not have been throttled: $client", throttleTimeMetricValue(clientId).isNaN)
   }
 
-  private def checkUnauthorizedRequestThrottle(apiKey: ApiKeys) {
+  private def checkUnauthorizedRequestThrottle(apiKey: ApiKeys): Unit = {
     val clientId = "unauthorized-" + apiKey.toString
     val client = Client(clientId, apiKey)
     val throttled = client.runUntil(response => throttleTimeMetricValue(clientId) > 0.0)
@@ -548,9 +692,11 @@ object RequestQuotaTest {
   // Principal used for all client connections. This is modified by tests which
   // check unauthorized code path
   var principal = KafkaPrincipal.ANONYMOUS
-  class TestAuthorizer extends SimpleAclAuthorizer {
-    override def authorize(session: Session, operation: Operation, resource: Resource): Boolean = {
-      session.principal != UnauthorizedPrincipal
+  class TestAuthorizer extends AclAuthorizer {
+    override def authorize(requestContext: AuthorizableRequestContext, actions: util.List[Action]): util.List[AuthorizationResult] = {
+      actions.asScala.map { _ =>
+        if (requestContext.principal != UnauthorizedPrincipal) AuthorizationResult.ALLOWED else AuthorizationResult.DENIED
+      }.asJava
     }
   }
   class TestPrincipalBuilder extends KafkaPrincipalBuilder {
diff --git a/core/src/test/scala/unit/kafka/server/SaslApiVersionsRequestTest.scala b/core/src/test/scala/unit/kafka/server/SaslApiVersionsRequestTest.scala
index 185a2f49e6343..b9510d483bd4f 100644
--- a/core/src/test/scala/unit/kafka/server/SaslApiVersionsRequestTest.scala
+++ b/core/src/test/scala/unit/kafka/server/SaslApiVersionsRequestTest.scala
@@ -36,7 +36,7 @@ class SaslApiVersionsRequestTest extends BaseRequestTest with SaslSetup {
   private val kafkaServerSaslMechanisms = List("PLAIN")
   protected override val serverSaslProperties = Some(kafkaServerSaslProperties(kafkaServerSaslMechanisms, kafkaClientSaslMechanism))
   protected override val clientSaslProperties = Some(kafkaClientSaslProperties(kafkaClientSaslMechanism))
-  override def numBrokers = 1
+  override def brokerCount = 1
 
   @Before
   override def setUp(): Unit = {
@@ -51,7 +51,7 @@ class SaslApiVersionsRequestTest extends BaseRequestTest with SaslSetup {
   }
 
   @Test
-  def testApiVersionsRequestBeforeSaslHandshakeRequest() {
+  def testApiVersionsRequestBeforeSaslHandshakeRequest(): Unit = {
     val plaintextSocket = connect(protocol = securityProtocol)
     try {
       val apiVersionsResponse = sendApiVersionsRequest(plaintextSocket, new ApiVersionsRequest.Builder().build(0))
@@ -63,24 +63,24 @@ class SaslApiVersionsRequestTest extends BaseRequestTest with SaslSetup {
   }
 
   @Test
-  def testApiVersionsRequestAfterSaslHandshakeRequest() {
+  def testApiVersionsRequestAfterSaslHandshakeRequest(): Unit = {
     val plaintextSocket = connect(protocol = securityProtocol)
     try {
       sendSaslHandshakeRequestValidateResponse(plaintextSocket)
       val response = sendApiVersionsRequest(plaintextSocket, new ApiVersionsRequest.Builder().build(0))
-      assertEquals(Errors.ILLEGAL_SASL_STATE, response.error)
+      assertEquals(Errors.ILLEGAL_SASL_STATE.code(), response.data.errorCode())
     } finally {
       plaintextSocket.close()
     }
   }
 
   @Test
-  def testApiVersionsRequestWithUnsupportedVersion() {
+  def testApiVersionsRequestWithUnsupportedVersion(): Unit = {
     val plaintextSocket = connect(protocol = securityProtocol)
     try {
-      val apiVersionsRequest = new ApiVersionsRequest(0)
+      val apiVersionsRequest = new ApiVersionsRequest.Builder().build(0)
       val apiVersionsResponse = sendApiVersionsRequest(plaintextSocket, apiVersionsRequest, Some(Short.MaxValue))
-      assertEquals(Errors.UNSUPPORTED_VERSION, apiVersionsResponse.error)
+      assertEquals(Errors.UNSUPPORTED_VERSION.code(), apiVersionsResponse.data.errorCode())
       val apiVersionsResponse2 = sendApiVersionsRequest(plaintextSocket, new ApiVersionsRequest.Builder().build(0))
       ApiVersionsRequestTest.validateApiVersionsResponse(apiVersionsResponse2)
       sendSaslHandshakeRequestValidateResponse(plaintextSocket)
@@ -95,7 +95,7 @@ class SaslApiVersionsRequestTest extends BaseRequestTest with SaslSetup {
     ApiVersionsResponse.parse(response, request.version)
   }
 
-  private def sendSaslHandshakeRequestValidateResponse(socket: Socket) {
+  private def sendSaslHandshakeRequestValidateResponse(socket: Socket): Unit = {
     val request = new SaslHandshakeRequest(new SaslHandshakeRequestData().setMechanism("PLAIN"))
     val response = sendAndReceive(request, ApiKeys.SASL_HANDSHAKE, socket)
     val handshakeResponse = SaslHandshakeResponse.parse(response, request.version)
diff --git a/core/src/test/scala/unit/kafka/server/ServerGenerateBrokerIdTest.scala b/core/src/test/scala/unit/kafka/server/ServerGenerateBrokerIdTest.scala
index 2fa66009739e2..8d0e7bb2d6b15 100755
--- a/core/src/test/scala/unit/kafka/server/ServerGenerateBrokerIdTest.scala
+++ b/core/src/test/scala/unit/kafka/server/ServerGenerateBrokerIdTest.scala
@@ -18,11 +18,14 @@ package kafka.server
 
 import java.util.Properties
 
+import scala.collection.Seq
+
 import kafka.zk.ZooKeeperTestHarness
 import kafka.utils.TestUtils
 import org.junit.{After, Before, Test}
 import org.junit.Assert._
 import java.io.File
+import org.scalatest.Assertions.intercept
 
 import org.apache.zookeeper.KeeperException.NodeExistsException
 
@@ -35,7 +38,7 @@ class ServerGenerateBrokerIdTest extends ZooKeeperTestHarness {
   var servers: Seq[KafkaServer] = Seq()
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     super.setUp()
     props1 = TestUtils.createBrokerConfig(-1, zkConnect)
     config1 = KafkaConfig.fromProps(props1)
@@ -44,32 +47,32 @@ class ServerGenerateBrokerIdTest extends ZooKeeperTestHarness {
   }
 
   @After
-  override def tearDown() {
+  override def tearDown(): Unit = {
     TestUtils.shutdownServers(servers)
     super.tearDown()
   }
 
   @Test
-  def testAutoGenerateBrokerId() {
+  def testAutoGenerateBrokerId(): Unit = {
     var server1 = new KafkaServer(config1, threadNamePrefix = Option(this.getClass.getName))
     server1.startup()
     server1.shutdown()
     assertTrue(verifyBrokerMetadata(config1.logDirs, 1001))
     // restart the server check to see if it uses the brokerId generated previously
-    server1 = TestUtils.createServer(config1)
+    server1 = TestUtils.createServer(config1, threadNamePrefix = Option(this.getClass.getName))
     servers = Seq(server1)
     assertEquals(server1.config.brokerId, 1001)
     server1.shutdown()
-    TestUtils.verifyNonDaemonThreadsStatus(this.getClass.getName)
+    TestUtils.assertNoNonDaemonThreads(this.getClass.getName)
   }
 
   @Test
-  def testUserConfigAndGeneratedBrokerId() {
+  def testUserConfigAndGeneratedBrokerId(): Unit = {
     // start the server with broker.id as part of config
     val server1 = new KafkaServer(config1, threadNamePrefix = Option(this.getClass.getName))
     val server2 = new KafkaServer(config2, threadNamePrefix = Option(this.getClass.getName))
     val props3 = TestUtils.createBrokerConfig(-1, zkConnect)
-    val server3 = new KafkaServer(KafkaConfig.fromProps(props3))
+    val server3 = new KafkaServer(KafkaConfig.fromProps(props3), threadNamePrefix = Option(this.getClass.getName))
     server1.startup()
     assertEquals(server1.config.brokerId, 1001)
     server2.startup()
@@ -81,26 +84,26 @@ class ServerGenerateBrokerIdTest extends ZooKeeperTestHarness {
     assertTrue(verifyBrokerMetadata(server1.config.logDirs, 1001))
     assertTrue(verifyBrokerMetadata(server2.config.logDirs, 0))
     assertTrue(verifyBrokerMetadata(server3.config.logDirs, 1002))
-    TestUtils.verifyNonDaemonThreadsStatus(this.getClass.getName)
+    TestUtils.assertNoNonDaemonThreads(this.getClass.getName)
   }
 
   @Test
-  def testDisableGeneratedBrokerId() {
+  def testDisableGeneratedBrokerId(): Unit = {
     val props3 = TestUtils.createBrokerConfig(3, zkConnect)
     props3.put(KafkaConfig.BrokerIdGenerationEnableProp, "false")
     // Set reserve broker ids to cause collision and ensure disabling broker id generation ignores the setting
     props3.put(KafkaConfig.MaxReservedBrokerIdProp, "0")
     val config3 = KafkaConfig.fromProps(props3)
-    val server3 = TestUtils.createServer(config3)
+    val server3 = TestUtils.createServer(config3, threadNamePrefix = Option(this.getClass.getName))
     servers = Seq(server3)
     assertEquals(server3.config.brokerId, 3)
     server3.shutdown()
     assertTrue(verifyBrokerMetadata(server3.config.logDirs, 3))
-    TestUtils.verifyNonDaemonThreadsStatus(this.getClass.getName)
+    TestUtils.assertNoNonDaemonThreads(this.getClass.getName)
   }
 
   @Test
-  def testMultipleLogDirsMetaProps() {
+  def testMultipleLogDirsMetaProps(): Unit = {
     // add multiple logDirs and check if the generate brokerId is stored in all of them
     val logDirs = props1.getProperty("log.dir")+ "," + TestUtils.tempDir().getAbsolutePath +
     "," + TestUtils.tempDir().getAbsolutePath
@@ -120,11 +123,11 @@ class ServerGenerateBrokerIdTest extends ZooKeeperTestHarness {
     servers = Seq(server1)
     server1.shutdown()
     assertTrue(verifyBrokerMetadata(config1.logDirs, 1001))
-    TestUtils.verifyNonDaemonThreadsStatus(this.getClass.getName)
+    TestUtils.assertNoNonDaemonThreads(this.getClass.getName)
   }
 
   @Test
-  def testConsistentBrokerIdFromUserConfigAndMetaProps() {
+  def testConsistentBrokerIdFromUserConfigAndMetaProps(): Unit = {
     // check if configured brokerId and stored brokerId are equal or throw InconsistentBrokerException
     var server1 = new KafkaServer(config1, threadNamePrefix = Option(this.getClass.getName)) //auto generate broker Id
     server1.startup()
@@ -137,20 +140,20 @@ class ServerGenerateBrokerIdTest extends ZooKeeperTestHarness {
       case _: kafka.common.InconsistentBrokerIdException => //success
     }
     server1.shutdown()
-    TestUtils.verifyNonDaemonThreadsStatus(this.getClass.getName)
+    TestUtils.assertNoNonDaemonThreads(this.getClass.getName)
   }
 
   @Test
-  def testBrokerMetadataOnIdCollision() {
+  def testBrokerMetadataOnIdCollision(): Unit = {
     // Start a good server
     val propsA = TestUtils.createBrokerConfig(1, zkConnect)
     val configA = KafkaConfig.fromProps(propsA)
-    val serverA = TestUtils.createServer(configA)
+    val serverA = TestUtils.createServer(configA, threadNamePrefix = Option(this.getClass.getName))
 
     // Start a server that collides on the broker id
     val propsB = TestUtils.createBrokerConfig(1, zkConnect)
     val configB = KafkaConfig.fromProps(propsB)
-    val serverB = new KafkaServer(configB)
+    val serverB = new KafkaServer(configB, threadNamePrefix = Option(this.getClass.getName))
     intercept[NodeExistsException] {
       serverB.startup()
     }
@@ -165,7 +168,7 @@ class ServerGenerateBrokerIdTest extends ZooKeeperTestHarness {
     // adjust the broker config and start again
     propsB.setProperty(KafkaConfig.BrokerIdProp, "2")
     val newConfigB = KafkaConfig.fromProps(propsB)
-    val newServerB = TestUtils.createServer(newConfigB)
+    val newServerB = TestUtils.createServer(newConfigB, threadNamePrefix = Option(this.getClass.getName))
     servers = Seq(serverA, newServerB)
 
     serverA.shutdown()
@@ -174,7 +177,7 @@ class ServerGenerateBrokerIdTest extends ZooKeeperTestHarness {
     // verify correct broker metadata was written
     assertTrue(verifyBrokerMetadata(serverA.config.logDirs, 1))
     assertTrue(verifyBrokerMetadata(newServerB.config.logDirs, 2))
-    TestUtils.verifyNonDaemonThreadsStatus(this.getClass.getName)
+    TestUtils.assertNoNonDaemonThreads(this.getClass.getName)
   }
 
   def verifyBrokerMetadata(logDirs: Seq[String], brokerId: Int): Boolean = {
diff --git a/core/src/test/scala/unit/kafka/server/ServerGenerateClusterIdTest.scala b/core/src/test/scala/unit/kafka/server/ServerGenerateClusterIdTest.scala
index e00e6c13cd0c6..66d95fce63c03 100755
--- a/core/src/test/scala/unit/kafka/server/ServerGenerateClusterIdTest.scala
+++ b/core/src/test/scala/unit/kafka/server/ServerGenerateClusterIdTest.scala
@@ -16,23 +16,31 @@
   */
 package kafka.server
 
+import java.io.File
+
+import kafka.common.{InconsistentBrokerMetadataException, InconsistentClusterIdException}
+
 import scala.concurrent._
 import ExecutionContext.Implicits._
 import scala.concurrent.duration._
 import kafka.utils.TestUtils
 import kafka.zk.ZooKeeperTestHarness
 import org.junit.Assert._
-import org.junit.{Before, After, Test}
+import org.junit.{After, Before, Test}
+import org.scalatest.Assertions.assertThrows
 import org.apache.kafka.test.TestUtils.isValidClusterId
 
+import scala.collection.Seq
+
 class ServerGenerateClusterIdTest extends ZooKeeperTestHarness {
   var config1: KafkaConfig = null
   var config2: KafkaConfig = null
   var config3: KafkaConfig = null
   var servers: Seq[KafkaServer] = Seq()
+  val brokerMetaPropsFile = "meta.properties"
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     super.setUp()
     config1 = KafkaConfig.fromProps(TestUtils.createBrokerConfig(1, zkConnect))
     config2 = KafkaConfig.fromProps(TestUtils.createBrokerConfig(2, zkConnect))
@@ -40,18 +48,18 @@ class ServerGenerateClusterIdTest extends ZooKeeperTestHarness {
   }
 
   @After
-  override def tearDown() {
+  override def tearDown(): Unit = {
     TestUtils.shutdownServers(servers)
     super.tearDown()
   }
 
 
   @Test
-  def testAutoGenerateClusterId() {
+  def testAutoGenerateClusterId(): Unit = {
     // Make sure that the cluster id doesn't exist yet.
     assertFalse(zkClient.getClusterId.isDefined)
 
-    var server1 = TestUtils.createServer(config1)
+    var server1 = TestUtils.createServer(config1, threadNamePrefix = Option(this.getClass.getName))
     servers = Seq(server1)
 
     // Validate the cluster id
@@ -65,7 +73,7 @@ class ServerGenerateClusterIdTest extends ZooKeeperTestHarness {
     assertEquals(zkClient.getClusterId, Some(clusterIdOnFirstBoot))
 
     // Restart the server check to confirm that it uses the clusterId generated previously
-    server1 = TestUtils.createServer(config1)
+    server1 = TestUtils.createServer(config1, threadNamePrefix = Option(this.getClass.getName))
     servers = Seq(server1)
 
     val clusterIdOnSecondBoot = server1.clusterId
@@ -77,18 +85,18 @@ class ServerGenerateClusterIdTest extends ZooKeeperTestHarness {
     assertTrue(zkClient.getClusterId.isDefined)
     assertEquals(zkClient.getClusterId, Some(clusterIdOnFirstBoot))
 
-    TestUtils.verifyNonDaemonThreadsStatus(this.getClass.getName)
+    TestUtils.assertNoNonDaemonThreads(this.getClass.getName)
   }
 
   @Test
-  def testAutoGenerateClusterIdForKafkaClusterSequential() {
-    val server1 = TestUtils.createServer(config1)
+  def testAutoGenerateClusterIdForKafkaClusterSequential(): Unit = {
+    val server1 = TestUtils.createServer(config1, threadNamePrefix = Option(this.getClass.getName))
     val clusterIdFromServer1 = server1.clusterId
 
-    val server2 = TestUtils.createServer(config2)
+    val server2 = TestUtils.createServer(config2, threadNamePrefix = Option(this.getClass.getName))
     val clusterIdFromServer2 = server2.clusterId
 
-    val server3 = TestUtils.createServer(config3)
+    val server3 = TestUtils.createServer(config3, threadNamePrefix = Option(this.getClass.getName))
     val clusterIdFromServer3 = server3.clusterId
     servers = Seq(server1, server2, server3)
 
@@ -107,12 +115,12 @@ class ServerGenerateClusterIdTest extends ZooKeeperTestHarness {
 
     servers.foreach(_.shutdown())
 
-    TestUtils.verifyNonDaemonThreadsStatus(this.getClass.getName)
+    TestUtils.assertNoNonDaemonThreads(this.getClass.getName)
   }
 
   @Test
-  def testAutoGenerateClusterIdForKafkaClusterParallel() {
-    val firstBoot = Future.traverse(Seq(config1, config2, config3))(config => Future(TestUtils.createServer(config)))
+  def testAutoGenerateClusterIdForKafkaClusterParallel(): Unit = {
+    val firstBoot = Future.traverse(Seq(config1, config2, config3))(config => Future(TestUtils.createServer(config, threadNamePrefix = Option(this.getClass.getName))))
     servers = Await.result(firstBoot, 100 second)
     val Seq(server1, server2, server3) = servers
 
@@ -134,7 +142,94 @@ class ServerGenerateClusterIdTest extends ZooKeeperTestHarness {
 
     servers.foreach(_.shutdown())
 
-    TestUtils.verifyNonDaemonThreadsStatus(this.getClass.getName)
+    TestUtils.assertNoNonDaemonThreads(this.getClass.getName)
+  }
+
+  @Test
+  def testConsistentClusterIdFromZookeeperAndFromMetaProps() = {
+    // Check at the first boot
+    val server = TestUtils.createServer(config1, threadNamePrefix = Option(this.getClass.getName))
+    val clusterId = server.clusterId
+
+    assertTrue(verifyBrokerMetadata(server.config.logDirs, clusterId))
+
+    server.shutdown()
+
+    // Check again after reboot
+    server.startup()
+
+    assertEquals(clusterId, server.clusterId)
+    assertTrue(verifyBrokerMetadata(server.config.logDirs, server.clusterId))
+
+    server.shutdown()
+
+    TestUtils.assertNoNonDaemonThreads(this.getClass.getName)
+  }
+
+  @Test
+  def testInconsistentClusterIdFromZookeeperAndFromMetaProps() = {
+    forgeBrokerMetadata(config1.logDirs, config1.brokerId, "aclusterid")
+
+    val server = new KafkaServer(config1, threadNamePrefix = Option(this.getClass.getName))
+
+    // Startup fails
+    assertThrows[InconsistentClusterIdException] {
+      server.startup()
+    }
+
+    server.shutdown()
+
+    TestUtils.assertNoNonDaemonThreads(this.getClass.getName)
+  }
+
+  @Test
+  def testInconsistentBrokerMetadataBetweenMultipleLogDirs(): Unit = {
+    // Add multiple logDirs with different BrokerMetadata
+    val logDir1 = TestUtils.tempDir().getAbsolutePath
+    val logDir2 = TestUtils.tempDir().getAbsolutePath
+    val logDirs = logDir1 + "," + logDir2
+
+    forgeBrokerMetadata(logDir1, 1, "ebwOKU-zSieInaFQh_qP4g")
+    forgeBrokerMetadata(logDir2, 1, "blaOKU-zSieInaFQh_qP4g")
+
+    val props = TestUtils.createBrokerConfig(1, zkConnect)
+    props.setProperty("log.dir", logDirs)
+    val config = KafkaConfig.fromProps(props)
+
+    val server = new KafkaServer(config, threadNamePrefix = Option(this.getClass.getName))
+
+    // Startup fails
+    assertThrows[InconsistentBrokerMetadataException] {
+      server.startup()
+    }
+
+    server.shutdown()
+
+    TestUtils.assertNoNonDaemonThreads(this.getClass.getName)
   }
 
+  def forgeBrokerMetadata(logDirs: Seq[String], brokerId: Int, clusterId: String): Unit = {
+    for (logDir <- logDirs) {
+      forgeBrokerMetadata(logDir, brokerId, clusterId)
+    }
+  }
+
+  def forgeBrokerMetadata(logDir: String, brokerId: Int, clusterId: String): Unit = {
+    val checkpoint = new BrokerMetadataCheckpoint(
+      new File(logDir + File.separator + brokerMetaPropsFile))
+    checkpoint.write(BrokerMetadata(brokerId, Option(clusterId)))
+  }
+
+  def verifyBrokerMetadata(logDirs: Seq[String], clusterId: String): Boolean = {
+    for (logDir <- logDirs) {
+      val brokerMetadataOpt = new BrokerMetadataCheckpoint(
+        new File(logDir + File.separator + brokerMetaPropsFile)).read()
+      brokerMetadataOpt match {
+        case Some(brokerMetadata) =>
+          if (brokerMetadata.clusterId.isDefined && brokerMetadata.clusterId.get != clusterId) return false
+        case _ => return false
+      }
+    }
+    true
+  }
 }
diff --git a/core/src/test/scala/unit/kafka/server/ServerMetricsTest.scala b/core/src/test/scala/unit/kafka/server/ServerMetricsTest.scala
index dc96680c95893..1883eddb1d72b 100755
--- a/core/src/test/scala/unit/kafka/server/ServerMetricsTest.scala
+++ b/core/src/test/scala/unit/kafka/server/ServerMetricsTest.scala
@@ -21,9 +21,9 @@ import kafka.utils.TestUtils
 import org.apache.kafka.common.metrics.Sensor
 import org.junit.Assert._
 import org.junit.Test
-import org.scalatest.junit.JUnitSuite
+import org.scalatest.Assertions.intercept
 
-class ServerMetricsTest extends JUnitSuite {
+class ServerMetricsTest {
 
   @Test
   def testMetricsConfig(): Unit = {
diff --git a/core/src/test/scala/unit/kafka/server/ServerShutdownTest.scala b/core/src/test/scala/unit/kafka/server/ServerShutdownTest.scala
index 7dddaabb5ae60..75fae948d6fc1 100755
--- a/core/src/test/scala/unit/kafka/server/ServerShutdownTest.scala
+++ b/core/src/test/scala/unit/kafka/server/ServerShutdownTest.scala
@@ -51,14 +51,14 @@ class ServerShutdownTest extends ZooKeeperTestHarness {
   val sent2 = List("more", "messages")
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     super.setUp()
     val props = TestUtils.createBrokerConfig(0, zkConnect)
     config = KafkaConfig.fromProps(props)
   }
 
   @Test
-  def testCleanShutdown() {
+  def testCleanShutdown(): Unit = {
 
     def createProducer(server: KafkaServer): KafkaProducer[Integer, String] =
       TestUtils.createProducer(
@@ -122,7 +122,7 @@ class ServerShutdownTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  def testCleanShutdownWithDeleteTopicEnabled() {
+  def testCleanShutdownWithDeleteTopicEnabled(): Unit = {
     val newProps = TestUtils.createBrokerConfig(0, zkConnect)
     newProps.setProperty("delete.topic.enable", "true")
     val newConfig = KafkaConfig.fromProps(newProps)
@@ -135,7 +135,7 @@ class ServerShutdownTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  def testCleanShutdownAfterFailedStartup() {
+  def testCleanShutdownAfterFailedStartup(): Unit = {
     val newProps = TestUtils.createBrokerConfig(0, zkConnect)
     newProps.setProperty(KafkaConfig.ZkConnectionTimeoutMsProp, "50")
     newProps.setProperty(KafkaConfig.ZkConnectProp, "some.invalid.hostname.foo.bar.local:65535")
@@ -144,7 +144,7 @@ class ServerShutdownTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  def testCleanShutdownAfterFailedStartupDueToCorruptLogs() {
+  def testCleanShutdownAfterFailedStartupDueToCorruptLogs(): Unit = {
     val server = new KafkaServer(config)
     server.startup()
     createTopic(zkClient, topic, numPartitions = 1, replicationFactor = 1, servers = Seq(server))
@@ -157,7 +157,7 @@ class ServerShutdownTest extends ZooKeeperTestHarness {
     verifyCleanShutdownAfterFailedStartup[KafkaStorageException](config)
   }
 
-  private def verifyCleanShutdownAfterFailedStartup[E <: Exception](config: KafkaConfig)(implicit exceptionClassTag: ClassTag[E]) {
+  private def verifyCleanShutdownAfterFailedStartup[E <: Exception](config: KafkaConfig)(implicit exceptionClassTag: ClassTag[E]): Unit = {
     val server = new KafkaServer(config, threadNamePrefix = Option(this.getClass.getName))
     try {
       server.startup()
@@ -184,14 +184,14 @@ class ServerShutdownTest extends ZooKeeperTestHarness {
     !t.isDaemon && t.isAlive && t.getName.startsWith(this.getClass.getName)
   }
 
-  def verifyNonDaemonThreadsStatus() {
+  def verifyNonDaemonThreadsStatus(): Unit = {
     assertEquals(0, Thread.getAllStackTraces.keySet.toArray
       .map(_.asInstanceOf[Thread])
       .count(isNonDaemonKafkaThread))
   }
 
   @Test
-  def testConsecutiveShutdown(){
+  def testConsecutiveShutdown(): Unit = {
     val server = new KafkaServer(config)
     server.startup()
     server.shutdown()
@@ -233,8 +233,8 @@ class ServerShutdownTest extends ZooKeeperTestHarness {
 
       // Initiate a sendRequest and wait until connection is established and one byte is received by the peer
       val requestBuilder = new LeaderAndIsrRequest.Builder(ApiKeys.LEADER_AND_ISR.latestVersion,
-        controllerId, 1, 0L, Map.empty.asJava, brokerAndEpochs.keys.map(_.node(listenerName)).toSet.asJava)
-      controllerChannelManager.sendRequest(1, ApiKeys.LEADER_AND_ISR, requestBuilder)
+        controllerId, 1, 0L, 0L, Seq.empty.asJava, brokerAndEpochs.keys.map(_.node(listenerName)).toSet.asJava)
+      controllerChannelManager.sendRequest(1, requestBuilder)
       receiveFuture.get(10, TimeUnit.SECONDS)
 
       // Shutdown controller. Request timeout is 30s, verify that shutdown completed well before that
diff --git a/core/src/test/scala/unit/kafka/server/ServerStartupTest.scala b/core/src/test/scala/unit/kafka/server/ServerStartupTest.scala
index 1bc025786367d..39795809587db 100755
--- a/core/src/test/scala/unit/kafka/server/ServerStartupTest.scala
+++ b/core/src/test/scala/unit/kafka/server/ServerStartupTest.scala
@@ -30,7 +30,7 @@ class ServerStartupTest extends ZooKeeperTestHarness {
   private var server: KafkaServer = null
 
   @After
-  override def tearDown() {
+  override def tearDown(): Unit = {
     if (server != null)
       TestUtils.shutdownServers(Seq(server))
     super.tearDown()
diff --git a/core/src/test/scala/unit/kafka/server/SimpleFetchTest.scala b/core/src/test/scala/unit/kafka/server/SimpleFetchTest.scala
index 94f9a16341323..26cd94158ace6 100644
--- a/core/src/test/scala/unit/kafka/server/SimpleFetchTest.scala
+++ b/core/src/test/scala/unit/kafka/server/SimpleFetchTest.scala
@@ -20,8 +20,8 @@ import java.io.File
 
 import kafka.api._
 import kafka.utils._
-import kafka.cluster.Replica
-import kafka.log.{Log, LogManager}
+import kafka.log.Log
+import kafka.log.LogManager
 import kafka.server.QuotaFactory.UnboundedQuota
 import kafka.zk.KafkaZkClient
 import org.apache.kafka.common.metrics.Metrics
@@ -68,7 +68,7 @@ class SimpleFetchTest {
   var replicaManager: ReplicaManager = _
 
   @Before
-  def setUp() {
+  def setUp(): Unit = {
     // create nice mock since we don't particularly care about zkclient calls
     val kafkaZkClient: KafkaZkClient = EasyMock.createNiceMock(classOf[KafkaZkClient])
     EasyMock.replay(kafkaZkClient)
@@ -82,25 +82,27 @@ class SimpleFetchTest {
     EasyMock.expect(log.logStartOffset).andReturn(0).anyTimes()
     EasyMock.expect(log.logEndOffset).andReturn(leaderLEO).anyTimes()
     EasyMock.expect(log.dir).andReturn(TestUtils.tempDir()).anyTimes()
-    EasyMock.expect(log.logEndOffsetMetadata).andReturn(new LogOffsetMetadata(leaderLEO)).anyTimes()
+    EasyMock.expect(log.logEndOffsetMetadata).andReturn(LogOffsetMetadata(leaderLEO)).anyTimes()
+    EasyMock.expect(log.maybeIncrementHighWatermark(EasyMock.anyObject[LogOffsetMetadata]))
+      .andReturn(Some(LogOffsetMetadata(partitionHW))).anyTimes()
+    EasyMock.expect(log.highWatermark).andReturn(partitionHW).anyTimes()
+    EasyMock.expect(log.lastStableOffset).andReturn(partitionHW).anyTimes()
     EasyMock.expect(log.read(
       startOffset = 0,
       maxLength = fetchSize,
-      maxOffset = Some(partitionHW),
-      minOneMessage = true,
-      includeAbortedTxns = false))
+      isolation = FetchHighWatermark,
+      minOneMessage = true))
       .andReturn(FetchDataInfo(
-        new LogOffsetMetadata(0L, 0L, 0),
+        LogOffsetMetadata(0L, 0L, 0),
         MemoryRecords.withRecords(CompressionType.NONE, recordToHW)
       )).anyTimes()
     EasyMock.expect(log.read(
       startOffset = 0,
       maxLength = fetchSize,
-      maxOffset = None,
-      minOneMessage = true,
-      includeAbortedTxns = false))
+      isolation = FetchLogEnd,
+      minOneMessage = true))
       .andReturn(FetchDataInfo(
-        new LogOffsetMetadata(0L, 0L, 0),
+        LogOffsetMetadata(0L, 0L, 0),
         MemoryRecords.withRecords(CompressionType.NONE, recordToLEO)
       )).anyTimes()
     EasyMock.replay(log)
@@ -117,33 +119,32 @@ class SimpleFetchTest {
       new MetadataCache(configs.head.brokerId), new LogDirFailureChannel(configs.head.logDirs.size))
 
     // add the partition with two replicas, both in ISR
-    val partition = replicaManager.getOrCreatePartition(new TopicPartition(topic, partitionId))
+    val partition = replicaManager.createPartition(new TopicPartition(topic, partitionId))
 
     // create the leader replica with the local log
-    val leaderReplica = new Replica(configs.head.brokerId, partition.topicPartition, time, 0, Some(log))
-    leaderReplica.highWatermark = new LogOffsetMetadata(partitionHW)
-    partition.leaderReplicaIdOpt = Some(leaderReplica.brokerId)
+    log.updateHighWatermark(partitionHW)
+    partition.leaderReplicaIdOpt = Some(configs.head.brokerId)
+    partition.setLog(log, false)
 
     // create the follower replica with defined log end offset
-    val followerReplica= new Replica(configs(1).brokerId, partition.topicPartition, time)
-    val leo = new LogOffsetMetadata(followerLEO, 0L, followerLEO.toInt)
-    followerReplica.updateLogReadResult(new LogReadResult(info = FetchDataInfo(leo, MemoryRecords.EMPTY),
-                                                          highWatermark = leo.messageOffset,
-                                                          leaderLogStartOffset = 0L,
-                                                          leaderLogEndOffset = leo.messageOffset,
-                                                          followerLogStartOffset = 0L,
-                                                          fetchTimeMs = time.milliseconds,
-                                                          readSize = -1,
-                                                          lastStableOffset = None))
-
-    // add both of them to ISR
-    val allReplicas = List(leaderReplica, followerReplica)
-    allReplicas.foreach(partition.addReplicaIfNotExists)
-    partition.inSyncReplicas = allReplicas.toSet
+    val followerId = configs(1).brokerId
+    val allReplicas = Seq(configs.head.brokerId, followerId)
+    partition.updateAssignmentAndIsr(
+      assignment = allReplicas,
+      isr = allReplicas.toSet
+    )
+    val leo = LogOffsetMetadata(followerLEO, 0L, followerLEO.toInt)
+    partition.updateFollowerFetchState(
+      followerId,
+      followerFetchOffsetMetadata = leo,
+      followerStartOffset = 0L,
+      followerFetchTimeMs= time.milliseconds,
+      leaderEndOffset = leo.messageOffset,
+      partition.localLogOrException.highWatermark)
   }
 
   @After
-  def tearDown() {
+  def tearDown(): Unit = {
     replicaManager.shutdown(false)
     metrics.close()
   }
@@ -165,7 +166,7 @@ class SimpleFetchTest {
    * This test also verifies counts of fetch requests recorded by the ReplicaManager
    */
   @Test
-  def testReadFromLog() {
+  def testReadFromLog(): Unit = {
     val brokerTopicStats = new BrokerTopicStats
     val initialTopicCount = brokerTopicStats.topicStats(topic).totalFetchRequestRate.count()
     val initialAllTopicsCount = brokerTopicStats.allTopicsStats.totalFetchRequestRate.count()
@@ -177,7 +178,8 @@ class SimpleFetchTest {
       fetchMaxBytes = Int.MaxValue,
       hardMaxBytesLimit = false,
       readPartitionInfo = fetchInfo,
-      quota = UnboundedQuota).find(_._1 == topicPartition)
+      quota = UnboundedQuota,
+      clientMetadata = None).find(_._1 == topicPartition)
     val firstReadRecord = readCommittedRecords.get._2.info.records.records.iterator.next()
     assertEquals("Reading committed data should return messages only up to high watermark", recordToHW,
       new SimpleRecord(firstReadRecord))
@@ -189,7 +191,8 @@ class SimpleFetchTest {
       fetchMaxBytes = Int.MaxValue,
       hardMaxBytesLimit = false,
       readPartitionInfo = fetchInfo,
-      quota = UnboundedQuota).find(_._1 == topicPartition)
+      quota = UnboundedQuota,
+      clientMetadata = None).find(_._1 == topicPartition)
 
     val firstRecord = readAllRecords.get._2.info.records.records.iterator.next()
     assertEquals("Reading any data can return messages up to the end of the log", recordToLEO,
diff --git a/core/src/test/scala/unit/kafka/server/StopReplicaRequestTest.scala b/core/src/test/scala/unit/kafka/server/StopReplicaRequestTest.scala
index 2d0f1db372755..1af4bd8ddb449 100644
--- a/core/src/test/scala/unit/kafka/server/StopReplicaRequestTest.scala
+++ b/core/src/test/scala/unit/kafka/server/StopReplicaRequestTest.scala
@@ -25,10 +25,9 @@ import org.junit.Assert._
 import org.junit.Test
 import collection.JavaConverters._
 
-
 class StopReplicaRequestTest extends BaseRequestTest {
   override val logDirCount = 2
-  override val numBrokers: Int = 1
+  override val brokerCount: Int = 1
 
   val topic = "topic"
   val partitionNum = 2
@@ -44,14 +43,17 @@ class StopReplicaRequestTest extends BaseRequestTest {
     val offlineDir = server.logManager.getLog(tp1).get.dir.getParent
     server.replicaManager.handleLogDirFailure(offlineDir, sendZkNotification = false)
 
-    for (i <- 1 to 2) {
+    for (_ <- 1 to 2) {
       val request1 = new StopReplicaRequest.Builder(1,
         server.config.brokerId, server.replicaManager.controllerEpoch, server.kafkaController.brokerEpoch,
+        server.kafkaController.brokerEpoch,
         true, Set(tp0, tp1).asJava).build()
       val response1 = connectAndSend(request1, ApiKeys.STOP_REPLICA, controllerSocketServer)
-      val partitionErrors1 = StopReplicaResponse.parse(response1, request1.version).responses()
-      assertEquals(Errors.NONE, partitionErrors1.get(tp0))
-      assertEquals(Errors.KAFKA_STORAGE_ERROR, partitionErrors1.get(tp1))
+      val partitionErrors1 = StopReplicaResponse.parse(response1, request1.version).partitionErrors.asScala
+      assertEquals(Some(Errors.NONE.code),
+        partitionErrors1.find(pe => pe.topicName == tp0.topic && pe.partitionIndex == tp0.partition).map(_.errorCode))
+      assertEquals(Some(Errors.KAFKA_STORAGE_ERROR.code),
+        partitionErrors1.find(pe => pe.topicName == tp1.topic && pe.partitionIndex == tp1.partition).map(_.errorCode))
     }
   }
 
diff --git a/core/src/test/scala/unit/kafka/server/ThrottledChannelExpirationTest.scala b/core/src/test/scala/unit/kafka/server/ThrottledChannelExpirationTest.scala
index c46404addca25..746ee93b392e2 100644
--- a/core/src/test/scala/unit/kafka/server/ThrottledChannelExpirationTest.scala
+++ b/core/src/test/scala/unit/kafka/server/ThrottledChannelExpirationTest.scala
@@ -68,14 +68,14 @@ class ThrottledChannelExpirationTest {
   }
 
   @Before
-  def beforeMethod() {
+  def beforeMethod(): Unit = {
     numCallbacksForStartThrottling = 0
     numCallbacksForEndThrottling = 0
   }
 
   @Test
-  def testCallbackInvocationAfterExpiration() {
-    val clientMetrics = new ClientQuotaManager(ClientQuotaManagerConfig(), metrics, QuotaType.Produce, time, "")
+  def testCallbackInvocationAfterExpiration(): Unit = {
+    val clientMetrics = new ClientQuotaManager(ClientQuotaManagerConfig(), metrics, QuotaType.Produce, time, None, "")
 
     val delayQueue = new DelayQueue[ThrottledChannel]()
     val reaper = new clientMetrics.ThrottledChannelReaper(delayQueue, "")
@@ -107,7 +107,7 @@ class ThrottledChannelExpirationTest {
   }
 
   @Test
-  def testThrottledChannelDelay() {
+  def testThrottledChannelDelay(): Unit = {
     val t1: ThrottledChannel = new ThrottledChannel(request, time, 10, callback)
     val t2: ThrottledChannel = new ThrottledChannel(request, time, 20, callback)
     val t3: ThrottledChannel = new ThrottledChannel(request, time, 20, callback)
diff --git a/core/src/test/scala/unit/kafka/server/checkpoints/LeaderEpochCheckpointFileTest.scala b/core/src/test/scala/unit/kafka/server/checkpoints/LeaderEpochCheckpointFileTest.scala
index 0c47f15a09ff2..b3f90bc506be0 100644
--- a/core/src/test/scala/unit/kafka/server/checkpoints/LeaderEpochCheckpointFileTest.scala
+++ b/core/src/test/scala/unit/kafka/server/checkpoints/LeaderEpochCheckpointFileTest.scala
@@ -22,9 +22,8 @@ import kafka.server.epoch.EpochEntry
 import kafka.utils.Logging
 import org.junit.Assert._
 import org.junit.Test
-import org.scalatest.junit.JUnitSuite
 
-class LeaderEpochCheckpointFileTest extends JUnitSuite with Logging{
+class LeaderEpochCheckpointFileTest extends Logging {
 
   @Test
   def shouldPersistAndOverwriteAndReloadFile(): Unit ={
@@ -68,4 +67,4 @@ class LeaderEpochCheckpointFileTest extends JUnitSuite with Logging{
     //The data should still be there
     assertEquals(epochs, checkpoint2.read())
   }
-}
\ No newline at end of file
+}
diff --git a/core/src/test/scala/unit/kafka/server/checkpoints/OffsetCheckpointFileTest.scala b/core/src/test/scala/unit/kafka/server/checkpoints/OffsetCheckpointFileTest.scala
index f4998f6397bc7..99a40b607d6f5 100644
--- a/core/src/test/scala/unit/kafka/server/checkpoints/OffsetCheckpointFileTest.scala
+++ b/core/src/test/scala/unit/kafka/server/checkpoints/OffsetCheckpointFileTest.scala
@@ -22,11 +22,12 @@ import org.apache.kafka.common.TopicPartition
 import org.apache.kafka.common.errors.KafkaStorageException
 import org.junit.Assert._
 import org.junit.Test
-import org.scalatest.junit.JUnitSuite
+import org.mockito.Mockito
+import org.scalatest.Assertions.assertThrows
 
 import scala.collection.Map
 
-class OffsetCheckpointFileTest extends JUnitSuite with Logging {
+class OffsetCheckpointFileTest extends Logging {
 
   @Test
   def shouldPersistAndOverwriteAndReloadFile(): Unit = {
@@ -99,4 +100,38 @@ class OffsetCheckpointFileTest extends JUnitSuite with Logging {
     new OffsetCheckpointFile(checkpointFile.file, logDirFailureChannel).read()
   }
 
+  @Test
+  def testLazyOffsetCheckpoint(): Unit = {
+    val logDir = "/tmp/kafka-logs"
+    val mockCheckpointFile = Mockito.mock(classOf[OffsetCheckpointFile])
+
+    val lazyCheckpoints = new LazyOffsetCheckpoints(Map(logDir -> mockCheckpointFile))
+    Mockito.verify(mockCheckpointFile, Mockito.never()).read()
+
+    val partition0 = new TopicPartition("foo", 0)
+    val partition1 = new TopicPartition("foo", 1)
+    val partition2 = new TopicPartition("foo", 2)
+
+    Mockito.when(mockCheckpointFile.read()).thenReturn(Map(
+      partition0 -> 1000L,
+      partition1 -> 2000L
+    ))
+
+    assertEquals(Some(1000L), lazyCheckpoints.fetch(logDir, partition0))
+    assertEquals(Some(2000L), lazyCheckpoints.fetch(logDir, partition1))
+    assertEquals(None, lazyCheckpoints.fetch(logDir, partition2))
+
+    Mockito.verify(mockCheckpointFile, Mockito.times(1)).read()
+  }
+
+  @Test
+  def testLazyOffsetCheckpointFileInvalidLogDir(): Unit = {
+    val logDir = "/tmp/kafka-logs"
+    val mockCheckpointFile = Mockito.mock(classOf[OffsetCheckpointFile])
+    val lazyCheckpoints = new LazyOffsetCheckpoints(Map(logDir -> mockCheckpointFile))
+    assertThrows[IllegalArgumentException] {
+      lazyCheckpoints.fetch("/invalid/kafka-logs", new TopicPartition("foo", 0))
+    }
+  }
+
 }
diff --git a/core/src/test/scala/unit/kafka/server/epoch/EpochDrivenReplicationProtocolAcceptanceTest.scala b/core/src/test/scala/unit/kafka/server/epoch/EpochDrivenReplicationProtocolAcceptanceTest.scala
index ac6dedcd3324c..bd9d02668d8fb 100644
--- a/core/src/test/scala/unit/kafka/server/epoch/EpochDrivenReplicationProtocolAcceptanceTest.scala
+++ b/core/src/test/scala/unit/kafka/server/epoch/EpochDrivenReplicationProtocolAcceptanceTest.scala
@@ -60,12 +60,12 @@ class EpochDrivenReplicationProtocolAcceptanceTest extends ZooKeeperTestHarness
   var consumer: KafkaConsumer[Array[Byte], Array[Byte]] = null
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     super.setUp()
   }
 
   @After
-  override def tearDown() {
+  override def tearDown(): Unit = {
     producer.close()
     TestUtils.shutdownServers(brokers)
     super.tearDown()
@@ -367,7 +367,7 @@ class EpochDrivenReplicationProtocolAcceptanceTest extends ZooKeeperTestHarness
     printSegments()
 
     def crcSeq(broker: KafkaServer, partition: Int = 0): Seq[Long] = {
-      val batches = getLog(broker, partition).activeSegment.read(0, None, Integer.MAX_VALUE)
+      val batches = getLog(broker, partition).activeSegment.read(0, Integer.MAX_VALUE)
         .records.batches().asScala.toSeq
       batches.map(_.checksum)
     }
@@ -438,13 +438,13 @@ class EpochDrivenReplicationProtocolAcceptanceTest extends ZooKeeperTestHarness
   private def epochCache(broker: KafkaServer): LeaderEpochFileCache = getLog(broker, 0).leaderEpochCache.get
 
   private def latestRecord(leader: KafkaServer, offset: Int = -1, partition: Int = 0): RecordBatch = {
-    getLog(leader, partition).activeSegment.read(0, None, Integer.MAX_VALUE)
+    getLog(leader, partition).activeSegment.read(0, Integer.MAX_VALUE)
       .records.batches().asScala.toSeq.last
   }
 
   private def awaitISR(tp: TopicPartition): Unit = {
     TestUtils.waitUntilTrue(() => {
-      leader.replicaManager.getPartition(tp).get.inSyncReplicas.map(_.brokerId).size == 2
+      leader.replicaManager.nonOfflinePartition(tp).get.inSyncReplicaIds.size == 2
     }, "Timed out waiting for replicas to join ISR")
   }
 
diff --git a/core/src/test/scala/unit/kafka/server/epoch/LeaderEpochFileCacheTest.scala b/core/src/test/scala/unit/kafka/server/epoch/LeaderEpochFileCacheTest.scala
index e3b96db5631ac..af98b4787f8ce 100644
--- a/core/src/test/scala/unit/kafka/server/epoch/LeaderEpochFileCacheTest.scala
+++ b/core/src/test/scala/unit/kafka/server/epoch/LeaderEpochFileCacheTest.scala
@@ -19,6 +19,9 @@ package kafka.server.epoch
 
 import java.io.File
 
+import scala.collection.Seq
+import scala.collection.mutable.ListBuffer
+
 import kafka.server.checkpoints.{LeaderEpochCheckpoint, LeaderEpochCheckpointFile}
 import org.apache.kafka.common.requests.EpochEndOffset.{UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET}
 import kafka.utils.TestUtils
@@ -26,8 +29,6 @@ import org.apache.kafka.common.TopicPartition
 import org.junit.Assert._
 import org.junit.Test
 
-import scala.collection.mutable.ListBuffer
-
 /**
   * Unit test for the LeaderEpochFileCache.
   */
@@ -119,13 +120,13 @@ class LeaderEpochFileCacheTest {
   }
 
   @Test
-  def shouldReturnUnsupportedIfNoEpochRecorded(){
+  def shouldReturnUnsupportedIfNoEpochRecorded(): Unit = {
     //Then
     assertEquals((UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET), cache.endOffsetFor(0))
   }
 
   @Test
-  def shouldReturnUnsupportedIfNoEpochRecordedAndUndefinedEpochRequested(){
+  def shouldReturnUnsupportedIfNoEpochRecordedAndUndefinedEpochRequested(): Unit = {
     logEndOffset = 73
 
     //When (say a follower on older message format version) sends request for UNDEFINED_EPOCH
@@ -137,7 +138,7 @@ class LeaderEpochFileCacheTest {
   }
 
   @Test
-  def shouldReturnFirstEpochIfRequestedEpochLessThanFirstEpoch(){
+  def shouldReturnFirstEpochIfRequestedEpochLessThanFirstEpoch(): Unit = {
     cache.assign(epoch = 5, startOffset = 11)
     cache.assign(epoch = 6, startOffset = 12)
     cache.assign(epoch = 7, startOffset = 13)
@@ -178,7 +179,7 @@ class LeaderEpochFileCacheTest {
   }
 
   @Test
-  def shouldReturnNextAvailableEpochIfThereIsNoExactEpochForTheOneRequested(){
+  def shouldReturnNextAvailableEpochIfThereIsNoExactEpochForTheOneRequested(): Unit = {
     //When
     cache.assign(epoch = 0, startOffset = 10)
     cache.assign(epoch = 2, startOffset = 13)
@@ -225,7 +226,7 @@ class LeaderEpochFileCacheTest {
   }
 
   @Test
-  def shouldPersistEpochsBetweenInstances(){
+  def shouldPersistEpochsBetweenInstances(): Unit = {
     val checkpointPath = TestUtils.tempFile().getAbsolutePath
     val checkpoint = new LeaderEpochCheckpointFile(new File(checkpointPath))
 
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 cb8996cc3e882..26de772a62fd2 100644
--- a/core/src/test/scala/unit/kafka/server/epoch/LeaderEpochIntegrationTest.scala
+++ b/core/src/test/scala/unit/kafka/server/epoch/LeaderEpochIntegrationTest.scala
@@ -52,7 +52,7 @@ class LeaderEpochIntegrationTest extends ZooKeeperTestHarness with Logging {
   var producer: KafkaProducer[Array[Byte], Array[Byte]] = null
 
   @After
-  override def tearDown() {
+  override def tearDown(): Unit = {
     if (producer != null)
       producer.close()
     TestUtils.shutdownServers(brokers)
@@ -60,7 +60,7 @@ class LeaderEpochIntegrationTest extends ZooKeeperTestHarness with Logging {
   }
 
   @Test
-  def shouldAddCurrentLeaderEpochToMessagesAsTheyAreWrittenToLeader() {
+  def shouldAddCurrentLeaderEpochToMessagesAsTheyAreWrittenToLeader(): Unit = {
     brokers ++= (0 to 1).map { id => createServer(fromProps(createBrokerConfig(id, zkConnect))) }
 
     // Given two topics with replication of a single partition
@@ -145,7 +145,7 @@ class LeaderEpochIntegrationTest extends ZooKeeperTestHarness with Logging {
 
     brokers += createServer(fromProps(createBrokerConfig(101, zkConnect)))
 
-    def leo() = brokers(1).replicaManager.localReplica(tp).get.logEndOffset
+    def leo() = brokers(1).replicaManager.localLog(tp).get.logEndOffset
 
     TestUtils.createTopic(zkClient, tp.topic, Map(tp.partition -> Seq(101)), brokers)
     producer = createProducer(getBrokerListStrFromServers(brokers), acks = -1)
@@ -231,10 +231,7 @@ class LeaderEpochIntegrationTest extends ZooKeeperTestHarness with Logging {
 
   private def waitForEpochChangeTo(topic: String, partition: Int, epoch: Int): Unit = {
     TestUtils.waitUntilTrue(() => {
-      brokers(0).metadataCache.getPartitionInfo(topic, partition) match {
-        case Some(m) => m.basePartitionState.leaderEpoch == epoch
-        case None => false
-      }
+      brokers(0).metadataCache.getPartitionInfo(topic, partition).exists(_.leaderEpoch == epoch)
     }, "Epoch didn't change")
   }
 
@@ -245,10 +242,10 @@ class LeaderEpochIntegrationTest extends ZooKeeperTestHarness with Logging {
       val leo = broker.getLogManager().getLog(tp).get.logEndOffset
       result = result && leo > 0 && brokers.forall { broker =>
         broker.getLogManager().getLog(tp).get.logSegments.iterator.forall { segment =>
-          if (segment.read(minOffset, None, Integer.MAX_VALUE) == null) {
+          if (segment.read(minOffset, Integer.MAX_VALUE) == null) {
             false
           } else {
-            segment.read(minOffset, None, Integer.MAX_VALUE)
+            segment.read(minOffset, Integer.MAX_VALUE)
               .records.batches().iterator().asScala.forall(
               expectedLeaderEpoch == _.partitionLeaderEpoch()
             )
@@ -278,9 +275,10 @@ class LeaderEpochIntegrationTest extends ZooKeeperTestHarness with Logging {
 
     def leaderOffsetsFor(partitions: Map[TopicPartition, Int]): Map[TopicPartition, EpochEndOffset] = {
       val partitionData = partitions.mapValues(
-        new OffsetsForLeaderEpochRequest.PartitionData(Optional.empty(), _))
-      val request = new OffsetsForLeaderEpochRequest.Builder(ApiKeys.OFFSET_FOR_LEADER_EPOCH.latestVersion,
-        partitionData.asJava)
+        new OffsetsForLeaderEpochRequest.PartitionData(Optional.empty(), _)).toMap
+
+      val request = OffsetsForLeaderEpochRequest.Builder.forFollower(
+        ApiKeys.OFFSET_FOR_LEADER_EPOCH.latestVersion, partitionData.asJava, 1)
       val response = sender.sendRequest(request)
       response.responseBody.asInstanceOf[OffsetsForLeaderEpochResponse].responses.asScala
     }
diff --git a/core/src/test/scala/unit/kafka/server/epoch/OffsetsForLeaderEpochTest.scala b/core/src/test/scala/unit/kafka/server/epoch/OffsetsForLeaderEpochTest.scala
index 3d3b3421546db..08be8a28f7967 100644
--- a/core/src/test/scala/unit/kafka/server/epoch/OffsetsForLeaderEpochTest.scala
+++ b/core/src/test/scala/unit/kafka/server/epoch/OffsetsForLeaderEpochTest.scala
@@ -20,7 +20,6 @@ import java.io.File
 import java.util.Optional
 import java.util.concurrent.atomic.AtomicBoolean
 
-import kafka.cluster.Replica
 import kafka.log.{Log, LogManager}
 import kafka.server._
 import kafka.utils.{MockTime, TestUtils}
@@ -57,9 +56,8 @@ class OffsetsForLeaderEpochTest {
     val replicaManager = new ReplicaManager(config, metrics, time, null, null, logManager, new AtomicBoolean(false),
       QuotaFactory.instantiate(config, metrics, time, ""), new BrokerTopicStats,
       new MetadataCache(config.brokerId), new LogDirFailureChannel(config.logDirs.size))
-    val partition = replicaManager.getOrCreatePartition(tp)
-    val leaderReplica = new Replica(config.brokerId, partition.topicPartition, time, 0, Some(mockLog))
-    partition.addReplicaIfNotExists(leaderReplica)
+    val partition = replicaManager.createPartition(tp)
+    partition.setLog(mockLog, isFutureLog = false)
     partition.leaderReplicaIdOpt = Some(config.brokerId)
 
     //When
@@ -79,7 +77,7 @@ class OffsetsForLeaderEpochTest {
     val replicaManager = new ReplicaManager(config, metrics, time, null, null, logManager, new AtomicBoolean(false),
       QuotaFactory.instantiate(config, metrics, time, ""), new BrokerTopicStats,
       new MetadataCache(config.brokerId), new LogDirFailureChannel(config.logDirs.size))
-    replicaManager.getOrCreatePartition(tp)
+    replicaManager.createPartition(tp)
 
     //Given
     val epochRequested: Integer = 5
diff --git a/core/src/test/scala/unit/kafka/server/epoch/util/ReplicaFetcherMockBlockingSend.scala b/core/src/test/scala/unit/kafka/server/epoch/util/ReplicaFetcherMockBlockingSend.scala
index f87e9a5e6c843..f8c528b29108d 100644
--- a/core/src/test/scala/unit/kafka/server/epoch/util/ReplicaFetcherMockBlockingSend.scala
+++ b/core/src/test/scala/unit/kafka/server/epoch/util/ReplicaFetcherMockBlockingSend.scala
@@ -16,13 +16,14 @@
   */
 package kafka.server.epoch.util
 
+import java.net.SocketTimeoutException
 import java.util
 import java.util.Collections
 
 import kafka.cluster.BrokerEndPoint
 import kafka.server.BlockingSend
 import org.apache.kafka.clients.MockClient.MockMetadataUpdater
-import org.apache.kafka.clients.{ClientRequest, ClientResponse, MockClient}
+import org.apache.kafka.clients.{ClientRequest, ClientResponse, MockClient, NetworkClientUtils}
 import org.apache.kafka.common.protocol.{ApiKeys, Errors}
 import org.apache.kafka.common.record.Records
 import org.apache.kafka.common.requests.AbstractRequest.Builder
@@ -38,7 +39,11 @@ import org.apache.kafka.common.{Node, TopicPartition}
   * OFFSET_FOR_LEADER_EPOCH with different offsets in response, it should update offsets using
   * setOffsetsForNextResponse
   */
-class ReplicaFetcherMockBlockingSend(offsets: java.util.Map[TopicPartition, EpochEndOffset], destination: BrokerEndPoint, time: Time) extends BlockingSend {
+class ReplicaFetcherMockBlockingSend(offsets: java.util.Map[TopicPartition, EpochEndOffset],
+                                     sourceBroker: BrokerEndPoint,
+                                     time: Time)
+  extends BlockingSend {
+
   private val client = new MockClient(new SystemTime, new MockMetadataUpdater {
     override def fetchNodes(): util.List[Node] = Collections.emptyList()
     override def isUpdateNeeded: Boolean = false
@@ -49,8 +54,9 @@ class ReplicaFetcherMockBlockingSend(offsets: java.util.Map[TopicPartition, Epoc
   var lastUsedOffsetForLeaderEpochVersion = -1
   var callback: Option[() => Unit] = None
   var currentOffsets: java.util.Map[TopicPartition, EpochEndOffset] = offsets
+  private val sourceNode = new Node(sourceBroker.id, sourceBroker.host, sourceBroker.port)
 
-  def setEpochRequestCallback(postEpochFunction: () => Unit){
+  def setEpochRequestCallback(postEpochFunction: () => Unit): Unit = {
     callback = Some(postEpochFunction)
   }
 
@@ -59,6 +65,8 @@ class ReplicaFetcherMockBlockingSend(offsets: java.util.Map[TopicPartition, Epoc
   }
 
   override def sendRequest(requestBuilder: Builder[_ <: AbstractRequest]): ClientResponse = {
+    if (!NetworkClientUtils.awaitReady(client, sourceNode, time, 500))
+      throw new SocketTimeoutException(s"Failed to connect within 500 ms")
 
     //Send the request to the mock client
     val clientRequest = request(requestBuilder)
@@ -82,13 +90,13 @@ class ReplicaFetcherMockBlockingSend(offsets: java.util.Map[TopicPartition, Epoc
     }
 
     //Use mock client to create the appropriate response object
-    client.respondFrom(response, new Node(destination.id, destination.host, destination.port))
+    client.respondFrom(response, sourceNode)
     client.poll(30, time.milliseconds()).iterator().next()
   }
 
   private def request(requestBuilder: Builder[_ <: AbstractRequest]): ClientRequest = {
     client.newClientRequest(
-      destination.id.toString,
+      sourceBroker.id.toString,
       requestBuilder,
       time.milliseconds(),
       true)
diff --git a/core/src/test/scala/unit/kafka/tools/ConsoleConsumerTest.scala b/core/src/test/scala/unit/kafka/tools/ConsoleConsumerTest.scala
index cdc146f366632..a33c07680ee42 100644
--- a/core/src/test/scala/unit/kafka/tools/ConsoleConsumerTest.scala
+++ b/core/src/test/scala/unit/kafka/tools/ConsoleConsumerTest.scala
@@ -43,7 +43,7 @@ class ConsoleConsumerTest {
   }
 
   @Test
-  def shouldResetUnConsumedOffsetsBeforeExit() {
+  def shouldResetUnConsumedOffsetsBeforeExit(): Unit = {
     val topic = "test"
     val maxMessages: Int = 123
     val totalMessages: Int = 700
@@ -75,7 +75,7 @@ class ConsoleConsumerTest {
   }
 
   @Test
-  def shouldLimitReadsToMaxMessageLimit() {
+  def shouldLimitReadsToMaxMessageLimit(): Unit = {
     val consumer = mock(classOf[ConsumerWrapper])
     val formatter = mock(classOf[MessageFormatter])
     val record = new ConsumerRecord("foo", 1, 1, Array[Byte](), Array[Byte]())
@@ -92,7 +92,7 @@ class ConsoleConsumerTest {
   }
 
   @Test
-  def shouldStopWhenOutputCheckErrorFails() {
+  def shouldStopWhenOutputCheckErrorFails(): Unit = {
     val consumer = mock(classOf[ConsumerWrapper])
     val formatter = mock(classOf[MessageFormatter])
     val printStream = mock(classOf[PrintStream])
@@ -113,7 +113,7 @@ class ConsoleConsumerTest {
   }
 
   @Test
-  def shouldParseValidConsumerValidConfig() {
+  def shouldParseValidConsumerValidConfig(): Unit = {
     //Given
     val args: Array[String] = Array(
       "--bootstrap-server", "localhost:9092",
@@ -170,7 +170,7 @@ class ConsoleConsumerTest {
   }
 
   @Test
-  def shouldParseValidSimpleConsumerValidConfigWithStringOffset() {
+  def shouldParseValidSimpleConsumerValidConfigWithStringOffset(): Unit = {
     //Given
     val args: Array[String] = Array(
       "--bootstrap-server", "localhost:9092",
@@ -192,7 +192,7 @@ class ConsoleConsumerTest {
   }
 
   @Test
-  def shouldParseValidConsumerConfigWithAutoOffsetResetLatest() {
+  def shouldParseValidConsumerConfigWithAutoOffsetResetLatest(): Unit = {
     //Given
     val args: Array[String] = Array(
       "--bootstrap-server", "localhost:9092",
@@ -211,7 +211,7 @@ class ConsoleConsumerTest {
   }
 
   @Test
-  def shouldParseValidConsumerConfigWithAutoOffsetResetEarliest() {
+  def shouldParseValidConsumerConfigWithAutoOffsetResetEarliest(): Unit = {
     //Given
     val args: Array[String] = Array(
       "--bootstrap-server", "localhost:9092",
@@ -230,7 +230,7 @@ class ConsoleConsumerTest {
   }
 
   @Test
-  def shouldParseValidConsumerConfigWithAutoOffsetResetAndMatchingFromBeginning() {
+  def shouldParseValidConsumerConfigWithAutoOffsetResetAndMatchingFromBeginning(): Unit = {
     //Given
     val args: Array[String] = Array(
       "--bootstrap-server", "localhost:9092",
@@ -250,7 +250,7 @@ class ConsoleConsumerTest {
   }
 
   @Test
-  def shouldParseValidConsumerConfigWithNoOffsetReset() {
+  def shouldParseValidConsumerConfigWithNoOffsetReset(): Unit = {
     //Given
     val args: Array[String] = Array(
       "--bootstrap-server", "localhost:9092",
@@ -268,7 +268,7 @@ class ConsoleConsumerTest {
   }
 
   @Test(expected = classOf[IllegalArgumentException])
-  def shouldExitOnInvalidConfigWithAutoOffsetResetAndConflictingFromBeginning() {
+  def shouldExitOnInvalidConfigWithAutoOffsetResetAndConflictingFromBeginning(): Unit = {
     Exit.setExitProcedure((_, message) => throw new IllegalArgumentException(message.orNull))
 
     //Given
@@ -287,7 +287,7 @@ class ConsoleConsumerTest {
   }
 
   @Test
-  def shouldParseConfigsFromFile() {
+  def shouldParseConfigsFromFile(): Unit = {
     val propsFile = TestUtils.tempFile()
     val propsStream = Files.newOutputStream(propsFile.toPath)
     propsStream.write("request.timeout.ms=1000\n".getBytes())
@@ -306,7 +306,7 @@ class ConsoleConsumerTest {
   }
 
   @Test
-  def groupIdsProvidedInDifferentPlacesMustMatch() {
+  def groupIdsProvidedInDifferentPlacesMustMatch(): Unit = {
     Exit.setExitProcedure((_, message) => throw new IllegalArgumentException(message.orNull))
 
     // different in all three places
@@ -433,7 +433,7 @@ class ConsoleConsumerTest {
   }
 
   @Test
-  def shouldParseGroupIdFromBeginningGivenTogether() {
+  def shouldParseGroupIdFromBeginningGivenTogether(): Unit = {
     // Start from earliest
     var args: Array[String] = Array(
       "--bootstrap-server", "localhost:9092",
@@ -462,7 +462,7 @@ class ConsoleConsumerTest {
   }
 
   @Test(expected = classOf[IllegalArgumentException])
-  def shouldExitOnGroupIdAndPartitionGivenTogether() {
+  def shouldExitOnGroupIdAndPartitionGivenTogether(): Unit = {
     Exit.setExitProcedure((_, message) => throw new IllegalArgumentException(message.orNull))
     //Given
     val args: Array[String] = Array(
@@ -480,7 +480,7 @@ class ConsoleConsumerTest {
   }
 
   @Test(expected = classOf[IllegalArgumentException])
-  def shouldExitOnOffsetWithoutPartition() {
+  def shouldExitOnOffsetWithoutPartition(): Unit = {
     Exit.setExitProcedure((_, message) => throw new IllegalArgumentException(message.orNull))
     //Given
     val args: Array[String] = Array(
diff --git a/core/src/test/scala/unit/kafka/tools/ConsoleProducerTest.scala b/core/src/test/scala/unit/kafka/tools/ConsoleProducerTest.scala
index e69a5e67ae820..00efc6a0ca8c7 100644
--- a/core/src/test/scala/unit/kafka/tools/ConsoleProducerTest.scala
+++ b/core/src/test/scala/unit/kafka/tools/ConsoleProducerTest.scala
@@ -23,6 +23,7 @@ import ConsoleProducer.LineMessageReader
 import org.apache.kafka.clients.producer.ProducerConfig
 import org.junit.{Assert, Test}
 import Assert.assertEquals
+import kafka.utils.Exit
 
 class ConsoleProducerTest {
 
@@ -43,20 +44,20 @@ class ConsoleProducerTest {
   )
 
   @Test
-  def testValidConfigs() {
+  def testValidConfigs(): Unit = {
     val config = new ConsoleProducer.ProducerConfig(validArgs)
     val producerConfig = new ProducerConfig(ConsoleProducer.producerProps(config))
     assertEquals(util.Arrays.asList("localhost:1001", "localhost:1002"),
       producerConfig.getList(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG))
   }
 
-  @Test
-  def testInvalidConfigs() {
+  @Test(expected = classOf[IllegalArgumentException])
+  def testInvalidConfigs(): Unit = {
+    Exit.setExitProcedure((_, message) => throw new IllegalArgumentException(message.orNull))
     try {
       new ConsoleProducer.ProducerConfig(invalidArgs)
-      Assert.fail("Should have thrown an UnrecognizedOptionException")
-    } catch {
-      case _: joptsimple.OptionException => // expected exception
+    } finally {
+      Exit.resetExitProcedure()
     }
   }
 
diff --git a/core/src/test/scala/unit/kafka/tools/DumpLogSegmentsTest.scala b/core/src/test/scala/unit/kafka/tools/DumpLogSegmentsTest.scala
index da830a88a45e6..7bfbeab028e99 100644
--- a/core/src/test/scala/unit/kafka/tools/DumpLogSegmentsTest.scala
+++ b/core/src/test/scala/unit/kafka/tools/DumpLogSegmentsTest.scala
@@ -20,7 +20,6 @@ package kafka.tools
 import java.io.{ByteArrayOutputStream, File}
 import java.util.Properties
 
-import scala.collection.mutable
 import kafka.log.{Log, LogConfig, LogManager}
 import kafka.server.{BrokerTopicStats, LogDirFailureChannel}
 import kafka.tools.DumpLogSegments.TimeIndexDumpErrors
@@ -29,9 +28,13 @@ import org.apache.kafka.common.record.{CompressionType, MemoryRecords, SimpleRec
 import org.apache.kafka.common.utils.Utils
 import org.junit.Assert._
 import org.junit.{After, Before, Test}
+import org.scalatest.Assertions.fail
 
+import scala.collection.mutable
 import scala.collection.mutable.ArrayBuffer
 
+case class BatchInfo(records: Seq[SimpleRecord], hasKeys: Boolean, hasValues: Boolean)
+
 class DumpLogSegmentsTest {
 
   val tmpDir = TestUtils.tempDir()
@@ -42,7 +45,7 @@ class DumpLogSegmentsTest {
   val timeIndexFilePath = s"$logDir/$segmentName.timeindex"
   val time = new MockTime(0, 0)
 
-  val batches = new ArrayBuffer[Seq[SimpleRecord]]
+  val batches = new ArrayBuffer[BatchInfo]
   var log: Log = _
 
   @Before
@@ -55,19 +58,19 @@ class DumpLogSegmentsTest {
       logDirFailureChannel = new LogDirFailureChannel(10))
 
     val now = System.currentTimeMillis()
-    val firstBatchRecords = (0 until 10).map { i => new SimpleRecord(now + i * 2, s"hello there $i".getBytes)}
-    batches += firstBatchRecords
-    log.appendAsLeader(MemoryRecords.withRecords(CompressionType.NONE, 0, firstBatchRecords: _*),
-      leaderEpoch = 0)
-    val secondBatchRecords = (10 until 30).map { i => new SimpleRecord(now + i * 3, s"hello there again $i".getBytes)}
-    batches += secondBatchRecords
-    log.appendAsLeader(MemoryRecords.withRecords(CompressionType.NONE, 0, secondBatchRecords: _*),
-      leaderEpoch = 0)
-    val thirdBatchRecords = (30 until 50).map { i => new SimpleRecord(now + i * 5, s"hello there one more time $i".getBytes)}
-    batches += thirdBatchRecords
-    log.appendAsLeader(MemoryRecords.withRecords(CompressionType.NONE, 0, thirdBatchRecords: _*),
-      leaderEpoch = 0)
-
+    val firstBatchRecords = (0 until 10).map { i => new SimpleRecord(now + i * 2, s"message key $i".getBytes, s"message value $i".getBytes)}
+    batches += BatchInfo(firstBatchRecords, true, true)
+    val secondBatchRecords = (10 until 30).map { i => new SimpleRecord(now + i * 3, s"message key $i".getBytes, null)}
+    batches += BatchInfo(secondBatchRecords, true, false)
+    val thirdBatchRecords = (30 until 50).map { i => new SimpleRecord(now + i * 5, null, s"message value $i".getBytes)}
+    batches += BatchInfo(thirdBatchRecords, false, true)
+    val fourthBatchRecords = (50 until 60).map { i => new SimpleRecord(now + i * 7, null)}
+    batches += BatchInfo(fourthBatchRecords, false, false)
+
+    batches.foreach { batchInfo =>
+      log.appendAsLeader(MemoryRecords.withRecords(CompressionType.NONE, 0, batchInfo.records: _*),
+        leaderEpoch = 0)
+    }
     // Flush, but don't close so that the indexes are not trimmed and contain some zero entries
     log.flush()
   }
@@ -81,7 +84,7 @@ class DumpLogSegmentsTest {
   @Test
   def testPrintDataLog(): Unit = {
 
-    def verifyRecordsInOutput(args: Array[String]): Unit = {
+    def verifyRecordsInOutput(checkKeysAndValues: Boolean, args: Array[String]): Unit = {
       def isBatch(index: Int): Boolean = {
         var i = 0
         batches.zipWithIndex.foreach { case (batch, batchIndex) =>
@@ -90,28 +93,39 @@ class DumpLogSegmentsTest {
 
           i += 1
 
-          batch.indices.foreach { recordIndex =>
+          batch.records.indices.foreach { recordIndex =>
             if (i == index)
               return false
             i += 1
           }
         }
-        TestUtils.fail(s"No match for index $index")
+        fail(s"No match for index $index")
       }
 
       val output = runDumpLogSegments(args)
       val lines = output.split("\n")
       assertTrue(s"Data not printed: $output", lines.length > 2)
-      val totalRecords = batches.map(_.size).sum
+      val totalRecords = batches.map(_.records.size).sum
       var offset = 0
+      val batchIterator = batches.iterator
+      var batch : BatchInfo = null;
       (0 until totalRecords + batches.size).foreach { index =>
         val line = lines(lines.length - totalRecords - batches.size + index)
         // The base offset of the batch is the offset of the first record in the batch, so we
         // only increment the offset if it's not a batch
-        if (isBatch(index))
+        if (isBatch(index)) {
           assertTrue(s"Not a valid batch-level message record: $line", line.startsWith(s"baseOffset: $offset lastOffset: "))
-        else {
+          batch = batchIterator.next
+        } else {
           assertTrue(s"Not a valid message record: $line", line.startsWith(s"${DumpLogSegments.RecordIndent} offset: $offset"))
+          if (checkKeysAndValues) {
+            var suffix = "headerKeys: []"
+            if (batch.hasKeys)
+              suffix += s" key: message key $offset"
+            if (batch.hasValues)
+              suffix += s" payload: message value $offset"
+            assertTrue(s"Message record missing key or value: $line", line.endsWith(suffix))
+          }
           offset += 1
         }
       }
@@ -123,15 +137,15 @@ class DumpLogSegmentsTest {
     }
 
     // Verify that records are printed with --print-data-log even if --deep-iteration is not specified
-    verifyRecordsInOutput(Array("--print-data-log", "--files", logFilePath))
+    verifyRecordsInOutput(true, Array("--print-data-log", "--files", logFilePath))
     // Verify that records are printed with --print-data-log if --deep-iteration is also specified
-    verifyRecordsInOutput(Array("--print-data-log", "--deep-iteration", "--files", logFilePath))
+    verifyRecordsInOutput(true, Array("--print-data-log", "--deep-iteration", "--files", logFilePath))
     // Verify that records are printed with --value-decoder even if --print-data-log is not specified
-    verifyRecordsInOutput(Array("--value-decoder-class", "kafka.serializer.StringDecoder", "--files", logFilePath))
+    verifyRecordsInOutput(true, Array("--value-decoder-class", "kafka.serializer.StringDecoder", "--files", logFilePath))
     // Verify that records are printed with --key-decoder even if --print-data-log is not specified
-    verifyRecordsInOutput(Array("--key-decoder-class", "kafka.serializer.StringDecoder", "--files", logFilePath))
+    verifyRecordsInOutput(true, Array("--key-decoder-class", "kafka.serializer.StringDecoder", "--files", logFilePath))
     // Verify that records are printed with --deep-iteration even if --print-data-log is not specified
-    verifyRecordsInOutput(Array("--deep-iteration", "--files", logFilePath))
+    verifyRecordsInOutput(false, Array("--deep-iteration", "--files", logFilePath))
 
     // Verify that records are not printed by default
     verifyNoRecordsInOutput(Array("--files", logFilePath))
diff --git a/core/src/test/scala/unit/kafka/tools/MirrorMakerTest.scala b/core/src/test/scala/unit/kafka/tools/MirrorMakerTest.scala
index a3abd4c10e060..de3016bedc9fc 100644
--- a/core/src/test/scala/unit/kafka/tools/MirrorMakerTest.scala
+++ b/core/src/test/scala/unit/kafka/tools/MirrorMakerTest.scala
@@ -26,7 +26,7 @@ import org.junit.Test
 class MirrorMakerTest {
 
   @Test
-  def testDefaultMirrorMakerMessageHandler() {
+  def testDefaultMirrorMakerMessageHandler(): Unit = {
     val now = 12345L
     val consumerRecord = BaseConsumerRecord("topic", 0, 1L, now, TimestampType.CREATE_TIME, "key".getBytes, "value".getBytes)
 
@@ -42,7 +42,7 @@ class MirrorMakerTest {
   }
 
   @Test
-  def testDefaultMirrorMakerMessageHandlerWithNoTimestampInSourceMessage() {
+  def testDefaultMirrorMakerMessageHandlerWithNoTimestampInSourceMessage(): Unit = {
     val consumerRecord = BaseConsumerRecord("topic", 0, 1L, RecordBatch.NO_TIMESTAMP, TimestampType.CREATE_TIME,
       "key".getBytes, "value".getBytes)
 
@@ -58,7 +58,7 @@ class MirrorMakerTest {
   }
 
   @Test
-  def testDefaultMirrorMakerMessageHandlerWithHeaders() {
+  def testDefaultMirrorMakerMessageHandlerWithHeaders(): Unit = {
     val now = 12345L
     val consumerRecord = BaseConsumerRecord("topic", 0, 1L, now, TimestampType.CREATE_TIME, "key".getBytes,
       "value".getBytes)
diff --git a/core/src/test/scala/unit/kafka/utils/CommandLineUtilsTest.scala b/core/src/test/scala/unit/kafka/utils/CommandLineUtilsTest.scala
index c573fce1bf7b8..2977a2b210c3c 100644
--- a/core/src/test/scala/unit/kafka/utils/CommandLineUtilsTest.scala
+++ b/core/src/test/scala/unit/kafka/utils/CommandLineUtilsTest.scala
@@ -27,21 +27,21 @@ class CommandLineUtilsTest {
 
 
   @Test(expected = classOf[java.lang.IllegalArgumentException])
-  def testParseEmptyArg() {
+  def testParseEmptyArg(): Unit = {
     val argArray = Array("my.empty.property=")
 
     CommandLineUtils.parseKeyValueArgs(argArray, acceptMissingValue = false)
   }
 
   @Test(expected = classOf[java.lang.IllegalArgumentException])
-  def testParseEmptyArgWithNoDelimiter() {
+  def testParseEmptyArgWithNoDelimiter(): Unit = {
     val argArray = Array("my.empty.property")
 
     CommandLineUtils.parseKeyValueArgs(argArray, acceptMissingValue = false)
   }
 
   @Test
-  def testParseEmptyArgAsValid() {
+  def testParseEmptyArgAsValid(): Unit = {
     val argArray = Array("my.empty.property=", "my.empty.property1")
     val props = CommandLineUtils.parseKeyValueArgs(argArray)
 
@@ -50,7 +50,7 @@ class CommandLineUtilsTest {
   }
 
   @Test
-  def testParseSingleArg() {
+  def testParseSingleArg(): Unit = {
     val argArray = Array("my.property=value")
     val props = CommandLineUtils.parseKeyValueArgs(argArray)
 
@@ -58,7 +58,7 @@ class CommandLineUtilsTest {
   }
 
   @Test
-  def testParseArgs() {
+  def testParseArgs(): Unit = {
     val argArray = Array("first.property=first","second.property=second")
     val props = CommandLineUtils.parseKeyValueArgs(argArray)
 
@@ -67,7 +67,7 @@ class CommandLineUtilsTest {
   }
 
   @Test
-  def testParseArgsWithMultipleDelimiters() {
+  def testParseArgsWithMultipleDelimiters(): Unit = {
     val argArray = Array("first.property==first", "second.property=second=", "third.property=thi=rd")
     val props = CommandLineUtils.parseKeyValueArgs(argArray)
 
diff --git a/core/src/test/scala/unit/kafka/utils/CoreUtilsTest.scala b/core/src/test/scala/unit/kafka/utils/CoreUtilsTest.scala
index 93578c616e682..dace3a5bb2395 100755
--- a/core/src/test/scala/unit/kafka/utils/CoreUtilsTest.scala
+++ b/core/src/test/scala/unit/kafka/utils/CoreUtilsTest.scala
@@ -24,7 +24,6 @@ import java.util.concurrent.locks.ReentrantLock
 import java.nio.ByteBuffer
 import java.util.regex.Pattern
 
-import org.scalatest.junit.JUnitSuite
 import org.junit.Assert._
 import kafka.utils.CoreUtils.inLock
 import org.apache.kafka.common.KafkaException
@@ -37,12 +36,12 @@ import scala.collection.mutable
 import scala.concurrent.duration.Duration
 import scala.concurrent.{Await, ExecutionContext, Future}
 
-class CoreUtilsTest extends JUnitSuite with Logging {
+class CoreUtilsTest extends Logging {
 
   val clusterIdPattern = Pattern.compile("[a-zA-Z0-9_\\-]+")
 
   @Test
-  def testSwallow() {
+  def testSwallow(): Unit = {
     CoreUtils.swallow(throw new KafkaException("test"), this, Level.INFO)
   }
 
@@ -97,7 +96,7 @@ class CoreUtilsTest extends JUnitSuite with Logging {
   }
 
   @Test
-  def testCircularIterator() {
+  def testCircularIterator(): Unit = {
     val l = List(1, 2)
     val itl = CoreUtils.circularIterator(l)
     assertEquals(1, itl.next())
@@ -116,7 +115,7 @@ class CoreUtilsTest extends JUnitSuite with Logging {
   }
 
   @Test
-  def testReadBytes() {
+  def testReadBytes(): Unit = {
     for(testCase <- List("", "a", "abcd")) {
       val bytes = testCase.getBytes
       assertTrue(Arrays.equals(bytes, Utils.readBytes(ByteBuffer.wrap(bytes))))
@@ -124,7 +123,7 @@ class CoreUtilsTest extends JUnitSuite with Logging {
   }
 
   @Test
-  def testAbs() {
+  def testAbs(): Unit = {
     assertEquals(0, Utils.abs(Integer.MIN_VALUE))
     assertEquals(1, Utils.abs(-1))
     assertEquals(0, Utils.abs(0))
@@ -133,7 +132,7 @@ class CoreUtilsTest extends JUnitSuite with Logging {
   }
 
   @Test
-  def testReplaceSuffix() {
+  def testReplaceSuffix(): Unit = {
     assertEquals("blah.foo.text", CoreUtils.replaceSuffix("blah.foo.txt", ".txt", ".text"))
     assertEquals("blah.foo", CoreUtils.replaceSuffix("blah.foo.txt", ".txt", ""))
     assertEquals("txt.txt", CoreUtils.replaceSuffix("txt.txt.txt", ".txt", ""))
@@ -141,7 +140,7 @@ class CoreUtilsTest extends JUnitSuite with Logging {
   }
 
   @Test
-  def testReadInt() {
+  def testReadInt(): Unit = {
     val values = Array(0, 1, -1, Byte.MaxValue, Short.MaxValue, 2 * Short.MaxValue, Int.MaxValue/2, Int.MinValue/2, Int.MaxValue, Int.MinValue, Int.MaxValue)
     val buffer = ByteBuffer.allocate(4 * values.size)
     for(i <- 0 until values.length) {
@@ -151,7 +150,7 @@ class CoreUtilsTest extends JUnitSuite with Logging {
   }
 
   @Test
-  def testCsvList() {
+  def testCsvList(): Unit = {
     val emptyString:String = ""
     val nullString:String = null
     val emptyList = CoreUtils.parseCsvList(emptyString)
@@ -164,7 +163,7 @@ class CoreUtilsTest extends JUnitSuite with Logging {
   }
 
   @Test
-  def testCsvMap() {
+  def testCsvMap(): Unit = {
     val emptyString: String = ""
     val emptyMap = CoreUtils.parseCsvMap(emptyString)
     val emptyStringMap = Map.empty[String, String]
@@ -199,7 +198,7 @@ class CoreUtilsTest extends JUnitSuite with Logging {
   }
 
   @Test
-  def testInLock() {
+  def testInLock(): Unit = {
     val lock = new ReentrantLock()
     val result = inLock(lock) {
       assertTrue("Should be in lock", lock.isHeldByCurrentThread)
@@ -210,7 +209,7 @@ class CoreUtilsTest extends JUnitSuite with Logging {
   }
 
   @Test
-  def testUrlSafeBase64EncodeUUID() {
+  def testUrlSafeBase64EncodeUUID(): Unit = {
 
     // Test a UUID that has no + or / characters in base64 encoding [a149b4a3-06e1-4b49-a8cb-8a9c4a59fa46 ->(base64)-> oUm0owbhS0moy4qcSln6Rg==]
     val clusterId1 = Base64.getUrlEncoder.withoutPadding.encodeToString(CoreUtils.getBytesFromUuid(UUID.fromString(
@@ -228,7 +227,7 @@ class CoreUtilsTest extends JUnitSuite with Logging {
   }
 
   @Test
-  def testGenerateUuidAsBase64() {
+  def testGenerateUuidAsBase64(): Unit = {
     val clusterId = CoreUtils.generateUuidAsBase64()
     assertEquals(clusterId.length, 22)
     assertTrue(clusterIdPattern.matcher(clusterId).matches())
diff --git a/core/src/test/scala/unit/kafka/utils/JaasTestUtils.scala b/core/src/test/scala/unit/kafka/utils/JaasTestUtils.scala
index 1870a4996fea7..cc9792e41619f 100644
--- a/core/src/test/scala/unit/kafka/utils/JaasTestUtils.scala
+++ b/core/src/test/scala/unit/kafka/utils/JaasTestUtils.scala
@@ -18,6 +18,9 @@ package kafka.utils
 
 import java.io.{File, BufferedWriter, FileWriter}
 import java.util.Properties
+
+import scala.collection.Seq
+
 import kafka.server.KafkaConfig
 import org.apache.kafka.common.utils.Java
 
@@ -152,10 +155,7 @@ object JaasTestUtils {
   val serviceName = "kafka"
 
   def saslConfigs(saslProperties: Option[Properties]): Properties = {
-    val result = saslProperties match {
-      case Some(properties) => properties
-      case None => new Properties
-    }
+    val result = saslProperties.getOrElse(new Properties)
     // IBM Kerberos module doesn't support the serviceName JAAS property, hence it needs to be
     // passed as a Kafka property
     if (Java.isIbmJdk && !result.contains(KafkaConfig.SaslKerberosServiceNameProp))
@@ -267,7 +267,7 @@ object JaasTestUtils {
   private def jaasSectionsToString(jaasSections: Seq[JaasSection]): String =
     jaasSections.mkString
 
-  private def writeToFile(file: File, jaasSections: Seq[JaasSection]) {
+  private def writeToFile(file: File, jaasSections: Seq[JaasSection]): Unit = {
     val writer = new BufferedWriter(new FileWriter(file))
     try writer.write(jaasSectionsToString(jaasSections))
     finally writer.close()
diff --git a/core/src/test/scala/unit/kafka/utils/JsonTest.scala b/core/src/test/scala/unit/kafka/utils/JsonTest.scala
index 209fdee2edc31..4e41bb36b05c2 100644
--- a/core/src/test/scala/unit/kafka/utils/JsonTest.scala
+++ b/core/src/test/scala/unit/kafka/utils/JsonTest.scala
@@ -37,7 +37,7 @@ object JsonTest {
 class JsonTest {
 
   @Test
-  def testJsonParse() {
+  def testJsonParse(): Unit = {
     val jnf = JsonNodeFactory.instance
 
     assertEquals(Json.parseFull("{}"), Some(JsonValue(new ObjectNode(jnf))))
@@ -66,7 +66,7 @@ class JsonTest {
   }
 
   @Test
-  def testLegacyEncodeAsString() {
+  def testLegacyEncodeAsString(): Unit = {
     assertEquals("null", Json.legacyEncodeAsString(null))
     assertEquals("1", Json.legacyEncodeAsString(1))
     assertEquals("1", Json.legacyEncodeAsString(1L))
@@ -88,7 +88,7 @@ class JsonTest {
   }
 
   @Test
-  def testEncodeAsString() {
+  def testEncodeAsString(): Unit = {
     assertEquals("null", Json.encodeAsString(null))
     assertEquals("1", Json.encodeAsString(1))
     assertEquals("1", Json.encodeAsString(1L))
@@ -111,7 +111,7 @@ class JsonTest {
   }
 
   @Test
-  def testEncodeAsBytes() {
+  def testEncodeAsBytes(): Unit = {
     assertEquals("null", new String(Json.encodeAsBytes(null), StandardCharsets.UTF_8))
     assertEquals("1", new String(Json.encodeAsBytes(1), StandardCharsets.UTF_8))
     assertEquals("1", new String(Json.encodeAsBytes(1L), StandardCharsets.UTF_8))
diff --git a/core/src/test/scala/unit/kafka/utils/LogCaptureAppender.scala b/core/src/test/scala/unit/kafka/utils/LogCaptureAppender.scala
index 80472e9b19f52..2d071452829ff 100644
--- a/core/src/test/scala/unit/kafka/utils/LogCaptureAppender.scala
+++ b/core/src/test/scala/unit/kafka/utils/LogCaptureAppender.scala
@@ -37,7 +37,7 @@ class LogCaptureAppender extends AppenderSkeleton {
     }
   }
 
-  override def close(): Unit =  {
+  override def close(): Unit = {
     events.synchronized {
       events.clear()
     }
diff --git a/core/src/test/scala/unit/kafka/utils/MockScheduler.scala b/core/src/test/scala/unit/kafka/utils/MockScheduler.scala
index 5ebdf406eb3a5..a5d51d066d98c 100644
--- a/core/src/test/scala/unit/kafka/utils/MockScheduler.scala
+++ b/core/src/test/scala/unit/kafka/utils/MockScheduler.scala
@@ -17,7 +17,7 @@
 package kafka.utils
 
 import scala.collection.mutable.PriorityQueue
-import java.util.concurrent.TimeUnit
+import java.util.concurrent.{Delayed, ScheduledFuture, TimeUnit}
 
 import org.apache.kafka.common.utils.Time
 
@@ -41,9 +41,9 @@ class MockScheduler(val time: Time) extends Scheduler {
   
   def isStarted = true
 
-  def startup() {}
+  def startup(): Unit = {}
   
-  def shutdown() {
+  def shutdown(): Unit = {
     this synchronized {
       tasks.foreach(_.fun())
       tasks.clear()
@@ -55,7 +55,7 @@ class MockScheduler(val time: Time) extends Scheduler {
    * when this method is called and the execution happens synchronously in the calling thread.
    * If you are using the scheduler associated with a MockTime instance this call be triggered automatically.
    */
-  def tick() {
+  def tick(): Unit = {
     this synchronized {
       val now = time.milliseconds
       while(tasks.nonEmpty && tasks.head.nextExecution <= now) {
@@ -71,11 +71,14 @@ class MockScheduler(val time: Time) extends Scheduler {
     }
   }
   
-  def schedule(name: String, fun: () => Unit, delay: Long = 0, period: Long = -1, unit: TimeUnit = TimeUnit.MILLISECONDS) {
+  def schedule(name: String, fun: () => Unit, delay: Long = 0, period: Long = -1, unit: TimeUnit = TimeUnit.MILLISECONDS): ScheduledFuture[Unit] = {
+    var task : MockTask = null
     this synchronized {
-      tasks += MockTask(name, fun, time.milliseconds + delay, period = period)
+      task = MockTask(name, fun, time.milliseconds + delay, period = period, time=time)
+      tasks += task
       tick()
     }
+    task
   }
 
   def clear(): Unit = {
@@ -86,7 +89,7 @@ class MockScheduler(val time: Time) extends Scheduler {
   
 }
 
-case class MockTask(name: String, fun: () => Unit, var nextExecution: Long, period: Long) extends Ordered[MockTask] {
+case class MockTask(name: String, fun: () => Unit, var nextExecution: Long, period: Long, time: Time) extends ScheduledFuture[Unit] {
   def periodic = period >= 0
   def compare(t: MockTask): Int = {
     if(t.nextExecution == nextExecution)
@@ -96,4 +99,42 @@ case class MockTask(name: String, fun: () => Unit, var nextExecution: Long, peri
     else
       1
   }
+
+  /**
+    * Not used, so not not fully implemented
+    */
+  def cancel(mayInterruptIfRunning: Boolean) : Boolean = {
+    false
+  }
+
+  def get(): Unit = {
+  }
+
+  def get(timeout: Long, unit: TimeUnit): Unit = {
+  }
+
+  def isCancelled: Boolean = {
+    false
+  }
+
+  def isDone: Boolean = {
+    false
+  }
+
+  def getDelay(unit: TimeUnit): Long = {
+    this synchronized {
+      time.milliseconds - nextExecution
+    }
+  }
+
+  def compareTo(o : Delayed) : Int = {
+    this.getDelay(TimeUnit.MILLISECONDS).compareTo(o.getDelay(TimeUnit.MILLISECONDS))
+  }
+}
+object MockTask {
+  implicit def MockTaskOrdering : Ordering[MockTask] = new Ordering[MockTask] {
+    def compare(x: MockTask, y: MockTask): Int = {
+      x.compare(y)
+    }
+  }
 }
diff --git a/core/src/test/scala/unit/kafka/utils/MockTime.scala b/core/src/test/scala/unit/kafka/utils/MockTime.scala
index 2d83d6547f574..bf0e7bd3816a8 100644
--- a/core/src/test/scala/unit/kafka/utils/MockTime.scala
+++ b/core/src/test/scala/unit/kafka/utils/MockTime.scala
@@ -32,7 +32,7 @@ class MockTime(currentTimeMs: Long, currentHiResTimeNs: Long) extends JMockTime(
 
   val scheduler = new MockScheduler(this)
 
-  override def sleep(ms: Long) {
+  override def sleep(ms: Long): Unit = {
     super.sleep(ms)
     scheduler.tick()
   }
diff --git a/core/src/test/scala/unit/kafka/utils/ReplicationUtilsTest.scala b/core/src/test/scala/unit/kafka/utils/ReplicationUtilsTest.scala
index 4bf747136841d..e8de8d9ecfb08 100644
--- a/core/src/test/scala/unit/kafka/utils/ReplicationUtilsTest.scala
+++ b/core/src/test/scala/unit/kafka/utils/ReplicationUtilsTest.scala
@@ -37,7 +37,7 @@ class ReplicationUtilsTest extends ZooKeeperTestHarness {
   private val isr = List(1, 2)
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     super.setUp()
     zkClient.makeSurePersistentPathExists(TopicZNode.path(topic))
     val topicPartition = new TopicPartition(topic, partition)
@@ -47,7 +47,7 @@ class ReplicationUtilsTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  def testUpdateLeaderAndIsr() {
+  def testUpdateLeaderAndIsr(): Unit = {
     val configs = TestUtils.createBrokerConfigs(1, zkConnect).map(KafkaConfig.fromProps)
     val log: Log = EasyMock.createMock(classOf[Log])
     EasyMock.expect(log.logEndOffset).andReturn(20).anyTimes()
diff --git a/core/src/test/scala/unit/kafka/utils/SchedulerTest.scala b/core/src/test/scala/unit/kafka/utils/SchedulerTest.scala
index dbb818c046948..201216c6d7291 100644
--- a/core/src/test/scala/unit/kafka/utils/SchedulerTest.scala
+++ b/core/src/test/scala/unit/kafka/utils/SchedulerTest.scala
@@ -18,8 +18,11 @@ package kafka.utils
 
 import org.junit.Assert._
 import java.util.concurrent.atomic._
-import org.junit.{Test, After, Before}
+import kafka.log.{Log, LogConfig, LogManager, ProducerStateManager}
+import kafka.server.{BrokerTopicStats, LogDirFailureChannel}
+import org.junit.{After, Before, Test}
 import kafka.utils.TestUtils.retry
+import java.util.Properties
 
 class SchedulerTest {
 
@@ -29,17 +32,17 @@ class SchedulerTest {
   val counter2 = new AtomicInteger(0)
   
   @Before
-  def setup() {
+  def setup(): Unit = {
     scheduler.startup()
   }
   
   @After
-  def teardown() {
+  def teardown(): Unit = {
     scheduler.shutdown()
   }
 
   @Test
-  def testMockSchedulerNonPeriodicTask() {
+  def testMockSchedulerNonPeriodicTask(): Unit = {
     mockTime.scheduler.schedule("test1", counter1.getAndIncrement _, delay=1)
     mockTime.scheduler.schedule("test2", counter2.getAndIncrement _, delay=100)
     assertEquals("Counter1 should not be incremented prior to task running.", 0, counter1.get)
@@ -53,7 +56,7 @@ class SchedulerTest {
   }
 
   @Test
-  def testMockSchedulerPeriodicTask() {
+  def testMockSchedulerPeriodicTask(): Unit = {
     mockTime.scheduler.schedule("test1", counter1.getAndIncrement _, delay=1, period=1)
     mockTime.scheduler.schedule("test2", counter2.getAndIncrement _, delay=100, period=100)
     assertEquals("Counter1 should not be incremented prior to task running.", 0, counter1.get)
@@ -67,14 +70,14 @@ class SchedulerTest {
   }
 
   @Test
-  def testReentrantTaskInMockScheduler() {
+  def testReentrantTaskInMockScheduler(): Unit = {
     mockTime.scheduler.schedule("test1", () => mockTime.scheduler.schedule("test2", counter2.getAndIncrement _, delay=0), delay=1)
     mockTime.sleep(1)
     assertEquals(1, counter2.get)
   }
 
   @Test
-  def testNonPeriodicTask() {
+  def testNonPeriodicTask(): Unit = {
     scheduler.schedule("test", counter1.getAndIncrement _, delay = 0)
     retry(30000) {
       assertEquals(counter1.get, 1)
@@ -84,7 +87,7 @@ class SchedulerTest {
   }
 
   @Test
-  def testPeriodicTask() {
+  def testPeriodicTask(): Unit = {
     scheduler.schedule("test", counter1.getAndIncrement _, delay = 0, period = 5)
     retry(30000){
       assertTrue("Should count to 20", counter1.get >= 20)
@@ -92,7 +95,7 @@ class SchedulerTest {
   }
 
   @Test
-  def testRestart() {
+  def testRestart(): Unit = {
     // schedule a task to increment a counter
     mockTime.scheduler.schedule("test1", counter1.getAndIncrement _, delay=1)
     mockTime.sleep(1)
@@ -107,4 +110,23 @@ class SchedulerTest {
     mockTime.sleep(1)
     assertEquals(2, counter1.get())
   }
+
+  @Test
+  def testUnscheduleProducerTask(): Unit = {
+    val tmpDir = TestUtils.tempDir()
+    val logDir = TestUtils.randomPartitionLogDir(tmpDir)
+    val logConfig = LogConfig(new Properties())
+    val brokerTopicStats = new BrokerTopicStats
+    val recoveryPoint = 0L
+    val maxProducerIdExpirationMs = 60 * 60 * 1000
+    val topicPartition = Log.parseTopicPartitionName(logDir)
+    val producerStateManager = new ProducerStateManager(topicPartition, logDir, maxProducerIdExpirationMs)
+    val log = new Log(logDir, logConfig, logStartOffset = 0, recoveryPoint = recoveryPoint, scheduler,
+      brokerTopicStats, mockTime, maxProducerIdExpirationMs, LogManager.ProducerIdExpirationCheckIntervalMs,
+      topicPartition, producerStateManager, new LogDirFailureChannel(10))
+    assertTrue(scheduler.taskRunning(log.producerExpireCheck))
+    log.close()
+    assertTrue(!(scheduler.taskRunning(log.producerExpireCheck)))
+  }
+
 }
\ No newline at end of file
diff --git a/core/src/test/scala/unit/kafka/utils/TestUtils.scala b/core/src/test/scala/unit/kafka/utils/TestUtils.scala
index 8b8b230a1c598..00c40209682a8 100755
--- a/core/src/test/scala/unit/kafka/utils/TestUtils.scala
+++ b/core/src/test/scala/unit/kafka/utils/TestUtils.scala
@@ -24,10 +24,12 @@ import java.nio.charset.{Charset, StandardCharsets}
 import java.nio.file.{Files, StandardOpenOption}
 import java.security.cert.X509Certificate
 import java.time.Duration
-import java.util.{Collections, Properties}
+import java.util.Arrays
+import java.util.Collections
+import java.util.Properties
 import java.util.concurrent.{Callable, ExecutionException, Executors, TimeUnit}
-import javax.net.ssl.X509TrustManager
 
+import javax.net.ssl.X509TrustManager
 import kafka.api._
 import kafka.cluster.{Broker, EndPoint}
 import kafka.log._
@@ -35,32 +37,41 @@ import kafka.security.auth.{Acl, Authorizer, Resource}
 import kafka.server._
 import kafka.server.checkpoints.OffsetCheckpointFile
 import Implicits._
+import com.yammer.metrics.Metrics
+import com.yammer.metrics.core.Meter
 import kafka.controller.LeaderIsrAndControllerEpoch
 import kafka.zk._
 import org.apache.kafka.clients.CommonClientConfigs
-import org.apache.kafka.clients.admin.{AdminClient, AlterConfigsResult, Config, ConfigEntry}
+import org.apache.kafka.clients.admin.AlterConfigOp.OpType
+import org.apache.kafka.clients.admin._
 import org.apache.kafka.clients.consumer._
 import org.apache.kafka.clients.producer.{KafkaProducer, ProducerConfig, ProducerRecord}
+import org.apache.kafka.common.acl.{AccessControlEntry, AccessControlEntryFilter, AclBindingFilter}
 import org.apache.kafka.common.{KafkaFuture, TopicPartition}
 import org.apache.kafka.common.config.ConfigResource
-import org.apache.kafka.common.errors.RetriableException
+import org.apache.kafka.common.errors.UnknownTopicOrPartitionException
 import org.apache.kafka.common.header.Header
 import org.apache.kafka.common.internals.Topic
 import org.apache.kafka.common.network.{ListenerName, Mode}
 import org.apache.kafka.common.record._
+import org.apache.kafka.common.resource.ResourcePattern
 import org.apache.kafka.common.security.auth.SecurityProtocol
-import org.apache.kafka.common.serialization.{ByteArrayDeserializer, ByteArraySerializer, Deserializer, Serializer}
+import org.apache.kafka.common.serialization.{ByteArrayDeserializer, ByteArraySerializer, Deserializer, IntegerSerializer, Serializer}
 import org.apache.kafka.common.utils.Time
 import org.apache.kafka.common.utils.Utils._
+import org.apache.kafka.server.authorizer.{Authorizer => JAuthorizer}
 import org.apache.kafka.test.{TestSslUtils, TestUtils => JTestUtils}
 import org.apache.zookeeper.KeeperException.SessionExpiredException
 import org.apache.zookeeper.ZooDefs._
 import org.apache.zookeeper.data.ACL
 import org.junit.Assert._
+import org.scalatest.Assertions.fail
 
 import scala.collection.JavaConverters._
-import scala.collection.{Map, mutable}
+import scala.collection.{Map, Seq, mutable}
 import scala.collection.mutable.{ArrayBuffer, ListBuffer}
+import scala.concurrent.duration.FiniteDuration
+import scala.concurrent.{Await, ExecutionContext, Future}
 
 /**
  * Utility functions to help with testing
@@ -136,7 +147,15 @@ object TestUtils extends Logging {
    * @param config The configuration of the server
    */
   def createServer(config: KafkaConfig, time: Time = Time.SYSTEM): KafkaServer = {
-    val server = new KafkaServer(config, time)
+    createServer(config, time, None)
+  }
+
+  def createServer(config: KafkaConfig, threadNamePrefix: Option[String]): KafkaServer = {
+    createServer(config, Time.SYSTEM, threadNamePrefix)
+  }
+
+  def createServer(config: KafkaConfig, time: Time, threadNamePrefix: Option[String]): KafkaServer = {
+    val server = new KafkaServer(config, time, threadNamePrefix = threadNamePrefix)
     server.startup()
     server
   }
@@ -170,11 +189,14 @@ object TestUtils extends Logging {
     enableSaslSsl: Boolean = false,
     rackInfo: Map[Int, String] = Map(),
     logDirCount: Int = 1,
-    enableToken: Boolean = false): Seq[Properties] = {
+    enableToken: Boolean = false,
+    numPartitions: Int = 1,
+    defaultReplicationFactor: Short = 1): Seq[Properties] = {
     (0 until numConfigs).map { node =>
       createBrokerConfig(node, zkConnect, enableControlledShutdown, enableDeleteTopic, RandomPort,
         interBrokerSecurityProtocol, trustStoreFile, saslProperties, enablePlaintext = enablePlaintext, enableSsl = enableSsl,
-        enableSaslPlaintext = enableSaslPlaintext, enableSaslSsl = enableSaslSsl, rack = rackInfo.get(node), logDirCount = logDirCount, enableToken = enableToken)
+        enableSaslPlaintext = enableSaslPlaintext, enableSaslSsl = enableSaslSsl, rack = rackInfo.get(node), logDirCount = logDirCount, enableToken = enableToken,
+        numPartitions = numPartitions, defaultReplicationFactor = defaultReplicationFactor)
     }
   }
 
@@ -197,11 +219,15 @@ object TestUtils extends Logging {
   /**
     * Shutdown `servers` and delete their log directories.
     */
-  def shutdownServers(servers: Seq[KafkaServer]) {
-    servers.par.foreach { s =>
-      s.shutdown()
-      CoreUtils.delete(s.config.logDirs)
+  def shutdownServers(servers: Seq[KafkaServer]): Unit = {
+    import ExecutionContext.Implicits._
+    val future = Future.traverse(servers) { s =>
+      Future {
+        s.shutdown()
+        CoreUtils.delete(s.config.logDirs)
+      }
     }
+    Await.result(future, FiniteDuration(5, TimeUnit.MINUTES))
   }
 
   /**
@@ -226,7 +252,9 @@ object TestUtils extends Logging {
                          saslSslPort: Int = RandomPort,
                          rack: Option[String] = None,
                          logDirCount: Int = 1,
-                         enableToken: Boolean = false): Properties = {
+                         enableToken: Boolean = false,
+                         numPartitions: Int = 1,
+                         defaultReplicationFactor: Short = 1): Properties = {
     def shouldEnable(protocol: SecurityProtocol) = interBrokerSecurityProtocol.fold(false)(_ == protocol)
 
     val protocolAndPorts = ArrayBuffer[(SecurityProtocol, Int)]()
@@ -286,6 +314,9 @@ object TestUtils extends Logging {
     if (enableToken)
       props.put(KafkaConfig.DelegationTokenMasterKeyProp, "masterkey")
 
+    props.put(KafkaConfig.NumPartitionsProp, numPartitions.toString)
+    props.put(KafkaConfig.DefaultReplicationFactorProp, defaultReplicationFactor.toString)
+
     props
   }
 
@@ -302,13 +333,13 @@ object TestUtils extends Logging {
                   topicConfig: Properties = new Properties): scala.collection.immutable.Map[Int, Int] = {
     val adminZkClient = new AdminZkClient(zkClient)
     // create topic
-    TestUtils.waitUntilTrue( () => {
+    waitUntilTrue( () => {
       var hasSessionExpirationException = false
       try {
         adminZkClient.createTopic(topic, numPartitions, replicationFactor, topicConfig)
       } catch {
         case _: SessionExpiredException => hasSessionExpirationException = true
-        case e => throw e // let other exceptions propagate
+        case e: Throwable => throw e // let other exceptions propagate
       }
       !hasSessionExpirationException},
       s"Can't create topic $topic")
@@ -344,13 +375,13 @@ object TestUtils extends Logging {
                   topicConfig: Properties): scala.collection.immutable.Map[Int, Int] = {
     val adminZkClient = new AdminZkClient(zkClient)
     // create topic
-    TestUtils.waitUntilTrue( () => {
+    waitUntilTrue( () => {
       var hasSessionExpirationException = false
       try {
         adminZkClient.createTopicWithAssignment(topic, topicConfig, partitionReplicaAssignment)
       } catch {
         case _: SessionExpiredException => hasSessionExpirationException = true
-        case e => throw e // let other exceptions propagate
+        case e: Throwable => throw e // let other exceptions propagate
       }
       !hasSessionExpirationException},
       s"Can't create topic $topic")
@@ -375,11 +406,6 @@ object TestUtils extends Logging {
       server.groupCoordinator.offsetsTopicConfigs)
   }
 
-  /**
-   * Fail a test case explicitly. Return Nothing so that we are not constrained by the return type.
-   */
-  def fail(msg: String): Nothing = throw new AssertionError(msg)
-
   /**
    * Wrap a single record log buffer.
    */
@@ -430,7 +456,7 @@ object TestUtils extends Logging {
   /**
    * Check that the buffer content from buffer.position() to buffer.limit() is equal
    */
-  def checkEquals(b1: ByteBuffer, b2: ByteBuffer) {
+  def checkEquals(b1: ByteBuffer, b2: ByteBuffer): Unit = {
     assertEquals("Buffers should have equal length", b1.limit() - b1.position(), b2.limit() - b2.position())
     for(i <- 0 until b1.limit() - b1.position())
       assertEquals("byte " + i + " byte not equal.", b1.get(b1.position() + i), b2.get(b1.position() + i))
@@ -440,7 +466,7 @@ object TestUtils extends Logging {
    * Throw an exception if the two iterators are of differing lengths or contain
    * different messages on their Nth element
    */
-  def checkEquals[T](expected: Iterator[T], actual: Iterator[T]) {
+  def checkEquals[T](expected: Iterator[T], actual: Iterator[T]): Unit = {
     var length = 0
     while(expected.hasNext && actual.hasNext) {
       length += 1
@@ -472,7 +498,7 @@ object TestUtils extends Logging {
    *  Throw an exception if an iterable has different length than expected
    *
    */
-  def checkLength[T](s1: Iterator[T], expectedLength:Int) {
+  def checkLength[T](s1: Iterator[T], expectedLength:Int): Unit = {
     var n = 0
     while (s1.hasNext) {
       n+=1
@@ -485,7 +511,7 @@ object TestUtils extends Logging {
    * Throw an exception if the two iterators are of differing lengths or contain
    * different messages on their Nth element
    */
-  def checkEquals[T](s1: java.util.Iterator[T], s2: java.util.Iterator[T]) {
+  def checkEquals[T](s1: java.util.Iterator[T], s2: java.util.Iterator[T]): Unit = {
     while(s1.hasNext && s2.hasNext)
       assertEquals(s1.next, s2.next)
     assertFalse("Iterators have uneven length--first has more", s1.hasNext)
@@ -582,6 +608,7 @@ object TestUtils extends Logging {
     producerProps.put(ProducerConfig.LINGER_MS_CONFIG, lingerMs.toString)
     producerProps.put(ProducerConfig.BATCH_SIZE_CONFIG, batchSize.toString)
     producerProps.put(ProducerConfig.COMPRESSION_TYPE_CONFIG, compressionType)
+    producerProps.put(ProducerConfig.ALLOW_AUTO_CREATE_TOPICS_CONFIG, boolean2Boolean(true))
     producerProps ++= producerSecurityConfigs(securityProtocol, trustStoreFile, saslProperties)
     new KafkaProducer[K, V](producerProps, keySerializer, valueSerializer)
   }
@@ -623,6 +650,7 @@ object TestUtils extends Logging {
     consumerProps.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, enableAutoCommit.toString)
     consumerProps.put(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, maxPollRecords.toString)
     consumerProps.put(ConsumerConfig.ISOLATION_LEVEL_CONFIG, if (readCommitted) "read_committed" else "read_uncommitted")
+    consumerProps.put(ConsumerConfig.ALLOW_AUTO_CREATE_TOPICS_CONFIG, boolean2Boolean(true))
     consumerProps ++= consumerSecurityConfigs(securityProtocol, trustStoreFile, saslProperties)
     new KafkaConsumer[K, V](consumerProps, keyDeserializer, valueDeserializer)
   }
@@ -658,7 +686,7 @@ object TestUtils extends Logging {
   def makeLeaderForPartition(zkClient: KafkaZkClient,
                              topic: String,
                              leaderPerPartitionMap: scala.collection.immutable.Map[Int, Int],
-                             controllerEpoch: Int) {
+                             controllerEpoch: Int): Unit = {
     val newLeaderIsrAndControllerEpochs = leaderPerPartitionMap.map { case (partition, leader) =>
       val topicPartition = new TopicPartition(topic, partition)
       val newLeaderAndIsr = zkClient.getTopicPartitionState(topicPartition)
@@ -729,7 +757,7 @@ object TestUtils extends Logging {
    * Execute the given block. If it throws an assert error, retry. Repeat
    * until no error is thrown or the time limit elapses
    */
-  def retry(maxWaitMs: Long)(block: => Unit) {
+  def retry(maxWaitMs: Long)(block: => Unit): Unit = {
     var wait = 1L
     val startTime = System.currentTimeMillis()
     while(true) {
@@ -755,7 +783,7 @@ object TestUtils extends Logging {
                     msg: => String,
                     waitTimeMs: Long = JTestUtils.DEFAULT_MAX_WAIT_MS): Unit = {
     waitUntilTrue(() => {
-      consumer.poll(Duration.ofMillis(50))
+      consumer.poll(Duration.ofMillis(100))
       action()
     }, msg = msg, pause = 0L, waitTimeMs = waitTimeMs)
   }
@@ -765,11 +793,39 @@ object TestUtils extends Logging {
                                  msg: => String,
                                  waitTimeMs: Long = JTestUtils.DEFAULT_MAX_WAIT_MS): Unit = {
     waitUntilTrue(() => {
-      val records = consumer.poll(Duration.ofMillis(50))
+      val records = consumer.poll(Duration.ofMillis(100))
       action(records)
     }, msg = msg, pause = 0L, waitTimeMs = waitTimeMs)
   }
 
+  def subscribeAndWaitForRecords(topic: String,
+                                 consumer: KafkaConsumer[Array[Byte], Array[Byte]],
+                                 waitTimeMs: Long = JTestUtils.DEFAULT_MAX_WAIT_MS): Unit = {
+    consumer.subscribe(Collections.singletonList(topic))
+    pollRecordsUntilTrue(
+      consumer,
+      (records: ConsumerRecords[Array[Byte], Array[Byte]]) => !records.isEmpty,
+      "Expected records",
+      waitTimeMs)
+  }
+
+  /**
+   * Wait for the presence of an optional value.
+   *
+   * @param func The function defining the optional value
+   * @param msg Error message in the case that the value never appears
+   * @param waitTimeMs Maximum time to wait
+   * @return The unwrapped value returned by the function
+   */
+  def awaitValue[T](func: () => Option[T], msg: => String, waitTimeMs: Long = JTestUtils.DEFAULT_MAX_WAIT_MS): T = {
+    var value: Option[T] = None
+    waitUntilTrue(() => {
+      value = func()
+      value.isDefined
+    }, msg, waitTimeMs)
+    value.get
+  }
+
   /**
     *  Wait until the given condition is true or throw an exception if the given wait time elapses.
     *
@@ -777,27 +833,18 @@ object TestUtils extends Logging {
     * @param msg error message
     * @param waitTimeMs maximum time to wait and retest the condition before failing the test
     * @param pause delay between condition checks
-    * @param maxRetries maximum number of retries to check the given condition if a retriable exception is thrown
     */
   def waitUntilTrue(condition: () => Boolean, msg: => String,
-                    waitTimeMs: Long = JTestUtils.DEFAULT_MAX_WAIT_MS, pause: Long = 100L, maxRetries: Int = 0): Unit = {
+                    waitTimeMs: Long = JTestUtils.DEFAULT_MAX_WAIT_MS, pause: Long = 100L): Unit = {
     val startTime = System.currentTimeMillis()
-    var retry = 0
     while (true) {
-      try {
-        if (condition())
-          return
-        if (System.currentTimeMillis() > startTime + waitTimeMs)
-          fail(msg)
-        Thread.sleep(waitTimeMs.min(pause))
-      }
-      catch {
-        case e: RetriableException if retry < maxRetries =>
-          debug("Retrying after error", e)
-          retry += 1
-        case e : Throwable => throw e
-      }
+      if (condition())
+        return
+      if (System.currentTimeMillis() > startTime + waitTimeMs)
+        fail(msg)
+      Thread.sleep(waitTimeMs.min(pause))
     }
+
     // should never hit here
     throw new RuntimeException("unexpected error")
   }
@@ -825,14 +872,14 @@ object TestUtils extends Logging {
   }
 
   def isLeaderLocalOnBroker(topic: String, partitionId: Int, server: KafkaServer): Boolean = {
-    server.replicaManager.getPartition(new TopicPartition(topic, partitionId)).exists(_.leaderReplicaIfLocal.isDefined)
+    server.replicaManager.nonOfflinePartition(new TopicPartition(topic, partitionId)).exists(_.leaderLogIfLocal.isDefined)
   }
 
   def findLeaderEpoch(brokerId: Int,
                       topicPartition: TopicPartition,
                       servers: Iterable[KafkaServer]): Int = {
     val leaderServer = servers.find(_.config.brokerId == brokerId)
-    val leaderPartition = leaderServer.flatMap(_.replicaManager.getPartition(topicPartition))
+    val leaderPartition = leaderServer.flatMap(_.replicaManager.nonOfflinePartition(topicPartition))
       .getOrElse(fail(s"Failed to find expected replica on broker $brokerId"))
     leaderPartition.getLeaderEpoch
   }
@@ -840,7 +887,7 @@ object TestUtils extends Logging {
   def findFollowerId(topicPartition: TopicPartition,
                      servers: Iterable[KafkaServer]): Int = {
     val followerOpt = servers.find { server =>
-      server.replicaManager.getPartition(topicPartition) match {
+      server.replicaManager.nonOfflinePartition(topicPartition) match {
         case Some(partition) => !partition.leaderReplicaIdOpt.contains(server.config.brokerId)
         case None => false
       }
@@ -859,7 +906,7 @@ object TestUtils extends Logging {
   def waitUntilBrokerMetadataIsPropagated(servers: Seq[KafkaServer],
                                           timeout: Long = JTestUtils.DEFAULT_MAX_WAIT_MS): Unit = {
     val expectedBrokerIds = servers.map(_.config.brokerId).toSet
-    TestUtils.waitUntilTrue(() => servers.forall(server =>
+    waitUntilTrue(() => servers.forall(server =>
       expectedBrokerIds == server.dataPlaneRequestProcessor.metadataCache.getAliveBrokers.map(_.id).toSet
     ), "Timed out waiting for broker metadata to propagate to all servers", timeout)
   }
@@ -877,16 +924,14 @@ object TestUtils extends Logging {
   def waitUntilMetadataIsPropagated(servers: Seq[KafkaServer], topic: String, partition: Int,
                                     timeout: Long = JTestUtils.DEFAULT_MAX_WAIT_MS): Int = {
     var leader: Int = -1
-    TestUtils.waitUntilTrue(() =>
-      servers.foldLeft(true) {
-        (result, server) =>
-          val partitionStateOpt = server.dataPlaneRequestProcessor.metadataCache.getPartitionInfo(topic, partition)
-          partitionStateOpt match {
-            case None => false
-            case Some(partitionState) =>
-              leader = partitionState.basePartitionState.leader
-              result && Request.isValidBrokerId(leader)
-          }
+    waitUntilTrue(
+      () => servers.forall { server =>
+        server.dataPlaneRequestProcessor.metadataCache.getPartitionInfo(topic, partition) match {
+          case Some(partitionState) if Request.isValidBrokerId(partitionState.leader) =>
+            leader = partitionState.leader
+            true
+          case _ => false
+        }
       },
       "Partition [%s,%d] metadata not propagated after %d ms".format(topic, partition, timeout),
       waitTimeMs = timeout)
@@ -906,11 +951,11 @@ object TestUtils extends Logging {
     def newLeaderExists: Option[Int] = {
       servers.find { server =>
         server.config.brokerId != oldLeader &&
-          server.replicaManager.getPartition(tp).exists(_.leaderReplicaIfLocal.isDefined)
+          server.replicaManager.nonOfflinePartition(tp).exists(_.leaderLogIfLocal.isDefined)
       }.map(_.config.brokerId)
     }
 
-    TestUtils.waitUntilTrue(() => newLeaderExists.isDefined,
+    waitUntilTrue(() => newLeaderExists.isDefined,
       s"Did not observe leader change for partition $tp after $timeout ms", waitTimeMs = timeout)
 
     newLeaderExists.get
@@ -921,17 +966,17 @@ object TestUtils extends Logging {
                              timeout: Long = JTestUtils.DEFAULT_MAX_WAIT_MS): Int = {
     def leaderIfExists: Option[Int] = {
       servers.find { server =>
-        server.replicaManager.getPartition(tp).exists(_.leaderReplicaIfLocal.isDefined)
+        server.replicaManager.nonOfflinePartition(tp).exists(_.leaderLogIfLocal.isDefined)
       }.map(_.config.brokerId)
     }
 
-    TestUtils.waitUntilTrue(() => leaderIfExists.isDefined,
+    waitUntilTrue(() => leaderIfExists.isDefined,
       s"Partition $tp leaders not made yet after $timeout ms", waitTimeMs = timeout)
 
     leaderIfExists.get
   }
 
-  def writeNonsenseToFile(fileName: File, position: Long, size: Int) {
+  def writeNonsenseToFile(fileName: File, position: Long, size: Int): Unit = {
     val file = new RandomAccessFile(fileName, "rw")
     file.seek(position)
     for (_ <- 0 until size)
@@ -939,7 +984,7 @@ object TestUtils extends Logging {
     file.close()
   }
 
-  def appendNonsenseToFile(file: File, size: Int) {
+  def appendNonsenseToFile(file: File, size: Int): Unit = {
     val outputStream = Files.newOutputStream(file.toPath(), StandardOpenOption.APPEND)
     try {
       for (_ <- 0 until size)
@@ -947,7 +992,7 @@ object TestUtils extends Logging {
     } finally outputStream.close()
   }
 
-  def checkForPhantomInSyncReplicas(zkClient: KafkaZkClient, topic: String, partitionToBeReassigned: Int, assignedReplicas: Seq[Int]) {
+  def checkForPhantomInSyncReplicas(zkClient: KafkaZkClient, topic: String, partitionToBeReassigned: Int, assignedReplicas: Seq[Int]): Unit = {
     val inSyncReplicas = zkClient.getInSyncReplicasForPartition(new TopicPartition(topic, partitionToBeReassigned))
     // in sync replicas should not have any replica that is not in the new assigned replicas
     val phantomInSyncReplicas = inSyncReplicas.get.toSet -- assignedReplicas.toSet
@@ -956,27 +1001,29 @@ object TestUtils extends Logging {
   }
 
   def ensureNoUnderReplicatedPartitions(zkClient: KafkaZkClient, topic: String, partitionToBeReassigned: Int, assignedReplicas: Seq[Int],
-                                                servers: Seq[KafkaServer]) {
+                                                servers: Seq[KafkaServer]): Unit = {
     val topicPartition = new TopicPartition(topic, partitionToBeReassigned)
-    TestUtils.waitUntilTrue(() => {
+    waitUntilTrue(() => {
         val inSyncReplicas = zkClient.getInSyncReplicasForPartition(topicPartition)
         inSyncReplicas.get.size == assignedReplicas.size
       },
       "Reassigned partition [%s,%d] is under replicated".format(topic, partitionToBeReassigned))
     var leader: Option[Int] = None
-    TestUtils.waitUntilTrue(() => {
+    waitUntilTrue(() => {
         leader = zkClient.getLeaderForPartition(topicPartition)
         leader.isDefined
       },
       "Reassigned partition [%s,%d] is unavailable".format(topic, partitionToBeReassigned))
-    TestUtils.waitUntilTrue(() => {
+    waitUntilTrue(() => {
         val leaderBroker = servers.filter(s => s.config.brokerId == leader.get).head
         leaderBroker.replicaManager.underReplicatedPartitionCount == 0
       },
       "Reassigned partition [%s,%d] is under-replicated as reported by the leader %d".format(topic, partitionToBeReassigned, leader.get))
   }
 
-  def verifyNonDaemonThreadsStatus(threadNamePrefix: String) {
+  // Note: Call this method in the test itself, rather than the @After method.
+  // Because of the assert, if assertNoNonDaemonThreads fails, nothing after would be executed.
+  def assertNoNonDaemonThreads(threadNamePrefix: String): Unit = {
     val threadCount = Thread.getAllStackTraces.keySet.asScala.count { t =>
       !t.isDaemon && t.isAlive && t.getName.startsWith(threadNamePrefix)
     }
@@ -1034,13 +1081,16 @@ object TestUtils extends Logging {
                                  numMessages: Int,
                                  acks: Int = -1): Seq[String] = {
     val values = (0 until numMessages).map(x =>  s"test-$x")
-    val records = values.map(v => new ProducerRecord[Array[Byte], Array[Byte]](topic, v.getBytes))
+    val intSerializer = new IntegerSerializer()
+    val records = values.zipWithIndex.map { case (v, i) =>
+      new ProducerRecord(topic, intSerializer.serialize(topic, i), v.getBytes)
+    }
     produceMessages(servers, records, acks)
     values
   }
 
   def produceMessage(servers: Seq[KafkaServer], topic: String, message: String,
-                     deliveryTimeoutMs: Int = 30 * 1000, requestTimeoutMs: Int = 20 * 1000) {
+                     deliveryTimeoutMs: Int = 30 * 1000, requestTimeoutMs: Int = 20 * 1000): Unit = {
     val producer = createProducer(TestUtils.getBrokerListStrFromServers(servers),
       deliveryTimeoutMs = deliveryTimeoutMs, requestTimeoutMs = requestTimeoutMs)
     try {
@@ -1050,39 +1100,39 @@ object TestUtils extends Logging {
     }
   }
 
-  def verifyTopicDeletion(zkClient: KafkaZkClient, topic: String, numPartitions: Int, servers: Seq[KafkaServer]) {
+  def verifyTopicDeletion(zkClient: KafkaZkClient, topic: String, numPartitions: Int, servers: Seq[KafkaServer]): Unit = {
     val topicPartitions = (0 until numPartitions).map(new TopicPartition(topic, _))
     // wait until admin path for delete topic is deleted, signaling completion of topic deletion
-    TestUtils.waitUntilTrue(() => !zkClient.isTopicMarkedForDeletion(topic),
-      "Admin path /admin/delete_topic/%s path not deleted even after a replica is restarted".format(topic))
-    TestUtils.waitUntilTrue(() => !zkClient.topicExists(topic),
-      "Topic path /brokers/topics/%s not deleted after /admin/delete_topic/%s path is deleted".format(topic, topic))
+    waitUntilTrue(() => !zkClient.isTopicMarkedForDeletion(topic),
+      "Admin path /admin/delete_topics/%s path not deleted even after a replica is restarted".format(topic))
+    waitUntilTrue(() => !zkClient.topicExists(topic),
+      "Topic path /brokers/topics/%s not deleted after /admin/delete_topics/%s path is deleted".format(topic, topic))
     // ensure that the topic-partition has been deleted from all brokers' replica managers
-    TestUtils.waitUntilTrue(() =>
-      servers.forall(server => topicPartitions.forall(tp => server.replicaManager.getPartition(tp).isEmpty)),
+    waitUntilTrue(() =>
+      servers.forall(server => topicPartitions.forall(tp => server.replicaManager.nonOfflinePartition(tp).isEmpty)),
       "Replica manager's should have deleted all of this topic's partitions")
     // ensure that logs from all replicas are deleted if delete topic is marked successful in ZooKeeper
     assertTrue("Replica logs not deleted after delete topic is complete",
       servers.forall(server => topicPartitions.forall(tp => server.getLogManager.getLog(tp).isEmpty)))
     // ensure that topic is removed from all cleaner offsets
-    TestUtils.waitUntilTrue(() => servers.forall(server => topicPartitions.forall { tp =>
+    waitUntilTrue(() => servers.forall(server => topicPartitions.forall { tp =>
       val checkpoints = server.getLogManager.liveLogDirs.map { logDir =>
         new OffsetCheckpointFile(new File(logDir, "cleaner-offset-checkpoint")).read()
       }
       checkpoints.forall(checkpointsPerLogDir => !checkpointsPerLogDir.contains(tp))
     }), "Cleaner offset for deleted partition should have been removed")
     import scala.collection.JavaConverters._
-    TestUtils.waitUntilTrue(() => servers.forall(server =>
+    waitUntilTrue(() => servers.forall(server =>
       server.config.logDirs.forall { logDir =>
         topicPartitions.forall { tp =>
           !new File(logDir, tp.topic + "-" + tp.partition).exists()
         }
       }
     ), "Failed to soft-delete the data to a delete directory")
-    TestUtils.waitUntilTrue(() => servers.forall(server =>
+    waitUntilTrue(() => servers.forall(server =>
       server.config.logDirs.forall { logDir =>
         topicPartitions.forall { tp =>
-          !java.util.Arrays.asList(new File(logDir).list()).asScala.exists { partitionDirectoryName =>
+          !Arrays.asList(new File(logDir).list()).asScala.exists { partitionDirectoryName =>
             partitionDirectoryName.startsWith(tp.topic + "-" + tp.partition) &&
               partitionDirectoryName.endsWith(Log.DeleteDirSuffix)
           }
@@ -1128,18 +1178,27 @@ object TestUtils extends Logging {
       override def getAcceptedIssuers: Array[X509Certificate] = {
         null
       }
-      override def checkClientTrusted(certs: Array[X509Certificate], authType: String) {
+      override def checkClientTrusted(certs: Array[X509Certificate], authType: String): Unit = {
       }
-      override def checkServerTrusted(certs: Array[X509Certificate], authType: String) {
+      override def checkServerTrusted(certs: Array[X509Certificate], authType: String): Unit = {
       }
     }
     trustManager
   }
 
+  def waitAndVerifyAcls(expected: Set[AccessControlEntry], authorizer: JAuthorizer, resource: ResourcePattern) = {
+    val newLine = scala.util.Properties.lineSeparator
+
+    val filter = new AclBindingFilter(resource.toFilter, AccessControlEntryFilter.ANY)
+    waitUntilTrue(() => authorizer.acls(filter).asScala.map(_.entry).toSet == expected,
+      s"expected acls:${expected.mkString(newLine + "\t", newLine + "\t", newLine)}" +
+        s"but got:${authorizer.acls(filter).asScala.map(_.entry).mkString(newLine + "\t", newLine + "\t", newLine)}", waitTimeMs = JTestUtils.DEFAULT_MAX_WAIT_MS)
+  }
+
   def waitAndVerifyAcls(expected: Set[Acl], authorizer: Authorizer, resource: Resource) = {
     val newLine = scala.util.Properties.lineSeparator
 
-    TestUtils.waitUntilTrue(() => authorizer.getAcls(resource) == expected,
+    waitUntilTrue(() => authorizer.getAcls(resource) == expected,
       s"expected acls:${expected.mkString(newLine + "\t", newLine + "\t", newLine)}" +
         s"but got:${authorizer.getAcls(resource).mkString(newLine + "\t", newLine + "\t", newLine)}", waitTimeMs = JTestUtils.DEFAULT_MAX_WAIT_MS)
   }
@@ -1182,7 +1241,7 @@ object TestUtils extends Logging {
   /**
    * Verifies that all secure paths in ZK are created with the expected ACL.
    */
-  def verifySecureZkAcls(zkClient: KafkaZkClient, usersWithAccess: Int) {
+  def verifySecureZkAcls(zkClient: KafkaZkClient, usersWithAccess: Int): Unit = {
     secureZkPaths(zkClient).foreach(path => {
       if (zkClient.pathExists(path)) {
         val sensitive = ZkData.sensitivePath(path)
@@ -1200,7 +1259,7 @@ object TestUtils extends Logging {
    * Verifies that secure paths in ZK have no access control. This is
    * the case when zookeeper.set.acl=false and no ACLs have been configured.
    */
-  def verifyUnsecureZkAcls(zkClient: KafkaZkClient) {
+  def verifyUnsecureZkAcls(zkClient: KafkaZkClient): Unit = {
     secureZkPaths(zkClient).foreach(path => {
       if (zkClient.pathExists(path)) {
         val acls = zkClient.getAcl(path)
@@ -1215,9 +1274,9 @@ object TestUtils extends Logging {
     * They all run at the same time in the assertConcurrent method; the chances of triggering a multithreading code error,
     * and thereby failing some assertion are greatly increased.
     */
-  def assertConcurrent(message: String, functions: Seq[() => Any], timeoutMs: Int) {
+  def assertConcurrent(message: String, functions: Seq[() => Any], timeoutMs: Int): Unit = {
 
-    def failWithTimeout() {
+    def failWithTimeout(): Unit = {
       fail(s"$message. Timed out, the concurrent functions took more than $timeoutMs milliseconds")
     }
 
@@ -1385,17 +1444,18 @@ object TestUtils extends Logging {
     offsetsToCommit.toMap
   }
 
-  def resetToCommittedPositions(consumer: KafkaConsumer[Array[Byte], Array[Byte]]) = {
+  def resetToCommittedPositions(consumer: KafkaConsumer[Array[Byte], Array[Byte]]) {
+    val committed = consumer.committed(consumer.assignment).asScala.filter(_._2 != null).mapValues(_.offset)
+
     consumer.assignment.asScala.foreach { topicPartition =>
-      val offset = consumer.committed(topicPartition)
-      if (offset != null)
-        consumer.seek(topicPartition, offset.offset)
+      if (committed.contains(topicPartition))
+        consumer.seek(topicPartition, committed(topicPartition))
       else
         consumer.seekToBeginning(Collections.singletonList(topicPartition))
     }
   }
 
-  def alterConfigs(servers: Seq[KafkaServer], adminClient: AdminClient, props: Properties,
+  def alterConfigs(servers: Seq[KafkaServer], adminClient: Admin, props: Properties,
                    perBrokerConfig: Boolean): AlterConfigsResult = {
     val configEntries = props.asScala.map { case (k, v) => new ConfigEntry(k, v) }.toList.asJava
     val newConfig = new Config(configEntries)
@@ -1410,13 +1470,82 @@ object TestUtils extends Logging {
     adminClient.alterConfigs(configs)
   }
 
-  def alterTopicConfigs(adminClient: AdminClient, topic: String, topicConfigs: Properties): AlterConfigsResult = {
+  def incrementalAlterConfigs(servers: Seq[KafkaServer], adminClient: Admin, props: Properties,
+                              perBrokerConfig: Boolean, opType: OpType = OpType.SET): AlterConfigsResult  = {
+    val configEntries = props.asScala.map { case (k, v) => new AlterConfigOp(new ConfigEntry(k, v), opType) }.toList.asJavaCollection
+    val configs = if (perBrokerConfig) {
+      servers.map { server =>
+        val resource = new ConfigResource(ConfigResource.Type.BROKER, server.config.brokerId.toString)
+        (resource, configEntries)
+      }.toMap.asJava
+    } else {
+      Map(new ConfigResource(ConfigResource.Type.BROKER, "") -> configEntries).asJava
+    }
+    adminClient.incrementalAlterConfigs(configs)
+  }
+
+  def alterTopicConfigs(adminClient: Admin, topic: String, topicConfigs: Properties): AlterConfigsResult = {
     val configEntries = topicConfigs.asScala.map { case (k, v) => new ConfigEntry(k, v) }.toList.asJava
     val newConfig = new Config(configEntries)
     val configs = Map(new ConfigResource(ConfigResource.Type.TOPIC, topic) -> newConfig).asJava
     adminClient.alterConfigs(configs)
   }
 
+  def assertLeader(client: Admin, topicPartition: TopicPartition, expectedLeader: Int): Unit = {
+    waitForLeaderToBecome(client, topicPartition, Some(expectedLeader))
+  }
+
+  def assertNoLeader(client: Admin, topicPartition: TopicPartition): Unit = {
+    waitForLeaderToBecome(client, topicPartition, None)
+  }
+
+  def waitForLeaderToBecome(client: Admin, topicPartition: TopicPartition, leader: Option[Int]): Unit = {
+    val topic = topicPartition.topic
+    val partition = topicPartition.partition
+
+    TestUtils.waitUntilTrue(() => {
+      try {
+        val topicResult = client.describeTopics(Arrays.asList(topic)).all.get.get(topic)
+        val partitionResult = topicResult.partitions.get(partition)
+        Option(partitionResult.leader).map(_.id) == leader
+      } catch {
+        case e: ExecutionException if e.getCause.isInstanceOf[UnknownTopicOrPartitionException] => false
+      }
+    }, "Timed out waiting for leader metadata")
+  }
+
+  def waitForBrokersOutOfIsr(client: Admin, partition: Set[TopicPartition], brokerIds: Set[Int]): Unit = {
+    TestUtils.waitUntilTrue(
+      () => {
+        val description = client.describeTopics(partition.map(_.topic).asJava).all.get.asScala
+        val isr = description
+          .values
+          .flatMap(_.partitions.asScala.flatMap(_.isr.asScala))
+          .map(_.id)
+          .toSet
+
+        brokerIds.intersect(isr).isEmpty
+      },
+      s"Expected brokers $brokerIds to no longer in the ISR for $partition"
+    )
+  }
+
+  def waitForBrokersInIsr(client: Admin, partition: TopicPartition, brokerIds: Set[Int]): Unit = {
+    TestUtils.waitUntilTrue(
+      () => {
+        val description = client.describeTopics(Set(partition.topic).asJava).all.get.asScala
+        val isr = description
+          .values
+          .flatMap(_.partitions.asScala.flatMap(_.isr.asScala))
+          .map(_.id)
+          .toSet
+
+        brokerIds.subsetOf(isr)
+      },
+      s"Expected brokers $brokerIds to be in the ISR for $partition"
+    )
+  }
+
   /**
    * Capture the console output during the execution of the provided function.
    */
@@ -1451,15 +1580,18 @@ object TestUtils extends Logging {
     (out.toString, err.toString)
   }
 
-  def assertFutureExceptionTypeEquals(future: KafkaFuture[_], clazz: Class[_ <: Throwable]): Unit = {
+  def assertFutureExceptionTypeEquals(future: KafkaFuture[_], clazz: Class[_ <: Throwable],
+                                      expectedErrorMessage: Option[String] = None): Unit = {
     try {
       future.get()
       fail("Expected CompletableFuture.get to return an exception")
     } catch {
       case e: ExecutionException =>
-        val cause = e.getCause()
+        val cause = e.getCause
         assertTrue("Expected an exception of type " + clazz.getName + "; got type " +
-            cause.getClass().getName, clazz.isInstance(cause))
+            cause.getClass.getName, clazz.isInstance(cause))
+        expectedErrorMessage.foreach(message => assertTrue(s"Received error message : ${cause.getMessage}" +
+          s" does not contain expected error message : $message", cause.getMessage.contains(message)))
     }
   }
 
@@ -1469,4 +1601,84 @@ object TestUtils extends Logging {
       .foldLeft(0.0)((total, metric) => total + metric.metricValue.asInstanceOf[Double])
     total.toLong
   }
+
+  def meterCount(metricName: String): Long = {
+    Metrics.defaultRegistry.allMetrics.asScala
+      .filterKeys(_.getMBeanName.endsWith(metricName))
+      .values
+      .headOption
+      .getOrElse(fail(s"Unable to find metric $metricName"))
+      .asInstanceOf[Meter]
+      .count
+  }
+
+  def clearYammerMetrics(): Unit = {
+    for (metricName <- Metrics.defaultRegistry.allMetrics.keySet.asScala)
+      Metrics.defaultRegistry.removeMetric(metricName)
+  }
+
+  def stringifyTopicPartitions(partitions: Set[TopicPartition]): String = {
+    Json.legacyEncodeAsString(
+      Map(
+        "partitions" -> partitions.map(tp => Map("topic" -> tp.topic, "partition" -> tp.partition))
+      )
+    )
+  }
+
+  def resource[R <: AutoCloseable, A](resource: R)(func: R => A): A = {
+    try {
+      func(resource)
+    } finally {
+      resource.close()
+    }
+  }
+
+  /**
+  * Throttles all replication across the cluster.
+  * @param adminClient is the adminClient to use for making connection with the cluster
+  * @param brokerIds all broker ids in the cluster
+  * @param throttleBytes is the target throttle
+  */
+  def throttleAllBrokersReplication(adminClient: Admin, brokerIds: Seq[Int], throttleBytes: Int): Unit = {
+    val throttleConfigs = Seq(
+      new AlterConfigOp(new ConfigEntry(DynamicConfig.Broker.LeaderReplicationThrottledRateProp, throttleBytes.toString), AlterConfigOp.OpType.SET),
+      new AlterConfigOp(new ConfigEntry(DynamicConfig.Broker.FollowerReplicationThrottledRateProp, throttleBytes.toString), AlterConfigOp.OpType.SET)
+    ).asJavaCollection
+
+    adminClient.incrementalAlterConfigs(
+      brokerIds.map { brokerId =>
+        new ConfigResource(ConfigResource.Type.BROKER, brokerId.toString) -> throttleConfigs
+      }.toMap.asJava
+    ).all().get()
+  }
+
+  def resetBrokersThrottle(adminClient: Admin, brokerIds: Seq[Int]): Unit =
+    throttleAllBrokersReplication(adminClient, brokerIds, Int.MaxValue)
+
+  def assignThrottledPartitionReplicas(adminClient: Admin, allReplicasByPartition: Map[TopicPartition, Seq[Int]]): Unit = {
+    val throttles = allReplicasByPartition.groupBy(_._1.topic()).map {
+      case (topic, replicasByPartition) =>
+        new ConfigResource(ConfigResource.Type.TOPIC, topic) -> Seq(
+          new AlterConfigOp(new ConfigEntry(LogConfig.LeaderReplicationThrottledReplicasProp, formatReplicaThrottles(replicasByPartition)), AlterConfigOp.OpType.SET),
+          new AlterConfigOp(new ConfigEntry(LogConfig.FollowerReplicationThrottledReplicasProp, formatReplicaThrottles(replicasByPartition)), AlterConfigOp.OpType.SET)
+        ).asJavaCollection
+    }
+    adminClient.incrementalAlterConfigs(throttles.asJava).all().get()
+  }
+
+  def removePartitionReplicaThrottles(adminClient: Admin, partitions: Set[TopicPartition]): Unit = {
+    val throttles = partitions.map {
+      tp =>
+        new ConfigResource(ConfigResource.Type.TOPIC, tp.topic()) -> Seq(
+          new AlterConfigOp(new ConfigEntry(LogConfig.LeaderReplicationThrottledReplicasProp, ""), AlterConfigOp.OpType.DELETE),
+          new AlterConfigOp(new ConfigEntry(LogConfig.FollowerReplicationThrottledReplicasProp, ""), AlterConfigOp.OpType.DELETE)
+        ).asJavaCollection
+    }.toMap
+    adminClient.incrementalAlterConfigs(throttles.asJava).all().get()
+  }
+
+  def formatReplicaThrottles(moves: Map[TopicPartition, Seq[Int]]): String =
+    moves.flatMap { case (tp, assignment) =>
+      assignment.map(replicaId => s"${tp.partition}:$replicaId")
+    }.mkString(",")
 }
diff --git a/core/src/test/scala/unit/kafka/utils/ThrottlerTest.scala b/core/src/test/scala/unit/kafka/utils/ThrottlerTest.scala
index d26e791ddf932..61a1e281a5ca5 100755
--- a/core/src/test/scala/unit/kafka/utils/ThrottlerTest.scala
+++ b/core/src/test/scala/unit/kafka/utils/ThrottlerTest.scala
@@ -25,7 +25,7 @@ import org.junit.Assert.{assertTrue, assertEquals}
 
 class ThrottlerTest {
   @Test
-  def testThrottleDesiredRate() {
+  def testThrottleDesiredRate(): Unit = {
     val throttleCheckIntervalMs = 100
     val desiredCountPerSec = 1000.0
     val desiredCountPerInterval = desiredCountPerSec * throttleCheckIntervalMs / 1000.0
diff --git a/core/src/test/scala/unit/kafka/utils/TopicFilterTest.scala b/core/src/test/scala/unit/kafka/utils/TopicFilterTest.scala
index 19cd1b65f8e98..4b623e978c707 100644
--- a/core/src/test/scala/unit/kafka/utils/TopicFilterTest.scala
+++ b/core/src/test/scala/unit/kafka/utils/TopicFilterTest.scala
@@ -20,12 +20,11 @@ package kafka.utils
 import org.apache.kafka.common.internals.Topic
 import org.junit.Assert._
 import org.junit.Test
-import org.scalatest.junit.JUnitSuite
 
-class TopicFilterTest extends JUnitSuite {
+class TopicFilterTest {
 
   @Test
-  def testWhitelists() {
+  def testWhitelists(): Unit = {
 
     val topicFilter1 = Whitelist("white1,white2")
     assertTrue(topicFilter1.isTopicAllowed("white2", excludeInternalTopics = true))
diff --git a/core/src/test/scala/unit/kafka/utils/ZkUtilsTest.scala b/core/src/test/scala/unit/kafka/utils/ZkUtilsTest.scala
deleted file mode 100755
index c0c3c6b8ef02b..0000000000000
--- a/core/src/test/scala/unit/kafka/utils/ZkUtilsTest.scala
+++ /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 kafka.utils
-
-import kafka.api.LeaderAndIsr
-import kafka.common.TopicAndPartition
-import kafka.controller.LeaderIsrAndControllerEpoch
-import kafka.zk.ZooKeeperTestHarness
-import org.apache.kafka.common.security.JaasUtils
-import org.junit.Assert._
-import org.junit.{After, Before, Test}
-
-@deprecated("Deprecated given that ZkUtils is deprecated", since = "2.0.0")
-class ZkUtilsTest extends ZooKeeperTestHarness {
-
-  val path = "/path"
-  var zkUtils: ZkUtils = _
-
-  @Before
-  override def setUp() {
-    super.setUp
-    zkUtils = ZkUtils(zkConnect, zkSessionTimeout, zkConnectionTimeout, zkAclsEnabled.getOrElse(JaasUtils.isZkSecurityEnabled))
-  }
-
-  @After
-  override def tearDown() {
-    if (zkUtils != null)
-     CoreUtils.swallow(zkUtils.close(), this)
-    super.tearDown
-  }
-
-  @Test
-  def testSuccessfulConditionalDeletePath() {
-    // Given an existing path
-    zkUtils.createPersistentPath(path)
-    val (_, statAfterCreation) = zkUtils.readData(path)
-
-    // Deletion is successful when the version number matches
-    assertTrue("Deletion should be successful", zkUtils.conditionalDeletePath(path, statAfterCreation.getVersion))
-    val (optionalData, _) = zkUtils.readDataMaybeNull(path)
-    assertTrue("Node should be deleted", optionalData.isEmpty)
-
-    // Deletion is successful when the node does not exist too
-    assertTrue("Deletion should be successful", zkUtils.conditionalDeletePath(path, 0))
-  }
-
-  // Verify behaviour of ZkUtils.createSequentialPersistentPath since PIDManager relies on it
-  @Test
-  def testPersistentSequentialPath() {
-    // Given an existing path
-    zkUtils.createPersistentPath(path)
-
-    var result = zkUtils.createSequentialPersistentPath(path + "/sequence_")
-
-    assertEquals("/path/sequence_0000000000", result)
-
-    result = zkUtils.createSequentialPersistentPath(path + "/sequence_")
-
-    assertEquals("/path/sequence_0000000001", result)
-  }
-
-  @Test
-  def testAbortedConditionalDeletePath() {
-    // Given an existing path that gets updated
-    zkUtils.createPersistentPath(path)
-    val (_, statAfterCreation) = zkUtils.readData(path)
-    zkUtils.updatePersistentPath(path, "data")
-
-    // Deletion is aborted when the version number does not match
-    assertFalse("Deletion should be aborted", zkUtils.conditionalDeletePath(path, statAfterCreation.getVersion))
-    val (optionalData, _) = zkUtils.readDataMaybeNull(path)
-    assertTrue("Node should still be there", optionalData.isDefined)
-  }
-
-  @Test
-  def testClusterIdentifierJsonParsing() {
-    val clusterId = "test"
-    assertEquals(zkUtils.ClusterId.fromJson(zkUtils.ClusterId.toJson(clusterId)), clusterId)
-  }
-
-  @Test
-  def testGetAllPartitionsTopicWithoutPartitions() {
-    val topic = "testtopic"
-    // Create a regular topic and a topic without any partitions
-    zkUtils.createPersistentPath(ZkUtils.getTopicPartitionPath(topic, 0))
-    zkUtils.createPersistentPath(ZkUtils.getTopicPath("nopartitions"))
-
-    assertEquals(Set(TopicAndPartition(topic, 0)), zkUtils.getAllPartitions())
-  }
-
-  @Test
-  def testGetLeaderIsrAndEpochForPartition() {
-    val topic = "my-topic-test"
-    val partition = 0
-    val leader = 1
-    val leaderEpoch = 1
-    val controllerEpoch = 1
-    val isr = List(1, 2)
-    val topicPath = s"/brokers/topics/$topic/partitions/$partition/state"
-    val topicData = Json.legacyEncodeAsString(Map("controller_epoch" -> controllerEpoch, "leader" -> leader,
-      "versions" -> 1, "leader_epoch" -> leaderEpoch, "isr" -> isr))
-    zkUtils.createPersistentPath(topicPath, topicData)
-
-    val leaderIsrAndControllerEpoch = zkUtils.getLeaderIsrAndEpochForPartition(topic, partition)
-    val topicDataLeaderIsrAndControllerEpoch = LeaderIsrAndControllerEpoch(LeaderAndIsr(leader, leaderEpoch, isr, 0),
-      controllerEpoch)
-    assertEquals(topicDataLeaderIsrAndControllerEpoch, leaderIsrAndControllerEpoch.get)
-    assertEquals(None, zkUtils.getLeaderIsrAndEpochForPartition(topic, partition + 1))
-  }
-
-  @Test
-  def testGetSequenceIdMethod() {
-    val path = "/test/seqid"
-    (1 to 10).foreach { seqid =>
-      assertEquals(seqid, zkUtils.getSequenceId(path))
-    }
-  }
-}
diff --git a/core/src/test/scala/unit/kafka/utils/json/JsonValueTest.scala b/core/src/test/scala/unit/kafka/utils/json/JsonValueTest.scala
index 640feedff46b7..04e5925d17146 100644
--- a/core/src/test/scala/unit/kafka/utils/json/JsonValueTest.scala
+++ b/core/src/test/scala/unit/kafka/utils/json/JsonValueTest.scala
@@ -17,6 +17,8 @@
 
 package kafka.utils.json
 
+import scala.collection.Seq
+
 import com.fasterxml.jackson.databind.{ObjectMapper, JsonMappingException}
 import org.junit.Test
 import org.junit.Assert._
diff --git a/core/src/test/scala/unit/kafka/utils/timer/MockTimer.scala b/core/src/test/scala/unit/kafka/utils/timer/MockTimer.scala
index 17ee578f6a56d..8805b11243fe8 100644
--- a/core/src/test/scala/unit/kafka/utils/timer/MockTimer.scala
+++ b/core/src/test/scala/unit/kafka/utils/timer/MockTimer.scala
@@ -25,7 +25,7 @@ class MockTimer extends Timer {
   val time = new MockTime
   private val taskQueue = mutable.PriorityQueue[TimerTaskEntry]()(Ordering[TimerTaskEntry].reverse)
 
-  def add(timerTask: TimerTask) {
+  def add(timerTask: TimerTask): Unit = {
     if (timerTask.delayMs <= 0)
       timerTask.run()
     else {
diff --git a/core/src/test/scala/unit/kafka/utils/timer/TimerTaskListTest.scala b/core/src/test/scala/unit/kafka/utils/timer/TimerTaskListTest.scala
index 64129e9f87b2c..0680ec202b6a7 100644
--- a/core/src/test/scala/unit/kafka/utils/timer/TimerTaskListTest.scala
+++ b/core/src/test/scala/unit/kafka/utils/timer/TimerTaskListTest.scala
@@ -33,7 +33,7 @@ class TimerTaskListTest {
   }
 
   @Test
-  def testAll() {
+  def testAll(): Unit = {
     val sharedCounter = new AtomicInteger(0)
     val list1 = new TimerTaskList(sharedCounter)
     val list2 = new TimerTaskList(sharedCounter)
diff --git a/core/src/test/scala/unit/kafka/utils/timer/TimerTest.scala b/core/src/test/scala/unit/kafka/utils/timer/TimerTest.scala
index e25a5cbdc7b3a..7b859b6db4a59 100644
--- a/core/src/test/scala/unit/kafka/utils/timer/TimerTest.scala
+++ b/core/src/test/scala/unit/kafka/utils/timer/TimerTest.scala
@@ -39,7 +39,7 @@ class TimerTest {
   private[this] var timer: Timer = null
 
   @Before
-  def setup() {
+  def setup(): Unit = {
     timer = new SystemTimer("test", tickMs = 1, wheelSize = 3)
   }
 
diff --git a/core/src/test/scala/unit/kafka/zk/AdminZkClientTest.scala b/core/src/test/scala/unit/kafka/zk/AdminZkClientTest.scala
index fa8635fd4da03..e81e0323f8bb8 100644
--- a/core/src/test/scala/unit/kafka/zk/AdminZkClientTest.scala
+++ b/core/src/test/scala/unit/kafka/zk/AdminZkClientTest.scala
@@ -19,6 +19,7 @@ package kafka.admin
 import java.util
 import java.util.Properties
 
+import kafka.controller.ReplicaAssignment
 import kafka.log._
 import kafka.server.DynamicConfig.Broker._
 import kafka.server.KafkaConfig._
@@ -35,22 +36,23 @@ import org.apache.kafka.test.{TestUtils => JTestUtils}
 import org.easymock.EasyMock
 import org.junit.Assert._
 import org.junit.{After, Test}
+import org.scalatest.Assertions.intercept
 
 import scala.collection.JavaConverters._
-import scala.collection.{Map, immutable}
+import scala.collection.{Map, Seq, immutable}
 
 class AdminZkClientTest extends ZooKeeperTestHarness with Logging with RackAwareTest {
 
   var servers: Seq[KafkaServer] = Seq()
 
   @After
-  override def tearDown() {
+  override def tearDown(): Unit = {
     TestUtils.shutdownServers(servers)
     super.tearDown()
   }
 
   @Test
-  def testManualReplicaAssignment() {
+  def testManualReplicaAssignment(): Unit = {
     val brokers = List(0, 1, 2, 3, 4)
     TestUtils.createBrokersInZk(zkClient, brokers)
 
@@ -86,11 +88,11 @@ class AdminZkClientTest extends ZooKeeperTestHarness with Logging with RackAware
                          1 -> List(1, 2, 3))
     adminZkClient.createTopicWithAssignment("test", topicConfig, assignment)
     val found = zkClient.getPartitionAssignmentForTopics(Set("test"))
-    assertEquals(assignment, found("test"))
+    assertEquals(assignment.mapValues(ReplicaAssignment(_, List(), List())).toMap, found("test"))
   }
 
   @Test
-  def testTopicCreationInZK() {
+  def testTopicCreationInZK(): Unit = {
     val expectedReplicaAssignment = Map(
       0  -> List(0, 1, 2),
       1  -> List(1, 2, 3),
@@ -138,7 +140,7 @@ class AdminZkClientTest extends ZooKeeperTestHarness with Logging with RackAware
   }
 
   @Test
-  def testTopicCreationWithCollision() {
+  def testTopicCreationWithCollision(): Unit = {
     val topic = "test.topic"
     val collidingTopic = "test_topic"
     TestUtils.createBrokersInZk(zkClient, List(0, 1, 2, 3, 4))
@@ -152,13 +154,13 @@ class AdminZkClientTest extends ZooKeeperTestHarness with Logging with RackAware
   }
 
   @Test
-  def testMockedConcurrentTopicCreation() {
+  def testMockedConcurrentTopicCreation(): Unit = {
     val topic = "test.topic"
 
     // simulate the ZK interactions that can happen when a topic is concurrently created by multiple processes
     val zkMock: KafkaZkClient = EasyMock.createNiceMock(classOf[KafkaZkClient])
     EasyMock.expect(zkMock.topicExists(topic)).andReturn(false)
-    EasyMock.expect(zkMock.getAllTopicsInCluster).andReturn(Seq("some.topic", topic, "some.other.topic"))
+    EasyMock.expect(zkMock.getAllTopicsInCluster).andReturn(Set("some.topic", topic, "some.other.topic"))
     EasyMock.replay(zkMock)
     val adminZkClient = new AdminZkClient(zkMock)
 
@@ -168,7 +170,7 @@ class AdminZkClientTest extends ZooKeeperTestHarness with Logging with RackAware
   }
 
   @Test
-  def testConcurrentTopicCreation() {
+  def testConcurrentTopicCreation(): Unit = {
     val topic = "test-concurrent-topic-creation"
     TestUtils.createBrokersInZk(zkClient, List(0, 1, 2, 3, 4))
     val props = new Properties
@@ -178,8 +180,9 @@ class AdminZkClientTest extends ZooKeeperTestHarness with Logging with RackAware
       catch { case _: TopicExistsException => () }
       val (_, partitionAssignment) = zkClient.getPartitionAssignmentForTopics(Set(topic)).head
       assertEquals(3, partitionAssignment.size)
-      partitionAssignment.foreach { case (partition, replicas) =>
-        assertEquals(s"Unexpected replication factor for $partition", 1, replicas.size)
+      partitionAssignment.foreach { case (partition, partitionReplicaAssignment) =>
+        assertEquals(s"Unexpected replication factor for $partition",
+          1, partitionReplicaAssignment.replicas.size)
       }
       val savedProps = zkClient.getEntityConfigs(ConfigType.Topic, topic)
       assertEquals(props, savedProps)
@@ -194,7 +197,7 @@ class AdminZkClientTest extends ZooKeeperTestHarness with Logging with RackAware
    * then changes the config and checks that the new values take effect.
    */
   @Test
-  def testTopicConfigChange() {
+  def testTopicConfigChange(): Unit = {
     val partitions = 3
     val topic = "my-topic"
     val server = TestUtils.createServer(KafkaConfig.fromProps(TestUtils.createBrokerConfig(0, zkConnect)))
@@ -209,7 +212,7 @@ class AdminZkClientTest extends ZooKeeperTestHarness with Logging with RackAware
       props
     }
 
-    def checkConfig(messageSize: Int, retentionMs: Long, throttledLeaders: String, throttledFollowers: String, quotaManagerIsThrottled: Boolean) {
+    def checkConfig(messageSize: Int, retentionMs: Long, throttledLeaders: String, throttledFollowers: String, quotaManagerIsThrottled: Boolean): Unit = {
       def checkList(actual: util.List[String], expected: String): Unit = {
         assertNotNull(actual)
         if (expected == "")
@@ -267,11 +270,11 @@ class AdminZkClientTest extends ZooKeeperTestHarness with Logging with RackAware
   }
 
   @Test
-  def shouldPropagateDynamicBrokerConfigs() {
+  def shouldPropagateDynamicBrokerConfigs(): Unit = {
     val brokerIds = Seq(0, 1, 2)
     servers = createBrokerConfigs(3, zkConnect).map(fromProps).map(createServer(_))
 
-    def checkConfig(limit: Long) {
+    def checkConfig(limit: Long): Unit = {
       retry(10000) {
         for (server <- servers) {
           assertEquals("Leader Quota Manager was not updated", limit, server.quotaManagers.leader.upperBound)
@@ -312,7 +315,7 @@ class AdminZkClientTest extends ZooKeeperTestHarness with Logging with RackAware
    * Basically, it asserts that notifications are bootstrapped from ZK
    */
   @Test
-  def testBootstrapClientIdConfig() {
+  def testBootstrapClientIdConfig(): Unit = {
     val clientId = "my-client"
     val props = new Properties()
     props.setProperty("producer_byte_rate", "1000")
@@ -333,7 +336,7 @@ class AdminZkClientTest extends ZooKeeperTestHarness with Logging with RackAware
   }
 
   @Test
-  def testGetBrokerMetadatas() {
+  def testGetBrokerMetadatas(): Unit = {
     // broker 4 has no rack information
     val brokerList = 0 to 5
     val rackInfo = Map(0 -> "rack1", 1 -> "rack2", 2 -> "rack2", 3 -> "rack1", 5 -> "rack3")
diff --git a/core/src/test/scala/unit/kafka/zk/EmbeddedZookeeper.scala b/core/src/test/scala/unit/kafka/zk/EmbeddedZookeeper.scala
index 32ca9690d6f03..28b592eaf7af8 100755
--- a/core/src/test/scala/unit/kafka/zk/EmbeddedZookeeper.scala
+++ b/core/src/test/scala/unit/kafka/zk/EmbeddedZookeeper.scala
@@ -48,8 +48,8 @@ class EmbeddedZookeeper() extends Logging {
   factory.startup(zookeeper)
   val port = zookeeper.getClientPort
 
-  def shutdown() {
-    CoreUtils.swallow(zookeeper.shutdown(), this)
+  def shutdown(): Unit = {
+    // Also shuts down ZooKeeperServer
     CoreUtils.swallow(factory.shutdown(), this)
 
     def isDown(): Boolean = {
@@ -60,6 +60,7 @@ class EmbeddedZookeeper() extends Logging {
     }
 
     Iterator.continually(isDown()).exists(identity)
+    CoreUtils.swallow(zookeeper.getZKDatabase().close(), this)
 
     Utils.delete(logDir)
     Utils.delete(snapshotDir)
diff --git a/core/src/test/scala/unit/kafka/zk/KafkaZkClientTest.scala b/core/src/test/scala/unit/kafka/zk/KafkaZkClientTest.scala
index 12790f5445962..f11b9ebdfe5bd 100644
--- a/core/src/test/scala/unit/kafka/zk/KafkaZkClientTest.scala
+++ b/core/src/test/scala/unit/kafka/zk/KafkaZkClientTest.scala
@@ -34,12 +34,13 @@ import org.apache.kafka.common.utils.{SecurityUtils, Time}
 import org.apache.zookeeper.KeeperException.{Code, NoNodeException, NodeExistsException}
 import org.junit.Assert._
 import org.junit.{After, Before, Test}
+import org.scalatest.Assertions.intercept
 
 import scala.collection.JavaConverters._
 import scala.collection.mutable.ArrayBuffer
 import scala.collection.{Seq, mutable}
 import scala.util.Random
-import kafka.controller.LeaderIsrAndControllerEpoch
+import kafka.controller.{LeaderIsrAndControllerEpoch, ReplicaAssignment}
 import kafka.zk.KafkaZkClient.UpdateLeaderAndIsrResult
 import kafka.zookeeper._
 import org.apache.kafka.common.errors.ControllerMovedException
@@ -85,7 +86,7 @@ class KafkaZkClientTest extends ZooKeeperTestHarness {
   private val topicPartition = new TopicPartition("topic", 0)
 
   @Test
-  def testSetAndGetConsumerOffset() {
+  def testSetAndGetConsumerOffset(): Unit = {
     val offset = 123L
     // None if no committed offsets
     assertTrue(zkClient.getConsumerOffset(group, topicPartition).isEmpty)
@@ -98,13 +99,13 @@ class KafkaZkClientTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  def testGetConsumerOffsetNoData() {
+  def testGetConsumerOffsetNoData(): Unit = {
     zkClient.createRecursive(ConsumerOffset.path(group, topicPartition.topic, topicPartition.partition))
     assertTrue(zkClient.getConsumerOffset(group, topicPartition).isEmpty)
   }
 
   @Test
-  def testDeleteRecursive() {
+  def testDeleteRecursive(): Unit = {
     zkClient.deleteRecursive("/delete/does-not-exist")
 
     zkClient.createRecursive("/delete/some/random/path")
@@ -132,7 +133,7 @@ class KafkaZkClientTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  def testCreateRecursive() {
+  def testCreateRecursive(): Unit = {
     zkClient.createRecursive("/create-newrootpath")
     assertTrue(zkClient.pathExists("/create-newrootpath"))
 
@@ -144,7 +145,7 @@ class KafkaZkClientTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  def testTopicAssignmentMethods() {
+  def testTopicAssignmentMethods(): Unit = {
     assertTrue(zkClient.getAllTopicsInCluster.isEmpty)
 
     // test with non-existing topic
@@ -168,7 +169,7 @@ class KafkaZkClientTest extends ZooKeeperTestHarness {
     val expectedAssignment = assignment map { topicAssignment =>
       val partition = topicAssignment._1.partition
       val assignment = topicAssignment._2
-      partition -> assignment
+      partition -> ReplicaAssignment(assignment, List(), List())
     }
 
     assertEquals(assignment.size, zkClient.getTopicPartitionCount(topic1).get)
@@ -178,7 +179,7 @@ class KafkaZkClientTest extends ZooKeeperTestHarness {
 
     val updatedAssignment = assignment - new TopicPartition(topic1, 2)
 
-    zkClient.setTopicAssignment(topic1, updatedAssignment)
+    zkClient.setTopicAssignment(topic1, updatedAssignment.mapValues { case v => ReplicaAssignment(v, List(), List()) }.toMap)
     assertEquals(updatedAssignment.size, zkClient.getTopicPartitionCount(topic1).get)
 
     // add second topic
@@ -193,7 +194,7 @@ class KafkaZkClientTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  def testGetDataAndVersion() {
+  def testGetDataAndVersion(): Unit = {
     val path = "/testpath"
 
     // test with non-existing path
@@ -217,7 +218,7 @@ class KafkaZkClientTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  def testConditionalUpdatePath() {
+  def testConditionalUpdatePath(): Unit = {
     val path = "/testconditionalpath"
 
     // test with non-existing path
@@ -388,7 +389,7 @@ class KafkaZkClientTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  def testSetGetAndDeletePartitionReassignment() {
+  def testSetGetAndDeletePartitionReassignment(): Unit = {
     zkClient.createRecursive(AdminZNode.path)
 
     assertEquals(Map.empty, zkClient.getPartitionReassignment)
@@ -418,7 +419,7 @@ class KafkaZkClientTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  def testGetDataAndStat() {
+  def testGetDataAndStat(): Unit = {
     val path = "/testpath"
 
     // test with non-existing path
@@ -442,7 +443,7 @@ class KafkaZkClientTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  def testGetChildren() {
+  def testGetChildren(): Unit = {
     val path = "/testpath"
 
     // test with non-existing path
@@ -460,7 +461,7 @@ class KafkaZkClientTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  def testAclManagementMethods() {
+  def testAclManagementMethods(): Unit = {
     ZkAclStore.stores.foreach(store => {
       assertFalse(zkClient.pathExists(store.aclPath))
       assertFalse(zkClient.pathExists(store.changeStore.aclChangePath))
@@ -564,7 +565,7 @@ class KafkaZkClientTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  def testDeleteTopicPathMethods() {
+  def testDeleteTopicPathMethods(): Unit = {
     assertFalse(zkClient.isTopicMarkedForDeletion(topic1))
     assertTrue(zkClient.getTopicDeletions.isEmpty)
 
@@ -598,7 +599,7 @@ class KafkaZkClientTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  def testEntityConfigManagementMethods() {
+  def testEntityConfigManagementMethods(): Unit = {
     assertTrue(zkClient.getEntityConfigs(ConfigType.Topic, topic1).isEmpty)
 
     zkClient.setOrCreateEntityConfigs(ConfigType.Topic, topic1, logProps)
@@ -647,19 +648,19 @@ class KafkaZkClientTest extends ZooKeeperTestHarness {
     val emptyConfig = LogConfig(Collections.emptyMap())
     assertEquals("Non existent config, no defaults",
       (Map(topic1 -> emptyConfig), Map.empty),
-      zkClient.getLogConfigs(Seq(topic1), Collections.emptyMap()))
+      zkClient.getLogConfigs(Set(topic1), Collections.emptyMap()))
 
     val logProps2 = createLogProps(2048)
 
     zkClient.setOrCreateEntityConfigs(ConfigType.Topic, topic1, logProps)
     assertEquals("One existing and one non-existent topic",
       (Map(topic1 -> LogConfig(logProps), topic2 -> emptyConfig), Map.empty),
-      zkClient.getLogConfigs(Seq(topic1, topic2), Collections.emptyMap()))
+      zkClient.getLogConfigs(Set(topic1, topic2), Collections.emptyMap()))
 
     zkClient.setOrCreateEntityConfigs(ConfigType.Topic, topic2, logProps2)
     assertEquals("Two existing topics",
       (Map(topic1 -> LogConfig(logProps), topic2 -> LogConfig(logProps2)), Map.empty),
-      zkClient.getLogConfigs(Seq(topic1, topic2), Collections.emptyMap()))
+      zkClient.getLogConfigs(Set(topic1, topic2), Collections.emptyMap()))
 
     val logProps1WithMoreValues = createLogProps(1024)
     logProps1WithMoreValues.put(LogConfig.SegmentJitterMsProp, "100")
@@ -667,7 +668,7 @@ class KafkaZkClientTest extends ZooKeeperTestHarness {
 
     assertEquals("Config with defaults",
       (Map(topic1 -> LogConfig(logProps1WithMoreValues)), Map.empty),
-      zkClient.getLogConfigs(Seq(topic1),
+      zkClient.getLogConfigs(Set(topic1),
         Map[String, AnyRef](LogConfig.SegmentJitterMsProp -> "100", LogConfig.SegmentBytesProp -> "128").asJava))
   }
 
@@ -784,23 +785,51 @@ class KafkaZkClientTest extends ZooKeeperTestHarness {
   val initialLeaderIsrAndControllerEpochs: Map[TopicPartition, LeaderIsrAndControllerEpoch] =
     leaderIsrAndControllerEpochs(0, 0)
 
-  val initialLeaderIsrs: Map[TopicPartition, LeaderAndIsr] = initialLeaderIsrAndControllerEpochs.mapValues(_.leaderAndIsr)
+  val initialLeaderIsrs: Map[TopicPartition, LeaderAndIsr] = initialLeaderIsrAndControllerEpochs.mapValues(_.leaderAndIsr).toMap
+
   private def leaderIsrs(state: Int, zkVersion: Int): Map[TopicPartition, LeaderAndIsr] =
-    leaderIsrAndControllerEpochs(state, zkVersion).mapValues(_.leaderAndIsr)
+    leaderIsrAndControllerEpochs(state, zkVersion).mapValues(_.leaderAndIsr).toMap
 
   private def checkUpdateLeaderAndIsrResult(
                   expectedSuccessfulPartitions: Map[TopicPartition, LeaderAndIsr],
                   expectedPartitionsToRetry: Seq[TopicPartition],
                   expectedFailedPartitions: Map[TopicPartition, (Class[_], String)],
                   actualUpdateLeaderAndIsrResult: UpdateLeaderAndIsrResult): Unit = {
-    val failedPartitionsExcerpt =
-      actualUpdateLeaderAndIsrResult.failedPartitions.mapValues(e => (e.getClass, e.getMessage))
+    val failedPartitionsExcerpt = mutable.Map.empty[TopicPartition, (Class[_], String)]
+    val successfulPartitions = mutable.Map.empty[TopicPartition, LeaderAndIsr]
+
+    actualUpdateLeaderAndIsrResult.finishedPartitions.foreach {
+      case (partition, Left(e)) => failedPartitionsExcerpt += partition -> (e.getClass, e.getMessage)
+      case (partition, Right(leaderAndIsr)) => successfulPartitions += partition -> leaderAndIsr
+    }
+
     assertEquals("Permanently failed updates do not match expected",
       expectedFailedPartitions, failedPartitionsExcerpt)
     assertEquals("Retriable updates (due to BADVERSION) do not match expected",
       expectedPartitionsToRetry, actualUpdateLeaderAndIsrResult.partitionsToRetry)
     assertEquals("Successful updates do not match expected",
-      expectedSuccessfulPartitions, actualUpdateLeaderAndIsrResult.successfulPartitions)
+      expectedSuccessfulPartitions, successfulPartitions)
+  }
+
+  @Test
+  def testTopicAssignments(): Unit = {
+    assertEquals(0, zkClient.getPartitionAssignmentForTopics(Set(topicPartition.topic())).size)
+    zkClient.createTopicAssignment(topicPartition.topic(),
+      Map(topicPartition -> Seq()))
+
+    val expectedAssignment = ReplicaAssignment(Seq(1,2,3), Seq(1), Seq(3))
+    val response = zkClient.setTopicAssignmentRaw(topicPartition.topic(),
+      Map(topicPartition -> expectedAssignment), controllerEpochZkVersion)
+    assertEquals(Code.OK, response.resultCode)
+
+    val topicPartitionAssignments = zkClient.getPartitionAssignmentForTopics(Set(topicPartition.topic()))
+    assertEquals(1, topicPartitionAssignments.size)
+    assertTrue(topicPartitionAssignments.contains(topicPartition.topic()))
+    val partitionAssignments = topicPartitionAssignments(topicPartition.topic())
+    assertEquals(1, partitionAssignments.size)
+    assertTrue(partitionAssignments.contains(topicPartition.partition()))
+    val assignment = partitionAssignments(topicPartition.partition())
+    assertEquals(expectedAssignment, assignment)
   }
 
   @Test
@@ -842,7 +871,7 @@ class KafkaZkClientTest extends ZooKeeperTestHarness {
       topicPartition20 -> LeaderAndIsr(leader = 0, leaderEpoch = 2, isr = List(3, 4), zkVersion = 0))
 
     checkUpdateLeaderAndIsrResult(
-      leaderIsrs(state = 2, zkVersion = 2).filterKeys{_ == topicPartition10},
+      leaderIsrs(state = 2, zkVersion = 2).filter { case (tp, _) => tp == topicPartition10 },
       ArrayBuffer(topicPartition11),
       Map(
         topicPartition20 -> (classOf[NoNodeException], "KeeperErrorCode = NoNode for /brokers/topics/topic2/partitions/0/state")),
@@ -1049,7 +1078,7 @@ class KafkaZkClientTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  def testClusterIdMethods() {
+  def testClusterIdMethods(): Unit = {
     val clusterId = CoreUtils.generateUuidAsBase64
 
     zkClient.createOrGetClusterId(clusterId)
@@ -1057,20 +1086,20 @@ class KafkaZkClientTest extends ZooKeeperTestHarness {
    }
 
   @Test
-  def testBrokerSequenceIdMethods() {
+  def testBrokerSequenceIdMethods(): Unit = {
     val sequenceId = zkClient.generateBrokerSequenceId()
     assertEquals(sequenceId + 1, zkClient.generateBrokerSequenceId)
   }
 
   @Test
-  def testCreateTopLevelPaths() {
+  def testCreateTopLevelPaths(): Unit = {
     zkClient.createTopLevelPaths()
 
     ZkData.PersistentZkPaths.foreach(path => assertTrue(zkClient.pathExists(path)))
   }
 
   @Test
-  def testPreferredReplicaElectionMethods() {
+  def testPreferredReplicaElectionMethods(): Unit = {
 
     assertTrue(zkClient.getPreferredReplicaElection.isEmpty)
 
@@ -1097,7 +1126,7 @@ class KafkaZkClientTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  def testDelegationTokenMethods() {
+  def testDelegationTokenMethods(): Unit = {
     assertFalse(zkClient.pathExists(DelegationTokensZNode.path))
     assertFalse(zkClient.pathExists(DelegationTokenChangeNotificationZNode.path))
 
diff --git a/core/src/test/scala/unit/kafka/zk/ReassignPartitionsZNodeTest.scala b/core/src/test/scala/unit/kafka/zk/ReassignPartitionsZNodeTest.scala
index 2f3456f4b0d0b..3ae237b2a8ac6 100644
--- a/core/src/test/scala/unit/kafka/zk/ReassignPartitionsZNodeTest.scala
+++ b/core/src/test/scala/unit/kafka/zk/ReassignPartitionsZNodeTest.scala
@@ -34,20 +34,20 @@ class ReassignPartitionsZNodeTest {
   private val reassignmentJson = """{"version":1,"partitions":[{"topic":"foo","partition":0,"replicas":[1,2]}]}"""
 
   @Test
-  def testEncode() {
+  def testEncode(): Unit = {
     val encodedJsonString = new String(ReassignPartitionsZNode.encode(reassignPartitionData), StandardCharsets.UTF_8)
     assertEquals(reassignmentJson, encodedJsonString)
   }
 
   @Test
-  def testDecodeInvalidJson() {
+  def testDecodeInvalidJson(): Unit = {
     val result = ReassignPartitionsZNode.decode("invalid json".getBytes)
     assertTrue(result.isLeft)
     assertTrue(result.left.get.isInstanceOf[JsonProcessingException])
   }
 
   @Test
-  def testDecodeValidJson() {
+  def testDecodeValidJson(): Unit = {
     val result = ReassignPartitionsZNode.decode(reassignmentJson.getBytes)
     assertTrue(result.isRight)
     val assignmentMap = result.right.get
diff --git a/core/src/test/scala/unit/kafka/zk/ZkFourLetterWords.scala b/core/src/test/scala/unit/kafka/zk/ZkFourLetterWords.scala
index bd0b257772afc..2930ac80c9314 100644
--- a/core/src/test/scala/unit/kafka/zk/ZkFourLetterWords.scala
+++ b/core/src/test/scala/unit/kafka/zk/ZkFourLetterWords.scala
@@ -28,7 +28,7 @@ import java.net.{SocketTimeoutException, Socket, InetAddress, InetSocketAddress}
   * clients, while "srvr" and "cons" give extended details on server and connections respectively.
   */
 object ZkFourLetterWords {
-  def sendStat(host: String, port: Int, timeout: Int) {
+  def sendStat(host: String, port: Int, timeout: Int): Unit = {
     val hostAddress =
       if (host != null) new InetSocketAddress(host, port)
       else new InetSocketAddress(InetAddress.getByName(null), port)
diff --git a/core/src/test/scala/unit/kafka/zk/ZooKeeperTestHarness.scala b/core/src/test/scala/unit/kafka/zk/ZooKeeperTestHarness.scala
index ebb5fb74dac45..3f995d946154b 100755
--- a/core/src/test/scala/unit/kafka/zk/ZooKeeperTestHarness.scala
+++ b/core/src/test/scala/unit/kafka/zk/ZooKeeperTestHarness.scala
@@ -22,7 +22,6 @@ import javax.security.auth.login.Configuration
 import kafka.utils.{CoreUtils, Logging, TestUtils}
 import org.junit.{After, AfterClass, Before, BeforeClass}
 import org.junit.Assert._
-import org.scalatest.junit.JUnitSuite
 import org.apache.kafka.common.security.JaasUtils
 import org.apache.kafka.test.IntegrationTest
 import org.junit.experimental.categories.Category
@@ -37,7 +36,7 @@ import org.apache.kafka.common.utils.Time
 import org.apache.zookeeper.{WatchedEvent, Watcher, ZooKeeper}
 
 @Category(Array(classOf[IntegrationTest]))
-abstract class ZooKeeperTestHarness extends JUnitSuite with Logging {
+abstract class ZooKeeperTestHarness extends Logging {
 
   val zkConnectionTimeout = 10000
   val zkSessionTimeout = 15000 // Allows us to avoid ZK session expiration due to GC up to 2/3 * 15000ms = 10 secs
@@ -54,7 +53,7 @@ abstract class ZooKeeperTestHarness extends JUnitSuite with Logging {
   def zkConnect: String = s"127.0.0.1:$zkPort"
   
   @Before
-  def setUp() {
+  def setUp(): Unit = {
     zookeeper = new EmbeddedZookeeper()
     zkClient = KafkaZkClient(zkConnect, zkAclsEnabled.getOrElse(JaasUtils.isZkSecurityEnabled), zkSessionTimeout,
       zkConnectionTimeout, zkMaxInFlightRequests, Time.SYSTEM)
@@ -62,7 +61,7 @@ abstract class ZooKeeperTestHarness extends JUnitSuite with Logging {
   }
 
   @After
-  def tearDown() {
+  def tearDown(): Unit = {
     if (zkClient != null)
      zkClient.close()
     if (zookeeper != null)
@@ -84,7 +83,7 @@ abstract class ZooKeeperTestHarness extends JUnitSuite with Logging {
 }
 
 object ZooKeeperTestHarness {
-  val ZkClientEventThreadPrefix = "ZkClient-EventThread"
+  val ZkClientEventThreadSuffix = "-EventThread"
 
   // Threads which may cause transient failures in subsequent tests if not shutdown.
   // These include threads which make connections to brokers and may cause issues
@@ -94,7 +93,7 @@ object ZooKeeperTestHarness {
                                   KafkaProducer.NETWORK_THREAD_PREFIX,
                                   AdminClientUnitTestEnv.kafkaAdminClientNetworkThreadPrefix(),
                                   AbstractCoordinator.HEARTBEAT_THREAD_PREFIX,
-                                  ZkClientEventThreadPrefix)
+                                  ZkClientEventThreadSuffix)
 
   /**
    * Verify that a previous test that doesn't use ZooKeeperTestHarness hasn't left behind an unexpected thread.
@@ -102,7 +101,7 @@ object ZooKeeperTestHarness {
    * which is true for core tests where this harness is used.
    */
   @BeforeClass
-  def setUpClass() {
+  def setUpClass(): Unit = {
     verifyNoUnexpectedThreads("@BeforeClass")
   }
 
@@ -110,7 +109,7 @@ object ZooKeeperTestHarness {
    * Verify that tests from the current test class using ZooKeeperTestHarness haven't left behind an unexpected thread
    */
   @AfterClass
-  def tearDownClass() {
+  def tearDownClass(): Unit = {
     verifyNoUnexpectedThreads("@AfterClass")
   }
 
@@ -118,7 +117,7 @@ object ZooKeeperTestHarness {
    * Verifies that threads which are known to cause transient failures in subsequent tests
    * have been shutdown.
    */
-  def verifyNoUnexpectedThreads(context: String) {
+  def verifyNoUnexpectedThreads(context: String): Unit = {
     def allThreads = Thread.getAllStackTraces.keySet.asScala.map(thread => thread.getName)
     val (threads, noUnexpected) = TestUtils.computeUntilTrue(allThreads) { threads =>
       threads.forall(t => unexpectedThreadNames.forall(s => !t.contains(s)))
diff --git a/core/src/test/scala/unit/kafka/zookeeper/ZooKeeperClientTest.scala b/core/src/test/scala/unit/kafka/zookeeper/ZooKeeperClientTest.scala
index fd3f59cefc2ef..0587def8c16d5 100644
--- a/core/src/test/scala/unit/kafka/zookeeper/ZooKeeperClientTest.scala
+++ b/core/src/test/scala/unit/kafka/zookeeper/ZooKeeperClientTest.scala
@@ -21,6 +21,8 @@ import java.util.UUID
 import java.util.concurrent.atomic.{AtomicBoolean, AtomicInteger}
 import java.util.concurrent.{ArrayBlockingQueue, ConcurrentLinkedQueue, CountDownLatch, Executors, Semaphore, TimeUnit}
 
+import scala.collection.Seq
+
 import com.yammer.metrics.Metrics
 import com.yammer.metrics.core.{Gauge, Meter, MetricName}
 import kafka.zk.ZooKeeperTestHarness
@@ -32,6 +34,7 @@ import org.apache.zookeeper.ZooKeeper.States
 import org.apache.zookeeper.{CreateMode, WatchedEvent, ZooDefs}
 import org.junit.Assert.{assertArrayEquals, assertEquals, assertFalse, assertTrue}
 import org.junit.{After, Before, Test}
+import org.scalatest.Assertions.{fail, intercept}
 
 import scala.collection.JavaConverters._
 
@@ -42,7 +45,8 @@ class ZooKeeperClientTest extends ZooKeeperTestHarness {
   private var zooKeeperClient: ZooKeeperClient = _
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
+    ZooKeeperTestHarness.verifyNoUnexpectedThreads("@Before")
     cleanMetricsRegistry()
     super.setUp()
     zooKeeperClient = new ZooKeeperClient(zkConnect, zkSessionTimeout, zkConnectionTimeout, zkMaxInFlightRequests,
@@ -50,11 +54,12 @@ class ZooKeeperClientTest extends ZooKeeperTestHarness {
   }
 
   @After
-  override def tearDown() {
+  override def tearDown(): Unit = {
     if (zooKeeperClient != null)
       zooKeeperClient.close()
     super.tearDown()
     System.clearProperty(JaasUtils.JAVA_LOGIN_CONFIG_PARAM)
+    ZooKeeperTestHarness.verifyNoUnexpectedThreads("@After")
   }
 
   @Test
@@ -82,8 +87,16 @@ class ZooKeeperClientTest extends ZooKeeperTestHarness {
 
   @Test
   def testConnection(): Unit = {
-    new ZooKeeperClient(zkConnect, zkSessionTimeout, zkConnectionTimeout, Int.MaxValue, time, "testMetricGroup",
-      "testMetricType").close()
+    val client = new ZooKeeperClient(zkConnect, zkSessionTimeout, zkConnectionTimeout, Int.MaxValue, time, "testMetricGroup",
+      "testMetricType")
+    try {
+      // Verify ZooKeeper event thread name. This is used in ZooKeeperTestHarness to verify that tests have closed ZK clients
+      val threads = Thread.getAllStackTraces.keySet.asScala.map(_.getName)
+      assertTrue(s"ZooKeeperClient event thread not found, threads=$threads",
+        threads.exists(_.contains(ZooKeeperTestHarness.ZkClientEventThreadSuffix)))
+    } finally {
+      client.close()
+    }
   }
 
   @Test
@@ -327,14 +340,15 @@ class ZooKeeperClientTest extends ZooKeeperTestHarness {
       }
     }
 
-    val client = new ZooKeeperClient(zkConnect, zkSessionTimeout, zkConnectionTimeout, Int.MaxValue, time,
+    zooKeeperClient.close()
+    zooKeeperClient = new ZooKeeperClient(zkConnect, zkSessionTimeout, zkConnectionTimeout, Int.MaxValue, time,
       "testMetricGroup", "testMetricType")
-    client.registerStateChangeHandler(stateChangeHandler)
+    zooKeeperClient.registerStateChangeHandler(stateChangeHandler)
 
     val requestThread = new Thread() {
       override def run(): Unit = {
         try
-          client.handleRequest(CreateRequest(mockPath, Array.empty[Byte],
+          zooKeeperClient.handleRequest(CreateRequest(mockPath, Array.empty[Byte],
             ZooDefs.Ids.OPEN_ACL_UNSAFE.asScala, CreateMode.PERSISTENT))
         finally
           latch.countDown()
@@ -343,7 +357,7 @@ class ZooKeeperClientTest extends ZooKeeperTestHarness {
 
     val reinitializeThread = new Thread() {
       override def run(): Unit = {
-        client.forceReinitialize()
+        zooKeeperClient.forceReinitialize()
       }
     }
 
@@ -375,12 +389,13 @@ class ZooKeeperClientTest extends ZooKeeperTestHarness {
       }
     }
 
-    val client = new ZooKeeperClient(zkConnect, zkSessionTimeout, zkConnectionTimeout, Int.MaxValue, time,
+    zooKeeperClient.close()
+    zooKeeperClient = new ZooKeeperClient(zkConnect, zkSessionTimeout, zkConnectionTimeout, Int.MaxValue, time,
       "testMetricGroup", "testMetricType")
-    client.registerStateChangeHandler(faultyHandler)
-    client.registerStateChangeHandler(goodHandler)
+    zooKeeperClient.registerStateChangeHandler(faultyHandler)
+    zooKeeperClient.registerStateChangeHandler(goodHandler)
 
-    client.forceReinitialize()
+    zooKeeperClient.forceReinitialize()
 
     assertEquals(1, goodCalls.get)
 
@@ -586,6 +601,7 @@ class ZooKeeperClientTest extends ZooKeeperTestHarness {
         }
       })
       assertFalse("Close completed without shutting down expiry scheduler gracefully", closeFuture.isDone)
+      assertTrue(zooKeeperClient.currentZooKeeper.getState.isAlive) // Client should be closed after expiry handler
       semaphore.release()
       closeFuture.get(10, TimeUnit.SECONDS)
       assertFalse("Expiry executor not shutdown", zooKeeperClient.expiryScheduler.isStarted)
@@ -598,8 +614,8 @@ class ZooKeeperClientTest extends ZooKeeperTestHarness {
     metricName.getName == name && metricName.getGroup == "testMetricGroup" && metricName.getType == "testMetricType"
 
   @Test
-  def testZooKeeperStateChangeRateMetrics() {
-    def checkMeterCount(name: String, expected: Long) {
+  def testZooKeeperStateChangeRateMetrics(): Unit = {
+    def checkMeterCount(name: String, expected: Long): Unit = {
       val meter = Metrics.defaultRegistry.allMetrics.asScala.collectFirst {
         case (metricName, meter: Meter) if isExpectedMetricName(metricName, name) => meter
       }.getOrElse(sys.error(s"Unable to find meter with name $name"))
@@ -637,7 +653,7 @@ class ZooKeeperClientTest extends ZooKeeperTestHarness {
     assertEquals(States.CLOSED, zooKeeperClient.connectionState)
   }
 
-  private def cleanMetricsRegistry() {
+  private def cleanMetricsRegistry(): Unit = {
     val metrics = Metrics.defaultRegistry
     metrics.allMetrics.keySet.asScala.foreach(metrics.removeMetric)
   }
diff --git a/doap_Kafka.rdf b/doap_Kafka.rdf
index 88e50a673183f..079d0c36619be 100644
--- a/doap_Kafka.rdf
+++ b/doap_Kafka.rdf
@@ -2,9 +2,9 @@
 
 
+         xmlns:rdf="https://www.w3.org/1999/02/22-rdf-syntax-ns#"
+         xmlns:asfext="https://projects.apache.org/ns/asfext#"
+         xmlns:foaf="https://xmlns.com/foaf/0.1/">
 
-  
+  
     2014-04-12
     
     Apache Kafka
-    
-    
+    
+    
     Apache Kafka is a distributed, fault tolerant, publish-subscribe messaging.
     A single Kafka broker can handle hundreds of megabytes of reads and writes per second from thousands of clients. Kafka is designed to allow a single cluster to serve as the central data backbone for a large organization. It can be elastically and transparently expanded without downtime. Data streams are partitioned and spread over a cluster of machines to allow data streams larger than the capability of any single machine and to allow clusters of co-ordinated consumers. Kafka has a modern cluster-centric design that offers strong durability and fault-tolerance guarantees. Messages are persisted on disk and replicated within the cluster to prevent data loss. Each broker can handle terabytes of messages without performance impact.
     
-    
-    
+    
+    
     Scala
-    
+    
     
       
-        
+        
         
       
     
diff --git a/docs/api.html b/docs/api.html
index 4545e9b95860b..b6ab1fa4c3ee1 100644
--- a/docs/api.html
+++ b/docs/api.html
@@ -21,7 +21,7 @@
 	
            2. The Consumer API allows applications to read streams of data from topics in the Kafka cluster.
 	
            3. The Streams API allows transforming streams of data from input topics to output topics.
 	
            4. The Connect API allows implementing connectors that continually pull from some source system or application into Kafka or push from Kafka into some sink system or application.
-	
            5. The AdminClient API allows managing and inspecting topics, brokers, and other Kafka objects.
+	
            6. The Admin API allows managing and inspecting topics, brokers, and other Kafka objects.
 	
            
 
 	Kafka exposes all its functionality over a language independent protocol which has clients available in many programming languages. However only the Java clients are maintained as part of the main Kafka project, the others are available as independent open source projects. A list of non-Java clients is available here.
@@ -100,11 +100,11 @@ 
            2.4 Connect API
            
 	Those who want to implement custom connectors can see the javadoc.
 	
            
 
-	
            2.5 AdminClient API
            
+	
            2.5 Admin API
            
 
-	The AdminClient API supports managing and inspecting topics, brokers, acls, and other Kafka objects.
+	The Admin API supports managing and inspecting topics, brokers, acls, and other Kafka objects.
 	
            
-	To use the AdminClient API, add the following Maven dependency:
+	To use the Admin API, add the following Maven dependency:
 	
             		<dependency>
 			<groupId>org.apache.kafka</groupId>
@@ -112,7 +112,7 @@ 
            2.5 AdminClient API
            
 			<version>{{fullDotVersion}}</version>
 		</dependency>
 	
            
-	For more information about the AdminClient APIs, see the javadoc.
+	For more information about the Admin APIs, see the javadoc.
 	
            
 
 
diff --git a/docs/configuration.html b/docs/configuration.html
index 06d7585c1ef5e..dc17333e9dbb3 100644
--- a/docs/configuration.html
+++ b/docs/configuration.html
@@ -156,6 +156,7 @@ 
            Updating Default Topic Configuration
            
     
          4. log.index.interval.bytes
            5. 
     
          5. log.cleaner.delete.retention.ms
            6. 
     
          6. log.cleaner.min.compaction.lag.ms
            7. 
+    
          7. log.cleaner.max.compaction.lag.ms
            8. 
     
          8. log.cleaner.min.cleanable.ratio
            9. 
     
          9. log.cleanup.policy
            10. 
     
          10. log.segment.delete.delay.ms
            11. 
@@ -284,7 +285,7 @@ 
            3.6 Kafka Streams Configs<
   Below is the configuration of the Kafka Streams client library.
   
 
-  
            3.7 AdminClient Configs
            
+  
            3.7 Admin Configs
            
   Below is the configuration of the Kafka Admin client library.
   
 
diff --git a/docs/connect.html b/docs/connect.html
index baefc1b417e4f..f1a9d05f2e4d5 100644
--- a/docs/connect.html
+++ b/docs/connect.html
@@ -56,7 +56,9 @@ 
            Running Kafka Connectoffset.storage.file.filename - File to store offset data in
     
    
 
-    
    The parameters that are configured here are intended for producers and consumers used by Kafka Connect to access the configuration, offset and status topics. For configuration of Kafka source and Kafka sink tasks, the same parameters can be used but need to be prefixed with consumer. and producer. respectively. The only parameter that is inherited from the worker configuration is bootstrap.servers, which in most cases will be sufficient, since the same cluster is often used for all purposes. A notable exception is a secured cluster, which requires extra parameters to allow connections. These parameters will need to be set up to three times in the worker configuration, once for management access, once for Kafka sinks and once for Kafka sources.
    
+    
    The parameters that are configured here are intended for producers and consumers used by Kafka Connect to access the configuration, offset and status topics. For configuration of the producers used by Kafka source tasks and the consumers used by Kafka sink tasks, the same parameters can be used but need to be prefixed with producer. and consumer. respectively. The only Kafka client parameter that is inherited without a prefix from the worker configuration is bootstrap.servers, which in most cases will be sufficient, since the same cluster is often used for all purposes. A notable exception is a secured cluster, which requires extra parameters to allow connections. These parameters will need to be set up to three times in the worker configuration, once for management access, once for Kafka sources and once for Kafka sinks.
     
+    
+    
    Starting with 2.3.0, client configuration overrides can be configured individually per connector by using the prefixes producer.override. and consumer.override. for Kafka sources or Kafka sinks respectively. These overrides are included with the rest of the connector's configuration properties.
    
 
     
    The remaining parameters are connector configuration files. You may include as many as you want, but all will execute within the same process (on different threads).
    
 
@@ -248,6 +250,12 @@ 
    REST API
    
         
  • PUT /connector-plugins/{connector-type}/config/validate - validate the provided configuration values against the configuration definition. This API performs per config validation, returns suggested values and error messages during validation.
    • 
     
 
+    
    The following is a supported REST request at the top-level (root) endpoint:
    
+
+    
      
+        
    • GET /- return basic information about the Kafka Connect cluster such as the version of the Connect worker that serves the REST request (including git commit ID of the source code) and the Kafka cluster ID that is connected to.
+    
    
+
     
    8.3 Connector Development Guide
    
 
     
    This guide describes how developers can write new connectors for Kafka Connect to move data between Kafka and other systems. It briefly reviews a few key concepts and then describes how to create a simple connector.
    
@@ -487,7 +495,7 @@ 
    Working with Schemas
    Kafka Connect
     
    
 
     
    
-    When a connector is first submitted to the cluster, the workers rebalance the full set of connectors in the cluster and their tasks so that each worker has approximately the same amount of work. This same rebalancing procedure is also used when connectors increase or decrease the number of tasks they require, or when a connector's configuration is changed. You can use the REST API to view the current status of a connector and its tasks, including the id of the worker to which each was assigned. For example, querying the status of a file source (using GET /connectors/file-source/status) might produce output like the following:
+    When a connector is first submitted to the cluster, a rebalance is triggered between the Connect workers in order to distribute the load that consists of the tasks of the new connector.
+    This same rebalancing procedure is also used when connectors increase or decrease the number of tasks they require, when a connector's configuration is changed, or when a
+    worker is added or removed from the group as part of an intentional upgrade of the Connect cluster or due to a failure.
+    
    
+
+    
    
+    In versions prior to 2.3.0, the Connect workers would rebalance the full set of connectors and their tasks in the cluster as a simple way to make sure that each worker has approximately the same amount of work.
+    This behavior can be still enabled by setting connect.protocol=eager.
+    
    
+
+    
    
+    Starting with 2.3.0, Kafka Connect is using by default a protocol that performs
+    incremental cooperative rebalancing
+    that incrementally balances the connectors and tasks across the Connect workers, affecting only tasks that are new, to be removed, or need to move from one worker to another.
+    Other tasks are not stopped and restarted during the rebalance, as they would have been with the old protocol.
+    
    
+
+    
    
+    If a Connect worker leaves the group, intentionally or due to a failure, Connect waits for scheduled.rebalance.max.delay.ms before triggering a rebalance.
+    This delay defaults to five minutes (300000ms) to tolerate failures or upgrades of workers without immediately redistributing the load of a departing worker.
+    If this worker returns within the configured delay, it gets its previously assigned tasks in full.
+    However, this means that the tasks will remain unassigned until the time specified by scheduled.rebalance.max.delay.ms elapses.
+    If a worker does not return within that time limit, Connect will reassign those tasks among the remaining workers in the Connect cluster.
+    
    
+
+    
    
+    The new Connect protocol is enabled when all the workers that form the Connect cluster are configured with connect.protocol=compatible, which is also the default value when this property is missing.
+    Therefore, upgrading to the new Connect protocol happens automatically when all the workers upgrade to 2.3.0.
+    A rolling upgrade of the Connect cluster will activate incremental cooperative rebalancing when the last worker joins on version 2.3.0.
+    
    
+
+    
    
+    You can use the REST API to view the current status of a connector and its tasks, including the ID of the worker to which each was assigned. For example, the GET /connectors/file-source/status request shows the status of a connector named file-source:
     
    
 
     
    @@ -537,6 +577,7 @@ 
    Kafka Connect
     
  • RUNNING: The connector/task is running.
    • 
     
  • PAUSED: The connector/task has been administratively paused.
    • 
     
  • FAILED: The connector/task has failed (usually by raising an exception, which is reported in the status output).
    • 
+    
  • DESTROYED: The connector/task has been administratively removed and will stop appearing in the Connect cluster.
    • 
     
 
     
    
diff --git a/docs/design.html b/docs/design.html
index ab8c002c9d3e9..4c4e796be3e04 100644
--- a/docs/design.html
+++ b/docs/design.html
@@ -125,7 +125,7 @@ 
    4.3 Efficiency
     
    
     This combination of pagecache and sendfile means that on a Kafka cluster where the consumers are mostly caught up you will see no read activity on the disks whatsoever as they will be serving data entirely from cache.
     
    
-    For more background on the sendfile and zero-copy support in Java, see this article.
+    For more background on the sendfile and zero-copy support in Java, see this article.
 
     
    End-to-end Batch Compression
    
     
    
@@ -216,6 +216,28 @@ 
    Offline Data Load<
     In the case of Hadoop we parallelize the data load by splitting the load over individual map tasks, one for each node/topic/partition combination, allowing full parallelism in the loading. Hadoop provides the task
     management, and tasks which fail can restart without danger of duplicate data—they simply restart from their original position.
 
+    
    Static Membership
    
+    Static membership aims to improve the availability of stream applications, consumer groups and other applications built on top of the group rebalance protocol.
+    The rebalance protocol relies on the group coordinator to allocate entity ids to group members. These generated ids are ephemeral and will change when members restart and rejoin.
+    For consumer based apps, this "dynamic membership" can cause a large percentage of tasks re-assigned to different instances during administrative operations
+    such as code deploys, configuration updates and periodic restarts. For large state applications, shuffled tasks need a long time to recover their local states before processing
+    and cause applications to be partially or entirely unavailable. Motivated by this observation, Kafka’s group management protocol allows group members to provide persistent entity ids.
+    Group membership remains unchanged based on those ids, thus no rebalance will be triggered.
+    
    
+    If you want to use static membership,
+    
      
+        
    • Upgrade both broker cluster and client apps to 2.3 or beyond, and also make sure the upgraded brokers are using inter.broker.protocol.version
+            of 2.3 or beyond as well.
      • 
+        
    • Set the config ConsumerConfig#GROUP_INSTANCE_ID_CONFIG to a unique value for each consumer instance under one group.
      • 
+        
    • For Kafka Streams applications, it is sufficient to set a unique ConsumerConfig#GROUP_INSTANCE_ID_CONFIG per KafkaStreams instance,
+            independent of the number of used threads for an instance.
      • 
+    
    
+    If your broker is on an older version than 2.3, but you choose to set ConsumerConfig#GROUP_INSTANCE_ID_CONFIG on the client side, the application will detect
+    the broker version and then throws an UnsupportedException. If you accidentally configure duplicate ids for different instances,
+    a fencing mechanism on broker side will inform your duplicate client to shutdown immediately by triggering a org.apache.kafka.common.errors.FencedInstanceIdException.
+    For more details, see
+    KIP-345
+
     
    4.6 Message Delivery Semantics
    
     
    
     Now that we understand a little about how producers and consumers work, let's discuss the semantic guarantees Kafka provides between producer and consumer. Clearly there are multiple possible message delivery
@@ -501,6 +523,7 @@ 
    What
     
      
     
    1. Any consumer that stays caught-up to within the head of the log will see every message that is written; these messages will have sequential offsets. The topic's min.compaction.lag.ms can be used to
     guarantee the minimum length of time must pass after a message is written before it could be compacted. I.e. it provides a lower bound on how long each message will remain in the (uncompacted) head.
+    The topic's max.compaction.lag.ms can be used to guarantee the maximum delay between the time a message is written and the time the message becomes eligible for compaction.
     
    2. Ordering of messages is always maintained.  Compaction will never re-order messages, just remove some.
     
    3. The offset for a message never changes.  It is the permanent identifier for a position in the log.
     
    4. Any consumer progressing from the start of the log will see at least the final state of all records in the order they were written.  Additionally, all delete markers for deleted records will be seen, provided
@@ -523,16 +546,27 @@ 
      Log Compac
     
      Configuring The Log Cleaner
      
 
     The log cleaner is enabled by default. This will start the pool of cleaner threads.
-    To enable log cleaning on a particular topic you can add the log-specific property
-    
        log.cleanup.policy=compact
      
-    This can be done either at topic creation time or using the alter topic command.
-    
      
+    To enable log cleaning on a particular topic, add the log-specific property
+    
       log.cleanup.policy=compact
      
+
+    The log.cleanup.policy property is a broker configuration setting defined
+    in the broker's server.properties file; it affects all of the topics
+    in the cluster that do not have a configuration override in place as documented
+    here.
+
     The log cleaner can be configured to retain a minimum amount of the uncompacted "head" of the log. This is enabled by setting the compaction time lag.
     
        log.cleaner.min.compaction.lag.ms
      
 
     This can be used to prevent messages newer than a minimum message age from being subject to compaction. If not set, all log segments are eligible for compaction except for the last segment, i.e. the one currently
     being written to. The active segment will not be compacted even if all of its messages are older than the minimum compaction time lag.
-    
      
+
+    The log cleaner can be configured to ensure a maximum delay after which the uncompacted "head" of the log becomes eligible for log compaction.
+    
        log.cleaner.max.compaction.lag.ms
      
+
+    This can be used to prevent log with low produce rate from remaining ineligible for compaction for an unbounded duration. If not set, logs that do not exceed min.cleanable.dirty.ratio are not compacted.
+    Note that this compaction deadline is not a hard guarantee since it is still subjected to the availability of log cleaner threads and the actual compaction time.
+    You will want to monitor the uncleanable-partitions-count, max-clean-time-secs and max-compaction-delay-secs metrics.
+
     
      
     Further cleaner configurations are described here.
 
diff --git a/docs/documentation.html b/docs/documentation.html
index 83c029bb3513b..8decc9bea5cf3 100644
--- a/docs/documentation.html
+++ b/docs/documentation.html
@@ -26,8 +26,8 @@
 	
      
 		
     
      Documentation
      
-    
      Kafka 2.2 Documentation
      
-    Prior releases: 0.7.x, 0.8.0, 0.8.1.X, 0.8.2.X, 0.9.0.X, 0.10.0.X, 0.10.1.X, 0.10.2.X, 0.11.0.X, 1.0.X, 1.1.X, 2.0.X, 2.1.X.
+    
      Kafka 2.4 Documentation
      
+    Prior releases: 0.7.x, 0.8.0, 0.8.1.X, 0.8.2.X, 0.9.0.X, 0.10.0.X, 0.10.1.X, 0.10.2.X, 0.11.0.X, 1.0.X, 1.1.X, 2.0.X, 2.1.X, 2.2.X, 2.3.X.
 
     
 
diff --git a/docs/documentation/streams/developer-guide/dsl-topology-naming.html b/docs/documentation/streams/developer-guide/dsl-topology-naming.html
new file mode 100644
index 0000000000000..db5eee368df7c
--- /dev/null
+++ b/docs/documentation/streams/developer-guide/dsl-topology-naming.html
@@ -0,0 +1,19 @@
+
+
+
+
diff --git a/docs/js/templateData.js b/docs/js/templateData.js
index e5d5f0b8d2461..66123b93d42be 100644
--- a/docs/js/templateData.js
+++ b/docs/js/templateData.js
@@ -17,8 +17,8 @@ limitations under the License.
 
 // Define variables for doc templates
 var context={
-    "version": "23",
-    "dotVersion": "2.3",
-    "fullDotVersion": "2.3.0",
+    "version": "24",
+    "dotVersion": "2.4",
+    "fullDotVersion": "2.4.2-SNAPSHOT",
     "scalaVersion": "2.12"
 };
diff --git a/docs/ops.html b/docs/ops.html
index 5ba9deff1203e..6f84186b299ae 100644
--- a/docs/ops.html
+++ b/docs/ops.html
@@ -453,7 +453,7 @@ 
      Limiting Bandwidth Usage during Da
   There are two interfaces that can be used to engage a throttle. The simplest, and safest, is to apply a throttle when invoking the kafka-reassign-partitions.sh, but kafka-configs.sh can also be used to view and alter the throttle values directly.
   
      
   So for example, if you were to execute a rebalance, with the below command, it would move partitions at no more than 50MB/s.
-  
      $ bin/kafka-reassign-partitions.sh --zookeeper myhost:2181--execute --reassignment-json-file bigger-cluster.json —throttle 50000000
      
+  
      $ bin/kafka-reassign-partitions.sh --zookeeper localhost:2181 --execute --reassignment-json-file bigger-cluster.json —throttle 50000000
      
   When you execute this script you will see the throttle engage:
   
         The throttle limit was set to 50000000 B/s
@@ -786,6 +786,16 @@ 
      6.6 Monitoring
      
   records-consumed-rate has a corresponding metric named records-consumed-total.
   
      
   The easiest way to see the available metrics is to fire up jconsole and point it at a running kafka client or server; this will allow browsing all metrics with JMX.
+
+  
      Security Considerations for Remote Monitoring using JMX
      
+  Apache Kafka disables remote JMX by default. You can enable remote monitoring using JMX by setting the environment variable
+  JMX_PORT for processes started using the CLI or standard Java system properties to enable remote JMX programmatically.
+  You must enable security when enabling remote JMX in production scenarios to ensure that unauthorized users cannot monitor or
+  control your broker or application as well as the platform on which these are running. Note that authentication is disabled for
+  JMX by default in Kafka and security configs must be overridden for production deployments by setting the environment variable
+  KAFKA_JMX_OPTS for processes started using the CLI or by setting appropriate Java system properties. See
+  Monitoring and Management Using JMX Technology
+  for details on securing JMX.
   
      
   We do graphing and alerting on the following metrics:
   
@@ -850,6 +860,26 @@ 
      6.6 Monitoring
      
+      
+        
+        
+        
+      
+      
+        
+        
+        
+      
+      
+        
+        
+        
+      
+      
+        
+        
+        
+      
@@ -865,6 +895,11 @@ 
      6.6 Monitoring
      
+      
+        
+        
+        
+      
@@ -885,6 +920,26 @@ 
      6.6 Monitoring
      
+      
+        
+        
+        
+      
+      
+        
+        
+        
+      
+      
+        
+        
+        
+      
+      
+        
+        
+        
+      
@@ -1014,6 +1069,26 @@ 
      6.6 Monitoring
      
+      
+        
+        
+        
+      
+      
+        
+        
+        
+      
+      
+        
+        
+        
+      
+      
+        
+        
+        
+      
      
         
      kafka.server:type=BrokerTopicMetrics,name=ReplicationBytesOutPerSec
         
       
      Message validation failure rate due to no key specified for compacted topickafka.server:type=BrokerTopicMetrics,name=NoKeyCompactedTopicRecordsPerSec
      Message validation failure rate due to invalid magic numberkafka.server:type=BrokerTopicMetrics,name=InvalidMagicNumberRecordsPerSec
      Message validation failure rate due to incorrect crc checksumkafka.server:type=BrokerTopicMetrics,name=InvalidMessageCrcRecordsPerSec
      Message validation failure rate due to non-continuous offset or sequence number in batchkafka.server:type=BrokerTopicMetrics,name=InvalidOffsetOrSequenceRecordsPerSec
      
       
      
         Log flush rate and time
         kafka.log:type=LogFlushStats,name=LogFlushRateAndTimeMs
         kafka.server:type=ReplicaManager,name=UnderMinIsrPartitionCount
         0
       
      # of at minIsr partitions (|ISR| = min.insync.replicas)kafka.server:type=ReplicaManager,name=AtMinIsrPartitionCount0
      
       
      
         # of offline log directories
         kafka.log:type=LogManager,name=OfflineLogDirectoryCount
         kafka.controller:type=ControllerStats,name=UncleanLeaderElectionsPerSec
         0
       
      Pending topic deleteskafka.controller:type=KafkaController,name=TopicsToDeleteCount
      Pending replica deleteskafka.controller:type=KafkaController,name=ReplicasToDeleteCount
      Ineligible pending topic deleteskafka.controller:type=KafkaController,name=TopicsIneligibleToDeleteCount
      Ineligible pending replica deleteskafka.controller:type=KafkaController,name=ReplicasIneligibleToDeleteCount
      
       
      
         Partition counts
         kafka.server:type=ReplicaManager,name=PartitionCount
         Connection status of broker's ZooKeeper session which may be one of
             Disconnected|SyncConnected|AuthFailed|ConnectedReadOnly|SaslAuthenticated|Expired.
       
      Max time to load group metadatakafka.server:type=group-coordinator-metrics,name=partition-load-time-maxmaximum time, in milliseconds, it took to load offsets and group metadata from the consumer offset partitions loaded in the last 30 seconds
      Avg time to load group metadatakafka.server:type=group-coordinator-metrics,name=partition-load-time-avgaverage time, in milliseconds, it took to load offsets and group metadata from the consumer offset partitions loaded in the last 30 seconds
      Max time to load transaction metadatakafka.server:type=transaction-coordinator-metrics,name=partition-load-time-maxmaximum time, in milliseconds, it took to load transaction metadata from the consumer offset partitions loaded in the last 30 seconds
      Avg time to load transaction metadatakafka.server:type=transaction-coordinator-metrics,name=partition-load-time-avgaverage time, in milliseconds, it took to load transaction metadata from the consumer offset partitions loaded in the last 30 seconds
      
   
      
 
   
      Common monitoring metrics for producer/consumer/connect/streams
      
@@ -1316,6 +1391,36 @@ 
      consumer monitoring<
 
   The following metrics are available on consumer instances.
 
+  
+    
+      
+        
+        
+        
+      
+      
+        
+        
+        
+      
+      
+        
+        
+        
+      
+      
+        
+        
+        
+      
+      
+        
+        
+        
+      
+    
+  
      Metric/Attribute nameDescriptionMbean name
      time-between-poll-avgThe average delay between invocations of poll().kafka.consumer:type=consumer-metrics,client-id=([-.\w]+)
      time-between-poll-maxThe max delay between invocations of poll().kafka.consumer:type=consumer-metrics,client-id=([-.\w]+)
      last-poll-seconds-agoThe number of seconds since the last poll() invocation.kafka.consumer:type=consumer-metrics,client-id=([-.\w]+)
      poll-idle-ratio-avgThe average fraction of time the consumer's poll() is idle as opposed to waiting for the user code to process records.kafka.consumer:type=consumer-metrics,client-id=([-.\w]+)
      
+
   
      Consumer Group Metrics
      
   
@@ -1404,11 +1509,81 @@ 
      Consumer
         
      
+      
+        
+        
+        
+      
+      
+        
+        
+        
+      
+      
+        
+        
+        
+      
+      
+        
+        
+        
+      
+      
+        
+        
+        
+      
+      
+        
+        
+        
+      
+      
+        
+        
+        
+      
+      
+        
+        
+        
+      
+      
+        
+        
+        
+      
+      
+        
+        
+        
+      
+      
+        
+        
+        
+      
+      
+        
+        
+        
+      
+      
+        
+        
+        
+      
+      
+        
+        
+        
+      
      
     
      The total number of group syncs
         kafka.consumer:type=consumer-coordinator-metrics,client-id=([-.\w]+)
       
      rebalance-latency-avgThe average time taken for a group rebalancekafka.consumer:type=consumer-coordinator-metrics,client-id=([-.\w]+)
      rebalance-latency-maxThe max time taken for a group rebalancekafka.consumer:type=consumer-coordinator-metrics,client-id=([-.\w]+)
      rebalance-latency-totalThe total time taken for group rebalances so farkafka.consumer:type=consumer-coordinator-metrics,client-id=([-.\w]+)
      rebalance-totalThe total number of group rebalances participatedkafka.consumer:type=consumer-coordinator-metrics,client-id=([-.\w]+)
      rebalance-rate-per-hourThe number of group rebalance participated per hourkafka.consumer:type=consumer-coordinator-metrics,client-id=([-.\w]+)
      failed-rebalance-totalThe total number of failed group rebalanceskafka.consumer:type=consumer-coordinator-metrics,client-id=([-.\w]+)
      failed-rebalance-rate-per-hourThe number of failed group rebalance event per hourkafka.consumer:type=consumer-coordinator-metrics,client-id=([-.\w]+)
      last-rebalance-seconds-agoThe number of seconds since the last rebalance eventkafka.consumer:type=consumer-coordinator-metrics,client-id=([-.\w]+)
      
       
      
         last-heartbeat-seconds-ago
         The number of seconds since the last controller heartbeat
         kafka.consumer:type=consumer-coordinator-metrics,client-id=([-.\w]+)
       
      partitions-revoked-latency-avgThe average time taken by the on-partitions-revoked rebalance listener callbackkafka.consumer:type=consumer-coordinator-metrics,client-id=([-.\w]+)
      partitions-revoked-latency-maxThe max time taken by the on-partitions-revoked rebalance listener callbackkafka.consumer:type=consumer-coordinator-metrics,client-id=([-.\w]+)
      partitions-assigned-latency-avgThe average time taken by the on-partitions-assigned rebalance listener callbackkafka.consumer:type=consumer-coordinator-metrics,client-id=([-.\w]+)
      partitions-assigned-latency-maxThe max time taken by the on-partitions-assigned rebalance listener callbackkafka.consumer:type=consumer-coordinator-metrics,client-id=([-.\w]+)
      partitions-lost-latency-avgThe average time taken by the on-partitions-lost rebalance listener callbackkafka.consumer:type=consumer-coordinator-metrics,client-id=([-.\w]+)
      partitions-lost-latency-maxThe max time taken by the on-partitions-lost rebalance listener callbackkafka.consumer:type=consumer-coordinator-metrics,client-id=([-.\w]+)
      
           
   
      
 
@@ -1430,7 +1605,8 @@ 
      Streams Mo
   the info level records only the thread-level metrics.
 
   
      
-    Note that the metrics have a 3-layer hierarchy. At the top level there are per-thread metrics. Each thread has tasks, with their
+    Note that the metrics have a 4-layer hierarchy. At the top level there are client-level metrics for each started
+    Kafka Streams client. Each client has stream threads, with their own metrics. Each stream thread has tasks, with their
     own metrics. Each task has a number of processor nodes, with their own metrics. Each task also has a number of state stores
     and record caches, all with their own metrics.
   
      
@@ -1440,6 +1616,43 @@ 
      Streams Mo
 
 
      metrics.recording.level="info"
      
 
+
      Client Metrics
      
+All the following metrics have a recording level of info:
+
+  
+  
+    
+    
+    
+  
+  
+    
+    
+    
+  
+  
+    
+    
+    
+  
+  
+    
+    
+    
+  
+  
+    
+    
+    
+  
+  
+    
+    
+    
+  
+  
+
      Metric/Attribute nameDescriptionMbean name
      versionThe version of the Kafka Streams client.kafka.streams:type=stream-metrics,client-id=([-.\w]+)
      commit-idThe version control commit ID of the Kafka Streams client.kafka.streams:type=stream-metrics,client-id=([-.\w]+)
      application-idThe application ID of the Kafka Streams client.kafka.streams:type=stream-metrics,client-id=([-.\w]+)
      topology-descriptionThe description of the topology executed in the Kafka Streams client.kafka.streams:type=stream-metrics,client-id=([-.\w]+)
      stateThe state of the Kafka Streams client.kafka.streams:type=stream-metrics,client-id=([-.\w]+)
      
+
 
      Thread Metrics
      
 All the following metrics have a recording level of info:
 
@@ -1924,6 +2137,108 @@ 
      RocksDB Metrics
      
+  All the following metrics have a recording level of debug.
+  The metrics are collected every minute from the RocksDB state stores.
+  If a state store consists of multiple RocksDB instances as it is the case for aggregations over time and session windows,
+  each metric reports an aggregation over the RocksDB instances of the state store.
+  Note that the store-scope for built-in RocksDB state stores are currently the following:
+  
        
+    
      • rocksdb-state (for RocksDB backed key-value store)
        • 
+    
      • rocksdb-window-state (for RocksDB backed window store)
        • 
+    
      • rocksdb-session-state (for RocksDB backed session store)
        • 
+  
      
+
+  
      
+    
+    
+      
+      
+      
+    
+    
+      
+      
+      
+    
+    
+      
+      
+      
+    
+    
+      
+      
+      
+    
+    
+      
+      
+      
+    
+    
+      
+      
+      
+    
+    
+      
+      
+      
+    
+    
+      
+      
+      
+    
+    
+      
+      
+      
+    
+    
+      
+      
+      
+    
+    
+      
+      
+      
+    
+    
+      
+      
+      
+    
+    
+      
+      
+      
+    
+    
+      
+      
+      
+    
+    
+      
+      
+      
+    
+    
+      
+      
+      
+    
+    
+      
+      
+      
+    
+    
+  
      Metric/Attribute nameDescriptionMbean name
      bytes-written-rateThe average number of bytes written per second to the RocksDB state store.kafka.streams:type=stream-state-metrics,client-id=([-.\w]+),task-id=([-.\w]+),[store-scope]-id=([-.\w]+)
      bytes-written-totalThe total number of bytes written to the RocksDB state store.kafka.streams:type=stream-state-metrics,client-id=([-.\w]+),task-id=([-.\w]+),[store-scope]-id=([-.\w]+)
      bytes-read-rateThe average number of bytes read per second from the RocksDB state store.kafka.streams:type=stream-state-metrics,client-id=([-.\w]+),task-id=([-.\w]+),[store-scope]-id=([-.\w]+)
      bytes-read-totalThe total number of bytes read from the RocksDB state store.kafka.streams:type=stream-state-metrics,client-id=([-.\w]+),task-id=([-.\w]+),[store-scope]-id=([-.\w]+)
      memtable-bytes-flushed-rateThe average number of bytes flushed per second from the memtable to disk.kafka.streams:type=stream-state-metrics,client-id=([-.\w]+),task-id=([-.\w]+),[store-scope]-id=([-.\w]+)
      memtable-bytes-flushed-totalThe total number of bytes flushed from the memtable to disk.kafka.streams:type=stream-state-metrics,client-id=([-.\w]+),task-id=([-.\w]+),[store-scope]-id=([-.\w]+)
      memtable-hit-ratioThe ratio of memtable hits relative to all lookups to the memtable.kafka.streams:type=stream-state-metrics,client-id=([-.\w]+),task-id=([-.\w]+),[store-scope]-id=([-.\w]+)
      block-cache-data-hit-ratioThe ratio of block cache hits for data blocks relative to all lookups for data blocks to the block cache.kafka.streams:type=stream-state-metrics,client-id=([-.\w]+),task-id=([-.\w]+),[store-scope]-id=([-.\w]+)
      block-cache-index-hit-ratioThe ratio of block cache hits for index blocks relative to all lookups for index blocks to the block cache.kafka.streams:type=stream-state-metrics,client-id=([-.\w]+),task-id=([-.\w]+),[store-scope]-id=([-.\w]+)
      block-cache-filter-hit-ratioThe ratio of block cache hits for filter blocks relative to all lookups for filter blocks to the block cache.kafka.streams:type=stream-state-metrics,client-id=([-.\w]+),task-id=([-.\w]+),[store-scope]-id=([-.\w]+)
      write-stall-duration-avgThe average duration of write stalls in ms.kafka.streams:type=stream-state-metrics,client-id=([-.\w]+),task-id=([-.\w]+),[store-scope]-id=([-.\w]+)
      write-stall-duration-totalThe total duration of write stalls in ms.kafka.streams:type=stream-state-metrics,client-id=([-.\w]+),task-id=([-.\w]+),[store-scope]-id=([-.\w]+)
      bytes-read-compaction-rateThe average number of bytes read per second during compaction.kafka.streams:type=stream-state-metrics,client-id=([-.\w]+),task-id=([-.\w]+),[store-scope]-id=([-.\w]+)
      bytes-written-compaction-rateThe average number of bytes written per second during compaction.kafka.streams:type=stream-state-metrics,client-id=([-.\w]+),task-id=([-.\w]+),[store-scope]-id=([-.\w]+)
      number-open-filesThe number of current open files.kafka.streams:type=stream-state-metrics,client-id=([-.\w]+),task-id=([-.\w]+),[store-scope]-id=([-.\w]+)
      number-file-errors-totalThe total number of file errors occurred.kafka.streams:type=stream-state-metrics,client-id=([-.\w]+),task-id=([-.\w]+),[store-scope]-id=([-.\w]+)
      
+
   
      Record Cache Metrics
      
   All the following metrics have a recording level of debug:
 
@@ -1999,7 +2314,7 @@ 
      Others
      
   
      6.7 ZooKeeper
      
 
   
      Stable version
      
-  The current stable branch is 3.4 and the latest release of that branch is 3.4.9.
+  The current stable branch is 3.5. Kafka is regularly updated to include the latest release in the 3.5 series.
 
   
      Operationalizing ZooKeeper
      
   Operationally, we do the following for a healthy ZooKeeper installation:
diff --git a/docs/protocol.html b/docs/protocol.html
index b884992a7bcdd..6d9bb5bef9b29 100644
--- a/docs/protocol.html
+++ b/docs/protocol.html
@@ -32,6 +32,8 @@ 
      Kafka protocol guide
      
             
    5. Partitioning Strategies
             
    6. Batching
             
    7. Versioning and Compatibility
+            
    8. Retrieving Supported API versions
+            
    9. SASL Authentication Sequence
         
     
      10. 
     
    10. The Protocol
@@ -42,6 +44,11 @@ 
      Kafka protocol guide
      
             
    11. Record Batch
         
     
      12. 
+    
    12. Evolving the Protocol
+        
     
    13. Constants
         
@@ -63,7 +63,7 @@ 
      Stream Processing Topology
     
 
     There are two special processors in the topology:
@@ -124,7 +124,7 @@ 
      Time
      
     
        
         
      •  When new output records are generated via processing some input record, for example, context.forward() triggered in the process() function call, output record timestamps are inherited from input record timestamps directly.
        • 
         
      •  When new output records are generated via periodic functions such as Punctuator#punctuate(), the output record timestamp is defined as the current internal time (obtained through context.timestamp()) of the stream task.
        • 
-        
      •  For aggregations, the timestamp of a resulting aggregate update record will be that of the latest arrived input record that triggered the update.
        • 
+        
      •  For aggregations, the timestamp of a result update record will be the maximum timestamp of all input records contributing to the result.
        • 
     
      
 
     
      
@@ -137,7 +137,7 @@ 
      
     
      
 
     
      
-        In the Kafka Streams DSL, an input stream of an aggregation can be a KStream or a KTable, but the output stream will always be a KTable. This allows Kafka Streams to update an aggregate value upon the late arrival of further records after the value was produced and emitted. When such late arrival happens, the aggregating KStream or KTable emits a new aggregate value. Because the output is a KTable, the new value is considered to overwrite the old value with the same key in subsequent processing steps.
+        In the Kafka Streams DSL, an input stream of an aggregation can be a KStream or a KTable, but the output stream will always be a KTable. This allows Kafka Streams to update an aggregate value upon the out-of-order arrival of further records after the value was produced and emitted. When such out-of-order arrival happens, the aggregating KStream or KTable emits a new aggregate value. Because the output is a KTable, the new value is considered to overwrite the old value with the same key in subsequent processing steps.
     
      
 
     
       Windowing
      
@@ -145,10 +145,10 @@ 
       Windo
         Windowing lets you control how to group records that have the same key for stateful operations such as aggregations or joins into so-called windows. Windows are tracked per record key.
     
      
     
      
-        Windowing operations are available in the Kafka Streams DSL. When working with windows, you can specify a retention period for the window. This retention period controls how long Kafka Streams will wait for out-of-order or late-arriving data records for a given window. If a record arrives after the retention period of a window has passed, the record is discarded and will not be processed in that window.
+        Windowing operations are available in the Kafka Streams DSL. When working with windows, you can specify a grace period for the window. This grace period controls how long Kafka Streams will wait for out-of-order data records for a given window. If a record arrives after the grace period of a window has passed, the record is discarded and will not be processed in that window. Specifically, a record is discarded if its timestamp dictates it belongs to a window, but the current stream time is greater than the end of the window plus the grace period.
     
      
     
      
-        Late-arriving records are always possible in the real world and should be properly accounted for in your applications. It depends on the effective time semantics  how late records are handled. In the case of processing-time, the semantics are "when the record is being processed", which means that the notion of late records is not applicable as, by definition, no record can be late. Hence, late-arriving records can only be considered as such (i.e. as arriving "late") for event-time or ingestion-time semantics. In both cases, Kafka Streams is able to properly handle late-arriving records.
+        Out-of-order records are always possible in the real world and should be properly accounted for in your applications. It depends on the effective time semantics  how out-of-order records are handled. In the case of processing-time, the semantics are "when the record is being processed", which means that the notion of out-of-order records is not applicable as, by definition, no record can be out-of-order. Hence, out-of-order records can only be considered as such for event-time. In both cases, Kafka Streams is able to properly handle out-of-order records.
     
      
 
     
      Duality of Streams and Tables
      
@@ -159,25 +159,24 @@ 
      Duality of
     
      
 
     
      
-        Any stream processing technology must therefore provide first-class support for streams and tables.
-        Kafka's Streams API provides such functionality through its core abstractions for
-        streams <streams_concepts_kstream> and
-        tables <streams_concepts_ktable>, which we will talk about in a minute.
-        Now, an interesting observation is that there is actually a close relationship between streams and tables,
-        the so-called stream-table duality.
-        And Kafka exploits this duality in many ways: for example, to make your applications
-        elastic <streams_developer-guide_execution-scaling>,
-        to support fault-tolerant stateful processing <streams_developer-guide_state-store_fault-tolerance>,
-        or to run interactive queries <streams_concepts_interactive-queries>
-        against your application's latest processing results. And, beyond its internal usage, the Kafka Streams API
-        also allows developers to exploit this duality in their own applications.
-    
      
-
-    
      
-        Before we discuss concepts such as aggregations <streams_concepts_aggregations>
-        in Kafka Streams we must first introduce tables in more detail, and talk about the aforementioned stream-table duality.
-        Essentially, this duality means that a stream can be viewed as a table, and a table can be viewed as a stream.
-    
      
+      Any stream processing technology must therefore provide first-class support for streams and tables.
+      Kafka's Streams API provides such functionality through its core abstractions for 
+            streams
+      and tables,
+      which we will talk about in a minute. Now, an interesting observation is that there is actually a close relationship between streams and tables,
+      the so-called stream-table duality. And Kafka exploits this duality in many ways: for example, to make your applications
+      elastic,
+      to support fault-tolerant stateful processing,
+      or to run interactive queries
+      against your application's latest processing results. And, beyond its internal usage, the Kafka Streams API
+      also allows developers to exploit this duality in their own applications.
+  
      
+
+  
      
+      Before we discuss concepts such as aggregations
+      in Kafka Streams, we must first introduce tables in more detail, and talk about the aforementioned stream-table duality.
+      Essentially, this duality means that a stream can be viewed as a table, and a table can be viewed as a stream.
+  
      
 
     
      States
      
 
@@ -185,7 +184,7 @@ 
      States
      
         Some stream processing applications don't require state, which means the processing of a message is independent from
         the processing of all other messages.
         However, being able to maintain state opens up many possibilities for sophisticated stream processing applications: you
-        can join input streams, or group and aggregate data records. Many such stateful operators are provided by the Kafka Streams DSL.
+        can join input streams, or group and aggregate data records. Many such stateful operators are provided by the Kafka Streams DSL.
     
      
     
      
         Kafka Streams provides so-called state stores, which can be used by stream processing applications to store and query data.
@@ -224,7 +223,7 @@ 
      Out-of-Order
 
     
      
         Besides the guarantee that each record will be processed exactly-once, another issue that many stream processing application will face is how to
-        handle out-of-order data that may impact their business logic. In Kafka Streams, there are two causes that could potentially
+        handle out-of-order data that may impact their business logic. In Kafka Streams, there are two causes that could potentially
         result in out-of-order data arrivals with respect to their timestamps:
     
      
 
diff --git a/docs/streams/developer-guide/config-streams.html b/docs/streams/developer-guide/config-streams.html
index 4b893ee7d1209..5084bc2aea597 100644
--- a/docs/streams/developer-guide/config-streams.html
+++ b/docs/streams/developer-guide/config-streams.html
@@ -284,12 +284,17 @@ 
      bootstrap.serversTimestamp extractor class that implements the TimestampExtractor interface.
             See Timestamp Extractor
           
-          value.serde
+          upgrade.from
+            Medium
+            The version you are upgrading from during a rolling upgrade.
+            See Upgrade From
+          
+          value.serde
             Medium
             Default serializer/deserializer class for record values, implements the Serde interface (see also key.serde).
             Serdes.ByteArray().getClass().getName()
           
-          windowstore.changelog.additional.retention.ms
+          windowstore.changelog.additional.retention.ms
             Low
             Added to a windows maintainMs to ensure data is not deleted from the log prematurely. Allows for clock drift.
             86400000 milliseconds = 1 day
@@ -474,7 +479,9 @@ 
      state.dir
             
      The state directory. Kafka Streams persists local states under the state directory. Each application has a subdirectory on its hosting
               machine that is located under the state directory. The name of the subdirectory is the application ID. The state stores associated
-              with the application are created under this subdirectory.
      
+              with the application are created under this subdirectory. When running multiple instances of the same application on a single machine,
+              this path must be unique for each such instance.
      14. 
+          
         
         
    
               
             
         
+        
    
+          
    upgrade.from
    
+          
    
+            
    
+              The version you are upgrading from. It is important to set this config when performing a rolling upgrade to certain versions, as described in the upgrade guide.
+              You should set this config to the appropriate version before bouncing your instances and upgrading them to the newer version. Once everyone is on the
+              newer version, you should remove this config and do a second rolling bounce. It is only necessary to set this config and follow the two-bounce upgrade path
+              when upgrading from below version 2.0, or when upgrading to 2.4+ from any version lower than 2.4.
+            
    
+          
    
+        
    
       
           
         
         
    
@@ -658,11 +684,7 @@ 
    Default ValuesConsumer
               1000
             
-            retries
-              Producer
-              10
-            
-            rocksdb.config.setter
+            rocksdb.config.setter
               Consumer
                
             
@@ -684,11 +706,14 @@ 
    Default ValuesHere is an example that adjusts the memory size consumed by RocksDB.
    
               
        public static class CustomRocksDBConfig implements RocksDBConfigSetter {
 
+       // This object should be a member variable so it can be closed in RocksDBConfigSetter#close.
+       private org.rocksdb.Cache cache = new org.rocksdb.LRUCache(16 * 1024L * 1024L);
+
        @Override
        public void setConfig(final String storeName, final Options options, final Map<String, Object> configs) {
          // See #1 below.
-         BlockBasedTableConfig tableConfig = new org.rocksdb.BlockBasedTableConfig();
-         tableConfig.setBlockCacheSize(16 * 1024 * 1024L);
+         BlockBasedTableConfig tableConfig = (BlockBasedTableConfig) options.tableFormatConfig();
+         tableConfig.setBlockCache(cache);
          // See #2 below.
          tableConfig.setBlockSize(16 * 1024L);
          // See #3 below.
@@ -697,6 +722,12 @@ 
    Default Values// See #4 below.
          options.setMaxWriteBufferNumber(2);
        }
+
+       @Override
+       public void close(final String storeName, final Options options) {
+         // See #5 below.
+         cache.close();
+       }
     }
 
 Properties streamsSettings = new Properties();
@@ -706,10 +737,11 @@ 
    Default Values
                 
    Notes for example:
    
                 
      
-                  
    1. BlockBasedTableConfig tableConfig = new org.rocksdb.BlockBasedTableConfig(); Reduce block cache size from the default, shown here,  as the total number of store RocksDB databases is partitions (40) * segments (3) = 120.
      2. 
-                  
    2. tableConfig.setBlockSize(16 * 1024L); Modify the default block size per these instructions from the RocksDB GitHub.
      3. 
+                  
    3. BlockBasedTableConfig tableConfig = (BlockBasedTableConfig) options.tableFormatConfig(); Get a reference to the existing table config rather than create a new one, so you don't accidentally overwrite defaults such as the BloomFilter, which is an important optimization.
+                  
    4. tableConfig.setBlockSize(16 * 1024L); Modify the default block size per these instructions from the RocksDB GitHub.
      5. 
                   
    5. tableConfig.setCacheIndexAndFilterBlocks(true); Do not let the index and filter blocks grow unbounded. For more information, see the RocksDB GitHub.
      6. 
                   
    6. options.setMaxWriteBufferNumber(2); See the advanced options in the RocksDB GitHub.
      7. 
+                  
    7. cache.close(); To avoid memory leaks, you must close any objects you constructed that extend org.rocksdb.RocksObject. See  RocksJava docs for more details.
      8. 
                 
    
                 
    
               
@@ -767,6 +799,7 @@ 
    replication.factordescription here.
    
           
    Properties streamsSettings = new Properties();
 streamsSettings.put(StreamsConfig.REPLICATION_FACTOR_CONFIG, 3);
+streamsSettings.put(StreamsConfig.topicPrefix(TopicConfig.MIN_IN_SYNC_REPLICAS_CONFIG), 2);
 streamsSettings.put(StreamsConfig.producerPrefix(ProducerConfig.ACKS_CONFIG), "all");
 
    
           
    
diff --git a/docs/streams/developer-guide/datatypes.html b/docs/streams/developer-guide/datatypes.html
index 6c7869c3ba143..ecdb48c732439 100644
--- a/docs/streams/developer-guide/datatypes.html
+++ b/docs/streams/developer-guide/datatypes.html
@@ -47,9 +47,8 @@
           
  • Available SerDes
    • 
           
           
  • Kafka Streams DSL for Scala Implicit SerDes
    • 
       
@@ -61,9 +60,9 @@ 
    Configuring SerDesProperties settings = new Properties();
 // Default serde for keys of data records (here: built-in serde for String type)
-settings.put(StreamsConfig.KEY_SERDE_CLASS_CONFIG, Serdes.String().getClass().getName());
+settings.put(StreamsConfig.DEFAULT_KEY_SERDE_CLASS_CONFIG, Serdes.String().getClass().getName());
 // Default serde for values of data records (here: built-in serde for Long type)
-settings.put(StreamsConfig.VALUE_SERDE_CLASS_CONFIG, Serdes.Long().getClass().getName());
+settings.put(StreamsConfig.DEFAULT_VALUE_SERDE_CLASS_CONFIG, Serdes.Long().getClass().getName());
    @@ -112,7 +111,7 @@

    Primitive and basic types</dependency> -

    This artifact provides the following serde implementations under the package org.apache.kafka.common.serialization, which you can leverage when e.g., defining default serializers in your Streams configuration.

    +

    This artifact provides the following serde implementations under the package org.apache.kafka.common.serialization, which you can leverage when e.g., defining default serializers in your Streams configuration.

    @@ -149,20 +148,17 @@

    Primitive and basic types

    Tip

    -

    Bytes is a wrapper for Java’s byte[] (byte array) that supports proper equality and ordering semantics. You may want to consider using Bytes instead of byte[] in your applications.

    +

    Bytes is a wrapper for Java’s byte[] (byte array) that supports proper equality and ordering semantics. You may want to consider using Bytes instead of byte[] in your applications.

    JSON

    -

    The code examples of Kafka Streams also include a basic serde implementation for JSON:

    +

    The Kafka Streams code examples also include a basic serde implementation for JSON:

    -

    You can construct a unified JSON serde from the JsonPOJOSerializer and JsonPOJODeserializer via - Serdes.serdeFrom(<serializerInstance>, <deserializerInstance>). The - PageViewTypedDemo - example demonstrates how to use this JSON serde.

    +

    As shown in the example, you can use JSONSerdes inner classes Serdes.serdeFrom(<serializerInstance>, <deserializerInstance>) to construct JSON compatible serializers and deserializers. +

    Implementing custom SerDes

    @@ -170,13 +166,13 @@

    JSON
    1. Write a serializer for your data type T by implementing - org.apache.kafka.common.serialization.Serializer.
      2. + org.apache.kafka.common.serialization.Serializer.
    2. Write a deserializer for T by implementing - org.apache.kafka.common.serialization.Deserializer.
      3. + org.apache.kafka.common.serialization.Deserializer.
    3. Write a serde for T by implementing - org.apache.kafka.common.serialization.Serde, + org.apache.kafka.common.serialization.Serde, which you either do manually (see existing SerDes in the previous section) or by leveraging helper functions in - Serdes + Serdes such as Serdes.serdeFrom(Serializer<T>, Deserializer<T>). Note that you will need to implement your own class (that has no generic types) if you want to use your custom serde in the configuration provided to KafkaStreams. If your serde class has generic types or you use Serdes.serdeFrom(Serializer<T>, Deserializer<T>), you can pass your serde only diff --git a/docs/streams/developer-guide/dsl-api.html b/docs/streams/developer-guide/dsl-api.html index 0ddcbc5bd193d..02f2441a57db7 100644 --- a/docs/streams/developer-guide/dsl-api.html +++ b/docs/streams/developer-guide/dsl-api.html @@ -48,7 +48,8 @@
    4. Joining @@ -66,6 +67,7 @@
    5. Applying processors and transformers (Processor API integration)
      6. +
    6. Naming Operators in a Streams DSL application
    7. Controlling KTable update rate
    8. Writing streams back to Kafka
    9. Testing a Streams application
      10. @@ -163,7 +165,7 @@

      KTable

      - We have already seen an example of a changelog stream in the section streams_concepts_duality. Another example are change data capture (CDC) records in the changelog of a relational database, representing which row in a database table was inserted, updated, or deleted. + We have already seen an example of a changelog stream in the section streams and tables. Another example are change data capture (CDC) records in the changelog of a relational database, representing which row in a database table was inserted, updated, or deleted.

      @@ -259,7 +261,7 @@

      You must specify SerDes explicitly if the key or value types of the records in the Kafka input topics do not match the configured default SerDes. For information about configuring default SerDes, available SerDes, and implementing your own custom SerDes see Data Types and Serialization.

      -

      Several variants of stream exist, for example to specify a regex pattern for input topics to read from).

      +

      Several variants of stream exist. For example, you can specify a regex pattern for input topics to read from (note that all matching topics will be part of the same input topic group, and the work will not be parallelized for different topics if subscribed to in this way).

    - + - + - + + + + + + + @@ -1651,13 +1659,11 @@

    Each case is explained in more detail in the subsequent sections.

    Join co-partitioning requirements
    -

    Input data must be co-partitioned when joining. This ensures that input records with the same key, from both sides of the +

    For equi-joins, input data must be co-partitioned when joining. This ensures that input records with the same key from both sides of the join, are delivered to the same stream task during processing. - It is the responsibility of the user to ensure data co-partitioning when joining.

    -
    -

    Tip

    -

    If possible, consider using global tables (GlobalKTable) for joining because they do not require data co-partitioning.

    -
    + It is your responsibility to ensure data co-partitioning when joining. + Co-partitioning is not required when performing KTable-KTable Foreign-Key joins and Global KTable joins. +

    The requirements for data co-partitioning are:

    • The input topics of the join (left side and right side) must have the same number of partitions.
      • @@ -1676,11 +1682,14 @@

      KStream-KTable joins are performed based on the keys of records (e.g., leftRecord.key == rightRecord.key), it is required that the input streams/tables of a join are co-partitioned by key.

      -

      The only exception are - KStream-GlobalKTable joins. Here, co-partitioning is - it not required because all partitions of the GlobalKTable‘s underlying changelog stream are made available to - each KafkaStreams instance, i.e. each instance has a full copy of the changelog stream. Further, a - KeyValueMapper allows for non-key based joins from the KStream to the GlobalKTable.

      +

      There are two exceptions where co-partitioning is not required. For + KStream-GlobalKTable joins joins, co-partitioning is + not required because all partitions of the GlobalKTable‘s underlying changelog stream are made available to + each KafkaStreams instance. That is, each instance has a full copy of the changelog stream. Further, a + KeyValueMapper allows for non-key based joins from the KStream to the GlobalKTable. + KTable-KTable Foreign-Key joins also do not require co-partitioning. Kafka Streams internally ensures co-partitioning for Foreign-Key joins. +

      +

      Note

      Kafka Streams partly verifies the co-partitioning requirement: @@ -2064,9 +2073,9 @@

    Table

    @@ -335,7 +337,7 @@

    Some KStream transformations may generate one or more KStream objects, for example: - filter and map on a KStream will generate another KStream - branch on KStream can generate multiple KStreams

    -

    Some others may generate a KTable object, for example an aggregation of a KStream also yields a KTable. This allows Kafka Streams to continuously update the computed value upon arrivals of late records after it +

    Some others may generate a KTable object, for example an aggregation of a KStream also yields a KTable. This allows Kafka Streams to continuously update the computed value upon arrivals of out-of-order records after it has already been produced to the downstream transformation operators.

    All KTable transformation operations can only generate another KTable. However, the Kafka Streams DSL does provide a special function that converts a KTable representation into a KStream. All of these transformation methods can be chained together to compose @@ -1628,19 +1630,25 @@

    		Supported Supported
    KStream-to-KTable
    KTable-to-KTable Foreign-Key Join Non-windowed Supported Supported Not Supported
    KStream-to-GlobalKTable
    KStream-to-KTable Non-windowed Supported Supported Not Supported
    KTable-to-GlobalKTable
    KStream-to-GlobalKTableNon-windowedSupportedSupportedNot Supported
    KTable-to-GlobalKTable N/A Not Supported Not Supported
    -
    -
    KTable-KTable Join
    -

    KTable-KTable joins are always non-windowed joins. They are designed to be consistent with their counterparts in +

    +
    KTable-KTable Equi-Join
    +

    KTable-KTable equi-joins are always non-windowed joins. They are designed to be consistent with their counterparts in relational databases. The changelog streams of both KTables are materialized into local state stores to represent the latest snapshot of their table duals. The join result is a new KTable that represents the changelog stream of the join operation.

    @@ -2240,8 +2249,8 @@

    -

    Semantics of table-table joins: - The semantics of the various table-table join variants are explained below. +

    Semantics of table-table equi-joins: + The semantics of the various table-table equi-join variants are explained below. To improve the readability of the table, you can assume that (1) all records have the same key (and thus the key in the table is omitted) and that (2) all records are processed in timestamp order. The columns INNER JOIN, LEFT JOIN, and OUTER JOIN denote what is passed as arguments to the user-supplied ValueJoiner for the join, leftJoin, and @@ -2374,6 +2383,309 @@

    +
    +
    KTable-KTable Foreign-Key + Join
    +

    KTable-KTable foreign-key joins are always non-windowed + joins. Foreign-key joins are analogous to joins in SQL. As a rough example: +

    + SELECT ... FROM {this KTable} + JOIN {other KTable} + ON {other.key} = {result of foreignKeyExtractor(this.value)} ... +
    + The output of the operation is a new KTable containing the join result. +

    +

    The changelog streams of + both KTables are materialized into local state stores to + represent the latest snapshot of their table duals. A foreign-key extractor + function is applied to the left record, with a new intermediate + record created and is used to lookup and join with the corresponding + primary key on the right hand side table. + The result is a new KTable that represents the changelog stream + of the join operation.

    + The left KTable can have multiple records which map to the same + key on the right KTable. An update to a single left KTable entry + may result in a single output event, provided the corresponding + key exists in the right KTable. Consequently, a single update to a + right KTable entry will result in an update for each record in the + left KTable that has the same foreign key.
    +
    +
    + + + + + + + + + + + + + + + +
    TransformationDescription
    +

    Inner Join

    +
      +
    • (KTable, KTable) → KTable
      • +
    +     		+

    Performs a foreign-key INNER JOIN of this + table with another table. The result is an ever-updating + KTable that represents the “current” result of the join. + (details)

    +
    +
    + 
    KTable<String, Long> left = ...;
+                KTable<Long, Double> right = ...;
    //This     foreignKeyExtractor simply uses the left-value to map to the right-key.
    Function<Long, Long> foreignKeyExtractor = (x) -> x;

    // Java 8+ example, using lambda expressions
+                KTable<String, String> joined = left.join(right,
        foreignKeyExtractor,
+                    (leftValue, rightValue) -> "left=" + leftValue + ", right=" + rightValue /* ValueJoiner */
+                  );
+
    +
    +
    +

    Detailed behavior:

    +
      +
    • +

      The join is key-based, i.e. + with the join predicate:

      + 
      foreignKeyExtractor.apply(leftRecord.value) == rightRecord.key
      +
      • +
    • +

      The join will be triggered under the + conditions listed below whenever new input is + received. When it is triggered, the user-supplied ValueJoiner + will be called to produce join output records.

      +
      +
      +
        +
      • + Records for which the foreignKeyExtractor produces null are ignored and do not trigger a join. + If you want to join with null foreign keys, use a suitable sentinel value to do so (i.e. "NULL" for a String field, or -1 for an auto-incrementing integer field). +
        • +
      • Input records with a null + value are interpreted as tombstones + for the corresponding key, which indicate the + deletion of the key from the table. When an input + tombstone is received, then an output + tombstone is forwarded directly to the join + result KTable if required (i.e. only if the + corresponding key actually exists already in + the join result KTable).
        • +
      +
      +
      +
      • +
    +

    See the semantics overview at the bottom + of this section for a detailed description.

    +
    +

    Left Join

    +
      +
    • (KTable, KTable) → KTable
      • +
    +     		+

    Performs a foreign-key LEFT JOIN of this + table with another table. (details)

    +
    +
    + 
    KTable<String, Long> left = ...;
+                KTable<Long, Double> right = ...;
    //This     foreignKeyExtractor simply uses the left-value to map to the right-key.
    Function<Long, Long> foreignKeyExtractor = (x) -> x;

    // Java 8+ example, using lambda expressions
+                KTable<String, String> joined = left.join(right,
        foreignKeyExtractor,
+                    (leftValue, rightValue) -> "left=" + leftValue + ", right=" + rightValue /* ValueJoiner */
+                  );
+
    +
    +
    +

    Detailed behavior:

    +
      +
    • +

      The join is key-based, i.e. + with the join predicate:

      + 
      foreignKeyExtractor.apply(leftRecord.value) == rightRecord.key
      +
      • +
    • +

      The join will be triggered under the + conditions listed below whenever new input is + received. When it is triggered, the user-supplied ValueJoiner + will be called to produce join output records.

      +
      +
      +
        +
      • + Records for which the foreignKeyExtractor produces null are ignored and do not trigger a join. + If you want to join with null foreign keys, use a suitable sentinel value to do so (i.e. "NULL" for a String field, or -1 for an auto-incrementing integer field). +
        • +
      • Input records with a null + value are interpreted as tombstones + for the corresponding key, which indicate the + deletion of the key from the table. When an input + tombstone is received, then an output + tombstone is forwarded directly to the join + result KTable if required (i.e. only if the + corresponding key actually exists already in + the join result KTable).
        • +
      +
      +
      +
      • +
    • +

      For each input record on the left + side that does not have any match on the right side, + the ValueJoiner will be + called with ValueJoiner#apply(leftRecord.value, + null); this + explains the row with timestamp=7 & 8 in the + table below, which lists (q,10,null) and (r,10,null) + in the LEFT JOIN column.

      +
      • +
    +

    See the semantics overview at the bottom + of this section for a detailed description.

    +
    +

    Semantics of table-table foreign-key joins: + The semantics of the table-table foreign-key INNER and LEFT JOIN + variants are demonstrated below. + The key is shown alongside the value for each record. + Records are processed in incrementing offset order. + The columns INNER JOIN and LEFT JOIN denote what is + passed as arguments to the user-supplied ValueJoiner + for the join + and leftJoin + methods, respectively, whenever a new input record is received + on either side of the join. An empty table cell denotes that the + ValueJoiner + is not called at all. For the purpose of this example, Function + foreignKeyExtractor simply uses the left-value + as the output.

    + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
    Record OffsetLeft KTable (K, extracted-FK)Right KTable (FK, VR)(INNER) JOINLEFT JOIN
    1(k,1) (1,foo) + + (k,1,foo)
    +     		+ + (k,1,foo)
    2(k,2) + + + +
    +     		(k,null)(k,2,null)
    +
    3(k,3)
    +     		 (k,null)(k,3,null)
    +
    4 (3,bar)
    +     		(k,3,bar)
    +     		(k,3,bar)
    +
    5(k,null)
    +     		 (k,null)
    +     		(k,null,null) + +
    6(k,1)
    +     		(k,1,foo)
    +     		(k,1,foo)
    +
    7(q,10)
    +     		 
    +     		(q,10,null)
    8(r,10)
    +     		 (r,10,null)
    9
    +     		(10,baz) (q,10,baz), (r,10,baz)
    +
    -

    If we then receive three additional records (including two late-arriving records), what would happen is that the two +

    If we then receive three additional records (including two out-of-order records), what would happen is that the two existing sessions for the green record key will be merged into a single session starting at time 0 and ending at time 6, consisting of a total of three records. The existing session for the blue record key will be extended to end at time 5, consisting of a total of two records. And, finally, there will be a new session for the blue key starting and ending at time 11.

    -

    Detected sessions after having received six input records. Note the two late-arriving data records at t=4 (green) and +

    Detected sessions after having received six input records. Note the two out-of-order data records at t=4 (green) and t=5 (blue), which lead to a merge of sessions and an extension of a session, respectively.

    @@ -3007,7 +3319,7 @@

    grace(ofMinutes(10))
    This allows us to bound the lateness of events the window will accept. - For example, the 09:00 to 10:00 window will accept late-arriving records until 10:10, at which point, the window is closed. + For example, the 09:00 to 10:00 window will accept out-of-order records until 10:10, at which point, the window is closed.
    .suppress(Suppressed.untilWindowCloses(...))
    This configures the suppression operator to emit nothing for a window until it closes, and then emit the final result. @@ -3204,7 +3516,10 @@

    - +
    +

    Naming Operators in a Streams DSL application

    + Kafka Streams allows you to name processors created via the Streams DSL +

    Controlling KTable emit rate

    A KTable is logically a continuously updated table. diff --git a/docs/streams/developer-guide/dsl-topology-naming.html b/docs/streams/developer-guide/dsl-topology-naming.html new file mode 100644 index 0000000000000..e0c1e1fc67daf --- /dev/null +++ b/docs/streams/developer-guide/dsl-topology-naming.html @@ -0,0 +1,370 @@ + + + + + + + + +

    + +
    + + +
    +
    +
    + + + + + + diff --git a/docs/streams/developer-guide/index.html b/docs/streams/developer-guide/index.html index 095ec829aa14a..eb96f7de0f583 100644 --- a/docs/streams/developer-guide/index.html +++ b/docs/streams/developer-guide/index.html @@ -43,6 +43,7 @@

    Developer Guide for Kafka Streams

  • Configuring a Streams Application
  • Streams DSL
  • Processor API
    • +
  • Naming Operators in a Streams DSL application
  • Data Types and Serialization
  • Testing a Streams Application
  • Interactive Queries
    • diff --git a/docs/streams/developer-guide/interactive-queries.html b/docs/streams/developer-guide/interactive-queries.html index 8d8143f04f2ce..e590a521731bc 100644 --- a/docs/streams/developer-guide/interactive-queries.html +++ b/docs/streams/developer-guide/interactive-queries.html @@ -59,7 +59,7 @@

    There are local and remote components to interactively querying the state of your application.

    Local state
    -
    An application instance can query the locally managed portion of the state and directly query its own local state stores. You can use the corresponding local data in other parts of your application code, as long as it doesn’t required calling the Kafka Streams API. Querying state stores is always read-only to guarantee that the underlying state stores will never be mutated out-of-band (e.g., you cannot add new entries). State stores should only be mutated by the corresponding processor topology and the input data it operates on. For more information, see Querying local state stores for an app instance.
    +
    An application instance can query the locally managed portion of the state and directly query its own local state stores. You can use the corresponding local data in other parts of your application code, as long as it doesn’t require calling the Kafka Streams API. Querying state stores is always read-only to guarantee that the underlying state stores will never be mutated out-of-band (e.g., you cannot add new entries). State stores should only be mutated by the corresponding processor topology and the input data it operates on. For more information, see Querying local state stores for an app instance.
    Remote state

    To query the full state of your application, you must connect the various fragments of the state, including:

    @@ -167,8 +168,59 @@
    -
    -

    Other memory usage

    +
    +

    RocksDB

    +

    Each instance of RocksDB allocates off-heap memory for a block cache, index and filter blocks, and memtable (write buffer). Critical configs (for RocksDB version 4.1.0) include + block_cache_size, write_buffer_size and max_write_buffer_number. These can be specified through the + rocksdb.config.setter configuration. +

    As of 2.3.0 the memory usage across all instances can be bounded, limiting the total off-heap memory of your Kafka Streams application. To do so you must configure RocksDB to cache the index and filter blocks in the block cache, limit the memtable memory through a shared WriteBufferManager and count its memory against the block cache, and then pass the same Cache object to each instance. See RocksDB Memory Usage for details. An example RocksDBConfigSetter implementing this is shown below:

    + + 
        public static class BoundedMemoryRocksDBConfig implements RocksDBConfigSetter {
+
+       private static org.rocksdb.Cache cache = new org.rocksdb.LRUCache(TOTAL_OFF_HEAP_MEMORY, -1, false, INDEX_FILTER_BLOCK_RATIO);1
+       private static org.rocksdb.WriteBufferManager writeBufferManager = new org.rocksdb.WriteBufferManager(TOTAL_MEMTABLE_MEMORY, cache);
+
+       @Override
+       public void setConfig(final String storeName, final Options options, final Map<String, Object> configs) {
+
+         BlockBasedTableConfig tableConfig = (BlockBasedTableConfig) options.tableFormatConfig();
+
+          // These three options in combination will limit the memory used by RocksDB to the size passed to the block cache (TOTAL_OFF_HEAP_MEMORY)
+         tableConfig.setBlockCache(cache);
+         tableConfig.setCacheIndexAndFilterBlocks(true);
+         options.setWriteBufferManager(writeBufferManager);
+
+          // These options are recommended to be set when bounding the total memory
+         tableConfig.setCacheIndexAndFilterBlocksWithHighPriority(true);2
+         tableConfig.setPinTopLevelIndexAndFilter(true);
+         tableConfig.setBlockSize(BLOCK_SIZE);3
+         options.setMaxWriteBufferNumber(N_MEMTABLES);
+         options.setWriteBufferSize(MEMTABLE_SIZE);
+
+         options.setTableFormatConfig(tableConfig);
+       }
+
+       @Override
+       public void close(final String storeName, final Options options) {
+         // Cache and WriteBufferManager should not be closed here, as the same objects are shared by every store instance.
+       }
+    }
+
    + 1. INDEX_FILTER_BLOCK_RATIO can be used to set a fraction of the block cache to set aside for "high priority" (aka index and filter) blocks, preventing them from being evicted by data blocks. See the full signature of the LRUCache constructor. +
    + 2. This must be set in order for INDEX_FILTER_BLOCK_RATIO to take effect (see footnote 1) as described in the RocksDB docs +
    + 3. You may want to modify the default block size per these instructions from the RocksDB GitHub. A larger block size means index blocks will be smaller, but the cached data blocks may contain more cold data that would otherwise be evicted. +
    +
    +
    Note:
    + While we recommend setting at least the above configs, the specific options that yield the best performance are workload dependent and you should consider experimenting with these to determine the best choices for your specific use case. Keep in mind that the optimal configs for one app may not apply to one with a different topology or input topic. + In addition to the recommended configs above, you may want to consider using partitioned index filters as described by the RocksDB docs. + +
    +
    +
    +

    Other memory usage

    There are other modules inside Apache Kafka that allocate memory during runtime. They include the following:

    Tip

    @@ -237,4 +286,4 @@

    Other memory usageimplement your own custom store type. It’s common practice to leverage an existing store type via the Stores factory.

    Note that, when using Kafka Streams, you normally don’t create or instantiate state stores directly in your code. - Rather, you define state stores indirectly by creating a so-called StoreBuilder. This buildeer is used by + Rather, you define state stores indirectly by creating a so-called StoreBuilder. This builder is used by Kafka Streams as a factory to instantiate the actual state stores locally in application instances when and where needed.

    The following store types are available out of the box.

    @@ -259,7 +259,7 @@

    disk space is either not available or local disk space is wiped in-between app instance restarts.
  • Available store variants: - time window key-value store
    • + time window key-value store, session window key-value store. 
    // Creating an in-memory key-value store:
 // here, we create a `KeyValueStore<String, Long>` named "inmemory-counts".
@@ -289,11 +289,11 @@ 
    
                     you experience machine failure, the state store and the application’s state can be fully restored from its changelog. You can
                     enable or disable this backup feature for a
                     state store.
    
-                
    By default, persistent key-value stores are fault-tolerant.  They are backed by a
+                
    Fault-tolerant state stores are backed by a
                     compacted changelog topic.  The purpose of compacting this
                     topic is to prevent the topic from growing indefinitely, to reduce the storage consumed in the associated Kafka cluster,
                     and to minimize recovery time if a state store needs to be restored from its changelog topic.
    
-                
    Similarly, persistent window stores are fault-tolerant.  They are backed by a topic that uses both compaction and
+                
    Fault-tolerant windowed state stores are backed by a topic that uses both compaction and
                     deletion. Because of the structure of the message keys that are being sent to the changelog topics, this combination of
                     deletion and compaction is required for the changelog topics of window stores. For window stores, the message keys are
                     composite keys that include the “normal” key and window timestamps.  For these types of composite keys it would not
@@ -309,7 +309,7 @@ 
    
                 
    Enable or Disable Fault Tolerance of State Stores (Store Changelogs)
    
                 
    You can enable or disable fault tolerance for a state store by enabling or disabling the change logging
                     of the store through enableLogging() and disableLogging().
-                    You can also fine-tune the associated topic’s configuration if needed.
    
+                    You can also fine-tune the associated topic’s configuration if needed.
    
                 
    Example for disabling fault-tolerance:
    
                 
    import org.apache.kafka.streams.state.StoreBuilder;
 import org.apache.kafka.streams.state.Stores;
@@ -326,14 +326,14 @@ 
    
                     
    If the changelog is disabled then the attached state store is no longer fault tolerant and it can’t have any standby replicas.
    
                 
    
                 
    Here is an example for enabling fault tolerance, with additional changelog-topic configuration:
-                    You can add any log config from kafka.log.LogConfig.
+                    You can add any log config from kafka.log.LogConfig.
                     Unrecognized configs will be ignored.
    
                 
    import org.apache.kafka.streams.state.StoreBuilder;
 import org.apache.kafka.streams.state.Stores;
 
 Map<String, String> changelogConfig = new HashMap();
 // override min.insync.replicas
-changelogConfig.put("min.insyc.replicas", "1")
+changelogConfig.put(TopicConfig.MIN_IN_SYNC_REPLICAS_CONFIG, "1")
 
 StoreBuilder<KeyValueStore<String, Long>> countStoreSupplier = Stores.keyValueStoreBuilder(
   Stores.persistentKeyValueStore("Counts"),
diff --git a/docs/streams/developer-guide/security.html b/docs/streams/developer-guide/security.html
index 821a34e44002c..29b7c2b8464c2 100644
--- a/docs/streams/developer-guide/security.html
+++ b/docs/streams/developer-guide/security.html
@@ -40,7 +40,7 @@
                 
  • Security example
    • 
             
         
    
-        
    Kafka Streams natively integrates with the Kafka’s security features and supports all of the
+        
    Kafka Streams natively integrates with the Kafka’s security features and supports all of the
             client-side security features in Kafka.  Streams leverages the Java Producer and Consumer API.
    
         
    To secure your Stream processing applications, configure the security settings in the corresponding Kafka producer
             and consumer clients, and then specify the corresponding configuration settings in your Kafka Streams application.
    
@@ -61,7 +61,7 @@
                 that only specific applications are allowed to read from a Kafka topic.  You can also restrict write access to Kafka
                 topics to prevent data pollution or fraudulent activities.

    -

    For more information about the security features in Apache Kafka, see Kafka Security.

    +

    For more information about the security features in Apache Kafka, see Kafka Security.

    Required ACL setting for secure Kafka clusters

    Kafka clusters can use ACLs to control access to resources (like the ability to create topics), and for such clusters each client, diff --git a/docs/streams/developer-guide/testing.html b/docs/streams/developer-guide/testing.html index 026b02b6d3a52..5be2255dc8bd3 100644 --- a/docs/streams/developer-guide/testing.html +++ b/docs/streams/developer-guide/testing.html @@ -92,30 +92,23 @@

    - To verify the output, the test driver produces ProducerRecords with key and value type - byte[]. - For result verification, you can specify corresponding deserializers when reading the output record from - the driver. - 

    -ProducerRecord<String, Integer> outputRecord = testDriver.readOutput("output-topic", new StringDeserializer(), new LongDeserializer());
-
    -

    - For result verification, you can use OutputVerifier. - It offers helper methods to compare only certain parts of the result record: - for example, you might only care about the key and value, but not the timestamp of the result record. + To verify the output, you can use TestOutputTopic + where you configure the topic and the corresponding deserializers during initialization. + It offers helper methods to read only certain parts of the result records or the collection of records. + For example, you can validate returned KeyValue with standard assertions + if you only care about the key and value, but not the timestamp of the result record. 

    -OutputVerifier.compareKeyValue(outputRecord, "key", 42L); // throws AssertionError if key or value does not match
+TestOutputTopic<String, Long> outputTopic = testDriver.createOutputTopic("result-topic", stringSerde.deserializer(), longSerde.deserializer());
+assertThat(outputTopic.readKeyValue(), equalTo(new KeyValue<>("key", 42L)));

    TopologyTestDriver supports punctuations, too. @@ -124,7 +117,7 @@

    Example

    @Test public void shouldFlushStoreForFirstInput() { - testDriver.pipeInput(recordFactory.create("input-topic", "a", 1L, 9999L)); - OutputVerifier.compareKeyValue(testDriver.readOutput("result-topic", stringDeserializer, longDeserializer), "a", 21L); - Assert.assertNull(testDriver.readOutput("result-topic", stringDeserializer, longDeserializer)); + inputTopic.pipeInput("a", 1L); + assertThat(outputTopic.readKeyValue(), equalTo(new KeyValue<>("a", 21L))); + assertThat(outputTopic.isEmpty(), is(true)); } @Test public void shouldNotUpdateStoreForSmallerValue() { - testDriver.pipeInput(recordFactory.create("input-topic", "a", 1L, 9999L)); - Assert.assertThat(store.get("a"), equalTo(21L)); - OutputVerifier.compareKeyValue(testDriver.readOutput("result-topic", stringDeserializer, longDeserializer), "a", 21L); - Assert.assertNull(testDriver.readOutput("result-topic", stringDeserializer, longDeserializer)); + inputTopic.pipeInput("a", 1L); + assertThat(store.get("a"), equalTo(21L)); + assertThat(outputTopic.readKeyValue(), equalTo(new KeyValue<>("a", 21L))); + assertThat(outputTopic.isEmpty(), is(true)); } @Test public void shouldNotUpdateStoreForLargerValue() { - testDriver.pipeInput(recordFactory.create("input-topic", "a", 42L, 9999L)); - Assert.assertThat(store.get("a"), equalTo(42L)); - OutputVerifier.compareKeyValue(testDriver.readOutput("result-topic", stringDeserializer, longDeserializer), "a", 42L); - Assert.assertNull(testDriver.readOutput("result-topic", stringDeserializer, longDeserializer)); + inputTopic.pipeInput("a", 42L); + assertThat(store.get("a"), equalTo(42L)); + assertThat(outputTopic.readKeyValue(), equalTo(new KeyValue<>("a", 42L))); + assertThat(outputTopic.isEmpty(), is(true)); } @Test public void shouldUpdateStoreForNewKey() { - testDriver.pipeInput(recordFactory.create("input-topic", "b", 21L, 9999L)); - Assert.assertThat(store.get("b"), equalTo(21L)); - OutputVerifier.compareKeyValue(testDriver.readOutput("result-topic", stringDeserializer, longDeserializer), "a", 21L); - OutputVerifier.compareKeyValue(testDriver.readOutput("result-topic", stringDeserializer, longDeserializer), "b", 21L); - Assert.assertNull(testDriver.readOutput("result-topic", stringDeserializer, longDeserializer)); + inputTopic.pipeInput("b", 21L); + assertThat(store.get("b"), equalTo(21L)); + assertThat(outputTopic.readKeyValue(), equalTo(new KeyValue<>("a", 21L))); + assertThat(outputTopic.readKeyValue(), equalTo(new KeyValue<>("b", 21L))); + assertThat(outputTopic.isEmpty(), is(true)); } @Test public void shouldPunctuateIfEvenTimeAdvances() { - testDriver.pipeInput(recordFactory.create("input-topic", "a", 1L, 9999L)); - OutputVerifier.compareKeyValue(testDriver.readOutput("result-topic", stringDeserializer, longDeserializer), "a", 21L); + final Instant recordTime = Instant.now(); + inputTopic.pipeInput("a", 1L, recordTime); + assertThat(outputTopic.readKeyValue(), equalTo(new KeyValue<>("a", 21L))); - testDriver.pipeInput(recordFactory.create("input-topic", "a", 1L, 9999L)); - Assert.assertNull(testDriver.readOutput("result-topic", stringDeserializer, longDeserializer)); + inputTopic.pipeInput("a", 1L, recordTime); + assertThat(outputTopic.isEmpty(), is(true)); - testDriver.pipeInput(recordFactory.create("input-topic", "a", 1L, 10000L)); - OutputVerifier.compareKeyValue(testDriver.readOutput("result-topic", stringDeserializer, longDeserializer), "a", 21L); - Assert.assertNull(testDriver.readOutput("result-topic", stringDeserializer, longDeserializer)); + inputTopic.pipeInput("a", 1L, recordTime.plusSeconds(10L)); + assertThat(outputTopic.readKeyValue(), equalTo(new KeyValue<>("a", 21L))); + assertThat(outputTopic.isEmpty(), is(true)); } @Test public void shouldPunctuateIfWallClockTimeAdvances() { - testDriver.advanceWallClockTime(60000); - OutputVerifier.compareKeyValue(testDriver.readOutput("result-topic", stringDeserializer, longDeserializer), "a", 21L); - Assert.assertNull(testDriver.readOutput("result-topic", stringDeserializer, longDeserializer)); + testDriver.advanceWallClockTime(Duration.ofSeconds(60)); + assertThat(outputTopic.readKeyValue(), equalTo(new KeyValue<>("a", 21L))); + assertThat(outputTopic.isEmpty(), is(true)); } public class CustomMaxAggregatorSupplier implements ProcessorSupplier<String, Long> { diff --git a/docs/streams/developer-guide/write-streams.html b/docs/streams/developer-guide/write-streams.html index ad320fadd932f..6844f84ef8961 100644 --- a/docs/streams/developer-guide/write-streams.html +++ b/docs/streams/developer-guide/write-streams.html @@ -53,7 +53,7 @@

    Libraries and Maven artifacts

    This section lists the Kafka Streams related libraries that are available for writing your Kafka Streams applications.

    You can define dependencies on the following libraries for your Kafka Streams applications.

    - +
    diff --git a/docs/streams/index.html b/docs/streams/index.html index 193a7b2798993..6872297c623c3 100644 --- a/docs/streams/index.html +++ b/docs/streams/index.html @@ -275,7 +275,7 @@

    Hello Kafka Streams

    val wordCounts: KTable[String, Long] = textLines .flatMapValues(textLine => textLine.toLowerCase.split("\\W+")) .groupBy((_, word) => word) - .count(Materialized.as("counts-store")) + .count()(Materialized.as("counts-store")) wordCounts.toStream.to("WordsWithCountsTopic") val streams: KafkaStreams = new KafkaStreams(builder.build(), props) diff --git a/docs/streams/quickstart.html b/docs/streams/quickstart.html index 717fe4699c62e..4c0d0c92b732f 100644 --- a/docs/streams/quickstart.html +++ b/docs/streams/quickstart.html @@ -165,10 +165,10 @@

    Step 3 
     > bin/kafka-topics.sh --bootstrap-server localhost:9092 --describe
 
-Topic:streams-plaintext-input	PartitionCount:1	ReplicationFactor:1	Configs:
-    Topic: streams-plaintext-input	Partition: 0	Leader: 0	Replicas: 0	Isr: 0
-Topic:streams-wordcount-output	PartitionCount:1	ReplicationFactor:1	Configs:cleanup.policy=compact
+Topic:streams-wordcount-output	PartitionCount:1	ReplicationFactor:1	Configs:cleanup.policy=compact,segment.bytes=1073741824
 	Topic: streams-wordcount-output	Partition: 0	Leader: 0	Replicas: 0	Isr: 0
+Topic:streams-plaintext-input	PartitionCount:1	ReplicationFactor:1	Configs:segment.bytes=1073741824
+	Topic: streams-plaintext-input	Partition: 0	Leader: 0	Replicas: 0	Isr: 0

    Step 4: Start the Wordcount Application

    @@ -276,10 +276,10 @@

    Step 5 Here the last printed lines kafka 2 and streams 2 indicate updates to the keys kafka and streams whose counts have been incremented from 1 to 2. Whenever you write further input messages to the input topic, you will observe new messages being added to the streams-wordcount-output topic, representing the most recent word counts as computed by the WordCount application. -Let's enter one final input text line "join kafka summit" and hit <RETURN> in the console producer to the input topic streams-wordcount-input before we wrap up this quickstart: +Let's enter one final input text line "join kafka summit" and hit <RETURN> in the console producer to the input topic streams-plaintext-input before we wrap up this quickstart: 
    -> bin/kafka-console-producer.sh --broker-list localhost:9092 --topic streams-wordcount-input
+> bin/kafka-console-producer.sh --broker-list localhost:9092 --topic streams-plaintext-input
 all streams lead to kafka
 hello kafka streams
 join kafka summit
diff --git a/docs/streams/tutorial.html b/docs/streams/tutorial.html
index 0006e3ef53515..eeb90d6abc858 100644
--- a/docs/streams/tutorial.html
+++ b/docs/streams/tutorial.html
@@ -74,8 +74,12 @@ 
    Setting up a Maven

    - The pom.xml file included in the project already has the Streams dependency defined, - and there are already several example programs written with Streams library under src/main/java. + The pom.xml file included in the project already has the Streams dependency defined. + Note, that the generated pom.xml targets Java 8, and does not work with higher Java versions. +

    + +

    + There are already several example programs written with Streams library under src/main/java. Since we are going to start writing such programs from scratch, we can now delete these examples:

    @@ -86,7 +90,7 @@

    Setting up a Maven

    Writing a first Streams application: Pipe

    - It's coding time now! Feel free to open your favorite IDE and import this Maven project, or simply open a text editor and create a java file under src/main/java. + It's coding time now! Feel free to open your favorite IDE and import this Maven project, or simply open a text editor and create a java file under src/main/java/myapps. Let's name it Pipe.java: 
    @@ -604,7 +608,7 @@ 
    Writing a th
                       .groupBy((key, value) -> value)
                       .count(Materialized.<String, Long, KeyValueStore<Bytes, byte[]>>as("counts-store"))
                       .toStream()
-                      .to("streams-wordcount-output", Produced.with(Serdes.String(), Serdes.Long());
+                      .to("streams-wordcount-output", Produced.with(Serdes.String(), Serdes.Long()));
 
                 final Topology topology = builder.build();
                 final KafkaStreams streams = new KafkaStreams(topology, props);
diff --git a/docs/streams/upgrade-guide.html b/docs/streams/upgrade-guide.html
index 9071dc2eefa15..4d47f709ced8f 100644
--- a/docs/streams/upgrade-guide.html
+++ b/docs/streams/upgrade-guide.html
@@ -34,10 +34,10 @@ 
    Upgrade Guide and API Changes
    
     
 
     
    
-        Upgrading from any older version to {{fullDotVersion}} is possible: (1) if you are upgrading from 2.0.x to {{fullDotVersion}} then a single rolling bounce is needed to swap in the new jar,
-        (2) if you are upgrading from older versions than 2.0.x in the online mode, you would need two rolling bounces where
-        the first rolling bounce phase you need to set config upgrade.from="older version" (possible values are "0.10.0", "0.10.1", "0.10.2", "0.11.0", "1.0", and "1.1")
-        (cf. KIP-268):
+        Upgrading from any older version to {{fullDotVersion}} is possible: you will need to do two rolling bounces, where during the first rolling bounce phase you set the config upgrade.from="older version"
+        (possible values are "0.10.0" - "2.3") and during the second you remove it. This is required to safely upgrade to the new cooperative rebalancing protocol of the embedded consumer. Note that you will remain using the old eager
+        rebalancing protocol if you skip or delay the second rolling bounce, but you can safely switch over to cooperative at any time once the entire group is on 2.4+ by removing the config value and bouncing. For more details please refer to
+        KIP-429:
     
    
     
      
         
    •  prepare your application instances for a rolling bounce and make sure that config upgrade.from is set to the version from which it is being upgrade.
      • 
@@ -53,11 +53,15 @@ 
      Upgrade Guide and API Changes
      
     
    
 
     
    
-        Note, that a brokers must be on version 0.10.1 or higher to run a Kafka Streams application version 0.10.1 or higher;
-        On-disk message format must be 0.10 or higher to run a Kafka Streams application version 1.0 or higher.
+        To run a Kafka Streams application version 2.2.1, 2.3.0, or higher a broker version 0.11.0 or higher is required
+        and the on-disk message format must be 0.11 or higher.
+        Brokers must be on version 0.10.1 or higher to run a Kafka Streams application version 0.10.1 to 2.2.0.
+        Additionally, on-disk message format must be 0.10 or higher to run a Kafka Streams application version 1.0 to 2.2.0.
         For Kafka Streams 0.10.0, broker version 0.10.0 or higher is required.
     
    
 
+    
    Since 2.2.0 release, Kafka Streams depends on a RocksDBs version that requires MacOS 10.13 or higher.
    
+
     
    
         Another important thing to keep in mind: in deprecated KStreamBuilder class, when a KTable is created from a source topic via KStreamBuilder.table(), its materialized state store
         will reuse the source topic as its changelog topic for restoring, and will disable logging to avoid appending new updates to the source topic; in the StreamsBuilder class introduced in 1.0, this behavior was changed
@@ -68,14 +72,170 @@ 
    Upgrade Guide and API Changes
    
         More details about the new config StreamsConfig#TOPOLOGY_OPTIMIZATION can be found in KIP-295.
     
    
 
+    
    Streams API changes in 2.4.0
    
+    
         
+         As of 2.4.0 Kafka Streams offers a KTable-KTable foreign-key join (as per KIP-213). 
+         This joiner allows for records to be joined between two KTables with different keys. 
+         Both INNER and LEFT foreign-key joins 
+         are supported.
+    
    
+    
    
+        In the 2.4 release, you now can name all operators in a Kafka Streams DSL topology via
+        KIP-307.
+        Giving your operators meaningful names makes it easier to understand the topology 
+        description (Topology#describe()#toString()) and 
+        understand the full context of what your Kafka Streams application is doing.
+        
    
+        There are new overloads on most KStream and KTable methods 
+        that accept a Named object.  Typically you'll provide a name for the DSL operation by 
+        using Named.as("my operator name").  Naming of repartition topics for aggregation 
+        operations will still use Grouped and join operations will use 
+        either Joined or the new StreamJoined object.
+
+    
    
+    
    
+        Before the 2.4.0 version of Kafka Streams, users of the DSL could not name the state stores involved in a stream-stream join.
+        If users changed their topology and added a operator before the
+        join, the internal names of the state stores would shift, requiring an application reset when redeploying.
+        In the 2.4.0 release, Kafka Streams adds the StreamJoined
+        class, which gives users the ability to name the join processor, repartition topic(s) (if a repartition is required),
+        and the state stores involved in the join.  Also, by naming the state stores, the changelog topics
+        backing the state stores are named as well.  It's important to note that naming the stores
+        will not make them queryable via Interactive Queries.
+        
    
+        Another feature delivered by StreamJoined is that you can now configure the type of state store used in the join.
+        You can elect to use in-memory stores or custom state stores for a stream-stream join.  Note that the provided stores
+        will not be available for querying via Interactive Queries.  With the addition
+        of StreamJoined, stream-stream join operations
+        using Joined have been deprecated. Please switch over to stream-stream join methods using the
+        new overloaded methods.  You can get more details from
+        KIP-479.
+    
    
+    
    
+        With the introduction of incremental cooperative rebalancing, Streams no longer requires all tasks be revoked at the beginning of a rebalance. Instead, at the completion of the rebalance only those tasks which are to be migrated to another consumer
+        for overall load balance will need to be closed and revoked. This changes the semantics of the StateListener a bit, as it will not necessarily transition to REBALANCING at the beginning of a rebalance anymore. Note that
+        this means IQ will now be available at all times except during state restoration, including while a rebalance is in progress. If restoration is occurring when a rebalance begins, we will continue to actively restore the state stores and/or process
+        standby tasks during a cooperative rebalance. Note that with this new rebalancing protocol, you may sometimes see a rebalance be followed by a second short rebalance that ensures all tasks are safely distributed. For details on please see
+        KIP-429.
+    
    
+
+    
    
+        The 2.4.0 release contains newly added and reworked metrics.
+        KIP-444
+        adds new client level (i.e., KafkaStreams instance level) metrics to the existing
+        thread-level, task-level, and processor-/state-store-level metrics.
+        For a full list of available client level metrics, see the
+        KafkaStreams monitoring
+        section in the operations guide.
+        
    
+        Furthermore, RocksDB metrics are exposed via
+        KIP-471.
+        For a full list of available RocksDB metrics, see the
+        RocksDB monitoring
+        section in the operations guide.
+    
    
+
+    
    
+        Kafka Streams test-utils got improved via
+        KIP-470
+        to simplify the process of using TopologyTestDriver to test your application code.
+        We deprecated ConsumerRecordFactory, TopologyTestDriver#pipeInput(),
+        OutputVerifier, as well as TopologyTestDriver#readOutput() and replace them with
+        TestInputTopic and TestOutputTopic, respectively.
+        We also introduced a new class TestRecord that simplifies assertion code.
+        For full details see the
+        Testing section in the developer guide.
+    
    
+
+    
    
+        In 2.4.0, we deprecated WindowStore#put(K key, V value) that should never be used.
+        Instead the existing WindowStore#put(K key, V value, long windowStartTimestamp) should be used
+        (KIP-474).
+    
    
+
+    
    
+        Furthermore, the PartitionGrouper interface and its corresponding configuration parameter
+        partition.grouper were deprecated
+        (KIP-528)
+        and will be removed in the next major release (KAFKA-7785.
+        Hence, this feature won't be supported in the future any longer and you need to updated your code accordingly.
+        If you use a custom PartitionGrouper and stop to use it, the created tasks might change.
+        Hence, you will need to reset your application to upgrade it.
+    
+
     
    Streams API changes in 2.3.0
    
+
+    
    Version 2.3.0 adds the Suppress operator to the kafka-streams-scala Ktable API.
    
+
+    
    
+        As of 2.3.0 Streams now offers an in-memory version of the window (KIP-428)
+        and the session (KIP-445) store, in addition to the persistent ones based on RocksDB.
+        The new public interfaces inMemoryWindowStore() and inMemorySessionStore() are added to Stores and provide the built-in in-memory window or session store.
+    
    
+
+    
    
+        As of 2.3.0 we've updated how to turn on optimizations. Now to enable optimizations, you need to do two things.
+        First add this line to your properties properties.setProperty(StreamsConfig.TOPOLOGY_OPTIMIZATION, StreamsConfig.OPTIMIZE);, as you have done before.
+        Second, when constructing your KafkaStreams instance, you'll need to pass your configuration properties when building your
+        topology by using the overloaded StreamsBuilder.build(Properties) method.
+        For example KafkaStreams myStream = new KafkaStreams(streamsBuilder.build(properties), properties).
+    
    
+
+    
    
+        In 2.3.0 we have added default implementation to close() and configure() for Serializer,
+        Deserializer and Serde so that they can be implemented by lambda expression.
+        For more details please read KIP-331.
+    
    
+
+    
    
+        To improve operator semantics, new store types are added that allow storing an additional timestamp per key-value pair or window.
+        Some DSL operators (for example KTables) are using those new stores.
+        Hence, you can now retrieve the last update timestamp via Interactive Queries if you specify
+        TimestampedKeyValueStoreType or TimestampedWindowStoreType as your QueryableStoreType.
+        While this change is mainly transparent, there are some corner cases that may require code changes:
+        Caution: If you receive an untyped store and use a cast, you might need to update your code to cast to the correct type.
+        Otherwise, you might get an exception similar to
+        java.lang.ClassCastException: class org.apache.kafka.streams.state.ValueAndTimestamp cannot be cast to class YOUR-VALUE-TYPE
+        upon getting a value from the store.
+        Additionally, TopologyTestDriver#getStateStore() only returns non-built-in stores and throws an exception if a built-in store is accessed.
+        For more details please read KIP-258.
+    
    
+
+    
    
+        To improve type safety, a new operator KStream#flatTransformValues is added.
+        For more details please read KIP-313.
+    
    
+
+    
    
+        Kafka Streams used to set the configuration parameter max.poll.interval.ms to Integer.MAX_VALUE.
+        This default value is removed and Kafka Streams uses the consumer default value now.
+        For more details please read KIP-442.
+    
    
+
+    
    
+        Default configuration for repartition topic was changed:
+        The segment size for index files (segment.index.bytes) is no longer 50MB, but uses the cluster default.
+        Similarly, the configuration segment.ms in no longer 10 minutes, but uses the cluster default configuration.
+        Lastly, the retention period (retention.ms) is changed from Long.MAX_VALUE to -1 (infinite).
+        For more details please read KIP-443.
+    
    
+
     
    
-        As of 2.3.0 Streams now offers an in-memory version of the window store, in addition to the persistent one based on RocksDB. The new public interface inMemoryWindowStore() is added to Stores that provides a built-in in-memory window store.
+        To avoid memory leaks, RocksDBConfigSetter has a new close() method that is called on shutdown.
+        Users should implement this method to release any memory used by RocksDB config objects, by closing those objects.
+        For more details please read KIP-453.
     
    
 
     
    
-        In 2.3.0 we have added default implementation to close() and configure() for Serializer, Deserializer and Serde so that they can be
-        implemented by lambda expression. For more details please read KIP-331.
+        RocksDB dependency was updated to version 5.18.3.
+        The new version allows to specify more RocksDB configurations, including WriteBufferManager which helps to limit RocksDB off-heap memory usage.
+        For more details please read KAFKA-8215.
+    
    
+
+    
    Notable changes in Kafka Streams 2.2.1
    
+    
    
+        As of Kafka Streams 2.2.1 a message format 0.11 or higher is required;
+        this implies that brokers must be on version 0.11.0 or higher.
     
    
 
     
    Streams API changes in 2.2.0
    
@@ -188,6 +348,13 @@ 
    Streams API
         details, see KIP-312
     
    
 
+    
    
+        We are introducing static membership towards Kafka Streams user. This feature reduces unnecessary rebalances during normal application upgrades or rolling bounces.
+        For more details on how to use it, checkout static membership design.
+        Note, Kafka Streams uses the same ConsumerConfig#GROUP_INSTANCE_ID_CONFIG, and you only need to make sure it is uniquely defined across
+        different stream instances in one application.
+    
    
+
     
    Streams API changes in 2.0.0
    
     
    
         In 2.0.0 we have added a few new APIs on the ReadOnlyWindowStore interface (for details please read Streams API changes below).
diff --git a/docs/toc.html b/docs/toc.html
index a8189216e9175..1223d20d3c25e 100644
--- a/docs/toc.html
+++ b/docs/toc.html
@@ -35,7 +35,7 @@
                 
  • 2.2 Consumer API
                 
  • 2.3 Streams API
                 
  • 2.4 Connect API
-                
  • 2.5 AdminClient API
+                
  • 2.5 Admin API
             
         
    • 
         
  • 3. Configuration
diff --git a/docs/upgrade.html b/docs/upgrade.html
index 2ba7fd3915737..6b3775cd3d5d0 100644
--- a/docs/upgrade.html
+++ b/docs/upgrade.html
@@ -19,7 +19,247 @@